diff --git a/.gitignore b/.gitignore index 059ddda3f..923cbc074 100644 --- a/.gitignore +++ b/.gitignore @@ -1,16 +1,249 @@ -*.egg-info +## vLLM Ascend Ignore +# VSCode +.vscode/ + +# egg-info +vllm_ascend.egg-info/ + +# DS Store +.DS_Store + +# Linting +actionlint +shellcheck*/ + + +# Python gitignore +## Adapted from: +## https://github.com/github/gitignore/blob/main/Python.gitignore + +# Byte-compiled / optimized / DLL files __pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python build/ +develop-eggs/ dist/ -vllm_fl/_version.py +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt -# Coverage +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ .coverage .coverage.* -htmlcov/ -coverage*.json +.cache +nosetests.xml coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/latest/usage/project/#working-with-version-control +.pdm.toml +.pdm-python +.pdm-build/ + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +.idea/ + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +kernel_meta/ + +# version file generated by setuptools-scm +/vllm_ascend/_version.py +# build info file generated by setup.py +/vllm_ascend/_build_info.py +/vllm_ascend/include/ + +# benchmark results generated by run-performance-benchmarks.sh +/benchmarks/results/ + +# _cann_ops_custom generated by build_aclnn.sh +/vllm_ascend/_cann_ops_custom/* +!/vllm_ascend/_cann_ops_custom/.gitkeep + +# generated by CANN +fusion_result.json +csrc/output/ + +# claude code skills +.claude/skills/* + +# vllm-plugin-FL generated artifacts +vllm_plugin_fl.egg-info/ +/vllm_fl/_cann_ops_custom/* +!/vllm_fl/_cann_ops_custom/.gitkeep +vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib/*.so + +# CANN framework operator build directories +csrc/ascend/build/ +csrc/ascend/build_out/ + +#benchmark result +benchmark_results + +#documents design +docs/* + +# Exception: these are source patches, not build artifacts. +# The parent directories must also be un-ignored because the generic +# 'build/' rule above ignores every directory named 'build'. +!csrc/ascend/cmake/third_party/build/ +!csrc/ascend/cmake/third_party/build/modules/ +!csrc/ascend/cmake/third_party/build/modules/patch/ +!csrc/ascend/cmake/third_party/build/modules/patch/** -# Reports -report*.json -test-results-* +# CANN framework operator downloaded third-party libraries. +# Keep only the catlass and pto-isa source submodules. +csrc/ascend/third_party/abseil-cpp/ +csrc/ascend/third_party/ascend_protobuf/ +csrc/ascend/third_party/json/ +csrc/ascend/third_party/makeself/ +csrc/ascend/third_party/pkg/ +# csrc/ascend/third_party/pto-isa/ +# csrc/ascend/third_party/catlass diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 000000000..93bbfe1a2 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,8 @@ +[submodule "csrc/ascend/third_party/catlass"] + path = csrc/ascend/third_party/catlass + url = https://gitcode.com/cann/catlass.git + branch = master +[submodule "csrc/ascend/third_party/pto-isa"] + path = csrc/ascend/third_party/pto-isa + url = https://gitcode.com/cann/pto-isa.git + branch = main diff --git a/README.md b/README.md index a7a323b8a..d5151fa61 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,7 @@ In theory, vllm-plugin-FL can support all models available in vLLM, as long as n | MiniCPM-o 4.5 | Supported | [example](./examples/minicpm/) | | GLM-5 | Supported | [example](./examples/glm_5_offline_inference.py) | | Qwen3.5-35B-A3B | Supported | [example](./examples/glm_5_offline_inference.py) | +| Qwen3.6-27B | Supported | [Ascend patch](./vllm_fl/dispatch/backends/vendor/ascend/patches/patch_qwen3_5.py) | | BAAI/bge-m3 | Supported | [implementation](./vllm_fl/models/bge_m3.py) | ### Supported Chips diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index e69de29bb..581f2fda0 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -0,0 +1,152 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-FL project +# +# vLLM-FL extension dispatcher. +# For Ascend, this file directly builds the torch extension (_C_ascend) and +# the kernels shared library. The CANN framework operators are built +# independently from csrc/ascend/ via build_aclnn.sh. + +cmake_minimum_required(VERSION 3.26) +project(vllm_fl_extensions LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +if(NOT DEFINED VLLM_VENDOR) + if(DEFINED ENV{VLLM_VENDOR}) + set(VLLM_VENDOR $ENV{VLLM_VENDOR}) + endif() +endif() + +if(NOT VLLM_VENDOR) + message(FATAL_ERROR + "VLLM_VENDOR is required but not specified.\n" + "Set VLLM_VENDOR=cuda or VLLM_VENDOR=ascend." + ) +endif() + +set(SUPPORTED_VENDORS cuda ascend) +if(NOT VLLM_VENDOR IN_LIST SUPPORTED_VENDORS) + message(FATAL_ERROR + "Unsupported vendor: ${VLLM_VENDOR}\n" + "Supported vendors: ${SUPPORTED_VENDORS}" + ) +endif() + +message(STATUS "vLLM-FL Extensions vendor: ${VLLM_VENDOR}") + +if(VLLM_PYTHON_EXECUTABLE) + set(Python_EXECUTABLE ${VLLM_PYTHON_EXECUTABLE}) +endif() + +find_package(Python REQUIRED COMPONENTS Interpreter Development.Module) +message(STATUS "Python: ${Python_EXECUTABLE} (${Python_VERSION})") + +execute_process( + COMMAND ${Python_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)" + OUTPUT_VARIABLE TORCH_CMAKE_PREFIX + OUTPUT_STRIP_TRAILING_WHITESPACE +) +list(APPEND CMAKE_PREFIX_PATH ${TORCH_CMAKE_PREFIX}) + +find_package(Torch REQUIRED) +message(STATUS "PyTorch: ${Torch_VERSION}") + +if(VLLM_VENDOR STREQUAL "ascend") + set(ASCEND_CANN_PACKAGE_PATH ${ASCEND_HOME_PATH}) + + if(NOT DEFINED TORCH_NPU_PATH) + message(FATAL_ERROR + "TORCH_NPU_PATH is required for Ascend builds. " + "It is normally supplied by setup.py." + ) + endif() + message(STATUS "torch_npu package path: ${TORCH_NPU_PATH}") + include_directories(${TORCH_NPU_PATH}/include) + link_directories(${TORCH_NPU_PATH}/lib) + + # Locate the AscendC kernel cmake helpers shipped with the CANN toolkit. + if(EXISTS ${ASCEND_HOME_PATH}/tools/tikcpp/ascendc_kernel_cmake) + set(ASCENDC_CMAKE_DIR ${ASCEND_HOME_PATH}/tools/tikcpp/ascendc_kernel_cmake) + elseif(EXISTS ${ASCEND_HOME_PATH}/compiler/tikcpp/ascendc_kernel_cmake) + set(ASCENDC_CMAKE_DIR ${ASCEND_HOME_PATH}/compiler/tikcpp/ascendc_kernel_cmake) + elseif(EXISTS ${ASCEND_HOME_PATH}/ascendc_devkit/tikcpp/samples/cmake) + set(ASCENDC_CMAKE_DIR ${ASCEND_HOME_PATH}/ascendc_devkit/tikcpp/samples/cmake) + else() + message(FATAL_ERROR "ascendc_kernel_cmake does not exist, please check whether the cann package is installed.") + endif() + include(${ASCENDC_CMAKE_DIR}/ascendc.cmake) + + set(VLLM_FL_ASCEND_SRCS + ascend/camem_allocator.cpp + ascend/torch_binding.cpp + ascend/torch_binding_meta.cpp + ascend/aclnn_torch_adapter/op_api_common.h + ascend/aclnn_torch_adapter/NPUBridge.cpp + ascend/aclnn_torch_adapter/NPUBridge.h + ascend/aclnn_torch_adapter/NPUStorageImpl.cpp + ascend/aclnn_torch_adapter/NPUStorageImpl.h + # Framework operator torch adapters (op_host/op_kernel/ are built by the + # CANN toolchain and are present here for reference / vendor packaging). + ascend/attention/fused_gdn_gating/fused_gdn_gating_torch_adpt.h + ascend/attention/kv_quant_sparse_flash_attention/kv_quant_sparse_flash_attention_torch_adpt.h + ascend/attention/lightning_indexer/lightning_indexer_torch_adpt.h + ascend/attention/lightning_indexer_quant/lightning_indexer_quant_torch_adpt.h + ascend/attention/lightning_indexer_vllm/lightning_indexer_vllm_torch_adpt.h + ascend/attention/ngram_spec_decode/ngram_spec_decode_torch_adpt.h + ascend/attention/recurrent_gated_delta_rule/recurrent_gated_delta_rule_torch_adpt.h + ascend/attention/sparse_flash_attention/sparse_flash_attention_torch_adpt.h + ascend/attention/store_kv_block/store_kv_block_torch_adpt.h + ascend/gmm/grouped_matmul_swiglu_quant_v2/grouped_matmul_swiglu_quant_v2_torch_adpt.h + ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/grouped_matmul_swiglu_quant_torch_adpt.h + ascend/mc2/dispatch_ffn_combine/dispatch_ffn_combine_torch_adpt.h + ascend/mc2/matmul_allreduce_add_rmsnorm/matmul_allreduce_add_rmsnorm_torch_adpt.h + ascend/moe/add_rms_norm_bias/add_rms_norm_bias_torch_adpt.h + ascend/moe/apply_top_k_top_p_custom/apply_top_k_top_p_custom_torch_adpt.h + ascend/moe/moe_gating_top_k/moe_gating_top_k_torch_adpt.h + ascend/moe/moe_init_routing_custom/moe_init_routing_custom_torch_adpt.h + ) + + file(GLOB VLLM_FL_ASCEND_KERNELS ascend/kernels/*.cpp) + ascendc_library(vllm_fl_kernels SHARED ${VLLM_FL_ASCEND_KERNELS}) + + Python_add_library(_C_ascend MODULE WITH_SOABI ${VLLM_FL_ASCEND_SRCS}) + target_compile_definitions(_C_ascend PRIVATE TORCH_EXTENSION_NAME=_C_ascend) + + target_include_directories(_C_ascend PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/ascend + ${CMAKE_CURRENT_SOURCE_DIR}/ascend/aclnn_torch_adapter + ${CMAKE_CURRENT_SOURCE_DIR} + ${TORCH_INCLUDE_DIRS} + ) + + target_link_libraries(_C_ascend PRIVATE + ${TORCH_LIBRARIES} + torch_npu + vllm_fl_kernels + ascendcl + opapi + tiling_api + register + platform + ascendalog + dl + ) + + target_compile_options(_C_ascend PRIVATE -O3) + target_link_options(_C_ascend PRIVATE "-Wl,-rpath,$ORIGIN") + + install(TARGETS _C_ascend vllm_fl_kernels LIBRARY DESTINATION vllm_fl COMPONENT _C_ascend) + +elseif(VLLM_VENDOR STREQUAL "cuda") + add_subdirectory(cuda) +else() + message(FATAL_ERROR "Unsupported vendor: ${VLLM_VENDOR}") +endif() + +# Optional PTO chunk-GDN megakernel pre-compilation (Ascend only). +option(BUILD_PTO_CHUNK_GDN "Pre-compile PTO GDN megakernel via Bisheng" OFF) +if(VLLM_VENDOR STREQUAL "ascend" AND BUILD_PTO_CHUNK_GDN) + add_subdirectory(ascend/pto_chunk_gdn EXCLUDE_FROM_ALL) +endif() diff --git a/csrc/ascend/CMakeLists.txt b/csrc/ascend/CMakeLists.txt index e69de29bb..7003bdcd6 100644 --- a/csrc/ascend/CMakeLists.txt +++ b/csrc/ascend/CMakeLists.txt @@ -0,0 +1,659 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +cmake_minimum_required(VERSION 3.16) + +project(cann_ops-transformer) + +option(BUILD_OPEN_PROJECT "Build open ascend ops project." ON) +option(BUILD_OPS_RTY_KERNEL "Build return yellow kernel." OFF) +option(ENABLE_CCACHE "Enable ccache capability" ON) +option(ENABLE_BUILT_IN "Enable built-in package" OFF) +option(ENABLE_STATIC "Enable Static" OFF) +option(ENABLE_EXPERIMENTAL "Enable experimental module" OFF) +option(ENABLE_TEST "Enable test" OFF) +option(ENABLE_UT_EXEC "Enable exec ut" OFF) +option(ENABLE_ASAN "Enable asan" OFF) +option(ENABLE_VALGRIND "Enable valgrind" OFF) +option(OP_HOST_UT "Enable ophost ut" OFF) +option(OP_API_UT "Enable opapi ut" OFF) +option(OP_GRAPH_UT "Enable graph ut" OFF) +option(OP_KERNEL_UT "Enable kernel ut" OFF) +option(OP_KERNEL_AICPU_UT "Enable aicpu kernel ut" OFF) +option(UT_TEST_ALL "Enable all ut" OFF) +option(ENABLE_OOM "Enable kernel oom" OFF) + +set(ASCEND_COMPUTE_UNIT "ascend910b" CACHE STRING "soc that need to be compiled") +set(ASCEND_OP_NAME "ALL" CACHE STRING "operators that need to be compiled") +set(ARCH_DIRECTORY "" CACHE STRING "arch directory that need to be compiled") +set(VENDOR_NAME "custom" CACHE STRING "vendor name") +set(ASCEND_ALL_COMPUTE_UNIT "ascend910b" CACHE STRING "all soc list") + +set(SOC_VERSION_LIST ascend310p ascend910b ascend910_93 ascend950 kirinx90) +set(ARCH_DIRECTORY_LIST arch22 arch32 arch32 arch35 arch32) + +if ("ascend950" IN_LIST ASCEND_COMPUTE_UNIT) +message(STATUS "build with 3~8 packages........") +set(BUILD_WITH_3_8_PACKAGE ON CACHE BOOL "build with 3~8 package and opsbase") +endif() + +foreach(SOC_VERSION ${ASCEND_COMPUTE_UNIT}) + list(FIND SOC_VERSION_LIST ${SOC_VERSION} INDEX) + if(NOT INDEX EQUAL -1) + list(GET ARCH_DIRECTORY_LIST ${INDEX} VAL) + list(APPEND ARCH_DIRECTORY ${VAL}) + else() + message(STATUS "unsupported chip type") + if ((NOT BUILD_OPS_RTY_KERNEL) AND (BUILD_OPEN_PROJECT)) + include(cmake/build_empty_package.cmake) + cpack_empty_package() + return() + endif() + endif() +endforeach() + +list(FIND ARCH_DIRECTORY "arch32" INDEX) +if(NOT INDEX EQUAL -1) + list(APPEND ARCH_DIRECTORY "arch22") +endif() + +if(PROJECT_SOURCE_DIR STREQUAL CMAKE_SOURCE_DIR) + message(STATUS "compile project with library") + option(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG "Build ops-transformer with cann pkg" ON) +else() + message(STATUS "compile project with src") + option(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG "Build ops-transformer with cann source" OFF) +endif() + +if(UNIX) + set(SYSTEM_PREFIX ${CMAKE_SYSTEM_PROCESSOR}-linux) +endif() + +#外部传参 +if(NOT ${CMAKE_BUILD_MODE} STREQUAL "FALSE") + if(ENABLE_DEBUG) + set(CMAKE_BUILD_MODE "${CMAKE_BUILD_MODE} -g") + endif() + set(COMPILE_OP_MODE ${CMAKE_BUILD_MODE}) +else() + if(ENABLE_TEST) + set(COMPILE_OP_MODE "-O0 -g") + endif() + if(ENABLE_DEBUG) + set(CMAKE_BUILD_MODE "-g") + endif() +endif() + +set(PKG_NAME transformer) +set(OPS_TRANSFORMER_DIR ${CMAKE_CURRENT_SOURCE_DIR}) +set(CMAKE_CXX_STANDARD 17 CACHE STRING "c++17 is needed for this project") +set_directory_properties(PROPERTIES + ADDITIONAL_MAKE_CLEAN_FILES "${CMAKE_BINARY_DIR}/_CPack_Packages" +) + +# Suppress warnings from catlass/tla third-party headers for CANN kernel compilation +set(VLLM_ASCEND_CANN_COMPAT_HEADER "${OPS_TRANSFORMER_DIR}/common/include/cann_compat.h") +list(APPEND OPS_COMPILE_OPTIONS -Wno-ignored-attributes) +list(APPEND OPS_COMPILE_OPTIONS -include${VLLM_ASCEND_CANN_COMPAT_HEADER}) +add_compile_options( + $<$:-include${VLLM_ASCEND_CANN_COMPAT_HEADER}> +) + +include(cmake/config.cmake) + +# Detect whether the current CANN toolchain exposes ASCEND950 in platform_ascendc::SocVersion. +# Operators that only compile against ASCEND950 are skipped automatically on older CANN. +set(ASCEND950_AVAILABLE FALSE) +set(_PLATFORM_ASCENDC_H_CANDIDATES + "${ASCEND_CANN_PACKAGE_PATH}/${SYSTEM_PREFIX}/asc/include/utils/tiling/platform/platform_ascendc.h" + "${ASCEND_CANN_PACKAGE_PATH}/asc/include/utils/tiling/platform/platform_ascendc.h" +) +foreach(_cand ${_PLATFORM_ASCENDC_H_CANDIDATES}) + if(EXISTS ${_cand}) + file(STRINGS ${_cand} _ASCEND950_LINES REGEX "ASCEND950") + if(_ASCEND950_LINES) + set(ASCEND950_AVAILABLE TRUE) + endif() + break() + endif() +endforeach() + +if(ASCEND950_AVAILABLE) + message(STATUS "ASCEND950 is available in CANN headers; building ASCEND950-only operators.") + set(ASCEND950_SKIP_OPS "") +else() + message(STATUS "ASCEND950 is NOT available in CANN headers; skipping ASCEND950-only operators.") + set(ASCEND950_SKIP_OPS + compressor + inplace_partial_rotary_mul + kv_quant_sparse_flash_attention + lightning_indexer + recurrent_gated_delta_rule + sparse_flash_attention + vllm_quant_lightning_indexer + add_rms_norm_bias + causal_conv1d + dequant_swiglu_quant + hc_post + hc_pre + hc_pre_inv_rms + hc_pre_sinkhorn + moe_gating_top_k + moe_gating_top_k_hash + scatter_nd_update_v2 + grouped_matmul_swiglu_quant_v2 + ) +endif() + +include(cmake/func.cmake) +include(cmake/third_party/json.cmake) +if (ENABLE_TEST) + include(${PROJECT_SOURCE_DIR}/cmake/third_party/gtest.cmake) +endif() +include(${OPS_ADV_CMAKE_DIR}/ut.cmake) + +if (BUILD_OPEN_PROJECT) + include(cmake/intf.cmake) + add_definitions(-DBUILD_OPEN_PROJECT) + if (BUILD_OPS_RTY_KERNEL) + message(STATUS "Build return yellow kernel.") + include(cmake/rty_obj_func.cmake) + else() + message(STATUS "Start building custom package.") + include(ExternalProject) + include(cmake/dependencies.cmake) + include(cmake/variables.cmake) + include(cmake/obj_func.cmake) + include(cmake/third_party/abseil-cpp.cmake) + include(cmake/third_party/ascend_protobuf.cmake) + + include(cmake/third_party/makeself-fetch.cmake) + include(cmake/opbuild.cmake) + include(cmake/custom_build.cmake) + message(STATUS "End building custom package.") + if (ENABLE_OPS_HOST) + gen_aclnn_with_opdef() + endif() + if (ENABLE_STATIC) + include(cmake/static.cmake) + endif() + if (ENABLE_AICPU) + include(cmake/symbol.cmake) + gen_cust_aicpu_json_symbol() + gen_cust_aicpu_kernel_symbol() + endif() + if (ENABLE_BUILT_IN) + message(STATUS "Start building built-in package.") + include(cmake/symbol.cmake) + gen_norm_symbol() + include(cmake/package.cmake) + pack_built_in() + else() + include(cmake/package.cmake) + pack_tiling_sink() + endif() + return() + endif() +else() + include(cmake/dependencies.cmake) + include(cmake/variables.cmake) + include(cmake/opbuild.cmake) + include(cmake/rty_obj_func.cmake) + include(cmake/intf_pub_linux.cmake) +endif() + +if (BUILD_OPS_RTY_KERNEL) + set(CMAKE_MODULE_PATH + ${CMAKE_MODULE_PATH} + ${CMAKE_CURRENT_LIST_DIR}/cmake/modules + ) + + set(CMAKE_PREFIX_PATH + ${CMAKE_PREFIX_PATH} + ${ASCEND_CANN_PACKAGE_PATH} + ) + + set(_op_host_aclnn_link + $ + exe_graph + register + c_sec + ) + + find_package(alog MODULE) + + if(NOT ${alog_FOUND}) + add_definitions(-DALOG_NOT_FOUND) + endif() + + add_library(op_host_aclnn SHARED EXCLUDE_FROM_ALL) + target_link_libraries(op_host_aclnn PRIVATE + ${_op_host_aclnn_link} + ) + target_compile_options(op_host_aclnn PRIVATE + $<$:-std=gnu++1z> + ) + + add_library(op_host_aclnnInner SHARED EXCLUDE_FROM_ALL) + target_link_libraries(op_host_aclnnInner PRIVATE + ${_op_host_aclnn_link} + ) + target_compile_options(op_host_aclnnInner PRIVATE + $<$:-std=gnu++1z> + ) + + add_library(op_host_aclnnExc SHARED EXCLUDE_FROM_ALL) + target_link_libraries(op_host_aclnnExc PRIVATE + ${_op_host_aclnn_link} + ) + target_compile_options(op_host_aclnnExc PRIVATE + $<$:-std=gnu++1z> + ) + + # op proto + add_library(opsproto SHARED) + target_compile_options(opsproto PRIVATE + $<$:-std=c++11> + -fvisibility=hidden + ) + target_compile_definitions(opsproto PRIVATE + LOG_CPP + PROCESS_LOG + ) + target_link_libraries(opsproto PRIVATE + $ + $ + $<$:$> + -Wl,--whole-archive + rt2_registry + -Wl,--no-whole-archive + -Wl,--no-as-needed + exe_graph + graph + graph_base + register + ascendalog + error_manager + platform + -Wl,--as-needed + c_sec + ) + set_target_properties(opsproto PROPERTIES OUTPUT_NAME + cust_opsproto_rt2.0 + ) + install(TARGETS opsproto + LIBRARY DESTINATION packages/vendors/${VENDOR_NAME}_transformer/op_proto/lib/linux/${CMAKE_SYSTEM_PROCESSOR} + ) + + add_ops_tiling_keys( + OP_NAME "ALL" + TILING_KEYS ${TILING_KEY} + ) + + add_opc_config( + OP_NAME "ALL" + CONFIG ${OP_DEBUG_CONFIG} + ) + + if(ADD_OPS_COMPILE_OPTION_V2) + add_ops_compile_options( + OP_NAME "ALL" + OPTIONS ${OPS_COMPILE_OPTIONS} + ) + endif() +endif () + +add_subdirectory(common) +if (NOT BUILD_OPS_RTY_KERNEL) + add_subdirectory(mc2) +endif() + +if(NOT DEFINED COMPILED_OPS) + set(COMPILED_OPS ${COMPILED_OPS} CACHE STRING "Comp") + set(COMPILED_OPS CACHE STRING "Compiled Ops" FORCE) +set(COMPILED_OP_DIRS CACHE STRING "Compiled Ops Dirs" FORCE) +endif() + +if(NOT DEFINED COMPILED_OP_DIRS) + set(COMPILED_OP_DIRS CACHE STRING "Compiled Ops Dirs" FORCE) +endif() + +set(OP_LIST) +set(OP_DIR_LIST) +op_add_subdirectory(OP_LIST OP_DIR_LIST) + +foreach (OP_DIR ${OP_DIR_LIST}) + if (EXISTS "${OP_DIR}/op_host") + add_subdirectory(${OP_DIR}/op_host) + else() + add_subdirectory(${OP_DIR}) + endif() +endforeach () + +add_subdirectory(moe) +list(APPEND OP_LIST "ffn") +add_subdirectory(attention) +list(APPEND OP_LIST ${COMPILED_OPS}) +list(REMOVE_DUPLICATES OP_LIST) +list(APPEND OP_DIR_LIST ${COMPILED_OP_DIRS}) +list(REMOVE_DUPLICATES OP_DIR_LIST) +add_subdirectory(gmm) +list(REMOVE_DUPLICATES OP_LIST) +list(APPEND OP_DIR_LIST ${COMPILED_OP_DIRS}) +list(REMOVE_DUPLICATES OP_DIR_LIST) +add_subdirectory(mc2) +list(REMOVE_DUPLICATES OP_LIST) +list(APPEND OP_DIR_LIST ${COMPILED_OP_DIRS}) +list(REMOVE_DUPLICATES OP_DIR_LIST) +list(APPEND OP_LIST "fused_gdn_gating") +list(APPEND OP_DIR_LIST ${CMAKE_CURRENT_SOURCE_DIR}/attention/fused_gdn_gating) + +set(OP_DEPEND_DIR_LIST) +op_add_depend_directory( + OP_LIST ${OP_LIST} + OP_DIR_LIST OP_DEPEND_DIR_LIST +) + +foreach (OP_DEPEND_DIR ${OP_DEPEND_DIR_LIST}) + if (EXISTS "${OP_DEPEND_DIR}/op_host") + add_subdirectory(${OP_DEPEND_DIR}/op_host) + else() + add_subdirectory(${OP_DEPEND_DIR}) + endif() +endforeach () + +install(DIRECTORY ${OPS_ADV_ACT}/ + DESTINATION ${IMPL_INSTALL_DIR}/ascendc/common/act +) + +install(DIRECTORY ${OPS_GROUPEDMATMUL_ACT}/ + DESTINATION ${IMPL_INSTALL_DIR}/ascendc/common/groupedmatmul_act +) + + +if (BUILD_OPS_RTY_KERNEL) + get_target_property(base_aclnn_srcs op_host_aclnn SOURCES) + get_target_property(base_aclnn_inner_srcs op_host_aclnnInner SOURCES) + get_target_property(base_aclnn_exclude_srcs op_host_aclnnExc SOURCES) + set(base_aclnn_binary_dir ${ASCEND_AUTOGEN_DIR}) + + set(generate_aclnn_srcs) + set(generate_aclnn_inner_srcs) + set(generate_aclnn_headers) + set(generate_proto_dir ${base_aclnn_binary_dir}) + set(generate_exclude_proto_srcs) + set(generate_proto_srcs) + set(generate_proto_headers) + + if (base_aclnn_srcs) + foreach (_src ${base_aclnn_srcs}) + string(REGEX MATCH "^${CMAKE_CURRENT_SOURCE_DIR}" is_match "${_src}") + if (is_match) + get_filename_component(name_without_ext ${_src} NAME_WE) + + string(REGEX REPLACE "_def$" "" _op_name ${name_without_ext}) + list(APPEND generate_aclnn_srcs ${base_aclnn_binary_dir}/aclnn_${_op_name}.cpp) + list(APPEND generate_aclnn_headers ${base_aclnn_binary_dir}/aclnn_${_op_name}.h) + list(APPEND generate_proto_srcs ${generate_proto_dir}/${_op_name}_proto.cpp) + list(APPEND generate_proto_headers ${generate_proto_dir}/${_op_name}_proto.h) + endif () + endforeach () + else () + add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_stub.cpp + COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_stub.cpp + ) + + target_sources(op_host_aclnn PRIVATE + ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_stub.cpp + ) + endif () + + if (base_aclnn_inner_srcs) + foreach (_src ${base_aclnn_inner_srcs}) + string(REGEX MATCH "^${CMAKE_CURRENT_SOURCE_DIR}" is_match "${_src}") + if (is_match) + get_filename_component(name_without_ext ${_src} NAME_WE) + string(REGEX REPLACE "_def$" "" _op_name ${name_without_ext}) + list(APPEND generate_aclnn_inner_srcs ${base_aclnn_binary_dir}/inner/aclnnInner_${_op_name}.cpp) + list(APPEND generate_proto_srcs ${generate_proto_dir}/inner/${_op_name}_proto.cpp) + list(APPEND generate_proto_headers ${generate_proto_dir}/inner/${_op_name}_proto.h) + endif () + endforeach () + else () + add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_inner_stub.cpp + COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_inner_stub.cpp + ) + + target_sources(op_host_aclnnInner PRIVATE + ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_inner_stub.cpp + ) + endif () + + if (base_aclnn_exclude_srcs) + foreach (_src ${base_aclnn_exclude_srcs}) + string(REGEX MATCH "^${CMAKE_CURRENT_SOURCE_DIR}" is_match "${_src}") + if (is_match) + get_filename_component(name_without_ext ${_src} NAME_WE) + string(REGEX REPLACE "_def$" "" _op_name ${name_without_ext}) + list(APPEND generate_exclude_proto_srcs ${generate_proto_dir}/exc/${_op_name}_proto.cpp) + list(APPEND generate_proto_srcs ${generate_proto_dir}/exc/${_op_name}_proto.cpp) + list(APPEND generate_proto_headers ${generate_proto_dir}/exc/${_op_name}_proto.h) + endif () + endforeach () + else() + add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_exc_stub.cpp + COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_exc_stub.cpp + ) + + target_sources(op_host_aclnnExc PRIVATE + ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_exc_stub.cpp + ) + endif () + + + if (generate_aclnn_srcs OR generate_aclnn_inner_srcs) + set(ops_aclnn_src ${generate_aclnn_srcs} ${generate_aclnn_inner_srcs}) + else () + set(ops_aclnn_src ${CMAKE_CURRENT_BINARY_DIR}/ops_aclnn_src_stub.cpp) + + add_custom_command(OUTPUT ${ops_aclnn_src} + COMMAND touch ${ops_aclnn_src} + ) + endif () + + set_source_files_properties(${ops_aclnn_src} + PROPERTIES GENERATED TRUE + ) + add_library(ops_aclnn STATIC + ${ops_aclnn_src} + ) + target_compile_options(ops_aclnn PRIVATE + $<$:-std=gnu++1z> + ) + target_link_libraries(ops_aclnn PRIVATE + $ + ) + add_dependencies(ops_aclnn opbuild_gen_default opbuild_gen_inner) + + set_source_files_properties(${generate_proto_srcs} + PROPERTIES GENERATED TRUE + ) + target_sources(opsproto PRIVATE + ${generate_proto_srcs} + ) + add_dependencies(opsproto ops_transformer_proto_headers) + + install(FILES ${generate_proto_headers} + DESTINATION packages/vendors/${VENDOR_NAME}_transformer/op_proto/inc OPTIONAL + ) + + add_library(ops_transformer_proto_headers INTERFACE) + + target_include_directories(ops_transformer_proto_headers INTERFACE + $ + $ + $ + $ + ) + + add_dependencies(ops_transformer_proto_headers opbuild_gen_default opbuild_gen_inner opbuild_gen_exc) + + if (NOT BUILD_OPEN_PROJECT) + if (generate_proto_srcs) + install_package( + PACKAGE ops_adv + TARGETS ops_proto_headers + FILES ${generate_proto_headers} + DESTINATION include/ops_adv/proto + ) + endif () + endif () + + if (generate_aclnn_srcs) + add_custom_command(OUTPUT ${generate_aclnn_srcs} ${generate_aclnn_headers} + COMMAND mkdir -p ${base_aclnn_binary_dir} + COMMAND OPS_PROTO_SEPARATE=1 + OPS_ACLNN_GEN=1 + OPS_PROJECT_NAME=aclnn + ${OP_BUILD_TOOL} + $ + ${base_aclnn_binary_dir} + ) + endif () + + add_custom_target(opbuild_gen_default + DEPENDS ${generate_aclnn_srcs} ${generate_aclnn_headers} op_host_aclnn + ) + + if (generate_aclnn_inner_srcs) + add_custom_command(OUTPUT ${generate_aclnn_inner_srcs} + COMMAND mkdir -p ${base_aclnn_binary_dir}/inner + COMMAND OPS_PROTO_SEPARATE=1 + OPS_ACLNN_GEN=1 + OPS_PROJECT_NAME=aclnnInner + ${OP_BUILD_TOOL} + $ + ${base_aclnn_binary_dir}/inner + ) + endif () + + add_custom_target(opbuild_gen_inner + DEPENDS ${generate_aclnn_inner_srcs} op_host_aclnnInner + ) + + if (generate_exclude_proto_srcs) + add_custom_command(OUTPUT ${generate_exclude_proto_srcs} + COMMAND mkdir -p ${base_aclnn_binary_dir}/exc + COMMAND OPS_PROTO_SEPARATE=1 + OPS_ACLNN_GEN=0 + OPS_PROJECT_NAME=aclnnExc + ${OP_BUILD_TOOL} + $ + ${base_aclnn_binary_dir}/exc + ) + endif () + + add_custom_target(opbuild_gen_exc + DEPENDS ${generate_exclude_proto_srcs} op_host_aclnnExc + ) + + add_custom_target(generate_transformer_adapt_py + COMMAND ${HI_PYTHON} ${CMAKE_CURRENT_SOURCE_DIR}/cmake/scripts/util/ascendc_impl_build.py + \"\" + \"\" + \"\" + \"\" + ${ASCEND_IMPL_OUT_DIR} + ${ASCEND_AUTOGEN_DIR} + --opsinfo-dir ${base_aclnn_binary_dir} ${base_aclnn_binary_dir}/inner ${base_aclnn_binary_dir}/exc + ) + + add_dependencies(generate_transformer_adapt_py opbuild_gen_default opbuild_gen_inner opbuild_gen_exc) + + foreach (_op_name ${OP_LIST}) + install(FILES ${ASCEND_IMPL_OUT_DIR}/dynamic/${_op_name}.py + DESTINATION ${IMPL_DYNAMIC_INSTALL_DIR} + OPTIONAL + ) + endforeach () + + install(DIRECTORY ${OPS_ADV_UTILS_KERNEL_INC}/ + DESTINATION ${IMPL_INSTALL_DIR}/ascendc/common + ) + + foreach (op_dir ${OP_DIR_LIST}) + get_filename_component(_op_name "${op_dir}" NAME) + + if (EXISTS "${op_dir}/op_kernel") + file(GLOB KERNEL_FILES + ${op_dir}/op_kernel/*.cpp + ${op_dir}/op_kernel/*.h + ) + else() + file(GLOB KERNEL_FILES + ${op_dir}/*.cpp + ${op_dir}/*.h + ) + endif() + + install(FILES ${KERNEL_FILES} + DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_name} + OPTIONAL + ) + + install(DIRECTORY ${op_dir}/arch32 + DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_name} + OPTIONAL + ) + + install(DIRECTORY ${op_dir}/arch35 + DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_name} + OPTIONAL + ) + + install(DIRECTORY ${op_dir}/arch38 + DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_name} + OPTIONAL + ) + + install(DIRECTORY ${op_dir}/regbase/opkernel + DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_name}/regbase + OPTIONAL + ) + endforeach () + + add_custom_target(prepare_build ALL) + add_custom_target(generate_compile_cmd ALL) + add_custom_target(generate_ops_info ALL) + add_dependencies(prepare_build generate_transformer_adapt_py generate_compile_cmd) + + foreach (compute_unit ${ASCEND_COMPUTE_UNIT}) + add_compile_cmd_target( + COMPUTE_UNIT ${compute_unit} + ) + + add_ops_info_target( + COMPUTE_UNIT ${compute_unit} + ) + endforeach () + + add_custom_target(ops_transformer_kernel ALL) + add_custom_target(ops_transformer_config ALL) + add_dependencies(ops_transformer_kernel ops_transformer_config) + + foreach (compute_unit ${ASCEND_COMPUTE_UNIT}) + add_bin_compile_target( + COMPUTE_UNIT + ${compute_unit} + OP_INFO + ${OP_DIR_LIST} + ) + endforeach () +endif () diff --git a/csrc/ascend/aclnn_torch_adapter/NPUBridge.cpp b/csrc/ascend/aclnn_torch_adapter/NPUBridge.cpp new file mode 100644 index 000000000..71574c580 --- /dev/null +++ b/csrc/ascend/aclnn_torch_adapter/NPUBridge.cpp @@ -0,0 +1,30 @@ +// Copyright (c) 2020, Huawei Technologies Co., Ltd +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include "NPUBridge.h" + +namespace vllm_fl +{ + NPUStorageImpl *NPUBridge::GetNpuStorageImpl(c10::StorageImpl *storageImpl) + { + return static_cast(storageImpl); + } + + NPUStorageImpl *NPUBridge::GetNpuStorageImpl(c10::Storage &&storage) + { + return static_cast(storage.unsafeGetStorageImpl()); + } + + NPUStorageImpl *NPUBridge::GetNpuStorageImpl(const at::Tensor &tensor) + { + return static_cast(tensor.storage().unsafeGetStorageImpl()); + } + + NPUStorageDesc &NPUBridge::GetNpuStorageImplDesc(const at::Tensor &tensor) + { + return static_cast(tensor.storage().unsafeGetStorageImpl())->npu_desc_; + } +} diff --git a/csrc/ascend/aclnn_torch_adapter/NPUBridge.h b/csrc/ascend/aclnn_torch_adapter/NPUBridge.h new file mode 100644 index 000000000..b9fdd844e --- /dev/null +++ b/csrc/ascend/aclnn_torch_adapter/NPUBridge.h @@ -0,0 +1,29 @@ +// Copyright (c) 2020, Huawei Technologies Co., Ltd +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once +#include +#include "NPUStorageImpl.h" + +namespace vllm_fl +{ + + class NPUBridge + { + public: + // at::tensor to NPUStorageImpl + static NPUStorageImpl *GetNpuStorageImpl(const at::Tensor &tensor); + + // c10::StorageImpl to NPUStorageImpl + static NPUStorageImpl *GetNpuStorageImpl(c10::StorageImpl *storageImpl); + + // c10::Storage to NPUStorageImpl + static NPUStorageImpl *GetNpuStorageImpl(c10::Storage &&storage); + + // tensor to NPUStorageDesc + static NPUStorageDesc &GetNpuStorageImplDesc(const at::Tensor &tensor); + }; +} diff --git a/csrc/ascend/aclnn_torch_adapter/NPUStorageImpl.cpp b/csrc/ascend/aclnn_torch_adapter/NPUStorageImpl.cpp new file mode 100644 index 000000000..80882b731 --- /dev/null +++ b/csrc/ascend/aclnn_torch_adapter/NPUStorageImpl.cpp @@ -0,0 +1,52 @@ +// Copyright (c) 2020, Huawei Technologies Co., Ltd +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include "NPUStorageImpl.h" + +namespace vllm_fl +{ + + NPUStorageImpl::NPUStorageImpl( + use_byte_size_t use_byte_size, + size_t size_bytes, + at::DataPtr data_ptr, + at::Allocator *allocator, + bool resizable) : c10::StorageImpl(use_byte_size, + size_bytes, + at::DataPtr(std::move(data_ptr)), + allocator, + resizable) + { + } + + void NPUStorageImpl::release_resources() + { + StorageImpl::release_resources(); + } + + c10::intrusive_ptr make_npu_storage_impl( + c10::StorageImpl::use_byte_size_t, + c10::SymInt size_bytes, + c10::DataPtr data_ptr, + c10::Allocator *allocator, + bool resizable) + { + if (data_ptr == nullptr) + { + data_ptr = allocator->allocate(size_bytes.as_int_unchecked()); + } + // Correctly create NPUStorageImpl object. + c10::intrusive_ptr npu_storage_impl = c10::make_intrusive( + c10::StorageImpl::use_byte_size_t(), + size_bytes.as_int_unchecked(), + std::move(data_ptr), + allocator, + resizable); + // There is no need to consider the NPUStorageDesc information, it will be carried out in the subsequent processing. + return npu_storage_impl; + } + +} diff --git a/csrc/ascend/aclnn_torch_adapter/NPUStorageImpl.h b/csrc/ascend/aclnn_torch_adapter/NPUStorageImpl.h new file mode 100644 index 000000000..858b96d63 --- /dev/null +++ b/csrc/ascend/aclnn_torch_adapter/NPUStorageImpl.h @@ -0,0 +1,67 @@ +// Copyright (c) 2020, Huawei Technologies Co., Ltd +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "acl/acl_rt.h" +#include "acl/acl_base.h" + +namespace vllm_fl +{ + + struct NPUStorageDesc + { + public: + struct use_byte_size_t + { + }; + + c10::SmallVector base_sizes_; + c10::SmallVector base_strides_; + c10::SmallVector storage_sizes_; + int64_t base_offset_ = 0; + use_byte_size_t base_dtype_ = {}; + aclFormat origin_format_ = ACL_FORMAT_UNDEFINED; + aclFormat npu_format_ = ACL_FORMAT_ND; + // used to make CANN GE tensor from storagImpl + caffe2::TypeMeta data_type_ = caffe2::TypeMeta::Make(); + }; + + struct NPUStorageImpl : public c10::StorageImpl + { + explicit NPUStorageImpl( + use_byte_size_t use_byte_size, + size_t size_bytes, + at::DataPtr data_ptr, + at::Allocator *allocator, + bool resizable); + ~NPUStorageImpl() override = default; + + void release_resources() override; + + NPUStorageDesc npu_desc_; + + NPUStorageDesc get_npu_desc() const + { + return npu_desc_; + } + }; + + c10::intrusive_ptr make_npu_storage_impl( + c10::StorageImpl::use_byte_size_t, + c10::SymInt size_bytes, + c10::DataPtr data_ptr, + c10::Allocator *allocator, + bool resizable); + +} diff --git a/csrc/ascend/aclnn_torch_adapter/op_api_common.h b/csrc/ascend/aclnn_torch_adapter/op_api_common.h new file mode 100644 index 000000000..b9dba4ee2 --- /dev/null +++ b/csrc/ascend/aclnn_torch_adapter/op_api_common.h @@ -0,0 +1,591 @@ +// Copyright (c) 2023 Huawei Technologies Co., Ltd +// All rights reserved. +// +// Licensed under the BSD 3-Clause License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://opensource.org/licenses/BSD-3-Clause +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef OP_API_COMMON_ADAPTER +#define OP_API_COMMON_ADAPTER + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include "torch_npu/csrc/aten/NPUNativeFunctions.h" +#include "torch_npu/csrc/core/npu/NPUStream.h" +#include "torch_npu/csrc/framework/OpCommand.h" +#include "torch_npu/csrc/framework/interface/EnvVariables.h" +#include "torch_npu/csrc/framework/utils/CalcuOpUtil.h" +#include "torch_npu/csrc/framework/utils/OpPreparation.h" +#include "NPUBridge.h" +#include "NPUStorageImpl.h" + +#define NPU_NAME_SPACE at_npu::native +using namespace at; + +typedef struct aclOpExecutor aclOpExecutor; +typedef struct aclTensor aclTensor; +typedef struct aclScalar aclScalar; +typedef struct aclIntArray aclIntArray; +typedef struct aclFloatArray aclFloatArray; +typedef struct aclBoolArray aclBoolArray; +typedef struct aclTensorList aclTensorList; + +typedef aclTensor *(*_aclCreateTensor)( + const int64_t *view_dims, uint64_t view_dims_num, aclDataType data_type, + const int64_t *stride, int64_t offset, aclFormat format, + const int64_t *storage_dims, uint64_t storage_dims_num, void *tensor_data); +typedef aclScalar *(*_aclCreateScalar)(void *value, aclDataType data_type); +typedef aclIntArray *(*_aclCreateIntArray)(const int64_t *value, uint64_t size); +typedef aclFloatArray *(*_aclCreateFloatArray)(const float *value, + uint64_t size); +typedef aclBoolArray *(*_aclCreateBoolArray)(const bool *value, uint64_t size); +typedef aclTensorList *(*_aclCreateTensorList)(const aclTensor *const *value, + uint64_t size); + +typedef int (*_aclDestroyTensor)(const aclTensor *tensor); +typedef int (*_aclDestroyScalar)(const aclScalar *scalar); +typedef int (*_aclDestroyIntArray)(const aclIntArray *array); +typedef int (*_aclDestroyFloatArray)(const aclFloatArray *array); +typedef int (*_aclDestroyBoolArray)(const aclBoolArray *array); +typedef int (*_aclDestroyTensorList)(const aclTensorList *array); + +constexpr int kHashBufSize = 8192; +constexpr int kHashBufMaxSize = kHashBufSize + 1024; +extern thread_local char g_hashBuf[kHashBufSize]; +extern thread_local int g_hashOffset; + +#ifdef MMCV_WITH_XLA +#define DEVICE_TYPE at_npu::key::NativeDeviceType +#else +#define DEVICE_TYPE c10::DeviceType::PrivateUse1 +#endif + +#define AT_ALL_SCALAR_TYPE_AND_ACL_DATATYPE_PAIR(_) \ + _(at::ScalarType::Byte, ACL_UINT8) \ + _(at::ScalarType::Char, ACL_INT8) \ + _(at::ScalarType::Short, ACL_INT16) \ + _(at::ScalarType::Int, ACL_INT32) \ + _(at::ScalarType::Long, ACL_INT64) \ + _(at::ScalarType::Half, ACL_FLOAT16) \ + _(at::ScalarType::Float, ACL_FLOAT) \ + _(at::ScalarType::Double, ACL_DOUBLE) \ + _(at::ScalarType::ComplexHalf, ACL_DT_UNDEFINED) \ + _(at::ScalarType::ComplexFloat, ACL_COMPLEX64) \ + _(at::ScalarType::ComplexDouble, ACL_COMPLEX128) \ + _(at::ScalarType::Bool, ACL_BOOL) \ + _(at::ScalarType::QInt8, ACL_DT_UNDEFINED) \ + _(at::ScalarType::QUInt8, ACL_DT_UNDEFINED) \ + _(at::ScalarType::QInt32, ACL_DT_UNDEFINED) \ + _(at::ScalarType::BFloat16, ACL_BF16) \ + _(at::ScalarType::QUInt4x2, ACL_DT_UNDEFINED) \ + _(at::ScalarType::QUInt2x4, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Undefined, ACL_DT_UNDEFINED) \ + _(at::ScalarType::NumOptions, ACL_DT_UNDEFINED) + +constexpr aclDataType kATenScalarTypeToAclDataTypeTable + [static_cast(at::ScalarType::NumOptions) + 1] = { +#define DEFINE_ENUM(_1, n) n, + AT_ALL_SCALAR_TYPE_AND_ACL_DATATYPE_PAIR(DEFINE_ENUM) +#undef DEFINE_ENUM +}; + +#define GET_OP_API_FUNC(apiName) \ + reinterpret_cast<_##apiName>(GetOpApiFuncAddr(#apiName)) + +#define MEMCPY_TO_BUF(data_expression, size_expression) \ + if (g_hashOffset + (size_expression) > kHashBufSize) { \ + g_hashOffset = kHashBufMaxSize; \ + return; \ + } \ + memcpy(g_hashBuf + g_hashOffset, data_expression, size_expression); \ + g_hashOffset += size_expression; + +bool IsOpInputBaseFormat(const at::Tensor &tensor) +{ + if (!tensor.is_privateuseone()) { + return true; + } + const auto format = vllm_fl::NPUBridge::GetNpuStorageImplDesc(tensor).npu_format_; + return (format == ACL_FORMAT_ND) || (format == ACL_FORMAT_NCHW) || (format == ACL_FORMAT_NHWC) || + (format == ACL_FORMAT_NCDHW); +} + +inline const char *GetOpApiLibName(void) { return "libopapi.so"; } + +inline const char *GetCustOpApiLibName(void) { return "libcust_opapi.so"; } + +inline void *GetOpApiFuncAddrInLib(void *handler, const char *libName, + const char *apiName) { + auto funcAddr = dlsym(handler, apiName); + return funcAddr; +} + +inline void *GetOpApiLibHandler(const char *libName) { + auto handler = dlopen(libName, RTLD_LAZY); + return handler; +} + +inline void *GetOpApiFuncAddr(const char *apiName) { + static auto custOpApiHandler = GetOpApiLibHandler(GetCustOpApiLibName()); + if (custOpApiHandler != nullptr) { + auto funcAddr = + GetOpApiFuncAddrInLib(custOpApiHandler, GetCustOpApiLibName(), apiName); + if (funcAddr != nullptr) { + return funcAddr; + } + } + + static auto opApiHandler = GetOpApiLibHandler(GetOpApiLibName()); + if (opApiHandler == nullptr) { + return nullptr; + } + return GetOpApiFuncAddrInLib(opApiHandler, GetOpApiLibName(), apiName); +} + +inline c10::Scalar ConvertTensorToScalar(const at::Tensor &tensor) { + c10::Scalar expScalar; + const at::Tensor *aclInput = &tensor; + if (aclInput->scalar_type() == at::ScalarType::Double) { + double value = *(double *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Long) { + int64_t value = *(int64_t *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Float) { + float value = *(float *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Int) { + int value = *(int *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Half) { + c10::Half value = *(c10::Half *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Bool) { + int8_t value = *(int8_t *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::ComplexDouble) { + c10::complex value = *(c10::complex *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::ComplexFloat) { + c10::complex value = *(c10::complex *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::BFloat16) { + c10::BFloat16 value = *(c10::BFloat16 *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } + return expScalar; +} + +inline at::Tensor CopyTensorHostToDevice(const at::Tensor &cpu_tensor) { + at::Tensor cpuPinMemTensor = cpu_tensor.pin_memory(); + int deviceIndex = 0; + return cpuPinMemTensor.to(c10::Device(DEVICE_TYPE, deviceIndex), + cpuPinMemTensor.scalar_type(), true, true); +} + +inline at::Tensor CopyScalarToDevice(const c10::Scalar &cpu_scalar, + at::ScalarType scalar_data_type) { + return CopyTensorHostToDevice( + scalar_to_tensor(cpu_scalar).to(scalar_data_type)); +} + +inline aclTensor *ConvertType(const at::Tensor &at_tensor) { + static const auto aclCreateTensor = GET_OP_API_FUNC(aclCreateTensor); + if (aclCreateTensor == nullptr) { + return nullptr; + } + + if (!at_tensor.defined()) { + return nullptr; + } + at::ScalarType scalar_data_type = at_tensor.scalar_type(); + aclDataType acl_data_type = + kATenScalarTypeToAclDataTypeTable[static_cast(scalar_data_type)]; + TORCH_CHECK( + acl_data_type != ACL_DT_UNDEFINED, + std::string(c10::toString(scalar_data_type)) + " has not been supported") + c10::SmallVector storageDims; + // if acl_data_type is ACL_STRING, storageDims is empty. + auto itemsize = at_tensor.itemsize(); + TORCH_CHECK(itemsize != 0, "When ConvertType, tensor item size cannot be zero."); + + const auto dimNum = at_tensor.sizes().size(); + aclFormat format = ACL_FORMAT_ND; + if (!IsOpInputBaseFormat(at_tensor)) { + format = vllm_fl::NPUBridge::GetNpuStorageImpl(at_tensor)->npu_desc_.npu_format_; + if (acl_data_type != ACL_STRING) { + storageDims = vllm_fl::NPUBridge::GetNpuStorageImpl(at_tensor)->npu_desc_.storage_sizes_; + } + } else { + switch (dimNum) { + case 3: + format = ACL_FORMAT_NCL; + break; + case 4: + format = ACL_FORMAT_NCHW; + break; + case 5: + format = ACL_FORMAT_NCDHW; + break; + default: + format = ACL_FORMAT_ND; + } + if (acl_data_type != ACL_STRING) { + storageDims.push_back(at_tensor.storage().nbytes() / itemsize); + } + } + + if (at_tensor.unsafeGetTensorImpl()->is_wrapped_number()) { + c10::Scalar expScalar = ConvertTensorToScalar(at_tensor); + at::Tensor aclInput = CopyScalarToDevice(expScalar, scalar_data_type); + return aclCreateTensor(aclInput.sizes().data(), aclInput.sizes().size(), + acl_data_type, aclInput.strides().data(), + aclInput.storage_offset(), format, + storageDims.data(), storageDims.size(), + const_cast(aclInput.storage().data())); + } + + auto acl_tensor = aclCreateTensor( + at_tensor.sizes().data(), at_tensor.sizes().size(), acl_data_type, + at_tensor.strides().data(), at_tensor.storage_offset(), format, + storageDims.data(), storageDims.size(), + const_cast(at_tensor.storage().data())); + return acl_tensor; +} + +inline aclScalar *ConvertType(const at::Scalar &at_scalar) { + static const auto aclCreateScalar = GET_OP_API_FUNC(aclCreateScalar); + if (aclCreateScalar == nullptr) { + return nullptr; + } + + at::ScalarType scalar_data_type = at_scalar.type(); + aclDataType acl_data_type = + kATenScalarTypeToAclDataTypeTable[static_cast(scalar_data_type)]; + TORCH_CHECK( + acl_data_type != ACL_DT_UNDEFINED, + std::string(c10::toString(scalar_data_type)) + " has not been supported") + aclScalar *acl_scalar = nullptr; + switch (scalar_data_type) { + case at::ScalarType::Double: { + double value = at_scalar.toDouble(); + acl_scalar = aclCreateScalar(&value, acl_data_type); + break; + } + case at::ScalarType::Long: { + int64_t value = at_scalar.toLong(); + acl_scalar = aclCreateScalar(&value, acl_data_type); + break; + } + case at::ScalarType::Bool: { + bool value = at_scalar.toBool(); + acl_scalar = aclCreateScalar(&value, acl_data_type); + break; + } + case at::ScalarType::ComplexDouble: { + auto value = at_scalar.toComplexDouble(); + acl_scalar = aclCreateScalar(&value, acl_data_type); + break; + } + default: + acl_scalar = nullptr; + break; + } + return acl_scalar; +} + +inline aclIntArray *ConvertType(const at::IntArrayRef &at_array) { + static const auto aclCreateIntArray = GET_OP_API_FUNC(aclCreateIntArray); + if (aclCreateIntArray == nullptr) { + return nullptr; + } + auto array = aclCreateIntArray(at_array.data(), at_array.size()); + return array; +} + +template +inline aclBoolArray *ConvertType(const std::array &value) { + static const auto aclCreateBoolArray = GET_OP_API_FUNC(aclCreateBoolArray); + if (aclCreateBoolArray == nullptr) { + return nullptr; + } + + auto array = aclCreateBoolArray(value.data(), value.size()); + return array; +} + +inline aclBoolArray *ConvertType(const at::ArrayRef &value) { + static const auto aclCreateBoolArray = GET_OP_API_FUNC(aclCreateBoolArray); + if (aclCreateBoolArray == nullptr) { + return nullptr; + } + + auto array = aclCreateBoolArray(value.data(), value.size()); + return array; +} + +inline aclTensorList *ConvertType(const at::TensorList &at_tensor_list) { + static const auto aclCreateTensorList = GET_OP_API_FUNC(aclCreateTensorList); + if (aclCreateTensorList == nullptr) { + return nullptr; + } + + std::vector tensor_list(at_tensor_list.size()); + for (size_t i = 0; i < at_tensor_list.size(); i++) { + tensor_list[i] = ConvertType(at_tensor_list[i]); + } + auto acl_tensor_list = + aclCreateTensorList(tensor_list.data(), tensor_list.size()); + return acl_tensor_list; +} + +inline aclTensor *ConvertType(const c10::optional &opt_tensor) { + if (opt_tensor.has_value() && opt_tensor.value().defined()) { + return ConvertType(opt_tensor.value()); + } + return nullptr; +} + +inline aclIntArray *ConvertType( + const c10::optional &opt_array) { + if (opt_array.has_value()) { + return ConvertType(opt_array.value()); + } + return nullptr; +} + +inline aclScalar *ConvertType(const c10::optional &opt_scalar) { + if (opt_scalar.has_value()) { + return ConvertType(opt_scalar.value()); + } + return nullptr; +} + +inline aclDataType ConvertType(const at::ScalarType scalarType) { + return kATenScalarTypeToAclDataTypeTable[static_cast(scalarType)]; +} + +template +T ConvertType(T value) { + return value; +} + +template +auto ConvertToOpApiFunc(const Tuple ¶ms, void *opApiAddr, + std::index_sequence) { + typedef int (*OpApiFunc)( + typename std::decay(params))>::type...); + auto func = reinterpret_cast(opApiAddr); + return func; +} + +template +auto ConvertToOpApiFunc(const Tuple ¶ms, void *opApiAddr) { + static constexpr auto size = std::tuple_size::value; + return ConvertToOpApiFunc(params, opApiAddr, + std::make_index_sequence{}); +} + +inline void Release(aclTensor *p) { + static const auto aclDestroyTensor = GET_OP_API_FUNC(aclDestroyTensor); + if (aclDestroyTensor == nullptr) { + return; + } + aclDestroyTensor(p); +} + +inline void Release(aclScalar *p) { + static const auto aclDestroyScalar = GET_OP_API_FUNC(aclDestroyScalar); + if (aclDestroyScalar == nullptr) { + return; + } + aclDestroyScalar(p); +} + +inline void Release(aclIntArray *p) { + static const auto aclDestroyIntArray = GET_OP_API_FUNC(aclDestroyIntArray); + if (aclDestroyIntArray == nullptr) { + return; + } + + aclDestroyIntArray(p); +} + +inline void Release(aclBoolArray *p) { + static const auto aclDestroyBoolArray = GET_OP_API_FUNC(aclDestroyBoolArray); + if (aclDestroyBoolArray == nullptr) { + return; + } + + aclDestroyBoolArray(p); +} + +inline void Release(aclTensorList *p) { + static const auto aclDestroyTensorList = + GET_OP_API_FUNC(aclDestroyTensorList); + if (aclDestroyTensorList == nullptr) { + return; + } + + aclDestroyTensorList(p); +} + +template +void Release(T value) { + (void)value; +} + +template +void CallRelease(Tuple t, std::index_sequence) { + (void)std::initializer_list{(Release(std::get(t)), 0)...}; +} + +template +void ReleaseConvertTypes(Tuple &t) { + static constexpr auto size = std::tuple_size::value; + CallRelease(t, std::make_index_sequence{}); +} + +template +constexpr auto ConvertTypes(Ts &... args) { + return std::make_tuple(ConvertType(args)...); +} + +template +auto call(Function f, Tuple t, std::index_sequence) { + return f(std::get(t)...); +} + +template +auto call(Function f, Tuple t) { + static constexpr auto size = std::tuple_size::value; + return call(f, t, std::make_index_sequence{}); +} + +template +void AddParamToBuf(const std::array &value) { + MEMCPY_TO_BUF(value.data(), value.size() * sizeof(bool)); +} + +template +void AddParamToBuf(const T &value) { + MEMCPY_TO_BUF(&value, sizeof(T)); +} + +void AddParamToBuf(const at::Tensor &); +void AddParamToBuf(const at::Scalar &); +void AddParamToBuf(const at::IntArrayRef &); +void AddParamToBuf(const at::ArrayRef &); +void AddParamToBuf(const at::TensorList &); +void AddParamToBuf(const c10::optional &); +void AddParamToBuf(const c10::optional &); +void AddParamToBuf(const c10::optional &); +void AddParamToBuf(const at::ScalarType); +void AddParamToBuf(const string &); +void AddParamToBuf(); + +template +void AddParamToBuf(const T &arg, Args &... args) { + AddParamToBuf(arg); + AddParamToBuf(args...); +} + +uint64_t CalcHashId(); +typedef int (*InitHugeMemThreadLocal)(void *, bool); +typedef void (*UnInitHugeMemThreadLocal)(void *, bool); +typedef void (*ReleaseHugeMem)(void *, bool); + +#define EXEC_NPU_CMD(aclnn_api, ...) \ + do { \ + static const auto getWorkspaceSizeFuncAddr = \ + GetOpApiFuncAddr(#aclnn_api "GetWorkspaceSize"); \ + static const auto opApiFuncAddr = GetOpApiFuncAddr(#aclnn_api); \ + static const auto initMemAddr = \ + GetOpApiFuncAddr("InitHugeMemThreadLocal"); \ + static const auto unInitMemAddr = \ + GetOpApiFuncAddr("UnInitHugeMemThreadLocal"); \ + static const auto releaseMemAddr = GetOpApiFuncAddr("ReleaseHugeMem"); \ + TORCH_CHECK( \ + getWorkspaceSizeFuncAddr != nullptr && opApiFuncAddr != nullptr, \ + #aclnn_api, " or ", #aclnn_api "GetWorkspaceSize", " not in ", \ + GetOpApiLibName(), ", or ", GetOpApiLibName(), "not found."); \ + auto acl_stream = c10_npu::getCurrentNPUStream().stream(false); \ + uint64_t workspace_size = 0; \ + uint64_t *workspace_size_addr = &workspace_size; \ + aclOpExecutor *executor = nullptr; \ + aclOpExecutor **executor_addr = &executor; \ + InitHugeMemThreadLocal initMemFunc = \ + reinterpret_cast(initMemAddr); \ + UnInitHugeMemThreadLocal unInitMemFunc = \ + reinterpret_cast(unInitMemAddr); \ + if (initMemFunc) { \ + initMemFunc(nullptr, false); \ + } \ + auto converted_params = \ + ConvertTypes(__VA_ARGS__, workspace_size_addr, executor_addr); \ + static auto getWorkspaceSizeFunc = \ + ConvertToOpApiFunc(converted_params, getWorkspaceSizeFuncAddr); \ + auto workspace_status = call(getWorkspaceSizeFunc, converted_params); \ + TORCH_CHECK(workspace_status == 0, \ + "call " #aclnn_api " failed, detail:", aclGetRecentErrMsg()); \ + void *workspace_addr = nullptr; \ + if (workspace_size != 0) { \ + at::TensorOptions options = \ + at::TensorOptions(torch_npu::utils::get_npu_device_type()); \ + auto workspace_tensor = \ + at::empty({workspace_size}, options.dtype(kByte)); \ + workspace_addr = const_cast(workspace_tensor.storage().data()); \ + } \ + auto acl_call = [converted_params, workspace_addr, workspace_size, \ + acl_stream, executor]() -> int { \ + typedef int (*OpApiFunc)(void *, uint64_t, aclOpExecutor *, \ + const aclrtStream); \ + OpApiFunc opApiFunc = reinterpret_cast(opApiFuncAddr); \ + auto api_ret = \ + opApiFunc(workspace_addr, workspace_size, executor, acl_stream); \ + TORCH_CHECK(api_ret == 0, "call " #aclnn_api " failed, detail:", \ + aclGetRecentErrMsg()); \ + ReleaseConvertTypes(converted_params); \ + ReleaseHugeMem releaseMemFunc = \ + reinterpret_cast(releaseMemAddr); \ + if (releaseMemFunc) { \ + releaseMemFunc(nullptr, false); \ + } \ + return api_ret; \ + }; \ + at_npu::native::OpCommand cmd; \ + cmd.Name(#aclnn_api); \ + cmd.SetCustomHandler(acl_call); \ + cmd.Run(); \ + if (unInitMemFunc) { \ + unInitMemFunc(nullptr, false); \ + } \ + } while (false) + +#endif diff --git a/csrc/ascend/attention/CMakeLists.txt b/csrc/ascend/attention/CMakeLists.txt new file mode 100644 index 000000000..37cea3a2c --- /dev/null +++ b/csrc/ascend/attention/CMakeLists.txt @@ -0,0 +1,35 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +set(OPTEST_NAME optest_${PKG_NAME}) + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +foreach(SUB_DIR ${CURRENT_DIRS}) + if (DEFINED ASCEND950_SKIP_OPS AND ${SUB_DIR} IN_LIST ASCEND950_SKIP_OPS) + message(STATUS "Skipping ${SUB_DIR} because ASCEND950 is not available.") + continue() + endif() + + if (DEFINED ASCEND_OP_NAME AND NOT "${ASCEND_OP_NAME}" STREQUAL "") + if (NOT "${ASCEND_OP_NAME}" STREQUAL "all" AND NOT "${ASCEND_OP_NAME}" STREQUAL "ALL") + if (NOT ${SUB_DIR} IN_LIST ASCEND_OP_NAME) + continue() + endif () + endif () + endif () + + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + else() + if (EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/op_host/CMakeLists.txt") + add_subdirectory(${SUB_DIR}/op_host) + endif() + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/attention/common/op_kernel/CopyInL1.h b/csrc/ascend/attention/common/op_kernel/CopyInL1.h new file mode 100644 index 000000000..e0aaf7826 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/CopyInL1.h @@ -0,0 +1,514 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file CopyInL1.h + * \brief + */ +#ifndef COPYINL1_H +#define COPYINL1_H + +enum class KVLAYOUT +{ + BNBD, // [blockNums, headNum, blockSize, headDim] + BBH, // [blockNums, blockSize, headNum * headDim] + NZ // [blockNums, headNum, d1, blockSize, d0], d1 = headDim / d0, d0 = 32 (block byte) / sizeof(KV_T) +}; + +struct CopyParam{ + uint32_t width; + uint32_t height; + uint32_t orgWidth; +}; + +struct PAShape{ + uint32_t blockNum; + uint32_t blockSize; + uint32_t headNum; // 一般为kv的head num + uint32_t headDim; // mla下rope为64， 非rope为512 + uint32_t maxblockNumPerBatch; // block table 每一行的最大个数 + uint32_t actHeadDim; // 实际拷贝col大小，考虑到N切块 s*d， 对应d + uint32_t copyRowNum; + uint32_t copyRowNumAlign; + uint32_t pageStride; +}; + +struct Position{ + uint32_t bIdx; + uint32_t n2Idx; + uint32_t s2Offset; + uint32_t dIdx; // N轴被切，对应D轴被切 +}; + +template +__aicore__ inline void GmCopyInToL1(LocalTensor& L1Tensor, GlobalTensor& GmTensor, const CopyParam& mmCopyParam) +{ + Nd2NzParams Gm2L1Nd2NzParams; + Gm2L1Nd2NzParams.ndNum = 1; // ND矩阵的个数 + Gm2L1Nd2NzParams.nValue = mmCopyParam.height; // 单个ND矩阵的实际行数，单位为元素个数 + Gm2L1Nd2NzParams.dValue = mmCopyParam.width; // 单个ND矩阵的实际列数(vD)，单位为元素个数 + Gm2L1Nd2NzParams.srcNdMatrixStride = 0; // 相邻ND矩阵起始地址之间的偏移， 单位为元素个数 + Gm2L1Nd2NzParams.srcDValue = mmCopyParam.orgWidth; // 同一个ND矩阵中相邻行起始地址之间的偏移， 单位为元素个数 + Gm2L1Nd2NzParams.dstNzC0Stride = (Gm2L1Nd2NzParams.nValue + 15) >> 4 << 4; // 转换为NZ矩阵后，相邻Block起始地址之间的偏移， 单位为Block个数 + Gm2L1Nd2NzParams.dstNzNStride = 1; // 转换为NZ矩阵后，ND之间相邻两行在NZ矩阵中起始地址之间的偏移， 单位为Block个数 + Gm2L1Nd2NzParams.dstNzMatrixStride = 0; // 两个NZ矩阵，起始地址之间的偏移， 单位为元素数量 + DataCopy(L1Tensor, GmTensor, Gm2L1Nd2NzParams); +} + +// 场景：key、value GM to L1 +// GM按ND格式存储 +// L1按NZ格式存储 +// GM的行、列、列的stride（D or ND）BNSD 和 BSH的区别 +template +__aicore__ inline void DataCopyGmNDToL1(LocalTensor& l1Tensor, GlobalTensor& gmTensor, + uint32_t rowAct, + uint32_t rowAlign, + uint32_t col, // D + uint32_t colStride) // D or N*D +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = rowAct; // 行数 + + nd2nzPara.dValue = col; + nd2nzPara.srcDValue = colStride; + nd2nzPara.dstNzC0Stride = rowAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(l1Tensor, gmTensor, nd2nzPara); +} + +template +__aicore__ inline void DataCopyGmScaleNDToL1(LocalTensor& l1Tensor, GlobalTensor& gmTensor, + uint32_t rowAct, + uint32_t rowAlign, + uint32_t col, // D + uint32_t colStride) // D or N*D +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = rowAct; + + nd2nzPara.dValue = col; + nd2nzPara.srcDValue = colStride; + nd2nzPara.dstNzC0Stride = rowAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = nd2nzPara.nValue; + + LocalTensor l1TensorCast = l1Tensor.template ReinterpretCast(); + GlobalTensor gmTensorCast; + gmTensorCast.SetGlobalBuffer(((__gm__ bfloat16_t*)(gmTensor.GetPhyAddr()))); + DataCopy(l1TensorCast, gmTensorCast, nd2nzPara); +} + +template +__aicore__ inline void DataCopyGmScaleDNToL1(LocalTensor& l1Tensor, GlobalTensor& gmTensor, + uint32_t rowAct, + uint32_t rowAlign, + uint32_t col, + uint32_t colStride) +{ + Dn2NzParams dn2nzPara; + dn2nzPara.dnNum = 1; + dn2nzPara.nValue = col / 2; + dn2nzPara.dValue = rowAct; + dn2nzPara.srcDValue = colStride / 2; + dn2nzPara.dstNzC0Stride = dn2nzPara.nValue; + dn2nzPara.dstNzNStride = 1; + dn2nzPara.srcDnMatrixStride = 0; + dn2nzPara.dstNzMatrixStride = dn2nzPara.nValue; + + LocalTensor l1TensorCast = l1Tensor.template ReinterpretCast(); + GlobalTensor gmTensorCast; + gmTensorCast.SetGlobalBuffer(((__gm__ bfloat16_t*)(gmTensor.GetPhyAddr()))); + DataCopy(l1TensorCast, gmTensorCast, dn2nzPara); +} + +template +__aicore__ inline void DataCopyGmNZToL1(LocalTensor& l1Tensor, GlobalTensor& gmTensor, + uint32_t rowAct, + uint32_t dstRowStride, + uint32_t srcRowStride, + uint32_t col) +{ + uint32_t blockElementCnt = 32U / sizeof(L1Type); + if constexpr (IsSameType::value) { + blockElementCnt = 64U; + } + DataCopyParams intriParams; + intriParams.blockCount = col / blockElementCnt; + intriParams.blockLen = rowAct; + intriParams.dstStride = dstRowStride; + intriParams.srcStride = srcRowStride; + DataCopy(l1Tensor, gmTensor, intriParams); +} + +template +__aicore__ inline void GmCopyInToL1HasRopePANoContinue(LocalTensor& nopeTensor, LocalTensor& ropeTensor, + GlobalTensor& nopeGmTensor, GlobalTensor& ropeGmTensor, + GlobalTensor& blockTableGm, KVLAYOUT kvLayout, + const PAShape &shape, + const PAShape &ropeShape, + const Position &startPos) +{ + uint32_t copyFinishRowCnt = 0; + uint64_t blockTableBaseOffset = startPos.bIdx * shape.maxblockNumPerBatch; // 块表的基偏移量 + uint32_t curS2Idx = startPos.s2Offset; + uint32_t blockElementCnt = 32U / sizeof(L1Type); // 每个块的元素数量 + // ropeshape的M方向与nopeshape保持一样， 此处只判断nopeshape的 + while(copyFinishRowCnt < shape.copyRowNum){ + uint64_t blockIdOffset = curS2Idx / shape.blockSize; // 获取block table上的索引 + uint64_t remainRowCnt = curS2Idx % shape.blockSize; // 获取在单个块上超出的行数 + uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset); // 从block table上获取的编号 + //计算可以拷贝行数 + uint32_t copyRowCnt = shape.blockSize - remainRowCnt; // 一次只能处理一个Block + if (copyFinishRowCnt + copyRowCnt > shape.copyRowNum){ + copyRowCnt = shape.copyRowNum - copyFinishRowCnt; // 一个block未拷满 + } + uint64_t offset = idInBlockTable * shape.blockSize * shape.headNum * shape.headDim; // PA的偏移 + if (shape.pageStride > 0) { + offset = idInBlockTable * shape.pageStride; + } + uint64_t keyRopeOffset = idInBlockTable * ropeShape.blockSize * ropeShape.headNum * ropeShape.headDim; + if (ropeShape.pageStride > 0) { + keyRopeOffset = idInBlockTable * ropeShape.pageStride; + } + + if (kvLayout == KVLAYOUT::NZ) { + offset += static_cast(startPos.n2Idx * shape.blockSize * shape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * shape.blockSize; + keyRopeOffset += static_cast(startPos.n2Idx * ropeShape.blockSize * ropeShape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * ropeShape.blockSize; + LocalTensor tmpNopeDstTensor = nopeTensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpNopeSrcTensor = nopeGmTensor[offset]; + DataCopyGmNZToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, (shape.copyRowNumAlign - copyRowCnt), (shape.blockSize - copyRowCnt), shape.actHeadDim); + + LocalTensor tmpRopeDstTensor = ropeTensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpRopeSrcTensor = ropeGmTensor[keyRopeOffset]; + DataCopyGmNZToL1(tmpRopeDstTensor, tmpRopeSrcTensor, copyRowCnt, (ropeShape.copyRowNumAlign - copyRowCnt), (ropeShape.blockSize - copyRowCnt), ropeShape.actHeadDim); + } else { + uint64_t dStride = shape.headDim; + uint64_t dRopeStride = ropeShape.headDim; + if (kvLayout == KVLAYOUT::BBH) { + offset += static_cast(startPos.n2Idx * shape.headDim) + remainRowCnt * shape.headDim * shape.headNum + startPos.dIdx; + keyRopeOffset += static_cast(startPos.n2Idx * ropeShape.headDim) + remainRowCnt * ropeShape.headDim * ropeShape.headNum; + dStride = shape.headDim * shape.headNum; + dRopeStride = ropeShape.headDim * ropeShape.headNum; + } else{ + offset += static_cast(startPos.n2Idx * shape.headDim * shape.blockSize) + remainRowCnt * shape.headDim + startPos.dIdx; + keyRopeOffset += static_cast(startPos.n2Idx * ropeShape.headDim * ropeShape.blockSize) + remainRowCnt * ropeShape.headDim; + } + + uint32_t dValue = shape.actHeadDim; + uint32_t srcDValue = dStride; + uint32_t dRopeValue = ropeShape.actHeadDim; + uint32_t srcRopeDValue = dRopeStride; + LocalTensor tmpNopeDstTensor = nopeTensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpNopeSrcTensor = nopeGmTensor[offset]; + DataCopyGmNDToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, shape.copyRowNumAlign, dValue, srcDValue); + + LocalTensor tmpRopeDstTensor = ropeTensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpRopeSrcTensor = ropeGmTensor[keyRopeOffset]; + DataCopyGmNDToL1(tmpRopeDstTensor, tmpRopeSrcTensor, copyRowCnt, shape.copyRowNumAlign, dRopeValue, srcRopeDValue); + } + copyFinishRowCnt += copyRowCnt; + curS2Idx += copyRowCnt; + } +} + +template +__aicore__ inline void GmCopyInToL1HasRopePA(LocalTensor& nopeTensor, LocalTensor& ropeTensor, + GlobalTensor& nopeGmTensor, GlobalTensor& ropeGmTensor, + GlobalTensor& blockTableGm, KVLAYOUT kvLayout, + const PAShape &shape, + const PAShape &ropeShape, + const Position &startPos) +{ + uint32_t copyFinishRowCnt = 0; + uint64_t blockTableBaseOffset = startPos.bIdx * shape.maxblockNumPerBatch; // 块表的基偏移量 + uint32_t curS2Idx = startPos.s2Offset; + uint32_t blockElementCnt = 32U / sizeof(L1Type); // 每个块的元素数量 + // ropeshape的M方向与nopeshape保持一样， 此处只判断nopeshape的 + while(copyFinishRowCnt < shape.copyRowNum){ + uint64_t blockIdOffset = curS2Idx / shape.blockSize; // 获取block table上的索引 + uint64_t remainRowCnt = curS2Idx % shape.blockSize; // 获取在单个块上超出的行数 + uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset); // 从block table上获取的编号 + //计算可以拷贝行数 + uint32_t copyRowCnt = shape.blockSize - remainRowCnt; // 一次只能处理一个Block + if (copyFinishRowCnt + copyRowCnt > shape.copyRowNum){ + copyRowCnt = shape.copyRowNum - copyFinishRowCnt; // 一个block未拷满 + } + uint64_t offset = idInBlockTable * shape.blockSize * shape.headNum * shape.headDim; // PA的偏移 + uint64_t keyRopeOffset = idInBlockTable * ropeShape.blockSize * ropeShape.headNum * ropeShape.headDim; + if (kvLayout == KVLAYOUT::NZ) { + offset += static_cast(startPos.n2Idx * shape.blockSize * shape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * shape.blockSize; + keyRopeOffset += static_cast(startPos.n2Idx * ropeShape.blockSize * ropeShape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * ropeShape.blockSize; + LocalTensor tmpNopeDstTensor = nopeTensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpNopeSrcTensor = nopeGmTensor[offset]; + DataCopyGmNZToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, (shape.copyRowNumAlign - copyRowCnt), (shape.blockSize - copyRowCnt), shape.actHeadDim); + + LocalTensor tmpRopeDstTensor = ropeTensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpRopeSrcTensor = ropeGmTensor[keyRopeOffset]; + DataCopyGmNZToL1(tmpRopeDstTensor, tmpRopeSrcTensor, copyRowCnt, (ropeShape.copyRowNumAlign - copyRowCnt), (ropeShape.blockSize - copyRowCnt), ropeShape.actHeadDim); + } else { + uint64_t dStride = shape.headDim; + uint64_t dRopeStride = ropeShape.headDim; + if (kvLayout == KVLAYOUT::BBH) { + offset += static_cast(startPos.n2Idx * shape.headDim) + remainRowCnt * shape.headDim * shape.headNum + startPos.dIdx; + keyRopeOffset += static_cast(startPos.n2Idx * ropeShape.headDim) + remainRowCnt * ropeShape.headDim * ropeShape.headNum; + dStride = shape.headDim * shape.headNum; + dRopeStride = ropeShape.headDim * ropeShape.headNum; + } else{ + offset += static_cast(startPos.n2Idx * shape.headDim * shape.blockSize) + remainRowCnt * shape.headDim + startPos.dIdx; + keyRopeOffset += static_cast(startPos.n2Idx * ropeShape.headDim * ropeShape.blockSize) + remainRowCnt * ropeShape.headDim; + } + + uint32_t dValue = shape.actHeadDim; + uint32_t srcDValue = dStride; + uint32_t dRopeValue = ropeShape.actHeadDim; + uint32_t srcRopeDValue = dRopeStride; + LocalTensor tmpNopeDstTensor = nopeTensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpNopeSrcTensor = nopeGmTensor[offset]; + DataCopyGmNDToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, shape.copyRowNumAlign, dValue, srcDValue); + + LocalTensor tmpRopeDstTensor = ropeTensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpRopeSrcTensor = ropeGmTensor[keyRopeOffset]; + DataCopyGmNDToL1(tmpRopeDstTensor, tmpRopeSrcTensor, copyRowCnt, shape.copyRowNumAlign, dRopeValue, srcRopeDValue); + } + copyFinishRowCnt += copyRowCnt; + curS2Idx += copyRowCnt; + } +} + +template +__aicore__ inline void GmCopyInToL1PA(LocalTensor& l1Tensor, GlobalTensor& gmTensor, + GlobalTensor& blockTableGm, KVLAYOUT kvLayout, + const PAShape &shape, const Position &startPos) +{ + uint32_t copyFinishRowCnt = 0; + uint64_t blockTableBaseOffset = startPos.bIdx * shape.maxblockNumPerBatch; // 块表的基偏移量 + uint32_t curS2Idx = startPos.s2Offset; + uint32_t blockElementCnt = 32U / sizeof(L1Type); // 每个块的元素数量 + while(copyFinishRowCnt < shape.copyRowNum){ + uint64_t blockIdOffset = curS2Idx / shape.blockSize; // 获取block table上的索引 + uint64_t remainRowCnt = curS2Idx % shape.blockSize; // 获取在单个块上超出的行数 + uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset); // 从block table上获取的编号 + //计算可以拷贝行数 + uint32_t copyRowCnt = shape.blockSize - remainRowCnt; // 一次只能处理一个Block + if (copyFinishRowCnt + copyRowCnt > shape.copyRowNum){ + copyRowCnt = shape.copyRowNum - copyFinishRowCnt; // 一个block未拷满 + } + uint64_t offset = idInBlockTable * shape.blockSize * shape.headNum * shape.headDim; // PA的偏移 + if (kvLayout == KVLAYOUT::NZ) { + offset += static_cast(startPos.n2Idx * shape.blockSize * shape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * shape.blockSize; + + LocalTensor tmpNopeDstTensor = l1Tensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpNopeSrcTensor = gmTensor[offset]; + DataCopyGmNZToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, (shape.copyRowNumAlign - copyRowCnt), (shape.blockSize - copyRowCnt), shape.actHeadDim); + } else { + uint64_t dStride = shape.headDim; + if (kvLayout == KVLAYOUT::BBH) { + offset += static_cast(startPos.n2Idx * shape.headDim) + remainRowCnt * shape.headDim * shape.headNum + startPos.dIdx; + dStride = shape.headDim * shape.headNum; + } else { + offset += static_cast(startPos.n2Idx * shape.headDim * shape.blockSize) + remainRowCnt * shape.headDim + startPos.dIdx; + } + + uint32_t dValue = shape.actHeadDim; + uint32_t srcDValue = dStride; + LocalTensor tmpNopeDstTensor = l1Tensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpNopeSrcTensor = gmTensor[offset]; + DataCopyGmNDToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, shape.copyRowNumAlign, dValue, srcDValue); + } + copyFinishRowCnt += copyRowCnt; + curS2Idx += copyRowCnt; + } +} + +template +__aicore__ inline void GmScaleCopyInToL1PAForND(LocalTensor& l1Tensor, GlobalTensor& gmTensor, + GlobalTensor& blockTableGm, KVLAYOUT kvLayout, + const PAShape &shape, const Position &startPos) +{ + uint32_t copyFinishRowCnt = 0; + uint64_t blockTableBaseOffset = startPos.bIdx * shape.maxblockNumPerBatch; + uint32_t curS2Idx = startPos.s2Offset; + constexpr uint32_t blockElementCnt = 32U / sizeof(L1Type); + while(copyFinishRowCnt < shape.copyRowNum) { + uint64_t blockIdOffset = curS2Idx / shape.blockSize; + uint64_t remainRowCnt = curS2Idx % shape.blockSize; + uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset); + uint32_t copyRowCnt = shape.blockSize - remainRowCnt; + if (copyFinishRowCnt + copyRowCnt > shape.copyRowNum) { + copyRowCnt = shape.copyRowNum - copyFinishRowCnt; + } + uint64_t offset = idInBlockTable * shape.blockSize * shape.headNum * shape.headDim; + if (kvLayout == KVLAYOUT::NZ) { + offset += static_cast(startPos.n2Idx * shape.blockSize * shape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * shape.blockSize; + + LocalTensor tmpNopeDstTensor = l1Tensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpNopeSrcTensor = gmTensor[offset]; + DataCopyGmNZToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, (shape.copyRowNumAlign - copyRowCnt), (shape.blockSize - copyRowCnt), shape.actHeadDim); + } else { + uint64_t dStride = shape.headDim; + if (kvLayout == KVLAYOUT::BBH) { + offset += static_cast(startPos.n2Idx * shape.headDim) + remainRowCnt * shape.headDim * shape.headNum + startPos.dIdx; + dStride = shape.headDim * shape.headNum; + } else { + offset += static_cast(startPos.n2Idx * shape.headDim * shape.blockSize) + remainRowCnt * shape.headDim + startPos.dIdx; + } + + uint32_t dValue = shape.actHeadDim; + uint32_t srcDValue = dStride; + + LocalTensor tmpNopeDstTensor = l1Tensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpNopeSrcTensor = gmTensor[offset * 2]; + DataCopyGmScaleNDToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, copyRowCnt, dValue, srcDValue); + } + copyFinishRowCnt += copyRowCnt; + curS2Idx += copyRowCnt; + } +} + +template +__aicore__ inline void GmScaleCopyInToL1PAForDN(LocalTensor& l1Tensor, GlobalTensor& gmTensor, + GlobalTensor& blockTableGm, KVLAYOUT kvLayout, + const PAShape &shape, const Position &startPos) +{ + uint32_t copyFinishRowCnt = 0; + uint64_t blockTableBaseOffset = startPos.bIdx * shape.maxblockNumPerBatch; + uint32_t curS2Idx = startPos.s2Offset; + constexpr uint32_t blockElementCnt = 32U / sizeof(L1Type); + while(copyFinishRowCnt < shape.copyRowNum) { + uint64_t blockIdOffset = curS2Idx / shape.blockSize; + uint64_t remainRowCnt = curS2Idx % shape.blockSize; + uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset); + uint32_t copyRowCnt = shape.blockSize - remainRowCnt; + if (copyFinishRowCnt + copyRowCnt > shape.copyRowNum) { + copyRowCnt = shape.copyRowNum - copyFinishRowCnt; + } + uint64_t offset = idInBlockTable * shape.blockSize * shape.headNum * shape.headDim; + if (kvLayout == KVLAYOUT::NZ) { + offset += static_cast(startPos.n2Idx * shape.blockSize * shape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * shape.blockSize; + + LocalTensor tmpNopeDstTensor = l1Tensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpNopeSrcTensor = gmTensor[offset]; + DataCopyGmNZToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, (shape.copyRowNumAlign - copyRowCnt), (shape.blockSize - copyRowCnt), shape.actHeadDim); + } else { + uint64_t dStride = shape.headDim; + if (kvLayout == KVLAYOUT::BBH) { + offset += static_cast(startPos.n2Idx * shape.headDim) + remainRowCnt * shape.headDim * shape.headNum + startPos.dIdx; + dStride = shape.headDim * shape.headNum; + } else { + offset += static_cast(startPos.n2Idx * shape.headDim * shape.blockSize) + remainRowCnt * shape.headDim + startPos.dIdx; + } + + uint32_t dValue = shape.actHeadDim; + uint32_t srcDValue = dStride; + LocalTensor tmpNopeDstTensor = l1Tensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpNopeSrcTensor = gmTensor[offset]; + + DataCopyGmScaleDNToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, copyRowCnt, dValue, srcDValue); + } + copyFinishRowCnt += copyRowCnt; + curS2Idx += copyRowCnt; + } +} + +template +__aicore__ inline void CopyToL1Nd2Nz(const LocalTensor &l1Tensor, const GlobalTensor &gmTensor, + uint32_t nValue, uint32_t dValue, uint32_t srcDValue) +{ + Nd2NzParams gm2L1Nd2NzParams; + gm2L1Nd2NzParams.ndNum = 1; // ND矩阵的个数 + gm2L1Nd2NzParams.nValue = nValue; // 单个ND矩阵的实际行数，单位为元素个数 + gm2L1Nd2NzParams.dValue = dValue; // 单个ND矩阵的实际列数，单位为元素个数 + gm2L1Nd2NzParams.srcNdMatrixStride = 0; // 相邻ND矩阵起始地址之间的偏移， 单位为元素个数 + gm2L1Nd2NzParams.srcDValue = srcDValue; // 同一个ND矩阵中相邻行起始地址之间的偏移， 单位为元素个数 +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) || (__NPU_ARCH__ == 5102) + if constexpr (IsSameType::value || IsSameType::value || + IsSameType::value || IsSameType::value) { + gm2L1Nd2NzParams.dstNzC0Stride = (nValue + 31) >> 5 << 5; + } else { + gm2L1Nd2NzParams.dstNzC0Stride = (nValue + 15) >> 4 << 4; + } +#else + gm2L1Nd2NzParams.dstNzC0Stride = (nValue + 15) >> 4 << 4; // NZ矩阵相邻Block起始地址之间的偏移， 单位为Block个数 +#endif + gm2L1Nd2NzParams.dstNzNStride = 1; // 转换为NZ矩阵后，ND之间相邻两行在NZ矩阵中起始地址之间的偏移， 单位为Block个数 + gm2L1Nd2NzParams.dstNzMatrixStride = 0; // 两个NZ矩阵，起始地址之间的偏移， 单位为元素数量 + DataCopy(l1Tensor, gmTensor, gm2L1Nd2NzParams); +} + +template +__aicore__ inline void CopyScaleToL1Nd2Nz(const LocalTensor &l1Tensor, const GlobalTensor &gmTensor, + uint32_t nValue, uint32_t dValue, uint32_t srcDValue) +{ + Nd2NzParams gm2L1Nd2NzParams; + gm2L1Nd2NzParams.ndNum = 1; // ND矩阵的个数 + gm2L1Nd2NzParams.nValue = nValue / 2; // 单个ND矩阵的实际行数，单位为元素个数 + gm2L1Nd2NzParams.dValue = dValue; // 单个ND矩阵的实际列数，单位为元素个数 + gm2L1Nd2NzParams.srcNdMatrixStride = 0; // 相邻ND矩阵起始地址之间的偏移， 单位为元素个数 + gm2L1Nd2NzParams.srcDValue = srcDValue; // 同一个ND矩阵中相邻行起始地址之间的偏移， 单位为元素个数 + gm2L1Nd2NzParams.dstNzC0Stride = nValue / 2; // NZ矩阵相邻Block起始地址之间的偏移， 单位为Block个数 + gm2L1Nd2NzParams.dstNzNStride = 1; // 转换为NZ矩阵后，ND之间相邻两行在NZ矩阵中起始地址之间的偏移， 单位为Block个数 + gm2L1Nd2NzParams.dstNzMatrixStride = gm2L1Nd2NzParams.nValue; // 两个NZ矩阵，起始地址之间的偏移， 单位为元素数量 + + LocalTensor l1TensorCast = l1Tensor.template ReinterpretCast(); + GlobalTensor gmTensorCast; + gmTensorCast.SetGlobalBuffer(((__gm__ bfloat16_t*)(gmTensor.GetPhyAddr()))); + DataCopy(l1TensorCast, gmTensorCast, gm2L1Nd2NzParams); +} + +template +__aicore__ inline void CopyScaleToL1Dn2Nz(const LocalTensor &l1Tensor, const GlobalTensor &gmTensor, + uint32_t nValue, uint32_t dValue, uint32_t srcDValue) +{ + Dn2NzParams gm2L1Dn2NzParams; + gm2L1Dn2NzParams.dnNum = 1; // ND矩阵的个数 + gm2L1Dn2NzParams.nValue = nValue / 2; // 单个DN矩阵的实际列数，单位为元素个数 + gm2L1Dn2NzParams.dValue = dValue; // 单个DN矩阵的实际行数，单位为元素个数 + gm2L1Dn2NzParams.srcDnMatrixStride = 0; // 相邻Dn矩阵起始地址之间的偏移， 单位为元素个数 + gm2L1Dn2NzParams.srcDValue = srcDValue / 2; // 同一个Dn矩阵中相邻行起始地址之间的偏移， 单位为元素个数 + gm2L1Dn2NzParams.dstNzC0Stride = nValue / 2; + gm2L1Dn2NzParams.dstNzNStride = 1; // 转换为NZ矩阵后，ND之间相邻两行在NZ矩阵中起始地址之间的偏移， 单位为Block个数 + gm2L1Dn2NzParams.dstNzMatrixStride = gm2L1Dn2NzParams.nValue; // 两个NZ矩阵，起始地址之间的偏移， 单位为元素数量 + + LocalTensor l1TensorCast = l1Tensor.template ReinterpretCast(); + GlobalTensor gmTensorCast; + gmTensorCast.SetGlobalBuffer(((__gm__ bfloat16_t*)(gmTensor.GetPhyAddr()))); + DataCopy(l1TensorCast, gmTensorCast, gm2L1Dn2NzParams); +} + +template +__aicore__ inline void CopyToL1Nd2NzGS1Merge(const LocalTensor &l1Tensor, const GlobalTensor &gmTensor, + uint32_t ndNum, uint32_t nValue, uint32_t dValue, uint32_t srcNdMatrixStride, uint32_t srcDValue, uint32_t dstNzC0Stride) // BSNGD 合轴拷贝 +{ + Nd2NzParams gm2L1Nd2NzParams; + gm2L1Nd2NzParams.ndNum = ndNum; // ND矩阵的个数 + gm2L1Nd2NzParams.nValue = nValue; // 单个ND矩阵的实际行数，单位为元素个数 + gm2L1Nd2NzParams.dValue = dValue; // 单个ND矩阵的实际列数，单位为元素个数 + gm2L1Nd2NzParams.srcNdMatrixStride = srcNdMatrixStride; // 相邻ND矩阵起始地址之间的偏移， 单位为元素个数 + gm2L1Nd2NzParams.srcDValue = srcDValue; // 同一个ND矩阵中相邻行起始地址之间的偏移， 单位为元素个数 +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) || (__NPU_ARCH__ == 5102) + if constexpr (IsSameType::value || IsSameType::value || + IsSameType::value || IsSameType::value) { + gm2L1Nd2NzParams.dstNzC0Stride = (dstNzC0Stride + 31) >> 5 << 5; // NZ矩阵相邻Block起始地址之间的偏移，单位为Block个数，32对齐 + } else { + gm2L1Nd2NzParams.dstNzC0Stride = (dstNzC0Stride + 15) >> 4 << 4; // NZ矩阵相邻Block起始地址之间的偏移，单位为Block个数，16对齐 + } +#else + gm2L1Nd2NzParams.dstNzC0Stride = (dstNzC0Stride + 15) >> 4 << 4; // NZ矩阵相邻Block起始地址之间的偏移，单位为Block个数，16对齐 +#endif + gm2L1Nd2NzParams.dstNzNStride = 1; // 转换为NZ矩阵后，ND之间相邻两行在NZ矩阵中起始地址之间的偏移， 单位为Block个数 + gm2L1Nd2NzParams.dstNzMatrixStride = nValue * 32 / sizeof(INPUT_T); // 两个NZ矩阵，起始地址之间的偏移， 单位为元素数量 + DataCopy(l1Tensor, gmTensor, gm2L1Nd2NzParams); +} +#endif diff --git a/csrc/ascend/attention/common/op_kernel/FixpipeOut.h b/csrc/ascend/attention/common/op_kernel/FixpipeOut.h new file mode 100644 index 000000000..fab78d20c --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/FixpipeOut.h @@ -0,0 +1,56 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file FixpipeOut.h + * \brief + */ +#ifndef FIXPIPEOUT_H +#define FIXPIPEOUT_H + +constexpr FixpipeConfig PFA_CFG_ROW_MAJOR_UB = {CO2Layout::ROW_MAJOR, true}; // ROW_MAJOR: 使能NZ2ND，输出数据格式为ND格式; true: 用于用户指定目的地址的位置是否是UB +constexpr FixpipeConfig PFA_CFG_ROW_MAJOR_GM = {CO2Layout::ROW_MAJOR, false}; // ROW_MAJOR: 使能NZ2ND，输出数据格式为ND格式; true: 用于用户指定目的地址的位置是否是UB +constexpr FixpipeConfig FA_CFG_NZ_UB = {CO2Layout::NZ, true}; // 不使能NZ2ND，输出数据格式为NZ格式; true: 用于用户指定目的地址的位置是否是UB + +struct fixpipeOutParams { + uint32_t fixpOutMSize; + uint32_t fixpOutNSize; +}; + +template +__aicore__ inline void FixpipeMmCopyOutToUB(LocalTensor& mmResUb, LocalTensor& L0CTensor, const fixpipeOutParams& fixpOutParam) +{ + FixpipeParamsC310 L0C2UbFixpParams; // L0C->UB + L0C2UbFixpParams.nSize = (fixpOutParam.fixpOutNSize + 7) >> 3 << 3; // L0C上的bmm1结果矩阵N方向的size大小；同mmadParams.n；8个元素（32B)对齐 + L0C2UbFixpParams.mSize = (fixpOutParam.fixpOutMSize + 1) >> 1 << 1; // 有效数据不足16行，只需输出部分行即可;L0C上的bmm1结果矩阵M方向的size大小必须是偶数 + L0C2UbFixpParams.srcStride = ((L0C2UbFixpParams.mSize + 15) >> 4) << 4; // L0C上matmul结果相邻连续数据片断间隔（前面一个数据块的头与后面数据块的头的间隔），单位为16 *sizeof(T) //源NZ矩阵中相邻Z排布的起始地址偏移 + L0C2UbFixpParams.dstStride = (L0C2UbFixpParams.nSize + 15) >> 4 << 4; // mmResUb上两行之间的间隔，单位：element。 // 128：根据比对dump文件得到，ND方案(S1 * S2)时脏数据用mask剔除 + L0C2UbFixpParams.dualDstCtl = 1; // 双目标模式，按M维度拆分， M / 2 * N写入每个UB，M必须为2的倍数 + L0C2UbFixpParams.params.ndNum = 1; + L0C2UbFixpParams.params.srcNdStride = 0; + L0C2UbFixpParams.params.dstNdStride = 0; + Fixpipe(mmResUb, L0CTensor, L0C2UbFixpParams); // 将matmul结果从L0C搬运到UB +} + +template +__aicore__ inline void FixpipeMmCopyOutToGm(GlobalTensor& mmResGm,LocalTensor& L0CTensor, const fixpipeOutParams& fixpOutParam) +{ + FixpipeParamsC310 L0C2GmFixpParams; // L0C->Gm + L0C2GmFixpParams.nSize = (fixpOutParam.fixpOutNSize + 7) >> 3 << 3; // L0C上的bmm1结果矩阵N方向的size大小；同mmadParams.n；8个元素（32B)对齐；分档计算且vector1中通过mask筛选出实际有效值 + L0C2GmFixpParams.mSize = (fixpOutParam.fixpOutMSize + 1) >> 1 << 1; // 有效数据不足16行，只需输出部分行即可;L0C上的bmm1结果矩阵M方向的size大小；同mmadParams.m + L0C2GmFixpParams.srcStride = ((L0C2GmFixpParams.mSize + 15) >> 4) << 4; // L0C上bmm1结果相邻连续数据片断间隔（前面一个数据块的头与后面数据块的头的间隔） + L0C2GmFixpParams.dstStride = (L0C2GmFixpParams.nSize + 15) >> 4 << 4; // mmResGm上两行之间的间隔 + L0C2GmFixpParams.dualDstCtl = 1; + L0C2GmFixpParams.params.ndNum = 1; + L0C2GmFixpParams.params.srcNdStride = 0; + L0C2GmFixpParams.params.dstNdStride = 0; + Fixpipe(mmResGm, L0CTensor, L0C2GmFixpParams); // 将matmul结果从L0C搬运到Gm +} +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_aligned128_no_update_sfa.h b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_aligned128_no_update_sfa.h new file mode 100644 index 000000000..48f83d581 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_aligned128_no_update_sfa.h @@ -0,0 +1,139 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_basic_block_aligned128_no_update_sfa.h + * \brief + */ +#ifndef VF_BASIC_BLOCK_ALIGNED128_NO_UPDATE_SFA_H +#define VF_BASIC_BLOCK_ALIGNED128_NO_UPDATE_SFA_H + +#include "vf_basic_block_utils.h" + +using namespace regbaseutil; + +namespace FaVectorApi { +// no update, originN == 128 +template +__simd_vf__ void ProcessVec1NoUpdateImpl128VF( + __ubuf__ T2 * expUb, __ubuf__ T * expSumUb, __ubuf__ T * maxUb, __ubuf__ T * maxUbStart, + __ubuf__ T * srcUb, const uint32_t blockStride, const uint32_t repeatStride, + const uint16_t m, const T scale, const T minValue) +{ + AscendC::MicroAPI::RegTensor vreg_input_x; + AscendC::MicroAPI::RegTensor vreg_input_x_unroll; + AscendC::MicroAPI::RegTensor vreg_max_tmp; + AscendC::MicroAPI::RegTensor vreg_input_max; + AscendC::MicroAPI::RegTensor vreg_max_brc; + AscendC::MicroAPI::RegTensor vreg_exp_sum; + AscendC::MicroAPI::RegTensor vreg_exp_even; + AscendC::MicroAPI::RegTensor vreg_exp_odd; + + // bfloat16_t + AscendC::MicroAPI::RegTensor vreg_exp_even_bf16; + AscendC::MicroAPI::RegTensor vreg_exp_odd_bf16; + AscendC::MicroAPI::RegTensor vreg_exp_bf16; + // half + AscendC::MicroAPI::RegTensor vreg_exp_even_fp16; + AscendC::MicroAPI::RegTensor vreg_exp_odd_fp16; + AscendC::MicroAPI::RegTensor vreg_exp_fp16; + + AscendC::MicroAPI::UnalignRegForStore ureg_max; + AscendC::MicroAPI::UnalignRegForStore ureg_exp_sum; + + AscendC::MicroAPI::MaskReg preg_all = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg preg_all_b16 = + AscendC::MicroAPI::CreateMask(); + + for (uint16_t i = 0; i < m; ++i) { + AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize); + AscendC::MicroAPI::LoadAlign(vreg_input_x_unroll, srcUb + floatRepSize + i * s2BaseSize); + + AscendC::MicroAPI::Muls(vreg_input_x, vreg_input_x, scale, preg_all); // Muls(scale) + AscendC::MicroAPI::Muls(vreg_input_x_unroll, vreg_input_x_unroll, scale, preg_all); + + AscendC::MicroAPI::StoreAlign( + (__ubuf__ T *&)srcUb + i * s2BaseSize, vreg_input_x, preg_all); + AscendC::MicroAPI::StoreAlign( + (__ubuf__ T *&)srcUb + floatRepSize + i * s2BaseSize, vreg_input_x_unroll, preg_all); + AscendC::MicroAPI::Max(vreg_max_tmp, vreg_input_x, vreg_input_x_unroll, preg_all); + + AscendC::MicroAPI::Reduce( + vreg_input_max, vreg_max_tmp, preg_all); + AscendC::MicroAPI::StoreUnAlign( + ((__ubuf__ T *&)maxUb), vreg_input_max, ureg_max, 1); + } + + AscendC::MicroAPI::StoreUnAlignPost( + ((__ubuf__ T *&)maxUb), ureg_max, 0); + AscendC::MicroAPI::LocalMemBar(); + + for (uint16_t i = 0; i < m; ++i) { + // maxUb is [S1, 1], BRC_B32 is reading one fp32 element and broadcast it to all 64 vreg element + AscendC::MicroAPI::LoadAlign( + vreg_max_brc, maxUbStart + i); + AscendC::MicroAPI::LoadAlign( + vreg_input_x, vreg_input_x_unroll, srcUb + i * s2BaseSize); + + AscendC::MicroAPI::ExpSub(vreg_exp_even, vreg_input_x, vreg_max_brc, preg_all); + AscendC::MicroAPI::ExpSub(vreg_exp_odd, vreg_input_x_unroll, vreg_max_brc, preg_all); + + // x_sum = sum(x_exp, axis=-1, keepdims=True) + AscendC::MicroAPI::Add(vreg_exp_sum, vreg_exp_even, vreg_exp_odd, preg_all); + AscendC::MicroAPI::Reduce( + vreg_exp_sum, vreg_exp_sum, preg_all); + AscendC::MicroAPI::StoreUnAlign( + ((__ubuf__ T *&)expSumUb), vreg_exp_sum, ureg_exp_sum, 1); + + if constexpr (IsSameType::value) { + AscendC::MicroAPI::Cast(vreg_exp_even_bf16, vreg_exp_even, preg_all); + AscendC::MicroAPI::Cast(vreg_exp_odd_bf16, vreg_exp_odd, preg_all); + AscendC::MicroAPI::Or((RegTensor&)vreg_exp_bf16, (RegTensor&)vreg_exp_even_bf16, + (RegTensor&)vreg_exp_odd_bf16, preg_all_b16); + AscendC::MicroAPI::StoreAlign( + ((__ubuf__ T2 *&)expUb), vreg_exp_bf16, blockStride, repeatStride, preg_all_b16); + } else if constexpr (IsSameType::value) { + AscendC::MicroAPI::Cast(vreg_exp_even_fp16, vreg_exp_even, preg_all); + AscendC::MicroAPI::Cast(vreg_exp_odd_fp16, vreg_exp_odd, preg_all); + AscendC::MicroAPI::Or((RegTensor&)vreg_exp_fp16, (RegTensor&)vreg_exp_even_fp16, + (RegTensor&)vreg_exp_odd_fp16, preg_all_b16); + AscendC::MicroAPI::StoreAlign( + ((__ubuf__ T2 *&)expUb), vreg_exp_fp16, blockStride, repeatStride, preg_all_b16); + } + } + AscendC::MicroAPI::StoreUnAlignPost( + ((__ubuf__ T *&)expSumUb), ureg_exp_sum, 0); +} + +// no update, originN == 128 +template +__aicore__ inline void ProcessVec1NoUpdateImpl128( + const LocalTensor& dstTensor, const LocalTensor& srcTensor, + const LocalTensor& expSumTensor, const LocalTensor& maxTensor, const LocalTensor& inMaxTensor, + const LocalTensor& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue) +{ + // 写的时候固定用65或者33的stride去写，因为正向目前使能settail之后mm2的s1方向必须算满128或者64行 + // stride, high 16bits: blockStride (m*16*2/32), low 16bits: repeatStride (1) + const uint32_t blockStride = s1BaseSize >> 1 | 0x1; + const uint32_t repeatStride = 1; + __ubuf__ T2 * expUb = (__ubuf__ T2*)dstTensor.GetPhyAddr(); + __ubuf__ T * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr(); + __ubuf__ T * maxUb = (__ubuf__ T*)maxTensor.GetPhyAddr(); + __ubuf__ T * maxUbStart = (__ubuf__ T*)maxTensor.GetPhyAddr(); + __ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr(); + + ProcessVec1NoUpdateImpl128VF( + expUb, expSumUb, maxUb, maxUbStart, srcUb, blockStride, repeatStride, m, scale, minValue); +} +} // namespace + +#endif // VF_BASIC_BLOCK_ALIGNED128_NO_UPDATE_SFA_H diff --git a/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_aligned128_update_sfa.h b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_aligned128_update_sfa.h new file mode 100644 index 000000000..07e35a3db --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_aligned128_update_sfa.h @@ -0,0 +1,143 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_basic_block_aligned128_update_sfa.h + * \brief + */ +#ifndef VF_BASIC_BLOCK_ALIGNED128_UPDATE_SFA_H +#define VF_BASIC_BLOCK_ALIGNED128_UPDATE_SFA_H + +#include "vf_basic_block_utils.h" + +using namespace regbaseutil; + +namespace FaVectorApi { +// update, originN == 128 +template +__simd_vf__ void ProcessVec1UpdateImpl128VF( + __ubuf__ T2 * expUb, __ubuf__ T * srcUb, __ubuf__ T * inMaxUb, + __ubuf__ T * tmpExpSumUb, __ubuf__ T * tmpMaxUb, __ubuf__ T * tmpMaxUb2, const uint32_t blockStride, + const uint32_t repeatStride, const uint16_t m, const T scale, const T minValue) +{ + AscendC::MicroAPI::RegTensor vreg_input_x; + AscendC::MicroAPI::RegTensor vreg_input_x_unroll; + AscendC::MicroAPI::RegTensor vreg_max_tmp; + AscendC::MicroAPI::RegTensor vreg_in_max; + AscendC::MicroAPI::RegTensor vreg_max_new; + AscendC::MicroAPI::RegTensor vreg_max_brc; + AscendC::MicroAPI::RegTensor vreg_cur_max; + AscendC::MicroAPI::RegTensor vreg_exp_sum; + AscendC::MicroAPI::RegTensor vreg_in_exp_sum; + AscendC::MicroAPI::RegTensor vreg_exp_even; + AscendC::MicroAPI::RegTensor vreg_exp_odd; + + // bfloat16_t + AscendC::MicroAPI::RegTensor vreg_exp_even_bf16; + AscendC::MicroAPI::RegTensor vreg_exp_odd_bf16; + AscendC::MicroAPI::RegTensor vreg_exp_bf16; + // half + AscendC::MicroAPI::RegTensor vreg_exp_even_fp16; + AscendC::MicroAPI::RegTensor vreg_exp_odd_fp16; + AscendC::MicroAPI::RegTensor vreg_exp_fp16; + + AscendC::MicroAPI::UnalignRegForStore ureg_max; + AscendC::MicroAPI::UnalignRegForStore ureg_exp_sum; + AscendC::MicroAPI::MaskReg preg_all = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg preg_all_b16 = + AscendC::MicroAPI::CreateMask(); + + // x_max = max(src, axis=-1, keepdims=True); x_max = Max(x_max, inMax) + for (uint16_t i = 0; i < m; ++i) { + AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize); + AscendC::MicroAPI::LoadAlign(vreg_input_x_unroll, srcUb + floatRepSize + i * s2BaseSize); + + AscendC::MicroAPI::Muls(vreg_input_x, vreg_input_x, scale, preg_all); // Muls(scale) + AscendC::MicroAPI::Muls(vreg_input_x_unroll, vreg_input_x_unroll, scale, preg_all); + AscendC::MicroAPI::StoreAlign( + (__ubuf__ T *&)srcUb + i * s2BaseSize, vreg_input_x, preg_all); + AscendC::MicroAPI::StoreAlign( + (__ubuf__ T *&)srcUb + floatRepSize + i * s2BaseSize, vreg_input_x_unroll, preg_all); + AscendC::MicroAPI::Max(vreg_max_tmp, vreg_input_x, vreg_input_x_unroll, preg_all); + AscendC::MicroAPI::Reduce( + vreg_max_tmp, vreg_max_tmp, preg_all); + AscendC::MicroAPI::StoreUnAlign( + ((__ubuf__ T *&)tmpMaxUb), vreg_max_tmp, ureg_max, 1); + } + AscendC::MicroAPI::StoreUnAlignPost( + ((__ubuf__ T *&)tmpMaxUb), ureg_max, 0); + AscendC::MicroAPI::LoadAlign(vreg_in_max, inMaxUb); + AscendC::MicroAPI::LocalMemBar(); + AscendC::MicroAPI::LoadAlign(vreg_cur_max, tmpMaxUb2); // 获取新的max[s1, 1] + AscendC::MicroAPI::Max(vreg_max_new, vreg_cur_max, vreg_in_max, preg_all); // 计算新、旧max的最大值 + AscendC::MicroAPI::StoreAlign( + (__ubuf__ T *&)tmpMaxUb2, vreg_max_new, preg_all); + AscendC::MicroAPI::LocalMemBar(); + + for (uint16_t i = 0; i < m; ++i) { + AscendC::MicroAPI::LoadAlign(vreg_max_brc, tmpMaxUb2 + i); + AscendC::MicroAPI::LoadAlign( + vreg_input_x, vreg_input_x_unroll, srcUb + i * s2BaseSize); + AscendC::MicroAPI::ExpSub(vreg_exp_even, vreg_input_x, vreg_max_brc, preg_all); + AscendC::MicroAPI::ExpSub(vreg_exp_odd, vreg_input_x_unroll, vreg_max_brc, preg_all); + + // x_sum = sum(x_exp, axis=-1, keepdims=True) + AscendC::MicroAPI::Add(vreg_exp_sum, vreg_exp_even, vreg_exp_odd, preg_all); + AscendC::MicroAPI::Reduce( + vreg_exp_sum, vreg_exp_sum, preg_all); + AscendC::MicroAPI::StoreUnAlign( + ((__ubuf__ T *&)tmpExpSumUb), vreg_exp_sum, ureg_exp_sum, 1); + + if constexpr (IsSameType::value) { + AscendC::MicroAPI::Cast(vreg_exp_even_bf16, vreg_exp_even, preg_all); + AscendC::MicroAPI::Cast(vreg_exp_odd_bf16, vreg_exp_odd, preg_all); + AscendC::MicroAPI::Or((RegTensor&)vreg_exp_bf16, (RegTensor&)vreg_exp_even_bf16, + (RegTensor&)vreg_exp_odd_bf16, preg_all_b16); + AscendC::MicroAPI::StoreAlign( + ((__ubuf__ T2 *&)expUb), vreg_exp_bf16, blockStride, repeatStride, preg_all_b16); + } else if constexpr (IsSameType::value) { + AscendC::MicroAPI::Cast(vreg_exp_even_fp16, vreg_exp_even, preg_all); + AscendC::MicroAPI::Cast(vreg_exp_odd_fp16, vreg_exp_odd, preg_all); + AscendC::MicroAPI::Or((RegTensor&)vreg_exp_fp16, (RegTensor&)vreg_exp_even_fp16, + (RegTensor&)vreg_exp_odd_fp16, preg_all_b16); + AscendC::MicroAPI::StoreAlign( + ((__ubuf__ T2 *&)expUb), vreg_exp_fp16, blockStride, repeatStride, preg_all_b16); + } + } + AscendC::MicroAPI::StoreUnAlignPost( + ((__ubuf__ T *&)tmpExpSumUb), ureg_exp_sum, 0); +} + +// update, originN == 128 +template +__aicore__ inline void ProcessVec1UpdateImpl128( + const LocalTensor& dstTensor, const LocalTensor& srcTensor, const LocalTensor& inMaxTensor, + const LocalTensor& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue) +{ + // 写的时候固定用65或者33的stride去写，因为正向目前使能settail之后mm2的s1方向必须算满128或者64行 + // stride, high 16bits: blockStride (m*16*2/32), low 16bits: repeatStride (1) + const uint32_t blockStride = s1BaseSize >> 1 | 0x1; + const uint32_t repeatStride = 1; + + __ubuf__ T2 * expUb = (__ubuf__ T2*)dstTensor.GetPhyAddr(); + __ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr(); + __ubuf__ T * inMaxUb = (__ubuf__ T*)inMaxTensor.GetPhyAddr(); + __ubuf__ T * tmpExpSumUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr(); + __ubuf__ T * tmpMaxUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64; + __ubuf__ T * tmpMaxUb2 = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64; + + ProcessVec1UpdateImpl128VF ( + expUb, srcUb, inMaxUb, tmpExpSumUb, tmpMaxUb, tmpMaxUb2, blockStride, repeatStride, m, scale, minValue); +} +} // namespace + +#endif // VF_BASIC_BLOCK_ALIGNED128_UPDATE_SFA_H diff --git a/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_unaligned128_no_update_sfa.h b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_unaligned128_no_update_sfa.h new file mode 100644 index 000000000..8e599df13 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_unaligned128_no_update_sfa.h @@ -0,0 +1,149 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_basic_block_unaligned128_no_update_sfa.h + * \brief + */ +#ifndef VF_BASIC_BLOCK_UNALIGNED128_NO_UPDATE_SFA_H +#define VF_BASIC_BLOCK_UNALIGNED128_NO_UPDATE_SFA_H + +#include "vf_basic_block_utils.h" + +using namespace regbaseutil; + +namespace FaVectorApi { + +template +__simd_vf__ void ProcessVec1NoUpdateGeneralImpl128VF( + __ubuf__ T2 * expUb, __ubuf__ T * expSumUb, __ubuf__ T * maxUb, __ubuf__ T * maxUbStart, + __ubuf__ T * srcUb, const uint32_t blockStride, const uint32_t repeatStride, + const uint16_t m, const T scale, const T minValue, uint32_t pltOriTailN, uint32_t pltTailN) +{ + AscendC::MicroAPI::RegTensor vreg_min; + AscendC::MicroAPI::RegTensor vreg_input_x; + AscendC::MicroAPI::RegTensor vreg_input_x_unroll; + AscendC::MicroAPI::RegTensor vreg_input_x_unroll_new; + AscendC::MicroAPI::RegTensor vreg_max_tmp; + AscendC::MicroAPI::RegTensor vreg_input_max; + AscendC::MicroAPI::RegTensor vreg_max_brc; + AscendC::MicroAPI::RegTensor vreg_exp_sum; + AscendC::MicroAPI::RegTensor vreg_exp_even; + AscendC::MicroAPI::RegTensor vreg_exp_odd; + + // bfloat16_t + AscendC::MicroAPI::RegTensor vreg_exp_even_bf16; + AscendC::MicroAPI::RegTensor vreg_exp_odd_bf16; + AscendC::MicroAPI::RegTensor vreg_exp_bf16; + // half + AscendC::MicroAPI::RegTensor vreg_exp_even_fp16; + AscendC::MicroAPI::RegTensor vreg_exp_odd_fp16; + AscendC::MicroAPI::RegTensor vreg_exp_fp16; + + AscendC::MicroAPI::UnalignRegForStore ureg_max; + AscendC::MicroAPI::UnalignRegForStore ureg_exp_sum; + + AscendC::MicroAPI::MaskReg preg_all = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg preg_all_b16 = + AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg preg_all_b8 = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg preg_tail_n = AscendC::MicroAPI::UpdateMask(pltTailN); + AscendC::MicroAPI::MaskReg preg_ori_tail_n = AscendC::MicroAPI::UpdateMask(pltOriTailN); + AscendC::MicroAPI::MaskReg preg_reduce_n = + AscendC::MicroAPI::CreateMask(); + + AscendC::MicroAPI::Duplicate(vreg_min, minValue); + for (uint16_t i = 0; i < m; ++i) { + AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize); + AscendC::MicroAPI::LoadAlign(vreg_input_x_unroll, srcUb + floatRepSize + i * s2BaseSize); + AscendC::MicroAPI::Muls(vreg_input_x, vreg_input_x, scale, preg_all); // Muls(scale) + AscendC::MicroAPI::Muls(vreg_input_x_unroll, vreg_input_x_unroll, scale, preg_ori_tail_n); + AscendC::MicroAPI::Select(vreg_input_x_unroll_new, vreg_input_x_unroll, vreg_min, preg_ori_tail_n); + AscendC::MicroAPI::StoreAlign( + (__ubuf__ T *&)srcUb + i * s2BaseSize, vreg_input_x, preg_all); + AscendC::MicroAPI::StoreAlign( + (__ubuf__ T *&)srcUb + floatRepSize + i * s2BaseSize, vreg_input_x_unroll_new, preg_tail_n); + + AscendC::MicroAPI::Max(vreg_max_tmp, vreg_input_x, vreg_input_x_unroll_new, preg_all); + AscendC::MicroAPI::Reduce( + vreg_input_max, vreg_max_tmp, preg_all); + AscendC::MicroAPI::StoreUnAlign( + ((__ubuf__ T *&)maxUb), vreg_input_max, ureg_max, 1); + } + + AscendC::MicroAPI::StoreUnAlignPost( + ((__ubuf__ T *&)maxUb), ureg_max, 0); + AscendC::MicroAPI::LocalMemBar(); + + for (uint16_t i = 0; i < m; ++i) { + AscendC::MicroAPI::LoadAlign(vreg_max_brc, maxUbStart + i); + AscendC::MicroAPI::LoadAlign( + vreg_input_x, vreg_input_x_unroll, srcUb + i * s2BaseSize); + AscendC::MicroAPI::ExpSub(vreg_exp_even, vreg_input_x, vreg_max_brc, preg_all); + AscendC::MicroAPI::ExpSub(vreg_exp_odd, vreg_input_x_unroll, vreg_max_brc, preg_all); + + // x_sum = sum(x_exp, axis=-1, keepdims=True) + AscendC::MicroAPI::Add(vreg_exp_sum, vreg_exp_even, vreg_exp_odd, preg_all); + AscendC::MicroAPI::Reduce( + vreg_exp_sum, vreg_exp_sum, preg_all); + AscendC::MicroAPI::StoreUnAlign( + ((__ubuf__ T *&)expSumUb), vreg_exp_sum, ureg_exp_sum, 1); + + if constexpr (IsSameType::value) { + AscendC::MicroAPI::Cast(vreg_exp_even_bf16, vreg_exp_even, preg_all); + AscendC::MicroAPI::Cast(vreg_exp_odd_bf16, vreg_exp_odd, preg_all); + AscendC::MicroAPI::Or((RegTensor&)vreg_exp_bf16, (RegTensor&)vreg_exp_even_bf16, + (RegTensor&)vreg_exp_odd_bf16, preg_all_b16); + AscendC::MicroAPI::StoreAlign( + ((__ubuf__ T2 *&)expUb), vreg_exp_bf16, blockStride, repeatStride, preg_all_b16); + } else if constexpr (IsSameType::value) { + AscendC::MicroAPI::Cast(vreg_exp_even_fp16, vreg_exp_even, preg_all); + AscendC::MicroAPI::Cast(vreg_exp_odd_fp16, vreg_exp_odd, preg_all); + AscendC::MicroAPI::Or((RegTensor&)vreg_exp_fp16, (RegTensor&)vreg_exp_even_fp16, + (RegTensor&)vreg_exp_odd_fp16, preg_all_b16); + AscendC::MicroAPI::StoreAlign( + ((__ubuf__ T2 *&)expUb), vreg_exp_fp16, blockStride, repeatStride, preg_all_b16); + } + } + AscendC::MicroAPI::StoreUnAlignPost( + ((__ubuf__ T *&)expSumUb), ureg_exp_sum, 0); +} + +// no update, 64 < originN <= 128 +template +__aicore__ inline void ProcessVec1NoUpdateGeneralImpl128( + const LocalTensor& dstTensor, const LocalTensor& srcTensor, + const LocalTensor& expSumTensor, const LocalTensor& maxTensor, const LocalTensor& inMaxTensor, + const LocalTensor& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue) +{ + // 写的时候固定用65或者33的stride去写，因为正向目前使能settail之后mm2的s1方向必须算满128或者64行 + // stride, high 16bits: blockStride (65*16*2/32)，单位block, low 16bits: repeatStride (1) + const uint32_t blockStride = s1BaseSize >> 1 | 0x1; + const uint32_t repeatStride = 1; + __ubuf__ T2 * expUb = (__ubuf__ T2*)dstTensor.GetPhyAddr(); + __ubuf__ T * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr(); + __ubuf__ T * maxUb = (__ubuf__ T*)maxTensor.GetPhyAddr(); + __ubuf__ T * maxUbStart = (__ubuf__ T*)maxTensor.GetPhyAddr(); + __ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr(); + + const uint32_t oriTailN = originN - floatRepSize; + const uint32_t tailN = s2BaseSize - floatRepSize; + uint32_t pltOriTailN = oriTailN; + uint32_t pltTailN = tailN; + + ProcessVec1NoUpdateGeneralImpl128VF( + expUb, expSumUb, maxUb, maxUbStart, srcUb, blockStride, repeatStride, m, scale, minValue, + pltOriTailN, pltTailN); +} +} // namespace + +#endif // VF_BASIC_BLOCK_UNALIGNED128_NO_UPDATE_SFA_H diff --git a/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_unaligned128_update_sfa.h b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_unaligned128_update_sfa.h new file mode 100644 index 000000000..17316bd36 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_unaligned128_update_sfa.h @@ -0,0 +1,159 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_basic_block_unaligned128_update_sfa.h + * \brief + */ +#ifndef VF_BASIC_BLOCK_UNALIGNED128_UPDATE_SFA_H +#define VF_BASIC_BLOCK_UNALIGNED128_UPDATE_SFA_H + +#include "vf_basic_block_utils.h" + +using namespace regbaseutil; + +namespace FaVectorApi { + +template +__simd_vf__ void ProcessVec1UpdateGeneralImpl128VF( + __ubuf__ T2 * expUb, __ubuf__ T * srcUb, __ubuf__ T * inMaxUb, + __ubuf__ T * tmpExpSumUb, __ubuf__ T * tmpMaxUb, __ubuf__ T * tmpMaxUb2, const uint32_t blockStride, + const uint32_t repeatStride, const uint16_t m, const T scale, const T minValue, uint32_t pltOriTailN, + uint32_t pltTailN, uint32_t pltN) +{ + AscendC::MicroAPI::RegTensor vreg_min; + AscendC::MicroAPI::RegTensor vreg_input_x; + AscendC::MicroAPI::RegTensor vreg_input_x_unroll; + AscendC::MicroAPI::RegTensor vreg_input_x_unroll_new; + AscendC::MicroAPI::RegTensor vreg_max_tmp; + AscendC::MicroAPI::RegTensor vreg_cur_max; + AscendC::MicroAPI::RegTensor vreg_max_new; + AscendC::MicroAPI::RegTensor vreg_exp_sum; + AscendC::MicroAPI::RegTensor vreg_in_max; + AscendC::MicroAPI::RegTensor vreg_max_brc; + AscendC::MicroAPI::RegTensor vreg_exp_even; + AscendC::MicroAPI::RegTensor vreg_exp_odd; + + // bfloat16_t + AscendC::MicroAPI::RegTensor vreg_exp_even_bf16; + AscendC::MicroAPI::RegTensor vreg_exp_odd_bf16; + AscendC::MicroAPI::RegTensor vreg_exp_bf16; + // half + AscendC::MicroAPI::RegTensor vreg_exp_even_fp16; + AscendC::MicroAPI::RegTensor vreg_exp_odd_fp16; + AscendC::MicroAPI::RegTensor vreg_exp_fp16; + + AscendC::MicroAPI::UnalignRegForStore ureg_max; + AscendC::MicroAPI::UnalignRegForStore ureg_exp_sum; + + AscendC::MicroAPI::MaskReg preg_all = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg preg_all_b16 = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg preg_n_b16 = AscendC::MicroAPI::UpdateMask(pltN); + AscendC::MicroAPI::MaskReg preg_tail_n = AscendC::MicroAPI::UpdateMask(pltTailN); + AscendC::MicroAPI::MaskReg preg_ori_tail_n = AscendC::MicroAPI::UpdateMask(pltOriTailN); + + AscendC::MicroAPI::Duplicate(vreg_min, minValue); + // x_max = max(src, axis=-1, keepdims=True); x_max = Max(x_max, inMax) + for (uint16_t i = 0; i < m; ++i) { + AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize); + AscendC::MicroAPI::LoadAlign(vreg_input_x_unroll, srcUb + floatRepSize + i * s2BaseSize); + AscendC::MicroAPI::Muls(vreg_input_x, vreg_input_x, scale, preg_all); // Muls(scale) + AscendC::MicroAPI::Muls(vreg_input_x_unroll, vreg_input_x_unroll, scale, preg_ori_tail_n); + AscendC::MicroAPI::Select(vreg_input_x_unroll_new, vreg_input_x_unroll, vreg_min, preg_ori_tail_n); + AscendC::MicroAPI::StoreAlign( + (__ubuf__ T *&)srcUb + i * s2BaseSize, vreg_input_x, preg_all); + AscendC::MicroAPI::StoreAlign( + (__ubuf__ T *&)srcUb + floatRepSize + i * s2BaseSize, vreg_input_x_unroll_new, preg_tail_n); + AscendC::MicroAPI::Max(vreg_max_tmp, vreg_input_x, vreg_input_x_unroll_new, preg_all); + AscendC::MicroAPI::Reduce( + vreg_cur_max, vreg_max_tmp, preg_all); + + AscendC::MicroAPI::StoreUnAlign( + ((__ubuf__ T *&)tmpMaxUb), vreg_cur_max, ureg_max, 1); + } + AscendC::MicroAPI::StoreUnAlignPost( + ((__ubuf__ T *&)tmpMaxUb), ureg_max, 0); + AscendC::MicroAPI::LoadAlign(vreg_in_max, inMaxUb); + AscendC::MicroAPI::LocalMemBar(); + AscendC::MicroAPI::LoadAlign(vreg_cur_max, tmpMaxUb2); // 获取新的max[s1, 1] + AscendC::MicroAPI::Max(vreg_max_new, vreg_cur_max, vreg_in_max, preg_all); // 计算新、旧max的最大值 + AscendC::MicroAPI::StoreAlign( + (__ubuf__ T *&)tmpMaxUb2, vreg_max_new, preg_all); + AscendC::MicroAPI::LocalMemBar(); + + for (uint16_t i = 0; i < m; ++i) { + AscendC::MicroAPI::LoadAlign( + vreg_max_brc, tmpMaxUb2 + i); + AscendC::MicroAPI::LoadAlign( + vreg_input_x, vreg_input_x_unroll, srcUb + i * s2BaseSize); + AscendC::MicroAPI::ExpSub(vreg_exp_even, vreg_input_x, vreg_max_brc, preg_all); + AscendC::MicroAPI::ExpSub(vreg_exp_odd, vreg_input_x_unroll, vreg_max_brc, preg_all); + + // x_sum = sum(x_exp, axis=-1, keepdims=True) + AscendC::MicroAPI::Add(vreg_exp_sum, vreg_exp_even, vreg_exp_odd, preg_all); + AscendC::MicroAPI::Reduce( + vreg_exp_sum, vreg_exp_sum, preg_all); + AscendC::MicroAPI::StoreUnAlign( + ((__ubuf__ T *&)tmpExpSumUb), vreg_exp_sum, ureg_exp_sum, 1); + + if constexpr (IsSameType::value) { + AscendC::MicroAPI::Cast(vreg_exp_even_bf16, vreg_exp_even, preg_all); + AscendC::MicroAPI::Cast(vreg_exp_odd_bf16, vreg_exp_odd, preg_all); + AscendC::MicroAPI::Or((RegTensor&)vreg_exp_bf16, (RegTensor&)vreg_exp_even_bf16, + (RegTensor&)vreg_exp_odd_bf16, preg_all_b16); + AscendC::MicroAPI::StoreAlign( + ((__ubuf__ T2 *&)expUb), vreg_exp_bf16, blockStride, repeatStride, preg_n_b16); + } else if constexpr (IsSameType::value) { + AscendC::MicroAPI::Cast(vreg_exp_even_fp16, vreg_exp_even, preg_all); + AscendC::MicroAPI::Cast(vreg_exp_odd_fp16, vreg_exp_odd, preg_all); + AscendC::MicroAPI::Or((RegTensor&)vreg_exp_fp16, (RegTensor&)vreg_exp_even_fp16, + (RegTensor&)vreg_exp_odd_fp16, preg_all_b16); + AscendC::MicroAPI::StoreAlign( + ((__ubuf__ T2 *&)expUb), vreg_exp_fp16, blockStride, repeatStride, preg_n_b16); + } + } + AscendC::MicroAPI::StoreUnAlignPost( + ((__ubuf__ T *&)tmpExpSumUb), ureg_exp_sum, 0); +} + + +// update, 64 < originN <= 128 +template +__aicore__ inline void ProcessVec1UpdateGeneralImpl128( + const LocalTensor& dstTensor, const LocalTensor& srcTensor, const LocalTensor& inMaxTensor, + const LocalTensor& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue) +{ + // 写的时候固定用65或者33的stride去写，因为正向目前使能settail之后mm2的s1方向必须算满128或者64行 + // stride, high 16bits: blockStride (m*16*2/32), low 16bits: repeatStride (1) + const uint32_t blockStride = s1BaseSize >> 1 | 0x1; + const uint32_t repeatStride = 1; + const uint32_t oriTailN = originN - floatRepSize; + const uint32_t tailN = s2BaseSize - floatRepSize; + uint32_t pltOriTailN = oriTailN; + uint32_t pltTailN = tailN; + uint32_t pltN = s2BaseSize; + + __ubuf__ T2 * expUb = (__ubuf__ T2*)dstTensor.GetPhyAddr(); + __ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr(); + __ubuf__ T * inMaxUb = (__ubuf__ T*)inMaxTensor.GetPhyAddr(); + __ubuf__ T * tmpExpSumUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr(); + __ubuf__ T * tmpMaxUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64; + __ubuf__ T * tmpMaxUb2 = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64; + + ProcessVec1UpdateGeneralImpl128VF( + expUb, srcUb, inMaxUb, tmpExpSumUb, tmpMaxUb, tmpMaxUb2, blockStride, repeatStride, + m, scale, minValue, pltOriTailN, pltTailN, pltN); +} +} // namespace + +#endif // VF_BASIC_BLOCK_UNALIGNED128_UPDATE_SFA_H diff --git a/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_unaligned64_no_update_sfa.h b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_unaligned64_no_update_sfa.h new file mode 100644 index 000000000..ebe58b638 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_unaligned64_no_update_sfa.h @@ -0,0 +1,129 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_basic_block_unaligned64_no_update_sfa.h + * \brief + */ +#ifndef VF_BASIC_BLOCK_UNALIGNED64_NO_UPDATE_SFA_H +#define VF_BASIC_BLOCK_UNALIGNED64_NO_UPDATE_SFA_H + +#include "vf_basic_block_utils.h" + +using namespace regbaseutil; + +namespace FaVectorApi { +template +__simd_vf__ void ProcessVec1NoUpdateImpl64VF( + __ubuf__ T2 * expUb, __ubuf__ T * expSumUb, __ubuf__ T * maxUb, __ubuf__ T * maxUbStart, + __ubuf__ T * srcUb, const uint32_t blockStride, const uint32_t repeatStride, + const uint16_t m, const T scale, const T minValue, uint32_t pltOriginalN, uint32_t pltSrcN) +{ + AscendC::MicroAPI::RegTensor vreg_min; + AscendC::MicroAPI::RegTensor vreg_input_x; + AscendC::MicroAPI::RegTensor vreg_input_max; + AscendC::MicroAPI::RegTensor vreg_max_brc; + AscendC::MicroAPI::RegTensor vreg_exp; + AscendC::MicroAPI::RegTensor vreg_exp_sum; + + // bfloat16_t + AscendC::MicroAPI::RegTensor vreg_exp_bf16; + AscendC::MicroAPI::RegTensor vreg_dst_even_bf16; + AscendC::MicroAPI::RegTensor vreg_dst_odd_bf16; + // half + AscendC::MicroAPI::RegTensor vreg_exp_fp16; + AscendC::MicroAPI::RegTensor vreg_dst_even_fp16; + AscendC::MicroAPI::RegTensor vreg_dst_odd_fp16; + + AscendC::MicroAPI::UnalignRegForStore ureg_max; + AscendC::MicroAPI::UnalignRegForStore ureg_exp_sum; + + AscendC::MicroAPI::MaskReg preg_all = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg preg_all_b16 = + AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg preg_src_n = AscendC::MicroAPI::UpdateMask(pltSrcN); + AscendC::MicroAPI::MaskReg preg_src_n_b16 = + AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg preg_ori_src_n = AscendC::MicroAPI::UpdateMask(pltOriginalN); + + // x_max = max(src, axis=-1, keepdims=True) + for (uint16_t i = 0; i < m; ++i) { + AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize); + AscendC::MicroAPI::Muls(vreg_input_x, vreg_input_x, scale, preg_ori_src_n); // Muls(scale) + AscendC::MicroAPI::StoreAlign( + (__ubuf__ T *&)srcUb + i * s2BaseSize, vreg_input_x, preg_src_n); + AscendC::MicroAPI::Reduce( + vreg_input_max, vreg_input_x, preg_ori_src_n); + AscendC::MicroAPI::StoreUnAlign( + ((__ubuf__ T *&)maxUb), vreg_input_max, ureg_max, 1); + } + AscendC::MicroAPI::StoreUnAlignPost( + ((__ubuf__ T *&)maxUb), ureg_max, 0); + AscendC::MicroAPI::LocalMemBar(); + + for (uint16_t i = 0; i < m; ++i) { + AscendC::MicroAPI::LoadAlign( + vreg_max_brc, maxUbStart + i); + AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize); + AscendC::MicroAPI::ExpSub(vreg_exp, vreg_input_x, vreg_max_brc, preg_ori_src_n); + + // x_sum = sum(x_exp, axis=-1, keepdims=True) + AscendC::MicroAPI::Reduce( + vreg_exp_sum, vreg_exp, preg_ori_src_n); + AscendC::MicroAPI::StoreUnAlign( + ((__ubuf__ T *&)expSumUb), vreg_exp_sum, ureg_exp_sum, 1); + + if constexpr (IsSameType::value) { + AscendC::MicroAPI::Cast(vreg_exp_bf16, vreg_exp, preg_all_b16); + AscendC::MicroAPI::DeInterleave(vreg_dst_even_bf16, vreg_dst_odd_bf16, + vreg_exp_bf16, vreg_exp_bf16); + AscendC::MicroAPI::StoreAlign( + ((__ubuf__ T2 *&)expUb), vreg_dst_even_bf16, blockStride, repeatStride, preg_src_n_b16); + } else if constexpr (IsSameType::value) { + AscendC::MicroAPI::Cast(vreg_exp_fp16, vreg_exp, preg_all_b16); + AscendC::MicroAPI::DeInterleave(vreg_dst_even_fp16, vreg_dst_odd_fp16, + vreg_exp_fp16, vreg_exp_fp16); + AscendC::MicroAPI::StoreAlign( + ((__ubuf__ T2 *&)expUb), vreg_dst_even_fp16, blockStride, repeatStride, preg_src_n_b16); + } + } + AscendC::MicroAPI::StoreUnAlignPost( + ((__ubuf__ T *&)expSumUb), ureg_exp_sum, 0); +} + +// no update, originN <= 64 +template +__aicore__ inline void ProcessVec1NoUpdateImpl64( + const LocalTensor& dstTensor, const LocalTensor& srcTensor, + const LocalTensor& expSumTensor, const LocalTensor& maxTensor, const LocalTensor& inMaxTensor, + const LocalTensor& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue) +{ + __ubuf__ T2 * expUb = (__ubuf__ T2*)dstTensor.GetPhyAddr(); + __ubuf__ T * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr(); + __ubuf__ T * maxUb = (__ubuf__ T*)maxTensor.GetPhyAddr(); + __ubuf__ T * maxUbStart = (__ubuf__ T*)maxTensor.GetPhyAddr(); + __ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr(); + + // 写的时候固定用65或者33的stride去写，因为正向目前使能settail之后mm2的s1方向必须算满128或者64行 + // stride, high 16bits: blockStride (m*16*2/32), low 16bits: repeatStride (1) + const uint32_t blockStride = s1BaseSize >> 1 | 0x1; + const uint32_t repeatStride = 1; + uint32_t pltOriginalN = originN; + uint32_t pltSrcN = s2BaseSize; + + ProcessVec1NoUpdateImpl64VF( + expUb, expSumUb, maxUb, maxUbStart, srcUb, blockStride, repeatStride, m, scale, minValue, + pltOriginalN, pltSrcN); +} +} // namespace + +#endif // VF_BASIC_BLOCK_UNALIGNED64_NO_UPDATE_SFA_H diff --git a/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_unaligned64_update_sfa.h b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_unaligned64_update_sfa.h new file mode 100644 index 000000000..0896b1374 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_unaligned64_update_sfa.h @@ -0,0 +1,141 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_basic_block_aligned64_update_sfa.h + * \brief + */ +#ifndef VF_BASIC_BLOCK_ALIGNED64_UPDATE_SFA_H +#define VF_BASIC_BLOCK_ALIGNED64_UPDATE_SFA_H + +#include "vf_basic_block_utils.h" + +using namespace regbaseutil; + +namespace FaVectorApi { +// update, originN <= 64 +template +__simd_vf__ void ProcessVec1UpdateImpl64VF( + __ubuf__ T2 * expUb, __ubuf__ T * srcUb, __ubuf__ T * inMaxUb, + __ubuf__ T * tmpExpSumUb, __ubuf__ T * tmpMaxUb, __ubuf__ T * tmpMaxUb2, const uint32_t blockStride, + const uint32_t repeatStride, const uint16_t m, const T scale, const T minValue, uint32_t pltOriginalN, + uint32_t pltSrcN) +{ + AscendC::MicroAPI::RegTensor vreg_input_x; + AscendC::MicroAPI::RegTensor vreg_max_tmp; + AscendC::MicroAPI::RegTensor vreg_in_max; + AscendC::MicroAPI::RegTensor vreg_max_new; + AscendC::MicroAPI::RegTensor vreg_max_brc; + AscendC::MicroAPI::RegTensor vreg_cur_max; + AscendC::MicroAPI::RegTensor vreg_exp; + AscendC::MicroAPI::RegTensor vreg_exp_sum; + + // bfloat16_t + AscendC::MicroAPI::RegTensor vreg_exp_bf16; + AscendC::MicroAPI::RegTensor vreg_dst_even_bf16; + AscendC::MicroAPI::RegTensor vreg_dst_odd_bf16; + // half + AscendC::MicroAPI::RegTensor vreg_exp_fp16; + AscendC::MicroAPI::RegTensor vreg_dst_even_fp16; + AscendC::MicroAPI::RegTensor vreg_dst_odd_fp16; + + AscendC::MicroAPI::UnalignRegForStore ureg_max; + AscendC::MicroAPI::UnalignRegForStore ureg_exp_sum; + + AscendC::MicroAPI::MaskReg preg_all = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg preg_all_b16 = + AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg preg_ori_src_n = AscendC::MicroAPI::UpdateMask(pltOriginalN); + AscendC::MicroAPI::MaskReg preg_src_n = AscendC::MicroAPI::UpdateMask(pltSrcN); + AscendC::MicroAPI::MaskReg preg_src_n_b16 = + AscendC::MicroAPI::CreateMask(); + + // x_max = max(src, axis=-1, keepdims=True) + for (uint16_t i = 0; i < m; ++i) { + AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize); + AscendC::MicroAPI::Muls(vreg_input_x, vreg_input_x, scale, preg_ori_src_n); + AscendC::MicroAPI::StoreAlign( + (__ubuf__ T *&)srcUb + i * s2BaseSize, vreg_input_x, preg_src_n); + AscendC::MicroAPI::Reduce( + vreg_cur_max, vreg_input_x, preg_ori_src_n); + AscendC::MicroAPI::StoreUnAlign( + ((__ubuf__ T *&)tmpMaxUb), vreg_cur_max, ureg_max, 1); + } + AscendC::MicroAPI::StoreUnAlignPost( + ((__ubuf__ T *&)tmpMaxUb), ureg_max, 0); + AscendC::MicroAPI::LoadAlign(vreg_in_max, inMaxUb); + AscendC::MicroAPI::LocalMemBar(); + AscendC::MicroAPI::LoadAlign(vreg_cur_max, tmpMaxUb2); + AscendC::MicroAPI::Max(vreg_max_new, vreg_cur_max, vreg_in_max, preg_all); // 计算新、旧的最大值 + AscendC::MicroAPI::StoreAlign( + (__ubuf__ T *&)tmpMaxUb2, vreg_max_new, preg_all); + + AscendC::MicroAPI::LocalMemBar(); + + for (uint16_t i = 0; i < m; ++i) { + AscendC::MicroAPI::LoadAlign( + vreg_max_brc, tmpMaxUb2 + i); + AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize); + AscendC::MicroAPI::ExpSub(vreg_exp, vreg_input_x, vreg_max_brc, preg_ori_src_n); + + // x_sum = sum(x_exp, axis=-1, keepdims=True) + AscendC::MicroAPI::Reduce( + vreg_exp_sum, vreg_exp, preg_ori_src_n); + AscendC::MicroAPI::StoreUnAlign( + ((__ubuf__ T *&)tmpExpSumUb), vreg_exp_sum, ureg_exp_sum, 1); + + if constexpr (IsSameType::value) { + AscendC::MicroAPI::Cast(vreg_exp_bf16, vreg_exp, preg_all_b16); + AscendC::MicroAPI::DeInterleave(vreg_dst_even_bf16, vreg_dst_odd_bf16, + vreg_exp_bf16, vreg_exp_bf16); + AscendC::MicroAPI::StoreAlign( + ((__ubuf__ T2 *&)expUb), vreg_dst_even_bf16, blockStride, repeatStride, preg_src_n_b16); + } else if constexpr (IsSameType::value) { + AscendC::MicroAPI::Cast(vreg_exp_fp16, vreg_exp, preg_all_b16); + AscendC::MicroAPI::DeInterleave(vreg_dst_even_fp16, vreg_dst_odd_fp16, + vreg_exp_fp16, vreg_exp_fp16); + AscendC::MicroAPI::StoreAlign( + ((__ubuf__ T2 *&)expUb), vreg_dst_even_fp16, blockStride, repeatStride, preg_src_n_b16); + } + } + AscendC::MicroAPI::StoreUnAlignPost( + ((__ubuf__ T *&)tmpExpSumUb), ureg_exp_sum, 0); +} + + +// update, originN <= 64 +template +__aicore__ inline void ProcessVec1UpdateImpl64( + const LocalTensor& dstTensor, const LocalTensor& srcTensor, const LocalTensor& inMaxTensor, + const LocalTensor& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue) +{ + // 写的时候固定用65或者33的stride去写，因为正向目前使能settail之后mm2的s1方向必须算满128或者64行 + // stride, high 16bits: blockStride (m*16*2/32), low 16bits: repeatStride (1) + const uint32_t blockStride = s1BaseSize >> 1 | 0x1; + const uint32_t repeatStride = 1; + uint32_t pltOriginalN = originN; + uint32_t pltSrcN = s2BaseSize; + + __ubuf__ T2 * expUb = (__ubuf__ T2*)dstTensor.GetPhyAddr(); + __ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr(); + __ubuf__ T * inMaxUb = (__ubuf__ T*)inMaxTensor.GetPhyAddr(); + __ubuf__ T * tmpExpSumUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr(); + __ubuf__ T * tmpMaxUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64; + __ubuf__ T * tmpMaxUb2 = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64; + + ProcessVec1UpdateImpl64VF ( + expUb, srcUb, inMaxUb, tmpExpSumUb, tmpMaxUb, tmpMaxUb2, blockStride, repeatStride, m, scale, minValue, + pltOriginalN, pltSrcN); +} +} // namespace + +#endif // VF_BASIC_BLOCK_ALIGNED64_UPDATE_SFA_H diff --git a/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_utils.h b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_utils.h new file mode 100644 index 000000000..50b95219a --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_basic_block_utils.h @@ -0,0 +1,112 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_basic_block_utils.h + * \brief + */ +#ifndef VF_BASIC_BLOCK_UTILS_H +#define VF_BASIC_BLOCK_UTILS_H + +#if ASC_DEVKIT_MAJOR >= 9 +#include "kernel_basic_intf.h" +#else +#include "kernel_operator.h" +#endif + +namespace FaVectorApi { +constexpr uint32_t floatRepSize = 64; +constexpr uint32_t halfRepSize = 128; +constexpr uint32_t blockBytesU8 = 32; +constexpr float fp8e4m3MaxValue = 448.0f; +constexpr float int8MaxValue = 127.0f; +constexpr float hifp8MaxValue = 32768.0f; +constexpr float floatEps = 2.220446049250313e-16; +/* ************************************************************************************************** + * Muls + Select(optional) + SoftmaxFlashV2 + Cast(fp32->fp16/bf16) + ND2NZ + * ************************************************************************************************* */ +using namespace MicroAPI; + +constexpr static AscendC::MicroAPI::CastTrait castTraitZero = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_ROUND, +}; + +constexpr static AscendC::MicroAPI::CastTrait castTraitOne = { + AscendC::MicroAPI::RegLayout::ONE, + AscendC::MicroAPI::SatMode::SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_ROUND, +}; + +constexpr static AscendC::MicroAPI::CastTrait castTraitTwo = { + AscendC::MicroAPI::RegLayout::TWO, + AscendC::MicroAPI::SatMode::SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_ROUND, +}; + +constexpr static AscendC::MicroAPI::CastTrait castTraitThree = { + AscendC::MicroAPI::RegLayout::THREE, + AscendC::MicroAPI::SatMode::SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_ROUND, +}; + +constexpr static AscendC::MicroAPI::CastTrait castTraitRintZero = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_RINT, +}; + +constexpr static AscendC::MicroAPI::CastTrait castTraitRintOne = { + AscendC::MicroAPI::RegLayout::ONE, + AscendC::MicroAPI::SatMode::SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_RINT, +}; + +constexpr static AscendC::MicroAPI::CastTrait castTraitRintTwo = { + AscendC::MicroAPI::RegLayout::TWO, + AscendC::MicroAPI::SatMode::SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_RINT, +}; + +constexpr static AscendC::MicroAPI::CastTrait castTraitRintThree = { + AscendC::MicroAPI::RegLayout::THREE, + AscendC::MicroAPI::SatMode::SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_RINT, +}; + +#define USE_MLA_FULLQUANT_V1_P(vreg_exp, vreg_rowmax_p, MaskReg) \ + do { \ + Muls(vreg_exp, vreg_exp, fp8e4m3MaxValue, MaskReg); \ + Div(vreg_exp, vreg_exp, vreg_rowmax_p, MaskReg); \ + } while (0) + +#define USE_MLA_FULLQUANT_V1_P_INT8(vreg_exp, vreg_rowmax_p, MaskReg) \ + do { \ + Muls(vreg_exp, vreg_exp, int8MaxValue, MaskReg); \ + Div(vreg_exp, vreg_exp, vreg_rowmax_p, MaskReg); \ + } while (0) + +#define USE_MLA_FULLQUANT_V1_P_HIFP8(vreg_exp, vreg_rowmax_p, MaskReg) \ + do { \ + Muls(vreg_exp, vreg_exp, hifp8MaxValue, MaskReg); \ + Div(vreg_exp, vreg_exp, vreg_rowmax_p, MaskReg); \ + } while (0) +} // namespace + +#endif // VF_BASIC_BLOCK_UTILS_H diff --git a/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_flashupdate_new.h b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_flashupdate_new.h new file mode 100644 index 000000000..0793fb8f5 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_flashupdate_new.h @@ -0,0 +1,727 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_flashupdate_new.h + * \brief + */ +#ifndef MY_FLASH_UPDATE_NEW_INTERFACE_H +#define MY_FLASH_UPDATE_NEW_INTERFACE_H + +#include "kernel_tensor.h" + +namespace FaVectorApi { +// bf16->fp32 +static constexpr MicroAPI::CastTrait castTraitFp16_32_update = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; +constexpr uint16_t REDUCE_SIZE = 1; +template +__simd_vf__ inline void FlashUpdateBasicVF(__ubuf__ float * dstUb, __ubuf__ float * curUb, __ubuf__ float * preUb, + __ubuf__ float * expMaxUb, __ubuf__ float * rowMaxUb, const uint16_t m, const uint16_t d, + const float deScaleV, const float deScaleVPre) +{ + constexpr uint16_t floatRepSize = 64; + constexpr uint16_t dLoops = srcD / floatRepSize; + RegTensor vreg_exp_max; + RegTensor vreg_row_max; + RegTensor vreg_input_pre; + RegTensor vreg_input_cur; + RegTensor vreg_mul; + RegTensor vreg_add; + + MaskReg preg_all = CreateMask(); + + // dstTensor = preTensor * expMaxTensor + curTensor + for (uint16_t i = 0; i < m; ++i) { + LoadAlign(vreg_exp_max, expMaxUb + i * reduceSize); // [m,8] + if constexpr (isMlaFullQuant) { + LoadAlign(vreg_row_max, rowMaxUb + i * reduceSize); + } + + for (uint16_t j = 0; j < dLoops; ++j) { + LoadAlign(vreg_input_pre, preUb + i * d + j * floatRepSize); + LoadAlign(vreg_input_cur, curUb + i * d + j * floatRepSize); + if constexpr (isMlaFullQuant) { + Mul(vreg_input_cur, vreg_input_cur, vreg_row_max, preg_all); + } + Mul(vreg_mul, vreg_exp_max, vreg_input_pre, preg_all); + if constexpr (IsSameType::value || + IsSameType::value || + IsSameType::value || + IsSameType::value) { + Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all); + if constexpr (isUpdatePre) { + Muls(vreg_mul, vreg_mul, deScaleVPre, preg_all); + } + } + Add(vreg_add, vreg_mul, vreg_input_cur, preg_all); + StoreAlign( + (__ubuf__ T *&)dstUb + i * d + j * floatRepSize, vreg_add, preg_all); + } + } +} +/* ************************************************************************************************** + * FlashUpdate, fp32 + * ************************************************************************************************* */ +template +__aicore__ inline void FlashUpdateBasic(const LocalTensor& dstTensor, const LocalTensor& curTensor, + const LocalTensor& preTensor, const LocalTensor& expMaxTensor, const LocalTensor& rowMaxTensor, + const uint16_t m, const uint16_t d, const float deScaleV, const float deScaleVPre) +{ + __ubuf__ float * dstUb = (__ubuf__ T*)dstTensor.GetPhyAddr(); + __ubuf__ float * curUb = (__ubuf__ T*)curTensor.GetPhyAddr(); + __ubuf__ float * preUb = (__ubuf__ T*)preTensor.GetPhyAddr(); + __ubuf__ float * expMaxUb = (__ubuf__ T*)expMaxTensor.GetPhyAddr(); + __ubuf__ float * rowMaxUb = (__ubuf__ T*)rowMaxTensor.GetPhyAddr(); + + FlashUpdateBasicVF( + dstUb, curUb, preUb, expMaxUb, rowMaxUb, m, d, deScaleV, deScaleVPre); +} + +template +__simd_vf__ inline void FlashUpdateGeneralVF(__ubuf__ float * dstUb, __ubuf__ float * curUb, __ubuf__ float * preUb, + __ubuf__ float * expMaxUb, const uint16_t m, const uint16_t d, + const float deScaleV, const float deScaleVPre, const uint32_t pltTailD, const uint16_t hasTail) +{ + RegTensor vreg_exp_max; + RegTensor vreg_input_pre; + RegTensor vreg_input_cur; + RegTensor vreg_mul; + RegTensor vreg_add; + + MaskReg preg_all = CreateMask(); + uint32_t tmpTailD = pltTailD; + MaskReg preg_tail_d = UpdateMask(tmpTailD); + constexpr uint16_t floatRepSize = 64; + const uint16_t dLoops = d / floatRepSize; + + // dstTensor = preTensor * expMaxTensor + curTensor + for (uint16_t i = 0; i < m; ++i) { + LoadAlign(vreg_exp_max, expMaxUb + i * reduceSize); // [m,8] + + for (uint16_t j = 0; j < dLoops; ++j) { + LoadAlign(vreg_input_pre, preUb + i * d + j * floatRepSize); + LoadAlign(vreg_input_cur, curUb + i * d + j * floatRepSize); + + Mul(vreg_mul, vreg_exp_max, vreg_input_pre, preg_all); + if constexpr (IsSameType::value || + IsSameType::value || + IsSameType::value || + IsSameType::value) { + Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all); + if constexpr (isUpdatePre) { + Muls(vreg_mul, vreg_mul, deScaleVPre, preg_all); + } + } + Add(vreg_add, vreg_mul, vreg_input_cur, preg_all); + StoreAlign( + (__ubuf__ T *&)dstUb + i * d + j * floatRepSize, vreg_add, preg_all); + } + for (uint16_t t = 0; t < hasTail; ++t) { + LoadAlign(vreg_input_pre, preUb + i * d + dLoops * floatRepSize); + LoadAlign(vreg_input_cur, curUb + i * d + dLoops * floatRepSize); + + Mul(vreg_mul, vreg_exp_max, vreg_input_pre, preg_tail_d); + if constexpr (IsSameType::value || + IsSameType::value || + IsSameType::value || + IsSameType::value) { + Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all); + if constexpr (isUpdatePre) { + Muls(vreg_mul, vreg_mul, deScaleVPre, preg_all); + } + } + Add(vreg_add, vreg_mul, vreg_input_cur, preg_tail_d); + + StoreAlign( + (__ubuf__ T *&)dstUb + i * d + dLoops * floatRepSize, vreg_add, preg_tail_d); + } + } +} + +template +__aicore__ inline void FlashUpdateGeneral(const LocalTensor& dstTensor, const LocalTensor& curTensor, + const LocalTensor& preTensor, const LocalTensor& expMaxTensor, const uint16_t m, const uint16_t d, + const float deScaleV, const float deScaleVPre) +{ + __ubuf__ float * dstUb = (__ubuf__ T*)dstTensor.GetPhyAddr(); + __ubuf__ float * curUb = (__ubuf__ T*)curTensor.GetPhyAddr(); + __ubuf__ float * preUb = (__ubuf__ T*)preTensor.GetPhyAddr(); + __ubuf__ float * expMaxUb = (__ubuf__ T*)expMaxTensor.GetPhyAddr(); + + constexpr uint16_t floatRepSize = 64; + const uint16_t tailD = d % floatRepSize; + uint32_t pltTailD = static_cast(tailD); + + uint16_t hasTail = 0; + if (tailD > 0) { + hasTail = 1; + } + + FlashUpdateGeneralVF( + dstUb, curUb, preUb, expMaxUb, m, d, deScaleV, deScaleVPre, pltTailD, hasTail); +} + +/* + * @ingroup FlashUpdate + * @brief compute, dstTensor = preTensor * expMaxTensor + curTensor + * @param [out] dstTensor, output LocalTensor + * @param [in] curTensor, input LocalTensor + * @param [in] preTensor, input LocalTensor + * @param [in] expMaxTensor, input LocalTensor + * @param [in] m, input rows + * @param [in] d, input columns, should be 32 bytes aligned + */ +template +__aicore__ inline void FlashUpdateNew(const LocalTensor& dstTensor, const LocalTensor& curTensor, + const LocalTensor& preTensor, const LocalTensor& expMaxTensor, const LocalTensor& rowMaxTensor, const uint16_t m, const uint16_t d, + const float deScaleV, const float deScaleVPre) +{ + static_assert(IsSameType::value, "VF FlashUpdate, T must be float"); + + constexpr uint16_t floatRepSize = 64; + if constexpr(srcD % floatRepSize == 0) { + FlashUpdateBasic(dstTensor, curTensor, preTensor, expMaxTensor, rowMaxTensor, + m, d, deScaleV, deScaleVPre); + } else { + + FlashUpdateGeneral(dstTensor, curTensor, preTensor, expMaxTensor, m, d, + deScaleV, deScaleVPre); + } +} + +template +__simd_vf__ inline void FlashUpdateLastBasicVF(__ubuf__ float * dstUb, __ubuf__ float * curUb, __ubuf__ float * preUb, + __ubuf__ float * expMaxUb, __ubuf__ float * expSumUb, __ubuf__ float * rowMaxUb, const uint16_t m, const uint16_t d, + const float deScaleV, const float deScaleVPre) +{ + RegTensor vreg_exp_max; + RegTensor vreg_row_max; + RegTensor vreg_input_pre; + RegTensor vreg_input_cur; + RegTensor vreg_mul; + RegTensor vreg_add; + RegTensor vreg_div; + RegTensor vreg_cast; + RegTensor vreg_exp_sum; + + MaskReg preg_all = CreateMask(); + constexpr uint16_t floatRepSize = 64; + constexpr uint16_t dLoops = srcD / floatRepSize; + constexpr float fp8e4m3MaxValueRec = 1 / 448.0f; + constexpr float int8MaxValueRec = 1 / 127.0f; + constexpr float hifp8MaxValueRec = 1 / 32768.0f; + + for (uint16_t i = 0; i < m; ++i) { + LoadAlign(vreg_exp_max, expMaxUb + i * reduceSize); + LoadAlign(vreg_exp_sum, expSumUb + i * reduceSize); + if constexpr (isMlaFullQuant) { + LoadAlign(vreg_row_max, rowMaxUb + i * reduceSize); + } + for (uint16_t j = 0; j < dLoops; ++j) { + LoadAlign(vreg_input_pre, preUb + i * d + j * floatRepSize); + LoadAlign(vreg_input_cur, curUb + i * d + j * floatRepSize); + if constexpr (isMlaFullQuant) { + Mul(vreg_input_cur, vreg_input_cur, vreg_row_max, preg_all); + } + Mul(vreg_mul, vreg_exp_max, vreg_input_pre, preg_all); + if constexpr (IsSameType::value || + IsSameType::value || + IsSameType::value || + IsSameType::value) { + Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all); + if constexpr (isUpdatePre) { + Muls(vreg_mul, vreg_mul, deScaleVPre, preg_all); + } + } + Add(vreg_add, vreg_mul, vreg_input_cur, preg_all); + Div(vreg_div, vreg_add, vreg_exp_sum, preg_all); + if constexpr (isMlaFullQuant) { + if constexpr (IsSameType::value) { + Muls(vreg_div, vreg_div, fp8e4m3MaxValueRec, preg_all); + } else if constexpr (IsSameType::value) { + Muls(vreg_div, vreg_div, int8MaxValueRec, preg_all); + } else { + Muls(vreg_div, vreg_div, hifp8MaxValueRec, preg_all); + } + } + StoreAlign( + (__ubuf__ T *&)dstUb + i * d + j * floatRepSize, vreg_div, preg_all); + } + } +} + +template +__aicore__ inline void FlashUpdateLastBasic(const LocalTensor& dstTensor, + const LocalTensor& curTensor, const LocalTensor& preTensor, + const LocalTensor& expMaxTensor, const LocalTensor& rowMaxTensor, const LocalTensor& expSumTensor, + const uint16_t m, const uint16_t d, const float deScaleV, const float deScaleVPre) +{ + __ubuf__ float * dstUb = (__ubuf__ T*)dstTensor.GetPhyAddr(); + __ubuf__ float * curUb = (__ubuf__ T*)curTensor.GetPhyAddr(); + __ubuf__ float * preUb = (__ubuf__ T*)preTensor.GetPhyAddr(); + __ubuf__ float * expMaxUb = (__ubuf__ T*)expMaxTensor.GetPhyAddr(); + __ubuf__ float * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr(); + __ubuf__ float * rowMaxUb = (__ubuf__ T*)rowMaxTensor.GetPhyAddr(); + + FlashUpdateLastBasicVF( + dstUb, curUb, preUb, expMaxUb, expSumUb, rowMaxUb, m, d, deScaleV, deScaleVPre); +} + +template +__simd_vf__ inline void FlashUpdateLastGeneralVF(__ubuf__ float * dstUb, __ubuf__ float * curUb, + __ubuf__ float * preUb, __ubuf__ float * expMaxUb, __ubuf__ float * expSumUb, const uint16_t m, const uint16_t d, + const float deScaleV, const float deScaleVPre, const uint32_t pltTailD, const uint16_t hasTail) +{ + RegTensor vreg_exp_max; + RegTensor vreg_input_pre; + RegTensor vreg_input_cur; + RegTensor vreg_mul; + RegTensor vreg_add; + RegTensor vreg_div; + RegTensor vreg_cast; + RegTensor vreg_exp_sum; + + MaskReg preg_all = CreateMask(); + uint32_t tmpTailD = pltTailD; + MaskReg preg_tail_d = UpdateMask(tmpTailD); + constexpr uint16_t floatRepSize = 64; + uint16_t dLoops = d / floatRepSize; + + for (uint16_t i = 0; i < m; ++i) { + LoadAlign(vreg_exp_max, expMaxUb + i * reduceSize); + LoadAlign(vreg_exp_sum, expSumUb + i * reduceSize); + for (uint16_t j = 0; j < dLoops; ++j) { + LoadAlign(vreg_input_pre, preUb + i * d + j * floatRepSize); + LoadAlign(vreg_input_cur, curUb + i * d + j * floatRepSize); + + Mul(vreg_mul, vreg_exp_max, vreg_input_pre, preg_all); + if constexpr (IsSameType::value || + IsSameType::value || + IsSameType::value || + IsSameType::value) { + Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all); + if constexpr (isUpdatePre) { + Muls(vreg_mul, vreg_mul, deScaleVPre, preg_all); + } + } + Add(vreg_add, vreg_mul, vreg_input_cur, preg_all); + Div(vreg_div, vreg_add, vreg_exp_sum, preg_all); + + StoreAlign( + (__ubuf__ T *&)dstUb + i * d + j * floatRepSize, vreg_div, preg_all); + } + + for (uint16_t t = 0; t < hasTail; ++t) { + LoadAlign(vreg_input_pre, preUb + i * d + dLoops * floatRepSize); + LoadAlign(vreg_input_cur, curUb + i * d + dLoops * floatRepSize); + Mul(vreg_mul, vreg_exp_max, vreg_input_pre, preg_tail_d); + if constexpr (IsSameType::value || + IsSameType::value || + IsSameType::value || + IsSameType::value) { + Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all); + if constexpr (isUpdatePre) { + Muls(vreg_mul, vreg_mul, deScaleVPre, preg_all); + } + } + Add(vreg_add, vreg_mul, vreg_input_cur, preg_tail_d); + Div(vreg_div, vreg_add, vreg_exp_sum, preg_tail_d); + + StoreAlign( + (__ubuf__ T *&)dstUb + i * d + dLoops * floatRepSize, vreg_div, preg_tail_d); + } + } +} + +template +__aicore__ inline void FlashUpdateLastGeneral(const LocalTensor& dstTensor, + const LocalTensor& curTensor, const LocalTensor& preTensor, + const LocalTensor& expMaxTensor, const LocalTensor& expSumTensor, + const uint16_t m, const uint16_t d, const float deScaleV, const float deScaleVPre) +{ + __ubuf__ float * dstUb = (__ubuf__ T*)dstTensor.GetPhyAddr(); + __ubuf__ float * curUb = (__ubuf__ T*)curTensor.GetPhyAddr(); + __ubuf__ float * preUb = (__ubuf__ T*)preTensor.GetPhyAddr(); + __ubuf__ float * expMaxUb = (__ubuf__ T*)expMaxTensor.GetPhyAddr(); + __ubuf__ float * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr(); + + constexpr uint16_t floatRepSize = 64; + uint16_t tailD = d % floatRepSize; + uint32_t pltTailD = tailD; + + uint16_t hasTail = 0; + if (tailD > 0) { + hasTail = 1; + } + + FlashUpdateLastGeneralVF( + dstUb, curUb, preUb, expMaxUb, expSumUb, m, d, deScaleV, deScaleVPre, pltTailD, hasTail); +} + +/* + * @ingroup FlashUpdateLast + * @brief compute, dstTensor = (preTensor * expMaxTensor + curTensor) / expSumTensor + * @param [out] dstTensor, output LocalTensor + * @param [in] curTensor, input LocalTensor + * @param [in] preTensor, input LocalTensor + * @param [in] expMaxTensor, input LocalTensor + * @param [in] expSumTensor, input LocalTensor + * @param [in] m, input rows + * @param [in] d, input columns, 32 bytes align + */ +template +__aicore__ inline void FlashUpdateLastNew(const LocalTensor& dstTensor, + const LocalTensor& curTensor, const LocalTensor& preTensor, + const LocalTensor& expMaxTensor, const LocalTensor& rowMaxTensor, const LocalTensor& expSumTensor, + uint16_t m, uint16_t d, const float deScaleV, const float deScaleVPre) +{ + static_assert(IsSameType::value, "VF FlashUpdateLast, T must be float"); + + constexpr uint16_t floatRepSize = 64; + if constexpr(srcD % floatRepSize == 0) { + FlashUpdateLastBasic( + dstTensor, curTensor, preTensor, expMaxTensor, rowMaxTensor, expSumTensor, m, d, deScaleV, deScaleVPre); + } else { + FlashUpdateLastGeneral( + dstTensor, curTensor, preTensor, expMaxTensor, expSumTensor, m, d, deScaleV, deScaleVPre); + } +} + +template +__simd_vf__ inline void LastDivNewVF(__ubuf__ float * dstUb, __ubuf__ float * curUb, __ubuf__ float * expSumUb, + const uint16_t m, const uint16_t d, const float deScaleV) +{ + RegTensor vreg_input_cur; + RegTensor vreg_div; + RegTensor vreg_exp_sum; + MaskReg preg_all = CreateMask(); + constexpr uint16_t floatRepSize = 64; + const uint16_t dLoops = d >> 6; + constexpr float fp8e4m3MaxValueRec = 1 / 448.0f; + constexpr float int8MaxValueRec = 1 / 127.0f; + constexpr float hifp8MaxValueRec = 1 / 32768.0f; + + for (uint16_t i = 0; i < m; ++i) { + uint32_t sreg_init = d; + LoadAlign(vreg_exp_sum, expSumUb + i * REDUCE_SIZE); + for (uint16_t j = 0; j < dLoops; ++j) { + MaskReg preg_update = UpdateMask(sreg_init); + + LoadAlign(vreg_input_cur, curUb + i * d + j * floatRepSize); + if constexpr (IsSameType::value || + IsSameType::value || + IsSameType::value || + IsSameType::value) { + Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all); + } + Div(vreg_div, vreg_input_cur, vreg_exp_sum, preg_update); + if constexpr (isMlaFullQuant) { + if constexpr (IsSameType::value) { + Muls(vreg_div, vreg_div, fp8e4m3MaxValueRec, preg_all); + } else if constexpr (IsSameType::value) { + Muls(vreg_div, vreg_div, int8MaxValueRec, preg_all); + } else { + Muls(vreg_div, vreg_div, hifp8MaxValueRec, preg_all); + } + } + StoreAlign( + (__ubuf__ T *&)dstUb + i * d + j * floatRepSize, vreg_div, preg_update); + } + } +} + +// dstTensor = curTensor / expSumTensor, curTensor: [64,128], expSumTensor: [64,8] +template +__aicore__ inline void LastDivNew(const LocalTensor& dstTensor, const LocalTensor& curTensor, + const LocalTensor& expSumTensor, const uint16_t m, const uint16_t d, const float deScaleV) +{ + __ubuf__ float * dstUb = (__ubuf__ T*)dstTensor.GetPhyAddr(); + __ubuf__ float * curUb = (__ubuf__ T*)curTensor.GetPhyAddr(); + __ubuf__ float * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr(); + + LastDivNewVF(dstUb, curUb, expSumUb, m, d, deScaleV); +} + +template +__simd_vf__ inline void InvalidLineUpdateVF(__ubuf__ T * dstUb, __ubuf__ T * srcUb, __ubuf__ T * maxUb, + const uint16_t m, const uint16_t d, const T minValue, const T invalidValue) +{ + RegTensor vreg_invalid_value; + RegTensor vreg_max; + RegTensor vreg_input; + RegTensor vreg_input_brc; + + MaskReg preg_all = CreateMask(); + MaskReg preg_compare; + const uint16_t dLoops = d >> 6; + + Duplicate(vreg_invalid_value, invalidValue); + for (uint16_t i = 0; i < m; ++i) { + LoadAlign(vreg_max, maxUb + i); + Compares(preg_compare, vreg_max, minValue, preg_all); + for (uint16_t j = 0; j < dLoops; ++j) { + LoadAlign(vreg_input, srcUb + i * d + j * floatRepSize); + Select(vreg_input_brc, vreg_invalid_value, vreg_input, preg_compare); + StoreAlign( + (__ubuf__ T *&)dstUb + i * d + j * floatRepSize, vreg_input_brc, preg_all); + } + } +} + +template +__aicore__ inline void InvalidLineUpdate(const LocalTensor& dstTensor, const LocalTensor& srcTensor, + const LocalTensor& maxTensor, const uint16_t m, const uint16_t d, const T minValue, const T invalidValue) +{ + __ubuf__ T * dstUb = (__ubuf__ T*)dstTensor.GetPhyAddr(); + __ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr(); + __ubuf__ T * maxUb = (__ubuf__ T*)maxTensor.GetPhyAddr(); + + constexpr uint16_t floatRepSize = 64; + uint16_t dLoops = d >> 6; + + InvalidLineUpdateVF(dstUb, srcUb, maxUb, m, d, minValue, invalidValue); +} + +template +__simd_vf__ inline void ComputeLseOutputVF(__ubuf__ T *srcSumUb, __ubuf__ T *srcMaxUb, __ubuf__ T *dstUb, const uint32_t dealCount) +{ + MicroAPI::RegTensor vregSum; + MicroAPI::RegTensor vregMax; + MicroAPI::RegTensor vregRes; + MicroAPI::RegTensor vregResFinal; + MicroAPI::RegTensor vregMinValue; + MicroAPI::RegTensor vregInfValue; + MicroAPI::MaskReg pregCompare; + constexpr uint32_t dealRows = 8; + constexpr uint32_t floatRepSize = 64; // 64: 一个寄存器存64个float + constexpr float infValue = 3e+99; // 3e+99 for float inf + constexpr uint32_t tmpMin = 0xFF7FFFFF; + float minValue = *((float*)&tmpMin); + uint16_t updateLoops = dealCount / dealRows; + uint16_t tailLSize = dealCount % dealRows * 8; + uint32_t pltTail = static_cast(tailLSize); + + MicroAPI::MaskReg pregAll = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregTail = MicroAPI::UpdateMask(pltTail); + MicroAPI::Duplicate(vregMinValue, minValue); + MicroAPI::Duplicate(vregInfValue, infValue); + + for (uint16_t i = 0; i < updateLoops; ++i) { + MicroAPI::LoadAlign(vregSum, srcSumUb + (i * dealRows)); + MicroAPI::LoadAlign(vregMax, srcMaxUb + (i * dealRows)); + + MicroAPI::Log(vregRes, vregSum, pregAll); + MicroAPI::Add(vregRes, vregRes, vregMax, pregAll); + + MicroAPI::Compare(pregCompare, vregMax, vregMinValue, pregAll); + MicroAPI::Select(vregResFinal, vregInfValue, vregRes, pregCompare); + + MicroAPI::StoreAlign(dstUb + (i * floatRepSize), vregResFinal, pregAll); + } + + if (tailLSize != 0) { + MicroAPI::LoadAlign(vregSum, srcSumUb + dealRows * updateLoops); + MicroAPI::LoadAlign(vregMax, srcMaxUb + dealRows * updateLoops); + + MicroAPI::Log(vregRes, vregSum, pregTail); + MicroAPI::Add(vregRes, vregRes, vregMax, pregTail); + + MicroAPI::Compare(pregCompare, vregMax, vregMinValue, pregTail); + MicroAPI::Select(vregResFinal, vregInfValue, vregRes, pregCompare); + + MicroAPI::StoreAlign(dstUb + floatRepSize * updateLoops, vregResFinal, pregTail); + } +} + +template +__aicore__ inline void ComputeLseOutputVF(const LocalTensor& dstTensor, const LocalTensor& softmaxSumTensor, + const LocalTensor& softmaxMaxTensor, uint32_t dealCount) +{ + __ubuf__ T * srcSumUb = (__ubuf__ T *)softmaxSumTensor.GetPhyAddr(); + __ubuf__ T * srcMaxUb = (__ubuf__ T *)softmaxMaxTensor.GetPhyAddr(); + __ubuf__ T * dstUb = (__ubuf__ T *)dstTensor.GetPhyAddr(); + + ComputeLseOutputVF(srcSumUb, srcMaxUb, dstUb, dealCount); +} + +template +__simd_vf__ inline void SinkSubExpAddVF(__ubuf__ T *srcSumUb, __ubuf__ T *srcMaxUb, const T sinkValue, const uint32_t dealCount) +{ + MicroAPI::RegTensor vregSum; + MicroAPI::RegTensor vregMax; + MicroAPI::RegTensor vregRes; + MicroAPI::RegTensor vregSink; + + constexpr uint32_t floatRepSize = 64; + + uint16_t updateLoops = dealCount / floatRepSize; + uint16_t tailSize = dealCount % floatRepSize; + uint32_t pltTail = static_cast(tailSize); + + //mask + MicroAPI::MaskReg pregAll = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregTail = MicroAPI::UpdateMask(pltTail); + + Duplicate(vregSink, sinkValue); + + for (uint16_t i = 0; i < updateLoops; ++i) { + MicroAPI::LoadAlign(vregSum, srcSumUb + (i * floatRepSize)); + MicroAPI::LoadAlign(vregMax, srcMaxUb + (i * floatRepSize)); + + MicroAPI::Sub(vregRes, vregSink, vregMax, pregAll); + MicroAPI::Exp(vregRes, vregRes, pregAll); + MicroAPI::Add(vregSum, vregSum, vregRes, pregAll); + + MicroAPI::StoreAlign(srcSumUb + (i * floatRepSize), vregSum, pregAll); + } + + for (uint16_t i = 0; i < tailSize; i = i + tailSize) { + MicroAPI::LoadAlign(vregSum, srcSumUb + (updateLoops * floatRepSize)); + MicroAPI::LoadAlign(vregMax, srcMaxUb + (updateLoops * floatRepSize)); + + MicroAPI::Sub(vregRes, vregSink, vregMax, pregTail); + MicroAPI::Exp(vregRes, vregRes, pregTail); + MicroAPI::Add(vregSum, vregSum, vregRes, pregTail); + + MicroAPI::StoreAlign(srcSumUb + (updateLoops * floatRepSize), vregSum, pregTail); + } +} + +template +__aicore__ inline void SinkSubExpAddVF(const LocalTensor& softmaxSumTensor, const LocalTensor& softmaxMaxTensor, + const T sinkValue, uint32_t dealCount) +{ + __ubuf__ T * srcSumUb = (__ubuf__ T *)softmaxSumTensor.GetPhyAddr(); + __ubuf__ T * srcMaxUb = (__ubuf__ T *)softmaxMaxTensor.GetPhyAddr(); + + SinkSubExpAddVF(srcSumUb, srcMaxUb, sinkValue, dealCount); +} + +template +__simd_vf__ inline void SinkSubExpAddGSFusedVF(__ubuf__ T *srcSumUb, __ubuf__ T *srcMaxUb, __ubuf__ uint16_t *sinkUb, const uint32_t dealCount) +{ + MicroAPI::RegTensor vregSum; + MicroAPI::RegTensor vregMax; + MicroAPI::RegTensor vregRes; + MicroAPI::RegTensor vregSink; + MicroAPI::RegTensor vregSinkCast; + + constexpr uint32_t floatRepSize = 64; + + uint16_t updateLoops = dealCount / floatRepSize; + uint16_t tailSize = dealCount % floatRepSize; + uint32_t pltTail = static_cast(tailSize); + + //mask + MicroAPI::MaskReg pregAll = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregTail = MicroAPI::UpdateMask(pltTail); + MicroAPI::MaskReg pregSinkAll = MicroAPI::CreateMask(); + + MicroAPI::LoadAlign((MicroAPI::RegTensor&)vregSink, sinkUb); + MicroAPI::Cast(vregSinkCast, vregSink, pregSinkAll); + + for (uint16_t i = 0; i < updateLoops; ++i) { + MicroAPI::LoadAlign(vregSum, srcSumUb + (i * floatRepSize)); + MicroAPI::LoadAlign(vregMax, srcMaxUb + (i * floatRepSize)); + + MicroAPI::Sub(vregRes, vregSinkCast, vregMax, pregAll); + MicroAPI::Exp(vregRes, vregRes, pregAll); + MicroAPI::Add(vregSum, vregSum, vregRes, pregAll); + + MicroAPI::StoreAlign(srcSumUb + (i * floatRepSize), vregSum, pregAll); + } + + if (tailSize != 0) { + MicroAPI::LoadAlign(vregSum, srcSumUb + (updateLoops * floatRepSize)); + MicroAPI::LoadAlign(vregMax, srcMaxUb + (updateLoops * floatRepSize)); + + MicroAPI::Sub(vregRes, vregSinkCast, vregMax, pregTail); + MicroAPI::Exp(vregRes, vregRes, pregTail); + MicroAPI::Add(vregSum, vregSum, vregRes, pregTail); + + MicroAPI::StoreAlign(srcSumUb + (updateLoops * floatRepSize), vregSum, pregTail); + } +} + +template +__aicore__ inline void SinkSubExpAddGSFusedVF(const LocalTensor& dstTensor, const LocalTensor& softmaxSumTensor, + const LocalTensor& softmaxMaxTensor, uint32_t dealCount) +{ + __ubuf__ T * srcSumUb = (__ubuf__ T *)softmaxSumTensor.GetPhyAddr(); + __ubuf__ T * srcMaxUb = (__ubuf__ T *)softmaxMaxTensor.GetPhyAddr(); + __ubuf__ uint16_t * dstUb = (__ubuf__ uint16_t *)dstTensor.GetPhyAddr(); + + SinkSubExpAddGSFusedVF(srcSumUb, srcMaxUb, dstUb, dealCount); +} + +template +__simd_vf__ inline void RowInvalidUpdateVF(__ubuf__ T *finalUb, __ubuf__ float *maxUb, const uint16_t m, + const uint16_t d, int64_t dSize, const uint32_t pltTailD, const uint16_t hasTail) +{ + constexpr uint16_t floatRepSize = 64; // 64: 一个寄存器可以存储64个float类型数据 + const uint16_t dLoops = d / floatRepSize; + + + constexpr uint32_t tmpZero = 0x00000000; // zero value of fp16 and fp32 + const T zeroValue = *((T*)&tmpZero); + constexpr uint32_t tmpMin = 0xFF7FFFFF; // min value of float + const float minValue = *((float*)&tmpMin); + MicroAPI::RegTensor vregMinValue; + MicroAPI::RegTensor vregZeroValue; + MicroAPI::RegTensor vregMax; + MicroAPI::RegTensor vregFinal; + MicroAPI::RegTensor vregFinalNew; + + MicroAPI::MaskReg pregAll = MicroAPI::CreateMask(); + uint32_t tmpTailD = pltTailD; + MicroAPI::MaskReg pregTailD = MicroAPI::UpdateMask(tmpTailD); + MicroAPI::MaskReg pregCompare; + + MicroAPI::Duplicate(vregMinValue, minValue); + MicroAPI::Duplicate(vregZeroValue, zeroValue); + for (uint16_t i = 0; i < m; ++i) { + MicroAPI::LoadAlign(vregMax, maxUb + i); + MicroAPI::Compare(pregCompare, vregMax, vregMinValue, pregAll); + for (uint16_t j = 0; j < dLoops; ++j) { + MicroAPI::LoadAlign(vregFinal, finalUb + i * dSize + j * floatRepSize); + MicroAPI::Select(vregFinalNew, vregZeroValue, vregFinal, pregCompare); + MicroAPI::StoreAlign(finalUb + i * dSize + j * floatRepSize, + vregFinalNew, pregAll); + } + for (uint16_t t = 0; t < hasTail; ++t) { + MicroAPI::LoadAlign(vregFinal, finalUb + i * dSize + dLoops * floatRepSize); + MicroAPI::Select(vregFinalNew, vregZeroValue, vregFinal, pregCompare); + MicroAPI::StoreAlign(finalUb + i * dSize + dLoops * floatRepSize, + vregFinalNew, pregTailD); + } + } +} + +template +__aicore__ inline void RowInvalidUpdateVF(const LocalTensor& finalTensor, const LocalTensor& maxTensor, + const uint16_t m, const uint16_t d, int64_t dSize) +{ + __ubuf__ T * finalUb = (__ubuf__ T*)finalTensor.GetPhyAddr(); + __ubuf__ float * maxUb = (__ubuf__ float*)maxTensor.GetPhyAddr(); + + constexpr uint16_t floatRepSize = 64; + const uint16_t tailD = d % floatRepSize; + uint32_t pltTailD = static_cast(tailD); + uint16_t hasTail = 0; + if (tailD > 0) { + hasTail = 1; + } + + RowInvalidUpdateVF(finalUb, maxUb, m, d, dSize, pltTailD, hasTail); +} +} // namespace + +#endif // MY_FLASH_UPDATE_INTERFACE_H diff --git a/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_mul_sel_softmaxflashv2_cast_nz_sfa.h b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_mul_sel_softmaxflashv2_cast_nz_sfa.h new file mode 100644 index 000000000..4c9997c06 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/arch35/vf/vf_mul_sel_softmaxflashv2_cast_nz_sfa.h @@ -0,0 +1,164 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_mul_sel_softmaxflashv2_cast_nz_sfa.h + * \brief + */ +#ifndef MUL_SEL_SOFTMAX_FLASH_V2_CAST_NZ_SFA_INTERFACE_H +#define MUL_SEL_SOFTMAX_FLASH_V2_CAST_NZ_SFA_INTERFACE_H + +#include "vf_basic_block_aligned128_no_update_sfa.h" +#include "vf_basic_block_aligned128_update_sfa.h" +#include "vf_basic_block_unaligned64_update_sfa.h" +#include "vf_basic_block_unaligned64_no_update_sfa.h" +#include "vf_basic_block_unaligned128_no_update_sfa.h" +#include "vf_basic_block_unaligned128_update_sfa.h" + +using namespace regbaseutil; + +namespace FaVectorApi { +/* ************************************************************************************************** + * Muls + Select(optional) + SoftmaxFlashV2 + Cast(fp32->fp16/bf16) + ND2NZ + * ************************************************************************************************* */ +using AscendC::LocalTensor; + +enum class OriginNRange { + EQ_128_SFA = 0, // originN == 128, better performance than GT_64_AND_LTE_128 (s2BaseSize=128) + GT_0_AND_LTE_64_SFA, // 0 < originN <= 64 (s2BaseSize <= 64 or tail s2) + GT_64_AND_LTE_128_SFA, // 64 < originN <= 128, support for non-alignment (s2BaseSize=128) + N_INVALID_SFA +}; +template +__aicore__ inline void ProcessVec1NoUpdate( + const LocalTensor& dstTensor, const LocalTensor& srcTensor, + const LocalTensor& expSumTensor, const LocalTensor& maxTensor, const LocalTensor& inMaxTensor, + const LocalTensor& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue) +{ + if constexpr (oriNRange == OriginNRange::EQ_128_SFA) { + ProcessVec1NoUpdateImpl128( + dstTensor, srcTensor, expSumTensor, maxTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue); + } else if constexpr (oriNRange == OriginNRange::GT_0_AND_LTE_64_SFA) { + ProcessVec1NoUpdateImpl64( + dstTensor, srcTensor, expSumTensor, maxTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue); + } else if constexpr (oriNRange == OriginNRange::GT_64_AND_LTE_128_SFA) { + ProcessVec1NoUpdateGeneralImpl128( + dstTensor, srcTensor, expSumTensor, maxTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue); + } +} + +template +__aicore__ inline void ProcessVec1Update( + const LocalTensor& dstTensor, const LocalTensor& srcTensor, + const LocalTensor& expSumTensor, const LocalTensor& maxTensor, const LocalTensor& inMaxTensor, + const LocalTensor& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue) +{ + if constexpr (oriNRange == OriginNRange::EQ_128_SFA) { + ProcessVec1UpdateImpl128( + dstTensor, srcTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue); + } else if constexpr (oriNRange == OriginNRange::GT_0_AND_LTE_64_SFA) { + ProcessVec1UpdateImpl64( + dstTensor, srcTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue); + } else if constexpr (oriNRange == OriginNRange::GT_64_AND_LTE_128_SFA) { + ProcessVec1UpdateGeneralImpl128( + dstTensor, srcTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue); + } +} + +template +__aicore__ inline void ProcessVec1Vf( + const LocalTensor& dstTensor, const LocalTensor& srcTensor, + const LocalTensor& expSumTensor, const LocalTensor& maxTensor, const LocalTensor& inMaxTensor, + const LocalTensor& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue) +{ + static_assert(IsSameType::value, "VF mul_sel_softmaxflashv2_cast_nz, T must be float"); + static_assert((IsSameType::value || IsSameType::value), + "VF mul_sel_softmaxflashv2_cast_nz, T2 must be half or bfloat16"); + + if constexpr (!isUpdate) { + ProcessVec1NoUpdate( + dstTensor, srcTensor, expSumTensor, maxTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue); + } else { + ProcessVec1Update( + dstTensor, srcTensor, expSumTensor, maxTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue); + } +} + +template +__simd_vf__ inline void UpdateExpSumAndExpMaxVF(__ubuf__ T * maxUb, __ubuf__ T * inMaxUb, __ubuf__ T * expMaxUb, + __ubuf__ T * expSumUb, __ubuf__ T * inExpSumUb, __ubuf__ T * tmpExpSumUb, __ubuf__ T * tmpMaxUb, const uint32_t m) +{ + RegTensor vreg_input_x; + RegTensor vreg_input_x_unroll; + RegTensor vreg_max; + RegTensor vreg_in_max; + RegTensor vreg_exp_sum; + RegTensor vreg_in_exp_sum; + RegTensor vreg_exp_max; + RegTensor vreg_exp_sum_brc; + RegTensor vreg_exp_sum_update; + MaskReg preg_all = CreateMask(); + // 注意：当m大于64的时候需要开启循环 + LoadAlign(vreg_max, tmpMaxUb); + LoadAlign(vreg_in_max, inMaxUb); + FusedExpSub(vreg_exp_max, vreg_in_max, vreg_max, preg_all); + StoreAlign( + (__ubuf__ T *&)expMaxUb, vreg_exp_max, preg_all); + StoreAlign( + (__ubuf__ T *&)maxUb, vreg_max, preg_all); + LoadAlign(vreg_in_exp_sum, inExpSumUb); + + // x_sum = exp_max * insum + x_sum + LoadAlign(vreg_exp_sum_brc, tmpExpSumUb); + Mul(vreg_exp_sum_update, vreg_exp_max, vreg_in_exp_sum, preg_all); + Add(vreg_exp_sum_update, vreg_exp_sum_update, vreg_exp_sum_brc, preg_all); + StoreAlign( + (__ubuf__ T *&)expSumUb, vreg_exp_sum_update, preg_all); +} + +template +__aicore__ inline void SFAUpdateExpSumAndExpMax( + const LocalTensor& expSumTensor, const LocalTensor& maxTensor, + const LocalTensor& expMaxTensor, const LocalTensor& inExpSumTensor, + const LocalTensor& inMaxTensor, const LocalTensor& sharedTmpBuffer, const uint32_t m) +{ + __ubuf__ T * maxUb = (__ubuf__ T*)maxTensor.GetPhyAddr(); + __ubuf__ T * inMaxUb = (__ubuf__ T*)inMaxTensor.GetPhyAddr(); + + __ubuf__ T * expMaxUb = (__ubuf__ T*)expMaxTensor.GetPhyAddr(); + __ubuf__ T * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr(); + __ubuf__ T * inExpSumUb = (__ubuf__ T*)inExpSumTensor.GetPhyAddr(); + + __ubuf__ T * tmpExpSumUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr(); + __ubuf__ T * tmpMaxUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64; + + UpdateExpSumAndExpMaxVF(maxUb, inMaxUb, expMaxUb, expSumUb, inExpSumUb, tmpExpSumUb, tmpMaxUb, m); +} + +template +__simd_vf__ inline void DuplicateSumWithR0VF(__ubuf__ T * sumUb, const T R0, uint32_t m) { + AscendC::MicroAPI::RegTensor vreg_sum; + AscendC::MicroAPI::MaskReg preg_m = AscendC::MicroAPI::UpdateMask(m); + AscendC::MicroAPI::UnalignRegForStore ureg; + AscendC::MicroAPI::Duplicate(vreg_sum, R0, preg_m); + AscendC::MicroAPI::StoreAlign(sumUb, vreg_sum, preg_m); +} + +template +__aicore__ inline void DuplicateSumWithR0(const LocalTensor& sumTensor, const T R0, uint32_t m) +{ + __ubuf__ T * sumUb = (__ubuf__ T*)sumTensor.GetPhyAddr(); + DuplicateSumWithR0VF(sumUb, R0, m); +} +} // namespace +#endif // MUL_SEL_SOFTMAX_FLASH_V2_CAST_NZ_SFA_INTERFACE_H diff --git a/csrc/ascend/attention/common/op_kernel/buffer.h b/csrc/ascend/attention/common/op_kernel/buffer.h new file mode 100644 index 000000000..cccbd59cc --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/buffer.h @@ -0,0 +1,292 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file buffer.h + * \brief同步管理 + */ +#ifndef BUFFER_H +#define BUFFER_H +#include +#include"lib/matmul_intf.h" +#if ASC_DEVKIT_MAJOR >= 9 +#include "kernel_basic_intf.h" +#else +#include "kernel_operator.h" +#endif +using namespace AscendC; +namespace fa_base_matmul { +__BLOCK_LOCAL__ __inline__ uint32_t idCounterNum; +#define MAKE_ID ((++idCounterNum) % 11) + +// 核间同步中，AIC(flagId 0-10)对应AIV0(flagId 0-10)，对应AIV1(flagId 16-26) +#define AIV0_AIV1_OFFSET 16 + +enum class BufferType { + L1 = 0, + L0A = 1, + L0B = 2, + L0C = 3, + UB = 4, + GM = 5, + C2 = 6, +}; + +enum class SyncType { + NO_SYNC, + INNER_CORE_SYNC, + CROSS_CORE_SYNC_FORWARD, + CROSS_CORE_SYNC_BOTH, + CROSS_CORE_SYNC_BACKWARD, +}; + +constexpr uint32_t INVALID_CROSS_CORE_EVENT_ID = 16; +static constexpr uint64_t CROSS_CORE_SYNC_MODE = 4; + +template +struct BufferInfo{ + // Cons 消费者，Prod 生产者 + __aicore__ const static constexpr HardEvent ConsWaitProdStatus() { + if constexpr (Type == BufferType::L1) { + return HardEvent::MTE2_MTE1; + } else if constexpr (Type == BufferType::L0A) { + return HardEvent::MTE1_M; + } else if constexpr (Type == BufferType::L0B) { + return HardEvent::MTE1_M; + } else if constexpr (Type == BufferType::L0C) { + return HardEvent::M_FIX; + } else if constexpr (Type == BufferType::C2) { + return HardEvent::MTE1_M; + } + } + + __aicore__ const static constexpr HardEvent ProdWaitConsStatus() { + if constexpr (Type == BufferType::L1) { + return HardEvent::MTE1_MTE2; + } else if constexpr (Type == BufferType::L0A) { + return HardEvent::M_MTE1; + } else if constexpr (Type == BufferType::L0B) { + return HardEvent::M_MTE1; + } else if constexpr (Type == BufferType::L0C) { + return HardEvent::FIX_M; + } else if constexpr (Type == BufferType::C2) { + return HardEvent::M_MTE1; + } + } + + __aicore__ const static constexpr TPosition GetTPosition() { + if constexpr (Type == BufferType::L1) { + return TPosition::A1; + } else if constexpr (Type == BufferType::L0A) { + return TPosition::A2; + } else if constexpr (Type == BufferType::L0B) { + return TPosition::B2; + } else if constexpr (Type == BufferType::L0C) { + return TPosition::CO1; + } else if constexpr (Type == BufferType::UB) { + return TPosition::VECIN; + } else if constexpr (Type == BufferType::GM) { + return TPosition::GM; + } else if constexpr (Type == BufferType::C2) { + return TPosition::C2; + } + } + + static constexpr HardEvent EventP2C = ConsWaitProdStatus(); // 生产者到消费者方向的HardEvent：消费者等生产者提供/生产者通知消费者已生成 + static constexpr HardEvent EventC2P = ProdWaitConsStatus(); // 消费者到生产者方向的HardEvent：生产者等消费者消耗/消费者通知生产者已消耗’ + static constexpr TPosition Position = GetTPosition(); +}; + +// buffer绑定生产者、消费者关系 +// L1 buffer的生产者为MTE2或者MTE3，消费者为MTE1 +// L0A buffer的生产者为MTE1，消费者为M +// L0B buffer的生产者为MTE1，消费者为M +// L0C buffer的生产者为M，消费者为FIX +template +class Buffer { + using TensorType = std::conditional_t, LocalTensor>; + + template + using TargetTensorType = std::conditional_t, LocalTensor>; +public: + __aicore__ inline Buffer() {} + __aicore__ inline Buffer(TensorType tensor, uint32_t size) { + tensor_ = tensor; + size_ = size; + if constexpr (syncType == SyncType::CROSS_CORE_SYNC_FORWARD) { + id0_ = MAKE_ID; + id1_ = INVALID_CROSS_CORE_EVENT_ID; + } else if constexpr (syncType == SyncType::CROSS_CORE_SYNC_BACKWARD) { + id0_ = INVALID_CROSS_CORE_EVENT_ID; + id1_ = MAKE_ID; + } else if constexpr (syncType == SyncType::CROSS_CORE_SYNC_BOTH) { + id0_ = MAKE_ID; + id1_ = MAKE_ID; + } else { + id0_ = INVALID_CROSS_CORE_EVENT_ID; + id1_ = INVALID_CROSS_CORE_EVENT_ID; + } + } + + __aicore__ inline void Init() { + if ASCEND_IS_AIC { + if constexpr (syncType == SyncType::INNER_CORE_SYNC) { + p2cEventId_ = GetTPipePtr()->AllocEventID::EventP2C>(); // 确保只能被调用一次 + c2pEventId_ = GetTPipePtr()->AllocEventID::EventC2P>(); + SetFlag::EventC2P>(c2pEventId_); + } + } + } + + __aicore__ inline void UnInit() { + if ASCEND_IS_AIC { + if constexpr (syncType == SyncType::INNER_CORE_SYNC) { + WaitFlag::EventC2P>(c2pEventId_); + GetTPipePtr()->ReleaseEventID::EventP2C>(p2cEventId_); // 确保只能被调用一次 + GetTPipePtr()->ReleaseEventID::EventC2P>(c2pEventId_); + } + } + } + + template + __aicore__ inline void Wait() { + if ASCEND_IS_AIC { + if constexpr (syncType == SyncType::INNER_CORE_SYNC) { + if constexpr (EventType == BufferInfo::EventP2C) { + WaitFlag::EventP2C>(p2cEventId_); // 消费者等待生产者完成生产 + } else { + WaitFlag::EventC2P>(c2pEventId_); // 生产者等待消费者完成消费 + } + } + } + } + + template + __aicore__ inline void Set() { + if ASCEND_IS_AIC { + if constexpr (syncType == SyncType::INNER_CORE_SYNC) { + if constexpr (EventType == BufferInfo::EventP2C) { + SetFlag::EventP2C>(p2cEventId_); // 生产者通知消费者已完成生产 + } else { + SetFlag::EventC2P>(c2pEventId_); // 消费者通知生产者已完成消费 + } + } + } + } + + __aicore__ inline void SetEventID() { + if ASCEND_IS_AIC { + p2cEventId_ = GetTPipePtr()->AllocEventID::EventP2C>(); // 确保只能被调用一次 + c2pEventId_ = GetTPipePtr()->AllocEventID::EventC2P>(); + } + } + + template + __aicore__ inline TEventID GetEventID() { + if ASCEND_IS_AIC { + if constexpr (EventType == BufferInfo::EventP2C) { + return p2cEventId_; // 生产者通知消费者已完成生产 + } else { + return c2pEventId_; // 消费者通知生产者已完成消费 + } + } + } + + template + __aicore__ inline void WaitCrossCore() { + if constexpr (bufferType == BufferType::GM && syncType == SyncType::CROSS_CORE_SYNC_BACKWARD) { + // AIC属于消费者，AIV属于生产者，且一个AIC对应两个AIV + if ASCEND_IS_AIC { + CrossCoreWaitFlag(id1_); + CrossCoreWaitFlag(id1_ + AIV0_AIV1_OFFSET); + } else { + CrossCoreWaitFlag(id0_); + } + } else if constexpr (bufferType == BufferType::UB || bufferType == BufferType::GM) { + // AIC属于生产者，AIV属于消费者，且一个AIC对应两个AIV + if ASCEND_IS_AIC { + CrossCoreWaitFlag(id1_); + CrossCoreWaitFlag(id1_ + AIV0_AIV1_OFFSET); + } else { + if constexpr (isReuse) { + CrossCoreWaitFlag(id0_); + } else { + CrossCoreWaitFlag(id0_); + } + } + } else if constexpr (bufferType == BufferType::L1) { + // AIC属于消费者，AIV属于生产者，且一个AIC对应两个AIV + if ASCEND_IS_AIC { + CrossCoreWaitFlag(id0_); + CrossCoreWaitFlag(id0_ + AIV0_AIV1_OFFSET); + } else { + if constexpr (syncType == SyncType::CROSS_CORE_SYNC_BOTH) { + CrossCoreWaitFlag(id1_); + } + } + } + } + + template + __aicore__ inline void SetCrossCore() { + if constexpr (bufferType == BufferType::GM && syncType == SyncType::CROSS_CORE_SYNC_BACKWARD) { + // AIC属于消费者，AIV属于生产者，且一个AIC对应两个AIV + if ASCEND_IS_AIC { + CrossCoreSetFlag(id0_); + CrossCoreSetFlag(id0_ + AIV0_AIV1_OFFSET); + } else { + CrossCoreSetFlag(id1_); + } + } else if constexpr (bufferType == BufferType::UB || bufferType == BufferType::GM) { + // AIC属于生产者，AIV属于消费者，且一个AIC对应两个AIV + if ASCEND_IS_AIC { + CrossCoreSetFlag(id0_); + CrossCoreSetFlag(id0_ + AIV0_AIV1_OFFSET); + } else { + if constexpr (isReuse) { + CrossCoreSetFlag(id1_); + } else { + CrossCoreSetFlag(id1_); + } + } + } else if constexpr (bufferType == BufferType::L1) { + // AIC属于消费者，AIV属于生产者，且一个AIC对应两个AIV + if ASCEND_IS_AIC { + if constexpr (syncType == SyncType::CROSS_CORE_SYNC_BOTH) { + CrossCoreSetFlag(id1_); + CrossCoreSetFlag(id1_ + AIV0_AIV1_OFFSET); + } + } else { + CrossCoreSetFlag(id0_); + } + } + } + + template + __aicore__ inline TargetTensorType GetTensor() { + return tensor_.template ReinterpretCast(); + } + + template + __aicore__ inline TargetTensorType GetTensor(uint64_t startindex) { + TargetTensorType tmpTensor = tensor_.template ReinterpretCast(); + return tmpTensor[startindex]; + } + +private: + TensorType tensor_; + uint32_t size_; + TEventID p2cEventId_; + TEventID c2pEventId_; + uint32_t id0_; // 用作正向同步：生产者通知消费者，或者消费者等待生产者； + uint32_t id1_; // 用作反向同步：消费者通知生产者，或者生产者等待消费者； +}; +} +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/common/op_kernel/buffer_manager.h b/csrc/ascend/attention/common/op_kernel/buffer_manager.h new file mode 100644 index 000000000..24d469a21 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/buffer_manager.h @@ -0,0 +1,61 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file buffer_manager.h + * \brief buffer内存管理 + */ +#ifndef BUFFER_MANAGER_H +#define BUFFER_MANAGER_H + +#if (__NPU_ARCH__ == 5102) +#include "buffer_mix_core.h" +#else +#include "buffer.h" +#endif + +// L1 TPosition::A1 +// L0A TPosition::A2 +// L0B TPosition::B2 +// L0C TPosition::CO1 +// UB TPosition::VECIN +namespace fa_base_matmul { +template +class BufferManager { + using TensorType = std::conditional_t, LocalTensor>; +public: + __aicore__ inline void Init(TPipe *pipe, uint32_t size) { + static_assert(bufferType != BufferType::GM, "GM should use workspace."); + TBuf::Position> tbuf; + pipe->InitBuffer(tbuf, size); + mem_ = tbuf.template Get(); + } + + __aicore__ inline void Init(__gm__ uint8_t* workspace) { + static_assert(bufferType == BufferType::GM, "BufferType should be GM."); + mem_.SetGlobalBuffer((__gm__ uint8_t*)workspace); + } + + template + __aicore__ inline Buffer AllocBuffer(uint32_t size) { + TensorType temp = mem_[offset_]; + offset_ += size; + return Buffer(temp, size); + } + + template + __aicore__ inline void FreeBuffer(Buffer &buffer){ + } +private: + uint32_t offset_ = 0; + TensorType mem_; +}; +} +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/common/op_kernel/buffer_mix_core.h b/csrc/ascend/attention/common/op_kernel/buffer_mix_core.h new file mode 100644 index 000000000..80a41316f --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/buffer_mix_core.h @@ -0,0 +1,216 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file buffer_mix_core.h + * \brief同步管理 + */ +#ifndef BUFFER_MIX_CORE_H +#define BUFFER_MIX_CORE_H +#include +#include "lib/matmul_intf.h" +#if ASC_DEVKIT_MAJOR >= 9 +#include "kernel_basic_intf.h" +#else +#include "kernel_operator.h" +#endif +using namespace AscendC; +namespace fa_base_matmul { +__BLOCK_LOCAL__ __inline__ uint32_t idCounterNum; +#define MAKE_ID ((++idCounterNum) % 11) + +// 核间同步中，AIC(flagId 0-10)对应AIV0(flagId 0-10)，对应AIV1(flagId 16-26) +#define AIV0_AIV1_OFFSET 16 + +enum class BufferType { + L1 = 0, + L0A = 1, + L0B = 2, + L0C = 3, + UB = 4, + GM = 5, +}; + +enum class SyncType { + NO_SYNC, + INNER_CORE_SYNC, + CROSS_CORE_SYNC_FORWARD, + CROSS_CORE_SYNC_BOTH, +}; + +constexpr uint32_t INVALID_CROSS_CORE_EVENT_ID = 16; +static constexpr uint64_t CROSS_CORE_SYNC_MODE = 4; + +template +struct BufferInfo { + // Cons 消费者，Prod 生产者 + __aicore__ const static constexpr HardEvent ConsWaitProdStatus() + { + if constexpr (Type == BufferType::L1) { + return HardEvent::MTE2_MTE1; + } else if constexpr (Type == BufferType::L0A) { + return HardEvent::MTE1_M; + } else if constexpr (Type == BufferType::L0B) { + return HardEvent::MTE1_M; + } else if constexpr (Type == BufferType::L0C) { + return HardEvent::M_FIX; + } + } + + __aicore__ const static constexpr HardEvent ProdWaitConsStatus() + { + if constexpr (Type == BufferType::L1) { + return HardEvent::MTE1_MTE2; + } else if constexpr (Type == BufferType::L0A) { + return HardEvent::M_MTE1; + } else if constexpr (Type == BufferType::L0B) { + return HardEvent::M_MTE1; + } else if constexpr (Type == BufferType::L0C) { + return HardEvent::FIX_M; + } + } + + __aicore__ const static constexpr TPosition GetTPosition() + { + if constexpr (Type == BufferType::L1) { + return TPosition::A1; + } else if constexpr (Type == BufferType::L0A) { + return TPosition::A2; + } else if constexpr (Type == BufferType::L0B) { + return TPosition::B2; + } else if constexpr (Type == BufferType::L0C) { + return TPosition::CO1; + } else if constexpr (Type == BufferType::UB) { + return TPosition::VECIN; + } else if constexpr (Type == BufferType::GM) { + return TPosition::GM; + } + } + + static constexpr HardEvent EventP2C = + ConsWaitProdStatus(); // 生产者到消费者方向的HardEvent：消费者等生产者提供/生产者通知消费者已生成 + static constexpr HardEvent EventC2P = + ProdWaitConsStatus(); // 消费者到生产者方向的HardEvent：生产者等消费者消耗/消费者通知生产者已消耗’ + static constexpr TPosition Position = GetTPosition(); +}; + +// buffer绑定生产者、消费者关系 +// L1 buffer的生产者为MTE2或者MTE3，消费者为MTE1 +// L0A buffer的生产者为MTE1，消费者为M +// L0B buffer的生产者为MTE1，消费者为M +// L0C buffer的生产者为M，消费者为FIX +template +class Buffer { + using TensorType = std::conditional_t, LocalTensor>; + + template + using TargetTensorType = std::conditional_t, LocalTensor>; + +public: + __aicore__ inline Buffer() + { + } + __aicore__ inline Buffer(TensorType tensor, uint32_t size) + { + tensor_ = tensor; + size_ = size; + if constexpr (syncType == SyncType::CROSS_CORE_SYNC_FORWARD) { + id0_ = MAKE_ID; + id1_ = INVALID_CROSS_CORE_EVENT_ID; + } else if constexpr (syncType == SyncType::CROSS_CORE_SYNC_BOTH) { + id0_ = MAKE_ID; + id1_ = MAKE_ID; + } else { + id0_ = INVALID_CROSS_CORE_EVENT_ID; + id1_ = INVALID_CROSS_CORE_EVENT_ID; + } + } + + __aicore__ inline void Init() + { + if constexpr (syncType == SyncType::INNER_CORE_SYNC) { + p2cEventId_ = GetTPipePtr()->AllocEventID::EventP2C>(); // 确保只能被调用一次 + c2pEventId_ = GetTPipePtr()->AllocEventID::EventC2P>(); + SetFlag::EventC2P>(c2pEventId_); + } + } + + __aicore__ inline void UnInit() + { + if constexpr (syncType == SyncType::INNER_CORE_SYNC) { + WaitFlag::EventC2P>(c2pEventId_); + GetTPipePtr()->ReleaseEventID::EventP2C>(p2cEventId_); // 确保只能被调用一次 + GetTPipePtr()->ReleaseEventID::EventC2P>(c2pEventId_); + } + } + + template + __aicore__ inline void Wait() + { + if constexpr (syncType == SyncType::INNER_CORE_SYNC) { + if constexpr (EventType == BufferInfo::EventP2C) { + WaitFlag::EventP2C>(p2cEventId_); // 消费者等待生产者完成生产 + } else { + WaitFlag::EventC2P>(c2pEventId_); // 生产者等待消费者完成消费 + } + } + } + + template + __aicore__ inline void Set() + { + if constexpr (syncType == SyncType::INNER_CORE_SYNC) { + if constexpr (EventType == BufferInfo::EventP2C) { + SetFlag::EventP2C>(p2cEventId_); // 生产者通知消费者已完成生产 + } else { + SetFlag::EventC2P>(c2pEventId_); // 消费者通知生产者已完成消费 + } + } + } + + __aicore__ inline void SetEventID() + { + p2cEventId_ = GetTPipePtr()->AllocEventID::EventP2C>(); // 确保只能被调用一次 + c2pEventId_ = GetTPipePtr()->AllocEventID::EventC2P>(); + } + + template + __aicore__ inline TEventID GetEventID() + { + if constexpr (EventType == BufferInfo::EventP2C) { + return p2cEventId_; // 生产者通知消费者已完成生产 + } else { + return c2pEventId_; // 消费者通知生产者已完成消费 + } + } + + template + __aicore__ inline TargetTensorType GetTensor() + { + return tensor_.template ReinterpretCast(); + } + + template + __aicore__ inline TargetTensorType GetTensor(uint64_t startindex) + { + TargetTensorType tmpTensor = tensor_.template ReinterpretCast(); + return tmpTensor[startindex]; + } + +private: + TensorType tensor_; + uint32_t size_; + TEventID p2cEventId_; + TEventID c2pEventId_; + uint32_t id0_; // 用作正向同步：生产者通知消费者，或者消费者等待生产者； + uint32_t id1_; // 用作反向同步：消费者通知生产者，或者生产者等待消费者； +}; +} // namespace fa_base_matmul +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/common/op_kernel/buffers_policy.h b/csrc/ascend/attention/common/op_kernel/buffers_policy.h new file mode 100644 index 000000000..7c5e924a9 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/buffers_policy.h @@ -0,0 +1,407 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file buffers_policy.h + * \brief 综合管理buffer的内存和同步 + */ +#ifndef BUFFERS_POLICY_H +#define BUFFERS_POLICY_H + +#include "buffer_manager.h" +#define NUM_2 2 +#define NUM_3 3 +#define NUM_4 4 +// Q复用 KV复用 +// 申请单块buffer +namespace fa_base_matmul { +template +class BuffersPolicySingleBuffer { +public: + __aicore__ inline void Init(BufferManager &bufferManager, uint32_t size){ + buffer_ = bufferManager.template AllocBuffer(size); + buffer_.Init(); + } + + __aicore__ inline void Uninit(BufferManager &bufferManager){ + buffer_.UnInit(); + bufferManager.FreeBuffer(buffer_); + } + + __aicore__ inline Buffer &Get(){ + return buffer_; + } + + __aicore__ inline Buffer &GetPre(){ + return Get(); + } + + __aicore__ inline Buffer &GetReused(){ + return Get(); + } +private: + Buffer buffer_; +}; + +// 申请2个buffer，乒乓轮转 +template +class BuffersPolicyDB { +public: + __aicore__ inline void Init(BufferManager &bufferManager, uint32_t size){ + ping_ = bufferManager.template AllocBuffer(size); + pong_ = bufferManager.template AllocBuffer(size); + + ping_.Init(); + pong_.Init(); + } + + __aicore__ inline void Uninit(BufferManager &bufferManager){ + ping_.UnInit(); + pong_.UnInit(); + + bufferManager.FreeBuffer(ping_); + bufferManager.FreeBuffer(pong_); + } + + __aicore__ inline Buffer &Get() { + if (flag1_) { // 1 + flag1_ = 0; + return ping_; + } else { // 0 + flag1_ = 1; + return pong_; + } + } + + // 需要与Get联用， 首次调用Get，第二次调用GetPre(Q复用) + __aicore__ inline Buffer &GetPre() { + if (flag1_) { // 0->1 + return pong_; + } else { // 1->0 + return ping_; + } + } + + // 需要与Get,GetPre联用， 首次调用Get，第二次调用GetPre,第三次复用时GetReused(KV复用) + __aicore__ inline Buffer &GetReused() { + if (flag2_ == 0) { + flag2_ = 1; + return pong_; + } else { + flag2_ = 0; + return ping_; + } + } + + __aicore__ inline Buffer &GetReused(bool isNextS2IdxNoChange) { + if (isNextS2IdxNoChange) { + if (flag2_ == 0) { + return pong_; + } else { + return ping_; + } + } else { + return GetReused(); + } + } + +private: + Buffer ping_; + Buffer pong_; + uint32_t flag1_ = 0; + uint32_t flag2_ = 0; +}; + +// 申请3个buffer, 轮转 +template +class BuffersPolicy3buff { +public: + __aicore__ inline void Init(BufferManager &bufferManager, uint32_t size) { + a_ = bufferManager.template AllocBuffer(size); + b_ = bufferManager.template AllocBuffer(size); + c_ = bufferManager.template AllocBuffer(size); + + a_.Init(); + b_.Init(); + c_.Init(); + } + + __aicore__ inline void Uninit(BufferManager &bufferManager) { + a_.UnInit(); + b_.UnInit(); + c_.UnInit(); + + bufferManager.FreeBuffer(a_); + bufferManager.FreeBuffer(b_); + bufferManager.FreeBuffer(c_); + } + + __aicore__ inline Buffer &Get() { + if (flag1_ == 0) { + flag1_ = 1; + return a_; + } else if (flag1_ == 1) { + flag1_ = NUM_2; + return b_; + } else { + flag1_ = 0; + return c_; + } + } + + __aicore__ inline Buffer &GetVec() { // mixcore architecture + if (flag1_vec1_ == 0) { + flag1_vec1_ = 1; + return a_; + } else if (flag1_vec1_ == 1) { + flag1_vec1_ = NUM_2; + return b_; + } else { + flag1_vec1_ = 0; + return c_; + } + } + + __aicore__ inline Buffer &GetCube() { // mixcore architecture + if (flag1_bmm2_ == 0) { + flag1_bmm2_ = 1; + return a_; + } else if (flag1_bmm2_ == 1) { + flag1_bmm2_ = NUM_2; + return b_; + } else { + flag1_bmm2_ = 0; + return c_; + } + } + + // Q复用 + __aicore__ inline Buffer &GetPre() { + if (flag1_ == 0) { + return c_; + } else if (flag1_ == 1) { + return a_; + } else { + return b_; + } + } + + // KV复用 + __aicore__ inline Buffer &GetReused() { + if (flag2_ == 0) { + flag2_ = 1; + return a_; + } else if (flag2_ == 1){ + flag2_ = NUM_2; + return b_; + } else { + flag2_ = 0; + return c_; + } + } +private: + Buffer a_; + Buffer b_; + Buffer c_; + uint32_t flag1_ = 0; + uint32_t flag1_vec1_ = 0; + uint32_t flag1_bmm2_ = 0; + uint32_t flag2_ = 0; +}; + +// 申请4个buffer + kv复用 +template +class BuffersPolicy4buff { +public: + __aicore__ inline void Init(BufferManager &bufferManager, uint32_t size) { + a_ = bufferManager.template AllocBuffer(size); + b_ = bufferManager.template AllocBuffer(size); + c_ = bufferManager.template AllocBuffer(size); + d_ = bufferManager.template AllocBuffer(size); + + a_.Init(); + b_.Init(); + c_.Init(); + d_.Init(); + } + + __aicore__ inline void Uninit(BufferManager &bufferManager) { + a_.UnInit(); + b_.UnInit(); + c_.UnInit(); + d_.UnInit(); + + bufferManager.FreeBuffer(a_); + bufferManager.FreeBuffer(b_); + bufferManager.FreeBuffer(c_); + bufferManager.FreeBuffer(d_); + } + + __aicore__ inline Buffer &Get(uint32_t id) { + uint32_t flag = id % 4; + if (flag == 0) { + return a_; + } else if (flag == 1) { + return b_; + } else if (flag == 2) { // 2:c_ + return c_; + } else { + return d_; + } + } + + __aicore__ inline Buffer &Get() { + auto& buffer = Get(head_); + head_++; + return buffer; + } + + __aicore__ inline Buffer &GetReused() { + auto& buffer = Get(used_); + used_ = (used_ - tail_ + 1) % (head_ - tail_) + tail_; + return buffer; + } + + __aicore__ inline Buffer &GetFree() { + if (tail_ == used_) { + used_++; + } + auto& buffer = Get(tail_); + tail_++; + return buffer; + } +private: + Buffer a_; + Buffer b_; + Buffer c_; + Buffer d_; + uint32_t tail_ = 0; // 表示当前正在使用的buffer队列队尾 + uint32_t head_ = 0; // 表示当前正在使用的buffer队列队首+1 + uint32_t used_ = 0; // 表示当前正在使用的buffer，于首尾间，左闭右开 +}; + +template +class Matrix2x2BufferPolicy { // 4buffer +// 二维buffer管理，地址行优先，使用列优先 +// MracBuffer:memory address with row first, alloc/use/free with column first +public: + __aicore__ inline void Init(BufferManager &bufferManager, uint32_t size) { + bufferM0k0_ = bufferManager.template AllocBuffer(size); + bufferM0k1_ = bufferManager.template AllocBuffer(size); + bufferM1k0_ = bufferManager.template AllocBuffer(size); + bufferM1k1_ = bufferManager.template AllocBuffer(size); + + bufferM0k0_.Init(); + bufferM0k1_.Init(); + bufferM1k0_.Init(); + bufferM1k1_.Init(); + } + + __aicore__ inline void Uninit(BufferManager &bufferManager) { + bufferM0k0_.UnInit(); + bufferM0k1_.UnInit(); + bufferM1k0_.UnInit(); + bufferM1k1_.UnInit(); + + bufferManager.FreeBuffer(bufferM0k0_); + bufferManager.FreeBuffer(bufferM0k1_); + bufferManager.FreeBuffer(bufferM1k0_); + bufferManager.FreeBuffer(bufferM1k1_); + } + + __aicore__ inline void SetMExtent(int32_t mExtent) { + aIdx_ = -1; + amIdx_ = (amIdx_ + mSize_ - 1) % mSize_; // 翻转 0->1, 1->0 + akIdx_ = 0; + + uIdx_ = -1; + umIdx_ = (umIdx_ + mSize_ - 1) % mSize_; + ukIdx_ = 0; + + fIdx_ = -1; + fmIdx_ = (fmIdx_ + mSize_ - 1) % mSize_; + fkIdx_ = 0; + + mExtent_ = mExtent; + } + + __aicore__ inline Buffer &AllocNext() { + aIdx_++; + return GetBuffer(aIdx_, amIdx_, akIdx_); + } + + __aicore__ inline Buffer &ReuseNext() { + uIdx_++; + return GetBuffer(uIdx_, umIdx_, ukIdx_); + } + + __aicore__ inline Buffer &FreeNext() { + fIdx_++; + return GetBuffer(fIdx_, fmIdx_, fkIdx_); + } + + __aicore__ inline Buffer &PeekNextK() { // 在Alloc阶段使用，k方向取下一个 + return PeekBuffer(amIdx_, (1 - akIdx_)); // k翻转 + } +private: + __aicore__ inline Buffer &GetBuffer(int32_t xIdx, int32_t &mIdx, int32_t &kIdx) { + // xIdx为入参，表示当前alloc/use/free的idx，mIdx和kIdx为下标出参，移动到下一个buffer并获取 + mIdx = (mIdx + mExtent_ - 1) % mExtent_; + kIdx = (xIdx / mExtent_) % kSize_; + if (mIdx == 0 && kIdx == 0) { + return bufferM0k0_; + } else if (mIdx == 0 && kIdx == 1) { + return bufferM0k1_; + } else if (mIdx == 1 && kIdx == 0) { + return bufferM1k0_; + } else { // 该分支条件为：mIdx == 1 && kIdx == 1 + return bufferM1k1_; + } + } + + __aicore__ inline Buffer &PeekBuffer(int32_t mIdx, int32_t kIdx) { + // 只访问buffer，不进行下标移动 + if (mIdx == 0 && kIdx == 0) { + return bufferM0k0_; + } else if (mIdx == 0 && kIdx == 1) { + return bufferM0k1_; + } else if ((mIdx == 1) && (kIdx == 0)) { + return bufferM1k0_; + } else { // mIdx == 1 && kIdx == 1 + return bufferM1k1_; + } + } + + Buffer bufferM0k0_; + Buffer bufferM0k1_; + Buffer bufferM1k0_; + Buffer bufferM1k1_; + int32_t mSize_ = 2; // m的总buffer数 + int32_t kSize_ = 2; // k的总buffer数 + + // Alloc + int32_t aIdx_ = -1; // 当前第几次Alloc Buffer + int32_t amIdx_ = 0; // 当前Alloc Buffer的m下标 + int32_t akIdx_ = 0; // 当前Alloc Buffer的k下标 + + // Reuse + int32_t uIdx_ = -1; + int32_t umIdx_ = 0; + int32_t ukIdx_ = 0; + + // Free + int32_t fIdx_ = -1; + int32_t fmIdx_ = 0; + int32_t fkIdx_ = 0; + + int32_t mExtent_ = 0; // m实际使用的大小，可以为1或者2 +}; +} +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/common/op_kernel/matmul.h b/csrc/ascend/attention/common/op_kernel/matmul.h new file mode 100644 index 000000000..d152d690f --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/matmul.h @@ -0,0 +1,1158 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file matmul.h + * \brief + */ +#ifndef MATMUL_H +#define MATMUL_H +#include "buffers_policy.h" +using namespace AscendC; + +namespace fa_base_matmul { + +constexpr uint32_t UNITFLAG_DISABLE = 0; +constexpr uint32_t UNITFLAG_ENABLE = 2; +constexpr uint32_t UNITFLAG_EN_OUTER_LAST = 3; +static constexpr uint32_t FP16_ONE_FRACTAL_ELEMENT = 16; // 一个分形512B,16*16个fp16 +static constexpr uint32_t INT4_ONE_FRACTAL_ELEMENT = 64; // 一个分形512B,16*64个fp16 +static constexpr uint32_t ONE_FRACTAL_H_ELEMENT = 16; // 一个分形512B,height方向为16个element +static constexpr uint32_t ONE_FRACTAL_W_BYTE = 32; // 一个分形512B,weight方向为32B +static constexpr uint32_t LOAD3D_L1W_SIZE = 16; +static constexpr uint32_t MMAD_MN_SIZE_10 = 10; +static constexpr uint8_t LOAD3D_STRIDE_W = 1; +static constexpr uint8_t LOAD3D_STRIDE_H = 1; +static constexpr uint8_t LOAD3D_FILTER_W = 1; +static constexpr uint8_t LOAD3D_FILTER_H = 1; +static constexpr uint8_t LOAD3D_DILA_FILTER_W = 1; +static constexpr uint8_t LOAD3D_DILA_FILTER_H = 1; +static constexpr uint32_t K_STEP_ALIGN_BASE = 2; +static constexpr uint32_t M_STEP_ALIGN_BASE = 2; +static constexpr uint32_t MX_FP4_PTG_PCG_SCALE_PARAM = 32; // mxfp4场景随路伪量化中32个元素共用一个scale +static constexpr uint32_t HI_FP4_PTG_PCG_SCALE_PARAM = 64; // hifp4场景随路伪量化中64个元素共用一个scale + +struct MMParam { + uint32_t singleM; + uint32_t singleN; + uint32_t singleK; + bool isLeftTranspose; + bool isRightTranspose; + bool cmatrixInitVal = true; + bool isOutKFisrt = true; // 默认值为true， true：在L1切K轴的场景中，表示首轮K + uint32_t unitFlag = 0; // 0：disable: 不配置unitFlag + // 2：enable: 行为在切K接口中（MatmulK），会将mmadParams.unitFlag设置为 0b10 + // 3：enable: 行为在切K接口中（MatmulK），在k的最后一轮循环，会将mmadParams.unitFlag设置为 0b11 + // 外部使用时，在外层k循环的最后一轮将该参数配置为3 + uint32_t realM = 0; // bmm2以s1realsize为M轴，不赋值时不影响现有代码逻辑 +}; + +__aicore__ inline MMParam MakeMMParam(uint32_t singleM, uint32_t singleN, uint32_t singleK, bool isLeftTranspose, + bool isRightTranspose, bool cmatrixInitVal = true, bool isOutKFisrt = true, + uint32_t unitFlag = 0, uint32_t realM = 0) +{ + return {.singleM = singleM, + .singleN = singleN, + .singleK = singleK, + .isLeftTranspose = isLeftTranspose, + .isRightTranspose = isRightTranspose, + .cmatrixInitVal = cmatrixInitVal, + .isOutKFisrt = isOutKFisrt, + .unitFlag = unitFlag, + .realM = realM + }; +} + +enum class ABLayout { + MK = 0, + KM = 1, + KN = 2, + NK = 3, +}; + +template +__aicore__ inline T AlignUp(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd) * (rnd))); +} + +#if ((__CCE_AICORE__ == 310) || (defined __DAV_310R6__) || (__NPU_ARCH__ == 5102)) +template +__aicore__ inline uint32_t GetBlockNum(uint32_t size) { + if constexpr (IsSameType::value) { + return ((size + 7) >> 3 << 3) >> 3; + } else if constexpr ((IsSameType::value || + IsSameType::value || + IsSameType::value || + IsSameType::value)) { + return ((size + 31) >> 5 << 5) >> 5; + } else { + return ((size + 15) >> 4 << 4) >> 4; + } +} +// L1->L0A + 切k/切M/全载 +template +__aicore__ inline void LoadDataToL0A(LocalTensor& aL0Tensor, const LocalTensor& aL1Tensor, + const MMParam& mmParam, uint64_t L1Aoffset, uint32_t kSplitSize, uint32_t mSplitSize) +{ + LoadData2DParamsV2 loadData2DParamsA; // 基础API LoadData的参数结构体 + loadData2DParamsA.mStartPosition = 0; // 以M*K矩阵为例，源矩阵M轴方向的起始位置，单位为16 element + loadData2DParamsA.kStartPosition = 0; // 以M*K矩阵为例，源矩阵K轴方向的起始位置，单位为32B + loadData2DParamsA.ifTranspose = mmParam.isLeftTranspose; // 是否启用转置功能，对每个分型矩阵进行转置 + if (loadData2DParamsA.ifTranspose) { + loadData2DParamsA.mStep = ((kSplitSize + 15) >> 4 << 4) >> 4; // 以M*K矩阵为例,源矩阵M轴方向搬运长度(S1向上对齐分形(512B),16*16个f16->向上对齐16)，单位为16 element,取值范围：mStep属于[0,255] + if constexpr (IsSameType::value || IsSameType::value || IsSameType::value || IsSameType::value) { + loadData2DParamsA.mStep = (loadData2DParamsA.mStep + 1) >> 1 << 1; + } + loadData2DParamsA.kStep = GetBlockNum(mSplitSize); // 以M*K矩阵为例,源矩阵K轴方向搬运长度(qkD个f16)，单位为32B,取值范围：nStep属于[0,255] + } else { + loadData2DParamsA.mStep = ((mSplitSize + 15) >> 4 << 4) >> 4; // 以M*K矩阵为例,源矩阵M轴方向搬运长度(S1向上对齐分形(512B),16*16个f16->向上对齐16)，单位为16 element,取值范围：mStep属于[0,255] + loadData2DParamsA.kStep = GetBlockNum(kSplitSize); // 以M*K矩阵为例,源矩阵K轴方向搬运长度(qkD个f16)，单位为32B,取值范围：nStep属于[0,255] + } + if constexpr (IsSameType::value) { + if (loadData2DParamsA.ifTranspose) { + loadData2DParamsA.kStep = CeilAlign(loadData2DParamsA.kStep, K_STEP_ALIGN_BASE); + } + } + if constexpr (IsSameType::value || IsSameType::value || IsSameType::value || IsSameType::value) { + // 配合ub->L1使用 256 * 32 / 256 + // 64搬运对齐 + loadData2DParamsA.srcStride = loadData2DParamsA.ifTranspose ? ((kSplitSize + 63) >> 6 << 6) >> 4 : ((mSplitSize + 31) >> 5 << 5) >> 4; // 以M*K矩阵为例，源矩阵K方向前一个分形起始地址与后一个分形起始地址的间隔，单位：512B + } else { + loadData2DParamsA.srcStride = loadData2DParamsA.ifTranspose ? ((mmParam.singleK + 15) >> 4 << 4) >> 4 : loadData2DParamsA.mStep; + } + if (mmParam.realM != 0) { + loadData2DParamsA.mStep = ((mmParam.realM + 15) >> 4 << 4) >> 4; + } + loadData2DParamsA.dstStride = loadData2DParamsA.ifTranspose ? (mSplitSize + 15) >> 4 : loadData2DParamsA.mStep; + if constexpr (IsSameType::value || IsSameType::value || IsSameType::value || IsSameType::value) { + if (loadData2DParamsA.ifTranspose && (loadData2DParamsA.dstStride & 1)) { + uint32_t l0bLoop = (loadData2DParamsA.mStep + 1) >> 1; + loadData2DParamsA.mStep = M_STEP_ALIGN_BASE; + uint64_t dstOffset = 0; + uint64_t dstAddrStride = (mSplitSize + 15) / 16 * 16 * 32; + uint16_t oriMStep = loadData2DParamsA.mStartPosition; + for (uint32_t idx = 0; idx < l0bLoop; ++idx) { + loadData2DParamsA.mStartPosition = oriMStep + M_STEP_ALIGN_BASE * idx; + LoadData(aL0Tensor[dstOffset], aL1Tensor[L1Aoffset], loadData2DParamsA); + dstOffset += dstAddrStride; + } + } else { + LoadData(aL0Tensor, aL1Tensor[L1Aoffset], loadData2DParamsA); + } + } else { + LoadData(aL0Tensor, aL1Tensor[L1Aoffset], loadData2DParamsA); + } +} + +// L1->L0A + 切k/切M/全载 +template +__aicore__ inline void LoadDataToL0AMx(LocalTensor& aL0Tensor, const LocalTensor& aL1Tensor, const LocalTensor& aScaleL1Tensor, + const MMParam& mmParam, uint64_t L1Aoffset, uint32_t kSplitSize, uint32_t mSplitSize) +{ + LoadData2DParamsV2 loadData2DParamsA; // 基础API LoadData的参数结构体 + loadData2DParamsA.mStartPosition = 0; // 以M*K矩阵为例，源矩阵M轴方向的起始位置，单位为16 element + loadData2DParamsA.kStartPosition = 0; // 以M*K矩阵为例，源矩阵K轴方向的起始位置，单位为32B + loadData2DParamsA.ifTranspose = mmParam.isLeftTranspose; // 是否启用转置功能，对每个分型矩阵进行转置 + if (loadData2DParamsA.ifTranspose) { + loadData2DParamsA.mStep = ((kSplitSize + 15) >> 4 << 4) >> 4; // 以M*K矩阵为例,源矩阵M轴方向搬运长度(S1向上对齐分形(512B),16*16个f16->向上对齐16)，单位为16 element,取值范围：mStep属于[0,255] + if constexpr (IsSameType::value || IsSameType::value || IsSameType::value) { + loadData2DParamsA.mStep = (loadData2DParamsA.mStep + 1) >> 1 << 1; + } + loadData2DParamsA.kStep = GetBlockNum(mSplitSize); // 以M*K矩阵为例,源矩阵K轴方向搬运长度(qkD个f16)，单位为32B,取值范围：nStep属于[0,255] + } else { + loadData2DParamsA.mStep = ((mSplitSize + 15) >> 4 << 4) >> 4; // 以M*K矩阵为例,源矩阵M轴方向搬运长度(S1向上对齐分形(512B),16*16个f16->向上对齐16)，单位为16 element,取值范围：mStep属于[0,255] + loadData2DParamsA.kStep = GetBlockNum(kSplitSize); // 以M*K矩阵为例,源矩阵K轴方向搬运长度(qkD个f16)，单位为32B,取值范围：nStep属于[0,255] + } + if constexpr (IsSameType::value || IsSameType::value || IsSameType::value) { + // 配合ub->L1使用 256 * 32 / 256 + // 64搬运对齐 + loadData2DParamsA.srcStride = loadData2DParamsA.ifTranspose ? 256 >> 4 : ((mSplitSize + 31) >> 5 << 5) >> 4; // 以M*K矩阵为例，源矩阵K方向前一个分形起始地址与后一个分形起始地址的间隔，单位：512B + } else { + loadData2DParamsA.srcStride = loadData2DParamsA.ifTranspose ? ((mmParam.singleK + 15) >> 4 << 4) >> 4 : loadData2DParamsA.mStep; + } + LoadData2DMxParams loadData2DMxParamsA; + loadData2DMxParamsA.xStartPosition = 0; + loadData2DMxParamsA.yStartPosition = 0; + loadData2DMxParamsA.xStep = ((mSplitSize + 15) >> 4 << 4) >> 4; + loadData2DMxParamsA.yStep = (kSplitSize + 63) >> 5 >> 1; + loadData2DMxParamsA.srcStride = loadData2DMxParamsA.yStep; + loadData2DMxParamsA.dstStride = loadData2DMxParamsA.yStep; + if (mmParam.realM != 0) { + loadData2DParamsA.mStep = ((mmParam.realM + 15) >> 4 << 4) >> 4; + loadData2DMxParamsA.xStep = ((mmParam.realM + 15) >> 4 << 4) >> 4; + } + loadData2DParamsA.dstStride = loadData2DParamsA.ifTranspose ? (mSplitSize + 15) >> 4 : loadData2DParamsA.mStep; + if constexpr (IsSameType::value || IsSameType::value || IsSameType::value) { + if (loadData2DParamsA.ifTranspose && (loadData2DParamsA.dstStride & 1)) { + uint32_t l0bLoop = (loadData2DParamsA.mStep + 1) >> 1; + loadData2DParamsA.mStep = M_STEP_ALIGN_BASE; + loadData2DMxParamsA.xStep = loadData2DParamsA.mStep ; + uint64_t dstOffset = 0; + uint64_t dstAddrStride = (mSplitSize + 15) / 16 * 16 * 32; + uint16_t oriMStep = loadData2DParamsA.mStartPosition; + uint16_t oriMScaleStep = loadData2DMxParamsA.xStartPosition; + for (uint32_t idx = 0; idx < l0bLoop; ++idx) { + loadData2DParamsA.mStartPosition = oriMStep + M_STEP_ALIGN_BASE * idx; + loadData2DMxParamsA.xStartPosition = oriMScaleStep + M_STEP_ALIGN_BASE * idx; + LoadData(aL0Tensor[dstOffset], aL1Tensor[L1Aoffset], aScaleL1Tensor[L1Aoffset >> 5], loadData2DParamsA, + loadData2DMxParamsA); + dstOffset += dstAddrStride; + } + } else { + LoadData(aL0Tensor, aL1Tensor[L1Aoffset], aScaleL1Tensor[L1Aoffset >> 5], loadData2DParamsA, + loadData2DMxParamsA); + } + } else { + LoadData(aL0Tensor, aL1Tensor[L1Aoffset], loadData2DParamsA); + } +} + +// L1->L0B + 切k/切M/全载 +template +__aicore__ inline void LoadDataToL0B(LocalTensor& bL0Tensor, const LocalTensor& bL1Tensor, + const MMParam& mmParam, uint64_t L1Boffset, uint32_t kSplitSize, uint32_t nSplitSize, int nLoops = 1) +{ + LoadData2DParamsV2 loadData2DParamsB; // 基础API LoadData的参数结构体 + loadData2DParamsB.mStartPosition = 0; // 以M*K矩阵为例，源矩阵M轴方向的起始位置，单位为16 element + loadData2DParamsB.kStartPosition = 0; // 以M*K矩阵为例，源矩阵K轴方向的起始位置，单位为32B + loadData2DParamsB.ifTranspose = !mmParam.isRightTranspose; // 是否启用转置功能，对每个分型矩阵进行转置 + if (loadData2DParamsB.ifTranspose) { + loadData2DParamsB.mStep = ((kSplitSize + 15) >> 4 << 4) >> 4; // 以M*K矩阵为例,源矩阵M轴方向搬运长度(S1向上对齐分形(512B),16*16个f16->向上对齐16)，单位为16 element,取值范围：mStep属于[0,255] + if constexpr (IsSameType::value || IsSameType::value || IsSameType::value || IsSameType::value) { + loadData2DParamsB.mStep = (loadData2DParamsB.mStep + 1) >> 1 << 1; + } + loadData2DParamsB.kStep = GetBlockNum(nSplitSize); // 以M*K矩阵为例,源矩阵K轴方向搬运长度(qkD个f16)，单位为32B,取值范围：nStep属于[0,255] + } else { + loadData2DParamsB.mStep = ((nSplitSize + 15) >> 4 << 4) >> 4; // 以M*K矩阵为例,源矩阵M轴方向搬运长度(S1向上对齐分形(512B),16*16个f16->向上对齐16)，单位为16 element,取值范围：mStep属于[0,255] + loadData2DParamsB.kStep = GetBlockNum(kSplitSize); // 以M*K矩阵为例,源矩阵K轴方向搬运长度(qkD个f16)，单位为32B,取值范围：nStep属于[0,255] + } + if constexpr (IsSameType::value) { + if (loadData2DParamsB.ifTranspose) { + loadData2DParamsB.kStep = CeilAlign(loadData2DParamsB.kStep, K_STEP_ALIGN_BASE); + } + } + if constexpr (IsSameType::value || IsSameType::value || IsSameType::value || IsSameType::value) { + if (loadData2DParamsB.ifTranspose) { + loadData2DParamsB.srcStride = ((kSplitSize + 31) >> 5 << 5) >> 4; + } else { + loadData2DParamsB.srcStride = ((nSplitSize + 31) >> 5 << 5) >> 4; + } + } else { + loadData2DParamsB.srcStride = loadData2DParamsB.ifTranspose ? (((mmParam.singleK + 15) >> 4 << 4) >> 4) : (((mmParam.singleN + 15 ) >> 4 << 4) >> 4); // 以M*K矩阵为例，源矩阵K方向前一个分形起始地址与后一个分形起始地址的间隔，单位：512B + } + loadData2DParamsB.dstStride = loadData2DParamsB.ifTranspose ? (nSplitSize + 15) >> 4 : loadData2DParamsB.mStep; // 以M*K矩阵为例，目标矩阵K方向前一个分形起始地址与后一个分形起始地址的间隔，单位：512B + if constexpr (IsSameType::value || IsSameType::value || IsSameType::value || IsSameType::value) { + if (loadData2DParamsB.ifTranspose && (loadData2DParamsB.dstStride & 1)) { + uint32_t l0bLoop = (loadData2DParamsB.mStep + 1) >> 1; + loadData2DParamsB.mStep = M_STEP_ALIGN_BASE; + uint64_t dstOffset = 0; + uint64_t dstAddrStride = (nSplitSize + 15) / 16 * 16 * 32; + uint16_t oriMStep = loadData2DParamsB.mStartPosition; + for (uint32_t idx = 0; idx < l0bLoop; ++idx) { + loadData2DParamsB.mStartPosition = oriMStep + M_STEP_ALIGN_BASE * idx; + LoadData(bL0Tensor[dstOffset], bL1Tensor[L1Boffset], loadData2DParamsB); + dstOffset += dstAddrStride; + } + } else { + LoadData(bL0Tensor, bL1Tensor[L1Boffset], loadData2DParamsB); + } + } else { + LoadData(bL0Tensor, bL1Tensor[L1Boffset], loadData2DParamsB); + } +} + +template +__aicore__ inline void LoadDataToL0BMx(LocalTensor& bL0Tensor, const LocalTensor& bL1Tensor, const LocalTensor& bScaleL1Tensor, + const MMParam& mmParam, uint64_t L1Boffset, uint32_t kSplitSize, uint32_t nSplitSize, int nLoops = 1) +{ + LoadData2DParamsV2 loadData2DParamsB; // 基础API LoadData的参数结构体 + loadData2DParamsB.mStartPosition = 0; // 以M*K矩阵为例，源矩阵M轴方向的起始位置，单位为16 element + loadData2DParamsB.kStartPosition = 0; // 以M*K矩阵为例，源矩阵K轴方向的起始位置，单位为32B + loadData2DParamsB.ifTranspose = !mmParam.isRightTranspose; // 是否启用转置功能，对每个分型矩阵进行转置 + if (loadData2DParamsB.ifTranspose) { + if constexpr (IsSameType::value || IsSameType::value) { + loadData2DParamsB.mStep = ((mmParam.singleK + 63) >> 6 << 6) / 16; + loadData2DParamsB.kStep = kSplitSize / 64; + loadData2DParamsB.srcStride = ((mmParam.singleK + 15) >> 4 << 4) / 16; + loadData2DParamsB.dstStride = kSplitSize / 16; + } else if constexpr (IsSameType::value) { + loadData2DParamsB.mStep = ((kSplitSize + 15) >> 4 << 4) >> 4; // 以M*K矩阵为例,源矩阵M轴方向搬运长度(S1向上对齐分形(512B),16*16个f16->向上对齐16)，单位为16 element,取值范围：mStep属于[0,255] + loadData2DParamsB.kStep = GetBlockNum(nSplitSize); // 以M*K矩阵为例,源矩阵K轴方向搬运长度(qkD个f16)，单位为32B,取值范围：nStep属于[0,255] + loadData2DParamsB.srcStride = ((kSplitSize + 31) >> 5 << 5) >> 4; + loadData2DParamsB.dstStride = (nSplitSize + 15) >> 4; + } + } else { + if constexpr (IsSameType::value || IsSameType::value) { + loadData2DParamsB.mStep = ((mmParam.singleN + 15) >> 4 << 4) / 16; + loadData2DParamsB.kStep = kSplitSize / 64; + loadData2DParamsB.srcStride = loadData2DParamsB.mStep; + loadData2DParamsB.dstStride = loadData2DParamsB.mStep; + } else if constexpr (IsSameType::value) { + loadData2DParamsB.mStep = ((nSplitSize + 15) >> 4 << 4) >> 4; // 以M*K矩阵为例,源矩阵M轴方向搬运长度(S1向上对齐分形(512B),16*16个f16->向上对齐16)，单位为16 element,取值范围：mStep属于[0,255] + loadData2DParamsB.kStep = GetBlockNum(kSplitSize); // 以M*K矩阵为例,源矩阵K轴方向搬运长度(qkD个f16)，单位为32B,取值范围：nStep属于[0,255] + loadData2DParamsB.srcStride = ((nSplitSize + 31) >> 5 << 5) >> 4; + loadData2DParamsB.dstStride = loadData2DParamsB.mStep; + } + } + + LoadData2DMxParams loadData2DMxParamsB; + if constexpr (IsSameType::value) { + loadData2DMxParamsB.xStartPosition = 0; + loadData2DMxParamsB.yStartPosition = 0; + loadData2DMxParamsB.xStep = ((nSplitSize + 15) >> 4 << 4) >> 4; + loadData2DMxParamsB.yStep = (kSplitSize + 63) >> 5 >> 1; + loadData2DMxParamsB.srcStride = loadData2DMxParamsB.yStep; + loadData2DMxParamsB.dstStride = loadData2DMxParamsB.yStep; + if (loadData2DParamsB.ifTranspose && (loadData2DParamsB.dstStride & 1)) { + uint32_t l0bLoop = (loadData2DParamsB.mStep + 1) >> 1; + loadData2DParamsB.mStep = M_STEP_ALIGN_BASE; + loadData2DMxParamsB.xStep = loadData2DParamsB.mStep; + uint64_t dstOffset = 0; + uint64_t dstAddrStride = (nSplitSize + 15) / 16 * 16 * 32; + uint16_t oriMStep = loadData2DParamsB.mStartPosition; + uint16_t oriMScaleStep = loadData2DMxParamsB.xStartPosition; + for (uint32_t idx = 0; idx < l0bLoop; ++idx) { + loadData2DParamsB.mStartPosition = oriMStep + M_STEP_ALIGN_BASE * idx; + loadData2DMxParamsB.xStartPosition = oriMScaleStep + M_STEP_ALIGN_BASE * idx; + LoadData(bL0Tensor[dstOffset], bL1Tensor[L1Boffset], bScaleL1Tensor[L1Boffset >> 5], loadData2DParamsB, + loadData2DMxParamsB); + dstOffset += dstAddrStride; + } + } else { + LoadData(bL0Tensor, bL1Tensor[L1Boffset], bScaleL1Tensor[L1Boffset >> 5], loadData2DParamsB, + loadData2DMxParamsB); + } + } else if constexpr (IsSameType::value) { + if (loadData2DParamsB.ifTranspose) { + loadData2DMxParamsB.xStartPosition= 0; // 矩阵X轴方向的起始位置，即M维度方向，单位为1个分形（1个单位代表一个32B的分形）。 + loadData2DMxParamsB.yStartPosition = mmParam.singleN / MX_FP4_PTG_PCG_SCALE_PARAM / 2; // 源矩阵Y轴方向的起始位置，即K维度方向，单位为32B。 + loadData2DMxParamsB.xStep = (kSplitSize + 15) / 16; // 源矩阵X轴方向搬运长度，即M维度方向，单位为1个分形（1个单位代表一个32B的分形）。取值范围：xStep∈[0, 255]。 + loadData2DMxParamsB.yStep = (mmParam.singleK + 63) / MX_FP4_PTG_PCG_SCALE_PARAM / 2; // 源矩阵Y轴方向搬运长度，即K维度方向，单位为32B。取值范围：yStep∈[0, 255]。 + // TODO: mmParam.singleM is ??? + loadData2DMxParamsB.srcStride = mmParam.singleM / MX_FP4_PTG_PCG_SCALE_PARAM / 2; // 源矩阵X方向前一个分形起始地址与后一个分形起始地址的间隔，单位为32B。 + loadData2DMxParamsB.dstStride = loadData2DMxParamsB.yStep; // 目标矩阵X方向前一个分形起始地址与后一个分形起始地址的间隔，单位为32B。 + LoadData(bL0Tensor, bL1Tensor, bScaleL1Tensor, loadData2DParamsB,loadData2DMxParamsB); + } else { + loadData2DMxParamsB.xStartPosition = mmParam.singleK / 16; // 矩阵X轴方向的起始位置，即M维度方向，单位为1个分形（1个单位代表一个32B的分形）。 + loadData2DMxParamsB.yStartPosition = 0; + loadData2DMxParamsB.xStep = (mmParam.singleN + 15) / 16; // 源矩阵X轴方向搬运长度，即M维度方向，单位为1个分形（1个单位代表一个32B的分形）。取值范围：xStep∈[0, 255]。 + loadData2DMxParamsB.yStep = (kSplitSize + 63) / MX_FP4_PTG_PCG_SCALE_PARAM / 2; // 源矩阵Y轴方向搬运长度，即K维度方向，单位为32B。取值范围：yStep∈[0, 255]。 + loadData2DMxParamsB.srcStride = loadData2DMxParamsB.yStep; + loadData2DMxParamsB.dstStride = loadData2DMxParamsB.yStep; // 目标矩阵X方向前一个分形起始地址与后一个分形起始地址的间隔，单位为32B。 + LoadData(bL0Tensor, bL1Tensor,bScaleL1Tensor, loadData2DParamsB,loadData2DMxParamsB); + } + } else if constexpr (IsSameType::value) { + if (loadData2DParamsB.ifTranspose) { + loadData2DMxParamsB.xStartPosition= 0; // 矩阵X轴方向的起始位置，即M维度方向，单位为1个分形（1个单位代表一个32B的分形）。 + loadData2DMxParamsB.yStartPosition = mmParam.singleN * 2 / HI_FP4_PTG_PCG_SCALE_PARAM; // 源矩阵Y轴方向的起始位置，即K维度方向，单位为32B。 + loadData2DMxParamsB.xStep = (kSplitSize + 15) / 16; // 源矩阵X轴方向搬运长度，即M维度方向，单位为1个分形（1个单位代表一个32B的分形）。取值范围：xStep∈[0, 255]。 + loadData2DMxParamsB.yStep = (mmParam.singleK + 31) * 2 / HI_FP4_PTG_PCG_SCALE_PARAM; // 1个fp32伪装为2个bf16，故最小分型数量翻倍。源矩阵Y轴方向搬运长度，即K维度方向，单位为32B。取值范围：yStep∈[0, 255]。 + // TODO: mmParam.singleM is ??? + loadData2DMxParamsB.srcStride = mmParam.singleM * 2 / HI_FP4_PTG_PCG_SCALE_PARAM; // 源矩阵X方向前一个分形起始地址与后一个分形起始地址的间隔，单位为32B。 + loadData2DMxParamsB.dstStride = loadData2DMxParamsB.yStep; // 目标矩阵X方向前一个分形起始地址与后一个分形起始地址的间隔，单位为32B。 + LoadData(bL0Tensor, bL1Tensor, bScaleL1Tensor, loadData2DParamsB,loadData2DMxParamsB); + } else { + loadData2DMxParamsB.xStartPosition= mmParam.singleK / 16; // 矩阵X轴方向的起始位置，即M维度方向，单位为1个分形（1个单位代表一个32B的分形）。 + loadData2DMxParamsB.yStartPosition = 0; // 源矩阵Y轴方向的起始位置，即K维度方向，单位为32B。 + loadData2DMxParamsB.xStep = (mmParam.singleN + 15) / 16; // 源矩阵X轴方向搬运长度，即M维度方向，单位为1个分形（1个单位代表一个32B的分形）。取值范围：xStep∈[0, 255]。 + loadData2DMxParamsB.yStep = (kSplitSize + 31) * 2 / HI_FP4_PTG_PCG_SCALE_PARAM; // 1个fp32伪装为2个bf16，故最小分型数量翻倍。源矩阵Y轴方向搬运长度，即K维度方向，单位为32B。取值范围：yStep∈[0, 255]。 + loadData2DMxParamsB.srcStride = loadData2DMxParamsB.yStep; // 源矩阵X方向前一个分形起始地址与后一个分形起始地址的间隔，单位为32B。 + loadData2DMxParamsB.dstStride = loadData2DMxParamsB.yStep; // 目标矩阵X方向前一个分形起始地址与后一个分形起始地址的间隔，单位为32B。 + LoadData(bL0Tensor, bL1Tensor,bScaleL1Tensor, loadData2DParamsB,loadData2DMxParamsB); + } + } +} + +// 全载 +// 外部L1切入K时，需要传入cmatrixInitVal的标记 +template +__aicore__ inline void MatmulFullMX(const LocalTensor &aL1Tensor, + const LocalTensor &bL1Tensor, + const LocalTensor &aScaleL1Tensor, + const LocalTensor &bScaleL1Tensor, + L0AType &aL0BuffsDb, + L0BType &bL0BuffsDb, + const LocalTensor &cL0Tensor, + const MMParam ¶m) +{ + Buffer l0aBuffer = aL0BuffsDb.Get(); + l0aBuffer.Wait(); + LocalTensor L0ATensor = l0aBuffer.GetTensor(); + if constexpr (IsSameType::value) { + LoadDataToL0AMx(L0ATensor, aL1Tensor, aScaleL1Tensor, param, 0, param.singleK, param.singleM); + } else if constexpr (IsSameType::value) { + LoadDataToL0A(L0ATensor, aL1Tensor, param, 0, param.singleK, param.singleM); + } + l0aBuffer.Set(); + + Buffer l0bBuffer = bL0BuffsDb.Get(); + l0bBuffer.Wait(); + LocalTensor L0BTensor = l0bBuffer.GetTensor(); + if constexpr (IsSameType::value) { + LoadDataToL0BMx(L0BTensor, bL1Tensor, bScaleL1Tensor, param, 0, param.singleK, param.singleN); + } else if constexpr (IsSameType::value) { + LoadDataToL0B(L0BTensor, bL1Tensor, param, 0, param.singleK, param.singleN); + } + l0bBuffer.Set(); + + l0aBuffer.Wait(); + l0bBuffer.Wait(); + + MmadParams mmadParams; + mmadParams.m = param.singleM; + if (param.realM != 0) { + mmadParams.m = param.realM; + } + mmadParams.n = param.singleN; + mmadParams.k = param.singleK; + mmadParams.cmatrixInitVal = param.isOutKFisrt; + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = param.unitFlag; + if (mmadParams.m == 1) { + mmadParams.m = 16; + } + + Mmad(cL0Tensor, L0ATensor, L0BTensor, mmadParams); + + l0aBuffer.Set(); + l0bBuffer.Set(); +} + +// 切K mx +template +__aicore__ inline void MatmulKMx(const LocalTensor &aL1Tensor, + const LocalTensor &bL1Tensor, + const LocalTensor &aScaleL1Tensor, + const LocalTensor &bScaleL1Tensor, + L0AType &aL0BuffsDb, + L0BType &bL0BuffsDb, + const LocalTensor &cL0Tensor, + const MMParam ¶m) +{ + uint32_t kLoops = (param.singleK + baseK - 1) / baseK; + uint32_t tailSize = param.singleK % baseK; + uint32_t tailK = tailSize ? tailSize : baseK; + uint64_t L1Aoffset = param.isLeftTranspose ? baseK << 4 : ((param.singleM + 15) >> 4 << 4) * baseK; + uint64_t L1Boffset = param.isRightTranspose ? ((param.singleN + 15) >> 4 << 4) * baseK : baseK << 4; +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) + if constexpr (IsSameType::value || IsSameType::value || IsSameType::value) { + L1Aoffset = ((param.singleM + 31) >> 5 << 5) * baseK; + L1Boffset = ((param.singleN + 31) >> 5 << 5) * baseK; + } + if constexpr (IsSameType::value) { + L1Aoffset = param.isLeftTranspose ? baseK << 3 : ((param.singleM + 15) >> 4 << 4) * baseK; + L1Boffset = param.isRightTranspose ? ((param.singleN + 15) >> 4 << 4) * baseK : baseK << 3; + } +#endif + for (uint32_t k = 0; k < kLoops; k++) { + uint32_t tileK = (k == (kLoops - 1)) ? tailK : baseK; + Buffer l0aBuffer = aL0BuffsDb.Get(); + l0aBuffer.Wait(); // mte1等Matmul：上一轮matmul完成后才能搬运新数据到L0A + LocalTensor L0ATensor = l0aBuffer.GetTensor(); + if constexpr (IsSameType::value) { + LoadDataToL0AMx(L0ATensor, aL1Tensor, aScaleL1Tensor, param, k * L1Aoffset, tileK, param.singleM); // s2, + } else if constexpr (IsSameType::value) { + LoadDataToL0A(L0ATensor, aL1Tensor, param, k * L1Aoffset, tileK, param.singleM); + } + l0aBuffer.Set(); // mte1搬运完后，通知可以开始matmul + + Buffer l0bBuffer = bL0BuffsDb.Get(); + l0bBuffer.Wait(); // mte1等Matmul：上一轮matmul完成后才能搬运新数据到L0B + LocalTensor L0BTensor = l0bBuffer.GetTensor(); + uint64_t loopNum = param.isRightTranspose ? 1 : kLoops; + if constexpr (IsSameType::value) { + LoadDataToL0BMx(L0BTensor, bL1Tensor, bScaleL1Tensor, param, k * L1Boffset, tileK, param.singleN, loopNum); // tileK.D + } else if constexpr (IsSameType::value) { + LoadDataToL0B(L0BTensor, bL1Tensor, param, k * L1Boffset, tileK, param.singleN, loopNum); + } + l0bBuffer.Set(); // mte1搬运完后，通知可以开始matmul + + l0aBuffer.Wait(); // matmul等mte1：L0A数据搬运完成后才能开始matmul + l0bBuffer.Wait(); // matmul等mte1：L0B数据搬运完成后才能开始matmul + + MmadParams mmadParams; + mmadParams.m = param.singleM; + if (param.realM != 0) { + mmadParams.m = param.realM; + } + mmadParams.n = param.singleN; + mmadParams.k = tileK; + if (mmadParams.m == 1) { // m等于1或默认开GEMV模式，文档上没有写怎么关闭GEMV，所以规避当做矩阵运算 + mmadParams.m = 16; + } + mmadParams.cmatrixInitVal = param.isOutKFisrt && (k == 0); + mmadParams.cmatrixSource = false; + if (param.unitFlag != 0) { + mmadParams.unitFlag = (param.unitFlag == UNITFLAG_EN_OUTER_LAST) && (k == kLoops - 1) ? + UNITFLAG_EN_OUTER_LAST : UNITFLAG_ENABLE; + } + + Mmad(cL0Tensor, L0ATensor, L0BTensor, mmadParams); + + l0aBuffer.Set(); // matmul完成后，通知mte1可以开始搬运新数据到L0A + l0bBuffer.Set(); // matmul完成后，通知mte1可以开始搬运新数据到L0B + } +} + +template +__aicore__ inline void MatmulMMx(const LocalTensor &aL1Tensor, + const LocalTensor &bL1Tensor, + const LocalTensor &aScaleL1Tensor, + const LocalTensor &bScaleL1Tensor, + L0AType &aL0BuffsDb, + L0BType &bL0BuffsDb, + const LocalTensor &cL0Tensor, + const MMParam ¶m) +{ + uint32_t mLoops = (param.singleM + baseM - 1) / baseM; + uint32_t tailSize = param.singleM % baseM; + uint32_t tailM = tailSize ? tailSize : baseM; + uint64_t L1Aoffset = param.isLeftTranspose ? baseM << 4 : ((param.singleK + 15) >> 4 << 4) * baseM; // 要对齐 +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) + if constexpr (IsSameType::value || IsSameType::value || IsSameType::value) { + L1Aoffset = ((param.singleK + 31) >> 5 << 5) * baseM; + } +#endif + + uint64_t L0Coffset = ((param.singleN + 31) >> 5 << 5) * baseM; + Buffer l0bBuffer = bL0BuffsDb.Get(); + l0bBuffer.Wait(); // mte1等Matmul：上一轮matmul完成后才能搬运新数据到L0A + LocalTensor L0BTensor = l0bBuffer.GetTensor(); + if constexpr (IsSameType::value) { + LoadDataToL0BMx(L0BTensor, bL1Tensor, bScaleL1Tensor, param, 0, param.singleK, param.singleN); + } else if constexpr (IsSameType::value) { + LoadDataToL0B(L0BTensor, bL1Tensor, param, 0, param.singleK, param.singleN); + } + l0bBuffer.Set(); // mte1搬运完后，通知可以matmul + l0bBuffer.Wait(); // matmul等mte1：L0A数据搬运完成后才能开始matmul + + for (uint32_t m = 0; m < mLoops; m++) { + uint32_t tileM = (m == (mLoops - 1)) ? tailM : baseM; + Buffer l0aBuffer = aL0BuffsDb.Get(); + l0aBuffer.Wait(); // mte1等Matmul：上一轮matmul完成后才能搬运新数据到L0A + LocalTensor L0ATensor = l0aBuffer.GetTensor(); + uint64_t loopNum = param.isRightTranspose ? mLoops : 1; + if constexpr (IsSameType::value) { + LoadDataToL0AMx(L0ATensor, aL1Tensor, aScaleL1Tensor, param, m * L1Aoffset, tileM, param.singleK); + } + l0aBuffer.Set(); // mte1搬运完后，通知可以开始matmul + l0aBuffer.Wait(); // matmul等mte1：L0A数据搬运完成后才能开始matmul + + MmadParams mmadParams; + mmadParams.m = tileM; + mmadParams.n = param.singleN; + mmadParams.k = param.singleK; + if (mmadParams.m == 1) { // m等于1或默认开GEMV模式，文档上没有写怎么关闭GEMV，所以规避当做矩阵运算 + mmadParams.m = 16; + } + mmadParams.cmatrixInitVal = param.isOutKFisrt && (m == 0); + mmadParams.cmatrixSource = false; + if (param.unitFlag != 0) { + mmadParams.unitFlag = (param.unitFlag == UNITFLAG_EN_OUTER_LAST) && (m == mLoops - 1) ? + UNITFLAG_EN_OUTER_LAST : UNITFLAG_ENABLE; + } + Mmad(cL0Tensor[m * L0Coffset], L0ATensor, L0BTensor, mmadParams); + l0aBuffer.Set(); // matmul完成后，通知mte1可以开始搬运新数据到L0B + } + l0bBuffer.Set(); // matmul完成后，通知mte1可以开始搬运新数据到L0A +} +#else +static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding; +template +__aicore__ inline void LoadDataToL0A(LocalTensor& aL0Tensor, const LocalTensor& aL1Tensor, + const MMParam& mmParam, uint64_t L1Aoffset, uint32_t kSplitSize, + uint32_t mSplitSize) +{ + if constexpr (AL == ABLayout::MK) { + LoadData3DParamsV2 loadData3DParams; + loadData3DParams.l1H = mSplitSize / LOAD3D_L1W_SIZE; // 源操作数height + loadData3DParams.l1W = LOAD3D_L1W_SIZE; // 源操作数weight + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + loadData3DParams.mExtension = mSplitSize; // 在目的操作数height维度的传输长度 + loadData3DParams.kExtension = kSplitSize; // 在目的操作数width维度的传输长度 + loadData3DParams.mStartPt = 0; // 卷积核在目的操作数width维度的起点 + loadData3DParams.kStartPt = 0; // 卷积核在目的操作数height维度的起点 + loadData3DParams.strideW = 1; // 卷积核在源操作数width维度滑动的步长 + loadData3DParams.strideH = 1; // 卷积核在源操作数height维度滑动的步长 + loadData3DParams.filterW = 1; // 卷积核width + loadData3DParams.filterSizeW = false; // 是否在filterW的基础上将卷积核width增加256个元素 + loadData3DParams.filterH = 1; // 卷积核height + loadData3DParams.filterSizeH = false; // 是否在filterH的基础上将卷积核height增加256个元素 + loadData3DParams.dilationFilterW = 1; // 卷积核width膨胀系数 + loadData3DParams.dilationFilterH = 1; // 卷积核height膨胀系数 + loadData3DParams.enTranspose = 0; // 是否启用转置功能，对整个目标矩阵进行转置 + loadData3DParams.fMatrixCtrl = 0; + loadData3DParams.channelSize = kSplitSize; // 源操作数的通道数。膨胀系数为1时，目的weight为filterW*filterH*channelSize + LoadData(aL0Tensor, aL1Tensor[L1Aoffset], loadData3DParams); + } else if constexpr (AL == ABLayout::KM) { + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; // 分型矩阵ID，表明搬运起始位置为源操作数中第0个分型 + loadData2DParams.repeatTimes = (kSplitSize / ONE_FRACTAL_H_ELEMENT) * (mmParam.singleM / + (ONE_FRACTAL_W_BYTE / sizeof(T))); // 迭代次数，每个迭代可以处理512B数据 + loadData2DParams.srcStride = 1; // 相邻迭代间，源操作数前一个分型和后一个分型起始地址的间隔（单位512B） + loadData2DParams.dstGap = 0; // 相邻迭代间，目的操作数前一个分型的结束地址和后一个分型起始地址的间隔（单位512B） + loadData2DParams.ifTranspose = true; + LoadData(aL0Tensor, aL1Tensor[L1Aoffset], loadData2DParams); + } +} + +// L1→L0B + 切K/切N/全载 +template +__aicore__ inline void LoadDataToL0B(LocalTensor& bL0Tensor, const LocalTensor& bL1Tensor, + const MMParam& mmParam, uint64_t L1Boffset, uint32_t kSplitSize, + uint32_t nSplitSize) +{ + if constexpr (BL == ABLayout::KN) { + LoadData3DParamsV2 loadData3DParams; + loadData3DParams.l1H = kSplitSize / LOAD3D_L1W_SIZE; // 源操作数height + loadData3DParams.l1W = LOAD3D_L1W_SIZE; // 源操作数weight=16，目的height=l1H*L1W + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + loadData3DParams.mExtension = kSplitSize; // 在目的操作数height维度的传输长度 + loadData3DParams.kExtension = nSplitSize; // 在目的操作数width维度的传输长度 + loadData3DParams.mStartPt = 0; // 卷积核在目的操作数width维度的起点 + loadData3DParams.kStartPt = 0; // 卷积核在目的操作数height维度的起点 + loadData3DParams.strideW = LOAD3D_STRIDE_W; + loadData3DParams.strideH = LOAD3D_STRIDE_H; + loadData3DParams.filterW = LOAD3D_FILTER_W; + loadData3DParams.filterSizeW = false; // 是否在filterW的基础上将卷积核width增加256个元素 + loadData3DParams.filterH = LOAD3D_FILTER_H; + loadData3DParams.filterSizeH = false; // 是否在filterH的基础上将卷积核height增加256个元素 + loadData3DParams.dilationFilterW = LOAD3D_DILA_FILTER_W; // 卷积核width膨胀系数 + loadData3DParams.dilationFilterH = LOAD3D_DILA_FILTER_H; // 卷积核height膨胀系数 + loadData3DParams.enTranspose = 1; // 是否启用转置功能 + loadData3DParams.fMatrixCtrl = 0; // 使用FMATRIX_LEFT还是使用FMATRIX_RIGHT，=0使用FMATRIX_LEFT，=1使用FMATRIX_RIGHT 1 + loadData3DParams.channelSize = nSplitSize; // 源操作数的通道数。膨胀系数为1时，目的weight为filterW*filterH*channelSize + LoadData(bL0Tensor, bL1Tensor[L1Boffset], loadData3DParams); + } else if constexpr (BL == ABLayout::NK) { + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = (nSplitSize + (ONE_FRACTAL_H_ELEMENT - 1)) / ONE_FRACTAL_H_ELEMENT * + (kSplitSize / (ONE_FRACTAL_W_BYTE / sizeof(T))); // 迭代次数，每个迭代可以处理512B数据 + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = false; + LoadData(bL0Tensor, bL1Tensor[L1Boffset], loadData2DParams); + } +} + +#endif +// 全载 +// 外部L1切入K时，需要传入cmatrixInitVal的标记 +template +__aicore__ inline void MatmulFull(const LocalTensor &aL1Tensor, + const LocalTensor &bL1Tensor, + L0AType &aL0BuffsDb, + L0BType &bL0BuffsDb, + const LocalTensor &cL0Tensor, + struct MMParam ¶m, + const LocalTensor &aScaleL1Tensor = LocalTensor(), + const LocalTensor &bScaleL1Tensor = LocalTensor()) +{ + Buffer l0aBuffer = aL0BuffsDb.Get(); + l0aBuffer.Wait(); + LocalTensor L0ATensor = l0aBuffer.GetTensor(); +#if ((__CCE_AICORE__ == 310) || (defined __DAV_310R6__) || (__NPU_ARCH__ == 5102)) + if constexpr (IsSameType::value) { + LoadDataToL0AMx(L0ATensor, aL1Tensor, aScaleL1Tensor, param, 0, param.singleK, param.singleM); // d,s2 + } else +#endif + { + LoadDataToL0A(L0ATensor, aL1Tensor, param, 0, param.singleK, param.singleM); // s2*d,d,s2 + } + l0aBuffer.Set(); + + Buffer l0bBuffer = bL0BuffsDb.Get(); + l0bBuffer.Wait(); + LocalTensor L0BTensor = l0bBuffer.GetTensor(); +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) + if constexpr (IsSameType::value) { + LoadDataToL0BMx(L0BTensor, bL1Tensor, bScaleL1Tensor, param, 0, param.singleK, param.singleN); + } else +#endif + { + LoadDataToL0B(L0BTensor, bL1Tensor, param, 0, param.singleK, param.singleN); + } + l0bBuffer.Set(); + + l0aBuffer.Wait(); + l0bBuffer.Wait(); + + MmadParams mmadParams; + mmadParams.m = param.singleM; + if (param.realM != 0) { + mmadParams.m = param.realM; + } + mmadParams.n = param.singleN; + mmadParams.k = param.singleK; + mmadParams.cmatrixInitVal = param.isOutKFisrt; + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = param.unitFlag; + if (mmadParams.m == 1) { + mmadParams.m = 16; + } + + Mmad(cL0Tensor, L0ATensor, L0BTensor, mmadParams); + + l0aBuffer.Set(); + l0bBuffer.Set(); +} + +// 切K +template +__aicore__ inline void MatmulK(const LocalTensor &aL1Tensor, + const LocalTensor &bL1Tensor, + L0AType &aL0BuffsDb, + L0BType &bL0BuffsDb, + const LocalTensor &cL0Tensor, + const MMParam ¶m, + const LocalTensor &aScaleL1Tensor = LocalTensor(), + const LocalTensor &bScaleL1Tensor = LocalTensor()) +{ + uint32_t kLoops = (param.singleK + baseK - 1) / baseK; + uint32_t tailSize = param.singleK % baseK; + uint32_t tailK = tailSize ? tailSize : baseK; + uint64_t L1Aoffset = param.isLeftTranspose ? baseK << 4 : ((param.singleM + 15) >> 4 << 4) * baseK; + uint64_t L1Boffset = param.isRightTranspose ? ((param.singleN + 15) >> 4 << 4) * baseK : baseK << 4; +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) + if constexpr (IsSameType::value || IsSameType::value || + IsSameType::value || IsSameType::value) { + L1Aoffset = ((param.singleM + 31) >> 5 << 5) * baseK; + L1Boffset = ((param.singleN + 31) >> 5 << 5) * baseK; + } + if constexpr (IsSameType::value) { + L1Aoffset = param.isLeftTranspose ? baseK << 3 : ((param.singleM + 15) >> 4 << 4) * baseK; + L1Boffset = param.isRightTranspose ? ((param.singleN + 15) >> 4 << 4) * baseK : baseK << 3; + } +#endif + + for (uint32_t k = 0; k < kLoops; k++) { + uint32_t tileK = (k == (kLoops - 1)) ? tailK : baseK; + Buffer l0aBuffer = aL0BuffsDb.Get(); + l0aBuffer.Wait(); // mte1等Matmul：上一轮matmul完成后才能搬运新数据到L0A + LocalTensor L0ATensor = l0aBuffer.GetTensor(); +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) + if constexpr (IsSameType::value) { + LoadDataToL0AMx(L0ATensor, aL1Tensor, aScaleL1Tensor, param, k * L1Aoffset, tileK, param.singleM); // s2, + } else +#endif + { + LoadDataToL0A(L0ATensor, aL1Tensor, param, k * L1Aoffset, tileK, param.singleM); // s2*d,d,s2 + } + + Buffer l0bBuffer = bL0BuffsDb.Get(); + l0bBuffer.Wait(); // mte1等Matmul：上一轮matmul完成后才能搬运新数据到L0B + LocalTensor L0BTensor = l0bBuffer.GetTensor(); + uint64_t loopNum = param.isRightTranspose ? 1 : kLoops; +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) + if constexpr (IsSameType::value) { + LoadDataToL0BMx(L0BTensor, bL1Tensor, bScaleL1Tensor, param, k * L1Boffset, tileK, param.singleN, loopNum); // tileK.D + } else +#endif + { + LoadDataToL0B(L0BTensor, bL1Tensor, param, k * L1Boffset, tileK, param.singleN, loopNum); + } + l0bBuffer.Set(); // mte1搬运完后，通知可以开始matmul + // l0aBuffer和l0bBuffer共用MTE1_M，在D=512场景减少同步指令数量，提升性能 + l0bBuffer.Wait(); // matmul等mte1：L0B数据搬运完成后才能开始matmul + + MmadParams mmadParams; + mmadParams.m = param.singleM; + if (param.realM != 0) { + mmadParams.m = param.realM; + } + mmadParams.n = param.singleN; + mmadParams.k = tileK; + if (mmadParams.m == 1) { // m等于1或默认开GEMV模式，文档上没有写怎么关闭GEMV，所以规避当做矩阵运算 + mmadParams.m = 16; + } + mmadParams.cmatrixInitVal = param.isOutKFisrt && (k == 0); + mmadParams.cmatrixSource = false; + if (param.unitFlag != 0) { + mmadParams.unitFlag = (param.unitFlag == UNITFLAG_EN_OUTER_LAST) && (k == kLoops - 1) ? + UNITFLAG_EN_OUTER_LAST : UNITFLAG_ENABLE; + } + Mmad(cL0Tensor, L0ATensor, L0BTensor, mmadParams); + + l0aBuffer.Set(); // matmul完成后，通知mte1可以开始搬运新数据到L0A + l0bBuffer.Set(); // matmul完成后，通知mte1可以开始搬运新数据到L0B + } +} + +// 切K---int8带偏置的实现 +template +__aicore__ inline void MatmulKbias(const LocalTensor &aL1Tensor, + const LocalTensor &bL1Tensor, + L0AType &aL0BuffsDb, + L0BType &bL0BuffsDb, + const LocalTensor &cL0Tensor, + const LocalTensor &biasTensor, + const MMParam ¶m) +{ + uint32_t kLoops = (param.singleK + baseK - 1) / baseK; // 尾块处理，如果dSize不可以被256整除的时候，也就是存在尾块的时候，循环次数需要+1，这里的尾块始终是针对一个基本块计算的视角来看的 + uint32_t tailSize = param.singleK % baseK; // 针对dsize维度来说，需要按照256切分，如果不可以整除的话，求出最后一个尾块的dsize，命名为tailSize + uint32_t tailK = tailSize ? tailSize : baseK; // 沿着K轴切分的尾块的K + uint64_t L1Aoffset = param.isLeftTranspose ? baseK << 4 : ((param.singleM + 15) >> 4 << 4) * baseK; // 给传入的s1realsize对齐到16的倍数 + uint64_t L1Boffset = param.isRightTranspose ? ((param.singleN + 15) >> 4 << 4) * baseK : baseK << 4; // 给传入的s2realsize对齐到16的倍数 +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) + if constexpr (IsSameType::value || IsSameType::value || IsSameType::value || IsSameType::value) { + L1Aoffset = ((param.singleM + 31) >> 5 << 5) * baseK; // 给传入的s1realsize对齐到32的倍数 + L1Boffset = ((param.singleN + 31) >> 5 << 5) * baseK; // 给传入的s2realsize对齐到32的倍数 + } + if constexpr (IsSameType::value) { + L1Aoffset = param.isLeftTranspose ? baseK << 3 : ((param.singleM + 15) >> 4 << 4) * baseK; + L1Boffset = param.isRightTranspose ? ((param.singleN + 15) >> 4 << 4) * baseK : baseK << 3; + } +#endif + + for (uint32_t k = 0; k < kLoops; k++) { + uint32_t tileK = (k == (kLoops - 1)) ? tailK : baseK; + Buffer l0aBuffer = aL0BuffsDb.Get(); + l0aBuffer.Wait(); // mte1等Matmul：上一轮matmul完成后才能搬运新数据到L0A + LocalTensor L0ATensor = l0aBuffer.GetTensor(); + LoadDataToL0A(L0ATensor, aL1Tensor, param, k * L1Aoffset, tileK, param.singleM); + l0aBuffer.Set(); // mte1搬运完后，通知可以开始matmul + + Buffer l0bBuffer = bL0BuffsDb.Get(); + l0bBuffer.Wait(); // mte1等Matmul：上一轮matmul完成后才能搬运新数据到L0B + LocalTensor L0BTensor = l0bBuffer.GetTensor(); + uint64_t loopNum = param.isRightTranspose ? 1 : kLoops; + LoadDataToL0B(L0BTensor, bL1Tensor, param, k * L1Boffset, tileK, param.singleN, loopNum); + l0bBuffer.Set(); // mte1搬运完后，通知可以开始matmul + + l0aBuffer.Wait(); // matmul等mte1：L0A数据搬运完成后才能开始matmul + l0bBuffer.Wait(); // matmul等mte1：L0B数据搬运完成后才能开始matmul + + MmadParams mmadParams; + mmadParams.m = param.singleM; + if (param.realM != 0) { + mmadParams.m = param.realM; + } + mmadParams.n = param.singleN; + mmadParams.k = tileK; + if (mmadParams.m == 1) { // m等于1或默认开GEMV模式，文档上没有写怎么关闭GEMV，所以规避当做矩阵运算 + mmadParams.m = 16; + } + mmadParams.cmatrixInitVal = false; + mmadParams.cmatrixSource = (k == 0); + if (param.unitFlag != 0) { + mmadParams.unitFlag = (param.unitFlag == UNITFLAG_EN_OUTER_LAST) && (k == kLoops - 1) ? + UNITFLAG_EN_OUTER_LAST : UNITFLAG_ENABLE; + } + + if (k == 0) { + Mmad(cL0Tensor, L0ATensor, L0BTensor, biasTensor, mmadParams); + } else { + Mmad(cL0Tensor, L0ATensor, L0BTensor, mmadParams); + } + + l0aBuffer.Set(); // matmul完成后，通知mte1可以开始搬运新数据到L0A + l0bBuffer.Set(); // matmul完成后，通知mte1可以开始搬运新数据到L0B + } +} + +// 切N +template +__aicore__ inline void MatmulN(const LocalTensor &aL1Tensor, + const LocalTensor &bL1Tensor, + L0AType &aL0BuffsDb, + L0BType &bL0BuffsDb, + const LocalTensor &cL0Tensor, + const MMParam ¶m, + const LocalTensor &aScaleL1Tensor = LocalTensor(), + const LocalTensor &bScaleL1Tensor = LocalTensor()) +{ + uint32_t nLoops = (param.singleN + baseN - 1) / baseN; // 尾块处理 + uint32_t tailSize = param.singleN % baseN; + uint32_t tailN = tailSize ? tailSize : baseN; + uint64_t L1Boffset = param.isRightTranspose ? (baseN << 4) : ((param.singleK + 15) >> 4 << 4) * baseN; +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) + if constexpr (IsSameType::value || IsSameType::value || + IsSameType::value || IsSameType::value) { + L1Boffset = ((param.singleK + 31) >> 5 << 5) * baseN; + } +#endif + uint64_t L0Coffset = ((param.singleM + 15) >> 4 << 4) * baseN; + if (param.realM != 0) { + L0Coffset = ((param.realM + 15) >> 4 << 4) * baseN; + } + + Buffer l0aBuffer = aL0BuffsDb.Get(); + l0aBuffer.Wait(); // mte1等Matmul：上一轮matmul完成后才能搬运新数据到L0A + LocalTensor L0ATensor = l0aBuffer.GetTensor(); +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) + if constexpr (IsSameType::value) { + LoadDataToL0AMx(L0ATensor, aL1Tensor, aScaleL1Tensor, param, 0, param.singleK, param.singleM); // d,s2 + } else +#endif + { + LoadDataToL0A(L0ATensor, aL1Tensor, param, 0, param.singleK, param.singleM); // s2*d,d,s2 + } + for (uint32_t n = 0; n < nLoops; n++) { + uint32_t tileN = (n == (nLoops - 1)) ? tailN : baseN; + + Buffer l0bBuffer = bL0BuffsDb.Get(); + l0bBuffer.Wait(); // mte1等Matmul：上一轮matmul完成后才能搬运新数据到L0B + LocalTensor L0BTensor = l0bBuffer.GetTensor(); + uint64_t loopNum = param.isRightTranspose ? nLoops : 1; +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) + if constexpr (IsSameType::value) { + LoadDataToL0BMx(L0BTensor, bL1Tensor, bScaleL1Tensor, param, n * L1Boffset, param.singleK, tileN, loopNum); // tileK.D + } else +#endif + { + LoadDataToL0B(L0BTensor, bL1Tensor, param, n * L1Boffset, param.singleK, tileN, loopNum); + } + l0bBuffer.Set(); // mte1搬运完后，通知可以开始matmul + // l0aBuffer和l0bBuffer共用MTE1_M，在D=512场景减少同步指令数量，提升性能 + l0bBuffer.Wait(); // matmul等mte1：L0B数据搬运完成后才能开始matmul + + MmadParams mmadParams; + mmadParams.m = param.singleM; + if (param.realM != 0) { + mmadParams.m = param.realM; + } + mmadParams.n = tileN; + mmadParams.k = param.singleK; + if (mmadParams.m == 1) { + mmadParams.m = FP16_ONE_FRACTAL_ELEMENT; + } + mmadParams.cmatrixInitVal = param.isOutKFisrt; + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = param.unitFlag; + Mmad(cL0Tensor[n * L0Coffset], L0ATensor, L0BTensor, mmadParams); + + l0bBuffer.Set(); // matmul完成后，通知mte1可以开始搬运新数据到L0B + } + l0aBuffer.Set(); // matmul完成后，通知mte1可以开始搬运新数据到L0A +} + +// 切M +template +__aicore__ inline void MatmulKM(const LocalTensor &aL1Tensor, + const LocalTensor &bL1Tensor, + BuffersPolicyDB &aL0BuffsDb, + BuffersPolicyDB &bL0BuffsDb, + const LocalTensor &cL0Tensor, + struct MMParam ¶m) +{ + uint32_t mLoops = (param.singleM + baseM - 1) / baseM; // 尾块处理 + uint32_t kLoops = (param.singleK + baseK - 1) / baseK; // 尾块处理 + uint32_t mSplitSize = (mLoops == 1) ? param.singleM : baseM; + uint32_t mplitTailSize = (param.singleM % baseM) ? (param.singleM % baseM) : mSplitSize; + uint32_t kSplitSize = (kLoops == 1) ? param.singleK : baseK; + uint32_t kSplitTailSize = (param.singleK % baseK) ? (param.singleK % baseK) : kSplitSize; + uint64_t L1Boffset = kSplitSize * param.singleN; + uint64_t L0Coffset = mSplitSize * param.singleN; + + for (uint32_t k = 0; k < kLoops; k++) { + kSplitSize = (k == (kLoops - 1)) ? kSplitTailSize : kSplitSize; + for (uint32_t m = 0; m < mLoops; m++){ + mSplitSize = (m == (mLoops - 1)) ? mplitTailSize : mSplitSize; + Buffer l0aBuffer = aL0BuffsDb.Get(); + l0aBuffer.Wait(); // 占用 + LocalTensor L0ATensor = l0aBuffer.GetTensor(); + LoadDataToL0A(L0ATensor, aL1Tensor, param, + k * param.singleM * kSplitSize + m * kSplitSize * mSplitSize, + kSplitSize, mSplitSize); + l0aBuffer.Set(); // 通知 + l0aBuffer.Wait(); // 等待L0A + + Buffer l0bBuffer = bL0BuffsDb.Get(); + l0bBuffer.Wait(); // 占用 + LocalTensor L0BTensor = l0bBuffer.GetTensor(); + LoadDataToL0B(L0BTensor, bL1Tensor, param, k * L1Boffset, kSplitSize, param.singleN); + l0bBuffer.Set(); // 通知 + l0bBuffer.Wait(); // 等待L0B + + MmadParams mmadParams; + mmadParams.m = mSplitSize; + mmadParams.n = param.singleN; + mmadParams.k = kSplitSize; + mmadParams.cmatrixInitVal = param.isOutKFisrt && (k == 0); // 配置C矩阵初始值是否为0。默认值ture + mmadParams.cmatrixSource = false; // 来源于CO1 + Mmad(cL0Tensor[m * L0Coffset], L0ATensor, L0BTensor, mmadParams); + + l0aBuffer.Set(); // 释放L0A + l0bBuffer.Set(); // 释放L0B + } + } +} + +template +__aicore__ inline void MatmulBase(const LocalTensor &aL1Tensor, + const LocalTensor &bL1Tensor, + L0AType &aL0BuffsDb, + L0BType &bL0BuffsDb, + const LocalTensor &cL0Tensor, + struct MMParam ¶m) +{ + if ((param.singleK + baseK - 1) / baseK > 1) { + MatmulK(aL1Tensor, bL1Tensor, aL0BuffsDb, bL0BuffsDb, cL0Tensor, param); + } else if ((param.singleN + baseN - 1) / baseN > 1) { + MatmulN(aL1Tensor, bL1Tensor, aL0BuffsDb, bL0BuffsDb, cL0Tensor, param); + } else { + MatmulFull(aL1Tensor, bL1Tensor, aL0BuffsDb, bL0BuffsDb, cL0Tensor, param); + } +} + +template +__aicore__ inline void MatmulKPP(const LocalTensor &aL1Tensor, + const LocalTensor &bL1Tensor, + BuffersPolicyDB &aL0BuffsDb, + BuffersPolicyDB &bL0BuffsDb, + const LocalTensor &cL0Tensor, + const MMParam ¶m) +{ + uint32_t kLoops = (param.singleK + baseK - 1) / baseK; + uint32_t kSplitSize = (kLoops == 1) ? param.singleK : baseK; + uint32_t kSplitSizeAlign = AlignUp(kSplitSize, FP16_ONE_FRACTAL_ELEMENT); + uint64_t L1Aoffset = AlignUp(param.singleM, FP16_ONE_FRACTAL_ELEMENT) * kSplitSize; + uint64_t L1Boffset = AlignUp(param.singleN, FP16_ONE_FRACTAL_ELEMENT) * kSplitSize; + for (uint32_t k = 0; k < kLoops; k++) { + if (k == kLoops - 1) { + kSplitSize = (param.singleK % baseK) ? (param.singleK % baseK) : kSplitSize; + kSplitSizeAlign = AlignUp(kSplitSize, FP16_ONE_FRACTAL_ELEMENT); + } + Buffer l0aBuffer = aL0BuffsDb.Get(); + l0aBuffer.Wait(); // mte1等Matmul:上一轮matmul完成后才能搬运新数据到L0A + LocalTensor L0ATensor = l0aBuffer.GetTensor(); + LoadDataToL0A(L0ATensor, aL1Tensor, param, k * L1Aoffset, kSplitSizeAlign, param.singleM); + Buffer l0bBuffer = bL0BuffsDb.Get(); + LocalTensor L0BTensor = l0bBuffer.GetTensor(); + LoadDataToL0B(L0BTensor, bL1Tensor, param, k * L1Boffset, kSplitSizeAlign, param.singleN); + + l0aBuffer.Set(); // mte1搬运完后，通知可以开始matmul + l0aBuffer.Wait(); // matmul等mte1：L0A数据搬运完成后才能开始matmul + + MmadParams mmadParams; + mmadParams.m = param.singleM; + mmadParams.n = param.singleN; + mmadParams.k = kSplitSize; + if (mmadParams.m == 1) { //m等于1会默认开GEMV模式，且不可关闭GEMV，所以规避当作矩阵计算 + mmadParams.m = FP16_ONE_FRACTAL_ELEMENT; + } + mmadParams.cmatrixInitVal = (param.isOutKFisrt == true) && (k == 0); + mmadParams.cmatrixSource = false; + if (param.unitFlag != 0) { + mmadParams.unitFlag = (param.unitFlag == UNITFLAG_EN_OUTER_LAST) && (k == kLoops - 1) ? + UNITFLAG_EN_OUTER_LAST : UNITFLAG_ENABLE; + } + + Mmad(cL0Tensor, L0ATensor, L0BTensor, mmadParams); + #if (__CCE_AICORE__ != 310) && (!(defined __DAV_310R6__)) + if ((mmadParams.m / FP16_ONE_FRACTAL_ELEMENT) * (mmadParams.n / FP16_ONE_FRACTAL_ELEMENT) < MMAD_MN_SIZE_10) { + AscendC::PipeBarrier(); + } + #endif + l0aBuffer.Set(); // matmul完成后，通知mte1可以开始搬运新数据到L0A + } +} +template +__aicore__ inline void LoadDataToL0A(LocalTensor& aL0Tensor, const LocalTensor& aL1Tensor, + uint32_t rowSize, uint32_t kSplitSize, uint32_t mSplitSize) +{ + uint32_t blockElementCnt = ONE_FRACTAL_W_BYTE / sizeof(T); + if constexpr (IsSameType::value) { + blockElementCnt = INT4_ONE_FRACTAL_ELEMENT; + } + if constexpr (AL == ABLayout::MK) { + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; // 分型矩阵ID，表明搬运起始位置为源操作数中第0个分型 + loadData2DParams.srcStride = 1; // 相邻迭代间，源操作数前一个分型和后一个分型起始地址的间隔（单位512B） + loadData2DParams.dstGap = kSplitSize / blockElementCnt - 1; // 相邻迭代间，目的操作数前一个分型的结束地址和后一个分型起始地址的间隔（单位512B） + loadData2DParams.repeatTimes = mSplitSize / ONE_FRACTAL_H_ELEMENT; // 迭代次数，每个迭代可以处理512B数据 + loadData2DParams.ifTranspose = false; + uint32_t loopTimes = kSplitSize / blockElementCnt; + uint64_t l1Offset = rowSize * blockElementCnt; + uint64_t l0Offset = ONE_FRACTAL_H_ELEMENT * blockElementCnt; + for(uint32_t loop = 0; loop < loopTimes; loop++) { + LoadData(aL0Tensor[loop * l0Offset], aL1Tensor[loop * l1Offset], loadData2DParams); + } + } else if constexpr (AL == ABLayout::KM) { + LoadData2dTransposeParams loadData2dTransposeParams; + loadData2dTransposeParams.startIndex = 0; + loadData2dTransposeParams.srcStride = 1; + loadData2dTransposeParams.dstFracGap = (kSplitSize + blockElementCnt -1) / blockElementCnt; + loadData2dTransposeParams.dstGap = mSplitSize / ONE_FRACTAL_H_ELEMENT - 1; + if(rowSize == kSplitSize) { + loadData2dTransposeParams.repeatTimes = (kSplitSize + blockElementCnt - 1) / blockElementCnt; + uint32_t loopTimes = mSplitSize / blockElementCnt; + uint64_t l1Offset = rowSize * blockElementCnt; + uint64_t l0Offset = kSplitSize * blockElementCnt; + for(uint32_t loop = 0; loop < loopTimes; loop++) { + LoadDataWithTranspose(aL0Tensor[loop * l0Offset], aL1Tensor[loop * l1Offset], loadData2dTransposeParams); + } + } else { + loadData2dTransposeParams.repeatTimes = ((kSplitSize + blockElementCnt - 1) / blockElementCnt) * (mSplitSize / blockElementCnt); + LoadDataWithTranspose(aL0Tensor, aL1Tensor, loadData2dTransposeParams); + } + } +} + +template +__aicore__ inline void LoadDataToL0B(LocalTensor& bL0Tensor, const LocalTensor& bL1Tensor, + uint32_t rowSize, uint32_t kSplitSize, uint32_t nSplitSize) +{ + uint32_t blockElementCnt = ONE_FRACTAL_W_BYTE / sizeof(T); + if constexpr (IsSameType::value) { + blockElementCnt = INT4_ONE_FRACTAL_ELEMENT; + } + if constexpr (BL == ABLayout::KN) { + LoadData2dTransposeParams loadData2dTransposeParams; + + loadData2dTransposeParams.startIndex = 0; + loadData2dTransposeParams.srcStride = 1; + loadData2dTransposeParams.dstFracGap = 0; + loadData2dTransposeParams.dstGap = nSplitSize / ONE_FRACTAL_H_ELEMENT - 1; + loadData2dTransposeParams.repeatTimes = (kSplitSize + blockElementCnt - 1) / blockElementCnt; + + uint32_t loopTimes = nSplitSize / blockElementCnt; + uint64_t l1Offset = rowSize * blockElementCnt; + uint64_t l0Offset = blockElementCnt * blockElementCnt; + for(uint32_t loop = 0; loop < loopTimes; loop++) { + LoadDataWithTranspose(bL0Tensor[loop * l0Offset], bL1Tensor[loop * l1Offset], loadData2dTransposeParams); + } + } else if constexpr (BL == ABLayout::NK) { + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; // 分型矩阵ID，表明搬运起始位置为源操作数中第0个分型 + loadData2DParams.srcStride = 1; // 相邻迭代间，源操作数前一个分型和后一个分型起始地址的间隔（单位512B） + loadData2DParams.dstGap = 0; // 相邻迭代间，目的操作数前一个分型的结束地址和后一个分型起始地址的间隔（单位512B） + loadData2DParams.ifTranspose = false; + if(rowSize == kSplitSize) { + loadData2DParams.repeatTimes = ((nSplitSize + ONE_FRACTAL_H_ELEMENT - 1) / ONE_FRACTAL_H_ELEMENT) * (kSplitSize / blockElementCnt);// 迭代次数，每个迭代可以处理512B数据 + LoadData(bL0Tensor, bL1Tensor, loadData2DParams); + } else { + loadData2DParams.repeatTimes = (nSplitSize + ONE_FRACTAL_H_ELEMENT - 1) / ONE_FRACTAL_H_ELEMENT;// 迭代次数，每个迭代可以处理512B数据 + uint32_t loopTimes = kSplitSize / blockElementCnt; + uint64_t l1Offset = nSplitSize * blockElementCnt; + uint64_t l0Offset = rowSize * blockElementCnt; + for (uint32_t loop = 0; loop < loopTimes; loop++) { + LoadData(bL0Tensor[loop * l0Offset], bL1Tensor[loop * l1Offset], loadData2DParams); + } + } + } +} +} +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/common/op_kernel/memcopy/fa_gm_tensor.h b/csrc/ascend/attention/common/op_kernel/memcopy/fa_gm_tensor.h new file mode 100644 index 000000000..2b13a55ac --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/memcopy/fa_gm_tensor.h @@ -0,0 +1,35 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file fa_gm_tensor.h + * \brief + */ +#ifndef FA_GM_TENSOR_H +#define FA_GM_TENSOR_H + +#if ASC_DEVKIT_MAJOR >= 9 +#include "kernel_vec_intf.h" +#include "kernel_cube_intf.h" +#else +#include "kernel_operator.h" +#endif +#include "gm_layout.h" +#include "offset_calculator_v2.h" + +using AscendC::GlobalTensor; + +template +struct FaGmTensor { + GlobalTensor gmTensor; + OffsetCalculator offsetCalculator; +}; + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/common/op_kernel/memcopy/fa_l1_tensor.h b/csrc/ascend/attention/common/op_kernel/memcopy/fa_l1_tensor.h new file mode 100644 index 000000000..d571f9360 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/memcopy/fa_l1_tensor.h @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file fa_l1_tensor.h + * \brief + */ +#ifndef FA_L1_TENSOR_H +#define FA_L1_TENSOR_H + +#if ASC_DEVKIT_MAJOR >= 9 +#include "kernel_vec_intf.h" +#include "kernel_cube_intf.h" +#else +#include "kernel_operator.h" +#endif + +using AscendC::LocalTensor; + +enum class L1Format { + NZ = 0 +}; + +enum class ScaleTrans { + NO_TRANS = 0, + ND2NZ = 1, + DN2NZ = 2 +}; + +template +struct FaL1Tensor { + LocalTensor tensor; + uint32_t rowCount; +}; + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/common/op_kernel/memcopy/gm_coord.h b/csrc/ascend/attention/common/op_kernel/memcopy/gm_coord.h new file mode 100644 index 000000000..98b4b6e59 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/memcopy/gm_coord.h @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file gm_coord.h + * \brief + */ +#ifndef GM_COORD_H +#define GM_COORD_H + +struct GmCoord { + uint32_t bIdx; + uint32_t n2Idx; + uint32_t gS1Idx; + uint32_t dIdx; + uint32_t gS1DealSize; + uint32_t dDealSize; +}; +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/common/op_kernel/memcopy/gm_layout.h b/csrc/ascend/attention/common/op_kernel/memcopy/gm_layout.h new file mode 100644 index 000000000..2058d576f --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/memcopy/gm_layout.h @@ -0,0 +1,427 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file gm_layout.h + * \brief + */ +#ifndef GM_LAYOUT_H +#define GM_LAYOUT_H + +#if ASC_DEVKIT_MAJOR >= 9 +#include "kernel_vec_intf.h" +#include "kernel_cube_intf.h" +#else +#include "kernel_operator.h" +#endif + +// ----------------------------------------------GmLayout-------------------------------- +enum class GmFormat { + BSNGD = 0, + BNGSD = 1, + NGBSD = 2, + TNGD = 3, + NGTD = 4, + BSND = 5, + BNSD = 6, + TND = 7, + NTD = 8, + PA_BnBsND = 9, + PA_BnNBsD = 10, + PA_NZ = 11, + NGD = 12, // post_quant + ND = 13, //antiquant no PA + BS2 = 14, + BNS2 = 15, + PA_BnBs = 16, //antiquant PA + PA_BnNBs = 17, + BN2GS1S2 = 18, //PSE_GmFormat + SBNGD = 19, + SBND = 20, + NTGD = 21, + PA_NZ_K_SCALE = 22, +}; + +template +struct GmLayout { +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t g, uint32_t s, uint32_t d) { + shape = AscendC::MakeShape(b, n, g, s, d); + uint64_t dStride = 1; + uint64_t gStride = dStride * d; + uint64_t nStride = gStride * g; + uint64_t sStride = nStride * n; + uint64_t bStride = sStride * s; + stride = AscendC::MakeStride(bStride, nStride, gStride, sStride, dStride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t g, uint32_t s, uint32_t d) { + shape = AscendC::MakeShape(b, n, g, s, d); + uint64_t dStride = 1; + uint64_t sStride = dStride * d; + uint64_t gStride = sStride * s; + uint64_t nStride = gStride * g; + uint64_t bStride = nStride * n; + stride = AscendC::MakeStride(bStride, nStride, gStride, sStride, dStride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t g, uint32_t s, uint32_t d) { + shape = AscendC::MakeShape(b, n, g, s, d); + uint64_t dStride = 1; + uint64_t sStride = dStride * d; + uint64_t bStride = sStride * s; + uint64_t gStride = bStride * b; + uint64_t nStride = gStride * g; + stride = AscendC::MakeStride(bStride, nStride, gStride, sStride, dStride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t t, uint32_t n, uint32_t g, uint32_t d) { + shape = AscendC::MakeShape(t, n, g, d); + uint64_t dStride = 1; + uint64_t gStride = dStride * d; + uint64_t nStride = gStride * g; + uint64_t tStride = nStride * n; + stride = AscendC::MakeStride(tStride, nStride, gStride, dStride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t t, uint32_t n, uint32_t g, uint32_t d) { + shape = AscendC::MakeShape(t, n, g, d); + uint64_t dStride = 1; + uint64_t tStride = dStride * d; + uint64_t gStride = tStride * t; + uint64_t nStride = gStride * g; + stride = AscendC::MakeStride(tStride, nStride, gStride, dStride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t t, uint32_t n, uint32_t g, uint32_t d) { + shape = AscendC::MakeShape(t, n, g, d); + uint64_t dStride = 1; + uint64_t gStride = dStride * d; + uint64_t tStride = gStride * g; + uint64_t nStride = tStride * t; + stride = AscendC::MakeStride(tStride, nStride, gStride, dStride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t s, uint32_t d) { + shape = AscendC::MakeShape(b, n, s, d); + uint64_t dStride = 1; + uint64_t nStride = dStride * d; + uint64_t sStride = nStride * n; + uint64_t bStride = sStride * s; + stride = AscendC::MakeStride(bStride, nStride, sStride, dStride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t s, uint32_t d) { + shape = AscendC::MakeShape(b, n, s, d); + uint64_t dStride = 1; + uint64_t sStride = dStride * d; + uint64_t nStride = sStride * s; + uint64_t bStride = nStride * n; + stride = AscendC::MakeStride(bStride, nStride, sStride, dStride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t t, uint32_t n, uint32_t d) { + shape = AscendC::MakeShape(t, n, d); + uint64_t dStride = 1; + uint64_t nStride = dStride * d; + uint64_t tStride = nStride * n; + stride = AscendC::MakeStride(tStride, nStride, dStride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t t, uint32_t n, uint32_t d) { + shape = AscendC::MakeShape(t, n, d); + uint64_t dStride = 1; + uint64_t tStride = dStride * d; + uint64_t nStride = tStride * t; + stride = AscendC::MakeStride(tStride, nStride, dStride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t n, uint32_t blockSize, uint32_t d) { + shape = AscendC::MakeShape(n, blockSize, d); + uint64_t dStride = 1; + uint64_t nStride = dStride * d; + uint64_t bsStride = nStride * n; + uint64_t bnStride = bsStride * blockSize; + stride = AscendC::MakeStride(bnStride, nStride, bsStride, dStride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t n, uint32_t blockSize, uint32_t d) { + shape = AscendC::MakeShape(n, blockSize, d); + uint64_t dStride = 1; + uint64_t bsStride = dStride * d; + uint64_t nStride = bsStride * blockSize; + uint64_t bnStride = nStride * n; + stride = AscendC::MakeStride(bnStride, nStride, bsStride, dStride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t n, uint32_t blockSize, uint32_t d1, uint32_t d0) { + shape = AscendC::MakeShape(n, d1, blockSize, d0); + uint64_t d0Stride = 1; + uint64_t bsStride = d0Stride * d0; + uint64_t d1Stride = bsStride * blockSize; + uint64_t nStride = d1Stride * d1; + uint64_t bnStride = nStride * n; + stride = AscendC::MakeStride(bnStride, nStride, d1Stride, bsStride, d0Stride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t n, uint32_t blockSize1, uint32_t d, uint32_t blockSize0) { + shape = AscendC::MakeShape(n, blockSize1, d, blockSize0); + uint64_t bs0Stride = 1; + uint64_t dStride = bs0Stride * blockSize0; + uint64_t bs1Stride = dStride * d; + uint64_t nStride = bs1Stride * blockSize1; + uint64_t bnStride = nStride * n; + stride = AscendC::MakeStride(bnStride, nStride, bs1Stride, dStride, bs0Stride); + } +}; + +// post_quant +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t n, uint32_t g, uint32_t d) { + shape = AscendC::MakeShape(n, g, d); + uint64_t dStride = 1; + uint64_t gStride = dStride * d; + uint64_t nStride = gStride * g; + stride = AscendC::MakeStride(nStride, gStride, dStride); + } +}; + +//antiquant +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t n, uint32_t d) { + shape = AscendC::MakeShape(n, d); + + uint64_t dStride = 1; + uint64_t nStride = dStride * d; //headDim + stride = AscendC::MakeStride(nStride, dStride); + } +}; +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t b, uint32_t s) { + shape = AscendC::MakeShape(b, s); + + uint64_t sStride = 1; + uint64_t bStride = sStride * s; + + stride = AscendC::MakeStride(bStride, sStride); + } +}; +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t s) { + shape = AscendC::MakeShape(b, n, s); + + uint64_t sStride = 1; + uint64_t nStride = sStride * s; + uint64_t bStride = nStride * n; + + stride = AscendC::MakeStride(bStride, nStride, sStride); + } +}; +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t blockSize) { + shape = AscendC::MakeShape(blockSize); + + uint64_t bsStride = 1; + uint64_t bnStride = bsStride * blockSize; + stride = AscendC::MakeStride(bnStride, bsStride); + } +}; +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t n, uint32_t blockSize) { + shape = AscendC::MakeShape(n, blockSize); + + uint64_t bsStride = 1; + uint64_t nStride = bsStride * blockSize; + uint64_t bnStride = nStride * n; //blockSize * kvHeadNum + stride = AscendC::MakeStride(bnStride, nStride, bsStride); + } +}; + +//PSE_GmLayout +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t g, uint32_t s1, uint32_t s2) + { + shape = AscendC::MakeShape(b, n, g, s1, s2); + uint64_t s2Stride = 1; + uint64_t s1Stride = s2Stride * s2; + uint64_t gStride = s1Stride * s1; + uint64_t nStride = gStride * g; + uint64_t bStride = nStride * n; + stride = AscendC::MakeStride(bStride, nStride, gStride, s1Stride, s2Stride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t g, uint32_t s, uint32_t d) { + shape = AscendC::MakeShape(b, n, g, s, d); + uint64_t dStride = 1; + uint64_t gStride = dStride * d; + uint64_t nStride = gStride * g; + uint64_t bStride = nStride * n; + uint64_t sStride = bStride * b; + stride = AscendC::MakeStride(bStride, nStride, gStride, sStride, dStride); + } +}; + +template <> +struct GmLayout { + AscendC::Shape shape; + AscendC::Stride stride; + + __aicore__ inline GmLayout() = default; + __aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t s, uint32_t d) { + shape = AscendC::MakeShape(b, n, s, d); + uint64_t dStride = 1; + uint64_t nStride = dStride * d; + uint64_t bStride = nStride * n; + uint64_t sStride = bStride * b; + stride = AscendC::MakeStride(bStride, nStride, sStride, dStride); + } +}; + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/common/op_kernel/memcopy/offset_calculator_v2.h b/csrc/ascend/attention/common/op_kernel/memcopy/offset_calculator_v2.h new file mode 100644 index 000000000..f3bb98233 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/memcopy/offset_calculator_v2.h @@ -0,0 +1,1104 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file offset_calculator_v2.h + * \brief + */ +#ifndef OFFSET_CALCULATOR_V2_H +#define OFFSET_CALCULATOR_V2_H + +#include "gm_layout.h" +#include "parser.h" + +using AscendC::GlobalTensor; + +// ----------------------------------------------GmLayoutParams-------------------------------- +enum class FormatCategory { + GM_Q_OUT_BNGSD = 0, + GM_Q_OUT_TND = 1, + GM_KV_BNSD = 2, + GM_KV_TND = 3, + GM_KV_PA_BNBD = 4, + GM_KV_PA_NZ = 5, + GM_POST_QUANT_NGD = 6, // post_quant + GM_ANTIQ_ND = 7, // antiquant no PA + GM_ANTIQ_BS = 8, + GM_ANTIQ_BNS = 9, + GM_ANTIQ_BnBs = 10, // antiquant PA + GM_ANTIQ_BnNBs = 11, + GM_PSE_BN2GS1S2 = 12, // PSE + GM_V_SCALE_TND = 13, + GM_K_SCALE_PA_NZ = 14, +}; + +template +struct GmLayoutParams {}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_Q_OUT_BNGSD; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_Q_OUT_BNGSD; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_Q_OUT_BNGSD; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_Q_OUT_TND; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_Q_OUT_TND; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_Q_OUT_TND; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_KV_BNSD; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_KV_BNSD; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_KV_TND; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_KV_TND; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_KV_PA_BNBD; +}; + + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_KV_PA_BNBD; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_KV_PA_NZ; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_K_SCALE_PA_NZ; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_Q_OUT_BNGSD; +}; + +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_KV_BNSD; +}; + +// post_quant +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_POST_QUANT_NGD; +}; + +// antiquant +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_ANTIQ_ND; +}; +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_ANTIQ_BS; +}; +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_ANTIQ_BNS; +}; +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_ANTIQ_BnBs; +}; +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_ANTIQ_BnNBs; +}; + +// pse +template <> +struct GmLayoutParams { + static constexpr FormatCategory CATEGORY = FormatCategory::GM_PSE_BN2GS1S2; +}; + +// ----------------------------------------------OffsetCalculator-------------------------------- +template +struct OffsetCalculatorImpl {}; + +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + ActualSeqLensParser actualSeqLensQParser; + bool isQPaddingFlag = false; + uint64_t qPaddingSize = 0; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t b, uint32_t n2, uint32_t g, uint32_t s1, uint32_t d) + { + gmLayout.MakeLayout(b, n2, g, s1, d); + } + + __aicore__ inline void Init(uint32_t b, uint32_t n2, uint32_t g, uint32_t s1, uint32_t d, + GlobalTensor actualSeqLengthsGmQ, uint32_t actualLenQDims, + bool isQPaddingFlag = false, uint64_t qPaddingSize = 0) + { + this->isQPaddingFlag = isQPaddingFlag; + this->qPaddingSize = qPaddingSize; + if (actualLenQDims != 0) { + actualSeqLensQParser.Init(actualSeqLengthsGmQ, actualLenQDims, 0); + } + gmLayout.MakeLayout(b, n2, g, s1, d); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t n2Idx, uint32_t gIdx, uint32_t s1Idx, uint32_t dIdx) + { + if (isQPaddingFlag) { + s1Idx += GetDimS1() - qPaddingSize - actualSeqLensQParser.GetActualSeqLength(bIdx); + } + uint64_t offset = bIdx * GetStrideB() + n2Idx * GetStrideN2() + gIdx * GetStrideG() + s1Idx * GetStrideS1() + + dIdx * GetStrideD(); + return offset; + } + + // Get Stride + __aicore__ inline uint64_t GetStrideB() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideN2() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideG() + { + return AscendC::Std::get<2>(gmLayout.stride); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideS1() + { + return AscendC::Std::get<3>(gmLayout.stride); // 3:代表第4个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideD() + { + return AscendC::Std::get<4>(gmLayout.stride); // 4:代表第5个维度，索引从0开始 + } + + // Get Dim + __aicore__ inline uint64_t GetDimB() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetDimN2() + { + return AscendC::Std::get<1>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetDimG() + { + return AscendC::Std::get<2>(gmLayout.shape); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetDimS1() + { + return AscendC::Std::get<3>(gmLayout.shape); // 3:代表第4个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetDimD() + { + return AscendC::Std::get<4>(gmLayout.shape); // 4:代表第5个维度，索引从0开始 + } +}; + +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + ActualSeqLensParser actualSeqLensQParser; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t n2, uint32_t g, uint32_t d, GlobalTensor actualSeqLengthsGmQ, + uint32_t actualLenQDims) + { + actualSeqLensQParser.Init(actualSeqLengthsGmQ, actualLenQDims); + gmLayout.MakeLayout(actualSeqLensQParser.GetTSize(), n2, g, d); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t n2Idx, uint32_t gIdx, uint32_t s1Idx, uint32_t dIdx) + { + uint64_t tIdx = actualSeqLensQParser.GetTBase(bIdx) + s1Idx; + uint64_t offset = tIdx * GetStrideT() + n2Idx * GetStrideN2() + gIdx * GetStrideG() + dIdx * GetStrideD(); + return offset; + } + + // Get Stride + __aicore__ inline uint64_t GetStrideT() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideN2() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideG() + { + return AscendC::Std::get<2>(gmLayout.stride); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideD() + { + return AscendC::Std::get<3>(gmLayout.stride); // 3:代表第4个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideS1() + { + return GetStrideT(); + } + + // Get Dim + __aicore__ inline uint64_t GetDimT() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetDimN2() + { + return AscendC::Std::get<1>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetDimG() + { + return AscendC::Std::get<2>(gmLayout.shape); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetDimD() + { + return AscendC::Std::get<3>(gmLayout.shape); // 3:代表第4个维度，索引从0开始 + } +}; + +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + ActualSeqLensParser actualSeqLensKVParser; + bool isKvPaddingFlag = false; + uint64_t kvPaddingSize = 0; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t b, uint32_t n2, uint32_t s2, uint32_t d) + { + gmLayout.MakeLayout(b, n2, s2, d); + } + + __aicore__ inline void Init(uint32_t b, uint32_t n2, uint32_t s2, uint32_t d, + GlobalTensor actualSeqLengthsGm, uint32_t actualLenKvDims, + bool isKvPaddingFlag = false, uint64_t kvPaddingSize = 0) + { + this->isKvPaddingFlag = isKvPaddingFlag; + this->kvPaddingSize = kvPaddingSize; + if (actualLenKvDims != 0) { + actualSeqLensKVParser.Init(actualSeqLengthsGm, actualLenKvDims, 0); + } + gmLayout.MakeLayout(b, n2, s2, d); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t n2Idx, uint32_t s2Idx, uint32_t dIdx) + { + if (isKvPaddingFlag) { + s2Idx += GetDimS2() - kvPaddingSize - actualSeqLensKVParser.GetActualSeqLength(bIdx); + } + + uint64_t offset = bIdx * GetStrideB() + n2Idx * GetStrideN2() + s2Idx * GetStrideS2() + dIdx * GetStrideD(); + return offset; + } + + // Get Stride + __aicore__ inline uint64_t GetStrideB() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideN2() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideS2() + { + return AscendC::Std::get<2>(gmLayout.stride); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideD() + { + return AscendC::Std::get<3>(gmLayout.stride); // 3:代表第4个维度，索引从0开始 + } + + // Get Dim + __aicore__ inline uint64_t GetDimB() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetDimN2() + { + return AscendC::Std::get<1>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetDimS2() + { + return AscendC::Std::get<2>(gmLayout.shape); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetDimD() + { + return AscendC::Std::get<3>(gmLayout.shape); // 3:代表第4个维度，索引从0开始 + } +}; + +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + ActualSeqLensParser actualSeqLensKVParser; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t n2, uint32_t d, GlobalTensor actualSeqLengthsGmKV, + uint32_t actualLenKVDims) + { + actualSeqLensKVParser.Init(actualSeqLengthsGmKV, actualLenKVDims); + gmLayout.MakeLayout(actualSeqLensKVParser.GetTSize(), n2, d); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t n2Idx, uint32_t s2Idx, uint32_t dIdx) + { + uint64_t tIdx = actualSeqLensKVParser.GetTBase(bIdx) + s2Idx; + uint64_t offset = tIdx * GetStrideT() + n2Idx * GetStrideN2() + dIdx * GetStrideD(); + return offset; + } + + // Get Stride + __aicore__ inline uint64_t GetStrideT() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideN2() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideD() + { + return AscendC::Std::get<2>(gmLayout.stride); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideS2() + { + return GetStrideT(); + } + + // Get Dim + __aicore__ inline uint64_t GetDimT() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetDimN2() + { + return AscendC::Std::get<1>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetDimD() + { + return AscendC::Std::get<2>(gmLayout.shape); // 2:代表第3个维度，索引从0开始 + } +}; + +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + ActualSeqLensParser actualSeqLensKVParser; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t n2, uint32_t d, GlobalTensor actualSeqLengthsGmKV, + uint32_t actualLenKVDims) + { + actualSeqLensKVParser.Init(actualSeqLengthsGmKV, actualLenKVDims); + gmLayout.MakeLayout(actualSeqLensKVParser.GetTSize(), n2, d); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t n2Idx, uint32_t s2Idx, uint32_t dIdx) + { + uint64_t tIdx = actualSeqLensKVParser.GetMxVscaleTBase(bIdx) + s2Idx; + uint64_t offset = tIdx * GetStrideT() + n2Idx * GetStrideN2() + dIdx * GetStrideD(); + return offset; + } + + // Get Stride + __aicore__ inline uint64_t GetStrideT() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideN2() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideD() + { + return AscendC::Std::get<2>(gmLayout.stride); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideS2() + { + return GetStrideT(); + } + + // Get Dim + __aicore__ inline uint64_t GetDimT() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetDimN2() + { + return AscendC::Std::get<1>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetDimD() + { + return AscendC::Std::get<2>(gmLayout.shape); // 2:代表第3个维度，索引从0开始 + } +}; + +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + BlockTableParser blockTableParser; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t n2, uint32_t blockSize, uint32_t d, GlobalTensor blockTableGm, + uint32_t maxblockNumPerBatch) + { + blockTableParser.Init(blockTableGm, maxblockNumPerBatch); + gmLayout.MakeLayout(n2, blockSize, d); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t n2Idx, uint32_t s2Idx, uint32_t dIdx) + { + uint64_t blockIdxInBatch = s2Idx / GetBlockSize(); // 获取block table上的索引 + uint64_t bsIdx = s2Idx % GetBlockSize(); // 获取在单个块上超出的行数 + int32_t blockIdx = blockTableParser.GetBlockIdx(bIdx, blockIdxInBatch); + + uint64_t offset = + blockIdx * GetStrideBlockNum() + n2Idx * GetStrideN2() + bsIdx * GetStrideBlockSize() + dIdx * GetStrideD(); + return offset; + } + + // Get Stride + __aicore__ inline uint64_t GetStrideBlockNum() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideN2() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideBlockSize() + { + return AscendC::Std::get<2>(gmLayout.stride); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideD() + { + return AscendC::Std::get<3>(gmLayout.stride); // 3:代表第4个维度，索引从0开始 + } + + // Get Dim + __aicore__ inline uint64_t GetN2() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetBlockSize() + { + return AscendC::Std::get<1>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetD() + { + return AscendC::Std::get<2>(gmLayout.shape); // 2:代表第3个维度，索引从0开始 + } +}; + +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + BlockTableParser blockTableParser; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t n2, uint32_t blockSize, uint32_t d1, uint32_t d0, + GlobalTensor blockTableGm, uint32_t maxblockNumPerBatch) + { + blockTableParser.Init(blockTableGm, maxblockNumPerBatch); + gmLayout.MakeLayout(n2, blockSize, d1, d0); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t n2Idx, uint32_t s2Idx, uint32_t dIdx) + { + uint64_t blockIdxInBatch = s2Idx / GetBlockSize(); // 获取block table上的索引 + uint64_t bsIdx = s2Idx % GetBlockSize(); // 获取在单个块上超出的行数 + int32_t blockIdx = blockTableParser.GetBlockIdx(bIdx, blockIdxInBatch); + + uint32_t d1Idx = dIdx / GetD0(); + uint32_t d0Idx = dIdx % GetD0(); + uint64_t offset = blockIdx * GetStrideBlockNum() + n2Idx * GetStrideN2() + d1Idx * GetStrideD1() + + bsIdx * GetStrideBlockSize() + d0Idx * GetStrideD0(); + return offset; + } + + // Get Stride + __aicore__ inline uint64_t GetStrideBlockNum() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideN2() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideD1() + { + return AscendC::Std::get<2>(gmLayout.stride); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideBlockSize() + { + return AscendC::Std::get<3>(gmLayout.stride); // 3:代表第4个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideD0() + { + return AscendC::Std::get<4>(gmLayout.stride); // 4:代表第5个维度，索引从0开始 + } + + // Get Dim + __aicore__ inline uint64_t GetN2() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetD1() + { + return AscendC::Std::get<1>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetBlockSize() + { + return AscendC::Std::get<2>(gmLayout.shape); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetD0() + { + return AscendC::Std::get<3>(gmLayout.shape); // 3:代表第4个维度，索引从0开始 + } +}; + +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + BlockTableParser blockTableParser; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t n2, uint32_t blockSize, uint32_t d1, uint32_t d0, + GlobalTensor blockTableGm, uint32_t maxblockNumPerBatch) + { + blockTableParser.Init(blockTableGm, maxblockNumPerBatch); + gmLayout.MakeLayout(n2, blockSize, d1, d0); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t n2Idx, uint32_t s2Idx, uint32_t dIdx) + { + uint64_t blockIdxInBatch = s2Idx / GetBlockSize(); // 获取block table上的索引 + uint64_t bsIdx = s2Idx % GetBlockSize(); // 获取在单个块上超出的行数 + int32_t blockIdx = blockTableParser.GetBlockIdx(bIdx, blockIdxInBatch); + + constexpr uint32_t bs0 = 16; + uint32_t bs1Idx = bsIdx / bs0; + uint32_t bs0Idx = bsIdx % bs0; + + uint64_t offset = blockIdx * GetStrideBlockNum() + n2Idx * GetStrideN2() + bs1Idx * GetStrideBlockSize1() + + dIdx * GetStrideD() + bs0Idx * GetStrideBlockSize0(); + return offset; + } + + // Get Stride + __aicore__ inline uint64_t GetStrideBlockNum() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideN2() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideBlockSize1() + { + return AscendC::Std::get<2>(gmLayout.stride); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideD() + { + return AscendC::Std::get<3>(gmLayout.stride); // 3:代表第4个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideBlockSize0() + { + return AscendC::Std::get<4>(gmLayout.stride); // 4:代表第5个维度，索引从0开始 + } + + // Get Dim + __aicore__ inline uint64_t GetN2() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetBlockSize1() + { + return AscendC::Std::get<1>(gmLayout.shape); + } + + __aicore__ inline uint64_t GetD() + { + return AscendC::Std::get<2>(gmLayout.shape); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetBlockSize0() + { + return AscendC::Std::get<3>(gmLayout.shape); // 3:代表第4个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetBlockSize() + { + return GetBlockSize1() * 16; + } +}; + +// post_quant +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t n2, uint32_t g, uint32_t d) + { + gmLayout.MakeLayout(n2, g, d); + } + + __aicore__ inline uint64_t GetOffset(uint32_t n2Idx, uint32_t gIdx, uint32_t dIdx) + { + uint64_t offset = n2Idx * GetStrideN2() + gIdx * GetStrideG() + dIdx * GetStrideD(); + return offset; + } + + // Get Stride + __aicore__ inline uint64_t GetStrideN2() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideG() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideD() + { + return AscendC::Std::get<2>(gmLayout.stride); // 2:代表第3个维度，索引从0开始 + } + + // Get Dim + __aicore__ inline uint32_t GetDimN2() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint32_t GetDimG() + { + return AscendC::Std::get<1>(gmLayout.shape); + } + + __aicore__ inline uint32_t GetDimD() + { + return AscendC::Std::get<2>(gmLayout.shape); // 2:代表第3个维度，索引从0开始 + } +}; + +// antiquant +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t n2, uint32_t d) + { + gmLayout.MakeLayout(n2, d); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t n2Idx, uint32_t s2Idx, uint32_t dIdx) + { + uint64_t offset = n2Idx * GetStrideN2() + dIdx * GetStrideD(); + return offset; + } + + // Get Stride + __aicore__ inline uint64_t GetStrideN2() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideD() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + // Get Dim + __aicore__ inline uint32_t GetDimN2() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint32_t GetDimD() + { + return AscendC::Std::get<1>(gmLayout.shape); + } +}; + +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t b, uint32_t s2) + { + gmLayout.MakeLayout(b, s2); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t n2Idx, uint32_t s2Idx, uint32_t dIdx) + { + uint64_t offset = bIdx * GetStrideB() + s2Idx * GetStrideS2(); + return offset; + } + + // Get Stride + __aicore__ inline uint64_t GetStrideB() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideS2() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + // Get Dim + __aicore__ inline uint32_t GetDimB() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint32_t GetDimS2() + { + return AscendC::Std::get<1>(gmLayout.shape); + } +}; + +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t b, uint32_t n2, uint32_t s2) + { + gmLayout.MakeLayout(b, n2, s2); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t n2Idx, uint32_t s2Idx, uint32_t dIdx) + { + uint64_t offset = bIdx * GetStrideB() + n2Idx * GetStrideN2() + s2Idx * GetStrideS2(); + return offset; + } + + // Get Stride + __aicore__ inline uint64_t GetStrideB() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideN2() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideS2() + { + return AscendC::Std::get<2>(gmLayout.stride); // 2:代表第3个维度，索引从0开始 + } + + // Get Dim + __aicore__ inline uint32_t GetDimB() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint32_t GetDimN2() + { + return AscendC::Std::get<1>(gmLayout.shape); + } + + __aicore__ inline uint32_t GetDimS2() + { + return AscendC::Std::get<2>(gmLayout.shape); // 2:代表第3个维度，索引从0开始 + } +}; + +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + BlockTableParser blockTableParser; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t blockSize, GlobalTensor blockTableGm, uint32_t maxblockNumPerBatch) + { + blockTableParser.Init(blockTableGm, maxblockNumPerBatch); + gmLayout.MakeLayout(blockSize); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t nIdx, uint32_t sIdx) + { + uint64_t blockIdxInBatch = sIdx / GetStrideBlockSize(); // 获取block table上的索引 + uint64_t bsIdx = sIdx % GetStrideBlockSize(); // 获取在单个块上超出的行数 + int32_t blockIdx = blockTableParser.GetBlockIdx(bIdx, blockIdxInBatch); + uint64_t offset = blockIdx * GetStrideBlockNum() + bsIdx * GetStrideBlockSize(); + + return offset; + } + + // Get Stride + + __aicore__ inline uint64_t GetStrideBlockNum() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideBlockSize() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + // Get Dim + __aicore__ inline uint32_t GetDimBlockSize() + { + return AscendC::Std::get<0>(gmLayout.shape); + } +}; + +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + BlockTableParser blockTableParser; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t n, uint32_t blockSize, GlobalTensor blockTableGm, + uint32_t maxblockNumPerBatch) + { + blockTableParser.Init(blockTableGm, maxblockNumPerBatch); + gmLayout.MakeLayout(n, blockSize); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t nIdx, uint32_t sIdx) + { + uint64_t blockIdxInBatch = sIdx / GetStrideBlockSize(); // 获取block table上的索引 + uint64_t bsIdx = sIdx % GetStrideBlockSize(); // 获取在单个块上超出的行数 + int32_t blockIdx = blockTableParser.GetBlockIdx(bIdx, blockIdxInBatch); + uint64_t offset = blockIdx * GetStrideBlockNum() + nIdx * GetStrideN() + bsIdx * GetStrideBlockSize(); + + return offset; + } + + // Get Stride + + __aicore__ inline uint64_t GetStrideBlockNum() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideN() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideBlockSize() + { + return AscendC::Std::get<2>(gmLayout.stride); // 2:代表第3个维度，索引从0开始 + } + + // Get Dim + __aicore__ inline uint32_t GetDimN() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint32_t GetDimBlockSize() + { + return AscendC::Std::get<1>(gmLayout.shape); + } +}; + +// PSE +template +struct OffsetCalculatorImpl { + GmLayout gmLayout; + ActualSeqLensParser actualSeqLensQParser; + bool isQPaddingFlag = false; + uint64_t qPaddingSize = 0; + + __aicore__ inline OffsetCalculatorImpl() = default; + + __aicore__ inline void Init(uint32_t b, uint32_t n2, uint32_t g, uint32_t s1, uint32_t s2, + GlobalTensor actualSeqLengthsGmQ, uint32_t actualLenQDims, + bool isQPaddingFlag = false, uint64_t qPaddingSize = 0) + { + this->isQPaddingFlag = isQPaddingFlag; + this->qPaddingSize = qPaddingSize; + if (actualLenQDims != 0) { + actualSeqLensQParser.Init(actualSeqLengthsGmQ, actualLenQDims, 0); + } + gmLayout.MakeLayout(b, n2, g, s1, s2); + } + + __aicore__ inline uint64_t GetOffset(uint32_t bIdx, uint32_t n2Idx, uint32_t gIdx, uint32_t s1Idx, uint32_t s2Idx) + { + if (isQPaddingFlag) { + s1Idx += GetDimS1() - qPaddingSize - actualSeqLensQParser.GetActualSeqLength(bIdx); + } + uint64_t offset = bIdx * GetStrideB() + n2Idx * GetStrideN2() + gIdx * GetStrideG() + s1Idx * GetStrideS1() + + s2Idx * GetStrideS2(); + return offset; + } + + // Get Stride + __aicore__ inline uint64_t GetStrideB() + { + return AscendC::Std::get<0>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideN2() + { + return AscendC::Std::get<1>(gmLayout.stride); + } + + __aicore__ inline uint64_t GetStrideG() + { + return AscendC::Std::get<2>(gmLayout.stride); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideS1() + { + return AscendC::Std::get<3>(gmLayout.stride); // 3:代表第4个维度，索引从0开始 + } + + __aicore__ inline uint64_t GetStrideS2() + { + return AscendC::Std::get<4>(gmLayout.stride); // 4:代表第5个维度，索引从0开始 + } + + // Get Dim + __aicore__ inline uint32_t GetDimB() + { + return AscendC::Std::get<0>(gmLayout.shape); + } + + __aicore__ inline uint32_t GetDimN2() + { + return AscendC::Std::get<1>(gmLayout.shape); + } + + __aicore__ inline uint32_t GetDimG() + { + return AscendC::Std::get<2>(gmLayout.shape); // 2:代表第3个维度，索引从0开始 + } + + __aicore__ inline uint32_t GetDimS1() + { + return AscendC::Std::get<3>(gmLayout.shape); // 3:代表第4个维度，索引从0开始 + } + + __aicore__ inline uint32_t GetDimS2() + { + return AscendC::Std::get<4>(gmLayout.shape); // 4:代表第5个维度，索引从0开始 + } +}; + +template +struct OffsetCalculator : public OffsetCalculatorImpl::CATEGORY, ACTLEN_T> {}; + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/common/op_kernel/memcopy/parser.h b/csrc/ascend/attention/common/op_kernel/memcopy/parser.h new file mode 100644 index 000000000..b5a0e724a --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/memcopy/parser.h @@ -0,0 +1,140 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file parser.h + * \brief + */ +#ifndef PARSER_H +#define PARSER_H + +#if ASC_DEVKIT_MAJOR >= 9 +#include "kernel_vec_intf.h" +#include "kernel_cube_intf.h" +#else +#include "kernel_operator.h" +#endif + +using AscendC::GlobalTensor; + +// ----------------------------------------------ActualSeqLensParser-------------------------------- +enum class ActualSeqLensMode +{ + BY_BATCH = 0, + ACCUM = 1, +}; + +template +class ActualSeqLensParser { +}; + +template +class ActualSeqLensParser { +public: + __aicore__ inline ActualSeqLensParser() = default; + + __aicore__ inline void Init(GlobalTensor actualSeqLengthsGm, uint32_t actualLenDims, + uint64_t defaultVal = 0) + { + this->actualSeqLengthsGm = actualSeqLengthsGm; + this->actualLenDims = actualLenDims; + } + + __aicore__ inline uint64_t GetTBase(uint32_t bIdx) const + { + if (bIdx == 0) { + return 0; + } + return actualSeqLengthsGm.GetValue(bIdx - 1); + } + + __aicore__ inline uint64_t GetMxVscaleTBase(uint32_t bIdx) const + { + if (bIdx == 0) { + return 0; + } + uint64_t vScaleTBaseOffset = 0; + for (uint32_t idx = 0; idx < bIdx; idx++) { + vScaleTBaseOffset += ((GetActualSeqLength(idx) + 63) >> 6); + } + return vScaleTBaseOffset; + } + + __aicore__ inline uint64_t GetActualSeqLength(uint32_t bIdx) const + { + if (bIdx == 0) { + return actualSeqLengthsGm.GetValue(0); + } + return (actualSeqLengthsGm.GetValue(bIdx) - actualSeqLengthsGm.GetValue(bIdx - 1)); + } + + __aicore__ inline uint64_t GetTSize() const + { + return actualSeqLengthsGm.GetValue(actualLenDims - 1); + } +private: + GlobalTensor actualSeqLengthsGm; + uint32_t actualLenDims; +}; + +template +class ActualSeqLensParser { +public: + __aicore__ inline ActualSeqLensParser() = default; + + __aicore__ inline void Init(GlobalTensor actualSeqLengthsGm, uint32_t actualLenDims, uint64_t defaultVal) + { + this->actualSeqLengthsGm = actualSeqLengthsGm; + this->actualLenDims = actualLenDims; + this->defaultVal = defaultVal; + } + + __aicore__ inline uint64_t GetActualSeqLength(uint32_t bIdx) const + { + if (actualLenDims == 0) { + return defaultVal; + } + if (actualLenDims == 1) { + return actualSeqLengthsGm.GetValue(0); + } + return actualSeqLengthsGm.GetValue(bIdx); + } + + __aicore__ inline uint32_t GetActualLenDims() const + { + return actualLenDims; + } +private: + GlobalTensor actualSeqLengthsGm; + uint32_t actualLenDims = 0; + uint64_t defaultVal = 0; +}; + +// ----------------------------------------------BlockTableParser-------------------------------- +class BlockTableParser { +public: + __aicore__ inline BlockTableParser() = default; + + __aicore__ inline void Init(GlobalTensor blockTableGm, uint32_t maxblockNumPerBatch) + { + this->blockTableGm = blockTableGm; + this->maxblockNumPerBatch = maxblockNumPerBatch; + } + + __aicore__ inline int32_t GetBlockIdx(uint32_t bIdx, uint32_t blockIdxInBatch) const + { + return blockTableGm.GetValue(bIdx * maxblockNumPerBatch + blockIdxInBatch); + } +private: + GlobalTensor blockTableGm; + uint32_t maxblockNumPerBatch; +}; + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/common/op_kernel/offset_calculator.h b/csrc/ascend/attention/common/op_kernel/offset_calculator.h new file mode 100644 index 000000000..d16dd7026 --- /dev/null +++ b/csrc/ascend/attention/common/op_kernel/offset_calculator.h @@ -0,0 +1,31 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file offset_calculator.h + * \brief + */ +#ifndef OFFSET_CALCULATOR_H +#define OFFSET_CALCULATOR_H + +#if ASC_DEVKIT_MAJOR >= 9 +#include "kernel_basic_intf.h" +#else +#include "kernel_operator.h" +#endif + +#include "memcopy/gm_layout.h" +#include "memcopy/parser.h" +#include "memcopy/offset_calculator_v2.h" +#include "memcopy/fa_gm_tensor.h" +#include "memcopy/fa_l1_tensor.h" +#include "memcopy/gm_coord.h" + +#endif diff --git a/csrc/ascend/attention/compressor/CMakeLists.txt b/csrc/ascend/attention/compressor/CMakeLists.txt new file mode 100644 index 000000000..a3ab34a93 --- /dev/null +++ b/csrc/ascend/attention/compressor/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/attention/compressor/README.md b/csrc/ascend/attention/compressor/README.md new file mode 100644 index 000000000..f6025cf02 --- /dev/null +++ b/csrc/ascend/attention/compressor/README.md @@ -0,0 +1,499 @@ +# Compressor + +## 产品支持情况 + +| 产品 | 是否支持 | +| ------------------------------------------------------------ | :------: | +|Ascend 950PR/Ascend 950DT| √ | +|Atlas A3 训练系列产品/Atlas A3 推理系列产品| √ | +|Atlas A2 训练系列产品/Atlas A2 推理系列产品| × | +|Atlas 200I/500 A2 推理产品| × | +|Atlas 推理系列加速卡产品| × | +|Atlas 训练系列产品| × | + +## 功能说明 + +- API功能：Compressor是推理场景下SAS和QLI的前处理算子，用于将每4或128个token的KV cache压缩成一个，然后每个token与这些压缩的KV cache进行DSA计算。在长序列的情况下，Compressor可以有效地减少计算开销。 + +- 计算公式： + + 压缩阶段： + 1. 计算矩阵乘法： + - C4A: $\left[kv\_state^a, score\_state^a\right] = X @ \left[W^{aKV}, W^{aGate}\right], \left[kv\_state^b, score\_state^b\right] = X @ \left[W^{bKV}, W^{bGate}\right];$ + - C128A: $\left[kv\_state, score\_state\right] = X @ \left[W^{KV}, W^{Gate}\right]$ + 2. 计算分组加法： + - C4A: $score\_state_i^\prime = \left[score\_state_{\left[4(i-1)+1:4i,:\right]}^a; score\_state_{\left[4i+1:4(i+1),:\right]}^b\right] + Ape,~i=1,2,\cdots, \frac{s}{4};$ + - C128A: $score\_state_i^\prime = score\_state_{\left[128(i-1)+1:128i,:\right]} + Ape,~i=1,2,\cdots, \frac{s}{128};$ + 3. 计算分组Softmax： + - C4A: $S_i^\prime = softmax(score\_state_i^\prime),~i=1,2,\cdots, \frac{s}{4};$ + - C128A: $S_i^\prime = softmax(score\_state_i^\prime),~i=1,2,\cdots, \frac{s}{128};$ + 4. 计算Hadamard乘积： + - C4A: $(S_H)_i = S_i^\prime \odot \left[kv\_state^a_{\left[4(i-1)+1:4i,:\right]} ; kv\_state^b_{\left[4i+1:4(i+1),:\right]}\right],~i=1,2,\cdots, \frac{s}{4};$ + - C128A: $S_H = S_i^\prime \odot kv\_state;$ + 5. 沿着压缩轴分组求和： + - C4A: $C_{i}^{\text{Comp}} = \left[1\right]_{1\times8} @ (S_H)_i, ~i=1,2,\cdots, \frac{s}{4};$ + - C128A: $C_{i}^{\text{Comp}} = \left[1\right]_{1\times128} @ (S_H)_i, ~i=1,2,\cdots, \frac{s}{128};$ + + 后处理阶段： + + 6. 计算RMSNorm： + - $\text{RMS}(C^{\text{Comp}}) = \sqrt{\frac{1}{N} \sum_{i=j* N}^{(j+1)* N} {(C_{i}^{\text{Comp}})}^{\text{2}} + norm\_eps} ,N=head\_dim, ~j=1,2,\cdots, \frac{s}{cmp\_ratio}$ + - $\text{RmsNorm}(C^{\text{Comp}}) = norm\_weight \cdot \frac{C_{i}^{\text{Comp}}}{\text{RMS}(C^{\text{Comp}})}$ + 7. 计算Rope； + +- 主要计算过程为： + 1. 将输入$X$与$W^{KV}$做Matmul运算得到$kv\_state$，将输入$X$与$W^{Gate}$做Matmul运算后再与$Ape$做Add运算得到$score\_state$，$kv\_state$与$score\_state$根据输入的start_pos及cu_seqlens完成更新。 + 2. 在coff为2的情况下对$kv\_state$和$score\_state$进行数据重排。 + 3. 对$score\_state$进行softmax运算将softmax结果与$kv\_state$做Mul计算，后进行ReduceSum运算。 + 4. 根据输入数据norm_weight、rope_sin、rope_cos，进行RMSNorm和Rope运算，得到$cmp\_kv$结果输出。 + +## 参数说明 + +| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 | +|----------------------------|-----------|----------------------------------------------------------------------|----------------|------------| +| x | 输入 | 公式中的$X$，表示原始不经压缩的数据。 | FLOAT16、BFLOAT16 | ND | +| wkv | 输入 | 公式中的$W^{KV}$，表示kv压缩权重。 | FLOAT16、BFLOAT16 | ND | +| wgate | 输入 | 公式中的$W^{Gate}$，表示gate压缩权重。 | FLOAT16、BFLOAT16 | ND | +| kv_state | 输入 | 公式中的$kv\_state$，表示kv\_state的历史数据。 | FLOAT32 | ND | +| score_state | 输入 | 公式中的$score\_state$，表示score\_state中的历史数据。 | FLOAT32 | ND | +| ape | 输入 | 公式中的$Ape$，表示positional biases。 | FLOAT32 | ND | +| norm\_weight | 输入 | 表示计算RmsNorm时的权重系数。 | FLOAT16、BFLOAT16 | ND | +| rope\_sin | 输入 | 表示Rope计算时sin的权重系数。 | FLOAT16、BFLOAT16 | ND | +| rope\_cos | 输入 | 表示Rope计算时cos的权重系数。 | FLOAT16、BFLOAT16 | ND | +| rope\_head\_dim | 属性 | 表示rope_cos和rope_sin的hidden层最小单元大小，当前仅支持64。 | INT32 | - | +| cmp\_ratio | 属性 | 用于稀疏计算，表示数据压缩率。 | INT32 | - | +| kv\_block\_table | 可选输入 | 表示kv\_state存储使用的block映射表。当其中元素的值为0时，表示当前位置无需进行更新kv\_state操作。 | INT32 | ND | +| score\_block\_table | 可选输入 | 表示score\_state存储使用的block映射表。当其中元素的值为0时，表示当前位置无需进行更新score\_state操作。 | INT32 | ND | +| cu\_seqlens | 可选输入 | 表示不同Batch中的有效token数。 | INT32 | ND | +| seqused | 可选输入 | 表示不同Batch中实际参与压缩的token数，如果指定为None时，表示和每个Batch上的Sequence Length长度相同。 | INT32 | ND | +| start\_pos | 可选输入 | 表示计算起始位置。 | INT32 | ND | +| coff | 可选属性 | 默认值1，支持1/2。当coff=1时，无需进行overlap数据重排。当coff=2时，需要进行overlap数据重排。 | INT32 | - | +| norm\_eps | 可选属性 | 表示RmsNorm计算的权重系数。默认值1e-6。 | FLOAT32 | - | +| rotary\_mode | 可选属性 | 表示Rop计算的模式。默认值1，支持1/2。rotary\_mode为1时，代表half模式。rotary\_mode为2时，代表interleave模式。 | INT32 | - | +| enabled\_grad | 可选属性 | 训练场景使用，表示是否参与反向更新。默认值false，支持false/true。**目前暂不支持输入true**。 | BOOL | - | +| cmp\_kv | 输出 | 表示压缩后的数据。 | FLOAT16、BFLOAT16 | ND | +| wkv\_proj | 可选输出 | 训练反向使用，表示wkv权重Matmul的计算结果，**目前暂不支持返回wkv\_proj**。 | FLOAT16、BFLOAT16 | ND | +| softmax\_res | 可选输出 | 训练反向使用，表示Softmax计算结果，**目前暂不支持返回softmax\_res**。 | FLOAT16、BFLOAT16 | ND | +| norm\_x | 可选输出 | 训练反向使用，表示Rms计算的输入，**目前暂不支持返回norm\_x**。 | FLOAT16、BFLOAT16 | ND | +| norm\_rstd | 可选输出 | 训练反向使用，表示Rms计算的中间结果，**目前暂不支持返回norm\_rstd**。 | FLOAT16、BFLOAT16 | ND | + +## 约束说明 + +- x参数维度含义：B（Batch Size）表示输入样本批量大小、S（Sequence Length）表示输入样本序列长度、H（Head Size）表示hidden层的大小、D（Head Dim）表示hidden层的最小单元大小、T表示所有Batch输入样本序列长度的累加和。 +- 输入shape限制： + - wkv支持输入shape[coff* D,H] + - wgate支持输入shape[coff* D,H] + - kv\_state、score\_state支持输入shape[block_num,block_size,coff* D]，要求block_num>0。 + - ape支持输入shape[cmp_ratio,coff* D] + - norm\_weight支持输入shape[D,] + - start\_pos支持输入shape[B,] + - 若x的维度采用BS合轴，即x的输入shape为[T,H] + - rope_sin、rope_cos要求输入shape为[min(T,T//cmp_ratio+B),rope_head_dim]。 + - cu\_seqlens输入shape必须为[B+1,]。该参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值，且第一位必须位0。 + - seqused，支持输入shape[B,]，要求每个Batch的有效token数要求小于等于对应Sequence Length长度，即seqused[n] <= cu\_seqlens[n+1] - cu\_seqlens[n]，且不小于0。 + - kv\_block\_table、score\_block\_table支持输入shape[B,ceil(Smax/block_size)]。Smax为每个Batch中最大的Sequence Length，即Smax=max(start\_pos)+max(cu\_seqlens[n+1] - cu\_seqlens[n])。 + - cmp\_kv，输出shape为[min(T,T//cmp_ratio+B),D]：compressed_tokens + compressed_tokens + ... + compressed_tokens + pad。 + - wkv\_proj，输出shape为[T,coff* D]。 + - norm\_x，输出shape为[min(T,T//cmp_ratio+B),D]。 + - norm\_rstd，输出shape为[min(T,T//cmp_ratio+B)]。 + - 若x的维度不采用BS合轴，即x的输入shape为[B,S,H] + - rope_sin、rope_cos要求输入shape为[B,ceil(S/cmp_ratio),rope_head_dim]。 + - cu\_seqlens，参数必须为空。 + - seqused，支持输入shape[B,]，要求每个Batch的有效token数要求小于等于对应Sequence Length长度，即要求seqused[n] <= S，且不小于0。 + - kv\_block\_table、score\_block\_table支持输入shape[B,ceil(Smax/block_size)]。Smax为每个Batch中最大的Sequence Length，即Smax=max(start\_pos)+S。 + - cmp\_kv，输出shape为[B,ceil(S/cmp_ratio),D]：(compressed_tokens+pad0) + (compressed_tokens+pad1) + ... + (compressed_tokens+padN)。 + - wkv\_proj，输出shape为[B,S,coff* D]。 + - norm\_x，输出shape为[B,ceil(S/cmp_ratio),D]。 + - norm\_rstd，输出shape为[B,ceil(S/cmp_ratio)]。 +- 输入值域限制： + - 该接口支持B、S泛化，且存在如下场景限制： + - 部分长序列场景下，如果计算量过大可能会导致出现超过NPU内存的报错，注：这里计算量会受x输入shape的影响，值越大计算量越大。典型的长序列（即B、S的乘积或T较大）场景包括但不限于： +
+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
BSH
100655254096
252611204096
1001310724096
1002611204096
+
+- 输入属性限制： + - 支持D为128/512。 + - 支持H为1K~10K，512对齐。 + - 泛化支持block_size小于等于1024，16对齐。 + - 支持cmp_ratio为4/128。支持如下三种情况： + - C4A: D=512, coff=2, cmp_ratio=4； + - C4Li: D=128, coff=2, cmp_ratio=4； + - C128A: D=512, coff=1, cmp_ratio=128。 + - 支持rotary_mode为2，Rope计算模式为interleave。 + +## Atlas A3 推理系列产品 调用说明 + +- 单算子模式调用 + + ```python + import torch + import torch_npu + import numpy as np + import custom_ops + import torch.nn as nn + import math + + def get_seq_used_by_batch(batch_idx, S, seqused, cu_seqlens): + if seqused is not None: + return seqused[batch_idx] + else: + if cu_seqlens is not None: + return cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx] + else: + return S + + data_type = torch.bfloat16 + hidden_size = 4096 + rope_head_dim = 64 + norm_eps = 1e-6 + coff = 1 # 1:no overlap 2:overlap + cmp_ratio = 128 + rotary_mode = 2 + head_dim = 512 + cu_seqlens = [0, 1] + # ------------- + B = 1 + S = 1 + S_max = 0 + block_size = 128 + start_pos = [8191] * B # (B,) + start_p=8191 + seqused = None # (B,), None时cu_seqlens的数据全部参与计算，否则按传参实际值计算 + + # BS是否合轴 + bs_combine_flag = True + update_flag = 1 + + if seqused is not None: + seqused = torch.tensor(seqused).to(torch.int32) + if start_pos is not None: + start_pos = torch.tensor(start_pos).to(torch.int32) + else: + start_pos = torch.full((B,), start_p, dtype=torch.int32) + + if bs_combine_flag: + if cu_seqlens is None: + T = B * S + if T !=0: + cu_seqlens = torch.arange(0, T + 1, S, dtype=torch.int32) + else: + cu_seqlens = torch.zeros((B+1), dtype=torch.int32) + else: + cu_seqlens = torch.tensor(cu_seqlens).to(torch.int32) + for i in range(B): + if start_pos[i] + cu_seqlens[i + 1] - cu_seqlens[i] > S_max: + S_max = start_pos[i] + cu_seqlens[i + 1] - cu_seqlens[i] + else: + cu_seqlens = None + S_max = max(start_pos) + S + ### ======================== gen input data start ============================= + # page state + max_block_num_per_batch = (S_max + block_size - 1) // block_size + block_num = B * max_block_num_per_batch + next_block_id = 1 + print(f"max_block_num_per_batch: {max_block_num_per_batch}") + block_table = torch.zeros(size=(B, max_block_num_per_batch), dtype=torch.int32) + for i in range(B): + # 需要读取state的范围 + cur_start = start_pos[i] // cmp_ratio * cmp_ratio - cmp_ratio + cur_end = start_pos[i] // cmp_ratio * cmp_ratio + cmp_ratio + if start_pos[i] % cmp_ratio == 0: + cur_end = start_pos[i] + cur_end = min(cur_end, start_pos[i] + S) + cur_start_block_id = (cur_start // block_size) if cur_start >= 0 else 0 + cur_end_block_id = (cur_end - 1) // block_size + for j in range(cur_start_block_id, cur_end_block_id + 1): + block_table[i][j] = next_block_id + next_block_id = next_block_id + 1 + # 需要写入state的范围 + end_pos = get_seq_used_by_batch(i, S, seqused, cu_seqlens) + next_start = (start_pos[i] + end_pos) // cmp_ratio * cmp_ratio - cmp_ratio + next_end = (start_pos[i] + end_pos) // cmp_ratio * cmp_ratio + cmp_ratio + if (start_pos[i] + end_pos) % cmp_ratio == 0: + next_end = start_pos[i] + end_pos + next_end = min(next_end, start_pos[i] + end_pos) + next_start_block_id = (next_start // block_size) if next_start >= 0 else 0 + next_end_block_id = (next_end - 1) // block_size + for j in range(next_start_block_id, next_end_block_id + 1): + if block_table[i][j] == 0: + block_table[i][j] = next_block_id + next_block_id = next_block_id + 1 + + if B==0: + kv_state = torch.tensor(np.random.uniform(-10, 10, (0, block_size, coff * head_dim))).to(torch.float32) + score_state = torch.tensor(np.random.uniform(-10, 10, (0, block_size, coff * head_dim))).to(torch.float32) + else: + kv_state = torch.tensor(np.random.uniform(-10, 10, (torch.max(block_table) + 1, block_size, coff * head_dim))).to(torch.float32) + score_state = torch.tensor(np.random.uniform(-10, 10, (torch.max(block_table) + 1, block_size, coff * head_dim))).to(torch.float32) + + # other input + if bs_combine_flag: + x_shape = (cu_seqlens[-1], hidden_size) + rope_sin_shape = (min(x_shape[0], x_shape[0] // cmp_ratio + B), rope_head_dim) + rope_cos_shape = rope_sin_shape + else: + x_shape = (B, S, hidden_size) + rope_sin_shape = (B, (S + cmp_ratio - 1) // cmp_ratio, rope_head_dim) + rope_cos_shape = rope_sin_shape + + x = torch.tensor(np.random.uniform(-10.0, 10.0, x_shape)).to(data_type).npu() + wkv = torch.tensor(np.random.uniform(-10, 10, (coff * head_dim, hidden_size))).to(data_type).npu() + wgate = torch.tensor(np.random.uniform(-10, 10, (coff * head_dim, hidden_size))).to(data_type).npu() + ape = torch.tensor(np.random.uniform(-10, 10, (cmp_ratio, coff * head_dim))).to(torch.float32).npu() + norm_weight = torch.tensor(np.random.uniform(-10, 10, (head_dim))).to(data_type).npu() + rope_sin = torch.tensor(np.random.uniform(-1, 1, rope_sin_shape)).to(data_type).npu() + rope_cos = torch.tensor(np.random.uniform(-1, 1, rope_cos_shape)).to(data_type).npu() + kv_state = kv_state.npu() + score_state = score_state.npu() + block_table = block_table.npu() + start_pos = torch.tensor(start_pos).to(torch.int32).npu() + if cu_seqlens is not None: + cu_seqlens = torch.tensor(cu_seqlens).to(torch.int32).npu() + if seqused is not None: + seqused = torch.tensor(seqused).to(torch.int32).npu() + + cmp_kv,_ ,_ ,_ ,_ = ( + torch.ops.custom.compressor( + x, + wkv, + wgate, + kv_state, + score_state, + ape, + norm_weight, + rope_sin, + rope_cos, + kv_block_table = block_table, + score_block_table = block_table, + cu_seqlens = cu_seqlens, + seqused = seqused, + start_pos = start_pos, + rope_head_dim = rope_head_dim, + cmp_ratio = cmp_ratio, + coff = coff, + norm_eps = norm_eps, + rotary_mode = rotary_mode + ) + ) + ``` +- aclgraph调用 + + ```python + import torch + import torch_npu + import numpy as np + import torch.nn as nn + import torchair + import custom_ops + import math + + def get_seq_used_by_batch(batch_idx, S, seqused, cu_seqlens): + if seqused is not None: + return seqused[batch_idx] + else: + if cu_seqlens is not None: + return cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx] + else: + return S + + data_type = torch.bfloat16 + hidden_size = 4096 + rope_head_dim = 64 + norm_eps = 1e-6 + coff = 1 # 1:no overlap 2:overlap + cmp_ratio = 128 + rotary_mode = 2 + head_dim = 512 + cu_seqlens = [0, 1] + # ------------- + B = 1 + S = 1 + S_max = 0 + block_size = 128 + start_pos = [8191] * B # (B,) + start_p=8191 + seqused = None # (B,), None时cu_seqlens的数据全部参与计算，否则按传参实际值计算 + + # BS是否合轴 + bs_combine_flag = True + update_flag = 1 + + if seqused is not None: + seqused = torch.tensor(seqused).to(torch.int32) + if start_pos is not None: + start_pos = torch.tensor(start_pos).to(torch.int32) + else: + start_pos = torch.full((B,), start_p, dtype=torch.int32) + + if bs_combine_flag: + if cu_seqlens is None: + T = B * S + if T !=0: + cu_seqlens = torch.arange(0, T + 1, S, dtype=torch.int32) + else: + cu_seqlens = torch.zeros((B+1), dtype=torch.int32) + else: + cu_seqlens = torch.tensor(cu_seqlens).to(torch.int32) + for i in range(B): + if start_pos[i] + cu_seqlens[i + 1] - cu_seqlens[i] > S_max: + S_max = start_pos[i] + cu_seqlens[i + 1] - cu_seqlens[i] + else: + cu_seqlens = None + S_max = max(start_pos) + S + ### ======================== gen input data start ============================= + # page state + max_block_num_per_batch = (S_max + block_size - 1) // block_size + block_num = B * max_block_num_per_batch + next_block_id = 1 + print(f"max_block_num_per_batch: {max_block_num_per_batch}") + block_table = torch.zeros(size=(B, max_block_num_per_batch), dtype=torch.int32) + for i in range(B): + # 需要读取state的范围 + cur_start = start_pos[i] // cmp_ratio * cmp_ratio - cmp_ratio + cur_end = start_pos[i] // cmp_ratio * cmp_ratio + cmp_ratio + if start_pos[i] % cmp_ratio == 0: + cur_end = start_pos[i] + cur_end = min(cur_end, start_pos[i] + S) + cur_start_block_id = (cur_start // block_size) if cur_start >= 0 else 0 + cur_end_block_id = (cur_end - 1) // block_size + for j in range(cur_start_block_id, cur_end_block_id + 1): + block_table[i][j] = next_block_id + next_block_id = next_block_id + 1 + # 需要写入state的范围 + end_pos = get_seq_used_by_batch(i, S, seqused, cu_seqlens) + next_start = (start_pos[i] + end_pos) // cmp_ratio * cmp_ratio - cmp_ratio + next_end = (start_pos[i] + end_pos) // cmp_ratio * cmp_ratio + cmp_ratio + if (start_pos[i] + end_pos) % cmp_ratio == 0: + next_end = start_pos[i] + end_pos + next_end = min(next_end, start_pos[i] + end_pos) + next_start_block_id = (next_start // block_size) if next_start >= 0 else 0 + next_end_block_id = (next_end - 1) // block_size + for j in range(next_start_block_id, next_end_block_id + 1): + if block_table[i][j] == 0: + block_table[i][j] = next_block_id + next_block_id = next_block_id + 1 + + if B==0: + kv_state = torch.tensor(np.random.uniform(-10, 10, (0, block_size, coff * head_dim))).to(torch.float32) + score_state = torch.tensor(np.random.uniform(-10, 10, (0, block_size, coff * head_dim))).to(torch.float32) + else: + kv_state = torch.tensor(np.random.uniform(-10, 10, (torch.max(block_table) + 1, block_size, coff * head_dim))).to(torch.float32) + score_state = torch.tensor(np.random.uniform(-10, 10, (torch.max(block_table) + 1, block_size, coff * head_dim))).to(torch.float32) + + # other input + if bs_combine_flag: + x_shape = (cu_seqlens[-1], hidden_size) + rope_sin_shape = (min(x_shape[0], x_shape[0] // cmp_ratio + B), rope_head_dim) + rope_cos_shape = rope_sin_shape + else: + x_shape = (B, S, hidden_size) + rope_sin_shape = (B, (S + cmp_ratio - 1) // cmp_ratio, rope_head_dim) + rope_cos_shape = rope_sin_shape + + x = torch.tensor(np.random.uniform(-10.0, 10.0, x_shape)).to(data_type).npu() + wkv = torch.tensor(np.random.uniform(-10, 10, (coff * head_dim, hidden_size))).to(data_type).npu() + wgate = torch.tensor(np.random.uniform(-10, 10, (coff * head_dim, hidden_size))).to(data_type).npu() + ape = torch.tensor(np.random.uniform(-10, 10, (cmp_ratio, coff * head_dim))).to(torch.float32).npu() + norm_weight = torch.tensor(np.random.uniform(-10, 10, (head_dim))).to(data_type).npu() + rope_sin = torch.tensor(np.random.uniform(-1, 1, rope_sin_shape)).to(data_type).npu() + rope_cos = torch.tensor(np.random.uniform(-1, 1, rope_cos_shape)).to(data_type).npu() + kv_state = kv_state.npu() + score_state = score_state.npu() + block_table = block_table.npu() + start_pos = torch.tensor(start_pos).to(torch.int32).npu() + if cu_seqlens is not None: + cu_seqlens = torch.tensor(cu_seqlens).to(torch.int32).npu() + if seqused is not None: + seqused = torch.tensor(seqused).to(torch.int32).npu() + + class CompressorNetwork(nn.Module): + def __init__(self): + super(CompressorNetwork, self).__init__() + + def forward(self, x, wkv, wgate, kv_state, score_state, ape, norm_weight, rope_sin, + rope_cos, rope_head_dim, cmp_ratio, kv_block_table = None, score_block_table = None, cu_seqlens = None, + seqused = None, start_pos = None, coff = 1, norm_eps = 1e-6, rotary_mode = 1): + cmp_kv,_ ,_ ,_ ,_ = ( + torch.ops.custom.compressor( + x, + wkv, + wgate, + kv_state, + score_state, + ape, + norm_weight, + rope_sin, + rope_cos, + kv_block_table = kv_block_table, + score_block_table = score_block_table, + cu_seqlens = cu_seqlens, + seqused = seqused, + start_pos = start_pos, + rope_head_dim = rope_head_dim, + cmp_ratio = cmp_ratio, + coff = coff, + norm_eps = norm_eps, + rotary_mode = rotary_mode + ) + ) + return cmp_kv + + from torchair.configs.compiler_config import CompilerConfig + config = CompilerConfig() + npu_backend = torchair.get_npu_backend(compiler_config=config) + torch._dynamo.reset() + npu_mode = torch.compile(CompressorNetwork(), fullgraph=True, backend=npu_backend, dynamic=False) + cmp_kv = npu_mode( + x, + wkv, + wgate, + kv_state, + score_state, + ape, + norm_weight, + rope_sin, + rope_cos, + kv_block_table = block_table, + score_block_table = block_table, + cu_seqlens = cu_seqlens, + seqused = seqused, + start_pos = start_pos, + rope_head_dim = rope_head_dim, + cmp_ratio = cmp_ratio, + coff = coff, + norm_eps = norm_eps, + rotary_mode = rotary_mode) + ``` + +更多使用示例见[pytest示例](./tests/pytest/README.md)。 diff --git a/csrc/ascend/attention/compressor/op_host/CMakeLists.txt b/csrc/ascend/attention/compressor/op_host/CMakeLists.txt new file mode 100644 index 000000000..1a06c6d9f --- /dev/null +++ b/csrc/ascend/attention/compressor/op_host/CMakeLists.txt @@ -0,0 +1,40 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + compressor_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME Compressor + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true +) + +if (NOT BUILD_OPS_RTY_KERNEL) + set(SUPPORTED_ARCHS arch32 arch35) + add_modules_sources(OPTYPE compressor ACLNNTYPE aclnn) + add_tiling_modules() + + foreach(ARCH ${ARCH_DIRECTORY}) + if(ARCH IN_LIST SUPPORTED_ARCHS) + target_sources(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/${ARCH}/compressor_tiling.cpp + ) + endif() + endforeach() + +endif() diff --git a/csrc/ascend/attention/compressor/op_host/arch32/compressor_tiling.cpp b/csrc/ascend/attention/compressor/op_host/arch32/compressor_tiling.cpp new file mode 100644 index 000000000..258c0aeb7 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_host/arch32/compressor_tiling.cpp @@ -0,0 +1,1027 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file compressor_tiling.cpp +* \file compressor_tiling.cpp +* \brief +*/ + +#include +#include +#include +#include +#include +#include "err/ops_err.h" +#include "register/op_def_registry.h" +#include "compressor_tiling.h" + +using namespace ge; +using namespace AscendC; +namespace optiling { + + + +void CompressorTiling::ConvertRequiredParams(gert::TilingContext &context, CompressorContext &compressorContext) +{ + compressorContext.x.desc = context.GetRequiredInputDesc(TOKEN_X_INPUT_INDEX); + compressorContext.x.shape = context.GetRequiredInputShape(TOKEN_X_INPUT_INDEX); + compressorContext.wkv.desc = context.GetRequiredInputDesc(WEIGHT_KV_INPUT_INDEX); + compressorContext.wkv.shape = context.GetRequiredInputShape(WEIGHT_KV_INPUT_INDEX); + compressorContext.wgate.desc = context.GetRequiredInputDesc(WEIGHT_WGATE_INPUT_INDEX); + compressorContext.wgate.shape = context.GetRequiredInputShape(WEIGHT_WGATE_INPUT_INDEX); + compressorContext.stateCache.desc = context.GetRequiredInputDesc(STATE_CACHE_INPUT_INDEX); + compressorContext.stateCache.shape = context.GetRequiredInputShape(STATE_CACHE_INPUT_INDEX); + compressorContext.ape.desc = context.GetRequiredInputDesc(APE_INPUT_INDEX); + compressorContext.ape.shape = context.GetRequiredInputShape(APE_INPUT_INDEX); + compressorContext.normWeight.desc = context.GetRequiredInputDesc(NORM_WEIGHT_INPUT_INDEX); + compressorContext.normWeight.shape = context.GetRequiredInputShape(NORM_WEIGHT_INPUT_INDEX); + compressorContext.ropeSin.desc = context.GetRequiredInputDesc(ROPE_SIN_INPUT_INDEX); + compressorContext.ropeSin.shape = context.GetRequiredInputShape(ROPE_SIN_INPUT_INDEX); + compressorContext.ropeCos.desc = context.GetRequiredInputDesc(ROPE_COS_INPUT_INDEX); + compressorContext.ropeCos.shape = context.GetRequiredInputShape(ROPE_COS_INPUT_INDEX); + + compressorContext.cmpKv.desc = context.GetOutputDesc(CMP_KV_OUTPUT_INDEX); + compressorContext.cmpKv.shape = context.GetOutputShape(CMP_KV_OUTPUT_INDEX); + + compressorContext.dtype = compressorContext.x.desc->GetDataType(); + auto xDimNum = compressorContext.x.shape->GetStorageShape().GetDimNum(); + if (xDimNum == COMPRESSOR_DIM_NUM_3) { + compressorContext.layout = LayoutType::LAYOUT_BSH; + } else if (xDimNum == COMPRESSOR_DIM_NUM_2) { + compressorContext.layout = LayoutType::LAYOUT_TH; + } +} + +void CompressorTiling::ConvertOptionalParams(gert::TilingContext &context, CompressorContext &compressorContext) +{ + compressorContext.stateBlockTable.desc = context.GetOptionalInputDesc(STATE_BLOCK_TABLE_INPUT_INDEX); + compressorContext.stateBlockTable.shape = context.GetOptionalInputShape(STATE_BLOCK_TABLE_INPUT_INDEX); + compressorContext.cuSeqlens.desc = context.GetOptionalInputDesc(CU_SEQ_LEN_INPUT_INDEX); + compressorContext.cuSeqlens.shape = context.GetOptionalInputShape(CU_SEQ_LEN_INPUT_INDEX); + compressorContext.seqUsed.desc = context.GetOptionalInputDesc(SEQ_USED_INPUT_INDEX); + compressorContext.seqUsed.shape = context.GetOptionalInputShape(SEQ_USED_INPUT_INDEX); + compressorContext.startPos.desc = context.GetOptionalInputDesc(START_POS_INPUT_INDEX); + compressorContext.startPos.shape = context.GetOptionalInputShape(START_POS_INPUT_INDEX); +} + +ge::graphStatus CompressorTiling::ConvertContext(gert::TilingContext &context, CompressorContext &compressorContext) +{ + if (context.GetNodeName() == nullptr) { + OP_LOGE("Compressor", "opName got from TilingContext is nullptr"); + return ge::GRAPH_FAILED; + } + + OP_LOGI("Getting Context"); + + compressorContext.opName = context.GetNodeName(); + compressorContext.opType = context.GetNodeType(); + compressorContext.platformInfo = context.GetPlatformInfo(); + ConvertRequiredParams(context, compressorContext); + ConvertOptionalParams(context, compressorContext); + + auto attrs = context.GetAttrs(); + OP_CHECK_IF(attrs == nullptr, OP_LOGE(context.GetNodeName(), "attrs got from ge is nullptr"), + return ge::GRAPH_FAILED); + compressorContext.ropeHeadDim = attrs->GetAttrPointer(ROPE_HEAD_DIM_ATTR_INDEX); + compressorContext.coff = attrs->GetAttrPointer(COFF_ATTR_INDEX); + compressorContext.cmpRatio = attrs->GetAttrPointer(CMP_RATIO_ATTR_INDEX); + compressorContext.normEps = attrs->GetAttrPointer(NORM_EPS_ATTR_INDEX); + compressorContext.rotaryMode = attrs->GetAttrPointer(ROTARY_MODE_ATTR_INDEX); + compressorContext.cacheMode = attrs->GetAttrPointer(CACHE_MODE_ATTR_INDEX); + compressorContext.stateCacheStrideDim0 = attrs->GetAttrPointer(STATE_CACHE_STRIDE_DIM0_ATTR_INDEX); + + OP_CHECK_IF(context.GetWorkspaceSizes(1) == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(context.GetNodeName(), "workSpaceSize got from ge is nullptr"), + return ge::GRAPH_FAILED); + compressorContext.workSpaces = context.GetWorkspaceSizes(1); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::GetNpuInfo() +{ + OP_CHECK_IF(context_->platformInfo == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(context_->opName, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->platformInfo); + socVersion_ = ascendcPlatform.GetSocVersion(); + + libapiSize_ = ascendcPlatform.GetLibApiWorkSpaceSize(); + + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize_); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L1, l1Size_); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_C, l0cSize_); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_B, l0bSize_); + + aivNum_ = ascendcPlatform.GetCoreNumAiv(); + aicNum_ = ascendcPlatform.GetCoreNumAic(); + + OP_CHECK_IF(aicNum_ == 0 || aivNum_ == 0, + OPS_REPORT_VECTOR_INNER_ERR(context_->opName, "num of core obtained is 0."), return GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::SetBaseInfo() +{ + if (context_->x.shape->GetStorageShape().GetDimNum() == COMPRESSOR_DIM_NUM_3) { + baseParams_->batchSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0); + baseParams_->seqSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1); + baseParams_->hiddenSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_2); + baseParams_->tokenSize = baseParams_->batchSize * baseParams_->seqSize; + baseParams_->cgSize = context_->ropeSin.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1); + } else { + baseParams_->batchSize = context_->cuSeqlens.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0) - 1; + baseParams_->tokenSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0); + baseParams_->hiddenSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1); + baseParams_->cgSize = context_->ropeSin.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0); + } + + baseParams_->headDim = context_->normWeight.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0); + baseParams_->cmpRatio = static_cast(*context_->cmpRatio); + baseParams_->csSize = baseParams_->seqSize - (baseParams_->seqSize % baseParams_->cmpRatio); + baseParams_->ropeHeadDim = static_cast(*context_->ropeHeadDim); + baseParams_->normEps = static_cast(*context_->normEps); + baseParams_->reciprocalD = 1.0 / baseParams_->headDim; + baseParams_->cgSize = + (baseParams_->seqSize + baseParams_->cmpRatio - 1) / baseParams_->cmpRatio; // number of token after compress + baseParams_->stateCacheStrideDim0 = static_cast(*context_->stateCacheStrideDim0); + coff = static_cast(*context_->coff); + baseParams_->nSize = 2; // 2:每个核处理两个基本块后做全核同步 + + OP_LOGI(context_->opName, "[TILING] bSize:%u tSize:%u cmpRatio:%u coff:%u", baseParams_->batchSize, baseParams_->tokenSize, baseParams_->cmpRatio, coff); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::SetPageAttentionInfo() +{ + pageAttentionParams_->blockNum = context_->stateCache.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0); + pageAttentionParams_->blockSize = context_->stateCache.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1); + if (static_cast(*context_->cacheMode) == static_cast(CACHE_MODE::CONTINUOUS)) { + pageAttentionParams_->maxBlockNumPerBatch = + context_->stateBlockTable.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::SetWorkSpaceInfo() +{ + workspaceParams_->dbWorkspaceRatio = 2; + workspaceParams_->mm1KvResSize = innerSplitParams_->mBaseSize * baseParams_->headDim * coff; + workspaceParams_->mm1ScoreResSize = innerSplitParams_->mBaseSize * baseParams_->headDim * coff; + if (coff == 2) { + workspaceParams_->vec1TailCacheSize = baseParams_->cmpRatio * baseParams_->headDim; + } + if (context_->templateId == TemplateId::PERF) { + workspaceParams_->vec1ResSize = innerSplitParams_->mBaseSize * baseParams_->headDim * baseParams_->nSize; + } else { + workspaceParams_->vec1ResSize = innerSplitParams_->mBaseSize / baseParams_->cmpRatio * innerSplitParams_->dBaseSize * baseParams_->nSize; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::SetScenarioInfo() +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::SetTemplateId() +{ + if (context_->templateId == TemplateId::EMPTY_X) { + return ge::GRAPH_SUCCESS; + } + // 设置高性能模板 + context_->templateId = TemplateId::PERF; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::SetInnerSplitInfo() +{ + innerSplitParams_->mBaseSize = 256; // 256:核间切分，M轴基本块大小 + innerSplitParams_->dBaseSize = 128 / coff; // 128：核间切分，D轴基本块大小 + if (context_->templateId == TemplateId::PERF) { + if (coff == 2) { + innerSplitParams_->mBaseSize = 128; + } else { + innerSplitParams_->mBaseSize = 256; + } + innerSplitParams_->dBaseSize = 64; + } else { + innerSplitParams_->mBaseSize = 256; // 256:核间切分，M轴基本块大小 + innerSplitParams_->dBaseSize = 128 / coff; // 128：核间切分，D轴基本块大小 + } + // a5 由于loc更大, mBaseSize x 2 + // if (socVersion_ == platform_ascendc::SocVersion::ASCEND910_95) { + // innerSplitParams_->mBaseSize *= 2; + // } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CalcWorkSpace() +{ + constexpr uint32_t MM1_RES_ELEM_SIZE = 4; // 4: fp32 + constexpr uint32_t V1_RES_ELEM_SIZE = 4; // 4: fp32 + uint32_t maxGroupNum = aicNum_ / (baseParams_->headDim / innerSplitParams_->dBaseSize); + workspaceSize_ = libapiSize_; + workspaceSize_ += workspaceParams_->mm1KvResSize * maxGroupNum * MM1_RES_ELEM_SIZE * workspaceParams_->dbWorkspaceRatio; + workspaceSize_ += workspaceParams_->mm1ScoreResSize * maxGroupNum * MM1_RES_ELEM_SIZE * workspaceParams_->dbWorkspaceRatio; + workspaceSize_ += workspaceParams_->vec1TailCacheSize * MM1_RES_ELEM_SIZE * workspaceParams_->dbWorkspaceRatio * 2; // 2 kv和score + workspaceSize_ += workspaceParams_->vec1ResSize * maxGroupNum * V1_RES_ELEM_SIZE * workspaceParams_->dbWorkspaceRatio; + + if (context_->workSpaces) { + context_->workSpaces[0] = workspaceSize_; + } + + OP_LOGI(context_->opName, "Tiling info: workspaceSize_ = %zu", workspaceSize_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckEmptyTensor() const +{ + if (context_->layout == LayoutType::LAYOUT_BSH && context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0) == 0 || + context_->layout == LayoutType::LAYOUT_BSH && context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1) == 0 || + context_->layout == LayoutType::LAYOUT_TH && context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0) == 0) { + context_->templateId = TemplateId::EMPTY_X; + } else { + if (context_->x.shape->GetStorageShape().GetShapeSize() == 0 || + context_->wkv.shape->GetStorageShape().GetShapeSize() == 0 || + context_->wgate.shape->GetStorageShape().GetShapeSize() == 0 || + context_->stateCache.shape->GetStorageShape().GetShapeSize() == 0 || + context_->ape.shape->GetStorageShape().GetShapeSize() == 0 || + context_->normWeight.shape->GetStorageShape().GetShapeSize() == 0 || + context_->ropeSin.shape->GetStorageShape().GetShapeSize() == 0 || + context_->ropeCos.shape->GetStorageShape().GetShapeSize() == 0 || + context_->stateBlockTable.shape->GetStorageShape().GetShapeSize() == 0) { + OP_LOGE(context_->opName, "Only input tensor x dim B or S or T supports to be 0"); + return ge::GRAPH_FAILED; + } + context_->templateId = TemplateId::NORMAL; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::RunBigKernelTiling(CompressorTilingData* tilingData) +{ + this->baseParams_ = &tilingData->baseParams; + this->pageAttentionParams_ = &tilingData->pageAttentionParams; + this->innerSplitParams_ = &tilingData->innerSplitParams; + this->workspaceParams_ = &tilingData->workspaceParams; + using StatusFunction = std::function; + std::vector requiredTilingFuncs { + std::bind(&CompressorTiling::GetNpuInfo, this), + std::bind(&CompressorTiling::CheckRequiredParaExistence, this), + std::bind(&CompressorTiling::CheckEmptyTensor, this), + std::bind(&CompressorTiling::CheckSinglePara, this), + std::bind(&CompressorTiling::SetBaseInfo, this), + std::bind(&CompressorTiling::SetPageAttentionInfo, this), + std::bind(&CompressorTiling::CheckFeature, this), + std::bind(&CompressorTiling::CheckMultiParaConsistency, this), + std::bind(&CompressorTiling::CheckBlockDimConstrain, this), + std::bind(&CompressorTiling::SetTemplateId, this), + std::bind(&CompressorTiling::SetInnerSplitInfo, this), + std::bind(&CompressorTiling::SetWorkSpaceInfo, this), + std::bind(&CompressorTiling::SetScenarioInfo, this) + }; + for (const auto &func: requiredTilingFuncs) { + if (func() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + } + + if (context_->templateId == TemplateId::EMPTY_X) { + workspaceSize_ = libapiSize_; + if (context_->workSpaces) { + context_->workSpaces[0] = workspaceSize_; + } + GenTilingKey(); + context_->blockDim = 1U; + return ge::GRAPH_SUCCESS; + } + std::vector optionalTilingFuncs { + std::bind(&CompressorTiling::CalcWorkSpace, this), + std::bind(&CompressorTiling::GenTilingKey, this) + }; + for (const auto &func : optionalTilingFuncs) { + if (func() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + } + + baseParams_->usedCoreNum = aicNum_; + + context_->blockDim = aicNum_; + + OP_LOGI("Run big kernel"); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::GenTilingKey() const +{ + // 0:BF16, 1:FP16 + uint8_t dtype = 0; + // 0: BSH 1:TH + uint8_t layout = 0; + uint8_t ropeDtype = 0; + uint8_t rotaryMode = static_cast(*context_->rotaryMode); + uint8_t templateId = static_cast(context_->templateId); + uint8_t cacheMode = static_cast(*context_->cacheMode); + + auto xDtype = context_->x.desc->GetDataType(); + if (xDtype == ge::DT_BF16) { + dtype = 0; + } else if (xDtype == ge::DT_FLOAT16) { + dtype = 1; + } + auto ropeSinDtype = context_->ropeSin.desc->GetDataType(); + auto ropeCosDtype = context_->ropeCos.desc->GetDataType(); + bool supportFp32Rope = socVersion_ == platform_ascendc::SocVersion::ASCEND910B || + socVersion_ == platform_ascendc::SocVersion::ASCEND910_93; + if (ropeSinDtype == ge::DT_FLOAT && ropeCosDtype == ge::DT_FLOAT && supportFp32Rope) { + ropeDtype = 1; + } + auto xDimNum = context_->x.shape->GetStorageShape().GetDimNum(); + if (xDimNum == COMPRESSOR_DIM_NUM_3) { + layout = 0; + } else { + layout = 1; + } + + context_->tilingKey = GET_TPL_TILING_KEY( + layout, + dtype, + coff, + rotaryMode, + cacheMode, + templateId, + ropeDtype + ); + OP_LOGI(context_->opName, + "Compressor dtype:%hhu layout:%hhu coff:%hhu rotary_mode:%hhu, cacheMode: %u, template_id:%hhu, rope_dtype:%hhu", + dtype, layout, coff, rotaryMode, cacheMode, templateId, ropeDtype); + OP_LOGI(context_->opName, "Compressor tilingKey:%lu", context_->tilingKey); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSinglePara() const +{ + if (ge::GRAPH_SUCCESS != CheckSingleParaX() || + ge::GRAPH_SUCCESS != CheckSingleParaWkv() || + ge::GRAPH_SUCCESS != CheckSingleParaWgate() || + ge::GRAPH_SUCCESS != CheckSingleParaStateCache() || + ge::GRAPH_SUCCESS != CheckSingleParaApe() || + ge::GRAPH_SUCCESS != CheckSingleParaNormWeight() || + ge::GRAPH_SUCCESS != CheckSingleParaRopeSin() || + ge::GRAPH_SUCCESS != CheckSingleParaRopeCos() || + ge::GRAPH_SUCCESS != CheckSingleParaStateBlockTable() || + ge::GRAPH_SUCCESS != CheckSingleParaCuSeqlens() || + ge::GRAPH_SUCCESS != CheckSingleParaSeqused() || + ge::GRAPH_SUCCESS != CheckSingleParaStartPos() || + ge::GRAPH_SUCCESS != CheckSingleParaCmpKv() || + ge::GRAPH_SUCCESS != CheckSingleParaRopeHeadDim() || + ge::GRAPH_SUCCESS != CheckSingleParaCmpRatio() || + ge::GRAPH_SUCCESS != CheckSingleParaCoff() || + ge::GRAPH_SUCCESS != CheckSingleParaNormEps() || + ge::GRAPH_SUCCESS != CheckSingleParaRotaryMode() || + ge::GRAPH_SUCCESS != CheckSingleParaCacheMode()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +template +ge::graphStatus CompressorTiling::CheckFeatureValueSupport(const T *featureValue, + const std::vector &expectFeatureValList, const std::string &name) const +{ + if (std::find(expectFeatureValList.begin(), expectFeatureValList.end(), *featureValue) == expectFeatureValList.end()) { + LogErrorNumberSupport(expectFeatureValList, *featureValue, name, "feature value"); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +template +ge::graphStatus CompressorTiling::CheckAttrValueSupport(const T *attrValue, + const std::vector &expectAttrValList, const std::string &name) const +{ + if (attrValue == nullptr) { + return ge::GRAPH_SUCCESS; + } + + if (std::find(expectAttrValList.begin(), expectAttrValList.end(), *attrValue) == expectAttrValList.end()) { + LogErrorNumberSupport(expectAttrValList, *attrValue, name, "attr value"); + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +template +std::string to_string(const T &value) { + if (std::is_same_v) { + return value ? "true" : "false"; + } else { + return std::to_string(value); + } +} + +template +void CompressorTiling::LogErrorNumberSupport(const std::vector &expectNumberList, + const T &actualValue, const std::string &name, const std::string subName) const +{ + std::ostringstream oss; + for (size_t i = 0; i < expectNumberList.size(); ++i) { + oss << to_string(expectNumberList[i]); + if (i < expectNumberList.size() - 1) { + oss << ", "; + } + } + + OP_LOGE(context_->opName, "%s %s only supports %s, but got %s", + name.c_str(), subName.c_str(), oss.str().c_str(), to_string(actualValue).c_str()); +} + +std::string LayoutTypeToStr(LayoutType layout) +{ + switch (layout) { + case LayoutType::LAYOUT_BSH: + return "BSH"; + case LayoutType::LAYOUT_TH: + return "TH"; + default: + return "UNKNOWN_LAYOUT"; + } +} + +ge::graphStatus CompressorTiling::CheckDimNumInLayoutSupport(const std::string &layout, const gert::StorageShape *shape, + const std::string &name) const +{ + const auto& dimIt = LAYOUT_DIM_MAP.find(layout); + OP_CHECK_IF(shape->GetStorageShape().GetDimNum() != dimIt->second, + OP_LOGE(context_->opName, "When layout is %s, %s dimension should be %zu, but it's %zu", + layout.c_str(), name.c_str(), dimIt->second, + shape->GetStorageShape().GetDimNum()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc, + const std::string &name) const +{ + if (desc != nullptr) { + const auto &it = DTYPE_SUPPORT_MAP.find(name); + OP_CHECK_IF(it == DTYPE_SUPPORT_MAP.end(), + OP_LOGE(context_->opName, "%s datatype support list should be specify in DTYPE_SUPPORT_MAP", name.c_str()), + return ge::GRAPH_FAILED); + auto &expectDtypeList = it->second; + OP_CHECK_IF(std::find(expectDtypeList.begin(), expectDtypeList.end(), desc->GetDataType()) == + expectDtypeList.end(), + LogErrorDtypeSupport(expectDtypeList, desc->GetDataType(), name), return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +void CompressorTiling::LogErrorDtypeSupport(const std::vector &expectDtypeList, + const ge::DataType &actualDtype, const std::string &name) const +{ + std::ostringstream oss; + for (size_t i = 0; i < expectDtypeList.size(); ++i) { + oss << DataTypeToSerialString(expectDtypeList[i]); + if (i < expectDtypeList.size() - 1) { + oss << ", "; + } + } + OP_LOGE(context_->opName, "Tensor %s only supports dtype %s, but got %s", name.c_str(), oss.str().c_str(), + DataTypeToSerialString(actualDtype).c_str()); +} + +static std::string DataTypeToSerialString(ge::DataType type) +{ + const auto it = DATATYPE_TO_STRING_MAP.find(type); + if (it != DATATYPE_TO_STRING_MAP.end()) { + return it->second; + } else { + OP_LOGE("Compressor", "datatype %d not support", type); + return "UNDEFINED"; + } +} + +ge::graphStatus CompressorTiling::CheckDimNumSupport(const gert::StorageShape *shape, const std::string &name) const +{ + if (shape == nullptr) { + return ge::GRAPH_SUCCESS; + } + const auto &it = DIM_NUM_MAP.find(name); + OP_CHECK_IF(it == DIM_NUM_MAP.end(), + OP_LOGE(context_->opName, "%s dim number support list should be specify in DIM_NUM_MAP", name.c_str()), + return ge::GRAPH_FAILED); + auto &expectDimNumList = it->second; + OP_CHECK_IF(std::find(expectDimNumList.begin(), expectDimNumList.end(), shape->GetStorageShape().GetDimNum()) == + expectDimNumList.end(), + LogErrorNumberSupport(expectDimNumList, static_cast(shape->GetStorageShape().GetDimNum()), + name, "dimension"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaX() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->x.desc, X_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->x.shape, X_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(LayoutTypeToStr(context_->layout), context_->x.shape, X_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaWkv() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->wkv.desc, WKV_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->wkv.shape, WKV_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaWgate() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->wgate.desc, WGATE_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->wgate.shape, WGATE_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaStateCache() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->stateCache.desc, STATE_CACHE_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->stateCache.shape, STATE_CACHE_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaApe() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->ape.desc, APE_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->ape.shape, APE_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaNormWeight() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->normWeight.desc, NORM_WEIGHT_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->normWeight.shape, NORM_WEIGHT_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaRopeSin() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->ropeSin.desc, ROPE_SIN_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->ropeSin.shape, ROPE_SIN_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(LayoutTypeToStr(context_->layout), context_->ropeSin.shape, ROPE_SIN_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaRopeCos() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->ropeCos.desc, ROPE_COS_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->ropeCos.shape, ROPE_COS_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(LayoutTypeToStr(context_->layout), context_->ropeCos.shape, ROPE_COS_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaStateBlockTable() const +{ + if (context_->stateBlockTable.desc == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->stateBlockTable.desc, STATE_BLOCK_TABLE_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->stateBlockTable.shape, STATE_BLOCK_TABLE_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaCuSeqlens() const +{ + if (context_->cuSeqlens.desc == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->cuSeqlens.desc, CU_SEQLENS_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->cuSeqlens.shape, CU_SEQLENS_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaSeqused() const +{ + if (context_->seqUsed.desc == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->seqUsed.desc, SEQUSED_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->seqUsed.shape, SEQUSED_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaStartPos() const +{ + if (context_->startPos.desc == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->startPos.desc, START_POS_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->startPos.shape, START_POS_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaCmpKv() const +{ + if (context_->cmpKv.desc == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->cmpKv.desc, CMP_KV_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->cmpKv.shape, CMP_KV_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaRopeHeadDim()const +{ + if (CheckAttrValueSupport(context_->ropeHeadDim, ROPE_HEAD_DIM, ROPE_HEAD_DIM_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaCmpRatio()const +{ + if (CheckAttrValueSupport(context_->cmpRatio, CMP_RATIO, CMP_RATIO_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaCoff()const +{ + if (CheckAttrValueSupport(context_->coff, COFF, COFF_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaNormEps()const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaRotaryMode()const +{ + if (ge::GRAPH_SUCCESS != CheckAttrValueSupport(context_->rotaryMode, ROTARY_MODE, ROTARY_MODE_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaCacheMode() const +{ + if (ge::GRAPH_SUCCESS != CheckAttrValueSupport(context_->cacheMode, CACHE_MODE, CACHE_MODE_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckRequiredParaExistence() const +{ + if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS || CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckRequiredInOutExistence() const +{ + OP_CHECK_IF(context_->x.shape == nullptr, OP_LOGE(context_->opName, "tensor x is nullptr"), return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->x.desc == nullptr, OP_LOGE(context_->opName, "tensor x is nullptr"), return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->wkv.shape == nullptr, OP_LOGE(context_->opName, "tensor wkv is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->wkv.desc == nullptr, OP_LOGE(context_->opName, "tensor wkv is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->wgate.shape == nullptr, OP_LOGE(context_->opName, "tensor wgate is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->wgate.desc == nullptr, OP_LOGE(context_->opName, "tensor wgate is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->stateCache.shape == nullptr, OP_LOGE(context_->opName, "tensor stateCache is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->stateCache.desc == nullptr, OP_LOGE(context_->opName, "tensor stateCache is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->ape.shape == nullptr, OP_LOGE(context_->opName, "tensor ape is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->ape.desc == nullptr, OP_LOGE(context_->opName, "tensor ape is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->normWeight.shape == nullptr, OP_LOGE(context_->opName, "tensor normWeight is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->normWeight.desc == nullptr, OP_LOGE(context_->opName, "tensor normWeight is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->ropeSin.shape == nullptr, OP_LOGE(context_->opName, "tensor ropeSin is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->ropeSin.desc == nullptr, OP_LOGE(context_->opName, "tensor ropeSin is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->ropeCos.shape == nullptr, OP_LOGE(context_->opName, "tensor ropeCos is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->ropeCos.desc == nullptr, OP_LOGE(context_->opName, "tensor ropeCos is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->stateBlockTable.shape == nullptr, + OP_LOGE(context_->opName, "tensor stateBlockTable is nullptr"), return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->stateBlockTable.desc == nullptr, + OP_LOGE(context_->opName, "tensor stateBlockTable is nullptr"), return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->cmpKv.shape == nullptr, OP_LOGE(context_->opName, "tensor cmpKv is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->cmpKv.desc == nullptr, OP_LOGE(context_->opName, "tensor cmpKv is nullptr"), + return ge::GRAPH_FAILED); + if (context_->layout == LayoutType::LAYOUT_TH) { + OP_CHECK_IF(context_->cuSeqlens.desc == nullptr, + OP_LOGE(context_->opName, "In TH layout, tensor cuSeqlens should not be nullptr"), return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->cuSeqlens.shape == nullptr, + OP_LOGE(context_->opName, "In TH layout, tensor cuSeqlens should not be nullptr"), return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF(context_->cuSeqlens.desc != nullptr, + OP_LOGE(context_->opName, "In BSH layout, tensor cuSeqlens must be nullptr"), return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->cuSeqlens.shape != nullptr, + OP_LOGE(context_->opName, "In TH layout, tensor cuSeqlens must be nullptr"), return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckRequiredAttrExistence() const +{ + OP_CHECK_IF(context_->ropeHeadDim == nullptr, OP_LOGE(context_->opName, "attr ropeHeadDim is nullptr"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(context_->cmpRatio == nullptr, OP_LOGE(context_->opName, "attr cmpRatio is nullptr"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckFeature() const +{ + if (ge::GRAPH_SUCCESS != CheckFeatureValueSupport(&baseParams_->headDim, HEAD_DIM, "headDim")) { + return ge::GRAPH_FAILED; + } + OP_CHECK_IF(baseParams_->hiddenSize > MAX_HIDDEN_SIZE || baseParams_->hiddenSize < MIN_HIDDEN_SIZE || + baseParams_->hiddenSize % ALIGN_FACTOR_HIDDEN_SIZE != 0, + OP_LOGE(context_->opName, "hiddenSize should be whthin [1k, 10k] and be 512-aligned, but got %u", + baseParams_->hiddenSize), + return ge::GRAPH_FAILED); + OP_CHECK_IF(pageAttentionParams_->blockSize < MIN_BLOCK_SIZE, + OP_LOGE(context_->opName, "blockSize should not be less than 1, but got %u", + pageAttentionParams_->blockSize), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::LogErrorShapeConsistency(const std::string &name, + const gert::StorageShape *shape, const uint32_t &dimNum, const std::string &subName, const uint32_t &expectNum) const +{ + if (shape == nullptr) { + return ge::GRAPH_SUCCESS; + } + + const uint32_t actualNum = shape->GetStorageShape().GetDim(dimNum); + OP_CHECK_IF(actualNum != expectNum, + OP_LOGE(context_->opName, + "%s shape dim %u, should be equal to %s: %u, but got %u", + name.c_str(), dimNum, subName.c_str(), expectNum, actualNum), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckShapeConsistency() const +{ + if (CheckShapeConsistencyRope() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + auto coffD = coff * baseParams_->headDim; + uint32_t stateNum = 2; + if (ge::GRAPH_SUCCESS != LogErrorShapeConsistency("stateBlockTable", context_->stateBlockTable.shape, + COMPRESSOR_DIM_INDEX_0, "batchSize", baseParams_->batchSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("cuSeqlens", context_->cuSeqlens.shape, COMPRESSOR_DIM_INDEX_0, + "batchSize+1", baseParams_->batchSize + 1) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("seqUsed", context_->seqUsed.shape, COMPRESSOR_DIM_INDEX_0, + "batchSize", baseParams_->batchSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("startPos", context_->startPos.shape, COMPRESSOR_DIM_INDEX_0, + "batchSize", baseParams_->batchSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("wkv", context_->wkv.shape, COMPRESSOR_DIM_INDEX_1, "hiddenSize", + baseParams_->hiddenSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("wgate", context_->wgate.shape, COMPRESSOR_DIM_INDEX_1, + "hiddenSize", baseParams_->hiddenSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("wkv", context_->wkv.shape, COMPRESSOR_DIM_INDEX_0, + "coff*headDim", static_cast(coffD)) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("wgate", context_->wgate.shape, COMPRESSOR_DIM_INDEX_0, + "coff*headDim", static_cast(coffD)) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("stateCache", context_->stateCache.shape, COMPRESSOR_DIM_INDEX_2, + "2*coff*headDim", stateNum * static_cast(coffD)) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ape", context_->ape.shape, COMPRESSOR_DIM_INDEX_1, + "coff*headDim", static_cast(coffD)) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ape", context_->ape.shape, COMPRESSOR_DIM_INDEX_0, "cmpRatio", + baseParams_->cmpRatio)) { + return ge::GRAPH_FAILED; + } + if (static_cast(*context_->cacheMode) == static_cast(CACHE_MODE::CONTINUOUS) && + (ge::GRAPH_SUCCESS != LogErrorShapeConsistency("stateCache", context_->stateCache.shape, COMPRESSOR_DIM_INDEX_0, + "blockNum", pageAttentionParams_->blockNum) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("stateCache", context_->stateCache.shape, COMPRESSOR_DIM_INDEX_1, + "blockSize", pageAttentionParams_->blockSize))) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckShapeConsistencyRope() const +{ + auto cmpT = std::min(baseParams_->tokenSize, baseParams_->tokenSize / baseParams_->cmpRatio + baseParams_->batchSize); + if (context_->layout == LayoutType::LAYOUT_BSH) { + if (ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeSin", context_->ropeSin.shape, COMPRESSOR_DIM_INDEX_0, "batchSize", baseParams_->batchSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeCos", context_->ropeCos.shape, COMPRESSOR_DIM_INDEX_0, "batchSize", baseParams_->batchSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeSin", context_->ropeSin.shape, COMPRESSOR_DIM_INDEX_1, "ceil(seqSize/cmpRatio)", baseParams_->cgSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeCos", context_->ropeCos.shape, COMPRESSOR_DIM_INDEX_1, "ceil(seqSize/cmpRatio)", baseParams_->cgSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeSin", context_->ropeSin.shape, COMPRESSOR_DIM_INDEX_2, "ropeHeadDim", baseParams_->ropeHeadDim) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeCos", context_->ropeCos.shape, COMPRESSOR_DIM_INDEX_2, "ropeHeadDim", baseParams_->ropeHeadDim)) { + return ge::GRAPH_FAILED; + } + } else { + if (ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeSin", context_->ropeSin.shape, COMPRESSOR_DIM_INDEX_0, "min(tokenSize, tokenSize/cmpRatio+batchSize)", static_cast(cmpT)) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeCos", context_->ropeCos.shape, COMPRESSOR_DIM_INDEX_0, "min(tokenSize, tokenSize/cmpRatio+batchSize)", static_cast(cmpT)) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeSin", context_->ropeSin.shape, COMPRESSOR_DIM_INDEX_1, "ropeHeadDim", baseParams_->ropeHeadDim) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeCos", context_->ropeCos.shape, COMPRESSOR_DIM_INDEX_1, "ropeHeadDim", baseParams_->ropeHeadDim)) { + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckDtypeConsistencyX(const gert::CompileTimeTensorDesc *desc, + const std::string &name) const +{ + const auto actualDtype = desc->GetDataType(); + OP_CHECK_IF( + actualDtype != context_->dtype, + OP_LOGE(context_->opName, "%s datatype should be same with x: %s, but got %s", name.c_str(), + DataTypeToSerialString(actualDtype).c_str(), DataTypeToSerialString(context_->dtype).c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckDtypeConsistencyRope() const +{ + auto sinDtype = context_->ropeSin.desc->GetDataType(); + auto cosDtype = context_->ropeCos.desc->GetDataType(); + OP_CHECK_IF( + sinDtype != cosDtype, + OP_LOGE(context_->opName, "%s datatype should be same with %s: %s, but got %s", ROPE_COS_NAME.c_str(), + ROPE_SIN_NAME.c_str(), DataTypeToSerialString(sinDtype).c_str(), + DataTypeToSerialString(cosDtype).c_str()), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + sinDtype != context_->dtype && sinDtype != ge::DT_FLOAT, + OP_LOGE(context_->opName, "rope datatype should be same with x or DT_FLOAT, x is %s, but got %s", + DataTypeToSerialString(context_->dtype).c_str(), DataTypeToSerialString(sinDtype).c_str()), + return ge::GRAPH_FAILED); + bool supportFp32Rope = socVersion_ == platform_ascendc::SocVersion::ASCEND910B || + socVersion_ == platform_ascendc::SocVersion::ASCEND910_93; + OP_CHECK_IF( + sinDtype == ge::DT_FLOAT && !supportFp32Rope, + OP_LOGE(context_->opName, "float32 rope is only enabled on ascend910b and ascend910_93."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckDtypeConsistency() const +{ + if (CheckDtypeConsistencyX(context_->wkv.desc, WKV_NAME) != ge::GRAPH_SUCCESS || + CheckDtypeConsistencyX(context_->wgate.desc, WGATE_NAME) != ge::GRAPH_SUCCESS || + CheckDtypeConsistencyX(context_->normWeight.desc, NORM_WEIGHT_NAME) != ge::GRAPH_SUCCESS || + CheckDtypeConsistencyRope() != ge::GRAPH_SUCCESS || + CheckDtypeConsistencyX(context_->cmpKv.desc, CMP_KV_NAME) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckDimNumConsistency() const +{ + auto xDimNum = context_->x.shape->GetStorageShape().GetDimNum(); + OP_CHECK_IF(xDimNum != context_->ropeSin.shape->GetStorageShape().GetDimNum(), + OP_LOGE(context_->opName, "ropeSin dim num should be equal to x: %u, but got %u", xDimNum, + context_->ropeSin.shape->GetStorageShape().GetDimNum()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(xDimNum != context_->ropeCos.shape->GetStorageShape().GetDimNum(), + OP_LOGE(context_->opName, "ropeCos dim num should be equal to x: %u, but got %u", xDimNum, + context_->ropeCos.shape->GetStorageShape().GetDimNum()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(xDimNum != context_->cmpKv.shape->GetStorageShape().GetDimNum(), + OP_LOGE(context_->opName, "cmpKv dim num should be equal to x: %u, but got %u", xDimNum, + context_->cmpKv.shape->GetStorageShape().GetDimNum()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckScenarioConsistency() const +{ + auto curCmpratio = baseParams_->cmpRatio; + auto curHeaddim = baseParams_->headDim; + auto curCoff = static_cast(*context_->coff); + std::vector curScenario{curCmpratio, curCoff, curHeaddim}; + const std::vector> allowdScenarios = {{4, 2, 512}, {4, 2, 128}, {128, 1, 512}}; + + OP_CHECK_IF(std::find(allowdScenarios.begin(), allowdScenarios.end(), curScenario) == allowdScenarios.end(), + OP_LOGE(context_->opName, "Cmpratio Coff Headdim should be equal to {4, 2, 512}, {4, 2, 128}, {128, 1, 512},\ + but now cmpratio=%u, coff=%u, headdim=%u", curCmpratio, curCoff, curHeaddim), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckBlockDimConstrain() const +{ + uint32_t minBlockNum = baseParams_->headDim / 64; // 64 is the largest dBaseSize + OP_CHECK_IF(aicNum_ < minBlockNum, OP_LOGE(context_->opName, "aicNum is %d, which should not be less than %d", + aicNum_, minBlockNum), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckMultiParaConsistency() const +{ + if (CheckShapeConsistency() != ge::GRAPH_SUCCESS || CheckDtypeConsistency() != ge::GRAPH_SUCCESS || + CheckDimNumConsistency() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } +#ifdef DAY0_SCOPE + if (CheckScenarioConsistency() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } +#endif + return ge::GRAPH_SUCCESS; +} + +CMP_EXTERN_C ge::graphStatus TilingCompressor(gert::TilingContext *context) +{ + OP_CHECK_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("Compressor", "Context is nullptr."), + return ge::GRAPH_FAILED); + + OP_LOGI("Getting Tiling"); + + CompressorContext compressorContext{}; + if (CompressorTiling::ConvertContext(*context, compressorContext) != ge::GRAPH_SUCCESS) { + OP_LOGE(context->GetNodeName(), "Error occurred while converting tilingContext to Compressor context"); + return ge::GRAPH_FAILED; + } + CompressorTiling compressorTiling(&compressorContext); + CompressorTilingData* tilingData = context->GetTilingData(); + OP_CHECK_IF(tilingData == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(compressorContext.opName, "TilingData is nullptr."), + return ge::GRAPH_FAILED); + // 使用SyncAll，需要设置为batchmode模式，所有核同时启动，否则多流方式下执行可能会卡死 + context->SetScheduleMode(BATCH_MODE_SCHEDULE); + if (compressorTiling.RunBigKernelTiling(tilingData) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + context->SetTilingKey(compressorContext.tilingKey); + context->SetBlockDim(compressorContext.blockDim); + OP_LOGI(compressorContext.opName, "block dim: %u.", compressorContext.blockDim); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingPrepareForCompressor(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(Compressor) + .Tiling(TilingCompressor) + .TilingParse(TilingPrepareForCompressor); +} // namespace optiling diff --git a/csrc/ascend/attention/compressor/op_host/arch32/compressor_tiling.h b/csrc/ascend/attention/compressor/op_host/arch32/compressor_tiling.h new file mode 100644 index 000000000..3fff53a5d --- /dev/null +++ b/csrc/ascend/attention/compressor/op_host/arch32/compressor_tiling.h @@ -0,0 +1,381 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_tiling.h + * \brief + */ + +#ifndef COMPRESSOR_TILING_H +#define COMPRESSOR_TILING_H + +#include +#include +#include +#include +#include +#include +#include +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "exe_graph/runtime/tiling_context.h" +#include "register/op_def_registry.h" +#include "../../op_kernel/arch32/compressor_template_tiling_key.h" +#include "../../op_kernel/arch32/compressor_tiling_data.h" +#include "platform/platform_info.h" + +#ifdef ASCENDC_OP_TEST +#define CMP_EXTERN_C extern "C" +#else +#define CMP_EXTERN_C +#endif +// #define DAY0_SCOPE + +namespace optiling { + +// INPUT +constexpr uint32_t TOKEN_X_INPUT_INDEX = 0; +constexpr uint32_t WEIGHT_KV_INPUT_INDEX = 1; +constexpr uint32_t WEIGHT_WGATE_INPUT_INDEX = 2; +constexpr uint32_t STATE_CACHE_INPUT_INDEX = 3; +constexpr uint32_t APE_INPUT_INDEX = 4; +constexpr uint32_t NORM_WEIGHT_INPUT_INDEX = 5; +constexpr uint32_t ROPE_SIN_INPUT_INDEX = 6; +constexpr uint32_t ROPE_COS_INPUT_INDEX = 7; + +// INPUT(OPTION) +constexpr uint32_t STATE_BLOCK_TABLE_INPUT_INDEX = 8; +constexpr uint32_t CU_SEQ_LEN_INPUT_INDEX = 9; +constexpr uint32_t SEQ_USED_INPUT_INDEX = 10; +constexpr uint32_t START_POS_INPUT_INDEX = 11; + +// ATTR +constexpr uint32_t ROPE_HEAD_DIM_ATTR_INDEX = 0; +constexpr uint32_t CMP_RATIO_ATTR_INDEX = 1; +constexpr uint32_t COFF_ATTR_INDEX = 2; +constexpr uint32_t NORM_EPS_ATTR_INDEX = 3; +constexpr uint32_t ROTARY_MODE_ATTR_INDEX = 4; +constexpr uint32_t CACHE_MODE_ATTR_INDEX = 5; +constexpr uint32_t STATE_CACHE_STRIDE_DIM0_ATTR_INDEX = 6; + +// OUTPUT +constexpr uint32_t CMP_KV_OUTPUT_INDEX = 0; + +constexpr uint32_t COMPRESSOR_DIM_NUM_1 = 1; +constexpr uint32_t COMPRESSOR_DIM_NUM_2 = 2; +constexpr uint32_t COMPRESSOR_DIM_NUM_3 = 3; +constexpr uint32_t COMPRESSOR_DIM_NUM_4 = 4; +constexpr uint32_t COMPRESSOR_DIM_INDEX_0 = 0; +constexpr uint32_t COMPRESSOR_DIM_INDEX_1 = 1; +constexpr uint32_t COMPRESSOR_DIM_INDEX_2 = 2; +constexpr uint32_t COMPRESSOR_DIM_INDEX_3 = 3; + +// CONSTRAINTS +constexpr uint32_t MAX_HIDDEN_SIZE = 10240; +constexpr uint32_t MIN_HIDDEN_SIZE = 1024; +constexpr uint32_t ALIGN_FACTOR_HIDDEN_SIZE = 512; +constexpr uint32_t MIN_BLOCK_SIZE = 1; + +constexpr uint32_t BATCH_MODE_SCHEDULE = 1; + +static const std::string X_NAME = "query"; +static const std::string WKV_NAME = "wkv"; +static const std::string WGATE_NAME = "wgate"; +static const std::string STATE_CACHE_NAME = "state_cache"; +static const std::string APE_NAME = "ape"; +static const std::string NORM_WEIGHT_NAME = "norm_weight"; +static const std::string ROPE_SIN_NAME = "rope_sin"; +static const std::string ROPE_COS_NAME = "rope_cos"; +static const std::string STATE_BLOCK_TABLE_NAME = "state_block_table"; +static const std::string CU_SEQLENS_NAME = "cu_seqlens"; +static const std::string SEQUSED_NAME = "seq_used"; +static const std::string START_POS_NAME = "start_pos"; +static const std::string ROPE_HEAD_DIM_NAME = "rope_head_dim"; +static const std::string CMP_RATIO_NAME = "cmp_ratio"; +static const std::string COFF_NAME = "coff"; +static const std::string NORM_EPS_NAME = "nrom_eps"; +static const std::string ROTARY_MODE_NAME = "rotary_mode"; +static const std::string CACHE_MODE_NAME = "cache_mode"; +static const std::string CMP_KV_NAME = "cmp_kv"; + +static std::string DataTypeToSerialString(ge::DataType type); + +const std::map> DTYPE_SUPPORT_MAP = { + {X_NAME, {ge::DT_BF16, ge::DT_FLOAT16}}, + {WKV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}}, + {WGATE_NAME, {ge::DT_BF16, ge::DT_FLOAT16}}, + {STATE_CACHE_NAME, {ge::DT_FLOAT}}, + {APE_NAME, {ge::DT_FLOAT}}, + {NORM_WEIGHT_NAME, {ge::DT_BF16, ge::DT_FLOAT16}}, + {ROPE_SIN_NAME, {ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT}}, + {ROPE_COS_NAME, {ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT}}, + {STATE_BLOCK_TABLE_NAME, {ge::DT_INT32}}, + {CU_SEQLENS_NAME, {ge::DT_INT32}}, + {SEQUSED_NAME, {ge::DT_INT32}}, + {START_POS_NAME, {ge::DT_INT32}}, + {CMP_KV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}} +}; + +const std::map> DIM_NUM_MAP = { + {X_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}}, + {WKV_NAME, {COMPRESSOR_DIM_NUM_2}}, + {WGATE_NAME, {COMPRESSOR_DIM_NUM_2}}, + {STATE_CACHE_NAME, {COMPRESSOR_DIM_NUM_3}}, + {APE_NAME, {COMPRESSOR_DIM_NUM_2}}, + {NORM_WEIGHT_NAME, {COMPRESSOR_DIM_NUM_1}}, + {ROPE_SIN_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}}, + {ROPE_COS_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}}, + {STATE_BLOCK_TABLE_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_1}}, + {CU_SEQLENS_NAME, {COMPRESSOR_DIM_NUM_1}}, + {SEQUSED_NAME, {COMPRESSOR_DIM_NUM_1}}, + {START_POS_NAME, {COMPRESSOR_DIM_NUM_1}}, + {CMP_KV_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}} +}; + +static const std::map LAYOUT_DIM_MAP = { + {"BSH", COMPRESSOR_DIM_NUM_3}, + {"TH", COMPRESSOR_DIM_NUM_2}, +}; + +const std::map DATATYPE_TO_STRING_MAP = { + {ge::DT_UNDEFINED, "DT_UNDEFINED"}, // Used to indicate a DataType field has not been set. + {ge::DT_FLOAT, "DT_FLOAT"}, // float type + {ge::DT_FLOAT16, "DT_FLOAT16"}, // fp16 type + {ge::DT_INT8, "DT_INT8"}, // int8 type + {ge::DT_INT16, "DT_INT16"}, // int16 type + {ge::DT_UINT16, "DT_UINT16"}, // uint16 type + {ge::DT_UINT8, "DT_UINT8"}, // uint8 type + {ge::DT_INT32, "DT_INT32"}, // uint32 type + {ge::DT_INT64, "DT_INT64"}, // int64 type + {ge::DT_UINT32, "DT_UINT32"}, // unsigned int32 + {ge::DT_UINT64, "DT_UINT64"}, // unsigned int64 + {ge::DT_BOOL, "DT_BOOL"}, // bool type + {ge::DT_DOUBLE, "DT_DOUBLE"}, // double type + {ge::DT_DUAL, "DT_DUAL"}, // dual output type + {ge::DT_DUAL_SUB_INT8, "DT_DUAL_SUB_INT8"}, // dual output int8 type + {ge::DT_DUAL_SUB_UINT8, "DT_DUAL_SUB_UINT8"}, // dual output uint8 type + {ge::DT_COMPLEX32, "DT_COMPLEX32"}, // complex32 type + {ge::DT_COMPLEX64, "DT_COMPLEX64"}, // complex64 type + {ge::DT_COMPLEX128, "DT_COMPLEX128"}, // complex128 type + {ge::DT_QINT8, "DT_QINT8"}, // qint8 type + {ge::DT_QINT16, "DT_QINT16"}, // qint16 type + {ge::DT_QINT32, "DT_QINT32"}, // qint32 type + {ge::DT_QUINT8, "DT_QUINT8"}, // quint8 type + {ge::DT_QUINT16, "DT_QUINT16"}, // quint16 type + {ge::DT_RESOURCE, "DT_RESOURCE"}, // resource type + {ge::DT_STRING_REF, "DT_STRING_REF"}, // string ref type + {ge::DT_STRING, "DT_STRING"}, // string type + {ge::DT_VARIANT, "DT_VARIANT"}, // dt_variant type + {ge::DT_BF16, "DT_BFLOAT16"}, // dt_bfloat16 type + {ge::DT_INT4, "DT_INT4"}, // dt_variant type + {ge::DT_UINT1, "DT_UINT1"}, // dt_variant type + {ge::DT_INT2, "DT_INT2"}, // dt_variant type + {ge::DT_UINT2, "DT_UINT2"} // dt_variant type +}; + +struct CompressorCompileInfo { + int64_t core_num; +}; + +struct RequiredParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; +}; + +struct OptionalParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; + const gert::Tensor *tensor; +}; + +enum class LayoutType { + LAYOUT_BSH, + LAYOUT_TH +}; + +enum class TemplateId:uint8_t { + NORMAL = 0, + EMPTY_X = 1, + PERF = 2 +}; + +CMP_EXTERN_C ge::graphStatus TilingCompressor(gert::TilingContext *context); +struct CompressorBaseShapeInfo { + uint32_t bSize = 0; // B + uint32_t sSize = 0; // S + uint32_t hSize = 0; // Hidden size + uint32_t tSize = 0; // T + uint32_t nSize = 0; // N + uint32_t dSize = 0; // D + uint32_t coffSize = 0; // Coff: 1 or 2 + uint32_t csSize = 0; // Compress sequence len + uint32_t rSize = 0; // Compress ratio + uint32_t cgSize = 0; // Compress group size + uint32_t drSize = 0; // Dr +}; + +const std::vector ROPE_HEAD_DIM {64}; +const std::vector COFF {1, 2}; +#ifdef DAY0_SCOPE +const std::vector CMP_RATIO {4, 128}; +const std::vector ROTARY_MODE {2}; +#else +const std::vector CMP_RATIO {2, 4, 8, 16, 32, 64, 128}; +const std::vector ROTARY_MODE {1, 2}; +#endif +const std::vector HEAD_DIM {128, 512}; +const std::vector CACHE_MODE {1}; + +enum class ROTARY_MODE:uint8_t { + HALF = 1, + INTERLEAVE = 2 +}; + +enum class CACHE_MODE:uint8_t { + CONTINUOUS = 1, + CYCLE = 2 +}; + +struct CompressorContext { + const char *opName; + const char *opType; + fe::PlatFormInfos *platformInfo; + + RequiredParaInfo x; + RequiredParaInfo wkv; + RequiredParaInfo wgate; + RequiredParaInfo stateCache; + RequiredParaInfo ape; + RequiredParaInfo normWeight; + RequiredParaInfo ropeSin; + RequiredParaInfo ropeCos; + OptionalParaInfo stateBlockTable; + OptionalParaInfo cuSeqlens; + OptionalParaInfo seqUsed; + OptionalParaInfo startPos; + RequiredParaInfo cmpKv; + + const int *ropeHeadDim; + const int *coff; + const int *cmpRatio; + const float *normEps; + const int *rotaryMode; + const int *cacheMode; + const int *stateCacheStrideDim0; + TemplateId templateId; + + ge::DataType dtype = ge::DT_BF16; + LayoutType layout = LayoutType::LAYOUT_BSH; + + size_t *workSpaces; + uint64_t tilingKey; + uint32_t blockDim; +}; + +class CompressorTiling { +public: + explicit CompressorTiling(CompressorContext *context) : context_(context) {} + ~CompressorTiling() = default; + + static ge::graphStatus ConvertContext(gert::TilingContext &context, CompressorContext &compressorContext); + ge::graphStatus RunBigKernelTiling(CompressorTilingData* tilingData); + +private: + static void ConvertRequiredParams(gert::TilingContext &context, CompressorContext &compressorContext); + + static void ConvertOptionalParams(gert::TilingContext &context, CompressorContext &compressorContext); + ge::graphStatus GetNpuInfo(); + ge::graphStatus SetBaseInfo(); + ge::graphStatus SetPageAttentionInfo(); + ge::graphStatus SetWorkSpaceInfo(); + ge::graphStatus SetScenarioInfo(); + ge::graphStatus SetTemplateId(); + ge::graphStatus SetInnerSplitInfo(); + ge::graphStatus CalcWorkSpace(); + ge::graphStatus CheckSinglePara() const; + ge::graphStatus GenTilingKey() const; + template + ge::graphStatus CheckFeatureValueSupport(const T *featureValue, const std::vector &expectFeatureValList, + const std::string &name) const; + template + ge::graphStatus CheckAttrValueSupport(const T *attrValue, const std::vector &expectAttrValList, + const std::string &name) const; + template + void LogErrorNumberSupport(const std::vector &expectNumberList, const T &actualValue, const std::string &name, + const std::string subName) const; + ge::graphStatus CheckDimNumInLayoutSupport(const std::string &layout, const gert::StorageShape *shape, + const std::string &name) const; + ge::graphStatus CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc, const std::string &name) const; + void LogErrorDtypeSupport(const std::vector &expectDtypeList, const ge::DataType &actualDtype, + const std::string &name) const; + ge::graphStatus CheckDimNumSupport(const gert::StorageShape *shape, const std::string &name) const; + ge::graphStatus LogErrorShapeConsistency(const std::string &name, const gert::StorageShape *shape, + const uint32_t &dimNum, const std::string &subName, + const uint32_t &expectNum) const; + ge::graphStatus CheckSingleParaX() const; + ge::graphStatus CheckSingleParaWkv() const; + ge::graphStatus CheckSingleParaWgate() const; + ge::graphStatus CheckSingleParaStateCache() const; + ge::graphStatus CheckSingleParaApe() const; + ge::graphStatus CheckSingleParaNormWeight() const; + ge::graphStatus CheckSingleParaRopeSin() const; + ge::graphStatus CheckSingleParaRopeCos() const; + ge::graphStatus CheckSingleParaStateBlockTable() const; + ge::graphStatus CheckSingleParaCuSeqlens() const; + ge::graphStatus CheckSingleParaSeqused() const; + ge::graphStatus CheckSingleParaStartPos() const; + ge::graphStatus CheckSingleParaCmpKv() const; + ge::graphStatus CheckSingleParaRopeHeadDim() const; + ge::graphStatus CheckSingleParaCmpRatio() const; + ge::graphStatus CheckSingleParaCoff() const; + ge::graphStatus CheckSingleParaNormEps() const; + ge::graphStatus CheckSingleParaRotaryMode() const; + ge::graphStatus CheckSingleParaCacheMode() const; + ge::graphStatus CheckRequiredParaExistence() const; + ge::graphStatus CheckRequiredInOutExistence() const; + ge::graphStatus CheckRequiredAttrExistence() const; + ge::graphStatus CheckFeature() const; + ge::graphStatus CheckShapeConsistency() const; + ge::graphStatus CheckShapeConsistencyRope() const; + ge::graphStatus CheckDtypeConsistencyX(const gert::CompileTimeTensorDesc *desc, const std::string &name) const; + ge::graphStatus CheckDtypeConsistencyRope() const; + ge::graphStatus CheckDtypeConsistency() const; + ge::graphStatus CheckMultiParaConsistency() const; + ge::graphStatus CheckDimNumConsistency() const; + ge::graphStatus CheckEmptyTensor() const; + ge::graphStatus CheckScenarioConsistency() const; + ge::graphStatus CheckBlockDimConstrain() const; + + size_t ubSize_ = 0; + size_t l1Size_ = 0; + size_t l0cSize_ = 0; + size_t l0bSize_ = 0; + uint32_t coreNum_ = 0; + uint32_t aicNum_ = 0; + uint32_t aivNum_ = 0; + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B; + size_t libapiSize_ = 0; + size_t workspaceSize_ = 0; + uint8_t coff = 1; + + uint32_t mBaseSize = 0; + uint32_t dbaseSize = 0; + + CompressorBaseShapeInfo baseShapeInfo_; + CompressorContext *context_ = nullptr; + CompressorBaseParams *baseParams_ = nullptr; + CompressorPageAttentionParams *pageAttentionParams_ = nullptr; + CompressorInnerSplitParams *innerSplitParams_ = nullptr; + CompressorWorkspaceParams *workspaceParams_ = nullptr; +}; + +} // optiling + +#endif diff --git a/csrc/ascend/attention/compressor/op_host/arch35/compressor_tiling.cpp b/csrc/ascend/attention/compressor/op_host/arch35/compressor_tiling.cpp new file mode 100644 index 000000000..edd04b3d4 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_host/arch35/compressor_tiling.cpp @@ -0,0 +1,1071 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_tiling.cpp + * \file compressor_tiling.cpp + * \brief + */ + +#include +#include +#include +#include +#include +#include "err/ops_err.h" +#include "register/op_def_registry.h" +#include "compressor_tiling.h" + +using namespace ge; +using namespace AscendC; +namespace optiling { + + +void CompressorTiling::ConvertRequiredParams(gert::TilingContext &context, CompressorContext &compressorContext) +{ + compressorContext.x.desc = context.GetRequiredInputDesc(TOKEN_X_INPUT_INDEX); + compressorContext.x.shape = context.GetRequiredInputShape(TOKEN_X_INPUT_INDEX); + compressorContext.wkv.desc = context.GetRequiredInputDesc(WEIGHT_KV_INPUT_INDEX); + compressorContext.wkv.shape = context.GetRequiredInputShape(WEIGHT_KV_INPUT_INDEX); + compressorContext.wgate.desc = context.GetRequiredInputDesc(WEIGHT_WGATE_INPUT_INDEX); + compressorContext.wgate.shape = context.GetRequiredInputShape(WEIGHT_WGATE_INPUT_INDEX); + compressorContext.stateCache.desc = context.GetRequiredInputDesc(STATE_CACHE_INPUT_INDEX); + compressorContext.stateCache.shape = context.GetRequiredInputShape(STATE_CACHE_INPUT_INDEX); + compressorContext.ape.desc = context.GetRequiredInputDesc(APE_INPUT_INDEX); + compressorContext.ape.shape = context.GetRequiredInputShape(APE_INPUT_INDEX); + compressorContext.normWeight.desc = context.GetRequiredInputDesc(NORM_WEIGHT_INPUT_INDEX); + compressorContext.normWeight.shape = context.GetRequiredInputShape(NORM_WEIGHT_INPUT_INDEX); + compressorContext.ropeSin.desc = context.GetRequiredInputDesc(ROPE_SIN_INPUT_INDEX); + compressorContext.ropeSin.shape = context.GetRequiredInputShape(ROPE_SIN_INPUT_INDEX); + compressorContext.ropeCos.desc = context.GetRequiredInputDesc(ROPE_COS_INPUT_INDEX); + compressorContext.ropeCos.shape = context.GetRequiredInputShape(ROPE_COS_INPUT_INDEX); + + compressorContext.cmpKv.desc = context.GetOutputDesc(CMP_KV_OUTPUT_INDEX); + compressorContext.cmpKv.shape = context.GetOutputShape(CMP_KV_OUTPUT_INDEX); + + compressorContext.dtype = compressorContext.x.desc->GetDataType(); + auto xDimNum = compressorContext.x.shape->GetStorageShape().GetDimNum(); + if (xDimNum == COMPRESSOR_DIM_NUM_3) { + compressorContext.layout = LayoutType::LAYOUT_BSH; + } else if (xDimNum == COMPRESSOR_DIM_NUM_2) { + compressorContext.layout = LayoutType::LAYOUT_TH; + } +} + +void CompressorTiling::ConvertOptionalParams(gert::TilingContext &context, CompressorContext &compressorContext) +{ + compressorContext.stateBlockTable.desc = context.GetOptionalInputDesc(STATE_BLOCK_TABLE_INPUT_INDEX); + compressorContext.stateBlockTable.shape = context.GetOptionalInputShape(STATE_BLOCK_TABLE_INPUT_INDEX); + compressorContext.cuSeqlens.desc = context.GetOptionalInputDesc(CU_SEQ_LEN_INPUT_INDEX); + compressorContext.cuSeqlens.shape = context.GetOptionalInputShape(CU_SEQ_LEN_INPUT_INDEX); + compressorContext.seqUsed.desc = context.GetOptionalInputDesc(SEQ_USED_INPUT_INDEX); + compressorContext.seqUsed.shape = context.GetOptionalInputShape(SEQ_USED_INPUT_INDEX); + compressorContext.startPos.desc = context.GetOptionalInputDesc(START_POS_INPUT_INDEX); + compressorContext.startPos.shape = context.GetOptionalInputShape(START_POS_INPUT_INDEX); +} + +ge::graphStatus CompressorTiling::ConvertContext(gert::TilingContext &context, CompressorContext &compressorContext) +{ + if (context.GetNodeName() == nullptr) { + OP_LOGE("Compressor", "opName got from TilingContext is nullptr"); + return ge::GRAPH_FAILED; + } + + OP_LOGI("Getting Context"); + + compressorContext.opName = context.GetNodeName(); + compressorContext.opType = context.GetNodeType(); + compressorContext.platformInfo = context.GetPlatformInfo(); + ConvertRequiredParams(context, compressorContext); + ConvertOptionalParams(context, compressorContext); + + auto attrs = context.GetAttrs(); + OP_CHECK_IF(attrs == nullptr, OP_LOGE(context.GetNodeName(), "attrs got from ge is nullptr"), + return ge::GRAPH_FAILED); + compressorContext.ropeHeadDim = attrs->GetAttrPointer(ROPE_HEAD_DIM_ATTR_INDEX); + compressorContext.coff = attrs->GetAttrPointer(COFF_ATTR_INDEX); + compressorContext.cmpRatio = attrs->GetAttrPointer(CMP_RATIO_ATTR_INDEX); + compressorContext.normEps = attrs->GetAttrPointer(NORM_EPS_ATTR_INDEX); + compressorContext.rotaryMode = attrs->GetAttrPointer(ROTARY_MODE_ATTR_INDEX); + compressorContext.cacheMode = attrs->GetAttrPointer(CACHE_MODE_ATTR_INDEX); + compressorContext.stateCacheStrideDim0 = attrs->GetAttrPointer(STATE_CACHE_STRIDE_DIM0_ATTR_INDEX); + + OP_CHECK_IF(context.GetWorkspaceSizes(1) == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(context.GetNodeName(), "workSpaceSize got from ge is nullptr"), + return ge::GRAPH_FAILED); + compressorContext.workSpaces = context.GetWorkspaceSizes(1); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::GetNpuInfo() +{ + OP_CHECK_IF(context_->platformInfo == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(context_->opName, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->platformInfo); + socVersion_ = ascendcPlatform.GetSocVersion(); + + libapiSize_ = ascendcPlatform.GetLibApiWorkSpaceSize(); + + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize_); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L1, l1Size_); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_C, l0cSize_); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_B, l0bSize_); + + aivNum_ = ascendcPlatform.GetCoreNumAiv(); + aicNum_ = ascendcPlatform.GetCoreNumAic(); + OP_CHECK_IF(aicNum_ == 0 || aivNum_ == 0, + OPS_REPORT_VECTOR_INNER_ERR(context_->opName, "num of core obtained is 0."), return GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::SetBaseInfo() +{ + if (context_->x.shape->GetStorageShape().GetDimNum() == COMPRESSOR_DIM_NUM_3) { + baseParams_->batchSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0); + baseParams_->seqSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1); + baseParams_->hiddenSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_2); + baseParams_->tokenSize = baseParams_->batchSize * baseParams_->seqSize; + baseParams_->cgSize = context_->ropeSin.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1); + } else { + baseParams_->batchSize = context_->cuSeqlens.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0) - 1; + baseParams_->tokenSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0); + baseParams_->hiddenSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1); + baseParams_->cgSize = context_->ropeSin.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0); + } + + baseParams_->headDim = context_->normWeight.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0); + baseParams_->cmpRatio = static_cast(*context_->cmpRatio); + baseParams_->csSize = baseParams_->seqSize - (baseParams_->seqSize % baseParams_->cmpRatio); + baseParams_->ropeHeadDim = static_cast(*context_->ropeHeadDim); + baseParams_->normEps = static_cast(*context_->normEps); + baseParams_->reciprocalD = 1.0 / baseParams_->headDim; + baseParams_->cgSize = + (baseParams_->seqSize + baseParams_->cmpRatio - 1) / baseParams_->cmpRatio; // number of token after compress + coff = static_cast(*context_->coff); + baseParams_->stateCacheStrideDim0 = static_cast(*context_->stateCacheStrideDim0); + baseParams_->nSize = 2; // 2:每个核处理两个基本块后做全核同步 + baseParams_->usedCoreNum = aicNum_; + OP_LOGI(context_->opName, "[TILING] bSize:%u tSize:%u cmpRatio:%u coff:%u, stateCacheStrideDim0:%u", + baseParams_->batchSize, baseParams_->tokenSize, baseParams_->cmpRatio, coff, + baseParams_->stateCacheStrideDim0); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::SetPageAttentionInfo() +{ + pageAttentionParams_->blockNum = context_->stateCache.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0); + pageAttentionParams_->blockSize = context_->stateCache.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1); + if (static_cast(*context_->cacheMode) == static_cast(CACHE_MODE::CONTINUOUS)) { + pageAttentionParams_->maxBlockNumPerBatch = + context_->stateBlockTable.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::SetWorkSpaceInfo() +{ + workspaceParams_->dbWorkspaceRatio = 2; + workspaceParams_->mm1KvResSize = innerSplitParams_->mBaseSize * baseParams_->headDim * coff; + workspaceParams_->mm1ScoreResSize = innerSplitParams_->mBaseSize * baseParams_->headDim * coff; + if (coff == 2) { + workspaceParams_->vec1TailCacheSize = baseParams_->cmpRatio * baseParams_->headDim; + } + workspaceParams_->vec1ResSize = innerSplitParams_->mBaseSize * baseParams_->headDim * baseParams_->nSize; + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::SetScenarioInfo() +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::SetTemplateId() +{ + if (context_->templateId == TemplateId::EMPTY_X) { + return ge::GRAPH_SUCCESS; + } + if (socVersion_ == platform_ascendc::SocVersion::ASCEND950) { + // 设置高性能模板 + if (context_->layout == LayoutType::LAYOUT_BSH && baseParams_->seqSize <= 4 && baseParams_->tokenSize <= 256) { + context_->templateId = TemplateId::FULL_LOAD; + } + return ge::GRAPH_SUCCESS; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::SetInnerSplitInfo() +{ + if (context_->templateId == TemplateId::FULL_LOAD) { + innerSplitParams_->mBaseSize = 256; // 256:核间切分，M轴基本块大小 + innerSplitParams_->dBaseSize = 256 / (coff * 2); // nBase = dBase * coff * 2 + uint32_t dBaseNum = baseParams_->headDim / innerSplitParams_->dBaseSize; + uint32_t mBaseNum = (baseParams_->tokenSize + innerSplitParams_->mBaseSize - 1) / innerSplitParams_->mBaseSize; + baseParams_->coreGroupNum = baseParams_->usedCoreNum / dBaseNum; + baseParams_->kBaseNum = 1; + baseParams_->kBaseSize = baseParams_->hiddenSize; + if ((dBaseNum * mBaseNum) < baseParams_->usedCoreNum) { + baseParams_->kBaseNum = baseParams_->usedCoreNum / dBaseNum; + uint32_t kAlignSize = (baseParams_->hiddenSize + baseParams_->kBaseNum - 1) / baseParams_->kBaseNum; + baseParams_->kBaseSize = kAlignSize / 16 * 16; // 切k的size需要16对齐 + } + for (uint32_t i = 0; i < baseParams_->usedCoreNum; i++) { + baseParams_->splitCoreParam[i].nStart = (i % dBaseNum) * innerSplitParams_->dBaseSize; + baseParams_->splitCoreParam[i].nEnd = baseParams_->splitCoreParam[i].nStart + innerSplitParams_->dBaseSize; + if (baseParams_->kBaseNum > 1) { + uint32_t kStartIdx = i / dBaseNum; + if (kStartIdx + 1 < baseParams_->coreGroupNum) { + uint32_t dealKSize = baseParams_->kBaseSize; + baseParams_->splitCoreParam[i].kStart = kStartIdx * baseParams_->kBaseSize; + baseParams_->splitCoreParam[i].kEnd = baseParams_->splitCoreParam[i].kStart + dealKSize; + } else { + uint32_t dealKSize = kStartIdx < baseParams_->coreGroupNum ? + baseParams_->hiddenSize - kStartIdx * baseParams_->kBaseSize : 0; + baseParams_->splitCoreParam[i].kStart = kStartIdx * baseParams_->kBaseSize; + baseParams_->splitCoreParam[i].kEnd = baseParams_->splitCoreParam[i].kStart + dealKSize; + } + baseParams_->splitCoreParam[i].mStart = 0; + baseParams_->splitCoreParam[i].mEnd = baseParams_->tokenSize; + baseParams_->mLoopNum = 1; + } else { + baseParams_->splitCoreParam[i].kStart = 0; + baseParams_->splitCoreParam[i].kEnd = baseParams_->splitCoreParam[i].kStart + baseParams_->kBaseSize; + baseParams_->splitCoreParam[i].mStart = (i / dBaseNum) * innerSplitParams_->mBaseSize; + baseParams_->splitCoreParam[i].mEnd = + baseParams_->splitCoreParam[i].mStart + innerSplitParams_->mBaseSize; + baseParams_->mLoopNum = mBaseNum / baseParams_->coreGroupNum; + } + } + } else { + innerSplitParams_->mBaseSize = 256; // 256:核间切分，M轴基本块大小 + innerSplitParams_->dBaseSize = 128 / coff; // 128：核间切分，D轴基本块大小 + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CalcWorkSpace() +{ + constexpr uint32_t MM1_RES_ELEM_SIZE = 4; // 4: fp32 + constexpr uint32_t V1_RES_ELEM_SIZE = 4; // 4: fp32 + uint32_t maxGroupNum = aicNum_ / (baseParams_->headDim / innerSplitParams_->dBaseSize); + workspaceSize_ = libapiSize_; + workspaceSize_ += + workspaceParams_->mm1KvResSize * maxGroupNum * MM1_RES_ELEM_SIZE * workspaceParams_->dbWorkspaceRatio; + workspaceSize_ += + workspaceParams_->mm1ScoreResSize * maxGroupNum * MM1_RES_ELEM_SIZE * workspaceParams_->dbWorkspaceRatio; + workspaceSize_ += + workspaceParams_->vec1TailCacheSize * MM1_RES_ELEM_SIZE * workspaceParams_->dbWorkspaceRatio * 2; // 2 kv和score + workspaceSize_ += + workspaceParams_->vec1ResSize * maxGroupNum * V1_RES_ELEM_SIZE * workspaceParams_->dbWorkspaceRatio; + + if (context_->workSpaces) { + context_->workSpaces[0] = workspaceSize_; + } + + OP_LOGI(context_->opName, "Tiling info: workspaceSize_ = %zu", workspaceSize_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckEmptyTensor() const +{ + if (context_->layout == LayoutType::LAYOUT_BSH && + context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0) == 0 || + context_->layout == LayoutType::LAYOUT_BSH && + context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1) == 0 || + context_->layout == LayoutType::LAYOUT_TH && + context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0) == 0) { + context_->templateId = TemplateId::EMPTY_X; + } else { + if (context_->x.shape->GetStorageShape().GetShapeSize() == 0 || + context_->wkv.shape->GetStorageShape().GetShapeSize() == 0 || + context_->wgate.shape->GetStorageShape().GetShapeSize() == 0 || + context_->stateCache.shape->GetStorageShape().GetShapeSize() == 0 || + context_->ape.shape->GetStorageShape().GetShapeSize() == 0 || + context_->normWeight.shape->GetStorageShape().GetShapeSize() == 0 || + context_->ropeSin.shape->GetStorageShape().GetShapeSize() == 0 || + context_->ropeCos.shape->GetStorageShape().GetShapeSize() == 0 || + context_->stateBlockTable.shape->GetStorageShape().GetShapeSize() == 0) { + OP_LOGE(context_->opName, "Only input tensor x dim B or S or T supports to be 0"); + return ge::GRAPH_FAILED; + } + context_->templateId = TemplateId::NORMAL; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::RunBigKernelTiling(CompressorTilingData *tilingData) +{ + this->baseParams_ = &tilingData->baseParams; + this->pageAttentionParams_ = &tilingData->pageAttentionParams; + this->innerSplitParams_ = &tilingData->innerSplitParams; + this->workspaceParams_ = &tilingData->workspaceParams; + using StatusFunction = std::function; + std::vector requiredTilingFuncs { + std::bind(&CompressorTiling::GetNpuInfo, this), + std::bind(&CompressorTiling::CheckRequiredParaExistence, this), + std::bind(&CompressorTiling::CheckEmptyTensor, this), + std::bind(&CompressorTiling::CheckSinglePara, this), + std::bind(&CompressorTiling::SetBaseInfo, this), + std::bind(&CompressorTiling::SetPageAttentionInfo, this), + std::bind(&CompressorTiling::CheckFeature, this), + std::bind(&CompressorTiling::CheckMultiParaConsistency, this), + std::bind(&CompressorTiling::CheckBlockDimConstrain, this), + std::bind(&CompressorTiling::SetTemplateId, this), + std::bind(&CompressorTiling::SetInnerSplitInfo, this), + std::bind(&CompressorTiling::SetWorkSpaceInfo, this), + std::bind(&CompressorTiling::SetScenarioInfo, this) + }; + for (const auto &func: requiredTilingFuncs) { + if (func() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + } + + if (context_->templateId == TemplateId::EMPTY_X) { + workspaceSize_ = libapiSize_; + if (context_->workSpaces) { + context_->workSpaces[0] = workspaceSize_; + } + GenTilingKey(); + context_->blockDim = 1U; + return ge::GRAPH_SUCCESS; + } + std::vector optionalTilingFuncs { + std::bind(&CompressorTiling::CalcWorkSpace, this), + std::bind(&CompressorTiling::GenTilingKey, this) + }; + for (const auto &func : optionalTilingFuncs) { + if (func() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + } + + context_->blockDim = aicNum_; + + OP_LOGI("Run big kernel"); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::GenTilingKey() const +{ + // 0:BF16, 1:FP16 + uint8_t dtype = 0; + // 0: BSH 1:TH + uint8_t layout = 0; + uint8_t rotaryMode = static_cast(*context_->rotaryMode); + uint8_t templateId = static_cast(context_->templateId); + uint8_t cacheMode = static_cast(*context_->cacheMode); + + auto xDtype = context_->x.desc->GetDataType(); + if (xDtype == ge::DT_BF16) { + dtype = 0; + } else if (xDtype == ge::DT_FLOAT16) { + dtype = 1; + } + auto xDimNum = context_->x.shape->GetStorageShape().GetDimNum(); + if (xDimNum == COMPRESSOR_DIM_NUM_3) { + layout = 0; + } else { + layout = 1; + } + + context_->tilingKey = GET_TPL_TILING_KEY( + layout, + dtype, + coff, + rotaryMode, + cacheMode, + templateId + ); + OP_LOGI(context_->opName, + "Compressor dtype:%hhu layout:%hhu coff:%hhu rotary_mode:%hhu, cacheMode: %u, template_id:%hhu", dtype, + layout, coff, rotaryMode, cacheMode, templateId); + OP_LOGI(context_->opName, "Compressor tilingKey:%lu", context_->tilingKey); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSinglePara() const +{ + if (ge::GRAPH_SUCCESS != CheckSingleParaX() || + ge::GRAPH_SUCCESS != CheckSingleParaWkv() || + ge::GRAPH_SUCCESS != CheckSingleParaWgate() || + ge::GRAPH_SUCCESS != CheckSingleParaStateCache() || + ge::GRAPH_SUCCESS != CheckSingleParaApe() || + ge::GRAPH_SUCCESS != CheckSingleParaNormWeight() || + ge::GRAPH_SUCCESS != CheckSingleParaRopeSin() || + ge::GRAPH_SUCCESS != CheckSingleParaRopeCos() || + ge::GRAPH_SUCCESS != CheckSingleParaStateBlockTable() || + ge::GRAPH_SUCCESS != CheckSingleParaCuSeqlens() || + ge::GRAPH_SUCCESS != CheckSingleParaSeqused() || + ge::GRAPH_SUCCESS != CheckSingleParaStartPos() || + ge::GRAPH_SUCCESS != CheckSingleParaCmpKv() || + ge::GRAPH_SUCCESS != CheckSingleParaRopeHeadDim() || + ge::GRAPH_SUCCESS != CheckSingleParaCmpRatio() || + ge::GRAPH_SUCCESS != CheckSingleParaCoff() || + ge::GRAPH_SUCCESS != CheckSingleParaNormEps() || + ge::GRAPH_SUCCESS != CheckSingleParaRotaryMode() || + ge::GRAPH_SUCCESS != CheckSingleParaCacheMode()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +template +ge::graphStatus CompressorTiling::CheckFeatureValueSupport(const T *featureValue, + const std::vector &expectFeatureValList, + const std::string &name) const +{ + if (std::find(expectFeatureValList.begin(), expectFeatureValList.end(), *featureValue) == + expectFeatureValList.end()) { + LogErrorNumberSupport(expectFeatureValList, *featureValue, name, "feature value"); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +template +ge::graphStatus CompressorTiling::CheckAttrValueSupport(const T *attrValue, const std::vector &expectAttrValList, + const std::string &name) const +{ + if (attrValue == nullptr) { + return ge::GRAPH_SUCCESS; + } + + if (std::find(expectAttrValList.begin(), expectAttrValList.end(), *attrValue) == expectAttrValList.end()) { + LogErrorNumberSupport(expectAttrValList, *attrValue, name, "attr value"); + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +template +std::string to_string(const T &value) +{ + if (std::is_same_v) { + return value ? "true" : "false"; + } else { + return std::to_string(value); + } +} + +template +void CompressorTiling::LogErrorNumberSupport(const std::vector &expectNumberList, const T &actualValue, + const std::string &name, const std::string subName) const +{ + std::ostringstream oss; + for (size_t i = 0; i < expectNumberList.size(); ++i) { + oss << to_string(expectNumberList[i]); + if (i < expectNumberList.size() - 1) { + oss << ", "; + } + } + + OP_LOGE(context_->opName, "%s %s only supports %s, but got %s", name.c_str(), subName.c_str(), oss.str().c_str(), + to_string(actualValue).c_str()); +} + +std::string LayoutTypeToStr(LayoutType layout) +{ + switch (layout) { + case LayoutType::LAYOUT_BSH: + return "BSH"; + case LayoutType::LAYOUT_TH: + return "TH"; + default: + return "UNKNOWN_LAYOUT"; + } +} + +ge::graphStatus CompressorTiling::CheckDimNumInLayoutSupport(const std::string &layout, const gert::StorageShape *shape, + const std::string &name) const +{ + const auto &dimIt = LAYOUT_DIM_MAP.find(layout); + OP_CHECK_IF(shape->GetStorageShape().GetDimNum() != dimIt->second, + OP_LOGE(context_->opName, "When layout is %s, %s dimension should be %zu, but it's %zu", layout.c_str(), + name.c_str(), dimIt->second, shape->GetStorageShape().GetDimNum()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc, + const std::string &name) const +{ + if (desc != nullptr) { + const auto &it = DTYPE_SUPPORT_MAP.find(name); + OP_CHECK_IF( + it == DTYPE_SUPPORT_MAP.end(), + OP_LOGE(context_->opName, "%s datatype support list should be specify in DTYPE_SUPPORT_MAP", name.c_str()), + return ge::GRAPH_FAILED); + auto &expectDtypeList = it->second; + OP_CHECK_IF(std::find(expectDtypeList.begin(), expectDtypeList.end(), desc->GetDataType()) == + expectDtypeList.end(), + LogErrorDtypeSupport(expectDtypeList, desc->GetDataType(), name), return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +void CompressorTiling::LogErrorDtypeSupport(const std::vector &expectDtypeList, + const ge::DataType &actualDtype, const std::string &name) const +{ + std::ostringstream oss; + for (size_t i = 0; i < expectDtypeList.size(); ++i) { + oss << DataTypeToSerialString(expectDtypeList[i]); + if (i < expectDtypeList.size() - 1) { + oss << ", "; + } + } + OP_LOGE(context_->opName, "Tensor %s only supports dtype %s, but got %s", name.c_str(), oss.str().c_str(), + DataTypeToSerialString(actualDtype).c_str()); +} + +static std::string DataTypeToSerialString(ge::DataType type) +{ + const auto it = DATATYPE_TO_STRING_MAP.find(type); + if (it != DATATYPE_TO_STRING_MAP.end()) { + return it->second; + } else { + OP_LOGE("Compressor", "datatype %d not support", type); + return "UNDEFINED"; + } +} + +ge::graphStatus CompressorTiling::CheckDimNumSupport(const gert::StorageShape *shape, const std::string &name) const +{ + if (shape == nullptr) { + return ge::GRAPH_SUCCESS; + } + const auto &it = DIM_NUM_MAP.find(name); + OP_CHECK_IF(it == DIM_NUM_MAP.end(), + OP_LOGE(context_->opName, "%s dim number support list should be specify in DIM_NUM_MAP", name.c_str()), + return ge::GRAPH_FAILED); + auto &expectDimNumList = it->second; + OP_CHECK_IF(std::find(expectDimNumList.begin(), expectDimNumList.end(), shape->GetStorageShape().GetDimNum()) == + expectDimNumList.end(), + LogErrorNumberSupport(expectDimNumList, static_cast(shape->GetStorageShape().GetDimNum()), + name, "dimension"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaX() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->x.desc, X_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->x.shape, X_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(LayoutTypeToStr(context_->layout), context_->x.shape, X_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaWkv() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->wkv.desc, WKV_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->wkv.shape, WKV_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaWgate() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->wgate.desc, WGATE_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->wgate.shape, WGATE_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaStateCache() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->stateCache.desc, STATE_CACHE_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->stateCache.shape, STATE_CACHE_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaApe() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->ape.desc, APE_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->ape.shape, APE_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaNormWeight() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->normWeight.desc, NORM_WEIGHT_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->normWeight.shape, NORM_WEIGHT_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaRopeSin() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->ropeSin.desc, ROPE_SIN_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->ropeSin.shape, ROPE_SIN_NAME) || + ge::GRAPH_SUCCESS != + CheckDimNumInLayoutSupport(LayoutTypeToStr(context_->layout), context_->ropeSin.shape, ROPE_SIN_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaRopeCos() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->ropeCos.desc, ROPE_COS_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->ropeCos.shape, ROPE_COS_NAME) || + ge::GRAPH_SUCCESS != + CheckDimNumInLayoutSupport(LayoutTypeToStr(context_->layout), context_->ropeCos.shape, ROPE_COS_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaStateBlockTable() const +{ + if (context_->stateBlockTable.desc == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->stateBlockTable.desc, STATE_BLOCK_TABLE_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->stateBlockTable.shape, STATE_BLOCK_TABLE_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaCuSeqlens() const +{ + if (context_->cuSeqlens.desc == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->cuSeqlens.desc, CU_SEQLENS_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->cuSeqlens.shape, CU_SEQLENS_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaSeqused() const +{ + if (context_->seqUsed.desc == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->seqUsed.desc, SEQUSED_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->seqUsed.shape, SEQUSED_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaStartPos() const +{ + if (context_->startPos.desc == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->startPos.desc, START_POS_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->startPos.shape, START_POS_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaCmpKv() const +{ + if (context_->cmpKv.desc == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->cmpKv.desc, CMP_KV_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->cmpKv.shape, CMP_KV_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaRopeHeadDim() const +{ + if (CheckAttrValueSupport(context_->ropeHeadDim, ROPE_HEAD_DIM, ROPE_HEAD_DIM_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaCmpRatio() const +{ + if (CheckAttrValueSupport(context_->cmpRatio, CMP_RATIO, CMP_RATIO_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaCoff() const +{ + if (CheckAttrValueSupport(context_->coff, COFF, COFF_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaNormEps() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaRotaryMode() const +{ + if (ge::GRAPH_SUCCESS != CheckAttrValueSupport(context_->rotaryMode, ROTARY_MODE, ROTARY_MODE_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckSingleParaCacheMode() const +{ + if (ge::GRAPH_SUCCESS != CheckAttrValueSupport(context_->cacheMode, CACHE_MODE, CACHE_MODE_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckRequiredParaExistence() const +{ + if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS || CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckRequiredInOutExistence() const +{ + OP_CHECK_IF(context_->x.shape == nullptr, OP_LOGE(context_->opName, "tensor x is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->x.desc == nullptr, OP_LOGE(context_->opName, "tensor x is nullptr"), return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->wkv.shape == nullptr, OP_LOGE(context_->opName, "tensor wkv is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->wkv.desc == nullptr, OP_LOGE(context_->opName, "tensor wkv is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->wgate.shape == nullptr, OP_LOGE(context_->opName, "tensor wgate is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->wgate.desc == nullptr, OP_LOGE(context_->opName, "tensor wgate is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->stateCache.shape == nullptr, OP_LOGE(context_->opName, "tensor stateCache is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->stateCache.desc == nullptr, OP_LOGE(context_->opName, "tensor stateCache is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->ape.shape == nullptr, OP_LOGE(context_->opName, "tensor ape is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->ape.desc == nullptr, OP_LOGE(context_->opName, "tensor ape is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->normWeight.shape == nullptr, OP_LOGE(context_->opName, "tensor normWeight is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->normWeight.desc == nullptr, OP_LOGE(context_->opName, "tensor normWeight is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->ropeSin.shape == nullptr, OP_LOGE(context_->opName, "tensor ropeSin is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->ropeSin.desc == nullptr, OP_LOGE(context_->opName, "tensor ropeSin is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->ropeCos.shape == nullptr, OP_LOGE(context_->opName, "tensor ropeCos is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->ropeCos.desc == nullptr, OP_LOGE(context_->opName, "tensor ropeCos is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->stateBlockTable.shape == nullptr, + OP_LOGE(context_->opName, "tensor stateBlockTable is nullptr"), return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->stateBlockTable.desc == nullptr, + OP_LOGE(context_->opName, "tensor stateBlockTable is nullptr"), return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->cmpKv.shape == nullptr, OP_LOGE(context_->opName, "tensor cmpKv is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->cmpKv.desc == nullptr, OP_LOGE(context_->opName, "tensor cmpKv is nullptr"), + return ge::GRAPH_FAILED); + if (context_->layout == LayoutType::LAYOUT_TH) { + OP_CHECK_IF(context_->cuSeqlens.desc == nullptr, + OP_LOGE(context_->opName, "In TH layout, tensor cuSeqlens should not be nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->cuSeqlens.shape == nullptr, + OP_LOGE(context_->opName, "In TH layout, tensor cuSeqlens should not be nullptr"), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF(context_->cuSeqlens.desc != nullptr, + OP_LOGE(context_->opName, "In BSH layout, tensor cuSeqlens must be nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->cuSeqlens.shape != nullptr, + OP_LOGE(context_->opName, "In TH layout, tensor cuSeqlens must be nullptr"), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckRequiredAttrExistence() const +{ + OP_CHECK_IF(context_->ropeHeadDim == nullptr, OP_LOGE(context_->opName, "attr ropeHeadDim is nullptr"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(context_->cmpRatio == nullptr, OP_LOGE(context_->opName, "attr cmpRatio is nullptr"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckFeature() const +{ + if (ge::GRAPH_SUCCESS != CheckFeatureValueSupport(&baseParams_->headDim, HEAD_DIM, "headDim")) { + return ge::GRAPH_FAILED; + } + OP_CHECK_IF(baseParams_->hiddenSize > MAX_HIDDEN_SIZE || baseParams_->hiddenSize < MIN_HIDDEN_SIZE || + baseParams_->hiddenSize % ALIGN_FACTOR_HIDDEN_SIZE != 0, + OP_LOGE(context_->opName, "hiddenSize should be whthin [1k, 10k] and be 512-aligned, but got %u", + baseParams_->hiddenSize), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + pageAttentionParams_->blockSize < MIN_BLOCK_SIZE, + OP_LOGE(context_->opName, "blockSize should not be less than 1, but got %u", pageAttentionParams_->blockSize), + return ge::GRAPH_FAILED); + if (static_cast(*context_->cacheMode) == static_cast(CACHE_MODE::CYCLE)) { + OP_CHECK_IF(pageAttentionParams_->blockNum < baseParams_->batchSize, + OP_LOGE(context_->opName, + "when cacheMode is %u, blockNum should not be less than batchSize(%u), but got %u", + static_cast(CACHE_MODE::CYCLE), baseParams_->batchSize, + pageAttentionParams_->blockSize), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::LogErrorShapeConsistency(const std::string &name, const gert::StorageShape *shape, + const uint32_t &dimNum, const std::string &subName, + const uint32_t &expectNum) const +{ + if (shape == nullptr) { + return ge::GRAPH_SUCCESS; + } + + const uint32_t actualNum = shape->GetStorageShape().GetDim(dimNum); + OP_CHECK_IF(actualNum != expectNum, + OP_LOGE(context_->opName, "%s shape dim %u, should be equal to %s: %u, but got %u", name.c_str(), + dimNum, subName.c_str(), expectNum, actualNum), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckShapeConsistency() const +{ + if (CheckShapeConsistencyRope() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + auto coffD = coff * baseParams_->headDim; + uint32_t stateNum = 2; + if (ge::GRAPH_SUCCESS != LogErrorShapeConsistency("stateBlockTable", context_->stateBlockTable.shape, + COMPRESSOR_DIM_INDEX_0, "batchSize", baseParams_->batchSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("cuSeqlens", context_->cuSeqlens.shape, COMPRESSOR_DIM_INDEX_0, + "batchSize+1", baseParams_->batchSize + 1) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("seqUsed", context_->seqUsed.shape, COMPRESSOR_DIM_INDEX_0, + "batchSize", baseParams_->batchSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("startPos", context_->startPos.shape, COMPRESSOR_DIM_INDEX_0, + "batchSize", baseParams_->batchSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("wkv", context_->wkv.shape, COMPRESSOR_DIM_INDEX_1, "hiddenSize", + baseParams_->hiddenSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("wgate", context_->wgate.shape, COMPRESSOR_DIM_INDEX_1, + "hiddenSize", baseParams_->hiddenSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("wkv", context_->wkv.shape, COMPRESSOR_DIM_INDEX_0, + "coff*headDim", static_cast(coffD)) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("wgate", context_->wgate.shape, COMPRESSOR_DIM_INDEX_0, + "coff*headDim", static_cast(coffD)) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("stateCache", context_->stateCache.shape, COMPRESSOR_DIM_INDEX_2, + "2*coff*headDim", stateNum * static_cast(coffD)) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ape", context_->ape.shape, COMPRESSOR_DIM_INDEX_1, + "coff*headDim", static_cast(coffD)) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ape", context_->ape.shape, COMPRESSOR_DIM_INDEX_0, "cmpRatio", + baseParams_->cmpRatio)) { + return ge::GRAPH_FAILED; + } + if (static_cast(*context_->cacheMode) == static_cast(CACHE_MODE::CONTINUOUS) && + (ge::GRAPH_SUCCESS != LogErrorShapeConsistency("stateCache", context_->stateCache.shape, COMPRESSOR_DIM_INDEX_0, + "blockNum", pageAttentionParams_->blockNum) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("stateCache", context_->stateCache.shape, COMPRESSOR_DIM_INDEX_1, + "blockSize", pageAttentionParams_->blockSize))) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckShapeConsistencyRope() const +{ + auto cmpT = + std::min(baseParams_->tokenSize, baseParams_->tokenSize / baseParams_->cmpRatio + baseParams_->batchSize); + if (context_->layout == LayoutType::LAYOUT_BSH) { + if (ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeSin", context_->ropeSin.shape, COMPRESSOR_DIM_INDEX_0, + "batchSize", baseParams_->batchSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeCos", context_->ropeCos.shape, COMPRESSOR_DIM_INDEX_0, + "batchSize", baseParams_->batchSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeSin", context_->ropeSin.shape, COMPRESSOR_DIM_INDEX_1, + "ceil(seqSize/cmpRatio)", baseParams_->cgSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeCos", context_->ropeCos.shape, COMPRESSOR_DIM_INDEX_1, + "ceil(seqSize/cmpRatio)", baseParams_->cgSize) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeSin", context_->ropeSin.shape, COMPRESSOR_DIM_INDEX_2, + "ropeHeadDim", baseParams_->ropeHeadDim) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeCos", context_->ropeCos.shape, COMPRESSOR_DIM_INDEX_2, + "ropeHeadDim", baseParams_->ropeHeadDim)) { + return ge::GRAPH_FAILED; + } + } else { + if (ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeSin", context_->ropeSin.shape, COMPRESSOR_DIM_INDEX_0, + "min(tokenSize, tokenSize/cmpRatio+batchSize)", + static_cast(cmpT)) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeCos", context_->ropeCos.shape, COMPRESSOR_DIM_INDEX_0, + "min(tokenSize, tokenSize/cmpRatio+batchSize)", + static_cast(cmpT)) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeSin", context_->ropeSin.shape, COMPRESSOR_DIM_INDEX_1, + "ropeHeadDim", baseParams_->ropeHeadDim) || + ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ropeCos", context_->ropeCos.shape, COMPRESSOR_DIM_INDEX_1, + "ropeHeadDim", baseParams_->ropeHeadDim)) { + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckDtypeConsistencyX(const gert::CompileTimeTensorDesc *desc, + const std::string &name) const +{ + const auto actualDtype = desc->GetDataType(); + OP_CHECK_IF(actualDtype != context_->dtype, + OP_LOGE(context_->opName, "%s datatype should be same with x: %s, but got %s", name.c_str(), + DataTypeToSerialString(actualDtype).c_str(), DataTypeToSerialString(context_->dtype).c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckDtypeConsistencyFp32(const gert::CompileTimeTensorDesc *desc, + const std::string &name) const +{ + const auto actualDtype = desc->GetDataType(); + OP_CHECK_IF(actualDtype != ge::DT_FLOAT, + OP_LOGE(context_->opName, "%s datatype should be DT_FLOAT, but got %s", name.c_str(), + DataTypeToSerialString(actualDtype).c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckDtypeConsistency() const +{ + if (CheckDtypeConsistencyX(context_->wkv.desc, WKV_NAME) != ge::GRAPH_SUCCESS || + CheckDtypeConsistencyX(context_->wgate.desc, WGATE_NAME) != ge::GRAPH_SUCCESS || + CheckDtypeConsistencyFp32(context_->normWeight.desc, NORM_WEIGHT_NAME) != ge::GRAPH_SUCCESS || + CheckDtypeConsistencyFp32(context_->ropeSin.desc, ROPE_SIN_NAME) != ge::GRAPH_SUCCESS || + CheckDtypeConsistencyFp32(context_->ropeCos.desc, ROPE_COS_NAME) != ge::GRAPH_SUCCESS || + CheckDtypeConsistencyX(context_->cmpKv.desc, CMP_KV_NAME) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckDimNumConsistency() const +{ + auto xDimNum = context_->x.shape->GetStorageShape().GetDimNum(); + OP_CHECK_IF(xDimNum != context_->ropeSin.shape->GetStorageShape().GetDimNum(), + OP_LOGE(context_->opName, "ropeSin dim num should be equal to x: %u, but got %u", xDimNum, + context_->ropeSin.shape->GetStorageShape().GetDimNum()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(xDimNum != context_->ropeCos.shape->GetStorageShape().GetDimNum(), + OP_LOGE(context_->opName, "ropeCos dim num should be equal to x: %u, but got %u", xDimNum, + context_->ropeCos.shape->GetStorageShape().GetDimNum()), + return ge::GRAPH_FAILED); + OP_CHECK_IF(xDimNum != context_->cmpKv.shape->GetStorageShape().GetDimNum(), + OP_LOGE(context_->opName, "cmpKv dim num should be equal to x: %u, but got %u", xDimNum, + context_->cmpKv.shape->GetStorageShape().GetDimNum()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckScenarioConsistency() const +{ + auto curCmpratio = baseParams_->cmpRatio; + auto curHeaddim = baseParams_->headDim; + auto curCoff = static_cast(*context_->coff); + std::vector curScenario{curCmpratio, curCoff, curHeaddim}; + const std::vector> allowdScenarios = {{4, 2, 512}, {4, 2, 128}, {128, 1, 512}}; + + OP_CHECK_IF(std::find(allowdScenarios.begin(), allowdScenarios.end(), curScenario) == allowdScenarios.end(), + OP_LOGE(context_->opName, + "Cmpratio Coff Headdim should be equal to {4, 2, 512}, {4, 2, 128}, {128, 1, 512}, \ + but now cmpratio=%u, coff=%u, headdim=%u", curCmpratio, curCoff, curHeaddim), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckBlockDimConstrain() const +{ + uint32_t minBlockNum = baseParams_->headDim / 64; // 64 is the largest dBaseSize + OP_CHECK_IF(aicNum_ < minBlockNum, + OP_LOGE(context_->opName, "aicNum is %d, which should not be less than %d", aicNum_, minBlockNum), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus CompressorTiling::CheckMultiParaConsistency() const +{ + if (CheckShapeConsistency() != ge::GRAPH_SUCCESS || CheckDtypeConsistency() != ge::GRAPH_SUCCESS || + CheckDimNumConsistency() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } +#ifdef DAY0_SCOPE + if (CheckScenarioConsistency() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } +#endif + return ge::GRAPH_SUCCESS; +} + +CMP_EXTERN_C ge::graphStatus TilingCompressor(gert::TilingContext *context) +{ + OP_CHECK_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("Compressor", "Context is nullptr."), + return ge::GRAPH_FAILED); + + OP_LOGI("Getting Tiling"); + + CompressorContext compressorContext{}; + if (CompressorTiling::ConvertContext(*context, compressorContext) != ge::GRAPH_SUCCESS) { + OP_LOGE(context->GetNodeName(), "Error occurred while converting tilingContext to Compressor context"); + return ge::GRAPH_FAILED; + } + CompressorTiling compressorTiling(&compressorContext); + CompressorTilingData *tilingData = context->GetTilingData(); + OP_CHECK_IF(tilingData == nullptr, OPS_REPORT_VECTOR_INNER_ERR(compressorContext.opName, "TilingData is nullptr."), + return ge::GRAPH_FAILED); + // 使用SyncAll，需要设置为batchmode模式，所有核同时启动，否则多流方式下执行可能会卡死 + context->SetScheduleMode(BATCH_MODE_SCHEDULE); + if (compressorTiling.RunBigKernelTiling(tilingData) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + context->SetTilingKey(compressorContext.tilingKey); + context->SetBlockDim(compressorContext.blockDim); + OP_LOGI(compressorContext.opName, "block dim: %u.", compressorContext.blockDim); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingPrepareForCompressor(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(Compressor) + .Tiling(TilingCompressor) + .TilingParse(TilingPrepareForCompressor); +} // namespace optiling diff --git a/csrc/ascend/attention/compressor/op_host/arch35/compressor_tiling.h b/csrc/ascend/attention/compressor/op_host/arch35/compressor_tiling.h new file mode 100644 index 000000000..b982ac5af --- /dev/null +++ b/csrc/ascend/attention/compressor/op_host/arch35/compressor_tiling.h @@ -0,0 +1,375 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_tiling.h + * \brief + */ + +#ifndef COMPRESSOR_TILING_H +#define COMPRESSOR_TILING_H + +#include +#include +#include +#include +#include +#include +#include +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "exe_graph/runtime/tiling_context.h" +#include "register/op_def_registry.h" +#include "../../op_kernel/arch35/compressor_template_tiling_key.h" +#include "../../op_kernel/arch35/compressor_tiling_data.h" +#include "platform/platform_info.h" + +#ifdef ASCENDC_OP_TEST +#define CMP_EXTERN_C extern "C" +#else +#define CMP_EXTERN_C +#endif + +namespace optiling { + +// INPUT +constexpr uint32_t TOKEN_X_INPUT_INDEX = 0; +constexpr uint32_t WEIGHT_KV_INPUT_INDEX = 1; +constexpr uint32_t WEIGHT_WGATE_INPUT_INDEX = 2; +constexpr uint32_t STATE_CACHE_INPUT_INDEX = 3; +constexpr uint32_t APE_INPUT_INDEX = 4; +constexpr uint32_t NORM_WEIGHT_INPUT_INDEX = 5; +constexpr uint32_t ROPE_SIN_INPUT_INDEX = 6; +constexpr uint32_t ROPE_COS_INPUT_INDEX = 7; + +// INPUT(OPTION) +constexpr uint32_t STATE_BLOCK_TABLE_INPUT_INDEX = 8; +constexpr uint32_t CU_SEQ_LEN_INPUT_INDEX = 9; +constexpr uint32_t SEQ_USED_INPUT_INDEX = 10; +constexpr uint32_t START_POS_INPUT_INDEX = 11; + +// ATTR +constexpr uint32_t ROPE_HEAD_DIM_ATTR_INDEX = 0; +constexpr uint32_t CMP_RATIO_ATTR_INDEX = 1; +constexpr uint32_t COFF_ATTR_INDEX = 2; +constexpr uint32_t NORM_EPS_ATTR_INDEX = 3; +constexpr uint32_t ROTARY_MODE_ATTR_INDEX = 4; +constexpr uint32_t CACHE_MODE_ATTR_INDEX = 5; +constexpr uint32_t STATE_CACHE_STRIDE_DIM0_ATTR_INDEX = 6; + +// OUTPUT +constexpr uint32_t CMP_KV_OUTPUT_INDEX = 0; + +constexpr uint32_t COMPRESSOR_DIM_NUM_1 = 1; +constexpr uint32_t COMPRESSOR_DIM_NUM_2 = 2; +constexpr uint32_t COMPRESSOR_DIM_NUM_3 = 3; +constexpr uint32_t COMPRESSOR_DIM_NUM_4 = 4; +constexpr uint32_t COMPRESSOR_DIM_INDEX_0 = 0; +constexpr uint32_t COMPRESSOR_DIM_INDEX_1 = 1; +constexpr uint32_t COMPRESSOR_DIM_INDEX_2 = 2; +constexpr uint32_t COMPRESSOR_DIM_INDEX_3 = 3; + +// CONSTRAINTS +constexpr uint32_t MAX_HIDDEN_SIZE = 10240; +constexpr uint32_t MIN_HIDDEN_SIZE = 1024; +constexpr uint32_t ALIGN_FACTOR_HIDDEN_SIZE = 512; +constexpr uint32_t MIN_BLOCK_SIZE = 1; + +constexpr uint32_t BATCH_MODE_SCHEDULE = 1; + +static const std::string X_NAME = "query"; +static const std::string WKV_NAME = "wkv"; +static const std::string WGATE_NAME = "wgate"; +static const std::string STATE_CACHE_NAME = "state_cache"; +static const std::string APE_NAME = "ape"; +static const std::string NORM_WEIGHT_NAME = "norm_weight"; +static const std::string ROPE_SIN_NAME = "rope_sin"; +static const std::string ROPE_COS_NAME = "rope_cos"; +static const std::string STATE_BLOCK_TABLE_NAME = "state_block_table"; +static const std::string CU_SEQLENS_NAME = "cu_seqlens"; +static const std::string SEQUSED_NAME = "seq_used"; +static const std::string START_POS_NAME = "start_pos"; +static const std::string ROPE_HEAD_DIM_NAME = "rope_head_dim"; +static const std::string CMP_RATIO_NAME = "cmp_ratio"; +static const std::string COFF_NAME = "coff"; +static const std::string NORM_EPS_NAME = "nrom_eps"; +static const std::string ROTARY_MODE_NAME = "rotary_mode"; +static const std::string CACHE_MODE_NAME = "cache_mode"; +static const std::string CMP_KV_NAME = "cmp_kv"; + +static std::string DataTypeToSerialString(ge::DataType type); + +const std::map> DTYPE_SUPPORT_MAP = { + {X_NAME, {ge::DT_BF16, ge::DT_FLOAT16}}, + {WKV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}}, + {WGATE_NAME, {ge::DT_BF16, ge::DT_FLOAT16}}, + {STATE_CACHE_NAME, {ge::DT_FLOAT}}, + {APE_NAME, {ge::DT_FLOAT}}, + {NORM_WEIGHT_NAME, {ge::DT_FLOAT}}, + {ROPE_SIN_NAME, {ge::DT_FLOAT}}, + {ROPE_COS_NAME, {ge::DT_FLOAT}}, + {STATE_BLOCK_TABLE_NAME, {ge::DT_INT32}}, + {CU_SEQLENS_NAME, {ge::DT_INT32}}, + {SEQUSED_NAME, {ge::DT_INT32}}, + {START_POS_NAME, {ge::DT_INT32}}, + {CMP_KV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}} +}; + +const std::map> DIM_NUM_MAP = { + {X_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}}, + {WKV_NAME, {COMPRESSOR_DIM_NUM_2}}, + {WGATE_NAME, {COMPRESSOR_DIM_NUM_2}}, + {STATE_CACHE_NAME, {COMPRESSOR_DIM_NUM_3}}, + {APE_NAME, {COMPRESSOR_DIM_NUM_2}}, + {NORM_WEIGHT_NAME, {COMPRESSOR_DIM_NUM_1}}, + {ROPE_SIN_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}}, + {ROPE_COS_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}}, + {STATE_BLOCK_TABLE_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_1}}, + {CU_SEQLENS_NAME, {COMPRESSOR_DIM_NUM_1}}, + {SEQUSED_NAME, {COMPRESSOR_DIM_NUM_1}}, + {START_POS_NAME, {COMPRESSOR_DIM_NUM_1}}, + {CMP_KV_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}} +}; + +static const std::map LAYOUT_DIM_MAP = { + {"BSH", COMPRESSOR_DIM_NUM_3}, + {"TH", COMPRESSOR_DIM_NUM_2}, +}; + +const std::map DATATYPE_TO_STRING_MAP = { + {ge::DT_UNDEFINED, "DT_UNDEFINED"}, // Used to indicate a DataType field has not been set. + {ge::DT_FLOAT, "DT_FLOAT"}, // float type + {ge::DT_FLOAT16, "DT_FLOAT16"}, // fp16 type + {ge::DT_INT8, "DT_INT8"}, // int8 type + {ge::DT_INT16, "DT_INT16"}, // int16 type + {ge::DT_UINT16, "DT_UINT16"}, // uint16 type + {ge::DT_UINT8, "DT_UINT8"}, // uint8 type + {ge::DT_INT32, "DT_INT32"}, // uint32 type + {ge::DT_INT64, "DT_INT64"}, // int64 type + {ge::DT_UINT32, "DT_UINT32"}, // unsigned int32 + {ge::DT_UINT64, "DT_UINT64"}, // unsigned int64 + {ge::DT_BOOL, "DT_BOOL"}, // bool type + {ge::DT_DOUBLE, "DT_DOUBLE"}, // double type + {ge::DT_DUAL, "DT_DUAL"}, // dual output type + {ge::DT_DUAL_SUB_INT8, "DT_DUAL_SUB_INT8"}, // dual output int8 type + {ge::DT_DUAL_SUB_UINT8, "DT_DUAL_SUB_UINT8"}, // dual output uint8 type + {ge::DT_COMPLEX32, "DT_COMPLEX32"}, // complex32 type + {ge::DT_COMPLEX64, "DT_COMPLEX64"}, // complex64 type + {ge::DT_COMPLEX128, "DT_COMPLEX128"}, // complex128 type + {ge::DT_QINT8, "DT_QINT8"}, // qint8 type + {ge::DT_QINT16, "DT_QINT16"}, // qint16 type + {ge::DT_QINT32, "DT_QINT32"}, // qint32 type + {ge::DT_QUINT8, "DT_QUINT8"}, // quint8 type + {ge::DT_QUINT16, "DT_QUINT16"}, // quint16 type + {ge::DT_RESOURCE, "DT_RESOURCE"}, // resource type + {ge::DT_STRING_REF, "DT_STRING_REF"}, // string ref type + {ge::DT_STRING, "DT_STRING"}, // string type + {ge::DT_VARIANT, "DT_VARIANT"}, // dt_variant type + {ge::DT_BF16, "DT_BFLOAT16"}, // dt_bfloat16 type + {ge::DT_INT4, "DT_INT4"}, // dt_variant type + {ge::DT_UINT1, "DT_UINT1"}, // dt_variant type + {ge::DT_INT2, "DT_INT2"}, // dt_variant type + {ge::DT_UINT2, "DT_UINT2"} // dt_variant type +}; + +struct CompressorCompileInfo { + int64_t core_num; +}; + +struct RequiredParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; +}; + +struct OptionalParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; + const gert::Tensor *tensor; +}; + +enum class LayoutType { + LAYOUT_BSH, + LAYOUT_TH +}; + +enum class TemplateId:uint8_t { + NORMAL = 0, + EMPTY_X = 1, + FULL_LOAD = 2 +}; + +CMP_EXTERN_C ge::graphStatus TilingCompressor(gert::TilingContext *context); +struct CompressorBaseShapeInfo { + uint32_t bSize = 0; // B + uint32_t sSize = 0; // S + uint32_t hSize = 0; // Hidden size + uint32_t tSize = 0; // T + uint32_t nSize = 0; // N + uint32_t dSize = 0; // D + uint32_t coffSize = 0; // Coff: 1 or 2 + uint32_t csSize = 0; // Compress sequence len + uint32_t rSize = 0; // Compress ratio + uint32_t cgSize = 0; // Compress group size + uint32_t drSize = 0; // Dr +}; + +const std::vector ROPE_HEAD_DIM {64}; +const std::vector COFF {1, 2}; +const std::vector CMP_RATIO {2, 4, 8, 16, 32, 64, 128}; +const std::vector ROTARY_MODE {1, 2}; +const std::vector HEAD_DIM {128, 512}; +const std::vector CACHE_MODE {1, 2}; + +enum class ROTARY_MODE:uint8_t { + HALF = 1, + INTERLEAVE = 2 +}; + +enum class CACHE_MODE:uint8_t { + CONTINUOUS = 1, + CYCLE = 2 +}; + +struct CompressorContext { + const char *opName; + const char *opType; + fe::PlatFormInfos *platformInfo; + + RequiredParaInfo x; + RequiredParaInfo wkv; + RequiredParaInfo wgate; + RequiredParaInfo stateCache; + RequiredParaInfo ape; + RequiredParaInfo normWeight; + RequiredParaInfo ropeSin; + RequiredParaInfo ropeCos; + OptionalParaInfo stateBlockTable; + OptionalParaInfo cuSeqlens; + OptionalParaInfo seqUsed; + OptionalParaInfo startPos; + RequiredParaInfo cmpKv; + + const int *ropeHeadDim; + const int *coff; + const int *cmpRatio; + const float *normEps; + const int *rotaryMode; + const int *cacheMode; + const int *stateCacheStrideDim0; + TemplateId templateId; + + ge::DataType dtype = ge::DT_BF16; + LayoutType layout = LayoutType::LAYOUT_BSH; + + size_t *workSpaces; + uint64_t tilingKey; + uint32_t blockDim; +}; + +class CompressorTiling { +public: + explicit CompressorTiling(CompressorContext *context) : context_(context) {} + ~CompressorTiling() = default; + + static ge::graphStatus ConvertContext(gert::TilingContext &context, CompressorContext &compressorContext); + ge::graphStatus RunBigKernelTiling(CompressorTilingData* tilingData); + +private: + static void ConvertRequiredParams(gert::TilingContext &context, CompressorContext &compressorContext); + + static void ConvertOptionalParams(gert::TilingContext &context, CompressorContext &compressorContext); + ge::graphStatus GetNpuInfo(); + ge::graphStatus SetBaseInfo(); + ge::graphStatus SetPageAttentionInfo(); + ge::graphStatus SetWorkSpaceInfo(); + ge::graphStatus SetScenarioInfo(); + ge::graphStatus SetTemplateId(); + ge::graphStatus SetInnerSplitInfo(); + ge::graphStatus CalcWorkSpace(); + ge::graphStatus CheckSinglePara() const; + ge::graphStatus GenTilingKey() const; + template + ge::graphStatus CheckFeatureValueSupport(const T *featureValue, const std::vector &expectFeatureValList, + const std::string &name) const; + template + ge::graphStatus CheckAttrValueSupport(const T *attrValue, const std::vector &expectAttrValList, + const std::string &name) const; + template + void LogErrorNumberSupport(const std::vector &expectNumberList, const T &actualValue, const std::string &name, + const std::string subName) const; + ge::graphStatus CheckDimNumInLayoutSupport(const std::string &layout, const gert::StorageShape *shape, + const std::string &name) const; + ge::graphStatus CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc, const std::string &name) const; + void LogErrorDtypeSupport(const std::vector &expectDtypeList, const ge::DataType &actualDtype, + const std::string &name) const; + ge::graphStatus CheckDimNumSupport(const gert::StorageShape *shape, const std::string &name) const; + ge::graphStatus LogErrorShapeConsistency(const std::string &name, const gert::StorageShape *shape, + const uint32_t &dimNum, const std::string &subName, + const uint32_t &expectNum) const; + ge::graphStatus CheckSingleParaX() const; + ge::graphStatus CheckSingleParaWkv() const; + ge::graphStatus CheckSingleParaWgate() const; + ge::graphStatus CheckSingleParaStateCache() const; + ge::graphStatus CheckSingleParaApe() const; + ge::graphStatus CheckSingleParaNormWeight() const; + ge::graphStatus CheckSingleParaRopeSin() const; + ge::graphStatus CheckSingleParaRopeCos() const; + ge::graphStatus CheckSingleParaStateBlockTable() const; + ge::graphStatus CheckSingleParaCuSeqlens() const; + ge::graphStatus CheckSingleParaSeqused() const; + ge::graphStatus CheckSingleParaStartPos() const; + ge::graphStatus CheckSingleParaCmpKv() const; + ge::graphStatus CheckSingleParaRopeHeadDim() const; + ge::graphStatus CheckSingleParaCmpRatio() const; + ge::graphStatus CheckSingleParaCoff() const; + ge::graphStatus CheckSingleParaNormEps() const; + ge::graphStatus CheckSingleParaRotaryMode() const; + ge::graphStatus CheckSingleParaCacheMode() const; + ge::graphStatus CheckRequiredParaExistence() const; + ge::graphStatus CheckRequiredInOutExistence() const; + ge::graphStatus CheckRequiredAttrExistence() const; + ge::graphStatus CheckFeature() const; + ge::graphStatus CheckShapeConsistency() const; + ge::graphStatus CheckShapeConsistencyRope() const; + ge::graphStatus CheckDtypeConsistencyX(const gert::CompileTimeTensorDesc *desc, const std::string &name) const; + ge::graphStatus CheckDtypeConsistencyFp32(const gert::CompileTimeTensorDesc *desc, const std::string &name) const; + ge::graphStatus CheckDtypeConsistency() const; + ge::graphStatus CheckMultiParaConsistency() const; + ge::graphStatus CheckDimNumConsistency() const; + ge::graphStatus CheckEmptyTensor() const; + ge::graphStatus CheckScenarioConsistency() const; + ge::graphStatus CheckBlockDimConstrain() const; + + size_t ubSize_ = 0; + size_t l1Size_ = 0; + size_t l0cSize_ = 0; + size_t l0bSize_ = 0; + uint32_t coreNum_ = 0; + uint32_t aicNum_ = 0; + uint32_t aivNum_ = 0; + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B; + size_t libapiSize_ = 0; + size_t workspaceSize_ = 0; + uint8_t coff = 1; + + uint32_t mBaseSize = 0; + uint32_t dbaseSize = 0; + + CompressorBaseShapeInfo baseShapeInfo_; + CompressorContext *context_ = nullptr; + CompressorBaseParams *baseParams_ = nullptr; + CompressorPageAttentionParams *pageAttentionParams_ = nullptr; + CompressorInnerSplitParams *innerSplitParams_ = nullptr; + CompressorWorkspaceParams *workspaceParams_ = nullptr; +}; + +} // optiling + +#endif diff --git a/csrc/ascend/attention/compressor/op_host/compressor_def.cpp b/csrc/ascend/attention/compressor/op_host/compressor_def.cpp new file mode 100644 index 000000000..a0a5bf818 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_host/compressor_def.cpp @@ -0,0 +1,190 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "register/op_def_registry.h" + +namespace ops { +class Compressor : public OpDef { +public: + static constexpr uint32_t ROPE_HEAD_DIM_VALUE = 64; + static constexpr uint32_t CMP_RATIO_VALUE = 4; + static constexpr uint32_t COFF_VALUE = 1; + static constexpr uint32_t ROTARY_MODE_VALUE = 1; + static constexpr uint32_t CACHE_MODE_VALUE = 1; + static constexpr uint32_t STATE_CACHE_STRIDE_DIM0 = 0; + + explicit Compressor(const char *name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("wkv") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("wgate") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("state_cache") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .IgnoreContiguous(); + this->Input("ape") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("norm_weight") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("rope_sin") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("rope_cos") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("state_block_table") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("cu_seqlens") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("seqused") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("start_pos") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("cmp_kv") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}); + this->Output("state_cache") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}); + this->Attr("rope_head_dim").AttrType(REQUIRED).Int(ROPE_HEAD_DIM_VALUE); + this->Attr("cmp_ratio").AttrType(REQUIRED).Int(CMP_RATIO_VALUE); + this->Attr("coff").AttrType(OPTIONAL).Int(COFF_VALUE); + this->Attr("norm_eps").AttrType(OPTIONAL).Float(1e-6f); + this->Attr("rotary_mode").AttrType(OPTIONAL).Int(ROTARY_MODE_VALUE); + this->Attr("cache_mode").AttrType(OPTIONAL).Int(CACHE_MODE_VALUE); + this->Attr("state_cache_stride_dim0").AttrType(OPTIONAL).Int(STATE_CACHE_STRIDE_DIM0); + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("aclnnSupport.value", "support_aclnn"); // set value of aclnn support + + OpAICoreConfig config910; + config910.Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + config910.Input("wkv") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + config910.Input("wgate") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + config910.Input("state_cache") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .IgnoreContiguous(); + config910.Input("ape") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + config910.Input("norm_weight") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + config910.Input("rope_sin") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + config910.Input("rope_cos") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + config910.Input("state_block_table") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + config910.Input("cu_seqlens") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + config910.Input("seqused") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + config910.Input("start_pos") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + config910.Output("cmp_kv") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}); + config910.Output("state_cache") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}); + config910.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("aclnnSupport.value", "support_aclnn"); + this->AICore().AddConfig("ascend910b", config910); + this->AICore().AddConfig("ascend910_93", config910); + } +}; +OP_ADD(Compressor, optiling::CompressorCompileInfo); +} // namespace ops diff --git a/csrc/ascend/attention/compressor/op_host/compressor_proto.cpp b/csrc/ascend/attention/compressor/op_host/compressor_proto.cpp new file mode 100644 index 000000000..1f944c073 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_host/compressor_proto.cpp @@ -0,0 +1,174 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include +#include +#include "log/ops_log.h" + +using namespace ge; + +namespace ops { + // INPUT + constexpr uint32_t TOKEN_X_INPUT_INDEX = 0; + constexpr uint32_t WEIGHT_KV_INPUT_INDEX = 1; + constexpr uint32_t WEIGHT_WGATE_INPUT_INDEX = 2; + + constexpr uint32_t STATE_CACHE_INPUT_INDEX = 3; + + constexpr uint32_t APE_INPUT_INDEX = 4; + constexpr uint32_t NORM_WEIGHT_INPUT_INDEX = 5; + constexpr uint32_t ROPE_SIN_INPUT_INDEX = 6; + constexpr uint32_t ROPE_COS_INPUT_INDEX = 7; + + // INPUT(OPTION) + constexpr uint32_t STATE_BLOCK_TABLE_INPUT_INDEX = 8; + + constexpr uint32_t CU_SEQ_LEN_INPUT_INDEX = 9; + constexpr uint32_t SEQ_USED_INPUT_INDEX = 10; + constexpr uint32_t START_POS_INPUT_INDEX = 11; + + // ATTR + constexpr uint32_t ROPE_HEAD_DIM_ATTR_INDEX = 0; + constexpr uint32_t CMP_RATIO_ATTR_INDEX = 1; + constexpr uint32_t COFF_ATTR_INDEX = 2; + constexpr uint32_t NORM_EPS_ATTR_INDEX = 3; + constexpr uint32_t ROTARY_MODE_ATTR_INDEX = 4; + constexpr uint32_t CACHE_MODE_ATTR_INDEX = 5; + constexpr uint32_t STATE_CACHE_STRIDE_DIM0_ATTR_INDEX = 6; + + // OUTPUT + constexpr uint32_t CMP_KV_OUTPUT_INDEX = 0; + + // ATTR DEFAULT VALUE + constexpr uint32_t CMP_RATIO_VALUE = 4; + constexpr uint32_t COFF_VALUE = 1; + +struct CompressorProtoShapeParam { + bool isBsMerge { false }; + int64_t B { 0 }; + int64_t T { 0 }; + int64_t S { 0 }; + int64_t Sr { 0 }; + int64_t H { 0 }; + int64_t D { 0 }; +}; + +// tmp +constexpr uint32_t DIM_NUM_1 = 1; +constexpr uint32_t DIM_NUM_2 = 2; +constexpr uint32_t DIM_NUM_3 = 3; +constexpr uint32_t DIM_NUM_4 = 4; +constexpr uint32_t DIM_INDEX_0 = 0; +constexpr uint32_t DIM_INDEX_1 = 1; +constexpr uint32_t DIM_INDEX_2 = 2; +constexpr uint32_t DIM_INDEX_3 = 3; + +ge::graphStatus GetCompressorShapeDim(const gert::InferShapeContext* context, CompressorProtoShapeParam &shapeParam) +{ + auto xShape = context->GetRequiredInputShape(TOKEN_X_INPUT_INDEX); // (B, S, H) | (T, H) + OPS_LOG_E_IF_NULL(context, xShape, return ge::GRAPH_FAILED) + auto wkvShape = context->GetRequiredInputShape(WEIGHT_KV_INPUT_INDEX); // (coff * D, H) + OPS_LOG_E_IF_NULL(context, wkvShape, return ge::GRAPH_FAILED) + auto wgateShape = context->GetRequiredInputShape(WEIGHT_WGATE_INPUT_INDEX); // (coff * D, H) + OPS_LOG_E_IF_NULL(context, wgateShape, return ge::GRAPH_FAILED) + + auto stateCacheShape = context->GetRequiredInputShape(STATE_CACHE_INPUT_INDEX); // (block_num, block_size, 2 * coff * D) | (B, tokrn_size, 2 * coff * D) + OPS_LOG_E_IF_NULL(context, stateCacheShape, return ge::GRAPH_FAILED) + + auto apeShape = context->GetRequiredInputShape(APE_INPUT_INDEX); // (r, coff * D) + OPS_LOG_E_IF_NULL(context, apeShape, return ge::GRAPH_FAILED) + auto normWeightShape = context->GetRequiredInputShape(NORM_WEIGHT_INPUT_INDEX); // (D) + OPS_LOG_E_IF_NULL(context, normWeightShape, return ge::GRAPH_FAILED) + auto ropeSinShape = context->GetRequiredInputShape(ROPE_SIN_INPUT_INDEX); // (B, ceil(S / r), rD) | (min(T, T/r + B), rD) + OPS_LOG_E_IF_NULL(context, ropeSinShape, return ge::GRAPH_FAILED) + auto ropeCosShape = context->GetRequiredInputShape(ROPE_COS_INPUT_INDEX); // (B, ceil(S / r), rD) | (min(T, T/r + B), rD) + OPS_LOG_E_IF_NULL(context, ropeCosShape, return ge::GRAPH_FAILED) + + auto stateBlockTableShape = context->GetRequiredInputShape(STATE_BLOCK_TABLE_INPUT_INDEX); // (B, sMax/block_size) | (B, ) + OPS_LOG_E_IF_NULL(context, stateBlockTableShape, return ge::GRAPH_FAILED) + + auto cuSeqlensShape = context->GetRequiredInputShape(CU_SEQ_LEN_INPUT_INDEX); // (B+1,) + OPS_LOG_E_IF_NULL(context, cuSeqlensShape, return ge::GRAPH_FAILED) + auto seqUsedShape = context->GetRequiredInputShape(SEQ_USED_INPUT_INDEX); // (B,) + OPS_LOG_E_IF_NULL(context, seqUsedShape, return ge::GRAPH_FAILED) + auto startPosShape = context->GetRequiredInputShape(START_POS_INPUT_INDEX); // (B,) + OPS_LOG_E_IF_NULL(context, startPosShape, return ge::GRAPH_FAILED) + + if (xShape->GetDimNum() == DIM_NUM_3) { // BS + shapeParam.isBsMerge = false; + shapeParam.B = xShape->GetDim(DIM_INDEX_0); + shapeParam.S = xShape->GetDim(DIM_INDEX_1); + shapeParam.H = xShape->GetDim(DIM_INDEX_2); + shapeParam.T = shapeParam.B * shapeParam.S; + } else { // T + shapeParam.isBsMerge = true; + shapeParam.T = xShape->GetDim(DIM_INDEX_0); + shapeParam.H = xShape->GetDim(DIM_INDEX_1); + } + + shapeParam.D = normWeightShape->GetDim(DIM_INDEX_0); + shapeParam.Sr = ropeSinShape->GetDim(DIM_INDEX_1); + + return GRAPH_SUCCESS; +} + +ge::graphStatus SetCompressorShapeDim(const CompressorProtoShapeParam &shapeParam, gert::InferShapeContext* context) +{ + auto cmpKvShape = context->GetOutputShape(CMP_KV_OUTPUT_INDEX); // query: (B, S, N, Hckv) | (T, N, Hckv) + OPS_LOG_E_IF_NULL(context, cmpKvShape, return ge::GRAPH_FAILED) + auto attr = context->GetAttrs(); + const uint32_t *cmpRatioPtr = attr->GetAttrPointer(CMP_RATIO_ATTR_INDEX); + uint32_t cmpRatio = (cmpRatioPtr != nullptr) ? *cmpRatioPtr : CMP_RATIO_VALUE; + const uint32_t *coffPtr = attr->GetAttrPointer(COFF_ATTR_INDEX); + uint32_t coff = (coffPtr != nullptr) ? *coffPtr : COFF_VALUE; + // Set output shape + if (!shapeParam.isBsMerge) { + cmpKvShape->SetDimNum(DIM_NUM_3); // (B, Sr, H) + cmpKvShape->SetDim(DIM_INDEX_0, shapeParam.B); + cmpKvShape->SetDim(DIM_INDEX_1, shapeParam.Sr); + cmpKvShape->SetDim(DIM_INDEX_2, shapeParam.H); + } else { + cmpKvShape->SetDimNum(DIM_NUM_2); // (T, N, Hckv) + cmpKvShape->SetDim(DIM_INDEX_0, shapeParam.Sr); + cmpKvShape->SetDim(DIM_INDEX_1, shapeParam.H); + } + + return GRAPH_SUCCESS; +} + +ge::graphStatus InferDataTypeCompressor(gert::InferDataTypeContext* context) +{ + OP_CHECK_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("Compressor", "Context is nullptr."), + return ge::GRAPH_FAILED); + OPS_LOG_I(context->GetNodeName(), "Enter Compressor inferDataType impl."); + + context->SetOutputDataType(CMP_KV_OUTPUT_INDEX, context->GetRequiredInputDataType(TOKEN_X_INPUT_INDEX)); + + return GRAPH_SUCCESS; +} + +ge::graphStatus InferShapeCompressor(gert::InferShapeContext* context) +{ + OP_CHECK_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("Compressor", "Context is nullptr."), + return ge::GRAPH_FAILED); + OPS_LOG_I(context->GetNodeName(), "Enter Compressor infershape impl."); + + CompressorProtoShapeParam shapeParam {}; + auto apiRet = GetCompressorShapeDim(context, shapeParam); + OPS_LOG_E_IF((apiRet != GRAPH_SUCCESS), context, return ge::GRAPH_FAILED, "Context get input shape failed"); + + apiRet = SetCompressorShapeDim(shapeParam, context); + OPS_LOG_E_IF((apiRet != GRAPH_SUCCESS), context, return ge::GRAPH_FAILED, "Context set output shape failed"); + + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(Compressor).InferShape(InferShapeCompressor).InferDataType(InferDataTypeCompressor); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_block_cube_perf.h b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_block_cube_perf.h new file mode 100644 index 000000000..4f01098b3 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_block_cube_perf.h @@ -0,0 +1,408 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_block_cube_perf.h + * \brief + */ + +#ifndef COMPRESSOR_BLOCK_CUBE_PERF_H +#define COMPRESSOR_BLOCK_CUBE_PERF_H + +#include "compressor_comm.h" +#include "compressor_tools.h" + +using namespace AscendC; + +namespace Compressor { + +template class CompressorBlockCubePerf { +using MM1_OUT_T = float; +public: + __aicore__ inline CompressorBlockCubePerf(){}; + __aicore__ inline void InitParams(const ConstInfo &constInfo, const CompressorTools &tools); + __aicore__ inline void Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitGlobalBuffers(const GlobalTensor& kvMm1ResGm, const GlobalTensor& scoreMm1ResGm); + __aicore__ inline void AllocEventID(TPipe *pipe); + __aicore__ inline void FreeEventID(TPipe *pipe); + __aicore__ inline void ComputeMm1(const RunInfo &info); + +private: + using T = float; + using X_T = typename AscendC::Conditional::type; + + __aicore__ inline uint32_t GetMSize(const RunInfo &info, uint32_t coffId); + __aicore__ inline void CopyXGmToL1(const RunInfo &info, LocalTensor xL1Tensor, uint32_t hIdx, uint32_t kBase); + __aicore__ inline void CopyWeightGmToL1(LocalTensor wL1Tensor, + uint32_t hIdx, uint32_t kBase, uint32_t coffId); + __aicore__ inline void LoadAToL0(const RunInfo &info, LocalTensor aL0Tensor, LocalTensor xL1Tensor, + uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize); + __aicore__ inline void LoadBToL0(LocalTensor bL0Tensor, LocalTensor wL1Tensor, + uint32_t kStart, uint32_t kBase); + __aicore__ inline void MatrixMmad(LocalTensor cL0Tensor, LocalTensor aL0Tensor, + LocalTensor bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C); + __aicore__ inline void CopyOutMm1Res(const RunInfo &info, LocalTensor cL0Tensor, + uint32_t coffId, uint32_t mStart, uint32_t mDealSize); + + ConstInfo constInfo_ = {}; + CompressorTools tools_; + + // GM + GlobalTensor xGm_; + GlobalTensor wkvGm_; + GlobalTensor wgateGm_; + GlobalTensorkvMm1ResGm; + GlobalTensorscoreMm1ResGm; + GlobalTensor cuSeqlensGm_; + GlobalTensor sequsedGm_; + GlobalTensor startPosGm_; + bool isExistSeqUsed = false; + + // =================================L1 Buffer================================= + static constexpr uint32_t L1_X_SIZE = 128 * 1024; + static constexpr uint32_t L1_W_SIZE = 64 * 1024; + // L1 Buffer + TBuf xBufL1; + TBuf wBufL1; + // =================================L0 Buffer================================= + // L0 buffer size + static constexpr uint32_t L0A_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k + static constexpr uint32_t L0B_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k + static constexpr uint32_t L0C_PP_SIZE = 64 * 1024; // (128 * 2) * 64 * 4 = 64k + // L0_A + TBuf tmpBufL0A; + // L0_B + TBuf tmpBufL0B; + // L0_C + TBuf tmpBufL0C; + // =================================Event&Buffer ID=========================== + // mte2 <> mte1 EventID + static constexpr uint32_t X_EVENT0 = EVENT_ID0; + static constexpr uint32_t X_EVENT1 = EVENT_ID1; + uint32_t xBufId = 0; // 用于DB计数 + static constexpr uint32_t W_EVENT0 = EVENT_ID4; + static constexpr uint32_t W_EVENT1 = EVENT_ID5; + static constexpr uint32_t W_EVENT2 = EVENT_ID6; + static constexpr uint32_t W_EVENT3 = EVENT_ID7; + uint32_t wBufId = 0; // 用于DB计数 + // mte1 <> mmad EventID + static constexpr uint32_t L0AB_EVENT0 = EVENT_ID3; + static constexpr uint32_t L0AB_EVENT1 = EVENT_ID4; + uint32_t l0abBufId = 0; + // mmad <> fixpipe EventID + static constexpr uint32_t L0C_EVENT0 = EVENT_ID0; // 每块L0C单独分配EVENT_ID + static constexpr uint32_t L0C_EVENT1 = EVENT_ID1; + uint32_t l0cBufId = 0; + + // =================================Loop====================================== + uint32_t curBIdx_ = 0; + uint32_t curSIdx_ = 0; + +}; + +template +__aicore__ inline void CompressorBlockCubePerf::InitParams(const ConstInfo &constInfo, const CompressorTools &tools) +{ + this->constInfo_ = constInfo; + this->tools_ = tools; +} + +template __aicore__ inline void CompressorBlockCubePerf::Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut) +{ + xGm_.SetGlobalBuffer((__gm__ X_T *)x); + wkvGm_.SetGlobalBuffer((__gm__ X_T *)wKv); + wgateGm_.SetGlobalBuffer((__gm__ X_T *)wGate); + startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos); + isExistSeqUsed = (seqUsed != nullptr); + if (isExistSeqUsed) { + sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed); + } + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens); + } +} + +template +__aicore__ inline void CompressorBlockCubePerf::InitBuffers(TPipe *pipe) +{ + // L1 + // 1. coff=1时, mBase=256, kL1=256, X单次拷贝到L1的数据量最大为mBase*kL1*sizeof(BF16/FP16)=256*256*2=128K + // 2. coff=2时, mBase=128, kL1=256, r最大为128, X单次拷贝到L1的最大数据量为(128+r)*kL1*sizeof(BF16/FP16)<=128K + pipe->InitBuffer(xBufL1, L1_X_SIZE * 2); + // dBaseSize<=64, wkv和wgate各一份, kL1=256, 右矩阵为dBaseSize*2*sizeof(BF16/FP16)<=64K + // cur和pre循环使用, 2份buffer就足够 + pipe->InitBuffer(wBufL1, L1_W_SIZE * 4); + + // L0 + pipe->InitBuffer(tmpBufL0A, L0A_PP_SIZE * 2); + pipe->InitBuffer(tmpBufL0B, L0B_PP_SIZE * 2); + pipe->InitBuffer(tmpBufL0C, L0C_PP_SIZE * 2); +} + +template +__aicore__ inline void CompressorBlockCubePerf::InitGlobalBuffers(const GlobalTensor& kvMm1ResGm, const GlobalTensor& scoreMm1ResGm) +{ + this->kvMm1ResGm = kvMm1ResGm; + this->scoreMm1ResGm = scoreMm1ResGm; +} + +template +__aicore__ inline void CompressorBlockCubePerf::AllocEventID(TPipe *pipe) +{ + SetFlag(X_EVENT0); + SetFlag(X_EVENT1); + + SetFlag(W_EVENT0); + SetFlag(W_EVENT1); + SetFlag(W_EVENT2); + SetFlag(W_EVENT3); + + SetFlag(L0AB_EVENT0); + SetFlag(L0AB_EVENT1); + + SetFlag(L0C_EVENT0); + SetFlag(L0C_EVENT1); +} + +template +__aicore__ inline void CompressorBlockCubePerf::FreeEventID(TPipe *pipe) +{ + WaitFlag(X_EVENT0); + WaitFlag(X_EVENT1); + + WaitFlag(W_EVENT0); + WaitFlag(W_EVENT1); + WaitFlag(W_EVENT2); + WaitFlag(W_EVENT3); + + WaitFlag(L0AB_EVENT0); + WaitFlag(L0AB_EVENT1); + + WaitFlag(L0C_EVENT0); + WaitFlag(L0C_EVENT1); +} + +template +__aicore__ inline void CompressorBlockCubePerf::CopyXGmToL1(const RunInfo &info, LocalTensor xL1Tensor, + uint32_t hIdx, uint32_t kBase) +{ + uint32_t tStart = tools_.GetTIdxByBatch(info.bStart) + info.sStart; // 此基本块在整个序列中的位置 + uint32_t copySeqCnt = info.dealSeqCnt; // 此基本块处理的长度 + + uint32_t xL1Offset = 0 * (32 / sizeof(X_T)); + uint64_t sIdx = tStart; // 起始s在整个T的起始点 + uint64_t gmOffset = sIdx * constInfo_.hSize + hIdx; + uint32_t nValue = copySeqCnt; + uint32_t dValue = kBase; // 拷贝的列数kBase + uint32_t srcDValue = constInfo_.hSize; + uint32_t dstNzC0Stride = (copySeqCnt + 15) / 16 * 16; // 1行变2行的行方向的偏移，需要16对齐 + CopySingleMatrixNDToNZ(xL1Tensor[xL1Offset], xGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride); +} + +template +__aicore__ inline void CompressorBlockCubePerf::CopyWeightGmToL1(LocalTensor wL1Tensor, + uint32_t hIdx, uint32_t kBase, uint32_t coffId) +{ + // coffId=0, 搬运左矩阵的数据; coffId=1, 搬运右矩阵的数据 + uint64_t gmOffset = coffId * constInfo_.headDim * constInfo_.hSize + constInfo_.dIdx * constInfo_.hSize + hIdx; + uint32_t wkvL1Offset = 0; + uint32_t wgateL1Offset = constInfo_.dBaseSize * (32 / sizeof(X_T)); // wgate与wkv的起始点相隔dBaseSize个32B + uint32_t nValue = constInfo_.dBaseSize; + uint32_t dValue = kBase; + uint32_t srcDValue = constInfo_.hSize; + uint32_t dstNzC0Stride = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数 + CopySingleMatrixNDToNZ(wL1Tensor[wkvL1Offset], wkvGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride); + CopySingleMatrixNDToNZ(wL1Tensor[wgateL1Offset], wgateGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride); +} + +template +__aicore__ inline void CompressorBlockCubePerf::LoadAToL0(const RunInfo &info, LocalTensor aL0Tensor, + LocalTensor xL1Tensor, uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize) +{ + uint32_t mSize = info.dealSeqCnt; + + uint32_t mSizeAlign = Align(mSize, 16U); + uint32_t xTensorOffset = kStart * mSizeAlign + mStart * (32 / sizeof(X_T)); + uint32_t mLoop = Align(mDealSize, 16U) / 16; + + for (uint32_t i = 0; i < mLoop; i++) { + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = i; + loadData2DParams.repeatTimes = kBase / (32 / sizeof(X_T)); + loadData2DParams.srcStride = mSizeAlign / 16; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = false; + LoadData(aL0Tensor[i * 16 * kBase], xL1Tensor[xTensorOffset], loadData2DParams); // 16: 一个分型的行数 + } +} + +template +__aicore__ inline void CompressorBlockCubePerf::LoadBToL0(LocalTensor bL0Tensor, LocalTensor wL1Tensor, + uint32_t kStart, uint32_t kBase) +{ + uint32_t rowCnt = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数 + uint64_t wTensorOffset = rowCnt * kStart; + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = (rowCnt / 16) * (kBase / (32 / sizeof(X_T))); + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = false; + LoadData(bL0Tensor, wL1Tensor[wTensorOffset], loadData2DParams); +} + +template +__aicore__ inline void CompressorBlockCubePerf::MatrixMmad(LocalTensor cL0Tensor, LocalTensor aL0Tensor, + LocalTensor bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C) +{ + MmadParams mmadParams; + mmadParams.m = (mActSize + 15) / 16 * 16; + mmadParams.n = nDealSize; + mmadParams.k = kActSize; + mmadParams.cmatrixInitVal = isInitL0C; + mmadParams.cmatrixSource = false; + Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams); + PipeBarrier(); +} + +template +__aicore__ inline void CompressorBlockCubePerf::CopyOutMm1Res(const RunInfo &info, LocalTensor cL0Tensor, + uint32_t coffId, uint32_t mStart, uint32_t mDealSize) +{ + // coffId=0, 存左矩阵的数据; coffId=1, 存右矩阵的数据 + FixpipeParamsV220 fixParams; + fixParams.mSize = mDealSize; + fixParams.nSize = constInfo_.dBaseSize; + fixParams.srcStride = (mDealSize + 15) / 16 * 16; // 需要16对齐 + fixParams.dstStride = (uint32_t)COMP::coff * constInfo_.headDim; + fixParams.ndNum = 1; + + uint64_t dbOffset = info.cubeDbIdx * constInfo_.dbSize; + uint64_t gmOffset = coffId * constInfo_.headDim + constInfo_.dIdx + mStart * fixParams.dstStride + dbOffset; + uint32_t kvOffset = 0; + uint32_t scoreOffset = (mDealSize + 15) / 16 * 16 * constInfo_.dBaseSize; + + Fixpipe(kvMm1ResGm[gmOffset], cL0Tensor[kvOffset], fixParams); + Fixpipe(scoreMm1ResGm[gmOffset], cL0Tensor[scoreOffset], fixParams); + +} + + +template +__aicore__ inline uint32_t CompressorBlockCubePerf::GetMSize(const RunInfo &info, uint32_t coffId) +{ + return info.dealSeqCnt; +} + +template +__aicore__ inline void CompressorBlockCubePerf::ComputeMm1(const RunInfo &info) +{ + static constexpr uint32_t K_SIZE = 512; + static constexpr uint32_t K_L1_BASE = 256; + static constexpr uint32_t M_L0_BASE = 128; + static constexpr uint32_t K_L0_BASE = 128; + uint32_t nCoff = (uint32_t)COMP::coff; + + // hSize为K_SIZE=512的倍数 + uint32_t hSize = constInfo_.hSize; + uint32_t hIdxStart = (constInfo_.aiCoreIdx % constInfo_.dBasicBlockNum) * K_L1_BASE; // 每组核内的h循环起始不同 + for (uint32_t h = 0; h < hSize; h += K_SIZE) { + for (uint32_t k = 0; k < K_SIZE; k += K_L1_BASE) { + bool isFirst = (h == 0 && k == 0); + bool isLast = ((h + K_SIZE >= hSize) && (k + K_L1_BASE >= K_SIZE)); + uint32_t hIdx = (h + k + hIdxStart) % hSize; // h方向错位搬运 + WaitFlag(X_EVENT0 + xBufId); + LocalTensor xL1Tensor = xBufL1.GetWithOffset(L1_X_SIZE / sizeof(X_T), xBufId * L1_X_SIZE); + CopyXGmToL1(info, xL1Tensor, hIdx, K_L1_BASE); + SetFlag(X_EVENT0 + xBufId); + WaitFlag(X_EVENT0 + xBufId); + for (uint32_t i = nCoff; i > 0; i--) { + // coffId=0, 计算pre数据; coffId=1, 计算cur数据 + uint32_t coffId = i - 1; + WaitFlag(W_EVENT0 + wBufId); + LocalTensor wL1Tensor = wBufL1.GetWithOffset(L1_W_SIZE / sizeof(X_T), wBufId * L1_W_SIZE); + CopyWeightGmToL1(wL1Tensor, hIdx, K_L1_BASE, coffId); + SetFlag(W_EVENT0 + wBufId); + WaitFlag(W_EVENT0 + wBufId); + + uint32_t mSize = GetMSize(info, coffId); + uint32_t actMDealSize = M_L0_BASE; + for (uint32_t mL0 = 0; mL0 < mSize; mL0 += M_L0_BASE) { + if (mL0 + M_L0_BASE > mSize) { + actMDealSize = mSize - mL0; + } + + l0cBufId = coffId + (mL0 / M_L0_BASE); + LocalTensor cL0Tensor = tmpBufL0C.GetWithOffset((L0C_PP_SIZE / sizeof(T)), l0cBufId * L0C_PP_SIZE); + if (isFirst) { + WaitFlag(L0C_EVENT0 + l0cBufId); + } + uint32_t nDealSize = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数 + for (uint32_t kL0 = 0; kL0 < K_L1_BASE; kL0 += K_L0_BASE) { + WaitFlag(L0AB_EVENT0 + l0abBufId); + LocalTensor aL0Tensor = tmpBufL0A.GetWithOffset(L0A_PP_SIZE / sizeof(X_T), l0abBufId * L0A_PP_SIZE); + LocalTensor bL0Tensor = tmpBufL0B.GetWithOffset(L0B_PP_SIZE / sizeof(X_T), l0abBufId * L0B_PP_SIZE); + LoadAToL0(info, aL0Tensor, xL1Tensor, kL0, K_L0_BASE, mL0, actMDealSize); + LoadBToL0(bL0Tensor, wL1Tensor, kL0, K_L0_BASE); + SetFlag(L0AB_EVENT0 + l0abBufId); + WaitFlag(L0AB_EVENT0 + l0abBufId); + bool isInitL0C = isFirst && (kL0 == 0); + MatrixMmad(cL0Tensor, aL0Tensor, bL0Tensor, actMDealSize, nDealSize, K_L0_BASE, isInitL0C); + SetFlag(L0AB_EVENT0 + l0abBufId); + l0abBufId = (l0abBufId + 1) % 2; + } + if (isLast) { + SetFlag(L0C_EVENT0 + l0cBufId); + WaitFlag(L0C_EVENT0 + l0cBufId); + CopyOutMm1Res(info, cL0Tensor, coffId, mL0, actMDealSize); + SetFlag(L0C_EVENT0 + l0cBufId); + } + } + + SetFlag(W_EVENT0 + wBufId); + wBufId = (wBufId + 1) % 4; + } + SetFlag(X_EVENT0 + xBufId); + xBufId = (xBufId + 1) % 2; + } + } + +} + +} // namespace Compressor + +#endif // COMPRESSOR_BLOCK_CUBE_PERF_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_block_vec_perf.h b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_block_vec_perf.h new file mode 100644 index 000000000..35d125046 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_block_vec_perf.h @@ -0,0 +1,1488 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_block_vec_perf.h + * \brief + */ + +#ifndef COMPRESSOR_BLOCK_VEC_PERF_H +#define COMPRESSOR_BLOCK_VEC_PERF_H + +#include "compressor_comm.h" +#include "compressor_tools.h" +#include "compressor_vector_comm.h" +#include "rms_norm.h" +#include "rope.h" +#include "soft_max.h" + + +using namespace AscendC; + +namespace Compressor { +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +struct LoopInfo { + uint32_t groupSize = 0U; + uint32_t groupNum = 0U; + uint32_t coreRowIdx = 0U; + uint32_t coreColIdx = 0U; + bool isCoreRowFirst = false; + bool isCoreRowLast = false; + bool isCoreLoopFirst = false; + bool isCoreLoopLast = false; +}; + +struct Vec1SplitInfo { + uint32_t dealSeqStartIdx = 0; + uint32_t dBaseSize = 0; + uint32_t vec1GroupSize = 0; + uint32_t vec1GroupNum = 0; + uint32_t dealTcSize = 0; + uint32_t preDealTcSize = 0; + uint32_t curBStart = 0; + uint32_t curSStart = 0; + uint32_t curCompressedCnt = 0; + uint32_t totalCompressedCnt = 0; + uint32_t tcSplitSize = 0; + uint32_t dSplitSize = 0; + uint32_t dLoopCount = 0; +}; + + +template +class CompressorBlockVectorPerf { +public: + static constexpr bool X_DTYPE = COMP::xDtype == X_DTYPE::BF16; + static constexpr uint64_t BLOCK_VEC_BASE_BUFFER_SIZE = 32 * 1024; // 32k + static constexpr uint32_t DATABLOCK_BYTES = 32; + static constexpr float FLOAT_ZERO = 0; + float SOFTMAX_MIN_NUM = static_cast(-1.0 / 0.0); + // =================================类型定义区================================= + // 中间计算数据类型为float，高精度模式 + using T = float; + using X_T = typename AscendC::Conditional::type; + using ROPE_T = typename AscendC::Conditional::type; + + __aicore__ inline CompressorBlockVectorPerf(){}; + // =================================设置参数================================= + __aicore__ inline void InitParams(const ConstInfo &constInfo, const CompressorTools &tools); + __aicore__ inline void Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut); + // =================================资源管理================================= + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + // =================================执行计算================================= + __aicore__ inline void ComputeVec1(const Vec1RunInfo &info); + __aicore__ inline uint32_t GetBasicNum(); + __aicore__ inline uint32_t GetScSize(); + __aicore__ inline void GetScIdxInfo(uint32_t bStart, uint32_t scStart, uint32_t dealScSize, uint32_t v2TcStart, + uint32_t v2TcEnd, uint32_t &outputBStart, uint32_t &outputSStart, + uint32_t &outputScSize); + __aicore__ inline void CalcScEndIdx(uint32_t bStart, uint32_t scStart, uint32_t dealScSize, uint32_t &bEnd, + uint32_t &scEnd); + __aicore__ inline void InitVec1GlobalTensor(GlobalTensor kvMm1ResGm, GlobalTensor scoreMm1ResGm, + GlobalTensor kvCacheTcGm, GlobalTensor scoreCacheTcGm, + GlobalTensor vec1ResGm, GlobalTensor vec2InputGm); + __aicore__ inline void ComputeVec2(const Vec2RunInfo &info); + +protected: + GlobalTensor vec1ResGm_; + GlobalTensor vec2InputGm_; + GlobalTensor scoreMm1ResGm_; + GlobalTensor kvMm1ResGm_; + GlobalTensor kvCacheTcGm_; + GlobalTensor scoreCacheTcGm_; + +private: + __aicore__ inline uint32_t GetSeqUsed(uint32_t bIdx); + __aicore__ inline uint32_t GetStartPos(uint32_t bIdx); + __aicore__ inline uint32_t GetSeqLength(uint32_t bIdx); + __aicore__ inline uint32_t GetBsLength(uint32_t index); + __aicore__ inline void CalcGlobalScStart(uint32_t bStart, uint32_t scStart, uint32_t bEnd, uint32_t scEnd, + uint64_t &globalScStart); + __aicore__ inline void UpdateOutputIdx(uint32_t &outputBStart, uint32_t &outputSStart, uint32_t &dealScSize, + uint32_t &curDealScSize); + __aicore__ inline void DealVec1BaseBlock(const Vec1RunInfo &info, CompressorVec1SliceIterator &sliceIterator, + const LoopInfo &loopInfo, uint32_t dStartIdx, uint32_t dDealSize, + uint32_t dBaseSize); + __aicore__ inline void CopyInApe(const LocalTensor &apeUb, uint32_t dStartIdx, uint32_t dDealSize); + __aicore__ inline void AddApeToScore(const LocalTensor &scoreLocal, const LocalTensor &apeUb, + const Vec1SliceInfo &sliceInfo, uint32_t dDealSize); + __aicore__ inline void AddSingleApeToScore(const LocalTensor &scoreLocal, const LocalTensor &apeUb, + const Vec1SliceInfo &sliceInfo, uint32_t dDealSize); + template + __aicore__ inline void DataCopyAlignUbToUb(const LocalTensor dstLocal, const LocalTensor srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, uint32_t srcSingleRowCount, + uint32_t dstSingleRowCount); + template + __aicore__ inline void DataCopyAlignGmToUb(const LocalTensor dstLocal, const GlobalTensor srcGm, + uint32_t copyRowCount, uint32_t copyColCount, uint32_t srcSingleRowCount, + uint32_t dstSingleRowCount); + template + __aicore__ inline void DataCopyAlignUbToGm(const GlobalTensor dstGm, const LocalTensor srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, uint32_t srcSingleRowCount, + uint32_t dstSingleRowCount); + template + __aicore__ inline void DataCopyWithOutputQue(const GlobalTensor dstGm, const LocalTensor srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount); + __aicore__ inline void PadAlign(const LocalTensor dstLocal, const LocalTensor srcLocal, + const Vec1SliceInfo &sliceInfo, uint32_t dStartIdx, uint32_t dDealSize); + template + __aicore__ inline void OverLap(const LocalTensor dstLocal, const LocalTensor srcLocal, + const GlobalTensor &srcGm, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const GlobalTensor &cacheTcGm, + const Vec1RunInfo &info, const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, uint32_t dStartIdx, + uint32_t globalSeqIdx, uint32_t dDealSize); + __aicore__ inline void FromWokrSpaceToUb(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + const Vec1SliceInfo &sliceInfo, const StatisticInfo &statisticInfo, + uint32_t dStartIdx, uint32_t dDealSize); + __aicore__ inline void WriteToCacheState(const GlobalTensor &state, const GlobalTensor &blockTableGm, + const LocalTensor &input, uint32_t batchIdx, uint32_t startSeqIdx, + uint32_t endSeqIdx, uint32_t dStartIdx, uint32_t dDealSize, uint32_t stateIdx); + __aicore__ inline void ReadFromCacheState(const LocalTensor &output, const GlobalTensor &state, + const GlobalTensor &blockTableGm, uint32_t batchIdx, + uint32_t startSeqIdx, uint32_t endSeqIdx, uint32_t dStartIdx, + uint32_t dDealSize, uint32_t stateIdx); + __aicore__ inline void SaveToWorkSpace(const LocalTensor srcLocal, const GlobalTensor &cacheTcGm, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, uint32_t dStartIdx, + uint32_t dDealSize); + __aicore__ inline void LoadFromWorkSpace(const LocalTensor dstLocal, const GlobalTensor &cacheTcGm, + const GlobalTensor &srcGm, const LocalTensor srcLocal, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, + uint32_t dStartIdx, uint32_t globalSeqIdx, uint32_t dDealSize); + __aicore__ inline void SoftmaxDN(const LocalTensor &scoreLocal, const LocalTensor &tmpUb, uint32_t tcDealSize, + uint32_t dDealSize); + __aicore__ inline void KvMulReduceScore(const LocalTensor &kvLocal, const LocalTensor &scoreLocal, + const LocalTensor &dstLocal, const LocalTensor &tmpUb, + uint32_t tcDealSize, uint32_t dDealSize); + __aicore__ inline void OverLapScoreKv(const LocalTensor &scoreLocal, const LocalTensor &kvLocal, + const Vec1RunInfo &info, + const LoopInfo &loopInfo, const StatisticInfo &statisticInfo, + const Vec1SliceInfo &originSliceInfo, uint32_t dStartIdx, uint32_t dDealSize, + uint32_t dBaseSize, uint32_t needDealTcSize); + __aicore__ inline void CopyOutVec1Res(const GlobalTensor &resGm, const Vec1RunInfo &info, + const LocalTensor comperssoredUb, uint32_t compressTcSize, + uint32_t dStartIdx, uint32_t dDealSize); + __aicore__ inline void CalcGroupInfo(const Vec1RunInfo &info, Vec1SplitInfo &splitInfo); + __aicore__ inline void CalcTaskDistribution(const Vec1RunInfo &info, Vec1SplitInfo &splitInfo); + __aicore__ inline void UpdateIteratorState(const Vec1RunInfo &info, Vec1SplitInfo &splitInfo); + __aicore__ inline void CalcTilingStrategy(Vec1SplitInfo &splitInfo); + __aicore__ inline Vec1SplitInfo SplitCoreV1(const Vec1RunInfo &info); + __aicore__ inline void SplitCoreV2(const Compressor::Vec2RunInfo &info); + __aicore__ inline void CopyFinalResultOut(const Compressor::Vec2RunInfo &info, const LocalTensor &cmpKvOutUb, + uint32_t startRow, uint32_t dealRowCount); + __aicore__ inline void DealVec2BaseBlock(const Compressor::Vec2RunInfo &info, uint32_t startRow, + uint32_t dealRowCount); + __aicore__ inline void MultRowRmsNorm(const LocalTensor &normResUb, const LocalTensor &vec1ResUb, + const LocalTensor &normWeightUb, const LocalTensor &tempLocal, + uint32_t dealRowCount); + __aicore__ inline void SingleCalRope(const LocalTensor &outputUb, const LocalTensor &normResUb, + uint32_t rowCnt, uint32_t curDealScSize, uint32_t globalScStart); + __aicore__ inline void CalRope(const LocalTensor &outputUb, const LocalTensor &normResUb, + uint32_t dealRowCount); + __aicore__ inline void SaveState(const LocalTensor &srcLocal, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const Vec1SliceInfo &sliceInfo, + uint32_t dStartIdx, uint32_t dDealSize, uint32_t stateIdx); + template + __aicore__ inline void DuplicateFirstBlock(const LocalTensor &dstLocal, uint32_t duplicateRowCount, + uint32_t duplicateColCount, uint32_t singleRowCount); + template + __aicore__ inline void ReadState(const LocalTensor &srcLocal, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const Vec1SliceInfo &sliceInfo, + uint32_t dStartIdx, uint32_t dDealSize, uint32_t stateIdx); + uint32_t cmpRatio_ = 0U; + uint32_t coff_ = 0U; + uint32_t curStartPos_ = 0; + uint32_t curActSeqLength_ = 0; + uint32_t compressedCnt_ = 0; + uint32_t v1SplitSize_ = 0; + uint32_t v1ScLoopTimes_ = 0; + uint32_t v1DLoopTimes_ = 0; + uint32_t dealTcNum_ = 0; + bool apeIsLoad_ = false; + bool isExistSeqUsed = false; + bool isExistStartPos = false; + // vec2 + uint32_t v2MBaseSize = 16; // Tc块数量：32 * 1024 / (512 * 4) + uint32_t v2TcStartIdx = 0U; + uint32_t v2TcEndIdx = 0U; + uint32_t mmResColSize_ = 128; + int64_t vec1ResGmStart = 0U; + uint32_t OutputBStartIdx, OutputSStartIdx, OutputSize; + CompressorTools tools_; + ConstInfo constInfo_ = {}; + MSplitInfo mSplitInfo = {}; + GlobalTensor startPosGm_; + GlobalTensor cuSeqlensGm_; + GlobalTensor sequsedGm_; + GlobalTensor stateBlockTableGm_; + GlobalTensor stateCacheGm_; + GlobalTensor apeGm_; + GlobalTensor normWeightGm_; + GlobalTensor ropeSinGm_; + GlobalTensor ropeCosGm_; + GlobalTensor cmpKvOutGm_; + + // ================================Local Buffer区==================================== + // TBuf mm1ResUb; + LocalTensor mm1ResTensor; + LocalTensor leftStateTensor; + LocalTensor rightStateTensor; + LocalTensor normWeightUb; + LocalTensor apeUb; + LocalTensor gatherOffsetCastUb; + // 临时tbuf + TBuf tmpBuff1; + TBuf tmpBuff2; + TBuf gatherOffsetBuf; + TBuf apeBuf; + // in queue + TQue inputQue1; + TBuf normWeightBuf; + // out queue + TQue outputQue1; +}; + + +template +__aicore__ inline void CompressorBlockVectorPerf::InitParams(const ConstInfo &constInfo, + const CompressorTools &tools) +{ + this->constInfo_ = constInfo; + this->tools_ = tools; + v2MBaseSize = BLOCK_VEC_BASE_BUFFER_SIZE / (constInfo_.headDim * sizeof(float)); +} + +template +__aicore__ inline void CompressorBlockVectorPerf::Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut) +{ + stateBlockTableGm_.SetGlobalBuffer((__gm__ int32_t *)stateBlockTable); + stateCacheGm_.SetGlobalBuffer((__gm__ T *)stateCache); + apeGm_.SetGlobalBuffer((__gm__ T *)ape); + normWeightGm_.SetGlobalBuffer((__gm__ X_T *)normWeight); + ropeSinGm_.SetGlobalBuffer((__gm__ ROPE_T *)ropeSin); + ropeCosGm_.SetGlobalBuffer((__gm__ ROPE_T *)ropeCos); + cmpKvOutGm_.SetGlobalBuffer((__gm__ X_T *)cmpKvOut); + isExistSeqUsed = (seqUsed != nullptr); + isExistStartPos = (startPos != nullptr); + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens); + } + if (isExistSeqUsed) { + sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed); + } + if (isExistStartPos) { + startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos); + } + coff_ = static_cast(COMP::coff); +} + +template +__aicore__ inline void CompressorBlockVectorPerf::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(inputQue1, 1, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(tmpBuff1, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(tmpBuff2, BUFFER_SIZE_BYTE_64K); + pipe->InitBuffer(outputQue1, 1, BUFFER_SIZE_BYTE_16K); + pipe->InitBuffer(normWeightBuf, BUFFER_SIZE_BYTE_4K); + pipe->InitBuffer(gatherOffsetBuf, BUFFER_SIZE_BYTE_1K); + pipe->InitBuffer(apeBuf, BUFFER_SIZE_BYTE_32K); + normWeightUb = normWeightBuf.Get(); + apeUb = apeBuf.Get(); + LocalTensor normweightInUb = inputQue1.AllocTensor(); + LocalTensor gatherOffsetUb = gatherOffsetBuf.Get(); + DataCopy(normweightInUb, normWeightGm_, constInfo_.headDim); // 获取normWeight，常驻 + inputQue1.EnQue(normweightInUb); + inputQue1.DeQue(); + Cast(normWeightUb, normweightInUb, RoundMode::CAST_NONE, constInfo_.headDim); + inputQue1.FreeTensor(normweightInUb); + if constexpr (COMP::rotaryMode == Compressor::ROTARY_MODE::INTERLEAVE) { + SetGatherSrcOffset(gatherOffsetUb, constInfo_.ropeHeadDim); + } + gatherOffsetCastUb = gatherOffsetUb.ReinterpretCast(); + PipeBarrier(); +} + +template +__aicore__ inline void CompressorBlockVectorPerf::AllocEventID() +{ +} + +template +__aicore__ inline void CompressorBlockVectorPerf::FreeEventID() +{ +} + +template +__aicore__ inline void +CompressorBlockVectorPerf::InitVec1GlobalTensor(GlobalTensor kvMm1ResGm, GlobalTensor scoreMm1ResGm, + GlobalTensor kvCacheTcGm, GlobalTensor scoreCacheTcGm, + GlobalTensor vec1ResGm, GlobalTensor vec2InputGm) +{ + this->kvMm1ResGm_ = kvMm1ResGm; + this->scoreMm1ResGm_ = scoreMm1ResGm; + this->kvCacheTcGm_ = kvCacheTcGm; + this->scoreCacheTcGm_ = scoreCacheTcGm; + this->vec1ResGm_ = vec1ResGm; + this->vec2InputGm_ = vec2InputGm; +} + +template +__aicore__ inline uint32_t CompressorBlockVectorPerf::GetSeqUsed(uint32_t bIdx) +{ + if (isExistSeqUsed) { + return (uint32_t)sequsedGm_.GetValue(bIdx); + } else { + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + return (uint32_t)(cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx)); + } else { + return constInfo_.sSize; + } + } +} + +template +__aicore__ inline uint32_t CompressorBlockVectorPerf::GetStartPos(uint32_t bIdx) +{ + if (isExistStartPos) { + return startPosGm_.GetValue(bIdx); + } + return 0; +} + +template +__aicore__ inline uint32_t CompressorBlockVectorPerf::GetSeqLength(uint32_t bIdx) +{ + if (isExistSeqUsed) { + return sequsedGm_.GetValue(bIdx); + } else if (COMP::xLayout == X_LAYOUT::TH) { + return cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx); + } else { + return constInfo_.sSize; + } +} + +template +__aicore__ inline uint32_t CompressorBlockVectorPerf::GetBsLength(uint32_t index) +{ + if (COMP::xLayout == X_LAYOUT::TH) { + return cuSeqlensGm_.GetValue(index); + } else { + return index * constInfo_.sSize; + } +} + +template +__aicore__ inline uint32_t CompressorBlockVectorPerf::GetBasicNum() +{ + // 获取 m方向上对应基本单元Tc的个数 + uint32_t curBasicNum = 0; + uint32_t headSize = 0; + if (curStartPos_ % constInfo_.cmpRatio != 0) { + headSize = constInfo_.cmpRatio - curStartPos_ % constInfo_.cmpRatio; + headSize = headSize > curActSeqLength_ ? curActSeqLength_ : headSize; + curBasicNum++; + } + // 加上中间整块及尾块 + curBasicNum += CeilDivT(curActSeqLength_ - headSize, constInfo_.cmpRatio); + return curBasicNum; +} + +template +__aicore__ inline uint32_t CompressorBlockVectorPerf::GetScSize() +{ + uint32_t curBasicNum = (curStartPos_ + curActSeqLength_) / constInfo_.cmpRatio - curStartPos_ / constInfo_.cmpRatio; + return curBasicNum; +} + +// 根据计算Tc开始结束索引 +template +__aicore__ inline void CompressorBlockVectorPerf::CalcScEndIdx(uint32_t bStart, uint32_t scStart, + uint32_t dealScSize, uint32_t &bEnd, + uint32_t &scEnd) +{ + uint32_t accScSize = 0; + for (int bIdx = bStart; bIdx < constInfo_.batchSize; ++bIdx) { + bEnd = bIdx; + // 计算起始batch的剩余块 + if (bIdx == bStart) { + curActSeqLength_ = GetSeqLength(bIdx); + curStartPos_ = GetStartPos(bIdx); + accScSize += GetScSize() - scStart; + if (accScSize >= dealScSize) { + scEnd = scStart + dealScSize; + return; + } + } else { + curActSeqLength_ = GetSeqLength(bIdx); + curStartPos_ = GetStartPos(bIdx); + uint32_t curBasicNum = GetScSize(); + uint32_t curBasicNumEnd = dealScSize - accScSize; + + if (accScSize + curBasicNum >= dealScSize) { + scEnd = curBasicNumEnd; + return; + } + accScSize += curBasicNum; + } + } +} + +// 根据sc的开始索引计算vec输出时的b、sc的索引 +template +__aicore__ inline void CompressorBlockVectorPerf::GetScIdxInfo(uint32_t bStart, uint32_t scStart, + uint32_t dealScSize, uint32_t v2TcStart, + uint32_t v2TcEnd, uint32_t &outputBStart, + uint32_t &outputSStart, uint32_t &outputScSize) +{ + outputScSize = v2TcEnd - v2TcStart; + uint32_t scEnd = 0; + uint32_t bEnd = 0; + CalcScEndIdx(bStart, scStart, v2TcStart, bEnd, scEnd); + outputSStart = scEnd; + outputBStart = bEnd; + // 处理跳batch + curActSeqLength_ = GetSeqLength(bEnd); + curStartPos_ = GetStartPos(bEnd); + uint32_t curScSize = GetScSize(); + if (curScSize == scEnd) { + outputSStart = 0; + outputBStart++; + } +} + +template +__aicore__ inline void CompressorBlockVectorPerf::CopyInApe(const LocalTensor &apeUb, uint32_t dStartIdx, + uint32_t dDealSize) +{ + LocalTensor apeUbTmp = inputQue1.AllocTensor(); + + uint32_t copyRowCount = coff_ * constInfo_.cmpRatio; + uint32_t copyColCount = dDealSize; + uint32_t dstSingleRowCount = dDealSize; + uint32_t srcSingleRowCount = constInfo_.headDim; + + uint64_t gmOffset = dStartIdx; + DataCopyAlignGmToUb(apeUbTmp, apeGm_[gmOffset], copyRowCount, copyColCount, srcSingleRowCount, dstSingleRowCount); + inputQue1.EnQue(apeUbTmp); + inputQue1.DeQue(); + DataCopy(apeUb, apeUbTmp, coff_ * dDealSize * constInfo_.cmpRatio); + inputQue1.FreeTensor(apeUbTmp); +} + +template +__aicore__ inline void +CompressorBlockVectorPerf::AddApeToScore(const LocalTensor &scoreLocal, const LocalTensor &apeUb, + const Vec1SliceInfo &sliceInfo, uint32_t dDealSize) +{ + uint32_t singleRowElemNum = dDealSize * coff_; + uint64_t scoreOffset = sliceInfo.dealedSeqCnt * singleRowElemNum; + + uint32_t tcDealSize = sliceInfo.dealTcSize; + if (sliceInfo.headHolderSeqCnt > 0) { + uint64_t apeOffset = sliceInfo.headHolderSeqCnt * singleRowElemNum; + uint32_t rCnt = tcDealSize == 1 ? sliceInfo.validSeqCnt * singleRowElemNum : + (constInfo_.cmpRatio - sliceInfo.headHolderSeqCnt) * singleRowElemNum; + Add(scoreLocal[scoreOffset], scoreLocal[scoreOffset], apeUb[apeOffset], rCnt); + scoreOffset += rCnt; + tcDealSize -= 1; + } + if (tcDealSize == 0) { + return; + } + if (sliceInfo.tailHolderSeqCnt > 0) { + tcDealSize -= 1; + uint64_t apeOffset = 0; + uint32_t rCnt = (constInfo_.cmpRatio - sliceInfo.tailHolderSeqCnt) * singleRowElemNum; + uint32_t tailScoreOffset = scoreOffset + tcDealSize * constInfo_.cmpRatio * singleRowElemNum; + Add(scoreLocal[tailScoreOffset], scoreLocal[tailScoreOffset], apeUb[apeOffset], rCnt); + } + if (tcDealSize == 0) { + return; + } + uint32_t rCnt = constInfo_.cmpRatio * singleRowElemNum; + for (uint32_t r = 0; r < tcDealSize; r++) { + Add(scoreLocal[scoreOffset + r * rCnt], scoreLocal[scoreOffset + r * rCnt], apeUb, rCnt); + } +} + +template +__aicore__ inline void +CompressorBlockVectorPerf::AddSingleApeToScore(const LocalTensor &scoreLocal, const LocalTensor &apeUb, + const Vec1SliceInfo &sliceInfo, uint32_t dDealSize) +{ + uint32_t SingleRowElemNum = dDealSize * coff_; + uint32_t dealRowCount = min(sliceInfo.sIdx, constInfo_.cmpRatio); + uint64_t scoreOffset = (constInfo_.cmpRatio - dealRowCount) * SingleRowElemNum; + uint64_t apeOffset = (constInfo_.cmpRatio - dealRowCount) * SingleRowElemNum; + for (uint32_t dOffset = 0; dOffset < dDealSize; dOffset += FP32_REPEAT_ELEMENT_NUM) { + uint32_t curAddColCount = min(dDealSize - dOffset, FP32_REPEAT_ELEMENT_NUM); + Add(scoreLocal[scoreOffset + dOffset], scoreLocal[scoreOffset + dOffset], apeUb[apeOffset + dOffset], + curAddColCount, dealRowCount, + {1, 1, 1, static_cast(SingleRowElemNum / FP32_BLOCK_ELEMENT_NUM), + static_cast(SingleRowElemNum / FP32_BLOCK_ELEMENT_NUM), + static_cast(SingleRowElemNum / FP32_BLOCK_ELEMENT_NUM)}); + } +} + +template +template +__aicore__ inline void +CompressorBlockVectorPerf::DataCopyAlignUbToUb(const LocalTensor dstLocal, const LocalTensor srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + DataCopyParams intriParams; + intriParams.blockCount = copyRowCount; + intriParams.blockLen = copyColCount / FP32_BLOCK_ELEMENT_NUM; + intriParams.dstGap = (dstSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + intriParams.srcGap = (srcSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + DataCopy(dstLocal, srcLocal, intriParams); +} + +template +template +__aicore__ inline void +CompressorBlockVectorPerf::DataCopyAlignGmToUb(const LocalTensor dstLocal, const GlobalTensor srcGm, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + DataCopyParams intriParams; + intriParams.blockCount = copyRowCount; + intriParams.blockLen = copyColCount / FP32_BLOCK_ELEMENT_NUM; + intriParams.dstGap = (dstSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + intriParams.srcGap = (srcSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + DataCopy(dstLocal, srcGm, intriParams); +} + +template +template +__aicore__ inline void +CompressorBlockVectorPerf::DataCopyAlignUbToGm(const GlobalTensor dstGm, const LocalTensor srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + DataCopyParams intriParams; + intriParams.blockCount = copyRowCount; + intriParams.blockLen = copyColCount / FP32_BLOCK_ELEMENT_NUM; + intriParams.dstGap = (dstSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + intriParams.srcGap = (srcSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + DataCopy(dstGm, srcLocal, intriParams); +} + +template +template +__aicore__ inline void +CompressorBlockVectorPerf::DataCopyWithOutputQue(const GlobalTensor dstGm, const LocalTensor srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + uint32_t singleCopyRowCount = BUFFER_SIZE_BYTE_16K / (copyColCount * sizeof(O)); + for (uint32_t rowCount = 0; rowCount < copyRowCount; rowCount += singleCopyRowCount) { + uint64_t srcOffset = rowCount * srcSingleRowCount; + uint64_t dstOffset = rowCount * dstSingleRowCount; + uint32_t curCopyRowCount = min(singleCopyRowCount, copyRowCount - rowCount); + + LocalTensor outputUb = outputQue1.AllocTensor(); + + DataCopyAlignUbToUb(outputUb, srcLocal[srcOffset], curCopyRowCount, copyColCount, srcSingleRowCount, + copyColCount); + PipeBarrier(); + + outputQue1.EnQue(outputUb); + outputQue1.DeQue(); + + DataCopyAlignUbToGm(dstGm[dstOffset], outputUb, curCopyRowCount, copyColCount, copyColCount, dstSingleRowCount); + + outputQue1.FreeTensor(outputUb); + } +} + +template +__aicore__ inline void +CompressorBlockVectorPerf::PadAlign(const LocalTensor dstLocal, const LocalTensor srcLocal, + const Vec1SliceInfo &sliceInfo, uint32_t dStartIdx, uint32_t dDealSize) +{ + // Ub data layout after overlap when r = 4 and coff = 2: + // Tc0_seq01: |--- --D_L--- -|------D_R-----| + // Tc0_seq02: |--- --D_L--- -|------D_R-----| + // Tc0_seq03: |--- --D_L--- -|------D_R-----| + // Tc0_seq04: |--- --D_L--- -|------D_R-----| + // Tc1_seq01: |--- --D_L--- -|------D_R-----| + // Tc1_seq02: |--- --D_L--- -|------D_R-----| + // Tc1_seq03: |--- --D_L--- -|------D_R-----| + // Tc1_seq04: |--- --D_L--- -|------D_R-----| + uint32_t srcSingleRowElemNum = dDealSize * coff_; + uint32_t copyRowCount = sliceInfo.compressTcSize * constInfo_.cmpRatio - sliceInfo.headHolderSeqCnt; + uint32_t copyColCount = dDealSize; + uint32_t srcSingleRowCount = srcSingleRowElemNum; + uint32_t dstSingleRowCount = srcSingleRowElemNum; // left和right在seq方向是交错存储的 + uint64_t srcLocalOffset = sliceInfo.dealedSeqCnt * srcSingleRowElemNum; + + uint64_t dstUbOffset = sliceInfo.compressoredScCnt * constInfo_.cmpRatio * dstSingleRowCount; + if constexpr (COMP::coff == COFF::OVERLAP) { + // 左侧 + uint64_t preSrcLocalOffset = srcLocalOffset; + uint64_t preDstUbOffset = dstUbOffset + (sliceInfo.headHolderSeqCnt + constInfo_.cmpRatio) * dstSingleRowCount; + DataCopyAlignUbToUb(dstLocal[preDstUbOffset], srcLocal[preSrcLocalOffset], + copyRowCount - min(copyRowCount, constInfo_.cmpRatio), copyColCount, srcSingleRowCount, + dstSingleRowCount); + dstUbOffset += dDealSize; + srcLocalOffset += dDealSize; + } + // 右侧 + dstUbOffset += sliceInfo.headHolderSeqCnt * dstSingleRowCount; + DataCopyAlignUbToUb(dstLocal[dstUbOffset], srcLocal[srcLocalOffset], copyRowCount, copyColCount, srcSingleRowCount, + dstSingleRowCount); +} + + +template +template +__aicore__ inline void +CompressorBlockVectorPerf::OverLap(const LocalTensor dstLocal, const LocalTensor srcLocal, + const GlobalTensor &srcGm, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const GlobalTensor &cacheTcGm, + const Vec1RunInfo &info, const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, uint32_t dStartIdx, + uint32_t globalSeqIdx, uint32_t dDealSize) +{ + if (sliceInfo.dealTcSize == 0) { + return; + } + + if constexpr (IS_SCORE) { + AddApeToScore(srcLocal, apeUb, sliceInfo, dDealSize); + PipeBarrier(); + } + SaveState(srcLocal, stateGm, blockTableGm, sliceInfo, dStartIdx, dDealSize, static_cast(IS_SCORE)); + + event_t eventId_V_MTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventId_V_MTE2); + WaitFlag(eventId_V_MTE2); + ReadState(dstLocal, stateGm, blockTableGm, sliceInfo, dStartIdx, dDealSize, static_cast(IS_SCORE)); + + if constexpr (COMP::coff == COFF::OVERLAP) { + uint32_t nextC1V1DbIdx = (info.c1v1DbIdx + 1) % constInfo_.dbWorkspaceRatio; + GlobalTensor nextCacheTcGm = cacheTcGm[nextC1V1DbIdx * constInfo_.cmpRatio * constInfo_.headDim]; + SaveToWorkSpace(srcLocal, nextCacheTcGm, sliceInfo, loopInfo, dStartIdx, dDealSize); + } + if (sliceInfo.compressTcSize > 0) { + PadAlign(dstLocal, srcLocal, sliceInfo, dStartIdx, dDealSize); + if constexpr (COMP::coff == COFF::OVERLAP) { + event_t eventId_MTE3_MTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventId_MTE3_MTE2); + WaitFlag(eventId_MTE3_MTE2); + GlobalTensor curCacheTcGm = cacheTcGm[info.c1v1DbIdx * constInfo_.cmpRatio * constInfo_.headDim]; + LoadFromWorkSpace(dstLocal, curCacheTcGm, srcGm, srcLocal, sliceInfo, loopInfo, dStartIdx, globalSeqIdx, + dDealSize); + } + } + event_t eventId_MTE2_V = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId_MTE2_V); + WaitFlag(eventId_MTE2_V); +} + +template +__aicore__ inline void +CompressorBlockVectorPerf::FromWokrSpaceToUb(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + const Vec1SliceInfo &sliceInfo, const StatisticInfo &statisticInfo, + uint32_t dStartIdx, uint32_t dDealSize) +{ + uint32_t srcSingleRowElemNum = constInfo_.headDim; + uint32_t copyRowCount = statisticInfo.dealSeqCnt * coff_; + uint32_t copyColCount = dDealSize; + uint32_t srcSingleRowCount = srcSingleRowElemNum; + uint32_t dstSingleRowCount = dDealSize; + uint64_t srcGmOffset = sliceInfo.dealedSeqCnt * srcSingleRowElemNum * coff_ + dStartIdx; + DataCopyAlignGmToUb(dstLocal, srcGm[srcGmOffset], copyRowCount, copyColCount, srcSingleRowCount, dstSingleRowCount); +} + +template +__aicore__ inline void +CompressorBlockVectorPerf::SaveToWorkSpace(const LocalTensor srcLocal, const GlobalTensor &cacheTcGm, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, + uint32_t dStartIdx, uint32_t dDealSize) +{ + uint32_t curSeqLen = sliceInfo.bStartPos + sliceInfo.sIdx + sliceInfo.validSeqCnt; + uint32_t totalSeqLen = sliceInfo.bStartPos + sliceInfo.sIdx + sliceInfo.bSeqUsed; + if (!loopInfo.isCoreRowLast || !loopInfo.isCoreLoopLast || !sliceInfo.isLast || totalSeqLen < constInfo_.cmpRatio || + curSeqLen > Trunc(totalSeqLen, constInfo_.cmpRatio) - constInfo_.cmpRatio) { + return; + } + uint32_t srcSingleRowElemNum = dDealSize * coff_; + uint64_t srcLocalOffset = + (sliceInfo.dealedSeqCnt + sliceInfo.validSeqCnt - min(sliceInfo.validSeqCnt, constInfo_.cmpRatio)) * + srcSingleRowElemNum; + DataCopyWithOutputQue(cacheTcGm[dStartIdx], srcLocal[srcLocalOffset], + curSeqLen - max(curSeqLen - constInfo_.cmpRatio, sliceInfo.bStartPos), dDealSize, + coff_ * dDealSize, constInfo_.headDim); +} + +template +__aicore__ inline void +CompressorBlockVectorPerf::LoadFromWorkSpace(const LocalTensor dstLocal, const GlobalTensor &cacheTcGm, + const GlobalTensor &srcGm, const LocalTensor srcLocal, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, + uint32_t dStartIdx, uint32_t globalSeqIdx, uint32_t dDealSize) +{ + if (sliceInfo.sIdx == 0) { + return; + } + uint32_t dstSingleRowElemNum = dDealSize * coff_; + uint32_t copyRowCount = min(sliceInfo.sIdx, constInfo_.cmpRatio); + uint64_t dstLocalOffset = + (sliceInfo.compressoredScCnt * constInfo_.cmpRatio + constInfo_.cmpRatio - copyRowCount) * dstSingleRowElemNum; + if (loopInfo.isCoreRowFirst && loopInfo.isCoreLoopFirst && sliceInfo.isFirst) { // 从cacheGm获取 + uint32_t srcSingleRowElemNum = constInfo_.headDim * coff_; + uint64_t srcLocalOffset = dStartIdx; + DataCopyAlignGmToUb(dstLocal[dstLocalOffset], cacheTcGm[srcLocalOffset], copyRowCount, dDealSize, + constInfo_.headDim, coff_ * dDealSize); + } else if (sliceInfo.isFirst) { // 从存放MatMul结果的WorkSpace中获取 + uint32_t srcSingleRowElemNum = constInfo_.headDim * coff_; + uint64_t srcLocalOffset = + (globalSeqIdx + sliceInfo.dealedSeqCnt - copyRowCount) * srcSingleRowElemNum + dStartIdx; + DataCopyAlignGmToUb(dstLocal[dstLocalOffset], srcGm[srcLocalOffset], copyRowCount, dDealSize, + coff_ * constInfo_.headDim, coff_ * dDealSize); + } else { // 从UB中获取 + uint32_t srcSingleRowElemNum = dDealSize * coff_; + uint64_t srcLocalOffset = (sliceInfo.dealedSeqCnt - copyRowCount) * srcSingleRowElemNum; + DataCopyAlignUbToUb(dstLocal[dstLocalOffset], srcLocal[srcLocalOffset], copyRowCount, dDealSize, + coff_ * dDealSize, coff_ * dDealSize); + } +} + +template +__aicore__ inline void CompressorBlockVectorPerf::ReadFromCacheState( + const LocalTensor &output, const GlobalTensor &state, const GlobalTensor &blockTableGm, + uint32_t batchIdx, uint32_t startSeqIdx, uint32_t endSeqIdx, uint32_t dStartIdx, uint32_t dDealSize, uint32_t stateIdx) +{ + uint64_t blockTablebaseOffset = batchIdx * constInfo_.maxBlockNumPerBatch; + uint32_t curSeqIdx = startSeqIdx; + uint32_t copyFinishRowCnt = 0; + uint32_t seqCnt = endSeqIdx - startSeqIdx; + while (copyFinishRowCnt < seqCnt) { + uint64_t blockIdOffset = curSeqIdx / constInfo_.blockSize; + uint64_t remainRowCnt = curSeqIdx % constInfo_.blockSize; + uint64_t idInBlockTable = blockTableGm.GetValue(blockTablebaseOffset + blockIdOffset); + uint32_t copyRowCount = constInfo_.blockSize - remainRowCnt; + if (copyFinishRowCnt + copyRowCount > seqCnt) { + copyRowCount = seqCnt - copyFinishRowCnt; + } + uint64_t stateOffset = idInBlockTable * constInfo_.stateCacheStrideDim0 + + remainRowCnt * 2 * coff_ * constInfo_.headDim + + stateIdx * coff_ * constInfo_.headDim + dStartIdx; + + DataCopyAlignGmToUb(output[copyFinishRowCnt * coff_ * dDealSize], state[stateOffset], copyRowCount, + dDealSize, coff_ * constInfo_.headDim * 2, coff_ * dDealSize); + copyFinishRowCnt += copyRowCount; + curSeqIdx += copyRowCount; + } +} + +template +__aicore__ inline void CompressorBlockVectorPerf::WriteToCacheState( + const GlobalTensor &state, const GlobalTensor &blockTableGm, const LocalTensor &input, + uint32_t batchIdx, uint32_t startSeqIdx, uint32_t endSeqIdx, uint32_t dStartIdx, uint32_t dDealSize, uint32_t stateIdx) +{ + uint64_t blockTablebaseOffset = batchIdx * constInfo_.maxBlockNumPerBatch; + uint32_t curSeqIdx = startSeqIdx; + uint32_t copyFinishRowCnt = 0; + uint32_t seqCnt = endSeqIdx - startSeqIdx; + while (copyFinishRowCnt < seqCnt) { + uint64_t blockIdOffset = curSeqIdx / constInfo_.blockSize; + uint64_t remainRowCnt = curSeqIdx % constInfo_.blockSize; + uint64_t idInBlockTable = blockTableGm.GetValue(blockTablebaseOffset + blockIdOffset); + uint32_t copyRowCount = constInfo_.blockSize - remainRowCnt; + if (copyFinishRowCnt + copyRowCount > seqCnt) { + copyRowCount = seqCnt - copyFinishRowCnt; + } + if (idInBlockTable != 0) { // 32 + uint64_t stateOffset = idInBlockTable * constInfo_.stateCacheStrideDim0 + + remainRowCnt * 2 * coff_ * constInfo_.headDim + + stateIdx * coff_ * constInfo_.headDim + dStartIdx; + DataCopyWithOutputQue(state[stateOffset], input[copyFinishRowCnt * coff_ * dDealSize], copyRowCount, + dDealSize, coff_ * dDealSize, coff_ * constInfo_.headDim * 2); + } + + copyFinishRowCnt += copyRowCount; + curSeqIdx += copyRowCount; + } +} + +template +__aicore__ inline void +CompressorBlockVectorPerf::SaveState(const LocalTensor &srcLocal, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const Vec1SliceInfo &sliceInfo, + uint32_t dStartIdx, uint32_t dDealSize, uint32_t stateIdx) +{ + uint32_t startSeqIdx = sliceInfo.bStartPos + sliceInfo.sIdx; + uint32_t endSeqIdx = startSeqIdx + sliceInfo.validSeqCnt; + uint64_t srcBaseOffset = sliceInfo.dealedSeqCnt * coff_ * dDealSize; + + if constexpr (COMP::coff == COFF::OVERLAP) { + WriteToCacheState(stateGm, blockTableGm, srcLocal[srcBaseOffset], sliceInfo.bIdx, startSeqIdx, endSeqIdx, + dStartIdx, dDealSize, stateIdx); + srcBaseOffset += dDealSize; + dStartIdx += constInfo_.headDim; + } + + WriteToCacheState(stateGm, blockTableGm, srcLocal[srcBaseOffset], sliceInfo.bIdx, startSeqIdx, endSeqIdx, dStartIdx, + dDealSize, stateIdx); +} + +template +template +__aicore__ inline void +CompressorBlockVectorPerf::DuplicateFirstBlock(const LocalTensor &dstLocal, uint32_t duplicateRowCount, + uint32_t duplicateColCount, uint32_t singleRowCount) +{ + for (uint32_t offset = 0; offset < duplicateColCount; offset += FP32_REPEAT_ELEMENT_NUM) { + uint32_t curDuplicateColCount = min(duplicateColCount - offset, FP32_REPEAT_ELEMENT_NUM); + if constexpr (IS_SCORE) { + Duplicate(dstLocal[offset], SOFTMAX_MIN_NUM, curDuplicateColCount, duplicateRowCount, 1, + singleRowCount / REPEAT_STRIDE_NUM); + } else { + Duplicate(dstLocal[offset], FLOAT_ZERO, curDuplicateColCount, duplicateRowCount, 1, + singleRowCount / REPEAT_STRIDE_NUM); + } + } +} + + +template +template +__aicore__ inline void +CompressorBlockVectorPerf::ReadState(const LocalTensor &dstLocal, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const Vec1SliceInfo &sliceInfo, + uint32_t dStartIdx, uint32_t dDealSize, uint32_t stateIdx) +{ + // 没有需要压缩的块时, 不需要读state的信息 + if (sliceInfo.compressTcSize == 0) { + return; + } + // 填充右边 + if (sliceInfo.headHolderSeqCnt > 0) { + // 整个batch的第一块 + uint32_t startSeqIdx = Trunc(sliceInfo.bStartPos + sliceInfo.sIdx, constInfo_.cmpRatio); + uint32_t endSeqIdx = sliceInfo.bStartPos; + uint64_t dstBaseOffset = sliceInfo.compressoredScCnt * constInfo_.cmpRatio * coff_ * dDealSize; + if constexpr (COMP::coff == Compressor::COFF::OVERLAP) { + dstBaseOffset += (coff_ - 1) * dDealSize; + } + ReadFromCacheState(dstLocal[dstBaseOffset], stateGm, blockTableGm, sliceInfo.bIdx, startSeqIdx, endSeqIdx, + dStartIdx + (coff_ - 1) * constInfo_.headDim, dDealSize, stateIdx); + } + + // 填充左边 + if constexpr (COMP::coff == Compressor::COFF::OVERLAP) { + bool isFirst = sliceInfo.bStartPos + sliceInfo.sIdx < constInfo_.cmpRatio; + if (isFirst) { + // 无历史数据 + // dDealSize必须为64 + uint64_t dstBaseOffset = sliceInfo.compressoredScCnt * constInfo_.cmpRatio * coff_ * dDealSize; + DuplicateFirstBlock(dstLocal[dstBaseOffset], constInfo_.cmpRatio, dDealSize, coff_ * dDealSize); + } + if (sliceInfo.sIdx < constInfo_.cmpRatio && (!isFirst || sliceInfo.compressTcSize > 1)) { + uint32_t startSeqIdx = + sliceInfo.bStartPos < constInfo_.cmpRatio ? + 0 : + Trunc(sliceInfo.bStartPos + sliceInfo.sIdx, constInfo_.cmpRatio) - constInfo_.cmpRatio; + uint32_t endSeqIdx = + min(Trunc(sliceInfo.bStartPos + sliceInfo.sIdx + sliceInfo.validSeqCnt, constInfo_.cmpRatio) - + constInfo_.cmpRatio, + sliceInfo.bStartPos); + uint64_t dstBaseOffset = sliceInfo.compressoredScCnt * constInfo_.cmpRatio * coff_ * dDealSize; + if (isFirst) { + dstBaseOffset += constInfo_.cmpRatio * coff_ * dDealSize; + } + ReadFromCacheState(dstLocal[dstBaseOffset], stateGm, blockTableGm, sliceInfo.bIdx, startSeqIdx, endSeqIdx, + dStartIdx, dDealSize, stateIdx); + } + } +} + +template +__aicore__ inline void CompressorBlockVectorPerf::SoftmaxDN(const LocalTensor &scoreLocal, + const LocalTensor &tmpUb, uint32_t tcDealSize, + uint32_t dDealSize) +{ + float minValue = -2e38; + uint32_t ReduceSize = coff_ * constInfo_.cmpRatio; + uint32_t rCnt = ReduceSize * dDealSize; + for (uint32_t r = 0; r < tcDealSize; r++) { + ColumnSoftMax(scoreLocal[r * rCnt], scoreLocal[r * rCnt], tmpUb[r * rCnt], ReduceSize, dDealSize); + } +} + +template +__aicore__ inline void +CompressorBlockVectorPerf::KvMulReduceScore(const LocalTensor &kvLocal, const LocalTensor &scoreLocal, + const LocalTensor &dstLocal, const LocalTensor &tmpUb, + uint32_t tcDealSize, uint32_t dDealSize) +{ + uint32_t ReduceSize = coff_ * constInfo_.cmpRatio; + uint32_t rCnt = ReduceSize * dDealSize; + Mul(kvLocal, kvLocal, scoreLocal, tcDealSize * rCnt); + PipeBarrier(); + for (uint32_t r = 0; r < tcDealSize; r++) { + ColumnSum(dstLocal[r * dDealSize], kvLocal[r * rCnt], tmpUb[r * rCnt], ReduceSize, dDealSize); + } +} + +template +__aicore__ inline void +CompressorBlockVectorPerf::CopyOutVec1Res(const GlobalTensor &resGm, const Vec1RunInfo &info, + const LocalTensor comperssoredUb, uint32_t compressTcSize, + uint32_t dStartIdx, uint32_t dDealSize) +{ + uint64_t outGmOffset = compressedCnt_ * constInfo_.headDim + dStartIdx; + DataCopyAlignUbToGm(resGm[outGmOffset], comperssoredUb, compressTcSize, dDealSize, dDealSize, constInfo_.headDim); +} + +template +__aicore__ inline void CompressorBlockVectorPerf::OverLapScoreKv( + const LocalTensor &scoreLocal, const LocalTensor &kvLocal, const Vec1RunInfo &info, + const LoopInfo &loopInfo, const StatisticInfo &statisticInfo, + const Vec1SliceInfo &originSliceInfo, uint32_t dStartIdx, uint32_t dDealSize, uint32_t dBaseSize, + uint32_t needDealTcSize) +{ + CompressorVec1SliceIterator overLapSliceIterator(tools_); + overLapSliceIterator.SetMaxBatchSize(constInfo_.batchSize); + Vec1SliceInfo &overLapSliceInfo = overLapSliceIterator.GetSlice(); + + GlobalTensor scoreDBMm1ResGm = scoreMm1ResGm_[info.c1v1DbIdx * constInfo_.dbSize]; + LocalTensor scoreUb = inputQue1.AllocTensor(); + FromWokrSpaceToUb(scoreUb, scoreDBMm1ResGm, originSliceInfo, statisticInfo, dStartIdx, dDealSize); + inputQue1.EnQue(scoreUb); + inputQue1.DeQue(); + overLapSliceIterator.Reset(originSliceInfo.bIdx, originSliceInfo.sIdx, 0U, 0U); + overLapSliceIterator.SetNeedDealTcSize(needDealTcSize); + while (!overLapSliceIterator.IsEnd()) { + overLapSliceIterator.GetSlice(); + OverLap(scoreLocal, scoreUb, scoreDBMm1ResGm, stateCacheGm_, stateBlockTableGm_, scoreCacheTcGm_, + info, overLapSliceInfo, loopInfo, dStartIdx, originSliceInfo.dealedSeqCnt, dDealSize); + overLapSliceIterator.IteratorSlice(); + } + inputQue1.FreeTensor(scoreUb); + + if constexpr (COMP::coff == COFF::OVERLAP) { + if (originSliceInfo.sIdx != 0 && originSliceInfo.compressTcSize > 0 && + (!loopInfo.isCoreRowFirst || !loopInfo.isCoreLoopFirst)) { + AddSingleApeToScore(scoreLocal, apeUb, originSliceInfo, dDealSize); + } + } + + GlobalTensor kvDBMm1ResGm = kvMm1ResGm_[info.c1v1DbIdx * constInfo_.dbSize]; + LocalTensor kvUb = inputQue1.AllocTensor(); + FromWokrSpaceToUb(kvUb, kvDBMm1ResGm, originSliceInfo, statisticInfo, dStartIdx, dDealSize); + + inputQue1.EnQue(kvUb); + inputQue1.DeQue(); + overLapSliceIterator.Reset(originSliceInfo.bIdx, originSliceInfo.sIdx, 0U, 0U); + overLapSliceIterator.SetNeedDealTcSize(needDealTcSize); + while (!overLapSliceIterator.IsEnd()) { + overLapSliceIterator.GetSlice(); + OverLap(kvLocal, kvUb, kvDBMm1ResGm, stateCacheGm_, stateBlockTableGm_, kvCacheTcGm_, info, overLapSliceInfo, + loopInfo, dStartIdx, originSliceInfo.dealedSeqCnt, dDealSize); + overLapSliceIterator.IteratorSlice(); + } + inputQue1.FreeTensor(kvUb); + PipeBarrier(); +} + +template +__aicore__ inline void CompressorBlockVectorPerf::DealVec1BaseBlock( + const Vec1RunInfo &info, CompressorVec1SliceIterator &sliceIterator, const LoopInfo &loopInfo, + uint32_t dStartIdx, uint32_t dDealSize, uint32_t dBaseSize) +{ + Vec1SliceInfo originSliceInfo = sliceIterator.GetSlice(); + uint32_t needDealTcSize = sliceIterator.GetNeedDealTcSize(); + StatisticInfo &statisticInfo = sliceIterator.template FullIteratorSlice(); + if (statisticInfo.actualTcCnt == 0) { + return; + } + LocalTensor scoreLocal = tmpBuff1.Get(); + LocalTensor kvLocal = tmpBuff2.Get(); + + OverLapScoreKv(scoreLocal, kvLocal, info, loopInfo, statisticInfo, originSliceInfo, dStartIdx, + dDealSize, dBaseSize, needDealTcSize); + + if (statisticInfo.compressorScCnt > 0) { + LocalTensor tmpUb = kvLocal[BUFFER_SIZE_BYTE_32K / sizeof(T)]; + SoftmaxDN(scoreLocal, tmpUb, statisticInfo.compressorScCnt, dDealSize); + LocalTensor comperssoredUb = outputQue1.AllocTensor(); + PipeBarrier(); + KvMulReduceScore(kvLocal, scoreLocal, comperssoredUb, tmpUb, statisticInfo.compressorScCnt, dDealSize); + PipeBarrier(); + outputQue1.EnQue(comperssoredUb); + outputQue1.DeQue(); + GlobalTensor resGm = vec1ResGm_[info.v1v2DbIdx * constInfo_.dbSize]; + CopyOutVec1Res(resGm, info, comperssoredUb, statisticInfo.compressorScCnt, dStartIdx, dDealSize); + outputQue1.FreeTensor(comperssoredUb); + } + compressedCnt_ += statisticInfo.compressorScCnt; +} + +template +__aicore__ inline void CompressorBlockVectorPerf::CalcGroupInfo(const Vec1RunInfo &info, Vec1SplitInfo &splitInfo) +{ + uint32_t aiCoreNum = constInfo_.usedCoreNum * 2; + splitInfo.dBaseSize = constInfo_.headDim / min(FloorPow2(aiCoreNum), CeilPow2(CeilDivT(aiCoreNum, info.dealTcNum))); + splitInfo.dBaseSize = max(splitInfo.dBaseSize, FP32_BLOCK_ELEMENT_NUM); + splitInfo.vec1GroupSize = constInfo_.headDim / splitInfo.dBaseSize; + splitInfo.vec1GroupNum = min(static_cast(aiCoreNum / splitInfo.vec1GroupSize), info.dealTcNum); +} + +template +__aicore__ inline void CompressorBlockVectorPerf::CalcTaskDistribution(const Vec1RunInfo &info, + Vec1SplitInfo &splitInfo) +{ + uint32_t blockIdx = GetBlockIdx(); + uint32_t groupSize = splitInfo.vec1GroupSize; + uint32_t groupNum = splitInfo.vec1GroupNum; + uint32_t dealTcNum = info.dealTcNum; + + if (blockIdx < groupSize * (dealTcNum % groupNum)) { + splitInfo.dealTcSize = dealTcNum / groupNum + 1; + splitInfo.preDealTcSize = splitInfo.dealTcSize * (blockIdx / groupSize); + } else if (blockIdx < groupSize * groupNum) { + splitInfo.dealTcSize = dealTcNum / groupNum; + splitInfo.preDealTcSize = splitInfo.dealTcSize * (blockIdx / groupSize) + dealTcNum % groupNum; + } else { + splitInfo.dealTcSize = 0; + splitInfo.preDealTcSize = dealTcNum; + } +} + +template +__aicore__ inline void CompressorBlockVectorPerf::UpdateIteratorState(const Vec1RunInfo &info, + Vec1SplitInfo &splitInfo) +{ + CompressorVec1SliceIterator sliceIterator(tools_); + sliceIterator.SetMaxBatchSize(constInfo_.batchSize); + sliceIterator.Reset(info.bStart, info.sStart, 0U, 0U); + Vec1SliceInfo &sliceInfo = sliceIterator.GetSlice(); + + // 处理前序任务量，更新起始索引 + if (splitInfo.preDealTcSize > 0) { + sliceIterator.SetNeedDealTcSize(splitInfo.preDealTcSize); + StatisticInfo &statisticInfo = sliceIterator.template FullIteratorSlice(); + splitInfo.curCompressedCnt = statisticInfo.compressorScCnt; + splitInfo.dealSeqStartIdx = sliceInfo.dealedSeqCnt; + splitInfo.curBStart = sliceInfo.bIdx; + splitInfo.curSStart = sliceInfo.sIdx; + } else { + splitInfo.curCompressedCnt = 0; + splitInfo.dealSeqStartIdx = 0; + splitInfo.curBStart = info.bStart; + splitInfo.curSStart = info.sStart; + } + + // 处理当前核实际要跑的任务量 + sliceIterator.SetNeedDealTcSize(info.dealTcNum - splitInfo.preDealTcSize); + StatisticInfo &statisticInfo = sliceIterator.template FullIteratorSlice(); + splitInfo.totalCompressedCnt = splitInfo.curCompressedCnt + statisticInfo.compressorScCnt; +} + +template +__aicore__ inline void CompressorBlockVectorPerf::CalcTilingStrategy(Vec1SplitInfo &splitInfo) +{ + // 计算headDim和Tc方向切分大小 + uint32_t maxDealColNum = BUFFER_SIZE_BYTE_32K / (constInfo_.cmpRatio * coff_ * sizeof(T)); + + // 切块逻辑 + if (maxDealColNum < splitInfo.dBaseSize) { + splitInfo.tcSplitSize = 1; + splitInfo.dLoopCount = CeilDivT(splitInfo.dBaseSize, maxDealColNum); + splitInfo.dSplitSize = splitInfo.dBaseSize / splitInfo.dLoopCount; + } else { + splitInfo.dSplitSize = splitInfo.dBaseSize; + splitInfo.dLoopCount = splitInfo.dBaseSize / splitInfo.dSplitSize; // 此处常等于1，保留原逻辑 + splitInfo.tcSplitSize = maxDealColNum / splitInfo.dBaseSize; + } +} + +template +__aicore__ inline Vec1SplitInfo CompressorBlockVectorPerf::SplitCoreV1(const Vec1RunInfo &info) +{ + Vec1SplitInfo splitInfo; + + // 1. 计算基础分组和分片大小 + CalcGroupInfo(info, splitInfo); + + // 2. 根据当前的 BlockIdx 计算任务分配（负载均衡） + CalcTaskDistribution(info, splitInfo); + + // 3. 刷新迭代器并获取当前核的起始位置状态 + UpdateIteratorState(info, splitInfo); + + if (splitInfo.dealTcSize == 0) { + return splitInfo; + } + + // 4. 计算具体在内存中的切块（Tiling）逻辑 + CalcTilingStrategy(splitInfo); + + return splitInfo; +} + +template +__aicore__ inline void CompressorBlockVectorPerf::ComputeVec1(const Vec1RunInfo &info) +{ + if (info.dealTcNum == 0) { + return; + } + if (info.resetResFlag) { + compressedCnt_ = 0; + } + uint32_t preCompressedCnt = compressedCnt_; + Vec1SplitInfo splitInfo = SplitCoreV1(info); + // 计算当前VecCore的任务量 + if (splitInfo.dealTcSize == 0) { + compressedCnt_ += splitInfo.totalCompressedCnt; + return; + } + + LoopInfo loopInfo; + loopInfo.groupSize = splitInfo.vec1GroupSize; + loopInfo.groupNum = splitInfo.vec1GroupNum; + loopInfo.coreRowIdx = GetBlockIdx() / splitInfo.vec1GroupSize; + loopInfo.coreColIdx = GetBlockIdx() % splitInfo.vec1GroupSize; + loopInfo.isCoreRowLast = loopInfo.coreRowIdx == splitInfo.vec1GroupNum - 1; + loopInfo.isCoreRowFirst = loopInfo.coreRowIdx == 0; + + + CompressorVec1SliceIterator sliceIterator(tools_); + sliceIterator.SetMaxBatchSize(constInfo_.batchSize); + // 切块循环 + uint64_t baseOffset = loopInfo.coreColIdx * splitInfo.dBaseSize; + for (uint32_t dLoopIdx = 0; dLoopIdx < splitInfo.dLoopCount; dLoopIdx++) { + uint64_t dBaseOffset = baseOffset + dLoopIdx * splitInfo.dSplitSize; + + CopyInApe(apeUb, dBaseOffset, splitInfo.dSplitSize); + + sliceIterator.Reset(splitInfo.curBStart, splitInfo.curSStart, splitInfo.dealSeqStartIdx, 0U); + compressedCnt_ = preCompressedCnt + splitInfo.curCompressedCnt; + for (uint32_t tcIdx = 0; tcIdx < splitInfo.dealTcSize; tcIdx += splitInfo.tcSplitSize) { + uint32_t actDealTcSize = min(splitInfo.tcSplitSize, splitInfo.dealTcSize - tcIdx); + + loopInfo.isCoreLoopFirst = tcIdx == 0; + loopInfo.isCoreLoopLast = tcIdx + splitInfo.tcSplitSize >= splitInfo.dealTcSize; + // 处理单个切块 + sliceIterator.SetNeedDealTcSize(actDealTcSize); + sliceIterator.SetDealedTcCnt(0U); + DealVec1BaseBlock(info, sliceIterator, loopInfo, dBaseOffset, splitInfo.dSplitSize, splitInfo.dBaseSize); + } + } + compressedCnt_ = preCompressedCnt + splitInfo.totalCompressedCnt; +} + +template +__aicore__ inline void CompressorBlockVectorPerf::ComputeVec2(const Compressor::Vec2RunInfo &info) +{ + SplitCoreV2(info); + uint32_t vec2DealM = v2TcEndIdx - v2TcStartIdx; + uint32_t loopCount = CeilDivT(vec2DealM, v2MBaseSize); + for (uint32_t v2LoopIdx = 0, dealSize = v2MBaseSize; v2LoopIdx < loopCount; ++v2LoopIdx) { + if (v2LoopIdx == loopCount - 1) { + dealSize = vec2DealM - v2LoopIdx * v2MBaseSize; + } + DealVec2BaseBlock(info, v2TcStartIdx + v2LoopIdx * v2MBaseSize, dealSize); + } + v2TcStartIdx = 0; + v2TcEndIdx = 0; +} + +template +__aicore__ inline void CompressorBlockVectorPerf::DealVec2BaseBlock(const Compressor::Vec2RunInfo &info, + uint32_t startRow, uint32_t dealRowCount) +{ + uint32_t computeSize = dealRowCount * constInfo_.headDim; + int64_t inGmOffset = startRow * constInfo_.headDim; + GlobalTensor vec2InputGm = vec2InputGm_[info.v2DbIdx * constInfo_.dbSize]; + // CopyIn + LocalTensor vec1ResUb = inputQue1.AllocTensor(); + DataCopy(vec1ResUb, vec2InputGm[inGmOffset], computeSize); + inputQue1.EnQue(vec1ResUb); + inputQue1.DeQue(); + + // RmsNorm + LocalTensor normResUb = tmpBuff1.Get(); + LocalTensor tempLocal = tmpBuff2.Get(); + PipeBarrier(); + MultRowRmsNorm(normResUb, vec1ResUb, normWeightUb, tempLocal, dealRowCount); + inputQue1.FreeTensor(vec1ResUb); + + + // rope: 只对后RD进行rope; 将normResUb每行前headDim - + // ropeHeadDim个元素cast到X_T，然后再与rope后的结果组合存到outputUb + LocalTensor outputUb = outputQue1.AllocTensor(); + PipeBarrier(); + CalRope(outputUb, normResUb, dealRowCount); + PipeBarrier(); + // CopyOut + outputQue1.EnQue(outputUb); + outputQue1.DeQue(); + CopyFinalResultOut(info, outputUb, startRow - v2TcStartIdx, dealRowCount); + outputQue1.FreeTensor(outputUb); +} + +template +__aicore__ inline void +CompressorBlockVectorPerf::MultRowRmsNorm(const LocalTensor &normResUb, const LocalTensor &vec1ResUb, + const LocalTensor &normWeightUb, const LocalTensor &tempLocal, + uint32_t dealRowCount) +{ + RmsNormParam rmsNormParams; + rmsNormParams.reciprocal = constInfo_.reciprocalD; + rmsNormParams.epsilon = constInfo_.normEps; + rmsNormParams.row = dealRowCount; + rmsNormParams.col = constInfo_.headDim; + RmsNorm(normResUb, vec1ResUb, normWeightUb, tempLocal, rmsNormParams); +} + + +template +__aicore__ inline void CompressorBlockVectorPerf::SingleCalRope(const LocalTensor &outputUb, + const LocalTensor &normResUb, uint32_t rowCnt, + uint32_t curDealScSize, uint32_t globalScStart) +{ + uint32_t computeSize = curDealScSize * constInfo_.ropeHeadDim; + uint64_t SinCosOffset = globalScStart * constInfo_.ropeHeadDim; + // sin/cos each reserves 16KB so fp32 rope can use the same compute tile. + LocalTensor cosUb = inputQue1.AllocTensor(); + LocalTensor sinUb = cosUb[BUFFER_SIZE_BYTE_16K / sizeof(ROPE_T)]; + DataCopy(cosUb, ropeCosGm_[SinCosOffset], computeSize); + DataCopy(sinUb, ropeSinGm_[SinCosOffset], computeSize); + inputQue1.EnQue(sinUb); + inputQue1.DeQue(); + + LocalTensor ropeCosFp32Local = tmpBuff2.Get(); + LocalTensor ropeSinFp32Local = ropeCosFp32Local[BUFFER_SIZE_BYTE_16K / sizeof(T)].template ReinterpretCast(); + LocalTensor tempLocal = ropeSinFp32Local[BUFFER_SIZE_BYTE_16K / sizeof(T)].template ReinterpretCast(); + PipeBarrier(); + if constexpr (IsSameType::value) { + DataCopy(ropeCosFp32Local, cosUb, computeSize); + DataCopy(ropeSinFp32Local, sinUb, computeSize); + } else { + Cast(ropeCosFp32Local, cosUb, RoundMode::CAST_NONE, computeSize); + Cast(ropeSinFp32Local, sinUb, RoundMode::CAST_NONE, computeSize); + } + PipeBarrier(); + inputQue1.FreeTensor(sinUb); + RotaryPosEmb(normResUb[rowCnt * constInfo_.headDim], normResUb[rowCnt * constInfo_.headDim], + ropeCosFp32Local, ropeSinFp32Local, tempLocal, gatherOffsetCastUb, curDealScSize, + constInfo_.ropeHeadDim, constInfo_.headDim, + constInfo_.headDim - constInfo_.ropeHeadDim); + PipeBarrier(); +} + +template +__aicore__ inline void CompressorBlockVectorPerf::CalRope(const LocalTensor &outputUb, + const LocalTensor &normResUb, uint32_t dealRowCount) +{ + uint32_t bStartIdx = OutputBStartIdx; + uint32_t sStartIdx = OutputSStartIdx; + uint64_t globalScStart = 0; + CalcGlobalScStart(0, 0, bStartIdx, sStartIdx, globalScStart); + uint32_t totalSize = dealRowCount * constInfo_.headDim; + uint32_t dealScSize = dealRowCount; + uint32_t curDealScSize = 0; + + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + // TH 模式逻辑：一次性执行核心计算，再空转更新 Index + curDealScSize = dealRowCount; + SingleCalRope(outputUb, normResUb, 0, curDealScSize, globalScStart); // rowOffset 传 0 + + while (dealScSize > 0) { + UpdateOutputIdx(bStartIdx, sStartIdx, dealScSize, curDealScSize); + } + } else { + // BSH 模式逻辑：分块（循环）执行核心计算 + uint32_t ubProcessedCount = 0; + uint32_t preOutputBStartIdx = 0; + uint32_t preOutputSStartIdx = 0; + + while (dealScSize > 0) { + preOutputBStartIdx = bStartIdx; + preOutputSStartIdx = sStartIdx; + UpdateOutputIdx(bStartIdx, sStartIdx, dealScSize, curDealScSize); + + if (curDealScSize > 0) { + uint32_t rowCnt = dealRowCount - dealScSize - curDealScSize; + SingleCalRope(outputUb, normResUb, rowCnt, curDealScSize, globalScStart); + } + CalcGlobalScStart(preOutputBStartIdx, preOutputSStartIdx, bStartIdx, sStartIdx, globalScStart); + ubProcessedCount += curDealScSize; + } + } + Cast(outputUb, normResUb, RoundMode::CAST_RINT, totalSize); + PipeBarrier(); +} + + +template +__aicore__ inline void CompressorBlockVectorPerf::SplitCoreV2(const Compressor::Vec2RunInfo &info) +{ + // 累积N个基本块数据后做vec2，N=2，传入的RunInfo包含该组核处理的数据块的bStart、bEnd、sStart、sEnd以及dealTcCount； + // 每组核切M方向将C1/V1后的数据分8 * 2个vec核上进行V2计算 + // 每次进行v2计算都会根据当前情况将workspace中的每组核处理的数据重新分到当前组的vec核 + + // Input: syncAll前每组cube核处理的实际数据块在batch及s方向的起止idx及实际数据量(m方向) + // Output: 每个vec核的处理数据块在m方向的起止位置及输出到Gm上的起始位置 + uint32_t coreNum = constInfo_.usedCoreNum * 2; // 总核数，vec*2 + uint32_t currCoreIdx = GetBlockIdx(); // 当前vec核ID + // 1.计算总vec2基本块数量 + uint32_t totalBaseNum = info.dealScSize; // 当前组核累积的实际数据量 + + uint32_t usedCoreNum = min(totalBaseNum, coreNum); + // 2.每个vec核上分到的数据量 + uint32_t avgBaseNum = CeilDivT(totalBaseNum, coreNum); + if (currCoreIdx % coreNum >= usedCoreNum) { + return; + } + // 3.计算每个vec核的起始结束位置 + uint32_t accumBaseNum = 0; // 当前累积的基本块数 + uint32_t targetBaseNum = (currCoreIdx % coreNum + 1) * avgBaseNum; // 当前vec核目标要达到的基本块数量 + uint32_t targetStartBaseNum = targetBaseNum - avgBaseNum; // 分当前vec核时前面已经完成分核的基本块数量 + bool setStart = false; + for (uint32_t i = 0; i < totalBaseNum; ++i) { + if (accumBaseNum >= totalBaseNum) { + return; + } + accumBaseNum += 1; + if (!setStart && (accumBaseNum > targetStartBaseNum)) { + v2TcStartIdx = i; + setStart = true; + } + if (setStart && (accumBaseNum >= targetBaseNum || i == (totalBaseNum - 1))) { + // 更新当前核的End分核信息 + v2TcEndIdx = i + 1; + GetScIdxInfo(info.bStart, info.bCompressedId, info.dealScSize, v2TcStartIdx, v2TcEndIdx, OutputBStartIdx, + OutputSStartIdx, OutputSize); + return; + } + } +} + +template +__aicore__ inline void CompressorBlockVectorPerf::CalcGlobalScStart(uint32_t bStart, uint32_t scStart, + uint32_t bEnd, uint32_t scEnd, + uint64_t &globalScStart) +{ + for (uint32_t bIdx = bStart; bIdx < bEnd; ++bIdx) { + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + curActSeqLength_ = GetSeqLength(bIdx); + curStartPos_ = GetStartPos(bIdx); + globalScStart += GetScSize(); + } else { + curActSeqLength_ = constInfo_.sSize; + globalScStart += CeilDivT(curActSeqLength_, constInfo_.cmpRatio); + } + } + globalScStart -= scStart; + globalScStart += scEnd; +} + +template +__aicore__ inline void CompressorBlockVectorPerf::UpdateOutputIdx(uint32_t &outputBStart, uint32_t &outputSStart, + uint32_t &dealScSize, uint32_t &curDealScSize) +{ + curActSeqLength_ = GetSeqLength(outputBStart); + curStartPos_ = GetStartPos(outputBStart); + uint32_t curBatchScSize = + (curStartPos_ + curActSeqLength_) / constInfo_.cmpRatio - curStartPos_ / constInfo_.cmpRatio; + uint32_t curBatchRemainScSize = curBatchScSize - outputSStart; + curDealScSize = curBatchRemainScSize > dealScSize ? dealScSize : curBatchRemainScSize; + dealScSize -= curDealScSize; + outputSStart += curDealScSize; + if (outputSStart == curBatchScSize) { + outputBStart++; + outputSStart = 0; + } +} + +template +__aicore__ inline void CompressorBlockVectorPerf::CopyFinalResultOut(const Compressor::Vec2RunInfo &info, + const LocalTensor &cmpKvOutUb, + uint32_t startRow, uint32_t dealRowCount) +{ + uint64_t globalScStart = 0; + CalcGlobalScStart(0, 0, OutputBStartIdx, OutputSStartIdx, globalScStart); + uint64_t outOffset = globalScStart * constInfo_.headDim; + uint32_t copySize = dealRowCount * constInfo_.headDim; + + uint32_t dealScSize = dealRowCount; + uint32_t curDealScSize = 0; + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + DataCopy(cmpKvOutGm_[outOffset], cmpKvOutUb, copySize); + while (dealScSize > 0) { + UpdateOutputIdx(OutputBStartIdx, OutputSStartIdx, dealScSize, curDealScSize); + } + } else { + // 处理BSH有效数据在内存上不连续（可能存在pad） + uint32_t ubProcessedCount = 0; + uint32_t preOutputBStartIdx = 0; + uint32_t preOutputSStartIdx = 0; + while (dealScSize > 0) { + // 逐batch计算写出索引 + preOutputBStartIdx = OutputBStartIdx; + preOutputSStartIdx = OutputSStartIdx; + UpdateOutputIdx(OutputBStartIdx, OutputSStartIdx, dealScSize, curDealScSize); + DataCopy(cmpKvOutGm_[globalScStart * constInfo_.headDim], cmpKvOutUb[ubProcessedCount * constInfo_.headDim], + curDealScSize * constInfo_.headDim); + CalcGlobalScStart(preOutputBStartIdx, preOutputSStartIdx, OutputBStartIdx, OutputSStartIdx, globalScStart); + ubProcessedCount += curDealScSize; + } + } +} +} // namespace Compressor +#endif // COMPRESSOR_BLOCK_VECTOR_PREF_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_comm.h b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_comm.h new file mode 100644 index 000000000..3e988ce04 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_comm.h @@ -0,0 +1,341 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_comm.h + * \brief + */ + +#ifndef COMPRESSOR_COMM_H +#define COMPRESSOR_COMM_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" + +using namespace AscendC; + +namespace Compressor { +template +__aicore__ inline T CeilDivT(T num1, T num2) +{ + if (num2 == 0) { + return static_cast(0); + } + return (num1 + num2 - 1) / num2; +} + +template +__aicore__ inline T Align(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd) * (rnd))); +} + +template +__aicore__ inline T Trunc(T num, T rnd) +{ + return ((rnd) == 0) ? 0 : (((num) / (rnd) * (rnd))); +} + +template +__aicore__ inline T FloorPow2(T num) +{ + if (num == 0) return 1; + for(uint32_t i = 1; i < sizeof(T) * 8; i <<= 1) { + num |= (num >> i); + } + return num - (num >> 1); +} + +template +__aicore__ inline T CeilPow2(T num) +{ + if (num <= 1) return 1; + num --; + for(uint32_t i = 1; i < sizeof(T) * 8; i <<= 1) { + num |= (num >> i); + } + num ++; + return num; +} + +enum class X_LAYOUT : std::uint8_t { + BSH = static_cast(0), + TH = static_cast(1) +}; + +enum class X_DTYPE : std::uint8_t { + BF16 = static_cast(0), + FP16 = static_cast(1) +}; + +enum class ROPE_DTYPE : std::uint8_t { + SAME_AS_X = static_cast(0), + FP32 = static_cast(1) +}; + +enum class COFF : std::uint8_t { + DISABLE = static_cast(1), + OVERLAP = static_cast(2) +}; + +enum class ROTARY_MODE : std::uint8_t { + HALF = static_cast(1), + INTERLEAVE = static_cast(2) +}; + +enum class CACHE_MODE : std::uint8_t { + CONTINUOUS = static_cast(1), + CYCLE = static_cast(2) +}; + +enum class TEMPLATE_ID : uint8_t { + NORMAL = 0, + EMPTY_X = 1, + PERF = 2 +}; + +template +struct COMPType { + static constexpr X_LAYOUT xLayout = X_L; + static constexpr X_DTYPE xDtype = X_T; + static constexpr ROPE_DTYPE ropeDtype = R_T; + static constexpr COFF coff = C; + static constexpr ROTARY_MODE rotaryMode = Rotary_Mode; +}; + +struct ConstInfo { + // 整个AICORE的任务信息, 左闭右开区间[ (bStart, s2Start), (bEnd, s2End) ) + uint32_t bStart = 0U; + uint32_t sStart = 0U; + uint32_t bEnd = 0U; + uint32_t sEnd = 0U; + + // 分核相关 + uint32_t usedCoreNum = 0; + uint32_t dBaseSize = 0; + uint32_t mBaseSize = 0; + uint32_t tcSize = 0; + uint32_t tcBaseSize = 0; + uint32_t tcBasicBlockNum = 0; + uint32_t dBasicBlockNum = 0; + uint32_t coreGroupNum = 0; + uint32_t singleCoreDealTcBasicNum = 0; + uint32_t dIdx = 0; + uint32_t bIdxOfLastTc = 0; + uint32_t sIdxOfLastTc = 0; + + // shape及参数 + uint32_t batchSize = 0; + uint32_t hSize = 0; + uint32_t sSize = 0; + uint32_t headDim = 0; + uint32_t ropeHeadDim = 0; + uint32_t cmpRatio = 0; + float normEps = 1e-6; + float reciprocalD = 0; + + uint32_t curGroupIdx = 0; + uint32_t tailGroupIdx = 0; + uint32_t tailBasicBlockNum = 0; + uint32_t realDealBasicBlockNum = 0; + + // pageAttention + uint32_t blockNum = 0; + uint32_t blockSize = 0; + uint32_t maxBlockNumPerBatch = 0; + uint64_t stateCacheStrideDim0 = 0; + + // workSpace + uint32_t dbWorkspaceRatio = 1; + uint32_t mm1KvResSize = 0; + uint32_t mm1ScoreResSize = 0; + uint32_t vec1TailCacheSize = 0; + uint32_t vec1ResSize = 0; + uint32_t mm1ResSize = 0; // 所有cube输出kv/score结果的总大小 + + uint32_t aiCoreIdx = 0; + uint32_t nSize = 0; + + uint32_t dbSize = 0; +}; + +struct RunInfo { + bool isValid = false; + uint32_t cubeDbIdx = 0; // kernel主循环索引 + + // 增加字段 + uint32_t dealTcNum = 0; + // 右边相关信息 + uint32_t bStart = 0; + uint32_t sStart = 0; + uint32_t dealSeqCnt = 0; + // 左边相关信息 + uint32_t preBStart = 0; + uint32_t preSStart = 0; + uint32_t preDealSeqCnt = 0; // 左边需要处理的s大小 + uint32_t preFirstSeqCnt = 0; // 左边首块大小 + + + uint32_t bEnd = 0; + uint32_t sEnd = 0; + uint32_t bStartSeqIdx = 0; + uint32_t bEndSeqIdx = 0; + + // v2分核信息 sc是左闭右开 + uint32_t scStart = 0; + uint32_t scEnd = 0; + uint32_t dealScSize = 0; + + // vec1Res offset + uint64_t vec1ResOffset = 0; +}; + +struct Vec1RunInfo { + // vec相关信息，一次syncAll需处理数据的起始索引 + bool resetResFlag = false; // v1积攒N轮 是否是N轮的起始轮 + uint32_t c1v1DbIdx = 0; // vec1 doubleBuffer索引 + uint32_t v1v2DbIdx = 0; // v1v2 doubleBuffer索引 + uint32_t bStart = 0; + uint32_t sStart = 0; + uint32_t dealTcNum = 0; + uint32_t dealScSize = 0; +}; + +struct Vec2RunInfo { + // uint32_t bStart = 0; + uint32_t v2DbIdx = 0; // v2 doubleBuffer索引 + uint32_t sStart = 0; + uint32_t bEnd = 0; + uint32_t sEnd = 0; + // v2分核信息 sc是左闭右开 + uint32_t scStart = 0; + uint32_t scEnd = 0; + // uint32_t dealScSize = 0; + + // 增加字段 + uint32_t bStart = 0; + uint32_t compressedId = 0; + uint32_t bCompressedId = 0; + uint32_t dealScSize = 0; +}; + +struct MSplitInfo { + uint32_t vecStartB = 0U; + uint32_t vecStartS = 0U; + uint32_t vecEndB = 0U; + uint32_t vecEndS = 0U; + uint32_t dealTcNum = 0U; + // vec1Res offset + uint64_t vec1StartOffset = 0; + uint64_t vec1ResOffset = 0; +}; + +struct BlockInfo { + __aicore__ inline BlockInfo(uint32_t bIdx, uint32_t sIdx, uint32_t dealSeqSize) : + bIdx(bIdx), sIdx(sIdx), dealSeqSize(dealSeqSize) {}; + uint32_t bIdx = 0U; + uint32_t sIdx = 0U; + uint32_t dealSeqSize = 0; + + uint32_t isFirst = true; + uint32_t bSeqUsed = 0U; + uint32_t bStartPos = 0U; + uint32_t headHolderSeqCnt = 0U; + uint32_t validSeqCnt = 0U; + uint32_t tailHolderSeqCnt = 0U; + uint32_t dealTcSize = 0U; + uint32_t tailValidSeqCnt = 0U; + uint32_t compressTcSize = 0U; +}; + +// BUFFER的字节数 +inline constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; +inline constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64; +inline constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256; +inline constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512; +inline constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024; +inline constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048; +inline constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096; +inline constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192; +inline constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384; +inline constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768; +inline constexpr uint32_t BUFFER_SIZE_BYTE_64K = 65536; + +// BLOCK和REPEAT的字节数 +inline constexpr uint64_t BYTE_BLOCK = 32UL; +inline constexpr uint32_t REPEAT_BLOCK_BYTE = 256U; +// BLOCK和REPEAT的FP32元素数 +inline constexpr uint32_t FP32_BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(float); // 8 +inline constexpr uint32_t FP32_REPEAT_ELEMENT_NUM = REPEAT_BLOCK_BYTE / sizeof(float); // 64 +inline constexpr uint32_t REPEAT_STRIDE_NUM = REPEAT_BLOCK_BYTE / BYTE_BLOCK; // 8 +inline constexpr uint32_t REPEAT_MAX_NUM = 255; +inline constexpr uint32_t BRCB_NUM = 8; +inline constexpr uint32_t MAX_R = 256; + +template +__aicore__ inline void CopySingleMatrixNDToNZ(LocalTensor l1Tensor, const GlobalTensor gmTensor, + uint32_t nValue, uint32_t dValue, uint32_t srcDValue, uint32_t dstNzC0Stride) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nValue; // nd矩阵的行数 + if constexpr (IsSameType::value) { + constexpr uint32_t HALF_SIZE_DIVISOR = 2; + nd2nzPara.dValue = dValue / HALF_SIZE_DIVISOR; + nd2nzPara.srcDValue = srcDValue / HALF_SIZE_DIVISOR; + } else { + nd2nzPara.dValue = dValue; // nd矩阵的列数 + nd2nzPara.srcDValue = srcDValue; // 同一nd矩阵相邻行起始地址间的偏移 + } + nd2nzPara.dstNzC0Stride = dstNzC0Stride; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(l1Tensor, gmTensor, nd2nzPara); +} +template +__aicore__ inline void DumpTensorForDim2(GlobalTensor tensor, uint32_t desc, uint32_t dumpSize, uint32_t row, uint32_t col) +{ + uint32_t array2[] = {static_cast(row), static_cast(col)}; + AscendC::ShapeInfo shapeInfo(2, array2); + // AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo); +} + +template +__aicore__ inline void DumpTensorForDim2(LocalTensor tensor, uint32_t desc, uint32_t dumpSize, uint32_t row, uint32_t col) +{ + uint32_t array2[] = {static_cast(row), static_cast(col)}; + AscendC::ShapeInfo shapeInfo(2, array2); + // AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo); +} + +template +__aicore__ inline void DumpTensorForDim2(LocalTensor tensor, uint32_t desc, uint32_t dumpSize) +{ + uint32_t col = 32 / sizeof(T); + uint32_t array2[] = {static_cast(dumpSize / col), static_cast(col)}; + AscendC::ShapeInfo shapeInfo(2, array2); + // AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo); +} + +template +__aicore__ inline void DumpTensorForDim2(GlobalTensor tensor, uint32_t desc, uint32_t dumpSize) +{ + uint32_t col = 32 / sizeof(T); + uint32_t array2[] = {static_cast(dumpSize / col), static_cast(col)}; + AscendC::ShapeInfo shapeInfo(2, array2); + // AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo); +} + +} // namespace Compressor +#endif diff --git a/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_kernel.h b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_kernel.h new file mode 100644 index 000000000..a066477e1 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_kernel.h @@ -0,0 +1,86 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_kernel.h + * \brief + */ + +#ifndef COMPRESSOR_KERNEL +#define COMPRESSOR_KERNEL + +#include "compressor_comm.h" +#include "compressor_template_tiling_key.h" +#include "compressor_kernel_perf.h" +#include "compressor_tiling_data.h" +#include "compressor_tools.h" + +using namespace AscendC; + +namespace Compressor { + +template +class CompressorKernel { +public: + __aicore__ inline CompressorKernel(TPipe* pipe, const optiling::CompressorTilingData* __restrict tilingData) + : pipe_(pipe), tilingData_(tilingData) {} + + __aicore__ inline void Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut, + __gm__ uint8_t *workspace); + __aicore__ inline void Process(); + + // ==============================TilingData&TPipe============================== + TPipe* pipe_; + const optiling::CompressorTilingData* __restrict tilingData_; +}; + +template +__aicore__ inline void CompressorKernel::Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut, + __gm__ uint8_t *workspace) +{ + +} + + +template +__aicore__ inline void CompressorKernel::Process() +{ + +} + +} // namespace Compressor + +#endif // COMPRESSOR_KERNEL diff --git a/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_kernel_perf.h b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_kernel_perf.h new file mode 100644 index 000000000..3beb71e10 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_kernel_perf.h @@ -0,0 +1,695 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_kernel_perf.h + * \brief + */ + +#ifndef COMPRESSOR_KERNEL_PERF_H +#define COMPRESSOR_KERNEL_PERF_H + +#include "compressor_comm.h" +#include "compressor_template_tiling_key.h" +#include "compressor_tiling_data.h" +#include "compressor_tools.h" +#include "compressor_block_cube_perf.h" +#include "compressor_block_vec_perf.h" + + +using namespace AscendC; + +namespace Compressor { + +struct CmpBlockInfo { + __aicore__ inline CmpBlockInfo() {}; + __aicore__ inline CmpBlockInfo(uint32_t bIdx, uint32_t sIdx, bool needReset = false) : bIdx(bIdx), sIdx(sIdx), needReset(needReset) {}; + + uint32_t bIdx = 0U; + uint32_t sIdx = 0U; + uint32_t bSeqUsed = 0U; + uint32_t bStartPos = 0U; + bool needReset = false; + bool isFirst = true; + + uint32_t headSeqCnt = 0U; + uint32_t validSeqCnt = 0U; + uint32_t tailSeqCnt = 0U; + bool isCompress = 0U; +}; + +struct BasicBlockInfo { + uint32_t bIdx = 0; + uint32_t sIdx = 0; + uint32_t compressedTcNum = 0; + uint32_t dealSeqCnt = 0; + uint32_t dealTcNum = 0; +}; + +struct BatchInfo { + uint32_t tcNum = 0; + uint32_t compressedTcNum = 0; + uint32_t remSeqCnt = 0; + uint32_t seqCnt = 0; + uint32_t seqUsedCnt = 0; + uint32_t headHolderSeq = 0; + uint32_t bStartPos = 0; + uint32_t bIdx = 0; + uint32_t sIdx = 0; +}; + +template +class CompressorKernelPerf { +public: + __aicore__ inline CompressorKernelPerf(TPipe* pipe, const optiling::CompressorTilingData* __restrict tilingData) + : pipe_(pipe), tilingData_(tilingData) {} + + __aicore__ inline void Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut, + __gm__ uint8_t *workspace); + __aicore__ inline void Process(); + +private: + // ================================Init functions================================== + __aicore__ inline void InitWorkspace(__gm__ uint8_t *workspace); + // ================================Process functions================================ + __aicore__ inline void InitTilingData(); + __aicore__ inline void SetBaseSize(); + // 获取基本块数量 + __aicore__ inline uint32_t GetLoopTimes(); + __aicore__ inline void SkipInvalidBatch(BatchInfo &batchInfo); + __aicore__ inline void UpdateCurGroup(BasicBlockInfo &basicBlockInfo, BatchInfo batchInfo, uint32_t &curGroupQuota, uint32_t curDealSeq); + __aicore__ inline BasicBlockInfo SkipOneLoop(BatchInfo &batchInfo); + // 计算分核基本信息 + __aicore__ inline void CalcSplitCoreInfo(); + + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void ComputeMm1(const RunInfo &info, bool isNeedExcute); + __aicore__ inline void ComputeVec1(const Vec1RunInfo &info); + __aicore__ inline void ComputeVec2(const Vec2RunInfo &info); + + __aicore__ inline bool IsNeedExcuteC1(RunInfo info); + __aicore__ inline bool IsNeedSyncAll(uint32_t curBasicBlockIdx); + __aicore__ inline void CalcC1V1Params(RunInfo &info, Vec1RunInfo &vec1Info, BatchInfo &batchInfo, uint32_t loopIdx); + __aicore__ inline void UpdateVec2Info(Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info); + __aicore__ inline bool IsNeedExcuteV2(Vec2RunInfo &vec2Info); + + using X_T = typename AscendC::Conditional::type; + using T = float; + using MM1_OUT_T = T; + using VEC1_OUT_T = T; + + // 常量 + static constexpr uint64_t SYNC_MODE0 = 0; + static constexpr uint64_t SYNC_MODE2 = 2; + static constexpr uint32_t SYNC_C1_FLAG = 3; + static constexpr uint32_t SYNC_V1_FLAG = 4; + static constexpr uint32_t SYNC_V1_FLAG2 = 5; + static constexpr uint32_t SYNC_C1_V1_FLAG = 6; + static constexpr uint32_t SYNC_V1_C1_FLAG = 8; + + // ==============================TilingData&TPipe============================== + TPipe* pipe_; + const optiling::CompressorTilingData* __restrict tilingData_; + // ===========================Workspace Global Tensor=========================== + GlobalTensor mm1KvResGm; + GlobalTensor mm1ScoreResGm; + GlobalTensor vec1KvCacheGm; + GlobalTensor vec1ScoreCacheGm; + GlobalTensor Vec1InputKvGm; + GlobalTensor Vec1InputScoreGm; + GlobalTensor vec1ResGm; + GlobalTensor vec2InputGm; + // ================================Task Info==================================== + CompressorTools tools_; + ConstInfo constInfo{}; + uint32_t aiCoreIdx = 0; + + // ==============================Service Define============================== + CompressorBlockCubePerf blockCube_; + CompressorBlockVectorPerf blockVec_; + + uint32_t allCompressedTcNum_ = 0; + uint32_t curCompressedTcNum_ = 0; + uint32_t accDealSize = 0; + uint32_t loopTimes = 0; + uint32_t cubeLoop = 0; + uint32_t vec1Loop = 0; + uint32_t vec2Loop = 0; + bool isFirstUpdateCurGroup = true; +}; + +template +__aicore__ inline void CompressorKernelPerf::Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut, + __gm__ uint8_t *workspace) +{ + if ASCEND_IS_AIV { + constInfo.aiCoreIdx = GetBlockIdx() / 2; + } else { + constInfo.aiCoreIdx = GetBlockIdx(); + } + InitTilingData(); + // init tools + tools_.toolParams_.seqSize = tilingData_->baseParams.seqSize; + tools_.toolParams_.cmpRatio = tilingData_->baseParams.cmpRatio; + tools_.Init(startPos, seqUsed, cuSeqlens); + + // 剔除尾部的无效batch + for (; constInfo.batchSize > 0; --constInfo.batchSize) { + uint32_t bSeqUsed = tools_.GetSeqLength(constInfo.batchSize - 1); + if (bSeqUsed > 0) { + break; + } + } + + // 所有batch的有效序列都为0时, 直接退出 + if (constInfo.batchSize == 0) { + return; + } + + // 0. 计算最后一个Tc块的起始位置 + constInfo.bIdxOfLastTc = constInfo.batchSize - 1; + // 1. 计算head_dim的切分大小, 构建ConstInfo的其他信息 + SetBaseSize(); // 设置基本块大小 + CalcSplitCoreInfo(); + // 2. 计算循环次数 + loopTimes = GetLoopTimes(); + // 3. 初始化workspace + InitWorkspace(workspace); + // 4. 初始化block层 + if ASCEND_IS_AIC { +#if __CCE_AICORE__ == 310 + blockCube_.InitParams(constInfo, tools_); +#else + blockCube_.InitParams(constInfo, tools_); +#endif + blockCube_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos, + stateBlockTable, cuSeqlens, seqUsed, startPos, cmpKvOut); + blockCube_.InitBuffers(pipe_); +#if __CCE_AICORE__ == 310 + blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm); +#else + blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm); +#endif + } else { + blockVec_.InitParams(constInfo, tools_); + blockVec_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos, stateBlockTable, + cuSeqlens, seqUsed, startPos, cmpKvOut); + blockVec_.InitBuffers(pipe_); +#if __CCE_AICORE__ == 310 + blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm, vec1ResGm, vec2InputGm); +#else + blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm, vec1ResGm, vec2InputGm); +#endif + } +} + +template +__aicore__ inline void CompressorKernelPerf::InitTilingData() { + constInfo.cmpRatio = tilingData_->baseParams.cmpRatio; + constInfo.batchSize = tilingData_->baseParams.batchSize; + constInfo.mBaseSize = tilingData_->innerSplitParams.mBaseSize; + constInfo.headDim = tilingData_->baseParams.headDim; + constInfo.hSize = tilingData_->baseParams.hiddenSize; + constInfo.sSize = tilingData_->baseParams.seqSize; + constInfo.ropeHeadDim = tilingData_->baseParams.ropeHeadDim; + constInfo.normEps = tilingData_->baseParams.normEps; + constInfo.reciprocalD = tilingData_->baseParams.reciprocalD; + constInfo.usedCoreNum = tilingData_->baseParams.usedCoreNum; + + constInfo.blockNum = tilingData_->pageAttentionParams.blockNum; + constInfo.blockSize = tilingData_->pageAttentionParams.blockSize; + constInfo.maxBlockNumPerBatch = tilingData_->pageAttentionParams.maxBlockNumPerBatch; + constInfo.stateCacheStrideDim0 = tilingData_->baseParams.stateCacheStrideDim0; + + constInfo.nSize = tilingData_->baseParams.nSize; + constInfo.vec1TailCacheSize = tilingData_->workspaceParams.vec1TailCacheSize; + constInfo.dbWorkspaceRatio = tilingData_->workspaceParams.dbWorkspaceRatio; +} + +template +__aicore__ inline void CompressorKernelPerf::SetBaseSize() +{ + uint32_t mSize = 0; + uint32_t minMBaseSize = 0; + bool sameSeqUsed = true; + uint32_t firstBatchSeqUsed = tools_.GetSeqLength(0); + for (uint32_t i = 0; i < constInfo.batchSize; i++) { + uint32_t bSeqUsed = tools_.GetSeqLength(i); + uint32_t bStartPos = tools_.GetStartPos(i); + // 获取m大小 + mSize += bSeqUsed; + // 获取是否等长 + if (sameSeqUsed && (bSeqUsed != firstBatchSeqUsed)) { + sameSeqUsed = false; + } + // 获取m轴最小切分大小 + if (minMBaseSize != constInfo.cmpRatio) { + uint32_t startCmpIdx = bStartPos / constInfo.cmpRatio; + uint32_t endCmpIdx = (bStartPos + bSeqUsed) / constInfo.cmpRatio; + if (startCmpIdx == endCmpIdx) { + if (bSeqUsed > minMBaseSize) { + minMBaseSize = bSeqUsed; + } + } else if (startCmpIdx + 1 == endCmpIdx) { + uint32_t startCmpValidSeqCnt = constInfo.cmpRatio - (bStartPos % constInfo.cmpRatio); + uint32_t endCmpValidSeqCnt = (bStartPos + bSeqUsed) % constInfo.cmpRatio; + if (startCmpValidSeqCnt > minMBaseSize) { + minMBaseSize = startCmpValidSeqCnt; + } + if (endCmpValidSeqCnt > minMBaseSize) { + minMBaseSize = endCmpValidSeqCnt; + } + } else { + minMBaseSize = constInfo.cmpRatio; + } + } + } + + uint32_t aiCoreNum = constInfo.usedCoreNum; + constInfo.dBaseSize = 64; + uint32_t dBaseBlockNum = constInfo.headDim / constInfo.dBaseSize; + if (sameSeqUsed && mSize <= (constInfo.mBaseSize * (aiCoreNum / dBaseBlockNum))) { + if constexpr (COMP::coff == COFF::OVERLAP) { + if (constInfo.headDim == 128) { + dBaseBlockNum = 8; + } else if (constInfo.headDim == 512) { + dBaseBlockNum = 16; + } + } else { + if (constInfo.headDim == 128) { + dBaseBlockNum = 8; + } else if (constInfo.headDim == 512) { + dBaseBlockNum = 16; + } + } + // 核数足够时, 修改才生效 + if (aiCoreNum >= dBaseBlockNum) { + constInfo.dBaseSize = constInfo.headDim / dBaseBlockNum; + // 开启全核 + uint32_t coreGroupNum = aiCoreNum / dBaseBlockNum; + uint32_t newMBaseSize = (constInfo.batchSize + coreGroupNum - 1) / coreGroupNum * firstBatchSeqUsed; + if (newMBaseSize > minMBaseSize && newMBaseSize < constInfo.mBaseSize) { + constInfo.mBaseSize = newMBaseSize; + } + } + } +} + + +template +__aicore__ inline void CompressorKernelPerf::SkipInvalidBatch(BatchInfo &batchInfo) +{ + for (; batchInfo.bIdx < constInfo.batchSize; ++batchInfo.bIdx) { + batchInfo.seqCnt = tools_.GetSeqLength(batchInfo.bIdx); + if (batchInfo.seqCnt > 0) { + break; + } + } + batchInfo.remSeqCnt = batchInfo.seqCnt; + if (tools_.isExistSeqUsed_) { + batchInfo.seqUsedCnt = tools_.GetSeqUsed(batchInfo.bIdx); + } else { + batchInfo.seqUsedCnt = batchInfo.seqCnt; + } + if (batchInfo.bIdx < constInfo.batchSize) { + batchInfo.bStartPos = tools_.GetStartPos(batchInfo.bIdx); + batchInfo.sIdx = 0; + batchInfo.headHolderSeq = batchInfo.bStartPos & (constInfo.cmpRatio - 1); + batchInfo.tcNum = (batchInfo.bStartPos + batchInfo.seqCnt + constInfo.cmpRatio - 1) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio; + batchInfo.compressedTcNum = (batchInfo.bStartPos + batchInfo.seqUsedCnt) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio; + } +} + + +template +__aicore__ inline void CompressorKernelPerf::UpdateCurGroup(BasicBlockInfo &basicBlockInfo, + BatchInfo batchInfo, uint32_t &curGroupQuota, uint32_t curDealSeq) +{ + // 更新当前组的信息 + if (curGroupQuota == 0 && !isFirstUpdateCurGroup) { + return; + } + isFirstUpdateCurGroup = false; + basicBlockInfo.bIdx = batchInfo.bIdx; + uint32_t curGroupDealSeq = curGroupQuota < curDealSeq ? curGroupQuota : curDealSeq; + basicBlockInfo.sIdx = batchInfo.sIdx + curGroupDealSeq; + basicBlockInfo.dealSeqCnt += curGroupDealSeq; + curGroupQuota -= curGroupDealSeq; + // 结尾需要跳batch，需要考虑在当前组起始为末尾，或者当前组起始大于整个M轴 + if ((curGroupQuota == 0 || basicBlockInfo.bIdx == constInfo.batchSize - 1) && basicBlockInfo.sIdx == batchInfo.seqCnt) { + basicBlockInfo.sIdx = 0; + for (basicBlockInfo.bIdx++; basicBlockInfo.bIdx < constInfo.batchSize; ++basicBlockInfo.bIdx) { + uint32_t seqCnt = tools_.GetSeqLength(basicBlockInfo.bIdx); + if (seqCnt > 0) { + break; + } + } + } +} + +template +__aicore__ inline BasicBlockInfo CompressorKernelPerf::SkipOneLoop(BatchInfo &batchInfo) +{ + BasicBlockInfo basicBlockInfo{}; + isFirstUpdateCurGroup = true; + uint32_t curGroupQuota = constInfo.mBaseSize * constInfo.curGroupIdx; // m轴当前组起始 + bool curGroupStartFlag = false; + uint32_t quota = constInfo.coreGroupNum * constInfo.mBaseSize; + + for (; batchInfo.bIdx < constInfo.batchSize;) { + uint32_t curDealSeq = 0; + uint32_t curDealTcNum = 0; + uint32_t curDealCompressedTcNum = 0; + // 无法处理完当前整个batch + if (quota < batchInfo.remSeqCnt) { + // 向下对齐r， + if (quota > constInfo.cmpRatio - batchInfo.headHolderSeq) { + uint32_t delta = (batchInfo.bStartPos + batchInfo.sIdx + quota) & (constInfo.cmpRatio - 1); // 超出对齐的部分 + curDealSeq = quota - delta; + quota -= curDealSeq; + curDealTcNum = (curDealSeq + constInfo.cmpRatio - 1) / constInfo.cmpRatio; + curDealCompressedTcNum = min(curDealTcNum, batchInfo.compressedTcNum); + // 更新当前组所需信息 + UpdateCurGroup(basicBlockInfo, batchInfo, curGroupQuota, curDealSeq); + // 更新batch信息 + batchInfo.remSeqCnt = batchInfo.remSeqCnt - curDealSeq; + batchInfo.sIdx = batchInfo.sIdx + curDealSeq; + batchInfo.compressedTcNum -= curDealCompressedTcNum; + batchInfo.tcNum -= curDealTcNum; + // 更新loop信息 + basicBlockInfo.dealTcNum += curDealTcNum; + basicBlockInfo.compressedTcNum += curDealCompressedTcNum; + } + break; + } else { + // 处理整个batch + quota -= batchInfo.remSeqCnt; + curDealSeq = batchInfo.remSeqCnt; + curDealTcNum = batchInfo.tcNum; + // 更新当前组所需信息 + UpdateCurGroup(basicBlockInfo, batchInfo, curGroupQuota, curDealSeq); + // 更新batch和loop信息 + batchInfo.remSeqCnt = 0; + basicBlockInfo.dealTcNum += batchInfo.tcNum; + basicBlockInfo.compressedTcNum += batchInfo.compressedTcNum; + batchInfo.bIdx++; + SkipInvalidBatch(batchInfo); + } + } + uint32_t totalDataSize = constInfo.coreGroupNum * constInfo.mBaseSize - quota; + // 2. 当前组的起始偏移 + uint32_t currentGroupStart = constInfo.curGroupIdx * constInfo.mBaseSize; + + // 3. 安全判断 + if (currentGroupStart >= totalDataSize) { + // 超出尾块 + basicBlockInfo.dealSeqCnt = 0; + } else { + // 还在有效范围内，计算剩余量 + uint32_t remaining = totalDataSize - currentGroupStart; + basicBlockInfo.dealSeqCnt = (remaining < constInfo.mBaseSize) ? remaining : constInfo.mBaseSize; + } + + return basicBlockInfo; +} + + +template +__aicore__ inline uint32_t CompressorKernelPerf::GetLoopTimes() +{ + // 计算主循环次数 + uint32_t loopTimes = 0; + BatchInfo batchInfo{}; + SkipInvalidBatch(batchInfo); + for (;batchInfo.bIdx < constInfo.batchSize; ++loopTimes) { + SkipOneLoop(batchInfo); + } + return loopTimes; +} + +template +__aicore__ inline void CompressorKernelPerf::CalcSplitCoreInfo() +{ + // D方向的基本块数量 + constInfo.dBasicBlockNum = constInfo.headDim / constInfo.dBaseSize; + // 核的组数 + constInfo.coreGroupNum = constInfo.usedCoreNum / constInfo.dBasicBlockNum; + // 每个核处理的d方向的索引 + constInfo.dIdx = (constInfo.aiCoreIdx % constInfo.dBasicBlockNum) * constInfo.dBaseSize; + // 当前组id + constInfo.curGroupIdx = constInfo.aiCoreIdx / constInfo.dBasicBlockNum; + + constInfo.mm1ResSize = constInfo.mBaseSize * constInfo.headDim * constInfo.coreGroupNum; + + uint32_t coff = (uint32_t)COMP::coff; + constInfo.mm1KvResSize = constInfo.mBaseSize * constInfo.headDim * coff; + constInfo.mm1ScoreResSize = constInfo.mBaseSize * constInfo.headDim * coff; + constInfo.vec1ResSize = constInfo.mBaseSize * constInfo.headDim * constInfo.nSize; + + constInfo.dbSize = constInfo.coreGroupNum * constInfo.mm1KvResSize; +} + +template +__aicore__ inline void CompressorKernelPerf::InitWorkspace(__gm__ uint8_t *workspace) { + uint64_t offset = 0; + uint64_t mm1KvResStartOffset = offset; + // mm1KvResGm + mm1KvResGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + offset + + constInfo.curGroupIdx * constInfo.mm1KvResSize * sizeof(MM1_OUT_T))); + offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1KvResSize * sizeof(MM1_OUT_T); + + uint64_t mm1ScoreResStartOffset = offset; + // mm1ScoreResGm + mm1ScoreResGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + offset + + constInfo.curGroupIdx * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T))); + offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T); + + Vec1InputKvGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + mm1KvResStartOffset)); + + Vec1InputScoreGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + mm1ScoreResStartOffset)); + + vec1KvCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset)); + offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T); + + vec1ScoreCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset)); + offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T); + + uint64_t beforeVecOffset = offset; + + // vec1Res + vec1ResGm.SetGlobalBuffer( + (__gm__ VEC1_OUT_T *)(workspace + offset)); + offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.vec1ResSize * sizeof(VEC1_OUT_T); + // vec2Input + vec2InputGm.SetGlobalBuffer( + (__gm__ VEC1_OUT_T *)(workspace + beforeVecOffset)); +} + +template +__aicore__ inline void CompressorKernelPerf::ComputeMm1(const RunInfo &info, bool isNeedExcute) { + CrossCoreWaitFlag(SYNC_V1_C1_FLAG + info.cubeDbIdx); + if (isNeedExcute) { + blockCube_.ComputeMm1(info); + } + CrossCoreSetFlag(SYNC_C1_FLAG); + CrossCoreWaitFlag(SYNC_C1_FLAG); + CrossCoreSetFlag(SYNC_C1_V1_FLAG + info.cubeDbIdx); +} + +template +__aicore__ inline void CompressorKernelPerf::ComputeVec1(const Vec1RunInfo &info) { + CrossCoreWaitFlag(SYNC_C1_V1_FLAG + info.c1v1DbIdx); + CrossCoreWaitFlag(SYNC_V1_FLAG2 + info.c1v1DbIdx); + blockVec_.ComputeVec1(info); + CrossCoreSetFlag(SYNC_V1_FLAG); + CrossCoreWaitFlag(SYNC_V1_FLAG); + CrossCoreSetFlag(SYNC_V1_C1_FLAG + info.c1v1DbIdx); + CrossCoreSetFlag(SYNC_V1_FLAG2 + (info.c1v1DbIdx + 1) % constInfo.dbWorkspaceRatio); +} + +template +__aicore__ inline void CompressorKernelPerf::ComputeVec2(const Vec2RunInfo &info) { + blockVec_.ComputeVec2(info); +} + +template +__aicore__ inline void CompressorKernelPerf::AllocEventID() +{ + if ASCEND_IS_AIC { + blockCube_.AllocEventID(pipe_); + } else { + blockVec_.AllocEventID(); + for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) { + CrossCoreSetFlag(SYNC_V1_C1_FLAG + i); + } + CrossCoreSetFlag(SYNC_V1_FLAG2); + } +} + +template +__aicore__ inline void CompressorKernelPerf::FreeEventID() +{ + if ASCEND_IS_AIC { + for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) { + CrossCoreWaitFlag(SYNC_V1_C1_FLAG + i); + } + blockCube_.FreeEventID(pipe_); + } else { + CrossCoreWaitFlag(SYNC_V1_FLAG2 + loopTimes % constInfo.dbWorkspaceRatio); + blockVec_.FreeEventID(); + } +} + +template +__aicore__ inline bool CompressorKernelPerf::IsNeedExcuteC1(RunInfo info) +{ + // B超出范围则cube不执行 + return info.bStart < constInfo.batchSize; +} + +template +__aicore__ inline void CompressorKernelPerf::CalcC1V1Params(RunInfo &info, Vec1RunInfo &vec1Info, BatchInfo &batchInfo, uint32_t loopIdx) +{ + vec1Info.bStart = batchInfo.bIdx; + vec1Info.sStart = batchInfo.sIdx; + vec1Info.resetResFlag = (loopIdx & (constInfo.nSize - 1)) == 0; + vec1Info.c1v1DbIdx = (vec1Loop++ & (constInfo.dbWorkspaceRatio - 1)); + vec1Info.v1v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1)); + BasicBlockInfo basicBlockInfo = SkipOneLoop(batchInfo); + info.cubeDbIdx = (cubeLoop++ & (constInfo.dbWorkspaceRatio - 1)); + info.dealSeqCnt = basicBlockInfo.dealSeqCnt; + info.dealTcNum = basicBlockInfo.dealTcNum; + info.bStart = basicBlockInfo.bIdx; + info.sStart = basicBlockInfo.sIdx; + vec1Info.dealTcNum = basicBlockInfo.dealTcNum; + vec1Info.dealScSize = basicBlockInfo.compressedTcNum; + allCompressedTcNum_ += basicBlockInfo.compressedTcNum; +} + +template +__aicore__ inline bool CompressorKernelPerf::IsNeedExcuteV2(Vec2RunInfo &vec2Info) +{ + return (vec2Info.dealScSize > 0); +} + +template +__aicore__ inline bool CompressorKernelPerf::IsNeedSyncAll(uint32_t curBasicBlockIdx) +{ + if (allCompressedTcNum_ == 0) { + return false; + } + + uint32_t cnt = curBasicBlockIdx + 1; + if ((cnt == loopTimes) || (cnt % constInfo.nSize == 0)) { + return true; + } + return false; +} + +template +__aicore__ inline void CompressorKernelPerf::UpdateVec2Info( + Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info) +{ + // nSize轮起始先重置v2Info信息 + if (curBasicBlockIdx % constInfo.nSize == 0) { + vec2Info.v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1)); + vec2Info.bStart = info.bStart; + vec2Info.sStart = info.sStart; + // 将sStart转成bCompressedId + uint32_t startPos = tools_.GetStartPos(info.bStart); + if (tools_.isExistSeqUsed_) { + uint32_t seqUsed = tools_.GetSeqUsed(info.bStart); + if (vec2Info.sStart >= seqUsed) { + vec2Info.bStart++; + vec2Info.sStart = 0; + } + } + vec2Info.bCompressedId = (startPos + vec2Info.sStart) / constInfo.cmpRatio - startPos / constInfo.cmpRatio; + + vec2Info.dealScSize = 0; + } else if ((curBasicBlockIdx + 1) % constInfo.nSize == 0) { + vec2Loop++; + } + vec2Info.dealScSize += info.dealScSize; + vec2Info.compressedId += info.dealScSize; +} + +template +__aicore__ inline void CompressorKernelPerf::Process() +{ + // 所有batch的有效序列都为0时, 直接退出 + if (constInfo.batchSize == 0) { + return; + } + AllocEventID(); + + BatchInfo batchInfo{}; + + RunInfo extraInfo[1]; + Vec1RunInfo vec1Info{}; + Vec2RunInfo vec2Info{}; + SkipInvalidBatch(batchInfo); + for (uint32_t i = 0; i < loopTimes; ++i) { + RunInfo &extraInfo0 = extraInfo[0]; + CalcC1V1Params(extraInfo0, vec1Info, batchInfo, i); + bool isNeedExcuteC1 = IsNeedExcuteC1(extraInfo0); + + if ASCEND_IS_AIC { + ComputeMm1(extraInfo0, isNeedExcuteC1); + } else { + ComputeVec1(vec1Info); + UpdateVec2Info(vec2Info, i, vec1Info); + + if (IsNeedSyncAll(i)) { + SyncAll(); + if (IsNeedExcuteV2(vec2Info)) { + ComputeVec2(vec2Info); + } + } + } + + } + FreeEventID(); +} + +} // namespace Compressor + +#endif // COMPRESSOR_KERNEL_PERF_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_template_tiling_key.h b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_template_tiling_key.h new file mode 100644 index 000000000..32db03e27 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_template_tiling_key.h @@ -0,0 +1,56 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file COMPRESSOR_template_tiling_key.h + * \brief + */ + +#ifndef COMPRESSOR_TEMPLATE_TILING_KEY_H +#define COMPRESSOR_TEMPLATE_TILING_KEY_H + +#include "ascendc/host_api/tiling/template_argument.h" + +#define ASCENDC_TPL_1_BW 1 // 每个参数占用1个bit位 +#define ASCENDC_TPL_2_BW 2 // 每个参数占用2个bit位 +#define ASCENDC_TPL_4_BW 4 // 每个参数占用4个bit位 + +// 可表示的tilingkey范围为64bit，注意不可超过限制 +ASCENDC_TPL_ARGS_DECL(compressor, // 算子唯一标识，与opType保持一致 + // 可能需要切分之后的headdim + // bit:0 LAYOUT 0:BSH 1:TH + ASCENDC_TPL_UINT_DECL(X_LAYOUT, ASCENDC_TPL_1_BW, ASCENDC_TPL_UI_LIST, 0, 1), + // bit:1-4 x的dtype 0:BF16 1:FP16 + ASCENDC_TPL_UINT_DECL(X_DTYPE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, 0, 1), + // bit:5-6 coff 1:无需overlap 2:需要overlap + ASCENDC_TPL_UINT_DECL(COFF, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2), + // bit:7-8 rotary_mode 1:half 2:interleave + ASCENDC_TPL_UINT_DECL(ROTARY_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2), + // bit:9-10 cache_mode 1:CONTINUOUS 2:cycle + ASCENDC_TPL_UINT_DECL(CACHE_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2), + // bit:11-12 template_id 0:empty_tensor 1:normal 2:full load + ASCENDC_TPL_UINT_DECL(TEMPLATE_ID, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 0, 1, 2), + // bit:13 rope dtype 0:same as x 1:fp32 + ASCENDC_TPL_UINT_DECL(ROPE_DTYPE, ASCENDC_TPL_1_BW, ASCENDC_TPL_UI_LIST, 0, 1), +); + +ASCENDC_TPL_SEL( + + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_UINT_SEL(X_LAYOUT, ASCENDC_TPL_UI_LIST, 0, 1), + ASCENDC_TPL_UINT_SEL(X_DTYPE, ASCENDC_TPL_UI_LIST, 0, 1), + ASCENDC_TPL_UINT_SEL(COFF, ASCENDC_TPL_UI_LIST, 1, 2), + ASCENDC_TPL_UINT_SEL(ROTARY_MODE, ASCENDC_TPL_UI_LIST, 1, 2), + ASCENDC_TPL_UINT_SEL(CACHE_MODE, ASCENDC_TPL_UI_LIST, 1, 2), + ASCENDC_TPL_UINT_SEL(TEMPLATE_ID, ASCENDC_TPL_UI_LIST, 0, 1, 2), + ASCENDC_TPL_UINT_SEL(ROPE_DTYPE, ASCENDC_TPL_UI_LIST, 0, 1), + ASCENDC_TPL_TILING_STRUCT_SEL(optiling::CompressorTilingData)), +); + +#endif // COMPRESSOR_TEMPLATE_TILING_KEY_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_tiling_data.h b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_tiling_data.h new file mode 100644 index 000000000..efebb9d69 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_tiling_data.h @@ -0,0 +1,69 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file COMPRESSOR_tiling_datay.h + * \brief + */ + +#ifndef COMPRESSOR_TILING_DATA_H +#define COMPRESSOR_TILING_DATA_H +#include +#include "kernel_tiling/kernel_tiling.h" + +const uint32_t CMP_MAX_AIC_CORE_NUM = 26; // 25 + 1 保证数组8字节对齐 + +namespace optiling { + // 1. 基础参数结构体 + struct CompressorBaseParams { + uint32_t batchSize = 0; // bastch size（批大小） + uint32_t seqSize = 0; // sequence size（kvs大小） + uint32_t hiddenSize = 0; // hidden size（隐藏层大小） + uint32_t tokenSize = 0; // token size = batchSize * seqSize(token总数：批大小x序列1长度) + uint32_t headDim = 0; // head size of kv + uint32_t ropeHeadDim = 64; // dim size per rope head 64（单个带RoPE头的维度） + uint32_t csSize = 0; // Compress sequence len + uint32_t cmpRatio = 4; // Compress ratio + uint32_t cgSize = 0; // Compress group size + float normEps = 1e-6; // RMSNorm eps + float reciprocalD = 0; // 1分之D + uint32_t usedCoreNum = 0; // 使用核数 + uint32_t nSize = 0; // 控制v2积攒的轮数 + uint64_t stateCacheStrideDim0 = 0; // stateCache第0维的stride + }; + + struct CompressorPageAttentionParams { + uint32_t blockNum = 0; + uint32_t blockSize = 1; + uint32_t maxBlockNumPerBatch = 1; + }; + + struct CompressorInnerSplitParams { + uint32_t mBaseSize; + uint32_t dBaseSize; + }; + + struct CompressorWorkspaceParams { + uint32_t mm1KvResSize; + uint32_t mm1ScoreResSize; + uint32_t vec1ResSize; + uint32_t vec1TailCacheSize; + uint32_t dbWorkspaceRatio = 1; + }; + + struct CompressorTilingData { + CompressorBaseParams baseParams; + CompressorPageAttentionParams pageAttentionParams; + CompressorInnerSplitParams innerSplitParams; + CompressorWorkspaceParams workspaceParams; + }; +} // optiling + +#endif // COMPRESSOR_TILING_DATA_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_tools.h b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_tools.h new file mode 100644 index 000000000..71f7eb8c9 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_tools.h @@ -0,0 +1,761 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_tools.h + * \brief 放算子都需要、与算子联系紧密、但是又不方便单独独立出来的公共工具 + */ + +#ifndef COMPRESSOR_TOOLS_H +#define COMPRESSOR_TOOLS_H + +#include "compressor_comm.h" + +using namespace AscendC; + +namespace Compressor { + +struct ToolsParams { + uint32_t seqSize = 0U; + uint32_t cmpRatio = 0U; +}; + +template +class CompressorTools { +public: + __aicore__ inline CompressorTools() + { + } + + __aicore__ inline void Init(__gm__ uint8_t *cuSeqlens, __gm__ uint8_t *seqUsed, __gm__ uint8_t *startPos); + + __aicore__ inline uint32_t GetSeqUsed(uint32_t bIdx); + __aicore__ inline uint32_t GetStartPos(uint32_t bIdx); + __aicore__ inline uint32_t GetSeqLength(uint32_t bIdx); + __aicore__ inline uint32_t GetTIdxByBatch(uint32_t bIdx); + +public: + ToolsParams toolParams_{}; + bool isExistSeqUsed_ = false; + +private: + bool isExistStartPos_ = false; + GlobalTensor cuSeqlensGm_; + GlobalTensor sequsedGm_; + GlobalTensor startPosGm_; +}; + +template +__aicore__ inline void CompressorTools::Init(__gm__ uint8_t *startPos, __gm__ uint8_t *seqUsed, + __gm__ uint8_t *cuSeqlens) +{ + isExistStartPos_ = (startPos != nullptr); + if (isExistStartPos_) { + startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos); + } + + isExistSeqUsed_ = (seqUsed != nullptr); + if (isExistSeqUsed_) { + sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed); + } + + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens); + } +} + +template +__aicore__ inline uint32_t CompressorTools::GetSeqUsed(uint32_t bIdx) +{ + if (isExistSeqUsed_) { + return (uint32_t)sequsedGm_.GetValue(bIdx); + } else { + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + return (uint32_t)(cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx)); + } else { + return toolParams_.seqSize; + } + } +} + +template +__aicore__ inline uint32_t CompressorTools::GetStartPos(uint32_t bIdx) +{ + if (isExistStartPos_) { + return (uint32_t)startPosGm_.GetValue(bIdx); + } else { + return 0; + } +} + +template +__aicore__ inline uint32_t CompressorTools::GetSeqLength(uint32_t bIdx) +{ + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + return cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx); + } else { + return toolParams_.seqSize; + } +} + +template +__aicore__ inline uint32_t CompressorTools::GetTIdxByBatch(uint32_t bIdx) +{ + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + return (uint32_t)(cuSeqlensGm_.GetValue(bIdx)); + } else { + return toolParams_.seqSize * bIdx; + } +} + +// iterator +struct SliceInfo { + __aicore__ inline SliceInfo(){}; + __aicore__ inline SliceInfo(uint32_t bIdx, uint32_t sIdx) : bIdx(bIdx), sIdx(sIdx){}; + + uint32_t bIdx = 0U; + uint32_t sIdx = 0U; + uint32_t bSeqUsed = 0U; + uint32_t bStartPos = 0U; + + uint32_t headHolderSeqCnt = 0U; + uint32_t validSeqCnt = 0U; + uint32_t tailHolderSeqCnt = 0U; + + uint32_t dealSeqCnt = 0; + uint32_t dealTcSize = 0U; + uint32_t compressTcSize = 0U; +}; + +template +class CompressorSliceIterator { +public: + __aicore__ inline CompressorSliceIterator(CompressorTools &tools) : tools_(tools) + { + } + + __aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx); + __aicore__ inline void SetMaxBatchSize(uint32_t batch_size); + __aicore__ inline void SetMaxDealSeqCnt(uint32_t maxDealSeqCnt); + __aicore__ inline bool IsEnd(); + __aicore__ inline void IteratorSlice(); + __aicore__ inline SliceInfo &GetSlice(); + __aicore__ inline SliceInfo &GetSliceByCmp(); + + bool isFirst_ = true; + SliceInfo sliceInfo_{}; + +private: + CompressorTools &tools_; + + // iterator + uint32_t maxDealSeqCnt_ = 0; + uint32_t batch_size_ = 0; +}; + +template +__aicore__ inline void CompressorSliceIterator::Reset(uint32_t bIdx, uint32_t sIdx) +{ + sliceInfo_.bIdx = bIdx; + sliceInfo_.sIdx = sIdx; + isFirst_ = true; +} + +template +__aicore__ inline void CompressorSliceIterator::SetMaxBatchSize(uint32_t batch_size) +{ + this->batch_size_ = batch_size; +} + +template +__aicore__ inline void CompressorSliceIterator::SetMaxDealSeqCnt(uint32_t maxDealSeqCnt) +{ + this->maxDealSeqCnt_ = maxDealSeqCnt; +} + +template +__aicore__ inline bool CompressorSliceIterator::IsEnd() +{ + return (sliceInfo_.bIdx >= batch_size_) || (maxDealSeqCnt_ == 0); +} + +template +__aicore__ inline void CompressorSliceIterator::IteratorSlice() +{ + bool isUpdateBatchInfo = false; + if (!isFirst_) { + // 更新剩余未处理的行数 + maxDealSeqCnt_ -= sliceInfo_.dealSeqCnt; + // 更新sIdx和bIdx、以及与bIdx相关的bStartPos和bSeqUsed + sliceInfo_.sIdx += sliceInfo_.validSeqCnt; + if (sliceInfo_.sIdx == sliceInfo_.bSeqUsed) { + sliceInfo_.sIdx = 0; + sliceInfo_.bIdx++; + isUpdateBatchInfo = true; + } + } else { + isUpdateBatchInfo = true; + isFirst_ = false; + } + + // 更新与bIdx相关的bStartPos和bSeqUsed + if (isUpdateBatchInfo) { + // SkipInvalidBatch + while (sliceInfo_.bIdx < batch_size_) { + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + if (sliceInfo_.bSeqUsed > 0) { + break; + } + sliceInfo_.bIdx++; + } + if (sliceInfo_.bIdx < batch_size_) { + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + } + } +} + +template +__aicore__ inline SliceInfo &CompressorSliceIterator::GetSliceByCmp() +{ + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + if (isFirst_) { + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + isFirst_ = false; + } + // 计算头部占位行数、有效数据行数、尾部占位行数 + sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio; + + sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx; + if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) { + sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt; + } + sliceInfo_.tailHolderSeqCnt = + cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio; + if (sliceInfo_.tailHolderSeqCnt == cmpRatio) { + sliceInfo_.tailHolderSeqCnt = 0; + } + + // 头和尾处理，否则需要处理的seq等于cmpRatio + if (sliceInfo_.validSeqCnt < cmpRatio) { + sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt; + if (sliceInfo_.sIdx == 0) { + sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt; + } + } else { + sliceInfo_.dealSeqCnt = cmpRatio; + } + sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt; + + // 计算本次可以处理的Tc个数 + sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio; + + // 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0 + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize; + if (sliceInfo_.tailHolderSeqCnt > 0) { + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时，其不需要压缩 + } + + return sliceInfo_; +} + +template +__aicore__ inline SliceInfo &CompressorSliceIterator::GetSlice() +{ + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + if (isFirst_) { + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + isFirst_ = false; + } + // 计算头部占位行数、有效数据行数、尾部占位行数 + sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio; + sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx; + if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) { + sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt; + } + sliceInfo_.tailHolderSeqCnt = + cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio; + if (sliceInfo_.tailHolderSeqCnt == cmpRatio) { + sliceInfo_.tailHolderSeqCnt = 0; + } + + sliceInfo_.dealSeqCnt = sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt + sliceInfo_.tailHolderSeqCnt; + // 计算本次可以处理的Tc个数 + sliceInfo_.dealTcSize = sliceInfo_.dealSeqCnt / cmpRatio; + + // 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0 + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize; + if (sliceInfo_.tailHolderSeqCnt > 0) { + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时，其不需要压缩 + } + + return sliceInfo_; +} + +struct SplitCoreSliceInfo : public SliceInfo { + __aicore__ inline SplitCoreSliceInfo(){}; + __aicore__ inline SplitCoreSliceInfo(uint32_t bIdx, uint32_t sIdx) : SliceInfo(bIdx, sIdx){}; + + uint32_t preFirstSeqCnt = 0U; // 左边每次迭代基本块的第一个seqCnt大小 +}; + +template +class CompressorSplitCoreSliceIterator { +public: + __aicore__ inline CompressorSplitCoreSliceIterator(CompressorTools &tools) : tools_(tools) + { + } + + __aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx); + __aicore__ inline void SetMaxBatchSize(uint32_t batch_size); + __aicore__ inline void SetMaxDealSeqCnt(uint32_t maxDealSeqCnt); + __aicore__ inline bool IsEnd(); + __aicore__ inline void IteratorSlice(); + __aicore__ inline SplitCoreSliceInfo &GetSlice(); + __aicore__ inline SplitCoreSliceInfo &GetSliceByCmp(); + __aicore__ inline uint32_t GetBIdx(); + __aicore__ inline SplitCoreSliceInfo &GetLeftNextCmpSeqCnt(); + __aicore__ inline SplitCoreSliceInfo &GetRightNextCmpSeqCnt(); + + bool isFirst_ = true; + bool isLeftFirstBath = false; + bool isMaxDealSeqCntFirst = false; + + SplitCoreSliceInfo sliceInfo_{}; + +private: + CompressorTools &tools_; + + // iterator + uint32_t maxDealSeqCnt_ = 0; + uint32_t batch_size_ = 0; +}; + +template +__aicore__ inline void CompressorSplitCoreSliceIterator::Reset(uint32_t bIdx, uint32_t sIdx) +{ + sliceInfo_.bIdx = bIdx; + sliceInfo_.sIdx = sIdx; + isFirst_ = true; +} + +template +__aicore__ inline void CompressorSplitCoreSliceIterator::SetMaxBatchSize(uint32_t batch_size) +{ + this->batch_size_ = batch_size; + isMaxDealSeqCntFirst = true; +} + +template +__aicore__ inline void CompressorSplitCoreSliceIterator::SetMaxDealSeqCnt(uint32_t maxDealSeqCnt) +{ + this->maxDealSeqCnt_ = maxDealSeqCnt; +} + +template +__aicore__ inline bool CompressorSplitCoreSliceIterator::IsEnd() +{ + return (sliceInfo_.bIdx >= batch_size_) || (maxDealSeqCnt_ == 0); +} + +template +__aicore__ inline uint32_t CompressorSplitCoreSliceIterator::GetBIdx() +{ + return sliceInfo_.bIdx; +} + +template +__aicore__ inline void CompressorSplitCoreSliceIterator::IteratorSlice() +{ + bool isUpdateBatchInfo = false; + if (isMaxDealSeqCntFirst) { + isMaxDealSeqCntFirst = false; + } + if (!isFirst_) { + // 更新剩余未处理的行数 + maxDealSeqCnt_ -= sliceInfo_.dealSeqCnt; + // 更新sIdx和bIdx、以及与bIdx相关的bStartPos和bSeqUsed + sliceInfo_.sIdx += sliceInfo_.validSeqCnt; + if (sliceInfo_.sIdx == sliceInfo_.bSeqUsed) { + sliceInfo_.sIdx = 0; + // 左边最后一块跳到b=0 s=0处理 + if (isLeftFirstBath) { + isLeftFirstBath = false; + } else { + sliceInfo_.bIdx++; + } + isUpdateBatchInfo = true; + } + } else { + isUpdateBatchInfo = true; + isFirst_ = false; + } + + // 更新与bIdx相关的bStartPos和bSeqUsed + if (isUpdateBatchInfo) { + // SkipInvalidBatch + while (sliceInfo_.bIdx < batch_size_) { + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + if (sliceInfo_.bSeqUsed > 0) { + break; + } + sliceInfo_.bIdx++; + } + if (sliceInfo_.bIdx < batch_size_) { + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + } + } +} + +template +__aicore__ inline SplitCoreSliceInfo &CompressorSplitCoreSliceIterator::GetLeftNextCmpSeqCnt() +{ + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + if (isFirst_) { + // 左边 T轴首次减去T轴最后一块 + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(batch_size_ - 1); + sliceInfo_.bStartPos = tools_.GetStartPos(batch_size_ - 1); + // 处理最后一块是中间整块或者尾块的情况 + uint32_t lastSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) % cmpRatio == 0 ? + cmpRatio : + (sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) % cmpRatio; + // 处理最后一块是头块的情况 + if (sliceInfo_.bSeqUsed < cmpRatio) { + lastSeqCnt = sliceInfo_.bSeqUsed; + } + + sliceInfo_.sIdx = sliceInfo_.bSeqUsed - lastSeqCnt; + isLeftFirstBath = true; + isFirst_ = false; + } + // 计算头部占位行数、有效数据行数、尾部占位行数 + sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio; + + sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx; + if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) { + sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt; + } + sliceInfo_.tailHolderSeqCnt = + cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio; + if (sliceInfo_.tailHolderSeqCnt == cmpRatio) { + sliceInfo_.tailHolderSeqCnt = 0; + } + + // 头和尾处理，否则需要处理的seq等于cmpRatio + if (sliceInfo_.validSeqCnt < cmpRatio) { + sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt; + if (sliceInfo_.sIdx == 0) { + sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt; + } + } else { + sliceInfo_.dealSeqCnt = cmpRatio; + } + sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt; + + // 计算本次可以处理的Tc个数 + sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio; + + // 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0 + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize; + if (sliceInfo_.tailHolderSeqCnt > 0) { + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时，其不需要压缩 + } + + // 记录左边第一个块 + if (isMaxDealSeqCntFirst) { + sliceInfo_.preFirstSeqCnt = sliceInfo_.dealSeqCnt; + } + + return sliceInfo_; +} + +template +__aicore__ inline SplitCoreSliceInfo &CompressorSplitCoreSliceIterator::GetRightNextCmpSeqCnt() +{ + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + if (isFirst_) { + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + isFirst_ = false; + } + // 计算头部占位行数、有效数据行数、尾部占位行数 + sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio; + + sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx; + if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) { + sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt; + } + sliceInfo_.tailHolderSeqCnt = + cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio; + if (sliceInfo_.tailHolderSeqCnt == cmpRatio) { + sliceInfo_.tailHolderSeqCnt = 0; + } + + // 头和尾处理，否则需要处理的seq等于cmpRatio + if (sliceInfo_.validSeqCnt < cmpRatio) { + sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt; + if (sliceInfo_.sIdx == 0) { + sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt; + } + } else { + sliceInfo_.dealSeqCnt = cmpRatio; + } + sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt; + + // 计算本次可以处理的Tc个数 + sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio; + + // 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0 + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize; + if (sliceInfo_.tailHolderSeqCnt > 0) { + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时，其不需要压缩 + } + + return sliceInfo_; +} + +struct Vec1SliceInfo : public SliceInfo { + __aicore__ inline Vec1SliceInfo(){}; + __aicore__ inline Vec1SliceInfo(uint32_t bIdx, uint32_t sIdx) : SliceInfo(bIdx, sIdx){}; + __aicore__ inline Vec1SliceInfo(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt) + : SliceInfo(bIdx, sIdx), dealedSeqCnt(dealedSeqCnt){}; + + uint32_t dealedSeqCnt = 0U; + uint32_t dealedTcCnt = 0U; + uint32_t bSeqLength = 0U; + uint32_t compressoredScCnt = 0U; + bool isFirst = false; + bool isLast = false; +}; + +struct StatisticInfo { + __aicore__ inline StatisticInfo(){}; + __aicore__ inline StatisticInfo(uint32_t actualTcCnt, uint32_t dealSeqCnt, uint32_t compressorScCnt) + : actualTcCnt(actualTcCnt), dealSeqCnt(dealSeqCnt), compressorScCnt(compressorScCnt){}; + + uint32_t actualTcCnt = 0U; + uint32_t dealSeqCnt = 0U; + uint32_t compressorScCnt = 0U; +}; + +template +class CompressorVec1SliceIterator { +public: + __aicore__ inline CompressorVec1SliceIterator(CompressorTools &tools) : tools_(tools) + { + } + + __aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx); + __aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt, uint32_t compressoredScCnt); + __aicore__ inline void SetMaxBatchSize(uint32_t batch_size); + __aicore__ inline void SetDealedSeqCnt(uint32_t dealedSeqCnt); + __aicore__ inline void SetDealedTcCnt(uint32_t dealedTcCnt); + __aicore__ inline void SetCompressoredScCnt(uint32_t compressoredScCnt); + __aicore__ inline void SetNeedDealTcSize(uint32_t needDealTcSize); + __aicore__ inline void SetNeedDealTcSize(uint32_t needDealTcSize, uint32_t canDealTcSize); + __aicore__ inline uint32_t GetNeedDealTcSize(); + __aicore__ inline bool IsEnd(); + template + __aicore__ inline void IteratorSlice(); + __aicore__ inline Vec1SliceInfo &GetSlice(); + template + __aicore__ inline StatisticInfo &FullIteratorSlice(); + +private: + CompressorTools &tools_; + + bool isFirst_ = true; + Vec1SliceInfo sliceInfo_{}; + StatisticInfo statisticInfo_{}; + uint32_t needDealTcSize_ = 0U; + uint32_t batch_size_ = 0U; +}; + +template +__aicore__ inline void CompressorVec1SliceIterator::Reset(uint32_t bIdx, uint32_t sIdx) +{ + sliceInfo_.bIdx = bIdx; + sliceInfo_.sIdx = sIdx; + while (tools_.GetSeqLength(sliceInfo_.bIdx) == 0) { + sliceInfo_.bIdx++; + if (sliceInfo_.bIdx == batch_size_) { + sliceInfo_.bIdx = 0; + } + } + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + sliceInfo_.bSeqLength = tools_.GetSeqLength(sliceInfo_.bIdx); + isFirst_ = true; +} + +template +__aicore__ inline void CompressorVec1SliceIterator::Reset(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt, + uint32_t compressoredScCnt) +{ + Reset(bIdx, sIdx); + SetDealedSeqCnt(dealedSeqCnt); + SetCompressoredScCnt(compressoredScCnt); +} + +template +__aicore__ inline void CompressorVec1SliceIterator::SetMaxBatchSize(uint32_t batch_size) +{ + this->batch_size_ = batch_size; +} + +template +__aicore__ inline void CompressorVec1SliceIterator::SetDealedSeqCnt(uint32_t dealedSeqCnt) +{ + this->sliceInfo_.dealedSeqCnt = dealedSeqCnt; +} + +template +__aicore__ inline void CompressorVec1SliceIterator::SetCompressoredScCnt(uint32_t compressoredScCnt) +{ + this->sliceInfo_.compressoredScCnt = compressoredScCnt; +} + +template +__aicore__ inline void CompressorVec1SliceIterator::SetDealedTcCnt(uint32_t dealedTcCnt) +{ + this->sliceInfo_.dealedTcCnt = dealedTcCnt; +} + +template +__aicore__ inline void CompressorVec1SliceIterator::SetNeedDealTcSize(uint32_t needDealTcSize) +{ + this->needDealTcSize_ = needDealTcSize; +} + +template +template +__aicore__ inline void CompressorVec1SliceIterator::IteratorSlice() +{ + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + if constexpr (IS_STATISTIC) { + statisticInfo_.actualTcCnt += sliceInfo_.dealTcSize; + statisticInfo_.compressorScCnt += sliceInfo_.compressTcSize; + } + needDealTcSize_ -= sliceInfo_.dealTcSize; + sliceInfo_.dealedSeqCnt += sliceInfo_.validSeqCnt; + sliceInfo_.compressoredScCnt += sliceInfo_.compressTcSize; + sliceInfo_.sIdx += sliceInfo_.validSeqCnt; + if (sliceInfo_.sIdx >= sliceInfo_.bSeqUsed) { + do { + uint32_t seqLength = tools_.GetSeqLength(sliceInfo_.bIdx); + if (sliceInfo_.bSeqUsed < seqLength) { + uint32_t nextAlignSIdx = Align(sliceInfo_.bStartPos + sliceInfo_.sIdx, cmpRatio) - sliceInfo_.bStartPos; + sliceInfo_.dealedSeqCnt += nextAlignSIdx - sliceInfo_.sIdx; + uint32_t tcGap = CeilDivT(static_cast(seqLength - nextAlignSIdx), + static_cast(cmpRatio)); + if (sliceInfo_.bSeqUsed == 0 && nextAlignSIdx > sliceInfo_.sIdx) { + // 此时bseqused所在压缩块未被纳入计算 + tcGap++; + } + sliceInfo_.sIdx = nextAlignSIdx; + if (needDealTcSize_ < tcGap) { + sliceInfo_.dealedSeqCnt += needDealTcSize_ * cmpRatio; + sliceInfo_.sIdx += needDealTcSize_ * cmpRatio; + needDealTcSize_ = 0; + break; + } + sliceInfo_.dealedSeqCnt += seqLength - sliceInfo_.sIdx; + needDealTcSize_ -= tcGap; + } + sliceInfo_.bIdx++; + if (sliceInfo_.bIdx == batch_size_) { + sliceInfo_.bIdx = 0; + } + sliceInfo_.sIdx = 0; + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + } while (sliceInfo_.bSeqUsed == 0); + sliceInfo_.bSeqLength = tools_.GetSeqLength(sliceInfo_.bIdx); + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + } + if (isFirst_) { + isFirst_ = false; + } +} + +template +__aicore__ inline uint32_t CompressorVec1SliceIterator::GetNeedDealTcSize() +{ + return needDealTcSize_; +} + + +template +__aicore__ inline bool CompressorVec1SliceIterator::IsEnd() +{ + return (needDealTcSize_ == 0); +} + +template +__aicore__ inline Vec1SliceInfo &CompressorVec1SliceIterator::GetSlice() +{ + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + if (sliceInfo_.bSeqUsed < sliceInfo_.sIdx) { + sliceInfo_.headHolderSeqCnt = 0; + sliceInfo_.validSeqCnt = 0; + sliceInfo_.tailHolderSeqCnt = 0; + sliceInfo_.dealTcSize = 0; + sliceInfo_.compressTcSize = 0; + } else { + // 计算头部占位行数、有效数据行数、尾部占位行数 + sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio; + sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx; + if (CeilDivT(sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt, cmpRatio) > needDealTcSize_) { + sliceInfo_.validSeqCnt = needDealTcSize_ * cmpRatio - sliceInfo_.headHolderSeqCnt; + } + uint32_t globalTotalSeqCnt = sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt; + sliceInfo_.tailHolderSeqCnt = Align(globalTotalSeqCnt, cmpRatio) - globalTotalSeqCnt; + + // 计算本次可以处理的Tc个数 + sliceInfo_.dealTcSize = + (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt + sliceInfo_.tailHolderSeqCnt) / cmpRatio; + + sliceInfo_.compressTcSize = + (sliceInfo_.headHolderSeqCnt + min(sliceInfo_.validSeqCnt, sliceInfo_.bSeqUsed - sliceInfo_.sIdx)) / + cmpRatio; + } + + sliceInfo_.isFirst = isFirst_; + sliceInfo_.isLast = + sliceInfo_.bSeqUsed > sliceInfo_.sIdx && + CeilDivT(sliceInfo_.headHolderSeqCnt + sliceInfo_.bSeqUsed - sliceInfo_.sIdx, cmpRatio) >= needDealTcSize_; + + return sliceInfo_; +} + +template +template +__aicore__ inline StatisticInfo &CompressorVec1SliceIterator::FullIteratorSlice() +{ + if constexpr (IS_STATISTIC) { + statisticInfo_ = {0U, 0U, 0U}; + Vec1SliceInfo tempSliceInfo = GetSlice(); + while (!IsEnd()) { + GetSlice(); + IteratorSlice(); + } + Vec1SliceInfo sliceInfo = GetSlice(); + statisticInfo_.dealSeqCnt = sliceInfo.dealedSeqCnt - tempSliceInfo.dealedSeqCnt; + } else { + while (!IsEnd()) { + GetSlice(); + IteratorSlice(); + } + } + return statisticInfo_; +} +} // namespace Compressor + +#endif diff --git a/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_vector_comm.h b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_vector_comm.h new file mode 100644 index 000000000..2495c269e --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch32/compressor_vector_comm.h @@ -0,0 +1,347 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_vector_comm.h + * \brief 存放各种vector的公共组件 + */ + +#ifndef COMPRESSOR_VECTOR_COMM_H +#define COMPRESSOR_VECTOR_COMM_H + +#include "compressor_comm.h" +namespace Compressor { + + +struct MatRpeatParam { + uint32_t row; + uint32_t col; + uint32_t dtypeMask; + uint32_t loopTimes; + uint32_t colRemain; + uint8_t repeatStride; +}; + +struct RmsNormParam { + float reciprocal; + float epsilon; + uint32_t row; + uint32_t col; +}; + +/** + * @brief ColumnSum 对矩阵按列进行求和 + * @param dstLocal 输出tensor [1, col]，支持和shareTmpUb是同一块空间 + * @param srcLocal 输入tensor [row, col] + * @param shareTmpUb 临时buffer 内部需要的空间为 [ceil(row / 2) * col * sizeof(float)] + * @param row 行数 + * @param col 列数 + */ +__aicore__ inline void ColumnSum(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + const LocalTensor &shareTmpUb, uint32_t row, uint32_t col) +{ + // 行数为1时，直接将srcLocal复制到dstLocal + if (unlikely(row == 1)) { + DataCopy(dstLocal, srcLocal, row * col); + PipeBarrier(); + return; + } + for (uint32_t mask = MAX_R << 1; mask > 1; mask >>= 1) { + if (row & mask) { + // 将输入对半求和后放进临时空间 + Add(shareTmpUb, srcLocal, srcLocal[mask * col / 2], mask * col / 2); // 2:对矩阵按列做计算 + PipeBarrier(); + // 将余量加到前一半上 + if (unlikely(row > mask)) { + if ((row - mask) > (mask >> 1)) { + Add(shareTmpUb, shareTmpUb, srcLocal[mask * col], mask * col / 2); // 2:对矩阵按列做计算 + PipeBarrier(); + Add(shareTmpUb, shareTmpUb, srcLocal[(mask + (mask >> 1)) * col], (row - mask - (mask >> 1)) * col); + PipeBarrier(); + } else { + Add(shareTmpUb, shareTmpUb, srcLocal[mask * col], (row - mask) * col); + PipeBarrier(); + } + } + // 每次将后一半行加到前一半上 + for (uint32_t i = mask >> 2; i > 1; i >>= 1) { + Add(shareTmpUb, shareTmpUb, shareTmpUb[i * col], i * col); + PipeBarrier(); + } + if (mask == 2) { // 2:最后一次矩阵运算处理 + DataCopy(dstLocal, shareTmpUb, col); + } else { + Add(dstLocal, shareTmpUb, shareTmpUb[col], col); + } + PipeBarrier(); + break; + } + } +} + +/** + * @brief ColumnMax 对矩阵按列进行求最大值 + * @param dstLocal 输出tensor [1, col]，支持和shareTmpUb是同一块空间 + * @param srcLocal 输入tensor [row, col] + * @param shareTmpUb 临时buffer 内部需要的空间为 [ceil(row / 2) * col * sizeof(float)] + * @param row 行数 + * @param col 列数 + */ +__aicore__ inline void ColumnMax(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + const LocalTensor &shareTmpUb, uint32_t row, uint32_t col) +{ + // 行数为1时，直接将srcLocal复制到dstLocal + if (unlikely(row == 1)) { + DataCopy(dstLocal, srcLocal, row * col); + PipeBarrier(); + return; + } + for (uint32_t mask = MAX_R << 1; mask > 1; mask >>= 1) { + if (row & mask) { + // 将输入对半求最大值后放进临时空间 + Max(shareTmpUb, srcLocal, srcLocal[mask * col / 2], mask * col / 2); // 2:对矩阵按列做计算 + PipeBarrier(); + // 将余量和前一半求最大值后加到前一半上 + if (unlikely(row > mask)) { + if ((row - mask) > (mask >> 1)) { + Max(shareTmpUb, shareTmpUb, srcLocal[mask * col], mask * col / 2); // 2:对矩阵按列做计算 + PipeBarrier(); + Max(shareTmpUb, shareTmpUb, srcLocal[(mask + (mask >> 1)) * col], (row - mask - (mask >> 1)) * col); + PipeBarrier(); + } else { + Max(shareTmpUb, shareTmpUb, srcLocal[mask * col], (row - mask) * col); + PipeBarrier(); + } + } + // 每次将后一半行和前一半最大值后加到前一半上 + for (uint32_t i = mask >> 2; i > 1; i >>= 1) { + Max(shareTmpUb, shareTmpUb, shareTmpUb[i * col], i * col); + PipeBarrier(); + } + if (mask == 2) { // 2:最后一次矩阵运算处理 + DataCopy(dstLocal, shareTmpUb, col); + } else { + Max(dstLocal, shareTmpUb, shareTmpUb[col], col); + } + PipeBarrier(); + break; + } + } +} + + +/** + * @brief MatSubVec 矩阵逐行减向量 + * @param dstLocal 输出tensor [row, col] + * @param src0Local 输入tensor [row, col] + * @param src1Local 输入tensor [1, col] + * @param repeatParam 描述待处理数据的排布，包括 + row 行数 + col 列数 + dtypeMask 一次迭代参与计算元素数 + loopTimes 循环次数 + colRemain 剩余列数 + repeatStride 循环步长（内存中实际列长度） + */ +__aicore__ inline void MatSubVec(const LocalTensor &dstLocal, const LocalTensor &src0Local, + const LocalTensor &src1Local, const MatRpeatParam &repeatParam) +{ + for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) { + uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM); + uint32_t offset = 0; + for (uint32_t i = 0; i < repeatParam.loopTimes; i++) { + Sub(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset], + repeatParam.dtypeMask, repeatRowTimes, + {1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0}); + offset += repeatParam.dtypeMask; + } + if (repeatParam.colRemain > 0) { + Sub(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset], + repeatParam.colRemain, repeatRowTimes, + {1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0}); + } + } +} + +/** + * @brief MatDivVec 矩阵逐行除以向量 + * @param dstLocal 输出tensor [row, col] + * @param src0Local 输入tensor [row, col] + * @param src1Local 输入tensor [1, col] + * @param repeatParam 描述待处理数据的排布，包括 + row 行数 + col 列数 + dtypeMask 一次迭代参与计算元素数 + loopTimes 循环次数 + colRemain 剩余列数 + repeatStride 循环步长（内存中实际列长度） + */ +__aicore__ inline void MatDivVec(const LocalTensor &dstLocal, const LocalTensor &src0Local, + const LocalTensor &src1Local, const MatRpeatParam &repeatParam) +{ + for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) { + uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM); + uint32_t offset = 0; + for (uint32_t i = 0; i < repeatParam.loopTimes; i++) { + Div(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset], + repeatParam.dtypeMask, repeatRowTimes, + {1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0}); + offset += repeatParam.dtypeMask; + } + if (repeatParam.colRemain > 0) { + Div(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset], + repeatParam.colRemain, repeatRowTimes, + {1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0}); + } + } +} + +/** + * @brief MatMulVec 矩阵逐行乘以向量 + * @param dstLocal 输出tensor [row, col] + * @param src0Local 输入tensor [row, col] + * @param src1Local 输入tensor [1, col] + * @param repeatParam 描述待处理数据的排布，包括 + row 行数 + col 列数 + dtypeMask 一次迭代参与计算元素数 + loopTimes 循环次数 + colRemain 剩余列数 + repeatStride 循环步长（内存中实际列长度） + */ +__aicore__ inline void MatMulVec(const LocalTensor &dstLocal, const LocalTensor &src0Local, + const LocalTensor &src1Local, const MatRpeatParam &repeatParam) +{ + for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) { + uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM); + uint32_t offset = 0; + for (uint32_t i = 0; i < repeatParam.loopTimes; i++) { + Mul(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset], + repeatParam.dtypeMask, repeatRowTimes, + {1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0}); + offset += repeatParam.dtypeMask; + } + if (repeatParam.colRemain > 0) { + Mul(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset], + repeatParam.colRemain, repeatRowTimes, + {1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0}); + } + } +} + +/** + * @brief RowSum 矩阵对每行求和 + * @param dstLocal 输出tensor [1, row] + * @param srcLocal 输入tensor [row, col] + * @param shareTmpUb 临时buffer 内部需要的空间为 [row, col]，支持和srcLocal是同一块空间 + * @param repeatParam 描述待处理数据的排布，包括 + row 行数 + col 列数 + dtypeMask 一次迭代参与计算元素数 + loopTimes 循环次数 + colRemain 剩余列数 + repeatStride 循环步长（内存中实际列长度） + */ +__aicore__ inline void RowSum(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + const LocalTensor &shareTmpUb, const MatRpeatParam &repeatParam) +{ + uint32_t blockCount = repeatParam.loopTimes; + if (blockCount > 0 && repeatParam.colRemain > 0) { + Add(shareTmpUb, srcLocal, srcLocal[blockCount * repeatParam.dtypeMask], repeatParam.colRemain, + repeatParam.row, + {1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, repeatParam.repeatStride}); + AscendC::PipeBarrier(); + } + + for (uint32_t loopCount = blockCount >> 1; loopCount > 0; loopCount = blockCount >> 1) { + blockCount = (blockCount + 1) >> 1; + for (uint32_t i = 0; i < loopCount; i++) { + Add(shareTmpUb[i * repeatParam.dtypeMask], srcLocal[i * repeatParam.dtypeMask], + srcLocal[(i + blockCount) * repeatParam.dtypeMask], repeatParam.dtypeMask, repeatParam.row, + {1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, repeatParam.repeatStride}); + } + AscendC::PipeBarrier(); + } + + WholeReduceSum(dstLocal, shareTmpUb, + (repeatParam.col < repeatParam.dtypeMask) ? repeatParam.col : + repeatParam.dtypeMask, + repeatParam.row, 1, 1, repeatParam.repeatStride); +} + +/** + * @brief RowDivs 矩阵每行除以对应元素 + * @param dstLocal 输出tensor [row, col] + * @param src0Local 输入tensor [row, col] + * @param src1Local 输入tensor [row, 1]，需要扩展到一个datablock中(实际内存需要为[row, FP32_BLOCK_ELEMENT_NUM]) + * @param repeatParam 描述待处理数据的排布，包括 + row 行数 + col 列数 + dtypeMask 一次迭代参与计算元素数 + loopTimes 循环次数 + colRemain 剩余列数 + repeatStride 循环步长（内存中实际列长度） + */ +__aicore__ inline void RowDivs(const LocalTensor &dstLocal, const LocalTensor &src0Local, + const LocalTensor &src1Local, const MatRpeatParam &repeatParam) +{ + for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) { + uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM); + uint32_t offset = 0; + for (uint32_t i = 0; i < repeatParam.loopTimes; i++) { + Div(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local, + repeatParam.dtypeMask, repeatRowTimes, + {1, 1, 0, repeatParam.repeatStride, repeatParam.repeatStride, 1}); + offset += repeatParam.dtypeMask; + } + if (repeatParam.colRemain > 0) { + Div(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local, + repeatParam.colRemain, repeatRowTimes, + {1, 1, 0, repeatParam.repeatStride, repeatParam.repeatStride, 1}); + } + } +} + + +/** + * @brief RowMuls 矩阵每行乘以相同元素 + * @param dstLocal 输出tensor [row, col] + * @param src0Local 输入tensor [row, col] + * @param src1Local 输入tensor [row, 1]，需要扩展到一个datablock中(实际内存需要为[row, FP32_BLOCK_ELEMENT_NUM]) + * @param repeatParam 描述待处理数据的排布，包括 + row 行数 + col 列数 + dtypeMask 一次迭代参与计算元素数 + loopTimes 循环次数 + colRemain 剩余列数 + repeatStride 循环步长（内存中实际列长度） + */ +__aicore__ inline void RowMuls(const LocalTensor &dstLocal, const LocalTensor &src0Local, + const LocalTensor &src1Local, const MatRpeatParam &repeatParam) +{ + for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) { + uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM); + uint32_t offset = 0; + for (uint32_t i = 0; i < repeatParam.loopTimes; i++) { + Mul(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local, + repeatParam.dtypeMask, repeatRowTimes, + {1, 1, 0, repeatParam.repeatStride, repeatParam.repeatStride, 1}); + offset += repeatParam.dtypeMask; + } + if (repeatParam.colRemain > 0) { + Mul(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local, + repeatParam.colRemain, repeatRowTimes, + {1, 1, 0, repeatParam.repeatStride, repeatParam.repeatStride, 1}); + } + } +} + +} // namespace Compressor +#endif // COMPRESSOR_VECTOR_COMM_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch32/rms_norm.h b/csrc/ascend/attention/compressor/op_kernel/arch32/rms_norm.h new file mode 100644 index 000000000..eff872d20 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch32/rms_norm.h @@ -0,0 +1,87 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + + +/*! + * \file rms_norm.h + * \brief + */ + +#ifndef RMS_NORM_H +#define RMS_NORM_H + +#include "compressor_comm.h" +#include "compressor_vector_comm.h" + +namespace Compressor { +/** + * @brief RmsNorm 对矩阵进行rmsnorm + * @param dstLocal 输出tensor [row, col]，支持和srcLocal是同一块空间 + * @param srcLocal 输入tensor [row, col] + * @param gammaLocal 系数gamma [1, col] + * @param shareTmpUb 临时buffer 内部需要的空间为 [(row * col + row) * sizeof(float)] + * @param rmsNormParams rms所需系数，包括 + reciprocal rmsnorm系数reciprocal + epsilon rmsnorm系数epsilon + row 处理的行数 + col 列数 + */ +template +__aicore__ inline void RmsNorm(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + const LocalTensor &gammaLocal, const LocalTensor &shareTmpUb, + const RmsNormParam &rmsNormParams) +{ + uint64_t cnt = rmsNormParams.row * rmsNormParams.col; + LocalTensor temp1Local = shareTmpUb.ReinterpretCast(); + LocalTensor temp2Local = temp1Local[cnt]; + + // temp1Local = srcLocal ^ 2 + Mul(temp1Local, srcLocal, srcLocal, cnt); + PipeBarrier(); + + MatRpeatParam repeatParams = { + rmsNormParams.row, // row + rmsNormParams.col, // col + FP32_REPEAT_ELEMENT_NUM, // dtypeMask + rmsNormParams.col / FP32_REPEAT_ELEMENT_NUM, // loopTimes + rmsNormParams.col % FP32_REPEAT_ELEMENT_NUM, // colsRemain + static_cast(rmsNormParams.col / FP32_BLOCK_ELEMENT_NUM), // repeatStride + }; + + // temp2Local[row] = Sum(temp1Local) + RowSum(temp2Local, temp1Local, temp1Local, repeatParams); + PipeBarrier(); + + + // temp2Local[row] = temp2Local[row] * reciprocal(1/N) + Muls(temp2Local, temp2Local, rmsNormParams.reciprocal, rmsNormParams.row); + PipeBarrier(); + + // temp2Local[row] = temp2Local[row] + epsilon + Adds(temp2Local, temp2Local, rmsNormParams.epsilon, rmsNormParams.row); + PipeBarrier(); + + // temp2Local[row] = Sqrt(temp2Local[row]) + Sqrt(temp2Local, temp2Local, rmsNormParams.row); + PipeBarrier(); + + // temp1Local[row, 8] = brc(temp2Local[row, 1]) + Brcb(temp1Local, temp2Local, CeilDivT(rmsNormParams.row, BRCB_NUM), {1, 8}); + PipeBarrier(); + + // dstLocal = srcLocal / temp1Local(sum) + RowDivs(dstLocal, srcLocal, temp1Local, repeatParams); + PipeBarrier(); + + // dstLocal = dstLocal * gammaLocal + MatMulVec(dstLocal, dstLocal, gammaLocal, repeatParams); +} +} // namespace Compressor +#endif // MLA_PROLOG_RMS_NORM_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch32/rope.h b/csrc/ascend/attention/compressor/op_kernel/arch32/rope.h new file mode 100644 index 000000000..f6eff71dc --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch32/rope.h @@ -0,0 +1,130 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rope.h + * \brief + */ + +#ifndef ROPE_H +#define ROPE_H + +#include "compressor_comm.h" +#include "compressor_vector_comm.h" + +namespace Compressor { + +/** + * @brief SetGatherSrcOffset 计算用于interleave模式的offset + * @param gatherOffsetLocal 输出tensor [count]，数据类型需要为int64_t，使用时要转换 + * @param count offset的元素个数，一般为列数 + */ +template +__aicore__ inline void SetGatherSrcOffset(const LocalTensor &gatherOffsetLocal, uint32_t count) +{ + for (uint32_t i = 0; i < 8; i++) { + gatherOffsetLocal.SetValue(i, i ^ 1); + } + + event_t eventId_S_V = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventId_S_V); + WaitFlag(eventId_S_V); + + int32_t scalarValue = 8; + while (scalarValue < count) { + int32_t nextValue = scalarValue * 2; + PipeBarrier(); + if (nextValue < count) { + Adds(gatherOffsetLocal[scalarValue], gatherOffsetLocal, scalarValue, scalarValue); + } else { + Adds(gatherOffsetLocal[scalarValue], gatherOffsetLocal, scalarValue, count - scalarValue); + break; + } + scalarValue = nextValue; + } + PipeBarrier(); + Muls(gatherOffsetLocal, gatherOffsetLocal, static_cast(sizeof(T)), count); +} + + +/** + * @brief RotaryPosEmb 同时做row行的RotaryPosEmb，每一行的元素为col + * @param dstLocal 输出tensor [row, actualCol]，支持和srcLocal是同一块空间 + * @param srcLocal 输入tensor [row, actualCol] + * @param cosLocal cos系数tensor [row, col] + * @param sinLocal sin系数tensor [row, col] + * @param shareTmpUb 临时buffer 内部需要的空间为 [row * col * sizeof(float)] + * @param gatherOffsetcastLocal 用于interleave模式的offset，数据类型需要为uint64_t + * @param row 待处理的行数 + * @param col 待处理的列数 + * @param actualCol 实际列数 + * @param baseAddr 计算基地址 + */ +template +__aicore__ inline void RotaryPosEmb(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + const LocalTensor &cosLocal, const LocalTensor &sinLocal, + const LocalTensor &shareTmpUb, + const LocalTensor &gatherOffsetcastLocal, uint32_t row, uint32_t col, + uint32_t actualCol, uint64_t baseAddr) +{ + uint64_t cnt = row * col; + uint32_t half_col = col >> 1; + uint64_t rsvdCnt = 0; + LocalTensor reArrLocal = shareTmpUb.ReinterpretCast(); + if constexpr (MODE == ROTARY_MODE::HALF) { + DataCopy(reArrLocal, srcLocal[baseAddr + half_col], + {static_cast(row), static_cast(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM)), + static_cast(CeilDivT(actualCol - half_col, FP32_BLOCK_ELEMENT_NUM)), + static_cast(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM))}); + DataCopy(reArrLocal[half_col], srcLocal[baseAddr], + {static_cast(row), static_cast(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM)), + static_cast(CeilDivT(actualCol - half_col, FP32_BLOCK_ELEMENT_NUM)), + static_cast(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM))}); + PipeBarrier(); + Muls(reArrLocal, reArrLocal, float(-1), half_col, row, + {1, 1, static_cast(CeilDivT(static_cast(col), FP32_BLOCK_ELEMENT_NUM)), + static_cast(CeilDivT(static_cast(col), FP32_BLOCK_ELEMENT_NUM))}); + } else if constexpr (MODE == ROTARY_MODE::INTERLEAVE) { + for (uint32_t i = 0; i < row; i++) { + Gather(reArrLocal[i * col], srcLocal[i * actualCol + baseAddr], gatherOffsetcastLocal, 0, col); + } + PipeBarrier(); + uint32_t repeatTimes = cnt / FP32_REPEAT_ELEMENT_NUM; + uint32_t remainder = cnt % FP32_REPEAT_ELEMENT_NUM; + uint64_t fullMask = 0x5555555555555555; + uint64_t partialMask = 0x55; + SetVectorMask(0, fullMask); + Muls(reArrLocal, reArrLocal, float(-1), MASK_PLACEHOLDER, repeatTimes, + {1, 1, FP32_BLOCK_ELEMENT_NUM, FP32_BLOCK_ELEMENT_NUM}); + + if (unlikely(remainder > 0)) { + SetVectorMask(0, partialMask); + Muls(reArrLocal[repeatTimes * FP32_REPEAT_ELEMENT_NUM], + reArrLocal[repeatTimes * FP32_REPEAT_ELEMENT_NUM], float(-1), MASK_PLACEHOLDER, + remainder / FP32_BLOCK_ELEMENT_NUM, {1, 1, 1, 1}); + } + ResetMask(); + } + + PipeBarrier(); + BinaryRepeatParams computeParams{1, + 1, + 1, + static_cast(CeilDivT(actualCol, FP32_BLOCK_ELEMENT_NUM)), + static_cast(CeilDivT(actualCol, FP32_BLOCK_ELEMENT_NUM)), + static_cast(CeilDivT(col, FP32_BLOCK_ELEMENT_NUM))}; + Mul(dstLocal[baseAddr], srcLocal[baseAddr], cosLocal, col, row, computeParams); + Mul(reArrLocal, reArrLocal, sinLocal, cnt); + PipeBarrier(); + Add(dstLocal[baseAddr], dstLocal[baseAddr], reArrLocal, col, row, computeParams); +} +} + +#endif diff --git a/csrc/ascend/attention/compressor/op_kernel/arch32/soft_max.h b/csrc/ascend/attention/compressor/op_kernel/arch32/soft_max.h new file mode 100644 index 000000000..d8ccd685e --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch32/soft_max.h @@ -0,0 +1,52 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + + +/*! + * \file soft_max.h + * \brief + */ + +#ifndef SOFT_MAX_H +#define SOFT_MAX_H + +#include "compressor_comm.h" +#include "compressor_vector_comm.h" + +namespace Compressor { +/** + * @brief ColumnSoftMax 对矩阵按列进行SoftMax + * @param dstLocal 输出tensor [row, col]，支持和srcLocal是同一块空间 + * @param srcLocal 输入tensor [row, col] + * @param shareTmpUb 临时buffer 内部需要的空间为 [floor(row / 2) * col * sizeof(float)] + * @param row 行数 + * @param col 列数 + */ +__aicore__ inline void ColumnSoftMax(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + const LocalTensor &shareTmpUb, uint32_t row, uint32_t col) +{ + uint32_t dtypeMask = FP32_REPEAT_ELEMENT_NUM; + uint32_t dLoop = col / dtypeMask; + uint32_t dRemain = col % dtypeMask; + uint8_t repeatStride = col / FP32_BLOCK_ELEMENT_NUM; + ColumnMax(shareTmpUb, srcLocal, shareTmpUb, row, col); + PipeBarrier(); + MatSubVec(dstLocal, srcLocal, shareTmpUb, {row, col, dtypeMask, dLoop, dRemain, repeatStride}); + PipeBarrier(); + Exp(dstLocal, dstLocal, row * col); + PipeBarrier(); + ColumnSum(shareTmpUb, dstLocal, shareTmpUb, row, col); + PipeBarrier(); + MatDivVec(dstLocal, dstLocal, shareTmpUb, {row, col, dtypeMask, dLoop, dRemain, repeatStride}); +} + +} // namespace Compressor + +#endif diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_block_cube.h b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_block_cube.h new file mode 100644 index 000000000..2a9d6f94b --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_block_cube.h @@ -0,0 +1,439 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_block_cube.h + * \brief + */ + +#ifndef COMPRESSOR_BLOCK_CUBE_H +#define COMPRESSOR_BLOCK_CUBE_H + +#include "compressor_comm.h" +#include "compressor_tools.h" + +using namespace AscendC; + +namespace Compressor { + +template class CompressorBlockCube { +using MM1_OUT_T = float; +public: + __aicore__ inline CompressorBlockCube(){}; + __aicore__ inline void InitParams(const ConstInfo &constInfo, const CompressorTools &tools); + __aicore__ inline void Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitGlobalBuffers(const GlobalTensor& kvMm1ResGm, const GlobalTensor& scoreMm1ResGm); + __aicore__ inline void AllocEventID(TPipe *pipe); + __aicore__ inline void FreeEventID(TPipe *pipe); + __aicore__ inline void ComputeMm1(const RunInfo &info); + +private: + using T = float; + using X_T = typename AscendC::Conditional::type; + + __aicore__ inline uint32_t GetMSize(const RunInfo &info, uint32_t coffId); + __aicore__ inline void CopyXGmToL1(const RunInfo &info, LocalTensor xL1Tensor, uint32_t hIdx, uint32_t kBase); + __aicore__ inline void CopyWeightGmToL1(LocalTensor wL1Tensor, + uint32_t hIdx, uint32_t kBase, uint32_t coffId); + __aicore__ inline void LoadAToL0(const RunInfo &info, LocalTensor aL0Tensor, LocalTensor xL1Tensor, + uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize); + __aicore__ inline void LoadBToL0(const RunInfo &info, LocalTensor bL0Tensor, LocalTensor wL1Tensor, + uint32_t kStart, uint32_t kBase, uint32_t nStart, uint32_t nDealSize); + __aicore__ inline void MatrixMmad(LocalTensor cL0Tensor, LocalTensor aL0Tensor, + LocalTensor bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C); + __aicore__ inline void CopyOutMm1Res(const RunInfo &info, LocalTensor cL0Tensor, + uint32_t coffId, uint32_t mStart, uint32_t mDealSize, uint32_t nStart, uint32_t nDealSize); + + ConstInfo constInfo_ = {}; + CompressorTools tools_; + + // GM + GlobalTensor xGm_; + GlobalTensor wkvGm_; + GlobalTensor wgateGm_; + GlobalTensorkvMm1ResGm; + GlobalTensorscoreMm1ResGm; + GlobalTensor cuSeqlensGm_; + GlobalTensor sequsedGm_; + GlobalTensor startPosGm_; + bool isExistSeqUsed = false; + + // =================================L1 Buffer================================= + static constexpr uint32_t L1_X_SIZE = 128 * 1024; + static constexpr uint32_t L1_W_SIZE = 128 * 1024; + // L1 Buffer + TBuf xBufL1; + TBuf wBufL1; + // =================================L0 Buffer================================= + // L0 buffer size + static constexpr uint32_t L0A_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k + static constexpr uint32_t L0B_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k + static constexpr uint32_t L0C_PP_SIZE = 64 * 1024; // (128 * 2) * 64 * 4 = 64k + // L0_A + TBuf tmpBufL0A; + // L0_B + TBuf tmpBufL0B; + // L0_C + TBuf tmpBufL0C; + // =================================Event&Buffer ID=========================== + // mte2 <> mte1 EventID + static constexpr uint32_t X_EVENT0 = EVENT_ID0; + static constexpr uint32_t X_EVENT1 = EVENT_ID1; + uint32_t xBufId = 0; // 用于DB计数 + static constexpr uint32_t W_EVENT0 = EVENT_ID4; + static constexpr uint32_t W_EVENT1 = EVENT_ID5; + static constexpr uint32_t W_EVENT2 = EVENT_ID6; + static constexpr uint32_t W_EVENT3 = EVENT_ID7; + uint32_t wBufId = 0; // 用于DB计数 + // mte1 <> mmad EventID + static constexpr uint32_t L0AB_EVENT0 = EVENT_ID3; + static constexpr uint32_t L0AB_EVENT1 = EVENT_ID4; + uint32_t l0abBufId = 0; + // mmad <> fixpipe EventID + static constexpr uint32_t L0C_EVENT0 = EVENT_ID0; // 每块L0C单独分配EVENT_ID + static constexpr uint32_t L0C_EVENT1 = EVENT_ID1; + static constexpr uint32_t L0C_EVENT2 = EVENT_ID2; + static constexpr uint32_t L0C_EVENT3 = EVENT_ID3; + uint32_t l0cBufId = 0; + + // =================================Loop====================================== + uint32_t curBIdx_ = 0; + uint32_t curSIdx_ = 0; +}; + +template +__aicore__ inline void CompressorBlockCube::InitParams(const ConstInfo &constInfo, const CompressorTools &tools) +{ + this->constInfo_ = constInfo; + this->tools_ = tools; +} + +template __aicore__ inline void CompressorBlockCube::Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut) +{ + xGm_.SetGlobalBuffer((__gm__ X_T *)x); + wkvGm_.SetGlobalBuffer((__gm__ X_T *)wKv); + wgateGm_.SetGlobalBuffer((__gm__ X_T *)wGate); + startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos); + isExistSeqUsed = (seqUsed != nullptr); + if (isExistSeqUsed) { + sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed); + } + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens); + } +} + +template +__aicore__ inline void CompressorBlockCube::InitBuffers(TPipe *pipe) +{ + // L1 + // 1. coff=1时, mBase=256, kL1=256, X单次拷贝到L1的数据量最大为mBase*kL1*sizeof(BF16/FP16)=256*256*2=128K + // 2. coff=2时, mBase=128, kL1=256, r最大为128, X单次拷贝到L1的最大数据量为(128+r)*kL1*sizeof(BF16/FP16)<=128K + pipe->InitBuffer(xBufL1, L1_X_SIZE * 2); + // dBaseSize<=64, wkv和wgate各一份, kL1=256, 右矩阵为dBaseSize*2*sizeof(BF16/FP16)<=64K + // cur和pre循环使用, 2份buffer就足够 + pipe->InitBuffer(wBufL1, L1_W_SIZE * 2); + + // L0 + pipe->InitBuffer(tmpBufL0A, L0A_PP_SIZE * 2); + pipe->InitBuffer(tmpBufL0B, L0B_PP_SIZE * 2); + pipe->InitBuffer(tmpBufL0C, L0C_PP_SIZE * 4); + +} + +template +__aicore__ inline void CompressorBlockCube::InitGlobalBuffers(const GlobalTensor& kvMm1ResGm, const GlobalTensor& scoreMm1ResGm) +{ + this->kvMm1ResGm = kvMm1ResGm; + this->scoreMm1ResGm = scoreMm1ResGm; +} + +template +__aicore__ inline void CompressorBlockCube::AllocEventID(TPipe *pipe) +{ + SetFlag(X_EVENT0); + SetFlag(X_EVENT1); + + SetFlag(W_EVENT0); + SetFlag(W_EVENT1); + SetFlag(W_EVENT2); + SetFlag(W_EVENT3); + + SetFlag(L0AB_EVENT0); + SetFlag(L0AB_EVENT1); + + SetFlag(L0C_EVENT0); + SetFlag(L0C_EVENT1); + SetFlag(L0C_EVENT2); + SetFlag(L0C_EVENT3); +} + +template +__aicore__ inline void CompressorBlockCube::FreeEventID(TPipe *pipe) +{ + WaitFlag(X_EVENT0); + WaitFlag(X_EVENT1); + + WaitFlag(W_EVENT0); + WaitFlag(W_EVENT1); + WaitFlag(W_EVENT2); + WaitFlag(W_EVENT3); + + WaitFlag(L0AB_EVENT0); + WaitFlag(L0AB_EVENT1); + + WaitFlag(L0C_EVENT0); + WaitFlag(L0C_EVENT1); + WaitFlag(L0C_EVENT2); + WaitFlag(L0C_EVENT3); +} + +template +__aicore__ inline void CompressorBlockCube::CopyXGmToL1(const RunInfo &info, LocalTensor xL1Tensor, + uint32_t hIdx, uint32_t kBase) +{ + uint32_t tStart = tools_.GetTIdxByBatch(info.bStart) + info.sStart; // 此基本块在整个序列中的位置 + uint32_t copySeqCnt = info.dealSeqCnt; // 此基本块处理的长度 + + uint32_t xL1Offset = 0 * (32 / sizeof(X_T)); + uint64_t sIdx = tStart; // 起始s在整个T的起始点 + uint64_t gmOffset = sIdx * constInfo_.hSize + hIdx; + uint32_t nValue = copySeqCnt; + uint32_t dValue = kBase; // 拷贝的列数kBase + uint32_t srcDValue = constInfo_.hSize; + uint32_t dstNzC0Stride = (copySeqCnt + 15) / 16 * 16; // 1行变2行的行方向的偏移，需要16对齐 + CopySingleMatrixNDToNZ(xL1Tensor[xL1Offset], xGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride); +} + +template +__aicore__ inline void CompressorBlockCube::CopyWeightGmToL1(LocalTensor wL1Tensor, + uint32_t hIdx, uint32_t kBase, uint32_t coffId) +{ + // coffId=0, 搬运左矩阵的数据; coffId=1, 搬运右矩阵的数据 + uint64_t gmOffset = coffId * constInfo_.headDim * constInfo_.hSize + constInfo_.dIdx * constInfo_.dBaseSize * constInfo_.hSize + hIdx; + uint32_t wkvL1Offset = 0; + uint32_t wgateL1Offset = constInfo_.dBaseSize * (32 / sizeof(X_T)); // wgate与wkv的起始点相隔dBaseSize个32B + uint32_t nValue = constInfo_.dBaseSize; + uint32_t dValue = kBase; + uint32_t srcDValue = constInfo_.hSize; + uint32_t dstNzC0Stride = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数 + CopySingleMatrixNDToNZ(wL1Tensor[wkvL1Offset], wkvGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride); + CopySingleMatrixNDToNZ(wL1Tensor[wgateL1Offset], wgateGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride); +} + +template +__aicore__ inline void CompressorBlockCube::LoadAToL0(const RunInfo &info, LocalTensor aL0Tensor, + LocalTensor xL1Tensor, uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize) +{ + uint32_t mSize = info.dealSeqCnt; + + uint32_t mSizeAlign = Align(mSize, 16U); + uint32_t xTensorOffset = kStart * mSizeAlign + mStart * (32 / sizeof(X_T)); + uint32_t mDealSizeAlign = Align(mDealSize, 16U); + + LoadData2DParamsV2 loadData2DParamsV2; + loadData2DParamsV2.mStartPosition = 0; + loadData2DParamsV2.kStartPosition = 0; + loadData2DParamsV2.mStep = mDealSizeAlign / 16; + loadData2DParamsV2.kStep = kBase / (32 / sizeof(X_T)); + loadData2DParamsV2.srcStride = mSizeAlign / 16; + loadData2DParamsV2.dstStride = loadData2DParamsV2.mStep; + loadData2DParamsV2.ifTranspose = false; + LoadData(aL0Tensor, xL1Tensor[xTensorOffset], loadData2DParamsV2); +} + +template +__aicore__ inline void CompressorBlockCube::LoadBToL0(const RunInfo &info, LocalTensor bL0Tensor, LocalTensor wL1Tensor, + uint32_t kStart, uint32_t kBase, uint32_t nStart, uint32_t nDealSize) +{ + uint32_t nSize = 2 * constInfo_.dBaseSize; + + uint32_t nSizeAlign = Align(nSize, 16U); + uint64_t wTensorOffset = nSizeAlign * kStart + nStart * (32 / sizeof(X_T)); + uint32_t nDealSizeAlign = Align(nDealSize, 16U); + + LoadData2DParamsV2 loadData2DParamsV2; + loadData2DParamsV2.mStartPosition = 0; + loadData2DParamsV2.kStartPosition = 0; + loadData2DParamsV2.mStep = nDealSizeAlign / 16; + loadData2DParamsV2.kStep = kBase / (32 / sizeof(X_T)); + loadData2DParamsV2.srcStride = nSizeAlign / 16; + loadData2DParamsV2.dstStride = loadData2DParamsV2.mStep; + loadData2DParamsV2.ifTranspose = false; + LoadData(bL0Tensor, wL1Tensor[wTensorOffset], loadData2DParamsV2); +} + +template +__aicore__ inline void CompressorBlockCube::MatrixMmad(LocalTensor cL0Tensor, LocalTensor aL0Tensor, + LocalTensor bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C) +{ + MmadParams mmadParams; + mmadParams.m = mActSize < 16 ? 16 : mActSize; + mmadParams.n = nDealSize; + mmadParams.k = kActSize; + mmadParams.cmatrixInitVal = isInitL0C; + mmadParams.cmatrixSource = false; + Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams); +} + +template +__aicore__ inline void CompressorBlockCube::CopyOutMm1Res(const RunInfo &info, LocalTensor cL0Tensor, + uint32_t coffId, uint32_t mStart, uint32_t mDealSize, uint32_t nStart, uint32_t nDealSize) +{ + // coffId=0, 存左矩阵的数据; coffId=1, 存右矩阵的数据 + FixpipeParamsV220 fixParams; + fixParams.mSize = mDealSize; + fixParams.srcStride = (mDealSize + 15) / 16 * 16; // 需要16对齐 + fixParams.dstStride = (uint32_t)COMP::coff * constInfo_.headDim; + fixParams.ndNum = 1; + + uint64_t dbOffset = info.cubeDbIdx * constInfo_.dbSize; + uint64_t gmOffset = constInfo_.dIdx * constInfo_.dBaseSize + coffId * constInfo_.headDim + mStart * fixParams.dstStride + dbOffset; + uint32_t kvOffset = (mDealSize + 15) / 16 * 16 * nStart; + uint32_t scoreOffset = (mDealSize + 15) / 16 * 16 * ((nStart + constInfo_.dBaseSize) % (2 * constInfo_.dBaseSize)); + if (nStart < constInfo_.dBaseSize) { + fixParams.nSize = min(constInfo_.dBaseSize - nStart, nDealSize); + Fixpipe(kvMm1ResGm[gmOffset], cL0Tensor[kvOffset], fixParams); + } + if (nStart + nDealSize > constInfo_.dBaseSize) { + fixParams.nSize = min(nStart + nDealSize - constInfo_.dBaseSize, nDealSize); + Fixpipe(scoreMm1ResGm[gmOffset], cL0Tensor[scoreOffset], fixParams); + } +} + +template +__aicore__ inline uint32_t CompressorBlockCube::GetMSize(const RunInfo &info, uint32_t coffId) +{ + return info.dealSeqCnt; +} + +template +__aicore__ inline void CompressorBlockCube::ComputeMm1(const RunInfo &info) +{ + static constexpr uint32_t K_L1_BASE = 256; + static constexpr uint32_t M_L0_BASE = 128; + static constexpr uint32_t K_L0_BASE = 128; + static constexpr uint32_t N_L0_BASE = 128; + uint32_t nCoff = (uint32_t)COMP::coff; + + // hSize为K_SIZE=512的倍数 + uint32_t hStart = info.hStart; + uint32_t hSize = info.dealKSize; + uint32_t hIdxStart = (constInfo_.aiCoreIdx % constInfo_.dBasicBlockNum) * K_L1_BASE; // 每组核内的h循环起始不同 + uint32_t kSize = K_L1_BASE; + for (uint32_t h = 0; h < hSize; h += K_L1_BASE) { + // h方向错位搬运 + uint32_t hIdx = (h + hIdxStart) % (CeilDivT(hSize, K_L1_BASE) * K_L1_BASE); + if (hIdx + K_L1_BASE > hSize) { + kSize = hSize - hIdx; + } else { + kSize = K_L1_BASE; + } + bool isFirst = (h == 0); + bool isLast = (h + K_L1_BASE >= hSize); + WaitFlag(X_EVENT0 + xBufId); + LocalTensor xL1Tensor = xBufL1.GetWithOffset(L1_X_SIZE / sizeof(X_T), xBufId * L1_X_SIZE); + CopyXGmToL1(info, xL1Tensor, hStart + hIdx, kSize); + SetFlag(X_EVENT0 + xBufId); + WaitFlag(X_EVENT0 + xBufId); + for (uint32_t i = nCoff; i > 0; i--) { + // coffId=0, 计算pre数据; coffId=1, 计算cur数据 + uint32_t coffId = i - 1; + WaitFlag(W_EVENT0 + wBufId); + LocalTensor wL1Tensor = wBufL1.GetWithOffset(L1_W_SIZE / sizeof(X_T), wBufId * L1_W_SIZE); + CopyWeightGmToL1(wL1Tensor, hStart + hIdx, kSize, coffId); + SetFlag(W_EVENT0 + wBufId); + WaitFlag(W_EVENT0 + wBufId); + + uint32_t mSize = GetMSize(info, coffId); + uint32_t actMDealSize = M_L0_BASE; + for (uint32_t mL0 = 0; mL0 < mSize; mL0 += M_L0_BASE) { + if (mL0 + M_L0_BASE > mSize) { + actMDealSize = mSize - mL0; + } + uint32_t nDealSize = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数 + uint32_t actNDealSize = N_L0_BASE; + for (uint32_t nL0 = 0; nL0 < nDealSize; nL0 += N_L0_BASE) { + if (nL0 + N_L0_BASE > nDealSize) { + actNDealSize = nDealSize - nL0; + } + l0cBufId = (mL0 / M_L0_BASE) * 2 + (nL0 / N_L0_BASE) + coffId; + LocalTensor cL0Tensor = + tmpBufL0C.GetWithOffset((L0C_PP_SIZE / sizeof(T)), l0cBufId * L0C_PP_SIZE); + if (isFirst) { + WaitFlag(L0C_EVENT0 + l0cBufId); + } + uint32_t actKDealSize = K_L0_BASE; + for (uint32_t kL0 = 0; kL0 < kSize; kL0 += K_L0_BASE) { + if (kL0 + K_L0_BASE > kSize) { + actKDealSize = kSize - kL0; + } + WaitFlag(L0AB_EVENT0 + l0abBufId); + LocalTensor aL0Tensor = + tmpBufL0A.GetWithOffset(L0A_PP_SIZE / sizeof(X_T), l0abBufId * L0A_PP_SIZE); + LocalTensor bL0Tensor = + tmpBufL0B.GetWithOffset(L0B_PP_SIZE / sizeof(X_T), l0abBufId * L0B_PP_SIZE); + LoadAToL0(info, aL0Tensor, xL1Tensor, kL0, actKDealSize, mL0, actMDealSize); + LoadBToL0(info, bL0Tensor, wL1Tensor, kL0, actKDealSize, nL0, actNDealSize); + SetFlag(L0AB_EVENT0 + l0abBufId); + WaitFlag(L0AB_EVENT0 + l0abBufId); + bool isInitL0C = isFirst && (kL0 == 0); + MatrixMmad(cL0Tensor, aL0Tensor, bL0Tensor, actMDealSize, actNDealSize, actKDealSize, + isInitL0C); + SetFlag(L0AB_EVENT0 + l0abBufId); + l0abBufId = (l0abBufId + 1) % 2; + } + if (isLast) { + SetFlag(L0C_EVENT0 + l0cBufId); + WaitFlag(L0C_EVENT0 + l0cBufId); + CopyOutMm1Res(info, cL0Tensor, coffId, mL0, actMDealSize, nL0, actNDealSize); + SetFlag(L0C_EVENT0 + l0cBufId); + } + } + } + + SetFlag(W_EVENT0 + wBufId); + wBufId = (wBufId + 1) % 2; + } + SetFlag(X_EVENT0 + xBufId); + xBufId = (xBufId + 1) % 2; + } +} + +} // namespace Compressor + +#endif // COMPRESSOR_BLOCK_CUBE_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_block_cube_full_load.h b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_block_cube_full_load.h new file mode 100644 index 000000000..1d296786b --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_block_cube_full_load.h @@ -0,0 +1,437 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_block_cube_full_load.h + * \brief + */ + +#ifndef COMPRESSOR_BLOCK_CUBE_FULL_LOAD_H +#define COMPRESSOR_BLOCK_CUBE_FULL_LOAD_H + +#include "compressor_comm.h" +#include "compressor_tools.h" + +using namespace AscendC; + +namespace Compressor { + +template class CompressorBlockCubeFullLoad { +using MM1_OUT_T = float; +public: + __aicore__ inline CompressorBlockCubeFullLoad(){}; + __aicore__ inline void InitParams(const ConstInfo &constInfo, const CompressorTools &tools); + __aicore__ inline void Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitGlobalBuffers(const GlobalTensor& kvMm1ResGm, const GlobalTensor& scoreMm1ResGm); + __aicore__ inline void AllocEventID(TPipe *pipe); + __aicore__ inline void FreeEventID(TPipe *pipe); + __aicore__ inline void ComputeMm1(const RunInfo &info); + +private: + using T = float; + using X_T = typename AscendC::Conditional::type; + + __aicore__ inline uint32_t GetMSize(const RunInfo &info, uint32_t coffId); + __aicore__ inline void CopyXGmToL1(LocalTensor xL1Tensor, uint32_t hIdx, uint32_t kBase); + __aicore__ inline void CopyWeightGmToL1(LocalTensor wL1Tensor, + uint32_t hIdx, uint32_t kBase, uint32_t coffId); + __aicore__ inline void LoadAToL0(LocalTensor aL0Tensor, LocalTensor xL1Tensor, + uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize); + __aicore__ inline void LoadBToL0(LocalTensor bL0Tensor, LocalTensor wL1Tensor, + uint32_t kStart, uint32_t kBase, uint32_t nStart, uint32_t nDealSize); + __aicore__ inline void MatrixMmad(LocalTensor cL0Tensor, LocalTensor aL0Tensor, + LocalTensor bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C); + __aicore__ inline void CopyOutMm1Res(const RunInfo &info, LocalTensor cL0Tensor, + uint32_t coffId, uint32_t mStart, uint32_t mDealSize, uint32_t nStart, uint32_t nDealSize); + + ConstInfo constInfo_ = {}; + CompressorTools tools_; + + // GM + GlobalTensor xGm_; + GlobalTensor wkvGm_; + GlobalTensor wgateGm_; + GlobalTensorkvMm1ResGm; + GlobalTensorscoreMm1ResGm; + GlobalTensor cuSeqlensGm_; + GlobalTensor sequsedGm_; + GlobalTensor startPosGm_; + bool isExistSeqUsed = false; + + // =================================L1 Buffer================================= + static constexpr uint32_t L1_X_SIZE = 128 * 1024; + static constexpr uint32_t L1_W_SIZE = 128 * 1024; + // L1 Buffer + TBuf xBufL1; + TBuf wBufL1; + // =================================L0 Buffer================================= + // L0 buffer size + static constexpr uint32_t L0A_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k + static constexpr uint32_t L0B_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k + static constexpr uint32_t L0C_PP_SIZE = 64 * 1024; // (128 * 2) * 64 * 4 = 64k + // L0_A + TBuf tmpBufL0A; + // L0_B + TBuf tmpBufL0B; + // L0_C + TBuf tmpBufL0C; + // =================================Event&Buffer ID=========================== + // mte2 <> mte1 EventID + static constexpr uint32_t X_EVENT0 = EVENT_ID0; + static constexpr uint32_t X_EVENT1 = EVENT_ID1; + uint32_t xBufId = 0; // 用于DB计数 + static constexpr uint32_t W_EVENT0 = EVENT_ID4; + static constexpr uint32_t W_EVENT1 = EVENT_ID5; + static constexpr uint32_t W_EVENT2 = EVENT_ID6; + static constexpr uint32_t W_EVENT3 = EVENT_ID7; + uint32_t wBufId = 0; // 用于DB计数 + // mte1 <> mmad EventID + static constexpr uint32_t L0AB_EVENT0 = EVENT_ID3; + static constexpr uint32_t L0AB_EVENT1 = EVENT_ID4; + uint32_t l0abBufId = 0; + // mmad <> fixpipe EventID + static constexpr uint32_t L0C_EVENT0 = EVENT_ID0; // 每块L0C单独分配EVENT_ID + static constexpr uint32_t L0C_EVENT1 = EVENT_ID1; + static constexpr uint32_t L0C_EVENT2 = EVENT_ID2; + static constexpr uint32_t L0C_EVENT3 = EVENT_ID3; + uint32_t l0cBufId = 0; + + // =================================Loop====================================== + uint32_t curBIdx_ = 0; + uint32_t curSIdx_ = 0; +}; + +template +__aicore__ inline void CompressorBlockCubeFullLoad::InitParams(const ConstInfo &constInfo, const CompressorTools &tools) +{ + this->constInfo_ = constInfo; + this->tools_ = tools; +} + +template __aicore__ inline void CompressorBlockCubeFullLoad::Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut) +{ + xGm_.SetGlobalBuffer((__gm__ X_T *)x); + wkvGm_.SetGlobalBuffer((__gm__ X_T *)wKv); + wgateGm_.SetGlobalBuffer((__gm__ X_T *)wGate); + startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos); + isExistSeqUsed = (seqUsed != nullptr); + if (isExistSeqUsed) { + sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed); + } + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens); + } +} + +template +__aicore__ inline void CompressorBlockCubeFullLoad::InitBuffers(TPipe *pipe) +{ + // L1 + // 1. coff=1时, mBase=256, kL1=256, X单次拷贝到L1的数据量最大为mBase*kL1*sizeof(BF16/FP16)=256*256*2=128K + // 2. coff=2时, mBase=128, kL1=256, r最大为128, X单次拷贝到L1的最大数据量为(128+r)*kL1*sizeof(BF16/FP16)<=128K + pipe->InitBuffer(xBufL1, L1_X_SIZE * 2); + // dBaseSize<=64, wkv和wgate各一份, kL1=256, 右矩阵为dBaseSize*2*sizeof(BF16/FP16)<=64K + // cur和pre循环使用, 2份buffer就足够 + pipe->InitBuffer(wBufL1, L1_W_SIZE * 2); + + // L0 + pipe->InitBuffer(tmpBufL0A, L0A_PP_SIZE * 2); + pipe->InitBuffer(tmpBufL0B, L0B_PP_SIZE * 2); + pipe->InitBuffer(tmpBufL0C, L0C_PP_SIZE * 4); + +} + +template +__aicore__ inline void CompressorBlockCubeFullLoad::InitGlobalBuffers(const GlobalTensor& kvMm1ResGm, const GlobalTensor& scoreMm1ResGm) +{ + this->kvMm1ResGm = kvMm1ResGm; + this->scoreMm1ResGm = scoreMm1ResGm; +} + +template +__aicore__ inline void CompressorBlockCubeFullLoad::AllocEventID(TPipe *pipe) +{ + SetFlag(X_EVENT0); + SetFlag(X_EVENT1); + + SetFlag(W_EVENT0); + SetFlag(W_EVENT1); + SetFlag(W_EVENT2); + SetFlag(W_EVENT3); + + SetFlag(L0AB_EVENT0); + SetFlag(L0AB_EVENT1); + + SetFlag(L0C_EVENT0); + SetFlag(L0C_EVENT1); + SetFlag(L0C_EVENT2); + SetFlag(L0C_EVENT3); +} + +template +__aicore__ inline void CompressorBlockCubeFullLoad::FreeEventID(TPipe *pipe) +{ + WaitFlag(X_EVENT0); + WaitFlag(X_EVENT1); + + WaitFlag(W_EVENT0); + WaitFlag(W_EVENT1); + WaitFlag(W_EVENT2); + WaitFlag(W_EVENT3); + + WaitFlag(L0AB_EVENT0); + WaitFlag(L0AB_EVENT1); + + WaitFlag(L0C_EVENT0); + WaitFlag(L0C_EVENT1); + WaitFlag(L0C_EVENT2); + WaitFlag(L0C_EVENT3); +} + +template +__aicore__ inline void CompressorBlockCubeFullLoad::CopyXGmToL1(LocalTensor xL1Tensor, + uint32_t hIdx, uint32_t kBase) +{ + uint32_t copySeqCnt = constInfo_.mEnd - constInfo_.mStart; //info.dealSeqCnt; // 此基本块处理的长度 + + uint32_t xL1Offset = 0 * (32 / sizeof(X_T)); + uint64_t sIdx = constInfo_.mStart; // 起始s在整个T的起始点 + uint64_t gmOffset = sIdx * constInfo_.hSize + hIdx; + uint32_t nValue = copySeqCnt; + uint32_t dValue = kBase; // 拷贝的列数kBase + uint32_t srcDValue = constInfo_.hSize; + uint32_t dstNzC0Stride = (copySeqCnt + 15) / 16 * 16; // 1行变2行的行方向的偏移，需要16对齐 + CopySingleMatrixNDToNZ(xL1Tensor[xL1Offset], xGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride); +} + +template +__aicore__ inline void CompressorBlockCubeFullLoad::CopyWeightGmToL1(LocalTensor wL1Tensor, + uint32_t hIdx, uint32_t kBase, uint32_t coffId) +{ + uint64_t gmOffset = coffId * constInfo_.headDim * constInfo_.hSize + constInfo_.nStart * constInfo_.hSize + hIdx; + uint32_t wkvL1Offset = 0; + uint32_t wgateL1Offset = constInfo_.dBaseSize * (32 / sizeof(X_T)); // wgate与wkv的起始点相隔dBaseSize个32B + uint32_t nValue = constInfo_.dBaseSize; + uint32_t dValue = kBase; + uint32_t srcDValue = constInfo_.hSize; + uint32_t dstNzC0Stride = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数 + CopySingleMatrixNDToNZ(wL1Tensor[wkvL1Offset], wkvGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride); + CopySingleMatrixNDToNZ(wL1Tensor[wgateL1Offset], wgateGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride); +} + +template +__aicore__ inline void CompressorBlockCubeFullLoad::LoadAToL0(LocalTensor aL0Tensor, + LocalTensor xL1Tensor, uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize) +{ + uint32_t mSize = constInfo_.mEnd - constInfo_.mStart; + + uint32_t mSizeAlign = Align(mSize, 16U); + uint32_t xTensorOffset = kStart * mSizeAlign + mStart * (32 / sizeof(X_T)); + uint32_t mDealSizeAlign = Align(mDealSize, 16U); + + LoadData2DParamsV2 loadData2DParamsV2; + loadData2DParamsV2.mStartPosition = 0; + loadData2DParamsV2.kStartPosition = 0; + loadData2DParamsV2.mStep = mDealSizeAlign / 16; + loadData2DParamsV2.kStep = kBase / (32 / sizeof(X_T)); + loadData2DParamsV2.srcStride = mSizeAlign / 16; + loadData2DParamsV2.dstStride = loadData2DParamsV2.mStep; + loadData2DParamsV2.ifTranspose = false; + LoadData(aL0Tensor, xL1Tensor[xTensorOffset], loadData2DParamsV2); +} + +template +__aicore__ inline void CompressorBlockCubeFullLoad::LoadBToL0(LocalTensor bL0Tensor, LocalTensor wL1Tensor, + uint32_t kStart, uint32_t kBase, uint32_t nStart, uint32_t nDealSize) +{ + uint32_t nSize = 2 * constInfo_.dBaseSize; + + uint32_t nSizeAlign = Align(nSize, 16U); + uint64_t wTensorOffset = nSizeAlign * kStart + nStart * (32 / sizeof(X_T)); + uint32_t nDealSizeAlign = Align(nDealSize, 16U); + + LoadData2DParamsV2 loadData2DParamsV2; + loadData2DParamsV2.mStartPosition = 0; + loadData2DParamsV2.kStartPosition = 0; + loadData2DParamsV2.mStep = nDealSizeAlign / 16; + loadData2DParamsV2.kStep = kBase / (32 / sizeof(X_T)); + loadData2DParamsV2.srcStride = nSizeAlign / 16; + loadData2DParamsV2.dstStride = loadData2DParamsV2.mStep; + loadData2DParamsV2.ifTranspose = false; + LoadData(bL0Tensor, wL1Tensor[wTensorOffset], loadData2DParamsV2); +} + +template +__aicore__ inline void CompressorBlockCubeFullLoad::MatrixMmad(LocalTensor cL0Tensor, LocalTensor aL0Tensor, + LocalTensor bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C) +{ + MmadParams mmadParams; + mmadParams.m = mActSize < 16 ? 16 : mActSize; + mmadParams.n = nDealSize; + mmadParams.k = kActSize; + mmadParams.cmatrixInitVal = isInitL0C; + mmadParams.cmatrixSource = false; + Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams); +} + +template +__aicore__ inline void CompressorBlockCubeFullLoad::CopyOutMm1Res(const RunInfo &info, LocalTensor cL0Tensor, + uint32_t coffId, uint32_t mStart, uint32_t mDealSize, uint32_t nStart, uint32_t nDealSize) +{ + // coffId=0, 存左矩阵的数据; coffId=1, 存右矩阵的数据 + FixpipeParamsV220 fixParams; + fixParams.mSize = mDealSize; + fixParams.srcStride = (mDealSize + 15) / 16 * 16; // 需要16对齐 + fixParams.dstStride = (uint32_t)COMP::coff * constInfo_.headDim; + fixParams.ndNum = 1; + + uint64_t dbOffset = info.cubeDbIdx * constInfo_.dbSize; + uint64_t gmOffset = constInfo_.nStart + coffId * constInfo_.headDim + mStart * fixParams.dstStride + dbOffset; + uint32_t kvOffset = (mDealSize + 15) / 16 * 16 * nStart; + uint32_t scoreOffset = (mDealSize + 15) / 16 * 16 * ((nStart + constInfo_.dBaseSize) % (2 * constInfo_.dBaseSize)); + + if (nStart < constInfo_.dBaseSize) { + fixParams.nSize = min(constInfo_.dBaseSize - nStart, nDealSize); + Fixpipe(kvMm1ResGm[gmOffset], cL0Tensor[kvOffset], fixParams); + } + if (nStart + nDealSize > constInfo_.dBaseSize) { + fixParams.nSize = min(nStart + nDealSize - constInfo_.dBaseSize, nDealSize); + Fixpipe(scoreMm1ResGm[gmOffset], cL0Tensor[scoreOffset], fixParams); + } +} + +template +__aicore__ inline uint32_t CompressorBlockCubeFullLoad::GetMSize(const RunInfo &info, uint32_t coffId) +{ + return info.dealSeqCnt; +} + +template +__aicore__ inline void CompressorBlockCubeFullLoad::ComputeMm1(const RunInfo &info) +{ + uint32_t mSize = info.dealSeqCnt; + if (mSize == 0) { + return; + } + static constexpr uint32_t K_L1_BASE = 128; + static constexpr uint32_t M_L0_BASE = 128; + static constexpr uint32_t K_L0_BASE = 128; + static constexpr uint32_t N_L0_BASE = 128; + uint32_t nCoff = (uint32_t)COMP::coff; + + // hSize为K_SIZE=512的倍数 + uint32_t hStart = constInfo_.kStart; + uint32_t hSize = constInfo_.kEnd - constInfo_.kStart; + uint32_t hIdxStart = (constInfo_.aiCoreIdx % constInfo_.dBasicBlockNum) * K_L1_BASE; // 每组核内的h循环起始不同 + uint32_t kSize = K_L1_BASE; + for (uint32_t h = 0; h < hSize; h += K_L1_BASE) { + // h方向错位搬运 + uint32_t hIdx = (h + hIdxStart) % (CeilDivT(hSize, K_L1_BASE) * K_L1_BASE); + if (hIdx + K_L1_BASE > hSize) { + kSize = hSize - hIdx; + } else { + kSize = K_L1_BASE; + } + bool isFirst = (h == 0); + bool isLast = (h + K_L1_BASE >= hSize); + WaitFlag(X_EVENT0 + xBufId); + LocalTensor xL1Tensor = xBufL1.GetWithOffset(L1_X_SIZE / sizeof(X_T), xBufId * L1_X_SIZE); + CopyXGmToL1(xL1Tensor, hStart + hIdx, kSize); + SetFlag(X_EVENT0 + xBufId); + WaitFlag(X_EVENT0 + xBufId); + for (uint32_t i = nCoff; i > 0; i--) { + // coffId=0, 计算pre数据; coffId=1, 计算cur数据 + uint32_t coffId = i - 1; + WaitFlag(W_EVENT0 + wBufId); + LocalTensor wL1Tensor = wBufL1.GetWithOffset(L1_W_SIZE / sizeof(X_T), wBufId * L1_W_SIZE); + CopyWeightGmToL1(wL1Tensor, hStart + hIdx, kSize, coffId); + SetFlag(W_EVENT0 + wBufId); + WaitFlag(W_EVENT0 + wBufId); + + uint32_t actMDealSize = M_L0_BASE; + for (uint32_t mL0 = 0; mL0 < mSize; mL0 += M_L0_BASE) { + if (mL0 + M_L0_BASE > mSize) { + actMDealSize = mSize - mL0; + } + uint32_t nDealSize = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数 + uint32_t actNDealSize = N_L0_BASE; + for (uint32_t nL0 = 0; nL0 < nDealSize; nL0 += N_L0_BASE) { + if (nL0 + N_L0_BASE > nDealSize) { + actNDealSize = nDealSize - nL0; + } + l0cBufId = (mL0 / M_L0_BASE) * 2 + (nL0 / N_L0_BASE) + coffId; + LocalTensor cL0Tensor = tmpBufL0C.GetWithOffset((L0C_PP_SIZE / sizeof(T)), l0cBufId * L0C_PP_SIZE); + if (isFirst) { + WaitFlag(L0C_EVENT0 + l0cBufId); + } + uint32_t actKDealSize = K_L0_BASE; + for (uint32_t kL0 = 0; kL0 < kSize; kL0 += K_L0_BASE) { + if (kL0 + K_L0_BASE > kSize) { + actKDealSize = kSize - kL0; + } + WaitFlag(L0AB_EVENT0 + l0abBufId); + LocalTensor aL0Tensor = tmpBufL0A.GetWithOffset(L0A_PP_SIZE / sizeof(X_T), l0abBufId * L0A_PP_SIZE); + LocalTensor bL0Tensor = tmpBufL0B.GetWithOffset(L0B_PP_SIZE / sizeof(X_T), l0abBufId * L0B_PP_SIZE); + LoadAToL0(aL0Tensor, xL1Tensor, kL0, actKDealSize, mL0, actMDealSize); + LoadBToL0(bL0Tensor, wL1Tensor, kL0, actKDealSize, nL0, actNDealSize); + SetFlag(L0AB_EVENT0 + l0abBufId); + WaitFlag(L0AB_EVENT0 + l0abBufId); + bool isInitL0C = isFirst && (kL0 == 0); + MatrixMmad(cL0Tensor, aL0Tensor, bL0Tensor, actMDealSize, actNDealSize, actKDealSize, isInitL0C); + SetFlag(L0AB_EVENT0 + l0abBufId); + l0abBufId = (l0abBufId + 1) % 2; + } + if (isLast) { + SetFlag(L0C_EVENT0 + l0cBufId); + WaitFlag(L0C_EVENT0 + l0cBufId); + CopyOutMm1Res(info, cL0Tensor, coffId, mL0, actMDealSize, nL0, actNDealSize); + SetFlag(L0C_EVENT0 + l0cBufId); + } + } + } + + SetFlag(W_EVENT0 + wBufId); + wBufId = (wBufId + 1) % 2; + } + SetFlag(X_EVENT0 + xBufId); + xBufId = (xBufId + 1) % 2; + } +} + +} // namespace Compressor + +#endif // COMPRESSOR_BLOCK_CUBE_FULL_LOAD_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_block_vec.h b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_block_vec.h new file mode 100644 index 000000000..c1e746ee2 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_block_vec.h @@ -0,0 +1,1369 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_block_vec.h + * \brief + */ + +#ifndef COMPRESSOR_BLOCK_VEC_H +#define COMPRESSOR_BLOCK_VEC_H + +#include "compressor_comm.h" +#include "compressor_tools.h" +#include "vf/vf_softmax.h" +#include "vf/vf_add.h" +#include "vf/vf_mul.h" +#include "vf/vf_rms_norm.h" +#include "vf/vf_rope.h" + + +using namespace AscendC; + +namespace Compressor { +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +template +class CompressorBlockVector { +public: + static constexpr bool X_DTYPE = COMP::xDtype == X_DTYPE::BF16; + static constexpr uint64_t BLOCK_VEC_BASE_BUFFER_SIZE = 32 * 1024; // 32k + static constexpr uint32_t DATABLOCK_BYTES = 32; + static constexpr float FLOAT_ZERO = 0; + float SOFTMAX_MIN_NUM = static_cast(-1.0 / 0.0); + // =================================类型定义区================================= + // 中间计算数据类型为float，高精度模式 + using T = float; + using X_T = typename AscendC::Conditional::type; + + __aicore__ inline CompressorBlockVector(){}; + // =================================设置参数================================= + __aicore__ inline void InitParams(const ConstInfo &constInfo, const CompressorTools &tools); + __aicore__ inline void Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut); + // =================================资源管理================================= + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + // =================================执行计算================================= + __aicore__ inline void ComputeVec1(const Vec1RunInfo &info); + __aicore__ inline void InitVec1GlobalTensor(GlobalTensor kvMm1ResGm, GlobalTensor scoreMm1ResGm, + GlobalTensor kvCacheTcGm, GlobalTensor scoreCacheTcGm, + GlobalTensor vec1ResGm, GlobalTensor vec2InputGm); + __aicore__ inline void ComputeVec2(const Vec2RunInfo &info); + +protected: + GlobalTensor vec1ResGm_; + GlobalTensor vec2InputGm_; + GlobalTensor scoreMm1ResGm_; + GlobalTensor kvMm1ResGm_; + GlobalTensor kvCacheTcGm_; + GlobalTensor scoreCacheTcGm_; + +private: + __aicore__ inline uint32_t GetSeqUsed(uint32_t bIdx); + __aicore__ inline uint32_t GetStartPos(uint32_t bIdx); + __aicore__ inline uint32_t GetSeqLength(uint32_t bIdx); + __aicore__ inline void CalcGlobalScStart(uint32_t bStart, uint32_t scStart, uint32_t bEnd, uint32_t scEnd, + uint64_t &globalScStart); + __aicore__ inline void UpdateOutputIdx(uint32_t &outputBStart, uint32_t &outputSStart, uint32_t &dealScSize, + uint32_t &curDealScSize); + __aicore__ inline void DealVec1BaseBlock(const Vec1RunInfo &info, CompressorVec1SliceIterator &sliceIterator, + const LoopInfo &loopInfo, uint32_t dStartIdx, uint32_t dDealSize, + uint32_t dBaseSize); + __aicore__ inline void CopyInApe(const LocalTensor &apeUb, uint32_t dStartIdx, uint32_t dDealSize); + __aicore__ inline void AddApeToScore(const LocalTensor &scoreLocal, const LocalTensor &apeUb, + const Vec1SliceInfo &sliceInfo, uint32_t dDealSize); + __aicore__ inline void AddSingleApeToScore(const LocalTensor &scoreLocal, const LocalTensor &apeUb, + const Vec1SliceInfo &sliceInfo, uint32_t dDealSize); + template + __aicore__ inline void DataCopyAlignUbToUb(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, uint32_t srcSingleRowCount, + uint32_t dstSingleRowCount); + template + __aicore__ inline void DataCopyAlignGmToUb(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + uint32_t copyRowCount, uint32_t copyColCount, uint32_t srcSingleRowCount, + uint32_t dstSingleRowCount); + template + __aicore__ inline void DataCopyAlignUbToGm(const GlobalTensor &dstGm, const LocalTensor &srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, uint32_t srcSingleRowCount, + uint32_t dstSingleRowCount); + template + __aicore__ inline void DataCopyWithOutputQue(const GlobalTensor &dstGm, const LocalTensor &srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount); + template + __aicore__ inline void DataCopyWithInputQue(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount); + template + __aicore__ inline void AddMultiDataToUb(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + uint32_t dealRowCount, uint32_t dealColCount, uint32_t srcSingleRowCount, + uint32_t dstSingleRowCount, uint32_t repeatTimes, uint64_t offset); + __aicore__ inline void PadAlign(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + const Vec1SliceInfo &sliceInfo, uint32_t dStartIdx, uint32_t dDealSize); + template + __aicore__ inline void OverLap(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + const GlobalTensor &srcGm, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const GlobalTensor &cacheTcGm, + const Vec1RunInfo &info, const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, + uint32_t dStartIdx, uint32_t globalSeqIdx, uint32_t dDealSize); + __aicore__ inline void FromWokrSpaceToUb(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + const Vec1SliceInfo &sliceInfo, const StatisticInfo &statisticInfo, + uint32_t dStartIdx, uint32_t dDealSize); + __aicore__ inline void WriteToCacheState(const GlobalTensor &state, const GlobalTensor &blockTableGm, + const LocalTensor &input, uint32_t batchIdx, uint32_t startSeqIdx, + uint32_t endSeqIdx, uint32_t dStartIdx, uint32_t dDealSize, + uint32_t stateIdx); + __aicore__ inline void ReadFromCacheState(const LocalTensor &output, const GlobalTensor &state, + const GlobalTensor &blockTableGm, uint32_t batchIdx, + uint32_t startSeqIdx, uint32_t endSeqIdx, uint32_t dStartIdx, + uint32_t dDealSize, uint32_t stateIdx); + __aicore__ inline void SaveToWorkSpace(const LocalTensor &srcLocal, const GlobalTensor &cacheTcGm, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, uint32_t dStartIdx, + uint32_t dDealSize); + __aicore__ inline void LoadFromWorkSpace(const LocalTensor &dstLocal, const GlobalTensor &cacheTcGm, + const GlobalTensor &srcGm, const LocalTensor &srcLocal, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, + uint32_t dStartIdx, uint32_t globalSeqIdx, uint32_t dDealSize); + __aicore__ inline void SoftmaxDN(const LocalTensor &scoreLocal, uint32_t tcDealSize, uint32_t dDealSize); + __aicore__ inline void KvMulReduceScore(const LocalTensor &kvLocal, const LocalTensor &scoreLocal, + const LocalTensor &dstLocal, uint32_t tcDealSize, uint32_t dDealSize); + __aicore__ inline void OverLapScoreKv(const LocalTensor &scoreLocal, const LocalTensor &kvLocal, + const Vec1RunInfo &info, const LoopInfo &loopInfo, + const StatisticInfo &statisticInfo, const Vec1SliceInfo &originSliceInfo, + uint32_t dStartIdx, uint32_t dDealSize, uint32_t dBaseSize, + uint32_t needDealTcSize); + __aicore__ inline void CopyOutVec1Res(const GlobalTensor &resGm, const Vec1RunInfo &info, + const LocalTensor &comperssoredUb, uint32_t compressTcSize, + uint32_t dStartIdx, uint32_t dDealSize); + __aicore__ inline void CalcGroupInfo(const Vec1RunInfo &info, Vec1SplitInfo &splitInfo); + __aicore__ inline void CalcTaskDistribution(const Vec1RunInfo &info, Vec1SplitInfo &splitInfo); + __aicore__ inline void UpdateIteratorState(const Vec1RunInfo &info, Vec1SplitInfo &splitInfo); + __aicore__ inline void CalcTilingStrategy(Vec1SplitInfo &splitInfo); + __aicore__ inline Vec1SplitInfo SplitCoreV1(const Vec1RunInfo &info); + __aicore__ inline Vec2SplitInfo SplitCoreV2(const Vec2RunInfo &info); + __aicore__ inline void MultRowRmsNorm(const LocalTensor &normResUb, const LocalTensor &vec1ResUb, + const LocalTensor &normWeightUb, const LocalTensor &tempLocal, + uint32_t dealRowCount); + __aicore__ inline void CalRope(const LocalTensor &outputUb, const LocalTensor &normResUb, + const Vec2SliceInfo &originSliceInfo, uint32_t dealRowCount); + __aicore__ inline void CopyFinalResultOut(const LocalTensor &cmpKvOutUb, + CompressorVec2SliceIterator &sliceIterator); + __aicore__ inline void DealVec2BaseBlock(const Vec2RunInfo &info, const Vec2SplitInfo &splitInfo, + CompressorVec2SliceIterator &sliceIterator); + __aicore__ inline void SaveState(const LocalTensor &srcLocal, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const Vec1SliceInfo &sliceInfo, + uint32_t dStartIdx, uint32_t dDealSize, uint32_t stateIdx); + template + __aicore__ inline void DuplicateFirstBlock(const LocalTensor &dstLocal, uint32_t duplicateRowCount, + uint32_t duplicateColCount, uint32_t singleRowCount); + template + __aicore__ inline void ReadState(const LocalTensor &srcLocal, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const Vec1SliceInfo &sliceInfo, + uint32_t dStartIdx, uint32_t dDealSize, uint32_t stateIdx); + uint32_t cmpRatio_ = 0U; + uint32_t coff_ = 0U; + uint32_t compressedCnt_ = 0; + uint32_t prevApeDStartIdx_ = 0; + uint32_t prevApeDDealSize_ = 0; + bool apeIsLoad_ = false; + bool isExistSeqUsed_ = false; + bool isExistStartPos_ = false; + CompressorTools tools_; + ConstInfo constInfo_ = {}; + MSplitInfo mSplitInfo = {}; + GlobalTensor startPosGm_; + GlobalTensor cuSeqlensGm_; + GlobalTensor sequsedGm_; + GlobalTensor stateBlockTableGm_; + GlobalTensor stateCacheGm_; + GlobalTensor apeGm_; + GlobalTensor normWeightGm_; + GlobalTensor ropeSinGm_; + GlobalTensor ropeCosGm_; + GlobalTensor cmpKvOutGm_; + + // ================================Local Buffer区==================================== + // TBuf mm1ResUb; + LocalTensor normWeightUb; + LocalTensor apeUb; + // 临时tbuf + TBuf tmpBuff1; + TBuf tmpBuff2; + TBuf apeBuf; + TBuf normWeightBuf; + // in queue + TQue inputQue1; + TQue inputQue2; + TQue inputQue3; + // out queue + TQue outputQue1; + TQue outputQue2; +}; + + +template +__aicore__ inline void CompressorBlockVector::InitParams(const ConstInfo &constInfo, + const CompressorTools &tools) +{ + this->constInfo_ = constInfo; + this->tools_ = tools; + coff_ = static_cast(COMP::coff); + cmpRatio_ = constInfo.cmpRatio; +} + +template +__aicore__ inline void CompressorBlockVector::Init(__gm__ uint8_t *x, __gm__ uint8_t *wKv, __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, __gm__ uint8_t *cmpKvOut) +{ + stateBlockTableGm_.SetGlobalBuffer((__gm__ int32_t *)stateBlockTable); + stateCacheGm_.SetGlobalBuffer((__gm__ T *)stateCache); + apeGm_.SetGlobalBuffer((__gm__ T *)ape); + normWeightGm_.SetGlobalBuffer((__gm__ T *)normWeight); + ropeSinGm_.SetGlobalBuffer((__gm__ T *)ropeSin); + ropeCosGm_.SetGlobalBuffer((__gm__ T *)ropeCos); + cmpKvOutGm_.SetGlobalBuffer((__gm__ X_T *)cmpKvOut); + isExistSeqUsed_ = (seqUsed != nullptr); + isExistStartPos_ = (startPos != nullptr); + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens); + } + if (isExistSeqUsed_) { + sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed); + } + if (isExistStartPos_) { + startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos); + } +} + +template +__aicore__ inline void CompressorBlockVector::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(inputQue1, 1, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(inputQue2, 1, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(inputQue3, 1, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(tmpBuff1, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(tmpBuff2, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(outputQue1, 1, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(outputQue2, 1, BUFFER_SIZE_BYTE_16K); + pipe->InitBuffer(normWeightBuf, BUFFER_SIZE_BYTE_4K); + pipe->InitBuffer(apeBuf, BUFFER_SIZE_BYTE_32K); + normWeightUb = normWeightBuf.Get(); + apeUb = apeBuf.Get(); + LocalTensor normweightInUb = inputQue2.AllocTensor(); + DataCopy(normweightInUb, normWeightGm_, constInfo_.headDim); // 获取normWeight，常驻 + inputQue2.EnQue(normweightInUb); + inputQue2.DeQue(); + DataCopy(normWeightUb, normweightInUb, constInfo_.headDim); + inputQue2.FreeTensor(normweightInUb); + PipeBarrier(); +} + +template +__aicore__ inline void CompressorBlockVector::AllocEventID() +{ +} + +template +__aicore__ inline void CompressorBlockVector::FreeEventID() +{ +} + +template +__aicore__ inline void +CompressorBlockVector::InitVec1GlobalTensor(GlobalTensor kvMm1ResGm, GlobalTensor scoreMm1ResGm, + GlobalTensor kvCacheTcGm, GlobalTensor scoreCacheTcGm, + GlobalTensor vec1ResGm, GlobalTensor vec2InputGm) +{ + this->kvMm1ResGm_ = kvMm1ResGm; + this->scoreMm1ResGm_ = scoreMm1ResGm; + this->kvCacheTcGm_ = kvCacheTcGm; + this->scoreCacheTcGm_ = scoreCacheTcGm; + this->vec1ResGm_ = vec1ResGm; + this->vec2InputGm_ = vec2InputGm; +} + +template +__aicore__ inline uint32_t CompressorBlockVector::GetSeqUsed(uint32_t bIdx) +{ + if (isExistSeqUsed_) { + return (uint32_t)sequsedGm_.GetValue(bIdx); + } else { + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + return (uint32_t)(cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx)); + } else { + return constInfo_.sSize; + } + } +} + +template +__aicore__ inline uint32_t CompressorBlockVector::GetStartPos(uint32_t bIdx) +{ + if (isExistStartPos_) { + return startPosGm_.GetValue(bIdx); + } + return 0; +} + +template +__aicore__ inline uint32_t CompressorBlockVector::GetSeqLength(uint32_t bIdx) +{ + if (COMP::xLayout == X_LAYOUT::TH) { + return cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx); + } else { + return constInfo_.sSize; + } +} + +template +__aicore__ inline void CompressorBlockVector::CopyInApe(const LocalTensor &apeUb, uint32_t dStartIdx, + uint32_t dDealSize) +{ + if (apeIsLoad_ && prevApeDStartIdx_ == dStartIdx && prevApeDDealSize_ == dDealSize) { + return; + } + + uint32_t copyRowCount = coff_ * cmpRatio_; + uint32_t copyColCount = dDealSize; + uint32_t dstSingleRowCount = dDealSize; + uint32_t srcSingleRowCount = constInfo_.headDim; + + uint64_t gmOffset = dStartIdx; + + DataCopyWithInputQue(apeUb, apeGm_[gmOffset], copyRowCount, copyColCount, srcSingleRowCount, dstSingleRowCount); + + prevApeDStartIdx_ = dStartIdx; + prevApeDDealSize_ = dDealSize; + apeIsLoad_ = true; +} + +template +__aicore__ inline void CompressorBlockVector::AddApeToScore(const LocalTensor &scoreLocal, + const LocalTensor &apeUb, + const Vec1SliceInfo &sliceInfo, uint32_t dDealSize) +{ + uint32_t singleRowElemNum = dDealSize * coff_; + uint64_t scoreOffset = sliceInfo.dealedSeqCnt * singleRowElemNum; + + uint32_t tcDealSize = sliceInfo.dealTcSize; + if (sliceInfo.headHolderSeqCnt > 0) { + uint64_t apeOffset = sliceInfo.headHolderSeqCnt * singleRowElemNum; + uint32_t row = tcDealSize == 1 ? sliceInfo.validSeqCnt : (cmpRatio_ - sliceInfo.headHolderSeqCnt); + AddVF(scoreLocal[scoreOffset], apeUb[apeOffset], coff_ * row, dDealSize, dDealSize); + scoreOffset += row * singleRowElemNum; + tcDealSize -= 1; + } + if (tcDealSize == 0) { + return; + } + if (sliceInfo.tailHolderSeqCnt > 0) { + tcDealSize -= 1; + uint64_t apeOffset = 0; + uint32_t row = cmpRatio_ - sliceInfo.tailHolderSeqCnt; + uint32_t tailScoreOffset = scoreOffset + tcDealSize * cmpRatio_ * singleRowElemNum; + AddVF(scoreLocal[tailScoreOffset], apeUb[apeOffset], coff_ * row, dDealSize, dDealSize); + } + if (tcDealSize == 0) { + return; + } + uint32_t row = cmpRatio_; + for (uint32_t r = 0; r < tcDealSize; r++) { + AddVF(scoreLocal[scoreOffset + r * row * singleRowElemNum], apeUb, coff_ * row, dDealSize, dDealSize); + } +} + +template +__aicore__ inline void +CompressorBlockVector::AddSingleApeToScore(const LocalTensor &scoreLocal, const LocalTensor &apeUb, + const Vec1SliceInfo &sliceInfo, uint32_t dDealSize) +{ + uint32_t singleRowElemNum = dDealSize * coff_; + uint32_t dealRowCount = min(sliceInfo.sIdx, cmpRatio_); + uint64_t scoreOffset = (cmpRatio_ - dealRowCount) * singleRowElemNum; + uint64_t apeOffset = (cmpRatio_ - dealRowCount) * singleRowElemNum; + AddVF(scoreLocal[scoreOffset], apeUb[apeOffset], dealRowCount, dDealSize, singleRowElemNum); +} + +template +template +__aicore__ inline void +CompressorBlockVector::DataCopyAlignUbToUb(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + DataCopyParams intriParams; + intriParams.blockCount = copyRowCount; + intriParams.blockLen = copyColCount / FP32_BLOCK_ELEMENT_NUM; + intriParams.dstGap = (dstSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + intriParams.srcGap = (srcSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + DataCopy(dstLocal, srcLocal, intriParams); +} + +template +template +__aicore__ inline void +CompressorBlockVector::DataCopyAlignGmToUb(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + DataCopyParams intriParams; + intriParams.blockCount = copyRowCount; + intriParams.blockLen = copyColCount / FP32_BLOCK_ELEMENT_NUM; + intriParams.dstGap = (dstSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + intriParams.srcGap = (srcSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + DataCopy(dstLocal, srcGm, intriParams); +} + +template +template +__aicore__ inline void +CompressorBlockVector::DataCopyAlignUbToGm(const GlobalTensor &dstGm, const LocalTensor &srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + DataCopyParams intriParams; + intriParams.blockCount = copyRowCount; + intriParams.blockLen = copyColCount / FP32_BLOCK_ELEMENT_NUM; + intriParams.dstGap = (dstSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + intriParams.srcGap = (srcSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + DataCopy(dstGm, srcLocal, intriParams); +} + +template +template +__aicore__ inline void +CompressorBlockVector::DataCopyWithOutputQue(const GlobalTensor &dstGm, const LocalTensor &srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + uint32_t singleCopyRowCount = BUFFER_SIZE_BYTE_32K / (copyColCount * sizeof(O)); + for (uint32_t rowCount = 0; rowCount < copyRowCount; rowCount += singleCopyRowCount) { + uint64_t srcOffset = rowCount * srcSingleRowCount; + uint64_t dstOffset = rowCount * dstSingleRowCount; + uint32_t curCopyRowCount = min(singleCopyRowCount, copyRowCount - rowCount); + + LocalTensor outputUb = outputQue1.AllocTensor(); + + DataCopyAlignUbToUb(outputUb, srcLocal[srcOffset], curCopyRowCount, copyColCount, srcSingleRowCount, + copyColCount); + + outputQue1.EnQue(outputUb); + outputQue1.DeQue(); + + DataCopyAlignUbToGm(dstGm[dstOffset], outputUb, curCopyRowCount, copyColCount, copyColCount, dstSingleRowCount); + + outputQue1.FreeTensor(outputUb); + } +} + +template +template +__aicore__ inline void +CompressorBlockVector::DataCopyWithInputQue(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + uint32_t singleCopyRowCount = BUFFER_SIZE_BYTE_32K / (copyColCount * sizeof(O)); + for (uint32_t rowCount = 0; rowCount < copyRowCount; rowCount += singleCopyRowCount) { + uint64_t srcOffset = rowCount * srcSingleRowCount; + uint64_t dstOffset = rowCount * dstSingleRowCount; + uint32_t curCopyRowCount = min(singleCopyRowCount, copyRowCount - rowCount); + + LocalTensor inputUb = inputQue2.AllocTensor(); + + DataCopyAlignGmToUb(inputUb, srcGm[srcOffset], curCopyRowCount, copyColCount, srcSingleRowCount, copyColCount); + + inputQue2.EnQue(inputUb); + inputQue2.DeQue(); + + DataCopyAlignUbToUb(dstLocal[dstOffset], inputUb, curCopyRowCount, copyColCount, copyColCount, + dstSingleRowCount); + + inputQue2.FreeTensor(inputUb); + } +} + +template +template +__aicore__ inline void +CompressorBlockVector::AddMultiDataToUb(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + uint32_t dealRowCount, uint32_t dealColCount, uint32_t srcSingleRowCount, + uint32_t dstSingleRowCount, uint32_t repeatTimes, uint64_t offset) +{ + uint32_t cnt = dealRowCount * dstSingleRowCount; + uint32_t groupSize = BUFFER_SIZE_BYTE_32K / (cnt * sizeof(O)); + uint32_t loopTimes = CeilDivT(repeatTimes, groupSize); + uint64_t srcGmOffset = 0; + for (uint32_t idx = 0; idx < loopTimes; idx++) { + auto &inputQue = idx % 2 == 0 ? inputQue2 : inputQue3; + uint32_t curGroupSize = min(groupSize, (repeatTimes - groupSize * idx)); + LocalTensor splitLocal = inputQue.AllocTensor(); + if (srcSingleRowCount == dstSingleRowCount && dstSingleRowCount == dealRowCount) { + for (uint32_t groupIdx = 0; groupIdx < curGroupSize; groupIdx++) { + DataCopy(splitLocal[groupIdx * cnt], srcGm[srcGmOffset], cnt); + srcGmOffset += offset; + } + } else { + for (uint32_t groupIdx = 0; groupIdx < curGroupSize; groupIdx++) { + DataCopyAlignGmToUb(splitLocal[groupIdx * cnt], srcGm[srcGmOffset], dealRowCount, dealColCount, + srcSingleRowCount, dstSingleRowCount); + srcGmOffset += offset; + } + } + + inputQue.EnQue(splitLocal); + inputQue.DeQue(); + + PipeBarrier(); + if (idx == 0) { + MultiAddVF(dstLocal, splitLocal, dealRowCount, dealColCount, dstSingleRowCount, curGroupSize, cnt); + } else { + MultiAddVF(dstLocal, splitLocal, dealRowCount, dealColCount, dstSingleRowCount, curGroupSize, cnt); + } + inputQue.FreeTensor(splitLocal); + } + PipeBarrier(); +} + +template +__aicore__ inline void +CompressorBlockVector::PadAlign(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + const Vec1SliceInfo &sliceInfo, uint32_t dStartIdx, uint32_t dDealSize) +{ + // Ub data layout after overlap when r = 4 and coff = 2: + // Tc0_seq01: |--- --D_L--- -|------D_R-----| + // Tc0_seq02: |--- --D_L--- -|------D_R-----| + // Tc0_seq03: |--- --D_L--- -|------D_R-----| + // Tc0_seq04: |--- --D_L--- -|------D_R-----| + // Tc1_seq01: |--- --D_L--- -|------D_R-----| + // Tc1_seq02: |--- --D_L--- -|------D_R-----| + // Tc1_seq03: |--- --D_L--- -|------D_R-----| + // Tc1_seq04: |--- --D_L--- -|------D_R-----| + uint32_t srcSingleRowElemNum = dDealSize * coff_; + uint32_t copyRowCount = sliceInfo.compressTcSize * cmpRatio_ - sliceInfo.headHolderSeqCnt; + uint32_t copyColCount = dDealSize; + uint32_t srcSingleRowCount = srcSingleRowElemNum; + uint32_t dstSingleRowCount = srcSingleRowElemNum; // left和right在seq方向是交错存储的 + uint64_t srcLocalOffset = sliceInfo.dealedSeqCnt * srcSingleRowElemNum; + + uint64_t dstUbOffset = sliceInfo.compressoredScCnt * cmpRatio_ * dstSingleRowCount; + if constexpr (COMP::coff == COFF::OVERLAP) { + // 左侧 + uint64_t preSrcLocalOffset = srcLocalOffset; + uint64_t preDstUbOffset = dstUbOffset + (sliceInfo.headHolderSeqCnt + cmpRatio_) * dstSingleRowCount; + DataCopyAlignUbToUb(dstLocal[preDstUbOffset], srcLocal[preSrcLocalOffset], + copyRowCount - min(copyRowCount, cmpRatio_), copyColCount, srcSingleRowCount, + dstSingleRowCount); + dstUbOffset += dDealSize; + srcLocalOffset += dDealSize; + } + // 右侧 + dstUbOffset += sliceInfo.headHolderSeqCnt * dstSingleRowCount; + DataCopyAlignUbToUb(dstLocal[dstUbOffset], srcLocal[srcLocalOffset], copyRowCount, copyColCount, srcSingleRowCount, + dstSingleRowCount); +} + + +template +template +__aicore__ inline void +CompressorBlockVector::OverLap(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + const GlobalTensor &srcGm, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const GlobalTensor &cacheTcGm, + const Vec1RunInfo &info, const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, + uint32_t dStartIdx, uint32_t globalSeqIdx, uint32_t dDealSize) +{ + if (sliceInfo.dealTcSize == 0) { + return; + } + + if constexpr (IS_SCORE) { + AddApeToScore(srcLocal, apeUb, sliceInfo, dDealSize); + PipeBarrier(); + } + SaveState(srcLocal, stateGm, blockTableGm, sliceInfo, dStartIdx, dDealSize, static_cast(IS_SCORE)); + ReadState(dstLocal, stateGm, blockTableGm, sliceInfo, dStartIdx, dDealSize, + static_cast(IS_SCORE)); + + if constexpr (COMP::coff == COFF::OVERLAP) { + uint32_t nextC1V1DbIdx = (info.c1v1DbIdx + 1) % constInfo_.dbWorkspaceRatio; + GlobalTensor nextCacheTcGm = cacheTcGm[nextC1V1DbIdx * cmpRatio_ * constInfo_.headDim]; + SaveToWorkSpace(srcLocal, nextCacheTcGm, sliceInfo, loopInfo, dStartIdx, dDealSize); + } + if (sliceInfo.compressTcSize > 0) { + PadAlign(dstLocal, srcLocal, sliceInfo, dStartIdx, dDealSize); + if constexpr (COMP::coff == COFF::OVERLAP) { + GlobalTensor curCacheTcGm = cacheTcGm[info.c1v1DbIdx * cmpRatio_ * constInfo_.headDim]; + LoadFromWorkSpace(dstLocal, curCacheTcGm, srcGm, srcLocal, sliceInfo, loopInfo, dStartIdx, globalSeqIdx, + dDealSize); + } + } +} + +template +__aicore__ inline void +CompressorBlockVector::FromWokrSpaceToUb(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + const Vec1SliceInfo &sliceInfo, const StatisticInfo &statisticInfo, + uint32_t dStartIdx, uint32_t dDealSize) +{ + uint32_t srcSingleRowElemNum = constInfo_.headDim; + uint32_t copyRowCount = statisticInfo.dealSeqCnt * coff_; + uint32_t copyColCount = dDealSize; + uint32_t srcSingleRowCount = srcSingleRowElemNum; + uint32_t dstSingleRowCount = dDealSize; + uint64_t srcGmOffset = sliceInfo.dealedSeqCnt * srcSingleRowElemNum * coff_ + dStartIdx; + if (constInfo_.kBaseNum == 1) { + DataCopyAlignGmToUb(dstLocal, srcGm[srcGmOffset], copyRowCount, copyColCount, srcSingleRowCount, + dstSingleRowCount); + } else { + AddMultiDataToUb(dstLocal, srcGm[srcGmOffset], copyRowCount, copyColCount, srcSingleRowCount, dstSingleRowCount, + constInfo_.kBaseNum, constInfo_.mm1KvResSize); + } +} + +template +__aicore__ inline void +CompressorBlockVector::SaveToWorkSpace(const LocalTensor &srcLocal, const GlobalTensor &cacheTcGm, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, + uint32_t dStartIdx, uint32_t dDealSize) +{ + uint32_t curSeqLen = sliceInfo.bStartPos + sliceInfo.sIdx + sliceInfo.validSeqCnt; + uint32_t totalSeqLen = sliceInfo.bStartPos + sliceInfo.sIdx + sliceInfo.bSeqUsed; + if (!loopInfo.isCoreRowLast || !loopInfo.isCoreLoopLast || !sliceInfo.isLast || totalSeqLen < cmpRatio_ || + curSeqLen > Trunc(totalSeqLen, cmpRatio_) - cmpRatio_) { + return; + } + uint32_t srcSingleRowElemNum = dDealSize * coff_; + uint64_t srcLocalOffset = + (sliceInfo.dealedSeqCnt + sliceInfo.validSeqCnt - min(sliceInfo.validSeqCnt, cmpRatio_)) * srcSingleRowElemNum; + DataCopyWithOutputQue(cacheTcGm[dStartIdx], srcLocal[srcLocalOffset], + curSeqLen - max(curSeqLen - cmpRatio_, sliceInfo.bStartPos), dDealSize, coff_ * dDealSize, + constInfo_.headDim); +} + +template +__aicore__ inline void +CompressorBlockVector::LoadFromWorkSpace(const LocalTensor &dstLocal, const GlobalTensor &cacheTcGm, + const GlobalTensor &srcGm, const LocalTensor &srcLocal, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, + uint32_t dStartIdx, uint32_t globalSeqIdx, uint32_t dDealSize) +{ + if (sliceInfo.sIdx == 0) { + return; + } + uint32_t dstSingleRowElemNum = dDealSize * coff_; + uint32_t copyRowCount = min(sliceInfo.sIdx, cmpRatio_); + uint64_t dstLocalOffset = + (sliceInfo.compressoredScCnt * cmpRatio_ + cmpRatio_ - copyRowCount) * dstSingleRowElemNum; + if (loopInfo.isCoreRowFirst && loopInfo.isCoreLoopFirst && sliceInfo.isFirst) { // 从cacheGm获取 + uint32_t srcSingleRowElemNum = constInfo_.headDim; + uint64_t srcLocalOffset = dStartIdx; + + DataCopyWithInputQue(dstLocal[dstLocalOffset], cacheTcGm[srcLocalOffset], copyRowCount, dDealSize, + srcSingleRowElemNum, coff_ * dDealSize); + } else if (sliceInfo.isFirst) { // 从存放MatMul结果的WorkSpace中获取 + uint32_t srcSingleRowElemNum = constInfo_.headDim * coff_; + uint64_t srcLocalOffset = + (globalSeqIdx + sliceInfo.dealedSeqCnt - copyRowCount) * srcSingleRowElemNum + dStartIdx; + + if (constInfo_.kBaseNum == 1) { + DataCopyWithInputQue(dstLocal[dstLocalOffset], srcGm[srcLocalOffset], copyRowCount, dDealSize, + srcSingleRowElemNum, coff_ * dDealSize); + } else { + AddMultiDataToUb(dstLocal[dstLocalOffset], srcGm[srcLocalOffset], copyRowCount, dDealSize, + srcSingleRowElemNum, coff_ * dDealSize, constInfo_.kBaseNum, constInfo_.mm1KvResSize); + } + } else { // 从UB中获取 + uint32_t srcSingleRowElemNum = dDealSize * coff_; + uint64_t srcLocalOffset = (sliceInfo.dealedSeqCnt - copyRowCount) * srcSingleRowElemNum; + DataCopyAlignUbToUb(dstLocal[dstLocalOffset], srcLocal[srcLocalOffset], copyRowCount, dDealSize, + srcSingleRowElemNum, coff_ * dDealSize); + } +} + +template +__aicore__ inline void +CompressorBlockVector::ReadFromCacheState(const LocalTensor &output, const GlobalTensor &state, + const GlobalTensor &blockTableGm, uint32_t batchIdx, + uint32_t startSeqIdx, uint32_t endSeqIdx, uint32_t dStartIdx, + uint32_t dDealSize, uint32_t stateIdx) +{ + if constexpr (COMP::cacheMode == CACHE_MODE::CONTINUOUS) { + uint64_t blockTablebaseOffset = batchIdx * constInfo_.maxBlockNumPerBatch; + uint32_t curSeqIdx = startSeqIdx; + uint32_t copyFinishRowCnt = 0; + uint32_t seqCnt = endSeqIdx - startSeqIdx; + while (copyFinishRowCnt < seqCnt) { + uint64_t blockIdOffset = curSeqIdx / constInfo_.blockSize; + uint64_t remainRowCnt = curSeqIdx % constInfo_.blockSize; + uint64_t idInBlockTable = blockTableGm.GetValue(blockTablebaseOffset + blockIdOffset); + uint32_t copyRowCount = constInfo_.blockSize - remainRowCnt; + if (copyFinishRowCnt + copyRowCount > seqCnt) { + copyRowCount = seqCnt - copyFinishRowCnt; + } + uint64_t stateOffset = idInBlockTable * constInfo_.stateCacheStrideDim0 + + remainRowCnt * 2 * coff_ * constInfo_.headDim + + stateIdx * coff_ * constInfo_.headDim + dStartIdx; + + DataCopyWithInputQue(output[copyFinishRowCnt * coff_ * dDealSize], state[stateOffset], copyRowCount, + dDealSize, coff_ * constInfo_.headDim * 2, coff_ * dDealSize); + copyFinishRowCnt += copyRowCount; + curSeqIdx += copyRowCount; + } + } else { + uint32_t curSeqIdx = startSeqIdx; + uint32_t copyFinishRowCnt = 0; + uint32_t seqCnt = endSeqIdx - startSeqIdx; + uint64_t idInBlockTable = blockTableGm.GetValue(batchIdx); + while (copyFinishRowCnt < seqCnt) { + uint64_t remainRowCnt = curSeqIdx % constInfo_.blockSize; + uint32_t copyRowCount = constInfo_.blockSize - remainRowCnt; + if (copyFinishRowCnt + copyRowCount > seqCnt) { + copyRowCount = seqCnt - copyFinishRowCnt; + } + uint64_t stateOffset = idInBlockTable * constInfo_.stateCacheStrideDim0 + + remainRowCnt * 2 * coff_ * constInfo_.headDim + + stateIdx * coff_ * constInfo_.headDim + dStartIdx; + + DataCopyWithInputQue(output[copyFinishRowCnt * coff_ * dDealSize], state[stateOffset], copyRowCount, + dDealSize, coff_ * constInfo_.headDim * 2, coff_ * dDealSize); + copyFinishRowCnt += copyRowCount; + curSeqIdx += copyRowCount; + } + } +} + +template +__aicore__ inline void +CompressorBlockVector::WriteToCacheState(const GlobalTensor &state, const GlobalTensor &blockTableGm, + const LocalTensor &input, uint32_t batchIdx, uint32_t startSeqIdx, + uint32_t endSeqIdx, uint32_t dStartIdx, uint32_t dDealSize, + uint32_t stateIdx) +{ + if constexpr (COMP::cacheMode == CACHE_MODE::CONTINUOUS) { + uint64_t blockTablebaseOffset = batchIdx * constInfo_.maxBlockNumPerBatch; + uint32_t curSeqIdx = startSeqIdx; + uint32_t copyFinishRowCnt = 0; + uint32_t seqCnt = endSeqIdx - startSeqIdx; + while (copyFinishRowCnt < seqCnt) { + uint64_t blockIdOffset = curSeqIdx / constInfo_.blockSize; + uint64_t remainRowCnt = curSeqIdx % constInfo_.blockSize; + uint64_t idInBlockTable = blockTableGm.GetValue(blockTablebaseOffset + blockIdOffset); + uint32_t copyRowCount = constInfo_.blockSize - remainRowCnt; + if (copyFinishRowCnt + copyRowCount > seqCnt) { + copyRowCount = seqCnt - copyFinishRowCnt; + } + if (idInBlockTable != 0) { // 32 + uint64_t stateOffset = idInBlockTable * constInfo_.stateCacheStrideDim0 + + remainRowCnt * 2 * coff_ * constInfo_.headDim + + stateIdx * coff_ * constInfo_.headDim + dStartIdx; + DataCopyWithOutputQue(state[stateOffset], input[copyFinishRowCnt * coff_ * dDealSize], copyRowCount, + dDealSize, coff_ * dDealSize, coff_ * constInfo_.headDim * 2); + } + + copyFinishRowCnt += copyRowCount; + curSeqIdx += copyRowCount; + } + } else { + uint32_t curSeqIdx = startSeqIdx; + uint32_t copyFinishRowCnt = 0; + uint32_t seqCnt = endSeqIdx - startSeqIdx; + uint64_t idInBlockTable = blockTableGm.GetValue(batchIdx); + while (copyFinishRowCnt < seqCnt) { + uint64_t remainRowCnt = curSeqIdx % constInfo_.blockSize; + uint32_t copyRowCount = constInfo_.blockSize - remainRowCnt; + if (copyFinishRowCnt + copyRowCount > seqCnt) { + copyRowCount = seqCnt - copyFinishRowCnt; + } + uint64_t stateOffset = idInBlockTable * constInfo_.stateCacheStrideDim0 + + remainRowCnt * 2 * coff_ * constInfo_.headDim + + stateIdx * coff_ * constInfo_.headDim + dStartIdx; + DataCopyWithOutputQue(state[stateOffset], input[copyFinishRowCnt * coff_ * dDealSize], copyRowCount, + dDealSize, coff_ * dDealSize, coff_ * constInfo_.headDim * 2); + + copyFinishRowCnt += copyRowCount; + curSeqIdx += copyRowCount; + } + } +} + +template +__aicore__ inline void +CompressorBlockVector::SaveState(const LocalTensor &srcLocal, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const Vec1SliceInfo &sliceInfo, + uint32_t dStartIdx, uint32_t dDealSize, uint32_t stateIdx) +{ + uint32_t startSeqIdx = sliceInfo.bStartPos + sliceInfo.sIdx; + uint32_t endSeqIdx = startSeqIdx + sliceInfo.validSeqCnt; + uint64_t srcBaseOffset = sliceInfo.dealedSeqCnt * coff_ * dDealSize; + + if constexpr (COMP::cacheMode == CACHE_MODE::CYCLE) { + uint32_t compressSeqIdx = Trunc(sliceInfo.bStartPos + sliceInfo.bSeqUsed, cmpRatio_); + uint32_t writeSeqStartIdx = compressSeqIdx > (coff_ - 1) * cmpRatio_ ? + compressSeqIdx - (coff_ - 1) * cmpRatio_ : 0; + if (endSeqIdx <= writeSeqStartIdx) { + return; + } + srcBaseOffset += (max(startSeqIdx, writeSeqStartIdx) - startSeqIdx) * coff_ * dDealSize; + startSeqIdx = max(startSeqIdx, writeSeqStartIdx); + } + + if constexpr (COMP::coff == COFF::OVERLAP) { + WriteToCacheState(stateGm, blockTableGm, srcLocal[srcBaseOffset], sliceInfo.bIdx, startSeqIdx, endSeqIdx, + dStartIdx, dDealSize, stateIdx); + srcBaseOffset += dDealSize; + dStartIdx += constInfo_.headDim; + } + + WriteToCacheState(stateGm, blockTableGm, srcLocal[srcBaseOffset], sliceInfo.bIdx, startSeqIdx, endSeqIdx, dStartIdx, + dDealSize, stateIdx); +} + +template +template +__aicore__ inline void +CompressorBlockVector::DuplicateFirstBlock(const LocalTensor &dstLocal, uint32_t duplicateRowCount, + uint32_t duplicateColCount, uint32_t singleRowCount) +{ + for (uint32_t offset = 0; offset < duplicateColCount; offset += FP32_REPEAT_ELEMENT_NUM) { + uint32_t curDuplicateColCount = min(duplicateColCount - offset, FP32_REPEAT_ELEMENT_NUM); + if constexpr (IS_SCORE) { + Duplicate(dstLocal[offset], SOFTMAX_MIN_NUM, curDuplicateColCount, duplicateRowCount, 1, + singleRowCount / REPEAT_STRIDE_NUM); + } else { + Duplicate(dstLocal[offset], FLOAT_ZERO, curDuplicateColCount, duplicateRowCount, 1, + singleRowCount / REPEAT_STRIDE_NUM); + } + } +} + + +template +template +__aicore__ inline void +CompressorBlockVector::ReadState(const LocalTensor &dstLocal, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const Vec1SliceInfo &sliceInfo, + uint32_t dStartIdx, uint32_t dDealSize, uint32_t stateIdx) +{ + // 没有需要压缩的块时, 不需要读state的信息 + if (sliceInfo.compressTcSize == 0) { + return; + } + // 填充右边 + if (sliceInfo.headHolderSeqCnt > 0) { + // 整个batch的第一块 + uint32_t startSeqIdx = Trunc(sliceInfo.bStartPos + sliceInfo.sIdx, cmpRatio_); + uint32_t endSeqIdx = sliceInfo.bStartPos; + uint64_t dstBaseOffset = sliceInfo.compressoredScCnt * cmpRatio_ * coff_ * dDealSize; + if constexpr (COMP::coff == Compressor::COFF::OVERLAP) { + dstBaseOffset += (coff_ - 1) * dDealSize; + } + ReadFromCacheState(dstLocal[dstBaseOffset], stateGm, blockTableGm, sliceInfo.bIdx, startSeqIdx, endSeqIdx, + dStartIdx + (coff_ - 1) * constInfo_.headDim, dDealSize, stateIdx); + } + + // 填充左边 + if constexpr (COMP::coff == Compressor::COFF::OVERLAP) { + bool isFirst = sliceInfo.bStartPos + sliceInfo.sIdx < cmpRatio_; + if (isFirst) { + // 无历史数据 + // dDealSize必须为64 + uint64_t dstBaseOffset = sliceInfo.compressoredScCnt * cmpRatio_ * coff_ * dDealSize; + DuplicateFirstBlock(dstLocal[dstBaseOffset], cmpRatio_, dDealSize, coff_ * dDealSize); + } + if (sliceInfo.sIdx < cmpRatio_ && (!isFirst || sliceInfo.compressTcSize > 1)) { + uint32_t startSeqIdx = sliceInfo.bStartPos < cmpRatio_ ? + 0 : + Trunc(sliceInfo.bStartPos + sliceInfo.sIdx, cmpRatio_) - cmpRatio_; + uint32_t endSeqIdx = + min(Trunc(sliceInfo.bStartPos + sliceInfo.sIdx + sliceInfo.validSeqCnt, cmpRatio_) - cmpRatio_, + sliceInfo.bStartPos); + uint64_t dstBaseOffset = sliceInfo.compressoredScCnt * cmpRatio_ * coff_ * dDealSize; + if (isFirst) { + dstBaseOffset += cmpRatio_ * coff_ * dDealSize; + } + ReadFromCacheState(dstLocal[dstBaseOffset], stateGm, blockTableGm, sliceInfo.bIdx, startSeqIdx, endSeqIdx, + dStartIdx, dDealSize, stateIdx); + } + } +} + +template +__aicore__ inline void CompressorBlockVector::SoftmaxDN(const LocalTensor &scoreLocal, uint32_t tcDealSize, + uint32_t dDealSize) +{ + float minValue = -2e38; + uint32_t ReduceSize = coff_ * cmpRatio_; + FaVectorApi::SoftmaxDnVF(scoreLocal, scoreLocal, dDealSize, ReduceSize, tcDealSize, minValue, dDealSize); +} + +template +__aicore__ inline void +CompressorBlockVector::KvMulReduceScore(const LocalTensor &kvLocal, const LocalTensor &scoreLocal, + const LocalTensor &dstLocal, uint32_t tcDealSize, uint32_t dDealSize) +{ + MulReduceSumbaseVF(kvLocal, scoreLocal, dstLocal, coff_, cmpRatio_, dDealSize, tcDealSize); +} + +template +__aicore__ inline void +CompressorBlockVector::CopyOutVec1Res(const GlobalTensor &resGm, const Vec1RunInfo &info, + const LocalTensor &comperssoredUb, uint32_t compressTcSize, + uint32_t dStartIdx, uint32_t dDealSize) +{ + uint64_t outGmOffset = compressedCnt_ * constInfo_.headDim + dStartIdx; + DataCopyAlignUbToGm(resGm[outGmOffset], comperssoredUb, compressTcSize, dDealSize, dDealSize, constInfo_.headDim); +} + +template +__aicore__ inline void CompressorBlockVector::OverLapScoreKv( + const LocalTensor &scoreLocal, const LocalTensor &kvLocal, const Vec1RunInfo &info, const LoopInfo &loopInfo, + const StatisticInfo &statisticInfo, const Vec1SliceInfo &originSliceInfo, uint32_t dStartIdx, uint32_t dDealSize, + uint32_t dBaseSize, uint32_t needDealTcSize) +{ + CompressorVec1SliceIterator overLapSliceIterator(tools_); + overLapSliceIterator.SetMaxBatchSize(constInfo_.batchSize); + Vec1SliceInfo &overLapSliceInfo = overLapSliceIterator.GetSlice(); + + GlobalTensor scoreDBMm1ResGm = scoreMm1ResGm_[info.c1v1DbIdx * constInfo_.dbSize]; + LocalTensor scoreUb = inputQue1.AllocTensor(); + FromWokrSpaceToUb(scoreUb, scoreDBMm1ResGm, originSliceInfo, statisticInfo, dStartIdx, dDealSize); + inputQue1.EnQue(scoreUb); + inputQue1.DeQue(); + overLapSliceIterator.Reset(originSliceInfo.bIdx, originSliceInfo.sIdx, 0U, 0U); + overLapSliceIterator.SetNeedDealTcSize(needDealTcSize); + while (!overLapSliceIterator.IsEnd()) { + overLapSliceIterator.GetSlice(); + OverLap(scoreLocal, scoreUb, scoreDBMm1ResGm, stateCacheGm_, stateBlockTableGm_, scoreCacheTcGm_, info, + overLapSliceInfo, loopInfo, dStartIdx, originSliceInfo.dealedSeqCnt, dDealSize); + overLapSliceIterator.IteratorSlice(); + } + inputQue1.FreeTensor(scoreUb); + + if constexpr (COMP::coff == COFF::OVERLAP) { + if (originSliceInfo.sIdx != 0 && originSliceInfo.compressTcSize > 0 && + (!loopInfo.isCoreRowFirst || !loopInfo.isCoreLoopFirst)) { + PipeBarrier(); + AddSingleApeToScore(scoreLocal, apeUb, originSliceInfo, dDealSize); + } + } + + GlobalTensor kvDBMm1ResGm = kvMm1ResGm_[info.c1v1DbIdx * constInfo_.dbSize]; + LocalTensor kvUb = inputQue1.AllocTensor(); + FromWokrSpaceToUb(kvUb, kvDBMm1ResGm, originSliceInfo, statisticInfo, dStartIdx, dDealSize); + inputQue1.EnQue(kvUb); + inputQue1.DeQue(); + overLapSliceIterator.Reset(originSliceInfo.bIdx, originSliceInfo.sIdx, 0U, 0U); + overLapSliceIterator.SetNeedDealTcSize(needDealTcSize); + + while (!overLapSliceIterator.IsEnd()) { + overLapSliceIterator.GetSlice(); + OverLap(kvLocal, kvUb, kvDBMm1ResGm, stateCacheGm_, stateBlockTableGm_, kvCacheTcGm_, info, + overLapSliceInfo, loopInfo, dStartIdx, originSliceInfo.dealedSeqCnt, dDealSize); + overLapSliceIterator.IteratorSlice(); + } + inputQue1.FreeTensor(kvUb); + + + PipeBarrier(); +} + +template +__aicore__ inline void CompressorBlockVector::DealVec1BaseBlock(const Vec1RunInfo &info, + CompressorVec1SliceIterator &sliceIterator, + const LoopInfo &loopInfo, uint32_t dStartIdx, + uint32_t dDealSize, uint32_t dBaseSize) +{ + Vec1SliceInfo originSliceInfo = sliceIterator.GetSlice(); + uint32_t needDealTcSize = sliceIterator.GetNeedDealTcSize(); + StatisticInfo &statisticInfo = sliceIterator.template FullIteratorSlice(); + if (statisticInfo.actualTcCnt == 0) { + return; + } + LocalTensor scoreLocal = tmpBuff1.Get(); + LocalTensor kvLocal = tmpBuff2.Get(); + + OverLapScoreKv(scoreLocal, kvLocal, info, loopInfo, statisticInfo, originSliceInfo, dStartIdx, dDealSize, dBaseSize, + needDealTcSize); + + if (statisticInfo.compressorScCnt > 0) { + SoftmaxDN(scoreLocal, statisticInfo.compressorScCnt, dDealSize); + LocalTensor comperssoredUb = outputQue2.AllocTensor(); + PipeBarrier(); + KvMulReduceScore(kvLocal, scoreLocal, comperssoredUb, statisticInfo.compressorScCnt, dDealSize); + PipeBarrier(); + outputQue2.EnQue(comperssoredUb); + outputQue2.DeQue(); + GlobalTensor resGm = vec1ResGm_[info.v1v2DbIdx * constInfo_.dbSize]; + CopyOutVec1Res(resGm, info, comperssoredUb, statisticInfo.compressorScCnt, dStartIdx, dDealSize); + outputQue2.FreeTensor(comperssoredUb); + } + compressedCnt_ += statisticInfo.compressorScCnt; +} + +template +__aicore__ inline void CompressorBlockVector::CalcGroupInfo(const Vec1RunInfo &info, Vec1SplitInfo &splitInfo) +{ + uint32_t aiCoreNum = constInfo_.usedCoreNum * 2; + splitInfo.dBaseSize = constInfo_.headDim / min(FloorPow2(aiCoreNum), CeilPow2(CeilDivT(aiCoreNum, info.dealTcNum))); + if (constInfo_.kBaseNum > 1) { + splitInfo.dBaseSize = max(splitInfo.dBaseSize, FP32_REPEAT_ELEMENT_NUM); + } + splitInfo.dBaseSize = max(splitInfo.dBaseSize, FP32_BLOCK_ELEMENT_NUM); + splitInfo.vec1GroupSize = constInfo_.headDim / splitInfo.dBaseSize; + splitInfo.vec1GroupNum = min(static_cast(aiCoreNum / splitInfo.vec1GroupSize), info.dealTcNum); +} + +template +__aicore__ inline void CompressorBlockVector::CalcTaskDistribution(const Vec1RunInfo &info, + Vec1SplitInfo &splitInfo) +{ + uint32_t blockIdx = GetBlockIdx(); + uint32_t groupSize = splitInfo.vec1GroupSize; + uint32_t groupNum = splitInfo.vec1GroupNum; + uint32_t dealTcNum = info.dealTcNum; + + if (blockIdx < groupSize * (dealTcNum % groupNum)) { + splitInfo.dealTcSize = dealTcNum / groupNum + 1; + splitInfo.preDealTcSize = splitInfo.dealTcSize * (blockIdx / groupSize); + } else if (blockIdx < groupSize * groupNum) { + splitInfo.dealTcSize = dealTcNum / groupNum; + splitInfo.preDealTcSize = splitInfo.dealTcSize * (blockIdx / groupSize) + dealTcNum % groupNum; + } else { + splitInfo.dealTcSize = 0; + splitInfo.preDealTcSize = dealTcNum; + } +} + +template +__aicore__ inline void CompressorBlockVector::UpdateIteratorState(const Vec1RunInfo &info, + Vec1SplitInfo &splitInfo) +{ + CompressorVec1SliceIterator sliceIterator(tools_); + sliceIterator.SetMaxBatchSize(constInfo_.batchSize); + sliceIterator.Reset(info.bStart, info.sStart, 0U, 0U); + Vec1SliceInfo &sliceInfo = sliceIterator.GetSlice(); + + // 处理前序任务量，更新起始索引 + if (splitInfo.preDealTcSize > 0) { + sliceIterator.SetNeedDealTcSize(splitInfo.preDealTcSize); + StatisticInfo &statisticInfo = sliceIterator.template FullIteratorSlice(); + splitInfo.curCompressedCnt = statisticInfo.compressorScCnt; + splitInfo.dealSeqStartIdx = sliceInfo.dealedSeqCnt; + splitInfo.curBStart = sliceInfo.bIdx; + splitInfo.curSStart = sliceInfo.sIdx; + } else { + splitInfo.curCompressedCnt = 0; + splitInfo.dealSeqStartIdx = 0; + splitInfo.curBStart = info.bStart; + splitInfo.curSStart = info.sStart; + } + + // 处理当前核实际要跑的任务量 + sliceIterator.SetNeedDealTcSize(info.dealTcNum - splitInfo.preDealTcSize); + StatisticInfo &statisticInfo = sliceIterator.template FullIteratorSlice(); + splitInfo.totalCompressedCnt = splitInfo.curCompressedCnt + statisticInfo.compressorScCnt; +} + +template +__aicore__ inline void CompressorBlockVector::CalcTilingStrategy(Vec1SplitInfo &splitInfo) +{ + // 计算headDim和Tc方向切分大小 + uint32_t maxDealColNum = BUFFER_SIZE_BYTE_32K / (cmpRatio_ * coff_ * sizeof(T)); + + // 切块逻辑 + if (maxDealColNum < splitInfo.dBaseSize) { + splitInfo.tcSplitSize = 1; + splitInfo.dLoopCount = CeilDivT(splitInfo.dBaseSize, maxDealColNum); + splitInfo.dSplitSize = splitInfo.dBaseSize / splitInfo.dLoopCount; + } else { + splitInfo.dSplitSize = splitInfo.dBaseSize; + splitInfo.dLoopCount = splitInfo.dBaseSize / splitInfo.dSplitSize; // 此处常等于1，保留原逻辑 + splitInfo.tcSplitSize = maxDealColNum / splitInfo.dBaseSize; + } +} + +template +__aicore__ inline Vec1SplitInfo CompressorBlockVector::SplitCoreV1(const Vec1RunInfo &info) +{ + Vec1SplitInfo splitInfo; + + // 1. 计算基础分组和分片大小 + CalcGroupInfo(info, splitInfo); + + // 2. 根据当前的 BlockIdx 计算任务分配（负载均衡） + CalcTaskDistribution(info, splitInfo); + + // 3. 刷新迭代器并获取当前核的起始位置状态 + UpdateIteratorState(info, splitInfo); + + if (splitInfo.dealTcSize == 0) { + return splitInfo; + } + + // 4. 计算具体在内存中的切块（Tiling）逻辑 + CalcTilingStrategy(splitInfo); + + return splitInfo; +} + +template +__aicore__ inline void CompressorBlockVector::ComputeVec1(const Vec1RunInfo &info) +{ + if (info.dealTcNum == 0) { + return; + } + if (info.resetResFlag) { + compressedCnt_ = 0; + } + uint32_t preCompressedCnt = compressedCnt_; + Vec1SplitInfo splitInfo = SplitCoreV1(info); + // 计算当前VecCore的任务量 + if (splitInfo.dealTcSize == 0) { + compressedCnt_ += splitInfo.totalCompressedCnt; + return; + } + + LoopInfo loopInfo; + loopInfo.groupSize = splitInfo.vec1GroupSize; + loopInfo.groupNum = splitInfo.vec1GroupNum; + loopInfo.coreRowIdx = GetBlockIdx() / splitInfo.vec1GroupSize; + loopInfo.coreColIdx = GetBlockIdx() % splitInfo.vec1GroupSize; + loopInfo.isCoreRowLast = loopInfo.coreRowIdx == splitInfo.vec1GroupNum - 1; + loopInfo.isCoreRowFirst = loopInfo.coreRowIdx == 0; + + + CompressorVec1SliceIterator sliceIterator(tools_); + sliceIterator.SetMaxBatchSize(constInfo_.batchSize); + // 切块循环 + uint64_t baseOffset = loopInfo.coreColIdx * splitInfo.dBaseSize; + for (uint32_t dLoopIdx = 0; dLoopIdx < splitInfo.dLoopCount; dLoopIdx++) { + uint64_t dBaseOffset = baseOffset + dLoopIdx * splitInfo.dSplitSize; + + CopyInApe(apeUb, dBaseOffset, splitInfo.dSplitSize); + + sliceIterator.Reset(splitInfo.curBStart, splitInfo.curSStart, splitInfo.dealSeqStartIdx, 0U); + compressedCnt_ = preCompressedCnt + splitInfo.curCompressedCnt; + for (uint32_t tcIdx = 0; tcIdx < splitInfo.dealTcSize; tcIdx += splitInfo.tcSplitSize) { + uint32_t actDealTcSize = min(splitInfo.tcSplitSize, splitInfo.dealTcSize - tcIdx); + + loopInfo.isCoreLoopFirst = tcIdx == 0; + loopInfo.isCoreLoopLast = tcIdx + splitInfo.tcSplitSize >= splitInfo.dealTcSize; + // 处理单个切块 + sliceIterator.SetNeedDealTcSize(actDealTcSize); + sliceIterator.SetDealedTcCnt(0U); + DealVec1BaseBlock(info, sliceIterator, loopInfo, dBaseOffset, splitInfo.dSplitSize, splitInfo.dBaseSize); + } + } + compressedCnt_ = preCompressedCnt + splitInfo.totalCompressedCnt; +} +template +__aicore__ inline Vec2SplitInfo CompressorBlockVector::SplitCoreV2(const Vec2RunInfo &info) +{ + Vec2SplitInfo splitInfo; + + uint32_t blockIdx = GetBlockIdx(); + uint32_t aiCoreNum = constInfo_.usedCoreNum * 2; + + if (blockIdx < (info.dealScSize % aiCoreNum)) { + splitInfo.dealScNum = info.dealScSize / aiCoreNum + 1; + splitInfo.preScCnt = splitInfo.dealScNum * blockIdx; + } else if (blockIdx < aiCoreNum) { + splitInfo.dealScNum = info.dealScSize / aiCoreNum; + splitInfo.preScCnt = splitInfo.dealScNum * blockIdx + info.dealScSize % aiCoreNum; + } else { + splitInfo.dealScNum = 0; + splitInfo.preScCnt = info.dealScSize; + } + if (splitInfo.dealScNum == 0) { + return splitInfo; + } + + uint32_t preScCnt = splitInfo.preScCnt; + for (uint32_t curB = info.bStart, curSc = info.bCompressedId; curB < constInfo_.batchSize; curB++, curSc = 0) { + uint32_t startPos = GetStartPos(curB); + uint32_t seqUsed = GetSeqUsed(curB); + uint32_t curScNum = ((startPos + seqUsed) / cmpRatio_) - (startPos / cmpRatio_) - curSc; + if (preScCnt < curScNum) { + splitInfo.curBStart = curB; + splitInfo.curScStart = curSc + preScCnt; + break; + } + preScCnt -= curScNum; + } + + splitInfo.dealedScCnt = info.compressedId - info.dealScSize; + + return splitInfo; +} + +template +__aicore__ inline void CompressorBlockVector::ComputeVec2(const Vec2RunInfo &info) +{ + Vec2SplitInfo splitInfo = SplitCoreV2(info); + if (splitInfo.dealScNum == 0) { + return; + } + CompressorVec2SliceIterator sliceIterator(tools_); + sliceIterator.SetMaxBatchSize(constInfo_.batchSize); + sliceIterator.Reset(splitInfo.curBStart, splitInfo.curScStart, splitInfo.dealedScCnt + splitInfo.preScCnt); + + uint32_t singleLoopScNum = BUFFER_SIZE_BYTE_32K / (constInfo_.headDim * sizeof(T)); + uint32_t loopTimes = CeilDivT(splitInfo.dealScNum, singleLoopScNum); + + for (uint32_t idx = 0; idx < loopTimes; idx++) { + uint32_t curLoopScNum = min(singleLoopScNum, splitInfo.dealScNum - singleLoopScNum * idx); + sliceIterator.SetNeedDealScSize(curLoopScNum); + sliceIterator.ResetLoopDealedScCnt(); + DealVec2BaseBlock(info, splitInfo, sliceIterator); + splitInfo.preScCnt += curLoopScNum; + } +} + +template +__aicore__ inline void CompressorBlockVector::DealVec2BaseBlock(const Vec2RunInfo &info, + const Vec2SplitInfo &splitInfo, + CompressorVec2SliceIterator &sliceIterator) +{ + Vec2SliceInfo &sliceInfo = sliceIterator.GetSlice(); + uint32_t needDealScSize = sliceIterator.GetNeedDealScSize(); + uint32_t computeSize = needDealScSize * constInfo_.headDim; + int64_t inGmOffset = splitInfo.preScCnt * constInfo_.headDim; + GlobalTensor vec2InputGm = vec2InputGm_[info.v2DbIdx * constInfo_.dbSize]; + // CopyIn + LocalTensor vec1ResUb = inputQue1.AllocTensor(); + DataCopy(vec1ResUb, vec2InputGm[inGmOffset], computeSize); + inputQue1.EnQue(vec1ResUb); + inputQue1.DeQue(); + + // RmsNorm + LocalTensor normResUb = tmpBuff1.Get(); + LocalTensor tempLocal = tmpBuff2.Get(); + PipeBarrier(); + MultRowRmsNorm(normResUb, vec1ResUb, normWeightUb, tempLocal, needDealScSize); + inputQue1.FreeTensor(vec1ResUb); + + // rope: 只对后RD进行rope + LocalTensor outputUb = outputQue1.AllocTensor(); + PipeBarrier(); + CalRope(outputUb, normResUb, sliceInfo, needDealScSize); + PipeBarrier(); + + // CopyOut + outputQue1.EnQue(outputUb); + outputQue1.DeQue(); + CopyFinalResultOut(outputUb, sliceIterator); + outputQue1.FreeTensor(outputUb); +} + +template +__aicore__ inline void +CompressorBlockVector::MultRowRmsNorm(const LocalTensor &normResUb, const LocalTensor &vec1ResUb, + const LocalTensor &normWeightUb, const LocalTensor &tempLocal, + uint32_t dealRowCount) +{ + uint32_t row = 1; + uint32_t col = constInfo_.headDim; + float reciprocal = 1.0f / col; + float epsilon = constInfo_.normEps; + for (uint32_t i = 0; i < dealRowCount; ++i) { + RmsNormVF(normResUb[i * col], vec1ResUb[i * col], normWeightUb, reciprocal, epsilon, row, col); + } +} + + +template +__aicore__ inline void CompressorBlockVector::CalRope(const LocalTensor &outputUb, + const LocalTensor &normResUb, + const Vec2SliceInfo &originSliceInfo, uint32_t dealRowCount) +{ + CompressorVec2SliceIterator ropeSliceIterator(tools_); + ropeSliceIterator.SetMaxBatchSize(constInfo_.batchSize); + ropeSliceIterator.Reset(originSliceInfo.bIdx, originSliceInfo.scIdx, originSliceInfo.dealedScCnt); + Vec2SliceInfo &sliceInfo = ropeSliceIterator.GetSlice(); + + ropeSliceIterator.SetNeedDealScSize(dealRowCount); + while (!ropeSliceIterator.IsEnd()) { + ropeSliceIterator.GetSlice(); + if (sliceInfo.curDealScNum > 0) { + uint32_t computeSize = sliceInfo.curDealScNum * constInfo_.ropeHeadDim; + uint64_t SinCosOffset = sliceInfo.padScIdx * constInfo_.ropeHeadDim; + + // sin与cos各占一半, 实际分别最多只会用16K,总占用32K + LocalTensor cosUb = inputQue2.AllocTensor(); + LocalTensor sinUb = cosUb[BUFFER_SIZE_BYTE_16K / sizeof(T)]; + DataCopy(cosUb, ropeCosGm_[SinCosOffset], computeSize); + DataCopy(sinUb, ropeSinGm_[SinCosOffset], computeSize); + inputQue2.EnQue(sinUb); + inputQue2.DeQue(); + + RopeVF(sinUb, cosUb, normResUb[sliceInfo.loopDealedScCnt * constInfo_.headDim], + outputUb[sliceInfo.loopDealedScCnt * constInfo_.headDim], sliceInfo.curDealScNum, + constInfo_.ropeHeadDim, constInfo_.headDim, + constInfo_.headDim - constInfo_.ropeHeadDim); + inputQue2.FreeTensor(sinUb); + } + ropeSliceIterator.IteratorSlice(); + } +} + + +template +__aicore__ inline void CompressorBlockVector::CopyFinalResultOut(const LocalTensor &cmpKvOutUb, + CompressorVec2SliceIterator &sliceIterator) +{ + Vec2SliceInfo &sliceInfo = sliceIterator.GetSlice(); + while (!sliceIterator.IsEnd()) { + sliceIterator.GetSlice(); + if (sliceInfo.curDealScNum > 0) { + DataCopy(cmpKvOutGm_[sliceInfo.padScIdx * constInfo_.headDim], + cmpKvOutUb[sliceInfo.loopDealedScCnt * constInfo_.headDim], + sliceInfo.curDealScNum * constInfo_.headDim); + } + sliceIterator.IteratorSlice(); + } +} + +} // namespace Compressor +#endif // COMPRESSOR_BLOCK_VECTOR_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_block_vec_full_load.h b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_block_vec_full_load.h new file mode 100644 index 000000000..1927a62de --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_block_vec_full_load.h @@ -0,0 +1,1415 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_block_vec_full_load.h + * \brief + */ + +#ifndef COMPRESSOR_BLOCK_VEC_FULL_LOAD_H +#define COMPRESSOR_BLOCK_VEC_FULL_LOAD_H + +#include "compressor_comm.h" +#include "compressor_tools.h" +#include "vf/vf_softmax.h" +#include "vf/vf_add.h" +#include "vf/vf_mul.h" +#include "vf/vf_rms_norm.h" +#include "vf/vf_rope.h" +#include + + +using namespace AscendC; + +namespace Compressor { +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +template +class CompressorBlockVectorFullLoad { +public: + static constexpr bool X_DTYPE = COMP::xDtype == X_DTYPE::BF16; + static constexpr float FLOAT_ZERO = 0; + static constexpr float SOFTMAX_MIN_NUM = -2e38; + // =================================类型定义区================================= + // 中间计算数据类型为float，高精度模式 + using T = float; + using X_T = typename AscendC::Conditional::type; + + __aicore__ inline CompressorBlockVectorFullLoad(){}; + // =================================设置参数================================= + __aicore__ inline void InitParams(const ConstInfo &constInfo, const CompressorTools &tools); + __aicore__ inline void Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut); + // =================================资源管理================================= + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + // =================================执行计算================================= + __aicore__ inline void InitVec1GlobalTensor(GlobalTensor kvMm1ResGm, GlobalTensor scoreMm1ResGm, + GlobalTensor kvCacheTcGm, GlobalTensor scoreCacheTcGm, + GlobalTensor vec1ResGm, GlobalTensor vec2InputGm); + __aicore__ inline void ComputeVec1(); + __aicore__ inline void ComputeVec2(); + +protected: + GlobalTensor vec1ResGm_; + GlobalTensor vec2InputGm_; + GlobalTensor scoreMm1ResGm_; + GlobalTensor kvMm1ResGm_; + GlobalTensor kvCacheTcGm_; + GlobalTensor scoreCacheTcGm_; + +private: + __aicore__ inline uint32_t GetSeqUsed(uint32_t bIdx); + __aicore__ inline uint32_t GetStartPos(uint32_t bIdx); + __aicore__ inline uint32_t GetSeqLength(uint32_t bIdx); + template + __aicore__ inline void DataCopyAlignUbToUb(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, uint32_t srcSingleRowCount, + uint32_t dstSingleRowCount); + template + __aicore__ inline void DataCopyAlignGmToUb(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + uint32_t copyRowCount, uint32_t copyColCount, uint32_t srcSingleRowCount, + uint32_t dstSingleRowCount); + template + __aicore__ inline void DataCopyAlignUbToGm(const GlobalTensor &dstGm, const LocalTensor &srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, uint32_t srcSingleRowCount, + uint32_t dstSingleRowCount); + template + __aicore__ inline void DataCopyWithOutputQue(const GlobalTensor &dstGm, const LocalTensor &srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount); + template + __aicore__ inline void DataCopyWithInputQue(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount); + template + __aicore__ inline void AddMultiDataToUb(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + uint32_t dealRowCount, uint32_t dealColCount, uint32_t srcSingleRowCount, + uint32_t dstSingleRowCount, uint32_t repeatTimes, uint64_t offset); + __aicore__ inline void CopyInApe(uint32_t dStartIdx, uint32_t dDealSize); + template + __aicore__ inline void AddApe(const LocalTensor &scoreLocal, uint32_t dealRowCount, uint32_t dealColCount, + uint32_t scoreSingleRowCount, uint32_t apeSingleRowCount, uint64_t scoreOffset, + uint64_t apeOffset); + __aicore__ inline void AddApeToScore(const LocalTensor &scoreLocal, const Vec1SliceInfo &sliceInfo, + uint32_t dDealSize, uint32_t dBaseSize, uint32_t dStartIdx, + bool isApeFullLoad); + __aicore__ inline void FromWokrSpaceToUb(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + uint32_t preDealSeqCnt, uint32_t dealSeqCnt, uint32_t dStartIdx, + uint32_t dDealSize); + + template + __aicore__ inline void DuplicateFirstBlock(const LocalTensor &dstLocal, uint32_t duplicateRowCount, + uint32_t duplicateColCount, uint32_t singleRowCount); + __aicore__ inline void WriteToCacheState(const GlobalTensor &state, const GlobalTensor &blockTableGm, + const LocalTensor &input, uint32_t batchIdx, uint32_t startSeqIdx, + uint32_t endSeqIdx, uint32_t dDealSize, uint32_t dBaseSize, + uint32_t stateIdx); + __aicore__ inline void ReadFromCacheState(const LocalTensor &output, const GlobalTensor &state, + const GlobalTensor &blockTableGm, uint32_t batchIdx, + uint32_t startSeqIdx, uint32_t endSeqIdx, uint32_t dStartIdx, + uint32_t dDealSize, uint32_t stateIdx); + __aicore__ inline void SaveState(const LocalTensor &srcLocal, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const Vec1SliceInfo &sliceInfo, + uint32_t dStartIdx, uint32_t dDealSize, uint32_t dBaseSize, + uint32_t stateIdx); + template + __aicore__ inline void ReadState(const LocalTensor &srcLocal, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const Vec1SliceInfo &sliceInfo, + uint32_t dStartIdx, uint32_t dDealSize, uint32_t stateIdx); + __aicore__ inline void PadAlign(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + const Vec1SliceInfo &sliceInfo, uint32_t dBaseOffset, uint32_t dDealSize, + uint32_t dBaseSize); + template + __aicore__ inline void OverLap(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + const GlobalTensor &srcGm, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, const GlobalTensor &cacheTcGm, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, uint32_t dStartIdx, + uint32_t dBaseOffset, uint32_t globalSeqIdx, uint32_t dDealSize, uint32_t dBaseSize); + __aicore__ inline void OverLapScoreKv(const LocalTensor &scoreLocal, const LocalTensor &kvLocal, + const LoopInfo &loopInfo, const StatisticInfo &statisticInfo, + const Vec1SliceInfo &originSliceInfo, uint32_t dStartIdx, + uint32_t dBaseOffset, uint32_t dDealSize, uint32_t dBaseSize, + uint32_t dealSeqStartIdx, uint32_t needDealTcSize); + __aicore__ inline void SaveToWorkSpace(const LocalTensor &srcLocal, const GlobalTensor &cacheTcGm, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, uint32_t dStartIdx, + uint32_t dDealSize); + __aicore__ inline void LoadFromWorkSpace(const LocalTensor &dstLocal, const GlobalTensor &cacheTcGm, + const GlobalTensor &srcGm, const LocalTensor &srcLocal, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, + uint32_t dStartIdx, uint32_t globalSeqIdx, uint32_t dDealSize); + __aicore__ inline void SoftmaxDN(const LocalTensor &scoreLocal, uint32_t tcDealSize, uint32_t dDealSize); + __aicore__ inline void KvMulReduceScore(const LocalTensor &kvLocal, const LocalTensor &scoreLocal, + const LocalTensor &dstLocal, uint32_t tcDealSize, uint32_t dDealSize); + __aicore__ inline void CopyOutVec1Res(const GlobalTensor &resGm, const LocalTensor &comperssoredUb, + uint32_t compressTcSize, uint32_t dStartIdx, uint32_t dDealSize); + __aicore__ inline void DealVec1BaseBlock(CompressorVec1SliceIterator &sliceIterator, const LoopInfo &loopInfo, + uint32_t dStartIdx, uint32_t dBaseOffset, uint32_t dDealSize, + uint32_t dBaseSize, uint32_t dealSeqStartIdx); + __aicore__ inline void MultRowRmsNorm(const LocalTensor &normResUb, const LocalTensor &vec1ResUb, + const LocalTensor &normWeightUb, const LocalTensor &tempLocal, + uint32_t dealRowCount); + __aicore__ inline void CalRope(const LocalTensor &outputUb, const LocalTensor &normResUb, + const Vec2SliceInfo &originSliceInfo, uint32_t dealRowCount); + __aicore__ inline void CopyFinalResultOut(const LocalTensor &cmpKvOutUb, + CompressorVec2SliceIterator &sliceIterator); + __aicore__ inline void DealVec2BaseBlock(const Vec2SplitInfo &splitInfo, + CompressorVec2SliceIterator &sliceIterator); + __aicore__ inline void CalcGroupInfo(Vec1SplitInfo &splitInfo); + __aicore__ inline void CalcTaskDistribution(Vec1SplitInfo &splitInfo); + __aicore__ inline void UpdateIteratorState(Vec1SplitInfo &splitInfo); + __aicore__ inline void CalcTilingStrategy(Vec1SplitInfo &splitInfo); + __aicore__ inline Vec1SplitInfo SplitCoreV1(); + __aicore__ inline Vec2SplitInfo SplitCoreV2(); + uint32_t cmpRatio_ = 0U; + uint32_t coff_ = 0U; + uint32_t compressedCnt_ = 0; + uint32_t totalCompressedCnt_ = 0; + uint32_t kvStateIdx_ = 0; + uint32_t scoreStateIdx_ = 1; + bool isExistSeqUsed_ = false; + bool isExistStartPos_ = false; + CompressorTools tools_; + ConstInfo constInfo_ = {}; + GlobalTensor startPosGm_; + GlobalTensor cuSeqlensGm_; + GlobalTensor sequsedGm_; + GlobalTensor stateBlockTableGm_; + GlobalTensor stateCacheGm_; + GlobalTensor apeGm_; + GlobalTensor normWeightGm_; + GlobalTensor ropeSinGm_; + GlobalTensor ropeCosGm_; + GlobalTensor cmpKvOutGm_; + + // ================================Local Buffer区==================================== + // TBuf mm1ResUb; + LocalTensor normWeightUb; + LocalTensor apeUb; + LocalTensor scoreUb; + LocalTensor kvUb; + // 临时tbuf + TBuf tmpBuf1; + TBuf tmpBuf2; + TBuf apeBuf; + TBuf normWeightBuf; + // in queue + TQue inputQue1; + TQue inputQue2; + TQue inputQue3; + TQue inputQueApe; + // out queue + TQue outputQue1; + TQue outputQue2; +}; + + +template +__aicore__ inline void CompressorBlockVectorFullLoad::InitParams(const ConstInfo &constInfo, + const CompressorTools &tools) +{ + this->constInfo_ = constInfo; + this->tools_ = tools; + coff_ = static_cast(COMP::coff); + cmpRatio_ = constInfo.cmpRatio; +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::Init( + __gm__ uint8_t *x, __gm__ uint8_t *wKv, __gm__ uint8_t *wGate, __gm__ uint8_t *stateCache, __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, __gm__ uint8_t *ropeSin, __gm__ uint8_t *ropeCos, __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, __gm__ uint8_t *seqUsed, __gm__ uint8_t *startPos, __gm__ uint8_t *cmpKvOut) +{ + stateBlockTableGm_.SetGlobalBuffer((__gm__ int32_t *)stateBlockTable); + stateCacheGm_.SetGlobalBuffer((__gm__ T *)stateCache); + apeGm_.SetGlobalBuffer((__gm__ T *)ape); + normWeightGm_.SetGlobalBuffer((__gm__ T *)normWeight); + ropeSinGm_.SetGlobalBuffer((__gm__ T *)ropeSin); + ropeCosGm_.SetGlobalBuffer((__gm__ T *)ropeCos); + cmpKvOutGm_.SetGlobalBuffer((__gm__ X_T *)cmpKvOut); + isExistSeqUsed_ = (seqUsed != nullptr); + isExistStartPos_ = (startPos != nullptr); + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens); + } + if (isExistSeqUsed_) { + sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed); + } + if (isExistStartPos_) { + startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos); + } +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(inputQue1, 1, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(inputQue2, 1, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(inputQue3, 1, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(outputQue1, 1, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(outputQue2, 1, BUFFER_SIZE_BYTE_16K); + pipe->InitBuffer(inputQueApe, 1, BUFFER_SIZE_BYTE_16K); + pipe->InitBuffer(normWeightBuf, BUFFER_SIZE_BYTE_4K); + pipe->InitBuffer(tmpBuf1, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(tmpBuf2, BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(apeBuf, BUFFER_SIZE_BYTE_16K); + normWeightUb = normWeightBuf.Get(); + LocalTensor normweightInUb = inputQue2.AllocTensor(); + DataCopy(normweightInUb, normWeightGm_, constInfo_.headDim); // 获取normWeight，常驻 + inputQue2.EnQue(normweightInUb); + inputQue2.DeQue(); + DataCopy(normWeightUb, normweightInUb, constInfo_.headDim); + inputQue2.FreeTensor(normweightInUb); + PipeBarrier(); +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::AllocEventID() +{ +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::FreeEventID() +{ +} + +template +__aicore__ inline void +CompressorBlockVectorFullLoad::InitVec1GlobalTensor(GlobalTensor kvMm1ResGm, GlobalTensor scoreMm1ResGm, + GlobalTensor kvCacheTcGm, GlobalTensor scoreCacheTcGm, + GlobalTensor vec1ResGm, GlobalTensor vec2InputGm) +{ + this->kvMm1ResGm_ = kvMm1ResGm; + this->scoreMm1ResGm_ = scoreMm1ResGm; + this->kvCacheTcGm_ = kvCacheTcGm; + this->scoreCacheTcGm_ = scoreCacheTcGm; + this->vec1ResGm_ = vec1ResGm; + this->vec2InputGm_ = vec2InputGm; +} + +template +__aicore__ inline uint32_t CompressorBlockVectorFullLoad::GetSeqUsed(uint32_t bIdx) +{ + if (isExistSeqUsed_) { + return (uint32_t)sequsedGm_.GetValue(bIdx); + } else { + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + return (uint32_t)(cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx)); + } else { + return constInfo_.sSize; + } + } +} + +template +__aicore__ inline uint32_t CompressorBlockVectorFullLoad::GetStartPos(uint32_t bIdx) +{ + if (isExistStartPos_) { + return startPosGm_.GetValue(bIdx); + } + return 0; +} + +template +__aicore__ inline uint32_t CompressorBlockVectorFullLoad::GetSeqLength(uint32_t bIdx) +{ + if (COMP::xLayout == X_LAYOUT::TH) { + return cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx); + } else { + return constInfo_.sSize; + } +} + + +template +template +__aicore__ inline void +CompressorBlockVectorFullLoad::DataCopyAlignUbToUb(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + DataCopyParams intriParams; + intriParams.blockCount = copyRowCount; + intriParams.blockLen = copyColCount / FP32_BLOCK_ELEMENT_NUM; + intriParams.dstGap = (dstSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + intriParams.srcGap = (srcSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + DataCopy(dstLocal, srcLocal, intriParams); +} + +template +template +__aicore__ inline void +CompressorBlockVectorFullLoad::DataCopyAlignGmToUb(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + DataCopyParams intriParams; + intriParams.blockCount = copyRowCount; + intriParams.blockLen = copyColCount / FP32_BLOCK_ELEMENT_NUM; + intriParams.dstGap = (dstSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + intriParams.srcGap = (srcSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + DataCopy(dstLocal, srcGm, intriParams); +} + +template +template +__aicore__ inline void +CompressorBlockVectorFullLoad::DataCopyAlignUbToGm(const GlobalTensor &dstGm, const LocalTensor &srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + DataCopyParams intriParams; + intriParams.blockCount = copyRowCount; + intriParams.blockLen = copyColCount / FP32_BLOCK_ELEMENT_NUM; + intriParams.dstGap = (dstSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + intriParams.srcGap = (srcSingleRowCount - copyColCount) / FP32_BLOCK_ELEMENT_NUM; + DataCopy(dstGm, srcLocal, intriParams); +} + +template +template +__aicore__ inline void +CompressorBlockVectorFullLoad::DataCopyWithOutputQue(const GlobalTensor &dstGm, const LocalTensor &srcLocal, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + uint32_t singleCopyRowCount = BUFFER_SIZE_BYTE_32K / (copyColCount * sizeof(O)); + for (uint32_t rowCount = 0; rowCount < copyRowCount; rowCount += singleCopyRowCount) { + uint64_t srcOffset = rowCount * srcSingleRowCount; + uint64_t dstOffset = rowCount * dstSingleRowCount; + uint32_t curCopyRowCount = min(singleCopyRowCount, copyRowCount - rowCount); + + LocalTensor outputUb = outputQue1.AllocTensor(); + + DataCopyAlignUbToUb(outputUb, srcLocal[srcOffset], curCopyRowCount, copyColCount, srcSingleRowCount, + copyColCount); + + outputQue1.EnQue(outputUb); + outputQue1.DeQue(); + + DataCopyAlignUbToGm(dstGm[dstOffset], outputUb, curCopyRowCount, copyColCount, copyColCount, dstSingleRowCount); + + outputQue1.FreeTensor(outputUb); + } +} + +template +template +__aicore__ inline void +CompressorBlockVectorFullLoad::DataCopyWithInputQue(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + uint32_t copyRowCount, uint32_t copyColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount) +{ + if (copyRowCount == 0) { + return; + } + uint32_t singleCopyRowCount = BUFFER_SIZE_BYTE_32K / (copyColCount * sizeof(O)); + for (uint32_t rowCount = 0; rowCount < copyRowCount; rowCount += singleCopyRowCount) { + uint64_t srcOffset = rowCount * srcSingleRowCount; + uint64_t dstOffset = rowCount * dstSingleRowCount; + uint32_t curCopyRowCount = min(singleCopyRowCount, copyRowCount - rowCount); + + LocalTensor inputUb = inputQue2.AllocTensor(); + + DataCopyAlignGmToUb(inputUb, srcGm[srcOffset], curCopyRowCount, copyColCount, srcSingleRowCount, copyColCount); + + inputQue2.EnQue(inputUb); + inputQue2.DeQue(); + + DataCopyAlignUbToUb(dstLocal[dstOffset], inputUb, curCopyRowCount, copyColCount, copyColCount, + dstSingleRowCount); + + inputQue2.FreeTensor(inputUb); + } +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::CopyInApe(uint32_t dStartIdx, + uint32_t dDealSize) +{ + apeUb = apeBuf.Get(); + + uint32_t copyRowCount = coff_ * cmpRatio_; + uint32_t copyColCount = dDealSize; + uint32_t dstSingleRowCount = dDealSize; + uint32_t srcSingleRowCount = constInfo_.headDim; + + uint64_t gmOffset = dStartIdx; + + DataCopyWithInputQue(apeUb, apeGm_[gmOffset], copyRowCount, copyColCount, srcSingleRowCount, dstSingleRowCount); + PipeBarrier(); +} + +template +template +__aicore__ inline void CompressorBlockVectorFullLoad::AddMultiDataToUb( + const LocalTensor &dstLocal, const GlobalTensor &srcGm, uint32_t dealRowCount, uint32_t dealColCount, + uint32_t srcSingleRowCount, uint32_t dstSingleRowCount, uint32_t repeatTimes, uint64_t offset) +{ + uint32_t cnt = dealRowCount * dealColCount; + uint32_t groupSize = BUFFER_SIZE_BYTE_32K / (cnt * sizeof(O)); + uint32_t loopTimes = CeilDivT(repeatTimes, groupSize); + uint64_t srcGmOffset = 0; + for (uint32_t idx = 0; idx < loopTimes; idx++) { + auto &inputQue = idx % 2 == 0 ? inputQue2 : inputQue3; + uint32_t curGroupSize = min(groupSize, (repeatTimes - groupSize * idx)); + LocalTensor splitLocal = inputQue.AllocTensor(); + if (srcSingleRowCount == dstSingleRowCount && dstSingleRowCount == dealRowCount) { + for (uint32_t groupIdx = 0; groupIdx < curGroupSize; groupIdx++) { + DataCopy(splitLocal[groupIdx * cnt], srcGm[srcGmOffset], cnt); + srcGmOffset += offset; + } + } else { + for (uint32_t groupIdx = 0; groupIdx < curGroupSize; groupIdx++) { + DataCopyAlignGmToUb(splitLocal[groupIdx * cnt], srcGm[srcGmOffset], dealRowCount, dealColCount, + srcSingleRowCount, dstSingleRowCount); + srcGmOffset += offset; + } + } + + inputQue.EnQue(splitLocal); + inputQue.DeQue(); + + PipeBarrier(); + if (idx == 0) { + MultiAddVF(dstLocal, splitLocal, dealRowCount, dealColCount, dealColCount, curGroupSize, cnt); + } else { + MultiAddVF(dstLocal, splitLocal, dealRowCount, dealColCount, dealColCount, curGroupSize, cnt); + } + inputQue.FreeTensor(splitLocal); + } + PipeBarrier(); +} + +template +template +__aicore__ inline void CompressorBlockVectorFullLoad::AddApe(const LocalTensor &scoreLocal, + uint32_t dealRowCount, uint32_t dealColCount, + uint32_t scoreSingleRowCount, + uint32_t apeSingleRowCount, + uint64_t scoreOffset, uint64_t apeOffset) +{ + if constexpr (IS_FULLLOAD) { + AddVF(scoreLocal[scoreOffset], apeUb[apeOffset], coff_ * dealRowCount, dealColCount, scoreSingleRowCount, + apeSingleRowCount); + } else { + apeUb = inputQueApe.AllocTensor(); + DataCopyAlignGmToUb(apeUb, apeGm_[apeOffset], coff_ * dealRowCount, dealColCount, constInfo_.headDim, + apeSingleRowCount); + inputQueApe.EnQue(apeUb); + inputQueApe.DeQue(); + AddVF(scoreLocal[scoreOffset], apeUb, coff_ * dealRowCount, dealColCount, scoreSingleRowCount, + apeSingleRowCount); + inputQueApe.FreeTensor(apeUb); + } +} + +template +__aicore__ inline void +CompressorBlockVectorFullLoad::AddApeToScore(const LocalTensor &scoreLocal, const Vec1SliceInfo &sliceInfo, + uint32_t dDealSize, uint32_t dBaseSize, uint32_t dStartIdx, + bool isApeFullLoad) +{ + uint32_t singleUbRowElemNum = dBaseSize * coff_; + uint32_t singleApeRowElemNum = isApeFullLoad ? singleUbRowElemNum : constInfo_.headDim * coff_; + uint64_t scoreOffset = sliceInfo.dealedSeqCnt * singleUbRowElemNum; + + uint32_t tcDealSize = sliceInfo.dealTcSize; + if (sliceInfo.headHolderSeqCnt > 0) { + uint32_t row = tcDealSize == 1 ? sliceInfo.validSeqCnt : (cmpRatio_ - sliceInfo.headHolderSeqCnt); + + if (isApeFullLoad) { + uint64_t apeOffset = sliceInfo.headHolderSeqCnt * singleApeRowElemNum; + AddApe(scoreLocal, row, dDealSize, dBaseSize, dBaseSize, scoreOffset, apeOffset); + + } else { + uint64_t apeOffset = sliceInfo.headHolderSeqCnt * singleApeRowElemNum + dStartIdx; + AddApe(scoreLocal, row, dDealSize, dBaseSize, dDealSize, scoreOffset, apeOffset); + } + scoreOffset += row * singleUbRowElemNum; + tcDealSize -= 1; + } + if (tcDealSize == 0) { + return; + } + if (sliceInfo.tailHolderSeqCnt > 0) { + tcDealSize -= 1; + uint32_t row = cmpRatio_ - sliceInfo.tailHolderSeqCnt; + uint32_t tailScoreOffset = scoreOffset + tcDealSize * cmpRatio_ * singleUbRowElemNum; + if (isApeFullLoad) { + uint64_t apeOffset = 0; + AddApe(scoreLocal, row, dDealSize, dBaseSize, dBaseSize, tailScoreOffset, apeOffset); + + } else { + uint64_t apeOffset = dStartIdx; + AddApe(scoreLocal, row, dDealSize, dBaseSize, dDealSize, tailScoreOffset, apeOffset); + } + } + if (tcDealSize == 0) { + return; + } + + if (isApeFullLoad) { + uint32_t row = cmpRatio_; + for (uint32_t r = 0; r < tcDealSize; r++) { + uint64_t curScoreOffset = scoreOffset + r * row * singleUbRowElemNum; + AddApe(scoreLocal, row, dDealSize, dBaseSize, dDealSize, curScoreOffset, 0U); + } + } + +} + + +template +__aicore__ inline void +CompressorBlockVectorFullLoad::FromWokrSpaceToUb(const LocalTensor &dstLocal, const GlobalTensor &srcGm, + uint32_t preDealSeqCnt, uint32_t dealSeqCnt, uint32_t dStartIdx, + uint32_t dDealSize) +{ + uint32_t srcSingleRowElemNum = constInfo_.headDim; + uint32_t copyRowCount = dealSeqCnt * coff_; + uint32_t copyColCount = dDealSize; + uint32_t srcSingleRowCount = srcSingleRowElemNum; + uint32_t dstSingleRowCount = dDealSize; + uint64_t srcGmOffset = preDealSeqCnt * srcSingleRowElemNum * coff_ + dStartIdx; + if (constInfo_.kBaseNum == 1) { + DataCopyAlignGmToUb(dstLocal, srcGm[srcGmOffset], copyRowCount, copyColCount, srcSingleRowCount, + dstSingleRowCount); + } else { + AddMultiDataToUb(dstLocal, srcGm[srcGmOffset], copyRowCount, copyColCount, srcSingleRowCount, dstSingleRowCount, + constInfo_.kBaseNum, constInfo_.mm1KvResSize); + } +} + +template +__aicore__ inline void +CompressorBlockVectorFullLoad::PadAlign(const LocalTensor &dstLocal, const LocalTensor &srcLocal, + const Vec1SliceInfo &sliceInfo, uint32_t dBaseOffset, uint32_t dDealSize, + uint32_t dBaseSize) +{ + // Ub data layout after overlap when r = 4 and coff = 2: + // Tc0_seq01: |--- --D_L--- -|------D_R-----| + // Tc0_seq02: |--- --D_L--- -|------D_R-----| + // Tc0_seq03: |--- --D_L--- -|------D_R-----| + // Tc0_seq04: |--- --D_L--- -|------D_R-----| + // Tc1_seq01: |--- --D_L--- -|------D_R-----| + // Tc1_seq02: |--- --D_L--- -|------D_R-----| + // Tc1_seq03: |--- --D_L--- -|------D_R-----| + // Tc1_seq04: |--- --D_L--- -|------D_R-----| + uint32_t srcSingleRowElemNum = dBaseSize * coff_; + uint32_t copyRowCount = sliceInfo.compressTcSize * cmpRatio_ - sliceInfo.headHolderSeqCnt; + uint32_t copyColCount = dDealSize; + uint32_t srcSingleRowCount = srcSingleRowElemNum; + uint32_t dstSingleRowCount = dDealSize * coff_; // left和right在seq方向是交错存储的 + uint64_t srcLocalOffset = sliceInfo.dealedSeqCnt * srcSingleRowElemNum + dBaseOffset; + + uint64_t dstUbOffset = sliceInfo.compressoredScCnt * cmpRatio_ * dstSingleRowCount; + if constexpr (COMP::coff == COFF::OVERLAP) { + // 左侧 + uint64_t preSrcLocalOffset = srcLocalOffset; + uint64_t preDstUbOffset = dstUbOffset + (sliceInfo.headHolderSeqCnt + cmpRatio_) * dstSingleRowCount; + DataCopyAlignUbToUb(dstLocal[preDstUbOffset], srcLocal[preSrcLocalOffset], + copyRowCount - min(copyRowCount, cmpRatio_), copyColCount, srcSingleRowCount, + dstSingleRowCount); + dstUbOffset += dDealSize; + srcLocalOffset += dBaseSize; + } + // 右侧 + dstUbOffset += sliceInfo.headHolderSeqCnt * dstSingleRowCount; + DataCopyAlignUbToUb(dstLocal[dstUbOffset], srcLocal[srcLocalOffset], copyRowCount, copyColCount, srcSingleRowCount, + dstSingleRowCount); +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::WriteToCacheState(const GlobalTensor &state, + const GlobalTensor &blockTableGm, + const LocalTensor &input, + uint32_t batchIdx, uint32_t startSeqIdx, + uint32_t endSeqIdx, uint32_t dDealSize, + uint32_t dBaseSize, uint32_t stateIdx) +{ + if constexpr (COMP::cacheMode == CACHE_MODE::CONTINUOUS) { + uint64_t blockTablebaseOffset = batchIdx * constInfo_.maxBlockNumPerBatch; + uint32_t curSeqIdx = startSeqIdx; + uint32_t copyFinishRowCnt = 0; + uint32_t seqCnt = endSeqIdx - startSeqIdx; + while (copyFinishRowCnt < seqCnt) { + uint64_t blockIdOffset = curSeqIdx / constInfo_.blockSize; + uint64_t remainRowCnt = curSeqIdx % constInfo_.blockSize; + uint64_t idInBlockTable = blockTableGm.GetValue(blockTablebaseOffset + blockIdOffset); + uint32_t copyRowCount = constInfo_.blockSize - remainRowCnt; + if (copyFinishRowCnt + copyRowCount > seqCnt) { + copyRowCount = seqCnt - copyFinishRowCnt; + } + // copyRowCount *= coff_; + if (idInBlockTable != 0) { // 32 + uint64_t stateOffset = + idInBlockTable * constInfo_.stateCacheStrideDim0 + remainRowCnt * 2 * coff_ * constInfo_.headDim + + stateIdx * coff_ * constInfo_.headDim; + uint64_t ubOffset = copyFinishRowCnt * coff_ * dBaseSize; + DataCopyWithOutputQue(state[stateOffset], input[ubOffset], copyRowCount, dDealSize, coff_ * dBaseSize, + coff_ * constInfo_.headDim * 2); + } + + copyFinishRowCnt += copyRowCount; + curSeqIdx += copyRowCount; + } + } else { + uint32_t curSeqIdx = startSeqIdx; + uint32_t copyFinishRowCnt = 0; + uint32_t seqCnt = endSeqIdx - startSeqIdx; + uint64_t idInBlockTable = blockTableGm.GetValue(batchIdx); + while (copyFinishRowCnt < seqCnt) { + uint64_t remainRowCnt = curSeqIdx % constInfo_.blockSize; + uint32_t copyRowCount = constInfo_.blockSize - remainRowCnt; + if (copyFinishRowCnt + copyRowCount > seqCnt) { + copyRowCount = seqCnt - copyFinishRowCnt; + } + uint64_t stateOffset = idInBlockTable * constInfo_.stateCacheStrideDim0 + + remainRowCnt * 2 * coff_ * constInfo_.headDim + + stateIdx * coff_ * constInfo_.headDim; + uint64_t ubOffset = copyFinishRowCnt * coff_ * dBaseSize; + DataCopyWithOutputQue(state[stateOffset], input[ubOffset], copyRowCount, + dDealSize, coff_ * dBaseSize, coff_ * constInfo_.headDim * 2); + + copyFinishRowCnt += copyRowCount; + curSeqIdx += copyRowCount; + } + } +} + +template +__aicore__ inline void +CompressorBlockVectorFullLoad::ReadFromCacheState(const LocalTensor &output, const GlobalTensor &state, + const GlobalTensor &blockTableGm, uint32_t batchIdx, + uint32_t startSeqIdx, uint32_t endSeqIdx, uint32_t dStartIdx, + uint32_t dDealSize, uint32_t stateIdx) +{ + if constexpr (COMP::cacheMode == CACHE_MODE::CONTINUOUS) { + uint64_t blockTablebaseOffset = batchIdx * constInfo_.maxBlockNumPerBatch; + uint32_t curSeqIdx = startSeqIdx; + uint32_t copyFinishRowCnt = 0; + uint32_t seqCnt = endSeqIdx - startSeqIdx; + while (copyFinishRowCnt < seqCnt) { + uint64_t blockIdOffset = curSeqIdx / constInfo_.blockSize; + uint64_t remainRowCnt = curSeqIdx % constInfo_.blockSize; + uint64_t idInBlockTable = blockTableGm.GetValue(blockTablebaseOffset + blockIdOffset); + uint32_t copyRowCount = constInfo_.blockSize - remainRowCnt; + if (copyFinishRowCnt + copyRowCount > seqCnt) { + copyRowCount = seqCnt - copyFinishRowCnt; + } + uint64_t stateOffset = idInBlockTable * constInfo_.stateCacheStrideDim0 + + remainRowCnt * 2 * coff_ * constInfo_.headDim + + stateIdx * coff_ * constInfo_.headDim + dStartIdx; + + DataCopyWithInputQue(output[copyFinishRowCnt * coff_ * dDealSize], state[stateOffset], copyRowCount, + dDealSize, coff_ * constInfo_.headDim * 2, coff_ * dDealSize); + copyFinishRowCnt += copyRowCount; + curSeqIdx += copyRowCount; + } + } else { + uint32_t curSeqIdx = startSeqIdx; + uint32_t copyFinishRowCnt = 0; + uint32_t seqCnt = endSeqIdx - startSeqIdx; + uint64_t idInBlockTable = blockTableGm.GetValue(batchIdx); + while (copyFinishRowCnt < seqCnt) { + uint64_t remainRowCnt = curSeqIdx % constInfo_.blockSize; + uint32_t copyRowCount = constInfo_.blockSize - remainRowCnt; + if (copyFinishRowCnt + copyRowCount > seqCnt) { + copyRowCount = seqCnt - copyFinishRowCnt; + } + uint64_t stateOffset = idInBlockTable * constInfo_.stateCacheStrideDim0 + + remainRowCnt * 2 * coff_ * constInfo_.headDim + + stateIdx * coff_ * constInfo_.headDim + dStartIdx; + + DataCopyWithInputQue(output[copyFinishRowCnt * coff_ * dDealSize], state[stateOffset], copyRowCount, + dDealSize, coff_ * constInfo_.headDim * 2, coff_ * dDealSize); + copyFinishRowCnt += copyRowCount; + curSeqIdx += copyRowCount; + } + } +} + + +template +template +__aicore__ inline void +CompressorBlockVectorFullLoad::DuplicateFirstBlock(const LocalTensor &dstLocal, uint32_t duplicateRowCount, + uint32_t duplicateColCount, uint32_t singleRowCount) +{ + for (uint32_t offset = 0; offset < duplicateColCount; offset += FP32_REPEAT_ELEMENT_NUM) { + uint32_t curDuplicateColCount = min(duplicateColCount - offset, FP32_REPEAT_ELEMENT_NUM); + if constexpr (IS_SCORE) { + Duplicate(dstLocal[offset], SOFTMAX_MIN_NUM, curDuplicateColCount, duplicateRowCount, 1, + singleRowCount / REPEAT_STRIDE_NUM); + } else { + Duplicate(dstLocal[offset], FLOAT_ZERO, curDuplicateColCount, duplicateRowCount, 1, + singleRowCount / REPEAT_STRIDE_NUM); + } + } +} + +template +__aicore__ inline void +CompressorBlockVectorFullLoad::SaveState(const LocalTensor &srcLocal, const GlobalTensor &stateGm, + const GlobalTensor &blockTableGm, + const Vec1SliceInfo &sliceInfo, uint32_t dStartIdx, + uint32_t dDealSize, uint32_t dBaseSize, uint32_t stateIdx) +{ + uint32_t startSeqIdx = sliceInfo.bStartPos + sliceInfo.sIdx; + uint32_t endSeqIdx = startSeqIdx + sliceInfo.validSeqCnt; + uint64_t srcBaseOffset = sliceInfo.dealedSeqCnt * coff_ * dBaseSize; + + if constexpr (COMP::cacheMode == CACHE_MODE::CYCLE) { + uint32_t compressSeqIdx = Trunc(sliceInfo.bStartPos + sliceInfo.bSeqUsed, cmpRatio_); + uint32_t writeSeqStartIdx = compressSeqIdx > (coff_ - 1) * cmpRatio_ ? + compressSeqIdx - (coff_ - 1) * cmpRatio_ : 0; + if (endSeqIdx <= writeSeqStartIdx) { + return; + } + srcBaseOffset += (max(startSeqIdx, writeSeqStartIdx) - startSeqIdx) * coff_ * dBaseSize; + startSeqIdx = max(startSeqIdx, writeSeqStartIdx); + } + + if constexpr (COMP::coff == COFF::OVERLAP) { + WriteToCacheState(stateGm[dStartIdx], blockTableGm, srcLocal[srcBaseOffset], sliceInfo.bIdx, + startSeqIdx, endSeqIdx, dDealSize, dBaseSize, stateIdx); + srcBaseOffset += dBaseSize; + dStartIdx += constInfo_.headDim; + } + + WriteToCacheState(stateGm[dStartIdx], blockTableGm, srcLocal[srcBaseOffset], sliceInfo.bIdx, + startSeqIdx, endSeqIdx, dDealSize, dBaseSize, stateIdx); +} + + +template +template +__aicore__ inline void CompressorBlockVectorFullLoad::ReadState( + const LocalTensor &dstLocal, const GlobalTensor &stateGm, const GlobalTensor &blockTableGm, + const Vec1SliceInfo &sliceInfo, uint32_t dStartIdx, uint32_t dDealSize, uint32_t stateIdx) +{ + // 没有需要压缩的块时, 不需要读state的信息 + if (sliceInfo.compressTcSize == 0) { + return; + } + // 填充右边 + if (sliceInfo.headHolderSeqCnt > 0) { + // 整个batch的第一块 + uint32_t startSeqIdx = Trunc(sliceInfo.bStartPos + sliceInfo.sIdx, cmpRatio_); + uint32_t endSeqIdx = sliceInfo.bStartPos; + uint64_t dstBaseOffset = sliceInfo.compressoredScCnt * cmpRatio_ * coff_ * dDealSize; + if constexpr (COMP::coff == Compressor::COFF::OVERLAP) { + dstBaseOffset += (coff_ - 1) * dDealSize; + } + ReadFromCacheState(dstLocal[dstBaseOffset], stateGm, blockTableGm, sliceInfo.bIdx, startSeqIdx, endSeqIdx, + dStartIdx + (coff_ - 1) * constInfo_.headDim, dDealSize, stateIdx); + } + + // 填充左边 + if constexpr (COMP::coff == Compressor::COFF::OVERLAP) { + bool isFirst = sliceInfo.bStartPos + sliceInfo.sIdx < cmpRatio_; + if (isFirst) { + // 无历史数据 + // dDealSize必须为64 + uint64_t dstBaseOffset = sliceInfo.compressoredScCnt * cmpRatio_ * coff_ * dDealSize; + DuplicateFirstBlock(dstLocal[dstBaseOffset], cmpRatio_, dDealSize, coff_ * dDealSize); + } + if (sliceInfo.sIdx < cmpRatio_ && (!isFirst || sliceInfo.compressTcSize > 1)) { + uint32_t startSeqIdx = + sliceInfo.bStartPos < cmpRatio_ ? + 0 : + Trunc(sliceInfo.bStartPos + sliceInfo.sIdx, cmpRatio_) - cmpRatio_; + uint32_t endSeqIdx = + min(Trunc(sliceInfo.bStartPos + sliceInfo.sIdx + sliceInfo.validSeqCnt, cmpRatio_) - + cmpRatio_, + sliceInfo.bStartPos); + uint64_t dstBaseOffset = sliceInfo.compressoredScCnt * cmpRatio_ * coff_ * dDealSize; + if (isFirst) { + dstBaseOffset += cmpRatio_ * coff_ * dDealSize; + } + ReadFromCacheState(dstLocal[dstBaseOffset], stateGm, blockTableGm, sliceInfo.bIdx, startSeqIdx, endSeqIdx, + dStartIdx, dDealSize, stateIdx); + } + } +} + + +template +template +__aicore__ inline void CompressorBlockVectorFullLoad::OverLap( + const LocalTensor &dstLocal, const LocalTensor &srcLocal, const GlobalTensor &srcGm, + const GlobalTensor &stateGm, const GlobalTensor &blockTableGm, const GlobalTensor &cacheTcGm, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, uint32_t dStartIdx, uint32_t dBaseOffset, + uint32_t globalSeqIdx, uint32_t dDealSize, uint32_t dBaseSize) +{ + if (sliceInfo.dealTcSize == 0) { + return; + } + + ReadState(dstLocal, stateGm, blockTableGm, sliceInfo, dStartIdx + dBaseOffset, dDealSize, + static_cast(IS_SCORE)); + + if (sliceInfo.compressTcSize > 0) { + PadAlign(dstLocal, srcLocal, sliceInfo, dBaseOffset, dDealSize, dBaseSize); + if constexpr (COMP::coff == COFF::OVERLAP) { + GlobalTensor curCacheTcGm = cacheTcGm; + LoadFromWorkSpace(dstLocal, curCacheTcGm, srcGm, srcLocal, sliceInfo, loopInfo, dStartIdx, globalSeqIdx, + dDealSize); + } + } +} + +template +__aicore__ inline void +CompressorBlockVectorFullLoad::SaveToWorkSpace(const LocalTensor &srcLocal, const GlobalTensor &cacheTcGm, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, + uint32_t dStartIdx, uint32_t dDealSize) +{ + uint32_t curSeqLen = sliceInfo.bStartPos + sliceInfo.sIdx + sliceInfo.validSeqCnt; + uint32_t totalSeqLen = sliceInfo.bStartPos + sliceInfo.sIdx + sliceInfo.bSeqUsed; + if (!loopInfo.isCoreRowLast || !loopInfo.isCoreLoopLast || !sliceInfo.isLast || totalSeqLen < cmpRatio_ || + curSeqLen > Trunc(totalSeqLen, cmpRatio_) - cmpRatio_) { + return; + } + uint32_t srcSingleRowElemNum = dDealSize * coff_; + uint64_t srcLocalOffset = + (sliceInfo.dealedSeqCnt + sliceInfo.validSeqCnt - min(sliceInfo.validSeqCnt, cmpRatio_)) * + srcSingleRowElemNum; + DataCopyWithOutputQue(cacheTcGm[dStartIdx], srcLocal[srcLocalOffset], + curSeqLen - max(curSeqLen - cmpRatio_, sliceInfo.bStartPos), dDealSize, + coff_ * dDealSize, constInfo_.headDim); +} + +template +__aicore__ inline void +CompressorBlockVectorFullLoad::LoadFromWorkSpace(const LocalTensor &dstLocal, const GlobalTensor &cacheTcGm, + const GlobalTensor &srcGm, const LocalTensor &srcLocal, + const Vec1SliceInfo &sliceInfo, const LoopInfo &loopInfo, + uint32_t dStartIdx, uint32_t globalSeqIdx, uint32_t dDealSize) +{ + if (sliceInfo.sIdx == 0) { + return; + } + uint32_t dstSingleRowElemNum = dDealSize * coff_; + uint32_t copyRowCount = min(sliceInfo.sIdx, cmpRatio_); + uint64_t dstLocalOffset = + (sliceInfo.compressoredScCnt * cmpRatio_ + cmpRatio_ - copyRowCount) * dstSingleRowElemNum; + if (loopInfo.isCoreRowFirst && loopInfo.isCoreLoopFirst && sliceInfo.isFirst) { // 从cacheGm获取 + uint32_t srcSingleRowElemNum = constInfo_.headDim; + uint64_t srcLocalOffset = dStartIdx; + + DataCopyWithInputQue(dstLocal[dstLocalOffset], cacheTcGm[srcLocalOffset], copyRowCount, dDealSize, + srcSingleRowElemNum, coff_ * dDealSize); + } else if (sliceInfo.isFirst) { // 从存放MatMul结果的WorkSpace中获取 + uint32_t srcSingleRowElemNum = constInfo_.headDim * coff_; + uint64_t srcLocalOffset = + (globalSeqIdx + sliceInfo.dealedSeqCnt - copyRowCount) * srcSingleRowElemNum + dStartIdx; + + if (constInfo_.kBaseNum == 1) { + DataCopyWithInputQue(dstLocal[dstLocalOffset], srcGm[srcLocalOffset], copyRowCount, dDealSize, + srcSingleRowElemNum, coff_ * dDealSize); + } else { + AddMultiDataToUb(dstLocal[dstLocalOffset], srcGm[srcLocalOffset], copyRowCount, dDealSize, + srcSingleRowElemNum, coff_ * dDealSize, constInfo_.kBaseNum, constInfo_.mm1KvResSize); + } + } else { // 从UB中获取 + uint32_t srcSingleRowElemNum = dDealSize * coff_; + uint64_t srcLocalOffset = (sliceInfo.dealedSeqCnt - copyRowCount) * srcSingleRowElemNum; + DataCopyAlignUbToUb(dstLocal[dstLocalOffset], srcLocal[srcLocalOffset], copyRowCount, dDealSize, + srcSingleRowElemNum, coff_ * dDealSize); + } +} + + +template +__aicore__ inline void CompressorBlockVectorFullLoad::OverLapScoreKv( + const LocalTensor &scoreLocal, const LocalTensor &kvLocal, const LoopInfo &loopInfo, + const StatisticInfo &statisticInfo, const Vec1SliceInfo &originSliceInfo, uint32_t dStartIdx, uint32_t dBaseOffset, + uint32_t dDealSize, uint32_t dBaseSize, uint32_t dealSeqStartIdx, uint32_t needDealTcSize) +{ + CompressorVec1SliceIterator overLapSliceIterator(tools_); + overLapSliceIterator.SetMaxBatchSize(constInfo_.batchSize); + Vec1SliceInfo &overLapSliceInfo = overLapSliceIterator.GetSlice(); + + GlobalTensor scoreDBMm1ResGm = scoreMm1ResGm_; + overLapSliceIterator.Reset(originSliceInfo.bIdx, originSliceInfo.sIdx, originSliceInfo.dealedSeqCnt, 0U); + overLapSliceIterator.SetNeedDealTcSize(needDealTcSize); + + while (!overLapSliceIterator.IsEnd()) { + overLapSliceIterator.GetSlice(); + OverLap(scoreLocal, scoreUb, scoreDBMm1ResGm, stateCacheGm_, stateBlockTableGm_, scoreCacheTcGm_, + overLapSliceInfo, loopInfo, dStartIdx, dBaseOffset, + originSliceInfo.dealedSeqCnt + dealSeqStartIdx, dDealSize, dBaseSize); + overLapSliceIterator.IteratorSlice(); + } + + GlobalTensor kvDBMm1ResGm = kvMm1ResGm_; + overLapSliceIterator.Reset(originSliceInfo.bIdx, originSliceInfo.sIdx, originSliceInfo.dealedSeqCnt, 0U); + overLapSliceIterator.SetNeedDealTcSize(needDealTcSize); + + while (!overLapSliceIterator.IsEnd()) { + overLapSliceIterator.GetSlice(); + OverLap(kvLocal, kvUb, kvDBMm1ResGm, stateCacheGm_, stateBlockTableGm_, kvCacheTcGm_, overLapSliceInfo, + loopInfo, dStartIdx, dBaseOffset, originSliceInfo.dealedSeqCnt + dealSeqStartIdx, dDealSize, + dBaseSize); + overLapSliceIterator.IteratorSlice(); + } + PipeBarrier(); +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::SoftmaxDN(const LocalTensor &scoreLocal, + uint32_t tcDealSize, uint32_t dDealSize) +{ + uint32_t ReduceSize = coff_ * cmpRatio_; + FaVectorApi::SoftmaxDnVF(scoreLocal, scoreLocal, dDealSize, ReduceSize, tcDealSize, SOFTMAX_MIN_NUM, dDealSize); +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::KvMulReduceScore(const LocalTensor &kvLocal, + const LocalTensor &scoreLocal, + const LocalTensor &dstLocal, + uint32_t tcDealSize, uint32_t dDealSize) +{ + MulReduceSumbaseVF(kvLocal, scoreLocal, dstLocal, coff_, cmpRatio_, dDealSize, tcDealSize); +} + +template +__aicore__ inline void +CompressorBlockVectorFullLoad::CopyOutVec1Res(const GlobalTensor &resGm, const LocalTensor &comperssoredUb, + uint32_t compressTcSize, uint32_t dStartIdx, uint32_t dDealSize) +{ + uint64_t outGmOffset = compressedCnt_ * constInfo_.headDim + dStartIdx; + DataCopyAlignUbToGm(resGm[outGmOffset], comperssoredUb, compressTcSize, dDealSize, dDealSize, constInfo_.headDim); +} + + +template +__aicore__ inline void CompressorBlockVectorFullLoad::DealVec1BaseBlock( + CompressorVec1SliceIterator &sliceIterator, const LoopInfo &loopInfo, uint32_t dStartIdx, + uint32_t dBaseOffset, uint32_t dDealSize, uint32_t dBaseSize, uint32_t dealSeqStartIdx) +{ + Vec1SliceInfo originSliceInfo = sliceIterator.GetSlice(); + uint32_t needDealTcSize = sliceIterator.GetNeedDealTcSize(); + StatisticInfo &statisticInfo = sliceIterator.template FullIteratorSlice(); + if (statisticInfo.actualTcCnt == 0) { + return; + } + LocalTensor scoreLocal = tmpBuf1.Get(); + LocalTensor kvLocal = tmpBuf2.Get(); + OverLapScoreKv(scoreLocal, kvLocal, loopInfo, statisticInfo, originSliceInfo, dStartIdx, dBaseOffset, dDealSize, + dBaseSize, dealSeqStartIdx, needDealTcSize); + if (statisticInfo.compressorScCnt > 0) { + SoftmaxDN(scoreLocal, statisticInfo.compressorScCnt, dDealSize); + LocalTensor comperssoredUb = outputQue2.AllocTensor(); + PipeBarrier(); + KvMulReduceScore(kvLocal, scoreLocal, comperssoredUb, statisticInfo.compressorScCnt, dDealSize); + PipeBarrier(); + outputQue2.EnQue(comperssoredUb); + outputQue2.DeQue(); + GlobalTensor resGm = vec1ResGm_; + CopyOutVec1Res(resGm, comperssoredUb, statisticInfo.compressorScCnt, dStartIdx + dBaseOffset, dDealSize); + outputQue2.FreeTensor(comperssoredUb); + } + compressedCnt_ += statisticInfo.compressorScCnt; +} + + +template +__aicore__ inline void +CompressorBlockVectorFullLoad::MultRowRmsNorm(const LocalTensor &normResUb, const LocalTensor &vec1ResUb, + const LocalTensor &normWeightUb, const LocalTensor &tempLocal, + uint32_t dealRowCount) +{ + uint32_t row = 1; + uint32_t col = constInfo_.headDim; + float reciprocal = 1.0f / col; + float epsilon = constInfo_.normEps; + for (uint32_t i = 0; i < dealRowCount; ++i) { + RmsNormVF(normResUb[i * col], vec1ResUb[i * col], normWeightUb, reciprocal, epsilon, row, col); + } +} + + +template +__aicore__ inline void +CompressorBlockVectorFullLoad::CalRope(const LocalTensor &outputUb, const LocalTensor &normResUb, + const Vec2SliceInfo &originSliceInfo, uint32_t dealRowCount) +{ + CompressorVec2SliceIterator ropeSliceIterator(tools_); + ropeSliceIterator.SetMaxBatchSize(constInfo_.batchSize); + ropeSliceIterator.Reset(originSliceInfo.bIdx, originSliceInfo.scIdx, originSliceInfo.dealedScCnt); + Vec2SliceInfo &sliceInfo = ropeSliceIterator.GetSlice(); + + ropeSliceIterator.SetNeedDealScSize(dealRowCount); + while (!ropeSliceIterator.IsEnd()) { + ropeSliceIterator.GetSlice(); + if (sliceInfo.curDealScNum > 0) { + uint32_t computeSize = sliceInfo.curDealScNum * constInfo_.ropeHeadDim; + uint64_t SinCosOffset = sliceInfo.padScIdx * constInfo_.ropeHeadDim; + + // sin与cos各占一半, 实际分别最多只会用16K,总占用32K + LocalTensor cosUb = inputQue2.AllocTensor(); + LocalTensor sinUb = cosUb[BUFFER_SIZE_BYTE_16K / sizeof(T)]; + DataCopy(cosUb, ropeCosGm_[SinCosOffset], computeSize); + DataCopy(sinUb, ropeSinGm_[SinCosOffset], computeSize); + inputQue2.EnQue(sinUb); + inputQue2.DeQue(); + RopeVF(sinUb, cosUb, normResUb[sliceInfo.loopDealedScCnt * constInfo_.headDim], + outputUb[sliceInfo.loopDealedScCnt * constInfo_.headDim], sliceInfo.curDealScNum, + constInfo_.ropeHeadDim, constInfo_.headDim, constInfo_.headDim - constInfo_.ropeHeadDim); + inputQue2.FreeTensor(sinUb); + } + ropeSliceIterator.IteratorSlice(); + } +} + + +template +__aicore__ inline void +CompressorBlockVectorFullLoad::CopyFinalResultOut(const LocalTensor &cmpKvOutUb, + CompressorVec2SliceIterator &sliceIterator) +{ + Vec2SliceInfo &sliceInfo = sliceIterator.GetSlice(); + while (!sliceIterator.IsEnd()) { + sliceIterator.GetSlice(); + if (sliceInfo.curDealScNum > 0) { + DataCopy(cmpKvOutGm_[sliceInfo.padScIdx * constInfo_.headDim], + cmpKvOutUb[sliceInfo.loopDealedScCnt * constInfo_.headDim], + sliceInfo.curDealScNum * constInfo_.headDim); + } + sliceIterator.IteratorSlice(); + } +} + + +template +__aicore__ inline void +CompressorBlockVectorFullLoad::DealVec2BaseBlock(const Vec2SplitInfo &splitInfo, + CompressorVec2SliceIterator &sliceIterator) +{ + Vec2SliceInfo &sliceInfo = sliceIterator.GetSlice(); + uint32_t needDealScSize = sliceIterator.GetNeedDealScSize(); + uint32_t computeSize = needDealScSize * constInfo_.headDim; + int64_t inGmOffset = splitInfo.preScCnt * constInfo_.headDim; + GlobalTensor vec2InputGm = vec2InputGm_; + // CopyIn + LocalTensor vec1ResUb = inputQue1.AllocTensor(); + DataCopy(vec1ResUb, vec2InputGm[inGmOffset], computeSize); + inputQue1.EnQue(vec1ResUb); + inputQue1.DeQue(); + + // RmsNorm + LocalTensor normResUb = tmpBuf1.Get(); + LocalTensor tempLocal = tmpBuf2.Get(); + PipeBarrier(); + MultRowRmsNorm(normResUb, vec1ResUb, normWeightUb, tempLocal, needDealScSize); + inputQue1.FreeTensor(vec1ResUb); + + + // rope: 只对后RD进行rope; 将normResUb每行前headDim - + // ropeHeadDim个元素cast到X_T，然后再与rope后的结果组合存到outputUb + LocalTensor outputUb = outputQue1.AllocTensor(); + PipeBarrier(); + CalRope(outputUb, normResUb, sliceInfo, needDealScSize); + PipeBarrier(); + // CopyOut + outputQue1.EnQue(outputUb); + outputQue1.DeQue(); + CopyFinalResultOut(outputUb, sliceIterator); + outputQue1.FreeTensor(outputUb); +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::CalcGroupInfo(Vec1SplitInfo &splitInfo) +{ + uint32_t aiCoreNum = constInfo_.usedCoreNum * 2; + splitInfo.dBaseSize = + constInfo_.headDim / min(FloorPow2(aiCoreNum), CeilPow2(CeilDivT(aiCoreNum, constInfo_.batchSize))); + if (constInfo_.kBaseNum > 1) { + splitInfo.dBaseSize = max(splitInfo.dBaseSize, FP32_REPEAT_ELEMENT_NUM); + } + splitInfo.vec1GroupSize = constInfo_.headDim / splitInfo.dBaseSize; + splitInfo.vec1GroupNum = min(static_cast(aiCoreNum / splitInfo.vec1GroupSize), constInfo_.batchSize); +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::CalcTaskDistribution(Vec1SplitInfo &splitInfo) +{ + uint32_t blockIdx = GetBlockIdx(); + uint32_t groupSize = splitInfo.vec1GroupSize; + uint32_t groupNum = splitInfo.vec1GroupNum; + uint32_t totalDealBatchNum = constInfo_.batchSize; + + if (blockIdx < groupSize * (totalDealBatchNum % groupNum)) { + splitInfo.dealBatchNum = totalDealBatchNum / groupNum + 1; + splitInfo.preDealBatchNum = splitInfo.dealBatchNum * (blockIdx / groupSize); + } else if (blockIdx < groupSize * groupNum) { + splitInfo.dealBatchNum = totalDealBatchNum / groupNum; + splitInfo.preDealBatchNum = splitInfo.dealBatchNum * (blockIdx / groupSize) + totalDealBatchNum % groupNum; + } else { + splitInfo.dealBatchNum = 0; + splitInfo.preDealBatchNum = totalDealBatchNum; + } +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::UpdateIteratorState(Vec1SplitInfo &splitInfo) +{ + splitInfo.preCompressedCnt = 0; + splitInfo.dealSeqStartIdx = splitInfo.preDealBatchNum * constInfo_.sSize; + splitInfo.curBStart = splitInfo.preDealBatchNum; + splitInfo.dealSeqCnt = splitInfo.dealBatchNum * constInfo_.sSize; + splitInfo.curSStart = 0; + totalCompressedCnt_ = 0; + uint32_t endB = splitInfo.preDealBatchNum + splitInfo.dealBatchNum; + for (uint32_t curB = 0; curB < constInfo_.batchSize; curB++) { + uint32_t startPos = GetStartPos(curB); + uint32_t seqLength = GetSeqLength(curB); + if (curB < splitInfo.curBStart) { + splitInfo.preCompressedCnt += (startPos + seqLength) / cmpRatio_ - startPos / cmpRatio_; + } else { + totalCompressedCnt_ += (startPos + seqLength) / cmpRatio_ - startPos / cmpRatio_; + } + } + totalCompressedCnt_ += splitInfo.preCompressedCnt; +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::CalcTilingStrategy(Vec1SplitInfo &splitInfo) +{ + // 计算headDim和Tc方向切分大小 + uint32_t maxDealColNum = BUFFER_SIZE_BYTE_32K / (cmpRatio_ * coff_ * sizeof(T)); + + // 切块逻辑 + if (maxDealColNum < splitInfo.dBaseSize) { + splitInfo.tcSplitSize = 1; + splitInfo.dLoopCount = CeilDivT(splitInfo.dBaseSize, maxDealColNum); + splitInfo.dSplitSize = splitInfo.dBaseSize / splitInfo.dLoopCount; + } else { + splitInfo.dSplitSize = splitInfo.dBaseSize; + splitInfo.dLoopCount = splitInfo.dBaseSize / splitInfo.dSplitSize; // 此处常等于1，保留原逻辑 + splitInfo.tcSplitSize = maxDealColNum / splitInfo.dBaseSize; + } +} + + +template +__aicore__ inline Vec1SplitInfo CompressorBlockVectorFullLoad::SplitCoreV1() +{ + Vec1SplitInfo splitInfo; + + // 1. 计算基础分组和分片大小 + CalcGroupInfo(splitInfo); + + // 2. 根据当前的 BlockIdx 计算任务分配（负载均衡） + CalcTaskDistribution(splitInfo); + + // 3. 刷新迭代器并获取当前核的起始位置状态 + UpdateIteratorState(splitInfo); + + if (splitInfo.dealBatchNum == 0) { + return splitInfo; + } + + // 4. 计算具体在内存中的切块（Tiling）逻辑 + CalcTilingStrategy(splitInfo); + + return splitInfo; +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::ComputeVec1() +{ + Vec1SplitInfo splitInfo = SplitCoreV1(); + // 计算当前VecCore的任务量 + if (splitInfo.dealBatchNum == 0) { + return; + } + + LoopInfo loopInfo; + loopInfo.groupSize = splitInfo.vec1GroupSize; + loopInfo.groupNum = splitInfo.vec1GroupNum; + loopInfo.coreRowIdx = GetBlockIdx() / splitInfo.vec1GroupSize; + loopInfo.coreColIdx = GetBlockIdx() % splitInfo.vec1GroupSize; + loopInfo.isCoreRowLast = loopInfo.coreRowIdx == splitInfo.vec1GroupNum - 1; + loopInfo.isCoreRowFirst = loopInfo.coreRowIdx == 0; + + + CompressorVec1SliceIterator sliceIterator(tools_); + sliceIterator.SetMaxBatchSize(constInfo_.batchSize); + // 切块循环 + uint64_t baseOffset = loopInfo.coreColIdx * splitInfo.dBaseSize; + + + uint32_t cnt = constInfo_.sSize * splitInfo.dBaseSize * coff_; + uint32_t singleLoopBatchNum = BUFFER_SIZE_BYTE_16K / (cnt * sizeof(T)); + uint32_t loopTimes = CeilDivT(splitInfo.dealBatchNum, singleLoopBatchNum); + bool isApeFullLoad = coff_ * cmpRatio_ * splitInfo.dBaseSize * sizeof(T) <= BUFFER_SIZE_BYTE_16K; + if (isApeFullLoad) { + CopyInApe(baseOffset, splitInfo.dBaseSize); + } + for (uint32_t idx = 0; idx < loopTimes; idx++) { + uint32_t curLoopBatchNum = min(singleLoopBatchNum, splitInfo.dealBatchNum - singleLoopBatchNum * idx); + scoreUb = inputQue1.AllocTensor(); + kvUb = scoreUb[BUFFER_SIZE_BYTE_16K / sizeof(T)]; + FromWokrSpaceToUb(scoreUb, scoreMm1ResGm_, splitInfo.dealSeqStartIdx, curLoopBatchNum * constInfo_.sSize, + baseOffset, splitInfo.dBaseSize); + FromWokrSpaceToUb(kvUb, kvMm1ResGm_, splitInfo.dealSeqStartIdx, curLoopBatchNum * constInfo_.sSize, baseOffset, + splitInfo.dBaseSize); + inputQue1.EnQue(scoreUb); + inputQue1.DeQue(); + splitInfo.dealTcNum = 0; + uint32_t curLoopCompressedCnt = 0; + for (uint32_t curB = splitInfo.curBStart; curB < splitInfo.curBStart + curLoopBatchNum; curB++) { + uint32_t startPos = GetStartPos(curB); + uint32_t seqLength = GetSeqLength(curB); + splitInfo.dealTcNum += + CeilDivT(startPos + seqLength, cmpRatio_) - (startPos / cmpRatio_); + curLoopCompressedCnt += (startPos + seqLength) / cmpRatio_ - startPos / cmpRatio_; + } + sliceIterator.Reset(splitInfo.curBStart, splitInfo.curSStart, 0U, 0U); + sliceIterator.SetNeedDealTcSize(splitInfo.dealTcNum); + sliceIterator.SetDealedTcCnt(0U); + Vec1SliceInfo &sliceInfo = sliceIterator.GetSlice(); + while (!sliceIterator.IsEnd()) { + sliceIterator.GetSlice(); + SaveState(kvUb, stateCacheGm_, stateBlockTableGm_, sliceInfo, baseOffset, splitInfo.dBaseSize, + splitInfo.dBaseSize, kvStateIdx_); + + AddApeToScore(scoreUb, sliceInfo, splitInfo.dBaseSize, splitInfo.dBaseSize, baseOffset, isApeFullLoad); + SaveState(scoreUb, stateCacheGm_, stateBlockTableGm_, sliceInfo, baseOffset, splitInfo.dBaseSize, + splitInfo.dBaseSize, scoreStateIdx_); + sliceIterator.IteratorSlice(); + } + + if (curLoopCompressedCnt == 0) { + inputQue1.FreeTensor(scoreUb); + continue; + } + for (uint32_t dLoopIdx = 0; dLoopIdx < splitInfo.dLoopCount; dLoopIdx++) { + uint64_t dBaseOffset = baseOffset + dLoopIdx * splitInfo.dSplitSize; + loopInfo.dLoopIdx = dLoopIdx; + + sliceIterator.Reset(splitInfo.curBStart, splitInfo.curSStart, 0U, 0U); + compressedCnt_ = splitInfo.preCompressedCnt; + for (uint32_t tcIdx = 0; tcIdx < splitInfo.dealTcNum; tcIdx += splitInfo.tcSplitSize) { + uint32_t actDealTcSize = min(splitInfo.tcSplitSize, splitInfo.dealTcNum - tcIdx); + + loopInfo.isCoreLoopFirst = tcIdx == 0; + loopInfo.isCoreLoopLast = tcIdx + splitInfo.tcSplitSize >= splitInfo.dealTcNum; + // 处理单个切块 + sliceIterator.SetNeedDealTcSize(actDealTcSize); + sliceIterator.SetDealedTcCnt(0U); + DealVec1BaseBlock(sliceIterator, loopInfo, baseOffset, dLoopIdx * splitInfo.dSplitSize, + splitInfo.dSplitSize, splitInfo.dBaseSize, splitInfo.dealSeqStartIdx); + } + } + inputQue1.FreeTensor(scoreUb); + splitInfo.curBStart += curLoopBatchNum; + splitInfo.dealSeqStartIdx += curLoopBatchNum * constInfo_.sSize; + splitInfo.preCompressedCnt += curLoopCompressedCnt; + } +} + + +template +__aicore__ inline Vec2SplitInfo CompressorBlockVectorFullLoad::SplitCoreV2() +{ + Vec2SplitInfo splitInfo; + + uint32_t blockIdx = GetBlockIdx(); + uint32_t aiCoreNum = constInfo_.usedCoreNum * 2; + + if (blockIdx < (totalCompressedCnt_ % aiCoreNum)) { + splitInfo.dealScNum = totalCompressedCnt_ / aiCoreNum + 1; + splitInfo.preScCnt = splitInfo.dealScNum * blockIdx; + } else if (blockIdx < aiCoreNum) { + splitInfo.dealScNum = totalCompressedCnt_ / aiCoreNum; + splitInfo.preScCnt = splitInfo.dealScNum * blockIdx + totalCompressedCnt_ % aiCoreNum; + } else { + splitInfo.dealScNum = 0; + splitInfo.preScCnt = totalCompressedCnt_; + } + + if (splitInfo.dealScNum == 0) { + return splitInfo; + } + + uint32_t preScCnt = splitInfo.preScCnt; + for (uint32_t curB = 0; curB < constInfo_.batchSize; curB++) { + uint32_t startPos = GetStartPos(curB); + uint32_t seqUsed = GetSeqUsed(curB); + uint32_t curScNum = (startPos + seqUsed) / cmpRatio_ - startPos / cmpRatio_; + if (preScCnt < curScNum) { + splitInfo.curBStart = curB; + splitInfo.curScStart = preScCnt; + break; + } + preScCnt -= curScNum; + } + + splitInfo.dealedScCnt = 0; + + return splitInfo; +} + +template +__aicore__ inline void CompressorBlockVectorFullLoad::ComputeVec2() +{ + Vec2SplitInfo splitInfo = SplitCoreV2(); + if (splitInfo.dealScNum == 0) { + return; + } + CompressorVec2SliceIterator sliceIterator(tools_); + sliceIterator.SetMaxBatchSize(constInfo_.batchSize); + sliceIterator.Reset(splitInfo.curBStart, splitInfo.curScStart, splitInfo.dealedScCnt + splitInfo.preScCnt); + + uint32_t singleLoopScNum = BUFFER_SIZE_BYTE_32K / (constInfo_.headDim * sizeof(T)); + uint32_t loopTimes = CeilDivT(splitInfo.dealScNum, singleLoopScNum); + + for (uint32_t idx = 0; idx < loopTimes; idx++) { + uint32_t curLoopScNum = min(singleLoopScNum, splitInfo.dealScNum - singleLoopScNum * idx); + sliceIterator.SetNeedDealScSize(curLoopScNum); + DealVec2BaseBlock(splitInfo, sliceIterator); + splitInfo.preScCnt += curLoopScNum; + } +} + + +} // namespace Compressor +#endif // COMPRESSOR_BLOCK_VECTOR_PREF_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_comm.h b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_comm.h new file mode 100644 index 000000000..636c5760b --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_comm.h @@ -0,0 +1,445 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_comm.h + * \brief + */ + +#ifndef COMPRESSOR_COMM_H +#define COMPRESSOR_COMM_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" + +using namespace AscendC; + +namespace Compressor { +template +__aicore__ inline T CeilDivT(T num1, T num2) +{ + if (num2 == 0) { + return static_cast(0); + } + return (num1 + num2 - 1) / num2; +} + +template +__aicore__ inline T Align(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd) * (rnd))); +} + +template +__aicore__ inline T Trunc(T num, T rnd) +{ + return ((rnd) == 0) ? 0 : (((num) / (rnd) * (rnd))); +} + +template +__aicore__ inline T FloorPow2(T num) +{ + if (num == 0) + return 1; + for (uint32_t i = 1; i < sizeof(T) * 8; i <<= 1) { + num |= (num >> i); + } + return num - (num >> 1); +} + +template +__aicore__ inline T CeilPow2(T num) +{ + if (num <= 1) + return 1; + num--; + for (uint32_t i = 1; i < sizeof(T) * 8; i <<= 1) { + num |= (num >> i); + } + num++; + return num; +} + +enum class X_LAYOUT : std::uint8_t { + BSH = static_cast(0), + TH = static_cast(1) +}; + +enum class X_DTYPE : std::uint8_t { + BF16 = static_cast(0), + FP16 = static_cast(1) +}; + +enum class COFF : std::uint8_t { + DISABLE = static_cast(1), + OVERLAP = static_cast(2) +}; + +enum class ROTARY_MODE : std::uint8_t { + HALF = static_cast(1), + INTERLEAVE = static_cast(2) +}; + +enum class CACHE_MODE : std::uint8_t { + CONTINUOUS = static_cast(1), + CYCLE = static_cast(2) +}; + +enum class TEMPLATE_ID : uint8_t { + NORMAL = 0, + EMPTY_X = 1, + FULL_LOAD = 2 +}; + +template +struct COMPType { + static constexpr X_LAYOUT xLayout = X_L; + static constexpr X_DTYPE xDtype = X_T; + static constexpr COFF coff = C; + static constexpr ROTARY_MODE rotaryMode = Rotary_Mode; + static constexpr CACHE_MODE cacheMode = Cache_Mode; +}; + +struct CmpBlockInfo { + __aicore__ inline CmpBlockInfo(){}; + __aicore__ inline CmpBlockInfo(uint32_t bIdx, uint32_t sIdx, bool needReset = false) + : bIdx(bIdx), sIdx(sIdx), needReset(needReset){}; + + uint32_t bIdx = 0U; + uint32_t sIdx = 0U; + uint32_t bSeqUsed = 0U; + uint32_t bStartPos = 0U; + bool needReset = false; + bool isFirst = true; + + uint32_t headSeqCnt = 0U; + uint32_t validSeqCnt = 0U; + uint32_t tailSeqCnt = 0U; + bool isCompress = 0U; +}; + +struct BasicBlockInfo { + uint32_t bIdx = 0; + uint32_t sIdx = 0; + uint32_t compressedTcNum = 0; + uint32_t dealSeqCnt = 0; + uint32_t dealTcNum = 0; +}; + +struct BatchInfo { + uint32_t tcNum = 0; + uint32_t compressedTcNum = 0; + uint32_t remSeqCnt = 0; + uint32_t seqCnt = 0; + uint32_t seqUsedCnt = 0; + uint32_t headHolderSeq = 0; + uint32_t bStartPos = 0; + uint32_t bIdx = 0; + uint32_t sIdx = 0; +}; + +struct ConstInfo { + // 整个AICORE的任务信息, 左闭右开区间[ (bStart, s2Start), (bEnd, s2End) ) + uint32_t bStart = 0U; + uint32_t sStart = 0U; + uint32_t bEnd = 0U; + uint32_t sEnd = 0U; + + // 分核相关 + uint32_t usedCoreNum = 0; + uint32_t dBaseSize = 0; + uint32_t mBaseSize = 0; + uint32_t kBaseSize = 0; + uint32_t kBaseNum = 0; + uint32_t tcSize = 0; + uint32_t tcBaseSize = 0; + uint32_t tcBasicBlockNum = 0; + uint32_t dBasicBlockNum = 0; + uint32_t coreGroupNum = 0; + uint32_t singleCoreDealTcBasicNum = 0; + uint32_t dIdx = 0; + uint32_t mStart = 0; + uint32_t mEnd = 0; + uint32_t nStart = 0; + uint32_t nEnd = 0; + uint32_t kStart = 0; + uint32_t kEnd = 0; + uint32_t mLoopNum = 0; + uint32_t bIdxOfLastTc = 0; + uint32_t sIdxOfLastTc = 0; + uint32_t mGroupNum = 0; + uint32_t mCurGroupIdx = 0; + + // shape及参数 + uint32_t batchSize = 0; + uint32_t hSize = 0; + uint32_t sSize = 0; + uint32_t headDim = 0; + uint32_t ropeHeadDim = 0; + uint32_t cmpRatio = 0; + float normEps = 1e-6; + float reciprocalD = 0; + uint64_t stateCacheStrideDim0 = 0; + + uint32_t curGroupIdx = 0; + uint32_t tailGroupIdx = 0; + uint32_t tailBasicBlockNum = 0; + uint32_t realDealBasicBlockNum = 0; + + // pageAttention + uint32_t blockNum = 0; + uint32_t blockSize = 0; + uint32_t maxBlockNumPerBatch = 0; + + // workSpace + uint32_t dbWorkspaceRatio = 1; + uint32_t mm1KvResSize = 0; + uint32_t mm1ScoreResSize = 0; + uint32_t vec1TailCacheSize = 0; + uint32_t vec1ResSize = 0; + uint32_t mm1ResSize = 0; // 所有cube输出kv/score结果的总大小 + + uint32_t aiCoreIdx = 0; + uint32_t nSize = 0; + + uint32_t dbSize = 0; +}; + +struct RunInfo { + bool isValid = false; + uint32_t cubeDbIdx = 0; // kernel主循环索引 + + // 增加字段 + uint32_t dealTcNum = 0; + // 右边相关信息 + uint32_t bStart = 0; + uint32_t sStart = 0; + uint32_t dealSeqCnt = 0; + // 左边相关信息 + uint32_t preBStart = 0; + uint32_t preSStart = 0; + uint32_t preDealSeqCnt = 0; // 左边需要处理的s大小 + uint32_t preFirstSeqCnt = 0; // 左边首块大小 + + uint32_t kStartIdx = 0; + uint32_t dealKSize = 0; + uint32_t hStart = 0; + + uint32_t bEnd = 0; + uint32_t sEnd = 0; + uint32_t bStartSeqIdx = 0; + uint32_t bEndSeqIdx = 0; + + // v2分核信息 sc是左闭右开 + uint32_t scStart = 0; + uint32_t scEnd = 0; + uint32_t dealScSize = 0; + + // vec1Res offset + uint64_t vec1ResOffset = 0; +}; + +struct Vec1RunInfo { + // vec相关信息，一次syncAll需处理数据的起始索引 + bool resetResFlag = false; // v1积攒N轮 是否是N轮的起始轮 + uint32_t c1v1DbIdx = 0; // vec1 doubleBuffer索引 + uint32_t v1v2DbIdx = 0; // v1v2 doubleBuffer索引 + uint32_t bStart = 0; + uint32_t sStart = 0; + uint32_t dealTcNum = 0; + uint32_t dealScSize = 0; +}; + +struct Vec2RunInfo { + // uint32_t bStart = 0; + uint32_t v2DbIdx = 0; // v2 doubleBuffer索引 + uint32_t sStart = 0; + uint32_t bEnd = 0; + uint32_t sEnd = 0; + // v2分核信息 sc是左闭右开 + uint32_t scStart = 0; + uint32_t scEnd = 0; + + // 增加字段 + uint32_t bStart = 0; + uint32_t compressedId = 0; + uint32_t bCompressedId = 0; + uint32_t dealScSize = 0; +}; + +struct MSplitInfo { + uint32_t vecStartB = 0U; + uint32_t vecStartS = 0U; + uint32_t vecEndB = 0U; + uint32_t vecEndS = 0U; + uint32_t dealTcNum = 0U; + // vec1Res offset + uint64_t vec1StartOffset = 0; + uint64_t vec1ResOffset = 0; +}; + +struct BlockInfo { + __aicore__ inline BlockInfo(uint32_t bIdx, uint32_t sIdx, uint32_t dealSeqSize) + : bIdx(bIdx), sIdx(sIdx), dealSeqSize(dealSeqSize){}; + uint32_t bIdx = 0U; + uint32_t sIdx = 0U; + uint32_t dealSeqSize = 0; + + uint32_t isFirst = true; + uint32_t bSeqUsed = 0U; + uint32_t bStartPos = 0U; + uint32_t headHolderSeqCnt = 0U; + uint32_t validSeqCnt = 0U; + uint32_t tailHolderSeqCnt = 0U; + uint32_t dealTcSize = 0U; + uint32_t tailValidSeqCnt = 0U; + uint32_t compressTcSize = 0U; +}; + +struct LoopInfo { + uint32_t groupSize = 0U; + uint32_t groupNum = 0U; + uint32_t coreRowIdx = 0U; + uint32_t coreColIdx = 0U; + uint32_t dLoopIdx = 0U; + bool isCoreRowFirst = false; + bool isCoreRowLast = false; + bool isCoreLoopFirst = false; + bool isCoreLoopLast = false; +}; + +struct Vec1SplitInfo { + uint32_t dealSeqStartIdx = 0; + uint32_t dealSeqCnt = 0; + uint32_t dBaseSize = 0; + uint32_t vec1GroupSize = 0; + uint32_t vec1GroupNum = 0; + uint32_t dealTcSize = 0; + uint32_t dealTcNum = 0; + uint32_t dealBatchNum = 0; + uint32_t preDealTcSize = 0; + uint32_t preDealBatchNum = 0; + uint32_t curBStart = 0; + uint32_t curSStart = 0; + uint32_t curCompressedCnt = 0; + uint32_t preCompressedCnt = 0; + uint32_t totalCompressedCnt = 0; + uint32_t tcSplitSize = 0; + uint32_t dSplitSize = 0; + uint32_t dLoopCount = 0; +}; + +struct Vec2SplitInfo { + uint32_t dealedScCnt = 0; + uint32_t preScCnt = 0; + uint32_t dealScNum = 0; + uint32_t curBStart = 0; + uint32_t curScStart = 0; +}; + +// BUFFER的字节数 +inline constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; +inline constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64; +inline constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256; +inline constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512; +inline constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024; +inline constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048; +inline constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096; +inline constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192; +inline constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384; +inline constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768; +inline constexpr uint32_t BUFFER_SIZE_BYTE_64K = 65536; + +// BLOCK和REPEAT的字节数 +inline constexpr uint64_t BYTE_BLOCK = 32UL; +inline constexpr uint32_t REPEAT_BLOCK_BYTE = 256U; + +template +__aicore__ inline constexpr T BlockElementNum() +{ + return BYTE_BLOCK / sizeof(T); +} + +template +__aicore__ inline constexpr T RepeatElementNum() +{ + return REPEAT_BLOCK_BYTE / sizeof(T); +} +// BLOCK和REPEAT的FP32元素数 +inline constexpr uint32_t FP32_BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(float); // 8 +inline constexpr uint32_t FP16_BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(bfloat16_t); // 16 +inline constexpr uint32_t FP32_REPEAT_ELEMENT_NUM = REPEAT_BLOCK_BYTE / sizeof(float); // 64 +inline constexpr uint32_t REPEAT_STRIDE_NUM = REPEAT_BLOCK_BYTE / BYTE_BLOCK; // 8 +inline constexpr uint32_t REPEAT_MAX_NUM = 255; +inline constexpr uint32_t BRCB_NUM = 8; +inline constexpr uint32_t MAX_R = 256; + +template +__aicore__ inline void CopySingleMatrixNDToNZ(LocalTensor l1Tensor, const GlobalTensor gmTensor, uint32_t nValue, + uint32_t dValue, uint32_t srcDValue, uint32_t dstNzC0Stride) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nValue; // nd矩阵的行数 + if constexpr (IsSameType::value) { + constexpr uint32_t HALF_SIZE_DIVISOR = 2; + nd2nzPara.dValue = dValue / HALF_SIZE_DIVISOR; + nd2nzPara.srcDValue = srcDValue / HALF_SIZE_DIVISOR; + } else { + nd2nzPara.dValue = dValue; // nd矩阵的列数 + nd2nzPara.srcDValue = srcDValue; // 同一nd矩阵相邻行起始地址间的偏移 + } + nd2nzPara.dstNzC0Stride = dstNzC0Stride; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(l1Tensor, gmTensor, nd2nzPara); +} +template +__aicore__ inline void DumpTensorForDim2(GlobalTensor tensor, uint32_t desc, uint32_t dumpSize, uint32_t row, + uint32_t col) +{ + uint32_t array2[] = {static_cast(row), static_cast(col)}; + AscendC::ShapeInfo shapeInfo(2, array2); + // AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo); +} + +template +__aicore__ inline void DumpTensorForDim2(LocalTensor tensor, uint32_t desc, uint32_t dumpSize, uint32_t row, + uint32_t col) +{ + uint32_t array2[] = {static_cast(row), static_cast(col)}; + AscendC::ShapeInfo shapeInfo(2, array2); + // AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo); +} + +template +__aicore__ inline void DumpTensorForDim2(LocalTensor tensor, uint32_t desc, uint32_t dumpSize) +{ + uint32_t col = 32 / sizeof(T); + uint32_t array2[] = {static_cast(dumpSize / col), static_cast(col)}; + AscendC::ShapeInfo shapeInfo(2, array2); + // AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo); +} + +template +__aicore__ inline void DumpTensorForDim2(GlobalTensor tensor, uint32_t desc, uint32_t dumpSize) +{ + uint32_t col = 32 / sizeof(T); + uint32_t array2[] = {static_cast(dumpSize / col), static_cast(col)}; + AscendC::ShapeInfo shapeInfo(2, array2); + // AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo); +} + +} // namespace Compressor +#endif diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_kernel.h b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_kernel.h new file mode 100644 index 000000000..3ef9a7302 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_kernel.h @@ -0,0 +1,628 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_kernel.h + * \brief + */ + +#ifndef COMPRESSOR_KERNEL_H +#define COMPRESSOR_KERNEL_H + +#include "compressor_comm.h" +#include "compressor_template_tiling_key.h" +#include "compressor_tiling_data.h" +#include "compressor_tools.h" +#include "compressor_block_cube.h" +#include "compressor_block_vec.h" + +using namespace AscendC; + +namespace Compressor { + +template +class CompressorKernel { +public: + __aicore__ inline CompressorKernel(TPipe* pipe, const optiling::CompressorTilingData* __restrict tilingData) + : pipe_(pipe), tilingData_(tilingData) {} + + __aicore__ inline void Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut, + __gm__ uint8_t *workspace); + __aicore__ inline void Process(); + +private: + // ================================Init functions================================== + __aicore__ inline void InitWorkspace(__gm__ uint8_t *workspace); + // ================================Process functions================================ + __aicore__ inline void InitTilingData(); + __aicore__ inline void SplitK(); + // 获取基本块数量 + __aicore__ inline uint32_t GetLoopTimes(); + __aicore__ inline void SkipInvalidBatch(BatchInfo &batchInfo); + __aicore__ inline void UpdateCurGroup(BasicBlockInfo &basicBlockInfo, BatchInfo batchInfo, uint32_t &curGroupQuota, uint32_t curDealSeq); + __aicore__ inline BasicBlockInfo SkipOneLoop(BatchInfo &batchInfo); + // 计算分核基本信息 + __aicore__ inline void CalcSplitCoreInfo(); + + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void ComputeMm1(const RunInfo &info, bool isNeedExcute); + __aicore__ inline void ComputeVec1(const Vec1RunInfo &info); + __aicore__ inline void ComputeVec2(const Vec2RunInfo &info); + + __aicore__ inline bool IsNeedExcuteC1(RunInfo info); + __aicore__ inline bool IsNeedSyncAll(uint32_t curBasicBlockIdx); + __aicore__ inline void CalcC1V1Params(RunInfo &info, Vec1RunInfo &vec1Info, BatchInfo &batchInfo, uint32_t loopIdx); + __aicore__ inline void UpdateVec2Info(Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info); + __aicore__ inline bool IsNeedExcuteV2(Vec2RunInfo &vec2Info); + + using X_T = typename AscendC::Conditional::type; + using T = float; + using MM1_OUT_T = T; + using VEC1_OUT_T = T; + + // 常量 + static constexpr uint64_t SYNC_MODE0 = 0; + static constexpr uint64_t SYNC_MODE2 = 2; + static constexpr uint32_t SYNC_C1_FLAG = 3; + static constexpr uint32_t SYNC_V1_FLAG = 4; + static constexpr uint32_t SYNC_V1_FLAG2 = 5; + static constexpr uint32_t SYNC_C1_V1_FLAG = 7; + static constexpr uint32_t SYNC_V1_C1_FLAG = 9; + + // ==============================TilingData&TPipe============================== + TPipe* pipe_; + const optiling::CompressorTilingData* __restrict tilingData_; + // ===========================Workspace Global Tensor=========================== + GlobalTensor mm1KvResGm; + GlobalTensor mm1ScoreResGm; + GlobalTensor vec1KvCacheGm; + GlobalTensor vec1ScoreCacheGm; + GlobalTensor Vec1InputKvGm; + GlobalTensor Vec1InputScoreGm; + GlobalTensor vec1ResGm; + GlobalTensor vec2InputGm; + // ================================Task Info==================================== + CompressorTools tools_; + ConstInfo constInfo{}; + uint32_t aiCoreIdx = 0; + + // ==============================Service Define============================== + CompressorBlockCube blockCube_; + CompressorBlockVector blockVec_; + + uint32_t allCompressedTcNum_ = 0; + uint32_t curCompressedTcNum_ = 0; + uint32_t accDealSize = 0; + uint32_t loopTimes = 0; + uint32_t cubeLoop = 0; + uint32_t vec1Loop = 0; + uint32_t vec2Loop = 0; + uint32_t kStartIdx_ = 0; + uint32_t dealKSize_ = 0; + uint32_t hStart_ = 0; + bool isFirstUpdateCurGroup = true; +}; + +template +__aicore__ inline void CompressorKernel::Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut, + __gm__ uint8_t *workspace) +{ + if ASCEND_IS_AIV { + constInfo.aiCoreIdx = GetBlockIdx() / 2; + } else { + constInfo.aiCoreIdx = GetBlockIdx(); + } + InitTilingData(); + // init tools + tools_.toolParams_.seqSize = tilingData_->baseParams.seqSize; + tools_.toolParams_.cmpRatio = tilingData_->baseParams.cmpRatio; + tools_.Init(startPos, seqUsed, cuSeqlens); + + // 剔除尾部的无效batch + for (; constInfo.batchSize > 0; --constInfo.batchSize) { + uint32_t bSeqUsed = tools_.GetSeqLength(constInfo.batchSize - 1); + if (bSeqUsed > 0) { + break; + } + } + + // 所有batch的有效序列都为0时, 直接退出 + if (constInfo.batchSize == 0) { + return; + } + + // 0. 计算最后一个Tc块的起始位置 + constInfo.bIdxOfLastTc = constInfo.batchSize - 1; + // 1. 计算head_dim的切分大小, 构建ConstInfo的其他信息 + CalcSplitCoreInfo(); + SplitK(); // K轴切分 + // 2. 计算循环次数 + loopTimes = GetLoopTimes(); + // 3. 初始化workspace + InitWorkspace(workspace); + // 4. 初始化block层 + if ASCEND_IS_AIC { + blockCube_.InitParams(constInfo, tools_); + blockCube_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos, + stateBlockTable, cuSeqlens, seqUsed, startPos, cmpKvOut); + blockCube_.InitBuffers(pipe_); + blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm); + } else { + blockVec_.InitParams(constInfo, tools_); + blockVec_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos, stateBlockTable, + cuSeqlens, seqUsed, startPos, cmpKvOut); + blockVec_.InitBuffers(pipe_); + blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm, vec1ResGm, vec2InputGm); + } +} + +template +__aicore__ inline void CompressorKernel::InitTilingData() { + constInfo.cmpRatio = tilingData_->baseParams.cmpRatio; + constInfo.batchSize = tilingData_->baseParams.batchSize; + constInfo.mBaseSize = tilingData_->innerSplitParams.mBaseSize; + constInfo.dBaseSize = tilingData_->innerSplitParams.dBaseSize; + constInfo.kBaseSize = tilingData_->baseParams.hiddenSize; + constInfo.kBaseNum = 1; + constInfo.headDim = tilingData_->baseParams.headDim; + constInfo.hSize = tilingData_->baseParams.hiddenSize; + constInfo.sSize = tilingData_->baseParams.seqSize; + constInfo.ropeHeadDim = tilingData_->baseParams.ropeHeadDim; + constInfo.normEps = tilingData_->baseParams.normEps; + constInfo.stateCacheStrideDim0 = tilingData_->baseParams.stateCacheStrideDim0; + constInfo.reciprocalD = tilingData_->baseParams.reciprocalD; + constInfo.usedCoreNum = tilingData_->baseParams.usedCoreNum; + + constInfo.blockNum = tilingData_->pageAttentionParams.blockNum; + constInfo.blockSize = tilingData_->pageAttentionParams.blockSize; + constInfo.maxBlockNumPerBatch = tilingData_->pageAttentionParams.maxBlockNumPerBatch; + + constInfo.nSize = tilingData_->baseParams.nSize; + constInfo.vec1TailCacheSize = tilingData_->workspaceParams.vec1TailCacheSize; + constInfo.dbWorkspaceRatio = tilingData_->workspaceParams.dbWorkspaceRatio; +} + +template +__aicore__ inline void CompressorKernel::SplitK() +{ + uint32_t mSize = 0; + for (uint32_t i = 0; i < constInfo.batchSize; i++) { + uint32_t bSeqUsed = tools_.GetSeqLength(i); + // 获取m大小 + mSize += bSeqUsed; + } + + uint32_t mBaseNum = CeilDivT(mSize, constInfo.mBaseSize); + if (constInfo.dBasicBlockNum * mBaseNum < constInfo.usedCoreNum) { + constInfo.kBaseNum = constInfo.usedCoreNum / constInfo.dBasicBlockNum; + uint32_t kAlignSize = + CeilDivT(Align(constInfo.hSize, static_cast(BUFFER_SIZE_BYTE_32B / sizeof(X_T))), + constInfo.kBaseNum); + constInfo.kBaseSize = Trunc(kAlignSize, static_cast(BUFFER_SIZE_BYTE_32B / sizeof(X_T))); + // 当切m轴无法满足开满核时，不切m轴(切m处理有点复杂) + constInfo.mGroupNum = 1; // 在m轴处理上所有核当一个组 + constInfo.mCurGroupIdx = 0; // 只有一个组 + } + // 每轮固定不变，预计算后主循环直接复用 + if (constInfo.kBaseNum > 1) { + kStartIdx_ = constInfo.aiCoreIdx / constInfo.dBasicBlockNum; + if (constInfo.curGroupIdx + 1 < constInfo.coreGroupNum) { + dealKSize_ = constInfo.kBaseSize; + hStart_ = kStartIdx_ * dealKSize_; + } else { + dealKSize_ = kStartIdx_ < constInfo.coreGroupNum ? + constInfo.hSize - kStartIdx_ * constInfo.kBaseSize : 0; + hStart_ = kStartIdx_ * constInfo.kBaseSize; + } + } else { + kStartIdx_ = 0; + dealKSize_ = constInfo.hSize; + hStart_ = kStartIdx_ * dealKSize_; + } + } + +template +__aicore__ inline void CompressorKernel::SkipInvalidBatch(BatchInfo &batchInfo) +{ + for (; batchInfo.bIdx < constInfo.batchSize; ++batchInfo.bIdx) { + batchInfo.seqCnt = tools_.GetSeqLength(batchInfo.bIdx); + if (batchInfo.seqCnt > 0) { + break; + } + } + batchInfo.remSeqCnt = batchInfo.seqCnt; + if (tools_.isExistSeqUsed_) { + batchInfo.seqUsedCnt = tools_.GetSeqUsed(batchInfo.bIdx); + } else { + batchInfo.seqUsedCnt = batchInfo.seqCnt; + } + if (batchInfo.bIdx < constInfo.batchSize) { + batchInfo.bStartPos = tools_.GetStartPos(batchInfo.bIdx); + batchInfo.sIdx = 0; + batchInfo.headHolderSeq = batchInfo.bStartPos & (constInfo.cmpRatio - 1); + batchInfo.tcNum = (batchInfo.bStartPos + batchInfo.seqCnt + constInfo.cmpRatio - 1) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio; + batchInfo.compressedTcNum = (batchInfo.bStartPos + batchInfo.seqUsedCnt) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio; + } +} + + +template +__aicore__ inline void CompressorKernel::UpdateCurGroup(BasicBlockInfo &basicBlockInfo, + BatchInfo batchInfo, uint32_t &curGroupQuota, uint32_t curDealSeq) +{ + // 更新当前组的信息 + if (curGroupQuota == 0 && !isFirstUpdateCurGroup) { + return; + } + isFirstUpdateCurGroup = false; + basicBlockInfo.bIdx = batchInfo.bIdx; + uint32_t curGroupDealSeq = curGroupQuota < curDealSeq ? curGroupQuota : curDealSeq; + basicBlockInfo.sIdx = batchInfo.sIdx + curGroupDealSeq; + basicBlockInfo.dealSeqCnt += curGroupDealSeq; + curGroupQuota -= curGroupDealSeq; + // 结尾需要跳batch，需要考虑在当前组起始为末尾，或者当前组起始大于整个M轴 + if ((curGroupQuota == 0 || basicBlockInfo.bIdx == constInfo.batchSize - 1) && basicBlockInfo.sIdx == batchInfo.seqCnt) { + basicBlockInfo.sIdx = 0; + for (basicBlockInfo.bIdx++; basicBlockInfo.bIdx < constInfo.batchSize; ++basicBlockInfo.bIdx) { + uint32_t seqCnt = tools_.GetSeqLength(basicBlockInfo.bIdx); + if (seqCnt > 0) { + break; + } + } + } +} + +template +__aicore__ inline BasicBlockInfo CompressorKernel::SkipOneLoop(BatchInfo &batchInfo) +{ + BasicBlockInfo basicBlockInfo{}; + isFirstUpdateCurGroup = true; + uint32_t curGroupQuota = constInfo.mBaseSize * constInfo.mCurGroupIdx; // m轴当前组起始 + bool curGroupStartFlag = false; + uint32_t quota = constInfo.mGroupNum * constInfo.mBaseSize; + + for (; batchInfo.bIdx < constInfo.batchSize;) { + uint32_t curDealSeq = 0; + uint32_t curDealTcNum = 0; + uint32_t curDealCompressedTcNum = 0; + // 无法处理完当前整个batch + if (quota < batchInfo.remSeqCnt) { + // 向下对齐r， + uint32_t alignSeq = constInfo.cmpRatio; + if (batchInfo.bIdx == 0) { + alignSeq = constInfo.cmpRatio - batchInfo.headHolderSeq; + } + if (quota > alignSeq) { + uint32_t delta = (batchInfo.bStartPos + batchInfo.sIdx + quota) & (constInfo.cmpRatio - 1); // 超出对齐的部分 + curDealSeq = quota - delta; + quota -= curDealSeq; + curDealTcNum = (curDealSeq + constInfo.cmpRatio - 1) / constInfo.cmpRatio; + curDealCompressedTcNum = min(curDealTcNum, batchInfo.compressedTcNum); + // 更新当前组所需信息 + UpdateCurGroup(basicBlockInfo, batchInfo, curGroupQuota, curDealSeq); + // 更新batch信息 + batchInfo.remSeqCnt = batchInfo.remSeqCnt - curDealSeq; + batchInfo.sIdx = batchInfo.sIdx + curDealSeq; + batchInfo.compressedTcNum -= curDealCompressedTcNum; + batchInfo.tcNum -= curDealTcNum; + // 更新loop信息 + basicBlockInfo.dealTcNum += curDealTcNum; + basicBlockInfo.compressedTcNum += curDealCompressedTcNum; + } + break; + } else { + // 处理整个batch + quota -= batchInfo.remSeqCnt; + curDealSeq = batchInfo.remSeqCnt; + curDealTcNum = batchInfo.tcNum; + // 更新当前组所需信息 + UpdateCurGroup(basicBlockInfo, batchInfo, curGroupQuota, curDealSeq); + // 更新batch和loop信息 + batchInfo.remSeqCnt = 0; + basicBlockInfo.dealTcNum += batchInfo.tcNum; + basicBlockInfo.compressedTcNum += batchInfo.compressedTcNum; + batchInfo.bIdx++; + SkipInvalidBatch(batchInfo); + } + } + uint32_t totalDataSize = constInfo.mGroupNum * constInfo.mBaseSize - quota; + // 2. 当前组的起始偏移 + uint32_t currentGroupStart = constInfo.mCurGroupIdx * constInfo.mBaseSize; + + // 3. 安全判断 + if (currentGroupStart >= totalDataSize) { + // 超出尾块 + basicBlockInfo.dealSeqCnt = 0; + } else { + // 还在有效范围内，计算剩余量 + uint32_t remaining = totalDataSize - currentGroupStart; + basicBlockInfo.dealSeqCnt = (remaining < constInfo.mBaseSize) ? remaining : constInfo.mBaseSize; + } + return basicBlockInfo; +} + + +template +__aicore__ inline uint32_t CompressorKernel::GetLoopTimes() +{ + // 计算主循环次数 + uint32_t loopTimes = 0; + BatchInfo batchInfo{}; + SkipInvalidBatch(batchInfo); + for (;batchInfo.bIdx < constInfo.batchSize; ++loopTimes) { + SkipOneLoop(batchInfo); + } + return loopTimes; +} + +template +__aicore__ inline void CompressorKernel::CalcSplitCoreInfo() +{ + // D方向的基本块数量 + constInfo.dBasicBlockNum = constInfo.headDim / constInfo.dBaseSize; + // 核的组数 + constInfo.coreGroupNum = constInfo.usedCoreNum / constInfo.dBasicBlockNum; + // 每个核处理的d方向的索引 + constInfo.dIdx = constInfo.aiCoreIdx % constInfo.dBasicBlockNum; + // 当前组id + constInfo.curGroupIdx = constInfo.aiCoreIdx / constInfo.dBasicBlockNum; + constInfo.mGroupNum = constInfo.coreGroupNum; + constInfo.mCurGroupIdx = constInfo.curGroupIdx; + + constInfo.mm1ResSize = constInfo.mBaseSize * constInfo.headDim * constInfo.coreGroupNum; + + uint32_t coff = (uint32_t)COMP::coff; + constInfo.mm1KvResSize = constInfo.mBaseSize * constInfo.headDim * coff; + constInfo.mm1ScoreResSize = constInfo.mBaseSize * constInfo.headDim * coff; + constInfo.vec1ResSize = constInfo.mBaseSize * constInfo.headDim * constInfo.nSize; + + constInfo.dbSize = constInfo.coreGroupNum * constInfo.mm1KvResSize; +} + +template +__aicore__ inline void CompressorKernel::InitWorkspace(__gm__ uint8_t *workspace) { + uint64_t offset = 0; + uint64_t mm1KvResStartOffset = offset; + // mm1KvResGm + mm1KvResGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + offset + + constInfo.curGroupIdx * constInfo.mm1KvResSize * sizeof(MM1_OUT_T))); + offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1KvResSize * sizeof(MM1_OUT_T); + + uint64_t mm1ScoreResStartOffset = offset; + // mm1ScoreResGm + mm1ScoreResGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + offset + + constInfo.curGroupIdx * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T))); + offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T); + + Vec1InputKvGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + mm1KvResStartOffset)); + + Vec1InputScoreGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + mm1ScoreResStartOffset)); + + vec1KvCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset)); + offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T); + + vec1ScoreCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset)); + offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T); + + uint64_t beforeVecOffset = offset; + + // vec1Res + vec1ResGm.SetGlobalBuffer( + (__gm__ VEC1_OUT_T *)(workspace + offset)); + offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.vec1ResSize * sizeof(VEC1_OUT_T); + // vec2Input + vec2InputGm.SetGlobalBuffer( + (__gm__ VEC1_OUT_T *)(workspace + beforeVecOffset)); +} + +template +__aicore__ inline void CompressorKernel::ComputeMm1(const RunInfo &info, bool isNeedExcute) { + CrossCoreWaitFlag(SYNC_V1_C1_FLAG + info.cubeDbIdx); + if (isNeedExcute) { + blockCube_.ComputeMm1(info); + } + CrossCoreSetFlag(SYNC_C1_FLAG); + CrossCoreWaitFlag(SYNC_C1_FLAG); + CrossCoreSetFlag(SYNC_C1_V1_FLAG + info.cubeDbIdx); +} + +template +__aicore__ inline void CompressorKernel::ComputeVec1(const Vec1RunInfo &info) { + CrossCoreWaitFlag(SYNC_C1_V1_FLAG + info.c1v1DbIdx); + CrossCoreWaitFlag(SYNC_V1_FLAG2 + info.c1v1DbIdx); + blockVec_.ComputeVec1(info); + CrossCoreSetFlag(SYNC_V1_FLAG); + CrossCoreWaitFlag(SYNC_V1_FLAG); + CrossCoreSetFlag(SYNC_V1_C1_FLAG + info.c1v1DbIdx); + CrossCoreSetFlag(SYNC_V1_FLAG2 + (info.c1v1DbIdx + 1) % constInfo.dbWorkspaceRatio); +} + +template +__aicore__ inline void CompressorKernel::ComputeVec2(const Vec2RunInfo &info) { + blockVec_.ComputeVec2(info); +} + +template +__aicore__ inline void CompressorKernel::AllocEventID() +{ + if ASCEND_IS_AIC { + blockCube_.AllocEventID(pipe_); + } else { + blockVec_.AllocEventID(); + for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) { + CrossCoreSetFlag(SYNC_V1_C1_FLAG + i); + } + CrossCoreSetFlag(SYNC_V1_FLAG2); + } +} + +template +__aicore__ inline void CompressorKernel::FreeEventID() +{ + if ASCEND_IS_AIC { + for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) { + CrossCoreWaitFlag(SYNC_V1_C1_FLAG + i); + } + blockCube_.FreeEventID(pipe_); + } else { + CrossCoreWaitFlag(SYNC_V1_FLAG2 + loopTimes % constInfo.dbWorkspaceRatio); + blockVec_.FreeEventID(); + } +} + +template +__aicore__ inline bool CompressorKernel::IsNeedExcuteC1(RunInfo info) +{ + // B超出范围则cube不执行 + return info.bStart < constInfo.batchSize; +} + +template +__aicore__ inline void CompressorKernel::CalcC1V1Params(RunInfo &info, Vec1RunInfo &vec1Info, BatchInfo &batchInfo, uint32_t loopIdx) +{ + vec1Info.bStart = batchInfo.bIdx; + vec1Info.sStart = batchInfo.sIdx; + vec1Info.resetResFlag = (loopIdx & (constInfo.nSize - 1)) == 0; + vec1Info.c1v1DbIdx = (vec1Loop++ & (constInfo.dbWorkspaceRatio - 1)); + vec1Info.v1v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1)); + BasicBlockInfo basicBlockInfo = SkipOneLoop(batchInfo); + info.cubeDbIdx = (cubeLoop++ & (constInfo.dbWorkspaceRatio - 1)); + info.dealSeqCnt = basicBlockInfo.dealSeqCnt; + info.dealTcNum = basicBlockInfo.dealTcNum; + info.bStart = basicBlockInfo.bIdx; + info.sStart = basicBlockInfo.sIdx; + info.kStartIdx = kStartIdx_; + info.dealKSize = dealKSize_; + info.hStart = hStart_; + vec1Info.dealTcNum = basicBlockInfo.dealTcNum; + vec1Info.dealScSize = basicBlockInfo.compressedTcNum; + allCompressedTcNum_ += basicBlockInfo.compressedTcNum; +} + +template +__aicore__ inline bool CompressorKernel::IsNeedExcuteV2(Vec2RunInfo &vec2Info) +{ + return (vec2Info.dealScSize > 0); +} + +template +__aicore__ inline bool CompressorKernel::IsNeedSyncAll(uint32_t curBasicBlockIdx) +{ + if (allCompressedTcNum_ == 0) { + return false; + } + + uint32_t cnt = curBasicBlockIdx + 1; + if ((cnt == loopTimes) || (cnt % constInfo.nSize == 0)) { + return true; + } + return false; +} + +template +__aicore__ inline void CompressorKernel::UpdateVec2Info( + Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info) +{ + // nSize轮起始先重置v2Info信息 + if (curBasicBlockIdx % constInfo.nSize == 0) { + vec2Info.v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1)); + vec2Info.bStart = info.bStart; + vec2Info.sStart = info.sStart; + // 将sStart转成bCompressedId + uint32_t startPos = tools_.GetStartPos(info.bStart); + if (tools_.isExistSeqUsed_) { + uint32_t seqUsed = tools_.GetSeqUsed(info.bStart); + if (vec2Info.sStart >= seqUsed) { + vec2Info.bStart++; + vec2Info.sStart = 0; + } + } + vec2Info.bCompressedId = (startPos + vec2Info.sStart) / constInfo.cmpRatio - startPos / constInfo.cmpRatio; + + vec2Info.dealScSize = 0; + } else if ((curBasicBlockIdx + 1) % constInfo.nSize == 0) { + vec2Loop++; + } + vec2Info.dealScSize += info.dealScSize; + vec2Info.compressedId += info.dealScSize; +} + +template +__aicore__ inline void CompressorKernel::Process() +{ + // 所有batch的有效序列都为0时, 直接退出 + if (constInfo.batchSize == 0) { + return; + } + AllocEventID(); + + BatchInfo batchInfo{}; + + RunInfo extraInfo[1]; + Vec1RunInfo vec1Info{}; + Vec2RunInfo vec2Info{}; + SkipInvalidBatch(batchInfo); + for (uint32_t i = 0; i < loopTimes; ++i) { + RunInfo &extraInfo0 = extraInfo[0]; + CalcC1V1Params(extraInfo0, vec1Info, batchInfo, i); + bool isNeedExcuteC1 = IsNeedExcuteC1(extraInfo0); + + if ASCEND_IS_AIC { + ComputeMm1(extraInfo0, isNeedExcuteC1); + } else { + ComputeVec1(vec1Info); + UpdateVec2Info(vec2Info, i, vec1Info); + + if (IsNeedSyncAll(i)) { + SyncAll(); + if (IsNeedExcuteV2(vec2Info)) { + ComputeVec2(vec2Info); + } + } + } + + } + FreeEventID(); +} + +} // namespace Compressor + +#endif // COMPRESSOR_KERNEL_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_kernel_full_load.h b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_kernel_full_load.h new file mode 100644 index 000000000..53e220a9b --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_kernel_full_load.h @@ -0,0 +1,474 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_kernel_full_load.h + * \brief + */ + +#ifndef COMPRESSOR_KERNEL_FULL_LOAD_H +#define COMPRESSOR_KERNEL_FULL_LOAD_H + +#include "compressor_comm.h" +#include "compressor_template_tiling_key.h" +#include "compressor_tiling_data.h" +#include "compressor_tools.h" +#include "compressor_block_cube_full_load.h" +#include "compressor_block_vec_full_load.h" + +using namespace AscendC; + +namespace Compressor { +template +class CompressorKernelFullLoad { +public: + __aicore__ inline CompressorKernelFullLoad(TPipe* pipe, const optiling::CompressorTilingData* __restrict tilingData) + : pipe_(pipe), tilingData_(tilingData) {} + + __aicore__ inline void Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut, + __gm__ uint8_t *workspace); + __aicore__ inline void Process(); + +private: + // ================================Init functions================================== + __aicore__ inline void InitWorkspace(__gm__ uint8_t *workspace); + // ================================Process functions================================ + __aicore__ inline void InitTilingData(); + // 获取基本块数量 + __aicore__ inline void SkipInvalidBatch(BatchInfo &batchInfo); + // 计算分核基本信息 + __aicore__ inline void CalcSplitCoreInfo(); + + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void ComputeMm1(const RunInfo &info); + __aicore__ inline void ComputeVec1(const Vec1RunInfo &info); + __aicore__ inline void ComputeVec2(const Vec2RunInfo &info); + + __aicore__ inline bool IsNeedExcuteC1(RunInfo info); + __aicore__ inline bool IsNeedSyncAll(uint32_t curBasicBlockIdx); + __aicore__ inline void UpdateVec2Info(Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info); + __aicore__ inline bool IsNeedExcuteV2(Vec2RunInfo &vec2Info); + __aicore__ inline void CalcCubeParams(RunInfo &info); + __aicore__ inline void CalcV1Params(Vec1RunInfo &vec1Info); + + using X_T = typename AscendC::Conditional::type; + using T = float; + using MM1_OUT_T = T; + using VEC1_OUT_T = T; + + // 常量 + static constexpr uint64_t SYNC_MODE0 = 0; + static constexpr uint64_t SYNC_MODE2 = 2; + static constexpr uint32_t SYNC_C1_FLAG = 3; + static constexpr uint32_t SYNC_V1_FLAG = 4; + static constexpr uint32_t SYNC_V1_FLAG2 = 5; + static constexpr uint32_t SYNC_C1_V1_FLAG = 7; + static constexpr uint32_t SYNC_V1_C1_FLAG = 9; + + // ==============================TilingData&TPipe============================== + TPipe* pipe_; + const optiling::CompressorTilingData* __restrict tilingData_; + // ===========================Workspace Global Tensor=========================== + GlobalTensor mm1KvResGm; + GlobalTensor mm1ScoreResGm; + GlobalTensor vec1KvCacheGm; + GlobalTensor vec1ScoreCacheGm; + GlobalTensor Vec1InputKvGm; + GlobalTensor Vec1InputScoreGm; + GlobalTensor vec1ResGm; + GlobalTensor vec2InputGm; + // ================================Task Info==================================== + CompressorTools tools_; + ConstInfo constInfo{}; + uint32_t aiCoreIdx = 0; + + // ==============================Service Define============================== + CompressorBlockCubeFullLoad blockCube_; + CompressorBlockVectorFullLoad blockVec_; + + uint32_t allCompressedTcNum_ = 0; + uint32_t curCompressedTcNum_ = 0; + uint32_t accDealSize = 0; + uint32_t loopTimes = 0; + uint32_t cubeLoop = 0; + uint32_t vec1Loop = 0; + uint32_t vec2Loop = 0; + uint32_t kStartIdx_ = 0; + uint32_t dealKSize_ = 0; +}; + +template +__aicore__ inline void CompressorKernelFullLoad::Init( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut, + __gm__ uint8_t *workspace) +{ + if ASCEND_IS_AIV { + constInfo.aiCoreIdx = GetBlockIdx() / 2; + } else { + constInfo.aiCoreIdx = GetBlockIdx(); + } + + InitTilingData(); + // init tools + tools_.toolParams_.seqSize = tilingData_->baseParams.seqSize; + tools_.toolParams_.cmpRatio = tilingData_->baseParams.cmpRatio; + tools_.Init(startPos, seqUsed, cuSeqlens); + + // 剔除尾部的无效batch + for (; constInfo.batchSize > 0; --constInfo.batchSize) { + uint32_t bSeqUsed = tools_.GetSeqLength(constInfo.batchSize - 1); + if (bSeqUsed > 0) { + break; + } + } + + // 所有batch的有效序列都为0时, 直接退出 + if (constInfo.batchSize == 0) { + return; + } + + // 0. 计算最后一个Tc块的起始位置 + constInfo.bIdxOfLastTc = constInfo.batchSize - 1; + // 1. 计算head_dim的切分大小, 构建ConstInfo的其他信息 + CalcSplitCoreInfo(); + // 2. 计算循环次数 + loopTimes = 1; + // 3. 初始化workspace + InitWorkspace(workspace); + // 4. 初始化block层 + if ASCEND_IS_AIC { + blockCube_.InitParams(constInfo, tools_); + blockCube_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos, + stateBlockTable, cuSeqlens, seqUsed, startPos, cmpKvOut); + blockCube_.InitBuffers(pipe_); + blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm); + } else { + blockVec_.InitParams(constInfo, tools_); + blockVec_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos, stateBlockTable, + cuSeqlens, seqUsed, startPos, cmpKvOut); + blockVec_.InitBuffers(pipe_); + blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm, vec1ResGm, vec2InputGm); + } +} + +template +__aicore__ inline void CompressorKernelFullLoad::InitTilingData() { + constInfo.cmpRatio = tilingData_->baseParams.cmpRatio; + constInfo.batchSize = tilingData_->baseParams.batchSize; + constInfo.mBaseSize = tilingData_->innerSplitParams.mBaseSize; + constInfo.dBaseSize = tilingData_->innerSplitParams.dBaseSize; + constInfo.headDim = tilingData_->baseParams.headDim; + constInfo.hSize = tilingData_->baseParams.hiddenSize; + constInfo.sSize = tilingData_->baseParams.seqSize; + constInfo.ropeHeadDim = tilingData_->baseParams.ropeHeadDim; + constInfo.normEps = tilingData_->baseParams.normEps; + constInfo.stateCacheStrideDim0 = tilingData_->baseParams.stateCacheStrideDim0; + constInfo.reciprocalD = tilingData_->baseParams.reciprocalD; + constInfo.usedCoreNum = tilingData_->baseParams.usedCoreNum; + + constInfo.blockNum = tilingData_->pageAttentionParams.blockNum; + constInfo.blockSize = tilingData_->pageAttentionParams.blockSize; + constInfo.maxBlockNumPerBatch = tilingData_->pageAttentionParams.maxBlockNumPerBatch; + + constInfo.nSize = tilingData_->baseParams.nSize; + constInfo.vec1TailCacheSize = tilingData_->workspaceParams.vec1TailCacheSize; + constInfo.dbWorkspaceRatio = tilingData_->workspaceParams.dbWorkspaceRatio; + + constInfo.mStart = tilingData_->baseParams.splitCoreParam[constInfo.aiCoreIdx].mStart; + constInfo.mEnd = tilingData_->baseParams.splitCoreParam[constInfo.aiCoreIdx].mEnd; + constInfo.nStart = tilingData_->baseParams.splitCoreParam[constInfo.aiCoreIdx].nStart; + constInfo.nEnd = tilingData_->baseParams.splitCoreParam[constInfo.aiCoreIdx].nEnd; + constInfo.kStart = tilingData_->baseParams.splitCoreParam[constInfo.aiCoreIdx].kStart; + constInfo.kEnd = tilingData_->baseParams.splitCoreParam[constInfo.aiCoreIdx].kEnd; + constInfo.mLoopNum = tilingData_->baseParams.mLoopNum; + constInfo.kBaseNum = tilingData_->baseParams.kBaseNum; + constInfo.kBaseSize = tilingData_->baseParams.kBaseSize; +} + +template +__aicore__ inline void CompressorKernelFullLoad::SkipInvalidBatch(BatchInfo &batchInfo) +{ + for (; batchInfo.bIdx < constInfo.batchSize; ++batchInfo.bIdx) { + batchInfo.seqCnt = tools_.GetSeqLength(batchInfo.bIdx); + if (batchInfo.seqCnt > 0) { + break; + } + } + batchInfo.remSeqCnt = batchInfo.seqCnt; + if (tools_.isExistSeqUsed_) { + batchInfo.seqUsedCnt = tools_.GetSeqUsed(batchInfo.bIdx); + } else { + batchInfo.seqUsedCnt = batchInfo.seqCnt; + } + if (batchInfo.bIdx < constInfo.batchSize) { + batchInfo.bStartPos = tools_.GetStartPos(batchInfo.bIdx); + batchInfo.sIdx = 0; + batchInfo.headHolderSeq = batchInfo.bStartPos & (constInfo.cmpRatio - 1); + batchInfo.tcNum = (batchInfo.bStartPos + batchInfo.seqCnt + constInfo.cmpRatio - 1) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio; + batchInfo.compressedTcNum = (batchInfo.bStartPos + batchInfo.seqUsedCnt) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio; + } +} + +template +__aicore__ inline void CompressorKernelFullLoad::CalcSplitCoreInfo() +{ + // D方向的基本块数量 + constInfo.dBasicBlockNum = constInfo.headDim / constInfo.dBaseSize; + // 核的组数 + constInfo.coreGroupNum = constInfo.usedCoreNum / constInfo.dBasicBlockNum; + // 当前组id + constInfo.curGroupIdx = constInfo.aiCoreIdx / constInfo.dBasicBlockNum; + constInfo.mGroupNum = constInfo.coreGroupNum; + constInfo.mCurGroupIdx = constInfo.curGroupIdx; + + uint32_t coff = (uint32_t)COMP::coff; + constInfo.mm1KvResSize = constInfo.mBaseSize * constInfo.headDim * coff; + constInfo.mm1ScoreResSize = constInfo.mBaseSize * constInfo.headDim * coff; + constInfo.vec1ResSize = constInfo.mBaseSize * constInfo.headDim * constInfo.nSize; + + constInfo.dbSize = constInfo.coreGroupNum * constInfo.mm1KvResSize; +} + +template +__aicore__ inline void CompressorKernelFullLoad::InitWorkspace(__gm__ uint8_t *workspace) { + uint64_t offset = 0; + uint64_t mm1KvResStartOffset = offset; + // mm1KvResGm + mm1KvResGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + offset + + constInfo.curGroupIdx * constInfo.mm1KvResSize * sizeof(MM1_OUT_T))); + offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1KvResSize * sizeof(MM1_OUT_T); + + uint64_t mm1ScoreResStartOffset = offset; + // mm1ScoreResGm + mm1ScoreResGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + offset + + constInfo.curGroupIdx * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T))); + offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T); + + Vec1InputKvGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + mm1KvResStartOffset)); + + Vec1InputScoreGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + mm1ScoreResStartOffset)); + + vec1KvCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset)); + offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T); + + vec1ScoreCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset)); + offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T); + + uint64_t beforeVecOffset = offset; + + // vec1Res + vec1ResGm.SetGlobalBuffer( + (__gm__ VEC1_OUT_T *)(workspace + offset)); + offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.vec1ResSize * sizeof(VEC1_OUT_T); + // vec2Input + vec2InputGm.SetGlobalBuffer( + (__gm__ VEC1_OUT_T *)(workspace + beforeVecOffset)); +} + +template +__aicore__ inline void CompressorKernelFullLoad::ComputeMm1(const RunInfo &info) { + CrossCoreWaitFlag(SYNC_V1_C1_FLAG + info.cubeDbIdx); + blockCube_.ComputeMm1(info); + CrossCoreSetFlag(SYNC_C1_FLAG); + CrossCoreWaitFlag(SYNC_C1_FLAG); + CrossCoreSetFlag(SYNC_C1_V1_FLAG + info.cubeDbIdx); +} + +template +__aicore__ inline void CompressorKernelFullLoad::ComputeVec1(const Vec1RunInfo &info) { + CrossCoreWaitFlag(SYNC_C1_V1_FLAG + info.c1v1DbIdx); + CrossCoreWaitFlag(SYNC_V1_FLAG2 + info.c1v1DbIdx); + blockVec_.ComputeVec1(); + CrossCoreSetFlag(SYNC_V1_FLAG); + CrossCoreWaitFlag(SYNC_V1_FLAG); + CrossCoreSetFlag(SYNC_V1_C1_FLAG + info.c1v1DbIdx); + CrossCoreSetFlag(SYNC_V1_FLAG2 + (info.c1v1DbIdx + 1) % constInfo.dbWorkspaceRatio); +} + +template +__aicore__ inline void CompressorKernelFullLoad::ComputeVec2(const Vec2RunInfo &info) { + blockVec_.ComputeVec2(); +} + +template +__aicore__ inline void CompressorKernelFullLoad::AllocEventID() +{ + if ASCEND_IS_AIC { + blockCube_.AllocEventID(pipe_); + } else { + blockVec_.AllocEventID(); + for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) { + CrossCoreSetFlag(SYNC_V1_C1_FLAG + i); + } + CrossCoreSetFlag(SYNC_V1_FLAG2); + } +} + +template +__aicore__ inline void CompressorKernelFullLoad::FreeEventID() +{ + if ASCEND_IS_AIC { + for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) { + CrossCoreWaitFlag(SYNC_V1_C1_FLAG + i); + } + blockCube_.FreeEventID(pipe_); + } else { + CrossCoreWaitFlag(SYNC_V1_FLAG2 + loopTimes % constInfo.dbWorkspaceRatio); + blockVec_.FreeEventID(); + } +} + +template +__aicore__ inline bool CompressorKernelFullLoad::IsNeedExcuteC1(RunInfo info) +{ + // B超出范围则cube不执行 + return info.bStart < constInfo.batchSize; +} + +template +__aicore__ inline void CompressorKernelFullLoad::CalcCubeParams(RunInfo &info) +{ + info.cubeDbIdx = (cubeLoop++ & (constInfo.dbWorkspaceRatio - 1)); + info.dealSeqCnt = constInfo.mEnd - constInfo.mStart; +} + +template +__aicore__ inline void CompressorKernelFullLoad::CalcV1Params(Vec1RunInfo &vec1Info) +{ + vec1Info.bStart = 0; + vec1Info.sStart = 0; + vec1Info.resetResFlag = false; + vec1Info.c1v1DbIdx = (vec1Loop++ & (constInfo.dbWorkspaceRatio - 1)); + vec1Info.v1v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1)); + for (uint32_t bIdx = 0; bIdx < constInfo.batchSize; ++bIdx) { + uint64_t bSeqCnt = tools_.GetSeqLength(bIdx); + if (bSeqCnt == 0) { + continue; + } + uint64_t bStartPos = tools_.GetStartPos(bIdx); + vec1Info.dealTcNum += (bStartPos + bSeqCnt + constInfo.cmpRatio - 1) / constInfo.cmpRatio - bStartPos / constInfo.cmpRatio; + vec1Info.dealScSize += (bStartPos + bSeqCnt) / constInfo.cmpRatio - bStartPos / constInfo.cmpRatio; + } +} + +template +__aicore__ inline bool CompressorKernelFullLoad::IsNeedExcuteV2(Vec2RunInfo &vec2Info) +{ + return (vec2Info.dealScSize > 0); +} + +template +__aicore__ inline bool CompressorKernelFullLoad::IsNeedSyncAll(uint32_t curBasicBlockIdx) +{ + if (allCompressedTcNum_ == 0) { + return false; + } + + uint32_t cnt = curBasicBlockIdx + 1; + if ((cnt == loopTimes) || (cnt % constInfo.nSize == 0)) { + return true; + } + return false; +} + +template +__aicore__ inline void CompressorKernelFullLoad::UpdateVec2Info( + Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info) +{ + // nSize轮起始先重置v2Info信息 + if (curBasicBlockIdx % constInfo.nSize == 0) { + vec2Info.v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1)); + vec2Info.bStart = info.bStart; + vec2Info.sStart = info.sStart; + // 将sStart转成bCompressedId + uint32_t startPos = tools_.GetStartPos(info.bStart); + if (tools_.isExistSeqUsed_) { + uint32_t seqUsed = tools_.GetSeqUsed(info.bStart); + if (vec2Info.sStart >= seqUsed) { + vec2Info.bStart++; + vec2Info.sStart = 0; + } + } + vec2Info.bCompressedId = (startPos + vec2Info.sStart) / constInfo.cmpRatio - startPos / constInfo.cmpRatio; + + vec2Info.dealScSize = 0; + } else if ((curBasicBlockIdx + 1) % constInfo.nSize == 0) { + vec2Loop++; + } + vec2Info.dealScSize += info.dealScSize; + vec2Info.compressedId += info.dealScSize; +} + +template +__aicore__ inline void CompressorKernelFullLoad::Process() +{ + // 所有batch的有效序列都为0时, 直接退出 + if (constInfo.batchSize == 0) { + return; + } + AllocEventID(); + + RunInfo extraInfo[1]; + Vec1RunInfo vec1Info{}; + Vec2RunInfo vec2Info{}; + for (uint32_t i = 0; i < loopTimes; ++i) { + RunInfo &extraInfo0 = extraInfo[0]; + if ASCEND_IS_AIV { + CalcV1Params(vec1Info); + } else { + CalcCubeParams(extraInfo0); + } + if ASCEND_IS_AIC { + ComputeMm1(extraInfo0); + } else { + ComputeVec1(vec1Info); + UpdateVec2Info(vec2Info, i, vec1Info); + SyncAll(); + if (IsNeedExcuteV2(vec2Info)) { + ComputeVec2(vec2Info); + } + } + } + FreeEventID(); +} + +} // namespace Compressor + +#endif // COMPRESSOR_KERNEL_FULL_LOAD_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_template_tiling_key.h b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_template_tiling_key.h new file mode 100644 index 000000000..03c8e1e84 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_template_tiling_key.h @@ -0,0 +1,54 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file COMPRESSOR_template_tiling_key.h + * \brief + */ + +#ifndef COMPRESSOR_TEMPLATE_TILING_KEY_H +#define COMPRESSOR_TEMPLATE_TILING_KEY_H + +#include "ascendc/host_api/tiling/template_argument.h" + +#define ASCENDC_TPL_1_BW 1 // 每个参数占用1个bit位 +#define ASCENDC_TPL_2_BW 2 // 每个参数占用2个bit位 +#define ASCENDC_TPL_4_BW 4 // 每个参数占用4个bit位 + +// 可表示的tilingkey范围为64bit，注意不可超过限制 +ASCENDC_TPL_ARGS_DECL(compressor, // 算子唯一标识，与opType保持一致 + // 可能需要切分之后的headdim + // bit:0 LAYOUT 0:BSH 1:TH + ASCENDC_TPL_UINT_DECL(X_LAYOUT, ASCENDC_TPL_1_BW, ASCENDC_TPL_UI_LIST, 0, 1), + // bit:1-4 x的dtype 0:BF16 1:FP16 + ASCENDC_TPL_UINT_DECL(X_DTYPE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, 0, 1), + // bit:5-6 coff 1:无需overlap 2:需要overlap + ASCENDC_TPL_UINT_DECL(COFF, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2), + // bit:7-8 rotary_mode 1:half 2:interleave + ASCENDC_TPL_UINT_DECL(ROTARY_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2), + // bit:9-10 cache_mode 1:CONTINUOUS 2:cycle + ASCENDC_TPL_UINT_DECL(CACHE_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2), + // bit:11-12 template_id 0:empty_tensor 1:normal 2:full load + ASCENDC_TPL_UINT_DECL(TEMPLATE_ID, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 0, 1, 2), + +); + +ASCENDC_TPL_SEL( + + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_UINT_SEL(X_LAYOUT, ASCENDC_TPL_UI_LIST, 0, 1), + ASCENDC_TPL_UINT_SEL(X_DTYPE, ASCENDC_TPL_UI_LIST, 0, 1), + ASCENDC_TPL_UINT_SEL(COFF, ASCENDC_TPL_UI_LIST, 1, 2), + ASCENDC_TPL_UINT_SEL(ROTARY_MODE, ASCENDC_TPL_UI_LIST, 1, 2), + ASCENDC_TPL_UINT_SEL(CACHE_MODE, ASCENDC_TPL_UI_LIST, 1, 2), + ASCENDC_TPL_UINT_SEL(TEMPLATE_ID, ASCENDC_TPL_UI_LIST, 0, 1, 2), + ASCENDC_TPL_TILING_STRUCT_SEL(optiling::CompressorTilingData)), +); + +#endif // COMPRESSOR_TEMPLATE_TILING_KEY_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_tiling_data.h b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_tiling_data.h new file mode 100644 index 000000000..8394ba197 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_tiling_data.h @@ -0,0 +1,83 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file COMPRESSOR_tiling_datay.h + * \brief + */ + +#ifndef COMPRESSOR_TILING_DATA_H +#define COMPRESSOR_TILING_DATA_H +#include +#include "kernel_tiling/kernel_tiling.h" + +const uint32_t CMP_MAX_AIC_CORE_NUM = 36; + +namespace optiling { + struct CompressorSplitCoreParams { + uint32_t mStart; + uint32_t mEnd; + uint32_t nStart; + uint32_t nEnd; + uint32_t kStart; + uint32_t kEnd; + }; + + // 1. 基础参数结构体 + struct CompressorBaseParams { + uint32_t batchSize = 0; // bastch size（批大小） + uint32_t seqSize = 0; // sequence size（kvs大小） + uint32_t hiddenSize = 0; // hidden size（隐藏层大小） + uint32_t tokenSize = 0; // token size = batchSize * seqSize(token总数：批大小x序列1长度) + uint32_t headDim = 0; // head size of kv + uint32_t ropeHeadDim = 64; // dim size per rope head 64（单个带RoPE头的维度） + uint32_t csSize = 0; // Compress sequence len + uint32_t cmpRatio = 4; // Compress ratio + uint32_t cgSize = 0; // Compress group size + float normEps = 1e-6; // RMSNorm eps + float reciprocalD = 0; // 1分之D + uint32_t usedCoreNum = 0; // 使用核数 + uint32_t nSize = 0; // 控制v2积攒的轮数 + uint64_t stateCacheStrideDim0 = 0; // stateCache第0维的stride + uint32_t kBaseNum = 0; + uint32_t kBaseSize = 0; + uint32_t coreGroupNum = 0; + uint32_t mLoopNum = 0; + CompressorSplitCoreParams splitCoreParam[CMP_MAX_AIC_CORE_NUM]; + }; + + struct CompressorPageAttentionParams { + uint32_t blockNum = 0; + uint32_t blockSize = 1; + uint32_t maxBlockNumPerBatch = 1; + }; + + struct CompressorInnerSplitParams { + uint32_t mBaseSize; + uint32_t dBaseSize; + }; + + struct CompressorWorkspaceParams { + uint32_t mm1KvResSize; + uint32_t mm1ScoreResSize; + uint32_t vec1ResSize; + uint32_t vec1TailCacheSize; + uint32_t dbWorkspaceRatio = 1; + }; + + struct CompressorTilingData { + CompressorBaseParams baseParams; + CompressorPageAttentionParams pageAttentionParams; + CompressorInnerSplitParams innerSplitParams; + CompressorWorkspaceParams workspaceParams; + }; +} // optiling + +#endif // COMPRESSOR_TILING_DATA_H diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_tools.h b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_tools.h new file mode 100644 index 000000000..ad613e634 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/compressor_tools.h @@ -0,0 +1,890 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor_tools.h + * \brief 放算子都需要、与算子联系紧密、但是又不方便单独独立出来的公共工具 + */ + +#ifndef COMPRESSOR_TOOLS_H +#define COMPRESSOR_TOOLS_H + +#include "compressor_comm.h" + +using namespace AscendC; + +namespace Compressor { + +struct ToolsParams { + uint32_t seqSize = 0U; + uint32_t cmpRatio = 0U; +}; + +template +class CompressorTools { +public: + __aicore__ inline CompressorTools() {} + + __aicore__ inline void Init(__gm__ uint8_t *cuSeqlens, __gm__ uint8_t *seqUsed, __gm__ uint8_t *startPos); + + __aicore__ inline uint32_t GetSeqUsed(uint32_t bIdx); + __aicore__ inline uint32_t GetStartPos(uint32_t bIdx); + __aicore__ inline uint32_t GetSeqLength(uint32_t bIdx); + __aicore__ inline uint32_t GetTIdxByBatch(uint32_t bIdx); + +public: + ToolsParams toolParams_ {}; + bool isExistSeqUsed_ = false; + +private: + bool isExistStartPos_ = false; + GlobalTensor cuSeqlensGm_; + GlobalTensor sequsedGm_; + GlobalTensor startPosGm_; +}; + +template +__aicore__ inline void CompressorTools::Init(__gm__ uint8_t *startPos, __gm__ uint8_t *seqUsed, + __gm__ uint8_t *cuSeqlens) +{ + isExistStartPos_ = (startPos != nullptr); + if (isExistStartPos_) { + startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos); + } + + isExistSeqUsed_ = (seqUsed != nullptr); + if (isExistSeqUsed_) { + sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed); + } + + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens); + } +} + +template +__aicore__ inline uint32_t CompressorTools::GetSeqUsed(uint32_t bIdx) +{ + if (isExistSeqUsed_) { + return (uint32_t)sequsedGm_.GetValue(bIdx); + } else { + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + return (uint32_t)(cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx)); + } else { + return toolParams_.seqSize; + } + } +} + +template +__aicore__ inline uint32_t CompressorTools::GetStartPos(uint32_t bIdx) +{ + if (isExistStartPos_) { + return (uint32_t)startPosGm_.GetValue(bIdx); + } else { + return 0; + } +} + +template +__aicore__ inline uint32_t CompressorTools::GetSeqLength(uint32_t bIdx) +{ + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + return cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx); + } else { + return toolParams_.seqSize; + } +} + +template +__aicore__ inline uint32_t CompressorTools::GetTIdxByBatch(uint32_t bIdx) +{ + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + return (uint32_t)(cuSeqlensGm_.GetValue(bIdx)); + } else { + return toolParams_.seqSize * bIdx; + } +} + +// iterator +struct SliceInfo { + __aicore__ inline SliceInfo(){}; + __aicore__ inline SliceInfo(uint32_t bIdx, uint32_t sIdx) : bIdx(bIdx), sIdx(sIdx) {}; + + uint32_t bIdx = 0U; + uint32_t sIdx = 0U; + uint32_t bSeqUsed = 0U; + uint32_t bStartPos = 0U; + + uint32_t headHolderSeqCnt = 0U; + uint32_t validSeqCnt = 0U; + uint32_t tailHolderSeqCnt = 0U; + + uint32_t dealSeqCnt = 0; + uint32_t dealTcSize = 0U; + uint32_t compressTcSize = 0U; +}; + +template +class CompressorSliceIterator { +public: + __aicore__ inline CompressorSliceIterator(CompressorTools &tools) : tools_(tools) {} + + __aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx); + __aicore__ inline void SetMaxBatchSize(uint32_t batch_size); + __aicore__ inline void SetMaxDealSeqCnt(uint32_t maxDealSeqCnt); + __aicore__ inline bool IsEnd(); + __aicore__ inline void IteratorSlice(); + __aicore__ inline SliceInfo& GetSlice(); + __aicore__ inline SliceInfo& GetSliceByCmp(); + + bool isFirst_ = true; + SliceInfo sliceInfo_{}; + +private: + CompressorTools &tools_; + + // iterator + uint32_t maxDealSeqCnt_ = 0; + uint32_t batch_size_ = 0; +}; + +template +__aicore__ inline void CompressorSliceIterator::Reset(uint32_t bIdx, uint32_t sIdx) +{ + sliceInfo_.bIdx = bIdx; + sliceInfo_.sIdx = sIdx; + isFirst_ = true; +} + +template +__aicore__ inline void CompressorSliceIterator::SetMaxBatchSize(uint32_t batch_size) +{ + this->batch_size_ = batch_size; +} + +template +__aicore__ inline void CompressorSliceIterator::SetMaxDealSeqCnt(uint32_t maxDealSeqCnt) +{ + this->maxDealSeqCnt_ = maxDealSeqCnt; +} + +template +__aicore__ inline bool CompressorSliceIterator::IsEnd() +{ + return (sliceInfo_.bIdx >= batch_size_) || (maxDealSeqCnt_ == 0); +} + +template +__aicore__ inline void CompressorSliceIterator::IteratorSlice() +{ + bool isUpdateBatchInfo = false; + if (!isFirst_) { + // 更新剩余未处理的行数 + maxDealSeqCnt_ -= sliceInfo_.dealSeqCnt; + // 更新sIdx和bIdx、以及与bIdx相关的bStartPos和bSeqUsed + sliceInfo_.sIdx += sliceInfo_.validSeqCnt; + if (sliceInfo_.sIdx == sliceInfo_.bSeqUsed) { + sliceInfo_.sIdx = 0; + sliceInfo_.bIdx++; + isUpdateBatchInfo = true; + } + } else { + isUpdateBatchInfo = true; + isFirst_ = false; + } + + // 更新与bIdx相关的bStartPos和bSeqUsed + if (isUpdateBatchInfo) { + // SkipInvalidBatch + while (sliceInfo_.bIdx < batch_size_) { + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + if (sliceInfo_.bSeqUsed > 0) { + break; + } + sliceInfo_.bIdx++; + } + if (sliceInfo_.bIdx < batch_size_) { + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + } + } +} + +template +__aicore__ inline SliceInfo& CompressorSliceIterator::GetSliceByCmp() +{ + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + if (isFirst_) { + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + isFirst_ = false; + } + // 计算头部占位行数、有效数据行数、尾部占位行数 + sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio; + + sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx; + if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) { + sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt; + } + sliceInfo_.tailHolderSeqCnt = + cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio; + if (sliceInfo_.tailHolderSeqCnt == cmpRatio) { + sliceInfo_.tailHolderSeqCnt = 0; + } + + // 头和尾处理，否则需要处理的seq等于cmpRatio + if (sliceInfo_.validSeqCnt < cmpRatio) { + sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt; + if (sliceInfo_.sIdx == 0) { + sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt; + } + } else { + sliceInfo_.dealSeqCnt = cmpRatio; + } + sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt; + + // 计算本次可以处理的Tc个数 + sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio; + + // 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0 + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize; + if (sliceInfo_.tailHolderSeqCnt > 0) { + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时，其不需要压缩 + } + + return sliceInfo_; +} + +template +__aicore__ inline SliceInfo& CompressorSliceIterator::GetSlice() +{ + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + if (isFirst_) { + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + isFirst_ = false; + } + // 计算头部占位行数、有效数据行数、尾部占位行数 + sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio; + sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx; + if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) { + sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt; + } + sliceInfo_.tailHolderSeqCnt = + cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio; + if (sliceInfo_.tailHolderSeqCnt == cmpRatio) { + sliceInfo_.tailHolderSeqCnt = 0; + } + + sliceInfo_.dealSeqCnt = sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt + sliceInfo_.tailHolderSeqCnt; + // 计算本次可以处理的Tc个数 + sliceInfo_.dealTcSize = sliceInfo_.dealSeqCnt / cmpRatio; + + // 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0 + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize; + if (sliceInfo_.tailHolderSeqCnt > 0) { + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时，其不需要压缩 + } + + return sliceInfo_; +} + +struct SplitCoreSliceInfo : public SliceInfo { + __aicore__ inline SplitCoreSliceInfo() {}; + __aicore__ inline SplitCoreSliceInfo(uint32_t bIdx, uint32_t sIdx) : SliceInfo(bIdx, sIdx) {}; + + uint32_t preFirstSeqCnt = 0U; // 左边每次迭代基本块的第一个seqCnt大小 +}; + +template +class CompressorSplitCoreSliceIterator { +public: + __aicore__ inline CompressorSplitCoreSliceIterator(CompressorTools &tools) : tools_(tools) {} + + __aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx); + __aicore__ inline void SetMaxBatchSize(uint32_t batch_size); + __aicore__ inline void SetMaxDealSeqCnt(uint32_t maxDealSeqCnt); + __aicore__ inline bool IsEnd(); + __aicore__ inline void IteratorSlice(); + __aicore__ inline SplitCoreSliceInfo& GetSlice(); + __aicore__ inline SplitCoreSliceInfo& GetSliceByCmp(); + __aicore__ inline uint32_t GetBIdx(); + __aicore__ inline SplitCoreSliceInfo& GetLeftNextCmpSeqCnt(); + __aicore__ inline SplitCoreSliceInfo& GetRightNextCmpSeqCnt(); + + bool isFirst_ = true; + bool isLeftFirstBath = false; + bool isMaxDealSeqCntFirst = false; + + SplitCoreSliceInfo sliceInfo_{}; + +private: + CompressorTools &tools_; + + // iterator + uint32_t maxDealSeqCnt_ = 0; + uint32_t batch_size_ = 0; +}; + +template +__aicore__ inline void CompressorSplitCoreSliceIterator::Reset(uint32_t bIdx, uint32_t sIdx) +{ + sliceInfo_.bIdx = bIdx; + sliceInfo_.sIdx = sIdx; + isFirst_ = true; +} + +template +__aicore__ inline void CompressorSplitCoreSliceIterator::SetMaxBatchSize(uint32_t batch_size) +{ + this->batch_size_ = batch_size; + isMaxDealSeqCntFirst = true; +} + +template +__aicore__ inline void CompressorSplitCoreSliceIterator::SetMaxDealSeqCnt(uint32_t maxDealSeqCnt) +{ + this->maxDealSeqCnt_ = maxDealSeqCnt; +} + +template +__aicore__ inline bool CompressorSplitCoreSliceIterator::IsEnd() +{ + return (sliceInfo_.bIdx >= batch_size_) || (maxDealSeqCnt_ == 0); +} + +template +__aicore__ inline uint32_t CompressorSplitCoreSliceIterator::GetBIdx() +{ + return sliceInfo_.bIdx; +} + +template +__aicore__ inline void CompressorSplitCoreSliceIterator::IteratorSlice() +{ + bool isUpdateBatchInfo = false; + if (isMaxDealSeqCntFirst) { + isMaxDealSeqCntFirst = false; + } + if (!isFirst_) { + // 更新剩余未处理的行数 + maxDealSeqCnt_ -= sliceInfo_.dealSeqCnt; + // 更新sIdx和bIdx、以及与bIdx相关的bStartPos和bSeqUsed + sliceInfo_.sIdx += sliceInfo_.validSeqCnt; + if (sliceInfo_.sIdx == sliceInfo_.bSeqUsed) { + sliceInfo_.sIdx = 0; + // 左边最后一块跳到b=0 s=0处理 + if (isLeftFirstBath) { + isLeftFirstBath = false; + } else { + sliceInfo_.bIdx++; + } + isUpdateBatchInfo = true; + } + } else { + isUpdateBatchInfo = true; + isFirst_ = false; + } + + // 更新与bIdx相关的bStartPos和bSeqUsed + if (isUpdateBatchInfo) { + // SkipInvalidBatch + while (sliceInfo_.bIdx < batch_size_) { + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + if (sliceInfo_.bSeqUsed > 0) { + break; + } + sliceInfo_.bIdx++; + } + if (sliceInfo_.bIdx < batch_size_) { + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + } + } +} + +template +__aicore__ inline SplitCoreSliceInfo& CompressorSplitCoreSliceIterator::GetLeftNextCmpSeqCnt() +{ + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + if (isFirst_) { + // 左边 T轴首次减去T轴最后一块 + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(batch_size_ - 1); + sliceInfo_.bStartPos = tools_.GetStartPos(batch_size_ - 1); + // 处理最后一块是中间整块或者尾块的情况 + uint32_t lastSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) % cmpRatio == 0 ? + cmpRatio : + (sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) % cmpRatio; + // 处理最后一块是头块的情况 + if (sliceInfo_.bSeqUsed < cmpRatio) { + lastSeqCnt = sliceInfo_.bSeqUsed; + } + + sliceInfo_.sIdx = sliceInfo_.bSeqUsed - lastSeqCnt; + isLeftFirstBath = true; + isFirst_ = false; + } + // 计算头部占位行数、有效数据行数、尾部占位行数 + sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio; + + sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx; + if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) { + sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt; + } + sliceInfo_.tailHolderSeqCnt = + cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio; + if (sliceInfo_.tailHolderSeqCnt == cmpRatio) { + sliceInfo_.tailHolderSeqCnt = 0; + } + + // 头和尾处理，否则需要处理的seq等于cmpRatio + if (sliceInfo_.validSeqCnt < cmpRatio) { + sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt; + if (sliceInfo_.sIdx == 0) { + sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt; + } + } else { + sliceInfo_.dealSeqCnt = cmpRatio; + } + sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt; + + // 计算本次可以处理的Tc个数 + sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio; + + // 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0 + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize; + if (sliceInfo_.tailHolderSeqCnt > 0) { + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时，其不需要压缩 + } + + // 记录左边第一个块 + if (isMaxDealSeqCntFirst) { + sliceInfo_.preFirstSeqCnt = sliceInfo_.dealSeqCnt; + } + + return sliceInfo_; +} + +template +__aicore__ inline SplitCoreSliceInfo& CompressorSplitCoreSliceIterator::GetRightNextCmpSeqCnt() +{ + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + if (isFirst_) { + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + isFirst_ = false; + } + // 计算头部占位行数、有效数据行数、尾部占位行数 + sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio; + + sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx; + if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) { + sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt; + } + sliceInfo_.tailHolderSeqCnt = + cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio; + if (sliceInfo_.tailHolderSeqCnt == cmpRatio) { + sliceInfo_.tailHolderSeqCnt = 0; + } + + // 头和尾处理，否则需要处理的seq等于cmpRatio + if (sliceInfo_.validSeqCnt < cmpRatio) { + sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt; + if (sliceInfo_.sIdx == 0) { + sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt; + } + } else { + sliceInfo_.dealSeqCnt = cmpRatio; + } + sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt; + + // 计算本次可以处理的Tc个数 + sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio; + + // 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0 + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize; + if (sliceInfo_.tailHolderSeqCnt > 0) { + sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时，其不需要压缩 + } + + return sliceInfo_; +} + +struct Vec1SliceInfo : public SliceInfo { + __aicore__ inline Vec1SliceInfo() {}; + __aicore__ inline Vec1SliceInfo(uint32_t bIdx, uint32_t sIdx) : SliceInfo(bIdx, sIdx) {}; + __aicore__ inline Vec1SliceInfo(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt) + : SliceInfo(bIdx, sIdx), dealedSeqCnt(dealedSeqCnt) {}; + + uint32_t dealedSeqCnt = 0U; + uint32_t dealedTcCnt = 0U; + uint32_t bSeqLength = 0U; + uint32_t compressoredScCnt = 0U; + bool isFirst = false; + bool isLast = false; +}; + +struct StatisticInfo { + __aicore__ inline StatisticInfo() {}; + __aicore__ inline StatisticInfo(uint32_t actualTcCnt, uint32_t dealSeqCnt, uint32_t compressorScCnt) + : actualTcCnt(actualTcCnt), dealSeqCnt(dealSeqCnt), compressorScCnt(compressorScCnt) {}; + + uint32_t actualTcCnt = 0U; + uint32_t dealSeqCnt = 0U; + uint32_t compressorScCnt = 0U; +}; + +template +class CompressorVec1SliceIterator { +public: + __aicore__ inline CompressorVec1SliceIterator(CompressorTools &tools) : tools_(tools) {} + + __aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx); + __aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt, uint32_t compressoredScCnt); + __aicore__ inline void SetMaxBatchSize(uint32_t batch_size); + __aicore__ inline void SetDealedSeqCnt(uint32_t dealedSeqCnt); + __aicore__ inline void SetDealedTcCnt(uint32_t dealedTcCnt); + __aicore__ inline void SetCompressoredScCnt(uint32_t compressoredScCnt); + __aicore__ inline void SetNeedDealTcSize(uint32_t needDealTcSize); + __aicore__ inline void SetNeedDealTcSize(uint32_t needDealTcSize, uint32_t canDealTcSize); + __aicore__ inline uint32_t GetNeedDealTcSize(); + __aicore__ inline bool IsEnd(); + template + __aicore__ inline void IteratorSlice(); + __aicore__ inline Vec1SliceInfo &GetSlice(); + template + __aicore__ inline StatisticInfo &FullIteratorSlice(); + +private: + CompressorTools &tools_; + + bool isFirst_ = true; + Vec1SliceInfo sliceInfo_{}; + StatisticInfo statisticInfo_{}; + uint32_t needDealTcSize_ = 0U; + uint32_t batch_size_ = 0U; +}; + +template +__aicore__ inline void CompressorVec1SliceIterator::Reset(uint32_t bIdx, uint32_t sIdx) +{ + sliceInfo_.bIdx = bIdx; + sliceInfo_.sIdx = sIdx; + while (tools_.GetSeqLength(sliceInfo_.bIdx) == 0) { + sliceInfo_.bIdx++; + if (sliceInfo_.bIdx == batch_size_) { + sliceInfo_.bIdx = 0; + } + } + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + sliceInfo_.bSeqLength = tools_.GetSeqLength(sliceInfo_.bIdx); + isFirst_ = true; + } + +template +__aicore__ inline void CompressorVec1SliceIterator::Reset(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt, + uint32_t compressoredScCnt) +{ + Reset(bIdx, sIdx); + SetDealedSeqCnt(dealedSeqCnt); + SetCompressoredScCnt(compressoredScCnt); +} + +template +__aicore__ inline void CompressorVec1SliceIterator::SetMaxBatchSize(uint32_t batch_size) +{ + this->batch_size_ = batch_size; +} + +template +__aicore__ inline void CompressorVec1SliceIterator::SetDealedSeqCnt(uint32_t dealedSeqCnt) +{ + this->sliceInfo_.dealedSeqCnt = dealedSeqCnt; +} + +template +__aicore__ inline void CompressorVec1SliceIterator::SetCompressoredScCnt(uint32_t compressoredScCnt) +{ + this->sliceInfo_.compressoredScCnt = compressoredScCnt; +} + +template +__aicore__ inline void CompressorVec1SliceIterator::SetDealedTcCnt(uint32_t dealedTcCnt) +{ + this->sliceInfo_.dealedTcCnt = dealedTcCnt; +} + +template +__aicore__ inline void CompressorVec1SliceIterator::SetNeedDealTcSize(uint32_t needDealTcSize) +{ + this->needDealTcSize_ = needDealTcSize; +} + +template +template +__aicore__ inline void CompressorVec1SliceIterator::IteratorSlice() +{ + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + if constexpr (IS_STATISTIC) { + statisticInfo_.actualTcCnt += sliceInfo_.dealTcSize; + statisticInfo_.compressorScCnt += sliceInfo_.compressTcSize; + } + needDealTcSize_ -= sliceInfo_.dealTcSize; + sliceInfo_.dealedSeqCnt += sliceInfo_.validSeqCnt; + sliceInfo_.compressoredScCnt += sliceInfo_.compressTcSize; + sliceInfo_.sIdx += sliceInfo_.validSeqCnt; + if (sliceInfo_.sIdx >= sliceInfo_.bSeqUsed) { + do { + uint32_t seqLength = tools_.GetSeqLength(sliceInfo_.bIdx); + if (sliceInfo_.bSeqUsed < seqLength) { + uint32_t nextAlignSIdx = Align(sliceInfo_.bStartPos + sliceInfo_.sIdx, cmpRatio) - sliceInfo_.bStartPos; + sliceInfo_.dealedSeqCnt += nextAlignSIdx - sliceInfo_.sIdx; + uint32_t tcGap = CeilDivT(static_cast(seqLength - nextAlignSIdx), + static_cast(cmpRatio)); + if (sliceInfo_.bSeqUsed == 0 && nextAlignSIdx > sliceInfo_.sIdx) { + // 此时bseqused所在压缩块未被纳入计算 + tcGap++; + } + sliceInfo_.sIdx = nextAlignSIdx; + if (needDealTcSize_ < tcGap) { + sliceInfo_.dealedSeqCnt += needDealTcSize_ * cmpRatio; + sliceInfo_.sIdx += needDealTcSize_ * cmpRatio; + needDealTcSize_ = 0; + break; + } + sliceInfo_.dealedSeqCnt += seqLength - sliceInfo_.sIdx; + needDealTcSize_ -= tcGap; + } + sliceInfo_.bIdx++; + if (sliceInfo_.bIdx == batch_size_) { + sliceInfo_.bIdx = 0; + } + sliceInfo_.sIdx = 0; + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + } while (sliceInfo_.bSeqUsed == 0); + sliceInfo_.bSeqLength = tools_.GetSeqLength(sliceInfo_.bIdx); + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + } + if (isFirst_) { + isFirst_ = false; + } +} + +template +__aicore__ inline uint32_t CompressorVec1SliceIterator::GetNeedDealTcSize() +{ + return needDealTcSize_; +} + + +template +__aicore__ inline bool CompressorVec1SliceIterator::IsEnd() +{ + return (needDealTcSize_ == 0); +} + +template +__aicore__ inline Vec1SliceInfo& CompressorVec1SliceIterator::GetSlice() +{ + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + if (sliceInfo_.bSeqUsed < sliceInfo_.sIdx) { + sliceInfo_.headHolderSeqCnt = 0; + sliceInfo_.validSeqCnt = 0; + sliceInfo_.tailHolderSeqCnt = 0; + sliceInfo_.dealTcSize = 0; + sliceInfo_.compressTcSize = 0; + } else { + // 计算头部占位行数、有效数据行数、尾部占位行数 + sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio; + sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx; + if (CeilDivT(sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt, cmpRatio) > needDealTcSize_) { + sliceInfo_.validSeqCnt = needDealTcSize_ * cmpRatio - sliceInfo_.headHolderSeqCnt; + } + uint32_t globalTotalSeqCnt = sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt; + sliceInfo_.tailHolderSeqCnt = Align(globalTotalSeqCnt, cmpRatio) - globalTotalSeqCnt; + + // 计算本次可以处理的Tc个数 + sliceInfo_.dealTcSize = + (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt + sliceInfo_.tailHolderSeqCnt) / cmpRatio; + + sliceInfo_.compressTcSize = + (sliceInfo_.headHolderSeqCnt + min(sliceInfo_.validSeqCnt, sliceInfo_.bSeqUsed - sliceInfo_.sIdx)) / + cmpRatio; + } + + sliceInfo_.isFirst = isFirst_; + sliceInfo_.isLast = + sliceInfo_.bSeqUsed > sliceInfo_.sIdx && + CeilDivT(sliceInfo_.headHolderSeqCnt + sliceInfo_.bSeqUsed - sliceInfo_.sIdx, cmpRatio) >= needDealTcSize_; + + return sliceInfo_; +} + +template +template +__aicore__ inline StatisticInfo& CompressorVec1SliceIterator::FullIteratorSlice() +{ + if constexpr (IS_STATISTIC) { + statisticInfo_ = {0U, 0U, 0U}; + Vec1SliceInfo tempSliceInfo = GetSlice(); + while (!IsEnd()) { + GetSlice(); + IteratorSlice(); + } + Vec1SliceInfo sliceInfo = GetSlice(); + statisticInfo_.dealSeqCnt = sliceInfo.dealedSeqCnt - tempSliceInfo.dealedSeqCnt; + } else { + while (!IsEnd()) { + GetSlice(); + IteratorSlice(); + } + } + return statisticInfo_; +} + +struct Vec2SliceInfo{ + __aicore__ inline Vec2SliceInfo(){}; + __aicore__ inline Vec2SliceInfo(uint32_t bIdx, uint32_t scIdx) : bIdx(bIdx), scIdx(scIdx) + { + } + + uint32_t bIdx = 0U; + uint32_t scIdx = 0U; + uint32_t scNum = 0U; + uint32_t remainScCnt = 0U; // 当前batch剩余sc数量 + uint32_t bStartPos = 0U; + uint32_t bSeqUsed = 0U; + uint32_t bSeqLength = 0U; + uint32_t dealedScCnt = 0U; // 全局的dealedScCnt（Reset刷新） + uint32_t curDealScNum = 0U; // 当前循环处理的sc数量（IteratorSlice刷新） + uint32_t bOutputScLen = 0U; // BSH场景每个batch填充后的输出长度 + uint32_t padScIdx = 0U; // 当前sc输出位置，TH场景为全局的dealedScCnt，BSH场景则为填充后全局的索引（Reset刷新） + uint32_t loopDealedScCnt = 0U; // 当前迭代已处理的sc数量（Reset刷新） +}; + + +template +class CompressorVec2SliceIterator { +public: + __aicore__ inline CompressorVec2SliceIterator(CompressorTools &tools) : tools_(tools) + { + } + __aicore__ inline void Reset(uint32_t bIdx, uint32_t scIdx, uint32_t dealedScCnt); + __aicore__ inline void SetMaxBatchSize(uint32_t batch_size); + __aicore__ inline void SetNeedDealScSize(uint32_t needDealScSize); + __aicore__ inline void ResetLoopDealedScCnt(); + __aicore__ inline uint32_t GetNeedDealScSize(); + __aicore__ inline bool IsEnd(); + __aicore__ inline void IteratorSlice(); + __aicore__ inline Vec2SliceInfo &GetSlice(); +private: + CompressorTools &tools_; + + Vec2SliceInfo sliceInfo_{}; + uint32_t needDealScSize_ = 0U; + uint32_t batch_size_ = 0U; +}; + + +template +__aicore__ inline void CompressorVec2SliceIterator::Reset(uint32_t bIdx, uint32_t scIdx, uint32_t dealedScCnt) +{ + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + sliceInfo_.bIdx = bIdx; + sliceInfo_.scIdx = scIdx; + sliceInfo_.dealedScCnt = dealedScCnt; + if constexpr (COMP::xLayout == X_LAYOUT::BSH) { + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + sliceInfo_.scNum = (sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) / cmpRatio - sliceInfo_.bStartPos / cmpRatio; + sliceInfo_.remainScCnt = sliceInfo_.scNum - sliceInfo_.scIdx; + sliceInfo_.bOutputScLen = CeilDivT(tools_.GetSeqLength(sliceInfo_.bIdx), cmpRatio); + sliceInfo_.padScIdx = sliceInfo_.bIdx * sliceInfo_.bOutputScLen + sliceInfo_.scIdx; + } else { + sliceInfo_.padScIdx = sliceInfo_.dealedScCnt; + } + sliceInfo_.loopDealedScCnt = 0U; +} + +template +__aicore__ inline void CompressorVec2SliceIterator::ResetLoopDealedScCnt() +{ + sliceInfo_.loopDealedScCnt = 0U; +} + + + +template +__aicore__ inline void CompressorVec2SliceIterator::SetMaxBatchSize(uint32_t batch_size) +{ + this->batch_size_ = batch_size; +} + + +template +__aicore__ inline void CompressorVec2SliceIterator::SetNeedDealScSize(uint32_t needDealScSize) +{ + this->needDealScSize_ = needDealScSize; +} + +template +__aicore__ inline void CompressorVec2SliceIterator::IteratorSlice() +{ + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + sliceInfo_.padScIdx += sliceInfo_.curDealScNum; + } else { + if (needDealScSize_ <= sliceInfo_.remainScCnt) { + sliceInfo_.scIdx += sliceInfo_.curDealScNum; + sliceInfo_.padScIdx += sliceInfo_.curDealScNum; + } else { + uint32_t cmpRatio = tools_.toolParams_.cmpRatio; + sliceInfo_.padScIdx += sliceInfo_.bOutputScLen - sliceInfo_.scIdx; + sliceInfo_.bIdx++; + sliceInfo_.scIdx = 0; + sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx); + sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx); + sliceInfo_.scNum = (sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) / cmpRatio - sliceInfo_.bStartPos / cmpRatio; + } + sliceInfo_.remainScCnt = sliceInfo_.scNum - sliceInfo_.scIdx; + } + sliceInfo_.dealedScCnt += sliceInfo_.curDealScNum; + needDealScSize_ -= sliceInfo_.curDealScNum; + sliceInfo_.loopDealedScCnt += sliceInfo_.curDealScNum; +} + +template +__aicore__ inline uint32_t CompressorVec2SliceIterator::GetNeedDealScSize() +{ + return needDealScSize_; +} + + +template +__aicore__ inline bool CompressorVec2SliceIterator::IsEnd() +{ + return (needDealScSize_ == 0); +} + +template +__aicore__ inline Vec2SliceInfo &CompressorVec2SliceIterator::GetSlice() +{ + if constexpr (COMP::xLayout == X_LAYOUT::TH) { + sliceInfo_.curDealScNum = needDealScSize_; + } else { + sliceInfo_.curDealScNum = min(sliceInfo_.remainScCnt, needDealScSize_); + } + return sliceInfo_; +} + + + +} // namespace Compressor + +#endif diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_add.h b/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_add.h new file mode 100644 index 000000000..7a28424bb --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_add.h @@ -0,0 +1,378 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_add.h + * \brief + */ + +#ifndef VF_ADD_H +#define VF_ADD_H + +#include "kernel_operator.h" +#include +using namespace AscendC; +constexpr uint32_t FLOAT_REP_SIZE = 64; +constexpr uint32_t BTYEALIGNSIZE = 32; +constexpr uint32_t REGSIZE = 256; +constexpr uint32_t HALFCORED = 128; + +template +struct AddRegList { + MicroAPI::RegTensor vreg; + MicroAPI::RegTensor vregape; +}; + + +template +__simd_callee__ void AddVFImpl(__ubuf__ T *inputAddr, __ubuf__ T *apeAddr, AddRegList ®List, uint32_t row, + uint32_t col, uint64_t offset0, uint64_t offset1) +{ + uint32_t maskValue = col; + MicroAPI::MaskReg mask = MicroAPI::UpdateMask(maskValue); + MicroAPI::LoadAlign(regList.vreg, inputAddr + offset0); + MicroAPI::LoadAlign(regList.vregape, apeAddr + offset1); + MicroAPI::Add(regList.vreg, regList.vreg, regList.vregape, mask); + MicroAPI::StoreAlign(inputAddr + offset0, regList.vreg, mask); +} + +template +__simd_callee__ void MultiAddVFImpl(__ubuf__ T *outputAddr, __ubuf__ T *inputAddr, AddRegList ®List, uint32_t row, + uint32_t col, uint64_t offset, uint32_t repeatNum, uint64_t repeatOffset) +{ + uint32_t maskValue = col; + uint32_t initialRepeatIdx = IS_FIRST ? 1 : 0; + __ubuf__ T *initialAddr = IS_FIRST ? inputAddr : outputAddr; + MicroAPI::MaskReg mask = MicroAPI::UpdateMask(maskValue); + MicroAPI::LoadAlign(regList.vreg, initialAddr + offset); + for (uint32_t repeatIdx = initialRepeatIdx; repeatIdx < repeatNum; repeatIdx++) { + uint64_t addOffset = offset + repeatIdx * repeatOffset; + MicroAPI::LoadAlign(regList.vregape, inputAddr + addOffset); + MicroAPI::Add(regList.vreg, regList.vreg, regList.vregape, mask); + } + MicroAPI::StoreAlign(outputAddr + offset, regList.vreg, mask); +} + +template +__simd_vf__ void Add64VFImpl(__ubuf__ T *inputAddr, __ubuf__ T *apeAddr, uint32_t row, uint32_t col, uint32_t actualCol0, uint32_t actualCol1) +{ + AddRegList regList[4]; + uint32_t loopTimes = row / 4; + for (uint32_t idx = 0; idx < loopTimes; idx++) { + uint64_t offset0 = idx * 4 * actualCol0; + uint64_t offset1 = idx * 4 * actualCol1; + AddVFImpl(inputAddr, apeAddr, regList[0], row, col, offset0, offset1); + AddVFImpl(inputAddr, apeAddr, regList[1], row, col, offset0 + actualCol0, offset1 + actualCol1); + AddVFImpl(inputAddr, apeAddr, regList[2], row, col, offset0 + 2 * actualCol0, offset1 + 2 * actualCol1); + AddVFImpl(inputAddr, apeAddr, regList[3], row, col, offset0 + 3 * actualCol0, offset1 + 3 * actualCol1); + } + + if (row % 4 > 0) { + AddVFImpl(inputAddr, apeAddr, regList[0], row, col, loopTimes * 4 * actualCol0, loopTimes * 4 * actualCol1); + } + + if (row % 4 > 1) { + AddVFImpl(inputAddr, apeAddr, regList[1], row, col, (loopTimes * 4 + 1) * actualCol0, (loopTimes * 4 + 1) * actualCol1); + } + + if (row % 4 > 2) { + AddVFImpl(inputAddr, apeAddr, regList[2], row, col, (loopTimes * 4 + 2) * actualCol0, (loopTimes * 4 + 2) * actualCol1); + } +} + +template +__simd_vf__ void Add128VFImpl(__ubuf__ T *inputAddr, __ubuf__ T *apeAddr, uint32_t row, uint32_t actualCol0, uint32_t actualCol1) +{ + AddRegList regList[4]; + uint32_t loopTimes = row / 2; + for (uint32_t idx = 0; idx < loopTimes; idx++) { + uint64_t offset0 = idx * 2 * actualCol0; + uint64_t offset1 = idx * 2 * actualCol1; + AddVFImpl(inputAddr, apeAddr, regList[0], row, FLOAT_REP_SIZE, offset0, offset1); + AddVFImpl(inputAddr, apeAddr, regList[1], row, FLOAT_REP_SIZE, offset0 + FLOAT_REP_SIZE, offset1 + FLOAT_REP_SIZE); + AddVFImpl(inputAddr, apeAddr, regList[2], row, FLOAT_REP_SIZE, offset0 + actualCol0, offset1 + actualCol1); + AddVFImpl(inputAddr, apeAddr, regList[3], row, FLOAT_REP_SIZE, offset0 + actualCol0 + FLOAT_REP_SIZE, offset1 + actualCol1 + FLOAT_REP_SIZE); + } + + if (row % 2 > 0) { + AddVFImpl(inputAddr, apeAddr, regList[0], row, FLOAT_REP_SIZE, loopTimes * 2 * actualCol0, loopTimes * 2 * actualCol1); + AddVFImpl(inputAddr, apeAddr, regList[1], row, FLOAT_REP_SIZE, loopTimes * 2 * actualCol0 + FLOAT_REP_SIZE, loopTimes * 2 * actualCol1 + FLOAT_REP_SIZE); + } +} + +template +__simd_vf__ void Add256VFImpl(__ubuf__ T *inputAddr, __ubuf__ T *apeAddr, uint32_t row, uint32_t actualCol0, uint32_t actualCol1) +{ + AddRegList regList[4]; + MicroAPI::MaskReg mask = MicroAPI::CreateMask(); + for (uint32_t idx = 0; idx < row; idx++) { + uint64_t offset0 = idx * actualCol0; + uint64_t offset1 = idx * actualCol1; + AddVFImpl(inputAddr, apeAddr, regList[0], row, FLOAT_REP_SIZE, offset0, offset1); + AddVFImpl(inputAddr, apeAddr, regList[1], row, FLOAT_REP_SIZE, offset0 + FLOAT_REP_SIZE, offset1 + FLOAT_REP_SIZE); + AddVFImpl(inputAddr, apeAddr, regList[2], row, FLOAT_REP_SIZE, offset0 + 2 * FLOAT_REP_SIZE, offset1 + 2 * FLOAT_REP_SIZE); + AddVFImpl(inputAddr, apeAddr, regList[3], row, FLOAT_REP_SIZE, offset0 + 3 * FLOAT_REP_SIZE, offset1 + 3 * FLOAT_REP_SIZE); + } +} + +template +__simd_vf__ void Add512VFImpl(__ubuf__ T *inputAddr, __ubuf__ T *apeAddr, uint32_t row, uint32_t actualCol0, uint32_t actualCol1) +{ + AddRegList regList[8]; + MicroAPI::MaskReg mask = MicroAPI::CreateMask(); + for (uint32_t idx = 0; idx < row; idx++) { + uint64_t offset0 = idx * actualCol0; + uint64_t offset1 = idx * actualCol1; + AddVFImpl(inputAddr, apeAddr, regList[0], row, FLOAT_REP_SIZE, offset0, offset1); + AddVFImpl(inputAddr, apeAddr, regList[1], row, FLOAT_REP_SIZE, offset0 + FLOAT_REP_SIZE, offset1 + FLOAT_REP_SIZE); + AddVFImpl(inputAddr, apeAddr, regList[2], row, FLOAT_REP_SIZE, offset0 + 2 * FLOAT_REP_SIZE, offset1 + 2 * FLOAT_REP_SIZE); + AddVFImpl(inputAddr, apeAddr, regList[3], row, FLOAT_REP_SIZE, offset0 + 3 * FLOAT_REP_SIZE, offset1 + 3 * FLOAT_REP_SIZE); + AddVFImpl(inputAddr, apeAddr, regList[4], row, FLOAT_REP_SIZE, offset0 + 4 * FLOAT_REP_SIZE, offset1 + 4 * FLOAT_REP_SIZE); + AddVFImpl(inputAddr, apeAddr, regList[5], row, FLOAT_REP_SIZE, offset0 + 5 * FLOAT_REP_SIZE, offset1 + 5 * FLOAT_REP_SIZE); + AddVFImpl(inputAddr, apeAddr, regList[6], row, FLOAT_REP_SIZE, offset0 + 6 * FLOAT_REP_SIZE, offset1 + 6 * FLOAT_REP_SIZE); + AddVFImpl(inputAddr, apeAddr, regList[7], row, FLOAT_REP_SIZE, offset0 + 7 * FLOAT_REP_SIZE, offset1 + 7 * FLOAT_REP_SIZE); + } +} + +template +__simd_vf__ void MultiAdd64VFImpl(__ubuf__ T *outputAddr, __ubuf__ T *inputAddr, uint32_t row, uint32_t col, + uint32_t actualCol, uint32_t repeatNum, uint64_t repeatOffset) +{ + AddRegList regList[4]; + uint32_t loopTimes = row / 4; + uint32_t maskValue = col; + uint32_t initialRepeatIdx = IS_FIRST ? 1 : 0; + __ubuf__ T *initialAddr = IS_FIRST ? inputAddr : outputAddr; + MicroAPI::MaskReg mask = MicroAPI::UpdateMask(maskValue); + for (uint32_t idx = 0; idx < loopTimes; idx++) { + uint64_t offset = idx * 4 * actualCol; + MicroAPI::LoadAlign(regList[0].vreg, initialAddr + offset); + MicroAPI::LoadAlign(regList[1].vreg, initialAddr + offset + actualCol); + MicroAPI::LoadAlign(regList[2].vreg, initialAddr + offset + 2 * actualCol); + MicroAPI::LoadAlign(regList[3].vreg, initialAddr + offset + 3 * actualCol); + for (uint32_t repeatIdx = initialRepeatIdx; repeatIdx < repeatNum; repeatIdx++) { + uint64_t addOffset = offset + repeatIdx * repeatOffset; + MicroAPI::LoadAlign(regList[0].vregape, inputAddr + addOffset); + MicroAPI::LoadAlign(regList[1].vregape, inputAddr + addOffset + actualCol); + MicroAPI::LoadAlign(regList[2].vregape, inputAddr + addOffset + 2 * actualCol); + MicroAPI::LoadAlign(regList[3].vregape, inputAddr + addOffset + 3 * actualCol); + MicroAPI::Add(regList[0].vreg, regList[0].vreg, regList[0].vregape, mask); + MicroAPI::Add(regList[1].vreg, regList[1].vreg, regList[1].vregape, mask); + MicroAPI::Add(regList[2].vreg, regList[2].vreg, regList[2].vregape, mask); + MicroAPI::Add(regList[3].vreg, regList[3].vreg, regList[3].vregape, mask); + } + MicroAPI::StoreAlign(outputAddr + offset, regList[0].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + actualCol, regList[1].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + 2 * actualCol, regList[2].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + 3 * actualCol, regList[3].vreg, mask); + } + + if (row % 4 > 0) { + MultiAddVFImpl(outputAddr, inputAddr, regList[0], row, col, loopTimes * 4 * actualCol, repeatNum, + repeatOffset); + } + + if (row % 4 > 1) { + MultiAddVFImpl(outputAddr, inputAddr, regList[1], row, col, (loopTimes * 4 + 1) * actualCol, + repeatNum, repeatOffset); + } + + if (row % 4 > 2) { + MultiAddVFImpl(outputAddr, inputAddr, regList[2], row, col, (loopTimes * 4 + 2) * actualCol, + repeatNum, repeatOffset); + } +} + +template +__simd_vf__ void MultiAdd128VFImpl(__ubuf__ T *outputAddr, __ubuf__ T *inputAddr, uint32_t row, uint32_t col, + uint32_t actualCol, uint32_t repeatNum, uint64_t repeatOffset) +{ + AddRegList regList[4]; + uint32_t loopTimes = row / 2; + uint32_t initialRepeatIdx = IS_FIRST ? 1 : 0; + __ubuf__ T *initialAddr = IS_FIRST ? inputAddr : outputAddr; + MicroAPI::MaskReg mask = MicroAPI::CreateMask(); + for (uint32_t idx = 0; idx < loopTimes; idx++) { + uint64_t offset = idx * actualCol * 2; + MicroAPI::LoadAlign(regList[0].vreg, initialAddr + offset); + MicroAPI::LoadAlign(regList[1].vreg, initialAddr + offset + FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[2].vreg, initialAddr + offset + actualCol); + MicroAPI::LoadAlign(regList[3].vreg, initialAddr + offset + actualCol + FLOAT_REP_SIZE); + for (uint32_t repeatIdx = initialRepeatIdx; repeatIdx < repeatNum; repeatIdx++) { + uint64_t addOffset = offset + repeatIdx * repeatOffset; + MicroAPI::LoadAlign(regList[0].vregape, inputAddr + addOffset); + MicroAPI::LoadAlign(regList[1].vregape, inputAddr + addOffset + FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[2].vregape, inputAddr + addOffset + actualCol); + MicroAPI::LoadAlign(regList[3].vregape, inputAddr + addOffset + actualCol + FLOAT_REP_SIZE); + MicroAPI::Add(regList[0].vreg, regList[0].vreg, regList[0].vregape, mask); + MicroAPI::Add(regList[1].vreg, regList[1].vreg, regList[1].vregape, mask); + MicroAPI::Add(regList[2].vreg, regList[2].vreg, regList[2].vregape, mask); + MicroAPI::Add(regList[3].vreg, regList[3].vreg, regList[3].vregape, mask); + } + MicroAPI::StoreAlign(outputAddr + offset, regList[0].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + FLOAT_REP_SIZE, regList[1].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + actualCol, regList[2].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + actualCol + FLOAT_REP_SIZE, regList[3].vreg, mask); + } + + if (row % 2 > 0) { + MultiAddVFImpl(outputAddr, inputAddr, regList[0], row, col, loopTimes * 2 * actualCol, repeatNum, + repeatOffset); + MultiAddVFImpl(outputAddr, inputAddr, regList[1], row, col, + loopTimes * 2 * actualCol + FLOAT_REP_SIZE, repeatNum, repeatOffset); + } +} + +template +__simd_vf__ void MultiAdd256VFImpl(__ubuf__ T *outputAddr, __ubuf__ T *inputAddr, uint32_t row, + uint32_t actualCol, uint32_t repeatNum, uint64_t repeatOffset) +{ + AddRegList regList[4]; + uint32_t loopTimes = row; + uint32_t initialRepeatIdx = IS_FIRST ? 1 : 0; + __ubuf__ T *initialAddr = IS_FIRST ? inputAddr : outputAddr; + MicroAPI::MaskReg mask = MicroAPI::CreateMask(); + for (uint32_t idx = 0; idx < loopTimes; idx++) { + uint64_t offset = idx * actualCol; + MicroAPI::LoadAlign(regList[0].vreg, initialAddr + offset); + MicroAPI::LoadAlign(regList[1].vreg, initialAddr + offset + FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[2].vreg, initialAddr + offset + 2 * FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[3].vreg, initialAddr + offset + 3 * FLOAT_REP_SIZE); + for (uint32_t repeatIdx = initialRepeatIdx; repeatIdx < repeatNum; repeatIdx++) { + uint64_t addOffset = offset + repeatIdx * repeatOffset; + MicroAPI::LoadAlign(regList[0].vregape, inputAddr + addOffset); + MicroAPI::LoadAlign(regList[1].vregape, inputAddr + addOffset + FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[2].vregape, inputAddr + addOffset + 2 * FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[3].vregape, inputAddr + addOffset + 3 * FLOAT_REP_SIZE); + MicroAPI::Add(regList[0].vreg, regList[0].vreg, regList[0].vregape, mask); + MicroAPI::Add(regList[1].vreg, regList[1].vreg, regList[1].vregape, mask); + MicroAPI::Add(regList[2].vreg, regList[2].vreg, regList[2].vregape, mask); + MicroAPI::Add(regList[3].vreg, regList[3].vreg, regList[3].vregape, mask); + } + MicroAPI::StoreAlign(outputAddr + offset, regList[0].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + FLOAT_REP_SIZE, regList[1].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + 2 * FLOAT_REP_SIZE, regList[2].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + 3 * FLOAT_REP_SIZE, regList[3].vreg, mask); + } +} + +template +__simd_vf__ void MultiAdd512VFImpl(__ubuf__ T *outputAddr, __ubuf__ T *inputAddr, uint32_t row, + uint32_t actualCol, uint32_t repeatNum, uint64_t repeatOffset) +{ + AddRegList regList[8]; + uint32_t loopTimes = row; + uint32_t initialRepeatIdx = IS_FIRST ? 1 : 0; + __ubuf__ T *initialAddr = IS_FIRST ? inputAddr : outputAddr; + MicroAPI::MaskReg mask = MicroAPI::CreateMask(); + for (uint32_t idx = 0; idx < loopTimes; idx++) { + uint64_t offset = idx * actualCol; + MicroAPI::LoadAlign(regList[0].vreg, initialAddr + offset); + MicroAPI::LoadAlign(regList[1].vreg, initialAddr + offset + FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[2].vreg, initialAddr + offset + 2 * FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[3].vreg, initialAddr + offset + 3 * FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[4].vreg, initialAddr + offset + 4 * FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[5].vreg, initialAddr + offset + 5 * FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[6].vreg, initialAddr + offset + 6 * FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[7].vreg, initialAddr + offset + 7 * FLOAT_REP_SIZE); + for (uint32_t repeatIdx = initialRepeatIdx; repeatIdx < repeatNum; repeatIdx++) { + uint64_t addOffset = offset + repeatIdx * row * actualCol; + MicroAPI::LoadAlign(regList[0].vregape, inputAddr + addOffset); + MicroAPI::LoadAlign(regList[1].vregape, inputAddr + addOffset + FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[2].vregape, inputAddr + addOffset + 2 * FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[3].vregape, inputAddr + addOffset + 3 * FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[4].vregape, inputAddr + addOffset + 4 * FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[5].vregape, inputAddr + addOffset + 5 * FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[6].vregape, inputAddr + addOffset + 6 * FLOAT_REP_SIZE); + MicroAPI::LoadAlign(regList[7].vregape, inputAddr + addOffset + 7 * FLOAT_REP_SIZE); + MicroAPI::Add(regList[0].vreg, regList[0].vreg, regList[0].vregape, mask); + MicroAPI::Add(regList[1].vreg, regList[1].vreg, regList[1].vregape, mask); + MicroAPI::Add(regList[2].vreg, regList[2].vreg, regList[2].vregape, mask); + MicroAPI::Add(regList[3].vreg, regList[3].vreg, regList[3].vregape, mask); + MicroAPI::Add(regList[4].vreg, regList[4].vreg, regList[4].vregape, mask); + MicroAPI::Add(regList[5].vreg, regList[5].vreg, regList[5].vregape, mask); + MicroAPI::Add(regList[6].vreg, regList[6].vreg, regList[6].vregape, mask); + MicroAPI::Add(regList[7].vreg, regList[7].vreg, regList[7].vregape, mask); + } + MicroAPI::StoreAlign(outputAddr + offset, regList[0].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + FLOAT_REP_SIZE, regList[1].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + 2 * FLOAT_REP_SIZE, regList[2].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + 3 * FLOAT_REP_SIZE, regList[3].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + 4 * FLOAT_REP_SIZE, regList[4].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + 5 * FLOAT_REP_SIZE, regList[5].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + 6 * FLOAT_REP_SIZE, regList[6].vreg, mask); + MicroAPI::StoreAlign(outputAddr + offset + 7 * FLOAT_REP_SIZE, regList[7].vreg, mask); + } +} + +/** + * @brief AddVF 输入与apt相加 + * @param rightLocal 输出tensor [] + * @param leftLocal 输入tensor [row, col] + * @param aptLocal apt输入tensor [r] + * @param apeIdx ape起始位置 + * @param d coff*d为ape的D轴大小 + * @param coreSplitD scoreleft大小，coff*coreSplitD为总大小 + * @param coreSplitS 核间d轴切分大小 + */ +template +__aicore__ inline void AddVF(const LocalTensor &scoreLocal, const LocalTensor &apeLocal, uint32_t row, + uint32_t col, uint32_t actualCol0, uint32_t actualCol1) +{ + __ubuf__ T *scoreAddr = (__ubuf__ T *)scoreLocal.GetPhyAddr(); + __ubuf__ T *apeAddr = (__ubuf__ T *)apeLocal.GetPhyAddr(); + + if (col <= 64) { + Add64VFImpl(scoreAddr, apeAddr, row, col, actualCol0, actualCol1); + } else if (col == 128) { + Add128VFImpl(scoreAddr, apeAddr, row, actualCol0, actualCol1); + } else if (col == 256) { + Add256VFImpl(scoreAddr, apeAddr, row, actualCol0, actualCol1); + } else if (col == 512) { + Add512VFImpl(scoreAddr, apeAddr, row, actualCol0, actualCol1); + } +} + +template +__aicore__ inline void AddVF(const LocalTensor &scoreLocal, const LocalTensor &apeLocal, uint32_t row, + uint32_t col, uint32_t actualCol) +{ + __ubuf__ T *scoreAddr = (__ubuf__ T *)scoreLocal.GetPhyAddr(); + __ubuf__ T *apeAddr = (__ubuf__ T *)apeLocal.GetPhyAddr(); + + if (col <= 64) { + Add64VFImpl(scoreAddr, apeAddr, row, col, actualCol, actualCol); + } else if (col == 128) { + Add128VFImpl(scoreAddr, apeAddr, row, actualCol, actualCol); + } else if (col == 256) { + Add256VFImpl(scoreAddr, apeAddr, row, actualCol, actualCol); + } else if (col == 512) { + Add512VFImpl(scoreAddr, apeAddr, row, actualCol, actualCol); + } +} + +template +__aicore__ inline void MultiAddVF(const LocalTensor &outputLocal, const LocalTensor &inputLocal, uint32_t row, + uint32_t col, uint32_t actualCol, uint32_t repeatNum, uint64_t repeatOffset) +{ + __ubuf__ T *outputAddr = (__ubuf__ T *)outputLocal.GetPhyAddr(); + __ubuf__ T *inputAddr = (__ubuf__ T *)inputLocal.GetPhyAddr(); + if (col <= 64) { + MultiAdd64VFImpl(outputAddr, inputAddr, row, col, actualCol, repeatNum, repeatOffset); + } else if (col == 128) { + MultiAdd128VFImpl(outputAddr, inputAddr, row, col, actualCol, repeatNum, repeatOffset); + } else if (col == 256) { + MultiAdd256VFImpl(outputAddr, inputAddr, row, actualCol, repeatNum, repeatOffset); + } else if (col == 512) { + MultiAdd512VFImpl(outputAddr, inputAddr, row, actualCol, repeatNum, repeatOffset); + } +} + +#endif diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_mul.h b/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_mul.h new file mode 100644 index 000000000..3528300d5 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_mul.h @@ -0,0 +1,318 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_mul.h + * \brief + */ + +#ifndef VF_MUL_H +#define VF_MUL_H + +#include "kernel_operator.h" +#include +using namespace AscendC; + +constexpr uint32_t FLOATBYTE = 4; +constexpr uint32_t baseD8 = 8; +constexpr uint32_t baseD16 = 16; +constexpr uint32_t baseD32 = 32; +constexpr uint32_t baseD64 = 64; +constexpr uint32_t baseD128 = 128; +constexpr uint32_t baseD256 = 256; +constexpr uint32_t baseD512 = 512; + + +template +__simd_callee__ inline T SimdCeilDivT(T num1, T num2) +{ + if (num2 == 0) { + return static_cast(0); + } + return (num1 + num2 - 1) / num2; +} + +template +struct ReduceMulRegList { + MicroAPI::RegTensor vreg0; + MicroAPI::RegTensor vreg1; + MicroAPI::RegTensor vregMul; + MicroAPI::RegTensor vregSum; +}; + + +template +__simd_callee__ void LoadMulAddVFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, ReduceMulRegList ®List, uint64_t offset, uint32_t maskValue) +{ + MicroAPI::MaskReg mask = MicroAPI::UpdateMask(maskValue); + MicroAPI::LoadAlign(regList.vreg0, kvAddr + offset); + MicroAPI::LoadAlign(regList.vreg1, scoreAddr + offset); + MicroAPI::Mul(regList.vregMul, regList.vreg0, regList.vreg1, mask); + MicroAPI::Add(regList.vregSum, regList.vregSum, regList.vregMul, mask); +} + + + +template +__simd_vf__ void MulReduceSumbase8VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr, + const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt, + const uint32_t baseD) +{ + ReduceMulRegList regList; + MicroAPI::RegTensor vregSum0; + MicroAPI::MaskReg mask = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskL32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskL16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskL8 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskH32; + MicroAPI::MaskReg maskH48; + MicroAPI::MaskReg maskH56; + MicroAPI::Not(maskH48, maskL16, mask); + MicroAPI::Not(maskH32, maskL32, mask); + MicroAPI::Not(maskH56, maskL8, mask); + uint32_t offset = 0; + uint32_t rCnt = coff * cmpRatio; + for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) { + MicroAPI::Duplicate(regList.vregSum, 0, mask); + // 当前仅支持coff * cmpRatio为2的幂的情况 + for (uint32_t rLoop = 0; rLoop < SimdCeilDivT(rCnt, 8U); rLoop++) { + uint32_t dealLen = min((rCnt - rLoop * 8) * baseD, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList, offset, dealLen); + offset += dealLen; + } + // 64 -> 32 + MicroAPI::Squeeze(vregSum0, regList.vregSum, maskH32); + MicroAPI::Add(regList.vregSum, regList.vregSum, vregSum0, maskL32); + + // 32 -> 16 + MicroAPI::Squeeze(vregSum0, regList.vregSum, maskH48); + MicroAPI::Add(regList.vregSum, regList.vregSum, vregSum0, maskL16); + + // 16 -> 8 + MicroAPI::Squeeze(vregSum0, regList.vregSum, maskH56); + MicroAPI::Add(regList.vregSum, regList.vregSum, vregSum0, maskL8); + + MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList.vregSum, maskL8); + } +} + +template +__simd_vf__ void MulReduceSumbase16VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr, + const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt, + const uint32_t baseD) +{ + ReduceMulRegList regList; + MicroAPI::RegTensor vregSum0; + MicroAPI::MaskReg mask = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskL32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskL16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskH32; + MicroAPI::MaskReg maskH48; + MicroAPI::Not(maskH48, maskL16, mask); + MicroAPI::Not(maskH32, maskL32, mask); + uint32_t offset = 0; + uint32_t rCnt = coff * cmpRatio; + for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) { + MicroAPI::Duplicate(regList.vregSum, 0, mask); + // 当前仅支持coff * cmpRatio为2的幂的情况 + for (uint32_t rLoop = 0; rLoop < SimdCeilDivT(rCnt, 4U); rLoop++) { + uint32_t dealLen = min((rCnt - rLoop * 4) * baseD, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList, offset, dealLen); + offset += dealLen; + } + // 64 -> 32 + MicroAPI::Squeeze(vregSum0, regList.vregSum, maskH32); + MicroAPI::Add(regList.vregSum, regList.vregSum, vregSum0, maskL32); + + // 32 -> 16 + MicroAPI::Squeeze(vregSum0, regList.vregSum, maskH48); + MicroAPI::Add(regList.vregSum, regList.vregSum, vregSum0, maskL16); + + MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList.vregSum, maskL16); + } +} + +template +__simd_vf__ void MulReduceSumbase32VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr, + const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt, + const uint32_t baseD) +{ + ReduceMulRegList regList; + MicroAPI::RegTensor vregSum0; + MicroAPI::RegTensor vregSum1; + MicroAPI::MaskReg mask = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskL32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskH32; + MicroAPI::Not(maskH32, maskL32, mask); + uint32_t offset = 0; + uint32_t rCnt = coff * cmpRatio; + for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) { + MicroAPI::Duplicate(regList.vregSum, 0, mask); + // 当前仅支持coff * cmpRatio为2的幂的情况 + for (uint32_t rLoop = 0; rLoop < SimdCeilDivT(rCnt, 2U); rLoop++) { + uint32_t dealLen = min((rCnt - rLoop * 2) * baseD, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList, offset, dealLen); + offset += dealLen; + } + // 64 -> 32 + MicroAPI::Squeeze(vregSum0, regList.vregSum, maskH32); + MicroAPI::Add(regList.vregSum, regList.vregSum, vregSum0, maskL32); + + MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList.vregSum, maskL32); + } +} + +template +__simd_vf__ void MulReduceSumbase64VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr, + const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt, + const uint32_t baseD) +{ + ReduceMulRegList regList; + MicroAPI::MaskReg mask = MicroAPI::CreateMask(); + uint32_t offset = 0; + uint32_t rCnt = coff * cmpRatio; + for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) { + MicroAPI::Duplicate(regList.vregSum, 0, mask); + for (uint32_t rLoop = 0; rLoop < rCnt; rLoop++) { + LoadMulAddVFImpl(kvAddr, scoreAddr, regList, offset, baseD64); + offset += baseD; + } + MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList.vregSum, mask); + } +} + +template +__simd_vf__ void MulReduceSumbase128VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr, + const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt, + const uint32_t baseD) +{ + ReduceMulRegList regList[2]; + MicroAPI::MaskReg mask = MicroAPI::CreateMask(); + uint32_t offset = 0; + uint32_t rCnt = coff * cmpRatio; + for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) { + MicroAPI::Duplicate(regList[0].vregSum, 0, mask); + MicroAPI::Duplicate(regList[1].vregSum, 0, mask); + for (uint32_t rLoop = 0; rLoop < rCnt; rLoop++) { + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[0], offset, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[1], offset + baseD64, baseD64); + offset += baseD; + } + MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList[0].vregSum, mask); + MicroAPI::StoreAlign(outputAddr + scLoop * baseD + baseD64, regList[1].vregSum, mask); + } +} + +template +__simd_vf__ void MulReduceSumbase256VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr, + const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt, + const uint32_t baseD) +{ + ReduceMulRegList regList[4]; + MicroAPI::MaskReg mask = MicroAPI::CreateMask(); + uint32_t offset = 0; + uint32_t rCnt = coff * cmpRatio; + for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) { + MicroAPI::Duplicate(regList[0].vregSum, 0, mask); + MicroAPI::Duplicate(regList[1].vregSum, 0, mask); + MicroAPI::Duplicate(regList[2].vregSum, 0, mask); + MicroAPI::Duplicate(regList[3].vregSum, 0, mask); + for (uint32_t rLoop = 0; rLoop < rCnt; rLoop++) { + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[0], offset, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[1], offset + baseD64, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[2], offset + 2 * baseD64, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[3], offset + 3 * baseD64, baseD64); + offset += baseD; + } + MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList[0].vregSum, mask); + MicroAPI::StoreAlign(outputAddr + scLoop * baseD + baseD64, regList[1].vregSum, mask); + MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 2 * baseD64, regList[2].vregSum, mask); + MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 3 * baseD64, regList[3].vregSum, mask); + } +} + +template +__simd_vf__ void MulReduceSumbase512VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr, + const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt, + const uint32_t baseD) +{ + ReduceMulRegList regList[8]; + MicroAPI::MaskReg mask = MicroAPI::CreateMask(); + uint32_t offset = 0; + uint32_t rCnt = coff * cmpRatio; + for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) { + MicroAPI::Duplicate(regList[0].vregSum, 0, mask); + MicroAPI::Duplicate(regList[1].vregSum, 0, mask); + MicroAPI::Duplicate(regList[2].vregSum, 0, mask); + MicroAPI::Duplicate(regList[3].vregSum, 0, mask); + MicroAPI::Duplicate(regList[4].vregSum, 0, mask); + MicroAPI::Duplicate(regList[5].vregSum, 0, mask); + MicroAPI::Duplicate(regList[6].vregSum, 0, mask); + MicroAPI::Duplicate(regList[7].vregSum, 0, mask); + for (uint32_t rLoop = 0; rLoop < rCnt; rLoop++) { + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[0], offset, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[1], offset + baseD64, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[2], offset + 2 * baseD64, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[3], offset + 3 * baseD64, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[4], offset + 4 * baseD64, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[5], offset + 5 * baseD64, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[6], offset + 6 * baseD64, baseD64); + LoadMulAddVFImpl(kvAddr, scoreAddr, regList[7], offset + 7 * baseD64, baseD64); + offset += baseD; + } + MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList[0].vregSum, mask); + MicroAPI::StoreAlign(outputAddr + scLoop * baseD + baseD64, regList[1].vregSum, mask); + MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 2 * baseD64, regList[2].vregSum, mask); + MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 3 * baseD64, regList[3].vregSum, mask); + MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 4 * baseD64, regList[4].vregSum, mask); + MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 5 * baseD64, regList[5].vregSum, mask); + MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 6 * baseD64, regList[6].vregSum, mask); + MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 7 * baseD64, regList[7].vregSum, mask); + } +} + +/** + * @brief MulReduceSumbaseVF 包含mul和reducesum + * @param outputLocal 输出tensor [] + * @param coff + * @param cmpRatio 压缩块大小 + * @param baseD 核内d轴切分大小 + * @param scLoopCnt sc数, + */ + +// 当前仅支持coff * cmpRatio为2的幂的情况 +template +__aicore__ inline void MulReduceSumbaseVF(const LocalTensor &kvLocal, const LocalTensor &scoreLocal, + const LocalTensor &outputLocal, const uint32_t coff, const uint32_t cmpRatio, + const uint32_t baseD, const uint32_t scLoopCnt) +{ + + __ubuf__ T *kvAddr = (__ubuf__ T *)kvLocal.GetPhyAddr(); + __ubuf__ T *scoreAddr = (__ubuf__ T *)scoreLocal.GetPhyAddr(); + __ubuf__ T *outputAddr = (__ubuf__ T *)outputLocal.GetPhyAddr(); + if (baseD == baseD8) { + MulReduceSumbase8VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD); + } else if (baseD == baseD16) { + MulReduceSumbase16VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD); + } else if (baseD == baseD32) { + MulReduceSumbase32VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD); + } else if (baseD == baseD64) { + MulReduceSumbase64VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD); + } else if (baseD == baseD128) { + MulReduceSumbase128VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD); + } else if (baseD == baseD256) { + MulReduceSumbase256VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD); + } else if (baseD == baseD512) { + MulReduceSumbase512VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD); + } +} + + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_rms_norm.h b/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_rms_norm.h new file mode 100644 index 000000000..2ddc1de3c --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_rms_norm.h @@ -0,0 +1,95 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_rms_norm.h + * \brief + */ + +#ifndef VF_RMS_NORM_H +#define VF_RMS_NORM_H +#include "kernel_tensor.h" + +//repeatTimes——D轴的分块数 +template +__simd_vf__ void RmsNormVFImpl(__ubuf__ T * inputBuf, __ubuf__ GammaType * gammaBuf, __ubuf__ T * outputBuf, + uint32_t repeatTimes, float reciprocal, float epsilon) +{ + MicroAPI::RegTensor vregSum; + MicroAPI::RegTensor vregSumReduce; + MicroAPI::RegTensor vregDiv; + MicroAPI::RegTensor vregSquareRoot; + + MicroAPI::MaskReg maskAll = MicroAPI::CreateMask(); + MicroAPI::MaskReg maskFirst = MicroAPI::CreateMask(); + + static constexpr MicroAPI::CastTrait castTraitB162B32 = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + MicroAPI::Duplicate(vregSum, 0.0f); + + for(uint32_t i = 0; i < repeatTimes; ++i){ + MicroAPI::RegTensor vregX; + MicroAPI::RegTensor vregXSquare; + uint64_t loopOffset = i * FLOAT_REP_SIZE; + + MicroAPI::LoadAlign(vregX, inputBuf + loopOffset); + MicroAPI::Mul(vregXSquare, vregX, vregX, maskAll); + MicroAPI::Add(vregSum, vregXSquare, vregSum, maskAll); + } + + MicroAPI::Reduce(vregSumReduce, vregSum, maskAll); + MicroAPI::Muls(vregSumReduce, vregSumReduce, reciprocal, maskFirst); + MicroAPI::Adds(vregSumReduce, vregSumReduce, epsilon, maskFirst); + MicroAPI::Sqrt(vregSquareRoot, vregSumReduce, maskFirst); + MicroAPI::Duplicate(vregDiv, vregSquareRoot, maskAll); + + for(uint32_t i = 0; i < repeatTimes; ++i){ + MicroAPI::RegTensor vregX; + MicroAPI::RegTensor vregGammaCast; + uint16_t loopOffset = i * FLOAT_REP_SIZE; + + MicroAPI::LoadAlign(vregX, inputBuf + loopOffset); + MicroAPI::LoadAlign(vregGammaCast, gammaBuf + loopOffset); + + MicroAPI::Div(vregX, vregX, vregDiv, maskAll); + MicroAPI::Mul(vregX, vregX, vregGammaCast, maskAll); + + MicroAPI::StoreAlign(outputBuf + loopOffset, vregX, maskAll); + } +} + +/** + * @brief RmsNormVF 对一行进行rmsnorm + * @param outputLocal 输出tensor [row, col]，row目前均为1 + * @param inputLocal 输入tensor [row, col] + * @param gammaLocal gamma参数tensor [row, col] + * @param rmsNormParams rmsNrom计算所需系数，包括 + row 行数 1 + col 列数，对应headSizeCq或headSizeCkv + reciprocal ，1/N + epsilon，防止除零极小数 + */ +template +__aicore__ inline void RmsNormVF(const LocalTensor outputLocal, const LocalTensor inputLocal, const LocalTensor gammaLocal, + float reciprocal, float epsilon, uint32_t row, uint32_t col) +{ + uint32_t cnt = row * col; + uint32_t repeatTimes = (cnt + FLOAT_REP_SIZE - 1) / FLOAT_REP_SIZE; + + __ubuf__ T * inputBuf = (__ubuf__ T *)inputLocal.GetPhyAddr(); + __ubuf__ GammaType * gammaBuf = (__ubuf__ GammaType *)gammaLocal.GetPhyAddr(); + __ubuf__ T * outputBuf = (__ubuf__ T *)outputLocal.GetPhyAddr(); + + RmsNormVFImpl(inputBuf, gammaBuf, outputBuf, repeatTimes, reciprocal, epsilon); +} + + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_rope.h b/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_rope.h new file mode 100644 index 000000000..afe421b9e --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_rope.h @@ -0,0 +1,158 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_rope.h + * \brief + */ + +#ifndef VF_ROPE_H +#define VF_ROPE_H + +#include "kernel_operator.h" +#include "../compressor_comm.h" + +using namespace AscendC; + +constexpr MicroAPI::CastTrait castTraitB162B32 = { + MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::UNKNOWN, +}; + +constexpr MicroAPI::CastTrait castTraitB322B16 = { + MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_RINT, +}; + + +template +__simd_vf__ void HalfModeRopeVF(__ubuf__ T *sinUb, __ubuf__ T *cosUb, __ubuf__ T *inUb, __ubuf__ ROPET *outUb, + uint32_t row, uint32_t col, uint32_t actualCol, uint64_t baseAddr) +{ + MicroAPI::RegTensor vregCos; + MicroAPI::RegTensor vregHalfCos; + MicroAPI::RegTensor vregSin; + MicroAPI::RegTensor vregHalfSin; + MicroAPI::RegTensor vregIn; + MicroAPI::RegTensor vregHalfIn; + MicroAPI::RegTensor vregOut; + MicroAPI::RegTensor vregHalfOut; + MicroAPI::RegTensor vregCastIn; + MicroAPI::RegTensor vregOutBf16; + MicroAPI::RegTensor vregOutHalfBf16; + MicroAPI::RegTensor vregCastOut; + uint32_t maskValue = col / 2; + MicroAPI::MaskReg mask = MicroAPI::UpdateMask(maskValue); + uint32_t halfCol = col / 2; + + + for (uint32_t rIdx = 0; rIdx < row; rIdx++) { + __ubuf__ T *curSinUb = sinUb + rIdx * col; + __ubuf__ T *curCosUb = cosUb + rIdx * col; + __ubuf__ T *curInUb = inUb + rIdx * actualCol; + __ubuf__ ROPET *curOutUb = outUb + rIdx * actualCol; + + MicroAPI::DataCopy(vregIn, curInUb + baseAddr); + MicroAPI::DataCopy(vregHalfIn, curInUb + baseAddr + halfCol); + MicroAPI::DataCopy(vregCos, curCosUb); + MicroAPI::DataCopy(vregHalfCos, curCosUb + halfCol); + MicroAPI::DataCopy(vregSin, curSinUb); + MicroAPI::DataCopy(vregHalfSin, curSinUb + halfCol); + MicroAPI::Mul(vregSin, vregSin, vregHalfIn, mask); + MicroAPI::Mul(vregHalfSin, vregHalfSin, vregIn, mask); + MicroAPI::Mul(vregCos, vregCos, vregIn, mask); + MicroAPI::Sub(vregOut, vregCos, vregSin, mask); + MicroAPI::Mul(vregHalfCos, vregHalfCos, vregHalfIn, mask); + MicroAPI::Add(vregHalfOut, vregHalfSin, vregHalfCos, mask); + MicroAPI::Cast(vregOutBf16, vregOut, mask); + MicroAPI::DataCopy(curOutUb + baseAddr, vregOutBf16, mask); + MicroAPI::Cast(vregOutHalfBf16, vregHalfOut, mask); + MicroAPI::DataCopy(curOutUb + baseAddr + halfCol, vregOutHalfBf16, + mask); + + for (uint64_t dOffset = 0; dOffset < baseAddr; dOffset += 64) { + uint32_t castMaskValue = min(baseAddr - dOffset, static_cast(64)); + MicroAPI::MaskReg castMask = MicroAPI::UpdateMask(castMaskValue); + MicroAPI::DataCopy(vregCastIn, curInUb + dOffset); + MicroAPI::Cast(vregCastOut, vregCastIn, castMask); + MicroAPI::DataCopy(curOutUb + dOffset, vregCastOut, castMask); + } + } +} + + +template +__simd_vf__ void InterleaveModeRopeVF(__ubuf__ T *sinUb, __ubuf__ T *cosUb, __ubuf__ T *inUb, __ubuf__ ROPET *outUb, + uint32_t row, uint32_t col, uint32_t actualCol, uint64_t baseAddr) +{ + MicroAPI::RegTensor vregCos; + MicroAPI::RegTensor vregSin; + MicroAPI::RegTensor vregIn; + MicroAPI::RegTensor vregOdd; + MicroAPI::RegTensor vregEven; + MicroAPI::RegTensor vregOut; + MicroAPI::RegTensor vregTemp; + MicroAPI::RegTensor vregCastIn; + MicroAPI::RegTensor vregOutBf16; + MicroAPI::RegTensor vregCastOut; + uint32_t maskValue = col; + MicroAPI::MaskReg mask = MicroAPI::UpdateMask(maskValue); + + + for (uint32_t rIdx = 0; rIdx < row; rIdx++) { + __ubuf__ T *curSinUb = sinUb + rIdx * col; + __ubuf__ T *curCosUb = cosUb + rIdx * col; + __ubuf__ T *curInUb = inUb + rIdx * actualCol; + __ubuf__ ROPET *curOutUb = outUb + rIdx * actualCol; + + MicroAPI::DataCopy(vregIn, curInUb + baseAddr); + MicroAPI::DataCopy(vregCos, curCosUb); + MicroAPI::DataCopy(vregSin, curSinUb); + MicroAPI::Mul(vregCos, vregCos, vregIn, mask); + MicroAPI::DeInterleave(vregEven, vregOdd, vregIn, vregTemp); + MicroAPI::Muls(vregOdd, vregOdd, static_cast(-1.0), mask); + MicroAPI::Interleave(vregIn, vregTemp, vregOdd, vregEven); + MicroAPI::Mul(vregSin, vregSin, vregIn, mask); + MicroAPI::Add(vregOut, vregCos, vregSin, mask); + MicroAPI::Cast(vregOutBf16, vregOut, mask); + MicroAPI::DataCopy(curOutUb + baseAddr, vregOutBf16, mask); + for (uint64_t dOffset = 0; dOffset < baseAddr; dOffset += 64) { + uint32_t castMaskValue = min(baseAddr - dOffset, static_cast(64)); + MicroAPI::MaskReg castMask = MicroAPI::UpdateMask(castMaskValue); + MicroAPI::DataCopy(vregCastIn, curInUb + dOffset); + MicroAPI::Cast(vregCastOut, vregCastIn, castMask); + MicroAPI::DataCopy(curOutUb + dOffset, vregCastOut, castMask); + } + } +} + + +template +__aicore__ inline void RopeVF(const LocalTensor &sinTensor, const LocalTensor &cosTensor, + const LocalTensor &inTensor, const LocalTensor &outTensor, uint32_t row, + uint32_t col, uint32_t actualCol, uint64_t baseAddr) +{ + __ubuf__ T *sinUb = (__ubuf__ T *)sinTensor.GetPhyAddr(); + __ubuf__ T *cosUb = (__ubuf__ T *)cosTensor.GetPhyAddr(); + __ubuf__ T *inUb = (__ubuf__ T *)inTensor.GetPhyAddr(); + __ubuf__ ROPET *outUb = (__ubuf__ ROPET *)outTensor.GetPhyAddr(); + + if constexpr (MODE == Compressor::ROTARY_MODE::HALF) { + HalfModeRopeVF(sinUb, cosUb, inUb, outUb, row, col, actualCol, baseAddr); + } else { + InterleaveModeRopeVF(sinUb, cosUb, inUb, outUb, row, col, actualCol, baseAddr); + } +} + +#endif diff --git a/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_softmax.h b/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_softmax.h new file mode 100644 index 000000000..b0c809116 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/arch35/vf/vf_softmax.h @@ -0,0 +1,1592 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vf_softmax.h + * \brief + */ + +#ifndef VF_SOFTMAX_H +#define VF_SOFTMAX_H +#include "kernel_tensor.h" +namespace FaVectorApi { +using AscendC::LocalTensor; +using namespace AscendC; +using namespace MicroAPI; + +template +__simd_vf__ inline void SoftmaxDndBase128(__ubuf__ T *inputAddr, __ubuf__ float *outputAddr, + const uint32_t RowSize, const uint32_t ReduceSize, const uint32_t vScRealSize, + const T minValue) +{ + RegTensor vregSum00; + RegTensor vregSum10; + RegTensor vregSum20; + RegTensor vregSum30; + RegTensor vregSum01; + RegTensor vregSum11; + RegTensor vregSum21; + RegTensor vregSum31; + RegTensor vregSum02; + RegTensor vregSum12; + RegTensor vregSum22; + RegTensor vregSum32; + RegTensor vregSum03; + RegTensor vregSum13; + RegTensor vregSum23; + RegTensor vregSum33; + + RegTensor vregExp00; + RegTensor vregExp10; + RegTensor vregExp20; + RegTensor vregExp30; + RegTensor vregExp01; + RegTensor vregExp11; + RegTensor vregExp21; + RegTensor vregExp31; + RegTensor vregExp02; + RegTensor vregExp12; + RegTensor vregExp22; + RegTensor vregExp32; + RegTensor vregExp03; + RegTensor vregExp13; + RegTensor vregExp23; + RegTensor vregExp33; + + RegTensor vregF32_00; + RegTensor vregF32_10; + RegTensor vregF32_20; + RegTensor vregF32_30; + RegTensor vregF32_01; + RegTensor vregF32_11; + RegTensor vregF32_21; + RegTensor vregF32_31; + RegTensor vregF32_02; + RegTensor vregF32_12; + RegTensor vregF32_22; + RegTensor vregF32_32; + RegTensor vregF32_03; + RegTensor vregF32_13; + RegTensor vregF32_23; + RegTensor vregF32_33; + + RegTensor vregStore00; + RegTensor vregStore10; + RegTensor vregStore20; + RegTensor vregStore30; + RegTensor vregStore01; + RegTensor vregStore11; + RegTensor vregStore21; + RegTensor vregStore31; + MaskReg pregAll; + pregAll = CreateMask(); + RegTensor src00, src10, src20, src30, src01, src11, src21, src31, + src02, src12, src22, src32, src03, src13, src23, src33; + RegTensor max00, max10, max20, max30, max01, max11, max21, max31, + max02, max12, max22, max32, max03, max13, max23, max33; + + __ubuf__ float *srcUb00 = outputAddr; + __ubuf__ float *srcUb01 = outputAddr + RowSize / 2; + __ubuf__ float *srcUb02 = outputAddr + RowSize; + __ubuf__ float *srcUb03 = outputAddr + RowSize + RowSize / 2; + __ubuf__ float *srcUb10 = srcUb00 + ReduceSize * RowSize; + __ubuf__ float *srcUb11 = srcUb00 + ReduceSize * RowSize + RowSize / 2; + __ubuf__ float *srcUb12 = srcUb00 + ReduceSize * RowSize + RowSize; + __ubuf__ float *srcUb13 = srcUb00 + ReduceSize * RowSize + RowSize + RowSize / 2; + __ubuf__ float *srcUb20 = srcUb00 + ReduceSize * RowSize * 2; + __ubuf__ float *srcUb21 = srcUb00 + ReduceSize * RowSize * 2 + RowSize / 2; + __ubuf__ float *srcUb22 = srcUb00 + ReduceSize * RowSize * 2 + RowSize; + __ubuf__ float *srcUb23 = srcUb00 + ReduceSize * RowSize * 2 + RowSize + RowSize / 2; + __ubuf__ float *srcUb30 = srcUb00 + ReduceSize * RowSize * 3; + __ubuf__ float *srcUb31 = srcUb00 + ReduceSize * RowSize * 3 + RowSize / 2; + __ubuf__ float *srcUb32 = srcUb00 + ReduceSize * RowSize * 3 + RowSize; + __ubuf__ float *srcUb33 = srcUb00 + ReduceSize * RowSize * 3 + RowSize + RowSize / 2; + + __ubuf__ float *inputAddr00 = inputAddr; + __ubuf__ float *inputAddr01 = inputAddr + RowSize / 2; + __ubuf__ float *inputAddr10 = inputAddr + (ReduceSize * RowSize); + __ubuf__ float *inputAddr11 = inputAddr + (ReduceSize * RowSize) + RowSize / 2; + __ubuf__ float *inputAddr20 = inputAddr + (ReduceSize * RowSize * 2); + __ubuf__ float *inputAddr21 = inputAddr + (ReduceSize * RowSize * 2) + RowSize / 2; + __ubuf__ float *inputAddr30 = inputAddr + (ReduceSize * RowSize * 3); + __ubuf__ float *inputAddr31 = inputAddr + (ReduceSize * RowSize * 3) + RowSize / 2; + + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize / 4); ++loopSc) { + Duplicate(max00, minValue); + Duplicate(max10, minValue); + Duplicate(max20, minValue); + Duplicate(max30, minValue); + Duplicate(max01, minValue); + Duplicate(max11, minValue); + Duplicate(max21, minValue); + Duplicate(max31, minValue); + Duplicate(max02, minValue); + Duplicate(max12, minValue); + Duplicate(max22, minValue); + Duplicate(max32, minValue); + Duplicate(max03, minValue); + Duplicate(max13, minValue); + Duplicate(max23, minValue); + Duplicate(max33, minValue); + + Duplicate(vregSum00, 0, pregAll); + Duplicate(vregSum10, 0, pregAll); + Duplicate(vregSum20, 0, pregAll); + Duplicate(vregSum30, 0, pregAll); + Duplicate(vregSum01, 0, pregAll); + Duplicate(vregSum11, 0, pregAll); + Duplicate(vregSum21, 0, pregAll); + Duplicate(vregSum31, 0, pregAll); + Duplicate(vregSum02, 0, pregAll); + Duplicate(vregSum12, 0, pregAll); + Duplicate(vregSum22, 0, pregAll); + Duplicate(vregSum32, 0, pregAll); + Duplicate(vregSum03, 0, pregAll); + Duplicate(vregSum13, 0, pregAll); + Duplicate(vregSum23, 0, pregAll); + Duplicate(vregSum33, 0, pregAll); + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 2); ++loopM) { + LoadAlign(src00, srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src01, srcUb01 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src02, srcUb02 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src03, srcUb03 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(src10, srcUb10 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src11, srcUb11 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src12, srcUb12 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src13, srcUb13 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(src20, srcUb20 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src21, srcUb21 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src22, srcUb22 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src23, srcUb23 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(src30, srcUb30 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src31, srcUb31 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src32, srcUb32 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src33, srcUb33 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + Max(max00, max00, src00, pregAll); + Max(max01, max01, src01, pregAll); + Max(max02, max02, src02, pregAll); + Max(max03, max03, src03, pregAll); + Max(max10, max10, src10, pregAll); + Max(max11, max11, src11, pregAll); + Max(max12, max12, src12, pregAll); + Max(max13, max13, src13, pregAll); + Max(max20, max20, src20, pregAll); + Max(max21, max21, src21, pregAll); + Max(max22, max22, src22, pregAll); + Max(max23, max23, src23, pregAll); + Max(max30, max30, src30, pregAll); + Max(max31, max31, src31, pregAll); + Max(max32, max32, src32, pregAll); + Max(max33, max33, src33, pregAll); + } + Max(max00, max00, max02, pregAll); + Max(max01, max01, max03, pregAll); + Max(max10, max10, max12, pregAll); + Max(max11, max11, max13, pregAll); + Max(max20, max20, max22, pregAll); + Max(max21, max21, max23, pregAll); + Max(max30, max30, max32, pregAll); + Max(max31, max31, max33, pregAll); + + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 2); ++loopM) { + LoadAlign(vregF32_00, srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_01, srcUb01 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_02, srcUb02 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_03, srcUb03 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(vregF32_10, srcUb10 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_11, srcUb11 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_12, srcUb12 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_13, srcUb13 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(vregF32_20, srcUb20 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_21, srcUb21 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_22, srcUb22 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_23, srcUb23 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(vregF32_30, srcUb30 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_31, srcUb31 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_32, srcUb32 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_33, srcUb33 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + FusedExpSub(vregExp00, vregF32_00, max00, pregAll); + FusedExpSub(vregExp01, vregF32_01, max01, pregAll); + FusedExpSub(vregExp02, vregF32_02, max00, pregAll); + FusedExpSub(vregExp03, vregF32_03, max01, pregAll); + FusedExpSub(vregExp10, vregF32_10, max10, pregAll); + FusedExpSub(vregExp11, vregF32_11, max11, pregAll); + FusedExpSub(vregExp12, vregF32_12, max10, pregAll); + FusedExpSub(vregExp13, vregF32_13, max11, pregAll); + FusedExpSub(vregExp20, vregF32_20, max20, pregAll); + FusedExpSub(vregExp21, vregF32_21, max21, pregAll); + FusedExpSub(vregExp22, vregF32_22, max20, pregAll); + FusedExpSub(vregExp23, vregF32_23, max21, pregAll); + FusedExpSub(vregExp30, vregF32_30, max30, pregAll); + FusedExpSub(vregExp31, vregF32_31, max31, pregAll); + FusedExpSub(vregExp32, vregF32_32, max30, pregAll); + FusedExpSub(vregExp33, vregF32_33, max31, pregAll); + + Add(vregSum00, vregExp00, vregSum00, pregAll); + Add(vregSum01, vregExp01, vregSum01, pregAll); + Add(vregSum02, vregExp02, vregSum02, pregAll); + Add(vregSum03, vregExp03, vregSum03, pregAll); + Add(vregSum10, vregExp10, vregSum10, pregAll); + Add(vregSum11, vregExp11, vregSum11, pregAll); + Add(vregSum12, vregExp12, vregSum12, pregAll); + Add(vregSum13, vregExp13, vregSum13, pregAll); + Add(vregSum20, vregExp20, vregSum20, pregAll); + Add(vregSum21, vregExp21, vregSum21, pregAll); + Add(vregSum22, vregExp22, vregSum22, pregAll); + Add(vregSum23, vregExp23, vregSum23, pregAll); + Add(vregSum30, vregExp30, vregSum30, pregAll); + Add(vregSum31, vregExp31, vregSum31, pregAll); + Add(vregSum32, vregExp32, vregSum32, pregAll); + Add(vregSum33, vregExp33, vregSum33, pregAll); + + StoreAlign(((__ubuf__ T *&)srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp00, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb01 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp01, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb02 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp02, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb03 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp03, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb10 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp10, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb11 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp11, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb12 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp12, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb13 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp13, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb20 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp20, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb21 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp21, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb22 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp22, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb23 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp23, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb30 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp30, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb31 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp31, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb32 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp32, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb33 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp33, pregAll); + } + Add(vregSum00, vregSum00, vregSum02, pregAll); + Add(vregSum01, vregSum01, vregSum03, pregAll); + Add(vregSum10, vregSum10, vregSum12, pregAll); + Add(vregSum11, vregSum11, vregSum13, pregAll); + Add(vregSum20, vregSum20, vregSum22, pregAll); + Add(vregSum21, vregSum21, vregSum23, pregAll); + Add(vregSum30, vregSum30, vregSum32, pregAll); + Add(vregSum31, vregSum31, vregSum33, pregAll); + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp00, srcUb00 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp01, srcUb01 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp10, srcUb10 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp11, srcUb11 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp20, srcUb20 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp21, srcUb21 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp30, srcUb30 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp31, srcUb31 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + Div(vregStore00, vregExp00, vregSum00, pregAll); + Div(vregStore01, vregExp01, vregSum01, pregAll); + Div(vregStore10, vregExp10, vregSum10, pregAll); + Div(vregStore11, vregExp11, vregSum11, pregAll); + Div(vregStore20, vregExp20, vregSum20, pregAll); + Div(vregStore21, vregExp21, vregSum21, pregAll); + Div(vregStore30, vregExp30, vregSum30, pregAll); + Div(vregStore31, vregExp31, vregSum31, pregAll); + + StoreAlign(((__ubuf__ T *&)inputAddr00 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore00, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr01 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore01, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr10 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore10, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr11 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore11, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr20 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore20, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr21 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore21, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr30 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore30, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr31 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore31, pregAll); + } + } + // 尾块处理 + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize % 4); ++loopSc) { + Duplicate(max00, minValue); + Duplicate(max01, minValue); + Duplicate(max02, minValue); + Duplicate(max03, minValue); + + Duplicate(vregSum00, 0, pregAll); + Duplicate(vregSum01, 0, pregAll); + Duplicate(vregSum02, 0, pregAll); + Duplicate(vregSum03, 0, pregAll); + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 2); ++loopM) { + LoadAlign(src00, srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + LoadAlign(src01, srcUb01 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + LoadAlign(src02, srcUb02 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + LoadAlign(src03, srcUb03 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + + Max(max00, max00, src00, pregAll); + Max(max01, max01, src01, pregAll); + Max(max02, max02, src02, pregAll); + Max(max03, max03, src03, pregAll); + } + Max(max00, max00, max02, pregAll); + Max(max01, max01, max03, pregAll); + + for (uint16_t loopM = 0; loopM < ReduceSize / 2; ++loopM) { + LoadAlign(vregF32_00, srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + LoadAlign(vregF32_01, srcUb01 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + LoadAlign(vregF32_02, srcUb02 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + LoadAlign(vregF32_03, srcUb03 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + + FusedExpSub(vregExp00, vregF32_00, max00, pregAll); + FusedExpSub(vregExp01, vregF32_01, max01, pregAll); + FusedExpSub(vregExp02, vregF32_02, max00, pregAll); + FusedExpSub(vregExp03, vregF32_03, max01, pregAll); + + Add(vregSum00, vregExp00, vregSum00, pregAll); + Add(vregSum01, vregExp01, vregSum01, pregAll); + Add(vregSum02, vregExp02, vregSum02, pregAll); + Add(vregSum03, vregExp03, vregSum03, pregAll); + + StoreAlign(((__ubuf__ T *&)srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregExp00, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb01 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregExp01, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb02 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregExp02, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb03 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregExp03, pregAll); + } + Add(vregSum00, vregSum00, vregSum02, pregAll); + Add(vregSum01, vregSum01, vregSum03, pregAll); + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp00, srcUb00 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + LoadAlign(vregExp01, srcUb01 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + + Div(vregStore00, vregExp00, vregSum00, pregAll); + Div(vregStore01, vregExp01, vregSum01, pregAll); + + StoreAlign(((__ubuf__ T *&)inputAddr00 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregStore00, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr01 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregStore01, pregAll); + } + } +} + +template +__simd_vf__ inline void SoftmaxDndBase64(__ubuf__ T *inputAddr, __ubuf__ float *outputAddr, + const uint32_t RowSize, const uint32_t ReduceSize, const uint32_t vScRealSize, + const T minValue) +{ + RegTensor vregSum00; + RegTensor vregSum10; + RegTensor vregSum20; + RegTensor vregSum30; + RegTensor vregSum01; + RegTensor vregSum11; + RegTensor vregSum21; + RegTensor vregSum31; + + RegTensor vregExp00; + RegTensor vregExp10; + RegTensor vregExp20; + RegTensor vregExp30; + RegTensor vregExp01; + RegTensor vregExp11; + RegTensor vregExp21; + RegTensor vregExp31; + + RegTensor vregF32_00; + RegTensor vregF32_10; + RegTensor vregF32_20; + RegTensor vregF32_30; + RegTensor vregF32_01; + RegTensor vregF32_11; + RegTensor vregF32_21; + RegTensor vregF32_31; + + RegTensor vregStore0; + RegTensor vregStore1; + RegTensor vregStore2; + RegTensor vregStore3; + MaskReg pregAll; + pregAll = CreateMask(); + RegTensor src00, src10, src20, src30, src01, src11, src21, src31; + RegTensor max00, max10, max20, max30, max01, max11, max21, max31; + + __ubuf__ float *srcUb00 = outputAddr; + __ubuf__ float *srcUb01 = outputAddr + RowSize; + __ubuf__ float *srcUb10 = srcUb00 + ReduceSize * RowSize; + __ubuf__ float *srcUb11 = srcUb00 + ReduceSize * RowSize + RowSize; + __ubuf__ float *srcUb20 = srcUb00 + ReduceSize * RowSize * 2; + __ubuf__ float *srcUb21 = srcUb00 + ReduceSize * RowSize * 2 + RowSize; + __ubuf__ float *srcUb30 = srcUb00 + ReduceSize * RowSize * 3; + __ubuf__ float *srcUb31 = srcUb00 + ReduceSize * RowSize * 3 + RowSize; + + __ubuf__ float *inputAddr0 = inputAddr; + __ubuf__ float *inputAddr1 = inputAddr + (ReduceSize * RowSize); + __ubuf__ float *inputAddr2 = inputAddr + (ReduceSize * RowSize * 2); + __ubuf__ float *inputAddr3 = inputAddr + (ReduceSize * RowSize * 3); + + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize / 4); ++loopSc) { + Duplicate(max00, minValue); + Duplicate(max10, minValue); + Duplicate(max20, minValue); + Duplicate(max30, minValue); + Duplicate(max01, minValue); + Duplicate(max11, minValue); + Duplicate(max21, minValue); + Duplicate(max31, minValue); + + Duplicate(vregSum00, 0, pregAll); + Duplicate(vregSum10, 0, pregAll); + Duplicate(vregSum20, 0, pregAll); + Duplicate(vregSum30, 0, pregAll); + Duplicate(vregSum01, 0, pregAll); + Duplicate(vregSum11, 0, pregAll); + Duplicate(vregSum21, 0, pregAll); + Duplicate(vregSum31, 0, pregAll); + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 2); ++loopM) { + LoadAlign(src00, srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src01, srcUb01 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(src10, srcUb10 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src11, srcUb11 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(src20, srcUb20 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src21, srcUb21 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(src30, srcUb30 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src31, srcUb31 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + Max(max00, max00, src00, pregAll); + Max(max01, max01, src01, pregAll); + Max(max10, max10, src10, pregAll); + Max(max11, max11, src11, pregAll); + Max(max20, max20, src20, pregAll); + Max(max21, max21, src21, pregAll); + Max(max30, max30, src30, pregAll); + Max(max31, max31, src31, pregAll); + } + Max(max00, max00, max01, pregAll); + Max(max10, max10, max11, pregAll); + Max(max20, max20, max21, pregAll); + Max(max30, max30, max31, pregAll); + + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 2); ++loopM) { + LoadAlign(vregF32_00, srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_01, srcUb01 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(vregF32_10, srcUb10 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_11, srcUb11 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(vregF32_20, srcUb20 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_21, srcUb21 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(vregF32_30, srcUb30 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_31, srcUb31 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + FusedExpSub(vregExp00, vregF32_00, max00, pregAll); + FusedExpSub(vregExp01, vregF32_01, max00, pregAll); + FusedExpSub(vregExp10, vregF32_10, max10, pregAll); + FusedExpSub(vregExp11, vregF32_11, max10, pregAll); + FusedExpSub(vregExp20, vregF32_20, max20, pregAll); + FusedExpSub(vregExp21, vregF32_21, max20, pregAll); + FusedExpSub(vregExp30, vregF32_30, max30, pregAll); + FusedExpSub(vregExp31, vregF32_31, max30, pregAll); + + Add(vregSum00, vregExp00, vregSum00, pregAll); + Add(vregSum01, vregExp01, vregSum01, pregAll); + Add(vregSum10, vregExp10, vregSum10, pregAll); + Add(vregSum11, vregExp11, vregSum11, pregAll); + Add(vregSum20, vregExp20, vregSum20, pregAll); + Add(vregSum21, vregExp21, vregSum21, pregAll); + Add(vregSum30, vregExp30, vregSum30, pregAll); + Add(vregSum31, vregExp31, vregSum31, pregAll); + + StoreAlign(((__ubuf__ T *&)srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp00, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb01 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp01, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb10 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp10, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb11 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp11, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb20 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp20, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb21 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp21, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb30 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp30, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb31 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp31, pregAll); + } + Add(vregSum00, vregSum00, vregSum01, pregAll); + Add(vregSum10, vregSum10, vregSum11, pregAll); + Add(vregSum20, vregSum20, vregSum21, pregAll); + Add(vregSum30, vregSum30, vregSum31, pregAll); + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp00, srcUb00 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp10, srcUb10 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp20, srcUb20 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp30, srcUb30 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + Div(vregStore0, vregExp00, vregSum00, pregAll); + Div(vregStore1, vregExp10, vregSum10, pregAll); + Div(vregStore2, vregExp20, vregSum20, pregAll); + Div(vregStore3, vregExp30, vregSum30, pregAll); + + StoreAlign(((__ubuf__ T *&)inputAddr0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore0, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore1, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore2, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore3, pregAll); + } + } + // 尾块处理 + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize % 4); ++loopSc) { + Duplicate(max00, minValue); + Duplicate(max01, minValue); + + Duplicate(vregSum00, 0, pregAll); + Duplicate(vregSum01, 0, pregAll); + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 2); ++loopM) { + LoadAlign(src00, srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + LoadAlign(src01, srcUb01 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + + Max(max00, max00, src00, pregAll); + Max(max01, max01, src01, pregAll); + } + Max(max00, max00, max01, pregAll); + + for (uint16_t loopM = 0; loopM < ReduceSize / 2; ++loopM) { + LoadAlign(vregF32_00, srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + LoadAlign(vregF32_01, srcUb01 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + + FusedExpSub(vregExp00, vregF32_00, max00, pregAll); + FusedExpSub(vregExp01, vregF32_01, max00, pregAll); + + Add(vregSum00, vregExp00, vregSum00, pregAll); + Add(vregSum01, vregExp01, vregSum01, pregAll); + + StoreAlign(((__ubuf__ T *&)srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregExp00, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb01 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregExp01, pregAll); + } + Add(vregSum00, vregSum00, vregSum01, pregAll); + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp00, srcUb00 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + Div(vregStore0, vregExp00, vregSum00, pregAll); + + StoreAlign(((__ubuf__ T *&)inputAddr0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregStore0, pregAll); + } + } +} + +template +__simd_vf__ inline void SoftmaxDndBase32(__ubuf__ T *inputAddr, __ubuf__ float *outputAddr, + const uint32_t RowSize, const uint32_t ReduceSize, const uint32_t vScRealSize, + const T minValue) +{ + RegTensor vregSum00; + RegTensor vregSum10; + RegTensor vregSum20; + RegTensor vregSum30; + RegTensor vregSum01; + RegTensor vregSum11; + RegTensor vregSum21; + RegTensor vregSum31; + + RegTensor vregExp00; + RegTensor vregExp10; + RegTensor vregExp20; + RegTensor vregExp30; + RegTensor vregExp01; + RegTensor vregExp11; + RegTensor vregExp21; + RegTensor vregExp31; + + RegTensor vregF32_00; + RegTensor vregF32_10; + RegTensor vregF32_20; + RegTensor vregF32_30; + RegTensor vregF32_01; + RegTensor vregF32_11; + RegTensor vregF32_21; + RegTensor vregF32_31; + + RegTensor vregStore0; + RegTensor vregStore1; + RegTensor vregStore2; + RegTensor vregStore3; + + MaskReg pregLHalf; + MaskReg pregHHalf; + MaskReg pregAll; + pregAll = CreateMask(); + pregLHalf = CreateMask(); + Not(pregHHalf, pregLHalf, pregAll); + RegTensor max0, max1, max2, max3; + RegTensor src00, src10, src20, src30, src01, src11, src21, src31; + RegTensor max00, max10, max20, max30, max01, max11, max21, max31; + + __ubuf__ float *srcUb00 = outputAddr; + __ubuf__ float *srcUb01 = outputAddr + RowSize * 2; + __ubuf__ float *srcUb10 = srcUb00 + ReduceSize * RowSize; + __ubuf__ float *srcUb11 = srcUb00 + ReduceSize * RowSize + RowSize * 2; + __ubuf__ float *srcUb20 = srcUb00 + ReduceSize * RowSize * 2; + __ubuf__ float *srcUb21 = srcUb00 + ReduceSize * RowSize * 2 + RowSize * 2; + __ubuf__ float *srcUb30 = srcUb00 + ReduceSize * RowSize * 3; + __ubuf__ float *srcUb31 = srcUb00 + ReduceSize * RowSize * 3 + RowSize * 2; + + __ubuf__ float *inputAddr0 = inputAddr; + __ubuf__ float *inputAddr1 = inputAddr + (ReduceSize * RowSize); + __ubuf__ float *inputAddr2 = inputAddr + (ReduceSize * RowSize * 2); + __ubuf__ float *inputAddr3 = inputAddr + (ReduceSize * RowSize * 3); + + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize / 4); ++loopSc) { + Duplicate(max0, minValue); + Duplicate(max1, minValue); + Duplicate(max2, minValue); + Duplicate(max3, minValue); + Duplicate(max00, minValue); + Duplicate(max10, minValue); + Duplicate(max20, minValue); + Duplicate(max30, minValue); + Duplicate(max01, minValue); + Duplicate(max11, minValue); + Duplicate(max21, minValue); + Duplicate(max31, minValue); + + Duplicate(vregSum00, 0, pregAll); + Duplicate(vregSum10, 0, pregAll); + Duplicate(vregSum20, 0, pregAll); + Duplicate(vregSum30, 0, pregAll); + Duplicate(vregSum01, 0, pregAll); + Duplicate(vregSum11, 0, pregAll); + Duplicate(vregSum21, 0, pregAll); + Duplicate(vregSum31, 0, pregAll); + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 4); ++loopM) { + LoadAlign(src00, srcUb00 + loopM * RowSize * 4 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src01, srcUb01 + loopM * RowSize * 4 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(src10, srcUb10 + loopM * RowSize * 4 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src11, srcUb11 + loopM * RowSize * 4 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(src20, srcUb20 + loopM * RowSize * 4 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src21, srcUb21 + loopM * RowSize * 4 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(src30, srcUb30 + loopM * RowSize * 4 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src31, srcUb31 + loopM * RowSize * 4 + ReduceSize * RowSize * loopSc * 4); + + Max(max00, max00, src00, pregAll); + Max(max01, max01, src01, pregAll); + Max(max10, max10, src10, pregAll); + Max(max11, max11, src11, pregAll); + Max(max20, max20, src20, pregAll); + Max(max21, max21, src21, pregAll); + Max(max30, max30, src30, pregAll); + Max(max31, max31, src31, pregAll); + } + Max(max0, max00, max01, pregAll); + Max(max1, max10, max11, pregAll); + Max(max2, max20, max21, pregAll); + Max(max3, max30, max31, pregAll); + + Squeeze(max00, max0, pregLHalf); + Squeeze(max01, max0, pregHHalf); + Max(max0, max00, max01, pregLHalf); + + Squeeze(max10, max1, pregLHalf); + Squeeze(max11, max1, pregHHalf); + Max(max1, max10, max11, pregLHalf); + + Squeeze(max20, max2, pregLHalf); + Squeeze(max21, max2, pregHHalf); + Max(max2, max20, max21, pregLHalf); + + Squeeze(max30, max3, pregLHalf); + Squeeze(max31, max3, pregHHalf); + Max(max3, max30, max31, pregLHalf); + + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 2); ++loopM) { + LoadAlign(vregF32_00, srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_01, (srcUb00 + RowSize) + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(vregF32_10, srcUb10 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_11, (srcUb10 + RowSize) + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(vregF32_20, srcUb20 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_21, (srcUb20 + RowSize) + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(vregF32_30, srcUb30 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_31, (srcUb30 + RowSize) + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4); + + FusedExpSub(vregExp00, vregF32_00, max0, pregLHalf); + FusedExpSub(vregExp01, vregF32_01, max0, pregLHalf); + FusedExpSub(vregExp10, vregF32_10, max1, pregLHalf); + FusedExpSub(vregExp11, vregF32_11, max1, pregLHalf); + FusedExpSub(vregExp20, vregF32_20, max2, pregLHalf); + FusedExpSub(vregExp21, vregF32_21, max2, pregLHalf); + FusedExpSub(vregExp30, vregF32_30, max3, pregLHalf); + FusedExpSub(vregExp31, vregF32_31, max3, pregLHalf); + + Add(vregSum00, vregExp00, vregSum00, pregLHalf); + Add(vregSum01, vregExp01, vregSum01, pregLHalf); + Add(vregSum10, vregExp10, vregSum10, pregLHalf); + Add(vregSum11, vregExp11, vregSum11, pregLHalf); + Add(vregSum20, vregExp20, vregSum20, pregLHalf); + Add(vregSum21, vregExp21, vregSum21, pregLHalf); + Add(vregSum30, vregExp30, vregSum30, pregLHalf); + Add(vregSum31, vregExp31, vregSum31, pregLHalf); + + StoreAlign(((__ubuf__ T *&)srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp00, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb00 + RowSize + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp01, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb10 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp10, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb10 + RowSize + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp11, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb20 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp20, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb20 + RowSize + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp21, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb30 + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp30, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb30 + RowSize + loopM * RowSize * 2 + ReduceSize * RowSize * loopSc * 4), + vregExp31, pregLHalf); + } + Add(vregSum00, vregSum00, vregSum01, pregLHalf); + Add(vregSum10, vregSum10, vregSum11, pregLHalf); + Add(vregSum20, vregSum20, vregSum21, pregLHalf); + Add(vregSum30, vregSum30, vregSum31, pregLHalf); + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp00, srcUb00 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp10, srcUb10 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp20, srcUb20 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp30, srcUb30 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + Div(vregStore0, vregExp00, vregSum00, pregLHalf); + Div(vregStore1, vregExp10, vregSum10, pregLHalf); + Div(vregStore2, vregExp20, vregSum20, pregLHalf); + Div(vregStore3, vregExp30, vregSum30, pregLHalf); + + StoreAlign(((__ubuf__ T *&)inputAddr0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore0, pregLHalf); + StoreAlign(((__ubuf__ T *&)inputAddr1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore1, pregLHalf); + StoreAlign(((__ubuf__ T *&)inputAddr2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore2, pregLHalf); + StoreAlign(((__ubuf__ T *&)inputAddr3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore3, pregLHalf); + } + } + // 尾块处理 + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize % 4); ++loopSc) { + Duplicate(max0, minValue); + Duplicate(max00, minValue); + Duplicate(max01, minValue); + + Duplicate(vregSum00, 0, pregAll); + Duplicate(vregSum01, 0, pregAll); + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 4); ++loopM) { + LoadAlign(src00, srcUb00 + loopM * RowSize * 4 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + LoadAlign(src01, srcUb01 + loopM * RowSize * 4 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + + Max(max00, max00, src00, pregAll); + Max(max01, max01, src01, pregAll); + } + Max(max0, max00, max01, pregAll); + + Squeeze(max00, max0, pregLHalf); + Squeeze(max01, max0, pregHHalf); + Max(max0, max00, max01, pregLHalf); + + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 2); ++loopM) { + LoadAlign(vregF32_00, srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + LoadAlign(vregF32_01, (srcUb00 + RowSize) + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + + FusedExpSub(vregExp00, vregF32_00, max0, pregLHalf); + FusedExpSub(vregExp01, vregF32_01, max0, pregLHalf); + + Add(vregSum00, vregExp00, vregSum00, pregLHalf); + Add(vregSum01, vregExp01, vregSum01, pregLHalf); + + StoreAlign(((__ubuf__ T *&)srcUb00 + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregExp00, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb00 + RowSize + loopM * RowSize * 2 + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregExp01, pregLHalf); + } + Add(vregSum00, vregSum00, vregSum01, pregLHalf); + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp00, srcUb00 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + + Div(vregStore0, vregExp00, vregSum00, pregLHalf); + + StoreAlign(((__ubuf__ T *&)inputAddr0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregStore0, pregLHalf); + } + } +} + +template +__simd_vf__ inline void SoftmaxDndBase8(__ubuf__ T *inputAddr, __ubuf__ float *outputAddr, + const uint32_t RowSize, const uint32_t ReduceSize, const uint32_t vScRealSize, + const T minValue) +{ + RegTensor vregSum0; + RegTensor vregSum1; + RegTensor vregSum2; + RegTensor vregSum3; + + RegTensor vregExp0; + RegTensor vregExp1; + RegTensor vregExp2; + RegTensor vregExp3; + + RegTensor vregF32_0; + RegTensor vregF32_1; + RegTensor vregF32_2; + RegTensor vregF32_3; + + RegTensor vregStore0; + RegTensor vregStore1; + RegTensor vregStore2; + RegTensor vregStore3; + + MaskReg pregL8; + pregL8 = CreateMask(); + RegTensor src0, src1, src2, src3; + RegTensor max0, max1, max2, max3; + + __ubuf__ float *srcUb0 = outputAddr; + __ubuf__ float *srcUb1 = srcUb0 + ReduceSize * RowSize; + __ubuf__ float *srcUb2 = srcUb0 + ReduceSize * RowSize * 2; + __ubuf__ float *srcUb3 = srcUb0 + ReduceSize * RowSize * 3; + + __ubuf__ float *inputAddr0 = inputAddr; + __ubuf__ float *inputAddr1 = inputAddr + (ReduceSize * RowSize); + __ubuf__ float *inputAddr2 = inputAddr + (ReduceSize * RowSize * 2); + __ubuf__ float *inputAddr3 = inputAddr + (ReduceSize * RowSize * 3); + + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize / 4); ++loopSc) { + Duplicate(max0, minValue); + Duplicate(max1, minValue); + Duplicate(max2, minValue); + Duplicate(max3, minValue); + + Duplicate(vregSum0, 0, pregL8); + Duplicate(vregSum1, 0, pregL8); + Duplicate(vregSum2, 0, pregL8); + Duplicate(vregSum3, 0, pregL8); + + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(src0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src1, srcUb1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src2, srcUb2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src3, srcUb3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + Max(max0, max0, src0, pregL8); + Max(max1, max1, src1, pregL8); + Max(max2, max2, src2, pregL8); + Max(max3, max3, src3, pregL8); + } + + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregF32_0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_1, srcUb1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_2, srcUb2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_3, srcUb3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + FusedExpSub(vregExp0, vregF32_0, max0, pregL8); + FusedExpSub(vregExp1, vregF32_1, max1, pregL8); + FusedExpSub(vregExp2, vregF32_2, max2, pregL8); + FusedExpSub(vregExp3, vregF32_3, max3, pregL8); + + Add(vregSum0, vregExp0, vregSum0, pregL8); + Add(vregSum1, vregExp1, vregSum1, pregL8); + Add(vregSum2, vregExp2, vregSum2, pregL8); + Add(vregSum3, vregExp3, vregSum3, pregL8); + + StoreAlign(((__ubuf__ T *&)srcUb0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp0, pregL8); + StoreAlign(((__ubuf__ T *&)srcUb1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp1, pregL8); + StoreAlign(((__ubuf__ T *&)srcUb2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp2, pregL8); + StoreAlign(((__ubuf__ T *&)srcUb3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp3, pregL8); + } + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp1, srcUb1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp2, srcUb2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp3, srcUb3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + Div(vregStore0, vregExp0, vregSum0, pregL8); + Div(vregStore1, vregExp1, vregSum1, pregL8); + Div(vregStore2, vregExp2, vregSum2, pregL8); + Div(vregStore3, vregExp3, vregSum3, pregL8); + + StoreAlign(((__ubuf__ T *&)inputAddr0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore0, pregL8); + StoreAlign(((__ubuf__ T *&)inputAddr1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore1, pregL8); + StoreAlign(((__ubuf__ T *&)inputAddr2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore2, pregL8); + StoreAlign(((__ubuf__ T *&)inputAddr3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore3, pregL8); + } + } + + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize % 4); ++loopSc) { + Duplicate(max0, minValue); + + Duplicate(vregSum0, 0, pregL8); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(src0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + Max(max0, max0, src0, pregL8); + } + + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregF32_0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + FusedExpSub(vregExp0, vregF32_0, max0, pregL8); + Add(vregSum0, vregExp0, vregSum0, pregL8); + + StoreAlign(((__ubuf__ T *&)srcUb0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregExp0, pregL8); + } + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + Div(vregStore0, vregExp0, vregSum0, pregL8); + + StoreAlign(((__ubuf__ T *&)inputAddr0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregStore0, pregL8); + } + } +} + +template +__simd_vf__ inline void SoftmaxDndBase16(__ubuf__ T *inputAddr, __ubuf__ float *outputAddr, + const uint32_t RowSize, const uint32_t ReduceSize, const uint32_t vScRealSize, + const T minValue) +{ + RegTensor vregSum00; + RegTensor vregSum10; + RegTensor vregSum20; + RegTensor vregSum30; + RegTensor vregSum01; + RegTensor vregSum11; + RegTensor vregSum21; + RegTensor vregSum31; + + RegTensor vregExp00; + RegTensor vregExp10; + RegTensor vregExp20; + RegTensor vregExp30; + RegTensor vregExp01; + RegTensor vregExp11; + RegTensor vregExp21; + RegTensor vregExp31; + + RegTensor vregF32_00; + RegTensor vregF32_10; + RegTensor vregF32_20; + RegTensor vregF32_30; + RegTensor vregF32_01; + RegTensor vregF32_11; + RegTensor vregF32_21; + RegTensor vregF32_31; + + RegTensor vregStore0; + RegTensor vregStore1; + RegTensor vregStore2; + RegTensor vregStore3; + + MaskReg pregLHalf; + MaskReg pregHHalf; + MaskReg pregAll; + pregAll = CreateMask(); + pregLHalf = CreateMask(); + Not(pregHHalf, pregLHalf, pregAll); + RegTensor max0, max1, max2, max3; + RegTensor src00, src10, src20, src30, src01, src11, src21, src31; + RegTensor max00, max10, max20, max30, max01, max11, max21, max31; + + __ubuf__ float *srcUb00 = outputAddr; + __ubuf__ float *srcUb01 = outputAddr + RowSize * (ReduceSize / 2); + __ubuf__ float *srcUb10 = srcUb00 + ReduceSize * RowSize; + __ubuf__ float *srcUb11 = srcUb00 + ReduceSize * RowSize + RowSize * (ReduceSize / 2); + __ubuf__ float *srcUb20 = srcUb00 + ReduceSize * RowSize * 2; + __ubuf__ float *srcUb21 = srcUb00 + ReduceSize * RowSize * 2 + RowSize * (ReduceSize / 2); + __ubuf__ float *srcUb30 = srcUb00 + ReduceSize * RowSize * 3; + __ubuf__ float *srcUb31 = srcUb00 + ReduceSize * RowSize * 3 + RowSize * (ReduceSize / 2); + + __ubuf__ float *inputAddr0 = inputAddr; + __ubuf__ float *inputAddr1 = inputAddr + (ReduceSize * RowSize); + __ubuf__ float *inputAddr2 = inputAddr + (ReduceSize * RowSize * 2); + __ubuf__ float *inputAddr3 = inputAddr + (ReduceSize * RowSize * 3); + + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize / 4); ++loopSc) { + Duplicate(max0, minValue); + Duplicate(max1, minValue); + Duplicate(max2, minValue); + Duplicate(max3, minValue); + Duplicate(max00, minValue); + Duplicate(max10, minValue); + Duplicate(max20, minValue); + Duplicate(max30, minValue); + Duplicate(max01, minValue); + Duplicate(max11, minValue); + Duplicate(max21, minValue); + Duplicate(max31, minValue); + + Duplicate(vregSum00, 0, pregAll); + Duplicate(vregSum10, 0, pregAll); + Duplicate(vregSum20, 0, pregAll); + Duplicate(vregSum30, 0, pregAll); + Duplicate(vregSum01, 0, pregAll); + Duplicate(vregSum11, 0, pregAll); + Duplicate(vregSum21, 0, pregAll); + Duplicate(vregSum31, 0, pregAll); + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 2); ++loopM) { + LoadAlign(src00, srcUb00 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src01, srcUb01 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(src10, srcUb10 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src11, srcUb11 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(src20, srcUb20 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src21, srcUb21 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(src30, srcUb30 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src31, srcUb31 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + Max(max00, max00, src00, pregLHalf); + Max(max01, max01, src01, pregLHalf); + Max(max10, max10, src10, pregLHalf); + Max(max11, max11, src11, pregLHalf); + Max(max20, max20, src20, pregLHalf); + Max(max21, max21, src21, pregLHalf); + Max(max30, max30, src30, pregLHalf); + Max(max31, max31, src31, pregLHalf); + } + Max(max0, max00, max01, pregAll); + Max(max1, max10, max11, pregAll); + Max(max2, max20, max21, pregAll); + Max(max3, max30, max31, pregAll); + + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 2); ++loopM) { + LoadAlign(vregF32_00, srcUb00 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_01, srcUb01 + + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(vregF32_10, srcUb10 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_11, srcUb11 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(vregF32_20, srcUb20 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_21, srcUb21 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + LoadAlign(vregF32_30, srcUb30 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_31, srcUb31 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + FusedExpSub(vregExp00, vregF32_00, max0, pregLHalf); + FusedExpSub(vregExp01, vregF32_01, max0, pregLHalf); + FusedExpSub(vregExp10, vregF32_10, max1, pregLHalf); + FusedExpSub(vregExp11, vregF32_11, max1, pregLHalf); + FusedExpSub(vregExp20, vregF32_20, max2, pregLHalf); + FusedExpSub(vregExp21, vregF32_21, max2, pregLHalf); + FusedExpSub(vregExp30, vregF32_30, max3, pregLHalf); + FusedExpSub(vregExp31, vregF32_31, max3, pregLHalf); + + Add(vregSum00, vregExp00, vregSum00, pregLHalf); + Add(vregSum01, vregExp01, vregSum01, pregLHalf); + Add(vregSum10, vregExp10, vregSum10, pregLHalf); + Add(vregSum11, vregExp11, vregSum11, pregLHalf); + Add(vregSum20, vregExp20, vregSum20, pregLHalf); + Add(vregSum21, vregExp21, vregSum21, pregLHalf); + Add(vregSum30, vregExp30, vregSum30, pregLHalf); + Add(vregSum31, vregExp31, vregSum31, pregLHalf); + + StoreAlign(((__ubuf__ T *&)srcUb00 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp00, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb01 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp01, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb10 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp10, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb11 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp11, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb20 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp20, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb21 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp21, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb30 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp30, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb31 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp31, pregLHalf); + } + Add(vregSum00, vregSum00, vregSum01, pregLHalf); + Add(vregSum10, vregSum10, vregSum11, pregLHalf); + Add(vregSum20, vregSum20, vregSum21, pregLHalf); + Add(vregSum30, vregSum30, vregSum31, pregLHalf); + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp00, srcUb00 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp10, srcUb10 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp20, srcUb20 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp30, srcUb30 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + Div(vregStore0, vregExp00, vregSum00, pregLHalf); + Div(vregStore1, vregExp10, vregSum10, pregLHalf); + Div(vregStore2, vregExp20, vregSum20, pregLHalf); + Div(vregStore3, vregExp30, vregSum30, pregLHalf); + + StoreAlign(((__ubuf__ T *&)inputAddr0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore0, pregLHalf); + StoreAlign(((__ubuf__ T *&)inputAddr1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore1, pregLHalf); + StoreAlign(((__ubuf__ T *&)inputAddr2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore2, pregLHalf); + StoreAlign(((__ubuf__ T *&)inputAddr3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore3, pregLHalf); + } + } + // 尾块处理 + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize % 4); ++loopSc) { + Duplicate(max0, minValue); + Duplicate(max00, minValue); + Duplicate(max01, minValue); + + Duplicate(vregSum00, 0, pregAll); + Duplicate(vregSum01, 0, pregAll); + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 2); ++loopM) { + LoadAlign(src00, srcUb00 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + LoadAlign(src01, srcUb01 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + + Max(max00, max00, src00, pregLHalf); + Max(max01, max01, src01, pregLHalf); + } + Max(max0, max00, max01, pregAll); + + for (uint16_t loopM = 0; loopM < uint16_t(ReduceSize / 2); ++loopM) { + LoadAlign(vregF32_00, srcUb00 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + LoadAlign(vregF32_01, srcUb01 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + + FusedExpSub(vregExp00, vregF32_00, max0, pregLHalf); + FusedExpSub(vregExp01, vregF32_01, max0, pregLHalf); + + Add(vregSum00, vregExp00, vregSum00, pregLHalf); + Add(vregSum01, vregExp01, vregSum01, pregLHalf); + + StoreAlign(((__ubuf__ T *&)srcUb00 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregExp00, pregLHalf); + StoreAlign(((__ubuf__ T *&)srcUb01 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregExp01, pregLHalf); + } + + Add(vregSum00, vregSum00, vregSum01, pregLHalf); + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp00, srcUb00 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + + Div(vregStore0, vregExp00, vregSum00, pregLHalf); + + StoreAlign(((__ubuf__ T *&)inputAddr0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregStore0, pregLHalf); + } + } +} + +template +__simd_vf__ inline void SoftmaxDndBase256(__ubuf__ T *inputAddr, __ubuf__ float *outputAddr, + const uint32_t RowSize, const uint32_t ReduceSize, const uint32_t vScRealSize, + const T minValue) +{ + RegTensor vregSum0; + RegTensor vregSum1; + RegTensor vregSum2; + RegTensor vregSum3; + + RegTensor vregExp0; + RegTensor vregExp1; + RegTensor vregExp2; + RegTensor vregExp3; + + RegTensor vregF32_0; + RegTensor vregF32_1; + RegTensor vregF32_2; + RegTensor vregF32_3; + + RegTensor vregStore0; + RegTensor vregStore1; + RegTensor vregStore2; + RegTensor vregStore3; + + RegTensor src0, src1, src2, src3; + RegTensor max0, max1, max2, max3; + MaskReg pregAll; + pregAll = CreateMask(); + + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize / 4); ++loopSc) { + for (uint16_t dChunk = 0; dChunk < 4; ++dChunk) { + uint32_t dOffset = dChunk * 64; + __ubuf__ float *srcUb0 = outputAddr + dOffset; + __ubuf__ float *srcUb1 = srcUb0 + ReduceSize * RowSize; + __ubuf__ float *srcUb2 = srcUb0 + ReduceSize * RowSize * 2; + __ubuf__ float *srcUb3 = srcUb0 + ReduceSize * RowSize * 3; + + __ubuf__ float *inputAddr0 = inputAddr + dOffset; + __ubuf__ float *inputAddr1 = inputAddr + dOffset + (ReduceSize * RowSize); + __ubuf__ float *inputAddr2 = inputAddr + dOffset + (ReduceSize * RowSize * 2); + __ubuf__ float *inputAddr3 = inputAddr + dOffset + (ReduceSize * RowSize * 3); + + Duplicate(max0, minValue); + Duplicate(max1, minValue); + Duplicate(max2, minValue); + Duplicate(max3, minValue); + + Duplicate(vregSum0, 0, pregAll); + Duplicate(vregSum1, 0, pregAll); + Duplicate(vregSum2, 0, pregAll); + Duplicate(vregSum3, 0, pregAll); + + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(src0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src1, srcUb1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src2, srcUb2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src3, srcUb3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + Max(max0, max0, src0, pregAll); + Max(max1, max1, src1, pregAll); + Max(max2, max2, src2, pregAll); + Max(max3, max3, src3, pregAll); + } + + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregF32_0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_1, srcUb1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_2, srcUb2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_3, srcUb3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + FusedExpSub(vregExp0, vregF32_0, max0, pregAll); + FusedExpSub(vregExp1, vregF32_1, max1, pregAll); + FusedExpSub(vregExp2, vregF32_2, max2, pregAll); + FusedExpSub(vregExp3, vregF32_3, max3, pregAll); + + Add(vregSum0, vregExp0, vregSum0, pregAll); + Add(vregSum1, vregExp1, vregSum1, pregAll); + Add(vregSum2, vregExp2, vregSum2, pregAll); + Add(vregSum3, vregExp3, vregSum3, pregAll); + + StoreAlign(((__ubuf__ T *&)srcUb0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp0, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp1, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp2, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp3, pregAll); + } + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp1, srcUb1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp2, srcUb2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp3, srcUb3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + Div(vregStore0, vregExp0, vregSum0, pregAll); + Div(vregStore1, vregExp1, vregSum1, pregAll); + Div(vregStore2, vregExp2, vregSum2, pregAll); + Div(vregStore3, vregExp3, vregSum3, pregAll); + + StoreAlign(((__ubuf__ T *&)inputAddr0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore0, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore1, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore2, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore3, pregAll); + } + } + } + + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize % 4); ++loopSc) { + for (uint16_t dChunk = 0; dChunk < 4; ++dChunk) { + uint32_t dOffset = dChunk * 64; + __ubuf__ float *srcUb0 = outputAddr + dOffset; + __ubuf__ float *inputAddr0 = inputAddr + dOffset; + + Duplicate(max0, minValue); + Duplicate(vregSum0, 0, pregAll); + + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(src0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + Max(max0, max0, src0, pregAll); + } + + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregF32_0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + FusedExpSub(vregExp0, vregF32_0, max0, pregAll); + Add(vregSum0, vregExp0, vregSum0, pregAll); + + StoreAlign(((__ubuf__ T *&)srcUb0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregExp0, pregAll); + } + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + Div(vregStore0, vregExp0, vregSum0, pregAll); + + StoreAlign(((__ubuf__ T *&)inputAddr0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregStore0, pregAll); + } + } + } +} + +template +__simd_vf__ inline void SoftmaxDndBase512(__ubuf__ T *inputAddr, __ubuf__ float *outputAddr, + const uint32_t RowSize, const uint32_t ReduceSize, const uint32_t vScRealSize, + const T minValue) +{ + RegTensor vregSum0; + RegTensor vregSum1; + RegTensor vregSum2; + RegTensor vregSum3; + + RegTensor vregExp0; + RegTensor vregExp1; + RegTensor vregExp2; + RegTensor vregExp3; + + RegTensor vregF32_0; + RegTensor vregF32_1; + RegTensor vregF32_2; + RegTensor vregF32_3; + + RegTensor vregStore0; + RegTensor vregStore1; + RegTensor vregStore2; + RegTensor vregStore3; + + RegTensor src0, src1, src2, src3; + RegTensor max0, max1, max2, max3; + MaskReg pregAll; + pregAll = CreateMask(); + + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize / 4); ++loopSc) { + for (uint16_t dChunk = 0; dChunk < 8; ++dChunk) { + uint32_t dOffset = dChunk * 64; + __ubuf__ float *srcUb0 = outputAddr + dOffset; + __ubuf__ float *srcUb1 = srcUb0 + ReduceSize * RowSize; + __ubuf__ float *srcUb2 = srcUb0 + ReduceSize * RowSize * 2; + __ubuf__ float *srcUb3 = srcUb0 + ReduceSize * RowSize * 3; + + __ubuf__ float *inputAddr0 = inputAddr + dOffset; + __ubuf__ float *inputAddr1 = inputAddr + dOffset + (ReduceSize * RowSize); + __ubuf__ float *inputAddr2 = inputAddr + dOffset + (ReduceSize * RowSize * 2); + __ubuf__ float *inputAddr3 = inputAddr + dOffset + (ReduceSize * RowSize * 3); + + Duplicate(max0, minValue); + Duplicate(max1, minValue); + Duplicate(max2, minValue); + Duplicate(max3, minValue); + + Duplicate(vregSum0, 0, pregAll); + Duplicate(vregSum1, 0, pregAll); + Duplicate(vregSum2, 0, pregAll); + Duplicate(vregSum3, 0, pregAll); + + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(src0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src1, srcUb1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src2, srcUb2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(src3, srcUb3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + Max(max0, max0, src0, pregAll); + Max(max1, max1, src1, pregAll); + Max(max2, max2, src2, pregAll); + Max(max3, max3, src3, pregAll); + } + + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregF32_0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_1, srcUb1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_2, srcUb2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregF32_3, srcUb3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + FusedExpSub(vregExp0, vregF32_0, max0, pregAll); + FusedExpSub(vregExp1, vregF32_1, max1, pregAll); + FusedExpSub(vregExp2, vregF32_2, max2, pregAll); + FusedExpSub(vregExp3, vregF32_3, max3, pregAll); + + Add(vregSum0, vregExp0, vregSum0, pregAll); + Add(vregSum1, vregExp1, vregSum1, pregAll); + Add(vregSum2, vregExp2, vregSum2, pregAll); + Add(vregSum3, vregExp3, vregSum3, pregAll); + + StoreAlign(((__ubuf__ T *&)srcUb0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp0, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp1, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp2, pregAll); + StoreAlign(((__ubuf__ T *&)srcUb3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregExp3, pregAll); + } + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp1, srcUb1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp2, srcUb2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + LoadAlign(vregExp3, srcUb3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4); + + Div(vregStore0, vregExp0, vregSum0, pregAll); + Div(vregStore1, vregExp1, vregSum1, pregAll); + Div(vregStore2, vregExp2, vregSum2, pregAll); + Div(vregStore3, vregExp3, vregSum3, pregAll); + + StoreAlign(((__ubuf__ T *&)inputAddr0 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore0, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr1 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore1, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr2 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore2, pregAll); + StoreAlign(((__ubuf__ T *&)inputAddr3 + loopM * RowSize + ReduceSize * RowSize * loopSc * 4), + vregStore3, pregAll); + } + } + } + + for (uint16_t loopSc = 0; loopSc < uint16_t(vScRealSize % 4); ++loopSc) { + for (uint16_t dChunk = 0; dChunk < 8; ++dChunk) { + uint32_t dOffset = dChunk * 64; + __ubuf__ float *srcUb0 = outputAddr + dOffset; + __ubuf__ float *inputAddr0 = inputAddr + dOffset; + + Duplicate(max0, minValue); + Duplicate(vregSum0, 0, pregAll); + + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(src0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + Max(max0, max0, src0, pregAll); + } + + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregF32_0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + FusedExpSub(vregExp0, vregF32_0, max0, pregAll); + Add(vregSum0, vregExp0, vregSum0, pregAll); + + StoreAlign(((__ubuf__ T *&)srcUb0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregExp0, pregAll); + } + + LocalMemBar(); + for (uint16_t loopM = 0; loopM < ReduceSize; ++loopM) { + LoadAlign(vregExp0, srcUb0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)); + Div(vregStore0, vregExp0, vregSum0, pregAll); + + StoreAlign(((__ubuf__ T *&)inputAddr0 + loopM * RowSize + ReduceSize * RowSize * (loopSc + vScRealSize / 4 * 4)), + vregStore0, pregAll); + } + } + } +} + +/* + * @ingroup ProcessVec1Vf + * @brief compute max = reducemax, exp(x-max)/sum(exp(x-max)) + * @param [out] dstTensor, output LocalTensor + * @param [in] srcTensor, input LocalTensor + * @param [in] RowSize, input rows + * @param [in] vScBaseSize, input columns, should be 256 bytes aligned, the value is originN aligned to 64 + * @param [in] vScRealSize, input origin columns, support range: 0 < originN <= 128 + * @param [in] scale, scale value + * @param [in] minValue, minimum value + */ + +template +__aicore__ inline void SoftmaxDnVF(const LocalTensor& dstTensor, const LocalTensor& srcTensor, + const uint32_t RowSize, const uint32_t ReduceSize, const uint32_t vScRealSize, + const T minValue, const uint32_t dDealSize) +{ + __ubuf__ T *inputAddr = (__ubuf__ T*) dstTensor.GetPhyAddr(); + __ubuf__ T *outputAddr = (__ubuf__ T*) srcTensor.GetPhyAddr(); + if (dDealSize == 8) { + SoftmaxDndBase8(inputAddr, outputAddr, RowSize, + ReduceSize, vScRealSize, minValue); + } else if (dDealSize == 16) { + SoftmaxDndBase16(inputAddr, outputAddr, RowSize, + ReduceSize, vScRealSize, minValue); + } else if (dDealSize == 32) { + SoftmaxDndBase32(inputAddr, outputAddr, RowSize, + ReduceSize, vScRealSize, minValue); + } else if (dDealSize == 64) { + SoftmaxDndBase64(inputAddr, outputAddr, RowSize, + ReduceSize, vScRealSize, minValue); + } else if (dDealSize == 128) { + SoftmaxDndBase128(inputAddr, outputAddr, RowSize, + ReduceSize, vScRealSize, minValue); + } else if (dDealSize == 256) { + SoftmaxDndBase256(inputAddr, outputAddr, RowSize, + ReduceSize, vScRealSize, minValue); + } else if (dDealSize == 512) { + SoftmaxDndBase512(inputAddr, outputAddr, RowSize, + ReduceSize, vScRealSize, minValue); + } +} +} +#endif // VF_SOFTMAX_H diff --git a/csrc/ascend/attention/compressor/op_kernel/compressor.cpp b/csrc/ascend/attention/compressor/op_kernel/compressor.cpp new file mode 100644 index 000000000..af25801e1 --- /dev/null +++ b/csrc/ascend/attention/compressor/op_kernel/compressor.cpp @@ -0,0 +1,87 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file compressor.cpp + * \brief + */ + +#if (__CCE_AICORE__ == 220) +#include "arch32/compressor_kernel.h" +#include "arch32/compressor_kernel_perf.h" +#else +#include "arch35/compressor_kernel.h" +#include "arch35/compressor_kernel_full_load.h" +#endif + +using namespace Compressor; + +#define INVOKE_COMPRESSOR_GENERAL_OP_IMPL(templateClass, ...) \ + do { \ + templateClass> op(&pipe, tilingData); \ + op.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos, stateBlockTable, \ + cuSeqlens, seqUsed, startPos, cmpKvOut, workspace); \ + op.Process(); \ + } while (0) + +#if (__CCE_AICORE__ == 220) +template +#else +template +#endif +__global__ __aicore__ void compressor( + __gm__ uint8_t *x, + __gm__ uint8_t *wKv, + __gm__ uint8_t *wGate, + __gm__ uint8_t *stateCache, + __gm__ uint8_t *ape, + __gm__ uint8_t *normWeight, + __gm__ uint8_t *ropeSin, + __gm__ uint8_t *ropeCos, + __gm__ uint8_t *stateBlockTable, + __gm__ uint8_t *cuSeqlens, + __gm__ uint8_t *seqUsed, + __gm__ uint8_t *startPos, + __gm__ uint8_t *cmpKvOut, + __gm__ uint8_t *stateCacheOut, + __gm__ uint8_t *workspace, + __gm__ uint8_t *tiling) { + REGISTER_TILING_DEFAULT(optiling::CompressorTilingData); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + GET_TILING_DATA_WITH_STRUCT(optiling::CompressorTilingData, tilingDataIn, tiling); + if constexpr (static_cast(TemplateId) == TEMPLATE_ID::EMPTY_X) { + return; + } + const optiling::CompressorTilingData *__restrict tilingData = &tilingDataIn; + TPipe pipe; + constexpr auto xLayout = static_cast(XLayout); + constexpr auto xDtype = static_cast(XDType); +#if (__CCE_AICORE__ == 220) + constexpr auto ropeDtype = static_cast(RopeDType); +#endif + constexpr auto coff = static_cast(Coff); + constexpr auto rotaryMode = static_cast(RotaryMode); +#if (__CCE_AICORE__ != 220) + constexpr auto cacheMode = static_cast(CacheMode); +#endif +#if (__CCE_AICORE__ == 220) + if constexpr (static_cast(TemplateId) == TEMPLATE_ID::PERF) { + INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernelPerf, xLayout, xDtype, ropeDtype, coff, rotaryMode); + } else { + INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernel, xLayout, xDtype, ropeDtype, coff, rotaryMode); + } +#else + if constexpr (static_cast(TemplateId) == TEMPLATE_ID::FULL_LOAD) { + INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernelFullLoad, xLayout, xDtype, coff, rotaryMode, cacheMode); + } else { + INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernel, xLayout, xDtype, coff, rotaryMode, cacheMode); + } +#endif +} diff --git a/csrc/ascend/attention/compressor_metadata/CMakeLists.txt b/csrc/ascend/attention/compressor_metadata/CMakeLists.txt new file mode 100644 index 000000000..e99a153f3 --- /dev/null +++ b/csrc/ascend/attention/compressor_metadata/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() diff --git a/csrc/ascend/attention/compressor_metadata/op_host/CMakeLists.txt b/csrc/ascend/attention/compressor_metadata/op_host/CMakeLists.txt new file mode 100644 index 000000000..0d8b29db5 --- /dev/null +++ b/csrc/ascend/attention/compressor_metadata/op_host/CMakeLists.txt @@ -0,0 +1,33 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + compressor_metadata_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME CompressorMetadata + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE compressor_metadata ACLNNTYPE aclnn) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/attention/compressor_metadata/op_host/compressor_metadata_def.cpp b/csrc/ascend/attention/compressor_metadata/op_host/compressor_metadata_def.cpp new file mode 100644 index 000000000..18e42d0f3 --- /dev/null +++ b/csrc/ascend/attention/compressor_metadata/op_host/compressor_metadata_def.cpp @@ -0,0 +1,69 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + */ + +#include "register/op_def_registry.h" + +namespace ops { +class CompressorMetadata : public OpDef { +public: + explicit CompressorMetadata(const char* name) : OpDef(name) + { + this->Input("ropeCos") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("ropeSin") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("cuSeqlens") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("startPos") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("kvBlockTable") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("compressCos") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("compressSin") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("slotMapping") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + this->Attr("kvBlockSize").Int(); + this->Attr("slotMappingFormat").Int(); + this->Attr("cmpRatio").Int(); + this->Attr("actualNumReqs").Int(); + + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + } +}; + +OP_ADD(CompressorMetadata); +} // namespace ops diff --git a/csrc/ascend/attention/compressor_metadata/op_host/compressor_metadata_tiling.cpp b/csrc/ascend/attention/compressor_metadata/op_host/compressor_metadata_tiling.cpp new file mode 100644 index 000000000..b9778a906 --- /dev/null +++ b/csrc/ascend/attention/compressor_metadata/op_host/compressor_metadata_tiling.cpp @@ -0,0 +1,289 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + */ + +#include "compressor_metadata_tiling.h" + +#include + +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" +#include "tiling_base/error_log.h" + +namespace optiling { +namespace { +constexpr uint32_t ROPE_COS_INDEX = 0; +constexpr uint32_t ROPE_SIN_INDEX = 1; +constexpr uint32_t CU_SEQLENS_INDEX = 2; +constexpr uint32_t START_POS_INDEX = 3; +constexpr uint32_t KV_BLOCK_TABLE_INDEX = 4; +constexpr uint32_t COMPRESS_COS_INDEX = 0; +constexpr uint32_t COMPRESS_SIN_INDEX = 1; +constexpr uint32_t SLOT_MAPPING_INDEX = 2; +constexpr uint32_t SLOT_MAPPING_FLAT = 1; +constexpr uint32_t SLOT_MAPPING_BLOCK_OFFSET = 2; +constexpr int64_t MAX_UINT32_VALUE = 0xFFFFFFFFLL; +constexpr int64_t MAX_INT32_VALUE = 0x7FFFFFFFLL; + +constexpr uint32_t TILING_KEY_FLOAT = 1; +constexpr uint32_t TILING_KEY_FLOAT16 = 2; +constexpr uint32_t TILING_KEY_BF16 = 3; +constexpr uint32_t ALIGN_BYTES = 32; +constexpr uint32_t BUFFER_NUM = 2; +constexpr uint32_t MAX_TILE_ROWS = 512; +constexpr uint32_t MAX_DATACOPY_BLOCK_COUNT = 4095; +constexpr uint32_t ROWS_PER_CORE_TARGET = 64; +constexpr uint32_t UB_RESERVED_BYTES = 16 * 1024; + +uint32_t AlignUp(uint64_t value, uint32_t align) +{ + return static_cast((value + align - 1) / align * align); +} + +uint32_t CeilDiv(uint64_t lhs, uint64_t rhs) +{ + return static_cast((lhs + rhs - 1) / rhs); +} +} // namespace + +static ge::graphStatus CompressorMetadataTilingFunc(gert::TilingContext* context) +{ + auto platformInfo = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + uint32_t aivCoreNum = ascendcPlatform.GetCoreNumAiv(); + if (aivCoreNum == 0) { + aivCoreNum = ascendcPlatform.GetCoreNum(); + } + if (aivCoreNum == 0) { + OP_LOGE(context->GetNodeName(), "Failed to get AIV core num."); + return ge::GRAPH_FAILED; + } + uint64_t ubSize = 0; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); + if (ubSize == 0) { + OP_LOGE(context->GetNodeName(), "Failed to get UB size."); + return ge::GRAPH_FAILED; + } + + auto outputShape = context->GetOutputShape(COMPRESS_COS_INDEX); + auto compressSinShape = context->GetOutputShape(COMPRESS_SIN_INDEX); + auto slotMappingShape = context->GetOutputShape(SLOT_MAPPING_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, outputShape); + OP_CHECK_NULL_WITH_CONTEXT(context, compressSinShape); + OP_CHECK_NULL_WITH_CONTEXT(context, slotMappingShape); + auto outputDimNum = outputShape->GetStorageShape().GetDimNum(); + if (outputDimNum < 2) { + OP_LOGE(context->GetNodeName(), "compressCos dim num should be at least 2."); + return ge::GRAPH_FAILED; + } + if (compressSinShape->GetStorageShape().GetDimNum() != outputDimNum) { + OP_LOGE(context->GetNodeName(), "compressCos and compressSin dim num mismatch."); + return ge::GRAPH_FAILED; + } + for (size_t dimIdx = 0; dimIdx < outputDimNum; ++dimIdx) { + if (compressSinShape->GetStorageShape().GetDim(dimIdx) != outputShape->GetStorageShape().GetDim(dimIdx)) { + OP_LOGE(context->GetNodeName(), "compressCos and compressSin shape mismatch."); + return ge::GRAPH_FAILED; + } + } + int64_t numRows = outputShape->GetStorageShape().GetDim(0); + int64_t ropeDim = outputShape->GetStorageShape().GetDim(outputDimNum - 1); + if (numRows <= 0 || ropeDim <= 0 || numRows > MAX_UINT32_VALUE || ropeDim > MAX_UINT32_VALUE) { + OP_LOGE(context->GetNodeName(), "compressCos shape is invalid."); + return ge::GRAPH_FAILED; + } + + auto ropeCosShape = context->GetInputShape(ROPE_COS_INDEX); + auto ropeSinShape = context->GetInputShape(ROPE_SIN_INDEX); + auto cuSeqlensShape = context->GetInputShape(CU_SEQLENS_INDEX); + auto startPosShape = context->GetInputShape(START_POS_INDEX); + auto kvBlockTableShape = context->GetInputShape(KV_BLOCK_TABLE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, ropeCosShape); + OP_CHECK_NULL_WITH_CONTEXT(context, ropeSinShape); + OP_CHECK_NULL_WITH_CONTEXT(context, cuSeqlensShape); + OP_CHECK_NULL_WITH_CONTEXT(context, startPosShape); + OP_CHECK_NULL_WITH_CONTEXT(context, kvBlockTableShape); + if (ropeCosShape->GetStorageShape().GetDimNum() != 2 || ropeSinShape->GetStorageShape().GetDimNum() != 2) { + OP_LOGE(context->GetNodeName(), "ropeCos and ropeSin should be 2D tensors."); + return ge::GRAPH_FAILED; + } + int64_t ropeRows = ropeCosShape->GetStorageShape().GetDim(0); + int64_t ropeCosDim = ropeCosShape->GetStorageShape().GetDim(1); + if (ropeRows <= 0 || ropeCosDim <= 0 || ropeRows > MAX_UINT32_VALUE || ropeCosDim != ropeDim || + ropeSinShape->GetStorageShape().GetDim(0) != ropeRows || + ropeSinShape->GetStorageShape().GetDim(1) != ropeCosDim) { + OP_LOGE(context->GetNodeName(), "ropeCos and ropeSin shape mismatch."); + return ge::GRAPH_FAILED; + } + int64_t cuSeqlensDim0 = cuSeqlensShape->GetStorageShape().GetDim(0); + if (cuSeqlensDim0 < 2 || cuSeqlensDim0 > MAX_UINT32_VALUE) { + OP_LOGE(context->GetNodeName(), "cuSeqlens dim0 should be at least 2."); + return ge::GRAPH_FAILED; + } + if (startPosShape->GetStorageShape().GetDimNum() != 1 || + startPosShape->GetStorageShape().GetDim(0) <= 0 || + startPosShape->GetStorageShape().GetDim(0) > MAX_UINT32_VALUE) { + OP_LOGE(context->GetNodeName(), "startPos should be a non-empty 1D tensor."); + return ge::GRAPH_FAILED; + } + if (kvBlockTableShape->GetStorageShape().GetDimNum() != 2) { + OP_LOGE(context->GetNodeName(), "kvBlockTable should be a 2D tensor."); + return ge::GRAPH_FAILED; + } + int64_t kvBlockTableRows = kvBlockTableShape->GetStorageShape().GetDim(0); + int64_t kvBlockTableStride = kvBlockTableShape->GetStorageShape().GetDim(1); + if (kvBlockTableRows <= 0 || kvBlockTableStride <= 0 || kvBlockTableRows > MAX_UINT32_VALUE || + kvBlockTableStride > MAX_UINT32_VALUE) { + OP_LOGE(context->GetNodeName(), "kvBlockTable shape is invalid."); + return ge::GRAPH_FAILED; + } + + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const int64_t* kvBlockSizePtr = attrs->GetInt(0); + const int64_t* slotMappingFormatPtr = attrs->GetInt(1); + const int64_t* cmpRatioPtr = attrs->GetInt(2); + const int64_t* actualNumReqsPtr = attrs->GetInt(3); + OP_CHECK_NULL_WITH_CONTEXT(context, kvBlockSizePtr); + OP_CHECK_NULL_WITH_CONTEXT(context, slotMappingFormatPtr); + OP_CHECK_NULL_WITH_CONTEXT(context, cmpRatioPtr); + OP_CHECK_NULL_WITH_CONTEXT(context, actualNumReqsPtr); + if (*kvBlockSizePtr <= 0 || *kvBlockSizePtr > MAX_INT32_VALUE) { + OP_LOGE(context->GetNodeName(), "kvBlockSize should be in (0, INT32_MAX]."); + return ge::GRAPH_FAILED; + } + if (*cmpRatioPtr <= 0 || *cmpRatioPtr > MAX_UINT32_VALUE) { + OP_LOGE(context->GetNodeName(), "cmpRatio should be in (0, UINT32_MAX]."); + return ge::GRAPH_FAILED; + } + if (*slotMappingFormatPtr != SLOT_MAPPING_FLAT && *slotMappingFormatPtr != SLOT_MAPPING_BLOCK_OFFSET) { + OP_LOGE(context->GetNodeName(), "slotMappingFormat should be 1(flat) or 2(block_offset)."); + return ge::GRAPH_FAILED; + } + auto slotMappingDimNum = slotMappingShape->GetStorageShape().GetDimNum(); + if ((*slotMappingFormatPtr == SLOT_MAPPING_FLAT && + (slotMappingDimNum != 1 || slotMappingShape->GetStorageShape().GetDim(0) != numRows)) || + (*slotMappingFormatPtr == SLOT_MAPPING_BLOCK_OFFSET && + (slotMappingDimNum != 2 || slotMappingShape->GetStorageShape().GetDim(0) != numRows || + slotMappingShape->GetStorageShape().GetDim(1) != 2))) { + OP_LOGE(context->GetNodeName(), "slotMapping shape does not match slotMappingFormat."); + return ge::GRAPH_FAILED; + } + if (*actualNumReqsPtr <= 0 || *actualNumReqsPtr >= cuSeqlensDim0 || + *actualNumReqsPtr > startPosShape->GetStorageShape().GetDim(0) || + *actualNumReqsPtr > kvBlockTableRows || + *actualNumReqsPtr > MAX_UINT32_VALUE) { + OP_LOGE(context->GetNodeName(), "actualNumReqs is invalid."); + return ge::GRAPH_FAILED; + } + + CompressorMetadataTilingData tilingData; + tilingData.set_numRows(static_cast(numRows)); + tilingData.set_numReqs(static_cast(cuSeqlensDim0 - 1)); + tilingData.set_actualNumReqs(static_cast(*actualNumReqsPtr)); + tilingData.set_ropeRows(static_cast(ropeRows)); + tilingData.set_ropeDim(static_cast(ropeDim)); + tilingData.set_kvBlockTableStride(static_cast(kvBlockTableStride)); + tilingData.set_kvBlockSize(static_cast(*kvBlockSizePtr)); + tilingData.set_slotMappingFormat(static_cast(*slotMappingFormatPtr)); + tilingData.set_cmpRatio(static_cast(*cmpRatioPtr)); + + auto ropeDesc = context->GetInputDesc(ROPE_COS_INDEX); + auto ropeSinDesc = context->GetInputDesc(ROPE_SIN_INDEX); + auto compressCosDesc = context->GetOutputDesc(COMPRESS_COS_INDEX); + auto compressSinDesc = context->GetOutputDesc(COMPRESS_SIN_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, ropeDesc); + OP_CHECK_NULL_WITH_CONTEXT(context, ropeSinDesc); + OP_CHECK_NULL_WITH_CONTEXT(context, compressCosDesc); + OP_CHECK_NULL_WITH_CONTEXT(context, compressSinDesc); + auto ropeDtype = ropeDesc->GetDataType(); + if (ropeSinDesc->GetDataType() != ropeDtype || + compressCosDesc->GetDataType() != ropeDtype || + compressSinDesc->GetDataType() != ropeDtype) { + OP_LOGE(context->GetNodeName(), "rope and compress output dtypes should match."); + return ge::GRAPH_FAILED; + } + + uint64_t tilingKey = 0; + uint32_t dtypeSize = 0; + if (ropeDtype == ge::DataType::DT_FLOAT) { + tilingKey = TILING_KEY_FLOAT; + dtypeSize = sizeof(float); + } else if (ropeDtype == ge::DataType::DT_FLOAT16) { + tilingKey = TILING_KEY_FLOAT16; + dtypeSize = sizeof(uint16_t); + } else if (ropeDtype == ge::DataType::DT_BF16) { + tilingKey = TILING_KEY_BF16; + dtypeSize = sizeof(uint16_t); + } else { + OP_LOGE(context->GetNodeName(), "Unsupported rope dtype."); + return ge::GRAPH_FAILED; + } + + uint32_t actualNumReqs = static_cast(*actualNumReqsPtr); + uint32_t cmpRatio = static_cast(*cmpRatioPtr); + if (static_cast(ropeDim) * dtypeSize > MAX_UINT32_VALUE || + (static_cast(actualNumReqs) + 1) * sizeof(int32_t) > MAX_UINT32_VALUE) { + OP_LOGE(context->GetNodeName(), "tiling byte size exceeds UINT32_MAX."); + return ge::GRAPH_FAILED; + } + uint32_t ropeRowBytes = static_cast(ropeDim) * dtypeSize; + if (static_cast(cmpRatio - 1) * ropeRowBytes > MAX_UINT32_VALUE) { + OP_LOGE(context->GetNodeName(), "rope stride exceeds UINT32_MAX."); + return ge::GRAPH_FAILED; + } + uint32_t ropeRowBytesAligned = AlignUp(ropeRowBytes, ALIGN_BYTES); + uint32_t slotCols = (*slotMappingFormatPtr == SLOT_MAPPING_FLAT) ? 1U : 2U; + uint32_t reqTableBytes = AlignUp((static_cast(actualNumReqs) + 1) * sizeof(int32_t), ALIGN_BYTES); + uint64_t fixedUbBytes = static_cast(reqTableBytes) * 3 + ALIGN_BYTES + UB_RESERVED_BYTES; + uint64_t rowUbBytes = + static_cast(BUFFER_NUM) * ropeRowBytesAligned * 2 + slotCols * sizeof(int32_t) + sizeof(int32_t); + if (rowUbBytes > MAX_UINT32_VALUE) { + OP_LOGE(context->GetNodeName(), "row UB footprint exceeds UINT32_MAX."); + return ge::GRAPH_FAILED; + } + uint64_t minUbBytes = static_cast(reqTableBytes) * 3 + ALIGN_BYTES + rowUbBytes; + if (ubSize <= minUbBytes) { + OP_LOGE(context->GetNodeName(), "UB size is insufficient for compressor metadata."); + return ge::GRAPH_FAILED; + } + uint32_t tileRows = 1; + if (ubSize > fixedUbBytes && rowUbBytes > 0) { + tileRows = static_cast((ubSize - fixedUbBytes) / rowUbBytes); + tileRows = std::max(tileRows, 1U); + } + tileRows = std::min(tileRows, MAX_TILE_ROWS); + tileRows = std::min(tileRows, MAX_DATACOPY_BLOCK_COUNT); + + uint32_t usedCoreNum = + std::min(aivCoreNum, std::max(1U, CeilDiv(static_cast(numRows), ROWS_PER_CORE_TARGET))); + tilingData.set_usedCoreNum(usedCoreNum); + tilingData.set_tileRows(tileRows); + tilingData.set_ropeRowBytes(ropeRowBytes); + tilingData.set_ropeRowBytesAligned(ropeRowBytesAligned); + tilingData.set_slotCols(slotCols); + + size_t* workspaceSize = context->GetWorkspaceSizes(1); + OP_CHECK_NULL_WITH_CONTEXT(context, workspaceSize); + *workspaceSize = 0; + context->SetBlockDim(usedCoreNum); + context->SetTilingKey(tilingKey); + + auto rawTilingData = context->GetRawTilingData(); + OP_CHECK_NULL_WITH_CONTEXT(context, rawTilingData); + tilingData.SaveToBuffer(rawTilingData->GetData(), rawTilingData->GetCapacity()); + rawTilingData->SetDataSize(tilingData.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus TilingParseForCompressorMetadata(gert::TilingParseContext* context) +{ + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(CompressorMetadata) + .Tiling(CompressorMetadataTilingFunc) + .TilingParse(TilingParseForCompressorMetadata); + +} // namespace optiling diff --git a/csrc/ascend/attention/compressor_metadata/op_host/compressor_metadata_tiling.h b/csrc/ascend/attention/compressor_metadata/op_host/compressor_metadata_tiling.h new file mode 100644 index 000000000..85baa2dc4 --- /dev/null +++ b/csrc/ascend/attention/compressor_metadata/op_host/compressor_metadata_tiling.h @@ -0,0 +1,32 @@ +#ifndef COMPRESSOR_METADATA_TILING_H +#define COMPRESSOR_METADATA_TILING_H + +#include "register/tilingdata_base.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(CompressorMetadataTilingData) + TILING_DATA_FIELD_DEF(uint32_t, numRows); + TILING_DATA_FIELD_DEF(uint32_t, numReqs); + TILING_DATA_FIELD_DEF(uint32_t, actualNumReqs); + TILING_DATA_FIELD_DEF(uint32_t, ropeRows); + TILING_DATA_FIELD_DEF(uint32_t, ropeDim); + TILING_DATA_FIELD_DEF(uint32_t, kvBlockTableStride); + TILING_DATA_FIELD_DEF(uint32_t, kvBlockSize); + TILING_DATA_FIELD_DEF(uint32_t, slotMappingFormat); + TILING_DATA_FIELD_DEF(uint32_t, cmpRatio); + TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum); + TILING_DATA_FIELD_DEF(uint32_t, tileRows); + TILING_DATA_FIELD_DEF(uint32_t, ropeRowBytes); + TILING_DATA_FIELD_DEF(uint32_t, ropeRowBytesAligned); + TILING_DATA_FIELD_DEF(uint32_t, slotCols); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(CompressorMetadata, CompressorMetadataTilingData) + +struct CompressorMetadataCompileInfo { + uint32_t coreNum; + uint64_t ubSizePlatForm; +}; +} // namespace optiling + +#endif diff --git a/csrc/ascend/attention/compressor_metadata/op_kernel/compressor_metadata.cpp b/csrc/ascend/attention/compressor_metadata/op_kernel/compressor_metadata.cpp new file mode 100644 index 000000000..76b74708b --- /dev/null +++ b/csrc/ascend/attention/compressor_metadata/op_kernel/compressor_metadata.cpp @@ -0,0 +1,38 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + */ + +#include "compressor_metadata.h" + +extern "C" __global__ __aicore__ void compressor_metadata( + GM_ADDR ropeCos, + GM_ADDR ropeSin, + GM_ADDR cuSeqlens, + GM_ADDR startPos, + GM_ADDR kvBlockTable, + GM_ADDR compressCos, + GM_ADDR compressSin, + GM_ADDR slotMapping, + GM_ADDR workspace, + GM_ADDR tiling) +{ + REGISTER_TILING_DEFAULT(CompressorMetadata::CompressorMetadataTilingData); + GET_TILING_DATA(tilingData, tiling); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); + + AscendC::TPipe pipe; + + if (TILING_KEY_IS(1)) { + CompressorMetadata::CompressorMetadataKernel op; + op.Init(&tilingData, &pipe); + op.Process(ropeCos, ropeSin, cuSeqlens, startPos, kvBlockTable, compressCos, compressSin, slotMapping, workspace); + } else if (TILING_KEY_IS(2)) { + CompressorMetadata::CompressorMetadataKernel op; + op.Init(&tilingData, &pipe); + op.Process(ropeCos, ropeSin, cuSeqlens, startPos, kvBlockTable, compressCos, compressSin, slotMapping, workspace); + } else if (TILING_KEY_IS(3)) { + CompressorMetadata::CompressorMetadataKernel op; + op.Init(&tilingData, &pipe); + op.Process(ropeCos, ropeSin, cuSeqlens, startPos, kvBlockTable, compressCos, compressSin, slotMapping, workspace); + } +} diff --git a/csrc/ascend/attention/compressor_metadata/op_kernel/compressor_metadata.h b/csrc/ascend/attention/compressor_metadata/op_kernel/compressor_metadata.h new file mode 100644 index 000000000..7e1c0209a --- /dev/null +++ b/csrc/ascend/attention/compressor_metadata/op_kernel/compressor_metadata.h @@ -0,0 +1,428 @@ +#ifndef COMPRESSOR_METADATA_H +#define COMPRESSOR_METADATA_H + +#include "kernel_operator.h" + +namespace CompressorMetadata { +using namespace AscendC; + +constexpr uint32_t SLOT_MAPPING_FLAT = 1; +constexpr uint32_t ALIGN_BYTES = 32; +constexpr uint32_t BUFFER_NUM = 2; +constexpr int64_t MAX_INT32_VALUE = 0x7FFFFFFFLL; + +__aicore__ inline uint32_t MinU32(uint32_t lhs, uint32_t rhs) +{ + return lhs < rhs ? lhs : rhs; +} + +__aicore__ inline uint32_t MaxU32(uint32_t lhs, uint32_t rhs) +{ + return lhs > rhs ? lhs : rhs; +} + +__aicore__ inline uint32_t AlignUpU32(uint32_t value, uint32_t align) +{ + return (value + align - 1) / align * align; +} + +__aicore__ inline uint32_t Int32BytesU32(uint32_t elems) +{ + return elems * static_cast(sizeof(int32_t)); +} + +__aicore__ inline void PipeMte2ToS() +{ + event_t eventID = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S)); + SetFlag(eventID); + WaitFlag(eventID); +} + +__aicore__ inline void PipeMte3ToS() +{ + event_t eventID = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S)); + SetFlag(eventID); + WaitFlag(eventID); +} + +__aicore__ inline void PipeSToMte3() +{ + event_t eventID = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3)); + SetFlag(eventID); + WaitFlag(eventID); +} + +__aicore__ inline void PipeVToMte3() +{ + event_t eventID = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventID); + WaitFlag(eventID); +} + +struct CompressorMetadataTilingData { + uint32_t numRows; + uint32_t numReqs; + uint32_t actualNumReqs; + uint32_t ropeRows; + uint32_t ropeDim; + uint32_t kvBlockTableStride; + uint32_t kvBlockSize; + uint32_t slotMappingFormat; + uint32_t cmpRatio; + uint32_t usedCoreNum; + uint32_t tileRows; + uint32_t ropeRowBytes; + uint32_t ropeRowBytesAligned; + uint32_t slotCols; +}; + +template +class CompressorMetadataKernel { +public: + __aicore__ inline CompressorMetadataKernel() {} + + __aicore__ inline void Init(CompressorMetadataTilingData* tilingData, TPipe* pipe) + { + numRows_ = tilingData->numRows; + actualNumReqs_ = tilingData->actualNumReqs; + ropeRows_ = tilingData->ropeRows; + ropeDim_ = tilingData->ropeDim; + kvBlockTableStride_ = tilingData->kvBlockTableStride; + kvBlockSize_ = tilingData->kvBlockSize; + slotMappingFormat_ = tilingData->slotMappingFormat; + cmpRatio_ = tilingData->cmpRatio; + tileRows_ = tilingData->tileRows; + ropeRowBytes_ = tilingData->ropeRowBytes; + ropeRowBytesAligned_ = tilingData->ropeRowBytesAligned; + slotCols_ = tilingData->slotCols; + reqTableBytes_ = AlignUpU32(Int32BytesU32(actualNumReqs_ + 1), ALIGN_BYTES); + ropeDimAligned_ = ropeRowBytesAligned_ / sizeof(T); + ropePadElems_ = ropeDimAligned_ - ropeDim_; + slotTileBytes_ = AlignUpU32(Int32BytesU32(tileRows_ * slotCols_), ALIGN_BYTES); + blockTableTileBytes_ = AlignUpU32(Int32BytesU32(tileRows_), ALIGN_BYTES); + + pipe->InitBuffer(prefixBuf_, reqTableBytes_); + pipe->InitBuffer(startPosBuf_, reqTableBytes_); + pipe->InitBuffer(cuSeqlensBuf_, reqTableBytes_); + pipe->InitBuffer(blockTableBuf_, blockTableTileBytes_); + pipe->InitBuffer(slotBuf_, slotTileBytes_); + pipe->InitBuffer(cosQueue_, BUFFER_NUM, tileRows_ * ropeRowBytesAligned_); + pipe->InitBuffer(sinQueue_, BUFFER_NUM, tileRows_ * ropeRowBytesAligned_); + } + + __aicore__ inline void Process( + GM_ADDR ropeCos, + GM_ADDR ropeSin, + GM_ADDR cuSeqlens, + GM_ADDR startPos, + GM_ADDR kvBlockTable, + GM_ADDR compressCos, + GM_ADDR compressSin, + GM_ADDR slotMapping, + GM_ADDR) + { + ropeCosGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(ropeCos)); + ropeSinGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(ropeSin)); + cuSeqlensGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(cuSeqlens)); + startPosGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(startPos)); + kvBlockTableGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(kvBlockTable)); + compressCosGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(compressCos)); + compressSinGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(compressSin)); + slotMappingGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(slotMapping)); + + LocalTensor prefixLocal = prefixBuf_.Get(); + LocalTensor startPosLocal = startPosBuf_.Get(); + LocalTensor cuSeqlensLocal = cuSeqlensBuf_.Get(); + BuildCompressedPrefix(prefixLocal, startPosLocal, cuSeqlensLocal); + + uint32_t validRows = static_cast(prefixLocal.GetValue(actualNumReqs_)); + validRows = MinU32(validRows, numRows_); + ProcessValidRows(prefixLocal, startPosLocal, validRows); + ProcessPaddingRows(validRows); + } + +private: + __aicore__ inline void BuildCompressedPrefix( + LocalTensor& prefixLocal, + LocalTensor& startPosLocal, + LocalTensor& cuSeqlensLocal) + { + DataCopyExtParams startCopyParams{1, Int32BytesU32(actualNumReqs_), 0, 0, 0}; + DataCopyExtParams cuCopyParams{1, Int32BytesU32(actualNumReqs_ + 1), 0, 0, 0}; + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(startPosLocal, startPosGm_, startCopyParams, padParams); + DataCopyPad(cuSeqlensLocal, cuSeqlensGm_, cuCopyParams, padParams); + PipeMte2ToS(); + + uint32_t prefix = 0; + prefixLocal.SetValue(0, 0); + for (uint32_t reqIdx = 0; reqIdx < actualNumReqs_; ++reqIdx) { + int64_t startPos = static_cast(startPosLocal.GetValue(reqIdx)); + int64_t seqLen = static_cast(cuSeqlensLocal.GetValue(reqIdx + 1)) - + static_cast(cuSeqlensLocal.GetValue(reqIdx)); + uint32_t compressedRows = 0; + if (startPos >= 0 && seqLen > 0) { + compressedRows = static_cast(((startPos + seqLen) / cmpRatio_) - (startPos / cmpRatio_)); + } + prefix += compressedRows; + prefixLocal.SetValue(reqIdx + 1, static_cast(prefix)); + } + } + + __aicore__ inline void SplitRange(uint32_t totalRows, uint32_t& begin, uint32_t& end) + { + uint32_t blockIdx = GetBlockIdx(); + uint32_t blockNum = MaxU32(GetBlockNum(), 1); + uint32_t rowsPerBlock = (totalRows + blockNum - 1) / blockNum; + begin = MinU32(blockIdx * rowsPerBlock, totalRows); + end = MinU32(begin + rowsPerBlock, totalRows); + } + + __aicore__ inline uint32_t FindRequest(LocalTensor& prefixLocal, uint32_t row) + { + uint32_t reqIdx = 0; + while (reqIdx < actualNumReqs_ && static_cast(prefixLocal.GetValue(reqIdx + 1)) <= row) { + ++reqIdx; + } + return reqIdx; + } + + __aicore__ inline void ProcessValidRows( + LocalTensor& prefixLocal, + LocalTensor& startPosLocal, + uint32_t validRows) + { + uint32_t begin = 0; + uint32_t end = 0; + SplitRange(validRows, begin, end); + if (begin >= end) { + return; + } + + uint32_t reqIdx = FindRequest(prefixLocal, begin); + uint32_t row = begin; + while (row < end && reqIdx < actualNumReqs_) { + uint32_t reqBegin = static_cast(prefixLocal.GetValue(reqIdx)); + uint32_t reqEnd = static_cast(prefixLocal.GetValue(reqIdx + 1)); + if (row >= reqEnd) { + ++reqIdx; + continue; + } + uint32_t rowsInReq = MinU32(end - row, reqEnd - row); + int64_t startPos = static_cast(startPosLocal.GetValue(reqIdx)); + uint32_t localCompressedIdx = row - reqBegin; + // KV slot uses compressed position; RoPE uses the original group-start position. + uint32_t compressedPos = static_cast(startPos / cmpRatio_) + localCompressedIdx; + ProcessRequestRows(reqIdx, row, compressedPos, rowsInReq); + row += rowsInReq; + } + } + + __aicore__ inline void ProcessRequestRows( + uint32_t reqIdx, + uint32_t outputRow, + uint32_t compressedPos, + uint32_t rows) + { + while (rows > 0) { + uint32_t blockOffset = compressedPos % kvBlockSize_; + uint32_t rowsToBlockEnd = kvBlockSize_ - blockOffset; + uint32_t curRows = MinU32(rows, tileRows_); + curRows = MinU32(curRows, rowsToBlockEnd); + ProcessTile(reqIdx, outputRow, compressedPos, curRows); + outputRow += curRows; + compressedPos += curRows; + rows -= curRows; + } + } + + __aicore__ inline void ProcessTile( + uint32_t reqIdx, + uint32_t outputRow, + uint32_t compressedPos, + uint32_t rows) + { + uint32_t blockIdOffset = compressedPos / kvBlockSize_; + if (blockIdOffset >= kvBlockTableStride_) { + WriteInvalidTile(outputRow, rows); + return; + } + + LocalTensor blockTableLocal = blockTableBuf_.Get(); + DataCopyExtParams blockCopyParams{1, Int32BytesU32(1), 0, 0, 0}; + DataCopyPadExtParams padParams{true, 0, 0, 0}; + uint64_t blockTableGmOffset = static_cast(reqIdx) * kvBlockTableStride_ + blockIdOffset; + DataCopyPad(blockTableLocal, kvBlockTableGm_[blockTableGmOffset], blockCopyParams, padParams); + PipeMte2ToS(); + + int32_t blockId = blockTableLocal.GetValue(0); + if (blockId < 0) { + WriteInvalidTile(outputRow, rows); + return; + } + + uint32_t blockOffset = compressedPos % kvBlockSize_; + if (slotMappingFormat_ == SLOT_MAPPING_FLAT) { + int64_t maxSlot = static_cast(blockId) * kvBlockSize_ + blockOffset + rows - 1; + if (maxSlot > MAX_INT32_VALUE) { + WriteInvalidTile(outputRow, rows); + return; + } + } + + uint64_t lastRopePos = (static_cast(compressedPos) + rows - 1) * cmpRatio_; + if (lastRopePos >= ropeRows_) { + WriteInvalidTile(outputRow, rows); + return; + } + + CopyRopeTile(outputRow, compressedPos, rows); + WriteSlotTile(outputRow, compressedPos, rows, blockId); + } + + __aicore__ inline void CopyRopeTile(uint32_t outputRow, uint32_t compressedPos, uint32_t rows) + { + LocalTensor cosLocal = cosQueue_.AllocTensor(); + LocalTensor sinLocal = sinQueue_.AllocTensor(); + uint64_t ropePos = static_cast(compressedPos) * cmpRatio_; + uint32_t srcStride = (cmpRatio_ - 1) * ropeRowBytes_; + + DataCopyExtParams copyInParams{ + static_cast(rows), ropeRowBytes_, srcStride, 0, 0}; + DataCopyPadExtParams padParams{true, 0, static_cast(ropePadElems_), 0}; + DataCopyPad(cosLocal, ropeCosGm_[ropePos * ropeDim_], copyInParams, padParams); + DataCopyPad(sinLocal, ropeSinGm_[ropePos * ropeDim_], copyInParams, padParams); + PipeMte2ToS(); + + DataCopyExtParams copyOutParams{ + static_cast(rows), ropeRowBytes_, 0, 0, 0}; + uint64_t outputBase = static_cast(outputRow) * ropeDim_; + DataCopyPad(compressCosGm_[outputBase], cosLocal, copyOutParams); + DataCopyPad(compressSinGm_[outputBase], sinLocal, copyOutParams); + PipeMte3ToS(); + + cosQueue_.FreeTensor(cosLocal); + sinQueue_.FreeTensor(sinLocal); + } + + __aicore__ inline void WriteSlotTile( + uint32_t outputRow, + uint32_t compressedPos, + uint32_t rows, + int32_t blockId) + { + LocalTensor slotLocal = slotBuf_.Get(); + int32_t blockOffset = static_cast(compressedPos % kvBlockSize_); + if (slotMappingFormat_ == SLOT_MAPPING_FLAT) { + int32_t slotBase = blockId * static_cast(kvBlockSize_) + blockOffset; + for (uint32_t row = 0; row < rows; ++row) { + slotLocal.SetValue(row, slotBase + static_cast(row)); + } + } else { + for (uint32_t row = 0; row < rows; ++row) { + uint32_t slotOffset = row * slotCols_; + slotLocal.SetValue(slotOffset, blockId); + slotLocal.SetValue(slotOffset + 1, blockOffset + static_cast(row)); + } + } + + DataCopyExtParams slotCopyParams{1, Int32BytesU32(rows * slotCols_), 0, 0, 0}; + PipeSToMte3(); + DataCopyPad(slotMappingGm_[static_cast(outputRow) * slotCols_], slotLocal, slotCopyParams); + PipeMte3ToS(); + } + + __aicore__ inline void WriteInvalidTile(uint32_t outputRow, uint32_t rows) + { + LocalTensor cosLocal = cosQueue_.AllocTensor(); + LocalTensor sinLocal = sinQueue_.AllocTensor(); + + Duplicate(cosLocal, static_cast(1.0f), rows * ropeDimAligned_); + Duplicate(sinLocal, static_cast(0.0f), rows * ropeDimAligned_); + PipeVToMte3(); + + DataCopyExtParams ropeCopyParams{ + static_cast(rows), ropeRowBytes_, 0, 0, 0}; + uint64_t outputBase = static_cast(outputRow) * ropeDim_; + DataCopyPad(compressCosGm_[outputBase], cosLocal, ropeCopyParams); + DataCopyPad(compressSinGm_[outputBase], sinLocal, ropeCopyParams); + PipeMte3ToS(); + + cosQueue_.FreeTensor(cosLocal); + sinQueue_.FreeTensor(sinLocal); + + LocalTensor slotLocal = slotBuf_.Get(); + if (slotMappingFormat_ == SLOT_MAPPING_FLAT) { + for (uint32_t row = 0; row < rows; ++row) { + slotLocal.SetValue(row, -1); + } + } else { + int32_t padOffset = static_cast(kvBlockSize_ - 1); + for (uint32_t row = 0; row < rows; ++row) { + uint32_t slotOffset = row * slotCols_; + slotLocal.SetValue(slotOffset, -1); + slotLocal.SetValue(slotOffset + 1, padOffset); + } + } + PipeSToMte3(); + DataCopyExtParams slotCopyParams{1, Int32BytesU32(rows * slotCols_), 0, 0, 0}; + DataCopyPad(slotMappingGm_[static_cast(outputRow) * slotCols_], slotLocal, slotCopyParams); + PipeMte3ToS(); + } + + __aicore__ inline void ProcessPaddingRows(uint32_t validRows) + { + if (validRows >= numRows_) { + return; + } + uint32_t padRows = numRows_ - validRows; + uint32_t begin = 0; + uint32_t end = 0; + SplitRange(padRows, begin, end); + uint32_t row = validRows + begin; + uint32_t padEnd = validRows + end; + while (row < padEnd) { + uint32_t curRows = MinU32(tileRows_, padEnd - row); + WriteInvalidTile(row, curRows); + row += curRows; + } + } + + uint32_t numRows_{0}; + uint32_t actualNumReqs_{0}; + uint32_t ropeRows_{0}; + uint32_t ropeDim_{0}; + uint32_t kvBlockTableStride_{0}; + uint32_t kvBlockSize_{0}; + uint32_t slotMappingFormat_{0}; + uint32_t cmpRatio_{1}; + uint32_t tileRows_{1}; + uint32_t ropeRowBytes_{0}; + uint32_t ropeRowBytesAligned_{0}; + uint32_t slotCols_{1}; + uint32_t reqTableBytes_{0}; + uint32_t ropeDimAligned_{0}; + uint32_t ropePadElems_{0}; + uint32_t slotTileBytes_{0}; + uint32_t blockTableTileBytes_{0}; + + TBuf prefixBuf_; + TBuf startPosBuf_; + TBuf cuSeqlensBuf_; + TBuf blockTableBuf_; + TBuf slotBuf_; + TQue cosQueue_; + TQue sinQueue_; + + GlobalTensor ropeCosGm_; + GlobalTensor ropeSinGm_; + GlobalTensor compressCosGm_; + GlobalTensor compressSinGm_; + GlobalTensor cuSeqlensGm_; + GlobalTensor startPosGm_; + GlobalTensor kvBlockTableGm_; + GlobalTensor slotMappingGm_; +}; +} // namespace CompressorMetadata + +#endif diff --git a/csrc/ascend/attention/fused_gdn_gating/fused_gdn_gating_torch_adpt.h b/csrc/ascend/attention/fused_gdn_gating/fused_gdn_gating_torch_adpt.h new file mode 100644 index 000000000..2e3e0956d --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/fused_gdn_gating_torch_adpt.h @@ -0,0 +1,59 @@ +/* + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project + */ + +#ifndef FUSED_GDN_GATING_TORCH_ADPT_H +#define FUSED_GDN_GATING_TORCH_ADPT_H + +#include + +namespace vllm_fl { + +std::tuple npu_fused_gdn_gating( + const at::Tensor& A_log, + const at::Tensor& a, + const at::Tensor& b, + const at::Tensor& dt_bias, + double beta = 1.0, + double threshold = 20.0) +{ + TORCH_CHECK(A_log.dim() == 1, "A_log should be 1-D [num_heads], got ", A_log.dim(), "D"); + TORCH_CHECK(dt_bias.dim() == 1, "dt_bias should be 1-D [num_heads], got ", dt_bias.dim(), "D"); + TORCH_CHECK(a.dim() == 2, "a should be 2-D [batch, num_heads], got ", a.dim(), "D"); + TORCH_CHECK(b.dim() == 2, "b should be 2-D [batch, num_heads], got ", b.dim(), "D"); + TORCH_CHECK(b.size(0) == a.size(0) && b.size(1) == a.size(1), + "a and b must have the same shape, got a=", a.sizes(), " b=", b.sizes()); + TORCH_CHECK(a.scalar_type() == b.scalar_type(), + "a and b must have the same dtype, got a=", a.scalar_type(), + " b=", b.scalar_type()); + TORCH_CHECK(A_log.scalar_type() == dt_bias.scalar_type(), + "A_log and dt_bias must have the same dtype, got A_log=", + A_log.scalar_type(), " dt_bias=", dt_bias.scalar_type()); + TORCH_CHECK(a.size(1) == A_log.size(0), + "a second dim (num_heads) must equal A_log first dim, got a.size(1)=", + a.size(1), " A_log.size(0)=", A_log.size(0)); + + int64_t batch = a.size(0); + int64_t num_heads = a.size(1); + + at::Tensor g = at::empty({1, batch, num_heads}, + a.options().dtype(c10::kFloat)); + at::Tensor beta_output = at::empty({1, batch, num_heads}, b.options()); + + float beta_val = static_cast(beta); + float threshold_val = static_cast(threshold); + + EXEC_NPU_CMD(aclnnFusedGdnGating, + A_log, a, b, dt_bias, + beta_val, + threshold_val, + g, beta_output); + + return std::make_tuple(g, beta_output); +} + +} // namespace vllm_fl + +#endif // FUSED_GDN_GATING_TORCH_ADPT_H diff --git a/csrc/ascend/attention/fused_gdn_gating/op_host/CMakeLists.txt b/csrc/ascend/attention/fused_gdn_gating/op_host/CMakeLists.txt new file mode 100644 index 000000000..ad7a8889b --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/op_host/CMakeLists.txt @@ -0,0 +1,23 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. + +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnExc PRIVATE + fused_gdn_gating_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME FusedGdnGating + OPTIONS --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE fused_gdn_gating ACLNNTYPE aclnn_exclude) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_def.cpp b/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_def.cpp new file mode 100644 index 000000000..e1748abe5 --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_def.cpp @@ -0,0 +1,68 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project + */ + +/*! + * \file fused_gdn_gating_def.cpp + * \brief OpDef registration for FusedGdnGating. + */ + +#include "register/op_def_registry.h" + +namespace ops { + +class FusedGdnGating : public OpDef { +public: + explicit FusedGdnGating(const char *name) : OpDef(name) + { + this->Input("a_log") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("a") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("b") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("dt_bias") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("g") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("beta_output") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + this->Attr("beta").AttrType(OPTIONAL).Float(1.0f); + this->Attr("threshold").AttrType(OPTIONAL).Float(20.0f); + + OpAICoreConfig aicConfig; + aicConfig.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .ExtendCfgInfo("softsync.flag", "true"); + this->AICore().AddConfig("ascend910b", aicConfig); + this->AICore().AddConfig("ascend910_93", aicConfig); + } +}; + +OP_ADD(FusedGdnGating); + +} // namespace ops diff --git a/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_infershape.cpp b/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_infershape.cpp new file mode 100644 index 000000000..18d67fb6a --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_infershape.cpp @@ -0,0 +1,78 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project + */ + +/*! + * \file fused_gdn_gating_infershape.cpp + * \brief Shape and data-type inference for FusedGdnGating. + */ + +#include "exe_graph/runtime/infer_shape_context.h" +#include "exe_graph/runtime/shape.h" +#include "exe_graph/runtime/storage_shape.h" +#include "register/op_impl_registry.h" + +using namespace gert; + +namespace ops { + +namespace { + +constexpr size_t INPUT_A_INDEX = 1; +constexpr size_t OUTPUT_G_INDEX = 0; +constexpr size_t OUTPUT_BETA_INDEX = 1; +constexpr size_t OUTPUT_DIM_NUM = 3; +constexpr int64_t OUTPUT_SEQ_LEN = 1; + +} // namespace + +static ge::graphStatus InferShapeFusedGdnGating(InferShapeContext *context) +{ + if (context == nullptr) { + return ge::GRAPH_FAILED; + } + + auto shapeA = context->GetInputShape(INPUT_A_INDEX); + auto shapeG = context->GetOutputShape(OUTPUT_G_INDEX); + auto shapeBeta = context->GetOutputShape(OUTPUT_BETA_INDEX); + if (shapeA == nullptr || shapeG == nullptr || shapeBeta == nullptr) { + return ge::GRAPH_FAILED; + } + if (shapeA->GetDimNum() < 2) { + return ge::GRAPH_FAILED; + } + + const int64_t batch = shapeA->GetDim(0); + const int64_t numHeads = shapeA->GetDim(1); + + shapeG->SetDimNum(OUTPUT_DIM_NUM); + shapeG->SetDim(0, OUTPUT_SEQ_LEN); + shapeG->SetDim(1, batch); + shapeG->SetDim(2, numHeads); + + shapeBeta->SetDimNum(OUTPUT_DIM_NUM); + shapeBeta->SetDim(0, OUTPUT_SEQ_LEN); + shapeBeta->SetDim(1, batch); + shapeBeta->SetDim(2, numHeads); + + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataTypeFusedGdnGating(gert::InferDataTypeContext *context) +{ + if (context == nullptr) { + return ge::GRAPH_FAILED; + } + ge::DataType inputADtype = context->GetInputDataType(INPUT_A_INDEX); + context->SetOutputDataType(OUTPUT_G_INDEX, ge::DT_FLOAT); + context->SetOutputDataType(OUTPUT_BETA_INDEX, inputADtype); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(FusedGdnGating) + .InferShape(InferShapeFusedGdnGating) + .InferDataType(InferDataTypeFusedGdnGating); + +} // namespace ops diff --git a/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_tiling.cpp b/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_tiling.cpp new file mode 100644 index 000000000..8f90c44bc --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_tiling.cpp @@ -0,0 +1,176 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project + */ + +/*! + * \file fused_gdn_gating_tiling.cpp + * \brief Tiling implementation for FusedGdnGating. + */ + +#include "fused_gdn_gating_tiling.h" +#include "fused_gdn_gating_tiling_utils.h" + +#include "register/op_impl_registry.h" +#include "securec.h" +#include "tiling/platform/platform_ascendc.h" +#include "tiling/tiling_api.h" + +#include "../op_kernel/fused_gdn_gating_tiling_data.h" + +using namespace FusedGdnGating; + +namespace optiling { + +namespace { + +constexpr uint64_t TILING_KEY_BF16 = 1; +constexpr uint64_t TILING_KEY_FP16 = 2; +constexpr uint64_t TILING_KEY_PARAM_BF16_OFFSET = 2; +constexpr uint64_t TILING_KEY_PARAM_FP16_OFFSET = 4; +constexpr size_t INPUT_INDEX_A_LOG = 0; +constexpr size_t INPUT_INDEX_A = 1; +constexpr size_t INPUT_INDEX_DT_BIAS = 3; + +} // namespace + +ge::graphStatus FusedGdnGatingTilingFunc(gert::TilingContext *context) +{ + if (context == nullptr) { + return ge::GRAPH_FAILED; + } + + auto platformInfoPtr = context->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + return ge::GRAPH_FAILED; + } + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + uint64_t ubSize = 0; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + if (aivNum == 0) { + aivNum = 1; + } + + auto *shapeA = context->GetInputShape(INPUT_INDEX_A); + if (shapeA == nullptr) { + return ge::GRAPH_FAILED; + } + const auto &storageShape = shapeA->GetStorageShape(); + if (storageShape.GetDimNum() < 2) { + return ge::GRAPH_FAILED; + } + int64_t numBatches = storageShape.GetDim(0); + int64_t numHeads = storageShape.GetDim(1); + if (numBatches <= 0 || numHeads <= 0) { + return ge::GRAPH_FAILED; + } + + float beta = 1.0f; + float threshold = 20.0f; + auto *attrs = context->GetAttrs(); + if (attrs != nullptr) { + const float *betaAttr = attrs->GetAttrPointer(0); + if (betaAttr != nullptr) { beta = *betaAttr; } + const float *thresholdAttr = attrs->GetAttrPointer(1); + if (thresholdAttr != nullptr) { threshold = *thresholdAttr; } + } + + auto *aDesc = context->GetInputDesc(INPUT_INDEX_A); + auto *aLogDesc = context->GetInputDesc(INPUT_INDEX_A_LOG); + auto *dtBiasDesc = context->GetInputDesc(INPUT_INDEX_DT_BIAS); + if (aDesc == nullptr || aLogDesc == nullptr || dtBiasDesc == nullptr) { + return ge::GRAPH_FAILED; + } + ge::DataType aDtype = aDesc->GetDataType(); + ge::DataType aLogDtype = aLogDesc->GetDataType(); + ge::DataType dtBiasDtype = dtBiasDesc->GetDataType(); + if (aLogDtype != dtBiasDtype) { + return ge::GRAPH_FAILED; + } + uint64_t tilingKey = TILING_KEY_BF16; + if (aDtype == ge::DT_FLOAT16) { + tilingKey = TILING_KEY_FP16; + } + if (aLogDtype == ge::DT_BF16) { + tilingKey += TILING_KEY_PARAM_BF16_OFFSET; + } else if (aLogDtype == ge::DT_FLOAT16) { + tilingKey += TILING_KEY_PARAM_FP16_OFFSET; + } + + uint32_t blockDim = static_cast(numBatches); + if (blockDim > aivNum) { + blockDim = aivNum; + } + + uint32_t numHeadsU32 = static_cast(numHeads); + uint32_t numBatchesU32 = static_cast(numBatches); + uint32_t rowsConservative = ComputeRowsPerIter(numHeadsU32, ubSize); + uint32_t rowsPerIter = rowsConservative; + + // Block utilization: ensure enough chunks for all AIV cores. + { + uint32_t totalChunksForRPI = (numBatchesU32 + rowsPerIter - 1) / rowsPerIter; + if (numBatchesU32 <= rowsPerIter || totalChunksForRPI < blockDim) { + uint32_t maxRPI = numBatchesU32 / blockDim; + if (maxRPI < 1) { maxRPI = 1; } + if (maxRPI >= 128) { rowsPerIter = 128; } + else if (maxRPI >= 64) { rowsPerIter = 64; } + else if (maxRPI >= 32) { rowsPerIter = 32; } + else if (maxRPI >= 16) { rowsPerIter = 16; } + else if (maxRPI >= 8) { rowsPerIter = 8; } + else if (maxRPI >= 4) { rowsPerIter = 4; } + else if (maxRPI >= 2) { rowsPerIter = 2; } + else { rowsPerIter = 1; } + if (rowsPerIter > rowsConservative) { rowsPerIter = rowsConservative; } + } + } + + const bool bulkDmaBatchOk = (numBatchesU32 > blockDim * rowsPerIter); + bool useBulkDma = bulkDmaBatchOk && CanUseBulkDma(numHeadsU32, rowsPerIter); + + FusedGdnGatingTilingData td{}; + td.numHeads = numHeadsU32; + td.numBatches = numBatchesU32; + td.rowsPerIter = rowsPerIter; + td.useBulkDma = useBulkDma ? 1u : 0u; + td.beta = beta; + td.threshold = threshold; + + const size_t tilingSize = sizeof(FusedGdnGatingTilingData); + auto *rawTilingData = context->GetRawTilingData(); + if (rawTilingData == nullptr || rawTilingData->GetCapacity() < tilingSize) { + return ge::GRAPH_FAILED; + } + errno_t rc = memcpy_s(rawTilingData->GetData(), rawTilingData->GetCapacity(), + &td, tilingSize); + if (rc != EOK) { + return ge::GRAPH_FAILED; + } + rawTilingData->SetDataSize(tilingSize); + + context->SetBlockDim(blockDim); + context->SetTilingKey(tilingKey); + + // No GM workspace needed. + size_t *workspaces = context->GetWorkspaceSizes(1); + if (workspaces != nullptr) { + workspaces[0] = 0; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingPrepareForFusedGdnGating(gert::TilingParseContext *context) +{ + // Required by CANN tiling framework for "_pattern" registration. + (void)context; + return ge::GRAPH_SUCCESS; +} + +} // namespace optiling + +IMPL_OP_OPTILING(FusedGdnGating) + .Tiling(optiling::FusedGdnGatingTilingFunc) + .TilingParse(optiling::TilingPrepareForFusedGdnGating); diff --git a/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_tiling.h b/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_tiling.h new file mode 100644 index 000000000..5e3e1af3e --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_tiling.h @@ -0,0 +1,29 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project + */ + +/*! + * \file fused_gdn_gating_tiling.h + * \brief Function-style tiling declaration for FusedGdnGating. + */ + +#ifndef FUSED_GDN_GATING_TILING_H +#define FUSED_GDN_GATING_TILING_H + +#include +#include +#include + +namespace optiling { + +// Required by CANN tiling framework. +struct FusedGdnGatingCompileInfo {}; + +ge::graphStatus FusedGdnGatingTilingFunc(gert::TilingContext *context); +ge::graphStatus TilingPrepareForFusedGdnGating(gert::TilingParseContext *context); + +} // namespace optiling + +#endif // FUSED_GDN_GATING_TILING_H diff --git a/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_tiling_utils.h b/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_tiling_utils.h new file mode 100644 index 000000000..bb76eb4fd --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/op_host/fused_gdn_gating_tiling_utils.h @@ -0,0 +1,108 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project + */ + +/*! + * \file fused_gdn_gating_tiling_utils.h + * \brief rowsPerIter and Bulk DMA helper functions. + */ + +#ifndef FUSED_GDN_GATING_TILING_UTILS_H +#define FUSED_GDN_GATING_TILING_UTILS_H + +#include + +namespace FusedGdnGating { + +// NPU hardware constants. +constexpr uint32_t VECTOR_BYTES_PER_ITER = 256; +constexpr uint32_t DATACOPY_MIN_BYTES = 32; +constexpr uint32_t BF16_PER_BLOCK = DATACOPY_MIN_BYTES / 2; // 16 +constexpr uint32_t MASK_ALIGN_ELEMS = 64; + +/// Align count to vector unit width (256 bytes) for given dtype size. +inline uint32_t AlignCountToVectorBytes(uint32_t count, uint32_t dtypeSize) +{ + uint32_t elemsPerIter = VECTOR_BYTES_PER_ITER / dtypeSize; + return ((count + elemsPerIter - 1) / elemsPerIter) * elemsPerIter; +} + +/// Check if Bulk DMA is viable: (R * nh) % 64 == 0, nh % 16 == 0. +inline bool CanUseBulkDma(uint32_t numHeads, uint32_t rowsPerIter) +{ + // Condition 1: (rows_per_iter * num_heads) % 64 == 0 + // fp32 vector unit processes 64 elements per repeat; the bulk operation + // must align with this granularity to avoid tail handling. + constexpr uint32_t fp32VecElems = VECTOR_BYTES_PER_ITER / 4; + if ((rowsPerIter * numHeads) % fp32VecElems != 0) { + return false; + } + + // Condition 2: num_heads % 16 == 0 + // DMA minimum transfer size is 32 bytes; for bf16/fp16 (2 bytes per element), + // this equals 16 elements. If num_heads is not a multiple of 16, the last + // few elements of each row require separate handling, negating the bulk benefit. + constexpr uint32_t bf16BlockElems = DATACOPY_MIN_BYTES / 2; + if (numHeads % bf16BlockElems != 0) { + return false; + } + + return true; +} + +/*! + * \brief Compute optimal rows_per_iter from UB budget. + * + * UB breakdown matches kernel Init(): 3 single-row fp32 constants + * + 2 multi-row fp32 constants (R * ubDim * 4 each) + * + 3 half + 6 fp32 per-row buffers (scaled by R). + * ubDim = ceil(numHeads / 16) * 16 (matching kernel DMA_ALIGN_ELEMS). + * Result clamped to power-of-2, max 128. + */ +inline uint32_t ComputeRowsPerIter(uint32_t numHeads, uint64_t ubBudget, + uint32_t ubDim = 0) +{ + if (ubDim == 0) { + // Match the kernel's fp32 compute/mask alignment. + ubDim = ((numHeads + MASK_ALIGN_ELEMS - 1) / MASK_ALIGN_ELEMS) * MASK_ALIGN_ELEMS; + } + uint32_t maskUbDim = ubDim; + + // 2 parameter input queues + 2 fp32 constant buffers, each 1 row. + // Use fp32 for the parameter queues as a conservative upper bound. + uint32_t sharedBytes = 4 * ubDim * static_cast(sizeof(float)); + + // Multi-row constant buffers (precomputed once, scaled by R): + // dtBiasMultiBuf_ + negExpMultiBuf_: 2 fp32 buffers. + uint32_t constPerRowBytes = 2 * ubDim * static_cast(sizeof(float)); + + // Per-row (per-chunk): 3 bf16/fp16 buffers + 5 fp32 buffers + 1 uint8 mask buffer. + uint32_t perRowBytes = 3 * ubDim * static_cast(sizeof(int16_t)) // a, b, betaOut + + 5 * ubDim * static_cast(sizeof(float)) // g, x, betaX, tmp, betaFp32 + + 1 * maskUbDim * static_cast(sizeof(uint8_t)); // threshold mask + + if (perRowBytes == 0) { + return 1; + } + + uint32_t maxRows = 1; + if (ubBudget > sharedBytes) { + maxRows = static_cast((ubBudget - sharedBytes) / (perRowBytes + constPerRowBytes)); + } + + // Round down to nearest power of 2 (128, 64, 32, ..., 1). + if (maxRows >= 128) { return 128; } + if (maxRows >= 64) { return 64; } + if (maxRows >= 32) { return 32; } + if (maxRows >= 16) { return 16; } + if (maxRows >= 8) { return 8; } + if (maxRows >= 4) { return 4; } + if (maxRows >= 2) { return 2; } + return 1; +} + +} // namespace FusedGdnGating + +#endif // FUSED_GDN_GATING_TILING_UTILS_H diff --git a/csrc/ascend/attention/fused_gdn_gating/op_host/op_api/aclnn_fused_gdn_gating.cpp b/csrc/ascend/attention/fused_gdn_gating/op_host/op_api/aclnn_fused_gdn_gating.cpp new file mode 100644 index 000000000..307fb5ce0 --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/op_host/op_api/aclnn_fused_gdn_gating.cpp @@ -0,0 +1,143 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project + */ + +/*! + * \file aclnn_fused_gdn_gating.cpp + * \brief ACLNN C-API (GetWorkspaceSize + Execute). + */ + +#include +#include "aclnn_fused_gdn_gating.h" +#include "fused_gdn_gating.h" + +#include "securec.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "opdev/common_types.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/platform.h" + +#include "aclnn_kernels/contiguous.h" + +using namespace op; + +#ifdef __cplusplus +extern "C" { +#endif + +namespace { + +struct FusedGdnGatingParams { + const aclTensor *aLog{nullptr}; + const aclTensor *a{nullptr}; + const aclTensor *b{nullptr}; + const aclTensor *dtBias{nullptr}; + float beta{1.0f}; + float threshold{20.0f}; + aclTensor *g{nullptr}; + aclTensor *betaOutput{nullptr}; +}; + +static const std::initializer_list AB_TYPE_SUPPORT_LIST = + {op::DataType::DT_BF16, op::DataType::DT_FLOAT16}; +static const std::initializer_list FP32_TYPE_SUPPORT_LIST = + {op::DataType::DT_FLOAT}; +static const std::initializer_list PARAM_TYPE_SUPPORT_LIST = + {op::DataType::DT_FLOAT, op::DataType::DT_BF16, op::DataType::DT_FLOAT16}; + +static inline bool CheckNotNull(const FusedGdnGatingParams ¶ms) +{ + OP_CHECK_NULL(params.aLog, return false); + OP_CHECK_NULL(params.a, return false); + OP_CHECK_NULL(params.b, return false); + OP_CHECK_NULL(params.dtBias, return false); + OP_CHECK_NULL(params.g, return false); + OP_CHECK_NULL(params.betaOutput, return false); + return true; +} + +static inline bool CheckDtype(const FusedGdnGatingParams ¶ms) +{ + OP_CHECK_DTYPE_NOT_SUPPORT(params.aLog, PARAM_TYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(params.dtBias, PARAM_TYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(params.a, AB_TYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(params.b, AB_TYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(params.g, FP32_TYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(params.betaOutput, AB_TYPE_SUPPORT_LIST, return false); + OP_CHECK(params.a->GetDataType() == params.b->GetDataType(), + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "a and b must have the same dtype."), + return false); + OP_CHECK(params.aLog->GetDataType() == params.dtBias->GetDataType(), + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "aLog and dtBias must have the same dtype."), + return false); + OP_CHECK(params.betaOutput->GetDataType() == params.b->GetDataType(), + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "betaOutput and b must have the same dtype."), + return false); + return true; +} + +static aclnnStatus CheckParams(const FusedGdnGatingParams ¶ms) +{ + CHECK_RET(CheckNotNull(params), ACLNN_ERR_PARAM_NULLPTR); + CHECK_RET(CheckDtype(params), ACLNN_ERR_PARAM_INVALID); + return ACLNN_SUCCESS; +} + +} // namespace + +aclnnStatus aclnnFusedGdnGatingGetWorkspaceSize( + const aclTensor *aLog, const aclTensor *a, const aclTensor *b, + const aclTensor *dtBias, float beta, float threshold, + aclTensor *g, aclTensor *betaOutput, + uint64_t *workspaceSize, aclOpExecutor **executor) +{ + L2_DFX_PHASE_1(aclnnFusedGdnGating, + DFX_IN(aLog, a, b, dtBias, beta, threshold), + DFX_OUT(g, betaOutput)); + + auto uniqueExecutor = CREATE_EXECUTOR(); + CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + + FusedGdnGatingParams params{aLog, a, b, dtBias, beta, threshold, g, betaOutput}; + CHECK_RET(CheckParams(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + + // Bring inputs to a contiguous form that the kernel expects. + auto aLogContig = l0op::Contiguous(aLog, uniqueExecutor.get()); + auto aContig = l0op::Contiguous(a, uniqueExecutor.get()); + auto bContig = l0op::Contiguous(b, uniqueExecutor.get()); + auto dtBiasContig = l0op::Contiguous(dtBias, uniqueExecutor.get()); + CHECK_RET(aLogContig != nullptr, ACLNN_ERR_INNER_NULLPTR); + CHECK_RET(aContig != nullptr, ACLNN_ERR_INNER_NULLPTR); + CHECK_RET(bContig != nullptr, ACLNN_ERR_INNER_NULLPTR); + CHECK_RET(dtBiasContig != nullptr, ACLNN_ERR_INNER_NULLPTR); + + auto result = l0op::FusedGdnGating(aLogContig, aContig, bContig, dtBiasContig, + beta, threshold, uniqueExecutor.get()); + CHECK_RET(result.g != nullptr && result.beta_output != nullptr, + ACLNN_ERR_INNER_NULLPTR); + + // Copy kernel results into the caller-provided output tensors. + auto vcG = l0op::ViewCopy(result.g, g, uniqueExecutor.get()); + CHECK_RET(vcG != nullptr, ACLNN_ERR_INNER_NULLPTR); + auto vcBeta = l0op::ViewCopy(result.beta_output, betaOutput, uniqueExecutor.get()); + CHECK_RET(vcBeta != nullptr, ACLNN_ERR_INNER_NULLPTR); + + *workspaceSize = uniqueExecutor->GetWorkspaceSize(); + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; +} + +aclnnStatus aclnnFusedGdnGating(void *workspace, uint64_t workspaceSize, + aclOpExecutor *executor, aclrtStream stream) +{ + L2_DFX_PHASE_2(aclnnFusedGdnGating); + return CommonOpExecutorRun(workspace, workspaceSize, executor, stream); +} + +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/attention/fused_gdn_gating/op_host/op_api/aclnn_fused_gdn_gating.h b/csrc/ascend/attention/fused_gdn_gating/op_host/op_api/aclnn_fused_gdn_gating.h new file mode 100644 index 000000000..74d935066 --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/op_host/op_api/aclnn_fused_gdn_gating.h @@ -0,0 +1,51 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project + */ + +/*! + * \file aclnn_fused_gdn_gating.h + * \brief ACLNN C-API for FusedGdnGating. + */ + +#ifndef OP_API_ACLNN_FUSED_GDN_GATING_H +#define OP_API_ACLNN_FUSED_GDN_GATING_H + +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @brief FusedGdnGating phase-1: compute required workspace size. + * @param [in] aLog : A_log, [num_heads], dtype fp32/bf16/fp16. + * @param [in] a : a, [batch, num_heads], dtype bf16/fp16. + * @param [in] b : b, [batch, num_heads], dtype bf16/fp16. + * @param [in] dtBias : dt_bias, [num_heads], same dtype as aLog. + * @param [in] beta : softplus beta (default 1.0). + * @param [in] threshold : softplus threshold (default 20.0). + * @param [out] g : output gate, [1, batch, num_heads], dtype fp32. + * @param [out] betaOutput : sigmoid(b), [1, batch, num_heads], same dtype as a/b. + * @param [out] workspaceSize: required workspace bytes on device. + * @param [out] executor : op executor handle. + */ +__attribute__((visibility("default"))) aclnnStatus aclnnFusedGdnGatingGetWorkspaceSize( + const aclTensor *aLog, const aclTensor *a, const aclTensor *b, + const aclTensor *dtBias, float beta, float threshold, + aclTensor *g, aclTensor *betaOutput, + uint64_t *workspaceSize, aclOpExecutor **executor); + +/** + * @brief FusedGdnGating phase-2: launch the kernel. + */ +__attribute__((visibility("default"))) aclnnStatus aclnnFusedGdnGating( + void *workspace, uint64_t workspaceSize, + aclOpExecutor *executor, aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif // OP_API_ACLNN_FUSED_GDN_GATING_H diff --git a/csrc/ascend/attention/fused_gdn_gating/op_host/op_api/fused_gdn_gating.cpp b/csrc/ascend/attention/fused_gdn_gating/op_host/op_api/fused_gdn_gating.cpp new file mode 100644 index 000000000..e38373f05 --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/op_host/op_api/fused_gdn_gating.cpp @@ -0,0 +1,65 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project + */ + +/*! + * \file fused_gdn_gating.cpp + * \brief L0-level API for FusedGdnGating. + */ + +#include "fused_gdn_gating.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "opdev/make_op_executor.h" +#include "opdev/op_def.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/shape_utils.h" + +using namespace op; + +namespace l0op { + +OP_TYPE_REGISTER(FusedGdnGating); + +static constexpr FusedGdnGatingOutput kNullOutput{nullptr, nullptr}; + +FusedGdnGatingOutput FusedGdnGating(const aclTensor *aLog, const aclTensor *a, + const aclTensor *b, const aclTensor *dtBias, + float beta, float threshold, + aclOpExecutor *executor) +{ + L0_DFX(FusedGdnGating, aLog, a, b, dtBias, beta, threshold); + + const DataType betaDtype = b->GetDataType(); + const Format format = Format::FORMAT_ND; + + auto g = executor->AllocTensor(DataType::DT_FLOAT, format, format); + OP_CHECK(g != nullptr, OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "g AllocTensor failed."), + return kNullOutput); + + auto betaOutput = executor->AllocTensor(betaDtype, format, format); + OP_CHECK(betaOutput != nullptr, + OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "beta_output AllocTensor failed."), + return kNullOutput); + + auto ret = INFER_SHAPE(FusedGdnGating, + OP_INPUT(aLog, a, b, dtBias), + OP_OUTPUT(g, betaOutput), + OP_ATTR(beta, threshold)); + OP_CHECK_INFERSHAPE(ret != ACLNN_SUCCESS, return kNullOutput, + "FusedGdnGating InferShape failed."); + + ret = ADD_TO_LAUNCHER_LIST_AICORE(FusedGdnGating, + OP_INPUT(aLog, a, b, dtBias), + OP_OUTPUT(g, betaOutput), + OP_ATTR(beta, threshold)); + OP_CHECK_ADD_TO_LAUNCHER_LIST_AICORE(ret != ACLNN_SUCCESS, return kNullOutput, + "FusedGdnGating ADD_TO_LAUNCHER_LIST_AICORE failed."); + + return FusedGdnGatingOutput{g, betaOutput}; +} + +} // namespace l0op diff --git a/csrc/ascend/attention/fused_gdn_gating/op_host/op_api/fused_gdn_gating.h b/csrc/ascend/attention/fused_gdn_gating/op_host/op_api/fused_gdn_gating.h new file mode 100644 index 000000000..2032151bf --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/op_host/op_api/fused_gdn_gating.h @@ -0,0 +1,27 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project + */ + +#ifndef PTA_NPU_OP_API_FUSED_GDN_GATING_H +#define PTA_NPU_OP_API_FUSED_GDN_GATING_H + +#include "opdev/op_executor.h" +#include "opdev/make_op_executor.h" + +namespace l0op { + +struct FusedGdnGatingOutput { + const aclTensor *g; + const aclTensor *beta_output; +}; + +FusedGdnGatingOutput FusedGdnGating(const aclTensor *aLog, const aclTensor *a, + const aclTensor *b, const aclTensor *dtBias, + float beta, float threshold, + aclOpExecutor *executor); + +} // namespace l0op + +#endif // PTA_NPU_OP_API_FUSED_GDN_GATING_H diff --git a/csrc/ascend/attention/fused_gdn_gating/op_kernel/fused_gdn_gating.cpp b/csrc/ascend/attention/fused_gdn_gating/op_kernel/fused_gdn_gating.cpp new file mode 100644 index 000000000..e9d483c5a --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/op_kernel/fused_gdn_gating.cpp @@ -0,0 +1,54 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project + */ + +/*! + * \file fused_gdn_gating.cpp + * \brief AscendC kernel entry for FusedGdnGating. + */ + +#include "fused_gdn_gating.h" +#include "fused_gdn_gating_tiling_data.h" + +using namespace AscendC; +using namespace FusedGdnGating; + +extern "C" __global__ __aicore__ void +fused_gdn_gating(GM_ADDR a_log, GM_ADDR a, GM_ADDR b, GM_ADDR dt_bias, + GM_ADDR g, GM_ADDR beta_output, + GM_ADDR workspace, GM_ADDR tiling_gm) +{ + REGISTER_TILING_DEFAULT(FusedGdnGatingTilingData); + GET_TILING_DATA(tilingData, tiling_gm); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); + + TPipe pipe; + + if (TILING_KEY_IS(1)) { + KernelFusedGdnGating op; + op.Init(a_log, a, b, dt_bias, g, beta_output, &tilingData, &pipe); + op.Process(); + } else if (TILING_KEY_IS(2)) { + KernelFusedGdnGating op; + op.Init(a_log, a, b, dt_bias, g, beta_output, &tilingData, &pipe); + op.Process(); + } else if (TILING_KEY_IS(3)) { + KernelFusedGdnGating op; + op.Init(a_log, a, b, dt_bias, g, beta_output, &tilingData, &pipe); + op.Process(); + } else if (TILING_KEY_IS(4)) { + KernelFusedGdnGating op; + op.Init(a_log, a, b, dt_bias, g, beta_output, &tilingData, &pipe); + op.Process(); + } else if (TILING_KEY_IS(5)) { + KernelFusedGdnGating op; + op.Init(a_log, a, b, dt_bias, g, beta_output, &tilingData, &pipe); + op.Process(); + } else if (TILING_KEY_IS(6)) { + KernelFusedGdnGating op; + op.Init(a_log, a, b, dt_bias, g, beta_output, &tilingData, &pipe); + op.Process(); + } +} diff --git a/csrc/ascend/attention/fused_gdn_gating/op_kernel/fused_gdn_gating.h b/csrc/ascend/attention/fused_gdn_gating/op_kernel/fused_gdn_gating.h new file mode 100644 index 000000000..108d035a1 --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/op_kernel/fused_gdn_gating.h @@ -0,0 +1,396 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project + */ + +/*! + * \file fused_gdn_gating.h + * \brief AscendC kernel for fused GDN gating. + * + * Per-row math: + * g = -exp(A_log) * softplus(cast(a,fp32) + dt_bias, beta, threshold) + * beta_output = sigmoid(cast(b, fp32)) -> cast back to InDtype + */ + +#ifndef FUSED_GDN_GATING_KERNEL_H +#define FUSED_GDN_GATING_KERNEL_H + +#include + +#include "kernel_operator.h" +#include "fused_gdn_gating_tiling_data.h" + +namespace FusedGdnGating { + +using namespace AscendC; + +// 32-byte alignment requirement for DataCopy on NPU. +constexpr uint32_t BYTES_PER_BLOCK = 32; +constexpr uint32_t BF16_PER_BLOCK = BYTES_PER_BLOCK / sizeof(int16_t); // 16 +constexpr uint32_t FP32_PER_BLOCK = BYTES_PER_BLOCK / sizeof(float); // 8 +constexpr uint32_t MASK_ALIGN_ELEMS = 64; + +// DMA-friendly alignment: 16 elements = 32 bytes = 1 DMA block. +// Vector ops use count=numHeads_ with partial-iteration masking, +// so there is no minimum-count constraint. +constexpr uint32_t DMA_ALIGN_ELEMS = BYTES_PER_BLOCK / sizeof(int16_t); // 16 + +template +__aicore__ inline T CeilDiv(T a, T b) { return (a + b - 1) / b; } + +template +__aicore__ inline T AlignUp(T a, T b) { return CeilDiv(a, b) * b; } + +template +class KernelFusedGdnGating { +public: + __aicore__ inline KernelFusedGdnGating() {} + + /*! + * \brief Init kernel with GM addresses and tiling data. + * + * Argument order matches OpDef: aLogGm, aGm, bGm, dtBiasGm, gGm, betaOutputGm. + */ + __aicore__ inline void Init(GM_ADDR aLogGm, GM_ADDR aGm, GM_ADDR bGm, GM_ADDR dtBiasGm, + GM_ADDR gGm, GM_ADDR betaOutputGm, + const FusedGdnGatingTilingData *tiling, TPipe *pipe) + { + pipe_ = pipe; + numHeads_ = tiling->numHeads; + numBatches_ = tiling->numBatches; + rowsPerIter_ = tiling->rowsPerIter; + useBulkDma_ = (tiling->useBulkDma != 0); + beta_ = tiling->beta; + threshold_ = tiling->threshold; + + // Aligned dimensions for UB tensors. + alignedHeadsHalf_ = AlignUp(numHeads_, MASK_ALIGN_ELEMS); + alignedHeadsFloat_ = AlignUp(numHeads_, MASK_ALIGN_ELEMS); + alignedHeadsMask_ = alignedHeadsFloat_; + constexpr uint32_t paramAlignElems = BYTES_PER_BLOCK / sizeof(ParamDtype); + alignedHeadsParam_ = AlignUp(numHeads_, paramAlignElems); + + aLogGm_.SetGlobalBuffer(reinterpret_cast<__gm__ ParamDtype *>(aLogGm), numHeads_); + dtBiasGm_.SetGlobalBuffer(reinterpret_cast<__gm__ ParamDtype *>(dtBiasGm), numHeads_); + aGm_.SetGlobalBuffer(reinterpret_cast<__gm__ InDtype *>(aGm), + static_cast(numBatches_) * numHeads_); + bGm_.SetGlobalBuffer(reinterpret_cast<__gm__ InDtype *>(bGm), + static_cast(numBatches_) * numHeads_); + gGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(gGm), + static_cast(numBatches_) * numHeads_); + betaGm_.SetGlobalBuffer(reinterpret_cast<__gm__ InDtype *>(betaOutputGm), + static_cast(numBatches_) * numHeads_); + + // I/O queues (depth=1). + pipe_->InitBuffer(aInQue_, 1, rowsPerIter_ * alignedHeadsHalf_ * sizeof(InDtype)); + pipe_->InitBuffer(bInQue_, 1, rowsPerIter_ * alignedHeadsHalf_ * sizeof(InDtype)); + pipe_->InitBuffer(gOutQue_, 1, rowsPerIter_ * alignedHeadsFloat_ * sizeof(float)); + pipe_->InitBuffer(betaOutQue_, 1, rowsPerIter_ * alignedHeadsHalf_ * sizeof(InDtype)); + + // Constant queues (single-row). + pipe_->InitBuffer(aLogInQue_, 1, 1 * alignedHeadsParam_ * sizeof(ParamDtype)); + pipe_->InitBuffer(dtBiasInQue_, 1, 1 * alignedHeadsParam_ * sizeof(ParamDtype)); + pipe_->InitBuffer(negExpInQue_, 1, 1 * alignedHeadsFloat_ * sizeof(float)); + pipe_->InitBuffer(dtBiasPreloadQue_, 1, 1 * alignedHeadsFloat_ * sizeof(float)); + + // Multi-row constants: dt_bias and neg_exp(A_log) replicated R times. + // Only allocated for R > 1; single-row kernels use per-row fallback. + if (rowsPerIter_ > 1) { + pipe_->InitBuffer(dtBiasMultiBuf_, rowsPerIter_ * alignedHeadsFloat_ * sizeof(float)); + pipe_->InitBuffer(negExpMultiBuf_, rowsPerIter_ * alignedHeadsFloat_ * sizeof(float)); + } + + // Scratch buffers (V-only access). + pipe_->InitBuffer(xBuf_, rowsPerIter_ * alignedHeadsFloat_ * sizeof(float)); + pipe_->InitBuffer(betaXBuf_, rowsPerIter_ * alignedHeadsFloat_ * sizeof(float)); + pipe_->InitBuffer(softplusTmpBuf_, rowsPerIter_ * alignedHeadsFloat_ * sizeof(float)); + pipe_->InitBuffer(thresholdMaskBuf_, rowsPerIter_ * alignedHeadsMask_ * sizeof(uint8_t)); + pipe_->InitBuffer(betaFp32Buf_, rowsPerIter_ * alignedHeadsFloat_ * sizeof(float)); + } + + __aicore__ inline void Process() + { + PreloadConstants(); + + uint32_t blockIdx = GetBlockIdx(); + uint32_t blockNum = GetBlockNum(); + if (blockNum == 0) { blockNum = 1; } + + // Chunk-based task distribution. + uint32_t totalChunks = CeilDiv(numBatches_, rowsPerIter_); + uint32_t chunksPerBlock = CeilDiv(totalChunks, blockNum); + uint32_t chunkStart = blockIdx * chunksPerBlock; + uint32_t chunkEnd = chunkStart + chunksPerBlock; + if (chunkEnd > totalChunks) { chunkEnd = totalChunks; } + + for (uint32_t chunk = chunkStart; chunk < chunkEnd; ++chunk) { + ProcessOneChunk(chunk); + } + } + +private: + /*! + * \brief Preload A_log, neg_exp(A_log), dt_bias, and multi-row replicas. + */ + __aicore__ inline void PreloadConstants() + { + LocalTensor tmpALog = aLogInQue_.template AllocTensor(); + dtBiasTensor_ = negExpInQue_.template AllocTensor(); + + DataCopyExtParams paramCopyParams{1, static_cast(numHeads_ * sizeof(ParamDtype)), + 0, 0, 0}; + DataCopyPadExtParams paramPadParams{false, 0, 0, static_cast(0)}; + + // Load A_log. + DataCopyPad(tmpALog, aLogGm_, paramCopyParams, paramPadParams); + aLogInQue_.template EnQue(tmpALog); + tmpALog = aLogInQue_.template DeQue(); + + if constexpr (std::is_same()) { + Adds(dtBiasTensor_, tmpALog, 0.0f, numHeads_); + } else { + Cast(dtBiasTensor_, tmpALog, RoundMode::CAST_NONE, numHeads_); + } + PipeBarrier(); + + // neg_exp(A_log). + Exp(dtBiasTensor_, dtBiasTensor_, numHeads_); + PipeBarrier(); + Muls(dtBiasTensor_, dtBiasTensor_, -1.0f, numHeads_); + PipeBarrier(); + + aLogInQue_.FreeTensor(tmpALog); + + negExpInQue_.template EnQue(dtBiasTensor_); + dtBiasTensor_ = negExpInQue_.template DeQue(); + + // Load dt_bias. + LocalTensor tmpDtBias = dtBiasInQue_.template AllocTensor(); + dtBiasPreloaded_ = dtBiasPreloadQue_.template AllocTensor(); + DataCopyPad(tmpDtBias, dtBiasGm_, paramCopyParams, paramPadParams); + dtBiasInQue_.template EnQue(tmpDtBias); + tmpDtBias = dtBiasInQue_.template DeQue(); + if constexpr (std::is_same()) { + Adds(dtBiasPreloaded_, tmpDtBias, 0.0f, numHeads_); + } else { + Cast(dtBiasPreloaded_, tmpDtBias, RoundMode::CAST_NONE, numHeads_); + } + PipeBarrier(); + dtBiasInQue_.FreeTensor(tmpDtBias); + dtBiasPreloadQue_.template EnQue(dtBiasPreloaded_); + dtBiasPreloaded_ = dtBiasPreloadQue_.template DeQue(); + + // Replicate to multi-row buffers (skip for single-row kernels). + if (rowsPerIter_ > 1) { + LocalTensor dtBiasMulti = dtBiasMultiBuf_.Get(); + LocalTensor negExpMulti = negExpMultiBuf_.Get(); + for (uint32_t r = 0; r < rowsPerIter_; ++r) { + const uint32_t off = r * alignedHeadsFloat_; + Adds(dtBiasMulti[off], dtBiasPreloaded_, 0.0f, numHeads_); + Adds(negExpMulti[off], dtBiasTensor_, 0.0f, numHeads_); + } + PipeBarrier(); + } + } + + __aicore__ inline void ProcessOneChunk(uint32_t chunkIdx) + { + const uint32_t baseRow = chunkIdx * rowsPerIter_; + if (baseRow >= numBatches_) { + return; + } + const uint32_t remaining = (numBatches_ > baseRow) ? (numBatches_ - baseRow) : 0; + const uint32_t validRows = (remaining >= rowsPerIter_) ? rowsPerIter_ : remaining; + const bool isFullChunk = (validRows == rowsPerIter_); + + LocalTensor aLocal = aInQue_.template AllocTensor(); + LocalTensor bLocal = bInQue_.template AllocTensor(); + LocalTensor gLocal = gOutQue_.template AllocTensor(); + LocalTensor betaLocal = betaOutQue_.template AllocTensor(); + + // MTE2: Load input. + if (useBulkDma_ && isFullChunk) { + const uint64_t rowOffset = static_cast(baseRow) * numHeads_; + const uint32_t rowBytesHalf = numHeads_ * static_cast(sizeof(InDtype)); + const uint32_t inputDstGap = + (alignedHeadsHalf_ - numHeads_) * static_cast(sizeof(InDtype)) / BYTES_PER_BLOCK; + DataCopyExtParams bulkCopyParams{static_cast(rowsPerIter_), + rowBytesHalf, 0, inputDstGap, 0}; + DataCopyPadExtParams bulkPadParams{false, 0, 0, static_cast(0)}; + DataCopyPad(aLocal, aGm_[rowOffset], bulkCopyParams, bulkPadParams); + DataCopyPad(bLocal, bGm_[rowOffset], bulkCopyParams, bulkPadParams); + } else { + for (uint32_t r = 0; r < validRows; ++r) { + const uint64_t rowOffset = static_cast(baseRow + r) * numHeads_; + DataCopyExtParams rowCopyParams{1, static_cast(numHeads_ * sizeof(InDtype)), 0, 0, 0}; + DataCopyPadExtParams rowPadParams{false, 0, 0, static_cast(0)}; + DataCopyPad(aLocal[r * alignedHeadsHalf_], aGm_[rowOffset], rowCopyParams, rowPadParams); + DataCopyPad(bLocal[r * alignedHeadsHalf_], bGm_[rowOffset], rowCopyParams, rowPadParams); + } + } + + aInQue_.template EnQue(aLocal); + bInQue_.template EnQue(bLocal); + aLocal = aInQue_.template DeQue(); + bLocal = bInQue_.template DeQue(); + + LocalTensor x = xBuf_.Get(); + LocalTensor betaX = betaXBuf_.Get(); + LocalTensor softplusTmp = softplusTmpBuf_.Get(); + LocalTensor thresholdMask = thresholdMaskBuf_.Get(); + LocalTensor betaFp32 = betaFp32Buf_.Get(); + + const uint32_t multiCount = validRows * alignedHeadsFloat_; + const uint32_t maskCount = validRows * alignedHeadsMask_; + + // Batch Cast a→fp32, b→fp32. + Cast(x, aLocal, RoundMode::CAST_NONE, multiCount); + Cast(betaFp32, bLocal, RoundMode::CAST_NONE, multiCount); + PipeBarrier(); + + if (rowsPerIter_ > 1) { + // Multi-row path: dt_bias and neg_exp from precomputed buffers. + LocalTensor dtBiasMulti = dtBiasMultiBuf_.Get(); + LocalTensor negExpMulti = negExpMultiBuf_.Get(); + Add(x, x, dtBiasMulti, multiCount); + PipeBarrier(); + Muls(betaX, x, beta_, multiCount); + PipeBarrier(); + Mins(softplusTmp, betaX, threshold_, multiCount); + PipeBarrier(); + Exp(softplusTmp, softplusTmp, multiCount); + PipeBarrier(); + Adds(softplusTmp, softplusTmp, 1.0f, multiCount); + PipeBarrier(); + Ln(softplusTmp, softplusTmp, multiCount); + PipeBarrier(); + Muls(softplusTmp, softplusTmp, 1.0f / beta_, multiCount); + PipeBarrier(); + CompareScalar(thresholdMask, betaX, threshold_, CMPMODE::LE, maskCount); + PipeBarrier(); + Select(gLocal, thresholdMask, softplusTmp, x, SELMODE::VSEL_TENSOR_TENSOR_MODE, multiCount); + PipeBarrier(); + Mul(gLocal, gLocal, negExpMulti, multiCount); + PipeBarrier(); + } else { + // Single-row fallback. + Add(x, x, dtBiasPreloaded_, numHeads_); + PipeBarrier(); + Muls(betaX, x, beta_, multiCount); + PipeBarrier(); + Mins(softplusTmp, betaX, threshold_, multiCount); + PipeBarrier(); + Exp(softplusTmp, softplusTmp, multiCount); + PipeBarrier(); + Adds(softplusTmp, softplusTmp, 1.0f, multiCount); + PipeBarrier(); + Ln(softplusTmp, softplusTmp, multiCount); + PipeBarrier(); + Muls(softplusTmp, softplusTmp, 1.0f / beta_, multiCount); + PipeBarrier(); + CompareScalar(thresholdMask, betaX, threshold_, CMPMODE::LE, maskCount); + PipeBarrier(); + Select(gLocal, thresholdMask, softplusTmp, x, SELMODE::VSEL_TENSOR_TENSOR_MODE, multiCount); + PipeBarrier(); + Mul(gLocal, gLocal, dtBiasTensor_, multiCount); + PipeBarrier(); + } + + // Numerically stable sigmoid: 1 / (1 + exp(-b)). + Muls(betaFp32, betaFp32, -1.0f, multiCount); + PipeBarrier(); + Exp(betaFp32, betaFp32, multiCount); + PipeBarrier(); + Duplicate(x, 1.0f, multiCount); + PipeBarrier(); + Add(betaFp32, betaFp32, x, multiCount); + PipeBarrier(); + Div(x, x, betaFp32, multiCount); + PipeBarrier(); + Cast(betaLocal, x, RoundMode::CAST_RINT, multiCount); + PipeBarrier(); + + aInQue_.FreeTensor(aLocal); + bInQue_.FreeTensor(bLocal); + + gOutQue_.template EnQue(gLocal); + betaOutQue_.template EnQue(betaLocal); + + // MTE3: Write output. + gLocal = gOutQue_.template DeQue(); + betaLocal = betaOutQue_.template DeQue(); + + if (useBulkDma_ && isFullChunk) { + const uint64_t rowOffset = static_cast(baseRow) * numHeads_; + const uint32_t gSrcGap = + (alignedHeadsFloat_ - numHeads_) * static_cast(sizeof(float)) / BYTES_PER_BLOCK; + const uint32_t bSrcGap = + (alignedHeadsHalf_ - numHeads_) * static_cast(sizeof(InDtype)) / BYTES_PER_BLOCK; + DataCopyExtParams gOutParams{static_cast(rowsPerIter_), + numHeads_ * static_cast(sizeof(float)), + gSrcGap, 0, 0}; + DataCopyExtParams bOutParams{static_cast(rowsPerIter_), + numHeads_ * static_cast(sizeof(InDtype)), + bSrcGap, 0, 0}; + DataCopyPad(gGm_[rowOffset], gLocal, gOutParams); + DataCopyPad(betaGm_[rowOffset], betaLocal, bOutParams); + } else { + for (uint32_t r = 0; r < validRows; ++r) { + const uint64_t rowOffset = static_cast(baseRow + r) * numHeads_; + DataCopyParams gOutParams{1, static_cast(numHeads_ * sizeof(float)), 0, 0}; + DataCopyParams bOutParams{1, static_cast(numHeads_ * sizeof(InDtype)), 0, 0}; + DataCopyPad(gGm_[rowOffset], gLocal[r * alignedHeadsFloat_], gOutParams); + DataCopyPad(betaGm_[rowOffset], betaLocal[r * alignedHeadsHalf_], bOutParams); + } + } + + gOutQue_.FreeTensor(gLocal); + betaOutQue_.FreeTensor(betaLocal); + } + +private: + TPipe *pipe_{nullptr}; + + GlobalTensor aLogGm_; + GlobalTensor dtBiasGm_; + GlobalTensor aGm_; + GlobalTensor bGm_; + GlobalTensor gGm_; + GlobalTensor betaGm_; + + TQue aInQue_; + TQue bInQue_; + TQue aLogInQue_; + TQue dtBiasInQue_; + TQue negExpInQue_; + TQue dtBiasPreloadQue_; + TQue gOutQue_; + TQue betaOutQue_; + + TBuf dtBiasMultiBuf_; + TBuf negExpMultiBuf_; + TBuf xBuf_; + TBuf betaXBuf_; + TBuf softplusTmpBuf_; + TBuf thresholdMaskBuf_; + TBuf betaFp32Buf_; + + LocalTensor dtBiasTensor_; // neg_exp(A_log), 1 row + LocalTensor dtBiasPreloaded_; // dt_bias, 1 row + + uint32_t numHeads_{0}; + uint32_t numBatches_{0}; + uint32_t rowsPerIter_{1}; + bool useBulkDma_{false}; + uint32_t alignedHeadsHalf_{0}; + uint32_t alignedHeadsFloat_{0}; + uint32_t alignedHeadsMask_{0}; + uint32_t alignedHeadsParam_{0}; + float beta_{1.0f}; + float threshold_{20.0f}; +}; + +} // namespace FusedGdnGating + +#endif // FUSED_GDN_GATING_KERNEL_H diff --git a/csrc/ascend/attention/fused_gdn_gating/op_kernel/fused_gdn_gating_tiling_data.h b/csrc/ascend/attention/fused_gdn_gating/op_kernel/fused_gdn_gating_tiling_data.h new file mode 100644 index 000000000..c8faed96b --- /dev/null +++ b/csrc/ascend/attention/fused_gdn_gating/op_kernel/fused_gdn_gating_tiling_data.h @@ -0,0 +1,32 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project + */ + +/*! + * \file fused_gdn_gating_tiling_data.h + * \brief Tiling data shared between host-side tiling and device-side kernel. + */ + +#ifndef FUSED_GDN_GATING_TILING_DATA_H +#define FUSED_GDN_GATING_TILING_DATA_H + +#include "kernel_tiling/kernel_tiling.h" + +namespace FusedGdnGating { + +#pragma pack(push, 8) +struct alignas(8) FusedGdnGatingTilingData { + uint32_t numHeads; + uint32_t numBatches; + uint32_t rowsPerIter; + uint32_t useBulkDma; + float beta; + float threshold; +}; +#pragma pack(pop) + +} // namespace FusedGdnGating + +#endif // FUSED_GDN_GATING_TILING_DATA_H diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/CMakeLists.txt b/csrc/ascend/attention/inplace_partial_rotary_mul/CMakeLists.txt new file mode 100644 index 000000000..86b308249 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/CMakeLists.txt b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/CMakeLists.txt new file mode 100644 index 000000000..031405ebd --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/CMakeLists.txt @@ -0,0 +1,75 @@ +# This program is free software, you can redistribute it and/or modify it. +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This file is a part of the CANN Open Software. +# Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ====================================================================================================================== + +# add_ops_compile_options( +# OP_NAME InplacePartialRotaryMul +# OPTIONS --cce-auto-sync=on +# -Wno-deprecated-declarations +# -Werror +# -mllvm -cce-aicore-hoist-movemask=false +# --op_relocatable_kernel_binary=true +# ) + +# set(inplace_partial_rotary_mul_depends transformer/posembedding/inplace_partial_rotary_mul PARENT_SCOPE) + +# target_sources(op_host_aclnn PRIVATE +# op_host/inplace_partial_rotary_mul_def.cpp +# ) + +# target_sources(optiling PRIVATE +# op_host/inplace_partial_rotary_mul_tiling.cpp +# op_host/inplace_partial_rotary_mul_a3_tiling.cpp +# op_host/rope_regbase_tiling_base.cpp +# op_host/rope_regbase_tiling_a_and_b.cpp +# op_host/rope_regbase_tiling_ab.cpp +# op_host/rope_regbase_tiling_aba_and_ba.cpp +# op_host/rope_regbase_tiling_bab.cpp +# ) + +# if (NOT BUILD_OPEN_PROJECT) +# target_sources(opmaster_ct PRIVATE +# op_host/inplace_partial_rotary_mul_tiling.cpp +# op_host/inplace_partial_rotary_mul_a3_tiling.cpp +# op_host/rope_regbase_tiling_base.cpp +# op_host/rope_regbase_tiling_a_and_b.cpp +# op_host/rope_regbase_tiling_ab.cpp +# op_host/rope_regbase_tiling_aba_and_ba.cpp +# op_host/rope_regbase_tiling_bab.cpp +# ) +# endif () + +# target_include_directories(optiling PRIVATE +# ${CMAKE_CURRENT_SOURCE_DIR}/op_host +# ) + +# target_sources(opsproto PRIVATE +# op_host/inplace_partial_rotary_mul_proto.cpp +# ) + +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + inplace_partial_rotary_mul_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME InplacePartialRotaryMul + OPTIONS --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true +) + + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE inplace_partial_rotary_mul ACLNNTYPE aclnn) +endif() \ No newline at end of file diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_a3_tiling.cpp b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_a3_tiling.cpp new file mode 100644 index 000000000..deea653d7 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_a3_tiling.cpp @@ -0,0 +1,591 @@ +/** + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/*! + * \file inplace_partial_rotary_mul_tiling.cpp + * \brief + */ + +#include "inplace_partial_rotary_mul_tiling.h" + +namespace optiling { +constexpr int64_t TILING_KEY_FLOAT16 = 0; +constexpr int64_t TILING_KEY_BFLOAT16 = 10; +constexpr int64_t TILING_KEY_FLOAT32 = 20; +constexpr int64_t TILING_KEY_UNPAD = 0; +constexpr int64_t TILING_KEY_PAD = 1; +constexpr int64_t TILING_KEY_SPLIT_S = 0; +constexpr int64_t TILING_KEY_SPLIT_BS = 100; +constexpr int64_t TILING_KEY_SPLIT_BSN = 200; +constexpr int64_t FP16_BF16_DTYPE_SIZE = 2; +constexpr int64_t FP32_DTYPE_SIZE = 4; +constexpr int64_t INT32_DTYPE_SIZE = 4; +constexpr int64_t REPEAT_FP32 = 64; +constexpr int64_t TBUF_SIZE = 0; +constexpr int64_t ALIGN_32 = 8; +constexpr int64_t ALIGN_16 = 16; +constexpr int64_t IO_NUM = 3; // sin、cos -> tri +constexpr int64_t BASE_KEY = 2000; +constexpr int64_t CONST_4 = 4; + +class InplacePartialRotaryMulTiling +{ +public: + explicit InplacePartialRotaryMulTiling(gert::TilingContext* context) : context_(context){}; + + ge::graphStatus Init(); + ge::graphStatus DoTiling(); + +private: + ge::graphStatus CheckInput(); + ge::graphStatus CalTilingData(); + ge::graphStatus TilingSplitN(int64_t numHeads, int64_t headDimAlign, int64_t ubSize, + ge::DataType dataDtype); + ge::graphStatus TilingSplitB(int64_t batchSize, int64_t numHeads, int64_t headDimAlign, + int64_t ubSize, ge::DataType dataDtype); + ge::graphStatus TilingSplitS(); + ge::graphStatus TilingSplit(); + void FillTilingData(); + void PrintTilingData() const; + void PrintInfo(); + +private: + int64_t coreNum_ = 0; + int64_t ubSize_=0; + int64_t dtypeX = 0; + int64_t repeatNum_ = 0; + bool isBrc_ = true; + int64_t dim0_ = 0; + int64_t dim1_ = 0; + int64_t dim2_ = 0; + int64_t end_ =0; + int64_t tilingKey_ =1; + bool isAlign_ = false; + bool isSpecial_ = false; + bool isFp32Rope_ = false; + int64_t oneBlockSize_ = 0; + int64_t dtypeSize_ = 2; + int64_t xdim0_ = 0; + int64_t xdim1_ = 0; + int64_t xdim2_ = 0; + int64_t xdim3_ = 0; + int64_t r1dim0_ = 0; + int64_t r1dim1_ = 0; + int64_t r1dim2_ = 0; + int64_t r1dim3_ = 0; + + // tiingdata + int64_t usedCoreNum_ = 0; + int64_t numHead_ = 0; + int64_t headDim_ = 0; + int64_t allHeadDim_ = 0; + int64_t coreTUbLoopTime_ = 0; + int64_t coreBUbLoopTime_ = 0; + int64_t coreTUbLoopTail_ = 0; + int64_t coreBUbLoopTail_ = 0; + int64_t ubFactor_ = 0; + int64_t start_=0; + int64_t blockFactor_=0; + gert::TilingContext* context_ = nullptr; + RopeRegbaseTilingData tilingData_; +}; +int64_t GetCeilInt(int64_t value1, int64_t value2) +{ + if (value2 == 0) + return value2; + return (value1 + value2 - 1) / value2; +} + +int64_t GetDiv(int64_t value1, int64_t value2) +{ + if (value2 == 0) + return value2; + return value1 / value2; +} + +int64_t GetDivRem(int64_t value1, int64_t value2) +{ + if (value2 == 0) + return value2; + return value1 % value2; +} +void InplacePartialRotaryMulTiling::FillTilingData() +{ + tilingData_.set_usedCoreNum(usedCoreNum_); + tilingData_.set_numHead(numHead_); + tilingData_.set_headDim(headDim_); + tilingData_.set_allHeadDim(allHeadDim_); + tilingData_.set_coreTUbLoopTime(coreTUbLoopTime_); + tilingData_.set_coreBUbLoopTime(coreBUbLoopTime_); + tilingData_.set_coreTUbLoopTail(coreTUbLoopTail_); + tilingData_.set_coreBUbLoopTail(coreBUbLoopTail_); + tilingData_.set_ubFactor(ubFactor_); + tilingData_.set_start(start_); + tilingData_.set_blockFactor(blockFactor_); +} +void InplacePartialRotaryMulTiling::PrintTilingData() const +{ + OPS_LOG_I(context_->GetNodeName(), "InplacePartialRotaryMulTiling begin print."); + OPS_LOG_I(context_->GetNodeName(), "usedCoreNum = %ld.", usedCoreNum_); + OPS_LOG_I(context_->GetNodeName(), "numHead_ = %ld.", numHead_); + OPS_LOG_I(context_->GetNodeName(), "headDim_ = %ld.", headDim_); + OPS_LOG_I(context_->GetNodeName(), "allHeadDim_ = %ld.", allHeadDim_); + OPS_LOG_I(context_->GetNodeName(), "coreTUbLoopTime_ = %ld.", coreTUbLoopTime_); + OPS_LOG_I(context_->GetNodeName(), "coreBUbLoopTime_ = %ld.", coreBUbLoopTime_); + OPS_LOG_I(context_->GetNodeName(), "coreTUbLoopTail_ = %ld.", coreTUbLoopTail_); + OPS_LOG_I(context_->GetNodeName(), "coreBUbLoopTail_ = %ld.", coreBUbLoopTail_); + OPS_LOG_I(context_->GetNodeName(), "ubFactor = %ld.", ubFactor_); + OPS_LOG_I(context_->GetNodeName(), "start_ = %ld.", start_); + OPS_LOG_I(context_->GetNodeName(), "blockFactor_ = %ld.", blockFactor_); + OPS_LOG_I(context_->GetNodeName(), "tilingKey = %ld.", tilingKey_); +} +void InplacePartialRotaryMulTiling::PrintInfo() +{ + OPS_LOG_I(context_->GetNodeName(), "usedCoreNum = %ld.", tilingData_.get_usedCoreNum()); + OPS_LOG_I(context_->GetNodeName(), "start = %ld.", tilingData_.get_start()); + OPS_LOG_I(context_->GetNodeName(), "allHeadDim = %ld.", tilingData_.get_allHeadDim()); + OPS_LOG_I(context_->GetNodeName(), " batchSize=%ld.", tilingData_.get_batchSize()); + OPS_LOG_I(context_->GetNodeName(), " seqLen=%ld.", tilingData_.get_seqLen()); + OPS_LOG_I(context_->GetNodeName(), " numHeads=%ld.", tilingData_.get_numHeads()); + OPS_LOG_I(context_->GetNodeName(), " headDim=%ld.", tilingData_.get_headDim()); + OPS_LOG_I(context_->GetNodeName(), " frontCoreNum=%ld.", tilingData_.get_frontCoreNum()); + OPS_LOG_I(context_->GetNodeName(), " tailCoreNum=%ld.", tilingData_.get_tailCoreNum()); + OPS_LOG_I(context_->GetNodeName(), " coreCalcNum=%ld.", tilingData_.get_coreCalcNum()); + OPS_LOG_I(context_->GetNodeName(), " coreCalcTail=%ld.", tilingData_.get_coreCalcTail()); + OPS_LOG_I(context_->GetNodeName(), " ubCalcNum=%ld.", tilingData_.get_ubCalcNum()); + OPS_LOG_I(context_->GetNodeName(), " ubCalcLoop=%ld.", tilingData_.get_ubCalcLoop()); + OPS_LOG_I(context_->GetNodeName(), " ubCalcTail=%ld.", tilingData_.get_ubCalcTail()); + OPS_LOG_I(context_->GetNodeName(), " ubCalcTailNum=%ld.", tilingData_.get_ubCalcTailNum()); + OPS_LOG_I(context_->GetNodeName(), " ubCalcTailLoop=%ld.", tilingData_.get_ubCalcTailLoop()); + OPS_LOG_I(context_->GetNodeName(), " ubCalcTailTail=%ld.", tilingData_.get_ubCalcTailTail()); + OPS_LOG_I(context_->GetNodeName(), " ubCalcBNum=%ld.", tilingData_.get_ubCalcBNum()); + OPS_LOG_I(context_->GetNodeName(), " ubCalcBLoop=%ld.", tilingData_.get_ubCalcBLoop()); + OPS_LOG_I(context_->GetNodeName(), " ubCalcBTail=%ld.", tilingData_.get_ubCalcBTail()); + OPS_LOG_I(context_->GetNodeName(), " ubCalcNNum=%ld.", tilingData_.get_ubCalcNNum()); + OPS_LOG_I(context_->GetNodeName(), " ubCalcNLoop=%ld.", tilingData_.get_ubCalcNLoop()); + OPS_LOG_I(context_->GetNodeName(), " ubCalcNTail=%ld.", tilingData_.get_ubCalcNTail()); + OPS_LOG_I(context_->GetNodeName(), "tilingKey = %ld.", tilingKey_); +} +ge::graphStatus InplacePartialRotaryMulTiling::CalTilingData() +{ + OPS_ERR_IF(!isSpecial_, OPS_LOG_I("Tiling4InplacePartialRotaryMul", "not special"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(!isAlign_, OPS_LOG_I("Tiling4InplacePartialRotaryMul", " d not align"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(xdim3_ > REPEAT_FP32, OPS_LOG_I("Tiling4InplacePartialRotaryMul", "D is repeat one repeat"), + return ge::GRAPH_FAILED); + int64_t ubNum = ubSize_ / sizeof(float); + int64_t last = ubNum - dim1_*dim2_; + int64_t preCoreNumFactor = (dim0_ + coreNum_ - 1) / coreNum_; + usedCoreNum_ = (dim0_ + preCoreNumFactor - 1) / preCoreNumFactor; + int64_t tailCoreNum = dim0_ - preCoreNumFactor *(usedCoreNum_ -1); + blockFactor_ = preCoreNumFactor; + ubFactor_ = last / (CONST_4*dim1_*dim2_ + CONST_4*dim2_); + if (ubFactor_ > preCoreNumFactor) { + ubFactor_ = preCoreNumFactor; + } + + OPS_LOG_I(context_->GetNodeName(), "ubFactor_ = %ld.", ubFactor_); + OPS_ERR_IF(ubFactor_ <= 0, OPS_LOG_I("Tiling4InplacePartialRotaryMul", " is large nout support"), + return ge::GRAPH_FAILED); + coreBUbLoopTime_ = (preCoreNumFactor + ubFactor_ -1) /ubFactor_; + coreBUbLoopTail_ = preCoreNumFactor % ubFactor_; + if (coreBUbLoopTail_ == 0) { + coreBUbLoopTail_ = ubFactor_; + } + coreTUbLoopTime_ = (tailCoreNum + ubFactor_ -1) / ubFactor_; + coreTUbLoopTail_ = tailCoreNum % ubFactor_; + if (coreTUbLoopTail_ == 0) { + coreTUbLoopTail_ = ubFactor_; + } + + return ge::GRAPH_SUCCESS; +} +ge::graphStatus InplacePartialRotaryMulTiling::Init() +{ + OPS_LOG_I(context_->GetNodeName(), "Tiling4InplacePartialRotaryMul Init running."); + OPS_ERR_IF(context_ == nullptr, OPS_REPORT_VECTOR_INNER_ERR("Tiling4InplacePartialRotaryMul", "Tiling context is null"), + return ge::GRAPH_FAILED); + auto platformInfo = context_->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, OPS_REPORT_VECTOR_INNER_ERR("Tiling4InplacePartialRotaryMul", "Tiling platformInfo is null"), + return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + + coreNum_ = ascendcPlatform.GetCoreNumAiv(); + OPS_ERR_IF( + coreNum_ <= 0, OPS_LOG_E(context_->GetNodeName(), "coreNum must be greater than 0."), + return ge::GRAPH_FAILED); + + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + ubSize_ = static_cast(ubSizePlatForm); + OPS_ERR_IF( + ubSize_ <= 0, OPS_LOG_E(context_->GetNodeName(), "ubSize must be greater than 0."), + return ge::GRAPH_FAILED); + OPS_LOG_I(context_->GetNodeName(),"coreNum_ is %ld, ubSize_ %ld ",coreNum_, ubSize_); + return ge::GRAPH_SUCCESS; +} +ge::graphStatus InplacePartialRotaryMulTiling::CheckInput() +{ + auto xInput = context_->GetInputShape(0); + auto inputR1 = context_->GetInputShape(1); + auto inputR2 = context_->GetInputShape(2); + auto xDesc = context_->GetInputDesc(0); + auto r1Desc = context_->GetInputDesc(1); + auto r2Desc = context_->GetInputDesc(2); + OPS_ERR_IF(xDesc == nullptr || r1Desc == nullptr || r2Desc == nullptr, + OPS_LOG_E(context_->GetNodeName(), "get input desc nullptr."), + return ge::GRAPH_FAILED); + auto dataDtype = xDesc->GetDataType(); + auto r1Dtype = r1Desc->GetDataType(); + auto r2Dtype = r2Desc->GetDataType(); + if (dataDtype == ge::DT_FLOAT16 || dataDtype == ge::DT_BF16) { + dtypeSize_ = FP16_BF16_DTYPE_SIZE; + oneBlockSize_ = ALIGN_16; + } else { + dtypeSize_ = FP32_DTYPE_SIZE; + oneBlockSize_ = ALIGN_32; + } + OPS_ERR_IF(r1Dtype != r2Dtype, + OPS_LOG_E(context_->GetNodeName(), "cos and sin dtype must be same."), + return ge::GRAPH_FAILED); + OPS_ERR_IF(r1Dtype != dataDtype && r1Dtype != ge::DT_FLOAT, + OPS_LOG_E(context_->GetNodeName(), "cos/sin dtype must be same as x or float32."), + return ge::GRAPH_FAILED); + isFp32Rope_ = (dataDtype != ge::DT_FLOAT && r1Dtype == ge::DT_FLOAT); + + OPS_ERR_IF(xInput == nullptr || inputR1 == nullptr || inputR2 == nullptr, OPS_LOG_E(context_->GetNodeName(), "get input nullptr."), + return ge::GRAPH_FAILED); + gert::Shape xShape = xInput->GetStorageShape(); + int64_t dimNum = xShape.GetDimNum(); + gert::Shape inputR1Shape = inputR1->GetStorageShape(); + gert::Shape inputR2Shape = inputR2->GetStorageShape(); + int64_t dimNumR1 = inputR1Shape.GetDimNum(); + int64_t dimNumR2 = inputR2Shape.GetDimNum(); + auto inputShape = xInput->GetStorageShape(); + OPS_ERR_IF(dimNum != CONST_4, + OPS_LOG_E(context_->GetNodeName(), "xInput dim:%ld, should be 4.", dimNum), + return ge::GRAPH_FAILED); + OPS_ERR_IF(dimNumR1 != CONST_4 || dimNumR2 != CONST_4, + OPS_LOG_E(context_->GetNodeName(), "dimNumR1:%ld dimNumR2 %ld, r1 r2 dim must be4", dimNumR1, dimNumR2), + return ge::GRAPH_FAILED); + for (int64_t i = 0 ; i < CONST_4; i++) { + int64_t r1dim = inputR1Shape.GetDim(i); + int64_t r2dim = inputR2Shape.GetDim(i); + if (r1dim != r2dim) { + OPS_LOG_E(context_->GetNodeName(), "i is %d r1dim is %ld, r2dim is %ld, not equal",i,r1dim,r2dim); + return ge::GRAPH_FAILED; + } + } + dim0_ = xShape.GetDim(0); + xdim0_ = dim0_; + xdim1_ = xShape.GetDim(1); + xdim2_ = xShape.GetDim(2); + allHeadDim_ = xShape.GetDim(dimNum - 1); + auto attrs = context_->GetAttrs(); + OPS_ERR_IF(attrs == nullptr, + OPS_LOG_E(context_->GetNodeName(), "attrs is nullptr"), + return ge::GRAPH_FAILED); + int64_t mode = *(attrs->GetAttrPointer(0)); + OPS_ERR_IF(mode != 1, + OPS_LOG_E(context_->GetNodeName(), "mode only support interleave"), + return ge::GRAPH_FAILED); + auto sliceListAttr = attrs->GetAttrPointer(1); + auto sliceData = static_cast(sliceListAttr->GetData()); + start_ = sliceData[0]; + end_ = sliceData[1]; + OPS_LOG_I(context_->GetNodeName(), "end_ %ld, end_ %ld",start_, end_); + + headDim_ = end_ - start_; + xdim3_ = headDim_; + OPS_ERR_IF(headDim_ <= 0, + OPS_LOG_E(context_->GetNodeName(), "slice not right"), + return ge::GRAPH_FAILED); + r1dim3_ = inputR1Shape.GetDim(3); + r1dim0_ = inputR1Shape.GetDim(0); + r1dim1_ = inputR1Shape.GetDim(1); + r1dim2_ = inputR1Shape.GetDim(2); + int64_t r1Dim2 = inputR1Shape.GetDim(2); + int64_t xDim1 = xShape.GetDim(1); + int64_t xDim2 = xShape.GetDim(2); + OPS_ERR_IF(headDim_ != r1dim3_, + OPS_LOG_E(context_->GetNodeName(), "slice not right, not equal r1 and r2 last dim num"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(r1dim0_ != dim0_, + OPS_LOG_E(context_->GetNodeName(), "dim0 must be equal"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(r1dim2_ != 1, + OPS_LOG_E(context_->GetNodeName(), "r1dim2_ must be 1"), + return ge::GRAPH_FAILED); + dim2_ = headDim_; + dim1_ = xShape.GetDim(1) * xShape.GetDim(2); + numHead_ = dim1_; + if (r1dim1_ == 1 && r1dim2_ == 1 && (xdim1_ == 1 || xdim2_ == 1)) { + tilingKey_ = 1; + isSpecial_ = true; + if (r1dim1_ == dim1_) { + isBrc_ = false; + tilingKey_ = tilingKey_ + 1; + } + if (isFp32Rope_) { + tilingKey_ += 10; + } + } + if (xdim3_ % oneBlockSize_ == 0){ + isAlign_ = true; + } + OPS_LOG_I(context_->GetNodeName(), "isSpecial_ %d", isSpecial_); + return ge::GRAPH_SUCCESS; +} +ge::graphStatus InplacePartialRotaryMulTiling::TilingSplitN(int64_t numHeads, int64_t headDimAlign, int64_t ubSize, + ge::DataType dataDtype) +{ + const int64_t bufferSize = ubSize - 0; + int64_t totalHeadNum1Size = headDimAlign * IO_NUM * dtypeSize_ + headDimAlign * INT32_DTYPE_SIZE; + if (dataDtype == ge::DT_BF16 || dataDtype == ge::DT_FLOAT16) { + totalHeadNum1Size += headDimAlign * FP32_DTYPE_SIZE * IO_NUM; + } + uint32_t ubCalcNNum{1}, ubCalcNLoop{numHeads}, ubCalcNTail{0}; + OPS_ERR_IF(bufferSize < totalHeadNum1Size, OPS_LOG_E(context_->GetNodeName(), "The D dimension of the input shape is too large."), + return ge::GRAPH_FAILED); + ubCalcNNum = GetDiv(bufferSize, totalHeadNum1Size); + ubCalcNLoop = GetCeilInt(numHeads, ubCalcNNum); + ubCalcNTail = GetDivRem(numHeads, ubCalcNNum) != 0 ? numHeads - (ubCalcNLoop - 1) * ubCalcNNum : 0; + tilingData_.set_ubCalcNNum(ubCalcNNum); + tilingData_.set_ubCalcNLoop(ubCalcNLoop); + tilingData_.set_ubCalcNTail(ubCalcNTail); + tilingKey_ += TILING_KEY_SPLIT_BSN; + context_->SetTilingKey(tilingKey_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus InplacePartialRotaryMulTiling::TilingSplitB(int64_t batchSize, int64_t numHeads, int64_t headDimAlign, + int64_t ubSize, ge::DataType dataDtype) +{ + const int64_t tBufferSize = numHeads * headDimAlign * FP32_DTYPE_SIZE; + const int64_t bufferSize = ubSize - tBufferSize; + int64_t totalBatch1Size = numHeads * headDimAlign * IO_NUM * dtypeSize_; + + if (dataDtype == ge::DT_BF16 || dataDtype == ge::DT_FLOAT16) { + totalBatch1Size += numHeads * headDimAlign * FP32_DTYPE_SIZE * IO_NUM; + } + int64_t ubCalcBNum{1}, ubCalcBLoop{batchSize}, ubCalcBTail{0}; + if (ubSize < tBufferSize || bufferSize < totalBatch1Size) { + OPS_ERR_IF(TilingSplitN(numHeads, headDimAlign, ubSize, dataDtype) != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "TilingSplitN fail."), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; + } + + ubCalcBNum = GetDiv(bufferSize, totalBatch1Size); + ubCalcBLoop = GetCeilInt(batchSize, ubCalcBNum); + ubCalcBTail = GetDivRem(batchSize, ubCalcBNum) != 0 ? batchSize - (ubCalcBLoop - 1) * ubCalcBNum : 0; + + tilingData_.set_ubCalcBNum(ubCalcBNum); + tilingData_.set_ubCalcBLoop(ubCalcBLoop); + tilingData_.set_ubCalcBTail(ubCalcBTail); + + tilingKey_ += TILING_KEY_SPLIT_S; + context_->SetTilingKey(tilingKey_); + return ge::GRAPH_SUCCESS; +} +ge::graphStatus InplacePartialRotaryMulTiling::TilingSplitS() +{ + auto xDesc = context_->GetInputDesc(0); + auto dataDtype = xDesc->GetDataType(); + int64_t batchSize = tilingData_.get_batchSize(); + int64_t seqLen = tilingData_.get_seqLen(); + int64_t numHeads = tilingData_.get_numHeads(); + int64_t headDim = tilingData_.get_headDim(); + // block split + int64_t frontCoreNum = GetDivRem(seqLen, coreNum_) != 0 ? GetDivRem(seqLen, coreNum_) : coreNum_; + int64_t tailCoreNum = seqLen <= coreNum_ ? 0 : coreNum_ - frontCoreNum; + usedCoreNum_ = frontCoreNum + tailCoreNum; + int64_t coreCalcNum = GetCeilInt(seqLen, coreNum_); + int64_t coreCalcTail = GetDiv(seqLen, coreNum_); + tilingData_.set_frontCoreNum(frontCoreNum); + tilingData_.set_tailCoreNum(tailCoreNum); + tilingData_.set_coreCalcNum(coreCalcNum); + tilingData_.set_coreCalcTail(coreCalcTail); + tilingData_.set_usedCoreNum(usedCoreNum_); + context_->SetBlockDim(usedCoreNum_); + int64_t headDimAlign = 0; + if (isAlign_) { + headDimAlign = headDim; + } else { + headDimAlign = GetCeilInt(headDim, oneBlockSize_) * oneBlockSize_; + tilingKey_ += 1; + } + // ub split + int64_t tBufferSize = numHeads * headDimAlign * FP32_DTYPE_SIZE; + int64_t bufferSize = ubSize_ - tBufferSize; + int64_t ioUbSize = batchSize * coreCalcNum * numHeads * headDimAlign * IO_NUM * dtypeSize_; + int64_t totalSeq1Size = batchSize * numHeads * headDimAlign * IO_NUM * dtypeSize_; + if (dataDtype == ge::DT_BF16 || dataDtype == ge::DT_FLOAT16) { + ioUbSize += batchSize * coreCalcNum * numHeads * headDimAlign * FP32_DTYPE_SIZE * IO_NUM; + totalSeq1Size += batchSize * numHeads * headDimAlign * FP32_DTYPE_SIZE * IO_NUM; + } + if (tBufferSize >= ubSize_) { + OPS_ERR_IF(TilingSplitN(numHeads, headDimAlign, ubSize_, dataDtype) != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "TilingSplitN fail."), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; + } + if (ubSize_ < tBufferSize || bufferSize < totalSeq1Size) { + OPS_ERR_IF(TilingSplitB(batchSize, numHeads, headDimAlign, ubSize_, dataDtype) != + ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "TilingSplitB fail."), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; + } + context_->SetTilingKey(tilingKey_); + int64_t ubCalcNum, ubCalcLoop, ubCalcTail; + if (bufferSize < ioUbSize) { + ubCalcNum = GetDiv(bufferSize, totalSeq1Size); + ubCalcLoop = GetCeilInt(coreCalcNum, ubCalcNum); + ubCalcTail = GetDivRem(coreCalcNum, ubCalcNum) != 0 ? coreCalcNum - (ubCalcLoop - 1) * ubCalcNum : 0; + } else { + ubCalcNum = coreCalcNum; + ubCalcLoop = 1; + ubCalcTail = 0; + } + tilingData_.set_ubCalcNum(ubCalcNum); + tilingData_.set_ubCalcLoop(ubCalcLoop); + tilingData_.set_ubCalcTail(ubCalcTail); + // ub split for tail core + int64_t ubCalcTailNum{0}, ubCalcTailLoop{0}, ubCalcTailTail{0}; + if (coreCalcTail != 0) { + ioUbSize = batchSize * coreCalcTail * numHeads * headDimAlign * IO_NUM * dtypeSize_; + totalSeq1Size = batchSize * numHeads * headDimAlign * IO_NUM * dtypeSize_; + if (dataDtype == ge::DT_BF16 || dataDtype == ge::DT_FLOAT16) { + ioUbSize += batchSize * coreCalcNum * numHeads * headDimAlign * FP32_DTYPE_SIZE * IO_NUM; + totalSeq1Size += batchSize * numHeads * headDimAlign * FP32_DTYPE_SIZE * IO_NUM; + } + if (bufferSize < ioUbSize) { + ubCalcTailNum = GetDiv(bufferSize, totalSeq1Size); + ubCalcTailLoop = GetCeilInt(coreCalcTail, ubCalcTailNum); + ubCalcTailTail = + GetDivRem(coreCalcTail, ubCalcTailNum) != 0 ? coreCalcTail - (ubCalcTailLoop - 1) * ubCalcTailNum : 0; + } else { + ubCalcTailNum = coreCalcTail; + ubCalcTailLoop = 1; + ubCalcTailTail = 0; + } + } + tilingData_.set_ubCalcTailNum(ubCalcTailNum); + tilingData_.set_ubCalcTailLoop(ubCalcTailLoop); + tilingData_.set_ubCalcTailTail(ubCalcTailTail); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus InplacePartialRotaryMulTiling::TilingSplit() +{ + int64_t batchSizeOut{1}, seqLenOut{1}, numHeadsOut{1}; + if (r1dim1_ == 1 && r1dim2_ == 1 && xdim0_ == r1dim0_) { + seqLenOut = r1dim0_; + numHeadsOut = xdim1_ * xdim2_; // SBND -> 1S(BN)D -> 1SND + } else if (r1dim0_ == 1 && r1dim2_ == 1 && xdim1_ == r1dim1_) { + seqLenOut = r1dim1_; + batchSizeOut = xdim0_; // BSND + numHeadsOut = xdim2_; + } else if (r1dim0_ == 1 && r1dim1_ == 1 && xdim2_ == r1dim2_) { + seqLenOut = r1dim2_; + batchSizeOut = xdim0_ * xdim1_; // BNSD -> (BN)S1D -> BS1D + } else if (xdim0_ == r1dim0_ && xdim1_ == r1dim1_) { + batchSizeOut = 1; + seqLenOut = r1dim0_ * r1dim1_; + numHeadsOut = xdim2_; // 1,BS,N,D cons/sin 1,BS,1,D + } else { + OPS_LOG_E(context_->GetNodeName(), "The shape of the input x, cos and sin is not supported."); + return ge::GRAPH_FAILED; + } + if (batchSizeOut != 1) { + OPS_LOG_E(context_->GetNodeName(), "batchSizeOut must be 1"); + return ge::GRAPH_FAILED; + } + tilingData_.set_batchSize(batchSizeOut); + tilingData_.set_seqLen(seqLenOut); + tilingData_.set_numHeads(numHeadsOut); + tilingData_.set_headDim(xdim3_); + + OPS_ERR_IF(TilingSplitS() != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "TilingSplitS fail."), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} +ge::graphStatus InplacePartialRotaryMulTiling::DoTiling() +{ + OPS_LOG_I(context_->GetNodeName(), "Enter InplacePartialRotaryMulTiling DoTiling"); + OPS_ERR_IF( + CheckInput() != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "CheckInputShapes is failed"), + return ge::GRAPH_FAILED); + ge::graphStatus calStatus = CalTilingData(); + if (calStatus == ge::GRAPH_SUCCESS) { + FillTilingData(); + PrintTilingData(); + context_->SetBlockDim(usedCoreNum_); + context_->SetTilingKey(tilingKey_); + + size_t* workspaces = context_->GetWorkspaceSizes(1); + workspaces[0] = static_cast(16 * 1024 * 1024); + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; + } else if (isFp32Rope_) { + OPS_LOG_E(context_->GetNodeName(), "float32 cos/sin only supports the special interleave tiling path."); + return ge::GRAPH_FAILED; + } else { + // 开始走原始的71的逻辑 + tilingKey_ = BASE_KEY; + auto dataDtype = context_->GetInputDesc(0)->GetDataType(); + if (dataDtype == ge::DT_FLOAT16) { + tilingKey_ += TILING_KEY_SPLIT_S; + } else if (dataDtype == ge::DT_BF16) { + tilingKey_ += TILING_KEY_BFLOAT16; + } else if (dataDtype == ge::DT_FLOAT) { + tilingKey_ += TILING_KEY_FLOAT32; + dtypeSize_ = FP32_DTYPE_SIZE; + } + tilingData_.set_allHeadDim(allHeadDim_); + tilingData_.set_start(start_); + OPS_ERR_IF(TilingSplit() != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "TilingSplit fail."), return ge::GRAPH_FAILED); + + OPS_LOG_I(context_->GetNodeName(), "[tilingKey]: %ld", tilingKey_); + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + size_t usrWorkspaceSize = 0; + size_t sysWorkspaceSize = 16 * 1024 * 1024; + size_t *currentWorkspace = context_->GetWorkspaceSizes(1); + currentWorkspace[0] = usrWorkspaceSize + sysWorkspaceSize; + + PrintInfo(); + return ge::GRAPH_SUCCESS; + } + +} +ge::graphStatus Tiling4InplacePartialRotaryMul(gert::TilingContext* context) +{ + InplacePartialRotaryMulTiling tilingImpl = InplacePartialRotaryMulTiling(context); + if (tilingImpl.Init() != ge::GRAPH_SUCCESS) { + OPS_LOG_E(context, "Tiling4InplacePartialRotaryMul init failed."); + return ge::GRAPH_FAILED; + } + if (tilingImpl.DoTiling() != ge::GRAPH_SUCCESS) { + OPS_LOG_E(context, "Tiling4InplacePartialRotaryMul do tiling failed."); + return ge::GRAPH_FAILED; + } + OPS_LOG_I(context->GetNodeName(), "end Tiling4InplacePartialRotaryMul"); + return ge::GRAPH_SUCCESS; +} +} // namespace optiling diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_def.cpp b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_def.cpp new file mode 100644 index 000000000..5d0436976 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_def.cpp @@ -0,0 +1,83 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file inplace_partial_rotary_mul_def.cpp + * \brief + */ +#include "register/op_def_registry.h" + +namespace ops { +class InplacePartialRotaryMul : public OpDef { +public: + explicit InplacePartialRotaryMul(const char *name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("cos") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("sin") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Attr("mode").AttrType(OPTIONAL).Int(0); + this->Attr("partial_slice").AttrType(OPTIONAL).ListInt({0, 0}); + + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + + OpAICoreConfig config950; + config950.Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + config950.Input("cos") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + config950.Input("sin") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + config950.Output("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + config950.DynamicCompileStaticFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .ExtendCfgInfo("opFile.value", "inplace_partial_rotary_mul"); + } +}; + +OP_ADD(InplacePartialRotaryMul); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_proto.cpp b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_proto.cpp new file mode 100644 index 000000000..30b069a86 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_proto.cpp @@ -0,0 +1,64 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file apply_rotary_pos_emb_proto.h + * \brief + */ +#ifndef OPS_OP_PROTO_INC_ROTARY_POSITION_EMBEDDING_OPS_H_ +#define OPS_OP_PROTO_INC_ROTARY_POSITION_EMBEDDING_OPS_H_ + +#include "graph/operator_reg.h" + +namespace ge { +/** + * @brief Apply rotary position embedding for a single tensor. + * @par Inputs: + * @li x: A 4D tensor which rotary position embedding is applied, format supports ND, and data type must be float16, + * float or bfloat16. + * @li cos: A 4D tensor which is "cos" in rotary position embedding, format supports ND, data type must be the same as + * "x" or float32, and shape must be the same as "sin". + * @li sin: A 4D tensor which is "sin" in rotary position embedding, format supports ND, data type must be the same as + * "cos". + * @par Outputs: + * y: A 4D tensor which is the result of rotary position embedding, format supports ND, data type must be the same as + * "x", and shape must be the same as "x". + * @par Attributes: + * mode: An optional attribute of type int, specifying the mode of rotary position embedding, must be 0-"half", + * 1-"interleave", 2-"quarter" or 3-"interleave-half". Defaults to 0. Atlas A2 Training Series Product/ Atlas 800I A2 + * Inference Product and Atlas A3 Training Series Product only support 0-"half" and 1-"interleave". + * @attention Constraints: + * Let (B, S, N, D) represents the shape of the 4-D input "x". Under this representation, the shape constraints of each + * parameter can be described as follows: + * @li The D of "x", "cos", "sin", "rotate" and "y" must be equal. For Ascend 950 AI Processor, D should be less or + * equal to 1024. For Atlas A2 Training Series Product/ Atlas 800I A2 Inference Product and Atlas A3 Training Series + * Product, D should be less or equal to 896. + * @li In half, interleave and interleave-half mode, D must be a multiple of 2. In quarter mode, D must be a multiple + * of 4. + * @li B, S, N of "cos" and "sin" must meet one of the following four conditions: + * - B, S, N are 1, means the shape is (1, 1, 1, D). + * - B, S, N are the same as that of "x", means the shape is (B, S, N, D). + * - One of S and N is 1, the remaining one dimension and B are the same as that of "x", means the shape is (B, 1, N, + * D) or (B, S, 1, D). + * - Two of B, S and N are 1, the remaining one dimension is the same as that of "x", means the shape is (1, 1, N, D), + * (1, S, 1, D) or (B, 1, 1, D). + */ +REG_OP(InplacePartialRotaryMul) + .INPUT(x, TensorType({DT_FLOAT16, DT_FLOAT, DT_BFLOAT16, DT_FLOAT16, DT_BFLOAT16})) + .INPUT(cos, TensorType({DT_FLOAT16, DT_FLOAT, DT_BFLOAT16, DT_FLOAT, DT_FLOAT})) + .INPUT(sin, TensorType({DT_FLOAT16, DT_FLOAT, DT_BFLOAT16, DT_FLOAT, DT_FLOAT})) + .OUTPUT(x, TensorType({DT_FLOAT16, DT_FLOAT, DT_BFLOAT16, DT_FLOAT16, DT_BFLOAT16})) + .ATTR(mode, Int, 0) + .ATTR(partial_slice, ListInt, {0, 0}) + .OP_END_FACTORY_REG(InplacePartialRotaryMul) + +} // namespace ge + +#endif diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_tiling.cpp b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_tiling.cpp new file mode 100644 index 000000000..f757b0ac2 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_tiling.cpp @@ -0,0 +1,118 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotary_position_embedding.cc + * \brief + */ +#include "inplace_partial_rotary_mul_tiling.h" +#include "register/op_def_registry.h" +// #include "log/log.h" +#include "tiling/tiling_api.h" +// #include "tiling_base/tiling_templates_registry.h" +#include +namespace optiling { +constexpr uint32_t MODE_ATTR_IDX = 0; + +ge::graphStatus RotaryPosEmbeddingMembaseTilingClass::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + if (platformInfo != nullptr) { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + aicoreParams_.blockDim = ascendcPlatform.GetCoreNumAiv(); + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + socVersion_ = ascendcPlatform.GetSocVersion(); + aicoreParams_.ubSize = ubSizePlatForm; + } else { + auto compileInfoPtr = reinterpret_cast(context_->GetCompileInfo()); + OPS_ERR_IF(compileInfoPtr == nullptr, OPS_LOG_E(context_, "compile info is null"), return ge::GRAPH_FAILED); + aicoreParams_.ubSize = compileInfoPtr->ubSize; + aicoreParams_.blockDim = compileInfoPtr->blockDim; + socVersion_ = compileInfoPtr->socVersion; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RotaryPosEmbeddingMembaseTilingClass::GetShapeAttrsInfo() +{ + auto attrs = context_->GetAttrs(); + OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED); + const uint32_t inputMode = *(attrs->GetAttrPointer(MODE_ATTR_IDX)); + OPS_LOG_I(context_->GetNodeName(), "[mode]: %d", inputMode); + inputMode_ = inputMode; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus Tiling4RotaryPositionEmbedding(gert::TilingContext *context) +{ + OPS_LOG_I(context, "Tiling4RotaryPositionEmbedding start"); + OPS_ERR_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("Tiling4RotaryPositionEmbedding", "Tiling context is null"), + return ge::GRAPH_FAILED); + + auto platformInfo = context->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, OPS_REPORT_VECTOR_INNER_ERR("Tiling4RotaryPositionEmbedding", "Tiling platformInfo is null"), + return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + auto socVersion = ascendcPlatform.GetSocVersion(); + auto xDesc = context->GetInputDesc(0); + auto cosDesc = context->GetInputDesc(1); + auto sinDesc = context->GetInputDesc(2); + OPS_ERR_IF(xDesc == nullptr || cosDesc == nullptr || sinDesc == nullptr, + OPS_REPORT_VECTOR_INNER_ERR("Tiling4RotaryPositionEmbedding", "input desc is null"), + return ge::GRAPH_FAILED); + bool useFp32Rope = xDesc->GetDataType() != ge::DT_FLOAT && + cosDesc->GetDataType() == ge::DT_FLOAT && + sinDesc->GetDataType() == ge::DT_FLOAT; + bool supportFp32Rope = socVersion == platform_ascendc::SocVersion::ASCEND910B || + socVersion == platform_ascendc::SocVersion::ASCEND910_93; + if (useFp32Rope && supportFp32Rope) { + return Tiling4InplacePartialRotaryMul(context); + } + if (socVersion == platform_ascendc::SocVersion::ASCEND950) + { + std::vector> regBaseTilingCases; + regBaseTilingCases.push_back(std::unique_ptr(new RopeRegBaseTilingClassAAndB(context))); + regBaseTilingCases.push_back(std::unique_ptr(new RopeRegBaseTilingClassAB(context))); + regBaseTilingCases.push_back(std::unique_ptr(new RopeRegBaseTilingClassABAAndBA(context))); + regBaseTilingCases.push_back(std::unique_ptr(new RopeRegBaseTilingClassBAB(context))); + OPS_LOG_I(context, "Using arch35 tiling for ASCEND950"); + + for (const auto& ptr : regBaseTilingCases) + { + if (ptr) + { + ge::graphStatus status = ptr->DoTiling(); + if (status != ge::GRAPH_PARAM_INVALID) + { + OPS_LOG_I(context, "Do general op tiling success priority"); + return status; + } + OPS_LOG_I(context, "Ignore general op tiling priority"); + } + } + OPS_LOG_I(context, "Using tiling for ASCEND910_71"); + RotaryPosEmbeddingMembaseTilingClass rotaryPosEmbeddingMembaseTilingClass(context); + return rotaryPosEmbeddingMembaseTilingClass.DoOpTiling(); + } else { + return Tiling4InplacePartialRotaryMul(context); + } +} + +ge::graphStatus TilingPrepareForRotaryPositionEmbedding(gert::TilingParseContext *context) +{ + OPS_LOG_I(context, "TilingPrepareForRotaryPositionEmbedding context success"); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(InplacePartialRotaryMul) + .Tiling(Tiling4RotaryPositionEmbedding) + .TilingParse(TilingPrepareForRotaryPositionEmbedding); +} // namespace optiling diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_tiling.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_tiling.h new file mode 100644 index 000000000..0bee33b05 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_tiling.h @@ -0,0 +1,491 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotary_position_embedding.h + * \brief + */ +#ifndef OPS_BUILD_IN_OP_TILING_RUNTIME_ROTARY_POSITION_EMBEDDING_H +#define OPS_BUILD_IN_OP_TILING_RUNTIME_ROTARY_POSITION_EMBEDDING_H + + +#include "register/tilingdata_base.h" +#include "register/op_def_registry.h" +// #include "tiling_base/tiling_templates_registry.h" +#include "tiling/tiling_api.h" +// #include "tiling_base/tiling_base.h" +// #include "tiling_base/tiling_util.h" +#include "platform/platform_info.h" +// #include "util/math_util.h" +#include "error/ops_error.h" +namespace optiling { + +BEGIN_TILING_DATA_DEF(RopeRegbaseTilingData) +TILING_DATA_FIELD_DEF(int64_t, B); +TILING_DATA_FIELD_DEF(int64_t, CosB); +TILING_DATA_FIELD_DEF(int64_t, S); +TILING_DATA_FIELD_DEF(int64_t, D); +TILING_DATA_FIELD_DEF(int64_t, N); +TILING_DATA_FIELD_DEF(int64_t, blockNumB); +TILING_DATA_FIELD_DEF(int64_t, blockFactorB); +TILING_DATA_FIELD_DEF(int64_t, blockNumS); +TILING_DATA_FIELD_DEF(int64_t, blockFactorS); +TILING_DATA_FIELD_DEF(int64_t, ubLoopNumS); +TILING_DATA_FIELD_DEF(int64_t, ubFactorS); +TILING_DATA_FIELD_DEF(int64_t, ubTailFactorS); +TILING_DATA_FIELD_DEF(int64_t, ubLoopNumB); +TILING_DATA_FIELD_DEF(int64_t, ubFactorB); +TILING_DATA_FIELD_DEF(int64_t, ubTailFactorB); +TILING_DATA_FIELD_DEF(int64_t, ubLoopNumN); +TILING_DATA_FIELD_DEF(int64_t, ubFactorN); +TILING_DATA_FIELD_DEF(int64_t, ubTailFactorN); +TILING_DATA_FIELD_DEF(int64_t, rotaryMode); +TILING_DATA_FIELD_DEF(int64_t, dAlign); +TILING_DATA_FIELD_DEF(int64_t, dSplitCoef); +TILING_DATA_FIELD_DEF(int64_t, blockNumBS); +TILING_DATA_FIELD_DEF(int64_t, blockFactorBS); +TILING_DATA_FIELD_DEF(int64_t, blockTailBS); +TILING_DATA_FIELD_DEF(int64_t, blockNumN); +TILING_DATA_FIELD_DEF(int64_t, blockFactorN); +TILING_DATA_FIELD_DEF(int64_t, blockTailN); +TILING_DATA_FIELD_DEF(int64_t, ubFactorBS); +TILING_DATA_FIELD_DEF(int64_t, sliceStart); +TILING_DATA_FIELD_DEF(int64_t, sliceEnd); +TILING_DATA_FIELD_DEF(int64_t, sliceLength); +// A3 +TILING_DATA_FIELD_DEF(int64_t, usedCoreNum); +TILING_DATA_FIELD_DEF(int64_t, numHead); +TILING_DATA_FIELD_DEF(int64_t, headDim); +TILING_DATA_FIELD_DEF(int64_t, allHeadDim); +TILING_DATA_FIELD_DEF(int64_t, coreTUbLoopTime); +TILING_DATA_FIELD_DEF(int64_t, coreBUbLoopTime); +TILING_DATA_FIELD_DEF(int64_t, coreTUbLoopTail); +TILING_DATA_FIELD_DEF(int64_t, coreBUbLoopTail); +TILING_DATA_FIELD_DEF(int64_t, ubFactor); +TILING_DATA_FIELD_DEF(int64_t, start); +TILING_DATA_FIELD_DEF(int64_t, blockFactor); +//A3 主线 +TILING_DATA_FIELD_DEF(int64_t, batchSize); +TILING_DATA_FIELD_DEF(int64_t, seqLen); +TILING_DATA_FIELD_DEF(int64_t, numHeads); +TILING_DATA_FIELD_DEF(int64_t, frontCoreNum); +TILING_DATA_FIELD_DEF(int64_t, tailCoreNum); +TILING_DATA_FIELD_DEF(int64_t, coreCalcNum); +TILING_DATA_FIELD_DEF(int64_t, coreCalcTail); +TILING_DATA_FIELD_DEF(int64_t, ubCalcNum); +TILING_DATA_FIELD_DEF(int64_t, ubCalcLoop); +TILING_DATA_FIELD_DEF(int64_t, ubCalcTail); +TILING_DATA_FIELD_DEF(int64_t, ubCalcTailNum); +TILING_DATA_FIELD_DEF(int64_t, ubCalcTailLoop); +TILING_DATA_FIELD_DEF(int64_t, ubCalcTailTail); +TILING_DATA_FIELD_DEF(int64_t, ubCalcBNum); +TILING_DATA_FIELD_DEF(int64_t, ubCalcBLoop); +TILING_DATA_FIELD_DEF(int64_t, ubCalcBTail); +TILING_DATA_FIELD_DEF(int64_t, ubCalcNNum); +TILING_DATA_FIELD_DEF(int64_t, ubCalcNLoop); +TILING_DATA_FIELD_DEF(int64_t, ubCalcNTail); + +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(InplacePartialRotaryMul, RopeRegbaseTilingData) + +ge::graphStatus Tiling4InplacePartialRotaryMul(gert::TilingContext* context); +struct RotaryPositionEmbeddingCompileInfo { + int64_t blockDim; + uint64_t ubSize; + platform_ascendc::SocVersion socVersion; +}; + +struct AiCoreParams { + uint64_t ubSize = 0; + uint64_t blockDim = 0; + uint64_t aicNum = 0; + uint64_t l1Size = 0; + uint64_t l0aSize = 0; + uint64_t l0bSize = 0; + uint64_t l0cSize = 0; +}; + +enum class RopeLayout : uint8_t { + NO_BROADCAST = 1, + BROADCAST_BSN = 2, + BSND = 3, + SBND = 4, + BNSD = 5 +}; + +enum class RotaryPosEmbeddingMode : uint8_t { + HALF = 0, + INTERLEAVE = 1, + QUARTER = 2, + DEEPSEEK_INTERLEAVE = 3 +}; + +template +static inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd))); +} + +template +static inline T CeilAlign(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)) * (rnd)); +} + +template +static inline T FloorDiv(T x, T y) +{ + if (y == 0) + { + return 0; + } + return x / y; +} + +template +static inline T FloorAlign(T x, T y) +{ + if (y == 0) + { + return 0; + } + return x / y * y; +} + +class RotaryPosEmbeddingMembaseTilingClass { +public: + explicit RotaryPosEmbeddingMembaseTilingClass(gert::TilingContext *context) : context_(context) + { + } + + void Reset(gert::TilingContext *context) + { + RotaryPosEmbeddingMembaseTilingClass::Reset(context); + } + + ge::graphStatus GetPlatformInfo(); + + ge::graphStatus GetWorkspaceSize() + { + return ge::GRAPH_SUCCESS; + } + + ge::graphStatus DoLibApiTiling() + { + return ge::GRAPH_SUCCESS; + } + + bool IsCapable() + { + return true; + } + // 3、计算数据切分TilingData + ge::graphStatus DoOpTiling() + { + return ge::GRAPH_SUCCESS; + } + // 7、保存Tiling数据 + ge::graphStatus PostTiling() + { + return ge::GRAPH_SUCCESS; + } + + ge::graphStatus GetShapeAttrsInfo(); + + uint64_t GetTilingKey() const + { + return context_->GetTilingKey(); + } + +protected: + static const uint32_t MODE_ROTATE_INTERLEAVED = 1; + uint32_t inputMode_ = 0; + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B; + gert::TilingContext* context_ = nullptr; + std::unique_ptr ascendcPlatform_{nullptr}; + uint32_t blockDim_{0}; + uint64_t workspaceSize_{0}; + uint64_t tilingKey_{0}; + AiCoreParams aicoreParams_; +}; + +class RopeRegBaseTilingClass { +public: + explicit RopeRegBaseTilingClass(gert::TilingContext *context) : context_(context) + { + } + + void Reset(gert::TilingContext *context) + { + RopeRegBaseTilingClass::Reset(context); + } + + bool IsRotaryPosEmbeddingMode(const int32_t mode) const; + ge::graphStatus CheckNullptr(); + ge::graphStatus CheckShape(); + ge::graphStatus CheckDtypeAndAttr(); + ge::graphStatus CheckParam(); + ge::graphStatus JudgeLayoutByShape(const gert::Shape &xShape, const gert::Shape &cosShape); + ge::graphStatus CheckRotaryModeShapeRelation(const int64_t d); + ge::graphStatus CheckShapeAllPositive(const int64_t idx) const; + ge::graphStatus CheckShapeAllPositive() const; + std::string rotaryModeStr_; + + ge::graphStatus GetPlatformInfo(); + ge::graphStatus GetShapeAttrsInfo(); + ge::graphStatus JudgeSliceInfo(); + virtual ge::graphStatus GetWorkspaceSize() + { + return ge::GRAPH_SUCCESS; + } + + virtual ge::graphStatus DoLibApiTiling() + { + return ge::GRAPH_SUCCESS; + } + + virtual ge::graphStatus DoOpTiling() + { + return ge::GRAPH_SUCCESS; + } + + virtual uint64_t GetTilingKey() const + { + return 0; + } + + virtual ge::graphStatus PostTiling() + { + return ge::GRAPH_SUCCESS; + } + + virtual bool IsCapable() + { + return true; + } + + bool IsRegbaseSocVersion() + { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo()); + auto socVersion = ascendcPlatform.GetSocVersion(); + return socVersion == platform_ascendc::SocVersion::ASCEND950; + } + + ge::graphStatus DoTiling() + { + auto ret = GetShapeAttrsInfo(); + if (ret != ge::GRAPH_SUCCESS) + { + return ret; + } + ret = GetPlatformInfo(); + if (ret != ge::GRAPH_SUCCESS) + { + return ret; + } + if (!IsCapable()) + { + return ge::GRAPH_PARAM_INVALID; + } + ret = DoOpTiling(); + if (ret != ge::GRAPH_SUCCESS) + { + return ret; + } + ret = DoLibApiTiling(); + if (ret != ge::GRAPH_SUCCESS) + { + return ret; + } + ret = GetWorkspaceSize(); + if (ret != ge::GRAPH_SUCCESS) + { + return ret; + } + ret = PostTiling(); + if (ret != ge::GRAPH_SUCCESS) + { + return ret; + } + context_->SetTilingKey(GetTilingKey()); + return ge::GRAPH_SUCCESS; + } + +protected: + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND950; + int64_t b_{0}; + int64_t s_{0}; + int64_t n_{0}; + int64_t d_{0}; + int64_t cosb_{0}; + ge::DataType dtype_; + RopeLayout layout_; + RotaryPosEmbeddingMode rotaryMode_; + + int64_t blockSize_; + int64_t dSplitCoef_; + bool is1snd_ = false; + gert::TilingContext* context_ = nullptr; + std::unique_ptr ascendcPlatform_{nullptr}; + uint32_t blockDim_{0}; + uint64_t workspaceSize_{0}; + uint64_t tilingKey_{0}; + AiCoreParams aicoreParams_; + int64_t sliceStart_{0}; + int64_t sliceEnd_{0}; + int64_t sliceLength_{0}; + int64_t cosd_{0}; + int64_t sind_{0}; +}; + +class RopeRegBaseTilingClassAAndB : public RopeRegBaseTilingClass { +public: + explicit RopeRegBaseTilingClassAAndB(gert::TilingContext *context) : RopeRegBaseTilingClass(context) + { + } + + bool IsCapable() override; + ge::graphStatus DoOpTiling() override; + ge::graphStatus DoLibApiTiling() override; + uint64_t GetTilingKey() const override; + ge::graphStatus GetWorkspaceSize() override; + ge::graphStatus PostTiling() override; + void SetTilingData(); + +private: + ge::graphStatus MergeDim(); + ge::graphStatus SplitCore(); + ge::graphStatus ComputeUbFactor(); + + int64_t blockNumB_{0}; + int64_t blockFactorB_{0}; + int64_t blockNumS_{0}; + int64_t blockFactorS_{0}; + int64_t ubFactorB_{0}; + int64_t ubLoopNumB_{0}; + int64_t ubTailFactorB_{0}; + int64_t ubFactorS_{0}; + int64_t ubLoopNumS_{0}; + int64_t ubTailFactorS_{0}; + int64_t ubFactorN_{0}; + int64_t ubLoopNumN_{0}; + int64_t ubTailFactorN_{0}; + RopeRegbaseTilingData tilingData_; +}; + +class RopeRegBaseTilingClassAB : public RopeRegBaseTilingClass { +public: + explicit RopeRegBaseTilingClassAB(gert::TilingContext *context) : RopeRegBaseTilingClass(context) + { + } + + bool IsCapable() override; + ge::graphStatus DoOpTiling() override; + ge::graphStatus PostTiling() override; + uint64_t GetTilingKey() const override; + +private: + int64_t blockNumBS_ = 0; + int64_t blockFactorBS_ = 0; + int64_t blockTailBS_ = 0; + int64_t blockNumN_ = 0; + int64_t blockFactorN_ = 0; + int64_t blockTailN_ = 0; + int64_t ubFactorBS_ = 0; + int64_t ubFactorN_ = 0; + int64_t dAlign_ = 0; + RopeRegbaseTilingData tilingData_; +}; + +class RopeRegBaseTilingClassABAAndBA : public RopeRegBaseTilingClass { +public: + explicit RopeRegBaseTilingClassABAAndBA(gert::TilingContext *context) : RopeRegBaseTilingClass(context) + { + } + + bool IsCapable() override; + ge::graphStatus DoOpTiling() override; + ge::graphStatus DoLibApiTiling() override; + uint64_t GetTilingKey() const override; + ge::graphStatus GetWorkspaceSize() override; + ge::graphStatus PostTiling() override; + void SetTilingData(); + +private: + ge::graphStatus SplitCore(); + ge::graphStatus ComputeUbFactor(); + + int64_t blockNumB_{0}; + int64_t blockFactorB_{0}; + int64_t blockNumS_{0}; + int64_t blockFactorS_{0}; + int64_t ubFactorB_{0}; + int64_t ubLoopNumB_{0}; + int64_t ubTailFactorB_{0}; + int64_t ubFactorS_{0}; + int64_t ubLoopNumS_{0}; + int64_t ubTailFactorS_{0}; + int64_t ubFactorN_{0}; + int64_t ubLoopNumN_{0}; + int64_t ubTailFactorN_{0}; + RopeRegbaseTilingData tilingData_; +}; + +class RopeRegBaseTilingClassBAB : public RopeRegBaseTilingClass { +public: + explicit RopeRegBaseTilingClassBAB(gert::TilingContext *context_) : RopeRegBaseTilingClass(context_) + { + } + + // 计算数据切分 + ge::graphStatus DoOpTiling() override; + // 计算TilingKey + uint64_t GetTilingKey() const override; + // 设置Tiling数据 + ge::graphStatus PostTiling() override; + + bool IsCapable() override + { + // BSND format, 1s1d模版，后续可扩展支持所有bab类型的boardcast + if (IsRegbaseSocVersion() && (layout_ == RopeLayout::BSND) && + (cosb_ == 1)) { + return true; + } + return false; + } + +private: + int64_t coreNum_ = 0; + int64_t blockNumB_ = 0; + int64_t blockFactorB_ = 0; + int64_t blockNumS_ = 0; + int64_t blockFactorS_ = 0; + int64_t usedCoreNum_ = 0; + int64_t ubLoopNumS_ = 0; + int64_t ubFactorS_ = 1; + int64_t ubTailFactorS_ = 0; + int64_t ubLoopNumB_ = 0; + int64_t ubFactorB_ = 1; + int64_t ubTailFactorB_ = 0; + int64_t ubLoopNumN_ = 0; // 核内处理N循环了多少次 + int64_t ubFactorN_ = 1; // 每次循环处理多少个N + int64_t ubTailFactorN_ = 0; // 最后一次循环处理多少N + int64_t ubSize_ = 0; + uint64_t tilingKey_ = 0; + + ge::graphStatus SplitCore(); + ge::graphStatus SplitUb(); + void PrintTilingData(); + RopeRegbaseTilingData tilingData_; +}; + +} // namespace optiling +#endif // OPS_BUILD_IN_OP_TILING_RUNTIME_ROTARY_POSITION_EMBEDDING_H diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_a_and_b.cpp b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_a_and_b.cpp new file mode 100644 index 000000000..c560bdf80 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_a_and_b.cpp @@ -0,0 +1,243 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rope_regbase_tiling_a_and_b.cc + * \brief + */ +#include "inplace_partial_rotary_mul_tiling.h" + +using namespace AscendC; + +namespace optiling { +constexpr uint64_t ROPE_A_AND_B_TILING_PRIORITY = 40000; +constexpr int64_t DOUBLE_BUFFER = 2; +constexpr int64_t UB_FACTOR = 4; +constexpr int64_t UB_X_FACTOR = 4; +constexpr int64_t UB_COS_SIN_FACTOR = 2; +constexpr int64_t MAX_COPY_BLOCK_COUNT = 4095; +constexpr int32_t WORKSPACE_SIZE = 16 * 1024 * 1024; +constexpr uint64_t TILING_KEY_A = 20040; +constexpr uint64_t TILING_KEY_B = 20041; +constexpr uint64_t TILING_KEY_A_BF16_FP32 = 20140; +constexpr uint64_t TILING_KEY_A_FP16_FP32 = 20240; +constexpr uint64_t TILING_KEY_B_BF16_FP32 = 20141; +constexpr uint64_t TILING_KEY_B_FP16_FP32 = 20241; + +bool RopeRegBaseTilingClassAAndB::IsCapable() +{ + // 处理全boardcast和不boardcast的情况 + return (IsRegbaseSocVersion()) && (layout_ == RopeLayout::NO_BROADCAST || layout_ == RopeLayout::BROADCAST_BSN); +} + +ge::graphStatus RopeRegBaseTilingClassAAndB::MergeDim() +{ + b_ = b_ * n_ * s_; + n_ = 1; + s_ = 1; + if (layout_ == RopeLayout::NO_BROADCAST) { + cosb_ = b_; + } else { + cosb_ = 1; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClassAAndB::SplitCore() +{ + blockFactorB_ = CeilDiv(static_cast(b_), aicoreParams_.blockDim); + blockNumB_ = CeilDiv(static_cast(b_), blockFactorB_); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClassAAndB::ComputeUbFactor() +{ + ubFactorB_ = 1; + + auto cosDtype = context_->GetInputDesc(1)->GetDataType(); + bool isMixedPrecision = (dtype_ == ge::DT_BF16 || dtype_ == ge::DT_FLOAT16) && cosDtype == ge::DT_FLOAT; + + int64_t dSizeX = CeilAlign(sliceLength_ * GetSizeByDataType(dtype_) / dSplitCoef_, this->blockSize_) * dSplitCoef_; + int64_t dSizeCosSin = + CeilAlign(sliceLength_ * GetSizeByDataType(cosDtype) / dSplitCoef_, this->blockSize_) * dSplitCoef_; + + if (layout_ == RopeLayout::NO_BROADCAST) { + int64_t totalPerBUnit; + if (isMixedPrecision) { + // UB: 4 queues x double-buffer = 8 total buffers + // 4 * ubFactorB * dSizeX + 4 * ubFactorB * dSizeCosSin + // Per-B unit: 4 * dSizeX + 4 * dSizeCosSin + totalPerBUnit = (dSizeX + dSizeCosSin) * UB_FACTOR; + } else { + // UB: 4 queues x double-buffer = 8 * ubFactorB * dSizeX + // Per-B unit: 8 * dSizeX + totalPerBUnit = dSizeX * UB_FACTOR * DOUBLE_BUFFER; + } + int64_t numOfDAvailable = FloorDiv(static_cast(aicoreParams_.ubSize), totalPerBUnit); + OPS_ERR_IF(numOfDAvailable < 1, + OPS_LOG_E(context_, + "D is too big to load in ub, ubSize is %ld bytes, loading requires %ld bytes.", + static_cast(aicoreParams_.ubSize), + totalPerBUnit), + return ge::GRAPH_FAILED); + ubFactorB_ = std::min(blockFactorB_, numOfDAvailable); + } else { + if (isMixedPrecision) { + // UB: UB_X_FACTOR * ubFactorB * dSizeX + UB_COS_SIN_FACTOR * dSizeCosSin + int64_t availableForX = static_cast(aicoreParams_.ubSize) - UB_COS_SIN_FACTOR * dSizeCosSin; + if (availableForX <= 0) { + OPS_LOG_E(context_, + "D is too big to load in ub, ubSize is %ld bytes, cos/sin requires %ld bytes.", + static_cast(aicoreParams_.ubSize), + UB_COS_SIN_FACTOR * dSizeCosSin); + return ge::GRAPH_FAILED; + } + int64_t numOfDAvailable = FloorDiv(availableForX, UB_X_FACTOR * dSizeX); + OPS_ERR_IF(numOfDAvailable < 1, + OPS_LOG_E(context_, + "D is too big to load in ub, ubSize is %ld bytes, loading requires %ld bytes.", + static_cast(aicoreParams_.ubSize), + UB_X_FACTOR * dSizeX + UB_COS_SIN_FACTOR * dSizeCosSin), + return ge::GRAPH_FAILED); + ubFactorB_ = std::min(blockFactorB_, numOfDAvailable); + } else { + int64_t numOfDAvailable = FloorDiv(static_cast(aicoreParams_.ubSize), DOUBLE_BUFFER * dSizeX); + OPS_ERR_IF(numOfDAvailable < UB_FACTOR, + OPS_LOG_E(context_, + "D is too big to load in ub, ubSize is %ld bytes, loading requires %ld bytes.", + static_cast(aicoreParams_.ubSize), + UB_FACTOR * dSizeX * (ubFactorB_ + 1)), + return ge::GRAPH_FAILED); + numOfDAvailable -= 1; + numOfDAvailable /= DOUBLE_BUFFER; + ubFactorB_ = std::min(blockFactorB_, numOfDAvailable); + } + } + + ubFactorB_ = std::min(ubFactorB_, MAX_COPY_BLOCK_COUNT / dSplitCoef_); + + return ge::GRAPH_SUCCESS; +} + +void RopeRegBaseTilingClassAAndB::SetTilingData() +{ + tilingData_.set_B(b_); + tilingData_.set_CosB(cosb_); + tilingData_.set_S(s_); + tilingData_.set_D(d_); + tilingData_.set_N(n_); + tilingData_.set_blockNumB(blockNumB_); + tilingData_.set_blockFactorB(blockFactorB_); + tilingData_.set_blockNumS(blockNumS_); + tilingData_.set_blockFactorS(blockFactorS_); + tilingData_.set_ubLoopNumS(ubLoopNumS_); + tilingData_.set_ubFactorS(ubFactorS_); + tilingData_.set_ubTailFactorS(ubTailFactorS_); + tilingData_.set_ubLoopNumB(ubLoopNumB_); + tilingData_.set_ubFactorB(ubFactorB_); + tilingData_.set_ubTailFactorB(ubTailFactorB_); + tilingData_.set_ubLoopNumN(ubLoopNumN_); + tilingData_.set_ubFactorN(ubFactorN_); + tilingData_.set_ubTailFactorN(ubTailFactorN_); + tilingData_.set_rotaryMode(static_cast(rotaryMode_)); + tilingData_.set_sliceStart(static_cast(sliceStart_)); + tilingData_.set_sliceEnd(static_cast(sliceEnd_)); + tilingData_.set_sliceLength(static_cast(sliceLength_)); + + OPS_LOG_I(context_->GetNodeName(), + "RopeRegBaseTilingClassAAndB tilingData: " + "B is %ld, CosB is %ld, S is %ld, D is %ld, N is %ld, blockNumB %ld," + "blockFactorB_ is %ld, blockNumS %ld, blockFactorS is %ld, ubLoopNumS is %ld," + "ubFactorS is %ld, ubTailFactorS %ld, ubLoopNumB is %ld, ubFactorB is %ld," + "ubTailFactorB is %ld, ubLoopNumN is %ld, ubFactorN is %ld, ubTailFactorN is %ld," + "rotaryMode is %ld, tilingKey is %ld, sliceStart is %ld, sliceEnd is %ld, sliceLength is %ld", + tilingData_.get_B(), + tilingData_.get_CosB(), + tilingData_.get_S(), + tilingData_.get_D(), + tilingData_.get_N(), + tilingData_.get_blockNumB(), + tilingData_.get_blockFactorB(), + tilingData_.get_blockNumS(), + tilingData_.get_blockFactorS(), + tilingData_.get_ubLoopNumS(), + tilingData_.get_ubFactorS(), + tilingData_.get_ubTailFactorS(), + tilingData_.get_ubLoopNumB(), + tilingData_.get_ubFactorB(), + tilingData_.get_ubTailFactorB(), + tilingData_.get_ubLoopNumN(), + tilingData_.get_ubFactorN(), + tilingData_.get_ubTailFactorN(), + tilingData_.get_rotaryMode(), + GetTilingKey(), + tilingData_.get_sliceStart(), + tilingData_.get_sliceEnd(), + tilingData_.get_sliceLength()); +} + +ge::graphStatus RopeRegBaseTilingClassAAndB::DoOpTiling() +{ + OPS_ERR_IF(MergeDim() != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "failed to merge dim."), + return ge::GRAPH_FAILED); + OPS_ERR_IF(SplitCore() != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "failed to split core."), + return ge::GRAPH_FAILED); + OPS_ERR_IF(ComputeUbFactor() != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "failed to compute ub factor."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClassAAndB::DoLibApiTiling() +{ + return ge::GRAPH_SUCCESS; +} + +uint64_t RopeRegBaseTilingClassAAndB::GetTilingKey() const +{ + auto xDtype = context_->GetInputDesc(0)->GetDataType(); + auto cosDtype = context_->GetInputDesc(1)->GetDataType(); + + bool isNoBroadcast = (layout_ == RopeLayout::NO_BROADCAST); + + if (xDtype == ge::DT_BF16 && cosDtype == ge::DT_FLOAT) { + return isNoBroadcast ? TILING_KEY_A_BF16_FP32 : TILING_KEY_B_BF16_FP32; + } else if (xDtype == ge::DT_FLOAT16 && cosDtype == ge::DT_FLOAT) { + return isNoBroadcast ? TILING_KEY_A_FP16_FP32 : TILING_KEY_B_FP16_FP32; + } + + return isNoBroadcast ? TILING_KEY_A : TILING_KEY_B; +} + +ge::graphStatus RopeRegBaseTilingClassAAndB::GetWorkspaceSize() +{ + workspaceSize_ = WORKSPACE_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClassAAndB::PostTiling() +{ + SetTilingData(); + uint64_t tilingKey = GetTilingKey(); + context_->SetTilingKey(tilingKey); + context_->SetBlockDim(blockNumB_); + size_t *workspaces = context_->GetWorkspaceSizes(1); + OPS_LOG_E_IF_NULL(context_, workspaces, return ge::GRAPH_FAILED); + workspaces[0] = workspaceSize_; + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +// REGISTER_OPS_TILING_TEMPLATE(InplacePartialRotaryMul, RopeRegBaseTilingClassAAndB, ROPE_A_AND_B_TILING_PRIORITY); +} // namespace optiling diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_ab.cpp b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_ab.cpp new file mode 100644 index 000000000..38fe2648c --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_ab.cpp @@ -0,0 +1,194 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rope_regbase_tiling_ab.cc + * \brief + */ + +#include "inplace_partial_rotary_mul_tiling.h" + +namespace optiling { + +constexpr size_t RESERVERD_WORKSPACE_SIZE = static_cast(16 * 1024 * 1024); +constexpr int64_t MAX_COPY_BLOCK_COUNT = 4095; +constexpr int64_t CONST_TWO = 2; +constexpr int64_t CONST_FOUR = 4; +constexpr int64_t DB_FLAG = 2; +constexpr int64_t TILING_KEY_AB = 20030; +constexpr int64_t TILING_KEY_AB_BF16_FP32 = 20130; +constexpr int64_t TILING_KEY_AB_FP16_FP32 = 20230; + +ge::graphStatus RopeRegBaseTilingClassAB::DoOpTiling() +{ + int64_t bs = b_ * s_; + if (cosb_ == 1) { + bs = s_; + n_ = b_ * n_; + } + int64_t typeSize = ge::GetSizeByDataType(dtype_); + if (typeSize == 0) { + OPS_LOG_I("RopeRegBaseTilingClassAB DoOpTiling error, typeSize == 0"); + return ge::GRAPH_FAILED; + } + dAlign_ = CeilAlign(sliceLength_ / dSplitCoef_, blockSize_ / typeSize) * dSplitCoef_; + + auto cosDtype = context_->GetInputDesc(1)->GetDataType(); + bool isMixedPrecision = (dtype_ == ge::DT_BF16 || dtype_ == ge::DT_FLOAT16) && cosDtype == ge::DT_FLOAT; + int64_t cosTypeSize = ge::GetSizeByDataType(cosDtype); + if (cosTypeSize == 0) { + OPS_LOG_I("RopeRegBaseTilingClassAB DoOpTiling error, cosTypeSize == 0"); + return ge::GRAPH_FAILED; + } + int64_t dAlignCosSin = CeilAlign(sliceLength_ / dSplitCoef_, blockSize_ / cosTypeSize) * dSplitCoef_; + + blockFactorBS_ = CeilDiv(bs, int64_t(aicoreParams_.blockDim)); + blockNumBS_ = CeilDiv(bs, blockFactorBS_); + blockTailBS_ = bs - (blockNumBS_ - 1) * blockFactorBS_; + + if (bs <= int64_t(aicoreParams_.blockDim) / CONST_TWO) { + if (blockNumBS_ == 0) { + OPS_LOG_I("RopeRegBaseTilingClassAB ComputeUbFactor error, blockNumBS_ == 0"); + return ge::GRAPH_FAILED; + } + blockNumN_ = aicoreParams_.blockDim / blockNumBS_; + blockFactorN_ = CeilDiv(n_, blockNumN_); + blockNumN_ = CeilDiv(n_, blockFactorN_); + blockTailN_ = n_ - (blockNumN_ - 1) * blockFactorN_; + } else { + blockNumN_ = 1; + blockFactorN_ = n_; + blockTailN_ = n_; + } + + int64_t dSizeX = dAlign_ * typeSize; + int64_t dSizeCosSin = dAlignCosSin * cosTypeSize; + int64_t baseBlockInUb; + + if (isMixedPrecision) { + baseBlockInUb = + FloorAlign(static_cast(aicoreParams_.ubSize / CONST_TWO / DB_FLAG), blockSize_) / dSizeX; + // UB: 4 * ubFactorBS * (dSizeX * ubFactorN + dSizeCosSin) <= ubSize + // => ubFactorBS * (ubFactorN + dSizeCosSin/dSizeX) <= ubSize/(4*dSizeX) + int64_t effectiveCosSinOverhead = CeilDiv(dSizeCosSin, dSizeX); + OPS_ERR_IF(baseBlockInUb < effectiveCosSinOverhead + 1, + OPS_LOG_I(context_->GetNodeName(), "ubSize can't load mixed precision, d = %ld.", d_), + return ge::GRAPH_FAILED); + + ubFactorN_ = std::min(blockFactorN_, baseBlockInUb - effectiveCosSinOverhead); + ubFactorN_ = std::min(ubFactorN_, MAX_COPY_BLOCK_COUNT / dSplitCoef_); + if (ubFactorN_ <= 0) { + ubFactorN_ = 1; + } + ubFactorBS_ = std::min(FloorDiv(baseBlockInUb, ubFactorN_ + effectiveCosSinOverhead), blockFactorBS_); + ubFactorBS_ = (ubFactorBS_ == 0) ? 1 : ubFactorBS_; + } else { + int64_t baseBufferSize = dSizeX; + baseBlockInUb = + FloorAlign(static_cast(aicoreParams_.ubSize / CONST_TWO / DB_FLAG), blockSize_) / baseBufferSize; + OPS_ERR_IF(baseBlockInUb < 1, + OPS_LOG_I(context_->GetNodeName(), "ubSize can't load 8 d size, d = %ld.", d_), + return ge::GRAPH_FAILED); + + ubFactorN_ = std::min(blockFactorN_, baseBlockInUb - 1); + ubFactorN_ = std::min(ubFactorN_, MAX_COPY_BLOCK_COUNT / dSplitCoef_); + + ubFactorBS_ = std::min(FloorDiv(baseBlockInUb, ubFactorN_ + 1), blockFactorBS_); + ubFactorBS_ = (ubFactorBS_ == 0) ? 1 : ubFactorBS_; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClassAB::PostTiling() +{ + tilingData_.set_B(b_); + tilingData_.set_CosB(cosb_); + tilingData_.set_S(s_); + tilingData_.set_D(d_); + tilingData_.set_N(n_); + tilingData_.set_dAlign(dAlign_); + tilingData_.set_dSplitCoef(dSplitCoef_); + tilingData_.set_blockNumBS(blockNumBS_); + tilingData_.set_blockFactorBS(blockFactorBS_); + tilingData_.set_blockTailBS(blockTailBS_); + tilingData_.set_blockNumN(blockNumN_); + tilingData_.set_blockFactorN(blockFactorN_); + tilingData_.set_blockTailN(blockTailN_); + tilingData_.set_ubFactorBS(ubFactorBS_); + tilingData_.set_ubFactorN(ubFactorN_); + tilingData_.set_rotaryMode(static_cast(rotaryMode_)); + tilingData_.set_sliceStart(static_cast(sliceStart_)); + tilingData_.set_sliceEnd(static_cast(sliceEnd_)); + tilingData_.set_sliceLength(static_cast(sliceLength_)); + + context_->SetTilingKey(GetTilingKey()); + context_->SetBlockDim(blockNumBS_ * blockNumN_); + size_t *workspaces = context_->GetWorkspaceSizes(1); + workspaces[0] = RESERVERD_WORKSPACE_SIZE; + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + + OPS_LOG_I(context_->GetNodeName(), + "RopeRegBaseTilingClassAB tilingData is B: %ld, CosB: %ld, S: %ld, D: %ld, N: %ld, kAlign: %ld, " + "dSplitCoef: %ld, BlockNumBS: %ld, BlockFactorBS: %ld, BlockTailBS: %ld, BlockNumN: %ld, " + "BlockFactorN: %ld, BlockTailN: %ld, UBFactorBS: %ld, UBFactorN: %ld, RotaryMode: %ld, TilingKey: %ld, " + "sliceStart is %ld, sliceEnd is %ld, sliceLength is %ld", + tilingData_.get_B(), + tilingData_.get_CosB(), + tilingData_.get_S(), + tilingData_.get_D(), + tilingData_.get_N(), + tilingData_.get_dAlign(), + tilingData_.get_dSplitCoef(), + tilingData_.get_blockNumBS(), + tilingData_.get_blockFactorBS(), + tilingData_.get_blockTailBS(), + tilingData_.get_blockNumN(), + tilingData_.get_blockFactorN(), + tilingData_.get_blockTailN(), + tilingData_.get_ubFactorBS(), + tilingData_.get_ubFactorN(), + tilingData_.get_rotaryMode(), + GetTilingKey(), + tilingData_.get_sliceStart(), + tilingData_.get_sliceEnd(), + tilingData_.get_sliceLength()); + + return ge::GRAPH_SUCCESS; +} + +uint64_t RopeRegBaseTilingClassAB::GetTilingKey() const +{ + auto xDtype = context_->GetInputDesc(0)->GetDataType(); + auto cosDtype = context_->GetInputDesc(1)->GetDataType(); + if (xDtype == ge::DT_BF16 && cosDtype == ge::DT_FLOAT) { + return TILING_KEY_AB_BF16_FP32; + } else if (xDtype == ge::DT_FLOAT16 && cosDtype == ge::DT_FLOAT) { + return TILING_KEY_AB_FP16_FP32; + } + + return TILING_KEY_AB; +} + +bool RopeRegBaseTilingClassAB::IsCapable() +{ + if (!IsRegbaseSocVersion()) { + return false; + } + + OPS_LOG_I(context_->GetNodeName(), "layout: %ld", static_cast(layout_)); + // 1. qk:bsnd, cos:bs1d 2. qk:sbnd, cos:sb1d 3. qk:sbnd, cos:s11d + return layout_ == RopeLayout::SBND; +} + +// REGISTER_OPS_TILING_TEMPLATE(InplacePartialRotaryMul, RopeRegBaseTilingClassAB, 25000); + +} // namespace optiling \ No newline at end of file diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_aba_and_ba.cpp b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_aba_and_ba.cpp new file mode 100644 index 000000000..792749818 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_aba_and_ba.cpp @@ -0,0 +1,264 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rope_regbase_tiling_aba_and_ba.cc + * \brief + */ +#include "inplace_partial_rotary_mul_tiling.h" + +using namespace AscendC; + +namespace optiling { + +constexpr uint64_t TILING_KEY_ABA = 20010; +constexpr uint64_t TILING_KEY_BA = 20011; +constexpr uint64_t TILING_KEY_ABA_BF16_FP32 = 20110; +constexpr uint64_t TILING_KEY_ABA_FP16_FP32 = 20210; +constexpr uint64_t TILING_KEY_BA_BF16_FP32 = 20111; +constexpr uint64_t TILING_KEY_BA_FP16_FP32 = 20211; +constexpr int64_t UB_FACTOR = 4; +constexpr int64_t MIXED_PRECISION_X_FACTOR = 4; +constexpr int64_t MIXED_PRECISION_COS_SIN_FACTOR = 4; +constexpr int64_t MAX_COPY_BLOCK_COUNT = 4095; +constexpr int32_t WORKSPACE_SIZE = 16 * 1024 * 1024; + +bool RopeRegBaseTilingClassABAAndBA::IsCapable() +{ + // BNSD对应11SD和B1SD两种brc模式 + return (IsRegbaseSocVersion()) && (layout_ == RopeLayout::BNSD); +} + +ge::graphStatus RopeRegBaseTilingClassABAAndBA::SplitCore() +{ + // B大于等于核数，且能被核数整除，则仅在B轴分核 + if (b_ % aicoreParams_.blockDim == 0) { + blockNumB_ = aicoreParams_.blockDim; + blockFactorB_ = b_ / aicoreParams_.blockDim; + blockNumS_ = 1; + blockFactorS_ = s_; + return ge::GRAPH_SUCCESS; + } + + // S大于等于核数，且能被核数整除，则仅在S轴分核 + if (s_ % aicoreParams_.blockDim == 0) { + blockNumS_ = aicoreParams_.blockDim; + blockFactorS_ = s_ / aicoreParams_.blockDim; + blockNumB_ = 1; + blockFactorB_ = b_; + return ge::GRAPH_SUCCESS; + } + + // 尝试优先对B分核，再尝试优先对S分核，比较二者切分后的总核数 + auto blockFactorB1 = CeilDiv(static_cast(b_), aicoreParams_.blockDim); + auto blockNumB1 = CeilDiv(static_cast(b_), blockFactorB1); + if (blockNumB1 == 0) { + OPS_LOG_I("RopeRegBaseTilingClassABAAndBA SplitCore error, blockNumB1 == 0"); + return ge::GRAPH_FAILED; + } + auto blockNumS1 = std::min(static_cast(s_), aicoreParams_.blockDim / blockNumB1); + auto blockFactorS1 = CeilDiv(static_cast(s_), blockNumS1); + blockNumS1 = CeilDiv(static_cast(s_), blockFactorS1); + auto usedCoreNum1 = blockNumB1 * blockNumS1; + + auto blockFactorS2 = CeilDiv(static_cast(s_), aicoreParams_.blockDim); + auto blockNumS2 = CeilDiv(static_cast(s_), blockFactorS2); + if (blockNumS2 == 0) { + OPS_LOG_I("RopeRegBaseTilingClassABAAndBA SplitCore error, blockNumS2 == 0"); + return ge::GRAPH_FAILED; + } + auto blockNumB2 = std::min(static_cast(b_), aicoreParams_.blockDim / blockNumS2); + auto blockFactorB2 = CeilDiv(static_cast(b_), blockNumB2); + blockNumB2 = CeilDiv(static_cast(b_), blockFactorB2); + auto usedCoreNum2 = blockNumB2 * blockNumS2; + + if (usedCoreNum1 >= usedCoreNum2) { + blockNumB_ = blockNumB1; + blockFactorB_ = blockFactorB1; + blockNumS_ = blockNumS1; + blockFactorS_ = blockFactorS1; + } else { + blockNumB_ = blockNumB2; + blockFactorB_ = blockFactorB2; + blockNumS_ = blockNumS2; + blockFactorS_ = blockFactorS2; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClassABAAndBA::ComputeUbFactor() +{ + ubFactorS_ = 1; + ubFactorB_ = 1; + ubFactorN_ = 1; + + auto cosDtype = context_->GetInputDesc(1)->GetDataType(); + bool isMixedPrecision = (dtype_ == ge::DT_BF16 || dtype_ == ge::DT_FLOAT16) && cosDtype == ge::DT_FLOAT; + + int64_t dSizeX = CeilAlign(sliceLength_ * GetSizeByDataType(dtype_) / dSplitCoef_, this->blockSize_) * dSplitCoef_; + int64_t dSizeCosSin = + CeilAlign(sliceLength_ * GetSizeByDataType(cosDtype) / dSplitCoef_, this->blockSize_) * dSplitCoef_; + + int64_t totalDSize; + if (isMixedPrecision) { + totalDSize = dSizeX * MIXED_PRECISION_X_FACTOR + dSizeCosSin * MIXED_PRECISION_COS_SIN_FACTOR; + } else { + totalDSize = dSizeX * UB_FACTOR; + } + + int64_t numOfDAvailable = FloorDiv(static_cast(aicoreParams_.ubSize), totalDSize); + OPS_ERR_IF(numOfDAvailable < ubFactorB_ + 1, + OPS_LOG_E(context_, + "D is too big to load in ub, ubSize is %ld bytes, loading requires %ld bytes.", + static_cast(aicoreParams_.ubSize), + totalDSize * (ubFactorB_ + 1)), + return ge::GRAPH_FAILED); + + ubFactorS_ = std::min(blockFactorS_, FloorDiv(numOfDAvailable, ubFactorB_ + 1)); + ubFactorS_ = std::min(ubFactorS_, MAX_COPY_BLOCK_COUNT / dSplitCoef_); + if (ubFactorS_ == 0) { + OPS_LOG_I("RopeRegBaseTilingClassABAAndBA ComputeUbFactor error, ubFactorS_ == 0"); + return ge::GRAPH_FAILED; + } + numOfDAvailable /= ubFactorS_; + if (numOfDAvailable <= ubFactorB_ + 1) { + return ge::GRAPH_SUCCESS; + } + + if (cosb_ == 1) { + numOfDAvailable -= 1; + ubFactorN_ = std::min(n_, numOfDAvailable); + if (ubFactorN_ == 0) { + OPS_LOG_I("RopeRegBaseTilingClassABAAndBA ComputeUbFactor error, ubFactorN_ == 0"); + return ge::GRAPH_FAILED; + } + numOfDAvailable /= ubFactorN_; + } else { + ubFactorN_ = std::min(n_, numOfDAvailable - 1); + numOfDAvailable /= (ubFactorN_ + 1); + } + + if (numOfDAvailable <= 1) { + return ge::GRAPH_SUCCESS; + } + ubFactorB_ = std::min(blockFactorB_, numOfDAvailable); + + return ge::GRAPH_SUCCESS; +} + +void RopeRegBaseTilingClassABAAndBA::SetTilingData() +{ + tilingData_.set_B(b_); + tilingData_.set_CosB(cosb_); + tilingData_.set_S(s_); + tilingData_.set_D(d_); + tilingData_.set_N(n_); + tilingData_.set_blockNumB(blockNumB_); + tilingData_.set_blockFactorB(blockFactorB_); + tilingData_.set_blockNumS(blockNumS_); + tilingData_.set_blockFactorS(blockFactorS_); + tilingData_.set_ubLoopNumS(ubLoopNumS_); + tilingData_.set_ubFactorS(ubFactorS_); + tilingData_.set_ubTailFactorS(ubTailFactorS_); + tilingData_.set_ubLoopNumB(ubLoopNumB_); + tilingData_.set_ubFactorB(ubFactorB_); + tilingData_.set_ubTailFactorB(ubTailFactorB_); + tilingData_.set_ubLoopNumN(ubLoopNumN_); + tilingData_.set_ubFactorN(ubFactorN_); + tilingData_.set_ubTailFactorN(ubTailFactorN_); + tilingData_.set_rotaryMode(static_cast(rotaryMode_)); + tilingData_.set_sliceStart(static_cast(sliceStart_)); + tilingData_.set_sliceEnd(static_cast(sliceEnd_)); + tilingData_.set_sliceLength(static_cast(sliceLength_)); + + OPS_LOG_I(context_->GetNodeName(), + "RopeRegBaseTilingClassABAAndBA tilingData: " + "B is %ld, CosB is %ld, S is %ld, D is %ld, N is %ld, blockNumB %ld," + "blockFactorB_ is %ld, blockNumS %ld, blockFactorS is %ld, ubLoopNumS is %ld," + "ubFactorS is %ld, ubTailFactorS %ld, ubLoopNumB is %ld, ubFactorB is %ld," + "ubTailFactorB is %ld, ubLoopNumN is %ld, ubFactorN is %ld, ubTailFactorN is %ld," + "rotaryMode is %ld, tilingKey is %ld, sliceStart is %ld, sliceEnd is %ld, sliceLength is %ld", + tilingData_.get_B(), + tilingData_.get_CosB(), + tilingData_.get_S(), + tilingData_.get_D(), + tilingData_.get_N(), + tilingData_.get_blockNumB(), + tilingData_.get_blockFactorB(), + tilingData_.get_blockNumS(), + tilingData_.get_blockFactorS(), + tilingData_.get_ubLoopNumS(), + tilingData_.get_ubFactorS(), + tilingData_.get_ubTailFactorS(), + tilingData_.get_ubLoopNumB(), + tilingData_.get_ubFactorB(), + tilingData_.get_ubTailFactorB(), + tilingData_.get_ubLoopNumN(), + tilingData_.get_ubFactorN(), + tilingData_.get_ubTailFactorN(), + tilingData_.get_rotaryMode(), + GetTilingKey(), + tilingData_.get_sliceStart(), + tilingData_.get_sliceEnd(), + tilingData_.get_sliceLength()); +} + +ge::graphStatus RopeRegBaseTilingClassABAAndBA::DoOpTiling() +{ + OPS_ERR_IF(SplitCore() != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "failed to split core."), + return ge::GRAPH_FAILED); + OPS_ERR_IF(ComputeUbFactor() != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "failed to compute ub factor."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClassABAAndBA::DoLibApiTiling() +{ + return ge::GRAPH_SUCCESS; +} + +uint64_t RopeRegBaseTilingClassABAAndBA::GetTilingKey() const +{ + auto xDtype = context_->GetInputDesc(0)->GetDataType(); + auto cosDtype = context_->GetInputDesc(1)->GetDataType(); + if (xDtype == ge::DT_BF16 && cosDtype == ge::DT_FLOAT) { + return (cosb_ == 1) ? TILING_KEY_BA_BF16_FP32 : TILING_KEY_ABA_BF16_FP32; + } else if (xDtype == ge::DT_FLOAT16 && cosDtype == ge::DT_FLOAT) { + return (cosb_ == 1) ? TILING_KEY_BA_FP16_FP32 : TILING_KEY_ABA_FP16_FP32; + } + + return (cosb_ == 1) ? TILING_KEY_BA : TILING_KEY_ABA; +} + +ge::graphStatus RopeRegBaseTilingClassABAAndBA::GetWorkspaceSize() +{ + workspaceSize_ = WORKSPACE_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClassABAAndBA::PostTiling() +{ + SetTilingData(); + uint64_t tilingKey = GetTilingKey(); + context_->SetTilingKey(tilingKey); + context_->SetBlockDim(blockNumB_ * blockNumS_); + size_t *workspaces = context_->GetWorkspaceSizes(1); + OPS_LOG_E_IF_NULL(context_, workspaces, return ge::GRAPH_FAILED); + workspaces[0] = workspaceSize_; + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +// REGISTER_OPS_TILING_TEMPLATE(InplacePartialRotaryMul, RopeRegBaseTilingClassABAAndBA, +// ROPE_ABA_AND_BA_TILING_PRIORITY); +} // namespace optiling \ No newline at end of file diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_bab.cpp b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_bab.cpp new file mode 100644 index 000000000..fac882a6f --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_bab.cpp @@ -0,0 +1,215 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file rope_regbase_tiling_bab.cc + * \brief + */ + +#include "inplace_partial_rotary_mul_tiling.h" + +namespace optiling { +constexpr uint64_t ROPE_BAB_TILING_PRIORITY = 20000; +constexpr uint32_t MIN_UB_LOAD_D_NUM = 4; // x, y或in, cos输入开doubleBuffer +constexpr uint32_t DOUBLE_BUFFER = 2; +constexpr int64_t MIN_COPY_BLOCK_COUNT = 4095; +constexpr size_t WORK_SPACE_SIZE = static_cast(16) * 1024 * 1024; +constexpr int64_t TILING_KEY_BAB = 20020; +constexpr int64_t TILING_KEY_BAB_BF16_FP32 = 20120; +constexpr int64_t TILING_KEY_BAB_FP16_FP32 = 20220; + +ge::graphStatus RopeRegBaseTilingClassBAB::DoOpTiling() +{ + ubSize_ = aicoreParams_.ubSize; + coreNum_ = aicoreParams_.blockDim; + ge::graphStatus status = SplitUb(); + if (status != ge::GRAPH_SUCCESS) { + OPS_LOG_E(context_->GetNodeName(), "SplitUb Failed."); + return ge::GRAPH_FAILED; + } + status = SplitCore(); + if (status != ge::GRAPH_SUCCESS) { + OPS_LOG_E(context_->GetNodeName(), "SplitCore Failed."); + return ge::GRAPH_FAILED; + } + if (blockNumB_ * blockNumS_ > coreNum_) { + OPS_LOG_E( + context_->GetNodeName(), "split coreNum [%ld] large than coreNum[%ld]", blockNumB_ * blockNumS_, coreNum_); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClassBAB::SplitCore() +{ + // 尝试先对B分核，再尝试优先对S分核，比较二者切分之后的总核数 + auto blockFactorB1 = CeilDiv(b_, coreNum_); + auto blockNumB1 = CeilDiv(b_, blockFactorB1); + if (blockNumB1 == 0) { + OPS_LOG_I("RopeRegBaseTilingClassBAB SplitCore error, blockNumB1 == 0"); + return ge::GRAPH_FAILED; + } + auto blockNumS1 = std::min(coreNum_ / blockNumB1, s_); + auto blockFactorS1 = CeilDiv(s_, blockNumS1); + blockNumS1 = CeilDiv(s_, blockFactorS1); + auto usedCoreNum1 = blockNumB1 * blockNumS1; + + auto blockFactorS2 = CeilDiv(s_, coreNum_); + auto blockNumS2 = CeilDiv(s_, blockFactorS2); + if (blockNumS2 == 0) { + OPS_LOG_I("RopeRegBaseTilingClassBAB SplitCore error, blockNumS2 == 0"); + return ge::GRAPH_FAILED; + } + auto blockNumB2 = std::min(coreNum_ / blockNumS2, b_); + auto blockFactorB2 = CeilDiv(b_, blockNumB2); + blockNumB2 = CeilDiv(b_, blockFactorB2); + auto usedCoreNum2 = blockNumB2 * blockNumS2; + + // ubFactorS 很小的时候，选择核数多的，ubFactorS大于blockFactorS时， 综合考虑分核和UB切分 + auto ubFactorS1 = std::min(ubFactorS_, blockFactorS1); + auto ubFactorS2 = std::min(ubFactorS_, blockFactorS2); + if (usedCoreNum1 * ubFactorS1 >= usedCoreNum2 * ubFactorS2) { + blockNumB_ = blockNumB1; + blockFactorB_ = blockFactorB1; + blockNumS_ = blockNumS1; + blockFactorS_ = blockFactorS1; + usedCoreNum_ = usedCoreNum1; + ubFactorS_ = ubFactorS1; + } else { + blockNumB_ = blockNumB2; + blockFactorB_ = blockFactorB2; + blockNumS_ = blockNumS2; + blockFactorS_ = blockFactorS2; + usedCoreNum_ = usedCoreNum2; + ubFactorS_ = ubFactorS2; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClassBAB::SplitUb() +{ + uint32_t typeSize = ge::GetSizeByDataType(dtype_); + auto cosDtype = context_->GetInputDesc(1)->GetDataType(); + uint32_t cosTypeSize = ge::GetSizeByDataType(cosDtype); + + // For mixed precision (x is BF16/FP16, cos/sin is FP32): + // UB needs to store: x(dtype_), cos(float32), sin(float32), output(dtype_) + // In mixed precision case, cos/sin use 4 bytes, x uses 2 bytes + int64_t dAlignX = CeilAlign(sliceLength_ * typeSize / dSplitCoef_, blockSize_) * dSplitCoef_; + int64_t dAlignCosSin = CeilAlign(sliceLength_ * cosTypeSize / dSplitCoef_, blockSize_) * dSplitCoef_; + + // Total buffer needed per element: x buffer + cos buffer + sin buffer + output buffer + // For double buffer mode, need to multiply by DOUBLE_BUFFER + int64_t totalDAlign = dAlignX * 2 + dAlignCosSin * 2; // x/y queue double buffer + cos/sin queue double buffer + + int64_t canLoadDNum = FloorDiv(ubSize_, totalDAlign); + if (canLoadDNum < MIN_UB_LOAD_D_NUM) { + OPS_LOG_E(context_->GetNodeName(), "ubSize_ can't load enough d_, d_ = %ld.", d_); + return ge::GRAPH_FAILED; + } + canLoadDNum = canLoadDNum / MIN_UB_LOAD_D_NUM; + int64_t ubLoopNum = CeilDiv(n_, (canLoadDNum - 1)); + ubFactorN_ = std::min(CeilDiv(n_, ubLoopNum), MIN_COPY_BLOCK_COUNT / dSplitCoef_); + ubLoopNumN_ = CeilDiv(n_, ubFactorN_); + if (ubFactorN_ == 0) { + OPS_LOG_I("RopeRegBaseTilingClassBAB SplitUb error, ubFactorN_ == 0"); + return ge::GRAPH_FAILED; + } + ubTailFactorN_ = (n_ % ubFactorN_ == 0) ? ubFactorN_ : n_ % ubFactorN_; + int64_t ubFactorS = FloorDiv(canLoadDNum, n_ + 1); + ubFactorS_ = (ubFactorS == 0) ? 1 : ubFactorS; + return ge::GRAPH_SUCCESS; +} + +void RopeRegBaseTilingClassBAB::PrintTilingData() +{ + OPS_LOG_I(context_->GetNodeName(), + "RopeRegBaseTilingClassBAB tilingData: useCoreNum is %ld," + "B is %ld, CosB is %ld, S is %ld, D is %ld, N is %ld, blockNumB %ld," + "blockFactorB_ is %ld, blockNumS %ld, blockFactorS is %ld, ubLoopNumS is %ld," + "ubFactorS is %ld, ubTailFactorS %ld, ubLoopNumB is %ld, ubFactorB is %ld," + "ubTailFactorB is %ld, ubLoopNumN is %ld, ubFactorN is %ld, ubTailFactorN is %ld," + "rotaryMode is %ld, tilingKey is %ld, sliceStart is %ld, sliceEnd is %ld, sliceLength is %ld", + usedCoreNum_, + tilingData_.get_B(), + tilingData_.get_CosB(), + tilingData_.get_S(), + tilingData_.get_D(), + tilingData_.get_N(), + tilingData_.get_blockNumB(), + tilingData_.get_blockFactorB(), + tilingData_.get_blockNumS(), + tilingData_.get_blockFactorS(), + tilingData_.get_ubLoopNumS(), + tilingData_.get_ubFactorS(), + tilingData_.get_ubTailFactorS(), + tilingData_.get_ubLoopNumB(), + tilingData_.get_ubFactorB(), + tilingData_.get_ubTailFactorB(), + tilingData_.get_ubLoopNumN(), + tilingData_.get_ubFactorN(), + tilingData_.get_ubTailFactorN(), + tilingData_.get_rotaryMode(), + tilingKey_, + tilingData_.get_sliceStart(), + tilingData_.get_sliceEnd(), + tilingData_.get_sliceLength()); + return; +} + +ge::graphStatus RopeRegBaseTilingClassBAB::PostTiling() +{ + tilingData_.set_B(b_); + tilingData_.set_CosB(0); + tilingData_.set_S(s_); + tilingData_.set_D(d_); + tilingData_.set_N(n_); + tilingData_.set_blockNumB(blockNumB_); + tilingData_.set_blockFactorB(blockFactorB_); + tilingData_.set_blockNumS(blockNumS_); + tilingData_.set_blockFactorS(blockFactorS_); + tilingData_.set_ubLoopNumS(ubLoopNumS_); + tilingData_.set_ubFactorS(ubFactorS_); + tilingData_.set_ubTailFactorS(ubTailFactorS_); + tilingData_.set_ubLoopNumB(ubLoopNumB_); + tilingData_.set_ubFactorB(ubFactorB_); + tilingData_.set_ubTailFactorB(ubTailFactorB_); + tilingData_.set_ubLoopNumN(ubLoopNumN_); + tilingData_.set_ubFactorN(ubFactorN_); + tilingData_.set_ubTailFactorN(ubTailFactorN_); + tilingData_.set_rotaryMode(static_cast(rotaryMode_)); + tilingData_.set_sliceStart(static_cast(sliceStart_)); + tilingData_.set_sliceEnd(static_cast(sliceEnd_)); + tilingData_.set_sliceLength(static_cast(sliceLength_)); + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + context_->SetBlockDim(usedCoreNum_); + context_->SetTilingKey(tilingKey_); + size_t *workspaces = context_->GetWorkspaceSizes(1); + workspaces[0] = WORK_SPACE_SIZE; + PrintTilingData(); + return ge::GRAPH_SUCCESS; +} + +uint64_t RopeRegBaseTilingClassBAB::GetTilingKey() const +{ + auto xDtype = context_->GetInputDesc(0)->GetDataType(); + auto cosDtype = context_->GetInputDesc(1)->GetDataType(); + if (xDtype == ge::DT_BF16 && cosDtype == ge::DT_FLOAT) { + return TILING_KEY_BAB_BF16_FP32; + } else if (xDtype == ge::DT_FLOAT16 && cosDtype == ge::DT_FLOAT) { + return TILING_KEY_BAB_FP16_FP32; + } + + return TILING_KEY_BAB; +} + +// REGISTER_OPS_TILING_TEMPLATE(InplacePartialRotaryMul, RopeRegBaseTilingClassBAB, ROPE_BAB_TILING_PRIORITY); +} // namespace optiling diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_base.cpp b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_base.cpp new file mode 100644 index 000000000..5109f8de8 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_host/rope_regbase_tiling_base.cpp @@ -0,0 +1,346 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rope_regbase_tiling_base.cc + * \brief + */ + +// #include "tiling_base/tiling_templates_registry.h" +#include "tiling/tiling_api.h" +// #include "tiling_base/tiling_base.h" +#include "platform/platform_info.h" +#include "inplace_partial_rotary_mul_tiling.h" +#include +// #include "log/log.h" + +namespace { +constexpr int64_t X_INDEX = 0; +constexpr int64_t COS_INDEX = 1; +constexpr int64_t SIN_INDEX = 2; +constexpr int64_t Y_INDEX = 0; +constexpr int64_t DIM_NUM = 4; +constexpr int64_t DIM_0 = 0; +constexpr int64_t DIM_1 = 1; +constexpr int64_t DIM_2 = 2; +constexpr int64_t DIM_3 = 3; +constexpr int64_t HALF_INTERLEAVE_MODE_COEF = 2; +constexpr int64_t QUARTER_MODE_COEF = 4; +constexpr int64_t BLOCK_SIZE = 32; +constexpr int64_t D_LIMIT = 1024; +const std::vector SUPPORT_DTYPE = {ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16}; +} // namespace + +namespace optiling { +ge::graphStatus RopeRegBaseTilingClass::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + if (platformInfo != nullptr) { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + aicoreParams_.blockDim = ascendcPlatform.GetCoreNumAiv(); + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + aicoreParams_.ubSize = ubSizePlatForm; + socVersion_ = ascendcPlatform.GetSocVersion(); + } else { + auto compileInfoPtr = reinterpret_cast(context_->GetCompileInfo()); + OPS_ERR_IF(compileInfoPtr == nullptr, OPS_LOG_E(context_, "compile info is null"), return ge::GRAPH_FAILED); + aicoreParams_.blockDim = compileInfoPtr->blockDim; + aicoreParams_.ubSize = compileInfoPtr->ubSize; + socVersion_ = compileInfoPtr->socVersion; + } + blockSize_ = BLOCK_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClass::CheckNullptr() +{ + for (int64_t i = 0; i <= SIN_INDEX; i++) { + auto desc = context_->GetInputDesc(i); + OPS_ERR_IF(desc == nullptr, OPS_LOG_E(context_, "input %ld desc is nullptr.", i), return ge::GRAPH_FAILED); + auto shape = context_->GetInputShape(i); + OPS_ERR_IF(shape == nullptr, OPS_LOG_E(context_, "input %ld shape is nullptr.", i), return ge::GRAPH_FAILED); + } + auto yDesc = context_->GetOutputDesc(Y_INDEX); + OPS_ERR_IF(yDesc == nullptr, OPS_LOG_E(context_, "output desc is nullptr."), return ge::GRAPH_FAILED); + auto yShape = context_->GetOutputShape(Y_INDEX); + OPS_ERR_IF(yShape == nullptr, OPS_LOG_E(context_, "output shape is nullptr."), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClass::CheckShapeAllPositive(const int64_t idx) const +{ + auto shape = context_->GetInputShape(idx)->GetStorageShape(); + for (size_t i = 0; i < shape.GetDimNum(); i++) { + OPS_ERR_IF( + shape.GetDim(i) <= 0, + OPS_LOG_E(context_, "input %ld has non positive shape, dim %lu actual %ld .", idx, i, shape.GetDim(i)), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} +bool RopeRegBaseTilingClass::IsRotaryPosEmbeddingMode(const int32_t mode) const +{ + switch (mode) { + case static_cast(RotaryPosEmbeddingMode::HALF): + case static_cast(RotaryPosEmbeddingMode::INTERLEAVE): + case static_cast(RotaryPosEmbeddingMode::QUARTER): + case static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE): + return true; + default: + return false; + } +} + +ge::graphStatus RopeRegBaseTilingClass::CheckShapeAllPositive() const +{ + OPS_ERR_IF(CheckShapeAllPositive(X_INDEX) != ge::GRAPH_SUCCESS, OPS_LOG_E(context_, "x has non positive shape."), + return ge::GRAPH_FAILED); + OPS_ERR_IF(CheckShapeAllPositive(COS_INDEX) != ge::GRAPH_SUCCESS, OPS_LOG_E(context_, "cos has non positive shape."), + return ge::GRAPH_FAILED); + OPS_ERR_IF(CheckShapeAllPositive(SIN_INDEX) != ge::GRAPH_SUCCESS, OPS_LOG_E(context_, "sin has non positive shape."), + return ge::GRAPH_FAILED); + auto yShape = context_->GetOutputShape(Y_INDEX)->GetStorageShape(); + for (size_t i = 0; i < yShape.GetDimNum(); i++) { + OPS_ERR_IF(yShape.GetDim(i) <= 0, + OPS_LOG_E(context_, "output has non positive shape, dim %lu actual %ld .", i, yShape.GetDim(i)), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClass::JudgeLayoutByShape(const gert::Shape &xShape, const gert::Shape &cosShape) +{ + uint64_t xShape0 = xShape.GetDim(DIM_0); + uint64_t xShape1 = xShape.GetDim(DIM_1); + uint64_t xShape2 = xShape.GetDim(DIM_2); + uint64_t cosShape0 = cosShape.GetDim(DIM_0); + uint64_t cosShape1 = cosShape.GetDim(DIM_1); + uint64_t cosShape2 = cosShape.GetDim(DIM_2); + if (xShape0 == cosShape0 && xShape1 == cosShape1 && xShape2 == cosShape2) { // BSND + layout_ = RopeLayout::NO_BROADCAST; + } else if (cosShape0 == 1 && cosShape1 == 1 && cosShape2 == 1) { // (111D) + layout_ = RopeLayout::BROADCAST_BSN; + } else if (cosShape2 == 1 && cosShape0 == 1 && xShape1 == cosShape1) { // BSND (1S1D) + layout_ = RopeLayout::BSND; + } else if (cosShape2 == 1 && xShape0 == cosShape0 && (cosShape1 == 1 || cosShape1 == xShape1)) { // SBND (S11D, + // SB1D), BSND + // (BS1D) + layout_ = RopeLayout::SBND; + } else if (cosShape1 == 1 && xShape2 == cosShape2 && (cosShape0 == 1 || cosShape0 == xShape0)) { // BNSD (11SD, + // B1SD) + layout_ = RopeLayout::BNSD; + } else if (cosShape0 == 1 && xShape1 == cosShape1 && xShape2 == cosShape2) { // 1SND + layout_ = RopeLayout::BNSD; + is1snd_ = true; + } else { + OPS_LOG_E(context_->GetNodeName(), "the shape of x and sin not satisfy the broadcast."); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClass::CheckShape() +{ + auto &xShape = context_->GetInputShape(X_INDEX)->GetStorageShape(); + auto &cosShape = context_->GetInputShape(COS_INDEX)->GetStorageShape(); + auto &sinShape = context_->GetInputShape(SIN_INDEX)->GetStorageShape(); + auto &yShape = context_->GetOutputShape(Y_INDEX)->GetStorageShape(); + OPS_ERR_IF(xShape.GetDimNum() != DIM_NUM, OPS_LOG_E(context_, "dim of x expect 4, actual %lu.", xShape.GetDimNum()), + return ge::GRAPH_FAILED); + OPS_ERR_IF(cosShape.GetDimNum() != DIM_NUM, + OPS_LOG_E(context_, "dim of cos expect 4, actual %lu.", cosShape.GetDimNum()), return ge::GRAPH_FAILED); + OPS_ERR_IF(sinShape.GetDimNum() != DIM_NUM, + OPS_LOG_E(context_, "dim of sin expect 4, actual %lu.", sinShape.GetDimNum()), return ge::GRAPH_FAILED); + OPS_ERR_IF(yShape.GetDimNum() != DIM_NUM, + OPS_LOG_E(context_, "dim of output expect 4, actual %lu.", yShape.GetDimNum()), return ge::GRAPH_FAILED); + OPS_ERR_IF(cosShape != sinShape, + OPS_LOG_E(context_, + "shape of cos and sin should be same, actual cos shape is (%ld, %ld, %ld, %ld), sin shape is " + "(%ld, %ld, %ld, %ld). ", + cosShape.GetDim(DIM_0), cosShape.GetDim(DIM_1), cosShape.GetDim(DIM_2), cosShape.GetDim(DIM_3), + sinShape.GetDim(DIM_0), sinShape.GetDim(DIM_1), sinShape.GetDim(DIM_2), sinShape.GetDim(DIM_3)), + return ge::GRAPH_FAILED); + OPS_ERR_IF(xShape != yShape, + OPS_LOG_E(context_, + "shape of x and output should be same, actual x shape is (%ld, " + "%ld, %ld, %ld), output shape is (%ld, %ld, %ld, %ld). ", + xShape.GetDim(DIM_0), xShape.GetDim(DIM_1), xShape.GetDim(DIM_2), xShape.GetDim(DIM_3), + yShape.GetDim(DIM_0), yShape.GetDim(DIM_1), yShape.GetDim(DIM_2), yShape.GetDim(DIM_3)), + return ge::GRAPH_FAILED); + // OPS_ERR_IF( + // (cosShape.GetDim(DIM_3) != xShape.GetDim(DIM_3)), + // OPS_LOG_E(context_, + // "D of x, cos, sin and output should be same, actual x is %ld, cos is %ld, sin is %ld, output is %ld. ", + // xShape.GetDim(DIM_3), cosShape.GetDim(DIM_3), sinShape.GetDim(DIM_3), yShape.GetDim(DIM_3)), + // return ge::GRAPH_FAILED); + OPS_ERR_IF(CheckRotaryModeShapeRelation(xShape.GetDim(DIM_3)) != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_, "D is invalid for rotary mode."), return ge::GRAPH_FAILED); + return CheckShapeAllPositive(); +} + +ge::graphStatus RopeRegBaseTilingClass::CheckDtypeAndAttr() +{ + dtype_ = context_->GetInputDesc(X_INDEX)->GetDataType(); + OPS_ERR_IF(std::find(SUPPORT_DTYPE.begin(), SUPPORT_DTYPE.end(), dtype_) == SUPPORT_DTYPE.end(), + OPS_LOG_E(context_->GetNodeName(), "Only support F32, BF16, F16 datetype for x, actual %s.", + ge::TypeUtils::DataTypeToSerialString(dtype_).c_str()), + return ge::GRAPH_FAILED); + + auto cosType = context_->GetInputDesc(COS_INDEX)->GetDataType(); + auto sinType = context_->GetInputDesc(SIN_INDEX)->GetDataType(); + + // Check cos/sin dtype: same type, and must be F32, BF16, or F16 + OPS_ERR_IF(cosType != sinType, + OPS_LOG_E(context_, "cos and sin datatype should be same, cos is %s, sin is %s.", + ge::TypeUtils::DataTypeToSerialString(cosType).c_str(), + ge::TypeUtils::DataTypeToSerialString(sinType).c_str()), + return ge::GRAPH_FAILED); + OPS_ERR_IF(std::find(SUPPORT_DTYPE.begin(), SUPPORT_DTYPE.end(), cosType) == SUPPORT_DTYPE.end(), + OPS_LOG_E(context_->GetNodeName(), "Only support F32, BF16, F16 datetype for cos/sin, actual %s.", + ge::TypeUtils::DataTypeToSerialString(cosType).c_str()), + return ge::GRAPH_FAILED); + + // Mixed precision: x is BF16/FP16, cos/sin are FP32 + bool isMixedPrecision = (dtype_ == ge::DT_BF16 || dtype_ == ge::DT_FLOAT16) && cosType == ge::DT_FLOAT; + bool isSamePrecision = (dtype_ == cosType); + + OPS_ERR_IF(!isSamePrecision && !isMixedPrecision, + OPS_LOG_E(context_, "Unsupported dtype combination: x=%s, cos=%s. " + "Supported: same type, or x=BF16/FP16 with cos/sin=FP32.", + ge::TypeUtils::DataTypeToSerialString(dtype_).c_str(), + ge::TypeUtils::DataTypeToSerialString(cosType).c_str()), + return ge::GRAPH_FAILED); + + auto outputType = context_->GetOutputDesc(Y_INDEX)->GetDataType(); + OPS_ERR_IF(outputType != dtype_, + OPS_LOG_E(context_, "output datatype expect %s, actual %s.", + ge::TypeUtils::DataTypeToSerialString(dtype_).c_str(), + ge::TypeUtils::DataTypeToSerialString(outputType).c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClass::CheckParam() +{ + auto platformInfo = context_->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, OPS_LOG_E(context_, "platform info is nullptr."), return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + if (!IsRegbaseSocVersion()) { + return ge::GRAPH_SUCCESS; + } + OPS_ERR_IF(CheckNullptr() != ge::GRAPH_SUCCESS, OPS_LOG_E(context_, "check nullptr fail."), return ge::GRAPH_FAILED); + OPS_ERR_IF(CheckDtypeAndAttr() != ge::GRAPH_SUCCESS, OPS_LOG_E(context_, "check dtype and attr fail."), + return ge::GRAPH_FAILED); + OPS_ERR_IF(CheckShape() != ge::GRAPH_SUCCESS, OPS_LOG_E(context_, "check shape fail."), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClass::CheckRotaryModeShapeRelation(const int64_t d) +{ + OPS_ERR_IF(d > D_LIMIT, OPS_LOG_E(context_, "D must be small than %ld, actual %ld.", D_LIMIT, d), + return ge::GRAPH_FAILED); + if (rotaryMode_ == RotaryPosEmbeddingMode::HALF || rotaryMode_ == RotaryPosEmbeddingMode::INTERLEAVE || + rotaryMode_ == RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE) { + OPS_ERR_IF( + d % HALF_INTERLEAVE_MODE_COEF != 0, + OPS_LOG_E(context_, "D must be multiples of 2 in half, interleave and interleave-half mode, actual %ld.", d), + return ge::GRAPH_FAILED); + } else if (rotaryMode_ == RotaryPosEmbeddingMode::QUARTER) { + OPS_ERR_IF(d % QUARTER_MODE_COEF != 0, + OPS_LOG_E(context_, "D must be multiples of 4 in quarter mode, actual %ld.", d), + return ge::GRAPH_FAILED); + } + if (rotaryMode_ == RotaryPosEmbeddingMode::HALF || rotaryMode_ == RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE) { + dSplitCoef_ = HALF_INTERLEAVE_MODE_COEF; + } else if (rotaryMode_ == RotaryPosEmbeddingMode::QUARTER) { + dSplitCoef_ = QUARTER_MODE_COEF; + } else { + dSplitCoef_ = 1; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClass::JudgeSliceInfo() { + if (sliceStart_ < 0 || sliceEnd_ < 0 || sliceLength_ <= 0 || sliceEnd_ > d_) { + OPS_LOG_E(context_, "slice info fail, sliceStart_ = %ld. sliceEnd_ = %ld", sliceStart_, sliceEnd_); + return ge::GRAPH_FAILED; + } + if (cosd_ != sind_ || cosd_ != sliceLength_) { + OPS_LOG_E(context_, "slice info fail, sliceLength_ = %ld. cosd_ = %ld", sliceLength_, cosd_); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RopeRegBaseTilingClass::GetShapeAttrsInfo() +{ + const gert::RuntimeAttrs *attrs = context_->GetAttrs(); + OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED); + const int32_t *mode = attrs->GetAttrPointer(0); + int32_t modeValue = (mode == nullptr) ? 0 : static_cast(*mode); + OPS_ERR_IF(IsRotaryPosEmbeddingMode(modeValue) != true, + OPS_LOG_E(context_->GetNodeName(), "mode only support 0, 1, 2 3, actual %d.", modeValue), + return ge::GRAPH_FAILED); + rotaryMode_ = static_cast(modeValue); + + OPS_ERR_IF(CheckParam() != ge::GRAPH_SUCCESS, OPS_LOG_E(context_, "check param fail."), return ge::GRAPH_FAILED); + + dtype_ = context_->GetInputDesc(X_INDEX)->GetDataType(); + auto &xShape = context_->GetInputShape(X_INDEX)->GetStorageShape(); + auto &cosShape = context_->GetInputShape(COS_INDEX)->GetStorageShape(); + auto &sinShape = context_->GetInputShape(SIN_INDEX)->GetStorageShape(); + OPS_ERR_IF(JudgeLayoutByShape(xShape, cosShape) != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_, "JudgeLayoutByShape fail."), return ge::GRAPH_FAILED); + + d_ = xShape.GetDim(DIM_3); + cosd_ = cosShape.GetDim(DIM_3); + sind_ = sinShape.GetDim(DIM_3); + if (layout_ == RopeLayout::BSND) { + b_ = xShape.GetDim(DIM_0); + cosb_ = cosShape.GetDim(DIM_0); + s_ = xShape.GetDim(DIM_1); + n_ = xShape.GetDim(DIM_2); + } else if (layout_ == RopeLayout::BNSD || layout_ == RopeLayout::NO_BROADCAST || + layout_ == RopeLayout::BROADCAST_BSN) { + b_ = xShape.GetDim(DIM_0); + cosb_ = cosShape.GetDim(DIM_0); + n_ = xShape.GetDim(DIM_1); + s_ = xShape.GetDim(DIM_2); + // 1XXX情况下，reshape成11XX + if (is1snd_ == true) { + s_ = s_ * n_; + n_ = 1; + } + } else if (layout_ == RopeLayout::SBND) { + s_ = xShape.GetDim(DIM_0); + b_ = xShape.GetDim(DIM_1); + cosb_ = cosShape.GetDim(DIM_1); + n_ = xShape.GetDim(DIM_2); + } + + // 获取slice + const gert::ContinuousVector *sliceRangeListPtr = attrs->GetAttrPointer(1); + if (sliceRangeListPtr->GetSize() == 0) { + sliceStart_ = 0; + sliceEnd_ = d_; + } + else { + const int64_t *expertRangeList = reinterpret_cast(sliceRangeListPtr->GetData()); + sliceStart_ = expertRangeList[0]; + sliceEnd_ = expertRangeList[1]; + } + sliceLength_ = sliceEnd_ - sliceStart_; + OPS_ERR_IF(JudgeSliceInfo() != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_, "JudgeSliceInfo fail."), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +} // namespace optiling diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/apply_rotary_pos_emb_common.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/apply_rotary_pos_emb_common.h new file mode 100644 index 000000000..d630741fb --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/apply_rotary_pos_emb_common.h @@ -0,0 +1,1009 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file apply_rotary_pos_emb_common.h + * \brief + */ +#ifndef APPLY_ROTARY_POS_EMB_COMMON_H +#define APPLY_ROTARY_POS_EMB_COMMON_H + +#include "inplace_partial_rotary_mul_common.h" + +using namespace AscendC; + +__aicore__ inline constexpr uint32_t GetVRegSize() +{ +#if defined(__DAV_C310__) + return AscendC::VECTOR_REG_WIDTH; +#else + return 256U; +#endif +} + +__aicore__ inline constexpr uint32_t GetUbBlockSize() +{ + return 32U; +} + +constexpr uint32_t VL_FLOAT32_SIZE = GetVRegSize() / sizeof(float); +constexpr uint32_t VL_FLOAT16_SIZE = GetVRegSize() / sizeof(half); +constexpr uint32_t BLOCK_TYPE_SIZE = GetUbBlockSize(); +constexpr uint32_t HALF_INTERLEAVE_COEF = 2; +constexpr uint32_t QUARTER_MODE_COEF = 4; +constexpr uint32_t DOUBLE_BUFFER = 2; + +struct RotaryPosEmbeddingMode { + static constexpr int64_t HALF = 0; + static constexpr int64_t INTERLEAVE = 1; + static constexpr int64_t QUARTER = 2; + static constexpr int64_t DEEPSEEK_INTERLEAVE = 3; +}; + +/* + qOut[0] = q[0] * cos[0] - q[1] * sin[0] + qOut[1] = q[1] * cos[1] + q[0] * sin[1] +*/ +template +__aicore__ inline void HalfAlignVF(const LocalTensor &sinTensor, const LocalTensor &cosTensor, + const LocalTensor &inTensor, const LocalTensor &outTensor, uint32_t dLen, uint32_t dAlign, uint16_t currSNum, + uint16_t currDNum) +{ + __local_mem__ T *sinUb = (__local_mem__ T *)sinTensor.GetPhyAddr(); + __local_mem__ T *cosUb = (__local_mem__ T *)cosTensor.GetPhyAddr(); + __local_mem__ T *inUb = (__local_mem__ T *)inTensor.GetPhyAddr(); + __local_mem__ T *outUb = (__local_mem__ T *)outTensor.GetPhyAddr(); + uint32_t halfD = dLen / HALF_INTERLEAVE_COEF; + uint32_t halfDAlign = ops::CeilAlign(halfD, static_cast(BLOCK_TYPE_SIZE / sizeof(T))); + uint16_t repeatTimes = ops::CeilDiv(halfD, VL_FLOAT32_SIZE); + __local_mem__ T *currInUb; + __local_mem__ T *currOutUb; + __local_mem__ T *currSinUb; + __local_mem__ T *currCosUb; + + __VEC_SCOPE__ + { + MicroAPI::RegTensor vregIn; + MicroAPI::RegTensor vregHalfIn; + MicroAPI::RegTensor vregSin; + MicroAPI::RegTensor vregHalfSin; + MicroAPI::RegTensor vregCos; + MicroAPI::RegTensor vregHalfCos; + MicroAPI::RegTensor vregOut; + MicroAPI::RegTensor vregHalfOut; + MicroAPI::MaskReg preg; + for (uint16_t sIdx = 0; sIdx < currSNum; sIdx++) { + currSinUb = sinUb + sIdx * dAlign; + currCosUb = cosUb + sIdx * dAlign; + for (uint16_t row = 0; row < currDNum; row++) { + currInUb = inUb + (sIdx * currDNum + row) * dAlign; + currOutUb = outUb + (sIdx * currDNum + row) * dAlign; + uint32_t updateCnt = halfD; + for (uint16_t i = 0; i < repeatTimes; i++) { + preg = MicroAPI::UpdateMask(updateCnt); + uint32_t offset = i * VL_FLOAT32_SIZE; + uint32_t halfOffset = offset + halfDAlign; + ops::LoadTwoTensorForDtypeT( + currInUb, currInUb, vregIn, vregHalfIn, preg, preg, offset, halfOffset); + ops::LoadTwoTensorForDtypeT( + currSinUb, currSinUb, vregSin, vregHalfSin, preg, preg, offset, halfOffset); + ops::LoadTwoTensorForDtypeT( + currCosUb, currCosUb, vregCos, vregHalfCos, preg, preg, offset, halfOffset); + + Mul(vregSin, vregSin, vregHalfIn, preg); + Mul(vregHalfOut, vregHalfSin, vregIn, preg); + Mul(vregCos, vregCos, vregIn, preg); + Sub(vregOut, vregCos, vregSin, preg); + Mul(vregHalfCos, vregHalfCos, vregHalfIn, preg); + Add(vregHalfOut, vregHalfOut, vregHalfCos, preg); + + ops::StoreOneTensorForDtypeT(currOutUb, vregOut, preg, offset); + ops::StoreOneTensorForDtypeT(currOutUb, vregHalfOut, preg, halfOffset); + } + } + } + } +} + +/* + qOut[0] = q[0] * cos[0] - q[1] * sin[0] + qOut[1] = q[1] * cos[1] + q[0] * sin[1] + qOut[2] = q[2] * cos[2] - q[3] * sin[2] + qOut[3] = q[3] * cos[3] + q[2] * sin[3] +*/ +template +__aicore__ inline void QuarterAlignVF(const LocalTensor &sinTensor, const LocalTensor &cosTensor, + const LocalTensor &inTensor, const LocalTensor &outTensor, uint32_t dLen, uint32_t dAlign, uint16_t currSNum, + uint16_t currDNum) +{ + __local_mem__ T *sinUb = (__local_mem__ T *)sinTensor.GetPhyAddr(); + __local_mem__ T *cosUb = (__local_mem__ T *)cosTensor.GetPhyAddr(); + __local_mem__ T *inUb = (__local_mem__ T *)inTensor.GetPhyAddr(); + __local_mem__ T *outUb = (__local_mem__ T *)outTensor.GetPhyAddr(); + uint32_t quarterD = dLen / QUARTER_MODE_COEF; + uint32_t quarterDAlign = ops::CeilAlign(quarterD, static_cast(BLOCK_TYPE_SIZE / sizeof(T))); + uint16_t repeatTimes = ops::CeilDiv(quarterD, VL_FLOAT32_SIZE); + __local_mem__ T *currInUb; + __local_mem__ T *currOutUb; + __local_mem__ T *currSinUb; + __local_mem__ T *currCosUb; + + __VEC_SCOPE__ + { + MicroAPI::RegTensor vregIn; + MicroAPI::RegTensor vregQ1In; + MicroAPI::RegTensor vregQ2In; + MicroAPI::RegTensor vregQ3In; + MicroAPI::RegTensor vregSin; + MicroAPI::RegTensor vregQ1Sin; + MicroAPI::RegTensor vregQ2Sin; + MicroAPI::RegTensor vregQ3Sin; + MicroAPI::RegTensor vregCos; + MicroAPI::RegTensor vregQ1Cos; + MicroAPI::RegTensor vregQ2Cos; + MicroAPI::RegTensor vregQ3Cos; + MicroAPI::RegTensor vregOut; + MicroAPI::RegTensor vregQ1Out; + MicroAPI::RegTensor vregQ2Out; + MicroAPI::RegTensor vregQ3Out; + MicroAPI::MaskReg preg; + for (uint16_t sIdx = 0; sIdx < currSNum; sIdx++) { + currSinUb = sinUb + sIdx * dAlign; + currCosUb = cosUb + sIdx * dAlign; + for (uint16_t row = 0; row < currDNum; row++) { + currInUb = inUb + (sIdx * currDNum + row) * dAlign; + currOutUb = outUb + (sIdx * currDNum + row) * dAlign; + uint32_t updateCnt = quarterD; + for (uint16_t i = 0; i < repeatTimes; i++) { + preg = MicroAPI::UpdateMask(updateCnt); + uint32_t offset = i * VL_FLOAT32_SIZE; + uint32_t q1Offset = offset + quarterDAlign; + uint32_t q2Offset = q1Offset + quarterDAlign; + uint32_t q3Offset = q2Offset + quarterDAlign; + ops::LoadTwoTensorForDtypeT(currInUb, currInUb, vregIn, vregQ1In, preg, preg, offset, q1Offset); + ops::LoadTwoTensorForDtypeT( + currInUb, currInUb, vregQ2In, vregQ3In, preg, preg, q2Offset, q3Offset); + ops::LoadTwoTensorForDtypeT( + currSinUb, currSinUb, vregSin, vregQ1Sin, preg, preg, offset, q1Offset); + ops::LoadTwoTensorForDtypeT( + currSinUb, currSinUb, vregQ2Sin, vregQ3Sin, preg, preg, q2Offset, q3Offset); + ops::LoadTwoTensorForDtypeT( + currCosUb, currCosUb, vregCos, vregQ1Cos, preg, preg, offset, q1Offset); + ops::LoadTwoTensorForDtypeT( + currCosUb, currCosUb, vregQ2Cos, vregQ3Cos, preg, preg, q2Offset, q3Offset); + + Mul(vregSin, vregSin, vregQ1In, preg); + Mul(vregQ1Out, vregQ1Sin, vregIn, preg); + Mul(vregQ2Sin, vregQ2Sin, vregQ3In, preg); + Mul(vregQ3Out, vregQ3Sin, vregQ2In, preg); + Mul(vregCos, vregCos, vregIn, preg); + Sub(vregOut, vregCos, vregSin, preg); + Mul(vregQ1Cos, vregQ1Cos, vregQ1In, preg); + Add(vregQ1Out, vregQ1Out, vregQ1Cos, preg); + Mul(vregQ2Cos, vregQ2Cos, vregQ2In, preg); + Sub(vregQ2Out, vregQ2Cos, vregQ2Sin, preg); + Mul(vregQ3Cos, vregQ3Cos, vregQ3In, preg); + Add(vregQ3Out, vregQ3Out, vregQ3Cos, preg); + + ops::StoreOneTensorForDtypeT(currOutUb, vregOut, preg, offset); + ops::StoreOneTensorForDtypeT(currOutUb, vregQ1Out, preg, q1Offset); + ops::StoreOneTensorForDtypeT(currOutUb, vregQ2Out, preg, q2Offset); + ops::StoreOneTensorForDtypeT(currOutUb, vregQ3Out, preg, q3Offset); + } + } + } + } +} + +template +__aicore__ inline void InterleaveModeVF(const LocalTensor &sinTensor, const LocalTensor &cosTensor, + const LocalTensor &inTensor, const LocalTensor &outTensor, uint32_t dLen, uint16_t currSNum, + uint16_t currDNum) +{ + __local_mem__ T *sinUb = (__local_mem__ T *)sinTensor.GetPhyAddr(); + __local_mem__ T *cosUb = (__local_mem__ T *)cosTensor.GetPhyAddr(); + __local_mem__ T *inUb = (__local_mem__ T *)inTensor.GetPhyAddr(); + __local_mem__ T *outUb = (__local_mem__ T *)outTensor.GetPhyAddr(); + uint16_t repeatTimes = dLen / VL_FLOAT32_SIZE; + uint32_t dAlignLen = ops::CeilAlign(dLen, static_cast(BLOCK_TYPE_SIZE / sizeof(T))); + uint16_t loopNum = repeatTimes / 2; + uint32_t tailNum = dLen - loopNum * 2 * VL_FLOAT32_SIZE; + uint16_t tailTwoVL = tailNum / VL_FLOAT32_SIZE; + uint16_t tailOneVL = (tailTwoVL == 1) ? 0 : 1; + uint32_t tailLen = tailNum % VL_FLOAT32_SIZE; + __local_mem__ T *currInUb; + __local_mem__ T *currOutUb; + __local_mem__ T *currSinUb; + __local_mem__ T *currCosUb; + __local_mem__ T *tailSinUb; + __local_mem__ T *tailCosUb; + + __VEC_SCOPE__ + { + MicroAPI::RegTensor vregFormerCos; + MicroAPI::RegTensor vregLatterCos; + MicroAPI::RegTensor vregFormerSin; + MicroAPI::RegTensor vregLatterSin; + MicroAPI::RegTensor vregFormerIn; + MicroAPI::RegTensor vregLatterIn; + MicroAPI::RegTensor vregOdd; + MicroAPI::RegTensor vregEven; + MicroAPI::RegTensor vregFormerOut; + MicroAPI::RegTensor vregLatterOut; + MicroAPI::MaskReg pregLoop; + MicroAPI::MaskReg pregTail; + for (uint16_t sIdx = 0; sIdx < currSNum; sIdx++) { + currSinUb = sinUb + sIdx * dAlignLen; + currCosUb = cosUb + sIdx * dAlignLen; + for (uint16_t idxD = 0; idxD < currDNum; idxD++) { + uint32_t updateCnt = dLen; + currInUb = inUb + (sIdx * currDNum + idxD) * dAlignLen; + currOutUb = outUb + (sIdx * currDNum + idxD) * dAlignLen; + pregLoop = MicroAPI::CreateMask(); + for (uint16_t i = 0; i < loopNum; i++) { + uint32_t evenOffSet = (i * 2) * VL_FLOAT32_SIZE; + uint32_t oddOffset = evenOffSet + VL_FLOAT32_SIZE; + ops::LoadOneTensorForDtypeT(currInUb, vregFormerIn, pregLoop, evenOffSet); + ops::LoadOneTensorForDtypeT(currInUb, vregLatterIn, pregLoop, oddOffset); + ops::LoadOneTensorForDtypeT(currCosUb, vregFormerCos, pregLoop, evenOffSet); + ops::LoadOneTensorForDtypeT(currCosUb, vregLatterCos, pregLoop, oddOffset); + ops::LoadOneTensorForDtypeT(currSinUb, vregFormerSin, pregLoop, evenOffSet); + ops::LoadOneTensorForDtypeT(currSinUb, vregLatterSin, pregLoop, oddOffset); + Mul(vregFormerCos, vregFormerCos, vregFormerIn, pregLoop); + Mul(vregLatterCos, vregLatterCos, vregLatterIn, pregLoop); + MicroAPI::DeInterleave(vregEven, vregOdd, vregFormerIn, vregLatterIn); + Muls(vregOdd, vregOdd, float(-1.0), pregLoop); + MicroAPI::Interleave(vregFormerIn, vregLatterIn, vregOdd, vregEven); + Mul(vregFormerSin, vregFormerSin, vregFormerIn, pregLoop); + Add(vregFormerCos, vregFormerCos, vregFormerSin, pregLoop); + Mul(vregLatterSin, vregLatterSin, vregLatterIn, pregLoop); + Add(vregLatterCos, vregLatterCos, vregLatterSin, pregLoop); + ops::StoreOneTensorForDtypeT(currOutUb, vregFormerCos, pregLoop, evenOffSet); + ops::StoreOneTensorForDtypeT(currOutUb, vregLatterCos, pregLoop, oddOffset); + } + + currInUb = inUb + (sIdx * currDNum + idxD) * dAlignLen + (loopNum * 2 * VL_FLOAT32_SIZE); + currOutUb = outUb + (sIdx * currDNum + idxD) * dAlignLen + (loopNum * 2 * VL_FLOAT32_SIZE); + tailSinUb = currSinUb + loopNum * 2 * VL_FLOAT32_SIZE; + tailCosUb = currCosUb + loopNum * 2 * VL_FLOAT32_SIZE; + // 尾块大于VL时,读取一个VL，读取尾块 + for (uint16_t i = 0; i < tailTwoVL; i++) { + uint32_t updateCnt = tailLen; + pregTail = MicroAPI::UpdateMask(updateCnt); + ops::LoadOneTensorForDtypeT(currInUb, vregFormerIn, pregLoop, 0); + ops::LoadOneTensorForDtypeT(currInUb, vregLatterIn, pregTail, VL_FLOAT32_SIZE); + ops::LoadOneTensorForDtypeT(tailCosUb, vregFormerCos, pregLoop, 0); + ops::LoadOneTensorForDtypeT(tailCosUb, vregLatterCos, pregTail, VL_FLOAT32_SIZE); + ops::LoadOneTensorForDtypeT(tailSinUb, vregFormerSin, pregLoop, 0); + ops::LoadOneTensorForDtypeT(tailSinUb, vregLatterSin, pregTail, VL_FLOAT32_SIZE); + Mul(vregFormerCos, vregFormerCos, vregFormerIn, pregLoop); + Mul(vregLatterCos, vregLatterCos, vregLatterIn, pregTail); + MicroAPI::DeInterleave(vregEven, vregOdd, vregFormerIn, vregLatterIn); + Muls(vregOdd, vregOdd, float(-1.0), pregLoop); + MicroAPI::Interleave(vregFormerIn, vregLatterIn, vregOdd, vregEven); + Mul(vregFormerSin, vregFormerSin, vregFormerIn, pregLoop); + Add(vregFormerCos, vregFormerCos, vregFormerSin, pregLoop); + Mul(vregLatterSin, vregLatterSin, vregLatterIn, pregTail); + Add(vregLatterCos, vregLatterCos, vregLatterSin, pregTail); + ops::StoreOneTensorForDtypeT(currOutUb, vregFormerCos, pregLoop, 0); + ops::StoreOneTensorForDtypeT(currOutUb, vregLatterCos, pregTail, VL_FLOAT32_SIZE); + } + + // 尾块小于VL时,只读取VL + for (uint16_t i = 0; i < tailOneVL; i++) { + uint32_t updateCnt = tailLen; + pregTail = MicroAPI::UpdateMask(updateCnt); + ops::LoadOneTensorForDtypeT(currInUb, vregFormerIn, pregTail, 0); + ops::LoadOneTensorForDtypeT(tailCosUb, vregFormerCos, pregTail, 0); + ops::LoadOneTensorForDtypeT(tailSinUb, vregFormerSin, pregTail, 0); + Mul(vregFormerCos, vregFormerCos, vregFormerIn, pregTail); + MicroAPI::DeInterleave(vregEven, vregOdd, vregFormerIn, vregLatterIn); + Muls(vregOdd, vregOdd, float(-1.0), pregTail); + MicroAPI::Interleave(vregFormerIn, vregLatterIn, vregOdd, vregEven); + Mul(vregFormerSin, vregFormerSin, vregFormerIn, pregTail); + Add(vregFormerCos, vregFormerCos, vregFormerSin, pregTail); + ops::StoreOneTensorForDtypeT(currOutUb, vregFormerCos, pregTail, 0); + } + } + } + } +} + +template +__aicore__ inline void DeepSeekInterleaveModeVF(const LocalTensor &sinTensor, const LocalTensor &cosTensor, + const LocalTensor &inTensor, const LocalTensor &outTensor, uint32_t dLen, uint16_t currSNum, + uint16_t currDNum) +{ + __local_mem__ T *sinUb = (__local_mem__ T *)sinTensor.GetPhyAddr(); + __local_mem__ T *cosUb = (__local_mem__ T *)cosTensor.GetPhyAddr(); + __local_mem__ T *inUb = (__local_mem__ T *)inTensor.GetPhyAddr(); + __local_mem__ T *outUb = (__local_mem__ T *)outTensor.GetPhyAddr(); + uint32_t dAlign = ops::CeilAlign(dLen, static_cast(BLOCK_TYPE_SIZE / sizeof(T))); + uint32_t halfD = dLen / HALF_INTERLEAVE_COEF; + uint32_t halfDAlign = ops::CeilAlign(halfD, static_cast(BLOCK_TYPE_SIZE / sizeof(T))); + uint16_t repeatTimes = halfD / VL_FLOAT32_SIZE; + uint32_t tailTwoNum = dLen - repeatTimes * VL_FLOAT32_SIZE * HALF_INTERLEAVE_COEF; + uint16_t tailTwoVL = tailTwoNum > VL_FLOAT32_SIZE ? 1 : 0; + uint16_t tailOneVL = tailTwoNum > 0 ? (1 - tailTwoVL) : 0; + uint32_t halfTailNum = tailTwoNum / HALF_INTERLEAVE_COEF; + uint32_t tailNum = tailTwoNum - tailTwoVL * VL_FLOAT32_SIZE; + __local_mem__ T *currInUb; + __local_mem__ T *currOutUb; + __local_mem__ T *currSinUb; + __local_mem__ T *currCosUb; + + __VEC_SCOPE__ + { + MicroAPI::RegTensor vregIn; + MicroAPI::RegTensor vregHalfIn; + MicroAPI::RegTensor vregSin; + MicroAPI::RegTensor vregHalfSin; + MicroAPI::RegTensor vregCos; + MicroAPI::RegTensor vregHalfCos; + MicroAPI::RegTensor vregOut; + MicroAPI::RegTensor vregHalfOut; + MicroAPI::MaskReg pregTail; + MicroAPI::MaskReg pregHalfTail; + MicroAPI::MaskReg pregFull = MicroAPI::CreateMask(); + for (uint16_t sIdx = 0; sIdx < currSNum; sIdx++) { + currSinUb = sinUb + sIdx * halfDAlign * HALF_INTERLEAVE_COEF; + currCosUb = cosUb + sIdx * halfDAlign * HALF_INTERLEAVE_COEF; + uint32_t updateTailNum = tailNum; + uint32_t updateHalfTailNum = halfTailNum; + pregTail = MicroAPI::UpdateMask(updateTailNum); + pregHalfTail = MicroAPI::UpdateMask(updateHalfTailNum); + for (uint16_t row = 0; row < currDNum; row++) { + currInUb = inUb + (sIdx * currDNum + row) * dAlign; + currOutUb = outUb + (sIdx * currDNum + row) * halfDAlign * HALF_INTERLEAVE_COEF; + for (uint16_t i = 0; i < repeatTimes; i++) { + uint32_t offset = i * VL_FLOAT32_SIZE; + uint32_t halfOffset = offset + halfDAlign; + uint32_t inOffset = offset * HALF_INTERLEAVE_COEF; + ops::LoadTwoTensorForDtypeT(currInUb, + currInUb, + vregIn, + vregHalfIn, + pregFull, + pregFull, + inOffset, + inOffset + VL_FLOAT32_SIZE); + ops::LoadTwoTensorForDtypeT( + currSinUb, currSinUb, vregSin, vregHalfSin, pregFull, pregFull, offset, halfOffset); + ops::LoadTwoTensorForDtypeT( + currCosUb, currCosUb, vregCos, vregHalfCos, pregFull, pregFull, offset, halfOffset); + + DeInterleave(vregIn, vregHalfIn, vregIn, vregHalfIn); + + Mul(vregOut, vregCos, vregIn, pregFull); + Mul(vregHalfOut, vregHalfCos, vregHalfIn, pregFull); + Muls(vregHalfIn, vregHalfIn, float(-1.0), pregFull); + Mul(vregSin, vregSin, vregHalfIn, pregFull); + Add(vregOut, vregOut, vregSin, pregFull); + Mul(vregHalfSin, vregHalfSin, vregIn, pregFull); + Add(vregHalfOut, vregHalfOut, vregHalfSin, pregFull); + + ops::StoreOneTensorForDtypeT(currOutUb, vregOut, pregFull, offset); + ops::StoreOneTensorForDtypeT(currOutUb, vregHalfOut, pregFull, halfOffset); + } + + for (uint16_t i = 0; i < tailTwoVL; i++) { + uint32_t offset = repeatTimes * VL_FLOAT32_SIZE; + uint32_t halfOffset = offset + halfDAlign; + uint32_t inOffset = offset * HALF_INTERLEAVE_COEF; + ops::LoadTwoTensorForDtypeT(currInUb, + currInUb, + vregIn, + vregHalfIn, + pregFull, + pregTail, + inOffset, + inOffset + VL_FLOAT32_SIZE); + ops::LoadTwoTensorForDtypeT( + currSinUb, currSinUb, vregSin, vregHalfSin, pregHalfTail, pregHalfTail, offset, halfOffset); + ops::LoadTwoTensorForDtypeT( + currCosUb, currCosUb, vregCos, vregHalfCos, pregHalfTail, pregHalfTail, offset, halfOffset); + + DeInterleave(vregIn, vregHalfIn, vregIn, vregHalfIn); + + Mul(vregOut, vregCos, vregIn, pregHalfTail); + Mul(vregHalfOut, vregHalfCos, vregHalfIn, pregHalfTail); + Muls(vregHalfIn, vregHalfIn, float(-1.0), pregHalfTail); + Mul(vregSin, vregSin, vregHalfIn, pregHalfTail); + Add(vregOut, vregOut, vregSin, pregHalfTail); + Mul(vregHalfSin, vregHalfSin, vregIn, pregHalfTail); + Add(vregHalfOut, vregHalfOut, vregHalfSin, pregHalfTail); + + ops::StoreOneTensorForDtypeT(currOutUb, vregOut, pregHalfTail, offset); + ops::StoreOneTensorForDtypeT(currOutUb, vregHalfOut, pregHalfTail, halfOffset); + } + + for (uint16_t i = 0; i < tailOneVL; i++) { + uint32_t offset = repeatTimes * VL_FLOAT32_SIZE; + uint32_t halfOffset = offset + halfDAlign; + uint32_t inOffset = offset * HALF_INTERLEAVE_COEF; + ops::LoadOneTensorForDtypeT(currInUb, vregIn, pregTail, inOffset); + ops::LoadTwoTensorForDtypeT( + currSinUb, currSinUb, vregSin, vregHalfSin, pregHalfTail, pregHalfTail, offset, halfOffset); + ops::LoadTwoTensorForDtypeT( + currCosUb, currCosUb, vregCos, vregHalfCos, pregHalfTail, pregHalfTail, offset, halfOffset); + + DeInterleave(vregIn, vregHalfIn, vregIn, vregHalfIn); + Mul(vregOut, vregCos, vregIn, pregHalfTail); + Mul(vregHalfOut, vregHalfCos, vregHalfIn, pregHalfTail); + Muls(vregHalfIn, vregHalfIn, float(-1.0), pregHalfTail); + Mul(vregSin, vregSin, vregHalfIn, pregHalfTail); + Add(vregOut, vregOut, vregSin, pregHalfTail); + Mul(vregHalfSin, vregHalfSin, vregIn, pregHalfTail); + Add(vregHalfOut, vregHalfOut, vregHalfSin, pregHalfTail); + ops::StoreOneTensorForDtypeT(currOutUb, vregOut, pregHalfTail, offset); + ops::StoreOneTensorForDtypeT(currOutUb, vregHalfOut, pregHalfTail, halfOffset); + } + } + } + } +} + +template +__aicore__ inline void BatchHalfAlignVF(__local_mem__ T *in, __local_mem__ T *cos, __local_mem__ T *sin, + __local_mem__ T *out, uint16_t sLength, uint16_t bLength, uint16_t nLength, int64_t d, int64_t dAlign, + int64_t ubFactorS, int64_t ubFactorN) +{ + uint32_t dHalfSize = d / HALF_INTERLEAVE_COEF; + uint16_t dLoopCount = (dHalfSize + VL_FLOAT32_SIZE - 1) / VL_FLOAT32_SIZE; + uint32_t dHalfOffset = dAlign / HALF_INTERLEAVE_COEF; + + // 计算循环参数 + int32_t bStepUb = ubFactorN * ubFactorS * dAlign; + int32_t nStepUb = ubFactorS * dAlign; + + __VEC_SCOPE__ + { + // 定义相关寄存器 + MicroAPI::RegTensor inPart1Reg; + MicroAPI::RegTensor inPart2Reg; + MicroAPI::RegTensor cosPart1Reg; + MicroAPI::RegTensor cosPart2Reg; + MicroAPI::RegTensor sinPart1Reg; + MicroAPI::RegTensor sinPart2Reg; + MicroAPI::MaskReg pregLoop; + __local_mem__ T *currInUb, *currOutUb, *currSinUb, *currCosUb; + for (uint16_t bIdx = 0; bIdx < bLength; bIdx++) { + for (uint16_t nIdx = 0; nIdx < nLength; nIdx++) { + for (uint16_t sIdx = 0; sIdx < sLength; sIdx++) { + uint32_t count = dHalfSize; + currInUb = in + bIdx * bStepUb + nIdx * nStepUb + sIdx * dAlign; + currOutUb = out + bIdx * bStepUb + nIdx * nStepUb + sIdx * dAlign; + if constexpr (IsBBoardcast) { + currCosUb = cos + sIdx * dAlign; + currSinUb = sin + sIdx * dAlign; + } else { + currCosUb = cos + bIdx * nStepUb + sIdx * dAlign; + currSinUb = sin + bIdx * nStepUb + sIdx * dAlign; + } + for (uint16_t i = 0; i < dLoopCount; i++) { + pregLoop = MicroAPI::UpdateMask(count); + // 拷贝到RegBase内 + ops::LoadOneTensorForDtypeT(currInUb, inPart1Reg, pregLoop, i * VL_FLOAT32_SIZE); + ops::LoadOneTensorForDtypeT( + currInUb, inPart2Reg, pregLoop, i * VL_FLOAT32_SIZE + dHalfOffset); + ops::LoadOneTensorForDtypeT(currCosUb, cosPart1Reg, pregLoop, i * VL_FLOAT32_SIZE); + ops::LoadOneTensorForDtypeT( + currCosUb, cosPart2Reg, pregLoop, i * VL_FLOAT32_SIZE + dHalfOffset); + ops::LoadOneTensorForDtypeT(currSinUb, sinPart1Reg, pregLoop, i * VL_FLOAT32_SIZE); + ops::LoadOneTensorForDtypeT( + currSinUb, sinPart2Reg, pregLoop, i * VL_FLOAT32_SIZE + dHalfOffset); + // 计算 + Mul(cosPart1Reg, inPart1Reg, cosPart1Reg, pregLoop); + Mul(sinPart1Reg, inPart2Reg, sinPart1Reg, pregLoop); + Sub(cosPart1Reg, cosPart1Reg, sinPart1Reg, pregLoop); + Mul(cosPart2Reg, inPart2Reg, cosPart2Reg, pregLoop); + Mul(sinPart2Reg, sinPart2Reg, inPart1Reg, pregLoop); + Add(cosPart2Reg, cosPart2Reg, sinPart2Reg, pregLoop); + // 拷贝回UB + ops::StoreOneTensorForDtypeT(currOutUb, cosPart1Reg, pregLoop, i * VL_FLOAT32_SIZE); + ops::StoreOneTensorForDtypeT( + currOutUb, cosPart2Reg, pregLoop, i * VL_FLOAT32_SIZE + dHalfOffset); + } + } + } + } + } +} + +template +__aicore__ inline void BatchQuarterAlignVF(__local_mem__ T *in, __local_mem__ T *cos, __local_mem__ T *sin, + __local_mem__ T *out, uint16_t sLength, uint16_t bLength, uint16_t nLength, int64_t d, int64_t dAlign, + int64_t ubFactorS, int64_t ubFactorN) +{ + uint32_t dQuarterSize = d / QUARTER_MODE_COEF; + uint16_t dLoopCount = (dQuarterSize + VL_FLOAT32_SIZE - 1) / VL_FLOAT32_SIZE; + uint32_t dQuarterOffset = dAlign / QUARTER_MODE_COEF; + uint32_t dHalfOffset = dAlign / HALF_INTERLEAVE_COEF; + uint32_t dThreeQuarterOffset = dQuarterOffset + dHalfOffset; + + // 计算循环参数 + int32_t bStepUb = ubFactorN * ubFactorS * dAlign; + int32_t nStepUb = ubFactorS * dAlign; + + __VEC_SCOPE__ + { + // 定义相关寄存器 + MicroAPI::RegTensor inPart1Reg; + MicroAPI::RegTensor inPart2Reg; + MicroAPI::RegTensor inPart3Reg; + MicroAPI::RegTensor inPart4Reg; + MicroAPI::RegTensor cosPart1Reg; + MicroAPI::RegTensor cosPart2Reg; + MicroAPI::RegTensor cosPart3Reg; + MicroAPI::RegTensor cosPart4Reg; + MicroAPI::RegTensor sinPart1Reg; + MicroAPI::RegTensor sinPart2Reg; + MicroAPI::RegTensor sinPart3Reg; + MicroAPI::RegTensor sinPart4Reg; + MicroAPI::MaskReg pregLoop; + __local_mem__ T *currInUb, *currOutUb, *currSinUb, *currCosUb; + for (uint16_t bIdx = 0; bIdx < bLength; bIdx++) { + for (uint16_t nIdx = 0; nIdx < nLength; nIdx++) { + for (uint16_t sIdx = 0; sIdx < sLength; sIdx++) { + uint32_t count = dQuarterSize; + currInUb = in + bIdx * bStepUb + nIdx * nStepUb + sIdx * dAlign; + currOutUb = out + bIdx * bStepUb + nIdx * nStepUb + sIdx * dAlign; + if constexpr (IsBBoardcast) { + currCosUb = cos + sIdx * dAlign; + currSinUb = sin + sIdx * dAlign; + } else { + currCosUb = cos + bIdx * nStepUb + sIdx * dAlign; + currSinUb = sin + bIdx * nStepUb + sIdx * dAlign; + } + for (uint16_t i = 0; i < dLoopCount; i++) { + pregLoop = MicroAPI::UpdateMask(count); + // 拷贝到RegBase内 + ops::LoadTwoTensorForDtypeT(currInUb, + currInUb, + inPart1Reg, + inPart2Reg, + pregLoop, + pregLoop, + i * VL_FLOAT32_SIZE, + i * VL_FLOAT32_SIZE + dQuarterOffset); + ops::LoadTwoTensorForDtypeT(currInUb, + currInUb, + inPart3Reg, + inPart4Reg, + pregLoop, + pregLoop, + i * VL_FLOAT32_SIZE + dHalfOffset, + i * VL_FLOAT32_SIZE + dThreeQuarterOffset); + ops::LoadTwoTensorForDtypeT(currCosUb, + currCosUb, + cosPart1Reg, + cosPart2Reg, + pregLoop, + pregLoop, + i * VL_FLOAT32_SIZE, + i * VL_FLOAT32_SIZE + dQuarterOffset); + ops::LoadTwoTensorForDtypeT(currCosUb, + currCosUb, + cosPart3Reg, + cosPart4Reg, + pregLoop, + pregLoop, + i * VL_FLOAT32_SIZE + dHalfOffset, + i * VL_FLOAT32_SIZE + dThreeQuarterOffset); + ops::LoadTwoTensorForDtypeT(currSinUb, + currSinUb, + sinPart1Reg, + sinPart2Reg, + pregLoop, + pregLoop, + i * VL_FLOAT32_SIZE, + i * VL_FLOAT32_SIZE + dQuarterOffset); + ops::LoadTwoTensorForDtypeT(currSinUb, + currSinUb, + sinPart3Reg, + sinPart4Reg, + pregLoop, + pregLoop, + i * VL_FLOAT32_SIZE + dHalfOffset, + i * VL_FLOAT32_SIZE + dThreeQuarterOffset); + // 计算 + Mul(cosPart1Reg, inPart1Reg, cosPart1Reg, pregLoop); + Mul(sinPart1Reg, inPart2Reg, sinPart1Reg, pregLoop); + Sub(cosPart1Reg, cosPart1Reg, sinPart1Reg, pregLoop); + Mul(cosPart2Reg, inPart2Reg, cosPart2Reg, pregLoop); + Mul(sinPart2Reg, sinPart2Reg, inPart1Reg, pregLoop); + Add(cosPart2Reg, cosPart2Reg, sinPart2Reg, pregLoop); + Mul(cosPart3Reg, inPart3Reg, cosPart3Reg, pregLoop); + Mul(sinPart3Reg, inPart4Reg, sinPart3Reg, pregLoop); + Sub(cosPart3Reg, cosPart3Reg, sinPart3Reg, pregLoop); + Mul(cosPart4Reg, inPart4Reg, cosPart4Reg, pregLoop); + Mul(sinPart4Reg, sinPart4Reg, inPart3Reg, pregLoop); + Add(cosPart4Reg, cosPart4Reg, sinPart4Reg, pregLoop); + // 拷贝回UB + ops::StoreOneTensorForDtypeT(currOutUb, cosPart1Reg, pregLoop, i * VL_FLOAT32_SIZE); + ops::StoreOneTensorForDtypeT( + currOutUb, cosPart2Reg, pregLoop, i * VL_FLOAT32_SIZE + dQuarterOffset); + ops::StoreOneTensorForDtypeT( + currOutUb, cosPart3Reg, pregLoop, i * VL_FLOAT32_SIZE + dHalfOffset); + ops::StoreOneTensorForDtypeT( + currOutUb, cosPart4Reg, pregLoop, i * VL_FLOAT32_SIZE + dThreeQuarterOffset); + } + } + } + } + } +} + +template +__aicore__ inline void BatchInterleaveModeVF(__local_mem__ T *in, __local_mem__ T *cos, __local_mem__ T *sin, + __local_mem__ T *out, uint16_t sLength, uint16_t bLength, uint16_t nLength, int64_t d, int64_t dAlign, + int64_t ubFactorS, int64_t ubFactorN) +{ + uint32_t loopSize = 2 * VL_FLOAT32_SIZE; + uint16_t dLoopCount = (d + loopSize - 1) / loopSize; + + // 计算Mask参数 + uint32_t halfNum = d / 2; + uint32_t part1Num = (dLoopCount - 1) * VL_FLOAT32_SIZE; + uint32_t part2Num = part1Num; + uint32_t tailNum = d - part1Num - part2Num; + if (tailNum > VL_FLOAT32_SIZE) { + part1Num += VL_FLOAT32_SIZE; + part2Num += (tailNum - VL_FLOAT32_SIZE); + } else { + part1Num += tailNum; + } + + // 计算循环参数 + int32_t bStepUb = ubFactorN * ubFactorS * dAlign; + int32_t nStepUb = ubFactorS * dAlign; + + __VEC_SCOPE__ + { + // 定义相关寄存器 + MicroAPI::RegTensor inPart1Reg; + MicroAPI::RegTensor inPart2Reg; + MicroAPI::RegTensor cosPart1Reg; + MicroAPI::RegTensor cosPart2Reg; + MicroAPI::RegTensor sinPart1Reg; + MicroAPI::RegTensor sinPart2Reg; + MicroAPI::MaskReg pregLoop; + MicroAPI::MaskReg pregPart1; + MicroAPI::MaskReg pregPart2; + __local_mem__ T *currInUb, *currOutUb, *currSinUb, *currCosUb; + for (uint16_t bIdx = 0; bIdx < bLength; bIdx++) { + for (uint16_t nIdx = 0; nIdx < nLength; nIdx++) { + for (uint16_t sIdx = 0; sIdx < sLength; sIdx++) { + uint32_t halfCnt = halfNum; + uint32_t part1Cnt = part1Num; + uint32_t part2Cnt = part2Num; + currInUb = in + bIdx * bStepUb + nIdx * nStepUb + sIdx * dAlign; + currOutUb = out + bIdx * bStepUb + nIdx * nStepUb + sIdx * dAlign; + if constexpr (IsBBoardcast) { + currCosUb = cos + sIdx * dAlign; + currSinUb = sin + sIdx * dAlign; + } else { + currCosUb = cos + bIdx * nStepUb + sIdx * dAlign; + currSinUb = sin + bIdx * nStepUb + sIdx * dAlign; + } + for (uint16_t i = 0; i < dLoopCount; i++) { + pregLoop = MicroAPI::UpdateMask(halfCnt); + pregPart1 = MicroAPI::UpdateMask(part1Cnt); + pregPart2 = MicroAPI::UpdateMask(part2Cnt); + ops::LoadOneTensorForDtypeT(currInUb, inPart1Reg, pregPart1, i * loopSize); + ops::LoadOneTensorForDtypeT(currInUb, inPart2Reg, pregPart2, i * loopSize + VL_FLOAT32_SIZE); + ops::LoadOneTensorForDtypeT(currCosUb, cosPart1Reg, pregPart1, i * loopSize); + ops::LoadOneTensorForDtypeT( + currCosUb, cosPart2Reg, pregPart2, i * loopSize + VL_FLOAT32_SIZE); + ops::LoadOneTensorForDtypeT(currSinUb, sinPart1Reg, pregPart1, i * loopSize); + ops::LoadOneTensorForDtypeT( + currSinUb, sinPart2Reg, pregPart2, i * loopSize + VL_FLOAT32_SIZE); + Mul(cosPart1Reg, cosPart1Reg, inPart1Reg, pregPart1); + Mul(cosPart2Reg, cosPart2Reg, inPart2Reg, pregPart2); + MicroAPI::DeInterleave(inPart1Reg, inPart2Reg, inPart1Reg, inPart2Reg); + Muls(inPart2Reg, inPart2Reg, float(-1.0), pregLoop); + MicroAPI::Interleave(inPart1Reg, inPart2Reg, inPart2Reg, inPart1Reg); + Mul(sinPart1Reg, sinPart1Reg, inPart1Reg, pregPart1); + Add(cosPart1Reg, cosPart1Reg, sinPart1Reg, pregPart1); + Mul(sinPart2Reg, sinPart2Reg, inPart2Reg, pregPart2); + Add(cosPart2Reg, cosPart2Reg, sinPart2Reg, pregPart2); + ops::StoreOneTensorForDtypeT(currOutUb, cosPart1Reg, pregPart1, i * loopSize); + ops::StoreOneTensorForDtypeT( + currOutUb, cosPart2Reg, pregPart2, i * loopSize + VL_FLOAT32_SIZE); + } + } + } + } + } +} + +template +__aicore__ inline void BatchDeepSeekInterleaveModeVF(__local_mem__ T *in, __local_mem__ T *cos, __local_mem__ T *sin, + __local_mem__ T *out, uint16_t sLength, uint16_t bLength, uint16_t nLength, int64_t d, int64_t dAlign, + int64_t ubFactorS, int64_t ubFactorN) +{ + uint32_t loopSize = 2 * VL_FLOAT32_SIZE; + uint16_t dLoopCount = (d + loopSize - 1) / loopSize; + uint32_t dHalfOffset = dAlign / HALF_INTERLEAVE_COEF; + + // 计算Mask参数 + uint32_t halfNum = d / 2; + uint32_t part1Num = (dLoopCount - 1) * VL_FLOAT32_SIZE; + uint32_t part2Num = part1Num; + uint32_t tailNum = d - part1Num - part2Num; + if (tailNum > VL_FLOAT32_SIZE) { + part1Num += VL_FLOAT32_SIZE; + part2Num += (tailNum - VL_FLOAT32_SIZE); + } else { + part1Num += tailNum; + } + + // 计算循环参数 + int32_t bStepUb = ubFactorN * ubFactorS * dAlign; + int32_t nStepUb = ubFactorS * dAlign; + + __VEC_SCOPE__ + { + // 定义相关寄存器 + MicroAPI::RegTensor inPart1Reg; + MicroAPI::RegTensor inPart2Reg; + MicroAPI::RegTensor cosPart1Reg; + MicroAPI::RegTensor cosPart2Reg; + MicroAPI::RegTensor sinPart1Reg; + MicroAPI::RegTensor sinPart2Reg; + MicroAPI::MaskReg pregLoop; + MicroAPI::MaskReg pregPart1; + MicroAPI::MaskReg pregPart2; + __local_mem__ T *currInUb, *currOutUb, *currSinUb, *currCosUb; + for (uint16_t bIdx = 0; bIdx < bLength; bIdx++) { + for (uint16_t nIdx = 0; nIdx < nLength; nIdx++) { + for (uint16_t sIdx = 0; sIdx < sLength; sIdx++) { + uint32_t halfCnt = halfNum; + uint32_t part1Cnt = part1Num; + uint32_t part2Cnt = part2Num; + currInUb = in + bIdx * bStepUb + nIdx * nStepUb + sIdx * dAlign; + currOutUb = out + bIdx * bStepUb + nIdx * nStepUb + sIdx * dAlign; + if constexpr (IsBBoardcast) { + currCosUb = cos + sIdx * dAlign; + currSinUb = sin + sIdx * dAlign; + } else { + currCosUb = cos + bIdx * nStepUb + sIdx * dAlign; + currSinUb = sin + bIdx * nStepUb + sIdx * dAlign; + } + for (uint16_t i = 0; i < dLoopCount; i++) { + pregLoop = MicroAPI::UpdateMask(halfCnt); + pregPart1 = MicroAPI::UpdateMask(part1Cnt); + pregPart2 = MicroAPI::UpdateMask(part2Cnt); + ops::LoadOneTensorForDtypeT(currInUb, inPart1Reg, pregPart1, i * loopSize); + ops::LoadOneTensorForDtypeT(currInUb, inPart2Reg, pregPart2, i * loopSize + VL_FLOAT32_SIZE); + ops::LoadOneTensorForDtypeT(currCosUb, cosPart1Reg, pregLoop, i * VL_FLOAT32_SIZE); + ops::LoadOneTensorForDtypeT( + currCosUb, cosPart2Reg, pregLoop, i * VL_FLOAT32_SIZE + dHalfOffset); + ops::LoadOneTensorForDtypeT(currSinUb, sinPart1Reg, pregLoop, i * VL_FLOAT32_SIZE); + ops::LoadOneTensorForDtypeT( + currSinUb, sinPart2Reg, pregLoop, i * VL_FLOAT32_SIZE + dHalfOffset); + MicroAPI::DeInterleave(inPart1Reg, inPart2Reg, inPart1Reg, inPart2Reg); + Mul(cosPart1Reg, cosPart1Reg, inPart1Reg, pregLoop); + Mul(cosPart2Reg, cosPart2Reg, inPart2Reg, pregLoop); + Muls(inPart2Reg, inPart2Reg, float(-1.0), pregLoop); + Mul(sinPart1Reg, sinPart1Reg, inPart2Reg, pregLoop); + Add(cosPart1Reg, cosPart1Reg, sinPart1Reg, pregLoop); + Mul(sinPart2Reg, sinPart2Reg, inPart1Reg, pregLoop); + Add(cosPart2Reg, cosPart2Reg, sinPart2Reg, pregLoop); + ops::StoreOneTensorForDtypeT(currOutUb, cosPart1Reg, pregLoop, i * VL_FLOAT32_SIZE); + ops::StoreOneTensorForDtypeT( + currOutUb, cosPart2Reg, pregLoop, i * VL_FLOAT32_SIZE + dHalfOffset); + } + } + } + } + } +} + +// Mixed precision: TX is half/bfloat16 for input, cos/sin are float +template +__aicore__ inline void InterleaveModeVFMixed(const LocalTensor &inTensor, const LocalTensor &cosTensor, + const LocalTensor &sinTensor, const LocalTensor &outTensor, uint32_t dLen, uint16_t currSNum, + uint16_t currDNum) +{ + __local_mem__ TX *inUb = (__local_mem__ TX *)inTensor.GetPhyAddr(); + __local_mem__ float *cosUb = (__local_mem__ float *)cosTensor.GetPhyAddr(); + __local_mem__ float *sinUb = (__local_mem__ float *)sinTensor.GetPhyAddr(); + __local_mem__ TX *outUb = (__local_mem__ TX *)outTensor.GetPhyAddr(); + uint16_t repeatTimes = dLen / VL_FLOAT32_SIZE; + uint32_t dAlignLen = ops::CeilAlign(dLen, static_cast(BLOCK_TYPE_SIZE / sizeof(TX))); + uint32_t dAlignLenFloat = ops::CeilAlign(dLen, static_cast(BLOCK_TYPE_SIZE / sizeof(float))); + uint16_t loopNum = repeatTimes / 2; + uint32_t tailNum = dLen - loopNum * 2 * VL_FLOAT32_SIZE; + uint16_t tailTwoVL = tailNum / VL_FLOAT32_SIZE; + uint16_t tailOneVL = (tailTwoVL == 1) ? 0 : 1; + uint32_t tailLen = tailNum % VL_FLOAT32_SIZE; + __local_mem__ TX *currInUb; + __local_mem__ TX *currOutUb; + __local_mem__ float *currSinUb; + __local_mem__ float *currCosUb; + + __VEC_SCOPE__ + { + MicroAPI::RegTensor vregFormerCos; + MicroAPI::RegTensor vregLatterCos; + MicroAPI::RegTensor vregFormerSin; + MicroAPI::RegTensor vregLatterSin; + MicroAPI::RegTensor vregFormerIn; + MicroAPI::RegTensor vregLatterIn; + MicroAPI::RegTensor vregOdd; + MicroAPI::RegTensor vregEven; + MicroAPI::MaskReg pregLoop; + MicroAPI::MaskReg pregTail; + for (uint16_t sIdx = 0; sIdx < currSNum; sIdx++) { + currSinUb = sinUb + sIdx * dAlignLenFloat; + currCosUb = cosUb + sIdx * dAlignLenFloat; + for (uint16_t idxD = 0; idxD < currDNum; idxD++) { + currInUb = inUb + (sIdx * currDNum + idxD) * dAlignLen; + currOutUb = outUb + (sIdx * currDNum + idxD) * dAlignLen; + pregLoop = MicroAPI::CreateMask(); + for (uint16_t i = 0; i < loopNum; i++) { + uint32_t evenOffSet = (i * 2) * VL_FLOAT32_SIZE; + uint32_t oddOffset = evenOffSet + VL_FLOAT32_SIZE; + ops::LoadOneTensorForDtypeT(currInUb, vregFormerIn, pregLoop, evenOffSet); + ops::LoadOneTensorForDtypeT(currInUb, vregLatterIn, pregLoop, oddOffset); + DataCopy(vregFormerCos, currCosUb + evenOffSet); + DataCopy(vregLatterCos, currCosUb + oddOffset); + DataCopy(vregFormerSin, currSinUb + evenOffSet); + DataCopy(vregLatterSin, currSinUb + oddOffset); + Mul(vregFormerCos, vregFormerCos, vregFormerIn, pregLoop); + Mul(vregLatterCos, vregLatterCos, vregLatterIn, pregLoop); + MicroAPI::DeInterleave(vregEven, vregOdd, vregFormerIn, vregLatterIn); + Muls(vregOdd, vregOdd, float(-1.0), pregLoop); + MicroAPI::Interleave(vregFormerIn, vregLatterIn, vregOdd, vregEven); + Mul(vregFormerSin, vregFormerSin, vregFormerIn, pregLoop); + Add(vregFormerCos, vregFormerCos, vregFormerSin, pregLoop); + Mul(vregLatterSin, vregLatterSin, vregLatterIn, pregLoop); + Add(vregLatterCos, vregLatterCos, vregLatterSin, pregLoop); + ops::StoreOneTensorForDtypeT(currOutUb, vregFormerCos, pregLoop, evenOffSet); + ops::StoreOneTensorForDtypeT(currOutUb, vregLatterCos, pregLoop, oddOffset); + } + + currInUb = inUb + (sIdx * currDNum + idxD) * dAlignLen + (loopNum * 2 * VL_FLOAT32_SIZE); + currOutUb = outUb + (sIdx * currDNum + idxD) * dAlignLen + (loopNum * 2 * VL_FLOAT32_SIZE); + __local_mem__ float *tailSinUb = currSinUb + loopNum * 2 * VL_FLOAT32_SIZE; + __local_mem__ float *tailCosUb = currCosUb + loopNum * 2 * VL_FLOAT32_SIZE; + for (uint16_t i = 0; i < tailTwoVL; i++) { + uint32_t updateCnt = tailLen; + pregTail = MicroAPI::UpdateMask(updateCnt); + ops::LoadOneTensorForDtypeT(currInUb, vregFormerIn, pregLoop, 0); + ops::LoadOneTensorForDtypeT(currInUb, vregLatterIn, pregTail, VL_FLOAT32_SIZE); + DataCopy(vregFormerCos, tailCosUb); + DataCopy(vregLatterCos, tailCosUb + VL_FLOAT32_SIZE); + DataCopy(vregFormerSin, tailSinUb); + DataCopy(vregLatterSin, tailSinUb + VL_FLOAT32_SIZE); + Mul(vregFormerCos, vregFormerCos, vregFormerIn, pregLoop); + Mul(vregLatterCos, vregLatterCos, vregLatterIn, pregTail); + MicroAPI::DeInterleave(vregEven, vregOdd, vregFormerIn, vregLatterIn); + Muls(vregOdd, vregOdd, float(-1.0), pregLoop); + MicroAPI::Interleave(vregFormerIn, vregLatterIn, vregOdd, vregEven); + Mul(vregFormerSin, vregFormerSin, vregFormerIn, pregLoop); + Add(vregFormerCos, vregFormerCos, vregFormerSin, pregLoop); + Mul(vregLatterSin, vregLatterSin, vregLatterIn, pregTail); + Add(vregLatterCos, vregLatterCos, vregLatterSin, pregTail); + ops::StoreOneTensorForDtypeT(currOutUb, vregFormerCos, pregLoop, 0); + ops::StoreOneTensorForDtypeT(currOutUb, vregLatterCos, pregTail, VL_FLOAT32_SIZE); + } + + for (uint16_t i = 0; i < tailOneVL; i++) { + uint32_t updateCnt = tailLen; + pregTail = MicroAPI::UpdateMask(updateCnt); + ops::LoadOneTensorForDtypeT(currInUb, vregFormerIn, pregTail, 0); + DataCopy(vregFormerCos, tailCosUb); + DataCopy(vregFormerSin, tailSinUb); + Mul(vregFormerCos, vregFormerCos, vregFormerIn, pregTail); + MicroAPI::DeInterleave(vregEven, vregOdd, vregFormerIn, vregLatterIn); + Muls(vregOdd, vregOdd, float(-1.0), pregTail); + MicroAPI::Interleave(vregFormerIn, vregLatterIn, vregOdd, vregEven); + Mul(vregFormerSin, vregFormerSin, vregFormerIn, pregTail); + Add(vregFormerCos, vregFormerCos, vregFormerSin, pregTail); + ops::StoreOneTensorForDtypeT(currOutUb, vregFormerCos, pregTail, 0); + } + } + } + } +} + +// Mixed precision BatchInterleaveModeVF for ABA layout +template +__aicore__ inline void BatchInterleaveModeVFMixed(__local_mem__ TX *in, __local_mem__ float *cos, + __local_mem__ float *sin, __local_mem__ TX *out, uint16_t sLength, uint16_t bLength, uint16_t nLength, int64_t d, + int64_t dAlign, int64_t dAlignFloat, int64_t ubFactorS, int64_t ubFactorN) +{ + uint32_t loopSize = 2 * VL_FLOAT32_SIZE; + uint32_t txLoopSize = loopSize * sizeof(float) / sizeof(TX); // Convert to TX element units + uint16_t dLoopCount = (d + loopSize - 1) / loopSize; + + uint32_t halfNum = d / 2; + uint32_t part1Num = (dLoopCount - 1) * VL_FLOAT32_SIZE; + uint32_t part2Num = part1Num; + uint32_t tailNum = d - part1Num - part2Num; + if (tailNum > VL_FLOAT32_SIZE) { + part1Num += VL_FLOAT32_SIZE; + part2Num += (tailNum - VL_FLOAT32_SIZE); + } else { + part1Num += tailNum; + } + + int32_t bStepUb = ubFactorN * ubFactorS * dAlign; + int32_t nStepUb = ubFactorS * dAlign; + int32_t cosSinBStepUb = ubFactorS * dAlignFloat; + int32_t cosSinSStepUb = dAlignFloat; + + __VEC_SCOPE__ + { + MicroAPI::RegTensor inPart1Reg; + MicroAPI::RegTensor inPart2Reg; + MicroAPI::RegTensor cosPart1Reg; + MicroAPI::RegTensor cosPart2Reg; + MicroAPI::RegTensor sinPart1Reg; + MicroAPI::RegTensor sinPart2Reg; + MicroAPI::MaskReg pregLoop; + MicroAPI::MaskReg pregPart1; + MicroAPI::MaskReg pregPart2; + __local_mem__ TX *currInUb, *currOutUb; + __local_mem__ float *currSinUb, *currCosUb; + for (uint16_t bIdx = 0; bIdx < bLength; bIdx++) { + for (uint16_t nIdx = 0; nIdx < nLength; nIdx++) { + for (uint16_t sIdx = 0; sIdx < sLength; sIdx++) { + uint32_t halfCnt = halfNum; + uint32_t part1Cnt = part1Num; + uint32_t part2Cnt = part2Num; + currInUb = in + bIdx * bStepUb + nIdx * nStepUb + sIdx * dAlign; + currOutUb = out + bIdx * bStepUb + nIdx * nStepUb + sIdx * dAlign; + if constexpr (IsBBoardcast) { + currCosUb = cos + sIdx * cosSinSStepUb; + currSinUb = sin + sIdx * cosSinSStepUb; + } else { + currCosUb = cos + bIdx * cosSinBStepUb + sIdx * cosSinSStepUb; + currSinUb = sin + bIdx * cosSinBStepUb + sIdx * cosSinSStepUb; + } + for (uint16_t i = 0; i < dLoopCount; i++) { + pregLoop = MicroAPI::UpdateMask(halfCnt); + pregPart1 = MicroAPI::UpdateMask(part1Cnt); + pregPart2 = MicroAPI::UpdateMask(part2Cnt); + ops::LoadOneTensorForDtypeT(currInUb, inPart1Reg, pregPart1, i * txLoopSize); + ops::LoadOneTensorForDtypeT(currInUb, + inPart2Reg, + pregPart2, + i * txLoopSize + VL_FLOAT32_SIZE * sizeof(float) / sizeof(TX)); + ops::LoadOneTensorForDtypeT(currCosUb, cosPart1Reg, pregPart1, i * loopSize); + ops::LoadOneTensorForDtypeT( + currCosUb, cosPart2Reg, pregPart2, i * loopSize + VL_FLOAT32_SIZE); + ops::LoadOneTensorForDtypeT(currSinUb, sinPart1Reg, pregPart1, i * loopSize); + ops::LoadOneTensorForDtypeT( + currSinUb, sinPart2Reg, pregPart2, i * loopSize + VL_FLOAT32_SIZE); + Mul(cosPart1Reg, cosPart1Reg, inPart1Reg, pregPart1); + Mul(cosPart2Reg, cosPart2Reg, inPart2Reg, pregPart2); + MicroAPI::DeInterleave(inPart1Reg, inPart2Reg, inPart1Reg, inPart2Reg); + Muls(inPart2Reg, inPart2Reg, float(-1.0), pregLoop); + MicroAPI::Interleave(inPart1Reg, inPart2Reg, inPart2Reg, inPart1Reg); + Mul(sinPart1Reg, sinPart1Reg, inPart1Reg, pregPart1); + Add(cosPart1Reg, cosPart1Reg, sinPart1Reg, pregPart1); + Mul(sinPart2Reg, sinPart2Reg, inPart2Reg, pregPart2); + Add(cosPart2Reg, cosPart2Reg, sinPart2Reg, pregPart2); + ops::StoreOneTensorForDtypeT(currOutUb, cosPart1Reg, pregPart1, i * txLoopSize); + ops::StoreOneTensorForDtypeT(currOutUb, + cosPart2Reg, + pregPart2, + i * txLoopSize + VL_FLOAT32_SIZE * sizeof(float) / sizeof(TX)); + } + } + } + } + } +} + +#endif // APPLY_ROTARY_POS_EMB_COMMON_H diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/inplace_partial_rotary_mul.cpp b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/inplace_partial_rotary_mul.cpp new file mode 100644 index 000000000..bb8a20d61 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/inplace_partial_rotary_mul.cpp @@ -0,0 +1,322 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotary_position_embedding.cpp + * \brief + */ +#if defined(__DAV_C310__) + #include "kernel_operator.h" + #include "rotary_position_embedding_reg_bab.h" + #include "rotary_position_embedding_reg_ab.h" + #include "rotary_position_embedding_reg_aba_and_ba.h" + #include "rotary_position_embedding_reg_a_and_b.h" + #include "rotary_position_embedding_reg_bab_mixed.h" + #include "rotary_position_embedding_reg_aba_and_ba_mixed.h" + #include "rotary_position_embedding_reg_a_and_b_mixed.h" + #include "rotary_position_embedding_reg_ab_mixed.h" +#else + #include "kernel_operator.h" + #include "kernel_tiling/kernel_tiling.h" + #include "inplace_partial_rotary_mul.h" + #include "rotate_interleaved_split_s.h" + #include "rotate_interleaved_split_bs.h" + #include "rotate_interleaved_split_bsn.h" + #include "rotate_interleaved_split_s_pad.h" + #include "rotate_interleaved_split_bs_pad.h" + #include "rotate_interleaved_split_bsn_pad.h" + using namespace AscendC; + using namespace RotateInterleavedN; +#endif + +#define TILING_KEY_ABA 20010 +#define TILING_KEY_BA 20011 +#define TILING_KEY_BAB 20020 +#define TILING_KEY_AB 20030 +#define TILING_KEY_A 20040 +#define TILING_KEY_B 20041 + +#define TILING_KEY_ABA_BF16_FP32_MIXED 20110 +#define TILING_KEY_ABA_FP16_FP32_MIXED 20210 +#define TILING_KEY_BA_BF16_FP32_MIXED 20111 +#define TILING_KEY_BA_FP16_FP32_MIXED 20211 +#define TILING_KEY_BAB_BF16_FP32_MIXED 20120 +#define TILING_KEY_BAB_FP16_FP32_MIXED 20220 +#define TILING_KEY_AB_BF16_FP32_MIXED 20130 +#define TILING_KEY_AB_FP16_FP32_MIXED 20230 +#define TILING_KEY_A_BF16_FP32_MIXED 20140 +#define TILING_KEY_A_FP16_FP32_MIXED 20240 +#define TILING_KEY_B_BF16_FP32_MIXED 20141 +#define TILING_KEY_B_FP16_FP32_MIXED 20241 + +#define TILING_KEY1 1 +#define TILING_KEY2 2 +#define TILING_KEY1_FP32_ROPE 11 +#define TILING_KEY2_FP32_ROPE 12 + +using namespace AscendC; +using namespace InplacePartialRotaryMul; + +extern "C" __global__ __aicore__ void inplace_partial_rotary_mul(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + GM_ADDR workspace, GM_ADDR tiling) +{ + AscendC::TPipe pipe; + #if defined(__DAV_C310__) + if (TILING_KEY_IS(TILING_KEY_ABA)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingABAAndBA op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + else if (TILING_KEY_IS(TILING_KEY_BA)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingABAAndBA op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + else if (TILING_KEY_IS(TILING_KEY_BAB)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingBAB op(&pipe, tilingData); + op.Init(x, cos, sin, y); + op.Process(); + } + else if (TILING_KEY_IS(TILING_KEY_AB)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingAB op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + else if (TILING_KEY_IS(TILING_KEY_A)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingAAndB op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + else if (TILING_KEY_IS(TILING_KEY_B)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingAAndB op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + // Mixed precision: x is half/bfloat16, cos/sin are float32 + else if (TILING_KEY_IS(TILING_KEY_BAB_FP16_FP32_MIXED)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingBABMixed op(&pipe, tilingData); + op.Init(x, cos, sin, y); + op.Process(); + } + else if (TILING_KEY_IS(TILING_KEY_BAB_BF16_FP32_MIXED)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingBABMixed op(&pipe, tilingData); + op.Init(x, cos, sin, y); + op.Process(); + } + // Mixed precision ABA/BA kernels + else if (TILING_KEY_IS(TILING_KEY_ABA_FP16_FP32_MIXED)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingABAAndBAMixed op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + else if (TILING_KEY_IS(TILING_KEY_ABA_BF16_FP32_MIXED)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingABAAndBAMixed op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + else if (TILING_KEY_IS(TILING_KEY_BA_FP16_FP32_MIXED)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingABAAndBAMixed op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + else if (TILING_KEY_IS(TILING_KEY_BA_BF16_FP32_MIXED)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingABAAndBAMixed op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + // Mixed precision AAndB kernels + else if (TILING_KEY_IS(TILING_KEY_A_FP16_FP32_MIXED)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingAAndBMixed op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + else if (TILING_KEY_IS(TILING_KEY_A_BF16_FP32_MIXED)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingAAndBMixed op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + else if (TILING_KEY_IS(TILING_KEY_B_FP16_FP32_MIXED)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingAAndBMixed op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + else if (TILING_KEY_IS(TILING_KEY_B_BF16_FP32_MIXED)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingAAndBMixed op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + // Mixed precision AB kernels + else if (TILING_KEY_IS(TILING_KEY_AB_FP16_FP32_MIXED)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingABMixed op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + else if (TILING_KEY_IS(TILING_KEY_AB_BF16_FP32_MIXED)) + { + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tiling_data_in, tiling); + const RopeRegbaseTilingData *__restrict tilingData = &tiling_data_in; + InplacePartialRotaryMul::RotaryPositionEmbeddingABMixed op; + op.Init(x, cos, sin, y, workspace, tilingData, &pipe); + op.Process(); + } + #else + GET_TILING_DATA_WITH_STRUCT(RopeRegbaseTilingData, tilingData, tiling); + const RopeRegbaseTilingData* __restrict__ tilingData1 = &tilingData; + if (TILING_KEY_IS(TILING_KEY1)) { + InplacePartialRotaryMul::InplacePartialRotaryMulABA op; + op.Init(x, cos, sin, y, workspace, tilingData1, &pipe); + op.Process(); + return; + } + if (TILING_KEY_IS(TILING_KEY2)) { + InplacePartialRotaryMul::InplacePartialRotaryMulABA op; + op.Init(x, cos, sin, y, workspace, tilingData1, &pipe); + op.Process(); + return; + } + if (TILING_KEY_IS(TILING_KEY1_FP32_ROPE)) { + InplacePartialRotaryMul::InplacePartialRotaryMulABA op; + op.Init(x, cos, sin, y, workspace, tilingData1, &pipe); + op.Process(); + return; + } + if (TILING_KEY_IS(TILING_KEY2_FP32_ROPE)) { + InplacePartialRotaryMul::InplacePartialRotaryMulABA op; + op.Init(x, cos, sin, y, workspace, tilingData1, &pipe); + op.Process(); + return; + } + // mode: rotate_interleaved + if (TILING_KEY_IS(2000)) { + InterleavedSplitS interleavedSplitS; + interleavedSplitS.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitS.Process(); + } else if (TILING_KEY_IS(2010)) { + InterleavedSplitS interleavedSplitS; + interleavedSplitS.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitS.Process(); + } else if (TILING_KEY_IS(2020)) { + InterleavedSplitS interleavedSplitS; + interleavedSplitS.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitS.Process(); + } else if (TILING_KEY_IS(2100)) { + InterleavedSplitBS interleavedSplitBS; + interleavedSplitBS.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitBS.Process(); + } else if (TILING_KEY_IS(2110)) { + InterleavedSplitBS interleavedSplitBS; + interleavedSplitBS.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitBS.Process(); + } else if (TILING_KEY_IS(2120)) { + InterleavedSplitBS interleavedSplitBS; + interleavedSplitBS.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitBS.Process(); + } else if (TILING_KEY_IS(2200)) { + InterleavedSplitBSN interleavedSplitBSN; + interleavedSplitBSN.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitBSN.Process(); + } else if (TILING_KEY_IS(2210)) { + InterleavedSplitBSN interleavedSplitBSN; + interleavedSplitBSN.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitBSN.Process(); + } else if (TILING_KEY_IS(2220)) { + InterleavedSplitBSN interleavedSplitBSN; + interleavedSplitBSN.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitBSN.Process(); + } else if (TILING_KEY_IS(2001)) { + InterleavedSplitSPad interleavedSplitSPad; + interleavedSplitSPad.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitSPad.Process(); + } else if (TILING_KEY_IS(2011)) { + InterleavedSplitSPad interleavedSplitSPad; + interleavedSplitSPad.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitSPad.Process(); + } else if (TILING_KEY_IS(2021)) { + InterleavedSplitSPad interleavedSplitSPad; + interleavedSplitSPad.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitSPad.Process(); + } else if (TILING_KEY_IS(2101)) { + InterleavedSplitBSPad interleavedSplitBSPad; + interleavedSplitBSPad.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitBSPad.Process(); + } else if (TILING_KEY_IS(2111)) { + InterleavedSplitBSPad interleavedSplitBSPad; + interleavedSplitBSPad.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitBSPad.Process(); + } else if (TILING_KEY_IS(2121)) { + InterleavedSplitBSPad interleavedSplitBSPad; + interleavedSplitBSPad.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitBSPad.Process(); + } else if (TILING_KEY_IS(2201)) { + InterleavedSplitBSNPad interleavedSplitBSNPad; + interleavedSplitBSNPad.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitBSNPad.Process(); + } else if (TILING_KEY_IS(2211)) { + InterleavedSplitBSNPad interleavedSplitBSNPad; + interleavedSplitBSNPad.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitBSNPad.Process(); + } else if (TILING_KEY_IS(2221)) { + InterleavedSplitBSNPad interleavedSplitBSNPad; + interleavedSplitBSNPad.Init(x, cos, sin, y, tilingData1, &pipe); + interleavedSplitBSNPad.Process(); + } + #endif +} diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/inplace_partial_rotary_mul.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/inplace_partial_rotary_mul.h new file mode 100644 index 000000000..b03616099 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/inplace_partial_rotary_mul.h @@ -0,0 +1,302 @@ +/** + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/*! + * \file inplace_partial_rotary_mul.h + * \brief + */ +#ifndef INPLACE_PARITAL_ROTARY_MUL_H +#define INPLACE_PARITAL_ROTARY_MUL_H + +#include "kernel_operator.h" + +namespace InplacePartialRotaryMul { +using namespace AscendC; + +template +class InplacePartialRotaryMulABA { +public: + __aicore__ inline InplacePartialRotaryMulABA() {}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR r1, GM_ADDR r2, GM_ADDR y, GM_ADDR workspace, + const RopeRegbaseTilingData *tilingData, TPipe *pipe); + __aicore__ inline void Process(); + __aicore__ inline void CopyInData(LocalTensor &xUb, GlobalTensor &xGm, int64_t blockCout, int64_t blockLen, int64_t gmOffset, int64_t ubOffset); + __aicore__ inline void CopyInDataR(LocalTensor &xUb, GlobalTensor &xGm, int64_t blockCout, int64_t blockLen, int64_t gmOffset, int64_t ubOffset); + __aicore__ inline void SetGatherSrcOffset(LocalTensor &idsUb, int64_t count); + __aicore__ inline void ComputeMul(LocalTensor &dtsUb, LocalTensor & src0Ub, LocalTensor &src1Ub, int64_t onceA, int64_t numHead,int64_t headDim); + __aicore__ inline void InterleavedInversion(int64_t count,LocalTensor &ub); + __aicore__ inline void DataCopyOut(LocalTensor &yUb, GlobalTensor &xGm, int64_t blockCout, int64_t blockLen, int64_t gmOffset, int64_t ubOffset); +private: + TPipe* pipe_; + const RopeRegbaseTilingData* tiling_; + int32_t blockIdx_ = 0; + + //需要的tilingdata + int64_t halfNumx_ = 0; + int64_t ropeUbOffset_ = 0; + int64_t xNum_= 0; + int64_t r1Num_= 0; + int64_t count_ = 0; + + static constexpr int32_t ONE_BLOCK_SIZE = 32; + int32_t perBlock32 = ONE_BLOCK_SIZE / sizeof(float); + + GlobalTensor xGm_; + GlobalTensor r1Gm_; + GlobalTensor r2Gm_; + GlobalTensor yGm_; + + TQue xQue_; + TQue r1Que_; + TQue r2Que_; + TQue yQue_; + TBuf idsBuf_; + BinaryRepeatParams repeatParams_{1, 1, 1, 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams_{false, 0, 0, 0}; + DataCopyPadExtParams dataCopyRPadParams_{false, 0, 0, 0}; + + int64_t usedCoreNum_=0; + int64_t numHead_ =0; + int64_t headDim_ =0; + int64_t allHeadDim_ = 0; + int64_t coreTUbLoopTime_ =0; + int64_t coreBUbLoopTime_ = 0; //b分核，每个核处理多少个b，b就是shape0 + int64_t coreTUbLoopTail_ =0; + int64_t coreBUbLoopTail_ = 0; + int64_t ubFactor_ = 0; + int64_t start_ = 0; + int64_t blockFactor_=0; + +}; + +template +__aicore__ inline void InplacePartialRotaryMulABA::Init(GM_ADDR x, GM_ADDR r1, GM_ADDR r2, GM_ADDR y, GM_ADDR workspace, + const RopeRegbaseTilingData *tilingData, TPipe *pipe) +{ + blockIdx_ = GetBlockIdx(); + tiling_ = tilingData; + pipe_ = pipe; + usedCoreNum_ =tiling_->usedCoreNum; + numHead_ = tiling_->numHead; + headDim_ = tiling_->headDim; + allHeadDim_ = tiling_->allHeadDim; + coreTUbLoopTime_ = tiling_->coreTUbLoopTime; + coreBUbLoopTime_ = tiling_->coreBUbLoopTime; //b分核，每个核处理多少个b，b就是shape0 + coreTUbLoopTail_ = tiling_->coreTUbLoopTail; + coreBUbLoopTail_ = tiling_->coreBUbLoopTail; + ubFactor_ = tiling_->ubFactor; + blockFactor_ = tiling_->blockFactor; + start_ = tiling_->start; + + xGm_.SetGlobalBuffer((__gm__ T *)x); + r1Gm_.SetGlobalBuffer((__gm__ R *)r1); + r2Gm_.SetGlobalBuffer((__gm__ R *)r2); + yGm_.SetGlobalBuffer((__gm__ T *)y); + count_ = numHead_ * headDim_; + xNum_ = ubFactor_ * count_; + r1Num_ = ubFactor_ * headDim_; + + pipe_->InitBuffer(xQue_, 2, xNum_ * sizeof(float)); + pipe_->InitBuffer(r1Que_, 2, r1Num_ * sizeof(float)); + pipe_->InitBuffer(r2Que_, 2, r1Num_ * sizeof(float)); + pipe_->InitBuffer(yQue_, 2, xNum_ * sizeof(float)); + pipe_->InitBuffer(idsBuf_, count_ *sizeof(uint32_t)); + if constexpr(sizeof(T) != sizeof(float)) { + halfNumx_ = xNum_; + } + if constexpr(sizeof(R) != sizeof(float)) { + ropeUbOffset_ = r1Num_; + } +} +template +__aicore__ inline void InplacePartialRotaryMulABA::CopyInData(LocalTensor &xUb, GlobalTensor &xGm, int64_t blockCout, int64_t blockLen, int64_t gmOffset, int64_t ubOffset) +{ + DataCopyExtParams copyParams; + copyParams.blockCount = blockCout; + copyParams.blockLen = blockLen * sizeof(T); + copyParams.srcStride = (allHeadDim_ - headDim_)*sizeof(T); //整个输入的大小 + copyParams.dstStride = 0; + DataCopyPad(xUb[ubOffset], xGm[gmOffset], copyParams, dataCopyPadParams_); +} +template +__aicore__ inline void InplacePartialRotaryMulABA::CopyInDataR(LocalTensor &xUb, GlobalTensor &xGm, int64_t blockCout, int64_t blockLen, int64_t gmOffset, int64_t ubOffset) +{ + DataCopyExtParams copyParams; + copyParams.blockCount = blockCout; + copyParams.blockLen = blockLen * sizeof(R); + copyParams.srcStride = 0; //整个输入的大小 + copyParams.dstStride = 0; + DataCopyPad(xUb[ubOffset], xGm[gmOffset], copyParams, dataCopyRPadParams_); +} +template +__aicore__ inline void InplacePartialRotaryMulABA::SetGatherSrcOffset(LocalTensor &idsUb, int64_t count) +{ + for (int32_t i = 0; i < 8; ++i) { + idsUb.SetValue(i, i ^ 1); // XOR with 1 to swap even and odd indices + } + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + int32_t scalarValue = 8; + int32_t onceNum = 8; + + while (scalarValue < count) { + int32_t nextValue = scalarValue * 2; + if (nextValue < count) { + Adds(idsUb[scalarValue], idsUb, scalarValue, scalarValue); + } else { + Adds(idsUb[scalarValue], idsUb, scalarValue, count - scalarValue); + break; + } + scalarValue = nextValue; + } + Muls(idsUb, idsUb, 4, count); +} +template +__aicore__ inline void InplacePartialRotaryMulABA::ComputeMul(LocalTensor &dtsUb, LocalTensor & src0Ub, LocalTensor &src1Ub, int64_t onceA, int64_t numHead,int64_t headDim) +{ + int64_t count = numHead *headDim; + if constexpr(!isBrc) { + int64_t xtotalNum = onceA * count; + Mul(dtsUb, src0Ub, src1Ub, xtotalNum); //x*cos 非brc elewise乘 + } else { + if (headDim <= 64) { + int32_t mask = headDim; + repeatParams_.dstBlkStride = 1; + repeatParams_.src0BlkStride = 1; + repeatParams_.src1BlkStride = 1; + repeatParams_.dstRepStride = 8; + repeatParams_.src0RepStride = 8; + repeatParams_.src1RepStride = 0; + for (int64_t j =0; j(); +} +template +__aicore__ inline void InplacePartialRotaryMulABA::InterleavedInversion(int64_t count,LocalTensor &ub) +{ + // 做奇数位的*-1 + SetMaskNorm(); + int64_t fp32Mask = 64; + int64_t repeatTimes = count / 64; + int64_t remain = count % 64; + uint64_t fullMask = 0x5555555555555555; //0101010101010101 + uint64_t tailMask = 0x55; + SetVectorMask(0, fullMask); + Muls(ub, ub, -1.0f, MASK_PLACEHOLDER, repeatTimes, {1,1,8,8}); + if (remain != 0) { + int32_t tailTimes = count % 64 / 8; + SetVectorMask(0, tailMask); + Muls(ub[repeatTimes * 64], ub[repeatTimes * 64], -1.0f, MASK_PLACEHOLDER, tailTimes, {1,1,1,1}); + } + ResetMask(); +} +template +__aicore__ inline void InplacePartialRotaryMulABA::DataCopyOut(LocalTensor &yUb, GlobalTensor &xGm, int64_t blockCout, int64_t blockLen, int64_t gmOffset, int64_t ubOffset) +{ + DataCopyExtParams copyParams; + copyParams.blockCount = blockCout; + copyParams.blockLen = blockLen * sizeof(T); + copyParams.srcStride = 0; + copyParams.dstStride = (allHeadDim_ - headDim_)*sizeof(T); + DataCopyPad(yGm_[gmOffset], yUb[ubOffset], copyParams); +} + +template +__aicore__ inline void InplacePartialRotaryMulABA::Process() +{ + if (blockIdx_ >= usedCoreNum_) { + return; + } + LocalTensor idsUb = idsBuf_.Get(); + int32_t count = numHead_ * headDim_; + SetGatherSrcOffset(idsUb, count); + LocalTensor idsUbUint32 = idsBuf_.Get(); + // A分核分ub, 一次搬入x个 cout, x就是ubFactor_ + int64_t ubLoopTimes = blockIdx_ == tiling_->usedCoreNum - 1 ? coreTUbLoopTime_ :coreBUbLoopTime_; //b分核，每个核处理多少个b，b就是shape0 + int64_t ubLoopTailNum = blockIdx_ == tiling_->usedCoreNum - 1 ? coreTUbLoopTail_ :coreBUbLoopTail_; + int64_t ysCount = numHead_ *allHeadDim_; + for(int64_t i=0; i< ubLoopTimes; i++) { + int64_t ubSize = i == ubLoopTimes -1 ? ubLoopTailNum : ubFactor_; + int64_t xtotalNum = ubSize * count; + int64_t r1totalNum = ubSize * headDim_; + LocalTensor xUb = xQue_.AllocTensor(); + int64_t gmOffset = blockIdx_ * blockFactor_ * ysCount + i * ubFactor_ * ysCount + start_; + int64_t blockCout = ubSize * numHead_; + CopyInData(xUb, xGm_, blockCout , headDim_, gmOffset, halfNumx_); + xQue_.EnQue(xUb); + + LocalTensor r1Ub = r1Que_.AllocTensor(); + int64_t r1Offset = blockIdx_ * blockFactor_ * headDim_ + i * ubFactor_ * headDim_; + CopyInDataR(r1Ub, r1Gm_, ubSize , headDim_, r1Offset, ropeUbOffset_); + r1Que_.EnQue(r1Ub); + + LocalTensor r2Ub = r2Que_.AllocTensor(); + CopyInDataR(r2Ub, r2Gm_, ubSize , headDim_, r1Offset, ropeUbOffset_); + r2Que_.EnQue(r2Ub); + + xUb = xQue_.DeQue(); + r1Ub = r1Que_.DeQue(); + r2Ub = r2Que_.DeQue(); + LocalTensor xUbFp32 = xUb.template ReinterpretCast(); + LocalTensor r1UbFp32 = r1Ub.template ReinterpretCast(); + LocalTensor r2UbFp32 = r2Ub.template ReinterpretCast(); + if constexpr(sizeof(T) != sizeof(float)) { + // 非fp32时需要做cast + Cast(xUbFp32, xUb[halfNumx_], RoundMode::CAST_NONE, xtotalNum); + } + if constexpr(sizeof(R) != sizeof(float)) { + Cast(r1UbFp32, r1Ub[ropeUbOffset_], RoundMode::CAST_NONE, r1totalNum); + Cast(r2UbFp32, r2Ub[ropeUbOffset_], RoundMode::CAST_NONE, r1totalNum); + } + PipeBarrier(); + LocalTensor yUb = yQue_.AllocTensor(); + LocalTensor yUbFp32 = yUb.template ReinterpretCast(); + ComputeMul(yUbFp32 , xUbFp32, r1UbFp32, ubSize, numHead_,headDim_); + // 开始做选择的取数，基偶取数 + for (int64_t k =0; k < ubSize; k++) { + Gather(xUbFp32[k*count], xUbFp32[k*count], idsUbUint32, uint32_t(0), uint32_t(count)); //奇偶交换完成 + } + PipeBarrier(); + ComputeMul(xUbFp32 , xUbFp32, r2UbFp32, ubSize, numHead_,headDim_); + r1Que_.FreeTensor(r1Ub); + r2Que_.FreeTensor(r2Ub); + InterleavedInversion(xtotalNum, xUbFp32); + // 做最后的Add + Add(yUbFp32, yUbFp32, xUbFp32, xtotalNum); + xQue_.FreeTensor(xUb); + PipeBarrier(); + if constexpr(sizeof(T) != sizeof(float)) { + Cast(yUb, yUbFp32, RoundMode::CAST_RINT, xtotalNum); + } + yQue_.EnQue(yUb); + yUb = yQue_.DeQue(); + DataCopyOut(yUb, yGm_, blockCout, headDim_, gmOffset, 0); + yQue_.FreeTensor(yUb); + } +} + +} +#endif diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/inplace_partial_rotary_mul_common.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/inplace_partial_rotary_mul_common.h new file mode 100644 index 000000000..9e5f8462b --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/inplace_partial_rotary_mul_common.h @@ -0,0 +1,214 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file common.h + * \brief + */ +#ifndef MOE_GATING_TOP_K_COMMON_H_V35 +#define MOE_GATING_TOP_K_COMMON_H_V35 + +#include "kernel_operator.h" + +namespace ops { +using namespace AscendC; +constexpr int32_t MIN_FP32 = 0xFF800000; +constexpr int64_t ONE_REPEAT_SORT_NUM = 32; +constexpr int64_t BLOCK_BYTES = 32; + +constexpr int64_t MERGE_LIST_TWO = 2; +constexpr int64_t MERGE_LIST_THREE = 3; +constexpr int64_t MERGE_LIST_FOUR = 4; + +constexpr int64_t MERGE_LIST_IDX_TWO = 2; +constexpr int64_t MERGE_LIST_IDX_THREE = 3; + +__aicore__ inline int64_t Ceil(int64_t a, int64_t b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes) +{ + if (bytes == 0) { + return 0; + } + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes; +} + +__aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes) +{ + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES; +} + +template +__aicore__ inline T Min(T a, T b) +{ + return a > b ? b : a; +} + +template +__aicore__ inline T Max(T a, T b) +{ + return a < b ? b : a; +} + +template +__aicore__ inline T1 CeilAlign(T1 a, T2 b) { + if (b == 0) { + return 0; + } + return (a + b - 1) / b * b; +} + +template +__aicore__ inline T1 CeilDiv(T1 x, T2 y) +{ + if (y != 0 && x != 0) { + const T1 quotient = x / y; + return (x % y != 0 && ((x ^ y) >= 0)) ? (quotient + 1) : quotient; + } + + return x; +} + +template +__aicore__ inline void SetWaitFlag(HardEvent evt) +{ + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(evt)); + SetFlag(eventId); + WaitFlag(eventId); +} + +constexpr AscendC::MicroAPI::CastTrait castTraitB162B32 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::UNKNOWN, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB322B16 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_RINT, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB322Int32 = { + AscendC::MicroAPI::RegLayout::UNKNOWN, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_TRUNC, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB322Int16 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_TRUNC, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB162Int8 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_TRUNC, +}; + +// load 对齐的 bfloat16,float16,bfloat32类型的 input(ub中)数据到 float32类型的dst(寄存器)中 +template +__aicore__ inline void LoadOneTensorForDtypeT(__local_mem__ T *input, MicroAPI::RegTensor &dst, + MicroAPI::MaskReg &preg, uint32_t offset) +{ + if constexpr (IsSameType::value) { + MicroAPI::RegTensor xFp16; + DataCopy(xFp16, ((__local_mem__ half *)(input) + (offset))); + Cast(dst, xFp16, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor xBf16; + DataCopy(xBf16, + ((__local_mem__ bfloat16_t *)(input) + (offset))); + Cast(dst, xBf16, preg); + } else { + DataCopy(dst, ((__local_mem__ float *)(input) + (offset))); + } +} + +// load 2个对齐的Tensor 到寄存器中 +template +__aicore__ inline void LoadTwoTensorForDtypeT(__local_mem__ T *src1, __local_mem__ T *src2, + MicroAPI::RegTensor &dst1, MicroAPI::RegTensor &dst2, + MicroAPI::MaskReg &dst1Preg, MicroAPI::MaskReg &dst2Preg, + uint32_t src1Offset, uint32_t src2Offset) +{ + if constexpr (IsSameType::value) { + MicroAPI::RegTensor xFp16Q; + MicroAPI::RegTensor xFp16R; + DataCopy(xFp16Q, ((__local_mem__ half *)(src1) + (src1Offset))); + DataCopy(xFp16R, ((__local_mem__ half *)(src2) + (src2Offset))); + Cast(dst1, xFp16Q, dst1Preg); + Cast(dst2, xFp16R, dst2Preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor xFp16Q; + MicroAPI::RegTensor xFp16R; + DataCopy(xFp16Q, ((__local_mem__ bfloat16_t *)(src1) + (src1Offset))); + DataCopy(xFp16R, ((__local_mem__ bfloat16_t *)(src2) + (src2Offset))); + Cast(dst1, xFp16Q, dst1Preg); + Cast(dst2, xFp16R, dst2Preg); + } else { + DataCopy(dst1, ((__local_mem__ float *)(src1) + (src1Offset))); + DataCopy(dst2, ((__local_mem__ float *)(src2) + (src2Offset))); + } +} + +// store 对齐的float32类型的src(寄存器)数据到output(ub)中，output数据类型支持bfloat16,float16,bfloat32,int32_t,int16_t,int8_t,uint8_t +template +__aicore__ inline void StoreOneTensorForDtypeT(__local_mem__ T *output, MicroAPI::RegTensor &src, + MicroAPI::MaskReg &preg, uint32_t offset) +{ + if constexpr (IsSameType::value) { + MicroAPI::RegTensor yFp16; + Cast(yFp16, src, preg); + DataCopy(((__local_mem__ half *)output + offset), yFp16, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor xBf16; + Cast(xBf16, src, preg); + DataCopy(((__local_mem__ bfloat16_t *)output + offset), + xBf16, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor zInt32; + MicroAPI::Cast(zInt32, src, preg); + DataCopy(((__local_mem__ int32_t *)output + offset), zInt32, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor zInt16; + MicroAPI::Cast(zInt16, src, preg); + DataCopy(((__local_mem__ int16_t *)output + offset), zInt16, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor yFp16; + MicroAPI::RegTensor zInt8; + Cast(yFp16, src, preg); + Cast(zInt8, yFp16, preg); + DataCopy(((__local_mem__ int8_t *)output + offset), zInt8, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor yFp16; + MicroAPI::RegTensor zUint8; + Cast(yFp16, src, preg); + Cast(zUint8, yFp16, preg); + DataCopy(((__local_mem__ uint8_t *)output + offset), zUint8, preg); + } else { + DataCopy(((__local_mem__ float *)output + offset), src, preg); + } +} + +} // namespace ops +#endif // MOE_GATING_TOP_K_COMMON_H_V35 \ No newline at end of file diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_a_and_b.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_a_and_b.h new file mode 100644 index 000000000..2846c0539 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_a_and_b.h @@ -0,0 +1,294 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotary_position_embedding_reg_a_and_b.h + * \brief + */ +#ifndef ROTARY_POSITION_EMBEDDING_REG_A_AND_B_H +#define ROTARY_POSITION_EMBEDDING_REG_A_AND_B_H + +// #include "op_kernel/math_util.h" +#include "apply_rotary_pos_emb_common.h" + +namespace InplacePartialRotaryMul { +using namespace AscendC; + +template +class RotaryPositionEmbeddingAAndB +{ +public: + __aicore__ inline RotaryPositionEmbeddingAAndB(){}; + + __aicore__ inline ~RotaryPositionEmbeddingAAndB(){}; + + __aicore__ inline void Init( + GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut, GM_ADDR workspace, const RopeRegbaseTilingData* tilingData, + TPipe* pipe); + + __aicore__ inline void Process(); + +private: + // Init过程中使用的内部函数 + __aicore__ inline void InitAllGlobalBuffer(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut); + __aicore__ inline void InitAllBuffer(); + __aicore__ inline void InitLoopParams(); + // 各个层级的Process函数 + __aicore__ inline void ProcessInLoop(LocalTensor& cos, LocalTensor& sin, int64_t bStart, int64_t bLength); + // 拷入拷出函数 + __aicore__ inline void CopyInCosAndSin(int64_t bStart, int64_t bLength); + __aicore__ inline void CopyInQ(GlobalTensor& source, int64_t bStart, int64_t bLength); + __aicore__ inline void CopyOutQ(GlobalTensor& target, int64_t bStart, int64_t bLength); + + // 计算函数 + __aicore__ inline void Compute(LocalTensor& cos, LocalTensor& sin, int64_t bLength); + +private: + constexpr static uint32_t COS_DB_BUFFER = IsBoardCast ? 1 : DOUBLE_BUFFER; + + TPipe* pipe_; + + // GlobalMemory + GlobalTensor qGm_; + GlobalTensor cosGm_; + GlobalTensor sinGm_; + GlobalTensor qOutGm_; + + // UB + TQue qInQueue_; + TQue cosInQueue_; + TQue sinInQueue_; + TQue qOutQueue_; + + // Split core info + int64_t blockIdx_ = 0; + int64_t bBlockStart_ = 0; + int64_t bBlockLength_ = 0; + + // TilingData + const RopeRegbaseTilingData* tilingData_; + int64_t ubFactorB_ = 0; + int64_t D_ = 0; + int64_t dAlign_ = 0; + + // 拷贝参数 + uint8_t dSplitCoef_ = 1; + uint8_t copyInQSplitCoef_ = 1; // 拷贝q时使用的splitCoef + uint64_t ubCopyInStride = 0; // 输入在ub中的stride，deepseek_interleave中不为0 +}; + +template +__aicore__ inline void RotaryPositionEmbeddingAAndB::Init( + GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut, GM_ADDR workspace, const RopeRegbaseTilingData* tilingData, + TPipe* pipe) +{ + this->tilingData_ = tilingData; + this->pipe_ = pipe; + this->blockIdx_ = GetBlockIdx(); + this->InitAllGlobalBuffer(q, cos, sin, qOut); + this->InitAllBuffer(); + this->InitLoopParams(); +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndB::InitAllGlobalBuffer( + GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut) +{ + this->qGm_.SetGlobalBuffer((__gm__ T*)q); + this->cosGm_.SetGlobalBuffer((__gm__ T*)cos); + this->sinGm_.SetGlobalBuffer((__gm__ T*)sin); + this->qOutGm_.SetGlobalBuffer((__gm__ T*)qOut); +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndB::InitAllBuffer() +{ + this->ubFactorB_ = this->tilingData_->ubFactorB; + this->D_ = this->tilingData_->D; + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::HALF) || + tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + this->dSplitCoef_ = HALF_INTERLEAVE_COEF; + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::QUARTER)) { + this->dSplitCoef_ = QUARTER_MODE_COEF; + } + this->copyInQSplitCoef_ = dSplitCoef_; + this->dAlign_ = ops::CeilAlign(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(T)) * dSplitCoef_; + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + this->copyInQSplitCoef_ = 1; + // 非boardcast时，使用批量计算API，需要拷贝时添加stride + if constexpr (!IsBoardCast) { + this->ubCopyInStride = + (this->dAlign_ * sizeof(T) - ops::CeilAlign(tilingData_->sliceLength * sizeof(T), BLOCK_TYPE_SIZE)) / + BLOCK_TYPE_SIZE; + } + } + + this->pipe_->InitBuffer(this->qInQueue_, DOUBLE_BUFFER, ubFactorB_ * dAlign_ * sizeof(T)); + this->pipe_->InitBuffer(this->qOutQueue_, DOUBLE_BUFFER, ubFactorB_ * dAlign_ * sizeof(T)); + if constexpr (IsBoardCast) { + this->pipe_->InitBuffer(this->cosInQueue_, COS_DB_BUFFER, dAlign_ * sizeof(T)); + this->pipe_->InitBuffer(this->sinInQueue_, COS_DB_BUFFER, dAlign_ * sizeof(T)); + } else { + this->pipe_->InitBuffer(this->cosInQueue_, COS_DB_BUFFER, ubFactorB_ * dAlign_ * sizeof(T)); + this->pipe_->InitBuffer(this->sinInQueue_, COS_DB_BUFFER, ubFactorB_ * dAlign_ * sizeof(T)); + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndB::InitLoopParams() +{ + this->bBlockLength_ = tilingData_->blockFactorB; + if (blockIdx_ == tilingData_->blockNumB - 1 && tilingData_->B % tilingData_->blockFactorB != 0) { + this->bBlockLength_ = tilingData_->B % tilingData_->blockFactorB; + } + this->bBlockStart_ = blockIdx_ * tilingData_->blockFactorB; +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndB::Process() +{ + // 在B轴进行循环 + int64_t ubLoopCount = ops::CeilDiv(bBlockLength_, ubFactorB_); + if constexpr (IsBoardCast) { + this->CopyInCosAndSin(0, 1); + LocalTensor cosUb = this->cosInQueue_.template DeQue(); + LocalTensor sinUb = this->sinInQueue_.template DeQue(); + for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) { + this->ProcessInLoop( + cosUb, sinUb, bBlockStart_ + ubLoopIdx * ubFactorB_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_); + } + this->cosInQueue_.FreeTensor(cosUb); + this->sinInQueue_.FreeTensor(sinUb); + } else { + for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) { + this->CopyInCosAndSin( + bBlockStart_ + ubLoopIdx * ubFactorB_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_); + LocalTensor cosUb = this->cosInQueue_.template DeQue(); + LocalTensor sinUb = this->sinInQueue_.template DeQue(); + this->ProcessInLoop( + cosUb, sinUb, bBlockStart_ + ubLoopIdx * ubFactorB_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_); + this->cosInQueue_.FreeTensor(cosUb); + this->sinInQueue_.FreeTensor(sinUb); + } + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndB::ProcessInLoop( + LocalTensor& cos, LocalTensor& sin, int64_t bUbStart, int64_t bUbLength) +{ + CopyInQ(qGm_, bUbStart, bUbLength); + Compute(cos, sin, bUbLength); + CopyOutQ(qOutGm_, bUbStart, bUbLength); +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndB::CopyInCosAndSin(int64_t bStart, int64_t bLength) +{ + LocalTensor cosUb = this->cosInQueue_.template AllocTensor(); + LocalTensor sinUb = this->sinInQueue_.template AllocTensor(); + DataCopyPadExtParams copyPadExtparams; + copyPadExtparams.isPad = false; + copyPadExtparams.leftPadding = 0; + copyPadExtparams.rightPadding = 0; + copyPadExtparams.paddingValue = 0; + DataCopyExtParams copyExtParams; + copyExtParams.blockCount = bLength * dSplitCoef_; + copyExtParams.blockLen = tilingData_->sliceLength * sizeof(T) / dSplitCoef_; + copyExtParams.srcStride = 0; + copyExtParams.dstStride = 0; + DataCopyPad(cosUb, this->cosGm_[bStart * tilingData_->sliceLength], copyExtParams, copyPadExtparams); + DataCopyPad(sinUb, this->sinGm_[bStart * tilingData_->sliceLength], copyExtParams, copyPadExtparams); + this->cosInQueue_.template EnQue(cosUb); + this->sinInQueue_.template EnQue(sinUb); +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndB::CopyInQ( + GlobalTensor& source, int64_t bStart, int64_t bLength) +{ + LocalTensor target = this->qInQueue_.template AllocTensor(); + DataCopyExtParams copyExtParams; + copyExtParams.blockCount = bLength * copyInQSplitCoef_; + copyExtParams.blockLen = tilingData_->sliceLength * sizeof(T) / copyInQSplitCoef_; + copyExtParams.srcStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(T); + copyExtParams.dstStride = ubCopyInStride; + DataCopyPadExtParams copyPadExtparams; + copyPadExtparams.isPad = false; + copyPadExtparams.leftPadding = 0; + copyPadExtparams.rightPadding = 0; + copyPadExtparams.paddingValue = 0; + DataCopyPad(target, source[bStart * D_ + tilingData_->sliceStart], copyExtParams, copyPadExtparams); + this->qInQueue_.template EnQue(target); +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndB::CopyOutQ( + GlobalTensor& target, int64_t bStart, int64_t bLength) +{ + LocalTensor source = this->qOutQueue_.template DeQue(); + DataCopyExtParams copyExtParams; + copyExtParams.blockCount = bLength * dSplitCoef_; + copyExtParams.blockLen = tilingData_->sliceLength * sizeof(T) / dSplitCoef_; + copyExtParams.srcStride = 0; + copyExtParams.dstStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(T); + DataCopyPad(target[bStart * D_+ tilingData_->sliceStart], source, copyExtParams); + this->qOutQueue_.FreeTensor(source); +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndB::Compute( + LocalTensor& cos, LocalTensor& sin, int64_t bLength) +{ + LocalTensor inUb = this->qInQueue_.template DeQue(); + LocalTensor outUb = this->qOutQueue_.template AllocTensor(); + if constexpr (IsBoardCast) { + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::HALF)) { + HalfAlignVF(sin, cos, inUb, outUb, tilingData_->sliceLength, dAlign_, 1, bLength); + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::INTERLEAVE)) { + InterleaveModeVF(sin, cos, inUb, outUb, tilingData_->sliceLength, 1, bLength); + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::QUARTER)) { + QuarterAlignVF(sin, cos, inUb, outUb, tilingData_->sliceLength, dAlign_, 1, bLength); + } else { + DeepSeekInterleaveModeVF(sin, cos, inUb, outUb, tilingData_->sliceLength, 1, bLength); + } + } else { + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::HALF)) { + BatchHalfAlignVF( + (__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(), + (__local_mem__ T*)sin.GetPhyAddr(), (__local_mem__ T*)outUb.GetPhyAddr(), bLength, 1, 1, tilingData_->sliceLength, dAlign_, + ubFactorB_, 1); + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::INTERLEAVE)) { + BatchInterleaveModeVF( + (__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(), + (__local_mem__ T*)sin.GetPhyAddr(), (__local_mem__ T*)outUb.GetPhyAddr(), bLength, 1, 1, tilingData_->sliceLength, dAlign_, + ubFactorB_, 1); + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::QUARTER)) { + BatchQuarterAlignVF( + (__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(), + (__local_mem__ T*)sin.GetPhyAddr(), (__local_mem__ T*)outUb.GetPhyAddr(), bLength, 1, 1, tilingData_->sliceLength, dAlign_, + ubFactorB_, 1); + } else { + BatchDeepSeekInterleaveModeVF( + (__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(), + (__local_mem__ T*)sin.GetPhyAddr(), (__local_mem__ T*)outUb.GetPhyAddr(), bLength, 1, 1, tilingData_->sliceLength, dAlign_, + ubFactorB_, 1); + } + } + + this->qInQueue_.FreeTensor(inUb); + this->qOutQueue_.template EnQue(outUb); +} +} // namespace InplacePartialRotaryMul + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_a_and_b_mixed.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_a_and_b_mixed.h new file mode 100644 index 000000000..0b7e6dc6f --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_a_and_b_mixed.h @@ -0,0 +1,273 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotary_position_embedding_reg_a_and_b_mixed.h + * \brief Mixed precision kernel for AAndB layout: x is half/bfloat16, cos/sin are float + */ +#ifndef ROTARY_POSITION_EMBEDDING_REG_A_AND_B_MIXED_H +#define ROTARY_POSITION_EMBEDDING_REG_A_AND_B_MIXED_H + +#include "apply_rotary_pos_emb_common.h" + +namespace InplacePartialRotaryMul { +using namespace AscendC; + +template +class RotaryPositionEmbeddingAAndBMixed { +public: + __aicore__ inline RotaryPositionEmbeddingAAndBMixed(){}; + + __aicore__ inline ~RotaryPositionEmbeddingAAndBMixed(){}; + + __aicore__ inline void Init(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut, GM_ADDR workspace, + const RopeRegbaseTilingData *tilingData, TPipe *pipe); + + __aicore__ inline void Process(); + +private: + __aicore__ inline void InitAllGlobalBuffer(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut); + __aicore__ inline void InitAllBuffer(); + __aicore__ inline void InitLoopParams(); + __aicore__ inline void ProcessInLoop( + LocalTensor &cos, LocalTensor &sin, int64_t bStart, int64_t bLength); + __aicore__ inline void CopyInCosAndSin(int64_t bStart, int64_t bLength); + __aicore__ inline void CopyInQ(GlobalTensor &source, int64_t bStart, int64_t bLength); + __aicore__ inline void CopyOutQ(GlobalTensor &target, int64_t bStart, int64_t bLength); + + __aicore__ inline void Compute(LocalTensor &cos, LocalTensor &sin, int64_t bLength); + +private: + static constexpr uint32_t COS_DB_BUFFER = IsBoardCast ? 1 : DOUBLE_BUFFER; + + TPipe *pipe_; + + GlobalTensor qGm_; + GlobalTensor cosGm_; + GlobalTensor sinGm_; + GlobalTensor qOutGm_; + + TQue qInQueue_; + TQue cosInQueue_; + TQue sinInQueue_; + TQue qOutQueue_; + + int64_t blockIdx_ = 0; + int64_t bBlockStart_ = 0; + int64_t bBlockLength_ = 0; + + const RopeRegbaseTilingData *tilingData_; + int64_t ubFactorB_ = 0; + int64_t D_ = 0; + int64_t dAlign_ = 0; + int64_t dAlignFloat_ = 0; + uint8_t dSplitCoef_ = 1; + uint8_t copyInQSplitCoef_ = 1; + uint64_t ubCopyInStride = 0; +}; + +template +__aicore__ inline void RotaryPositionEmbeddingAAndBMixed::Init(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, + GM_ADDR qOut, GM_ADDR workspace, const RopeRegbaseTilingData *tilingData, TPipe *pipe) +{ + this->tilingData_ = tilingData; + this->pipe_ = pipe; + this->blockIdx_ = GetBlockIdx(); + this->InitAllGlobalBuffer(q, cos, sin, qOut); + this->InitAllBuffer(); + this->InitLoopParams(); +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndBMixed::InitAllGlobalBuffer( + GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut) +{ + this->qGm_.SetGlobalBuffer((__gm__ TX *)q); + this->cosGm_.SetGlobalBuffer((__gm__ float *)cos); + this->sinGm_.SetGlobalBuffer((__gm__ float *)sin); + this->qOutGm_.SetGlobalBuffer((__gm__ TX *)qOut); +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndBMixed::InitAllBuffer() +{ + this->ubFactorB_ = this->tilingData_->ubFactorB; + this->D_ = this->tilingData_->D; + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::HALF) || + tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + this->dSplitCoef_ = HALF_INTERLEAVE_COEF; + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::QUARTER)) { + this->dSplitCoef_ = QUARTER_MODE_COEF; + } + this->copyInQSplitCoef_ = dSplitCoef_; + this->dAlign_ = + ops::CeilAlign(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(TX)) * dSplitCoef_; + this->dAlignFloat_ = + ops::CeilAlign(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(float)) * dSplitCoef_; + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + this->copyInQSplitCoef_ = 1; + if constexpr (!IsBoardCast) { + this->ubCopyInStride = + (this->dAlign_ * sizeof(TX) - + ops::CeilAlign(tilingData_->sliceLength * sizeof(TX), BLOCK_TYPE_SIZE)) / + BLOCK_TYPE_SIZE; + } + } + + this->pipe_->InitBuffer(this->qInQueue_, DOUBLE_BUFFER, ubFactorB_ * dAlign_ * sizeof(TX)); + this->pipe_->InitBuffer(this->qOutQueue_, DOUBLE_BUFFER, ubFactorB_ * dAlign_ * sizeof(TX)); + if constexpr (IsBoardCast) { + this->pipe_->InitBuffer(this->cosInQueue_, COS_DB_BUFFER, dAlignFloat_ * sizeof(float)); + this->pipe_->InitBuffer(this->sinInQueue_, COS_DB_BUFFER, dAlignFloat_ * sizeof(float)); + } else { + this->pipe_->InitBuffer(this->cosInQueue_, COS_DB_BUFFER, ubFactorB_ * dAlignFloat_ * sizeof(float)); + this->pipe_->InitBuffer(this->sinInQueue_, COS_DB_BUFFER, ubFactorB_ * dAlignFloat_ * sizeof(float)); + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndBMixed::InitLoopParams() +{ + this->bBlockLength_ = tilingData_->blockFactorB; + if (blockIdx_ == tilingData_->blockNumB - 1 && tilingData_->B % tilingData_->blockFactorB != 0) { + this->bBlockLength_ = tilingData_->B % tilingData_->blockFactorB; + } + this->bBlockStart_ = blockIdx_ * tilingData_->blockFactorB; +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndBMixed::Process() +{ + int64_t ubLoopCount = ops::CeilDiv(bBlockLength_, ubFactorB_); + if constexpr (IsBoardCast) { + this->CopyInCosAndSin(0, 1); + LocalTensor cosUb = this->cosInQueue_.template DeQue(); + LocalTensor sinUb = this->sinInQueue_.template DeQue(); + for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) { + this->ProcessInLoop(cosUb, + sinUb, + bBlockStart_ + ubLoopIdx * ubFactorB_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_); + } + this->cosInQueue_.FreeTensor(cosUb); + this->sinInQueue_.FreeTensor(sinUb); + } else { + for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) { + this->CopyInCosAndSin(bBlockStart_ + ubLoopIdx * ubFactorB_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_); + LocalTensor cosUb = this->cosInQueue_.template DeQue(); + LocalTensor sinUb = this->sinInQueue_.template DeQue(); + this->ProcessInLoop(cosUb, + sinUb, + bBlockStart_ + ubLoopIdx * ubFactorB_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_); + this->cosInQueue_.FreeTensor(cosUb); + this->sinInQueue_.FreeTensor(sinUb); + } + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndBMixed::ProcessInLoop( + LocalTensor &cos, LocalTensor &sin, int64_t bUbStart, int64_t bUbLength) +{ + CopyInQ(qGm_, bUbStart, bUbLength); + Compute(cos, sin, bUbLength); + CopyOutQ(qOutGm_, bUbStart, bUbLength); +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndBMixed::CopyInCosAndSin( + int64_t bStart, int64_t bLength) +{ + LocalTensor cosUb = this->cosInQueue_.template AllocTensor(); + LocalTensor sinUb = this->sinInQueue_.template AllocTensor(); + DataCopyPadExtParams copyPadExtparams; + copyPadExtparams.isPad = false; + copyPadExtparams.leftPadding = 0; + copyPadExtparams.rightPadding = 0; + copyPadExtparams.paddingValue = 0; + DataCopyExtParams copyExtParams; + if constexpr (IsBoardCast) { + copyExtParams.blockCount = 1 * dSplitCoef_; + } else { + copyExtParams.blockCount = bLength * dSplitCoef_; + } + copyExtParams.blockLen = tilingData_->sliceLength * sizeof(float) / dSplitCoef_; + copyExtParams.srcStride = 0; + copyExtParams.dstStride = 0; + DataCopyPad(cosUb, this->cosGm_[bStart * tilingData_->sliceLength], copyExtParams, copyPadExtparams); + DataCopyPad(sinUb, this->sinGm_[bStart * tilingData_->sliceLength], copyExtParams, copyPadExtparams); + this->cosInQueue_.template EnQue(cosUb); + this->sinInQueue_.template EnQue(sinUb); +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndBMixed::CopyInQ( + GlobalTensor &source, int64_t bStart, int64_t bLength) +{ + LocalTensor target = this->qInQueue_.template AllocTensor(); + DataCopyExtParams copyExtParams; + copyExtParams.blockCount = bLength * copyInQSplitCoef_; + copyExtParams.blockLen = tilingData_->sliceLength * sizeof(TX) / copyInQSplitCoef_; + copyExtParams.srcStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(TX); + copyExtParams.dstStride = ubCopyInStride; + DataCopyPadExtParams copyPadExtparams; + copyPadExtparams.isPad = false; + copyPadExtparams.leftPadding = 0; + copyPadExtparams.rightPadding = 0; + copyPadExtparams.paddingValue = 0; + DataCopyPad(target, source[bStart * D_ + tilingData_->sliceStart], copyExtParams, copyPadExtparams); + this->qInQueue_.template EnQue(target); +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndBMixed::CopyOutQ( + GlobalTensor &target, int64_t bStart, int64_t bLength) +{ + LocalTensor source = this->qOutQueue_.template DeQue(); + DataCopyExtParams copyExtParams; + copyExtParams.blockCount = bLength * dSplitCoef_; + copyExtParams.blockLen = tilingData_->sliceLength * sizeof(TX) / dSplitCoef_; + copyExtParams.srcStride = 0; + copyExtParams.dstStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(TX); + DataCopyPad(target[bStart * D_ + tilingData_->sliceStart], source, copyExtParams); + this->qOutQueue_.FreeTensor(source); +} + +template +__aicore__ inline void RotaryPositionEmbeddingAAndBMixed::Compute( + LocalTensor &cos, LocalTensor &sin, int64_t bLength) +{ + LocalTensor inUb = this->qInQueue_.template DeQue(); + LocalTensor outUb = this->qOutQueue_.template AllocTensor(); + if constexpr (IsBoardCast) { + InterleaveModeVFMixed(inUb, cos, sin, outUb, tilingData_->sliceLength, 1, bLength); + } else { + BatchInterleaveModeVFMixed((__local_mem__ TX *)inUb.GetPhyAddr(), + (__local_mem__ float *)cos.GetPhyAddr(), + (__local_mem__ float *)sin.GetPhyAddr(), + (__local_mem__ TX *)outUb.GetPhyAddr(), + bLength, + 1, + 1, + tilingData_->sliceLength, + dAlign_, + dAlignFloat_, + ubFactorB_, + 1); + } + + this->qInQueue_.FreeTensor(inUb); + this->qOutQueue_.template EnQue(outUb); +} +} // namespace InplacePartialRotaryMul + +#endif diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_ab.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_ab.h new file mode 100644 index 000000000..bd6a91738 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_ab.h @@ -0,0 +1,164 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotary_position_embedding_reg_ab.h + * \brief + */ +#ifndef ROTARY_POSITION_EMBEDDING_REG_AB_H +#define ROTARY_POSITION_EMBEDDING_REG_AB_H + +// #include "op_kernel/math_util.h" +// #include "op_kernel/platform_util.h" +// #include "op_kernel/load_store_utils.h" +#include "apply_rotary_pos_emb_common.h" + +namespace InplacePartialRotaryMul { +using namespace AscendC; + +template +class RotaryPositionEmbeddingAB +{ +public: + __aicore__ inline RotaryPositionEmbeddingAB(){}; + __aicore__ inline void Init( + GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, GM_ADDR workspace, const RopeRegbaseTilingData* tilingData, + TPipe* pipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void ProcessLoop( + int64_t xGmOffset, LocalTensor cosBuffer, LocalTensor sinBuffer, int64_t ubIdx, int64_t bsCount, + int64_t nCount); + +private: + TPipe* pipe_; + TQue xInQueue_; + TQue cosInQueue_; + TQue sinInQueue_; + TQue yOutQueue_; + + GlobalTensor xGm_; + GlobalTensor cosGm_; + GlobalTensor sinGm_; + GlobalTensor yGm_; + const RopeRegbaseTilingData* tilingData_; + DataCopyPadExtParams padParams_ = {false, 0, 0, static_cast(0)}; + uint8_t DB_FLAG = 2; + uint32_t dSplitSize_ = 0; + int64_t bsBlockCount_ = 0; + int64_t nBlockCount_ = 0; + int64_t sliceAlign_ = 0; +}; + +template +__aicore__ inline void RotaryPositionEmbeddingAB::Init( + GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, GM_ADDR workspace, const RopeRegbaseTilingData* tilingData, + TPipe* pipe) +{ + pipe_ = pipe; + tilingData_ = tilingData; + dSplitSize_ = tilingData_->sliceLength / tilingData_->dSplitCoef * sizeof(T); + int64_t blockDimBS = GetBlockIdx() / tilingData_->blockNumN; + int64_t blockDimN = GetBlockIdx() % tilingData_->blockNumN; + bsBlockCount_ = (blockDimBS == tilingData_->blockNumBS - 1) ? tilingData_->blockTailBS : tilingData_->blockFactorBS; + nBlockCount_ = (blockDimN == tilingData_->blockNumN - 1) ? tilingData_->blockTailN : tilingData_->blockFactorN; + + int64_t cosOffset = blockDimBS * tilingData_->blockFactorBS * tilingData_->sliceLength; + int64_t offset = blockDimBS * tilingData_->blockFactorBS * tilingData_->D; + int64_t xOffset = offset * tilingData_->N + blockDimN * tilingData_->blockFactorN * tilingData_->D + tilingData_->sliceStart; + this->cosGm_.SetGlobalBuffer((__gm__ T*)cos + cosOffset); + this->sinGm_.SetGlobalBuffer((__gm__ T*)sin + cosOffset); + this->xGm_.SetGlobalBuffer((__gm__ T*)x + xOffset); + this->yGm_.SetGlobalBuffer((__gm__ T*)y + xOffset); + + sliceAlign_ = ops::CeilDiv(tilingData_->sliceLength * sizeof(T), GetUbBlockSize()) * GetUbBlockSize() / sizeof(T); + int64_t bufferSize = sliceAlign_ * sizeof(T) * tilingData_->ubFactorBS; + pipe_->InitBuffer(xInQueue_, DB_FLAG, bufferSize * tilingData_->ubFactorN); + pipe_->InitBuffer(cosInQueue_, DB_FLAG, bufferSize); + pipe_->InitBuffer(sinInQueue_, DB_FLAG, bufferSize); + pipe_->InitBuffer(yOutQueue_, DB_FLAG, bufferSize * tilingData_->ubFactorN); +} + +template +__aicore__ inline void RotaryPositionEmbeddingAB::Process() +{ + uint32_t bsLoopCnt = ops::CeilDiv(bsBlockCount_, tilingData_->ubFactorBS); + uint32_t nLoopCnt = ops::CeilDiv(nBlockCount_, tilingData_->ubFactorN); + for (uint32_t bsLoopIdx = 0; bsLoopIdx < bsLoopCnt; bsLoopIdx++) { + int64_t xGmOffset = bsLoopIdx * tilingData_->ubFactorBS * tilingData_->N * tilingData_->D; + uint32_t currBSNum = (bsLoopIdx != bsLoopCnt - 1) ? tilingData_->ubFactorBS : + bsBlockCount_ - (bsLoopIdx * tilingData_->ubFactorBS); + + DataCopyExtParams cosParams = { + static_cast(currBSNum * tilingData_->dSplitCoef), dSplitSize_, 0, 0, 0}; + + LocalTensor cosBuffer = cosInQueue_.AllocTensor(); + LocalTensor sinBuffer = sinInQueue_.AllocTensor(); + DataCopyPad(cosBuffer, cosGm_[bsLoopIdx * tilingData_->ubFactorBS * tilingData_->sliceLength], cosParams, padParams_); + cosInQueue_.EnQue(cosBuffer); + cosBuffer = cosInQueue_.DeQue(); + DataCopyPad(sinBuffer, sinGm_[bsLoopIdx * tilingData_->ubFactorBS * tilingData_->sliceLength], cosParams, padParams_); + sinInQueue_.EnQue(sinBuffer); + sinBuffer = sinInQueue_.DeQue(); + + for (int64_t nLoopIdx = 0; nLoopIdx < nLoopCnt; nLoopIdx++) { + int64_t currNNum = (nLoopIdx != nLoopCnt - 1) ? tilingData_->ubFactorN : + nBlockCount_ - (nLoopIdx * tilingData_->ubFactorN); + ProcessLoop(xGmOffset, cosBuffer, sinBuffer, nLoopIdx, currBSNum, currNNum); + } + + cosInQueue_.FreeTensor(cosBuffer); + sinInQueue_.FreeTensor(sinBuffer); + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingAB::ProcessLoop( + int64_t xGmOffset, LocalTensor cosBuffer, LocalTensor sinBuffer, int64_t ubIdx, int64_t bsCount, + int64_t nCount) +{ + int64_t totalCount = bsCount * nCount; + DataCopyExtParams inParams = {static_cast(totalCount * tilingData_->dSplitCoef), dSplitSize_, static_cast((tilingData_->D - tilingData_->sliceLength) * sizeof(T)), 0, 0}; + DataCopyExtParams outParams = {static_cast(totalCount * tilingData_->dSplitCoef), dSplitSize_, 0, static_cast((tilingData_->D - tilingData_->sliceLength) * sizeof(T)), 0}; + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + inParams = {static_cast(totalCount), tilingData_->D * sizeof(T), static_cast((tilingData_->D - tilingData_->sliceLength) * sizeof(T)), 0, 0}; + } + + LocalTensor inBuffer = xInQueue_.AllocTensor(); + LocalTensor outBuffer = yOutQueue_.AllocTensor(); + + DataCopyPad(inBuffer, xGm_[xGmOffset + ubIdx * tilingData_->ubFactorN * tilingData_->D], inParams, padParams_); + + xInQueue_.EnQue(inBuffer); + inBuffer = xInQueue_.DeQue(); + + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::HALF)) { + HalfAlignVF(sinBuffer, cosBuffer, inBuffer, outBuffer, tilingData_->sliceLength, sliceAlign_, bsCount, nCount); + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::INTERLEAVE)) { + InterleaveModeVF(sinBuffer, cosBuffer, inBuffer, outBuffer, tilingData_->sliceLength, bsCount, nCount); + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::QUARTER)) { + QuarterAlignVF(sinBuffer, cosBuffer, inBuffer, outBuffer, tilingData_->sliceLength, sliceAlign_, bsCount, nCount); + } else { + DeepSeekInterleaveModeVF(sinBuffer, cosBuffer, inBuffer, outBuffer, tilingData_->sliceLength, bsCount, nCount); + } + + yOutQueue_.EnQue(outBuffer); + outBuffer = yOutQueue_.DeQue(); + xInQueue_.FreeTensor(inBuffer); + + DataCopyPad(yGm_[xGmOffset + ubIdx * tilingData_->ubFactorN * tilingData_->D], outBuffer, outParams); + + yOutQueue_.FreeTensor(outBuffer); +} + +} // namespace InplacePartialRotaryMul + +#endif // ROTARY_POSITION_EMBEDDING_REG_AB_H \ No newline at end of file diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_ab_mixed.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_ab_mixed.h new file mode 100644 index 000000000..7685b3ce9 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_ab_mixed.h @@ -0,0 +1,175 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotary_position_embedding_reg_ab_mixed.h + * \brief Mixed precision kernel for AB layout: x is half/bfloat16, cos/sin are float + */ +#ifndef ROTARY_POSITION_EMBEDDING_REG_AB_MIXED_H +#define ROTARY_POSITION_EMBEDDING_REG_AB_MIXED_H + +#include "apply_rotary_pos_emb_common.h" + +namespace InplacePartialRotaryMul { +using namespace AscendC; + +template +class RotaryPositionEmbeddingABMixed { +public: + __aicore__ inline RotaryPositionEmbeddingABMixed(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, GM_ADDR workspace, + const RopeRegbaseTilingData *tilingData, TPipe *pipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void ProcessLoop(int64_t xGmOffset, LocalTensor cosBuffer, LocalTensor sinBuffer, + int64_t ubIdx, int64_t bsCount, int64_t nCount); + +private: + TPipe *pipe_; + TQue xInQueue_; + TQue cosInQueue_; + TQue sinInQueue_; + TQue yOutQueue_; + + GlobalTensor xGm_; + GlobalTensor cosGm_; + GlobalTensor sinGm_; + GlobalTensor yGm_; + const RopeRegbaseTilingData *tilingData_; + DataCopyPadExtParams padParams_ = {false, 0, 0, static_cast(0)}; + DataCopyPadExtParams padParamsFloat_ = {false, 0, 0, 0}; + uint8_t DB_FLAG = 2; + uint32_t dSplitSizeTX_ = 0; + uint32_t dSplitSizeFloat_ = 0; + int64_t bsBlockCount_ = 0; + int64_t nBlockCount_ = 0; + int64_t sliceAlignTX_ = 0; + int64_t sliceAlignFloat_ = 0; +}; + +template +__aicore__ inline void RotaryPositionEmbeddingABMixed::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + GM_ADDR workspace, const RopeRegbaseTilingData *tilingData, TPipe *pipe) +{ + pipe_ = pipe; + tilingData_ = tilingData; + dSplitSizeTX_ = tilingData_->sliceLength / tilingData_->dSplitCoef * sizeof(TX); + dSplitSizeFloat_ = tilingData_->sliceLength / tilingData_->dSplitCoef * sizeof(float); + int64_t blockDimBS = GetBlockIdx() / tilingData_->blockNumN; + int64_t blockDimN = GetBlockIdx() % tilingData_->blockNumN; + bsBlockCount_ = (blockDimBS == tilingData_->blockNumBS - 1) ? tilingData_->blockTailBS : tilingData_->blockFactorBS; + nBlockCount_ = (blockDimN == tilingData_->blockNumN - 1) ? tilingData_->blockTailN : tilingData_->blockFactorN; + + int64_t cosOffset = blockDimBS * tilingData_->blockFactorBS * tilingData_->sliceLength; + int64_t offset = blockDimBS * tilingData_->blockFactorBS * tilingData_->D; + int64_t xOffset = + offset * tilingData_->N + blockDimN * tilingData_->blockFactorN * tilingData_->D + tilingData_->sliceStart; + this->cosGm_.SetGlobalBuffer((__gm__ float *)cos + cosOffset); + this->sinGm_.SetGlobalBuffer((__gm__ float *)sin + cosOffset); + this->xGm_.SetGlobalBuffer((__gm__ TX *)x + xOffset); + this->yGm_.SetGlobalBuffer((__gm__ TX *)y + xOffset); + + sliceAlignTX_ = + ops::CeilDiv(tilingData_->sliceLength * sizeof(TX), GetUbBlockSize()) * GetUbBlockSize() / sizeof(TX); + sliceAlignFloat_ = + ops::CeilDiv(tilingData_->sliceLength * sizeof(float), GetUbBlockSize()) * GetUbBlockSize() / sizeof(float); + int64_t bufferSizeTX = sliceAlignTX_ * sizeof(TX) * tilingData_->ubFactorBS; + int64_t bufferSizeFloat = sliceAlignFloat_ * sizeof(float) * tilingData_->ubFactorBS; + pipe_->InitBuffer(xInQueue_, DB_FLAG, bufferSizeTX * tilingData_->ubFactorN); + pipe_->InitBuffer(yOutQueue_, DB_FLAG, bufferSizeTX * tilingData_->ubFactorN); + pipe_->InitBuffer(cosInQueue_, DB_FLAG, bufferSizeFloat); + pipe_->InitBuffer(sinInQueue_, DB_FLAG, bufferSizeFloat); +} + +template +__aicore__ inline void RotaryPositionEmbeddingABMixed::Process() +{ + uint32_t bsLoopCnt = ops::CeilDiv(bsBlockCount_, tilingData_->ubFactorBS); + uint32_t nLoopCnt = ops::CeilDiv(nBlockCount_, tilingData_->ubFactorN); + for (uint32_t bsLoopIdx = 0; bsLoopIdx < bsLoopCnt; bsLoopIdx++) { + int64_t xGmOffset = bsLoopIdx * tilingData_->ubFactorBS * tilingData_->N * tilingData_->D; + uint32_t currBSNum = (bsLoopIdx != bsLoopCnt - 1) ? tilingData_->ubFactorBS + : bsBlockCount_ - (bsLoopIdx * tilingData_->ubFactorBS); + + DataCopyExtParams cosParams = { + static_cast(currBSNum * tilingData_->dSplitCoef), dSplitSizeFloat_, 0, 0, 0}; + + LocalTensor cosBuffer = cosInQueue_.AllocTensor(); + LocalTensor sinBuffer = sinInQueue_.AllocTensor(); + DataCopyPad(cosBuffer, + cosGm_[bsLoopIdx * tilingData_->ubFactorBS * tilingData_->sliceLength], + cosParams, + padParamsFloat_); + cosInQueue_.EnQue(cosBuffer); + cosBuffer = cosInQueue_.DeQue(); + DataCopyPad(sinBuffer, + sinGm_[bsLoopIdx * tilingData_->ubFactorBS * tilingData_->sliceLength], + cosParams, + padParamsFloat_); + sinInQueue_.EnQue(sinBuffer); + sinBuffer = sinInQueue_.DeQue(); + + for (int64_t nLoopIdx = 0; nLoopIdx < nLoopCnt; nLoopIdx++) { + int64_t currNNum = (nLoopIdx != nLoopCnt - 1) ? tilingData_->ubFactorN + : nBlockCount_ - (nLoopIdx * tilingData_->ubFactorN); + ProcessLoop(xGmOffset, cosBuffer, sinBuffer, nLoopIdx, currBSNum, currNNum); + } + + cosInQueue_.FreeTensor(cosBuffer); + sinInQueue_.FreeTensor(sinBuffer); + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingABMixed::ProcessLoop(int64_t xGmOffset, LocalTensor cosBuffer, + LocalTensor sinBuffer, int64_t ubIdx, int64_t bsCount, int64_t nCount) +{ + int64_t totalCount = bsCount * nCount; + DataCopyExtParams inParams = {static_cast(totalCount * tilingData_->dSplitCoef), + dSplitSizeTX_, + static_cast((tilingData_->D - tilingData_->sliceLength) * sizeof(TX)), + 0, + 0}; + DataCopyExtParams outParams = {static_cast(totalCount * tilingData_->dSplitCoef), + dSplitSizeTX_, + 0, + static_cast((tilingData_->D - tilingData_->sliceLength) * sizeof(TX)), + 0}; + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + inParams = {static_cast(totalCount), + tilingData_->D * sizeof(TX), + static_cast((tilingData_->D - tilingData_->sliceLength) * sizeof(TX)), + 0, + 0}; + } + + LocalTensor inBuffer = xInQueue_.AllocTensor(); + LocalTensor outBuffer = yOutQueue_.AllocTensor(); + + DataCopyPad(inBuffer, xGm_[xGmOffset + ubIdx * tilingData_->ubFactorN * tilingData_->D], inParams, padParams_); + + xInQueue_.EnQue(inBuffer); + inBuffer = xInQueue_.DeQue(); + + InterleaveModeVFMixed(inBuffer, cosBuffer, sinBuffer, outBuffer, tilingData_->sliceLength, bsCount, nCount); + + yOutQueue_.EnQue(outBuffer); + outBuffer = yOutQueue_.DeQue(); + xInQueue_.FreeTensor(inBuffer); + + DataCopyPad(yGm_[xGmOffset + ubIdx * tilingData_->ubFactorN * tilingData_->D], outBuffer, outParams); + + yOutQueue_.FreeTensor(outBuffer); +} + +} // namespace InplacePartialRotaryMul + +#endif // ROTARY_POSITION_EMBEDDING_REG_AB_MIXED_H diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_aba_and_ba.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_aba_and_ba.h new file mode 100644 index 000000000..d26dcaaae --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_aba_and_ba.h @@ -0,0 +1,364 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotary_position_embedding_reg_aba_and_ba.h + * \brief + */ +#ifndef ROTARY_POSITION_EMBEDDING_REG_ABA_AND_BA_H +#define ROTARY_POSITION_EMBEDDING_REG_ABA_AND_BA_H + +// #include "op_kernel/math_util.h" +#include "apply_rotary_pos_emb_common.h" + +namespace InplacePartialRotaryMul { +using namespace AscendC; + +template +class RotaryPositionEmbeddingABAAndBA +{ +public: + __aicore__ inline RotaryPositionEmbeddingABAAndBA(){}; + + __aicore__ inline ~RotaryPositionEmbeddingABAAndBA(){}; + + __aicore__ inline void Init( + GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut, GM_ADDR workspace, const RopeRegbaseTilingData* tilingData, + TPipe* pipe); + + __aicore__ inline void Process(); + +private: + // Init过程中使用的内部函数 + __aicore__ inline void InitAllGlobalBuffer(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut); + __aicore__ inline void InitAllBuffer(); + __aicore__ inline void InitLoopParams(); + // 各个层级的Process函数 + __aicore__ inline void ProcessInSLoop( + int64_t sUbStart, + int64_t sUbLength); // 第一重循环体，给定S范围，沿B轴进行遍历处理 + __aicore__ inline void ProcessInSBLoop( + int64_t sUbStart, int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, LocalTensor& cos, + LocalTensor& sin); // 第二重循环体，给定BS范围，沿Q和K的N轴进行遍历处理 + __aicore__ inline void ProcessInSBNLoop( + int64_t sUbStart, int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, int64_t nUbStart, int64_t nUbLength, + int64_t nTotalSize, LocalTensor& cos, LocalTensor& sin, GlobalTensor& in, + GlobalTensor& out); // 第三重循环体，给定BSN范围，计算其中数据的rope + // 拷入拷出函数 + __aicore__ inline void CopyInCosAndSin(int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength); + __aicore__ inline void CopyInQ( + GlobalTensor& source, int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength, int64_t nStart, + int64_t nLength, int64_t nTotalSize); + __aicore__ inline void CopyOutQ( + GlobalTensor& target, int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength, int64_t nStart, + int64_t nLength, int64_t nTotalSize); + + // 计算函数 + __aicore__ inline void Compute( + LocalTensor& cos, LocalTensor& sin, int64_t sLength, int64_t bLength, int64_t nLength); + +private: + TPipe* pipe_; + + // GlobalMemory + GlobalTensor qGm_; + GlobalTensor cosGm_; + GlobalTensor sinGm_; + GlobalTensor qOutGm_; + + // UB + TQue qInQueue_; + TQue cosInQueue_; + TQue sinInQueue_; + TQue qOutQueue_; + + // Split core info + int64_t blockIdx_ = 0; + int64_t bBlockStart_ = 0; + int64_t bBlockLength_ = 0; + int64_t sBlockStart_ = 0; + int64_t sBlockLength_ = 0; + + // TilingData + const RopeRegbaseTilingData* tilingData_; + int64_t ubFactorB_ = 0; + int64_t ubFactorS_ = 0; + int64_t ubFactorN_ = 0; + int64_t D_ = 0; + int64_t dAlign_ = 0; + + // 拷贝参数 + uint8_t dSplitCoef_ = 1; + uint8_t copyInQSplitCoef_ = 1; // 拷贝q时使用的splitCoef + uint64_t ubCopyInStride = 0; // 输入在ub中的stride，deepseek_interleave中不为0 +}; + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBA::Init( + GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut, GM_ADDR workspace, const RopeRegbaseTilingData* tilingData, + TPipe* pipe) +{ + this->tilingData_ = tilingData; + this->blockIdx_ = GetBlockIdx(); + this->pipe_ = pipe; + this->InitAllGlobalBuffer(q, cos, sin, qOut); + this->InitAllBuffer(); + this->InitLoopParams(); +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBA::InitAllGlobalBuffer( + GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut) +{ + this->qGm_.SetGlobalBuffer((__gm__ T*)q); + this->cosGm_.SetGlobalBuffer((__gm__ T*)cos); + this->sinGm_.SetGlobalBuffer((__gm__ T*)sin); + this->qOutGm_.SetGlobalBuffer((__gm__ T*)qOut); +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBA::InitAllBuffer() +{ + this->ubFactorB_ = this->tilingData_->ubFactorB; + this->ubFactorS_ = this->tilingData_->ubFactorS; + this->ubFactorN_ = this->tilingData_->ubFactorN; + this->D_ = this->tilingData_->D; + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::HALF) || + tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + this->dSplitCoef_ = HALF_INTERLEAVE_COEF; + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::QUARTER)) { + this->dSplitCoef_ = QUARTER_MODE_COEF; + } + this->copyInQSplitCoef_ = dSplitCoef_; + this->dAlign_ = ops::CeilAlign(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(T)) * dSplitCoef_; + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + this->copyInQSplitCoef_ = 1; + this->ubCopyInStride = + (this->dAlign_ * sizeof(T) - ops::CeilAlign(tilingData_->sliceLength * sizeof(T), BLOCK_TYPE_SIZE)) / BLOCK_TYPE_SIZE; + } + this->pipe_->InitBuffer(this->qInQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(T)); + this->pipe_->InitBuffer( + this->qOutQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(T)); + if constexpr (IsBBoardcast) { + this->pipe_->InitBuffer(this->cosInQueue_, DOUBLE_BUFFER, ubFactorS_ * dAlign_ * sizeof(T)); + this->pipe_->InitBuffer(this->sinInQueue_, DOUBLE_BUFFER, ubFactorS_ * dAlign_ * sizeof(T)); + } else { + this->pipe_->InitBuffer(this->cosInQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * dAlign_ * sizeof(T)); + this->pipe_->InitBuffer(this->sinInQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * dAlign_ * sizeof(T)); + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBA::InitLoopParams() +{ + int64_t bIdx = blockIdx_ % tilingData_->blockNumB; + int64_t sIdx = blockIdx_ / tilingData_->blockNumB; + this->bBlockLength_ = tilingData_->blockFactorB; + this->sBlockLength_ = tilingData_->blockFactorS; + if (bIdx == tilingData_->blockNumB - 1 && tilingData_->B % tilingData_->blockFactorB != 0) { + this->bBlockLength_ = tilingData_->B % tilingData_->blockFactorB; + } + if (sIdx == tilingData_->blockNumS - 1 && tilingData_->S % tilingData_->blockFactorS != 0) { + this->sBlockLength_ = tilingData_->S % tilingData_->blockFactorS; + } + this->bBlockStart_ = bIdx * tilingData_->blockFactorB; + this->sBlockStart_ = sIdx * tilingData_->blockFactorS; +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBA::Process() +{ + // 在S轴进行循环 + int64_t ubLoopCount = ops::CeilDiv(sBlockLength_, ubFactorS_); + for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) { + this->ProcessInSLoop( + sBlockStart_ + ubLoopIdx * ubFactorS_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorS_ : sBlockLength_ - ubLoopIdx * ubFactorS_); + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBA::ProcessInSLoop( + int64_t sUbStart, int64_t sUbLength) +{ + // 在B轴进行循环 + int64_t ubLoopCount = ops::CeilDiv(bBlockLength_, ubFactorB_); + if constexpr (IsBBoardcast) { + // cos和sin需要在B轴广播的情况 + this->CopyInCosAndSin(sUbStart, sUbLength, 0, 1); + LocalTensor cosUb = this->cosInQueue_.template DeQue(); + LocalTensor sinUb = this->sinInQueue_.template DeQue(); + for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) { + this->ProcessInSBLoop( + sUbStart, sUbLength, bBlockStart_ + ubLoopIdx * ubFactorB_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_, cosUb, sinUb); + } + this->sinInQueue_.FreeTensor(cosUb); + this->cosInQueue_.FreeTensor(sinUb); + } else { + // sin和cos无需在B轴广播的情况 + for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) { + this->CopyInCosAndSin( + sUbStart, sUbLength, bBlockStart_ + ubLoopIdx * ubFactorB_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_); + LocalTensor cosUb = this->cosInQueue_.template DeQue(); + LocalTensor sinUb = this->sinInQueue_.template DeQue(); + this->ProcessInSBLoop( + sUbStart, sUbLength, bBlockStart_ + ubLoopIdx * ubFactorB_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_, cosUb, sinUb); + this->cosInQueue_.FreeTensor(cosUb); + this->sinInQueue_.FreeTensor(sinUb); + } + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBA::ProcessInSBLoop( + int64_t sUbStart, int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, LocalTensor& cos, LocalTensor& sin) +{ + // 循环处理Q + int64_t qUbLoopCount = ops::CeilDiv(tilingData_->N, ubFactorN_); + for (int64_t ubLoopIdx = 0; ubLoopIdx < qUbLoopCount; ubLoopIdx++) { + this->ProcessInSBNLoop( + sUbStart, sUbLength, bUbStart, bUbLength, ubLoopIdx * ubFactorN_, + ubLoopIdx != qUbLoopCount - 1 ? ubFactorN_ : tilingData_->N - ubLoopIdx * ubFactorN_, tilingData_->N, cos, + sin, qGm_, qOutGm_); + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBA::ProcessInSBNLoop( + int64_t sUbStart, int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, int64_t nUbStart, int64_t nUbLength, + int64_t nTotalSize, LocalTensor& cos, LocalTensor& sin, GlobalTensor& in, GlobalTensor& out) +{ + CopyInQ(in, sUbStart, sUbLength, bUbStart, bUbLength, nUbStart, nUbLength, nTotalSize); + Compute(cos, sin, sUbLength, bUbLength, nUbLength); + CopyOutQ(out, sUbStart, sUbLength, bUbStart, bUbLength, nUbStart, nUbLength, nTotalSize); +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBA::CopyInCosAndSin( + int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength) +{ + LocalTensor cosUb = this->cosInQueue_.template AllocTensor(); + LocalTensor sinUb = this->sinInQueue_.template AllocTensor(); + LoopModeParams loopParams; + loopParams.loop2Size = 1; + loopParams.loop1Size = bLength; + loopParams.loop2SrcStride = 0; + loopParams.loop2DstStride = 0; + loopParams.loop1SrcStride = tilingData_->S * tilingData_->sliceLength * sizeof(T); + loopParams.loop1DstStride = ubFactorS_ * dAlign_ * sizeof(T); + SetLoopModePara(loopParams, DataCopyMVType::OUT_TO_UB); + DataCopyPadExtParams copyPadExtparams; + copyPadExtparams.isPad = false; + copyPadExtparams.leftPadding = 0; + copyPadExtparams.rightPadding = 0; + copyPadExtparams.paddingValue = 0; + DataCopyExtParams copyExtParams; + copyExtParams.blockCount = sLength * dSplitCoef_; + copyExtParams.blockLen = tilingData_->sliceLength * sizeof(T) / dSplitCoef_; + copyExtParams.srcStride = 0; + copyExtParams.dstStride = 0; + DataCopyPad(cosUb, this->cosGm_[bStart * tilingData_->S * tilingData_->sliceLength + sStart * tilingData_->sliceLength], copyExtParams, copyPadExtparams); + DataCopyPad(sinUb, this->sinGm_[bStart * tilingData_->S * tilingData_->sliceLength + sStart * tilingData_->sliceLength], copyExtParams, copyPadExtparams); + ResetLoopModePara(DataCopyMVType::OUT_TO_UB); + this->cosInQueue_.template EnQue(cosUb); + this->sinInQueue_.template EnQue(sinUb); +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBA::CopyInQ( + GlobalTensor& source, int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength, int64_t nStart, + int64_t nLength, int64_t nTotalSize) +{ + LocalTensor target = this->qInQueue_.template AllocTensor(); + // 数据格式为BNSD，B->N->S->D + LoopModeParams loopParams; + loopParams.loop2Size = bLength; + loopParams.loop1Size = nLength; + loopParams.loop2SrcStride = nTotalSize * tilingData_->S * D_ * sizeof(T); + loopParams.loop2DstStride = ubFactorN_ * ubFactorS_ * dAlign_ * sizeof(T); + loopParams.loop1SrcStride = tilingData_->S * D_ * sizeof(T); + loopParams.loop1DstStride = ubFactorS_ * dAlign_ * sizeof(T); + SetLoopModePara(loopParams, DataCopyMVType::OUT_TO_UB); + DataCopyExtParams copyExtParams; + copyExtParams.blockCount = sLength * copyInQSplitCoef_; + copyExtParams.blockLen = tilingData_->sliceLength * sizeof(T) / copyInQSplitCoef_; + copyExtParams.srcStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(T); + copyExtParams.dstStride = ubCopyInStride; + DataCopyPadExtParams copyPadExtparams; + copyPadExtparams.isPad = false; + copyPadExtparams.leftPadding = 0; + copyPadExtparams.rightPadding = 0; + copyPadExtparams.paddingValue = 0; + DataCopyPad( + target, source[bStart * nTotalSize * tilingData_->S * D_ + nStart * tilingData_->S * D_ + sStart * D_ + tilingData_->sliceStart], + copyExtParams, copyPadExtparams); + ResetLoopModePara(DataCopyMVType::OUT_TO_UB); + this->qInQueue_.template EnQue(target); +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBA::CopyOutQ( + GlobalTensor& target, int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength, int64_t nStart, + int64_t nLength, int64_t nTotalSize) +{ + LocalTensor source = this->qOutQueue_.template DeQue(); + // 数据格式为BNSD，B->N->S->D + LoopModeParams loopParams; + loopParams.loop2Size = bLength; + loopParams.loop1Size = nLength; + loopParams.loop2DstStride = nTotalSize * tilingData_->S * D_ * sizeof(T); + loopParams.loop2SrcStride = ubFactorN_ * ubFactorS_ * dAlign_ * sizeof(T); + loopParams.loop1DstStride = tilingData_->S * D_ * sizeof(T); + loopParams.loop1SrcStride = ubFactorS_ * dAlign_ * sizeof(T); + SetLoopModePara(loopParams, DataCopyMVType::UB_TO_OUT); + DataCopyExtParams copyExtParams; + copyExtParams.blockCount = sLength * dSplitCoef_; + copyExtParams.blockLen = tilingData_->sliceLength * sizeof(T) / dSplitCoef_; + copyExtParams.srcStride = 0; + copyExtParams.dstStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(T); + DataCopyPad( + target[bStart * nTotalSize * tilingData_->S * D_ + nStart * tilingData_->S * D_ + sStart * D_ + tilingData_->sliceStart], source, + copyExtParams); + ResetLoopModePara(DataCopyMVType::UB_TO_OUT); + this->qOutQueue_.FreeTensor(source); +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBA::Compute( + LocalTensor& cos, LocalTensor& sin, int64_t sLength, int64_t bLength, int64_t nLength) +{ + LocalTensor inUb = this->qInQueue_.template DeQue(); + LocalTensor outUb = this->qOutQueue_.template AllocTensor(); + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::HALF)) { + BatchHalfAlignVF( + (__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(), (__local_mem__ T*)sin.GetPhyAddr(), + (__local_mem__ T*)outUb.GetPhyAddr(), sLength, bLength, nLength, tilingData_->sliceLength, dAlign_, ubFactorS_, ubFactorN_); + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::INTERLEAVE)) { + BatchInterleaveModeVF( + (__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(), (__local_mem__ T*)sin.GetPhyAddr(), + (__local_mem__ T*)outUb.GetPhyAddr(), sLength, bLength, nLength, tilingData_->sliceLength, dAlign_, ubFactorS_, ubFactorN_); + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::QUARTER)) { + BatchQuarterAlignVF( + (__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(), (__local_mem__ T*)sin.GetPhyAddr(), + (__local_mem__ T*)outUb.GetPhyAddr(), sLength, bLength, nLength, tilingData_->sliceLength, dAlign_, ubFactorS_, ubFactorN_); + } else { + BatchDeepSeekInterleaveModeVF( + (__local_mem__ T*)inUb.GetPhyAddr(), (__local_mem__ T*)cos.GetPhyAddr(), (__local_mem__ T*)sin.GetPhyAddr(), + (__local_mem__ T*)outUb.GetPhyAddr(), sLength, bLength, nLength, tilingData_->sliceLength, dAlign_, ubFactorS_, ubFactorN_); + } + this->qInQueue_.FreeTensor(inUb); + this->qOutQueue_.template EnQue(outUb); +} +} // namespace InplacePartialRotaryMul + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_aba_and_ba_mixed.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_aba_and_ba_mixed.h new file mode 100644 index 000000000..33d0c6b5c --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_aba_and_ba_mixed.h @@ -0,0 +1,364 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotary_position_embedding_reg_aba_and_ba_mixed.h + * \brief Mixed precision kernel for ABA/BA layout: x is half/bfloat16, cos/sin are float + */ + +#ifndef ROTARY_POSITION_EMBEDDING_REG_ABA_AND_BA_MIXED_H +#define ROTARY_POSITION_EMBEDDING_REG_ABA_AND_BA_MIXED_H + +#include "apply_rotary_pos_emb_common.h" + +namespace InplacePartialRotaryMul { +using namespace AscendC; + +template +class RotaryPositionEmbeddingABAAndBAMixed { +public: + __aicore__ inline RotaryPositionEmbeddingABAAndBAMixed(){}; + __aicore__ inline ~RotaryPositionEmbeddingABAAndBAMixed(){}; + + __aicore__ inline void Init(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut, GM_ADDR workspace, + const RopeRegbaseTilingData *tilingData, TPipe *pipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void InitAllGlobalBuffer(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut); + __aicore__ inline void InitAllBuffer(); + __aicore__ inline void InitLoopParams(); + __aicore__ inline void ProcessInSLoop(int64_t sUbStart, int64_t sUbLength); + __aicore__ inline void ProcessInSBLoop(int64_t sUbStart, int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, + LocalTensor &cos, LocalTensor &sin); + __aicore__ inline void ProcessInSBNLoop(int64_t sUbStart, int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, + int64_t nUbStart, int64_t nUbLength, int64_t nTotalSize, LocalTensor &cosFloat, + LocalTensor &sinFloat, GlobalTensor &in, GlobalTensor &out); + __aicore__ inline void CopyInCosAndSin(int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength); + __aicore__ inline void CopyInQ(GlobalTensor &source, int64_t sStart, int64_t sLength, int64_t bStart, + int64_t bLength, int64_t nStart, int64_t nLength, int64_t nTotalSize); + __aicore__ inline void CopyOutQ(GlobalTensor &target, int64_t sStart, int64_t sLength, int64_t bStart, + int64_t bLength, int64_t nStart, int64_t nLength, int64_t nTotalSize); + __aicore__ inline void Compute( + LocalTensor &cos, LocalTensor &sin, int64_t sLength, int64_t bLength, int64_t nLength); + +private: + TPipe *pipe_; + + GlobalTensor qGm_; + GlobalTensor cosGm_; + GlobalTensor sinGm_; + GlobalTensor qOutGm_; + + TQue qInQueue_; + TQue cosInQueue_; + TQue sinInQueue_; + TQue qOutQueue_; + + int64_t blockIdx_ = 0; + int64_t bBlockStart_ = 0; + int64_t bBlockLength_ = 0; + int64_t sBlockStart_ = 0; + int64_t sBlockLength_ = 0; + + const RopeRegbaseTilingData *tilingData_; + int64_t ubFactorB_ = 0; + int64_t ubFactorS_ = 0; + int64_t ubFactorN_ = 0; + int64_t D_ = 0; + int64_t dAlign_ = 0; + int64_t dAlignFloat_ = 0; + uint8_t dSplitCoef_ = 1; + uint8_t copyInQSplitCoef_ = 1; + uint64_t ubCopyInStride = 0; +}; + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed::Init(GM_ADDR q, GM_ADDR cos, GM_ADDR sin, + GM_ADDR qOut, GM_ADDR workspace, const RopeRegbaseTilingData *tilingData, TPipe *pipe) +{ + this->tilingData_ = tilingData; + this->blockIdx_ = GetBlockIdx(); + this->pipe_ = pipe; + this->InitAllGlobalBuffer(q, cos, sin, qOut); + this->InitAllBuffer(); + this->InitLoopParams(); +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed::InitAllGlobalBuffer( + GM_ADDR q, GM_ADDR cos, GM_ADDR sin, GM_ADDR qOut) +{ + this->qGm_.SetGlobalBuffer((__gm__ TX *)q); + this->cosGm_.SetGlobalBuffer((__gm__ float *)cos); + this->sinGm_.SetGlobalBuffer((__gm__ float *)sin); + this->qOutGm_.SetGlobalBuffer((__gm__ TX *)qOut); +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed::InitAllBuffer() +{ + this->ubFactorB_ = this->tilingData_->ubFactorB; + this->ubFactorS_ = this->tilingData_->ubFactorS; + this->ubFactorN_ = this->tilingData_->ubFactorN; + this->D_ = this->tilingData_->D; + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::HALF) || + tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + this->dSplitCoef_ = HALF_INTERLEAVE_COEF; + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::QUARTER)) { + this->dSplitCoef_ = QUARTER_MODE_COEF; + } + this->copyInQSplitCoef_ = dSplitCoef_; + this->dAlign_ = + ops::CeilAlign(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(TX)) * dSplitCoef_; + this->dAlignFloat_ = + ops::CeilAlign(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(float)) * dSplitCoef_; + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + this->copyInQSplitCoef_ = 1; + this->ubCopyInStride = (this->dAlign_ * sizeof(TX) - + ops::CeilAlign(tilingData_->sliceLength * sizeof(TX), BLOCK_TYPE_SIZE)) / + BLOCK_TYPE_SIZE; + } + this->pipe_->InitBuffer( + this->qInQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(TX)); + this->pipe_->InitBuffer( + this->qOutQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(TX)); + if constexpr (IsBBoardcast) { + this->pipe_->InitBuffer(this->cosInQueue_, DOUBLE_BUFFER, ubFactorS_ * dAlignFloat_ * sizeof(float)); + this->pipe_->InitBuffer(this->sinInQueue_, DOUBLE_BUFFER, ubFactorS_ * dAlignFloat_ * sizeof(float)); + } else { + this->pipe_->InitBuffer( + this->cosInQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * dAlignFloat_ * sizeof(float)); + this->pipe_->InitBuffer( + this->sinInQueue_, DOUBLE_BUFFER, ubFactorB_ * ubFactorS_ * dAlignFloat_ * sizeof(float)); + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed::InitLoopParams() +{ + int64_t bIdx = blockIdx_ % tilingData_->blockNumB; + int64_t sIdx = blockIdx_ / tilingData_->blockNumB; + this->bBlockLength_ = tilingData_->blockFactorB; + this->sBlockLength_ = tilingData_->blockFactorS; + if (bIdx == tilingData_->blockNumB - 1 && tilingData_->B % tilingData_->blockFactorB != 0) { + this->bBlockLength_ = tilingData_->B % tilingData_->blockFactorB; + } + if (sIdx == tilingData_->blockNumS - 1 && tilingData_->S % tilingData_->blockFactorS != 0) { + this->sBlockLength_ = tilingData_->S % tilingData_->blockFactorS; + } + this->bBlockStart_ = bIdx * tilingData_->blockFactorB; + this->sBlockStart_ = sIdx * tilingData_->blockFactorS; +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed::Process() +{ + int64_t ubLoopCount = ops::CeilDiv(sBlockLength_, ubFactorS_); + for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) { + this->ProcessInSLoop(sBlockStart_ + ubLoopIdx * ubFactorS_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorS_ : sBlockLength_ - ubLoopIdx * ubFactorS_); + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed::ProcessInSLoop( + int64_t sUbStart, int64_t sUbLength) +{ + int64_t ubLoopCount = ops::CeilDiv(bBlockLength_, ubFactorB_); + if constexpr (IsBBoardcast) { + this->CopyInCosAndSin(sUbStart, sUbLength, 0, 1); + LocalTensor cosUbFloat = this->cosInQueue_.template DeQue(); + LocalTensor sinUbFloat = this->sinInQueue_.template DeQue(); + LocalTensor cosUb = cosUbFloat.template ReinterpretCast(); + LocalTensor sinUb = sinUbFloat.template ReinterpretCast(); + for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) { + this->ProcessInSBLoop(sUbStart, + sUbLength, + bBlockStart_ + ubLoopIdx * ubFactorB_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_, + cosUb, + sinUb); + } + this->sinInQueue_.FreeTensor(sinUbFloat); + this->cosInQueue_.FreeTensor(cosUbFloat); + } else { + for (int64_t ubLoopIdx = 0; ubLoopIdx < ubLoopCount; ubLoopIdx++) { + this->CopyInCosAndSin(sUbStart, + sUbLength, + bBlockStart_ + ubLoopIdx * ubFactorB_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_); + LocalTensor cosUbFloat = this->cosInQueue_.template DeQue(); + LocalTensor sinUbFloat = this->sinInQueue_.template DeQue(); + LocalTensor cosUb = cosUbFloat.template ReinterpretCast(); + LocalTensor sinUb = sinUbFloat.template ReinterpretCast(); + this->ProcessInSBLoop(sUbStart, + sUbLength, + bBlockStart_ + ubLoopIdx * ubFactorB_, + ubLoopIdx != ubLoopCount - 1 ? ubFactorB_ : bBlockLength_ - ubLoopIdx * ubFactorB_, + cosUb, + sinUb); + this->cosInQueue_.FreeTensor(cosUbFloat); + this->sinInQueue_.FreeTensor(sinUbFloat); + } + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed::ProcessInSBLoop(int64_t sUbStart, + int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, LocalTensor &cos, LocalTensor &sin) +{ + int64_t qUbLoopCount = ops::CeilDiv(tilingData_->N, ubFactorN_); + LocalTensor cosFloat = cos.template ReinterpretCast(); + LocalTensor sinFloat = sin.template ReinterpretCast(); + for (int64_t ubLoopIdx = 0; ubLoopIdx < qUbLoopCount; ubLoopIdx++) { + this->ProcessInSBNLoop(sUbStart, + sUbLength, + bUbStart, + bUbLength, + ubLoopIdx * ubFactorN_, + ubLoopIdx != qUbLoopCount - 1 ? ubFactorN_ : tilingData_->N - ubLoopIdx * ubFactorN_, + tilingData_->N, + cosFloat, + sinFloat, + qGm_, + qOutGm_); + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed::ProcessInSBNLoop(int64_t sUbStart, + int64_t sUbLength, int64_t bUbStart, int64_t bUbLength, int64_t nUbStart, int64_t nUbLength, int64_t nTotalSize, + LocalTensor &cosFloat, LocalTensor &sinFloat, GlobalTensor &in, GlobalTensor &out) +{ + CopyInQ(in, sUbStart, sUbLength, bUbStart, bUbLength, nUbStart, nUbLength, nTotalSize); + Compute(cosFloat, sinFloat, sUbLength, bUbLength, nUbLength); + CopyOutQ(out, sUbStart, sUbLength, bUbStart, bUbLength, nUbStart, nUbLength, nTotalSize); +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed::CopyInCosAndSin( + int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength) +{ + LocalTensor cosUb = this->cosInQueue_.template AllocTensor(); + LocalTensor sinUb = this->sinInQueue_.template AllocTensor(); + LoopModeParams loopParams; + loopParams.loop2Size = 1; + loopParams.loop1Size = bLength; + loopParams.loop2SrcStride = 0; + loopParams.loop2DstStride = 0; + loopParams.loop1SrcStride = tilingData_->S * tilingData_->sliceLength * sizeof(float); + loopParams.loop1DstStride = ubFactorS_ * dAlignFloat_ * sizeof(float); + SetLoopModePara(loopParams, DataCopyMVType::OUT_TO_UB); + DataCopyPadExtParams copyPadExtparams; + copyPadExtparams.isPad = false; + copyPadExtparams.leftPadding = 0; + copyPadExtparams.rightPadding = 0; + copyPadExtparams.paddingValue = 0; + DataCopyExtParams copyExtParams; + copyExtParams.blockCount = sLength * dSplitCoef_; + copyExtParams.blockLen = tilingData_->sliceLength * sizeof(float) / dSplitCoef_; + copyExtParams.srcStride = 0; + copyExtParams.dstStride = 0; + DataCopyPad(cosUb, + this->cosGm_[bStart * tilingData_->S * tilingData_->sliceLength + sStart * tilingData_->sliceLength], + copyExtParams, + copyPadExtparams); + DataCopyPad(sinUb, + this->sinGm_[bStart * tilingData_->S * tilingData_->sliceLength + sStart * tilingData_->sliceLength], + copyExtParams, + copyPadExtparams); + ResetLoopModePara(DataCopyMVType::OUT_TO_UB); + this->cosInQueue_.template EnQue(cosUb); + this->sinInQueue_.template EnQue(sinUb); +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed::CopyInQ(GlobalTensor &source, + int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength, int64_t nStart, int64_t nLength, + int64_t nTotalSize) +{ + LocalTensor target = this->qInQueue_.template AllocTensor(); + LoopModeParams loopParams; + loopParams.loop2Size = bLength; + loopParams.loop1Size = nLength; + loopParams.loop2SrcStride = nTotalSize * tilingData_->S * D_ * sizeof(TX); + loopParams.loop2DstStride = ubFactorN_ * ubFactorS_ * dAlign_ * sizeof(TX); + loopParams.loop1SrcStride = tilingData_->S * D_ * sizeof(TX); + loopParams.loop1DstStride = ubFactorS_ * dAlign_ * sizeof(TX); + SetLoopModePara(loopParams, DataCopyMVType::OUT_TO_UB); + DataCopyExtParams copyExtParams; + copyExtParams.blockCount = sLength * copyInQSplitCoef_; + copyExtParams.blockLen = tilingData_->sliceLength * sizeof(TX) / copyInQSplitCoef_; + copyExtParams.srcStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(TX); + copyExtParams.dstStride = ubCopyInStride; + DataCopyPadExtParams copyPadExtparams; + copyPadExtparams.isPad = false; + int64_t offset = bStart * nTotalSize * tilingData_->S * D_ + nStart * tilingData_->S * D_ + sStart * D_ + + tilingData_->sliceStart; + DataCopyPad(target, source[offset], copyExtParams, copyPadExtparams); + ResetLoopModePara(DataCopyMVType::OUT_TO_UB); + this->qInQueue_.template EnQue(target); +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed::CopyOutQ(GlobalTensor &target, + int64_t sStart, int64_t sLength, int64_t bStart, int64_t bLength, int64_t nStart, int64_t nLength, + int64_t nTotalSize) +{ + LocalTensor source = this->qOutQueue_.template DeQue(); + LoopModeParams loopParams; + loopParams.loop2Size = bLength; + loopParams.loop1Size = nLength; + loopParams.loop2SrcStride = ubFactorN_ * ubFactorS_ * dAlign_ * sizeof(TX); + loopParams.loop2DstStride = nTotalSize * tilingData_->S * D_ * sizeof(TX); + loopParams.loop1SrcStride = ubFactorS_ * dAlign_ * sizeof(TX); + loopParams.loop1DstStride = tilingData_->S * D_ * sizeof(TX); + SetLoopModePara(loopParams, DataCopyMVType::UB_TO_OUT); + DataCopyExtParams copyExtParams; + copyExtParams.blockCount = sLength * dSplitCoef_; + copyExtParams.blockLen = tilingData_->sliceLength * sizeof(TX) / dSplitCoef_; + copyExtParams.srcStride = 0; + copyExtParams.dstStride = (tilingData_->D - tilingData_->sliceLength) * sizeof(TX); + int64_t offset = bStart * nTotalSize * tilingData_->S * D_ + nStart * tilingData_->S * D_ + sStart * D_ + + tilingData_->sliceStart; + DataCopyPad(target[offset], source, copyExtParams); + ResetLoopModePara(DataCopyMVType::UB_TO_OUT); + this->qOutQueue_.FreeTensor(source); +} + +template +__aicore__ inline void RotaryPositionEmbeddingABAAndBAMixed::Compute( + LocalTensor &cos, LocalTensor &sin, int64_t sLength, int64_t bLength, int64_t nLength) +{ + LocalTensor inUb = this->qInQueue_.template DeQue(); + LocalTensor outUb = this->qOutQueue_.template AllocTensor(); + int64_t totalLength = sLength * bLength * nLength * tilingData_->sliceLength; + + BatchInterleaveModeVFMixed((__local_mem__ TX *)inUb.GetPhyAddr(), + (__local_mem__ float *)cos.GetPhyAddr(), + (__local_mem__ float *)sin.GetPhyAddr(), + (__local_mem__ TX *)outUb.GetPhyAddr(), + sLength, + bLength, + nLength, + tilingData_->sliceLength, + dAlign_, + dAlignFloat_, + ubFactorS_, + ubFactorN_); + + this->qInQueue_.FreeTensor(inUb); + this->qOutQueue_.template EnQue(outUb); +} + +} // namespace InplacePartialRotaryMul + +#endif // ROTARY_POSITION_EMBEDDING_REG_ABA_AND_BA_MIXED_H \ No newline at end of file diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_bab.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_bab.h new file mode 100644 index 000000000..4f1b71a17 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_bab.h @@ -0,0 +1,192 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotary_position_embedding_reg_bab.h + * \brief + */ + +#ifndef ROTARY_POSITION_EMBEDDING_REG_BAB_H +#define ROTARY_POSITION_EMBEDDING_REG_BAB_H + +// #include "op_kernel/math_util.h" +#include "apply_rotary_pos_emb_common.h" + +namespace InplacePartialRotaryMul { +using namespace AscendC; + +template +class RotaryPositionEmbeddingBAB +{ +public: + __aicore__ inline RotaryPositionEmbeddingBAB(TPipe* pipe, const RopeRegbaseTilingData* tiling) + : pipe_(pipe), tilingData_(tiling){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y); + __aicore__ inline void Process(); + +private: + constexpr static int32_t bufferNum = 2; + const RopeRegbaseTilingData* tilingData_; + TPipe* pipe_; + int64_t blockIdx_ = 0; + int64_t dSplitCoef_ = 1; // 切分系数初始化为1 + uint32_t dSplitSize_ = 0; + int64_t dAlign_ = 0; + int64_t bIdx_ = 0; + int64_t sIdx_ = 0; + int64_t bNum_ = 0; + int64_t sNum_ = 0; + int64_t ubFactorS_ = 0; + int64_t ubFactorN_ = 0; + GlobalTensor xGm_; + GlobalTensor cosGm_; + GlobalTensor sinGm_; + GlobalTensor yOutGm_; + + TQue xInQue_; + TQue cosInQue_; + TQue sinInQue_; + TQue yOutQue_; + +private: + __aicore__ inline void PrePareParams(); + __aicore__ inline void ProcessNLoop(const uint32_t bIdx, const uint32_t sIdx, const uint32_t currSNum); + __aicore__ inline void Compute( + const LocalTensor& sinTensor, const LocalTensor& cosTensor, const LocalTensor& inTensor, + const LocalTensor& outTensor, const uint32_t currSNum, const uint32_t currDNum); + __aicore__ inline void ProcessN( + const LocalTensor& sinTensor, const LocalTensor& cosTensor, const uint32_t bIdx, const uint32_t sIdx, + const uint32_t currSNum); +}; + +template +__aicore__ inline void RotaryPositionEmbeddingBAB::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y) +{ + this->blockIdx_ = GetBlockIdx(); + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::HALF) || + tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + this->dSplitCoef_ = HALF_INTERLEAVE_COEF; + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::QUARTER)) { + this->dSplitCoef_ = QUARTER_MODE_COEF; + } + this->dSplitSize_ = tilingData_->sliceLength / dSplitCoef_ * sizeof(T); + this->dAlign_ = ops::CeilAlign(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(T)) * dSplitCoef_; + ubFactorN_ = tilingData_->ubFactorN; + ubFactorS_ = tilingData_->ubFactorS; + this->xGm_.SetGlobalBuffer((__gm__ T*)x); + this->cosGm_.SetGlobalBuffer((__gm__ T*)cos); + this->sinGm_.SetGlobalBuffer((__gm__ T*)sin); + this->yOutGm_.SetGlobalBuffer((__gm__ T*)y); + this->pipe_->InitBuffer(xInQue_, bufferNum, ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(T)); + this->pipe_->InitBuffer(cosInQue_, bufferNum, ubFactorS_ * dAlign_ * sizeof(T)); + this->pipe_->InitBuffer(sinInQue_, bufferNum, ubFactorS_ * dAlign_ * sizeof(T)); + this->pipe_->InitBuffer(yOutQue_, bufferNum, ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(T)); +} + +template +__aicore__ inline void RotaryPositionEmbeddingBAB::PrePareParams() +{ + bIdx_ = blockIdx_ % tilingData_->blockNumB; + sIdx_ = blockIdx_ / tilingData_->blockNumB; + bNum_ = tilingData_->blockFactorB; + sNum_ = tilingData_->blockFactorS; + if (bIdx_ == tilingData_->blockNumB - 1 && tilingData_->B % tilingData_->blockFactorB != 0) { + bNum_ = tilingData_->B % tilingData_->blockFactorB; + } + if (sIdx_ == tilingData_->blockNumS - 1 && tilingData_->S % tilingData_->blockFactorS != 0) { + sNum_ = tilingData_->S % tilingData_->blockFactorS; + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingBAB::Process() +{ + PrePareParams(); + uint32_t bIdxStart = bIdx_ * tilingData_->blockFactorB; + for (uint32_t bIdx = bIdxStart; bIdx < bIdxStart + bNum_; bIdx++) { + uint32_t sIdxStart = sIdx_ * tilingData_->blockFactorS; + uint32_t sLoopCnt = ops::CeilDiv(sNum_, ubFactorS_); + for (uint32_t loopIdx = 0; loopIdx < sLoopCnt; loopIdx++) { + uint32_t currSNum = (loopIdx != sLoopCnt - 1) ? ubFactorS_ : sNum_ - loopIdx * ubFactorS_; + ProcessNLoop(bIdx, sIdxStart + loopIdx * ubFactorS_, currSNum); + } + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingBAB::ProcessNLoop( + const uint32_t bIdx, const uint32_t sIdx, const uint32_t currSNum) +{ + LocalTensor sinTensor = sinInQue_.AllocTensor(); + LocalTensor cosTensor = cosInQue_.AllocTensor(); + int64_t offset = sIdx * tilingData_->sliceLength; + DataCopyExtParams copyParams{static_cast(currSNum * dSplitCoef_), dSplitSize_, 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(sinTensor, sinGm_[offset], copyParams, padParams); + DataCopyPad(cosTensor, cosGm_[offset], copyParams, padParams); + sinInQue_.EnQue(sinTensor); + cosInQue_.EnQue(cosTensor); + sinTensor = sinInQue_.DeQue(); + cosTensor = cosInQue_.DeQue(); + ProcessN(sinTensor, cosTensor, bIdx, sIdx, currSNum); + sinInQue_.FreeTensor(sinTensor); + cosInQue_.FreeTensor(cosTensor); +} + +template +__aicore__ inline void RotaryPositionEmbeddingBAB::ProcessN( + const LocalTensor& sinTensor, const LocalTensor& cosTensor, const uint32_t bIdx, const uint32_t sIdx, + const uint32_t currSNum) +{ + LocalTensor xTensor; + LocalTensor yTensor; + int64_t baseOffset = (bIdx * tilingData_->S + sIdx) * tilingData_->N * tilingData_->D + tilingData_->sliceStart; + for (uint32_t idxN = 0; idxN < tilingData_->ubLoopNumN; idxN++) { + int64_t currDNum = (idxN == tilingData_->ubLoopNumN - 1) ? tilingData_->ubTailFactorN : ubFactorN_; + int64_t offset = baseOffset + idxN * ubFactorN_ * tilingData_->D; + xTensor = xInQue_.AllocTensor(); + DataCopyExtParams copyInParams{static_cast(currSNum * currDNum * dSplitCoef_), dSplitSize_, static_cast((tilingData_->D - tilingData_->sliceLength) * sizeof(T)), 0, 0}; + DataCopyExtParams copyOutParams{static_cast(currSNum * currDNum * dSplitCoef_), dSplitSize_, 0, static_cast((tilingData_->D - tilingData_->sliceLength) * sizeof(T)), 0}; + // deepSeekInterleave 场景，copyIn时按照D对齐，copyOut 时，按照 d/2 对齐 + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + copyInParams = {static_cast(currSNum * currDNum), tilingData_->sliceLength * sizeof(T), static_cast((tilingData_->D - tilingData_->sliceLength) * sizeof(T)), 0, 0}; + } + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(xTensor, xGm_[offset], copyInParams, padParams); + xInQue_.EnQue(xTensor); + xTensor = xInQue_.DeQue(); + yTensor = yOutQue_.AllocTensor(); + Compute(sinTensor, cosTensor, xTensor, yTensor, currSNum, currDNum); + xInQue_.FreeTensor(xTensor); + yOutQue_.EnQue(yTensor); + yTensor = yOutQue_.DeQue(); + DataCopyPad(yOutGm_[offset], yTensor, copyOutParams); + yOutQue_.FreeTensor(yTensor); + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingBAB::Compute( + const LocalTensor& sinTensor, const LocalTensor& cosTensor, const LocalTensor& inTensor, + const LocalTensor& outTensor, const uint32_t currSNum, const uint32_t currDNum) +{ + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::HALF)) { + HalfAlignVF(sinTensor, cosTensor, inTensor, outTensor, tilingData_->sliceLength, dAlign_, currSNum, currDNum); + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::INTERLEAVE)) { + InterleaveModeVF(sinTensor, cosTensor, inTensor, outTensor, tilingData_->sliceLength, currSNum, currDNum); + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::QUARTER)) { + QuarterAlignVF(sinTensor, cosTensor, inTensor, outTensor, tilingData_->sliceLength, dAlign_, currSNum, currDNum); + } else { + DeepSeekInterleaveModeVF(sinTensor, cosTensor, inTensor, outTensor, tilingData_->sliceLength, currSNum, currDNum); + } +} + +} // namespace InplacePartialRotaryMul +#endif // ROTARY_POSITION_EMBEDDING_REG_BAB_H diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_bab_mixed.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_bab_mixed.h new file mode 100644 index 000000000..4a61a5322 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotary_position_embedding_reg_bab_mixed.h @@ -0,0 +1,207 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotary_position_embedding_reg_bab_mixed.h + * \brief Mixed precision kernel: x is half/bfloat16, cos/sin are float + */ + +#ifndef ROTARY_POSITION_EMBEDDING_REG_BAB_MIXED_H +#define ROTARY_POSITION_EMBEDDING_REG_BAB_MIXED_H + +#include "apply_rotary_pos_emb_common.h" + +namespace InplacePartialRotaryMul { +using namespace AscendC; + +template +class RotaryPositionEmbeddingBABMixed { +public: + __aicore__ inline RotaryPositionEmbeddingBABMixed(TPipe *pipe, const RopeRegbaseTilingData *tiling) + : pipe_(pipe), tilingData_(tiling){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y); + __aicore__ inline void Process(); + +private: + constexpr static int32_t bufferNum = 2; + const RopeRegbaseTilingData *tilingData_; + TPipe *pipe_; + int64_t blockIdx_ = 0; + int64_t dSplitCoef_ = 1; + uint32_t dSplitSize_ = 0; + int64_t dAlign_ = 0; + int64_t dAlignFloat_ = 0; + int64_t bIdx_ = 0; + int64_t sIdx_ = 0; + int64_t bNum_ = 0; + int64_t sNum_ = 0; + int64_t ubFactorS_ = 0; + int64_t ubFactorN_ = 0; + GlobalTensor xGm_; + GlobalTensor cosGm_; + GlobalTensor sinGm_; + GlobalTensor yOutGm_; + + TQue xInQue_; + TQue cosInQue_; + TQue sinInQue_; + TQue yOutQue_; + +private: + __aicore__ inline void PrePareParams(); + __aicore__ inline void ProcessNLoop(const uint32_t bIdx, const uint32_t sIdx, const uint32_t currSNum); + __aicore__ inline void Compute(const LocalTensor &sinTensor, const LocalTensor &cosTensor, + const LocalTensor &inTensor, const LocalTensor &outTensor, const uint32_t currSNum, + const uint32_t currDNum); + __aicore__ inline void ProcessN(const LocalTensor &sinTensor, const LocalTensor &cosTensor, + const uint32_t bIdx, const uint32_t sIdx, const uint32_t currSNum); +}; + +template +__aicore__ inline void RotaryPositionEmbeddingBABMixed::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y) +{ + this->blockIdx_ = GetBlockIdx(); + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::HALF) || + tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + this->dSplitCoef_ = HALF_INTERLEAVE_COEF; + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::QUARTER)) { + this->dSplitCoef_ = QUARTER_MODE_COEF; + } + this->dSplitSize_ = tilingData_->sliceLength / dSplitCoef_ * sizeof(TX); + this->dAlign_ = + ops::CeilAlign(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(TX)) * dSplitCoef_; + this->dAlignFloat_ = + ops::CeilAlign(tilingData_->sliceLength / dSplitCoef_, BLOCK_TYPE_SIZE / sizeof(float)) * dSplitCoef_; + ubFactorN_ = tilingData_->ubFactorN; + ubFactorS_ = tilingData_->ubFactorS; + this->xGm_.SetGlobalBuffer((__gm__ TX *)x); + this->cosGm_.SetGlobalBuffer((__gm__ float *)cos); + this->sinGm_.SetGlobalBuffer((__gm__ float *)sin); + this->yOutGm_.SetGlobalBuffer((__gm__ TX *)y); + this->pipe_->InitBuffer(xInQue_, bufferNum, ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(TX)); + this->pipe_->InitBuffer(cosInQue_, bufferNum, ubFactorS_ * dAlignFloat_ * sizeof(float)); + this->pipe_->InitBuffer(sinInQue_, bufferNum, ubFactorS_ * dAlignFloat_ * sizeof(float)); + this->pipe_->InitBuffer(yOutQue_, bufferNum, ubFactorS_ * ubFactorN_ * dAlign_ * sizeof(TX)); +} + +template +__aicore__ inline void RotaryPositionEmbeddingBABMixed::PrePareParams() +{ + bIdx_ = blockIdx_ % tilingData_->blockNumB; + sIdx_ = blockIdx_ / tilingData_->blockNumB; + bNum_ = tilingData_->blockFactorB; + sNum_ = tilingData_->blockFactorS; + if (bIdx_ == tilingData_->blockNumB - 1 && tilingData_->B % tilingData_->blockFactorB != 0) { + bNum_ = tilingData_->B % tilingData_->blockFactorB; + } + if (sIdx_ == tilingData_->blockNumS - 1 && tilingData_->S % tilingData_->blockFactorS != 0) { + sNum_ = tilingData_->S % tilingData_->blockFactorS; + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingBABMixed::Process() +{ + PrePareParams(); + uint32_t bIdxStart = bIdx_ * tilingData_->blockFactorB; + for (uint32_t bIdx = bIdxStart; bIdx < bIdxStart + bNum_; bIdx++) { + uint32_t sIdxStart = sIdx_ * tilingData_->blockFactorS; + uint32_t sLoopCnt = ops::CeilDiv(sNum_, ubFactorS_); + for (uint32_t loopIdx = 0; loopIdx < sLoopCnt; loopIdx++) { + uint32_t currSNum = (loopIdx != sLoopCnt - 1) ? ubFactorS_ : sNum_ - loopIdx * ubFactorS_; + ProcessNLoop(bIdx, sIdxStart + loopIdx * ubFactorS_, currSNum); + } + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingBABMixed::ProcessNLoop( + const uint32_t bIdx, const uint32_t sIdx, const uint32_t currSNum) +{ + LocalTensor sinTensor = sinInQue_.AllocTensor(); + LocalTensor cosTensor = cosInQue_.AllocTensor(); + int64_t offset = sIdx * tilingData_->sliceLength; + uint32_t dSplitSizeFloat = tilingData_->sliceLength / dSplitCoef_ * sizeof(float); + DataCopyExtParams copyParams{static_cast(currSNum * dSplitCoef_), dSplitSizeFloat, 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(sinTensor, sinGm_[offset], copyParams, padParams); + DataCopyPad(cosTensor, cosGm_[offset], copyParams, padParams); + sinInQue_.EnQue(sinTensor); + cosInQue_.EnQue(cosTensor); + sinTensor = sinInQue_.DeQue(); + cosTensor = cosInQue_.DeQue(); + ProcessN(sinTensor, cosTensor, bIdx, sIdx, currSNum); + sinInQue_.FreeTensor(sinTensor); + cosInQue_.FreeTensor(cosTensor); +} + +template +__aicore__ inline void RotaryPositionEmbeddingBABMixed::ProcessN(const LocalTensor &sinTensor, + const LocalTensor &cosTensor, const uint32_t bIdx, const uint32_t sIdx, const uint32_t currSNum) +{ + LocalTensor xTensor; + LocalTensor yTensor; + int64_t baseOffset = (bIdx * tilingData_->S + sIdx) * tilingData_->N * tilingData_->D + tilingData_->sliceStart; + for (uint32_t idxN = 0; idxN < tilingData_->ubLoopNumN; idxN++) { + int64_t currDNum = (idxN == tilingData_->ubLoopNumN - 1) ? tilingData_->ubTailFactorN : ubFactorN_; + int64_t offset = baseOffset + idxN * ubFactorN_ * tilingData_->D; + xTensor = xInQue_.AllocTensor(); + DataCopyExtParams copyInParams{static_cast(currSNum * currDNum * dSplitCoef_), + dSplitSize_, + static_cast((tilingData_->D - tilingData_->sliceLength) * sizeof(TX)), + 0, + 0}; + DataCopyExtParams copyOutParams{static_cast(currSNum * currDNum * dSplitCoef_), + dSplitSize_, + 0, + static_cast((tilingData_->D - tilingData_->sliceLength) * sizeof(TX)), + 0}; + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::DEEPSEEK_INTERLEAVE)) { + copyInParams = {static_cast(currSNum * currDNum), + tilingData_->sliceLength * sizeof(TX), + static_cast((tilingData_->D - tilingData_->sliceLength) * sizeof(TX)), + 0, + 0}; + } + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(xTensor, xGm_[offset], copyInParams, padParams); + xInQue_.EnQue(xTensor); + xTensor = xInQue_.DeQue(); + yTensor = yOutQue_.AllocTensor(); + Compute(sinTensor, cosTensor, xTensor, yTensor, currSNum, currDNum); + xInQue_.FreeTensor(xTensor); + yOutQue_.EnQue(yTensor); + yTensor = yOutQue_.DeQue(); + DataCopyPad(yOutGm_[offset], yTensor, copyOutParams); + yOutQue_.FreeTensor(yTensor); + } +} + +template +__aicore__ inline void RotaryPositionEmbeddingBABMixed::Compute(const LocalTensor &sinTensor, + const LocalTensor &cosTensor, const LocalTensor &inTensor, const LocalTensor &outTensor, + const uint32_t currSNum, const uint32_t currDNum) +{ + if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::INTERLEAVE)) { + InterleaveModeVFMixed( + inTensor, cosTensor, sinTensor, outTensor, tilingData_->sliceLength, currSNum, currDNum); + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::HALF)) { + // For HALF mode, need to implement HalfAlignVFMixed + InterleaveModeVFMixed( + inTensor, cosTensor, sinTensor, outTensor, tilingData_->sliceLength, currSNum, currDNum); + } else if (tilingData_->rotaryMode == static_cast(RotaryPosEmbeddingMode::QUARTER)) { + // For QUARTER mode, need to implement QuarterAlignVFMixed + InterleaveModeVFMixed( + inTensor, cosTensor, sinTensor, outTensor, tilingData_->sliceLength, currSNum, currDNum); + } +} + +} // namespace InplacePartialRotaryMul +#endif // ROTARY_POSITION_EMBEDDING_REG_BAB_MIXED_H \ No newline at end of file diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_half.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_half.h new file mode 100644 index 000000000..b6402ffb0 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_half.h @@ -0,0 +1,350 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotate_half.h + * \brief + */ +#ifndef ROTATE_HALF_H +#define ROTATE_HALF_H + +#include "rotate_half_base.h" + +namespace RotateHalfN { +using namespace AscendC; + +template +class RotateHalf : public RotateHalfBase { +public: + __aicore__ inline RotateHalf(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RotaryPositionEmbeddingTilingData &tilingData); + __aicore__ inline void Process(); + +protected: + TPipe pipe; + TQue inQueueX; + TQue inQueueCos; + TQue inQueueSin; + TQue outQueueY; + GlobalTensor xGm; + GlobalTensor cosGm; + GlobalTensor sinGm; + GlobalTensor yGm; + + __aicore__ inline void NormalProcess(); + __aicore__ inline void RB1sdProcess(); + __aicore__ inline void BndProcess(); + __aicore__ inline void SingleStepProcess(uint32_t progress, uint32_t sLines, uint64_t copyLength, + uint64_t calcLength); + __aicore__ inline void RB1sdSingleStepProcess(uint32_t progress, uint32_t sLines, uint64_t xBatchStartOffset, + uint64_t rBatchStartOffset, uint64_t copyLength, uint64_t calcLength); + __aicore__ inline void Compute(LocalTensor &cos, LocalTensor &sin, uint32_t sLines, uint32_t calcLength); + __aicore__ inline void CopyInR(uint64_t rStartOffset, uint16_t sLines, uint32_t copyLength); + __aicore__ inline void CopyInX(uint64_t xStartOffset, uint16_t sLines, uint32_t copyLength); + __aicore__ inline void CopyOut(uint64_t yStartOffset, uint16_t sLines, uint32_t copyLength); +}; + +template +__aicore__ inline void RotateHalf::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RotaryPositionEmbeddingTilingData &tilingData) +{ + this->BaseMemberInit(tilingData); + + xGm.SetGlobalBuffer((__gm__ T *)x + this->xOffset + this->xCoreOffset * this->coreRelativeIdx, this->xAllocLength); + yGm.SetGlobalBuffer((__gm__ T *)y + this->xOffset + this->xCoreOffset * this->coreRelativeIdx, this->xAllocLength); + cosGm.SetGlobalBuffer((__gm__ T *)cos + this->rOffset + this->rCoreOffset * this->coreRelativeIdx, + this->rAllocLength); + sinGm.SetGlobalBuffer((__gm__ T *)sin + this->rOffset + this->rCoreOffset * this->coreRelativeIdx, + this->rAllocLength); + + pipe.InitBuffer(inQueueX, DOUBLE_BUFFER, this->storePadDataLength * sizeof(T)); + pipe.InitBuffer(outQueueY, DOUBLE_BUFFER, this->storePadDataLength * sizeof(T)); + pipe.InitBuffer(inQueueCos, DOUBLE_BUFFER, this->storePadDataLength * sizeof(T)); + pipe.InitBuffer(inQueueSin, DOUBLE_BUFFER, this->storePadDataLength * sizeof(T)); +} + +template +__aicore__ inline void RotateHalf::Process() +{ + if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_BSND || this->layout == LAYOUT_SBND || + this->layout == LAYOUT_NO_BROADCAST) { + NormalProcess(); + } else if (this->layout == LAYOUT_R_B1SD) { + RB1sdProcess(); + } else if (this->layout == LAYOUT_BND) { + BndProcess(); + } +} + +template +__aicore__ inline void RotateHalf::NormalProcess() +{ + for (uint32_t progress = 0; progress < this->ubLoop; progress++) { + SingleStepProcess(progress, this->storeSLines, this->storeDataLength, this->storePadDataLength); + } + if (this->ubLast > 0) { + SingleStepProcess(this->ubLoop, this->ubLast, this->ubLastDataLength, this->ubLastPadDataLength); + } +} + +template +__aicore__ inline void RotateHalf::RB1sdProcess() +{ + uint64_t totalSdLength = this->totalSLines * this->dLength; + uint64_t totalNsdLength = totalSdLength * this->bcSecondDim; + uint64_t xBatchOffset, rBatchOffset; + for (uint32_t batchLoop = 0; batchLoop < this->bcFirstDim; batchLoop++) { + xBatchOffset = batchLoop * totalNsdLength; + rBatchOffset = batchLoop * totalSdLength; + for (uint32_t progress = 0; progress < this->ubLoop; progress++) { + RB1sdSingleStepProcess(progress, this->storeSLines, xBatchOffset, rBatchOffset, this->storeDataLength, + this->storePadDataLength); + } + if (this->ubLast > 0) { + RB1sdSingleStepProcess(this->ubLoop, this->ubLast, xBatchOffset, rBatchOffset, this->ubLastDataLength, + this->ubLastPadDataLength); + } + } +} + +template +__aicore__ inline void RotateHalf::BndProcess() +{ + CopyInR(0, 1, this->dLength); + LocalTensor cosLocal = inQueueCos.DeQue(); + LocalTensor sinLocal = inQueueSin.DeQue(); + Muls(sinLocal, sinLocal, (T)(-1.0), this->halfDPadLength); + uint32_t broadcastLines = this->ubLoop > 0 ? this->storeSLines - 1 : this->ubLast - 1; + if (broadcastLines > 0) { + this->RBroadCast(cosLocal, sinLocal, broadcastLines); + } + uint64_t xOffset; + for (uint32_t progress = 0; progress < this->ubLoop; progress++) { + xOffset = progress * this->storeDataLength; + CopyInX(xOffset, this->storeSLines, this->storeDataLength); + Compute(cosLocal, sinLocal, this->storeSLines, this->storePadDataLength); + CopyOut(xOffset, this->storeSLines, this->storeDataLength); + } + if (this->ubLast > 0) { + xOffset = this->ubLoop * this->storeDataLength; + CopyInX(xOffset, this->ubLast, this->ubLastDataLength); + Compute(cosLocal, sinLocal, this->ubLast, this->ubLastPadDataLength); + CopyOut(xOffset, this->ubLast, this->ubLastDataLength); + } + inQueueCos.FreeTensor(cosLocal); + inQueueSin.FreeTensor(sinLocal); +} + +template +__aicore__ inline void RotateHalf::SingleStepProcess(uint32_t progress, uint32_t sLines, uint64_t copyLength, + uint64_t calcLength) +{ + uint64_t xOffset, rOffset, bnLoopXStartOffset, progressOffset, batchOffset; + rOffset = progress * this->storeDataLength; + CopyInR(rOffset, sLines, copyLength); + LocalTensor cosLocal = inQueueCos.DeQue(); + LocalTensor sinLocal = inQueueSin.DeQue(); + this->SinCompute(sinLocal, sLines); + + if (this->layout == LAYOUT_BNSD) { + uint64_t totalSdSize = this->totalSLines * this->dLength; + bnLoopXStartOffset = progress * this->storeDataLength; + for (uint32_t bnLoop = 0; bnLoop < this->bnSize; bnLoop++) { + xOffset = bnLoopXStartOffset + bnLoop * totalSdSize; + CopyInX(xOffset, sLines, copyLength); + Compute(cosLocal, sinLocal, sLines, calcLength); + CopyOut(xOffset, sLines, copyLength); + } + } else if (this->layout == LAYOUT_BSND) { + uint64_t totalSndSize = this->totalSLines * this->ndSize; + progressOffset = progress * this->bcSecondDim * this->storeDataLength; + for (uint32_t bLoop = 0; bLoop < this->bcFirstDim; bLoop++) { + batchOffset = bLoop * totalSndSize; + for (uint32_t nLoop = 0; nLoop < this->bcSecondDim; nLoop++) { + xOffset = nLoop * this->dLength + batchOffset + progressOffset; + CopyInX(xOffset, sLines, copyLength); + Compute(cosLocal, sinLocal, sLines, calcLength); + CopyOut(xOffset, sLines, copyLength); + } + } + } else if (this->layout == LAYOUT_SBND) { + bnLoopXStartOffset = progress * this->storeDataLength * this->bnSize; + for (uint32_t bnLoop = 0; bnLoop < this->bnSize; bnLoop++) { + xOffset = bnLoopXStartOffset + bnLoop * this->dLength; + CopyInX(xOffset, sLines, copyLength); + Compute(cosLocal, sinLocal, sLines, calcLength); + CopyOut(xOffset, sLines, copyLength); + } + } else if (this->layout == LAYOUT_NO_BROADCAST) { + CopyInX(rOffset, sLines, copyLength); + Compute(cosLocal, sinLocal, sLines, calcLength); + CopyOut(rOffset, sLines, copyLength); + } + inQueueCos.FreeTensor(cosLocal); + inQueueSin.FreeTensor(sinLocal); +} + +template +__aicore__ inline void RotateHalf::RB1sdSingleStepProcess(uint32_t progress, uint32_t sLines, + uint64_t xBatchStartOffset, uint64_t rBatchStartOffset, + uint64_t copyLength, uint64_t calcLength) +{ + CopyInR(progress * this->storeDataLength + rBatchStartOffset, sLines, copyLength); + LocalTensor cosLocal = inQueueCos.DeQue(); + LocalTensor sinLocal = inQueueSin.DeQue(); + this->SinCompute(sinLocal, sLines); + + uint64_t xOffset, progressXOffset; + progressXOffset = progress * this->storeDataLength + xBatchStartOffset; + for (uint32_t nLoop = 0; nLoop < this->bcSecondDim; nLoop++) { + xOffset = nLoop * this->totalSLines * this->dLength + progressXOffset; + CopyInX(xOffset, sLines, copyLength); + Compute(cosLocal, sinLocal, sLines, calcLength); + CopyOut(xOffset, sLines, copyLength); + } + inQueueCos.FreeTensor(cosLocal); + inQueueSin.FreeTensor(sinLocal); +} + +template +__aicore__ inline void RotateHalf::CopyInR(uint64_t rStartOffset, uint16_t sLines, uint32_t copyLength) +{ + LocalTensor cosLocal = inQueueCos.AllocTensor(); + LocalTensor sinLocal = inQueueSin.AllocTensor(); + if (this->isAligned == true) { + DataCopy(cosLocal, cosGm[rStartOffset], copyLength); + DataCopy(sinLocal, sinGm[rStartOffset], copyLength); + } else { + DataCopyExtParams copyParams{(uint16_t)(2 * sLines), // blockCount + this->halfDBytes, // blockLen + 0, // srcStride(bytes) + 0, // dstStride(block) + 0}; + DataCopyPad(cosLocal, cosGm[rStartOffset], copyParams, this->noPadParams); + DataCopyPad(sinLocal, sinGm[rStartOffset], copyParams, this->noPadParams); + } + inQueueCos.EnQue(cosLocal); + inQueueSin.EnQue(sinLocal); +} + +template +__aicore__ inline void RotateHalf::CopyInX(uint64_t xStartOffset, uint16_t sLines, uint32_t copyLength) +{ + LocalTensor xLocal = inQueueX.AllocTensor(); + DataCopyExtParams copyParams; + + if (this->isAligned == true) { + if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND || + this->layout == LAYOUT_R_B1SD) { + DataCopy(xLocal, xGm[xStartOffset], copyLength); + } else if (this->layout == LAYOUT_BSND) { + copyParams.blockCount = sLines; + copyParams.blockLen = this->dBytes; + copyParams.srcStride = (this->bcSecondDim - 1) * this->dBytes; + copyParams.dstStride = 0; + DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams); + } else if (this->layout == LAYOUT_SBND) { + copyParams.blockCount = sLines; + copyParams.blockLen = this->dBytes; + copyParams.srcStride = (this->bnSize - 1) * this->dBytes; + copyParams.dstStride = 0; + DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams); + } + } else { + if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND || + this->layout == LAYOUT_R_B1SD) { + copyParams.blockCount = (uint16_t)(2 * sLines); + copyParams.blockLen = this->halfDBytes; + copyParams.srcStride = 0; + copyParams.dstStride = 0; + DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams); + } else if (this->layout == LAYOUT_BSND) { + copyParams.blockCount = sLines; + copyParams.blockLen = this->halfDBytes; + copyParams.srcStride = (2 * this->bcSecondDim - 1) * this->halfDBytes; + copyParams.dstStride = this->halfDPadBlocks; + DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams); + DataCopyPad(xLocal[this->halfDPadLength], xGm[xStartOffset + this->halfDLength], copyParams, + this->noPadParams); + } else if (this->layout == LAYOUT_SBND) { + copyParams.blockCount = sLines; + copyParams.blockLen = this->halfDBytes; + copyParams.srcStride = (2 * this->bnSize - 1) * this->halfDBytes; + copyParams.dstStride = this->halfDPadBlocks; + DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams); + DataCopyPad(xLocal[this->halfDPadLength], xGm[xStartOffset + this->halfDLength], copyParams, + this->noPadParams); + } + } + inQueueX.EnQue(xLocal); +} + +template +__aicore__ inline void RotateHalf::CopyOut(uint64_t yStartOffset, uint16_t sLines, uint32_t copyLength) +{ + LocalTensor yLocal = outQueueY.DeQue(); + DataCopyExtParams copyParams; + + if (this->isAligned == true) { + copyParams.blockCount = sLines; + copyParams.blockLen = this->dBytes; + copyParams.srcStride = 0; + if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND || + this->layout == LAYOUT_R_B1SD) { + DataCopy(yGm[yStartOffset], yLocal, copyLength); + } else if (this->layout == LAYOUT_BSND) { + copyParams.dstStride = (this->bcSecondDim - 1) * this->dBytes; + DataCopyPad(yGm[yStartOffset], yLocal, copyParams); + } else if (this->layout == LAYOUT_SBND) { + copyParams.dstStride = (this->bnSize - 1) * this->dBytes; + DataCopyPad(yGm[yStartOffset], yLocal, copyParams); + } + } else { + if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND || + this->layout == LAYOUT_R_B1SD) { + copyParams.blockCount = (uint16_t)(2 * sLines); + copyParams.blockLen = this->halfDBytes; + copyParams.srcStride = 0; + copyParams.dstStride = 0; + DataCopyPad(yGm[yStartOffset], yLocal, copyParams); + } else if (this->layout == LAYOUT_BSND) { + copyParams.blockCount = sLines; + copyParams.blockLen = this->halfDBytes; + copyParams.srcStride = this->halfDPadBlocks; + copyParams.dstStride = (2 * this->bcSecondDim - 1) * this->halfDBytes; + DataCopyPad(yGm[yStartOffset], yLocal, copyParams); + DataCopyPad(yGm[yStartOffset + this->halfDLength], yLocal[this->halfDPadLength], copyParams); + } else if (this->layout == LAYOUT_SBND) { + copyParams.blockCount = sLines; + copyParams.blockLen = this->halfDBytes; + copyParams.srcStride = this->halfDPadBlocks; + copyParams.dstStride = (2 * this->bnSize - 1) * this->halfDBytes; + DataCopyPad(yGm[yStartOffset], yLocal, copyParams); + DataCopyPad(yGm[yStartOffset + this->halfDLength], yLocal[this->halfDPadLength], copyParams); + } + } + outQueueY.FreeTensor(yLocal); +} + +template +__aicore__ inline void RotateHalf::Compute(LocalTensor &cos, LocalTensor &sin, uint32_t sLines, + uint32_t calcLength) +{ + LocalTensor xLocal = inQueueX.DeQue(); + LocalTensor yLocal = outQueueY.AllocTensor(); + this->XNewCopy(xLocal, yLocal, sLines); + this->ComputeInner(xLocal, yLocal, cos, sin, calcLength); + outQueueY.EnQue(yLocal); + inQueueX.FreeTensor(xLocal); +} + +} // namespace RotateHalfN +#endif // ROTATE_HALF_H diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_half_base.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_half_base.h new file mode 100644 index 000000000..ad438b5c9 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_half_base.h @@ -0,0 +1,283 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotate_half_base.h + * \brief + */ + +#ifndef ROTATE_HALF_BASE_H +#define ROTATE_HALF_BASE_H + +#include "kernel_operator.h" + +namespace RotateHalfN { +using namespace AscendC; + +constexpr uint8_t REPEAT_MAX = 255; +constexpr int32_t SINGLE_BUFFER = 1; +constexpr int32_t DOUBLE_BUFFER = 2; +constexpr uint32_t BYTE_OF_BLOCK = 32; +constexpr uint32_t BYTE_OF_REPEAT = 256; + +constexpr uint16_t LAYOUT_BNSD = 1; +constexpr uint16_t LAYOUT_BSND = 2; +constexpr uint16_t LAYOUT_SBND = 3; +constexpr uint16_t LAYOUT_NO_BROADCAST = 4; +constexpr uint16_t LAYOUT_BND = 5; +constexpr uint16_t LAYOUT_R_B1SD = 6; + +template +class RotateHalfBase { +public: + __aicore__ inline RotateHalfBase(){}; + __aicore__ inline void BaseMemberInit(const RotaryPositionEmbeddingTilingData &tilingData); + +protected: + __aicore__ inline void GetTilingData(const RotateHalfParams &tiling); + __aicore__ inline void SinCompute(LocalTensor &sin, uint32_t sLines); + __aicore__ inline void ComputeInner(LocalTensor &x, LocalTensor &xNew, LocalTensor &cos, + LocalTensor &sin, uint32_t calcLength); + __aicore__ inline void RBroadCast(LocalTensor &cos, LocalTensor &sin, uint32_t broadcastLines); + __aicore__ inline void XNewCopy(LocalTensor &x, LocalTensor &xNew, uint16_t sLines); + + bool isAligned; + uint16_t layout; + uint64_t gmLength; + uint64_t bcFirstDim; + uint64_t bcSecondDim; + uint64_t dLength; + uint64_t dPadLength; + uint64_t halfDLength; + uint64_t halfDPadLength; + uint64_t totalSLines; + uint64_t storeSLines; + uint64_t storeDataLength; + uint64_t storePadDataLength; + uint64_t ubLoop; + uint64_t ubLast; + uint64_t formerCoreNum; + uint64_t tailCoreNum; + uint64_t formerSLines; + uint64_t tailSLines; + uint64_t coreSLines; + uint64_t xDataLength; + uint64_t rDataLength; + uint64_t ubLastDataLength; + uint64_t ubLastPadDataLength; + uint64_t xOffset; + uint64_t rOffset; + + uint8_t repeatStride; + uint32_t dataEachRepeat; + uint32_t bnSize; + uint32_t ndSize; + uint32_t bndSize; + uint32_t dBytes; + uint32_t halfDBytes; + uint32_t halfDPadBlocks; + uint32_t innerHalfLoop; + uint32_t innerHalfLast; + uint64_t xAllocLength; + uint64_t rAllocLength; + uint64_t xCoreOffset; + uint64_t rCoreOffset; + uint64_t coreRelativeIdx; + + DataCopyPadExtParams noPadParams{false, 0, 0, 0}; +}; + +template +__aicore__ inline void RotateHalfBase::BaseMemberInit(const RotaryPositionEmbeddingTilingData &tilingData) +{ + const RotateHalfParams &tiling = tilingData.rotateHalfParams; + GetTilingData(tiling); + + // intermediate variable + repeatStride = dPadLength * sizeof(CmpT) / BYTE_OF_BLOCK; + dataEachRepeat = BYTE_OF_REPEAT / sizeof(CmpT); + innerHalfLoop = halfDLength / dataEachRepeat; + innerHalfLast = halfDLength % dataEachRepeat; + bnSize = bcFirstDim * bcSecondDim; + ndSize = bcSecondDim * dLength; + bndSize = bnSize * dLength; + dBytes = dLength * sizeof(OriT); + halfDBytes = halfDLength * sizeof(OriT); + halfDPadBlocks = halfDPadLength * sizeof(OriT) / BYTE_OF_BLOCK; + + // gm and ub space params + xAllocLength = gmLength; + rAllocLength = rDataLength; + rCoreOffset = rDataLength; + if (layout == LAYOUT_BNSD || layout == LAYOUT_NO_BROADCAST) { + xCoreOffset = coreSLines * dLength; + } else if (layout == LAYOUT_BSND) { + xCoreOffset = coreSLines * ndSize; + } else if (layout == LAYOUT_SBND) { + xCoreOffset = coreSLines * bndSize; + } else if (layout == LAYOUT_BND) { + xCoreOffset = xDataLength; + xAllocLength = xDataLength; + rCoreOffset = 0; + rOffset = 0; + } else if (layout == LAYOUT_R_B1SD) { + xCoreOffset = coreSLines * dLength; + rCoreOffset = xCoreOffset; + rAllocLength = bcFirstDim * totalSLines * dLength; + } +} + +template +__aicore__ inline void RotateHalfBase::GetTilingData(const RotateHalfParams &tiling) +{ + isAligned = tiling.isAligned == 1; + layout = tiling.tilingMode; + gmLength = tiling.gmLength; + bcFirstDim = tiling.broadcastFirstDim; + bcSecondDim = tiling.broadcastSecondDim; + dLength = tiling.dLength; + dPadLength = tiling.dPadLength; + halfDLength = tiling.halfDLength; + halfDPadLength = tiling.halfDPadLength; + totalSLines = tiling.totalSLines; + storeSLines = tiling.storeSLines; + storeDataLength = tiling.storeDataLength; + storePadDataLength = tiling.storePadDataLength; + formerCoreNum = tiling.formerCoreNum; + tailCoreNum = tiling.tailCoreNum; + formerSLines = tiling.formerSLines; + tailSLines = tiling.tailSLines; + + if (GetBlockIdx() < formerCoreNum) { + coreRelativeIdx = GetBlockIdx(); + coreSLines = formerSLines; + ubLoop = tiling.formerUbLoop; + ubLast = tiling.formerUbLast; + xDataLength = tiling.formerXDataLength; + rDataLength = tiling.formerRDataLength; + ubLastDataLength = tiling.formerUbLastDataLength; + ubLastPadDataLength = tiling.formerUbLastPadDataLength; + xOffset = 0; + rOffset = 0; + } else { + coreRelativeIdx = GetBlockIdx() - formerCoreNum; + coreSLines = tailSLines; + ubLoop = tiling.tailUbLoop; + ubLast = tiling.tailUbLast; + xDataLength = tiling.tailXDataLength; + rDataLength = tiling.tailRDataLength; + ubLastDataLength = tiling.tailUbLastDataLength; + ubLastPadDataLength = tiling.tailUbLastPadDataLength; + xOffset = tiling.formerXCoreOffset; + rOffset = tiling.formerRCoreOffset; + } +} + +/* sin_l = -1 * sin_l */ +template +__aicore__ inline void RotateHalfBase::SinCompute(LocalTensor &sin, uint32_t sLines) +{ + uint32_t repeatOffset, innerOffset; + uint32_t repeatLoop = sLines / REPEAT_MAX; + uint8_t repeatLast = sLines % REPEAT_MAX; + UnaryRepeatParams repeatParams{1, 1, repeatStride, repeatStride}; + + for (uint32_t i = 0; i < repeatLoop; i++) { + repeatOffset = i * REPEAT_MAX * dPadLength; + for (uint32_t j = 0; j < innerHalfLoop; j++) { + innerOffset = j * dataEachRepeat + repeatOffset; + Muls(sin[innerOffset], sin[innerOffset], (CmpT)(-1.0), dataEachRepeat, REPEAT_MAX, repeatParams); + } + if (innerHalfLast > 0) { + innerOffset = innerHalfLoop * dataEachRepeat + repeatOffset; + Muls(sin[innerOffset], sin[innerOffset], (CmpT)(-1.0), innerHalfLast, REPEAT_MAX, repeatParams); + } + } + if (repeatLast > 0) { + repeatOffset = repeatLoop * REPEAT_MAX * dPadLength; + for (uint32_t j = 0; j < innerHalfLoop; j++) { + innerOffset = j * dataEachRepeat + repeatOffset; + Muls(sin[innerOffset], sin[innerOffset], (CmpT)(-1.0), dataEachRepeat, repeatLast, repeatParams); + } + if (innerHalfLast > 0) { + innerOffset = innerHalfLoop * dataEachRepeat + repeatOffset; + Muls(sin[innerOffset], sin[innerOffset], (CmpT)(-1.0), innerHalfLast, repeatLast, repeatParams); + } + } +} + +/* x = x * cos, xNew = xNew * sin, y = x + xNew */ +template +__aicore__ inline void RotateHalfBase::ComputeInner(LocalTensor &x, LocalTensor &xNew, + LocalTensor &cos, LocalTensor &sin, + uint32_t calcLength) +{ + Mul(x, x, cos, calcLength); + Mul(xNew, xNew, sin, calcLength); + Add(xNew, xNew, x, calcLength); +} + +/* broadcast cos, sin from (1, D) to (storeSLines, D) or (ubLast, D) shape */ +template +__aicore__ inline void RotateHalfBase::RBroadCast(LocalTensor &cos, LocalTensor &sin, + uint32_t broadcastLines) +{ + uint32_t repeatOffset, copySrcOffset, copyDstOffset; + uint32_t innerLoop = dPadLength / dataEachRepeat; + uint32_t innerLast = dPadLength % dataEachRepeat; + uint32_t repeatLoop = broadcastLines / REPEAT_MAX; + uint8_t repeatLast = broadcastLines % REPEAT_MAX; + CopyRepeatParams repParams{1, 1, repeatStride, 0}; + + for (uint32_t i = 0; i < repeatLoop; i++) { + repeatOffset = i * REPEAT_MAX * dPadLength; + for (uint32_t j = 0; j < innerLoop; j++) { + copySrcOffset = j * dataEachRepeat; + copyDstOffset = copySrcOffset + repeatOffset + dPadLength; + Copy(cos[copyDstOffset], cos[copySrcOffset], dataEachRepeat, REPEAT_MAX, repParams); + Copy(sin[copyDstOffset], sin[copySrcOffset], dataEachRepeat, REPEAT_MAX, repParams); + } + if (innerLast > 0) { + copySrcOffset = innerLoop * dataEachRepeat; + copyDstOffset = copySrcOffset + repeatOffset + dPadLength; + Copy(cos[copyDstOffset], cos[copySrcOffset], innerLast, REPEAT_MAX, repParams); + Copy(sin[copyDstOffset], sin[copySrcOffset], innerLast, REPEAT_MAX, repParams); + } + } + if (repeatLast > 0) { + repeatOffset = repeatLoop * REPEAT_MAX * dPadLength; + for (uint32_t j = 0; j < innerLoop; j++) { + copySrcOffset = j * dataEachRepeat; + copyDstOffset = copySrcOffset + repeatOffset + dPadLength; + Copy(cos[copyDstOffset], cos[copySrcOffset], dataEachRepeat, repeatLast, repParams); + Copy(sin[copyDstOffset], sin[copySrcOffset], dataEachRepeat, repeatLast, repParams); + } + if (innerLast > 0) { + copySrcOffset = innerLoop * dataEachRepeat; + copyDstOffset = copySrcOffset + repeatOffset + dPadLength; + Copy(cos[copyDstOffset], cos[copySrcOffset], innerLast, repeatLast, repParams); + Copy(sin[copyDstOffset], sin[copySrcOffset], innerLast, repeatLast, repParams); + } + } +} + +/* copy x to xNew: x_l --> xNew_r, x_r --> xNew_l */ +template +__aicore__ inline void RotateHalfBase::XNewCopy(LocalTensor &x, LocalTensor &xNew, + uint16_t sLines) +{ + uint16_t stride = this->halfDPadLength * sizeof(CmpT) / BYTE_OF_BLOCK; + DataCopyParams copyParams{sLines, stride, stride, stride}; + DataCopy(xNew, x[this->halfDPadLength], copyParams); + DataCopy(xNew[this->halfDPadLength], x, copyParams); +} + +} // namespace RotateHalfN +#endif // ROTATE_HALF_BASE_H \ No newline at end of file diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_half_bf16.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_half_bf16.h new file mode 100644 index 000000000..ec4bb49e6 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_half_bf16.h @@ -0,0 +1,385 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotate_half_bf16.h + * \brief + */ +#ifndef ROTATE_HALF_BF16_H +#define ROTATE_HALF_BF16_H + +#include "rotate_half_base.h" + +namespace RotateHalfN { +using namespace AscendC; + +template +class RotateHalfBf16 : public RotateHalfBase { +public: + __aicore__ inline RotateHalfBf16(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RotaryPositionEmbeddingTilingData &tilingData); + __aicore__ inline void Process(); + +protected: + TPipe pipe; + TQue inQueueX; + TQue inQueueCos; + TQue inQueueSin; + TQue outQueueY; + TBuf xBuf; + TBuf xNewBuf; + TBuf cosBuf; + TBuf sinBuf; + GlobalTensor xGm; + GlobalTensor cosGm; + GlobalTensor sinGm; + GlobalTensor yGm; + + __aicore__ inline void NormalProcessBf16(); + __aicore__ inline void RB1sdProcessBf16(); + __aicore__ inline void BndProcessBf16(); + __aicore__ inline void SingleStepProcess(uint32_t progress, uint32_t sLines, uint64_t copyLength, + uint64_t calcLength); + __aicore__ inline void RB1sdSingleStepProcessBf16(uint32_t progress, uint64_t sLines, uint64_t xBatchStartOffset, + uint64_t rBatchStartOffset, uint64_t copyLength, + uint64_t calcLength); + __aicore__ inline void Compute(LocalTensor &cos, LocalTensor &sin, uint32_t sLines, + uint32_t calcLength); + __aicore__ inline void CopyInR(uint64_t rStartOffset, uint16_t sLines, uint32_t copyLength); + __aicore__ inline void CopyInX(uint64_t xStartOffset, uint16_t sLines, uint32_t copyLength); + __aicore__ inline void CopyOut(uint64_t yOffset, uint16_t sLines, uint32_t copyLength); +}; + +template +__aicore__ inline void RotateHalfBf16::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RotaryPositionEmbeddingTilingData &tilingData) +{ + this->BaseMemberInit(tilingData); + + xGm.SetGlobalBuffer((__gm__ OriT *)x + this->xOffset + this->xCoreOffset * this->coreRelativeIdx, + this->xAllocLength); + yGm.SetGlobalBuffer((__gm__ OriT *)y + this->xOffset + this->xCoreOffset * this->coreRelativeIdx, + this->xAllocLength); + cosGm.SetGlobalBuffer((__gm__ OriT *)cos + this->rOffset + this->rCoreOffset * this->coreRelativeIdx, + this->rAllocLength); + sinGm.SetGlobalBuffer((__gm__ OriT *)sin + this->rOffset + this->rCoreOffset * this->coreRelativeIdx, + this->rAllocLength); + + pipe.InitBuffer(inQueueX, DOUBLE_BUFFER, this->storePadDataLength * sizeof(OriT)); + pipe.InitBuffer(outQueueY, DOUBLE_BUFFER, this->storePadDataLength * sizeof(OriT)); + pipe.InitBuffer(inQueueCos, DOUBLE_BUFFER, this->storePadDataLength * sizeof(OriT)); + pipe.InitBuffer(inQueueSin, DOUBLE_BUFFER, this->storePadDataLength * sizeof(OriT)); + + pipe.InitBuffer(xBuf, this->storePadDataLength * sizeof(CmpT)); + pipe.InitBuffer(cosBuf, this->storePadDataLength * sizeof(CmpT)); + pipe.InitBuffer(sinBuf, this->storePadDataLength * sizeof(CmpT)); + pipe.InitBuffer(xNewBuf, this->storePadDataLength * sizeof(CmpT)); +} + +template +__aicore__ inline void RotateHalfBf16::Process() +{ + if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_BSND || this->layout == LAYOUT_SBND || + this->layout == LAYOUT_NO_BROADCAST) { + NormalProcessBf16(); + } else if (this->layout == LAYOUT_R_B1SD) { + RB1sdProcessBf16(); + } else if (this->layout == LAYOUT_BND) { + BndProcessBf16(); + } +} + +template +__aicore__ inline void RotateHalfBf16::NormalProcessBf16() +{ + for (uint32_t progress = 0; progress < this->ubLoop; progress++) { + SingleStepProcess(progress, this->storeSLines, this->storeDataLength, this->storePadDataLength); + } + if (this->ubLast > 0) { + SingleStepProcess(this->ubLoop, this->ubLast, this->ubLastDataLength, this->ubLastPadDataLength); + } +} + +template +__aicore__ inline void RotateHalfBf16::RB1sdProcessBf16() +{ + uint64_t totalSdLength = this->totalSLines * this->dLength; + uint64_t totalNsdLength = totalSdLength * this->bcSecondDim; + uint64_t xBatchOffset, rBatchOffset; + for (uint32_t loop = 0; loop < this->bcFirstDim; loop++) { + xBatchOffset = loop * totalNsdLength; + rBatchOffset = loop * totalSdLength; + for (uint32_t progress = 0; progress < this->ubLoop; progress++) { + RB1sdSingleStepProcessBf16(progress, this->storeSLines, xBatchOffset, rBatchOffset, this->storeDataLength, + this->storePadDataLength); + } + if (this->ubLast > 0) { + RB1sdSingleStepProcessBf16(this->ubLoop, this->ubLast, xBatchOffset, rBatchOffset, this->ubLastDataLength, + this->ubLastPadDataLength); + } + } +} + +template +__aicore__ inline void RotateHalfBf16::BndProcessBf16() +{ + CopyInR(0, 1, this->dLength); + LocalTensor cosLocal = inQueueCos.DeQue(); + LocalTensor sinLocal = inQueueSin.DeQue(); + LocalTensor cosFp32 = cosBuf.Get(); + LocalTensor sinFp32 = sinBuf.Get(); + + Cast(cosFp32, cosLocal, RoundMode::CAST_NONE, this->dPadLength); + Cast(sinFp32, sinLocal, RoundMode::CAST_NONE, this->dPadLength); + inQueueCos.FreeTensor(cosLocal); + inQueueSin.FreeTensor(sinLocal); + Muls(sinFp32, sinFp32, (CmpT)(-1.0), this->halfDPadLength); + uint32_t broadcastLines = this->ubLoop > 0 ? this->storeSLines - 1 : this->ubLast - 1; + if (broadcastLines > 0) { + this->RBroadCast(cosFp32, sinFp32, broadcastLines); + } + uint64_t xOffset; + for (uint32_t progress = 0; progress < this->ubLoop; progress++) { + xOffset = progress * this->storeDataLength; + CopyInX(xOffset, this->storeSLines, this->storeDataLength); + Compute(cosFp32, sinFp32, this->storeSLines, this->storePadDataLength); + CopyOut(xOffset, this->storeSLines, this->storeDataLength); + } + if (this->ubLast > 0) { + xOffset = this->ubLoop * this->storeDataLength; + CopyInX(xOffset, this->ubLast, this->ubLastDataLength); + Compute(cosFp32, sinFp32, this->ubLast, this->ubLastPadDataLength); + CopyOut(xOffset, this->ubLast, this->ubLastDataLength); + } +} + +template +__aicore__ inline void RotateHalfBf16::SingleStepProcess(uint32_t progress, uint32_t sLines, + uint64_t copyLength, uint64_t calcLength) +{ + uint64_t xOffset, rOffset, bnLoopXStartOffset, progressOffset, batchOffset; + rOffset = progress * this->storeDataLength; + CopyInR(rOffset, sLines, copyLength); + LocalTensor cosLocal = inQueueCos.DeQue(); + LocalTensor sinLocal = inQueueSin.DeQue(); + LocalTensor cosFp32 = cosBuf.Get(); + LocalTensor sinFp32 = sinBuf.Get(); + + Cast(cosFp32, cosLocal, RoundMode::CAST_NONE, calcLength); + Cast(sinFp32, sinLocal, RoundMode::CAST_NONE, calcLength); + inQueueCos.FreeTensor(cosLocal); + inQueueSin.FreeTensor(sinLocal); + this->SinCompute(sinFp32, sLines); + + if (this->layout == LAYOUT_BNSD) { + uint64_t totalSdSize = this->totalSLines * this->dLength; + bnLoopXStartOffset = progress * this->storeDataLength; + for (uint32_t bnLoop = 0; bnLoop < this->bnSize; bnLoop++) { + xOffset = bnLoopXStartOffset + bnLoop * totalSdSize; + CopyInX(xOffset, sLines, copyLength); + Compute(cosFp32, sinFp32, sLines, calcLength); + CopyOut(xOffset, sLines, copyLength); + } + } else if (this->layout == LAYOUT_BSND) { + uint64_t totalSndSize = this->totalSLines * this->ndSize; + progressOffset = progress * this->bcSecondDim * this->storeDataLength; + for (uint32_t batchLoop = 0; batchLoop < this->bcFirstDim; batchLoop++) { + batchOffset = batchLoop * totalSndSize; + for (uint32_t nLoop = 0; nLoop < this->bcSecondDim; nLoop++) { + xOffset = nLoop * this->dLength + batchOffset + progressOffset; + CopyInX(xOffset, sLines, copyLength); + Compute(cosFp32, sinFp32, sLines, calcLength); + CopyOut(xOffset, sLines, copyLength); + } + } + } else if (this->layout == LAYOUT_SBND) { + bnLoopXStartOffset = progress * this->storeDataLength * this->bnSize; + for (uint32_t bnLoop = 0; bnLoop < this->bnSize; bnLoop++) { + xOffset = bnLoopXStartOffset + bnLoop * this->dLength; + CopyInX(xOffset, sLines, copyLength); + Compute(cosFp32, sinFp32, sLines, calcLength); + CopyOut(xOffset, sLines, copyLength); + } + } else if (this->layout == LAYOUT_NO_BROADCAST) { + CopyInX(rOffset, sLines, copyLength); + Compute(cosFp32, sinFp32, sLines, calcLength); + CopyOut(rOffset, sLines, copyLength); + } +} + +template +__aicore__ inline void RotateHalfBf16::RB1sdSingleStepProcessBf16(uint32_t progress, uint64_t sLines, + uint64_t xBatchStartOffset, + uint64_t rBatchStartOffset, + uint64_t copyLength, uint64_t calcLength) +{ + CopyInR(progress * this->storeDataLength + rBatchStartOffset, sLines, copyLength); + LocalTensor cosLocal = inQueueCos.DeQue(); + LocalTensor sinLocal = inQueueSin.DeQue(); + LocalTensor cosFp32 = cosBuf.Get(); + LocalTensor sinFp32 = sinBuf.Get(); + + Cast(cosFp32, cosLocal, RoundMode::CAST_NONE, calcLength); + Cast(sinFp32, sinLocal, RoundMode::CAST_NONE, calcLength); + inQueueCos.FreeTensor(cosLocal); + inQueueSin.FreeTensor(sinLocal); + this->SinCompute(sinFp32, sLines); + + uint64_t xOffset, progressXOffset; + progressXOffset = progress * this->storeDataLength + xBatchStartOffset; + for (uint32_t nLoop = 0; nLoop < this->bcSecondDim; nLoop++) { + xOffset = nLoop * this->totalSLines * this->dLength + progressXOffset; + CopyInX(xOffset, sLines, copyLength); + Compute(cosFp32, sinFp32, sLines, calcLength); + CopyOut(xOffset, sLines, copyLength); + } +} + +template +__aicore__ inline void RotateHalfBf16::CopyInR(uint64_t rStartOffset, uint16_t sLines, uint32_t copyLength) +{ + LocalTensor sinLocal = inQueueSin.AllocTensor(); + LocalTensor cosLocal = inQueueCos.AllocTensor(); + if (this->isAligned == true) { + DataCopy(sinLocal, sinGm[rStartOffset], copyLength); + DataCopy(cosLocal, cosGm[rStartOffset], copyLength); + } else { + DataCopyExtParams copyParams{(uint16_t)(2 * sLines), // blockCount + this->halfDBytes, // blockLen + 0, // srcStride(bytes) + 0, // dstStride(block) + 0}; + DataCopyPad(sinLocal, sinGm[rStartOffset], copyParams, this->noPadParams); + DataCopyPad(cosLocal, cosGm[rStartOffset], copyParams, this->noPadParams); + } + inQueueSin.EnQue(sinLocal); + inQueueCos.EnQue(cosLocal); +} + +template +__aicore__ inline void RotateHalfBf16::CopyInX(uint64_t xStartOffset, uint16_t storeSLines, + uint32_t copyLength) +{ + LocalTensor xLocal = inQueueX.AllocTensor(); + DataCopyExtParams copyParams; + + if (this->isAligned == true) { + if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND || + this->layout == LAYOUT_R_B1SD) { + DataCopy(xLocal, xGm[xStartOffset], copyLength); + } else if (this->layout == LAYOUT_BSND) { + copyParams.blockCount = storeSLines; + copyParams.blockLen = this->dBytes; + copyParams.srcStride = (this->bcSecondDim - 1) * this->dBytes; + copyParams.dstStride = 0; + DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams); + } else if (this->layout == LAYOUT_SBND) { + copyParams.blockCount = storeSLines; + copyParams.blockLen = this->dBytes; + copyParams.srcStride = (this->bnSize - 1) * this->dBytes; + copyParams.dstStride = 0; + DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams); + } + } else { + if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND || + this->layout == LAYOUT_R_B1SD) { + copyParams.blockCount = (uint16_t)(2 * storeSLines); + copyParams.blockLen = this->halfDBytes; + copyParams.srcStride = 0; + copyParams.dstStride = 0; + DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams); + } else if (this->layout == LAYOUT_BSND) { + copyParams.blockCount = storeSLines; + copyParams.blockLen = this->halfDBytes; + copyParams.srcStride = (2 * this->bcSecondDim - 1) * this->halfDBytes; + copyParams.dstStride = this->halfDPadBlocks; + DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams); + DataCopyPad(xLocal[this->halfDPadLength], xGm[xStartOffset + this->halfDLength], copyParams, + this->noPadParams); + } else if (this->layout == LAYOUT_SBND) { + copyParams.blockCount = storeSLines; + copyParams.blockLen = this->halfDBytes; + copyParams.srcStride = (this->bnSize * 2 - 1) * this->halfDBytes; + copyParams.dstStride = this->halfDPadBlocks; + DataCopyPad(xLocal, xGm[xStartOffset], copyParams, this->noPadParams); + DataCopyPad(xLocal[this->halfDPadLength], xGm[xStartOffset + this->halfDLength], copyParams, + this->noPadParams); + } + } + inQueueX.EnQue(xLocal); +} + +template +__aicore__ inline void RotateHalfBf16::CopyOut(uint64_t yOffset, uint16_t storeSLines, uint32_t copyLength) +{ + LocalTensor yLocal = outQueueY.DeQue(); + DataCopyExtParams copyParams; + + if (this->isAligned == true) { + copyParams.blockCount = storeSLines; + copyParams.blockLen = this->dBytes; + copyParams.srcStride = 0; + if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND || + this->layout == LAYOUT_R_B1SD) { + DataCopy(yGm[yOffset], yLocal, copyLength); + } else if (this->layout == LAYOUT_BSND) { + copyParams.dstStride = (this->bcSecondDim - 1) * this->dBytes; + DataCopyPad(yGm[yOffset], yLocal, copyParams); + } else if (this->layout == LAYOUT_SBND) { + copyParams.dstStride = (this->bnSize - 1) * this->dBytes; + DataCopyPad(yGm[yOffset], yLocal, copyParams); + } + } else { + if (this->layout == LAYOUT_BNSD || this->layout == LAYOUT_NO_BROADCAST || this->layout == LAYOUT_BND || + this->layout == LAYOUT_R_B1SD) { + copyParams.blockCount = (uint16_t)(2 * storeSLines); + copyParams.blockLen = this->halfDBytes; + copyParams.srcStride = 0; + copyParams.dstStride = 0; + DataCopyPad(yGm[yOffset], yLocal, copyParams); + } else if (this->layout == LAYOUT_BSND) { + copyParams.blockCount = storeSLines; + copyParams.blockLen = this->halfDBytes; + copyParams.srcStride = this->halfDPadBlocks; + copyParams.dstStride = (2 * this->bcSecondDim - 1) * this->halfDBytes; + DataCopyPad(yGm[yOffset], yLocal, copyParams); + DataCopyPad(yGm[yOffset + this->halfDLength], yLocal[this->halfDPadLength], copyParams); + } else if (this->layout == LAYOUT_SBND) { + copyParams.blockCount = storeSLines; + copyParams.blockLen = this->halfDBytes; + copyParams.srcStride = this->halfDPadBlocks; + copyParams.dstStride = (2 * this->bnSize - 1) * this->halfDBytes; + DataCopyPad(yGm[yOffset], yLocal, copyParams); + DataCopyPad(yGm[yOffset + this->halfDLength], yLocal[this->halfDPadLength], copyParams); + } + } + outQueueY.FreeTensor(yLocal); +} + +template +__aicore__ inline void RotateHalfBf16::Compute(LocalTensor &cos, LocalTensor &sin, + uint32_t sLines, uint32_t calcLength) +{ + LocalTensor xLocal = inQueueX.DeQue(); + LocalTensor yLocal = outQueueY.AllocTensor(); + LocalTensor xFp32 = xBuf.Get(); + LocalTensor xNewFp32 = xNewBuf.Get(); + + Cast(xFp32, xLocal, RoundMode::CAST_NONE, calcLength); + inQueueX.FreeTensor(xLocal); + + this->XNewCopy(xFp32, xNewFp32, sLines); + this->ComputeInner(xFp32, xNewFp32, cos, sin, calcLength); + Cast(yLocal, xNewFp32, RoundMode::CAST_RINT, calcLength); + outQueueY.EnQue(yLocal); +} +} // namespace RotateHalfN +#endif // ROTATE_HALF_BF16_H diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_common.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_common.h new file mode 100644 index 000000000..756f201d7 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_common.h @@ -0,0 +1,145 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotate_interleaved_common.h + * \brief + */ +#ifndef ROTATE_INTERLEAVED_COMMON_H +#define ROTATE_INTERLEAVED_COMMON_H +#include "kernel_operator.h" +#include "impl/dav_c220/kernel_operator_reg_others_impl.h" + +namespace RotateInterleavedN { +using namespace AscendC; +constexpr int32_t BUFFER_NUM = 1; +constexpr int32_t NUM_8 = 8; +constexpr uint8_t REPEAT_MAX = 255; +constexpr int32_t BLOCK_SIZE = 32; +constexpr int32_t MASK_FP16 = 128; +constexpr int32_t MASK_FP32 = 64; +constexpr int32_t ALIGN_16 = 16; +constexpr int32_t ALIGN_32 = 8; + +// SD -> BSND +template +__aicore__ inline void BroadCastTriToBsnd(LocalTensor &tri, uint32_t batchSize, uint32_t calcLen, uint32_t numHeads, + uint32_t headDimAlign) +{ + DataCopyParams intriParams; + intriParams.blockCount = static_cast(calcLen); + intriParams.blockLen = static_cast(headDimAlign * sizeof(T) / BLOCK_SIZE); + intriParams.srcStride = static_cast((numHeads - 1) * headDimAlign * sizeof(T) / BLOCK_SIZE); + intriParams.dstStride = static_cast((numHeads - 1) * headDimAlign * sizeof(T) / BLOCK_SIZE); + // SD -> SND + for (uint32_t numHeadsIdx = 1; numHeadsIdx < numHeads; ++numHeadsIdx) { + DataCopy(tri[numHeadsIdx * headDimAlign], tri, intriParams); + } + + intriParams.blockCount = 1; + intriParams.blockLen = static_cast(calcLen * numHeads * headDimAlign * sizeof(T) / BLOCK_SIZE); + intriParams.srcStride = 0; + intriParams.dstStride = 0; + // SND -> BSND + for (uint32_t batchIdx = 1; batchIdx < batchSize; ++batchIdx) { + DataCopy(tri[batchIdx * calcLen * numHeads * headDimAlign], tri, intriParams); + } +} + +// D -> BND +template +__aicore__ inline void BroadCastTriToB1nd(LocalTensor &tri, uint32_t calcLen, uint32_t numHeads, + uint32_t headDimAlign) +{ + using ElementType = + typename std::conditional::value || std::is_same::value, half, T>::type; + LocalTensor triNew = tri.template ReinterpretCast(); + + const int32_t mask = (std::is_same::value || std::is_same::value) ? MASK_FP16 : MASK_FP32; + const int32_t count = headDimAlign / mask; + const int32_t remain = headDimAlign % mask; + const int32_t repeatTimes = calcLen * numHeads - 1; + const int32_t repeatTimesLoop = repeatTimes / REPEAT_MAX; + const int32_t repeatTimesRemain = repeatTimes % REPEAT_MAX; + + CopyRepeatParams repeatParams; + repeatParams.dstStride = 1; + repeatParams.srcStride = 1; + repeatParams.dstRepeatSize = headDimAlign * sizeof(ElementType) / BLOCK_SIZE; + repeatParams.srcRepeatSize = 0; + + for (uint32_t loopIdx = 0; loopIdx < count; ++loopIdx) { + for (uint32_t i = 0; i < repeatTimesLoop; ++i) { + Copy(triNew[headDimAlign + loopIdx * mask + i * REPEAT_MAX * headDimAlign], triNew[loopIdx * mask], mask, + REPEAT_MAX, repeatParams); + } + Copy(triNew[headDimAlign * (repeatTimesLoop * REPEAT_MAX + 1) + loopIdx * mask], triNew[loopIdx * mask], mask, + repeatTimesRemain, repeatParams); + } + if (remain != 0) { + for (uint32_t i = 0; i < repeatTimesLoop; ++i) { + Copy(triNew[headDimAlign + count * mask + i * REPEAT_MAX * headDimAlign], triNew[count * mask], remain, + REPEAT_MAX, repeatParams); + } + Copy(triNew[headDimAlign * (repeatTimesLoop * REPEAT_MAX + 1) + count * mask], triNew[count * mask], remain, + repeatTimesRemain, repeatParams); + } +} + +// The minimum amount of data set by offset is 8 +__aicore__ inline void SetGatherSrcOffset(LocalTensor &gatherOffset, int32_t count, int32_t srcSizeof) +{ + for (int32_t i = 0; i < NUM_8; ++i) { + gatherOffset.SetValue(i, i ^ 1); // XOR with 1 to swap even and odd indices + } + + int32_t scalarValue = 8; + while (scalarValue < count) { + int32_t nextValue = scalarValue * 2; + if (nextValue < count) { + Adds(gatherOffset[scalarValue], gatherOffset, scalarValue, scalarValue); + } else { + Adds(gatherOffset[scalarValue], gatherOffset, scalarValue, count - scalarValue); + break; + } + scalarValue = nextValue; + } + Muls(gatherOffset, gatherOffset, srcSizeof, count); +} + +// count < 256 * 64 and count % 8 == 0 +__aicore__ inline void InterleavedInversion(LocalTensor &srcInversion, int32_t count, bool isOffset = false) +{ + SetMaskNorm(); + + const int32_t mask = MASK_FP32; + const int32_t repeatTimes = count / mask; + const int32_t remainder = count % mask; + + // Define masks based on the 'isOffset' flag + const uint64_t fullMask = isOffset ? 0xAAAAAAAAAAAAAAAA : 0x5555555555555555; + const uint64_t partialMask = isOffset ? 0xAA : 0x55; + + // Apply the mask and multiplication for the full + SetVectorMask(0, fullMask); + Muls(srcInversion, srcInversion, float(-1), MASK_PLACEHOLDER, repeatTimes, {1, 1, 8, 8}); + + // Apply the mask and multiplication for the remainder if needed + if (remainder) { + SetVectorMask(0, partialMask); + Muls(srcInversion[repeatTimes * MASK_FP32], srcInversion[repeatTimes * MASK_FP32], float(-1), + MASK_PLACEHOLDER, count % MASK_FP32 / NUM_8, {1, 1, 1, 1}); + } + ResetMask(); +} + +} // namespace RotateInterleavedN + +#endif // ROTATE_INTERLEAVED_COMMON_H diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bs.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bs.h new file mode 100644 index 000000000..ccae90ee2 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bs.h @@ -0,0 +1,316 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotate_interleaved_split_bs.h + * \brief + */ +#ifndef ROTATE_INTERLEAVED_SPLIT_BS_H +#define ROTATE_INTERLEAVED_SPLIT_BS_H +#include "rotate_interleaved_common.h" + +namespace RotateInterleavedN { +using namespace AscendC; + +template +class InterleavedSplitBS { +public: + __aicore__ inline InterleavedSplitBS(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RopeRegbaseTilingData *tiling, TPipe *pipe); + __aicore__ inline void Process(); + +protected: + GlobalTensor xGm; + GlobalTensor cosGm; + GlobalTensor sinGm; + GlobalTensor yGm; + TQue inQueX; + TQue inQueCos; + TQue outQueY; + TBuf tmpFp32Buf1; + TBuf tmpFp32Buf2; + TBuf tmpFp32Buf3; + TBuf gatherOffsetBuf; + const RopeRegbaseTilingData* tiling_; + + // tilingdata + uint64_t batchSize; + uint64_t seqLen; + uint64_t numHeads; + uint64_t headDim; + uint64_t frontCoreNum; + uint64_t tailCoreNum; + uint64_t coreCalcNum; + uint64_t coreCalcTail; + uint64_t ubCalcBNum; + uint64_t ubCalcBLoop; + uint64_t ubCalcBTail; + uint64_t allHeadDim; + uint64_t start; + uint64_t ioOffsetAll; + uint64_t bufferNdSizeAll; + + // init tmp data + uint32_t blockIdx; + uint32_t ubCalcSeqLoop; + uint64_t ioOffset; + uint64_t triOffset; + uint64_t bufferBsndSize; + uint64_t bufferSdSize; + uint64_t bufferNdSize; + uint64_t bufferLenSize; + uint64_t gatherOffsetLenSize; + + __aicore__ inline void InitData(const RopeRegbaseTilingData *tiling); + __aicore__ inline void CopyInX(LocalTensor &x, uint32_t seqIdx, uint32_t batchIdx, uint32_t calcLen); + __aicore__ inline void CopyInCos(LocalTensor &cos, uint32_t seqIdx, uint32_t calcLen); + __aicore__ inline void CopyInSin(LocalTensor &sin, uint32_t seqIdx, uint32_t calcLen); + __aicore__ inline void CopyOut(uint32_t seqIdx, uint32_t batchIdx, uint32_t calcLen); + __aicore__ inline void Compute(uint32_t seqIdx, uint32_t batchIdx, LocalTensor &gatherOffsetCast, + uint32_t calcLen); + __aicore__ inline void ComputeCastFp32(uint32_t seqIdx, uint32_t batchIdx, LocalTensor &gatherOffsetCast, + uint32_t calcLen); +}; + +template +__aicore__ inline void InterleavedSplitBS::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RopeRegbaseTilingData *tiling, TPipe *pipe) +{ + InitData(tiling); + + blockIdx = GetBlockIdx(); + bufferSdSize = seqLen * headDim; + bufferNdSize = numHeads * headDim; + bufferNdSizeAll = numHeads * allHeadDim; + + if (blockIdx < frontCoreNum) { + ubCalcSeqLoop = coreCalcNum; + ioOffset = blockIdx * coreCalcNum * bufferNdSize; + ioOffsetAll = blockIdx * coreCalcNum * bufferNdSizeAll; + triOffset = blockIdx * coreCalcNum * headDim; + } else if (coreCalcTail != 0) { + ubCalcSeqLoop = coreCalcTail; + ioOffset = frontCoreNum * coreCalcNum * bufferNdSize + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSize; + ioOffsetAll = frontCoreNum * coreCalcNum * bufferNdSizeAll + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSizeAll; + triOffset = frontCoreNum * coreCalcNum * headDim + (blockIdx - frontCoreNum) * coreCalcTail * headDim; + } + + bufferBsndSize = batchSize * seqLen * bufferNdSizeAll; + xGm.SetGlobalBuffer((__gm__ T *)x + ioOffsetAll, bufferBsndSize); + yGm.SetGlobalBuffer((__gm__ T *)y + ioOffsetAll, bufferBsndSize); + cosGm.SetGlobalBuffer((__gm__ T *)cos + triOffset, bufferSdSize); + sinGm.SetGlobalBuffer((__gm__ T *)sin + triOffset, bufferSdSize); + + bufferLenSize = ubCalcBNum * bufferNdSize * sizeof(T); + pipe->InitBuffer(inQueX, BUFFER_NUM, bufferLenSize); + pipe->InitBuffer(inQueCos, BUFFER_NUM, bufferLenSize); + pipe->InitBuffer(outQueY, BUFFER_NUM, bufferLenSize); + + if constexpr (std::is_same::value || std::is_same::value) { + bufferLenSize = ubCalcBNum * bufferNdSize * sizeof(float); + pipe->InitBuffer(tmpFp32Buf1, bufferLenSize); + pipe->InitBuffer(tmpFp32Buf2, bufferLenSize); + pipe->InitBuffer(tmpFp32Buf3, bufferLenSize); + } + + gatherOffsetLenSize = bufferNdSize * sizeof(int32_t); + pipe->InitBuffer(gatherOffsetBuf, gatherOffsetLenSize); +} + +template +__aicore__ inline void InterleavedSplitBS::InitData(const RopeRegbaseTilingData *tiling) +{ + tiling_ = tiling; + batchSize = tiling_->batchSize; + seqLen = tiling_->seqLen; + numHeads = tiling_->numHeads; + headDim = tiling_->headDim; + frontCoreNum = tiling_->frontCoreNum; + tailCoreNum = tiling_->tailCoreNum; + coreCalcNum = tiling_->coreCalcNum; + coreCalcTail = tiling_->coreCalcTail; + ubCalcBNum = tiling_->ubCalcBNum; + ubCalcBLoop = tiling_->ubCalcBLoop; + ubCalcBTail = tiling_->ubCalcBTail; + allHeadDim = tiling_->allHeadDim; + start = tiling_->start; +} + +template +__aicore__ inline void InterleavedSplitBS::CopyInX(LocalTensor &x, uint32_t seqIdx, uint32_t batchIdx, + uint32_t calcLen) +{ + uint64_t startOffset = batchIdx * ubCalcBNum * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll; + DataCopyExtParams dataCopyParams; + for (uint32_t loopIdx = 0; loopIdx < calcLen; ++loopIdx) { + dataCopyParams.blockCount = numHeads; + dataCopyParams.blockLen = headDim * sizeof(T); + dataCopyParams.srcStride = (allHeadDim - headDim) * sizeof(T); + dataCopyParams.dstStride = 0; + DataCopyPad(x[loopIdx * numHeads * headDim], xGm[startOffset + loopIdx * seqLen * bufferNdSizeAll + start], + dataCopyParams, {false, 0, 0, 0}); + } + event_t eventIdMTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMTE2ToV); + WaitFlag(eventIdMTE2ToV); +} + +template +__aicore__ inline void InterleavedSplitBS::CopyInCos(LocalTensor &cos, uint32_t seqIdx, uint32_t calcLen) +{ + DataCopyExtParams bsDataCopyTriParams; + bsDataCopyTriParams.blockCount = 1; + bsDataCopyTriParams.blockLen = headDim * sizeof(T); + bsDataCopyTriParams.srcStride = 0; + bsDataCopyTriParams.dstStride = 0; + DataCopyPad(cos, cosGm[seqIdx * headDim], bsDataCopyTriParams, {false, 0, 0, 0}); + event_t eventId2MTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId2MTE2ToV); + WaitFlag(eventId2MTE2ToV); + BroadCastTriToB1nd(cos, calcLen, numHeads, headDim); +} + +template +__aicore__ inline void InterleavedSplitBS::CopyInSin(LocalTensor &sin, uint32_t seqIdx, uint32_t calcLen) +{ + event_t eventIdVToMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventIdVToMTE2); + WaitFlag(eventIdVToMTE2); + + DataCopyExtParams bsDataCopyTriParams; + bsDataCopyTriParams.blockCount = 1; + bsDataCopyTriParams.blockLen = headDim * sizeof(T); + bsDataCopyTriParams.srcStride = 0; + bsDataCopyTriParams.dstStride = 0; + DataCopyPad(sin, sinGm[seqIdx * headDim], bsDataCopyTriParams, {false, 0, 0, 0}); + event_t eventId3MTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId3MTE2ToV); + WaitFlag(eventId3MTE2ToV); + BroadCastTriToB1nd(sin, calcLen, numHeads, headDim); +} + +template +__aicore__ inline void InterleavedSplitBS::CopyOut(uint32_t seqIdx, uint32_t batchIdx, uint32_t calcLen) +{ + DataCopyExtParams dataCopyParams; + LocalTensor y = outQueY.DeQue(); + uint64_t startOffset = batchIdx * ubCalcBNum * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll; + for (uint32_t loopIdx = 0; loopIdx < calcLen; ++loopIdx) { + dataCopyParams.blockCount = numHeads; + dataCopyParams.blockLen = headDim * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (allHeadDim - headDim) * sizeof(T); + DataCopyPad(yGm[startOffset + loopIdx * seqLen * bufferNdSizeAll + start], y[loopIdx * numHeads * headDim], + dataCopyParams); + } + outQueY.FreeTensor(y); +} + +template +__aicore__ inline void InterleavedSplitBS::Process() +{ + LocalTensor gatherOffset = gatherOffsetBuf.Get(); + SetGatherSrcOffset(gatherOffset, headDim * numHeads, static_cast(sizeof(float))); + LocalTensor gatherOffsetCast = gatherOffset.ReinterpretCast(); + + if constexpr (std::is_same::value || std::is_same::value) { + for (uint32_t i = 0; i < ubCalcSeqLoop; ++i) { + for (uint32_t j = 0; j < (ubCalcBTail == 0 ? ubCalcBLoop : ubCalcBLoop - 1); ++j) { + ComputeCastFp32(i, j, gatherOffsetCast, ubCalcBNum); + CopyOut(i, j, ubCalcBNum); + } + if (ubCalcBTail != 0) { + ComputeCastFp32(i, ubCalcBLoop - 1, gatherOffsetCast, ubCalcBTail); + CopyOut(i, ubCalcBLoop - 1, ubCalcBTail); + } + } + } else { + for (uint32_t i = 0; i < ubCalcSeqLoop; ++i) { + for (uint32_t j = 0; j < (ubCalcBTail == 0 ? ubCalcBLoop : ubCalcBLoop - 1); ++j) { + Compute(i, j, gatherOffsetCast, ubCalcBNum); + CopyOut(i, j, ubCalcBNum); + } + if (ubCalcBTail != 0) { + Compute(i, ubCalcBLoop - 1, gatherOffsetCast, ubCalcBTail); + CopyOut(i, ubCalcBLoop - 1, ubCalcBTail); + } + } + } +} + +template +__aicore__ inline void InterleavedSplitBS::Compute(uint32_t seqIdx, uint32_t batchIdx, + LocalTensor &gatherOffsetCast, uint32_t calcLen) +{ + uint64_t calcTotalNum = calcLen * bufferNdSize; + + LocalTensor x = inQueX.AllocTensor(); + CopyInX(x, seqIdx, batchIdx, calcLen); + + LocalTensor cos = inQueCos.AllocTensor(); + CopyInCos(cos, seqIdx, calcLen); + + LocalTensor y = outQueY.AllocTensor(); + Mul(y, x, cos, calcTotalNum); + for (uint32_t i = 0; i < calcLen; ++i) { + Gather(x[i * bufferNdSize], x[i * bufferNdSize], gatherOffsetCast, 0, bufferNdSize); + } + + CopyInSin(cos, seqIdx, calcLen); + + Mul(x, x, cos, calcTotalNum); + inQueCos.FreeTensor(cos); + InterleavedInversion(x, calcTotalNum); + Add(y, y, x, calcTotalNum); + inQueX.FreeTensor(x); + outQueY.EnQue(y); +} + +template +__aicore__ inline void InterleavedSplitBS::ComputeCastFp32(uint32_t seqIdx, uint32_t batchIdx, + LocalTensor &gatherOffsetCast, uint32_t calcLen) +{ + uint64_t calcTotalNum = calcLen * bufferNdSize; + + LocalTensor x = inQueX.AllocTensor(); + CopyInX(x, seqIdx, batchIdx, calcLen); + LocalTensor tmp32BSBuf1 = tmpFp32Buf1.Get(); + Cast(tmp32BSBuf1, x, RoundMode::CAST_NONE, calcTotalNum); + inQueX.FreeTensor(x); + + LocalTensor cos = inQueCos.AllocTensor(); + CopyInCos(cos, seqIdx, calcLen); + LocalTensor tmp32Buf2 = tmpFp32Buf2.Get(); + Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum); + + LocalTensor tmp32Buf3 = tmpFp32Buf3.Get(); + Mul(tmp32Buf3, tmp32BSBuf1, tmp32Buf2, calcTotalNum); + + for (uint32_t i = 0; i < calcLen; ++i) { + Gather(tmp32BSBuf1[i * bufferNdSize], tmp32BSBuf1[i * bufferNdSize], gatherOffsetCast, 0, bufferNdSize); + } + + CopyInSin(cos, seqIdx, calcLen); + Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum); + inQueCos.FreeTensor(cos); + + Mul(tmp32BSBuf1, tmp32BSBuf1, tmp32Buf2, calcTotalNum); + InterleavedInversion(tmp32BSBuf1, calcTotalNum); + Add(tmp32Buf3, tmp32Buf3, tmp32BSBuf1, calcTotalNum); + + LocalTensor y = outQueY.AllocTensor(); + Cast(y, tmp32Buf3, RoundMode::CAST_RINT, calcTotalNum); + outQueY.EnQue(y); +} + +} // namespace RotateInterleavedN + +#endif // ROTATE_INTERLEAVED_SPLIT_BS_H diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bs_pad.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bs_pad.h new file mode 100644 index 000000000..37109d37d --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bs_pad.h @@ -0,0 +1,328 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotate_interleaved_split_bs_pad.h + * \brief + */ +#ifndef ROTATE_INTERLEAVED_SPLIT_BS_PAD_H +#define ROTATE_INTERLEAVED_SPLIT_BS_PAD_H +#include "rotate_interleaved_common.h" + +namespace RotateInterleavedN { +using namespace AscendC; + +template +class InterleavedSplitBSPad { +public: + __aicore__ inline InterleavedSplitBSPad(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RopeRegbaseTilingData *tiling, TPipe *pipe); + __aicore__ inline void Process(); + +protected: + GlobalTensor xGm; + GlobalTensor cosGm; + GlobalTensor sinGm; + GlobalTensor yGm; + TQue inQueX; + TQue inQueCos; + TQue outQueY; + TBuf tmpFp32Buf1; + TBuf tmpFp32Buf2; + TBuf tmpFp32Buf3; + TBuf gatherOffsetBuf; + const RopeRegbaseTilingData* tiling_; + uint64_t allHeadDim; + uint64_t start; + uint64_t ioOffsetAll; + uint64_t bufferNdSizeAll; + + // tilingdata + uint64_t batchSize; + uint64_t seqLen; + uint64_t numHeads; + uint64_t headDim; + uint64_t frontCoreNum; + uint64_t tailCoreNum; + uint64_t coreCalcNum; + uint64_t coreCalcTail; + uint64_t ubCalcBNum; + uint64_t ubCalcBLoop; + uint64_t ubCalcBTail; + + // init tmp data + uint32_t alignLen; + uint32_t headDimAlign; + uint32_t blockIdx; + uint32_t ubCalcSeq; + uint32_t ubCalcSeqTail; + uint32_t ubCalcSeqLoop; + uint64_t ioOffset; + uint64_t triOffset; + uint64_t bufferBsndSize; + uint64_t bufferSdSize; + uint64_t bufferNdSize; + uint64_t bufferLenSize; + uint64_t gatherOffsetLenSize; + uint32_t blockNum = BLOCK_SIZE / sizeof(T); + + __aicore__ inline void InitData(const RopeRegbaseTilingData *tiling); + __aicore__ inline void CopyInX(LocalTensor &x, uint32_t seqIdx, uint32_t batchIdx, uint32_t calcLen); + __aicore__ inline void CopyInCos(LocalTensor &cos, uint32_t seqIdx, uint32_t calcLen); + __aicore__ inline void CopyInSin(LocalTensor &sin, uint32_t seqIdx, uint32_t calcLen); + __aicore__ inline void CopyOut(uint32_t seqIdx, uint32_t batchIdx, uint32_t calcLen); + __aicore__ inline void Compute(uint32_t seqIdx, uint32_t batchIdx, LocalTensor &gatherOffsetCast, + uint32_t calcLen); + __aicore__ inline void ComputeCastFp32(uint32_t seqIdx, uint32_t batchIdx, LocalTensor &gatherOffsetCast, + uint32_t calcLen); +}; + +template +__aicore__ inline void InterleavedSplitBSPad::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RopeRegbaseTilingData *tiling, TPipe *pipe) +{ + InitData(tiling); + + blockIdx = GetBlockIdx(); + bufferSdSize = seqLen * headDim; + bufferNdSize = numHeads * headDim; + bufferNdSizeAll = numHeads * allHeadDim; + + if (blockIdx < frontCoreNum) { + ubCalcSeqLoop = coreCalcNum; + ioOffset = blockIdx * coreCalcNum * bufferNdSize; + ioOffsetAll = blockIdx * coreCalcNum * bufferNdSizeAll; + triOffset = blockIdx * coreCalcNum * headDim; + } else if (coreCalcTail != 0) { + ubCalcSeqLoop = coreCalcTail; + ioOffset = frontCoreNum * coreCalcNum * bufferNdSize + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSize; + ioOffsetAll = frontCoreNum * coreCalcNum * bufferNdSizeAll + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSizeAll; + triOffset = frontCoreNum * coreCalcNum * headDim + (blockIdx - frontCoreNum) * coreCalcTail * headDim; + } + + bufferBsndSize = batchSize * seqLen * bufferNdSizeAll; + xGm.SetGlobalBuffer((__gm__ T *)x + ioOffsetAll, bufferBsndSize); + yGm.SetGlobalBuffer((__gm__ T *)y + ioOffsetAll, bufferBsndSize); + cosGm.SetGlobalBuffer((__gm__ T *)cos + triOffset, bufferSdSize); + sinGm.SetGlobalBuffer((__gm__ T *)sin + triOffset, bufferSdSize); + + bufferLenSize = ubCalcBNum * numHeads * headDimAlign * sizeof(T); + pipe->InitBuffer(inQueX, BUFFER_NUM, bufferLenSize); + pipe->InitBuffer(inQueCos, BUFFER_NUM, bufferLenSize); + pipe->InitBuffer(outQueY, BUFFER_NUM, bufferLenSize); + + if constexpr (std::is_same::value || std::is_same::value) { + bufferLenSize = ubCalcBNum * numHeads * headDimAlign * sizeof(float); + pipe->InitBuffer(tmpFp32Buf1, bufferLenSize); + pipe->InitBuffer(tmpFp32Buf2, bufferLenSize); + pipe->InitBuffer(tmpFp32Buf3, bufferLenSize); + } + + gatherOffsetLenSize = numHeads * headDimAlign * sizeof(int32_t); + pipe->InitBuffer(gatherOffsetBuf, gatherOffsetLenSize); +} + +template +__aicore__ inline void InterleavedSplitBSPad::InitData(const RopeRegbaseTilingData *tiling) +{ + tiling_ = tiling; + batchSize = tiling_->batchSize; + seqLen = tiling_->seqLen; + numHeads = tiling_->numHeads; + headDim = tiling_->headDim; + frontCoreNum = tiling_->frontCoreNum; + tailCoreNum = tiling_->tailCoreNum; + coreCalcNum = tiling_->coreCalcNum; + coreCalcTail = tiling_->coreCalcTail; + ubCalcBNum = tiling_->ubCalcBNum; + ubCalcBLoop = tiling_->ubCalcBLoop; + ubCalcBTail = tiling_->ubCalcBTail; + allHeadDim = tiling_->allHeadDim; + start = tiling_->start; + + alignLen = (std::is_same::value || std::is_same::value) ? ALIGN_16 : ALIGN_32; + headDimAlign = (headDim + alignLen - 1) / alignLen * alignLen; +} + +template +__aicore__ inline void InterleavedSplitBSPad::CopyInX(LocalTensor &x, uint32_t seqIdx, uint32_t batchIdx, + uint32_t calcLen) +{ + DataCopyExtParams dataCopyParams; + uint64_t startOffset = batchIdx * ubCalcBNum * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll; + + for (uint32_t loopIdx = 0; loopIdx < calcLen; ++loopIdx) { + dataCopyParams.blockCount = numHeads; + dataCopyParams.blockLen = headDim * sizeof(T); + dataCopyParams.srcStride = (allHeadDim - headDim) * sizeof(T); + dataCopyParams.dstStride = 0; + DataCopyPad(x[loopIdx * numHeads * headDimAlign], xGm[startOffset + loopIdx * seqLen * bufferNdSizeAll + start], + dataCopyParams, {false, 0, 0, 0}); + } + + event_t eventIdMTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMTE2ToV); + WaitFlag(eventIdMTE2ToV); +} + +template +__aicore__ inline void InterleavedSplitBSPad::CopyInCos(LocalTensor &cos, uint32_t seqIdx, uint32_t calcLen) +{ + DataCopyExtParams bsPadDataCopyTriParams; + bsPadDataCopyTriParams.blockCount = 1; + bsPadDataCopyTriParams.blockLen = headDim * sizeof(T); + bsPadDataCopyTriParams.srcStride = 0; + bsPadDataCopyTriParams.dstStride = 0; + DataCopyPad(cos, cosGm[seqIdx * headDim], bsPadDataCopyTriParams, {false, 0, 0, 0}); + event_t eventId2MTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId2MTE2ToV); + WaitFlag(eventId2MTE2ToV); + BroadCastTriToB1nd(cos, calcLen, numHeads, headDimAlign); +} + +template +__aicore__ inline void InterleavedSplitBSPad::CopyInSin(LocalTensor &sin, uint32_t seqIdx, uint32_t calcLen) +{ + event_t eventIdVToMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventIdVToMTE2); + WaitFlag(eventIdVToMTE2); + DataCopyExtParams bsPadDataCopyTriParams; + bsPadDataCopyTriParams.blockCount = 1; + bsPadDataCopyTriParams.blockLen = headDim * sizeof(T); + bsPadDataCopyTriParams.srcStride = 0; + bsPadDataCopyTriParams.dstStride = 0; + DataCopyPad(sin, sinGm[seqIdx * headDim], bsPadDataCopyTriParams, {false, 0, 0, 0}); + event_t eventId3MTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId3MTE2ToV); + WaitFlag(eventId3MTE2ToV); + BroadCastTriToB1nd(sin, calcLen, numHeads, headDimAlign); +} + +template +__aicore__ inline void InterleavedSplitBSPad::CopyOut(uint32_t seqIdx, uint32_t batchIdx, uint32_t calcLen) +{ + LocalTensor y = outQueY.DeQue(); + DataCopyExtParams dataCopyParams; + uint64_t startOffset = batchIdx * ubCalcBNum * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll; + for (uint32_t loopIdx = 0; loopIdx < calcLen; ++loopIdx) { + dataCopyParams.blockCount = numHeads; + dataCopyParams.blockLen = headDim * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride =(allHeadDim - headDim) * sizeof(T); + DataCopyPad(yGm[startOffset + loopIdx * seqLen * bufferNdSizeAll + start], y[loopIdx * numHeads * headDimAlign], + dataCopyParams); + } + outQueY.FreeTensor(y); +} + +template +__aicore__ inline void InterleavedSplitBSPad::Process() +{ + LocalTensor gatherOffset = gatherOffsetBuf.Get(); + SetGatherSrcOffset(gatherOffset, headDimAlign * numHeads, static_cast(sizeof(float))); + LocalTensor gatherOffsetCast = gatherOffset.ReinterpretCast(); + + if constexpr (std::is_same::value || std::is_same::value) { + for (uint32_t loop = 0; loop < ubCalcSeqLoop; ++loop) { + for (uint32_t j = 0; j < (ubCalcBTail == 0 ? ubCalcBLoop : ubCalcBLoop - 1); ++j) { + ComputeCastFp32(loop, j, gatherOffsetCast, ubCalcBNum); + CopyOut(loop, j, ubCalcBNum); + } + if (ubCalcBTail != 0) { + ComputeCastFp32(loop, ubCalcBLoop - 1, gatherOffsetCast, ubCalcBTail); + CopyOut(loop, ubCalcBLoop - 1, ubCalcBTail); + } + } + } else { + for (uint32_t loop = 0; loop < ubCalcSeqLoop; ++loop) { + for (uint32_t j = 0; j < (ubCalcBTail == 0 ? ubCalcBLoop : ubCalcBLoop - 1); ++j) { + Compute(loop, j, gatherOffsetCast, ubCalcBNum); + CopyOut(loop, j, ubCalcBNum); + } + if (ubCalcBTail != 0) { + Compute(loop, ubCalcBLoop - 1, gatherOffsetCast, ubCalcBTail); + CopyOut(loop, ubCalcBLoop - 1, ubCalcBTail); + } + } + } +} + +template +__aicore__ inline void InterleavedSplitBSPad::Compute(uint32_t seqIdx, uint32_t batchIdx, + LocalTensor &gatherOffsetCast, uint32_t calcLen) +{ + uint64_t calcTotalNum = calcLen * numHeads * headDimAlign; + + LocalTensor x = inQueX.AllocTensor(); + CopyInX(x, seqIdx, batchIdx, calcLen); + LocalTensor cos = inQueCos.AllocTensor(); + CopyInCos(cos, seqIdx, calcLen); + + LocalTensor yTensor = outQueY.AllocTensor(); + Mul(yTensor, x, cos, calcTotalNum); + for (uint32_t i = 0; i < calcLen; ++i) { + Gather(x[i * numHeads * headDimAlign], x[i * numHeads * headDimAlign], gatherOffsetCast, 0, + numHeads * headDimAlign); + } + + CopyInSin(cos, seqIdx, calcLen); + Mul(x, x, cos, calcTotalNum); + + inQueCos.FreeTensor(cos); + InterleavedInversion(x, calcTotalNum); + Add(yTensor, yTensor, x, calcTotalNum); + + inQueX.FreeTensor(x); + outQueY.EnQue(yTensor); +} + +template +__aicore__ inline void InterleavedSplitBSPad::ComputeCastFp32(uint32_t seqIdx, uint32_t batchIdx, + LocalTensor &gatherOffsetCast, + uint32_t calcLen) +{ + uint64_t calcTotalNum = calcLen * numHeads * headDimAlign; + + LocalTensor xTensor = inQueX.AllocTensor(); + CopyInX(xTensor, seqIdx, batchIdx, calcLen); + LocalTensor tmp32BsPadBuf1 = tmpFp32Buf1.Get(); + Cast(tmp32BsPadBuf1, xTensor, RoundMode::CAST_NONE, calcTotalNum); + inQueX.FreeTensor(xTensor); + + LocalTensor cos = inQueCos.AllocTensor(); + CopyInCos(cos, seqIdx, calcLen); + LocalTensor tmp32Buf2 = tmpFp32Buf2.Get(); + Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum); + + LocalTensor tmp32Buf3 = tmpFp32Buf3.Get(); + Mul(tmp32Buf3, tmp32BsPadBuf1, tmp32Buf2, calcTotalNum); + + for (uint32_t i = 0; i < calcLen; ++i) { + Gather(tmp32BsPadBuf1[i * numHeads * headDimAlign], tmp32BsPadBuf1[i * numHeads * headDimAlign], + gatherOffsetCast, 0, numHeads * headDimAlign); + } + + CopyInSin(cos, seqIdx, calcLen); + Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum); + inQueCos.FreeTensor(cos); + + Mul(tmp32BsPadBuf1, tmp32BsPadBuf1, tmp32Buf2, calcTotalNum); + InterleavedInversion(tmp32BsPadBuf1, calcTotalNum); + Add(tmp32Buf3, tmp32Buf3, tmp32BsPadBuf1, calcTotalNum); + + LocalTensor y = outQueY.AllocTensor(); + Cast(y, tmp32Buf3, RoundMode::CAST_RINT, calcTotalNum); + outQueY.EnQue(y); +} + +} // namespace RotateInterleavedN + +#endif // ROTATE_INTERLEAVED_SPLIT_BS_PAD_H diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bsn.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bsn.h new file mode 100644 index 000000000..d8f2dfb22 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bsn.h @@ -0,0 +1,314 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotate_interleaved_split_bsn.h + * \brief + */ +#ifndef ROTATE_INTERLEAVED_SPLIT_BSN_H +#define ROTATE_INTERLEAVED_SPLIT_BSN_H +#include "rotate_interleaved_common.h" + +namespace RotateInterleavedN { +using namespace AscendC; + +template +class InterleavedSplitBSN { +public: + __aicore__ inline InterleavedSplitBSN(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RopeRegbaseTilingData *tiling, TPipe *pipe); + __aicore__ inline void Process(); + +protected: + GlobalTensor xGm; + GlobalTensor cosGm; + GlobalTensor sinGm; + GlobalTensor yGm; + TQue inQueX; + TQue inQueCos; + TQue outQueY; + TBuf tmpFp32Buf1; + TBuf tmpFp32Buf2; + TBuf tmpFp32Buf3; + TBuf gatherOffsetBuf; + const RopeRegbaseTilingData* tiling_; + + // tilingdata + uint64_t batchSize; + uint64_t seqLen; + uint64_t numHeads; + uint64_t headDim; + uint64_t frontCoreNum; + uint64_t tailCoreNum; + uint64_t coreCalcNum; + uint64_t coreCalcTail; + uint64_t ubCalcNNum; + uint64_t ubCalcNLoop; + uint64_t ubCalcNTail; + uint64_t allHeadDim; + uint64_t start; + uint64_t ioOffsetAll; + uint64_t bufferNdSizeAll; + + // init tmp data + uint32_t blockIdx; + uint32_t ubCalcSeqLoop; + uint64_t ioOffset; + uint64_t triOffset; + uint64_t bufferBsndSize; + uint64_t bufferSdSize; + uint64_t bufferNdSize; + uint64_t bufferLenSize; + uint64_t gatherOffsetLenSize; + + __aicore__ inline void InitData(const RopeRegbaseTilingData *tiling); + __aicore__ inline void CopyInX(LocalTensor &x, uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, + uint32_t calcLen); + __aicore__ inline void CopyInCos(LocalTensor &cos, uint32_t seqIdx, uint32_t calcLen); + __aicore__ inline void CopyInSin(LocalTensor &sin, uint32_t seqIdx, uint32_t calcLen); + __aicore__ inline void CopyOut(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, uint32_t calcLen); + __aicore__ inline void Compute(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, + LocalTensor &gatherOffsetCast, uint32_t calcLen); + __aicore__ inline void ComputeCastFp32(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, + LocalTensor &gatherOffsetCast, uint32_t calcLen); +}; + +template +__aicore__ inline void InterleavedSplitBSN::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RopeRegbaseTilingData *tiling, TPipe *pipe) +{ + InitData(tiling); + + blockIdx = GetBlockIdx(); + bufferSdSize = seqLen * headDim; + bufferNdSize = numHeads * headDim; + bufferNdSizeAll = numHeads * allHeadDim; + + if (blockIdx < frontCoreNum) { + ubCalcSeqLoop = coreCalcNum; + ioOffset = blockIdx * coreCalcNum * bufferNdSize; + ioOffsetAll = blockIdx * coreCalcNum * bufferNdSizeAll; + triOffset = blockIdx * coreCalcNum * headDim; + } else if (coreCalcTail != 0) { + ubCalcSeqLoop = coreCalcTail; + ioOffset = frontCoreNum * coreCalcNum * bufferNdSize + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSize; + ioOffsetAll = frontCoreNum * coreCalcNum * bufferNdSizeAll + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSizeAll; + triOffset = frontCoreNum * coreCalcNum * headDim + (blockIdx - frontCoreNum) * coreCalcTail * headDim; + } + + bufferBsndSize = batchSize * seqLen * bufferNdSizeAll; + xGm.SetGlobalBuffer((__gm__ T *)x + ioOffsetAll, bufferBsndSize); + yGm.SetGlobalBuffer((__gm__ T *)y + ioOffsetAll, bufferBsndSize); + cosGm.SetGlobalBuffer((__gm__ T *)cos + triOffset, bufferSdSize); + sinGm.SetGlobalBuffer((__gm__ T *)sin + triOffset, bufferSdSize); + + bufferLenSize = ubCalcNNum * headDim * sizeof(T); + pipe->InitBuffer(inQueX, BUFFER_NUM, bufferLenSize); + pipe->InitBuffer(inQueCos, BUFFER_NUM, bufferLenSize); + pipe->InitBuffer(outQueY, BUFFER_NUM, bufferLenSize); + + if constexpr (std::is_same::value || std::is_same::value) { + bufferLenSize = ubCalcNNum * headDim * sizeof(float); + pipe->InitBuffer(tmpFp32Buf1, bufferLenSize); + pipe->InitBuffer(tmpFp32Buf2, bufferLenSize); + pipe->InitBuffer(tmpFp32Buf3, bufferLenSize); + } + + gatherOffsetLenSize = ubCalcNNum * headDim * sizeof(int32_t); + pipe->InitBuffer(gatherOffsetBuf, gatherOffsetLenSize); +} + +template +__aicore__ inline void InterleavedSplitBSN::InitData(const RopeRegbaseTilingData *tiling) +{ + tiling_ = tiling; + batchSize = tiling_->batchSize; + seqLen = tiling_->seqLen; + numHeads = tiling_->numHeads; + headDim = tiling_->headDim; + frontCoreNum = tiling_->frontCoreNum; + tailCoreNum = tiling_->tailCoreNum; + coreCalcNum = tiling_->coreCalcNum; + coreCalcTail = tiling_->coreCalcTail; + ubCalcNNum = tiling_->ubCalcNNum; + ubCalcNLoop = tiling_->ubCalcNLoop; + ubCalcNTail = tiling_->ubCalcNTail; + allHeadDim = tiling_->allHeadDim; + start = tiling_->start; +} + +template +__aicore__ inline void InterleavedSplitBSN::CopyInX(LocalTensor &x, uint32_t batchIdx, uint32_t seqIdx, + uint32_t numHeadsIdx, uint32_t calcLen) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = calcLen; + dataCopyParams.blockLen = headDim * sizeof(T); + dataCopyParams.srcStride = (allHeadDim - headDim)* sizeof(T); + dataCopyParams.dstStride = 0; + DataCopyPad(x, xGm[batchIdx * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll + numHeadsIdx * ubCalcNNum * allHeadDim + start], + dataCopyParams, {false, 0, 0, 0}); + event_t eventIdMTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMTE2ToV); + WaitFlag(eventIdMTE2ToV); +} + +template +__aicore__ inline void InterleavedSplitBSN::CopyInCos(LocalTensor &cos, uint32_t seqIdx, uint32_t calcLen) +{ + DataCopyExtParams bsnDataCopyTriParams; + bsnDataCopyTriParams.blockCount = 1; + bsnDataCopyTriParams.blockLen = headDim * sizeof(T); + bsnDataCopyTriParams.srcStride = 0; + bsnDataCopyTriParams.dstStride = 0; + DataCopyPad(cos, cosGm[seqIdx * headDim], bsnDataCopyTriParams, {false, 0, 0, 0}); + event_t eventId2MTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId2MTE2ToV); + WaitFlag(eventId2MTE2ToV); + BroadCastTriToB1nd(cos, 1, calcLen, headDim); +} + +template +__aicore__ inline void InterleavedSplitBSN::CopyInSin(LocalTensor &sin, uint32_t seqIdx, uint32_t calcLen) +{ + event_t eventIdVToMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventIdVToMTE2); + WaitFlag(eventIdVToMTE2); + + DataCopyExtParams bsnDataCopyTriParams; + bsnDataCopyTriParams.blockCount = 1; + bsnDataCopyTriParams.blockLen = headDim * sizeof(T); + bsnDataCopyTriParams.srcStride = 0; + bsnDataCopyTriParams.dstStride = 0; + DataCopyPad(sin, sinGm[seqIdx * headDim], bsnDataCopyTriParams, {false, 0, 0, 0}); + event_t eventId3MTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId3MTE2ToV); + WaitFlag(eventId3MTE2ToV); + BroadCastTriToB1nd(sin, 1, calcLen, headDim); +} + +template +__aicore__ inline void InterleavedSplitBSN::CopyOut(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, + uint32_t calcLen) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = calcLen; + dataCopyParams.blockLen = headDim * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (allHeadDim - headDim)* sizeof(T); + LocalTensor y = outQueY.DeQue(); + DataCopyPad(yGm[batchIdx * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll + numHeadsIdx * ubCalcNNum * allHeadDim + start], y, + dataCopyParams); + outQueY.FreeTensor(y); +} + +template +__aicore__ inline void InterleavedSplitBSN::Process() +{ + LocalTensor gatherOffset = gatherOffsetBuf.Get(); + SetGatherSrcOffset(gatherOffset, ubCalcNNum * headDim, static_cast(sizeof(float))); + LocalTensor gatherOffsetCast = gatherOffset.ReinterpretCast(); + + if constexpr (std::is_same::value || std::is_same::value) { + for (uint32_t i = 0; i < batchSize; ++i) { + for (uint32_t j = 0; j < ubCalcSeqLoop; ++j) { + for (uint32_t z = 0; z < (ubCalcNTail == 0 ? ubCalcNLoop : ubCalcNLoop - 1); ++z) { + ComputeCastFp32(i, j, z, gatherOffsetCast, ubCalcNNum); + CopyOut(i, j, z, ubCalcNNum); + } + if (ubCalcNTail != 0) { + ComputeCastFp32(i, j, ubCalcNLoop - 1, gatherOffsetCast, ubCalcNTail); + CopyOut(i, j, ubCalcNLoop - 1, ubCalcNTail); + } + } + } + } else { + for (uint32_t i = 0; i < batchSize; ++i) { + for (uint32_t j = 0; j < ubCalcSeqLoop; ++j) { + for (uint32_t z = 0; z < (ubCalcNTail == 0 ? ubCalcNLoop : ubCalcNLoop - 1); ++z) { + Compute(i, j, z, gatherOffsetCast, ubCalcNNum); + CopyOut(i, j, z, ubCalcNNum); + } + if (ubCalcNTail != 0) { + Compute(i, j, ubCalcNLoop - 1, gatherOffsetCast, ubCalcNTail); + CopyOut(i, j, ubCalcNLoop - 1, ubCalcNTail); + } + } + } + } +} + +template +__aicore__ inline void InterleavedSplitBSN::Compute(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, + LocalTensor &gatherOffsetCast, uint32_t calcLen) +{ + uint64_t calcTotalNum = calcLen * headDim; + + LocalTensor x = inQueX.AllocTensor(); + CopyInX(x, batchIdx, seqIdx, numHeadsIdx, calcLen); + + LocalTensor cos = inQueCos.AllocTensor(); + CopyInCos(cos, seqIdx, calcLen); + + LocalTensor y = outQueY.AllocTensor(); + Mul(y, x, cos, calcTotalNum); + + Gather(x, x, gatherOffsetCast, 0, calcTotalNum); + + CopyInSin(cos, seqIdx, calcLen); + + Mul(x, x, cos, calcTotalNum); + inQueCos.FreeTensor(cos); + InterleavedInversion(x, calcTotalNum); + Add(y, y, x, calcTotalNum); + inQueX.FreeTensor(x); + outQueY.EnQue(y); +} + +template +__aicore__ inline void InterleavedSplitBSN::ComputeCastFp32(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, + LocalTensor &gatherOffsetCast, + uint32_t calcLen) +{ + uint64_t calcTotalNum = calcLen * headDim; + + LocalTensor x = inQueX.AllocTensor(); + CopyInX(x, batchIdx, seqIdx, numHeadsIdx, calcLen); + LocalTensor tmp32BSNBuf1 = tmpFp32Buf1.Get(); + Cast(tmp32BSNBuf1, x, RoundMode::CAST_NONE, calcTotalNum); + inQueX.FreeTensor(x); + + LocalTensor cos = inQueCos.AllocTensor(); + CopyInCos(cos, seqIdx, calcLen); + LocalTensor tmp32Buf2 = tmpFp32Buf2.Get(); + Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum); + + LocalTensor tmp32Buf3 = tmpFp32Buf3.Get(); + Mul(tmp32Buf3, tmp32BSNBuf1, tmp32Buf2, calcTotalNum); + + Gather(tmp32BSNBuf1, tmp32BSNBuf1, gatherOffsetCast, 0, calcTotalNum); + + CopyInSin(cos, seqIdx, calcLen); + Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum); + inQueCos.FreeTensor(cos); + + Mul(tmp32BSNBuf1, tmp32BSNBuf1, tmp32Buf2, calcTotalNum); + InterleavedInversion(tmp32BSNBuf1, calcTotalNum); + Add(tmp32Buf3, tmp32Buf3, tmp32BSNBuf1, calcTotalNum); + + LocalTensor y = outQueY.AllocTensor(); + Cast(y, tmp32Buf3, RoundMode::CAST_RINT, calcTotalNum); + outQueY.EnQue(y); +} + +} // namespace RotateInterleavedN + +#endif // ROTATE_INTERLEAVED_SPLIT_BSN_H diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bsn_pad.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bsn_pad.h new file mode 100644 index 000000000..5942b7242 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_bsn_pad.h @@ -0,0 +1,321 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotate_interleaved_split_bsn_pad.h + * \brief + */ +#ifndef ROTATE_INTERLEAVED_SPLIT_BSN_PAD_H +#define ROTATE_INTERLEAVED_SPLIT_BSN_PAD_H +#include "rotate_interleaved_common.h" + +namespace RotateInterleavedN { +using namespace AscendC; + +template +class InterleavedSplitBSNPad { +public: + __aicore__ inline InterleavedSplitBSNPad(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RopeRegbaseTilingData *tiling, TPipe *pipe); + __aicore__ inline void Process(); + +protected: + GlobalTensor xGm; + GlobalTensor cosGm; + GlobalTensor sinGm; + GlobalTensor yGm; + TQue inQueX; + TQue inQueCos; + TQue outQueY; + TBuf tmpFp32Buf1; + TBuf tmpFp32Buf2; + TBuf tmpFp32Buf3; + TBuf gatherOffsetBuf; + const RopeRegbaseTilingData* tiling_; + + // tilingdata + uint64_t batchSize; + uint64_t seqLen; + uint64_t numHeads; + uint64_t headDim; + uint64_t frontCoreNum; + uint64_t tailCoreNum; + uint64_t coreCalcNum; + uint64_t coreCalcTail; + uint64_t ubCalcNNum; + uint64_t ubCalcNLoop; + uint64_t ubCalcNTail; + uint64_t allHeadDim; + uint64_t start; + uint64_t ioOffsetAll; + uint64_t bufferNdSizeAll; + + // init tmp data + uint32_t alignLen; + uint32_t headDimAlign; + uint32_t blockIdx; + uint32_t ubCalcSeqLoop; + uint64_t ioOffset; + uint64_t triOffset; + uint64_t bufferBsndSize; + uint64_t bufferSdSize; + uint64_t bufferNdSize; + uint64_t bufferLenSize; + uint64_t gatherOffsetLenSize; + + __aicore__ inline void InitData(const RopeRegbaseTilingData *tiling); + __aicore__ inline void CopyInX(LocalTensor &x, uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, + uint32_t calcLen); + __aicore__ inline void CopyInCos(LocalTensor &cos, uint32_t seqIdx, uint32_t calcLen); + __aicore__ inline void CopyInSin(LocalTensor &sin, uint32_t seqIdx, uint32_t calcLen); + __aicore__ inline void CopyOut(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, uint32_t calcLen); + __aicore__ inline void Compute(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, + LocalTensor &gatherOffsetCast, uint32_t calcLen); + __aicore__ inline void ComputeCastFp32(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, + LocalTensor &gatherOffsetCast, uint32_t calcLen); +}; + +template +__aicore__ inline void InterleavedSplitBSNPad::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RopeRegbaseTilingData *tiling, TPipe *pipe) +{ + InitData(tiling); + + blockIdx = GetBlockIdx(); + bufferSdSize = seqLen * headDim; + bufferNdSize = numHeads * headDim; + bufferNdSizeAll = numHeads * allHeadDim; + + if (blockIdx < frontCoreNum) { + ubCalcSeqLoop = coreCalcNum; + ioOffset = blockIdx * coreCalcNum * bufferNdSize; + ioOffsetAll = blockIdx * coreCalcNum * bufferNdSizeAll; + triOffset = blockIdx * coreCalcNum * headDim; + } else if (coreCalcTail != 0) { + ubCalcSeqLoop = coreCalcTail; + ioOffset = frontCoreNum * coreCalcNum * bufferNdSize + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSize; + ioOffsetAll = frontCoreNum * coreCalcNum * bufferNdSizeAll + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSizeAll; + triOffset = frontCoreNum * coreCalcNum * headDim + (blockIdx - frontCoreNum) * coreCalcTail * headDim; + } + + bufferBsndSize = batchSize * seqLen * bufferNdSizeAll; + xGm.SetGlobalBuffer((__gm__ T *)x + ioOffsetAll, bufferBsndSize); + yGm.SetGlobalBuffer((__gm__ T *)y + ioOffsetAll, bufferBsndSize); + cosGm.SetGlobalBuffer((__gm__ T *)cos + triOffset, bufferSdSize); + sinGm.SetGlobalBuffer((__gm__ T *)sin + triOffset, bufferSdSize); + + bufferLenSize = ubCalcNNum * headDimAlign * sizeof(T); + pipe->InitBuffer(inQueX, BUFFER_NUM, bufferLenSize); + pipe->InitBuffer(inQueCos, BUFFER_NUM, bufferLenSize); + pipe->InitBuffer(outQueY, BUFFER_NUM, bufferLenSize); + + if constexpr (std::is_same::value || std::is_same::value) { + bufferLenSize = ubCalcNNum * headDimAlign * sizeof(float); + pipe->InitBuffer(tmpFp32Buf1, bufferLenSize); + pipe->InitBuffer(tmpFp32Buf2, bufferLenSize); + pipe->InitBuffer(tmpFp32Buf3, bufferLenSize); + } + + gatherOffsetLenSize = ubCalcNNum * headDimAlign * sizeof(int32_t); + pipe->InitBuffer(gatherOffsetBuf, gatherOffsetLenSize); +} + +template +__aicore__ inline void InterleavedSplitBSNPad::InitData(const RopeRegbaseTilingData *tiling) +{ + tiling_ = tiling; + batchSize = tiling_->batchSize; + seqLen = tiling_->seqLen; + numHeads = tiling_->numHeads; + headDim = tiling_->headDim; + frontCoreNum = tiling_->frontCoreNum; + tailCoreNum = tiling_->tailCoreNum; + coreCalcNum = tiling_->coreCalcNum; + coreCalcTail = tiling_->coreCalcTail; + ubCalcNNum = tiling_->ubCalcNNum; + ubCalcNLoop = tiling_->ubCalcNLoop; + ubCalcNTail = tiling_->ubCalcNTail; + allHeadDim = tiling_->allHeadDim; + start = tiling_->start; + + alignLen = (std::is_same::value || std::is_same::value) ? ALIGN_16 : ALIGN_32; + headDimAlign = (headDim + alignLen - 1) / alignLen * alignLen; +} + +template +__aicore__ inline void InterleavedSplitBSNPad::CopyInX(LocalTensor &xTensor, uint32_t batchIdx, uint32_t seqIdx, + uint32_t numHeadsIdx, uint32_t calcLen) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = calcLen; + dataCopyParams.blockLen = headDim * sizeof(T); + dataCopyParams.srcStride = (allHeadDim - headDim) * sizeof(T); + dataCopyParams.dstStride = 0; + DataCopyPad(xTensor, + xGm[batchIdx * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll + numHeadsIdx * ubCalcNNum * allHeadDim + start], + dataCopyParams, {false, 0, 0, 0}); + event_t eventIdMTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMTE2ToV); + WaitFlag(eventIdMTE2ToV); +} + +template +__aicore__ inline void InterleavedSplitBSNPad::CopyInCos(LocalTensor &cos, uint32_t seqIdx, uint32_t calcLen) +{ + DataCopyExtParams bsnPadDataCopyTriParams; + bsnPadDataCopyTriParams.blockCount = 1; + bsnPadDataCopyTriParams.blockLen = headDim * sizeof(T); + bsnPadDataCopyTriParams.srcStride = 0; + bsnPadDataCopyTriParams.dstStride = 0; + DataCopyPad(cos, cosGm[seqIdx * headDim], bsnPadDataCopyTriParams, {false, 0, 0, 0}); + event_t eventId2MTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId2MTE2ToV); + WaitFlag(eventId2MTE2ToV); + BroadCastTriToB1nd(cos, 1, calcLen, headDimAlign); +} + +template +__aicore__ inline void InterleavedSplitBSNPad::CopyInSin(LocalTensor &sinTensor, uint32_t seqIdx, + uint32_t calcLen) +{ + event_t eventIdVToMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventIdVToMTE2); + WaitFlag(eventIdVToMTE2); + + DataCopyExtParams bsnPadDataCopyTriParams; + bsnPadDataCopyTriParams.blockCount = 1; + bsnPadDataCopyTriParams.blockLen = headDim * sizeof(T); + bsnPadDataCopyTriParams.srcStride = 0; + bsnPadDataCopyTriParams.dstStride = 0; + DataCopyPad(sinTensor, sinGm[seqIdx * headDim], bsnPadDataCopyTriParams, {false, 0, 0, 0}); + event_t eventId3MTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId3MTE2ToV); + WaitFlag(eventId3MTE2ToV); + BroadCastTriToB1nd(sinTensor, 1, calcLen, headDimAlign); +} + +template +__aicore__ inline void InterleavedSplitBSNPad::CopyOut(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, + uint32_t calcLen) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = calcLen; + dataCopyParams.blockLen = headDim * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (allHeadDim - headDim) * sizeof(T); + LocalTensor y = outQueY.DeQue(); + DataCopyPad(yGm[batchIdx * seqLen * bufferNdSizeAll + seqIdx * bufferNdSizeAll + numHeadsIdx * ubCalcNNum * allHeadDim + start], y, + dataCopyParams); + outQueY.FreeTensor(y); +} + +template +__aicore__ inline void InterleavedSplitBSNPad::Process() +{ + LocalTensor gatherOffset = gatherOffsetBuf.Get(); + SetGatherSrcOffset(gatherOffset, ubCalcNNum * headDimAlign, static_cast(sizeof(float))); + LocalTensor gatherOffsetCast = gatherOffset.ReinterpretCast(); + + if constexpr (std::is_same::value || std::is_same::value) { + for (uint32_t batch = 0; batch < batchSize; ++batch) { + for (uint32_t j = 0; j < ubCalcSeqLoop; ++j) { + for (uint32_t z = 0; z < (ubCalcNTail == 0 ? ubCalcNLoop : ubCalcNLoop - 1); ++z) { + ComputeCastFp32(batch, j, z, gatherOffsetCast, ubCalcNNum); + CopyOut(batch, j, z, ubCalcNNum); + } + if (ubCalcNTail != 0) { + ComputeCastFp32(batch, j, ubCalcNLoop - 1, gatherOffsetCast, ubCalcNTail); + CopyOut(batch, j, ubCalcNLoop - 1, ubCalcNTail); + } + } + } + } else { + for (uint32_t batch = 0; batch < batchSize; ++batch) { + for (uint32_t j = 0; j < ubCalcSeqLoop; ++j) { + for (uint32_t z = 0; z < (ubCalcNTail == 0 ? ubCalcNLoop : ubCalcNLoop - 1); ++z) { + Compute(batch, j, z, gatherOffsetCast, ubCalcNNum); + CopyOut(batch, j, z, ubCalcNNum); + } + if (ubCalcNTail != 0) { + Compute(batch, j, ubCalcNLoop - 1, gatherOffsetCast, ubCalcNTail); + CopyOut(batch, j, ubCalcNLoop - 1, ubCalcNTail); + } + } + } + } +} + +template +__aicore__ inline void InterleavedSplitBSNPad::Compute(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, + LocalTensor &gatherOffsetCast, uint32_t calcLen) +{ + uint64_t calcTotalNum = calcLen * headDimAlign; + + LocalTensor xTensor = inQueX.AllocTensor(); + CopyInX(xTensor, batchIdx, seqIdx, numHeadsIdx, calcLen); + + LocalTensor cos = inQueCos.AllocTensor(); + CopyInCos(cos, seqIdx, calcLen); + + LocalTensor yTensor = outQueY.AllocTensor(); + Mul(yTensor, xTensor, cos, calcTotalNum); + + Gather(xTensor, xTensor, gatherOffsetCast, 0, calcTotalNum); + + CopyInSin(cos, seqIdx, calcLen); + + Mul(xTensor, xTensor, cos, calcTotalNum); + inQueCos.FreeTensor(cos); + InterleavedInversion(xTensor, calcTotalNum); + Add(yTensor, yTensor, xTensor, calcTotalNum); + inQueX.FreeTensor(xTensor); + outQueY.EnQue(yTensor); +} + +template +__aicore__ inline void +InterleavedSplitBSNPad::ComputeCastFp32(uint32_t batchIdx, uint32_t seqIdx, uint32_t numHeadsIdx, + LocalTensor &gatherOffsetCast, uint32_t calcLen) +{ + uint64_t totalCount = calcLen * headDimAlign; + + LocalTensor xTensor = inQueX.AllocTensor(); + CopyInX(xTensor, batchIdx, seqIdx, numHeadsIdx, calcLen); + LocalTensor tmp32Buf1 = tmpFp32Buf1.Get(); + Cast(tmp32Buf1, xTensor, RoundMode::CAST_NONE, totalCount); + inQueX.FreeTensor(xTensor); + + LocalTensor cos = inQueCos.AllocTensor(); + CopyInCos(cos, seqIdx, calcLen); + LocalTensor tmp32Buf2 = tmpFp32Buf2.Get(); + Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, totalCount); + + LocalTensor tmp32Buf3 = tmpFp32Buf3.Get(); + Mul(tmp32Buf3, tmp32Buf1, tmp32Buf2, totalCount); + + Gather(tmp32Buf1, tmp32Buf1, gatherOffsetCast, 0, totalCount); + + CopyInSin(cos, seqIdx, calcLen); + Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, totalCount); + inQueCos.FreeTensor(cos); + + Mul(tmp32Buf1, tmp32Buf1, tmp32Buf2, totalCount); + InterleavedInversion(tmp32Buf1, totalCount); + Add(tmp32Buf3, tmp32Buf3, tmp32Buf1, totalCount); + + LocalTensor y = outQueY.AllocTensor(); + Cast(y, tmp32Buf3, RoundMode::CAST_RINT, totalCount); + outQueY.EnQue(y); +} + +} // namespace RotateInterleavedN + +#endif // ROTATE_INTERLEAVED_SPLIT_BSN_PAD_H diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_s.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_s.h new file mode 100644 index 000000000..f9f24f508 --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_s.h @@ -0,0 +1,313 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotate_interleaved_split_s.h + * \brief + */ +#ifndef ROTATE_INTERLEAVED_SPLIT_S_H +#define ROTATE_INTERLEAVED_SPLIT_S_H +#include "rotate_interleaved_common.h" + +namespace RotateInterleavedN { +using namespace AscendC; + +template +class InterleavedSplitS { +public: + __aicore__ inline InterleavedSplitS(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RopeRegbaseTilingData *tiling, TPipe *pipe); + __aicore__ inline void Process(); + +protected: + GlobalTensor xGm; + GlobalTensor cosGm; + GlobalTensor sinGm; + GlobalTensor yGm; + TQue inQueX; + TQue inQueCos; + TQue outQueY; + TBuf tmpFp32Buf1; + TBuf tmpFp32Buf2; + TBuf tmpFp32Buf3; + TBuf gatherOffsetBuf; + const RopeRegbaseTilingData* tiling_; + + // tilingdata + uint64_t batchSize; + uint64_t seqLen; + uint64_t numHeads; + uint64_t headDim; + uint64_t frontCoreNum; + uint64_t tailCoreNum; + uint64_t coreCalcNum; + uint64_t coreCalcTail; + uint64_t ubCalcNum; + uint64_t ubCalcLoop; + uint64_t ubCalcTail; + uint64_t ubCalcTailNum; + uint64_t ubCalcTailLoop; + uint64_t ubCalcTailTail; + uint64_t allHeadDim; + uint64_t start; + + // init tmp data + uint32_t blockIdx; + uint32_t ubCalcSeq; + uint32_t ubCalcSeqTail; + uint32_t ubCalcSeqLoop; + uint64_t ioOffset; + uint64_t ioOffsetAll; + uint64_t triOffset; + uint64_t bufferBsndSize; + uint64_t bufferBsndSizeAll; + uint64_t bufferSdSize; + uint64_t bufferNdSize; + uint64_t bufferNdSizeAll; + uint64_t bufferLenSize; + uint64_t gatherOffsetLenSize; + + __aicore__ inline void InitData(const RopeRegbaseTilingData *tiling); + __aicore__ inline void CopyInX(LocalTensor &x, uint32_t loopIdx, uint32_t calcLen); + __aicore__ inline void CopyInCos(LocalTensor &cos, uint32_t loopIdx, uint32_t calcLen); + __aicore__ inline void CopyInSin(LocalTensor &sin, uint32_t loopIdx, uint32_t calcLen); + __aicore__ inline void CopyOut(uint32_t loopIdx, uint32_t calcLen); + __aicore__ inline void Compute(uint32_t loopIdx, LocalTensor &gatherOffsetCast, uint32_t calcLen); + __aicore__ inline void ComputeCastFp32(uint32_t loopIdx, LocalTensor &gatherOffsetCast, uint32_t calcLen); +}; + +template +__aicore__ inline void InterleavedSplitS::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RopeRegbaseTilingData *tiling, TPipe *pipe) +{ + InitData(tiling); + + blockIdx = GetBlockIdx(); + bufferSdSize = seqLen * headDim; + bufferNdSize = numHeads * headDim; + bufferNdSizeAll = numHeads * allHeadDim; + + if (blockIdx < frontCoreNum) { + ubCalcSeq = ubCalcNum; + ubCalcSeqTail = ubCalcTail; + ubCalcSeqLoop = ubCalcLoop; + ioOffset = blockIdx * coreCalcNum * bufferNdSize; + ioOffsetAll = blockIdx * coreCalcNum * bufferNdSizeAll; + triOffset = blockIdx * coreCalcNum * headDim; + } else if (coreCalcTail != 0) { + ubCalcSeq = ubCalcTailNum; + ubCalcSeqTail = ubCalcTailTail; + ubCalcSeqLoop = ubCalcTailLoop; + ioOffset = frontCoreNum * coreCalcNum * bufferNdSize + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSize; + ioOffsetAll = frontCoreNum * coreCalcNum * bufferNdSizeAll + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSizeAll; + triOffset = frontCoreNum * coreCalcNum * headDim + (blockIdx - frontCoreNum) * coreCalcTail * headDim; + } + + bufferBsndSize = batchSize * seqLen * bufferNdSize; + bufferBsndSizeAll = batchSize * seqLen * bufferNdSizeAll; + xGm.SetGlobalBuffer((__gm__ T *)x + ioOffsetAll, bufferBsndSizeAll); + yGm.SetGlobalBuffer((__gm__ T *)y + ioOffsetAll, bufferBsndSizeAll); + cosGm.SetGlobalBuffer((__gm__ T *)cos + triOffset, bufferSdSize); + sinGm.SetGlobalBuffer((__gm__ T *)sin + triOffset, bufferSdSize); + + bufferLenSize = batchSize * ubCalcSeq * bufferNdSize * sizeof(T); + pipe->InitBuffer(inQueX, BUFFER_NUM, bufferLenSize); + pipe->InitBuffer(inQueCos, BUFFER_NUM, bufferLenSize); + pipe->InitBuffer(outQueY, BUFFER_NUM, bufferLenSize); + + if constexpr (std::is_same::value || std::is_same::value) { + bufferLenSize = batchSize * ubCalcSeq * bufferNdSize * sizeof(float); + pipe->InitBuffer(tmpFp32Buf1, bufferLenSize); + pipe->InitBuffer(tmpFp32Buf2, bufferLenSize); + pipe->InitBuffer(tmpFp32Buf3, bufferLenSize); + } + + gatherOffsetLenSize = bufferNdSize * sizeof(int32_t); + pipe->InitBuffer(gatherOffsetBuf, gatherOffsetLenSize); +} + +template +__aicore__ inline void InterleavedSplitS::InitData(const RopeRegbaseTilingData *tiling) +{ + tiling_ = tiling; + batchSize = tiling_->batchSize; //B + seqLen = tiling_->seqLen; // S + numHeads = tiling_->numHeads; // N + headDim = tiling_->headDim; // D + frontCoreNum = tiling_->frontCoreNum; + tailCoreNum = tiling_->tailCoreNum; + coreCalcNum = tiling_->coreCalcNum; + coreCalcTail = tiling_->coreCalcTail; + ubCalcNum = tiling_->ubCalcNum; + ubCalcLoop = tiling_->ubCalcLoop; + ubCalcTail = tiling_->ubCalcTail; + ubCalcTailNum = tiling_->ubCalcTailNum; + ubCalcTailLoop = tiling_->ubCalcTailLoop; + ubCalcTailTail = tiling_->ubCalcTailTail; + allHeadDim = tiling_->allHeadDim; + start = tiling_->start; +} + +template +__aicore__ inline void InterleavedSplitS::CopyInX(LocalTensor &x, uint32_t loopIdx, uint32_t calcLen) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = calcLen*numHeads; + dataCopyParams.blockLen = headDim * sizeof(T); + dataCopyParams.srcStride = (allHeadDim - headDim)* sizeof(T); + dataCopyParams.dstStride = 0; + DataCopyPad(x, xGm[loopIdx * ubCalcSeq * bufferNdSizeAll + start], dataCopyParams, {false, 0, 0, 0}); + event_t eventIdMTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMTE2ToV); + WaitFlag(eventIdMTE2ToV); +} + +template +__aicore__ inline void InterleavedSplitS::CopyInCos(LocalTensor &cos, uint32_t loopIdx, uint32_t calcLen) +{ + DataCopyExtParams dataCopyTriParams; + dataCopyTriParams.blockCount = calcLen; + dataCopyTriParams.blockLen = headDim * sizeof(T); + dataCopyTriParams.srcStride = 0; + dataCopyTriParams.dstStride = static_cast((numHeads - 1) * headDim * sizeof(T) / BLOCK_SIZE); + DataCopyPad(cos, cosGm[loopIdx * ubCalcSeq * headDim], dataCopyTriParams, {false, 0, 0, 0}); + event_t eventId2MTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId2MTE2ToV); + WaitFlag(eventId2MTE2ToV); + BroadCastTriToBsnd(cos, batchSize, calcLen, numHeads, headDim); +} + +template +__aicore__ inline void InterleavedSplitS::CopyInSin(LocalTensor &sin, uint32_t loopIdx, uint32_t calcLen) +{ + event_t eventIdVToMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventIdVToMTE2); + WaitFlag(eventIdVToMTE2); + DataCopyExtParams dataCopyTriParams; + dataCopyTriParams.blockCount = calcLen; + dataCopyTriParams.blockLen = headDim * sizeof(T); + dataCopyTriParams.srcStride = 0; + dataCopyTriParams.dstStride = static_cast((numHeads - 1) * headDim * sizeof(T) / BLOCK_SIZE); + DataCopyPad(sin, sinGm[loopIdx * ubCalcSeq * headDim], dataCopyTriParams, {false, 0, 0, 0}); + event_t eventId3MTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId3MTE2ToV); + WaitFlag(eventId3MTE2ToV); + BroadCastTriToBsnd(sin, batchSize, calcLen, numHeads, headDim); +} + +template +__aicore__ inline void InterleavedSplitS::CopyOut(uint32_t loopIdx, uint32_t calcLen) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = calcLen*numHeads; + dataCopyParams.blockLen = headDim * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (allHeadDim - headDim) * sizeof(T); + LocalTensor y = outQueY.DeQue(); + DataCopyPad(yGm[loopIdx * ubCalcSeq * bufferNdSizeAll + start], y, dataCopyParams); + outQueY.FreeTensor(y); +} + +template +__aicore__ inline void InterleavedSplitS::Process() +{ + LocalTensor gatherOffset = gatherOffsetBuf.Get(); + SetGatherSrcOffset(gatherOffset, headDim * numHeads, static_cast(sizeof(float))); + LocalTensor gatherOffsetCast = gatherOffset.ReinterpretCast(); + if constexpr (std::is_same::value || std::is_same::value) { + for (uint32_t i = 0; i < (ubCalcSeqTail == 0 ? ubCalcSeqLoop : ubCalcSeqLoop - 1); ++i) { + ComputeCastFp32(i, gatherOffsetCast, ubCalcSeq); + CopyOut(i, ubCalcSeq); + } + if (ubCalcSeqTail != 0) { + ComputeCastFp32(ubCalcSeqLoop - 1, gatherOffsetCast, ubCalcSeqTail); + CopyOut(ubCalcSeqLoop - 1, ubCalcSeqTail); + } + } else { + for (uint32_t i = 0; i < (ubCalcSeqTail == 0 ? ubCalcSeqLoop : ubCalcSeqLoop - 1); ++i) { + Compute(i, gatherOffsetCast, ubCalcSeq); + CopyOut(i, ubCalcSeq); + } + if (ubCalcSeqTail != 0) { + Compute(ubCalcSeqLoop - 1, gatherOffsetCast, ubCalcSeqTail); + CopyOut(ubCalcSeqLoop - 1, ubCalcSeqTail); + } + } +} + +template +__aicore__ inline void InterleavedSplitS::Compute(uint32_t loopIdx, LocalTensor &gatherOffsetCast, + uint32_t calcLen) +{ + uint64_t calcTotalNum = calcLen * batchSize * bufferNdSize; + + LocalTensor x = inQueX.AllocTensor(); + CopyInX(x, loopIdx, calcLen); + + LocalTensor cos = inQueCos.AllocTensor(); + CopyInCos(cos, loopIdx, calcLen); + + LocalTensor y = outQueY.AllocTensor(); + Mul(y, x, cos, calcTotalNum); + for (uint32_t i = 0; i < batchSize * calcLen; ++i) { + Gather(x[i * bufferNdSize], x[i * bufferNdSize], gatherOffsetCast, 0, bufferNdSize); + } + + CopyInSin(cos, loopIdx, calcLen); + + Mul(x, x, cos, calcTotalNum); + inQueCos.FreeTensor(cos); + InterleavedInversion(x, calcTotalNum); + Add(y, y, x, calcTotalNum); + inQueX.FreeTensor(x); + outQueY.EnQue(y); +} + +template +__aicore__ inline void InterleavedSplitS::ComputeCastFp32(uint32_t loopIdx, LocalTensor &gatherOffsetCast, + uint32_t calcLen) +{ + uint64_t calcTotalNum = calcLen * batchSize * bufferNdSize; + + LocalTensor x = inQueX.AllocTensor(); + CopyInX(x, loopIdx, calcLen); + LocalTensor tmp32Buf1 = tmpFp32Buf1.Get(); + Cast(tmp32Buf1, x, RoundMode::CAST_NONE, calcTotalNum); + inQueX.FreeTensor(x); + + LocalTensor cos = inQueCos.AllocTensor(); + CopyInCos(cos, loopIdx, calcLen); + LocalTensor tmp32Buf2 = tmpFp32Buf2.Get(); + Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum); + + LocalTensor tmp32Buf3 = tmpFp32Buf3.Get(); + Mul(tmp32Buf3, tmp32Buf1, tmp32Buf2, calcTotalNum); + + for (uint32_t i = 0; i < batchSize * calcLen; ++i) { + Gather(tmp32Buf1[i * bufferNdSize], tmp32Buf1[i * bufferNdSize], gatherOffsetCast, 0, bufferNdSize); + } + + CopyInSin(cos, loopIdx, calcLen); + Cast(tmp32Buf2, cos, RoundMode::CAST_NONE, calcTotalNum); + inQueCos.FreeTensor(cos); + + Mul(tmp32Buf1, tmp32Buf1, tmp32Buf2, calcTotalNum); + InterleavedInversion(tmp32Buf1, calcTotalNum); + Add(tmp32Buf3, tmp32Buf3, tmp32Buf1, calcTotalNum); + + LocalTensor yTensor = outQueY.AllocTensor(); + Cast(yTensor, tmp32Buf3, RoundMode::CAST_RINT, calcTotalNum); + outQueY.EnQue(yTensor); +} + +} // namespace RotateInterleavedN + +#endif // ROTATE_INTERLEAVED_SPLIT_S_H diff --git a/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_s_pad.h b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_s_pad.h new file mode 100644 index 000000000..be9db2b2f --- /dev/null +++ b/csrc/ascend/attention/inplace_partial_rotary_mul/op_kernel/rotate_interleaved_split_s_pad.h @@ -0,0 +1,328 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rotate_interleaved_split_s_pad.h + * \brief + */ +#ifndef ROTATE_INTERLEAVED_SPLIT_S_PAD_H +#define ROTATE_INTERLEAVED_SPLIT_S_PAD_H +#include "rotate_interleaved_common.h" + +namespace RotateInterleavedN { +using namespace AscendC; + +template +class InterleavedSplitSPad { +public: + __aicore__ inline InterleavedSplitSPad(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RopeRegbaseTilingData *tiling, TPipe *pipe); + __aicore__ inline void Process(); + +protected: + GlobalTensor xGm; + GlobalTensor cosGm; + GlobalTensor sinGm; + GlobalTensor yGm; + TQue inQueX; + TQue inQueCos; + TQue outQueY; + TBuf tmpFp32Buf1; + TBuf tmpFp32Buf2; + TBuf tmpFp32Buf3; + TBuf gatherOffsetBuf; + const RopeRegbaseTilingData* tiling_; + + // tilingdata + uint64_t batchSize; + uint64_t seqLen; + uint64_t numHeads; + uint64_t headDim; + uint64_t frontCoreNum; + uint64_t tailCoreNum; + uint64_t coreCalcNum; + uint64_t coreCalcTail; + uint64_t ubCalcNum; + uint64_t ubCalcLoop; + uint64_t ubCalcTail; + uint64_t ubCalcTailNum; + uint64_t ubCalcTailLoop; + uint64_t ubCalcTailTail; + uint64_t bufferNdSizeAll; + uint64_t allHeadDim; + uint64_t start; + + // init tmp data + uint32_t alignLen; + uint32_t headDimAlign; + uint32_t allHeadDimAlign; + uint32_t blockIdx; + uint32_t ubCalcSeq; + uint32_t ubCalcSeqTail; + uint32_t ubCalcSeqLoop; + uint64_t ioOffset; + uint64_t ioOffsetAll; + uint64_t triOffset; + uint64_t bufferBsndSize; + uint64_t bufferSdSize; + uint64_t bufferNdSize; + uint64_t bufferLenSize; + uint64_t gatherOffsetLenSize; + uint32_t blockNum = BLOCK_SIZE / sizeof(T); + + __aicore__ inline void InitData(const RopeRegbaseTilingData *tiling); + __aicore__ inline void CopyInX(LocalTensor &x, uint32_t loopIdx, uint32_t calcLen); + __aicore__ inline void CopyInCos(LocalTensor &cos, uint32_t loopIdx, uint32_t calcLen); + __aicore__ inline void CopyInSin(LocalTensor &sin, uint32_t loopIdx, uint32_t calcLen); + __aicore__ inline void CopyOut(uint32_t loopIdx, uint32_t calcLen); + __aicore__ inline void Compute(uint32_t loopIdx, LocalTensor &gatherOffsetCast, uint32_t calcLen); + __aicore__ inline void ComputeCastFp32(uint32_t loopIdx, LocalTensor &gatherOffsetCast, uint32_t calcLen); +}; + +template +__aicore__ inline void InterleavedSplitSPad::Init(GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR y, + const RopeRegbaseTilingData *tiling, TPipe *pipe) +{ + InitData(tiling); + + blockIdx = GetBlockIdx(); + bufferSdSize = seqLen * headDim; + bufferNdSize = numHeads * headDim; + bufferNdSizeAll = numHeads * allHeadDim; + + if (blockIdx < frontCoreNum) { + ubCalcSeq = ubCalcNum; + ubCalcSeqTail = ubCalcTail; + ubCalcSeqLoop = ubCalcLoop; + ioOffset = blockIdx * coreCalcNum * bufferNdSize; + ioOffsetAll = blockIdx * coreCalcNum * bufferNdSizeAll; + triOffset = blockIdx * coreCalcNum * headDim; + } else if (coreCalcTail != 0) { + ubCalcSeq = ubCalcTailNum; + ubCalcSeqTail = ubCalcTailTail; + ubCalcSeqLoop = ubCalcTailLoop; + ioOffset = frontCoreNum * coreCalcNum * bufferNdSize + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSize; + ioOffsetAll = frontCoreNum * coreCalcNum * bufferNdSizeAll + (blockIdx - frontCoreNum) * coreCalcTail * bufferNdSizeAll; + triOffset = frontCoreNum * coreCalcNum * headDim + (blockIdx - frontCoreNum) * coreCalcTail * headDim; + } + + bufferBsndSize = batchSize * seqLen * bufferNdSizeAll; + xGm.SetGlobalBuffer((__gm__ T *)x + ioOffsetAll, bufferBsndSize); + yGm.SetGlobalBuffer((__gm__ T *)y + ioOffsetAll, bufferBsndSize); + cosGm.SetGlobalBuffer((__gm__ T *)cos + triOffset, bufferSdSize); + sinGm.SetGlobalBuffer((__gm__ T *)sin + triOffset, bufferSdSize); + + bufferLenSize = batchSize * ubCalcSeq * numHeads * headDimAlign * sizeof(T); + pipe->InitBuffer(inQueX, BUFFER_NUM, bufferLenSize); + pipe->InitBuffer(inQueCos, BUFFER_NUM, bufferLenSize); + pipe->InitBuffer(outQueY, BUFFER_NUM, bufferLenSize); + + if constexpr (std::is_same::value || std::is_same::value) { + bufferLenSize = batchSize * ubCalcSeq * numHeads * headDimAlign * sizeof(float); + pipe->InitBuffer(tmpFp32Buf1, bufferLenSize); + pipe->InitBuffer(tmpFp32Buf2, bufferLenSize); + pipe->InitBuffer(tmpFp32Buf3, bufferLenSize); + } + + gatherOffsetLenSize = numHeads * headDimAlign * sizeof(int32_t); + pipe->InitBuffer(gatherOffsetBuf, gatherOffsetLenSize); +} + +template +__aicore__ inline void InterleavedSplitSPad::InitData(const RopeRegbaseTilingData *tiling) +{ + tiling_ = tiling; + batchSize = tiling_->batchSize; + seqLen = tiling_->seqLen; + numHeads = tiling_->numHeads; + headDim = tiling_->headDim; + frontCoreNum = tiling_->frontCoreNum; + tailCoreNum = tiling_->tailCoreNum; + coreCalcNum = tiling_->coreCalcNum; + coreCalcTail = tiling_->coreCalcTail; + ubCalcNum = tiling_->ubCalcNum; + ubCalcLoop = tiling_->ubCalcLoop; + ubCalcTail = tiling_->ubCalcTail; + ubCalcTailNum = tiling_->ubCalcTailNum; + ubCalcTailLoop = tiling_->ubCalcTailLoop; + ubCalcTailTail = tiling_->ubCalcTailTail; + allHeadDim = tiling_->allHeadDim; + start = tiling_->start; + + alignLen = (std::is_same::value || std::is_same::value) ? ALIGN_16 : ALIGN_32; + headDimAlign = (headDim + alignLen - 1) / alignLen * alignLen; + allHeadDimAlign = (allHeadDim + alignLen - 1) / alignLen * alignLen; +} + +template +__aicore__ inline void InterleavedSplitSPad::CopyInX(LocalTensor &x, uint32_t loopIdx, uint32_t calcLen) +{ + DataCopyExtParams dataCopyParams; + for (uint32_t batchIdx = 0; batchIdx < batchSize; ++batchIdx) { + dataCopyParams.blockCount = calcLen * numHeads; + dataCopyParams.blockLen = headDim * sizeof(T); + dataCopyParams.srcStride = (allHeadDim - headDim)* sizeof(T); + dataCopyParams.dstStride = 0; + DataCopyPad(x[batchIdx * calcLen * numHeads * headDimAlign], + xGm[batchIdx * seqLen * bufferNdSizeAll + loopIdx * ubCalcSeq * bufferNdSizeAll + start], dataCopyParams, + {false, 0, 0, 0}); + } + event_t eventIdMTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMTE2ToV); + WaitFlag(eventIdMTE2ToV); +} + +template +__aicore__ inline void InterleavedSplitSPad::CopyInCos(LocalTensor &cos, uint32_t loopIdx, uint32_t calcLen) +{ + DataCopyExtParams dataCopyTriParams; + dataCopyTriParams.blockCount = calcLen; + dataCopyTriParams.blockLen = headDim * sizeof(T); + dataCopyTriParams.srcStride = 0; + dataCopyTriParams.dstStride = static_cast((numHeads - 1) * headDimAlign / blockNum); + DataCopyPad(cos, cosGm[loopIdx * ubCalcSeq * headDim], dataCopyTriParams, {false, 0, 0, 0}); + event_t eventId2MTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId2MTE2ToV); + WaitFlag(eventId2MTE2ToV); + BroadCastTriToBsnd(cos, batchSize, calcLen, numHeads, headDimAlign); +} + +template +__aicore__ inline void InterleavedSplitSPad::CopyInSin(LocalTensor &sin, uint32_t loopIdx, uint32_t calcLen) +{ + event_t eventIdVToMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventIdVToMTE2); + WaitFlag(eventIdVToMTE2); + DataCopyExtParams dataCopyTriParams; + dataCopyTriParams.blockCount = calcLen; + dataCopyTriParams.blockLen = headDim * sizeof(T); + dataCopyTriParams.srcStride = 0; + dataCopyTriParams.dstStride = static_cast((numHeads - 1) * headDimAlign / blockNum); + DataCopyPad(sin, sinGm[loopIdx * ubCalcSeq * headDim], dataCopyTriParams, {false, 0, 0, 0}); + event_t eventId3MTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId3MTE2ToV); + WaitFlag(eventId3MTE2ToV); + BroadCastTriToBsnd(sin, batchSize, calcLen, numHeads, headDimAlign); +} + +template +__aicore__ inline void InterleavedSplitSPad::CopyOut(uint32_t loopIdx, uint32_t calcLen) +{ + LocalTensor y = outQueY.DeQue(); + DataCopyExtParams dataCopyParams; + for (uint32_t batchIdx = 0; batchIdx < batchSize; ++batchIdx) { + dataCopyParams.blockCount = calcLen * numHeads; + dataCopyParams.blockLen = headDim * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (allHeadDim - headDim)* sizeof(T); + DataCopyPad(yGm[batchIdx * seqLen * bufferNdSizeAll + loopIdx * ubCalcSeq * bufferNdSizeAll + start], + y[batchIdx * calcLen * numHeads * headDimAlign], dataCopyParams); + } + outQueY.FreeTensor(y); +} + +template +__aicore__ inline void InterleavedSplitSPad::Process() +{ + LocalTensor gatherOffset = gatherOffsetBuf.Get(); + SetGatherSrcOffset(gatherOffset, headDimAlign * numHeads, static_cast(sizeof(float))); + LocalTensor gatherOffsetCast = gatherOffset.ReinterpretCast(); + if constexpr (std::is_same::value || std::is_same::value) { + for (uint32_t loop = 0; loop < (ubCalcSeqTail == 0 ? ubCalcSeqLoop : ubCalcSeqLoop - 1); ++loop) { + ComputeCastFp32(loop, gatherOffsetCast, ubCalcSeq); + CopyOut(loop, ubCalcSeq); + } + if (ubCalcSeqTail != 0) { + ComputeCastFp32(ubCalcSeqLoop - 1, gatherOffsetCast, ubCalcSeqTail); + CopyOut(ubCalcSeqLoop - 1, ubCalcSeqTail); + } + } else { + for (uint32_t loop = 0; loop < (ubCalcSeqTail == 0 ? ubCalcSeqLoop : ubCalcSeqLoop - 1); ++loop) { + Compute(loop, gatherOffsetCast, ubCalcSeq); + CopyOut(loop, ubCalcSeq); + } + if (ubCalcSeqTail != 0) { + Compute(ubCalcSeqLoop - 1, gatherOffsetCast, ubCalcSeqTail); + CopyOut(ubCalcSeqLoop - 1, ubCalcSeqTail); + } + } +} + +template +__aicore__ inline void InterleavedSplitSPad::Compute(uint32_t loopIdx, LocalTensor &gatherOffsetCast, + uint32_t calcLen) +{ + uint64_t totalCount = calcLen * batchSize * numHeads * headDimAlign; + + LocalTensor x = inQueX.AllocTensor(); + CopyInX(x, loopIdx, calcLen); + + LocalTensor cos = inQueCos.AllocTensor(); + CopyInCos(cos, loopIdx, calcLen); + + LocalTensor y = outQueY.AllocTensor(); + Mul(y, x, cos, totalCount); + for (uint32_t i = 0; i < batchSize * calcLen; ++i) { + Gather(x[i * numHeads * headDimAlign], x[i * numHeads * headDimAlign], gatherOffsetCast, 0, + numHeads * headDimAlign); + } + + CopyInSin(cos, loopIdx, calcLen); + + Mul(x, x, cos, totalCount); + inQueCos.FreeTensor(cos); + InterleavedInversion(x, totalCount); + Add(y, y, x, totalCount); + inQueX.FreeTensor(x); + outQueY.EnQue(y); +} + +template +__aicore__ inline void +InterleavedSplitSPad::ComputeCastFp32(uint32_t loopIdx, LocalTensor &gatherOffsetCast, uint32_t calcLen) +{ + uint64_t totalCount = calcLen * batchSize * numHeads * headDimAlign; + + LocalTensor xTensor = inQueX.AllocTensor(); + CopyInX(xTensor, loopIdx, calcLen); + LocalTensor tmp32SPadBuf1 = tmpFp32Buf1.Get(); + Cast(tmp32SPadBuf1, xTensor, RoundMode::CAST_NONE, totalCount); + inQueX.FreeTensor(xTensor); + + LocalTensor cos = inQueCos.AllocTensor(); + CopyInCos(cos, loopIdx, calcLen); + LocalTensor tmp32SPadBuf2 = tmpFp32Buf2.Get(); + Cast(tmp32SPadBuf2, cos, RoundMode::CAST_NONE, totalCount); + + LocalTensor tmp32Buf3 = tmpFp32Buf3.Get(); + Mul(tmp32Buf3, tmp32SPadBuf1, tmp32SPadBuf2, totalCount); + + for (uint32_t i = 0; i < batchSize * calcLen; ++i) { + Gather(tmp32SPadBuf1[i * numHeads * headDimAlign], tmp32SPadBuf1[i * numHeads * headDimAlign], gatherOffsetCast, + 0, numHeads * headDimAlign); + } + + CopyInSin(cos, loopIdx, calcLen); + Cast(tmp32SPadBuf2, cos, RoundMode::CAST_NONE, totalCount); + inQueCos.FreeTensor(cos); + + Mul(tmp32SPadBuf1, tmp32SPadBuf1, tmp32SPadBuf2, totalCount); + InterleavedInversion(tmp32SPadBuf1, totalCount); + Add(tmp32Buf3, tmp32Buf3, tmp32SPadBuf1, totalCount); + + LocalTensor y = outQueY.AllocTensor(); + Cast(y, tmp32Buf3, RoundMode::CAST_RINT, totalCount); + outQueY.EnQue(y); +} + +} // namespace RotateInterleavedN + +#endif // ROTATE_INTERLEAVED_SPLIT_S_PAD_H diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/CMakeLists.txt b/csrc/ascend/attention/kv_quant_sparse_flash_attention/CMakeLists.txt new file mode 100644 index 000000000..86b308249 --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/README.md b/csrc/ascend/attention/kv_quant_sparse_flash_attention/README.md new file mode 100644 index 000000000..d18f9968b --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/README.md @@ -0,0 +1,222 @@ +# KvQuantSparseFlashAttention + +## 产品支持情况 + +|产品 | 是否支持 | +|:----------------------------|:-----------:| +|Ascend 950PR/Ascend 950DT| √ | +|Atlas A3 训练系列产品/Atlas A3 推理系列产品| √ | +|Atlas A2 训练系列产品/Atlas A2 推理系列产品| √ | +|Atlas 200I/500 A2 推理产品| × | +|Atlas 推理系列加速卡产品| × | +|Atlas 训练系列产品| × | + +## 功能说明 + +- API功能：`kv_quant_sparse_flash_attention`在`sparse_flash_attention`的基础上支持了[Per-Token-Head-Tile-128量化]输入。随着大模型上下文长度的增加，Sparse Attention的重要性与日俱增，这一技术通过“只计算关键部分”大幅减少计算量，然而会引入大量的离散访存，造成数据搬运时间增加，进而影响整体性能。 + +- 计算公式： + + $$ + Attention=\text{softmax}(\frac{Q @ \text{Dequant}({\tilde{K}^{INT8}},{Scale_K})^T}{\sqrt{d_k}})@\text{Dequant}(\tilde{V}^{INT8},{Scale_V}), + $$ + + 其中$\tilde{K},\tilde{V}$为基于某种选择算法（如`LightningIndexer`）得到的重要性较高的Key和Value，一般具有稀疏或分块稀疏的特征，$d_k$为$Q,\tilde{K}$每一个头的维度，$\text{Dequant}(\cdot,\cdot)$为反量化函数。 +本次公布的`kv_quant_sparse_flash_attention`是面向Sparse Attention的全新算子，针对离散访存进行了指令缩减及搬运聚合的细致优化。 + +## 参数说明 + +> **说明：**
+> 参数维度含义：B表示Batch Size、Q_S和KV_S分别表示query和key/value的Sequence Length、Q_N和KV_N分别表示query和key/value的Head Num、Q_D和KV_D分别表示query和key/value的Head Dim、Q_T和KV_T分别表示query和key/value的Total Tokens、sparse_size表示一次离散选取的block数、block_num和block_size分别表示PageAttention场景下的block总数和每个block的token数。 + ++ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数名输入/输出/属性描述数据类型数据格式
query输入attention结构的Q输入，不支持非连续。query由相同数据类型的q_nope和q_rope按D维度拼接得到。layout_query为"BSND"时shape为[B, Q_S, Q_N, Q_D]。layout_query为"TND"时shape为[Q_T, Q_N, Q_D]。其中Q_D值仅支持576，即q_nope+q_rope=512+64；Q_N值支持1/2/4/8/16/32/48/64/128。FLOAT16、BFLOAT16ND
key输入attention结构的K输入，不支持非连续。k_nope、query相同数据类型的k_rope和float32的量化参数按D维度拼接得到。layout_kv为"BSND"时shape为[B, KV_S, KV_N, KV_D]。layout_kv为"TND"时shape为[KV_T, KV_N, KV_D]。layout_kv为"PA_BSND"时shape为[block_num, block_size, KV_N, KV_D]，其中block_num为PageAttention时block总数，block_size为一个block的token数，block_size取值为16的整数倍，最大支持到1024。KV_N仅支持1；KV_D值仅支持656，即nope+rope*2+dequant_scale*4=512+64*2+4*4。FLOAT8_E4M3、INT8、HIFLOAT8ND
value输入attention结构的V输入，不支持非连续。FLOAT8_E4M3、INT8、HIFLOAT8ND
sparse_indices输入代表离散取kvCache的索引，不支持非连续。layout_query为"BSND"时shape为[B, Q_S, KV_N, sparse_size]。layout_query为"TND"时shape为[Q_T, KV_N, sparse_size]。其中sparse_size为一次离散选取的block数，需要保证每行有效值均在前半部分，无效值均在后半部分，且需要满足sparse_size大于0。当key和value的数据类型为hifloat8时，sparse_size仅支持2048。INT32ND
scale_value属性公式中dk开根号的倒数，代表缩放系数，作为query和key矩阵乘后Muls的scalar值。FLOAT-
key_quant_mode属性代表key的量化模式，仅支持传入2，代表per_tile量化模式。INT64-
value_quant_mode属性代表value的量化模式，仅支持传入2，代表per_tile量化模式。INT64-
key_dequant_scale输入预留参数。--
value_dequant_scale输入预留参数。--
block_table输入表示PageAttention中kvCache存储使用的block映射表。shape为[B, KV_S_max/block_size]，其中第一维长度为B，第二维长度不小于所有batch中最大的KV_S对应的block数量，即KV_S_max / block_size向上取整。INT32ND
actual_seq_lengths_query输入表示不同Batch中query的有效token数。如果不指定seqlen可传入None，表示和query shape的Q_S长度相同。shape为[B,]。每个Batch的有效token数不超过query中的Q_S大小且不小于0。当layout_query为"TND"时，该入参必须传入，且以该入参元素的数量作为B值，该入参中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。INT32ND
actual_seq_lengths_kv输入表示不同Batch中key和value的有效token数。如果不指定None，表示和key的shape的KV_S长度相同。shape为[B,]。每个Batch的有效token数不超过key/value中的KV_S大小且不小于0。当layout_kv为"TND"或"PA_BSND"时，该入参必须传入，layout_kv为"TND"时，该参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。INT32ND
sparse_block_size属性代表sparse阶段的block大小。sparse_block_size为1时，为Token-wise稀疏化场景；sparse_block_size大于1且小于等于128时，为Block-wise稀疏化场景，块内token共享相同的稀疏化决策。INT64-
layout_query属性用于标识输入query的数据排布格式，默认值"BSND"，支持传入BSND和TND。STRING-
layout_kv属性用于标识输入key的数据排布格式，默认值"BSND"，支持传入BSND、TND和PA_BSND，PA_BSND在开启PageAttention时使用。STRING-
sparse_mode属性表示sparse的模式。sparse_mode为0时，代表全部计算。sparse_mode为3时，代表rightDownCausal模式的mask，对应以右下顶点往左上为划分线的下三角场景。INT64-
pre_tokens属性用于稀疏计算，表示attention需要和前几个Token计算关联，仅支持2^63-1。INT64-
next_tokens属性用于稀疏计算，表示attention需要和后几个Token计算关联，仅支持2^63-1。INT64-
attention_mode属性表示attention的模式，仅支持传入2，表示MLA-absorb模式，即QK的D包含rope和nope两部分，且KV是同一份。INT64-
quant_scale_repo_mode属性表示量化参数的存放模式，仅支持传入1，表示combine模式，即量化参数和数据混合存放。INT64-
tile_size属性表示per_tile时每个参数对应的数据块大小，仅在per_tile时有效，仅支持128。INT64-
rope_head_dim属性表示MLA架构下的rope_head_dim大小，仅在attention_mode为2时有效，仅支持64。INT64-
output输出代表公式中的输出Attention。输出shape与入参query的shape保持一致，layout_query为"BSND"时shape为[B, Q_S, Q_N, Q_out_D]，layout_query为"TND"时shape为[Q_T, Q_N, Q_out_D]，其中Q_out_D = Q_D - rope_head_dim。FLOAT16、BFLOAT16ND
+ +## 约束说明 + +- 该接口支持图模式。 +- 参数query shape中：Atlas A3 训练系列产品/Atlas A3 推理系列产品、Atlas A2 训练系列产品/Atlas A2 推理系列产品：Q_N不支持48。 +- 参数key、value数据类型要求： + - Ascend 950PR/Ascend 950DT：仅支持float8_e4m3、int8、hifloat8数据类型。 + - Atlas A3 训练系列产品/Atlas A3 推理系列产品、Atlas A2 训练系列产品/Atlas A2 推理系列产品：仅支持int8数据类型。 +- 参数sparse\_block\_size： + - Ascend 950PR/Ascend 950DT：只支持sparse\_block\_size为1。 + - Atlas A3 训练系列产品/Atlas A3 推理系列产品、Atlas A2 训练系列产品/Atlas A2 推理系列产品：支持[1,16]，且要求是2的幂次方，在PageAttention场景下要求sparse\_block\_size整除block\_size +- 非PageAttention场景layout\_query和layout\_kv取值需要保持一致。 diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/kv_quant_sparse_flash_attention_torch_adpt.h b/csrc/ascend/attention/kv_quant_sparse_flash_attention/kv_quant_sparse_flash_attention_torch_adpt.h new file mode 100644 index 000000000..d6373a6ea --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/kv_quant_sparse_flash_attention_torch_adpt.h @@ -0,0 +1,168 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef KV_QUANT_SPARSE_FLASH_ATTENTION_TORCH_ADPT_H +#define KV_QUANT_SPARSE_FLASH_ATTENTION_TORCH_ADPT_H + +namespace vllm_fl { + +namespace { + +std::tuple +construct_kv_quant_sparse_flash_attention_output_tensor( + const at::Tensor &query, + const at::Tensor &key, + const std::string &layout_query_str, + const std::string &layout_kv_str, + int64_t rope_head_dim, + bool return_softmax_lse) +{ + constexpr int64_t SIZE = 8; + constexpr int64_t DIM_0 = 0; + constexpr int64_t DIM_1 = 1; + constexpr int64_t DIM_2 = 2; + constexpr int64_t DIM_3 = 3; + constexpr int64_t DIM_4 = 4; + + TORCH_CHECK(layout_query_str == "BSND" || layout_query_str == "TND", + "The layout of query only support BSND and TND, but got ", + layout_query_str); + for (size_t i = 0; i < query.sizes().size(); i++) { + TORCH_CHECK(query.size(i) > 0, + "All values within query's shape should be greater than 0, but shape[", + i, "] is ", query.size(i)); + } + + at::SmallVector output_size; + if (layout_query_str == "BSND") { + TORCH_CHECK(query.dim() == DIM_4, + "When the layout of query is BSND, the query dimension must be 4, but got ", + query.dim()); + output_size = {query.size(DIM_0), query.size(DIM_1), query.size(DIM_2), + query.size(DIM_3) - rope_head_dim}; + } else { + TORCH_CHECK(query.dim() == DIM_3, + "When the layout of query is TND, the query dimension must be 3, but got ", + query.dim()); + output_size = {query.size(DIM_0), query.size(DIM_1), + query.size(DIM_2) - rope_head_dim}; + } + + at::Tensor attention_output = + at::empty(output_size, query.options().dtype(query.dtype())); + at::SmallVector softmax_size; + if (return_softmax_lse) { + if (query.dim() == DIM_3) { + const int64_t kv_head_dim = + layout_kv_str == "PA_BSND" ? key.size(DIM_2) : key.size(DIM_1); + softmax_size = {kv_head_dim, query.size(DIM_0), + query.size(DIM_1) / kv_head_dim}; + } else { + softmax_size = {query.size(DIM_0), key.size(DIM_2), + query.size(DIM_1), + query.size(DIM_2) / key.size(DIM_2)}; + } + } else { + softmax_size = {0}; + } + + at::Tensor softmax_max = + at::empty(softmax_size, query.options().dtype(at::kFloat)); + at::Tensor softmax_sum = + at::empty(softmax_size, query.options().dtype(at::kFloat)); + return std::tuple( + attention_output, softmax_max, softmax_sum); +} + +} // namespace + +std::tuple +npu_kv_quant_sparse_flash_attention( + const at::Tensor &query, + const at::Tensor &key, + const at::Tensor &value, + const at::Tensor &sparse_indices, + double scale_value, + int64_t key_quant_mode, + int64_t value_quant_mode, + const c10::optional &key_dequant_scale, + const c10::optional &value_dequant_scale, + const c10::optional &block_table, + const c10::optional &actual_seq_lengths_query, + const c10::optional &actual_seq_lengths_kv, + int64_t sparse_block_size, + c10::string_view layout_query, + c10::string_view layout_kv, + int64_t sparse_mode, + int64_t pre_tokens, + int64_t next_tokens, + int64_t attention_mode, + int64_t quant_scale_repo_mode, + int64_t tile_size, + int64_t rope_head_dim, + bool return_softmax_lse) +{ + TORCH_CHECK(query.numel() > 0, "Tensor query is empty."); + TORCH_CHECK(key.numel() > 0, "Tensor key is empty."); + TORCH_CHECK(value.numel() > 0, "Tensor value is empty."); + TORCH_CHECK(sparse_indices.numel() > 0, "Tensor sparse_indices is empty."); + + std::string layout_query_str = std::string(layout_query); + std::string layout_kv_str = std::string(layout_kv); + + auto output = construct_kv_quant_sparse_flash_attention_output_tensor( + query, key, layout_query_str, layout_kv_str, rope_head_dim, + return_softmax_lse); + at::Tensor attention_output = std::get<0>(output); + at::Tensor softmax_max = std::get<1>(output); + at::Tensor softmax_sum = std::get<2>(output); + + char *layout_query_ptr = const_cast(layout_query_str.c_str()); + char *layout_kv_ptr = const_cast(layout_kv_str.c_str()); + + EXEC_NPU_CMD( + aclnnKvQuantSparseFlashAttention, + query, + key, + value, + sparse_indices, + key_dequant_scale, + value_dequant_scale, + block_table, + actual_seq_lengths_query, + actual_seq_lengths_kv, + scale_value, + key_quant_mode, + value_quant_mode, + sparse_block_size, + layout_query_ptr, + layout_kv_ptr, + sparse_mode, + pre_tokens, + next_tokens, + attention_mode, + quant_scale_repo_mode, + tile_size, + rope_head_dim, + return_softmax_lse, + attention_output, + softmax_max, + softmax_sum); + return std::tuple( + attention_output, softmax_max, softmax_sum); +} +} // namespace vllm_fl + +#endif // KV_QUANT_SPARSE_FLASH_ATTENTION_TORCH_ADPT_H diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/CMakeLists.txt b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/CMakeLists.txt new file mode 100644 index 000000000..e834b76b7 --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/CMakeLists.txt @@ -0,0 +1,32 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +add_op_to_compiled_list() + +set(KV_QUANT_SPARSE_FLASH_ATTENTION_SKIP_HEADER TRUE CACHE INTERNAL "Skip packaging header for this operator") + +if (BUILD_OPEN_PROJECT) + set(kv_quant_sparse_flash_attention_depends attention/common attention/sparse_flash_attention CACHE INTERNAL "Dependencies for kv_quant_sparse_flash_attention") + target_sources(op_host_aclnn PRIVATE + kv_quant_sparse_flash_attention_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME KvQuantSparseFlashAttention + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-vf-remove-membar=false + -mllvm -cce-aicore-hoist-movemask=false +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE kv_quant_sparse_flash_attention ACLNNTYPE aclnn) +endif() diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/kv_quant_sparse_flash_attention_def.cpp b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/kv_quant_sparse_flash_attention_def.cpp new file mode 100644 index 000000000..375323168 --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/kv_quant_sparse_flash_attention_def.cpp @@ -0,0 +1,170 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_def.cpp + * \brief + */ + +#include "register/op_def_registry.h" + +namespace ops { +class KvQuantSparseFlashAttention : public OpDef { +public: + explicit KvQuantSparseFlashAttention(const char *name) : OpDef(name) + { + this->Input("query") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("key") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_INT8}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("value") + .ParamType(REQUIRED) + .Follow("key") + .AutoContiguous(); + this->Input("sparse_indices") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("key_dequant_scale") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("value_dequant_scale") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("block_table") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("actual_seq_lengths_query") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("actual_seq_lengths_kv") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("attention_out") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("softmax_max") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("softmax_sum") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Attr("scale_value").AttrType(REQUIRED).Float(1.0); + this->Attr("key_quant_mode").AttrType(REQUIRED).Int(1); + this->Attr("value_quant_mode").AttrType(REQUIRED).Int(1); + this->Attr("sparse_block_size").AttrType(OPTIONAL).Int(1); + this->Attr("layout_query").AttrType(OPTIONAL).String("BSND"); + this->Attr("layout_kv").AttrType(OPTIONAL).String("BSND"); + this->Attr("sparse_mode").AttrType(OPTIONAL).Int(3); // 3:默认值，只计算下三角 + this->Attr("pre_tokens").AttrType(OPTIONAL).Int(INT64_MAX); + this->Attr("next_tokens").AttrType(OPTIONAL).Int(INT64_MAX); + this->Attr("attention_mode").AttrType(OPTIONAL).Int(0); + this->Attr("quant_scale_repo_mode").AttrType(OPTIONAL).Int(1); + this->Attr("tile_size").AttrType(OPTIONAL).Int(128); // 128:默认值 + this->Attr("rope_head_dim").AttrType(OPTIONAL).Int(64); // 64:默认值 + this->Attr("return_softmax_lse").AttrType(OPTIONAL).Bool(false); + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true); + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + + OpAICoreConfig aicore_config_95; + aicore_config_95.Input("query") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("key") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT8_E4M3FN, ge::DT_HIFLOAT8, ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_HIFLOAT8, ge::DT_INT8}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("value") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT8_E4M3FN, ge::DT_HIFLOAT8, ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_HIFLOAT8, ge::DT_INT8}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("sparse_indices") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("key_dequant_scale") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("value_dequant_scale") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("block_table") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("actual_seq_lengths_query") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Input("actual_seq_lengths_kv") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_95.Output("attention_out") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}); + aicore_config_95.Output("softmax_max") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}); + aicore_config_95.Output("softmax_sum") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}); + aicore_config_95.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true); + } +}; +OP_ADD(KvQuantSparseFlashAttention); +} // namespace ops diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/kv_quant_sparse_flash_attention_infershape.cpp b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/kv_quant_sparse_flash_attention_infershape.cpp new file mode 100644 index 000000000..930959324 --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/kv_quant_sparse_flash_attention_infershape.cpp @@ -0,0 +1,153 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_infershape.cpp + * \brief + */ + +#include +#include +#include "err/ops_err.h" + +using namespace ge; + +#ifdef OP_LOGE_WITH_INVALID_INPUT +#undef OP_LOGE_WITH_INVALID_INPUT +#endif +#define OP_LOGE_WITH_INVALID_INPUT(opname, param) \ + OP_LOGE(opname, "Invalid input: %s.", param) + +namespace ops { +constexpr size_t QUERY_INPUT_INDEX = 0; +constexpr size_t KEY_INPUT_INDEX = 1; +constexpr uint32_t LAYOUT_QUERY_ATTR_INDEX = 4; +constexpr uint32_t LAYOUT_KV_ATTR_INDEX = 5; +constexpr uint32_t ROPE_HEAD_DIM_ATTR_INDEX = 12; +constexpr uint32_t RETURN_SOFTMAX_LSE_INDEX = 13; +constexpr uint32_t DIM_INDEX_0 = 0; +constexpr uint32_t DIM_INDEX_1 = 1; +constexpr uint32_t DIM_INDEX_2 = 2; +constexpr uint32_t DIM_INDEX_3 = 3; +constexpr uint32_t DIM_NUM_1 = 1; +constexpr uint32_t DIM_NUM_3 = 3; +constexpr uint32_t DIM_NUM_4 = 4; +constexpr uint32_t OUTPUT_INDEX_0 = 0; +constexpr uint32_t OUTPUT_INDEX_1 = 1; +constexpr uint32_t OUTPUT_INDEX_2 = 2; + +ge::graphStatus InferShapeKvQuantSparseFlashAttention(gert::InferShapeContext *context) +{ + OP_CHECK_IF(context == nullptr, OP_LOGE_WITH_INVALID_INPUT("KvQuantSparseFlashAttention", "InferShapeContext"), + return ge::GRAPH_FAILED); + const gert::Shape *queryShape = context->GetInputShape(QUERY_INPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, queryShape); + const gert::Shape *keyShape = context->GetInputShape(KEY_INPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, keyShape); + + gert::Shape *attentionOutShape = context->GetOutputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context, attentionOutShape); + gert::Shape *softmaxMaxShape = context->GetOutputShape(OUTPUT_INDEX_1); + OP_CHECK_NULL_WITH_CONTEXT(context, softmaxMaxShape); + gert::Shape *softmaxSumShape = context->GetOutputShape(OUTPUT_INDEX_2); + OP_CHECK_NULL_WITH_CONTEXT(context, softmaxSumShape); + + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const char *inputLayoutQueryPtr = attrs->GetAttrPointer(LAYOUT_QUERY_ATTR_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, inputLayoutQueryPtr); + std::string inputLayoutQueryPtrStr = std::string(inputLayoutQueryPtr); + const char *inputLayoutKvPtr = attrs->GetAttrPointer(LAYOUT_KV_ATTR_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, inputLayoutKvPtr); + std::string inputLayoutKvPtrStr = std::string(inputLayoutKvPtr); + const int64_t ropeHeadDim = *attrs->GetAttrPointer(ROPE_HEAD_DIM_ATTR_INDEX); + const bool *lse_flag = attrs->GetAttrPointer(RETURN_SOFTMAX_LSE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, lse_flag); + bool return_softmax_lse = (lse_flag != nullptr) ? *lse_flag : false; + + *attentionOutShape = *queryShape; + if (inputLayoutQueryPtrStr == "BSND") { + attentionOutShape->SetDimNum(DIM_NUM_4); + attentionOutShape->SetDim(DIM_INDEX_0, queryShape->GetDim(DIM_INDEX_0)); + attentionOutShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_1)); + attentionOutShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_2)); // 2:dim2 + if(queryShape->GetDim(DIM_INDEX_3) != -1){ + attentionOutShape->SetDim(DIM_INDEX_3, queryShape->GetDim(DIM_INDEX_3) - ropeHeadDim); // 3:dim3 + } + } else { // TND + attentionOutShape->SetDimNum(DIM_NUM_3); + attentionOutShape->SetDim(DIM_INDEX_0, queryShape->GetDim(DIM_INDEX_0)); + attentionOutShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_1)); + if(queryShape->GetDim(DIM_INDEX_2) != -1){ + attentionOutShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_2) - ropeHeadDim); // 2:dim2 + } + } + + if (return_softmax_lse) { + if (queryShape->GetDimNum() == DIM_NUM_3) { + if (inputLayoutKvPtrStr == "PA_BSND") { + softmaxMaxShape->SetDimNum(DIM_NUM_3); + softmaxMaxShape->SetDim(DIM_INDEX_0, keyShape->GetDim(DIM_INDEX_2)); + softmaxMaxShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_0)); + softmaxMaxShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1) / keyShape->GetDim(DIM_INDEX_2)); + + softmaxSumShape->SetDimNum(DIM_NUM_3); + softmaxSumShape->SetDim(DIM_INDEX_0, keyShape->GetDim(DIM_INDEX_2)); + softmaxSumShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_0)); + softmaxSumShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1) / keyShape->GetDim(DIM_INDEX_2)); + } else { + softmaxMaxShape->SetDimNum(DIM_NUM_3); + softmaxMaxShape->SetDim(DIM_INDEX_0, keyShape->GetDim(DIM_INDEX_1)); + softmaxMaxShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_0)); + softmaxMaxShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1) / keyShape->GetDim(DIM_INDEX_1)); + + softmaxSumShape->SetDimNum(DIM_NUM_3); + softmaxSumShape->SetDim(DIM_INDEX_0, keyShape->GetDim(DIM_INDEX_1)); + softmaxSumShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_0)); + softmaxSumShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1) / keyShape->GetDim(DIM_INDEX_1)); + } + } else { + softmaxMaxShape->SetDimNum(DIM_NUM_4); + softmaxMaxShape->SetDim(DIM_INDEX_0, queryShape->GetDim(DIM_INDEX_0)); + softmaxMaxShape->SetDim(DIM_INDEX_1, keyShape->GetDim(DIM_INDEX_2)); + softmaxMaxShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1)); + softmaxMaxShape->SetDim(DIM_INDEX_3, queryShape->GetDim(DIM_INDEX_2) / keyShape->GetDim(DIM_INDEX_2)); + + softmaxSumShape->SetDimNum(DIM_NUM_4); + softmaxSumShape->SetDim(DIM_INDEX_0, queryShape->GetDim(DIM_INDEX_0)); + softmaxSumShape->SetDim(DIM_INDEX_1, keyShape->GetDim(DIM_INDEX_2)); + softmaxSumShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1)); + softmaxSumShape->SetDim(DIM_INDEX_3, queryShape->GetDim(DIM_INDEX_2) / keyShape->GetDim(DIM_INDEX_2)); + } + } else { + softmaxMaxShape->SetDimNum(DIM_NUM_1); + softmaxMaxShape->SetDim(DIM_INDEX_0, 0); + softmaxSumShape->SetDimNum(DIM_NUM_1); + softmaxSumShape->SetDim(DIM_INDEX_0, 0); + } + return GRAPH_SUCCESS; +} + +ge::graphStatus InferDataTypeKvQuantSparseFlashAttention(gert::InferDataTypeContext *context) +{ + OP_CHECK_IF(context == nullptr, OP_LOGE_WITH_INVALID_INPUT("KvQuantSparseFlashAttention", "InferShapeContext"), + return ge::GRAPH_FAILED); + const auto inputDataType = context->GetInputDataType(QUERY_INPUT_INDEX); + context->SetOutputDataType(OUTPUT_INDEX_0, inputDataType); + context->SetOutputDataType(OUTPUT_INDEX_1, ge::DT_FLOAT); + context->SetOutputDataType(OUTPUT_INDEX_2, ge::DT_FLOAT); + context->SetOutputDataType(0, inputDataType); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(KvQuantSparseFlashAttention) + .InferShape(InferShapeKvQuantSparseFlashAttention) + .InferDataType(InferDataTypeKvQuantSparseFlashAttention); +} // namespace ops diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/kv_quant_sparse_flash_attention_tiling.cpp b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/kv_quant_sparse_flash_attention_tiling.cpp new file mode 100644 index 000000000..ad3594555 --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/kv_quant_sparse_flash_attention_tiling.cpp @@ -0,0 +1,2073 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_tiling.cpp + * \brief + */ + +#include +#include +#include +#include +#include +#include +#include "err/ops_err.h" +#include "register/op_def_registry.h" +#include "../op_kernel/kv_quant_sparse_flash_attention_template_tiling_key.h" +#include "kv_quant_sparse_flash_attention_tiling.h" + +using std::map; +using std::string; +using std::pair; + +using namespace ge; +using namespace AscendC; +namespace optiling { + +inline std::string QSFAErrorToString(const char *value) +{ + return value == nullptr ? std::string() : std::string(value); +} + +inline std::string QSFAErrorToString(char *value) +{ + return QSFAErrorToString(static_cast(value)); +} + +inline std::string QSFAErrorToString(const std::string &value) +{ + return value; +} + +template +std::string QSFAErrorToString(const T &value) +{ + std::ostringstream oss; + oss << value; + return oss.str(); +} + +#define QSFA_LOG_INVALID_WITH_EXPECTED(opname, kind, param, actual, expected) \ + do { \ + const auto qsfaParam = ::optiling::QSFAErrorToString(param); \ + const auto qsfaActual = ::optiling::QSFAErrorToString(actual); \ + const auto qsfaExpected = ::optiling::QSFAErrorToString(expected); \ + OP_LOGE(opname, "Invalid %s for %s, actual: %s, expected: %s.", kind, \ + qsfaParam.c_str(), qsfaActual.c_str(), qsfaExpected.c_str()); \ + } while (0) + +#define QSFA_LOG_INVALID_WITH_REASON(opname, kind, param, actual, reason) \ + do { \ + const auto qsfaParam = ::optiling::QSFAErrorToString(param); \ + const auto qsfaActual = ::optiling::QSFAErrorToString(actual); \ + const auto qsfaReason = ::optiling::QSFAErrorToString(reason); \ + OP_LOGE(opname, "Invalid %s for %s, actual: %s, reason: %s.", kind, \ + qsfaParam.c_str(), qsfaActual.c_str(), qsfaReason.c_str()); \ + } while (0) + +#ifdef OP_LOGE_WITH_INVALID_INPUT +#undef OP_LOGE_WITH_INVALID_INPUT +#endif +#define OP_LOGE_WITH_INVALID_INPUT(opname, param) \ + do { \ + const auto qsfaParam = ::optiling::QSFAErrorToString(param); \ + OP_LOGE(opname, "Invalid input: %s.", qsfaParam.c_str()); \ + } while (0) + +#ifndef OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON +#define OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(opname, param, actual, reason) \ + QSFA_LOG_INVALID_WITH_REASON(opname, "dtype", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON +#define OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(opname, param, actual, reason) \ + QSFA_LOG_INVALID_WITH_REASON(opname, "dtype", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPE +#define OP_LOGE_FOR_INVALID_SHAPE(opname, param, actual, expected) \ + QSFA_LOG_INVALID_WITH_EXPECTED(opname, "shape", param, actual, expected) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON +#define OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(opname, param, actual, reason) \ + QSFA_LOG_INVALID_WITH_REASON(opname, "shape", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON +#define OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(opname, param, actual, reason) \ + QSFA_LOG_INVALID_WITH_REASON(opname, "shape dim", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON +#define OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON(opname, param, actual, reason) \ + QSFA_LOG_INVALID_WITH_REASON(opname, "shape size", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_VALUE +#define OP_LOGE_FOR_INVALID_VALUE(opname, param, actual, expected) \ + QSFA_LOG_INVALID_WITH_EXPECTED(opname, "value", param, actual, expected) +#endif + +#ifndef OP_LOGE_FOR_INVALID_VALUE_WITH_REASON +#define OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opname, param, actual, reason) \ + QSFA_LOG_INVALID_WITH_REASON(opname, "value", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_VALUES_WITH_REASON +#define OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(opname, param, actual, reason) \ + QSFA_LOG_INVALID_WITH_REASON(opname, "values", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_FORMAT +#define OP_LOGE_FOR_INVALID_FORMAT(opname, param, actual, expected) \ + QSFA_LOG_INVALID_WITH_EXPECTED(opname, "format", param, actual, expected) +#endif + +#ifndef OP_LOGE_FOR_INVALID_FORMATS_WITH_REASON +#define OP_LOGE_FOR_INVALID_FORMATS_WITH_REASON(opname, param, actual, reason) \ + QSFA_LOG_INVALID_WITH_REASON(opname, "format", param, actual, reason) +#endif + +constexpr uint32_t PRE_LOAD_NUM = 2; +constexpr uint32_t BLOCK_TABLE_ELEM_BYTE = 4; +constexpr int32_t SPARSE_MODE_BAND = 4; + +static const std::string QUERY_NAME = "query"; +static const std::string KEY_NAME = "key"; +static const std::string VALUE_NAME = "value"; +static const std::string SPARSE_INDICES_NAME = "sparse_indices"; +static const std::string BLOCK_TABLE_NAME = "block_table"; +static const std::string ATTEN_OUT_NAME = "attention_out"; + +const std::map> DTYPE_SUPPORT_MAP = { + {QUERY_NAME, {ge::DT_FLOAT16, ge::DT_BF16}}, + {KEY_NAME, {ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_HIFLOAT8}}, + {VALUE_NAME, {ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_HIFLOAT8}}, + {ATTEN_OUT_NAME, {ge::DT_FLOAT16, ge::DT_BF16}}, + {SPARSE_INDICES_NAME, {ge::DT_INT32}} +}; + +const std::map> LAYOUT_SUPPORT_MAP = { + {QUERY_NAME, {QSFALayout::BSND, QSFALayout::TND}}, + {KEY_NAME, {QSFALayout::BSND, QSFALayout::TND, QSFALayout::PA_BSND}}, + {VALUE_NAME, {QSFALayout::BSND, QSFALayout::TND, QSFALayout::PA_BSND}}, + {ATTEN_OUT_NAME, {QSFALayout::BSND, QSFALayout::TND}}, +}; + +const std::map DATATYPE_TO_STRING_MAP = { + {ge::DT_FLOAT, "DT_FLOAT"}, // float type + {ge::DT_UNDEFINED, "DT_UNDEFINED"}, // Used to indicate a DataType field has not been set. + {ge::DT_FLOAT16, "DT_FLOAT16"}, // fp16 type + {ge::DT_INT8, "DT_INT8"}, // int8 type + {ge::DT_INT16, "DT_INT16"}, // int16 type + {ge::DT_FLOAT8_E4M3FN, "DT_FLOAT8_E4M3FN"}, // fp8_e4m3 type + {ge::DT_HIFLOAT8, "DT_HIFLOAT8"}, // hifloat8 type + {ge::DT_UINT16, "DT_UINT16"}, // uint16 type + {ge::DT_UINT8, "DT_UINT8"}, // uint8 type + {ge::DT_INT64, "DT_INT64"}, // int64 type + {ge::DT_INT32, "DT_INT32"}, // int32 type + {ge::DT_UINT64, "DT_UINT64"}, // unsigned int64 + {ge::DT_UINT32, "DT_UINT32"}, // unsigned int32 + {ge::DT_BOOL, "DT_BOOL"}, // bool type + {ge::DT_DOUBLE, "DT_DOUBLE"}, // double type + {ge::DT_DUAL, "DT_DUAL"}, // dual output type + {ge::DT_COMPLEX32, "DT_COMPLEX32"}, // complex32 type + {ge::DT_COMPLEX64, "DT_COMPLEX64"}, // complex64 type + {ge::DT_COMPLEX128, "DT_COMPLEX128"}, // complex128 type + {ge::DT_DUAL_SUB_INT8, "DT_DUAL_SUB_INT8"}, // dual output int8 type + {ge::DT_DUAL_SUB_UINT8, "DT_DUAL_SUB_UINT8"}, // dual output uint8 type + {ge::DT_QUINT8, "DT_QUINT8"}, // quint8 type + {ge::DT_QUINT16, "DT_QUINT16"}, // quint16 type + {ge::DT_QINT8, "DT_QINT8"}, // qint8 type + {ge::DT_QINT16, "DT_QINT16"}, // qint16 type + {ge::DT_QINT32, "DT_QINT32"}, // qint32 type + {ge::DT_RESOURCE, "DT_RESOURCE"}, // resource type + {ge::DT_STRING_REF, "DT_STRING_REF"}, // string ref type + {ge::DT_BF16, "DT_BFLOAT16"}, // dt_bfloat16 type + {ge::DT_STRING, "DT_STRING"}, // string type + {ge::DT_VARIANT, "DT_VARIANT"}, // dt_variant type + {ge::DT_INT2, "DT_INT2"}, // dt_variant type + {ge::DT_UINT2, "DT_UINT2"}, // dt_variant type + {ge::DT_INT4, "DT_INT4"}, // dt_variant type + {ge::DT_UINT1, "DT_UINT1"} // dt_variant type +}; + +struct KvQuantSparseFlashAttentionCompileInfo { + int64_t coreNum; +}; + +static const std::map> QSFA_LAYOUT_AXIS_MAP = { + {QSFALayout::BSND, {QSFAAxis::B, QSFAAxis::S, QSFAAxis::N, QSFAAxis::D}}, + {QSFALayout::TND, {QSFAAxis::T, QSFAAxis::N, QSFAAxis::D}}, + {QSFALayout::PA_BSND, {QSFAAxis::Bn, QSFAAxis::Bs, QSFAAxis::N, QSFAAxis::D}}, +}; + +static const std::map QSFA_LAYOUT_DIM_MAP = { + {QSFALayout::BSND, DIM_NUM_FOUR}, + {QSFALayout::TND, DIM_NUM_THREE}, + {QSFALayout::PA_BSND, DIM_NUM_FOUR}, +}; + +template +static std::string GetShapeStr(const T &shape) +{ + std::ostringstream qsfaOss; + qsfaOss << "["; + if (shape.GetDimNum() > 0) { + for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) { + qsfaOss << shape.GetDim(i) << ", "; + } + qsfaOss << shape.GetDim(shape.GetDimNum() - 1); + } + qsfaOss << "]"; + return qsfaOss.str(); +} + +static std::string QSFADataTypeToSerialString(ge::DataType type) +{ + const auto qsfaIt = DATATYPE_TO_STRING_MAP.find(type); + if (qsfaIt != DATATYPE_TO_STRING_MAP.end()) { + return qsfaIt->second; + } else { + OP_LOGE("SparseFlashAttention", "datatype %d not support", type); + return "UNDEFINED"; + } +} + +string QSFATensorDesc2String(const gert::StorageShape *shape, const gert::CompileTimeTensorDesc *tensor) +{ + if (shape == nullptr || tensor == nullptr) { + return "nil "; + } + + std::ostringstream qsfaOss; + qsfaOss << "(dtype: " << ge::TypeUtils::DataTypeToAscendString(tensor->GetDataType()).GetString() << "),"; + qsfaOss << "(shape:" << GetShapeStr(shape->GetStorageShape()) << "),"; + qsfaOss << "(ori_shape:" << GetShapeStr(shape->GetOriginShape()) << "),"; + qsfaOss << "(format: " + << ge::TypeUtils::FormatToAscendString( + static_cast(ge::GetPrimaryFormat(tensor->GetStorageFormat()))) + .GetString() + << "),"; + qsfaOss << "(ori_format: " << ge::TypeUtils::FormatToAscendString(tensor->GetOriginFormat()).GetString() << ") "; + + return qsfaOss.str(); +} + +string QSFADebugTilingContext(const gert::TilingContext *context) +{ + std::ostringstream qsfaOss; + for (size_t i = 0; i < context->GetComputeNodeInfo()->GetInputsNum(); ++i) { + qsfaOss << "input" << i << ": "; + qsfaOss << QSFATensorDesc2String(context->GetInputShape(i), context->GetInputDesc(i)); + } + + for (size_t i = 0; i < context->GetComputeNodeInfo()->GetOutputsNum(); ++i) { + qsfaOss << "output" << i << ": "; + qsfaOss << QSFATensorDesc2String(context->GetOutputShape(i), context->GetOutputDesc(i)); + } + return qsfaOss.str(); +} + +std::string QSFALayoutToSerialString(QSFALayout layout) +{ + switch (layout) { + case QSFALayout::BSND: return "BSND"; + case QSFALayout::TND: return "TND"; + case QSFALayout::PA_BSND: return "PA_BSND"; + default: return "UNKNOWN"; + } +} + +static uint32_t GetTypeSize(ge::DataType dtype) +{ + uint32_t qsfaTypeSize = NUM_BYTES_FLOAT16; + switch (dtype) { + case ge::DT_FLOAT16: + qsfaTypeSize = NUM_BYTES_FLOAT16; + break; + case ge::DT_BF16: + qsfaTypeSize = NUM_BYTES_BF16; + break; + default: + qsfaTypeSize = NUM_BYTES_FLOAT16; + } + return qsfaTypeSize; +} + +ge::graphStatus QSFAMlaTiling::SetBlockDim(uint32_t blockDim) const +{ + context_->SetBlockDim(blockDim); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAMlaTiling::SetTilingKey(uint64_t tilingKey) const +{ + context_->SetTilingKey(tilingKey); + context_->SetScheduleMode(1); // 1: batchmode模式 + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAMlaTiling::SetWorkspaceSize(uint64_t workspaceSize) const +{ + OP_CHECK_IF(context_->GetWorkspaceSizes(1) == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "workSpaceSize got from ge is nullptr"), + return ge::GRAPH_FAILED); + size_t *workSpaces = context_->GetWorkspaceSizes(1); + workSpaces[0] = workspaceSize; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAMlaTiling::SetTilingData(TilingDef &tilingData) const +{ + OP_CHECK_IF(context_->GetRawTilingData() == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "RawTilingData got from GE context is nullptr."), + return ge::GRAPH_FAILED); + + tilingData.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData.GetDataSize()); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAMlaTiling::GetPlatformInfo() +{ + OP_CHECK_IF(qsfaInfo_->platformInfo == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(qsfaInfo_->opName, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED); + + auto qsfaAscendcPlatform = platform_ascendc::PlatformAscendC(qsfaInfo_->platformInfo); + libapiSize_ = qsfaAscendcPlatform.GetLibApiWorkSpaceSize(); + aivNum_ = qsfaAscendcPlatform.GetCoreNumAiv(); + aicNum_ = qsfaAscendcPlatform.GetCoreNumAic(); + + OP_CHECK_IF(aicNum_ == 0 || aivNum_ == 0, + OPS_REPORT_VECTOR_INNER_ERR(qsfaInfo_->opName, "num of core obtained is 0."), return GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void QSFAMlaTiling::GenTilingKey() +{ + uint32_t layoutQuery = static_cast(qsfaInfo_->qLayout); + uint32_t layoutKV = static_cast(qsfaInfo_->kvLayout); + uint32_t pageAttention = 0U; + if (qsfaInfo_->kvLayout == QSFALayout::PA_BSND) { + pageAttention = 1U; + } + + tilingKey_ = GET_TPL_TILING_KEY(0U, pageAttention, layoutQuery, layoutKV, \ + perfMode_ == QSFAPerfMode::V_TEMPLATE_MODE, static_cast(qsfaInfo_->gSize > 64)); // G大于64时核间切G + + OP_LOGI(qsfaInfo_->opName, "QSFA tilingKey_: %lu.", tilingKey_); +} + +void QSFAMlaTiling::ZeroTensorProcess() const +{ + if (qsfaInfo_->s2Size == 0) { + /* + * 1024，空tensor场景下，作为默认值完成后续计算 + * 避免matmal tiling softmax tiling异常 + * kernel计算使用真实的seqSize=0, 与actuseq_len流程归一 + */ + qsfaInfo_->s2Size = 1024; + } +} + +void QSFAMlaTiling::InitParams() +{ + perfMode_ = QSFAPerfMode::V_TEMPLATE_MODE; + coreNum_ = aicNum_; + + headDimAlign_ = Align(qsfaInfo_->qHeadDim, BYTE_BLOCK); // 元素个数按照基本块大小对齐 + ZeroTensorProcess(); +} + +void QSFAMlaTiling::CalcUbBmm() +{ + uint32_t qsfaCubeMSize = qsfaInfo_->gSize * qsfaInfo_->s1Size; + uint32_t qsfaMaxMSize = mBaseSize_; + if (qsfaCubeMSize > qsfaMaxMSize) { + qsfaCubeMSize = qsfaMaxMSize; + } + mmResUbSize_ = sInnerSizeAlign_ * Align(qsfaCubeMSize, 16U); // kernel按照16对齐写出，tiling按照这个原则分配内存 + bmm2ResUbSize_ = headDimAlign_ * Align(qsfaCubeMSize, 16U); // kernel按照16对齐写出，tiling按照这个原则分配内存 + + qPreSizeMla_ = qsfaInfo_->gSize * (headDimAlign_ + 64U) * qsfaInfo_->s1Size; +} + +void QSFAMlaTiling::CheckUbSpace() +{ + CalcUbBmm(); +} + +void QSFAMlaTiling::CalcInnerSize(uint32_t qsfaS2Size) +{ + sInnerSize_ = 512; // 512:s2默认切分大小 + // FlashDecode时，如果S2的计算量>=256(确保切分后不小于128)但又不足以分2次计算时，则修改sInnerSize_，均分为2份进行计算，确保Nbuffer=2 + if (splitKVFlag_ && qsfaInfo_->qLayout != QSFALayout::TND) { + if (qsfaS2Size == 256) { // 256:s2Size的阈值，判断sInnerSize_是否切分 + sInnerSize_ = 128; // 128:sInnerSize_值为s2Size的一半，均分为2份进行计算， + } else if (qsfaS2Size > 256 && qsfaS2Size <= sInnerSize_) { // 256:s2Size的阈值，判断sInnerSize_是否切分 + sInnerSize_ = (sInnerSize_ + 1) / 2; // 2:减半 + } + } + + sInnerLoopTimes_ = (qsfaS2Size + sInnerSize_ - 1) / sInnerSize_; + sInnerSizeTail_ = qsfaS2Size - (sInnerLoopTimes_ - 1) * sInnerSize_; + if (sInnerSize_ > qsfaS2Size) { + sInnerSize_ = qsfaS2Size; + } + sInnerSizeAlign_ = + Align(sInnerSize_, BYTE_BLOCK); // 元素个数按照基本块大小对齐 + CheckUbSpace(); +} + +void QSFAMlaTiling::SplitBalanced() +{ + CalcInnerSize(qsfaInfo_->s2Size); + InnerSplitParams qsfaInnerSplitParams; + qsfaInnerSplitParams.s1GBaseSize = qsfaInfo_->gSize; + tilingData_.innerSplitParams.set_mBaseSize(qsfaInnerSplitParams.s1GBaseSize); + + qsfaInnerSplitParams.s2BaseSize = sInnerSize_; + tilingData_.innerSplitParams.set_s2BaseSize(qsfaInnerSplitParams.s2BaseSize); + + usedCoreNum_ = aicNum_; +} + +void QSFAMlaTiling::Split() +{ + SplitBalanced(); +} + +void QSFAMlaTiling::FillTilingBaseParamsMla() +{ + tilingData_.baseParams.set_batchSize(qsfaInfo_->bSize); + tilingData_.baseParams.set_seqSize(qsfaInfo_->s2Size); + tilingData_.baseParams.set_qSeqSize(qsfaInfo_->s1Size); + tilingData_.baseParams.set_blockSize(qsfaInfo_->blockSize); + tilingData_.baseParams.set_maxBlockNumPerBatch(qsfaInfo_->maxBlockNumPerBatch); + tilingData_.baseParams.set_scaleValue(qsfaInfo_->scaleValue); + tilingData_.baseParams.set_nNumOfQInOneGroup(qsfaInfo_->n1Size / qsfaInfo_->n2Size); + tilingData_.baseParams.set_actualLenDimsQ(qsfaInfo_->actualLenDimsQ); + tilingData_.baseParams.set_actualLenDimsKV(qsfaInfo_->actualLenDimsKV); + tilingData_.baseParams.set_outputLayout(static_cast(qsfaInfo_->outLayout)); + tilingData_.baseParams.set_sparseMode(qsfaInfo_->sparseMode); + tilingData_.baseParams.set_sparseBlockSize(qsfaInfo_->sparseBlockSize); + tilingData_.baseParams.set_sparseBlockCount(qsfaInfo_->sparseBlockCount); + tilingData_.baseParams.set_dSizeVInput(qsfaInfo_->dSizeVInput); + tilingData_.baseParams.set_isActualLenDimsNull(qsfaInfo_->actualQSeqLenFlag ? 0U : 1U); + tilingData_.baseParams.set_isActualLenDimsKVNull(qsfaInfo_->actualSeqLenFlag ? 0U : 1U); + tilingData_.baseParams.set_returnSoftmaxLse(qsfaInfo_->returnSoftmaxLse); +} + +// for flash decode +void QSFAMlaTiling::FillTilingSplitKVMla() +{ + tilingData_.splitKVParams.set_s2(kvSplitPart_); + // 2:每个核可能有头规约和尾规约，一共两份规约信息 + tilingData_.splitKVParams.set_accumOutSize(aicNum_ * 2 * qsfaInfo_->n2Size * mBaseSize_ * headDimAlign_); + // 2:每个核可能有头规约和尾规约，一共两份规约信息;sum + max + tilingData_.splitKVParams.set_logSumExpSize(2 * aicNum_ * 2 * qsfaInfo_->n2Size * mBaseSize_ * + (BYTE_BLOCK / BLOCK_TABLE_ELEM_BYTE)); + + if (!splitKVFlag_) { + tilingData_.splitKVParams.set_s2(0); + } +} + +void QSFAMlaTiling::FillTilingSingleCoreParamsMla() +{ + tilingData_.singleCoreParams.set_usedCoreNum(usedCoreNum_); +} + +void QSFAMlaTiling::FillTilingSingleCoreTensorSizeMla() +{ + tilingData_.singleCoreTensorSize.set_mmResUbSize(mmResUbSize_); + tilingData_.singleCoreTensorSize.set_bmm2ResUbSize(bmm2ResUbSize_); +} + +void QSFAMlaTiling::FillTiling() +{ + FillTilingBaseParamsMla(); + FillTilingSplitKVMla(); + FillTilingSingleCoreParamsMla(); + FillTilingSingleCoreTensorSizeMla(); +} + +uint32_t QSFAMlaTiling::CalcBalanceFDParamNums(const uint32_t actCoreNum) const +{ + return actCoreNum * 2 * qsfaInfo_->n2Size * mBaseSize_; // 2:每个核可能有头规约和尾规约，一共两份规约信息 +} + +void QSFAMlaTiling::NormalCalcFDWorkSpace(const uint32_t actCoreNum) +{ + if (splitKVFlag_) { + uint32_t accumOutSize = 0; + uint32_t logSumExpSize = 0; + uint32_t FDParamNums = CalcBalanceFDParamNums(actCoreNum); + accumOutSize = FDParamNums * headDimAlign_; + logSumExpSize = 2 * FDParamNums * (BYTE_BLOCK / qsfaInfo_->blockTypeSize); // log和sum的存储空间一致，共需要2份内存 + workspaceSize_ += (accumOutSize + logSumExpSize) * qsfaInfo_->blockTypeSize; + if (qsfaInfo_->npuArch == NpuArch::DAV_2002) { // 310P + workspaceSize_ += static_cast(actCoreNum) * 32; // 每个核SyncAll软同步需要32Byte记录状态 + } + } +} + +void QSFAMlaTiling::CalcFDWorkSpace(const uint32_t actCoreNum) +{ + NormalCalcFDWorkSpace(actCoreNum); +} + +void QSFAMlaTiling::GetWorkspaceSize() +{ + uint32_t actCoreNum = coreNum_; + if (qsfaInfo_->isA5) { + workspaceSize_ = libapiSize_; + constexpr uint32_t TRIPLE_BUFFER_NUM = 3; + constexpr uint32_t S2_BASE_SIZE = 128; // S2轴基本块大小 + constexpr uint32_t D_SIZE = 576; + auto ascendcPlatform = platform_ascendc::PlatformAscendC(qsfaInfo_->platformInfo); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + if (qsfaInfo_->gSize > 64) { // G大于64时核间切G，V0结果出核，需申请GM空间 + workspaceSize_ += (S2_BASE_SIZE * D_SIZE * GetTypeSize(qsfaInfo_->inputQType) \ + * TRIPLE_BUFFER_NUM * (aicNum >> 1)); + } + } else { + uint32_t mmResElemSize = 4; // 4:fp32 + uint32_t vec1ResElemSize = 2; // 2:fp16/bf16 + uint32_t bmm2ResElemSize = 4; // 4:fp32 + uint32_t qPreProcResElemSize = 0; // 普通场景不涉及Q预处理 + uint32_t softmaxSumElemSize = 4; // 4:int32 + float kvDtypeRatio = 1.0; + + workspaceSize_ = libapiSize_; + uint32_t preLoadNum = 1; + preLoadNum = PRE_LOAD_NUM; + + workspaceSize_ += preLoadNum * (mmResUbSize_ * actCoreNum * mmResElemSize); + workspaceSize_ += preLoadNum * static_cast(static_cast( + mmResUbSize_ * actCoreNum * vec1ResElemSize) * kvDtypeRatio); + workspaceSize_ += preLoadNum * bmm2ResUbSize_ * actCoreNum * bmm2ResElemSize; + workspaceSize_ += preLoadNum * static_cast(static_cast( + qPreSizeMla_ * actCoreNum * qPreProcResElemSize) * kvDtypeRatio); + workspaceSize_ += preLoadNum * mBaseSize_ * actCoreNum * softmaxSumElemSize; + workspaceSize_ += preLoadNum * bmm2ResUbSize_ * actCoreNum * bmm2ResElemSize; // vec2ResGm + // topk BlkSize == 1场景, 需要额外空间缓存离散聚合的值 + // bufNum s2Base D dRope sizeOf(half) + // 4:bufNum 512:s2Base 512:D 64:dRope 2:sizeOf(half) + workspaceSize_ += 4 * 512 * (512 + 64) * 2 * actCoreNum; + // 缓存有效mte2 size的长度 份数 512B对齐的长度 sizeof(int32_t) aiv核数 + workspaceSize_ += 4 * 128 * 4 * (2 * actCoreNum); // 4:缓存有效mte2 size的长度 128:份数 4:512B对齐的长度 2:aiv核数 + } + + CalcFDWorkSpace(actCoreNum); +} + +void QSFAMlaTiling::CalcBlockDim() +{ + auto ascendcPlatform = platform_ascendc::PlatformAscendC(qsfaInfo_->platformInfo); + auto aicNum = usedCoreNum_; + auto aivNum = 2 * usedCoreNum_; + + blockDim_ = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum); + OP_LOGI(qsfaInfo_->opName, "QSFA block dim: %u aiv Num: %u aic Num: %u.", blockDim_, aivNum, aicNum); +} + +ge::graphStatus QSFAMlaTiling::DoOpTiling(QSFATilingInfo *qsfaInfo) +{ + qsfaInfo_ = qsfaInfo; + if (GetPlatformInfo() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + InitParams(); + Split(); + FillTiling(); + CalcBlockDim(); + GetWorkspaceSize(); + GenTilingKey(); + + if ((SetBlockDim(blockDim_) != ge::GRAPH_SUCCESS) || + (SetTilingKey(tilingKey_) != ge::GRAPH_SUCCESS) || + (SetWorkspaceSize(workspaceSize_) != ge::GRAPH_SUCCESS) || + (SetTilingData(tilingData_) != ge::GRAPH_SUCCESS)) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingKvQuantSparseFlashAttention(gert::TilingContext *context) +{ + QSFATilingInfo qsfaInfo; + QSFAInfoParser qsfaInfoParser(context); + if (qsfaInfoParser.Parse(qsfaInfo) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + QSFATilingCheck tilingChecker(qsfaInfo); + if (tilingChecker.Process() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + QSFAMlaTiling tiling(context); + return tiling.DoOpTiling(&qsfaInfo); +} + +ge::graphStatus TilingPrepareForKvQuantSparseFlashAttention(gert::TilingParseContext* const context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::GetExpectedShape(gert::Shape &shapeExpected, + const QSFATilingShapeCompareParam ¶m, const QSFALayout &layout) const +{ + if (layout == QSFALayout::BSND) { + shapeExpected = gert::Shape({param.B, param.S, param.N, param.D}); + } else if (layout == QSFALayout::TND) { + shapeExpected = gert::Shape({param.T, param.N, param.D}); + } else if (layout == QSFALayout::PA_BSND) { + shapeExpected = gert::Shape({param.Bn, param.Bs, param.N, param.D}); + } else { + OP_LOGE(opName_, "layout %s is unsupported", QSFALayoutToSerialString(layout).c_str()); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CompareShape(QSFATilingShapeCompareParam ¶m, + const gert::Shape &shape, const QSFALayout &layout, const std::string &name) const +{ + gert::Shape qsfaShapeExpected; + if (GetExpectedShape(qsfaShapeExpected, param, layout) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + if (shape.GetDimNum() != qsfaShapeExpected.GetDimNum()) { + OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(opName_, name.c_str(), + std::to_string(shape.GetDimNum()).c_str(), + "The shape dim of " + name + " must be " + std::to_string(qsfaShapeExpected.GetDimNum())); + + return ge::GRAPH_FAILED; + } + + for (size_t i = 0; i < shape.GetDimNum(); i++) { + if (shape.GetDim(i) != qsfaShapeExpected.GetDim(i)) { + OP_LOGE_FOR_INVALID_SHAPE(opName_, name.c_str(), + GetShapeStr(shape).c_str(), GetShapeStr(qsfaShapeExpected).c_str()); + return ge::GRAPH_FAILED; + } + } + + return ge::GRAPH_SUCCESS; +} + +void QSFATilingCheck::LogErrorDtypeSupport(const std::vector &expectDtypeList, + const ge::DataType &actualDtype, const std::string &name) const +{ + std::ostringstream qsfaOss; + for (size_t i = 0; i < expectDtypeList.size(); ++i) { + qsfaOss << QSFADataTypeToSerialString(expectDtypeList[i]); + if (i < expectDtypeList.size() - 1) { + qsfaOss << ", "; + } + } + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(opName_, name.c_str(), + QSFADataTypeToSerialString(actualDtype).c_str(), + "The dtype of " + name + " must be " + qsfaOss.str()); +} + +ge::graphStatus QSFATilingCheck::CheckDtypeSupport(const gert::CompileTimeTensorDesc *qsfaDesc, + const std::string &name) const +{ + if (qsfaDesc != nullptr) { + const auto& qsfaIt = DTYPE_SUPPORT_MAP.find(name); + OP_CHECK_IF(qsfaIt == DTYPE_SUPPORT_MAP.end(), + OP_LOGE(opName_, "%s datatype support list should be specify in DTYPE_SUPPORT_MAP", name.c_str()), + return ge::GRAPH_FAILED); + auto &qsfaExpectDtypeList = qsfaIt->second; + OP_CHECK_IF(std::find( + qsfaExpectDtypeList.begin(), qsfaExpectDtypeList.end(), + qsfaDesc->GetDataType()) == qsfaExpectDtypeList.end(), + LogErrorDtypeSupport(qsfaExpectDtypeList, qsfaDesc->GetDataType(), name), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +template +void QSFATilingCheck::LogErrorNumberSupport(const std::vector &expectNumberList, + const T &actualValue, const std::string &name, const std::string subName) const +{ + std::ostringstream qsfaOssNum; + for (size_t i = 0; i < expectNumberList.size(); ++i) { + qsfaOssNum << std::to_string(expectNumberList[i]); + if (i < expectNumberList.size() - 1) { + qsfaOssNum << ", "; + } + } + + OP_LOGE_FOR_INVALID_VALUE(opName_, (name + " " + subName).c_str(), + std::to_string(actualValue).c_str(), qsfaOssNum.str()); +} + +template +void QSFATilingCheck::LogErrorDimNumSupport(const std::vector &expectNumberList, + const T &actualValue, const std::string &name) const +{ + LogErrorNumberSupport(expectNumberList, actualValue, name, "dimension"); +} + +ge::graphStatus QSFATilingCheck::CheckDimNumInLayoutSupport(const QSFALayout &layout, + const gert::StorageShape *shape, const std::string &name) const +{ + const auto& qsfaDimIt = QSFA_LAYOUT_DIM_MAP.find(layout); + OP_CHECK_IF(shape->GetStorageShape().GetDimNum() != qsfaDimIt->second, + OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(opName_, name.c_str(), + std::to_string(shape->GetStorageShape().GetDimNum()).c_str(), + "When layout is " + QSFALayoutToSerialString(layout) + + ", the shape dim of " + name + "should be " + std::to_string(qsfaDimIt->second)), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckDimNumSupport(const gert::StorageShape *shape, + const std::vector &qsfaExpectDimNumList, const std::string &name) const +{ + if (shape == nullptr) { + return ge::GRAPH_SUCCESS; + } + + if (std::find(qsfaExpectDimNumList.begin(), qsfaExpectDimNumList.end(), + shape->GetStorageShape().GetDimNum()) == qsfaExpectDimNumList.end()) { + LogErrorDimNumSupport(qsfaExpectDimNumList, shape->GetStorageShape().GetDimNum(), name); + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +void QSFATilingCheck::LogErrorLayoutSupport(const std::vector &expectLayoutList, + const QSFALayout &actualLayout, const std::string &name) const +{ + std::ostringstream qsfaOssLayout; + for (size_t i = 0; i < expectLayoutList.size(); ++i) { + qsfaOssLayout << QSFALayoutToSerialString(expectLayoutList[i]); + if (i < expectLayoutList.size() - 1) { + qsfaOssLayout << ", "; + } + } + OP_LOGE_FOR_INVALID_FORMAT(opName_, name.c_str(), + QSFALayoutToSerialString(actualLayout).c_str(), qsfaOssLayout.str().c_str()); +} + +ge::graphStatus QSFATilingCheck::CheckLayoutSupport(const QSFALayout &actualLayout, const std::string &name) const +{ + const auto& qsfaItLayout = LAYOUT_SUPPORT_MAP.find(name); + OP_CHECK_IF(qsfaItLayout == LAYOUT_SUPPORT_MAP.end(), + OP_LOGE(opName_, "%s layout support list should be specify in LAYOUT_SUPPORT_MAP", name.c_str()), + return ge::GRAPH_FAILED); + auto &qsfaExpectLayoutList = qsfaItLayout->second; + OP_CHECK_IF(std::find( + qsfaExpectLayoutList.begin(), qsfaExpectLayoutList.end(), actualLayout) == qsfaExpectLayoutList.end(), + LogErrorLayoutSupport(qsfaExpectLayoutList, actualLayout, name), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckSingleParaQuery() const +{ + const std::vector qsfaQueryDimNumList = {DIM_NUM_THREE, DIM_NUM_FOUR}; + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.query.desc, QUERY_NAME) || + ge::GRAPH_SUCCESS != CheckLayoutSupport(qLayout_, QUERY_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(opParamInfo_.query.shape, qsfaQueryDimNumList, QUERY_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(qLayout_, opParamInfo_.query.shape, QUERY_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckSingleParaKey() const +{ + const std::vector qsfaKeyDimNumList = {DIM_NUM_THREE, DIM_NUM_FOUR}; + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.key.desc, KEY_NAME) || + ge::GRAPH_SUCCESS != CheckLayoutSupport(kvLayout_, KEY_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(opParamInfo_.key.shape, qsfaKeyDimNumList, KEY_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(kvLayout_, opParamInfo_.key.shape, KEY_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckSingleParaNumHeads() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckSingleParaKvHeadNums() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckSingleParaSparseMode() const +{ + OP_CHECK_IF((*opParamInfo_.sparseMode != 3 && *opParamInfo_.sparseMode != 0), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "sparseMode", + std::to_string(*opParamInfo_.sparseMode).c_str(), "sparseMode must be 0 or 3."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckSingleParaSparseBlockSize() const +{ + OP_CHECK_IF((npuArch_ == NpuArch::DAV_2201) && + ((*opParamInfo_.sparseBlockSize <= 0 || *opParamInfo_.sparseBlockSize > 16) || + (static_cast(*opParamInfo_.sparseBlockSize) & + static_cast(*opParamInfo_.sparseBlockSize - 1L)) != 0UL), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "sparseBlockSize", + std::to_string(*opParamInfo_.sparseBlockSize).c_str(), + "sparseBlockSize must be in range [1, 16] and be a power of 2."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckSingleParaSparseIndices() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.sparseIndices.desc, SPARSE_INDICES_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckSinglePara() const +{ + if (ge::GRAPH_SUCCESS != CheckSingleParaQuery() || + ge::GRAPH_SUCCESS != CheckSingleParaKey() || + ge::GRAPH_SUCCESS != CheckSingleParaSparseIndices() || + ge::GRAPH_SUCCESS != CheckSingleParaNumHeads() || + ge::GRAPH_SUCCESS != CheckSingleParaKvHeadNums() || + ge::GRAPH_SUCCESS != CheckSingleParaSparseMode() || + ge::GRAPH_SUCCESS != CheckSingleParaSparseBlockSize()) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckDequantScaleNotExistence() +{ + if (quantScaleRepoMode_ == 1) { + OP_CHECK_IF((opParamInfo_.keyDequantScale.tensor == nullptr || + opParamInfo_.valueDequantScale.tensor == nullptr), + OP_LOGE_WITH_INVALID_INPUT(opName_, "key_dequant_scale and value_dequant_scale"), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +template +ge::graphStatus QSFATilingCheck::CheckAttrValueByMap(std::map> &attrMap) const +{ + for (auto const &kv : attrMap) { + const std::string &qsfaAttrName = kv.first; + const std::pair &qsfaPointerValuePair = kv.second; + if (qsfaPointerValuePair.first == nullptr) { + OP_LOGE_WITH_INVALID_INPUT(opName_, qsfaAttrName.c_str()); + return ge::GRAPH_FAILED; + } + + if (*(qsfaPointerValuePair.first) != qsfaPointerValuePair.second) { + std::ostringstream qsfaOssExpect; + qsfaOssExpect << std::to_string(qsfaPointerValuePair.second); + std::ostringstream qsfaOssActual; + qsfaOssActual << std::to_string(*(qsfaPointerValuePair.first)); + OP_LOGE(opName_, + "%s value should be %s, but got %s", + qsfaAttrName.c_str(), + qsfaOssExpect.str().c_str(), + qsfaOssActual.str().c_str()); + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckParaExistenceMlaAntiquant() const +{ + if (kvLayout_ == QSFALayout::BSND) { + return ge::GRAPH_SUCCESS; + } else if (kvLayout_ == QSFALayout::TND) { + OP_CHECK_IF(opParamInfo_.actualSeqLengths.tensor == nullptr, + OP_LOGE_WITH_INVALID_INPUT(opName_, "actualSeqLengthsKv"), + return ge::GRAPH_FAILED); + } else if (kvLayout_ == QSFALayout::PA_BSND) { + OP_CHECK_IF(opParamInfo_.actualSeqLengths.tensor == nullptr, + OP_LOGE_WITH_INVALID_INPUT(opName_, "actualSeqLengthsKv"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.blockTable.tensor == nullptr, + OP_LOGE_WITH_INVALID_INPUT(opName_, "blockTable"), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckParaExistenceMla() const +{ + return CheckParaExistenceMlaAntiquant(); +} + +ge::graphStatus QSFATilingCheck::CheckParaExistence() +{ + if (ge::GRAPH_SUCCESS != CheckDequantScaleNotExistence()) { + return ge::GRAPH_FAILED; + } + + return CheckParaExistenceMla(); +} + +static ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + const std::string &name, const char *opName) +{ + if (tensor == nullptr) { + OP_LOGE_WITH_INVALID_INPUT(opName, name.c_str()); + return ge::GRAPH_FAILED; + } + int64_t qsfaShapeSize = tensor->GetShapeSize(); + if (qsfaShapeSize <= 0) { + OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON(opName, name.c_str(), + std::to_string(qsfaShapeSize).c_str(), + "The shape size of " + name + " should be greater than 0."); + return ge::GRAPH_FAILED; + } + size = static_cast(qsfaShapeSize); + return ge::GRAPH_SUCCESS; +} + +void QSFATilingCheck::SetQSFAShapeCompare() +{ + queryShapeCmp_ = opParamInfo_.query.shape->GetStorageShape(); + topkShapeCmp_ = opParamInfo_.sparseIndices.shape->GetStorageShape(); + keyShapeCmp_ = opParamInfo_.key.shape->GetStorageShape(); + valueShapeCmp_ = opParamInfo_.value.shape->GetStorageShape(); + attenOutShapeCmp_ = opParamInfo_.attenOut.shape->GetStorageShape(); +} + +ge::graphStatus QSFATilingCheck::CheckBlockTable() const +{ + if (kvStorageMode_ != KvStorageMode::PAGE_ATTENTION) { + OP_CHECK_IF(opParamInfo_.blockTable.tensor != nullptr, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(opName_, BLOCK_TABLE_NAME.c_str(), + Ops::Base::ToString(opParamInfo_.blockTable.tensor->GetStorageShape()).c_str(), + "When the layout_kv is " + QSFALayoutToSerialString(kvLayout_) + ", block_table should be null."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; + } + + uint32_t blockTableBatch = opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0); + OP_CHECK_IF(blockTableBatch != bSize_, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(opName_, BLOCK_TABLE_NAME.c_str(), + Ops::Base::ToString(opParamInfo_.blockTable.tensor->GetStorageShape()).c_str(), + "The first dim of " + BLOCK_TABLE_NAME + " should be equal to batch size " + std::to_string(bSize_)), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckDTypeConsistency(const ge::DataType &actualDtype, + const ge::DataType &expectDtype, const std::string &name) const +{ + if (actualDtype != expectDtype) { + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(opName_, name.c_str(), + QSFADataTypeToSerialString(actualDtype).c_str(), + "The dtype of " + name + " must be " + QSFADataTypeToSerialString(expectDtype)); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckTopkShape() +{ + QSFATilingShapeCompareParam qsfaShapeParams; + qsfaShapeParams.B = bSize_; + qsfaShapeParams.N = n2Size_; + qsfaShapeParams.S = s1Size_; + qsfaShapeParams.D = sparseBlockCount_; + qsfaShapeParams.T = qTSize_; + return CompareShape(qsfaShapeParams, topkShapeCmp_, topkLayout_, SPARSE_INDICES_NAME); +} + +ge::graphStatus QSFATilingCheck::CheckAttenOutShape() +{ + QSFATilingShapeCompareParam shapeParams; + shapeParams.B = bSize_; + shapeParams.N = n1Size_; + shapeParams.S = s1Size_; + shapeParams.D = 512; // 512:输出的head_dim + shapeParams.T = qTSize_; + if (CompareShape(shapeParams, attenOutShapeCmp_, outLayout_, ATTEN_OUT_NAME) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckAttenOut() +{ + if (ge::GRAPH_SUCCESS != CheckDTypeConsistency(opParamInfo_.attenOut.desc->GetDataType(), + inputQType_, ATTEN_OUT_NAME) || + ge::GRAPH_SUCCESS != CheckAttenOutShape()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckTopK() +{ + if (ge::GRAPH_SUCCESS != CheckTopkShape()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckKVShapeForBatchContinuous() +{ + QSFATilingShapeCompareParam shapeParams; + shapeParams.B = bSize_; + shapeParams.N = n2Size_; + shapeParams.S = s2Size_; + shapeParams.D = vHeadDim_; + shapeParams.T = kvTSize_; + if (CompareShape(shapeParams, valueShapeCmp_, kvLayout_, VALUE_NAME) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckKVShapeForPageAttention() +{ + int64_t blockNum = keyShapeCmp_.GetDim(0); + QSFATilingShapeCompareParam shapeParams; + shapeParams.Bn = blockNum; + shapeParams.N = n2Size_; + shapeParams.Bs = blockSize_; + shapeParams.T = kvTSize_; + shapeParams.D = vHeadDim_; + if (CompareShape(shapeParams, valueShapeCmp_, kvLayout_, VALUE_NAME) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckKVShape() +{ + if (kvStorageMode_ == KvStorageMode::BATCH_CONTINUOUS) { + return CheckKVShapeForBatchContinuous(); + } + + if (kvStorageMode_ == KvStorageMode::PAGE_ATTENTION) { + return CheckKVShapeForPageAttention(); + } + + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "kvStorageMode_", + std::to_string(static_cast(kvStorageMode_)).c_str(), + "storage mode of key and value must be BATCH_CONTINUOUS or PAGE_ATTENTION."); + return ge::GRAPH_FAILED; +} + +ge::graphStatus QSFATilingCheck::CheckKV() +{ + if (ge::GRAPH_SUCCESS != CheckDTypeConsistency(opParamInfo_.value.desc->GetDataType(), + inputKvType_, VALUE_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckActualSeqLensQ() +{ + if (opParamInfo_.actualSeqLengthsQ.tensor == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (ge::GRAPH_SUCCESS != CheckActualSeqLensQDType() || + ge::GRAPH_SUCCESS != CheckActualSeqLensQShape()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckActualSeqLensQDType() +{ + if (opParamInfo_.actualSeqLengthsQ.desc == nullptr) { + OP_LOGE_WITH_INVALID_INPUT(opName_, "actualSeqLengthsQ's dtype"); + return ge::GRAPH_FAILED; + } + + if (opParamInfo_.actualSeqLengthsQ.desc->GetDataType() != ge::DT_INT32) { + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(opName_, "actualSeqLengthsQ", + QSFADataTypeToSerialString(opParamInfo_.actualSeqLengthsQ.desc->GetDataType()).c_str(), + "The dtype of actualSeqLengthsQ must be DT_INT32."); + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckActualSeqLensQShape() +{ + uint32_t qsfaShapeSize = 0; + if (GetActualSeqLenSize(qsfaShapeSize, opParamInfo_.actualSeqLengthsQ.tensor, + "actualSeqLengthsQ", opName_) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + if (qsfaShapeSize != bSize_) { + OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON(opName_, "actualSeqLengthsQ", + std::to_string(qsfaShapeSize).c_str(), + "The shape size of actualSeqLengthsQ should be equal to batch size " + std::to_string(bSize_)); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckActualSeqLens() +{ + if (ge::GRAPH_SUCCESS != CheckActualSeqLensDType() || + ge::GRAPH_SUCCESS != CheckActualSeqLensShape()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckActualSeqLensDType() +{ + if (opParamInfo_.actualSeqLengths.tensor == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (opParamInfo_.actualSeqLengths.desc == nullptr) { + OP_LOGE_WITH_INVALID_INPUT(opName_, "actualSeqLengths's dtype"); + return ge::GRAPH_FAILED; + } + if (opParamInfo_.actualSeqLengths.desc->GetDataType() != ge::DT_INT32) { + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(opName_, "actualSeqLengths", + QSFADataTypeToSerialString(opParamInfo_.actualSeqLengthsQ.desc->GetDataType()).c_str(), + "The dtype of actualSeqLengths must be DT_INT32."); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckActualSeqLensShape() +{ + if (opParamInfo_.actualSeqLengths.tensor == nullptr) { + return ge::GRAPH_SUCCESS; + } + uint32_t qsfaShapeSizeKv = 0; + if (GetActualSeqLenSize(qsfaShapeSizeKv, opParamInfo_.actualSeqLengths.tensor, + "actualSeqLengths", opName_) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + if (qsfaShapeSizeKv != bSize_) { + OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON(opName_, "actualSeqLengths", + std::to_string(qsfaShapeSizeKv).c_str(), + "The shape size of actualSeqLengths should be equal to batch size " + std::to_string(bSize_)); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckMultiParaConsistency() +{ + SetQSFAShapeCompare(); + if (ge::GRAPH_SUCCESS != CheckKV() || + ge::GRAPH_SUCCESS != CheckTopK() || + ge::GRAPH_SUCCESS != CheckAttenOut() || + ge::GRAPH_SUCCESS != CheckActualSeqLensQ() || + ge::GRAPH_SUCCESS != CheckActualSeqLens() || + ge::GRAPH_SUCCESS != CheckBlockTable()) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckFeatureMlaAntiquantShape() const +{ + if (ge::GRAPH_SUCCESS != CheckFeatureMlaAntiquantShapeSizes() || + ge::GRAPH_SUCCESS != CheckFeatureMlaAntiquantShapeSparseAndHeadDim()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckFeatureMlaAntiquantShapeSizes() const +{ + OP_CHECK_IF(bSize_ <= 0, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "batch_size", std::to_string(bSize_).c_str(), + "batch_size should be greater than 0."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(qTSize_ <= 0 && (qLayout_ == QSFALayout::TND), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "T_size of query", std::to_string(qTSize_).c_str(), + "T_size of query should be greater than 0."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(n1Size_ <= 0, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "q_head_num", std::to_string(n1Size_).c_str(), + "q_head_num should be greater than 0."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(n2Size_ != 1, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "kv_head_num", std::to_string(n2Size_).c_str(), + "kv_head_num should be 1."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(n1Size_ % n2Size_ != 0, + OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(opName_, "q_head_num and kv_head_num", + std::to_string(n1Size_) + " and " + std::to_string(n2Size_), + "q_head_num must be divisible by kv_head_num."), + return ge::GRAPH_FAILED); + + if (isA5_) { + std::vector gSizeSupportList = {1, 2, 4, 8, 16, 32, 48, 64, 128}; + OP_CHECK_IF(std::find(gSizeSupportList.begin(), gSizeSupportList.end(), gSize_) == gSizeSupportList.end(), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "group num", std::to_string(gSize_).c_str(), + "group num should be in 1, 2, 4, 8, 16, 32, 48, 64, 128."), + return ge::GRAPH_FAILED); + } else { + std::vector gSizeSupportList = {1, 2, 4, 8, 16, 32, 64, 128}; + OP_CHECK_IF(std::find(gSizeSupportList.begin(), gSizeSupportList.end(), gSize_) == gSizeSupportList.end(), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "group num", std::to_string(gSize_).c_str(), + "group num should be in 1, 2, 4, 8, 16, 32, 64, 128."), + return ge::GRAPH_FAILED); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckFeatureMlaAntiquantShapeSparseAndHeadDim() const +{ + if (isA5_) { + if (inputKvType_ == ge::DT_HIFLOAT8) { + OP_CHECK_IF(sparseBlockCount_ != 2048, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "sparse block count", + std::to_string(sparseBlockCount_).c_str(), + "when key and value dtype use hifloat8, sparse block count must be 2048."), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(sparseBlockSize_ != 1, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "sparse block size", + std::to_string(sparseBlockSize_).c_str(), "sparse block size must be 1"), + return ge::GRAPH_FAILED); + } else { + std::vector sparseBlockSizeSupportList = {1, 2, 4, 8, 16}; + OP_CHECK_IF(std::find(sparseBlockSizeSupportList.begin(), sparseBlockSizeSupportList.end(), + sparseBlockSize_) == sparseBlockSizeSupportList.end(), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "sparse block size", + std::to_string(sparseBlockSize_).c_str(), + "group num should be in 1, 2, 4, 8, 16."), + return ge::GRAPH_FAILED); + } + + OP_CHECK_IF(qHeadDim_ != 576, // 576:当前不泛化 + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "qHeadDim_", + std::to_string(qHeadDim_).c_str(), "q_head_dim only support 576."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(kHeadDim_ != 656, // 656:当前不泛化 + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "kHeadDim_", + std::to_string(kHeadDim_).c_str(), "k_head_dim only support 656."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckFeatureMlaAntiquantLayout() const +{ + const std::vector qsfaLayoutSupportList = { + "BSND", + "TND" + }; + std::string layoutQuery = opParamInfo_.layoutQuery; + OP_CHECK_IF(std::find(qsfaLayoutSupportList.begin(), + qsfaLayoutSupportList.end(), layoutQuery) == qsfaLayoutSupportList.end(), + OP_LOGE_FOR_INVALID_FORMAT(opName_, "query", layoutQuery.c_str(), "BSND or TND"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckFeatureMlaAntiquantDtype() const +{ + OP_CHECK_IF(inputQType_ != ge::DT_BF16 && inputQType_ != ge::DT_FLOAT16, + OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(opName_, "query", + QSFADataTypeToSerialString(inputQType_).c_str(), + "The dtype of query must be " + QSFADataTypeToSerialString(ge::DT_BF16) + + " and " + QSFADataTypeToSerialString(ge::DT_FLOAT16)), + return ge::GRAPH_FAILED); + + if (isA5_) { + OP_CHECK_IF(inputKvType_ != ge::DT_FLOAT8_E4M3FN && + inputKvType_ != ge::DT_HIFLOAT8 && + inputKvType_ != ge::DT_INT8, + OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(opName_, "key and value", + QSFADataTypeToSerialString(inputKvType_).c_str(), + "The dtype of key and value must be " + + QSFADataTypeToSerialString(ge::DT_FLOAT8_E4M3FN) + ", " + + QSFADataTypeToSerialString(ge::DT_HIFLOAT8) + ", " + + QSFADataTypeToSerialString(ge::DT_INT8)), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF(inputKvType_ != ge::DT_INT8, + OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(opName_, "key and value", + QSFADataTypeToSerialString(inputKvType_).c_str(), + "The dtype of key and value must be " + QSFADataTypeToSerialString(ge::DT_INT8)), + return ge::GRAPH_FAILED); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckFeatureMlaAntiquantAttr() const +{ + OP_CHECK_IF(attentionMode_ != 2, // 2:MLA-absorb + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "attention_mode", + std::to_string(attentionMode_).c_str(), "attention_mode should be 2(MLA-absorb)."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(keyQuantMode_ != 2, // 2:per-tile + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "key_quant_mode", + std::to_string(keyQuantMode_).c_str(), "key_quant_mode should be 2(per-tile)."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(valueQuantMode_ != 2, // 2:per-tile + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "value_quant_mode", + std::to_string(valueQuantMode_).c_str(), "value_quant_mode should be 2(per-tile)."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(quantScaleRepoMode_ != 1, // 1:combine + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "quant_scale_repo_mode", + std::to_string(quantScaleRepoMode_).c_str(), "quant_scale_repo_mode should be 1(combine)."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(preTokens_ != INT64_MAX, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "preTokens_", + std::to_string(preTokens_).c_str(), "preTokens_ should be INT64_MAX."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(nextTokens_ != INT64_MAX, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "nextTokens_", + std::to_string(nextTokens_).c_str(), "nextTokens_ should be INT64_MAX."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(tileSize_ != 128, // 128:当前不泛化 + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "tile_size", + std::to_string(tileSize_).c_str(), "tile_size should be 128."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(ropeHeadDim_ != 64, // 64:当前不泛化 + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "rope", + std::to_string(ropeHeadDim_).c_str(), "rope_head_dim should be 64."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckFeatureMlaAntiquantPa() const +{ + if (kvStorageMode_ != KvStorageMode::PAGE_ATTENTION) { + return ge::GRAPH_SUCCESS; + } + + OP_CHECK_IF(blockSize_ <= 0 || blockSize_ > static_cast(MAX_BLOCK_SIZE), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "block_size", std::to_string(blockSize_).c_str(), + "when page attention is enabled, block_size should be in range (0, " + + std::to_string(MAX_BLOCK_SIZE) + "]."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(blockSize_ % 16 > 0, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "block_size", std::to_string(blockSize_).c_str(), + "when page attention is enabled, block_size should be 16-aligned."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(blockSize_ % sparseBlockSize_ > 0, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opName_, "block_size", std::to_string(blockSize_).c_str(), + "when page attention is enabled, " + "block_size must be divided by sparse_block_size, but now the remainder is " + + std::to_string(blockSize_ % sparseBlockSize_)), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckFeatureMlaAntiquant() const +{ + if (ge::GRAPH_SUCCESS != CheckFeatureMlaAntiquantAttr() || + ge::GRAPH_SUCCESS != CheckFeatureMlaAntiquantShape() || + ge::GRAPH_SUCCESS != CheckFeatureMlaAntiquantLayout() || + ge::GRAPH_SUCCESS != CheckFeatureMlaAntiquantDtype() || + ge::GRAPH_SUCCESS != CheckFeatureMlaAntiquantPa()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFATilingCheck::CheckFeatureMla() const +{ + return CheckFeatureMlaAntiquant(); +} + +ge::graphStatus QSFATilingCheck::CheckFeature() const +{ + return CheckFeatureMla(); +} + +void QSFATilingCheck::Init() +{ + opName_ = qsfaInfo_.opName; + platformInfo_ = qsfaInfo_.platformInfo; + opParamInfo_ = qsfaInfo_.opParamInfo; + npuArch_ = qsfaInfo_.npuArch; + isA5_ = qsfaInfo_.isA5; + + bSize_ = qsfaInfo_.bSize; + n1Size_ = qsfaInfo_.n1Size; + n2Size_ = qsfaInfo_.n2Size; + s1Size_ = qsfaInfo_.s1Size; + s2Size_ = qsfaInfo_.s2Size; + gSize_ = qsfaInfo_.gSize; + qHeadDim_ = qsfaInfo_.qHeadDim; + kHeadDim_ = qsfaInfo_.kHeadDim; + vHeadDim_ = qsfaInfo_.vHeadDim; + ropeHeadDim_ = qsfaInfo_.ropeHeadDim; + maxBlockNumPerBatch_ = qsfaInfo_.maxBlockNumPerBatch; + qTSize_ = qsfaInfo_.qTSize; + kvTSize_ = qsfaInfo_.kvTSize; + blockSize_ = qsfaInfo_.blockSize; + sparseBlockCount_ = qsfaInfo_.sparseBlockCount; + sparseBlockSize_ = qsfaInfo_.sparseBlockSize; + + attentionMode_ = qsfaInfo_.attentionMode; + keyQuantMode_ = qsfaInfo_.keyQuantMode; + valueQuantMode_ = qsfaInfo_.valueQuantMode; + quantScaleRepoMode_ = qsfaInfo_.quantScaleRepoMode; + tileSize_ = qsfaInfo_.tileSize; + preTokens_ = qsfaInfo_.preTokens; + nextTokens_ = qsfaInfo_.nextTokens; + + inputQType_ = qsfaInfo_.inputQType; + inputKvType_ = qsfaInfo_.inputKvType; + outputType_ = qsfaInfo_.outputType; + + qLayout_ = qsfaInfo_.qLayout; + topkLayout_ = qsfaInfo_.topkLayout; + kvLayout_ = qsfaInfo_.kvLayout; + outLayout_ = qsfaInfo_.outLayout; + + kvStorageMode_ = qsfaInfo_.kvStorageMode; + l2CacheSize_ = qsfaInfo_.l2CacheSize; +} + +ge::graphStatus QSFATilingCheck::Process() +{ + Init(); + if (CheckSinglePara() != ge::GRAPH_SUCCESS || + CheckParaExistence() != ge::GRAPH_SUCCESS || + CheckFeature() != ge::GRAPH_SUCCESS || + CheckMultiParaConsistency() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +static constexpr int64_t kInvalidDimValue = std::numeric_limits::min(); + +static bool HasAxis(const QSFAAxis &axis, const QSFALayout &layout, const gert::Shape &shape) +{ + const auto& qsfaLayoutIt = QSFA_LAYOUT_AXIS_MAP.find(layout); + if (qsfaLayoutIt == QSFA_LAYOUT_AXIS_MAP.end()) { + return false; + } + + const std::vector& qsfaAxes = qsfaLayoutIt->second; + const auto& qsfaAxisIt = std::find(qsfaAxes.begin(), qsfaAxes.end(), axis); + if (qsfaAxisIt == qsfaAxes.end()) { + return false; + } + + const auto& qsfaDimIt = QSFA_LAYOUT_DIM_MAP.find(layout); + if (qsfaDimIt == QSFA_LAYOUT_DIM_MAP.end() || qsfaDimIt->second != shape.GetDimNum()) { + return false; + } + + return true; +} + +static size_t GetAxisIdx(const QSFAAxis &axis, const QSFALayout &layout) +{ + const std::vector& axes = QSFA_LAYOUT_AXIS_MAP.find(layout)->second; + const auto& axisIt = std::find(axes.begin(), axes.end(), axis); + + return std::distance(axes.begin(), axisIt); +} + +static uint32_t GetAxisNum(const gert::Shape &shape, const QSFAAxis &axis, const QSFALayout &layout) +{ + return HasAxis(axis, layout, shape) ? shape.GetDim(GetAxisIdx(axis, layout)) : kInvalidDimValue; +} + +ge::graphStatus QSFAInfoParser::CheckRequiredInOutExistence() const +{ + OP_CHECK_IF(opParamInfo_.query.shape == nullptr, OP_LOGE_WITH_INVALID_INPUT(opName_, "Shape of tensor query"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.query.desc == nullptr, OP_LOGE_WITH_INVALID_INPUT(opName_, "Desc of tensor query"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key.shape == nullptr, OP_LOGE_WITH_INVALID_INPUT(opName_, "Shape of tensor k"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key.desc == nullptr, OP_LOGE_WITH_INVALID_INPUT(opName_, "Desc of tensor k"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.value.shape == nullptr, OP_LOGE_WITH_INVALID_INPUT(opName_, "Shape of tensor value"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.value.desc == nullptr, OP_LOGE_WITH_INVALID_INPUT(opName_, "Desc of tensor value"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.sparseIndices.shape == nullptr, + OP_LOGE_WITH_INVALID_INPUT(opName_, "Shape of tensor sparseIndices"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.sparseIndices.desc == nullptr, + OP_LOGE_WITH_INVALID_INPUT(opName_, "Desc of tensor sparseIndices"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.attenOut.shape == nullptr, + OP_LOGE_WITH_INVALID_INPUT(opName_, "Shape of tensor output"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.attenOut.desc == nullptr, + OP_LOGE_WITH_INVALID_INPUT(opName_, "Desc of tensor output"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::CheckRequiredAttrExistence() const +{ + OP_CHECK_IF(opParamInfo_.layoutQuery == nullptr, OP_LOGE_WITH_INVALID_INPUT(opName_, "layoutQuery"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.layoutKV == nullptr, OP_LOGE_WITH_INVALID_INPUT(opName_, "layoutKV"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.sparseBlockSize == nullptr, OP_LOGE_WITH_INVALID_INPUT(opName_, "sparseBlockSize"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.scaleValue == nullptr, OP_LOGE_WITH_INVALID_INPUT(opName_, "scaleValue"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.sparseMode == nullptr, OP_LOGE_WITH_INVALID_INPUT(opName_, "sparseMode"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::CheckRequiredParaExistence() const +{ + if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS || + CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetActualSeqLenQSize(uint32_t &size) +{ + return GetActualSeqLenSize(size, opParamInfo_.actualSeqLengthsQ.tensor, "actualSeqLengthsQ", opName_); +} + +ge::graphStatus QSFAInfoParser::GetOpName() +{ + if (context_->GetNodeName() == nullptr) { + OP_LOGE_WITH_INVALID_INPUT("KvQuantSparseFlashAttention", "opName"); + return ge::GRAPH_FAILED; + } + opName_ = context_->GetNodeName(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetNpuInfo() +{ + platformInfo_ = context_->GetPlatformInfo(); + OP_CHECK_IF(platformInfo_ == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(opName_, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED); + + auto qsfaAscendcPlat = platform_ascendc::PlatformAscendC(platformInfo_); + uint32_t qsfaAivNum = qsfaAscendcPlat.GetCoreNumAiv(); + uint32_t qsfaAicNum = qsfaAscendcPlat.GetCoreNumAic(); + OP_CHECK_IF(qsfaAicNum == 0 || qsfaAivNum == 0, + OPS_REPORT_VECTOR_INNER_ERR(opName_, "num of core obtained is 0."), return GRAPH_FAILED); + + npuArch_ = qsfaAscendcPlat.GetCurNpuArch(); + isA5_ = (npuArch_ == NpuArch::DAV_3510); + if (npuArch_ != NpuArch::DAV_2201 && npuArch_ != NpuArch::DAV_3510) { + OPS_REPORT_VECTOR_INNER_ERR(opName_, "Npu Arch Version[%d] is not support.", static_cast(npuArch_)); + return GRAPH_FAILED; + } + + qsfaAscendcPlat.GetCoreMemSize(platform_ascendc::CoreMemType::L2, l2CacheSize_); + + return ge::GRAPH_SUCCESS; +} + +void QSFAInfoParser::GetOptionalInputParaInfo() +{ + opParamInfo_.blockTable.tensor = context_->GetOptionalInputTensor(BLOCK_TABLE_INPUT_INDEX); + opParamInfo_.actualSeqLengthsQ.tensor = context_->GetOptionalInputTensor(ACT_SEQ_LEN_Q_INPUT_INDEX); + opParamInfo_.actualSeqLengthsQ.desc = context_->GetOptionalInputDesc(ACT_SEQ_LEN_Q_INPUT_INDEX); + opParamInfo_.actualSeqLengths.tensor = context_->GetOptionalInputTensor(ACT_SEQ_LEN_KV_INPUT_INDEX); + opParamInfo_.actualSeqLengths.desc = context_->GetOptionalInputDesc(ACT_SEQ_LEN_KV_INPUT_INDEX); + opParamInfo_.keyDequantScale.tensor = context_->GetOptionalInputTensor(KEY_DEQUANT_SCALE_INPUT_INDEX); + opParamInfo_.valueDequantScale.tensor = context_->GetOptionalInputTensor(VALUE_DEQUANT_SCALE_INPUT_INDEX); +} + +void QSFAInfoParser::GetInputParaInfo() +{ + opParamInfo_.query.desc = context_->GetInputDesc(QUERY_INPUT_INDEX); + opParamInfo_.query.shape = context_->GetInputShape(QUERY_INPUT_INDEX); + opParamInfo_.key.desc = context_->GetInputDesc(KEY_INPUT_INDEX); + opParamInfo_.key.shape = context_->GetInputShape(KEY_INPUT_INDEX); + opParamInfo_.value.desc = context_->GetInputDesc(VALUE_INPUT_INDEX); + opParamInfo_.value.shape = context_->GetInputShape(VALUE_INPUT_INDEX); + opParamInfo_.sparseIndices.desc = context_->GetInputDesc(SPARSE_INDICES_INPUT_INDEX); + opParamInfo_.sparseIndices.shape = context_->GetInputShape(SPARSE_INDICES_INPUT_INDEX); + GetOptionalInputParaInfo(); +} + +void QSFAInfoParser::GetOutputParaInfo() +{ + opParamInfo_.attenOut.desc = context_->GetOutputDesc(OUTPUT_INDEX); + opParamInfo_.attenOut.shape = context_->GetOutputShape(OUTPUT_INDEX); +} + +ge::graphStatus QSFAInfoParser::GetAttrParaInfo() +{ + auto attrs = context_->GetAttrs(); + OP_CHECK_IF(attrs == nullptr, OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "attrs got from ge is nullptr"), + return ge::GRAPH_FAILED); + + opParamInfo_.layoutQuery = attrs->GetStr(LAYOUT_QUERY_ATTR_INDEX); + opParamInfo_.layoutKV = attrs->GetStr(LAYOUT_KV_ATTR_INDEX); + opParamInfo_.sparseBlockSize = attrs->GetAttrPointer(SPARSE_BLOCK_SIZE_ATTR_INDEX); + opParamInfo_.scaleValue = attrs->GetAttrPointer(SCALE_VALUE_ATTR_INDEX); + opParamInfo_.sparseMode = attrs->GetAttrPointer(SPARSE_MODE_ATTR_INDEX); + opParamInfo_.keyQuantMode = attrs->GetAttrPointer(KEY_QUANT_MODE_ATTR_INDEX); + opParamInfo_.valueQuantMode = attrs->GetAttrPointer(VALUE_QUANT_MODE_ATTR_INDEX); + opParamInfo_.attentionMode = attrs->GetAttrPointer(ATTENTION_MODE_ATTR_INDEX); + opParamInfo_.preTokens = attrs->GetAttrPointer(PRE_TOKENS_ATTR_INDEX); + opParamInfo_.nextTokens = attrs->GetAttrPointer(NEXT_TOKENS_ATTR_INDEX); + opParamInfo_.quantScaleRepoMode = attrs->GetAttrPointer(QUANT_SCALE_REPO_MODE_ATTR_INDEX); + opParamInfo_.tileSize = attrs->GetAttrPointer(TILE_SIZE_ATTR_INDEX); + opParamInfo_.ropeHeadDim = attrs->GetAttrPointer(ROPE_HEAD_DIM_ATTR_INDEX); + opParamInfo_.returnSoftmaxLse = attrs->GetAttrPointer(RETURN_SOFTMAX_LSE_ATTR_INDEX); + OP_LOGE("ccccc"," ============ tiling入口 ============"); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetOpParaInfo() +{ + GetInputParaInfo(); + GetOutputParaInfo(); + if (ge::GRAPH_SUCCESS != GetAttrParaInfo()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetInOutDataType() +{ + inputQType_ = opParamInfo_.query.desc->GetDataType(); + inputKvType_ = opParamInfo_.key.desc->GetDataType(); + outputType_ = opParamInfo_.attenOut.desc->GetDataType(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetBatchSize() +{ + // 获取B基准值 + // 1、非TND时, 以query的batch_size维度为基准; + // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小 + if (qLayout_ == QSFALayout::TND) { + return GetActualSeqLenQSize(bSize_); + } else { // BSND + bSize_ = GetAxisNum(queryShape_, QSFAAxis::B, qLayout_); + return ge::GRAPH_SUCCESS; + } +} + +ge::graphStatus QSFAInfoParser::GetQTSize() +{ + // 获取query的T基准值 + // 1、非TND时, 以query的batch_size维度为基准; + // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小 + qTSize_ = (qLayout_ == QSFALayout::TND) ? GetAxisNum(queryShape_, QSFAAxis::T, qLayout_) : 0; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetKVTSize() +{ + // 获取query的T基准值 + // 1、非TND时, 以key的batch_size维度为基准; + // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小 + kvTSize_ = (kvLayout_ == QSFALayout::TND) ? GetAxisNum(keyShape_, QSFAAxis::T, kvLayout_) : 0; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetQHeadDim() +{ + // 获取qHeadDim基准值 + // 以query的D维度为基准 + qHeadDim_ = GetAxisNum(queryShape_, QSFAAxis::D, qLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetKHeadDim() +{ + // 获取kHeadDim基准值 + // 以key的D维度为基准 + kHeadDim_ = GetAxisNum(keyShape_, QSFAAxis::D, kvLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetS1Size() +{ + // 获取S1基准值 + // 1、非TND时, 以query的S维度为基准; + // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组中的最大值为基准 + if (qLayout_ == QSFALayout::TND) { + s1Size_ = GetAxisNum(queryShape_, QSFAAxis::T, qLayout_); + return ge::GRAPH_SUCCESS; + } else { // BSND + s1Size_ = GetAxisNum(queryShape_, QSFAAxis::S, qLayout_); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetKvStorageMode() +{ + if (kvLayout_ == QSFALayout::PA_BSND) { + kvStorageMode_ = KvStorageMode::PAGE_ATTENTION; + } else { + kvStorageMode_ = KvStorageMode::BATCH_CONTINUOUS; + } + // kv存储模式基准值 + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetKvLayout() +{ + const map layoutKVMap = { + {"BSND", QSFALayout::BSND}, + {"PA_BSND", QSFALayout::PA_BSND}, + {"TND", QSFALayout::TND} + }; + + std::string layout(opParamInfo_.layoutKV); + auto it = layoutKVMap.find(layout); + if (it != layoutKVMap.end()) { + kvLayout_ = it->second; + } else { + OP_LOGE_FOR_INVALID_FORMAT(opName_, "KV", + layout.c_str(), "BSND or PA_BSND or TND"); + return ge::GRAPH_FAILED; + } + if (kvLayout_ != QSFALayout::PA_BSND && qLayout_ != kvLayout_) { + OP_LOGE_FOR_INVALID_FORMATS_WITH_REASON(opName_, "KV and Q", + QSFALayoutToSerialString(kvLayout_) + " and " + QSFALayoutToSerialString(qLayout_), + "When layoutKV is not PA_BSND, layoutKV and layoutQ must be same."); + return ge::GRAPH_FAILED; + } + uint32_t keyDimNum = opParamInfo_.key.shape->GetStorageShape().GetDimNum(); + if (kvLayout_ == QSFALayout::PA_BSND && keyDimNum != 4U) { + OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(opName_, "key", + std::to_string(keyDimNum).c_str(), "When layoutKV is PA_BSND, kvDimNum must be 4"); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetS2SizeForBatchContinuous() +{ + if (kvLayout_ == QSFALayout::BSND) { // BSND + s2Size_ = GetAxisNum(keyShape_, QSFAAxis::S, kvLayout_); + } else if (kvLayout_ == QSFALayout::TND) { // TND + s2Size_ = GetAxisNum(keyShape_, QSFAAxis::T, kvLayout_); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetMaxBlockNumPerBatch() +{ + if (opParamInfo_.blockTable.tensor == nullptr) { + OP_LOGE_WITH_INVALID_INPUT(opName_, "blockTable"); + return ge::GRAPH_FAILED; + } + uint32_t qsfaDimNum = opParamInfo_.blockTable.tensor->GetStorageShape().GetDimNum(); + if (qsfaDimNum != DIM_NUM_TWO) { + OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(opName_, "block_table", + std::to_string(qsfaDimNum), + "The shape dim of block_table must be " + std::to_string(DIM_NUM_TWO)); + return ge::GRAPH_FAILED; + } + if (opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1) <= 0) { + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(opName_, "block_table", + Ops::Base::ToString(opParamInfo_.blockTable.tensor->GetStorageShape()).c_str(), + "block_table's second dim should be greater than 0."); + return ge::GRAPH_FAILED; + } + maxBlockNumPerBatch_ = opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetSparseBlockCount() +{ + sparseBlockCount_ = GetAxisNum(sparseIndicesShape_, QSFAAxis::K, qLayout_); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetBlockSize() +{ + blockSize_ = GetAxisNum(keyShape_, QSFAAxis::Bs, kvLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetS2SizeForPageAttention() +{ + if (GetMaxBlockNumPerBatch() != ge::GRAPH_SUCCESS || GetBlockSize() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + s2Size_ = maxBlockNumPerBatch_ * blockSize_; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetS2Size() +{ + // 获取S2基准值 + // 1、BATCH_CONTINUOUS时, 从key的S轴获取 + // 2、PAGE_ATTENTION时, S2 = block_table.dim1 * block_size + if (kvStorageMode_ == KvStorageMode::BATCH_CONTINUOUS) { + return GetS2SizeForBatchContinuous(); + } + return GetS2SizeForPageAttention(); +} + +ge::graphStatus QSFAInfoParser::GetValueHeadDim() +{ + // 获取vHeadDim基准值 + // 以value的D维度为基准 + vHeadDim_ = GetAxisNum(valueShape_, QSFAAxis::D, kvLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetDSizeKV() +{ + dSizeKV_ = GetAxisNum(keyShape_, QSFAAxis::D, kvLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetQueryAndOutLayout() +{ + // 获取query和attentionOut的Layout基准值 + // layoutQuery: {qLayout, outLayout} + const std::map> qsfaLayoutMap = { + {"BSND", {QSFALayout::BSND, QSFALayout::BSND}}, + {"TND", {QSFALayout::TND, QSFALayout::TND }}, + }; + + std::string qsfaLayout(opParamInfo_.layoutQuery); + auto qsfaLayoutIt = qsfaLayoutMap.find(qsfaLayout); + if (qsfaLayoutIt != qsfaLayoutMap.end()) { + qLayout_ = qsfaLayoutIt->second.first; + outLayout_ = qsfaLayoutIt->second.second; + } else { + OP_LOGE_FOR_INVALID_FORMAT(opName_, "query", qsfaLayout.c_str(), "BSND or TND"); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetTopkLayout() +{ + topkLayout_ = qLayout_; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetN1Size() +{ + n1Size_ = GetAxisNum(queryShape_, QSFAAxis::N, qLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetN2Size() +{ + n2Size_ = GetAxisNum(keyShape_, QSFAAxis::N, kvLayout_); + return ge::GRAPH_SUCCESS; +} + +void QSFAInfoParser::SetQSFAShape() +{ + queryShape_ = opParamInfo_.query.shape->GetStorageShape(); + keyShape_ = opParamInfo_.key.shape->GetStorageShape(); + + valueShape_ = opParamInfo_.value.shape->GetStorageShape(); + sparseIndicesShape_ = opParamInfo_.sparseIndices.shape->GetStorageShape(); +} + +ge::graphStatus QSFAInfoParser::GetGSize() +{ + if (n2Size_ != 0) { + gSize_ = n1Size_ / n2Size_; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetActualseqInfo() +{ + maxActualseq_ = static_cast(s2Size_); + if (opParamInfo_.actualSeqLengths.tensor != nullptr) { + actualLenDimsKV_ = opParamInfo_.actualSeqLengths.tensor->GetShapeSize(); + } + if (opParamInfo_.actualSeqLengthsQ.tensor != nullptr) { + actualLenDimsQ_ = opParamInfo_.actualSeqLengthsQ.tensor->GetShapeSize(); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QSFAInfoParser::GetShapeAndSizeInfo() +{ + SetQSFAShape(); + if (ge::GRAPH_SUCCESS != GetN1Size() || + ge::GRAPH_SUCCESS != GetN2Size() || + ge::GRAPH_SUCCESS != GetGSize() || + ge::GRAPH_SUCCESS != GetBatchSize() || + ge::GRAPH_SUCCESS != GetQTSize() || + ge::GRAPH_SUCCESS != GetKVTSize() || + ge::GRAPH_SUCCESS != GetS1Size() || + ge::GRAPH_SUCCESS != GetQHeadDim() || + ge::GRAPH_SUCCESS != GetKHeadDim() || + ge::GRAPH_SUCCESS != GetS2Size() || + ge::GRAPH_SUCCESS != GetValueHeadDim() || + ge::GRAPH_SUCCESS != GetDSizeKV() || + ge::GRAPH_SUCCESS != GetSparseBlockCount()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +void QSFAInfoParser::GenerateInfo(QSFATilingInfo &qsfaInfo) +{ + qsfaInfo.opName = opName_; + qsfaInfo.platformInfo = platformInfo_; + qsfaInfo.opParamInfo = opParamInfo_; + qsfaInfo.npuArch = npuArch_; + qsfaInfo.isA5 = isA5_; + + qsfaInfo.bSize = bSize_; + qsfaInfo.n1Size = n1Size_; + qsfaInfo.n2Size = n2Size_; + qsfaInfo.s1Size = s1Size_; + qsfaInfo.s2Size = s2Size_; + qsfaInfo.gSize = gSize_; + qsfaInfo.qHeadDim = qHeadDim_; + qsfaInfo.kHeadDim = kHeadDim_; + qsfaInfo.vHeadDim = vHeadDim_; + qsfaInfo.qTSize = qTSize_; + qsfaInfo.kvTSize = kvTSize_; + qsfaInfo.sparseBlockSize = *opParamInfo_.sparseBlockSize; + qsfaInfo.sparseBlockCount = sparseBlockCount_; + + qsfaInfo.inputQType = inputQType_; + qsfaInfo.inputKvType = inputKvType_; + qsfaInfo.outputType = outputType_; + + qsfaInfo.kvStorageMode = kvStorageMode_; + qsfaInfo.l2CacheSize = l2CacheSize_; + + qsfaInfo.totalBlockNum = opParamInfo_.key.shape->GetStorageShape().GetDim(0); + qsfaInfo.scaleValue = *opParamInfo_.scaleValue; + qsfaInfo.pageAttentionFlag = (kvStorageMode_ == KvStorageMode::PAGE_ATTENTION); + qsfaInfo.blockSize = blockSize_; + qsfaInfo.blockTypeSize = sizeof(float); + qsfaInfo.maxBlockNumPerBatch = maxBlockNumPerBatch_; + + FillTilingInfoAttrsAndLayouts(qsfaInfo); +} + +void QSFAInfoParser::FillTilingInfoAttrsAndLayouts(QSFATilingInfo &qsfaInfo) +{ + qsfaInfo.actualLenDimsQ = actualLenDimsQ_; + qsfaInfo.actualLenDimsKV = actualLenDimsKV_; + qsfaInfo.maxActualseq = maxActualseq_; + + qsfaInfo.actualQSeqLenFlag = (opParamInfo_.actualSeqLengthsQ.tensor != nullptr); + qsfaInfo.actualSeqLenFlag = (opParamInfo_.actualSeqLengths.tensor != nullptr); + + qsfaInfo.isSameSeqAllKVTensor = isSameSeqAllKVTensor_; + qsfaInfo.isSameActualseq = isSameActualseq_; + + qsfaInfo.sparseMode = *opParamInfo_.sparseMode; + qsfaInfo.attentionMode = *opParamInfo_.attentionMode; + qsfaInfo.keyQuantMode = *opParamInfo_.keyQuantMode; + qsfaInfo.valueQuantMode = *opParamInfo_.valueQuantMode; + qsfaInfo.quantScaleRepoMode = *opParamInfo_.quantScaleRepoMode; + qsfaInfo.preTokens = *opParamInfo_.preTokens; + qsfaInfo.nextTokens = *opParamInfo_.nextTokens; + qsfaInfo.tileSize = *opParamInfo_.tileSize; + qsfaInfo.ropeHeadDim = *opParamInfo_.ropeHeadDim; + qsfaInfo.returnSoftmaxLse = *opParamInfo_.returnSoftmaxLse; + + qsfaInfo.qLayout = qLayout_; + qsfaInfo.topkLayout = topkLayout_; + qsfaInfo.kvLayout = kvLayout_; + qsfaInfo.outLayout = outLayout_; + qsfaInfo.dSizeVInput = dSizeKV_; +} + +ge::graphStatus QSFAInfoParser::Parse(QSFATilingInfo &qsfaInfo) +{ + if (context_ == nullptr) { + OP_LOGE_WITH_INVALID_INPUT("KvQuantSparseFlashAttention", "tiling context"); + return ge::GRAPH_FAILED; + } + if (ge::GRAPH_SUCCESS != GetOpName() || + ge::GRAPH_SUCCESS != GetNpuInfo() || + ge::GRAPH_SUCCESS != GetOpParaInfo() || + ge::GRAPH_SUCCESS != CheckRequiredParaExistence()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetInOutDataType() || + ge::GRAPH_SUCCESS != GetQueryAndOutLayout() || + ge::GRAPH_SUCCESS != GetTopkLayout() || + ge::GRAPH_SUCCESS != GetKvLayout() || + ge::GRAPH_SUCCESS != GetKvStorageMode()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetShapeAndSizeInfo()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetActualseqInfo()) { + return ge::GRAPH_FAILED; + } + + GenerateInfo(qsfaInfo); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(KvQuantSparseFlashAttention) + .Tiling(TilingKvQuantSparseFlashAttention) + .TilingParse(TilingPrepareForKvQuantSparseFlashAttention); +} // namespace optiling diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/kv_quant_sparse_flash_attention_tiling.h b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/kv_quant_sparse_flash_attention_tiling.h new file mode 100644 index 000000000..e1ce7f3fe --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_host/kv_quant_sparse_flash_attention_tiling.h @@ -0,0 +1,614 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_tiling.h + * \brief + */ +#ifndef KV_QUANT_SPARSE_FLASH_ATTENTION_TILING_H +#define KV_QUANT_SPARSE_FLASH_ATTENTION_TILING_H + +#include +#include +#include +#include +#include "register/tilingdata_base.h" +#include "exe_graph/runtime/tiling_context.h" +#include "platform/soc_spec.h" +namespace optiling { +// ------------------算子原型索引常量定义---------------- +// Inputs Index +constexpr uint32_t QUERY_INPUT_INDEX = 0; +constexpr uint32_t KEY_INPUT_INDEX = 1; +constexpr uint32_t VALUE_INPUT_INDEX = 2; +constexpr uint32_t SPARSE_INDICES_INPUT_INDEX = 3; +constexpr uint32_t KEY_DEQUANT_SCALE_INPUT_INDEX = 3; +constexpr uint32_t VALUE_DEQUANT_SCALE_INPUT_INDEX = 3; +constexpr uint32_t BLOCK_TABLE_INPUT_INDEX = 6; +constexpr uint32_t ACT_SEQ_LEN_Q_INPUT_INDEX = 7; +constexpr uint32_t ACT_SEQ_LEN_KV_INPUT_INDEX = 8; +// Outputs Index +constexpr uint32_t OUTPUT_INDEX = 0; +constexpr uint32_t SOFTMAXMAX_INDEX = 1; +constexpr uint32_t SOFTMAXSUM_INDEX = 2; +// Attributes Index +constexpr uint32_t SCALE_VALUE_ATTR_INDEX = 0; +constexpr uint32_t KEY_QUANT_MODE_ATTR_INDEX = 1; +constexpr uint32_t VALUE_QUANT_MODE_ATTR_INDEX = 2; +constexpr uint32_t SPARSE_BLOCK_SIZE_ATTR_INDEX = 3; +constexpr uint32_t LAYOUT_QUERY_ATTR_INDEX = 4; +constexpr uint32_t LAYOUT_KV_ATTR_INDEX = 5; +constexpr uint32_t SPARSE_MODE_ATTR_INDEX = 6; +constexpr uint32_t PRE_TOKENS_ATTR_INDEX = 7; +constexpr uint32_t NEXT_TOKENS_ATTR_INDEX = 8; +constexpr uint32_t ATTENTION_MODE_ATTR_INDEX = 9; +constexpr uint32_t QUANT_SCALE_REPO_MODE_ATTR_INDEX = 10; +constexpr uint32_t TILE_SIZE_ATTR_INDEX = 11; +constexpr uint32_t ROPE_HEAD_DIM_ATTR_INDEX = 12; +constexpr uint32_t RETURN_SOFTMAX_LSE_ATTR_INDEX = 13; +// Dim Num +constexpr size_t DIM_NUM_TWO = 2; +constexpr size_t DIM_NUM_THREE = 3; +constexpr size_t DIM_NUM_FOUR = 4; +// 常量 +constexpr uint32_t MAX_BLOCK_SIZE = 1024; +constexpr uint32_t COPYND2NZ_SRC_STRIDE_LIMITATION = 65535; +constexpr uint32_t NUM_BYTES_FLOAT = 4; +constexpr uint32_t NUM_BYTES_FLOAT16 = 2; +constexpr uint32_t NUM_BYTES_BF16 = 2; +constexpr uint32_t BYTE_BLOCK = 32; +const uint32_t QSFA_MAX_AIC_CORE_NUM = 26; // 25 + 1 保证数组8字节对齐 + +// ------------------公共定义-------------------------- +enum class QSFALayout : uint32_t { + BSND = 0, + TND = 1, + PA_BSND = 2, +}; + +struct QSFATilingShapeCompareParam { + int64_t B = 1; + int64_t S = 1; + int64_t N = 1; + int64_t D = 1; + int64_t T = 1; + // PA + int64_t Bs = 1; + int64_t Bn = 1; +}; + +enum class KvStorageMode : uint32_t { + BATCH_CONTINUOUS = 0, + PAGE_ATTENTION = 1 +}; + +enum class QSFAPerfMode : uint32_t { + C_TEMPLATE_MODE = 0, + V_TEMPLATE_MODE +}; + +enum class QSFAAxis : uint32_t { + B = 0, + S = 1, + N = 2, + D = 3, + K = 3, // sparse_indices的K和key的D枚举值相同，表达相同位置, 最后一维 + T = 5, + Bn = 6, // block number + Bs = 7, // block size +}; + +struct QSFARequiredParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; +}; + +struct QSFAOptionalParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::Tensor *tensor; +}; + +// -----------算子Tiling入参结构体定义--------------- +struct QSFAParaInfo { + QSFARequiredParaInfo query = {nullptr, nullptr}; + QSFARequiredParaInfo key = {nullptr, nullptr}; + QSFARequiredParaInfo value = {nullptr, nullptr}; + QSFARequiredParaInfo sparseIndices = {nullptr, nullptr}; + QSFAOptionalParaInfo blockTable = {nullptr, nullptr}; + QSFAOptionalParaInfo actualSeqLengthsQ = {nullptr, nullptr}; + QSFAOptionalParaInfo actualSeqLengths = {nullptr, nullptr}; + QSFAOptionalParaInfo queryRope = {nullptr, nullptr}; + QSFAOptionalParaInfo keyRope = {nullptr, nullptr}; + QSFAOptionalParaInfo keyDequantScale = {nullptr, nullptr}; + QSFAOptionalParaInfo valueDequantScale = {nullptr, nullptr}; + QSFARequiredParaInfo attenOut = {nullptr, nullptr}; + QSFARequiredParaInfo softmaxMax = {nullptr, nullptr}; + QSFARequiredParaInfo softmaxSum = {nullptr, nullptr}; + + const char *layoutQuery = nullptr; + const char *layoutKV = nullptr; + const int64_t *sparseBlockSize = nullptr; + const uint32_t *sparseBlockCount = nullptr; + const uint32_t *blockSize = nullptr; + const float *scaleValue = nullptr; + const int64_t *sparseMode = nullptr; + const int64_t *attentionMode = nullptr; + const int64_t *keyQuantMode = nullptr; + const int64_t *valueQuantMode = nullptr; + const int64_t *quantScaleRepoMode = nullptr; + const int64_t *tileSize = nullptr; + const int64_t *ropeHeadDim = nullptr; + const int64_t *preTokens = nullptr; + const int64_t *nextTokens = nullptr; + const bool *returnSoftmaxLse = nullptr; +}; + +struct InnerSplitParams { + uint32_t s1GBaseSize = 1; + uint32_t s2BaseSize = 1; +}; + +// -----------算子TilingData定义--------------- +BEGIN_TILING_DATA_DEF(KvQuantSparseFlashAttentionBaseParamsMla) +TILING_DATA_FIELD_DEF(uint32_t, batchSize) +TILING_DATA_FIELD_DEF(uint32_t, seqSize) +TILING_DATA_FIELD_DEF(uint32_t, qSeqSize) +TILING_DATA_FIELD_DEF(int64_t, blockSize) +TILING_DATA_FIELD_DEF(uint32_t, maxBlockNumPerBatch) +TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsQ) +TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsKV) +TILING_DATA_FIELD_DEF(float, scaleValue) +TILING_DATA_FIELD_DEF(uint32_t, nNumOfQInOneGroup) +TILING_DATA_FIELD_DEF(uint32_t, outputLayout) +TILING_DATA_FIELD_DEF(uint32_t, sparseMode) +TILING_DATA_FIELD_DEF(int64_t, sparseBlockSize) +TILING_DATA_FIELD_DEF(uint32_t, sparseBlockCount) +TILING_DATA_FIELD_DEF(int64_t, dSizeVInput) +TILING_DATA_FIELD_DEF(uint32_t, isActualLenDimsNull) +TILING_DATA_FIELD_DEF(uint32_t, isActualLenDimsKVNull) +TILING_DATA_FIELD_DEF(uint32_t, returnSoftmaxLse) +END_TILING_DATA_DEF + +REGISTER_TILING_DATA_CLASS(KvQuantSparseFlashAttentionBaseParamsMlaOp, KvQuantSparseFlashAttentionBaseParamsMla) + +BEGIN_TILING_DATA_DEF(KvQuantSparseFlashAttentionSingleCoreParamsMla) +TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum); +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(KvQuantSparseFlashAttentionSingleCoreParamsMlaOp, + KvQuantSparseFlashAttentionSingleCoreParamsMla) + +BEGIN_TILING_DATA_DEF(KvQuantSparseFlashAttentionSingleCoreTensorSizeMla) +TILING_DATA_FIELD_DEF(uint32_t, mmResUbSize); +TILING_DATA_FIELD_DEF(uint32_t, bmm2ResUbSize); +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(KvQuantSparseFlashAttentionSingleCoreTensorSizeMlaOp, + KvQuantSparseFlashAttentionSingleCoreTensorSizeMla) + +BEGIN_TILING_DATA_DEF(KvQuantSparseFlashAttentionSplitKVParamsMla) +TILING_DATA_FIELD_DEF(uint32_t, s2) // S2切分份数 +TILING_DATA_FIELD_DEF(uint32_t, accumOutSize) // FD workspace +TILING_DATA_FIELD_DEF(uint32_t, logSumExpSize) // FD workspace +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(KvQuantSparseFlashAttentionSplitKVParamsMlaOp, + KvQuantSparseFlashAttentionSplitKVParamsMla) + +// 内切基本块参数 +BEGIN_TILING_DATA_DEF(KvQuantSparseFlashAttentionInnerSplitParams) +TILING_DATA_FIELD_DEF(uint32_t, mBaseSize) +TILING_DATA_FIELD_DEF(uint32_t, s2BaseSize) +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(KvQuantSparseFlashAttentionInnerSplitParamsOp, + KvQuantSparseFlashAttentionInnerSplitParams) + +BEGIN_TILING_DATA_DEF(KvQuantSparseFlashAttentionTilingDataMla) +TILING_DATA_FIELD_DEF_STRUCT(KvQuantSparseFlashAttentionBaseParamsMla, baseParams); +TILING_DATA_FIELD_DEF_STRUCT(KvQuantSparseFlashAttentionSplitKVParamsMla, splitKVParams); +TILING_DATA_FIELD_DEF_STRUCT(KvQuantSparseFlashAttentionSingleCoreParamsMla, singleCoreParams); +TILING_DATA_FIELD_DEF_STRUCT(KvQuantSparseFlashAttentionSingleCoreTensorSizeMla, singleCoreTensorSize); +TILING_DATA_FIELD_DEF_STRUCT(KvQuantSparseFlashAttentionInnerSplitParams, innerSplitParams); +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(KvQuantSparseFlashAttention, KvQuantSparseFlashAttentionTilingDataMla) + +template inline T Align(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd) * (rnd))); +} + +static std::string QSFADataTypeToSerialString(ge::DataType type); +std::string QSFATensorDesc2String(const gert::StorageShape *shape, const gert::CompileTimeTensorDesc *tensor); +std::string QSFADebugTilingContext(const gert::TilingContext *context); +std::string QSFALayoutToSerialString(QSFALayout layout); + +// -----------算子Tiling入参信息类--------------- +struct QSFATilingInfo { + const char *opName = nullptr; + fe::PlatFormInfos *platformInfo = nullptr; + QSFAParaInfo opParamInfo; + + // Base Param + NpuArch npuArch = NpuArch::DAV_2201; + bool isA5 = false; + uint32_t bSize = 0; + uint32_t n1Size = 0; + uint32_t n2Size = 0; + uint32_t s1Size = 0; + int64_t s2Size = 0; + uint32_t qHeadDim = 0; + uint32_t kHeadDim = 0; + uint32_t vHeadDim = 0; + uint32_t gSize = 0; + uint32_t ropeHeadDim = 0; + uint32_t qTSize = 0; // 仅TND时生效 + uint32_t kvTSize = 0; // 仅TND时生效 + float scaleValue = 0; + uint32_t innerPrecise = 0; + uint32_t l2CacheOffFlag = 0; + int64_t sparseBlockSize = 0; + int64_t sparseBlockCount = 0; + + bool pageAttentionFlag = false; + int64_t blockSize = 0; + uint32_t blockTypeSize = 0; + uint32_t maxBlockNumPerBatch = 0; + uint32_t totalBlockNum = 0; + + uint32_t actualLenDimsQ = 0; + uint32_t maxActualseq = 0; + + bool actualQSeqLenFlag = false; + bool actualSeqLenFlag = false; + bool isSameSeqAllKVTensor = true; + bool isSameActualseq = true; + uint32_t actualLenDimsKV = 0; + std::vector kvListSeqLens {}; + + uint32_t sparseMode = 0; + bool returnSoftmaxLse = false; + + int64_t attentionMode = 0; + int64_t keyQuantMode = 0; + int64_t valueQuantMode = 0; + int64_t quantScaleRepoMode = 0; + int64_t tileSize = 0; + int64_t preTokens = 0; + int64_t nextTokens = 0; + + ge::DataType inputQType = ge::DT_FLOAT16; + ge::DataType inputKvType = ge::DT_FLOAT16; + ge::DataType outputType = ge::DT_FLOAT16; + + KvStorageMode kvStorageMode = KvStorageMode::BATCH_CONTINUOUS; + + QSFALayout qLayout = QSFALayout::BSND; + QSFALayout topkLayout = QSFALayout::BSND; + QSFALayout outLayout = QSFALayout::BSND; + QSFALayout kvLayout = QSFALayout::BSND; + + ge::DataType inputQRopeType = ge::DT_FLOAT16; + ge::DataType inputKRopeType = ge::DT_FLOAT16; + + uint64_t l2CacheSize = 0; + int64_t dSizeVInput = 0; +}; + +// ---------------算子Tiling类--------------- +class QSFAMlaTiling { +public: + explicit QSFAMlaTiling(gert::TilingContext *context) : context_(context) {} + ge::graphStatus DoOpTiling(QSFATilingInfo *qsfaInfo); + +private: + ge::graphStatus SetBlockDim(uint32_t blockDim) const; + ge::graphStatus SetTilingKey(uint64_t tilingKey) const; + ge::graphStatus SetWorkspaceSize(uint64_t workspaceSize) const; + ge::graphStatus SetTilingData(TilingDef &tilingData) const; + gert::TilingContext *context_ = nullptr; + ge::graphStatus GetPlatformInfo(); + void GenTilingKey(); + bool DealSameSeqEachBatch(); + + void ZeroTensorProcess() const; + void InitParams(); + + void Split(); + bool IsBalanceSplitCore(); + + void SplitBalanced(); + void CalcInnerSize(uint32_t qsfaS2Size); + + bool IsFlashDecode(uint32_t coreNum); + + void FillTilingBaseParamsMla(); + void FillTilingSplitKVMla(); + + void FillTilingSingleCoreParamsMla(); + void FillTilingSingleCoreTensorSizeMla(); + void FillTiling(); + + void CalcUbBmm(); + void CheckUbSpace(); + void NormalCalcFDWorkSpace(const uint32_t actCoreNum); + void CalcFDWorkSpace(const uint32_t actCoreNum); + void GetWorkspaceSize(); + + uint32_t CalcBalanceFDParamNums(const uint32_t actCoreNum) const; + + void CalcBlockDim(); + + bool balanceModeFlag_ = false; + bool splitKVFlag_ = false; + + uint32_t coreNum_ = 0; + QSFAPerfMode perfMode_ = QSFAPerfMode::V_TEMPLATE_MODE; + uint32_t kvSplitPart_ = 1; + size_t mmResUbSize_ = 0; + size_t bmm2ResUbSize_ = 0; + size_t qPreSizeMla_ = 0; + uint32_t sInnerLoopTimes_ = 0; + uint32_t sInnerSize_ = 0; + uint32_t sInnerSizeTail_ = 0; + uint32_t sInnerSizeAlign_ = 0; + uint32_t kvSplit_ = 0; + uint32_t usedCoreNum_ = 0; + uint32_t formerCoreNum_ = 0; + uint32_t blockSplitBn2Range_ = 0; + uint32_t tailSplitedBatchRange_ = 0; + + uint32_t aicNum_ = 0; + uint32_t aivNum_ = 0; + size_t libapiSize_ = 0; + + KvQuantSparseFlashAttentionTilingDataMla tilingData_; + uint32_t blockDim_{0}; + uint64_t workspaceSize_{0}; + uint64_t tilingKey_{0}; + + uint32_t headDimAlign_ = 0; + uint32_t mBaseSize_ = 128; + uint32_t mFdBaseSize_ = 8; + + QSFATilingInfo *qsfaInfo_ = nullptr; +}; + +// -----------算子Tiling入参信息解析及Check类--------------- +class QSFATilingCheck { +public: + explicit QSFATilingCheck(const QSFATilingInfo &qsfaInfo) : qsfaInfo_(qsfaInfo) {}; + ~QSFATilingCheck() = default; + ge::graphStatus Process(); +private: + void Init(); + void LogErrorDtypeSupport(const std::vector &expectDtypeList, + const ge::DataType &actualDtype, const std::string &name) const; + ge::graphStatus CheckDtypeSupport(const gert::CompileTimeTensorDesc *qsfaDesc, + const std::string &name) const; + template void LogErrorNumberSupport(const std::vector &expectNumberList, + const T &actualValue, const std::string &name, const std::string subName) const; + template void LogErrorDimNumSupport(const std::vector &expectNumberList, + const T &actualValue, const std::string &name) const; + ge::graphStatus CheckDimNumSupport(const gert::StorageShape *shape, + const std::vector &qsfaExpectDimNumList, const std::string &name) const; + ge::graphStatus CheckDimNumInLayoutSupport(const QSFALayout &layout, + const gert::StorageShape *shape, const std::string &name) const; + void LogErrorLayoutSupport(const std::vector &expectLayoutList, + const QSFALayout &actualLayout, const std::string &name) const; + ge::graphStatus GetExpectedShape(gert::Shape &shapeExpected, + const QSFATilingShapeCompareParam ¶m, const QSFALayout &layout) const; + ge::graphStatus CompareShape(QSFATilingShapeCompareParam ¶m, + const gert::Shape &shape, const QSFALayout &layout, const std::string &name) const; + ge::graphStatus CheckLayoutSupport(const QSFALayout &actualLayout, const std::string &name) const; + ge::graphStatus CheckSingleParaQuery() const; + ge::graphStatus CheckSingleParaKey() const; + ge::graphStatus CheckSingleParaValue() const; + ge::graphStatus CheckSingleParaAttenOut() const; + ge::graphStatus CheckSingleParaNumHeads() const; + ge::graphStatus CheckSingleParaKvHeadNums() const; + ge::graphStatus CheckSingleParaLayout() const; + ge::graphStatus CheckSingleParaSparseMode() const; + ge::graphStatus CheckSingleParaSparseBlockSize() const; + ge::graphStatus CheckSingleParaSparseIndices() const; + ge::graphStatus CheckSinglePara() const; + ge::graphStatus CheckMultiParaConsistency() const; + ge::graphStatus CheckDequantScaleNotExistence(); + template ge::graphStatus CheckAttrValueByMap( + std::map> &attrMap) const; + ge::graphStatus CheckParaExistenceMlaAntiquant() const; + ge::graphStatus CheckParaExistenceGqaAntiquant() const; + ge::graphStatus CheckParaExistenceMla() const; + ge::graphStatus CheckParaExistence(); + void SetQSFAShapeCompare(); + ge::graphStatus CheckKVDType(); + ge::graphStatus CheckKVShapeForBatchContinuous(); + ge::graphStatus CheckKVShapeForPageAttention(); + ge::graphStatus CheckKVShape(); + ge::graphStatus CheckKV(); + ge::graphStatus CheckTopK(); + ge::graphStatus CheckTopkShape(); + ge::graphStatus CheckBlockTable() const; + ge::graphStatus CheckDTypeConsistency(const ge::DataType &actualDtype, + const ge::DataType &expectDtype, const std::string &name) const; + + ge::graphStatus CheckAttenOut(); + ge::graphStatus CheckAttenOutShape(); + ge::graphStatus CheckActualSeqLensQ(); + ge::graphStatus CheckActualSeqLensQShape(); + ge::graphStatus CheckActualSeqLensQDType(); + ge::graphStatus CheckActualSeqLens(); + ge::graphStatus CheckActualSeqLensDType(); + ge::graphStatus CheckActualSeqLensShape(); + ge::graphStatus CheckMultiParaConsistency(); + + ge::graphStatus CheckFeatureMlaAntiquantShape() const; + ge::graphStatus CheckFeatureMlaAntiquantShapeSizes() const; + ge::graphStatus CheckFeatureMlaAntiquantShapeSparseAndHeadDim() const; + ge::graphStatus CheckFeatureMlaAntiquantLayout() const; + ge::graphStatus CheckFeatureMlaAntiquantDtype() const; + ge::graphStatus CheckFeatureMlaAntiquantAttr() const; + ge::graphStatus CheckFeatureMlaAntiquantPa() const; + ge::graphStatus CheckFeatureMlaAntiquant() const; + ge::graphStatus CheckFeatureMla() const; + ge::graphStatus CheckFeature() const; + +private: + const char *opName_; + fe::PlatFormInfos *platformInfo_; + QSFAParaInfo opParamInfo_; + const QSFATilingInfo &qsfaInfo_; + + uint32_t bSize_ = 0; + uint32_t n1Size_ = 0; + uint32_t n2Size_ = 0; + uint32_t gSize_ = 0; + uint32_t s1Size_ = 0; + int64_t s2Size_ = 0; + uint32_t qHeadDim_ = 0; + uint32_t kHeadDim_ = 0; + uint32_t vHeadDim_ = 0; + uint32_t qTSize_ = 0; // 仅TND时生效 + uint32_t kvTSize_ = 0; // 仅TND时生效 + KvStorageMode kvStorageMode_ = KvStorageMode::BATCH_CONTINUOUS; + uint32_t sparseBlockCount_ = 0; + int64_t sparseBlockSize_ = 0; + int32_t attentionMode_ = 0; + int32_t keyQuantMode_ = 0; + int32_t valueQuantMode_ = 0; + int32_t quantScaleRepoMode_ = 0; + int64_t tileSize_ = 0; + int64_t preTokens_ = 0; + int64_t nextTokens_ = 0; + int32_t ropeHeadDim_ = 0; + + QSFALayout qLayout_ = QSFALayout::BSND; + QSFALayout topkLayout_ = QSFALayout::BSND; + QSFALayout outLayout_ = QSFALayout::BSND; + QSFALayout kvLayout_ = QSFALayout::BSND; + + uint32_t maxBlockNumPerBatch_ = 0; + int64_t blockSize_ = 0; + + uint32_t aicNum_ = 0; + uint32_t aivNum_ = 0; + NpuArch npuArch_ = NpuArch::DAV_2201; + bool isA5_ = false; + uint64_t l2CacheSize_ = 0; + + ge::DataType inputQType_ = ge::DT_FLOAT16; + ge::DataType inputKvType_ = ge::DT_FLOAT16; + ge::DataType outputType_ = ge::DT_FLOAT16; + + gert::Shape queryShapeCmp_{}; + gert::Shape keyShapeCmp_{}; + gert::Shape valueShapeCmp_{}; + gert::Shape topkShapeCmp_{}; + gert::Shape attenOutShapeCmp_{}; +}; + +class QSFAInfoParser { +public: + explicit QSFAInfoParser(const gert::TilingContext *context) : context_(context) {} + ~QSFAInfoParser() = default; + + ge::graphStatus CheckRequiredInOutExistence() const; + ge::graphStatus CheckRequiredAttrExistence() const; + ge::graphStatus CheckRequiredParaExistence() const; + + ge::graphStatus GetActualSeqLenQSize(uint32_t &size); + ge::graphStatus GetNpuInfo(); + ge::graphStatus GetOpName(); + void GetOptionalInputParaInfo(); + void GetInputParaInfo(); + void GetOutputParaInfo(); + ge::graphStatus GetAttrParaInfo(); + ge::graphStatus GetOpParaInfo(); + ge::graphStatus GetKvCache(); + + ge::graphStatus GetInOutDataType(); + ge::graphStatus GetQTSize(); + ge::graphStatus GetBatchSize(); + ge::graphStatus GetKVTSize(); + ge::graphStatus GetQHeadDim(); + ge::graphStatus GetKHeadDim(); + ge::graphStatus GetS1Size(); + ge::graphStatus GetKvStorageMode(); + ge::graphStatus GetKvLayout(); + void SetQSFAShape(); + ge::graphStatus GetS2SizeForBatchContinuous(); + ge::graphStatus GetMaxBlockNumPerBatch(); + ge::graphStatus GetBlockSize(); + ge::graphStatus GetS2SizeForPageAttention(); + ge::graphStatus GetS2Size(); + ge::graphStatus GetValueHeadDim(); + ge::graphStatus GetDSizeKV(); + ge::graphStatus GetRopeHeadDim(); + ge::graphStatus GetQueryAndOutLayout(); + ge::graphStatus GetTopkLayout(); + ge::graphStatus GetN1Size(); + ge::graphStatus GetN2Size(); + ge::graphStatus GetGSize(); + ge::graphStatus GetSparseBlockCount(); + ge::graphStatus GetActualseqInfo(); + ge::graphStatus GetShapeAndSizeInfo(); + void GenerateInfo(QSFATilingInfo &qsfaInfo); + void FillTilingInfoAttrsAndLayouts(QSFATilingInfo &qsfaInfo); + ge::graphStatus Parse(QSFATilingInfo &qsfaInfo); + + const gert::TilingContext *context_ = nullptr; + + const char *opName_; + fe::PlatFormInfos *platformInfo_; + QSFAParaInfo opParamInfo_; + + uint32_t bSize_ = 0; + uint32_t n1Size_ = 0; + uint32_t n2Size_ = 0; + uint32_t gSize_ = 0; + uint32_t s1Size_ = 0; + int64_t s2Size_ = 0; + uint32_t qHeadDim_ = 0; + uint32_t kHeadDim_ = 0; + uint32_t vHeadDim_ = 0; + int32_t ropeHeadDim_ = 0; + int64_t dSizeKV_ = 0; + uint32_t qTSize_ = 0; // 仅TND时生效 + uint32_t kvTSize_ = 0; // 仅TND时生效 + KvStorageMode kvStorageMode_ = KvStorageMode::BATCH_CONTINUOUS; + uint32_t sparseBlockCount_ = 0; + + QSFALayout qLayout_ = QSFALayout::BSND; + QSFALayout topkLayout_ = QSFALayout::BSND; + QSFALayout outLayout_ = QSFALayout::BSND; + QSFALayout kvLayout_ = QSFALayout::BSND; + + uint32_t maxBlockNumPerBatch_ = 0; + uint32_t blockSize_ = 0; + + NpuArch npuArch_ = NpuArch::DAV_2201; + bool isA5_ = false; + + ge::DataType inputQType_ = ge::DT_FLOAT16; + ge::DataType inputKvType_ = ge::DT_FLOAT16; + ge::DataType outputType_ = ge::DT_FLOAT16; + + uint64_t l2CacheSize_ = 0; + + bool isSameSeqAllKVTensor_ = true; + bool isSameActualseq_ = true; + uint32_t maxActualseq_ = 0; + + uint32_t actualLenDimsQ_ = 0; + uint32_t actualLenDimsKV_ = 0; + + gert::Shape queryShape_{}; + gert::Shape keyShape_{}; + gert::Shape valueShape_{}; + gert::Shape sparseIndicesShape_{}; +}; +} // namespace optiling +#endif // KV_QUANT_SPARSE_FLASH_ATTENTION_TILING_H diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_common_arch35.h b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_common_arch35.h new file mode 100644 index 000000000..c1bb5d231 --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_common_arch35.h @@ -0,0 +1,130 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_common_arch35.h + * \brief + */ +#ifndef KV_QUANT_SPARSE_FLASH_ATTENTION_COMMON_ARCH35_H +#define KV_QUANT_SPARSE_FLASH_ATTENTION_COMMON_ARCH35_H +#include +#include "kernel_tiling/kernel_tiling.h" + +#if __has_include("../../sparse_flash_attention/arch35/common/util_regbase.h") +#include "../../sparse_flash_attention/arch35/common/util_regbase.h" +#else +#include "../../../sparse_flash_attention/op_kernel/arch35/common/util_regbase.h" +#endif + +#if __has_include("../../common/op_kernel/buffer.h") +#include "../../common/op_kernel/buffer.h" +#else +#include "../../common/buffer.h" +#endif +#if __has_include("../../common/op_kernel/buffer_manager.h") +#include "../../common/op_kernel/buffer_manager.h" +#else +#include "../../common/buffer_manager.h" +#endif +#if __has_include("../../common/op_kernel/buffers_policy.h") +#include "../../common/op_kernel/buffers_policy.h" +#else +#include "../../common/buffers_policy.h" +#endif + +constexpr uint64_t BLOCK_BYTE = 32; +constexpr uint32_t NEGATIVE_MIN_VALUE_FP32 = 0xFF7FFFFF; + +constexpr uint32_t BUFFER_SIZE_16K = 16384; // 16384表示16 * 1024 +constexpr uint32_t BUFFER_SIZE_32K = 32768; // 32768表示32 * 1024 +constexpr uint32_t BUFFER_SIZE_128K = 131072; // 131072表示128 * 1024 + +constexpr uint32_t L0AB_SHARED_SIZE_64K = 65536; // 65536表示64*1024 +constexpr uint32_t L0C_SHARED_SIZE_256K = 262144; // 262144表示256 * 1024 + +constexpr uint32_t CV_RATIO = 2; +constexpr uint64_t SYNC_MODE = 4; + +static constexpr uint32_t QSFA_SYNC_MODE0 = 0; + +enum class QSFA_LAYOUT { + BSND = 0, + TND = 1, + PA_BSND = 2, +}; + +enum class QSFATemplateMode { + SWA_TEMPLATE_MODE = 0, + CFA_TEMPLATE_MODE = 1, + SCFA_TEMPLATE_MODE = 2 +}; + +namespace BaseApi { +__aicore__ constexpr uint64_t Align2Func(uint64_t data) { + return (data + 1UL) >> 1UL << 1UL; // 向上2对齐, +1移位2 +} + +__aicore__ constexpr uint64_t Align8Func(uint64_t data) { + return (data + 7UL) >> 3UL << 3UL; // 向上8对齐, +7移位3 +} + +__aicore__ constexpr uint64_t Align16Func(uint64_t data) { + return (data + 15UL) >> 4UL << 4UL; // 向上16对齐, +15移位4 +} + +__aicore__ constexpr uint64_t Align64Func(uint64_t data) { + return (data + 63UL) >> 6UL << 6UL; // 向上64对齐, +63移位6 +} +} + +#define TEMPLATE_INTF \ + template + +#define TEMPLATE_INTF_ARGS \ + Q_T, KV_T, T, OUTPUT_T, isFd, isPa, LAYOUT_T, KV_LAYOUT_T, TEMPLATE_MODE, IS_SPLIT_G + +#define QSFA_CUBE_BLOCK_TRAITS_TYPE_FIELDS(X) \ + X(Q_T) \ + X(KV_T) \ + X(T) \ + X(OUTPUT_T) \ + +#define QSFA_CUBE_BLOCK_TRAITS_CONST_FIELDS(X) \ + X(isFd, bool, false) \ + X(isPa, bool, true) \ + X(LAYOUT_T, QSFA_LAYOUT, QSFA_LAYOUT::BSND) \ + X(KV_LAYOUT_T, QSFA_LAYOUT, QSFA_LAYOUT::PA_BSND) \ + X(TEMPLATE_MODE, QSFATemplateMode, QSFATemplateMode::SCFA_TEMPLATE_MODE) \ + X(IS_SPLIT_G, bool, false) + + +/* 1. 生成带默认值的模版Template */ +#define GEN_TYPE_PARAM(name) typename name, +#define GEN_CONST_PARAM(name, type, default_val) type name = default_val, + +#define TEMPLATES_DEF \ +template + +/* 2. 生成不带默认值的模版Template */ +#define GEN_TEMPLATE_TYPE_NODEF(name) typename name, +#define GEN_TEMPLATE_CONST_NODEF(name, type, default_val) type name, +#define TEMPLATES_DEF_NO_DEFAULT \ +template + +/* 3. 生成有默认值的Args */ +#define GEN_ARG_NAME(name, ...) name, +#define TEMPLATE_ARGS \ + QSFA_CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_ARG_NAME) \ + QSFA_CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_ARG_NAME) end + +#endif //KV_QUANT_SPARSE_FLASH_ATTENTION_COMMON_ARCH35_H diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_kernel_mla.h b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_kernel_mla.h new file mode 100644 index 000000000..90f34dc39 --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_kernel_mla.h @@ -0,0 +1,707 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_kernel_mla.h + * \brief + */ + +#ifndef KV_QUANT_SPARSE_FLASH_ATTENTION_KERNEL_MLA_H +#define KV_QUANT_SPARSE_FLASH_ATTENTION_KERNEL_MLA_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "kv_quant_sparse_flash_attention_service_cube_mla.h" +#include "kv_quant_sparse_flash_attention_service_vector_mla.h" +#include "kv_quant_sparse_flash_attention_common_arch35.h" +#include "kv_quant_sparse_flash_attention_kvcache.h" +#if __has_include("../../common/op_kernel/CopyInL1.h") +#include "../../common/op_kernel/CopyInL1.h" +#else +#include "../common/CopyInL1.h" +#endif +#if __has_include("../../common/op_kernel/matmul.h") +#include "../../common/op_kernel/matmul.h" +#else +#include "../common/matmul.h" +#endif +#if __has_include("../../common/op_kernel/FixpipeOut.h") +#include "../../common/op_kernel/FixpipeOut.h" +#else +#include "../common/FixpipeOut.h" +#endif + +using matmul::MatmulType; +using namespace AscendC; +using namespace AscendC::Impl::Detail; +using namespace regbaseutil; + +namespace BaseApi { +template class KvQuantSparseFlashAttentionMla { +public: + ARGS_TRAITS; + + __aicore__ inline KvQuantSparseFlashAttentionMla(){}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t* keyScale, + __gm__ uint8_t* valueScale, __gm__ uint8_t *blockTable, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths, + __gm__ uint8_t *attentionOut, __gm__ uint8_t *workspace, + const KvQuantSparseFlashAttentionTilingDataMla *__restrict tiling, + TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void ProcessMainLoop(); + __aicore__ inline void InitGlobalBuffer(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *blockTable, __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths, + __gm__ uint8_t *workspace, const KvQuantSparseFlashAttentionTilingDataMla *__restrict tiling, TPipe *tPipe); + __aicore__ inline void InitLocalBuffer(); + __aicore__ inline void ComputeConstexpr(); + __aicore__ inline void InitMMResBuf(__gm__ uint8_t *workspace); + __aicore__ inline void SetRunInfo(RunInfo &runInfo, RunParamStr &runParam, int64_t taskId, int64_t s2LoopCount, + int64_t s2LoopLimit, int64_t multiCoreInnerIdx); + __aicore__ inline void ComputeBmm1Tail(RunInfo &runInfo, RunParamStr &runParam); + __aicore__ inline void InitUniqueConstInfo(); + __aicore__ inline void InitUniqueRunInfo(const RunParamStr &runParam, RunInfo &runInfo); + __aicore__ inline void ComputeAxisIdxByBnAndGs1(int64_t bnIndex, int64_t gS1Index, RunParamStr &runParam); + __aicore__ inline void InitCalcParamsEach(); + __aicore__ inline uint64_t GetBalanceActualSeqLengths(GlobalTensor &actualSeqLengths, uint32_t bIdx); + __aicore__ inline void GetAxisStartIdx(uint32_t bN2EndPrev, uint32_t s1GEndPrev, uint32_t s2EndPrev); + + TPipe *pipe; + + const KvQuantSparseFlashAttentionTilingDataMla *__restrict tilingData; + static constexpr uint64_t SYNC_MODE = 4; + static constexpr uint32_t PRELOAD_NUM = 2; + /* 核间通道 */ + BufferManager ubBufferManager; + BuffersPolicyDB bmm1Buffers; + BuffersPolicySingleBuffer bmm2Buffers; + BufferManager gmBufferManager; + + // mm2左矩阵P + BufferManager l1BufferManager; + BuffersPolicy3buff l1RightBuffers; + CVSharedParams sharedParams; + /* GM信息 */ + __gm__ int32_t *actualSeqKvlenAddr = nullptr; + __gm__ int32_t *actualSeqQlenAddr = nullptr; + + GlobalTensor actualSeqLengthsQGm; + uint32_t usedCoreNum = 0U; + + /* workspace 空间 */ + BuffersPolicy3buff v0ResGmBuffers; + + /* 核Index信息 */ + int32_t aicIdx; + + /* 切G时最大s2Loop */ + int64_t maxS2LoopCnt; + + /* 初始化后不变的信息 */ + ConstInfo constInfo; + + /* 模板库Block */ + CubeBlockType cubeBlock; + VecBlockType vecBlock; + + uint32_t crossCoreSyncBufId = 0; +}; + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::Init( + __gm__ uint8_t *query, + __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t* keyScale, + __gm__ uint8_t* valueScale, __gm__ uint8_t *blockTable, __gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengths, __gm__ uint8_t *attentionOut, __gm__ uint8_t *workspace, + const KvQuantSparseFlashAttentionTilingDataMla *__restrict tiling, + TPipe *tPipe) +{ + fa_base_matmul::idCounterNum = 0; + constInfo.subBlockIdx = GetSubBlockIdx(); + if ASCEND_IS_AIC { + this->aicIdx = GetBlockIdx(); + constInfo.aivIdx = 0; + } else { + constInfo.aivIdx = GetBlockIdx(); + this->aicIdx = constInfo.aivIdx >> 1; + this->tilingData = tiling; + } + + constInfo.s1BaseSize = 64; + constInfo.s2BaseSize = 128; + + this->pipe = tPipe; + vecBlock.InitVecBlock(tPipe, this->tilingData, this->sharedParams, this->aicIdx, constInfo.subBlockIdx, actualSeqLengthsQ, actualSeqLengths); + if ASCEND_IS_AIV { + constInfo.bSize = this->sharedParams.bSize; + constInfo.gSize = this->sharedParams.gSize; + constInfo.s1Size = this->sharedParams.s1Size; + constInfo.needInit = this->sharedParams.needInit; + constInfo.dSizeV = 512; + } + vecBlock.CleanOutput(attentionOut, constInfo); + /* cube侧不依赖sharedParams的scalar前置 */ + InitMMResBuf(workspace); + if ASCEND_IS_AIC { + cubeBlock.InitCubeBlock(pipe, &l1BufferManager, query); + /* wait kfc message */ + CrossCoreWaitFlag(15); + auto tempTilingSSbuf = reinterpret_cast<__ssbuf__ uint32_t*>(0); // 从ssbuf的0地址开始拷贝 + auto tempTiling = reinterpret_cast(&sharedParams); + #pragma unroll + for (int i = 0; i < sizeof(CVSharedParams) / sizeof(uint32_t); ++i, ++tempTilingSSbuf, ++tempTiling) { + *tempTiling = *tempTilingSSbuf; + } + } + this->ComputeConstexpr(); + this->InitGlobalBuffer(query, key, value, sparseIndices, blockTable, actualSeqLengthsQ, actualSeqLengths, + workspace, tiling, tPipe); // gm设置 + this->InitCalcParamsEach(); + this->InitLocalBuffer(); +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::InitCalcParamsEach() +{ + // 计算总的基本块 + maxS2LoopCnt = 0; // 所有核中最大累计s2Loop + uint32_t qsfaTotalBaseNum = 0; + uint32_t actBatchS2 = 1; + uint32_t coreNum = GetBlockNum(); // G128时相邻两个cube核处理一个s1，coreNum减半 + uint32_t currCoreIdx = aicIdx; + + if constexpr (IS_SPLIT_G) { + currCoreIdx = currCoreIdx >> 1; + coreNum = coreNum >> 1; + } + + uint32_t actBatchS1 = 1; + for (uint32_t bIdx = 0; bIdx < constInfo.bSize; bIdx++) { + uint32_t actBatchS1 = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bIdx); //不切S2，只关注S1 + qsfaTotalBaseNum += actBatchS1 * actBatchS2; + } + + uint32_t avgBaseNum = 1; + if (qsfaTotalBaseNum > coreNum) { + avgBaseNum = (qsfaTotalBaseNum + coreNum - 1) / coreNum; + if constexpr (IS_SPLIT_G) { + usedCoreNum = ((qsfaTotalBaseNum + avgBaseNum - 1) / avgBaseNum) << 1; + } + } else { + if constexpr (IS_SPLIT_G) { + usedCoreNum = qsfaTotalBaseNum << 1; + } else { + usedCoreNum = qsfaTotalBaseNum; + } + } + + if constexpr (IS_SPLIT_G) { + maxS2LoopCnt = avgBaseNum * (Min(constInfo.sparseBlockCount, constInfo.s2Size) + + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + } + + if (aicIdx >= usedCoreNum) { + return; + } + // 计算当前核的基本块 + uint32_t qsfaAccumBaseNum = 0; // qsfa当前累积的基本块数 + uint32_t targetBaseNum = 0; + uint32_t qsfaLastValidBIdx = 0; + uint32_t lastValidactBatchS1 = 0; + bool setStart = false; + targetBaseNum = (currCoreIdx + 1) * avgBaseNum; // 计算当前的目标权重 + uint32_t targetStartBaseNum = targetBaseNum - avgBaseNum; + for (uint32_t bN2Idx = 0; bN2Idx < constInfo.bSize * constInfo.n2Size; bN2Idx++) { + uint32_t bIdx = bN2Idx / constInfo.n2Size; + actBatchS1 = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bIdx); + for (uint32_t s1GIdx = 0; s1GIdx < actBatchS1; s1GIdx++) { + qsfaAccumBaseNum += 1; + if (!setStart && qsfaAccumBaseNum >= targetStartBaseNum) { + constInfo.bN2Start = bN2Idx; + constInfo.gS1Start = s1GIdx; + setStart = true; + } + if (qsfaAccumBaseNum >= targetBaseNum) { + // 更新当前核的End分核信息 + constInfo.s2End = 0; + constInfo.bN2End = bN2Idx; + constInfo.gS1End = s1GIdx; + + if (currCoreIdx != 0) { + GetAxisStartIdx(constInfo.bN2Start, constInfo.gS1Start, 0); + } + + return; + } + } + + if ((actBatchS1 > 0) && (actBatchS2 > 0)) { + qsfaLastValidBIdx = bIdx; + lastValidactBatchS1 = actBatchS1; + } + } + if (!setStart) { + constInfo.bN2Start = qsfaLastValidBIdx; + constInfo.gS1Start = lastValidactBatchS1 - 1; + } + if (qsfaAccumBaseNum < targetBaseNum) { + // 更新最后一个核的End分核信息 + constInfo.bN2End = qsfaLastValidBIdx; + constInfo.gS1End = lastValidactBatchS1 - 1; + constInfo.s2End = 0; + if (currCoreIdx != 0) { + GetAxisStartIdx(constInfo.bN2Start, constInfo.gS1Start, 0); + } + return; + } +} + +template +__aicore__ inline uint64_t KvQuantSparseFlashAttentionMla::\ + GetBalanceActualSeqLengths(GlobalTensor &actualSeqLengths, uint32_t bIdx) +{ + if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + if (bIdx == 0) { + return actualSeqQlenAddr[0]; + } else if (bIdx > 0) { + return actualSeqQlenAddr[bIdx] - actualSeqQlenAddr[bIdx - 1]; + } else { + return 0; + } + } else { + if (constInfo.isActualLenDimsNull == 0) { + return actualSeqQlenAddr[bIdx]; + } else { + return constInfo.s1Size; + } + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::GetAxisStartIdx(uint32_t bN2EndPrev, + uint32_t s1GEndPrev, + uint32_t s2EndPrev) +{ + uint32_t qsfaBEndPrev = bN2EndPrev / constInfo.n2Size; + uint32_t actualSeqQPrev = GetBalanceActualSeqLengths(actualSeqLengthsQGm, qsfaBEndPrev); + uint32_t s1GPrevBaseNum = actualSeqQPrev; + constInfo.bN2Start = bN2EndPrev; + constInfo.gS1Start = s1GEndPrev; + constInfo.s2Start = 0; + if (s1GEndPrev >= s1GPrevBaseNum - 1) { // 上个核把S1G处理完了 + constInfo.bN2Start++; + constInfo.gS1Start = 0; + } else { + constInfo.gS1Start++; + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::InitGlobalBuffer( + __gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *blockTable, __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths, + __gm__ uint8_t *workspace, const KvQuantSparseFlashAttentionTilingDataMla *__restrict tiling, TPipe *tPipe) +{ + if (actualSeqLengthsQ != nullptr) { + actualSeqQlenAddr = (__gm__ int32_t *)actualSeqLengthsQ; + } + + if (actualSeqLengths != nullptr) { + actualSeqKvlenAddr = (__gm__ int32_t *)actualSeqLengths; + } + + vecBlock.InitGlobalBuffer(key, value, sparseIndices, blockTable); + cubeBlock.InitCubeInput(actualSeqLengthsQ, constInfo); +} + + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::InitMMResBuf( + __gm__ uint8_t *workspace) +{ + uint32_t mm1RightSize = constInfo.s2BaseSize * 576 * sizeof(Q_T); + l1BufferManager.Init(pipe, 524288); // 512 * 1024 + l1RightBuffers.Init(l1BufferManager, mm1RightSize); + l1RightBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID); + crossCoreSyncBufId++; + l1RightBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID); + crossCoreSyncBufId++; + l1RightBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID); + crossCoreSyncBufId++; + + if ASCEND_IS_AIC { + l1RightBuffers.Get().SetCrossCore(); + l1RightBuffers.Get().SetCrossCore(); + l1RightBuffers.Get().SetCrossCore(); + } + uint32_t mm1ResultSize = constInfo.s1BaseSize / CV_RATIO * constInfo.s2BaseSize * sizeof(T); + uint32_t mm2ResultSize = constInfo.s1BaseSize / CV_RATIO * 512 * sizeof(T); + ubBufferManager.Init(pipe, mm1ResultSize * 2 + mm2ResultSize); + + bmm1Buffers.Init(ubBufferManager, mm1ResultSize); + bmm1Buffers.Get().SetCrossCoreID(crossCoreSyncBufId, crossCoreSyncBufId); + crossCoreSyncBufId++; + bmm1Buffers.Get().SetCrossCoreID(crossCoreSyncBufId, crossCoreSyncBufId); + crossCoreSyncBufId++; + if ASCEND_IS_AIV { + bmm1Buffers.Get().SetCrossCore(); + bmm1Buffers.Get().SetCrossCore(); + } + + bmm2Buffers.Init(ubBufferManager, mm2ResultSize); + bmm2Buffers.Get().SetCrossCoreID(crossCoreSyncBufId, crossCoreSyncBufId); + crossCoreSyncBufId++; + + if ASCEND_IS_AIV { + bmm2Buffers.Get().SetCrossCore(); + } + + if constexpr (IS_SPLIT_G) { + uint32_t v0ResSize = constInfo.s2BaseSize * 576U * sizeof(Q_T); + int64_t totalOffset = v0ResSize * 3 * (aicIdx >> 1U); + gmBufferManager.Init(workspace + totalOffset); + v0ResGmBuffers.Init(gmBufferManager, v0ResSize); + v0ResGmBuffers.Get().SetCrossCoreID(INVALID_CROSS_CORE_EVENT_ID, crossCoreSyncBufId); + crossCoreSyncBufId++; + v0ResGmBuffers.Get().SetCrossCoreID(INVALID_CROSS_CORE_EVENT_ID, crossCoreSyncBufId); + crossCoreSyncBufId++; + v0ResGmBuffers.Get().SetCrossCoreID(INVALID_CROSS_CORE_EVENT_ID, crossCoreSyncBufId); + crossCoreSyncBufId++; + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::InitLocalBuffer() +{ + vecBlock.InitLocalBuffer(pipe, constInfo); +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::ComputeConstexpr() +{ + // 计算轴的乘积 + usedCoreNum = sharedParams.usedCoreNum; + + if ASCEND_IS_AIC { + constInfo.bSize = this->sharedParams.bSize; + constInfo.gSize = this->sharedParams.gSize; + constInfo.s1Size = this->sharedParams.s1Size; + constInfo.needInit = this->sharedParams.needInit; + constInfo.dSizeV = 512; + } + constInfo.n2Size = sharedParams.n2Size; + constInfo.s2Size = sharedParams.s2Size; + constInfo.dSize = sharedParams.dSize; + constInfo.dSizeVInput = sharedParams.dSizeVInput; + constInfo.dSizeRope = sharedParams.dSizeRope; + constInfo.dSizeNope = constInfo.dSize - constInfo.dSizeRope; + constInfo.tileSize = sharedParams.tileSize; + constInfo.sparseBlockCount = sharedParams.sparseBlockCount; + constInfo.sparseBlockSize = 1; + + constInfo.sparseMode = sharedParams.maskMode; + constInfo.n2G = constInfo.n2Size * constInfo.gSize; + + constInfo.s1Dv = constInfo.s1Size * constInfo.dSizeV; + constInfo.s2Dv = constInfo.s2Size * constInfo.dSizeV; + constInfo.n2Dv = constInfo.n2Size * constInfo.dSizeV; + + constInfo.gDv = constInfo.gSize * constInfo.dSizeV; + constInfo.n2S2Dv = constInfo.n2Size * constInfo.s2Dv; + constInfo.n2GDv = constInfo.n2Size * constInfo.gDv; + constInfo.s2BaseN2Dv = constInfo.s2BaseSize * constInfo.n2Dv; + constInfo.layoutType = sharedParams.layoutType; + + constInfo.isActualLenDimsNull = sharedParams.isActualSeqLengthsNull; + constInfo.isActualLenDimsKVNull = sharedParams.isActualSeqLengthsKVNull; + + if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + // (BS)ND + constInfo.s1BaseN2GDv = constInfo.s1BaseSize * constInfo.n2GDv; + constInfo.mm1Ka = constInfo.n2Size * constInfo.dSize; + if ASCEND_IS_AIV { + constInfo.attentionOutStride = \ + (constInfo.n2G - constInfo.gSize) * constInfo.dSizeV * sizeof(OUTPUT_T); + } + } else if constexpr (LAYOUT_T == QSFA_LAYOUT::BSND) { + // BSH/BSNGD + constInfo.s1BaseN2GDv = constInfo.s1BaseSize * constInfo.n2GDv; + constInfo.mm1Ka = constInfo.n2Size * constInfo.dSize; + + if ASCEND_IS_AIV { + constInfo.attentionOutStride = \ + (constInfo.n2G - constInfo.gSize) * constInfo.dSizeV * sizeof(OUTPUT_T); + } + } + + if ASCEND_IS_AIV { + constInfo.blockSize = sharedParams.blockSize; + constInfo.softmaxScale = sharedParams.softmaxScale; + constInfo.maxBlockNumPerBatch = sharedParams.maxBlockNumPerBatch; + } + + InitUniqueConstInfo(); +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::InitUniqueConstInfo() +{ + // bsize + 1-> bsize + this->constInfo.actualSeqLenSize = this->sharedParams.bSize; + this->constInfo.actualSeqLenKVSize = this->sharedParams.bSize; +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::Process() +{ + // SyncAll Cube和Vector都需要调用 + if (this->sharedParams.needInit) { + SyncAll(); + } + + ProcessMainLoop(); +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::ProcessMainLoop() +{ + bool hasLoad = aicIdx < usedCoreNum; + if (!hasLoad) { + if ASCEND_IS_AIV { + if constexpr (IS_SPLIT_G) { + for (int64_t loopCnt = 0; loopCnt < maxS2LoopCnt; loopCnt++) { + CrossCoreSetFlag(15); + CrossCoreWaitFlag(15); + } + } + } + return; + } + + // 适配分核左闭右开 + uint32_t bIdx = constInfo.bN2End / constInfo.n2Size; + uint32_t qsfaActS1Size = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bIdx); + uint32_t gS1max = qsfaActS1Size; + if (constInfo.gS1End + 1 < gS1max) { + /* constInfo.gS1End != gS1max时，gS1End需要往后加一格, bN2End不变 */ + constInfo.gS1End = constInfo.gS1End + 1; + } else { + /* constInfo.gS1End == gS1max，bN2End需要往后加一格，bN2End变为0，以代表末尾 */ + constInfo.bN2End = constInfo.bN2End + 1; + constInfo.gS1End = 0; + } + + // 分核信息 + uint32_t qsfaBN2StartIdx = constInfo.bN2Start; + uint32_t bN2EndIdx = constInfo.bN2End; + uint32_t gS1StartIdx = constInfo.gS1Start; + uint32_t nextGs1Idx = constInfo.gS1End; + uint32_t s2StartIdx = 0; + uint32_t s2EndIdx = 0; + + uint32_t s2LoopLimit = 0; + if (nextGs1Idx != 0) { + bN2EndIdx++; + } + + RunInfo runInfo[3]; + RunParamStr runParam; + int64_t taskId = 0; + bool notLast = true; + int64_t multiCoreInnerIdx = 1; + for (int64_t qsfaBnIdx = qsfaBN2StartIdx; qsfaBnIdx < bN2EndIdx; qsfaBnIdx++) { + bool lastBN = (qsfaBnIdx == bN2EndIdx - 1); + runParam.boIdx = qsfaBnIdx; + runParam.n2oIdx = 0; + ComputeParamBatch(runParam, this->constInfo, + this->actualSeqQlenAddr, this->actualSeqKvlenAddr); + ComputeS1LoopInfo(runParam, this->constInfo, lastBN, nextGs1Idx, gS1StartIdx); + + int64_t gS1LoopEnd = lastBN ? (runParam.gs1LoopEndIdx + PRELOAD_NUM) : runParam.gs1LoopEndIdx; + for (int64_t gS1Index = runParam.gs1LoopStartIdx; gS1Index < gS1LoopEnd; gS1Index++) { + bool notLastTwoLoop = true; + if (lastBN) { + int32_t qsfaExtraGS1 = gS1Index - runParam.gs1LoopEndIdx; + switch (qsfaExtraGS1) { + case 0: + notLastTwoLoop = false; + break; + case 1: + notLastTwoLoop = false; + notLast = false; + break; + default: + break; + } + } + + if (notLastTwoLoop) { + this->ComputeAxisIdxByBnAndGs1(qsfaBnIdx, gS1Index, runParam); + bool s1NoNeedCalc = ComputeParamS1( + runParam, this->constInfo, gS1Index, this->actualSeqQlenAddr); + bool s2NoNeedCalc = + ComputeS2LoopInfo(runParam, this->constInfo); + // s1和s2有任意一个不需要算, 则continue, 如果是当前核最后一次循环，则补充计算taskIdx+2的部分 + if (s1NoNeedCalc || s2NoNeedCalc) { + continue; + } + if constexpr (IS_SPLIT_G) { + maxS2LoopCnt -= runParam.s2LoopEndIdx; + } + s2LoopLimit = runParam.s2LoopEndIdx - 1; + } else { + s2LoopLimit = 0; + } + + for (int64_t s2LoopCount = 0; s2LoopCount <= s2LoopLimit; ++s2LoopCount) { + if (notLastTwoLoop) { + RunInfo &runInfo1 = runInfo[taskId % 3]; + this->SetRunInfo(runInfo1, runParam, taskId, s2LoopCount, s2LoopLimit, multiCoreInnerIdx); + if ASCEND_IS_AIC { + this->cubeBlock.IterateBmm1(this->bmm1Buffers.Get(), this->l1RightBuffers.Get(), + this->v0ResGmBuffers.Get(), runInfo1, this->constInfo); + } else { + this->vecBlock.ProcessVec0(this->l1RightBuffers.Get(), this->v0ResGmBuffers.Get(), + runInfo1, this->constInfo); + } + } else { + if ASCEND_IS_AIV { + if constexpr (IS_SPLIT_G) { + if (maxS2LoopCnt > 0) { + maxS2LoopCnt--; + CrossCoreSetFlag(15); + CrossCoreWaitFlag(15); + } + } + } + } + if (taskId > 0 && notLast) { + auto &runInfo2 = runInfo[(taskId + 2) % 3]; + if ASCEND_IS_AIV { + this->vecBlock.ProcessVec1(this->l1RightBuffers.GetReused(), this->bmm1Buffers.Get(), runInfo2, + this->constInfo); + } else { + RunInfo &runInfo2 = runInfo[(taskId + 2) % 3]; + this->cubeBlock.IterateBmm2(this->bmm2Buffers.Get(), this->l1RightBuffers, this->l1RightBuffers.GetReused(), runInfo2, + this->constInfo); + } + } + if (taskId > 1) { + if ASCEND_IS_AIV { + RunInfo &qsfaRunInfo3 = runInfo[(taskId + 1) % 3]; + this->vecBlock.ProcessVec2(this->bmm2Buffers.Get(), qsfaRunInfo3, this->constInfo); + } + } + ++taskId; + } + ++multiCoreInnerIdx; + } + gS1StartIdx = 0; + } + + if ASCEND_IS_AIV { + if constexpr (IS_SPLIT_G) { + for (int64_t qsfaLoopCnt = 0; qsfaLoopCnt < maxS2LoopCnt; qsfaLoopCnt++) { + CrossCoreSetFlag(15); + CrossCoreWaitFlag(15); + } + } + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::ComputeAxisIdxByBnAndGs1( + int64_t bnIndex, int64_t gS1Index, RunParamStr &runParam) +{ + // GS1合轴, 不切G, 只切S1 + runParam.s1oIdx = gS1Index * runParam.qSNumInOneBlock; + if constexpr (IS_SPLIT_G) { + runParam.goIdx = (aicIdx % 2 == 0) ? 0 : 64; // N1=128场景，相邻cube核处理一个s1，第一个cube核承担0-63行g，第二个cube核承担后64行g + } else { + runParam.goIdx = 0; + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::SetRunInfo( + RunInfo &runInfo, RunParamStr &runParam, int64_t taskId, int64_t s2LoopCount, int64_t s2LoopLimit, int64_t multiCoreInnerIdx) +{ + if (s2LoopCount < runParam.kvLoopEndIdx) { + runInfo.s2StartIdx = runParam.s2LineStartIdx; + runInfo.s2EndIdx = runParam.s2LineEndIdx; + } + + runInfo.s2LoopCount = s2LoopCount; + + if (runInfo.multiCoreInnerIdx != multiCoreInnerIdx) { + runInfo.boIdx = runParam.boIdx; + runInfo.s1oIdx = runParam.s1oIdx; + runInfo.n2oIdx = runParam.n2oIdx; + runInfo.goIdx = runParam.goIdx; + + runInfo.multiCoreInnerIdx = multiCoreInnerIdx; + runInfo.multiCoreIdxMod2 = multiCoreInnerIdx & 1; + runInfo.multiCoreIdxMod3 = multiCoreInnerIdx % 3; + } + + runInfo.s2LoopLimit = s2LoopLimit; + runInfo.taskId = taskId; + runInfo.taskIdMod2 = taskId & 1; + runInfo.taskIdMod3 = taskId % 3; + + runInfo.sOuterOffset = runParam.sOuterOffset; + runInfo.actualS1Size = runParam.actualS1Size; + runInfo.actualS2Size = runParam.actualS2Size; + runInfo.attentionOutOffset = runParam.attentionOutOffset; + this->ComputeBmm1Tail(runInfo, runParam); + InitUniqueRunInfo(runParam, runInfo); +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::InitUniqueRunInfo( + const RunParamStr &runParam, RunInfo &runInfo) +{ + InitTaskParamByRun(runParam, runInfo); +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::ComputeBmm1Tail( + RunInfo &runInfo, RunParamStr &runParam) +{ + // ------------------------S1 Base Related--------------------------- + runInfo.s1RealSize = runParam.s1RealSize; + runInfo.halfS1RealSize = runParam.halfS1RealSize; + runInfo.firstHalfS1RealSize = runParam.firstHalfS1RealSize; + + runInfo.halfMRealSize = runParam.halfMRealSize; + runInfo.firstHalfMRealSize = runParam.firstHalfMRealSize; + runInfo.mRealSize = runParam.mRealSize; + + runInfo.vec2S1BaseSize = runInfo.halfS1RealSize; + runInfo.vec2MBaseSize = runInfo.halfMRealSize; + + // ------------------------S2 Base Related---------------------------- + runInfo.s2RealSize = constInfo.s2BaseSize; + runInfo.s2AlignedSize = runInfo.s2RealSize; + + if (runInfo.s2StartIdx + (runInfo.s2LoopCount + 1) * runInfo.s2RealSize > runInfo.s2EndIdx) { + runInfo.s2RealSize = runInfo.s2EndIdx - runInfo.s2LoopCount * runInfo.s2RealSize - runInfo.s2StartIdx; + runInfo.s2AlignedSize = Align(runInfo.s2RealSize); + } +} +} +#endif // KV_QUANT_SPARSE_FLASH_ATTENTION_KERNEL_MLA_H \ No newline at end of file diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_kvcache.h b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_kvcache.h new file mode 100644 index 000000000..9a40fb584 --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_kvcache.h @@ -0,0 +1,256 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_kvcache.h + * \brief + */ +#ifndef KV_QUANT_SPARSE_FLASH_ATTENTION_KVCACHE_H +#define KV_QUANT_SPARSE_FLASH_ATTENTION_KVCACHE_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kv_quant_sparse_flash_attention_common_arch35.h" + +using namespace matmul; +using namespace regbaseutil; +using namespace AscendC; +using namespace AscendC::Impl::Detail; +static constexpr uint32_t sparseModeThree = 3; +static constexpr uint32_t sparseModeZero = 0; + +TEMPLATE_INTF +__aicore__ inline void GetSingleCoreParam(RunParamStr& runParam, const ConstInfo &constInfo, + __gm__ int32_t *actualSeqQlenAddr, __gm__ int32_t * actualSeqKvlenAddr) +{ + int32_t qsfaActualS1Size = 0; + int32_t qsfaActualS2Size = 0; + int32_t actualSeqMin = 1; + int32_t actualSeqKVMin = 1; + int32_t sIdx = runParam.boIdx; + if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + // actual seq length first + if (actualSeqQlenAddr != nullptr) { + qsfaActualS1Size = (sIdx == 0) ? actualSeqQlenAddr[0] : + actualSeqQlenAddr[sIdx] - actualSeqQlenAddr[sIdx - 1]; + } else { + qsfaActualS1Size = constInfo.s1Size; + } + } else { + qsfaActualS1Size = (actualSeqQlenAddr == nullptr) ? constInfo.s1Size : + actualSeqQlenAddr[sIdx]; + } + + if (constInfo.isActualLenDimsKVNull) { + qsfaActualS2Size = constInfo.s2Size; + } else { + if constexpr (isPa) { + if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + qsfaActualS2Size = actualSeqKvlenAddr[sIdx]; + } else { + qsfaActualS2Size = (constInfo.actualSeqLenKVSize == actualSeqKVMin) ? + actualSeqKvlenAddr[0] : actualSeqKvlenAddr[sIdx]; + } + } else { + if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + qsfaActualS2Size = (sIdx == 0) ? actualSeqKvlenAddr[0] : + actualSeqKvlenAddr[sIdx] - actualSeqKvlenAddr[sIdx - 1]; + } else { + qsfaActualS2Size = (constInfo.actualSeqLenKVSize == actualSeqKVMin) ? + actualSeqKvlenAddr[0] : actualSeqKvlenAddr[sIdx]; + } + } + } + + runParam.actualS1Size = qsfaActualS1Size; + runParam.actualS2Size = qsfaActualS2Size; + runParam.preTokensPerBatch = runParam.actualS1Size; + if (constInfo.sparseMode == sparseModeZero) { + runParam.nextTokensPerBatch = MAX_PRE_NEXT_TOKENS; + } else { + runParam.nextTokensPerBatch = runParam.actualS2Size - runParam.actualS1Size; + } +} + +TEMPLATE_INTF +__aicore__ inline void ComputeParamBatch(RunParamStr& runParam, + const ConstInfo &constInfo, __gm__ int32_t *actualSeqQlenAddr, __gm__ int32_t *actualSeqKvlenAddr) +{ + GetSingleCoreParam(runParam, constInfo, actualSeqQlenAddr, actualSeqKvlenAddr); +} + +TEMPLATE_INTF +__aicore__ inline void ComputeS1LoopInfo(RunParamStr& runParam, const ConstInfo &constInfo, + bool lastBN, int64_t nextGs1Idx, int64_t gS1StartIdx) +{ + runParam.gs1LoopStartIdx = gS1StartIdx; + runParam.qSNumInOneBlock = 1; // qsfa 不切G轴, 计算每个基本块可以拷贝多少行s + + if (runParam.nextTokensPerBatch < 0) { + uint64_t invalidTokenCount = static_cast(-(runParam.nextTokensPerBatch + 1)) + 1ULL; + int64_t gs1LoopStartIdx = + invalidTokenCount / runParam.qSNumInOneBlock * runParam.qSNumInOneBlock; + if (gs1LoopStartIdx > gS1StartIdx) { + runParam.gs1LoopStartIdx = gs1LoopStartIdx; + } + } + + int32_t qsfaGs1LoopEndIdx = runParam.actualS1Size; // qsfa 不切G轴, 每次拷贝一行的topk，只算一行的qs + + // 不是最后一个bn, 赋值souterBlockNum + if (!lastBN) { + runParam.gs1LoopEndIdx = qsfaGs1LoopEndIdx; + } else { // 最后一个bn, 从数组下一个元素取值 + runParam.gs1LoopEndIdx = nextGs1Idx == 0 ? qsfaGs1LoopEndIdx : nextGs1Idx; + } + + if (runParam.gs1LoopStartIdx > runParam.gs1LoopEndIdx) { + runParam.gs1LoopStartIdx = runParam.gs1LoopEndIdx; + } +} + +TEMPLATE_INTF +__aicore__ inline void ComputeSouterParam(RunParamStr& runParam, const ConstInfo &constInfo, + uint32_t sOuterLoopIdx) +{ + int64_t qsfaCubeSOuterOffset = sOuterLoopIdx * runParam.qSNumInOneBlock; + if (runParam.actualS1Size == 0) { + runParam.s1RealSize = 0; + runParam.mRealSize = 0; + } else { + runParam.s1RealSize = Min(runParam.qSNumInOneBlock, runParam.actualS1Size - qsfaCubeSOuterOffset); + runParam.mRealSize = runParam.s1RealSize * constInfo.gSize; + if constexpr (IS_SPLIT_G) { + runParam.mRealSize = runParam.mRealSize >> 1; + } + } + + runParam.cubeMOuterOffset = qsfaCubeSOuterOffset * constInfo.gSize; + runParam.halfMRealSize = (runParam.mRealSize + 1) >> 1; + runParam.firstHalfMRealSize = runParam.halfMRealSize; + if (constInfo.subBlockIdx == 0) { + runParam.mOuterOffset = runParam.cubeMOuterOffset; + } else { + runParam.halfMRealSize = runParam.mRealSize - runParam.halfMRealSize; + runParam.mOuterOffset = runParam.cubeMOuterOffset + runParam.firstHalfMRealSize; + } + runParam.halfS1RealSize = (runParam.s1RealSize + 1) >> 1; + runParam.firstHalfS1RealSize = runParam.halfS1RealSize; + + if (constInfo.subBlockIdx == 1) { + runParam.halfS1RealSize = runParam.s1RealSize - runParam.halfS1RealSize; + runParam.sOuterOffset = qsfaCubeSOuterOffset + runParam.halfMRealSize / constInfo.gSize; + } else { + runParam.sOuterOffset = qsfaCubeSOuterOffset; + } + runParam.cubeSOuterOffset = qsfaCubeSOuterOffset; +} + +TEMPLATE_INTF +__aicore__ inline void LoopSOuterOffsetInit(RunParamStr& runParam, const ConstInfo &constInfo, + int32_t sIdx, __gm__ int32_t *cuSeqlensQAddr) +{ + if ASCEND_IS_AIV { + int64_t qsfaSeqOffset = 0; + if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + qsfaSeqOffset = sIdx == 0 ? 0 : cuSeqlensQAddr[sIdx - 1]; + } else { + qsfaSeqOffset = sIdx * constInfo.s1Size; + } + + int64_t attentionOutSeqOffset = qsfaSeqOffset * constInfo.n2GDv; + if constexpr (LAYOUT_T == QSFA_LAYOUT::BSND || LAYOUT_T == QSFA_LAYOUT::TND) { + runParam.attentionOutOffset = attentionOutSeqOffset + + runParam.sOuterOffset * constInfo.n2GDv + runParam.n2oIdx * constInfo.gDv + + runParam.goIdx * constInfo.dSizeV; + } + if (constInfo.subBlockIdx == 1) { + runParam.attentionOutOffset += runParam.firstHalfMRealSize * constInfo.dSizeV; + } + } +} + +TEMPLATE_INTF +__aicore__ inline bool ComputeParamS1(RunParamStr& runParam, const ConstInfo &constInfo, + uint32_t sOuterLoopIdx, __gm__ int32_t *cuSeqlensQAddr) +{ + if (runParam.nextTokensPerBatch < 0) { + uint64_t invalidTokenCount = static_cast(-(runParam.nextTokensPerBatch + 1)) + 1ULL; + if (runParam.s1oIdx < + invalidTokenCount / runParam.qSNumInOneBlock * runParam.qSNumInOneBlock) { + return true; + } + } + ComputeSouterParam(runParam, constInfo, sOuterLoopIdx); + LoopSOuterOffsetInit(runParam, constInfo, + runParam.boIdx, cuSeqlensQAddr); + return false; +} + +TEMPLATE_INTF +__aicore__ inline bool ComputeLastBN(RunParamStr& runParam, __gm__ int32_t *cuSeqlensQAddr) +{ + if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + // TND格式下 相邻Batch中当actualSeqQlen相等时则返回true + if (runParam.boIdx > 0 && ((runParam.boIdx == 0 && cuSeqlensQAddr[runParam.boIdx] == 0) || (cuSeqlensQAddr[runParam.boIdx] - cuSeqlensQAddr[runParam.boIdx - 1] == 0))) { + return true; + } + } + return false; +} + +TEMPLATE_INTF +__aicore__ inline int64_t ClipSInnerTokenCube(int64_t qsfaSInnerToken, int64_t minValue, int64_t maxValue) +{ + qsfaSInnerToken = qsfaSInnerToken > minValue ? qsfaSInnerToken : minValue; + qsfaSInnerToken = qsfaSInnerToken < maxValue ? qsfaSInnerToken : maxValue; + return qsfaSInnerToken; +} + +TEMPLATE_INTF +__aicore__ inline bool ComputeS2LoopInfo(RunParamStr& runParam, const ConstInfo &constInfo) +{ + if (runParam.actualS2Size == 0) { + runParam.kvLoopEndIdx = 0; + runParam.s2LoopEndIdx = 0; + return true; + } + uint32_t qsfaS2BaseSize = constInfo.s2BaseSize; + + if (constInfo.sparseMode == sparseModeZero) { + runParam.s2LineStartIdx = 0; + runParam.s2LineEndIdx = Min(runParam.actualS2Size, constInfo.sparseBlockCount); + } else if (constInfo.sparseMode == sparseModeThree) { + runParam.s2LineStartIdx = ClipSInnerTokenCube(runParam.cubeSOuterOffset - runParam.preTokensPerBatch, + 0, runParam.actualS2Size); + runParam.s2LineEndIdx = ClipSInnerTokenCube(runParam.cubeSOuterOffset + runParam.nextTokensPerBatch + + runParam.s1RealSize, 0, runParam.actualS2Size); + runParam.s2LineEndIdx = Min(runParam.s2LineEndIdx, constInfo.sparseBlockCount); // 当前LI输出的block size只可能是1 + } + + runParam.kvLoopEndIdx = (runParam.s2LineEndIdx + qsfaS2BaseSize - 1) / qsfaS2BaseSize; + runParam.s2LoopEndIdx = runParam.kvLoopEndIdx; + return false; +} + +TEMPLATE_INTF +__aicore__ inline void InitTaskParamByRun(const RunParamStr& runParam, RunInfo &runInfo) +{ + runInfo.boIdx = runParam.boIdx; + runInfo.actualS1Size = runParam.actualS1Size; + runInfo.actualS2Size = runParam.actualS2Size; + runInfo.preTokensPerBatch = runParam.preTokensPerBatch; + runInfo.nextTokensPerBatch = runParam.nextTokensPerBatch; + runInfo.softmaxLseOffset = runParam.softmaxLseOffset; + runInfo.qSNumInOneBlock = runParam.qSNumInOneBlock; + runInfo.kvLoopEndIdx = runParam.kvLoopEndIdx; +} + +#endif // KV_QUANT_SPARSE_FLASH_ATTENTION_KVCACHE_H diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_service_cube_mla.h b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_service_cube_mla.h new file mode 100644 index 000000000..3f7013ce3 --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_service_cube_mla.h @@ -0,0 +1,388 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_service_cube_mla.h + */ +#ifndef KV_QUANT_SPARSE_FLASH_ATTENTION_SERVICE_CUBE_MLA_H +#define KV_QUANT_SPARSE_FLASH_ATTENTION_SERVICE_CUBE_MLA_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "kv_quant_sparse_flash_attention_common_arch35.h" + +#if __has_include("../../common/op_kernel/offset_calculator.h") +#include "../../common/op_kernel/offset_calculator.h" +#else +#include "../common/offset_calculator.h" +#endif + +#if __has_include("../../common/op_kernel/matmul.h") +#include "../../common/op_kernel/matmul.h" +#else +#include "../common/matmul.h" +#endif +#if __has_include("../../common/op_kernel/CopyInL1.h") +#include "../../common/op_kernel/CopyInL1.h" +#else +#include "../common/CopyInL1.h" +#endif +#if __has_include("../../common/op_kernel/FixpipeOut.h") +#include "../../common/op_kernel/FixpipeOut.h" +#else +#include "../common/FixpipeOut.h" +#endif + +using namespace AscendC; +using namespace AscendC::Impl::Detail; + +using namespace fa_base_matmul; +using namespace regbaseutil; + +namespace BaseApi { +struct CubeCoordInfo { + uint32_t curBIdx; + uint32_t s1Coord; + uint32_t s2Coord; +}; + +template +__aicore__ inline constexpr GmFormat GetQueryGmFormat() +{ + if constexpr (LAYOUT == QSFA_LAYOUT::BSND) { + return GmFormat::BSNGD; + } else { + return GmFormat::TNGD; + } +} + +TEMPLATES_DEF +class QSFAMatmulService { +public: + /* =================编译期常量的基本块信息================= */ + static constexpr uint32_t s1BaseSize = 64; + static constexpr uint32_t s2BaseSize = 128; + static constexpr uint32_t dBaseSize = 576; + static constexpr uint32_t dBaseMatmulSize = 128; + + __aicore__ inline QSFAMatmulService() {}; + __aicore__ inline void InitCubeBlock(TPipe *pipe, BufferManager *qsfaL1BufferManagerPtr, + __gm__ uint8_t *query); + __aicore__ inline void InitCubeInput(__gm__ uint8_t *cuSeqlensQ, const ConstInfo& constInfo); + __aicore__ inline void IterateBmm1(Buffer &output, + Buffer &inputRightBuf, + Buffer &v0ResGm, + RunInfo &runInfo, ConstInfo &constInfo); + + __aicore__ inline void IterateBmm2(Buffer &outputBuf, + BuffersPolicy3buff &inputLeftBuffers, + Buffer &inputRightBuf, RunInfo &runInfo, + ConstInfo &constInfo); + +private: + __aicore__ inline void InitLocalBuffer(); + __aicore__ inline void InitGmTensor(__gm__ uint8_t *cuSeqlensQ, const ConstInfo& constInfo); + __aicore__ inline void CalcS1Coord(RunInfo &runInfo, ConstInfo &constInfo); + + __aicore__ inline void IterateBmm1QSFA(Buffer &outputBuf, + Buffer &inputRightBuf, + Buffer &v0ResGm, + RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline void PrepareLeftMatrixBmm1QSFA(Buffer &inputLeftBuf, + RunInfo &runInfo, ConstInfo &constInfo); + + // --------------------Bmm2-------------------------- + __aicore__ inline void IterateBmm2QSFA(Buffer &outputBuf, + BuffersPolicy3buff &inputLeftBuffers, + Buffer &inputRightBuf, RunInfo &runInfo, + ConstInfo &constInfo); + TPipe *tPipe; + /* =====================GM变量==================== */ + static constexpr GmFormat Q_FORMAT = GetQueryGmFormat(); + FaGmTensor queryGm; + + /* =====================运行时变量==================== */ + CubeCoordInfo coordInfo[3]; + TEventID mte1ToMte2Id[3]; + TEventID mte2ToMte1Id[3]; + + /* =====================LocalBuffer变量==================== */ + // D小于等于256 mm1左矩阵Q，GS1循环内左矩阵复用, GS1循环间开pingpong；D大于256使用单块Buffer，S1循环间驻留；fp32场景单块不驻留 + BuffersPolicySingleBuffer l1QBuffers; + // L0空间buffer manager + BufferManager *qsfaL1BufferManagerPtr; + BufferManager l0aBufferManager; + BufferManager l0bBufferManager; + BufferManager l0cBufferManager; + // L0A + BuffersPolicyDB mmL0ABuffers; + // L0B + BuffersPolicyDB mmL0BBuffers; + // L0C + BuffersPolicyDB mmL0CBuffers; +}; + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void QSFAMatmulService::InitCubeBlock( + TPipe *pipe, BufferManager *qsfaL1BuffMgr, __gm__ uint8_t *query) +{ + if ASCEND_IS_AIC { + tPipe = pipe; + qsfaL1BufferManagerPtr = qsfaL1BuffMgr; + this->queryGm.gmTensor.SetGlobalBuffer((__gm__ Q_T *)query); + InitLocalBuffer(); + } +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void +QSFAMatmulService::InitCubeInput(__gm__ uint8_t *qsfaActualSeqLengthsQ, const ConstInfo& constInfo) +{ + if ASCEND_IS_AIC { + InitGmTensor(qsfaActualSeqLengthsQ, constInfo); + if constexpr (IS_SPLIT_G) { + mte1ToMte2Id[0] = GetTPipePtr()->AllocEventID(); + mte1ToMte2Id[1] = GetTPipePtr()->AllocEventID(); + mte1ToMte2Id[2] = GetTPipePtr()->AllocEventID(); + mte2ToMte1Id[0] = GetTPipePtr()->AllocEventID(); + mte2ToMte1Id[1] = GetTPipePtr()->AllocEventID(); + mte2ToMte1Id[2] = GetTPipePtr()->AllocEventID(); + } + } +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void +QSFAMatmulService::InitLocalBuffer() +{ + constexpr uint32_t mm1LeftSize = s1BaseSize * dBaseSize * sizeof(Q_T); + l1QBuffers.Init((*qsfaL1BufferManagerPtr), mm1LeftSize); + + // L0A B C 当前写死，能否通过基础api获取 + l0aBufferManager.Init(tPipe, L0AB_SHARED_SIZE_64K); + l0bBufferManager.Init(tPipe, L0AB_SHARED_SIZE_64K); + l0cBufferManager.Init(tPipe, L0C_SHARED_SIZE_256K); + + mmL0ABuffers.Init(l0aBufferManager, BUFFER_SIZE_16K); // db类型，填入数值是总大小的一半 + mmL0BBuffers.Init(l0bBufferManager, BUFFER_SIZE_32K); + mmL0CBuffers.Init(l0cBufferManager, BUFFER_SIZE_128K); +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void +QSFAMatmulService::InitGmTensor(__gm__ uint8_t *qsfaActualSeqLengthsQ, const ConstInfo& constInfo) +{ + if constexpr (LAYOUT_T == QSFA_LAYOUT::BSND) { + this->queryGm.offsetCalculator.Init(constInfo.bSize, constInfo.n2Size, constInfo.gSize, + constInfo.s1Size, constInfo.dSize); + } else { // QSFA_LAYOUT::TND + GlobalTensor actualSeqQLen; + actualSeqQLen.SetGlobalBuffer((__gm__ int32_t *)qsfaActualSeqLengthsQ); + this->queryGm.offsetCalculator.Init(constInfo.n2Size, constInfo.gSize, constInfo.dSize, + actualSeqQLen, constInfo.actualSeqLenSize); + } +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void QSFAMatmulService::CalcS1Coord(RunInfo &runInfo, + ConstInfo &constInfo) +{ + // 计算s1方向偏移 + coordInfo[runInfo.taskIdMod3].s1Coord = runInfo.s1oIdx * runInfo.qSNumInOneBlock; +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void QSFAMatmulService::IterateBmm1( + Buffer &outputBuf, + Buffer &inputRightBuf, + Buffer &v0ResGm, + RunInfo &runInfo, ConstInfo &constInfo) +{ + CalcS1Coord(runInfo, constInfo); + + IterateBmm1QSFA(outputBuf, inputRightBuf, v0ResGm, runInfo, constInfo); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void QSFAMatmulService::IterateBmm2(Buffer &outputBuf, + BuffersPolicy3buff &inputLeftBuffers, + Buffer &inputRightBuf, RunInfo &runInfo, + ConstInfo &constInfo) +{ + IterateBmm2QSFA(outputBuf, inputLeftBuffers, inputRightBuf, runInfo, constInfo); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void QSFAMatmulService::IterateBmm1QSFA( + Buffer &outputBuf, + Buffer &inputRightBuf, + Buffer &v0ResGm, + RunInfo &runInfo, ConstInfo &constInfo) +{ + Buffer inputLeftBuf; + PrepareLeftMatrixBmm1QSFA(inputLeftBuf, runInfo, constInfo); + + // 加载当前轮的右矩阵到L1 + inputRightBuf.WaitCrossCore(); // 核间同步，这里需要根据V0操作处理同步，确保取tensor时，数据已经准备好 + + if constexpr (IS_SPLIT_G) { + SetFlag(mte2ToMte1Id[runInfo.taskIdMod3]); + WaitFlag(mte2ToMte1Id[runInfo.taskIdMod3]); + LocalTensor dst = inputRightBuf.GetTensor(); + v0ResGm.WaitCrossCore(); + GlobalTensor v0ResGmTensor = v0ResGm.template GetTensor(); + DataCopy(dst, v0ResGmTensor, Align16Func(runInfo.s2RealSize) * constInfo.dSize); + SetFlag(mte1ToMte2Id[runInfo.taskIdMod3]); + WaitFlag(mte1ToMte2Id[runInfo.taskIdMod3]); + } + + inputLeftBuf.Wait(); // 等待L1A + Buffer mm1ResL0C = mmL0CBuffers.Get(); + mm1ResL0C.Wait(); // 占用 + + MMParam param = {static_cast(runInfo.mRealSize), // singleM + static_cast(runInfo.s2RealSize), // singleN + static_cast(constInfo.dSize), // singleK + 0, // isLeftTranspose + 1 // isRightTranspose + }; + + MatmulK( // m,n不切，k切128 + inputLeftBuf.GetTensor(), inputRightBuf.GetTensor(), // mm1B直接用tensor的数据 + mmL0ABuffers, mmL0BBuffers, mm1ResL0C.GetTensor(), param); + + if (unlikely(runInfo.s2LoopCount == runInfo.s2LoopLimit)) { + inputLeftBuf.Set(); // 释放L1A + } + + mm1ResL0C.Set(); // 通知 + mm1ResL0C.Wait(); // 等待L0C + + outputBuf.WaitCrossCore(); + FixpipeParamsC310 fixpipeParams; // L0C→UB + fixpipeParams.mSize = Align2Func(runInfo.mRealSize); // 有效数据不足16行，只需要输出部分行即可; + fixpipeParams.nSize = Align8Func(runInfo.s2RealSize); // L0C上的bmm1结果矩阵N方向的size大小; 同mmadParams.n; 为什么要8个元素对齐(32B对齐) // 128 + fixpipeParams.srcStride = Align16Func(fixpipeParams.mSize); // L0C上bmm1结果相邻连续数据片段间隔(前面一个数据块的头与后面数据块的头的间隔), 单位为16*sizeof(T) // 源Nz矩阵中相邻大Z排布的起始地址偏移 + fixpipeParams.dstStride = s2BaseSize; // mmResUb上两行之间的间隔，单位：element。 // 128:根据比对dump文件得到, ND方案(S1*S2)时脏数据用mask剔除 + fixpipeParams.dualDstCtl = 1; // 双目标模式，按M维度拆分，M / 2 * N写入每个UB, M必须为2的倍数 + fixpipeParams.params.srcNdStride = 0; + fixpipeParams.params.dstNdStride = 0; + fixpipeParams.params.ndNum = 1; + + Fixpipe(outputBuf.template GetTensor(), mm1ResL0C.GetTensor(), fixpipeParams); // 将matmul结果从L0C搬运到UB + mm1ResL0C.Set(); // 释放L0C + outputBuf.SetCrossCore(); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void QSFAMatmulService::PrepareLeftMatrixBmm1QSFA( + Buffer &inputLeftBuf, + RunInfo &runInfo, ConstInfo &constInfo) +{ + // 左矩阵复用，S2的第一次循环加载左矩阵 + // 加载左矩阵到L1, 全载 + if (unlikely(runInfo.s2LoopCount == 0)) { // sOuter循环第一个基本块：搬运Q + inputLeftBuf = l1QBuffers.Get(); + inputLeftBuf.Wait(); // 占用L1A + LocalTensor inputLeftTensor = inputLeftBuf.GetTensor(); + uint64_t gmOffset = this->queryGm.offsetCalculator.GetOffset(runInfo.boIdx, runInfo.n2oIdx, runInfo.goIdx, + coordInfo[runInfo.taskIdMod3].s1Coord, 0); + CopyToL1Nd2Nz(inputLeftTensor, this->queryGm.gmTensor[gmOffset], runInfo.mRealSize, constInfo.dSize, + constInfo.mm1Ka); + + inputLeftBuf.Set(); // 通知 + } else { // 非S2的第一次循环直接复用Q + inputLeftBuf = l1QBuffers.GetPre(); + // 左矩阵复用时，sinner循环内不需要MTE2同步等待 + inputLeftBuf.Set(); // 通知 + } +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void QSFAMatmulService::IterateBmm2QSFA( + Buffer &outputBuf, + BuffersPolicy3buff &inputLeftBuffers, + Buffer &inputRightBuf, + RunInfo &runInfo, ConstInfo &constInfo) +{ + inputRightBuf.WaitCrossCore(); + Buffer mm2ResL0C = mmL0CBuffers.Get(); + mm2ResL0C.Wait(); // 占用 + + MMParam qsfaParam = {static_cast(runInfo.mRealSize), // singleM 64 + static_cast(constInfo.dSizeNope), // singleN 576->512 + static_cast(runInfo.s2RealSize), // singleK 128 + 0, 0}; + + MatmulN( + inputRightBuf.GetTensor(s2BaseSize * constInfo.dSizeNope), // 左矩阵P 来自rope位置 + inputRightBuf.GetTensor(), // 右矩阵V nope + mmL0ABuffers, mmL0BBuffers, + mm2ResL0C.GetTensor(), qsfaParam); + + inputRightBuf.SetCrossCore(); // bmm2才释放KV，在这里释放 + mm2ResL0C.Set(); // 通知 + mm2ResL0C.Wait(); // 等待 + + outputBuf.WaitCrossCore(); //占用 + + FixpipeParamsC310 fixpipeParams; // L0C→UB;FixpipeParamsM300:L0C→UB + fixpipeParams.mSize = Align2Func(runInfo.mRealSize); // 有效数据不足16行，只需要输出部分行即可; + fixpipeParams.nSize = Align8Func(constInfo.dSizeNope); // L0C上的bmm1结果矩阵N方向的size大小, 分档计算且vector2中通过mask筛选出实际有效值 + fixpipeParams.srcStride = Align16Func(fixpipeParams.mSize); // L0C上bmm1结果相邻连续数据片段间隔（前面一个数据块的头与后面数据块的头的间隔） + fixpipeParams.dstStride = Align16Func(constInfo.dSizeNope); + fixpipeParams.dualDstCtl = 1; + fixpipeParams.params.srcNdStride = 0; + fixpipeParams.params.dstNdStride = 0; + fixpipeParams.params.ndNum = 1; + + Fixpipe(outputBuf.template GetTensor(), mm2ResL0C.GetTensor(), fixpipeParams); // 将matmul结果从L0C搬运到UB + mm2ResL0C.Set(); // 释放 + + outputBuf.SetCrossCore(); +} + +TEMPLATES_DEF +class QSFAMatmulServiceDummy { +public: + __aicore__ inline QSFAMatmulServiceDummy() {}; + __aicore__ inline void InitCubeBlock(TPipe *pipe, BufferManager *qsfaL1BufferManagerPtr, + __gm__ uint8_t *query) {} + __aicore__ inline void InitCubeInput(__gm__ uint8_t *cuSeqlensQ, const ConstInfo& constInfo) {} + __aicore__ inline void IterateBmm1(Buffer &outputBuf, + Buffer &inputRightBuf, RunInfo &runInfo, + ConstInfo &constInfo) {} + + __aicore__ inline void IterateBmm2(Buffer &outputBuf, + BuffersPolicyDB &inputLeftBuffers, + Buffer &inputRightBuf, RunInfo &runInfo, + ConstInfo &constInfo) {} +}; + +template +struct CubeBlockTraits; // 声明 +/* 生成CubeBlockTraits */ +#define GEN_TRAIT_TYPE(name, ...) using name##_TRAITS = name; +#define GEN_TRAIT_CONST(name, type, ...) static constexpr type name##Traits = name; + +#define DEFINE_QSFA_CUBE_BLOCK_TRAITS(CUBE_BLOCK_CLASS) \ + TEMPLATES_DEF_NO_DEFAULT \ + struct CubeBlockTraits> { \ + QSFA_CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_TRAIT_TYPE) \ + QSFA_CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_TRAIT_CONST) \ + } + +DEFINE_QSFA_CUBE_BLOCK_TRAITS(QSFAMatmulService); +DEFINE_QSFA_CUBE_BLOCK_TRAITS(QSFAMatmulServiceDummy); + +// /* 生成Arg Traits, kernel中只需要调用ARGS_TRAITS就可以获取所有CubeBlock中的模板参数 */ +#define GEN_ARGS_TYPE(name, ...) using name = typename CubeBlockTraits::name##_TRAITS; +#define GEN_ARGS_CONST(name, type, ...) static constexpr type name = CubeBlockTraits::name##Traits; +#define ARGS_TRAITS \ + QSFA_CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_ARGS_TYPE) \ + QSFA_CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_ARGS_CONST) +} +#endif // KV_QUANT_SPARSE_FLASH_ATTENTION_SERVICE_CUBE_MLA_H \ No newline at end of file diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_service_vector_mla.h b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_service_vector_mla.h new file mode 100644 index 000000000..ebff32a68 --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/arch35/kv_quant_sparse_flash_attention_service_vector_mla.h @@ -0,0 +1,894 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_service_vector_mla.h + * \brief + */ +#ifndef KV_QUANT_SPARSE_FLASH_ATTENTION_SERVICE_VECTOR_MLA_H +#define KV_QUANT_SPARSE_FLASH_ATTENTION_SERVICE_VECTOR_MLA_H + +#include "kv_quant_sparse_flash_attention_common_arch35.h" +#include "kernel_operator_list_tensor_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "lib/matmul_intf.h" +#if __has_include("../../common/op_kernel/arch35/vf/vf_mul_sel_softmaxflashv2_cast_nz_sfa.h") +#include "../../common/op_kernel/arch35/vf/vf_mul_sel_softmaxflashv2_cast_nz_sfa.h" +#else +#include "../../common/arch35/vf/vf_mul_sel_softmaxflashv2_cast_nz_sfa.h" +#endif +#if __has_include("../../common/op_kernel/arch35/vf/vf_flashupdate_new.h") +#include "../../common/op_kernel/arch35/vf/vf_flashupdate_new.h" +#else +#include "../../common/arch35/vf/vf_flashupdate_new.h" +#endif + +using namespace AscendC; +using namespace FaVectorApi; +using namespace AscendC::Impl::Detail; +using namespace regbaseutil; +using namespace matmul; + +namespace BaseApi { + +TEMPLATES_DEF +class QSFAVectorService { +public: + // BUFFER的字节数 + static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; + /* =================编译期常量的基本块信息================= */ + static constexpr uint32_t s1BaseSize = 64; + static constexpr uint32_t s2BaseSize = 128; + static constexpr uint32_t vec1Srcstride = (s1BaseSize >> 1) + 1; + static constexpr uint32_t dVTemplateType = 512; + static constexpr uint32_t qsfaDTemplateAlign64 = Align64Func(dVTemplateType); + static constexpr uint32_t dVTemplateTypeInput = 672; + static constexpr float R0 = 1.0f; + static constexpr uint64_t SYNC_SINKS_BUF_FLAG = 6; + + // ==================== Functions ====================== + __aicore__ inline QSFAVectorService() {}; + __aicore__ inline void InitVecBlock(TPipe *pipe, const KvQuantSparseFlashAttentionTilingDataMla *__restrict tiling, + CVSharedParams &sharedParams, int32_t aicIdx, uint8_t subBlockIdx, __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths) + { + if ASCEND_IS_AIV { + tilingData = tiling; + tPipe = pipe; + if (actualSeqLengths != nullptr) { + actualSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengths); + } + if (actualSeqLengthsQ != nullptr) { + cuSeqlensQGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsQ); + } + + this->InitCubeVecSharedParams(sharedParams, aicIdx, subBlockIdx); + this->GetExtremeValue(this->negativeFloatScalar); + } + } + + // 初始化LocalTensor + __aicore__ inline void InitLocalBuffer(TPipe *pipe, ConstInfo &constInfo); + // 初始化attentionOutGM + __aicore__ inline void CleanOutput(__gm__ uint8_t *attentionOut, ConstInfo &constInfo); + __aicore__ inline void InitGlobalBuffer(__gm__ uint8_t *key, __gm__ uint8_t *value, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *blockTable); + __aicore__ inline void InitOutputSingleCore(ConstInfo &constInfo); + __aicore__ inline void ProcessVec0(Buffer &outputL1, + Buffer &v0ResGm, + const RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline void ProcessVec1(Buffer &outputBuf, + Buffer &bmm1ResBuf, RunInfo &runInfo, + ConstInfo &constInfo); + using mm2ResPos = Buffer; + __aicore__ inline void ProcessVec2(mm2ResPos &bmm2ResBuf, RunInfo &runInfo, + ConstInfo &constInfo); + +private: + __aicore__ inline void ProcessVec1SoftmaxDispatchQSFA(LocalTensor &stage1CastTensor, + LocalTensor &mmRes, LocalTensor &sumUb, LocalTensor &maxUb, + LocalTensor &apiTmpBuffer, RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline void ProcessSparseKv(Buffer &outputL1, + Buffer &v0ResGm, + const RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline void CalSparseCalSize(const RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline int64_t GetkeyOffset(int64_t s2Idx, const RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline void GetRealCmpS2Idx(int64_t &token0Idx, int64_t &token1Idx, int64_t s2IdxInBase, + const RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline void CopyInKvNotSparse(LocalTensor kvMergUb, int64_t v0Loop, int64_t dealRow, + int64_t s2StartIdx, const RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline uint32_t CopyInKvSparse(LocalTensor kvInUb , int64_t startRow, int64_t token0Idx, + int64_t token1Idx, const RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline void DequantKv(LocalTensor antiKvTensorAsB16, LocalTensor srcTensor, int64_t dealRow, + ConstInfo &constInfo); + __aicore__ inline void CopyOutKvUb2L1(Buffer &outputL1, + LocalTensor antiKvTensorAsB16, int64_t dealRow, int64_t s2StartIdx, + const RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline void CopyOutKvUb2Gm(Buffer &v0ResGm, + LocalTensor antiKvTensorAsB16, int64_t dealRow, int64_t s2StartIdx, const RunInfo &runInfo, + ConstInfo &constInfo); + __aicore__ inline void CopyOutMrgeResult(Buffer &outputL1, + int64_t mte2Size, int64_t mte3Size, int64_t s2keyOffset, int64_t mergeMte3Idx, const RunInfo &runInfo); + __aicore__ inline void CopyInSingleKv(LocalTensor kvInUb, int64_t startRow, + int64_t keyOffset, uint32_t combineBytes); + /* VEC2_RES_T 表示bmm2ResUb当前的类型，VEC2_RES_T = Q_T那么不需要做Cast。另外，无效行场景当前默认需要做Cast */ + using VEC2_RES_T = T; + template + __aicore__ inline void Bmm2DataCopyOut(RunInfo &runInfo, ConstInfo &constInfo, + LocalTensor &vec2ResUb, int64_t vec2S1Idx, int64_t qsfaVec2CalcSize = 0); + template + __aicore__ inline void CopyOutAttentionOut( + RunInfo &runInfo, ConstInfo &constInfo, LocalTensor &vec2ResUb, int64_t vec2S1Idx, + int64_t qsfaVec2CalcSize); + __aicore__ inline void SoftmaxInitBuffer(); + __aicore__ inline void InitCubeVecSharedParams(CVSharedParams &sharedParams, int32_t aicIdx, uint8_t subBlockIdx); + __aicore__ inline void ComputeNeedInitQSFA(CVSharedParams &sharedParams) const; + __aicore__ inline void GetExtremeValue(T &negativeScalar); + + TPipe *tPipe; + const KvQuantSparseFlashAttentionTilingDataMla *__restrict tilingData; + + GlobalTensor attentionOutGm; + GlobalTensor keyGm; + GlobalTensor SparseIndicesGm; + GlobalTensor blockTableGm; + GlobalTensor cuSeqlensQGm; + GlobalTensor actualSeqLengthsKVGm; + + TBuf<> commonTBuf; // common的复用空间 + TQue stage1OutQue[2]; // 2份表示可能存在pingpong + TQue stage0InQue; // for v0 input, 2份表示可能存在pingpong + TQue stage0OutQue; // for v0 output, 2份表示可能存在pingpong + TBuf<> stage2OutBuf; + TEventID mte3ToVId[2]; // 存放MTE3_V的eventId, 2份表示可能存在pingpong + TEventID vToMte3Id[2]; // 存放V_MTE3的eventId, 2份表示可能存在pingpong + TBuf<> softmaxMaxBuf[2]; + TBuf<> softmaxSumBuf[2]; + TBuf<> softmaxExpBuf[2]; + + T negativeFloatScalar; + uint32_t maxBlockNumPerBatch; + uint32_t blockSize; + int64_t qsfaSparseCalSize; + int64_t sparseS2Start; + int64_t sparseS2End; +}; + + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void QSFAVectorService::GetRealCmpS2Idx(int64_t &token0Idx, int64_t &token1Idx, + int64_t s2IdxInBase, const RunInfo &runInfo, ConstInfo &constInfo) +{ + int64_t topkBS1Idx = 0; + if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + uint64_t actualSeqQPrefixSum = runInfo.boIdx == 0 ? 0 : cuSeqlensQGm.GetValue(runInfo.boIdx - 1); + topkBS1Idx += (actualSeqQPrefixSum + runInfo.s1oIdx) * constInfo.sparseBlockCount; // T, N2(1), K + } else { + topkBS1Idx += runInfo.boIdx * constInfo.s1Size * constInfo.sparseBlockCount + + runInfo.s1oIdx * constInfo.sparseBlockCount; // B, S1, N2(1), K + } + + int64_t qsfaCmpS2LoopCnt = runInfo.s2LoopCount; + int64_t qsfaTopkIdx = s2IdxInBase + qsfaCmpS2LoopCnt * constInfo.s2BaseSize; + + if (unlikely(qsfaTopkIdx >= constInfo.sparseBlockCount)) { + token0Idx = -1; + } else { + token0Idx = SparseIndicesGm.GetValue(topkBS1Idx + qsfaTopkIdx) + runInfo.s2StartIdx; + } + qsfaTopkIdx += 1; + if (unlikely((qsfaTopkIdx >= constInfo.sparseBlockCount) || (s2IdxInBase + 1 >= sparseS2End))) { + token1Idx = -1; + } else { + token1Idx = SparseIndicesGm.GetValue(topkBS1Idx + qsfaTopkIdx) + runInfo.s2StartIdx; + } +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline int64_t QSFAVectorService::GetkeyOffset(int64_t s2Idx, const RunInfo &runInfo, ConstInfo &constInfo) +{ + if (s2Idx < 0) { + return -1; + } + int64_t realkeyOffset = 0; + if constexpr (isPa) { + int64_t blkTableIdx = s2Idx / blockSize; + int64_t blkTableOffset = s2Idx % blockSize; + realkeyOffset = blockTableGm.GetValue(runInfo.boIdx * maxBlockNumPerBatch + blkTableIdx) * + static_cast(blockSize) * constInfo.dSizeVInput + + blkTableOffset * constInfo.dSizeVInput; // BlockNum, BlockSize, N(1), D + } else { + if constexpr (LAYOUT_T == QSFA_LAYOUT::BSND) { + realkeyOffset = (runInfo.boIdx * constInfo.s2Size + s2Idx) * constInfo.dSizeVInput; // BSN(1)D + } else if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + int64_t batchKvStart = (runInfo.boIdx == 0) ? 0 : actualSeqLengthsKVGm.GetValue(runInfo.boIdx - 1); + realkeyOffset = (batchKvStart + s2Idx) * constInfo.dSizeVInput; + } + } + return realkeyOffset; +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void +QSFAVectorService::CopyInSingleKv(LocalTensor kvInUb, int64_t startRow, + int64_t keyOffset, uint32_t combineBytes) +{ + if (keyOffset < 0) { + return; + } + DataCopyExtParams intriParams; + + intriParams.blockCount = 1; + intriParams.dstStride = 0; + intriParams.srcStride = 0; + DataCopyPadExtParams padParams; + // 当前仅支持COMBINE模式 + intriParams.blockLen = combineBytes; + uint32_t combineDim = combineBytes / sizeof(KV_T); + uint32_t combineDimAlign = CeilAlign(combineBytes, BUFFER_SIZE_BYTE_32B) / sizeof(KV_T); + padParams.isPad = true; + padParams.leftPadding = 0; + padParams.rightPadding = combineDimAlign - combineDim; + padParams.paddingValue = 0; + DataCopyPad(kvInUb[startRow * combineDimAlign], keyGm[keyOffset], intriParams, padParams); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline uint32_t QSFAVectorService::CopyInKvSparse(LocalTensor kvInUb , int64_t startRow, + int64_t token0Idx, int64_t token1Idx, const RunInfo &runInfo, ConstInfo &constInfo) +{ + int64_t keyOffset0 = GetkeyOffset(token0Idx, runInfo, constInfo); + int64_t keyOffset1 = GetkeyOffset(token1Idx, runInfo, constInfo); + if (unlikely(keyOffset0 < 0 && keyOffset1 < 0)) { + return 0; + } + uint32_t combineBytes = constInfo.dSizeVInput * sizeof(KV_T); + int64_t keySrcStride = (keyOffset0 > keyOffset1 ? (keyOffset0 - keyOffset1) * sizeof(KV_T): + (keyOffset1 - keyOffset0)) * sizeof(KV_T) - combineBytes; + if (keySrcStride >= INT32_MAX || keySrcStride < 0 || constInfo.sparseBlockSize > 1) { + // stride溢出、stride为负数、s2超长等异常场景，还原成2条搬运指令 + CopyInSingleKv(kvInUb, startRow, keyOffset0, combineBytes); + CopyInSingleKv(kvInUb, startRow + 1, keyOffset1, combineBytes); + } else { + DataCopyExtParams intriParams; + intriParams.blockCount = (keyOffset0 >= 0) + (keyOffset1 >= 0); + intriParams.blockLen = combineBytes; + intriParams.dstStride = 0; + intriParams.srcStride = keySrcStride; + DataCopyPadExtParams padParams; + + int64_t keyOffset = keyOffset0 > -1 ? keyOffset0 : keyOffset1; + if (keyOffset1 > -1 && keyOffset1 < keyOffset0) { + keyOffset = keyOffset1; + } + + // 当前仅支持COMBINE模式 + uint32_t combineDim = combineBytes / sizeof(KV_T); + uint32_t combineDimAlign = CeilAlign(combineBytes, BUFFER_SIZE_BYTE_32B) / sizeof(KV_T); + padParams.isPad = true; + padParams.leftPadding = 0; + padParams.rightPadding = combineDimAlign - combineDim; + padParams.paddingValue = 0; + DataCopyPad(kvInUb[startRow * combineDimAlign], keyGm[keyOffset], intriParams, padParams); + } + return (keyOffset0 > -1) + (keyOffset1 > -1); +} + +// fp8->fp32 +static constexpr MicroAPI::CastTrait castTraitFp8_1 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + +// fp32->fp16 +static constexpr MicroAPI::CastTrait castTraitFp8_3 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT}; + +// int8->half +static constexpr MicroAPI::CastTrait castTraitint8_1 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + +// half->fp32 +static constexpr MicroAPI::CastTrait castTraithalf_1 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + +template +__simd_vf__ void AntiquantVFImplFp8D448(__ubuf__ int8_t* ubSrcAddr, __ubuf__ Q_T* ubDstAddr, // output first + __ubuf__ float* ubScaleSrcAddr, uint32_t dealRowCount) +{ + uint32_t combineDim = 672; // 128对齐 640->672 + MicroAPI::RegTensor vKvData0; + MicroAPI::RegTensor vKvData1; + MicroAPI::RegTensor vKvDataHalf0; + MicroAPI::RegTensor vKvDataHalf1; + MicroAPI::RegTensor vCastHalfRes0; + MicroAPI::RegTensor vCastHalfRes1; + MicroAPI::RegTensor vCastFp32Res0; + MicroAPI::RegTensor vCastFp32Res1; + MicroAPI::RegTensor vMulRes0; + MicroAPI::RegTensor vMulRes1; + MicroAPI::RegTensor vScale0; + MicroAPI::RegTensor vScale1; + MicroAPI::RegTensor vCastRes0; + MicroAPI::RegTensor vCastRes1; + MicroAPI::RegTensor vCastResPack0; + MicroAPI::RegTensor vCastResPack1; + + MicroAPI::MaskReg kvTypeMaskAll = MicroAPI::CreateMask(); + MicroAPI::MaskReg kvRopeTypeMaskAll = MicroAPI::CreateMask(); + MicroAPI::MaskReg int8MaskAll = MicroAPI::CreateMask(); + MicroAPI::MaskReg fp32MaskAll = MicroAPI::CreateMask(); + uint32_t blockStride = 17; // +1 to solve bank conflict + uint32_t repeatStride = 1; + const uint32_t nopeDim = 512; // 448->512 64 + const uint32_t kvNumPerLoop = 128; + const uint32_t scaleNumPerLoop = 1; + const uint32_t tileSize = 128; + static constexpr bool isKvInt8 = (IsSameType::value); + // tilesize is 128, deal 128 b8 kv, deal 1 fp32 scale + for (uint16_t j = 0; j < (nopeDim / kvNumPerLoop); j++) { + __ubuf__ int8_t* ubSrcTemp = ubSrcAddr + j * kvNumPerLoop; + __ubuf__ float* ubScaleSrcAddrTemp = ubScaleSrcAddr + j * scaleNumPerLoop; + __ubuf__ Q_T* ubDstAddrTmp = ubDstAddr + j * kvNumPerLoop * blockStride; + for (uint16_t i = 0; i < static_cast(dealRowCount); i++) { + // load scale + MicroAPI::LoadAlign( + (MicroAPI::RegTensor&)vKvData0, ubSrcTemp, tileSize / 2); + MicroAPI::LoadAlign( + (MicroAPI::RegTensor&)vKvData1, ubSrcTemp, combineDim - tileSize / 2); + + MicroAPI::LoadAlign( + (MicroAPI::RegTensor&)vScale0, ubScaleSrcAddrTemp, combineDim / 4); + + if constexpr (isKvInt8) { + // int8 -> half + MicroAPI::Cast(vCastHalfRes0, vKvData0, int8MaskAll); + MicroAPI::Cast(vCastHalfRes1, vKvData1, int8MaskAll); + // half -> float + MicroAPI::Cast(vCastFp32Res0, vCastHalfRes0, fp32MaskAll); + MicroAPI::Cast(vCastFp32Res1, vCastHalfRes1, fp32MaskAll); + } else { + MicroAPI::Cast(vCastFp32Res0, vKvData0, fp32MaskAll); + MicroAPI::Cast(vCastFp32Res1, vKvData1, fp32MaskAll); + } + + MicroAPI::Mul(vMulRes0, vCastFp32Res0, vScale0, fp32MaskAll); + MicroAPI::Mul(vMulRes1, vCastFp32Res1, vScale0, fp32MaskAll); + + MicroAPI::Cast(vCastRes0, vMulRes0, fp32MaskAll); + MicroAPI::Cast(vCastRes1, vMulRes1, fp32MaskAll); + + MicroAPI::DeInterleave(vCastResPack0, vCastResPack1, vCastRes0, vCastRes1); + + MicroAPI::StoreAlign( + ubDstAddrTmp, vCastResPack0, blockStride, repeatStride, kvRopeTypeMaskAll); + } + } +} + +template +__aicore__ inline void AntiquantVFFp8D448(LocalTensor& outputUb, LocalTensor& inputUb, uint32_t dealRowCount) +{ + __ubuf__ int8_t* ubSrcAddr = (__ubuf__ int8_t*)(inputUb.GetPhyAddr()); // nope改成在左，所以起始位置是0 + __ubuf__ Q_T* ubDstAddr = (__ubuf__ Q_T*)(outputUb.GetPhyAddr()); + __ubuf__ float* ubScaleAddr = (__ubuf__ float*)(inputUb[512 + 64 * 2].GetPhyAddr()); + + AntiquantVFImplFp8D448(ubSrcAddr, ubDstAddr, ubScaleAddr, dealRowCount); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void QSFAVectorService::DequantKv(LocalTensor antiKvTensorAsB16, + LocalTensor srcTensor, int64_t dealRow, ConstInfo &constInfo) +{ + // srcTensor是nope(512) + nope(64) + scale + pad, dstTensor是nope(512) + rope(64) + AntiquantVFFp8D448(antiKvTensorAsB16, srcTensor, dealRow); + + LocalTensor kRopeUb = srcTensor[constInfo.dSizeNope].template ReinterpretCast(); + LocalTensor kRopeUbNz = antiKvTensorAsB16[constInfo.dSizeNope * (16 + 1)]; // V0单次处理16行数据 + Copy(kRopeUbNz, kRopeUb, + constInfo.dSizeRope, // mask 处理多少列数据 + static_cast(dealRow), // repeatTime, 每次处理多少个block + { + 17, // dst stride + 1, // src stride + 1, // dst repeat stride + 21 // src repeat stride, 640 / 32 // 640 -> 672 : 20 -> 21 + }); +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void QSFAVectorService::CopyOutKvUb2L1( + Buffer &outputL1, + LocalTensor antiKvTensorAsB16, int64_t dealRow, int64_t s2StartIdx, + const RunInfo &runInfo, ConstInfo &constInfo) +{ + uint64_t blockElementNum = 16; + DataCopyParams dataCopyParams; + dataCopyParams.blockCount = (constInfo.dSizeNope + constInfo.dSizeRope) / blockElementNum; + dataCopyParams.blockLen = dealRow; + dataCopyParams.srcGap = blockElementNum + 1 - dealRow; + dataCopyParams.dstGap = Align16Func(runInfo.s2RealSize) - dealRow; + + LocalTensor dst = outputL1.GetTensor(); + DataCopy(dst[s2StartIdx * blockElementNum], antiKvTensorAsB16, dataCopyParams); +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void QSFAVectorService::CopyOutKvUb2Gm( + Buffer &v0ResGm, LocalTensor antiKvTensorAsB16, + int64_t dealRow, int64_t s2StartIdx, const RunInfo &runInfo, ConstInfo &constInfo) +{ + GlobalTensor v0ResGmTensor = v0ResGm.template GetTensor(); + uint64_t blockElementNum = 16; + DataCopyParams dataCopyParams; + dataCopyParams.blockCount = (constInfo.dSizeNope + constInfo.dSizeRope) / blockElementNum; + dataCopyParams.blockLen = dealRow; + dataCopyParams.srcGap = blockElementNum + 1 - dealRow; + dataCopyParams.dstGap = Align16Func(runInfo.s2RealSize) - dealRow; + DataCopy(v0ResGmTensor[s2StartIdx * blockElementNum], antiKvTensorAsB16, dataCopyParams); +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void QSFAVectorService::CalSparseCalSize(const RunInfo &runInfo, ConstInfo &constInfo) +{ + if constexpr (IS_SPLIT_G) { + uint32_t aicIdx = constInfo.aivIdx >> 1U; + uint32_t v0S2SizeFirstCore = CeilDiv(runInfo.s2RealSize, 2); + uint32_t v0S2SizeSecondCore = runInfo.s2RealSize - v0S2SizeFirstCore; + if (aicIdx % 2U == 0) { + if (GetSubBlockIdx() == 0) { + qsfaSparseCalSize = CeilDiv(v0S2SizeFirstCore, 2); // 2: Vector split size for first core (first half) + sparseS2Start = 0; + } else { + // 2: Vector split size for first core (second half) + qsfaSparseCalSize = v0S2SizeFirstCore - CeilDiv(v0S2SizeFirstCore, 2); + sparseS2Start = CeilDiv(v0S2SizeFirstCore, 2); // 2: Start offset for second half of first core + } + } else { + if (GetSubBlockIdx() == 0) { + qsfaSparseCalSize = CeilDiv(v0S2SizeSecondCore, 2); // 2: Same as above + sparseS2Start = v0S2SizeFirstCore; + } else { + qsfaSparseCalSize = v0S2SizeSecondCore - CeilDiv(v0S2SizeSecondCore, 2); // 2: Same as above + sparseS2Start = v0S2SizeFirstCore + CeilDiv(v0S2SizeSecondCore, 2); // 2: Same as above + } + } + sparseS2End = sparseS2Start + qsfaSparseCalSize; + } else { + int64_t s2PerVecLoop = 2LL; + int64_t vecNum = 2LL; + int64_t s2Loops = CeilDiv(CeilDiv(runInfo.s2RealSize, vecNum), s2PerVecLoop); + sparseS2Start = GetSubBlockIdx() == 0 ? 0 : s2Loops * s2PerVecLoop; + sparseS2End = GetSubBlockIdx() == 0 ? s2Loops * s2PerVecLoop : runInfo.s2RealSize; + qsfaSparseCalSize = sparseS2End - sparseS2Start; + } +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void QSFAVectorService::ProcessVec0( + Buffer &outputL1, + Buffer &v0ResGm, + const RunInfo &runInfo, ConstInfo &constInfo) +{ + outputL1.WaitCrossCore(); // 核间同步 + blockSize = constInfo.blockSize; + maxBlockNumPerBatch = constInfo.maxBlockNumPerBatch; + + CalSparseCalSize(runInfo, constInfo); + ProcessSparseKv(outputL1, v0ResGm, runInfo, constInfo); + + if constexpr (IS_SPLIT_G) { + CrossCoreSetFlag(15); // 15: 跨核同步标志位值 + CrossCoreWaitFlag(15); // 15: 跨核同步标志位值 + } + + outputL1.SetCrossCore(); // 核间同步 + if constexpr (IS_SPLIT_G) { + v0ResGm.SetCrossCore(); + } +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void QSFAVectorService::ProcessSparseKv( + Buffer &outputL1, + Buffer &v0ResGm, const RunInfo &runInfo, ConstInfo &constInfo) +{ + if (qsfaSparseCalSize == 0) { + return; + } + // Left-closed, right-open interval + // 4x = 2x + 2x + // 4x + 1 = (2x + 2) + (2x - 1) + // 4x + 2 = (2x + 2) + (2x) + // 4x + 3 = (2x + 2) + (2x + 1) + int64_t s2Start = sparseS2Start; + int64_t s2 = sparseS2Start; + bool meetEnd = false; + int64_t token0Idx, token1Idx; // 拷贝进入的两个token的index + // 处理一个s2的base块 + while ((s2 < sparseS2End) && !meetEnd) { // 拷贝到s2End或者遇到-1 + int64_t dealRow = 0; + // 1、copy kv in, gm ->ub + LocalTensor kvInUb = stage0InQue.AllocTensor(); + while (dealRow < Min(16, qsfaSparseCalSize) && s2(); + + // 2、dequant by vf + LocalTensor kvDequantOutUb = stage0OutQue.AllocTensor(); + DequantKv(kvDequantOutUb, kvInUb, dealRow, constInfo); + stage0InQue.FreeTensor(kvInUb); + stage0OutQue.EnQue(kvDequantOutUb); + kvDequantOutUb = stage0OutQue.DeQue(); + + // 3、copy kv out, ub -> l1 + if constexpr (IS_SPLIT_G) { + CopyOutKvUb2Gm(v0ResGm, kvDequantOutUb, dealRow, s2Start, runInfo, constInfo); + } else { + CopyOutKvUb2L1(outputL1, kvDequantOutUb, dealRow, s2Start, runInfo, constInfo); + } + s2Start += dealRow; + stage0OutQue.FreeTensor(kvDequantOutUb); + } +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void QSFAVectorService::ProcessVec1( + Buffer &outputBuf, + Buffer &bmm1ResBuf, RunInfo &runInfo, + ConstInfo &constInfo) +{ + bmm1ResBuf.WaitCrossCore(); + + LocalTensor sumUb = this->softmaxSumBuf[runInfo.multiCoreIdxMod2].template Get(); + LocalTensor maxUb = this->softmaxMaxBuf[runInfo.multiCoreIdxMod2].template Get(); + LocalTensor qsfaExpUb = this->softmaxExpBuf[runInfo.taskIdMod2].template Get(); + int64_t stage1Offset = runInfo.taskIdMod2; + auto stage1CastTensor = this->stage1OutQue[stage1Offset].template AllocTensor(); + + LocalTensor apiTmpBuffer = this->commonTBuf.template Get(); + LocalTensor mmRes = bmm1ResBuf.template GetTensor(); + + ProcessVec1SoftmaxDispatchQSFA(stage1CastTensor, mmRes, sumUb, maxUb, apiTmpBuffer, runInfo, constInfo); + + bmm1ResBuf.SetCrossCore(); + // ===================DataCopy to L1 ==================== + this->stage1OutQue[stage1Offset].template EnQue(stage1CastTensor); + this->stage1OutQue[stage1Offset].template DeQue(); + + LocalTensor mm2AL1Tensor = + outputBuf.GetTensor(s2BaseSize * constInfo.dSizeV); + + if (likely(runInfo.halfMRealSize != 0)) { + DataCopy(mm2AL1Tensor[constInfo.subBlockIdx * (BLOCK_BYTE / sizeof(Q_T)) * (runInfo.mRealSize - runInfo.halfMRealSize)], + stage1CastTensor, {s2BaseSize / 16, (uint16_t)runInfo.halfMRealSize, + (uint16_t)(vec1Srcstride - runInfo.halfMRealSize), + (uint16_t)(Align16Func(runInfo.mRealSize) - runInfo.halfMRealSize)}); + } + + this->stage1OutQue[stage1Offset].template FreeTensor(stage1CastTensor); + + outputBuf.SetCrossCore(); + if (runInfo.s2LoopCount != 0) { + SFAUpdateExpSumAndExpMax(sumUb, maxUb, qsfaExpUb, sumUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize); + } +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void QSFAVectorService::ProcessVec1SoftmaxDispatchQSFA( + LocalTensor &stage1CastTensor, LocalTensor &mmRes, LocalTensor &sumUb, + LocalTensor &maxUb, LocalTensor &apiTmpBuffer, RunInfo &runInfo, + ConstInfo &constInfo) +{ + if (runInfo.s2LoopCount == 0) { + if (likely(runInfo.s2RealSize == 128)) { // s2RealSize等于128分档, VF内常量化减少if判断 + ProcessVec1Vf( + stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize, runInfo.s2RealSize, + static_cast(constInfo.softmaxScale), negativeFloatScalar); + } else if (runInfo.s2RealSize <= 64) { // s2RealSize小于等于64分档, VF内常量化减少if判断 + ProcessVec1Vf( + stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize, runInfo.s2RealSize, + static_cast(constInfo.softmaxScale), negativeFloatScalar); + } else if (runInfo.s2RealSize < 128) { // s2RealSize小于128分档, VF内常量化减少if判断 + ProcessVec1Vf( + stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize, + runInfo.s2RealSize, static_cast(constInfo.softmaxScale), negativeFloatScalar); + } + } else { + if (likely(runInfo.s2RealSize == 128)) { // s2RealSize等于128分档, VF内常量化减少if判断 + ProcessVec1Vf( + stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize, + runInfo.s2RealSize, static_cast(constInfo.softmaxScale), negativeFloatScalar); + } else if (runInfo.s2RealSize <= 64) { // s2RealSize小于等于64分档, VF内常量化减少if判断 + ProcessVec1Vf( + stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize, + runInfo.s2RealSize, static_cast(constInfo.softmaxScale), negativeFloatScalar); + } else if (runInfo.s2RealSize < 128) { // s2RealSize小于128分档, VF内常量化减少if判断 + ProcessVec1Vf( + stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize, + runInfo.s2RealSize, static_cast(constInfo.softmaxScale), negativeFloatScalar); + } + } +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void QSFAVectorService::ProcessVec2( + Buffer &bmm2ResBuf, RunInfo &runInfo, + ConstInfo &constInfo) +{ + bmm2ResBuf.WaitCrossCore(); + + if (unlikely(runInfo.vec2MBaseSize == 0)) { + bmm2ResBuf.SetCrossCore(); + return; + } + runInfo.vec2MRealSize = runInfo.vec2MBaseSize; + runInfo.vec2S1RealSize = runInfo.vec2S1BaseSize; + int64_t qsfaVec2CalcSize = runInfo.vec2MRealSize * qsfaDTemplateAlign64; + + LocalTensor vec2ResUb = this->stage2OutBuf.template Get(); + LocalTensor mmRes = bmm2ResBuf.template GetTensor(); + + WaitFlag(mte3ToVId[0]); + if (unlikely(runInfo.s2LoopCount == 0)) { + DataCopy(vec2ResUb, mmRes, qsfaVec2CalcSize); + } else { + LocalTensor qsfaExpUb = softmaxExpBuf[runInfo.taskIdMod2].template Get(); + if (runInfo.s2LoopCount < runInfo.s2LoopLimit) { + FlashUpdateNew( + vec2ResUb, mmRes, vec2ResUb, qsfaExpUb, qsfaExpUb, runInfo.vec2MRealSize, + qsfaDTemplateAlign64, 1.0, 1.0); + } else { + LocalTensor sumUb = this->softmaxSumBuf[runInfo.multiCoreIdxMod2].template Get(); + FlashUpdateLastNew( + vec2ResUb, mmRes, vec2ResUb, qsfaExpUb, qsfaExpUb, sumUb, runInfo.vec2MRealSize, + qsfaDTemplateAlign64, 1.0, 1.0); + } + } + + bmm2ResBuf.SetCrossCore(); + if (runInfo.s2LoopCount == runInfo.s2LoopLimit) { + if (unlikely(runInfo.s2LoopCount == 0)) { + LocalTensor sumUb = this->softmaxSumBuf[runInfo.multiCoreIdxMod2].template Get(); + LastDivNew( + vec2ResUb, vec2ResUb, sumUb, runInfo.vec2MRealSize, qsfaDTemplateAlign64, 1.0); + } + + this->CopyOutAttentionOut(runInfo, constInfo, vec2ResUb, 0, qsfaVec2CalcSize); + } + SetFlag(mte3ToVId[0]); +} + +TEMPLATES_DEF_NO_DEFAULT +template +__aicore__ inline void QSFAVectorService::Bmm2DataCopyOut (RunInfo &runInfo, ConstInfo &constInfo, + LocalTensor &vec2ResUb, int64_t vec2S1Idx, int64_t qsfaVec2CalcSize) +{ + LocalTensor attenOut; + int64_t dSizeAligned64 = (int64_t)qsfaDTemplateAlign64; + + attenOut.SetAddr(vec2ResUb.address_); + Cast(attenOut, vec2ResUb, RoundMode::CAST_ROUND, qsfaVec2CalcSize); + SetFlag(vToMte3Id[0]); + WaitFlag(vToMte3Id[0]); + + DataCopyExtParams dataCopyParams; + dataCopyParams.blockLen = constInfo.dSizeV * sizeof(OUTPUT_T); + dataCopyParams.srcStride = (dSizeAligned64 - constInfo.dSizeV) >> 4; // 以32B为单位偏移，bf16类型即偏移16个数，右移4 + dataCopyParams.dstStride = constInfo.attentionOutStride; + dataCopyParams.blockCount = runInfo.vec2MRealSize; + + DataCopyPad(this->attentionOutGm[runInfo.attentionOutOffset], attenOut, dataCopyParams); +} + +TEMPLATES_DEF_NO_DEFAULT +template +__aicore__ inline void QSFAVectorService::CopyOutAttentionOut( + RunInfo &runInfo, ConstInfo &constInfo, LocalTensor &vec2ResUb, + int64_t vec2S1Idx, int64_t qsfaVec2CalcSize) +{ + this->Bmm2DataCopyOut(runInfo, constInfo, vec2ResUb, vec2S1Idx, qsfaVec2CalcSize); +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void QSFAVectorService::InitOutputSingleCore(ConstInfo &constInfo) +{ + uint32_t coreNum = GetBlockNum(); + uint64_t totalOutputSize = 0; + + // n2 = 1, n1 = gn2 = gSize + if constexpr (LAYOUT_T == QSFA_LAYOUT::BSND) { + totalOutputSize = constInfo.bSize * constInfo.gSize * constInfo.s1Size * constInfo.dSizeV; + } else if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + totalOutputSize = constInfo.s1Size * constInfo.gSize * constInfo.dSizeV; + } + + if (coreNum != 0) { + uint64_t singleCoreSize = (totalOutputSize + (CV_RATIO * coreNum) - 1) / (CV_RATIO * coreNum); + uint64_t tailSize = totalOutputSize - constInfo.aivIdx * singleCoreSize; + uint64_t singleInitOutputSize = tailSize < singleCoreSize ? tailSize : singleCoreSize; + if (singleInitOutputSize > 0) { + matmul::InitOutput(this->attentionOutGm[constInfo.aivIdx * singleCoreSize], singleInitOutputSize, 0); + } + } + SyncAll(); +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void QSFAVectorService::CleanOutput(__gm__ uint8_t *attentionOut, ConstInfo &constInfo) +{ + if ASCEND_IS_AIV { + this->attentionOutGm.SetGlobalBuffer((__gm__ OUTPUT_T *)attentionOut); + if (constInfo.needInit == 1) { + InitOutputSingleCore(constInfo); + } + } +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void QSFAVectorService::InitGlobalBuffer(__gm__ uint8_t *key, +__gm__ uint8_t *value, __gm__ uint8_t *sparseIndices, __gm__ uint8_t *blockTable) +{ + keyGm.SetGlobalBuffer((__gm__ KV_T *)(key)); + SparseIndicesGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices); + if constexpr (isPa) { + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + } +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void QSFAVectorService::SoftmaxInitBuffer() +{ + constexpr uint32_t softmaxBufSize = 256; // VF单次操作256Byte + tPipe->InitBuffer(softmaxSumBuf[0], softmaxBufSize); + tPipe->InitBuffer(softmaxSumBuf[1], softmaxBufSize); + tPipe->InitBuffer(softmaxMaxBuf[0], softmaxBufSize); + tPipe->InitBuffer(softmaxMaxBuf[1], softmaxBufSize); + tPipe->InitBuffer(softmaxExpBuf[0], softmaxBufSize); + tPipe->InitBuffer(softmaxExpBuf[1], softmaxBufSize); +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void QSFAVectorService::InitLocalBuffer(TPipe *pipe, ConstInfo &constInfo) +{ + // ub buffer + SoftmaxInitBuffer(); + + tPipe->InitBuffer(commonTBuf, 512); // commonTBuf内存申请512B + tPipe->InitBuffer(stage0InQue, 2, dVTemplateTypeInput * 16 * sizeof(KV_T)); // V0阶段每次处理16个seq, 开2 buffer + // 576: 模型特征维度(dSize) + tPipe->InitBuffer(stage0OutQue, 2, 576 * (16 + 1) * sizeof(Q_T)); // kv输入D轴640, V0阶段每次处理16个seq, 开2 buffer + + tPipe->InitBuffer(stage1OutQue[0], 1, vec1Srcstride * s2BaseSize * sizeof(Q_T)); + tPipe->InitBuffer(stage1OutQue[1], 1, vec1Srcstride * s2BaseSize * sizeof(Q_T)); + tPipe->InitBuffer(stage2OutBuf, (s1BaseSize / CV_RATIO) * qsfaDTemplateAlign64 * sizeof(T)); + + mte3ToVId[0] = GetTPipePtr()->AllocEventID(); + mte3ToVId[1] = GetTPipePtr()->AllocEventID(); + + vToMte3Id[0] = GetTPipePtr()->AllocEventID(); + vToMte3Id[1] = GetTPipePtr()->AllocEventID(); + SetFlag(mte3ToVId[0]); + SetFlag(mte3ToVId[1]); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void QSFAVectorService::InitCubeVecSharedParams( + CVSharedParams &sharedParams, int32_t aicIdx, uint8_t subBlockIdx) +{ + auto &sparseAttnSharedkvBaseParams = this->tilingData->baseParams; + sharedParams.bSize = sparseAttnSharedkvBaseParams.batchSize; + sharedParams.n2Size = 1; + sharedParams.s1Size = sparseAttnSharedkvBaseParams.qSeqSize; + sharedParams.s2Size = sparseAttnSharedkvBaseParams.seqSize; + sharedParams.gSize = sparseAttnSharedkvBaseParams.nNumOfQInOneGroup; + + sharedParams.sparseBlockCount = sparseAttnSharedkvBaseParams.sparseBlockCount; + sharedParams.maskMode = sparseAttnSharedkvBaseParams.sparseMode; + sharedParams.layoutType = sparseAttnSharedkvBaseParams.outputLayout; + sharedParams.dSizeRope = 64; // 64: 编码维度 + sharedParams.softmaxScale = sparseAttnSharedkvBaseParams.scaleValue; + sharedParams.dSize = 576; // 576: 模型特征维度(dSize) + sharedParams.dSizeVInput = sparseAttnSharedkvBaseParams.dSizeVInput; + sharedParams.usedCoreNum = this->tilingData->singleCoreParams.usedCoreNum; + if constexpr (isPa) { + sharedParams.blockSize = sparseAttnSharedkvBaseParams.blockSize; + sharedParams.maxBlockNumPerBatch = sparseAttnSharedkvBaseParams.maxBlockNumPerBatch; + } + + sharedParams.isActualSeqLengthsNull = sparseAttnSharedkvBaseParams.isActualLenDimsNull; + sharedParams.isActualSeqLengthsKVNull = sparseAttnSharedkvBaseParams.isActualLenDimsKVNull; + + ComputeNeedInitQSFA(sharedParams); + + if ASCEND_IS_AIV { + if (subBlockIdx == 0) { + auto qsfaTempTilingSSbuf = reinterpret_cast<__ssbuf__ uint32_t*>(0); // 从ssbuf的0地址开始拷贝 + auto tempTiling = reinterpret_cast(&sharedParams); + + #pragma unroll + for (int i = 0; i < sizeof(CVSharedParams) / sizeof(uint32_t); ++i, ++qsfaTempTilingSSbuf, ++tempTiling) { + *qsfaTempTilingSSbuf = *tempTiling; + } + + CrossCoreSetFlag(15); + } + } +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void QSFAVectorService::ComputeNeedInitQSFA( + CVSharedParams &sharedParams) const +{ + sharedParams.needInit = 0; + for (uint32_t bIdx = 0; bIdx < sharedParams.bSize; bIdx++) { + int64_t s2Size; + if constexpr (KV_LAYOUT_T == QSFA_LAYOUT::TND) { + s2Size = (bIdx == 0) ? actualSeqLengthsKVGm.GetValue(bIdx) : \ + (actualSeqLengthsKVGm.GetValue(bIdx) - actualSeqLengthsKVGm.GetValue(bIdx - 1)); + } else { + if (sharedParams.isActualSeqLengthsKVNull) { + s2Size = sharedParams.s2Size; + } else { + s2Size = actualSeqLengthsKVGm.GetValue(bIdx); + } + } + + int64_t s1Size; + if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + s1Size = (bIdx == 0) ? cuSeqlensQGm.GetValue(bIdx) : \ + (cuSeqlensQGm.GetValue(bIdx) - cuSeqlensQGm.GetValue(bIdx - 1)); + } else { + if (sharedParams.isActualSeqLengthsNull) { + s1Size = sharedParams.s1Size; + } else { + s1Size = cuSeqlensQGm.GetValue(bIdx); + } + } + if (s1Size > s2Size || (LAYOUT_T == QSFA_LAYOUT::BSND && s1Size < sharedParams.s1Size)) { + sharedParams.needInit = 1; + break; + } + } +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void QSFAVectorService::GetExtremeValue( + T &negativeScalar) +{ + uint32_t tmp1 = NEGATIVE_MIN_VALUE_FP32; + negativeScalar = *((float *)&tmp1); +} + +TEMPLATES_DEF class QSFAVectorServiceDummy { +public: + __aicore__ inline QSFAVectorServiceDummy() {}; + __aicore__ inline void CleanOutput(__gm__ uint8_t *attentionOut, ConstInfo &constInfo) {} + __aicore__ inline void InitGlobalBuffer(__gm__ uint8_t *key, __gm__ uint8_t *value, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *blockTable) {} + __aicore__ inline void InitVecBlock(TPipe *pipe, const KvQuantSparseFlashAttentionTilingDataMla *__restrict tiling, + CVSharedParams &sharedParams, int32_t aicIdx, uint8_t subBlockIdx, __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths) {}; + __aicore__ inline void InitLocalBuffer(TPipe *pipe, ConstInfo &constInfo) {} + + __aicore__ inline void ProcessVec1(Buffer &outputBuf, + Buffer &bmm1ResBuf, + RunInfo &runInfo, + ConstInfo &constInfo) {} + using mm2ResPos = Buffer; + __aicore__ inline void ProcessVec2(mm2ResPos &bmm2ResBuf, RunInfo &runInfo, + ConstInfo &constInfo) {} +}; +} +#endif // KV_QUANT_SPARSE_FLASH_ATTENTION_SERVICE_VECTOR_MLA_H diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention.cpp b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention.cpp new file mode 100644 index 000000000..57a00f5c5 --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention.cpp @@ -0,0 +1,147 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + + /*! + * \file kv_quant_sparse_flash_attention.cpp + * \brief + */ + +#include "kernel_operator.h" +#include "kv_quant_sparse_flash_attention_template_tiling_key.h" +#if (__CCE_AICORE__ == 310) +#include "arch35/kv_quant_sparse_flash_attention_kernel_mla.h" +#else +#include "kv_quant_sparse_flash_attention_kernel_mla.h" +#endif + +using namespace AscendC; + +#if (__CCE_AICORE__ == 310) +#if defined(__DAV_C310_CUBE__) +#define QSFA_OP_IMPL(templateClass, tilingdataClass, ...) \ + do { \ + using CubeBlockType = typename std::conditional, BaseApi::QSFAMatmulServiceDummy<__VA_ARGS__>>::type; \ + using VecBlockType = typename std::conditional, BaseApi::QSFAVectorService<__VA_ARGS__>>::type; \ + templateClass op; \ + op.Init(query, key, value, sparseIndices, keyScale, valueScale, blocktable, \ + actualSeqLengthsQuery, actualSeqLengthsKV, \ + attentionOut, user, nullptr, &tPipe); \ + op.Process(); \ + } while (0) +#else +#define QSFA_OP_IMPL(templateClass, tilingdataClass, ...) \ + do { \ + using CubeBlockType = typename std::conditional, BaseApi::QSFAMatmulServiceDummy<__VA_ARGS__>>::type; \ + using VecBlockType = typename std::conditional, BaseApi::QSFAVectorService<__VA_ARGS__>>::type; \ + templateClass op; \ + GET_TILING_DATA_WITH_STRUCT(tilingdataClass, tilingDataIn, tiling); \ + const tilingdataClass *__restrict tilingData = &tilingDataIn; \ + op.Init(query, key, value, sparseIndices, keyScale, valueScale, blocktable, \ + actualSeqLengthsQuery, actualSeqLengthsKV, \ + attentionOut, user, tilingData, &tPipe); \ + op.Process(); \ + } while (0) +#endif +#else +#define QSFA_OP_IMPL(templateClass, tilingdataClass, ...) \ + do { \ + templateClass> op; \ + GET_TILING_DATA_WITH_STRUCT(tilingdataClass, tiling_data_in, tiling); \ + const tilingdataClass *__restrict tiling_data = &tiling_data_in; \ + op.Init(query, key, value, sparseIndices, keyScale, valueScale, blocktable, \ + actualSeqLengthsQuery, actualSeqLengthsKV, \ + attentionOut, softmaxMax, softmaxSum, user, tiling_data, tiling, &tPipe); \ + op.Process(); \ + } while (0) +#endif + +#if (__CCE_AICORE__ == 310) +template +__aicore__ inline void DispatchKernelDtype310( + __gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *keyScale, __gm__ uint8_t *valueScale, + __gm__ uint8_t *blocktable, __gm__ uint8_t *actualSeqLengthsQuery, + __gm__ uint8_t *actualSeqLengthsKV, __gm__ uint8_t *attentionOut, + __gm__ uint8_t *user, __gm__ uint8_t *tiling, TPipe &tPipe) +{ + if constexpr (ORIG_DTYPE_QUERY == DT_BF16 && ORIG_DTYPE_KEY == DT_FLOAT8_E4M3FN && + ORIG_DTYPE_ATTENTION_OUT == DT_BF16) { + QSFA_OP_IMPL(BaseApi::KvQuantSparseFlashAttentionMla, KvQuantSparseFlashAttentionTilingDataMla, + bfloat16_t, fp8_e4m3fn_t, float, bfloat16_t, FLASH_DECODE, PAGE_ATTENTION, + static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), + static_cast(TEMPLATE_MODE), IS_SPLIT_G); + } else if constexpr (ORIG_DTYPE_QUERY == DT_BF16 && ORIG_DTYPE_KEY == DT_HIFLOAT8 && + ORIG_DTYPE_ATTENTION_OUT == DT_BF16) { + QSFA_OP_IMPL(BaseApi::KvQuantSparseFlashAttentionMla, KvQuantSparseFlashAttentionTilingDataMla, + bfloat16_t, hifloat8_t, float, bfloat16_t, FLASH_DECODE, PAGE_ATTENTION, + static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), + static_cast(TEMPLATE_MODE), IS_SPLIT_G); + } else if constexpr (ORIG_DTYPE_QUERY == DT_BF16 && ORIG_DTYPE_KEY == DT_INT8 && + ORIG_DTYPE_ATTENTION_OUT == DT_BF16) { + QSFA_OP_IMPL(BaseApi::KvQuantSparseFlashAttentionMla, KvQuantSparseFlashAttentionTilingDataMla, + bfloat16_t, int8_t, float, bfloat16_t, FLASH_DECODE, PAGE_ATTENTION, + static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), + static_cast(TEMPLATE_MODE), IS_SPLIT_G); + } else if constexpr (ORIG_DTYPE_QUERY == DT_FLOAT16 && ORIG_DTYPE_KEY == DT_FLOAT8_E4M3FN && + ORIG_DTYPE_ATTENTION_OUT == DT_FLOAT16) { + QSFA_OP_IMPL(BaseApi::KvQuantSparseFlashAttentionMla, KvQuantSparseFlashAttentionTilingDataMla, + half, fp8_e4m3fn_t, float, half, FLASH_DECODE, PAGE_ATTENTION, + static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), + static_cast(TEMPLATE_MODE), IS_SPLIT_G); + } else if constexpr (ORIG_DTYPE_QUERY == DT_FLOAT16 && ORIG_DTYPE_KEY == DT_HIFLOAT8 && + ORIG_DTYPE_ATTENTION_OUT == DT_FLOAT16) { + QSFA_OP_IMPL(BaseApi::KvQuantSparseFlashAttentionMla, KvQuantSparseFlashAttentionTilingDataMla, + half, hifloat8_t, float, half, FLASH_DECODE, PAGE_ATTENTION, + static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), + static_cast(TEMPLATE_MODE), IS_SPLIT_G); + } else if constexpr (ORIG_DTYPE_QUERY == DT_FLOAT16 && ORIG_DTYPE_KEY == DT_INT8 && + ORIG_DTYPE_ATTENTION_OUT == DT_FLOAT16) { + QSFA_OP_IMPL(BaseApi::KvQuantSparseFlashAttentionMla, KvQuantSparseFlashAttentionTilingDataMla, + half, int8_t, float, half, FLASH_DECODE, PAGE_ATTENTION, + static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), + static_cast(TEMPLATE_MODE), IS_SPLIT_G); + } +} +#endif + +template + __global__ __aicore__ void +kv_quant_sparse_flash_attention(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t* keyScale, __gm__ uint8_t* valueScale, + __gm__ uint8_t *blocktable, __gm__ uint8_t *actualSeqLengthsQuery, + __gm__ uint8_t *actualSeqLengthsKV, __gm__ uint8_t *attentionOut, + __gm__ uint8_t *softmaxMax, __gm__ uint8_t *softmaxSum, + __gm__ uint8_t *workspace, __gm__ uint8_t *tiling) +{ + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + + TPipe tPipe; + __gm__ uint8_t *user = GetUserWorkspace(workspace); +#if (__CCE_AICORE__ == 310) + DispatchKernelDtype310( + query, key, value, sparseIndices, keyScale, valueScale, blocktable, + actualSeqLengthsQuery, actualSeqLengthsKV, attentionOut, user, tiling, tPipe); +#else + if constexpr (ORIG_DTYPE_QUERY == DT_FLOAT16 && ORIG_DTYPE_KEY == DT_INT8 && + ORIG_DTYPE_ATTENTION_OUT == DT_FLOAT16) { + QSFA_OP_IMPL(KvQuantSparseFlashAttentionMla, KvQuantSparseFlashAttentionTilingDataMla, half, int8_t, + half, FLASH_DECODE, static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), + TEMPLATE_MODE); + } else { // bf16 + QSFA_OP_IMPL(KvQuantSparseFlashAttentionMla, KvQuantSparseFlashAttentionTilingDataMla, bfloat16_t, int8_t, + bfloat16_t, FLASH_DECODE, static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), + TEMPLATE_MODE); + } +#endif +} \ No newline at end of file diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_common.h b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_common.h new file mode 100644 index 000000000..7ea38dbc5 --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_common.h @@ -0,0 +1,225 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_common.h + * \brief + */ + +#ifndef KV_QUANT_SPARSE_FLASH_ATTENTION_COMMON_H +#define KV_QUANT_SPARSE_FLASH_ATTENTION_COMMON_H + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" + +using namespace AscendC; +// 将isCheckTiling设置为false, 输入输出的max&sum&exp的shape为(m, 1) +constexpr SoftmaxConfig QSFA_SOFTMAX_FLASHV2_CFG_WITHOUT_BRC = {false, 0, 0, SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC}; + +enum class QSFA_LAYOUT { + BSND = 0, + TND = 1, + PA_BSND = 2, +}; + +enum class QUANT_MODE { + PER_CHANNEL = 0, // GQA支持 + PER_TOKEN_HEAD = 1, // GQA支持 + PER_TILE = 2, // MLA支持 +}; + +enum class ATTENTION_MODE { + GQA_MHA = 0, // QKV headDim相等 + MLA_NATIVE = 1, // Dn=128, Dr=64 + MLA_ABSORB = 2, // Dn=512, Dr=64 +}; + +enum class QUANT_SCALE_REPO_MODE { + SEPARATE = 0, // 分开存储 + COMBINE = 1, // 合并存储，量化模式是PER_TOKEN_HEAD/PER_TILE时支持COMBINE模式，参数顺序为：Nope+Rope+DequantScale +}; + +template +struct QSFAType { + using queryType = Q_T; + using kvType = KV_T; + using kRopeType = Q_T; + using outputType = OUT_T; + static constexpr bool flashDecode = FLASH_DECODE; + static constexpr QSFA_LAYOUT layout = LAYOUT_T; + static constexpr QSFA_LAYOUT kvLayout = KV_LAYOUT_T; + static constexpr int templateMode = TEMPLATE_MODE; + static constexpr bool pageAttention = (KV_LAYOUT_T == QSFA_LAYOUT::PA_BSND); +}; + +// ================================Util functions================================== +template __aicore__ inline T QSFAAlign(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd) * (rnd))); +} + +template __aicore__ inline size_t BlockAlign(size_t s) +{ + if constexpr (IsSameType::value) { + return (s + 63) / 64 * 64; + } + size_t n = (32 / sizeof(T)); + return (s + n - 1) / n * n; +} + +template __aicore__ inline T1 Min(T1 a, T2 b) +{ + return (a > b) ? (b) : (a); +} + +struct RunInfo { + uint32_t loop; + uint32_t bIdx; + uint32_t gIdx; + uint32_t s1Idx; + uint32_t s2Idx; + uint32_t bn2IdxInCurCore; + uint32_t curSInnerLoopTimes; + uint32_t s2BatchOffset; + + uint64_t tndBIdxOffsetForQ; + uint64_t tndBIdxOffsetForKV; + uint64_t tensorAOffset; + uint64_t tensorBOffset; + uint64_t tensorARopeOffset; + uint64_t tensorBRopeOffset; + uint64_t attenOutOffset; + uint64_t topKBaseOffset; + uint64_t attenMaskOffset; + + uint32_t actualSingleProcessSInnerSize; + uint32_t actualSingleProcessSInnerSizeAlign; + uint32_t gSize; + uint32_t s1Size; + uint32_t s2Size; + uint32_t mSize; + uint32_t mSizeV; + uint32_t mSizeVStart; + uint32_t tndIsS2SplitCore; + uint32_t tndCoreStartKVSplitPos; + bool isBmm2Output; + bool isValid = false; + bool isFirstSInnerLoop; + bool isChangeBatch; + static constexpr uint32_t n2Idx = 0; + + uint64_t actS1Size = 1; + uint64_t curActualSeqLenOri = 0ULL; + uint64_t actS2Size = 1; + + uint32_t gS1Idx; + uint32_t actMBaseSize; + int32_t nextTokensPerBatch = 0; + bool isLastS2Loop; + uint8_t resv[3]; + int64_t threshold; +}; + +struct ConstInfo { + // CUBE与VEC核间同步的模式 + static constexpr uint32_t QSFA_SYNC_MODE2 = 2; + // BUFFER的字节数 + static constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024; + static constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048; + static constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096; + static constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192; + static constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384; + static constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768; + static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; + static constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64; + static constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256; + static constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512; + // FP32的0值和极大值 + static constexpr float FLOAT_ZERO = 0; + static constexpr float FLOAT_MAX = 3.402823466e+38F; + + // preLoad的总次数 + uint32_t preLoadNum = 0U; + uint32_t nBufferMBaseSize = 0U; + // CUBE和VEC的核间同步EventID + uint32_t syncV0C1 = 0U; + uint32_t syncC1V1 = 0U; + uint32_t syncV1C2 = 0U; + uint32_t syncC2V2 = 0U; + uint32_t syncC2V1 = 0U; + uint32_t syncV1NupdateC2 = 0U; + + uint32_t mmResUbSize = 0U; // Matmul1输出结果GM上的大小 + uint32_t vec1ResUbSize = 0U; // Vector1输出结果GM上的大小 + uint32_t bmm2ResUbSize = 0U; // Matmul2输出结果GM上的大小 + uint64_t gSize = 0ULL; + uint64_t batchSize = 0ULL; + uint64_t qHeadNum = 0ULL; + uint64_t kvHeadNum; + uint64_t headDim; + uint64_t headDimRope; + uint64_t combineHeadDim; // quantScaleRepoMode为Combine模式时=headDim+headDimRope, 否则=headDim + uint64_t kvSeqSize = 0ULL; // kv最大S长度 + uint64_t qSeqSize = 1ULL; // q最大S长度 + int64_t kvCacheBlockSize = 0; // PA场景的block size + uint32_t maxBlockNumPerBatch = 0; // PA场景的最大单batch block number + uint32_t splitKVNum = 0U; // S2核间切分的切分份数 + QSFA_LAYOUT outputLayout; // 输出的Transpose格式 + uint32_t sparseMode = 0; + bool returnSoftmaxLse = false; + bool needInit = false; + + // FlashDecoding + uint64_t combineLseOffset = 0ULL; + uint64_t combineAccumOutOffset = 0ULL; + uint32_t actualCombineLoopSize = 0U; // FlashDecoding场景, S2在核间切分的最大份数 + + uint32_t actualLenDimsQ = 0U; // query的actualSeqLength 的维度 + uint32_t actualLenDimsKV = 0U; // KV 的actualSeqLength 的维度 + + // TND + uint32_t s2Start = 0U; // TND场景下，S2的起始位置 + uint32_t s2End = 0U; // 单核TND场景下S2循环index上限 + + uint32_t bN2Start = 0U; + uint32_t bN2End = 0U; + uint32_t gS1Start = 0U; + uint32_t gS1End = 0U; + + uint32_t mBaseSize = 1ULL; + uint32_t s2BaseSize = 1ULL; + + uint32_t tndFDCoreArrLen = 0U; // TNDFlashDecoding相关分核信息array的长度 + uint32_t coreStartKVSplitPos = 0U; // TNDFlashDecoding kv起始位置 + + // sparse attr + uint32_t sparseBlockCount = 0; + int64_t sparseBlockSize = 0; + + // attention模式与量化模式 + ATTENTION_MODE attentionMode = ATTENTION_MODE::MLA_ABSORB; + QUANT_MODE keyQuantMode = QUANT_MODE::PER_TILE; + QUANT_MODE valueQuantMode = QUANT_MODE::PER_TILE; + QUANT_SCALE_REPO_MODE quantScaleRepoMode = QUANT_SCALE_REPO_MODE::COMBINE; + uint64_t tileSize = 128ULL; +}; + +struct MSplitInfo { + uint32_t nBufferIdx = 0U; + uint32_t nBufferStartM = 0U; + uint32_t nBufferDealM = 0U; + uint32_t vecStartM = 0U; + uint32_t vecDealM = 0U; +}; + +#endif // KV_QUANT_SPARSE_FLASH_ATTENTION_COMMON_H \ No newline at end of file diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_kernel_mla.h b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_kernel_mla.h new file mode 100644 index 000000000..174b432ba --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_kernel_mla.h @@ -0,0 +1,1003 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_kernel_mla.h + * \brief + */ + +#ifndef KV_QUANT_SPARSE_FLASH_ATTENTION_KERNEL_MLA_H +#define KV_QUANT_SPARSE_FLASH_ATTENTION_KERNEL_MLA_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "kv_quant_sparse_flash_attention_common.h" +#include "kv_quant_sparse_flash_attention_service_cube_mla.h" +#include "kv_quant_sparse_flash_attention_service_vector_mla.h" + +using namespace matmul; +using AscendC::CacheMode; +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +// 由于S2循环前，RunInfo还没有赋值，使用Bngs1Param临时存放B、N、S1轴相关的信息；同时减少重复计算 +struct TempLoopInfo { + uint32_t bn2IdxInCurCore = 0; + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint32_t s2LoopTimes = 0U; // S2方向循环的总次数，无论TND还是BXXD都是等于实际次数，不用减1 + uint64_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小 + uint64_t curActualSeqLen = 0ULL; + uint64_t curActualSeqLenOri = 0ULL; + uint64_t actS1Size = 1ULL; // TND场景下当前Batch循环处理的S1轴的大小，非TND场景下不要用这个字段 + uint64_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小 + + int32_t nextTokensPerBatch = 0; + uint32_t tndCoreStartKVSplitPos; + uint32_t gS1Idx = 0U; + bool tndIsS2SplitCore; + bool curActSeqLenIsZero = false; +}; + +template class KvQuantSparseFlashAttentionMla { +public: + // 中间计算数据类型为float，高精度模式 + using T = float; + using Q_T = typename QSFAT::queryType; + using KV_T = typename QSFAT::kvType; + using OUT_T = typename QSFAT::outputType; + using Q_ROPE_T = Q_T; + using K_ROPE_T = typename QSFAT::kRopeType; + using UPDATE_T = T; + using MM1_OUT_T = T; + using MM2_OUT_T = T; + + __aicore__ inline KvQuantSparseFlashAttentionMla(){}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t* keyScale, + __gm__ uint8_t* valueScale, __gm__ uint8_t *blockTable, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths, + __gm__ uint8_t *attentionOut, __gm__ uint8_t *softmaxMax, + __gm__ uint8_t *softmaxSum, __gm__ uint8_t *workspace, + const KvQuantSparseFlashAttentionTilingDataMla *__restrict tiling, + __gm__ uint8_t *gmTiling, TPipe *tPipe); + + __aicore__ inline void Process(); + +private: + static constexpr bool PAGE_ATTENTION = QSFAT::pageAttention; + static constexpr int TEMPLATE_MODE = QSFAT::templateMode; + static constexpr bool FLASH_DECODE = QSFAT::flashDecode; + static constexpr QSFA_LAYOUT LAYOUT_T = QSFAT::layout; + static constexpr QSFA_LAYOUT KV_LAYOUT_T = QSFAT::kvLayout; + + static constexpr uint32_t PRELOAD_NUM = 2; + static constexpr uint32_t N_BUFFER_M_BASIC_SIZE = 256; + static constexpr uint32_t QSFA_PRELOAD_TASK_CACHE_SIZE = 3; + + static constexpr uint32_t SYNC_V0_C1_FLAG = 6; + static constexpr uint32_t SYNC_C1_V1_FLAG = 7; + static constexpr uint32_t SYNC_V1_C2_FLAG = 8; + static constexpr uint32_t SYNC_C2_V2_FLAG = 9; + static constexpr uint32_t SYNC_C2_V1_FLAG = 4; + static constexpr uint32_t SYNC_V1_NUPDATE_C2_FLAG = 5; + + static constexpr uint64_t SYNC_MM2RES_BUF1_FLAG = 10; + static constexpr uint64_t SYNC_MM2RES_BUF2_FLAG = 11; + static constexpr uint64_t SYNC_FDOUTPUT_BUF_FLAG = 12; + + static constexpr uint32_t BLOCK_ELEMENT_NUM = QSFAVectorService::BYTE_BLOCK / sizeof(T); + + static constexpr uint64_t kvHeadNum = 1ULL; + static constexpr uint64_t headDim = 512ULL; + static constexpr uint64_t headDimAlign = 512ULL; + static constexpr uint64_t headDimRope = 64ULL; + static constexpr uint32_t msdIterNum = 2U; + + static constexpr uint32_t dbWorkspaceRatio = PRELOAD_NUM; + + const KvQuantSparseFlashAttentionTilingDataMla *__restrict tilingData = nullptr; + + TPipe *pipe = nullptr; + + uint64_t mSizeVStart = 0ULL; + int64_t threshold = 0; + uint64_t topKBaseOffset = 0ULL; + uint64_t s2BatchBaseOffset = 0; + uint64_t tensorACoreOffset = 0ULL; + uint64_t tensorARopeCoreOffset = 0ULL; + uint64_t tensorBCoreOffset = 0ULL; + uint64_t tensorBRopeCoreOffset = 0ULL; + uint64_t attenOutOffset = 0ULL; + + uint32_t tmpBlockIdx = 0U; + uint32_t aiCoreIdx = 0U; + uint32_t usedCoreNum = 0U; + + __gm__ uint8_t *keyPtr = nullptr; + __gm__ uint8_t *valuePtr = nullptr; + + ConstInfo constInfo{}; + TempLoopInfo tempLoopInfo{}; + + QSFAMatmulService matmulService; + QSFAVectorService vectorService; + + GlobalTensor queryGm; + GlobalTensor keyGm; + GlobalTensor valueGm; + GlobalTensor qRopeGm; + GlobalTensor kRopeGm; + + GlobalTensor attentionOutGm; + GlobalTensor softmaxMaxGm; + GlobalTensor softmaxSumGm; + GlobalTensor blockTableGm; + GlobalTensor topKGm; + + GlobalTensor actualSeqLengthsQGm; + GlobalTensor actualSeqLengthsKVGm; + + // workspace + GlobalTensor mm1ResGm; + GlobalTensor vec1ResGm; + GlobalTensor mm2ResGm; + GlobalTensor kvMergeGm_; + GlobalTensor kvValidSizeGm_; + + GlobalTensor mm2ResInt32Gm; + GlobalTensor vec2ResGm; + + GlobalTensor accumOutGm; + GlobalTensor lseSumFdGm; + GlobalTensor lseMaxFdGm; + + // ================================Init functions================================== + __aicore__ inline void InitTilingData(); + __aicore__ inline void InitCalcParamsEach(); + __aicore__ inline void InitBuffers(); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths); + __aicore__ inline void InitOutputSingleCore(); + // ================================Process functions================================ + __aicore__ inline void ProcessBalance(); + __aicore__ inline void PreloadPipeline(uint32_t loop, uint64_t s2Start, uint64_t s2LoopIdx, + RunInfo extraInfo[QSFA_PRELOAD_TASK_CACHE_SIZE]); + // ================================Offset Calc===================================== + __aicore__ inline void GetActualSeqLen(uint32_t bIdx, uint32_t s1Idx = 0); + __aicore__ inline void GetSparseActualSeqLen(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx); + __aicore__ inline void UpdateInnerLoopCond(); + __aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx); + __aicore__ inline void CalcParams(uint32_t loop, uint64_t s2Start, uint32_t s2LoopIdx, RunInfo &info); + __aicore__ inline void CalcMSizeInfo(RunInfo &info); + __aicore__ inline void CalcFirstTensorOffsets(RunInfo &info, uint64_t qsfaActualSeqQPrefixSum, + uint64_t actualSeqKVPrefixSum); + __aicore__ inline void GetAxisStartIdx(uint32_t bN2EndPrev, uint32_t gS1EndPrev, uint32_t s2EndPrev); + __aicore__ inline uint64_t GetBalanceActualSeqLengths(GlobalTensor &actualSeqLengths, uint32_t bIdx); + __aicore__ inline uint32_t GetActualSeqLenKV(uint32_t bIdx); + __aicore__ inline void GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx); + __aicore__ inline void GetPreNextTokensLeftUp(); + __aicore__ inline void UpdateInner(uint32_t &s2End, uint32_t &curS2End, uint32_t s1Idx, bool isEnd); + // ================================Mm1============================================== + __aicore__ inline void ComputeMm1(const RunInfo &info); + // ================================Mm2============================================== + __aicore__ inline void InitAllZeroOutput(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx); + __aicore__ inline void ComputeMm2(const RunInfo &info); + __aicore__ inline void Bmm2DataCopyOut(uint64_t attenOutOffset, LocalTensor &attenOutUb, uint32_t startRow, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); +}; + +template __aicore__ inline void KvQuantSparseFlashAttentionMla::InitTilingData() +{ + usedCoreNum = tilingData->singleCoreParams.usedCoreNum; + constInfo.splitKVNum = tilingData->splitKVParams.s2; + constInfo.mmResUbSize = tilingData->singleCoreTensorSize.mmResUbSize; + constInfo.bmm2ResUbSize = tilingData->singleCoreTensorSize.bmm2ResUbSize; + constInfo.vec1ResUbSize = constInfo.mmResUbSize * msdIterNum; + + constInfo.qHeadNum = constInfo.gSize = tilingData->baseParams.nNumOfQInOneGroup; + constInfo.batchSize = tilingData->baseParams.batchSize; + constInfo.kvSeqSize = tilingData->baseParams.seqSize; + constInfo.qSeqSize = tilingData->baseParams.qSeqSize; + constInfo.maxBlockNumPerBatch = tilingData->baseParams.maxBlockNumPerBatch; + constInfo.kvCacheBlockSize = tilingData->baseParams.blockSize; + constInfo.outputLayout = static_cast(tilingData->baseParams.outputLayout); + constInfo.mBaseSize = tilingData->innerSplitParams.mBaseSize; + constInfo.s2BaseSize = tilingData->innerSplitParams.s2BaseSize; + constInfo.kvHeadNum = kvHeadNum; + constInfo.headDim = headDim; + constInfo.headDimRope = headDimRope; + constInfo.sparseBlockSize = tilingData->baseParams.sparseBlockSize; + constInfo.sparseBlockCount = tilingData->baseParams.sparseBlockCount; + constInfo.sparseMode = tilingData->baseParams.sparseMode; + constInfo.returnSoftmaxLse = tilingData->baseParams.returnSoftmaxLse; + constInfo.quantScaleRepoMode = QUANT_SCALE_REPO_MODE::COMBINE; + constInfo.attentionMode = ATTENTION_MODE::MLA_ABSORB; + constInfo.combineHeadDim = (constInfo.quantScaleRepoMode == + QUANT_SCALE_REPO_MODE::COMBINE) ? headDim + headDimRope : headDim; + + constInfo.preLoadNum = PRELOAD_NUM; + constInfo.nBufferMBaseSize = N_BUFFER_M_BASIC_SIZE; + constInfo.syncV0C1 = SYNC_V0_C1_FLAG; + constInfo.syncC1V1 = SYNC_C1_V1_FLAG; + constInfo.syncV1C2 = SYNC_V1_C2_FLAG; + constInfo.syncC2V2 = SYNC_C2_V2_FLAG; + // constInfo.syncC2V1 = SYNC_C2_V1_FLAG; + constInfo.syncV1NupdateC2 = SYNC_V1_NUPDATE_C2_FLAG; +} + +template __aicore__ inline void KvQuantSparseFlashAttentionMla::InitBuffers() +{ + if ASCEND_IS_AIV { + vectorService.InitBuffers(pipe); + } else { + matmulService.InitBuffers(pipe); + } +} + +template +__aicore__ inline void +KvQuantSparseFlashAttentionMla::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengths) +{ + constInfo.actualLenDimsQ = tilingData->baseParams.actualLenDimsQ; + constInfo.actualLenDimsKV = tilingData->baseParams.actualLenDimsKV; + if (constInfo.actualLenDimsQ != 0) { + actualSeqLengthsQGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsQ, constInfo.actualLenDimsQ); + } + if (constInfo.actualLenDimsKV != 0) { + actualSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengths, constInfo.actualLenDimsKV); + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::InitAllZeroOutput(uint32_t bIdx, uint32_t s1Idx, + uint32_t n2Idx) +{ + if (constInfo.outputLayout == QSFA_LAYOUT::TND) { + uint32_t tBase = bIdx == 0 ? 0 : actualSeqLengthsQGm.GetValue(bIdx - 1); + uint32_t s1Count = tempLoopInfo.actS1Size; + + uint64_t attenOutOffset = (tBase + s1Idx) * kvHeadNum * constInfo.gSize * headDim + // T轴、s1轴偏移 + n2Idx * constInfo.gSize * headDim; // N2轴偏移 + matmul::InitOutput(attentionOutGm[attenOutOffset], constInfo.gSize * headDim, 0); + if (constInfo.returnSoftmaxLse) { //N2T1G + uint64_t softmaxSumOffset = n2Idx * actualSeqLengthsQGm.GetValue(constInfo.batchSize - 1) * constInfo.gSize + + (tBase + s1Idx) * constInfo.gSize; + uint64_t softmaxMaxOffset = softmaxSumOffset; + matmul::InitOutput(softmaxSumGm[softmaxSumOffset], constInfo.gSize, 0); + matmul::InitOutput(softmaxMaxGm[softmaxMaxOffset], constInfo.gSize, 0); + } + } else if (constInfo.outputLayout == QSFA_LAYOUT::BSND) { + uint64_t attenOutOffset = bIdx * constInfo.qSeqSize * kvHeadNum * constInfo.gSize * headDim + + s1Idx * kvHeadNum * constInfo.gSize * headDim + // B轴、S1轴偏移 + n2Idx * constInfo.gSize * headDim; // N2轴偏移 + matmul::InitOutput(attentionOutGm[attenOutOffset], constInfo.gSize * headDim, 0); + if (constInfo.returnSoftmaxLse) { //B, N2, S1, G + uint64_t softmaxSumOffset = bIdx * kvHeadNum * constInfo.qSeqSize * constInfo.gSize + + n2Idx * constInfo.qSeqSize * constInfo.gSize + + s1Idx * constInfo.gSize; + uint64_t softmaxMaxOffset = softmaxSumOffset; + matmul::InitOutput(softmaxSumGm[softmaxSumOffset], constInfo.gSize, 0); + matmul::InitOutput(softmaxMaxGm[softmaxMaxOffset], constInfo.gSize, 0); + } + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::InitOutputSingleCore() +{ + uint32_t qsfaCoreNum = GetBlockNum(); + if (qsfaCoreNum != 0) { + uint64_t qsfaTotalOutputSize = constInfo.batchSize * constInfo.qHeadNum * constInfo.qSeqSize * + constInfo.headDim; + // 2 means c:v = 1:2 + uint64_t qsfaSingleCoreSize = (qsfaTotalOutputSize + (2 * qsfaCoreNum) - 1) / (2 * qsfaCoreNum); + uint64_t qsfaTailSize = qsfaTotalOutputSize - tmpBlockIdx * qsfaSingleCoreSize; + uint64_t qsfaSingleInitOutputSize = qsfaTailSize < qsfaSingleCoreSize ? qsfaTailSize : qsfaSingleCoreSize; + if (tmpBlockIdx * qsfaSingleCoreSize < qsfaTotalOutputSize && qsfaSingleInitOutputSize > 0) { + matmul::InitOutput(attentionOutGm[tmpBlockIdx * qsfaSingleCoreSize], qsfaSingleInitOutputSize, 0); + } + if (constInfo.returnSoftmaxLse) { + uint64_t totalReturnSoftmaxSize = constInfo.batchSize * constInfo.kvHeadNum * constInfo.qSeqSize * constInfo.gSize; + uint64_t singleCoreReturnSoftmaxSize = (totalReturnSoftmaxSize + (2 * qsfaCoreNum) - 1) / (2 * qsfaCoreNum); + uint64_t tailReturnSoftmaxSize = totalReturnSoftmaxSize - tmpBlockIdx * singleCoreReturnSoftmaxSize; + uint64_t singleInitReturnSoftmaxSize = tailReturnSoftmaxSize < singleCoreReturnSoftmaxSize ? + tailReturnSoftmaxSize : singleCoreReturnSoftmaxSize; + if (tmpBlockIdx * singleCoreReturnSoftmaxSize < totalReturnSoftmaxSize && singleInitReturnSoftmaxSize > 0) { + matmul::InitOutput(softmaxSumGm[tmpBlockIdx * singleCoreReturnSoftmaxSize], singleInitReturnSoftmaxSize, 0); + matmul::InitOutput(softmaxMaxGm[tmpBlockIdx * singleCoreReturnSoftmaxSize], singleInitReturnSoftmaxSize, 0); + } + } + SyncAll(); + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::GetActualSeqLen(uint32_t bIdx, uint32_t s1Idx) +{ + tempLoopInfo.curActualSeqLenOri = GetActualSeqLenKV(bIdx); + tempLoopInfo.actS1Size = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bIdx); +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::GetSparseActualSeqLen(uint32_t bIdx, uint32_t s1Idx, + uint32_t n2Idx) +{ + if (tempLoopInfo.nextTokensPerBatch < 0 && s1Idx < (-tempLoopInfo.nextTokensPerBatch)) { // 存在行无效 + tempLoopInfo.curActualSeqLen = 0; + return; + } + int64_t threshold = tempLoopInfo.curActualSeqLenOri; + if (constInfo.sparseMode == 3) { + threshold = static_cast(tempLoopInfo.nextTokensPerBatch) + s1Idx + 1; + } + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + tempLoopInfo.curActualSeqLen = (constInfo.sparseBlockCount * constInfo.sparseBlockSize > threshold) ? + threshold : + constInfo.sparseBlockCount * constInfo.sparseBlockSize; + } else { + uint64_t topKBaseOffset = 0; + if constexpr (LAYOUT_T == QSFA_LAYOUT::BSND) { // B,S1,N2 K + topKBaseOffset = bIdx * constInfo.qSeqSize * kvHeadNum * constInfo.sparseBlockCount + + s1Idx * kvHeadNum * constInfo.sparseBlockCount + n2Idx * constInfo.sparseBlockCount; + } else if (LAYOUT_T == QSFA_LAYOUT::TND) { // T N2 K + uint64_t actualSeqQPrefixSum = (bIdx <= 0) ? 0 : actualSeqLengthsQGm.GetValue(bIdx - 1); + topKBaseOffset = actualSeqQPrefixSum * kvHeadNum * constInfo.sparseBlockCount + + s1Idx * kvHeadNum * constInfo.sparseBlockCount + + n2Idx * constInfo.sparseBlockCount; + } else { // B N2 S1 K + topKBaseOffset = bIdx * kvHeadNum * constInfo.qSeqSize * constInfo.sparseBlockCount + + n2Idx * constInfo.qSeqSize * constInfo.sparseBlockCount + + s1Idx * constInfo.sparseBlockCount; + } + + uint64_t sparseLen = 0; + + for (uint64_t topkIdx = 0; topkIdx < constInfo.sparseBlockCount; topkIdx++) { + int32_t sparseIndices = topKGm.GetValue(topKBaseOffset + topkIdx); + uint64_t blockBegin = sparseIndices * constInfo.sparseBlockSize; + if (blockBegin >= threshold) { + continue; + } + uint64_t blockEnd = (blockBegin + constInfo.sparseBlockSize > tempLoopInfo.curActualSeqLenOri) ? + tempLoopInfo.curActualSeqLenOri : + blockBegin + constInfo.sparseBlockSize; + uint64_t blockLen = (blockEnd <= threshold) ? blockEnd - blockBegin : threshold - blockBegin; + sparseLen += blockLen; + } + tempLoopInfo.curActualSeqLen = sparseLen; + } +} + +template +__aicore__ inline uint32_t KvQuantSparseFlashAttentionMla::GetActualSeqLenKV(uint32_t bIdx) +{ + if constexpr (KV_LAYOUT_T == QSFA_LAYOUT::TND) { + if (bIdx > 0) { + int32_t curActualSeqLen = actualSeqLengthsKVGm.GetValue(bIdx); + int32_t prevActualSeqLen = actualSeqLengthsKVGm.GetValue(bIdx - 1); + return (curActualSeqLen >= prevActualSeqLen) ? + static_cast(curActualSeqLen - prevActualSeqLen) : 0U; + } else if (bIdx == 0) { + return actualSeqLengthsKVGm.GetValue(0); + } else { + return 0; + } + } else { + if (constInfo.actualLenDimsKV == 0) { + return constInfo.kvSeqSize; + } else if (constInfo.actualLenDimsKV == 1) { + return actualSeqLengthsKVGm.GetValue(0); + } else { + return actualSeqLengthsKVGm.GetValue(bIdx); + } + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::DealActSeqLenIsZero(uint32_t bIdx, uint32_t s1Idx, + uint32_t n2Idx) +{ + if ASCEND_IS_AIV { + InitAllZeroOutput(bIdx, s1Idx, n2Idx); + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::GetPreNextTokensLeftUp() +{ + if (constInfo.sparseMode == 3) { + tempLoopInfo.nextTokensPerBatch = + static_cast(tempLoopInfo.curActualSeqLenOri) - static_cast(tempLoopInfo.actS1Size); + } +} + +template __aicore__ inline void KvQuantSparseFlashAttentionMla::UpdateInnerLoopCond() +{ + if ((tempLoopInfo.curActualSeqLen == 0) || (tempLoopInfo.actS1Size == 0)) { + tempLoopInfo.curActSeqLenIsZero = true; + return; + } + tempLoopInfo.curActSeqLenIsZero = false; + tempLoopInfo.s2BasicSizeTail = tempLoopInfo.curActualSeqLen % constInfo.s2BaseSize; + tempLoopInfo.s2BasicSizeTail = + (tempLoopInfo.s2BasicSizeTail == 0) ? constInfo.s2BaseSize : tempLoopInfo.s2BasicSizeTail; + tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize; + tempLoopInfo.mBasicSizeTail = + (tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail; + tempLoopInfo.s2LoopTimes = 0; +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::UpdateInner(uint32_t &s2End, uint32_t &curS2End, + uint32_t s1Idx, bool isEnd) +{ + uint32_t s1BaseSize = 1; + int64_t s1Offset = s1BaseSize * s1Idx; + int64_t s2LastToken = Min(s1Offset + tempLoopInfo.nextTokensPerBatch + s1BaseSize, tempLoopInfo.curActualSeqLenOri); + s2LastToken = Min(constInfo.sparseBlockSize * constInfo.sparseBlockCount, s2LastToken); + curS2End = (s2LastToken + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + tempLoopInfo.s2LoopTimes = isEnd ? constInfo.s2End + 1 : curS2End; +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::Init(__gm__ uint8_t *query, + __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t* keyScale, + __gm__ uint8_t* valueScale, __gm__ uint8_t *blockTable, __gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengths, __gm__ uint8_t *attentionOut, __gm__ uint8_t *softmaxMax, + __gm__ uint8_t *softmaxSum, __gm__ uint8_t *workspace, + const KvQuantSparseFlashAttentionTilingDataMla *__restrict tiling, + __gm__ uint8_t *gmTiling, TPipe *tPipe) +{ + if ASCEND_IS_AIC { + tmpBlockIdx = GetBlockIdx(); // cube:0-23 + aiCoreIdx = tmpBlockIdx; + } else { + tmpBlockIdx = GetBlockIdx(); // vec:0-47 + aiCoreIdx = tmpBlockIdx / 2; + } + + // init tiling data + tilingData = tiling; + + InitTilingData(); + InitActualSeqLen(actualSeqLengthsQ, actualSeqLengths); + + // 初始化计算参数 + InitCalcParamsEach(); + keyPtr = key; + valuePtr = value; + pipe = tPipe; + + // init global buffer + queryGm.SetGlobalBuffer((__gm__ Q_T *)query); + keyGm.SetGlobalBuffer((__gm__ KV_T *)keyPtr); + valueGm.SetGlobalBuffer((__gm__ KV_T *)valuePtr); + + attentionOutGm.SetGlobalBuffer((__gm__ OUT_T *)attentionOut); + softmaxMaxGm.SetGlobalBuffer((__gm__ T *)softmaxMax); + softmaxSumGm.SetGlobalBuffer((__gm__ T *)softmaxSum); + + if ASCEND_IS_AIV { + if (constInfo.needInit && LAYOUT_T != QSFA_LAYOUT::TND) { + InitOutputSingleCore(); + } + } + + if constexpr (PAGE_ATTENTION) { + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + } + topKGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices); + + // workspace 内存排布 + // |Q--|mm1ResGm(存S)|vec1ResGm(存A1,A2)|mm2ResGm(存O)|vec2ResGm + // |Core0_Q1-Core0_Q2-Core1_Q1-Core1_Q2....Core32_Q1-Core32_Q2|Core0_mmRes + uint64_t qsfaOffset = 0; + mm1ResGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + qsfaOffset + + aiCoreIdx * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(MM1_OUT_T))); + qsfaOffset += GetBlockNum() * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(MM1_OUT_T); + + vec1ResGm.SetGlobalBuffer( + (__gm__ K_ROPE_T *)(workspace + qsfaOffset + aiCoreIdx * dbWorkspaceRatio * constInfo.mmResUbSize * + sizeof(K_ROPE_T))); + qsfaOffset += GetBlockNum() * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(K_ROPE_T); + + mm2ResGm.SetGlobalBuffer( + (__gm__ MM2_OUT_T *)(workspace + qsfaOffset + + aiCoreIdx * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(MM2_OUT_T))); + qsfaOffset += GetBlockNum() * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(MM2_OUT_T); + mm2ResInt32Gm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(mm2ResGm.GetPhyAddr(0))); + + vec2ResGm.SetGlobalBuffer((__gm__ T *)(workspace + qsfaOffset + + aiCoreIdx * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(T))); + qsfaOffset += GetBlockNum() * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(MM2_OUT_T); + + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + // s2 d+rope bufNum + kvMergeGm_.SetGlobalBuffer((__gm__ K_ROPE_T *)(workspace + qsfaOffset + aiCoreIdx * 512 * 576 * 4 * + sizeof(K_ROPE_T))); + qsfaOffset += GetBlockNum() * 512 * 576 * 4 * sizeof(K_ROPE_T); + + kvValidSizeGm_.SetGlobalBuffer( + (__gm__ int32_t *)(workspace + qsfaOffset + (aiCoreIdx * 2) * 128 * 4 * sizeof(int32_t))); + } + + if constexpr (FLASH_DECODE) { + accumOutGm.SetGlobalBuffer((__gm__ float *)(workspace + qsfaOffset)); + qsfaOffset = qsfaOffset + tilingData->splitKVParams.accumOutSize * sizeof(float); + lseSumFdGm.SetGlobalBuffer((__gm__ float *)(workspace + qsfaOffset)); + lseMaxFdGm.SetGlobalBuffer((__gm__ float *)(workspace + qsfaOffset) + + tilingData->splitKVParams.logSumExpSize / 2); + qsfaOffset = qsfaOffset + tilingData->splitKVParams.logSumExpSize * sizeof(float); + } + + if ASCEND_IS_AIV { + vectorService.InitParams(constInfo, tilingData); + vectorService.InitMm2ResInt32GmGlobalTensor(mm2ResInt32Gm); + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + vectorService.InitVec0GlobalTensor(kvValidSizeGm_, kvMergeGm_, kRopeGm, keyGm, blockTableGm); + } + vectorService.InitVec1GlobalTensor(mm1ResGm, vec1ResGm, actualSeqLengthsQGm, + actualSeqLengthsKVGm, lseMaxFdGm, lseSumFdGm, topKGm, + softmaxMaxGm, softmaxSumGm); + vectorService.InitVec2GlobalTensor(accumOutGm, vec2ResGm, mm2ResGm, attentionOutGm); + } + + if ASCEND_IS_AIC { + matmulService.InitParams(constInfo); + matmulService.InitMm1GlobalTensor(queryGm, qRopeGm, keyGm, kRopeGm, mm1ResGm); + matmulService.InitMm2GlobalTensor(vec1ResGm, valueGm, mm2ResGm, attentionOutGm); + matmulService.InitPageAttentionInfo(kvMergeGm_, blockTableGm, topKGm, + constInfo.kvCacheBlockSize, constInfo.maxBlockNumPerBatch); + } + // 要在InitParams之后执行 + if (pipe != nullptr) { + InitBuffers(); + } +} + +template __aicore__ inline void KvQuantSparseFlashAttentionMla::InitCalcParamsEach() +{ + // 计算总的基本块 + uint32_t totalBaseNum = 0; + uint32_t s1GBaseSize = constInfo.gSize; + uint32_t actBatchS2 = 1; + uint32_t coreNum = GetBlockNum(); + uint32_t actBatchS1 = 1; + uint32_t currCoreIdx = aiCoreIdx; + for (uint32_t bIdx = 0; bIdx < constInfo.batchSize; bIdx++) { + uint32_t actBatchS1 = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bIdx); + if (actBatchS1 < constInfo.qSeqSize) { + constInfo.needInit = true; + } + totalBaseNum += actBatchS1 * actBatchS2 ; + } + uint32_t avgBaseNum = 1; + if (totalBaseNum > coreNum) { + avgBaseNum = (totalBaseNum + coreNum - 1) / coreNum; + } else { + usedCoreNum = totalBaseNum; + } + if (aiCoreIdx >= usedCoreNum) { + return; + } + // 计算当前核的基本块 + uint32_t accumBaseNum = 0; // 当前累积的基本块数 + uint32_t targetBaseNum = 0; + uint32_t lastValidBIdx = 0; + uint32_t lastValidactBatchS1 = 0; + bool setStart = false; + targetBaseNum = (currCoreIdx + 1) * avgBaseNum; // 计算当前的目标权重 + uint32_t targetStartBaseNum = targetBaseNum-avgBaseNum; + for (uint32_t bN2Idx = 0; bN2Idx < constInfo.batchSize * constInfo.kvHeadNum; bN2Idx++) { + uint32_t bIdx = bN2Idx / constInfo.kvHeadNum; + actBatchS1 = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bIdx); + for (uint32_t s1GIdx = 0; s1GIdx < actBatchS1; s1GIdx++) { + accumBaseNum += 1; + if (!setStart && accumBaseNum >= targetStartBaseNum) { + constInfo.bN2Start = bN2Idx; + constInfo.gS1Start = s1GIdx; + setStart = true; + } + if (accumBaseNum >= targetBaseNum) { + // 更新当前核的End分核信息 + constInfo.bN2End = bN2Idx; + constInfo.gS1End = s1GIdx; + constInfo.coreStartKVSplitPos = 0; + constInfo.s2End = 0; + if (aiCoreIdx != 0) { + GetAxisStartIdx(constInfo.bN2Start, constInfo.gS1Start, 0); + } + return; + } + } + if ((actBatchS1 > 0) && (actBatchS2 > 0)) { + lastValidactBatchS1 = actBatchS1; + lastValidBIdx = bIdx; + } + } + if (!setStart) { + constInfo.bN2Start = lastValidBIdx; + constInfo.gS1Start = lastValidactBatchS1 - 1; + } + if (accumBaseNum < targetBaseNum) { + // 更新最后一个核的End分核信息 + constInfo.bN2End = lastValidBIdx; + constInfo.gS1End = lastValidactBatchS1-1; + constInfo.s2End = 0; + constInfo.coreStartKVSplitPos = 0; + if (aiCoreIdx != 0) { + GetAxisStartIdx(constInfo.bN2Start, constInfo.gS1Start, 0); + } + return; + } +} + +template +__aicore__ inline void +KvQuantSparseFlashAttentionMla::Bmm2DataCopyOut(uint64_t attenOutOffset, LocalTensor &attenOutUb, + uint32_t startRow, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = dealRowCount; + dataCopyParams.blockLen = actualColumnCount * sizeof(OUT_T); + dataCopyParams.srcStride = (columnCount - actualColumnCount) / (QSFAVectorService::BYTE_BLOCK / + sizeof(OUT_T)); + dataCopyParams.dstStride = 0; + DataCopyPad(attentionOutGm[attenOutOffset + (mSizeVStart + startRow) * actualColumnCount], attenOutUb, + dataCopyParams); +} + + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::CalcParams(uint32_t loop, uint64_t s2Start, + uint32_t s2LoopIdx, RunInfo &info) +{ + info.loop = loop; + info.bIdx = tempLoopInfo.bIdx; + info.gS1Idx = tempLoopInfo.gS1Idx; + info.s2Idx = s2LoopIdx; + info.curSInnerLoopTimes = tempLoopInfo.s2LoopTimes; + + info.isBmm2Output = false; + info.tndIsS2SplitCore = tempLoopInfo.tndIsS2SplitCore; + info.tndCoreStartKVSplitPos = tempLoopInfo.tndCoreStartKVSplitPos; + + info.actS1Size = tempLoopInfo.actS1Size; + info.actS2Size = tempLoopInfo.curActualSeqLen; + + info.actMBaseSize = constInfo.mBaseSize; + uint32_t qsfaRemainedGS1Size = tempLoopInfo.actS1Size * constInfo.gSize - tempLoopInfo.gS1Idx; + if (qsfaRemainedGS1Size <= constInfo.mBaseSize && qsfaRemainedGS1Size > 0) { + info.actMBaseSize = tempLoopInfo.mBasicSizeTail; + } + + info.isValid = s2LoopIdx < tempLoopInfo.s2LoopTimes; + CalcMSizeInfo(info); + + info.isChangeBatch = false; + + info.isFirstSInnerLoop = (s2LoopIdx == s2Start); + if (info.isFirstSInnerLoop) { + tempLoopInfo.bn2IdxInCurCore++; + } + info.isLastS2Loop = (s2LoopIdx == tempLoopInfo.s2LoopTimes - 1); + info.bn2IdxInCurCore = tempLoopInfo.bn2IdxInCurCore - 1; + uint64_t qsfaActualSeqQPrefixSum; + if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + qsfaActualSeqQPrefixSum = (info.bIdx <= 0) ? 0 : actualSeqLengthsQGm.GetValue(info.bIdx - 1); + } else { + qsfaActualSeqQPrefixSum = (info.bIdx <= 0) ? 0 : info.bIdx * constInfo.qSeqSize; + } + info.tndBIdxOffsetForQ = qsfaActualSeqQPrefixSum * constInfo.qHeadNum * constInfo.combineHeadDim; + + uint64_t actualSeqKVPrefixSum; + if constexpr (KV_LAYOUT_T == QSFA_LAYOUT::TND) { + actualSeqKVPrefixSum = (info.bIdx <= 0) ? 0 : actualSeqLengthsKVGm.GetValue(info.bIdx - 1); + } else { + actualSeqKVPrefixSum = (info.bIdx <= 0) ? 0 : info.bIdx * constInfo.kvSeqSize; + } + info.tndBIdxOffsetForKV = actualSeqKVPrefixSum * constInfo.kvHeadNum * constInfo.combineHeadDim; + + CalcFirstTensorOffsets(info, qsfaActualSeqQPrefixSum, actualSeqKVPrefixSum); + + uint64_t sInnerOffsetDataSize = info.s2Idx * constInfo.s2BaseSize; + info.s2BatchOffset = s2BatchBaseOffset + sInnerOffsetDataSize; + + info.curActualSeqLenOri = tempLoopInfo.curActualSeqLenOri; + if (tempLoopInfo.curActualSeqLen > sInnerOffsetDataSize) { + info.actualSingleProcessSInnerSize = tempLoopInfo.curActualSeqLen - sInnerOffsetDataSize; + info.actualSingleProcessSInnerSize = info.actualSingleProcessSInnerSize > constInfo.s2BaseSize ? + constInfo.s2BaseSize : info.actualSingleProcessSInnerSize; + } else { + info.actualSingleProcessSInnerSize = 0; + } + info.actualSingleProcessSInnerSizeAlign = + QSFAAlign((uint32_t)info.actualSingleProcessSInnerSize, (uint32_t)QSFAVectorService::BYTE_BLOCK); +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::CalcMSizeInfo(RunInfo &info) +{ + if ASCEND_IS_AIV { + info.mSize = info.actMBaseSize; + info.mSizeV = (info.mSize <= 16) ? \ + info.mSize : (((info.mSize + 15) / 16 + 1) / 2 * 16); + info.mSizeVStart = 0; + if (tmpBlockIdx % 2 == 1) { + info.mSizeVStart = info.mSizeV; + info.mSizeV = info.mSize - info.mSizeV; + } + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::CalcFirstTensorOffsets( + RunInfo &info, uint64_t qsfaActualSeqQPrefixSum, uint64_t actualSeqKVPrefixSum) +{ + if (info.isFirstSInnerLoop) { + tensorACoreOffset = info.tndBIdxOffsetForQ + info.gS1Idx * constInfo.combineHeadDim; + tensorBCoreOffset = info.tndBIdxOffsetForKV + info.n2Idx * constInfo.combineHeadDim; + if (constInfo.quantScaleRepoMode == QUANT_SCALE_REPO_MODE::COMBINE) { + attenOutOffset = (qsfaActualSeqQPrefixSum * constInfo.qHeadNum + info.gS1Idx) * headDim; + } else { + uint64_t tndBIdxRopeOffsetForQ = qsfaActualSeqQPrefixSum * constInfo.qHeadNum * headDimRope; + tensorARopeCoreOffset = tndBIdxRopeOffsetForQ + info.gS1Idx * headDimRope; + uint64_t tndBIdxRopeOffsetForK = actualSeqKVPrefixSum * constInfo.kvHeadNum * headDimRope; + tensorBRopeCoreOffset = tndBIdxRopeOffsetForK + info.n2Idx * headDimRope; + attenOutOffset = tensorACoreOffset; + } + if (constInfo.sparseMode == 3) { + threshold = static_cast(tempLoopInfo.nextTokensPerBatch) + info.gS1Idx / constInfo.gSize + 1; + } else { + threshold = tempLoopInfo.curActualSeqLenOri; + } + if constexpr (LAYOUT_T == QSFA_LAYOUT::BSND) { + topKBaseOffset = info.bIdx * constInfo.qSeqSize * constInfo.kvHeadNum * constInfo.sparseBlockCount + + info.gS1Idx / constInfo.gSize * constInfo.kvHeadNum * constInfo.sparseBlockCount + + info.n2Idx * constInfo.sparseBlockCount; + } else if (LAYOUT_T == QSFA_LAYOUT::TND) { + topKBaseOffset = info.tndBIdxOffsetForQ / constInfo.gSize / constInfo.combineHeadDim * + constInfo.kvHeadNum * constInfo.sparseBlockCount + + info.n2Idx * constInfo.sparseBlockCount + + info.gS1Idx / constInfo.gSize * constInfo.kvHeadNum * constInfo.sparseBlockCount; + } else { + topKBaseOffset = info.bIdx * constInfo.kvHeadNum * constInfo.qSeqSize * constInfo.sparseBlockCount + + info.n2Idx * constInfo.qSeqSize * constInfo.sparseBlockCount + + info.gS1Idx / constInfo.gSize * constInfo.sparseBlockCount; + } + } + info.topKBaseOffset = topKBaseOffset; + info.threshold = threshold; + info.tensorAOffset = tensorACoreOffset; + info.tensorARopeOffset = tensorARopeCoreOffset; + info.tensorBOffset = tensorBCoreOffset; + info.tensorBRopeOffset = tensorBRopeCoreOffset; + info.attenOutOffset = attenOutOffset; +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::ComputeMm1(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + matmulService.ComputeMm1(info, mSplitInfo); + CrossCoreSetFlag(constInfo.syncC1V1); + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::ComputeMm2(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + CrossCoreWaitFlag(constInfo.syncV1C2); + matmulService.ComputeMm2(info, mSplitInfo); + CrossCoreSetFlag(constInfo.syncC2V2); + // CrossCoreSetFlag(constInfo.syncC2V1); + } +} + +template __aicore__ inline void KvQuantSparseFlashAttentionMla::Process() +{ + if (aiCoreIdx < usedCoreNum) { + if ASCEND_IS_AIC { + matmulService.AllocEventID(); + } else { + vectorService.AllocEventID(); + vectorService.InitSoftmaxDefaultBuffer(); + } + ProcessBalance(); + + if ASCEND_IS_AIC { + matmulService.FreeEventID(); + } else { + vectorService.FreeEventID(); + } + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, + uint32_t &n2Idx) +{ + bIdx = bN2Idx / kvHeadNum; + n2Idx = bN2Idx % kvHeadNum; +} + +template __aicore__ inline void KvQuantSparseFlashAttentionMla::ProcessBalance() +{ + RunInfo extraInfo[QSFA_PRELOAD_TASK_CACHE_SIZE]; + uint32_t gloop = 0; + int gS1LoopEnd; + bool globalLoopStart = true; + if ASCEND_IS_AIC { + // CrossCoreSetFlag(constInfo.syncC2V1); + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + CrossCoreSetFlag(3); + CrossCoreSetFlag(3); + CrossCoreSetFlag(3); + CrossCoreSetFlag(3); + } + } + for (uint32_t qsfaBN2LoopIdx = constInfo.bN2Start; qsfaBN2LoopIdx <= constInfo.bN2End; qsfaBN2LoopIdx++) { + GetBN2Idx(qsfaBN2LoopIdx, tempLoopInfo.bIdx, tempLoopInfo.n2Idx); + GetActualSeqLen(tempLoopInfo.bIdx); // 获取actualSeqLength及ActualSeqLengthKV + GetPreNextTokensLeftUp(); + if (tempLoopInfo.actS1Size == 0) { + continue; + } + int gS1SplitNum = (tempLoopInfo.actS1Size * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize; + gS1LoopEnd = (qsfaBN2LoopIdx == constInfo.bN2End) ? constInfo.gS1End : gS1SplitNum - 1; + for (uint32_t qsfaGS1LoopIdx = constInfo.gS1Start; qsfaGS1LoopIdx <= gS1LoopEnd; qsfaGS1LoopIdx++) { + tempLoopInfo.gS1Idx = qsfaGS1LoopIdx * constInfo.mBaseSize; + // TopK值sparse完后的ActualSeqLengthKV + GetSparseActualSeqLen(tempLoopInfo.bIdx, qsfaGS1LoopIdx, tempLoopInfo.n2Idx); + UpdateInnerLoopCond(); + + if (tempLoopInfo.curActSeqLenIsZero) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, qsfaGS1LoopIdx, tempLoopInfo.n2Idx); + } + int s2SplitNum = + (tempLoopInfo.curActualSeqLen + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; // S2切分份数 + bool qsfaIsEnd = (qsfaBN2LoopIdx == constInfo.bN2End) && (qsfaGS1LoopIdx == constInfo.gS1End); + tempLoopInfo.s2LoopTimes = s2SplitNum; + // 分核修改后需要打开 + // 当前s2是否被切，决定了输出是否要写到attenOut上 + tempLoopInfo.tndIsS2SplitCore = + ((constInfo.s2Start == 0) && (tempLoopInfo.s2LoopTimes == s2SplitNum)) ? false : true; + tempLoopInfo.tndCoreStartKVSplitPos = globalLoopStart ? constInfo.coreStartKVSplitPos : 0; + uint32_t qsfaExtraLoop = qsfaIsEnd ? 2 : 0; + for (int s2LoopIdx = constInfo.s2Start; s2LoopIdx < (tempLoopInfo.s2LoopTimes + qsfaExtraLoop); + s2LoopIdx++) { + // PreloadPipeline loop初始值要求为 PRELOAD_NUM + PreloadPipeline(gloop, constInfo.s2Start, s2LoopIdx, extraInfo); + ++gloop; + } + globalLoopStart = false; + constInfo.s2Start = 0; + } + constInfo.gS1Start = 0; + } + if ASCEND_IS_AIV { + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + CrossCoreWaitFlag(3); + CrossCoreWaitFlag(3); + CrossCoreWaitFlag(3); + CrossCoreWaitFlag(3); + } + } +} + +template +__aicore__ inline void +KvQuantSparseFlashAttentionMla::PreloadPipeline(uint32_t loop, uint64_t s2Start, uint64_t s2LoopIdx, + RunInfo extraInfo[QSFA_PRELOAD_TASK_CACHE_SIZE]) +{ + RunInfo &extraInfo0 = extraInfo[loop % QSFA_PRELOAD_TASK_CACHE_SIZE]; // 本轮任务 + RunInfo &extraInfo2 = extraInfo[(loop + 2) % QSFA_PRELOAD_TASK_CACHE_SIZE]; // 上一轮任务 + RunInfo &extraInfo1 = extraInfo[(loop + 1) % QSFA_PRELOAD_TASK_CACHE_SIZE]; // 上两轮任务 + + CalcParams(loop, s2Start, s2LoopIdx, extraInfo0); + + if (extraInfo0.isValid) { + if ASCEND_IS_AIC { + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + CrossCoreWaitFlag(constInfo.syncV0C1); + } + ComputeMm1(extraInfo0); + } else { + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + CrossCoreWaitFlag(3); + vectorService.MergeKv(extraInfo0); + CrossCoreSetFlag(constInfo.syncV0C1); + } + } + } + if (extraInfo2.isValid) { + if ASCEND_IS_AIV { + vectorService.ProcessVec1L(extraInfo2); + } + if ASCEND_IS_AIC { + ComputeMm2(extraInfo2); + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + CrossCoreSetFlag(3); + } + } + } + if (extraInfo1.isValid) { + if ASCEND_IS_AIV { + vectorService.ProcessVec2L(extraInfo1); + } + extraInfo1.isValid = false; + } +} + +template +__aicore__ inline uint64_t +KvQuantSparseFlashAttentionMla::GetBalanceActualSeqLengths(GlobalTensor &actualSeqLengths, + uint32_t bIdx) +{ + if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + if (bIdx > 0) { + int32_t curActualSeqLen = actualSeqLengths.GetValue(bIdx); + int32_t prevActualSeqLen = actualSeqLengths.GetValue(bIdx - 1); + return (curActualSeqLen >= prevActualSeqLen) ? + static_cast(curActualSeqLen - prevActualSeqLen) : 0ULL; + } else if (bIdx == 0) { + return actualSeqLengths.GetValue(0); + } else { + return 0; + } + } else { + if (constInfo.actualLenDimsQ == 1) { + return actualSeqLengths.GetValue(0); + } else if (constInfo.actualLenDimsQ == 0) { + return constInfo.qSeqSize; + } else { + return actualSeqLengths.GetValue(bIdx); + } + } +} + +template +__aicore__ inline void KvQuantSparseFlashAttentionMla::GetAxisStartIdx(uint32_t bN2EndPrev, + uint32_t s1GEndPrev, uint32_t s2EndPrev) +{ + uint32_t qsfaBEndPrev = bN2EndPrev / kvHeadNum; + uint32_t qsfaActualSeqQPrev = GetBalanceActualSeqLengths(actualSeqLengthsQGm, qsfaBEndPrev); + uint32_t qsfaS1GPrevBaseNum = + (qsfaActualSeqQPrev * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize; + constInfo.bN2Start = bN2EndPrev; + constInfo.gS1Start = s1GEndPrev; + + constInfo.s2Start = 0; + if (s1GEndPrev >= qsfaS1GPrevBaseNum - 1) { // 上个核把S1G处理完了 + constInfo.gS1Start = 0; + constInfo.bN2Start++; + } else { + constInfo.gS1Start++; + } +} +#endif // KV_QUANT_SPARSE_FLASH_ATTENTION_KERNEL_MLA_H diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_service_cube_mla.h b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_service_cube_mla.h new file mode 100644 index 000000000..1614a6e4e --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_service_cube_mla.h @@ -0,0 +1,943 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_service_cube_mla.h + * \brief use 7 buffer for matmul l1, better pipeline + */ +#ifndef KV_QUANT_SPARSE_FLASH_ATTENTION_SERVICE_CUBE_MLA_H +#define KV_QUANT_SPARSE_FLASH_ATTENTION_SERVICE_CUBE_MLA_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "kv_quant_sparse_flash_attention_common.h" + +struct Position { + uint32_t bIdx; + uint32_t n2Idx; + uint32_t s2Idx; + uint32_t dIdx; +}; + +struct PAShape { + uint32_t blockSize; + uint32_t headNum; // 一般为kv的head num，对应n2 + uint32_t headDim; // mla下rope为64，nope为512, 对应d + uint32_t maxblockNumPerBatch; // block table 每一行的最大个数 + uint32_t actHeadDim; // 实际拷贝col大小,考虑到N切块 s*d, 对应d + uint32_t copyRowNum; // 总共要拷贝的行数 + uint32_t copyRowNumAlign; +}; + +// 场景：query、queryRope、key、value GM to L1 +// GM按ND格式存储 +// L1按NZ格式存储 +// GM的行、列、列的stride +template +__aicore__ inline void DataCopyGmNDToL1(LocalTensor &l1Tensor, GlobalTensor &gmTensor, + uint32_t rowAct, uint32_t rowAlign, + uint32_t col, // D + uint32_t colStride) // D or N*D +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = rowAct; // nd矩阵的行数 + // T为int4场景下，dValue = col / 2，srcDValue = colStride / 2 + nd2nzPara.srcDValue = colStride; // 同一nd矩阵相邻行起始地址间的偏移 + nd2nzPara.dValue = col; // nd矩阵的列数 + nd2nzPara.dstNzC0Stride = rowAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.dstNzMatrixStride = 0; + nd2nzPara.srcNdMatrixStride = 0; + DataCopy(l1Tensor, gmTensor, nd2nzPara); +} + +/* + 适用PA数据从GM拷贝到L1，支持ND、NZ数据； + PA的layout分 BNBD（blockNum,N,blockSize,D） BBH（blockNum,blockSize,N*D + BSH\BSND\TND 为BBH + shape.copyRowNumAlign 需要16字节对齐，如拷贝k矩阵，一次拷贝128*512，遇到尾块 10*512 需对齐到16*512 +*/ +template +__aicore__ inline void DataCopyPA(LocalTensor &dstTensor, // l1 + GlobalTensor &srcTensor, // gm + GlobalTensor &blockTableGm, + const PAShape &shape, // blockSize, headNum, headDim + const Position &startPos) // bacthIdx nIdx curSeqIdx +{ + uint32_t copyFinishRowCnt = 0; + uint64_t blockTableBaseOffset = startPos.bIdx * shape.maxblockNumPerBatch; + uint32_t curS2Idx = startPos.s2Idx; + uint32_t blockElementCnt = 32 / sizeof(T); + while (copyFinishRowCnt < shape.copyRowNum) { + uint64_t blockIdOffset = curS2Idx / shape.blockSize; // 获取block table上的索引 + uint64_t reaminRowCnt = curS2Idx % shape.blockSize; // 获取在单个块上超出的行数 + // 从block table上的获取编号 + uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset); + // 计算可以拷贝行数 + uint32_t copyRowCnt = shape.blockSize - reaminRowCnt; // 一次只能处理一个Block + if (copyFinishRowCnt + copyRowCnt > shape.copyRowNum) { + copyRowCnt = shape.copyRowNum - copyFinishRowCnt; // 一个block未拷满 + } + uint64_t offset = idInBlockTable * shape.blockSize * shape.headNum * shape.headDim ; // PA的偏移 + + uint64_t dStride = shape.headDim; + if constexpr (SRC_LAYOUT == QSFA_LAYOUT::BSND || SRC_LAYOUT == QSFA_LAYOUT::TND) { + offset += (uint64_t)(startPos.n2Idx * shape.headDim) + + reaminRowCnt * shape.headDim * shape.headNum + startPos.dIdx; + dStride = shape.headDim * shape.headNum; + } else { + offset += (uint64_t)(startPos.n2Idx * shape.headDim * shape.blockSize) + + reaminRowCnt * shape.headDim + startPos.dIdx; + } + + uint32_t srcDValue = dStride; + uint32_t dValue = shape.actHeadDim; + LocalTensor tmpDstTensor = dstTensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpSrcTensor = srcTensor[offset]; + + DataCopyGmNDToL1(tmpDstTensor, tmpSrcTensor, copyRowCnt, shape.copyRowNumAlign, dValue, srcDValue); + copyFinishRowCnt += copyRowCnt; + curS2Idx += copyRowCnt; + } +} + +template class QSFAMatmulService { +public: + // 中间计算数据类型为float, 高精度模式 + using T = float; + using Q_T = typename QSFAT::queryType; + using KV_T = typename QSFAT::kvType; + using K_ROPE_T = typename QSFAT::kRopeType; + using OUT_T = typename QSFAT::outputType; + using MM_OUT_T = T; + + __aicore__ inline QSFAMatmulService(){}; + __aicore__ inline void InitParams(const ConstInfo &constInfo); + __aicore__ inline void InitMm1GlobalTensor(GlobalTensor queryGm, GlobalTensor qRopeGm, + GlobalTensor keyGm, GlobalTensor kRopeGm, + GlobalTensor mm1ResGm); + __aicore__ inline void InitMm2GlobalTensor(GlobalTensor vec1ResGm, GlobalTensor valueGm, + GlobalTensor mm2ResGm, GlobalTensor attentionOutGm); + __aicore__ inline void InitPageAttentionInfo(const GlobalTensor& kvMergeGm, + GlobalTensor blockTableGm, GlobalTensor topKGm, + uint32_t blockSize, uint32_t maxBlockNumPerBatch); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void UpdateKey(GlobalTensor keyGm); + __aicore__ inline void UpdateValue(GlobalTensor valueGm); + + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void CalcTopKBlockInfo(const RunInfo &info, uint32_t &curTopKIdx, + uint64_t &curOffsetInSparseBlock, uint32_t curSeqIdx, + uint32_t ©RowCnt, uint64_t &idInTopK); + __aicore__ inline void ComputeMm1(const RunInfo &info, const MSplitInfo mSplitInfo); + __aicore__ inline void ComputeMm2(const RunInfo &info, const MSplitInfo mSplitInfo); + +private: + static constexpr bool PAGE_ATTENTION = QSFAT::pageAttention; + static constexpr int TEMPLATE_MODE = QSFAT::templateMode; + static constexpr bool FLASH_DECODE = QSFAT::flashDecode; + static constexpr QSFA_LAYOUT LAYOUT_T = QSFAT::layout; + static constexpr QSFA_LAYOUT KV_LAYOUT_T = QSFAT::kvLayout; + + static constexpr uint32_t M_SPLIT_SIZE = 128; // m方向切分 + static constexpr uint32_t N_SPLIT_SIZE = 128; // n方向切分 + static constexpr uint32_t N_WORKSPACE_SIZE = 512; // n方向切分 + static constexpr uint32_t K_SPLIT_SIZE = 288; // K方向切分 + + static constexpr uint32_t L1_BLOCK_SIZE = (64 * (512 + 64) * sizeof(Q_T)); + static constexpr uint32_t L1_BLOCK_OFFSET = 64 * (512 + 64); // 72K的元素个数 + + static constexpr uint32_t L0A_PP_SIZE = (32 * 1024); + static constexpr uint32_t L0B_PP_SIZE = (32 * 1024); + static constexpr uint32_t L0C_PP_SIZE = (64 * 1024); + + // m <> mte1 EventID + static constexpr uint32_t L0AB_EVENT0 = EVENT_ID3; + static constexpr uint32_t L0AB_EVENT1 = EVENT_ID4; + + // mte2 <> mte1 EventID + // L1 3buf, 使用3个eventId + static constexpr uint32_t L1_EVENT0 = EVENT_ID2; + static constexpr uint32_t L1_EVENT1 = EVENT_ID3; + static constexpr uint32_t L1_EVENT2 = EVENT_ID4; + static constexpr uint32_t L1_EVENT3 = EVENT_ID5; + static constexpr uint32_t L1_EVENT4 = EVENT_ID6; + static constexpr uint32_t L1_EVENT5 = EVENT_ID7; + static constexpr uint32_t L1_EVENT6 = EVENT_ID1; + + static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding; + static constexpr uint32_t mte21QPIds[4] = {L1_EVENT0, L1_EVENT1, L1_EVENT2, L1_EVENT3}; // mte12复用 + static constexpr uint32_t mte21KVIds[3] = {L1_EVENT4, L1_EVENT5, L1_EVENT6}; + + static constexpr uint32_t BLOCK_ELEMENT_NUM = ConstInfo::BUFFER_SIZE_BYTE_32B / sizeof(K_ROPE_T); + + uint32_t kvCacheBlockSize = 0; + uint32_t maxBlockNumPerBatch = 0; + ConstInfo constInfo{}; + + // L1分成3块buf, 用于记录 + uint32_t qpL1BufIter = 0; + uint32_t kvL1BufIter = -1; + uint32_t abL0BufIter = 0; + uint32_t cL0BufIter = 0; + + // mm1 + GlobalTensor queryGm; + GlobalTensor qRopeGm; + GlobalTensor keyGm; + GlobalTensor kRopeGm; + GlobalTensor mm1ResGm; + GlobalTensor kvMergeGm_; + + // mm2 + GlobalTensor vec1ResGm; + GlobalTensor valueGm; + GlobalTensor mm2ResGm; + GlobalTensor attentionOutGm; + + // block_table + GlobalTensor topKGm; + GlobalTensor blockTableGm; + + TBuf bufQPL1; + TBuf bufKVL1; + TBuf tmpBufL0A; + TBuf tmpBufL0B; + TBuf tmpBufL0C; + + LocalTensor aL0TensorPingPong; + LocalTensor bL0TensorPingPong; + LocalTensor cL0TensorPingPong; + LocalTensor l1QPTensor; + LocalTensor l1KVTensor; + + // L0AB m <> mte1 EventID + __aicore__ inline uint32_t Mte1MmABEventId(uint32_t idx) + { + return (L0AB_EVENT0 + idx); + } + + __aicore__ inline uint32_t GetQPL1RealIdx(uint32_t mIdx, uint32_t k1Idx) + { + uint32_t idxMap[] = {0, 2}; // 确保0块和1块连在一起, 2和3块连在一起, 来保证同一m块的地址相连 + return idxMap[mIdx % 2] + k1Idx; + } + + __aicore__ inline void CopyGmToL1(LocalTensor &l1Tensor, GlobalTensor &gmSrcTensor, + uint32_t srcN, uint32_t srcD, uint32_t srcDstride); + __aicore__ inline void CopyInMm1AToL1(LocalTensor &aL1Tensor, const RunInfo &info, uint32_t mSeqIdx, + uint32_t mSizeAct, uint32_t headSize, uint32_t headOffset); + __aicore__ inline void CopyInMm1ARopeToL1(LocalTensor &aL1Tensor, const RunInfo &info, uint32_t mSeqIdx, + uint32_t mSizeAct); + __aicore__ inline void CopyInMm1BToL1(LocalTensor &bL1Tensor, const uint64_t keyGmBaseOffset, + uint32_t copyTotalRowCntAlign, uint32_t copyStartRowCnt, + uint32_t nActCopyRowCount, uint32_t headSize); + __aicore__ inline void CopyInMm1BRopeToL1(LocalTensor &bL1Tensor, const uint64_t keyGmBaseOffset, + uint32_t copyTotalRowCntAlign, uint32_t copyStartRowCnt, + uint32_t nActCopyRowCount, uint32_t headSize); + __aicore__ inline void CopyInMm2AToL1(LocalTensor &aL1Tensor, const RunInfo &info, uint32_t mSeqIdx, + uint32_t subMSizeAct, uint32_t nSize, uint32_t nOffset); + __aicore__ inline void CopyInMm2BToL1(LocalTensor &bL1Tensor, const uint64_t valueGmBaseOffset, + uint32_t copyTotalRowCntAlign, uint32_t copyStartRowCnt, + uint32_t nActCopyRowCount, uint32_t copyStartColumnCount, + uint32_t copyColumnCount); + __aicore__ inline void LoadDataMm1A(LocalTensor &aL0Tensor, LocalTensor &aL1Tensor, + uint32_t idx, uint32_t kSplitSize, uint32_t mSize, uint32_t kSize); + __aicore__ inline void LoadDataMm1B(LocalTensor &bL0Tensor, LocalTensor &bL1Tensor, + uint32_t idx, uint32_t kSplitSize, uint32_t kSize, uint32_t nSize); +}; + +template __aicore__ inline void QSFAMatmulService::InitParams(const ConstInfo &constInfo) +{ + this->constInfo = constInfo; +} + +template +__aicore__ inline void +QSFAMatmulService::InitMm1GlobalTensor(GlobalTensor queryGm, GlobalTensor qRopeGm, + GlobalTensor keyGm, GlobalTensor kRopeGm, + GlobalTensor mm1ResGm) +{ + // mm1 + this->queryGm = queryGm; + this->qRopeGm = qRopeGm; + this->keyGm = keyGm; + this->kRopeGm = kRopeGm; + this->mm1ResGm = mm1ResGm; +} + +template +__aicore__ inline void +QSFAMatmulService::InitMm2GlobalTensor(GlobalTensor vec1ResGm, GlobalTensor valueGm, + GlobalTensor mm2ResGm, GlobalTensor attentionOutGm) +{ + // mm2 + this->vec1ResGm = vec1ResGm; + this->valueGm = valueGm; + this->mm2ResGm = mm2ResGm; + this->attentionOutGm = attentionOutGm; +} + +template +__aicore__ inline void +QSFAMatmulService::InitPageAttentionInfo(const GlobalTensor& kvMergeGm, + GlobalTensor blockTableGm, GlobalTensor topKGm, + uint32_t blockSize, uint32_t maxBlockNumPerBatch) +{ + this->blockTableGm = blockTableGm; + this->topKGm = topKGm; + this->kvCacheBlockSize = blockSize; + this->maxBlockNumPerBatch = maxBlockNumPerBatch; + this->kvMergeGm_ = kvMergeGm; +} + +template __aicore__ inline void QSFAMatmulService::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(bufQPL1, L1_BLOCK_SIZE * 4); // (64K + 8K) * 4 + l1QPTensor = bufQPL1.Get(); + pipe->InitBuffer(bufKVL1, L1_BLOCK_SIZE * 3); // (64K + 8K) * 3 + l1KVTensor = bufKVL1.Get(); + + // L0A + pipe->InitBuffer(tmpBufL0A, L0A_PP_SIZE * 2); // 64K + aL0TensorPingPong = tmpBufL0A.Get(); + // L0B + pipe->InitBuffer(tmpBufL0B, L0B_PP_SIZE * 2); // 64K + bL0TensorPingPong = tmpBufL0B.Get(); + // L0C + pipe->InitBuffer(tmpBufL0C, L0C_PP_SIZE * 2); // 128K + cL0TensorPingPong = tmpBufL0C.Get(); +} + +template __aicore__ inline void QSFAMatmulService::UpdateKey(GlobalTensor keyGm) +{ + this->keyGm = keyGm; +} + +template __aicore__ inline void QSFAMatmulService::UpdateValue(GlobalTensor valueGm) +{ + this->valueGm = valueGm; +} + +template __aicore__ inline void QSFAMatmulService::AllocEventID() +{ + SetFlag(L0AB_EVENT0); + SetFlag(L0AB_EVENT1); + SetFlag(L1_EVENT0); + SetFlag(L1_EVENT1); + SetFlag(L1_EVENT2); + SetFlag(L1_EVENT3); + SetFlag(L1_EVENT4); + SetFlag(L1_EVENT5); + SetFlag(L1_EVENT6); +} + +template __aicore__ inline void QSFAMatmulService::FreeEventID() +{ + WaitFlag(L0AB_EVENT0); + WaitFlag(L0AB_EVENT1); + WaitFlag(L1_EVENT0); + WaitFlag(L1_EVENT1); + WaitFlag(L1_EVENT2); + WaitFlag(L1_EVENT3); + WaitFlag(L1_EVENT4); + WaitFlag(L1_EVENT5); + WaitFlag(L1_EVENT6); +} + +template +__aicore__ inline void QSFAMatmulService::CopyGmToL1(LocalTensor &l1Tensor, + GlobalTensor &gmSrcTensor, uint32_t srcN, + uint32_t srcD, uint32_t srcDstride) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.dValue = srcD; + nd2nzPara.nValue = srcN; // 行数 + nd2nzPara.srcDValue = srcDstride; + nd2nzPara.dstNzC0Stride = (srcN + 15) / 16 * 16; // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.dstNzMatrixStride = 0; + nd2nzPara.srcNdMatrixStride = 0; + DataCopy(l1Tensor, gmSrcTensor, nd2nzPara); +} + +template +__aicore__ inline void QSFAMatmulService::CopyInMm1AToL1(LocalTensor &l1Tensor, const RunInfo &info, + uint32_t mSeqIdx, uint32_t mSizeAct, + uint32_t headSize, uint32_t headOffset) +{ + auto srcGm = queryGm[info.tensorAOffset + mSeqIdx * constInfo.combineHeadDim + headOffset]; + CopyGmToL1(l1Tensor, srcGm, mSizeAct, headSize, headSize); +} + +template +__aicore__ inline void QSFAMatmulService::CopyInMm1ARopeToL1(LocalTensor &l1Tensor, + const RunInfo &info, uint32_t mSeqIdx, + uint32_t mSizeAct) +{ + auto srcGm = qRopeGm[info.tensorARopeOffset + mSeqIdx * constInfo.headDimRope]; + CopyGmToL1(l1Tensor, srcGm, mSizeAct, constInfo.headDimRope, constInfo.headDimRope); +} + +template +__aicore__ inline void +QSFAMatmulService::CopyInMm1BToL1(LocalTensor &bL1Tensor, const uint64_t keyGmBaseOffset, + uint32_t copyTotalRowCntAlign, uint32_t copyStartRowCnt, + uint32_t nActCopyRowCount, uint32_t headSize) +{ + uint64_t dStride = constInfo.headDim; + if constexpr (LAYOUT_T == QSFA_LAYOUT::BSND || LAYOUT_T == QSFA_LAYOUT::TND) { + dStride = constInfo.headDim * constInfo.kvHeadNum; + } + + uint32_t blockElementCnt = 32 / sizeof(K_ROPE_T); + + Nd2NzParams mm1Nd2NzParamsForB; + mm1Nd2NzParamsForB.ndNum = 1; + mm1Nd2NzParamsForB.nValue = nActCopyRowCount; + mm1Nd2NzParamsForB.dValue = headSize; + mm1Nd2NzParamsForB.srcDValue = dStride; + mm1Nd2NzParamsForB.dstNzNStride = 1; + mm1Nd2NzParamsForB.dstNzC0Stride = copyTotalRowCntAlign; + mm1Nd2NzParamsForB.srcNdMatrixStride = 0; + mm1Nd2NzParamsForB.dstNzMatrixStride = 0; + DataCopy(bL1Tensor[copyStartRowCnt * blockElementCnt], keyGm[keyGmBaseOffset], mm1Nd2NzParamsForB); +} + +template +__aicore__ inline void +QSFAMatmulService::CopyInMm1BRopeToL1(LocalTensor &bL1Tensor, const uint64_t kRopeGmBaseOffset, + uint32_t copyTotalRowCntAlign, uint32_t copyStartRowCnt, + uint32_t nActCopyRowCount, uint32_t headSize) +{ + uint64_t dStride = constInfo.headDimRope; + if constexpr (LAYOUT_T == QSFA_LAYOUT::BSND || LAYOUT_T == QSFA_LAYOUT::TND) { + dStride = constInfo.headDimRope * constInfo.kvHeadNum; + } + + uint32_t blockElementCnt = 32 / sizeof(K_ROPE_T); + + Nd2NzParams mm1Nd2NzParamsForB; + mm1Nd2NzParamsForB.nValue = nActCopyRowCount; + mm1Nd2NzParamsForB.dValue = headSize; + mm1Nd2NzParamsForB.ndNum = 1; + mm1Nd2NzParamsForB.srcDValue = dStride; + mm1Nd2NzParamsForB.srcNdMatrixStride = 0; + mm1Nd2NzParamsForB.dstNzMatrixStride = 0; + mm1Nd2NzParamsForB.dstNzNStride = 1; + mm1Nd2NzParamsForB.dstNzC0Stride = copyTotalRowCntAlign; + DataCopy(bL1Tensor[copyStartRowCnt * blockElementCnt], kRopeGm[kRopeGmBaseOffset], mm1Nd2NzParamsForB); +} + +template +__aicore__ inline void QSFAMatmulService::LoadDataMm1A(LocalTensor &aL0Tensor, + LocalTensor &aL1Tensor, uint32_t idx, + uint32_t kSplitSize, uint32_t mSize, uint32_t kSize) +{ + LocalTensor srcTensor = aL1Tensor[mSize * kSplitSize * idx]; + LoadData3DParamsV2 loadData3DParams; + // SetFmatrixParams + loadData3DParams.l1H = mSize / 16; // Hin=M1=8 + loadData3DParams.l1W = 16; // Win=M0 + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + // SetLoadToA0Params + loadData3DParams.mExtension = mSize; // M + loadData3DParams.kExtension = kSize; // K + loadData3DParams.mStartPt = 0; + loadData3DParams.kStartPt = 0; + loadData3DParams.strideH = 1; + loadData3DParams.strideW = 1; + loadData3DParams.filterW = 1; + loadData3DParams.filterSizeW = (1 >> 8) & 255; + loadData3DParams.filterH = 1; + loadData3DParams.filterSizeH = (1 >> 8) & 255; + loadData3DParams.dilationFilterH = 1; + loadData3DParams.dilationFilterW = 1; + loadData3DParams.fMatrixCtrl = 0; + loadData3DParams.channelSize = kSize; // Cin=K + loadData3DParams.enTranspose = 0; + LoadData(aL0Tensor, srcTensor, loadData3DParams); +} + +template +__aicore__ inline void QSFAMatmulService::LoadDataMm1B(LocalTensor &l0Tensor, + LocalTensor &l1Tensor, uint32_t idx, + uint32_t kSplitSize, uint32_t kSize, uint32_t nSize) +{ + // N 方向全载 + LocalTensor srcTensor = l1Tensor[nSize * kSplitSize * idx]; + + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = (nSize + 15) / 16 * kSize / (32 / sizeof(K_ROPE_T)); + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = false; + LoadData(l0Tensor, srcTensor, loadData2DParams); +} + +template +__aicore__ inline void QSFAMatmulService::CopyInMm2AToL1(LocalTensor &aL1Tensor, const RunInfo &info, + uint32_t mSeqIdx, uint32_t subMSizeAct, + uint32_t nSize, uint32_t nOffset) +{ + auto srcGm = vec1ResGm[(info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize + + mSeqIdx * info.actualSingleProcessSInnerSizeAlign + nOffset]; + CopyGmToL1(aL1Tensor, srcGm, subMSizeAct, nSize, info.actualSingleProcessSInnerSizeAlign); +} + +template +__aicore__ inline void QSFAMatmulService::CopyInMm2BToL1( + LocalTensor &bL1Tensor, const uint64_t valueGmBaseOffset, uint32_t copyTotalRowCntAlign, + uint32_t copyStartRowCnt, uint32_t nActCopyRowCount, uint32_t copyStartColumnCount, uint32_t copyColumnCount) +{ + uint64_t step = constInfo.headDim; + if constexpr (LAYOUT_T == QSFA_LAYOUT::BSND || LAYOUT_T == QSFA_LAYOUT::TND) { + step = constInfo.headDim * constInfo.kvHeadNum; + } + + uint32_t blockElementCnt = 32 / sizeof(K_ROPE_T); + + Nd2NzParams mm1Nd2NzParamsForB; + mm1Nd2NzParamsForB.ndNum = 1; + mm1Nd2NzParamsForB.nValue = nActCopyRowCount; + mm1Nd2NzParamsForB.dValue = copyColumnCount; + mm1Nd2NzParamsForB.srcDValue = step; + mm1Nd2NzParamsForB.dstNzNStride = 1; + mm1Nd2NzParamsForB.dstNzC0Stride = copyTotalRowCntAlign; + mm1Nd2NzParamsForB.srcNdMatrixStride = 0; + mm1Nd2NzParamsForB.dstNzMatrixStride = 0; + DataCopy(bL1Tensor[copyStartRowCnt * blockElementCnt], valueGm[valueGmBaseOffset + copyStartColumnCount], + mm1Nd2NzParamsForB); +} + +template +__aicore__ inline void QSFAMatmulService::CalcTopKBlockInfo( + const RunInfo &info, uint32_t &curTopKIdx, uint64_t &curOffsetInSparseBlock, + uint32_t curSeqIdx, uint32_t ©RowCnt, uint64_t &idInTopK) +{ + if (curTopKIdx == 0 && curOffsetInSparseBlock == 0 && copyRowCnt == 0) { + uint64_t sparseLen = 0; + for (uint64_t qsfaTopkidx = 0; qsfaTopkidx < constInfo.sparseBlockCount; qsfaTopkidx++) { + int32_t qsfaSparseIndices = topKGm.GetValue(info.topKBaseOffset + qsfaTopkidx); + if (qsfaSparseIndices == -1) { + break; + } + uint64_t qsfaBlockBegin = qsfaSparseIndices * constInfo.sparseBlockSize; + if (qsfaBlockBegin >= info.threshold) { + continue; + } + uint64_t qsfaBlockEnd = (qsfaBlockBegin + constInfo.sparseBlockSize > info.curActualSeqLenOri) ? + info.curActualSeqLenOri : qsfaBlockBegin + constInfo.sparseBlockSize; + uint64_t qsfaBlockLen = (qsfaBlockEnd <= info.threshold) ? \ + qsfaBlockEnd - qsfaBlockBegin : info.threshold - qsfaBlockBegin; + sparseLen += qsfaBlockLen; + if (sparseLen >= curSeqIdx + 1) { + curTopKIdx = qsfaTopkidx; + idInTopK = qsfaSparseIndices; + curOffsetInSparseBlock = qsfaBlockLen - (sparseLen - curSeqIdx); + copyRowCnt = sparseLen - curSeqIdx; + break; + } + } + return; + } + uint64_t qsfaBlockBegin = idInTopK * constInfo.sparseBlockSize; + uint64_t qsfaBlockEnd = (qsfaBlockBegin + constInfo.sparseBlockSize > info.threshold) ? + info.threshold : qsfaBlockBegin + constInfo.sparseBlockSize; + uint64_t qsfaBlockLen = qsfaBlockEnd - qsfaBlockBegin; + if (curOffsetInSparseBlock + copyRowCnt < qsfaBlockLen) { + curOffsetInSparseBlock += copyRowCnt; + copyRowCnt = qsfaBlockLen - curOffsetInSparseBlock; + } else { + for (uint64_t qsfaTopkidx = curTopKIdx + 1; qsfaTopkidx < constInfo.sparseBlockCount; qsfaTopkidx++) { + int64_t qsfaSparseIndices = topKGm.GetValue(info.topKBaseOffset + qsfaTopkidx); + if (qsfaSparseIndices == -1) { + break; + } + + uint64_t qsfaBlockBegin = qsfaSparseIndices * constInfo.sparseBlockSize; + if (qsfaBlockBegin >= info.threshold) { + continue; + } + uint64_t qsfaBlockEnd = (qsfaBlockBegin + constInfo.sparseBlockSize > info.threshold) ? + info.threshold : qsfaBlockBegin + constInfo.sparseBlockSize; + uint64_t qsfaBlockLen = qsfaBlockEnd - qsfaBlockBegin; + curTopKIdx = qsfaTopkidx; + idInTopK = qsfaSparseIndices; + curOffsetInSparseBlock = 0; + copyRowCnt = qsfaBlockLen; + break; + } + } +} + +template +__aicore__ inline void QSFAMatmulService::ComputeMm1(const RunInfo &info, const MSplitInfo mSplitInfo) +{ + // 最外层还需要一层m的循环 + uint32_t mSize = mSplitInfo.nBufferDealM; + uint32_t mL1Size = M_SPLIT_SIZE; + uint32_t mL1SizeAlign = QSFAAlign(M_SPLIT_SIZE, 16U); + uint32_t mL1Loops = (mSize + M_SPLIT_SIZE - 1) / M_SPLIT_SIZE; + + uint32_t nSize = info.actualSingleProcessSInnerSize; + uint32_t nL1Size = N_SPLIT_SIZE; + uint32_t nL1SizeAlign = QSFAAlign(N_SPLIT_SIZE, 16U); + uint32_t nL1Loops = (nSize + N_SPLIT_SIZE - 1) / N_SPLIT_SIZE; + + uint32_t kSize = 576; + uint32_t kL1Size = 288; + uint32_t kL1Loops = 2; // 2 : 576/288, mla专用 这里不考虑d泛化 + + uint32_t kL0Size = 96; + uint32_t kL0Loops = (kL1Size + kL0Size - 1) / kL0Size; // 288 / 96 = 3 kloops + + // ka表示左矩阵4buf选择哪一块buf, kb表示右矩阵3buf选择哪一块buf + uint32_t ka = 0, kb = 0; + for (uint32_t mL1 = 0; mL1 < mL1Loops; mL1++) { + mL1Size = M_SPLIT_SIZE; + mL1SizeAlign = QSFAAlign(M_SPLIT_SIZE, 16U); + if (mL1 == (mL1Loops - 1)) { + // 尾块重新计算size + mL1Size = mSize - (mL1Loops - 1) * M_SPLIT_SIZE; + mL1SizeAlign = QSFAAlign(mL1Size, 16U); + } + + // 左矩阵L1选择12块还是34块的index, 由m l1 index决定 + // 左矩阵L1选择12块或34块的前一块还是后一块, 由k l1 index决定 + uint32_t mIdx = qpL1BufIter + mL1; + ka = GetQPL1RealIdx(mIdx, 0); + LocalTensor aL1Tensor = l1QPTensor[ka * L1_BLOCK_OFFSET]; + WaitFlag(mte21QPIds[ka]); + CopyInMm1AToL1(aL1Tensor, info, mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE, mL1Size, 576, 0); + SetFlag(mte21QPIds[ka]); + WaitFlag(mte21QPIds[ka]); + for (uint32_t nL1 = 0; nL1 < nL1Loops; nL1++) { // L1切n, 512/128=4 + if (nL1 == (nL1Loops - 1)) { + // 尾块重新计算size + nL1Size = nSize - (nL1Loops - 1) * N_SPLIT_SIZE; + nL1SizeAlign = QSFAAlign(nL1Size, 16U); + } + + // 使用unitflag同步 + // 需要保证cL0BufIter和m步调一致 + LocalTensor cL0Tensor = cL0TensorPingPong[(cL0BufIter % 2) * (L0C_PP_SIZE / sizeof(MM_OUT_T))]; + for (uint32_t kL1 = 0; kL1 < kL1Loops; kL1++) { // L1切k, 576/288, 这里不考虑d泛化 + kvL1BufIter++; + uint32_t kb = kvL1BufIter % 3; + WaitFlag(mte21KVIds[kb]); + // 从k当中取当前的块 + LocalTensor bL1Tensor = l1KVTensor[kb * L1_BLOCK_OFFSET]; + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + if (kL1 == 0) { + DataCopyParams copyParams; + copyParams.blockCount = 288 / BLOCK_ELEMENT_NUM; + copyParams.blockLen = nL1Size; + copyParams.srcStride = constInfo.s2BaseSize - nL1Size; + copyParams.dstStride = nL1SizeAlign - nL1Size; + DataCopy(bL1Tensor, kvMergeGm_[info.loop % 4 * N_WORKSPACE_SIZE * kSize + + nL1 * N_SPLIT_SIZE * BLOCK_ELEMENT_NUM], copyParams); + } else { + DataCopyParams copyParams; + copyParams.blockCount = 224 / BLOCK_ELEMENT_NUM; + copyParams.blockLen = nL1Size; + copyParams.srcStride = constInfo.s2BaseSize - nL1Size; + copyParams.dstStride = nL1SizeAlign - nL1Size; + DataCopy(bL1Tensor, kvMergeGm_[info.loop % 4 * N_WORKSPACE_SIZE * kSize + + 288 * constInfo.s2BaseSize + nL1 * N_SPLIT_SIZE * BLOCK_ELEMENT_NUM], copyParams); + copyParams.blockCount = constInfo.headDimRope / BLOCK_ELEMENT_NUM; + DataCopy( + bL1Tensor[224 * nL1SizeAlign], + kvMergeGm_[info.loop % 4 * N_WORKSPACE_SIZE * kSize + N_WORKSPACE_SIZE * constInfo.headDim + + nL1 * N_SPLIT_SIZE * BLOCK_ELEMENT_NUM], + copyParams); + } + } + SetFlag(mte21KVIds[kb]); + WaitFlag(mte21KVIds[kb]); + + aL1Tensor = l1QPTensor[ka * L1_BLOCK_OFFSET + kL1 * mL1SizeAlign * K_SPLIT_SIZE]; + for (uint32_t kL0 = 0; kL0 < kL0Loops; kL0++) { + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + LocalTensor aL0Tensor = aL0TensorPingPong[(abL0BufIter % 2) * (L0A_PP_SIZE / + sizeof(K_ROPE_T))]; + LoadDataMm1A(aL0Tensor, aL1Tensor, kL0, kL0Size, mL1SizeAlign, kL0Size); + LocalTensor bL0Tensor = bL0TensorPingPong[(abL0BufIter % 2) * (L0B_PP_SIZE / + sizeof(K_ROPE_T))]; + LoadDataMm1B(bL0Tensor, bL1Tensor, kL0, kL0Size, kL0Size, nL1SizeAlign); + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + + // m == 1的时候需要特殊处理 + MmadParams mmadParams; + mmadParams.m = mL1SizeAlign; + mmadParams.n = nL1SizeAlign; + mmadParams.k = kL0Size; + mmadParams.cmatrixSource = false; + mmadParams.cmatrixInitVal = (kL1 == 0 && kL0 == 0); + mmadParams.unitFlag = + (kL1 == 1 && kL0 == (kL0Loops - 1)) ? 0b11 : 0b10; // 累加最后一次翻转flag, 表示可以搬出 + Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams); + + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + abL0BufIter++; + } + SetFlag(mte21KVIds[kb]); // 反向同步, 表示L1已经被mte1消费完 + } + FixpipeParamsV220 fixParams; + fixParams.mSize = mL1SizeAlign; + fixParams.nSize = nL1SizeAlign; + fixParams.srcStride = mL1SizeAlign; + fixParams.ndNum = 1; // 输出ND + // 改成nSizeAlign + fixParams.dstStride = info.actualSingleProcessSInnerSizeAlign; // mm1ResGm两行之间的间隔 + fixParams.unitFlag = 0b11; + + // 输出偏移info.loop % (constInfo.preLoadNum)) * mmResUbSize是否在matmul里计算 + Fixpipe(mm1ResGm[(info.loop % (constInfo.preLoadNum)) * constInfo.mmResUbSize + nL1 * N_SPLIT_SIZE + + (mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE) * + info.actualSingleProcessSInnerSizeAlign], + cL0Tensor, fixParams); + cL0BufIter++; + } + SetFlag(mte21QPIds[ka]); // 反向同步, 表示L1中的A已经被mte1消费完 + } + qpL1BufIter += mL1Loops; +} + +template +__aicore__ inline void QSFAMatmulService::ComputeMm2(const RunInfo &info, const MSplitInfo mSplitInfo) +{ + uint32_t mSize = mSplitInfo.nBufferDealM; + uint32_t mSizeAlign = (mSize + 16 - 1) / 16; + uint32_t mL1Loops = (mSize + M_SPLIT_SIZE - 1) / M_SPLIT_SIZE; + uint32_t mL1SizeAlign = M_SPLIT_SIZE; // 16对齐 + uint32_t mL1Size = M_SPLIT_SIZE; // m的实际大小 + + uint32_t nSize = BlockAlign(constInfo.headDim); + uint32_t nL1Loops = (nSize + N_SPLIT_SIZE - 1) / N_SPLIT_SIZE; + uint32_t nL1SizeAlign = N_SPLIT_SIZE; // 16对齐 + uint32_t nL1Size = N_SPLIT_SIZE; // n的实际大小 + + uint32_t kSize = info.actualSingleProcessSInnerSize; + uint32_t kL1Size = 256; + uint32_t kL1SizeAlign = QSFAAlign(kL1Size, 16U); + uint32_t kL1Loops = (kSize + kL1Size - 1) / kL1Size; + uint32_t kL0Size = 128; + uint32_t kL0Loops = (kL1Size + kL0Size - 1) / kL0Size; + uint32_t kL0SizeAlign = kL0Size; + LocalTensor bL1Tensor; + LocalTensor subvTensor; + + // ka表示左矩阵4buf选择哪一块buf, kb表示右矩阵3buf选择哪一块buf + uint32_t ka = 0, qsfaKb = 0; + uint32_t mBaseIdx = qpL1BufIter; + for (uint32_t nL1 = 0; nL1 < nL1Loops; nL1++) { // n切L1 + if (nL1 == (nL1Loops - 1)) { + // 尾块 + nL1Size = nSize - (nL1Loops - 1) * N_SPLIT_SIZE; + nL1SizeAlign = QSFAAlign(nL1Size, 16U); + } + + // k l1写成一个循环, 和mm1保持一致 + kL1Size = 256; + kL1SizeAlign = QSFAAlign(kL1Size, 16U); + for (uint32_t k1 = 0; k1 < kL1Loops; k1++) { // k切L1, 这里套了一层l0来操作 + if (k1 == (kL1Loops - 1)) { + // 尾块 + kL1Size = kSize - (kL1Loops - 1) * 256; + kL1SizeAlign = QSFAAlign(kL1Size, 16U); + } + kvL1BufIter++; + uint32_t qsfaKb = kvL1BufIter % 3; + WaitFlag(mte21KVIds[qsfaKb]); + bL1Tensor = l1KVTensor[qsfaKb * L1_BLOCK_OFFSET]; + uint32_t qsfaKOffset = k1 * kL0Loops; + kL0Size = 128; + // 此处必须先初始化kL0Size, 再求kL0Loops, 否则由于循环会改变kL0Size大小, 导致kL0Loops错误 + kL0Loops = (kL1Size + kL0Size - 1) / kL0Size; + kL0SizeAlign = kL0Size; + for (uint32_t qsfaKL1 = qsfaKOffset; qsfaKL1 < kL0Loops + qsfaKOffset; qsfaKL1++) { // 128 循环搬pa + if (qsfaKL1 == qsfaKOffset + kL0Loops - 1) { + // 尾块 + kL0Size = kL1Size - (kL0Loops - 1) * kL0Size; + kL0SizeAlign = QSFAAlign(kL0Size, 16U); + } + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + DataCopyParams copyParams; + copyParams.blockLen = kL0Size; + copyParams.blockCount = nL1Size / BLOCK_ELEMENT_NUM; + copyParams.srcStride = constInfo.s2BaseSize - kL0Size; + copyParams.dstStride = kL0SizeAlign - kL0Size; + DataCopy(bL1Tensor[(qsfaKL1 - qsfaKOffset) * 128 * N_SPLIT_SIZE], kvMergeGm_[info.loop % 4 * + N_WORKSPACE_SIZE * 576 + qsfaKL1 * 128 * BLOCK_ELEMENT_NUM + nL1 * N_SPLIT_SIZE * + constInfo.s2BaseSize], copyParams); + } + } + SetFlag(mte21KVIds[qsfaKb]); + WaitFlag(mte21KVIds[qsfaKb]); + mL1SizeAlign = M_SPLIT_SIZE; + mL1Size = M_SPLIT_SIZE; // m的实际大小 + for (uint32_t qsfaML1 = 0; qsfaML1 < mL1Loops; qsfaML1++) { + if (qsfaML1 == (mL1Loops - 1)) { + // 尾块 + mL1Size = mSize - (mL1Loops - 1) * M_SPLIT_SIZE; + mL1SizeAlign = QSFAAlign(mL1Size, 16U); + } + + uint32_t mIdx = mBaseIdx + qsfaML1; + ka = GetQPL1RealIdx(mIdx, k1); + LocalTensor aL1Tensor = l1QPTensor[ka * L1_BLOCK_OFFSET]; + if (nL1 == 0) { + WaitFlag(mte21QPIds[ka]); + CopyInMm2AToL1(aL1Tensor, info, mSplitInfo.nBufferStartM + qsfaML1 * M_SPLIT_SIZE, mL1Size, kL1Size, + 256 * k1); + SetFlag(mte21QPIds[ka]); + WaitFlag(mte21QPIds[ka]); + } + + LocalTensor cL0Tensor = + cL0TensorPingPong[(cL0BufIter % 2) * + (L0C_PP_SIZE / sizeof(MM_OUT_T))]; // 需要保证cL0BufIter和m步调一致 + uint32_t qsfaBaseK = 128; + uint32_t qsfaBaseN = 128; + kL0Size = 128; + kL0SizeAlign = kL0Size; + for (uint32_t qsfaKL0 = 0; qsfaKL0 < kL0Loops; qsfaKL0++) { + if (qsfaKL0 + 1 == kL0Loops) { + kL0Size = kL1Size - (kL0Loops - 1) * kL0Size; + kL0SizeAlign = QSFAAlign(kL0Size, 16U); + } + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + LocalTensor bL0Tensor = bL0TensorPingPong[(abL0BufIter % 2) * (L0B_PP_SIZE / + sizeof(K_ROPE_T))]; + LoadData3DParamsV2 loadData3DParamsForB; + loadData3DParamsForB.l1H = kL0SizeAlign / 16; // 源操作数height + loadData3DParamsForB.l1W = 16; // 源操作数weight=16，目的height=l1H*L1W + loadData3DParamsForB.padList[0] = 0; + loadData3DParamsForB.padList[1] = 0; + loadData3DParamsForB.padList[2] = 0; + loadData3DParamsForB.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + loadData3DParamsForB.mExtension = kL0SizeAlign; // 在目的操作数height维度的传输长度 + loadData3DParamsForB.kExtension = nL1SizeAlign; // 在目的操作数width维度的传输长度 + loadData3DParamsForB.mStartPt = 0; // 卷积核在目的操作数width维度的起点 + loadData3DParamsForB.kStartPt = 0; // 卷积核在目的操作数height维度的起点 + loadData3DParamsForB.strideH = 1; + loadData3DParamsForB.strideW = 1; + loadData3DParamsForB.filterW = 1; + loadData3DParamsForB.filterSizeW = false; // 是否在filterW的基础上将卷积核width增加256个元素 + loadData3DParamsForB.filterH = 1; + loadData3DParamsForB.filterSizeH = false; // 是否在filterH的基础上将卷积核height增加256个元素 + loadData3DParamsForB.dilationFilterH = 1; // 卷积核height膨胀系数 + loadData3DParamsForB.dilationFilterW = 1; // 卷积核width膨胀系数 + loadData3DParamsForB.enTranspose = 1; // 是否启用转置功能 + // 使用FMATRIX_LEFT还是使用FMATRIX_RIGHT，=0使用FMATRIX_LEFT，=1使用FMATRIX_RIGHT 1 + loadData3DParamsForB.fMatrixCtrl = 0; + // 源操作数的通道数。膨胀系数为1时，目的weight为filterW*filterH*channelSize + loadData3DParamsForB.channelSize = nL1SizeAlign; + LoadData(bL0Tensor, bL1Tensor[qsfaKL0 * qsfaBaseK * qsfaBaseN], + loadData3DParamsForB); + LocalTensor aL0Tensor = aL0TensorPingPong[(abL0BufIter % 2) * (L0A_PP_SIZE / + sizeof(K_ROPE_T))]; + LoadData3DParamsV2 loadData3DParamsForA; + loadData3DParamsForA.l1H = mL1SizeAlign / 16; // 源操作数height + loadData3DParamsForA.l1W = 16; // 源操作数weight + loadData3DParamsForA.padList[0] = 0; + loadData3DParamsForA.padList[1] = 0; + loadData3DParamsForA.padList[2] = 0; + loadData3DParamsForA.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + loadData3DParamsForA.mExtension = mL1SizeAlign; // 在目的操作数height维度的传输长度 + loadData3DParamsForA.kExtension = kL0SizeAlign; // 在目的操作数width维度的传输长度 + loadData3DParamsForA.mStartPt = 0; // 卷积核在目的操作数width维度的起点 + loadData3DParamsForA.kStartPt = 0; // 卷积核在目的操作数height维度的起点 + loadData3DParamsForA.strideW = 1; // 卷积核在源操作数width维度滑动的步长 + loadData3DParamsForA.strideH = 1; // 卷积核在源操作数height维度滑动的步长 + loadData3DParamsForA.filterW = 1; // 卷积核width + loadData3DParamsForA.filterSizeW = false; // 是否在filterW的基础上将卷积核width增加256个元素 + loadData3DParamsForA.filterH = 1; // 卷积核height + loadData3DParamsForA.filterSizeH = false; // 是否在filterH的基础上将卷积核height增加256个元素 + loadData3DParamsForA.dilationFilterW = 1; // 卷积核width膨胀系数 + loadData3DParamsForA.dilationFilterH = 1; // 卷积核height膨胀系数 + loadData3DParamsForA.enTranspose = 0; // 是否启用转置功能，对整个目标矩阵进行转置 + loadData3DParamsForA.fMatrixCtrl = 0; + // 源操作数的通道数。膨胀系数为1时，目的weight为filterW*filterH*channelSize + loadData3DParamsForA.channelSize = kL0SizeAlign; + LoadData(aL0Tensor, aL1Tensor[qsfaKL0 * qsfaBaseK * mL1SizeAlign], + loadData3DParamsForA); + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + + MmadParams mmadParams; + mmadParams.m = mL1SizeAlign; + mmadParams.n = nL1SizeAlign; + mmadParams.k = kL0Size; + mmadParams.cmatrixInitVal = (qsfaKL0 == 0 && k1 == 0); + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = ((k1 == (kL1Loops - 1)) && (qsfaKL0 == (kL0Loops - 1))) ? 0b11 : 0b10; + + Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + abL0BufIter++; + } + + if (nL1 == (nL1Loops - 1)) { // nL1最后一轮, 需要将B驻留在L1中, 用于下一轮的计算？ + SetFlag(mte21QPIds[ka]); // 反向同步, 表示L1中的A已经被mte1消费完 + } + + if (k1 == (kL1Loops - 1)) { + // ND + FixpipeParamsV220 fixParams; + fixParams.nSize = nL1SizeAlign; + fixParams.mSize = mL1SizeAlign; + fixParams.srcStride = mL1SizeAlign; + fixParams.dstStride = nSize; // mm2ResGm两行之间的间隔 + fixParams.ndNum = 1; // 输出ND + fixParams.unitFlag = 0b11; + + uint64_t qsfaMm2Offset = (mSplitInfo.nBufferStartM + qsfaML1 * M_SPLIT_SIZE) * nSize + + nL1 * N_SPLIT_SIZE; + Fixpipe(mm2ResGm[(info.loop % (constInfo.preLoadNum)) * + constInfo.bmm2ResUbSize + qsfaMm2Offset], cL0Tensor, fixParams); + } + + if (mL1Loops == 2) { + cL0BufIter++; + } + } + SetFlag(mte21KVIds[qsfaKb]); // 反向同步, 表示L1已经被mte1消费完 + } + // cL0BufIter已经不在使用 + if (mL1Loops == 1) { + cL0BufIter++; + } + } + qpL1BufIter += mL1Loops; +} + +#endif // KV_QUANT_SPARSE_FLASH_ATTENTION_SERVICE_CUBE_MLA_H diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_service_vector_mla.h b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_service_vector_mla.h new file mode 100644 index 000000000..4185f5531 --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_service_vector_mla.h @@ -0,0 +1,1367 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_service_vector_mla.h + * \brief + */ +#ifndef KV_QUANT_SPARSE_FLASH_ATTENTION_SERVICE_VECTOR_MLA_H +#define KV_QUANT_SPARSE_FLASH_ATTENTION_SERVICE_VECTOR_MLA_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "kv_quant_sparse_flash_attention_common.h" + +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +template class QSFAVectorService { +public: + // 中间计算数据类型为float，高精度模式 + using T = float; + using KV_T = typename QSFAT::kvType; + using K_ROPE_T = typename QSFAT::kRopeType; + using OUT_T = typename QSFAT::outputType; + using UPDATE_T = T; + using MM1_OUT_T = float; + using MM2_OUT_T = float; + bool NO_AMLA = true; + + __aicore__ inline QSFAVectorService(){}; + __aicore__ inline void ProcessVec1L(const RunInfo &info); + __aicore__ inline void ProcessVec2L(const RunInfo &info); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitParams(const struct ConstInfo &constInfo, + const KvQuantSparseFlashAttentionTilingDataMla *__restrict tilingData); + __aicore__ inline void InitMm2ResInt32GmGlobalTensor(GlobalTensor mm2ResInt32Gm); + __aicore__ inline void InitVec0GlobalTensor(const GlobalTensor &kvValidSizeGm, + const GlobalTensor &kvMergeGm, + const GlobalTensor &keyRopeGm, + const GlobalTensor &keyGm, + const GlobalTensor &blkTableGm); + __aicore__ inline void InitVec1GlobalTensor(GlobalTensor mm1ResGm, GlobalTensor vec1ResGm, + GlobalTensor actualSeqLengthsQGm, + GlobalTensor actualSeqLengthsKVGm, GlobalTensor lseMaxFdGm, + GlobalTensor lseSumFdGm, GlobalTensor topKGm, + GlobalTensor softmaxMaxGm, GlobalTensor softmaxSumGm); + __aicore__ inline void InitVec2GlobalTensor(GlobalTensor accumOutGm, GlobalTensor vec2ResGm, + GlobalTensor mm2ResGm, GlobalTensor attentionOutGm); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void InitSoftmaxDefaultBuffer(); + // ================================Base Vector========================================== + __aicore__ inline void RowDivs(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline void RowMuls(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + // ================================Vector0========================================== + __aicore__ inline void MergeKv(const RunInfo &runInfo); + __aicore__ inline int64_t GetKeyBNBOffset(int64_t realS2Idx, const RunInfo &runInfo, int64_t s2IdLimit); + __aicore__ inline void GetRealS2Idx(int64_t s2GmOffset, int64_t &realS2Idx, int64_t topkGmBaseOffset, + const RunInfo &runInfo); + __aicore__ inline void SetInfInBlk(const LocalTensor &mmResUb, uint32_t dealRowCount, uint32_t columnCount, + uint64_t startId, uint64_t endId); + __aicore__ inline void SetMidInf(const LocalTensor &mmResUb, uint32_t dealRowCount, uint32_t columnCount, + uint64_t startId, uint64_t endId); + __aicore__ inline void CopyInKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, int64_t realS2Idx1, + int64_t realS2Idx2, const RunInfo &runInfo); + __aicore__ inline void CopyOutMrgeResult(int64_t mte2Size, int64_t mte3Size, int64_t s2StartGmOffset, + int64_t mergeMte3Idx, const RunInfo &runInfo); + __aicore__ inline void CopyInSingleKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, int64_t realS2Idx, + int64_t keyBNBOffset, int64_t s2IdLimit, const RunInfo &runInfo); + // ================================Vector1========================================== + __aicore__ inline void ProcessVec1SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo); + __aicore__ inline void DealBmm1ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t startRow, + uint32_t dealRowCount, uint32_t columnCount, uint32_t loopId); + __aicore__ inline void SoftmaxFlashV2Compute(const RunInfo &info, const MSplitInfo &mSplitInfo, + LocalTensor &mmResUb, LocalTensor &softmaxTmpUb, + uint32_t startRow, uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + __aicore__ inline void ElewiseCompute(const RunInfo &info, const LocalTensor &mmResUb, uint32_t dealRowCount, + uint32_t columnCount); + __aicore__ inline void ComputeLogSumExpAndCopyToGm(const RunInfo &info, const MSplitInfo &mSplitInfo, + LocalTensor &softmaxSumUb, LocalTensor &softmaxMaxUb); + __aicore__ inline void CopyFALseToGm(const RunInfo &info, const MSplitInfo &mSplitInfo, + LocalTensor &softmaxSumUb, LocalTensor &softmaxMaxUb); + // ================================Vecotr2========================================== + __aicore__ inline void ProcessVec2SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo); + __aicore__ inline void DealBmm2ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t startRow, + uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + __aicore__ inline void ProcessVec2Inner(const RunInfo &info, const MSplitInfo &mSplitInfo, + uint32_t mStartRow, uint32_t mDealSize); + __aicore__ inline void Bmm2DataCopyOutTrans(const RunInfo &info, LocalTensor &attenOutUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + __aicore__ inline void Bmm2ResCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + __aicore__ inline void Bmm2CastAndCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline void Bmm2FDDataCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + __aicore__ inline uint64_t CalcAccumOffset(uint32_t bN2Idx, uint32_t gS1Idx); + __aicore__ inline void GetConfusionTransposeTiling(int64_t numR, int64_t numC, const uint32_t stackBufferSize, + const uint32_t typeSize, ConfusionTransposeTiling &tiling); + + // BLOCK和REPEAT的字节数 + static constexpr uint64_t BYTE_BLOCK = 32UL; + static constexpr uint32_t REPEAT_BLOCK_BYTE = 256U; + // BLOCK和REPEAT的FP32元素数 + static constexpr uint32_t FP32_BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(float); + static constexpr uint32_t FP32_REPEAT_ELEMENT_NUM = REPEAT_BLOCK_BYTE / sizeof(float); + // repeat stride不能超过256 + static constexpr uint32_t REPEATE_STRIDE_UP_BOUND = 256; + +private: + static constexpr bool PAGE_ATTENTION = QSFAT::pageAttention; + static constexpr int TEMPLATE_MODE = QSFAT::templateMode; + static constexpr bool FLASH_DECODE = QSFAT::flashDecode; + static constexpr QSFA_LAYOUT LAYOUT_T = QSFAT::layout; + static constexpr QSFA_LAYOUT KV_LAYOUT_T = QSFAT::kvLayout; + + static constexpr uint64_t MERGE_CACHE_GM_BUF_NUM = 4; + static constexpr uint64_t SYNC_INPUT_BUF1_FLAG = 2; + static constexpr uint64_t SYNC_INPUT_BUF1_PONG_FLAG = 3; + static constexpr uint64_t SYNC_INPUT_BUF2_FLAG = 4; + static constexpr uint64_t SYNC_OUTPUT_BUF1_FLAG = 4; + static constexpr uint64_t SYNC_OUTPUT_BUF2_FLAG = 5; + static constexpr uint32_t INPUT1_BUFFER_OFFSET = ConstInfo::BUFFER_SIZE_BYTE_32K; + static constexpr uint32_t SOFTMAX_TMP_BUFFER_OFFSET = ConstInfo::BUFFER_SIZE_BYTE_512B / sizeof(T); + static constexpr uint32_t BASE_BLOCK_MAX_ELEMENT_NUM = ConstInfo::BUFFER_SIZE_BYTE_32K / sizeof(T); // 32768/4=8096 + static constexpr uint32_t BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(T); // 32/4=8 + static constexpr uint32_t LIMIT_DEAL_ROW = 16U; + static constexpr T FLOAT_E_SCALAR = 8388608; + static constexpr T LN2 = 0.6931471805599453094172; + static constexpr T RECIP_OF_LN2 = 1 / LN2; + static constexpr T SOFTMAX_MIN_NUM = -2e38; + + const KvQuantSparseFlashAttentionTilingDataMla *__restrict tilingData; + + uint32_t pingpongFlag = 0U; + ConstInfo constInfo = {}; + + GlobalTensor mm2ResInt32Gm; + GlobalTensor mm1ResGm; + GlobalTensor vec1ResGm; + GlobalTensor lseSumFdGm; + GlobalTensor lseMaxFdGm; + + GlobalTensor actualSeqLengthsQGm; + GlobalTensor actualSeqLengthsKVGm; + GlobalTensor vec2ResGm; + GlobalTensor mm2ResGm; + GlobalTensor accumOutGm; + GlobalTensor attentionOutGm; + GlobalTensor softmaxMaxGm; + GlobalTensor softmaxSumGm; + GlobalTensor blkTableGm_; + + GlobalTensor kvMergeGm_; + GlobalTensor keyRopeGm_; + GlobalTensor keyGm_; + GlobalTensor topkGm_; + GlobalTensor kvValidSizeGm_; + + // ================================Local Buffer区==================================== + TBuf<> inputBuff1; // 32K * 2 + TBuf<> inputBuff2; // 32K + TBuf<> outputBuff1; // 32K + TBuf<> outputBuff2; // 4K + + TBuf<> tmpBuff1; // 32K + TBuf<> tmpBuff2; // 8K + TBuf<> v0ValidSizeBuff; // 8K + + TBuf<> softmaxMaxBuff; // PRE_LOAD_NUM * 1K + TBuf<> softmaxExpBuff; // PRE_LOAD_NUM * 1K + TBuf<> softmaxSumBuff; // PRE_LOAD_NUM * 1K + TBuf<> softmaxMaxDefaultBuff; // 1K + TBuf<> softmaxSumDefaultBuff; // 1K + + LocalTensor softmaxMaxDefaultUb; + LocalTensor softmaxSumDefaultUb; + + LocalTensor softmaxMaxUb; + LocalTensor softmaxSumUb; + LocalTensor softmaxExpUb; + LocalTensor kvMergUb_; + LocalTensor v0ValidSizeUb_; +}; + +template __aicore__ inline void QSFAVectorService::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(inputBuff1, ConstInfo::BUFFER_SIZE_BYTE_32K * 2); // 2:pingpong + pipe->InitBuffer(inputBuff2, ConstInfo::BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(outputBuff1, ConstInfo::BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(outputBuff2, ConstInfo::BUFFER_SIZE_BYTE_4K); + + pipe->InitBuffer(tmpBuff1, ConstInfo::BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(tmpBuff2, ConstInfo::BUFFER_SIZE_BYTE_8K); + pipe->InitBuffer(v0ValidSizeBuff, ConstInfo::BUFFER_SIZE_BYTE_8K); + + pipe->InitBuffer(softmaxMaxBuff, ConstInfo::BUFFER_SIZE_BYTE_512B * constInfo.preLoadNum); + pipe->InitBuffer(softmaxExpBuff, ConstInfo::BUFFER_SIZE_BYTE_512B * constInfo.preLoadNum); + pipe->InitBuffer(softmaxSumBuff, ConstInfo::BUFFER_SIZE_BYTE_512B * constInfo.preLoadNum); + + pipe->InitBuffer(softmaxMaxDefaultBuff, ConstInfo::BUFFER_SIZE_BYTE_512B); + pipe->InitBuffer(softmaxSumDefaultBuff, ConstInfo::BUFFER_SIZE_BYTE_512B); + + softmaxMaxUb = softmaxMaxBuff.Get(); + softmaxSumUb = softmaxSumBuff.Get(); + softmaxExpUb = softmaxExpBuff.Get(); + + softmaxMaxDefaultUb = softmaxMaxDefaultBuff.Get(); + softmaxSumDefaultUb = softmaxSumDefaultBuff.Get(); + + kvMergUb_ = inputBuff1.Get(); + + v0ValidSizeUb_ = v0ValidSizeBuff.Get(); +} + +template +__aicore__ inline void +QSFAVectorService::InitParams(const struct ConstInfo &constInfo, + const KvQuantSparseFlashAttentionTilingDataMla *__restrict tilingData) +{ + this->constInfo = constInfo; + this->tilingData = tilingData; +} + +template +__aicore__ inline void +QSFAVectorService::InitMm2ResInt32GmGlobalTensor(GlobalTensor mm2ResInt32Gm) +{ + this->mm2ResInt32Gm = mm2ResInt32Gm; +} + +template +__aicore__ inline void QSFAVectorService::InitVec0GlobalTensor( + const GlobalTensor &kvValidSizeGm, const GlobalTensor &kvMergeGm, + const GlobalTensor &keyRopeGm, const GlobalTensor &keyGm, const GlobalTensor &blkTableGm) +{ + this->kvMergeGm_ = kvMergeGm; + this->keyRopeGm_ = keyRopeGm; + this->keyGm_ = keyGm; + this->blkTableGm_ = blkTableGm; + this->kvValidSizeGm_ = kvValidSizeGm; +} + +template +__aicore__ inline void QSFAVectorService::InitVec1GlobalTensor( + GlobalTensor mm1ResGm, GlobalTensor vec1ResGm, + GlobalTensor actualSeqLengthsQGm, GlobalTensor actualSeqLengthsKVGm, GlobalTensor lseMaxFdGm, + GlobalTensor lseSumFdGm, GlobalTensor topKGm, + GlobalTensor softmaxMaxGm, GlobalTensor softmaxSumGm) +{ + this->mm1ResGm = mm1ResGm; + this->vec1ResGm = vec1ResGm; + this->actualSeqLengthsQGm = actualSeqLengthsQGm; + this->actualSeqLengthsKVGm = actualSeqLengthsKVGm; + this->lseMaxFdGm = lseMaxFdGm; + this->lseSumFdGm = lseSumFdGm; + this->topkGm_ = topKGm; + this->softmaxMaxGm = softmaxMaxGm; + this->softmaxSumGm = softmaxSumGm; +} + +template +__aicore__ inline void QSFAVectorService::InitVec2GlobalTensor(GlobalTensor accumOutGm, + GlobalTensor vec2ResGm, + GlobalTensor mm2ResGm, + GlobalTensor attentionOutGm) +{ + this->accumOutGm = accumOutGm; + this->vec2ResGm = vec2ResGm; + this->mm2ResGm = mm2ResGm; + this->attentionOutGm = attentionOutGm; +} + +template __aicore__ inline void QSFAVectorService::AllocEventID() +{ + SetFlag(SYNC_INPUT_BUF1_FLAG); + SetFlag(SYNC_INPUT_BUF1_PONG_FLAG); + SetFlag(SYNC_INPUT_BUF2_FLAG); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); +} + +template __aicore__ inline void QSFAVectorService::FreeEventID() +{ + WaitFlag(SYNC_INPUT_BUF1_FLAG); + WaitFlag(SYNC_INPUT_BUF1_PONG_FLAG); + WaitFlag(SYNC_INPUT_BUF2_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); +} + +template __aicore__ inline void QSFAVectorService::InitSoftmaxDefaultBuffer() +{ + Duplicate(softmaxMaxDefaultUb, SOFTMAX_MIN_NUM, SOFTMAX_TMP_BUFFER_OFFSET); + Duplicate(softmaxSumDefaultUb, ConstInfo::FLOAT_ZERO, SOFTMAX_TMP_BUFFER_OFFSET); +} + +template +__aicore__ inline void QSFAVectorService::ComputeLogSumExpAndCopyToGm(const RunInfo &info, + const MSplitInfo &mSplitInfo, + LocalTensor &softmaxSumUb, + LocalTensor &softmaxMaxUb) +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + uint64_t qsfaBaseOffset = mSplitInfo.nBufferStartM / 2; + size_t qsfaSize = mSplitInfo.vecDealM * FP32_BLOCK_ELEMENT_NUM; + uint64_t qsfaAccumTmpOutNum = CalcAccumOffset(info.bIdx, info.gS1Idx); + uint64_t qsfaOffset = (qsfaAccumTmpOutNum * constInfo.kvHeadNum * constInfo.mBaseSize + // taskoffset + info.tndCoreStartKVSplitPos * constInfo.kvHeadNum * constInfo.mBaseSize + // 份数offset + mSplitInfo.nBufferStartM + mSplitInfo.vecStartM) * + FP32_BLOCK_ELEMENT_NUM; // m轴offset + if (info.actualSingleProcessSInnerSize != 0) { + LocalTensor qsfaTmp = outputBuff2.Get(); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + Brcb(qsfaTmp, softmaxSumUb[qsfaBaseOffset], (mSplitInfo.vecDealM + 7) / 8, {1, 8}); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + DataCopy(lseSumFdGm[qsfaOffset], qsfaTmp, qsfaSize); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + + qsfaTmp = outputBuff2.Get(); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + Brcb(qsfaTmp, softmaxMaxUb[qsfaBaseOffset], (mSplitInfo.vecDealM + 7) / 8, {1, 8}); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + DataCopy(lseMaxFdGm[qsfaOffset], qsfaTmp, qsfaSize); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + } else { + matmul::InitOutput(lseSumFdGm[qsfaOffset], qsfaSize, ConstInfo::FLOAT_ZERO); + matmul::InitOutput(lseMaxFdGm[qsfaOffset], qsfaSize, SOFTMAX_MIN_NUM); + } +} + +template +__aicore__ inline void QSFAVectorService::CopyFALseToGm(const RunInfo &info, const MSplitInfo &mSplitInfo, + LocalTensor &softmaxSumUb, LocalTensor &softmaxMaxUb) +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + uint64_t baseOffset = mSplitInfo.nBufferStartM / 2; + size_t size = mSplitInfo.vecDealM; + + int64_t offset = 0; + if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + uint64_t actualSeqQTotal = (info.bIdx <= 0) ? 0 : actualSeqLengthsQGm.GetValue(constInfo.batchSize - 1); + uint64_t actualSeqQPrefixSum = (info.bIdx <= 0) ? 0 : actualSeqLengthsQGm.GetValue(info.bIdx - 1); + offset += info.n2Idx * actualSeqQTotal * constInfo.gSize + + (actualSeqQPrefixSum + info.gS1Idx / constInfo.gSize) * constInfo.gSize + + mSplitInfo.nBufferStartM + mSplitInfo.vecStartM; + } else { + offset += info.bIdx * constInfo.kvHeadNum * constInfo.qSeqSize * constInfo.gSize + + info.n2Idx * constInfo.qSeqSize * constInfo.gSize + + info.gS1Idx / constInfo.gSize * constInfo.gSize + + mSplitInfo.nBufferStartM + mSplitInfo.vecStartM; + } + + if (info.actualSingleProcessSInnerSize != 0) { + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = 1; + dataCopyParams.blockLen = sizeof(T) * size; + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = 0; + size_t alignedSize = (sizeof(T) * size + 31) / 32 * 32 / sizeof(T); + LocalTensor tmp = outputBuff2.Get(); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + DataCopy(tmp, softmaxMaxUb[baseOffset], alignedSize); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + DataCopyPad(softmaxMaxGm[offset], tmp, dataCopyParams); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + + tmp = outputBuff2.Get(); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + DataCopy(tmp, softmaxSumUb[baseOffset], alignedSize); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + DataCopyPad(softmaxSumGm[offset], tmp, dataCopyParams); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + } else { + matmul::InitOutput(softmaxSumGm[offset], size, ConstInfo::FLOAT_ZERO); + matmul::InitOutput(softmaxMaxGm[offset], size, SOFTMAX_MIN_NUM); + } +} + +template +__aicore__ inline void QSFAVectorService::ElewiseCompute(const RunInfo &info, + const LocalTensor &mmResUb, + uint32_t dealRowCount, uint32_t columnCount) +{ + Muls(mmResUb, mmResUb, static_cast(tilingData->baseParams.scaleValue), dealRowCount * columnCount); + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + // v0的无效值判断 + uint64_t qsfaS2ValidSizeFirstPart = v0ValidSizeUb_.GetValue(128 + info.loop % MERGE_CACHE_GM_BUF_NUM); + uint64_t qsfaS2ValidSizeSecondPart = v0ValidSizeUb_.GetValue(256 + info.loop % MERGE_CACHE_GM_BUF_NUM); + + int64_t qsfaS2ProcessSize = info.actualSingleProcessSInnerSize; + int64_t qsfaS2Pair = CeilDiv(qsfaS2ProcessSize, 2L * constInfo.sparseBlockSize); + int64_t qsfaS2Mid = CeilDiv(qsfaS2Pair, 2L) * 2 * constInfo.sparseBlockSize; + if (qsfaS2Mid > qsfaS2ProcessSize) { + qsfaS2Mid = qsfaS2ProcessSize; + } + if (unlikely(qsfaS2ValidSizeFirstPart < qsfaS2Mid)) { + int64_t qsfaS2StartCeilAlign = CeilAlign(qsfaS2ValidSizeFirstPart, 8); + int64_t qsfaS2MidFloorAlign = qsfaS2Mid / 8 * 8; + // 场景一 s2Mid > s2ValidSizeFirstPart + oneBlk + // 可以推导出s2StartCeilAlign < s2Mid 第一阶段取到s2StartCeilAlign + // s2StartCeilAlign <= s2MidFloorAlign 第二阶段取到s2MidFloorAlign + // 场景二 s2Mid <= s2ValidSizeFirstPart + oneBlk + // 可以推导出 s2StartCeilAlign >= s2Mid 第一阶段取到mid + // s2StartCeilAlign > s2MidFloorAlign 第二阶段取到s2StartCeilAlign + SetInfInBlk(mmResUb, dealRowCount, columnCount, qsfaS2ValidSizeFirstPart, + qsfaS2StartCeilAlign >= qsfaS2Mid ? qsfaS2Mid : qsfaS2StartCeilAlign); + SetMidInf(mmResUb, dealRowCount, columnCount, qsfaS2StartCeilAlign, qsfaS2MidFloorAlign); + SetInfInBlk(mmResUb, dealRowCount, columnCount, + qsfaS2StartCeilAlign <= qsfaS2MidFloorAlign ? \ + qsfaS2MidFloorAlign : qsfaS2StartCeilAlign, qsfaS2Mid); + } + if (unlikely(qsfaS2ValidSizeSecondPart < qsfaS2ProcessSize - qsfaS2Mid)) { + // 场景一 s2Mid + s2ValidSizeSecondPart > s2ProcessSize + oneBlk + // 可以推导出 s2StartCeilAlign < s2ProcessSize 第一阶段取到s2StartCeilAlign + // s2StartCeilAlign <= s2EndFloorAlign 第二阶段取到s2EndFloorAlign + // 场景二 s2Mid + s2ValidSizeSecondPart <= s2ProcessSize + oneBlk + // 可以推导出 s2StartCeilAlign >= s2ProcessSize 第一阶段取到s2ProcessSize + // s2StartCeilAlign > s2EndFloorAlign 第二阶段取到s2StartCeilAlign + int64_t qsfaS2StartCeilAlign = CeilAlign(qsfaS2Mid + qsfaS2ValidSizeSecondPart, 8); + int64_t qsfaS2EndFloorAlign = qsfaS2ProcessSize / 8 * 8; + SetInfInBlk(mmResUb, dealRowCount, columnCount, qsfaS2Mid + qsfaS2ValidSizeSecondPart, + qsfaS2StartCeilAlign >= qsfaS2ProcessSize ? qsfaS2ProcessSize : qsfaS2StartCeilAlign); + SetMidInf(mmResUb, dealRowCount, columnCount, qsfaS2StartCeilAlign, qsfaS2EndFloorAlign); + SetInfInBlk(mmResUb, dealRowCount, columnCount, + qsfaS2StartCeilAlign <= qsfaS2EndFloorAlign ? qsfaS2EndFloorAlign : qsfaS2StartCeilAlign, + qsfaS2ProcessSize); + } + } +} + +template +__aicore__ inline void QSFAVectorService::SetInfInBlk(const LocalTensor &mmResUb, + uint32_t dealRowCount, uint32_t columnCount, + uint64_t startId, uint64_t endId) +{ + // startId endId + // x x x 0 0 0 x x x + // 从startId到endId部分置-inf, endId、startId为endId一个blk内部的下标 + if (startId >= endId) { + return; + } + + uint64_t qsfaStartFloorAlignSize = startId / BLOCK_ELEMENT_NUM * BLOCK_ELEMENT_NUM; + uint64_t qsfaNotComputePreMaskOneBlk = (1 << (startId - qsfaStartFloorAlignSize)) - 1; + uint64_t qsfaNotComputePostMaskOneBlk = ~((1 << (endId - qsfaStartFloorAlignSize)) - 1); + uint64_t qsfaNotComputeMaskOneBlk = qsfaNotComputePreMaskOneBlk ^ qsfaNotComputePostMaskOneBlk; + + uint64_t qsfaMaskOneBlk = ~qsfaNotComputeMaskOneBlk; + uint64_t mask[1] = {qsfaMaskOneBlk}; + for (int i = 1; i < 8; i++) { + mask[0] = mask[0] | (qsfaMaskOneBlk << (i * 8)); + } + for (uint64_t qsfaRowId = 0; qsfaRowId < dealRowCount; qsfaRowId += 8) { + Duplicate(mmResUb[qsfaRowId * columnCount + qsfaStartFloorAlignSize], SOFTMAX_MIN_NUM, mask, + 1, CeilDiv(columnCount, 8), 0); + } +} + +template +__aicore__ inline void QSFAVectorService::SetMidInf(const LocalTensor &mmResUb, + uint32_t dealRowCount, uint32_t columnCount, + uint64_t startId, uint64_t endId) +{ + if (startId >= endId) { + return; + } + // startId endId + // 0 ... 0 + // 从startId到endId部分置-inf, startId、endId为32B对齐的下标 + for (uint64_t qsfaRowId = 0; qsfaRowId < dealRowCount; qsfaRowId++) { + Duplicate(mmResUb[qsfaRowId * columnCount + startId], SOFTMAX_MIN_NUM, endId - startId); + } +} + +template +__aicore__ inline void QSFAVectorService::SoftmaxFlashV2Compute( + const RunInfo &info, const MSplitInfo &mSplitInfo, LocalTensor &mmResUb, LocalTensor &softmaxTmpUb, + uint32_t startRow, uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount) +{ + LocalTensor inSumTensor; + LocalTensor inMaxTensor; + uint32_t baseOffset = mSplitInfo.nBufferStartM / 2 + startRow; + uint32_t outIdx = info.loop % (constInfo.preLoadNum); + uint32_t softmaxOutOffset = outIdx * SOFTMAX_TMP_BUFFER_OFFSET + baseOffset; + if (info.isFirstSInnerLoop) { + inMaxTensor = softmaxMaxDefaultUb; + inSumTensor = softmaxSumDefaultUb; + } else { + uint32_t inIdx = (info.loop - 1) % (constInfo.preLoadNum); + inMaxTensor = softmaxMaxUb[inIdx * SOFTMAX_TMP_BUFFER_OFFSET + baseOffset]; + inSumTensor = softmaxSumUb[inIdx * SOFTMAX_TMP_BUFFER_OFFSET + baseOffset]; + } + if (actualColumnCount !=0) { + SoftMaxShapeInfo srcShape{dealRowCount, columnCount, dealRowCount, actualColumnCount}; + SoftMaxTiling newTiling = + SoftMaxFlashV2TilingFunc(srcShape, sizeof(T), sizeof(T), softmaxTmpUb.GetSize(), true, false); + SoftmaxFlashV2( + mmResUb, softmaxSumUb[softmaxOutOffset], softmaxMaxUb[softmaxOutOffset], mmResUb, + softmaxExpUb[softmaxOutOffset], inSumTensor, inMaxTensor, softmaxTmpUb, newTiling, srcShape); + } else { + DataCopy(softmaxSumUb[softmaxOutOffset], inSumTensor, dealRowCount); + PipeBarrier(); + DataCopy(softmaxMaxUb[softmaxOutOffset], inMaxTensor, dealRowCount); + } +} + +template +__aicore__ inline void QSFAVectorService::DealBmm1ResBaseBlock( + const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t startRow, uint32_t dealRowCount, + uint32_t columnCount, uint32_t loopId) +{ + uint32_t qsfaComputeSize = dealRowCount * columnCount; + uint64_t qsfaInOutGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize + + (mSplitInfo.nBufferStartM + mSplitInfo.vecStartM + startRow) * columnCount; + LocalTensor qsfaMmResUb = inputBuff1.Get(); + qsfaMmResUb = qsfaMmResUb[pingpongFlag * INPUT1_BUFFER_OFFSET / sizeof(MM1_OUT_T)]; + WaitFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + + DataCopy(qsfaMmResUb, mm1ResGm[qsfaInOutGmOffset], qsfaComputeSize); + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + if (loopId == 0) { + WaitFlag(0); + } + } + SetFlag(SYNC_INPUT_BUF1_FLAG); + WaitFlag(SYNC_INPUT_BUF1_FLAG); + + ElewiseCompute(info, qsfaMmResUb, dealRowCount, columnCount); + + PipeBarrier(); + LocalTensor qsfaTmpAFloorUb = tmpBuff1.Get(); + LocalTensor qsfaSoftmaxTmpUb = qsfaTmpAFloorUb.template ReinterpretCast(); + + SoftmaxFlashV2Compute(info, mSplitInfo, qsfaMmResUb, qsfaSoftmaxTmpUb, startRow, dealRowCount, columnCount, + info.actualSingleProcessSInnerSize); + + PipeBarrier(); + LocalTensor tmpMMResCastTensor = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + + Cast(tmpMMResCastTensor, qsfaMmResUb, AscendC::RoundMode::CAST_ROUND, qsfaComputeSize); + SetFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + DataCopy(vec1ResGm[qsfaInOutGmOffset], tmpMMResCastTensor, qsfaComputeSize); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); +} + +template +__aicore__ inline void QSFAVectorService::ProcessVec1SingleBuf(const RunInfo &info, + const MSplitInfo &mSplitInfo) +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + uint32_t qsfaMSplitSize = info.actualSingleProcessSInnerSize == 0 ? + 16 : (BASE_BLOCK_MAX_ELEMENT_NUM / info.actualSingleProcessSInnerSizeAlign); + // 1. 向下8对齐是因为UB操作至少32B + // 2. info.actualSingleProcessSInnerSizeAlign最大512, mSplitSize可以确保最小为16 + qsfaMSplitSize = qsfaMSplitSize >> 3U << 3U; + + if (qsfaMSplitSize > mSplitInfo.vecDealM) { + qsfaMSplitSize = mSplitInfo.vecDealM; + } + uint32_t qsfaLoopCount = (mSplitInfo.vecDealM + qsfaMSplitSize - 1) / qsfaMSplitSize; + uint32_t qsfaTailSplitSize = mSplitInfo.vecDealM - (qsfaLoopCount - 1) * qsfaMSplitSize; + + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = 1; + dataCopyParams.blockLen = 256 * sizeof(int32_t); + dataCopyParams.dstStride = 0; + dataCopyParams.srcStride = 0; + DataCopyPadExtParams padParams; + // 额外偏移128个元素，避免不同loop下v0和v1互相影响 + DataCopyPad(v0ValidSizeUb_[128], kvValidSizeGm_[info.loop % MERGE_CACHE_GM_BUF_NUM * (128 * 2)], + dataCopyParams, padParams); + SetFlag(0); + if (unlikely(qsfaLoopCount == 0)) { + // scalar同步影响较大，挪到循环内部进行 + WaitFlag(0); + } + } + for (uint32_t qsfaI = 0, dealSize = qsfaMSplitSize; qsfaI < qsfaLoopCount; qsfaI++) { + if (qsfaI == (qsfaLoopCount - 1)) { + dealSize = qsfaTailSplitSize; + } + DealBmm1ResBaseBlock(info, mSplitInfo, qsfaI * qsfaMSplitSize, dealSize, + info.actualSingleProcessSInnerSizeAlign, qsfaI); + pingpongFlag ^= 1; // pingpong 0 1切换 + } +} + +template +__aicore__ inline void QSFAVectorService::GetRealS2Idx(int64_t s2GmOffset, int64_t &realS2Idx, + int64_t topkGmBaseOffset, const RunInfo &runInfo) +{ + int64_t qsfaTopkGmIdx = (s2GmOffset + runInfo.s2Idx * constInfo.s2BaseSize) / constInfo.sparseBlockSize; + if (unlikely(qsfaTopkGmIdx >= constInfo.sparseBlockCount)) { + realS2Idx = -1; + return; + } + realS2Idx = topkGm_.GetValue(topkGmBaseOffset + qsfaTopkGmIdx) * static_cast(constInfo.sparseBlockSize) + + static_cast((s2GmOffset + runInfo.s2Idx * constInfo.s2BaseSize) % constInfo.sparseBlockSize); +} + +template +__aicore__ inline int64_t QSFAVectorService::GetKeyBNBOffset(int64_t realS2Idx, + const RunInfo &runInfo, int64_t s2IdLimit) +{ + if (realS2Idx < 0 || realS2Idx >= s2IdLimit) { + return -1; + } + int64_t realKeyBNBOffset = 0; + if constexpr (PAGE_ATTENTION) { + int64_t blkTableIdx = realS2Idx / constInfo.kvCacheBlockSize; + int64_t blkTableOffset = realS2Idx % constInfo.kvCacheBlockSize; + realKeyBNBOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo.maxBlockNumPerBatch + blkTableIdx) * + static_cast(constInfo.kvCacheBlockSize) * + static_cast(constInfo.kvHeadNum) + + blkTableOffset; + } else { + realKeyBNBOffset = (runInfo.tensorBOffset + + realS2Idx * constInfo.kvHeadNum * constInfo.combineHeadDim) / + constInfo.combineHeadDim; + } + return realKeyBNBOffset; +} + +template +__aicore__ inline void +QSFAVectorService::CopyInSingleKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, int64_t realS2Idx, + int64_t keyBNBOffset, int64_t s2IdLimit, const RunInfo &runInfo) +{ + if (keyBNBOffset < 0) { + return; + } + int64_t validS2Count = + ((realS2Idx + constInfo.sparseBlockSize > s2IdLimit) ? (s2IdLimit - realS2Idx) : constInfo.sparseBlockSize); + DataCopyExtParams intriParams; + + intriParams.blockCount = validS2Count; + intriParams.dstStride = 0; + intriParams.srcStride = 0; + DataCopyPadExtParams padParams; + // 当前仅支持COMBINE模式 + if (constInfo.quantScaleRepoMode == QUANT_SCALE_REPO_MODE::COMBINE) { + uint32_t combineBytes = (constInfo.headDim * sizeof(KV_T) + constInfo.headDimRope * sizeof(K_ROPE_T) + + constInfo.headDim / constInfo.tileSize * sizeof(T)); + intriParams.blockLen = combineBytes; + uint32_t combineDim = combineBytes / sizeof(KV_T); + uint32_t combineDimAlign = CeilAlign(combineBytes, ConstInfo::BUFFER_SIZE_BYTE_32B) / sizeof(KV_T); + padParams.isPad = true; + padParams.leftPadding = 0; + padParams.rightPadding = combineDimAlign - combineDim; + padParams.paddingValue = 0; + DataCopyPad(kvMergUb_[mergeMte3Idx % 2 * INPUT1_BUFFER_OFFSET / sizeof(KV_T) + (mte2Size - mte3Size) * + combineDimAlign], keyGm_[keyBNBOffset * combineDim], intriParams, padParams); + } + mte2Size += validS2Count; +} + +template +__aicore__ inline void QSFAVectorService::CopyInKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, + int64_t realS2Idx1, int64_t realS2Idx2, + const RunInfo &runInfo) +{ + int64_t s2IdLimit = runInfo.curActualSeqLenOri; + if (constInfo.sparseMode == 3) { + s2IdLimit = runInfo.curActualSeqLenOri - runInfo.actS1Size + runInfo.gS1Idx / constInfo.gSize + 1; + } + + int64_t keyBNBOffset1 = GetKeyBNBOffset(realS2Idx1, runInfo, s2IdLimit); + int64_t keyBNBOffset2 = GetKeyBNBOffset(realS2Idx2, runInfo, s2IdLimit); + if (unlikely(keyBNBOffset1 < 0 && keyBNBOffset2 < 0)) { + return; + } + + int64_t sparseBlockSrcStride = + ((keyBNBOffset1 > keyBNBOffset2 ? (keyBNBOffset1 - keyBNBOffset2) : + (keyBNBOffset2 - keyBNBOffset1)) - constInfo.sparseBlockSize); + uint32_t combineBytes = (constInfo.headDim * sizeof(KV_T) + + constInfo.headDimRope * sizeof(K_ROPE_T) + + constInfo.headDim / constInfo.tileSize * sizeof(T)); + int64_t keySrcStride = sparseBlockSrcStride * combineBytes; + if (unlikely(keySrcStride >= INT32_MAX || keySrcStride < 0 || + realS2Idx1 + constInfo.sparseBlockSize >= s2IdLimit || + realS2Idx2 + constInfo.sparseBlockSize >= s2IdLimit) || + constInfo.sparseBlockSize > 1) { + // stride溢出、stride为负数、s2超长等异常场景，还原成2条搬运指令 + CopyInSingleKv(mte2Size, mte3Size, mergeMte3Idx, realS2Idx1, keyBNBOffset1, s2IdLimit, runInfo); + CopyInSingleKv(mte2Size, mte3Size, mergeMte3Idx, realS2Idx2, keyBNBOffset2, s2IdLimit, runInfo); + } else { + DataCopyExtParams intriParams; + intriParams.blockCount = (keyBNBOffset1 >= 0) + (keyBNBOffset2 >= 0); + intriParams.dstStride = 0; + intriParams.srcStride = keySrcStride; + DataCopyPadExtParams padParams; + + int64_t startGmOffset = keyBNBOffset1 > -1 ? keyBNBOffset1 : keyBNBOffset2; + if (keyBNBOffset2 > -1 && keyBNBOffset2 < keyBNBOffset1) { + startGmOffset = keyBNBOffset2; + } + + // 当前仅支持COMBINE模式 + if (constInfo.quantScaleRepoMode == QUANT_SCALE_REPO_MODE::COMBINE) { + intriParams.blockLen = constInfo.sparseBlockSize * combineBytes; + uint32_t combineDim = combineBytes / sizeof(KV_T); + uint32_t combineDimAlign = CeilAlign(combineBytes, ConstInfo::BUFFER_SIZE_BYTE_32B) / sizeof(KV_T); + padParams.isPad = true; + padParams.leftPadding = 0; + padParams.rightPadding = combineDimAlign - combineDim; + padParams.paddingValue = 0; + DataCopyPad(kvMergUb_[mergeMte3Idx % 2 * INPUT1_BUFFER_OFFSET / sizeof(KV_T) + (mte2Size - mte3Size) * + combineDimAlign], keyGm_[startGmOffset * combineDim], intriParams, padParams); + } + mte2Size += ((keyBNBOffset1 > -1) + (keyBNBOffset2 > -1)) * constInfo.sparseBlockSize; + } +} + +template +__aicore__ inline void QSFAVectorService::CopyOutMrgeResult(int64_t mte2Size, int64_t mte3Size, + int64_t s2GmStartOffset, int64_t mergeMte3Idx, + const RunInfo &runInfo) +{ + if (mte2Size <= mte3Size) { + return; + } + int32_t dealRow = mte2Size - mte3Size; + SetFlag(0); + WaitFlag(0); + LocalTensor kvTensorAsFp16 = tmpBuff1.Get(); + uint64_t mask = ConstInfo::BUFFER_SIZE_BYTE_256B / sizeof(half); + LocalTensor srcTensor = kvMergUb_[mergeMte3Idx % 2 * INPUT1_BUFFER_OFFSET / sizeof(KV_T)]; + if (dealRow == 1) { + Cast(kvTensorAsFp16, srcTensor, RoundMode::CAST_NONE, mask, 4, {1, 1, 8, 4}); + } else { + uint8_t repeatTimes = static_cast(dealRow); + Cast(kvTensorAsFp16, srcTensor, RoundMode::CAST_NONE, mask, repeatTimes, {1, 1, 32, 21}); // 21=(512+64*2+32)/32 + Cast(kvTensorAsFp16[128], srcTensor[128], RoundMode::CAST_NONE, mask, repeatTimes, {1, 1, 32, 21}); + Cast(kvTensorAsFp16[256], srcTensor[256], RoundMode::CAST_NONE, mask, repeatTimes, {1, 1, 32, 21}); + Cast(kvTensorAsFp16[384], srcTensor[384], RoundMode::CAST_NONE, mask, repeatTimes, {1, 1, 32, 21}); + } + PipeBarrier(); + LocalTensor antiQuantScale = tmpBuff2.Get(); + LocalTensor oriQuantScaleTensor = srcTensor[640].template ReinterpretCast(); + if (dealRow == 1) { + Brcb(antiQuantScale, oriQuantScaleTensor, 1, {1, 4}); + } else { + DataCopyParams params; + params.blockCount = dealRow; + params.blockLen = 1; + params.dstStride = 0; + params.srcStride = (constInfo.headDim * sizeof(KV_T) + constInfo.headDimRope * sizeof(K_ROPE_T)) / + ConstInfo::BUFFER_SIZE_BYTE_32B; + LocalTensor tmpAntiQuantScale = antiQuantScale[ConstInfo::BUFFER_SIZE_BYTE_1K]; + DataCopy(tmpAntiQuantScale, oriQuantScaleTensor, params); + PipeBarrier(); + Brcb(antiQuantScale, tmpAntiQuantScale, dealRow, {1, 4}); + } + PipeBarrier(); + uint32_t dealLoop = CeilDiv(dealRow, LIMIT_DEAL_ROW); + uint32_t dealRowFp32 = LIMIT_DEAL_ROW; + uint32_t element = LIMIT_DEAL_ROW * constInfo.headDim; + LocalTensor kvTensorAsFp32 = inputBuff2.Get(); + LocalTensor antiKvTensorAsB16 = tmpBuff1.Get(); + for (uint32_t i = 0; i < dealLoop; i++) { + if (i == dealLoop - 1) { + dealRowFp32 = dealRow - i * LIMIT_DEAL_ROW; + } + Cast(kvTensorAsFp32, kvTensorAsFp16[i * element], RoundMode::CAST_NONE, + static_cast(dealRowFp32 * constInfo.headDim)); + PipeBarrier(); + for (uint32_t j = 0; j < constInfo.tileSize / FP32_REPEAT_ELEMENT_NUM; j++) { + Mul(kvTensorAsFp32[j * FP32_REPEAT_ELEMENT_NUM], kvTensorAsFp32[j * FP32_REPEAT_ELEMENT_NUM], + antiQuantScale[i * LIMIT_DEAL_ROW * 32], + FP32_REPEAT_ELEMENT_NUM, 4 * dealRowFp32, {1, 1, 0, 16, 16, 1}); + } + PipeBarrier(); + if constexpr (IsSameType::value) { // bf16 采取四舍六入五成双模式 + Cast(antiKvTensorAsB16[i * element], kvTensorAsFp32, RoundMode::CAST_RINT, + static_cast(dealRowFp32 * constInfo.headDim)); + } else { + Cast(antiKvTensorAsB16[i * element], kvTensorAsFp32, RoundMode::CAST_ROUND, + static_cast(dealRowFp32 * constInfo.headDim)); + } + PipeBarrier(); + } + + LocalTensor antiKvTensorAsB16Nz = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + int dataBlocks = REPEAT_BLOCK_BYTE / BYTE_BLOCK; + int loops = CeilDiv(dealRow, dataBlocks); + uint64_t tail = dealRow - (loops - 1) * dataBlocks; + uint64_t repeatElementNum = FP32_REPEAT_ELEMENT_NUM * 2; + uint64_t blockElementNum = FP32_BLOCK_ELEMENT_NUM * 2; + uint8_t repeatTimes = static_cast(constInfo.headDim / blockElementNum); + for (int i = 0; i < loops; i++) { + mask = (i == loops - 1) ? tail * blockElementNum : repeatElementNum; + Copy(antiKvTensorAsB16Nz[i * repeatElementNum], antiKvTensorAsB16[i * dataBlocks * constInfo.headDim], mask, + repeatTimes, {1, 32, static_cast(dealRow), 1}); + } + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = constInfo.headDim / blockElementNum; + dataCopyParams.blockLen = dealRow * blockElementNum * sizeof(K_ROPE_T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (constInfo.s2BaseSize - dealRow) * blockElementNum * sizeof(K_ROPE_T); + DataCopyPad(kvMergeGm_[runInfo.loop % MERGE_CACHE_GM_BUF_NUM * 512 * 576 + (s2GmStartOffset + mte3Size) * + blockElementNum], antiKvTensorAsB16Nz, dataCopyParams); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + + LocalTensor kRopeUb = srcTensor[512].template ReinterpretCast(); + LocalTensor kRopeUbNz = outputBuff2.Get(); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + Copy(kRopeUbNz, kRopeUb, constInfo.headDimRope, static_cast(dealRow), {static_cast(dealRow), 1, + 1, 21}); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + dataCopyParams.blockCount = constInfo.headDimRope / blockElementNum; + DataCopyPad(kvMergeGm_[runInfo.loop % MERGE_CACHE_GM_BUF_NUM * 512 * 576 + 512 * 512 + (s2GmStartOffset + + mte3Size) * blockElementNum], kRopeUbNz, dataCopyParams); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); +} + +// b s1 k +template +__aicore__ inline void QSFAVectorService::MergeKv(const RunInfo &runInfo) +{ + int64_t s2ProcessSize = runInfo.actualSingleProcessSInnerSize; + int64_t s2Pair = CeilDiv(s2ProcessSize, 2L * constInfo.sparseBlockSize); + int64_t topkGmBaseOffset = 0; + + if constexpr (LAYOUT_T == QSFA_LAYOUT::TND) { + uint64_t qsfaActualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsQGm.GetValue(runInfo.bIdx - 1); + topkGmBaseOffset += (qsfaActualSeqQPrefixSum + runInfo.gS1Idx / constInfo.gSize) * constInfo.kvHeadNum * + constInfo.sparseBlockCount + runInfo.n2Idx * constInfo.sparseBlockCount; + } else { + topkGmBaseOffset += runInfo.bIdx * constInfo.qSeqSize * constInfo.sparseBlockCount + + runInfo.gS1Idx / constInfo.gSize * constInfo.sparseBlockCount; + } + int64_t qsfaMergeMte3Idx = 0; + int64_t qsfaMte2Size = 0; + int64_t qsfaMte3Size = 0; + int64_t qsfaS2IdxArray0 = -1; + int64_t qsfaS2IdxArray1 = -1; + bool qsfaNeedWaitMte3ToMte2 = true; + SetFlag(0); + SetFlag(1); + int64_t qsfaS2GmStartOffset = GetSubBlockIdx() == 0 ? 0 : CeilDiv(s2Pair, 2L) * 2 * constInfo.sparseBlockSize; + int64_t qsfaS2GmLimit = GetSubBlockIdx() == 0 ? CeilDiv(s2Pair, 2L) * 2 * constInfo.sparseBlockSize: s2ProcessSize; + if (qsfaS2GmLimit > s2ProcessSize) { + qsfaS2GmLimit = s2ProcessSize; + } + for (int64_t s2GmOffsetArray = qsfaS2GmStartOffset; s2GmOffsetArray < qsfaS2GmLimit; s2GmOffsetArray += 2 * + constInfo.sparseBlockSize) { + if (qsfaNeedWaitMte3ToMte2) { + WaitFlag(qsfaMergeMte3Idx % 2); + qsfaNeedWaitMte3ToMte2 = false; + } + GetRealS2Idx(s2GmOffsetArray, qsfaS2IdxArray0, topkGmBaseOffset, runInfo); + if (unlikely(qsfaS2IdxArray0 < 0)) { + CopyOutMrgeResult(qsfaMte2Size, qsfaMte3Size, qsfaS2GmStartOffset, qsfaMergeMte3Idx, runInfo); + SetFlag(qsfaMergeMte3Idx % 2); + qsfaMergeMte3Idx++; + break; + } + GetRealS2Idx(s2GmOffsetArray + constInfo.sparseBlockSize, qsfaS2IdxArray1, topkGmBaseOffset, runInfo); + CopyInKv(qsfaMte2Size, qsfaMte3Size, qsfaMergeMte3Idx, qsfaS2IdxArray0, qsfaS2IdxArray1, runInfo); + if ((qsfaMte2Size - qsfaMte3Size + 2 * constInfo.sparseBlockSize > 32) || + s2GmOffsetArray + 2 * constInfo.sparseBlockSize >= qsfaS2GmLimit) { + CopyOutMrgeResult(qsfaMte2Size, qsfaMte3Size, qsfaS2GmStartOffset, qsfaMergeMte3Idx, runInfo); + qsfaMte3Size = qsfaMte2Size; + SetFlag(qsfaMergeMte3Idx % 2); + qsfaMergeMte3Idx++; + qsfaNeedWaitMte3ToMte2 = true; + } + } + + if (unlikely(qsfaS2GmStartOffset + qsfaMte2Size < qsfaS2GmLimit)) { + uint64_t blockElementNum = FP32_BLOCK_ELEMENT_NUM * 2; + SetFlag(0); + WaitFlag(0); + WaitFlag(qsfaMergeMte3Idx & 1); + LocalTensor mergeUb = kvMergUb_.template ReinterpretCast(); + Duplicate(mergeUb, static_cast(0.0), constInfo.headDim); + SetFlag(0); + WaitFlag(0); + + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = constInfo.headDim / blockElementNum; + dataCopyParams.blockLen = blockElementNum * sizeof(K_ROPE_T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (constInfo.s2BaseSize - 1) * blockElementNum * sizeof(K_ROPE_T); + for (int64_t s2GmOffset = qsfaS2GmStartOffset + qsfaMte2Size; s2GmOffset < qsfaS2GmLimit; s2GmOffset++) { + DataCopyPad(kvMergeGm_[runInfo.loop % MERGE_CACHE_GM_BUF_NUM * 512 * 576 + s2GmOffset * blockElementNum], + mergeUb, dataCopyParams); + } + dataCopyParams.blockCount = constInfo.headDimRope / blockElementNum; + for (int64_t s2GmOffset = qsfaS2GmStartOffset + qsfaMte2Size; s2GmOffset < qsfaS2GmLimit; s2GmOffset++) { + DataCopyPad(kvMergeGm_[runInfo.loop % MERGE_CACHE_GM_BUF_NUM * 512 * 576 + 512 * constInfo.headDim + + s2GmOffset * blockElementNum], + mergeUb, dataCopyParams); + } + SetFlag(qsfaMergeMte3Idx & 1); + qsfaMergeMte3Idx++; + } + WaitFlag(0); + WaitFlag(1); + v0ValidSizeUb_.SetValue(runInfo.loop % MERGE_CACHE_GM_BUF_NUM, qsfaMte2Size); + SetFlag(1); + WaitFlag(1); + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = 1; + dataCopyParams.blockLen = 128 * sizeof(int32_t); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = 0; + DataCopyPad(kvValidSizeGm_[runInfo.loop % MERGE_CACHE_GM_BUF_NUM * (128 * 2) + GetSubBlockIdx() * 128], + v0ValidSizeUb_, dataCopyParams); + return; +} + +template +__aicore__ inline void QSFAVectorService::ProcessVec1L(const RunInfo &info) +{ + uint32_t qsfaNBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t qsfaNBufferTail = info.actMBaseSize - (qsfaNBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t qsfaI = 0; qsfaI < qsfaNBufferLoopTimes; qsfaI++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferIdx = qsfaI; + mSplitInfo.nBufferStartM = qsfaI * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (qsfaI + 1 != qsfaNBufferLoopTimes) ? constInfo.nBufferMBaseSize : qsfaNBufferTail; + + mSplitInfo.vecDealM = (mSplitInfo.nBufferDealM <= 16) ? mSplitInfo.nBufferDealM : + (((mSplitInfo.nBufferDealM + 15) / 16 + 1) / 2 * 16); + mSplitInfo.vecStartM = 0; + if (GetBlockIdx() % 2 == 1) { + mSplitInfo.vecStartM = mSplitInfo.vecDealM; + mSplitInfo.vecDealM = mSplitInfo.nBufferDealM - mSplitInfo.vecDealM; + } + + CrossCoreWaitFlag(constInfo.syncC1V1); + // vec1 compute + ProcessVec1SingleBuf(info, mSplitInfo); + CrossCoreSetFlag(constInfo.syncV1C2); + // move lse for flash decode + if (info.s2Idx == info.curSInnerLoopTimes - 1 && (constInfo.returnSoftmaxLse || info.tndIsS2SplitCore)) { + uint32_t outIdx = info.loop % (constInfo.preLoadNum); + auto sumTensor = softmaxSumUb[outIdx * SOFTMAX_TMP_BUFFER_OFFSET]; + auto maxTensor = softmaxMaxUb[outIdx * SOFTMAX_TMP_BUFFER_OFFSET]; + if (constInfo.returnSoftmaxLse) { + CopyFALseToGm(info, mSplitInfo, sumTensor, maxTensor); + } + if (info.tndIsS2SplitCore) { + if constexpr (FLASH_DECODE) { + ComputeLogSumExpAndCopyToGm(info, mSplitInfo, sumTensor, maxTensor); + } + } + } + } +} + +template +__aicore__ inline uint64_t QSFAVectorService::CalcAccumOffset(uint32_t bN2Idx, uint32_t gS1Idx) +{ + return 0; +} + +template +__aicore__ inline void QSFAVectorService::ProcessVec2SingleBuf(const RunInfo &info, + const MSplitInfo &mSplitInfo) +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + + uint32_t gPreSplitSize = BASE_BLOCK_MAX_ELEMENT_NUM / constInfo.headDim; + if (gPreSplitSize > mSplitInfo.vecDealM) { + gPreSplitSize = mSplitInfo.vecDealM; + } + uint32_t loopCount = (mSplitInfo.vecDealM + gPreSplitSize - 1) / gPreSplitSize; + uint32_t tailSplitSize = mSplitInfo.vecDealM - (loopCount - 1) * gPreSplitSize; + + for (uint32_t i = 0, dealSize = gPreSplitSize; i < loopCount; i++) { + if (i == (loopCount - 1)) { + dealSize = tailSplitSize; + } + DealBmm2ResBaseBlock(info, mSplitInfo, i * gPreSplitSize, dealSize, constInfo.headDim, constInfo.headDim); + pingpongFlag ^= 1; // pingpong 0 1切换 + } +} + +template +__aicore__ inline void QSFAVectorService::DealBmm2ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, + uint32_t startRow, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + uint32_t vec2ComputeSize = dealRowCount * columnCount; + uint32_t baseOffset = startRow; + LocalTensor bmm2ResUb = tmpBuff1.Get(); + bmm2ResUb.SetSize(vec2ComputeSize); + + size_t batchBase = 0; + uint64_t inOutBaseOffset = (mSplitInfo.vecStartM + startRow) * columnCount; + uint64_t srcGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + inOutBaseOffset; + + LocalTensor tmpBmm2ResUb = inputBuff1.Get(); + tmpBmm2ResUb = tmpBmm2ResUb[pingpongFlag * INPUT1_BUFFER_OFFSET / sizeof(MM2_OUT_T)]; + WaitFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + + DataCopy(tmpBmm2ResUb, mm2ResGm[srcGmOffset + batchBase], vec2ComputeSize); + SetFlag(SYNC_INPUT_BUF1_FLAG); + WaitFlag(SYNC_INPUT_BUF1_FLAG); + DataCopy(bmm2ResUb, tmpBmm2ResUb, vec2ComputeSize); + SetFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + + // 除第一个循环外，均需要更新中间计算结果 + if (info.s2Idx > 0) { + event_t eventIdMte2WaitMte3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventIdMte2WaitMte3); + WaitFlag(eventIdMte2WaitMte3); + LocalTensor bmm2ResPreUb = inputBuff2.Get(); + WaitFlag(SYNC_INPUT_BUF2_FLAG); + uint64_t vecPre2ResGmOffset = ((info.loop - 1) % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + + inOutBaseOffset; + DataCopy(bmm2ResPreUb, vec2ResGm[vecPre2ResGmOffset + batchBase], vec2ComputeSize); + SetFlag(SYNC_INPUT_BUF2_FLAG); + WaitFlag(SYNC_INPUT_BUF2_FLAG); + LocalTensor softmaxExpBrcb = tmpBuff2.Get(); + Brcb(softmaxExpBrcb, softmaxExpUb[(info.loop % constInfo.preLoadNum) * SOFTMAX_TMP_BUFFER_OFFSET + baseOffset], + (mSplitInfo.vecDealM + 7) / 8, {1, 8}); + PipeBarrier(); + RowMuls(bmm2ResPreUb, bmm2ResPreUb, softmaxExpBrcb, dealRowCount, columnCount, actualColumnCount); + PipeBarrier(); + Add(bmm2ResUb, bmm2ResUb, bmm2ResPreUb, vec2ComputeSize); + SetFlag(SYNC_INPUT_BUF2_FLAG); + } + // 最后一次输出计算结果，否则将中间结果暂存至workspace + if (info.s2Idx + 1 == info.curSInnerLoopTimes) { + LocalTensor softmaxSumBrcb = tmpBuff2.Get(); + Brcb(softmaxSumBrcb, softmaxSumUb[(info.loop % constInfo.preLoadNum) * SOFTMAX_TMP_BUFFER_OFFSET + baseOffset], + (mSplitInfo.vecDealM + 7) / 8, {1, 8}); + PipeBarrier(); + RowDivs(bmm2ResUb, bmm2ResUb, softmaxSumBrcb, dealRowCount, columnCount, actualColumnCount); + + PipeBarrier(); + Bmm2ResCopyOut(info, bmm2ResUb, mSplitInfo.vecStartM + startRow, dealRowCount, columnCount, actualColumnCount); + } else { + PipeBarrier(); + LocalTensor tmpBmm2Res = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + DataCopy(tmpBmm2Res, bmm2ResUb, dealRowCount * columnCount); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + + uint64_t vecPre2ResGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + inOutBaseOffset; + DataCopy(vec2ResGm[vecPre2ResGmOffset + batchBase], tmpBmm2Res, vec2ComputeSize); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + } +} + +template __aicore__ inline void QSFAVectorService::ProcessVec2L(const RunInfo &info) +{ + uint32_t qsfaNBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t qsfaNBufferTail = info.actMBaseSize - (qsfaNBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t qsfaI = 0; qsfaI < qsfaNBufferLoopTimes; qsfaI++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferIdx = qsfaI; + mSplitInfo.nBufferDealM = (qsfaI + 1 != qsfaNBufferLoopTimes) ? constInfo.nBufferMBaseSize : qsfaNBufferTail; + mSplitInfo.nBufferStartM = qsfaI * constInfo.nBufferMBaseSize; + + mSplitInfo.vecDealM = (mSplitInfo.nBufferDealM <= 16) ? mSplitInfo.nBufferDealM : + (((mSplitInfo.nBufferDealM + 15) / 16 + 1) / 2 * 16); + mSplitInfo.vecStartM = 0; + if (GetBlockIdx() % 2 == 1) { + mSplitInfo.vecStartM = mSplitInfo.vecDealM; + mSplitInfo.vecDealM = mSplitInfo.nBufferDealM - mSplitInfo.vecDealM; + } + CrossCoreWaitFlag(constInfo.syncC2V2); + ProcessVec2SingleBuf(info, mSplitInfo); + } +} + +template +__aicore__ inline void QSFAVectorService::ProcessVec2Inner(const RunInfo &info, + const MSplitInfo &mSplitInfo, + uint32_t mStartRow, uint32_t mDealSize) +{ + uint32_t qsfaMSplitSize = BASE_BLOCK_MAX_ELEMENT_NUM / constInfo.headDim; + if (qsfaMSplitSize > mDealSize) { + qsfaMSplitSize = mDealSize; + } + + uint32_t qsfaLoopCount = (mDealSize + qsfaMSplitSize - 1) / qsfaMSplitSize; + uint32_t qsfaTailSplitSize = mDealSize - (qsfaLoopCount - 1) * qsfaMSplitSize; + for (uint32_t qsfaI = 0, dealSize = qsfaMSplitSize; qsfaI < qsfaLoopCount; qsfaI++) { + if (qsfaI == (qsfaLoopCount - 1)) { + dealSize = qsfaTailSplitSize; + } + DealBmm2ResBaseBlock(info, mSplitInfo, qsfaI * qsfaMSplitSize + mStartRow, dealSize, + constInfo.headDim, constInfo.headDim); + pingpongFlag ^= 1; // pingpong 0 1切换 + } +} + + +template +__aicore__ inline void QSFAVectorService::GetConfusionTransposeTiling( + int64_t numR, int64_t numC, const uint32_t stackBufferSize, const uint32_t typeSize, + ConfusionTransposeTiling &tiling) +{ + (void)stackBufferSize; + uint32_t qsfaBlockSize = ONE_BLK_SIZE / typeSize; + uint32_t qsfaHeight = numC; + uint32_t qsfaWidth = numR; + uint32_t qsfaHighBlock = qsfaHeight / BLOCK_CUBE; + uint32_t qsfaStride = qsfaHeight * qsfaBlockSize * typeSize / ONE_BLK_SIZE; + uint32_t qsfaRepeat = qsfaWidth / qsfaBlockSize; + + tiling.param0 = qsfaBlockSize; + tiling.param1 = qsfaHeight; + tiling.param2 = qsfaWidth; + tiling.param3 = qsfaHighBlock; + tiling.param4 = qsfaStride; + tiling.param5 = qsfaRepeat; +} + +template +__aicore__ inline void +QSFAVectorService::Bmm2FDDataCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, + uint32_t wsMStart, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + LocalTensor tmp = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + DataCopy(tmp, bmm2ResUb, columnCount * dealRowCount); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + uint64_t accumTmpOutNum = CalcAccumOffset(info.bIdx, info.gS1Idx); + uint64_t offset = accumTmpOutNum * constInfo.kvHeadNum * constInfo.mBaseSize * constInfo.headDim + // taskoffset + info.tndCoreStartKVSplitPos * constInfo.kvHeadNum * constInfo.mBaseSize * constInfo.headDim + // 份数offset + wsMStart * actualColumnCount; // m轴offset + GlobalTensor dst = accumOutGm[offset]; + if (info.actualSingleProcessSInnerSize == 0) { + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = dealRowCount; + dataCopyParams.blockLen = actualColumnCount * sizeof(T); + dataCopyParams.dstStride = 0; + dataCopyParams.srcStride = (columnCount - actualColumnCount) / (BYTE_BLOCK / sizeof(T)); + DataCopyPad(dst, tmp, dataCopyParams); + } else { + matmul::InitOutput(dst, dealRowCount * actualColumnCount, ConstInfo::FLOAT_ZERO); + } + SetFlag(SYNC_OUTPUT_BUF1_FLAG); +} + +template +__aicore__ inline void +QSFAVectorService::Bmm2DataCopyOutTrans(const RunInfo &info, LocalTensor &attenOutUb, + uint32_t wsMStart, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = dealRowCount; + dataCopyParams.blockLen = actualColumnCount * sizeof(OUT_T); + dataCopyParams.srcStride = (columnCount - actualColumnCount) / (BYTE_BLOCK / sizeof(OUT_T)); + dataCopyParams.dstStride = 0; + DataCopyPad(attentionOutGm[info.attenOutOffset + wsMStart * actualColumnCount], attenOutUb, dataCopyParams); + return; +} + +template +__aicore__ inline void +QSFAVectorService::Bmm2CastAndCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, + uint32_t wsMStart, uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount) +{ + LocalTensor qsfaTmpBmm2ResCastTensor = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + if constexpr (IsSameType::value) { // bf16 采取四舍六入五成双模式 + Cast(qsfaTmpBmm2ResCastTensor, bmm2ResUb, AscendC::RoundMode::CAST_RINT, dealRowCount * columnCount); + } else { + Cast(qsfaTmpBmm2ResCastTensor, bmm2ResUb, AscendC::RoundMode::CAST_ROUND, dealRowCount * columnCount); + } + + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + Bmm2DataCopyOutTrans(info, qsfaTmpBmm2ResCastTensor, wsMStart, dealRowCount, columnCount, actualColumnCount); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); +} + +template +__aicore__ inline void +QSFAVectorService::Bmm2ResCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount) +{ + if constexpr (!FLASH_DECODE) { + Bmm2CastAndCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount); + } else { + if (info.tndIsS2SplitCore) { + Bmm2FDDataCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount); + } else { + Bmm2CastAndCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount); + } + } +} + +template +__aicore__ inline void +QSFAVectorService::RowDivs(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount) +{ + // divs by row, 每行的元素除以相同的元素 + // dstUb[i, (j * 8) : (j * 8 + 7)] = src0Ub[i, (j * 8) : (j * 8 + 7)] / src1Ub[i, 0 : 7] + // src0Ub:[dealRowCount, columnCount], src1Ub:[dealRowCount, FP32_BLOCK_ELEMENT_NUM] dstUb:[dealRowCount, + // columnCount] + uint32_t qsfaDtypeMask = FP32_REPEAT_ELEMENT_NUM; + uint32_t qsfaDLoop = actualColumnCount / qsfaDtypeMask; + uint32_t qsfaDRemain = actualColumnCount % qsfaDtypeMask; + + BinaryRepeatParams qsfaRepeatParamsDiv; + qsfaRepeatParamsDiv.src0BlkStride = 1; + qsfaRepeatParamsDiv.src1BlkStride = 0; + qsfaRepeatParamsDiv.dstBlkStride = 1; + qsfaRepeatParamsDiv.src0RepStride = columnCount / FP32_BLOCK_ELEMENT_NUM; + qsfaRepeatParamsDiv.src1RepStride = 1; + qsfaRepeatParamsDiv.dstRepStride = columnCount / FP32_BLOCK_ELEMENT_NUM; + uint32_t qsfaColumnRepeatCount = qsfaDLoop; + if (qsfaColumnRepeatCount <= dealRowCount) { + uint32_t qsfaOffset = 0; + for (uint32_t qsfaI = 0; qsfaI < qsfaDLoop; qsfaI++) { + Div(dstUb[qsfaOffset], src0Ub[qsfaOffset], src1Ub, qsfaDtypeMask, dealRowCount, qsfaRepeatParamsDiv); + qsfaOffset += qsfaDtypeMask; + } + } else { + BinaryRepeatParams qsfaColumnRepeatParams; + qsfaColumnRepeatParams.src0BlkStride = 1; + qsfaColumnRepeatParams.src1BlkStride = 0; + qsfaColumnRepeatParams.dstBlkStride = 1; + qsfaColumnRepeatParams.src0RepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + qsfaColumnRepeatParams.src1RepStride = 0; + qsfaColumnRepeatParams.dstRepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + uint32_t qsfaOffset = 0; + for (uint32_t qsfaI = 0; qsfaI < dealRowCount; qsfaI++) { + Div(dstUb[qsfaOffset], src0Ub[qsfaOffset], src1Ub[qsfaI * FP32_BLOCK_ELEMENT_NUM], qsfaDtypeMask, + qsfaColumnRepeatCount, qsfaColumnRepeatParams); + qsfaOffset += columnCount; + } + } + if (qsfaDRemain > 0) { + Div(dstUb[qsfaDLoop * qsfaDtypeMask], src0Ub[qsfaDLoop * qsfaDtypeMask], src1Ub, qsfaDRemain, + dealRowCount, qsfaRepeatParamsDiv); + } +} + +template +__aicore__ inline void +QSFAVectorService::RowMuls(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount) +{ + // muls by row, 每行的元素乘以相同的元素 + // dstUb[i, (j * 8) : (j * 8 + 7)] = src0Ub[i, (j * 8) : (j * 8 + 7)] * src1Ub[i, 0 : 7] + // src0Ub:[dealRowCount, columnCount] src1Ub:[dealRowCount, FP32_BLOCK_ELEMENT_NUM] dstUb:[dealRowCount, + // columnCount] + // dealRowCount is repeat times, must be less 256 + uint32_t qsfaRepeatElementNum = FP32_REPEAT_ELEMENT_NUM; + uint32_t qsfaBlockElementNum = FP32_BLOCK_ELEMENT_NUM; + + if constexpr (std::is_same::value) { + // 此限制由于每个repeat至多连续读取256B数据 + qsfaRepeatElementNum = FP32_REPEAT_ELEMENT_NUM * 2; // 256/4 * 2=128 + qsfaBlockElementNum = FP32_BLOCK_ELEMENT_NUM * 2; // 32/4 * 2 = 16 + } + + // 每次只能连续读取256B的数据进行计算，故每次只能处理256B/sizeof(dType)= + // 列方向分dLoop次，每次处理8列数据 + uint32_t qsfaDLoop = actualColumnCount / qsfaRepeatElementNum; + uint32_t qsfaDRemain = actualColumnCount % qsfaRepeatElementNum; + // REPEATE_STRIDE_UP_BOUND=256， 此限制由于src0RepStride数据类型为uint8之多256个datablock间距 + if (columnCount < REPEATE_STRIDE_UP_BOUND * qsfaBlockElementNum) { + BinaryRepeatParams qsfaRepeatParams; + qsfaRepeatParams.src0BlkStride = 1; + qsfaRepeatParams.src1BlkStride = 0; + qsfaRepeatParams.dstBlkStride = 1; + qsfaRepeatParams.src0RepStride = columnCount / qsfaBlockElementNum; + qsfaRepeatParams.src1RepStride = 1; + qsfaRepeatParams.dstRepStride = columnCount / qsfaBlockElementNum; + + // 如果以列为repeat所处理的次数小于行处理次数，则以列方式处理。反之则以行进行repeat处理 + if (qsfaDLoop <= dealRowCount) { + uint32_t qsfaOffset = 0; + for (uint32_t qsfaI = 0; qsfaI < qsfaDLoop; qsfaI++) { + Mul(dstUb[qsfaOffset], src0Ub[qsfaOffset], src1Ub, qsfaRepeatElementNum, dealRowCount, + qsfaRepeatParams); + qsfaOffset += qsfaRepeatElementNum; + } + } else { + BinaryRepeatParams qsfaColumnRepeatParams; + qsfaColumnRepeatParams.src0BlkStride = 1; + qsfaColumnRepeatParams.src1BlkStride = 0; + qsfaColumnRepeatParams.dstBlkStride = 1; + qsfaColumnRepeatParams.src0RepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + qsfaColumnRepeatParams.src1RepStride = 0; + qsfaColumnRepeatParams.dstRepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + for (uint32_t qsfaI = 0; qsfaI < dealRowCount; qsfaI++) { + Mul(dstUb[qsfaI * columnCount], src0Ub[qsfaI * columnCount], src1Ub[qsfaI * qsfaBlockElementNum], + qsfaRepeatElementNum, qsfaDLoop, qsfaColumnRepeatParams); + } + } + + // 最后一次完成[dealRowCount, dRemain] * [dealRowCount, blockElementNum] 只计算有效部分 + if (qsfaDRemain > 0) { + Mul(dstUb[qsfaDLoop * qsfaRepeatElementNum], src0Ub[qsfaDLoop * qsfaRepeatElementNum], src1Ub, + qsfaDRemain, dealRowCount, qsfaRepeatParams); + } + } else { + BinaryRepeatParams qsfaRepeatParams; + qsfaRepeatParams.src0RepStride = 8; // 每个repeat为256B数据，正好8个datablock + qsfaRepeatParams.src0BlkStride = 1; + qsfaRepeatParams.src1RepStride = 0; + qsfaRepeatParams.src1BlkStride = 0; + qsfaRepeatParams.dstRepStride = 8; + qsfaRepeatParams.dstBlkStride = 1; + // 每次计算一行，共计算dealRowCount行 + for (uint32_t qsfaI = 0; qsfaI < dealRowCount; qsfaI++) { + // 计算一行中的dLoop个repeat, 每个repeat计算256/block_size 个data_block + Mul(dstUb[qsfaI * columnCount], src0Ub[qsfaI * columnCount], src1Ub[qsfaI * qsfaBlockElementNum], + qsfaRepeatElementNum, qsfaDLoop, qsfaRepeatParams); + // 计算一行中的尾块 + if (qsfaDRemain > 0) { + Mul(dstUb[qsfaI * columnCount + qsfaDLoop * qsfaRepeatElementNum], + src0Ub[qsfaI * columnCount + qsfaDLoop * qsfaRepeatElementNum], + src1Ub[qsfaI * qsfaBlockElementNum], qsfaDRemain, 1, qsfaRepeatParams); + } + } + } +} + +#endif // KV_QUANT_SPARSE_FLASH_ATTENTION_SERVICE_VECTOR_MLA_H diff --git a/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_template_tiling_key.h b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_template_tiling_key.h new file mode 100644 index 000000000..ccfba717c --- /dev/null +++ b/csrc/ascend/attention/kv_quant_sparse_flash_attention/op_kernel/kv_quant_sparse_flash_attention_template_tiling_key.h @@ -0,0 +1,82 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file kv_quant_sparse_flash_attention_template_tiling_key.h + * \brief + */ + +#ifndef KV_QUANT_SPARSE_FLASH_ATTENTION_TEMPLATE_TILING_KEY_H +#define KV_QUANT_SPARSE_FLASH_ATTENTION_TEMPLATE_TILING_KEY_H + +#include "ascendc/host_api/tiling/template_argument.h" + +#define QSFA_LAYOUT_BSND 0 +#define QSFA_LAYOUT_TND 1 +#define QSFA_LAYOUT_PA_BSND 2 + +#define ASCENDC_TPL_4_BW 4 + +#define C_TEMPLATE 0 +#define V_TEMPLATE 1 + +// 模板参数支持的范围定义 +ASCENDC_TPL_ARGS_DECL(KvQuantSparseFlashAttention, // 算子OpType +ASCENDC_TPL_BOOL_DECL(FLASH_DECODE, 0, 1), +ASCENDC_TPL_BOOL_DECL(PAGE_ATTENTION, 0, 1), +ASCENDC_TPL_UINT_DECL(LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, + QSFA_LAYOUT_BSND, QSFA_LAYOUT_TND), +ASCENDC_TPL_UINT_DECL(KV_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, + QSFA_LAYOUT_BSND, QSFA_LAYOUT_TND, QSFA_LAYOUT_PA_BSND), +ASCENDC_TPL_UINT_DECL(TEMPLATE_MODE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, C_TEMPLATE, V_TEMPLATE), +ASCENDC_TPL_BOOL_DECL(IS_SPLIT_G, 0, 1), +); + +// 支持的模板参数组合 +// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法 +ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, QSFA_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, QSFA_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, V_TEMPLATE), + ASCENDC_TPL_BOOL_SEL(IS_SPLIT_G, 0, 1), + ), + + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, QSFA_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, QSFA_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, V_TEMPLATE), + ASCENDC_TPL_BOOL_SEL(IS_SPLIT_G, 0, 1), + ), + + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, QSFA_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, QSFA_LAYOUT_PA_BSND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, V_TEMPLATE), + ASCENDC_TPL_BOOL_SEL(IS_SPLIT_G, 0, 1), + ), + + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, QSFA_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, QSFA_LAYOUT_PA_BSND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, V_TEMPLATE), + ASCENDC_TPL_BOOL_SEL(IS_SPLIT_G, 0, 1), + ), +); + +#endif // TEMPLATE_TILING_KEY \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer/CMakeLists.txt b/csrc/ascend/attention/lightning_indexer/CMakeLists.txt new file mode 100644 index 000000000..549b5c14c --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() diff --git a/csrc/ascend/attention/lightning_indexer/README.md b/csrc/ascend/attention/lightning_indexer/README.md new file mode 100644 index 000000000..23ae7ec47 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/README.md @@ -0,0 +1,269 @@ +# LightningIndexer + +## 产品支持情况 + +|产品 | 是否支持 | +|:----------------------------|:-----------:| +|Ascend 950PR/Ascend 950DT| √ | +|Atlas A3 训练系列产品/Atlas A3 推理系列产品| √ | +|Atlas A2 训练系列产品/Atlas A2 推理系列产品| √ | +|Atlas 200I/500 A2 推理产品| × | +|Atlas 推理系列产品| × | +|Atlas 训练系列产品| × | + +## 功能说明 + +- 算子功能：`lightning_indexer`基于一系列操作得到每一个token对应的Top-$k$个位置。 + +- 计算公式： + + $$ + Indices=\text{Top-}k\left\{[1]_{1\times g}@\left[(W@[1]_{1\times S_{k}})\odot\text{ReLU}\left(Q_{index}@K_{index}^T\right)\right]\right\} + $$ + + 对于某个token对应的Index Query $Q_{index}\in\R^{g\times d}$，给定上下文Index Key $K_{index}\in\R^{S_{k}\times d},W\in\R^{g\times 1}$，其中$g$为GQA对应的group size，$d$为每一个头的维度，$S_{k}$是上下文的长度。 + +## 参数说明 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数名输入/输出/属性描述数据类型数据格式
query输入 +
    +
  • 公式中的输入Q。
    • +
  • 不支持空tensor和非连续。
    • +
  • layout_query为BSND时，shape为(B,S1,N1,D)；layout_query为TND时，shape为(T1,N1,D)。
    • +
+ 		FLOAT16、BFLOAT16ND
key输入 +
    +
  • 公式中的输入K。
    • +
  • 不支持空tensor和非连续。
    • +
  • layout_key为PA_BSND时，shape为(block_num, block_size, N2, D)，其中block_num为PageAttention时block总数、block_size为一个block的token数；layout_key为BSND时，shape为(B, S2, N2, D)；layout_key为TND时，shape为(T2, N2, D)。
    • +
+ 		FLOAT16、BFLOAT16ND
weights输入 +
    +
  • 公式中的输入W。
    • +
  • 不支持空tensor和非连续。
    • +
  • layout_query为BSND时，shape为(B,S1,N1)；layout_query为TND时，shape为(T1,N1)。
    • +
+ 		FLOAT16、BFLOAT16、FLOATND
actual_seq_lengths_query输入 +
    +
  • 每个Batch中Query的有效token数。
    • +
  • 不支持空tensor和非连续。
    • +
  • 可传入None表示与query的S长度相同。
    • +
  • 支持长度为B的一维tensor，且每个Batch的有效token数不超过query中的维度S大小且不小于0。layout_query为TND时该入参必须传入，并以元素数量作为B值。
    • +
  • 每个元素表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。
    • +
+ 		INT32ND
actual_seq_lengths_key输入 +
    +
  • 每个Batch中Key的有效token数。
    • +
  • 不支持空tensor和非连续。
    • +
  • 可传入None表示与key的S长度相同。
    • +
  • 支持长度为B的一维tensor，且每个Batch的有效token数不超过key/value中的维度S大小且不小于0。
    • +
  • layout_key为TND或PA_BSND时该入参必须传入；其中layout_key为TND时，每个元素表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。
    • +
+ 		INT32ND
block_table输入 +
    +
  • 表示PageAttention中KV存储使用的block映射表。
    • +
  • 不支持空tensor和非连续。
    • +
  • PageAttention场景下，block_table必须为二维，第一维长度需要等于B，第二维长度不能小于maxBlockNumPerSeq（每个batch中最大actual_seq_lengths_key对应的block数量）。
    • +
  • shape支持(B,S2/block_size)。
    • +
+ 		INT32ND
layout_query属性 +
    +
  • 用于标识输入Query的数据排布格式。
    • +
  • 默认值为"BSND"，当前支持BSND、TND。
    • +
+ 		STRING-
layout_key属性 +
    +
  • 用于标识输入Key的数据排布格式。
    • +
  • 默认值为"BSND"，当前支持PA_BSND、BSND、TND。
    • +
+ 		STRING-
sparse_count属性 +
    +
  • topK阶段需要保留的block数量。
    • +
  • 支持[1, 2048]，以及3072、4096、5120、6144、7168、8192。
    • +
  • 默认值为2048。
    • +
+ 		INT32-
sparse_mode属性 +
    +
  • 表示sparse的模式。
    • +
  • sparse_mode为0时代表defaultMask模式。
    • +
  • sparse_mode为3时代表rightDownCausal模式的mask，对应以右顶点为划分的下三角场景。
    • +
  • 默认值为3。
    • +
+ 		INT32-
pre_tokens属性用于稀疏计算，表示attention需要和前几个Token计算关联，仅支持默认值2^63-1。INT64-
next_tokens属性用于稀疏计算，表示attention需要和后几个Token计算关联，仅支持默认值2^63-1。INT64-
return_values属性 +
    +
  • 表示是否输出sparseValuesOut。
    • +
  • True表示输出，False表示不输出，默认值为False。仅在训练且layout_key不为PA_BSND场景支持。
    • +
+ 		BOOL-
sparse_indices输出 +
    +
  • 公式中的Indices输出。
    • +
  • 不支持空tensor和非连续。
    • +
  • layout_query为BSND时输出shape为[B, S1, N2, sparseCount]；layout_query为TND时输出shape为[T1, N2, sparseCount]。
    • +
+ 		INT32-
sparse_values输出 +
    +
  • 公式中的Indices输出对应的value值。
    • +
  • 不支持空tensor和非连续。
    • +
  • shape与sparseIndicesOut保持一致。
    • +
+ 		FLOAT16、BFLOAT16ND
+ +## 约束说明 + +- 该接口支持图模式。 +- 参数key中的N2支持1。 +- headdim支持128。 +- block_size取值为16的倍数，最大支持1024。 +- 参数query、key的数据类型应保持一致。 +- 参数weights不为`float32`时，参数query、key、weights的数据类型应保持一致。 +- Ascend 950PR/Ascend 950DT： + - query N1仅支持8、16、24、32、64。 + - 参数weights不支持`float32`类型。 +- A3 训练系列产品/Atlas A3 推理系列产品： + - query N1支持小于等于64。 + +## 调用示例 + + + + + + + + + + + + + +
调用方式样例代码说明
aclnn接口 + test_aclnn_lightning_indexer + + + 通过 + aclnnLightningIndexer + + 接口方式调用算子 +
diff --git a/csrc/ascend/attention/lightning_indexer/docs/aclnnLightningIndexer.md b/csrc/ascend/attention/lightning_indexer/docs/aclnnLightningIndexer.md new file mode 100644 index 000000000..9f2ff533e --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/docs/aclnnLightningIndexer.md @@ -0,0 +1,740 @@ +# aclnnLightningIndexer + +[📄 查看源码](https://gitcode.com/cann/ops-transformer/tree/master/attention/lightning_indexer) + +## 产品支持情况 + +|产品 | 是否支持 | +|:----------------------------|:-----------:| +|Ascend 950PR/Ascend 950DT| × | +|Atlas A3 训练系列产品/Atlas A3 推理系列产品| √ | +|Atlas A2 训练系列产品/Atlas A2 推理系列产品| √ | +|Atlas 200I/500 A2 推理产品| × | +|Atlas 推理系列产品| × | +|Atlas 训练系列产品| × | + +## 功能说明 + +- 接口功能：`lightning_indexer`基于一系列操作得到每一个token对应的Top-$k$个位置。 + +- 计算公式： + +$$ +Indices=\text{Top-}k\left\{[1]_{1\times g}@\left[(W@[1]_{1\times S_{k}})\odot\text{ReLU}\left(Q_{index}@K_{index}^T\right)\right]\right\} +$$ + +对于某个token对应的Index Query $Q_{index}\in\R^{g\times d}$，给定上下文Index Key $K_{index}\in\R^{S_{k}\times d},W\in\R^{g\times 1}$，其中$g$为GQA对应的group size，$d$为每一个头的维度，$S_{k}$是上下文的长度。 + +## 函数原型 + +每个算子分为[两段式接口](../../../docs/zh/context/两段式接口.md)，必须先调用“aclnnLightningIndexerGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnLightningIndexer”接口执行计算。 + +```Cpp +aclnnStatus aclnnLightningIndexerGetWorkspaceSize( + const aclTensor *query, + const aclTensor *key, + const aclTensor *weights, + const aclTensor *actualSeqLengthsQueryOptional, + const aclTensor *actualSeqLengthsKeyOptional, + const aclTensor *blockTableOptional, + char *layoutQueryOptional, + char *layoutKeyOptional, + int64_t sparseCount, + int64_t sparseMode, + int64_t preTokens, + int64_t nextTokens, + bool returnValues, + const aclTensor *sparseIndicesOut, + const aclTensor *sparseValuesOut, + uint64_t *workspaceSize, + aclOpExecutor **executor) +``` + +```Cpp +aclnnStatus aclnnLightningIndexer( + void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + const aclrtStream stream) +``` + +## aclnnLightningIndexerGetWorkspaceSize + +- **参数说明：** + +> [!NOTE] +> +> - query、key、weights参数维度含义：B（Batch Size）表示输入样本批量大小、S（Sequence Length）表示输入样本序列长度、H（Head Size）表示hidden层的大小、N（Head Num）表示多头数、D（Head Dim）表示hidden层最小的单元尺寸，且满足D=H/N、T表示所有Batch输入样本序列长度的累加和。 +> - S1表示query shape中的S，S2表示key shape中的S，T1表示query shape中的T，T2表示key shape中的T，N1表示query shape中的N，N2表示key shape中的N。 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensor
query输入公式中的输入Q。不支持空tensor。FLOAT16、BFLOAT16ND +
    +
  • layout_query为BSND时，shape为(B,S1,N1,D)。
    • +
  • layout_query为TND时，shape为(T1,N1,D)。
    • +
+ 		x
key输入公式中的输入K。 +
    +
  • 不支持空tensor。
    • +
  • block_num为PageAttention时block总数，block_size为一个block的token数。
    • +
+ 		FLOAT16、BFLOAT16ND +
    +
  • layout_key为PA_BSND时，shape为(block_num, block_size, N2, D)。
    • +
  • layout_kv为BSND时，shape为(B, S2, N2, D)。
    • +
  • layout_kv为TND时，shape为(T2, N2, D)。
    • +
+ 		x
weights输入公式中的输入W。不支持空tensor。FLOAT16、BFLOAT16、FLOATND +
    +
  • layout_query为BSND时，shape为(B,S1,N1)。
    • +
  • layout_query为TND时，shape为(T1,N1)。
    • +
+ 		x
actualSeqLengthsQueryOptional输入每个Batch中，Query的有效token数。 +
    +
  • 不支持空tensor。
    • +
  • 如果不指定seqlen可传入None，表示和`query`的shape的S长度相同。
    • +
  • 该入参中每个Batch的有效token数不超过`query`中的维度S大小且不小于0，支持长度为B的一维tensor。
    • +
  • 当`layout_query`为TND时，该入参必须传入，且以该入参元素的数量作为B值，该入参中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。
    • +
+ 		INT32ND(B,)x
actualSeqLengthsKeyOptional输入每个Batch中，Key的有效token数。 +
    +
  • 不支持空tensor。
    • +
  • 如果不指定seqlen可传入None，表示和key的shape的S长度相同。
    • +
  • 该参数中每个Batch的有效token数不超过`key/value`中的维度S大小且不小于0，支持长度为B的一维tensor。
    • +
  • 当`layout_key`为TND或PA_BSND时，该入参必须传入，`layout_key`为TND，该参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。
    • +
+ 		INT32ND(B,)x
blockTableOptional输入表示PageAttention中KV存储使用的block映射表。 +
    +
  • 不支持空tensor。
    • +
  • PageAttention场景下，block\_table必须为二维，第一维长度需要等于B，第二维长度不能小于maxBlockNumPerSeq（maxBlockNumPerSeq为每个batch中最大actual\_seq\_lengths\_key对应的block数量）
    • +
+ 		INT32NDshape支持(B,S2/block_size)x
layoutQueryOptional输入用于标识输入Query的数据排布格式。 +
    +
  • 用户不特意指定时可传入默认值"BSND"。
    • +
  • 当前支持BSND、TND。
    • +
+ 		STRING---
layoutKeyOptional输入用于标识输入Key的数据排布格式。 +
    +
  • 用户不特意指定时可传入默认值"BSND"。
    • +
  • 当前支持PA_BSND、BSND、TND。
    • +
+ 		STRING---
sparseCount输入topK阶段需要保留的block数量。支持[1, 2048]，以及3072、4096、5120、6144、7168、8192INT32---
sparseMode输入表示sparse的模式。 +
    +
  • sparse_mode为0时，代表defaultMask模式。
    • +
  • sparse_mode为3时，代表rightDownCausal模式的mask，对应以右顶点为划分的下三角场景。
    • +
+ 		INT32---
preTokens输入用于稀疏计算，表示attention需要和前几个Token计算关联。仅支持默认值2^63-1。INT64---
nextTokens输入用于稀疏计算，表示attention需要和后几个Token计算关联。仅支持默认值2^63-1。INT64---
returnValues输入表示是否输出sparseValuesOut。 +
    +
  • True表示输出，但图模式下不支持，False表示不输出；默认值为False
    • +
  • 仅在训练且layout_key不为PA_BSND场景支持
    • +
+ 		BOOL---
sparseIndicesOut输出公式中的Indices输出。不支持空tensor。INT32- +
    +
  • layout_query为"BSND"时输出shape为[B, S1, N2, sparseCount]。
    • +
  • layout_query为"TND"时输出shape为[T1, N2, sparseCount]。
    • +
+ 		x
sparseValuesOut输出公式中的Indices输出对应的value值。不支持空tensor。FLOAT16、BFLOAT16NDshape与sparseIndicesOut保持一致x
workspaceSize输出返回需要在Device侧申请的workspace大小。-----
executor输出返回op执行器，包含了算子计算流程。-----
+ +- **返回值：** + + aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn返回码.md)。 + + 第一段接口会完成入参校验，出现以下场景时报错： + + + + + + + + + + + + + + + + + + + + + + + +
返回值错误码描述
ACLNN_ERR_PARAM_NULLPTR161001如果传入参数是必选输入，输出或者必选属性，且是空指针，则返回161001。
ACLNN_ERR_PARAM_INVALID161002query、key、weights、actualSeqLengthsQueryOptional、actualSeqLengthsKeyOptional、layoutQueryOptional、layoutKeyOptional、sparseCount、sparseMode、returnValues、sparseIndicesOut、sparseValuesOut的数据类型和数据格式不在支持的范围内。
+ +## aclnnLightningIndexer + +- **参数说明：** + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数名输入/输出描述
workspace输入在Device侧申请的workspace内存地址。
workspaceSize输入在Device侧申请的workspace大小，由第一段接口aclnnLightningIndexerGetWorkspaceSize获取。
executor输入op执行器，包含了算子计算流程。
stream输入指定执行任务的Stream。
+ +- **返回值：** + + aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn返回码.md)。 + +## 约束说明 + +- 参数query中的N支持小于等于64，key的N支持1。 +- headdim支持128。 +- block_size取值为16的倍数，最大支持1024。 +- 参数query、key的数据类型应保持一致。 +- 参数weights不为`float32`时，参数query、key、weights的数据类型应保持一致。 + +## 调用示例 + +示例代码如下，仅供参考，具体编译和执行过程请参考[编译与运行样例](../../../docs/zh/context/编译与运行样例.md)。 + +```Cpp +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file test_incre_flash_attention_v4.cpp + * \brief + */ +//testci +#include +#include +#include +#include +#include "securec.h" +#include "acl/acl.h" +#include "aclnnop/aclnn_lightning_indexer.h" + +using namespace std; + +namespace { + +#define CHECK_RET(cond) ((cond) ? true :(false)) + +#define LOG_PRINT(message, ...) \ + do { \ + (void)printf(message, ##__VA_ARGS__); \ + } while (0) + +int64_t GetShapeSize(const std::vector& shape) { + int64_t shapeSize = 1; + for (auto i : shape) { + shapeSize *= i; + } + return shapeSize; +} + +int Init(int32_t deviceId, aclrtStream* stream) { + auto ret = aclInit(nullptr); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclInit failed. ERROR: %d\n", ret); + return ret; + } + ret = aclrtSetDevice(deviceId); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); + return ret; + } + ret = aclrtCreateStream(stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); + return ret; + } + return 0; +} + +template +int CreateAclTensor(const std::vector& hostData, const std::vector& shape, void** deviceAddr, + aclDataType dataType, aclTensor** tensor) { + auto size = GetShapeSize(shape) * sizeof(T); + auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); + return ret; + } + + ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); + return ret; + } + + std::vector strides(shape.size(), 1); + for (int64_t i = shape.size() - 2; i >= 0; i--) { + strides[i] = shape[i + 1] * strides[i + 1]; + } + + *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, + shape.data(), shape.size(), *deviceAddr); + return 0; +} + +struct TensorResources { + void* queryDeviceAddr = nullptr; + void* keyDeviceAddr = nullptr; + void* weightsDeviceAddr = nullptr; + void* sparseIndicesDeviceAddr = nullptr; + void* sparseValuesDeviceAddr = nullptr; + + aclTensor* queryTensor = nullptr; + aclTensor* keyTensor = nullptr; + aclTensor* weightsTensor = nullptr; + aclTensor* sparseIndicesTensor = nullptr; + aclTensor* sparseValuesTensor = nullptr; +}; + +int InitializeTensors(TensorResources& resources) { + std::vector queryShape = {1, 2, 1, 128}; + std::vector keyShape = {1, 2, 1, 128}; + std::vector weightsShape = {1, 2, 1}; + std::vector sparseIndicesShape = {1, 2, 1, 2048}; + std::vector sparseValuesShape = {1, 2, 1, 2048}; + + int64_t queryShapeSize = GetShapeSize(queryShape); + int64_t keyShapeSize = GetShapeSize(keyShape); + int64_t weightsShapeSize = GetShapeSize(weightsShape); + int64_t sparseIndicesShapeSize = GetShapeSize(sparseIndicesShape); + int64_t sparseValuesShapeSize = GetShapeSize(sparseValuesShape); + + std::vector queryHostData(queryShapeSize, 1); + std::vector keyHostData(keyShapeSize, 1); + std::vector weightsHostData(weightsShapeSize, 1); + std::vector sparseIndicesHostData(sparseIndicesShapeSize, 1); + std::vector sparseValuesHostData(sparseValuesShapeSize, 1); + + int ret = CreateAclTensor(queryHostData, queryShape, &resources.queryDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.queryTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + ret = CreateAclTensor(keyHostData, keyShape, &resources.keyDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.keyTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + ret = CreateAclTensor(weightsHostData, weightsShape, &resources.weightsDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.weightsTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + ret = CreateAclTensor(sparseIndicesHostData, sparseIndicesShape, &resources.sparseIndicesDeviceAddr, + aclDataType::ACL_INT32, &resources.sparseIndicesTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + ret = CreateAclTensor(sparseValuesHostData, sparseValuesShape, &resources.sparseValuesDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.sparseValuesTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + return ACL_SUCCESS; +} + +int ExecuteLightningIndexer(TensorResources& resources, aclrtStream stream, + void** workspaceAddr, uint64_t* workspaceSize) { + int64_t sparseCount = 2048; + int64_t sparseMode = 3; + int64_t preTokens = 9223372036854775807; + int64_t nextTokens = 9223372036854775807; + bool returnValue = true; + constexpr const char layerOutStr[] = "BSND"; + constexpr size_t layerOutLen = sizeof(layerOutStr); + char layoutQuery[layerOutLen]; + char layoutKey[layerOutLen]; + errno_t memcpyRet = memcpy_s(layoutQuery, sizeof(layoutQuery), layerOutStr, layerOutLen); + if (!CHECK_RET(memcpyRet == 0)) { + LOG_PRINT("memcpy_s layoutQuery failed. ERROR: %d\n", memcpyRet); + return -1; + } + memcpyRet = memcpy_s(layoutKey, sizeof(layoutKey), layerOutStr, layerOutLen); + if (!CHECK_RET(memcpyRet == 0)) { + LOG_PRINT("memcpy_s layoutKey failed. ERROR: %d\n", memcpyRet); + return -1; + } + aclOpExecutor* executor; + + int ret = aclnnLightningIndexerGetWorkspaceSize(resources.queryTensor, resources.keyTensor, resources.weightsTensor, nullptr, nullptr, nullptr, + layoutQuery, layoutKey, sparseCount, sparseMode, preTokens, nextTokens,returnValue, + resources.sparseIndicesTensor, resources.sparseValuesTensor, workspaceSize, &executor); + + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclnnLightningIndexerGetWorkspaceSize failed. ERROR: %d\n", ret); + return ret; + } + + if (*workspaceSize > 0ULL) { + ret = aclrtMalloc(workspaceAddr, *workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); + return ret; + } + } + + ret = aclnnLightningIndexer(*workspaceAddr, *workspaceSize, executor, stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclnnLightningIndexer failed. ERROR: %d\n", ret); + return ret; + } + + return ACL_SUCCESS; +} + +int PrintValueOutResult(std::vector &shape, void** deviceAddr) { + auto size = GetShapeSize(shape); + std::vector resultData(size, 0); + auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), + *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); + return ret; + } + for (int64_t i = 0; i < size; i++) { + LOG_PRINT("mean result[%ld] is: %f\n", i, aclFloat16ToFloat(resultData[i])); + } + return ACL_SUCCESS; +} + +int PrintIndicesOutResult(std::vector &shape, void** deviceAddr) { + auto size = GetShapeSize(shape); + std::vector resultData(size, 0); + auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), + *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); + return ret; + } + for (int64_t i = 0; i < size; i++) { + LOG_PRINT("mean result[%ld] is: %d\n", i, resultData[i]); + } + return ACL_SUCCESS; +} + +void CleanupResources(TensorResources& resources, void* workspaceAddr, + aclrtStream stream, int32_t deviceId) { + if (resources.queryTensor) { + aclDestroyTensor(resources.queryTensor); + } + if (resources.keyTensor) { + aclDestroyTensor(resources.keyTensor); + } + if (resources.weightsTensor) { + aclDestroyTensor(resources.weightsTensor); + } + if (resources.sparseIndicesTensor) { + aclDestroyTensor(resources.sparseIndicesTensor); + } + if (resources.sparseValuesTensor) { + aclDestroyTensor(resources.sparseValuesTensor); + } + + if (resources.queryDeviceAddr) { + aclrtFree(resources.queryDeviceAddr); + } + if (resources.keyDeviceAddr) { + aclrtFree(resources.keyDeviceAddr); + } + if (resources.weightsDeviceAddr) { + aclrtFree(resources.weightsDeviceAddr); + } + if (resources.sparseIndicesDeviceAddr) { + aclrtFree(resources.sparseIndicesDeviceAddr); + } + if (resources.sparseValuesDeviceAddr) { + aclrtFree(resources.sparseValuesDeviceAddr); + } + + if (workspaceAddr) { + aclrtFree(workspaceAddr); + } + if (stream) { + aclrtDestroyStream(stream); + } + aclrtResetDevice(deviceId); + aclFinalize(); +} + +} // namespace + +int main() { + int32_t deviceId = 0; + aclrtStream stream = nullptr; + TensorResources resources = {}; + void* workspaceAddr = nullptr; + uint64_t workspaceSize = 0; + std::vector sparseIndicesShape = {1, 2, 1, 2048}; + std::vector sparseValuesShape = {1, 2, 1, 2048}; + int ret = ACL_SUCCESS; + + // 1. Initialize device and stream + ret = Init(deviceId, &stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("Init acl failed. ERROR: %d\n", ret); + return ret; + } + + // 2. Initialize tensors + ret = InitializeTensors(resources); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + // 3. Execute the operation + ret = ExecuteLightningIndexer(resources, stream, &workspaceAddr, &workspaceSize); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + // 4. Synchronize stream + ret = aclrtSynchronizeStream(stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + // 5. Process results + PrintIndicesOutResult(sparseIndicesShape, &resources.sparseIndicesDeviceAddr); + PrintValueOutResult(sparseValuesShape, &resources.sparseValuesDeviceAddr); + + // 6. Cleanup resources + CleanupResources(resources, workspaceAddr, stream, deviceId); + return 0; +} +``` diff --git a/csrc/ascend/attention/lightning_indexer/lightning_indexer_torch_adpt.h b/csrc/ascend/attention/lightning_indexer/lightning_indexer_torch_adpt.h new file mode 100644 index 000000000..c2ebe11cb --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/lightning_indexer_torch_adpt.h @@ -0,0 +1,106 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef LIGHTNING_INDEXER_TORCH_ADPT_H +#define LIGHTNING_INDEXER_TORCH_ADPT_H + +namespace vllm_fl { + +std::tuple construct_lightning_indexer_output_tensor( + const at::Tensor& query, const at::Tensor& key, int64_t sparse_count, + const std::string& query_layout_str, const std::string& key_layout_str, + bool return_value) +{ + constexpr int64_t SIZE = 8; + constexpr int64_t DIM_0 = 0; + constexpr int64_t DIM_1 = 1; + constexpr int64_t DIM_2 = 2; + + at::SmallVector output_size; + for (size_t i = 0; i < query.sizes().size(); i++) { + TORCH_CHECK(query.size(i) > 0, + "All values within query's shape should be greater " + "than 0, but shape[", + i, "] is ", query.size(i)); + } + for (size_t i = 0; i < key.sizes().size(); i++) { + TORCH_CHECK(key.size(i) > 0, + "All values within key's shape should be greater " + "than 0, but shape[", + i, "] is ", key.size(i)); + } + TORCH_CHECK(sparse_count > 0, + "sparse count should be greater than 0, but now is ", + sparse_count); + + if (query_layout_str == "BSND") { + output_size = {query.size(DIM_0), query.size(DIM_1), + key.size(DIM_2), sparse_count}; + } else { + int64_t n_dim_index = (key_layout_str == "TND") ? DIM_1 : DIM_2; + output_size = {query.size(DIM_0), key.size(n_dim_index), + sparse_count}; + } + + at::Tensor sparse_indices_out = + at::empty(output_size, query.options().dtype(at::kInt)); + at::Tensor sparse_values_out; + if (return_value) { + sparse_values_out = + at::empty(output_size, query.options().dtype(query.dtype())); + } else { + sparse_values_out = at::empty({0}, query.options().dtype(query.dtype())); + } + + return std::tuple(sparse_indices_out, + sparse_values_out); +} + +std::tuple npu_lightning_indexer( + const at::Tensor& query, const at::Tensor& key, const at::Tensor& weights, + const c10::optional& actual_seq_lengths_query, + const c10::optional& actual_seq_lengths_key, + const c10::optional& block_table, c10::string_view layout_query, + c10::string_view layout_key, int64_t sparse_count, int64_t sparse_mode, + int64_t pre_tokens, int64_t next_tokens, bool return_value) +{ + TORCH_CHECK(query.numel() > 0, "Tensor query is empty."); + TORCH_CHECK(key.numel() > 0, "Tensor key is empty."); + TORCH_CHECK(weights.numel() > 0, "Tensor weights is empty."); + + std::string query_layout_str = std::string(layout_query); + std::string key_layout_str = std::string(layout_key); + + auto lightning_indexer_output = construct_lightning_indexer_output_tensor( + query, key, sparse_count, query_layout_str, key_layout_str, + return_value); + at::Tensor sparse_indices_out = std::get<0>(lightning_indexer_output); + at::Tensor sparse_values_out = std::get<1>(lightning_indexer_output); + + char* query_layout_ptr = const_cast(query_layout_str.c_str()); + char* key_layout_ptr = const_cast(key_layout_str.c_str()); + + EXEC_NPU_CMD(aclnnLightningIndexer, query, key, weights, + actual_seq_lengths_query, actual_seq_lengths_key, block_table, + query_layout_ptr, key_layout_ptr, sparse_count, sparse_mode, + pre_tokens, next_tokens, return_value, sparse_indices_out, + sparse_values_out); + + return std::tuple(sparse_indices_out, + sparse_values_out); +} +} // namespace vllm_fl + +#endif // LIGHTNING_INDEXER_TORCH_ADPT_H diff --git a/csrc/ascend/attention/lightning_indexer/op_host/CMakeLists.txt b/csrc/ascend/attention/lightning_indexer/op_host/CMakeLists.txt new file mode 100644 index 000000000..9d5dcb5d6 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_host/CMakeLists.txt @@ -0,0 +1,29 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnInner PRIVATE + lightning_indexer_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME LightningIndexer + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-vf-remove-membar=false + -mllvm -cce-aicore-hoist-movemask=false +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE lightning_indexer ACLNNTYPE aclnn_inner) +endif() diff --git a/csrc/ascend/attention/lightning_indexer/op_host/lightning_indexer_def.cpp b/csrc/ascend/attention/lightning_indexer/op_host/lightning_indexer_def.cpp new file mode 100644 index 000000000..9014d9691 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_host/lightning_indexer_def.cpp @@ -0,0 +1,80 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_def.cpp + * \brief + */ +#include +#include "register/op_def_registry.h" + +namespace ops { +class LightningIndexer : public OpDef { +public: + explicit LightningIndexer(const char *name) : OpDef(name) + { + this->Input("query") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("key") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("weights") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_FLOAT}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("actual_seq_lengths_query") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("actual_seq_lengths_key") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("block_table") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("sparse_indices") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}); + this->Output("sparse_values") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}); + this->Attr("layout_query").AttrType(OPTIONAL).String("BSND"); + this->Attr("layout_key").AttrType(OPTIONAL).String("BSND"); + this->Attr("sparse_count").AttrType(OPTIONAL).Int(2048); // 2048:默认值，筛选前2048 + this->Attr("sparse_mode").AttrType(OPTIONAL).Int(3); // 3:默认值，只计算下三角 + this->Attr("pre_tokens").AttrType(OPTIONAL).Int(INT64_MAX); + this->Attr("next_tokens").AttrType(OPTIONAL).Int(INT64_MAX); + this->Attr("return_values").AttrType(OPTIONAL).Bool(false); + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true); + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + } +}; +OP_ADD(LightningIndexer); +} // namespace ops diff --git a/csrc/ascend/attention/lightning_indexer/op_host/lightning_indexer_infershape.cpp b/csrc/ascend/attention/lightning_indexer/op_host/lightning_indexer_infershape.cpp new file mode 100644 index 000000000..5b938ae7c --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_host/lightning_indexer_infershape.cpp @@ -0,0 +1,112 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_infershape.cpp + * \brief + */ +#include +#include +#include "err/ops_err.h" + + +using namespace ge; + +namespace ops { +constexpr uint32_t QUERY_INDEX = 0; +constexpr uint32_t KEY_INDEX = 1; +constexpr uint32_t ACTUAL_SEQ_K_INDEX = 4; +constexpr uint32_t ATTR_QUERY_LAYOUT_INDEX = 0; +constexpr uint32_t ATTR_KEY_LAYOUT_INDEX = 1; +constexpr uint32_t ATTR_SPARSE_COUNT_INDEX = 2; +constexpr uint32_t ATTR_RETURN_VALUE_INDEX = 6; + +static ge::graphStatus InferShapeLightningIndexer(gert::InferShapeContext *context) +{ + OP_CHECK_IF(context == nullptr, OP_LOGE("LightningIndexer", "InferShapeContext is nullptr!"), + return ge::GRAPH_FAILED); + const gert::Shape *queryShape = context->GetInputShape(QUERY_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, queryShape); + const gert::Shape *keyShape = context->GetInputShape(KEY_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, keyShape); + + gert::Shape *sparseIndicesShape = context->GetOutputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context, sparseIndicesShape); + gert::Shape *sparseValuesShape = context->GetOutputShape(1); + OP_CHECK_NULL_WITH_CONTEXT(context, sparseValuesShape); + + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const char *inputLayoutQueryPtr = attrs->GetAttrPointer(ATTR_QUERY_LAYOUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, inputLayoutQueryPtr); + const char *inputLayoutKeyPtr = attrs->GetAttrPointer(ATTR_KEY_LAYOUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, inputLayoutKeyPtr); + const int64_t *seleced_count = attrs->GetInt(ATTR_SPARSE_COUNT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, seleced_count); + std::string inputLayoutQueryPtrStr = std::string(inputLayoutQueryPtr); + std::string inputLayoutKeyPtrStr = std::string(inputLayoutKeyPtr); + OP_CHECK_IF( + inputLayoutQueryPtrStr != "TND" && inputLayoutQueryPtrStr != "BSND", + OP_LOGE(context, "The attr layout_query should be TND or BSND, but got %s.", inputLayoutQueryPtrStr.c_str()), + return ge::GRAPH_FAILED); + + sparseIndicesShape->SetDimNum(queryShape->GetDimNum()); + if (inputLayoutQueryPtrStr == "BSND") { + OP_CHECK_IF( + queryShape->GetDimNum() != 4, + OP_LOGE(context, "Layout BSND, queryDims (%zu) must be 4!", queryShape->GetDimNum()), + return ge::GRAPH_FAILED); + sparseIndicesShape->SetDim(0, queryShape->GetDim(0)); // 0:Dim B + sparseIndicesShape->SetDim(1, queryShape->GetDim(1)); // 1:Dim S + sparseIndicesShape->SetDim(2, keyShape->GetDim(2)); // 2:Dim N + sparseIndicesShape->SetDim(3, *seleced_count); // 3:Dim K + } else { + OP_CHECK_IF( + queryShape->GetDimNum() != 3, + OP_LOGE(context, "Layout TND, queryDims (%zu) must be 3!", queryShape->GetDimNum()), + return ge::GRAPH_FAILED); + sparseIndicesShape->SetDim(0, queryShape->GetDim(0)); // 0:Dim T + int32_t nDimIndex = (inputLayoutKeyPtrStr == "PA_BSND") ? 2 : 1; // 2:Key Dim N + sparseIndicesShape->SetDim(1, keyShape->GetDim(nDimIndex)); // 1:Dim N + sparseIndicesShape->SetDim(2, *seleced_count); // 2:Dim K + } + + const bool *return_value = attrs->GetAttrPointer(ATTR_RETURN_VALUE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, return_value); + bool returnValueFlag = (return_value != nullptr) ? *return_value : false; + if (returnValueFlag) { + *sparseValuesShape = *sparseIndicesShape; + } else { + sparseValuesShape->SetDimNum(1); + sparseValuesShape->SetDim(0, 0); + } + OP_LOGI(context->GetNodeName(), "LightningIndexer InferShape end."); + + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataTypeLightningIndexer(gert::InferDataTypeContext *context) +{ + OP_CHECK_IF(context == nullptr, OP_LOGE("LightningIndexer", "InferDataTypeContext is nullptr!"), + return ge::GRAPH_FAILED); + OP_LOGI(context->GetNodeName(), "Enter LightningIndexer InferDataType impl."); + // default set q's dtype as fia's output type + ge::DataType outputType = ge::DT_INT32; + // attention_out, outidx:0 + context->SetOutputDataType(0, outputType); + context->SetOutputDataType(1, context->GetInputDataType(QUERY_INDEX)); + OP_LOGI(context->GetNodeName(), "LightningIndexer InferDataType end."); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(LightningIndexer) + .InferShape(InferShapeLightningIndexer) + .InferDataType(InferDataTypeLightningIndexer); +} // namespace ops diff --git a/csrc/ascend/attention/lightning_indexer/op_host/lightning_indexer_tiling.cpp b/csrc/ascend/attention/lightning_indexer/op_host/lightning_indexer_tiling.cpp new file mode 100644 index 000000000..37b41636d --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_host/lightning_indexer_tiling.cpp @@ -0,0 +1,832 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_tiling.cpp + * \brief + */ + +#include "lightning_indexer_tiling.h" +#include "../op_kernel/lightning_indexer_template_tiling_key.h" + +using namespace ge; +using namespace AscendC; +using std::map; +using std::string; +namespace optiling { +// --------------------------LIInfoParser类成员函数定义------------------------------------- +ge::graphStatus LIInfoParser::CheckRequiredInOutExistence() const +{ + OP_CHECK_IF(opParamInfo_.query.shape == nullptr, OP_LOGE(opName_, "Shape of tensor query is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.query.desc == nullptr, OP_LOGE(opName_, "Desc of tensor query is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key.shape == nullptr, OP_LOGE(opName_, "Shape of tensor k is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key.desc == nullptr, OP_LOGE(opName_, "Desc of tensor k is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.weights.shape == nullptr, OP_LOGE(opName_, "Shape of tensor value is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.weights.desc == nullptr, OP_LOGE(opName_, "Desc of tensor value is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.attenOut.shape == nullptr, OP_LOGE(opName_, "Shape of tensor output is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.attenOut.desc == nullptr, OP_LOGE(opName_, "Desc of tensor output is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.valuesOut.shape == nullptr, OP_LOGE(opName_, "Shape of tensor output values is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.valuesOut.desc == nullptr, OP_LOGE(opName_, "Desc of tensor output values is nullptr"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::CheckRequiredAttrExistence() const +{ + OP_CHECK_IF(opParamInfo_.layOut == nullptr, OP_LOGE(opName_, "attr layout_query is nullptr"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(opParamInfo_.layOutKey == nullptr, OP_LOGE(opName_, "attr layout_key is nullptr"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(opParamInfo_.sparseCount == nullptr, OP_LOGE(opName_, "attr sparse_count is nullptr"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(opParamInfo_.sparseMode == nullptr, OP_LOGE(opName_, "attr sparse_mode is nullptr"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::CheckRequiredParaExistence() const +{ + if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS || CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetOpName() +{ + if (context_->GetNodeName() == nullptr) { + OP_LOGE("LightningIndexer", "opName got from TilingContext is nullptr"); + return ge::GRAPH_FAILED; + } + opName_ = context_->GetNodeName(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetNpuInfo() +{ + platformInfo_ = context_->GetPlatformInfo(); + OP_CHECK_IF(platformInfo_ == nullptr, OP_LOGE(opName_, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo_); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + OP_CHECK_IF(aicNum == 0 || aivNum == 0, OP_LOGE(opName_, "num of core obtained is 0."), return GRAPH_FAILED); + + socVersion_ = ascendcPlatform.GetSocVersion(); + if ((socVersion_ != platform_ascendc::SocVersion::ASCEND910B) && + (socVersion_ != platform_ascendc::SocVersion::ASCEND910_93) && + (socVersion_ != platform_ascendc::SocVersion::ASCEND950)) { + OP_LOGE(opName_, "SOC Version[%d] is not support.", static_cast(socVersion_)); + return GRAPH_FAILED; + } + OP_CHECK_IF(context_->GetWorkspaceSizes(1) == nullptr, OP_LOGE(opName_, "workSpaceSize got from ge is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->GetRawTilingData() == nullptr, + OP_LOGE(context_->GetNodeName(), "RawTilingData got from GE context is nullptr."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void LIInfoParser::GetOptionalInputParaInfo() +{ + opParamInfo_.actualSeqLengthsQ.tensor = context_->GetOptionalInputTensor(ACTUAL_SEQ_Q_INDEX); + opParamInfo_.actualSeqLengthsQ.desc = context_->GetOptionalInputDesc(ACTUAL_SEQ_Q_INDEX); + opParamInfo_.actualSeqLengths.tensor = context_->GetOptionalInputTensor(ACTUAL_SEQ_K_INDEX); + opParamInfo_.actualSeqLengths.desc = context_->GetOptionalInputDesc(ACTUAL_SEQ_K_INDEX); + opParamInfo_.blockTable.tensor = context_->GetOptionalInputTensor(BLOCK_TABLE_INDEX); + opParamInfo_.blockTable.desc = context_->GetOptionalInputDesc(BLOCK_TABLE_INDEX); +} + +void LIInfoParser::GetInputParaInfo() +{ + opParamInfo_.query.desc = context_->GetInputDesc(QUERY_INDEX); + opParamInfo_.query.shape = context_->GetInputShape(QUERY_INDEX); + opParamInfo_.key.desc = context_->GetInputDesc(KEY_INDEX); + opParamInfo_.key.shape = context_->GetInputShape(KEY_INDEX); + opParamInfo_.weights.desc = context_->GetInputDesc(WEIGTHS_INDEX); + opParamInfo_.weights.shape = context_->GetInputShape(WEIGTHS_INDEX); + GetOptionalInputParaInfo(); +} + +void LIInfoParser::GetOutputParaInfo() +{ + opParamInfo_.attenOut.desc = context_->GetOutputDesc(LIGHTNING_INDEXER); + opParamInfo_.attenOut.shape = context_->GetOutputShape(LIGHTNING_INDEXER); + opParamInfo_.valuesOut.desc = context_->GetOutputDesc(LIGHTNING_VALUES); + opParamInfo_.valuesOut.shape = context_->GetOutputShape(LIGHTNING_VALUES); +} + +ge::graphStatus LIInfoParser::GetAndCheckAttrParaInfo() +{ + auto attrs = context_->GetAttrs(); + OP_CHECK_IF(attrs == nullptr, OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "attrs got from ge is nullptr"), + return ge::GRAPH_FAILED); + OP_LOGI(context_->GetNodeName(), "GetAndCheckAttrParaInfo start"); + opParamInfo_.layOut = attrs->GetStr(ATTR_QUERY_LAYOUT_INDEX); + opParamInfo_.layOutKey = attrs->GetStr(ATTR_KEY_LAYOUT_INDEX); + opParamInfo_.sparseCount = attrs->GetAttrPointer(ATTR_SPARSE_COUNT_INDEX); + opParamInfo_.sparseMode = attrs->GetAttrPointer(ATTR_SPARSE_MODE_INDEX); + opParamInfo_.preTokens = attrs->GetAttrPointer(ATTR_PRE_TOKENS_INDEX); + opParamInfo_.nextTokens = attrs->GetAttrPointer(ATTR_NEXT_TOKENS_INDEX); + opParamInfo_.returnValue = attrs->GetAttrPointer(ATTR_RETURN_VALUE_INDEX); + if (opParamInfo_.layOut != nullptr) { + OP_LOGI(context_->GetNodeName(), "layout_query is:%s", opParamInfo_.layOut); + } + if (opParamInfo_.layOutKey != nullptr) { + OP_LOGI(context_->GetNodeName(), "layout_key is:%s", opParamInfo_.layOutKey); + } + if (opParamInfo_.sparseCount != nullptr) { + OP_LOGI(context_->GetNodeName(), "selscted count is:%d", *opParamInfo_.sparseCount); + } + if (opParamInfo_.sparseMode != nullptr) { + OP_LOGI(context_->GetNodeName(), "sparse mode is:%d", *opParamInfo_.sparseMode); + } + if (opParamInfo_.preTokens != nullptr) { + OP_LOGI(context_->GetNodeName(), "pre tokens is:%d", *opParamInfo_.preTokens); + } + if (opParamInfo_.nextTokens != nullptr) { + OP_LOGI(context_->GetNodeName(), "next tokens is:%d", *opParamInfo_.nextTokens); + } + if (opParamInfo_.returnValue != nullptr) { + OP_LOGI(context_->GetNodeName(), "return value is:%d", *opParamInfo_.returnValue); + } + OP_LOGI(context_->GetNodeName(), "GetAndCheckAttrParaInfo end"); + OP_CHECK_IF( + ((std::string(opParamInfo_.layOutKey) != "PA_BSND") + && (std::string(opParamInfo_.layOut) != std::string(opParamInfo_.layOutKey))), + OP_LOGE(opName_, "under non-PA conditions, layout_query and layout_key should be equal."), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + ((std::string(opParamInfo_.layOutKey) != "PA_BSND") && (std::string(opParamInfo_.layOutKey) != "BSND") + && (std::string(opParamInfo_.layOutKey) != "TND")), + OP_LOGE(opName_, "input attr layout_key only supported PA_BSND, BSND or TND"), return ge::GRAPH_FAILED); + OP_CHECK_IF(((std::string(opParamInfo_.layOut) != "BSND") && (std::string(opParamInfo_.layOut) != "TND")), + OP_LOGE(opName_, "input attr layout_query only supported BSND or TND."), return ge::GRAPH_FAILED); + OP_CHECK_IF((!((*opParamInfo_.sparseCount > 0) && (*opParamInfo_.sparseCount <= SPARSE_LIMIT)) && + *opParamInfo_.sparseCount % 1024 != 0), + OP_LOGE(opName_, "input attr sparse_count must > 0 and <= 8192. And when sparse_count > 2048, sparse_count must be an integer multiple of 1024."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(!((*opParamInfo_.sparseMode == 0) || (*opParamInfo_.sparseMode == SPARSE_MODE_LOWER)), + OP_LOGE(opName_, "input attr sparse_mode only supported 0 or 3."), return ge::GRAPH_FAILED); + OP_CHECK_IF(*opParamInfo_.preTokens != INT64_MAX, + OP_LOGE(opName_, "input attr pre_tokens only supported INT64_MAX."), return ge::GRAPH_FAILED); + OP_CHECK_IF(*opParamInfo_.nextTokens != INT64_MAX, + OP_LOGE(opName_, "input attr nextTokens only supported INT64_MAX."), return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetOpParaInfo() +{ + GetInputParaInfo(); + GetOutputParaInfo(); + if (ge::GRAPH_SUCCESS != GetAndCheckAttrParaInfo()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetAndCheckInOutDataType() +{ + inputQType_ = opParamInfo_.query.desc->GetDataType(); + inputKType_ = opParamInfo_.key.desc->GetDataType(); + weightsType_ = opParamInfo_.weights.desc->GetDataType(); + outputType_ = opParamInfo_.attenOut.desc->GetDataType(); + valuesOutType_ = opParamInfo_.valuesOut.desc->GetDataType(); + + bool inDTypeAllEqual = (inputQType_ == inputKType_); + OP_CHECK_IF(!inDTypeAllEqual, + OP_LOGE(opName_, "The data types of the input query and key must be the same."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(((inputQType_ != ge::DT_FLOAT16) && (inputQType_ != ge::DT_BF16)), + OP_LOGE(opName_, "The data types of the input query, key must be float16 or bfloat16."), + return ge::GRAPH_FAILED); + if (socVersion_ == platform_ascendc::SocVersion::ASCEND950) { + OP_CHECK_IF((inputQType_ != weightsType_), + OP_LOGE(opName_, "The data types of the input query, key, and weights must be the same."), + return ge::GRAPH_FAILED); + } else { + if (weightsType_ != ge::DT_FLOAT) { + OP_CHECK_IF((inputQType_ != weightsType_), + OP_LOGE(opName_, "The data types of the input query, key, and weights must be the same."), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF((weightsType_ != ge::DT_FLOAT), + OP_LOGE(opName_, "The data types of the input weights must be float32."), + return ge::GRAPH_FAILED); + } + } + OP_CHECK_IF(outputType_ != ge::DT_INT32, + OP_LOGE(opName_, "The data types of the output sparse_indices must be int32."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(valuesOutType_ != inputQType_, + OP_LOGE(opName_, "The data types of the output sparse_values must be same as inputQType."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetQueryKeyAndOutLayout() +{ + // 获取query,key的Layout基准值 + const map layoutMap = { + {"BSND", DataLayout::BSND}, + {"TND", DataLayout::TND}, + {"PA_BSND", DataLayout::BnBsND} + }; + + std::string layout(opParamInfo_.layOut); + auto it = layoutMap.find(layout); + if (it != layoutMap.end()) { + qLayout_ = it->second; + } + + std::string layoutKey(opParamInfo_.layOutKey); + auto itKey = layoutMap.find(layoutKey); + if (itKey != layoutMap.end()) { + kLayout_ = itKey->second; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetAndCheckOptionalInput() +{ + if (kLayout_ == DataLayout::BnBsND) { + OP_CHECK_IF(opParamInfo_.blockTable.tensor == nullptr, + OP_LOGE(opName_, "when layout_key is PA_BSND, input block_table must not be null"), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + opParamInfo_.actualSeqLengths.tensor == nullptr, + OP_LOGE(opName_, "when layout_key is PA_BSND, input actual_seq_lengths_key must not be null"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.blockTable.desc->GetDataType() != ge::DT_INT32, + OP_LOGE(opName_, "input block_table data type only support int32"), return ge::GRAPH_FAILED); + } else if (kLayout_ == DataLayout::TND) { + OP_CHECK_IF(opParamInfo_.actualSeqLengths.tensor == nullptr, + OP_LOGE(opName_, "when layout_key is TND, input actual_seq_lengths_key must not be null"), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(opParamInfo_.actualSeqLengths.tensor != nullptr && + opParamInfo_.actualSeqLengths.desc->GetDataType() != ge::DT_INT32, + OP_LOGE(opName_, "input actual_seq_lengths_key data type only support int32"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.actualSeqLengths.tensor != nullptr && + opParamInfo_.actualSeqLengths.desc->GetDataType() != ge::DT_INT32, + OP_LOGE(opName_, "input actual_seq_lengths_key data type only support int32"), + return ge::GRAPH_FAILED); + if (qLayout_ == DataLayout::TND) { + OP_CHECK_IF(opParamInfo_.actualSeqLengthsQ.tensor == nullptr, + OP_LOGE(opName_, "when layout_query is TND, input actual_seq_lengths_query must not be null"), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(opParamInfo_.actualSeqLengthsQ.tensor != nullptr && + opParamInfo_.actualSeqLengthsQ.desc->GetDataType() != ge::DT_INT32, + OP_LOGE(opName_, "input actual_seq_lengths_query data type only support int32"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(kLayout_ != DataLayout::BnBsND && opParamInfo_.blockTable.tensor != nullptr, + OP_LOGE(opName_, "when key layout is not PA_BSND, input block_table must be null"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::CheckShapeDim() +{ + OP_CHECK_IF((opParamInfo_.blockTable.tensor != nullptr) && + (opParamInfo_.blockTable.tensor->GetStorageShape().GetDimNum() != DIM_NUM_TWO), + OP_LOGE(opName_, "the dim num of block_table's shape should be 2"), return ge::GRAPH_FAILED); + + uint32_t kShapeDim = opParamInfo_.key.shape->GetStorageShape().GetDimNum(); + uint32_t qShapeDim = opParamInfo_.query.shape->GetStorageShape().GetDimNum(); + uint32_t weightsShapeDim = opParamInfo_.weights.shape->GetStorageShape().GetDimNum(); + uint32_t outShapeDim = opParamInfo_.attenOut.shape->GetStorageShape().GetDimNum(); + uint32_t valuesOutShapeDim = opParamInfo_.valuesOut.shape->GetStorageShape().GetDimNum(); + uint32_t qExpectShapeDim = DIM_NUM_FOUR; + uint32_t kExpectShapeDim = DIM_NUM_FOUR; + if (qLayout_ == DataLayout::TND) { + qExpectShapeDim = DIM_NUM_THREE; + } + if (kLayout_ == DataLayout::TND) { + kExpectShapeDim = DIM_NUM_THREE; + } + OP_CHECK_IF(kShapeDim != kExpectShapeDim, + OP_LOGE(opName_, "the dim num of key's shape should be %u, but now is %u", kExpectShapeDim, kShapeDim), + return ge::GRAPH_FAILED); + OP_CHECK_IF(qShapeDim != qExpectShapeDim, + OP_LOGE(opName_, "the dim num of query's shape should be %u, but now is %u", + qExpectShapeDim, qShapeDim), + return ge::GRAPH_FAILED); + OP_CHECK_IF(outShapeDim != qExpectShapeDim, + OP_LOGE(opName_, "the dim num of sparse_indices's shape should be %u, but now is %u", + qExpectShapeDim, outShapeDim), + return ge::GRAPH_FAILED); + OP_CHECK_IF(valuesOutShapeDim != qExpectShapeDim && (*opParamInfo_.returnValue), + OP_LOGE(opName_, "the dim num of sparse_values's shape should be %u, but now is %u", + qExpectShapeDim, valuesOutShapeDim), + return ge::GRAPH_FAILED); + OP_CHECK_IF(!(weightsShapeDim == qExpectShapeDim - 1), + OP_LOGE(opName_, "the dim num of weights's shape should be %u, but now is %u", qExpectShapeDim - 1, + weightsShapeDim), + return ge::GRAPH_FAILED); + if (opParamInfo_.valuesOut.shape->GetStorageShape().GetShapeSize() != 0 && !(*opParamInfo_.returnValue)) { + OP_LOGW(opName_, "when returnValue is false, valuesOut must be null."); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetN1Size() +{ + if (qLayout_ == DataLayout::BSND) { + n1Size_ = static_cast(opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_TWO)); + } else { + // TND + n1Size_ = static_cast(opParamInfo_.query.shape->GetStorageShape().GetDim(1)); + } + OP_LOGI(context_->GetNodeName(), "n1Size is %d", n1Size_); + + OP_CHECK_IF(n1Size_ > QUERY_HEAD_NUM_LIMIT, OP_LOGE(opName_, "N1 is %u, but N1 must be no greater than %u.", + n1Size_, QUERY_HEAD_NUM_LIMIT), return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + const std::string &actualSeqLenName) const +{ + size = static_cast(tensor->GetShapeSize()); + if (size <= 0) { + OP_LOGE(opName_, "%s's shape size is %u, it should be greater than 0.", actualSeqLenName.c_str(), size); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetAndCheckN2Size() +{ + uint32_t n2Index = (kLayout_ == DataLayout::TND) ? DIM_IDX_ONE : DIM_IDX_TWO; + n2Size_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(n2Index)); + OP_LOGI(context_->GetNodeName(), "n2Size_ is %d", n2Size_); + OP_CHECK_IF(n2Size_ != 1, OP_LOGE(opName_, "key shape[%u] is numhead, only support 1.", n2Index), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetGSize() +{ + if (n1Size_ % n2Size_ != 0) { + OP_LOGE(opName_, "input query's head_num %u can not be a multiple of key's head_num %u.", n1Size_, n2Size_); + return ge::GRAPH_FAILED; + } + gSize_ = n1Size_ / n2Size_; + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetBatchSize() +{ + // 获取B基准值 + // 1、非TND/NTD时, 以query的batch_size维度为基准; + // 2、TND/NTD时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小 + if ((qLayout_ == DataLayout::TND)) { + return GetActualSeqLenSize(bSize_, opParamInfo_.actualSeqLengthsQ.tensor, "input actual_seq_lengths_query"); + } else { // BSND + bSize_ = opParamInfo_.query.shape->GetStorageShape().GetDim(0); + return ge::GRAPH_SUCCESS; + } +} + +ge::graphStatus LIInfoParser::GetHeadDim() +{ + // 以query的D维度为基准 + uint32_t dIndex = DIM_IDX_TWO; + // 根据layout确定D维度在shape中的位置 + switch (qLayout_) { + case DataLayout::TND: + // TND格式: [Total, N, D] -> D是第2维(索引2) + dIndex = DIM_IDX_TWO; + break; + case DataLayout::BSND: + // BSND格式: [Batch, SeqLen, N, D] -> D是第3维(索引3) + dIndex = DIM_IDX_THREE; + break; + default: + OP_LOGE(opName_, "unsupported layout for getting head dim."); + return ge::GRAPH_FAILED; + } + headDim_ = opParamInfo_.query.shape->GetStorageShape().GetDim(dIndex); + OP_CHECK_IF(headDim_ != HEAD_DIM_LIMIT, OP_LOGE(opName_, "input query's last dim head_dim only support 128."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetS1Size() +{ + if (qLayout_ == DataLayout::BSND) { + s1Size_ = opParamInfo_.query.shape->GetStorageShape().GetDim(1); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetAndCheckBlockSize() +{ + blockSize_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(1)); + OP_LOGI(context_->GetNodeName(), "blockSize_ is %d", blockSize_); + + OP_CHECK_IF(((blockSize_ % 16 != 0) || (blockSize_ == 0) || (blockSize_ > 1024)), + OP_LOGE(opName_, "input key's block_size must be a multiple of 16 and belong to (0, 1024]."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::CheckBlockCount() +{ + int32_t blockCount_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(0)); + OP_CHECK_IF((blockCount_ == 0), + OP_LOGE(opName_, "input key's block_count cannot be 0."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetS2SizeForPageAttention() +{ + if (GetAndCheckBlockSize() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + if (CheckBlockCount() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + maxBlockNumPerBatch_ = opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1); + s2Size_ = maxBlockNumPerBatch_ * blockSize_; + OP_LOGI(context_->GetNodeName(), "maxBlockNumPerBatch_ is %d, blockSize_ is %d, s2Size_ is %d", + maxBlockNumPerBatch_, blockSize_, s2Size_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetS2Size() +{ + // 获取S2基准值 + // 1、BATCH_CONTINUOUS时, 从key的S轴获取 + // 3、PAGE_ATTENTION时, S2 = block_table.dim1 * block_size + if (kLayout_ == DataLayout::BnBsND) { + return GetS2SizeForPageAttention(); + } else if (kLayout_ == DataLayout::TND) { + s2Size_ = opParamInfo_.key.shape->GetStorageShape().GetDim(0); + } else if (kLayout_ == DataLayout::BSND) { + s2Size_ = opParamInfo_.key.shape->GetStorageShape().GetDim(1); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::ValidateInputShapesMatchQtnd() +{ + // -----------------------check BatchSize------------------- + // bSize_ 来源于act_seq_q + if (kLayout_ == DataLayout::TND) { + OP_CHECK_IF( + (opParamInfo_.actualSeqLengths.tensor->GetShapeSize() != bSize_), + OP_LOGE(opName_, + "TND case input actual_seq_lengths_query, actual_seq_lengths_key are %u, %ld respectively, they must be same.", + bSize_, opParamInfo_.actualSeqLengths.tensor->GetShapeSize()), + return ge::GRAPH_FAILED); + } else { // kLayout_ PA_BSND + OP_CHECK_IF( + (opParamInfo_.actualSeqLengths.tensor->GetShapeSize() != bSize_) || + (opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) != bSize_), + OP_LOGE( + opName_, + "TND case input actual_seq_lengths_query, actual_seq_lengths_key, block_table dim 0 are %u, %ld, %ld respectively, they must be same.", + bSize_, opParamInfo_.actualSeqLengths.tensor->GetShapeSize(), + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + } + // -----------------------check T------------------- + uint32_t qTsize = opParamInfo_.query.shape->GetStorageShape().GetDim(0); + OP_CHECK_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != qTsize) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != qTsize), + OP_LOGE(opName_, "TND case input query, weights and sparse_indices dim 0 are %u, %ld, %ld respectively, they must be same.", + qTsize, opParamInfo_.weights.shape->GetStorageShape().GetDim(0), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + + OP_CHECK_IF((opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(0) != qTsize && + (*opParamInfo_.returnValue)), + OP_LOGE(opName_, "TND case input query and sparse_values dim 0 are %u, %ld respectively, they must be same.", + qTsize, opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::ValidateInputShapesMatchQbsnd() +{ + // -----------------------check BatchSize------------------- + // bSize_ 来源于query + if (kLayout_ == DataLayout::BnBsND) { + OP_CHECK_IF((opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) != bSize_) || + (opParamInfo_.actualSeqLengths.tensor->GetShapeSize() != bSize_), + OP_LOGE(opName_, "BSND case input query, actual_seq_lengths_key, block_table dim 0 are %u, %ld, %ld respectively, they must be same.", + bSize_, opParamInfo_.actualSeqLengths.tensor->GetShapeSize(), + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + } else if (kLayout_ == DataLayout::BSND) { + OP_CHECK_IF(opParamInfo_.key.shape->GetStorageShape().GetDim(0) != bSize_, + OP_LOGE(opName_, "BSND case input query, key dim 0 are %u, %ld respectively, they must be same.", + bSize_, opParamInfo_.key.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + OP_CHECK_IF((opParamInfo_.actualSeqLengths.tensor != nullptr) && + (opParamInfo_.actualSeqLengths.tensor->GetShapeSize() != bSize_), + OP_LOGE(opName_, "BSND case input query, actual_seq_lengths_key dim 0 are %u, %ld respectively, they must be same.", + bSize_, opParamInfo_.actualSeqLengths.tensor->GetShapeSize()), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != bSize_) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != bSize_), + OP_LOGE(opName_, "BSND case input query, weight and sparse_indices dim 0 are %u, %ld, %ld respectively, they must be same.", + bSize_, opParamInfo_.weights.shape->GetStorageShape().GetDim(0), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + OP_CHECK_IF((opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(0) != bSize_ && + (*opParamInfo_.returnValue)), + OP_LOGE(opName_, "BSND case input query, sparse_values dim 0 are %u, %ld respectively, they must be same.", + bSize_, opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + OP_CHECK_IF((opParamInfo_.actualSeqLengthsQ.tensor != nullptr) && + (opParamInfo_.actualSeqLengthsQ.tensor->GetShapeSize() != bSize_), + OP_LOGE(opName_, "BSND case input query, actual_seq_lengths_query dim 0 are %u, %ld respectively, they must be same", + bSize_, opParamInfo_.actualSeqLengthsQ.tensor->GetShapeSize()), + return ge::GRAPH_FAILED); + // -----------------------check S1------------------- + OP_CHECK_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(1) != s1Size_) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(1) != s1Size_), + OP_LOGE(opName_, "BSND case input query, weight and sparse_indices dim 1 are %u, %ld, %ld, they must be same.", + s1Size_, opParamInfo_.weights.shape->GetStorageShape().GetDim(1), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(1)), + return ge::GRAPH_FAILED); + OP_CHECK_IF((opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(1) != s1Size_ && + (*opParamInfo_.returnValue)), + OP_LOGE(opName_, "BSND case input query and sparse_values dim 1 are %u, %ld, they must be same.", + s1Size_, opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(1)), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::ValidateInputShapesMatch() +{ + /* + TND: + query [T,N1,D], + key [BlockNum,BlockSize,N2,D], + weight [T,N1], + block_table [BatchSize, BatchMaxBlockNum], + act_seq_k [BatchSize] + act_seq_q [BatchSize], + out [T,N2,topk] + ---------------------- + BSND: + query [BatchSize,S1,N1,D], + key [BlockNum,BlockSize,N2,D], + weight [BatchSize,S1,N1], + block_table [BatchSize, BatchMaxBlockNum], + act_seq_k [BatchSize] + act_seq_q [BatchSize] 可选 + out [BatchSize,S1,N2,topk] + */ + uint32_t queryWeightsN1Dim = 1; + uint32_t outN2Dim = 1; + if (qLayout_ == DataLayout::TND) { + if (ValidateInputShapesMatchQtnd() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + } else { // qLayout_ BSND + if (ValidateInputShapesMatchQbsnd() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + queryWeightsN1Dim = DIM_IDX_TWO; + outN2Dim = DIM_IDX_TWO; + } + // -----------------------check N1------------------- + OP_CHECK_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(queryWeightsN1Dim) != n1Size_), + OP_LOGE(opName_, "input query, weight shape dim N1 must be same."), return ge::GRAPH_FAILED); + // -----------------------check D------------------- + uint32_t keyDDim = kLayout_ == DataLayout::TND ? DIM_IDX_TWO : DIM_IDX_THREE; + OP_CHECK_IF((opParamInfo_.key.shape->GetStorageShape().GetDim(keyDDim) != headDim_), + OP_LOGE(opName_, "input query, key shape last dim must be same."), return ge::GRAPH_FAILED); + // -----------------------check N2------------------- + OP_CHECK_IF((opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim) != n2Size_), + OP_LOGE(opName_, "input query and output sparse_indices shape n2 dim must be same," + "but now they are %u, %ld respectively.", + n2Size_, opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim)), + return ge::GRAPH_FAILED); + OP_CHECK_IF((opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(outN2Dim) != n2Size_ && + (*opParamInfo_.returnValue)), + OP_LOGE(opName_, "input query and sparse_values shape n2 dim must be same," + "but now they are %u, %ld respectively.", + n2Size_, opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(outN2Dim)), + return ge::GRAPH_FAILED); + // -----------------------check sparse_count------------------- + OP_CHECK_IF((opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim + 1) != *opParamInfo_.sparseCount), + OP_LOGE(opName_, "output sparse_indices shape last dim must be same as attr sparse_count," + "but now they are %u, %ld respectively.", *opParamInfo_.sparseCount, + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim + 1)), + return ge::GRAPH_FAILED); + OP_CHECK_IF((opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(outN2Dim + 1) != *opParamInfo_.sparseCount && + (*opParamInfo_.returnValue)), + OP_LOGE(opName_, "output sparse_values shape last dim must be same as attr sparse_count," + "but now they are %u, %ld respectively.", *opParamInfo_.sparseCount, + opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(outN2Dim + 1)), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void LIInfoParser::GenerateInfo(LITilingInfo &liInfo) +{ + liInfo.opName = opName_; + liInfo.platformInfo = platformInfo_; + liInfo.opParamInfo = opParamInfo_; + liInfo.socVersion = socVersion_; + + liInfo.bSize = bSize_; + liInfo.n1Size = n1Size_; + liInfo.n2Size = n2Size_; + liInfo.s1Size = s1Size_; + liInfo.s2Size = s2Size_; + liInfo.gSize = gSize_; + + liInfo.inputQType = inputQType_; + liInfo.inputKType = inputKType_; + liInfo.weightsType = weightsType_; + liInfo.outputType = outputType_; + + liInfo.blockSize = blockSize_; + liInfo.maxBlockNumPerBatch = maxBlockNumPerBatch_; + + std::string layOutKeyStr(opParamInfo_.layOutKey); + liInfo.pageAttentionFlag = layOutKeyStr == "PA_BSND" ? true : false; + liInfo.sparseMode = *opParamInfo_.sparseMode; + liInfo.sparseCount = *opParamInfo_.sparseCount; + liInfo.preTokens = *opParamInfo_.preTokens; + liInfo.nextTokens = *opParamInfo_.nextTokens; + liInfo.returnValue = *opParamInfo_.returnValue; + + liInfo.inputQLayout = qLayout_; + liInfo.inputKLayout = kLayout_; +} + +ge::graphStatus LIInfoParser::ParseAndCheck(LITilingInfo &liInfo) +{ + if (ge::GRAPH_SUCCESS != GetOpName() || ge::GRAPH_SUCCESS != GetNpuInfo() || ge::GRAPH_SUCCESS != GetOpParaInfo() || + ge::GRAPH_SUCCESS != CheckRequiredParaExistence()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetAndCheckInOutDataType() || ge::GRAPH_SUCCESS != GetQueryKeyAndOutLayout() || + ge::GRAPH_SUCCESS != GetAndCheckOptionalInput()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != CheckShapeDim() || ge::GRAPH_SUCCESS != GetN1Size() || + ge::GRAPH_SUCCESS != GetAndCheckN2Size() || ge::GRAPH_SUCCESS != GetGSize()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetBatchSize() || ge::GRAPH_SUCCESS != GetS1Size() || ge::GRAPH_SUCCESS != GetHeadDim() || + ge::GRAPH_SUCCESS != GetS2Size()) { + return ge::GRAPH_FAILED; + } + if (ge::GRAPH_SUCCESS != ValidateInputShapesMatch()) { + return ge::GRAPH_FAILED; + } + + GenerateInfo(liInfo); + + return ge::GRAPH_SUCCESS; +} + +// --------------------------TilingPrepare函数定义------------------------------------- +static ge::graphStatus TilingPrepareForLightningIndexer(gert::TilingParseContext * /* context */) +{ + return ge::GRAPH_SUCCESS; +} + +// --------------------------LightningIndexerTiling类成员函数定义----------------------- +ge::graphStatus LightningIndexerTiling::DoTiling(LITilingInfo *tilingInfo) +{ + // -------------set blockdim----------------- + auto ascendcPlatform = platform_ascendc::PlatformAscendC(tilingInfo->platformInfo); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + uint32_t blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum); + context_->SetBlockDim(blockDim); + + // -------------set workspacesize----------------- + constexpr uint32_t MM1_RES_ELEM_SIZE = 4; // 4: fp32 + constexpr uint32_t DOUBLE_BUFFER = 2; // 双Buffer + constexpr uint32_t M_BASE_SIZE = 512; // m轴基本块大小 + constexpr uint32_t S2_BASE_SIZE = 512; // S2轴基本块大小 + constexpr uint32_t V1_RES_ELEM_SIZE = 4; // 4: int32 + constexpr uint32_t V1_RES_ELEM_TYPE = 2; // 保留Index和Value 2种数据 + constexpr uint32_t V1_DECODE_PARAM_ELEM_SIZE = 8; // 8: int64 + constexpr uint32_t V1_DECODE_PARAM_NUM = 16; // Decode参数个数 + constexpr uint32_t V1_DECODE_DATA_NUM = 2; // Decode每个核需要存储头和尾部两块数据 + constexpr uint32_t S1_BASE_SIZE = 8; // S1轴基本块的大小 + constexpr uint32_t TOPK_MAX_SIZE = 2048; // TopK选取个数 + uint32_t workspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize(); + // 主流程需Workspace大小 + if (ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_3510) { + constexpr uint32_t s1BaseSize = 4; + constexpr uint32_t s2BaseSize = 128; + workspaceSize += + s1BaseSize * ((tilingInfo->s2Size + s2BaseSize - 1) / s2BaseSize) * s2BaseSize * sizeof(uint16_t) * aicNum; + } else { + constexpr uint32_t mm1ResSize = M_BASE_SIZE * S2_BASE_SIZE; + workspaceSize += mm1ResSize * MM1_RES_ELEM_SIZE * DOUBLE_BUFFER * aicNum; + // Decode流程(LD)需要Workspace大小 + // 临时存储Decode中间结果大小: 2(头/尾)*8(s1Base)*2(idx/value)*2048(K)*sizeof(int32)*24=6M + workspaceSize += + V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_RES_ELEM_TYPE * TOPK_MAX_SIZE * V1_RES_ELEM_SIZE * aicNum; + // 临时存储Decode中间参数信息大小: 2(头/尾)*8(s1Base)*16(paramNum)*sizeof(int64_t)*24=48k + workspaceSize += + V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_DECODE_PARAM_NUM * V1_DECODE_PARAM_ELEM_SIZE * aicNum; + } + size_t *workSpaces = context_->GetWorkspaceSizes(1); + workSpaces[0] = workspaceSize; + + // -------------set tilingdata----------------- + tilingData_.set_bSize(tilingInfo->bSize); + tilingData_.set_s2Size(tilingInfo->s2Size); + tilingData_.set_s1Size(tilingInfo->s1Size); + tilingData_.set_sparseCount(tilingInfo->sparseCount); + tilingData_.set_gSize(tilingInfo->gSize); + tilingData_.set_blockSize(tilingInfo->blockSize); + tilingData_.set_maxBlockNumPerBatch(tilingInfo->maxBlockNumPerBatch); + tilingData_.set_sparseMode(tilingInfo->sparseMode); + tilingData_.set_preTokens(tilingInfo->preTokens); + tilingData_.set_nextTokens(tilingInfo->nextTokens); + tilingData_.set_returnValue(tilingInfo->returnValue); + tilingData_.set_usedCoreNum(blockDim); + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + + // -------------set tilingkey----------------- + // int DT_W_FLAG, DT_Q, DT_KV, DT_OUT, PAGE_ATTENTION, FLASH_DECODE, LAYOUT_T, KV_LAYOUT_T + uint32_t inputQType = static_cast(tilingInfo->inputQType); + uint32_t inputKType = static_cast(tilingInfo->inputKType); + uint32_t weightsType = static_cast(tilingInfo->weightsType); + uint32_t outputType = static_cast(tilingInfo->outputType); + uint32_t pageAttentionFlag = static_cast(tilingInfo->pageAttentionFlag); + uint32_t inputQLayout = static_cast(tilingInfo->inputQLayout); + uint32_t inputKLayout = static_cast(tilingInfo->inputKLayout); + uint32_t weightTypeFlag = (weightsType == ge::DT_FLOAT) ? 1 : 0; + uint64_t tilingKey = + GET_TPL_TILING_KEY(inputQType, inputKType, outputType, pageAttentionFlag, inputQLayout, inputKLayout, weightTypeFlag); + context_->SetTilingKey(tilingKey); + context_->SetScheduleMode(1); // 1: batchmode模式 + + return ge::GRAPH_SUCCESS; +} + +// --------------------------Tiling函数定义--------------------------- +ge::graphStatus TilingForLightningIndexer(gert::TilingContext *context) +{ + OP_CHECK_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("LightningIndexer", "Tiling context is null."), + return ge::GRAPH_FAILED); + LITilingInfo liInfo; + LIInfoParser LIInfoParser(context); + if (LIInfoParser.ParseAndCheck(liInfo) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + LightningIndexerTiling liTiling(context); + return liTiling.DoTiling(&liInfo); +} + +// --------------------------Tiling函数及TilingPrepare函数注册-------- +IMPL_OP_OPTILING(LightningIndexer) + .Tiling(TilingForLightningIndexer) + .TilingParse(TilingPrepareForLightningIndexer); + +} // namespace optiling diff --git a/csrc/ascend/attention/lightning_indexer/op_host/lightning_indexer_tiling.h b/csrc/ascend/attention/lightning_indexer/op_host/lightning_indexer_tiling.h new file mode 100644 index 000000000..a4e3bfa26 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_host/lightning_indexer_tiling.h @@ -0,0 +1,240 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_tiling.h + * \brief + */ + +#ifndef LIGHTNING_INDEXER_TILING_H_ +#define LIGHTNING_INDEXER_TILING_H_ + +#include "exe_graph/runtime/tiling_context.h" +#include "tiling/platform/platform_ascendc.h" +#include "register/op_def_registry.h" +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "err/ops_err.h" +#include "platform/platform_info.h" + +namespace optiling { +// ------------------公共定义-------------------------- +struct TilingRequiredParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; +}; + +struct TilingOptionalParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::Tensor *tensor; +}; + +enum class DataLayout : uint32_t { + BSND = 0, + TND = 1, + BnBsND = 2 +}; + +// ------------------算子原型索引常量定义---------------- +// Inputs Index +constexpr uint32_t QUERY_INDEX = 0; +constexpr uint32_t KEY_INDEX = 1; +constexpr uint32_t WEIGTHS_INDEX = 2; +constexpr uint32_t ACTUAL_SEQ_Q_INDEX = 3; +constexpr uint32_t ACTUAL_SEQ_K_INDEX = 4; +constexpr uint32_t BLOCK_TABLE_INDEX = 5; +//Outputs Index +constexpr uint32_t LIGHTNING_INDEXER = 0; +constexpr uint32_t LIGHTNING_VALUES = 1; +// Attributes Index +constexpr uint32_t ATTR_QUERY_LAYOUT_INDEX = 0; +constexpr uint32_t ATTR_KEY_LAYOUT_INDEX = 1; +constexpr uint32_t ATTR_SPARSE_COUNT_INDEX = 2; +constexpr uint32_t ATTR_SPARSE_MODE_INDEX = 3; +constexpr uint32_t ATTR_PRE_TOKENS_INDEX = 4; +constexpr uint32_t ATTR_NEXT_TOKENS_INDEX = 5; +constexpr uint32_t ATTR_RETURN_VALUE_INDEX = 6; +// Dim Index +constexpr uint32_t DIM_IDX_ONE = 1; +constexpr uint32_t DIM_IDX_TWO = 2; +constexpr uint32_t DIM_IDX_THREE = 3; +// Dim Num +constexpr uint32_t DIM_NUM_TWO = 2; +constexpr uint32_t DIM_NUM_THREE = 3; +constexpr uint32_t DIM_NUM_FOUR = 4; +// 入参限制常量 +constexpr uint32_t HEAD_DIM_LIMIT = 128; +constexpr uint32_t SPARSE_LIMIT = 2048; +constexpr uint32_t SPARSE_MODE_LOWER = 3; +constexpr uint32_t QUERY_HEAD_NUM_LIMIT = 64; + +// -----------算子TilingData定义--------------- +BEGIN_TILING_DATA_DEF(LITilingData) +TILING_DATA_FIELD_DEF(uint32_t, bSize) +TILING_DATA_FIELD_DEF(uint32_t, n2Size) +TILING_DATA_FIELD_DEF(uint32_t, gSize) +TILING_DATA_FIELD_DEF(uint32_t, s1Size) +TILING_DATA_FIELD_DEF(uint32_t, s2Size) +TILING_DATA_FIELD_DEF(uint32_t, sparseCount) +TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum) +TILING_DATA_FIELD_DEF(uint32_t, blockSize) +TILING_DATA_FIELD_DEF(uint32_t, maxBlockNumPerBatch) +TILING_DATA_FIELD_DEF(uint32_t, sparseMode) +TILING_DATA_FIELD_DEF(int64_t, preTokens) +TILING_DATA_FIELD_DEF(int64_t, nextTokens) +TILING_DATA_FIELD_DEF(uint32_t, returnValue) +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(LightningIndexer, LITilingData) + +// -----------算子CompileInfo定义------------------- +struct LICompileInfo {}; + +// -----------算子Tiling入参结构体定义--------------- +struct LiParaInfo { + TilingRequiredParaInfo query = {nullptr, nullptr}; + TilingRequiredParaInfo key = {nullptr, nullptr}; + TilingRequiredParaInfo weights = {nullptr, nullptr}; + TilingOptionalParaInfo actualSeqLengthsQ = {nullptr, nullptr}; + TilingOptionalParaInfo actualSeqLengths = {nullptr, nullptr}; + TilingOptionalParaInfo blockTable = {nullptr, nullptr}; + TilingRequiredParaInfo attenOut = {nullptr, nullptr}; + TilingRequiredParaInfo valuesOut = {nullptr, nullptr}; + + const char *layOut = nullptr; + const char *layOutKey = nullptr; + const int32_t *blockSize = nullptr; + const int32_t *sparseMode = nullptr; + const int32_t *sparseCount = nullptr; + const int64_t *preTokens = nullptr; + const int64_t *nextTokens = nullptr; + const bool *returnValue = nullptr; +}; + +// -----------算子Tiling入参信息类--------------- +class LITilingInfo { +public: + const char *opName = nullptr; + fe::PlatFormInfos *platformInfo = nullptr; + LiParaInfo opParamInfo; + // Base Param + platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND910B; + uint32_t bSize = 0; + uint32_t n1Size = 0; + uint32_t n2Size = 0; + uint32_t s1Size = 0; + int64_t s2Size = 0; + uint32_t qkHeadDim = 0; + uint32_t gSize = 0; + // PageAttention + bool pageAttentionFlag = false; + int32_t blockSize = 0; + uint32_t maxBlockNumPerBatch = 0; + // Mask + int32_t sparseMode = 0; + // Others Flag + uint32_t sparseCount = 0; + int64_t preTokens = INT64_MAX; + int64_t nextTokens = INT64_MAX; + bool returnValue = false; + // DType + ge::DataType inputQType = ge::DT_FLOAT16; + ge::DataType inputKType = ge::DT_FLOAT16; + ge::DataType weightsType = ge::DT_FLOAT16; + ge::DataType outputType = ge::DT_INT32; + // Layout + DataLayout inputQLayout = DataLayout::BSND; + DataLayout inputKLayout = DataLayout::BnBsND; +}; + +// -----------算子Tiling入参信息解析及Check类--------------- +class LIInfoParser { +public: + explicit LIInfoParser(gert::TilingContext *context) : context_(context) + { + } + ~LIInfoParser() = default; + + ge::graphStatus CheckRequiredInOutExistence() const; + ge::graphStatus CheckRequiredAttrExistence() const; + ge::graphStatus CheckRequiredParaExistence() const; + ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + const std::string &actualSeqLenName) const; + ge::graphStatus GetOpName(); + ge::graphStatus GetNpuInfo(); + void GetOptionalInputParaInfo(); + void GetInputParaInfo(); + void GetOutputParaInfo(); + ge::graphStatus GetAndCheckAttrParaInfo(); + ge::graphStatus GetOpParaInfo(); + ge::graphStatus ValidateInputShapesMatchQbsnd(); + ge::graphStatus ValidateInputShapesMatchQtnd(); + ge::graphStatus ValidateInputShapesMatch(); + ge::graphStatus GetAndCheckInOutDataType(); + ge::graphStatus GetBatchSize(); + ge::graphStatus GetHeadDim(); + ge::graphStatus GetS1Size(); + ge::graphStatus GetAndCheckOptionalInput(); + ge::graphStatus CheckShapeDim(); + ge::graphStatus GetAndCheckBlockSize(); + ge::graphStatus CheckBlockCount(); + ge::graphStatus GetS2SizeForPageAttention(); + ge::graphStatus GetS2Size(); + ge::graphStatus GetQueryKeyAndOutLayout(); + ge::graphStatus GetN1Size(); + ge::graphStatus GetAndCheckN2Size(); + ge::graphStatus GetGSize(); + ge::graphStatus GetAttenMaskInfo(); + ge::graphStatus GetActualSeqInfo(); + void GenerateInfo(LITilingInfo &liInfo); + ge::graphStatus ParseAndCheck(LITilingInfo &liInfo); + +public: + gert::TilingContext *context_ = nullptr; + const char *opName_; + fe::PlatFormInfos *platformInfo_; + LiParaInfo opParamInfo_; + + // BaseParams + uint32_t bSize_ = 0; + uint32_t n1Size_ = 0; + uint32_t n2Size_ = 0; + uint32_t gSize_ = 0; + uint32_t s1Size_ = 0; + int64_t s2Size_ = 0; + uint32_t headDim_ = 0; + // Layout + DataLayout qLayout_ = DataLayout::BSND; + DataLayout kLayout_ = DataLayout::BnBsND; + // PageAttention + uint32_t maxBlockNumPerBatch_ = 0; + int32_t blockSize_ = 0; + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B; + ge::DataType inputQType_ = ge::DT_FLOAT16; + ge::DataType inputKType_ = ge::DT_FLOAT16; + ge::DataType weightsType_ = ge::DT_FLOAT16; + ge::DataType blockTableType_ = ge::DT_FLOAT16; + ge::DataType inputKRopeType_ = ge::DT_FLOAT16; + ge::DataType outputType_ = ge::DT_FLOAT16; + ge::DataType valuesOutType_ = ge::DT_FLOAT16; +}; + +// ---------------算子Tiling类--------------- +class LightningIndexerTiling { +public: + explicit LightningIndexerTiling(gert::TilingContext *context) : context_(context){}; + ge::graphStatus DoTiling(LITilingInfo *tilingInfo); + +private: + gert::TilingContext *context_ = nullptr; + LITilingData tilingData_; +}; + +} // namespace optiling +#endif // LIGHTNING_INDEXER_TILING_H_ diff --git a/csrc/ascend/attention/lightning_indexer/op_host/op_api/aclnn_lightning_indexer.cpp b/csrc/ascend/attention/lightning_indexer/op_host/op_api/aclnn_lightning_indexer.cpp new file mode 100644 index 000000000..60c1eb5fd --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_host/op_api/aclnn_lightning_indexer.cpp @@ -0,0 +1,138 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include +#include "graph/types.h" +#include "aclnn_lightning_indexer.h" + +#include "opdev/make_op_executor.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/tensor_view_utils.h" +#include "opdev/op_def.h" +#include "opdev/op_log.h" +#include "opdev/shape_utils.h" +#include "opdev/common_types.h" +#include "opdev/data_type_utils.h" +#include "opdev/format_utils.h" + +using namespace op; + +#ifdef __cplusplus +extern "C" { +#endif + +namespace { + +extern aclnnStatus aclnnInnerLightningIndexerGetWorkspaceSize( + const aclTensor *query, const aclTensor *key, const aclTensor *weights, + const aclTensor *actualSeqLengthsQueryOptional, const aclTensor *actualSeqLengthsKeyOptional, + const aclTensor *blockTableOptional, char *layoutQueryOptional, + char *layoutKeyOptional, int64_t sparseCount, int64_t sparseMode, + int64_t preTokens, int64_t nextTokens, bool returnValues, + const aclTensor *sparseIndicesOut, const aclTensor *sparseValuesOut, + uint64_t *workspaceSize, aclOpExecutor **executor); + +extern aclnnStatus aclnnInnerLightningIndexer(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, + const aclrtStream stream); + +class TensorHolder { +public: + TensorHolder(const aclTensor *&output, aclDataType dataType, std::string varName) { + inner_ = nullptr; + name_ = varName; + if (output == nullptr) { + std::vector shape = {0}; + int64_t addr = 0xff; + inner_ = aclCreateTensor(shape.data(), shape.size(), + dataType, shape.data(), 0, ACL_FORMAT_ND, + shape.data(), shape.size(), static_cast(&addr)); + output = inner_; + } + } + + ~TensorHolder() { + if (inner_) { + aclDestroyTensor(inner_); + inner_ = nullptr; + } + } + + void CheckTensorConditionalNotNull(bool conditional) const { + if (inner_ && conditional) { + OP_LOGW("Check %s != nullptr failed!", name_.c_str()); + } else if (!inner_ && !conditional) { + OP_LOGW("Check %s == nullptr failed!", name_.c_str()); + } + } + + bool IsTensorNotNull() const { + return inner_ == nullptr; + } + +private: + const aclTensor *inner_; + std::string name_; +}; + +aclnnStatus aclnnLightningIndexerGetWorkspaceSize( + const aclTensor *query, + const aclTensor *key, + const aclTensor *weights, + const aclTensor *actualSeqLengthsQueryOptional, + const aclTensor *actualSeqLengthsKeyOptional, + const aclTensor *blockTableOptional, + char *layoutQueryOptional, + char *layoutKeyOptional, + int64_t sparseCount, + int64_t sparseMode, + int64_t preTokens, + int64_t nextTokens, + bool returnValues, + const aclTensor *sparseIndicesOut, + const aclTensor *sparseValuesOut, + uint64_t *workspaceSize, + aclOpExecutor **executor) +{ + if (query == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Query pointer is null, cannot get data type!"); + return ge::GRAPH_FAILED; + } + DataType queryDataType = query->GetDataType(); + aclDataType queryAclDataType = ToAclDataType(queryDataType); + if (returnValues) { + if (sparseValuesOut == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "sparseValuesOut cannot be nullptr."); + return ge::GRAPH_FAILED; + } + } + auto sparseValuesOutHolder = TensorHolder(sparseValuesOut, queryAclDataType, std::string("sparseValuesOut")); + if (sparseValuesOut == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Failed to create the holder of tensor sparseValuesOut!"); + return ge::GRAPH_FAILED; + } + + return aclnnInnerLightningIndexerGetWorkspaceSize( + query, key, weights, actualSeqLengthsQueryOptional, actualSeqLengthsKeyOptional, blockTableOptional, + layoutQueryOptional, layoutKeyOptional, sparseCount, sparseMode, preTokens, nextTokens, returnValues, + sparseIndicesOut, sparseValuesOut, workspaceSize, executor); +} + +aclnnStatus aclnnLightningIndexer(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, + const aclrtStream stream) +{ + return aclnnInnerLightningIndexer(workspace, workspaceSize, executor, stream); +} + +} // namespace + +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/attention/lightning_indexer/op_host/op_api/aclnn_lightning_indexer.h b/csrc/ascend/attention/lightning_indexer/op_host/op_api/aclnn_lightning_indexer.h new file mode 100644 index 000000000..4b50f0944 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_host/op_api/aclnn_lightning_indexer.h @@ -0,0 +1,59 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef ACLNN_LIGHTNING_INDEXER_H +#define ACLNN_LIGHTNING_INDEXER_H + +#include "aclnn/acl_meta.h" +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @brief The first interface of aclnnLightningIndexerGetWorkspaceSize calculates the workspace size based on the specific calculation process. + * @domain aclnn_ops_infer + */ +__attribute__((visibility("default"))) +aclnnStatus aclnnLightningIndexerGetWorkspaceSize( + const aclTensor *query, + const aclTensor *key, + const aclTensor *weights, + const aclTensor *actualSeqLengthsQueryOptional, + const aclTensor *actualSeqLengthsKeyOptional, + const aclTensor *blockTableOptional, + char *layoutQueryOptional, + char *layoutKeyOptional, + int64_t sparseCount, + int64_t sparseMode, + int64_t preTokens, + int64_t nextTokens, + bool returnValues, + const aclTensor *sparseIndicesOut, + const aclTensor *sparseValuesOut, + uint64_t *workspaceSize, + aclOpExecutor **executor); + +/** + * @brief The second interface of aclnnLightningIndexer is used to perform calculations. + */ +__attribute__((visibility("default"))) +aclnnStatus aclnnLightningIndexer( + void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + const aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif // ACLNN_LIGHTNING_INDEXER_H diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/arch22/lightning_indexer_kernel.h b/csrc/ascend/attention/lightning_indexer/op_kernel/arch22/lightning_indexer_kernel.h new file mode 100644 index 000000000..8cdae321e --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/arch22/lightning_indexer_kernel.h @@ -0,0 +1,692 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. +  */ + +/*! + * \file lightning_indexer_kernel.h + * \brief + */ + +#ifndef LIGHTNING_INDEXER_KERNEL_H +#define LIGHTNING_INDEXER_KERNEL_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../lightning_indexer_common.h" +#include "lightning_indexer_service_vector.h" +#include "lightning_indexer_service_cube.h" + +namespace LIKernel { +using namespace LICommon; +using namespace LIServiceVec; +using namespace matmul; +using AscendC::CacheMode; +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +// 由于S2循环前，RunInfo还没有赋值，使用TempLoopInfo临时存放B、N、S1轴相关的信息；同时减少重复计算 +struct TempLoopInfo { + uint32_t bN2Idx = 0; + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint32_t gS1Idx = 0U; + uint32_t gS1LoopEnd = 0U; // gS1方向循环的结束Idx + uint32_t s2LoopEnd = 0U; // S2方向循环的结束Idx + uint32_t actS1Size = 1ULL; // 当前Batch循环处理的S1轴的实际大小 + uint32_t actS2Size = 0ULL; + bool curActSeqLenIsZero = false; + bool needDealActS1LessThanS1 = false; // S1的实际长度小于shape的S1长度时，是否需要清理输出 + uint32_t actMBaseSize = 0U; // m轴(gS1)方向实际大小 + uint32_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小 + uint32_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小 +}; + +template +class LightningIndexerKernel { +public: + __aicore__ inline LightningIndexerKernel(){}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths, + __gm__ uint8_t *blockTable, __gm__ uint8_t *sparseIndices, __gm__ uint8_t *sparseValues, + __gm__ uint8_t *workspace, const LITilingData *__restrict tiling, TPipe *tPipe); + __aicore__ inline void Process(); + + // =================================类型定义区================================= + static constexpr bool DT_W_FLAG = LIT::weightsTypeFlag; + using Q_T = typename LIT::queryType; + using K_T = typename LIT::keyType; + using OUT_T = typename LIT::outputType; + static constexpr bool PAGE_ATTENTION = LIT::pageAttention; + static constexpr LI_LAYOUT LAYOUT_T = LIT::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = LIT::keyLayout; + // 编译期条件选择模板第二个参数的类型，直接声明W_T + // 第一个模板参数：固定为Q_T；第二个模板参数：编译期选float/void + using W_T = typename LightningIndexerTypeTraits::type>::weightsType; + + using MM1_OUT_T = float; + + LightningIndexerServiceCube matmulService; + LightningIndexerServiceVector vectorService; + + // =================================常量区================================= + static constexpr uint32_t SYNC_C1_V1_FLAG = 4; + static constexpr uint32_t SYNC_V1_C1_FLAG = 5; + + static constexpr uint32_t M_BASE_SIZE = 512; + static constexpr uint32_t S2_BASE_SIZE = 512; + static constexpr uint32_t HEAD_DIM = 128; + static constexpr uint32_t K_HEAD_NUM = 1; + static constexpr uint32_t GM_ALIGN_BYTES = 512; + static constexpr uint32_t SPARSE_COUNT_8K = 8192; + static constexpr uint32_t BLOCK_CUBE_SIZE = 16; + + static constexpr int64_t LD_PREFETCH_LEN = 2; + // for workspace double + static constexpr uint32_t WS_DOUBLE = 2; + +protected: + TPipe *pipe = nullptr; + + // offset + uint64_t queryCoreOffset = 0ULL; + uint64_t keyCoreOffset = 0ULL; + uint64_t weightsCoreOffset = 0ULL; + uint64_t indiceOutCoreOffset = 0ULL; + + // ================================Global Buffer区================================= + GlobalTensor queryGm; + GlobalTensor keyGm; + GlobalTensor weightsGm; + + GlobalTensor indiceOutGm; + GlobalTensor valueOutGm; + GlobalTensor blockTableGm; + + GlobalTensor actualSeqLengthsGmQ; + GlobalTensor actualSeqLengthsGm; + // workspace + GlobalTensor mm1ResGm; // 存放S + GlobalTensor vec1ResGm; // 存放TopK计算中间结果 + GlobalTensor vec1ParamGm; // 存放LD参数信息 + + // ================================类成员变量==================================== + // aic、aiv核信息 + uint32_t tmpBlockIdx = 0U; + uint32_t aiCoreIdx = 0U; + uint32_t usedCoreNum = 0U; + + LICommon::ConstInfo constInfo{}; + TempLoopInfo tempLoopInfo{}; + LICommon::SplitCoreInfo splitCoreInfo{}; + + // ================================Init functions================================== + __aicore__ inline void InitTilingData(const LITilingData *__restrict tilingData); + __aicore__ inline void InitBuffers(); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths); + // ================================Split Core================================ + __aicore__ inline void SplitCore(uint32_t curCoreIdx, uint32_t &coreNum, LICommon::SplitCoreInfo &info); + __aicore__ inline uint32_t GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, uint32_t actS2Size); + __aicore__ inline uint32_t GetTotalBaseBlockNum(); + // ================================Process functions================================ + __aicore__ inline void ProcessMain(); + __aicore__ inline void ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, LICommon::RunInfo &runInfo); + __aicore__ inline void ProcessDecode(); + __aicore__ inline void ProcessInvalid(); + // ================================Params Calc===================================== + __aicore__ inline void CalcGS1LoopParams(uint32_t bN2Idx); + __aicore__ inline void GetBN2Idx(uint32_t bN2Idx); + __aicore__ inline uint32_t GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, uint32_t defaultSeqLen); + __aicore__ inline void GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, uint32_t &actS2Size); + __aicore__ inline void CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx); + __aicore__ inline void CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, LICommon::RunInfo &runInfo); + __aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start); +}; + +template +__aicore__ inline void LightningIndexerKernel::InitTilingData(const LITilingData *__restrict tilingData) +{ + usedCoreNum = tilingData->usedCoreNum; + constInfo.batchSize = tilingData->bSize; + constInfo.qHeadNum = constInfo.gSize = tilingData->gSize; + constInfo.kSeqSize = tilingData->s2Size; + constInfo.qSeqSize = tilingData->s1Size; + constInfo.attenMaskFlag = (tilingData->sparseMode == 3); + constInfo.kCacheBlockSize = tilingData->blockSize; + constInfo.maxBlockNumPerBatch = tilingData->maxBlockNumPerBatch; + constInfo.sparseCount = tilingData->sparseCount; + constInfo.preTokens = tilingData->preTokens; + constInfo.nextTokens = tilingData->nextTokens; + constInfo.returnValue = tilingData->returnValue; + + constInfo.outputLayout = LAYOUT_T; // 输出和输入形状一致 + if (LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS1 = true; + } + if (K_LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS2 = true; + } + + constInfo.kHeadNum = K_HEAD_NUM; + constInfo.headDim = HEAD_DIM; + constInfo.s2BaseSize = S2_BASE_SIZE; + constInfo.isSparseCountOver2K = (constInfo.sparseCount <= BASE_TOPK) ? false : true; + + constInfo.s1BaseSize = constInfo.isSparseCountOver2K ? SPARSE_COUNT_8K / constInfo.sparseCount * 2 : 8; + constInfo.mBaseSize = constInfo.s1BaseSize * constInfo.gSize; + constInfo.mBaseSizeAlign = LICommon::Align(constInfo.mBaseSize, BLOCK_CUBE_SIZE); +} + +template +__aicore__ inline void LightningIndexerKernel::InitBuffers() +{ + if ASCEND_IS_AIV { + vectorService.InitBuffers(pipe); + } else { + matmulService.InitBuffers(pipe); + } +} + +template +__aicore__ inline void LightningIndexerKernel::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengths) +{ + if (actualSeqLengthsQ == nullptr) { + constInfo.actualLenQDims = 0; + } else { + constInfo.actualLenQDims = constInfo.batchSize; + actualSeqLengthsGmQ.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsQ, constInfo.actualLenQDims); + } + if (actualSeqLengths == nullptr) { + constInfo.actualLenDims = 0; + } else { + constInfo.actualLenDims = constInfo.batchSize; + actualSeqLengthsGm.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengths, constInfo.actualLenDims); + } +} + +template +__aicore__ inline uint32_t LightningIndexerKernel::GetActualSeqLen(uint32_t bIdx, + uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, + uint32_t defaultSeqLen) +{ + if (actualLenDims == 0) { + return defaultSeqLen; + } else if (isAccumSeq && bIdx > 0) { + return actualSeqLengthsGm.GetValue(bIdx) - actualSeqLengthsGm.GetValue(bIdx - 1); + } else { + return actualSeqLengthsGm.GetValue(bIdx); + } +} + +template +__aicore__ inline void LightningIndexerKernel::GetS1S2ActualSeqLen(uint32_t bIdx, + uint32_t &actS1Size, uint32_t &actS2Size) +{ + actS1Size = GetActualSeqLen(bIdx, constInfo.actualLenQDims, constInfo.isAccumSeqS1, actualSeqLengthsGmQ, + constInfo.qSeqSize); + actS2Size = + GetActualSeqLen(bIdx, constInfo.actualLenDims, constInfo.isAccumSeqS2, actualSeqLengthsGm, constInfo.kSeqSize); +} + +template +__aicore__ inline uint32_t LightningIndexerKernel::GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, + uint32_t actS2Size) +{ + if (actS2Size == 0) { + return 0; + } + uint32_t s1Offset = constInfo.s1BaseSize * s1gIdx; + int32_t validS2LenBase = static_cast(actS2Size) - static_cast(actS1Size); + int32_t validS2Len = s1Offset + validS2LenBase + constInfo.s1BaseSize; + validS2Len = Min(validS2Len, static_cast(actS2Size)); + validS2Len = Max(validS2Len, 1); + return (validS2Len + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; +} + +template +__aicore__ inline uint32_t LightningIndexerKernel::GetTotalBaseBlockNum() +{ + uint32_t totalBlockNum = 0; + uint32_t actS1Size, actS2Size; + uint32_t s1GBaseNum, s2BaseNum; + for (uint32_t bIdx = 0; bIdx < constInfo.batchSize; bIdx++) { + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size); + s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize); + if (!constInfo.attenMaskFlag) { + s2BaseNum = constInfo.isSparseCountOver2K + ? (actS2Size > 0 ? 1 : 0) + : CeilDiv(actS2Size, constInfo.s2BaseSize); + totalBlockNum += s1GBaseNum * s2BaseNum * constInfo.kHeadNum; + continue; + } + for (uint32_t s1gIdx = 0; s1gIdx < s1GBaseNum; s1gIdx++) { + s2BaseNum = constInfo.isSparseCountOver2K + ? (actS2Size > 0 ? 1 : 0) + : GetS2BaseBlockNumOnMask(s1gIdx, actS1Size, actS2Size); + totalBlockNum += s2BaseNum * constInfo.kHeadNum; + } + } + return totalBlockNum; +} + +// 多核版本，双闭区间 +template +__aicore__ void inline LightningIndexerKernel::SplitCore(uint32_t curCoreIdx, + uint32_t &coreNum, LICommon::SplitCoreInfo &info) +{ + // 计算每个核最少处理的块数, 剩余的部分前面的核每个核多处理一块 + uint32_t totalBlockNum = GetTotalBaseBlockNum(); + uint32_t minBlockPerCore = totalBlockNum / coreNum; + uint32_t deal1MoreBlockCoreNum = totalBlockNum % coreNum; + uint32_t coreIdx = 0; + uint32_t lastGS1RemainBlockCnt = 0; + uint32_t coreDealBlockCnt = coreIdx < deal1MoreBlockCoreNum ? minBlockPerCore + 1 : minBlockPerCore; + coreNum = minBlockPerCore == 0 ? deal1MoreBlockCoreNum : coreNum; + + bool findLastCoreEnd = true; + uint32_t actS1Size, actS2Size; + uint32_t s1GBaseNum, s2BaseNum, s2Loop; + for (uint32_t bN2Idx = 0; bN2Idx < constInfo.batchSize * constInfo.kHeadNum; bN2Idx++) { + uint32_t bIdx = bN2Idx / constInfo.kHeadNum; + if (bN2Idx % constInfo.kHeadNum == 0) { + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size); + s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize); + s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize); + } + if constexpr (LAYOUT_T == LI_LAYOUT::BSND) { + if (findLastCoreEnd && (s1GBaseNum == 0U || s2BaseNum == 0U)) { + info.bN2Start = bN2Idx; + info.gS1Start = 0; + info.s2Start = 0; + findLastCoreEnd = false; + } + } + for (uint32_t gS1Idx = 0; gS1Idx < s1GBaseNum; gS1Idx++) { + if (constInfo.attenMaskFlag) { + s2BaseNum = GetS2BaseBlockNumOnMask(gS1Idx, actS1Size, actS2Size); + } + if (findLastCoreEnd && s2BaseNum == 0U) { + info.bN2Start = bN2Idx; + info.gS1Start = gS1Idx; + info.s2Start = 0; + findLastCoreEnd = false; + } + s2Loop = constInfo.isSparseCountOver2K ? (actS2Size > 0 ? 1 : 0) : s2BaseNum; + for (uint32_t s2Idx = 0; s2Idx < s2Loop;) { + if (findLastCoreEnd) { + info.bN2Start = bN2Idx; + info.gS1Start = gS1Idx; + info.s2Start = s2Idx; + findLastCoreEnd = false; + } + uint32_t s2RemainBaseNum = s2Loop - s2Idx; + if (lastGS1RemainBlockCnt + s2RemainBaseNum >= coreDealBlockCnt) { + info.bN2End = bN2Idx; + info.gS1End = gS1Idx; + info.s2End = constInfo.isSparseCountOver2K + ? s2BaseNum - 1 + : s2Idx + coreDealBlockCnt - lastGS1RemainBlockCnt - 1; + + if (coreIdx == curCoreIdx) { + // S2被切N核，那么只有第一个核需要处理LD，其他核不用 + if (s2Idx == 0 && info.s2End + 1 < s2BaseNum) { + info.isLD = true; + } + // 最后一个核处理的不是最后一个Batch，表明后面的Batch为空块(S2=0), 调整终点坐标以便清理输出 + if (coreIdx == coreNum - 1 && info.bN2End != constInfo.batchSize -1) { + info.bN2End = constInfo.batchSize -1; + info.gS1End = 0; + info.s2End = 0; + } + return; + } + coreIdx++; + findLastCoreEnd = true; + s2Idx = info.s2End + 1; + lastGS1RemainBlockCnt = 0; + coreDealBlockCnt = coreIdx < deal1MoreBlockCoreNum ? minBlockPerCore + 1 : minBlockPerCore; + } else { + lastGS1RemainBlockCnt += s2RemainBaseNum; + break; + } + } + } + } +} + +template +__aicore__ inline void LightningIndexerKernel::DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start) +{ + if ASCEND_IS_AIV { + if (constInfo.outputLayout == LI_LAYOUT::TND) { + uint32_t tSize = actualSeqLengthsGmQ.GetValue(constInfo.batchSize - 1); + uint32_t tBase = bIdx == 0 ? 0 : actualSeqLengthsGmQ.GetValue(bIdx - 1); + uint32_t s1Count = tempLoopInfo.actS1Size; + + for (uint32_t s1Idx = s1Start; s1Idx < s1Count; s1Idx++) { + uint64_t indiceOutOffset = + (tBase + s1Idx) * constInfo.kHeadNum * constInfo.sparseCount + // T轴、s1轴偏移 + n2Idx * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } else if (constInfo.outputLayout == LI_LAYOUT::BSND) { + for (uint32_t s1Idx = s1Start; s1Idx < constInfo.qSeqSize; s1Idx++) { + // B,S1,N2,K + uint64_t indiceOutOffset = bIdx * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount + + s1Idx * constInfo.kHeadNum * constInfo.sparseCount + // B轴、S1轴偏移 + n2Idx * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } + } +} + +template +__aicore__ inline void LightningIndexerKernel::Init(__gm__ uint8_t *query, + __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths, + __gm__ uint8_t *blockTable, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *sparseValues, + __gm__ uint8_t *workspace, const LITilingData *__restrict tiling, + TPipe *tPipe) +{ + if ASCEND_IS_AIV { + tmpBlockIdx = GetBlockIdx(); // vec:0-47 + aiCoreIdx = tmpBlockIdx / 2; + } else { + tmpBlockIdx = GetBlockIdx(); // cube:0-23 + aiCoreIdx = tmpBlockIdx; + } + + InitTilingData(tiling); + InitActualSeqLen(actualSeqLengthsQ, actualSeqLengths); + + // 计算分核 + SplitCore(aiCoreIdx, usedCoreNum, splitCoreInfo); + + pipe = tPipe; + // workspace 内存排布 + // |mm1ResGm(存S)|vec1ResGm(存LD中间结果)|vec1ParamGm(存LD参数) + // |Core0_mm1ResDB0-Core0_mm1ResDB1-Core1_mm1ResDB0....Core23_mm1ResDB0-Core23_mm1ResDB1|Core0_vec1Res... + uint64_t offset = 0; + + // mm1开DoubleBuffer + uint64_t singleCoreMm1ResSize = WS_DOUBLE * constInfo.mBaseSizeAlign * constInfo.s2BaseSize * sizeof(MM1_OUT_T); + mm1ResGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset + aiCoreIdx * singleCoreMm1ResSize)); + offset += GetBlockNum() * singleCoreMm1ResSize; + + // ld流程需要ws大小: [aicnum, 2, CeilDiv(constInfo.mBaseSize, constInfo.gSize), topkOut_*2] + // (aic, 8, 2, 2, 2048) + // (aic, s1_cube, 头尾, idx/value, K) + vec1ResGm.SetGlobalBuffer((__gm__ float *)(workspace + offset)); + offset += GetBlockNum() * constInfo.s1BaseSize * WS_DOUBLE * WS_DOUBLE * BASE_TOPK * sizeof(float); + + // (aic, 8, 2, 16) + // (aic, s1_cube, 头尾，16ele) + vec1ParamGm.SetGlobalBuffer((__gm__ int64_t *)(workspace + offset)); + offset += GetBlockNum() * constInfo.s1BaseSize * WS_DOUBLE * LD_PARAM_NUM * sizeof(int64_t); + + if ASCEND_IS_AIV { + vectorService.InitParams(constInfo, tiling); + indiceOutGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices); + valueOutGm.SetGlobalBuffer((__gm__ K_T *)sparseValues); + weightsGm.SetGlobalBuffer((__gm__ W_T *)weights); + vectorService.InitVec1GlobalTensor(mm1ResGm, vec1ResGm, vec1ParamGm, weightsGm, indiceOutGm, valueOutGm); + } else { + matmulService.InitParams(constInfo); + queryGm.SetGlobalBuffer((__gm__ Q_T *)query); + if constexpr (PAGE_ATTENTION) { + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + } + keyGm.SetGlobalBuffer((__gm__ K_T *)key); + matmulService.InitMm1GlobalTensor(blockTableGm, keyGm, queryGm, mm1ResGm); + } + InitBuffers(); +} + +template +__aicore__ inline void LightningIndexerKernel::GetBN2Idx(uint32_t bN2Idx) +{ + tempLoopInfo.bN2Idx = bN2Idx; + tempLoopInfo.bIdx = bN2Idx / constInfo.kHeadNum; + tempLoopInfo.n2Idx = bN2Idx % constInfo.kHeadNum; +} + +template +__aicore__ inline void LightningIndexerKernel::CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx) +{ + tempLoopInfo.gS1Idx = gS1LoopIdx; + tempLoopInfo.actMBaseSize = constInfo.mBaseSize; + uint32_t remainedGS1Size = tempLoopInfo.actS1Size * constInfo.gSize - tempLoopInfo.gS1Idx * constInfo.mBaseSize; + if (remainedGS1Size <= constInfo.mBaseSize && remainedGS1Size > 0) { + tempLoopInfo.actMBaseSize = tempLoopInfo.mBasicSizeTail; + } + + bool isEnd = (bN2LoopIdx == splitCoreInfo.bN2End) && (gS1LoopIdx == splitCoreInfo.gS1End); + uint32_t s2BlockNum; + if (constInfo.attenMaskFlag) { + s2BlockNum = GetS2BaseBlockNumOnMask(gS1LoopIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2Size); + } else { + s2BlockNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + } + tempLoopInfo.s2LoopEnd = isEnd ? splitCoreInfo.s2End : s2BlockNum - 1; +} + +template +__aicore__ inline void LightningIndexerKernel::CalcGS1LoopParams(uint32_t bN2LoopIdx) +{ + GetBN2Idx(bN2LoopIdx); + GetS1S2ActualSeqLen(tempLoopInfo.bIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2Size); + if ((tempLoopInfo.actS2Size == 0) || (tempLoopInfo.actS1Size == 0)) { + tempLoopInfo.curActSeqLenIsZero = true; + return; + } + tempLoopInfo.curActSeqLenIsZero = false; + tempLoopInfo.s2BasicSizeTail = tempLoopInfo.actS2Size % constInfo.s2BaseSize; + tempLoopInfo.s2BasicSizeTail = + (tempLoopInfo.s2BasicSizeTail == 0) ? constInfo.s2BaseSize : tempLoopInfo.s2BasicSizeTail; + tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize; + tempLoopInfo.mBasicSizeTail = + (tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail; + + uint32_t gS1SplitNum = (tempLoopInfo.actS1Size * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize; + tempLoopInfo.gS1LoopEnd = (bN2LoopIdx == splitCoreInfo.bN2End) ? splitCoreInfo.gS1End : gS1SplitNum - 1; + if constexpr (LAYOUT_T == LI_LAYOUT::BSND) { + if (tempLoopInfo.gS1LoopEnd == gS1SplitNum - 1 && constInfo.qSeqSize > tempLoopInfo.actS1Size) { + tempLoopInfo.needDealActS1LessThanS1 = true; + } + } +} + +template +__aicore__ inline void LightningIndexerKernel::CalcRunInfo(uint32_t loop, + uint32_t s2LoopIdx, LICommon::RunInfo &runInfo) +{ + runInfo.loop = loop; + runInfo.bIdx = tempLoopInfo.bIdx; + runInfo.gS1Idx = tempLoopInfo.gS1Idx; + runInfo.s2Idx = s2LoopIdx; + runInfo.bN2Idx = tempLoopInfo.bN2Idx; + + runInfo.actS1Size = tempLoopInfo.actS1Size; + runInfo.actS2Size = tempLoopInfo.actS2Size; + // 计算实际基本块size + runInfo.actMBaseSize = tempLoopInfo.actMBaseSize; + runInfo.actualSingleProcessSInnerSize = constInfo.s2BaseSize; + uint32_t s2SplitNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + if (runInfo.s2Idx == s2SplitNum - 1) { + runInfo.actualSingleProcessSInnerSize = tempLoopInfo.s2BasicSizeTail; + } + runInfo.actualSingleProcessSInnerSizeAlign = + LICommon::Align((uint32_t)runInfo.actualSingleProcessSInnerSize, LICommon::ConstInfo::BUFFER_SIZE_BYTE_32B); + + runInfo.isFirstS2InnerLoop = s2LoopIdx == splitCoreInfo.s2Start; + runInfo.isLastS2InnerLoop = s2LoopIdx == tempLoopInfo.s2LoopEnd; + runInfo.isAllLoopEnd = (runInfo.bN2Idx == splitCoreInfo.bN2End) && (runInfo.gS1Idx == splitCoreInfo.gS1End) && + (runInfo.s2Idx == splitCoreInfo.s2End); + + if (runInfo.isFirstS2InnerLoop) { + uint64_t actualSeqQPrefixSum; + uint64_t actualSeqKPrefixSum; + if constexpr (LAYOUT_T == LI_LAYOUT::TND) { + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGmQ.GetValue(runInfo.bIdx - 1); + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGm.GetValue(runInfo.bIdx - 1); + } else { // BSND + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.qSeqSize; + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.kSeqSize; + } + uint64_t tndBIdxOffset = actualSeqQPrefixSum * constInfo.qHeadNum * constInfo.headDim; + uint64_t tndKeyBIdxOffset = actualSeqKPrefixSum * constInfo.kHeadNum * constInfo.headDim; + // B,S1,N1(N2,G),D + queryCoreOffset = tndBIdxOffset + runInfo.gS1Idx * constInfo.mBaseSize * constInfo.headDim; + keyCoreOffset = tndKeyBIdxOffset + runInfo.n2Idx * constInfo.headDim; + // B,S1,N1(N2,G)/T,N1(N2,G) + weightsCoreOffset = actualSeqQPrefixSum * constInfo.qHeadNum + runInfo.n2Idx * constInfo.gSize; + // B,S1,N2,k/T,N2,k + indiceOutCoreOffset = actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + + runInfo.n2Idx * constInfo.sparseCount; + } + runInfo.tensorQueryOffset = queryCoreOffset; + runInfo.tensorKeyOffset = keyCoreOffset + runInfo.s2Idx * constInfo.s2BaseSize * constInfo.kHeadNum + * constInfo.headDim; + runInfo.tensorWeightsOffset = weightsCoreOffset; + runInfo.indiceOutOffset = indiceOutCoreOffset; +} + +template +__aicore__ inline void LightningIndexerKernel::Process() +{ + if (usedCoreNum == 0) { + // 没有计算任务，直接清理输出 + ProcessInvalid(); + return; + } + ProcessMain(); + ProcessDecode(); +} + +template +__aicore__ inline void LightningIndexerKernel::ProcessInvalid() +{ + if ASCEND_IS_AIV { + uint32_t aivCoreNum = GetBlockNum() * 2; // 2 means c:v = 1:2 + uint64_t totalOutputSize = + constInfo.batchSize * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount; + uint64_t singleCoreSize = + LICommon::Align((totalOutputSize + aivCoreNum - 1) / aivCoreNum, GM_ALIGN_BYTES / sizeof(OUT_T)); + uint64_t baseSize = tmpBlockIdx * singleCoreSize; + if (baseSize < totalOutputSize) { + uint64_t dealSize = + (baseSize + singleCoreSize <= totalOutputSize) ? singleCoreSize : totalOutputSize - baseSize; + GlobalTensor output = indiceOutGm[baseSize]; + AscendC::InitGlobalMemory(output, dealSize, constInfo.INVALID_IDX); + if (constInfo.returnValue) { + event_t eventIDMTE3ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V)); + SetFlag(eventIDMTE3ToV); + WaitFlag(eventIDMTE3ToV); + + GlobalTensor valueOutGmTmp; + valueOutGmTmp.SetGlobalBuffer((__gm__ uint16_t *)valueOutGm.GetPhyAddr()); + GlobalTensor valueOut = valueOutGmTmp[baseSize]; + + uint16_t negInf = 0; + if constexpr(std::is_same::value) { + negInf = 0xFC00; + } else { + negInf = 0xFF80; + } + AscendC::InitGlobalMemory(valueOut, dealSize, negInf); + } + } + } +} + +template +__aicore__ inline void LightningIndexerKernel::ProcessMain() +{ + if (aiCoreIdx >= usedCoreNum) { + // 无任务核直接返回 + return; + } + + if ASCEND_IS_AIV { + vectorService.AllocEventID(); + CrossCoreSetFlag(constInfo.syncV1C1); + CrossCoreSetFlag(constInfo.syncV1C1); + } else { + matmulService.AllocEventID(); + } + + LICommon::RunInfo runInfo; + uint32_t gloop = 0; + for (uint32_t bN2LoopIdx = splitCoreInfo.bN2Start; bN2LoopIdx <= splitCoreInfo.bN2End; bN2LoopIdx++) { + CalcGS1LoopParams(bN2LoopIdx); + if (tempLoopInfo.curActSeqLenIsZero) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, 0U); + continue; + } + for (uint32_t gS1LoopIdx = splitCoreInfo.gS1Start; gS1LoopIdx <= tempLoopInfo.gS1LoopEnd; gS1LoopIdx++) { + CalcS2LoopParams(bN2LoopIdx, gS1LoopIdx); + for (int s2LoopIdx = splitCoreInfo.s2Start; s2LoopIdx <= tempLoopInfo.s2LoopEnd; s2LoopIdx++) { + ProcessBaseBlock(gloop, s2LoopIdx, runInfo); + ++gloop; + } + splitCoreInfo.s2Start = 0; + } + if (tempLoopInfo.needDealActS1LessThanS1) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, tempLoopInfo.actS1Size); + } + splitCoreInfo.gS1Start = 0; + } + + if ASCEND_IS_AIV { + vectorService.FreeEventID(); + } else { + matmulService.FreeEventID(); + CrossCoreWaitFlag(constInfo.syncV1C1); + CrossCoreWaitFlag(constInfo.syncV1C1); + } +} + +template +__aicore__ inline void LightningIndexerKernel::ProcessBaseBlock(uint32_t loop, + uint64_t s2LoopIdx, LICommon::RunInfo &runInfo) +{ + CalcRunInfo(loop, s2LoopIdx, runInfo); + if ASCEND_IS_AIC { + CrossCoreWaitFlag(constInfo.syncV1C1); + matmulService.ComputeMm1(runInfo); + CrossCoreSetFlag(constInfo.syncC1V1); + } else { + CrossCoreWaitFlag(constInfo.syncC1V1); + vectorService.ProcessVec(runInfo); + CrossCoreSetFlag(constInfo.syncV1C1); + } +} + +template +__aicore__ inline void LightningIndexerKernel::ProcessDecode() +{ + if ASCEND_IS_AIV { + vectorService.InitLDBuffers(pipe); + ICachePreLoad(LD_PREFETCH_LEN); + SyncAll(); + if (splitCoreInfo.isLD) { + vectorService.ProcessLD(); + } + } +} +} // namespace LIKernel +#endif // LIGHTNING_INDEXER_KERNEL_H diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/arch22/lightning_indexer_service_cube.h b/csrc/ascend/attention/lightning_indexer/op_kernel/arch22/lightning_indexer_service_cube.h new file mode 100644 index 000000000..c14335662 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/arch22/lightning_indexer_service_cube.h @@ -0,0 +1,425 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. +  */ + +/*! + * \file lightning_indexer_service_cube.h + * \brief use 5 buffer for matmul l1, better pipeline + */ +#ifndef LIGHTNING_INDEXER_SERVICE_CUBE_H +#define LIGHTNING_INDEXER_SERVICE_CUBE_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../lightning_indexer_common.h" + +namespace LIKernel { +using namespace LICommon; +template +class LightningIndexerServiceCube { +public: + using Q_T = typename LIT::queryType; + using K_T = typename LIT::keyType; + + __aicore__ inline LightningIndexerServiceCube(){}; + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitMm1GlobalTensor(const GlobalTensor &blkTableGm, const GlobalTensor &keyGm, + const GlobalTensor &queryGm, const GlobalTensor &mm1ResGm); + __aicore__ inline void InitParams(const ConstInfo &constInfo); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void ComputeMm1(const LICommon::RunInfo &runInfo); + + static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding; + static constexpr uint64_t KEY_BUF_NUM = 3; + static constexpr uint64_t QUERY_BUF_NUM = 2; + static constexpr uint64_t L0_BUF_NUM = 2; + + static constexpr uint32_t KEY_MTE1_MTE2_EVENT = EVENT_ID2; + static constexpr uint32_t QUERY_MTE1_MTE2_EVENT = EVENT_ID5; // KEY_MTE1_MTE2_EVENT + KEY_BUF_NUM; + static constexpr uint32_t M_MTE1_EVENT = EVENT_ID3; + + static constexpr uint32_t MTE2_MTE1_EVENT = EVENT_ID2; + static constexpr uint32_t MTE1_M_EVENT = EVENT_ID2; + + static constexpr uint64_t M_BASIC_BLOCK = 256; + static constexpr uint64_t D_BASIC_BLOCK = 128; + static constexpr uint64_t S2_BASIC_BLOCK = 256; + + static constexpr uint64_t M_BASIC_BLOCK_L0 = 128; + static constexpr uint64_t D_BASIC_BLOCK_L0 = 128; + static constexpr uint64_t S2_BASIC_BLOCK_L0 = 128; + + static constexpr uint64_t QUERY_BUFFER_OFFSET = M_BASIC_BLOCK * D_BASIC_BLOCK; + static constexpr uint64_t KEY_BUFFER_OFFSET = S2_BASIC_BLOCK * D_BASIC_BLOCK; + static constexpr uint64_t L0AB_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0; + static constexpr uint64_t L0C_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0; + +protected: + __aicore__ inline void Fixp(uint64_t s1gGmOffset, uint64_t s2GmOffset, uint64_t s1gL0RealSize, + uint64_t s2L0RealSize, const LICommon::RunInfo &runInfo); + __aicore__ inline void ComuteL0c(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, const LICommon::RunInfo &runInfo); + __aicore__ inline void LoadKeyToL0b(uint64_t s2L0Offset, uint64_t s2L1RealSize, uint64_t s2L0RealSize, + const LICommon::RunInfo &runInfo); + __aicore__ inline void LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL0Offset, uint64_t s1gL1RealSize, + uint64_t s1gL0RealSize, const LICommon::RunInfo &runInfo); + __aicore__ inline void QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gL1Offset, const LICommon::RunInfo &runInfo); + __aicore__ inline void KeyNd2Nz(uint64_t s2L1RealSize, uint64_t s2GmOffset, const LICommon::RunInfo &runInfo); + __aicore__ inline void KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, const LICommon::RunInfo &runInfo); + GlobalTensor blkTableGm_; + GlobalTensor keyGm_; + GlobalTensor queryGm_; + GlobalTensor mm1ResGm_; + + TBuf bufQL1_; + LocalTensor queryL1_; + TBuf bufKeyL1_; + LocalTensor keyL1_; + + TBuf bufQL0_; + LocalTensor queryL0_; + TBuf bufKeyL0_; + LocalTensor keyL0_; + + TBuf bufL0C_; + LocalTensor cL0_; + + uint64_t keyL1BufIdx_ = 0; + uint64_t queryL1Mte2BufIdx_ = 0; + uint64_t queryL1Mte1BufIdx_ = 0; + uint64_t l0BufIdx_ = 0; + + ConstInfo constInfo_; + +private: + static constexpr bool PAGE_ATTENTION = LIT::pageAttention; +}; + +template +__aicore__ inline void LightningIndexerServiceCube::InitParams(const ConstInfo &constInfo) +{ + constInfo_ = constInfo; +} + +template +__aicore__ inline void LightningIndexerServiceCube::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(bufQL1_, QUERY_BUF_NUM * M_BASIC_BLOCK * D_BASIC_BLOCK * sizeof(Q_T)); + queryL1_ = bufQL1_.Get(); + pipe->InitBuffer(bufKeyL1_, KEY_BUF_NUM * S2_BASIC_BLOCK * D_BASIC_BLOCK * sizeof(K_T)); + keyL1_ = bufKeyL1_.Get(); + + pipe->InitBuffer(bufQL0_, L0_BUF_NUM * M_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0 * sizeof(Q_T)); + queryL0_ = bufQL0_.Get(); + pipe->InitBuffer(bufKeyL0_, L0_BUF_NUM * D_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0 * sizeof(K_T)); + keyL0_ = bufKeyL0_.Get(); + + pipe->InitBuffer(bufL0C_, L0_BUF_NUM * M_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0 * sizeof(float)); + cL0_ = bufL0C_.Get(); +} + +template +__aicore__ inline void +LightningIndexerServiceCube::InitMm1GlobalTensor(const GlobalTensor &blkTableGm, + const GlobalTensor &keyGm, + const GlobalTensor &queryGm, const GlobalTensor &mm1ResGm) +{ + blkTableGm_ = blkTableGm; + keyGm_ = keyGm; + queryGm_ = queryGm; + mm1ResGm_ = mm1ResGm; +} + +template +__aicore__ inline void LightningIndexerServiceCube::ComputeMm1(const LICommon::RunInfo &runInfo) +{ + uint64_t s2GmBaseOffset = runInfo.s2Idx * constInfo_.s2BaseSize; + uint64_t s1gProcessSize = runInfo.actMBaseSize; + uint64_t s2ProcessSize = runInfo.actualSingleProcessSInnerSize; + for (uint64_t s2GmOffset = 0; s2GmOffset < s2ProcessSize; s2GmOffset += S2_BASIC_BLOCK) { + WaitFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM); + uint64_t s2L1RealSize = + s2GmOffset + S2_BASIC_BLOCK > s2ProcessSize ? s2ProcessSize - s2GmOffset : S2_BASIC_BLOCK; + if (PAGE_ATTENTION) { + KeyNd2NzForPA(s2L1RealSize, s2GmBaseOffset + s2GmOffset, runInfo); + }else { + KeyNd2Nz(s2L1RealSize, s2GmOffset, runInfo); + } + + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + // s1gProcessSize当前必定不会超过2倍的s1g basic block + for (uint64_t s1gGmOffset = 0; s1gGmOffset < s1gProcessSize; s1gGmOffset += M_BASIC_BLOCK) { + uint64_t s1gL1RealSize = + s1gGmOffset + M_BASIC_BLOCK > s1gProcessSize ? s1gProcessSize - s1gGmOffset : M_BASIC_BLOCK; + if (runInfo.isFirstS2InnerLoop && s2GmOffset == 0) { + queryL1Mte2BufIdx_++; + queryL1Mte1BufIdx_ = queryL1Mte2BufIdx_; + WaitFlag(QUERY_MTE1_MTE2_EVENT + queryL1Mte2BufIdx_ % QUERY_BUF_NUM); + QueryNd2Nz(s1gL1RealSize, s1gGmOffset, runInfo); + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + } else { + queryL1Mte1BufIdx_ = + queryL1Mte2BufIdx_ - (CeilDiv(s1gProcessSize, M_BASIC_BLOCK) - 1 - (s1gGmOffset > 0)); + } + for (uint64_t s2L1Offset = 0; s2L1Offset < s2L1RealSize; s2L1Offset += S2_BASIC_BLOCK_L0) { + uint64_t s2L0RealSize = + s2L1Offset + S2_BASIC_BLOCK_L0 > s2L1RealSize ? s2L1RealSize - s2L1Offset : S2_BASIC_BLOCK_L0; + for (uint64_t s1gL1Offset = 0; s1gL1Offset < s1gL1RealSize; s1gL1Offset += M_BASIC_BLOCK_L0) { + WaitFlag(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM); + uint64_t s1gL0RealSize = + s1gL1Offset + M_BASIC_BLOCK_L0 > s1gL1RealSize ? s1gL1RealSize - s1gL1Offset : M_BASIC_BLOCK_L0; + LoadQueryToL0a(s1gGmOffset, s1gL1Offset, s1gL1RealSize, s1gL0RealSize, runInfo); + LoadKeyToL0b(s2L1Offset, s2L1RealSize, s2L0RealSize, runInfo); + + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + + ComuteL0c(s1gL0RealSize, s2L0RealSize, runInfo); + + SetFlag(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM); + + Fixp(s1gGmOffset + s1gL1Offset, s2GmOffset + s2L1Offset, s1gL0RealSize, s2L0RealSize, runInfo); + l0BufIdx_++; + } + } + if (s2GmOffset + S2_BASIC_BLOCK >= s2ProcessSize && runInfo.isLastS2InnerLoop) { + SetFlag(QUERY_MTE1_MTE2_EVENT + queryL1Mte1BufIdx_ % QUERY_BUF_NUM); + } + } + + SetFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM); + keyL1BufIdx_++; + } +} + +template +__aicore__ inline void LightningIndexerServiceCube::KeyNd2Nz(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const LICommon::RunInfo &runInfo) +{ + uint64_t s2L1Offset = 0; + while (s2L1Offset < s2L1RealSize) { + uint64_t keyGmOffset = runInfo.tensorKeyOffset + (s2GmOffset + s2L1Offset) * constInfo_.headDim; + // 搬运按照S2_BASIC_BLOCK_L0*D_BASIC_BLOCK_L0的方式在l1上排布, 方便后续mte1 + // 根据s2的offset判断当前属于前一个L0分型还是后一个L0分型，暂时只支持两个分型 + uint64_t s2Mte2Size = (s2L1RealSize <= S2_BASIC_BLOCK_L0 || s2L1Offset >= S2_BASIC_BLOCK_L0) ? + s2L1RealSize - s2L1Offset : + S2_BASIC_BLOCK_L0 - s2L1Offset; + + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2Mte2Size; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = s2L1Offset >= S2_BASIC_BLOCK_L0 ? + CeilAlign(s2L1RealSize - S2_BASIC_BLOCK_L0, (uint64_t)BLOCK_CUBE) : + (s2L1RealSize > S2_BASIC_BLOCK_L0 ? + S2_BASIC_BLOCK_L0 : + CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE)); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET + + (s2L1Offset >= S2_BASIC_BLOCK_L0 ? + S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0 + (s2L1Offset - S2_BASIC_BLOCK_L0) * BLOCK_CUBE : + s2L1Offset * BLOCK_CUBE)], + keyGm_[keyGmOffset], nd2nzPara); + + s2L1Offset += s2Mte2Size; + } +} + +// blkNum, blkSize, N2, D +template +__aicore__ inline void LightningIndexerServiceCube::KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const LICommon::RunInfo &runInfo) +{ + uint64_t s2L1Offset = 0; + while (s2L1Offset < s2L1RealSize) { + uint64_t s2BlkId = (s2L1Offset + s2GmOffset) / constInfo_.kCacheBlockSize; + uint64_t s2BlkOffset = (s2L1Offset + s2GmOffset) % constInfo_.kCacheBlockSize; + uint64_t keyGmOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId) * + constInfo_.kCacheBlockSize * constInfo_.kHeadNum * constInfo_.headDim + + s2BlkOffset * constInfo_.headDim; + // 搬运按照S2_BASIC_BLOCK_L0*D_BASIC_BLOCK_L0的方式在l1上排布, 方便后续mte1 + // 根据s2的offset判断当前属于前一个L0分型还是后一个L0分型，暂时只支持两个分型 + uint64_t s2Mte2Size = (s2L1RealSize <= S2_BASIC_BLOCK_L0 || s2L1Offset >= S2_BASIC_BLOCK_L0) ? + s2L1RealSize - s2L1Offset : + S2_BASIC_BLOCK_L0 - s2L1Offset; + s2Mte2Size = s2BlkOffset + s2Mte2Size >= constInfo_.kCacheBlockSize ? constInfo_.kCacheBlockSize - s2BlkOffset : + s2Mte2Size; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2Mte2Size; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = s2L1Offset >= S2_BASIC_BLOCK_L0 ? + CeilAlign(s2L1RealSize - S2_BASIC_BLOCK_L0, (uint64_t)BLOCK_CUBE) : + (s2L1RealSize > S2_BASIC_BLOCK_L0 ? + S2_BASIC_BLOCK_L0 : + CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE)); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET + + (s2L1Offset >= S2_BASIC_BLOCK_L0 ? + S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0 + (s2L1Offset - S2_BASIC_BLOCK_L0) * BLOCK_CUBE : + s2L1Offset * BLOCK_CUBE)], + keyGm_[keyGmOffset], nd2nzPara); + + s2L1Offset += s2Mte2Size; + } +} + +// batch, s1, n2, g, d +template +__aicore__ inline void LightningIndexerServiceCube::QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gGmOffset, + const LICommon::RunInfo &runInfo) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s1gL1RealSize; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(s1gL1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + // 默认一块buf最多放两份 + DataCopy(queryL1_[(queryL1Mte2BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET], + queryGm_[runInfo.tensorQueryOffset + s1gGmOffset * constInfo_.headDim], nd2nzPara); +} + +template +__aicore__ inline void LightningIndexerServiceCube::LoadQueryToL0a(uint64_t s1gGmOffset, + uint64_t s1gL1Offset, uint64_t s1gL1RealSize, + uint64_t s1gL0RealSize, const LICommon::RunInfo &runInfo) +{ + LoadData3DParamsV2 loadData3DParams; + // SetFmatrixParams + loadData3DParams.l1H = CeilDiv(s1gL1RealSize, BLOCK_CUBE); // Hin=M1=8 + loadData3DParams.l1W = BLOCK_CUBE; // Win=M0 + loadData3DParams.channelSize = constInfo_.headDim; // Cin=K + + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + // SetLoadToA0Params + loadData3DParams.mExtension = CeilAlign(s1gL0RealSize, BLOCK_CUBE); // M height维度目的 + loadData3DParams.kExtension = constInfo_.headDim; // K width维度目的 + loadData3DParams.mStartPt = s1gL1Offset; + loadData3DParams.kStartPt = 0; + loadData3DParams.strideW = 1; + loadData3DParams.strideH = 1; + loadData3DParams.filterW = 1; + loadData3DParams.filterSizeW = (1 >> 8) & 255; + loadData3DParams.filterH = 1; + loadData3DParams.filterSizeH = (1 >> 8) & 255; + loadData3DParams.dilationFilterW = 1; + loadData3DParams.dilationFilterH = 1; + loadData3DParams.enTranspose = 0; + loadData3DParams.fMatrixCtrl = 0; + + LoadData(queryL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], + queryL1_[(queryL1Mte1BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET], + loadData3DParams); +} + +template +__aicore__ inline void LightningIndexerServiceCube::LoadKeyToL0b(uint64_t s2L1Offset, + uint64_t s2L1RealSize, uint64_t s2L0RealSize, + const LICommon::RunInfo &runInfo) +{ + uint64_t keyL1Offset = s2L1Offset >= S2_BASIC_BLOCK_L0 ? S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0 : 0; + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = CeilDiv(s2L0RealSize, BLOCK_CUBE) * CeilDiv(constInfo_.headDim, BLOCK_CUBE); + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = false; + LoadData(keyL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], + keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET + keyL1Offset], loadData2DParams); +} + +template +__aicore__ inline void LightningIndexerServiceCube::ComuteL0c(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, + const LICommon::RunInfo &runInfo) +{ + MmadParams mmadParams; + mmadParams.m = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + mmadParams.n = s2L0RealSize; + mmadParams.k = constInfo_.headDim; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = 0b11; + Mmad(cL0_[(l0BufIdx_ % L0_BUF_NUM) * L0C_BUFFER_OFFSET], queryL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], + keyL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } +} + +template +__aicore__ inline void LightningIndexerServiceCube::Fixp(uint64_t s1gGmOffset, + uint64_t s2GmOffset, uint64_t s1gL0RealSize, + uint64_t s2L0RealSize, const LICommon::RunInfo &runInfo) +{ + AscendC::DataCopyCO12DstParams intriParams; + intriParams.mSize = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + intriParams.nSize = s2L0RealSize; + intriParams.dstStride = runInfo.actualSingleProcessSInnerSizeAlign; + intriParams.srcStride = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + // set mode according to dtype + intriParams.quantPre = QuantMode_t::NoQuant; + intriParams.nz2ndEn = true; + intriParams.unitFlag = 0b11; // 3 unitflag + intriParams.reluPre = 1; + AscendC::SetFixpipeNz2ndFlag(1, 1, 1); + AscendC::DataCopy(mm1ResGm_[(runInfo.loop % 2) * constInfo_.mBaseSizeAlign * constInfo_.s2BaseSize + + s1gGmOffset * intriParams.dstStride + s2GmOffset], + cL0_[(l0BufIdx_ % L0_BUF_NUM) * L0C_BUFFER_OFFSET], intriParams); +} + +template +__aicore__ inline void LightningIndexerServiceCube::AllocEventID() +{ + SetMMLayoutTransform(true); + SetFlag(KEY_MTE1_MTE2_EVENT + 0); + SetFlag(KEY_MTE1_MTE2_EVENT + 1); + SetFlag(KEY_MTE1_MTE2_EVENT + 2); + + SetFlag(QUERY_MTE1_MTE2_EVENT + 0); + SetFlag(QUERY_MTE1_MTE2_EVENT + 1); + + SetFlag(M_MTE1_EVENT + 0); + SetFlag(M_MTE1_EVENT + 1); +} + +template +__aicore__ inline void LightningIndexerServiceCube::FreeEventID() +{ + SetMMLayoutTransform(false); + WaitFlag(KEY_MTE1_MTE2_EVENT + 0); + WaitFlag(KEY_MTE1_MTE2_EVENT + 1); + WaitFlag(KEY_MTE1_MTE2_EVENT + 2); + + WaitFlag(QUERY_MTE1_MTE2_EVENT + 0); + WaitFlag(QUERY_MTE1_MTE2_EVENT + 1); + + WaitFlag(M_MTE1_EVENT + 0); + WaitFlag(M_MTE1_EVENT + 1); +} +} // namespace LIKernel +#endif diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/arch22/lightning_indexer_service_vector.h b/csrc/ascend/attention/lightning_indexer/op_kernel/arch22/lightning_indexer_service_vector.h new file mode 100644 index 000000000..60a881e34 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/arch22/lightning_indexer_service_vector.h @@ -0,0 +1,706 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. +  */ + +/*! + * \file lightning_indexer_service_vector.h + * \brief + */ +#ifndef LIGHTNING_INDEXER_SERVICE_VECTOR_H +#define LIGHTNING_INDEXER_SERVICE_VECTOR_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../lightning_indexer_common.h" +#include "lightning_indexer_vector.h" + +namespace LIKernel { +using namespace LICommon; +using namespace LIServiceVec; +constexpr uint32_t BASE_TOPK = 2048; +constexpr uint32_t SPARSE_COUNT_4K = 4096; +constexpr uint32_t LD_PARAM_NUM = 16; +constexpr uint32_t EVENTID_V_TO_MTE2_PING = 0; +constexpr uint32_t EVENTID_V_TO_MTE2_PONG = 1; +constexpr uint32_t EVENTID_V_TO_MTE2_TMPUB = 2; + +// 主模板：Q_T必选，W_T可选（默认void），无论W_T传什么，默认weightsType=Q_T +template +struct LightningIndexerTypeTraits { + using weightsType = Q_T; // 默认：weightsType绑定Q_T +}; + +// 偏特化1：固定第二个参数W_T=float，Q_T保留泛型 +template +struct LightningIndexerTypeTraits { + using weightsType = float; // W_T=float时，强制weightsType为float +}; + +template +class LightningIndexerServiceVector { +public: + // =================================类型定义区================================= + // 中间计算数据类型为float，高精度模式 + static constexpr bool DT_W_FLAG = LIT::weightsTypeFlag; + using Q_T = typename LIT::queryType; + using K_T = typename LIT::keyType; + static constexpr LI_LAYOUT LAYOUT_T = LIT::layout; + using W_T = typename LightningIndexerTypeTraits::type>::weightsType; + + // MM输出数据类型, 当前只支持float + using MM1_OUT_T = float; + + __aicore__ inline LightningIndexerServiceVector(){}; + __aicore__ inline void ProcessVec(const LICommon::RunInfo &info); + __aicore__ inline void ProcessLD(); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitParams(const struct LICommon::ConstInfo &constInfo, + const LITilingData *__restrict tilingData); + __aicore__ inline void InitVec1GlobalTensor(GlobalTensor mm1ResGm, GlobalTensor vec1ResGm, + GlobalTensor vec1ParamGm, GlobalTensor weightsGm, + GlobalTensor indiceOutGm, GlobalTensor valueOutGm); + __aicore__ inline void CleanInvalidOutput(int64_t invalidS1offset); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void InitLDBuffers(TPipe *pipe); + +protected: + GlobalTensor mm1ResGm; + GlobalTensor vec1ResGm; + GlobalTensor vec1ParamGm; + GlobalTensor weightsGm; + GlobalTensor indiceOutGm; + GlobalTensor valueOutGm; + // =================================常量区================================= + +private: + // ================================Local Buffer区==================================== + // queue + TQue outQueue_; + + // tmp buff for vector + TBuf sortOutBuf_; + TBuf tmpBuf_; + TBuf indexBuf_; + TBuf reduceOutBuf_; + TBuf brcBuf_; + TBuf paramBuf_; + + // tmp buff for LD + TBuf<> ldToBeMrgBuf_; + TBuf<> ldTmpBuf_; + TBuf<> ldOutValueBuf_; + TBuf<> ldOutIdxBuf_; + + LocalTensor tmpUb_; + LocalTensor globalTopkIndice_; + LocalTensor globalTopkUb_; + LocalTensor SortedBasicBlock_; + + int32_t blockId_ = -1; + // para for vector + int32_t groupInner_ = 0; + int32_t globalTopkNum_ = 0; + int64_t blockS2StartIdx_ = 0; + int32_t gSize_ = 0; + int32_t kHeadNum_ = 0; + int32_t s1BaseSize_ = 0; + int32_t s2BaseSize_ = 0; + + // para for LD + uint32_t mrgListNum_ = 4; + uint32_t paramNum_ = 16; + int32_t virTopK = 0; + + constexpr static uint32_t REDUCE_BANK_CONFLICT_OFFSETS = 256; + constexpr static uint32_t REDUCE_BANK_CONFLICT_NUM = REDUCE_BANK_CONFLICT_OFFSETS / sizeof(float); + + struct LICommon::ConstInfo constInfo_; +}; + +template +__aicore__ inline void LightningIndexerServiceVector::InitBuffers(TPipe *pipe) +{ + uint32_t outNeedBufSize = (BASE_TOPK * 2) * 2 * sizeof(float); + uint32_t reduceCacheSize = REDUCE_BANK_CONFLICT_OFFSETS + groupInner_ * s2BaseSize_ * sizeof(float); + outNeedBufSize = reduceCacheSize > outNeedBufSize ? reduceCacheSize : outNeedBufSize; + virTopK = constInfo_.isSparseCountOver2K ? constInfo_.sparseCount : BASE_TOPK; + + pipe->InitBuffer(outQueue_, 1, outNeedBufSize); // 32KB extract + // 68KB 在搬运cube核计算得到的结果和weight时，分成两块34KB，用于db；在mrgsort时，用作临时UB + pipe->InitBuffer(tmpBuf_, (groupInner_ * s2BaseSize_ + s2BaseSize_) * 2 * sizeof(float)); + pipe->InitBuffer(sortOutBuf_, CeilDiv(s1BaseSize_, 2) * virTopK * 2 * sizeof(float)); // 64KB + pipe->InitBuffer(indexBuf_, s2BaseSize_ * sizeof(int32_t)); // 2KB + pipe->InitBuffer(reduceOutBuf_, s2BaseSize_ * 2 * sizeof(float)); // 4KB + pipe->InitBuffer(brcBuf_, groupInner_ * 8 * sizeof(float)); + pipe->InitBuffer(paramBuf_, LD_PARAM_NUM * sizeof(int64_t)); + + tmpUb_ = tmpBuf_.Get(); + globalTopkIndice_ = indexBuf_.Get(); + globalTopkUb_ = sortOutBuf_.Get(); + SortedBasicBlock_ = globalTopkUb_[virTopK * 2 * 2]; + globalTopkNum_ = 0; + + // 基本块执行前初始化UB和GM + // step1. 初始化一个有序索引 0 - s2BaseSize_ + ArithProgression(globalTopkIndice_, 0, 1, s2BaseSize_); + // step2. globalTopkUb_ [CeilDiv(s1BaseSize_, 2), BASE_TOPK, 2] -inf,-1 + InitSortOutBuf(globalTopkUb_, CeilDiv(s1BaseSize_, 2) * virTopK * 2); + + // step3. 初始化vec1ParamGm，是否进行LD的标志位设为-1(needFd=-1) + // vec1ResIn32Gm = [aic, 2, s1BaseSize_, 16] int32 + // ws清零 [needFd, s2AcSeq, s2Start, s2End, isS2End, bn2idx, s1Idx, ......] + LocalTensor tmpBuff = outQueue_.AllocTensor(); + Duplicate(tmpBuff.template ReinterpretCast(), -1, 2 * (s1BaseSize_ / 2) * paramNum_ * 2); + outQueue_.EnQue(tmpBuff); + tmpBuff = outQueue_.DeQue(); + int64_t wsInfoOffset = (blockId_ / 2) * s1BaseSize_ * 2 * paramNum_ + // 2个AIV共同地址偏移 + (blockId_ % 2) * (s1BaseSize_ / 2) * 2 * paramNum_; // 每个AIV的地址偏移，S1方向 + DataCopyPad(vec1ParamGm[wsInfoOffset], tmpBuff.template ReinterpretCast(), + {1, static_cast((s1BaseSize_ / 2) * 2 * paramNum_ * sizeof(int64_t)), 0, 0}); + outQueue_.FreeTensor(tmpBuff); +} + +template +__aicore__ inline void LightningIndexerServiceVector::InitLDBuffers(TPipe *pipe) +{ + pipe->Reset(); + pipe->InitBuffer(ldToBeMrgBuf_, 2 * BASE_TOPK * mrgListNum_ * sizeof(float)); // 2：value + index + pipe->InitBuffer(ldTmpBuf_, 2 * BASE_TOPK * mrgListNum_ * sizeof(float)); // 2：value + index + pipe->InitBuffer(ldOutValueBuf_, BASE_TOPK * sizeof(float)); + pipe->InitBuffer(ldOutIdxBuf_, BASE_TOPK * sizeof(int32_t)); +} + +template +__aicore__ inline void LightningIndexerServiceVector::InitParams(const struct LICommon::ConstInfo &constInfo, + const LITilingData *__restrict tilingData) +{ + this->constInfo_ = constInfo; + blockS2StartIdx_ = 0; + gSize_ = constInfo.gSize; + // define N2 para + kHeadNum_ = constInfo.kHeadNum; + // define MMBase para + s1BaseSize_ = constInfo.s1BaseSize; + s2BaseSize_ = constInfo.s2BaseSize; + + // group ub 切分因子当前按照UB空间强制为16 + groupInner_ = 16; + + blockId_ = GetBlockIdx(); +} + +template +__aicore__ inline void +LightningIndexerServiceVector::InitVec1GlobalTensor(GlobalTensor mm1ResGm, + GlobalTensor vec1ResGm, + GlobalTensor vec1ParamGm, GlobalTensor weightsGm, + GlobalTensor indiceOutGm, GlobalTensor valueOutGm) +{ + this->mm1ResGm = mm1ResGm; + this->vec1ResGm = vec1ResGm; + this->vec1ParamGm = vec1ParamGm; + this->weightsGm = weightsGm; + this->indiceOutGm = indiceOutGm; + this->valueOutGm = valueOutGm; +} + +template +__aicore__ inline void LightningIndexerServiceVector::AllocEventID() +{ + SetFlag(EVENTID_V_TO_MTE2_PING); + SetFlag(EVENTID_V_TO_MTE2_PONG); + SetFlag(EVENTID_V_TO_MTE2_TMPUB); +} + +template +__aicore__ inline void LightningIndexerServiceVector::FreeEventID() +{ + WaitFlag(EVENTID_V_TO_MTE2_PING); + WaitFlag(EVENTID_V_TO_MTE2_PONG); + WaitFlag(EVENTID_V_TO_MTE2_TMPUB); +} + +template +__aicore__ inline void LightningIndexerServiceVector::CleanInvalidOutput(int64_t invalidS1offset) +{ + // init -1 and copy to output + LocalTensor valueULocal = outQueue_.AllocTensor(); + LocalTensor idxULocal1 = valueULocal.template ReinterpretCast(); + Duplicate(idxULocal1, constInfo_.INVALID_IDX, constInfo_.sparseCount); + outQueue_.EnQue(valueULocal); + valueULocal = outQueue_.DeQue(); + LIServiceVec::CopyOut(indiceOutGm[invalidS1offset], idxULocal1, constInfo_.sparseCount); + outQueue_.FreeTensor(valueULocal); + + if (constInfo_.returnValue) { + uint16_t negInf = 0; + if constexpr(std::is_same::value) { + negInf = 0xFC00; + } else { + negInf = 0xFF80; + } + LocalTensor valueULocal = outQueue_.AllocTensor(); + Duplicate(valueULocal, negInf, constInfo_.sparseCount); + outQueue_.EnQue(valueULocal); + valueULocal = outQueue_.DeQue(); + GlobalTensor valueOutGmTmp; + valueOutGmTmp.SetGlobalBuffer((__gm__ uint16_t *)valueOutGm.GetPhyAddr()); + LIServiceVec::CopyOut(valueOutGmTmp[invalidS1offset], valueULocal, constInfo_.sparseCount); + outQueue_.FreeTensor(valueULocal); + } +} + +template +__aicore__ inline void LightningIndexerServiceVector::ProcessVec(const LICommon::RunInfo &info) +{ + int32_t cuBaseS1Idx = info.gS1Idx * s1BaseSize_; + int32_t cuBaseS2Idx = info.s2Idx * s2BaseSize_; + + // 计算基本块基地址偏移 偶数循环 -> 0 + aic_offset 奇数循环 -> 512*512 + aic_offset + int64_t mmGmOffset = (info.loop % 2) * (constInfo_.mBaseSizeAlign * s2BaseSize_); + // (B,S1,N1,1);(T,N1,1) -> (B,S1,N2,G,1) 当前只切分到S1轴 + int64_t weightGmOffset = info.tensorWeightsOffset + cuBaseS1Idx * kHeadNum_ * gSize_; + + PipeBarrier(); + // cuS1BeginIdxPerAiv: 每个AIV的S1起始偏移 + int32_t cuS1BeginIdxPerAiv = cuBaseS1Idx; + int32_t cuS1ProcNum = + cuS1BeginIdxPerAiv + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + // cuS1ProcNumPerAiv: 每个AIv的S1计算量 + int32_t cuS1ProcNumPerAiv = blockId_ % 2 == 0 ? CeilDiv(cuS1ProcNum, 2) : (cuS1ProcNum / 2); + cuS1BeginIdxPerAiv += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2); + + // 基本块基地址偏移奇数核加一个S1地址偏移 + weightGmOffset += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2) * kHeadNum_ * gSize_; + mmGmOffset += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2) * gSize_ * info.actualSingleProcessSInnerSizeAlign; + + // cut G + int32_t outerG = CeilDiv(gSize_, groupInner_); + + // 非首个基本块, M(S1)轴发生切换需要初始化 + if (info.loop != 0 && info.s2Idx == 0) { + // globalTopkUb_ value,index=-inf,-1 + InitSortOutBuf(globalTopkUb_, CeilDiv(s1BaseSize_, 2) * virTopK * 2); + blockS2StartIdx_ = 0; + } else if (info.loop == 0) { + blockS2StartIdx_ = info.s2Idx; + } + // cuRealAcSeq: 当前基本块S1对应的AcSeq + int32_t cuRealAcSeq = info.actS2Size; + if (constInfo_.attenMaskFlag) { + // attenMask true场景 + cuRealAcSeq = info.actS2Size - (info.actS1Size - cuS1BeginIdxPerAiv); + } + LocalTensor reduceOutBuff = reduceOutBuf_.Get(); + LocalTensor brcBuf = brcBuf_.Get(); + // LD输出S1方向偏移，保证2个Vector输出的内容连续 + uint32_t ldS1Offset = (blockId_ % 2 == 0) ? s1BaseSize_ / 2 - cuS1ProcNumPerAiv : 0; + for (int innerS1Idx = 0; innerS1Idx < cuS1ProcNumPerAiv; innerS1Idx++) { + if (constInfo_.attenMaskFlag) { + cuRealAcSeq += 1; + } + int32_t cuS2Len = cuBaseS2Idx + s2BaseSize_ >= cuRealAcSeq ? cuRealAcSeq - cuBaseS2Idx : s2BaseSize_; + int32_t cuS1Idx = cuS1BeginIdxPerAiv + innerS1Idx; + if (cuRealAcSeq > 0 && cuS2Len > 0) { + int32_t cuS2LenVecAlign = CeilDiv(cuS2Len, s2BaseSize_) * s2BaseSize_; + int32_t mmUbStride = (cuS2LenVecAlign - info.actualSingleProcessSInnerSizeAlign) / B32_BLOCK_ALIGN_NUM; + LocalTensor reduceOutInner = reduceOutBuff[s2BaseSize_]; + PipeBarrier(); + LocalTensor reduceCacheBuf = outQueue_.AllocTensor(); + if (constInfo_.isSparseCountOver2K) { + WaitFlag(EVENTID_V_TO_MTE2_TMPUB); + } + for (int outerGidx = 0; outerGidx < outerG; outerGidx++) { + int32_t procGnum = outerGidx != outerG - 1 ? groupInner_ : gSize_ - outerGidx * groupInner_; + + int32_t pingpong = outerGidx % 2; + LocalTensor dbTmpUb = tmpUb_[pingpong * (groupInner_ * s2BaseSize_ + s2BaseSize_)]; + LocalTensor weightsInUb = dbTmpUb[procGnum * s2BaseSize_]; + WaitFlag(pingpong); + LocalTensor weightsInTUb = weightsInUb.template ReinterpretCast(); + if constexpr (!IsSameType::value) { + weightsInTUb = weightsInTUb[groupInner_]; + } + int64_t mmGmAllOffet = mmGmOffset + innerS1Idx * gSize_ * info.actualSingleProcessSInnerSizeAlign + + outerGidx * groupInner_ * info.actualSingleProcessSInnerSizeAlign; + int64_t weightGmAllOffset = weightGmOffset + innerS1Idx * gSize_ + outerGidx * groupInner_; + + LIServiceVec::CopyIn(dbTmpUb, weightsInTUb, mm1ResGm, weightsGm, mmGmAllOffet, weightGmAllOffset, + procGnum, info.actualSingleProcessSInnerSizeAlign, mmUbStride); + + SetFlag(pingpong); + WaitFlag(pingpong); + LIServiceVec::DoScale(reduceCacheBuf[REDUCE_BANK_CONFLICT_NUM], dbTmpUb, weightsInUb, weightsInTUb, + brcBuf, procGnum, s2BaseSize_, outerGidx); + // confused reduceOp in DoScale + // neednot use LIServiceVec::doReduce(mmInUb, reduceOutInner, procGnum, (s2BaseSize_+8)); + SetFlag(pingpong); + } + + int32_t gRedCnt = groupInner_ > gSize_ ? gSize_ : groupInner_; + bool isS2End = cuBaseS2Idx + s2BaseSize_ >= cuRealAcSeq; + LIServiceVec::DoReduce(reduceCacheBuf[REDUCE_BANK_CONFLICT_NUM], reduceOutInner, gRedCnt, s2BaseSize_); + outQueue_.FreeTensor(reduceCacheBuf); + + LocalTensor sortScoreUb = reduceOutBuff; + LocalTensor sortIndiceUb = reduceOutBuff[cuS2LenVecAlign]; + Duplicate(sortScoreUb.template ReinterpretCast(), LIServiceVec::NEG_INF, cuS2LenVecAlign); + PipeBarrier(); + Adds(sortScoreUb, reduceOutInner, 0.0f, cuS2Len); + PipeBarrier(); + LocalTensor sortIndiceUbInt = sortIndiceUb.template ReinterpretCast(); + // 无效数据索引填充为-1 + if (cuS2LenVecAlign != cuS2Len) { + Duplicate(sortIndiceUbInt, -1, cuS2LenVecAlign); + } + PipeBarrier(); + Adds(sortIndiceUbInt, globalTopkIndice_, static_cast(cuBaseS2Idx), cuS2Len); + PipeBarrier(); + + LocalTensor tmpSortBuf = outQueue_.AllocTensor(); + if (info.actS1Size > 4 || constInfo_.isSparseCountOver2K) { + // info.actS1Size > 4 则单个vector核内处理的 s1>2，缓存方案无法处理 + LIServiceVec::SortAll(reduceOutBuff, tmpSortBuf, + cuS2LenVecAlign); // cuS2LenVecAlign <= s2BaseSize_, fill -inf + PipeBarrier(); + LocalTensor UbTmpSort = constInfo_.isSparseCountOver2K ? tmpUb_ : tmpSortBuf; + LIServiceVec::MergeSort(globalTopkUb_[innerS1Idx * virTopK * 2], virTopK, reduceOutBuff, + cuS2LenVecAlign, UbTmpSort); + } else { + int64_t globalTopkUbCacheIdx = (info.s2Idx - blockS2StartIdx_) % 4; + Sort( + SortedBasicBlock_[innerS1Idx * BASE_TOPK * 2 + globalTopkUbCacheIdx * s2BaseSize_ * 2], + reduceOutBuff, sortIndiceUbInt.template ReinterpretCast(), tmpSortBuf, + cuS2LenVecAlign / 32); + AscendC::PipeBarrier(); + // 缓存4块512或者S2结束, 需要进行精排 + if (globalTopkUbCacheIdx == 3 || isS2End || info.isAllLoopEnd) { + LocalTensor tt = SortedBasicBlock_[innerS1Idx * BASE_TOPK * 2]; + // 前4块直接精排覆盖到globalTopkUb_ + if (info.s2Idx - blockS2StartIdx_ < 4) { + MrgBasicBlock(globalTopkUb_[innerS1Idx * BASE_TOPK * 2], tt, + static_cast(globalTopkUbCacheIdx + 1), s2BaseSize_); + } else { // 后面缓存在 SortedBasicBlock_, 先精排, 再merge到globalTopkUb_ + if (globalTopkUbCacheIdx > 0) { + MrgBasicBlock(tmpSortBuf, tt, static_cast(globalTopkUbCacheIdx + 1), s2BaseSize_); + PipeBarrier(); + DataCopy(SortedBasicBlock_[innerS1Idx * BASE_TOPK * 2], tmpSortBuf, + (globalTopkUbCacheIdx + 1) * s2BaseSize_ * 2); + } + PipeBarrier(); + SparseTopK(globalTopkUb_[innerS1Idx * BASE_TOPK * 2], + SortedBasicBlock_[innerS1Idx * BASE_TOPK * 2], tmpSortBuf, BASE_TOPK, + s2BaseSize_ * (globalTopkUbCacheIdx + 1)); + } + } + } + if (constInfo_.isSparseCountOver2K) { + SetFlag(EVENTID_V_TO_MTE2_TMPUB); + } + + PipeBarrier(); + outQueue_.FreeTensor(tmpSortBuf); + + bool needCopyOutGm = blockS2StartIdx_ == 0 && isS2End; + + // 中间结果保存 + bool needCopyWsGm = info.isAllLoopEnd || isS2End; + + if (needCopyOutGm) { + int64_t offset = (constInfo_.sparseCount <= SPARSE_COUNT_4K) ? virTopK : constInfo_.sparseCount / 2; + int64_t copyLen = (constInfo_.sparseCount <= SPARSE_COUNT_4K) + ? constInfo_.sparseCount + : constInfo_.sparseCount / 2; + int64_t copyNum = (constInfo_.sparseCount <= SPARSE_COUNT_4K) ? 1 : 2; + for (int64_t i = 0; i < copyNum; i++) { + LocalTensor outValueUb = outQueue_.AllocTensor(); + LocalTensor outIdxUb = outValueUb[offset].template ReinterpretCast(); + Extract(outValueUb, outIdxUb, + globalTopkUb_[innerS1Idx * virTopK * 2 + 2 * i * offset], (offset /32)); + + LocalTensor valueULocal1 = outValueUb.template ReinterpretCast(); + if (constInfo_.returnValue) { + PipeBarrier(); + Cast(valueULocal1, outValueUb, RoundMode::CAST_ROUND, copyLen); + } + + LocalTensor idxULocal1 = outValueUb[offset].template ReinterpretCast(); + outQueue_.EnQue(outValueUb); + outValueUb = outQueue_.DeQue(); + + LIServiceVec::CopyOut(indiceOutGm[info.indiceOutOffset + cuS1Idx * + constInfo_.sparseCount + i * offset], + idxULocal1, copyLen); + if (constInfo_.returnValue) { + LIServiceVec::CopyOut(valueOutGm[info.indiceOutOffset + cuS1Idx * + constInfo_.sparseCount + i * offset], + valueULocal1, copyLen); + } + outQueue_.FreeTensor(outValueUb); + } + } else if (needCopyWsGm) { + // vec1Res Gm = [aic, s1BaseSize_, 2, 2, topkOut_] float32 + // vec1Param Gm = [aic, s1BaseSize_, 2, 16] int64 + // 16 = [needFd, s2AcSeq, s2Start, s2End, isS2End, bn2idx, s1Idx, S1ProcNum, ......] + + int64_t wsOffset = (blockId_ / 2) * s1BaseSize_ * 2 * 2 * BASE_TOPK + // 2个AIV共同地址偏移 + (blockId_ % 2) * (s1BaseSize_ / 2) * 2 * 2 * BASE_TOPK + // 每个AIV的地址偏移，S1方向 + (ldS1Offset + innerS1Idx) * 2 * 2 * BASE_TOPK; + int64_t wsInfoOffset = (blockId_ / 2) * s1BaseSize_ * 2 * paramNum_ + // 2个AIV共同地址偏移 + (blockId_ % 2) * (s1BaseSize_ / 2) * 2 * paramNum_ + // 每个AIV的地址偏移，S1方向 + (ldS1Offset + innerS1Idx) * 2 * paramNum_; + + LocalTensor tmpiBuff = paramBuf_.Get(); + SetWaitFlag(HardEvent::MTE3_S); + tmpiBuff.SetValue(0, static_cast(1)); + tmpiBuff.SetValue(1, static_cast(cuRealAcSeq)); + tmpiBuff.SetValue(2, static_cast(blockS2StartIdx_)); + tmpiBuff.SetValue(3, static_cast(cuBaseS2Idx + cuS2Len)); + tmpiBuff.SetValue(4, static_cast(isS2End)); + tmpiBuff.SetValue(5, static_cast(info.bN2Idx)); + tmpiBuff.SetValue(6, static_cast(cuS1Idx)); + tmpiBuff.SetValue(7, static_cast(cuS1ProcNum)); + tmpiBuff.SetValue(8, static_cast(info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount)); + // 写入头尾判断 + // [head, tail] + // head: 与前面规约，与前后规约 + // tail: 与后面规约 + bool isTailReduce = blockS2StartIdx_ == 0; // 一定是isLastTile + // WS偏移规则 blockS2StartIdx_ != 0 + // 跟前面块做规约 写到0偏移 不用做计算 blockS2StartIdx_ == 0 and !isS2End + // 跟后面块做规约 写到1偏移 需要 + s1BaseSize_, BASE_TOPK*2 + if (isTailReduce) { // S2不是最后结束的数据就需要往后做规约，放入第二块ws + wsInfoOffset += paramNum_; + wsOffset += 2 * BASE_TOPK; + } + SetWaitFlag(HardEvent::S_MTE3); + LIServiceVec::CopyOut(vec1ParamGm[wsInfoOffset], tmpiBuff, 16); + SetWaitFlag(HardEvent::V_MTE3); + LIServiceVec::CopyOut(vec1ResGm[wsOffset], globalTopkUb_[innerS1Idx * BASE_TOPK * 2], 2 * BASE_TOPK); + SetWaitFlag(HardEvent::MTE3_V); + } + } else if (cuRealAcSeq <= 0) { + CleanInvalidOutput(info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount); + } + } + + // BNSD场景无效S1 输出-1 + if (LAYOUT_T == LI_LAYOUT::BSND) { + // 最后一个S1的基本块, 需要 >= info.actS1Size + bool isS1LoopEnd = (cuBaseS1Idx + s1BaseSize_) >= info.actS1Size; + int32_t invalidS1Num = constInfo_.qSeqSize - info.actS1Size; + // blockS2StartIdx_ == 0 控制S2从开始的核去做冗余清理 + if (invalidS1Num > 0 && isS1LoopEnd && blockS2StartIdx_ == 0) { + int32_t s1NumPerAiv = blockId_ % 2 == 0 ? CeilDiv(invalidS1Num, 2) : (invalidS1Num / 2); + int32_t s1OffsetPerAiv = info.actS1Size + (blockId_ % 2) * CeilDiv(invalidS1Num, 2); + for (int innerS1Idx = 0; innerS1Idx < s1NumPerAiv; innerS1Idx++) { + CleanInvalidOutput(info.indiceOutOffset + (s1OffsetPerAiv + innerS1Idx) * constInfo_.sparseCount); + } + } + + int32_t invalidS1Num2 = info.actS1Size - info.actS2Size; + if (invalidS1Num2 > 0 && isS1LoopEnd && blockS2StartIdx_ == 0 && constInfo_.attenMaskFlag) { + int32_t s1NumPerAiv = blockId_ % 2 == 0 ? CeilDiv(invalidS1Num2, 2) : (invalidS1Num2 / 2); + int32_t s1OffsetPerAiv = (blockId_ % 2) * CeilDiv(invalidS1Num2, 2); + for (int innerS1Idx = 0; innerS1Idx < s1NumPerAiv; innerS1Idx++) { + CleanInvalidOutput((info.bN2Idx * constInfo_.qSeqSize + s1OffsetPerAiv + innerS1Idx) * + constInfo_.sparseCount); + } + } + } + + if (info.isLastS2InnerLoop) { + // S2最后一个Loop后, 下一个基本块初始从0开始 + blockS2StartIdx_ = 0; + } +} + +template +__aicore__ inline void LightningIndexerServiceVector::ProcessLD() +{ + int32_t curCubeId = blockId_ / 2; + int32_t tmpCubeId = curCubeId; + + int64_t s2ActSeq; + int64_t s2Start; + int64_t s2End; + int64_t isS2End; + int64_t bn2Idx; + int64_t s1Idx; + uint32_t acc_list_num = 0; + int64_t bIdx = 0; + int64_t needFd; + int64_t wsOffset; + int64_t wsInfoOffset = 0; + int64_t nextneedFd; + int64_t valueOffset = 0; + int64_t outOffset = 0; + + LocalTensor curValueIdxUb = ldToBeMrgBuf_.Get(); + LocalTensor tmpUb = ldTmpBuf_.Get(); + + // S2开头信息 + // 开始必然没有头规约，因此从尾规约开始处理，while循环读取下一个核的头规约 + // 存满4个list或者遇到S2结尾，则做merge，直到做完S2 + // 每个核都忽略自己的头规约，因为必然由前面的核做完 + uint32_t s1LdStartIdx = 0; + uint32_t s1ProcNum = 0; + uint64_t paramGmCoreOffset = tmpCubeId * s1BaseSize_ * 2 * paramNum_; + for (uint32_t innerS1Idx = 0; innerS1Idx < s1BaseSize_; innerS1Idx++) { + needFd = vec1ParamGm.GetValue(paramGmCoreOffset + innerS1Idx * 2 * paramNum_ + paramNum_); + if (needFd == 1) { + s1LdStartIdx = (s1ProcNum == 0) ? innerS1Idx : s1LdStartIdx; + s1ProcNum++; + } + } + + if (s1ProcNum == 0) { + return; + } + + // S1逐行计算 + uint32_t s1VecNum = CeilDiv(s1ProcNum, 2); + if (blockId_ % 2 == 1) { + s1LdStartIdx = s1LdStartIdx + s1VecNum; + s1VecNum = s1ProcNum - s1VecNum; + } + for (uint32_t innerS1Idx = s1LdStartIdx; innerS1Idx < s1LdStartIdx + s1VecNum; innerS1Idx++) { + // 重置偏移 + tmpCubeId = curCubeId; + acc_list_num = 0; + valueOffset = 0; + + // 搬入数据 + wsOffset = tmpCubeId * s1BaseSize_ * 2 * 2 * BASE_TOPK + // 2个AIV共同地址偏移 + innerS1Idx * 2 * 2 * BASE_TOPK + 2 * BASE_TOPK; + SetWaitFlag(HardEvent::V_MTE2); + SetWaitFlag(HardEvent::S_MTE2); + DataCopyPad(curValueIdxUb, vec1ResGm[wsOffset], + {1, static_cast(2 * BASE_TOPK * sizeof(int32_t)), 0, 0}, {true, 0, 0, 0}); + acc_list_num++; + valueOffset += 2 * BASE_TOPK; + + // 获取下一个核规约信息 + tmpCubeId++; + wsInfoOffset = tmpCubeId * s1BaseSize_ * 2 * paramNum_ + innerS1Idx * 2 * paramNum_; + needFd = vec1ParamGm.GetValue(wsInfoOffset); + isS2End = vec1ParamGm.GetValue(wsInfoOffset + 4); + s1Idx = vec1ParamGm.GetValue(wsInfoOffset + 6); + outOffset = vec1ParamGm.GetValue(wsInfoOffset + 8); + + while (needFd == 1) { + // 搬入头规约数据 + wsOffset = tmpCubeId * s1BaseSize_ * 2 * 2 * BASE_TOPK + // 2个AIV共同地址偏移 + innerS1Idx * 2 * 2 * BASE_TOPK; + SetWaitFlag(HardEvent::V_MTE2); + SetWaitFlag(HardEvent::S_MTE2); + DataCopyPad(curValueIdxUb[valueOffset], vec1ResGm[wsOffset], + {1, static_cast(2 * BASE_TOPK * sizeof(int32_t)), 0, 0}, {true, 0, 0, 0}); + valueOffset += 2 * BASE_TOPK; + acc_list_num++; + + // 每满4个list，聚合 前2K为mrg结果 + if (acc_list_num == mrgListNum_) { + // MrgSort 四条2048的队列，Mrg成一条 + AscendC::MrgSort4Info params; + params.elementLengths[0] = BASE_TOPK; + params.elementLengths[1] = BASE_TOPK; + params.elementLengths[2] = BASE_TOPK; + params.elementLengths[3] = BASE_TOPK; + params.ifExhaustedSuspension = true; + params.validBit = 0b1111; + params.repeatTimes = 1; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = curValueIdxUb[0]; + srcList.src2 = curValueIdxUb[2 * BASE_TOPK]; + srcList.src3 = curValueIdxUb[4 * BASE_TOPK]; + srcList.src4 = curValueIdxUb[6 * BASE_TOPK]; + SetWaitFlag(HardEvent::MTE2_V); + MrgSort(tmpUb, srcList, params); + PipeBarrier(); + DataCopy(curValueIdxUb, tmpUb, 2 * BASE_TOPK); + PipeBarrier(); + acc_list_num = 1; + valueOffset = 2 * BASE_TOPK; + } + + // reduce到S2末尾，则跳出 + if (isS2End == 1) { + break; + } + + tmpCubeId++; + wsInfoOffset = tmpCubeId * s1BaseSize_ * 2 * paramNum_ + innerS1Idx * 2 * paramNum_; + needFd = vec1ParamGm.GetValue(wsInfoOffset); + isS2End = vec1ParamGm.GetValue(wsInfoOffset + 4); + } + + // mrg不足4个list的数据 + if (acc_list_num != 1) { + AscendC::MrgSort4Info params; + params.elementLengths[0] = BASE_TOPK; + params.elementLengths[1] = BASE_TOPK; + params.elementLengths[2] = BASE_TOPK; + params.elementLengths[3] = BASE_TOPK; + params.ifExhaustedSuspension = true; + if (acc_list_num == 2) { + params.validBit = 0b0011; + } else if (acc_list_num == 3) { + params.validBit = 0b0111; + } + params.repeatTimes = 1; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = curValueIdxUb[0]; + srcList.src2 = curValueIdxUb[2 * BASE_TOPK]; + srcList.src3 = curValueIdxUb[4 * BASE_TOPK]; + srcList.src4 = curValueIdxUb[6 * BASE_TOPK]; + SetWaitFlag(HardEvent::MTE2_V); + MrgSort(tmpUb, srcList, params); + PipeBarrier(); + DataCopy(curValueIdxUb, tmpUb, 2 * BASE_TOPK); + PipeBarrier(); + } + + // 搬出 + LocalTensor outValueUb = ldOutValueBuf_.Get(); + LocalTensor outIdxUb = ldOutIdxBuf_.Get(); + if (!constInfo_.returnValue) { + Extract(outValueUb, outIdxUb, curValueIdxUb, (BASE_TOPK / 32)); + LocalTensor idxULocal1 = outIdxUb.template ReinterpretCast(); + SetWaitFlag(HardEvent::V_MTE3); + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(indiceOutGm[outOffset], idxULocal1, + {1, static_cast(constInfo_.sparseCount * sizeof(int32_t)), 0, 0}); + SetWaitFlag(HardEvent::MTE3_V); + } else { + Extract(outValueUb, outIdxUb, curValueIdxUb, (BASE_TOPK / 32)); + PipeBarrier(); + LocalTensor idxULocal1 = outIdxUb.template ReinterpretCast(); + LocalTensor valueULocal1 = outValueUb.template ReinterpretCast(); + Cast(valueULocal1, outValueUb, RoundMode::CAST_ROUND, constInfo_.sparseCount); + PipeBarrier(); + SetWaitFlag(HardEvent::V_MTE3); + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(indiceOutGm[outOffset], idxULocal1, + {1, static_cast(constInfo_.sparseCount * sizeof(int32_t)), 0, 0}); + DataCopyPad(valueOutGm[outOffset], valueULocal1, + {1, static_cast(constInfo_.sparseCount * sizeof(K_T)), 0, 0}); + SetWaitFlag(HardEvent::MTE3_V); + } + } +} +} // namespace LIKernel +#endif diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/arch22/lightning_indexer_vector.h b/csrc/ascend/attention/lightning_indexer/op_kernel/arch22/lightning_indexer_vector.h new file mode 100644 index 000000000..84d534c42 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/arch22/lightning_indexer_vector.h @@ -0,0 +1,424 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. +  */ + +/*! + * \file lightning_indexer_vector.h + * \brief + */ +#ifndef LIGHTNING_INDEXER_VECTOR_H +#define LIGHTNING_INDEXER_VECTOR_H + +#include "lightning_indexer_vector.h" +#include "kernel_operator.h" + +namespace LIServiceVec { +using namespace AscendC; + +constexpr int32_t NEG_INF = 0xFF800000; +constexpr int32_t INVALID_INDEX = -1; +constexpr uint8_t VEC_REPEAT_MAX = 255; +constexpr uint8_t B32_VEC_ELM_NUM = 64; +constexpr uint8_t B32_BLOCK_ALIGN_NUM = 8; +constexpr uint8_t B32_VEC_REPEAT_STRIDE = 8; +constexpr uint64_t VEC_REPEAT_BYTES = 256; +constexpr int32_t CONST_TWO = 2; +constexpr int64_t VALUE_AND_INDEX_NUM = 2; +constexpr int64_t BLOCK_BYTES = 32; +constexpr int64_t MRG_QUE_0 = 0; +constexpr int64_t MRG_QUE_1 = 1; +constexpr int64_t MRG_QUE_2 = 2; +constexpr int64_t MRG_QUE_3 = 3; +constexpr int64_t MRG_BLOCK_2 = 2; +constexpr int64_t MRG_BLOCK_3 = 3; +constexpr int64_t MRG_BLOCK_4 = 4; + +template +__aicore__ inline void CopyIn(LocalTensor &mmOutUb, LocalTensor &weightsUb, GlobalTensor &mMoutGm, + GlobalTensor &weightScaleGm, int64_t MMout_gmoffset, int64_t weights_gmoffset, + int64_t groupInner, int64_t s2Inner, int64_t mmUbStride) +{ + // 将MMout_gmoffset copy到UB上 + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams dataCopymMoutParams; + dataCopymMoutParams.blockCount = groupInner; + dataCopymMoutParams.blockLen = s2Inner * sizeof(float); + dataCopymMoutParams.srcStride = 0; + dataCopymMoutParams.dstStride = mmUbStride; + dataCopymMoutParams.rsv = 0; + AscendC::DataCopyPad(mmOutUb, mMoutGm[MMout_gmoffset], dataCopymMoutParams, padParams); + + // 将weights_gmoffset copy到UB + AscendC::DataCopyPadExtParams padTParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams dataCopyweightParams; + dataCopyweightParams.blockCount = 1; + dataCopyweightParams.blockLen = groupInner * sizeof(T); + dataCopyweightParams.srcStride = 0; + dataCopyweightParams.dstStride = 0; + dataCopyweightParams.rsv = 0; + AscendC::DataCopyPad(weightsUb, weightScaleGm[weights_gmoffset], dataCopyweightParams, padTParams); +} + + +template +__aicore__ inline void CopyOut(const GlobalTensor &dstGm, const LocalTensor &srcUb, int64_t copyCount) +{ + AscendC::DataCopyParams dataCopyOutyParams; + dataCopyOutyParams.blockCount = 1; + dataCopyOutyParams.blockLen = copyCount * sizeof(T); + dataCopyOutyParams.srcStride = 0; + dataCopyOutyParams.dstStride = 0; + AscendC::DataCopyPad(dstGm, srcUb, dataCopyOutyParams); +} + + +template +__aicore__ inline void DoScale(const LocalTensor &reduceCacheBuf, LocalTensor &mmOutUb, + LocalTensor &weightsUb, LocalTensor &weightsTUb, LocalTensor &tmpBuff, + int64_t groupInner, int64_t s2Inner, int32_t outerGidx) +{ + // cast bfloat16_t to float + if constexpr (!IsSameType::value) { + AscendC::Cast(weightsUb, weightsTUb, RoundMode::CAST_NONE, groupInner); + AscendC::PipeBarrier(); + } + + // weight broadcast: [groupInner, 1] -> [groupInner, 8] + AscendC::Brcb(tmpBuff, weightsUb, LICommon::CeilDiv(groupInner, static_cast(B32_BLOCK_ALIGN_NUM)), + {1, B32_VEC_REPEAT_STRIDE}); + AscendC::PipeBarrier(); + + // do scale: [groupInner, 8] * [groupInner, s2Inner] + uint64_t countPerRepeat = VEC_REPEAT_BYTES / sizeof(float); + uint64_t repeatTimes = s2Inner / countPerRepeat; + for (int32_t i = 0; i < groupInner; i++) { + if (outerGidx == 0) { + AscendC::Mul(reduceCacheBuf[i * s2Inner], mmOutUb[i * s2Inner], tmpBuff[i * B32_BLOCK_ALIGN_NUM], + countPerRepeat, repeatTimes, {1, 1, 0, B32_VEC_REPEAT_STRIDE, B32_VEC_REPEAT_STRIDE, 0}); + } else { + AscendC::Mul(mmOutUb[i * s2Inner], mmOutUb[i * s2Inner], tmpBuff[i * B32_BLOCK_ALIGN_NUM], countPerRepeat, + repeatTimes, {1, 1, 0, B32_VEC_REPEAT_STRIDE, B32_VEC_REPEAT_STRIDE, 0}); + } + } + + if (outerGidx != 0) { + AscendC::PipeBarrier(); + AscendC::Add(reduceCacheBuf, mmOutUb, reduceCacheBuf, groupInner * s2Inner); + } + AscendC::PipeBarrier(); +} + + +__aicore__ inline uint64_t FindNearestPower2(uint64_t value) +{ + if (value <= CONST_TWO) { + return value; + } else { + const uint64_t pow = 63 - AscendC::ScalarCountLeadingZero(value); // 返回前导0的个数，对于64位整数，最大有效位位置 = 63 - 前导0个数 + return (1 << pow); + } +} + + +// dstTensor 需要初始化0 +__aicore__ inline void DoReduce(const LocalTensor &srcTensor, LocalTensor &dstTensor, int32_t rNum, + int32_t aNum) +{ + if (rNum == 1) { + AscendC::Adds(dstTensor, srcTensor, 0, aNum); + AscendC::PipeBarrier(); + return; + } + + uint32_t dichotomizeAddPow = FindNearestPower2(rNum); + uint32_t dichotomizeAddDiffSize = rNum - dichotomizeAddPow; + if (dichotomizeAddDiffSize != 0) { + AscendC::Add(srcTensor, srcTensor, srcTensor[dichotomizeAddPow * aNum], dichotomizeAddDiffSize * aNum); + AscendC::PipeBarrier(); + } + int32_t nowRows = dichotomizeAddPow; + while (nowRows > CONST_TWO) { + nowRows = nowRows / CONST_TWO; + AscendC::Add(srcTensor, srcTensor, srcTensor[nowRows * aNum], nowRows * aNum); + AscendC::PipeBarrier(); + } + AscendC::Add(dstTensor, srcTensor, srcTensor[aNum], aNum); + AscendC::PipeBarrier(); +} + + +/** + src: 传入的初始化空间 + eleNum: 需要初始化的元素个数需为64整数倍，元素将被初始化为交错排布的-inf，-1 + */ +__aicore__ inline void InitSortOutBuf(const LocalTensor &src, int64_t eleNum) +{ + uint64_t mask1[2] = {0x5555555555555555, 0}; + uint64_t mask0[2] = {0xaaaaaaaaaaaaaaaa, 0}; + int64_t repeatNum = eleNum / B32_VEC_ELM_NUM; + int64_t forLoop = repeatNum / VEC_REPEAT_MAX; + int64_t forRemain = repeatNum % VEC_REPEAT_MAX; + for (int i = 0; i < forLoop; i++) { + AscendC::Duplicate(src.template ReinterpretCast(), NEG_INF, mask1, VEC_REPEAT_MAX, 1, + B32_VEC_REPEAT_STRIDE); + AscendC::PipeBarrier(); + AscendC::Duplicate(src.template ReinterpretCast(), INVALID_INDEX, mask0, VEC_REPEAT_MAX, 1, + B32_VEC_REPEAT_STRIDE); + } + if (forRemain > 0) { + AscendC::Duplicate(src.template ReinterpretCast()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM], NEG_INF, + mask1, forRemain, 1, B32_VEC_REPEAT_STRIDE); + AscendC::PipeBarrier(); + AscendC::Duplicate(src.template ReinterpretCast()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM], + INVALID_INDEX, mask0, forRemain, 1, B32_VEC_REPEAT_STRIDE); + } + AscendC::PipeBarrier(); +} + + +/** + src: logits和索引，前logitsNum为logits，后logitsNum为索引 + tmp: 计算使用到的临时空间，大小与src一致 + logitsNum: 排序的元素个数, 暂只支持[128,256,384,512,1024,2048] + */ +__aicore__ inline void SortAll(LocalTensor &src, LocalTensor &tmp, int64_t logitsNum) +{ + int64_t sort32Repeats = logitsNum / BLOCK_BYTES; + AscendC::Sort32(tmp, src, src[logitsNum].ReinterpretCast(), sort32Repeats); + AscendC::PipeBarrier(); + + int64_t mrgGroups = sort32Repeats; + int64_t mrgElements = BLOCK_BYTES; + int64_t i = 0; + AscendC::LocalTensor srcTensor; + AscendC::LocalTensor dstTensor; + while (true) { + if (i % CONST_TWO == 0) { + srcTensor = tmp; + dstTensor = src; + } else { + srcTensor = src; + dstTensor = tmp; + } + AscendC::MrgSort4Info params; + params.elementLengths[0] = mrgElements; + params.elementLengths[MRG_QUE_1] = mrgElements; + params.elementLengths[MRG_QUE_2] = mrgElements; + params.elementLengths[MRG_QUE_3] = mrgElements; + params.ifExhaustedSuspension = false; + params.validBit = 0b1111; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = srcTensor[0]; + srcList.src2 = srcTensor[MRG_QUE_1 * VALUE_AND_INDEX_NUM * mrgElements]; + srcList.src3 = srcTensor[MRG_QUE_2 * VALUE_AND_INDEX_NUM * mrgElements]; + srcList.src4 = srcTensor[MRG_QUE_3 * VALUE_AND_INDEX_NUM * mrgElements]; + if (mrgGroups <= MRG_BLOCK_4) { + params.repeatTimes = 1; + if (mrgGroups == 1) { + break; + } else if (mrgGroups == MRG_BLOCK_2) { + params.validBit = 0b0011; + } else if (mrgGroups == MRG_BLOCK_3) { + params.validBit = 0b0111; + } else if (mrgGroups == MRG_BLOCK_4) { + params.validBit = 0b1111; + } + AscendC::MrgSort(dstTensor, srcList, params); + i += 1; + AscendC::PipeBarrier(); + break; + } else { + params.repeatTimes = mrgGroups / MRG_BLOCK_4; + AscendC::MrgSort(dstTensor, srcList, params); + i += 1; + mrgElements = mrgElements * MRG_BLOCK_4; + mrgGroups = mrgGroups / MRG_BLOCK_4; + } + AscendC::PipeBarrier(); + } + if (i % CONST_TWO == 0) { + AscendC::DataCopy(src, tmp, logitsNum * VALUE_AND_INDEX_NUM); + AscendC::PipeBarrier(); + } +} + + +/** + dst: 输出全排序的结果，排布方式为value，index + srcValue：输入的待排序浮点数 + srcIndex：浮点数的索引 + tmp: 计算使用到的临时空间，大小为srcValue+srcIndex + logitsNum: 排序的元素个数 + */ +__aicore__ inline void SortAll(LocalTensor &dst, LocalTensor &srcValue, LocalTensor &srcIndex, + LocalTensor &tmpTensor, int64_t logitsNum) +{ + int64_t sort32Repeats = logitsNum / BLOCK_BYTES; + AscendC::Sort(dst, srcValue, srcIndex, tmpTensor, sort32Repeats); + AscendC::PipeBarrier(); +} + + +/** + mrgDst: 合并进的Tensor + mrgSrc: 待合并的Tensor + tmpTensor：空间为mrgDst+mrgSrc + */ +__aicore__ inline void MergeSort(const LocalTensor &mrgDst, int32_t mrgDstNum, LocalTensor &mrgSrc, + int32_t mrgSrcNum, LocalTensor &tmpTensor) +{ + if (mrgDstNum <= 3072) { // 3072: threshold of data size for different processing strategy + AscendC::MrgSort4Info params; + params.elementLengths[MRG_QUE_0] = mrgSrcNum; + params.elementLengths[MRG_QUE_1] = mrgDstNum; + params.ifExhaustedSuspension = false; + params.validBit = 0b0011; + params.repeatTimes = 1; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = mrgSrc; + srcList.src2 = mrgDst; + + AscendC::MrgSort(tmpTensor, srcList, params); + AscendC::PipeBarrier(); + AscendC::DataCopy(mrgDst, tmpTensor, mrgDstNum * VALUE_AND_INDEX_NUM); + AscendC::PipeBarrier(); + } else { + int64_t unitElements = 1024; + int64_t segNum = mrgDstNum / unitElements; + int64_t mrgQuelen_1 = (segNum + 2) / 3; + int64_t mrgQuelen_2 = ((segNum - mrgQuelen_1) + 1) / 2; + int64_t mrgQuelen_3 = segNum - mrgQuelen_1 - mrgQuelen_2; + + AscendC::MrgSort4Info params; + params.elementLengths[MRG_QUE_0] = mrgQuelen_1 * unitElements; + params.elementLengths[MRG_QUE_1] = mrgQuelen_2 * unitElements; + params.elementLengths[MRG_QUE_2] = mrgQuelen_3 * unitElements; + params.elementLengths[MRG_QUE_3] = mrgSrcNum; + + params.ifExhaustedSuspension = false; + params.validBit = 0b1111; + params.repeatTimes = 1; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = mrgDst[0]; + srcList.src2 = mrgDst[mrgQuelen_1 * VALUE_AND_INDEX_NUM * unitElements]; + srcList.src3 = mrgDst[(mrgQuelen_1 + mrgQuelen_2) * VALUE_AND_INDEX_NUM * unitElements]; + srcList.src4 = mrgSrc; + + AscendC::MrgSort(tmpTensor, srcList, params); + AscendC::PipeBarrier(); + AscendC::DataCopy(mrgDst, tmpTensor, mrgDstNum * VALUE_AND_INDEX_NUM); + AscendC::PipeBarrier(); + } +} + + +/** + * @brief 合并基础块函数 + * @param dst 归并后的输出, 大小为blockNum * basicBlockSize * 2 * sizeof(float) + * @param src 基本块输入 + * @param blockNum 基本块的数量 + * @param basicBlockSize 基础块的大小 + * @return 无 + */ +__aicore__ inline void MrgBasicBlock(const LocalTensor &dst, const LocalTensor &src, int64_t blockNum, + int64_t basicBlockSize) +{ + // 初始化合并排序参数 + AscendC::MrgSort4Info params; + params.elementLengths[MRG_QUE_0] = basicBlockSize; + params.elementLengths[MRG_QUE_1] = basicBlockSize; + params.elementLengths[MRG_QUE_2] = basicBlockSize; + params.elementLengths[MRG_QUE_3] = basicBlockSize; + params.ifExhaustedSuspension = false; + // 根据块的数量设置有效位 + if (blockNum == MRG_BLOCK_2) { + params.validBit = 0b0011; + } else if (blockNum == MRG_BLOCK_3) { + params.validBit = 0b0111; + } else if (blockNum == MRG_BLOCK_4) { + params.validBit = 0b1111; + } else { + AscendC::DataCopy(dst, src, basicBlockSize * VALUE_AND_INDEX_NUM); + return; + } + // 初始化源列表 + AscendC::MrgSortSrcList srcList; + srcList.src1 = src[0]; + srcList.src2 = src[basicBlockSize * VALUE_AND_INDEX_NUM * MRG_QUE_1]; + srcList.src3 = src[basicBlockSize * VALUE_AND_INDEX_NUM * MRG_QUE_2]; + srcList.src4 = src[basicBlockSize * VALUE_AND_INDEX_NUM * MRG_QUE_3]; + // 执行合并排序 + AscendC::MrgSort(dst, srcList, params); +} + + +/** + * @brief 从两个队列中选择topk + * @param dst 已经归并好的topk数据 + * @param needsMerging 需要合并的有序数据 + * @param tmp 临时空间 + * @param topk topk的元素个数 + * @param mergSize 待合并的元素个数 + * @return 无 + */ +template +__aicore__ inline void SparseTopK(const LocalTensor &dst, const LocalTensor &needsMerging, + const LocalTensor &tmp, int64_t topk, int64_t mergSize) +{ + // 如果不需要合并，则直接复制数据 + if (!needMrg) { + AscendC::DataCopy(dst, needsMerging, mergSize * VALUE_AND_INDEX_NUM); + return; + } + // 初始化合并排序参数 + AscendC::MrgSort4Info params; + params.elementLengths[0] = topk; + params.elementLengths[1] = mergSize; + params.ifExhaustedSuspension = (topk == mergSize); + params.validBit = 0b0011; + // 初始化源列表 + AscendC::MrgSortSrcList srcList; + srcList.src1 = dst; + srcList.src2 = needsMerging; + // 执行合并排序 + AscendC::MrgSort(tmp, srcList, params); + AscendC::PipeBarrier(); + // 将结果复制到目标张量 + AscendC::DataCopy(dst, tmp, topk * VALUE_AND_INDEX_NUM); +} + + +__aicore__ inline void ExtractIndex(const LocalTensor &idxULocal, const LocalTensor &sortLocal, + int64_t extractNum) +{ + AscendC::GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = Ceil(extractNum * sizeof(float) * VALUE_AND_INDEX_NUM, VEC_REPEAT_BYTES); + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = B32_VEC_REPEAT_STRIDE; + gatherMaskParams.src1RepeatStride = 0; + uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数 + uint8_t src1Pattern = 2; // 固定模式2,表示筛选出奇数索引的数 + AscendC::GatherMask(idxULocal, sortLocal, src1Pattern, false, static_cast(0), gatherMaskParams, rsvdCnt); + AscendC::PipeBarrier(); +} + + +template +__aicore__ inline void SetWaitFlag(HardEvent evt) +{ + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(evt)); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); +} + +} // namespace LIServiceVec +#endif // LIGHTNING_INDEXER_VECTOR_H diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/lightning_indexer_kernel.h b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/lightning_indexer_kernel.h new file mode 100644 index 000000000..91f07a7d4 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/lightning_indexer_kernel.h @@ -0,0 +1,676 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. +  */ + +/*! + * \file lightning_indexer_kernel.h + * \brief + */ + +#ifndef LIGHTNING_INDEXER_KERNEL_H +#define LIGHTNING_INDEXER_KERNEL_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../lightning_indexer_common.h" +#include "lightning_indexer_service_vector.h" +#include "lightning_indexer_service_cube.h" + +namespace LIKernel { +using namespace LICommon; +using namespace matmul; +using AscendC::CacheMode; +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +// 由于S2循环前，RunInfo还没有赋值，使用TempLoopInfo临时存放B、N、S1轴相关的信息；同时减少重复计算 +struct TempLoopInfo { + uint32_t bN2Idx = 0; + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint32_t gS1Idx = 0U; + uint32_t gS1LoopEnd = 0U; // gS1方向循环的结束Idx + uint32_t s2LoopEnd = 0U; // S2方向循环的结束Idx + uint32_t actS1Size = 1U; // 当前Batch循环处理的S1轴的实际大小 + uint32_t actS2Size = 0U; + uint32_t actS2SizeOrig = 0U; // 压缩前s2 + bool curActSeqLenIsZero = false; + bool needDealActS1LessThanS1 = false; // S1的实际长度小于shape的S1长度时，是否需要清理输出 + uint32_t actMBaseSize = 0U; // m轴(gS1)方向实际大小 + uint32_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小 + uint32_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小 +}; + +template +class LightningIndexerKernel { +public: + __aicore__ inline LightningIndexerKernel(){}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK, + __gm__ uint8_t *blockTable, __gm__ uint8_t *sparseIndices, __gm__ uint8_t *sparseValues, + __gm__ uint8_t *workspace, const LITilingData *__restrict tiling, TPipe *tPipe); + __aicore__ inline void Process(); + + // =================================类型定义区================================= + static constexpr bool DT_W_FLAG = LIT::weightsTypeFlag; + using Q_T = typename LIT::queryType; + using K_T = typename LIT::keyType; + using OUT_T = typename LIT::outputType; + static constexpr bool PAGE_ATTENTION = LIT::pageAttention; + static constexpr LI_LAYOUT LAYOUT_T = LIT::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = LIT::keyLayout; + using W_T = + typename LightningIndexerTypeTraits::type>::weightsType; + + LightningIndexerServiceCube matmulService; + LightningIndexerServiceVector vectorService; + + // =================================常量区================================= + static constexpr uint32_t SYNC_C1_V1_FLAG = 4; + static constexpr uint32_t SYNC_V1_C1_FLAG = 5; + + static constexpr uint32_t M_BASE_SIZE = 256; + static constexpr uint32_t S1_BASE_SIZE = 4; + static constexpr uint32_t S1_BASE_SIZE_SMALL = 2; + static constexpr uint32_t S2_BASE_SIZE = 128; + static constexpr uint32_t HEAD_DIM = 128; + static constexpr uint32_t K_HEAD_NUM = 1; + static constexpr uint32_t GM_ALIGN_BYTES = 512; + + static constexpr int64_t LD_PREFETCH_LEN = 2; + +protected: + TPipe *pipe = nullptr; + + // offset + uint64_t queryCoreOffset = 0ULL; + uint64_t keyCoreOffset = 0ULL; + uint64_t weightsCoreOffset = 0ULL; + uint64_t indiceOutCoreOffset = 0ULL; + uint64_t valueOutCoreOffset = 0ULL; + // ================================Global Buffer区================================= + GlobalTensor queryGm; + GlobalTensor keyGm; + GlobalTensor weightsGm; + + GlobalTensor indiceOutGm; + GlobalTensor valueOutGm; + GlobalTensor blockTableGm; + + GlobalTensor actualSeqLengthsGmQ; + GlobalTensor actualSeqLengthsGmKv; + + // ================================类成员变量==================================== + // aic、aiv核信息 + uint32_t tmpBlockIdx = 0U; + uint32_t aiCoreIdx = 0U; + uint32_t usedCoreNum = 0U; + + LICommon::ConstInfo constInfo{}; + TempLoopInfo tempLoopInfo{}; + LICommon::SplitCoreInfo splitCoreInfo{}; + + // ================================Init functions================================== + __aicore__ inline void InitTilingData(const LITilingData *__restrict tilingData); + __aicore__ inline void InitBuffers(); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK); + // ================================Split Core================================ + __aicore__ inline void SplitCore(uint32_t curCoreIdx, uint32_t &coreNum, LICommon::SplitCoreInfo &info); + __aicore__ inline uint32_t GetTotalBaseBlockNum(); + __aicore__ inline uint32_t GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, uint32_t actS2SizeOrig); + // ================================Process functions================================ + __aicore__ inline void ProcessMain(); + __aicore__ inline void ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, + LICommon::RunInfo runInfo); + __aicore__ inline void ProcessInvalid(); + // ================================Params Calc===================================== + __aicore__ inline void CalcGS1LoopParams(uint32_t bN2Idx); + __aicore__ inline void GetBN2Idx(uint32_t bN2Idx); + __aicore__ inline uint32_t GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGmKv, uint32_t defaultSeqLen); + __aicore__ inline void GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, + uint32_t &actS2Size, uint32_t &actS2SizeOrig); + __aicore__ inline void CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx); + __aicore__ inline void CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, LICommon::RunInfo &runInfo); + __aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start); +}; + +template +__aicore__ inline void LightningIndexerKernel::InitTilingData(const LITilingData *__restrict tilingData) +{ + usedCoreNum = tilingData->usedCoreNum; + constInfo.batchSize = tilingData->bSize; + constInfo.qHeadNum = constInfo.gSize = tilingData->gSize; + constInfo.kSeqSize = tilingData->s2Size; + constInfo.qSeqSize = tilingData->s1Size; + constInfo.attenMaskFlag = (tilingData->sparseMode == 3); + constInfo.kCacheBlockSize = tilingData->blockSize; + constInfo.maxBlockNumPerBatch = tilingData->maxBlockNumPerBatch; + constInfo.sparseCount = tilingData->sparseCount; + constInfo.outputLayout = LAYOUT_T; // 输出和输入形状一致 + if constexpr (std::is_same_v) { + constInfo.INVALID_VAL = 0xFC00; + }else { + constInfo.INVALID_VAL = 0xFF80; + } + if constexpr (LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS1 = true; + } + if constexpr (K_LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS2 = true; + } + + constInfo.kHeadNum = K_HEAD_NUM; + constInfo.headDim = HEAD_DIM; + + if (constInfo.sparseCount > 2048) { + constInfo.mBaseSize = S1_BASE_SIZE_SMALL * constInfo.gSize; + constInfo.s1BaseSize = S1_BASE_SIZE_SMALL; + } else { + constInfo.mBaseSize = S1_BASE_SIZE * constInfo.gSize; + constInfo.s1BaseSize = S1_BASE_SIZE; + } + constInfo.s2BaseSize = S2_BASE_SIZE; + constInfo.returnValueFlag = tilingData->returnValue; + constInfo.splitMFlag = (constInfo.gSize == 64 && constInfo.sparseCount <= 2048); +} + +template +__aicore__ inline void LightningIndexerKernel::InitBuffers() +{ + if ASCEND_IS_AIV { + vectorService.InitBuffers(pipe); + } else { + matmulService.InitBuffers(pipe); + } +} + +template +__aicore__ inline void LightningIndexerKernel::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengthsK) +{ + if (actualSeqLengthsQ == nullptr) { + constInfo.actualLenQDims = 0; + } else { + constInfo.actualLenQDims = constInfo.batchSize; + actualSeqLengthsGmQ.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsQ, constInfo.actualLenQDims); + } + if (actualSeqLengthsK == nullptr) { + constInfo.actualLenDims = 0; + } else { + constInfo.actualLenDims = constInfo.batchSize; + actualSeqLengthsGmKv.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsK, constInfo.actualLenDims); + } +} + +template +__aicore__ inline uint32_t LightningIndexerKernel::GetActualSeqLen(uint32_t bIdx, + uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGmKv, + uint32_t defaultSeqLen) +{ + if (actualLenDims == 0) { + return defaultSeqLen; + } else if (isAccumSeq && bIdx > 0) { + return actualSeqLengthsGmKv.GetValue(bIdx) - actualSeqLengthsGmKv.GetValue(bIdx - 1); + } else { + return actualSeqLengthsGmKv.GetValue(bIdx); + } +} + +template +__aicore__ inline void LightningIndexerKernel::GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, + uint32_t &actS2Size, uint32_t &actS2SizeOrig) +{ + actS1Size = GetActualSeqLen(bIdx, constInfo.actualLenQDims, constInfo.isAccumSeqS1, actualSeqLengthsGmQ, + constInfo.qSeqSize); + actS2SizeOrig = + GetActualSeqLen(bIdx, constInfo.actualLenDims, + constInfo.isAccumSeqS2, actualSeqLengthsGmKv, constInfo.kSeqSize); + actS2Size = actS2SizeOrig; +} + +template +__aicore__ inline uint32_t LightningIndexerKernel::GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, + uint32_t actS2SizeOrig) +{ + if (actS2SizeOrig == 0) { + return 0; + } + uint32_t s1Offset = constInfo.s1BaseSize * s1gIdx; + int32_t validS2LenBase = static_cast(actS2SizeOrig) - static_cast(actS1Size); + int32_t validS2Len = (static_cast(s1Offset) + validS2LenBase + static_cast(constInfo.s1BaseSize)); + validS2Len = Min(validS2Len, static_cast(actS2SizeOrig)); + validS2Len = Max(validS2Len, 1); + return (validS2Len + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; +} + +template +__aicore__ inline uint32_t LightningIndexerKernel::GetTotalBaseBlockNum() +{ + uint32_t totalBlockNum = 0; + uint32_t actS1Size, actS2Size, actS2SizeOrig; + uint32_t s1GBaseNum, s2BaseNum; + for (uint32_t bIdx = 0; bIdx < constInfo.batchSize; bIdx++) { + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size, actS2SizeOrig); + s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize); + if (!constInfo.attenMaskFlag) { + s2BaseNum = constInfo.isLDOpen ? CeilDiv(actS2Size, constInfo.s2BaseSize) : (actS2Size > 0 ? 1 : 0); + totalBlockNum += s1GBaseNum * s2BaseNum * constInfo.kHeadNum; + continue; + } + for (uint32_t s1gIdx = 0; s1gIdx < s1GBaseNum; s1gIdx++) { + s2BaseNum = constInfo.isLDOpen + ? GetS2BaseBlockNumOnMask(s1gIdx, actS1Size, actS2SizeOrig) + : (actS2Size > 0 ? 1 : 0); + totalBlockNum += s2BaseNum * constInfo.kHeadNum; + } + } + return totalBlockNum; +} + + +// 多核版本，双闭区间。基本原则：计算每个核最少处理的块数, 剩余的部分前面的核每个核多处理一块 +template +__aicore__ void inline LightningIndexerKernel::SplitCore(uint32_t curCoreIdx, uint32_t &coreNum, + LICommon::SplitCoreInfo &info) +{ + uint32_t totalBlockNum = GetTotalBaseBlockNum(); + uint32_t minBlockPerCore = totalBlockNum / coreNum; + uint32_t deal1MoreBlockCoreNum = totalBlockNum % coreNum; + uint32_t coreIdx = 0; + uint32_t lastGS1RemainBlockCnt = 0; + uint32_t coreDealBlockCnt = coreIdx < deal1MoreBlockCoreNum ? minBlockPerCore + 1 : minBlockPerCore; + coreNum = minBlockPerCore == 0 ? deal1MoreBlockCoreNum : coreNum; + if (curCoreIdx < coreNum) { + splitCoreInfo.isCoreEnable = true; + } else { + splitCoreInfo.isCoreEnable = false; + return; + } + + bool findLastCoreEnd = true; + uint32_t actS1Size, actS2Size, actS2SizeOrig; + uint32_t s1GBaseNum, s2BaseNum, s2Loop; + for (uint32_t bN2Idx = 0; bN2Idx < constInfo.batchSize * constInfo.kHeadNum; bN2Idx++) { + uint32_t bIdx = bN2Idx / constInfo.kHeadNum; + if (bN2Idx % constInfo.kHeadNum == 0) { + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size, actS2SizeOrig); + s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize); + s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize); + } + if constexpr (LAYOUT_T == LI_LAYOUT::BSND) { + if (findLastCoreEnd && (s1GBaseNum == 0U || s2BaseNum == 0U)) { + info.bN2Start = bN2Idx; + info.gS1Start = 0; + info.s2Start = 0; + findLastCoreEnd = false; + } + } + for (uint32_t gS1Idx = 0; gS1Idx < s1GBaseNum; gS1Idx++) { + if (constInfo.attenMaskFlag) { + s2BaseNum = GetS2BaseBlockNumOnMask(gS1Idx, actS1Size, actS2SizeOrig); + } + if (findLastCoreEnd && s2BaseNum == 0U) { + info.bN2Start = bN2Idx; + info.gS1Start = gS1Idx; + info.s2Start = 0; + findLastCoreEnd = false; + } + s2Loop = constInfo.isLDOpen ? s2BaseNum : (actS2Size > 0 ? 1 : 0); + for (uint32_t s2Idx = 0; s2Idx < s2Loop;) { + if (findLastCoreEnd) { + info.bN2Start = bN2Idx; + info.gS1Start = gS1Idx; + info.s2Start = s2Idx; + findLastCoreEnd = false; + } + uint32_t s2RemainBaseNum = s2Loop - s2Idx; + if (lastGS1RemainBlockCnt + s2RemainBaseNum >= coreDealBlockCnt) { + info.bN2End = bN2Idx; + info.gS1End = gS1Idx; + info.s2End = constInfo.isLDOpen + ? s2Idx + coreDealBlockCnt - lastGS1RemainBlockCnt - 1 + : s2BaseNum - 1; + + if (coreIdx == curCoreIdx) { + // S2被切N核，那么只有第一个核需要处理LD，其他核不用 + if (s2Idx == 0 && info.s2End + 1 < s2BaseNum) { + info.isLD = true; + } + // 最后一个核处理的不是最后一个Batch，表明后面的Batch为空块(S2=0), 调整终点坐标以便清理输出 + if (coreIdx == coreNum - 1 && info.bN2End != constInfo.batchSize - 1) { + info.bN2End = constInfo.batchSize - 1; + info.gS1End = 0; + info.s2End = 0; + } + return; + } + coreIdx++; + findLastCoreEnd = true; + s2Idx = info.s2End + 1; + lastGS1RemainBlockCnt = 0; + coreDealBlockCnt = coreIdx < deal1MoreBlockCoreNum ? minBlockPerCore + 1 : minBlockPerCore; + } else { + lastGS1RemainBlockCnt += s2RemainBaseNum; + break; + } + } + } + } +} + +template +__aicore__ inline void LightningIndexerKernel::DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start) +{ + if ASCEND_IS_AIV { + if (constInfo.outputLayout == LI_LAYOUT::TND) { + uint32_t tSize = actualSeqLengthsGmQ.GetValue(constInfo.batchSize - 1); + uint32_t tBase = bIdx == 0 ? 0 : actualSeqLengthsGmQ.GetValue(bIdx - 1); + uint32_t s1Count = tempLoopInfo.actS1Size; + + for (uint32_t s1Idx = s1Start; s1Idx < s1Count; s1Idx++) { + uint64_t indiceOutOffset = + (tBase + s1Idx) * constInfo.kHeadNum * constInfo.sparseCount + // T轴、s1轴偏移 + n2Idx * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } else if (constInfo.outputLayout == LI_LAYOUT::BSND) { + for (uint32_t s1Idx = s1Start; s1Idx < constInfo.qSeqSize; s1Idx++) { + // B,S1,N2,K + uint64_t indiceOutOffset = bIdx * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount + + s1Idx * constInfo.kHeadNum * constInfo.sparseCount + // B轴、S1轴偏移 + n2Idx * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } + } +} + +template +__aicore__ inline void LightningIndexerKernel::Init(__gm__ uint8_t *query, + __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK, + __gm__ uint8_t *blockTable, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *sparseValues, + __gm__ uint8_t *workspace, const LITilingData *__restrict tiling, + TPipe *tPipe) +{ + if ASCEND_IS_AIV { + tmpBlockIdx = GetBlockIdx(); // vec:0-47 + aiCoreIdx = tmpBlockIdx / 2; + } else { + tmpBlockIdx = GetBlockIdx(); // cube:0-23 + aiCoreIdx = tmpBlockIdx; + } + + InitTilingData(tiling); + InitActualSeqLen(actualSeqLengthsQ, actualSeqLengthsK); + + // 获取分核信息 + SplitCore(aiCoreIdx, usedCoreNum, splitCoreInfo); + + pipe = tPipe; + + uint64_t offset = 0; + // vec 把整个s2的score存储在GM，大小为s1BaseSize * 16K * 4 + GlobalTensor scoreGm; // 存放vec核写出的score + uint64_t singleCoreScoreSize = constInfo.s1BaseSize * + LICommon::Align( + (uint64_t)constInfo.kSeqSize, + (uint64_t)constInfo.s2BaseSize) * + sizeof(uint16_t); + scoreGm.SetGlobalBuffer((__gm__ uint16_t *)(workspace + aiCoreIdx * singleCoreScoreSize)); + offset += GetBlockNum() * singleCoreScoreSize; + + if ASCEND_IS_AIV { + vectorService.InitParams(constInfo, tiling); + indiceOutGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices); + valueOutGm.SetGlobalBuffer((__gm__ K_T *)sparseValues); + weightsGm.SetGlobalBuffer((__gm__ W_T *)weights); + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + vectorService.InitVecInputTensor(weightsGm, indiceOutGm, valueOutGm, blockTableGm); + vectorService.InitVecWorkspaceTensor(scoreGm); + } else { + matmulService.InitParams(constInfo); + queryGm.SetGlobalBuffer((__gm__ Q_T *)query); + if constexpr (PAGE_ATTENTION) { + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + } + keyGm.SetGlobalBuffer((__gm__ K_T *)key); + matmulService.InitMm1GlobalTensor(blockTableGm, keyGm, queryGm); + } + InitBuffers(); +} + +template +__aicore__ inline void LightningIndexerKernel::GetBN2Idx(uint32_t bN2Idx) +{ + tempLoopInfo.bN2Idx = bN2Idx; + tempLoopInfo.bIdx = bN2Idx / constInfo.kHeadNum; + tempLoopInfo.n2Idx = bN2Idx % constInfo.kHeadNum; +} + +template +__aicore__ inline void LightningIndexerKernel::CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx) +{ + tempLoopInfo.gS1Idx = gS1LoopIdx; + tempLoopInfo.actMBaseSize = constInfo.mBaseSize; + uint32_t remainedGS1Size = tempLoopInfo.actS1Size * constInfo.gSize - tempLoopInfo.gS1Idx * constInfo.mBaseSize; + if (remainedGS1Size <= constInfo.mBaseSize && remainedGS1Size > 0) { + tempLoopInfo.actMBaseSize = tempLoopInfo.mBasicSizeTail; + } + + bool isEnd = (bN2LoopIdx == splitCoreInfo.bN2End) && (gS1LoopIdx == splitCoreInfo.gS1End); + uint32_t s2BlockNum; + if (constInfo.attenMaskFlag) { + s2BlockNum = GetS2BaseBlockNumOnMask(gS1LoopIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2SizeOrig); + } else { + s2BlockNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + } + tempLoopInfo.s2LoopEnd = isEnd ? splitCoreInfo.s2End : s2BlockNum - 1; +} + +template +__aicore__ inline void LightningIndexerKernel::CalcGS1LoopParams(uint32_t bN2LoopIdx) +{ + GetBN2Idx(bN2LoopIdx); + GetS1S2ActualSeqLen(tempLoopInfo.bIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2Size, tempLoopInfo.actS2SizeOrig); + if ((tempLoopInfo.actS2Size == 0) || (tempLoopInfo.actS1Size == 0)) { + tempLoopInfo.curActSeqLenIsZero = true; + return; + } + tempLoopInfo.curActSeqLenIsZero = false; + tempLoopInfo.s2BasicSizeTail = tempLoopInfo.actS2Size % constInfo.s2BaseSize; + tempLoopInfo.s2BasicSizeTail = + (tempLoopInfo.s2BasicSizeTail == 0) ? constInfo.s2BaseSize : tempLoopInfo.s2BasicSizeTail; + tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize; + tempLoopInfo.mBasicSizeTail = + (tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail; + + uint32_t gS1SplitNum = (tempLoopInfo.actS1Size * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize; + tempLoopInfo.gS1LoopEnd = (bN2LoopIdx == splitCoreInfo.bN2End) ? splitCoreInfo.gS1End : gS1SplitNum - 1; + if constexpr (LAYOUT_T == LI_LAYOUT::BSND) { + if (tempLoopInfo.gS1LoopEnd == gS1SplitNum - 1 && constInfo.qSeqSize > tempLoopInfo.actS1Size) { + tempLoopInfo.needDealActS1LessThanS1 = true; + } + } +} + +template +__aicore__ inline void LightningIndexerKernel::CalcRunInfo(uint32_t loop, + uint32_t s2LoopIdx, LICommon::RunInfo &runInfo) +{ + runInfo.loop = loop; + runInfo.bIdx = tempLoopInfo.bIdx; + runInfo.gS1Idx = tempLoopInfo.gS1Idx; + runInfo.s2Idx = s2LoopIdx; + runInfo.bN2Idx = tempLoopInfo.bN2Idx; + runInfo.isValid = s2LoopIdx <= tempLoopInfo.s2LoopEnd; + + if (!runInfo.isValid) { + return; + } + + runInfo.actS1Size = tempLoopInfo.actS1Size; + runInfo.actS2Size = tempLoopInfo.actS2Size; + runInfo.actS2SizeOrig = tempLoopInfo.actS2SizeOrig; + // 计算实际基本块size + runInfo.actMBaseSize = tempLoopInfo.actMBaseSize; + runInfo.actualSingleProcessSInnerSize = constInfo.s2BaseSize; + uint32_t s2SplitNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + if (runInfo.s2Idx == s2SplitNum - 1) { + runInfo.actualSingleProcessSInnerSize = tempLoopInfo.s2BasicSizeTail; + } + runInfo.actualSingleProcessSInnerSizeAlign = + LICommon::Align((uint32_t)runInfo.actualSingleProcessSInnerSize, LICommon::ConstInfo::BUFFER_SIZE_BYTE_32B); + + runInfo.isFirstS2InnerLoop = s2LoopIdx == splitCoreInfo.s2Start; + runInfo.isLastS2InnerLoop = s2LoopIdx == tempLoopInfo.s2LoopEnd; + runInfo.isAllLoopEnd = (runInfo.bN2Idx == splitCoreInfo.bN2End) && (runInfo.gS1Idx == splitCoreInfo.gS1End) && + (runInfo.s2Idx == splitCoreInfo.s2End); + + if (runInfo.isFirstS2InnerLoop) { + uint64_t actualSeqQPrefixSum; + if constexpr (LAYOUT_T == LI_LAYOUT::TND) { + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGmQ.GetValue(runInfo.bIdx - 1); + } else { // BSND + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.qSeqSize; + } + uint64_t tndBIdxOffset = actualSeqQPrefixSum * constInfo.qHeadNum * constInfo.headDim; + // B,S1,N1(N2,G),D + queryCoreOffset = tndBIdxOffset + runInfo.gS1Idx * constInfo.mBaseSize * constInfo.headDim; + // B,S1,N1(N2,G)/T,N1(N2,G) + weightsCoreOffset = actualSeqQPrefixSum * constInfo.qHeadNum + runInfo.n2Idx * constInfo.gSize; + // B,S1,N2,k/T,N2,k + indiceOutCoreOffset = + actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + runInfo.n2Idx * constInfo.sparseCount; + // B,S1,N2,k/T,N2,k + valueOutCoreOffset = + actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + runInfo.n2Idx * constInfo.sparseCount; + } + uint64_t actualSeqKPrefixSum; + if constexpr (K_LAYOUT_T == LI_LAYOUT::TND) { // T N2 D + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGmKv.GetValue(runInfo.bIdx - 1); + } else { + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.kSeqSize; + } + uint64_t tndBIdxOffsetForK = actualSeqKPrefixSum * constInfo.kHeadNum * constInfo.headDim; + keyCoreOffset = tndBIdxOffsetForK + runInfo.s2Idx * constInfo.s2BaseSize * constInfo.kHeadNum * constInfo.headDim; + runInfo.tensorQueryOffset = queryCoreOffset; + runInfo.tensorKeyOffset = keyCoreOffset; + runInfo.tensorWeightsOffset = weightsCoreOffset; + runInfo.indiceOutOffset = indiceOutCoreOffset; + runInfo.valueOutOffset = valueOutCoreOffset; +} + +template +__aicore__ inline void LightningIndexerKernel::Process() +{ + if (usedCoreNum == 0) { + // 没有计算任务，直接清理输出 + ProcessInvalid(); + return; + } + + ProcessMain(); +} + +template +__aicore__ inline void LightningIndexerKernel::ProcessInvalid() +{ + if ASCEND_IS_AIV { + uint32_t aivCoreNum = GetBlockNum() * 2; // 2 means c:v = 1:2 + uint64_t totalOutputSize = + constInfo.batchSize * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount; + uint64_t singleCoreSize = + LICommon::Align((totalOutputSize + aivCoreNum - 1) / aivCoreNum, GM_ALIGN_BYTES / sizeof(OUT_T)); + uint64_t baseSize = tmpBlockIdx * singleCoreSize; + if (baseSize < totalOutputSize) { + uint64_t dealSize = + (baseSize + singleCoreSize <= totalOutputSize) ? singleCoreSize : totalOutputSize - baseSize; + GlobalTensor output = indiceOutGm[baseSize]; + AscendC::InitGlobalMemory(output, dealSize, constInfo.INVALID_IDX); + if (constInfo.returnValueFlag) { + GlobalTensor valueOutGmTmp; + valueOutGmTmp.SetGlobalBuffer((__gm__ uint16_t *)valueOutGm.GetPhyAddr()); + GlobalTensor valueOut = valueOutGmTmp[baseSize]; + AscendC::InitGlobalMemory(valueOut, dealSize, constInfo.INVALID_VAL); + } + } + } +} + +template +__aicore__ inline void LightningIndexerKernel::ProcessMain() +{ + if (!splitCoreInfo.isCoreEnable) { + return; + } + + if ASCEND_IS_AIV { + vectorService.AllocEventID(); + CrossCoreSetFlag(LICommon::ConstInfo::CROSS_VC_EVENT + 0); + CrossCoreSetFlag(LICommon::ConstInfo::CROSS_VC_EVENT + 1); + } else { + matmulService.AllocEventID(); + } + + LICommon::RunInfo runInfo; + uint32_t gloop = 0; + for (uint32_t bN2LoopIdx = splitCoreInfo.bN2Start; bN2LoopIdx <= splitCoreInfo.bN2End; bN2LoopIdx++) { + CalcGS1LoopParams(bN2LoopIdx); + if (tempLoopInfo.curActSeqLenIsZero) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, 0U); + continue; + } + for (uint32_t gS1LoopIdx = splitCoreInfo.gS1Start; gS1LoopIdx <= tempLoopInfo.gS1LoopEnd; gS1LoopIdx++) { + CalcS2LoopParams(bN2LoopIdx, gS1LoopIdx); + for (int s2LoopIdx = splitCoreInfo.s2Start; s2LoopIdx <= tempLoopInfo.s2LoopEnd; s2LoopIdx++) { + ProcessBaseBlock(gloop, s2LoopIdx, runInfo); + ++gloop; + } + splitCoreInfo.s2Start = 0; + } + if (tempLoopInfo.needDealActS1LessThanS1) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, tempLoopInfo.actS1Size); + } + splitCoreInfo.gS1Start = 0; + } + + if ASCEND_IS_AIV { + vectorService.FreeEventID(); + } else { + matmulService.FreeEventID(); + CrossCoreWaitFlag(LICommon::ConstInfo::CROSS_VC_EVENT + 0); + CrossCoreWaitFlag(LICommon::ConstInfo::CROSS_VC_EVENT + 1); + } +} + +template +__aicore__ inline void LightningIndexerKernel::ProcessBaseBlock(uint32_t loop, + uint64_t s2LoopIdx, LICommon::RunInfo runInfo) +{ + CalcRunInfo(loop, s2LoopIdx, runInfo); + if ASCEND_IS_AIC { + matmulService.ComputeMm1(runInfo); + } else { + vectorService.ProcessVec1(runInfo); + if (runInfo.isLastS2InnerLoop) { // 本核s2last + vectorService.ProcessTopK(runInfo); + } + } +} + +} // namespace LIKernel +#endif // LIGHTNING_INDEXER_KERNEL_H diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/lightning_indexer_service_cube.h b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/lightning_indexer_service_cube.h new file mode 100644 index 000000000..ba39973a6 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/lightning_indexer_service_cube.h @@ -0,0 +1,470 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. +  */ + +/*! + * \file lightning_indexer_service_cube.h + * \brief use 5 buffer for matmul l1, better pipeline + */ +#ifndef LIGHTNING_INDEXER_SERVICE_CUBE_H +#define LIGHTNING_INDEXER_SERVICE_CUBE_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../lightning_indexer_common.h" + +namespace LIKernel { +using namespace LICommon; +template +class LightningIndexerServiceCube { +public: + using Q_T = typename LIT::queryType; + using K_T = typename LIT::keyType; + + __aicore__ inline LightningIndexerServiceCube(){}; + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitMm1GlobalTensor(const GlobalTensor &blkTableGm, const GlobalTensor &keyGm, + const GlobalTensor &queryGm); + __aicore__ inline void InitParams(const ConstInfo &constInfo); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void ComputeMm1(const LICommon::RunInfo &runInfo); + + static constexpr uint64_t KEY_BUF_NUM = 3; + static constexpr uint64_t QUERY_BUF_NUM = 2; + static constexpr uint64_t L0_BUF_NUM = 2; + + static constexpr uint32_t KEY_MTE1_MTE2_EVENT = EVENT_ID2; + static constexpr uint32_t QUERY_MTE1_MTE2_EVENT = EVENT_ID5; // KEY_MTE1_MTE2_EVENT + KEY_BUF_NUM; + static constexpr uint32_t M_MTE1_EVENT = EVENT_ID3; + + static constexpr uint32_t MTE2_MTE1_EVENT = EVENT_ID2; + static constexpr uint32_t MTE1_M_EVENT = EVENT_ID2; + static constexpr uint32_t FIX_M_EVENT = EVENT_ID2; + static constexpr uint32_t M_FIX_EVENT = EVENT_ID3; + + static constexpr uint64_t M_BASIC_BLOCK = 256; + static constexpr uint64_t D_BASIC_BLOCK = 128; + static constexpr uint64_t S2_BASIC_BLOCK = 128; + + static constexpr uint64_t M_BASIC_BLOCK_L0 = 128; + static constexpr uint64_t D_BASIC_BLOCK_L0 = 128; + static constexpr uint64_t S2_BASIC_BLOCK_L0 = 128; + + static constexpr uint64_t FP16_BLOCK_CUBE = 16; + static constexpr FixpipeConfig QLI_CFG_ROW_MAJOR_UB = {CO2Layout::ROW_MAJOR, true}; + + static constexpr uint64_t QUERY_BUFFER_OFFSET = M_BASIC_BLOCK * D_BASIC_BLOCK; + static constexpr uint64_t KEY_BUFFER_OFFSET = S2_BASIC_BLOCK * D_BASIC_BLOCK; + static constexpr uint64_t L0AB_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0; + static constexpr uint64_t L0C_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0; + +protected: + __aicore__ inline void Fixp(uint64_t s1gGmOffset, uint64_t s2GmOffset, uint64_t s1gL0RealSize, + uint64_t s2L0RealSize, const LICommon::RunInfo &runInfo); + __aicore__ inline void ComputeL0c(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, const LICommon::RunInfo &runInfo); + __aicore__ inline void LoadKeyToL0b(uint64_t s2L0Offset, uint64_t s2L1RealSize, uint64_t s2L0RealSize, + const LICommon::RunInfo &runInfo); + __aicore__ inline void LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL0Offset, uint64_t s1gL1RealSize, + uint64_t s1gL0RealSize, const LICommon::RunInfo &runInfo); + __aicore__ inline void QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gL1Offset, const LICommon::RunInfo &runInfo); + __aicore__ inline void KeyNd2Nz(uint64_t s2L1RealSize, uint64_t s2GmOffset, const LICommon::RunInfo &runInfo); + __aicore__ inline void KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, const LICommon::RunInfo &runInfo); + GlobalTensor blkTableGm_; + GlobalTensor keyGm_; + GlobalTensor queryGm_; + + TBuf bufQL1_; + LocalTensor queryL1_; + TBuf bufKeyL1_; + LocalTensor keyL1_; + + TBuf bufQL0_; + LocalTensor queryL0_; + TBuf bufKeyL0_; + LocalTensor keyL0_; + + TBuf bufL0C_; + LocalTensor cL0_; + + TBuf bufUB_; + LocalTensor mm1ResUB_; + + uint64_t keyL1BufIdx_ = 0; + uint64_t queryL1Mte2BufIdx_ = 0; + uint64_t queryL1Mte1BufIdx_ = 0; + uint64_t l0BufIdx_ = 0; + uint64_t kl0BufIdx_ = 0; + + ConstInfo constInfo_; + +private: + static constexpr bool PAGE_ATTENTION = LIT::pageAttention; +}; + +template +__aicore__ inline void LightningIndexerServiceCube::InitParams(const ConstInfo &constInfo) +{ + constInfo_ = constInfo; +} + +template +__aicore__ inline void LightningIndexerServiceCube::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(bufUB_, 2 *CeilDiv(constInfo_.mBaseSize, 2) * + constInfo_.s2BaseSize * + sizeof(float)); // 大小：2(开dB) * 2 * 64 * 128 * 4 = 128KB + mm1ResUB_ = bufUB_.Get(); + pipe->InitBuffer(bufQL1_, QUERY_BUF_NUM * M_BASIC_BLOCK * D_BASIC_BLOCK * sizeof(Q_T)); + queryL1_ = bufQL1_.Get(); + pipe->InitBuffer(bufKeyL1_, KEY_BUF_NUM * S2_BASIC_BLOCK * D_BASIC_BLOCK * sizeof(K_T)); + keyL1_ = bufKeyL1_.Get(); + + pipe->InitBuffer(bufQL0_, L0_BUF_NUM * M_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0 * sizeof(Q_T)); + queryL0_ = bufQL0_.Get(); + pipe->InitBuffer(bufKeyL0_, L0_BUF_NUM * D_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0 * sizeof(K_T)); + keyL0_ = bufKeyL0_.Get(); + + pipe->InitBuffer(bufL0C_, L0_BUF_NUM * M_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0 * sizeof(float)); + cL0_ = bufL0C_.Get(); +} + +template +__aicore__ inline void +LightningIndexerServiceCube::InitMm1GlobalTensor(const GlobalTensor &blkTableGm, + const GlobalTensor &keyGm, const GlobalTensor &queryGm) +{ + blkTableGm_ = blkTableGm; + keyGm_ = keyGm; + queryGm_ = queryGm; +} + +template +__aicore__ inline void LightningIndexerServiceCube::ComputeMm1(const LICommon::RunInfo &runInfo) +{ + CrossCoreWaitFlag( + LICommon::ConstInfo::CROSS_VC_EVENT + runInfo.loop % 2); + CrossCoreWaitFlag( + LICommon::ConstInfo::CROSS_VC_EVENT + + runInfo.loop % 2 + + LICommon::ConstInfo::AIV0_AIV1_OFFSET); + uint64_t s2GmBaseOffset = runInfo.s2Idx * constInfo_.s2BaseSize; + uint64_t s1gProcessSize = runInfo.actMBaseSize; + uint64_t s2ProcessSize = runInfo.actualSingleProcessSInnerSize; + for (uint64_t s2GmOffset = 0; s2GmOffset < s2ProcessSize; s2GmOffset += S2_BASIC_BLOCK) { + WaitFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM); + uint64_t s2L1RealSize = + s2GmOffset + S2_BASIC_BLOCK > s2ProcessSize ? s2ProcessSize - s2GmOffset : S2_BASIC_BLOCK; + if (PAGE_ATTENTION) { + KeyNd2NzForPA(s2L1RealSize, s2GmBaseOffset + s2GmOffset, runInfo); + }else { + KeyNd2Nz(s2L1RealSize, s2GmOffset, runInfo); + } + + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + // s1gProcessSize当前必定不会超过2倍的s1g basic block + for (uint64_t s1gGmOffset = 0; s1gGmOffset < s1gProcessSize; s1gGmOffset += constInfo_.mBaseSize) { + uint64_t s1gL1RealSize = + s1gGmOffset + constInfo_.mBaseSize > s1gProcessSize + ? s1gProcessSize - s1gGmOffset + : constInfo_.mBaseSize; + uint64_t s1gL1SizeAlign2G = CeilAlign(s1gL1RealSize, 2 * constInfo_.gSize); + if (runInfo.isFirstS2InnerLoop && s2GmOffset == 0) { + queryL1Mte2BufIdx_++; + queryL1Mte1BufIdx_ = queryL1Mte2BufIdx_; + WaitFlag(QUERY_MTE1_MTE2_EVENT + queryL1Mte2BufIdx_ % QUERY_BUF_NUM); + QueryNd2Nz(s1gL1RealSize, s1gGmOffset, runInfo); + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + } else { + queryL1Mte1BufIdx_ = + queryL1Mte2BufIdx_ - (CeilDiv(s1gProcessSize, constInfo_.mBaseSize) - 1 - (s1gGmOffset > 0)); + } + for (uint64_t s2L1Offset = 0; s2L1Offset < s2L1RealSize; s2L1Offset += S2_BASIC_BLOCK_L0) { + uint64_t s2L0RealSize = + s2L1Offset + S2_BASIC_BLOCK_L0 > s2L1RealSize ? s2L1RealSize - s2L1Offset : S2_BASIC_BLOCK_L0; + + uint64_t l0Stride = constInfo_.mBaseSize; + if (constInfo_.splitMFlag) { + l0Stride /= 2; + } + + for (uint64_t s1gL1Offset = 0; s1gL1Offset < s1gL1SizeAlign2G; s1gL1Offset += l0Stride) { + WaitFlag(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM); + uint64_t s1gL0RealSize = + s1gL1Offset + constInfo_.mBaseSize > s1gL1SizeAlign2G + ? s1gL1SizeAlign2G - s1gL1Offset + : constInfo_.mBaseSize; + if (constInfo_.splitMFlag) { + s1gL0RealSize = 128; // g=64, topK=2k时固定m=128 + } + LoadQueryToL0a(s1gGmOffset, s1gL1Offset, s1gL1SizeAlign2G, s1gL0RealSize, runInfo); + if (s1gL1Offset == 0) { + LoadKeyToL0b(s2L1Offset, s2L1RealSize, s2L0RealSize, runInfo); + } + + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + + WaitFlag(FIX_M_EVENT + l0BufIdx_ % L0_BUF_NUM); + ComputeL0c(s1gL0RealSize, s2L0RealSize, runInfo); + + SetFlag(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM); + + bool lastIter = s1gL1Offset + l0Stride >= s1gL1SizeAlign2G; + if (lastIter) { + kl0BufIdx_++; + } + + Fixp(s1gGmOffset + s1gL1Offset, s2GmOffset + s2L1Offset, s1gL0RealSize, s2L0RealSize, runInfo); + SetFlag(FIX_M_EVENT + l0BufIdx_ % L0_BUF_NUM); + l0BufIdx_++; + } + } + if (s2GmOffset + S2_BASIC_BLOCK >= s2ProcessSize && runInfo.isLastS2InnerLoop) { + SetFlag(QUERY_MTE1_MTE2_EVENT + queryL1Mte1BufIdx_ % QUERY_BUF_NUM); + } + } + SetFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM); + keyL1BufIdx_++; + } + CrossCoreSetFlag( + LICommon::ConstInfo::CROSS_CV_EVENT + + runInfo.loop % 2); + CrossCoreSetFlag( + LICommon::ConstInfo::CROSS_CV_EVENT + + runInfo.loop % 2 + + LICommon::ConstInfo::AIV0_AIV1_OFFSET); +} + +template +__aicore__ inline void LightningIndexerServiceCube::KeyNd2Nz(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const LICommon::RunInfo &runInfo) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2L1RealSize; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + // 默认一块buf最多放两份 + DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET], + keyGm_[runInfo.tensorKeyOffset + s2GmOffset * constInfo_.headDim], nd2nzPara); +} + +// blkNum, blkSize, N2, D +template +__aicore__ inline void LightningIndexerServiceCube::KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const LICommon::RunInfo &runInfo) +{ + uint64_t s2L1Offset = 0; + while (s2L1Offset < s2L1RealSize) { + uint64_t s2BlkId = (s2L1Offset + s2GmOffset) / constInfo_.kCacheBlockSize; + uint64_t s2BlkOffset = (s2L1Offset + s2GmOffset) % constInfo_.kCacheBlockSize; + uint64_t keyGmOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId) * + constInfo_.kCacheBlockSize * constInfo_.kHeadNum * constInfo_.headDim + + s2BlkOffset * constInfo_.headDim; + + uint64_t s2Mte2Size = s2L1RealSize - s2L1Offset; + s2Mte2Size = s2BlkOffset + s2Mte2Size >= constInfo_.kCacheBlockSize + ? constInfo_.kCacheBlockSize - s2BlkOffset + : s2Mte2Size; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2Mte2Size; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET + s2L1Offset * FP16_BLOCK_CUBE], + keyGm_[keyGmOffset], nd2nzPara); + + s2L1Offset += s2Mte2Size; + } +} + +// batch, s1, n2, g, d +template +__aicore__ inline void LightningIndexerServiceCube::QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gGmOffset, + const LICommon::RunInfo &runInfo) +{ + uint64_t dstNzC0Stride = CeilAlign(s1gL1RealSize, 2 * constInfo_.gSize); + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s1gL1RealSize; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(dstNzC0Stride, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + // 默认一块buf最多放两份 + DataCopy(queryL1_[(queryL1Mte2BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET], + queryGm_[runInfo.tensorQueryOffset + s1gGmOffset * constInfo_.headDim], nd2nzPara); +} + +template +__aicore__ inline void LightningIndexerServiceCube::LoadQueryToL0a(uint64_t s1gGmOffset, + uint64_t s1gL1Offset, uint64_t s1gL1RealSize, + uint64_t s1gL0RealSize, const LICommon::RunInfo &runInfo) +{ + LoadData2DParamsV2 loadData2DParamsV2; + if (constInfo_.splitMFlag && runInfo.actMBaseSize > 128) { // 非尾块，切M + uint64_t dstOffset = 0; + loadData2DParamsV2.kStartPosition = 0; + loadData2DParamsV2.mStep = CeilDiv(64, BLOCK_CUBE); + loadData2DParamsV2.kStep = CeilDiv(constInfo_.headDim, FP16_BLOCK_CUBE); + loadData2DParamsV2.srcStride = CeilDiv(s1gL1RealSize, BLOCK_CUBE); + loadData2DParamsV2.dstStride = CeilDiv(s1gL0RealSize, BLOCK_CUBE); + loadData2DParamsV2.ifTranspose = false; + for (int i = 0; i < 2; i++) { + loadData2DParamsV2.mStartPosition = CeilDiv((s1gL1Offset / 2) + i * 128, BLOCK_CUBE); + dstOffset = i * 64 * 16; + + LoadData(queryL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET + dstOffset], + queryL1_[(queryL1Mte1BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET], loadData2DParamsV2); + } + } else { + loadData2DParamsV2.mStartPosition = CeilDiv(s1gL1Offset, BLOCK_CUBE); + loadData2DParamsV2.kStartPosition = 0; + loadData2DParamsV2.mStep = CeilDiv(s1gL0RealSize, BLOCK_CUBE); + loadData2DParamsV2.kStep = CeilDiv(constInfo_.headDim, FP16_BLOCK_CUBE); + loadData2DParamsV2.srcStride = CeilDiv(s1gL1RealSize, BLOCK_CUBE); + loadData2DParamsV2.dstStride = CeilDiv(s1gL0RealSize, BLOCK_CUBE); + loadData2DParamsV2.ifTranspose = false; + + LoadData(queryL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], + queryL1_[(queryL1Mte1BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET], loadData2DParamsV2); + } +} + +template +__aicore__ inline void LightningIndexerServiceCube::LoadKeyToL0b(uint64_t s2L1Offset, + uint64_t s2L1RealSize, uint64_t s2L0RealSize, + const LICommon::RunInfo &runInfo) +{ + LoadData2DParamsV2 loadData2DParamsV2; + loadData2DParamsV2.mStartPosition = CeilDiv(s2L1Offset, BLOCK_CUBE); + loadData2DParamsV2.kStartPosition = 0; + loadData2DParamsV2.mStep = CeilDiv(s2L0RealSize, BLOCK_CUBE); + loadData2DParamsV2.kStep = CeilDiv(constInfo_.headDim, FP16_BLOCK_CUBE); + loadData2DParamsV2.srcStride = CeilDiv(s2L1RealSize, BLOCK_CUBE); + loadData2DParamsV2.dstStride = CeilDiv(s2L0RealSize, BLOCK_CUBE); + loadData2DParamsV2.ifTranspose = false; + + LoadData(keyL0_[(kl0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], + keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET], loadData2DParamsV2); +} + +template +__aicore__ inline void LightningIndexerServiceCube::ComputeL0c(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, + const LICommon::RunInfo &runInfo) +{ + MmadParams mmadParams; + mmadParams.m = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + mmadParams.n = s2L0RealSize; + mmadParams.k = constInfo_.headDim; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + Mmad(cL0_[(l0BufIdx_ % L0_BUF_NUM) * L0C_BUFFER_OFFSET], queryL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], + keyL0_[(kl0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } +} + +template +__aicore__ inline void LightningIndexerServiceCube::Fixp(uint64_t s1gGmOffset, + uint64_t s2GmOffset, uint64_t s1gL0RealSize, + uint64_t s2L0RealSize, const LICommon::RunInfo &runInfo) +{ + SetFlag(M_FIX_EVENT + l0BufIdx_ % L0_BUF_NUM); + WaitFlag(M_FIX_EVENT + l0BufIdx_ % L0_BUF_NUM); + + static_assert(S2_BASIC_BLOCK == S2_BASIC_BLOCK_L0 && S2_BASIC_BLOCK_L0 == 128); + // s1gL0RealSize：2*gSize(128)对齐, 最大256 + // s2L0RealSize <= S2_BASIC_BLOCK_L0, 未约束 + uint32_t nSize = (s2L0RealSize + 7) >> 3 << 3; // 32B对齐 + uint32_t mSize = (s1gL0RealSize + 1) >> 1 << 1; + FixpipeParamsC310 fixpipeParams; + // 固定参数 + fixpipeParams.mSize = mSize; + fixpipeParams.srcStride = mSize; // 已16对齐 + fixpipeParams.dstStride = UB_BANK_DEPTH_STRIDE / sizeof(float); // 落到同一个bank + fixpipeParams.dualDstCtl = 1; // 双目标模式，按M维度拆分， M / 2 * N写入每个UB，M必须为2的倍数 + + uint64_t dstOffset = 0; + if (constInfo_.splitMFlag && runInfo.actMBaseSize > 128) { // 非尾块，切M + dstOffset = s1gGmOffset * 64; + } + + // nSize已保证N方向32B对齐 + if (nSize <= (256 / sizeof(float))) { + // N方向小于一个bank(256B), 只需搬一个ND块, 且不用补齐 + fixpipeParams.nSize = nSize; + fixpipeParams.params.ndNum = 1; + fixpipeParams.params.srcNdStride = 0; + fixpipeParams.params.dstNdStride = 0; + } else { + // N方向在(256B, 512B]范围， 直接按512B搬, 注意此时不能开unitflag + fixpipeParams.nSize = S2_BASIC_BLOCK_L0 / 2; // 分2个ND搬, S2_BASIC_BLOCK_L0不为128会有问题 + fixpipeParams.params.ndNum = 2; + fixpipeParams.params.srcNdStride = ((fixpipeParams.mSize + 15) / 16) * fixpipeParams.nSize; + fixpipeParams.params.dstNdStride = constInfo_.s2BaseSize * constInfo_.mBaseSize / 2; + } + Fixpipe(mm1ResUB_[(runInfo.loop % 2) * constInfo_.s2BaseSize / 2 + dstOffset], + cL0_[(l0BufIdx_ % L0_BUF_NUM) * L0C_BUFFER_OFFSET], fixpipeParams); +} + +template +__aicore__ inline void LightningIndexerServiceCube::AllocEventID() +{ + SetMMLayoutTransform(true); + SetFlag(KEY_MTE1_MTE2_EVENT + 0); + SetFlag(KEY_MTE1_MTE2_EVENT + 1); + SetFlag(KEY_MTE1_MTE2_EVENT + 2); + + SetFlag(QUERY_MTE1_MTE2_EVENT + 0); + SetFlag(QUERY_MTE1_MTE2_EVENT + 1); + + SetFlag(M_MTE1_EVENT + 0); + SetFlag(M_MTE1_EVENT + 1); + + SetFlag(FIX_M_EVENT + 0); + SetFlag(FIX_M_EVENT + 1); +} + +template +__aicore__ inline void LightningIndexerServiceCube::FreeEventID() +{ + SetMMLayoutTransform(false); + WaitFlag(KEY_MTE1_MTE2_EVENT + 0); + WaitFlag(KEY_MTE1_MTE2_EVENT + 1); + WaitFlag(KEY_MTE1_MTE2_EVENT + 2); + + WaitFlag(QUERY_MTE1_MTE2_EVENT + 0); + WaitFlag(QUERY_MTE1_MTE2_EVENT + 1); + + WaitFlag(M_MTE1_EVENT + 0); + WaitFlag(M_MTE1_EVENT + 1); + + WaitFlag(FIX_M_EVENT + 0); + WaitFlag(FIX_M_EVENT + 1); +} +} // namespace LIKernel +#endif diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/lightning_indexer_service_vector.h b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/lightning_indexer_service_vector.h new file mode 100644 index 000000000..d23fbc7e9 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/lightning_indexer_service_vector.h @@ -0,0 +1,575 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_service_vector.h + * \brief + */ +#ifndef LIGHTNING_INDEXER_SERVICE_VECTOR_H +#define LIGHTNING_INDEXER_SERVICE_VECTOR_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../lightning_indexer_common.h" +#include "../arch35/vf/lightning_indexer_vector1.h" +#include "../arch35/vf/lightning_indexer_topk.h" + +namespace LIKernel { +using namespace LICommon; +constexpr uint32_t TRUNK_LEN_16K = 16384; +constexpr uint32_t TRUNK_LEN_8K = 8192; +constexpr uint32_t TOPK_LEN_4K = 4096; + +template +struct LightningIndexerTypeTraits { + using weightsType = Q_T; // 默认：weightsType绑定Q_T +}; + +template +struct LightningIndexerTypeTraits { + using weightsType = float; // W_T=float时，强制weightsType为float +}; +template +class LightningIndexerServiceVector { +public: + // =================================类型定义区================================= + static constexpr LI_LAYOUT LAYOUT_T = LIT::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = LIT::keyLayout; + static constexpr bool PAGE_ATTENTION = LIT::pageAttention; + static constexpr bool DT_W_FLAG = LIT::weightsTypeFlag; + using Q_T = typename LIT::queryType; + using K_T = typename LIT::keyType; + using W_T = typename LightningIndexerTypeTraits::type>::weightsType; + + __aicore__ inline LightningIndexerServiceVector(){}; + __aicore__ inline void ProcessVec1(const LICommon::RunInfo &info); + __aicore__ inline void ProcessTopK(const LICommon::RunInfo &info); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitParams(const struct LICommon::ConstInfo &constInfo, + const LITilingData *__restrict tilingData); + __aicore__ inline void InitVecWorkspaceTensor(GlobalTensor scoreGm); + __aicore__ inline void InitVecInputTensor(GlobalTensor weightsGm, GlobalTensor indiceOutGm, + GlobalTensor valueOutGm, GlobalTensor blockTableGm); + __aicore__ inline void CleanInvalidOutput(int64_t invalidS1offset); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + +protected: + GlobalTensor scoreGm; + GlobalTensor weightsGm; + GlobalTensor indiceOutGm; + GlobalTensor valueOutGm; + GlobalTensor blockTableGm; + // =================================常量区================================= + static constexpr uint32_t VEC1_V_MTE2_EVENT = EVENT_ID0; + static constexpr uint32_t VEC1_MTE2_V_EVENT = EVENT_ID1; + static constexpr uint32_t VEC1_V_MTE3_EVENT = EVENT_ID2; + static constexpr uint32_t VEC1_MTE3_V_EVENT = EVENT_ID3; + + static constexpr uint32_t TOPK_V_MTE2_EVENT = EVENT_ID4; + static constexpr uint32_t TOPK_MTE2_V_EVENT = EVENT_ID5; + static constexpr uint32_t TOPK_V_MTE3_EVENT = EVENT_ID6; + static constexpr uint32_t TOPK_MTE3_V_EVENT = EVENT_ID7; + + static constexpr uint32_t MTE3_MTE2_EVENT = EVENT_ID0; + static constexpr uint32_t V_MTE2_EVENT = EVENT_ID7; + static constexpr uint32_t V_MTE2_EVENT1 = EVENT_ID2; + static constexpr uint32_t V_MTE2_EVENT2 = EVENT_ID3; + static constexpr uint32_t V_MTE2_EVENT3 = EVENT_ID5; + +private: + // ================================Local Buffer区==================================== + + // tmp buff for vector + TBuf resMm1Buf_; + LocalTensor resMm1UB_; + // tmp buff for weight + TBuf weightBuf_; + LocalTensor weightUB_; + // tmp buff for weight cast float + TBuf weightFloatBuf_; + LocalTensor weightFloatUB_; + + // tmp buff for out + TBuf outBuf_; + LocalTensor vec1OutUB_; + + // tmp buff for returnValue K_T + TBuf valueOutBuf_; + LocalTensor valueOutLocal_; + + // tmp buff for topk + TBuf mrgValueBuf_; + LocalTensor mrgValueLocal_; + + TBuf indicesOutBuf_; + LocalTensor indicesOutLocal_; + + TBuf scoreOutBuf_; + LocalTensor scoreOutLocal_; + + TBuf topkSharedTmpBuf_; + LocalTensor topkSharedTmpLocal_; + + int32_t blockId_ = -1; + // para for vector + int32_t groupInner_ = 0; + int32_t globalTopkNum_ = 0; + int64_t blockS2StartIdx_ = 0; + int32_t gSize_ = 0; + int32_t kSeqSize_ = 0; + int32_t kHeadNum_ = 0; + int32_t qHeadNum_ = 0; + int32_t s1BaseSize_ = 0; + int32_t s2BaseSize_ = 0; + int32_t kCacheBlockSize_ = 0; + int32_t maxBlockNumPerBatch_ = 0; + uint32_t topkCount_ = 0; + uint32_t topkCountAlign256_ = 0; // topkCount对齐到256(直方图需要)，支持topk泛化 + uint32_t trunkLen_ = 0; + bool returnValueFlag = false; + + struct LICommon::ConstInfo constInfo_; + topk::LITopk topkOp_; +}; + +template +__aicore__ inline void LightningIndexerServiceVector::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(resMm1Buf_, 2 * CeilDiv(constInfo_.mBaseSize, 2) * s2BaseSize_ * sizeof(float)); + resMm1UB_ = resMm1Buf_.Get(); + + pipe->InitBuffer(weightBuf_, 2 * CeilDiv(s1BaseSize_, 2) * UB_BANK_DEPTH_STRIDE); + weightUB_ = weightBuf_.Get(); + pipe->InitBuffer(weightFloatBuf_, 2 * CeilDiv(s1BaseSize_, 2) * UB_BANK_DEPTH_STRIDE); + weightFloatUB_ = weightFloatBuf_.Get(); + pipe->InitBuffer(outBuf_, + 2 * CeilDiv(s1BaseSize_, 2) * s2BaseSize_ * sizeof(uint16_t)); // 大小：2(开dB) * 2 * 128 * 4 = 2KB + vec1OutUB_ = outBuf_.Get(); // out + + // Topk + pipe->InitBuffer(mrgValueBuf_, + (topkCountAlign256_ + trunkLen_) * sizeof(uint16_t)); + mrgValueLocal_ = mrgValueBuf_.Get(); + // returnvalue + if (topkCount_ <= 2048) { + pipe->InitBuffer(valueOutBuf_, topkCountAlign256_ * sizeof(K_T)); + valueOutLocal_ = valueOutBuf_.Get(); + } else { // sparseCount > 2k时，复用return value相关UB + valueOutLocal_ = mrgValueBuf_.Get(); // returnValue float + } + + // 大小：(topkCountAlign256_ + 64) * 4 64:duplicate刷-1需要额外空间 + pipe->InitBuffer(indicesOutBuf_, + (topkCountAlign256_ + 64) * sizeof(uint32_t)); + indicesOutLocal_ = indicesOutBuf_.Get(); + + pipe->InitBuffer(scoreOutBuf_, topkCountAlign256_ * sizeof(uint16_t)); + scoreOutLocal_ = scoreOutBuf_.Get(); + + uint64_t topkSharedTmpSize = topkOp_.GetSharedTmpBufferSize(); + pipe->InitBuffer(topkSharedTmpBuf_, topkSharedTmpSize); + topkSharedTmpLocal_ = topkSharedTmpBuf_.Get(); + topkOp_.InitBuffers(topkSharedTmpLocal_); +} + +template +__aicore__ inline void LightningIndexerServiceVector::InitParams(const struct LICommon::ConstInfo &constInfo, + const LITilingData *__restrict tilingData) +{ + this->constInfo_ = constInfo; + blockS2StartIdx_ = 0; + gSize_ = constInfo.gSize; + kSeqSize_ = constInfo.kSeqSize; + // define N2 para + kHeadNum_ = constInfo.kHeadNum; + qHeadNum_ = constInfo.qHeadNum; + // define MMBase para + s1BaseSize_ = constInfo.s1BaseSize; // 4 + s2BaseSize_ = constInfo.s2BaseSize; // 128 + kCacheBlockSize_ = constInfo.kCacheBlockSize; + maxBlockNumPerBatch_ = constInfo.maxBlockNumPerBatch; + returnValueFlag = constInfo.returnValueFlag; + blockId_ = GetBlockIdx(); + trunkLen_ = constInfo.sparseCount >= TOPK_LEN_4K ? TRUNK_LEN_8K : TRUNK_LEN_16K; + topkCount_ = constInfo.sparseCount; + topkOp_.Init(topkCount_, trunkLen_); + topkCountAlign256_ = LICommon::Align(constInfo.sparseCount, (uint64_t)256); // topkCount对齐到256 +} + +template +__aicore__ inline void LightningIndexerServiceVector::InitVecInputTensor(GlobalTensor weightsGm, + GlobalTensor indiceOutGm, + GlobalTensor valueOutGm, + GlobalTensor blockTableGm) +{ + this->weightsGm = weightsGm; + this->indiceOutGm = indiceOutGm; + this->valueOutGm = valueOutGm; + this->blockTableGm = blockTableGm; +} + +template +__aicore__ inline void LightningIndexerServiceVector::InitVecWorkspaceTensor(GlobalTensor scoreGm) +{ + this->scoreGm = scoreGm; // resucesum*k +} + +template +__aicore__ inline void LightningIndexerServiceVector::AllocEventID() +{ + SetFlag(VEC1_V_MTE2_EVENT + 0); + SetFlag(VEC1_V_MTE2_EVENT + 1); + SetFlag(VEC1_MTE3_V_EVENT + 0); + SetFlag(VEC1_MTE3_V_EVENT + 1); + + SetFlag(TOPK_V_MTE2_EVENT); + SetFlag(TOPK_MTE3_V_EVENT); + SetFlag(V_MTE2_EVENT1); +} + +template +__aicore__ inline void LightningIndexerServiceVector::FreeEventID() +{ + WaitFlag(VEC1_V_MTE2_EVENT + 0); + WaitFlag(VEC1_V_MTE2_EVENT + 1); + WaitFlag(VEC1_MTE3_V_EVENT + 0); + WaitFlag(VEC1_MTE3_V_EVENT + 1); + + WaitFlag(TOPK_V_MTE2_EVENT); + WaitFlag(TOPK_MTE3_V_EVENT); + WaitFlag(V_MTE2_EVENT1); +} + +template +__aicore__ inline void LightningIndexerServiceVector::CleanInvalidOutput(int64_t invalidS1Offset) +{ + // init -1 and copy to output + uint64_t dealSize = constInfo_.sparseCount; + GlobalTensor indexOutput = indiceOutGm[invalidS1Offset]; + AscendC::InitGlobalMemory(indexOutput, dealSize, constInfo_.INVALID_IDX); + if (returnValueFlag) { + SetFlag(TOPK_MTE3_V_EVENT); + WaitFlag(TOPK_MTE3_V_EVENT); + Duplicate(valueOutLocal_.template ReinterpretCast(), constInfo_.INVALID_VAL, constInfo_.sparseCount); + + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + + AscendC::DataCopyParams copyOutValueParams; + copyOutValueParams.blockCount = 1; + copyOutValueParams.blockLen = constInfo_.sparseCount * sizeof(K_T); + copyOutValueParams.srcStride = 0; + copyOutValueParams.dstStride = 0; + AscendC::DataCopyPad(valueOutGm[invalidS1Offset], valueOutLocal_, copyOutValueParams); + } +} + +template +__aicore__ inline void LightningIndexerServiceVector::ProcessVec1(const LICommon::RunInfo &info) +{ + auto pingpong = (info.loop % 2); + auto s1BaseSizePerAIV = CeilDiv(s1BaseSize_, 2); + int64_t curS1Idx = info.gS1Idx * s1BaseSize_; + int64_t curS2Idx = info.s2Idx * s2BaseSize_; + int64_t curS1ProcNum = curS1Idx + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + int64_t curAivS1Idx = curS1Idx + (blockId_ % 2) * CeilDiv(curS1ProcNum, 2); + int64_t curAivS1ProcNum = (blockId_ % 2 == 0) ? CeilDiv(curS1ProcNum, 2) : curS1ProcNum / 2; + if (curAivS1ProcNum == 0) { + CrossCoreWaitFlag( + LICommon::ConstInfo::CROSS_CV_EVENT + pingpong + ); // V核等C核计算完mm1，mm1Res已搬运到UB + CrossCoreSetFlag( + LICommon::ConstInfo::CROSS_VC_EVENT + pingpong + ); // V核处理完，通知C核可以把mm1Res搬运到UB + return; + } + WaitFlag(VEC1_V_MTE2_EVENT + pingpong); + // weightsGm --> weightUB_ + int64_t weightGmOffset = info.tensorWeightsOffset + curAivS1Idx * kHeadNum_ * gSize_; + DataCopyPadExtParams padWeightsParams{false, 0, 0, 0}; + DataCopyExtParams wDataCopyExtParams; + wDataCopyExtParams.blockCount = curAivS1ProcNum; + wDataCopyExtParams.blockLen = gSize_ * sizeof(W_T); + wDataCopyExtParams.srcStride = 0; + wDataCopyExtParams.dstStride = (UB_BANK_DEPTH_STRIDE - wDataCopyExtParams.blockLen) / 32; + DataCopyPad(weightUB_[pingpong * (UB_BANK_STRIDE / sizeof(W_T))], + weightsGm[weightGmOffset], wDataCopyExtParams, padWeightsParams); + + SetFlag(VEC1_MTE2_V_EVENT + pingpong); + WaitFlag(VEC1_MTE2_V_EVENT + pingpong); + WaitFlag(VEC1_MTE3_V_EVENT + pingpong); + + // CV同步 + CrossCoreWaitFlag( + LICommon::ConstInfo::CROSS_CV_EVENT + info.loop % 2 + ); // V核等C核计算完mm1，mm1Res已搬运到UB + + auto outBase = vec1OutUB_[pingpong * (UB_BANK_STRIDE / sizeof(uint16_t))]; + auto weightBase = weightUB_[pingpong * (UB_BANK_STRIDE / sizeof(W_T))]; + auto weightFloatBase = weightFloatUB_[pingpong * (UB_BANK_STRIDE / sizeof(float))]; + auto qkBase = resMm1UB_[pingpong * (UB_BANK_STRIDE / sizeof(float))]; + auto qkVLstride = (UB_BANK_DEPTH_STRIDE / sizeof(float)) / 2 * constInfo_.mBaseSize; + + vector1::BatchMulWeightAndReduceSum(outBase, UB_BANK_DEPTH_STRIDE / sizeof(uint16_t), + qkBase, qkVLstride, (uint32_t)(gSize_ * UB_BANK_DEPTH_STRIDE / sizeof(float)), + weightBase, UB_BANK_DEPTH_STRIDE / sizeof(W_T), weightFloatBase, + gSize_, curAivS1ProcNum); + SetFlag(VEC1_V_MTE2_EVENT + pingpong); + SetFlag(VEC1_V_MTE3_EVENT + pingpong); + WaitFlag(VEC1_V_MTE3_EVENT + pingpong); + // outUB_ ---> scoreGm + int64_t vec1OutGmOffset = blockId_ % 2 == 0 + ? curS2Idx + : s1BaseSizePerAIV * LICommon::Align( + (uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_ + ) + curS2Idx; + DataCopyExtParams copyOutParams; + copyOutParams.blockCount = curAivS1ProcNum; + copyOutParams.blockLen = s2BaseSize_ * sizeof(uint16_t); + copyOutParams.srcStride = (UB_BANK_DEPTH_STRIDE - UB_BANK_STRIDE) / 32; + copyOutParams.dstStride = (LICommon::Align( + (uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_ + ) - s2BaseSize_) * sizeof(uint16_t); + DataCopyPad(scoreGm[vec1OutGmOffset], outBase, copyOutParams); + SetFlag(VEC1_MTE3_V_EVENT + pingpong); + CrossCoreSetFlag( + LICommon::ConstInfo::CROSS_VC_EVENT + pingpong + ); // V核处理完，通知C核可以把mm1Res搬运到UB +} + +template +__aicore__ inline void LightningIndexerServiceVector::ProcessTopK(const LICommon::RunInfo &info) +{ + SetFlag(MTE3_MTE2_EVENT); + WaitFlag(MTE3_MTE2_EVENT); + + int64_t curS1Idx = info.gS1Idx * s1BaseSize_; + int64_t curS2Idx = info.s2Idx * s2BaseSize_; + int64_t curS1ProcNum = curS1Idx + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + int64_t curAivS1Idx = curS1Idx + (blockId_ % 2) * CeilDiv(curS1ProcNum, 2); + int64_t curAivS1ProcNum = (blockId_ % 2 == 0) ? CeilDiv(curS1ProcNum, 2) : curS1ProcNum / 2; + + AscendC::DataCopyExtParams copyInParams; + copyInParams.blockCount = 1; + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + + AscendC::DataCopyParams copyOutParams; + copyOutParams.blockCount = 1; + copyOutParams.blockLen = topkCount_ * sizeof(uint32_t); // bytes + copyOutParams.srcStride = 0; + copyOutParams.dstStride = 0; + + int32_t cuRealAcSeq = info.actS2Size; + if (constInfo_.attenMaskFlag) { + cuRealAcSeq = info.actS2SizeOrig - info.actS1Size + curAivS1Idx + 1; + } + + int32_t validS2Len = cuRealAcSeq; + for (uint32_t i = 0; i < curAivS1ProcNum; i++) { + uint32_t rowIdx = blockId_ % 2 * CeilDiv(curS1ProcNum, 2) + i; + uint32_t vecOffset = blockId_ % 2 * CeilDiv(s1BaseSize_, 2) + i; + + uint16_t zero = 0; + int32_t neg = -1; + if (constInfo_.attenMaskFlag) { + validS2Len = (int32_t)i + cuRealAcSeq; + } + if (validS2Len <= 0) { + WaitFlag(TOPK_MTE3_V_EVENT); + Duplicate(indicesOutLocal_.ReinterpretCast(), neg, topkCount_); + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + AscendC::DataCopyPad(indiceOutGm[info.indiceOutOffset + (curS1Idx + rowIdx) * topkCount_], + indicesOutLocal_.ReinterpretCast(), + copyOutParams); + SetFlag(TOPK_MTE3_V_EVENT); + if (returnValueFlag) { + WaitFlag(TOPK_MTE3_V_EVENT); + Duplicate(valueOutLocal_.template ReinterpretCast(), constInfo_.INVALID_VAL, topkCount_); + + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + + AscendC::DataCopyParams copyOutValueParams; + copyOutValueParams.blockCount = 1; + copyOutValueParams.blockLen = topkCount_ * sizeof(K_T); + copyOutValueParams.srcStride = 0; + copyOutValueParams.dstStride = 0; + AscendC::DataCopyPad( + valueOutGm[info.valueOutOffset + (curS1Idx + rowIdx) * topkCount_], + valueOutLocal_, + copyOutValueParams); + SetFlag(TOPK_MTE3_V_EVENT); + } + continue; + } + + WaitFlag(TOPK_V_MTE2_EVENT); + WaitFlag(TOPK_MTE3_V_EVENT); + + AscendC::DataCopyPadExtParams padParams{true, 0, 0, 0}; + if (validS2Len >= topkCount_) { + uint32_t s2LoopNum = (validS2Len + trunkLen_ - 1) / trunkLen_; + if (s2LoopNum == 1) { + uint32_t validS2LenAlign = LICommon::Align(validS2Len, (int32_t)256); + Duplicate(mrgValueLocal_[validS2Len / 256 * 256], zero, validS2LenAlign - validS2Len / 256 * 256); + SetFlag(V_MTE2_EVENT); + WaitFlag(V_MTE2_EVENT); + copyInParams.blockLen = validS2Len * sizeof(uint16_t); // byte + AscendC::DataCopyPadExtParams padParams{true, 0, 0, 0}; + AscendC::DataCopyPad( + mrgValueLocal_, + scoreGm[vecOffset * LICommon::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_)], + copyInParams, padParams); + SetFlag(TOPK_MTE2_V_EVENT); + WaitFlag(TOPK_MTE2_V_EVENT); + topkOp_(mrgValueLocal_, indicesOutLocal_, scoreOutLocal_, validS2LenAlign, 0, 1, returnValueFlag); + } else { + for (uint32_t loopIdx = 0; loopIdx < s2LoopNum; loopIdx++) { + if (loopIdx == 0) { + copyInParams.blockLen = trunkLen_ * sizeof(uint16_t); // byte + AscendC::DataCopyPad( + mrgValueLocal_, + scoreGm[vecOffset * LICommon::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_)], + copyInParams, padParams); + SetFlag(TOPK_MTE2_V_EVENT); + WaitFlag(TOPK_MTE2_V_EVENT); + topkOp_(mrgValueLocal_, indicesOutLocal_, + scoreOutLocal_, trunkLen_, loopIdx, + s2LoopNum, returnValueFlag); + continue; + } + SetFlag(V_MTE2_EVENT2); + WaitFlag(V_MTE2_EVENT2); + uint32_t validTrunkLen = (loopIdx * trunkLen_ + trunkLen_) > validS2Len + ? validS2Len % trunkLen_ + :trunkLen_; + uint32_t offset = vecOffset * + LICommon::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_) + + loopIdx * trunkLen_; + AscendC::DataCopy(mrgValueLocal_, scoreOutLocal_, topkCountAlign256_); + // topk如果没有对齐到256，则把topkCountAlign256_ - topkCount_部分刷0 + if (topkCountAlign256_ != topkCount_) { + uint64_t mask[1]; + mask[0] = ~0; + mask[0] = mask[0] << (topkCount_ % 64); + PipeBarrier(); + // 把topkCount_对齐到64刷0，此处由于duplicate的限制mask[0]刷64个数 + Duplicate(mrgValueLocal_[topkCount_ / 64 * 64], zero, mask, 1, 1, 0); + PipeBarrier(); + // 把topk剩余对齐到256的部分刷0 + Duplicate(mrgValueLocal_[topkCount_ / 64 * 64 + 64], zero, + topkCountAlign256_ - (topkCount_ / 64 * 64 + 64)); + SetFlag(V_MTE2_EVENT3); + WaitFlag(V_MTE2_EVENT3); + } + copyInParams.blockLen = validTrunkLen * sizeof(uint16_t); // byte + // TOPK 直方图一次必须计算256，输入处理数据需要和256对齐 + if ((topkCountAlign256_ + validTrunkLen) % 256 != 0) { + Duplicate(mrgValueLocal_[topkCountAlign256_ + validTrunkLen / 256 * 256], + zero, LICommon::Align(validTrunkLen, + (uint32_t)256) - validTrunkLen / 256 * 256); + SetFlag(V_MTE2_EVENT); + WaitFlag(V_MTE2_EVENT); + } + WaitFlag(V_MTE2_EVENT1); + AscendC::DataCopyPad(mrgValueLocal_[topkCountAlign256_], scoreGm[offset], copyInParams, padParams); + SetFlag(TOPK_MTE2_V_EVENT); + WaitFlag(TOPK_MTE2_V_EVENT); + topkOp_(mrgValueLocal_, indicesOutLocal_, + scoreOutLocal_, + LICommon::Align(topkCountAlign256_ + validTrunkLen, (uint32_t)256), + loopIdx, s2LoopNum, returnValueFlag); + SetFlag(V_MTE2_EVENT1); + } + } + } else { + AscendC::CreateVecIndex(indicesOutLocal_.ReinterpretCast(), (int32_t)zero, validS2Len); + if (returnValueFlag) { + copyInParams.blockLen = LICommon::Align(validS2Len, (int32_t)32) * sizeof(uint16_t); + AscendC::DataCopyPad(scoreOutLocal_, + scoreGm[vecOffset * LICommon::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_)], + copyInParams, padParams); + SetFlag(TOPK_MTE2_V_EVENT); + WaitFlag(TOPK_MTE2_V_EVENT); + } + } + + if (validS2Len < topkCount_) { + uint64_t mask[1]; + mask[0] = ~0; + mask[0] = mask[0] << (validS2Len % 8); + PipeBarrier(); + Duplicate(indicesOutLocal_.ReinterpretCast()[validS2Len / 8 * 8], neg, mask, 1, 1, 0); + } + + if (validS2Len / 8 * 8 + 64 < topkCount_) { + PipeBarrier(); + Duplicate(indicesOutLocal_.ReinterpretCast()[validS2Len / 8 * 8 + 64], + neg, topkCount_ - (validS2Len / 8 * 8 + 64)); + } + + SetFlag(TOPK_V_MTE2_EVENT); + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + AscendC::DataCopyPad(indiceOutGm[info.indiceOutOffset + (curS1Idx + rowIdx) * topkCount_], + indicesOutLocal_.ReinterpretCast(), copyOutParams); + + + // 是否返回Value值 + if (returnValueFlag) { + WaitFlag(TOPK_V_MTE2_EVENT); + // uint16_t -> bfloat16 + if (std::is_same_v) { + vector1::UIntToFloatReturnValue(valueOutLocal_.template ReinterpretCast(), + scoreOutLocal_, topkCountAlign256_); + } else { + vector1::UIntToFloatReturnValue(valueOutLocal_.template ReinterpretCast(), + scoreOutLocal_, topkCountAlign256_); + } + + if (validS2Len < topkCount_) { + uint64_t mask[1]; + mask[0] = ~0; + mask[0] = mask[0] << (validS2Len % 16); + PipeBarrier(); + Duplicate(valueOutLocal_.template ReinterpretCast()[validS2Len / 16 * 16], + constInfo_.INVALID_VAL, mask, 1, 1, 0); + } + if (validS2Len / 16 * 16 + 64 < topkCount_) { + PipeBarrier(); + Duplicate(valueOutLocal_.template ReinterpretCast()[validS2Len / 16 * 16 + 64], + constInfo_.INVALID_VAL, topkCount_ - (validS2Len / 16 * 16 + 64)); + } + SetFlag(TOPK_V_MTE2_EVENT); + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + AscendC::DataCopyParams copyOutValueParams; + copyOutValueParams.blockCount = 1; + copyOutValueParams.blockLen = topkCount_ * sizeof(K_T); // bytes + copyOutValueParams.srcStride = 0; + copyOutValueParams.dstStride = 0; + // 搬运到GM + AscendC::DataCopyPad( + valueOutGm[info.valueOutOffset + (curS1Idx + rowIdx) * topkCount_], + valueOutLocal_, copyOutValueParams); + } + SetFlag(TOPK_MTE3_V_EVENT); + } +} +} // namespace LIKernel +#endif diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/vf/lightning_indexer_topk.h b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/vf/lightning_indexer_topk.h new file mode 100644 index 000000000..2335b99da --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/vf/lightning_indexer_topk.h @@ -0,0 +1,181 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. +  */ + +/*! + * \file lightning_indexer_topk.h + * \brief + */ +#ifndef LIGHTNING_INDEXER_TOPK_H +#define LIGHTNING_INDEXER_TOPK_H + +#include "kernel_operator.h" +#include "vf_topk.h" +#include "vf_topk_16_gather.h" + +namespace topk { +template +class LITopk { +public: + __aicore__ inline void operator()(LocalTensor& outputIdxLocal, + LocalTensor& inputLocal, + uint32_t s2SeqLen) + { + } +}; + +template<> +class LITopk { +public: + static __aicore__ inline uint32_t GetSharedTmpBufferSize(uint32_t topK) + { + return 2 * topK * sizeof(uint32_t) + 5 * 256 * sizeof(uint32_t) + 64 * sizeof(uint32_t) + + (topK + 64) * sizeof(uint32_t); // for output value tensor + } + + static __aicore__ inline uint32_t GetIndexBufferSize(uint32_t topK) + { + return (topK + 64) * sizeof(uint32_t); + } + + __aicore__ inline void Init(uint32_t topK) + { + this->topK = topK; + } + + __aicore__ inline void InitBuffers(LocalTensor& sharedTmpBuffer) + { + tmpIdxLocal = sharedTmpBuffer[0]; + tmpValueLocal = tmpIdxLocal[topK]; + histogramsLocal = tmpValueLocal[topK]; + idx0Local = histogramsLocal[256]; + idx1Local = idx0Local[256]; + idx2Local = idx1Local[256]; + idx3Local = idx2Local[256]; + nkValueLocal = idx3Local[256]; + outputValueLocal = nkValueLocal[64]; + } + + __aicore__ inline void operator()(LocalTensor& outputIdxLocal, + LocalTensor& inputLocal, + uint32_t s2SeqLen) + { + topkb32::LiTopKVF(outputIdxLocal, // filter阶段使用输出value Buf topK * 4B + outputValueLocal, // filter阶段使用输出 Idx Buf topK * 4B + inputLocal, // 输入 s2SeqLen * 4B + tmpIdxLocal, // filter阶段使用暂存index Buf topK * 4B + tmpValueLocal, // filter阶段使用暂存value Buf topK * 4B + histogramsLocal, // 直方图的临时Buf 256 * 4B + idx0Local, // 输入数据第1个8位Buf 256 * 4B + idx1Local, // 输入数据第2个8位Buf 256 * 4B + idx2Local, // 输入数据第3个8位Buf 256 * 4B + idx3Local, // 输入数据第4个8位Buf 256 * 4B + nkValueLocal, // next_k 暂存Buf 64 * 4B + topK, // topk数量 + s2SeqLen); // 输入元素总数 + } +private: + LocalTensor tmpIdxLocal; // filter阶段使用暂存index Buf topK * 4B + LocalTensor tmpValueLocal; // filter阶段使用暂存value Buf topK * 4B + LocalTensor histogramsLocal; // 直方图的临时Buf 256 * 4B + LocalTensor idx0Local; // 输入数据第1个8位Buf 256 * 4B + LocalTensor idx1Local; // 输入数据第2个8位Buf 256 * 4B + LocalTensor idx2Local; // 输入数据第3个8位Buf 256 * 4B + LocalTensor idx3Local; // 输入数据第4个8位Buf 256 * 4B + LocalTensor nkValueLocal; // next_k 暂存Buf 64 * 4B + LocalTensor outputValueLocal; // 输出value tensor + uint32_t topK; +}; + +template<> +class LITopk { +public: + __aicore__ inline uint32_t GetSharedTmpBufferSize() + { + // 2 * LICommon::Align(topK, (uint32_t)256):两块hisIndexLocal； + // 3 * 256：histogramsLocal idxHighLocal idxLowLocal；64：nkValueLocal + uint64_t bufferSize1 = (2 * LICommon::Align(topK, (uint32_t)256) + 3 * 256 + 64) * sizeof(uint32_t); + // LICommon::Align(topK, (uint32_t)256) + trunkLen：tmpIndexLocal + uint64_t bufferSize2 = (LICommon::Align(topK, (uint32_t)256) + trunkLen) * sizeof(uint16_t); + return bufferSize1 + bufferSize2; + } + + __aicore__ inline void Init(uint32_t topK, uint32_t trunkLen) + { + this->topK = topK; + this->trunkLen = trunkLen; + } + + __aicore__ inline void InitBuffers(LocalTensor& sharedTmpBuffer) + { + LocalTensor hisIndexLocal1 = sharedTmpBuffer[0]; + LocalTensor hisIndexLocal2 = hisIndexLocal1[LICommon::Align(topK, (uint32_t)256)]; + hisIndexLocal[0] = hisIndexLocal1; + hisIndexLocal[1] = hisIndexLocal2; + histogramsLocal = hisIndexLocal2[LICommon::Align(topK, (uint32_t)256)]; + idxHighLocal = histogramsLocal[256]; + idxLowLocal = idxHighLocal[256]; + nkValueLocal = idxLowLocal[256]; + LocalTensor tmpIndexLocalTmp = nkValueLocal[64]; + tmpIndexLocal = tmpIndexLocalTmp.template ReinterpretCast(); + } + + __aicore__ inline void operator()(LocalTensor& mrgValueLocal, LocalTensor& indicesOutLocal, + LocalTensor& hisValueLocal, uint32_t s2SeqLen, uint32_t loopIdx, + uint32_t s2LoopNum, bool returnValueFlag) + { + if (s2LoopNum == 1) { + if (returnValueFlag) { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, + mrgValueLocal, histogramsLocal, idxHighLocal, + idxLowLocal, nkValueLocal, topK, s2SeqLen); + } else { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, + mrgValueLocal, histogramsLocal, idxHighLocal, + idxLowLocal, nkValueLocal, topK, s2SeqLen); + } + PipeBarrier(); + Cast(indicesOutLocal, tmpIndexLocal, RoundMode::CAST_NONE, topK); + return; + } + + if (loopIdx == 0) { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, + mrgValueLocal, histogramsLocal, idxHighLocal, + idxLowLocal, nkValueLocal, topK, s2SeqLen); + PipeBarrier(); + Cast(hisIndexLocal[(loopIdx + 1) % 2], tmpIndexLocal, RoundMode::CAST_NONE, topK); + } else { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, + mrgValueLocal, histogramsLocal, idxHighLocal, + idxLowLocal, nkValueLocal, topK, s2SeqLen); + PipeBarrier(); + topkb16gather::LiTopKGatherVF(hisIndexLocal[(loopIdx + 1) % 2], hisValueLocal, mrgValueLocal, tmpIndexLocal, + hisIndexLocal[loopIdx % 2], topK, + loopIdx * trunkLen - LICommon::Align(topK, (uint32_t)256), s2SeqLen); + if (loopIdx == s2LoopNum - 1) { + PipeBarrier(); + AscendC::DataCopy(indicesOutLocal, + hisIndexLocal[(loopIdx + 1) % 2], + LICommon::Align(topK, (uint32_t)256)); + } + } + } +private: + LocalTensor hisIndexLocal[2]; // 每trunkLen长度的s2选出的topK个索引 + LocalTensor histogramsLocal; // 直方图的临时Buf 256 * 4B + LocalTensor idxHighLocal; // 输入数据高8位Buf 256 * 4B + LocalTensor idxLowLocal; // 输入数据低8位Buf 256 * 4B + LocalTensor nkValueLocal; // next_k 暂存Buf 64 * 4B + LocalTensor tmpIndexLocal; // 每trunkLen + topK的临时index + uint32_t topK = 512; + uint32_t trunkLen = 16384; +}; +} +#endif diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/vf/lightning_indexer_vector1.h b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/vf/lightning_indexer_vector1.h new file mode 100644 index 000000000..d623587f3 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/vf/lightning_indexer_vector1.h @@ -0,0 +1,830 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_vector1.h + * \brief + */ +#ifndef LIGHTNING_INDEXER_VECTOR1_H +#define LIGHTNING_INDEXER_VECTOR1_H + +#include "kernel_operator.h" + +namespace vector1 { + +template +struct FloatSortTraits; + +template +struct UIntSortTraits; + +// fp32 +template <> +struct FloatSortTraits { + using UInt = uint32_t; + static constexpr UInt ZERO = 0x00000000; + static constexpr UInt SIGN_MASK = 0x80000000; + static constexpr UInt NAN_MASK = 0x7FC00000; + static constexpr UInt ALL_ONE = 0xFFFFFFFF; +}; + +// bf16 +template <> +struct FloatSortTraits { + using UInt = uint16_t; + static constexpr UInt ZERO = 0x0000; + static constexpr UInt SIGN_MASK = 0x8000; + static constexpr UInt NAN_MASK = 0x7FC0; + static constexpr UInt ALL_ONE = 0xFFFF; +}; + + +template +struct FloatSortConstCtx { + using Traits = FloatSortTraits; + using UInt = typename Traits::UInt; + AscendC::MicroAPI::RegTensor zeros; + AscendC::MicroAPI::RegTensor all_one; + AscendC::MicroAPI::RegTensor signMask; + AscendC::MicroAPI::RegTensor nan; +}; + + +template +__simd_callee__ inline void InitFloatSortConstCtx(FloatSortConstCtx& ctx, AscendC::MicroAPI::MaskReg& maskAll) +{ + using Traits = FloatSortTraits; + AscendC::MicroAPI::Duplicate(ctx.zeros, Traits::ZERO, maskAll); + AscendC::MicroAPI::Duplicate(ctx.all_one, Traits::ALL_ONE, maskAll); + AscendC::MicroAPI::Duplicate(ctx.signMask, Traits::SIGN_MASK, maskAll); + AscendC::MicroAPI::Duplicate(ctx.nan, Traits::NAN_MASK, maskAll); +} + + +template +__simd_callee__ inline void FloatToSortableKey( + AscendC::MicroAPI::RegTensor::UInt>& + outKey, + AscendC::MicroAPI::RegTensor& inVal, + FloatSortConstCtx& ctx, + AscendC::MicroAPI::MaskReg& maskAll) +{ + using Traits = FloatSortTraits; + using UInt = typename Traits::UInt; + + AscendC::MicroAPI::RegTensor regTemp; + AscendC::MicroAPI::RegTensor regMask; + AscendC::MicroAPI::MaskReg regSelectNan; + AscendC::MicroAPI::MaskReg regSelectSign; + + auto& inBits = (AscendC::MicroAPI::RegTensor&)inVal; + + // 1. NaN check + AscendC::MicroAPI::Compare(regSelectNan, inBits, ctx.nan, maskAll); + + // 2. NaN -> ALL_ONE + AscendC::MicroAPI::Select(outKey, ctx.all_one, inBits, regSelectNan); + + // 3. sign bit + AscendC::MicroAPI::And(regTemp, outKey, ctx.signMask, maskAll); + + AscendC::MicroAPI::Compare(regSelectSign, regTemp, ctx.zeros, maskAll); + + // 4. xor mask + AscendC::MicroAPI::Select(regMask, ctx.all_one, ctx.signMask, regSelectSign); + AscendC::MicroAPI::Xor(outKey, outKey, regMask, maskAll); +} + +// uint16-bf16 +template <> +struct UIntSortTraits { + using UInt = uint16_t; + static constexpr UInt ZERO = 0x0000; + static constexpr UInt SIGN_MASK = 0x8000; + static constexpr UInt NAN_MASK = 0xFFC0; + static constexpr UInt ALL_ONE = 0xFFFF; +}; + +template +struct UIntSortConstCtx { + using Traits = UIntSortTraits; + using UInt = typename Traits::UInt; + AscendC::MicroAPI::RegTensor zeros; + AscendC::MicroAPI::RegTensor all_one; + AscendC::MicroAPI::RegTensor signMask; + AscendC::MicroAPI::RegTensor nan; +}; + +template +__simd_callee__ inline void InitUIntSortConstCtx(UIntSortConstCtx& ctx, AscendC::MicroAPI::MaskReg& maskAll) +{ + using Traits = UIntSortTraits; + AscendC::MicroAPI::Duplicate(ctx.zeros, Traits::ZERO, maskAll); + AscendC::MicroAPI::Duplicate(ctx.all_one, Traits::ALL_ONE, maskAll); + AscendC::MicroAPI::Duplicate(ctx.signMask, Traits::SIGN_MASK, maskAll); + AscendC::MicroAPI::Duplicate(ctx.nan, Traits::NAN_MASK, maskAll); +} + +template +__simd_callee__ inline void UIntToSortableKey(AscendC::MicroAPI::RegTensor& outKey, + AscendC::MicroAPI::RegTensor::UInt>& + inVal, + UIntSortConstCtx& ctx, + AscendC::MicroAPI::MaskReg& maskAll) +{ + using Traits = UIntSortTraits; + using UInt = typename Traits::UInt; + + AscendC::MicroAPI::RegTensor regTemp; + AscendC::MicroAPI::RegTensor regMask; + AscendC::MicroAPI::MaskReg regSelectZero; + AscendC::MicroAPI::MaskReg regSelectSign; + + auto& inBits = inVal; + + // 1. 0 check + AscendC::MicroAPI::Compare(regSelectZero, inBits, ctx.zeros, maskAll); + + // 2. 0 -> -NAN + AscendC::MicroAPI::Select((AscendC::MicroAPI::RegTensor&)outKey, ctx.nan, inBits, regSelectZero); + + // 3. sign bit + AscendC::MicroAPI::And(regTemp, (AscendC::MicroAPI::RegTensor&)outKey, ctx.signMask, maskAll); + + AscendC::MicroAPI::Compare(regSelectSign, regTemp, ctx.zeros, maskAll); + + // 4. xor mask + AscendC::MicroAPI::Select(regMask, ctx.signMask, ctx.all_one, regSelectSign); + AscendC::MicroAPI::Xor((AscendC::MicroAPI::RegTensor&)outKey, + (AscendC::MicroAPI::RegTensor&)outKey, regMask, maskAll); +} + +__aicore__ inline void UIntToFloatReturnValue(const LocalTensor &out_, + const LocalTensor &in, + const uint32_t topK) +{ + auto outBuf = (__local_mem__ bfloat16_t*)out_.GetPhyAddr(); + auto inBuf = (__local_mem__ uint16_t*)in.GetPhyAddr(); + + const uint16_t repeatSize16 = 128; + uint16_t topkLoopNum = (topK + repeatSize16 - 1) / repeatSize16; + + __VEC_SCOPE__ + { + AscendC::MicroAPI::RegTensor regIn; + AscendC::MicroAPI::RegTensor regOut; + AscendC::MicroAPI::MaskReg maskAllB16 = + AscendC::MicroAPI::CreateMask(); + + for (uint16_t i = 0; i < topkLoopNum; ++i) { + AscendC::MicroAPI::LoadAlign(regIn, inBuf + i * 128); + + UIntSortConstCtx uint16Ctx; + InitUIntSortConstCtx(uint16Ctx, maskAllB16); + + UIntToSortableKey(regOut, regIn, uint16Ctx, maskAllB16); + + AscendC::MicroAPI::StoreAlign( + outBuf + i * 128, + regOut, + maskAllB16); + } + } +} + +__aicore__ inline void UIntToFloatReturnValue(const LocalTensor &out_, + const LocalTensor &in, + const uint32_t topK) +{ + auto outBuf = (__local_mem__ half*)out_.GetPhyAddr(); + auto inBuf = (__local_mem__ uint16_t*)in.GetPhyAddr(); + + const uint16_t repeatSize16 = 128; + uint16_t topkLoopNum = (topK + repeatSize16 - 1) / repeatSize16; + + __VEC_SCOPE__ + { + AscendC::MicroAPI::RegTensor regIn; + AscendC::MicroAPI::RegTensor regOut; + AscendC::MicroAPI::RegTensor regOutHalf; + AscendC::MicroAPI::MaskReg maskAllB16 = + AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg maskAllHalf = + AscendC::MicroAPI::CreateMask(); + constexpr static MicroAPI::CastTrait castTraitBF16ToHalf = + {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_RINT}; + + for (uint16_t i = 0; i < topkLoopNum; ++i) { + AscendC::MicroAPI::LoadAlign(regIn, inBuf + i * repeatSize16); + + UIntSortConstCtx uint16Ctx; + InitUIntSortConstCtx(uint16Ctx, maskAllB16); + + UIntToSortableKey(regOut, regIn, uint16Ctx, maskAllB16); + + AscendC::MicroAPI::Cast(regOutHalf, regOut, maskAllB16); + + AscendC::MicroAPI::StoreAlign( + outBuf + i * repeatSize16, + regOutHalf, + maskAllHalf); + } + } +} + +template +__simd_callee__ inline void FloatX2ToSortableKey(AscendC::MicroAPI::RegTensor::UInt>& + outKey0, + AscendC::MicroAPI::RegTensor::UInt>& + outKey1, + AscendC::MicroAPI::RegTensor& inVal0, + AscendC::MicroAPI::RegTensor& inVal1, + FloatSortConstCtx& ctx, + AscendC::MicroAPI::MaskReg& maskAll) +{ + using Traits = FloatSortTraits; + using UInt = typename Traits::UInt; + + AscendC::MicroAPI::RegTensor regTemp[2]; + AscendC::MicroAPI::RegTensor regMask[2]; + AscendC::MicroAPI::MaskReg regSelectNan[2]; + AscendC::MicroAPI::MaskReg regSelectSign[2]; + + auto& inBits0 = (AscendC::MicroAPI::RegTensor&)inVal0; + auto& inBits1 = (AscendC::MicroAPI::RegTensor&)inVal1; + + // 1. NaN check + AscendC::MicroAPI::Compare(regSelectNan[0], inBits0, ctx.nan, maskAll); + AscendC::MicroAPI::Compare(regSelectNan[1], inBits1, ctx.nan, maskAll); + + // 2. NaN -> ALL_ONE + AscendC::MicroAPI::Select(outKey0, ctx.all_one, inBits0, regSelectNan[0]); + AscendC::MicroAPI::Select(outKey1, ctx.all_one, inBits1, regSelectNan[1]); + + // 3. sign bit + AscendC::MicroAPI::And(regTemp[0], outKey0, ctx.signMask, maskAll); + AscendC::MicroAPI::And(regTemp[1], outKey1, ctx.signMask, maskAll); + + AscendC::MicroAPI::Compare(regSelectSign[0], regTemp[0], ctx.zeros, maskAll); + AscendC::MicroAPI::Compare(regSelectSign[1], regTemp[1], ctx.zeros, maskAll); + + // 4. xor mask + AscendC::MicroAPI::Select(regMask[0], ctx.all_one, ctx.signMask, regSelectSign[0]); + AscendC::MicroAPI::Select(regMask[1], ctx.all_one, ctx.signMask, regSelectSign[1]); + AscendC::MicroAPI::Xor(outKey0, outKey0, regMask[0], maskAll); + AscendC::MicroAPI::Xor(outKey1, outKey1, regMask[1], maskAll); +} + + +template +__simd_callee__ inline void DuplicateZero(AscendC::MicroAPI::RegTensor (®Array)[N], + AscendC::MicroAPI::MaskReg& mask) +{ + static_assert(N <= 4, "N must be <= 4"); + // 不能用循环, 会导致fatal error: error in backend: Unsupported Inst must be hoisted. + if constexpr (N >= 1) { + AscendC::MicroAPI::Duplicate(regArray[0], static_cast(0), mask); + } + if constexpr (N >= 2) { + AscendC::MicroAPI::Duplicate(regArray[1], static_cast(0), mask); + } + if constexpr (N >= 3) { + AscendC::MicroAPI::Duplicate(regArray[2], static_cast(0), mask); + } + if constexpr (N >= 4) { + AscendC::MicroAPI::Duplicate(regArray[3], static_cast(0), mask); + } +} + + +template +__simd_callee__ inline void WeightedAccum(AscendC::MicroAPI::RegTensor (&accum)[N], + AscendC::MicroAPI::RegTensor (&input)[N], + AscendC::MicroAPI::RegTensor& weight, + AscendC::MicroAPI::MaskReg& mask) +{ + static_assert(N <= 2, "N must be <= 2"); + // ---- Relu block ---- + if constexpr (ApplyRelu) { + if constexpr (N >= 1) { + AscendC::MicroAPI::Relu(input[0], input[0], mask); + } + if constexpr (N >= 2) { + AscendC::MicroAPI::Relu(input[1], input[1], mask); + } + } + // ---- MulAdd block ---- + if constexpr (N >= 1) { + AscendC::MicroAPI::MulAddDst(accum[0], input[0], weight, mask); + } + if constexpr (N >= 2) { + AscendC::MicroAPI::MulAddDst(accum[1], input[1], weight, mask); + } +} + + +__simd_callee__ inline void BroadcastLane(AscendC::MicroAPI::RegTensor& dst, + AscendC::MicroAPI::RegTensor& src, + uint16_t laneIdx) +{ + AscendC::MicroAPI::RegTensor brcGatherIndex; + AscendC::MicroAPI::Duplicate(brcGatherIndex, laneIdx); + AscendC::MicroAPI::Gather(dst, src, brcGatherIndex); +} + +__simd_callee__ inline void BroadcastLane(AscendC::MicroAPI::RegTensor& dst, + __local_mem__ float* src, + uint16_t laneIdx) +{ + AscendC::MicroAPI::LoadAlign(dst, src + laneIdx); +} + +// float in uint16 out +__simd_vf__ inline void MulWeightAndReduceSum(__ubuf__ uint16_t* out_, + __ubuf__ float* qk_, + const uint32_t qkVLStride, + __ubuf__ float* weight_, + const int gSize) +{ + AscendC::MicroAPI::RegTensor regwBrc; + AscendC::MicroAPI::RegTensor regQK[2]; + AscendC::MicroAPI::RegTensor regW; + + AscendC::MicroAPI::RegTensor regSum0[2]; + AscendC::MicroAPI::RegTensor regSum1[2]; + AscendC::MicroAPI::MaskReg maskAllB32 = + AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg maskAllB16 = + AscendC::MicroAPI::CreateMask(); + + FloatSortConstCtx bf16Ctx; + InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = + {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = + {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_ROUND}; + + AscendC::MicroAPI::LoadAlign(regW, weight_); + DuplicateZero(regSum0, maskAllB32); + DuplicateZero(regSum1, maskAllB32); + + // unroll2 + for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i += 2) { + MicroAPI::LoadAlign(regQK[0], qk_ + 128 * i); // RowStride是128, 行都落在一个bank上 + MicroAPI::LoadAlign(regQK[1], qk_ + 128 * i + qkVLStride); + BroadcastLane(regwBrc, regW, i); + WeightedAccum(regSum0, regQK, regwBrc, maskAllB32); + + MicroAPI::LoadAlign(regQK[0], qk_ + 128 * i + 128); + MicroAPI::LoadAlign(regQK[1], qk_ + 128 * i + 128 + qkVLStride); + BroadcastLane(regwBrc, regW, i + 1); + WeightedAccum(regSum1, regQK, regwBrc, maskAllB32); + } + + AscendC::MicroAPI::Add(regSum0[0], regSum0[0], regSum1[0], maskAllB32); + AscendC::MicroAPI::Add(regSum0[1], regSum0[1], regSum1[1], maskAllB32); + + AscendC::MicroAPI::RegTensor regSumBF16; + // interleave cast ==> regSum[1] high regSum[0] low + AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + AscendC::MicroAPI::Cast(regSumBF16, regSum0[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16, regSum0[0], maskAllB32); + + AscendC::MicroAPI::RegTensor regOut; + FloatToSortableKey(regOut, regSumBF16, bf16Ctx, maskAllB16); + // normal store + AscendC::MicroAPI::StoreAlign(out_, regOut, maskAllB16); +} + +// float in uint16 out +__simd_vf__ inline void MulWeightAndReduceSum(__ubuf__ uint16_t* out_, + __ubuf__ float* qk_, + const uint32_t qkVLStride, + __ubuf__ bfloat16_t* weight_, + const int gSize) +{ + AscendC::MicroAPI::RegTensor regwBrc; + AscendC::MicroAPI::RegTensor regQK[2]; + AscendC::MicroAPI::RegTensor regWBF16; + AscendC::MicroAPI::RegTensor regW; + + AscendC::MicroAPI::RegTensor regSum0[2]; + AscendC::MicroAPI::RegTensor regSum1[2]; + AscendC::MicroAPI::MaskReg maskAllB32 = + AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg maskAllB16 = + AscendC::MicroAPI::CreateMask(); + + FloatSortConstCtx bf16Ctx; + InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitBF16ToFP32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::UNKNOWN}; + + AscendC::MicroAPI::LoadAlign(regWBF16, weight_); + AscendC::MicroAPI::Cast(regW, regWBF16, maskAllB16); + + DuplicateZero(regSum0, maskAllB32); + DuplicateZero(regSum1, maskAllB32); + + // unroll2 + for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i += 2) { + MicroAPI::LoadAlign(regQK[0], qk_ + 128 * i); // RowStride是128, 行都落在一个bank上 + MicroAPI::LoadAlign(regQK[1], qk_ + 128 * i + qkVLStride); + BroadcastLane(regwBrc, regW, i); + WeightedAccum(regSum0, regQK, regwBrc, maskAllB32); + + MicroAPI::LoadAlign(regQK[0], qk_ + 128 * i + 128); + MicroAPI::LoadAlign(regQK[1], qk_ + 128 * i + 128 + qkVLStride); + BroadcastLane(regwBrc, regW, i + 1); + WeightedAccum(regSum1, regQK, regwBrc, maskAllB32); + } + + AscendC::MicroAPI::Add(regSum0[0], regSum0[0], regSum1[0], maskAllB32); + AscendC::MicroAPI::Add(regSum0[1], regSum0[1], regSum1[1], maskAllB32); + + AscendC::MicroAPI::RegTensor regSumBF16; + // interleave cast ==> regSum[1] high regSum[0] low + AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + AscendC::MicroAPI::Cast(regSumBF16, regSum0[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16, regSum0[0], maskAllB32); + + AscendC::MicroAPI::RegTensor regOut; + FloatToSortableKey(regOut, regSumBF16, bf16Ctx, maskAllB16); + // normal store + AscendC::MicroAPI::StoreAlign(out_, regOut, maskAllB16); +} + +// float in uint16 out +__simd_vf__ inline void MulWeightAndReduceSum(__ubuf__ uint16_t* out_, + __ubuf__ float* qk_, + const uint32_t qkVLStride, + __ubuf__ half* weight_, + const int gSize) +{ + AscendC::MicroAPI::RegTensor regwBrc; + AscendC::MicroAPI::RegTensor regQK[2]; + AscendC::MicroAPI::RegTensor regW; + AscendC::MicroAPI::RegTensor regWFP16; + AscendC::MicroAPI::RegTensor regSum0[2]; + AscendC::MicroAPI::RegTensor regSum1[2]; + AscendC::MicroAPI::MaskReg maskAllB32 = + AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg maskAllB16 = + AscendC::MicroAPI::CreateMask(); + + FloatSortConstCtx bf16Ctx; + InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitFP16ToFP32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::UNKNOWN}; + + AscendC::MicroAPI::LoadAlign(regWFP16, weight_); + AscendC::MicroAPI::Cast(regW, regWFP16, maskAllB16); + + DuplicateZero(regSum0, maskAllB32); + DuplicateZero(regSum1, maskAllB32); + + // unroll2 + for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i += 2) { + MicroAPI::LoadAlign(regQK[0], qk_ + 128 * i); // RowStride是128, 行都落在一个bank上 + MicroAPI::LoadAlign(regQK[1], qk_ + 128 * i + qkVLStride); + BroadcastLane(regwBrc, regW, i); + WeightedAccum(regSum0, regQK, regwBrc, maskAllB32); + + MicroAPI::LoadAlign(regQK[0], qk_ + 128 * i + 128); + MicroAPI::LoadAlign(regQK[1], qk_ + 128 * i + 128 + qkVLStride); + BroadcastLane(regwBrc, regW, i + 1); + WeightedAccum(regSum1, regQK, regwBrc, maskAllB32); + } + + AscendC::MicroAPI::Add(regSum0[0], regSum0[0], regSum1[0], maskAllB32); + AscendC::MicroAPI::Add(regSum0[1], regSum0[1], regSum1[1], maskAllB32); + + AscendC::MicroAPI::RegTensor regSumBF16; + // interleave cast ==> regSum[1] high regSum[0] low + AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + AscendC::MicroAPI::Cast(regSumBF16, regSum0[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16, regSum0[0], maskAllB32); + + AscendC::MicroAPI::RegTensor regOut; + FloatToSortableKey(regOut, regSumBF16, bf16Ctx, maskAllB16); + // normal store + AscendC::MicroAPI::StoreAlign(out_, regOut, maskAllB16); +} + +// 计算S1=2 +// float in uint16 out +__simd_vf__ inline void MulWeightAndReduceSum2(__ubuf__ uint16_t* out0_, + __ubuf__ uint16_t* out1_, + uint32_t outStride, + __ubuf__ float* qk0_, + __ubuf__ float* qk1_, + uint32_t qkVLStride, + uint32_t qkStride, + __ubuf__ float* weight0_, + __ubuf__ float* weight1_, + uint32_t weightStride, + __ubuf__ float* weightFloat_, + const int gSize) +{ + AscendC::MicroAPI::RegTensor regwBrc[2]; + AscendC::MicroAPI::RegTensor regQK0[2]; + AscendC::MicroAPI::RegTensor regQK1[2]; + AscendC::MicroAPI::RegTensor regW[2]; + + AscendC::MicroAPI::RegTensor regSum0[2]; + AscendC::MicroAPI::RegTensor regSum1[2]; + AscendC::MicroAPI::MaskReg maskAllB32 = + AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg maskAllB16 = + AscendC::MicroAPI::CreateMask(); + + FloatSortConstCtx bf16Ctx; + InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_ROUND}; + + AscendC::MicroAPI::LoadAlign(regW[0], weight0_); + AscendC::MicroAPI::LoadAlign(regW[1], weight1_); + // regW[0]与weight1混合使用 + AscendC::MicroAPI::StoreAlign(weight1_, regW[1], maskAllB32); + AscendC::MicroAPI::LocalMemBar(); + DuplicateZero(regSum0, maskAllB32); + DuplicateZero(regSum1, maskAllB32); + + for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i++) { + MicroAPI::LoadAlign(regQK0[0], qk0_ + 128 * i); + MicroAPI::LoadAlign(regQK0[1], qk0_ + 128 * i + qkVLStride); + MicroAPI::LoadAlign(regQK1[0], qk1_ + 128 * i); + MicroAPI::LoadAlign(regQK1[1], qk1_ + 128 * i + qkVLStride); + // 混合使用对整体性能更好 + BroadcastLane(regwBrc[0], regW[0], i); + // Weight无bank冲突，用LoadAlign来提取weight标量 + BroadcastLane(regwBrc[1], weight1_, i); + AscendC::MicroAPI::Relu(regQK0[0], regQK0[0], maskAllB32); + AscendC::MicroAPI::Relu(regQK0[1], regQK0[1], maskAllB32); + AscendC::MicroAPI::Relu(regQK1[0], regQK1[0], maskAllB32); + AscendC::MicroAPI::Relu(regQK1[1], regQK1[1], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum0[0], regQK0[0], regwBrc[0], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum0[1], regQK0[1], regwBrc[0], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum1[0], regQK1[0], regwBrc[1], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum1[1], regQK1[1], regwBrc[1], maskAllB32); + } + + // Convert to bfloat16 and store output channel + AscendC::MicroAPI::RegTensor regSumBF16[2]; + AscendC::MicroAPI::RegTensor regOut[2]; + AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + AscendC::MicroAPI::DeInterleave(regSum1[0], regSum1[1], regSum1[0], regSum1[1]); + AscendC::MicroAPI::Cast(regSumBF16[0], regSum0[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[1], regSum1[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[0], regSum0[0], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[1], regSum1[0], maskAllB32); + + FloatX2ToSortableKey(regOut[0], regOut[1], regSumBF16[0], regSumBF16[1], bf16Ctx, maskAllB16); + AscendC::MicroAPI::StoreAlign(out0_, regOut[0], maskAllB16); + AscendC::MicroAPI::StoreAlign(out1_, regOut[1], maskAllB16); +} + +// 计算S1=2 +// float in uint16 out +__simd_vf__ inline void MulWeightAndReduceSum2(__ubuf__ uint16_t* out0_, + __ubuf__ uint16_t* out1_, + uint32_t outStride, + __ubuf__ float* qk0_, + __ubuf__ float* qk1_, + uint32_t qkVLStride, + uint32_t qkStride, + __ubuf__ bfloat16_t* weight0_, + __ubuf__ bfloat16_t* weight1_, + uint32_t weightStride, + __ubuf__ float* weightFloat_, + const int gSize) +{ + AscendC::MicroAPI::RegTensor regwBrc[2]; + AscendC::MicroAPI::RegTensor regQK0[2]; + AscendC::MicroAPI::RegTensor regQK1[2]; + AscendC::MicroAPI::RegTensor regW[2]; + AscendC::MicroAPI::RegTensor regWBF16[2]; + + AscendC::MicroAPI::RegTensor regSum0[2]; + AscendC::MicroAPI::RegTensor regSum1[2]; + AscendC::MicroAPI::MaskReg maskAllB32 = + AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg maskAllB16 = + AscendC::MicroAPI::CreateMask(); + + FloatSortConstCtx bf16Ctx; + InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitBF16ToFP32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::UNKNOWN}; + + AscendC::MicroAPI::LoadAlign(regWBF16[0], weight0_); + AscendC::MicroAPI::LoadAlign(regWBF16[1], weight1_); + AscendC::MicroAPI::Cast(regW[0], regWBF16[0], maskAllB16); + AscendC::MicroAPI::Cast(regW[1], regWBF16[1], maskAllB16); + + // regW[0]与weight1混合使用 + AscendC::MicroAPI::StoreAlign(weightFloat_, regW[1], maskAllB32); + AscendC::MicroAPI::LocalMemBar(); + DuplicateZero(regSum0, maskAllB32); + DuplicateZero(regSum1, maskAllB32); + + for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i++) { + MicroAPI::LoadAlign(regQK0[0], qk0_ + 128 * i); + MicroAPI::LoadAlign(regQK0[1], qk0_ + 128 * i + qkVLStride); + MicroAPI::LoadAlign(regQK1[0], qk1_ + 128 * i); + MicroAPI::LoadAlign(regQK1[1], qk1_ + 128 * i + qkVLStride); + // 混合使用对整体性能更好 + BroadcastLane(regwBrc[0], regW[0], i); + // Weight无bank冲突，用LoadAlign来提取weight标量 + BroadcastLane(regwBrc[1], weightFloat_, i); + AscendC::MicroAPI::Relu(regQK0[0], regQK0[0], maskAllB32); + AscendC::MicroAPI::Relu(regQK0[1], regQK0[1], maskAllB32); + AscendC::MicroAPI::Relu(regQK1[0], regQK1[0], maskAllB32); + AscendC::MicroAPI::Relu(regQK1[1], regQK1[1], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum0[0], regQK0[0], regwBrc[0], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum0[1], regQK0[1], regwBrc[0], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum1[0], regQK1[0], regwBrc[1], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum1[1], regQK1[1], regwBrc[1], maskAllB32); + } + + // Convert to bfloat16 and store output channel + AscendC::MicroAPI::RegTensor regSumBF16[2]; + AscendC::MicroAPI::RegTensor regOut[2]; + AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + AscendC::MicroAPI::DeInterleave(regSum1[0], regSum1[1], regSum1[0], regSum1[1]); + AscendC::MicroAPI::Cast(regSumBF16[0], regSum0[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[1], regSum1[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[0], regSum0[0], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[1], regSum1[0], maskAllB32); + + FloatX2ToSortableKey(regOut[0], regOut[1], regSumBF16[0], regSumBF16[1], bf16Ctx, maskAllB16); + AscendC::MicroAPI::StoreAlign(out0_, regOut[0], maskAllB16); + AscendC::MicroAPI::StoreAlign(out1_, regOut[1], maskAllB16); +} + +// 计算S1=2 +// float in uint16 out +__simd_vf__ inline void MulWeightAndReduceSum2(__ubuf__ uint16_t* out0_, + __ubuf__ uint16_t* out1_, + uint32_t outStride, + __ubuf__ float* qk0_, + __ubuf__ float* qk1_, + uint32_t qkVLStride, + uint32_t qkStride, + __ubuf__ half* weight0_, + __ubuf__ half* weight1_, + uint32_t weightStride, + __ubuf__ float* weightFloat_, + const int gSize) +{ + AscendC::MicroAPI::RegTensor regwBrc[2]; + AscendC::MicroAPI::RegTensor regQK0[2]; + AscendC::MicroAPI::RegTensor regQK1[2]; + AscendC::MicroAPI::RegTensor regW[2]; + AscendC::MicroAPI::RegTensor regWFP16[2]; + + AscendC::MicroAPI::RegTensor regSum0[2]; + AscendC::MicroAPI::RegTensor regSum1[2]; + AscendC::MicroAPI::MaskReg maskAllB32 = + AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg maskAllB16 = + AscendC::MicroAPI::CreateMask(); + + FloatSortConstCtx bf16Ctx; + InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitFP16ToFP32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::UNKNOWN}; + + AscendC::MicroAPI::LoadAlign(regWFP16[0], weight0_); + AscendC::MicroAPI::LoadAlign(regWFP16[1], weight1_); + AscendC::MicroAPI::Cast(regW[0], regWFP16[0], maskAllB16); + AscendC::MicroAPI::Cast(regW[1], regWFP16[1], maskAllB16); + + // regW[0]与weight1混合使用 + AscendC::MicroAPI::StoreAlign(weightFloat_, regW[1], maskAllB32); + AscendC::MicroAPI::LocalMemBar(); + DuplicateZero(regSum0, maskAllB32); + DuplicateZero(regSum1, maskAllB32); + + for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i++) { + MicroAPI::LoadAlign(regQK0[0], qk0_ + 128 * i); + MicroAPI::LoadAlign(regQK0[1], qk0_ + 128 * i + qkVLStride); + MicroAPI::LoadAlign(regQK1[0], qk1_ + 128 * i); + MicroAPI::LoadAlign(regQK1[1], qk1_ + 128 * i + qkVLStride); + // 混合使用对整体性能更好 + BroadcastLane(regwBrc[0], regW[0], i); + // Weight无bank冲突，用LoadAlign来提取weight标量 + BroadcastLane(regwBrc[1], weightFloat_, i); + AscendC::MicroAPI::Relu(regQK0[0], regQK0[0], maskAllB32); + AscendC::MicroAPI::Relu(regQK0[1], regQK0[1], maskAllB32); + AscendC::MicroAPI::Relu(regQK1[0], regQK1[0], maskAllB32); + AscendC::MicroAPI::Relu(regQK1[1], regQK1[1], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum0[0], regQK0[0], regwBrc[0], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum0[1], regQK0[1], regwBrc[0], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum1[0], regQK1[0], regwBrc[1], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum1[1], regQK1[1], regwBrc[1], maskAllB32); + } + + // Convert to bfloat16 and store output channel + AscendC::MicroAPI::RegTensor regSumBF16[2]; + AscendC::MicroAPI::RegTensor regOut[2]; + AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + AscendC::MicroAPI::DeInterleave(regSum1[0], regSum1[1], regSum1[0], regSum1[1]); + AscendC::MicroAPI::Cast(regSumBF16[0], regSum0[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[1], regSum1[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[0], regSum0[0], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[1], regSum1[0], maskAllB32); + + FloatX2ToSortableKey(regOut[0], regOut[1], regSumBF16[0], regSumBF16[1], bf16Ctx, maskAllB16); + AscendC::MicroAPI::StoreAlign(out0_, regOut[0], maskAllB16); + AscendC::MicroAPI::StoreAlign(out1_, regOut[1], maskAllB16); +} + +template +__aicore__ inline void BatchMulWeightAndReduceSum(const LocalTensor &out_, // out [S2Base] [128 ] + uint32_t outStride, + const LocalTensor &qk_, // q*k^t [G, S2Base] [64 128] + uint32_t qkVLStride, + uint32_t qkStride, + const LocalTensor &weight_, // w [G] [64 ] + uint32_t weightStride, + const LocalTensor &weightFloat_, + const int gSize, // G 64 + const int batch) +{ + // 暂只支持这两种情况, 后续改成循环 + if (batch != 2 && batch != 1) { + return; + } + auto weight = (__ubuf__ W_T *)weight_.GetPhyAddr(); + auto weightFloat = (__ubuf__ float *)weightFloat_.GetPhyAddr(); + auto qk = (__ubuf__ float *)qk_.GetPhyAddr(); + auto out = (__ubuf__ uint16_t *)out_.GetPhyAddr(); + if (batch == 2) { + auto weight1 = weight + weightStride; + auto qk1 = qk + qkStride; + auto out1 = out + outStride; + MulWeightAndReduceSum2(out, out1, outStride, + qk, qk1, qkVLStride, qkStride, + weight, weight1, weightStride, weightFloat, + gSize); + } else { + MulWeightAndReduceSum(out, qk, qkVLStride, weight, gSize); + } +} + +} + +#endif diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/vf/vf_topk.h b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/vf/vf_topk.h new file mode 100644 index 000000000..4a777b7ca --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/vf/vf_topk.h @@ -0,0 +1,736 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. +  */ + +/*! +* \file vf_top_k.h +* \brief +*/ + +#ifndef VF_TOP_K_H +#define VF_TOP_K_H + +namespace topkb32 { +template +__simd_vf__ void HistogramsFirstVFImpl(__ubuf__ uint32_t* histogramsBuf, + __ubuf__ uint32_t* inputBuf, + uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图cout0 0-127 cout1 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + // 32bit 高16bit + MicroAPI::RegTensor vreg0U16; + // 32bit 低16bit + MicroAPI::RegTensor vreg1U16; + MicroAPI::RegTensor vreg2U16; + MicroAPI::RegTensor vreg3U16; + + MicroAPI::RegTensor vreg0; + MicroAPI::RegTensor vreg1; + MicroAPI::RegTensor vreg2; + MicroAPI::RegTensor vreg3; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vreg1U16, vreg0U16, inputBuf + i * 256); + MicroAPI::LoadAlign( + vreg3U16, vreg2U16, inputBuf + (i * 256) + 128); + + MicroAPI::DeInterleave(vreg1, vreg0, + (MicroAPI::RegTensor&)vreg0U16, + (MicroAPI::RegTensor&)vreg2U16); + + MicroAPI::Histograms(cout0, vreg0, pregB8); + MicroAPI::Histograms(cout1, vreg0, pregB8); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign( + histogramsBuf, cout0U32Even, cout0U32Odd, pregB32); + MicroAPI::StoreAlign( + histogramsBuf + 128, cout1U32Even, cout1U32Odd, pregB32); +} + +__simd_vf__ void FindFirstTargetBinVFImpl(__ubuf__ uint32_t* idx0Buf, + __ubuf__ uint32_t* nkValueBuf, __ubuf__ uint32_t* + histogramsBuf, uint32_t bottomK) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx0; + + MicroAPI::RegTensor btmK; + MicroAPI::Duplicate(btmK, bottomK); + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdx0; + + MicroAPI::MaskReg pregGE = MicroAPI::CreateMask(); + + MicroAPI::Arange(idxC, i * 64); + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + MicroAPI::Compare(pregGE, cout, btmK, pregB32); + MicroAPI::Squeeze( + sqzIdx0, (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idx0Buf, sqzIdx0, alignIdx0); + } + MicroAPI::StoreUnAlignPost(idx0Buf, alignIdx0); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idx0; + MicroAPI::LoadAlign(idx0, idx0Buf); + + MicroAPI::RegTensor idxAll1; + MicroAPI::RegTensor idxPrev0; + MicroAPI::RegTensor prevBinValue; + MicroAPI::Duplicate(idxAll1, 1); + + MicroAPI::RegTensor zeroAll; + MicroAPI::Duplicate(zeroAll, 0); + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + MicroAPI::Compare(preg0, idx0, zeroAll, pregB32); + MicroAPI::Sub(idxPrev0, idx0, (MicroAPI::RegTensor&)idxAll1, pregB32); + MicroAPI::ShiftRights(idxPrev0, idxPrev0, (int16_t)24, pregB32); + + MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev0, pregB32); + MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, preg0); + + MicroAPI::RegTensor nextK; + MicroAPI::Sub(nextK, btmK, prevBinValue, pregB32); + MicroAPI::StoreAlign(nkValueBuf, nextK, pregB32); +} + +template +__simd_vf__ void HistogramsSecondVFImpl(__ubuf__ uint32_t* histogramsBuf, + __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* idx0Buf, + uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图0-127 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + MicroAPI::RegTensor idx0; + // 0x000000fc -> 0xfcfcfcfc + MicroAPI::LoadAlign(idx0, idx0Buf); + + MicroAPI::RegTensor vreg0U16; + MicroAPI::RegTensor vreg1U16; + MicroAPI::RegTensor vreg2U16; + MicroAPI::RegTensor vreg3U16; + + MicroAPI::RegTensor vreg0; + MicroAPI::RegTensor vreg1; + MicroAPI::RegTensor vreg2; + MicroAPI::RegTensor vreg3; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vreg1U16, + vreg0U16, inputBuf + i * 256); + MicroAPI::LoadAlign(vreg3U16, + vreg2U16, inputBuf + (i * 256) + 128); + + MicroAPI::DeInterleave(vreg1, vreg0, + (MicroAPI::RegTensor&)vreg0U16, + (MicroAPI::RegTensor&)vreg2U16); + + MicroAPI::MaskReg pregEQ = MicroAPI::CreateMask(); + MicroAPI::Compare(pregEQ, vreg0, (MicroAPI::RegTensor&)idx0, pregB8); + + MicroAPI::Histograms(cout0, vreg1, pregEQ); + MicroAPI::Histograms(cout1, vreg1, pregEQ); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, + cout0U32Even, cout0U32Odd, pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, + cout1U32Even, cout1U32Odd, pregB32); +} + +// kValue新的bottomK +__simd_vf__ void FindSecondTargetBinVFImpl(__ubuf__ uint32_t* idx1Buf, + __ubuf__ uint32_t* nkValueBuf, __ubuf__ uint32_t* kValue, + __ubuf__ uint32_t* histogramsBuf) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx1; + + MicroAPI::RegTensor btmK1; + MicroAPI::LoadAlign(btmK1, kValue); + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdx1; + + MicroAPI::MaskReg pregGE = MicroAPI::CreateMask(); + + MicroAPI::Arange(idxC, i * 64); + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + MicroAPI::Compare(pregGE, cout, btmK1, pregB32); + MicroAPI::Squeeze(sqzIdx1, + (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idx1Buf, sqzIdx1, alignIdx1); + } + MicroAPI::StoreUnAlignPost(idx1Buf, alignIdx1); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idx1; + MicroAPI::LoadAlign(idx1, idx1Buf); + + MicroAPI::RegTensor idxAll1; + MicroAPI::RegTensor idxPrev1; + MicroAPI::RegTensor prevBinValue; + MicroAPI::Duplicate(idxAll1, 1); + + MicroAPI::RegTensor zeroAll; + MicroAPI::Duplicate(zeroAll, 0); + + MicroAPI::MaskReg preg1 = MicroAPI::CreateMask(); + MicroAPI::Compare(preg1, idx1, zeroAll, pregB32); + MicroAPI::Sub(idxPrev1, idx1, (MicroAPI::RegTensor&)idxAll1, pregB32); + MicroAPI::ShiftRights(idxPrev1, idxPrev1, (int16_t)24, pregB32); + + MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev1, pregB32); + MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, preg1); + + MicroAPI::RegTensor nextK; + MicroAPI::Sub(nextK, btmK1, prevBinValue, pregB32); + MicroAPI::StoreAlign(nkValueBuf, nextK, pregB32); +} + +template +__simd_vf__ void HistogramsThirdVFImpl(__ubuf__ uint32_t* histogramsBuf, + __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* idx0Buf, + __ubuf__ uint32_t* idx1Buf, uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图0-127 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + MicroAPI::RegTensor idx0; + MicroAPI::RegTensor idx1; + // 0x000000fc -> 0xfcfcfcfc + MicroAPI::LoadAlign(idx0, idx0Buf); + MicroAPI::LoadAlign(idx1, idx1Buf); + + MicroAPI::RegTensor vreg0U16; + MicroAPI::RegTensor vreg1U16; + MicroAPI::RegTensor vreg2U16; + MicroAPI::RegTensor vreg3U16; + + MicroAPI::RegTensor vreg0; + MicroAPI::RegTensor vreg1; + MicroAPI::RegTensor vreg2; + MicroAPI::RegTensor vreg3; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vreg1U16, + vreg0U16, inputBuf + i * 256); + MicroAPI::LoadAlign(vreg3U16, + vreg2U16, inputBuf + (i * 256) + 128); + + MicroAPI::DeInterleave(vreg1, vreg0, (MicroAPI::RegTensor&)vreg0U16, + (MicroAPI::RegTensor&)vreg2U16); + MicroAPI::DeInterleave(vreg3, vreg2, (MicroAPI::RegTensor&)vreg1U16, + (MicroAPI::RegTensor&)vreg3U16); + + MicroAPI::MaskReg pregEQ0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregEQ1 = MicroAPI::CreateMask(); + MicroAPI::Compare(pregEQ0, vreg0, (MicroAPI::RegTensor&)idx0, pregB8); + MicroAPI::Compare(pregEQ1, vreg1, (MicroAPI::RegTensor&)idx1, pregB8); + + MicroAPI::MaskReg pregEQ = MicroAPI::CreateMask(); + MicroAPI::And(pregEQ, pregEQ0, pregEQ1, pregB8); + + MicroAPI::Histograms(cout0, vreg2, pregEQ); + MicroAPI::Histograms(cout1, vreg2, pregEQ); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, + cout0U32Even, cout0U32Odd, pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, + cout1U32Even, cout1U32Odd, pregB32); +} + +__simd_vf__ void FindThirdTargetBinVFImpl(__ubuf__ uint32_t* idx2Buf, + __ubuf__ uint32_t* nkValueBuf, __ubuf__ uint32_t* kValue, + __ubuf__ uint32_t* histogramsBuf) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx2; + + MicroAPI::RegTensor btmK2; + MicroAPI::LoadAlign(btmK2, kValue); + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdx2; + + MicroAPI::MaskReg pregGE = MicroAPI::CreateMask(); + + MicroAPI::Arange(idxC, i * 64); + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + MicroAPI::Compare(pregGE, cout, btmK2, pregB32); + MicroAPI::Squeeze( + sqzIdx2, + (MicroAPI::RegTensor&)idxC, + pregGE); + MicroAPI::StoreUnAlign(idx2Buf, sqzIdx2, alignIdx2); + } + MicroAPI::StoreUnAlignPost(idx2Buf, alignIdx2); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idx2; + MicroAPI::LoadAlign(idx2, idx2Buf); + + MicroAPI::RegTensor idxAll1; + MicroAPI::RegTensor idxPrev2; + MicroAPI::RegTensor prevBinValue; + MicroAPI::Duplicate(idxAll1, 1); + + MicroAPI::RegTensor zeroAll; + MicroAPI::Duplicate(zeroAll, 0); + + MicroAPI::MaskReg preg2 = MicroAPI::CreateMask(); + MicroAPI::Compare(preg2, idx2, zeroAll, pregB32); + MicroAPI::Sub(idxPrev2, idx2, (MicroAPI::RegTensor&)idxAll1, pregB32); + MicroAPI::ShiftRights(idxPrev2, idxPrev2, (int16_t)24, pregB32); + + MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev2, pregB32); + MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, preg2); + + MicroAPI::RegTensor nextK; + MicroAPI::Sub(nextK, btmK2, prevBinValue, pregB32); + MicroAPI::StoreAlign(nkValueBuf, nextK, pregB32); +} + +template +__simd_vf__ void HistogramsLastVFImpl(__ubuf__ uint32_t* histogramsBuf, + __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* idx0Buf, + __ubuf__ uint32_t* idx1Buf, __ubuf__ uint32_t* idx2Buf, + uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图0-127 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + MicroAPI::RegTensor idx0; + MicroAPI::RegTensor idx1; + MicroAPI::RegTensor idx2; + // 0x000000fc -> 0xfcfcfcfc + MicroAPI::LoadAlign(idx0, idx0Buf); + MicroAPI::LoadAlign(idx1, idx1Buf); + MicroAPI::LoadAlign(idx2, idx2Buf); + + MicroAPI::RegTensor vreg0U16; + MicroAPI::RegTensor vreg1U16; + MicroAPI::RegTensor vreg2U16; + MicroAPI::RegTensor vreg3U16; + + MicroAPI::RegTensor vreg0; + MicroAPI::RegTensor vreg1; + MicroAPI::RegTensor vreg2; + MicroAPI::RegTensor vreg3; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vreg1U16, vreg0U16, inputBuf + i * 256); + MicroAPI::LoadAlign(vreg3U16, + vreg2U16, inputBuf + (i * 256) + 128); + + MicroAPI::DeInterleave(vreg1, vreg0, + (MicroAPI::RegTensor&)vreg0U16, + (MicroAPI::RegTensor&)vreg2U16); + MicroAPI::DeInterleave(vreg3, vreg2, + (MicroAPI::RegTensor&)vreg1U16, + (MicroAPI::RegTensor&)vreg3U16); + + MicroAPI::MaskReg pregEQ0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregEQ1 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregEQ2 = MicroAPI::CreateMask(); + MicroAPI::Compare(pregEQ0, vreg0, (MicroAPI::RegTensor&)idx0, pregB8); + MicroAPI::Compare(pregEQ1, vreg1, (MicroAPI::RegTensor&)idx1, pregB8); + MicroAPI::Compare(pregEQ2, vreg2, (MicroAPI::RegTensor&)idx2, pregB8); + + MicroAPI::MaskReg pregEQ0And1 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregEQAll = MicroAPI::CreateMask(); + MicroAPI::And(pregEQ0And1, pregEQ0, pregEQ1, pregB8); + MicroAPI::And(pregEQAll, pregEQ0And1, pregEQ2, pregB8); + + MicroAPI::Histograms(cout0, vreg3, pregEQAll); + MicroAPI::Histograms(cout1, vreg3, pregEQAll); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, + cout0U32Even, cout0U32Odd, pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, + cout1U32Even, cout1U32Odd, pregB32); +} + +__simd_vf__ void FindKthVFImpl(__ubuf__ uint32_t* kValue, + __ubuf__ uint32_t* histogramsBuf, __ubuf__ uint32_t* idx0Buf, + __ubuf__ uint32_t* idx1Buf, __ubuf__ uint32_t* idx2Buf, + __ubuf__ uint32_t* idx3Buf) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx3; + + MicroAPI::RegTensor btmK3; + MicroAPI::LoadAlign(btmK3, kValue); + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdx3; + + MicroAPI::MaskReg pregGE = MicroAPI::CreateMask(); + + MicroAPI::Arange(idxC, i * 64); + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + MicroAPI::Compare(pregGE, cout, btmK3, pregB32); + MicroAPI::Squeeze(sqzIdx3, + (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idx3Buf, sqzIdx3, alignIdx3); + } + MicroAPI::StoreUnAlignPost(idx3Buf, alignIdx3); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idx0; + MicroAPI::RegTensor idx1; + MicroAPI::RegTensor idx2; + MicroAPI::RegTensor idx3; + MicroAPI::LoadAlign(idx0, idx0Buf); + MicroAPI::LoadAlign(idx1, idx1Buf); + MicroAPI::LoadAlign(idx2, idx2Buf); + MicroAPI::LoadAlign(idx3, idx3Buf); + + MicroAPI::ShiftLefts(idx0, idx0, (int16_t)24, pregB32); + MicroAPI::ShiftLefts(idx1, idx1, (int16_t)16, pregB32); + MicroAPI::ShiftLefts(idx2, idx2, (int16_t)8, pregB32); + + // ADD + MicroAPI::Add(idx0, idx0, idx1, pregB32); + MicroAPI::Add(idx0, idx0, idx2, pregB32); + MicroAPI::Add(idx0, idx0, idx3, pregB32); + + MicroAPI::StoreAlign(kValue, idx0, pregB32); +} + +__simd_vf__ void FindIdxGTOutputVFImpl(__ubuf__ uint32_t* outputIdxBuf, + __ubuf__ uint32_t* inputBuf, uint32_t beginIdx, + __ubuf__ uint32_t* kValue, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::RegTensor idxC; + MicroAPI::Arange(idxC, beginIdx + i * 64); + + MicroAPI::LoadAlign(vregInput, inputBuf + i * 64); + + MicroAPI::MaskReg poutGT = MicroAPI::CreateMask(); + + MicroAPI::RegTensor sqzIdxOut; + MicroAPI::Compare(poutGT, vregInput, kthValue, pregB32); + + MicroAPI::Squeeze(sqzIdxOut, + (MicroAPI::RegTensor&)idxC, poutGT); + MicroAPI::StoreUnAlign(outputIdxBuf, sqzIdxOut, alignIdx); + } + MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx); +} + +__simd_vf__ void FindIdxEQOutputVFImpl(__ubuf__ uint32_t* outputIdxBuf, + __ubuf__ uint32_t* inputBuf, uint32_t beginIdx, + __ubuf__ uint32_t* kValue) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::UnalignRegForStore alignIdx; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + + MicroAPI::RegTensor idxC; + MicroAPI::Arange(idxC, beginIdx); + + MicroAPI::LoadAlign(vregInput, inputBuf); + + MicroAPI::MaskReg poutEQ = MicroAPI::CreateMask(); + + MicroAPI::RegTensor sqzIdxOut; + MicroAPI::Compare(poutEQ, vregInput, kthValue, pregB32); + + MicroAPI::Squeeze(sqzIdxOut, + (MicroAPI::RegTensor&)idxC, poutEQ); + MicroAPI::StoreUnAlign(outputIdxBuf, sqzIdxOut, alignIdx); + MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx); +} + +__simd_vf__ void FindValueGTOutputVFImpl(__ubuf__ uint32_t* outputValueBuf, + __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* kValue, + uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignValue; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::LoadAlign(vregInput, inputBuf + i * 64); + + MicroAPI::MaskReg poutGT = MicroAPI::CreateMask(); + + MicroAPI::RegTensor sqzValueOut; + MicroAPI::Compare(poutGT, vregInput, kthValue, pregB32); + + MicroAPI::Squeeze(sqzValueOut, vregInput, poutGT); + MicroAPI::StoreUnAlign(outputValueBuf, + sqzValueOut, alignValue); + } + MicroAPI::StoreUnAlignPost(outputValueBuf, alignValue); +} + +__simd_vf__ void FindValueEQOutputVFImpl(__ubuf__ uint32_t* outputValueBuf, + __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* kValue) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::UnalignRegForStore alignValue; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + + MicroAPI::LoadAlign(vregInput, inputBuf); + + MicroAPI::MaskReg poutEQ = MicroAPI::CreateMask(); + + MicroAPI::RegTensor sqzValueOut; + MicroAPI::Compare(poutEQ, vregInput, kthValue, pregB32); + + MicroAPI::Squeeze(sqzValueOut, vregInput, poutEQ); + MicroAPI::StoreUnAlign(outputValueBuf, sqzValueOut, alignValue); + MicroAPI::StoreUnAlignPost(outputValueBuf, alignValue); +} + +__aicore__ inline void LiTopKVF(const LocalTensor& outputIdxLocal, + const LocalTensor& outputValueLocal, + const LocalTensor& inputLocal, + const LocalTensor& tmpIdxLocal, + const LocalTensor& tmpValueLocal, + const LocalTensor& histogramsLocal, + const LocalTensor& idx0Local, + const LocalTensor& idx1Local, + const LocalTensor& idx2Local, + const LocalTensor& idx3Local, + const LocalTensor& nkValueLocal, + uint32_t topK, + uint32_t s2SeqLen) +{ + __ubuf__ uint32_t* outputIdxBuf = (__ubuf__ uint32_t*)outputIdxLocal.GetPhyAddr(); + __ubuf__ uint32_t* outputValueBuf = (__ubuf__ uint32_t*)outputValueLocal.GetPhyAddr(); + __ubuf__ uint32_t* inputBuf = (__ubuf__ uint32_t*)inputLocal.GetPhyAddr(); + __ubuf__ uint32_t* tmpIdxBuf = (__ubuf__ uint32_t*)tmpIdxLocal.GetPhyAddr(); + __ubuf__ uint32_t* tmpValueBuf = (__ubuf__ uint32_t*)tmpValueLocal.GetPhyAddr(); + __ubuf__ uint32_t* histogramsBuf = (__ubuf__ uint32_t*)histogramsLocal.GetPhyAddr(); + __ubuf__ uint32_t* idx0Buf = (__ubuf__ uint32_t*)idx0Local.GetPhyAddr(); + __ubuf__ uint32_t* idx1Buf = (__ubuf__ uint32_t*)idx1Local.GetPhyAddr(); + __ubuf__ uint32_t* idx2Buf = (__ubuf__ uint32_t*)idx2Local.GetPhyAddr(); + __ubuf__ uint32_t* idx3Buf = (__ubuf__ uint32_t*)idx3Local.GetPhyAddr(); + __ubuf__ uint32_t* nkValueBuf = (__ubuf__ uint32_t*)nkValueLocal.GetPhyAddr(); + + uint32_t bottomK = s2SeqLen - topK + 1; + uint32_t beginIdx = 0; + bool flag = true; + + const uint16_t repeatSize8 = 256; + const uint16_t repeatSize32 = 64; + + uint16_t histogramsLoopNum = (s2SeqLen + repeatSize8 - 1) / repeatSize8; + uint16_t inputLoopNum = (s2SeqLen + repeatSize32 - 1) / repeatSize32; + uint16_t topkLoopNum = (topK + 64 - 1) / 64; + + // find kth-value + HistogramsFirstVFImpl(histogramsBuf, inputBuf, histogramsLoopNum, flag); + FindFirstTargetBinVFImpl(idx0Buf, nkValueBuf, histogramsBuf, bottomK); + HistogramsSecondVFImpl(histogramsBuf, inputBuf, idx0Buf, histogramsLoopNum, flag); + FindSecondTargetBinVFImpl(idx1Buf, nkValueBuf, nkValueBuf, histogramsBuf); + HistogramsThirdVFImpl(histogramsBuf, inputBuf, idx0Buf, idx1Buf, histogramsLoopNum, flag); + FindThirdTargetBinVFImpl(idx2Buf, nkValueBuf, nkValueBuf, histogramsBuf); + HistogramsLastVFImpl(histogramsBuf, inputBuf, idx0Buf, idx1Buf, idx2Buf, histogramsLoopNum, flag); + FindKthVFImpl(nkValueBuf, histogramsBuf, idx0Buf, idx1Buf, idx2Buf, idx3Buf); + + // filter + // 输出大于k-value的值value + FindValueGTOutputVFImpl(outputValueBuf, inputBuf, nkValueBuf, inputLoopNum); + // value-当前偏移大于k-value的值在AR特殊寄存器中的有效字节数 + int64_t arValueNum = AscendC::GetSpr(); + // value-剩余需要输出等于k-value的数量 + int64_t remainValueNum = topK - (arValueNum / sizeof(uint32_t)); + for (uint16_t i = 0; i < inputLoopNum; ++i) { + int64_t arValueNumPerLoop = AscendC::GetSpr(); + if (((arValueNumPerLoop - arValueNum) / sizeof(uint32_t)) < remainValueNum) { + // 调用一次查找等于k-value情况的过程 + FindValueEQOutputVFImpl(outputValueBuf, inputBuf + i * 64, nkValueBuf); + } else { + break; + } + } + + // 输出大于k-value的值idx + FindIdxGTOutputVFImpl(outputIdxBuf, inputBuf, (uint32_t)(0), nkValueBuf, inputLoopNum); + // idx-当前偏移大于k-value的值在AR特殊寄存器中的有效字节数 + int64_t arIdxNum = AscendC::GetSpr(); + int64_t remainIdxNum = topK - (arIdxNum / sizeof(uint32_t)); + for (uint16_t i = 0; i < inputLoopNum; ++i) { + int64_t arIdxNumPerLoop = AscendC::GetSpr(); + if (((arIdxNumPerLoop - arIdxNum) / sizeof(uint32_t)) < remainIdxNum) { + // 调用一次查找等于k-value情况的过程 + beginIdx = i * 64; + FindIdxEQOutputVFImpl(outputIdxBuf, inputBuf + i * 64, beginIdx, nkValueBuf); + } else { + break; + } + } +} +} +#endif diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/vf/vf_topk_16_gather.h b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/vf/vf_topk_16_gather.h new file mode 100644 index 000000000..343c00d20 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/arch35/vf/vf_topk_16_gather.h @@ -0,0 +1,465 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. +  */ + +/*! +* \file vf_top_k_16_gather.h +* \brief +*/ + +#ifndef VF_TOP_K_16_GATHER_H +#define VF_TOP_K_16_GATHER_H + +namespace topkb16gather { + +template +__simd_vf__ void HistogramsHighVFImpl(__ubuf__ uint32_t* histogramsBuf, + __ubuf__ uint16_t* inputBuf, + uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图cout0 0-127 cout1 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + MicroAPI::RegTensor vregHigh; + MicroAPI::RegTensor vregLow; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vregLow, vregHigh, inputBuf + i * 256); + + MicroAPI::Histograms(cout0, + (MicroAPI::RegTensor&)vregHigh, + pregB8); + MicroAPI::Histograms(cout1, + (MicroAPI::RegTensor&)vregHigh, + pregB8); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, + cout0U32Even, cout0U32Odd, pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, + cout1U32Even, cout1U32Odd, pregB32); +} + +__simd_vf__ void FindHighTargetBinVFImpl(__ubuf__ uint32_t* idxHighBuf, + __ubuf__ uint32_t* nkValueBuf, + __ubuf__ uint32_t* histogramsBuf, + uint32_t bottomK) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg pregGE; + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdxHigh; + + MicroAPI::RegTensor btmK; + MicroAPI::Duplicate(btmK, bottomK); + + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdxHigh; + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::Arange(idxC, i * 64); + + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + + MicroAPI::Compare(pregGE, cout, btmK, pregB32); + + MicroAPI::Squeeze( + sqzIdxHigh, (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idxHighBuf, sqzIdxHigh, alignIdxHigh); + } + MicroAPI::StoreUnAlignPost(idxHighBuf, alignIdxHigh); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idxHigh; + MicroAPI::LoadAlign(idxHigh, idxHighBuf); + + MicroAPI::RegTensor idxAll1; + MicroAPI::RegTensor idxPrev0; + MicroAPI::RegTensor prevBinValue; + MicroAPI::Duplicate(idxAll1, 1); + + MicroAPI::RegTensor zeroAll; + MicroAPI::Duplicate(zeroAll, 0); + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + MicroAPI::Compare(preg0, idxHigh, zeroAll, pregB32); + MicroAPI::Sub(idxPrev0, idxHigh, (MicroAPI::RegTensor&)idxAll1, pregB32); + MicroAPI::ShiftRights(idxPrev0, idxPrev0, (int16_t)24, pregB32); + + MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev0, pregB32); + MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, preg0); + + MicroAPI::RegTensor nextK; + MicroAPI::Sub(nextK, btmK, prevBinValue, pregB32); + MicroAPI::StoreAlign(nkValueBuf, nextK, pregB32); +} + +template +__simd_vf__ void HistogramsLowVFImpl(__ubuf__ uint32_t* histogramsBuf, + __ubuf__ uint16_t* inputBuf, __ubuf__ uint32_t* idxHighBuf, + uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg pregEQ; + + // 计算直方图0-127 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + MicroAPI::RegTensor idxHigh; + MicroAPI::LoadAlign(idxHigh, idxHighBuf); + + MicroAPI::RegTensor vregHigh; + MicroAPI::RegTensor vregLow; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vregLow, vregHigh, inputBuf + i * 256); + + MicroAPI::Compare(pregEQ, + (MicroAPI::RegTensor&)vregHigh, + (MicroAPI::RegTensor&)idxHigh, pregB8); + + MicroAPI::Histograms(cout0, + (MicroAPI::RegTensor&)vregLow, pregEQ); + MicroAPI::Histograms(cout1, + (MicroAPI::RegTensor&)vregLow, pregEQ); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, + cout0U32Even, cout0U32Odd, pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, + cout1U32Even, cout1U32Odd, pregB32); +} + +__simd_vf__ void FindKthVFImpl(__ubuf__ uint32_t* kValue, + __ubuf__ uint32_t* histogramsBuf, __ubuf__ uint32_t* idxHighBuf, + __ubuf__ uint32_t* idxLowBuf) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg pregGE; + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdxLow; + + MicroAPI::RegTensor btmK; + MicroAPI::LoadAlign(btmK, kValue); + + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdxLow; + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::Arange(idxC, i * 64); + + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + + MicroAPI::Compare(pregGE, cout, btmK, pregB32); + + MicroAPI::Squeeze(sqzIdxLow, + (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idxLowBuf, sqzIdxLow, alignIdxLow); + } + MicroAPI::StoreUnAlignPost(idxLowBuf, alignIdxLow); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idxHigh; + MicroAPI::RegTensor idxLow; + MicroAPI::LoadAlign(idxHigh, idxHighBuf); + MicroAPI::LoadAlign(idxLow, idxLowBuf); + + MicroAPI::RegTensor idxTmp; + MicroAPI::Duplicate(idxTmp, 0xff00); + + MicroAPI::And(idxHigh, idxHigh, (MicroAPI::RegTensor&)idxTmp, pregB32); + + MicroAPI::RegTensor idxK; + MicroAPI::Add(idxK, idxHigh, idxLow, pregB16); + + MicroAPI::StoreAlign(kValue, idxK, pregB32); +} + +/** + 输出所有大于的kth-value的Index + */ +__simd_vf__ void FindIdxGTOutputVFImpl(__ubuf__ uint16_t* outputIdxBuf, + __ubuf__ uint16_t* inputValueBuf, uint16_t beginIdx, + __ubuf__ uint32_t* kValue, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg poutGT; + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor sqzIdxOut; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::Arange(idxC, beginIdx + i * 128); + + MicroAPI::LoadAlign(vregInput, inputValueBuf + i * 128); + + MicroAPI::Compare(poutGT, vregInput, (MicroAPI::RegTensor&)kthValue, pregB16); + + MicroAPI::Squeeze(sqzIdxOut, + (MicroAPI::RegTensor&)idxC, poutGT); + MicroAPI::StoreUnAlign(outputIdxBuf, sqzIdxOut, alignIdx); + } + MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx); +} + +/** + 输出所有等于的kth-value的Index + */ +__simd_vf__ void FindIdxEQOutputVFImpl(__ubuf__ uint16_t* outputIdxBuf, + __ubuf__ uint16_t* inputValueBuf, uint16_t beginIdx, + __ubuf__ uint32_t* kValue, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg poutEQ; + + MicroAPI::UnalignRegForStore alignIdx; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor sqzIdxOut; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::Arange(idxC, beginIdx + i * 128); + + MicroAPI::LoadAlign(vregInput, inputValueBuf + i * 128); + + MicroAPI::Compare(poutEQ, vregInput, (MicroAPI::RegTensor&)kthValue, pregB16); + + MicroAPI::Squeeze(sqzIdxOut, + (MicroAPI::RegTensor&)idxC, poutEQ); + MicroAPI::StoreUnAlign(outputIdxBuf, sqzIdxOut, alignIdx); + } + MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx); +} + +/** + 输出最终的Value + */ +__simd_vf__ void FindValueOutputVFImpl(__ubuf__ uint16_t* outputValueBuf, + __ubuf__ uint16_t* inputValueBuf, + __ubuf__ uint16_t* tmpIdxBuf, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + + MicroAPI::RegTensor tmpIdx; + MicroAPI::RegTensor outputValue; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::LoadAlign(tmpIdx, tmpIdxBuf + i * 128); + + MicroAPI::Gather(outputValue, inputValueBuf, tmpIdx, pregB16); + + MicroAPI::StoreAlign(outputValueBuf + i * 128, outputValue, pregB16); + } +} + +/** + 输出最终的Idx + */ +__simd_vf__ void FindRealIndexVFImpl(__ubuf__ uint32_t* outputIdxBuf, + __ubuf__ uint16_t* tmpIdxBuf, __ubuf__ uint32_t* hisIdxBuf, + uint32_t topK, uint32_t loopIndex, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg pregNow; + MicroAPI::MaskReg pregHis; + + MicroAPI::RegTensor tmpIdx; + MicroAPI::RegTensor outputGatherIdx; + MicroAPI::RegTensor outputAddsIdx; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::LoadAlign(tmpIdx, tmpIdxBuf + i * 64); + + MicroAPI::Compares(pregNow, (MicroAPI::RegTensor&)tmpIdx, topK - 1, pregB32); + MicroAPI::Xor(pregHis, pregNow, pregB32, pregB32); + + MicroAPI::Gather(outputGatherIdx, hisIdxBuf, (MicroAPI::RegTensor&)tmpIdx, pregHis); + MicroAPI::Adds(outputAddsIdx, (MicroAPI::RegTensor&)tmpIdx, loopIndex, pregNow); + + MicroAPI::Add(outputGatherIdx, outputGatherIdx, outputAddsIdx, pregB32); + + MicroAPI::StoreAlign(outputIdxBuf + i * 64, outputGatherIdx, pregB32); + } +} + +/** + * @brief LiTopKVF 对一个validLen的输入进行topk算法，输出idx_tmp + * @param tmpIdxLocal Temp阶段输出的TopKIndex;如果s2SeqLen < 16K作为最终输出 validLen * 2B + * @param outputValueLocal 如果s2SeqLen > 16K并且是首轮输出Value topK * 2B + * @param inputValueLocal 输入Value validLen * 2B + * @param histogramsLocal 直方图 256 * 4B + * @param idxHighLocal 目标桶高八位 256 * 4B + * @param idxLowLocal 目标桶低八位 256 * 4B + * @param nkValueLocal 存储next_k的值 64 * 4B + * @param topK topK元素 + * @param validLen 有效元素个数:LICommon::Align(topkCountAlign256_ + validTrunkLen, (uint32_t)256) + */ +template // 是否输出VALUE +__aicore__ inline void LiTopKVF(const LocalTensor& tmpIdxLocal, + const LocalTensor& outputValueLocal, + const LocalTensor& inputValueLocal, + const LocalTensor& histogramsLocal, + const LocalTensor& idxHighLocal, + const LocalTensor& idxLowLocal, + const LocalTensor& nkValueLocal, + uint32_t topK, + uint32_t validLen) +{ + __ubuf__ uint16_t* tmpIdxBuf = (__ubuf__ uint16_t*)tmpIdxLocal.GetPhyAddr(); + __ubuf__ uint16_t* outputValueBuf = (__ubuf__ uint16_t*)outputValueLocal.GetPhyAddr(); + __ubuf__ uint16_t* inputValueBuf = (__ubuf__ uint16_t*)inputValueLocal.GetPhyAddr(); + __ubuf__ uint32_t* histogramsBuf = (__ubuf__ uint32_t*)histogramsLocal.GetPhyAddr(); + __ubuf__ uint32_t* idxHighBuf = (__ubuf__ uint32_t*)idxHighLocal.GetPhyAddr(); + __ubuf__ uint32_t* idxLowBuf = (__ubuf__ uint32_t*)idxLowLocal.GetPhyAddr(); + __ubuf__ uint32_t* nkValueBuf = (__ubuf__ uint32_t*)nkValueLocal.GetPhyAddr(); + + uint32_t bottomK = validLen - topK + 1; + uint32_t beginIdx = 0; + bool flag = true; + + const uint16_t repeatSize8 = 256; + const uint16_t repeatSize16 = 128; + const uint16_t repeatSize32 = 64; + + uint16_t histogramsLoopNum = (validLen + repeatSize8 - 1) / repeatSize8; + uint16_t inputLoopNum = (validLen + repeatSize16 - 1) / repeatSize16; + uint16_t topkLoopNum = (topK + repeatSize32 - 1) / repeatSize32; + uint16_t topkLoopNum16 = (topK + repeatSize16 - 1) / repeatSize16; + + // find kth-value + HistogramsHighVFImpl(histogramsBuf, inputValueBuf, histogramsLoopNum, flag); + FindHighTargetBinVFImpl(idxHighBuf, nkValueBuf, histogramsBuf, bottomK); + + HistogramsLowVFImpl(histogramsBuf, inputValueBuf, idxHighBuf, histogramsLoopNum, flag); + FindKthVFImpl(nkValueBuf, histogramsBuf, idxHighBuf, idxLowBuf); + + // filter + int32_t count = LICommon::Align(topK, (uint32_t)128) - topK / 128 * 128; + AscendC::Duplicate(tmpIdxLocal[topK / 128 * 128], (uint16_t)(0), count); + // 输出大于k-value的值idx + FindIdxGTOutputVFImpl(tmpIdxBuf, inputValueBuf, (uint32_t)(0), nkValueBuf, inputLoopNum); + // 输出等于k-value的值idx + FindIdxEQOutputVFImpl(tmpIdxBuf, inputValueBuf, (uint32_t)(0), nkValueBuf, inputLoopNum); + + // 是否输出Value + if constexpr (ISOUTVALUE) { + FindValueOutputVFImpl(outputValueBuf, inputValueBuf, tmpIdxBuf, topkLoopNum16); + } +} + +/** + * @brief 通过idx_tmp gather出实际的TopKIndex，s2SeqLen > 16K才会执行 + * @param outputIdxLocal 输出Idx 有效:topK * 2B + * @param outputValueLocal 输出Value topK * 2B(以后需要输出实际value使用) + * @param inputValueLocal 输入Value validLen * 2B + * @param tmpIdxLocal 本轮tmpIdx输入 validLen * 2B (0 ~ validLen - 1) + * @param hisIdxLocal 上一轮实际Idx输入 有效:topK * 4B + * @param topK topK元素个数 + * @param loopBasicIdx 当前循环需要加上得基准Index + * @param validLen 有效元素个数 + */ +__aicore__ inline void LiTopKGatherVF(const LocalTensor& outputIdxLocal, + const LocalTensor& outputValueLocal, + const LocalTensor& inputValueLocal, + const LocalTensor& tmpIdxLocal, + const LocalTensor& hisIdxLocal, + uint32_t topK, + uint32_t loopBasicIdx, + uint32_t validLen) +{ + __ubuf__ uint32_t* outputIdxBuf = (__ubuf__ uint32_t*)outputIdxLocal.GetPhyAddr(); + __ubuf__ uint16_t* outputValueBuf = (__ubuf__ uint16_t*)outputValueLocal.GetPhyAddr(); + __ubuf__ uint16_t* inputValueBuf = (__ubuf__ uint16_t*)inputValueLocal.GetPhyAddr(); + __ubuf__ uint16_t* tmpIdxBuf = (__ubuf__ uint16_t*)tmpIdxLocal.GetPhyAddr(); + __ubuf__ uint32_t* hisIdxBuf = (__ubuf__ uint32_t*)hisIdxLocal.GetPhyAddr(); + + const uint16_t repeatSize32 = 64; + const uint16_t repeatSize16 = 128; + uint16_t topkLoopNum16 = (topK + repeatSize16 - 1) / repeatSize16; + uint16_t topkLoopNum32 = (topK + repeatSize32 - 1) / repeatSize32; + + FindRealIndexVFImpl(outputIdxBuf, tmpIdxBuf, hisIdxBuf, topK, loopBasicIdx, topkLoopNum32); +} +} +#endif diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/lightning_indexer.cpp b/csrc/ascend/attention/lightning_indexer/op_kernel/lightning_indexer.cpp new file mode 100644 index 000000000..ddb3f29fc --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/lightning_indexer.cpp @@ -0,0 +1,66 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer.cpp + * \brief + */ + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" +#include "lightning_indexer_template_tiling_key.h" + +#if (__CCE_AICORE__ == 310) + #include "arch35/lightning_indexer_kernel.h" + +#else + #include "arch22/lightning_indexer_kernel.h" +#endif + +using namespace LIKernel; + +#define INVOKE_LI_NO_KFC_OP_IMPL(templateClass, ...) \ + do { \ + templateClass> op; \ + GET_TILING_DATA_WITH_STRUCT(LITilingData, tiling_data_in, tiling); \ + const LITilingData *__restrict tiling_data = &tiling_data_in; \ + op.Init(query, key, weights, actualSeqLengthsQ, actualSeqLengths, blocktable, sparseIndices, sparseValues, user, \ + tiling_data, &tPipe); \ + op.Process(); \ + } while (0) + +template +__global__ __aicore__ void lightning_indexer(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths, + __gm__ uint8_t *blocktable, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *sparseValues, __gm__ uint8_t *workspace, + __gm__ uint8_t *tiling) +{ + TPipe tPipe; + __gm__ uint8_t *user = GetUserWorkspace(workspace); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) || (__CCE_AICORE__ == 200) + if (ORIG_DTYPE_QUERY == DT_BF16) { + INVOKE_LI_NO_KFC_OP_IMPL(LightningIndexerKernel, bfloat16_t, bfloat16_t, int32_t, + PAGE_ATTENTION, LI_LAYOUT(LAYOUT_T), LI_LAYOUT(K_LAYOUT_T), DT_W_FLAG); + } else if (ORIG_DTYPE_QUERY == DT_FLOAT16) { + INVOKE_LI_NO_KFC_OP_IMPL(LightningIndexerKernel, half, half, int32_t, + PAGE_ATTENTION, LI_LAYOUT(LAYOUT_T), LI_LAYOUT(K_LAYOUT_T), DT_W_FLAG); + } +#else + if constexpr (DT_Q == LI_TPL_FP16 && DT_K == LI_TPL_FP16 && DT_OUT == LI_TPL_INT32) { + INVOKE_LI_NO_KFC_OP_IMPL(LightningIndexerKernel, half, half, int32_t, PAGE_ATTENTION, + LI_LAYOUT(LAYOUT_T), LI_LAYOUT(K_LAYOUT_T), DT_W_FLAG); + } else { + INVOKE_LI_NO_KFC_OP_IMPL(LightningIndexerKernel, bfloat16_t, bfloat16_t, int32_t, PAGE_ATTENTION, + LI_LAYOUT(LAYOUT_T), LI_LAYOUT(K_LAYOUT_T), DT_W_FLAG); + } +#endif +} diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/lightning_indexer_common.h b/csrc/ascend/attention/lightning_indexer/op_kernel/lightning_indexer_common.h new file mode 100644 index 000000000..9ae75b89a --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/lightning_indexer_common.h @@ -0,0 +1,175 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_common.h + * \brief + */ +#ifndef LIGHTNING_INDEXER_COMMON_H +#define LIGHTNING_INDEXER_COMMON_H +using namespace AscendC; +namespace LICommon { + +// 与tiling的layout保持一致 +enum class LI_LAYOUT { + BSND = 0, + TND = 1, + PA_BSND = 2 +}; + +template +struct LIType { + static constexpr bool weightsTypeFlag = DT_W_FLAG; // weight的dtype是否为FP32 + using queryType = Q_T; + using keyType = K_T; + using outputType = OUT_T; + static constexpr bool pageAttention = PAGE_ATTENTION; + static constexpr LI_LAYOUT layout = LAYOUT_T; + static constexpr LI_LAYOUT keyLayout = K_LAYOUT_T; +}; + +struct RunInfo { + uint32_t loop; + uint32_t bN2Idx; + uint32_t bIdx; + uint32_t n2Idx = 0; + uint32_t gS1Idx; + uint32_t s2Idx; + + uint32_t actS1Size = 1; + uint32_t actS2Size = 1; + uint32_t actS2SizeOrig = 1; + uint32_t actMBaseSize; + uint32_t actualSingleProcessSInnerSize; + uint32_t actualSingleProcessSInnerSizeAlign; + + uint64_t tensorQueryOffset; + uint64_t tensorKeyOffset; + uint64_t tensorWeightsOffset; + uint64_t indiceOutOffset; + uint64_t valueOutOffset; + + bool isFirstS2InnerLoop; + bool isLastS2InnerLoop; + bool isAllLoopEnd = false; + bool isValid = false; +}; + +struct ConstInfo { + // CUBE与VEC核间同步的模式 + static constexpr uint32_t FIA_SYNC_MODE2 = 2; + static constexpr uint32_t QLI_SYNC_MODE4 = 4; + static constexpr uint32_t AIV0_AIV1_OFFSET = 16; + static constexpr uint32_t CROSS_VC_EVENT = 0; + static constexpr uint32_t CROSS_CV_EVENT = 2; + // BUFFER的字节数 + static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; + static constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64; + static constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256; + static constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512; + static constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024; + static constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048; + static constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096; + static constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192; + static constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384; + static constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768; + // 无效索引 + static constexpr int INVALID_IDX = -1; + uint16_t INVALID_VAL = 0; + // CUBE和VEC的核间同步EventID + uint32_t syncC1V1 = 0U; + uint32_t syncC1V0 = 2U; + uint32_t syncV1C1 = 0U; + uint32_t syncV0C1 = 1U; + + // 基本块大小 + uint32_t mBaseSize = 1ULL; + uint32_t mBaseSizeAlign = 1ULL; + uint32_t s1BaseSize = 1ULL; + uint32_t s2BaseSize = 1ULL; + + uint64_t batchSize = 0ULL; + uint64_t gSize = 0ULL; + uint64_t qHeadNum = 0ULL; + uint64_t kHeadNum; + uint64_t headDim; + uint64_t sparseCount; // topK选取大小 + uint64_t kSeqSize = 0ULL; // kv最大S长度 + uint64_t qSeqSize = 1ULL; // q最大S长度 + uint32_t kCacheBlockSize = 0; // PA场景的block size + uint32_t maxBlockNumPerBatch = 0; // PA场景的最大单batch block number + LI_LAYOUT outputLayout; // 输出的格式 + bool attenMaskFlag = false; + int64_t preTokens = INT64_MAX; + int64_t nextTokens = INT64_MAX; + bool returnValue = false; + + uint32_t actualLenQDims = 0U; // query的actualSeqLength 的维度 + uint32_t actualLenDims = 0U; // KV 的actualSeqLength 的维度 + bool isAccumSeqS1 = false; // 是否累加模式 + bool isAccumSeqS2 = false; // 是否累加模式 + bool isSparseCountOver2K = false; //sparseCount小于等于2048为false + bool isLDOpen = false; + bool returnValueFlag = false; + bool splitMFlag = false; +}; + +struct SplitCoreInfo { + uint32_t s2Start = 0U; // S2的起始位置 + uint32_t s2End = 0U; // S2循环index上限 + uint32_t bN2Start = 0U; + uint32_t bN2End = 0U; + uint32_t gS1Start = 0U; + uint32_t gS1End = 0U; + bool isLD = false; // 当前核是否需要进行Decode归约任务 + bool isCoreEnable = false; +}; + +template +__aicore__ inline T Align(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd) * (rnd))); +} + +template +__aicore__ inline T1 Min(T1 a, T2 b) +{ + return (a > b) ? (b) : (a); +} + +template +__aicore__ inline T1 Max(T1 a, T2 b) +{ + return (a > b) ? (a) : (b); +} + +template +__aicore__ inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd))); +} +} // namespace LICommon + +// bank冲突优化 +// david 256KB bank layout +// shape ( bank_depth ( banks bank_groups block)) (512 ( 2 8 32)) +// stride (banks*bank_groups*block (bank_groups*block block 1)) (512 (256 32 1)) +#define UB_BLOCK 32 // 32B +#define UB_BANK_GROUPS 8 +#define UB_BANKS 2 +#define UB_BANK_DEPTH 512 + +#define UB_BANK_GROUP_STRIDE UB_BLOCK // 32B +#define UB_BANK_STRIDE (UB_BANK_GROUPS * UB_BLOCK) // 256B +#define UB_BANK_DEPTH_STRIDE (UB_BANKS * UB_BANK_GROUPS * UB_BLOCK) // 512B + +#endif // LIGHTNING_INDEXER_COMMON_H diff --git a/csrc/ascend/attention/lightning_indexer/op_kernel/lightning_indexer_template_tiling_key.h b/csrc/ascend/attention/lightning_indexer/op_kernel/lightning_indexer_template_tiling_key.h new file mode 100644 index 000000000..7b3022107 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer/op_kernel/lightning_indexer_template_tiling_key.h @@ -0,0 +1,88 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_template_tiling_key.h + * \brief + */ + +#ifndef TEMPLATE_TILING_KEY_LI_H_ +#define TEMPLATE_TILING_KEY_LI_H_ + +#include "ascendc/host_api/tiling/template_argument.h" + +#define LI_TPL_FP32 0 +#define LI_TPL_FP16 1 +#define LI_TPL_INT32 3 +#define LI_TPL_BF16 27 + +#define LI_LAYOUT_BSND 0 +#define LI_LAYOUT_TND 1 +#define LI_LAYOUT_PA_BSND 2 + +#define ASCENDC_TPL_4_BW 4 + +// 模板参数支持的范围定义 +ASCENDC_TPL_ARGS_DECL(LightningIndexer, // 算子OpType + ASCENDC_TPL_DTYPE_DECL(DT_Q, LI_TPL_FP16, LI_TPL_BF16), + ASCENDC_TPL_DTYPE_DECL(DT_K, LI_TPL_FP16, LI_TPL_BF16), + ASCENDC_TPL_DTYPE_DECL(DT_OUT, LI_TPL_INT32), ASCENDC_TPL_BOOL_DECL(PAGE_ATTENTION, 0, 1), + ASCENDC_TPL_UINT_DECL(LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, LI_LAYOUT_BSND, + LI_LAYOUT_TND), + ASCENDC_TPL_UINT_DECL(K_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, + LI_LAYOUT_BSND, LI_LAYOUT_TND, LI_LAYOUT_PA_BSND), + ASCENDC_TPL_BOOL_DECL(DT_W_FLAG, 0, 1), ); + +// 支持的模板参数组合 +// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法 +ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, LI_TPL_FP16), ASCENDC_TPL_DTYPE_SEL(DT_K, LI_TPL_FP16), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, LI_TPL_INT32), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_BSND, LI_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_PA_BSND), + ASCENDC_TPL_BOOL_SEL(DT_W_FLAG, 0, 1), ), + + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, LI_TPL_BF16), ASCENDC_TPL_DTYPE_SEL(DT_K, LI_TPL_BF16), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, LI_TPL_INT32), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_BSND, LI_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_PA_BSND), + ASCENDC_TPL_BOOL_SEL(DT_W_FLAG, 0, 1), ), + + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, LI_TPL_FP16), ASCENDC_TPL_DTYPE_SEL(DT_K, LI_TPL_FP16), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, LI_TPL_INT32), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_BSND), + ASCENDC_TPL_BOOL_SEL(DT_W_FLAG, 0, 1), ), + + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, LI_TPL_FP16), ASCENDC_TPL_DTYPE_SEL(DT_K, LI_TPL_FP16), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, LI_TPL_INT32), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_TND), + ASCENDC_TPL_BOOL_SEL(DT_W_FLAG, 0, 1), ), + + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, LI_TPL_BF16), ASCENDC_TPL_DTYPE_SEL(DT_K, LI_TPL_BF16), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, LI_TPL_INT32), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_BSND), + ASCENDC_TPL_BOOL_SEL(DT_W_FLAG, 0, 1), ), + + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, LI_TPL_BF16), ASCENDC_TPL_DTYPE_SEL(DT_K, LI_TPL_BF16), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, LI_TPL_INT32), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_TND), + ASCENDC_TPL_BOOL_SEL(DT_W_FLAG, 0, 1), ), ); + +#endif diff --git a/csrc/ascend/attention/lightning_indexer_quant/CMakeLists.txt b/csrc/ascend/attention/lightning_indexer_quant/CMakeLists.txt new file mode 100644 index 000000000..298ed5677 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_quant/lightning_indexer_quant_torch_adpt.h b/csrc/ascend/attention/lightning_indexer_quant/lightning_indexer_quant_torch_adpt.h new file mode 100644 index 000000000..f14570f40 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/lightning_indexer_quant_torch_adpt.h @@ -0,0 +1,81 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef LIGHTING_INDEXER_QUANT_VLLM_TORCH_ADPT_H +#define LIGHTING_INDEXER_QUANT_VLLM_TORCH_ADPT_H +namespace vllm_fl { + +at::Tensor npu_lightning_indexer_quant( + const at::Tensor &query, const at::Tensor &key, const at::Tensor &weights, + const at::Tensor &query_dequant_scale, const at::Tensor &key_dequant_scale, + const c10::optional &actual_seq_lengths_query, + const c10::optional &actual_seq_lengths_key, + const c10::optional &block_table, int64_t query_quant_mode, int64_t key_quant_mode, + c10::string_view layout_query, c10::string_view layout_key, int64_t sparse_count, int64_t sparse_mode) +{ + std::string query_layout_str = std::string(layout_query); + std::string key_layout_str = std::string(layout_key); + + const int SIZE = 8; + const int DIM_0 = 0; + const int DIM_1 = 1; + const int DIM_2 = 2; + const int DIM_3 = 3; + + at::SmallVector output_size; + for (size_t i = 0; i < query.sizes().size(); i++) { + TORCH_CHECK(query.size(i) > 0, "All values within query's shape should be greater " + "than 0, but shape[", i, "] is ", query.size(i)); + } + for (size_t i = 0; i < key.sizes().size(); i++) { + TORCH_CHECK(key.size(i) > 0, "All values within key's shape should be greater " + "than 0, but shape[", i, "] is ", key.size(i)); + } + TORCH_CHECK(sparse_count > 0, "sparse count should be greater than 0, but now is ", sparse_count); + int64_t keyHeadNum = (key_layout_str == "TND")? key.size(DIM_1) : key.size(DIM_2); + if (query_layout_str == "BSND") { + output_size = {query.size(DIM_0), query.size(DIM_1), keyHeadNum, sparse_count}; + } else { + output_size = {query.size(DIM_0), keyHeadNum, sparse_count}; + } + at::Tensor lightning_indexer_quant_output = at::empty(output_size, query.options().dtype(at::kInt)); + + // convert str + char *query_layout_ptr = const_cast(query_layout_str.c_str()); + char *key_layout_ptr = const_cast(key_layout_str.c_str()); + + EXEC_NPU_CMD(aclnnLightningIndexerQuant, + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + actual_seq_lengths_query, + actual_seq_lengths_key, + block_table, + query_quant_mode, + key_quant_mode, + query_layout_ptr, + key_layout_ptr, + sparse_count, + sparse_mode, + lightning_indexer_quant_output + ); + + return lightning_indexer_quant_output; + +} +} +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_quant/op_host/CMakeLists.txt b/csrc/ascend/attention/lightning_indexer_quant/op_host/CMakeLists.txt new file mode 100644 index 000000000..7c747ed78 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/op_host/CMakeLists.txt @@ -0,0 +1,24 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + lightning_indexer_quant_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME LightningIndexerQuant + OPTIONS + --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE lightning_indexer_quant ACLNNTYPE aclnn) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/attention/lightning_indexer_quant/op_host/lightning_indexer_quant_def.cpp b/csrc/ascend/attention/lightning_indexer_quant/op_host/lightning_indexer_quant_def.cpp new file mode 100644 index 000000000..7049c2d03 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/op_host/lightning_indexer_quant_def.cpp @@ -0,0 +1,85 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_quant_def.cpp + * \brief + */ +#include + +#include "register/op_def_registry.h" + +namespace ops { +class LightningIndexerQuant : public OpDef { +public: + explicit LightningIndexerQuant(const char *name) : OpDef(name) + { + this->Input("query") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("key") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("weights") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("query_dequant_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("key_dequant_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("actual_seq_lengths_query") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("actual_seq_lengths_key") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("block_table") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("sparse_indices").ParamType(REQUIRED).DataType({ge::DT_INT32}).Format({ge::FORMAT_ND}); + this->Attr("query_quant_mode").AttrType(REQUIRED).Int(0); // 0: 默认值，per-token-head + this->Attr("key_quant_mode").AttrType(REQUIRED).Int(0); // 0: 默认值，per-token-head + this->Attr("layout_query").AttrType(OPTIONAL).String("BSND"); + this->Attr("layout_key").AttrType(OPTIONAL).String("PA_BSND"); + this->Attr("sparse_count").AttrType(OPTIONAL).Int(2048); // 2048: 默认值，筛选前2048 + this->Attr("sparse_mode").AttrType(OPTIONAL).Int(3); // 3: 默认值，只计算下三角 + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("aclnnSupport.value", "support_aclnn") + .ExtendCfgInfo("jitCompile.flag", "static_false,dynamic_false"); + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + } +}; +OP_ADD(LightningIndexerQuant); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_quant/op_host/lightning_indexer_quant_proto.cpp b/csrc/ascend/attention/lightning_indexer_quant/op_host/lightning_indexer_quant_proto.cpp new file mode 100644 index 000000000..fb4539d61 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/op_host/lightning_indexer_quant_proto.cpp @@ -0,0 +1,91 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_quant_proto.cpp + * \brief + */ +#include +#include + +#include "error/ops_error.h" + +using namespace ge; + +namespace ops { +constexpr uint32_t QUERY_INDEX = 0; +constexpr uint32_t KEY_INDEX = 1; +constexpr uint32_t ATTR_QUERY_LAYOUT_INDEX = 2; +constexpr uint32_t ATTR_KV_LAYOUT_INDEX = 3; +constexpr uint32_t ATTR_SPARSE_COUNT_INDEX = 4; + +static ge::graphStatus InferShapeLightningIndexerQuant(gert::InferShapeContext *context) +{ + if (context == nullptr) { + OPS_LOG_E("LightningIndexerQuant", "context is nullptr!"); + return ge::GRAPH_FAILED; + } + const gert::Shape *queryShape = context->GetInputShape(QUERY_INDEX); + OPS_LOG_E_IF_NULL(context, queryShape, return ge::GRAPH_FAILED); + const gert::Shape *keyShape = context->GetInputShape(KEY_INDEX); + OPS_LOG_E_IF_NULL(context, keyShape, return ge::GRAPH_FAILED); + gert::Shape *outShape = context->GetOutputShape(0); + + auto attrs = context->GetAttrs(); + OPS_LOG_E_IF_NULL(context, attrs, return ge::GRAPH_FAILED); + const char *inputLayoutQueryPtr = attrs->GetAttrPointer(ATTR_QUERY_LAYOUT_INDEX); + OPS_LOG_E_IF_NULL(context, inputLayoutQueryPtr, return ge::GRAPH_FAILED); + const char *inputLayoutKeyPtr = attrs->GetAttrPointer(ATTR_KV_LAYOUT_INDEX); + OPS_LOG_E_IF_NULL(context, inputLayoutKeyPtr, return ge::GRAPH_FAILED); + const int64_t *sparse_count = attrs->GetInt(ATTR_SPARSE_COUNT_INDEX); + OPS_LOG_E_IF_NULL(context, sparse_count, return ge::GRAPH_FAILED); + + std::string inputLayoutQueryPtrStr = std::string(inputLayoutQueryPtr); + std::string inputLayoutKeyPtrStr = std::string(inputLayoutKeyPtr); + if (inputLayoutQueryPtrStr != "TND" && inputLayoutQueryPtrStr != "BSND") { + OPS_LOG_E(context, "The input layout query should be TND or BSND, but got %s.", inputLayoutQueryPtrStr.c_str()); + return GRAPH_FAILED; + } + + outShape->SetDimNum(queryShape->GetDimNum()); + int64_t keyHeadNum = (inputLayoutKeyPtrStr == "TND") ? keyShape->GetDim(1) : keyShape->GetDim(2); + if (inputLayoutQueryPtrStr == "BSND") { + outShape->SetDim(0, queryShape->GetDim(0)); // 0:Dim B + outShape->SetDim(1, queryShape->GetDim(1)); // 1:Dim S + outShape->SetDim(2, keyHeadNum); // 2:Dim N + outShape->SetDim(3, *sparse_count); // 3:Dim K + } else { + outShape->SetDim(0, queryShape->GetDim(0)); // 0:Dim T + outShape->SetDim(1, keyHeadNum); // 1:output shape's N Dim, 2: key shape's N Dim + outShape->SetDim(2, *sparse_count); // 2:Dim K + } + + OPS_LOG_D(context->GetNodeName(), "LightningIndexerQuant InferShape end."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataTypeLightningIndexerQuant(gert::InferDataTypeContext *context) +{ + if (context == nullptr) { + OPS_LOG_E("LightningIndexerQuant", "InferDataTypeContext context is nullptr!"); + return ge::GRAPH_FAILED; + } + OPS_LOG_D(context->GetNodeName(), "Enter LightningIndexerQuant InferDataType impl."); + // default index data type is int32 + ge::DataType outputType = ge::DT_INT32; + context->SetOutputDataType(0, outputType); + OPS_LOG_D(context->GetNodeName(), "LightningIndexerQuant InferDataType end."); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(LightningIndexerQuant) + .InferShape(InferShapeLightningIndexerQuant) + .InferDataType(InferDataTypeLightningIndexerQuant); +} // namespace ops diff --git a/csrc/ascend/attention/lightning_indexer_quant/op_host/lightning_indexer_quant_tiling.cpp b/csrc/ascend/attention/lightning_indexer_quant/op_host/lightning_indexer_quant_tiling.cpp new file mode 100644 index 000000000..042f0271c --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/op_host/lightning_indexer_quant_tiling.cpp @@ -0,0 +1,828 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_quant_tiling.cpp + * \brief + */ + +#include "lightning_indexer_quant_tiling.h" + +#include "../op_kernel/lightning_indexer_quant_template_tiling_key.h" + +using namespace ge; +using namespace AscendC; +using std::map; +using std::string; +namespace optiling { +// --------------------------LIQInfoParser类成员函数定义------------------------------------- +ge::graphStatus LIQInfoParser::CheckRequiredInOutExistence() const +{ + OPS_ERR_IF(opParamInfo_.query.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor query is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.query.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor query is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.key.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor key is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.key.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor key is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.weights.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor weights is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.weights.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor weights is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.query_dequant_scale.shape == nullptr, + OPS_LOG_E(opName_, "Shape of tensor query_dequant_scale is nullptr"), return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.query_dequant_scale.desc == nullptr, + OPS_LOG_E(opName_, "Desc of tensor query_dequant_scale is nullptr"), return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.key_dequant_scale.shape == nullptr, + OPS_LOG_E(opName_, "Shape of tensor key_dequant_scale is nullptr"), return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.key_dequant_scale.desc == nullptr, + OPS_LOG_E(opName_, "Desc of tensor key_dequant_scale is nullptr"), return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.attenOut.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor output is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.attenOut.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor output is nullptr"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::CheckRequiredAttrExistence() const +{ + OPS_ERR_IF(opParamInfo_.layOutQuery == nullptr, OPS_LOG_E(opName_, "attr layout_query is nullptr"), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(opParamInfo_.layOutKey == nullptr, OPS_LOG_E(opName_, "attr layout_key is nullptr"), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(opParamInfo_.sparseCount == nullptr, OPS_LOG_E(opName_, "attr sparse_count is nullptr"), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(opParamInfo_.sparseMode == nullptr, OPS_LOG_E(opName_, "attr sparse_mode is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.queryQuantMode == nullptr, OPS_LOG_E(opName_, "query_quant_mode is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.keyQuantMode == nullptr, OPS_LOG_E(opName_, "key_quant_mode is nullptr"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::CheckRequiredParaExistence() const +{ + if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS || CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetOpName() +{ + if (context_->GetNodeName() == nullptr) { + OPS_LOG_E("LightningIndexerQuant", "opName got from TilingContext is nullptr"); + return ge::GRAPH_FAILED; + } + opName_ = context_->GetNodeName(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetNpuInfo() +{ + platformInfo_ = context_->GetPlatformInfo(); + OPS_ERR_IF(platformInfo_ == nullptr, OPS_LOG_E(opName_, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo_); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + OPS_ERR_IF(aicNum == 0 || aivNum == 0, OPS_LOG_E(opName_, "num of core obtained is 0."), return GRAPH_FAILED); + + socVersion_ = ascendcPlatform.GetSocVersion(); + if ((socVersion_ != platform_ascendc::SocVersion::ASCEND910B) && + (socVersion_ != platform_ascendc::SocVersion::ASCEND910_93)) { + OPS_LOG_E(opName_, "SOC Version[%d] is not support.", (int32_t)socVersion_); + return GRAPH_FAILED; + } + OPS_ERR_IF(context_->GetWorkspaceSizes(1) == nullptr, OPS_LOG_E(opName_, "workSpaceSize got from ge is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(context_->GetRawTilingData() == nullptr, + OPS_LOG_E(context_->GetNodeName(), "RawTilingData got from GE context is nullptr."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void LIQInfoParser::GetOptionalInputParaInfo() +{ + opParamInfo_.actualSeqLengthsQ.tensor = context_->GetOptionalInputTensor(ACTUAL_SEQ_Q_INDEX); + opParamInfo_.actualSeqLengthsQ.desc = context_->GetOptionalInputDesc(ACTUAL_SEQ_Q_INDEX); + opParamInfo_.actualSeqLengthsK.tensor = context_->GetOptionalInputTensor(ACTUAL_SEQ_K_INDEX); + opParamInfo_.actualSeqLengthsK.desc = context_->GetOptionalInputDesc(ACTUAL_SEQ_K_INDEX); + opParamInfo_.blockTable.tensor = context_->GetOptionalInputTensor(BLOCK_TABLE_INDEX); + opParamInfo_.blockTable.desc = context_->GetOptionalInputDesc(BLOCK_TABLE_INDEX); +} + +void LIQInfoParser::GetInputParaInfo() +{ + opParamInfo_.query.desc = context_->GetInputDesc(QUERY_INDEX); + opParamInfo_.query.shape = context_->GetInputShape(QUERY_INDEX); + opParamInfo_.key.desc = context_->GetInputDesc(KEY_INDEX); + opParamInfo_.key.shape = context_->GetInputShape(KEY_INDEX); + opParamInfo_.weights.desc = context_->GetInputDesc(WEIGTHS_INDEX); + opParamInfo_.weights.shape = context_->GetInputShape(WEIGTHS_INDEX); + opParamInfo_.query_dequant_scale.desc = context_->GetInputDesc(QUERY_DEQUANT_SCALE_INDEX); + opParamInfo_.query_dequant_scale.shape = context_->GetInputShape(QUERY_DEQUANT_SCALE_INDEX); + opParamInfo_.key_dequant_scale.desc = context_->GetInputDesc(KEY_DEQUANT_SCALE_INDEX); + opParamInfo_.key_dequant_scale.shape = context_->GetInputShape(KEY_DEQUANT_SCALE_INDEX); + GetOptionalInputParaInfo(); +} + +void LIQInfoParser::GetOutputParaInfo() +{ + opParamInfo_.attenOut.desc = context_->GetOutputDesc(LIGHTNING_INDEXER_QUANT); + opParamInfo_.attenOut.shape = context_->GetOutputShape(LIGHTNING_INDEXER_QUANT); +} + +ge::graphStatus LIQInfoParser::GetAttrParaInfo() +{ + auto attrs = context_->GetAttrs(); + OPS_ERR_IF(attrs == nullptr, OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "attrs got from ge is nullptr"), + return ge::GRAPH_FAILED); + + OPS_LOG_I(context_->GetNodeName(), "GetAttrParaInfo start"); + opParamInfo_.layOutQuery = attrs->GetStr(ATTR_QUERY_LAYOUT_INDEX); + opParamInfo_.layOutKey = attrs->GetStr(ATTR_KEY_LAYOUT_INDEX); + + opParamInfo_.queryQuantMode = attrs->GetAttrPointer(ATTR_QUERY_QUANT_MODE_INDEX); + opParamInfo_.keyQuantMode = attrs->GetAttrPointer(ATTR_KEY_QUANT_MODE_INDEX); + opParamInfo_.layOutQuery = attrs->GetStr(ATTR_QUERY_LAYOUT_INDEX); + opParamInfo_.layOutKey = attrs->GetStr(ATTR_KEY_LAYOUT_INDEX); + opParamInfo_.sparseCount = attrs->GetAttrPointer(ATTR_SPARSE_COUNT_INDEX); + opParamInfo_.sparseMode = attrs->GetAttrPointer(ATTR_SPARSE_MODE_INDEX); + + if (opParamInfo_.layOutQuery != nullptr) { + OPS_LOG_I(context_->GetNodeName(), "layout_query is:%s", opParamInfo_.layOutQuery); + } + if (opParamInfo_.layOutKey != nullptr) { + OPS_LOG_I(context_->GetNodeName(), "layout_key is:%s", opParamInfo_.layOutKey); + } + if (opParamInfo_.sparseCount != nullptr) { + OPS_LOG_I(context_->GetNodeName(), "selscted count is:%d", *opParamInfo_.sparseCount); + } + if (opParamInfo_.sparseMode != nullptr) { + OPS_LOG_I(context_->GetNodeName(), "sparse mode is:%d", *opParamInfo_.sparseMode); + } + if (opParamInfo_.queryQuantMode != nullptr) { + OPS_LOG_I(context_->GetNodeName(), "query_quant_mode mode is:%d", *opParamInfo_.queryQuantMode); + } + if (opParamInfo_.keyQuantMode != nullptr) { + OPS_LOG_I(context_->GetNodeName(), "key_quant_mode mode is:%d", *opParamInfo_.keyQuantMode); + } + OPS_LOG_I(context_->GetNodeName(), "GetAttrParaInfo end"); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::CheckAttrParaInfo() +{ + std::string layout_key(opParamInfo_.layOutKey); + std::string layout_query(opParamInfo_.layOutQuery); + OPS_ERR_IF( + ((std::string(opParamInfo_.layOutKey) == "BNSD") || (std::string(opParamInfo_.layOutKey) == "PA_BBND")), + OPS_LOG_E(opName_, "input attr layout_key only supported PA_BSND, PA_BBND, BSND or TND" + "but now layout_key is %s.", layout_key.c_str()), + return ge::GRAPH_FAILED); + OPS_ERR_IF(((std::string(opParamInfo_.layOutQuery) != "BSND") && (std::string(opParamInfo_.layOutQuery) != "TND")), + OPS_LOG_E(opName_, "input attr layout_query only supported BSND or TND."), return ge::GRAPH_FAILED); + OPS_ERR_IF( + ((std::string(opParamInfo_.layOutKey) != "PA_BSND") && + (std::string(opParamInfo_.layOutQuery)) != (std::string(opParamInfo_.layOutKey))), + OPS_LOG_E(opName_, "outside of PA, input attr layout_query and input attr layout_key must be the same, but now layout_key is %s, layout_query is %s.", + layout_key.c_str(), layout_query.c_str()), return ge::GRAPH_FAILED); + OPS_ERR_IF(!((*opParamInfo_.sparseCount > 0) && (*opParamInfo_.sparseCount <= SPARSE_LIMIT)), + OPS_LOG_E(opName_, "input attr sparse_count must > 0 and <= 2048."), return ge::GRAPH_FAILED); + OPS_ERR_IF(!((*opParamInfo_.sparseMode == 0) || (*opParamInfo_.sparseMode == SPARSE_MODE_LOWER)), + OPS_LOG_E(opName_, "input attr sparse_mode only supported 0 or 3, but now is %u.", + *opParamInfo_.sparseMode), return ge::GRAPH_FAILED); + + OPS_ERR_IF(*opParamInfo_.queryQuantMode != 0, OPS_LOG_E(opName_, "input attr query_quant_mode only supported 0."), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(*opParamInfo_.keyQuantMode != 0, OPS_LOG_E(opName_, "input attr key_quant_mode only supported 0."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetOpParaInfo() +{ + GetInputParaInfo(); + GetOutputParaInfo(); + if (ge::GRAPH_SUCCESS != GetAttrParaInfo()) { + return ge::GRAPH_FAILED; + } + if (ge::GRAPH_SUCCESS != CheckAttrParaInfo()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetAndCheckInOutDataType() +{ + inputQType_ = opParamInfo_.query.desc->GetDataType(); + inputKType_ = opParamInfo_.key.desc->GetDataType(); + weightsType_ = opParamInfo_.weights.desc->GetDataType(); + inputQueryScaleType_ = opParamInfo_.query_dequant_scale.desc->GetDataType(); + inputKeyScaleType_ = opParamInfo_.key_dequant_scale.desc->GetDataType(); + outputType_ = opParamInfo_.attenOut.desc->GetDataType(); + + OPS_ERR_IF(!(inputQType_ == inputKType_), + OPS_LOG_E(opName_, "The data types of the input query and key must be the same, but now is %s, %s respectively.", + inputQType_, inputKType_), + return ge::GRAPH_FAILED); + + OPS_ERR_IF( + !(inputQueryScaleType_ == inputKeyScaleType_), + OPS_LOG_E(opName_, "The data types of the input query_dequant_scale and key_dequant_scale must be the same."), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(inputQType_ != ge::DT_INT8, + OPS_LOG_E(opName_, "The data types of the input query and key must be int8."), return ge::GRAPH_FAILED); + + OPS_ERR_IF(weightsType_ != ge::DT_FLOAT16, + OPS_LOG_E(opName_, "The data types of the input weights must be float16."), return ge::GRAPH_FAILED); + + OPS_ERR_IF( + inputQueryScaleType_ != ge::DT_FLOAT16, + OPS_LOG_E(opName_, "The data types of the input query_dequant_scale and key_dequant_scale must be float16."), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(outputType_ != ge::DT_INT32, + OPS_LOG_E(opName_, "The data types of the output sparse_indices must be int32."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetQueryKeyAndOutLayout() +{ + // 获取query,key的Layout基准值 + const map layoutQueryMap = {{"BSND", DataLayout::BSND}, {"TND", DataLayout::TND}}; + + std::string layout_query(opParamInfo_.layOutQuery); + auto QLayout_ = layoutQueryMap.find(layout_query); + if (QLayout_ != layoutQueryMap.end()) { + qLayout_ = QLayout_->second; + } + + const map layoutKeyMap = { + {"BSND", DataLayout::BSND}, {"TND", DataLayout::TND}, + {"PA_BSND", DataLayout::PA_BSND}, {"PA_BBND", DataLayout::PA_BSND}}; + std::string layout_key(opParamInfo_.layOutKey); + auto KLayout = layoutKeyMap.find(layout_key); + if (KLayout != layoutKeyMap.end()) { + kLayout_ = KLayout->second; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetAndCheckOptionalInput() +{ + if (kLayout_ == DataLayout::PA_BSND) { + OPS_ERR_IF(opParamInfo_.blockTable.tensor == nullptr, + OPS_LOG_E(opName_, "key layout only supported PA_BSND, input block_table must not be null"), + return ge::GRAPH_FAILED); + OPS_ERR_IF( + opParamInfo_.actualSeqLengthsK.tensor == nullptr, + OPS_LOG_E(opName_, "key layout only supported PA_BSND, input actual_seq_lengths_key must not be null"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.blockTable.desc->GetDataType() != ge::DT_INT32, + OPS_LOG_E(opName_, "input block_table data type only support int32"), return ge::GRAPH_FAILED); + } else { + OPS_ERR_IF(opParamInfo_.blockTable.tensor != nullptr, + OPS_LOG_E(opName_, "key layout is not PA_BSND, input block_table must be null"), + return ge::GRAPH_FAILED); + } + + if (kLayout_ == DataLayout::TND) { + OPS_ERR_IF(opParamInfo_.actualSeqLengthsK.tensor == nullptr, + OPS_LOG_E(opName_, "when layout_key is TND, input actual_seq_lengths_key must not be null"), + return ge::GRAPH_FAILED); + } + OPS_ERR_IF(opParamInfo_.actualSeqLengthsK.tensor != nullptr && + opParamInfo_.actualSeqLengthsK.desc->GetDataType() != ge::DT_INT32, + OPS_LOG_E(opName_, "input actual_seq_lengths_key data type only support int32"), + return ge::GRAPH_FAILED); + if (qLayout_ == DataLayout::TND) { + OPS_ERR_IF(opParamInfo_.actualSeqLengthsQ.tensor == nullptr, + OPS_LOG_E(opName_, "when layout_query is TND, input actual_seq_lengths_query must not be null"), + return ge::GRAPH_FAILED); + } + OPS_ERR_IF(opParamInfo_.actualSeqLengthsQ.tensor != nullptr && + opParamInfo_.actualSeqLengthsQ.desc->GetDataType() != ge::DT_INT32, + OPS_LOG_E(opName_, "input actual_seq_lengths_query data type only support int32"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::CheckShapeDim() +{ + OPS_ERR_IF((opParamInfo_.blockTable.tensor != nullptr) && + (opParamInfo_.blockTable.tensor->GetStorageShape().GetDimNum() != DIM_NUM_TWO), + OPS_LOG_E(opName_, "the dim num of block_table's shape should be 2, but now is %u", + opParamInfo_.blockTable.tensor->GetStorageShape().GetDimNum()), return ge::GRAPH_FAILED); + OPS_ERR_IF( + (kLayout_ == DataLayout::PA_BSND) && (opParamInfo_.key.shape->GetStorageShape().GetDimNum() != DIM_NUM_FOUR), + OPS_LOG_E(opName_, "the dim num of key's shape should be 4, but now is %u", + opParamInfo_.key.shape->GetStorageShape().GetDimNum()), return ge::GRAPH_FAILED); + + uint32_t qShapeDim = opParamInfo_.query.shape->GetStorageShape().GetDimNum(); + uint32_t weightsShapeDim = opParamInfo_.weights.shape->GetStorageShape().GetDimNum(); + uint32_t outShapeDim = opParamInfo_.attenOut.shape->GetStorageShape().GetDimNum(); + uint32_t expectShapeDim = DIM_NUM_FOUR; + if (qLayout_ == DataLayout::TND) { + expectShapeDim = DIM_NUM_THREE; + } + OPS_ERR_IF( + qShapeDim != expectShapeDim, + OPS_LOG_E(opName_, "the dim num of query's shape should be %u, but now is %u", expectShapeDim, qShapeDim), + return ge::GRAPH_FAILED); + OPS_ERR_IF(outShapeDim != expectShapeDim, + OPS_LOG_E(opName_, "the dim num of sparse_indices's shape should be %u, but now is %u", expectShapeDim, + outShapeDim), + return ge::GRAPH_FAILED); + OPS_ERR_IF(!(weightsShapeDim == expectShapeDim - 1), + OPS_LOG_E(opName_, "the dim num of weights's shape should be %u, but now is %u", expectShapeDim - 1, + weightsShapeDim), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetN1Size() +{ + if (qLayout_ == DataLayout::BSND) { + n1Size_ = static_cast(opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_TWO)); + } else { + // TND + n1Size_ = static_cast(opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_ONE)); + } + OPS_LOG_I(context_->GetNodeName(), "n1Size is %d", n1Size_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + const std::string &actualSeqLenName) +{ + size = static_cast(tensor->GetShapeSize()); + if (size <= 0) { + OPS_LOG_E(opName_, "%s's shape size is %u, it should be greater than 0.", actualSeqLenName.c_str(), size); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetAndCheckN2Size() +{ + // PA_BSND + if (kLayout_ == DataLayout::TND) { + n2Size_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_ONE)); + } else { + n2Size_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_TWO)); + } + OPS_LOG_I(context_->GetNodeName(), "N2 is %d", n2Size_); + OPS_ERR_IF(n2Size_ != 1, OPS_LOG_E(opName_, "key shape[2] is numhead, only support 1."), return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetGSize() +{ + if (n1Size_ % n2Size_ != 0) { + OPS_LOG_E(opName_, "input query's head_num %u can not be a multiple of key's head_num %u.", n1Size_, n2Size_); + return ge::GRAPH_FAILED; + } + gSize_ = n1Size_ / n2Size_; + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetBatchSize() +{ + // 获取B基准值 + // 1、非TND/NTD时, 以query的batch_size维度为基准; + // 2、TND/NTD时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小 + if (qLayout_ == DataLayout::TND) { + return GetActualSeqLenSize(bSize_, opParamInfo_.actualSeqLengthsQ.tensor, "input actual_seq_lengths_query"); + } else { // BSND + bSize_ = opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_ZERO); + OPS_LOG_I(context_->GetNodeName(), "b: %d, s: %d, n: %d,d :%d", + opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_ZERO), + opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_ONE), + opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_TWO), + opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_THREE)); + return ge::GRAPH_SUCCESS; + } +} + +ge::graphStatus LIQInfoParser::GetHeadDim() +{ + // 以query的D维度为基准 + uint32_t dIndex = DIM_IDX_TWO; + // 根据layout确定D维度在shape中的位置 + switch (qLayout_) { + case DataLayout::TND: + // TND格式: [Total, N, D] -> D是第2维(索引2) + dIndex = DIM_IDX_TWO; + break; + case DataLayout::BSND: + // BSND格式: [Batch, SeqLen, N, D] -> D是第3维(索引3) + dIndex = DIM_IDX_THREE; + break; + default: + OPS_LOG_E(opName_, "unsupported layout for getting head dim."); + return ge::GRAPH_FAILED; + } + headDim_ = opParamInfo_.query.shape->GetStorageShape().GetDim(dIndex); + OPS_ERR_IF(headDim_ != HEAD_DIM_LIMIT, OPS_LOG_E(opName_, "input query's last dim head_dim only support 128, but now is %u.", headDim_), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetS1Size() +{ + if (qLayout_ == DataLayout::BSND) { + s1Size_ = opParamInfo_.query.shape->GetStorageShape().GetDim(1); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetAndCheckBlockSize() +{ + blockSize_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(1)); + OPS_LOG_I(context_->GetNodeName(), "blockSize_ is %d", blockSize_); + + OPS_ERR_IF( + ((blockSize_ % BLOCK_SIZE_FACTOR != 0) || (blockSize_ == 0) || (blockSize_ > BLOCK_SIZE_LIMIT)), + OPS_LOG_E(opName_, "input key's block_size must be a multiple of 16 and belong to (0, 1024], but now is %u.", + blockSize_), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetS2SizeForPageAttention() +{ + if (GetAndCheckBlockSize() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + int32_t blockCount_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(0)); + OPS_ERR_IF((blockCount_ == 0), OPS_LOG_E(opName_, "input key's block_count cannot be 0."), return ge::GRAPH_FAILED); + + maxBlockNumPerBatch_ = opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1); + s2Size_ = maxBlockNumPerBatch_ * blockSize_; + OPS_LOG_I(context_->GetNodeName(), "maxBlockNumPerBatch_ is %d, blockSize_ is %d, s2Size_ is %d", + maxBlockNumPerBatch_, blockSize_, s2Size_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetS2SizeForBatchContinuous() +{ + std::string layout_key(opParamInfo_.layOutKey); + if (kLayout_ == DataLayout::BSND) { + s2Size_ = opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_ONE); + } else if (kLayout_ == DataLayout::TND) { + s2Size_ = opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_ZERO); + } + OPS_ERR_IF((kLayout_ != DataLayout::BSND) && (kLayout_ != DataLayout::TND), + OPS_LOG_E(opName_, "the layout of key is %s, it is unsupported.", layout_key.c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::GetS2Size() +{ + // 获取S2基准值 + // 1、BATCH_CONTINUOUS时, 从key的S轴获取 + // 3、PAGE_ATTENTION时, S2 = block_table.dim1 * block_size + if (kLayout_ == DataLayout::PA_BSND) { + return GetS2SizeForPageAttention(); + } + return GetS2SizeForBatchContinuous(); +} + +ge::graphStatus LIQInfoParser::ValidateInputShapesMatch() +{ + /* + TND: + query [T,N1,D], + key [BlockNum,BlockSize,N2,D], + weight [T,N1], + block_table [BatchSize, BatchMaxBlockNum], + act_seq_k [BatchSize] + act_seq_q [BatchSize], + out [T,N2,topk] + ---------------------- + BSND: + query [BatchSize,S1,N1,D], + key [BlockNum,BlockSize,N2,D], + weight [BatchSize,S1,N1], + block_table [BatchSize, BatchMaxBlockNum], + act_seq_k [BatchSize] + act_seq_q [BatchSize] 可选 + out [BatchSize,S1,N2,topk] + */ + uint32_t queryWeightsN1Dim = 1; + uint32_t outN2Dim = 1; + + if (qLayout_ == DataLayout::TND) { + // -----------------------check BatchSize------------------- + // bSize_ 来源于act_seq_q + OPS_ERR_IF((kLayout_ == DataLayout::PA_BSND) && + ((opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize() != bSize_) || + (opParamInfo_.blockTable.tensor != nullptr && + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) != bSize_)), + OPS_LOG_E( + opName_, + "TND case input actual_seq_lengths_query, actual_seq_lengths_key, block_table dim 0 are %u, %u, %u respectively, they must be same.", + bSize_, opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize(), + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + OPS_ERR_IF((kLayout_ != DataLayout::PA_BSND) && + (opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize() != bSize_), + OPS_LOG_E( + opName_, + "TND case input actual_seq_lengths_query, actual_seq_lengths_key, are %u, %u respectively, they must be same.", + bSize_, opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize()), + return ge::GRAPH_FAILED); + // -----------------------check T------------------- + uint32_t qTsize = opParamInfo_.query.shape->GetStorageShape().GetDim(0); + OPS_ERR_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != qTsize) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != qTsize), + OPS_LOG_E(opName_, + "TND case input query, weights, sparse_indices dim 0 are %u, %u, %u respectively, they must be same.", + qTsize, opParamInfo_.weights.shape->GetStorageShape().GetDim(0), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + } else { + // -----------------------check BatchSize------------------- + // bSize_ 来源于query + OPS_ERR_IF((kLayout_ == DataLayout::PA_BSND) && + ((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != bSize_) || + (opParamInfo_.blockTable.tensor != nullptr && + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) != bSize_) || + (opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize() != bSize_) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != bSize_)), + OPS_LOG_E(opName_, + "BSND case input query, weight, actual_seq_lengths_key, block_table, sparse_indices dim 0 are %u, %u, %u, %u, %u respectively, they must be same.", + bSize_, opParamInfo_.weights.shape->GetStorageShape().GetDim(0), + opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize(), + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + OPS_ERR_IF((kLayout_ != DataLayout::PA_BSND) && + ((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != bSize_) || + (opParamInfo_.actualSeqLengthsK.tensor != nullptr && + opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize() != bSize_) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != bSize_)), + OPS_LOG_E(opName_, + "BSND case input query, weight, actual_seq_lengths_key, sparse_indices dim 0 are %u, %u, %u, %u respectively, they must be same.", + bSize_, opParamInfo_.weights.shape->GetStorageShape().GetDim(0), + opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize(), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + OPS_ERR_IF( + (opParamInfo_.actualSeqLengthsQ.tensor != nullptr) && + (opParamInfo_.actualSeqLengthsQ.tensor->GetShapeSize() != bSize_), + OPS_LOG_E( + opName_, + "BSND case input query, actual_seq_lengths_query dim 0 are %u, %u respectively, they must be same", + bSize_, opParamInfo_.actualSeqLengthsQ.tensor->GetShapeSize()), + return ge::GRAPH_FAILED); + // -----------------------check S1------------------- + OPS_ERR_IF( + (opParamInfo_.weights.shape->GetStorageShape().GetDim(1) != s1Size_) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(1) != s1Size_), + OPS_LOG_E(opName_, "BSND case input query, weight, sparse_indices dim 1 are %u, %u, %u, they must be same.", + s1Size_, opParamInfo_.weights.shape->GetStorageShape().GetDim(1), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(1)), + return ge::GRAPH_FAILED); + queryWeightsN1Dim = DIM_IDX_TWO; + outN2Dim = DIM_IDX_TWO; + } + // -----------------------check N1------------------- + OPS_ERR_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(queryWeightsN1Dim) != n1Size_), + OPS_LOG_E(opName_, "input query, weight shape dim N1 must be same, but now are %u, %u respectively.", + opParamInfo_.weights.shape->GetStorageShape().GetDim(queryWeightsN1Dim), n1Size_), + return ge::GRAPH_FAILED); + // -----------------------check D------------------- + OPS_ERR_IF( + ((kLayout_ != DataLayout::TND && opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_THREE) != headDim_) + || (kLayout_ == DataLayout::TND && opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_TWO) != headDim_)), + OPS_LOG_E(opName_, "input query, key shape last dim must be same, now are %u, %u respectively.", + headDim_, opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_THREE)), + return ge::GRAPH_FAILED); + // -----------------------check N2------------------- + OPS_ERR_IF((opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim) != n2Size_), + OPS_LOG_E(opName_, "input query and output sparse_indices shape n2 dim must be same."), + return ge::GRAPH_FAILED); + // -----------------------check sparse_count------------------- + OPS_ERR_IF((opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim + 1) != *opParamInfo_.sparseCount), + OPS_LOG_E(opName_, "output sparse_indices shape last dim must be same as attr sparse_count."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIQInfoParser::CheckScaleShape() +{ + uint32_t qShapeDim = opParamInfo_.query.shape->GetStorageShape().GetDimNum(); + uint32_t kShapeDim = opParamInfo_.key.shape->GetStorageShape().GetDimNum(); + uint32_t qDequantScaleShapeDim = opParamInfo_.query_dequant_scale.shape->GetStorageShape().GetDimNum(); + uint32_t kDequantScaleShapeDim = opParamInfo_.key_dequant_scale.shape->GetStorageShape().GetDimNum(); + OPS_ERR_IF(qDequantScaleShapeDim != (qShapeDim - 1), + OPS_LOG_E(opName_, "the dim num of query_dequant_scale's shape should be %u, but now is %u", + qShapeDim - 1, qDequantScaleShapeDim), + return ge::GRAPH_FAILED); + OPS_ERR_IF(kDequantScaleShapeDim != (kShapeDim - 1), + OPS_LOG_E(opName_, "the dim num of key_dequant_scale's shape should be %u, but now is %u", kShapeDim - 1, + kDequantScaleShapeDim), + return ge::GRAPH_FAILED); + // check q scale + for (uint32_t i = 0; i < (qShapeDim - 1); i++) { + uint32_t dimValueQueryScale = opParamInfo_.query_dequant_scale.shape->GetStorageShape().GetDim(i); + uint32_t dimValueQuery = opParamInfo_.query.shape->GetStorageShape().GetDim(i); + OPS_ERR_IF(dimValueQueryScale != dimValueQuery, + OPS_LOG_E(opName_, "query_dequant_scale's shape[%u] %u and query's shape[%u] %u is not same", i, + dimValueQueryScale, i, dimValueQuery), + return ge::GRAPH_FAILED); + } + // check k scale + for (uint32_t i = 0; i < (kShapeDim - 1); i++) { + uint32_t dimValueKeyScale = opParamInfo_.key_dequant_scale.shape->GetStorageShape().GetDim(i); + uint32_t dimValueKey = opParamInfo_.key.shape->GetStorageShape().GetDim(i); + OPS_ERR_IF(dimValueKeyScale != dimValueKey, + OPS_LOG_E(opName_, "key_dequant_scale's shape[%u] %u and key's shape[%u] %u is not same", i, + dimValueKeyScale, i, dimValueKey), + return ge::GRAPH_FAILED); + } + + return ge::GRAPH_SUCCESS; +} + +void LIQInfoParser::GenerateInfo(LIQTilingInfo &liqInfo) +{ + liqInfo.opName = opName_; + liqInfo.platformInfo = platformInfo_; + liqInfo.opParamInfo = opParamInfo_; + liqInfo.socVersion = socVersion_; + + liqInfo.bSize = bSize_; + liqInfo.n1Size = n1Size_; + liqInfo.n2Size = n2Size_; + liqInfo.s1Size = s1Size_; + liqInfo.s2Size = s2Size_; + liqInfo.gSize = gSize_; + + liqInfo.inputQType = inputQType_; + liqInfo.inputKType = inputKType_; + liqInfo.outputType = outputType_; + + liqInfo.blockSize = blockSize_; + liqInfo.maxBlockNumPerBatch = maxBlockNumPerBatch_; + + liqInfo.pageAttentionFlag = (kLayout_ == DataLayout::PA_BSND); + liqInfo.sparseMode = *opParamInfo_.sparseMode; + liqInfo.sparseCount = *opParamInfo_.sparseCount; + + liqInfo.inputQLayout = qLayout_; + liqInfo.inputKLayout = kLayout_; +} + +ge::graphStatus LIQInfoParser::ParseAndCheck(LIQTilingInfo &liqInfo) +{ + if (ge::GRAPH_SUCCESS != GetOpName() || ge::GRAPH_SUCCESS != GetNpuInfo() || ge::GRAPH_SUCCESS != GetOpParaInfo() || + ge::GRAPH_SUCCESS != CheckRequiredParaExistence()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetAndCheckInOutDataType() || ge::GRAPH_SUCCESS != GetQueryKeyAndOutLayout() || + ge::GRAPH_SUCCESS != GetAndCheckOptionalInput()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != CheckShapeDim() || ge::GRAPH_SUCCESS != GetN1Size() || + ge::GRAPH_SUCCESS != GetAndCheckN2Size() || ge::GRAPH_SUCCESS != GetGSize()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetBatchSize() || ge::GRAPH_SUCCESS != GetS1Size() || ge::GRAPH_SUCCESS != GetHeadDim() || + ge::GRAPH_SUCCESS != GetS2Size()) { + return ge::GRAPH_FAILED; + } + if (ge::GRAPH_SUCCESS != ValidateInputShapesMatch() || ge::GRAPH_SUCCESS != CheckScaleShape()) { + return ge::GRAPH_FAILED; + } + + GenerateInfo(liqInfo); + + return ge::GRAPH_SUCCESS; +} + +// --------------------------TilingPrepare函数定义------------------------------------- +static ge::graphStatus TilingPrepareForLightningIndexerQuant(gert::TilingParseContext * /* context */) +{ + return ge::GRAPH_SUCCESS; +} + +// --------------------------LightningIndexerQuantTiling类成员函数定义----------------------- +ge::graphStatus LightningIndexerQuantTiling::DoTiling(LIQTilingInfo *tilingInfo) +{ + // -------------set blockdim----------------- + auto ascendcPlatform = platform_ascendc::PlatformAscendC(tilingInfo->platformInfo); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + uint32_t blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum); + context_->SetBlockDim(blockDim); + + // -------------set workspacesize----------------- + constexpr uint32_t MM1_RES_ELEM_SIZE = 4; // 4: fp32 + constexpr uint32_t DOUBLE_BUFFER = 2; // 双Buffer + constexpr uint32_t M_BASE_SIZE = 512; // m轴基本块大小 + constexpr uint32_t S2_BASE_SIZE = 512; // S2轴基本块大小 + constexpr uint32_t V1_RES_ELEM_SIZE = 4; // 4: int32 + constexpr uint32_t V1_RES_ELEM_TYPE = 2; // 保留Index和Value 2种数据 + constexpr uint32_t V1_DECODE_PARAM_ELEM_SIZE = 8; // 8: int64 + constexpr uint32_t V1_DECODE_PARAM_NUM = 16; // Decode参数个数 + constexpr uint32_t V1_DECODE_DATA_NUM = 2; // Decode每个核需要存储头和尾部两块数据 + constexpr uint32_t S1_BASE_SIZE = 8; // S1轴基本块的大小 + constexpr uint32_t TOPK_MAX_SIZE = 2048; // TopK选取个数 + uint32_t workspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize(); + // 主流程需Workspace大小 + uint32_t mm1ResSize = M_BASE_SIZE * S2_BASE_SIZE; + workspaceSize += mm1ResSize * MM1_RES_ELEM_SIZE * DOUBLE_BUFFER * aicNum; + // Decode流程(LD)需要Workspace大小 + // 临时存储Decode中间结果大小: 2(头/尾)*8(s1Base)*2(idx/value)*2048(K)*sizeof(int32)*24=6M + workspaceSize += V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_RES_ELEM_TYPE * TOPK_MAX_SIZE * V1_RES_ELEM_SIZE * aicNum; + // 临时存储Decode中间参数信息大小: 2(头/尾)*8(s1Base)*16(paramNum)*sizeof(int64_t)*24=48k + workspaceSize += V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_DECODE_PARAM_NUM * V1_DECODE_PARAM_ELEM_SIZE * aicNum; + size_t *workSpaces = context_->GetWorkspaceSizes(1); + workSpaces[0] = workspaceSize; + + // -------------set tilingdata----------------- + tilingData_.set_bSize(tilingInfo->bSize); + tilingData_.set_s2Size(tilingInfo->s2Size); + tilingData_.set_s1Size(tilingInfo->s1Size); + tilingData_.set_sparseCount(tilingInfo->sparseCount); + tilingData_.set_gSize(tilingInfo->gSize); + tilingData_.set_blockSize(tilingInfo->blockSize); + tilingData_.set_maxBlockNumPerBatch(tilingInfo->maxBlockNumPerBatch); + tilingData_.set_sparseMode(tilingInfo->sparseMode); + tilingData_.set_usedCoreNum(blockDim); + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + + // -------------set tilingkey----------------- + // DT_Q, DT_KV, DT_OUT, PAGE_ATTENTION, FLASH_DECODE, LAYOUT_T, KV_LAYOUT_T + uint32_t inputQType = static_cast(tilingInfo->inputQType); + uint32_t inputKType = static_cast(tilingInfo->inputKType); + uint32_t outputType = static_cast(tilingInfo->outputType); + uint32_t pageAttentionFlag = static_cast(tilingInfo->pageAttentionFlag); + uint32_t inputQLayout = static_cast(tilingInfo->inputQLayout); + uint32_t inputKLayout = static_cast(tilingInfo->inputKLayout); + uint32_t tilingKey = + GET_TPL_TILING_KEY(inputQType, inputKType, outputType, pageAttentionFlag, inputQLayout, inputKLayout); + context_->SetTilingKey(tilingKey); + + return ge::GRAPH_SUCCESS; +} + +// --------------------------Tiling函数定义--------------------------- +ge::graphStatus TilingForLightningIndexerQuant(gert::TilingContext *context) +{ + OPS_ERR_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("LightningIndexerQuant", "Tiling context is null."), + return ge::GRAPH_FAILED); + LIQTilingInfo liqInfo; + LIQInfoParser LIQInfoParser(context); + if (LIQInfoParser.ParseAndCheck(liqInfo) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + LightningIndexerQuantTiling liqTiling(context); + return liqTiling.DoTiling(&liqInfo); +} + +// --------------------------Tiling及函数TilingPrepare函数注册-------- +IMPL_OP_OPTILING(LightningIndexerQuant) + .Tiling(TilingForLightningIndexerQuant) + .TilingParse(TilingPrepareForLightningIndexerQuant); + +} // namespace optiling diff --git a/csrc/ascend/attention/lightning_indexer_quant/op_host/lightning_indexer_quant_tiling.h b/csrc/ascend/attention/lightning_indexer_quant/op_host/lightning_indexer_quant_tiling.h new file mode 100644 index 000000000..d51a51e7d --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/op_host/lightning_indexer_quant_tiling.h @@ -0,0 +1,234 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_quant_tiling.h + * \brief + */ + +#ifndef LIGHTNING_INDEXER_QUANT_TILING_H_ +#define LIGHTNING_INDEXER_QUANT_TILING_H_ + +#include "error/ops_error.h" +#include "exe_graph/runtime/tiling_context.h" +#include "platform/platform_info.h" +#include "register/op_def_registry.h" +#include "register/tilingdata_base.h" +#include "tiling/platform/platform_ascendc.h" +#include "tiling/tiling_api.h" + +namespace optiling { +// ------------------公共定义-------------------------- +struct TilingRequiredParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; +}; + +struct TilingOptionalParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::Tensor *tensor; +}; + +enum class DataLayout : uint32_t { + BSND = 0, + TND = 1, + PA_BSND = 2 +}; + +// ------------------算子原型索引常量定义---------------- +// Inputs Index +constexpr uint32_t QUERY_INDEX = 0; +constexpr uint32_t KEY_INDEX = 1; +constexpr uint32_t WEIGTHS_INDEX = 2; +constexpr uint32_t QUERY_DEQUANT_SCALE_INDEX = 3; +constexpr uint32_t KEY_DEQUANT_SCALE_INDEX = 4; +constexpr uint32_t ACTUAL_SEQ_Q_INDEX = 5; +constexpr uint32_t ACTUAL_SEQ_K_INDEX = 6; +constexpr uint32_t BLOCK_TABLE_INDEX = 7; +constexpr uint32_t LIGHTNING_INDEXER_QUANT = 0; +// Attributes Index +constexpr uint32_t ATTR_QUERY_QUANT_MODE_INDEX = 0; +constexpr uint32_t ATTR_KEY_QUANT_MODE_INDEX = 1; +constexpr uint32_t ATTR_QUERY_LAYOUT_INDEX = 2; +constexpr uint32_t ATTR_KEY_LAYOUT_INDEX = 3; +constexpr uint32_t ATTR_SPARSE_COUNT_INDEX = 4; +constexpr uint32_t ATTR_SPARSE_MODE_INDEX = 5; +// Dim Index +constexpr uint32_t DIM_IDX_ZERO = 0; +constexpr uint32_t DIM_IDX_ONE = 1; +constexpr uint32_t DIM_IDX_TWO = 2; +constexpr uint32_t DIM_IDX_THREE = 3; +// Dim Num +constexpr uint32_t DIM_NUM_TWO = 2; +constexpr uint32_t DIM_NUM_THREE = 3; +constexpr uint32_t DIM_NUM_FOUR = 4; +// 入参限制常量 +constexpr uint32_t HEAD_DIM_LIMIT = 128; +constexpr uint32_t SPARSE_LIMIT = 2048; +constexpr uint32_t G_SIZE_LIMIT = 64; +constexpr uint32_t BLOCK_SIZE_LIMIT = 1024; +constexpr uint32_t BLOCK_SIZE_FACTOR = 16; +constexpr uint32_t SPARSE_MODE_LOWER = 3; + +// -----------算子TilingData定义--------------- +BEGIN_TILING_DATA_DEF(LIQTilingData) +TILING_DATA_FIELD_DEF(uint32_t, bSize) +TILING_DATA_FIELD_DEF(uint32_t, n2Size) +TILING_DATA_FIELD_DEF(uint32_t, gSize) +TILING_DATA_FIELD_DEF(uint32_t, s1Size) +TILING_DATA_FIELD_DEF(uint32_t, s2Size) +TILING_DATA_FIELD_DEF(uint32_t, sparseCount) +TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum) +TILING_DATA_FIELD_DEF(uint32_t, blockSize) +TILING_DATA_FIELD_DEF(uint32_t, maxBlockNumPerBatch) +TILING_DATA_FIELD_DEF(uint32_t, sparseMode) +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(LightningIndexerQuant, LIQTilingData) + +// -----------算子CompileInfo定义------------------- +struct LIQCompileInfo {}; + +// -----------算子Tiling入参结构体定义--------------- +struct LIQParaInfo { + TilingRequiredParaInfo query = {nullptr, nullptr}; + TilingRequiredParaInfo key = {nullptr, nullptr}; + TilingRequiredParaInfo weights = {nullptr, nullptr}; + TilingRequiredParaInfo query_dequant_scale = {nullptr, nullptr}; + TilingRequiredParaInfo key_dequant_scale = {nullptr, nullptr}; + TilingOptionalParaInfo actualSeqLengthsQ = {nullptr, nullptr}; + TilingOptionalParaInfo actualSeqLengthsK = {nullptr, nullptr}; + TilingOptionalParaInfo blockTable = {nullptr, nullptr}; + TilingRequiredParaInfo attenOut = {nullptr, nullptr}; + + const int32_t *queryQuantMode = nullptr; + const int32_t *keyQuantMode = nullptr; + const char *layOutQuery = nullptr; + const char *layOutKey = nullptr; + const int32_t *blockSize = nullptr; + const int32_t *sparseMode = nullptr; + const int32_t *sparseCount = nullptr; +}; + +// -----------算子Tiling入参信息类--------------- +class LIQTilingInfo { +public: + const char *opName = nullptr; + fe::PlatFormInfos *platformInfo = nullptr; + LIQParaInfo opParamInfo; + // Base Param + platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND910B; + uint32_t bSize = 0; + uint32_t n1Size = 0; + uint32_t n2Size = 0; + uint32_t s1Size = 0; + int64_t s2Size = 0; + uint32_t qkHeadDim = 0; + uint32_t gSize = 0; + // PageAttention + bool pageAttentionFlag = false; + int32_t blockSize = 0; + uint32_t maxBlockNumPerBatch = 0; + // Mask + int32_t sparseMode = 0; + // Others Flag + uint32_t sparseCount = 0; + // DType + ge::DataType inputQType = ge::DT_FLOAT16; + ge::DataType inputKType = ge::DT_FLOAT16; + ge::DataType outputType = ge::DT_INT32; + // Layout + DataLayout inputQLayout = DataLayout::BSND; + DataLayout inputKLayout = DataLayout::PA_BSND; +}; + +// -----------算子Tiling入参信息解析及Check类--------------- +class LIQInfoParser { +public: + explicit LIQInfoParser(gert::TilingContext *context) : context_(context) {} + ~LIQInfoParser() = default; + + ge::graphStatus CheckRequiredInOutExistence() const; + ge::graphStatus CheckRequiredAttrExistence() const; + ge::graphStatus CheckRequiredParaExistence() const; + ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + const std::string &actualSeqLenName); + ge::graphStatus GetOpName(); + ge::graphStatus GetNpuInfo(); + void GetOptionalInputParaInfo(); + void GetInputParaInfo(); + void GetOutputParaInfo(); + ge::graphStatus GetAttrParaInfo(); + ge::graphStatus CheckAttrParaInfo(); + ge::graphStatus GetOpParaInfo(); + ge::graphStatus ValidateInputShapesMatch(); + ge::graphStatus CheckScaleShape(); + ge::graphStatus GetAndCheckInOutDataType(); + ge::graphStatus GetBatchSize(); + ge::graphStatus GetHeadDim(); + ge::graphStatus GetS1Size(); + ge::graphStatus GetAndCheckOptionalInput(); + ge::graphStatus CheckShapeDim(); + ge::graphStatus GetAndCheckBlockSize(); + ge::graphStatus GetS2SizeForPageAttention(); + ge::graphStatus GetS2SizeForBatchContinuous(); + ge::graphStatus GetS2Size(); + ge::graphStatus GetQueryKeyAndOutLayout(); + ge::graphStatus GetN1Size(); + ge::graphStatus GetAndCheckN2Size(); + ge::graphStatus GetGSize(); + ge::graphStatus GetAttenMaskInfo(); + ge::graphStatus GetActualSeqInfo(); + void GenerateInfo(LIQTilingInfo &liqInfo); + ge::graphStatus ParseAndCheck(LIQTilingInfo &liqInfo); + +public: + gert::TilingContext *context_ = nullptr; + const char *opName_; + fe::PlatFormInfos *platformInfo_; + LIQParaInfo opParamInfo_; + + // BaseParams + uint32_t bSize_ = 0; + uint32_t n1Size_ = 0; + uint32_t n2Size_ = 0; + uint32_t gSize_ = 0; + uint32_t s1Size_ = 0; + int64_t s2Size_ = 0; + uint32_t headDim_ = 0; + // Layout + DataLayout qLayout_ = DataLayout::BSND; + DataLayout kLayout_ = DataLayout::PA_BSND; + // PageAttention + uint32_t maxBlockNumPerBatch_ = 0; + int32_t blockSize_ = 0; + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B; + ge::DataType inputQType_ = ge::DT_FLOAT16; + ge::DataType inputKType_ = ge::DT_FLOAT16; + ge::DataType weightsType_ = ge::DT_FLOAT16; + ge::DataType inputQueryScaleType_ = ge::DT_FLOAT16; + ge::DataType inputKeyScaleType_ = ge::DT_FLOAT16; + ge::DataType blockTableType_ = ge::DT_FLOAT16; + ge::DataType inputKRopeType_ = ge::DT_FLOAT16; + ge::DataType outputType_ = ge::DT_FLOAT16; +}; + +// ---------------算子Tiling类--------------- +class LightningIndexerQuantTiling { +public: + explicit LightningIndexerQuantTiling(gert::TilingContext *context) : context_(context) {}; + ge::graphStatus DoTiling(LIQTilingInfo *tilingInfo); + +private: + gert::TilingContext *context_ = nullptr; + LIQTilingData tilingData_; +}; + +} // namespace optiling +#endif // LIGHTNING_INDEXER_QUANT_TILING_H_ \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant.cpp b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant.cpp new file mode 100644 index 000000000..a9513daf4 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant.cpp @@ -0,0 +1,50 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_quant.cpp + * \brief + */ + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" +#include "lightning_indexer_quant_kernel.h" +#include "lightning_indexer_quant_template_tiling_key.h" + +using namespace LIQKernel; + +#define INVOKE_LI_NO_KFC_OP_IMPL(templateClass, ...) \ + do { \ + templateClass> op; \ + GET_TILING_DATA_WITH_STRUCT(LIQTilingData, tiling_data_in, tiling); \ + const LIQTilingData *__restrict tiling_data = &tiling_data_in; \ + op.Init(query, key, weights, queryScale, keyScale, actualSeqLengthsQ, actualSeqLengthsK, blocktable, \ + sparseIndices, user, tiling_data, &tPipe); \ + op.Process(); \ + } while (0) + +template +__global__ __aicore__ void lightning_indexer_quant(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK, + __gm__ uint8_t *blocktable, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *workspace, __gm__ uint8_t *tiling) +{ +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) || (__CCE_AICORE__ == 200) + +#else + TPipe tPipe; + __gm__ uint8_t *user = GetUserWorkspace(workspace); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + + INVOKE_LI_NO_KFC_OP_IMPL(LIQPreload, int8_t, int8_t, int32_t, + PAGE_ATTENTION, LI_LAYOUT(Q_LAYOUT_T), LI_LAYOUT(K_LAYOUT_T)); +#endif +} diff --git a/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_common.h b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_common.h new file mode 100644 index 000000000..a0f0eb7fe --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_common.h @@ -0,0 +1,146 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_quant_common.h + * \brief + */ +#ifndef LIGHTNING_INDEXER_QUANT_COMMON_H +#define LIGHTNING_INDEXER_QUANT_COMMON_H + +namespace LIQCommon { + +// 与tiling的layout保持一致 +enum class LI_LAYOUT : uint32_t { + BSND = 0, + TND = 1, + PA_BSND = 2 +}; + +template +struct LIQType { + using queryType = Q_T; + using keyType = K_T; + using outputType = OUT_T; + static constexpr bool pageAttention = PAGE_ATTENTION; + static constexpr LI_LAYOUT layout = Q_LAYOUT_T; + static constexpr LI_LAYOUT keyLayout = K_LAYOUT_T; +}; + +struct RunInfo { + uint32_t loop; + uint32_t bN2Idx; + uint32_t bIdx; + uint32_t n2Idx = 0; + uint32_t gS1Idx; + uint32_t s2Idx; + + uint32_t actS1Size = 1; + uint32_t actS2Size = 1; + uint32_t actMBaseSize; + uint32_t actualSingleProcessSInnerSize; + uint32_t actualSingleProcessSInnerSizeAlign; + + uint64_t tensorQueryOffset; + uint64_t tensorKeyOffset; + uint64_t tensorKeyScaleOffset; + uint64_t tensorWeightsOffset; + uint64_t indiceOutOffset; + + bool isFirstS2InnerLoop; + bool isLastS2InnerLoop; + bool isAllLoopEnd = false; + bool isValid = false; +}; + +struct ConstInfo { + // CUBE与VEC核间同步的模式 + static constexpr uint32_t FIA_SYNC_MODE2 = 2; + // BUFFER的字节数 + static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; + static constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64; + static constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256; + static constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512; + static constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024; + static constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048; + static constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096; + static constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192; + static constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384; + static constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768; + // 无效索引 + static constexpr int INVALID_IDX = -1; + + // CUBE和VEC的核间同步EventID + uint32_t syncC1V1 = 0U; + uint32_t syncC1V0 = 2U; + uint32_t syncV1C1 = 0U; + uint32_t syncV0C1 = 1U; + + // 基本块大小 + uint32_t mBaseSize = 1ULL; + uint32_t s1BaseSize = 1ULL; + uint32_t s2BaseSize = 1ULL; + + uint64_t batchSize = 0ULL; + uint64_t gSize = 0ULL; + uint64_t qHeadNum = 0ULL; + uint64_t kHeadNum; + uint64_t headDim; + uint64_t sparseCount; // topK选取大小 + uint64_t kSeqSize = 0ULL; // kv最大S长度 + uint64_t qSeqSize = 1ULL; // q最大S长度 + uint32_t kCacheBlockSize = 0; // PA场景的block size + uint32_t maxBlockNumPerBatch = 0; // PA场景的最大单batch block number + LI_LAYOUT outputLayout; // 输出的格式 + bool attenMaskFlag = false; + + uint32_t actualLenQDims = 0U; // query的actualSeqLength 的维度 + uint32_t actualLenDims = 0U; // KV 的actualSeqLength 的维度 + bool isAccumSeqS1 = false; // 是否累加模式 + bool isAccumSeqS2 = false; // 是否累加模式 +}; + +struct SplitCoreInfo { + uint32_t s2Start = 0U; // S2的起始位置 + uint32_t s2End = 0U; // S2循环index上限 + uint32_t bN2Start = 0U; + uint32_t bN2End = 0U; + uint32_t gS1Start = 0U; + uint32_t gS1End = 0U; + bool isLD = false; // 当前核是否需要进行Decode归约任务 +}; + +template +__aicore__ inline T Align(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd) * (rnd))); +} + +template +__aicore__ inline T1 Min(T1 a, T2 b) +{ + return (a > b) ? (b) : (a); +} + +template +__aicore__ inline T1 Max(T1 a, T2 b) +{ + return (a > b) ? (a) : (b); +} + +template +__aicore__ inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd))); +} +} // namespace LIQCommon + +#endif // LIGHTNING_INDEXER_QUANT_COMMON_H \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_kernel.h b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_kernel.h new file mode 100644 index 000000000..723255d3c --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_kernel.h @@ -0,0 +1,714 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_quant_kernel.h + * \brief + */ + +#ifndef LIGHTNING_INDEXER_QUANT_KERNEL_H +#define LIGHTNING_INDEXER_QUANT_KERNEL_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "lightning_indexer_quant_common.h" +#include "lightning_indexer_quant_service_vector.h" +#include "lightning_indexer_quant_service_cube.h" + +namespace LIQKernel { +using namespace LIQCommon; +using namespace LIQServiceVec; +using namespace matmul; +using AscendC::CacheMode; +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +// 由于S2循环前，RunInfo还没有赋值，使用TempLoopInfo临时存放B、N、S1轴相关的信息；同时减少重复计算 +struct TempLoopInfo { + uint32_t bN2Idx = 0; + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint32_t gS1Idx = 0U; + uint32_t gS1LoopEnd = 0U; // gS1方向循环的结束Idx + uint32_t s2LoopEnd = 0U; // S2方向循环的结束Idx + uint32_t actS1Size = 1ULL; // 当前Batch循环处理的S1轴的实际大小 + uint32_t actS2Size = 0ULL; + bool curActSeqLenIsZero = false; + bool needDealActS1LessThanS1 = false; // S1的实际长度小于shape的S1长度时，是否需要清理输出 + uint32_t actMBaseSize = 0U; // m轴(gS1)方向实际大小 + uint32_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小 + uint32_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小 +}; + +template +class LIQPreload { +public: + __aicore__ inline LIQPreload(){}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale, __gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengthsK, __gm__ uint8_t *blockTable, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *workspace, + const LIQTilingData *__restrict tiling, TPipe *tPipe); + __aicore__ inline void Process(); + + // =================================类型定义区================================= + using Q_T = typename LIQT::queryType; + using K_T = typename LIQT::keyType; + using OUT_T = typename LIQT::outputType; + static constexpr bool PAGE_ATTENTION = LIQT::pageAttention; + static constexpr LI_LAYOUT Q_LAYOUT_T = LIQT::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = LIQT::keyLayout; + + using MM1_OUT_T = float; + + LIQMatmul matmulService; + LIQVector vectorService; + + // =================================常量区================================= + static constexpr uint32_t SYNC_C1_V1_FLAG = 4; + static constexpr uint32_t SYNC_V1_C1_FLAG = 5; + + static constexpr uint32_t M_BASE_SIZE = 256; + static constexpr uint32_t S2_BASE_SIZE = 2048; + static constexpr uint32_t HEAD_DIM = 128; + static constexpr uint32_t K_HEAD_NUM = 1; + static constexpr uint32_t GM_ALIGN_BYTES = 512; + static constexpr uint32_t LI_QUANT_PRELOAD_TASK_CACHE_SIZE = 2; + + static constexpr int64_t LD_PREFETCH_LEN = 2; + // for workspace double + static constexpr uint32_t WS_DOBULE = 2; + +protected: + TPipe *pipe = nullptr; + + // offset + uint64_t queryCoreOffset = 0ULL; + uint64_t keyCoreOffset = 0ULL; + uint64_t keyScaleCoreOffset = 0ULL; + uint64_t weightsCoreOffset = 0ULL; + uint64_t indiceOutCoreOffset = 0ULL; + + // ================================Global Buffer区================================= + GlobalTensor queryGm; + GlobalTensor keyGm; + GlobalTensor weightsGm; + + GlobalTensor indiceOutGm; + GlobalTensor blockTableGm; + + GlobalTensor actualSeqLengthsGmQ; + GlobalTensor actualSeqLengthsGm; + + // ================================类成员变量==================================== + // aic、aiv核信息 + uint32_t tmpBlockIdx = 0U; + uint32_t aiCoreIdx = 0U; + uint32_t usedCoreNum = 0U; + + LIQCommon::ConstInfo constInfo{}; + TempLoopInfo tempLoopInfo{}; + LIQCommon::SplitCoreInfo splitCoreInfo{}; + + // ================================Init functions================================== + __aicore__ inline void InitTilingData(const LIQTilingData *__restrict tilingData); + __aicore__ inline void InitBuffers(); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK); + // ================================Split Core================================ + __aicore__ inline void SplitCore(uint32_t curCoreIdx, uint32_t &coreNum, LIQCommon::SplitCoreInfo &info); + __aicore__ inline uint32_t GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, uint32_t actS2Size); + __aicore__ inline uint32_t GetTotalBaseBlockNum(); + // ================================Process functions================================ + __aicore__ inline void ProcessMain(); + __aicore__ inline void ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, + LIQCommon::RunInfo runInfo[LI_QUANT_PRELOAD_TASK_CACHE_SIZE]); + __aicore__ inline void ProcessDecode(); + __aicore__ inline void ProcessInvalid(); + // ================================Params Calc===================================== + __aicore__ inline void CalcGS1LoopParams(uint32_t bN2Idx); + __aicore__ inline void GetBN2Idx(uint32_t bN2Idx); + __aicore__ inline uint32_t GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, uint32_t defaultSeqLen); + __aicore__ inline void GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, uint32_t &actS2Size); + __aicore__ inline void CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx); + __aicore__ inline void CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, LIQCommon::RunInfo &runInfo); + __aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start); +}; + +template +__aicore__ inline void LIQPreload::InitTilingData(const LIQTilingData *__restrict tilingData) +{ + usedCoreNum = tilingData->usedCoreNum; + constInfo.batchSize = tilingData->bSize; + constInfo.qHeadNum = constInfo.gSize = tilingData->gSize; + constInfo.kSeqSize = tilingData->s2Size; + constInfo.qSeqSize = tilingData->s1Size; + constInfo.attenMaskFlag = (tilingData->sparseMode == 3); + constInfo.kCacheBlockSize = tilingData->blockSize; + constInfo.maxBlockNumPerBatch = tilingData->maxBlockNumPerBatch; + constInfo.sparseCount = tilingData->sparseCount; + constInfo.outputLayout = Q_LAYOUT_T; // 输出和输入形状一致 + if (Q_LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS1 = true; + } + if (K_LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS2 = true; + } + + constInfo.kHeadNum = K_HEAD_NUM; + constInfo.headDim = HEAD_DIM; + + constInfo.mBaseSize = M_BASE_SIZE; + constInfo.s2BaseSize = S2_BASE_SIZE; + constInfo.s1BaseSize = (constInfo.mBaseSize + constInfo.gSize - 1) / constInfo.gSize; +} + +template +__aicore__ inline void LIQPreload::InitBuffers() +{ + if ASCEND_IS_AIV { + vectorService.InitBuffers(pipe); + } else { + matmulService.InitBuffers(pipe); + } +} + +template +__aicore__ inline void LIQPreload::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengthsK) +{ + if (actualSeqLengthsQ == nullptr) { + constInfo.actualLenQDims = 0; + } else { + constInfo.actualLenQDims = constInfo.batchSize; + actualSeqLengthsGmQ.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsQ, constInfo.actualLenQDims); + } + if (actualSeqLengthsK == nullptr) { + constInfo.actualLenDims = 0; + } else { + constInfo.actualLenDims = constInfo.batchSize; + actualSeqLengthsGm.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsK, constInfo.actualLenDims); + } +} + +template +__aicore__ inline uint32_t LIQPreload::GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, + uint32_t defaultSeqLen) +{ + if (actualLenDims == 0) { + return defaultSeqLen; + } else if (isAccumSeq && bIdx > 0) { + return actualSeqLengthsGm.GetValue(bIdx) - actualSeqLengthsGm.GetValue(bIdx - 1); + } else { + return actualSeqLengthsGm.GetValue(bIdx); + } +} + +template +__aicore__ inline void LIQPreload::GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, uint32_t &actS2Size) +{ + actS1Size = GetActualSeqLen(bIdx, constInfo.actualLenQDims, constInfo.isAccumSeqS1, actualSeqLengthsGmQ, + constInfo.qSeqSize); + actS2Size = + GetActualSeqLen(bIdx, constInfo.actualLenDims, constInfo.isAccumSeqS2, actualSeqLengthsGm, constInfo.kSeqSize); +} + +template +__aicore__ inline uint32_t LIQPreload::GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, + uint32_t actS2Size) +{ + if (actS2Size == 0) { + return 0; + } + uint32_t s1Offset = constInfo.s1BaseSize * s1gIdx; + int32_t validS2LenBase = static_cast(actS2Size) - static_cast(actS1Size); + int32_t validS2Len = s1Offset + validS2LenBase + constInfo.s1BaseSize; + validS2Len = Min(validS2Len, static_cast(actS2Size)); + validS2Len = Max(validS2Len, 1); + return (validS2Len + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; +} + +template +__aicore__ inline uint32_t LIQPreload::GetTotalBaseBlockNum() +{ + uint32_t totalBlockNum = 0; + uint32_t actS1Size, actS2Size; + uint32_t s1GBaseNum, s2BaseNum; + for (uint32_t bIdx = 0; bIdx < constInfo.batchSize; bIdx++) { + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size); + s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize); + if (!constInfo.attenMaskFlag) { + s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize); + totalBlockNum += s1GBaseNum * s2BaseNum * constInfo.kHeadNum; + continue; + } + for (uint32_t s1gIdx = 0; s1gIdx < s1GBaseNum; s1gIdx++) { + s2BaseNum = GetS2BaseBlockNumOnMask(s1gIdx, actS1Size, actS2Size); + totalBlockNum += s2BaseNum * constInfo.kHeadNum; + } + } + return totalBlockNum; +} + +// 多核版本，双闭区间。基本原则：计算每个核最少处理的块数, 剩余的部分前面的核每个核多处理一块 +template +__aicore__ void inline LIQPreload::SplitCore(uint32_t curCoreIdx, uint32_t &coreNum, + LIQCommon::SplitCoreInfo &info) +{ + uint32_t totalBlockNum = GetTotalBaseBlockNum(); + uint32_t minBlockPerCore = totalBlockNum / coreNum; + uint32_t deal1MoreBlockCoreNum = totalBlockNum % coreNum; + uint32_t coreIdx = 0; + uint32_t lastGS1RemainBlockCnt = 0; + uint32_t coreDealBlockCnt = coreIdx < deal1MoreBlockCoreNum ? minBlockPerCore + 1 : minBlockPerCore; + coreNum = minBlockPerCore == 0 ? deal1MoreBlockCoreNum : coreNum; + + bool findLastCoreEnd = true; + uint32_t actS1Size, actS2Size; + uint32_t s1GBaseNum, s2BaseNum; + for (uint32_t bN2Idx = 0; bN2Idx < constInfo.batchSize * constInfo.kHeadNum; bN2Idx++) { + uint32_t bIdx = bN2Idx / constInfo.kHeadNum; + if (bN2Idx % constInfo.kHeadNum == 0) { + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size); + s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize); + s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize); + } + if constexpr (Q_LAYOUT_T == LI_LAYOUT::BSND) { + if (findLastCoreEnd && (s1GBaseNum == 0U || s2BaseNum == 0U)) { + info.bN2Start = bN2Idx; + info.gS1Start = 0; + info.s2Start = 0; + findLastCoreEnd = false; + } + } + for (uint32_t gS1Idx = 0; gS1Idx < s1GBaseNum; gS1Idx++) { + if (constInfo.attenMaskFlag) { + s2BaseNum = GetS2BaseBlockNumOnMask(gS1Idx, actS1Size, actS2Size); + } + if (findLastCoreEnd && s2BaseNum == 0U) { + info.bN2Start = bN2Idx; + info.gS1Start = gS1Idx; + info.s2Start = 0; + findLastCoreEnd = false; + } + for (uint32_t s2Idx = 0; s2Idx < s2BaseNum;) { + if (findLastCoreEnd) { + info.bN2Start = bN2Idx; + info.gS1Start = gS1Idx; + info.s2Start = s2Idx; + findLastCoreEnd = false; + } + uint32_t s2RemainBaseNum = s2BaseNum - s2Idx; + if (lastGS1RemainBlockCnt + s2RemainBaseNum >= coreDealBlockCnt) { + info.bN2End = bN2Idx; + info.gS1End = gS1Idx; + info.s2End = s2Idx + coreDealBlockCnt - lastGS1RemainBlockCnt - 1; + + if (coreIdx == curCoreIdx) { + // S2被切N核，那么只有第一个核需要处理LD，其他核不用 + if (s2Idx == 0 && info.s2End + 1 < s2BaseNum) { + info.isLD = true; + } + // 最后一个核处理的不是最后一个Batch，表明后面的Batch为空块(S2=0), 调整终点坐标以便清理输出 + if (coreIdx == coreNum - 1 && info.bN2End != constInfo.batchSize - 1) { + info.bN2End = constInfo.batchSize - 1; + info.gS1End = 0; + info.s2End = 0; + } + return; + } + coreIdx++; + findLastCoreEnd = true; + s2Idx = info.s2End + 1; + lastGS1RemainBlockCnt = 0; + coreDealBlockCnt = coreIdx < deal1MoreBlockCoreNum ? minBlockPerCore + 1 : minBlockPerCore; + } else { + lastGS1RemainBlockCnt += s2RemainBaseNum; + break; + } + } + } + } +} + +template +__aicore__ inline void LIQPreload::DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start) +{ + if ASCEND_IS_AIV { + if (constInfo.outputLayout == LI_LAYOUT::TND) { + uint32_t tSize = actualSeqLengthsGmQ.GetValue(constInfo.batchSize - 1); + uint32_t tBase = bIdx == 0 ? 0 : actualSeqLengthsGmQ.GetValue(bIdx - 1); + uint32_t s1Count = tempLoopInfo.actS1Size; + + for (uint32_t s1Idx = s1Start; s1Idx < s1Count; s1Idx++) { + uint64_t indiceOutOffset = + (tBase + s1Idx) * constInfo.kHeadNum * constInfo.sparseCount + // T轴、s1轴偏移 + n2Idx * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } else if (constInfo.outputLayout == LI_LAYOUT::BSND) { + for (uint32_t s1Idx = s1Start; s1Idx < constInfo.qSeqSize; s1Idx++) { + // B,S1,N2,K + uint64_t indiceOutOffset = bIdx * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount + + s1Idx * constInfo.kHeadNum * constInfo.sparseCount + // B轴、S1轴偏移 + n2Idx * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } + } +} + +template +__aicore__ inline void LIQPreload::Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK, + __gm__ uint8_t *blockTable, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *workspace, const LIQTilingData *__restrict tiling, + TPipe *tPipe) +{ + if ASCEND_IS_AIV { + tmpBlockIdx = GetBlockIdx(); // vec:0-47 + aiCoreIdx = tmpBlockIdx / 2; + } else { + tmpBlockIdx = GetBlockIdx(); // cube:0-23 + aiCoreIdx = tmpBlockIdx; + } + + InitTilingData(tiling); + InitActualSeqLen(actualSeqLengthsQ, actualSeqLengthsK); + + // 计算分核 + SplitCore(aiCoreIdx, usedCoreNum, splitCoreInfo); + + pipe = tPipe; + // workspace 内存排布 + // |mm1ResGm(存S)|vec1ResGm(存LD中间结果)|vec1ParamGm(存LD参数) + // |Core0_mm1ResDB0-Core0_mm1ResDB1-Core1_mm1ResDB0....Core23_mm1ResDB0-Core23_mm1ResDB1|Core0_vec1Res... + uint64_t offset = 0; + + // mm1开DoubleBuffer + GlobalTensor mm1ResGm; // 存放S + uint64_t singleCoreMm1ResSize = WS_DOBULE * constInfo.s1BaseSize * constInfo.s2BaseSize * sizeof(MM1_OUT_T); + mm1ResGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + aiCoreIdx * singleCoreMm1ResSize)); + offset += GetBlockNum() * singleCoreMm1ResSize; + + // ld流程需要ws大小: [aicnum, 2, CeilDiv(constInfo.mBaseSize, constInfo.gSize), topkOut_*2] + // (aic, 8, 2, 2, 2048) + // (aic, s1_cube, 头尾, idx/value, K) + GlobalTensor vec1ResGm; // 存放TopK计算中间结果 + vec1ResGm.SetGlobalBuffer((__gm__ float *)(workspace + offset)); + offset += GetBlockNum() * constInfo.s1BaseSize * WS_DOBULE * WS_DOBULE * BASE_TOPK * sizeof(float); + + // (aic, 8, 2, 16) + // (aic, s1_cube, 头尾，16ele) + GlobalTensor vec1ParamGm; // 存放LD参数信息 + vec1ParamGm.SetGlobalBuffer((__gm__ int64_t *)(workspace + offset)); + offset += GetBlockNum() * constInfo.s1BaseSize * WS_DOBULE * LD_PARAM_NUM * sizeof(int64_t); + + GlobalTensor weightWorkspaceGm; // v1阶段处理w*scale后的结果 + uint64_t weightMemSize = BLOCK_CUBE * constInfo.mBaseSize * WS_DOBULE * sizeof(half); + weightWorkspaceGm.SetGlobalBuffer((__gm__ half *)(workspace + offset + aiCoreIdx * weightMemSize)); + offset += GetBlockNum() * weightMemSize; + + GlobalTensor qScaleGm; + GlobalTensor kScaleGm; + if ASCEND_IS_AIV { + vectorService.InitParams(constInfo, tiling); + indiceOutGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices); + weightsGm.SetGlobalBuffer((__gm__ half *)weights); + qScaleGm.SetGlobalBuffer((__gm__ half *)queryScale); + kScaleGm.SetGlobalBuffer((__gm__ half *)keyScale); + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + vectorService.InitVecInputTensor(weightsGm, qScaleGm, kScaleGm, indiceOutGm, blockTableGm); + vectorService.InitVecWorkspaceTensor(weightWorkspaceGm, mm1ResGm, vec1ResGm, vec1ParamGm); + } else { + matmulService.InitParams(constInfo); + queryGm.SetGlobalBuffer((__gm__ Q_T *)query); + if constexpr (PAGE_ATTENTION) { + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + } + keyGm.SetGlobalBuffer((__gm__ K_T *)key); + matmulService.InitMm1GlobalTensor(blockTableGm, keyGm, queryGm, mm1ResGm, weightWorkspaceGm); + } + InitBuffers(); +} + +template +__aicore__ inline void LIQPreload::GetBN2Idx(uint32_t bN2Idx) +{ + tempLoopInfo.bN2Idx = bN2Idx; + tempLoopInfo.bIdx = bN2Idx / constInfo.kHeadNum; + tempLoopInfo.n2Idx = bN2Idx % constInfo.kHeadNum; +} + +template +__aicore__ inline void LIQPreload::CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx) +{ + tempLoopInfo.gS1Idx = gS1LoopIdx; + tempLoopInfo.actMBaseSize = constInfo.mBaseSize; + uint32_t remainedGS1Size = tempLoopInfo.actS1Size * constInfo.gSize - tempLoopInfo.gS1Idx * constInfo.mBaseSize; + if (remainedGS1Size <= constInfo.mBaseSize && remainedGS1Size > 0) { + tempLoopInfo.actMBaseSize = tempLoopInfo.mBasicSizeTail; + } + + bool isEnd = (bN2LoopIdx == splitCoreInfo.bN2End) && (gS1LoopIdx == splitCoreInfo.gS1End); + uint32_t s2BlockNum; + if (constInfo.attenMaskFlag) { + s2BlockNum = GetS2BaseBlockNumOnMask(gS1LoopIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2Size); + } else { + s2BlockNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + } + tempLoopInfo.s2LoopEnd = isEnd ? splitCoreInfo.s2End : s2BlockNum - 1; +} + +template +__aicore__ inline void LIQPreload::CalcGS1LoopParams(uint32_t bN2LoopIdx) +{ + GetBN2Idx(bN2LoopIdx); + GetS1S2ActualSeqLen(tempLoopInfo.bIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2Size); + if ((tempLoopInfo.actS2Size == 0) || (tempLoopInfo.actS1Size == 0)) { + tempLoopInfo.curActSeqLenIsZero = true; + return; + } + tempLoopInfo.curActSeqLenIsZero = false; + tempLoopInfo.s2BasicSizeTail = tempLoopInfo.actS2Size % constInfo.s2BaseSize; + tempLoopInfo.s2BasicSizeTail = + (tempLoopInfo.s2BasicSizeTail == 0) ? constInfo.s2BaseSize : tempLoopInfo.s2BasicSizeTail; + tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize; + tempLoopInfo.mBasicSizeTail = + (tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail; + + uint32_t gS1SplitNum = (tempLoopInfo.actS1Size * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize; + tempLoopInfo.gS1LoopEnd = (bN2LoopIdx == splitCoreInfo.bN2End) ? splitCoreInfo.gS1End : gS1SplitNum - 1; + if constexpr (Q_LAYOUT_T == LI_LAYOUT::BSND) { + if (tempLoopInfo.gS1LoopEnd == gS1SplitNum - 1 && constInfo.qSeqSize > tempLoopInfo.actS1Size) { + tempLoopInfo.needDealActS1LessThanS1 = true; + } + } +} + +template +__aicore__ inline void LIQPreload::CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, LIQCommon::RunInfo &runInfo) +{ + runInfo.loop = loop; + runInfo.bIdx = tempLoopInfo.bIdx; + runInfo.gS1Idx = tempLoopInfo.gS1Idx; + runInfo.s2Idx = s2LoopIdx; + runInfo.bN2Idx = tempLoopInfo.bN2Idx; + runInfo.isValid = s2LoopIdx <= tempLoopInfo.s2LoopEnd; + + if (!runInfo.isValid) { + return; // 需要验证， v1 时候需要runInfo + } + + runInfo.actS1Size = tempLoopInfo.actS1Size; + runInfo.actS2Size = tempLoopInfo.actS2Size; + // 计算实际基本块size + runInfo.actMBaseSize = tempLoopInfo.actMBaseSize; + runInfo.actualSingleProcessSInnerSize = constInfo.s2BaseSize; + uint32_t s2SplitNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + if (runInfo.s2Idx == s2SplitNum - 1) { + runInfo.actualSingleProcessSInnerSize = tempLoopInfo.s2BasicSizeTail; + } + runInfo.actualSingleProcessSInnerSizeAlign = + LIQCommon::Align((uint32_t)runInfo.actualSingleProcessSInnerSize, LIQCommon::ConstInfo::BUFFER_SIZE_BYTE_32B); + + runInfo.isFirstS2InnerLoop = s2LoopIdx == splitCoreInfo.s2Start; + runInfo.isLastS2InnerLoop = s2LoopIdx == tempLoopInfo.s2LoopEnd; + runInfo.isAllLoopEnd = (runInfo.bN2Idx == splitCoreInfo.bN2End) && (runInfo.gS1Idx == splitCoreInfo.gS1End) && + (runInfo.s2Idx == splitCoreInfo.s2End); + + if (runInfo.isFirstS2InnerLoop) { + uint64_t actualSeqQPrefixSum; + if constexpr (Q_LAYOUT_T == LI_LAYOUT::TND) { + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGmQ.GetValue(runInfo.bIdx - 1); + } else { // BSND + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.qSeqSize; + } + uint64_t tndBIdxOffset = actualSeqQPrefixSum * constInfo.qHeadNum * constInfo.headDim; + // B,S1,N1(N2,G),D + queryCoreOffset = tndBIdxOffset + runInfo.gS1Idx * constInfo.mBaseSize * constInfo.headDim; + // B,S1,N1(N2,G)/T,N1(N2,G) + weightsCoreOffset = actualSeqQPrefixSum * constInfo.qHeadNum + runInfo.n2Idx * constInfo.gSize; + // B,S1,N2,k/T,N2,k + indiceOutCoreOffset = + actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + runInfo.n2Idx * constInfo.sparseCount; + } + uint64_t actualSeqKPrefixSum; + if constexpr (K_LAYOUT_T == LI_LAYOUT::TND) { // T N2 D + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGm.GetValue(runInfo.bIdx - 1); + } else { + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.kSeqSize; + } + uint64_t tndBIdxOffsetForK = actualSeqKPrefixSum * constInfo.kHeadNum * constInfo.headDim; + keyCoreOffset = tndBIdxOffsetForK + runInfo.s2Idx * constInfo.s2BaseSize * constInfo.kHeadNum * constInfo.headDim; + keyScaleCoreOffset = (actualSeqKPrefixSum + runInfo.s2Idx * constInfo.s2BaseSize) * constInfo.kHeadNum; + runInfo.tensorQueryOffset = queryCoreOffset; + runInfo.tensorKeyOffset = keyCoreOffset; + runInfo.tensorKeyScaleOffset = keyScaleCoreOffset; + runInfo.tensorWeightsOffset = weightsCoreOffset; + runInfo.indiceOutOffset = indiceOutCoreOffset; +} + +template +__aicore__ inline void LIQPreload::Process() +{ + if (usedCoreNum == 0) { + // 没有计算任务，直接清理输出 + ProcessInvalid(); + return; + } + + ProcessMain(); + + ProcessDecode(); +} + +template +__aicore__ inline void LIQPreload::ProcessInvalid() +{ + if ASCEND_IS_AIV { + uint32_t aivCoreNum = GetBlockNum() * 2; // 2 means c:v = 1:2 + uint64_t totalOutputSize = + constInfo.batchSize * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount; + uint64_t singleCoreSize = + LIQCommon::Align((totalOutputSize + aivCoreNum - 1) / aivCoreNum, GM_ALIGN_BYTES / sizeof(OUT_T)); + uint64_t baseSize = tmpBlockIdx * singleCoreSize; + if (baseSize < totalOutputSize) { + uint64_t dealSize = + (baseSize + singleCoreSize <= totalOutputSize) ? singleCoreSize : totalOutputSize - baseSize; + GlobalTensor output = indiceOutGm[baseSize]; + AscendC::InitGlobalMemory(output, dealSize, constInfo.INVALID_IDX); + } + } +} + +template +__aicore__ inline void LIQPreload::ProcessMain() +{ + if (aiCoreIdx >= usedCoreNum) { + // 无任务核直接返回 + return; + } + + if ASCEND_IS_AIV { + vectorService.AllocEventID(); + CrossCoreSetFlag(constInfo.syncV1C1); + CrossCoreSetFlag(constInfo.syncV1C1); + } else { + matmulService.AllocEventID(); + CrossCoreSetFlag(constInfo.syncC1V0); + CrossCoreSetFlag(constInfo.syncC1V0); + } + + LIQCommon::RunInfo runInfo[LI_QUANT_PRELOAD_TASK_CACHE_SIZE]; + + uint32_t gloop = 0; + for (uint32_t bN2LoopIdx = splitCoreInfo.bN2Start; bN2LoopIdx <= splitCoreInfo.bN2End; bN2LoopIdx++) { + CalcGS1LoopParams(bN2LoopIdx); + if (tempLoopInfo.curActSeqLenIsZero) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, 0U); + + if ASCEND_IS_AIV { + if (bN2LoopIdx == splitCoreInfo.bN2End && gloop > 0) { + CrossCoreWaitFlag(constInfo.syncC1V1); + vectorService.ProcessVec1(runInfo[1 - gloop % LI_QUANT_PRELOAD_TASK_CACHE_SIZE]); + CrossCoreSetFlag( + constInfo.syncV1C1); // 反向同步 1 + } + } + continue; + } + for (uint32_t gS1LoopIdx = splitCoreInfo.gS1Start; gS1LoopIdx <= tempLoopInfo.gS1LoopEnd; gS1LoopIdx++) { + CalcS2LoopParams(bN2LoopIdx, gS1LoopIdx); + bool isEnd = (bN2LoopIdx == splitCoreInfo.bN2End) && (gS1LoopIdx == splitCoreInfo.gS1End); + uint32_t extraLoop = isEnd ? LI_QUANT_PRELOAD_TASK_CACHE_SIZE - 1 : 0; + for (int s2LoopIdx = splitCoreInfo.s2Start; s2LoopIdx <= (tempLoopInfo.s2LoopEnd + extraLoop); + s2LoopIdx++) { + ProcessBaseBlock(gloop, s2LoopIdx, runInfo); + ++gloop; + } + splitCoreInfo.s2Start = 0; + } + if (tempLoopInfo.needDealActS1LessThanS1) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, tempLoopInfo.actS1Size); + } + splitCoreInfo.gS1Start = 0; + } + + if ASCEND_IS_AIV { + vectorService.FreeEventID(); + CrossCoreWaitFlag(constInfo.syncC1V0); + CrossCoreWaitFlag(constInfo.syncC1V0); + } else { + matmulService.FreeEventID(); + CrossCoreWaitFlag(constInfo.syncV1C1); + CrossCoreWaitFlag(constInfo.syncV1C1); + } +} + +template +__aicore__ inline void LIQPreload::ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, + LIQCommon::RunInfo runInfo[LI_QUANT_PRELOAD_TASK_CACHE_SIZE]) +{ + int32_t curTaskId = loop % LI_QUANT_PRELOAD_TASK_CACHE_SIZE; + LIQCommon::RunInfo &curRunInfo = runInfo[curTaskId]; + LIQCommon::RunInfo &lastRunInfo = runInfo[1 - curTaskId]; + + CalcRunInfo(loop, s2LoopIdx, curRunInfo); + + if (curRunInfo.isValid) { + if ASCEND_IS_AIC { + if (curRunInfo.isFirstS2InnerLoop) { + CrossCoreWaitFlag(constInfo.syncV0C1); + } + CrossCoreWaitFlag(constInfo.syncV1C1); // 反向同步 1 + matmulService.ComputeMm1(curRunInfo); + CrossCoreSetFlag(constInfo.syncC1V1); + if (curRunInfo.isLastS2InnerLoop) { + CrossCoreSetFlag(constInfo.syncC1V0); // 反向同步 0 + } + } else { + if (curRunInfo.isFirstS2InnerLoop) { + CrossCoreWaitFlag(constInfo.syncC1V0); // 反向同步 0 + vectorService.ProcessVec0(curRunInfo); + CrossCoreSetFlag(constInfo.syncV0C1); + } + } + } + + if (lastRunInfo.isValid) { + if ASCEND_IS_AIV { + CrossCoreWaitFlag(constInfo.syncC1V1); + vectorService.ProcessVec1(lastRunInfo); + CrossCoreSetFlag(constInfo.syncV1C1); // 反向同步 1 + } + lastRunInfo.isValid = false; + } +} + +template +__aicore__ inline void LIQPreload::ProcessDecode() +{ + if ASCEND_IS_AIV { + vectorService.InitLDBuffers(pipe); + ICachePreLoad(LD_PREFETCH_LEN); + SyncAll(); + if (splitCoreInfo.isLD) { + vectorService.ProcessLD(); + } + } +} +} // namespace LIQKernel +#endif // LIGHTNING_INDEXER_QUANT_KERNEL_H \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_service_cube.h b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_service_cube.h new file mode 100644 index 000000000..2f58a9e1b --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_service_cube.h @@ -0,0 +1,613 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_quant_service_cube.h + * \brief use 5 buffer for matmul l1, better pipeline + */ +#ifndef LIGHTNING_INDEXER_QUANT_SERVICE_CUBE_H +#define LIGHTNING_INDEXER_QUANT_SERVICE_CUBE_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "lightning_indexer_quant_common.h" + +namespace LIQKernel { +using namespace LIQCommon; +struct MmInfo { + int64_t s2L0LoopId; + int64_t s1gL0LoopId; + int64_t s2L0RealSize; + int64_t s2GmOffset; +}; + +template +class LIQMatmul { +public: + using Q_T = typename LIQT::queryType; + using K_T = typename LIQT::keyType; + + __aicore__ inline LIQMatmul(){}; + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitMm1GlobalTensor(const GlobalTensor &blkTableGm, const GlobalTensor &keyGm, + const GlobalTensor &queryGm, const GlobalTensor &mm1ResGm, + const GlobalTensor &weightWorkspaceGm); + __aicore__ inline void InitParams(const ConstInfo &constInfo); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void ComputeMm1(const LIQCommon::RunInfo &runInfo); + + static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding; + static constexpr uint64_t DOUBLE_BUF_NUM = 2; + static constexpr uint64_t L0AB_BUF_NUM = 4; + + static constexpr uint32_t KEY_MTE1_MTE2_EVENT = EVENT_ID2; + static constexpr uint32_t QW_MTE1_MTE2_EVENT = EVENT_ID5; // KEY_MTE1_MTE2_EVENT + DOUBLE_BUF_NUM; + static constexpr uint32_t M_MTE1_EVENT = EVENT_ID3; + static constexpr uint32_t M_FIX_EVENT = EVENT_ID0; + static constexpr uint32_t FIX_M_EVENT = EVENT_ID2; + static constexpr uint32_t FIX_MTE1_EVENT = EVENT_ID4; + + static constexpr uint64_t S8_BLOCK_CUBE = 32; + + static constexpr uint32_t MTE2_MTE1_EVENT = EVENT_ID2; + static constexpr uint32_t MTE1_M_EVENT = EVENT_ID2; + + static constexpr uint64_t D_BASIC_BLOCK = 128; + static constexpr uint64_t S1G_BASIC_BLOCK_L1 = 256; + + static constexpr uint64_t S1G_BASIC_BLOCK_L0 = 128; + static constexpr uint64_t S2_BASIC_BLOCK_L0 = 128; + + static constexpr uint64_t QUERY_BUFFER_OFFSET = S1G_BASIC_BLOCK_L1 * D_BASIC_BLOCK; + static constexpr uint64_t SL1_BUFFER_OFFSET = S1G_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0; + static constexpr uint64_t KEY_BUFFER_OFFSET = S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK; + static constexpr uint64_t WEIGHT_BUFFER_OFFSET = S1G_BASIC_BLOCK_L1 * BLOCK_CUBE; + static constexpr uint64_t L0AB_BUFFER_OFFSET_S8_16K = 16 * 1024; + static constexpr uint64_t L0AB_BUFFER_OFFSET_FP16_16K = 16 * 512; + static constexpr uint64_t L0C_BUFFER_OFFSET = 64 * 256; + +private: + __aicore__ inline void WeightDmaCopy(uint64_t s1gL1RealSize, const LIQCommon::RunInfo &runInfo); + __aicore__ inline void LoadKeyToL0b(uint64_t s2L0RealSize); + __aicore__ inline void LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL1RealSize, uint64_t s1gL0RealSize); + __aicore__ inline void QueryNd2Nz(uint64_t s1gL1RealSize, const LIQCommon::RunInfo &runInfo); + __aicore__ inline void KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, const LIQCommon::RunInfo &runInfo); + __aicore__ inline void KeyNd2Nz(uint64_t s2L1RealSize, const MmInfo &mmInfo, const LIQCommon::RunInfo &runInfo); + __aicore__ inline void FixpSToL1(uint64_t s1gL0RealSize, uint64_t s2L0RealSize); + __aicore__ inline void LoadSToL0b(uint64_t s1gL1RealSize, uint64_t s2L0RealSize, uint64_t sL1BufIdx, + int64_t mStartPt); + __aicore__ inline void LoadWeightToL0a(uint64_t s1gL1Offset); + __aicore__ inline void ComputeWs(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, int64_t s1gOffset); + __aicore__ inline void FixpResToGm(uint64_t s1L0RealCount, uint64_t s2L0RealSize, uint64_t s1GmOffset, + uint64_t s2GmOffset, const LIQCommon::RunInfo &runInfo); + __aicore__ inline void ComputeQk(uint64_t s1gL0RealSize, uint64_t s2L0RealSize); + __aicore__ inline void ProcessWs(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t sL1BufIdx, + const MmInfo &mmInfo, const LIQCommon::RunInfo &runInfo); + __aicore__ inline void ProcessQk(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t s1L0LoopCnt, + const MmInfo &mmInfo, const LIQCommon::RunInfo &runInfo); + __aicore__ inline void CalcMmInfo(MmInfo &mmInfo, uint64_t loopIdx, uint64_t s1L0LoopCnt, const MmInfo &lastMmInfo, + const LIQCommon::RunInfo &runInfo); + static constexpr LI_LAYOUT Q_LAYOUT_T = LIQT::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = LIQT::keyLayout; + GlobalTensor blkTableGm_; + GlobalTensor keyGm_; + GlobalTensor queryGm_; + GlobalTensor weightGm_; + GlobalTensor mm1ResGm_; + + TBuf bufQL1_; + LocalTensor queryL1_; + TBuf bufKeyL1_; + LocalTensor keyL1_; + TBuf bufWeightL1_; + LocalTensor weightL1_; + TBuf bufSL1_; + LocalTensor sL1_; + + TBuf bufL0A_; + LocalTensor l0a_; + TBuf bufL0B_; + LocalTensor l0b_; + + TBuf bufL0C_; + LocalTensor cL0_; + + uint64_t keyL1BufIdx_ = 0; + uint64_t qwL1Mte2BufIdx_ = 0; + uint64_t sL1BufIdx_ = 0; + uint64_t l0BufIdx_ = 0; + uint64_t l0cBufIdx_ = 0; + + ConstInfo constInfo_; +}; + +template +__aicore__ inline void LIQMatmul::InitParams(const ConstInfo &constInfo) +{ + constInfo_ = constInfo; +} + +template +__aicore__ inline void LIQMatmul::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(bufQL1_, DOUBLE_BUF_NUM * S1G_BASIC_BLOCK_L1 * D_BASIC_BLOCK * sizeof(Q_T)); + queryL1_ = bufQL1_.Get(); + pipe->InitBuffer(bufKeyL1_, DOUBLE_BUF_NUM * S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK * sizeof(K_T)); + keyL1_ = bufKeyL1_.Get(); + + pipe->InitBuffer(bufWeightL1_, DOUBLE_BUF_NUM * S1G_BASIC_BLOCK_L1 * BLOCK_CUBE * sizeof(half)); + weightL1_ = bufWeightL1_.Get(); + pipe->InitBuffer(bufSL1_, DOUBLE_BUF_NUM * S2_BASIC_BLOCK_L0 * S1G_BASIC_BLOCK_L0 * sizeof(half)); + sL1_ = bufSL1_.Get(); + + pipe->InitBuffer(bufL0A_, 64 * 1024); + l0a_ = bufL0A_.Get(); + pipe->InitBuffer(bufL0B_, 64 * 1024); + l0b_ = bufL0B_.Get(); + + pipe->InitBuffer(bufL0C_, 128 * 1024); + cL0_ = bufL0C_.Get(); +} + +template +__aicore__ inline void LIQMatmul::InitMm1GlobalTensor(const GlobalTensor &blkTableGm, + const GlobalTensor &keyGm, + const GlobalTensor &queryGm, + const GlobalTensor &mm1ResGm, + const GlobalTensor &weightWorkspaceGm) +{ + blkTableGm_ = blkTableGm; + keyGm_ = keyGm; + queryGm_ = queryGm; + mm1ResGm_ = mm1ResGm; + weightGm_ = weightWorkspaceGm; +} + +template +__aicore__ inline void LIQMatmul::ProcessWs(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t sL1BufIdx, + const MmInfo &mmInfo, const LIQCommon::RunInfo &runInfo) +{ + WaitFlag(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM); + for (int64_t s1gOffset = 0; s1gOffset < s1gL0RealSize; s1gOffset += constInfo_.gSize) { + WaitFlag(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM); + LoadSToL0b(s1gL0RealSize, mmInfo.s2L0RealSize, sL1BufIdx, s1gOffset); + LoadWeightToL0a(s1gOffset + s1gL1Offset); + + ComputeWs(s1gL0RealSize, mmInfo.s2L0RealSize, s1gOffset); + + SetFlag(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM); + l0BufIdx_++; + } + + FixpResToGm(s1gL0RealSize / constInfo_.gSize, mmInfo.s2L0RealSize, s1gL1Offset / constInfo_.gSize, + mmInfo.s2L0LoopId * S2_BASIC_BLOCK_L0, runInfo); + SetFlag(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM); + l0cBufIdx_++; +} + +template +__aicore__ inline void LIQMatmul::ProcessQk(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t s1L0LoopCnt, + const MmInfo &mmInfo, const LIQCommon::RunInfo &runInfo) +{ + if (mmInfo.s1gL0LoopId == 0) { + WaitFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % DOUBLE_BUF_NUM); + if constexpr (K_LAYOUT_T == LI_LAYOUT::PA_BSND) { + KeyNd2NzForPA(mmInfo.s2L0RealSize, runInfo.s2Idx * constInfo_.s2BaseSize + mmInfo.s2GmOffset, runInfo); + } else { + KeyNd2Nz(mmInfo.s2L0RealSize, mmInfo, runInfo); + } + + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + } + + WaitFlag(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM); + LoadQueryToL0a(s1gL1Offset, runInfo.actMBaseSize, s1gL0RealSize); + LoadKeyToL0b(mmInfo.s2L0RealSize); + + if (mmInfo.s1gL0LoopId + 1 >= s1L0LoopCnt) { + SetFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % DOUBLE_BUF_NUM); + keyL1BufIdx_++; + } + + WaitFlag(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM); + ComputeQk(s1gL0RealSize, mmInfo.s2L0RealSize); + SetFlag(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM); + + FixpSToL1(s1gL0RealSize, mmInfo.s2L0RealSize); + SetFlag(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM); + l0BufIdx_++; + l0cBufIdx_++; +} + +template +__aicore__ inline void LIQMatmul::CalcMmInfo(MmInfo &mmInfo, uint64_t loopIdx, uint64_t s1L0LoopCnt, + const MmInfo &lastMmInfo, const LIQCommon::RunInfo &runInfo) +{ + mmInfo.s2L0LoopId = loopIdx / s1L0LoopCnt; + mmInfo.s1gL0LoopId = loopIdx % s1L0LoopCnt; + + if (mmInfo.s1gL0LoopId == 0) { + mmInfo.s2GmOffset = mmInfo.s2L0LoopId * S2_BASIC_BLOCK_L0; + mmInfo.s2L0RealSize = mmInfo.s2GmOffset + S2_BASIC_BLOCK_L0 > runInfo.actualSingleProcessSInnerSize + ? runInfo.actualSingleProcessSInnerSize - mmInfo.s2GmOffset + : S2_BASIC_BLOCK_L0; + } else { + mmInfo.s2L0RealSize = lastMmInfo.s2L0RealSize; + } +} + +template +__aicore__ inline void LIQMatmul::ComputeMm1(const LIQCommon::RunInfo &runInfo) +{ + if (runInfo.isFirstS2InnerLoop) { + WaitFlag(QW_MTE1_MTE2_EVENT + qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM); + QueryNd2Nz(runInfo.actMBaseSize, runInfo); // 256 * 128 // L1BasicBlock + WeightDmaCopy(runInfo.actMBaseSize, runInfo); + } + int64_t loopIdx = 0; + int64_t s2L0LoopCnt = CeilDiv(runInfo.actualSingleProcessSInnerSize, S2_BASIC_BLOCK_L0); // 2048取128 + int64_t s1L0LoopCnt = CeilDiv(runInfo.actMBaseSize, S1G_BASIC_BLOCK_L0); // 256取128 + int64_t s1gL1Offset[2] = {0, static_cast(S1G_BASIC_BLOCK_L0)}; + int64_t s1gL0RealSize[2] = {s1L0LoopCnt > 1 ? static_cast(S1G_BASIC_BLOCK_L0) : runInfo.actMBaseSize, + runInfo.actMBaseSize - s1gL1Offset[1]}; + MmInfo mmInfo[2]; + CalcMmInfo(mmInfo[loopIdx & 1], loopIdx, s1L0LoopCnt, mmInfo[(loopIdx + 1) & 1], runInfo); + + ProcessQk(s1gL0RealSize[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], + s1gL1Offset[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], s1L0LoopCnt, mmInfo[loopIdx & 1], + runInfo); + + SetFlag(FIX_MTE1_EVENT + sL1BufIdx_ % DOUBLE_BUF_NUM); + sL1BufIdx_++; + loopIdx++; + + while (loopIdx < s2L0LoopCnt * s1L0LoopCnt) { + CalcMmInfo(mmInfo[loopIdx & 1], loopIdx, s1L0LoopCnt, mmInfo[(loopIdx + 1) & 1], runInfo); + + ProcessQk(s1gL0RealSize[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], + s1gL1Offset[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], s1L0LoopCnt, mmInfo[loopIdx & 1], + runInfo); + + SetFlag(FIX_MTE1_EVENT + sL1BufIdx_ % DOUBLE_BUF_NUM); + sL1BufIdx_++; + + WaitFlag(FIX_MTE1_EVENT + sL1BufIdx_ % DOUBLE_BUF_NUM); + + ProcessWs(s1gL0RealSize[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], + s1gL1Offset[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], sL1BufIdx_, + mmInfo[(loopIdx + 1) & 1], runInfo); + loopIdx++; + } + + WaitFlag(FIX_MTE1_EVENT + (sL1BufIdx_ + 1) % DOUBLE_BUF_NUM); + + ProcessWs(s1gL0RealSize[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], + s1gL1Offset[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], sL1BufIdx_ - 1, + mmInfo[(loopIdx + 1) & 1], runInfo); + + if (runInfo.isLastS2InnerLoop) { + SetFlag(QW_MTE1_MTE2_EVENT + qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM); + qwL1Mte2BufIdx_++; + } +} + +// blkNum, blkSize, N2, D +template +__aicore__ inline void LIQMatmul::KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const LIQCommon::RunInfo &runInfo) +{ + uint64_t s2L1Offset = 0; + while (s2L1Offset < s2L1RealSize) { + uint64_t s2BlkId = (s2L1Offset + s2GmOffset) / constInfo_.kCacheBlockSize; + uint64_t s2BlkOffset = (s2L1Offset + s2GmOffset) % constInfo_.kCacheBlockSize; + uint64_t keyGmOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId) * + constInfo_.kCacheBlockSize * constInfo_.kHeadNum * constInfo_.headDim + + s2BlkOffset * constInfo_.headDim; + uint64_t s2Mte2Size = s2L1RealSize - s2L1Offset; + s2Mte2Size = s2BlkOffset + s2Mte2Size >= constInfo_.kCacheBlockSize ? constInfo_.kCacheBlockSize - s2BlkOffset + : s2Mte2Size; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2Mte2Size; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(keyL1_[(keyL1BufIdx_ % DOUBLE_BUF_NUM) * KEY_BUFFER_OFFSET + s2L1Offset * S8_BLOCK_CUBE], + keyGm_[keyGmOffset], nd2nzPara); + + s2L1Offset += s2Mte2Size; + } +} + +template +__aicore__ inline void LIQMatmul::KeyNd2Nz(uint64_t s2L1RealSize, const MmInfo &mmInfo, + const LIQCommon::RunInfo &runInfo) +{ + uint64_t dStride = constInfo_.headDim; + if constexpr (K_LAYOUT_T == LI_LAYOUT::BSND || K_LAYOUT_T == LI_LAYOUT::TND) { + dStride = constInfo_.headDim * constInfo_.kHeadNum; // constInfo_.kHeadNum + } + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2L1RealSize; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = dStride; + nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + // 默认一块buf最多放两份 + DataCopy(keyL1_[(keyL1BufIdx_ % DOUBLE_BUF_NUM) * KEY_BUFFER_OFFSET], + keyGm_[runInfo.tensorKeyOffset + mmInfo.s2GmOffset * constInfo_.headDim], nd2nzPara); +} + +// batch, s1, g, 1 +template +__aicore__ inline void LIQMatmul::WeightDmaCopy(uint64_t s1gL1RealSize, const LIQCommon::RunInfo &runInfo) +{ + DataCopyParams copyInParams; + copyInParams.blockCount = 1; + copyInParams.blockLen = s1gL1RealSize; + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + DataCopy(weightL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * WEIGHT_BUFFER_OFFSET], + weightGm_[runInfo.loop % DOUBLE_BUF_NUM * BLOCK_CUBE * constInfo_.mBaseSize], copyInParams); +} + +// batch, s1, n2, g, d +template +__aicore__ inline void LIQMatmul::QueryNd2Nz(uint64_t s1gL1RealSize, const LIQCommon::RunInfo &runInfo) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s1gL1RealSize; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(s1gL1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + // 默认一块buf最多放两份 + DataCopy(queryL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * QUERY_BUFFER_OFFSET], queryGm_[runInfo.tensorQueryOffset], + nd2nzPara); +} + +// s1g, d +template +__aicore__ inline void LIQMatmul::LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL1RealSize, + uint64_t s1gL0RealSize) +{ + LoadData3DParamsV2 loadData3DParams; + // SetFmatrixParams + loadData3DParams.l1H = CeilDiv(s1gL1RealSize, BLOCK_CUBE); // Hin=M1=8 + loadData3DParams.l1W = BLOCK_CUBE; // Win=M0 + loadData3DParams.channelSize = constInfo_.headDim; // Cin=K + + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + // SetLoadToA0Params + loadData3DParams.mExtension = CeilAlign(s1gL0RealSize, BLOCK_CUBE); // M height维度目的 + loadData3DParams.kExtension = constInfo_.headDim; // K width维度目的 + loadData3DParams.mStartPt = s1gL1Offset; + loadData3DParams.kStartPt = 0; + loadData3DParams.strideW = 1; + loadData3DParams.strideH = 1; + loadData3DParams.filterW = 1; + loadData3DParams.filterSizeW = (1 >> 8) & 255; + loadData3DParams.filterH = 1; + loadData3DParams.filterSizeH = (1 >> 8) & 255; + loadData3DParams.dilationFilterW = 1; + loadData3DParams.dilationFilterH = 1; + loadData3DParams.enTranspose = 0; + loadData3DParams.fMatrixCtrl = 0; + + LoadData(l0a_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K], + queryL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * QUERY_BUFFER_OFFSET], + loadData3DParams); +} + +// s1, g, s2 --> 2 * 64* 128 +template +__aicore__ inline void LIQMatmul::LoadSToL0b(uint64_t s1gL1RealSize, uint64_t s2L0RealSize, uint64_t sL1BufIdx, + int64_t mStartPt) +{ + LoadData3DParamsV2 loadData3DParams; + // SetFmatrixParams + loadData3DParams.l1H = S1G_BASIC_BLOCK_L0 / BLOCK_CUBE; // Hin=M1=8 + loadData3DParams.l1W = BLOCK_CUBE; // Win=M0 + loadData3DParams.channelSize = CeilAlign(s2L0RealSize, BLOCK_CUBE); // Cin=K + + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + // SetLoadToA0Params + loadData3DParams.mExtension = constInfo_.gSize; // M height维度目的 + loadData3DParams.kExtension = CeilAlign(s2L0RealSize, BLOCK_CUBE); // K width维度目的 + loadData3DParams.kStartPt = 0; + loadData3DParams.strideW = 1; + loadData3DParams.strideH = 1; + loadData3DParams.filterW = 1; + loadData3DParams.filterSizeW = (1 >> 8) & 255; + loadData3DParams.filterH = 1; + loadData3DParams.filterSizeH = (1 >> 8) & 255; + loadData3DParams.dilationFilterW = 1; + loadData3DParams.dilationFilterH = 1; + loadData3DParams.enTranspose = 1; + loadData3DParams.fMatrixCtrl = 0; + + loadData3DParams.mStartPt = mStartPt; + LoadData( + l0b_.template ReinterpretCast()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K], + sL1_[(sL1BufIdx % DOUBLE_BUF_NUM) * SL1_BUFFER_OFFSET], loadData3DParams); +} + +// s1,g,1(16), 2,64,16 +template +__aicore__ inline void LIQMatmul::LoadWeightToL0a(uint64_t s1gL1Offset) +{ + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = CeilDiv(constInfo_.gSize, BLOCK_CUBE); + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = true; + LoadData(l0a_.template ReinterpretCast()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K], + weightL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * WEIGHT_BUFFER_OFFSET + s1gL1Offset* BLOCK_CUBE], + loadData2DParams); +} + +// s2, d -> 128,128 +template +__aicore__ inline void LIQMatmul::LoadKeyToL0b(uint64_t s2L0RealSize) +{ + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = CeilDiv(s2L0RealSize, BLOCK_CUBE) * CeilDiv(constInfo_.headDim, S8_BLOCK_CUBE); + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = false; + LoadData(l0b_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K], + keyL1_[(keyL1BufIdx_ % DOUBLE_BUF_NUM) * KEY_BUFFER_OFFSET], loadData2DParams); +} + +// A: s1,g,1(16) B: s1,g,s2 C: s1, 1(16), s2 +template +__aicore__ inline void LIQMatmul::ComputeWs(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, int64_t s1gOffset) +{ + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + MmadParams mmadParams; + mmadParams.m = BLOCK_CUBE; + mmadParams.n = s2L0RealSize; + mmadParams.k = constInfo_.gSize; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + Mmad(cL0_.template ReinterpretCast()[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET + + s1gOffset * S2_BASIC_BLOCK_L0], + l0a_.template ReinterpretCast()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K], + l0b_.template ReinterpretCast()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K], + mmadParams); +} + +template +__aicore__ inline void LIQMatmul::ComputeQk(uint64_t s1gL0RealSize, uint64_t s2L0RealSize) +{ + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + + MmadParams mmadParams; + mmadParams.m = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + mmadParams.n = s2L0RealSize; + mmadParams.k = constInfo_.headDim; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + Mmad(cL0_[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET], + l0a_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K], + l0b_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K], mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } +} + +template +__aicore__ inline void LIQMatmul::FixpSToL1(uint64_t s1gL0RealSize, uint64_t s2L0RealSize) +{ + SetFlag(M_FIX_EVENT); + WaitFlag(M_FIX_EVENT); + DataCopyCO12DstParams params; + params.mSize = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + params.nSize = CeilAlign(s2L0RealSize, BLOCK_CUBE); + params.dstStride = S1G_BASIC_BLOCK_L0; + params.srcStride = params.mSize; + params.quantPre = QuantMode_t::DEQF16; + params.reluPre = 1; + params.channelSplit = 0; + params.nz2ndEn = 0; + SetFixpipePreQuantFlag(0x3a800000); + DataCopy(sL1_[(sL1BufIdx_ % DOUBLE_BUF_NUM) * SL1_BUFFER_OFFSET], + cL0_[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET], params); +} + +template +__aicore__ inline void LIQMatmul::FixpResToGm(uint64_t s1L0RealCount, uint64_t s2L0RealSize, uint64_t s1GmOffset, + uint64_t s2GmOffset, const LIQCommon::RunInfo &runInfo) +{ + SetFlag(M_FIX_EVENT); + WaitFlag(M_FIX_EVENT); + + AscendC::DataCopyCO12DstParams intriParams; + intriParams.mSize = 1; + intriParams.nSize = s2L0RealSize; + intriParams.dstStride = constInfo_.s2BaseSize; + intriParams.srcStride = 16; + // set mode according to dtype + intriParams.quantPre = QuantMode_t::NoQuant; + intriParams.nz2ndEn = true; + intriParams.reluPre = 0; + AscendC::SetFixpipeNz2ndFlag(s1L0RealCount, CeilDiv(constInfo_.gSize, BLOCK_CUBE) * S2_BASIC_BLOCK_L0 / BLOCK_CUBE, + 2048); + AscendC::DataCopy(mm1ResGm_[(runInfo.loop % 2) * constInfo_.mBaseSize / constInfo_.gSize * constInfo_.s2BaseSize + + s1GmOffset * intriParams.dstStride + s2GmOffset], + cL0_.template ReinterpretCast()[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET], + intriParams); +} + +template +__aicore__ inline void LIQMatmul::AllocEventID() +{ + SetFlag(KEY_MTE1_MTE2_EVENT + 0); + SetFlag(KEY_MTE1_MTE2_EVENT + 1); + SetFlag(KEY_MTE1_MTE2_EVENT + 2); + + SetFlag(QW_MTE1_MTE2_EVENT + 0); + SetFlag(QW_MTE1_MTE2_EVENT + 1); + + SetFlag(M_MTE1_EVENT + 0); + SetFlag(M_MTE1_EVENT + 1); + SetFlag(M_MTE1_EVENT + 2); + SetFlag(M_MTE1_EVENT + 3); + + SetFlag(FIX_M_EVENT + 0); + SetFlag(FIX_M_EVENT + 1); +} + +template +__aicore__ inline void LIQMatmul::FreeEventID() +{ + WaitFlag(KEY_MTE1_MTE2_EVENT + 0); + WaitFlag(KEY_MTE1_MTE2_EVENT + 1); + WaitFlag(KEY_MTE1_MTE2_EVENT + 2); + + WaitFlag(QW_MTE1_MTE2_EVENT + 0); + WaitFlag(QW_MTE1_MTE2_EVENT + 1); + + WaitFlag(M_MTE1_EVENT + 0); + WaitFlag(M_MTE1_EVENT + 1); + WaitFlag(M_MTE1_EVENT + 2); + WaitFlag(M_MTE1_EVENT + 3); + + WaitFlag(FIX_M_EVENT + 0); + WaitFlag(FIX_M_EVENT + 1); +} +} // namespace LIQKernel +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_service_vector.h b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_service_vector.h new file mode 100644 index 000000000..2588998c9 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_service_vector.h @@ -0,0 +1,665 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_quant_service_vector.h + * \brief + */ +#ifndef LIGHTNING_INDEXER_QUANT_SERVICE_VECTOR_H +#define LIGHTNING_INDEXER_QUANT_SERVICE_VECTOR_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "lightning_indexer_quant_common.h" +#include "lightning_indexer_quant_vector.h" + +namespace LIQKernel { +using namespace LIQCommon; +using namespace LIQServiceVec; +constexpr uint32_t BASE_TOPK = 2048; +constexpr uint32_t BASE_TOPK_VALUE_IDX_SIZE = 4096; +constexpr uint32_t LD_PARAM_NUM = 16; + +template +class LIQVector { +public: + // =================================类型定义区================================= + static constexpr LI_LAYOUT Q_LAYOUT_T = LIQT::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = LIQT::keyLayout; + static constexpr bool PAGE_ATTENTION = LIQT::pageAttention; + // MM输出数据类型, 当前只支持float + using MM1_OUT_T = float; + + __aicore__ inline LIQVector(){}; + __aicore__ inline void ProcessVec0(const LIQCommon::RunInfo &info); + __aicore__ inline void ProcessVec1(const LIQCommon::RunInfo &info); + __aicore__ inline void ProcessLD(); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitParams(const struct LIQCommon::ConstInfo &constInfo, + const LIQTilingData *__restrict tilingData); + __aicore__ inline void InitVecWorkspaceTensor(GlobalTensor vec0OutGm, GlobalTensor mm1ResGm, + GlobalTensor vec1ResGm, GlobalTensor vec1ParamGm); + __aicore__ inline void InitVecInputTensor(GlobalTensor weightsGm, GlobalTensor qScaleGm, + GlobalTensor kScaleGm, GlobalTensor indiceOutGm, + GlobalTensor blockTableGm); + __aicore__ inline void CleanInvalidOutput(int64_t invalidS1offset); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void InitLDBuffers(TPipe *pipe); + +protected: + GlobalTensor mm1ResGm; + GlobalTensor vec1ResGm; + GlobalTensor vec1ParamGm; + GlobalTensor weightsGm; + GlobalTensor qScaleGm; + GlobalTensor kScaleGm; + GlobalTensor vec0OutGm; + GlobalTensor indiceOutGm; + GlobalTensor blockTableGm; + // =================================常量区================================= + +private: + __aicore__ inline void GetKeyScale(const LIQCommon::RunInfo &runInfo, const LocalTensor &resUb, + int64_t batchId, int64_t startS2, int64_t getLen); + // ================================Local Buffer区==================================== + // queue + TQue inQueue_; + TQue outQueue_; + + // tmp buff for vector + TBuf sortOutBuf_; + TBuf indexBuf_; + TBuf paramBuf_; + TBuf tmpBuf_; + + // tmp buff for LD + TBuf<> ldToBeMrgBuf_; + TBuf<> ldTmpBuf_; + TBuf<> ldOutValueBuf_; + TBuf<> ldOutIdxBuf_; + + LocalTensor globalTopkIndice_; + LocalTensor globalTopkUb_; + + int32_t blockId_ = -1; + // para for vector + int32_t groupInner_ = 0; + int32_t globalTopkNum_ = 0; + int64_t blockS2StartIdx_ = 0; + int32_t gSize_ = 0; + int32_t kSeqSize_ = 0; + int32_t kHeadNum_ = 0; + int32_t qHeadNum_ = 0; + int32_t s1BaseSize_ = 0; + int32_t s2BaseSize_ = 0; + int32_t kCacheBlockSize_ = 0; + int32_t maxBlockNumPerBatch_ = 0; + + // para for LD + uint32_t mrgListNum_ = 4; + uint32_t paramNum_ = 16; + + struct LIQCommon::ConstInfo constInfo_; +}; + +template +__aicore__ inline void LIQVector::GetKeyScale(const LIQCommon::RunInfo &runInfo, const LocalTensor &resUb, + int64_t batchId, int64_t startS2, int64_t getLen) +{ + // startS2一定能整除kCacheBlockSize_ + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams copyInParams; + if constexpr (PAGE_ATTENTION) { + int32_t startBlockTableIdx = startS2 / kCacheBlockSize_; + int32_t startBlockTableOffset = startS2 % kCacheBlockSize_; + int32_t blockTableBatchOffset = batchId * maxBlockNumPerBatch_; + copyInParams.blockCount = 1; + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + int32_t resUbBaseOffset = 0; + if (startBlockTableOffset > 0) { + int32_t firstPartLen = + kCacheBlockSize_ - startBlockTableOffset > getLen ? getLen : kCacheBlockSize_ - startBlockTableOffset; + copyInParams.blockLen = firstPartLen * sizeof(half); + int32_t blockId = blockTableGm.GetValue(blockTableBatchOffset + startBlockTableIdx); + SetWaitFlag(HardEvent::S_MTE2); + AscendC::DataCopyPad(resUb, kScaleGm[blockId * kCacheBlockSize_ + startBlockTableOffset], + copyInParams, padParams); + startBlockTableIdx++; + getLen = getLen - firstPartLen; + resUbBaseOffset = firstPartLen; + } + int32_t getLoopNum = CeilDiv(getLen, kCacheBlockSize_); + copyInParams.blockLen = kCacheBlockSize_ * sizeof(half); + for (int32_t i = 0; i < getLoopNum; i++) { + if (i == getLoopNum - 1) { + copyInParams.blockLen = (getLen - i * kCacheBlockSize_) * sizeof(half); + } + int32_t blockId = blockTableGm.GetValue(blockTableBatchOffset + startBlockTableIdx + i); + SetWaitFlag(HardEvent::S_MTE2); + AscendC::DataCopyPad(resUb[resUbBaseOffset + i * kCacheBlockSize_], kScaleGm[blockId * kCacheBlockSize_], + copyInParams, padParams); + } + } else { + copyInParams.blockCount = 1; + copyInParams.blockLen = getLen * sizeof(half); + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + AscendC::DataCopyPad(resUb, kScaleGm[runInfo.tensorKeyScaleOffset], copyInParams, padParams); + } +} + +template +__aicore__ inline void LIQVector::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(paramBuf_, LD_PARAM_NUM * sizeof(int64_t)); // 1 KB + pipe->InitBuffer(inQueue_, 2, s2BaseSize_ * sizeof(float) * 2); // 32KB + pipe->InitBuffer(outQueue_, 1, BASE_TOPK * sizeof(float)); // 8 KB + pipe->InitBuffer(indexBuf_, s2BaseSize_ * sizeof(int32_t)); // 8 KB + pipe->InitBuffer(tmpBuf_, 64 * 1024); // 64KB + pipe->InitBuffer(sortOutBuf_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK_VALUE_IDX_SIZE * sizeof(float)); // 32KB + + globalTopkIndice_ = indexBuf_.Get(); + globalTopkUb_ = sortOutBuf_.Get(); + globalTopkNum_ = 0; + + // 基本块执行前初始化UB和GM + // step1. 初始化一个有序索引 0 - s2BaseSize_ + ArithProgression(globalTopkIndice_, 0, 1, s2BaseSize_); + // step2. globalTopkUb_ [CeilDiv(s1BaseSize_, 2), BASE_TOPK, 2] -inf,-1 + InitSortOutBuf(globalTopkUb_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK_VALUE_IDX_SIZE); + + // step3. 初始化vec1ParamGm，是否进行LD的标志位设为-1(needFd=-1) + // vec1ResIn32Gm = [aic, 2, s1BaseSize_, 16] int32 + // ws清零 [needFd, s2AcSeq, s2Start, s2End, isS2End, bn2idx, s1Idx, ......] + LocalTensor tmpfBuff = outQueue_.AllocTensor(); + Duplicate(tmpfBuff.template ReinterpretCast(), -1, 2 * (s1BaseSize_ / 2) * paramNum_ * 2); + SetWaitFlag(HardEvent::V_MTE3); + int64_t wsInfoOffset = (blockId_ / 2) * s1BaseSize_ * 2 * paramNum_ + // 2个AIV共同地址偏移 + (blockId_ % 2) * (s1BaseSize_ / 2) * 2 * paramNum_; // 每个AIV的地址偏移，S1方向 + DataCopyPad(vec1ParamGm[wsInfoOffset], tmpfBuff.template ReinterpretCast(), + {1, static_cast((s1BaseSize_ / 2) * 2 * paramNum_ * sizeof(int64_t)), 0, 0}); + SetWaitFlag(HardEvent::MTE3_V); + outQueue_.FreeTensor(tmpfBuff); +} + +template +__aicore__ inline void LIQVector::InitLDBuffers(TPipe *pipe) +{ + pipe->Reset(); + pipe->InitBuffer(ldToBeMrgBuf_, BASE_TOPK_VALUE_IDX_SIZE * mrgListNum_ * sizeof(float)); + pipe->InitBuffer(ldTmpBuf_, BASE_TOPK_VALUE_IDX_SIZE * mrgListNum_ * sizeof(float)); + pipe->InitBuffer(ldOutValueBuf_, BASE_TOPK * sizeof(float)); + pipe->InitBuffer(ldOutIdxBuf_, BASE_TOPK * sizeof(int32_t)); +} + +template +__aicore__ inline void LIQVector::InitParams(const struct LIQCommon::ConstInfo &constInfo, + const LIQTilingData *__restrict tilingData) +{ + this->constInfo_ = constInfo; + blockS2StartIdx_ = 0; + gSize_ = constInfo.gSize; + kSeqSize_ = constInfo.kSeqSize; + // define N2 para + kHeadNum_ = constInfo.kHeadNum; + qHeadNum_ = constInfo.qHeadNum; + // define MMBase para + s1BaseSize_ = constInfo.s1BaseSize; // 4 + s2BaseSize_ = constInfo.s2BaseSize; // 2048 + kCacheBlockSize_ = constInfo.kCacheBlockSize; + maxBlockNumPerBatch_ = constInfo.maxBlockNumPerBatch; + blockId_ = GetBlockIdx(); +} + +template +__aicore__ inline void LIQVector::InitVecInputTensor(GlobalTensor weightsGm, GlobalTensor qScaleGm, + GlobalTensor kScaleGm, + GlobalTensor indiceOutGm, + GlobalTensor blockTableGm) +{ + this->weightsGm = weightsGm; + this->qScaleGm = qScaleGm; + this->kScaleGm = kScaleGm; + this->indiceOutGm = indiceOutGm; + this->blockTableGm = blockTableGm; +} + +template +__aicore__ inline void LIQVector::InitVecWorkspaceTensor(GlobalTensor vec0OutGm, + GlobalTensor mm1ResGm, + GlobalTensor vec1ResGm, + GlobalTensor vec1ParamGm) +{ + this->mm1ResGm = mm1ResGm; + this->vec1ResGm = vec1ResGm; + this->vec0OutGm = vec0OutGm; + this->vec1ParamGm = vec1ParamGm; +} + +template +__aicore__ inline void LIQVector::AllocEventID() +{ +} + +template +__aicore__ inline void LIQVector::FreeEventID() +{ +} + +template +__aicore__ inline void LIQVector::CleanInvalidOutput(int64_t invalidS1offset) +{ + // init -1 and copy to output + LocalTensor valueULocal = outQueue_.AllocTensor(); + LocalTensor idxULocal1 = valueULocal.template ReinterpretCast(); + Duplicate(idxULocal1, constInfo_.INVALID_IDX, constInfo_.sparseCount); + outQueue_.EnQue(valueULocal); + valueULocal = outQueue_.DeQue(); + LIQServiceVec::CopyOut(indiceOutGm[invalidS1offset], idxULocal1, constInfo_.sparseCount); + outQueue_.FreeTensor(valueULocal); +} + +template +__aicore__ inline void LIQVector::ProcessVec0(const LIQCommon::RunInfo &info) +{ + // 只需要一个v核做 + if (blockId_ % 2 != 0) { + return; + } + int32_t cuBaseS1Idx = info.gS1Idx * s1BaseSize_; + // 计算输出w基地址偏移 偶数循环 -> 0 + aic_offset 奇数循环 -> 4*64 + aic_offset + int64_t vec0OutGmOffset = (info.loop % 2) * ((s1BaseSize_ * gSize_ * BLOCK_CUBE)); + // 计算输入weight的地址偏移，qScale的地址偏移与weight相同 + int64_t weightGmOffset = info.tensorWeightsOffset + cuBaseS1Idx * qHeadNum_; + // 当前需要计算的S1行数，处理尾块场景 + int32_t cuS1ProcNum = cuBaseS1Idx + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + int32_t cuProcEleNum = cuS1ProcNum * gSize_; + + LocalTensor inWeightsUb = inQueue_.AllocTensor(); + LocalTensor inQScaleUb = inWeightsUb[cuProcEleNum]; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams copyInParams; + copyInParams.blockCount = 1; + copyInParams.blockLen = cuProcEleNum * sizeof(half); + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + AscendC::DataCopyPad(inWeightsUb, weightsGm[weightGmOffset], copyInParams, padParams); + AscendC::DataCopyPad(inQScaleUb, qScaleGm[weightGmOffset], copyInParams, padParams); + + inQueue_.EnQue(inWeightsUb); + inWeightsUb = inQueue_.DeQue(); + AscendC::Mul(inWeightsUb, inWeightsUb, inQScaleUb, cuProcEleNum); + PipeBarrier(); + LocalTensor resUb = outQueue_.AllocTensor(); + AscendC::Brcb(resUb, inWeightsUb, static_cast(cuProcEleNum / 8), {1, 8}); + inQueue_.FreeTensor(inWeightsUb); + + outQueue_.EnQue(resUb); + resUb = outQueue_.DeQue(); + AscendC::DataCopyParams copyOutParams; + copyOutParams.blockCount = 1; + copyOutParams.blockLen = cuProcEleNum * BLOCK_CUBE * sizeof(half); + copyOutParams.srcStride = 0; + copyOutParams.dstStride = 0; + AscendC::DataCopyPad(vec0OutGm[vec0OutGmOffset], resUb, copyOutParams); + outQueue_.FreeTensor(resUb); +} + +template +__aicore__ inline void LIQVector::ProcessVec1(const LIQCommon::RunInfo &info) +{ + int32_t cuBaseS1Idx = info.gS1Idx * s1BaseSize_; + int32_t cuBaseS2Idx = info.s2Idx * s2BaseSize_; + + // 计算基本块基地址偏移 偶数循环 -> 0 + aic_offset 奇数循环 -> 4*2048 + aic_offset + int64_t mmGmOffset = (info.loop % 2) * (s1BaseSize_ * s2BaseSize_); + + // cuS1BeginIdxPerAiv: 每个AIV的S1起始偏移 + int32_t cuS1BeginIdxPerAiv = cuBaseS1Idx; + int32_t cuS1ProcNum = + cuS1BeginIdxPerAiv + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + // cuS1ProcNumPerAiv: 每个AIv的S1计算量 + int32_t cuS1ProcNumPerAiv = blockId_ % 2 == 0 ? CeilDiv(cuS1ProcNum, 2) : (cuS1ProcNum / 2); + cuS1BeginIdxPerAiv += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2); + // 基本块基地址偏移奇数核加一个S1地址偏移 + mmGmOffset += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2) * s2BaseSize_; + // 非首个基本块, M(S1)轴发生切换需要初始化 + if (info.loop != 0 && info.s2Idx == 0) { + // globalTopkUb_ value,index=-inf,-1 + InitSortOutBuf(globalTopkUb_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK_VALUE_IDX_SIZE); + blockS2StartIdx_ = 0; + } else if (info.loop == 0) { + blockS2StartIdx_ = info.s2Idx; + } + // cuRealAcSeq: 当前基本块S1对应的AcSeq + int32_t cuRealAcSeq = info.actS2Size; + if (constInfo_.attenMaskFlag) { + // attenMask true场景 + cuRealAcSeq = info.actS2Size - (info.actS1Size - cuS1BeginIdxPerAiv); + } + + // LD输出S1方向偏移，保证2个Vector输出的内容连续 + uint32_t ldS1Offset = (blockId_ % 2 == 0) ? s1BaseSize_ / 2 - cuS1ProcNumPerAiv : 0; + for (int innerS1Idx = 0; innerS1Idx < cuS1ProcNumPerAiv; innerS1Idx++) { + if (constInfo_.attenMaskFlag) { + cuRealAcSeq += 1; + } + int32_t cuS2Len = cuBaseS2Idx + s2BaseSize_ >= cuRealAcSeq ? cuRealAcSeq - cuBaseS2Idx : s2BaseSize_; + int32_t cuS1Idx = cuS1BeginIdxPerAiv + innerS1Idx; + if (cuRealAcSeq > 0 && cuS2Len > 0) { + int32_t cuS2LenVecAlign = CeilDiv(cuS2Len, s2BaseSize_) * s2BaseSize_; + LocalTensor mmInUb = inQueue_.AllocTensor(); + LocalTensor kScaleUb = mmInUb[cuS2LenVecAlign]; + LocalTensor kScaleTUb = kScaleUb.template ReinterpretCast()[cuS2LenVecAlign]; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPadExtParams padTParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams copyInParams; + copyInParams.blockCount = 1; + copyInParams.blockLen = cuS2Len * sizeof(float); + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + AscendC::DataCopyPad(mmInUb, mm1ResGm[mmGmOffset + innerS1Idx * s2BaseSize_], copyInParams, padParams); + GetKeyScale(info, kScaleTUb, info.bIdx, cuBaseS2Idx, cuS2Len); + inQueue_.EnQue(mmInUb); + mmInUb = inQueue_.DeQue(); + AscendC::Cast(kScaleUb, kScaleTUb, RoundMode::CAST_NONE, cuS2Len); + PipeBarrier(); + AscendC::Mul(mmInUb, mmInUb, kScaleUb, cuS2Len); + PipeBarrier(); + LocalTensor sortBuff = tmpBuf_.Get(); + LocalTensor sortScoreUb = sortBuff; + LocalTensor sortIndiceUb = sortBuff[cuS2LenVecAlign]; + PipeBarrier(); + Duplicate(sortScoreUb.template ReinterpretCast(), LIQServiceVec::NEG_INF, cuS2LenVecAlign); + PipeBarrier(); + Adds(sortScoreUb, mmInUb, 0.0f, cuS2Len); + PipeBarrier(); + inQueue_.FreeTensor(mmInUb); + LocalTensor sortIndiceUbInt = sortIndiceUb.template ReinterpretCast(); + // 无效数据索引填充为-1 + if (cuS2LenVecAlign != cuS2Len) { + Duplicate(sortIndiceUbInt, -1, cuS2LenVecAlign); + PipeBarrier(); + } + Adds(sortIndiceUbInt, globalTopkIndice_, static_cast(cuBaseS2Idx), cuS2Len); + PipeBarrier(); + LocalTensor tmpSortBuf = sortBuff[2 * cuS2LenVecAlign]; + LIQServiceVec::SortAll(sortBuff, tmpSortBuf, cuS2LenVecAlign); + PipeBarrier(); + LIQServiceVec::MergeSort(globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], BASE_TOPK, sortBuff, + cuS2LenVecAlign, tmpSortBuf); + PipeBarrier(); + bool isS2End = cuBaseS2Idx + s2BaseSize_ >= cuRealAcSeq; + bool needCopyOutGm = blockS2StartIdx_ == 0 && isS2End; + // 中间结果保存 + bool needCopyWsGm = info.isAllLoopEnd || isS2End; + if (needCopyOutGm) { + LocalTensor idxULocal = outQueue_.AllocTensor(); + ExtractIndex(idxULocal, + globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE].template ReinterpretCast(), + BASE_TOPK); + PipeBarrier(); + InitSortOutBuf(globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], BASE_TOPK_VALUE_IDX_SIZE); + outQueue_.EnQue(idxULocal); + idxULocal = outQueue_.DeQue(); + LIQServiceVec::CopyOut(indiceOutGm[info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount], + idxULocal.template ReinterpretCast(), constInfo_.sparseCount); + outQueue_.FreeTensor(idxULocal); + } else if (needCopyWsGm) { + // vec1Res Gm = [aic, s1BaseSize_, 2, 2, topkOut_] float32 + // vec1Param Gm = [aic, s1BaseSize_, 2, 16] int64 + // 16 = [needFd, s2AcSeq, s2Start, s2End, isS2End, bn2idx, s1Idx, S1ProcNum, ......] + + int64_t wsOffset = + (blockId_ / 2) * s1BaseSize_ * 2 * BASE_TOPK_VALUE_IDX_SIZE + // 2个AIV共同地址偏移 + (blockId_ % 2) * (s1BaseSize_ / 2) * 2 * BASE_TOPK_VALUE_IDX_SIZE + // 每个AIV的地址偏移，S1方向 + (ldS1Offset + innerS1Idx) * 2 * BASE_TOPK_VALUE_IDX_SIZE; + int64_t wsInfoOffset = + (blockId_ / 2) * s1BaseSize_ * 2 * paramNum_ + // 2个AIV共同地址偏移 + (blockId_ % 2) * (s1BaseSize_ / 2) * 2 * paramNum_ + // 每个AIV的地址偏移，S1方向 + (ldS1Offset + innerS1Idx) * 2 * paramNum_; + + LocalTensor tmpiBuff = paramBuf_.Get(); + SetWaitFlag(HardEvent::MTE3_S); + tmpiBuff.SetValue(0, static_cast(1)); + tmpiBuff.SetValue(1, static_cast(cuRealAcSeq)); + tmpiBuff.SetValue(2, static_cast(blockS2StartIdx_)); + tmpiBuff.SetValue(3, static_cast(cuBaseS2Idx + cuS2Len)); + tmpiBuff.SetValue(4, static_cast(isS2End)); + tmpiBuff.SetValue(5, static_cast(info.bN2Idx)); + tmpiBuff.SetValue(6, static_cast(cuS1Idx)); + tmpiBuff.SetValue(7, static_cast(cuS1ProcNum)); + tmpiBuff.SetValue(8, static_cast(info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount)); + // 写入头尾判断 + // [head, tail] + // head: 与前面规约，与前后规约 + // tail: 与后面规约 + bool isTailReduce = blockS2StartIdx_ == 0; // 一定是isLastTile + // WS偏移规则 blockS2StartIdx_ != 0 + // 跟前面块做规约 写到0偏移 不用做计算 blockS2StartIdx_ == 0 and !isS2End + // 跟后面块做规约 写到1偏移 需要 + s1BaseSize_, BASE_TOPK*2 + if (isTailReduce) { // S2不是最后结束的数据就需要往后做规约，放入第二块ws + wsInfoOffset += paramNum_; + wsOffset += BASE_TOPK_VALUE_IDX_SIZE; + } + SetWaitFlag(HardEvent::S_MTE3); + LIQServiceVec::CopyOut(vec1ParamGm[wsInfoOffset], tmpiBuff, 16); + SetWaitFlag(HardEvent::V_MTE3); + LIQServiceVec::CopyOut(vec1ResGm[wsOffset], globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], + BASE_TOPK_VALUE_IDX_SIZE); + SetWaitFlag(HardEvent::MTE3_V); + } + } else if (cuRealAcSeq <= 0) { + CleanInvalidOutput(info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount); + } + } + + // BNSD场景无效S1 输出-1 + if (Q_LAYOUT_T == LI_LAYOUT::BSND) { + // 最后一个S1的基本块, 需要 >= info.actS1Size + bool isS1LoopEnd = (cuBaseS1Idx + s1BaseSize_) >= info.actS1Size; + int32_t invalidS1Num = constInfo_.qSeqSize - info.actS1Size; + // blockS2StartIdx_ == 0 控制S2从开始的核去做冗余清理 + if (invalidS1Num > 0 && isS1LoopEnd && blockS2StartIdx_ == 0) { + int32_t s1NumPerAiv = blockId_ % 2 == 0 ? CeilDiv(invalidS1Num, 2) : (invalidS1Num / 2); + int32_t s1OffsetPerAiv = info.actS1Size + (blockId_ % 2) * CeilDiv(invalidS1Num, 2); + for (int innerS1Idx = 0; innerS1Idx < s1NumPerAiv; innerS1Idx++) { + CleanInvalidOutput(info.indiceOutOffset + (s1OffsetPerAiv + innerS1Idx) * constInfo_.sparseCount); + } + } + + int32_t invalidS1Num2 = info.actS1Size - info.actS2Size; + if (invalidS1Num2 > 0 && isS1LoopEnd && blockS2StartIdx_ == 0 && constInfo_.attenMaskFlag) { + int32_t s1NumPerAiv = blockId_ % 2 == 0 ? CeilDiv(invalidS1Num2, 2) : (invalidS1Num2 / 2); + int32_t s1OffsetPerAiv = (blockId_ % 2) * CeilDiv(invalidS1Num2, 2); + for (int innerS1Idx = 0; innerS1Idx < s1NumPerAiv; innerS1Idx++) { + CleanInvalidOutput((info.bN2Idx * constInfo_.qSeqSize + s1OffsetPerAiv + innerS1Idx) * + constInfo_.sparseCount); + } + } + } + + if (info.isLastS2InnerLoop) { + // S2最后一个Loop后, 下一个基本块初始从0开始 + blockS2StartIdx_ = 0; + } +} + +template +__aicore__ inline void LIQVector::ProcessLD() +{ + int32_t curCubeId = blockId_ / 2; + int32_t tmpCubeId = curCubeId; + + int64_t s2ActSeq; + int64_t s2Start; + int64_t s2End; + int64_t isS2End; + int64_t bn2Idx; + int64_t s1Idx; + uint32_t acc_list_num = 0; + int64_t bIdx = 0; + int64_t needFd; + int64_t wsOffset; + int64_t wsInfoOffset = 0; + int64_t nextneedFd; + int64_t valueOffset = 0; + int64_t outOffset = 0; + + LocalTensor curValueIdxUb = ldToBeMrgBuf_.Get(); + LocalTensor tmpUb = ldTmpBuf_.Get(); + + // S2开头信息 + // 开始必然没有头规约，因此从尾规约开始处理，while循环读取下一个核的头规约 + // 存满4个list或者遇到S2结尾，则做merge，直到做完S2 + // 每个核都忽略自己的头规约，因为必然由前面的核做完 + uint32_t s1LdStartIdx = 0; + uint32_t s1ProcNum = 0; + uint64_t paramGmCoreOffset = tmpCubeId * s1BaseSize_ * 2 * paramNum_; + for (uint32_t innerS1Idx = 0; innerS1Idx < s1BaseSize_; innerS1Idx++) { + needFd = vec1ParamGm.GetValue(paramGmCoreOffset + innerS1Idx * 2 * paramNum_ + paramNum_); + if (needFd == 1) { + s1LdStartIdx = (s1ProcNum == 0) ? innerS1Idx : s1LdStartIdx; + s1ProcNum++; + } + } + + if (s1ProcNum == 0) { + return; + } + + // S1逐行计算 + uint32_t s1VecNum = CeilDiv(s1ProcNum, 2); + if (blockId_ % 2 == 1) { + s1LdStartIdx = s1LdStartIdx + s1VecNum; + s1VecNum = s1ProcNum - s1VecNum; + } + for (uint32_t innerS1Idx = s1LdStartIdx; innerS1Idx < s1LdStartIdx + s1VecNum; innerS1Idx++) { + // 重置偏移 + tmpCubeId = curCubeId; + acc_list_num = 0; + valueOffset = 0; + + // 搬入数据 + wsOffset = tmpCubeId * s1BaseSize_ * 2 * BASE_TOPK_VALUE_IDX_SIZE + // 2个AIV共同地址偏移 + innerS1Idx * 2 * BASE_TOPK_VALUE_IDX_SIZE + BASE_TOPK_VALUE_IDX_SIZE; + SetWaitFlag(HardEvent::V_MTE2); + SetWaitFlag(HardEvent::S_MTE2); + DataCopyPad(curValueIdxUb, vec1ResGm[wsOffset], + {1, static_cast(BASE_TOPK_VALUE_IDX_SIZE * sizeof(int32_t)), 0, 0}, {true, 0, 0, 0}); + acc_list_num++; + valueOffset += BASE_TOPK_VALUE_IDX_SIZE; + + // 获取下一个核规约信息 + tmpCubeId++; + wsInfoOffset = tmpCubeId * s1BaseSize_ * 2 * paramNum_ + innerS1Idx * 2 * paramNum_; + needFd = vec1ParamGm.GetValue(wsInfoOffset); + isS2End = vec1ParamGm.GetValue(wsInfoOffset + 4); + s1Idx = vec1ParamGm.GetValue(wsInfoOffset + 6); + outOffset = vec1ParamGm.GetValue(wsInfoOffset + 8); + + while (needFd == 1) { + // 搬入头规约数据 + wsOffset = tmpCubeId * s1BaseSize_ * 2 * BASE_TOPK_VALUE_IDX_SIZE + // 2个AIV共同地址偏移 + innerS1Idx * 2 * BASE_TOPK_VALUE_IDX_SIZE; + SetWaitFlag(HardEvent::V_MTE2); + SetWaitFlag(HardEvent::S_MTE2); + DataCopyPad(curValueIdxUb[valueOffset], vec1ResGm[wsOffset], + {1, static_cast(BASE_TOPK_VALUE_IDX_SIZE * sizeof(int32_t)), 0, 0}, {true, 0, 0, 0}); + valueOffset += BASE_TOPK_VALUE_IDX_SIZE; + acc_list_num++; + + // 每满4个list，聚合 前2K为mrg结果 + if (acc_list_num == mrgListNum_) { + // MrgSort 四条2048的队列，Mrg成一条 + AscendC::MrgSort4Info params; + params.elementLengths[0] = BASE_TOPK; + params.elementLengths[1] = BASE_TOPK; + params.elementLengths[2] = BASE_TOPK; + params.elementLengths[3] = BASE_TOPK; + params.ifExhaustedSuspension = true; + params.validBit = 0b1111; + params.repeatTimes = 1; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = curValueIdxUb[0]; + srcList.src2 = curValueIdxUb[BASE_TOPK_VALUE_IDX_SIZE]; + srcList.src3 = curValueIdxUb[2 * BASE_TOPK_VALUE_IDX_SIZE]; + srcList.src4 = curValueIdxUb[3 * BASE_TOPK_VALUE_IDX_SIZE]; + SetWaitFlag(HardEvent::MTE2_V); + MrgSort(tmpUb, srcList, params); + PipeBarrier(); + DataCopy(curValueIdxUb, tmpUb, BASE_TOPK_VALUE_IDX_SIZE); + PipeBarrier(); + acc_list_num = 1; + valueOffset = BASE_TOPK_VALUE_IDX_SIZE; + } + + // reduce到S2末尾，则跳出 + if (isS2End == 1) { + break; + } + + tmpCubeId++; + wsInfoOffset = tmpCubeId * s1BaseSize_ * 2 * paramNum_ + innerS1Idx * 2 * paramNum_; + needFd = vec1ParamGm.GetValue(wsInfoOffset); + isS2End = vec1ParamGm.GetValue(wsInfoOffset + 4); + } + + // mrg不足4个list的数据 + if (acc_list_num != 1) { + AscendC::MrgSort4Info params; + params.elementLengths[0] = BASE_TOPK; + params.elementLengths[1] = BASE_TOPK; + params.elementLengths[2] = BASE_TOPK; + params.elementLengths[3] = BASE_TOPK; + params.ifExhaustedSuspension = true; + if (acc_list_num == 2) { + params.validBit = 0b0011; + } else if (acc_list_num == 3) { + params.validBit = 0b0111; + } + params.repeatTimes = 1; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = curValueIdxUb[0]; + srcList.src2 = curValueIdxUb[BASE_TOPK_VALUE_IDX_SIZE]; + srcList.src3 = curValueIdxUb[2 * BASE_TOPK_VALUE_IDX_SIZE]; + srcList.src4 = curValueIdxUb[3 * BASE_TOPK_VALUE_IDX_SIZE]; + SetWaitFlag(HardEvent::MTE2_V); + MrgSort(tmpUb, srcList, params); + PipeBarrier(); + DataCopy(curValueIdxUb, tmpUb, BASE_TOPK_VALUE_IDX_SIZE); + PipeBarrier(); + } + + // 搬出 + LocalTensor outValueUb = ldOutValueBuf_.Get(); + LocalTensor outIdxUb = ldOutIdxBuf_.Get(); + Extract(outValueUb, outIdxUb, curValueIdxUb, (BASE_TOPK / 32)); + LocalTensor idxULocal1 = outIdxUb.template ReinterpretCast(); + SetWaitFlag(HardEvent::V_MTE3); + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(indiceOutGm[outOffset], idxULocal1, + {1, static_cast(constInfo_.sparseCount * sizeof(int32_t)), 0, 0}); + SetWaitFlag(HardEvent::MTE3_V); + } +} +} // namespace LIQKernel +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_template_tiling_key.h b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_template_tiling_key.h new file mode 100644 index 000000000..165e62156 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_template_tiling_key.h @@ -0,0 +1,53 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_quant_template_tiling_key.h + * \brief + */ + +#ifndef TEMPLATE_TILING_KEY_LI_H_ +#define TEMPLATE_TILING_KEY_LI_H_ + +#include "ascendc/host_api/tiling/template_argument.h" + +#define LI_TPL_FP16 1 +#define LI_TPL_IN8 2 +#define LI_TPL_INT32 3 +#define LI_TPL_BF16 27 + +#define LIQ_LAYOUT_BSND 0 +#define LIQ_LAYOUT_TND 1 +#define LIQ_LAYOUT_PA_BSND 2 + +#define ASCENDC_TPL_4_BW 4 + +// 模板参数支持的范围定义 +ASCENDC_TPL_ARGS_DECL(LightningIndexerQuant, // 算子OpType + ASCENDC_TPL_DTYPE_DECL(DT_Q, LI_TPL_IN8), ASCENDC_TPL_DTYPE_DECL(DT_K, LI_TPL_IN8), + ASCENDC_TPL_DTYPE_DECL(DT_OUT, LI_TPL_INT32), ASCENDC_TPL_BOOL_DECL(PAGE_ATTENTION, 1, 0), + ASCENDC_TPL_UINT_DECL(Q_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, LIQ_LAYOUT_BSND, + LIQ_LAYOUT_TND), + ASCENDC_TPL_UINT_DECL(K_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, + LIQ_LAYOUT_PA_BSND, LIQ_LAYOUT_BSND, LIQ_LAYOUT_TND), ); + +// 支持的模板参数组合 +// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法 +ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, LI_TPL_IN8), ASCENDC_TPL_DTYPE_SEL(DT_K, LI_TPL_IN8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, LI_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, LIQ_LAYOUT_BSND, LIQ_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, LIQ_LAYOUT_PA_BSND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, LI_TPL_IN8), ASCENDC_TPL_DTYPE_SEL(DT_K, LI_TPL_IN8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, LI_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, LIQ_LAYOUT_BSND, LIQ_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, LIQ_LAYOUT_BSND, LIQ_LAYOUT_TND), ), ); + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_vector.h b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_vector.h new file mode 100644 index 000000000..d6a2e277e --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_quant/op_kernel/lightning_indexer_quant_vector.h @@ -0,0 +1,193 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_quant_vector.h + * \brief + */ +#ifndef LIGHTNING_INDEXER_QUANT_VECTOR_H +#define LIGHTNING_INDEXER_QUANT_VECTOR_H + +#include "kernel_operator.h" +#include "lightning_indexer_quant_vector.h" + +namespace LIQServiceVec { +using namespace AscendC; + +constexpr int32_t NEG_INF = 0xFF800000; +constexpr int32_t INVALID_INDEX = -1; +constexpr uint8_t VEC_REPEAT_MAX = 255; +constexpr uint8_t B32_VEC_ELM_NUM = 64; +constexpr uint8_t B32_BLOCK_ALIGN_NUM = 8; +constexpr uint8_t B32_VEC_REPEAT_STRIDE = 8; +constexpr uint64_t VEC_REPEAT_BYTES = 256; +constexpr int32_t CONST_TWO = 2; +constexpr int64_t VALUE_AND_INDEX_NUM = 2; +constexpr int64_t BLOCK_BYTES = 32; +constexpr int64_t MRG_QUE_0 = 0; +constexpr int64_t MRG_QUE_1 = 1; +constexpr int64_t MRG_QUE_2 = 2; +constexpr int64_t MRG_QUE_3 = 3; +constexpr int64_t MRG_BLOCK_2 = 2; +constexpr int64_t MRG_BLOCK_3 = 3; +constexpr int64_t MRG_BLOCK_4 = 4; + +template +__aicore__ inline void CopyOut(const GlobalTensor &dstGm, const LocalTensor &srcUb, int64_t copyCount) +{ + AscendC::DataCopyParams dataCopyOutyParams; + dataCopyOutyParams.blockCount = 1; + dataCopyOutyParams.blockLen = copyCount * sizeof(T); + dataCopyOutyParams.srcStride = 0; + dataCopyOutyParams.dstStride = 0; + AscendC::DataCopyPad(dstGm, srcUb, dataCopyOutyParams); +} + +/** + src: 传入的初始化空间 + eleNum: 需要初始化的元素个数需为64整数倍，元素将被初始化为交错排布的-inf，-1 + */ +__aicore__ inline void InitSortOutBuf(const LocalTensor &src, int64_t eleNum) +{ + uint64_t mask1[2] = {0x5555555555555555, 0}; + uint64_t mask0[2] = {0xaaaaaaaaaaaaaaaa, 0}; + int64_t repeatNum = eleNum / B32_VEC_ELM_NUM; + int64_t forLoop = repeatNum / VEC_REPEAT_MAX; + int64_t forRemain = repeatNum % VEC_REPEAT_MAX; + for (int i = 0; i < forLoop; i++) { + AscendC::Duplicate(src.template ReinterpretCast(), NEG_INF, mask1, VEC_REPEAT_MAX, 1, + B32_VEC_REPEAT_STRIDE); + AscendC::Duplicate(src.template ReinterpretCast(), INVALID_INDEX, mask0, VEC_REPEAT_MAX, 1, + B32_VEC_REPEAT_STRIDE); + } + if (forRemain > 0) { + AscendC::Duplicate(src.template ReinterpretCast()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM], NEG_INF, + mask1, forRemain, 1, B32_VEC_REPEAT_STRIDE); + AscendC::Duplicate(src.template ReinterpretCast()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM], + INVALID_INDEX, mask0, forRemain, 1, B32_VEC_REPEAT_STRIDE); + } + AscendC::PipeBarrier(); +} + +/** + src: logits和索引，前logitsNum为logits，后logitsNum为索引 + tmp: 计算使用到的临时空间，大小与src一致 + logitsNum: 排序的元素个数, 暂只支持[128,256,384,512,1024,2048] + */ +__aicore__ inline void SortAll(LocalTensor &src, LocalTensor &tmp, int64_t logitsNum) +{ + int64_t sort32Repeats = logitsNum / BLOCK_BYTES; + AscendC::Sort32(tmp, src, src[logitsNum].ReinterpretCast(), sort32Repeats); + AscendC::PipeBarrier(); + + int64_t mrgGroups = sort32Repeats; + int64_t mrgElements = BLOCK_BYTES; + int64_t i = 0; + AscendC::LocalTensor srcTensor; + AscendC::LocalTensor dstTensor; + while (true) { + if (i % CONST_TWO == 0) { + srcTensor = tmp; + dstTensor = src; + } else { + srcTensor = src; + dstTensor = tmp; + } + AscendC::MrgSort4Info params; + params.elementLengths[0] = mrgElements; + params.elementLengths[MRG_QUE_1] = mrgElements; + params.elementLengths[MRG_QUE_2] = mrgElements; + params.elementLengths[MRG_QUE_3] = mrgElements; + params.ifExhaustedSuspension = false; + params.validBit = 0b1111; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = srcTensor[0]; + srcList.src2 = srcTensor[MRG_QUE_1 * VALUE_AND_INDEX_NUM * mrgElements]; + srcList.src3 = srcTensor[MRG_QUE_2 * VALUE_AND_INDEX_NUM * mrgElements]; + srcList.src4 = srcTensor[MRG_QUE_3 * VALUE_AND_INDEX_NUM * mrgElements]; + if (mrgGroups <= MRG_BLOCK_4) { + params.repeatTimes = 1; + if (mrgGroups == 1) { + break; + } else if (mrgGroups == MRG_BLOCK_2) { + params.validBit = 0b0011; + } else if (mrgGroups == MRG_BLOCK_3) { + params.validBit = 0b0111; + } else if (mrgGroups == MRG_BLOCK_4) { + params.validBit = 0b1111; + } + AscendC::MrgSort(dstTensor, srcList, params); + i += 1; + break; + } else { + params.repeatTimes = mrgGroups / MRG_BLOCK_4; + AscendC::MrgSort(dstTensor, srcList, params); + i += 1; + mrgElements = mrgElements * MRG_BLOCK_4; + mrgGroups = mrgGroups / MRG_BLOCK_4; + } + AscendC::PipeBarrier(); + } + if (i % CONST_TWO == 0) { + AscendC::DataCopy(src, tmp, logitsNum * VALUE_AND_INDEX_NUM); + AscendC::PipeBarrier(); + } +} + +/** + mrgDst: 合并进的Tensor + mrgSrc: 待合并的Tensor + tmpTensor：空间为mrgDst+mrgSrc + */ +__aicore__ inline void MergeSort(const LocalTensor &mrgDst, int32_t mrgDstNum, LocalTensor &mrgSrc, + int32_t mrgSrcNum, LocalTensor &tmpTensor) +{ + AscendC::MrgSort4Info params; + params.elementLengths[0] = mrgSrcNum; + params.elementLengths[1] = mrgDstNum; + params.ifExhaustedSuspension = false; + params.validBit = 0b0011; + params.repeatTimes = 1; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = mrgSrc; + srcList.src2 = mrgDst; + + AscendC::MrgSort(tmpTensor, srcList, params); + AscendC::PipeBarrier(); + AscendC::DataCopy(mrgDst, tmpTensor, mrgDstNum * VALUE_AND_INDEX_NUM); + AscendC::PipeBarrier(); +} + +__aicore__ inline void ExtractIndex(const LocalTensor &idxULocal, const LocalTensor &sortLocal, + int64_t extractNum) +{ + AscendC::GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = Ceil(extractNum * sizeof(float) * VALUE_AND_INDEX_NUM, VEC_REPEAT_BYTES); + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = B32_VEC_REPEAT_STRIDE; + gatherMaskParams.src1RepeatStride = 0; + uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数 + uint8_t src1Pattern = 2; // 固定模式2,表示筛选出奇数索引的数 + AscendC::GatherMask(idxULocal, sortLocal, src1Pattern, false, static_cast(0), gatherMaskParams, rsvdCnt); + AscendC::PipeBarrier(); +} + +template +__aicore__ inline void SetWaitFlag(HardEvent evt) +{ + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(evt)); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); +} + +} // namespace LIQServiceVec +#endif // LIGHTNING_INDEXER_QUANT_VECTOR_H \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_vllm/CMakeLists.txt b/csrc/ascend/attention/lightning_indexer_vllm/CMakeLists.txt new file mode 100644 index 000000000..298ed5677 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_vllm/lightning_indexer_vllm_torch_adpt.h b/csrc/ascend/attention/lightning_indexer_vllm/lightning_indexer_vllm_torch_adpt.h new file mode 100644 index 000000000..537964582 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/lightning_indexer_vllm_torch_adpt.h @@ -0,0 +1,73 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef LIGHTING_INDEXER_VLLM_TORCH_ADPT_H +#define LIGHTING_INDEXER_VLLM_TORCH_ADPT_H +namespace vllm_fl { + +at::Tensor npu_lightning_indexer( + const at::Tensor &query, const at::Tensor &key, const at::Tensor &weights, + const c10::optional &actual_seq_lengths_query, + const c10::optional &actual_seq_lengths_key, + const c10::optional &block_table, c10::string_view layout_query, + c10::string_view layout_key, int64_t sparse_count, int64_t sparse_mode) +{ + // npu tensor max size + constexpr int32_t SIZE = 8; + constexpr int32_t DIM_0 = 0; + constexpr int32_t DIM_1 = 1; + constexpr int32_t DIM_2 = 2; + constexpr int32_t DIM_3 = 3; + + TORCH_CHECK(query.numel() > 0, "Query is empty."); + TORCH_CHECK(key.numel() > 0, "Key is empty."); + TORCH_CHECK(weights.numel() > 0, "Weights is empty."); + for (size_t i = 0; i < query.sizes().size(); i++) { + TORCH_CHECK(query.size(i) > 0, "All values within query's shape should be greater " + "than 0, but shape[", i, "] is ", query.size(i)); + } + TORCH_CHECK(sparse_count > 0, "sparse count should be greater than 0, but now is ", sparse_count); + + at::SmallVector output_size; + std::string query_layout_str = std::string(layout_query); + std::string key_layout_str = std::string(layout_key); + if (query_layout_str == "BSND") { + output_size = {query.size(DIM_0), query.size(DIM_1), key.size(DIM_2), sparse_count}; + } else { + int n_dim_index = 0; + n_dim_index = (key_layout_str == "TND") ? DIM_1 : DIM_2; + output_size = {query.size(DIM_0), key.size(n_dim_index), sparse_count}; + } + at::Tensor lightning_indexer_output = at::empty(output_size, query.options().dtype(at::kInt)); + // convert str + char *query_layout_ptr = const_cast(query_layout_str.c_str()); + char *key_layout_ptr = const_cast(key_layout_str.c_str()); + EXEC_NPU_CMD( + aclnnLightningIndexerVllm, + query, + key, + weights, + actual_seq_lengths_query, + actual_seq_lengths_key, + block_table, + query_layout_ptr, + key_layout_ptr, + sparse_count, + sparse_mode, + lightning_indexer_output); + return lightning_indexer_output; +} +} +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_vllm/op_host/CMakeLists.txt b/csrc/ascend/attention/lightning_indexer_vllm/op_host/CMakeLists.txt new file mode 100644 index 000000000..d8711d492 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/op_host/CMakeLists.txt @@ -0,0 +1,24 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + lightning_indexer_vllm_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME LightningIndexerVllm + OPTIONS + --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE lightning_indexer_vllm ACLNNTYPE aclnn) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/attention/lightning_indexer_vllm/op_host/lightning_indexer_vllm_def.cpp b/csrc/ascend/attention/lightning_indexer_vllm/op_host/lightning_indexer_vllm_def.cpp new file mode 100644 index 000000000..3df7dd64b --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/op_host/lightning_indexer_vllm_def.cpp @@ -0,0 +1,72 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_def.cpp + * \brief + */ +#include +#include "register/op_def_registry.h" + +namespace ops { +class LightningIndexerVllm : public OpDef { +public: + explicit LightningIndexerVllm(const char *name) : OpDef(name) + { + this->Input("query") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("key") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("weights") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("actual_seq_lengths_query") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("actual_seq_lengths_key") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("block_table") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("sparse_indices").ParamType(REQUIRED).DataTypeList({ge::DT_INT32}).FormatList({ge::FORMAT_ND}); + this->Attr("layout_query").AttrType(OPTIONAL).String("BSND"); + this->Attr("layout_key").AttrType(OPTIONAL).String("PA_BSND"); + this->Attr("sparse_count").AttrType(OPTIONAL).Int(2048); // 2048: Default value, filter the top 2048 + this->Attr("sparse_mode").AttrType(OPTIONAL).Int(3); // 3: Default value, only calculate the lower triangular matrix + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("aclnnSupport.value", "support_aclnn") + .ExtendCfgInfo("jitCompile.flag", "static_false,dynamic_false"); + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + } +}; +OP_ADD(LightningIndexerVllm); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_vllm/op_host/lightning_indexer_vllm_proto.cpp b/csrc/ascend/attention/lightning_indexer_vllm/op_host/lightning_indexer_vllm_proto.cpp new file mode 100644 index 000000000..8761d9cb8 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/op_host/lightning_indexer_vllm_proto.cpp @@ -0,0 +1,96 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_proto.cpp + * \brief + */ +#include +#include +#include "error/ops_error.h" + + +using namespace ge; + +namespace ops { +constexpr uint32_t QUERY_INDEX = 0; +constexpr uint32_t KEY_INDEX = 1; +constexpr uint32_t ACTUAL_SEQ_K_INDEX = 4; +constexpr uint32_t ATTR_QUERY_LAYOUT_INDEX = 0; +constexpr uint32_t ATTR_KEY_LAYOUT_INDEX = 1; +constexpr uint32_t ATTR_SPARSE_COUNT_INDEX = 2; + +static ge::graphStatus InferShapeLightningIndexer(gert::InferShapeContext *context) +{ + OPS_ERR_IF(context == nullptr, OPS_LOG_E("LightningIndexer", "InferShapeContext is nullptr!"), + return ge::GRAPH_FAILED); + const gert::Shape *queryShape = context->GetInputShape(QUERY_INDEX); + OPS_LOG_E_IF_NULL(context, queryShape, return ge::GRAPH_FAILED); + const gert::Shape *keyShape = context->GetInputShape(KEY_INDEX); + OPS_LOG_E_IF_NULL(context, keyShape, return ge::GRAPH_FAILED); + gert::Shape *outShape = context->GetOutputShape(0); + + auto attrs = context->GetAttrs(); + OPS_LOG_E_IF_NULL(context, attrs, return ge::GRAPH_FAILED); + const char *inputLayoutQueryPtr = attrs->GetAttrPointer(ATTR_QUERY_LAYOUT_INDEX); + OPS_LOG_E_IF_NULL(context, inputLayoutQueryPtr, return ge::GRAPH_FAILED); + const char *inputLayoutKeyPtr = attrs->GetAttrPointer(ATTR_KEY_LAYOUT_INDEX); + OPS_LOG_E_IF_NULL(context, inputLayoutKeyPtr, return ge::GRAPH_FAILED); + const int64_t *seleced_count = attrs->GetInt(ATTR_SPARSE_COUNT_INDEX); + OPS_LOG_E_IF_NULL(context, seleced_count, return ge::GRAPH_FAILED); + std::string inputLayoutQueryPtrStr = std::string(inputLayoutQueryPtr); + std::string inputLayoutKeyPtrStr = std::string(inputLayoutKeyPtr); + OPS_ERR_IF( + inputLayoutQueryPtrStr != "TND" && inputLayoutQueryPtrStr != "BSND", + OPS_LOG_E(context, "The attr layout_query should be TND or BSND, but got %s.", inputLayoutQueryPtrStr.c_str()), + return ge::GRAPH_FAILED); + + outShape->SetDimNum(queryShape->GetDimNum()); + if (inputLayoutQueryPtrStr == "BSND") { + OPS_ERR_IF( + queryShape->GetDimNum() != 4, + OPS_LOG_E(context, "Layout BSND, queryDims (%zu) must be 4!", queryShape->GetDimNum()), + return ge::GRAPH_FAILED); + outShape->SetDim(0, queryShape->GetDim(0)); // 0:Dim B + outShape->SetDim(1, queryShape->GetDim(1)); // 1:Dim S + outShape->SetDim(2, keyShape->GetDim(2)); // 2:Dim N + outShape->SetDim(3, *seleced_count); // 3:Dim K + } else { + OPS_ERR_IF( + queryShape->GetDimNum() != 3, + OPS_LOG_E(context, "Layout TND, queryDims (%zu) must be 3!", queryShape->GetDimNum()), + return ge::GRAPH_FAILED); + outShape->SetDim(0, queryShape->GetDim(0)); // 0:Dim T + int32_t nDimIndex = (inputLayoutKeyPtrStr == "PA_BSND") ? 2 : 1; // 2:Key Dim N + outShape->SetDim(1, keyShape->GetDim(nDimIndex)); // 1:Dim N + outShape->SetDim(2, *seleced_count); // 2:Dim K + } + OPS_LOG_D(context->GetNodeName(), "LightningIndexer InferShape end."); + + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataTypeLightningIndexer(gert::InferDataTypeContext *context) +{ + OPS_ERR_IF(context == nullptr, OPS_LOG_E("LightningIndexer", "InferDataTypeContext is nullptr!"), + return ge::GRAPH_FAILED); + OPS_LOG_D(context->GetNodeName(), "Enter LightningIndexer InferDataType impl."); + // default set q's dtype as fia's output type + ge::DataType outputType = ge::DT_INT32; + // attention_out, outidx:0 + context->SetOutputDataType(0, outputType); + OPS_LOG_D(context->GetNodeName(), "LightningIndexer InferDataType end."); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(LightningIndexerVllm) + .InferShape(InferShapeLightningIndexer) + .InferDataType(InferDataTypeLightningIndexer); +} // namespace ops diff --git a/csrc/ascend/attention/lightning_indexer_vllm/op_host/lightning_indexer_vllm_tiling.cpp b/csrc/ascend/attention/lightning_indexer_vllm/op_host/lightning_indexer_vllm_tiling.cpp new file mode 100644 index 000000000..1355ff5bf --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/op_host/lightning_indexer_vllm_tiling.cpp @@ -0,0 +1,694 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_tiling.cpp + * \brief + */ + +#include "lightning_indexer_vllm_tiling.h" +#include "../op_kernel/lightning_indexer_template_tiling_key.h" + +using namespace ge; +using namespace AscendC; +using std::map; +using std::string; +namespace optiling { +ge::graphStatus LIInfoParser::CheckRequiredInOutExistence() const +{ + OPS_ERR_IF(opParamInfo_.query.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor query is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.query.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor query is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.key.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor k is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.key.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor k is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.weights.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor value is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.weights.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor value is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.attenOut.shape == nullptr, OPS_LOG_E(opName_, "Shape of tensor output is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.attenOut.desc == nullptr, OPS_LOG_E(opName_, "Desc of tensor output is nullptr"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::CheckRequiredAttrExistence() const +{ + OPS_ERR_IF(opParamInfo_.layOut == nullptr, OPS_LOG_E(opName_, "attr layout_query is nullptr"), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(opParamInfo_.layOutKey == nullptr, OPS_LOG_E(opName_, "attr layout_key is nullptr"), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(opParamInfo_.sparseCount == nullptr, OPS_LOG_E(opName_, "attr sparse_count is nullptr"), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(opParamInfo_.sparseMode == nullptr, OPS_LOG_E(opName_, "attr sparse_mode is nullptr"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::CheckRequiredParaExistence() const +{ + if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS || CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetOpName() +{ + if (context_->GetNodeName() == nullptr) { + OPS_LOG_E("LightningIndexer", "opName got from TilingContext is nullptr"); + return ge::GRAPH_FAILED; + } + opName_ = context_->GetNodeName(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetNpuInfo() +{ + platformInfo_ = context_->GetPlatformInfo(); + OPS_ERR_IF(platformInfo_ == nullptr, OPS_LOG_E(opName_, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo_); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + OPS_ERR_IF(aicNum == 0 || aivNum == 0, OPS_LOG_E(opName_, "num of core obtained is 0."), return GRAPH_FAILED); + + socVersion_ = ascendcPlatform.GetSocVersion(); + if ((socVersion_ != platform_ascendc::SocVersion::ASCEND910B) && + (socVersion_ != platform_ascendc::SocVersion::ASCEND910_93)) { + OPS_LOG_E(opName_, "SOC Version[%d] is not support.", (int32_t)socVersion_); + return GRAPH_FAILED; + } + OPS_ERR_IF(context_->GetWorkspaceSizes(1) == nullptr, OPS_LOG_E(opName_, "workSpaceSize got from ge is nullptr"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(context_->GetRawTilingData() == nullptr, + OPS_LOG_E(context_->GetNodeName(), "RawTilingData got from GE context is nullptr."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void LIInfoParser::GetOptionalInputParaInfo() +{ + opParamInfo_.actualSeqLengthsQ.tensor = context_->GetOptionalInputTensor(ACTUAL_SEQ_Q_INDEX); + opParamInfo_.actualSeqLengthsQ.desc = context_->GetOptionalInputDesc(ACTUAL_SEQ_Q_INDEX); + opParamInfo_.actualSeqLengths.tensor = context_->GetOptionalInputTensor(ACTUAL_SEQ_K_INDEX); + opParamInfo_.actualSeqLengths.desc = context_->GetOptionalInputDesc(ACTUAL_SEQ_K_INDEX); + opParamInfo_.blockTable.tensor = context_->GetOptionalInputTensor(BLOCK_TABLE_INDEX); + opParamInfo_.blockTable.desc = context_->GetOptionalInputDesc(BLOCK_TABLE_INDEX); +} + +void LIInfoParser::GetInputParaInfo() +{ + opParamInfo_.query.desc = context_->GetInputDesc(QUERY_INDEX); + opParamInfo_.query.shape = context_->GetInputShape(QUERY_INDEX); + opParamInfo_.key.desc = context_->GetInputDesc(KEY_INDEX); + opParamInfo_.key.shape = context_->GetInputShape(KEY_INDEX); + opParamInfo_.weights.desc = context_->GetInputDesc(WEIGTHS_INDEX); + opParamInfo_.weights.shape = context_->GetInputShape(WEIGTHS_INDEX); + GetOptionalInputParaInfo(); +} + +void LIInfoParser::GetOutputParaInfo() +{ + opParamInfo_.attenOut.desc = context_->GetOutputDesc(LIGHTNING_INDEXER); + opParamInfo_.attenOut.shape = context_->GetOutputShape(LIGHTNING_INDEXER); +} + +ge::graphStatus LIInfoParser::GetAndCheckAttrParaInfo() +{ + auto attrs = context_->GetAttrs(); + OPS_ERR_IF(attrs == nullptr, OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "attrs got from ge is nullptr"), + return ge::GRAPH_FAILED); + + OPS_LOG_I(context_->GetNodeName(), "GetAndCheckAttrParaInfo start"); + opParamInfo_.layOut = attrs->GetStr(ATTR_QUERY_LAYOUT_INDEX); + opParamInfo_.layOutKey = attrs->GetStr(ATTR_KEY_LAYOUT_INDEX); + opParamInfo_.sparseCount = attrs->GetAttrPointer(ATTR_SPARSE_COUNT_INDEX); + opParamInfo_.sparseMode = attrs->GetAttrPointer(ATTR_SPARSE_MODE_INDEX); + + if (opParamInfo_.layOut != nullptr) { + OPS_LOG_I(context_->GetNodeName(), "layout_query is:%s", opParamInfo_.layOut); + } + if (opParamInfo_.layOutKey != nullptr) { + OPS_LOG_I(context_->GetNodeName(), "layout_key is:%s", opParamInfo_.layOutKey); + } + if (opParamInfo_.sparseCount != nullptr) { + OPS_LOG_I(context_->GetNodeName(), "selscted count is:%d", *opParamInfo_.sparseCount); + } + if (opParamInfo_.sparseMode != nullptr) { + OPS_LOG_I(context_->GetNodeName(), "sparse mode is:%d", *opParamInfo_.sparseMode); + } + OPS_LOG_I(context_->GetNodeName(), "GetAndCheckAttrParaInfo end"); + + OPS_ERR_IF( + ((std::string(opParamInfo_.layOutKey) != "PA_BSND") + && (std::string(opParamInfo_.layOut) != std::string(opParamInfo_.layOutKey))), + OPS_LOG_E(opName_, "under non-PA conditions, layout_query and layout_key should be equal."), + return ge::GRAPH_FAILED); + OPS_ERR_IF( + ((std::string(opParamInfo_.layOutKey) != "PA_BSND") && (std::string(opParamInfo_.layOutKey) != "BSND") + && (std::string(opParamInfo_.layOutKey) != "TND")), + OPS_LOG_E(opName_, "input attr layout_key only supported PA_BSND, BSND or TND"), return ge::GRAPH_FAILED); + OPS_ERR_IF(((std::string(opParamInfo_.layOut) != "BSND") && (std::string(opParamInfo_.layOut) != "TND")), + OPS_LOG_E(opName_, "input attr layout_query only supported BSND or TND."), return ge::GRAPH_FAILED); + OPS_ERR_IF(!((*opParamInfo_.sparseCount > 0) && (*opParamInfo_.sparseCount <= SPARSE_LIMIT)), + OPS_LOG_E(opName_, "input attr sparse_count must > 0 and <= 2048."), return ge::GRAPH_FAILED); + OPS_ERR_IF(!((*opParamInfo_.sparseMode == 0) || (*opParamInfo_.sparseMode == SPARSE_MODE_LOWER)), + OPS_LOG_E(opName_, "input attr sparse_mode only supported 0 or 3."), return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetOpParaInfo() +{ + GetInputParaInfo(); + GetOutputParaInfo(); + if (ge::GRAPH_SUCCESS != GetAndCheckAttrParaInfo()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetAndCheckInOutDataType() +{ + inputQType_ = opParamInfo_.query.desc->GetDataType(); + inputKType_ = opParamInfo_.key.desc->GetDataType(); + weightsType_ = opParamInfo_.weights.desc->GetDataType(); + outputType_ = opParamInfo_.attenOut.desc->GetDataType(); + + bool inDTypeAllEqual = (inputQType_ == inputKType_) && (inputKType_ == weightsType_); + OPS_ERR_IF(!inDTypeAllEqual, + OPS_LOG_E(opName_, "The data types of the input query, key, and weights must be the same."), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(((inputQType_ != ge::DT_FLOAT16) && (inputQType_ != ge::DT_BF16)), + OPS_LOG_E(opName_, "The data types of the input query, key, and weights must be float16 or bfloat16."), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(outputType_ != ge::DT_INT32, + OPS_LOG_E(opName_, "The data types of the output sparse_indices must be int32."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetQueryKeyAndOutLayout() +{ + const map layoutMap = { + {"BSND", DataLayout::BSND}, + {"TND", DataLayout::TND}, + {"PA_BSND", DataLayout::BnBsND} + }; + + std::string layout(opParamInfo_.layOut); + auto it = layoutMap.find(layout); + if (it != layoutMap.end()) { + qLayout_ = it->second; + } + + std::string layoutKey(opParamInfo_.layOutKey); + auto itKey = layoutMap.find(layoutKey); + if (itKey != layoutMap.end()) { + kLayout_ = itKey->second; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetAndCheckOptionalInput() +{ + if (kLayout_ == DataLayout::BnBsND) { + OPS_ERR_IF(opParamInfo_.blockTable.tensor == nullptr, + OPS_LOG_E(opName_, "key layout only supported PA_BSND, input block_table must not be null"), + return ge::GRAPH_FAILED); + OPS_ERR_IF( + opParamInfo_.actualSeqLengths.tensor == nullptr, + OPS_LOG_E(opName_, "key layout only supported PA_BSND, input actual_seq_lengths_key must not be null"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.blockTable.desc->GetDataType() != ge::DT_INT32, + OPS_LOG_E(opName_, "input block_table data type only support int32"), return ge::GRAPH_FAILED); + } else if (kLayout_ == DataLayout::TND) { + OPS_ERR_IF(opParamInfo_.actualSeqLengths.tensor == nullptr, + OPS_LOG_E(opName_, "when layout_key is TND, input actual_seq_lengths_key must not be null"), + return ge::GRAPH_FAILED); + } + OPS_ERR_IF(opParamInfo_.actualSeqLengths.tensor != nullptr && + opParamInfo_.actualSeqLengths.desc->GetDataType() != ge::DT_INT32, + OPS_LOG_E(opName_, "input actual_seq_lengths_key data type only support int32"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(opParamInfo_.actualSeqLengths.tensor != nullptr && + opParamInfo_.actualSeqLengths.desc->GetDataType() != ge::DT_INT32, + OPS_LOG_E(opName_, "input actual_seq_lengths_key data type only support int32"), + return ge::GRAPH_FAILED); + if (qLayout_ == DataLayout::TND) { + OPS_ERR_IF(opParamInfo_.actualSeqLengthsQ.tensor == nullptr, + OPS_LOG_E(opName_, "when layout_query is TND, input actual_seq_lengths_query must not be null"), + return ge::GRAPH_FAILED); + } + OPS_ERR_IF(opParamInfo_.actualSeqLengthsQ.tensor != nullptr && + opParamInfo_.actualSeqLengthsQ.desc->GetDataType() != ge::DT_INT32, + OPS_LOG_E(opName_, "input actual_seq_lengths_query data type only support int32"), + return ge::GRAPH_FAILED); + OPS_ERR_IF(kLayout_ != DataLayout::BnBsND && opParamInfo_.blockTable.tensor != nullptr, + OPS_LOG_E(opName_, "when key layout is not PA_BSND, input block_table must be null"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::CheckShapeDim() +{ + OPS_ERR_IF((opParamInfo_.blockTable.tensor != nullptr) && + (opParamInfo_.blockTable.tensor->GetStorageShape().GetDimNum() != DIM_NUM_TWO), + OPS_LOG_E(opName_, "the dim num of block_table's shape should be 2"), return ge::GRAPH_FAILED); + + uint32_t kShapeDim = opParamInfo_.key.shape->GetStorageShape().GetDimNum(); + uint32_t qShapeDim = opParamInfo_.query.shape->GetStorageShape().GetDimNum(); + uint32_t weightsShapeDim = opParamInfo_.weights.shape->GetStorageShape().GetDimNum(); + uint32_t outShapeDim = opParamInfo_.attenOut.shape->GetStorageShape().GetDimNum(); + uint32_t qExpectShapeDim = DIM_NUM_FOUR; + uint32_t kExpectShapeDim = DIM_NUM_FOUR; + if (qLayout_ == DataLayout::TND) { + qExpectShapeDim = DIM_NUM_THREE; + } + if (kLayout_ == DataLayout::TND) { + kExpectShapeDim = DIM_NUM_THREE; + } + OPS_ERR_IF(kShapeDim != kExpectShapeDim, + OPS_LOG_E(opName_, "the dim num of key's shape should be %u, but now is %u", kExpectShapeDim, kShapeDim), + return ge::GRAPH_FAILED); + OPS_ERR_IF(qShapeDim != qExpectShapeDim, + OPS_LOG_E(opName_, "the dim num of query's shape should be %u, but now is %u", + qExpectShapeDim, qShapeDim), + return ge::GRAPH_FAILED); + OPS_ERR_IF(outShapeDim != qExpectShapeDim, + OPS_LOG_E(opName_, "the dim num of sparse_indices's shape should be %u, but now is %u", + qExpectShapeDim, outShapeDim), + return ge::GRAPH_FAILED); + OPS_ERR_IF(!(weightsShapeDim == qExpectShapeDim - 1), + OPS_LOG_E(opName_, "the dim num of weights's shape should be %u, but now is %u", qExpectShapeDim - 1, + weightsShapeDim), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetN1Size() +{ + if (qLayout_ == DataLayout::BSND) { + n1Size_ = static_cast(opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_TWO)); + } else { + // TND + n1Size_ = static_cast(opParamInfo_.query.shape->GetStorageShape().GetDim(1)); + } + OPS_LOG_I(context_->GetNodeName(), "n1Size is %d", n1Size_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + const std::string &actualSeqLenName) +{ + size = static_cast(tensor->GetShapeSize()); + if (size <= 0) { + OPS_LOG_E(opName_, "%s's shape size is %u, it should be greater than 0.", actualSeqLenName.c_str(), size); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetAndCheckN2Size() +{ + uint32_t n2Index = (kLayout_ == DataLayout::TND) ? DIM_IDX_ONE : DIM_IDX_TWO; + n2Size_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(n2Index)); + OPS_LOG_I(context_->GetNodeName(), "n2Size_ is %d", n2Size_); + OPS_ERR_IF(n2Size_ != 1, OPS_LOG_E(opName_, "key shape[%u] is numhead, only support 1.", n2Index), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetGSize() +{ + if (n1Size_ % n2Size_ != 0) { + OPS_LOG_E(opName_, "input query's head_num %u can not be a multiple of key's head_num %u.", n1Size_, n2Size_); + return ge::GRAPH_FAILED; + } + gSize_ = n1Size_ / n2Size_; + OPS_ERR_IF(gSize_ != 64, OPS_LOG_E(opName_, "N1 is %u, N2 is %u, N1 divided by N2 must equal 64.", + n1Size_, n2Size_), return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetBatchSize() +{ + if ((qLayout_ == DataLayout::TND)) { + return GetActualSeqLenSize(bSize_, opParamInfo_.actualSeqLengthsQ.tensor, "input actual_seq_lengths_query"); + } else { // BSND + bSize_ = opParamInfo_.query.shape->GetStorageShape().GetDim(0); + return ge::GRAPH_SUCCESS; + } +} + +ge::graphStatus LIInfoParser::GetHeadDim() +{ + uint32_t dIndex = DIM_IDX_TWO; + switch (qLayout_) { + case DataLayout::TND: + // TND: [Total, N, D] -> D is the 2nd dimension + dIndex = DIM_IDX_TWO; + break; + case DataLayout::BSND: + // BSND: [Batch, SeqLen, N, D] -> D is the 3rd dimension + dIndex = DIM_IDX_THREE; + break; + default: + OPS_LOG_E(opName_, "unsupported layout for getting head dim."); + return ge::GRAPH_FAILED; + } + headDim_ = opParamInfo_.query.shape->GetStorageShape().GetDim(dIndex); + OPS_ERR_IF(headDim_ != HEAD_DIM_LIMIT, OPS_LOG_E(opName_, "input query's last dim head_dim only support 128."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetS1Size() +{ + if (qLayout_ == DataLayout::BSND) { + s1Size_ = opParamInfo_.query.shape->GetStorageShape().GetDim(1); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetAndCheckBlockSize() +{ + blockSize_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(1)); + OPS_LOG_I(context_->GetNodeName(), "blockSize_ is %d", blockSize_); + + OPS_ERR_IF(((blockSize_ % 16 != 0) || (blockSize_ == 0) || (blockSize_ > 1024)), + OPS_LOG_E(opName_, "input key's block_size must be a multiple of 16 and belong to (0, 1024]."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::CheckBlockCount() +{ + int32_t blockCount_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(0)); + OPS_ERR_IF((blockCount_ == 0), + OPS_LOG_E(opName_, "input key's block_count cannot be 0."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetS2SizeForPageAttention() +{ + if (GetAndCheckBlockSize() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + if (CheckBlockCount() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + maxBlockNumPerBatch_ = opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1); + s2Size_ = maxBlockNumPerBatch_ * blockSize_; + OPS_LOG_I(context_->GetNodeName(), "maxBlockNumPerBatch_ is %d, blockSize_ is %d, s2Size_ is %d", + maxBlockNumPerBatch_, blockSize_, s2Size_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::GetS2Size() +{ + if (kLayout_ == DataLayout::BnBsND) { + return GetS2SizeForPageAttention(); + } else if (kLayout_ == DataLayout::TND) { + s2Size_ = opParamInfo_.key.shape->GetStorageShape().GetDim(0); + } else if (kLayout_ == DataLayout::BSND) { + s2Size_ = opParamInfo_.key.shape->GetStorageShape().GetDim(1); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::ValidateInputShapesMatchQTnd() +{ + // -----------------------check BatchSize------------------- + if (kLayout_ == DataLayout::TND) { + OPS_ERR_IF( + (opParamInfo_.actualSeqLengths.tensor->GetShapeSize() != bSize_), + OPS_LOG_E(opName_, + "TND case input actual_seq_lengths_query, actual_seq_lengths_key are %u, %ld respectively, they must be same.", + bSize_, opParamInfo_.actualSeqLengths.tensor->GetShapeSize()), + return ge::GRAPH_FAILED); + } else { // kLayout_ PA_BSND + OPS_ERR_IF( + (opParamInfo_.actualSeqLengths.tensor->GetShapeSize() != bSize_) || + (opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) != bSize_), + OPS_LOG_E( + opName_, + "TND case input actual_seq_lengths_query, actual_seq_lengths_key, block_table dim 0 are %u, %ld, %ld respectively, they must be same.", + bSize_, opParamInfo_.actualSeqLengths.tensor->GetShapeSize(), + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + } + // -----------------------check T------------------- + uint32_t qTsize = opParamInfo_.query.shape->GetStorageShape().GetDim(0); + OPS_ERR_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != qTsize) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != qTsize), + OPS_LOG_E(opName_, "TND case input query, weights, sparse_indices dim 0 are %u, %ld, %ld respectively, they must be same.", + qTsize, opParamInfo_.weights.shape->GetStorageShape().GetDim(0), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::ValidateInputShapesMatchQBsnd() +{ + // -----------------------check BatchSize------------------- + if (kLayout_ == DataLayout::BnBsND) { + OPS_ERR_IF((opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) != bSize_) || + (opParamInfo_.actualSeqLengths.tensor->GetShapeSize() != bSize_), + OPS_LOG_E(opName_, "BSND case input query, actual_seq_lengths_key, block_table dim 0 are %u, %ld, %ld respectively, they must be same.", + bSize_, opParamInfo_.actualSeqLengths.tensor->GetShapeSize(), + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + } else if (kLayout_ == DataLayout::BSND) { + OPS_ERR_IF(opParamInfo_.key.shape->GetStorageShape().GetDim(0) != bSize_, + OPS_LOG_E(opName_, "BSND case input query, key dim 0 are %u, %ld respectively, they must be same.", + bSize_, opParamInfo_.key.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + OPS_ERR_IF((opParamInfo_.actualSeqLengths.tensor != nullptr) && + (opParamInfo_.actualSeqLengths.tensor->GetShapeSize() != bSize_), + OPS_LOG_E(opName_, "BSND case input query, actual_seq_lengths_key dim 0 are %u, %ld respectively, they must be same.", + bSize_, opParamInfo_.actualSeqLengths.tensor->GetShapeSize()), + return ge::GRAPH_FAILED); + } + OPS_ERR_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != bSize_) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != bSize_), + OPS_LOG_E(opName_, "BSND case input query, weight, sparse_indices dim 0 are %u, %ld, %ld respectively, they must be same.", + bSize_, opParamInfo_.weights.shape->GetStorageShape().GetDim(0), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + OPS_ERR_IF((opParamInfo_.actualSeqLengthsQ.tensor != nullptr) && + (opParamInfo_.actualSeqLengthsQ.tensor->GetShapeSize() != bSize_), + OPS_LOG_E(opName_, "BSND case input query, actual_seq_lengths_query dim 0 are %u, %ld respectively, they must be same.", + bSize_, opParamInfo_.actualSeqLengthsQ.tensor->GetShapeSize()), + return ge::GRAPH_FAILED); + // -----------------------check S1------------------- + OPS_ERR_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(1) != s1Size_) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(1) != s1Size_), + OPS_LOG_E(opName_, "BSND case input query, weight, sparse_indices dim 1 are %u, %ld, %ld, they must be same.", + s1Size_, opParamInfo_.weights.shape->GetStorageShape().GetDim(1), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(1)), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LIInfoParser::ValidateInputShapesMatch() +{ + uint32_t queryWeightsN1Dim = 1; + uint32_t outN2Dim = 1; + if (qLayout_ == DataLayout::TND) { + if (ValidateInputShapesMatchQTnd() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + } else { + if (ValidateInputShapesMatchQBsnd() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + queryWeightsN1Dim = DIM_IDX_TWO; + outN2Dim = DIM_IDX_TWO; + } + // -----------------------check N1------------------- + OPS_ERR_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(queryWeightsN1Dim) != n1Size_), + OPS_LOG_E(opName_, "input query, weight shape dim N1 must be same."), return ge::GRAPH_FAILED); + // -----------------------check D------------------- + uint32_t keyDDim = kLayout_ == DataLayout::TND ? DIM_IDX_TWO : DIM_IDX_THREE; + OPS_ERR_IF((opParamInfo_.key.shape->GetStorageShape().GetDim(keyDDim) != headDim_), + OPS_LOG_E(opName_, "input query, key shape last dim must be same."), return ge::GRAPH_FAILED); + // -----------------------check N2------------------- + OPS_ERR_IF((opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim) != n2Size_), + OPS_LOG_E(opName_, "input query and output sparse_indices shape n2 dim must be same."), + return ge::GRAPH_FAILED); + // -----------------------check sparse_count------------------- + OPS_ERR_IF((opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim + 1) != *opParamInfo_.sparseCount), + OPS_LOG_E(opName_, "output sparse_indices shape last dim must be same as attr sparse_count."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void LIInfoParser::GenerateInfo(LITilingInfo &liInfo) +{ + liInfo.opName = opName_; + liInfo.platformInfo = platformInfo_; + liInfo.opParamInfo = opParamInfo_; + liInfo.socVersion = socVersion_; + + liInfo.bSize = bSize_; + liInfo.n1Size = n1Size_; + liInfo.n2Size = n2Size_; + liInfo.s1Size = s1Size_; + liInfo.s2Size = s2Size_; + liInfo.gSize = gSize_; + + liInfo.inputQType = inputQType_; + liInfo.inputKType = inputKType_; + liInfo.outputType = outputType_; + + liInfo.blockSize = blockSize_; + liInfo.maxBlockNumPerBatch = maxBlockNumPerBatch_; + + std::string layOutKeyStr(opParamInfo_.layOutKey); + liInfo.pageAttentionFlag = layOutKeyStr == "PA_BSND" ? true : false; + liInfo.sparseMode = *opParamInfo_.sparseMode; + liInfo.sparseCount = *opParamInfo_.sparseCount; + + liInfo.inputQLayout = qLayout_; + liInfo.inputKLayout = kLayout_; +} + +ge::graphStatus LIInfoParser::ParseAndCheck(LITilingInfo &liInfo) +{ + if (ge::GRAPH_SUCCESS != GetOpName() || ge::GRAPH_SUCCESS != GetNpuInfo() || ge::GRAPH_SUCCESS != GetOpParaInfo() || + ge::GRAPH_SUCCESS != CheckRequiredParaExistence()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetAndCheckInOutDataType() || ge::GRAPH_SUCCESS != GetQueryKeyAndOutLayout() || + ge::GRAPH_SUCCESS != GetAndCheckOptionalInput()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != CheckShapeDim() || ge::GRAPH_SUCCESS != GetN1Size() || + ge::GRAPH_SUCCESS != GetAndCheckN2Size() || ge::GRAPH_SUCCESS != GetGSize()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetBatchSize() || ge::GRAPH_SUCCESS != GetS1Size() || ge::GRAPH_SUCCESS != GetHeadDim() || + ge::GRAPH_SUCCESS != GetS2Size()) { + return ge::GRAPH_FAILED; + } + if (ge::GRAPH_SUCCESS != ValidateInputShapesMatch()) { + return ge::GRAPH_FAILED; + } + + GenerateInfo(liInfo); + + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus TilingPrepareForLightningIndexer(gert::TilingParseContext * /* context */) +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus LightningIndexerTiling::DoTiling(LITilingInfo *tilingInfo) +{ + // -------------set blockdim----------------- + auto ascendcPlatform = platform_ascendc::PlatformAscendC(tilingInfo->platformInfo); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + uint32_t blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum); + context_->SetBlockDim(blockDim); + + // -------------set workspacesize----------------- + constexpr uint32_t MM1_RES_ELEM_SIZE = 4; + constexpr uint32_t DOUBLE_BUFFER = 2; + constexpr uint32_t M_BASE_SIZE = 512; + constexpr uint32_t S2_BASE_SIZE = 512; + constexpr uint32_t V1_RES_ELEM_SIZE = 4; + constexpr uint32_t V1_RES_ELEM_TYPE = 2; + constexpr uint32_t V1_DECODE_PARAM_ELEM_SIZE = 8; + constexpr uint32_t V1_DECODE_PARAM_NUM = 16; + constexpr uint32_t V1_DECODE_DATA_NUM = 2; + constexpr uint32_t S1_BASE_SIZE = 8; + constexpr uint32_t TOPK_MAX_SIZE = 2048; + uint32_t workspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize(); + uint32_t mm1ResSize = M_BASE_SIZE * S2_BASE_SIZE; + workspaceSize += mm1ResSize * MM1_RES_ELEM_SIZE * DOUBLE_BUFFER * aicNum; + workspaceSize += V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_RES_ELEM_TYPE * TOPK_MAX_SIZE * V1_RES_ELEM_SIZE * aicNum; + workspaceSize += V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_DECODE_PARAM_NUM * V1_DECODE_PARAM_ELEM_SIZE * aicNum; + size_t *workSpaces = context_->GetWorkspaceSizes(1); + workSpaces[0] = workspaceSize; + + // -------------set tilingdata----------------- + tilingData_.set_bSize(tilingInfo->bSize); + tilingData_.set_s2Size(tilingInfo->s2Size); + tilingData_.set_s1Size(tilingInfo->s1Size); + tilingData_.set_sparseCount(tilingInfo->sparseCount); + tilingData_.set_gSize(tilingInfo->gSize); + tilingData_.set_blockSize(tilingInfo->blockSize); + tilingData_.set_maxBlockNumPerBatch(tilingInfo->maxBlockNumPerBatch); + tilingData_.set_sparseMode(tilingInfo->sparseMode); + tilingData_.set_usedCoreNum(blockDim); + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + + // -------------set tilingkey----------------- + // DT_Q, DT_KV, DT_OUT, PAGE_ATTENTION, FLASH_DECODE, LAYOUT_T, KV_LAYOUT_T + uint32_t inputQType = static_cast(tilingInfo->inputQType); + uint32_t inputKType = static_cast(tilingInfo->inputKType); + uint32_t outputType = static_cast(tilingInfo->outputType); + uint32_t pageAttentionFlag = static_cast(tilingInfo->pageAttentionFlag); + uint32_t inputQLayout = static_cast(tilingInfo->inputQLayout); + uint32_t inputKLayout = static_cast(tilingInfo->inputKLayout); + uint32_t tilingKey = + GET_TPL_TILING_KEY(inputQType, inputKType, outputType, pageAttentionFlag, inputQLayout, inputKLayout); + context_->SetTilingKey(tilingKey); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingForLightningIndexer(gert::TilingContext *context) +{ + OPS_ERR_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("LightningIndexer", "Tiling context is null."), + return ge::GRAPH_FAILED); + LITilingInfo liInfo; + LIInfoParser LIInfoParser(context); + if (LIInfoParser.ParseAndCheck(liInfo) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + LightningIndexerTiling liTiling(context); + return liTiling.DoTiling(&liInfo); +} + +IMPL_OP_OPTILING(LightningIndexerVllm) + .Tiling(TilingForLightningIndexer) + .TilingParse(TilingPrepareForLightningIndexer); + +} // namespace optiling diff --git a/csrc/ascend/attention/lightning_indexer_vllm/op_host/lightning_indexer_vllm_tiling.h b/csrc/ascend/attention/lightning_indexer_vllm/op_host/lightning_indexer_vllm_tiling.h new file mode 100644 index 000000000..bf5156a9b --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/op_host/lightning_indexer_vllm_tiling.h @@ -0,0 +1,215 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_tiling.h + * \brief + */ + +#ifndef LIGHTNING_INDEXER_TILING_H_ +#define LIGHTNING_INDEXER_TILING_H_ + +#include "exe_graph/runtime/tiling_context.h" +#include "tiling/platform/platform_ascendc.h" +#include "register/op_def_registry.h" +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "error/ops_error.h" +#include "platform/platform_info.h" + +namespace optiling { + +struct TilingRequiredParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; +}; + +struct TilingOptionalParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::Tensor *tensor; +}; + +enum class DataLayout : uint32_t { + BSND = 0, + TND = 1, + BnBsND = 2 +}; + +// Inputs Index +constexpr uint32_t QUERY_INDEX = 0; +constexpr uint32_t KEY_INDEX = 1; +constexpr uint32_t WEIGTHS_INDEX = 2; +constexpr uint32_t ACTUAL_SEQ_Q_INDEX = 3; +constexpr uint32_t ACTUAL_SEQ_K_INDEX = 4; +constexpr uint32_t BLOCK_TABLE_INDEX = 5; +constexpr uint32_t LIGHTNING_INDEXER = 0; +// Attributes Index +constexpr uint32_t ATTR_QUERY_LAYOUT_INDEX = 0; +constexpr uint32_t ATTR_KEY_LAYOUT_INDEX = 1; +constexpr uint32_t ATTR_SPARSE_COUNT_INDEX = 2; +constexpr uint32_t ATTR_SPARSE_MODE_INDEX = 3; +// Dim Index +constexpr uint32_t DIM_IDX_ONE = 1; +constexpr uint32_t DIM_IDX_TWO = 2; +constexpr uint32_t DIM_IDX_THREE = 3; +// Dim Num +constexpr uint32_t DIM_NUM_TWO = 2; +constexpr uint32_t DIM_NUM_THREE = 3; +constexpr uint32_t DIM_NUM_FOUR = 4; +// Input Parameter Limit Constant +constexpr uint32_t HEAD_DIM_LIMIT = 128; +constexpr uint32_t SPARSE_LIMIT = 2048; +constexpr uint32_t SPARSE_MODE_LOWER = 3; + +BEGIN_TILING_DATA_DEF(LITilingData) +TILING_DATA_FIELD_DEF(uint32_t, bSize) +TILING_DATA_FIELD_DEF(uint32_t, n2Size) +TILING_DATA_FIELD_DEF(uint32_t, gSize) +TILING_DATA_FIELD_DEF(uint32_t, s1Size) +TILING_DATA_FIELD_DEF(uint32_t, s2Size) +TILING_DATA_FIELD_DEF(uint32_t, sparseCount) +TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum) +TILING_DATA_FIELD_DEF(uint32_t, blockSize) +TILING_DATA_FIELD_DEF(uint32_t, maxBlockNumPerBatch) +TILING_DATA_FIELD_DEF(uint32_t, sparseMode) +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(LightningIndexerVllm, LITilingData) + +struct LICompileInfo {}; + +struct LiParaInfo { + TilingRequiredParaInfo query = {nullptr, nullptr}; + TilingRequiredParaInfo key = {nullptr, nullptr}; + TilingRequiredParaInfo weights = {nullptr, nullptr}; + TilingOptionalParaInfo actualSeqLengthsQ = {nullptr, nullptr}; + TilingOptionalParaInfo actualSeqLengths = {nullptr, nullptr}; + TilingOptionalParaInfo blockTable = {nullptr, nullptr}; + TilingRequiredParaInfo attenOut = {nullptr, nullptr}; + + const char *layOut = nullptr; + const char *layOutKey = nullptr; + const int32_t *blockSize = nullptr; + const int32_t *sparseMode = nullptr; + const int32_t *sparseCount = nullptr; +}; + +class LITilingInfo { +public: + const char *opName = nullptr; + fe::PlatFormInfos *platformInfo = nullptr; + LiParaInfo opParamInfo; + // Base Param + platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND910B; + uint32_t bSize = 0; + uint32_t n1Size = 0; + uint32_t n2Size = 0; + uint32_t s1Size = 0; + int64_t s2Size = 0; + uint32_t qkHeadDim = 0; + uint32_t gSize = 0; + // PageAttention + bool pageAttentionFlag = false; + int32_t blockSize = 0; + uint32_t maxBlockNumPerBatch = 0; + // Mask + int32_t sparseMode = 0; + // Others Flag + uint32_t sparseCount = 0; + // DType + ge::DataType inputQType = ge::DT_FLOAT16; + ge::DataType inputKType = ge::DT_FLOAT16; + ge::DataType outputType = ge::DT_INT32; + // Layout + DataLayout inputQLayout = DataLayout::BSND; + DataLayout inputKLayout = DataLayout::BnBsND; +}; + +class LIInfoParser { +public: + explicit LIInfoParser(gert::TilingContext *context) : context_(context) + { + } + ~LIInfoParser() = default; + + ge::graphStatus CheckRequiredInOutExistence() const; + ge::graphStatus CheckRequiredAttrExistence() const; + ge::graphStatus CheckRequiredParaExistence() const; + ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + const std::string &actualSeqLenName); + ge::graphStatus GetOpName(); + ge::graphStatus GetNpuInfo(); + void GetOptionalInputParaInfo(); + void GetInputParaInfo(); + void GetOutputParaInfo(); + ge::graphStatus GetAndCheckAttrParaInfo(); + ge::graphStatus GetOpParaInfo(); + ge::graphStatus ValidateInputShapesMatchQBsnd(); + ge::graphStatus ValidateInputShapesMatchQTnd(); + ge::graphStatus ValidateInputShapesMatch(); + ge::graphStatus GetAndCheckInOutDataType(); + ge::graphStatus GetBatchSize(); + ge::graphStatus GetHeadDim(); + ge::graphStatus GetS1Size(); + ge::graphStatus GetAndCheckOptionalInput(); + ge::graphStatus CheckShapeDim(); + ge::graphStatus GetAndCheckBlockSize(); + ge::graphStatus CheckBlockCount(); + ge::graphStatus GetS2SizeForPageAttention(); + ge::graphStatus GetS2Size(); + ge::graphStatus GetQueryKeyAndOutLayout(); + ge::graphStatus GetN1Size(); + ge::graphStatus GetAndCheckN2Size(); + ge::graphStatus GetGSize(); + ge::graphStatus GetAttenMaskInfo(); + ge::graphStatus GetActualSeqInfo(); + void GenerateInfo(LITilingInfo &liInfo); + ge::graphStatus ParseAndCheck(LITilingInfo &liInfo); + +public: + gert::TilingContext *context_ = nullptr; + const char *opName_; + fe::PlatFormInfos *platformInfo_; + LiParaInfo opParamInfo_; + + // BaseParams + uint32_t bSize_ = 0; + uint32_t n1Size_ = 0; + uint32_t n2Size_ = 0; + uint32_t gSize_ = 0; + uint32_t s1Size_ = 0; + int64_t s2Size_ = 0; + uint32_t headDim_ = 0; + // Layout + DataLayout qLayout_ = DataLayout::BSND; + DataLayout kLayout_ = DataLayout::BnBsND; + // PageAttention + uint32_t maxBlockNumPerBatch_ = 0; + int32_t blockSize_ = 0; + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B; + ge::DataType inputQType_ = ge::DT_FLOAT16; + ge::DataType inputKType_ = ge::DT_FLOAT16; + ge::DataType weightsType_ = ge::DT_FLOAT16; + ge::DataType blockTableType_ = ge::DT_FLOAT16; + ge::DataType inputKRopeType_ = ge::DT_FLOAT16; + ge::DataType outputType_ = ge::DT_FLOAT16; +}; + +class LightningIndexerTiling { +public: + explicit LightningIndexerTiling(gert::TilingContext *context) : context_(context){}; + ge::graphStatus DoTiling(LITilingInfo *tilingInfo); + +private: + gert::TilingContext *context_ = nullptr; + LITilingData tilingData_; +}; + +} // namespace optiling +#endif // LIGHTNING_INDEXER_TILING_H_ diff --git a/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_common.h b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_common.h new file mode 100644 index 000000000..4c693140d --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_common.h @@ -0,0 +1,135 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_common.h + * \brief + */ +#ifndef LIGHTNING_INDEXER_COMMON_H +#define LIGHTNING_INDEXER_COMMON_H + +namespace LICommon { +enum class LI_LAYOUT { + BSND = 0, + TND = 1, + PA_BSND = 2 +}; + +template +struct LIType { + using queryType = Q_T; + using keyType = K_T; + using outputType = OUT_T; + static constexpr bool pageAttention = PAGE_ATTENTION; + static constexpr LI_LAYOUT layout = LAYOUT_T; + static constexpr LI_LAYOUT keyLayout = K_LAYOUT_T; +}; + +struct RunInfo { + uint32_t loop; + uint32_t bN2Idx; + uint32_t bIdx; + uint32_t n2Idx = 0; + uint32_t gS1Idx; + uint32_t s2Idx; + + uint32_t actS1Size = 1; + uint32_t actS2Size = 1; + uint32_t actMBaseSize; + uint32_t actualSingleProcessSInnerSize; + uint32_t actualSingleProcessSInnerSizeAlign; + + uint64_t tensorQueryOffset; + uint64_t tensorKeyOffset; + uint64_t tensorWeightsOffset; + uint64_t indiceOutOffset; + + bool isFirstS2InnerLoop; + bool isLastS2InnerLoop; + bool isAllLoopEnd = false; +}; + +struct ConstInfo { + static constexpr uint32_t FIA_SYNC_MODE2 = 2; + static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; + static constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64; + static constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256; + static constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512; + static constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024; + static constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048; + static constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096; + static constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192; + static constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384; + static constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768; + static constexpr int INVALID_IDX = -1; + + uint32_t syncC1V1 = 0U; + uint32_t syncV1C1 = 0U; + + uint32_t mBaseSize = 1ULL; + uint32_t s1BaseSize = 1ULL; + uint32_t s2BaseSize = 1ULL; + + uint64_t batchSize = 0ULL; + uint64_t gSize = 0ULL; + uint64_t qHeadNum = 0ULL; + uint64_t kHeadNum; + uint64_t headDim; + uint64_t sparseCount; + uint64_t kSeqSize = 0ULL; + uint64_t qSeqSize = 1ULL; + uint32_t kCacheBlockSize = 0; + uint32_t maxBlockNumPerBatch = 0; + LI_LAYOUT outputLayout; + bool attenMaskFlag = false; + + uint32_t actualLenQDims = 0U; + uint32_t actualLenDims = 0U; + bool isAccumSeqS1 = false; + bool isAccumSeqS2 = false; +}; + +struct SplitCoreInfo { + uint32_t s2Start = 0U; + uint32_t s2End = 0U; + uint32_t bN2Start = 0U; + uint32_t bN2End = 0U; + uint32_t gS1Start = 0U; + uint32_t gS1End = 0U; + bool isLD = false; +}; + +template +__aicore__ inline T Align(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd) * (rnd))); +} + +template +__aicore__ inline T1 Min(T1 a, T2 b) +{ + return (a > b) ? (b) : (a); +} + +template +__aicore__ inline T1 Max(T1 a, T2 b) +{ + return (a > b) ? (a) : (b); +} + +template +__aicore__ inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd))); +} +} // namespace LICommon + +#endif // LIGHTNING_INDEXER_COMMON_H \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_kernel.h b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_kernel.h new file mode 100644 index 000000000..14ef69781 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_kernel.h @@ -0,0 +1,623 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_kernel.h + * \brief + */ + +#ifndef LIGHTNING_INDEXER_KERNEL_H +#define LIGHTNING_INDEXER_KERNEL_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "lightning_indexer_common.h" +#include "lightning_indexer_service_vector.h" +#include "lightning_indexer_service_cube.h" + +namespace LIKernel { +using namespace LICommon; +using namespace LIServiceVec; +using namespace matmul; +using AscendC::CacheMode; +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +struct TempLoopInfo { + uint32_t bN2Idx = 0; + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint32_t gS1Idx = 0U; + uint32_t gS1LoopEnd = 0U; + uint32_t s2LoopEnd = 0U; + uint32_t actS1Size = 1ULL; + uint32_t actS2Size = 0ULL; + bool curActSeqLenIsZero = false; + bool needDealActS1LessThanS1 = false; + uint32_t actMBaseSize = 0U; + uint32_t mBasicSizeTail = 0U; + uint32_t s2BasicSizeTail = 0U; +}; + +template +class LIPreload { +public: + __aicore__ inline LIPreload(){}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths, + __gm__ uint8_t *blockTable, __gm__ uint8_t *sparseIndices, __gm__ uint8_t *workspace, + const LITilingData *__restrict tiling, TPipe *tPipe); + __aicore__ inline void Process(); + + using Q_T = typename LIT::queryType; + using K_T = typename LIT::keyType; + using OUT_T = typename LIT::outputType; + static constexpr bool PAGE_ATTENTION = LIT::pageAttention; + static constexpr LI_LAYOUT LAYOUT_T = LIT::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = LIT::keyLayout; + + using MM1_OUT_T = float; + + LIMatmul matmulService; + LIVector vectorService; + + static constexpr uint32_t SYNC_C1_V1_FLAG = 4; + static constexpr uint32_t SYNC_V1_C1_FLAG = 5; + + static constexpr uint32_t M_BASE_SIZE = 512; + static constexpr uint32_t S2_BASE_SIZE = 512; + static constexpr uint32_t HEAD_DIM = 128; + static constexpr uint32_t K_HEAD_NUM = 1; + static constexpr uint32_t GM_ALIGN_BYTES = 512; + + static constexpr int64_t LD_PREFETCH_LEN = 2; + // for workspace double + static constexpr uint32_t WS_DOBULE = 2; + +protected: + TPipe *pipe = nullptr; + + // offset + uint64_t queryCoreOffset = 0ULL; + uint64_t keyCoreOffset = 0ULL; + uint64_t weightsCoreOffset = 0ULL; + uint64_t indiceOutCoreOffset = 0ULL; + + GlobalTensor queryGm; + GlobalTensor keyGm; + GlobalTensor weightsGm; + + GlobalTensor indiceOutGm; + GlobalTensor blockTableGm; + + GlobalTensor actualSeqLengthsGmQ; + GlobalTensor actualSeqLengthsGm; + // workspace + GlobalTensor mm1ResGm; + GlobalTensor vec1ResGm; + GlobalTensor vec1ParamGm; + + // aic、aiv kernel info + uint32_t tmpBlockIdx = 0U; + uint32_t aiCoreIdx = 0U; + uint32_t usedCoreNum = 0U; + + LICommon::ConstInfo constInfo{}; + TempLoopInfo tempLoopInfo{}; + LICommon::SplitCoreInfo splitCoreInfo{}; + + // ================================Init functions================================== + __aicore__ inline void InitTilingData(const LITilingData *__restrict tilingData); + __aicore__ inline void InitBuffers(); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths); + // ================================Split Core================================ + __aicore__ inline void SplitCore(uint32_t curCoreIdx, uint32_t &coreNum, LICommon::SplitCoreInfo &info); + __aicore__ inline uint32_t GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, uint32_t actS2Size); + __aicore__ inline uint32_t GetTotalBaseBlockNum(); + // ================================Process functions================================ + __aicore__ inline void ProcessMain(); + __aicore__ inline void ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, LICommon::RunInfo &runInfo); + __aicore__ inline void ProcessDecode(); + __aicore__ inline void ProcessInvalid(); + // ================================Params Calc===================================== + __aicore__ inline void CalcGS1LoopParams(uint32_t bN2Idx); + __aicore__ inline void GetBN2Idx(uint32_t bN2Idx); + __aicore__ inline uint32_t GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, uint32_t defaultSeqLen); + __aicore__ inline void GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, uint32_t &actS2Size); + __aicore__ inline void CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx); + __aicore__ inline void CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, LICommon::RunInfo &runInfo); + __aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start); +}; + +template +__aicore__ inline void LIPreload::InitTilingData(const LITilingData *__restrict tilingData) +{ + usedCoreNum = tilingData->usedCoreNum; + constInfo.batchSize = tilingData->bSize; + constInfo.qHeadNum = constInfo.gSize = tilingData->gSize; + constInfo.kSeqSize = tilingData->s2Size; + constInfo.qSeqSize = tilingData->s1Size; + constInfo.attenMaskFlag = (tilingData->sparseMode == 3); + constInfo.kCacheBlockSize = tilingData->blockSize; + constInfo.maxBlockNumPerBatch = tilingData->maxBlockNumPerBatch; + constInfo.sparseCount = tilingData->sparseCount; + constInfo.outputLayout = LAYOUT_T; + if (LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS1 = true; + } + if (K_LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS2 = true; + } + + constInfo.kHeadNum = K_HEAD_NUM; + constInfo.headDim = HEAD_DIM; + + constInfo.mBaseSize = M_BASE_SIZE; + constInfo.s2BaseSize = S2_BASE_SIZE; + constInfo.s1BaseSize = (constInfo.mBaseSize + constInfo.gSize - 1) / constInfo.gSize; +} + +template +__aicore__ inline void LIPreload::InitBuffers() +{ + if ASCEND_IS_AIV { + vectorService.InitBuffers(pipe); + } else { + matmulService.InitBuffers(pipe); + } +} + +template +__aicore__ inline void LIPreload::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengths) +{ + if (actualSeqLengthsQ == nullptr) { + constInfo.actualLenQDims = 0; + } else { + constInfo.actualLenQDims = constInfo.batchSize; + actualSeqLengthsGmQ.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsQ, constInfo.actualLenQDims); + } + if (actualSeqLengths == nullptr) { + constInfo.actualLenDims = 0; + } else { + constInfo.actualLenDims = constInfo.batchSize; + actualSeqLengthsGm.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengths, constInfo.actualLenDims); + } +} + +template +__aicore__ inline uint32_t LIPreload::GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, + uint32_t defaultSeqLen) +{ + if (actualLenDims == 0) { + return defaultSeqLen; + } else if (isAccumSeq && bIdx > 0) { + return actualSeqLengthsGm.GetValue(bIdx) - actualSeqLengthsGm.GetValue(bIdx - 1); + } else { + return actualSeqLengthsGm.GetValue(bIdx); + } +} + +template +__aicore__ inline void LIPreload::GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, uint32_t &actS2Size) +{ + actS1Size = GetActualSeqLen(bIdx, constInfo.actualLenQDims, constInfo.isAccumSeqS1, actualSeqLengthsGmQ, + constInfo.qSeqSize); + actS2Size = + GetActualSeqLen(bIdx, constInfo.actualLenDims, constInfo.isAccumSeqS2, actualSeqLengthsGm, constInfo.kSeqSize); +} + +template +__aicore__ inline uint32_t LIPreload::GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, + uint32_t actS2Size) +{ + if (actS2Size == 0) { + return 0; + } + uint32_t s1Offset = constInfo.s1BaseSize * s1gIdx; + int32_t validS2LenBase = static_cast(actS2Size) - static_cast(actS1Size); + int32_t validS2Len = s1Offset + validS2LenBase + constInfo.s1BaseSize; + validS2Len = Min(validS2Len, static_cast(actS2Size)); + validS2Len = Max(validS2Len, 1); + return (validS2Len + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; +} + +template +__aicore__ inline uint32_t LIPreload::GetTotalBaseBlockNum() +{ + uint32_t totalBlockNum = 0; + uint32_t actS1Size, actS2Size; + uint32_t s1GBaseNum, s2BaseNum; + for (uint32_t bIdx = 0; bIdx < constInfo.batchSize; bIdx++) { + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size); + s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize); + if (!constInfo.attenMaskFlag) { + s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize); + totalBlockNum += s1GBaseNum * s2BaseNum * constInfo.kHeadNum; + continue; + } + for (uint32_t s1gIdx = 0; s1gIdx < s1GBaseNum; s1gIdx++) { + s2BaseNum = GetS2BaseBlockNumOnMask(s1gIdx, actS1Size, actS2Size); + totalBlockNum += s2BaseNum * constInfo.kHeadNum; + } + } + return totalBlockNum; +} + +template +__aicore__ void inline LIPreload::SplitCore(uint32_t curCoreIdx, uint32_t &coreNum, LICommon::SplitCoreInfo &info) +{ + uint32_t totalBlockNum = GetTotalBaseBlockNum(); + uint32_t minBlockPerCore = totalBlockNum / coreNum; + uint32_t deal1MoreBlockCoreNum = totalBlockNum % coreNum; + uint32_t coreIdx = 0; + uint32_t lastGS1RemainBlockCnt = 0; + uint32_t coreDealBlockCnt = coreIdx < deal1MoreBlockCoreNum ? minBlockPerCore + 1 : minBlockPerCore; + coreNum = minBlockPerCore == 0 ? deal1MoreBlockCoreNum : coreNum; + + bool findLastCoreEnd = true; + uint32_t actS1Size, actS2Size; + uint32_t s1GBaseNum, s2BaseNum; + for (uint32_t bN2Idx = 0; bN2Idx < constInfo.batchSize * constInfo.kHeadNum; bN2Idx++) { + uint32_t bIdx = bN2Idx / constInfo.kHeadNum; + if (bN2Idx % constInfo.kHeadNum == 0) { + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size); + s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize); + s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize); + } + if constexpr (LAYOUT_T == LI_LAYOUT::BSND) { + if (findLastCoreEnd && (s1GBaseNum == 0U || s2BaseNum == 0U)) { + info.bN2Start = bN2Idx; + info.gS1Start = 0; + info.s2Start = 0; + findLastCoreEnd = false; + } + } + for (uint32_t gS1Idx = 0; gS1Idx < s1GBaseNum; gS1Idx++) { + if (constInfo.attenMaskFlag) { + s2BaseNum = GetS2BaseBlockNumOnMask(gS1Idx, actS1Size, actS2Size); + } + if (findLastCoreEnd && s2BaseNum == 0U) { + info.bN2Start = bN2Idx; + info.gS1Start = gS1Idx; + info.s2Start = 0; + findLastCoreEnd = false; + } + for (uint32_t s2Idx = 0; s2Idx < s2BaseNum;) { + if (findLastCoreEnd) { + info.bN2Start = bN2Idx; + info.gS1Start = gS1Idx; + info.s2Start = s2Idx; + findLastCoreEnd = false; + } + uint32_t s2RemainBaseNum = s2BaseNum - s2Idx; + if (lastGS1RemainBlockCnt + s2RemainBaseNum >= coreDealBlockCnt) { + info.bN2End = bN2Idx; + info.gS1End = gS1Idx; + info.s2End = s2Idx + coreDealBlockCnt - lastGS1RemainBlockCnt - 1; + + if (coreIdx == curCoreIdx) { + if (s2Idx == 0 && info.s2End + 1 < s2BaseNum) { + info.isLD = true; + } + if (coreIdx == coreNum - 1 && info.bN2End != constInfo.batchSize -1) { + info.bN2End = constInfo.batchSize -1; + info.gS1End = 0; + info.s2End = 0; + } + return; + } + coreIdx++; + findLastCoreEnd = true; + s2Idx = info.s2End + 1; + lastGS1RemainBlockCnt = 0; + coreDealBlockCnt = coreIdx < deal1MoreBlockCoreNum ? minBlockPerCore + 1 : minBlockPerCore; + } else { + lastGS1RemainBlockCnt += s2RemainBaseNum; + break; + } + } + } + } +} + +template +__aicore__ inline void LIPreload::DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start) +{ + if ASCEND_IS_AIV { + if (constInfo.outputLayout == LI_LAYOUT::TND) { + uint32_t tSize = actualSeqLengthsGmQ.GetValue(constInfo.batchSize - 1); + uint32_t tBase = bIdx == 0 ? 0 : actualSeqLengthsGmQ.GetValue(bIdx - 1); + uint32_t s1Count = tempLoopInfo.actS1Size; + + for (uint32_t s1Idx = s1Start; s1Idx < s1Count; s1Idx++) { + uint64_t indiceOutOffset = + (tBase + s1Idx) * constInfo.kHeadNum * constInfo.sparseCount + + n2Idx * constInfo.sparseCount; + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } else if (constInfo.outputLayout == LI_LAYOUT::BSND) { + for (uint32_t s1Idx = s1Start; s1Idx < constInfo.qSeqSize; s1Idx++) { + // B,S1,N2,K + uint64_t indiceOutOffset = bIdx * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount + + s1Idx * constInfo.kHeadNum * constInfo.sparseCount + + n2Idx * constInfo.sparseCount; + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } + } +} + +template +__aicore__ inline void LIPreload::Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths, + __gm__ uint8_t *blockTable, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *workspace, const LITilingData *__restrict tiling, + TPipe *tPipe) +{ + if ASCEND_IS_AIV { + tmpBlockIdx = GetBlockIdx(); // vec:0-47 + aiCoreIdx = tmpBlockIdx / 2; + } else { + tmpBlockIdx = GetBlockIdx(); // cube:0-23 + aiCoreIdx = tmpBlockIdx; + } + + InitTilingData(tiling); + InitActualSeqLen(actualSeqLengthsQ, actualSeqLengths); + + SplitCore(aiCoreIdx, usedCoreNum, splitCoreInfo); + + pipe = tPipe; + uint64_t offset = 0; + uint64_t singleCoreMm1ResSize = WS_DOBULE * constInfo.mBaseSize * constInfo.s2BaseSize * sizeof(MM1_OUT_T); + mm1ResGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset + aiCoreIdx * singleCoreMm1ResSize)); + offset += GetBlockNum() * singleCoreMm1ResSize; + + vec1ResGm.SetGlobalBuffer((__gm__ float *)(workspace + offset)); + offset += GetBlockNum() * constInfo.s1BaseSize * WS_DOBULE * WS_DOBULE * BASE_TOPK * sizeof(float); + + vec1ParamGm.SetGlobalBuffer((__gm__ int64_t *)(workspace + offset)); + offset += GetBlockNum() * constInfo.s1BaseSize * WS_DOBULE * LD_PARAM_NUM * sizeof(int64_t); + + if ASCEND_IS_AIV { + vectorService.InitParams(constInfo, tiling); + indiceOutGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices); + weightsGm.SetGlobalBuffer((__gm__ K_T *)weights); + vectorService.InitVec1GlobalTensor(mm1ResGm, vec1ResGm, vec1ParamGm, weightsGm, indiceOutGm); + } else { + matmulService.InitParams(constInfo); + queryGm.SetGlobalBuffer((__gm__ Q_T *)query); + if constexpr (PAGE_ATTENTION) { + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + } + keyGm.SetGlobalBuffer((__gm__ K_T *)key); + matmulService.InitMm1GlobalTensor(blockTableGm, keyGm, queryGm, mm1ResGm); + } + InitBuffers(); +} + +template +__aicore__ inline void LIPreload::GetBN2Idx(uint32_t bN2Idx) +{ + tempLoopInfo.bN2Idx = bN2Idx; + tempLoopInfo.bIdx = bN2Idx / constInfo.kHeadNum; + tempLoopInfo.n2Idx = bN2Idx % constInfo.kHeadNum; +} + +template +__aicore__ inline void LIPreload::CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx) +{ + tempLoopInfo.gS1Idx = gS1LoopIdx; + tempLoopInfo.actMBaseSize = constInfo.mBaseSize; + uint32_t remainedGS1Size = tempLoopInfo.actS1Size * constInfo.gSize - tempLoopInfo.gS1Idx * constInfo.mBaseSize; + if (remainedGS1Size <= constInfo.mBaseSize && remainedGS1Size > 0) { + tempLoopInfo.actMBaseSize = tempLoopInfo.mBasicSizeTail; + } + + bool isEnd = (bN2LoopIdx == splitCoreInfo.bN2End) && (gS1LoopIdx == splitCoreInfo.gS1End); + uint32_t s2BlockNum; + if (constInfo.attenMaskFlag) { + s2BlockNum = GetS2BaseBlockNumOnMask(gS1LoopIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2Size); + } else { + s2BlockNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + } + tempLoopInfo.s2LoopEnd = isEnd ? splitCoreInfo.s2End : s2BlockNum - 1; +} + +template +__aicore__ inline void LIPreload::CalcGS1LoopParams(uint32_t bN2LoopIdx) +{ + GetBN2Idx(bN2LoopIdx); + GetS1S2ActualSeqLen(tempLoopInfo.bIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2Size); + if ((tempLoopInfo.actS2Size == 0) || (tempLoopInfo.actS1Size == 0)) { + tempLoopInfo.curActSeqLenIsZero = true; + return; + } + tempLoopInfo.curActSeqLenIsZero = false; + tempLoopInfo.s2BasicSizeTail = tempLoopInfo.actS2Size % constInfo.s2BaseSize; + tempLoopInfo.s2BasicSizeTail = + (tempLoopInfo.s2BasicSizeTail == 0) ? constInfo.s2BaseSize : tempLoopInfo.s2BasicSizeTail; + tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize; + tempLoopInfo.mBasicSizeTail = + (tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail; + + uint32_t gS1SplitNum = (tempLoopInfo.actS1Size * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize; + tempLoopInfo.gS1LoopEnd = (bN2LoopIdx == splitCoreInfo.bN2End) ? splitCoreInfo.gS1End : gS1SplitNum - 1; + if constexpr (LAYOUT_T == LI_LAYOUT::BSND) { + if (tempLoopInfo.gS1LoopEnd == gS1SplitNum - 1 && constInfo.qSeqSize > tempLoopInfo.actS1Size) { + tempLoopInfo.needDealActS1LessThanS1 = true; + } + } +} + +template +__aicore__ inline void LIPreload::CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, LICommon::RunInfo &runInfo) +{ + runInfo.loop = loop; + runInfo.bIdx = tempLoopInfo.bIdx; + runInfo.gS1Idx = tempLoopInfo.gS1Idx; + runInfo.s2Idx = s2LoopIdx; + runInfo.bN2Idx = tempLoopInfo.bN2Idx; + + runInfo.actS1Size = tempLoopInfo.actS1Size; + runInfo.actS2Size = tempLoopInfo.actS2Size; + runInfo.actMBaseSize = tempLoopInfo.actMBaseSize; + runInfo.actualSingleProcessSInnerSize = constInfo.s2BaseSize; + uint32_t s2SplitNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + if (runInfo.s2Idx == s2SplitNum - 1) { + runInfo.actualSingleProcessSInnerSize = tempLoopInfo.s2BasicSizeTail; + } + runInfo.actualSingleProcessSInnerSizeAlign = + LICommon::Align((uint32_t)runInfo.actualSingleProcessSInnerSize, LICommon::ConstInfo::BUFFER_SIZE_BYTE_32B); + + runInfo.isFirstS2InnerLoop = s2LoopIdx == splitCoreInfo.s2Start; + runInfo.isLastS2InnerLoop = s2LoopIdx == tempLoopInfo.s2LoopEnd; + runInfo.isAllLoopEnd = (runInfo.bN2Idx == splitCoreInfo.bN2End) && (runInfo.gS1Idx == splitCoreInfo.gS1End) && + (runInfo.s2Idx == splitCoreInfo.s2End); + + if (runInfo.isFirstS2InnerLoop) { + uint64_t actualSeqQPrefixSum; + uint64_t actualSeqKPrefixSum; + if constexpr (LAYOUT_T == LI_LAYOUT::TND) { + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGmQ.GetValue(runInfo.bIdx - 1); + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGm.GetValue(runInfo.bIdx - 1); + } else { // BSND + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.qSeqSize; + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.kSeqSize; + } + uint64_t tndBIdxOffset = actualSeqQPrefixSum * constInfo.qHeadNum * constInfo.headDim; + uint64_t tndKeyBIdxOffset = actualSeqKPrefixSum * constInfo.kHeadNum * constInfo.headDim; + // B,S1,N1(N2,G),D + queryCoreOffset = tndBIdxOffset + runInfo.gS1Idx * constInfo.mBaseSize * constInfo.headDim; + keyCoreOffset = tndKeyBIdxOffset + runInfo.n2Idx * constInfo.headDim; + // B,S1,N1(N2,G)/T,N1(N2,G) + weightsCoreOffset = actualSeqQPrefixSum * constInfo.qHeadNum + runInfo.n2Idx * constInfo.gSize; + // B,S1,N2,k/T,N2,k + indiceOutCoreOffset = actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + + runInfo.n2Idx * constInfo.sparseCount; + } + runInfo.tensorQueryOffset = queryCoreOffset; + runInfo.tensorKeyOffset = keyCoreOffset + runInfo.s2Idx * constInfo.s2BaseSize * constInfo.kHeadNum + * constInfo.headDim; + runInfo.tensorWeightsOffset = weightsCoreOffset; + runInfo.indiceOutOffset = indiceOutCoreOffset; +} + +template +__aicore__ inline void LIPreload::Process() +{ + if (usedCoreNum == 0) { + ProcessInvalid(); + return; + } + ProcessMain(); + ProcessDecode(); +} + +template +__aicore__ inline void LIPreload::ProcessInvalid() +{ + if ASCEND_IS_AIV { + uint32_t aivCoreNum = GetBlockNum() * 2; // 2 means c:v = 1:2 + uint64_t totalOutputSize = + constInfo.batchSize * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount; + uint64_t singleCoreSize = + LICommon::Align((totalOutputSize + aivCoreNum - 1) / aivCoreNum, GM_ALIGN_BYTES / sizeof(OUT_T)); + uint64_t baseSize = tmpBlockIdx * singleCoreSize; + if (baseSize < totalOutputSize) { + uint64_t dealSize = + (baseSize + singleCoreSize > totalOutputSize) ? singleCoreSize : totalOutputSize - baseSize; + GlobalTensor output = indiceOutGm[baseSize]; + AscendC::InitGlobalMemory(output, dealSize, constInfo.INVALID_IDX); + } + } +} + +template +__aicore__ inline void LIPreload::ProcessMain() +{ + if (aiCoreIdx >= usedCoreNum) { + return; + } + + if ASCEND_IS_AIV { + vectorService.AllocEventID(); + CrossCoreSetFlag(constInfo.syncV1C1); + CrossCoreSetFlag(constInfo.syncV1C1); + } else { + matmulService.AllocEventID(); + } + + LICommon::RunInfo runInfo; + uint32_t gloop = 0; + for (uint32_t bN2LoopIdx = splitCoreInfo.bN2Start; bN2LoopIdx <= splitCoreInfo.bN2End; bN2LoopIdx++) { + CalcGS1LoopParams(bN2LoopIdx); + if (tempLoopInfo.curActSeqLenIsZero) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, 0U); + continue; + } + for (uint32_t gS1LoopIdx = splitCoreInfo.gS1Start; gS1LoopIdx <= tempLoopInfo.gS1LoopEnd; gS1LoopIdx++) { + CalcS2LoopParams(bN2LoopIdx, gS1LoopIdx); + for (int s2LoopIdx = splitCoreInfo.s2Start; s2LoopIdx <= tempLoopInfo.s2LoopEnd; s2LoopIdx++) { + ProcessBaseBlock(gloop, s2LoopIdx, runInfo); + ++gloop; + } + splitCoreInfo.s2Start = 0; + } + if (tempLoopInfo.needDealActS1LessThanS1) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, tempLoopInfo.actS1Size); + } + splitCoreInfo.gS1Start = 0; + } + + if ASCEND_IS_AIV { + vectorService.FreeEventID(); + } else { + matmulService.FreeEventID(); + CrossCoreWaitFlag(constInfo.syncV1C1); + CrossCoreWaitFlag(constInfo.syncV1C1); + } +} + +template +__aicore__ inline void LIPreload::ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, LICommon::RunInfo &runInfo) +{ + CalcRunInfo(loop, s2LoopIdx, runInfo); + if ASCEND_IS_AIC { + CrossCoreWaitFlag(constInfo.syncV1C1); + matmulService.ComputeMm1(runInfo); + CrossCoreSetFlag(constInfo.syncC1V1); + } else { + CrossCoreWaitFlag(constInfo.syncC1V1); + vectorService.ProcessVec(runInfo); + CrossCoreSetFlag(constInfo.syncV1C1); + } +} + +template +__aicore__ inline void LIPreload::ProcessDecode() +{ + if ASCEND_IS_AIV { + vectorService.InitLDBuffers(pipe); + ICachePreLoad(LD_PREFETCH_LEN); + SyncAll(); + if (splitCoreInfo.isLD) { + vectorService.ProcessLD(); + } + } +} +} // namespace LIKernel +#endif // LIGHTNING_INDEXER_KERNEL_H \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_service_cube.h b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_service_cube.h new file mode 100644 index 000000000..aa1888764 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_service_cube.h @@ -0,0 +1,415 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_service_cube.h + * \brief use 5 buffer for matmul l1, better pipeline + */ +#ifndef LIGHTNING_INDEXER_SERVICE_CUBE_H +#define LIGHTNING_INDEXER_SERVICE_CUBE_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "lightning_indexer_common.h" + +namespace LIKernel { +using namespace LICommon; +template +class LIMatmul { +public: + using Q_T = typename LIT::queryType; + using K_T = typename LIT::keyType; + + __aicore__ inline LIMatmul(){}; + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitMm1GlobalTensor(const GlobalTensor &blkTableGm, const GlobalTensor &keyGm, + const GlobalTensor &queryGm, const GlobalTensor &mm1ResGm); + __aicore__ inline void InitParams(const ConstInfo &constInfo); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void ComputeMm1(const LICommon::RunInfo &runInfo); + + static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding; + static constexpr uint64_t KEY_BUF_NUM = 3; + static constexpr uint64_t QUERY_BUF_NUM = 2; + static constexpr uint64_t L0_BUF_NUM = 2; + + static constexpr uint32_t KEY_MTE1_MTE2_EVENT = EVENT_ID2; + static constexpr uint32_t QUERY_MTE1_MTE2_EVENT = EVENT_ID5; // KEY_MTE1_MTE2_EVENT + KEY_BUF_NUM; + static constexpr uint32_t M_MTE1_EVENT = EVENT_ID3; + + static constexpr uint32_t MTE2_MTE1_EVENT = EVENT_ID2; + static constexpr uint32_t MTE1_M_EVENT = EVENT_ID2; + + static constexpr uint64_t M_BASIC_BLOCK = 256; + static constexpr uint64_t D_BASIC_BLOCK = 128; + static constexpr uint64_t S2_BASIC_BLOCK = 256; + + static constexpr uint64_t M_BASIC_BLOCK_L0 = 128; + static constexpr uint64_t D_BASIC_BLOCK_L0 = 128; + static constexpr uint64_t S2_BASIC_BLOCK_L0 = 128; + + static constexpr uint64_t QUERY_BUFFER_OFFSET = M_BASIC_BLOCK * D_BASIC_BLOCK; + static constexpr uint64_t KEY_BUFFER_OFFSET = S2_BASIC_BLOCK * D_BASIC_BLOCK; + static constexpr uint64_t L0AB_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0; + static constexpr uint64_t L0C_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0; + +protected: + __aicore__ inline void Fixp(uint64_t s1gGmOffset, uint64_t s2GmOffset, uint64_t s1gL0RealSize, + uint64_t s2L0RealSize, const LICommon::RunInfo &runInfo); + __aicore__ inline void ComuteL0c(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, const LICommon::RunInfo &runInfo); + __aicore__ inline void LoadKeyToL0b(uint64_t s2L0Offset, uint64_t s2L1RealSize, uint64_t s2L0RealSize, + const LICommon::RunInfo &runInfo); + __aicore__ inline void LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL0Offset, uint64_t s1gL1RealSize, + uint64_t s1gL0RealSize, const LICommon::RunInfo &runInfo); + __aicore__ inline void QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gL1Offset, const LICommon::RunInfo &runInfo); + __aicore__ inline void KeyNd2Nz(uint64_t s2L1RealSize, uint64_t s2GmOffset, const LICommon::RunInfo &runInfo); + __aicore__ inline void KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, const LICommon::RunInfo &runInfo); + GlobalTensor blkTableGm_; + GlobalTensor keyGm_; + GlobalTensor queryGm_; + GlobalTensor mm1ResGm_; + + TBuf bufQL1_; + LocalTensor queryL1_; + TBuf bufKeyL1_; + LocalTensor keyL1_; + + TBuf bufQL0_; + LocalTensor queryL0_; + TBuf bufKeyL0_; + LocalTensor keyL0_; + + TBuf bufL0C_; + LocalTensor cL0_; + + uint64_t keyL1BufIdx_ = 0; + uint64_t queryL1Mte2BufIdx_ = 0; + uint64_t queryL1Mte1BufIdx_ = 0; + uint64_t l0BufIdx_ = 0; + + ConstInfo constInfo_; + +private: + static constexpr bool PAGE_ATTENTION = LIT::pageAttention; +}; + +template +__aicore__ inline void LIMatmul::InitParams(const ConstInfo &constInfo) +{ + constInfo_ = constInfo; +} + +template +__aicore__ inline void LIMatmul::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(bufQL1_, QUERY_BUF_NUM * M_BASIC_BLOCK * D_BASIC_BLOCK * sizeof(Q_T)); + queryL1_ = bufQL1_.Get(); + pipe->InitBuffer(bufKeyL1_, KEY_BUF_NUM * S2_BASIC_BLOCK * D_BASIC_BLOCK * sizeof(K_T)); + keyL1_ = bufKeyL1_.Get(); + + pipe->InitBuffer(bufQL0_, L0_BUF_NUM * M_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0 * sizeof(Q_T)); + queryL0_ = bufQL0_.Get(); + pipe->InitBuffer(bufKeyL0_, L0_BUF_NUM * D_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0 * sizeof(K_T)); + keyL0_ = bufKeyL0_.Get(); + + pipe->InitBuffer(bufL0C_, L0_BUF_NUM * M_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0 * sizeof(float)); + cL0_ = bufL0C_.Get(); +} + +template +__aicore__ inline void +LIMatmul::InitMm1GlobalTensor(const GlobalTensor &blkTableGm, const GlobalTensor &keyGm, + const GlobalTensor &queryGm, const GlobalTensor &mm1ResGm) +{ + blkTableGm_ = blkTableGm; + keyGm_ = keyGm; + queryGm_ = queryGm; + mm1ResGm_ = mm1ResGm; +} + +template +__aicore__ inline void LIMatmul::ComputeMm1(const LICommon::RunInfo &runInfo) +{ + uint64_t s2GmBaseOffset = runInfo.s2Idx * constInfo_.s2BaseSize; + uint64_t s1gProcessSize = runInfo.actMBaseSize; + uint64_t s2ProcessSize = runInfo.actualSingleProcessSInnerSize; + for (uint64_t s2GmOffset = 0; s2GmOffset < s2ProcessSize; s2GmOffset += S2_BASIC_BLOCK) { + WaitFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM); + uint64_t s2L1RealSize = + s2GmOffset + S2_BASIC_BLOCK > s2ProcessSize ? s2ProcessSize - s2GmOffset : S2_BASIC_BLOCK; + if (PAGE_ATTENTION) { + KeyNd2NzForPA(s2L1RealSize, s2GmBaseOffset + s2GmOffset, runInfo); + }else { + KeyNd2Nz(s2L1RealSize, s2GmOffset, runInfo); + } + + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + for (uint64_t s1gGmOffset = 0; s1gGmOffset < s1gProcessSize; s1gGmOffset += M_BASIC_BLOCK) { + uint64_t s1gL1RealSize = + s1gGmOffset + M_BASIC_BLOCK > s1gProcessSize ? s1gProcessSize - s1gGmOffset : M_BASIC_BLOCK; + if (runInfo.isFirstS2InnerLoop && s2GmOffset == 0) { + queryL1Mte2BufIdx_++; + queryL1Mte1BufIdx_ = queryL1Mte2BufIdx_; + WaitFlag(QUERY_MTE1_MTE2_EVENT + queryL1Mte2BufIdx_ % QUERY_BUF_NUM); + QueryNd2Nz(s1gL1RealSize, s1gGmOffset, runInfo); + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + } else { + queryL1Mte1BufIdx_ = + queryL1Mte2BufIdx_ - (CeilDiv(s1gProcessSize, M_BASIC_BLOCK) - 1 - (s1gGmOffset > 0)); + } + for (uint64_t s2L1Offset = 0; s2L1Offset < s2L1RealSize; s2L1Offset += S2_BASIC_BLOCK_L0) { + uint64_t s2L0RealSize = + s2L1Offset + S2_BASIC_BLOCK_L0 > s2L1RealSize ? s2L1RealSize - s2L1Offset : S2_BASIC_BLOCK_L0; + for (uint64_t s1gL1Offset = 0; s1gL1Offset < s1gL1RealSize; s1gL1Offset += M_BASIC_BLOCK_L0) { + WaitFlag(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM); + uint64_t s1gL0RealSize = + s1gL1Offset + M_BASIC_BLOCK_L0 > s1gL1RealSize ? s1gL1RealSize - s1gL1Offset : M_BASIC_BLOCK_L0; + LoadQueryToL0a(s1gGmOffset, s1gL1Offset, s1gL1RealSize, s1gL0RealSize, runInfo); + LoadKeyToL0b(s2L1Offset, s2L1RealSize, s2L0RealSize, runInfo); + + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + + ComuteL0c(s1gL0RealSize, s2L0RealSize, runInfo); + + SetFlag(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM); + + Fixp(s1gGmOffset + s1gL1Offset, s2GmOffset + s2L1Offset, s1gL0RealSize, s2L0RealSize, runInfo); + l0BufIdx_++; + } + } + if (s2GmOffset + S2_BASIC_BLOCK >= s2ProcessSize && runInfo.isLastS2InnerLoop) { + SetFlag(QUERY_MTE1_MTE2_EVENT + queryL1Mte1BufIdx_ % QUERY_BUF_NUM); + } + } + + SetFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM); + keyL1BufIdx_++; + } +} + +template +__aicore__ inline void LIMatmul::KeyNd2Nz(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const LICommon::RunInfo &runInfo) +{ + uint64_t s2L1Offset = 0; + while (s2L1Offset < s2L1RealSize) { + uint64_t keyGmOffset = runInfo.tensorKeyOffset + (s2GmOffset + s2L1Offset) * constInfo_.headDim; + uint64_t s2Mte2Size = (s2L1RealSize <= S2_BASIC_BLOCK_L0 || s2L1Offset >= S2_BASIC_BLOCK_L0) ? + s2L1RealSize - s2L1Offset : + S2_BASIC_BLOCK_L0 - s2L1Offset; + + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2Mte2Size; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = s2L1Offset >= S2_BASIC_BLOCK_L0 ? + CeilAlign(s2L1RealSize - S2_BASIC_BLOCK_L0, (uint64_t)BLOCK_CUBE) : + (s2L1RealSize > S2_BASIC_BLOCK_L0 ? + S2_BASIC_BLOCK_L0 : + CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE)); + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET + + (s2L1Offset >= S2_BASIC_BLOCK_L0 ? + S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0 + (s2L1Offset - S2_BASIC_BLOCK_L0) * BLOCK_CUBE : + s2L1Offset * BLOCK_CUBE)], + keyGm_[keyGmOffset], nd2nzPara); + + s2L1Offset += s2Mte2Size; + } +} + +// blkNum, blkSize, N2, D +template +__aicore__ inline void LIMatmul::KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const LICommon::RunInfo &runInfo) +{ + uint64_t s2L1Offset = 0; + while (s2L1Offset < s2L1RealSize) { + uint64_t s2BlkId = (s2L1Offset + s2GmOffset) / constInfo_.kCacheBlockSize; + uint64_t s2BlkOffset = (s2L1Offset + s2GmOffset) % constInfo_.kCacheBlockSize; + uint64_t keyGmOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId) * + constInfo_.kCacheBlockSize * constInfo_.kHeadNum * constInfo_.headDim + + s2BlkOffset * constInfo_.headDim; + uint64_t s2Mte2Size = (s2L1RealSize <= S2_BASIC_BLOCK_L0 || s2L1Offset >= S2_BASIC_BLOCK_L0) ? + s2L1RealSize - s2L1Offset : + S2_BASIC_BLOCK_L0 - s2L1Offset; + s2Mte2Size = s2BlkOffset + s2Mte2Size >= constInfo_.kCacheBlockSize ? constInfo_.kCacheBlockSize - s2BlkOffset : + s2Mte2Size; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2Mte2Size; + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = s2L1Offset >= S2_BASIC_BLOCK_L0 ? + CeilAlign(s2L1RealSize - S2_BASIC_BLOCK_L0, (uint64_t)BLOCK_CUBE) : + (s2L1RealSize > S2_BASIC_BLOCK_L0 ? + S2_BASIC_BLOCK_L0 : + CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE)); + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET + + (s2L1Offset >= S2_BASIC_BLOCK_L0 ? + S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0 + (s2L1Offset - S2_BASIC_BLOCK_L0) * BLOCK_CUBE : + s2L1Offset * BLOCK_CUBE)], + keyGm_[keyGmOffset], nd2nzPara); + + s2L1Offset += s2Mte2Size; + } +} + +// batch, s1, n2, g, d +template +__aicore__ inline void LIMatmul::QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gGmOffset, + const LICommon::RunInfo &runInfo) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s1gL1RealSize; + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(s1gL1RealSize, (uint64_t)BLOCK_CUBE); + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(queryL1_[(queryL1Mte2BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET], + queryGm_[runInfo.tensorQueryOffset + s1gGmOffset * constInfo_.headDim], nd2nzPara); +} + +template +__aicore__ inline void LIMatmul::LoadQueryToL0a(uint64_t s1gGmOffset, uint64_t s1gL1Offset, uint64_t s1gL1RealSize, + uint64_t s1gL0RealSize, const LICommon::RunInfo &runInfo) +{ + LoadData3DParamsV2 loadData3DParams; + // SetFmatrixParams + loadData3DParams.l1H = CeilDiv(s1gL1RealSize, BLOCK_CUBE); // Hin=M1=8 + loadData3DParams.l1W = BLOCK_CUBE; // Win=M0 + loadData3DParams.channelSize = constInfo_.headDim; // Cin=K + + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; + + // SetLoadToA0Params + loadData3DParams.mExtension = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + loadData3DParams.kExtension = constInfo_.headDim; + loadData3DParams.mStartPt = s1gL1Offset; + loadData3DParams.kStartPt = 0; + loadData3DParams.strideW = 1; + loadData3DParams.strideH = 1; + loadData3DParams.filterW = 1; + loadData3DParams.filterSizeW = (1 >> 8) & 255; + loadData3DParams.filterH = 1; + loadData3DParams.filterSizeH = (1 >> 8) & 255; + loadData3DParams.dilationFilterW = 1; + loadData3DParams.dilationFilterH = 1; + loadData3DParams.enTranspose = 0; + loadData3DParams.fMatrixCtrl = 0; + + LoadData(queryL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], + queryL1_[(queryL1Mte1BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET], + loadData3DParams); +} + +template +__aicore__ inline void LIMatmul::LoadKeyToL0b(uint64_t s2L1Offset, uint64_t s2L1RealSize, uint64_t s2L0RealSize, + const LICommon::RunInfo &runInfo) +{ + uint64_t keyL1Offset = s2L1Offset >= S2_BASIC_BLOCK_L0 ? S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0 : 0; + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = CeilDiv(s2L0RealSize, BLOCK_CUBE) * CeilDiv(constInfo_.headDim, BLOCK_CUBE); + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = false; + LoadData(keyL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], + keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET + keyL1Offset], loadData2DParams); +} + +template +__aicore__ inline void LIMatmul::ComuteL0c(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, + const LICommon::RunInfo &runInfo) +{ + MmadParams mmadParams; + mmadParams.m = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + mmadParams.n = s2L0RealSize; + mmadParams.k = constInfo_.headDim; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = 0b11; + Mmad(cL0_[(l0BufIdx_ % L0_BUF_NUM) * L0C_BUFFER_OFFSET], queryL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], + keyL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } +} + +template +__aicore__ inline void LIMatmul::Fixp(uint64_t s1gGmOffset, uint64_t s2GmOffset, uint64_t s1gL0RealSize, + uint64_t s2L0RealSize, const LICommon::RunInfo &runInfo) +{ + AscendC::DataCopyCO12DstParams intriParams; + intriParams.mSize = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + intriParams.nSize = s2L0RealSize; + intriParams.dstStride = runInfo.actualSingleProcessSInnerSizeAlign; + intriParams.srcStride = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + // set mode according to dtype + intriParams.quantPre = QuantMode_t::NoQuant; + intriParams.nz2ndEn = true; + intriParams.unitFlag = 0b11; // 3 unitflag + intriParams.reluPre = 1; + AscendC::SetFixpipeNz2ndFlag(1, 1, 1); + AscendC::DataCopy(mm1ResGm_[(runInfo.loop % 2) * constInfo_.mBaseSize * constInfo_.s2BaseSize + + s1gGmOffset * intriParams.dstStride + s2GmOffset], + cL0_[(l0BufIdx_ % L0_BUF_NUM) * L0C_BUFFER_OFFSET], intriParams); +} + +template +__aicore__ inline void LIMatmul::AllocEventID() +{ + SetMMLayoutTransform(true); + SetFlag(KEY_MTE1_MTE2_EVENT + 0); + SetFlag(KEY_MTE1_MTE2_EVENT + 1); + SetFlag(KEY_MTE1_MTE2_EVENT + 2); + + SetFlag(QUERY_MTE1_MTE2_EVENT + 0); + SetFlag(QUERY_MTE1_MTE2_EVENT + 1); + + SetFlag(M_MTE1_EVENT + 0); + SetFlag(M_MTE1_EVENT + 1); +} + +template +__aicore__ inline void LIMatmul::FreeEventID() +{ + SetMMLayoutTransform(false); + WaitFlag(KEY_MTE1_MTE2_EVENT + 0); + WaitFlag(KEY_MTE1_MTE2_EVENT + 1); + WaitFlag(KEY_MTE1_MTE2_EVENT + 2); + + WaitFlag(QUERY_MTE1_MTE2_EVENT + 0); + WaitFlag(QUERY_MTE1_MTE2_EVENT + 1); + + WaitFlag(M_MTE1_EVENT + 0); + WaitFlag(M_MTE1_EVENT + 1); +} +} // namespace LIKernel +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_service_vector.h b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_service_vector.h new file mode 100644 index 000000000..1ed25b4cb --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_service_vector.h @@ -0,0 +1,559 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_service_vector.h + * \brief + */ +#ifndef LIGHTNING_INDEXER_SERVICE_VECTOR_H +#define LIGHTNING_INDEXER_SERVICE_VECTOR_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "lightning_indexer_common.h" +#include "lightning_indexer_vector.h" + +namespace LIKernel { +using namespace LICommon; +using namespace LIServiceVec; +constexpr uint32_t BASE_TOPK = 2048; +constexpr uint32_t LD_PARAM_NUM = 16; + +template +class LIVector { +public: + using K_T = typename LIT::keyType; + static constexpr LI_LAYOUT LAYOUT_T = LIT::layout; + + using MM1_OUT_T = float; + + __aicore__ inline LIVector(){}; + __aicore__ inline void ProcessVec(const LICommon::RunInfo &info); + __aicore__ inline void ProcessLD(); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitParams(const struct LICommon::ConstInfo &constInfo, + const LITilingData *__restrict tilingData); + __aicore__ inline void InitVec1GlobalTensor(GlobalTensor mm1ResGm, GlobalTensor vec1ResGm, + GlobalTensor vec1ParamGm, GlobalTensor weightsGm, + GlobalTensor indiceOutGm); + __aicore__ inline void CleanInvalidOutput(int64_t invalidS1offset); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void InitLDBuffers(TPipe *pipe); + +protected: + GlobalTensor mm1ResGm; + GlobalTensor vec1ResGm; + GlobalTensor vec1ParamGm; + GlobalTensor weightsGm; + GlobalTensor indiceOutGm; + +private: + // queue + TQue inQueue_; + TQue outQueue_; + + // tmp buff for vector + TBuf sortOutBuf_; + TBuf indexBuf_; + TBuf reduceOutBuf_; + TBuf brcBuf_; + TBuf paramBuf_; + + // tmp buff for LD + TBuf<> ldToBeMrgBuf_; + TBuf<> ldTmpBuf_; + TBuf<> ldOutValueBuf_; + TBuf<> ldOutIdxBuf_; + + LocalTensor globalTopkIndice_; + LocalTensor globalTopkUb_; + LocalTensor SortedBasicBlock_; + + int32_t blockId_ = -1; + // para for vector + int32_t groupInner_ = 0; + int32_t globalTopkNum_ = 0; + int64_t blockS2StartIdx_ = 0; + int32_t gSize_ = 0; + int32_t kHeadNum_ = 0; + int32_t s1BaseSize_ = 0; + int32_t s2BaseSize_ = 0; + + // para for LD + uint32_t mrgListNum_ = 4; + uint32_t paramNum_ = 16; + + constexpr static uint32_t REDUCE_BANK_CONFLICT_OFFSETS = 256; + constexpr static uint32_t REDUCE_BANK_CONFLICT_NUM = REDUCE_BANK_CONFLICT_OFFSETS / sizeof(float); + + struct LICommon::ConstInfo constInfo_; +}; + +template +__aicore__ inline void LIVector::InitBuffers(TPipe *pipe) +{ + uint32_t outNeedBufSize = (BASE_TOPK * 2) * 2 * sizeof(float); + uint32_t reduceCacheSize = REDUCE_BANK_CONFLICT_OFFSETS + groupInner_ * s2BaseSize_ * sizeof(float); + outNeedBufSize = reduceCacheSize > outNeedBufSize ? reduceCacheSize : outNeedBufSize; + + pipe->InitBuffer(inQueue_, 2, + groupInner_ * s2BaseSize_ * sizeof(float) + s2BaseSize_ * sizeof(float)); // 69KB mm_out_ub + pipe->InitBuffer(outQueue_, 1, outNeedBufSize); // 32KB extract + pipe->InitBuffer(sortOutBuf_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK * 2 * sizeof(float)); // 64KB + pipe->InitBuffer(indexBuf_, s2BaseSize_ * sizeof(int32_t)); // 2KB + pipe->InitBuffer(reduceOutBuf_, s2BaseSize_ * 2 * sizeof(float)); // 4KB + pipe->InitBuffer(brcBuf_, groupInner_ * 8 * sizeof(float)); + pipe->InitBuffer(paramBuf_, LD_PARAM_NUM * sizeof(int64_t)); + + // + globalTopkIndice_ = indexBuf_.Get(); + globalTopkUb_ = sortOutBuf_.Get(); + SortedBasicBlock_ = globalTopkUb_[BASE_TOPK * 2 * 2]; + globalTopkNum_ = 0; + + ArithProgression(globalTopkIndice_, 0, 1, s2BaseSize_); + InitSortOutBuf(globalTopkUb_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK * 2); + LocalTensor tmpfBuff = outQueue_.AllocTensor(); + Duplicate(tmpfBuff.template ReinterpretCast(), -1, 2 * (s1BaseSize_ / 2) * paramNum_ * 2); + SetWaitFlag(HardEvent::V_MTE3); + int64_t wsInfoOffset = (blockId_ / 2) * s1BaseSize_ * 2 * paramNum_ + + (blockId_ % 2) * (s1BaseSize_ / 2) * 2 * paramNum_; + DataCopyPad(vec1ParamGm[wsInfoOffset], tmpfBuff.template ReinterpretCast(), + {1, static_cast((s1BaseSize_ / 2) * 2 * paramNum_ * sizeof(int64_t)), 0, 0}); + SetWaitFlag(HardEvent::MTE3_V); + outQueue_.FreeTensor(tmpfBuff); +} + +template +__aicore__ inline void LIVector::InitLDBuffers(TPipe *pipe) +{ + pipe->Reset(); + pipe->InitBuffer(ldToBeMrgBuf_, 2 * BASE_TOPK * mrgListNum_ * sizeof(float)); // 2：value + index + pipe->InitBuffer(ldTmpBuf_, 2 * BASE_TOPK * mrgListNum_ * sizeof(float)); // 2：value + index + pipe->InitBuffer(ldOutValueBuf_, BASE_TOPK * sizeof(float)); + pipe->InitBuffer(ldOutIdxBuf_, BASE_TOPK * sizeof(int32_t)); +} + +template +__aicore__ inline void LIVector::InitParams(const struct LICommon::ConstInfo &constInfo, + const LITilingData *__restrict tilingData) +{ + this->constInfo_ = constInfo; + blockS2StartIdx_ = 0; + gSize_ = constInfo.gSize; + // define N2 para + kHeadNum_ = constInfo.kHeadNum; + // define MMBase para + s1BaseSize_ = constInfo.s1BaseSize; + s2BaseSize_ = constInfo.s2BaseSize; + + groupInner_ = 16; + blockId_ = GetBlockIdx(); +} + +template +__aicore__ inline void +LIVector::InitVec1GlobalTensor(GlobalTensor mm1ResGm, GlobalTensor vec1ResGm, + GlobalTensor vec1ParamGm, GlobalTensor weightsGm, + GlobalTensor indiceOutGm) +{ + this->mm1ResGm = mm1ResGm; + this->vec1ResGm = vec1ResGm; + this->vec1ParamGm = vec1ParamGm; + this->weightsGm = weightsGm; + this->indiceOutGm = indiceOutGm; +} + +template +__aicore__ inline void LIVector::AllocEventID() +{ +} + +template +__aicore__ inline void LIVector::FreeEventID() +{ +} + +template +__aicore__ inline void LIVector::CleanInvalidOutput(int64_t invalidS1offset) +{ + // init -1 and copy to output + LocalTensor valueULocal = outQueue_.AllocTensor(); + LocalTensor idxULocal1 = valueULocal.template ReinterpretCast(); + Duplicate(idxULocal1, constInfo_.INVALID_IDX, constInfo_.sparseCount); + outQueue_.EnQue(valueULocal); + valueULocal = outQueue_.DeQue(); + LIServiceVec::CopyOut(indiceOutGm[invalidS1offset], idxULocal1, constInfo_.sparseCount); + outQueue_.FreeTensor(valueULocal); +} + +template +__aicore__ inline void LIVector::ProcessVec(const LICommon::RunInfo &info) +{ + int32_t cuBaseS1Idx = info.gS1Idx * s1BaseSize_; + int32_t cuBaseS2Idx = info.s2Idx * s2BaseSize_; + + int64_t mmGmOffset = (info.loop % 2) * ((s1BaseSize_ * gSize_) * s2BaseSize_); + int64_t weightGmOffset = info.tensorWeightsOffset + cuBaseS1Idx * kHeadNum_ * gSize_; + + PipeBarrier(); + int32_t cuS1BeginIdxPerAiv = cuBaseS1Idx; + int32_t cuS1ProcNum = + cuS1BeginIdxPerAiv + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + int32_t cuS1ProcNumPerAiv = blockId_ % 2 == 0 ? CeilDiv(cuS1ProcNum, 2) : (cuS1ProcNum / 2); + cuS1BeginIdxPerAiv += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2); + + weightGmOffset += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2) * kHeadNum_ * gSize_; + mmGmOffset += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2) * gSize_ * info.actualSingleProcessSInnerSizeAlign; + + // cut G + int32_t outerG = CeilDiv(gSize_, groupInner_); + + if (info.loop != 0 && info.s2Idx == 0) { + // globalTopkUb_ value,index=-inf,-1 + InitSortOutBuf(globalTopkUb_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK * 2); + blockS2StartIdx_ = 0; + } else if (info.loop == 0) { + blockS2StartIdx_ = info.s2Idx; + } + int32_t cuRealAcSeq = info.actS2Size; + if (constInfo_.attenMaskFlag) { + cuRealAcSeq = info.actS2Size - (info.actS1Size - cuS1BeginIdxPerAiv); + } + LocalTensor reduceOutBuff = reduceOutBuf_.Get(); + LocalTensor brcBuf = brcBuf_.Get(); + uint32_t ldS1Offset = (blockId_ % 2 == 0) ? s1BaseSize_ / 2 - cuS1ProcNumPerAiv : 0; + for (int innerS1Idx = 0; innerS1Idx < cuS1ProcNumPerAiv; innerS1Idx++) { + if (constInfo_.attenMaskFlag) { + cuRealAcSeq += 1; + } + int32_t cuS2Len = cuBaseS2Idx + s2BaseSize_ >= cuRealAcSeq ? cuRealAcSeq - cuBaseS2Idx : s2BaseSize_; + int32_t cuS1Idx = cuS1BeginIdxPerAiv + innerS1Idx; + if (cuRealAcSeq > 0 && cuS2Len > 0) { + int32_t cuS2LenVecAlign = CeilDiv(cuS2Len, s2BaseSize_) * s2BaseSize_; + int32_t mmUbStride = (cuS2LenVecAlign - info.actualSingleProcessSInnerSizeAlign) / B32_BLOCK_ALIGN_NUM; + LocalTensor reduceOutInner = reduceOutBuff[s2BaseSize_]; + PipeBarrier(); + LocalTensor reduceCacheBuf = outQueue_.AllocTensor(); + for (int outerGidx = 0; outerGidx < outerG; outerGidx++) { + int32_t procGnum = outerGidx != outerG - 1 ? groupInner_ : gSize_ - outerGidx * groupInner_; + LocalTensor mmInUb = inQueue_.AllocTensor(); + LocalTensor weightsInUb = mmInUb[procGnum * s2BaseSize_]; + LocalTensor weightsInTUb = weightsInUb.template ReinterpretCast(); + if constexpr (!IsSameType::value) { + weightsInTUb = weightsInTUb[groupInner_]; + } + LIServiceVec::CopyIn(mmInUb, weightsInTUb, mm1ResGm, weightsGm, + mmGmOffset + innerS1Idx * gSize_ * info.actualSingleProcessSInnerSizeAlign + + outerGidx * groupInner_ * info.actualSingleProcessSInnerSizeAlign, + weightGmOffset + innerS1Idx * gSize_ + outerGidx * groupInner_, procGnum, + info.actualSingleProcessSInnerSizeAlign, mmUbStride); + + inQueue_.EnQue(mmInUb); + mmInUb = inQueue_.DeQue(); + weightsInUb = mmInUb[procGnum * s2BaseSize_]; + LIServiceVec::DoScale(reduceCacheBuf[REDUCE_BANK_CONFLICT_NUM], mmInUb, weightsInUb, weightsInTUb, + brcBuf, procGnum, s2BaseSize_, outerGidx); + // confused reduceOp in DoScale + // neednot use LIServiceVec::doReduce(mmInUb, reduceOutInner, procGnum, (s2BaseSize_+8)); + inQueue_.FreeTensor(mmInUb); + } + + int32_t gRedCnt = groupInner_ > gSize_ ? gSize_ : groupInner_; + bool isS2End = cuBaseS2Idx + s2BaseSize_ >= cuRealAcSeq; + LIServiceVec::DoReduce(reduceCacheBuf[REDUCE_BANK_CONFLICT_NUM], reduceOutInner, gRedCnt, s2BaseSize_); + outQueue_.FreeTensor(reduceCacheBuf); + + LocalTensor sortScoreUb = reduceOutBuff; + LocalTensor sortIndiceUb = reduceOutBuff[cuS2LenVecAlign]; + PipeBarrier(); + Duplicate(sortScoreUb.template ReinterpretCast(), LIServiceVec::NEG_INF, cuS2LenVecAlign); + PipeBarrier(); + Adds(sortScoreUb, reduceOutInner, 0.0f, cuS2Len); + PipeBarrier(); + LocalTensor sortIndiceUbInt = sortIndiceUb.template ReinterpretCast(); + if (cuS2LenVecAlign != cuS2Len) { + Duplicate(sortIndiceUbInt, -1, cuS2LenVecAlign); + } + PipeBarrier(); + Adds(sortIndiceUbInt, globalTopkIndice_, static_cast(cuBaseS2Idx), cuS2Len); + PipeBarrier(); + + LocalTensor tmpSortBuf = outQueue_.AllocTensor(); + if (info.actS1Size > 4) { + LIServiceVec::SortAll(reduceOutBuff, tmpSortBuf, + cuS2LenVecAlign); // cuS2LenVecAlign <= s2BaseSize_, fill -inf + PipeBarrier(); + LIServiceVec::MergeSort(globalTopkUb_[innerS1Idx * BASE_TOPK * 2], BASE_TOPK, reduceOutBuff, + cuS2LenVecAlign, tmpSortBuf); + } else { + int64_t globalTopkUbCacheIdx = (info.s2Idx - blockS2StartIdx_) % 4; + Sort( + SortedBasicBlock_[innerS1Idx * BASE_TOPK * 2 + globalTopkUbCacheIdx * s2BaseSize_ * 2], + reduceOutBuff, sortIndiceUbInt.template ReinterpretCast(), tmpSortBuf, + cuS2LenVecAlign / 32); + if (globalTopkUbCacheIdx == 3 || isS2End || info.isAllLoopEnd) { + LocalTensor tt = SortedBasicBlock_[innerS1Idx * BASE_TOPK * 2]; + if (info.s2Idx - blockS2StartIdx_ < 4) { + MrgBasicBlock(globalTopkUb_[innerS1Idx * BASE_TOPK * 2], tt, + static_cast(globalTopkUbCacheIdx + 1), s2BaseSize_); + } else { + if (globalTopkUbCacheIdx > 0) { + MrgBasicBlock(tmpSortBuf, tt, static_cast(globalTopkUbCacheIdx + 1), s2BaseSize_); + PipeBarrier(); + DataCopy(SortedBasicBlock_[innerS1Idx * BASE_TOPK * 2], tmpSortBuf, + (globalTopkUbCacheIdx + 1) * s2BaseSize_ * 2); + } + PipeBarrier(); + SparseTopK(globalTopkUb_[innerS1Idx * BASE_TOPK * 2], + SortedBasicBlock_[innerS1Idx * BASE_TOPK * 2], tmpSortBuf, BASE_TOPK, + s2BaseSize_ * (globalTopkUbCacheIdx + 1)); + } + } + } + + PipeBarrier(); + outQueue_.FreeTensor(tmpSortBuf); + + bool needCopyOutGm = blockS2StartIdx_ == 0 && isS2End; + bool needCopyWsGm = info.isAllLoopEnd || isS2End; + + if (needCopyOutGm) { + LocalTensor valueULocal = outQueue_.AllocTensor(); + LocalTensor idxULocal = valueULocal.template ReinterpretCast()[BASE_TOPK]; + ExtractIndex(idxULocal, globalTopkUb_[innerS1Idx * BASE_TOPK * 2].template ReinterpretCast(), + BASE_TOPK); + PipeBarrier(); + InitSortOutBuf(globalTopkUb_[innerS1Idx * BASE_TOPK * 2], BASE_TOPK * 2); + outQueue_.EnQue(valueULocal); + valueULocal = outQueue_.DeQue(); + LocalTensor idxULocal1 = valueULocal.template ReinterpretCast()[BASE_TOPK]; + LIServiceVec::CopyOut(indiceOutGm[info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount], + idxULocal1, constInfo_.sparseCount); + outQueue_.FreeTensor(valueULocal); + } else if (needCopyWsGm) { + // vec1Res Gm = [aic, s1BaseSize_, 2, 2, topkOut_] float32 + // vec1Param Gm = [aic, s1BaseSize_, 2, 16] int64 + // 16 = [needFd, s2AcSeq, s2Start, s2End, isS2End, bn2idx, s1Idx, S1ProcNum, ......] + + int64_t wsOffset = (blockId_ / 2) * s1BaseSize_ * 2 * 2 * BASE_TOPK + + (blockId_ % 2) * (s1BaseSize_ / 2) * 2 * 2 * BASE_TOPK + + (ldS1Offset + innerS1Idx) * 2 * 2 * BASE_TOPK; + int64_t wsInfoOffset = (blockId_ / 2) * s1BaseSize_ * 2 * paramNum_ + + (blockId_ % 2) * (s1BaseSize_ / 2) * 2 * paramNum_ + + (ldS1Offset + innerS1Idx) * 2 * paramNum_; + + LocalTensor tmpiBuff = paramBuf_.Get(); + SetWaitFlag(HardEvent::MTE3_S); + tmpiBuff.SetValue(0, static_cast(1)); + tmpiBuff.SetValue(1, static_cast(cuRealAcSeq)); + tmpiBuff.SetValue(2, static_cast(blockS2StartIdx_)); + tmpiBuff.SetValue(3, static_cast(cuBaseS2Idx + cuS2Len)); + tmpiBuff.SetValue(4, static_cast(isS2End)); + tmpiBuff.SetValue(5, static_cast(info.bN2Idx)); + tmpiBuff.SetValue(6, static_cast(cuS1Idx)); + tmpiBuff.SetValue(7, static_cast(cuS1ProcNum)); + tmpiBuff.SetValue(8, static_cast(info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount)); + bool isTailReduce = blockS2StartIdx_ == 0; + if (isTailReduce) { + wsInfoOffset += paramNum_; + wsOffset += 2 * BASE_TOPK; + } + SetWaitFlag(HardEvent::S_MTE3); + LIServiceVec::CopyOut(vec1ParamGm[wsInfoOffset], tmpiBuff, 16); + SetWaitFlag(HardEvent::V_MTE3); + LIServiceVec::CopyOut(vec1ResGm[wsOffset], globalTopkUb_[innerS1Idx * BASE_TOPK * 2], 2 * BASE_TOPK); + SetWaitFlag(HardEvent::MTE3_V); + } + } else if (cuRealAcSeq <= 0) { + CleanInvalidOutput(info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount); + } + } + + if (LAYOUT_T == LI_LAYOUT::BSND) { + bool isS1LoopEnd = (cuBaseS1Idx + s1BaseSize_) >= info.actS1Size; + int32_t invalidS1Num = constInfo_.qSeqSize - info.actS1Size; + if (invalidS1Num > 0 && isS1LoopEnd && blockS2StartIdx_ == 0) { + int32_t s1NumPerAiv = blockId_ % 2 == 0 ? CeilDiv(invalidS1Num, 2) : (invalidS1Num / 2); + int32_t s1OffsetPerAiv = info.actS1Size + (blockId_ % 2) * CeilDiv(invalidS1Num, 2); + for (int innerS1Idx = 0; innerS1Idx < s1NumPerAiv; innerS1Idx++) { + CleanInvalidOutput(info.indiceOutOffset + (s1OffsetPerAiv + innerS1Idx) * constInfo_.sparseCount); + } + } + + int32_t invalidS1Num2 = info.actS1Size - info.actS2Size; + if (invalidS1Num2 > 0 && isS1LoopEnd && blockS2StartIdx_ == 0 && constInfo_.attenMaskFlag) { + int32_t s1NumPerAiv = blockId_ % 2 == 0 ? CeilDiv(invalidS1Num2, 2) : (invalidS1Num2 / 2); + int32_t s1OffsetPerAiv = (blockId_ % 2) * CeilDiv(invalidS1Num2, 2); + for (int innerS1Idx = 0; innerS1Idx < s1NumPerAiv; innerS1Idx++) { + CleanInvalidOutput((info.bN2Idx * constInfo_.qSeqSize + s1OffsetPerAiv + innerS1Idx) * + constInfo_.sparseCount); + } + } + } + + if (info.isLastS2InnerLoop) { + blockS2StartIdx_ = 0; + } +} + +template +__aicore__ inline void LIVector::ProcessLD() +{ + int32_t curCubeId = blockId_ / 2; + int32_t tmpCubeId = curCubeId; + + int64_t s2ActSeq; + int64_t s2Start; + int64_t s2End; + int64_t isS2End; + int64_t bn2Idx; + int64_t s1Idx; + uint32_t acc_list_num = 0; + int64_t bIdx = 0; + int64_t needFd; + int64_t wsOffset; + int64_t wsInfoOffset = 0; + int64_t nextneedFd; + int64_t valueOffset = 0; + int64_t outOffset = 0; + + LocalTensor curValueIdxUb = ldToBeMrgBuf_.Get(); + LocalTensor tmpUb = ldTmpBuf_.Get(); + + uint32_t s1LdStartIdx = 0; + uint32_t s1ProcNum = 0; + uint64_t paramGmCoreOffset = tmpCubeId * s1BaseSize_ * 2 * paramNum_; + for (uint32_t innerS1Idx = 0; innerS1Idx < s1BaseSize_; innerS1Idx++) { + needFd = vec1ParamGm.GetValue(paramGmCoreOffset + innerS1Idx * 2 * paramNum_ + paramNum_); + if (needFd == 1) { + s1LdStartIdx = (s1ProcNum == 0) ? innerS1Idx : s1LdStartIdx; + s1ProcNum++; + } + } + + if (s1ProcNum == 0) { + return; + } + + uint32_t s1VecNum = CeilDiv(s1ProcNum, 2); + if (blockId_ % 2 == 1) { + s1LdStartIdx = s1LdStartIdx + s1VecNum; + s1VecNum = s1ProcNum - s1VecNum; + } + for (uint32_t innerS1Idx = s1LdStartIdx; innerS1Idx < s1LdStartIdx + s1VecNum; innerS1Idx++) { + tmpCubeId = curCubeId; + acc_list_num = 0; + valueOffset = 0; + + wsOffset = tmpCubeId * s1BaseSize_ * 2 * 2 * BASE_TOPK + + innerS1Idx * 2 * 2 * BASE_TOPK + 2 * BASE_TOPK; + SetWaitFlag(HardEvent::V_MTE2); + SetWaitFlag(HardEvent::S_MTE2); + DataCopyPad(curValueIdxUb, vec1ResGm[wsOffset], + {1, static_cast(2 * BASE_TOPK * sizeof(int32_t)), 0, 0}, {true, 0, 0, 0}); + acc_list_num++; + valueOffset += 2 * BASE_TOPK; + + tmpCubeId++; + wsInfoOffset = tmpCubeId * s1BaseSize_ * 2 * paramNum_ + innerS1Idx * 2 * paramNum_; + needFd = vec1ParamGm.GetValue(wsInfoOffset); + isS2End = vec1ParamGm.GetValue(wsInfoOffset + 4); + s1Idx = vec1ParamGm.GetValue(wsInfoOffset + 6); + outOffset = vec1ParamGm.GetValue(wsInfoOffset + 8); + + while (needFd == 1) { + wsOffset = tmpCubeId * s1BaseSize_ * 2 * 2 * BASE_TOPK + + innerS1Idx * 2 * 2 * BASE_TOPK; + SetWaitFlag(HardEvent::V_MTE2); + SetWaitFlag(HardEvent::S_MTE2); + DataCopyPad(curValueIdxUb[valueOffset], vec1ResGm[wsOffset], + {1, static_cast(2 * BASE_TOPK * sizeof(int32_t)), 0, 0}, {true, 0, 0, 0}); + valueOffset += 2 * BASE_TOPK; + acc_list_num++; + + if (acc_list_num == mrgListNum_) { + AscendC::MrgSort4Info params; + params.elementLengths[0] = BASE_TOPK; + params.elementLengths[1] = BASE_TOPK; + params.elementLengths[2] = BASE_TOPK; + params.elementLengths[3] = BASE_TOPK; + params.ifExhaustedSuspension = true; + params.validBit = 0b1111; + params.repeatTimes = 1; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = curValueIdxUb[0]; + srcList.src2 = curValueIdxUb[2 * BASE_TOPK]; + srcList.src3 = curValueIdxUb[4 * BASE_TOPK]; + srcList.src4 = curValueIdxUb[6 * BASE_TOPK]; + SetWaitFlag(HardEvent::MTE2_V); + MrgSort(tmpUb, srcList, params); + PipeBarrier(); + DataCopy(curValueIdxUb, tmpUb, 2 * BASE_TOPK); + PipeBarrier(); + acc_list_num = 1; + valueOffset = 2 * BASE_TOPK; + } + + if (isS2End == 1) { + break; + } + + tmpCubeId++; + wsInfoOffset = tmpCubeId * s1BaseSize_ * 2 * paramNum_ + innerS1Idx * 2 * paramNum_; + needFd = vec1ParamGm.GetValue(wsInfoOffset); + isS2End = vec1ParamGm.GetValue(wsInfoOffset + 4); + } + + if (acc_list_num != 1) { + AscendC::MrgSort4Info params; + params.elementLengths[0] = BASE_TOPK; + params.elementLengths[1] = BASE_TOPK; + params.elementLengths[2] = BASE_TOPK; + params.elementLengths[3] = BASE_TOPK; + params.ifExhaustedSuspension = true; + if (acc_list_num == 2) { + params.validBit = 0b0011; + } else if (acc_list_num == 3) { + params.validBit = 0b0111; + } + params.repeatTimes = 1; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = curValueIdxUb[0]; + srcList.src2 = curValueIdxUb[2 * BASE_TOPK]; + srcList.src3 = curValueIdxUb[4 * BASE_TOPK]; + srcList.src4 = curValueIdxUb[6 * BASE_TOPK]; + SetWaitFlag(HardEvent::MTE2_V); + MrgSort(tmpUb, srcList, params); + PipeBarrier(); + DataCopy(curValueIdxUb, tmpUb, 2 * BASE_TOPK); + PipeBarrier(); + } + + LocalTensor outValueUb = ldOutValueBuf_.Get(); + LocalTensor outIdxUb = ldOutIdxBuf_.Get(); + + Extract(outValueUb, outIdxUb, curValueIdxUb, (BASE_TOPK / 32)); + LocalTensor idxULocal1 = outIdxUb.template ReinterpretCast(); + SetWaitFlag(HardEvent::V_MTE3); + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(indiceOutGm[outOffset], idxULocal1, + {1, static_cast(constInfo_.sparseCount * sizeof(int32_t)), 0, 0}); + SetWaitFlag(HardEvent::MTE3_V); + } +} +} // namespace LIKernel +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_template_tiling_key.h b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_template_tiling_key.h new file mode 100644 index 000000000..b6f0a484d --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_template_tiling_key.h @@ -0,0 +1,66 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_template_tiling_key.h + * \brief + */ + +#ifndef TEMPLATE_TILING_KEY_LI_H_ +#define TEMPLATE_TILING_KEY_LI_H_ + +#include "ascendc/host_api/tiling/template_argument.h" + +#define LI_TPL_FP16 1 +#define LI_TPL_INT32 3 +#define LI_TPL_BF16 27 + +#define LI_LAYOUT_BSND 0 +#define LI_LAYOUT_TND 1 +#define LI_LAYOUT_PA_BSND 2 + +#define ASCENDC_TPL_4_BW 4 + +ASCENDC_TPL_ARGS_DECL(LightningIndexerVllm, + ASCENDC_TPL_DTYPE_DECL(DT_Q, LI_TPL_FP16, LI_TPL_BF16), + ASCENDC_TPL_DTYPE_DECL(DT_K, LI_TPL_FP16, LI_TPL_BF16), + ASCENDC_TPL_DTYPE_DECL(DT_OUT, LI_TPL_INT32), ASCENDC_TPL_BOOL_DECL(PAGE_ATTENTION, 0, 1), + ASCENDC_TPL_UINT_DECL(LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, LI_LAYOUT_BSND, + LI_LAYOUT_TND), + ASCENDC_TPL_UINT_DECL(K_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, + LI_LAYOUT_PA_BSND, LI_LAYOUT_BSND, LI_LAYOUT_TND), ); + +ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, LI_TPL_FP16), ASCENDC_TPL_DTYPE_SEL(DT_K, LI_TPL_FP16), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, LI_TPL_INT32), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_BSND, LI_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_PA_BSND), ), + + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, LI_TPL_BF16), ASCENDC_TPL_DTYPE_SEL(DT_K, LI_TPL_BF16), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, LI_TPL_INT32), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_BSND, LI_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_PA_BSND), ), + + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, LI_TPL_FP16), ASCENDC_TPL_DTYPE_SEL(DT_K, LI_TPL_FP16), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, LI_TPL_INT32), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_BSND, LI_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, + LI_LAYOUT_BSND, LI_LAYOUT_TND), ), + + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, LI_TPL_BF16), ASCENDC_TPL_DTYPE_SEL(DT_K, LI_TPL_BF16), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, LI_TPL_INT32), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_BSND, LI_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, LI_LAYOUT_BSND, LI_LAYOUT_TND), ), ); + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_vector.h b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_vector.h new file mode 100644 index 000000000..962901275 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_vector.h @@ -0,0 +1,335 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_vector.h + * \brief + */ +#ifndef LIGHTNING_INDEXER_VECTOR_H +#define LIGHTNING_INDEXER_VECTOR_H + +#include "lightning_indexer_vector.h" +#include "kernel_operator.h" + +namespace LIServiceVec { +using namespace AscendC; + +constexpr int32_t NEG_INF = 0xFF800000; +constexpr int32_t INVALID_INDEX = -1; +constexpr uint8_t VEC_REPEAT_MAX = 255; +constexpr uint8_t B32_VEC_ELM_NUM = 64; +constexpr uint8_t B32_BLOCK_ALIGN_NUM = 8; +constexpr uint8_t B32_VEC_REPEAT_STRIDE = 8; +constexpr uint64_t VEC_REPEAT_BYTES = 256; +constexpr int32_t CONST_TWO = 2; +constexpr int64_t VALUE_AND_INDEX_NUM = 2; +constexpr int64_t BLOCK_BYTES = 32; +constexpr int64_t MRG_QUE_0 = 0; +constexpr int64_t MRG_QUE_1 = 1; +constexpr int64_t MRG_QUE_2 = 2; +constexpr int64_t MRG_QUE_3 = 3; +constexpr int64_t MRG_BLOCK_2 = 2; +constexpr int64_t MRG_BLOCK_3 = 3; +constexpr int64_t MRG_BLOCK_4 = 4; + +template +__aicore__ inline void CopyIn(LocalTensor &mmOutUb, LocalTensor &weightsUb, GlobalTensor &mMoutGm, + GlobalTensor &weightScaleGm, int64_t MMout_gmoffset, int64_t weights_gmoffset, + int64_t groupInner, int64_t s2Inner, int64_t mmUbStride) +{ + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams dataCopymMoutParams; + dataCopymMoutParams.blockCount = groupInner; + dataCopymMoutParams.blockLen = s2Inner * sizeof(float); + dataCopymMoutParams.srcStride = 0; + dataCopymMoutParams.dstStride = mmUbStride; + dataCopymMoutParams.rsv = 0; + AscendC::DataCopyPad(mmOutUb, mMoutGm[MMout_gmoffset], dataCopymMoutParams, padParams); + + AscendC::DataCopyPadExtParams padTParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams dataCopyweightParams; + dataCopyweightParams.blockCount = 1; + dataCopyweightParams.blockLen = groupInner * sizeof(T); + dataCopyweightParams.srcStride = 0; + dataCopyweightParams.dstStride = 0; + dataCopyweightParams.rsv = 0; + AscendC::DataCopyPad(weightsUb, weightScaleGm[weights_gmoffset], dataCopyweightParams, padTParams); +} + + +template +__aicore__ inline void CopyOut(const GlobalTensor &dstGm, const LocalTensor &srcUb, int64_t copyCount) +{ + AscendC::DataCopyParams dataCopyOutyParams; + dataCopyOutyParams.blockCount = 1; + dataCopyOutyParams.blockLen = copyCount * sizeof(T); + dataCopyOutyParams.srcStride = 0; + dataCopyOutyParams.dstStride = 0; + AscendC::DataCopyPad(dstGm, srcUb, dataCopyOutyParams); +} + + +template +__aicore__ inline void DoScale(const LocalTensor &reduceCacheBuf, LocalTensor &mmOutUb, + LocalTensor &weightsUb, LocalTensor &weightsTUb, LocalTensor &tmpBuff, + int64_t groupInner, int64_t s2Inner, int32_t outerGidx) +{ + // cast bfloat16_t to float + if constexpr (!IsSameType::value) { + AscendC::Cast(weightsUb, weightsTUb, RoundMode::CAST_NONE, groupInner); + AscendC::PipeBarrier(); + } + + // weight broadcast: [groupInner, 1] -> [groupInner, 8] + AscendC::Brcb(tmpBuff, weightsUb, LICommon::CeilDiv(groupInner, static_cast(B32_BLOCK_ALIGN_NUM)), + {1, B32_VEC_REPEAT_STRIDE}); + AscendC::PipeBarrier(); + + // do scale: [groupInner, 8] * [groupInner, s2Inner] + uint64_t countPerRepeat = VEC_REPEAT_BYTES / sizeof(float); + uint64_t repeatTimes = s2Inner / countPerRepeat; + for (int32_t i = 0; i < groupInner; i++) { + if (outerGidx == 0) { + AscendC::Mul(reduceCacheBuf[i * s2Inner], mmOutUb[i * s2Inner], tmpBuff[i * B32_BLOCK_ALIGN_NUM], + countPerRepeat, repeatTimes, {1, 1, 0, B32_VEC_REPEAT_STRIDE, B32_VEC_REPEAT_STRIDE, 0}); + } else { + AscendC::Mul(mmOutUb[i * s2Inner], mmOutUb[i * s2Inner], tmpBuff[i * B32_BLOCK_ALIGN_NUM], countPerRepeat, + repeatTimes, {1, 1, 0, B32_VEC_REPEAT_STRIDE, B32_VEC_REPEAT_STRIDE, 0}); + } + } + + if (outerGidx != 0) { + AscendC::PipeBarrier(); + AscendC::Add(reduceCacheBuf, mmOutUb, reduceCacheBuf, groupInner * s2Inner); + } + AscendC::PipeBarrier(); +} + + +__aicore__ inline uint64_t FindNearestPower2(uint64_t value) +{ + if (value <= CONST_TWO) { + return value; + } else { + const uint64_t pow = 63 - clz(value); + return (1 << pow); + } +} + + +__aicore__ inline void DoReduce(const LocalTensor &srcTensor, LocalTensor &dstTensor, int32_t rNum, + int32_t aNum) +{ + if (rNum == 1) { + AscendC::Adds(dstTensor, srcTensor, 0, aNum); + AscendC::PipeBarrier(); + return; + } + + uint32_t dichotomizeAddPow = FindNearestPower2(rNum); + uint32_t dichotomizeAddDiffSize = rNum - dichotomizeAddPow; + if (dichotomizeAddDiffSize != 0) { + AscendC::Add(srcTensor, srcTensor, srcTensor[dichotomizeAddPow * aNum], dichotomizeAddDiffSize * aNum); + AscendC::PipeBarrier(); + } + int32_t nowRows = dichotomizeAddPow; + while (nowRows > CONST_TWO) { + nowRows = nowRows / CONST_TWO; + AscendC::Add(srcTensor, srcTensor, srcTensor[nowRows * aNum], nowRows * aNum); + AscendC::PipeBarrier(); + } + AscendC::Add(dstTensor, srcTensor, srcTensor[aNum], aNum); + AscendC::PipeBarrier(); +} + +__aicore__ inline void InitSortOutBuf(const LocalTensor &src, int64_t eleNum) +{ + uint64_t mask1[2] = {0x5555555555555555, 0}; + uint64_t mask0[2] = {0xaaaaaaaaaaaaaaaa, 0}; + int64_t repeatNum = eleNum / B32_VEC_ELM_NUM; + int64_t forLoop = repeatNum / VEC_REPEAT_MAX; + int64_t forRemain = repeatNum % VEC_REPEAT_MAX; + for (int i = 0; i < forLoop; i++) { + AscendC::Duplicate(src.template ReinterpretCast(), NEG_INF, mask1, VEC_REPEAT_MAX, 1, + B32_VEC_REPEAT_STRIDE); + AscendC::Duplicate(src.template ReinterpretCast(), INVALID_INDEX, mask0, VEC_REPEAT_MAX, 1, + B32_VEC_REPEAT_STRIDE); + } + if (forRemain > 0) { + AscendC::Duplicate(src.template ReinterpretCast()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM], NEG_INF, + mask1, forRemain, 1, B32_VEC_REPEAT_STRIDE); + AscendC::Duplicate(src.template ReinterpretCast()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM], + INVALID_INDEX, mask0, forRemain, 1, B32_VEC_REPEAT_STRIDE); + } + AscendC::PipeBarrier(); +} + +__aicore__ inline void SortAll(LocalTensor &src, LocalTensor &tmp, int64_t logitsNum) +{ + int64_t sort32Repeats = logitsNum / BLOCK_BYTES; + AscendC::Sort32(tmp, src, src[logitsNum].ReinterpretCast(), sort32Repeats); + AscendC::PipeBarrier(); + + int64_t mrgGroups = sort32Repeats; + int64_t mrgElements = BLOCK_BYTES; + int64_t i = 0; + AscendC::LocalTensor srcTensor; + AscendC::LocalTensor dstTensor; + while (true) { + if (i % CONST_TWO == 0) { + srcTensor = tmp; + dstTensor = src; + } else { + srcTensor = src; + dstTensor = tmp; + } + AscendC::MrgSort4Info params; + params.elementLengths[0] = mrgElements; + params.elementLengths[MRG_QUE_1] = mrgElements; + params.elementLengths[MRG_QUE_2] = mrgElements; + params.elementLengths[MRG_QUE_3] = mrgElements; + params.ifExhaustedSuspension = false; + params.validBit = 0b1111; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = srcTensor[0]; + srcList.src2 = srcTensor[MRG_QUE_1 * VALUE_AND_INDEX_NUM * mrgElements]; + srcList.src3 = srcTensor[MRG_QUE_2 * VALUE_AND_INDEX_NUM * mrgElements]; + srcList.src4 = srcTensor[MRG_QUE_3 * VALUE_AND_INDEX_NUM * mrgElements]; + if (mrgGroups <= MRG_BLOCK_4) { + params.repeatTimes = 1; + if (mrgGroups == 1) { + break; + } else if (mrgGroups == MRG_BLOCK_2) { + params.validBit = 0b0011; + } else if (mrgGroups == MRG_BLOCK_3) { + params.validBit = 0b0111; + } else if (mrgGroups == MRG_BLOCK_4) { + params.validBit = 0b1111; + } + AscendC::MrgSort(dstTensor, srcList, params); + i += 1; + break; + } else { + params.repeatTimes = mrgGroups / MRG_BLOCK_4; + AscendC::MrgSort(dstTensor, srcList, params); + i += 1; + mrgElements = mrgElements * MRG_BLOCK_4; + mrgGroups = mrgGroups / MRG_BLOCK_4; + } + AscendC::PipeBarrier(); + } + if (i % CONST_TWO == 0) { + AscendC::DataCopy(src, tmp, logitsNum * VALUE_AND_INDEX_NUM); + AscendC::PipeBarrier(); + } +} + +__aicore__ inline void SortAll(LocalTensor &dst, LocalTensor &srcValue, LocalTensor &srcIndex, + LocalTensor &tmpTensor, int64_t logitsNum) +{ + int64_t sort32Repeats = logitsNum / BLOCK_BYTES; + AscendC::Sort(dst, srcValue, srcIndex, tmpTensor, sort32Repeats); + AscendC::PipeBarrier(); +} + +__aicore__ inline void MergeSort(const LocalTensor &mrgDst, int32_t mrgDstNum, LocalTensor &mrgSrc, + int32_t mrgSrcNum, LocalTensor &tmpTensor) +{ + AscendC::MrgSort4Info params; + params.elementLengths[0] = mrgDstNum; + params.elementLengths[1] = mrgSrcNum; + params.ifExhaustedSuspension = false; + params.validBit = 0b0011; + params.repeatTimes = 1; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = mrgDst; + srcList.src2 = mrgSrc; + + AscendC::MrgSort(tmpTensor, srcList, params); + AscendC::PipeBarrier(); + AscendC::DataCopy(mrgDst, tmpTensor, mrgDstNum * VALUE_AND_INDEX_NUM); + AscendC::PipeBarrier(); +} + +__aicore__ inline void MrgBasicBlock(const LocalTensor &dst, const LocalTensor &src, int64_t blockNum, + int64_t basicBlockSize) +{ + AscendC::MrgSort4Info params; + params.elementLengths[MRG_QUE_0] = basicBlockSize; + params.elementLengths[MRG_QUE_1] = basicBlockSize; + params.elementLengths[MRG_QUE_2] = basicBlockSize; + params.elementLengths[MRG_QUE_3] = basicBlockSize; + params.ifExhaustedSuspension = false; + if (blockNum == MRG_BLOCK_2) { + params.validBit = 0b0011; + } else if (blockNum == MRG_BLOCK_3) { + params.validBit = 0b0111; + } else if (blockNum == MRG_BLOCK_4) { + params.validBit = 0b1111; + } else { + AscendC::DataCopy(dst, src, basicBlockSize * VALUE_AND_INDEX_NUM); + return; + } + AscendC::MrgSortSrcList srcList; + srcList.src1 = src[0]; + srcList.src2 = src[basicBlockSize * VALUE_AND_INDEX_NUM * MRG_QUE_1]; + srcList.src3 = src[basicBlockSize * VALUE_AND_INDEX_NUM * MRG_QUE_2]; + srcList.src4 = src[basicBlockSize * VALUE_AND_INDEX_NUM * MRG_QUE_3]; + AscendC::MrgSort(dst, srcList, params); +} + +template +__aicore__ inline void SparseTopK(const LocalTensor &dst, const LocalTensor &needsMerging, + const LocalTensor &tmp, int64_t topk, int64_t mergSize) +{ + if (!needMrg) { + AscendC::DataCopy(dst, needsMerging, mergSize * VALUE_AND_INDEX_NUM); + return; + } + AscendC::MrgSort4Info params; + params.elementLengths[0] = topk; + params.elementLengths[1] = mergSize; + params.ifExhaustedSuspension = (topk == mergSize); + params.validBit = 0b0011; + AscendC::MrgSortSrcList srcList; + srcList.src1 = dst; + srcList.src2 = needsMerging; + AscendC::MrgSort(tmp, srcList, params); + AscendC::DataCopy(dst, tmp, topk * VALUE_AND_INDEX_NUM); +} + + +__aicore__ inline void ExtractIndex(const LocalTensor &idxULocal, const LocalTensor &sortLocal, + int64_t extractNum) +{ + AscendC::GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = Ceil(extractNum * sizeof(float) * VALUE_AND_INDEX_NUM, VEC_REPEAT_BYTES); + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = B32_VEC_REPEAT_STRIDE; + gatherMaskParams.src1RepeatStride = 0; + uint64_t rsvdCnt = 0; + uint8_t src1Pattern = 2; + AscendC::GatherMask(idxULocal, sortLocal, src1Pattern, false, static_cast(0), gatherMaskParams, rsvdCnt); + AscendC::PipeBarrier(); +} + + +template +__aicore__ inline void SetWaitFlag(HardEvent evt) +{ + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(evt)); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); +} + +} // namespace LIServiceVec +#endif // LIGHTNING_INDEXER_VECTOR_H \ No newline at end of file diff --git a/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_vllm.cpp b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_vllm.cpp new file mode 100644 index 000000000..8f2981a25 --- /dev/null +++ b/csrc/ascend/attention/lightning_indexer_vllm/op_kernel/lightning_indexer_vllm.cpp @@ -0,0 +1,58 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer.cpp + * \brief + */ + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" +#include "lightning_indexer_template_tiling_key.h" +#include "lightning_indexer_kernel.h" + +using namespace LIKernel; + +#define INVOKE_LI_NO_KFC_OP_IMPL(templateClass, ...) \ + do { \ + templateClass> op; \ + LI_COPY_TILING_DATA(LITilingData, tiling); \ + op.Init(query, key, weights, actualSeqLengthsQ, actualSeqLengths, blocktable, sparseIndices, user, \ + tiling_data, &tPipe); \ + op.Process(); \ + } while (0) + +#define LI_COPY_TILING_DATA(tilingDataStruct, tiling) \ + GET_TILING_DATA_WITH_STRUCT(tilingDataStruct, tiling_data_in, tiling); \ + const tilingDataStruct *__restrict tiling_data = &tiling_data_in; + + +template +__global__ __aicore__ void lightning_indexer_vllm(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths, + __gm__ uint8_t *blocktable, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *workspace, __gm__ uint8_t *tiling) +{ +#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) || (__CCE_AICORE__ == 200) + +#else + TPipe tPipe; + __gm__ uint8_t *user = GetUserWorkspace(workspace); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + + if constexpr (DT_Q == LI_TPL_FP16 && DT_K == LI_TPL_FP16 && DT_OUT == LI_TPL_INT32) { + INVOKE_LI_NO_KFC_OP_IMPL(LIPreload, half, half, int32_t, PAGE_ATTENTION, + LI_LAYOUT(LAYOUT_T), LI_LAYOUT(K_LAYOUT_T)); + } else { + INVOKE_LI_NO_KFC_OP_IMPL(LIPreload, bfloat16_t, bfloat16_t, int32_t, PAGE_ATTENTION, + LI_LAYOUT(LAYOUT_T), LI_LAYOUT(K_LAYOUT_T)); + } +#endif +} diff --git a/csrc/ascend/attention/ngram_spec_decode/CMakeLists.txt b/csrc/ascend/attention/ngram_spec_decode/CMakeLists.txt new file mode 100644 index 000000000..549b5c14c --- /dev/null +++ b/csrc/ascend/attention/ngram_spec_decode/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() diff --git a/csrc/ascend/attention/ngram_spec_decode/ngram_spec_decode_torch_adpt.h b/csrc/ascend/attention/ngram_spec_decode/ngram_spec_decode_torch_adpt.h new file mode 100644 index 000000000..f42756f50 --- /dev/null +++ b/csrc/ascend/attention/ngram_spec_decode/ngram_spec_decode_torch_adpt.h @@ -0,0 +1,82 @@ +/* + * Licensed under the BSD 3-Clause License (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + */ +#ifndef NGRAM_SPEC_DECODE_TORCH_ADPT_H +#define NGRAM_SPEC_DECODE_TORCH_ADPT_H + +#include +#include + +namespace vllm_fl { + +// N-gram spec decode op +// inputs： +// token_ids: [batch_size, max_seq_len], int32, +// num_tokens_no_spec: [batch_size], int32 +// sampled_token_ids: [batch_size, max_new_tokens], int32 +// discard_request_mask: [batch_size], int32 +// vocab_size, min_n, max_n, k +// outputs: +// token_ids (in-place change), next_token_ids, draft_token_ids, num_valid_draft_tokens +inline std::tuple npu_ngram_spec_decode( + at::Tensor &token_ids, + const at::Tensor &num_tokens_no_spec, + const at::Tensor &sampled_token_ids, + const at::Tensor &discard_request_mask, + int64_t vocab_size, + int64_t min_n, + int64_t max_n, + int64_t k) +{ + int64_t batch_size = token_ids.size(0); + auto device = token_ids.device(); + + at::Tensor discard_mask_int = discard_request_mask.dtype() == at::kBool + ? discard_request_mask.to(at::kInt) + : discard_request_mask; + + // Allocate outputs with a trailing over-write cushion. The kernel's + // CopyOut path issues DataCopyPad GM writes whose burst length can + // be smaller than the NPU's 32-byte MTE alignment; under that + // alignment the underlying MTE3 burst can write past the apparent + // tensor end on the last row. Tightly-sized allocations (the original + // ``at::empty({batch_size}, ...)``) leave no room for that + // alignment-driven over-write, surfacing as a multi-core MTE OOB on + // device (CI signature: fixp_error0 = 0x30266b9 across cores). + // + // We therefore allocate ``batch_size + OVER_WRITE_MARGIN`` rows / + // ``(batch_size + OVER_WRITE_MARGIN) * k`` elements and ``narrow`` + // back to the user-visible shape. The narrowed view shares storage + // with the larger allocation, so any kernel-side alignment + // over-write lands inside owned memory rather than off the end. + constexpr int64_t OVER_WRITE_MARGIN = 8; // 32 bytes / sizeof(int32) = 8 ints + + at::Tensor next_token_ids_storage = at::empty( + {batch_size + OVER_WRITE_MARGIN}, + at::dtype(at::kInt).device(device)); + at::Tensor next_token_ids = next_token_ids_storage.narrow(0, 0, batch_size); + + at::Tensor draft_token_ids_storage = at::empty( + {batch_size + OVER_WRITE_MARGIN, k}, + at::dtype(at::kInt).device(device)); + at::Tensor draft_token_ids = draft_token_ids_storage.narrow(0, 0, batch_size); + + at::Tensor num_valid_draft_tokens_storage = at::empty( + {batch_size + OVER_WRITE_MARGIN}, + at::dtype(at::kInt).device(device)); + at::Tensor num_valid_draft_tokens = + num_valid_draft_tokens_storage.narrow(0, 0, batch_size); + + EXEC_NPU_CMD(aclnnNgramSpecDecode, + token_ids, num_tokens_no_spec, sampled_token_ids, discard_mask_int, + vocab_size, min_n, max_n, k, + next_token_ids, draft_token_ids, num_valid_draft_tokens); + + return std::make_tuple(token_ids, next_token_ids, draft_token_ids, num_valid_draft_tokens); +} + +} // namespace vllm_fl + +#endif // NGRAM_SPEC_DECODE_TORCH_ADPT_H diff --git a/csrc/ascend/attention/ngram_spec_decode/op_host/CMakeLists.txt b/csrc/ascend/attention/ngram_spec_decode/op_host/CMakeLists.txt new file mode 100644 index 000000000..b6c7851d5 --- /dev/null +++ b/csrc/ascend/attention/ngram_spec_decode/op_host/CMakeLists.txt @@ -0,0 +1,22 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnInner PRIVATE + ngram_spec_decode_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME NgramSpecDecode + OPTIONS + --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE ngram_spec_decode ACLNNTYPE aclnn_inner) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/attention/ngram_spec_decode/op_host/ngram_spec_decode_def.cpp b/csrc/ascend/attention/ngram_spec_decode/op_host/ngram_spec_decode_def.cpp new file mode 100644 index 000000000..b48b4de7f --- /dev/null +++ b/csrc/ascend/attention/ngram_spec_decode/op_host/ngram_spec_decode_def.cpp @@ -0,0 +1,72 @@ +#include "register/op_def_registry.h" + +namespace ops { +class NgramSpecDecode : public OpDef { +public: + explicit NgramSpecDecode(const char *name) : OpDef(name) + { + this->Input("tokenIds") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + + this->Input("numTokensNoSpec") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + + this->Input("sampledTokenIds") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + + this->Input("discardRequestMask") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + + this->Attr("vocab_size").Int(); + this->Attr("min_n").Int(); + this->Attr("max_n").Int(); + this->Attr("k").Int(); + + this->Output("nextTokenIds") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + + this->Output("draftTokenIds") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + + this->Output("numValidDraftTokens") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("aclnnSupport.value", "support_aclnn") + .ExtendCfgInfo("jitCompile.flag", "static_true") + .ExtendCfgInfo("multiKernelSupportDynamicGraph.value", "multi_kernel"); + + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + } +}; + +OP_ADD(NgramSpecDecode); +} // namespace ops diff --git a/csrc/ascend/attention/ngram_spec_decode/op_host/ngram_spec_decode_tiling.cpp b/csrc/ascend/attention/ngram_spec_decode/op_host/ngram_spec_decode_tiling.cpp new file mode 100644 index 000000000..1ca239f1c --- /dev/null +++ b/csrc/ascend/attention/ngram_spec_decode/op_host/ngram_spec_decode_tiling.cpp @@ -0,0 +1,118 @@ +#include +#include +#include "log/ops_log.h" +#include "graph/utils/type_utils.h" +#include "register/op_def_registry.h" +#include "../op_kernel/ngram_spec_decode.h" +#include "tiling/platform/platform_ascendc.h" +#include "platform/platform_infos_def.h" + +using namespace ge; +namespace { +constexpr uint32_t INPUT_TOKEN_IDS_INDEX = 0; +constexpr uint32_t INPUT_NUM_TOKENS_INDEX = 1; +constexpr uint32_t INPUT_SAMPLED_INDEX = 2; +constexpr uint32_t INPUT_DISCARD_INDEX = 3; + +constexpr uint32_t ATTR_VOCAB_SIZE_INDEX = 0; +constexpr uint32_t ATTR_MIN_N_INDEX = 1; +constexpr uint32_t ATTR_MAX_N_INDEX = 2; +constexpr uint32_t ATTR_K_INDEX = 3; + +constexpr int64_t ELEM_SIZE = 4; // int32 +} // namespace + +namespace optiling { + +static ge::graphStatus NgramSpecDecodeTilingFunc(gert::TilingContext *context) +{ + const char *nodeName = context->GetNodeName(); + NgramSpecDecodeTilingData *tilingData = context->GetTilingData(); + OPS_CHECK(tilingData == nullptr, + OPS_LOG_E(nodeName, "tilingData is nullptr."), return ge::GRAPH_FAILED); + + auto attrs = context->GetAttrs(); + OPS_CHECK(attrs == nullptr, + OPS_LOG_E(nodeName, "attrs is nullptr."), return ge::GRAPH_FAILED); + + auto vocabSizePtr = attrs->GetAttrPointer(static_cast(ATTR_VOCAB_SIZE_INDEX)); + auto minNPtr = attrs->GetAttrPointer(static_cast(ATTR_MIN_N_INDEX)); + auto maxNPtr = attrs->GetAttrPointer(static_cast(ATTR_MAX_N_INDEX)); + auto kPtr = attrs->GetAttrPointer(static_cast(ATTR_K_INDEX)); + + OPS_CHECK(vocabSizePtr == nullptr, OPS_LOG_E(nodeName, "vocabSizePtr is null."), return ge::GRAPH_FAILED); + OPS_CHECK(minNPtr == nullptr, OPS_LOG_E(nodeName, "minNPtr is null."), return ge::GRAPH_FAILED); + OPS_CHECK(maxNPtr == nullptr, OPS_LOG_E(nodeName, "maxNPtr is null."), return ge::GRAPH_FAILED); + OPS_CHECK(kPtr == nullptr, OPS_LOG_E(nodeName, "kPtr is null."), return ge::GRAPH_FAILED); + + int64_t vocab_size = *vocabSizePtr; + int64_t min_n = *minNPtr; + int64_t max_n = *maxNPtr; + int64_t k = *kPtr; + + const gert::StorageShape *tokenIdsShape = context->GetInputShape(INPUT_TOKEN_IDS_INDEX); + const gert::StorageShape *sampledShape = context->GetInputShape(INPUT_SAMPLED_INDEX); + OPS_CHECK(tokenIdsShape == nullptr, OPS_LOG_E(nodeName, "tokenIdsShape is null."), return ge::GRAPH_FAILED); + OPS_CHECK(sampledShape == nullptr, OPS_LOG_E(nodeName, "sampledShape is null."), return ge::GRAPH_FAILED); + + int64_t batch_size = tokenIdsShape->GetStorageShape().GetDim(0); + int64_t max_seq_len = tokenIdsShape->GetStorageShape().GetDim(1); + int64_t max_new_tokens = sampledShape->GetStorageShape().GetDim(1); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint64_t ubSize = 0UL; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); + int64_t ub_size_limit = static_cast(ubSize); + + int64_t align_elems = 32 / ELEM_SIZE; + int64_t max_seq_len_align = ((max_seq_len + align_elems - 1) / align_elems) * align_elems; + int64_t max_new_tokens_align = ((max_new_tokens + align_elems - 1) / align_elems) * align_elems; + int64_t k_align = ((k + align_elems - 1) / align_elems) * align_elems; + + int64_t ub_per_row = (max_seq_len_align + max_new_tokens_align + k_align) * ELEM_SIZE; + int64_t ub_overhead = 4 * 32 + static_cast(max_n) * ELEM_SIZE + + ((max_seq_len_align + 7) / 8); // maskBuf + int64_t ub_available = ub_size_limit - ub_overhead; + int64_t max_block_rows = (ub_available > 0) ? (ub_available / ub_per_row) : 1; + max_block_rows = std::max(max_block_rows, static_cast(1)); + + int64_t block_dim = std::min(batch_size, static_cast(aivNum)); + int64_t rows_per_core = (block_dim > 0) ? (batch_size / block_dim) : 0; + int64_t former_num = (block_dim > 0) ? (block_dim - 1) : 0; + int64_t tail_rows = batch_size - former_num * rows_per_core; + int64_t block_rows = std::min(rows_per_core, max_block_rows); + + tilingData->ngramInfo.batchSize = static_cast(batch_size); + tilingData->ngramInfo.maxSeqLen = static_cast(max_seq_len); + tilingData->ngramInfo.maxNewTokens = static_cast(max_new_tokens); + tilingData->ngramInfo.vocabSize = static_cast(vocab_size); + tilingData->ngramInfo.minN = static_cast(min_n); + tilingData->ngramInfo.maxN = static_cast(max_n); + tilingData->ngramInfo.k = static_cast(k); + tilingData->ngramInfo.formerNum = static_cast(former_num); + tilingData->ngramInfo.rowsPerCore = static_cast(rows_per_core); + tilingData->ngramInfo.tailRows = static_cast(tail_rows); + tilingData->ngramInfo.blockRows = static_cast(block_rows); + + context->SetBlockDim(static_cast(block_dim)); + + OPS_LOG_D(nodeName, "batchSize=%lu, maxSeqLen=%lu, maxNewTokens=%lu, k=%lu, blockDim=%lu, blockRows=%lu", + batch_size, max_seq_len, max_new_tokens, k, block_dim, block_rows); + + return ge::GRAPH_SUCCESS; +} + +struct NgramSpecDecodeCompileInfo {}; + +ge::graphStatus TilingParseForNgramSpecDecode(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(NgramSpecDecode) + .Tiling(NgramSpecDecodeTilingFunc) + .TilingParse(TilingParseForNgramSpecDecode); + +} // namespace optiling diff --git a/csrc/ascend/attention/ngram_spec_decode/op_host/ngram_spec_decode_tiling.h b/csrc/ascend/attention/ngram_spec_decode/op_host/ngram_spec_decode_tiling.h new file mode 100644 index 000000000..278473914 --- /dev/null +++ b/csrc/ascend/attention/ngram_spec_decode/op_host/ngram_spec_decode_tiling.h @@ -0,0 +1,26 @@ +#ifndef NGRAM_SPEC_DECODE_TILING_H +#define NGRAM_SPEC_DECODE_TILING_H + +#include "kernel_tiling/kernel_tiling.h" + +struct NgramSpecDecodeInfo { + uint32_t batchSize; + uint32_t maxSeqLen; + uint32_t maxNewTokens; + uint32_t vocabSize; + uint32_t minN; + uint32_t maxN; + uint32_t k; + uint32_t formerNum; + uint32_t rowsPerCore; + uint32_t tailRows; + uint32_t blockRows; +}; + +struct NgramSpecDecodeTilingData { + Mc2InitTiling mc2InitTiling; + Mc2CcTiling mc2CcTiling1; + NgramSpecDecodeInfo ngramInfo; +}; + +#endif // NGRAM_SPEC_DECODE_TILING_H diff --git a/csrc/ascend/attention/ngram_spec_decode/op_host/op_api/aclnn_ngram_spec_decode.cpp b/csrc/ascend/attention/ngram_spec_decode/op_host/op_api/aclnn_ngram_spec_decode.cpp new file mode 100644 index 000000000..3bc617fce --- /dev/null +++ b/csrc/ascend/attention/ngram_spec_decode/op_host/op_api/aclnn_ngram_spec_decode.cpp @@ -0,0 +1,73 @@ +#include +#include "graph/types.h" +#include "aclnn_ngram_spec_decode.h" + +enum NnopbaseHcclServerType { + NNOPBASE_HCCL_SERVER_TYPE_AICPU = 0, + NNOPBASE_HCCL_SERVER_TYPE_MTE, + NNOPBASE_HCCL_SERVER_TYPE_END +}; +extern "C" void __attribute__((weak)) NnopbaseSetHcclServerType(void *executor, NnopbaseHcclServerType sType); + +#ifdef __cplusplus +extern "C" { +#endif + +extern aclnnStatus aclnnInnerNgramSpecDecodeGetWorkspaceSize( + const aclTensor *tokenIds, + const aclTensor *numTokensNoSpec, + const aclTensor *sampledTokenIds, + const aclTensor *discardRequestMask, + int64_t vocabSize, + int64_t minN, + int64_t maxN, + int64_t k, + const aclTensor *nextTokenIds, + const aclTensor *draftTokenIds, + const aclTensor *numValidDraftTokens, + uint64_t *workspaceSize, + aclOpExecutor **executor); + +extern aclnnStatus aclnnInnerNgramSpecDecode( + void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream); + +aclnnStatus aclnnNgramSpecDecodeGetWorkspaceSize( + const aclTensor *tokenIds, + const aclTensor *numTokensNoSpec, + const aclTensor *sampledTokenIds, + const aclTensor *discardRequestMask, + int64_t vocabSize, + int64_t minN, + int64_t maxN, + int64_t k, + const aclTensor *nextTokenIds, + const aclTensor *draftTokenIds, + const aclTensor *numValidDraftTokens, + uint64_t *workspaceSize, + aclOpExecutor **executor) +{ + return aclnnInnerNgramSpecDecodeGetWorkspaceSize( + tokenIds, numTokensNoSpec, sampledTokenIds, discardRequestMask, + vocabSize, minN, maxN, k, + nextTokenIds, draftTokenIds, numValidDraftTokens, + workspaceSize, executor); +} + +aclnnStatus aclnnNgramSpecDecode( + void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream) +{ + if (NnopbaseSetHcclServerType) { + NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE); + } + return aclnnInnerNgramSpecDecode(workspace, workspaceSize, executor, stream); +} + +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/attention/ngram_spec_decode/op_host/op_api/aclnn_ngram_spec_decode.h b/csrc/ascend/attention/ngram_spec_decode/op_host/op_api/aclnn_ngram_spec_decode.h new file mode 100644 index 000000000..5e52d5fdd --- /dev/null +++ b/csrc/ascend/attention/ngram_spec_decode/op_host/op_api/aclnn_ngram_spec_decode.h @@ -0,0 +1,56 @@ +#ifndef ACLNN_NGRAM_SPEC_DECODE_H_ +#define ACLNN_NGRAM_SPEC_DECODE_H_ + +#include "aclnn/acl_meta.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/* function: aclnnNgramSpecDecodeGetWorkspaceSize + * tokenIds : required, [batch_size, max_seq_len], int32 + * numTokensNoSpec : required, [batch_size], int32 + * sampledTokenIds : required, [batch_size, max_new_tokens], int32 + * discardRequestMask : required, [batch_size], int32 + * vocabSize : required, int + * minN : required, int + * maxN : required, int + * k : required, int + * nextTokenIds : required, [batch_size], int32 + * draftTokenIds : required, [batch_size, k], int32 + * numValidDraftTokens : required, [batch_size], int32 + * workspaceSize : size of workspace(output). + * executor : executor context(output). + */ +__attribute__((visibility("default"))) aclnnStatus aclnnNgramSpecDecodeGetWorkspaceSize( + const aclTensor *tokenIds, + const aclTensor *numTokensNoSpec, + const aclTensor *sampledTokenIds, + const aclTensor *discardRequestMask, + int64_t vocabSize, + int64_t minN, + int64_t maxN, + int64_t k, + const aclTensor *nextTokenIds, + const aclTensor *draftTokenIds, + const aclTensor *numValidDraftTokens, + uint64_t *workspaceSize, + aclOpExecutor **executor); + +/* function: aclnnNgramSpecDecode + * workspace : workspace memory addr(input). + * workspaceSize : size of workspace(input). + * executor : executor context(input). + * stream : acl stream. + */ +__attribute__((visibility("default"))) aclnnStatus aclnnNgramSpecDecode( + void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif // ACLNN_NGRAM_SPEC_DECODE_H_ diff --git a/csrc/ascend/attention/ngram_spec_decode/op_kernel/ngram_spec_decode.cpp b/csrc/ascend/attention/ngram_spec_decode/op_kernel/ngram_spec_decode.cpp new file mode 100644 index 000000000..830cafa6b --- /dev/null +++ b/csrc/ascend/attention/ngram_spec_decode/op_kernel/ngram_spec_decode.cpp @@ -0,0 +1,655 @@ +// Licensed under the BSD 3-Clause License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "kernel_operator.h" +#include "ngram_spec_decode.h" + +constexpr int32_t ELEM_SIZE = sizeof(int32_t); // 4 bytes +// Safety UB buffer size：32768(128KB) +constexpr uint32_t SAFE_CHUNK = 32768u; + +class KernelNgramSpecDecode { +public: + __aicore__ inline KernelNgramSpecDecode() {} + + __aicore__ inline void Init( + GM_ADDR token_ids_gm, GM_ADDR num_tokens_gm, GM_ADDR sampled_gm, + GM_ADDR discard_gm, GM_ADDR next_tokens_gm, GM_ADDR draft_tokens_gm, + GM_ADDR num_valid_gm, GM_ADDR workspace, GM_ADDR tiling) + { + REGISTER_TILING_DEFAULT(NgramSpecDecodeTilingData); + GET_TILING_DATA_WITH_STRUCT(NgramSpecDecodeTilingData, tilingData, tiling); + + this->batch_size = static_cast(tilingData.ngramInfo.batchSize); + this->max_seq_len = static_cast(tilingData.ngramInfo.maxSeqLen); + this->max_new_tokens = static_cast(tilingData.ngramInfo.maxNewTokens); + this->vocab_size_val = static_cast(tilingData.ngramInfo.vocabSize); + this->min_n_val = static_cast(tilingData.ngramInfo.minN); + this->max_n_val = static_cast(tilingData.ngramInfo.maxN); + this->k_val = static_cast(tilingData.ngramInfo.k); + this->former_num = static_cast(tilingData.ngramInfo.formerNum); + this->rows_per_core = static_cast(tilingData.ngramInfo.rowsPerCore); + this->tail_rows = static_cast(tilingData.ngramInfo.tailRows); + this->block_rows = static_cast(tilingData.ngramInfo.blockRows); + + int32_t align_elems = 32 / ELEM_SIZE; // = 8 + this->max_seq_len_align = ((this->max_seq_len + align_elems - 1) / align_elems) * align_elems; + this->max_new_tokens_align = ((this->max_new_tokens + align_elems - 1) / align_elems) * align_elems; + this->k_align = ((this->k_val + align_elems - 1) / align_elems) * align_elems; + + this->is_large_row = (this->max_seq_len_align > static_cast(SAFE_CHUNK)); + + uint32_t blockIdx = AscendC::GetBlockIdx(); + if (blockIdx < static_cast(this->former_num)) { + this->my_rows = static_cast(this->rows_per_core); + this->row_offset = static_cast(this->rows_per_core) * blockIdx; + } else { + this->my_rows = static_cast(this->tail_rows); + this->row_offset = static_cast(this->rows_per_core) * static_cast(this->former_num); + } + + tokenGm.SetGlobalBuffer((__gm__ int32_t *)token_ids_gm, + static_cast(this->batch_size) * this->max_seq_len); + numTokensGm.SetGlobalBuffer((__gm__ int32_t *)num_tokens_gm, + static_cast(this->batch_size)); + sampledGm.SetGlobalBuffer((__gm__ int32_t *)sampled_gm, + static_cast(this->batch_size) * this->max_new_tokens); + discardGm.SetGlobalBuffer((__gm__ int32_t *)discard_gm, + static_cast(this->batch_size)); + nextTokensGm.SetGlobalBuffer((__gm__ int32_t *)next_tokens_gm, + static_cast(this->batch_size)); + draftTokensGm.SetGlobalBuffer((__gm__ int32_t *)draft_tokens_gm, + static_cast(this->batch_size) * this->k_val); + numValidGm.SetGlobalBuffer((__gm__ int32_t *)num_valid_gm, + static_cast(this->batch_size)); + + uint32_t br = static_cast(this->block_rows); + uint32_t br_align = ((br * ELEM_SIZE + 31) / 32) * 32 / ELEM_SIZE; + + if (!this->is_large_row) { + pipe.InitBuffer(tokenTileBuf, br * static_cast(this->max_seq_len_align) * ELEM_SIZE); + } else { + uint32_t chunk_ub = SAFE_CHUNK + static_cast(this->max_n_val); + uint32_t chunk_ub_align = ((chunk_ub + 7u) / 8u) * 8u; + pipe.InitBuffer(tokenTileBuf, chunk_ub_align * ELEM_SIZE); + } + + uint32_t mask_bytes = ((SAFE_CHUNK + 7u) / 8u); + pipe.InitBuffer(maskBuf, mask_bytes); + + pipe.InitBuffer(sampledTileBuf, br * static_cast(this->max_new_tokens_align) * ELEM_SIZE); + pipe.InitBuffer(numTokensBuf, br_align * ELEM_SIZE); + pipe.InitBuffer(discardTileBuf, br_align * ELEM_SIZE); + pipe.InitBuffer(nextTokenBuf, br_align * ELEM_SIZE); + pipe.InitBuffer(draftBuf, br * static_cast(this->k_align) * ELEM_SIZE); + pipe.InitBuffer(numValidBuf, br_align * ELEM_SIZE); + pipe.InitBuffer(suffixBuf, static_cast(this->max_n_val) * ELEM_SIZE); + } + + __aicore__ inline void Process() + { + uint32_t remaining = this->my_rows; + uint32_t cur_offset = 0; + while (remaining > 0) { + uint32_t cur_rows = (remaining > static_cast(this->block_rows)) + ? static_cast(this->block_rows) : remaining; + if (this->is_large_row) { + ProcessChunkedRows(this->row_offset + cur_offset, cur_rows); + } else { + CopyIn(this->row_offset + cur_offset, cur_rows); + Compute(cur_rows); + CopyOut(this->row_offset + cur_offset, cur_rows); + } + cur_offset += cur_rows; + remaining -= cur_rows; + } + } + +private: + + __aicore__ inline void ProcessChunkedRows(uint32_t start_row, uint32_t rows) + { + uint32_t msl = static_cast(this->max_seq_len); + uint32_t mnta = static_cast(this->max_new_tokens_align); + uint32_t ka = static_cast(this->k_align); + + auto sampledLocal = sampledTileBuf.Get(); + auto numTokensLocal = numTokensBuf.Get(); + auto discardLocal = discardTileBuf.Get(); + auto nextLocal = nextTokenBuf.Get(); + auto draftLocal = draftBuf.Get(); + auto numValidLocal = numValidBuf.Get(); + auto suffixLocal = suffixBuf.Get(); + auto tokenLocal = tokenTileBuf.Get(); + auto maskLocal = maskBuf.Get(); + + uint32_t metaBytes = rows * ELEM_SIZE; + AscendC::DataCopyExtParams metaParams{1, metaBytes, 0, metaBytes, 0}; + AscendC::DataCopyPadExtParams noPadT{false, 0, 0, 0}; + AscendC::DataCopyPad(numTokensLocal, numTokensGm[start_row], metaParams, noPadT); + AscendC::DataCopyPad(discardLocal, discardGm[start_row], metaParams, noPadT); + + uint32_t srcRowBytes2 = static_cast(this->max_new_tokens) * ELEM_SIZE; + uint32_t dstRowBytes2 = mnta * ELEM_SIZE; + AscendC::DataCopyExtParams sampledParams{1, srcRowBytes2, 0, dstRowBytes2, 0}; + AscendC::DataCopyPadExtParams sampledPad{ + false, 0, static_cast(mnta - this->max_new_tokens), 0}; + for (uint32_t r = 0; r < rows; ++r) { + AscendC::DataCopyPad(sampledLocal[static_cast(r) * mnta], + sampledGm[static_cast(start_row + r) * this->max_new_tokens], + sampledParams, sampledPad); + } + + for (uint32_t i = 0; i < rows; ++i) { + uint64_t gmRow = static_cast(start_row + i) * msl; + int32_t seq_len = numTokensLocal.GetValue(i); + int32_t discard = discardLocal.GetValue(i); + int32_t valid_count = 0; + + int32_t backup_pos = (seq_len > 0) ? (seq_len - 1) : 0; + + for (int32_t j = 0; j < this->max_new_tokens; ++j) { + int32_t val = sampledLocal.GetValue(i * mnta + j); + if (discard != 0) { + sampledLocal.SetValue(i * mnta + j, -1); + } else if (val != -1 && val < this->vocab_size_val) { + valid_count++; + } else { + sampledLocal.SetValue(i * mnta + j, -1); + } + } + + int32_t avail_space = this->max_seq_len - seq_len; + if (avail_space < 0) avail_space = 0; + if (valid_count > avail_space) valid_count = avail_space; + + LoadGmElements(gmRow + backup_pos, 1); + int32_t backup_token = tokenLocal.GetValue(0); + + if (valid_count > 0) { + nextLocal.SetValue(i, sampledLocal.GetValue(i * mnta + valid_count - 1)); + } else { + nextLocal.SetValue(i, backup_token); + } + + int32_t nt = seq_len + valid_count; + if (valid_count > 0) { + for (int32_t j = 0; j < valid_count; ++j) { + tokenLocal.SetValue(j, sampledLocal.GetValue(i * mnta + j)); + } + StoreGmElements(gmRow + seq_len, valid_count); + } + + int32_t best_match_pos = -1; + int32_t best_ngram_len = 0; + + if (valid_count > 0 && nt >= this->min_n_val) { + int32_t suffix_gm_start = nt - this->max_n_val; + if (suffix_gm_start < 0) suffix_gm_start = 0; + LoadGmElements(gmRow + suffix_gm_start, this->max_n_val); + for (int32_t s = 0; s < this->max_n_val; ++s) { + suffixLocal.SetValue(static_cast(s), tokenLocal.GetValue(static_cast(s))); + } + + for (int32_t ngram_len = this->min_n_val; ngram_len <= this->max_n_val; ++ngram_len) { + if (ngram_len > nt) break; + int32_t wc = nt - ngram_len; + if (wc <= 0) break; + + int32_t suffix_offset = this->max_n_val - ngram_len; + int32_t suffix0 = suffixLocal.GetValue(static_cast(suffix_offset)); + + for (int32_t chunk_start = 0; chunk_start < wc; chunk_start += SAFE_CHUNK) { + int32_t chunk_count = (chunk_start + SAFE_CHUNK <= wc) ? SAFE_CHUNK : (wc - chunk_start); + int32_t load_count = chunk_count + (ngram_len - 1); + if (chunk_start + load_count > nt) load_count = nt - chunk_start; + LoadGmElements(gmRow + chunk_start, load_count); + + uint32_t cmp_count = ((static_cast(chunk_count) + 63u) / 64u) * 64u; + uint32_t max_cmp = SAFE_CHUNK > 8192u ? 8192u : SAFE_CHUNK; + if (cmp_count > max_cmp) cmp_count = max_cmp; + if (cmp_count > static_cast(load_count)) { + cmp_count = ((static_cast(load_count) + 63u) / 64u) * 64u; + } + + for (uint32_t cmp_off = 0; cmp_off < static_cast(chunk_count); cmp_off += cmp_count) { + uint32_t rem = static_cast(chunk_count) - cmp_off; + uint32_t elements = (rem >= cmp_count) ? cmp_count : rem; + uint32_t aligned = ((elements + 63u) / 64u) * 64u; + + AscendC::CompareScalar( + maskLocal, tokenLocal[cmp_off], + suffix0, AscendC::CMPMODE::EQ, aligned); + + for (uint32_t p = 0; p < elements; ++p) { + uint8_t bv = maskLocal.GetValue(p >> 3); + if (bv & (1u << (p & 7u))) { + bool all_match = true; + for (int32_t s = 1; s < ngram_len; ++s) { + int32_t sv = suffixLocal.GetValue(static_cast(suffix_offset + s)); + if (cmp_off + p + s < static_cast(load_count)) { + int32_t tv = tokenLocal.GetValue(cmp_off + p + static_cast(s)); + if (tv != sv) { all_match = false; break; } + } else { + all_match = false; break; + } + } + if (all_match) { + best_match_pos = chunk_start + static_cast(cmp_off + p); + best_ngram_len = ngram_len; + break; + } + } + } + if (best_match_pos >= 0) break; + } + if (best_match_pos >= 0) break; + } + if (best_match_pos >= 0) break; + } + } + + if (best_match_pos >= 0) { + int32_t draft_start = best_match_pos + best_ngram_len; + int32_t tokens_available = nt - draft_start; + int32_t draft_load = (tokens_available < this->k_val) ? tokens_available : this->k_val; + if (draft_load > 0) { + LoadGmElements(gmRow + draft_start, draft_load); + for (int32_t j = 0; j < this->k_val; ++j) { + if (j < draft_load) { + draftLocal.SetValue(i * ka + j, tokenLocal.GetValue(static_cast(j))); + } else { + draftLocal.SetValue(i * ka + j, -1); + } + } + } else { + for (int32_t j = 0; j < this->k_val; ++j) { + draftLocal.SetValue(i * ka + j, -1); + } + } + } else { + for (int32_t j = 0; j < this->k_val; ++j) { + draftLocal.SetValue(i * ka + j, -1); + } + } + + int32_t valid_draft_count = 0; + for (int32_t j = 0; j < this->k_val; ++j) { + if (draftLocal.GetValue(i * ka + j) != -1) { + valid_draft_count++; + } else { + break; + } + } + numValidLocal.SetValue(i, valid_draft_count); + } + + uint32_t metaBytes32 = static_cast(rows) * ELEM_SIZE; + AscendC::DataCopyExtParams nextParams{1, metaBytes32, 0, 0, 0}; + AscendC::DataCopyPad(nextTokensGm[start_row], nextLocal, nextParams); + + uint32_t kBytes = static_cast(this->k_val) * ELEM_SIZE; + for (uint32_t r = 0; r < rows; ++r) { + AscendC::DataCopyExtParams draftRowParams{1, kBytes, 0, 0, 0}; + AscendC::DataCopyPad( + draftTokensGm[static_cast(start_row + r) * this->k_val], + draftLocal[static_cast(r) * this->k_align], draftRowParams); + } + + AscendC::DataCopyPad(numValidGm[start_row], numValidLocal, nextParams); + } + + __aicore__ inline void LoadGmElements(uint64_t gm_offset, int32_t count) + { + if (count <= 0) return; + auto tokenLocal = tokenTileBuf.Get(); + uint32_t c = static_cast(count); + uint32_t aligned = ((c + 7u) / 8u) * 8u; + uint8_t pad = static_cast(aligned - c); + AscendC::DataCopyExtParams p{1, c * ELEM_SIZE, 0, aligned * ELEM_SIZE, 0}; + AscendC::DataCopyPadExtParams pp{false, 0, pad, 0}; + AscendC::DataCopyPad(tokenLocal[0], tokenGm[gm_offset], p, pp); + } + + __aicore__ inline void StoreGmElements(uint64_t gm_offset, int32_t count) + { + if (count <= 0) return; + auto tokenLocal = tokenTileBuf.Get(); + constexpr uint32_t STORE_MAX = 16383u; + uint32_t c = static_cast(count); + for (uint32_t off = 0; off < c; off += STORE_MAX) { + uint32_t chunk = (off + STORE_MAX <= c) ? STORE_MAX : (c - off); + AscendC::DataCopyExtParams p{1, chunk * ELEM_SIZE, 0, 0, 0}; + AscendC::DataCopyPad(tokenGm[gm_offset + off], tokenLocal[off], p); + } + } + + + __aicore__ inline void CopyIn(uint32_t start_row, uint32_t rows) + { + uint32_t msa = static_cast(this->max_seq_len_align); + uint32_t mnta = static_cast(this->max_new_tokens_align); + constexpr uint32_t MAX_CHUNK_ELEMS = 8192u; + + auto tokenLocal = tokenTileBuf.Get(); + uint32_t msl = static_cast(this->max_seq_len); + for (uint32_t r = 0; r < rows; ++r) { + uint64_t gmRow = static_cast(start_row + r) * msl; + uint32_t ubRow = r * msa; + for (uint32_t off = 0; off < msl; off += MAX_CHUNK_ELEMS) { + uint32_t chunk = (off + MAX_CHUNK_ELEMS <= msl) ? MAX_CHUNK_ELEMS : (msl - off); + uint32_t isLast = (off + chunk >= msl) ? 1u : 0u; + uint32_t dstChunk = isLast ? (msa - off) : MAX_CHUNK_ELEMS; + uint8_t pad = static_cast(dstChunk - chunk); + AscendC::DataCopyExtParams p{1, chunk * ELEM_SIZE, 0, dstChunk * ELEM_SIZE, 0}; + AscendC::DataCopyPadExtParams pp{false, 0, pad, 0}; + AscendC::DataCopyPad(tokenLocal[ubRow + off], tokenGm[gmRow + off], p, pp); + } + } + + auto sampledLocal = sampledTileBuf.Get(); + uint32_t srcRowBytes2 = static_cast(this->max_new_tokens) * ELEM_SIZE; + uint32_t dstRowBytes2 = mnta * ELEM_SIZE; + AscendC::DataCopyExtParams sampledParams{1, srcRowBytes2, 0, dstRowBytes2, 0}; + AscendC::DataCopyPadExtParams sampledPad{ + false, 0, static_cast(mnta - this->max_new_tokens), 0}; + for (uint32_t r = 0; r < rows; ++r) { + AscendC::DataCopyPad(sampledLocal[static_cast(r) * mnta], + sampledGm[static_cast(start_row + r) * this->max_new_tokens], + sampledParams, sampledPad); + } + + auto numTokensLocal = numTokensBuf.Get(); + uint32_t metaBytes = static_cast(rows) * ELEM_SIZE; + AscendC::DataCopyExtParams metaParams{1, metaBytes, 0, metaBytes, 0}; + AscendC::DataCopyPadExtParams noPadT{false, 0, 0, 0}; + AscendC::DataCopyPad(numTokensLocal, numTokensGm[start_row], metaParams, noPadT); + + auto discardLocal = discardTileBuf.Get(); + AscendC::DataCopyPad(discardLocal, discardGm[start_row], metaParams, noPadT); + } + + __aicore__ inline void Compute(uint32_t rows) + { + auto tokenLocal = tokenTileBuf.Get(); + auto sampledLocal = sampledTileBuf.Get(); + auto numTokensLocal = numTokensBuf.Get(); + auto discardLocal = discardTileBuf.Get(); + auto nextLocal = nextTokenBuf.Get(); + auto draftLocal = draftBuf.Get(); + auto numValidLocal = numValidBuf.Get(); + auto suffixLocal = suffixBuf.Get(); + auto maskLocal = maskBuf.Get(); + + for (uint32_t i = 0; i < rows; ++i) { + ComputeOneRow(i, tokenLocal, sampledLocal, numTokensLocal, + discardLocal, nextLocal, draftLocal, numValidLocal, + suffixLocal, maskLocal); + } + } + + __aicore__ inline void ComputeOneRow( + uint32_t idx, + AscendC::LocalTensor &tokenLocal, + AscendC::LocalTensor &sampledLocal, + AscendC::LocalTensor &numTokensLocal, + AscendC::LocalTensor &discardLocal, + AscendC::LocalTensor &nextLocal, + AscendC::LocalTensor &draftLocal, + AscendC::LocalTensor &numValidLocal, + AscendC::LocalTensor &suffixLocal, + AscendC::LocalTensor &maskLocal) + { + uint32_t msa = this->max_seq_len_align; + uint32_t mnta = this->max_new_tokens_align; + uint32_t ka = this->k_align; + + int32_t seq_len = numTokensLocal.GetValue(idx); + int32_t discard = discardLocal.GetValue(idx); + int32_t valid_count = 0; + + int32_t backup_pos = (seq_len > 0) ? (seq_len - 1) : 0; + int32_t backup_token = tokenLocal.GetValue(idx * msa + backup_pos); + + for (int32_t j = 0; j < this->max_new_tokens; ++j) { + int32_t val = sampledLocal.GetValue(idx * mnta + j); + if (discard != 0) { + sampledLocal.SetValue(idx * mnta + j, -1); + } else if (val != -1 && val < this->vocab_size_val) { + valid_count++; + } else { + sampledLocal.SetValue(idx * mnta + j, -1); + } + } + + int32_t avail_space = this->max_seq_len - seq_len; + if (avail_space < 0) avail_space = 0; + if (valid_count > avail_space) valid_count = avail_space; + + if (valid_count > 0) { + nextLocal.SetValue(idx, sampledLocal.GetValue(idx * mnta + valid_count - 1)); + } else { + nextLocal.SetValue(idx, backup_token); + } + + int32_t num_tokens_tmp = seq_len + valid_count; + for (int32_t j = 0; j < valid_count; ++j) { + tokenLocal.SetValue(idx * msa + seq_len + j, sampledLocal.GetValue(idx * mnta + j)); + } + + int32_t best_match_pos = -1; + int32_t best_ngram_len = 0; + + if (valid_count > 0 && num_tokens_tmp >= this->min_n_val) { + if (this->block_rows <= 1) { + int32_t nt = num_tokens_tmp; + constexpr uint32_t CMP_MAX = 8192u; + + for (int32_t ngram_len = this->min_n_val; ngram_len <= this->max_n_val; ++ngram_len) { + if (ngram_len > nt) break; + int32_t wc = nt - ngram_len; + if (wc <= 0) break; + + int32_t suffix0 = tokenLocal.GetValue(static_cast(nt - ngram_len)); + uint32_t msa_cmp = static_cast(msa); + + for (int32_t cmp_off = 0; cmp_off < wc; cmp_off += CMP_MAX) { + uint32_t remaining = static_cast(wc - cmp_off); + uint32_t elements = (remaining >= CMP_MAX) ? CMP_MAX : remaining; + uint32_t count_aligned = ((elements + 63u) / 64u) * 64u; + uint32_t buf_avail = msa_cmp - static_cast(cmp_off); + if (count_aligned > buf_avail) { + count_aligned = (buf_avail / 64u) * 64u; + } + + if (count_aligned == 0) { + for (int32_t p = 0; p < static_cast(elements); ++p) { + if (tokenLocal.GetValue(static_cast(cmp_off + p)) == suffix0) { + bool all_match = true; + for (int32_t s = 1; s < ngram_len; ++s) { + int32_t sv = tokenLocal.GetValue(static_cast(nt - ngram_len + s)); + int32_t tv = tokenLocal.GetValue(static_cast(cmp_off + p + s)); + if (tv != sv) { all_match = false; break; } + } + if (all_match) { + best_match_pos = cmp_off + p; + best_ngram_len = ngram_len; + break; + } + } + } + } else { + AscendC::CompareScalar( + maskLocal, tokenLocal[static_cast(cmp_off)], + suffix0, AscendC::CMPMODE::EQ, count_aligned); + + for (int32_t p = 0; p < static_cast(elements); ++p) { + uint8_t byte_val = maskLocal.GetValue(static_cast(p) >> 3); + if (byte_val & (1u << (static_cast(p) & 7u))) { + bool all_match = true; + for (int32_t s = 1; s < ngram_len; ++s) { + int32_t sv = tokenLocal.GetValue(static_cast(nt - ngram_len + s)); + int32_t tv = tokenLocal.GetValue(static_cast(cmp_off + p + s)); + if (tv != sv) { all_match = false; break; } + } + if (all_match) { + best_match_pos = cmp_off + p; + best_ngram_len = ngram_len; + break; + } + } + } + } + if (best_match_pos >= 0) break; + } + } + } else { + int32_t row_base = static_cast(idx) * static_cast(msa); + + for (int32_t ngram_len = this->min_n_val; ngram_len <= this->max_n_val; ++ngram_len) { + if (ngram_len > num_tokens_tmp) break; + + for (int32_t s = 0; s < ngram_len; ++s) { + suffixLocal.SetValue(static_cast(s), + tokenLocal.GetValue(static_cast( + row_base + num_tokens_tmp - ngram_len + s))); + } + + int32_t max_pos = num_tokens_tmp - ngram_len - 1; + for (int32_t pos = 0; pos <= max_pos; ++pos) { + bool match = true; + for (int32_t s = 0; s < ngram_len; ++s) { + if (tokenLocal.GetValue(static_cast(row_base + pos + s)) + != suffixLocal.GetValue(static_cast(s))) { + match = false; + break; + } + } + if (match) { + best_match_pos = pos; + best_ngram_len = ngram_len; + break; + } + } + } + } + } + + if (best_match_pos >= 0) { + int32_t draft_start = best_match_pos + best_ngram_len; + int32_t tokens_available = num_tokens_tmp - draft_start; + for (int32_t j = 0; j < this->k_val; ++j) { + if (j < tokens_available) { + draftLocal.SetValue(idx * ka + j, tokenLocal.GetValue(idx * msa + draft_start + j)); + } else { + draftLocal.SetValue(idx * ka + j, -1); + } + } + } else { + for (int32_t j = 0; j < this->k_val; ++j) { + draftLocal.SetValue(idx * ka + j, -1); + } + } + + int32_t valid_draft_count = 0; + for (int32_t j = 0; j < this->k_val; ++j) { + if (draftLocal.GetValue(idx * ka + j) != -1) { + valid_draft_count++; + } else { + break; + } + } + numValidLocal.SetValue(idx, valid_draft_count); + } + + __aicore__ inline void CopyOut(uint32_t start_row, uint32_t rows) + { + uint32_t msa = static_cast(this->max_seq_len_align); + uint32_t msl = static_cast(this->max_seq_len); + constexpr uint32_t OUT_CHUNK_ELEMS = 8192u; + + auto tokenLocal = tokenTileBuf.Get(); + for (uint32_t r = 0; r < rows; ++r) { + uint64_t gmRow = static_cast(start_row + r) * msl; + uint32_t ubRow = r * msa; + for (uint32_t off = 0; off < msl; off += OUT_CHUNK_ELEMS) { + uint32_t chunk = (off + OUT_CHUNK_ELEMS <= msl) ? OUT_CHUNK_ELEMS : (msl - off); + AscendC::DataCopyExtParams p{1, chunk * ELEM_SIZE, 0, 0, 0}; + AscendC::DataCopyPad(tokenGm[gmRow + off], tokenLocal[ubRow + off], p); + } + } + + auto nextLocal = nextTokenBuf.Get(); + uint32_t metaBytes32 = static_cast(rows) * ELEM_SIZE; + AscendC::DataCopyExtParams nextParams{1, metaBytes32, 0, 0, 0}; + AscendC::DataCopyPad(nextTokensGm[start_row], nextLocal, nextParams); + + auto draftLocal = draftBuf.Get(); + uint32_t kBytes = static_cast(this->k_val) * ELEM_SIZE; + for (uint32_t r = 0; r < rows; ++r) { + AscendC::DataCopyExtParams draftRowParams{1, kBytes, 0, 0, 0}; + AscendC::DataCopyPad( + draftTokensGm[static_cast(start_row + r) * this->k_val], + draftLocal[static_cast(r) * this->k_align], draftRowParams); + } + + auto numValidLocal = numValidBuf.Get(); + AscendC::DataCopyPad(numValidGm[start_row], numValidLocal, nextParams); + } + +private: + AscendC::TPipe pipe; + AscendC::TBuf tokenTileBuf; + AscendC::TBuf sampledTileBuf; + AscendC::TBuf numTokensBuf; + AscendC::TBuf discardTileBuf; + AscendC::TBuf nextTokenBuf; + AscendC::TBuf draftBuf; + AscendC::TBuf numValidBuf; + AscendC::TBuf suffixBuf; + AscendC::TBuf maskBuf; + + AscendC::GlobalTensor tokenGm; + AscendC::GlobalTensor numTokensGm; + AscendC::GlobalTensor sampledGm; + AscendC::GlobalTensor discardGm; + AscendC::GlobalTensor nextTokensGm; + AscendC::GlobalTensor draftTokensGm; + AscendC::GlobalTensor numValidGm; + + int32_t batch_size; + int32_t max_seq_len; + int32_t max_seq_len_align; + int32_t max_new_tokens; + int32_t max_new_tokens_align; + int32_t k_val; + int32_t k_align; + int32_t vocab_size_val; + int32_t min_n_val; + int32_t max_n_val; + int32_t former_num; + int32_t rows_per_core; + int32_t tail_rows; + int32_t block_rows; + uint32_t my_rows; + uint32_t row_offset; + bool is_large_row; +}; + +extern "C" __global__ __aicore__ void ngram_spec_decode( + GM_ADDR token_ids, GM_ADDR num_tokens, GM_ADDR sampled, + GM_ADDR discard, GM_ADDR next_tokens, GM_ADDR draft_tokens, + GM_ADDR num_valid, GM_ADDR workspace, GM_ADDR tiling) +{ + KernelNgramSpecDecode op; + op.Init(token_ids, num_tokens, sampled, discard, next_tokens, + draft_tokens, num_valid, workspace, tiling); + op.Process(); +} diff --git a/csrc/ascend/attention/ngram_spec_decode/op_kernel/ngram_spec_decode.h b/csrc/ascend/attention/ngram_spec_decode/op_kernel/ngram_spec_decode.h new file mode 100644 index 000000000..f8eacb51d --- /dev/null +++ b/csrc/ascend/attention/ngram_spec_decode/op_kernel/ngram_spec_decode.h @@ -0,0 +1,26 @@ +#ifndef NGRAM_SPEC_DECODE_H +#define NGRAM_SPEC_DECODE_H + +#include "kernel_tiling/kernel_tiling.h" + +struct NgramSpecDecodeInfo { + uint32_t batchSize; + uint32_t maxSeqLen; + uint32_t maxNewTokens; + uint32_t vocabSize; + uint32_t minN; + uint32_t maxN; + uint32_t k; + uint32_t formerNum; + uint32_t rowsPerCore; + uint32_t tailRows; + uint32_t blockRows; +}; + +struct NgramSpecDecodeTilingData { + Mc2InitTiling mc2InitTiling; + Mc2CcTiling mc2CcTiling1; + NgramSpecDecodeInfo ngramInfo; +}; + +#endif // NGRAM_SPEC_DECODE_H diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/CMakeLists.txt b/csrc/ascend/attention/recurrent_gated_delta_rule/CMakeLists.txt new file mode 100644 index 000000000..549b5c14c --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/CMakeLists.txt b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/CMakeLists.txt new file mode 100644 index 000000000..0dcaab210 --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/CMakeLists.txt @@ -0,0 +1,22 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnExc PRIVATE + recurrent_gated_delta_rule_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME RecurrentGatedDeltaRule + OPTIONS + --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE recurrent_gated_delta_rule ACLNNTYPE aclnn_exclude) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/arch35/recurrent_gated_delta_rule_tiling_arch35.cpp b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/arch35/recurrent_gated_delta_rule_tiling_arch35.cpp new file mode 100644 index 000000000..05e3b7cc9 --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/arch35/recurrent_gated_delta_rule_tiling_arch35.cpp @@ -0,0 +1,207 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file recurrent_gated_delta_rule_tiling_arch35.cpp + * \brief + */ +#include "recurrent_gated_delta_rule_tiling.h" + +#include + +#include "platform/platform_ascendc.h" +#include "tiling_base/tiling_templates_registry.h" + +namespace optiling { +namespace { + +constexpr uint64_t RGDR_ASCEND_950_TEMPLATE_PRIORITY = 1000; + +constexpr size_t QUERY_INDEX = 0; +constexpr size_t KEY_INDEX = 1; +constexpr size_t VALUE_INDEX = 2; +constexpr size_t BETA_INDEX = 3; +constexpr size_t STATE_INDEX = 4; +constexpr size_t CUSEQLENS_INDEX = 5; +constexpr size_t SSM_STATE_INDICES_INDEX = 6; + +constexpr size_t DIM_0 = 0; +constexpr size_t DIM_1 = 1; +constexpr size_t DIM_2 = 2; + +class RecurrentGatedDeltaRuleTilingArch35 final : public RecurrentGatedDeltaRuleTiling { +public: + explicit RecurrentGatedDeltaRuleTilingArch35(gert::TilingContext *context) + : RecurrentGatedDeltaRuleTiling(context) + { + } + +protected: + bool IsCapable() override + { + auto platformInfo = context_->GetPlatformInfo(); + if (platformInfo == nullptr) { + return false; + } + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + return ascendcPlatform.GetSocVersion() == platform_ascendc::SocVersion::ASCEND950; + } + + ge::graphStatus GetShapeAttrsInfo() override + { + OP_CHECK_IF(CheckContext() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "Invalid context."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(AnalyzeDtype() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "Invalid dtypes."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(AnalyzeShapesArch35() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "Invalid shapes."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(GetScale() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "Invalid GetScale."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(GetOptionalInput() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "Invalid GetOptionalInput."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(AnalyzeFormat() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "Invalid Format."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; + } + + ge::graphStatus DoOpTiling() override + { + OP_CHECK_IF(CalUbSizeArch35() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "CalUbSize failed."), + return ge::GRAPH_FAILED); + + PrintTilingData(); + return ge::GRAPH_SUCCESS; + } + +private: + ge::graphStatus AnalyzeShapesArch35() + { + const auto &queryShape = context_->GetInputShape(QUERY_INDEX)->GetOriginShape(); + const auto &keyShape = context_->GetInputShape(KEY_INDEX)->GetOriginShape(); + const auto &valueShape = context_->GetInputShape(VALUE_INDEX)->GetOriginShape(); + const auto &betaShape = context_->GetInputShape(BETA_INDEX)->GetOriginShape(); + const auto &stateShape = context_->GetInputShape(STATE_INDEX)->GetOriginShape(); + const auto &cuSeqlensShape = context_->GetInputShape(CUSEQLENS_INDEX)->GetOriginShape(); + const auto &ssmStateShape = context_->GetInputShape(SSM_STATE_INDICES_INDEX)->GetOriginShape(); + + OP_CHECK_IF(CheckShapeDimAndRelation(queryShape, keyShape, valueShape, betaShape, stateShape, cuSeqlensShape, + ssmStateShape) != ge::GRAPH_SUCCESS, + OP_LOGE(inputParams_.opName, "AnalyzeShapes rule failed: CheckShapeDimAndRelation"), + return ge::GRAPH_FAILED); + + tilingData_.t = queryShape.GetDim(DIM_0); + tilingData_.nk = queryShape.GetDim(DIM_1); + tilingData_.dk = queryShape.GetDim(DIM_2); + tilingData_.nv = valueShape.GetDim(DIM_1); + tilingData_.dv = valueShape.GetDim(DIM_2); + tilingData_.sBlockNum = stateShape.GetDim(DIM_0); + tilingData_.b = cuSeqlensShape.GetDim(DIM_0) - 1; + + OP_CHECK_IF(CheckShapeValueRangeAndRule() != ge::GRAPH_SUCCESS, + OP_LOGE(inputParams_.opName, "AnalyzeShapes rule failed: CheckShapeValueRangeAndRule"), + return ge::GRAPH_FAILED); + + UpdateDynamicBlockDimByTaskUnits(); + return ge::GRAPH_SUCCESS; + } + + ge::graphStatus CalUbSizeArch35() + { + struct RuleItem { + const char *name; + HostRuleFn fn; + }; + + OP_CHECK_IF(RuleInitUbCalcContext() != ge::GRAPH_SUCCESS, + OP_LOGE(inputParams_.opName, "CalUbSize rule failed: RuleInitUbCalcContext"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(RuleCalcFixedUbBytes() != ge::GRAPH_SUCCESS, + OP_LOGE(inputParams_.opName, "CalUbSize rule failed: RuleCalcFixedUbBytes"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(RuleCalcWorkingUbBytes() != ge::GRAPH_SUCCESS, + OP_LOGE(inputParams_.opName, "CalUbSize rule failed: RuleCalcWorkingUbBytes"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(RuleCalcVStepCoeff() != ge::GRAPH_SUCCESS, + OP_LOGE(inputParams_.opName, "CalUbSize rule failed: RuleCalcVStepCoeff"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(FinalizeVStepFromUbArch35() != ge::GRAPH_SUCCESS, + OP_LOGE(inputParams_.opName, "CalUbSize rule failed: FinalizeVStepFromUbArch35"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; + } + + ge::graphStatus FinalizeVStepFromUbArch35() + { + BufferProfile selected; + const std::array candidates = {{ + BufferProfile(1u, 1u, 0u, 0u, false), + BufferProfile(1u, 2u, 0u, 0u, false), + BufferProfile(2u, 2u, 0u, 0u, false), + }}; + + for (const auto &candidate : candidates) { + BufferProfile profile; + if (!EvaluateBufferProfile(ubCalcCtx_.ubSize, ubCalcCtx_.workingUbBytes, ubCalcCtx_.aDk, + candidate.stateOutBufferNum, candidate.attnOutBufferNum, profile)) { + continue; + } + if (IsBetterProfile(profile, selected)) { + selected = profile; + } + } + + OP_LOGD(context_->GetNodeName(), + "selected profile: stateOutBufferNum=[%u], attnOutBufferNum=[%u], vStep=[%u], repeatTime=[%u], " + "valid=[%d]", + selected.stateOutBufferNum, selected.attnOutBufferNum, selected.vStep, selected.repeatTime, + selected.valid); + + if (!selected.valid) { + OP_LOGE(context_->GetNodeName(), "vStep should be bigger than 8, shape is too big"); + return ge::GRAPH_FAILED; + } + + auto stateDtype = context_->GetInputDesc(STATE_INDEX)->GetDataType(); + int64_t stateDtypeSize = (stateDtype == ge::DT_FLOAT) ? 4 : 2; + int64_t queueCoeff = + (stateDtypeSize + static_cast(stateDtypeSize * selected.stateOutBufferNum)) * ubCalcCtx_.aDk + + static_cast(4 * selected.attnOutBufferNum); + int64_t ubRestBytes = + ubCalcCtx_.ubSize - ubCalcCtx_.fixedUbBytes - queueCoeff * static_cast(selected.vStep); + if (ubRestBytes < 0) { + OP_LOGE(context_->GetNodeName(), "ubRestBytes should be non-negative, but got %ld", ubRestBytes); + return ge::GRAPH_FAILED; + } + + tilingData_.ubCalSize = compileInfo_.ubSize; + tilingData_.vStep = selected.vStep; + tilingData_.stateOutBufferNum = selected.stateOutBufferNum; + tilingData_.attnOutBufferNum = selected.attnOutBufferNum; + tilingData_.ubRestBytes = static_cast(ubRestBytes); + return ge::GRAPH_SUCCESS; + } +}; + +} // namespace + +REGISTER_OPS_TILING_TEMPLATE(RecurrentGatedDeltaRule, + RecurrentGatedDeltaRuleTilingArch35, + RGDR_ASCEND_950_TEMPLATE_PRIORITY); + +} // namespace optiling diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/math_util.h b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/math_util.h new file mode 100644 index 000000000..14f90fc4c --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/math_util.h @@ -0,0 +1,61 @@ +/** +* Copyright (c) 2025 Huawei Technologies Co., Ltd. +* This program is free software, you can redistribute it and/or modify it under the terms and conditions of +* CANN Open Software License Agreement Version 2.0 (the "License"). +* Please refer to the License for details. You may not use this file except in compliance with the License. +* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +* See LICENSE in the root of the software repository for the full text of the License. +*/ + +/*! + * \file math_util.h + * \brief + */ + +#ifndef TILING_MATMUL_MATH_UTIL_H +#define TILING_MATMUL_MATH_UTIL_H + +#include +#include +#include +#include +namespace matmul_tiling { +class MathUtil { +public: + static bool IsEqual(float leftValue, float rightValue); + template + static auto CeilDivision(T num1, T num2) -> T + { + if (num2 == 0) { + return 0; + } + return static_cast((static_cast(num1) + static_cast(num2) - 1) / + static_cast(num2)); + } + template + static auto Align(T num1, T num2) -> T + { + return CeilDivision(num1, num2) * num2; + } + static int32_t AlignDown(int32_t num1, int32_t num2); + static bool CheckMulOverflow(int32_t a, int32_t b, int32_t &c); + static int32_t MapShape(int32_t shape, bool roundUpFlag = true); + static void AddFactor(std::vector &dimsFactors, int32_t dim); + static void GetFactorCnt(const int32_t shape, int32_t &factorCnt, const int32_t factorStart, + const int32_t factorEnd); + static void GetFactorLayerCnt(const int32_t shape, int32_t &factorCnt, const int32_t factorStart, + const int32_t factorEnd); + static bool CheckFactorNumSatisfy(const int32_t dim); + static int32_t FindBestSingleCore(const int32_t oriShape, const int32_t mappedShape, const int32_t coreNum, + bool isKDim); + static void GetFactors(std::vector &factorList, int32_t srcNum, int32_t minFactor, int32_t maxFactor); + static void GetFactors(std::vector &factorList, int32_t srcNum, int32_t maxFactor); + static void GetBlockFactors(std::vector &factorList, const int32_t oriShape, const int32_t mpShape, + const int32_t coreNum, const int32_t maxNum); + static int32_t GetNonFactorMap(std::vector &factorList, int32_t srcNum, int32_t maxFactor); + static std::vector> GetFactorPairs(int32_t num); + static std::pair DivideIntoMainAndTail(int32_t num, int32_t divisor); +}; +} // namespace matmul_tiling +#endif // _MATH_UTIL_H_ \ No newline at end of file diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/op_api/aclnn_recurrent_gated_delta_rule.cpp b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/op_api/aclnn_recurrent_gated_delta_rule.cpp new file mode 100644 index 000000000..8347a9615 --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/op_api/aclnn_recurrent_gated_delta_rule.cpp @@ -0,0 +1,207 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file aclnn_recurrent_gated_delta_rule.cpp + * \brief + */ +#include +#include "aclnn_recurrent_gated_delta_rule.h" +#include "../recurrent_gated_delta_rule.h" + +#include "securec.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "opdev/common_types.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/platform.h" + +#include "aclnn_kernels/transdata.h" +#include "aclnn_kernels/transpose.h" +#include "aclnn_kernels/contiguous.h" +#include "aclnn_kernels/reshape.h" + +using namespace op; + +#ifdef __cplusplus +extern "C" { +#endif + +namespace { +constexpr size_t QUERY_DIM_NUM = 3; +constexpr size_t KEY_DIM_NUM = 3; +constexpr size_t VALUE_DIM_NUM = 3; +constexpr size_t BETA_DIM_NUM = 2; +constexpr size_t STATE_DIM_NUM = 4; + +struct RecurrentGatedDeltaRuleParams { + // mandatory + const aclTensor *query {nullptr}; + const aclTensor *key {nullptr}; + const aclTensor *value {nullptr}; + const aclTensor *beta {nullptr}; + const aclTensor *state {nullptr}; + const aclTensor *actual_seq_lengths {nullptr}; + const aclTensor *ssm_state_indices {nullptr}; + // optional + const aclTensor *g {nullptr}; + const aclTensor *gk {nullptr}; + const aclTensor *num_accepted_tokens {nullptr}; + // attrs + float scale {1.0f}; + //output + const aclTensor *out {nullptr}; +}; + +// support dtype +static const std::initializer_list QKV_TYPE_SUPPORT_LIST = {op::DataType::DT_BF16}; +static const std::initializer_list STATE_TYPE_SUPPORT_LIST = {op::DataType::DT_BF16,op::DataType::DT_FLOAT}; +static const std::initializer_list BETA_TYPE_SUPPORT_LIST = {op::DataType::DT_BF16}; +static const std::initializer_list SEQ_LENS_TYPE_SUPPORT_LIST = {op::DataType::DT_INT32}; +static const std::initializer_list SSM_TYPE_SUPPORT_LIST = {op::DataType::DT_INT32}; +static const std::initializer_list G_TYPE_SUPPORT_LIST = {op::DataType::DT_FLOAT}; +static const std::initializer_list ACC_TO_TYPE_SUPPORT_LIST = {op::DataType::DT_INT32}; +static const std::initializer_list OUT_TYPE_SUPPORT_LIST = {op::DataType::DT_BF16}; + +static inline bool CheckNotNull(const RecurrentGatedDeltaRuleParams ¶ms) +{ + // 必选参数 + OP_CHECK_NULL(params.query, return false); + OP_CHECK_NULL(params.key, return false); + OP_CHECK_NULL(params.value, return false); + OP_CHECK_NULL(params.state, return false); + OP_CHECK_NULL(params.beta, return false); + OP_CHECK_NULL(params.actual_seq_lengths, return false); + OP_CHECK_NULL(params.ssm_state_indices, return false); + OP_CHECK_NULL(params.out, return false); + + return true; +} + +static inline bool CheckDtypeVaild(const RecurrentGatedDeltaRuleParams ¶ms) +{ + // 检查必选参数数据类型 + OP_CHECK_DTYPE_NOT_SUPPORT(params.query, QKV_TYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(params.key, QKV_TYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(params.value, QKV_TYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(params.state, STATE_TYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(params.beta, BETA_TYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(params.actual_seq_lengths, SEQ_LENS_TYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(params.ssm_state_indices, SSM_TYPE_SUPPORT_LIST, return false); + + // 检查可选参数数据类型 + if (params.g != nullptr) { + OP_CHECK_DTYPE_NOT_SUPPORT(params.g, G_TYPE_SUPPORT_LIST, return false); + } + if (params.gk != nullptr) { + OP_CHECK_DTYPE_NOT_SUPPORT(params.gk, G_TYPE_SUPPORT_LIST, return false); + } + if (params.num_accepted_tokens != nullptr) { + OP_CHECK_DTYPE_NOT_SUPPORT(params.num_accepted_tokens, ACC_TO_TYPE_SUPPORT_LIST, return false); + } + + OP_CHECK_DTYPE_NOT_SUPPORT(params.out, OUT_TYPE_SUPPORT_LIST, return false); + return true; +} + +static aclnnStatus CheckParams(RecurrentGatedDeltaRuleParams ¶ms) +{ + // 检查输入参数是否在支持的数据类型范围内 + CHECK_RET(CheckDtypeVaild(params), ACLNN_ERR_PARAM_INVALID); + + OP_LOGD("RecurrentGatedDeltaRule check params success."); + + return ACLNN_SUCCESS; +} + +static aclnnStatus PreProcess(RecurrentGatedDeltaRuleParams ¶ms) +{ + params.query->SetOriginalShape(params.query->GetViewShape()); + params.key->SetOriginalShape(params.key->GetViewShape()); + params.value->SetOriginalShape(params.value->GetViewShape()); + params.beta->SetOriginalShape(params.beta->GetViewShape()); + params.state->SetOriginalShape(params.state->GetViewShape()); + params.actual_seq_lengths->SetOriginalShape(params.actual_seq_lengths->GetViewShape()); + params.ssm_state_indices->SetOriginalShape(params.ssm_state_indices->GetViewShape()); + + return ACLNN_SUCCESS; +} +} // namespace + +aclnnStatus aclnnRecurrentGatedDeltaRuleGetWorkspaceSize(const aclTensor *query, const aclTensor *key, + const aclTensor *value, const aclTensor *beta, + aclTensor *stateRef, const aclTensor *actualSeqLengths, + const aclTensor *ssmStateIndices, const aclTensor *g, + const aclTensor *gk, const aclTensor *numAcceptedTokens, + float scaleValue, aclTensor *out, uint64_t *workspaceSize, + aclOpExecutor **executor) +{ + L2_DFX_PHASE_1(aclnnRecurrentGatedDeltaRule, + DFX_IN(query, key, value, beta, stateRef, actualSeqLengths, ssmStateIndices, g, gk, + numAcceptedTokens, scaleValue), + DFX_OUT(out, stateRef)); + + auto uniqueExecutor = CREATE_EXECUTOR(); + CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + + RecurrentGatedDeltaRuleParams params {query, key, value, beta, stateRef, actualSeqLengths, ssmStateIndices, g, gk, numAcceptedTokens,scaleValue, out}; + + CHECK_RET(CheckNotNull(params), ACLNN_ERR_PARAM_INVALID); + CHECK_RET(CheckParams(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + auto ret = PreProcess(params); + CHECK_RET(ret == ACLNN_SUCCESS, ret); + + auto query_ = l0op::Contiguous(query, uniqueExecutor.get()); + auto key_ = l0op::Contiguous(key, uniqueExecutor.get()); + auto value_ = l0op::Contiguous(value, uniqueExecutor.get()); + auto beta_ = l0op::Contiguous(beta, uniqueExecutor.get()); + auto actualSeqLengths_ = l0op::Contiguous(actualSeqLengths, uniqueExecutor.get()); + auto ssmStateIndices_ = l0op::Contiguous(ssmStateIndices, uniqueExecutor.get()); + if (g != nullptr) { + g = l0op::Contiguous(g, uniqueExecutor.get()); + } + if (gk != nullptr) { + gk = l0op::Contiguous(gk, uniqueExecutor.get()); + } + if (numAcceptedTokens != nullptr) { + numAcceptedTokens = l0op::Contiguous(numAcceptedTokens, uniqueExecutor.get()); + } + + auto out_ = l0op::Contiguous(out, uniqueExecutor.get()); + + // 调用l0接口 + auto outRet = + l0op::RecurrentGatedDeltaRule(query_, key_, value_, beta_, stateRef, actualSeqLengths_, ssmStateIndices_, g, gk, + numAcceptedTokens, scaleValue, uniqueExecutor.get()); + if (outRet == nullptr) { + return ACLNN_ERR_INNER_NULLPTR; + } + + auto ViewCopyResult = l0op::ViewCopy(outRet, out_, uniqueExecutor.get()); + if (ViewCopyResult == nullptr) { + return ACLNN_ERR_INNER_NULLPTR; + } + + // 获取计算过程中需要使用的workspace大小。 + *workspaceSize = uniqueExecutor->GetWorkspaceSize(); + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; +} + +aclnnStatus aclnnRecurrentGatedDeltaRule(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, + aclrtStream stream) +{ + L2_DFX_PHASE_2(aclnnRecurrentGatedDeltaRule); + return CommonOpExecutorRun(workspace, workspaceSize, executor, stream); +} + +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/op_api/aclnn_recurrent_gated_delta_rule.h b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/op_api/aclnn_recurrent_gated_delta_rule.h new file mode 100644 index 000000000..42c33d6ea --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/op_api/aclnn_recurrent_gated_delta_rule.h @@ -0,0 +1,59 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_API_ACLNN_RECURRENT_GETED_DELTA_RULE_H +#define OP_API_ACLNN_RECURRENT_GETED_DELTA_RULE_H + +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @brief RecurrentGatedDeltaRule 的第一段接口，根据具体的计算流程，计算workspace大小。 + * @param [in] query: 数据类型支持：bfloat16。 + * @param [in] key: 数据类型支持：bfloat16。 + * @param [in] value: 数据类型支持：bfloat16。 + * @param [in] beta: 数据类型支持：bfloat16。 + * @param [in] state: 数据类型支持：bfloat16。 + * @param [in] actualSeqLengths: 数据类型支持：int32。 + * @param [in] ssmStateIndices: 数据类型支持：int32。 + * @param [in] g: 数据类型支持：float32。 + * @param [in] gk: 数据类型支持：float32。 + * @param [in] numAcceptedTokens: 数据类型支持：int32。 + * @param [in] scaleValue: 数据类型支持：float32。 + * @param [out] out: 数据类型支持：bfloat16。 + * @param [out] 返回需要在npu device侧申请的workspace大小。 + * @param [out] executor: 返回op执行器，包含了算子计算流程。 + * @return aclnnStatus: 返回状态码 + */ +__attribute__((visibility("default"))) aclnnStatus aclnnRecurrentGatedDeltaRuleGetWorkspaceSize( + const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *beta, aclTensor *stateRef, + const aclTensor *actualSeqLengths, const aclTensor *ssmStateIndices, const aclTensor *g, const aclTensor *gk, + const aclTensor *numAcceptedTokens, float scaleValue, aclTensor *out, uint64_t *workspaceSize, + aclOpExecutor **executor); + +/** + * @brief + * @param [in] workspace: 在npu device侧申请的workspace内存起址。 + * @param [in] workspace_size: 在npu + * device侧申请的workspace大小，由第一段接口aclnnRecurrentGatedDeltaRuleGetWorkspaceSize获取。 + * @param [in] executor: op执行器，包含了算子计算流程。 + * @param [in] stream: acl stream流。 + * @return aclnnStatus: 返回状态码 + */ +__attribute__((visibility("default"))) aclnnStatus aclnnRecurrentGatedDeltaRule(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, + aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif // OP_API_ACLNN_RECURRENT_GETED_DELTA_RULE_H \ No newline at end of file diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/op_api/recurrent_gated_delta_rule.cpp b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/op_api/recurrent_gated_delta_rule.cpp new file mode 100644 index 000000000..55ad3c133 --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/op_api/recurrent_gated_delta_rule.cpp @@ -0,0 +1,61 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file recurrent_gated_delta_rule.cpp + * \brief + */ +#include "../recurrent_gated_delta_rule.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "opdev/make_op_executor.h" +#include "opdev/op_def.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/shape_utils.h" + +using namespace op; + +namespace l0op { + +OP_TYPE_REGISTER(RecurrentGatedDeltaRule); + +const aclTensor *RecurrentGatedDeltaRule(const aclTensor *query, const aclTensor *key, const aclTensor *value, + const aclTensor *beta, aclTensor *stateRef, const aclTensor *actualSeqLengths, + const aclTensor *ssmStateIndices, const aclTensor *g, const aclTensor *gk, + const aclTensor *numAcceptedTokens, float scaleValue, aclOpExecutor *executor) +{ + L0_DFX(RecurrentGatedDeltaRule, query, key, value, beta, stateRef, actualSeqLengths, ssmStateIndices, g, gk, + numAcceptedTokens, scaleValue); + + DataType outType = DataType::DT_BF16; + Format format = Format::FORMAT_ND; + + auto out = executor->AllocTensor(outType, format, format); + + OP_CHECK(out != nullptr, OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "out AllocTensor failed."), + return nullptr); + + // infershape + auto ret = INFER_SHAPE( + RecurrentGatedDeltaRule, + OP_INPUT(query, key, value, beta, stateRef, actualSeqLengths, ssmStateIndices, g, gk, numAcceptedTokens), + OP_OUTPUT(out, stateRef), OP_ATTR(scaleValue)); + OP_CHECK_INFERSHAPE(ret != ACLNN_SUCCESS, return nullptr, "RecurrentGatedDeltaRule InferShape failed."); + + ret = ADD_TO_LAUNCHER_LIST_AICORE( + RecurrentGatedDeltaRule, + OP_INPUT(query, key, value, beta, stateRef, actualSeqLengths, ssmStateIndices, g, gk, numAcceptedTokens), + OP_OUTPUT(out, stateRef), OP_ATTR(scaleValue)); + OP_CHECK_ADD_TO_LAUNCHER_LIST_AICORE(ret != ACLNN_SUCCESS, return nullptr, + "RecurrentGatedDeltaRule ADD_TO_LAUNCHER_LIST_AICORE failed."); + + return out; +} +} // namespace l0op diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule.h b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule.h new file mode 100644 index 000000000..15454fc7e --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule.h @@ -0,0 +1,23 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef PTA_NPU_OP_API_COMMON_INC_LEVEL0_OP_RECURRENT_GETED_DELTA_RULE +#define PTA_NPU_OP_API_COMMON_INC_LEVEL0_OP_RECURRENT_GETED_DELTA_RULE + +#include "opdev/op_executor.h" +#include "opdev/make_op_executor.h" + +namespace l0op { +const aclTensor *RecurrentGatedDeltaRule(const aclTensor *query, const aclTensor *key, const aclTensor *value, + const aclTensor *beta, aclTensor *stateRef, const aclTensor *actualSeqLengths, + const aclTensor *ssmStateIndices, const aclTensor *g, const aclTensor *gk, + const aclTensor *numAcceptedTokens, float scaleValue, aclOpExecutor *executor); +} + +#endif // PTA_NPU_OP_API_COMMON_INC_LEVEL0_OP_RECURRENT_GETED_DELTA_RULE \ No newline at end of file diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule_def.cpp b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule_def.cpp new file mode 100644 index 000000000..d7cc1af99 --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule_def.cpp @@ -0,0 +1,97 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file recurrent_gated_delta_rule.h.cpp + * \brief + */ +#include "register/op_def_registry.h" + +namespace ops { +class RecurrentGatedDeltaRule : public OpDef { +public: + explicit RecurrentGatedDeltaRule(const char *name) : OpDef(name) + { + this->Input("query") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("key") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("value") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("beta") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("state") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("actual_seq_lengths") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("ssm_state_indices") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("g") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("gk") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("num_accepted_tokens") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("out") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("state") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Attr("scale_value").AttrType(OPTIONAL).Float(1.0); + + OpAICoreConfig aicConfig; + aicConfig.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .ExtendCfgInfo("softsync.flag", "true"); + this->AICore().AddConfig("ascend910b", aicConfig); + this->AICore().AddConfig("ascend910_93", aicConfig); + } +}; + +OP_ADD(RecurrentGatedDeltaRule); + +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule_infershape.cpp b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule_infershape.cpp new file mode 100644 index 000000000..e5eae6814 --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule_infershape.cpp @@ -0,0 +1,86 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file recurrent_gated_delta_rule_infershape.cpp + * \brief + */ +#include +#include +#include +#include + +#include "exe_graph/runtime/infer_shape_context.h" +#include "exe_graph/runtime/shape.h" +#include "exe_graph/runtime/storage_shape.h" +#include "register/op_impl_registry.h" +#include "tiling_base/error_log.h" + +using namespace gert; +namespace ops { + +const size_t VALUE_INDEX = 2; +const size_t STATE_INDEX = 4; +const size_t VALUE_DIM = 3; +const size_t STATE_DIM = 4; + +const size_t DIM_0 = 0; +const size_t DIM_1 = 1; +const size_t DIM_2 = 2; +const size_t DIM_3 = 3; + +static ge::graphStatus InferShapeRecurrentGatedDeltaRule(InferShapeContext *context) +{ + if (context == nullptr) { + OP_LOGE("RecurrentGatedDeltaRule", "inference context is null"); + return ge::GRAPH_FAILED; + } + + auto opName = context->GetNodeName(); + auto shapeValue = context->GetInputShape(VALUE_INDEX); + auto shapeInitialState = context->GetInputShape(STATE_INDEX); + auto shapeOut = context->GetOutputShape(DIM_0); + auto shapeFinalState = context->GetOutputShape(DIM_1); + if (shapeValue == nullptr || shapeInitialState == nullptr || shapeOut == nullptr || shapeFinalState == nullptr) { + OP_LOGE(opName, "[InferShape] shape is null"); + return ge::GRAPH_FAILED; + } + + shapeOut->SetDimNum(VALUE_DIM); + int64_t outDim0 = shapeValue->GetDim(DIM_0); + int64_t outDim1 = shapeValue->GetDim(DIM_1); + int64_t outDim2 = shapeValue->GetDim(DIM_2); + shapeOut->SetDim(DIM_0, outDim0); + shapeOut->SetDim(DIM_1, outDim1); + shapeOut->SetDim(DIM_2, outDim2); + + shapeFinalState->SetDimNum(STATE_DIM); + int64_t stateDim0 = shapeInitialState->GetDim(DIM_0); + int64_t stateDim1 = shapeInitialState->GetDim(DIM_1); + int64_t stateDim2 = shapeInitialState->GetDim(DIM_2); + int64_t stateDim3 = shapeInitialState->GetDim(DIM_3); + shapeFinalState->SetDim(DIM_0, stateDim0); + shapeFinalState->SetDim(DIM_1, stateDim1); + shapeFinalState->SetDim(DIM_2, stateDim2); + shapeFinalState->SetDim(DIM_3, stateDim3); + + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataTypeRecurrentGatedDeltaRule(gert::InferDataTypeContext *context) +{ + context->SetOutputDataType(0, ge::DT_BF16); + context->SetOutputDataType(1, ge::DT_BF16); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(RecurrentGatedDeltaRule) + .InferShape(InferShapeRecurrentGatedDeltaRule) + .InferDataType(InferDataTypeRecurrentGatedDeltaRule); +} // namespace ops diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule_tiling.cpp b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule_tiling.cpp new file mode 100644 index 000000000..a66754ba2 --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule_tiling.cpp @@ -0,0 +1,681 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file recurrent_gated_delta_rule_tiling.cpp + * \brief + */ +#include "recurrent_gated_delta_rule_tiling.h" + +#include "tiling_base/tiling_templates_registry.h" +#include "register/op_def_registry.h" +#include "platform/platform_infos_def.h" +#include "tiling_base/error_log.h" +#include "tiling/platform/platform_ascendc.h" +#include "math_util.h" +#include "error/ops_error.h" +#include + +namespace optiling { + +REGISTER_OPS_TILING_TEMPLATE(RecurrentGatedDeltaRule, RecurrentGatedDeltaRuleTiling, 0); + +const size_t QUERY_INDEX = 0; +const size_t KEY_INDEX = 1; +const size_t VALUE_INDEX = 2; +const size_t BETA_INDEX = 3; +const size_t STATE_INDEX = 4; +const size_t CUSEQLENS_INDEX = 5; +const size_t SSM_STATE_INDICES_INDEX = 6; +const size_t G_INDEX = 7; +const size_t GK_INDEX = 8; +const size_t ACC_TO_INDEX = 9; + +const size_t QKV_DIM_NUM = 3; +const size_t BETA_DIM_NUM = 2; +const size_t STATE_DIM_NUM = 4; +const size_t CUSEQLENS_DIM_NUM = 1; +const size_t SSM_STATE_INDICES_DIM_NUM = 1; +const size_t G_DIM_NUM = 2; + +const size_t DIM_0 = 0; +const size_t DIM_1 = 1; +const size_t DIM_2 = 2; +const size_t DIM_3 = 3; + +const size_t MAX_MTP = 16; + +template +static T1 CeilDiv(T1 a, T2 b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +template +typename std::enable_if ::value, T>::type CeilAlign(T x, T align) { + return CeilDiv(x, align) * align; +} + +void RecurrentGatedDeltaRuleTiling::InitCompileInfo() +{ + auto platformInfoPtr = context_->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + OP_LOGE(context_->GetNodeName(), "platformInfoPtr is null"); + return; + } + const auto &ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, compileInfo_.ubSize); + compileInfo_.aivNum = ascendcPlatform.GetCoreNumAiv(); + + if (compileInfo_.aivNum <= 0) { + OP_LOGE(context_->GetNodeName(), "aivNum <= 0"); + return; + } + tilingData_.vectorCoreNum = compileInfo_.aivNum; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::GetPlatformInfo() +{ + return ge::GRAPH_SUCCESS; +}; + +ge::graphStatus RecurrentGatedDeltaRuleTiling::GetShapeAttrsInfo() +{ + OP_CHECK_IF(CheckContext() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "Invalid context."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(AnalyzeDtype() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "Invalid dtypes."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(AnalyzeShapes() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "Invalid shapes."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(GetScale() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "Invalid GetScale."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(GetOptionalInput() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "Invalid GetOptionalInput."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(AnalyzeFormat() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "Invalid Format."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::DoOpTiling() +{ + OP_CHECK_IF(CalUbSize() != ge::GRAPH_SUCCESS, OP_LOGE(inputParams_.opName, "CalUbSize failed."), + return ge::GRAPH_FAILED); + + PrintTilingData(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::DoLibApiTiling() +{ + tilingKey_ = 0; + return ge::GRAPH_SUCCESS; +}; + +uint64_t RecurrentGatedDeltaRuleTiling::GetTilingKey() const +{ + return tilingKey_; +}; + +ge::graphStatus RecurrentGatedDeltaRuleTiling::GetWorkspaceSize() +{ + // system workspace size is 16 * 1024 * 1024 = 16M; + constexpr int64_t sysWorkspaceSize = 16777216; + workspaceSize_ = sysWorkspaceSize; + + return ge::GRAPH_SUCCESS; +}; + +ge::graphStatus RecurrentGatedDeltaRuleTiling::PostTiling() +{ + context_->SetBlockDim(tilingData_.vectorCoreNum); + auto tilingDataSize = sizeof(RecurrentGatedDeltaRuleTilingData); + errno_t ret = memcpy_s(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity(), + reinterpret_cast(&tilingData_), tilingDataSize); + if (ret != EOK) { + OP_LOGE(context_->GetNodeName(), "memcpy_s failed, ret=%d", ret); + return ge::GRAPH_FAILED; + } + context_->GetRawTilingData()->SetDataSize(tilingDataSize); + + size_t *workspaces = context_->GetWorkspaceSizes(1); // set workspace + OP_CHECK_IF(workspaces == nullptr, OPS_REPORT_CUBE_INNER_ERR(context_->GetNodeName(), "workspaces is null"), + return ge::GRAPH_FAILED); + workspaces[0] = workspaceSize_; + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::CheckContext() +{ + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputShape(QUERY_INDEX)); + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputDesc(QUERY_INDEX)); + + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputShape(KEY_INDEX)); + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputDesc(KEY_INDEX)); + + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputShape(VALUE_INDEX)); + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputDesc(VALUE_INDEX)); + + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputShape(BETA_INDEX)); + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputDesc(BETA_INDEX)); + + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputShape(STATE_INDEX)); + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputDesc(STATE_INDEX)); + + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputShape(CUSEQLENS_INDEX)); + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputDesc(CUSEQLENS_INDEX)); + + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputShape(SSM_STATE_INDICES_INDEX)); + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetInputDesc(SSM_STATE_INDICES_INDEX)); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::AnalyzeDtype() +{ + auto queryDtype = context_->GetInputDesc(QUERY_INDEX)->GetDataType(); + auto keyDtype = context_->GetInputDesc(KEY_INDEX)->GetDataType(); + auto valueDtype = context_->GetInputDesc(VALUE_INDEX)->GetDataType(); + OP_CHECK_IF(queryDtype != ge::DT_BF16 || keyDtype != ge::DT_BF16 || valueDtype != ge::DT_BF16, + OP_LOGE(context_->GetNodeName(), "query dtype, key dtype and value dtype should be bfloat16"), + return ge::GRAPH_FAILED); + + auto betaDtype = context_->GetInputDesc(BETA_INDEX)->GetDataType(); + auto stateDtype = context_->GetInputDesc(STATE_INDEX)->GetDataType(); + OP_CHECK_IF(betaDtype != ge::DT_BF16 , + OP_LOGE(context_->GetNodeName(), "beta dtype should be bfloat16"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(stateDtype != ge::DT_FLOAT && stateDtype != ge::DT_BF16, + OP_LOGE(context_->GetNodeName(), "state dtype should be bfloat16 or float32"), + return ge::GRAPH_FAILED); + auto cuSeqlensDtype = context_->GetInputDesc(CUSEQLENS_INDEX)->GetDataType(); + auto ssmStateIndicesDtype = context_->GetInputDesc(SSM_STATE_INDICES_INDEX)->GetDataType(); + OP_CHECK_IF(cuSeqlensDtype != ge::DT_INT32 || ssmStateIndicesDtype != ge::DT_INT32, + OP_LOGE(context_->GetNodeName(), "cuSeqlens dtype and ssmStateIndices dtype should be int32"), + return ge::GRAPH_FAILED); + + if (context_->GetOptionalInputDesc(G_INDEX) != nullptr) { + auto gamaDtype = context_->GetOptionalInputDesc(G_INDEX)->GetDataType(); + OP_CHECK_IF(gamaDtype != ge::DT_FLOAT, OP_LOGE(context_->GetNodeName(), "gama dtype should be float32"), + return ge::GRAPH_FAILED); + } + + if (context_->GetOptionalInputDesc(GK_INDEX) != nullptr) { + auto gamaKDtype = context_->GetOptionalInputDesc(GK_INDEX)->GetDataType(); + OP_CHECK_IF(gamaKDtype != ge::DT_FLOAT, OP_LOGE(context_->GetNodeName(), "gamaK dtype should be float32"), + return ge::GRAPH_FAILED); + } + + if (context_->GetOptionalInputDesc(ACC_TO_INDEX) != nullptr) { + auto numAcceptedTokensDtype = context_->GetOptionalInputDesc(ACC_TO_INDEX)->GetDataType(); + OP_CHECK_IF(numAcceptedTokensDtype != ge::DT_INT32, + OP_LOGE(context_->GetNodeName(), "numAcceptedTokens dtype should be int32"), + return ge::GRAPH_FAILED); + } + + return ge::GRAPH_SUCCESS; +} + + +bool RecurrentGatedDeltaRuleTiling::CheckDimEqual(const gert::Shape a, const int64_t dimA, gert::Shape b, const int64_t dimB, + const std::string &nameA, const std::string &nameB, + const std::string &dimDesc) +{ + if (a.GetDim(dimA) != b.GetDim(dimB)) { + OP_LOGE(context_->GetNodeName(), "The %s of %s and %s should be the same, but %s is %ld while %s is %ld", + dimDesc.c_str(), nameA.c_str(), nameB.c_str(), nameA.c_str(), a.GetDim(dimA), nameB.c_str(), + b.GetDim(dimB)); + return false; + } + return true; +} + +bool RecurrentGatedDeltaRuleTiling::CheckDim(const gert::Shape shape, const size_t dim, const std::string &dimDesc) +{ + if (shape.GetDimNum() != dim) { + OP_LOGE(context_->GetNodeName(), "The number of dimensions of %s should be %zu, but it is %zu", + dimDesc.c_str(), dim, shape.GetDimNum()); + return false; + } + return true; +} + +// Split shape checks/fill/scheduling decisions to improve readability and maintenance. +ge::graphStatus RecurrentGatedDeltaRuleTiling::CheckShapeDimAndRelation(const gert::Shape &queryShape, + const gert::Shape &keyShape, + const gert::Shape &valueShape, + const gert::Shape &betaShape, + const gert::Shape &stateShape, + const gert::Shape &cuSeqlensShape, + const gert::Shape &ssmStateShape) +{ + if (!CheckDim(queryShape, QKV_DIM_NUM, "query") || !CheckDim(keyShape, QKV_DIM_NUM, "key") || + !CheckDim(valueShape, QKV_DIM_NUM, "value") || !CheckDim(betaShape, BETA_DIM_NUM, "beta") || + !CheckDim(stateShape, STATE_DIM_NUM, "state") || + !CheckDim(cuSeqlensShape, CUSEQLENS_DIM_NUM, "actual_seq_lengths") || + !CheckDim(ssmStateShape, SSM_STATE_INDICES_DIM_NUM, "ssm_state_indices")) { + return ge::GRAPH_FAILED; + } + + if (!CheckDimEqual(queryShape, DIM_0, keyShape, DIM_0, "query", "key", "T dimension") || + !CheckDimEqual(queryShape, DIM_1, keyShape, DIM_1, "query", "key", "Nk dimension") || + !CheckDimEqual(queryShape, DIM_2, keyShape, DIM_2, "query", "key", "Dk dimension") || + !CheckDimEqual(stateShape, DIM_1, valueShape, DIM_1, "state", "value", "Nv dimension") || + !CheckDimEqual(stateShape, DIM_2, valueShape, DIM_2, "state", "value", "Dv dimension") || + !CheckDimEqual(valueShape, DIM_0, queryShape, DIM_0, "value", "query", "T dimension") || + !CheckDimEqual(betaShape, DIM_0, queryShape, DIM_0, "beta", "query", "T dimension") || + !CheckDimEqual(betaShape, DIM_1, valueShape, DIM_1, "beta", "value", "Nv dimension") || + !CheckDimEqual(stateShape, DIM_3, queryShape, DIM_2, "state", "query", "Dk dimension")) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +void RecurrentGatedDeltaRuleTiling::FillTilingShapeData(const gert::Shape &queryShape, const gert::Shape &valueShape, + const gert::Shape &stateShape, + const gert::Shape &cuSeqlensShape) +{ + tilingData_.t = queryShape.GetDim(DIM_0); + tilingData_.nk = queryShape.GetDim(DIM_1); + tilingData_.dk = queryShape.GetDim(DIM_2); + tilingData_.nv = valueShape.GetDim(DIM_1); + tilingData_.dv = valueShape.GetDim(DIM_2); + tilingData_.sBlockNum = stateShape.GetDim(DIM_0); + tilingData_.b = cuSeqlensShape.GetDim(DIM_0) - 1; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::CheckShapeValueRangeAndRule() +{ + OP_CHECK_IF(tilingData_.nk > 256 || tilingData_.nv > 256 || tilingData_.dk > 512 || tilingData_.dv > 512, + OP_LOGE(inputParams_.opName, + "nk and nv should no bigger than 256, dk and dv should no bigger than 512, but nk is %u, nv is " + "%u, dk is %u, dv is %u", + tilingData_.nk, tilingData_.nv, tilingData_.dk, tilingData_.dv), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(tilingData_.nv % tilingData_.nk != 0, + OP_LOGE(inputParams_.opName, + "nv should be an integer multiple of nk, but nv is %u, nk is %u", + tilingData_.nv, tilingData_.nk), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void RecurrentGatedDeltaRuleTiling::UpdateDynamicBlockDimByTaskUnits() +{ + // Dynamic blockDim: do not launch more cores than effective (batch, head) task units. + uint64_t taskUnits = static_cast(tilingData_.b) * static_cast(tilingData_.nv); + if (taskUnits == 0) { + taskUnits = 1; + } + uint64_t maxCoreNum = (compileInfo_.aivNum > 0) ? compileInfo_.aivNum : 1; + uint64_t selectedCoreNum = (taskUnits < maxCoreNum) ? taskUnits : maxCoreNum; + tilingData_.vectorCoreNum = static_cast(selectedCoreNum); + OP_LOGD(context_->GetNodeName(), "taskUnits: [%llu], selected vectorCoreNum: [%u]", + static_cast(taskUnits), tilingData_.vectorCoreNum); +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::RuleCheckShapeDimAndRelation() +{ + const auto &queryShape = context_->GetInputShape(QUERY_INDEX)->GetOriginShape(); + const auto &keyShape = context_->GetInputShape(KEY_INDEX)->GetOriginShape(); + const auto &valueShape = context_->GetInputShape(VALUE_INDEX)->GetOriginShape(); + const auto &betaShape = context_->GetInputShape(BETA_INDEX)->GetOriginShape(); + const auto &stateShape = context_->GetInputShape(STATE_INDEX)->GetOriginShape(); + const auto &cuSeqlensShape = context_->GetInputShape(CUSEQLENS_INDEX)->GetOriginShape(); + const auto &ssmStateShape = context_->GetInputShape(SSM_STATE_INDICES_INDEX)->GetOriginShape(); + return CheckShapeDimAndRelation(queryShape, keyShape, valueShape, betaShape, stateShape, cuSeqlensShape, ssmStateShape); +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::RuleFillTilingShapeData() +{ + const auto &queryShape = context_->GetInputShape(QUERY_INDEX)->GetOriginShape(); + const auto &valueShape = context_->GetInputShape(VALUE_INDEX)->GetOriginShape(); + const auto &stateShape = context_->GetInputShape(STATE_INDEX)->GetOriginShape(); + const auto &cuSeqlensShape = context_->GetInputShape(CUSEQLENS_INDEX)->GetOriginShape(); + FillTilingShapeData(queryShape, valueShape, stateShape, cuSeqlensShape); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::RuleCheckShapeValueRangeAndRule() +{ + return CheckShapeValueRangeAndRule(); +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::RuleUpdateDynamicBlockDimByTaskUnits() +{ + UpdateDynamicBlockDimByTaskUnits(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::RuleInitUbCalcContext() +{ + ubCalcCtx_.ubSize = compileInfo_.ubSize; + ubCalcCtx_.aNv = CeilAlign(tilingData_.nv, static_cast(16)); // 16 * 2 = 32B + ubCalcCtx_.aDv = CeilAlign(tilingData_.dv, static_cast(16)); // 16 * 2 = 32B + ubCalcCtx_.aDk = CeilAlign(tilingData_.dk, static_cast(16)); // 16 * 2 = 32B + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::RuleCalcFixedUbBytes() +{ + ubCalcCtx_.fixedUbBytes = CalcFixedUbBytes(ubCalcCtx_.aNv, ubCalcCtx_.aDv, ubCalcCtx_.aDk); + tilingData_.ubRestBytes = ubCalcCtx_.ubSize - ubCalcCtx_.fixedUbBytes; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::RuleCalcWorkingUbBytes() +{ + ubCalcCtx_.workingUbBytes = CalcWorkingUbBytes(ubCalcCtx_.aNv, ubCalcCtx_.aDv, ubCalcCtx_.aDk); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::RuleCalcVStepCoeff() +{ + ubCalcCtx_.coeff = CalcVStepCoeff(ubCalcCtx_.aDk, 1, 1); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::RuleFinalizeVStepFromUb() +{ + return FinalizeVStepFromUb(ubCalcCtx_.ubSize, ubCalcCtx_.workingUbBytes, ubCalcCtx_.coeff); +} + +// AnalyzeShapes now executes a deterministic rule-chain, easier to extend/maintain. +ge::graphStatus RecurrentGatedDeltaRuleTiling::AnalyzeShapes() +{ + struct RuleItem { + const char *name; + HostRuleFn fn; + }; + const std::array shapeRules = {{ + {"RuleCheckShapeDimAndRelation", &RecurrentGatedDeltaRuleTiling::RuleCheckShapeDimAndRelation}, + {"RuleFillTilingShapeData", &RecurrentGatedDeltaRuleTiling::RuleFillTilingShapeData}, + {"RuleCheckShapeValueRangeAndRule", &RecurrentGatedDeltaRuleTiling::RuleCheckShapeValueRangeAndRule}, + {"RuleUpdateDynamicBlockDimByTaskUnits", &RecurrentGatedDeltaRuleTiling::RuleUpdateDynamicBlockDimByTaskUnits}, + }}; + for (const auto &rule : shapeRules) { + OP_CHECK_IF((this->*(rule.fn))() != ge::GRAPH_SUCCESS, + OP_LOGE(inputParams_.opName, "AnalyzeShapes rule failed: %s", rule.name), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + + +bool RecurrentGatedDeltaRuleTiling::CheckFormat(ge::Format format, const std::string &Desc) +{ + if (format == ge::FORMAT_FRACTAL_NZ) { + OP_LOGE(context_->GetNodeName(), "%s format not support NZ", Desc.c_str()); + return false; + } + return true; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::AnalyzeFormat() +{ + if (!CheckFormat(context_->GetInputDesc(QUERY_INDEX)->GetStorageFormat(), "query") || + !CheckFormat(context_->GetInputDesc(KEY_INDEX)->GetStorageFormat(), "key") || + !CheckFormat(context_->GetInputDesc(VALUE_INDEX)->GetStorageFormat(), "value") || + !CheckFormat(context_->GetInputDesc(STATE_INDEX)->GetStorageFormat(), "state") || + !CheckFormat(context_->GetInputDesc(CUSEQLENS_INDEX)->GetStorageFormat(), "actual_seq_lengths") || + !CheckFormat(context_->GetInputDesc(SSM_STATE_INDICES_INDEX)->GetStorageFormat(), "ssm_state_indices")) { + return ge::GRAPH_FAILED; + } + + if (context_->GetOptionalInputDesc(G_INDEX) != nullptr) { + auto gamaFormat = context_->GetOptionalInputDesc(G_INDEX)->GetStorageFormat(); + OP_CHECK_IF(gamaFormat == ge::FORMAT_FRACTAL_NZ, OP_LOGE(context_->GetNodeName(), "gama format not support NZ"), + return ge::GRAPH_FAILED); + } + if (context_->GetOptionalInputDesc(GK_INDEX) != nullptr) { + auto gamaKFormat = context_->GetOptionalInputDesc(GK_INDEX)->GetStorageFormat(); + OP_CHECK_IF(gamaKFormat == ge::FORMAT_FRACTAL_NZ, OP_LOGE(context_->GetNodeName(), "gamaK format not support NZ"), + return ge::GRAPH_FAILED); + } + if (context_->GetOptionalInputDesc(ACC_TO_INDEX) != nullptr) { + auto numAcceptedTokensFormat = context_->GetOptionalInputDesc(ACC_TO_INDEX)->GetStorageFormat(); + OP_CHECK_IF(numAcceptedTokensFormat == ge::FORMAT_FRACTAL_NZ, + OP_LOGE(context_->GetNodeName(), "numAcceptedTokens format not support NZ"), return ge::GRAPH_FAILED); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::GetScale() +{ + auto attrs = context_->GetAttrs(); + float scaleValue = *attrs->GetAttrPointer(0); + tilingData_.scale = scaleValue; + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::GetOptionalInput() +{ + if (context_->GetOptionalInputDesc(G_INDEX) == nullptr) { + tilingData_.hasGama = 0; + } else { + tilingData_.hasGama = 1; + } + if (context_->GetOptionalInputDesc(GK_INDEX) == nullptr) { + tilingData_.hasGamaK = 0; + } else { + tilingData_.hasGamaK = 1; + } + if (context_->GetOptionalInputDesc(ACC_TO_INDEX) == nullptr) { + tilingData_.hasAcceptedTokens = 0; + } else { + tilingData_.hasAcceptedTokens = 1; + } + + return ge::GRAPH_SUCCESS; +} + +void RecurrentGatedDeltaRuleTiling::PrintTilingData() +{ + OP_LOGD(context_->GetNodeName(), "vectorCoreNum: [%u]", tilingData_.vectorCoreNum); + OP_LOGD(context_->GetNodeName(), "ubCalSize: [%u]", tilingData_.ubCalSize); + OP_LOGD(context_->GetNodeName(), "ubRestBytes: [%u]", tilingData_.ubRestBytes); + OP_LOGD(context_->GetNodeName(), "t: [%u]", tilingData_.t); + OP_LOGD(context_->GetNodeName(), "nk: [%u]", tilingData_.nk); + OP_LOGD(context_->GetNodeName(), "dk: [%u]", tilingData_.dk); + OP_LOGD(context_->GetNodeName(), "nv: [%u]", tilingData_.nv); + OP_LOGD(context_->GetNodeName(), "dv: [%u]", tilingData_.dv); + OP_LOGD(context_->GetNodeName(), "sBlockNum: [%u]", tilingData_.sBlockNum); + OP_LOGD(context_->GetNodeName(), "b: [%u]", tilingData_.b); + OP_LOGD(context_->GetNodeName(), "vStep: [%u]", tilingData_.vStep); + OP_LOGD(context_->GetNodeName(), "stateOutBufferNum: [%u]", tilingData_.stateOutBufferNum); + OP_LOGD(context_->GetNodeName(), "attnOutBufferNum: [%u]", tilingData_.attnOutBufferNum); + OP_LOGD(context_->GetNodeName(), "scale: [%f]", tilingData_.scale); + OP_LOGD(context_->GetNodeName(), "hasGama: [%u]", tilingData_.hasGama); + OP_LOGD(context_->GetNodeName(), "hasGamaK: [%u]", tilingData_.hasGamaK); + OP_LOGD(context_->GetNodeName(), "hasAcceptedTokens: [%u]", tilingData_.hasAcceptedTokens); +} + +int64_t RecurrentGatedDeltaRuleTiling::CalcFixedUbBytes(int64_t aNv, int64_t aDv, int64_t aDk) const +{ + int64_t usedUbBytes = MAX_MTP * (4 * aDk + 2 * aDv); // 4 for qInQueue_ & kInQueue_, 2 for vInQueue_ + usedUbBytes += 128; // reserve 128 Bytes + if (tilingData_.hasGamaK) { + usedUbBytes += MAX_MTP * 4 * aDk; // 4 for gk gamaInQueue_ + } + if (tilingData_.hasGama) { + usedUbBytes += MAX_MTP * 4 * aNv; // 4 for g gamaInQueue_ + } + usedUbBytes += MAX_MTP * 2 * aNv; // 2 for betaInQueue_ + return usedUbBytes; +} + +int64_t RecurrentGatedDeltaRuleTiling::CalcWorkingUbBytes(int64_t aNv, int64_t aDv, int64_t aDk) const +{ + int64_t usedUbBytes = CalcFixedUbBytes(aNv, aDv, aDk); + usedUbBytes += MAX_MTP * (8 * aDk + 4 * aDv + 4 * aNv); // 8 for qk in ub, 4 for v in ub, 4 for beta in ub + return usedUbBytes; +} + +int64_t RecurrentGatedDeltaRuleTiling::CalcVStepCoeff(int64_t aDk, uint32_t stateOutBufferNum, + uint32_t attnOutBufferNum) const +{ + auto stateDtype = context_->GetInputDesc(STATE_INDEX)->GetDataType(); + int64_t stateDtypeSize = (stateDtype == ge::DT_FLOAT) ? 4 : 2; + int64_t coeff = (stateDtypeSize + static_cast(stateDtypeSize * stateOutBufferNum)) * aDk + + static_cast(4 * attnOutBufferNum); // stateIn/stateOut/attnOut queues + coeff += (4 + 4) * aDk + 4 + 4; // qInUb/kInUb/vInUb/deltaInUb/attnInUb + return coeff; +} + +bool RecurrentGatedDeltaRuleTiling::EvaluateBufferProfile(int64_t ubSize, int64_t usedUbBytes, int64_t aDk, + uint32_t stateOutBufferNum, uint32_t attnOutBufferNum, + BufferProfile &profile) const +{ + int64_t coeff = CalcVStepCoeff(aDk, stateOutBufferNum, attnOutBufferNum); + int64_t vStep = (ubSize - usedUbBytes) / coeff / 8 * 8; // 8 * sizeof(float) = 32 + if (vStep < 8) { + return false; + } + int64_t repeatTime = CeilDiv(tilingData_.dv, static_cast(vStep)); + vStep = CeilAlign(CeilDiv(tilingData_.dv, static_cast(repeatTime)), + static_cast(8)); + if (vStep < 8) { + return false; + } + profile.stateOutBufferNum = stateOutBufferNum; + profile.attnOutBufferNum = attnOutBufferNum; + profile.vStep = static_cast(vStep); + profile.repeatTime = static_cast(repeatTime); + profile.valid = true; + return true; +} + +bool RecurrentGatedDeltaRuleTiling::IsBetterProfile(const BufferProfile &candidate, const BufferProfile ¤t) const +{ + if (!current.valid) { + return true; + } + if (candidate.repeatTime != current.repeatTime) { + return candidate.repeatTime < current.repeatTime; + } + uint32_t candidateDepth = candidate.stateOutBufferNum + candidate.attnOutBufferNum; + uint32_t currentDepth = current.stateOutBufferNum + current.attnOutBufferNum; + if (candidateDepth != currentDepth) { + return candidateDepth > currentDepth; + } + return candidate.vStep > current.vStep; +} + +ge::graphStatus RecurrentGatedDeltaRuleTiling::FinalizeVStepFromUb(int64_t ubSize, int64_t usedUbBytes, int64_t coeff) +{ + (void)coeff; + int64_t aDk = CeilAlign(tilingData_.dk, static_cast(16)); // 16 * 2 = 32B + BufferProfile selected; + const std::array candidates = {{ + BufferProfile(1u, 1u, 0u, 0u, false), + BufferProfile(1u, 2u, 0u, 0u, false), + BufferProfile(2u, 2u, 0u, 0u, false), + BufferProfile(3u, 3u, 0u, 0u, false) + }}; + for (const auto &candidate : candidates) { + BufferProfile profile; + if (!EvaluateBufferProfile(ubSize, usedUbBytes, aDk, candidate.stateOutBufferNum, candidate.attnOutBufferNum, + profile)) { + continue; + } + if (IsBetterProfile(profile, selected)) { + selected = profile; + } + } + + OP_LOGD(context_->GetNodeName(), "selected profile: stateOutBufferNum=[%u], attnOutBufferNum=[%u], vStep=[%u], repeatTime=[%u], valid=[%d]", + selected.stateOutBufferNum, selected.attnOutBufferNum, selected.vStep, selected.repeatTime, selected.valid); + + if (!selected.valid) { + OP_LOGE(context_->GetNodeName(), "vStep should be bigger than 8, shape is too big"); + return ge::GRAPH_FAILED; + } + auto stateDtype = context_->GetInputDesc(STATE_INDEX)->GetDataType(); + + int64_t stateDtypeSize = (stateDtype == ge::DT_FLOAT) ? 4 : 2; + + int64_t queueCoeff = (stateDtypeSize + static_cast(stateDtypeSize * selected.stateOutBufferNum)) * aDk + + static_cast(4 * selected.attnOutBufferNum); + int64_t ubRestBytes = ubSize - ubCalcCtx_.fixedUbBytes - queueCoeff * static_cast(selected.vStep); + if (ubRestBytes < 0) { + OP_LOGE(context_->GetNodeName(), "ubRestBytes should be non-negative, but got %ld", ubRestBytes); + return ge::GRAPH_FAILED; + } + tilingData_.ubCalSize = compileInfo_.ubSize; + tilingData_.vStep = selected.vStep; + tilingData_.stateOutBufferNum = selected.stateOutBufferNum; + tilingData_.attnOutBufferNum = selected.attnOutBufferNum; + tilingData_.ubRestBytes = static_cast(ubRestBytes); + return ge::GRAPH_SUCCESS; +} + +// CalUbSize now runs an ordered UB rule-chain with explicit intermediate states. +ge::graphStatus RecurrentGatedDeltaRuleTiling::CalUbSize() +{ + struct RuleItem { + const char *name; + HostRuleFn fn; + }; + const std::array ubRules = {{ + {"RuleInitUbCalcContext", &RecurrentGatedDeltaRuleTiling::RuleInitUbCalcContext}, + {"RuleCalcFixedUbBytes", &RecurrentGatedDeltaRuleTiling::RuleCalcFixedUbBytes}, + {"RuleCalcWorkingUbBytes", &RecurrentGatedDeltaRuleTiling::RuleCalcWorkingUbBytes}, + {"RuleCalcVStepCoeff", &RecurrentGatedDeltaRuleTiling::RuleCalcVStepCoeff}, + {"RuleFinalizeVStepFromUb", &RecurrentGatedDeltaRuleTiling::RuleFinalizeVStepFromUb}, + }}; + for (const auto &rule : ubRules) { + OP_CHECK_IF((this->*(rule.fn))() != ge::GRAPH_SUCCESS, + OP_LOGE(inputParams_.opName, "CalUbSize rule failed: %s", rule.name), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus RecurrentGatedDeltaRuleTilingFunc(gert::TilingContext *context) +{ + OP_CHECK_IF(context == nullptr, OPS_REPORT_CUBE_INNER_ERR("RecurrentGatedDeltaRule", "context is null"), + return ge::GRAPH_FAILED); + return Ops::Transformer::OpTiling::TilingRegistry::GetInstance().DoTilingImpl(context); +} + +static ge::graphStatus TilingPrepareForRecurrentGatedDeltaRule(gert::TilingParseContext *context) +{ + OP_CHECK_IF(context == nullptr, OPS_REPORT_CUBE_INNER_ERR("RecurrentGatedDeltaRule", "context is null"), + return ge::GRAPH_FAILED); + + fe::PlatFormInfos *platformInfo = context->GetPlatformInfo(); + OP_CHECK_IF(platformInfo == nullptr, OPS_REPORT_CUBE_INNER_ERR(context->GetNodeName(), "platformInfoPtr is null"), + return ge::GRAPH_FAILED); + + auto compileInfoPtr = context->GetCompiledInfo(); + OP_CHECK_IF(compileInfoPtr == nullptr, OPS_REPORT_CUBE_INNER_ERR(context->GetNodeName(), "compileInfoPtr is null"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(RecurrentGatedDeltaRule) + .Tiling(RecurrentGatedDeltaRuleTilingFunc) + .TilingParse(TilingPrepareForRecurrentGatedDeltaRule); +} // namespace optiling + diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule_tiling.h b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule_tiling.h new file mode 100644 index 000000000..d502231ac --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule_tiling.h @@ -0,0 +1,137 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file recurrent_gated_delta_rule_tiling.h + * \brief + */ +#ifndef __OP_HOST_RECURRENT_GETED_DELTA_RULE_TILING_H__ +#define __OP_HOST_RECURRENT_GETED_DELTA_RULE_TILING_H__ +#include +#include "register/tilingdata_base.h" +#include "tiling_base/tiling_base.h" +#include "tiling_base/error_log.h" +#include "../op_kernel/recurrent_gated_delta_rule_tiling_data.h" + +namespace optiling { +using namespace RecurrentGatedDeltaRule; + +struct RecurrentGatedDeltaRuleCompileInfo { + uint64_t aivNum{0UL}; + uint64_t ubSize{0UL}; +}; + +struct RecurrentGatedDeltaRuleInfo { +public: + int64_t usedCoreNum = 0; + const char *opName = "RecurrentGatedDeltaRule"; +}; + +class RecurrentGatedDeltaRuleTiling : public Ops::Transformer::OpTiling::TilingBaseClass { +public: + explicit RecurrentGatedDeltaRuleTiling(gert::TilingContext *context) : Ops::Transformer::OpTiling::TilingBaseClass(context) + { + InitCompileInfo(); + }; + ~RecurrentGatedDeltaRuleTiling() override = default; + +protected: + bool IsCapable() override + { + return true; + } + // 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小 + ge::graphStatus GetPlatformInfo() override; + // 2、获取INPUT/OUTPUT/ATTR信息 + ge::graphStatus GetShapeAttrsInfo() override; + // 3、计算数据切分TilingData + ge::graphStatus DoOpTiling() override; + // 4、计算高阶API的TilingData + ge::graphStatus DoLibApiTiling() override; + // 5、计算TilingKey + uint64_t GetTilingKey() const override; + // 6、计算Workspace 大小 + ge::graphStatus GetWorkspaceSize() override; + // 7、保存Tiling数据 + ge::graphStatus PostTiling() override; + +protected: + void InitCompileInfo(); + void PrintTilingData(); + + //Host tiling rule-chain engine: compose shape/UB steps by ordered rules. + using HostRuleFn = ge::graphStatus (RecurrentGatedDeltaRuleTiling::*)(); + struct UbCalcContext { + int64_t ubSize = 0; + int64_t aNv = 0; + int64_t aDv = 0; + int64_t aDk = 0; + int64_t fixedUbBytes = 0; + int64_t workingUbBytes = 0; + int64_t coeff = 0; + }; + + struct BufferProfile { + BufferProfile() = default; + BufferProfile(uint32_t s, uint32_t a, uint32_t v, uint32_t r, bool val) + : stateOutBufferNum(s), attnOutBufferNum(a), vStep(v), repeatTime(r), valid(val) {} + + uint32_t stateOutBufferNum = 1; + uint32_t attnOutBufferNum = 1; + uint32_t vStep = 0; + uint32_t repeatTime = 0; + bool valid = false; + }; + + ge::graphStatus CheckContext(); + ge::graphStatus AnalyzeDtype(); + ge::graphStatus AnalyzeShapes(); + ge::graphStatus CalUbSize(); + ge::graphStatus GetScale(); + ge::graphStatus GetOptionalInput(); + ge::graphStatus AnalyzeFormat(); + //Host tiling refactor helpers: split shape validation/fill and UB calculation. + ge::graphStatus CheckShapeDimAndRelation(const gert::Shape &queryShape, const gert::Shape &keyShape, + const gert::Shape &valueShape, const gert::Shape &betaShape, + const gert::Shape &stateShape, const gert::Shape &cuSeqlensShape, + const gert::Shape &ssmStateShape); + void FillTilingShapeData(const gert::Shape &queryShape, const gert::Shape &valueShape, const gert::Shape &stateShape, + const gert::Shape &cuSeqlensShape); + ge::graphStatus CheckShapeValueRangeAndRule(); + void UpdateDynamicBlockDimByTaskUnits(); + int64_t CalcFixedUbBytes(int64_t aNv, int64_t aDv, int64_t aDk) const; + int64_t CalcWorkingUbBytes(int64_t aNv, int64_t aDv, int64_t aDk) const; + int64_t CalcVStepCoeff(int64_t aDk, uint32_t stateOutBufferNum, uint32_t attnOutBufferNum) const; + bool EvaluateBufferProfile(int64_t ubSize, int64_t usedUbBytes, int64_t aDk, uint32_t stateOutBufferNum, + uint32_t attnOutBufferNum, BufferProfile &profile) const; + bool IsBetterProfile(const BufferProfile &candidate, const BufferProfile ¤t) const; + ge::graphStatus FinalizeVStepFromUb(int64_t ubSize, int64_t usedUbBytes, int64_t coeff); + ge::graphStatus RuleCheckShapeDimAndRelation(); + ge::graphStatus RuleFillTilingShapeData(); + ge::graphStatus RuleCheckShapeValueRangeAndRule(); + ge::graphStatus RuleUpdateDynamicBlockDimByTaskUnits(); + ge::graphStatus RuleInitUbCalcContext(); + ge::graphStatus RuleCalcFixedUbBytes(); + ge::graphStatus RuleCalcWorkingUbBytes(); + ge::graphStatus RuleCalcVStepCoeff(); + ge::graphStatus RuleFinalizeVStepFromUb(); + + bool CheckDimEqual(const gert::Shape a, const int64_t dimA, gert::Shape b, const int64_t dimB, const std::string &nameA, + const std::string &nameB, const std::string &dimDesc); + bool CheckDim(const gert::Shape shape, const size_t dim, const std::string &dimDesc); + bool CheckFormat(ge::Format format, const std::string &Desc); + + RecurrentGatedDeltaRuleCompileInfo compileInfo_; + RecurrentGatedDeltaRuleTilingData tilingData_; + RecurrentGatedDeltaRuleInfo inputParams_; + UbCalcContext ubCalcCtx_; +}; + +} // namespace optiling +#endif // __OP_HOST_RECURRENT_GETED_DELTA_RULE_TILING_H__ diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_kernel/arch35/recurrent_gated_delta_rule.h b/csrc/ascend/attention/recurrent_gated_delta_rule/op_kernel/arch35/recurrent_gated_delta_rule.h new file mode 100644 index 000000000..fd62841a3 --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_kernel/arch35/recurrent_gated_delta_rule.h @@ -0,0 +1,633 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +/*! + * \file recurrent_gated_delta_rule.h + * \brief + */ + +#ifndef __RECURRENT_GATED_DELTA_RULE_KERNEL_H_ +#define __RECURRENT_GATED_DELTA_RULE_KERNEL_H_ + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" +#include "../recurrent_gated_delta_rule_tiling_data.h" + +namespace RecurrentGatedDeltaRule { + +using namespace matmul; +using namespace AscendC; +using namespace AscendC::MicroAPI; +constexpr uint64_t BUFFER_NUM = 1; +constexpr uint32_t MAX_OUT_BUFFER_NUM = 2; +constexpr uint64_t MAX_MTP = 16; +constexpr uint64_t BF16_NUM_PER_BLOCK = 16; +constexpr uint64_t FP32_NUM_PER_BLOCK = 8; +constexpr uint32_t REPEAT_LENTH = 64; // 256Byte for float +constexpr uint32_t MAX_REPEAT_TIME = 255; +constexpr uint32_t ADD_FOLD_REDUCE_MIN_K = 128; +constexpr uint16_t V_LENGTH = VECTOR_REG_WIDTH / sizeof(float); +constexpr uint16_t TWO_V_LENGTH = 2 * V_LENGTH; + +constexpr CastTrait castTraitB16ToB32 = { + RegLayout::ZERO, SatMode::UNKNOWN, MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + +#ifndef RGDR_ENABLE_ADD_FOLD_REDUCE +#define RGDR_ENABLE_ADD_FOLD_REDUCE 1 +#endif +struct RGDRInitParams { + GM_ADDR query; + GM_ADDR key; + GM_ADDR value; + GM_ADDR gama; + GM_ADDR gamaK; + GM_ADDR beta; + GM_ADDR initState; + GM_ADDR cuSeqlens; + GM_ADDR ssmStateIndices; + GM_ADDR numAcceptedTokens; + GM_ADDR attnOut; + GM_ADDR finalState; +}; + +template +class RGDR { +public: + __aicore__ inline RGDR(const RecurrentGatedDeltaRuleTilingData *tilingData) + { + B_ = tilingData->b; + T_ = tilingData->t; + NK_ = tilingData->nk; + realK_ = tilingData->dk; + NV_ = tilingData->nv; + realV_ = tilingData->dv; + scale_ = tilingData->scale; + hasAcceptedTokens_ = (tilingData->hasAcceptedTokens == 1); + hasGama_ = (tilingData->hasGama == 1); + hasGamaK_ = (tilingData->hasGamaK == 1); + useAddFoldReduce_ = (RGDR_ENABLE_ADD_FOLD_REDUCE != 0); + vStep_ = tilingData->vStep; + stateOutBufferNum_ = (tilingData->stateOutBufferNum == MAX_OUT_BUFFER_NUM) ? MAX_OUT_BUFFER_NUM : BUFFER_NUM; + attnOutBufferNum_ = (tilingData->attnOutBufferNum == MAX_OUT_BUFFER_NUM) ? MAX_OUT_BUFFER_NUM : BUFFER_NUM; + restUbSize_ = tilingData->ubRestBytes; + alignK_ = Ceil(tilingData->dk, BF16_NUM_PER_BLOCK) * BF16_NUM_PER_BLOCK; + alignV_ = Ceil(tilingData->dv, BF16_NUM_PER_BLOCK) * BF16_NUM_PER_BLOCK; + load = 0; + usedblk = 0; + } + + __aicore__ inline void Init(const RGDRInitParams &initParams, TPipe *pipe) + { + uint64_t blockDim = GetBlockNum(); + blockIdx = GetBlockIdx(); + if (blockIdx >= blockDim) { + return; + } + pipe_ = pipe; + SetGlobalTensors(initParams); + InitLocalBuffers(); + } + + __aicore__ inline void SetGlobalTensors(const RGDRInitParams &initParams) + { + queryGm_.SetGlobalBuffer((__gm__ inType *)initParams.query); + keyGm_.SetGlobalBuffer((__gm__ inType *)initParams.key); + valueGm_.SetGlobalBuffer((__gm__ inType *)initParams.value); + gamaGm_.SetGlobalBuffer((__gm__ float *)initParams.gama); + gamaKGm_.SetGlobalBuffer((__gm__ float *)initParams.gamaK); + betaGm_.SetGlobalBuffer((__gm__ inType *)initParams.beta); + initStateGm_.SetGlobalBuffer((__gm__ stateType *)initParams.initState); + cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)initParams.cuSeqlens); + ssmStateIndicesGm_.SetGlobalBuffer((__gm__ int32_t *)initParams.ssmStateIndices); + numAcceptedTokensGm_.SetGlobalBuffer((__gm__ int32_t *)initParams.numAcceptedTokens); + finalStateGm_.SetGlobalBuffer((__gm__ stateType *)initParams.finalState); + attnOutGm_.SetGlobalBuffer((__gm__ outType *)initParams.attnOut); + } + + __aicore__ inline void InitLocalBuffers() + { + uint32_t cubeSize = alignK_ * vStep_ * sizeof(float); + uint32_t singleVSize = vStep_ * sizeof(float); + uint32_t vSize = MAX_MTP * alignV_ * sizeof(float); + uint32_t kSize = MAX_MTP * alignK_ * sizeof(float); + uint32_t betaNumAlign = Ceil(MAX_MTP * NV_, BF16_NUM_PER_BLOCK) * BF16_NUM_PER_BLOCK; + pipe_->InitBuffer(qInQueue_, BUFFER_NUM, MAX_MTP * alignK_ * sizeof(inType)); + pipe_->InitBuffer(kInQueue_, BUFFER_NUM, MAX_MTP * alignK_ * sizeof(inType)); + pipe_->InitBuffer(vInQueue_, BUFFER_NUM, MAX_MTP * alignV_ * sizeof(inType)); + pipe_->InitBuffer(stateInQueue_, BUFFER_NUM, alignK_ * vStep_ * sizeof(stateType)); + if (hasGama_) { + pipe_->InitBuffer(gamaInQueue_, BUFFER_NUM, MAX_MTP * NV_ * sizeof(float)); + } + if (hasGamaK_) { + pipe_->InitBuffer(gamaKInQueue_, BUFFER_NUM, MAX_MTP * alignK_ * sizeof(float)); + } + pipe_->InitBuffer(betaInQueue_, BUFFER_NUM, MAX_MTP * NV_ * sizeof(inType)); + pipe_->InitBuffer(stateOutQueue_, stateOutBufferNum_, alignK_ * vStep_ * sizeof(stateType)); + pipe_->InitBuffer(attnOutQueue_, attnOutBufferNum_, vStep_ * sizeof(outType)); + pipe_->InitBuffer(tmpBuff, restUbSize_); + uint32_t buffOffset = 0; + deltaInUb = tmpBuff.GetWithOffset(static_cast(vStep_), buffOffset); + buffOffset += singleVSize; + attnInUb = tmpBuff.GetWithOffset(static_cast(vStep_), buffOffset); + buffOffset += singleVSize; + vInUb = tmpBuff.GetWithOffset(static_cast(MAX_MTP * alignV_), buffOffset); + buffOffset += vSize; + qInUb = tmpBuff.GetWithOffset(static_cast(MAX_MTP * alignK_), buffOffset); + buffOffset += kSize; + kInUb = tmpBuff.GetWithOffset(static_cast(MAX_MTP * alignK_), buffOffset); + buffOffset += kSize; + stateInUb = tmpBuff.GetWithOffset(static_cast(alignK_ * vStep_), buffOffset); + buffOffset += cubeSize; + broadTmpInUb = tmpBuff.GetWithOffset(static_cast(alignK_ * vStep_), buffOffset); + buffOffset += cubeSize; + betaInUb = tmpBuff.GetWithOffset(static_cast(betaNumAlign), buffOffset); + // gamaInUb is NOT carved from tmpBuff. It reuses the gamaInQueue_ tensor + // directly (see CopyInGamaBeta), matching the generic kernel. Otherwise + // the host-side UB accounting (CalcWorkingUbBytes, which reserves beta + // but not gama in tmpBuff) under-counts by betaNumAlign floats, and the + // shortfall doubles with MAX_MTP -> risk of tmpBuff overflow. + } + + __aicore__ inline void ComputeAvgload() + { + uint64_t realT = 0; + for (uint64_t batch_i = 1; batch_i < B_ + 1; batch_i++) { + realT += cuSeqlensGm_.GetValue(batch_i); + } + avgload = Ceil(realT * NV_, GetBlockNum()); + } + + __aicore__ inline void Process() + { + ComputeAvgload(); + int32_t seq1 = cuSeqlensGm_.GetValue(0); + for (uint64_t batch_i = 0; batch_i < B_; batch_i++) { + int32_t seqLen = cuSeqlensGm_.GetValue(batch_i+1); + if (seqLen <= 0) { + continue; + } + if (seqLen > static_cast(MAX_MTP)) { + return; + } + if (seq1 < 0 || seq1 > static_cast(T_) || (seq1 + seqLen) > static_cast(T_)) { + return; + } + int32_t seq0 = seq1; + seq1 += seqLen; + uint32_t copyFlag = 0; + uint64_t stateOffset; + for (uint64_t head_i = 0; head_i < NV_; head_i++) { + if (!IsCurrentBlock(seq1 - seq0)) { + continue; + } + copyFlag++; + if (copyFlag == 1) { + int32_t stateTokenIdx = seq0; + if (hasAcceptedTokens_) { + int32_t acceptedTokenNum = numAcceptedTokensGm_.GetValue(batch_i); + if (acceptedTokenNum <= 0 || acceptedTokenNum > seqLen) { + return; + } + stateTokenIdx = seq0 + acceptedTokenNum - 1; + } + stateOffset = ssmStateIndicesGm_.GetValue(stateTokenIdx); + CopyInGamaBeta(seq0, seq1); + } + ProcessHead(seq0, seq1, head_i, stateOffset); + } + if (hasGama_ && copyFlag != 0) { + gamaInQueue_.FreeTensor(gamaInUb); + } + } + } + +private: + __aicore__ inline void CopyInQKV(uint64_t vOffset, uint64_t qkOffset, int32_t seqLen) + { + LocalTensor qLocal = qInQueue_.AllocTensor(); + LocalTensor kLocal = kInQueue_.AllocTensor(); + LocalTensor vLocal = vInQueue_.AllocTensor(); + DataCopyExtParams qkInParams{static_cast(seqLen), static_cast(realK_ * sizeof(inType)), + static_cast((NK_ - 1) * realK_ * sizeof(inType)), 0, 0}; + DataCopyExtParams vInParams{static_cast(seqLen), static_cast(realV_ * sizeof(inType)), + static_cast((NV_ - 1) * realV_ * sizeof(inType)), 0, 0}; + DataCopyPadExtParams qkPadParams{true, 0, static_cast(alignK_ - realK_), 0}; + DataCopyPadExtParams vPadParams{true, 0, static_cast(alignV_ - realV_), 0}; + if (hasGamaK_) { + uint32_t alignKGamma = Ceil(realK_, FP32_NUM_PER_BLOCK) * FP32_NUM_PER_BLOCK; + uint32_t stride = alignKGamma < alignK_ ? 1 : 0; + DataCopyExtParams gkInParams{static_cast(seqLen), static_cast(realK_ * sizeof(float)), + static_cast((NV_ - 1) * realK_ * sizeof(float)), stride, 0}; + DataCopyPadExtParams gkPadParams{true, 0, static_cast(alignKGamma - realK_), 0}; + LocalTensor gamaKLocal = gamaKInQueue_.AllocTensor(); + Duplicate(gamaKLocal, 0, alignK_ * seqLen); + TEventID evevtIdVtoMte2 = GetTPipePtr()->FetchEventID(HardEvent::V_MTE2); + SetFlag(evevtIdVtoMte2); + WaitFlag(evevtIdVtoMte2); + DataCopyPad(gamaKLocal, gamaKGm_[vOffset / realV_ * realK_], gkInParams, gkPadParams); + gamaKInQueue_.EnQue(gamaKLocal); + gamaKInUb = gamaKInQueue_.DeQue(); + Exp(gamaKInUb, gamaKInUb, alignK_ * seqLen); + AscendC::PipeBarrier(); + } + DataCopyPad(qLocal, queryGm_[qkOffset], qkInParams, qkPadParams); + DataCopyPad(kLocal, keyGm_[qkOffset], qkInParams, qkPadParams); + DataCopyPad(vLocal, valueGm_[vOffset], vInParams, vPadParams); + qInQueue_.EnQue(qLocal); + kInQueue_.EnQue(kLocal); + vInQueue_.EnQue(vLocal); + qLocal = qInQueue_.DeQue(); + kLocal = kInQueue_.DeQue(); + vLocal = vInQueue_.DeQue(); + Cast(qInUb, qLocal, AscendC::RoundMode::CAST_NONE, alignK_ * seqLen); + Cast(kInUb, kLocal, AscendC::RoundMode::CAST_NONE, alignK_ * seqLen); + Cast(vInUb, vLocal, AscendC::RoundMode::CAST_NONE, alignV_ * seqLen); + AscendC::PipeBarrier(); + Muls(qInUb, qInUb, scale_, seqLen * alignK_); + qInQueue_.FreeTensor(qLocal); + kInQueue_.FreeTensor(kLocal); + vInQueue_.FreeTensor(vLocal); + } + + __aicore__ inline void PrefetchState(uint64_t stateOffest, uint32_t curSingleV) + { + LocalTensor stateLocal = stateInQueue_.AllocTensor(); + DataCopyExtParams stateInParams{static_cast(curSingleV), + static_cast(realK_ * sizeof(stateType)), 0, 0, 0}; + DataCopyPadExtParams padParams{true, 0, static_cast(alignK_ - realK_), 0}; + DataCopyPad(stateLocal, initStateGm_[stateOffest], stateInParams, padParams); + stateInQueue_.EnQue(stateLocal); + } + + __aicore__ inline void LoadPrefetchedState(uint32_t curSingleV) + { + LocalTensor stateLocal = stateInQueue_.DeQue(); + if constexpr (std::is_same()) { + DataCopy(stateInUb, stateLocal, alignK_ * curSingleV); + } else { + Cast(stateInUb, stateLocal, AscendC::RoundMode::CAST_NONE, alignK_ * curSingleV); + } + stateInQueue_.FreeTensor(stateLocal); + } + + __aicore__ inline void MatVecMul(const LocalTensor &cubeTensor, const LocalTensor &vecTensor, + LocalTensor &dstTensor, uint32_t rows) + { + __ubuf__ float* cubeAddr = (__ubuf__ float*)cubeTensor.GetPhyAddr(); + __ubuf__ float* vecAddr = (__ubuf__ float*)vecTensor.GetPhyAddr(); + __ubuf__ float* dstAddr = (__ubuf__ float*)dstTensor.GetPhyAddr(); + + uint16_t rowNum = static_cast(rows); + uint16_t colLoopTimes = static_cast(Ceil(alignK_, V_LENGTH)); + uint32_t colLength = alignK_; + __VEC_SCOPE__ + { + RegTensor cube; + RegTensor vec; + RegTensor dst; + MaskReg pregLoop; + for (uint16_t j = 0; j < colLoopTimes; j++) { + pregLoop = UpdateMask(colLength); + DataCopy(vec, vecAddr + j * V_LENGTH); + for (uint16_t i = 0; i < rowNum; i ++) { + DataCopy(cube, cubeAddr + i * alignK_ + j * V_LENGTH); + Mul(dst, cube, vec, pregLoop); + DataCopy(dstAddr + i * alignK_ + j * V_LENGTH, dst, pregLoop); + } + } + } + } + + __aicore__ inline void ProcessKQ(const LocalTensor &cubeTensor, const LocalTensor &vec1Tensor, + LocalTensor &dst1Tensor, const LocalTensor &vec2Tensor, + LocalTensor &dst2Tensor, uint32_t rows) + { + __ubuf__ float* cubeAddr = (__ubuf__ float*)cubeTensor.GetPhyAddr(); + __ubuf__ float* vec1Addr = (__ubuf__ float*)vec1Tensor.GetPhyAddr(); + __ubuf__ float* vec2Addr = (__ubuf__ float*)vec2Tensor.GetPhyAddr(); + __ubuf__ float* dst1Addr = (__ubuf__ float*)dst1Tensor.GetPhyAddr(); + __ubuf__ float* dst2Addr = (__ubuf__ float*)dst2Tensor.GetPhyAddr(); + + uint16_t rowNum = static_cast(rows); + uint16_t colLoopTimes = static_cast(Ceil(alignK_, V_LENGTH)); + uint32_t colLength = alignK_; + __VEC_SCOPE__ + { + RegTensor cube; + RegTensor vec1; + RegTensor vec2; + RegTensor dst1; + RegTensor dst2; + MaskReg pregLoop; + for (uint16_t j = 0; j < colLoopTimes; j++) { + pregLoop = UpdateMask(colLength); + DataCopy(vec1, vec1Addr + j * V_LENGTH); + DataCopy(vec2, vec2Addr + j * V_LENGTH); + for (uint16_t i = 0; i < rowNum; i ++) { + DataCopy(cube, cubeAddr + i); + DataCopy(dst1, dst1Addr + i * alignK_ + j * V_LENGTH); + Mul(cube, cube, vec1, pregLoop); + Add(dst1, dst1, cube, pregLoop); + Mul(dst2, dst1, vec2, pregLoop); + DataCopy(dst1Addr + i * alignK_ + j * V_LENGTH, dst1, pregLoop); + DataCopy(dst2Addr + i * alignK_ + j * V_LENGTH, dst2, pregLoop); + } + } + } + } + + __aicore__ inline void ReduceSum64(__ubuf__ float* dstAddr, __ubuf__ float* srcAddr, uint16_t rowNum) + { + uint32_t colLength = alignK_; + __VEC_SCOPE__ + { + RegTensor src; + RegTensor sum; + MaskReg pregLoop = UpdateMask(colLength); + for (uint16_t i = 0;i < rowNum;i ++) { + DataCopy(src, srcAddr + i * alignK_); + ReduceSum(sum, src, pregLoop); + DataCopy(dstAddr + i, sum, pregLoop); + } + } + } + + __aicore__ inline void ReduceSum128(__ubuf__ float* dstAddr, __ubuf__ float* srcAddr, uint16_t rowNum) + { + uint32_t colLength = alignK_ - V_LENGTH; + __VEC_SCOPE__ + { + RegTensor src1; + RegTensor src2; + RegTensor sum; + MaskReg pregFull = CreateMask(); + MaskReg pregLoop = UpdateMask(colLength); + for (uint16_t i = 0;i < rowNum;i ++) { + DataCopy(src1, srcAddr + i * alignK_); + DataCopy(src2, srcAddr + i * alignK_ + V_LENGTH); + Add(src1, src1, src2, pregLoop); + ReduceSum(sum, src1, pregFull); + DataCopy(dstAddr + i, sum, pregFull); + } + } + } + + __aicore__ inline void ReduceSumVF(__ubuf__ float* dstAddr, __ubuf__ float* srcAddr, uint16_t rowNum) + { + uint16_t colLoopTimes = static_cast(Ceil(alignK_, V_LENGTH)); + __VEC_SCOPE__ + { + RegTensor src; + RegTensor tmp; + RegTensor sum; + MaskReg pregFull = CreateMask(); + MaskReg pregLoop; + for (uint16_t i = 0;i < rowNum;i ++) { + uint32_t colLength = alignK_; + Duplicate(tmp, 0.0f); + for (uint16_t j = 0; j < colLoopTimes; j++) { + pregLoop = UpdateMask(colLength); + DataCopy(src, srcAddr + i * alignK_ + j * V_LENGTH); + Add(tmp, tmp, src, pregLoop); + } + ReduceSum(sum, tmp, pregFull); + DataCopy(dstAddr + i, sum, pregFull); + } + } + } + + __aicore__ inline void ReduceSumDispatch(LocalTensor &dstTensor, LocalTensor &srcTensor, + uint32_t rows) + { + __ubuf__ float* srcAddr = (__ubuf__ float*)srcTensor.GetPhyAddr(); + __ubuf__ float* dstAddr = (__ubuf__ float*)dstTensor.GetPhyAddr(); + uint16_t rowNum = static_cast(rows); + if (alignK_ <= V_LENGTH) { + ReduceSum64(dstAddr, srcAddr, rowNum); + } else if (alignK_ <= TWO_V_LENGTH) { + ReduceSum128(dstAddr, srcAddr, rowNum); + } else { + ReduceSumVF(dstAddr, srcAddr, rowNum); + } + } + + __aicore__ inline void Compute(uint32_t curSingleV, uint64_t curQKOffset, uint64_t curVOffset) + { + if (hasGama_) { + Muls(stateInUb, stateInUb, gama_, alignK_ * curSingleV); + } + if (hasGamaK_) { + MatVecMul(stateInUb, gamaKInUb[curQKOffset], stateInUb, curSingleV); + } + if (hasGama_ || hasGamaK_) { + AscendC::PipeBarrier(); + } + MatVecMul(stateInUb, kInUb[curQKOffset], broadTmpInUb, curSingleV); + AscendC::PipeBarrier(); + ReduceSumDispatch(deltaInUb, broadTmpInUb, curSingleV); + AscendC::PipeBarrier(); + Sub(deltaInUb, vInUb[curVOffset], deltaInUb, curSingleV); + AscendC::PipeBarrier(); + Muls(deltaInUb, deltaInUb, beta_, curSingleV); + AscendC::PipeBarrier(); + ProcessKQ(deltaInUb, kInUb[curQKOffset], stateInUb, qInUb[curQKOffset], broadTmpInUb, curSingleV); + AscendC::PipeBarrier(); + ReduceSumDispatch(attnInUb, broadTmpInUb, curSingleV); + LocalTensor stateOutLocal = stateOutQueue_.AllocTensor(); + LocalTensor attnOutLocal = attnOutQueue_.AllocTensor(); + if constexpr (std::is_same()) { + DataCopy(stateOutLocal, stateInUb, alignK_ * curSingleV); + } else { + Cast(stateOutLocal, stateInUb, AscendC::RoundMode::CAST_RINT, alignK_ * curSingleV); + } + stateOutQueue_.EnQue(stateOutLocal); + Cast(attnOutLocal, attnInUb, AscendC::RoundMode::CAST_RINT, curSingleV); + attnOutQueue_.EnQue(attnOutLocal); + } + + __aicore__ inline void CopyOutAttn(uint64_t attnOffset, uint32_t curSingleV) + { + LocalTensor attnLocal = attnOutQueue_.DeQue(); + DataCopyParams attnOutParams{1, static_cast(curSingleV * sizeof(outType)), 0, 0}; + DataCopyPad(attnOutGm_[attnOffset], attnLocal, attnOutParams); + attnOutQueue_.FreeTensor(attnLocal); + } + + __aicore__ inline void CopyOutState(uint64_t stateOffset, uint32_t curSingleV) + { + LocalTensor stateOutLocal = stateOutQueue_.DeQue(); + DataCopyParams stateOutParams{static_cast(curSingleV), + static_cast(realK_ * sizeof(stateType)), 0, 0}; + DataCopyPad(finalStateGm_[stateOffset], stateOutLocal, stateOutParams); + stateOutQueue_.FreeTensor(stateOutLocal); + } + + __aicore__ inline void CopyInGamaBeta(int32_t seq0, int32_t seq1) + { + int32_t seqLen = seq1 - seq0; + LocalTensor betaLocal = betaInQueue_.AllocTensor(); + DataCopyParams betaInParams{1, static_cast(seqLen * NV_ * sizeof(inType)), 0, 0}; + DataCopyPadParams padParams; + DataCopyPad(betaLocal, betaGm_[seq0 * NV_], betaInParams, padParams); + betaInQueue_.EnQue(betaLocal); + betaLocal = betaInQueue_.DeQue(); + Cast(betaInUb, betaLocal, AscendC::RoundMode::CAST_NONE, seqLen * NV_); + betaInQueue_.FreeTensor(betaLocal); + if (hasGama_) { + LocalTensor gamaLocal = gamaInQueue_.AllocTensor(); + DataCopyParams gamaInParams{1, static_cast(seqLen * NV_ * sizeof(float)), 0, 0}; + DataCopyPad(gamaLocal, gamaGm_[seq0 * NV_], gamaInParams, padParams); + gamaInQueue_.EnQue(gamaLocal); + gamaInUb = gamaInQueue_.DeQue(); + Exp(gamaInUb, gamaInUb, seqLen * NV_); + AscendC::PipeBarrier(); + // gamaInUb (the queue tensor) stays live until the batch-boundary + // FreeTensor in Process(), mirroring the generic kernel. + } + } + + __aicore__ inline void ProcessHead(int32_t seq0, int32_t seq1, uint64_t head_i, uint64_t stateOffset) + { + uint64_t vOffset = (seq0 * NV_ + head_i) * realV_; + uint64_t qkOffset = (seq0 * NK_ + head_i / (NV_ / NK_)) * realK_; + CopyInQKV(vOffset, qkOffset, seq1 - seq0); + if (realV_ == 0) { + if (hasGamaK_) { + gamaKInQueue_.FreeTensor(gamaKInUb); + } + return; + } + uint64_t nextVOffset = 0; + uint32_t nextSingleV = realV_ > vStep_ ? vStep_ : realV_; + uint64_t nextStateOffset = ((stateOffset * NV_ + head_i) * realV_) * realK_; + PrefetchState(nextStateOffset, nextSingleV); + for (uint64_t v_i = 0; v_i < realV_; v_i += vStep_) { + uint32_t curSingleV = v_i + vStep_ > realV_ ? realV_ - v_i : vStep_; + LoadPrefetchedState(curSingleV); + nextVOffset = v_i + vStep_; + if (nextVOffset < realV_) { + nextSingleV = nextVOffset + vStep_ > realV_ ? realV_ - nextVOffset : vStep_; + nextStateOffset = ((stateOffset * NV_ + head_i) * realV_ + nextVOffset) * realK_; + PrefetchState(nextStateOffset, nextSingleV); + } + uint64_t pendingAttnOffset = 0; + uint64_t pendingStateOffset = 0; + bool hasPendingAttn = false; + bool hasPendingState = false; + for (uint64_t seq_i = seq0; seq_i < seq1; seq_i++) { + uint64_t gbOffset = head_i + (seq_i - seq0) * NV_; + uint64_t curQKOffset = (seq_i - seq0) * alignK_; + uint64_t curVOffset = (seq_i - seq0) * alignV_ + v_i; + uint64_t attnOffset = (seq_i * NV_ + head_i) * realV_ + v_i; + uint64_t curStateOutOffset = + ((ssmStateIndicesGm_.GetValue(seq_i) * NV_ + head_i) * realV_ + v_i) * realK_; + gama_ = hasGama_ ? gamaInUb.GetValue(gbOffset) : 1; + beta_ = betaInUb.GetValue(gbOffset); + Compute(curSingleV, curQKOffset, curVOffset); + if (attnOutBufferNum_ == BUFFER_NUM) { + CopyOutAttn(attnOffset, curSingleV); + } else { + if (hasPendingAttn) { + CopyOutAttn(pendingAttnOffset, curSingleV); + } + pendingAttnOffset = attnOffset; + hasPendingAttn = true; + } + if (stateOutBufferNum_ == BUFFER_NUM) { + CopyOutState(curStateOutOffset, curSingleV); + } else { + if (hasPendingState) { + CopyOutState(pendingStateOffset, curSingleV); + } + pendingStateOffset = curStateOutOffset; + hasPendingState = true; + } + } + if (hasPendingAttn) { + CopyOutAttn(pendingAttnOffset, curSingleV); + } + if (hasPendingState) { + CopyOutState(pendingStateOffset, curSingleV); + } + } + if (hasGamaK_) { + gamaKInQueue_.FreeTensor(gamaKInUb); + } + } + + __aicore__ inline bool IsCurrentBlock(int32_t seqlen) + { + load += seqlen; + bool ret = (blockIdx == usedblk && seqlen > 0); + if (load >= avgload) { + load = 0; + usedblk++; + } + return ret; + } + +private: + GlobalTensor queryGm_; + GlobalTensor keyGm_; + GlobalTensor valueGm_; + GlobalTensor betaGm_; + GlobalTensor gamaGm_; + GlobalTensor gamaKGm_; + GlobalTensor initStateGm_; + GlobalTensor cuSeqlensGm_; + GlobalTensor ssmStateIndicesGm_; + GlobalTensor numAcceptedTokensGm_; + GlobalTensor finalStateGm_; + GlobalTensor attnOutGm_; + TPipe *pipe_; + TQue qInQueue_; + TQue kInQueue_; + TQue vInQueue_; + TQue gamaInQueue_; + TQue gamaKInQueue_; + TQue betaInQueue_; + TQue stateInQueue_; + TQue attnOutQueue_; + TQue stateOutQueue_; + TBuf tmpBuff; + LocalTensor qInUb; + LocalTensor kInUb; + LocalTensor vInUb; + LocalTensor gamaInUb; + LocalTensor gamaKInUb; + LocalTensor betaInUb; + LocalTensor deltaInUb; + LocalTensor broadTmpInUb; + LocalTensor attnInUb; + LocalTensor stateInUb; + uint32_t B_; + uint32_t T_; + uint32_t NK_; + uint32_t alignK_; + uint32_t realK_; + uint32_t NV_; + uint32_t alignV_; + uint32_t realV_; + uint32_t vStep_; + uint32_t stateOutBufferNum_; + uint32_t attnOutBufferNum_; + uint32_t restUbSize_; + uint32_t load; + uint32_t usedblk; + uint32_t avgload; + bool hasAcceptedTokens_; + bool hasGama_; + bool hasGamaK_; + bool useAddFoldReduce_; + float gama_; + float beta_; + float scale_; + uint64_t blockIdx; +}; +} // namespace RecurrentGatedDeltaRule +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_kernel/recurrent_gated_delta_rule.cpp b/csrc/ascend/attention/recurrent_gated_delta_rule/op_kernel/recurrent_gated_delta_rule.cpp new file mode 100644 index 000000000..e1464314c --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_kernel/recurrent_gated_delta_rule.cpp @@ -0,0 +1,41 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file recurrent_gated_delta_rule.cpp + * \brief + */ +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#include "arch35/recurrent_gated_delta_rule.h" +#else +#include "recurrent_gated_delta_rule.h" +#endif +#include "recurrent_gated_delta_rule_tiling_data.h" + + +using namespace AscendC; +using namespace matmul; +using namespace RecurrentGatedDeltaRule; + + +extern "C" __global__ __aicore__ void +recurrent_gated_delta_rule(GM_ADDR query, GM_ADDR key, GM_ADDR value, GM_ADDR beta, GM_ADDR state, GM_ADDR cuSeqlens, + GM_ADDR ssmStateIndices, GM_ADDR g, GM_ADDR gk, GM_ADDR numAcceptedTokens, GM_ADDR out, + GM_ADDR stateOut, GM_ADDR workspaceGM, GM_ADDR tilingGM) +{ + REGISTER_TILING_DEFAULT(RecurrentGatedDeltaRuleTilingData); + GET_TILING_DATA(tilingData, tilingGM); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); + TPipe pipe; + RGDR op(&tilingData); + RGDRInitParams initParams{query, key, value, g, gk, beta, state, cuSeqlens, + ssmStateIndices, numAcceptedTokens, out, stateOut}; + op.Init(initParams, &pipe); + op.Process(); +} diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_kernel/recurrent_gated_delta_rule.h b/csrc/ascend/attention/recurrent_gated_delta_rule/op_kernel/recurrent_gated_delta_rule.h new file mode 100644 index 000000000..f60b9dec8 --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_kernel/recurrent_gated_delta_rule.h @@ -0,0 +1,581 @@ +/** +?* Copyright (c) 2025 Huawei Technologies Co., Ltd. +?* This program is free software, you can redistribute it and/or modify it under the terms and conditions of +?* Please refer to the License for details. You may not use this file except in compliance with the License. +?* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +?* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +?* See LICENSE in the root of the software repository for the full text of the License. +?*/ + +/*! + * \file grouped_matmul_finalize_routing.h + * \brief + */ + +#ifndef __RECURRENT_GATED_DELTA_RULE_KERNEL_H_ +#define __RECURRENT_GATED_DELTA_RULE_KERNEL_H_ + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" +#include "recurrent_gated_delta_rule_tiling_data.h" + +namespace RecurrentGatedDeltaRule { + +using namespace matmul; +using namespace AscendC; +constexpr uint64_t BUFFER_NUM = 1; +constexpr uint32_t MAX_OUT_BUFFER_NUM = 2; +constexpr uint64_t MAX_MTP = 16; +constexpr uint64_t BF16_NUM_PER_BLOCK = 16; +constexpr uint64_t FP32_NUM_PER_BLOCK = 8; +constexpr uint32_t REPEAT_LENTH = 64; // 256Byte for float +constexpr uint32_t MAX_REPEAT_TIME = 255; +constexpr uint32_t ADD_FOLD_REDUCE_MIN_K = 128; + +#ifndef RGDR_ENABLE_ADD_FOLD_REDUCE +#define RGDR_ENABLE_ADD_FOLD_REDUCE 1 +#endif +struct RGDRInitParams { + GM_ADDR query; + GM_ADDR key; + GM_ADDR value; + GM_ADDR gama; + GM_ADDR gamaK; + GM_ADDR beta; + GM_ADDR initState; + GM_ADDR cuSeqlens; + GM_ADDR ssmStateIndices; + GM_ADDR numAcceptedTokens; + GM_ADDR attnOut; + GM_ADDR finalState; +}; + +template +class RGDR { +public: + __aicore__ inline RGDR(const RecurrentGatedDeltaRuleTilingData *tilingData) + { + B_ = tilingData->b; + T_ = tilingData->t; + NK_ = tilingData->nk; + realK_ = tilingData->dk; + NV_ = tilingData->nv; + realV_ = tilingData->dv; + scale_ = tilingData->scale; + hasAcceptedTokens_ = (tilingData->hasAcceptedTokens == 1); + hasGama_ = (tilingData->hasGama == 1); + hasGamaK_ = (tilingData->hasGamaK == 1); + useAddFoldReduce_ = (RGDR_ENABLE_ADD_FOLD_REDUCE != 0); + vStep_ = tilingData->vStep; + stateOutBufferNum_ = (tilingData->stateOutBufferNum == MAX_OUT_BUFFER_NUM) ? MAX_OUT_BUFFER_NUM : BUFFER_NUM; + attnOutBufferNum_ = (tilingData->attnOutBufferNum == MAX_OUT_BUFFER_NUM) ? MAX_OUT_BUFFER_NUM : BUFFER_NUM; + restUbSize_ = tilingData->ubRestBytes; + alignK_ = Ceil(tilingData->dk, BF16_NUM_PER_BLOCK) * BF16_NUM_PER_BLOCK; + alignV_ = Ceil(tilingData->dv, BF16_NUM_PER_BLOCK) * BF16_NUM_PER_BLOCK; + load = 0; + usedblk = 0; + } + + __aicore__ inline void Init(const RGDRInitParams &initParams, TPipe *pipe) + { + uint64_t blockDim = GetBlockNum(); + blockIdx = GetBlockIdx(); + if (blockIdx >= blockDim) { + return; + } + pipe_ = pipe; + SetGlobalTensors(initParams); + InitLocalBuffers(); + } + + __aicore__ inline void SetGlobalTensors(const RGDRInitParams &initParams) + { + queryGm_.SetGlobalBuffer((__gm__ inType *)initParams.query); + keyGm_.SetGlobalBuffer((__gm__ inType *)initParams.key); + valueGm_.SetGlobalBuffer((__gm__ inType *)initParams.value); + gamaGm_.SetGlobalBuffer((__gm__ float *)initParams.gama); + gamaKGm_.SetGlobalBuffer((__gm__ float *)initParams.gamaK); + betaGm_.SetGlobalBuffer((__gm__ inType *)initParams.beta); + initStateGm_.SetGlobalBuffer((__gm__ stateType *)initParams.initState); + cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)initParams.cuSeqlens); + ssmStateIndicesGm_.SetGlobalBuffer((__gm__ int32_t *)initParams.ssmStateIndices); + numAcceptedTokensGm_.SetGlobalBuffer((__gm__ int32_t *)initParams.numAcceptedTokens); + finalStateGm_.SetGlobalBuffer((__gm__ stateType *)initParams.finalState); + attnOutGm_.SetGlobalBuffer((__gm__ outType *)initParams.attnOut); + } + + __aicore__ inline void InitLocalBuffers() + { + uint32_t cubeSize = alignK_ * vStep_ * sizeof(float); + uint32_t singleVSize = vStep_ * sizeof(float); + uint32_t vSize = MAX_MTP * alignV_ * sizeof(float); + uint32_t kSize = MAX_MTP * alignK_ * sizeof(float); + uint32_t betaUbSize = + Ceil(MAX_MTP * NV_, BF16_NUM_PER_BLOCK) * BF16_NUM_PER_BLOCK * sizeof(float); // 8: 8 * 4 = 32B; + pipe_->InitBuffer(qInQueue_, BUFFER_NUM, MAX_MTP * alignK_ * sizeof(inType)); + pipe_->InitBuffer(kInQueue_, BUFFER_NUM, MAX_MTP * alignK_ * sizeof(inType)); + pipe_->InitBuffer(vInQueue_, BUFFER_NUM, MAX_MTP * alignV_ * sizeof(inType)); + pipe_->InitBuffer(stateInQueue_, BUFFER_NUM, alignK_ * vStep_ * sizeof(stateType)); + if (hasGama_) { + pipe_->InitBuffer(gamaInQueue_, BUFFER_NUM, MAX_MTP * NV_ * sizeof(float)); + } + if (hasGamaK_) { + pipe_->InitBuffer(gamaKInQueue_, BUFFER_NUM, MAX_MTP * alignK_ * sizeof(float)); + } + pipe_->InitBuffer(betaInQueue_, BUFFER_NUM, MAX_MTP * NV_ * sizeof(inType)); + pipe_->InitBuffer(stateOutQueue_, stateOutBufferNum_, alignK_ * vStep_ * sizeof(stateType)); + pipe_->InitBuffer(attnOutQueue_, attnOutBufferNum_, vStep_ * sizeof(outType)); + pipe_->InitBuffer(tmpBuff, restUbSize_); + uint32_t buffOffset = 0; + deltaInUb = tmpBuff.GetWithOffset(static_cast(vStep_), buffOffset); + buffOffset += singleVSize; + attnInUb = tmpBuff.GetWithOffset(static_cast(vStep_), buffOffset); + buffOffset += singleVSize; + vInUb = tmpBuff.GetWithOffset(static_cast(MAX_MTP * alignV_), buffOffset); + buffOffset += vSize; + qInUb = tmpBuff.GetWithOffset(static_cast(MAX_MTP * alignK_), buffOffset); + buffOffset += kSize; + kInUb = tmpBuff.GetWithOffset(static_cast(MAX_MTP * alignK_), buffOffset); + buffOffset += kSize; + stateInUb = tmpBuff.GetWithOffset(static_cast(alignK_ * vStep_), buffOffset); + buffOffset += cubeSize; + broadTmpInUb = tmpBuff.GetWithOffset(static_cast(alignK_ * vStep_), buffOffset); + buffOffset += cubeSize; + betaInUb = tmpBuff.GetWithOffset(static_cast(betaUbSize), buffOffset); + } + + __aicore__ inline void ComputeAvgload() + { + uint64_t realT = 0; + for (uint64_t batch_i = 1; batch_i < B_ + 1; batch_i++) { + realT += cuSeqlensGm_.GetValue(batch_i); + } + avgload = Ceil(realT * NV_, GetBlockNum()); + } + + __aicore__ inline void Process() + { + ComputeAvgload(); + int32_t seq1 = cuSeqlensGm_.GetValue(0); + for (uint64_t batch_i = 0; batch_i < B_; batch_i++) { + int32_t seqLen = cuSeqlensGm_.GetValue(batch_i+1); + if (seqLen <= 0) { + continue; + } + if (seqLen > static_cast(MAX_MTP)) { + return; + } + if (seq1 < 0 || seq1 > static_cast(T_) || (seq1 + seqLen) > static_cast(T_)) { + return; + } + int32_t seq0 = seq1; + seq1 += seqLen; + uint32_t copyFlag = 0; + uint64_t stateOffset; + for (uint64_t head_i = 0; head_i < NV_; head_i++) { + if (!IsCurrentBlock(seq1 - seq0)) { + continue; + } + copyFlag++; + if (copyFlag == 1) { + int32_t stateTokenIdx = seq0; + if (hasAcceptedTokens_) { + int32_t acceptedTokenNum = numAcceptedTokensGm_.GetValue(batch_i); + if (acceptedTokenNum <= 0 || acceptedTokenNum > seqLen) { + return; + } + stateTokenIdx = seq0 + acceptedTokenNum - 1; + } + stateOffset = ssmStateIndicesGm_.GetValue(stateTokenIdx); + CopyInGamaBeta(seq0, seq1); + } + ProcessHead(seq0, seq1, head_i, stateOffset); + } + if (hasGama_ && copyFlag != 0) { + gamaInQueue_.FreeTensor(gamaInUb); + } + } + } + +private: + __aicore__ inline void CopyInQKV(uint64_t vOffset, uint64_t qkOffset, int32_t seqLen) + { + LocalTensor qLocal = qInQueue_.AllocTensor(); + LocalTensor kLocal = kInQueue_.AllocTensor(); + LocalTensor vLocal = vInQueue_.AllocTensor(); + DataCopyExtParams qkInParams{static_cast(seqLen), static_cast(realK_ * sizeof(inType)), + static_cast((NK_ - 1) * realK_ * sizeof(inType)), 0, 0}; + DataCopyExtParams vInParams{static_cast(seqLen), static_cast(realV_ * sizeof(inType)), + static_cast((NV_ - 1) * realV_ * sizeof(inType)), 0, 0}; + DataCopyPadExtParams qkPadParams{true, 0, static_cast(alignK_ - realK_), 0}; + DataCopyPadExtParams vPadParams{true, 0, static_cast(alignV_ - realV_), 0}; + if (hasGamaK_) { + uint32_t alignKGamma = Ceil(realK_, FP32_NUM_PER_BLOCK) * FP32_NUM_PER_BLOCK; + uint32_t stride = alignKGamma < alignK_ ? 1 : 0; + DataCopyExtParams gkInParams{static_cast(seqLen), static_cast(realK_ * sizeof(float)), + static_cast((NV_ - 1) * realK_ * sizeof(float)), stride, 0}; + DataCopyPadExtParams gkPadParams{true, 0, static_cast(alignKGamma - realK_), 0}; + LocalTensor gamaKLocal = gamaKInQueue_.AllocTensor(); + Duplicate(gamaKLocal, 0, alignK_ * seqLen); + TEventID evevtIdVtoMte2 = GetTPipePtr()->FetchEventID(HardEvent::V_MTE2); + SetFlag(evevtIdVtoMte2); + WaitFlag(evevtIdVtoMte2); + DataCopyPad(gamaKLocal, gamaKGm_[vOffset / realV_ * realK_], gkInParams, gkPadParams); + gamaKInQueue_.EnQue(gamaKLocal); + gamaKInUb = gamaKInQueue_.DeQue(); + Exp(gamaKInUb, gamaKInUb, alignK_ * seqLen); + AscendC::PipeBarrier(); + } + DataCopyPad(qLocal, queryGm_[qkOffset], qkInParams, qkPadParams); + DataCopyPad(kLocal, keyGm_[qkOffset], qkInParams, qkPadParams); + DataCopyPad(vLocal, valueGm_[vOffset], vInParams, vPadParams); + qInQueue_.EnQue(qLocal); + kInQueue_.EnQue(kLocal); + vInQueue_.EnQue(vLocal); + qLocal = qInQueue_.DeQue(); + kLocal = kInQueue_.DeQue(); + vLocal = vInQueue_.DeQue(); + Cast(qInUb, qLocal, AscendC::RoundMode::CAST_NONE, alignK_ * seqLen); + Cast(kInUb, kLocal, AscendC::RoundMode::CAST_NONE, alignK_ * seqLen); + Cast(vInUb, vLocal, AscendC::RoundMode::CAST_NONE, alignV_ * seqLen); + AscendC::PipeBarrier(); + Muls(qInUb, qInUb, scale_, seqLen * alignK_); + qInQueue_.FreeTensor(qLocal); + kInQueue_.FreeTensor(kLocal); + vInQueue_.FreeTensor(vLocal); + } + + __aicore__ inline void PrefetchState(uint64_t stateOffest, uint32_t curSingleV) + { + LocalTensor stateLocal = stateInQueue_.AllocTensor(); + DataCopyExtParams stateInParams{static_cast(curSingleV), + static_cast(realK_ * sizeof(stateType)), 0, 0, 0}; + DataCopyPadExtParams padParams{true, 0, static_cast(alignK_ - realK_), 0}; + DataCopyPad(stateLocal, initStateGm_[stateOffest], stateInParams, padParams); + stateInQueue_.EnQue(stateLocal); + } + + __aicore__ inline void LoadPrefetchedState(uint32_t curSingleV) + { + LocalTensor stateLocal = stateInQueue_.DeQue(); + if constexpr (std::is_same()) { + DataCopy(stateInUb, stateLocal, alignK_ * curSingleV); + } else { + Cast(stateInUb, stateLocal, AscendC::RoundMode::CAST_NONE, alignK_ * curSingleV); + } + stateInQueue_.FreeTensor(stateLocal); + } + + __aicore__ inline void MatVecMul(const LocalTensor &cubeTensor, const LocalTensor &vecTensor, + LocalTensor &dstTensor, uint32_t cols, bool isAdd) + { + uint8_t repeatStride = alignK_ / FP32_NUM_PER_BLOCK; + for (uint32_t i = 0; i < alignK_; i += REPEAT_LENTH) { + uint64_t mask = Std::min(REPEAT_LENTH, alignK_ - i); + for (uint32_t j = 0; j < cols; j += MAX_REPEAT_TIME) { + uint64_t repeatTime = Std::min(MAX_REPEAT_TIME, cols - j); + if (isAdd) { + MulAddDst(dstTensor[j * alignK_ + i], cubeTensor[j * alignK_ + i], vecTensor[i], mask, repeatTime, + {1, 1, 1, repeatStride, repeatStride, 0}); + } else { + Mul(dstTensor[j * alignK_ + i], cubeTensor[j * alignK_ + i], vecTensor[i], mask, repeatTime, + {1, 1, 1, repeatStride, repeatStride, 0}); + } + } + } + } + + __aicore__ inline void ReduceSumBaseline(LocalTensor &dstTensor, const LocalTensor &srcTensor, + uint32_t rows) + { + uint32_t stateShape[2] = {rows, alignK_}; + ReduceSum(dstTensor, srcTensor, stateShape, true); + } + + __aicore__ inline bool CanUseK128AddFoldFastPath(uint32_t rows) const + { + if (alignK_ != ADD_FOLD_REDUCE_MIN_K) { + return false; + } + if (rows == 0 || rows > MAX_REPEAT_TIME) { + return false; + } + return true; + } + + __aicore__ inline void ReduceSumAddFoldK128(LocalTensor &dstTensor, LocalTensor &srcTensor, + uint32_t rows) + { + const uint8_t repeatTime = static_cast(rows); + const uint8_t rowRepStride = static_cast(alignK_ / FP32_NUM_PER_BLOCK); + + // Write the folded result to the upper half to avoid the multi-repeat src0/dst overlap case. + Add(srcTensor[REPEAT_LENTH], srcTensor, srcTensor[REPEAT_LENTH], REPEAT_LENTH, repeatTime, + {1, 1, 1, rowRepStride, rowRepStride, rowRepStride}); + AscendC::PipeBarrier(); + WholeReduceSum(dstTensor, srcTensor[REPEAT_LENTH], REPEAT_LENTH, repeatTime, 1, 1, rowRepStride); + } + + __aicore__ inline void ReduceSumAddFold(LocalTensor &dstTensor, LocalTensor &srcTensor, + uint32_t rows) + { + if (alignK_ < REPEAT_LENTH) { + ReduceSumBaseline(dstTensor, srcTensor, rows); + return; + } + + if ((alignK_ & (alignK_ - 1)) != 0) { + ReduceSumBaseline(dstTensor, srcTensor, rows); + return; + } + + if (CanUseK128AddFoldFastPath(rows)) { + ReduceSumAddFoldK128(dstTensor, srcTensor, rows); + return; + } + + for (uint32_t row = 0; row < rows; ++row) { + uint32_t rowOffset = row * alignK_; + uint32_t activeLen = alignK_; + while (activeLen > REPEAT_LENTH) { + uint32_t half = activeLen >> 1; + Add(srcTensor[rowOffset], srcTensor[rowOffset], srcTensor[rowOffset + half], half); + AscendC::PipeBarrier(); + activeLen = half; + } + + WholeReduceSum(dstTensor[row], srcTensor[rowOffset], REPEAT_LENTH, 1, 1, 1, FP32_NUM_PER_BLOCK); + } + } + + __aicore__ inline void ReduceSumDispatch(LocalTensor &dstTensor, LocalTensor &srcTensor, + uint32_t rows) + { + if (useAddFoldReduce_ && alignK_ >= ADD_FOLD_REDUCE_MIN_K) { + ReduceSumAddFold(dstTensor, srcTensor, rows); + return; + } + ReduceSumBaseline(dstTensor, srcTensor, rows); + } + + __aicore__ inline void Compute(uint32_t curSingleV, uint64_t curQKOffset, uint64_t curVOffset) + { + uint32_t stateShape[2] = {curSingleV, alignK_}; + uint32_t ktShape[2] = {1, alignK_}; + uint32_t deltaShape[2] = {curSingleV, 1}; + if (hasGama_) { + Muls(stateInUb, stateInUb, gama_, alignK_ * curSingleV); + } + if (hasGamaK_) { + MatVecMul(stateInUb, gamaKInUb[curQKOffset], stateInUb, curSingleV, false); + } + if (hasGama_ || hasGamaK_) { + AscendC::PipeBarrier(); + } + MatVecMul(stateInUb, kInUb[curQKOffset], broadTmpInUb, curSingleV, false); + AscendC::PipeBarrier(); + ReduceSumDispatch(deltaInUb, broadTmpInUb, curSingleV); + AscendC::PipeBarrier(); + deltaInUb = vInUb[curVOffset] - deltaInUb; + AscendC::PipeBarrier(); + Muls(deltaInUb, deltaInUb, beta_, curSingleV); + AscendC::PipeBarrier(); + Broadcast(broadTmpInUb, deltaInUb, stateShape, deltaShape); // 2: Dim Number 1: Second Dim + AscendC::PipeBarrier(); + MatVecMul(broadTmpInUb, kInUb[curQKOffset], stateInUb, curSingleV, true); + AscendC::PipeBarrier(); + MatVecMul(stateInUb, qInUb[curQKOffset], broadTmpInUb, curSingleV, false); + AscendC::PipeBarrier(); + ReduceSumDispatch(attnInUb, broadTmpInUb, curSingleV); + LocalTensor stateOutLocal = stateOutQueue_.AllocTensor(); + LocalTensor attnOutLocal = attnOutQueue_.AllocTensor(); + if constexpr (std::is_same()) { + DataCopy(stateOutLocal, stateInUb, alignK_ * curSingleV); + } else { + Cast(stateOutLocal, stateInUb, AscendC::RoundMode::CAST_RINT, alignK_ * curSingleV); + } + stateOutQueue_.EnQue(stateOutLocal); + Cast(attnOutLocal, attnInUb, AscendC::RoundMode::CAST_RINT, curSingleV); + attnOutQueue_.EnQue(attnOutLocal); + } + + __aicore__ inline void CopyOutAttn(uint64_t attnOffset, uint32_t curSingleV) + { + LocalTensor attnLocal = attnOutQueue_.DeQue(); + DataCopyParams attnOutParams{1, static_cast(curSingleV * sizeof(outType)), 0, 0}; + DataCopyPad(attnOutGm_[attnOffset], attnLocal, attnOutParams); + attnOutQueue_.FreeTensor(attnLocal); + } + + __aicore__ inline void CopyOutState(uint64_t stateOffset, uint32_t curSingleV) + { + LocalTensor stateOutLocal = stateOutQueue_.DeQue(); + DataCopyParams stateOutParams{static_cast(curSingleV), + static_cast(realK_ * sizeof(stateType)), 0, 0}; + DataCopyPad(finalStateGm_[stateOffset], stateOutLocal, stateOutParams); + stateOutQueue_.FreeTensor(stateOutLocal); + } + + __aicore__ inline void CopyInGamaBeta(int32_t seq0, int32_t seq1) + { + int32_t seqLen = seq1 - seq0; + uint64_t bBatchSize = Ceil(seqLen * NV_, BF16_NUM_PER_BLOCK) * BF16_NUM_PER_BLOCK; + LocalTensor betaLocal = betaInQueue_.AllocTensor(); + DataCopyParams betaInParams{1, static_cast(seqLen * NV_ * sizeof(inType)), 0, 0}; + DataCopyPadParams padParams; + DataCopyPad(betaLocal, betaGm_[seq0 * NV_], betaInParams, padParams); + betaInQueue_.EnQue(betaLocal); + betaLocal = betaInQueue_.DeQue(); + Cast(betaInUb, betaLocal, AscendC::RoundMode::CAST_NONE, bBatchSize); + betaInQueue_.FreeTensor(betaLocal); + if (hasGama_) { + LocalTensor gamaLocal = gamaInQueue_.AllocTensor(); + DataCopyParams gamaInParams{1, static_cast(seqLen * NV_ * sizeof(float)), 0, 0}; + DataCopyPad(gamaLocal, gamaGm_[seq0 * NV_], gamaInParams, padParams); + gamaInQueue_.EnQue(gamaLocal); + gamaInUb = gamaInQueue_.DeQue(); + Exp(gamaInUb, gamaInUb, seqLen * NV_); + AscendC::PipeBarrier(); + } + } + + __aicore__ inline void ProcessHead(int32_t seq0, int32_t seq1, uint64_t head_i, uint64_t stateOffset) + { + uint64_t vOffset = (seq0 * NV_ + head_i) * realV_; + uint64_t qkOffset = (seq0 * NK_ + head_i / (NV_ / NK_)) * realK_; + CopyInQKV(vOffset, qkOffset, seq1 - seq0); + if (realV_ == 0) { + if (hasGamaK_) { + gamaKInQueue_.FreeTensor(gamaKInUb); + } + return; + } + uint64_t nextVOffset = 0; + uint32_t nextSingleV = realV_ > vStep_ ? vStep_ : realV_; + uint64_t nextStateOffset = ((stateOffset * NV_ + head_i) * realV_) * realK_; + PrefetchState(nextStateOffset, nextSingleV); + for (uint64_t v_i = 0; v_i < realV_; v_i += vStep_) { + uint32_t curSingleV = v_i + vStep_ > realV_ ? realV_ - v_i : vStep_; + LoadPrefetchedState(curSingleV); + nextVOffset = v_i + vStep_; + if (nextVOffset < realV_) { + nextSingleV = nextVOffset + vStep_ > realV_ ? realV_ - nextVOffset : vStep_; + nextStateOffset = ((stateOffset * NV_ + head_i) * realV_ + nextVOffset) * realK_; + PrefetchState(nextStateOffset, nextSingleV); + } + uint64_t pendingAttnOffset = 0; + uint64_t pendingStateOffset = 0; + bool hasPendingAttn = false; + bool hasPendingState = false; + for (uint64_t seq_i = seq0; seq_i < seq1; seq_i++) { + uint64_t gbOffset = head_i + (seq_i - seq0) * NV_; + uint64_t curQKOffset = (seq_i - seq0) * alignK_; + uint64_t curVOffset = (seq_i - seq0) * alignV_ + v_i; + uint64_t attnOffset = (seq_i * NV_ + head_i) * realV_ + v_i; + uint64_t curStateOutOffset = + ((ssmStateIndicesGm_.GetValue(seq_i) * NV_ + head_i) * realV_ + v_i) * realK_; + gama_ = hasGama_ ? gamaInUb.GetValue(gbOffset) : 1; + beta_ = betaInUb.GetValue(gbOffset); + Compute(curSingleV, curQKOffset, curVOffset); + if (attnOutBufferNum_ == BUFFER_NUM) { + CopyOutAttn(attnOffset, curSingleV); + } else { + if (hasPendingAttn) { + CopyOutAttn(pendingAttnOffset, curSingleV); + } + pendingAttnOffset = attnOffset; + hasPendingAttn = true; + } + if (stateOutBufferNum_ == BUFFER_NUM) { + CopyOutState(curStateOutOffset, curSingleV); + } else { + if (hasPendingState) { + CopyOutState(pendingStateOffset, curSingleV); + } + pendingStateOffset = curStateOutOffset; + hasPendingState = true; + } + } + if (hasPendingAttn) { + CopyOutAttn(pendingAttnOffset, curSingleV); + } + if (hasPendingState) { + CopyOutState(pendingStateOffset, curSingleV); + } + } + if (hasGamaK_) { + gamaKInQueue_.FreeTensor(gamaKInUb); + } + } + + __aicore__ inline bool IsCurrentBlock(int32_t seqlen) + { + load += seqlen; + bool ret = (blockIdx == usedblk && seqlen > 0); + if (load >= avgload) { + load = 0; + usedblk++; + } + return ret; + } + +private: + GlobalTensor queryGm_; + GlobalTensor keyGm_; + GlobalTensor valueGm_; + GlobalTensor betaGm_; + GlobalTensor gamaGm_; + GlobalTensor gamaKGm_; + GlobalTensor initStateGm_; + GlobalTensor cuSeqlensGm_; + GlobalTensor ssmStateIndicesGm_; + GlobalTensor numAcceptedTokensGm_; + GlobalTensor finalStateGm_; + GlobalTensor attnOutGm_; + TPipe *pipe_; + TQue qInQueue_; + TQue kInQueue_; + TQue vInQueue_; + TQue gamaInQueue_; + TQue gamaKInQueue_; + TQue betaInQueue_; + TQue stateInQueue_; + TQue attnOutQueue_; + TQue stateOutQueue_; + TBuf tmpBuff; + LocalTensor qInUb; + LocalTensor kInUb; + LocalTensor vInUb; + LocalTensor gamaInUb; + LocalTensor gamaKInUb; + LocalTensor betaInUb; + LocalTensor deltaInUb; + LocalTensor broadTmpInUb; + LocalTensor attnInUb; + LocalTensor stateInUb; + uint32_t B_; + uint32_t T_; + uint32_t NK_; + uint32_t alignK_; + uint32_t realK_; + uint32_t NV_; + uint32_t alignV_; + uint32_t realV_; + uint32_t vStep_; + uint32_t stateOutBufferNum_; + uint32_t attnOutBufferNum_; + uint32_t restUbSize_; + uint32_t load; + uint32_t usedblk; + uint32_t avgload; + bool hasAcceptedTokens_; + bool hasGama_; + bool hasGamaK_; + bool useAddFoldReduce_; + float gama_; + float beta_; + float scale_; + uint64_t blockIdx; +}; +} // namespace RecurrentGatedDeltaRule +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/op_kernel/recurrent_gated_delta_rule_tiling_data.h b/csrc/ascend/attention/recurrent_gated_delta_rule/op_kernel/recurrent_gated_delta_rule_tiling_data.h new file mode 100644 index 000000000..f3205cf4e --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/op_kernel/recurrent_gated_delta_rule_tiling_data.h @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file recurrent_gated_delta_rule.cpp + * \brief + */ +#ifndef RECURRENT_GATED_DELTA_RULE_TILING_DATA_H +#define RECURRENT_GATED_DELTA_RULE_TILING_DATA_H + +#include "kernel_tiling/kernel_tiling.h" + +namespace RecurrentGatedDeltaRule { +#pragma pack(push, 8) +struct alignas(8) RecurrentGatedDeltaRuleTilingData { // alignas(8)确保8字节对齐 + uint32_t vectorCoreNum; + uint32_t ubCalSize; + uint32_t ubRestBytes; + uint32_t t; + uint32_t nk; + uint32_t dk; + uint32_t nv; + uint32_t dv; + uint32_t sBlockNum; + uint32_t b; + uint32_t vStep; + uint32_t stateOutBufferNum; + uint32_t attnOutBufferNum; + float scale; + uint32_t hasGama; + uint32_t hasGamaK; + uint32_t hasAcceptedTokens; +}; +#pragma pack(pop) +} // RecurrentGatedDeltaRule + +#endif // RECURRENT_GATED_DELTA_RULE_TILING_DATA_H diff --git a/csrc/ascend/attention/recurrent_gated_delta_rule/recurrent_gated_delta_rule_torch_adpt.h b/csrc/ascend/attention/recurrent_gated_delta_rule/recurrent_gated_delta_rule_torch_adpt.h new file mode 100644 index 000000000..5d18dd354 --- /dev/null +++ b/csrc/ascend/attention/recurrent_gated_delta_rule/recurrent_gated_delta_rule_torch_adpt.h @@ -0,0 +1,56 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef RECURRENT_GATED_DELTA_RULE_TORCH_ADPT_H +#define RECURRENT_GATED_DELTA_RULE_TORCH_ADPT_H + +namespace vllm_fl { + +at::Tensor npu_recurrent_gated_delta_rule( + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + at::Tensor& state, + const c10::optional& beta, + const c10::optional scale, + const c10::optional& actual_seq_lengths, + const c10::optional& ssm_state_indices, + const c10::optional& num_accepted_tokens, + const c10::optional& g, + const c10::optional& gk) +{ + TORCH_CHECK(scale.has_value(), "scale cannot be empty."); + + auto options = value.options().dtype(at::ScalarType::BFloat16); + at::Tensor output = at::empty(value.sizes(), options); + float scale_real = static_cast(scale.value()); + EXEC_NPU_CMD(aclnnRecurrentGatedDeltaRule, + query, + key, + value, + beta, + state, + actual_seq_lengths, + ssm_state_indices, + g, + gk, + num_accepted_tokens, + scale_real, + output); + return output; +} + +} // namespace vllm_fl +#endif diff --git a/csrc/ascend/attention/reshape_and_cache_bnsd/CMakeLists.txt b/csrc/ascend/attention/reshape_and_cache_bnsd/CMakeLists.txt new file mode 100644 index 000000000..549b5c14c --- /dev/null +++ b/csrc/ascend/attention/reshape_and_cache_bnsd/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() diff --git a/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/CMakeLists.txt b/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/CMakeLists.txt new file mode 100644 index 000000000..e5eb61691 --- /dev/null +++ b/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/CMakeLists.txt @@ -0,0 +1,24 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + reshape_and_cache_bnsd.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME ReshapeAndCacheBnsd + OPTIONS + --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE reshape_and_cache_bnsd ACLNNTYPE aclnn) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/reshape_and_cache_bnsd.cpp b/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/reshape_and_cache_bnsd.cpp new file mode 100644 index 000000000..adbf88174 --- /dev/null +++ b/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/reshape_and_cache_bnsd.cpp @@ -0,0 +1,66 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file reshape_and_cache_bnsd.cpp + * \brief + */ +#include +#include "register/op_def_registry.h" + +namespace ops { +class ReshapeAndCacheBnsd : public OpDef { +public: + explicit ReshapeAndCacheBnsd(const char* name) : OpDef(name) + { + this->Input("keyIn") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_UINT8, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("keyCacheIn") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_UINT8, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("slotMapping") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("seqLen") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("keyCacheOut") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_UINT8, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("aclnnSupport.value", "support_aclnn") + .ExtendCfgInfo("jitCompile.flag", "static_false,dynamic_false"); + + this->AICore().AddConfig("ascend910_93", aicore_config); + this->AICore().AddConfig("ascend910b", aicore_config); + + } +}; + +OP_ADD(ReshapeAndCacheBnsd); +} diff --git a/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/reshape_and_cache_bnsd_proto.cpp b/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/reshape_and_cache_bnsd_proto.cpp new file mode 100644 index 000000000..aca0c1d2f --- /dev/null +++ b/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/reshape_and_cache_bnsd_proto.cpp @@ -0,0 +1,40 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hamming_dist_top_k_proto.cpp + * \brief + */ +#include +#include +#include "error/ops_error.h" + +using namespace ge; + +namespace ops { +static ge::graphStatus InferShapeReshapeAndCacheBnsd(gert::InferShapeContext *context) +{ + gert::Shape *outShape = context->GetOutputShape(0); + const gert::Shape *inputShape = context->GetInputShape(1); + *outShape = *inputShape; + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataTypeReshapeAndCacheBnsd(gert::InferDataTypeContext *context) +{ + const auto inputDataType = context->GetInputDataType(1); + context->SetOutputDataType(0, inputDataType); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(ReshapeAndCacheBnsd) + .InferShape(InferShapeReshapeAndCacheBnsd) + .InferDataType(InferDataTypeReshapeAndCacheBnsd); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/reshape_and_cache_bnsd_tiling.cpp b/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/reshape_and_cache_bnsd_tiling.cpp new file mode 100644 index 000000000..66480114e --- /dev/null +++ b/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/reshape_and_cache_bnsd_tiling.cpp @@ -0,0 +1,60 @@ + +#include "reshape_and_cache_bnsd_tiling.h" +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" + +namespace optiling { +static ge::graphStatus TilingFunc(gert::TilingContext* context) +{ + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + + ReshapeAndCacheBNSDTilingData tiling; + auto keyShape = context->GetInputShape(0)->GetStorageShape(); + auto keyCacheShape = context->GetInputShape(1)->GetStorageShape(); + auto slotMappingShape = context->GetInputShape(2)->GetStorageShape(); + auto seqLenShape = context->GetInputShape(3)->GetStorageShape(); + int64_t numRow = 1; + + for (size_t i = 0; i < keyCacheShape.GetDimNum() - 1; ++i) { + numRow *= keyCacheShape.GetDim(i); + } + + uint32_t numTokens = static_cast(keyShape.GetDim(0)); + uint32_t headDim = static_cast(keyShape.GetDim(1)); + uint32_t numBlocks = static_cast(keyCacheShape.GetDim(0)); + uint32_t numHeads = static_cast(keyCacheShape.GetDim(1)); + uint32_t blockSize = static_cast(keyCacheShape.GetDim(2)); + uint32_t batchSeqLen = static_cast(slotMappingShape.GetDim(0)); + uint32_t batch = static_cast(seqLenShape.GetDim(0)); + uint32_t numCore = ascendcPlatform.GetCoreNumAiv(); + + tiling.set_numTokens(numTokens); + tiling.set_headDim(headDim); + tiling.set_numBlocks(numBlocks); + tiling.set_numHeads(numHeads); + tiling.set_blockSize(blockSize); + tiling.set_batchSeqLen(batchSeqLen); + tiling.set_batch(batch); + tiling.set_numCore(numCore); + + context->SetTilingKey(0); + context->SetBlockDim(numCore); + + size_t *workspaces = context->GetWorkspaceSizes(1); // get second variable + workspaces[0] = 16 * 1024 * 1024; + + tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus TilingPrepareForReshapeAndCacheBnsd(gert::TilingParseContext *context) +{ + return ge::GRAPH_SUCCESS; +} + + +IMPL_OP_OPTILING(ReshapeAndCacheBnsd) + .Tiling(TilingFunc) + .TilingParse(TilingPrepareForReshapeAndCacheBnsd); +} \ No newline at end of file diff --git a/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/reshape_and_cache_bnsd_tiling.h b/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/reshape_and_cache_bnsd_tiling.h new file mode 100644 index 000000000..28cb47c83 --- /dev/null +++ b/csrc/ascend/attention/reshape_and_cache_bnsd/op_host/reshape_and_cache_bnsd_tiling.h @@ -0,0 +1,21 @@ + +#include "register/tilingdata_base.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(ReshapeAndCacheBNSDTilingData) + TILING_DATA_FIELD_DEF(uint32_t, numTokens); + TILING_DATA_FIELD_DEF(uint32_t, headDim); + TILING_DATA_FIELD_DEF(uint32_t, numBlocks); + TILING_DATA_FIELD_DEF(uint32_t, numHeads); + TILING_DATA_FIELD_DEF(uint32_t, blockSize); + TILING_DATA_FIELD_DEF(uint32_t, batchSeqLen); + TILING_DATA_FIELD_DEF(uint32_t, batch); + TILING_DATA_FIELD_DEF(uint32_t, numCore); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(ReshapeAndCacheBnsd, ReshapeAndCacheBNSDTilingData) +} + +struct reshapeAndCacheBnsdCompileInfo {}; + + diff --git a/csrc/ascend/attention/reshape_and_cache_bnsd/op_kernel/kernel_utils.h b/csrc/ascend/attention/reshape_and_cache_bnsd/op_kernel/kernel_utils.h new file mode 100644 index 000000000..6db56134b --- /dev/null +++ b/csrc/ascend/attention/reshape_and_cache_bnsd/op_kernel/kernel_utils.h @@ -0,0 +1,270 @@ +/* +* Copyright (c) 2024 Huawei Technologies Co., Ltd. +* This file is a part of the CANN Open Software. +* Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). +* Please refer to the License for details. You may not use this file except in compliance with the License. +* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +* See LICENSE in the root of the software repository for the full text of the License. +*/ +#ifndef ASCEND_OPS_UTILS_COMMON_KERNEL_KERNEL_UTILS_H +#define ASCEND_OPS_UTILS_COMMON_KERNEL_KERNEL_UTILS_H +#include "kernel_operator.h" + +using AscendC::HardEvent; + +__aicore__ inline uint32_t CeilDiv(uint32_t x, uint32_t y) +{ + return y == 0 ? 0 : ((x + y - 1) / y); +} + +__aicore__ inline uint32_t RoundUp(uint32_t x, uint32_t y = 16) +{ + return (x + y - 1) / y * y; +} + +__aicore__ inline uint32_t Min(uint32_t x, uint32_t y) +{ + return x < y ? x : y; +} + +__aicore__ inline uint32_t Max(uint32_t x, uint32_t y) +{ + return x > y ? x : y; +} + +template +__aicore__ inline void CopyIn(const AscendC::GlobalTensor &gm, Q &queue, uint64_t offset, uint32_t count) +{ + AscendC::LocalTensor local = queue.template AllocTensor(); + DataCopy(local, gm[offset], count); + queue.EnQue(local); +} + +template +__aicore__ inline void CopyOut(const AscendC::GlobalTensor &gm, Q &queue, uint64_t offset, uint32_t count) +{ + AscendC::LocalTensor local = queue.template DeQue(); + DataCopy(gm[offset], local, count); + queue.FreeTensor(local); +} + +template +__aicore__ inline void CastFrom16To32(const AscendC::LocalTensor &out, const AscendC::LocalTensor &in, + uint32_t count) +{ + Cast(out, in, AscendC::RoundMode::CAST_NONE, count); + AscendC::PipeBarrier(); +} + +template +__aicore__ inline void CastFrom32To16(const AscendC::LocalTensor &out, const AscendC::LocalTensor &in, + uint32_t count) +{ + if constexpr (AscendC::IsSameType::value) { + Cast(out, in, AscendC::RoundMode::CAST_NONE, count); // 310p cast fp32->half 只能用CAST_NONE，这里拉齐310p和910b + } else { // bf16 + Cast(out, in, AscendC::RoundMode::CAST_RINT, count); + } + AscendC::PipeBarrier(); +} + +__aicore__ inline void CastFromF16ToI8(const AscendC::LocalTensor &out, const AscendC::LocalTensor &in, + half quantMin, uint32_t count) +{ + Maxs(in, in, quantMin, count); + AscendC::PipeBarrier(); + Mins(in, in, (half)127, count); // 127: limit + AscendC::PipeBarrier(); +#if defined(__CCE_KT_TEST__) || (__CCE_AICORE__ == 220) + Cast(out, in, AscendC::RoundMode::CAST_RINT, count); +#else + Cast(out, in, AscendC::RoundMode::CAST_NONE, count); +#endif + AscendC::PipeBarrier(); +} + +template +__aicore__ inline void CopyInAndCastF32(const AscendC::LocalTensor &out, const AscendC::GlobalTensor &gm, + Q &queue, uint64_t offset, uint32_t count) +{ + CopyIn(gm, queue, offset, count); + AscendC::LocalTensor local = queue.template DeQue(); + Cast(out, local, AscendC::RoundMode::CAST_NONE, count); + queue.FreeTensor(local); + AscendC::PipeBarrier(); +} + +template +__aicore__ inline void Cast16AndCopyOut(const AscendC::LocalTensor &in, const AscendC::GlobalTensor &gm, + Q &queue, uint64_t offset, uint32_t count) +{ + AscendC::LocalTensor local = queue.template AllocTensor(); + CastFrom32To16(local, in, count); + queue.EnQue(local); + CopyOut(gm, queue, offset, count); + AscendC::PipeBarrier(); +} + +template +__aicore__ inline T ComputeSum(const AscendC::LocalTensor &in, const AscendC::LocalTensor &tmp, + uint32_t count) +{ + ReduceSum(tmp, in, tmp, count); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + return tmp.GetValue(0); +} + +__aicore__ inline float ComputeSliceSquareSum(const AscendC::LocalTensor &in, + const AscendC::LocalTensor &tmp, uint32_t count) +{ + Mul(tmp, in, in, count); + AscendC::PipeBarrier(); + return ComputeSum(tmp, tmp, count); +} +template +__aicore__ inline void ComputeRmsNorm(const AscendC::LocalTensor &out, const AscendC::LocalTensor &in, + float rms, const AscendC::LocalTensor &gamma, uint32_t count, uint32_t precisionMode, uint32_t gemmaMode, + const AscendC::LocalTensor &tmp) +{ + float value = 1.0; + Duplicate(tmp, rms, count); + AscendC::PipeBarrier(); + Div(tmp, in, tmp, count); + AscendC::PipeBarrier(); + + if (precisionMode == 0) { + CastFrom16To32(in, gamma, count); + AscendC::PipeBarrier(); + if (gemmaMode == 1) { + Adds(in, in, value, count); + AscendC::PipeBarrier(); + } + Mul(in, in, tmp, count); + AscendC::PipeBarrier(); + CastFrom32To16(out, in, count); + return; + } + if constexpr (std::is_same::value) { + CastFrom32To16(out, tmp, count); + Mul(out, out, gamma, count); + AscendC::PipeBarrier(); + } +} + +template +__aicore__ inline void ComputeRmsNorm(const AscendC::LocalTensor &out, const AscendC::LocalTensor &in, + float rms, const AscendC::LocalTensor &gamma, const AscendC::LocalTensor &beta, + const AscendC::LocalTensor &tmp, uint32_t count) +{ + Duplicate(tmp, rms, count); + AscendC::PipeBarrier(); + Div(out, in, tmp, count); + AscendC::PipeBarrier(); + CastFrom16To32(tmp, gamma, count); + Mul(out, out, tmp, count); + AscendC::PipeBarrier(); + if constexpr (WITH_BETA) { + CastFrom16To32(tmp, beta, count); + Add(out, out, tmp, count); + AscendC::PipeBarrier(); + } +} + +template +__aicore__ inline void ComputeResidualAdd(const AscendC::LocalTensor &out, + const AscendC::LocalTensor &in, const AscendC::LocalTensor &resIn, uint32_t count) +{ + Add(out, in, resIn, count); + AscendC::PipeBarrier(); +} + +template +__aicore__ inline void ComputeMean(const AscendC::LocalTensor &out, const AscendC::LocalTensor &in, + T aveNum, uint32_t count) +{ + Duplicate(out, aveNum, count); + AscendC::PipeBarrier(); + Mul(out, in, out, count); + AscendC::PipeBarrier(); + T sum = ComputeSum(out, out, count); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + Duplicate(out, sum, count); + AscendC::PipeBarrier(); +} + +__aicore__ inline void ComputeLayerNorm(const AscendC::LocalTensor &out, const AscendC::LocalTensor &in, + const AscendC::LocalTensor &mean, float eps, float aveNum, const AscendC::LocalTensor &gamma, + const AscendC::LocalTensor &beta, uint32_t count) +{ + Sub(in, in, mean, count); + AscendC::PipeBarrier(); + Mul(out, in, in, count); + AscendC::PipeBarrier(); + Muls(out, out, aveNum, count); + AscendC::PipeBarrier(); + ReduceSum(out, out, out, count); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + float var = out.GetValue(0); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + Duplicate(out, var, count); + AscendC::PipeBarrier(); + Adds(out, out, eps, count); + AscendC::PipeBarrier(); + Sqrt(out, out, count); + AscendC::PipeBarrier(); + + Div(out, in, out, count); + AscendC::PipeBarrier(); + + Cast(in, gamma, AscendC::RoundMode::CAST_NONE, count); + AscendC::PipeBarrier(); + Mul(out, out, in, count); + AscendC::PipeBarrier(); + Cast(in, beta, AscendC::RoundMode::CAST_NONE, count); + AscendC::PipeBarrier(); + Add(out, out, in, count); + AscendC::PipeBarrier(); +} + +__aicore__ inline void ComputeFp16ToI8Quant(const AscendC::LocalTensor &out, + const AscendC::LocalTensor &in, const AscendC::LocalTensor &tmp, half scale, half offset, + half quantMin, uint32_t count) +{ + Muls(tmp, in, scale, count); + AscendC::PipeBarrier(); + Adds(tmp, tmp, offset, count); + AscendC::PipeBarrier(); + CastFromF16ToI8(out, tmp, quantMin, count); +} + +__aicore__ inline void ComputeFp32ToI8Quant(const AscendC::LocalTensor &out, + const AscendC::LocalTensor &in, const AscendC::LocalTensor &tmp, half scale, half offset, + half quantMin, uint32_t count) +{ + CastFrom32To16(tmp, in, count); + AscendC::PipeBarrier(); + ComputeFp16ToI8Quant(out, tmp, tmp, scale, offset, quantMin, count); +} + +__aicore__ inline void CopyGmTilingToUb(__ubuf__ uint8_t *tilingInUb, const __gm__ uint8_t *tilingInGm, + size_t tilingSize, AscendC::TPipe *pipe) +{ + uint32_t roundTilingSize = RoundUp(tilingSize, 32); + AscendC::TBuf tilingBuf; + AscendC::GlobalTensor tilingGm; + + tilingGm.SetGlobalBuffer((__gm__ uint8_t *)tilingInGm); + pipe->InitBuffer(tilingBuf, roundTilingSize); + + AscendC::LocalTensor tilingUb = tilingBuf.Get(); + AscendC::DataCopy(tilingUb, tilingGm, roundTilingSize); + + tilingInUb = (__ubuf__ uint8_t *)tilingUb.GetPhyAddr(); +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/reshape_and_cache_bnsd/op_kernel/reshape_and_cache_bnsd.cpp b/csrc/ascend/attention/reshape_and_cache_bnsd/op_kernel/reshape_and_cache_bnsd.cpp new file mode 100644 index 000000000..9cffb0cc6 --- /dev/null +++ b/csrc/ascend/attention/reshape_and_cache_bnsd/op_kernel/reshape_and_cache_bnsd.cpp @@ -0,0 +1,229 @@ +#include "kernel_utils.h" + +constexpr int32_t ALIGN = 32; +using namespace AscendC; + +#define YF_LOG(format, ...) \ + if (false) { \ + printf("CoreIdx: %d on CoreType %d, " format, GetBlockIdx(), g_coreType, ##__VA_ARGS__); \ + } + + +class ReshapeAndCacheBnsd { +public: + __aicore__ inline ReshapeAndCacheBnsd(ReshapeAndCacheBNSDTilingData tilingData) + : batchNum_(tilingData.batch), blockSize_(tilingData.blockSize), + coreNum_(tilingData.numCore), headNum_(tilingData.numHeads), headDim_(tilingData.headDim) + {} + + __aicore__ inline void Init(GM_ADDR keyIn, GM_ADDR keyCacheIn, GM_ADDR slotMapping, GM_ADDR seqLen, + GM_ADDR keyCacheOut) + { + AscendC::TPipe pipe; + pipe.InitBuffer(ubBuf_, RoundUp(blockSize_ * headDim_, ALIGN)); + tmpTensor_ = ubBuf_.Get(); + keyInGm_.SetGlobalBuffer((__gm__ uint8_t *)keyIn); + keyCacheInGm_.SetGlobalBuffer((__gm__ uint8_t *)keyCacheIn); + slotMappingGm_.SetGlobalBuffer((__gm__ int32_t *)slotMapping); + seqLenGm_.SetGlobalBuffer((__gm__ int32_t *)seqLen); + keyCacheOutGm_.SetGlobalBuffer((__gm__ uint8_t *)keyCacheOut); + } + + __aicore__ inline void Process() + { + // Calculate the total number of pages + uint32_t totalBlockNum = 0; + uint32_t offsetInSlotmapping = 0; + for (uint32_t batchIdx = 0; batchIdx < batchNum_; batchIdx++) { + uint32_t seqLen = seqLenGm_.GetValue(batchIdx); + int32_t slotValue = slotMappingGm_.GetValue(offsetInSlotmapping); + uint32_t offsetInBlock = slotValue % blockSize_; + uint32_t leftTokenNum = blockSize_ - offsetInBlock; + uint32_t blockNumForCurrBatch = seqLen < leftTokenNum ? 1 : + (CeilDiv(seqLen - leftTokenNum, blockSize_) + 1); + totalBlockNum += blockNumForCurrBatch; + offsetInSlotmapping += seqLen; + + //YF_LOG("batchIdx: %d, totalBlockNum: %d, offsetInSlotmapping: %d\n", batchIdx, totalBlockNum, offsetInSlotmapping); + } + + uint32_t blockIdx_ = GetBlockIdx(); + uint32_t actualCoreNum = totalBlockNum <= coreNum_ ? totalBlockNum : coreNum_; + // How many pages each core transfers + uint32_t blockNumPerCore = totalBlockNum / actualCoreNum; + uint32_t leftBlockNum = totalBlockNum - blockNumPerCore * actualCoreNum; + uint32_t blockNum = blockIdx_ < leftBlockNum ? blockNumPerCore + 1 : blockNumPerCore; + uint32_t startBlockOffset_ = blockIdx_ < leftBlockNum ? (blockNumPerCore * blockIdx_ + blockIdx_) : + (blockNumPerCore * blockIdx_ + leftBlockNum); + if (blockIdx_ >= actualCoreNum) { + return; + } + + // Position of keyIn and KeyCache corresponding to startBlockOffset_ for each core + uint32_t startBatchIdx = 0; + uint32_t accuBlockNum = 0; + uint32_t startTokenOffsetInBatch = 0; + offsetInSlotmapping = 0; + bool copyFromBatchStart = true; + for (uint32_t batchIdx = 0; batchIdx < batchNum_; batchIdx++) { + uint32_t seqLen = seqLenGm_.GetValue(batchIdx); + int32_t slotValue = slotMappingGm_.GetValue(offsetInSlotmapping); + uint32_t offsetInBlock = slotValue % blockSize_; + uint32_t leftTokenNum = blockSize_ - offsetInBlock; + uint32_t blockNumForCurrBatch = seqLen < leftTokenNum ? 1 : + (CeilDiv(seqLen - leftTokenNum, blockSize_) + 1); + accuBlockNum += blockNumForCurrBatch; + + if (startBlockOffset_ == 0) { + break; + } else if (accuBlockNum == startBlockOffset_) { + startBatchIdx = batchIdx + 1; + startTokenOffsetInBatch = 0; + copyFromBatchStart = true; + offsetInSlotmapping = offsetInSlotmapping + seqLen; + break; + } else if (accuBlockNum > startBlockOffset_) { + startBatchIdx = batchIdx; + startTokenOffsetInBatch = (startBlockOffset_ - (accuBlockNum - blockNumForCurrBatch + 1)) * + blockSize_ +leftTokenNum; + copyFromBatchStart = false; + offsetInSlotmapping = offsetInSlotmapping + startTokenOffsetInBatch; + break; + } + offsetInSlotmapping += seqLen; + } + + uint32_t batchIdx = startBatchIdx; + for (uint32_t blockIdx = 0; blockIdx < blockNum; blockIdx++) { + uint32_t seqLen = seqLenGm_.GetValue(batchIdx); + int32_t slotValue = slotMappingGm_.GetValue(offsetInSlotmapping); + uint32_t blockId = static_cast(slotValue) / blockSize_; + uint32_t slotId = static_cast(slotValue) % blockSize_; + + if (startTokenOffsetInBatch + blockSize_ - slotId > seqLen) { + //YF_LOG("batchIdx: %d, true\n", batchIdx); + uint32_t currCopyTokenNum = seqLen - startTokenOffsetInBatch; + uint32_t copyBlocks = CeilDiv(currCopyTokenNum * headDim_, 32); + //YF_LOG("batchIdx: %d, currCopyTokenNum: %d, copyBlocks: %d from %d\n", batchIdx, currCopyTokenNum, copyBlocks, currCopyTokenNum * headDim_); + AscendC::DataCopyParams copyInParams = {1, static_cast(copyBlocks), 0, 0}; + AscendC::DataCopyParams copyOutParams = {1, static_cast(copyBlocks), 0, 0}; + int64_t dstOffset = blockId * headNum_ * blockSize_ * headDim_ + slotId * headDim_; + int64_t srcOffset = (offsetInSlotmapping - startTokenOffsetInBatch) * headNum_ * headDim_ + + startTokenOffsetInBatch * headDim_; + //YF_LOG("batchIdx: %d, srcOffset[%d] -> dstOffset[%d], size: %d\n", batchIdx, srcOffset, dstOffset, static_cast(copyBlocks)); + + for (uint32_t headId = 0; headId < headNum_; headId++) { + DataCopy(tmpTensor_, keyInGm_[srcOffset + headId * seqLen * headDim_], copyInParams); + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + DataCopy(keyCacheOutGm_[dstOffset + headId * blockSize_* headDim_], tmpTensor_, copyOutParams); + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + //YF_LOG("batchIdx: %d, src[%d] -> dst[%d], size: %d\n", batchIdx, srcOffset + headId * seqLen * headDim_, dstOffset + headId * blockSize_* headDim_, static_cast(copyBlocks)); + + } + batchIdx += 1; + startTokenOffsetInBatch = 0; + offsetInSlotmapping += currCopyTokenNum; + } else { + uint32_t currCopyTokenNum = blockSize_ - slotId; + uint32_t copyBlocks = currCopyTokenNum * headDim_ / ALIGN; + //YF_LOG("batchIdx: %d, currCopyTokenNum: %d, currCopyTokenNum * headDim_: %d\n", batchIdx, currCopyTokenNum, currCopyTokenNum * headDim_); + uint32_t leftBytes = currCopyTokenNum * headDim_ - copyBlocks * ALIGN; + AscendC::DataCopyParams copyInParams = {1, static_cast(copyBlocks), 0, 0}; + AscendC::DataCopyParams copyOutParams = {1, static_cast(copyBlocks), 0, 0}; + int64_t dstOffset = blockId * headNum_ * blockSize_ * headDim_ + slotId * headDim_; + int64_t srcOffset = (offsetInSlotmapping - startTokenOffsetInBatch) * headNum_ * headDim_ + + startTokenOffsetInBatch * headDim_; + if (copyBlocks != 0) { + for (uint32_t headId = 0; headId < headNum_; headId++) { + DataCopy(tmpTensor_, keyInGm_[srcOffset + headId * seqLen * headDim_], copyInParams); + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + DataCopy(keyCacheOutGm_[dstOffset + headId * blockSize_* headDim_], tmpTensor_, copyOutParams); + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + //YF_LOG("batchIdx: %d, src[%d] -> dst[%d], size: %d\n", batchIdx, srcOffset + headId * seqLen * headDim_, dstOffset + headId * blockSize_* headDim_, static_cast(copyBlocks)); + } + } + + if (currCopyTokenNum + startTokenOffsetInBatch == seqLen) { + batchIdx += 1; + startTokenOffsetInBatch = 0; + offsetInSlotmapping += currCopyTokenNum; + } else { + startTokenOffsetInBatch += currCopyTokenNum; + offsetInSlotmapping += currCopyTokenNum; + } + if (leftBytes == 0) { + continue; + } + // If there is a tail block, process it; it is less than 32 bytes. + for (uint32_t headId = 0; headId < headNum_; headId++) { + for (uint32_t dimId = 0; dimId < leftBytes; dimId++) { + uint8_t cacheValue = keyInGm_.GetValue(srcOffset + headId * seqLen * headDim_ + + copyBlocks * ALIGN + dimId); + keyCacheOutGm_.SetValue(dstOffset + headId * blockSize_ * headDim_ + + copyBlocks * ALIGN + dimId, cacheValue); + } + } + // TODO: Move DataCacheCleanAndInvalid outside the loop, + // to resolve the issue where partial data cannot be read correctly. + AscendC::DataCacheCleanAndInvalid(keyCacheOutGm_); + } + } + } + +private: + GlobalTensor keyInGm_; + GlobalTensor keyCacheInGm_; + GlobalTensor slotMappingGm_; + GlobalTensor seqLenGm_; + GlobalTensor keyCacheOutGm_; + TBuf ubBuf_; + LocalTensor tmpTensor_; + LocalTensor keyIn_; + LocalTensor keyCacheIn_; + LocalTensor slotMapping_; + LocalTensor seqLen_; + LocalTensor keyCacheOut_; + + uint32_t batchNum_{0}; + uint32_t blockSize_{0}; + uint32_t coreNum_{0}; + uint32_t headNum_{0}; + uint32_t headDim_{0}; +}; + +inline __aicore__ void InitTilingData(const __gm__ uint8_t *p_tilingdata, + ReshapeAndCacheBNSDTilingData *tilingdata) { + tilingdata->numTokens = (*(const __gm__ uint32_t *)(p_tilingdata + 0)); + tilingdata->headDim = (*(const __gm__ uint32_t *)(p_tilingdata + 4)); + tilingdata->numBlocks = (*(const __gm__ uint32_t *)(p_tilingdata + 8)); + tilingdata->numHeads = (*(const __gm__ uint32_t *)(p_tilingdata + 12)); + tilingdata->blockSize = (*(const __gm__ uint32_t *)(p_tilingdata + 16)); + tilingdata->batchSeqLen = (*(const __gm__ uint32_t *)(p_tilingdata + 20)); + tilingdata->batch = (*(const __gm__ uint32_t *)(p_tilingdata + 24)); + tilingdata->numCore = (*(const __gm__ uint32_t *)(p_tilingdata + 28)); + + //YF_LOG("numTokens: %d\n", tilingdata->numTokens); + //YF_LOG("headDim: %d\n", tilingdata->headDim); + //YF_LOG("numBlocks: %d\n", tilingdata->numBlocks); + //YF_LOG("numHeads: %d\n", tilingdata->numHeads); + //YF_LOG("blockSize: %d\n", tilingdata->blockSize); + //YF_LOG("batchSeqLen: %d\n", tilingdata->batchSeqLen); + //YF_LOG("batch: %d\n", tilingdata->batch); + //YF_LOG("numCore: %d\n", tilingdata->numCore); +} + +extern "C" __global__ __aicore__ void reshape_and_cache_bnsd(GM_ADDR keyIn, GM_ADDR keyCacheIn, GM_ADDR slotMapping, GM_ADDR seqLen, GM_ADDR keyCacheOut, GM_ADDR workspace, GM_ADDR tiling) { + // ReshapeAndCacheBNSDTilingData tilingData; + // InitTilingData(tiling, &tilingData); + + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); + GET_TILING_DATA(tilingData, tiling); + + ReshapeAndCacheBnsd op(tilingData); + op.Init(keyIn, keyCacheIn, slotMapping, seqLen, keyCacheOut); + op.Process(); +} \ No newline at end of file diff --git a/csrc/ascend/attention/rms_norm_dynamic_quant/CMakeLists.txt b/csrc/ascend/attention/rms_norm_dynamic_quant/CMakeLists.txt new file mode 100644 index 000000000..549b5c14c --- /dev/null +++ b/csrc/ascend/attention/rms_norm_dynamic_quant/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() diff --git a/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/CMakeLists.txt b/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/CMakeLists.txt new file mode 100644 index 000000000..45bd3235d --- /dev/null +++ b/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/CMakeLists.txt @@ -0,0 +1,61 @@ +# This program is free software, you can redistribute it and/or modify it. +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This file is a part of the CANN Open Software. +# Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ====================================================================================================================== + +# add_ops_compile_options( +# OP_NAME RmsNormDynamicQuant +# OPTIONS --cce-auto-sync=off +# -Wno-deprecated-declarations +# -Werror +# -mllvm -cce-aicore-hoist-movemask=false +# --op_relocatable_kernel_binary=true +# ) + +# target_sources(op_host_aclnn PRIVATE +# op_host/rms_norm_dynamic_quant_def.cpp +# ) + +# target_sources(optiling PRIVATE +# op_host/rms_norm_dynamic_quant_tiling.cpp +# ) + +# if (NOT BUILD_OPEN_PROJECT) +# target_sources(opmaster_ct PRIVATE +# op_host/rms_norm_dynamic_quant_tiling.cpp +# ) +# endif () + +# target_include_directories(optiling PRIVATE +# ${CMAKE_CURRENT_SOURCE_DIR}/op_host +# ) + +# target_sources(opsproto PRIVATE +# op_host/rms_norm_dynamic_quant_proto.cpp +# ) + +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + rms_norm_dynamic_quant_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME RmsNormDynamicQuant + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true +) + + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE rms_norm_dynamic_quant ACLNNTYPE aclnn) +endif() diff --git a/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/rms_norm_dynamic_quant_def.cpp b/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/rms_norm_dynamic_quant_def.cpp new file mode 100644 index 000000000..e4e6d0feb --- /dev/null +++ b/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/rms_norm_dynamic_quant_def.cpp @@ -0,0 +1,85 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rms_norm_dynamic_quant_def.cpp + * \brief + */ +#include "register/op_def_registry.h" + +namespace ops { +class RmsNormDynamicQuant : public OpDef { +public: + explicit RmsNormDynamicQuant(const char* name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("gamma") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("smooth_scale1") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("smooth_scale2") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("beta") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("y1") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT4, ge::DT_INT4}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("y2") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT4, ge::DT_INT4}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("scale1") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("scale2") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Attr("epsilon").AttrType(OPTIONAL).Float(1e-6); + this->Attr("output_mask").AttrType(OPTIONAL).ListBool({}); + this->Attr("dst_type").AttrType(OPTIONAL).Int(ge::DT_INT8); + + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + } +}; +OP_ADD(RmsNormDynamicQuant); +} // namespace ops diff --git a/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/rms_norm_dynamic_quant_proto.cpp b/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/rms_norm_dynamic_quant_proto.cpp new file mode 100644 index 000000000..796154db6 --- /dev/null +++ b/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/rms_norm_dynamic_quant_proto.cpp @@ -0,0 +1,82 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file rms_norm_dynamic_quant_proto.h + * \brief + */ +#ifndef OPS_RMS_NORM_DYNAMIC_QUANT_PROTO_H_ +#define OPS_RMS_NORM_DYNAMIC_QUANT_PROTO_H_ + +#include "graph/operator_reg.h" + +namespace ge { +/** +* @brief Fused Operator of RmsNorm and DynamicQuant. +* Calculating input: x, gamma, smooth_scale1, smooth_scale2 \n +* Calculating process: \n +* rstd = np.rsqrt(np.mean(np.power(x, 2), reduce_axis, keepdims=True) + epsilon)) \n +* rmsnorm_out = x * rstd * gamma \n +* if smooth_scales1 exist: \n +* scale1 = row_max(abs(rmsnorm_out * smooth_scale1)) / 127 \n +* if smooth_scales1 not exist: \n +* scale1 = row_max(abs(rmsnorm_out)) / 127 \n +* y1 = round(rmsnorm_out / scale1) \n +* if smooth_scales2 exist: \n +* scale2 = row_max(abs(rmsnorm_out * smooth_scale2)) / 127 \n +* y2 = round(rmsnorm_out / scale2) \n +* if smooth_scales2 not exist: \n +* not calculate scale2 and y2. \n + +* @par Inputs +* @li x: A tensor. Input x for the operation. +* Support dtype: float16/bfloat16, support format: ND. +* @li gamma: A tensor. Describing the weight of the rmsnorm operation. +* Support dtype: float16/bfloat16, support format: ND. +* @li smooth_scale1: A tensor. Describing the weight of the first dynamic quantization. +* Support dtype: float16/bfloat16, support format: ND. +* @li smooth_scale2: An optional input tensor. Describing the weight of the secend dynamic quantization. +* Support dtype: float16/bfloat16, support format: ND. +* @li beta: An optional input tensor. Describing the offset value of dynamic quantization. +* Support dtype: float16/bfloat16, support format: ND. Has the same dtype and shape as "gamma". +* @par Attributes +* @li epsilon: An optional attribute. Describing the epsilon of the rmsnorm operation. +* The type is float. Defaults to 1e-6. +* @li dst_type: An optional int32. Output y data type enum value. Support DT_INT8, DT_INT4, DT_HIFLOAT8, DT_FLOAT8_E5M2, +* DT_FLOAT8_E4M3FN. Defaults to DT_INT8. + +* @par Outputs +* @li y1: A tensor. Describing the output of the first dynamic quantization. +* Support dtype: int8/hifloat8/float8e5m2/float8e4m3fn, support format: ND. +* @li y2: A tensor. Describing the output of the second dynamic quantization. +* Support dtype: int8/hifloat8/float8e5m2/float8e4m3fn, support format: ND. +* @li scale1: A tensor. Describing of the factor for the first dynamic quantization. +* Support dtype: float32, support format: ND. +* @li scale2: A tensor. Describing of the factor for the second dynamic quantization. +* Support dtype: float32, support format: ND. +*/ + +REG_OP(RmsNormDynamicQuant) + .INPUT(x, TensorType({DT_FLOAT16, DT_BF16})) + .INPUT(gamma, TensorType({DT_FLOAT16, DT_BF16})) + .OPTIONAL_INPUT(smooth_scale1, TensorType({DT_FLOAT16, DT_BF16})) + .OPTIONAL_INPUT(smooth_scale2, TensorType({DT_FLOAT16, DT_BF16})) + .OPTIONAL_INPUT(beta, TensorType({DT_FLOAT16, DT_BF16})) + .OUTPUT(y1, TensorType({DT_INT8, DT_HIFLOAT8, DT_FP8_E5M2, DT_FP8_E4M3FN, DT_INT4})) + .OUTPUT(y2, TensorType({DT_INT8, DT_HIFLOAT8, DT_FP8_E5M2, DT_FP8_E4M3FN, DT_INT4})) + .OUTPUT(scale1, TensorType({DT_FLOAT, DT_FLOAT})) + .OUTPUT(scale2, TensorType({DT_FLOAT, DT_FLOAT})) + .ATTR(epsilon, Float, 1e-6) + .ATTR(output_mask, ListBool, {}) + .ATTR(dst_type, Int, DT_INT8) + .OP_END_FACTORY_REG(RmsNormDynamicQuant) +} // namespace ge + +#endif // OPS_RMS_NORM_DYNAMIC_QUANT_PROTO_H_ diff --git a/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/rms_norm_dynamic_quant_tiling.cpp b/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/rms_norm_dynamic_quant_tiling.cpp new file mode 100644 index 000000000..04651047c --- /dev/null +++ b/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/rms_norm_dynamic_quant_tiling.cpp @@ -0,0 +1,539 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_dynamic_quant_tiling.cpp + * \brief + */ +#include "rms_norm_dynamic_quant_tiling.h" + +namespace optiling { + +constexpr int X_IDX = 0; +constexpr int GAMMA_IDX = 1; +constexpr int SMOOTH1_IDX = 2; +constexpr int SMOOTH2_IDX = 3; +constexpr int BETA_IDX = 4; + +constexpr int Y1_IDX = 0; +constexpr int Y2_IDX = 1; +constexpr int SCALE1_IDX = 2; +constexpr int SCALE2_IDX = 3; + +constexpr int NUM_WITH_BETA = 4; +constexpr int NUM_WITHOUT_BETA = 3; + +constexpr int EPS_IDX = 0; +constexpr int OUT_QUANT_1_IDX = 1; +constexpr int OUT_QUANT_2_IDX = 2; +constexpr int DST_TYPE_IDX = 2; + +constexpr uint64_t USR_WORKSPACE_SIZE_910B = 1; + +constexpr uint32_t SIZEOF_B16 = 2; +constexpr uint32_t BLOCK_SIZE = 32; +constexpr uint64_t ROW_FACTOR = 128; +constexpr uint64_t UB_RESERVED_BYTE = 768; +constexpr uint32_t MAX_ROW_STEP = 16; +constexpr uint32_t INT4_ALIGN_SIZE = 64; + +constexpr uint32_t UB_TILING_POLICY_NORMAL = 1; +constexpr uint32_t UB_TILING_POLICY_SINGLE_ROW = 2; +constexpr uint32_t UB_TILING_POLICY_SLICE_D = 3; + +constexpr uint32_t SLICE_COL_LEN = 8864; +constexpr uint32_t SLICE_COL_LEN_INT4 = 8832; + +constexpr int32_t INT_NEGATIVE_ONE = -1; +constexpr int32_t INT_ZERO = 0; +constexpr int32_t INT_ONE = 1; +constexpr int32_t INT_TWO = 2; + +template +static inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd))); +} + +template +static inline T CeilAlign(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)) * (rnd)); +} + +bool CheckOptionalShapeExisting(const gert::StorageShape* smoothShape) +{ + OPS_CHECK(nullptr == smoothShape, OPS_LOG_D("CheckOptionalShapeExisting", "Get nullptr smoothShape"), return false); + int64_t smoothShapeSize = smoothShape->GetOriginShape().GetShapeSize(); + OPS_CHECK((smoothShapeSize <= 0), OPS_LOG_D("CheckOptionalShapeExisting", "Get empty smoothShape"), return false); + return true; +} + +bool CheckOptionalBetaExisting(const gert::StorageShape* betaShape) +{ + OPS_CHECK(nullptr == betaShape, OPS_LOG_D("CheckOptionalBetaExisting", "Get nullptr betaShape"), return false); + int64_t betaShapeSize = betaShape->GetOriginShape().GetShapeSize(); + OPS_CHECK((betaShapeSize <= 0), OPS_LOG_D("CheckOptionalBetaExisting", "Get empty betaShape"), return false); + return true; +} + +size_t GetworkspaceRowsNum(int32_t outQuant1Flag, int32_t outQuant2Flag, uint32_t smoothNum1_, uint32_t smoothNum2_) +{ + size_t workspaceRowsNum = INT_ZERO; + if ((outQuant1Flag == INT_NEGATIVE_ONE && outQuant2Flag == INT_NEGATIVE_ONE)) { + workspaceRowsNum = (smoothNum1_ == INT_ZERO && smoothNum2_ == INT_ZERO) ? INT_ONE : INT_TWO; + } else { + workspaceRowsNum = (outQuant1Flag == INT_ONE || outQuant2Flag == INT_ONE) ? INT_TWO : INT_ONE; + } + return workspaceRowsNum; +} + +void RmsNormDynamicQuantTilingHelper::SetTilingDataAndTilingKeyAndWorkSpace(RmsNormDynamicQuantTilingData* tiling) +{ + context_->SetBlockDim(this->useCore_); + tiling->set_useCore(this->useCore_); + tiling->set_numFirstDim(this->numFirstDim_); + tiling->set_numLastDim(this->numLastDim_); + tiling->set_numLastDimAligned(this->numLastDimAligned_); + tiling->set_firstDimPerCore(this->firstDimPerCore_); + tiling->set_firstDimPerCoreTail(this->firstDimPerCoreTail_); + tiling->set_firstDimPerLoop(this->firstDimPerLoop_); + tiling->set_lastDimSliceLen(this->lastDimSliceLen_); + tiling->set_lastDimLoopNum(this->lastDimLoopNum_); + tiling->set_lastDimSliceLenTail(this->lastDimSliceLenTail_); + tiling->set_smoothNum1(this->smoothNum1_); + tiling->set_smoothNum2(this->smoothNum2_); + tiling->set_epsilon(this->eps_); + tiling->set_outQuant1Flag(this->outQuant1Flag); + tiling->set_outQuant2Flag(this->outQuant2Flag); + tiling->set_avgFactor(this->avgFactor_); + tiling->set_betaFlag(this->betaFlag_); + uint32_t tilingKey = 0; + size_t usrSize = USR_WORKSPACE_SIZE_910B; + + if (this->ubTilingPolicy_ == UB_TILING_POLICY::NORMAL) { + tilingKey += UB_TILING_POLICY_NORMAL; + } else if (this->ubTilingPolicy_ == UB_TILING_POLICY::SINGLE_ROW) { + tilingKey += UB_TILING_POLICY_SINGLE_ROW; + } else { + tilingKey += UB_TILING_POLICY_SLICE_D; + size_t workspaceRowsNum = + GetworkspaceRowsNum(this->outQuant1Flag, this->outQuant2Flag, this->smoothNum1_, this->smoothNum2_); + usrSize = this->useCore_ * this->numLastDim_ * sizeof(float) * workspaceRowsNum; + } + + context_->SetTilingKey(tilingKey); + + tiling->SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tiling->GetDataSize()); + + // set workspace + size_t* currentWorkspace = context_->GetWorkspaceSizes(1); + currentWorkspace[0] = this->sysWorkspaceSize_ + usrSize; + + OPS_LOG_I( + "SetTilingDataAndTilingKeyAndWorkSpace", "Tilingdata useCore_: %lu, smoothNum1_: %u, smoothNum2_: %u", + this->useCore_, this->smoothNum1_, this->smoothNum2_); + OPS_LOG_I( + "SetTilingDataAndTilingKeyAndWorkSpace", "Tilingdata N: %lu, D:%lu, DAligned: %lu", numFirstDim_, numLastDim_, + numLastDimAligned_); + OPS_LOG_I( + "SetTilingDataAndTilingKeyAndWorkSpace", "Tilingdata firstDimPerCore_: %lu, firstDimPerCoreTail_: %lu", + firstDimPerCore_, firstDimPerCoreTail_); + OPS_LOG_I("SetTilingDataAndTilingKeyAndWorkSpace", "Tilingdata firstDimPerLoop_: %lu", firstDimPerLoop_); + OPS_LOG_I( + "SetTilingDataAndTilingKeyAndWorkSpace", + "Tilingdata lastDimSliceLen_: %lu, lastDimLoopNum_: %lu, lastDimSliceLenTail_: %lu", lastDimSliceLen_, + lastDimLoopNum_, lastDimSliceLenTail_); + OPS_LOG_I("SetTilingDataAndTilingKeyAndWorkSpace", "Tilingdata eps_: %f, avgFactor_: %f", eps_, avgFactor_); + OPS_LOG_I( + "SetTilingDataAndTilingKeyAndWorkSpace", "Tilingdata tilingKey = %u, usr Workspace: %zu", tilingKey, usrSize); +} + +bool RmsNormDynamicQuantTilingHelper::DoTiling() +{ + OPS_CHECK( + (nullptr == context_), OPS_LOG_E("AddRmsNormDynamicQuantTiling", "Helper context_ get nullptr, return failed."), + return false); + OPS_CHECK(!GetBaseInfo(), OPS_LOG_E(context_->GetNodeName(), "GetBaseInfo failed, return false"), return false); + OPS_CHECK( + !GetShapeInfo(), OPS_LOG_E(context_->GetNodeName(), "GetShapeInfo failed, return false"), return false); + OPS_CHECK( + !DoBlockTiling(), OPS_LOG_E(context_->GetNodeName(), "DoBlockTiling failed, return false"), return false); + OPS_CHECK(!DoUbTiling(), OPS_LOG_E(context_->GetNodeName(), "DoUbTiling failed, return false"), return false); + return true; +} + +bool RmsNormDynamicQuantTilingHelper::DoBlockTiling() +{ + // Block Tiling, Cut N + this->firstDimPerCore_ = CeilDiv(this->numFirstDim_, this->socCoreNums_); + this->useCore_ = CeilDiv(this->numFirstDim_, this->firstDimPerCore_); + this->firstDimPerCore_ = CeilDiv(this->numFirstDim_, this->useCore_); + this->firstDimPerCoreTail_ = this->numFirstDim_ - this->firstDimPerCore_ * (this->useCore_ - 1); + OPS_LOG_I( + "DoBlockTiling", "BlockTiling Factor: useCore_: %lu, firstDimPerCore_: %lu, firstDimPerCoreTail_: %lu", + this->useCore_, this->firstDimPerCore_, this->firstDimPerCoreTail_); + return true; +} + +bool RmsNormDynamicQuantTilingHelper::InitializePlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + // OP_CHECK_NULL_WITH_CONTEXT(context_, platformInfo); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + this->socCoreNums_ = ascendcPlatform.GetCoreNumAiv(); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, this->ubSize_); + this->sysWorkspaceSize_ = ascendcPlatform.GetLibApiWorkSpaceSize(); + return true; +} + +bool RmsNormDynamicQuantTilingHelper::GetBaseInfo() +{ + if (!InitializePlatformInfo()) { + return false; + } + + auto attrs = context_->GetAttrs(); + OPS_CHECK( + nullptr == attrs, OPS_LOG_E(context_->GetNodeName(), "Get attrs nullptr, return false."), return false); + + const float* epsPtr = attrs->GetFloat(EPS_IDX); + if (epsPtr != nullptr) { + this->eps_ = *epsPtr; + } + + const gert::ContinuousVector* outputMaskAttr = attrs->GetAttrPointer(OUT_QUANT_1_IDX); + if (outputMaskAttr != nullptr && outputMaskAttr->GetSize() == INT_TWO) { + const bool* scalesArray = static_cast(outputMaskAttr->GetData()); + this->outQuant1Flag = (scalesArray[0] == true) ? 1 : 0; + this->outQuant2Flag = (scalesArray[1] == true) ? 1 : 0; + } else { + this->outQuant1Flag = -1; + this->outQuant2Flag = -1; + } + OPS_LOG_I("outputMask", "outQuant1Flag: %u, outQuant2Flag: %u", this->outQuant1Flag, this->outQuant2Flag); + if (!ValidateBaseParameters()) { + return false; + } + OPS_LOG_I( + "GetBaseInfo", "socCoreNum: %lu, ubSize: %lu, sysWorkspaceSize: %lu, epsilon: %f", this->socCoreNums_, + this->ubSize_, this->sysWorkspaceSize_, this->eps_); + + return true; +} + +bool RmsNormDynamicQuantTilingHelper::ValidateBaseParameters() +{ + OPS_CHECK( + this->eps_ <= 0, + OPS_LOG_E(context_->GetNodeName(), "Epsilon less or equal than precision threshold, please check."), + return false); + OPS_CHECK( + (this->ubSize_ <= 0), OPS_LOG_E(context_->GetNodeName(), "ubSize less or equal than zero, please check."), + return false); + OPS_CHECK( + (this->socCoreNums_ <= 0), + OPS_LOG_E(context_->GetNodeName(), "socCoreNums_ less or equal than zero, please check."), return false); + + return true; +} + +ge::graphStatus CheckDtypeVaild(ge::DataType& srcDtype, std::vector& supportDtypeList) +{ + for (const auto& supportedDtype : supportDtypeList) { + if (supportedDtype == srcDtype) { + return ge::GRAPH_SUCCESS; + } + } + return ge::GRAPH_FAILED; +} + +bool RmsNormDynamicQuantTilingHelper::ValidateInputOutput() +{ + // 检查输入输出形状 + OPS_CHECK( + CheckInputOutputShape() == false, OPS_LOG_E(context_->GetNodeName(), "Check tensor shape failed."), return false); + + // 验证输出数据类型 + auto y1DataType = context_->GetOutputDesc(Y1_IDX)->GetDataType(); + auto y2DataType = context_->GetOutputDesc(Y2_IDX)->GetDataType(); + std::vector supportedYDtypes = {ge::DataType::DT_INT8, ge::DataType::DT_INT4}; + if ((ge::GRAPH_SUCCESS != CheckDtypeVaild(y1DataType, supportedYDtypes)) || + (ge::GRAPH_SUCCESS != CheckDtypeVaild(y2DataType, supportedYDtypes)) || (y1DataType != y2DataType)) { + OPS_LOG_E(context_->GetNodeName(), "Output dtype should be int8 int4 hifp8 and y1DataType y2DataType need same."); + return false; + } + + return true; +} + +bool RmsNormDynamicQuantTilingHelper::CalculateShapeParameters() +{ + // 设置数据类型大小 + this->dtSize_ = SIZEOF_B16; + + // 获取输入形状 + auto xShape = context_->GetInputShape(X_IDX)->GetStorageShape(); + auto gammaShape = context_->GetInputShape(GAMMA_IDX)->GetStorageShape(); + size_t xDimNum = xShape.GetDimNum(); + size_t gammaDimNum = gammaShape.GetDimNum(); + + // 计算numRow和numCol + uint64_t numRow = 1; + uint64_t numCol = 1; + for (size_t i = 0; i < xDimNum - gammaDimNum; i++) { + numRow *= xShape.GetDim(i); + } + for (size_t i = 0; i < gammaDimNum; i++) { + numCol *= gammaShape.GetDim(i); + } + + // 设置对齐大小和目标类型 + this->numFirstDim_ = numRow; + this->numLastDim_ = numCol; + auto y1DataType = context_->GetOutputDesc(Y1_IDX)->GetDataType(); + uint32_t alignSize = y1DataType == ge::DT_INT4 ? INT4_ALIGN_SIZE : BLOCK_SIZE; + this->dstType_ = static_cast(y1DataType); + this->numLastDimAligned_ = + CeilDiv(numCol, static_cast(alignSize)) * static_cast(alignSize); + + // 计算平均因子 + this->avgFactor_ = 1.0 / ((float)this->numLastDim_); + + return true; +} + +bool RmsNormDynamicQuantTilingHelper::SetFlagsAndCheckConsistency() +{ + // 检查可选输入是否存在 + const gert::StorageShape* smooth1Shape = this->context_->GetOptionalInputShape(SMOOTH1_IDX); + const gert::StorageShape* smooth2Shape = this->context_->GetOptionalInputShape(SMOOTH2_IDX); + const gert::StorageShape* betaShape = this->context_->GetOptionalInputShape(BETA_IDX); + bool smooth1Exist = CheckOptionalShapeExisting(smooth1Shape); + bool smooth2Exist = CheckOptionalShapeExisting(smooth2Shape); + bool betaExist = CheckOptionalBetaExisting(betaShape); + + // 设置标志位 + this->smoothNum1_ = (smooth1Exist) ? 1 : 0; + this->smoothNum2_ = (smooth2Exist) ? 1 : 0; + this->betaFlag_ = (betaExist) ? 1 : 0; + + // 检查形状匹配性 + auto gammaShape = context_->GetInputShape(GAMMA_IDX)->GetStorageShape(); + OPS_CHECK( + (smooth1Exist && smooth1Shape->GetStorageShape() != gammaShape), + OPS_LOG_E(context_->GetNodeName(), "GammaShape is not same to smooth1Shape."), return false); + OPS_CHECK( + (smooth2Exist && smooth2Shape->GetStorageShape() != gammaShape), + OPS_LOG_E(context_->GetNodeName(), "GammaShape is not same to smooth2Shape."), return false); + + // 检查量化标志和可选输入的一致性 + if (this->outQuant1Flag == INT_NEGATIVE_ONE && this->outQuant2Flag == INT_NEGATIVE_ONE) { + OPS_CHECK( + (!smooth1Exist) && (smooth2Exist), + OPS_LOG_E(context_->GetNodeName(), "Smooth2 exist but smooth1 not exist, bad input."), return false); + } + + return true; +} + +bool RmsNormDynamicQuantTilingHelper::GetShapeInfo() +{ + // 验证输入输出 + if (!ValidateInputOutput()) { + return false; + } + + // 计算形状参数 + if (!CalculateShapeParameters()) { + return false; + } + + // 设置标志和检查一致性 + if (!SetFlagsAndCheckConsistency()) { + return false; + } + + // 打印日志 + OPS_LOG_I("GetShapeInfo", "[N, D] = [%lu, %lu]", this->numFirstDim_, this->numLastDim_); + OPS_LOG_I("GetShapeInfo", "dtSize_=%lu, avgFactor_=%f", this->dtSize_, this->avgFactor_); + return true; +} + +bool RmsNormDynamicQuantTilingHelper::DoUbTiling() +{ + OPS_CHECK(CheckUbNormalTiling(), OPS_LOG_I(context_->GetNodeName(), "Ub Tiling: Normal."), return true); + OPS_CHECK(CheckUbSingleRowTiling(), OPS_LOG_I(context_->GetNodeName(), "Ub Tiling: SingleRow."), return true); + OPS_CHECK(CheckUbSliceDTiling(), OPS_LOG_I(context_->GetNodeName(), "Ub Tiling: SliceD."), return true); + return false; +} + +bool RmsNormDynamicQuantTilingHelper::CheckUbNormalTiling() +{ + // 3 weights tensor required. + int64_t ubConst = 0; + if (this->betaFlag_ == 1) { + ubConst = this->numLastDimAligned_ * this->dtSize_ * NUM_WITH_BETA + UB_RESERVED_BYTE; + } else { + ubConst = this->numLastDimAligned_ * this->dtSize_ * NUM_WITHOUT_BETA + UB_RESERVED_BYTE; + } + int64_t ubAvaliable = this->ubSize_ - ubConst; + // 2 rows for tmpBuffer. + int64_t coexistingRowsNum = 2 * (this->dtSize_) + 2 * (this->dtSize_) + 1 * sizeof(float) + 1 * sizeof(float); + // 2 buffers for out_scale. + int64_t rowCommons = coexistingRowsNum * this->numLastDimAligned_ + 2 * sizeof(float); + int64_t rowStep = ubAvaliable / rowCommons; + bool ret = (rowStep >= 1); + OPS_LOG_I( + this->context_->GetNodeName(), + "CheckUbNormalTiling, ret:%d, ubConst: %ld, ubAvaliable=%ld, coexistingRowsNum: %ld, rowStep: %ld, " + "rowCommons: %ld", + ret, ubConst, ubAvaliable, coexistingRowsNum, rowStep, rowCommons); + if (ret) { + // No mutilN now. max RowStep = 16 + this->firstDimPerLoop_ = (rowStep <= MAX_ROW_STEP) ? rowStep : MAX_ROW_STEP; + this->lastDimSliceLen_ = this->numLastDimAligned_; + this->lastDimLoopNum_ = 1; + this->lastDimSliceLenTail_ = 0; + this->ubTilingPolicy_ = UB_TILING_POLICY::NORMAL; + } + return ret; +} + +bool RmsNormDynamicQuantTilingHelper::CheckUbSingleRowTiling() +{ + // 2 tmp buffer, 2 rows copy in and 1 rows copy out + int64_t ubRequired = ((2 + 1 + 1) * this->dtSize_ + 2 * sizeof(float)) * this->numLastDimAligned_; + ubRequired = ubRequired + 2L * ROW_FACTOR * sizeof(float); + bool ret = (((int64_t)this->ubSize_) >= ubRequired); + OPS_LOG_I(this->context_->GetNodeName(), "CheckUbSingleRowTiling, ret:%d, ubRequired: %ld", ret, ubRequired); + if (ret) { + this->firstDimPerLoop_ = 1; + this->lastDimSliceLen_ = this->numLastDimAligned_; + this->lastDimLoopNum_ = 1; + this->lastDimSliceLenTail_ = 0; + this->ubTilingPolicy_ = UB_TILING_POLICY::SINGLE_ROW; + } + return ret; +} + +bool RmsNormDynamicQuantTilingHelper::CheckUbSliceDTiling() +{ + OPS_LOG_I(this->context_->GetNodeName(), "CheckUbSliceDTiling success. Compute tiling by yourself."); + this->ubTilingPolicy_ = UB_TILING_POLICY::SLICE_D; + this->firstDimPerLoop_ = 1; + if (this->dstType_ == 29) { + this->lastDimSliceLen_ = SLICE_COL_LEN_INT4; + } else { + this->lastDimSliceLen_ = SLICE_COL_LEN; + } + this->lastDimSliceLenTail_ = (this->numLastDim_ % this->lastDimSliceLen_ == 0) ? + this->lastDimSliceLen_ : + this->numLastDim_ % this->lastDimSliceLen_; + this->lastDimLoopNum_ = (this->numLastDim_ - this->lastDimSliceLenTail_) / this->lastDimSliceLen_; + return true; +} + +ge::graphStatus Tiling4AddRmsNormDynamicQuant(gert::TilingContext* context) +{ + OPS_CHECK(nullptr == context, OPS_LOG_E("AddRmsNormDynamicQuant", "Context is null"), return ge::GRAPH_FAILED); + OPS_LOG_I(context->GetNodeName(), "Enter Tiling4AddRmsNormDynamicQuant"); + auto colShape = context->GetInputShape(GAMMA_IDX); + // OP_CHECK_NULL_WITH_CONTEXT(context, colShape); + auto colStorageShape = optiling::EnsureNotScalar(colShape->GetStorageShape()); + uint32_t col_val = colStorageShape.GetDim(0); + bool isEmptyTensor = (col_val == 0); + auto ptrCompileInfo = reinterpret_cast(context->GetCompileInfo()); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + platform_ascendc::SocVersion curSocVersion = + (ptrCompileInfo) == nullptr ? ascendcPlatform.GetSocVersion() : ptrCompileInfo->curSocVersion; + RmsNormDynamicQuantTilingData tiling; + RmsNormDynamicQuantTilingHelper instanceNormV3TilingHelper(context); + bool status = instanceNormV3TilingHelper.DoTiling(); + OPS_CHECK( + !status, OPS_LOG_E(context->GetNodeName(), "DoTiling Failed, return Failed."), return ge::GRAPH_FAILED); + instanceNormV3TilingHelper.SetTilingDataAndTilingKeyAndWorkSpace(&tiling); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingPrepare4AddRmsNormDynamicQuant(gert::TilingParseContext* context) +{ + OPS_CHECK(nullptr == context, OPS_LOG_E("AddRmsNormDynamicQuant", "Context is null"), return ge::GRAPH_FAILED); + OPS_LOG_D(context->GetNodeName(), "Enter TilingPrepare4AddRmsNormDynamicQuant."); + fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo(); + // OP_CHECK_NULL_WITH_CONTEXT(context, platformInfoPtr); + + auto compileInfoPtr = context->GetCompiledInfo(); + // OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + compileInfoPtr->curSocVersion = ascendcPlatform.GetSocVersion(); + compileInfoPtr->totalCoreNum = ascendcPlatform.GetCoreNumAiv(); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, compileInfoPtr->maxUbSize); + return ge::GRAPH_SUCCESS; +} + +bool RmsNormDynamicQuantTilingHelper::CheckInputOutputShape() +{ + // Check Shape Not NULL + const gert::StorageShape* xShape = this->context_->GetInputShape(X_IDX); + const gert::StorageShape* gammaShape = this->context_->GetInputShape(GAMMA_IDX); + + const gert::StorageShape* y1Shape = this->context_->GetOutputShape(Y1_IDX); + const gert::StorageShape* y2Shape = this->context_->GetOutputShape(Y2_IDX); + const gert::StorageShape* scale1Shape = this->context_->GetOutputShape(SCALE1_IDX); + const gert::StorageShape* scale2Shape = this->context_->GetOutputShape(SCALE2_IDX); + + // OP_CHECK_NULL_WITH_CONTEXT(this->context_, xShape); + // OP_CHECK_NULL_WITH_CONTEXT(this->context_, gammaShape); + // OP_CHECK_NULL_WITH_CONTEXT(this->context_, y1Shape); + // OP_CHECK_NULL_WITH_CONTEXT(this->context_, y2Shape); + // OP_CHECK_NULL_WITH_CONTEXT(this->context_, scale1Shape); + // OP_CHECK_NULL_WITH_CONTEXT(this->context_, scale2Shape); + + // Check Shape relations + size_t xDimNum = xShape->GetStorageShape().GetDimNum(); + size_t gammaDimNum = gammaShape->GetStorageShape().GetDimNum(); + size_t y1DimNum = y1Shape->GetStorageShape().GetDimNum(); + size_t y2DimNum = y2Shape->GetStorageShape().GetDimNum(); + size_t scale1DimNum = scale1Shape->GetStorageShape().GetDimNum(); + size_t scale2DimNum = scale2Shape->GetStorageShape().GetDimNum(); + + OPS_LOG_I( + this->context_->GetNodeName(), + "ShapeDim info: x.dim=%zu, gamma.dim=%zu, y1.dim=%zu, y2.dim=%zu, scale1.dim=%zu, " + "scale2.dim=%zu", + xDimNum, gammaDimNum, y1DimNum, y2DimNum, scale1DimNum, scale2DimNum); + + bool hasZeroDimTensor = xDimNum <= 0 || gammaDimNum <= 0; + OPS_CHECK( + (hasZeroDimTensor), + OPS_LOG_E( + this->context_->GetNodeName(), + "Input x/y1/scale1DimNum shape invalid, dim num should not be smaller or equal to zero."), + return false); + OPS_CHECK( + ((gammaDimNum != 1)), OPS_LOG_E(this->context_->GetNodeName(), "gamma shape dims not equal to 1. Tiling failed."), + return false); + gert::Shape shapeOfX = xShape->GetStorageShape(); + gert::Shape shapeOfGamma = gammaShape->GetStorageShape(); + OPS_CHECK( + (shapeOfX[xDimNum - 1] != shapeOfGamma[gammaDimNum - 1]), + OPS_LOG_E(context_->GetNodeName(), "gammaShape isn't consistent with the last dimension of x."), return false); + return true; +} + +IMPL_OP_OPTILING(RmsNormDynamicQuant) + .Tiling(Tiling4AddRmsNormDynamicQuant) + .TilingParse(TilingPrepare4AddRmsNormDynamicQuant); + +} // namespace optiling diff --git a/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/rms_norm_dynamic_quant_tiling.h b/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/rms_norm_dynamic_quant_tiling.h new file mode 100644 index 000000000..2c393dd63 --- /dev/null +++ b/csrc/ascend/attention/rms_norm_dynamic_quant/op_host/rms_norm_dynamic_quant_tiling.h @@ -0,0 +1,133 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +/*! + * \file add_rms_norm_dynamic_quant_tiling.h + */ +#ifndef OPS_BUILT_IN_OP_TILING_RUNTIME_ADD_RMS_NORM_DYN_QUANT_TILING_H +#define OPS_BUILT_IN_OP_TILING_RUNTIME_ADD_RMS_NORM_DYN_QUANT_TILING_H +#include "register/op_impl_registry.h" +#include "platform/platform_infos_def.h" +#include "exe_graph/runtime/tiling_context.h" +#include "tiling/platform/platform_ascendc.h" +#include "register/op_def_registry.h" +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "error/ops_error.h" +#include "platform/platform_info.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(RmsNormDynamicQuantTilingData) +TILING_DATA_FIELD_DEF(uint64_t, useCore); +TILING_DATA_FIELD_DEF(uint64_t, numFirstDim); +TILING_DATA_FIELD_DEF(uint64_t, numLastDim); +TILING_DATA_FIELD_DEF(uint64_t, numLastDimAligned); +TILING_DATA_FIELD_DEF(uint64_t, firstDimPerCore); +TILING_DATA_FIELD_DEF(uint64_t, firstDimPerCoreTail); +TILING_DATA_FIELD_DEF(uint64_t, firstDimPerLoop); +TILING_DATA_FIELD_DEF(uint64_t, lastDimLoopNum); +TILING_DATA_FIELD_DEF(uint64_t, lastDimSliceLen); +TILING_DATA_FIELD_DEF(uint64_t, lastDimSliceLenTail); +TILING_DATA_FIELD_DEF(uint32_t, smoothNum1); +TILING_DATA_FIELD_DEF(uint32_t, smoothNum2); +TILING_DATA_FIELD_DEF(float, epsilon); +TILING_DATA_FIELD_DEF(int32_t, outQuant1Flag); +TILING_DATA_FIELD_DEF(int32_t, outQuant2Flag); +TILING_DATA_FIELD_DEF(float, avgFactor); +TILING_DATA_FIELD_DEF(uint32_t, betaFlag); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(RmsNormDynamicQuant, RmsNormDynamicQuantTilingData); + +constexpr uint32_t TILING_TYPE_NORMAL = 0; +constexpr uint32_t TILING_TYPE_SPILT = 1; +constexpr uint32_t TILING_OFFSET_HAS_QUANT = 10; +constexpr uint32_t TILING_OFFSET_REGBASE = 100; +constexpr uint64_t TILING_KEY_UNRUN = 199; + +struct RmsNormDynamicQuantCompileInfo { + platform_ascendc::SocVersion curSocVersion = platform_ascendc::SocVersion::ASCEND910B; + uint64_t totalCoreNum = 0; + uint64_t maxUbSize = 0; +}; + +enum class UB_TILING_POLICY : std::int32_t +{ + NORMAL, + SINGLE_ROW, + SLICE_D +}; + +static const gert::Shape g_vec_1_shape = {1}; + +inline const gert::Shape& EnsureNotScalar(const gert::Shape& inShape) +{ + if (inShape.IsScalar()) { + return g_vec_1_shape; + } + return inShape; +} + +class RmsNormDynamicQuantTilingHelper { +public: + explicit RmsNormDynamicQuantTilingHelper(gert::TilingContext* context) : context_(context) + {} + + ~RmsNormDynamicQuantTilingHelper() = default; + bool DoTiling(); + void SetTilingDataAndTilingKeyAndWorkSpace(RmsNormDynamicQuantTilingData* tiling); + +private: + bool GetBaseInfo(); + bool GetShapeInfo(); + bool DoBlockTiling(); + bool DoUbTiling(); + bool CheckInputOutputShape(); + + bool CheckUbNormalTiling(); + bool CheckUbSingleRowTiling(); + bool CheckUbSliceDTiling(); + bool ValidateBaseParameters(); + bool InitializePlatformInfo(); + bool ValidateInputOutput(); + bool CalculateShapeParameters(); + bool SetFlagsAndCheckConsistency(); + + gert::TilingContext* context_; + + ge::DataType xDtype_{ge::DataType::DT_FLOAT16}; + uint64_t dtSize_{2}; + uint64_t socCoreNums_{1}; + uint64_t ubSize_{1}; + uint64_t sysWorkspaceSize_{1}; + + uint64_t useCore_{1}; + uint64_t numFirstDim_{1}; + uint64_t numLastDim_{1}; + uint64_t numLastDimAligned_{1}; + uint64_t firstDimPerCore_{1}; + uint64_t firstDimPerCoreTail_{1}; + uint64_t firstDimPerLoop_{1}; + uint64_t lastDimSliceLen_{1}; + uint64_t lastDimLoopNum_{1}; + uint64_t lastDimSliceLenTail_{1}; + float eps_{1e-6}; + int32_t outQuant1Flag{0}; + int32_t outQuant2Flag{0}; + float avgFactor_{0.0}; + uint32_t smoothNum1_{0}; + uint32_t smoothNum2_{0}; + uint32_t betaFlag_{0}; + uint32_t dstType_{2}; + + UB_TILING_POLICY ubTilingPolicy_{UB_TILING_POLICY::SINGLE_ROW}; +}; +} // namespace optiling + +#endif // OPS_BUILT_IN_OP_TILING_RUNTIME_ADD_RMS_NORM_DYN_QUANT_TILING_H diff --git a/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/reduce_common.h b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/reduce_common.h new file mode 100644 index 000000000..9e9f7ac2c --- /dev/null +++ b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/reduce_common.h @@ -0,0 +1,167 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +/*! + * \file reduce_common.h + */ +#ifndef REDUCE_COMMON_H_RMS_NORM +#define REDUCE_COMMON_H_RMS_NORM +#include "kernel_operator.h" +using namespace AscendC; + +constexpr uint32_t MAX_REP_NUM = 255; +constexpr uint32_t ELEM_PER_REP_FP32 = 64; +constexpr uint32_t ELEM_PER_BLK_FP32 = 8; +constexpr float ZERO = 0; +constexpr int32_t HALf_INTERVAL = 2; +constexpr int32_t INDEX_TWO = 2; +constexpr int32_t INDEX_FOUR = 4; +constexpr int32_t INDEX_EIGHT = 8; +constexpr int32_t INDEX_SIXTEEN = 16; + +__aicore__ inline void ReduceSumForSmallReduceDimPreRepeat( + const LocalTensor& dstLocal, const LocalTensor& srcLocal, const LocalTensor& tmpLocal, + const uint32_t elemNum, const uint32_t numLastDim, const uint32_t tailCount, const uint32_t repeat, + const uint8_t repStride) +{ + uint32_t elemIndex = 0; + for (; elemIndex + ELEM_PER_REP_FP32 <= numLastDim; elemIndex += ELEM_PER_REP_FP32) { + Add(tmpLocal, srcLocal[elemIndex], tmpLocal, elemNum, repeat, + {1, 1, 1, ELEM_PER_BLK_FP32, repStride, ELEM_PER_BLK_FP32}); + PipeBarrier(); + } + if (unlikely(tailCount != 0)) { + Add(tmpLocal, srcLocal[elemIndex], tmpLocal, tailCount, repeat, + {1, 1, 1, ELEM_PER_BLK_FP32, repStride, ELEM_PER_BLK_FP32}); + } + PipeBarrier(); + AscendCUtils::SetMask(ELEM_PER_REP_FP32); // set mask = 64 + WholeReduceSum(dstLocal, tmpLocal, MASK_PLACEHOLDER, repeat, 1, 1, ELEM_PER_BLK_FP32); +} + +/* + * reduce dim form (N, D) to (N, 1) + * this reduce sum is for small reduce dim. + */ +__aicore__ inline void ReduceSumForSmallReduceDim( + const LocalTensor& dstLocal, const LocalTensor& srcLocal, const LocalTensor& tmpLocal, + const uint32_t numLastDimAligned, const uint32_t numLastDim, const uint32_t tailCount, const uint32_t repeat, + const uint8_t repStride) +{ + uint32_t repeatTimes = repeat / MAX_REP_NUM; + if (repeatTimes == 0) { + ReduceSumForSmallReduceDimPreRepeat( + dstLocal, srcLocal, tmpLocal, ELEM_PER_REP_FP32, numLastDim, tailCount, repeat, repStride); + } else { + uint32_t repTailNum = repeat % MAX_REP_NUM; + uint32_t repIndex = 0; + uint32_t repElem; + for (; repIndex + MAX_REP_NUM <= repeat; repIndex += MAX_REP_NUM) { + ReduceSumForSmallReduceDimPreRepeat( + dstLocal[repIndex], srcLocal[repIndex * numLastDimAligned], tmpLocal[repIndex * ELEM_PER_REP_FP32], + ELEM_PER_REP_FP32, numLastDim, tailCount, MAX_REP_NUM, repStride); + } + if (repTailNum != 0) { + ReduceSumForSmallReduceDimPreRepeat( + dstLocal[repIndex], srcLocal[repIndex * numLastDimAligned], tmpLocal[repIndex * ELEM_PER_REP_FP32], + ELEM_PER_REP_FP32, numLastDim, tailCount, repTailNum, repStride); + } + } +} + +/* + * reduce dim form (N, D) to (N, 1) + * this reduce sum is for small reduce dim, require D < 255 * 8. + * size of tmpLocal: (N, 64) + */ +__aicore__ inline void ReduceSumMultiN( + const LocalTensor& dstLocal, const LocalTensor& srcLocal, const LocalTensor& tmpLocal, + const uint32_t numRow, const uint32_t numCol, const uint32_t numColAlign) +{ + const uint32_t tailCount = numCol % ELEM_PER_REP_FP32; + const uint32_t repeat = numRow; + const uint8_t repStride = numColAlign / ELEM_PER_BLK_FP32; + Duplicate(tmpLocal, ZERO, numRow * ELEM_PER_REP_FP32); + PipeBarrier(); + ReduceSumForSmallReduceDim(dstLocal, srcLocal, tmpLocal, numColAlign, numCol, tailCount, repeat, repStride); +} + +__aicore__ inline int32_t findPowerTwo(int32_t n) +{ + // find max power of 2 no more than n (32 bit) + n |= n >> 1; // Set the first digit of n's binary to 1 + n |= n >> INDEX_TWO; + n |= n >> INDEX_FOUR; + n |= n >> INDEX_EIGHT; + n |= n >> INDEX_SIXTEEN; + return (n + 1) >> 1; +} + +__aicore__ inline void ReduceSumHalfInterval( + const LocalTensor& dst_local, const LocalTensor& src_local, int32_t count) +{ + if (likely(count > ELEM_PER_REP_FP32)) { + int32_t bodyCount = findPowerTwo(count); + int32_t tailCount = count - bodyCount; + if (tailCount > 0) { + Add(src_local, src_local, src_local[bodyCount], tailCount); + PipeBarrier(); + } + while (bodyCount > ELEM_PER_REP_FP32) { + bodyCount = bodyCount / HALf_INTERVAL; + Add(src_local, src_local, src_local[bodyCount], bodyCount); + PipeBarrier(); + } + + AscendCUtils::SetMask(ELEM_PER_REP_FP32); + } else { + AscendCUtils::SetMask(count); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + if (g_coreType == AIV) { + WholeReduceSum(dst_local, src_local, MASK_PLACEHOLDER, 1, 0, 1, 0); + } +#else + WholeReduceSum(dst_local, src_local, MASK_PLACEHOLDER, 1, 1, 1, DEFAULT_REPEAT_STRIDE); +#endif + PipeBarrier(); +} + +__aicore__ inline float ReduceSumHalfInterval(const LocalTensor& src_local, int32_t count) +{ + if (likely(count > ELEM_PER_REP_FP32)) { + int32_t bodyCount = findPowerTwo(count); + int32_t tailCount = count - bodyCount; + if (tailCount > 0) { + Add(src_local, src_local, src_local[bodyCount], tailCount); + PipeBarrier(); + } + while (bodyCount > ELEM_PER_REP_FP32) { + bodyCount = bodyCount / HALf_INTERVAL; + Add(src_local, src_local, src_local[bodyCount], bodyCount); + PipeBarrier(); + } + + AscendCUtils::SetMask(ELEM_PER_REP_FP32); + } else { + AscendCUtils::SetMask(count); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + if (g_coreType == AIV) { + WholeReduceSum(src_local, src_local, MASK_PLACEHOLDER, 1, 0, 1, 0); + } +#else + WholeReduceSum(src_local, src_local, MASK_PLACEHOLDER, 1, 1, 1, DEFAULT_REPEAT_STRIDE); +#endif + event_t event_v_s = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(event_v_s); + WaitFlag(event_v_s); + return src_local.GetValue(0); +} +#endif // _REDUCE_COMMON_H_ diff --git a/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant.cpp b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant.cpp new file mode 100644 index 000000000..2d784a8a4 --- /dev/null +++ b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant.cpp @@ -0,0 +1,42 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_dynamic_quant.cpp + * \brief + */ +#include "rms_norm_dynamic_quant_normal_kernel.h" +#include "rms_norm_dynamic_quant_single_row_kernel.h" +#include "rms_norm_dynamic_quant_cut_d_kernel.h" + +extern "C" __global__ __aicore__ void rms_norm_dynamic_quant( + GM_ADDR x, GM_ADDR gamma, GM_ADDR smooth1, GM_ADDR smooth2, GM_ADDR beta, GM_ADDR y1, GM_ADDR y2, + GM_ADDR outScale1, GM_ADDR outScale2, GM_ADDR workspace, GM_ADDR tiling) +{ + TPipe pipe; + GET_TILING_DATA(tilingData, tiling); + GM_ADDR usrWorkspace = AscendC::GetUserWorkspace(workspace); + +#define INIT_AND_PROCESS \ + op.Init(x, gamma, smooth1, smooth2, beta, y1, y2, outScale1, outScale2, usrWorkspace, &tilingData); \ + op.Process() + if (TILING_KEY_IS(0)) { + // 0 Tiling, Do Nothing. + } else if (TILING_KEY_IS(1)) { + KernelAddRmsNormDynamicQuantNormal op(&pipe); + INIT_AND_PROCESS; + } else if (TILING_KEY_IS(2)) { + KernelAddRmsNormDynamicQuantSingleRow op(&pipe); + INIT_AND_PROCESS; + } else if (TILING_KEY_IS(3)) { + KernelAddRmsNormDynamicQuantSliceD op(&pipe); + INIT_AND_PROCESS; + } +} diff --git a/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_base.h b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_base.h new file mode 100644 index 000000000..35a635bf5 --- /dev/null +++ b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_base.h @@ -0,0 +1,146 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_dynamic_quant_base.h + * \brief + */ + +#ifndef ADD_RMS_NORM_DYNAMIC_QUANT_BASE_CLASS_H_ +#define ADD_RMS_NORM_DYNAMIC_QUANT_BASE_CLASS_H_ + +#include "rms_norm_dynamic_quant_helper.h" + +template +class KernelAddRmsNormDynamicQuantBase { +public: + __aicore__ inline KernelAddRmsNormDynamicQuantBase() + {} + + __aicore__ inline void InitBaseParams(const RmsNormDynamicQuantTilingData* tiling) + { + this->numCore = tiling->useCore; + this->numFirstDim = tiling->numFirstDim; + this->numLastDim = tiling->numLastDim; + this->numLastDimAligned = tiling->numLastDimAligned; // Quantize better be aligned to 32 elements + + this->firstDimPerCore = tiling->firstDimPerCore; + this->firstDimPerCoreTail = tiling->firstDimPerCoreTail; + this->firstDimPerLoop = tiling->firstDimPerLoop; + + this->lastDimSliceLen = tiling->lastDimSliceLen; + this->lastDimLoopNum = tiling->lastDimLoopNum; + this->lastDimSliceLenTail = tiling->lastDimSliceLenTail; + this->betaFlag = tiling->betaFlag; + this->eps = tiling->epsilon; + this->aveNum = tiling->avgFactor; + + blockIdx_ = GetBlockIdx(); + if (blockIdx_ != this->numCore - 1) { + this->rowWork = this->firstDimPerCore; + this->rowStep = this->firstDimPerLoop; + } else { + this->rowWork = this->firstDimPerCoreTail; + this->rowStep = TWO_NUMS_MIN(this->firstDimPerLoop, this->rowWork); + } + this->rowTail_ = (this->rowWork % this->rowStep == 0) ? this->rowStep : (this->rowWork % this->rowStep); + this->gmOffset_ = this->firstDimPerCore * this->numLastDim; + + this->smooth1Exist = tiling->smoothNum1; + // 2 dynamic quant operator required 2 scale buffer. + this->smooth2Exist = tiling->smoothNum2; + + // dynamic quant max value + if constexpr (IsSameType::value) { + this->quantMaxVal = DYNAMIC_QUANT_DIVIDEND; + } else { + this->quantMaxVal = DYNAMIC_QUANT_DIVIDEND_INT4; + } + this->outQuant1Flag = tiling->outQuant1Flag; + this->outQuant2Flag = tiling->outQuant2Flag; + + this->isOld = (this->outQuant1Flag == -1) && (this->outQuant2Flag == -1); + this->oldDouble = this->isOld && this->smooth1Exist && this->smooth2Exist; + this->newSingleFirst = this->smooth1Exist && (this->outQuant1Flag == 1); + this->newSingleSecond = this->smooth2Exist && (this->outQuant2Flag == 1); + } + + __aicore__ inline void InitInGlobalTensors( + GM_ADDR x, GM_ADDR gamma, GM_ADDR smooth1, GM_ADDR smooth2, GM_ADDR beta) + { + xGm.SetGlobalBuffer((__gm__ T*)(x) + blockIdx_ * this->gmOffset_); + gammaGm.SetGlobalBuffer((__gm__ T*)gamma); + smooth1Gm.SetGlobalBuffer((__gm__ T*)smooth1); + smooth2Gm.SetGlobalBuffer((__gm__ T*)smooth2); + if (this->betaFlag == 1) { + betaGm.SetGlobalBuffer((__gm__ T*)beta); + } + } + + __aicore__ inline void InitOutGlobalTensors(GM_ADDR y1, GM_ADDR y2, GM_ADDR outScale1, GM_ADDR outScale2) + { + int64_t yBufferSize = blockIdx_ * this->gmOffset_; + if constexpr (IsSameType::value) { + yBufferSize = yBufferSize / 2; + } + y1Gm.SetGlobalBuffer((__gm__ T_Y*)(y1) + yBufferSize); + y2Gm.SetGlobalBuffer((__gm__ T_Y*)(y2) + yBufferSize); + outScale1Gm.SetGlobalBuffer((__gm__ float*)outScale1 + blockIdx_ * this->firstDimPerCore); + outScale2Gm.SetGlobalBuffer((__gm__ float*)outScale2 + blockIdx_ * this->firstDimPerCore); + } + + __aicore__ inline void InitWorkSpaceGlobalTensors(GM_ADDR workspace) + {} + +protected: + GlobalTensor xGm; + GlobalTensor gammaGm; + GlobalTensor smooth1Gm; + GlobalTensor smooth2Gm; + GlobalTensor betaGm; + GlobalTensor y1Gm; + GlobalTensor y2Gm; + GlobalTensor outScale1Gm; + GlobalTensor outScale2Gm; + + uint32_t betaFlag; + uint64_t numCore; + uint64_t numFirstDim; + uint64_t numLastDim; + uint64_t numLastDimAligned; + uint64_t firstDimPerCore; + uint64_t firstDimPerCoreTail; + uint64_t firstDimPerLoop; + uint64_t lastDimSliceLen; + uint64_t lastDimLoopNum; + uint64_t lastDimSliceLenTail; + + float eps; + float aveNum; + + uint64_t blockIdx_; + uint64_t gmOffset_; + uint64_t rowTail_; + uint64_t rowStep; + uint64_t rowWork; + + bool smooth1Exist; + bool smooth2Exist; + int32_t outQuant1Flag; + int32_t outQuant2Flag; + + bool isOld; + bool oldDouble; + bool newSingleFirst; + bool newSingleSecond; + float quantMaxVal; +}; + +#endif // __ADD_RMS_NORM_DYNAMIC_QUANT_BASE_CLASS_H_ diff --git a/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_cut_d_kernel.h b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_cut_d_kernel.h new file mode 100644 index 000000000..b9538238c --- /dev/null +++ b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_cut_d_kernel.h @@ -0,0 +1,407 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_dynamic_quant_cut_d_kernel.h + * \brief + */ + +#ifndef ADD_RMS_NORM_DYNAMIC_QUANT_SLICE_D_H_ +#define ADD_RMS_NORM_DYNAMIC_QUANT_SLICE_D_H_ + +#include "rms_norm_dynamic_quant_base.h" + +template +class KernelAddRmsNormDynamicQuantSliceD : public KernelAddRmsNormDynamicQuantBase { +public: + __aicore__ inline KernelAddRmsNormDynamicQuantSliceD(TPipe* pipe) + { + Ppipe = pipe; + } + + __aicore__ inline void Init( + GM_ADDR x, GM_ADDR gamma, GM_ADDR smooth1, GM_ADDR smooth2, GM_ADDR beta, GM_ADDR y1, GM_ADDR y2, + GM_ADDR outScale1, GM_ADDR outScale2, GM_ADDR workspace, const RmsNormDynamicQuantTilingData* tiling) + { + this->InitBaseParams(tiling); + this->InitInGlobalTensors(x, gamma, smooth1, smooth2, beta); + this->InitOutGlobalTensors(y1, y2, outScale1, outScale2); + + if (this->oldDouble || (this->outQuant2Flag == 1 || this->outQuant1Flag == 1)) { + workspaceGm.SetGlobalBuffer((__gm__ float*)(workspace) + 2 * this->blockIdx_ * this->numLastDim); + } else { + workspaceGm.SetGlobalBuffer((__gm__ float*)(workspace) + this->blockIdx_ * this->numLastDim); + } + + /* + colFactor = 8864 + UB = 3 * colFactor * sizeof(T) + 1 * colFactor * sizeof(float) + + 3 * colFactor * sizeof(float) + + 256B_for_reduce + 64B_for_scale + */ + Ppipe->InitBuffer(inRowsQue, BUFFER_NUM, 2 * this->lastDimSliceLen * sizeof(T)); // 2 * D * 2 + Ppipe->InitBuffer(outRowQue, BUFFER_NUM, this->lastDimSliceLen * sizeof(T)); // D * 2 + Ppipe->InitBuffer(tmpOutQue, BUFFER_NUM, this->lastDimSliceLen * sizeof(float)); // D * 4 + Ppipe->InitBuffer(xBufFp32, this->lastDimSliceLen * sizeof(float)); // D * 4 + Ppipe->InitBuffer(yBufFp32, this->lastDimSliceLen * sizeof(float)); // D * 4 + Ppipe->InitBuffer(zBufFp32, this->lastDimSliceLen * sizeof(float)); // D * 4 + // 2 dynamic quant operator required 2 scale buffer. + Ppipe->InitBuffer(scalesQue, BUFFER_NUM, 2 * ELEM_PER_BLK_FP32 * sizeof(float)); + } + + __aicore__ inline void Process() + { + uint32_t baseGmOffset = 0; + uint32_t rowGmOffset = 0; + for (int32_t rowIdx = 0; rowIdx < this->rowWork; ++rowIdx) { + rowGmOffset = 0; + this->localSum = ZERO; + this->localMax1 = ZERO; + this->localMax2 = ZERO; + for (int32_t colIdx = 0; colIdx < this->lastDimLoopNum; ++colIdx) { + CopyInX(baseGmOffset, rowGmOffset, this->lastDimSliceLen); + this->localSum += ReduceSquareSumSlice(this->lastDimSliceLen); + PipeBarrier(); + rowGmOffset += this->lastDimSliceLen; + } + + { + CopyInX(baseGmOffset, rowGmOffset, this->lastDimSliceLenTail); + this->localSum += ReduceSquareSumSlice(this->lastDimSliceLenTail); + PipeBarrier(); + } + float rstdLocalTemp = 1 / sqrt(this->localSum * this->aveNum + this->eps); + PIPE_S_V(); + PIPE_MTE3_MTE2(); + + rowGmOffset = 0; + for (int32_t colIdx = 0; colIdx < this->lastDimLoopNum; ++colIdx) { + ComputeRmsNormAndSmoothMax(rowGmOffset, this->lastDimSliceLen, rstdLocalTemp); + rowGmOffset += this->lastDimSliceLen; + } + { + ComputeRmsNormAndSmoothMax(rowGmOffset, this->lastDimSliceLenTail, rstdLocalTemp); + } + if (this->isOld || (this->outQuant1Flag == 1)) { + this->localMax1 = this->quantMaxVal / this->localMax1; + } + if (this->outQuant2Flag == 1 || this->oldDouble) { + this->localMax2 = this->quantMaxVal / this->localMax2; + } + PIPE_S_V(); + PIPE_MTE3_MTE2(); + + rowGmOffset = 0; + for (int32_t colIdx = 0; colIdx < this->lastDimLoopNum; ++colIdx) { + ComputeDynamicQuant(rowGmOffset, this->lastDimSliceLen); + CopyOutQuant(baseGmOffset, rowGmOffset, this->lastDimSliceLen); + rowGmOffset += this->lastDimSliceLen; + } + { + ComputeDynamicQuant(rowGmOffset, this->lastDimSliceLenTail); + CopyOutQuant(baseGmOffset, rowGmOffset, this->lastDimSliceLenTail); + } + LocalTensor scalesTensor = scalesQue.template AllocTensor(); + if (this->isOld || this->outQuant1Flag == 1) { + scalesTensor.SetValue(0, 1 / this->localMax1); + } + if (this->oldDouble || (this->outQuant2Flag == 1)) { + scalesTensor.SetValue(ELEM_PER_BLK_FP32, 1 / this->localMax2); + } + scalesQue.EnQue(scalesTensor); + CopyOutScale(rowIdx); + + baseGmOffset += this->numLastDim; + } + } + +private: + __aicore__ inline void ComputeDynamicQuant(int32_t rowGmOffset, int32_t elementCount) + { + LocalTensor xLocalFp32 = xBufFp32.Get(); + LocalTensor yLocalFp32 = yBufFp32.Get(); + LocalTensor y12Local = outRowQue.template AllocTensor(); + if (this->outQuant1Flag == 1 || this->isOld) { + CopyInSmoothNorm(xLocalFp32, 0, rowGmOffset, elementCount, this->localMax1); + auto y1Local = y12Local[0]; + RoundFloat2IntQuant(y1Local, xLocalFp32, elementCount); + } + if ((this->outQuant2Flag == 1) || this->oldDouble) { + CopyInSmoothNorm(yLocalFp32, this->numLastDim, rowGmOffset, elementCount, this->localMax2); + auto y2Local = y12Local[this->lastDimSliceLen]; + RoundFloat2IntQuant(y2Local, yLocalFp32, elementCount); + } + outRowQue.template EnQue(y12Local); + } + + __aicore__ inline void CopyOutQuant(int32_t baseGmOffset, int32_t rowGmOffset, int32_t elementCount) + { + LocalTensor yOut = outRowQue.template DeQue(); + if (this->isOld || this->outQuant1Flag == 1) { + DataCopyEx(this->y1Gm[baseGmOffset + rowGmOffset], yOut, elementCount); + } + if (this->oldDouble || (this->outQuant2Flag == 1)) { + DataCopyEx(this->y2Gm[baseGmOffset + rowGmOffset], yOut[this->lastDimSliceLen], elementCount); + } + outRowQue.FreeTensor(yOut); + } + + __aicore__ inline void CopyOutScale(int32_t idx) + { + LocalTensor scalesOut = scalesQue.template DeQue(); + if (this->isOld || this->outQuant1Flag == 1) { + DataCopyEx(this->outScale1Gm[idx], scalesOut[0], 1); + } + + if (this->oldDouble || (this->outQuant2Flag == 1)) { + DataCopyEx(this->outScale2Gm[idx], scalesOut[ELEM_PER_BLK_FP32], 1); + } + scalesQue.FreeTensor(scalesOut); + } + + __aicore__ inline void CopyInSmoothNorm( + LocalTensor& dstLocal, int32_t workspaceOffset, int32_t rowGmOffset, int32_t elementCount, + float scaleNum) + { + LocalTensor smoothYLocalIn = inRowsQue.template AllocTensor(); + DataCopyEx(smoothYLocalIn, this->workspaceGm[workspaceOffset + rowGmOffset], elementCount); + inRowsQue.EnQue(smoothYLocalIn); + LocalTensor smoothYLocal = inRowsQue.template DeQue(); + Muls(dstLocal, smoothYLocal, scaleNum, elementCount); + PipeBarrier(); + inRowsQue.FreeTensor(smoothYLocal); + } + + __aicore__ inline void ComputeRmsNormAndSmoothMax(int32_t rowGmOffset, int32_t elementCount, float rstdLocalTemp) + { + CopyInTmpX(rowGmOffset, elementCount, rstdLocalTemp); + CopyInGamma(rowGmOffset, elementCount); + ComputeNormAndSmooth(rowGmOffset, elementCount); + UpdateLocalMax(elementCount); + } + + __aicore__ inline void CopyInTmpX(int32_t rowGmOffset, int32_t elementCount, float rstdLocalTemp) + { + LocalTensor yLocalFp32 = yBufFp32.Get(); + LocalTensor xLocalIn = inRowsQue.template AllocTensor(); + DataCopyEx(xLocalIn, this->workspaceGm[rowGmOffset], elementCount); + inRowsQue.EnQue(xLocalIn); + LocalTensor xLocal = inRowsQue.template DeQue(); + Muls(yLocalFp32, xLocal, rstdLocalTemp, elementCount); + PipeBarrier(); + inRowsQue.FreeTensor(xLocal); + } + + __aicore__ inline void CopyInGamma(int32_t rowGmOffset, int32_t elementCount) + { + LocalTensor zLocalFp32 = zBufFp32.Get(); + LocalTensor gammaLocalIn = inRowsQue.template AllocTensor(); + DataCopyEx(gammaLocalIn, this->gammaGm[rowGmOffset], elementCount); + inRowsQue.EnQue(gammaLocalIn); + LocalTensor gammaLocal = inRowsQue.template DeQue(); + Cast(zLocalFp32, gammaLocal, RoundMode::CAST_NONE, elementCount); // xLocalFp32 <- gammaFp32 + PipeBarrier(); + inRowsQue.FreeTensor(gammaLocal); + } + + __aicore__ inline void CopyInBeta(int32_t rowGmOffset, int32_t elementCount) + { + LocalTensor zLocalFp32 = zBufFp32.Get(); + LocalTensor betaLocalIn = inRowsQue.template AllocTensor(); + DataCopyEx(betaLocalIn, this->betaGm[rowGmOffset], elementCount); + inRowsQue.EnQue(betaLocalIn); + LocalTensor betaLocal = inRowsQue.template DeQue(); + Cast(zLocalFp32, betaLocalIn, RoundMode::CAST_NONE, elementCount); // xLocalFp32 <- betaFp32 + PipeBarrier(); + inRowsQue.FreeTensor(betaLocalIn); + } + + __aicore__ inline void CopyInSmooth(int32_t rowGmOffset, int32_t elementCount) + { + if (this->newSingleFirst || this->newSingleSecond || + (this->isOld && (this->smooth1Exist || this->smooth2Exist))) { + LocalTensor smooth12CopyIn = inRowsQue.template AllocTensor(); + if (this->newSingleFirst || (this->isOld && this->smooth1Exist)) { + LocalTensor smooth1In = smooth12CopyIn[0]; + DataCopyEx(smooth1In, this->smooth1Gm[rowGmOffset], elementCount); + } + + if (this->newSingleSecond || this->oldDouble) { + LocalTensor smooth2In = smooth12CopyIn[this->lastDimSliceLen]; + DataCopyEx(smooth2In, this->smooth2Gm[rowGmOffset], elementCount); + } + inRowsQue.EnQue(smooth12CopyIn); + } + } + + __aicore__ inline void ComputeNormAndSmooth(int32_t rowGmOffset, int32_t elementCount) + { + LocalTensor xLocalFp32 = xBufFp32.Get(); + LocalTensor yLocalFp32 = yBufFp32.Get(); + LocalTensor zLocalFp32 = zBufFp32.Get(); + + Mul(xLocalFp32, yLocalFp32, zLocalFp32, elementCount); // yLocalFp32 <- x * rstd * gamma + PipeBarrier(); + if (this->betaFlag == 1) { + CopyInBeta(rowGmOffset, elementCount); + LocalTensor zLocalFp32 = zBufFp32.Get(); + Add(xLocalFp32, xLocalFp32, zLocalFp32, elementCount); + PipeBarrier(); + } + CopyInSmooth(rowGmOffset, elementCount); + ComputeSmoothWithFlag(yLocalFp32, zLocalFp32, xLocalFp32, rowGmOffset, elementCount); + } + + __aicore__ inline void ComputeSmoothWithFlag( + LocalTensor yLocalFp32, LocalTensor zLocalFp32, LocalTensor xLocalFp32, + int32_t rowGmOffset, int32_t elementCount) + { + if (this->newSingleFirst || this->newSingleSecond || + (this->isOld && (this->smooth1Exist || this->smooth2Exist))) { + LocalTensor smooth12Local = inRowsQue.template DeQue(); + if (this->newSingleFirst || (this->isOld && this->smooth1Exist)) { + LocalTensor smooth1Local = smooth12Local[0]; + Cast(yLocalFp32, smooth1Local, RoundMode::CAST_NONE, elementCount); // yLocalFp32 <- smooth1 + } + if (this->newSingleSecond || this->oldDouble) { + LocalTensor smooth2Local = smooth12Local[this->lastDimSliceLen]; + Cast(zLocalFp32, smooth2Local, RoundMode::CAST_NONE, elementCount); // zLocalFp32 <- smooth2 + } + inRowsQue.FreeTensor(smooth12Local); + PipeBarrier(); + } + if (this->outQuant1Flag == 1 || this->isOld) { + if (this->smooth1Exist) { + Mul(yLocalFp32, xLocalFp32, yLocalFp32, elementCount); // yLocalFp32 <- norm * smooth1 + } else { + Muls(yLocalFp32, xLocalFp32, 1.0f, elementCount); // yLocalFp32 <- norm * smooth1 + } + PipeBarrier(); + CopyOutSmoothNorm(yLocalFp32, 0, rowGmOffset, elementCount); + } + if (this->outQuant2Flag == 1 || this->oldDouble) { + if (this->smooth2Exist) { + Mul(zLocalFp32, xLocalFp32, zLocalFp32, elementCount); // zLocalFp32 <- norm * smooth2 + } else { + Muls(zLocalFp32, xLocalFp32, 1.0f, elementCount); // zLocalFp32 <- norm * smooth2 + } + PipeBarrier(); + CopyOutSmoothNorm(zLocalFp32, this->numLastDim, rowGmOffset, elementCount); + } + } + + __aicore__ inline void UpdateLocalMax(int32_t elementCount) + { + LocalTensor xLocalFp32 = xBufFp32.Get(); + LocalTensor yLocalFp32 = yBufFp32.Get(); + LocalTensor zLocalFp32 = zBufFp32.Get(); + if (this->outQuant2Flag == 1 || this->oldDouble) { + float tmpMax2 = FindSliceMax(zLocalFp32, xLocalFp32, elementCount); + this->localMax2 = (tmpMax2 > this->localMax2) ? tmpMax2 : localMax2; + } + if (this->outQuant1Flag == 1 || (this->isOld)) { + float tmpMax1 = FindSliceMax(yLocalFp32, xLocalFp32, elementCount); + this->localMax1 = (tmpMax1 > this->localMax1) ? tmpMax1 : localMax1; + } + } + + __aicore__ inline float FindSliceMax( + LocalTensor& srcTensor, LocalTensor& tmpTensor, int32_t elementCount) + { + Abs(tmpTensor, srcTensor, elementCount); // tmpLocal <-- |y * smooth| + PipeBarrier(); + ReduceMaxInplace(tmpTensor, elementCount); + PIPE_V_S(); + float maxTemp = tmpTensor.GetValue(0); + return maxTemp; + } + + __aicore__ inline void CopyOutSmoothNorm( + LocalTensor& smoothNormTensor, int32_t workspaceOffset, int32_t rowGmOffset, int32_t elementCount) + { + LocalTensor ySmoothLocal = tmpOutQue.template AllocTensor(); + Adds(ySmoothLocal, smoothNormTensor, ZERO, elementCount); + tmpOutQue.template EnQue(ySmoothLocal); + LocalTensor ySmooth = tmpOutQue.template DeQue(); + DataCopyEx(this->workspaceGm[workspaceOffset + rowGmOffset], ySmooth, elementCount); + tmpOutQue.FreeTensor(ySmooth); + } + + __aicore__ inline void CopyInX(int32_t baseGmOffset, int32_t rowGmOffset, int32_t elementCount) + { + LocalTensor xLocalIn = inRowsQue.template AllocTensor(); + DataCopyEx(xLocalIn[0], this->xGm[baseGmOffset + rowGmOffset], elementCount); + inRowsQue.EnQue(xLocalIn); + } + + __aicore__ inline void CopyOutX(int32_t baseGmOffset, int32_t rowGmOffset, int32_t elementCount) + { + LocalTensor x = outRowQue.template DeQue(); + DataCopyEx(this->xGm[baseGmOffset + rowGmOffset], x, elementCount); + outRowQue.FreeTensor(x); + LocalTensor xFp32 = tmpOutQue.template DeQue(); + DataCopyEx(this->workspaceGm[rowGmOffset], xFp32, elementCount); + tmpOutQue.FreeTensor(xFp32); + } + + __aicore__ inline float ReduceSquareSumSlice(int32_t elementCount) + { + LocalTensor xLocalFp32 = xBufFp32.Get(); + LocalTensor xInputLocal = inRowsQue.template DeQue(); + LocalTensor yLocalFp32 = yBufFp32.Get(); + + Cast(xLocalFp32, xInputLocal, RoundMode::CAST_NONE, elementCount); + Mul(yLocalFp32, xLocalFp32, xLocalFp32, elementCount); // yLocalFp32 <- x ** 2 + inRowsQue.FreeTensor(xInputLocal); + PipeBarrier(); + return ReduceSumHalfInterval(yLocalFp32, elementCount); // aveLocalTemp <-- E(x**2) + } + + __aicore__ inline void PIPE_MTE3_MTE2() + { + event_t eventMTE3MTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventMTE3MTE2); + WaitFlag(eventMTE3MTE2); + } + + __aicore__ inline void PIPE_S_V() + { + event_t eventSV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventSV); + WaitFlag(eventSV); + } + + __aicore__ inline void PIPE_V_S() + { + event_t eventVS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventVS); + WaitFlag(eventVS); + } + +private: + TPipe* Ppipe = nullptr; + GlobalTensor workspaceGm; + + TQue inRowsQue; + TQue outRowQue; + TQue tmpOutQue; + TQue scalesQue; + + TBuf xBufFp32; + TBuf yBufFp32; + TBuf zBufFp32; + TBuf reduceBuf; + + float localMax1; + float localMax2; + float localSum; +}; + +#endif // __ADD_RMS_NORM_DYNAMIC_QUANT_SLICE_D_H_ diff --git a/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_helper.h b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_helper.h new file mode 100644 index 000000000..265009f34 --- /dev/null +++ b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_helper.h @@ -0,0 +1,191 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_dynamic_quant_helper.h + * \brief + */ + +#ifndef ADD_RMS_NORM_DYNAMIC_QUANT_HELPER_H_ +#define ADD_RMS_NORM_DYNAMIC_QUANT_HELPER_H_ + +#include "reduce_common.h" +#if __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) +#include "impl/dav_c220/kernel_operator_reg_others_impl.h" +#endif + +using namespace AscendC; +constexpr uint32_t FLOAT_BLOCK_ELEM = 8; +constexpr int32_t ROW_FACTOR = 128; +constexpr uint32_t ELEM_PER_BLK_FP16 = 16; +constexpr float DYNAMIC_QUANT_DIVIDEND = 127.0; +constexpr float DYNAMIC_QUANT_DIVIDEND_INT4 = 7.0; + +template +struct integral_constant { + static constexpr Tp value = v; +}; +using true_type = integral_constant; +using false_type = integral_constant; +template +struct is_same : public false_type {}; +template +struct is_same : public true_type {}; + +__aicore__ inline uint32_t CEIL_DIV(uint32_t x, uint32_t y) +{ + if (y > 0) { + return (x + y - 1) / y; + } + return 0; +} + +__aicore__ inline uint32_t ROUND_UP32(uint32_t x) +{ + return (x + ONE_BLK_SIZE - 1) / ONE_BLK_SIZE * ONE_BLK_SIZE; +} + +__aicore__ inline uint32_t TWO_NUMS_MIN(uint32_t x, uint32_t y) +{ + return x < y ? x : y; +} + +__aicore__ inline uint32_t TWO_NUMS_MAX(uint32_t x, uint32_t y) +{ + return x > y ? x : y; +} + +template +__aicore__ inline uint32_t CalculateBlockLen(const uint32_t len) +{ + if constexpr (std::is_same_v) { + return len / 2; + } else { + return len * sizeof(T); + } +} + +template typename R, template typename S> +__aicore__ inline void DataCopyEx( + const R& dst, const S& src, const uint32_t len, const uint32_t count = 1, const bool ubAligned = false) +{ + DataCopyExtParams copyParams; + copyParams.blockCount = count; + copyParams.blockLen = CalculateBlockLen(len); + + if constexpr (is_same, AscendC::LocalTensor>::value) { + copyParams.srcStride = 0; + copyParams.dstStride = (ubAligned) ? 1 : 0; + DataCopyPad(dst, src, copyParams, {}); + } else { + copyParams.srcStride = (ubAligned) ? 1 : 0; + copyParams.dstStride = 0; + DataCopyPad(dst, src, copyParams); + } +} + +template typename R, template typename S> +__aicore__ inline void DataCopyExStride( + const R& dst, const S& src, const uint32_t len, const uint32_t count = 1, const uint32_t ubAligned = 0) +{ + DataCopyExtParams copyParams; + copyParams.blockCount = count; + copyParams.blockLen = CalculateBlockLen(len); + + if constexpr (is_same, AscendC::LocalTensor>::value) { + copyParams.srcStride = 0; + copyParams.dstStride = ubAligned; + DataCopyPad(dst, src, copyParams, {}); + } else { + copyParams.srcStride = ubAligned; + copyParams.dstStride = 0; + DataCopyPad(dst, src, copyParams); + } +} + +/* + * only support count in (128, 255 * 64) + * about 20us faster than above in case fp16:(1024, 11264) on 910B + */ +__aicore__ inline void ReduceMaxInplace(const LocalTensor& srcLocal, int32_t count) +{ + uint64_t repsFp32 = count >> 6; // 6 is count / ELEM_PER_REP_FP32 + uint64_t offsetsFp32 = repsFp32 << 6; // 6 is repsFp32 * ELEM_PER_REP_FP32 + uint64_t remsFp32 = count & 0x3f; // 0x3f 63, count % ELEM_PER_REP_FP32 + + if (likely(repsFp32 > 1)) { + // 8 is rep stride + Max(srcLocal, srcLocal[ELEM_PER_REP_FP32], srcLocal, ELEM_PER_REP_FP32, repsFp32 - 1, {1, 1, 1, 0, 8, 0}); + PipeBarrier(); + } + if (unlikely(remsFp32 > 0)) { + Max(srcLocal, srcLocal[offsetsFp32], srcLocal, remsFp32, 1, {1, 1, 1, 0, 8, 0}); + PipeBarrier(); + } + uint32_t mask = (repsFp32 > 0) ? ELEM_PER_REP_FP32 : count; + // 8 is rep stride + WholeReduceMax(srcLocal, srcLocal, mask, 1, 8, 1, 8); + PipeBarrier(); +} + +/* + * only support count in (128, 255 * 64) + * about 6us slower than above in case fp16:(1024, 11264) on 910B + */ +__aicore__ inline void ReduceSumInplace(const LocalTensor& srcLocal, int32_t count) +{ + uint64_t repsFp32 = count >> 6; // 6 is count / ELEM_PER_REP_FP32 + uint64_t offsetsFp32 = repsFp32 << 6; // 6 is repsFp32 * ELEM_PER_REP_FP32 + uint64_t remsFp32 = count & 0x3f; // 0x3f 63, count % ELEM_PER_REP_FP32 + + if (likely(repsFp32 > 1)) { + // 8 is rep stride + Add(srcLocal, srcLocal[ELEM_PER_REP_FP32], srcLocal, ELEM_PER_REP_FP32, repsFp32 - 1, {1, 1, 1, 0, 8, 0}); + PipeBarrier(); + } + if (unlikely(remsFp32 > 0)) { + Add(srcLocal, srcLocal[offsetsFp32], srcLocal, remsFp32, 1, {1, 1, 1, 0, 8, 0}); + PipeBarrier(); + } + uint32_t mask = (repsFp32 > 0) ? ELEM_PER_REP_FP32 : count; + // 8 is rep stride + WholeReduceSum(srcLocal, srcLocal, mask, 1, 8, 1, 8); + PipeBarrier(); +} + +__aicore__ inline void DivScalarFP32( + LocalTensor& dstTensor, LocalTensor& dividendTensor, LocalTensor& tmpTensor, + float divisorScalar, uint32_t count) +{ + uint32_t repsFp32 = count >> 6; // 6 is divide 64 + uint32_t offsetsFp32 = count & 0xffffffc0; // 0xffffffc0 is floor by 64 + uint32_t remsFp32 = count & 0x3f; // 0x3f is mod(64) + Duplicate(tmpTensor, divisorScalar, FLOAT_BLOCK_ELEM); // FLOAT_BLOCK_ELEM); + PipeBarrier(); + Div(dstTensor, dividendTensor, tmpTensor, ELEM_PER_REP_FP32, repsFp32, {1, 1, 0, 8, 8, 0}); + if ((remsFp32 > 0)) { + Div(dstTensor[offsetsFp32], dividendTensor[offsetsFp32], tmpTensor, remsFp32, 1, {1, 1, 0, 8, 8, 0}); + } + PipeBarrier(); +} + +template +__aicore__ inline void RoundFloat2IntQuant(LocalTensor& dstTensor, LocalTensor& srcTensor, int32_t size) +{ + Cast(srcTensor.ReinterpretCast(), srcTensor, RoundMode::CAST_RINT, size); + PipeBarrier(); + SetDeqScale((half)1.000000e+00f); + PipeBarrier(); + Cast(srcTensor.ReinterpretCast(), srcTensor.ReinterpretCast(), RoundMode::CAST_NONE, size); + PipeBarrier(); + Cast(dstTensor, srcTensor.ReinterpretCast(), RoundMode::CAST_TRUNC, size); + PipeBarrier(); +} +#endif // __ADD_RMS_NORM_DYNAMIC_QUANT_HELPER_H_ diff --git a/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_normal_kernel.h b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_normal_kernel.h new file mode 100644 index 000000000..c2b5aaba2 --- /dev/null +++ b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_normal_kernel.h @@ -0,0 +1,337 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_dynamic_quant_normal_kernel.h + * \brief + */ + +#ifndef ADD_RMS_NORM_DYNAMIC_QUANT_NORMAL_KERNEL_H_ +#define ADD_RMS_NORM_DYNAMIC_QUANT_NORMAL_KERNEL_H_ + +#include "rms_norm_dynamic_quant_base.h" + +template +class KernelAddRmsNormDynamicQuantNormal : public KernelAddRmsNormDynamicQuantBase { +public: + __aicore__ inline KernelAddRmsNormDynamicQuantNormal(TPipe* pipe) + { + Ppipe = pipe; + } + + __aicore__ inline void Init( + GM_ADDR x, GM_ADDR gamma, GM_ADDR smooth1, GM_ADDR smooth2, GM_ADDR beta, GM_ADDR y1, GM_ADDR y2, + GM_ADDR outScale1, GM_ADDR outScale2, GM_ADDR workspace, const RmsNormDynamicQuantTilingData* tiling) + { + this->InitBaseParams(tiling); + this->InitInGlobalTensors(x, gamma, smooth1, smooth2, beta); + this->InitOutGlobalTensors(y1, y2, outScale1, outScale2); + this->numRowsAligned = (this->rowStep + ELEM_PER_BLK_FP32 - 1) / ELEM_PER_BLK_FP32 * ELEM_PER_BLK_FP32; + this->ubAligned = static_cast((this->numLastDimAligned - this->numLastDim) / ELEM_PER_BLK_FP16); + /* + UB = 3 * this->rowStep * alignedCol * sizeof(T) + + 2 * this->rowStep * alignedCol * sizeof(float) + + Count(gamma,beta,bias) * alignedCol * sizeof(T) + + 512Bytes(256 + reduceOut) + */ + Ppipe->InitBuffer(inRowsQue, BUFFER_NUM, 2 * this->rowStep * this->numLastDimAligned * sizeof(T)); // 2 * D * 2 + Ppipe->InitBuffer(outRowsQue, BUFFER_NUM, 2 * this->rowStep * this->numLastDimAligned * sizeof(T)); // D * 2 + Ppipe->InitBuffer(xBufFp32, this->rowStep * this->numLastDimAligned * sizeof(float)); // D * 4 + Ppipe->InitBuffer(yBufFp32, this->rowStep * this->numLastDimAligned * sizeof(float)); // D * 4 + Ppipe->InitBuffer(weightBuf01, this->numLastDimAligned * sizeof(T)); // D * 2 + Ppipe->InitBuffer(weightBuf02, this->numLastDimAligned * sizeof(T)); // D * 2 + Ppipe->InitBuffer(weightBuf03, this->numLastDimAligned * sizeof(T)); // D * 2 + if (this->betaFlag == 1) { + Ppipe->InitBuffer(weightBuf04, this->numLastDimAligned * sizeof(T)); + } + // 2 dynamic quant operator required 2 scale buffer. + Ppipe->InitBuffer(scalesBuf, 2 * this->numRowsAligned * sizeof(float)); + } + + __aicore__ inline void Process() + { + int32_t rowMoveCnt = CEIL_DIV(this->rowWork, this->rowStep); + CopyInWeights(); + + LocalTensor gammaLocal = weightBuf01.template Get(); + + int32_t gmOffset = 0; + int32_t gmOffsetScale = 0; + int32_t elementCount = this->numLastDimAligned * this->rowStep; + + for (int32_t rowIdx = 0; rowIdx < rowMoveCnt - 1; ++rowIdx) { + CopyInX(gmOffset, this->rowStep, elementCount); + ComputeRmsNorm(this->rowStep, elementCount, gammaLocal); + ComputeDynamicQuant(this->rowStep, elementCount); + CopyOut(gmOffset, gmOffsetScale, this->rowStep); + gmOffset += this->rowStep * this->numLastDim; + gmOffsetScale += this->rowStep; + } + { + elementCount = this->numLastDimAligned * this->rowTail_; + int32_t rowIdx = rowMoveCnt - 1; + CopyInX(gmOffset, this->rowTail_, elementCount); + ComputeRmsNorm(this->rowTail_, elementCount, gammaLocal); + ComputeDynamicQuant(this->rowTail_, elementCount); + CopyOut(gmOffset, gmOffsetScale, this->rowTail_); + } + } + +private: + __aicore__ inline void CopyInX(int32_t gmOffset, int32_t rowCount, int32_t elementCount) + { + LocalTensor xLocalIn = inRowsQue.template AllocTensor(); + DataCopyExStride(xLocalIn, this->xGm[gmOffset], this->numLastDim, rowCount, this->ubAligned); + inRowsQue.EnQue(xLocalIn); + } + + __aicore__ inline void CopyOutY(int32_t gmOffset, int32_t rowCount, int32_t elementCount) + { + PipeBarrier(); + LocalTensor yLocal = xBufFp32.Get(); + LocalTensor yOut = yBufFp32.Get(); + PipeBarrier(); + if constexpr (is_same::value) { + Cast(yOut, yLocal, RoundMode::CAST_NONE, elementCount); + } else { // BF16 + Cast(yOut, yLocal, RoundMode::CAST_RINT, elementCount); + } + PipeBarrier(); + DataCopyExStride(this->xGm[gmOffset], yOut, this->numLastDim, rowCount, this->ubAligned); + PipeBarrier(); + } + + __aicore__ inline void CopyInWeights() + { + LocalTensor gammaLocal = weightBuf01.template Get(); + DataCopyEx(gammaLocal, this->gammaGm, this->numLastDim); + if ((this->isOld && this->smooth1Exist) || this->newSingleFirst) { + LocalTensor smooth1Local = weightBuf02.template Get(); + DataCopyEx(smooth1Local, this->smooth1Gm, this->numLastDim); + } + if (this->oldDouble || this->newSingleSecond) { + LocalTensor smooth2Local = weightBuf03.template Get(); + DataCopyEx(smooth2Local, this->smooth2Gm, this->numLastDim); + } + if (this->betaFlag == 1) { + LocalTensor betaLocal = weightBuf04.template Get(); + DataCopyEx(betaLocal, this->betaGm, this->numLastDim); + } + } + + __aicore__ inline void ComputeRmsNorm(int32_t nums, int32_t elementCount, LocalTensor& gammaLocal) + { + LocalTensor xLocalFp32 = xBufFp32.Get(); // xLocalFp32 <-- x + LocalTensor xInputLocal = inRowsQue.template DeQue(); + LocalTensor yLocalFp32 = yBufFp32.Get(); + Cast(xLocalFp32, xInputLocal, RoundMode::CAST_NONE, elementCount); + + Mul(yLocalFp32, xLocalFp32, xLocalFp32, elementCount); // yLocalFp32 <- x ** 2 + PipeBarrier(); + + // reduce#1 for mean + for (int32_t rid = 0; rid < nums; ++rid) { + auto roundOffset = rid * this->numLastDimAligned; + float squareSumTemp = + ReduceSumHalfInterval(yLocalFp32[roundOffset], this->numLastDim); // aveLocalTemp <-- E(x**2) + float rstdLocalTemp = 1 / sqrt(squareSumTemp * this->aveNum + this->eps); + event_t eventSV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventSV); + WaitFlag(eventSV); + Muls( + xLocalFp32[roundOffset], xLocalFp32[roundOffset], rstdLocalTemp, + this->numLastDim); // xLocalFp32 <- x * rstd + } + PipeBarrier(); + + Cast(yLocalFp32, gammaLocal, RoundMode::CAST_NONE, this->numLastDim); // yLocalFp32 <- gamma + PipeBarrier(); + for (int32_t rid = 0; rid < nums; ++rid) { + auto roundOffset = rid * this->numLastDimAligned; + Mul(xLocalFp32[roundOffset], xLocalFp32[roundOffset], yLocalFp32, + this->numLastDim); // xLocalFp32 <- x * rstd * gamma + PipeBarrier(); + } + + if (this->betaFlag == 1) { + LocalTensor betaLocal = weightBuf04.template Get(); + Cast(yLocalFp32, betaLocal, RoundMode::CAST_NONE, this->numLastDim); // yLocalFp32 <- gamma + for (int32_t rid = 0; rid < nums; ++rid) { + auto roundOffset = rid * this->numLastDimAligned; + PipeBarrier(); + Add(xLocalFp32[roundOffset], xLocalFp32[roundOffset], yLocalFp32, this->numLastDim); + PipeBarrier(); + } + } + inRowsQue.FreeTensor(xInputLocal); + } + + __aicore__ inline void ComputeDynamicQuant(int32_t nums, int32_t elementCount) + { + LocalTensor xLocalFp32 = xBufFp32.Get(); // xLocalFp32 <-- y + LocalTensor scaleLocal = scalesBuf.Get(); + LocalTensor zLocalFp32 = outRowsQue.template AllocTensor(); + LocalTensor outQuant01 = zLocalFp32.ReinterpretCast(); + doQuant1withFlag(scaleLocal, xLocalFp32, outQuant01, nums, elementCount); + doQuant2withFlag(scaleLocal, xLocalFp32, outQuant01, nums, elementCount); + outRowsQue.EnQue(zLocalFp32); + } + + __aicore__ inline void doQuant1withFlag( + LocalTensor scaleLocal, LocalTensor xLocalFp32, LocalTensor outQuant01, int32_t nums, + int32_t elementCount) + { + if (this->outQuant1Flag == 0 && !this->isOld) { + return; + } + LocalTensor tmpFp32 = inRowsQue.template AllocTensor(); + LocalTensor yLocalFp32 = yBufFp32.Get(); + LocalTensor scale1Local = scaleLocal[0]; + if (this->smooth1Exist) { + // compute smooth1 + LocalTensor smooth1Local = weightBuf02.Get(); + LocalTensor smooth1Fp32 = yLocalFp32[(nums - 1) * this->numLastDimAligned]; + Cast(smooth1Fp32, smooth1Local, RoundMode::CAST_NONE, this->numLastDim); + PipeBarrier(); + for (int32_t rid = 0; rid < nums; ++rid) { + Mul(yLocalFp32[rid * this->numLastDimAligned], xLocalFp32[rid * this->numLastDimAligned], smooth1Fp32, + this->numLastDim); // yLocalFp32 <-- y * smooth1 + } + PipeBarrier(); + } else { + for (int32_t rid = 0; rid < nums; ++rid) { + Muls( + yLocalFp32[rid * this->numLastDimAligned], xLocalFp32[rid * this->numLastDimAligned], (float)(1.0), + this->numLastDim); // yLocalFp32 <-- y * 1 + } + PipeBarrier(); + } + ScaleTensor(yLocalFp32, tmpFp32, scale1Local, elementCount, nums); + PipeBarrier(); + Cast(yLocalFp32.ReinterpretCast(), yLocalFp32, RoundMode::CAST_RINT, elementCount); + PipeBarrier(); + SetDeqScale((half)1.000000e+00f); + PipeBarrier(); + Cast( + yLocalFp32.ReinterpretCast(), yLocalFp32.ReinterpretCast(), RoundMode::CAST_NONE, + elementCount); + PipeBarrier(); + Cast(outQuant01, yLocalFp32.ReinterpretCast(), RoundMode::CAST_TRUNC, elementCount); + PipeBarrier(); + inRowsQue.FreeTensor(tmpFp32); + } + + __aicore__ inline void doQuant2withFlag( + LocalTensor scaleLocal, LocalTensor xLocalFp32, LocalTensor outQuant01, int32_t nums, + int32_t elementCount) + { + if (this->outQuant2Flag == 0 && !this->oldDouble) { + return; + } + LocalTensor tmpFp32 = inRowsQue.template AllocTensor(); + LocalTensor scale2Local = scaleLocal[this->numRowsAligned]; + LocalTensor yLocalFp32 = yBufFp32.Get(); + LocalTensor outQuant02 = outQuant01[elementCount]; + if (this->smooth2Exist) { + LocalTensor smooth2Local = weightBuf03.Get(); + Cast(tmpFp32, smooth2Local, RoundMode::CAST_NONE, this->numLastDim); + PipeBarrier(); + for (int32_t rid = 0; rid < nums; ++rid) { + Mul(xLocalFp32[rid * this->numLastDimAligned], xLocalFp32[rid * this->numLastDimAligned], tmpFp32, + this->numLastDim); // yLocalFp32 <-- y * smooth2 + } + PipeBarrier(); + } else { + for (int32_t rid = 0; rid < nums; ++rid) { + Muls( + xLocalFp32[rid * this->numLastDimAligned], xLocalFp32[rid * this->numLastDimAligned], (float)(1.0), + this->numLastDim); // yLocalFp32 <-- y * 1 + } + PipeBarrier(); + } + ScaleTensor(xLocalFp32, tmpFp32, scale2Local, elementCount, nums); + PipeBarrier(); + Cast(xLocalFp32.ReinterpretCast(), xLocalFp32, RoundMode::CAST_RINT, elementCount); + PipeBarrier(); + SetDeqScale((half)1.000000e+00f); + PipeBarrier(); + Cast( + xLocalFp32.ReinterpretCast(), xLocalFp32.ReinterpretCast(), RoundMode::CAST_NONE, + elementCount); + PipeBarrier(); + Cast(outQuant02, xLocalFp32.ReinterpretCast(), RoundMode::CAST_TRUNC, elementCount); + PipeBarrier(); + inRowsQue.FreeTensor(tmpFp32); + } + + __aicore__ inline void CopyOut(int32_t gmOffset, int32_t gmOffsetScale, int32_t rowCount) + { + LocalTensor outY12 = outRowsQue.template DeQue(); + LocalTensor scaleLocal = scalesBuf.Get(); + if (this->isOld || (this->outQuant1Flag == 1)) { + LocalTensor outQuant01 = outY12[0]; + LocalTensor scale1Local = scaleLocal[0]; + DataCopyEx(this->y1Gm[gmOffset], outQuant01, this->numLastDim, rowCount); + DataCopyEx(this->outScale1Gm[gmOffsetScale], scale1Local, rowCount); + } + if (this->oldDouble || (this->outQuant2Flag == 1)) { + LocalTensor outQuant02 = outY12[rowCount * this->numLastDimAligned]; + LocalTensor scale2Local = scaleLocal[this->numRowsAligned]; + DataCopyEx(this->y2Gm[gmOffset], outQuant02, this->numLastDim, rowCount); + DataCopyEx(this->outScale2Gm[gmOffsetScale], scale2Local, rowCount); + } + outRowsQue.FreeTensor(outY12); + } + + __aicore__ inline void ScaleTensor( + LocalTensor& srcTensor, LocalTensor& tmpTensor, LocalTensor& scaleTensor, int32_t size, + int32_t nums) + { + float maxTemp; + float scaleTemp; + event_t eventVS; + event_t eventSV; + Abs(tmpTensor, srcTensor, size); // tmpLocal <-- |y * smooth1| + PipeBarrier(); + for (int32_t rid = 0; rid < nums; ++rid) { + ReduceMaxInplace(tmpTensor[rid * this->numLastDimAligned], this->numLastDim); + eventVS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventVS); + WaitFlag(eventVS); + maxTemp = tmpTensor[rid * this->numLastDimAligned].GetValue(0); // Reduce + scaleTemp = this->quantMaxVal / maxTemp; + scaleTensor.SetValue(rid, 1 / scaleTemp); + eventSV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventSV); + WaitFlag(eventSV); + auto srcSlice = srcTensor[rid * this->numLastDimAligned]; + Muls(srcSlice, srcSlice, scaleTemp, this->numLastDim); + } + } + +private: + TPipe* Ppipe = nullptr; + TQue inRowsQue; + TQue outRowsQue; + + TBuf xBufFp32; + TBuf yBufFp32; + + TBuf weightBuf01; + TBuf weightBuf02; + TBuf weightBuf03; + TBuf weightBuf04; + TBuf scalesBuf; + + uint32_t numRowsAligned; + uint32_t ubAligned; +}; + +#endif // __ADD_RMS_NORM_DYNAMIC_QUANT_NORMAL_KERNEL_H_ diff --git a/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_single_row_kernel.h b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_single_row_kernel.h new file mode 100644 index 000000000..712d3489a --- /dev/null +++ b/csrc/ascend/attention/rms_norm_dynamic_quant/op_kernel/rms_norm_dynamic_quant_single_row_kernel.h @@ -0,0 +1,274 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_dynamic_quant_single_row_kernel.h + * \brief + */ + +#ifndef ADD_RMS_NORM_DYNAMIC_QUANT_SINGLE_ROW_KERNEL_H_ +#define ADD_RMS_NORM_DYNAMIC_QUANT_SINGLE_ROW_KERNEL_H_ + +#include "rms_norm_dynamic_quant_base.h" + +template +class KernelAddRmsNormDynamicQuantSingleRow : public KernelAddRmsNormDynamicQuantBase { +public: + __aicore__ inline KernelAddRmsNormDynamicQuantSingleRow(TPipe* pipe) + { + Ppipe = pipe; + } + + __aicore__ inline void Init( + GM_ADDR x, GM_ADDR gamma, GM_ADDR smooth1, GM_ADDR smooth2, GM_ADDR beta, GM_ADDR y1, GM_ADDR y2, + GM_ADDR outScale1, GM_ADDR outScale2, GM_ADDR workspace, const RmsNormDynamicQuantTilingData* tiling) + { + this->InitBaseParams(tiling); + this->InitInGlobalTensors(x, gamma, smooth1, smooth2, beta); + this->InitOutGlobalTensors(y1, y2, outScale1, outScale2); + + /* + UB = 3 * alignedCol * sizeof(T) + + 2 * alignedCol * sizeof(float) + + Count(bias) * alignedCol * sizeof(T) + + 512Btyes(256 + reduceOut) + */ + Ppipe->InitBuffer(inRowsQue, BUFFER_NUM, 2 * this->numLastDimAligned * sizeof(T)); // 2 * D * 2 + Ppipe->InitBuffer(yQue, BUFFER_NUM, this->numLastDimAligned * sizeof(T)); // D * 2 + + Ppipe->InitBuffer(xBufFp32, this->numLastDimAligned * sizeof(float)); // D * 4 + Ppipe->InitBuffer(yBufFp32, this->numLastDimAligned * sizeof(float)); // D * 4 + Ppipe->InitBuffer(smoothBuf, this->numLastDimAligned * sizeof(T)); // D * 2 + + // 2 dynamic quant operator required 2 scale buffer. + Ppipe->InitBuffer(scalesQue, BUFFER_NUM, 2 * ROW_FACTOR * sizeof(float)); + } + + __aicore__ inline void Process() + { + if ((this->isOld && this->smooth1Exist) || this->newSingleFirst) { + LocalTensor smooth1Local = smoothBuf.template Get(); + DataCopyEx(smooth1Local, this->smooth1Gm, this->numLastDim); + } + + int32_t outLoopCount = this->rowWork / ROW_FACTOR; + int32_t outLoopTail = this->rowWork % ROW_FACTOR; + uint32_t gmOffset = 0; + uint32_t gmOffsetReduce = 0; + + LocalTensor scalesLocalOut; + + for (int32_t loopIdx = 0; loopIdx < outLoopCount; ++loopIdx) { + scalesLocalOut = scalesQue.template AllocTensor(); + for (int32_t innerIdx = 0; innerIdx < ROW_FACTOR; ++innerIdx) { + CopyInXAndGamma(gmOffset); + ComputeRmsNorm(gmOffset); + CopyInSmooth(); + ComputeDynamicQuant(innerIdx, scalesLocalOut, gmOffset); + CopyOut(gmOffset); + gmOffset += this->numLastDim; + } + scalesQue.EnQue(scalesLocalOut); + CopyOutScale(gmOffsetReduce, ROW_FACTOR); + gmOffsetReduce += ROW_FACTOR; + } + { + scalesLocalOut = scalesQue.template AllocTensor(); + for (int32_t innerIdx = 0; innerIdx < outLoopTail; ++innerIdx) { + CopyInXAndGamma(gmOffset); + ComputeRmsNorm(gmOffset); + CopyInSmooth(); + ComputeDynamicQuant(innerIdx, scalesLocalOut, gmOffset); + CopyOut(gmOffset); + gmOffset += this->numLastDim; + } + scalesQue.EnQue(scalesLocalOut); + CopyOutScale(gmOffsetReduce, outLoopTail); + } + } + +private: + __aicore__ inline void ComputeRmsNorm(int32_t gmOffset) + { + LocalTensor xLocalFp32 = xBufFp32.Get(); + LocalTensor iputLocal = inRowsQue.template DeQue(); + LocalTensor yLocalFp32 = yBufFp32.Get(); + LocalTensor yLocalB16 = yBufFp32.Get(); + + Cast(xLocalFp32, iputLocal, RoundMode::CAST_NONE, this->numLastDim); + + Mul(yLocalFp32, xLocalFp32, xLocalFp32, this->numLastDim); // yLocalFp32 <- x ** 2 + PipeBarrier(); + + float squareSumTemp = ReduceSumHalfInterval(yLocalFp32, this->numLastDim); + float rstdLocalTemp = 1 / sqrt(squareSumTemp * this->aveNum + this->eps); + event_t eventSV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventSV); + WaitFlag(eventSV); + Muls(xLocalFp32, xLocalFp32, rstdLocalTemp, this->numLastDim); // xLocalFp32 <- x * rstd + PipeBarrier(); + LocalTensor gammaLocal = xLocalFp32[this->numLastDimAligned]; + + inRowsQue.FreeTensor(iputLocal); + Mul(xLocalFp32, xLocalFp32, gammaLocal, this->numLastDim); // xLocalFp32 <- x * rstd * gamma + PipeBarrier(); + if (this->betaFlag == 1) { + CopyInBeta(); + LocalTensor betaLocal = inRowsQue.template DeQue(); + Cast(yLocalFp32, betaLocal, RoundMode::CAST_NONE, this->numLastDim); // yLocalB16 <- Cast(beta) + PipeBarrier(); + Add(xLocalFp32, xLocalFp32, yLocalFp32, this->numLastDim); + PipeBarrier(); + inRowsQue.FreeTensor(betaLocal); + } + } + + __aicore__ inline void ComputeDynamicQuant(int32_t idx, LocalTensor& scalesLocalOut, int32_t gmOffset) + { + LocalTensor xLocalFp32 = xBufFp32.Get(); + LocalTensor yLocalFp32 = yBufFp32.Get(); + LocalTensor yLocal = yQue.template AllocTensor(); + + LocalTensor smooth1Local = smoothBuf.template Get(); + LocalTensor smooth2Local = inRowsQue.template DeQue(); + LocalTensor tmpTensor = smooth2Local.template ReinterpretCast(); + auto y1Local = yLocal[0]; + auto y2Local = yLocal[this->numLastDimAligned]; + + if ((this->outQuant2Flag == 1) || this->oldDouble) { + if (this->smooth2Exist) { + Cast(yLocalFp32, smooth2Local, RoundMode::CAST_NONE, this->numLastDim); // yLocalFp32 <-- smooth2 + PipeBarrier(); + Mul(yLocalFp32, xLocalFp32, yLocalFp32, this->numLastDim); // yLocalFp32 <-- y * smooth2 + PipeBarrier(); + } else { + Muls(yLocalFp32, xLocalFp32, (float)1.0, this->numLastDim); // yLocalFp32 <-- y * 1 + PipeBarrier(); + } + ScaleTensor( + yLocalFp32, tmpTensor, scalesLocalOut, + idx + ROW_FACTOR); // yLocalFp32 <-- yLocalFp32 / max(abs(yLocalFp32)) + PipeBarrier(); + inRowsQue.FreeTensor(tmpTensor); + RoundFloat2IntQuant(y2Local, yLocalFp32, this->numLastDim); + } + + if ((this->outQuant1Flag == 1) || this->isOld) { + if (this->smooth1Exist) { + Cast(yLocalFp32, smooth1Local, RoundMode::CAST_NONE, this->numLastDim); // yLocalFp32 <-- smooth1 + PipeBarrier(); + Mul(yLocalFp32, xLocalFp32, yLocalFp32, this->numLastDim); // yLocalFp32 <-- y * smooth1 + PipeBarrier(); + } else { + Muls(yLocalFp32, xLocalFp32, (float)1.0, this->numLastDim); // yLocalFp32 <-- y * smooth1 + PipeBarrier(); + } + ScaleTensor( + yLocalFp32, xLocalFp32, scalesLocalOut, idx); // yLocalFp32 <-- yLocalFp32 / max(abs(yLocalFp32)) + PipeBarrier(); + RoundFloat2IntQuant(y1Local, yLocalFp32, this->numLastDim); + } + PipeBarrier(); + yQue.EnQue(yLocal); + } + + // srcTensor <- srcTensor / max(abs(srcTensor)) + __aicore__ inline void ScaleTensor( + LocalTensor& srcTensor, LocalTensor& tmpTensor, LocalTensor& scaleTensor, int32_t idx) + { + Abs(tmpTensor, srcTensor, this->numLastDim); // tmpLocal <-- |y * smooth| + PipeBarrier(); + ReduceMaxInplace(tmpTensor, this->numLastDim); + event_t eventVS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventVS); + WaitFlag(eventVS); + float maxTemp = tmpTensor.GetValue(0); + float scaleTemp = this->quantMaxVal / maxTemp; + scaleTensor.SetValue(idx, 1 / scaleTemp); + event_t eventSV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventSV); + WaitFlag(eventSV); + Muls(srcTensor, srcTensor, scaleTemp, this->numLastDim); + PipeBarrier(); + } + + __aicore__ inline void CopyOut(int32_t gmOffset) + { + LocalTensor res12 = yQue.template DeQue(); + auto res1 = res12[0]; + auto res2 = res12[this->numLastDimAligned]; + if (this->isOld || (this->outQuant1Flag == 1)) { + DataCopyEx(this->y1Gm[gmOffset], res1, this->numLastDim); + } + + if (this->oldDouble || (this->outQuant2Flag == 1)) { + DataCopyEx(this->y2Gm[gmOffset], res2, this->numLastDim); + } + yQue.FreeTensor(res12); + } + + __aicore__ inline void CopyOutScale(int32_t gmOffset, int32_t copyInNums) + { + LocalTensor outScalesLocal = scalesQue.template DeQue(); + LocalTensor outScales1Local = outScalesLocal[0]; + LocalTensor outScales2Local = outScalesLocal[ROW_FACTOR]; + if (this->isOld || (this->outQuant1Flag == 1)) { + DataCopyEx(this->outScale1Gm[gmOffset], outScales1Local, copyInNums); + } + if (this->oldDouble || (this->outQuant2Flag == 1)) { + DataCopyEx(this->outScale2Gm[gmOffset], outScales2Local, copyInNums); + } + scalesQue.FreeTensor(outScalesLocal); + } + + __aicore__ inline void CopyInXAndGamma(int32_t gmOffset) + { + LocalTensor xLocalIn = inRowsQue.template AllocTensor(); + DataCopyEx(xLocalIn[0], this->xGm[gmOffset], this->numLastDim); + DataCopyEx(xLocalIn[this->numLastDimAligned], this->gammaGm, this->numLastDim); + inRowsQue.EnQue(xLocalIn); + } + + __aicore__ inline void CopyInSmooth() + { + if (this->oldDouble || this->newSingleSecond) { + LocalTensor smoothCopyIn = inRowsQue.template AllocTensor(); + DataCopyEx(smoothCopyIn[0], this->smooth2Gm, this->numLastDim); + inRowsQue.EnQue(smoothCopyIn); + } + } + + __aicore__ inline void CopyInGamma() + { + LocalTensor gammaCopyIn = inRowsQue.template AllocTensor(); + DataCopyEx(gammaCopyIn[0], this->gammaGm, this->numLastDim); + inRowsQue.EnQue(gammaCopyIn); + } + + __aicore__ inline void CopyInBeta() + { + LocalTensor betaCopyIn = inRowsQue.template AllocTensor(); + DataCopyEx(betaCopyIn[0], this->betaGm, this->numLastDim); + inRowsQue.EnQue(betaCopyIn); + } + +private: + TPipe* Ppipe = nullptr; + TQue inRowsQue; + TQue yQue; + TQue scalesQue; + + TBuf xBufFp32; + TBuf yBufFp32; + + TBuf smoothBuf; +}; + +#endif // __ADD_RMS_NORM_DYNAMIC_QUANT_SINGLE_ROW_KERNEL_H_ diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/CMakeLists.txt b/csrc/ascend/attention/sparse_attn_sharedkv/CMakeLists.txt new file mode 100644 index 000000000..a3ab34a93 --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/README.md b/csrc/ascend/attention/sparse_attn_sharedkv/README.md new file mode 100644 index 000000000..18519e891 --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/README.md @@ -0,0 +1,324 @@ +# SparseAttnSharedkv + +## 产品支持情况 + +| 产品 | 是否支持 | +| ------------------------------------------------------------ | :------: | +|Ascend 950PR/Ascend 950DT | × | +|Atlas A3 训练系列产品/Atlas A3 推理系列产品 | √ | +|Atlas A2 训练系列产品/Atlas A2 推理系列产品 | √ | +|Atlas 200I/500 A2 推理系列产品 | × | +|Atlas 推理系列产品 | × | +|Atlas 训练系列产品 | × | + +## 功能说明 + +- API功能：`SparseAttnSharedKV`算子旨在完成以下公式描述的Attention计算，支持Sliding Window Attention、Compressed Attention以及Sparse Compressed Attention。 + +- 计算公式： + + $$ + O = \text{softmax}(Q@\tilde{K}^T \cdot \text{softmax\_scale})@\tilde{V} + $$ + + 其中$\tilde{K}=\tilde{V}$为基于ori_kv、cmp_kv以及cmp_ratio等入参控制的实际参与计算的 $KV$。 + +## 参数说明 + +| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 | +|----------------------------|-----------|----------------------------------------------------------------------|----------------|------------| +| q | 输入 | 对应公式中的$Q$。 | BFLOAT16、FLOAT16 | ND | +| ori\_kv | 可选输入 | 对应公式中的$\tilde{K}和\tilde{V}$的一部分，为原始不经压缩的KV。 | BFLOAT16、FLOAT16 | ND | +| cmp\_kv | 可选输入 | 对应公式中的$\tilde{K}和\tilde{V}$的一部分，为经过压缩的KV。 | BFLOAT16、FLOAT16 | ND | +| ori\_sparse\_indices | 可选输入 | 代表离散取oriKvCache的索引。 | INT32 | ND | +| cmp\_sparse\_indices | 可选输入 | 代表离散取cmpKvCache的索引。 | INT32 | ND | +| ori\_block\_table | 可选输入 | 表示PageAttention中oriKvCache存储使用的block映射表。 | INT32 | ND | +| cmp\_block\_table | 可选输入 | 表示PageAttention中cmpKvCache存储使用的block映射表。 | INT32 | ND | +| cu\_seqlens\_q | 可选输入 | 表示不同Batch中`q`的有效token数。 | INT32 | ND | +| cu\_seqlens\_ori\_kv | 可选输入 | 表示不同Batch中`ori_kv`的有效token数。 | INT32 | ND | +| cu\_seqlens\_cmp\_kv | 可选输入 | 表示不同Batch中`cmp_kv`的有效token数。 | INT32 | ND | +| seqused\_q | 可选输入 | 表示不同Batch中`q`实际参与运算的token数。 | INT32 | ND | +| seqused\_kv | 可选输入 | 表示不同Batch中`ori_kv`实际参与运算的token数。 | INT32 | ND | +| sinks | 可选输入 | 注意力下沉tensor。 | FLOAT32 | ND | +| metadata | 可选输入 | aicpu算子（npu\_sparse\_attn\_sharedkv\_metadata）的分核结果。 | INT32 | ND | +| softmax\_scale | 可选属性 | 代表缩放系数，对应公式中的$\text{softmax\_scale}$，默认值为None。 | FLOAT32 | - | +| cmp_ratio | 可选属性 | 表示对`ori_kv`的压缩率，仅支持输入4或128，默认值为None。 | INT32 | - | +| ori\_mask\_mode | 可选属性 | 表示`q`和`ori_kv`计算的mask模式，仅支持输入默认值4。 | INT32 | - | +| cmp\_mask\_mode | 可选属性 | 表示`q`和`cmp_kv`计算的mask模式，仅支持输入默认值3。 | INT32 | - | +| ori\_win\_left | 可选属性 | 表示`q`和`ori_kv`计算中q对过去token计算的数量，仅支持输入默认值127。 | INT32 | - | +| ori\_win\_right | 可选属性 | 表示`q`和`ori_kv`计算中q对未来token计算的数量，仅支持输入默认值0。 | INT32 | - | +| layout\_q | 可选属性 | 用于标识输入`q`的数据排布格式，支持输入"TND"和"BSND"，默认值为"BSND"。 | STRING | - | +| layout\_kv | 可选属性 | 用于标识输入`ori_kv`和`cmp_kv`的数据排布格式，支持输入"PA_ND"和"BSND"。 | STRING | - | +| return\_softmax_lse | 可选属性 | 表示是否返回`softmax_lse`。True表示返回，False表示不返回，默认值为False。 | BOOL | - | +| attention\_out | 输出 | 公式中的输出。 | BFLOAT16、FLOAT16 | ND | +| softmax\_lse | 输出 | 返回的`softmax_lse`。 | FLOAT32 | ND | + +## 约束说明 + +- 该接口支持推理场景下使用。 +- 该接口支持aclgraph模式。 +- 该接口当前支持三种计算场景：场景一，仅传入`ori_kv`时为Sliding Window Attention计算；场景二，传入`ori_kv`及`cmp_kv`时为Sliding Window Attention + Compressed Attention计算；场景三，传入`ori_kv`、`cmp_kv`及`cmp_sparse_indices`时为Sliding Window Attention + Sparse Compressed Attention计算。 + +- 当`layout_q`为TND时，功能使用限制如下： + - `q`的shape需要为[T1,N1,D]，其中N1仅支持64。 + - `ori_sparse_indices`的shape需要为[Q\_T, KV\_N, K1]，其中K1为对`ori_kv`一次离散选取的token数，K1仅支持512。 + - `cmp_sparse_indices`的shape需要为[Q\_T, KV\_N, K2]，其中K2为对`cmp_kv`一次离散选取的token数，K2仅支持512。 + - `cu_seqlens_q`必须传入，输入维度为B+1，大小为参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须>=前一个元素的值。 + +- 当`layout_q`为BSND时，功能使用限制如下： + - `q`的shape需要为[B, Q\_S,N1,D]，其中N1仅支持64。 + - `ori_sparse_indices`的shape需要为[B, Q\_S, KV\_N, K1]，其中K1为对`ori_kv`一次离散选取的token数，K1仅支持512。 + - `cmp_sparse_indices`的shape需要为[B, Q\_S, KV\_N, K2]，其中K2为对`cmp_kv`一次离散选取的token数，K2仅支持512。 + +- PageAttention场景下，功能使用限制如下： + - `ori_kv`和`cmp_kv`的shape分别为[ori\_block\_num, ori\_block\_size, KV\_N, D]和[cmp\_block\_num, cmp\_block\_size, KV\_N, D]，其中ori\_block\_num和cmp\_block\_num为PageAttention时block总数，ori\_block\_size和cmp\_block\_size为一个block的token数，ori\_block\_size和cmp\_block\_size取值为16的倍数，最大支持1024，KV_N仅支持1。 + - `ori_block_table`和`cmp_block_table`的shape为2维，其中第一维长度为B，第二维长度不小于所有batch中最大的S2和S3对应的block数量，即S2\_max / block\_size和S3\_max / block\_size向上取整。 +- `metadata`为算子实际需要使用的分核结果，目前该参数必传，shape大小固定为[1024]。 +- `layout_kv`仅支持输入PA_ND，故设置`cu_seqlens_ori_kv`和`cu_seqlens_cmp_kv`无效。 +- 目前暂不支持返回`softmax_lse`，`return_softmax_lse`仅支持输入False，返回值`softmax_lse`为无效值。 +- ori_mask_mode及cmp_mask_mode所表示的mask模式的详细介绍见[sparse_mode参数说明](../../../docs/zh/context/sparse_mode参数说明.md)。 +- 目前暂不支持指定`q`中参与运算的token数，因此设置`seqused_q`无效。 +- 目前暂不支持对`ori_kv`进行稀疏计算，因此设置`ori_sparse_indices`无效。 +- 目前所有输入不支持传入空tensor。 +- `q`、`ori_kv`、`cmp_kv`数据排布格式支持从多种维度解读，B（Batch）表示输入样本批量大小、S（Seq-Length）表示输入样本序列长度、H（Hidden-Size）表示隐藏层的大小、N（Head-Num）表示多头数、D（Head-Dim）表示hidden层最小的单元尺寸，且满足D=H/N、T表示所有Batch输入样本序列长度的累加和。 +- Q\_S和S1表示q shape中的S，S2表示ori_kv shape中的S，S3表示cmp_kv shape中的S；Q\_N和N1表示num\_q\_heads，KV\_N和N2表示num\_ori_kv\_heads和num\_cmp_kv\_heads；Q\_T和T1表示q shape中的输入样本序列长度的累加和。 + +- 当`layout_kv`为BSND时，功能使用限制如下： + - `ori_kv`和`cmp_kv`的layout都必须为BSND，ori_kv的shape为[B, S2, N2,D]，cmp_kv的shape为[B, S3, N2,D]。 + +## Atlas A3 推理系列产品 调用说明 + +- 单算子模式调用 + + ```python + import torch + import torch_npu + import numpy as np + import random + import math + import custom_ops + + data_type = torch.bfloat16 + softmax_scale = 0.041666666666666664 + b = 4 + s1 = 128 + s2 = 8192 + n1 = 64 + n2 = 1 + dn = 512 + k = 512 + ori_block_size = 128 + cmp_block_size = 128 + s2_act = 4096 + cmp_ratio = 4 + ori_win_left = 127 + ori_win_right = 0 + layout_q = 'TND' + layout_kv = 'PA_ND' + ori_mask_mode = 4 + cmp_mask_mode = 3 + q = torch.tensor(np.random.uniform(-10, 10, (b*s1, n1, dn))).to(data_type).npu() + + cu_seqlens_q = torch.arange(0, (b + 1) * s1, step=s1).to(torch.int32).npu() + t = cu_seqlens_q[-1].item() + seqused_kv = torch.tensor([s2_act]*b).to(torch.int32).npu() + + cmp_kv_len = s2_act // cmp_ratio + idxs = random.sample(range(cmp_kv_len - s1 + 1), k) + cmp_sparse_indices = torch.tensor([idxs for _ in range(t * n2)]).reshape(t, n2, k). \ + to(torch.int32).npu() + + ori_block_num = math.ceil(s2_act/ori_block_size) * b + ori_block_table = torch.tensor(np.random.permutation(range(ori_block_num))).to(torch.int32).reshape(b, -1).npu() + ori_kv = torch.tensor(np.random.uniform(-5, 10, (ori_block_num, ori_block_size, n2, dn))).to(data_type).npu() + + block_num2 = math.ceil(cmp_kv_len/ori_block_size) * b + cmp_block_table = torch.tensor(np.random.permutation(range(block_num2))).to(torch.int32).reshape(b, -1).npu() + cmp_kv = torch.tensor(np.random.uniform(-5, 10, (block_num2, cmp_block_size, n2, dn))).to(data_type).npu() + sinks = torch.rand(n1).to(torch.float32).npu() + metadata = torch.ops.custom.npu_sparse_attn_sharedkv_metadata( + num_heads_q=n1, + num_heads_kv=n2, + head_dim=dn, + cu_seqlens_q=cu_seqlens_q, + seqused_kv=seqused_kv, + batch_size=b, + max_seqlen_q=s1, + max_seqlen_kv=s2, + cmp_topk=k, + cmp_ratio=cmp_ratio, + ori_mask_mode=ori_mask_mode, + cmp_mask_mode=cmp_mask_mode, + ori_win_left=ori_win_left, + ori_win_right=ori_win_right, + layout_q=layout_q, + layout_kv=layout_kv, + has_ori_kv=True, + has_cmp_kv=True + ) + attn_out, softmax_lse = torch.ops.custom.npu_sparse_attn_sharedkv( + q, + ori_kv=ori_kv, + cmp_kv=cmp_kv, + ori_sparse_indices=None, + cmp_sparse_indices=cmp_sparse_indices, + ori_block_table=ori_block_table, + cmp_block_table=cmp_block_table, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_ori_kv=None, + cu_seqlens_cmp_kv=None, + seqused_q=None, + seqused_kv=seqused_kv, + sinks=sinks, + metadata=metadata, + softmax_scale=softmax_scale, + cmp_ratio=cmp_ratio, + ori_mask_mode=ori_mask_mode, + cmp_mask_mode=cmp_mask_mode, + ori_win_left=ori_win_left, + ori_win_right=ori_win_right, + layout_q=layout_q, + layout_kv=layout_kv, + return_softmax_lse=False) + ``` + +- aclgraph模式调用 + + ```python + import torch + import torch_npu + import numpy as np + import random + import math + import npugraph_ex as nge + import custom_ops + + data_type = torch.bfloat16 + softmax_scale = 0.041666666666666664 + b = 4 + s1 = 128 + s2 = 8192 + n1 = 64 + n2 = 1 + dn = 512 + k = 512 + ori_block_size = 128 + cmp_block_size = 128 + s2_act = 4096 + cmp_ratio = 4 + ori_win_left = 127 + ori_win_right = 0 + layout_q = 'TND' + layout_kv = 'PA_ND' + ori_mask_mode = 4 + cmp_mask_mode = 3 + q = torch.tensor(np.random.uniform(-10, 10, (b*s1, n1, dn))).to(data_type).npu() + + cu_seqlens_q = torch.arange(0, (b + 1) * s1, step=s1).to(torch.int32).npu() + t = cu_seqlens_q[-1].item() + seqused_kv = torch.tensor([s2_act]*b).to(torch.int32).npu() + + cmp_kv_len = s2_act // cmp_ratio + idxs = random.sample(range(cmp_kv_len - s1 + 1), k) + cmp_sparse_indices = torch.tensor([idxs for _ in range(t * n2)]).reshape(t, n2, k). \ + to(torch.int32).npu() + + ori_block_num = math.ceil(s2_act/ori_block_size) * b + ori_block_table = torch.tensor(np.random.permutation(range(ori_block_num))).to(torch.int32).reshape(b, -1).npu() + ori_kv = torch.tensor(np.random.uniform(-5, 10, (ori_block_num, ori_block_size, n2, dn))).to(data_type).npu() + + block_num2 = math.ceil(cmp_kv_len/ori_block_size) * b + cmp_block_table = torch.tensor(np.random.permutation(range(block_num2))).to(torch.int32).reshape(b, -1).npu() + cmp_kv = torch.tensor(np.random.uniform(-5, 10, (block_num2, cmp_block_size, n2, dn))).to(data_type).npu() + sinks = torch.rand(n1).to(torch.float32).npu() + + from npugraph_ex.configs.compiler_config import CompilerConfig + config = CompilerConfig() + npu_backend = nge.get_npu_backend(compiler_config=config) + + class Network(torch.nn.Module): + def __init__(self): + super(Network, self).__init__() + + def forward(self, num_heads_q, num_heads_kv, head_dim, batch_size, max_seqlen_q, max_seqlen_kv, + topk, has_ori_kv, has_cmp_kv, q, ori_kv, cmp_kv, cmp_sparse_indices, ori_block_table, + cmp_block_table, cu_seqlens_q, seqused_kv, softmax_scale, cmp_ratio, sinks, + ori_mask_mode, cmp_mask_mode, ori_win_left, ori_win_right, layout_q, layout_kv): + metadata = torch.ops.custom.npu_sparse_attn_sharedkv_metadata( + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + head_dim=head_dim, + cu_seqlens_q=cu_seqlens_q, + seqused_kv=seqused_kv, + batch_size=batch_size, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + cmp_topk=topk, + cmp_ratio=cmp_ratio, + ori_mask_mode=ori_mask_mode, + cmp_mask_mode=cmp_mask_mode, + ori_win_left=ori_win_left, + ori_win_right=ori_win_right, + layout_q=layout_q, + layout_kv=layout_kv, + has_ori_kv=has_ori_kv, + has_cmp_kv=has_cmp_kv, + device="npu:0" + ) + npu_out = torch.ops.custom.npu_sparse_attn_sharedkv( + q, + ori_kv=ori_kv, + cmp_kv=cmp_kv, + ori_sparse_indices=None, + cmp_sparse_indices=cmp_sparse_indices, + ori_block_table=ori_block_table, + cmp_block_table=cmp_block_table, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_ori_kv=None, + cu_seqlens_cmp_kv=None, + seqused_q=None, + seqused_kv=seqused_kv, + sinks=sinks, + metadata=metadata, + softmax_scale=softmax_scale, + cmp_ratio=cmp_ratio, + ori_mask_mode=ori_mask_mode, + cmp_mask_mode=cmp_mask_mode, + ori_win_left=ori_win_left, + ori_win_right=ori_win_right, + layout_q=layout_q, + layout_kv=layout_kv, + return_softmax_lse=False) + return npu_out + + mod = torch.compile(Network().npu(), backend=npu_backend, fullgraph=True) + attn_out, softmax_lse = mod( + num_heads_q=n1, + num_heads_kv=n2, + head_dim=dn, + batch_size=b, + max_seqlen_q=s1, + max_seqlen_kv=s2, + topk=k, + has_ori_kv=True, + has_cmp_kv=True, + q=q, + ori_kv=ori_kv, + cmp_kv=cmp_kv, + cmp_sparse_indices=cmp_sparse_indices, + ori_block_table=ori_block_table, + cmp_block_table=cmp_block_table, + cu_seqlens_q=cu_seqlens_q, + seqused_kv=seqused_kv, + softmax_scale=softmax_scale, + cmp_ratio=cmp_ratio, + sinks=sinks, + ori_mask_mode=ori_mask_mode, + cmp_mask_mode=cmp_mask_mode, + ori_win_left=ori_win_left, + ori_win_right=ori_win_right, + layout_q=layout_q, + layout_kv=layout_kv) + ``` + +更多使用示例见[pytest示例](./tests/pytest/README.md)。 diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_host/CMakeLists.txt b/csrc/ascend/attention/sparse_attn_sharedkv/op_host/CMakeLists.txt new file mode 100644 index 000000000..98d12adb3 --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_host/CMakeLists.txt @@ -0,0 +1,29 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + sparse_attn_sharedkv_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME SparseAttnSharedkv + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-vf-remove-membar=false + -mllvm -cce-aicore-hoist-movemask=false +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE sparse_attn_sharedkv ACLNNTYPE aclnn) +endif() \ No newline at end of file diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_def.cpp b/csrc/ascend/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_def.cpp new file mode 100644 index 000000000..d098493c0 --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_def.cpp @@ -0,0 +1,126 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_def.cpp + * \brief + */ + +#include "register/op_def_registry.h" + +namespace ops { +class SparseAttnSharedkv : public OpDef { +public: + explicit SparseAttnSharedkv(const char *name) : OpDef(name) + { + this->Input("q") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("ori_kv") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .IgnoreContiguous(); + this->Input("cmp_kv") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .IgnoreContiguous(); + this->Input("ori_sparse_indices") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("cmp_sparse_indices") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("ori_block_table") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("cmp_block_table") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("cu_seqlens_q") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("cu_seqlens_ori_kv") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("cu_seqlens_cmp_kv") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("seqused_q") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("seqused_kv") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("sinks") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("metadata") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("attn_out") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("softmax_lse") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Attr("softmax_scale").AttrType(REQUIRED).Float(1.0); + this->Attr("cmp_ratio").AttrType(REQUIRED).Int(1); + this->Attr("ori_mask_mode").AttrType(REQUIRED).Int(4); // ori_mask_mode默认值4 + this->Attr("cmp_mask_mode").AttrType(REQUIRED).Int(3); // cmp_mask_mode默认值3 + this->Attr("ori_kv_stride").AttrType(REQUIRED).Int(0); // ori_mask_mode默认值4 + this->Attr("cmp_kv_stride").AttrType(REQUIRED).Int(0); // cmp_mask_mode默认值3 + this->Attr("ori_win_left").AttrType(OPTIONAL).Int(127); // ori_win_left默认值127 + this->Attr("ori_win_right").AttrType(OPTIONAL).Int(0); + this->Attr("layout_q").AttrType(OPTIONAL).String("BSND"); + this->Attr("layout_kv").AttrType(OPTIONAL).String("PA_ND"); + this->Attr("return_softmax_lse").AttrType(OPTIONAL).Bool(false); + + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("aclnnSupport.value", "support_aclnn"); + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + } +}; +OP_ADD(SparseAttnSharedkv); +} // namespace ops diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_proto.cpp b/csrc/ascend/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_proto.cpp new file mode 100644 index 000000000..f899dae28 --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_proto.cpp @@ -0,0 +1,62 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_proto.cpp + * \brief + */ + +#include +#include +#include "error/ops_error.h" + +using namespace ge; + +namespace ops { +constexpr uint32_t QUERY_INPUT_INDEX = 0; +constexpr uint32_t RETURN_SOFTMAX_INDEX = 8; + +ge::graphStatus InferShapeSparseAttnSharedkv(gert::InferShapeContext *context) +{ + OPS_ERR_IF(context == nullptr, OPS_LOG_E("SparseAttnSharedkv", "InferShapeContext is nullptr"), + return ge::GRAPH_FAILED); + const gert::Shape *queryShape = context->GetInputShape(QUERY_INPUT_INDEX); + OPS_LOG_E_IF_NULL(context, queryShape, return ge::GRAPH_FAILED) + gert::Shape *attentionOutShape = context->GetOutputShape(0); + OPS_LOG_E_IF_NULL(context, attentionOutShape, return ge::GRAPH_FAILED) + *attentionOutShape = *queryShape; + + gert::Shape *softmaxLseShape = context->GetOutputShape(1); + OPS_LOG_E_IF_NULL(context, attentionOutShape, return ge::GRAPH_FAILED) + auto attr = context->GetAttrs(); + const bool *returnSoftmaxLsePtr = attr->GetAttrPointer(RETURN_SOFTMAX_INDEX); + bool returnSoftmaxLse = (returnSoftmaxLsePtr != nullptr) ? *returnSoftmaxLsePtr : false; + if (returnSoftmaxLse) { + *softmaxLseShape = *queryShape; + auto lastDimIdx = softmaxLseShape->GetDimNum() - 1; + softmaxLseShape->SetDim(lastDimIdx, 1); + } else { + softmaxLseShape->SetDimNum(1); + softmaxLseShape->SetDim(0, 0); + } + return GRAPH_SUCCESS; +} + +ge::graphStatus InferDataTypeSparseFlashAttention(gert::InferDataTypeContext *context) +{ + OPS_ERR_IF(context == nullptr, OPS_LOG_E("SparseAttnSharedkv", "InferShapeContext is nullptr"), + return ge::GRAPH_FAILED); + const auto inputDataType = context->GetInputDataType(QUERY_INPUT_INDEX); + context->SetOutputDataType(0, inputDataType); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP(SparseAttnSharedkv).InferShape(InferShapeSparseAttnSharedkv).InferDataType(InferDataTypeSparseAttnSharedkv); +} // namespace ops diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_tiling.cpp b/csrc/ascend/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_tiling.cpp new file mode 100644 index 000000000..746f8bca3 --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_tiling.cpp @@ -0,0 +1,1654 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_tiling.cpp + * \brief + */ + +#include "sparse_attn_sharedkv_tiling.h" +#include "../op_kernel/sparse_attn_sharedkv_template_tiling_key.h" + +using namespace ge; +using namespace AscendC; +using std::map; +using std::string; +using std::pair; +namespace optiling { + +static const std::string QUERY_NAME = "query"; +static const std::string ORI_KV_NAME = "ori_kv"; +static const std::string CMP_KV_NAME = "cmp_kv"; +static const std::string ORI_SPARSE_INDICES = "ori_sparse_indices"; +static const std::string CMP_SPARSE_INDICES = "cmp_sparse_indices"; +static const std::string ORI_BLOCK_TABLE_NAME = "ori_block_table"; +static const std::string CMP_BLOCK_TABLE_NAME = "cmp_block_table"; +static const std::string SINKS_NAME = "sinks"; +static const std::string METADATA_NAME = "metadata"; +static const std::string ATTEN_OUT_NAME = "attn_out"; +const std::map> DTYPE_SUPPORT_MAP = { + {QUERY_NAME, {ge::DT_FLOAT16, ge::DT_BF16}}, + {ORI_KV_NAME, {ge::DT_FLOAT16, ge::DT_BF16}}, + {CMP_KV_NAME, {ge::DT_FLOAT16, ge::DT_BF16}}, + {ORI_SPARSE_INDICES, {ge::DT_INT32}}, + {CMP_SPARSE_INDICES, {ge::DT_INT32}}, + {ATTEN_OUT_NAME, {ge::DT_FLOAT16, ge::DT_BF16}}, + {ORI_BLOCK_TABLE_NAME, {ge::DT_INT32}}, + {CMP_BLOCK_TABLE_NAME, {ge::DT_INT32}}, + {SINKS_NAME, {ge::DT_FLOAT}}, + {METADATA_NAME, {ge::DT_INT32}} +}; + +const std::map> LAYOUT_SUPPORT_MAP = { + {QUERY_NAME, {SASLayout::BSND, SASLayout::TND}}, + {ORI_KV_NAME, {SASLayout::PA_ND, SASLayout::BSND, SASLayout::TND}}, + {CMP_KV_NAME, {SASLayout::PA_ND, SASLayout::BSND, SASLayout::TND}}, + {ATTEN_OUT_NAME, {SASLayout::BSND, SASLayout::TND}}, + {ORI_SPARSE_INDICES, {SASLayout::BSND, SASLayout::TND}}, + {CMP_SPARSE_INDICES, {SASLayout::BSND, SASLayout::TND}}, +}; + +const std::map DATATYPE_TO_STRING_MAP = { + {ge::DT_UNDEFINED, "DT_UNDEFINED"}, // Used to indicate a DataType field has not been set. + {ge::DT_FLOAT, "DT_FLOAT"}, // float type + {ge::DT_FLOAT16, "DT_FLOAT16"}, // fp16 type + {ge::DT_INT8, "DT_INT8"}, // int8 type + {ge::DT_INT16, "DT_INT16"}, // int16 type + {ge::DT_UINT16, "DT_UINT16"}, // uint16 type + {ge::DT_UINT8, "DT_UINT8"}, // uint8 type + {ge::DT_INT32, "DT_INT32"}, // uint32 type + {ge::DT_INT64, "DT_INT64"}, // int64 type + {ge::DT_UINT32, "DT_UINT32"}, // unsigned int32 + {ge::DT_UINT64, "DT_UINT64"}, // unsigned int64 + {ge::DT_BOOL, "DT_BOOL"}, // bool type + {ge::DT_DOUBLE, "DT_DOUBLE"}, // double type + {ge::DT_DUAL, "DT_DUAL"}, // dual output type + {ge::DT_DUAL_SUB_INT8, "DT_DUAL_SUB_INT8"}, // dual output int8 type + {ge::DT_DUAL_SUB_UINT8, "DT_DUAL_SUB_UINT8"}, // dual output uint8 type + {ge::DT_COMPLEX32, "DT_COMPLEX32"}, // complex32 type + {ge::DT_COMPLEX64, "DT_COMPLEX64"}, // complex64 type + {ge::DT_COMPLEX128, "DT_COMPLEX128"}, // complex128 type + {ge::DT_QINT8, "DT_QINT8"}, // qint8 type + {ge::DT_QINT16, "DT_QINT16"}, // qint16 type + {ge::DT_QINT32, "DT_QINT32"}, // qint32 type + {ge::DT_QUINT8, "DT_QUINT8"}, // quint8 type + {ge::DT_QUINT16, "DT_QUINT16"}, // quint16 type + {ge::DT_RESOURCE, "DT_RESOURCE"}, // resource type + {ge::DT_STRING_REF, "DT_STRING_REF"}, // string ref type + {ge::DT_STRING, "DT_STRING"}, // string type + {ge::DT_VARIANT, "DT_VARIANT"}, // dt_variant type + {ge::DT_BF16, "DT_BFLOAT16"}, // dt_bfloat16 type + {ge::DT_INT4, "DT_INT4"}, // dt_variant type + {ge::DT_UINT1, "DT_UINT1"}, // dt_variant type + {ge::DT_INT2, "DT_INT2"}, // dt_variant type + {ge::DT_UINT2, "DT_UINT2"} // dt_variant type +}; + +std::string SASLayoutToSerialString(SASLayout layout) +{ + switch (layout) { + case SASLayout::BSND: return "BSND"; + case SASLayout::TND: return "TND"; + case SASLayout::PA_ND: return "PA_ND"; + default: return "UNKNOWN"; + } +} + +struct SASCompileInfo { + int64_t core_num; +}; + +static const std::map> SAS_LAYOUT_AXIS_MAP = { + {SASLayout::BSND, {SASAxis::B, SASAxis::S, SASAxis::N, SASAxis::D}}, + {SASLayout::TND, {SASAxis::T, SASAxis::N, SASAxis::D}}, + {SASLayout::PA_ND, {SASAxis::Bn, SASAxis::Bs, SASAxis::N, SASAxis::D}}, +}; + +static const std::map SAS_LAYOUT_DIM_MAP = { + {SASLayout::BSND, DIM_NUM_FOUR}, + {SASLayout::TND, DIM_NUM_THREE}, + {SASLayout::PA_ND, DIM_NUM_FOUR}, +}; + +static std::string SASDataTypeToSerialString(ge::DataType type) +{ + const auto it = DATATYPE_TO_STRING_MAP.find(type); + if (it != DATATYPE_TO_STRING_MAP.end()) { + return it->second; + } else { + OP_LOGE("sparseAttnSharedkv", "datatype %d not support", type); + return "UNDEFINED"; + } +} + +// --------------------------SASInfoParser类成员函数定义------------------------------------- +ge::graphStatus SASInfoParser::CheckRequiredInOutExistence() const +{ + OP_CHECK_IF(opParamInfo_.q.shape == nullptr, OP_LOGE(opName_, "Shape of tensor q is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.q.desc == nullptr, OP_LOGE(opName_, "Desc of tensor q is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.oriKv.tensor == nullptr, OP_LOGE(opName_, "tensor of ori_Kv is nullptr"), + return ge::GRAPH_FAILED); + if (kvLayout_ == SASLayout::PA_ND) { + OP_CHECK_IF(opParamInfo_.oriBlockTable.tensor == nullptr, OP_LOGE(opName_, "tensor of ori_block_table is nullptr"), + return ge::GRAPH_FAILED); + } + if (perfMode_ == SASTemplateMode::CFA_TEMPLATE_MODE){ + OP_CHECK_IF(opParamInfo_.cmpKv.tensor == nullptr, OP_LOGE(opName_, "tensor of cmp_kv is nullptr"), + return ge::GRAPH_FAILED); + } + if (perfMode_ == SASTemplateMode::SCFA_TEMPLATE_MODE){ + OP_CHECK_IF(opParamInfo_.cmpKv.tensor == nullptr, OP_LOGE(opName_, "tensor of cmp_kv is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.cmpSparseIndices.tensor == nullptr, OP_LOGE(opName_, "cmp_sparse_indices is nullptr"), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::CheckRequiredAttrExistence() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::CheckRequiredParaExistence() const +{ + if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS || + CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::CheckUnrequiredParaExistence() const +{ + OP_CHECK_IF(opParamInfo_.oriSparseIndices.tensor != nullptr || opParamInfo_.oriSparseIndices.desc != nullptr, + OP_LOGE(opName_, "Currently, ori_sparse_indices must be a nullptr"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetOpName() +{ + if (context_->GetNodeName() == nullptr) { + OP_LOGE("SparseAttnSharedkv", "opName got from TilingContext is nullptr"); + return ge::GRAPH_FAILED; + } + opName_ = context_->GetNodeName(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetNpuInfo() +{ + platformInfo_ = context_->GetPlatformInfo(); + OP_CHECK_IF(platformInfo_ == nullptr, OP_LOGE(opName_, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo_); + aivNum_ = ascendcPlatform.GetCoreNumAiv(); + aicNum_ = ascendcPlatform.GetCoreNumAic(); + OP_CHECK_IF(aicNum_ == 0 || aivNum_ == 0, OP_LOGE(opName_, "num of core obtained is 0."), return ge::GRAPH_FAILED); + + socVersion_ = ascendcPlatform.GetSocVersion(); + if ((socVersion_ != platform_ascendc::SocVersion::ASCEND910B) && + (socVersion_ != platform_ascendc::SocVersion::ASCEND910_93)) { + OP_LOGE(opName_, "SOC Version[%d] is not support.", (int32_t)socVersion_); + return GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +void SASInfoParser::GetOptionalInputParaInfo() +{ + opParamInfo_.oriKv.tensor = context_->GetOptionalInputTensor(ORI_KV_INDEX); + opParamInfo_.oriKv.desc = context_->GetOptionalInputDesc(ORI_KV_INDEX); + opParamInfo_.cmpKv.tensor = context_->GetOptionalInputTensor(CMP_KV_INDEX); + opParamInfo_.cmpKv.desc = context_->GetOptionalInputDesc(CMP_KV_INDEX); + opParamInfo_.oriSparseIndices.tensor = context_->GetOptionalInputTensor(ORI_SPARSE_INDICES_INDEX); + opParamInfo_.oriSparseIndices.desc = context_->GetOptionalInputDesc(ORI_SPARSE_INDICES_INDEX); + opParamInfo_.cmpSparseIndices.tensor = context_->GetOptionalInputTensor(CMP_SPARSE_INDICES_INDEX); + opParamInfo_.cmpSparseIndices.desc = context_->GetOptionalInputDesc(CMP_SPARSE_INDICES_INDEX); + opParamInfo_.oriBlockTable.tensor = context_->GetOptionalInputTensor(ORI_BLOCK_TABLE_INDEX); + opParamInfo_.oriBlockTable.desc = context_->GetOptionalInputDesc(ORI_BLOCK_TABLE_INDEX); + opParamInfo_.cmpBlockTable.tensor = context_->GetOptionalInputTensor(CMP_BLOCK_TABLE_INDEX); + opParamInfo_.cmpBlockTable.desc = context_->GetOptionalInputDesc(CMP_BLOCK_TABLE_INDEX); + opParamInfo_.sinks.tensor = context_->GetOptionalInputTensor(SINKS_INDEX); + opParamInfo_.sinks.desc = context_->GetOptionalInputDesc(SINKS_INDEX); + opParamInfo_.cuSeqLensQ.tensor = context_->GetOptionalInputTensor(CU_SEQLENS_Q_INDEX); + opParamInfo_.cuSeqLensQ.desc = context_->GetOptionalInputDesc(CU_SEQLENS_Q_INDEX); + opParamInfo_.seqUsedQ.tensor = context_->GetOptionalInputTensor(SEQUSED_Q_INDEX); + opParamInfo_.seqUsedQ.desc = context_->GetOptionalInputDesc(SEQUSED_Q_INDEX); + opParamInfo_.cuSeqLensKv.tensor = context_->GetOptionalInputTensor(CU_SEQLENS_KV_INDEX); + opParamInfo_.cuSeqLensKv.desc = context_->GetOptionalInputDesc(CU_SEQLENS_KV_INDEX); + opParamInfo_.cuSeqLensCmpKv.tensor = context_->GetOptionalInputTensor(CU_SEQLENS_CMP_KV_INDEX); + opParamInfo_.cuSeqLensCmpKv.desc = context_->GetOptionalInputDesc(CU_SEQLENS_CMP_KV_INDEX); + opParamInfo_.sequsedKv.tensor = context_->GetOptionalInputTensor(SEQUSED_KV_INDEX); + opParamInfo_.sequsedKv.desc = context_->GetOptionalInputDesc(SEQUSED_KV_INDEX); + opParamInfo_.metadata.desc = context_->GetOptionalInputDesc(METADATA_INDEX); + opParamInfo_.metadata.tensor = context_->GetOptionalInputTensor(METADATA_INDEX); +} + +void SASInfoParser::GetInputParaInfo() +{ + opParamInfo_.q.desc = context_->GetInputDesc(Q_INDEX); + opParamInfo_.q.shape = context_->GetInputShape(Q_INDEX); + GetOptionalInputParaInfo(); +} + +void SASInfoParser::GetOutputParaInfo() +{ + opParamInfo_.attnOut.desc = context_->GetOutputDesc(ATTN_OUT_INDEX); + opParamInfo_.attnOut.shape = context_->GetOutputShape(ATTN_OUT_INDEX); +} + +ge::graphStatus SASInfoParser::GetAttrParaInfo() +{ + auto attrs = context_->GetAttrs(); + OP_CHECK_IF(attrs == nullptr, OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "attrs got from ge is nullptr"), + return ge::GRAPH_FAILED); + OP_LOGI(context_->GetNodeName(), "GetAttrParaInfo start"); + opParamInfo_.softmaxScale = attrs->GetAttrPointer(ATTR_SOFTMAX_SCALE_INDEX); + opParamInfo_.cmpRatio = attrs->GetAttrPointer(ATTR_CMP_RATIO_INDEX); + opParamInfo_.oriMaskMode = attrs->GetAttrPointer(ATTR_ORI_MASK_MODE_INDEX); + opParamInfo_.cmpMaskMode = attrs->GetAttrPointer(ATTR_CMP_MASK_MODE_INDEX); + opParamInfo_.oriKvStride = attrs->GetAttrPointer(ATTR_ORI_KV_STRIDE_INDEX); + opParamInfo_.cmpKvStride = attrs->GetAttrPointer(ATTR_CMP_KV_STRIDE_INDEX); + opParamInfo_.oriWinLeft = attrs->GetAttrPointer(ATTR_ORI_WIN_LEFT_INDEX); + opParamInfo_.oriWinRight = attrs->GetAttrPointer(ATTR_ORI_WIN_RIGHT_INDEX); + opParamInfo_.layoutQ = attrs->GetStr(ATTR_LAYOUT_Q_INDEX); + opParamInfo_.layoutKv = attrs->GetStr(ATTR_LAYOUT_KV_INDEX); + opParamInfo_.returnSoftmaxLse = attrs->GetAttrPointer(ATTR_RETURN_SOFTMAX_LSE); + + OP_LOGI(context_->GetNodeName(), "GetAttrParaInfo end"); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetOpParaInfo() +{ + GetInputParaInfo(); + GetOutputParaInfo(); + if (ge::GRAPH_SUCCESS != GetAttrParaInfo()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetInOutDataType() +{ + qType_ = opParamInfo_.q.desc->GetDataType(); + outputType_ = opParamInfo_.attnOut.desc->GetDataType(); + if (opParamInfo_.oriKv.desc != nullptr) { + oriKvType_ = opParamInfo_.oriKv.desc->GetDataType(); + } + if (opParamInfo_.cmpKv.desc != nullptr) { + cmpKvType_ = opParamInfo_.cmpKv.desc->GetDataType(); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetSASTemplateMode(SASTilingInfo &sasInfo) +{ + if (opParamInfo_.oriKv.desc != nullptr) { + if (opParamInfo_.cmpKv.desc != nullptr && opParamInfo_.cmpSparseIndices.tensor != nullptr) { + perfMode_ = SASTemplateMode::SCFA_TEMPLATE_MODE; + } else if (opParamInfo_.cmpKv.desc != nullptr && opParamInfo_.cmpSparseIndices.tensor == nullptr) { + perfMode_ = SASTemplateMode::CFA_TEMPLATE_MODE; + } else if (opParamInfo_.cmpKv.desc == nullptr && opParamInfo_.cmpSparseIndices.tensor == nullptr) { + perfMode_ = SASTemplateMode::SWA_TEMPLATE_MODE; + } else { + OP_LOGE(opName_, "When cmp_sparse_indices is not nullptr, cmp_kv cannot be nullptr."); + return ge::GRAPH_FAILED; + } + if (sasInfo.perfMode == SASTemplateMode::CFA_TEMPLATE_MODE || sasInfo.perfMode == SASTemplateMode::SCFA_TEMPLATE_MODE) { + if (kvLayout_ == SASLayout::TND && opParamInfo_.cuSeqLensCmpKv.tensor == nullptr) { + OP_LOGE(opName_, "the layout_kv is %s, seqlens_cmp_kv must be provided.", SASLayoutToSerialString(kvLayout_).c_str()); + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; + } else { + OP_LOGE(opName_, "ori_kv is nullptr"); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetQueryAndOutLayout() +{ + // 获取q和attnOut的Layout基准值 + // layoutQuery: {qLayout, outLayout} + const map> layoutMap = { + {"BSND", {SASLayout::BSND, SASLayout::BSND}}, + {"TND", {SASLayout::TND, SASLayout::TND }}, + }; + std::string layout(opParamInfo_.layoutQ); + auto it = layoutMap.find(layout); + if (it != layoutMap.end()) { + qLayout_ = it->second.first; + outLayout_ = it->second.second; + oriSparseIndicesLayout_ = qLayout_; + cmpSparseIndicesLayout_ = qLayout_; + } else { + OP_LOGE(opName_, "layout of q is %s, it is unsupported.", layout.c_str()); + return ge::GRAPH_FAILED; + } + if (qLayout_ == SASLayout::BSND){ + OP_CHECK_IF(opParamInfo_.cuSeqLensQ.tensor != nullptr, + OP_LOGE(opName_, "when q's layout is BSND, cu_seqlens_q is null."), + return ge::GRAPH_FAILED); + } + if (qLayout_ == SASLayout::TND){ + OP_CHECK_IF(opParamInfo_.seqUsedQ.tensor != nullptr, + OP_LOGE(opName_, "when q's layout is TND, seqused_q is null."), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetKvLayout() +{ + const map layoutKVMap = { + {"PA_ND", SASLayout::PA_ND}, + {"BSND", SASLayout::BSND}, + {"TND", SASLayout::TND}, + }; + std::string layout(opParamInfo_.layoutKv); + auto it = layoutKVMap.find(layout); + if (it != layoutKVMap.end()) { + kvLayout_ = it->second; + } else { + OP_LOGE(opName_, "layout_kv is %s, it is unsupported.", layout.c_str()); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +// =============Parser function==================== +bool SASInfoParser::HasAxis(const SASAxis &axis, const SASLayout &layout, const gert::Shape &shape) const +{ + const auto& layoutIt = SAS_LAYOUT_AXIS_MAP.find(layout); + if (layoutIt == SAS_LAYOUT_AXIS_MAP.end()) { + return false; + } + + const std::vector& axes = layoutIt->second; + const auto& axisIt = std::find(axes.begin(), axes.end(), axis); + if (axisIt == axes.end()) { + return false; + } + const auto& dimIt = SAS_LAYOUT_DIM_MAP.find(layout); + if (dimIt == SAS_LAYOUT_DIM_MAP.end() || dimIt->second != shape.GetDimNum()) { + return false; + } + return true; +} + +size_t SASInfoParser::GetAxisIdx(const SASAxis &axis, const SASLayout &layout) const +{ + const std::vector& axes = SAS_LAYOUT_AXIS_MAP.find(layout)->second; + const auto& axisIt = std::find(axes.begin(), axes.end(), axis); + return std::distance(axes.begin(), axisIt); +} + +uint32_t SASInfoParser::GetAxisNum(const gert::Shape &shape, const SASAxis &axis,const SASLayout &layout) const +{ + return HasAxis(axis, layout, shape) ? shape.GetDim(GetAxisIdx(axis, layout)) : invalidDimValue_; +} + +void SASInfoParser::SetSASShape() +{ + qShape_ = opParamInfo_.q.shape->GetStorageShape(); + if (opParamInfo_.oriKv.tensor != nullptr) { + oriKvShape_ = opParamInfo_.oriKv.tensor->GetStorageShape(); + } else { + OP_LOGE(opName_, "q tensor is nullptr, please check input parameters."); + } + if (opParamInfo_.cmpKv.tensor != nullptr) { + cmpKvShape_ = opParamInfo_.cmpKv.tensor->GetStorageShape(); + } + if (perfMode_ == SASTemplateMode::SCFA_TEMPLATE_MODE) + { + if (opParamInfo_.cmpSparseIndices.tensor != nullptr) { + cmpSparseIndicesShape_ = opParamInfo_.cmpSparseIndices.tensor->GetStorageShape(); + uint32_t cmpSparseIndicesT = GetAxisNum(cmpSparseIndicesShape_, SASAxis::T, cmpSparseIndicesLayout_); + } else { + OP_LOGE(opName_, "cmp_sparse_indices tensor is nullptr, please check input parameters."); + } + } +} + +ge::graphStatus SASInfoParser::GetN1Size() +{ + n1Size_ = GetAxisNum(qShape_, SASAxis::N, qLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetN2Size() +{ + if (opParamInfo_.oriKv.tensor != nullptr) { + n2Size_ = GetAxisNum(oriKvShape_, SASAxis::N, kvLayout_); + } + if (opParamInfo_.cmpKv.tensor != nullptr) { + uint32_t cmpKvN2Size_ = GetAxisNum(cmpKvShape_, SASAxis::N, kvLayout_); + if (perfMode_ == SASTemplateMode::SCFA_TEMPLATE_MODE){ + uint32_t cmpSparseIndicesN2Size_ = GetAxisNum(cmpSparseIndicesShape_, SASAxis::N, cmpSparseIndicesLayout_); + OP_CHECK_IF(cmpKvN2Size_ != n2Size_ || n2Size_ != cmpSparseIndicesN2Size_, + OP_LOGE(opName_, "N2 size check failed! Expected ori_kv's N2(%u) == cmp_sparse_indices's N2(%u).", n2Size_, cmpSparseIndicesN2Size_), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(cmpKvN2Size_ != n2Size_, + OP_LOGE(opName_, "N2 size check failed! Expected cmp_kv's N2(%u) ==ori_kv's N2(%u).", cmpKvN2Size_, n2Size_), + return ge::GRAPH_FAILED); + n2Size_ = cmpKvN2Size_; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetGSize() +{ + if (n2Size_ != 0) { + gSize_ = n1Size_ / n2Size_; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + SASLayout &layout, const std::string &name) const +{ + if ((tensor == nullptr)) { + OP_LOGE(opName_, "when layout of q is %s, %s must be provided.", + SASLayoutToSerialString(layout).c_str(), name.c_str()); + return ge::GRAPH_FAILED; + } + int64_t shapeSize = tensor->GetShapeSize(); + if (shapeSize <= 0) { + OP_LOGE(opName_, "the shape size of %s is %ld, it should be greater than 0.", + name.c_str(), shapeSize); + return ge::GRAPH_FAILED; + } + size = static_cast(shapeSize); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetActualSeqLenQSize(uint32_t &size) +{ + return GetActualSeqLenSize(size, opParamInfo_.cuSeqLensQ.tensor, qLayout_, "cuSeqLensQ"); +} + +ge::graphStatus SASInfoParser::GetBatchSize() +{ + // 获取B基准值 + // 1、非TND时, 以query的batch_size维度为基准; + // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小 + if (qLayout_ == SASLayout::TND) { + return GetActualSeqLenQSize(bSize_); + } else { // BSND + bSize_ = GetAxisNum(qShape_, SASAxis::B, qLayout_); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetQTSize() +{ + // 获取query的T基准值 + // 1、非TND时, 以query的batch_size维度为基准; + // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小 + qTSize_ = (qLayout_ == SASLayout::TND) ? GetAxisNum(qShape_, SASAxis::T, qLayout_) : 0; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetKVTSize() +{ + // 获取KV的T基准值 + // 1、非TND时, 以KV的batch_size维度为基准; + // 2、TND时, actual_seq_lens_ori_kv和actual_seq_lens_cmp_kv必须传入, 以actual_seq_lens_ori_kv数组的长度为B轴大小(当前接口只传入oriseq，先以oriseq算出cmpseq) + orikvTSize_ = (kvLayout_ == SASLayout::TND) ? GetAxisNum(oriKvShape_, SASAxis::T, kvLayout_) : 0; + // 入参接口信息可以从GetOptionalInputParaInfo()函数中获取 + // cmpkvTSize_ = (kvLayout_ == SASLayout::TND) ? GetAxisNum(cmpKvShape_, SASAxis::T, kvLayout_) : 0; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetS1Size() +{ + // 获取S1基准值 + // 1、非TND时, 以query的S维度为基准; + // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组中的最大值为基准 + if (qLayout_ == SASLayout::TND) { + s1Size_ = GetAxisNum(qShape_, SASAxis::T, qLayout_); + } else { // BSND + s1Size_ = GetAxisNum(qShape_, SASAxis::S, qLayout_); + } + if (perfMode_ == SASTemplateMode::SCFA_TEMPLATE_MODE){ + if (cmpSparseIndicesLayout_ == SASLayout::TND) { + uint32_t cmpSparseIndicesT = GetAxisNum(cmpSparseIndicesShape_, SASAxis::T, cmpSparseIndicesLayout_); + OP_CHECK_IF(cmpSparseIndicesT != s1Size_, + OP_LOGE(opName_, "T size check failed !"), + return ge::GRAPH_FAILED); + } else{ + uint32_t cmpSparseIndicesS1 = GetAxisNum(cmpSparseIndicesShape_, SASAxis::S, cmpSparseIndicesLayout_); + OP_CHECK_IF(cmpSparseIndicesS1 != s1Size_, + OP_LOGE(opName_, "s1 size check failed !"), + return ge::GRAPH_FAILED); + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetMaxBlockNumPerBatch() +{ + if (opParamInfo_.oriBlockTable.tensor == nullptr) { + OP_LOGE(opName_, "the layout_kv is %s, block_table must be provided.", SASLayoutToSerialString(kvLayout_).c_str()); + return ge::GRAPH_FAILED; + } + uint32_t oriDimNum = opParamInfo_.oriBlockTable.tensor->GetStorageShape().GetDimNum(); + if (oriDimNum != DIM_NUM_TWO) { + OP_LOGE(opName_, "the dim num of ori_block_table is %u, it should be %u.", oriDimNum, DIM_NUM_TWO); + return ge::GRAPH_FAILED; + } + if (opParamInfo_.oriBlockTable.tensor->GetStorageShape().GetDim(1) < 0) { + OP_LOGE(opName_, "%s's second dimension(%lld) should be non-negative number.", + ORI_BLOCK_TABLE_NAME.c_str(), opParamInfo_.oriBlockTable.tensor->GetStorageShape().GetDim(1)); + return ge::GRAPH_FAILED; + } + oriMaxBlockNumPerBatch_ = opParamInfo_.oriBlockTable.tensor->GetStorageShape().GetDim(1); + + if (opParamInfo_.cmpBlockTable.tensor != nullptr) { + uint32_t cmpDimNum = opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDimNum(); + if (cmpDimNum != DIM_NUM_TWO) { + OP_LOGE(opName_, "the dim num of cmp_block_table is %u, it should be %u.", cmpDimNum, DIM_NUM_TWO); + return ge::GRAPH_FAILED; + } + if (qLayout_ == SASLayout::TND) { + if (opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(0) != bSize_ - 1) { + OP_LOGE(opName_, "cmp_block_table's first dimension(%u) should be equal to query's B(%u).", + opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(1), bSize_ - 1); + return ge::GRAPH_FAILED; + } + } else if (qLayout_ == SASLayout::BSND) { + if (opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(0) != bSize_) { + OP_LOGE(opName_, "cmp_block_table's first dimension(%u) should be equal to query's B(%u).", + opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(1), bSize_); + return ge::GRAPH_FAILED; + } + } + if (opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(1) <= 0) { + OP_LOGE(opName_, "%s's second dimension(%lld) should be greater than 0", + CMP_BLOCK_TABLE_NAME.c_str(), opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(1)); + return ge::GRAPH_FAILED; + } + cmpMaxBlockNumPerBatch_ = opParamInfo_.cmpBlockTable.tensor->GetStorageShape().GetDim(1); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetBlockSize() +{ + oriBlockSize_ = GetAxisNum(oriKvShape_, SASAxis::Bs, kvLayout_); + cmpBlockSize_ = GetAxisNum(cmpKvShape_, SASAxis::Bs, kvLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetS2SizeForPageAttention() +{ + if (GetMaxBlockNumPerBatch() != ge::GRAPH_SUCCESS || GetBlockSize() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + s2Size_ = oriMaxBlockNumPerBatch_ * oriBlockSize_; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetS2SizeForTND() +{ + if (opParamInfo_.cuSeqLensKv.tensor == nullptr) { + OP_LOGE(opName_, "the layout_kv is %s, seqlens_ori_kv must be provided.", SASLayoutToSerialString(kvLayout_).c_str()); + return ge::GRAPH_FAILED; + } + // if (opParamInfo_.sequsedKv.tensor == nullptr) { + // OP_LOGE(opName_, "the layout_kv is %s, sequsedKv must be provided.", SASLayoutToSerialString(kvLayout_).c_str()); + // return ge::GRAPH_FAILED; + // } + // 这里返回累加和的最大值 + s2Size_ = GetAxisNum(oriKvShape_, SASAxis::T, qLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetS2Size() +{ + // 获取S2基准值:PAGE_ATTENTION时, S2 = block_table.dim1 * block_size + // 1、PAGE_ATTENTION时, S2 = block_table.dim1 * block_size + // 2、BSND时, S2直接获取 + if (kvLayout_ == SASLayout::BSND) { + if (opParamInfo_.oriKv.tensor != nullptr) { + s2Size_ = GetAxisNum(oriKvShape_, SASAxis::S, kvLayout_); + return ge::GRAPH_SUCCESS; + } + if (opParamInfo_.cmpKv.tensor != nullptr) { + s2Size_ = GetAxisNum(cmpKvShape_, SASAxis::S, kvLayout_); + return ge::GRAPH_SUCCESS; + } + return ge::GRAPH_FAILED; + } + return (kvLayout_ == SASLayout::PA_ND) ? GetS2SizeForPageAttention() : GetS2SizeForTND(); +} + +ge::graphStatus SASInfoParser::GetQHeadDim() +{ + qHeadDim_ = GetAxisNum(qShape_, SASAxis::D, qLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetValueHeadDim() +{ + if (opParamInfo_.oriKv.tensor != nullptr) { + oriKvHeadDim_ = GetAxisNum(oriKvShape_, SASAxis::D, kvLayout_); + } + if (opParamInfo_.cmpKv.tensor != nullptr) { + cmpKvHeadDim_ = GetAxisNum(cmpKvShape_, SASAxis::D, kvLayout_); + } + return ge::GRAPH_SUCCESS; +} + + +ge::graphStatus SASInfoParser::GetSparseBlockCount() +{ + if (opParamInfo_.cmpSparseIndices.tensor != nullptr) { + sparseBlockCount_ = GetAxisNum(cmpSparseIndicesShape_, SASAxis::K, cmpSparseIndicesLayout_); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetSinks() +{ + if (opParamInfo_.sinks.tensor == nullptr) { + OP_LOGE(opName_, "%s must be provided!", SINKS_NAME.c_str()); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASInfoParser::GetActualseqInfo() +{ + maxActualseq_ = static_cast(s2Size_); + if (qLayout_ == SASLayout::TND) { + if (opParamInfo_.cuSeqLensQ.tensor != nullptr) { + if (opParamInfo_.cuSeqLensQ.tensor->GetShapeSize() != bSize_) { + OP_LOGE(opName_, "cu_seqlens_q's dimension should be equal to %u.", bSize_); + return ge::GRAPH_FAILED; + } + actualLenDimsQ_ = opParamInfo_.cuSeqLensQ.tensor->GetShapeSize() - 1; // cuSeqLensQ shape is B+1 + OP_CHECK_IF(actualLenDimsQ_ == 0, + OP_LOGE(opName_, "cu_seqlens_q cannot be empty tensor."), + return ge::GRAPH_FAILED); + } else { + OP_LOGE(opName_, "When layout_q is TND, input cu_seqlens_q must be provided"); + return ge::GRAPH_FAILED; + } + } else { + if (opParamInfo_.seqUsedQ.tensor != nullptr) { + actualLenDimsQ_ = opParamInfo_.seqUsedQ.tensor->GetShapeSize(); + } + } + if (kvLayout_ != SASLayout::PA_ND && kvLayout_ != SASLayout::BSND && kvLayout_ != SASLayout::TND) { + OP_LOGE(opName_, "ori_kv and cmp_kv only support PA_ND, BSND and TND layout."); + return ge::GRAPH_FAILED; + } + if (kvLayout_ == SASLayout::PA_ND) { + if (opParamInfo_.sequsedKv.tensor != nullptr) { + if (qLayout_ == SASLayout::BSND){ + if (opParamInfo_.sequsedKv.tensor->GetShapeSize() != bSize_) { + OP_LOGE(opName_, "seqused_kv's dimension should be equal to %u, but got %ld.", + bSize_, opParamInfo_.sequsedKv.tensor->GetShapeSize()); + return ge::GRAPH_FAILED; + } + } else { + if (opParamInfo_.sequsedKv.tensor->GetShapeSize() != (bSize_ - 1)) { + OP_LOGE(opName_, "seqused_kv's dimension should be equal to %u (bSize - 1), but got %ld.", + (bSize_ - 1), opParamInfo_.sequsedKv.tensor->GetShapeSize()); + return ge::GRAPH_FAILED; + } + } + OP_CHECK_IF(opParamInfo_.sequsedKv.desc->GetDataType() != ge::DT_INT32, + OP_LOGE(opName_, "seqused_kv's dtype must be DT_INT32."), + return ge::GRAPH_FAILED); + actualLenDimsKV_ = opParamInfo_.sequsedKv.tensor->GetShapeSize(); + OP_CHECK_IF(actualLenDimsKV_ == 0, + OP_LOGE(opName_, "seqused_kv cannot be empty tensor."), + return ge::GRAPH_FAILED); + } else { + OP_LOGE(opName_, "When kv layout is PA_ND, input sequsedKv must be provided"); + return ge::GRAPH_FAILED; + } + } else if (kvLayout_ == SASLayout::TND) { + if (opParamInfo_.cuSeqLensKv.tensor != nullptr) { + if (qLayout_ == SASLayout::BSND){ + if (opParamInfo_.cuSeqLensKv.tensor->GetShapeSize() != bSize_ + 1) { + OP_LOGE(opName_, "cuSeqLensKv's dimension should be equal to %u (bSize + 1), but got %ld.", + (bSize_ + 1), opParamInfo_.sequsedKv.tensor->GetShapeSize()); + return ge::GRAPH_FAILED; + } + } else { + if (opParamInfo_.cuSeqLensKv.tensor->GetShapeSize() != (bSize_)) { + OP_LOGE(opName_, "cuSeqLensKv's dimension should be equal to %u, but got %ld.", + bSize_, opParamInfo_.sequsedKv.tensor->GetShapeSize()); + return ge::GRAPH_FAILED; + } + } + actualLenDimsKV_ = opParamInfo_.cuSeqLensKv.tensor->GetShapeSize(); + } else { + OP_LOGE(opName_, "When kv layout is TND, input cuSeqLensKv must be provided"); + return ge::GRAPH_FAILED; + } + } + if (opParamInfo_.seqUsedQ.tensor != nullptr) { + actualLenDimsQ_ = opParamInfo_.seqUsedQ.tensor->GetShapeSize(); + } else if (opParamInfo_.cuSeqLensQ.tensor != nullptr) { + actualLenDimsQ_ = opParamInfo_.cuSeqLensQ.tensor->GetShapeSize() - 1; // cuSeqLensQ shape is B+1 + } + return ge::GRAPH_SUCCESS; +} + +void SASInfoParser::GenerateInfo(SASTilingInfo &sasInfo) +{ + sasInfo.opName = opName_; + sasInfo.platformInfo = platformInfo_; + sasInfo.opParamInfo = opParamInfo_; + sasInfo.socVersion = socVersion_; + + sasInfo.bSize = bSize_; + sasInfo.n1Size = n1Size_; + sasInfo.n2Size = n2Size_; + sasInfo.s1Size = s1Size_; + sasInfo.s2Size = s2Size_; + sasInfo.gSize = gSize_; + sasInfo.qHeadDim = qHeadDim_; + sasInfo.oriKvHeadDim = oriKvHeadDim_; + sasInfo.cmpKvHeadDim = cmpKvHeadDim_; + sasInfo.qTSize = qTSize_; + sasInfo.sparseBlockCount = sparseBlockCount_; + sasInfo.oriWinLeft = oriWinLeft_; + sasInfo.oriWinRight = oriWinRight_; + sasInfo.qType = qType_; + sasInfo.oriKvType = oriKvType_; + sasInfo.cmpKvType = cmpKvType_; + sasInfo.outputType = outputType_; + sasInfo.perfMode = perfMode_; + + if (kvLayout_ == SASLayout::PA_ND) { + sasInfo.totalBlockNum = (opParamInfo_.oriKv.tensor != nullptr) ? + opParamInfo_.oriKv.tensor->GetStorageShape().GetDim(0) : 0; + } + sasInfo.sparseBlockSize = 1; + sasInfo.oriBlockSize = oriBlockSize_; + sasInfo.cmpBlockSize = cmpBlockSize_; + sasInfo.blockTypeSize = sizeof(float); + sasInfo.oriMaxBlockNumPerBatch = oriMaxBlockNumPerBatch_; + sasInfo.cmpMaxBlockNumPerBatch = cmpMaxBlockNumPerBatch_; + + sasInfo.actualLenDimsQ = actualLenDimsQ_; + sasInfo.actualLenDimsKV = actualLenDimsKV_; + sasInfo.maxActualseq = maxActualseq_; + sasInfo.actualSeqLenFlag = (opParamInfo_.sequsedKv.tensor != nullptr); + sasInfo.isSameSeqAllKVTensor = isSameSeqAllKVTensor_; + + sasInfo.softmaxScale = *opParamInfo_.softmaxScale; + sasInfo.cmpRatio = *opParamInfo_.cmpRatio; + sasInfo.oriMaskMode = *opParamInfo_.oriMaskMode; + sasInfo.cmpMaskMode = *opParamInfo_.cmpMaskMode; + sasInfo.oriKvStride = *opParamInfo_.oriKvStride; + sasInfo.cmpKvStride = *opParamInfo_.cmpKvStride; + sasInfo.oriWinLeft = *opParamInfo_.oriWinLeft; + sasInfo.oriWinRight = *opParamInfo_.oriWinRight; + + sasInfo.qLayout = qLayout_; + sasInfo.oriSparseIndicesLayout = oriSparseIndicesLayout_; + sasInfo.cmpSparseIndicesLayout = cmpSparseIndicesLayout_; + sasInfo.kvLayout = kvLayout_; + sasInfo.outLayout = outLayout_; + sasInfo.returnSoftmaxLse = *opParamInfo_.returnSoftmaxLse; +} + +ge::graphStatus SASInfoParser::Parse(SASTilingInfo &sasInfo) +{ + if (context_ == nullptr) { + OP_LOGE("SparseFlashAttention", "tiling context is nullptr!"); + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetOpName() || + ge::GRAPH_SUCCESS != GetNpuInfo() || + ge::GRAPH_SUCCESS != GetOpParaInfo() || + ge::GRAPH_SUCCESS != GetKvLayout() || + ge::GRAPH_SUCCESS != CheckRequiredParaExistence() || + ge::GRAPH_SUCCESS != CheckUnrequiredParaExistence()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetInOutDataType() || + ge::GRAPH_SUCCESS != GetQueryAndOutLayout() || + ge::GRAPH_SUCCESS != GetSASTemplateMode(sasInfo)) { + return ge::GRAPH_FAILED; + } + + SetSASShape(); + if ( + ge::GRAPH_SUCCESS != GetN1Size() || + ge::GRAPH_SUCCESS != GetN2Size() || + ge::GRAPH_SUCCESS != GetGSize() || + ge::GRAPH_SUCCESS != GetBatchSize() || + ge::GRAPH_SUCCESS != GetQTSize() || + ge::GRAPH_SUCCESS != GetKVTSize() || + ge::GRAPH_SUCCESS != GetS1Size() || + ge::GRAPH_SUCCESS != GetS2Size() || + ge::GRAPH_SUCCESS != GetQHeadDim() || + ge::GRAPH_SUCCESS != GetValueHeadDim() || + ge::GRAPH_SUCCESS != GetSparseBlockCount() || + ge::GRAPH_SUCCESS != GetSinks() + ) { + return ge::GRAPH_FAILED; + } + if (ge::GRAPH_SUCCESS != GetActualseqInfo()) { + return ge::GRAPH_FAILED; + } + GenerateInfo(sasInfo); + return ge::GRAPH_SUCCESS; +} + +void SASTilingCheck::Init() +{ + opName_ = sasInfo_.opName; + platformInfo_ = sasInfo_.platformInfo; + opParamInfo_ = sasInfo_.opParamInfo; + socVersion_ = sasInfo_.socVersion; + bSize_ = sasInfo_.bSize; + n1Size_ = sasInfo_.n1Size; + n2Size_ = sasInfo_.n2Size; + s1Size_ = sasInfo_.s1Size; + s2Size_ = sasInfo_.s2Size; + gSize_ = sasInfo_.gSize; + qHeadDim_ = sasInfo_.qHeadDim; + oriKvHeadDim_ = sasInfo_.oriKvHeadDim; + cmpKvHeadDim_ = sasInfo_.cmpKvHeadDim; + oriBlockSize_ = sasInfo_.oriBlockSize; + cmpBlockSize_ = sasInfo_.cmpBlockSize; + qTSize_ = sasInfo_.qTSize; + qType_ = sasInfo_.qType; + oriKvType_ = sasInfo_.oriKvType; + cmpKvType_ = sasInfo_.cmpKvType; + outputType_ = sasInfo_.outputType; + cmpRatio_ = sasInfo_.cmpRatio; + qLayout_ = sasInfo_.qLayout; + oriSparseIndicesLayout_ = sasInfo_.oriSparseIndicesLayout; + cmpSparseIndicesLayout_ = sasInfo_.cmpSparseIndicesLayout; + oriWinLeft_ = sasInfo_.oriWinLeft; + oriWinRight_ = sasInfo_.oriWinRight; + kvLayout_ = sasInfo_.kvLayout; + outLayout_ = sasInfo_.outLayout; +} + +void SASTilingCheck::LogErrorDtypeSupport(const std::vector &expectDtypeList, + const ge::DataType &actualDtype, const std::string &name) const +{ + std::ostringstream oss; + for (size_t i = 0; i < expectDtypeList.size(); ++i) { + oss << SASDataTypeToSerialString(expectDtypeList[i]); + if (i < expectDtypeList.size() - 1) { + oss << ", "; + } + } + OP_LOGE(opName_, "Tensor %s only supports dtype %s, but got %s", + name.c_str(), oss.str().c_str(), SASDataTypeToSerialString(actualDtype).c_str()); +} + +ge::graphStatus SASTilingCheck::CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc, + const std::string &name) const +{ + if (desc != nullptr) { + const auto& it = DTYPE_SUPPORT_MAP.find(name); + OP_CHECK_IF(it == DTYPE_SUPPORT_MAP.end(), + OP_LOGE(opName_, "%s datatype support list should be specify in DTYPE_SUPPORT_MAP", name.c_str()), + return ge::GRAPH_FAILED); + auto &expectDtypeList = it->second; + OP_CHECK_IF(std::find( + expectDtypeList.begin(), expectDtypeList.end(), desc->GetDataType()) == expectDtypeList.end(), + LogErrorDtypeSupport(expectDtypeList, desc->GetDataType(), name), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +void SASTilingCheck::LogErrorLayoutSupport(const std::vector &expectLayoutList, + const SASLayout &actualLayout, const std::string &name) const +{ + std::ostringstream oss; + for (size_t i = 0; i < expectLayoutList.size(); ++i) { + oss << SASLayoutToSerialString(expectLayoutList[i]); + if (i < expectLayoutList.size() - 1) { + oss << ", "; + } + } + OP_LOGE(opName_, "Tensor %s only supports layout %s, but got %s", + name.c_str(), oss.str().c_str(), SASLayoutToSerialString(actualLayout).c_str()); +} + +ge::graphStatus SASTilingCheck::CheckLayoutSupport(const SASLayout &actualLayout, const std::string &name) const +{ + const auto& it = LAYOUT_SUPPORT_MAP.find(name); + OP_CHECK_IF(it == LAYOUT_SUPPORT_MAP.end(), + OP_LOGE(opName_, "%s layout support list should be specify in LAYOUT_SUPPORT_MAP", name.c_str()), + return ge::GRAPH_FAILED); + auto &expectLayoutList = it->second; + OP_CHECK_IF(std::find( + expectLayoutList.begin(), expectLayoutList.end(), actualLayout) == expectLayoutList.end(), + LogErrorLayoutSupport(expectLayoutList, actualLayout, name), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +template +void SASTilingCheck::LogErrorNumberSupport(const std::vector &expectNumberList, + const T &actualValue, const std::string &name, const std::string subName) const +{ + std::ostringstream oss; + for (size_t i = 0; i < expectNumberList.size(); ++i) { + oss << std::to_string(expectNumberList[i]); + if (i < expectNumberList.size() - 1) { + oss << ", "; + } + } + OP_LOGE(opName_, "%s %s only supports %s, but got %s", + name.c_str(), subName.c_str(), oss.str().c_str(), std::to_string(actualValue).c_str()); +} + +template +void SASTilingCheck::LogErrorDimNumSupport(const std::vector &expectNumberList, + const T &actualValue, const std::string &name) const +{ + LogErrorNumberSupport(expectNumberList, actualValue, name, "dimension"); +} + +ge::graphStatus SASTilingCheck::CheckDimNumSupport(const gert::StorageShape *shape, + const std::vector &expectDimNumList, const std::string &name) const +{ + if (shape == nullptr) { + return ge::GRAPH_SUCCESS; + } + + if (std::find(expectDimNumList.begin(), expectDimNumList.end(), + shape->GetStorageShape().GetDimNum()) == expectDimNumList.end()) { + LogErrorDimNumSupport(expectDimNumList, shape->GetStorageShape().GetDimNum(), name); + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckDimNumInLayoutSupport(const SASLayout &layout, + const gert::StorageShape *shape, const std::string &name) const +{ + const auto& dimIt = SAS_LAYOUT_DIM_MAP.find(layout); + OP_CHECK_IF(shape->GetStorageShape().GetDimNum() != dimIt->second, + OP_LOGE(opName_, "When layout is %s, %s dimension should be %zu, but it's %zu", + SASLayoutToSerialString(layout).c_str(), name.c_str(), dimIt->second, + shape->GetStorageShape().GetDimNum()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaQuery() const +{ + OP_CHECK_IF(opParamInfo_.q.shape->GetStorageShape().GetShapeSize() == 0, + OP_LOGE(opName_, "q cannot be empty tensor."), + return ge::GRAPH_FAILED); + if (opParamInfo_.q.desc == nullptr) { + OP_LOGE(opName_, "%s must be provided!", QUERY_NAME.c_str()); + return ge::GRAPH_FAILED; + } + const std::vector queryDimNumList = {DIM_NUM_THREE, DIM_NUM_FOUR}; + if ( + ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.q.desc, QUERY_NAME) || + ge::GRAPH_SUCCESS != CheckLayoutSupport(qLayout_, QUERY_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(opParamInfo_.q.shape, queryDimNumList, QUERY_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(qLayout_, opParamInfo_.q.shape, QUERY_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaOriKv() const +{ + const std::vector oriKvDimNumList = {DIM_NUM_THREE, DIM_NUM_FOUR}; + if ( + ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.oriKv.desc, ORI_KV_NAME) || + ge::GRAPH_SUCCESS != CheckLayoutSupport(kvLayout_, ORI_KV_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(&opParamInfo_.oriKv.tensor->GetShape(), oriKvDimNumList, ORI_KV_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(kvLayout_, &opParamInfo_.oriKv.tensor->GetShape(), ORI_KV_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaCmpKv() const +{ + if (sasInfo_.perfMode == SASTemplateMode::SCFA_TEMPLATE_MODE || + sasInfo_.perfMode == SASTemplateMode::CFA_TEMPLATE_MODE) { + const std::vector cmpKvDimNumList = {DIM_NUM_THREE, DIM_NUM_FOUR}; + if ( + ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.cmpKv.desc, CMP_KV_NAME) || + ge::GRAPH_SUCCESS != CheckLayoutSupport(kvLayout_, CMP_KV_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(&opParamInfo_.cmpKv.tensor->GetShape(), cmpKvDimNumList, CMP_KV_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(kvLayout_, &opParamInfo_.cmpKv.tensor->GetShape(), CMP_KV_NAME)) { + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaNumHeads() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaKvHeadNums() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaCmpSparseIndices() const +{ + if (sasInfo_.perfMode == optiling::SASTemplateMode::SCFA_TEMPLATE_MODE){ + OP_CHECK_IF(opParamInfo_.cmpSparseIndices.tensor->GetStorageShape().GetShapeSize() == 0, + OP_LOGE(opName_, "when cmp_sparse_indices is not nullptr(SCFA), cmp_sparse_indices cannot be empty tensor."), + return ge::GRAPH_FAILED); + const std::vector cmpSparseIndicesDimNumList = {DIM_NUM_THREE, DIM_NUM_FOUR}; + if ( + ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.cmpSparseIndices.desc, CMP_SPARSE_INDICES) || + ge::GRAPH_SUCCESS != CheckLayoutSupport(cmpSparseIndicesLayout_, CMP_SPARSE_INDICES) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(&opParamInfo_.cmpSparseIndices.tensor->GetShape(), cmpSparseIndicesDimNumList, CMP_SPARSE_INDICES) || + ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(cmpSparseIndicesLayout_, &opParamInfo_.cmpSparseIndices.tensor->GetShape(), CMP_SPARSE_INDICES)) { return ge::GRAPH_FAILED; + } + if (cmpSparseIndicesLayout_ == SASLayout::TND) + { + OP_CHECK_IF(!(opParamInfo_.cmpSparseIndices.tensor->GetStorageShape().GetDim(DIM_NUM_THREE - 1) != 512 || \ + opParamInfo_.cmpSparseIndices.tensor->GetStorageShape().GetDim(DIM_NUM_THREE - 1) != 1024), + OP_LOGE(opName_, "K should be 512 or 1024, but got: %lld ", + opParamInfo_.cmpSparseIndices.tensor->GetStorageShape().GetDim(DIM_NUM_THREE - 1)), + return ge::GRAPH_FAILED); + } else{ + OP_CHECK_IF(!(opParamInfo_.cmpSparseIndices.tensor->GetStorageShape().GetDim(DIM_NUM_THREE - 1) != 512 || \ + opParamInfo_.cmpSparseIndices.tensor->GetStorageShape().GetDim(DIM_NUM_THREE - 1) != 1024), + OP_LOGE(opName_, "K should be 512 or 1024, but got: %lld ", + opParamInfo_.cmpSparseIndices.tensor->GetStorageShape().GetDim(DIM_NUM_FOUR - 1)), + return ge::GRAPH_FAILED); + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaOriBlockTable() const +{ + if (kvLayout_ == SASLayout::BSND) { + return ge::GRAPH_SUCCESS; // BSND 场景不需要使用oriBlockTable + } + if(kvLayout_ == SASLayout::TND) { + return ge::GRAPH_SUCCESS; + } + OP_CHECK_IF(opParamInfo_.oriBlockTable.tensor->GetStorageShape().GetShapeSize() == 0, + OP_LOGE(opName_, "ori_block_table cannot be empty tensor."), + return ge::GRAPH_FAILED); + const std::vector oriBlockTableDimNumList = {DIM_NUM_TWO}; + if ( + ge::GRAPH_SUCCESS != CheckDtypeSupport( opParamInfo_.oriBlockTable.desc, ORI_BLOCK_TABLE_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(&opParamInfo_.oriBlockTable.tensor->GetShape(), oriBlockTableDimNumList, ORI_BLOCK_TABLE_NAME)) { + return ge::GRAPH_FAILED; + } + OP_CHECK_IF((oriBlockSize_ <= 0 || oriBlockSize_ > BLOCK_SIZE_LIMIT || + (static_cast(oriBlockSize_) % 16 != 0UL)), + OP_LOGE(opName_, "ori_block_size should be in range [1, 1024], and be aligned to 16, but got: %d.", + oriBlockSize_), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaCmpBlockTable() const +{ + if (kvLayout_ == SASLayout::BSND) { + return ge::GRAPH_SUCCESS; // BSND 场景不需要使用oriBlockTable + } + if (kvLayout_ == SASLayout::TND) { + return ge::GRAPH_SUCCESS; + } + if (sasInfo_.perfMode == optiling::SASTemplateMode::SCFA_TEMPLATE_MODE || + sasInfo_.perfMode == optiling::SASTemplateMode::CFA_TEMPLATE_MODE){ + const std::vector cmpBlockTableDimNumList = {DIM_NUM_TWO}; + if ( + ge::GRAPH_SUCCESS != CheckDtypeSupport( opParamInfo_.cmpBlockTable.desc, CMP_BLOCK_TABLE_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(&opParamInfo_.cmpBlockTable.tensor->GetShape(), + cmpBlockTableDimNumList, CMP_BLOCK_TABLE_NAME)) { + return ge::GRAPH_FAILED; + } + OP_CHECK_IF((cmpBlockSize_ <= 0 || cmpBlockSize_ > BLOCK_SIZE_LIMIT || + (static_cast(cmpBlockSize_) % 16 != 0UL)), + OP_LOGE(opName_, "cmp_block_size should be in [1, 1024], and be aligned to 16, but got: %d.", + cmpBlockSize_), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaSinks() const +{ + OP_CHECK_IF(opParamInfo_.sinks.tensor->GetStorageShape().GetShapeSize() == 0, + OP_LOGE(opName_, "sinks cannot be empty tensor."), + return ge::GRAPH_FAILED); + if (opParamInfo_.sinks.tensor->GetStorageShape().GetDimNum() != DIM_NUM_ONE) { + OP_LOGE(opName_, "the dim num of %s is %u, it should be %u.", SINKS_NAME.c_str(), + opParamInfo_.sinks.tensor->GetStorageShape().GetDimNum(), DIM_NUM_ONE); + return ge::GRAPH_FAILED; + } + if (opParamInfo_.sinks.tensor->GetStorageShape().GetDim(0) != n1Size_) { + OP_LOGE(opName_, "%s's dimension(%ld) should be equal to query head num(%u).", SINKS_NAME.c_str(), + opParamInfo_.sinks.tensor->GetStorageShape().GetDim(0), n1Size_); + return ge::GRAPH_FAILED; + } + OP_CHECK_IF(opParamInfo_.sinks.desc->GetDataType() != ge::DT_FLOAT, + OP_LOGE(opName_, "sinks's dtype must be DT_FLOAT."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaMetadata() const +{ + if (opParamInfo_.metadata.tensor == nullptr) { + OP_LOGE(opName_, "%s must be provided!", METADATA_NAME.c_str()); + return ge::GRAPH_FAILED; + } + OP_CHECK_IF((opParamInfo_.metadata.tensor->GetShapeSize() != METADATA_LIMIT), + OP_LOGE(opName_, "input metadata dim 0 must be %u.", METADATA_LIMIT), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.metadata.desc->GetDataType() != ge::DT_INT32, + OP_LOGE(opName_, "metadata's dtype must be DT_INT32."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaCmpRatio() const +{ + if (sasInfo_.perfMode == optiling::SASTemplateMode::CFA_TEMPLATE_MODE || sasInfo_.perfMode == optiling::SASTemplateMode::SCFA_TEMPLATE_MODE){ + OP_CHECK_IF(cmpRatio_ != 128 && cmpRatio_ != 4, + OP_LOGE(opName_, "cmp_ratio should be 128 or 4, but got %u", cmpRatio_), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaOriMaskMode() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaCmpMaskMode() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaOriKvStride() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaCmpKvStride() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaOriWinLeft() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSingleParaOriWinRight() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckSinglePara() const +{ + if ( + ge::GRAPH_SUCCESS != CheckSingleParaQuery() || + ge::GRAPH_SUCCESS != CheckSingleParaOriKv() || + ge::GRAPH_SUCCESS != CheckSingleParaCmpKv() || + ge::GRAPH_SUCCESS != CheckSingleParaNumHeads() || + ge::GRAPH_SUCCESS != CheckSingleParaKvHeadNums() || + ge::GRAPH_SUCCESS != CheckSingleParaCmpSparseIndices() || + ge::GRAPH_SUCCESS != CheckSingleParaOriBlockTable() || + ge::GRAPH_SUCCESS != CheckSingleParaCmpBlockTable() || + ge::GRAPH_SUCCESS != CheckSingleParaSinks() || + ge::GRAPH_SUCCESS != CheckSingleParaMetadata() || + ge::GRAPH_SUCCESS != CheckSingleParaCmpRatio() || + ge::GRAPH_SUCCESS != CheckSingleParaOriMaskMode() || + ge::GRAPH_SUCCESS != CheckSingleParaCmpMaskMode() || + ge::GRAPH_SUCCESS != CheckSingleParaOriKvStride() || + ge::GRAPH_SUCCESS != CheckSingleParaCmpKvStride() || + ge::GRAPH_SUCCESS != CheckSingleParaOriWinLeft() || + ge::GRAPH_SUCCESS != CheckSingleParaOriWinRight()) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckExists(const void *pointer, const std::string &name) const +{ + OP_CHECK_IF(pointer == nullptr, + OP_LOGE(opName_, "%s should not be null", name.c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckNotExists(const void *pointer, const std::string &name) const +{ + OP_CHECK_IF(pointer != nullptr, + OP_LOGE(opName_, "%s should be null", name.c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckExistsByMap(const std::map ¶mMap) const +{ + for (const auto& kv : paramMap) { + if (CheckExists(kv.second, kv.first) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckNotExistsByMap(const std::map ¶mMap) const +{ + for (const auto& kv : paramMap) { + if (CheckNotExists(kv.second, kv.first) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckExistenceByMap(std::map &existMap, + std::map ¬ExistMap) const +{ + if (CheckExistsByMap(existMap) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + if (CheckNotExistsByMap(notExistMap) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckParaExistence() const +{ + if (kvLayout_ != SASLayout::PA_ND) { + return ge::GRAPH_SUCCESS; + } + std::map ParamExistMap = { + {"actualSeqLengths", opParamInfo_.sequsedKv.tensor}, + {"oriBlockTable", opParamInfo_.oriBlockTable.tensor}, + }; + std::map ParamNotExistMap = {}; + if (CheckExistenceByMap(ParamExistMap, ParamNotExistMap) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckFeatureShape() const +{ + OP_CHECK_IF(bSize_ <= 0, + OP_LOGE(opName_, "batch_size should be greater than 0, but got %u", bSize_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(qTSize_ <= 0 && (qLayout_ == SASLayout::TND), + OP_LOGE(opName_, "T_size of query should be greater than 0, but got %u", qTSize_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(n1Size_ % 4 != 0, + OP_LOGE(opName_, "q_head_num should be multiple of 4, but got %u", n1Size_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(n2Size_ != 1, + OP_LOGE(opName_, "kv_head_num should be 1, but got %u", n2Size_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(n1Size_ % n2Size_ != 0, + OP_LOGE(opName_, "q_head_num(%u) must be divisible by kv_head_num(%u)", n1Size_, n2Size_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(gSize_ % 4 != 0, + OP_LOGE(opName_, "group num should be multiple of 4, but got %u", gSize_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(qHeadDim_ != DIM_LIMIT, + OP_LOGE(opName_, "q_head_dim only support %u, but got %u", DIM_LIMIT, qHeadDim_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(oriKvHeadDim_ != DIM_LIMIT, + OP_LOGE(opName_, "ori_kv_head_dim only support %u, but got %u", DIM_LIMIT, oriKvHeadDim_), + return ge::GRAPH_FAILED); + if (!(sasInfo_.perfMode == SASTemplateMode::SWA_TEMPLATE_MODE)){ + OP_CHECK_IF(cmpKvHeadDim_ != DIM_LIMIT, + OP_LOGE(opName_, "cmp_kv_head_dim only support %u, but got %u", DIM_LIMIT, cmpKvHeadDim_), + return ge::GRAPH_FAILED); + } + + OP_CHECK_IF(!(qType_ == oriKvType_), + OP_LOGE(opName_, "Head dimension data type check failed! qType[%s] must be the same with oriKvType[%s].", + SASDataTypeToSerialString(qType_).c_str(), + SASDataTypeToSerialString(oriKvType_).c_str()), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(*opParamInfo_.oriMaskMode != 4, + OP_LOGE(opName_, "ori_mask_mode should be 4, but got %d", *opParamInfo_.oriMaskMode), + return ge::GRAPH_FAILED); + OP_CHECK_IF(*opParamInfo_.cmpMaskMode != 3, + OP_LOGE(opName_, "cmp_mask_mode should be 3, but got %d", *opParamInfo_.cmpMaskMode), + return ge::GRAPH_FAILED); + OP_CHECK_IF(oriWinLeft_ != 127, + OP_LOGE(opName_, "ori_win_left should be 127, but got %d", oriWinLeft_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(oriWinRight_ != 0, + OP_LOGE(opName_, "ori_win_right should be 0, but got %d", oriWinRight_), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckFeatureLayout() const +{ + const std::vector layoutQuerySupportList = { + "BSND", + "TND" + }; + std::string layoutQuery = opParamInfo_.layoutQ; + OP_CHECK_IF(std::find(layoutQuerySupportList.begin(), layoutQuerySupportList.end(), layoutQuery) == + layoutQuerySupportList.end(), + OP_LOGE(opName_, "layout_q only supports BSND/TND, but got %s", layoutQuery.c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckFeatureDtype() const +{ + OP_CHECK_IF(qType_ != ge::DT_BF16 && qType_ != ge::DT_FLOAT16, + OP_LOGE(opName_, "q dtype only support %s and %s, but got %s", + SASDataTypeToSerialString(ge::DT_BF16).c_str(), SASDataTypeToSerialString(ge::DT_FLOAT16).c_str(), + SASDataTypeToSerialString(qType_).c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckFeaturePa() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckFeature() const +{ + if (ge::GRAPH_SUCCESS != CheckFeatureShape() || + ge::GRAPH_SUCCESS != CheckFeatureLayout() || + ge::GRAPH_SUCCESS != CheckFeatureDtype() || + ge::GRAPH_SUCCESS != CheckFeaturePa()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +void SASTilingCheck::SetSASShapeCompare() +{ + queryShapeCmp_ = opParamInfo_.q.shape->GetStorageShape(); + oriKvShapeCmp_= opParamInfo_.oriKv.tensor->GetShape().GetStorageShape(); + attenOutShapeCmp_ = opParamInfo_.attnOut.shape->GetStorageShape(); + if (sasInfo_.perfMode == SASTemplateMode::CFA_TEMPLATE_MODE || + sasInfo_.perfMode == SASTemplateMode::SCFA_TEMPLATE_MODE) { + cmpKvShapeCmp_= opParamInfo_.cmpKv.tensor->GetShape().GetStorageShape(); + } + if (sasInfo_.perfMode == SASTemplateMode::SCFA_TEMPLATE_MODE) { + cmpKvSparseIndicesCmp_ = opParamInfo_.cmpSparseIndices.tensor->GetShape().GetStorageShape(); + } +} + +ge::graphStatus SASTilingCheck::CheckDTypeConsistency(const ge::DataType &actualDtype, + const ge::DataType &expectDtype, const std::string &name) const +{ + if (actualDtype != expectDtype) { + OP_LOGE(opName_, "%s dtype should be the same to %s, but it's %s.", name.c_str(), + SASDataTypeToSerialString(expectDtype).c_str(), + SASDataTypeToSerialString(actualDtype).c_str()); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckOriAndCmpKv() const +{ + OP_CHECK_IF(opParamInfo_.oriKv.tensor->GetStorageShape().GetShapeSize() == 0, + OP_LOGE(opName_, "ori_kv cannot be empty tensor."), + return ge::GRAPH_FAILED); + if (sasInfo_.perfMode == SASTemplateMode::CFA_TEMPLATE_MODE || + sasInfo_.perfMode == SASTemplateMode::SCFA_TEMPLATE_MODE) + { + if (opParamInfo_.cmpKv.tensor->GetStorageShape().GetDim(0) != 0 ) { + OP_CHECK_IF(opParamInfo_.cmpKv.tensor->GetStorageShape().GetShapeSize() == 0, + OP_LOGE(opName_, "cmp_kv cannot be empty tensor."), + return ge::GRAPH_FAILED); + } + if (ge::GRAPH_SUCCESS != CheckDTypeConsistency(cmpKvType_, + oriKvType_, CMP_KV_NAME)) { + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckAttenOut() const +{ + if (opParamInfo_.attnOut.desc != nullptr && opParamInfo_.attnOut.shape != nullptr) { + OP_CHECK_IF(opParamInfo_.attnOut.shape->GetStorageShape().GetShapeSize() == 0, + OP_LOGE(opName_, "attn_out cannot be empty tensor."), + return ge::GRAPH_FAILED); + } else{ + OP_LOGE(opName_, "attn_out cannot be nullptr."); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckActualSeqLensQ() const +{ + if (qLayout_ == SASLayout::TND) { + OP_CHECK_IF(opParamInfo_.cuSeqLensQ.tensor->GetStorageShape().GetShapeSize() == 0, + OP_LOGE(opName_, "when q's is TND, cu_seqlens_q cannot be empty tensor."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.cuSeqLensQ.desc->GetDataType() != ge::DT_INT32, + OP_LOGE(opName_, "when q's is TND, cu_seqlens_q's dtype msut be DT_INT32."), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckActualSeqLens() const +{ + return ge::GRAPH_SUCCESS; +} +ge::graphStatus SASTilingCheck::CheckBlockTable() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::CheckMultiParaConsistency() +{ + SetSASShapeCompare(); + if ( + ge::GRAPH_SUCCESS != CheckOriAndCmpKv() || + ge::GRAPH_SUCCESS != CheckAttenOut() || + ge::GRAPH_SUCCESS != CheckActualSeqLensQ() || + ge::GRAPH_SUCCESS != CheckActualSeqLens() || + ge::GRAPH_SUCCESS != CheckBlockTable()) + { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SASTilingCheck::Process() +{ + Init(); + if ( + CheckSinglePara() != ge::GRAPH_SUCCESS || + CheckParaExistence() != ge::GRAPH_SUCCESS || + CheckFeature() != ge::GRAPH_SUCCESS || + CheckMultiParaConsistency() != ge::GRAPH_SUCCESS + ) + { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +// --------------------------TilingPrepare函数定义------------------------------------- +static ge::graphStatus TilingPrepareForSparseAttnSharedkv(gert::TilingParseContext * /* context */) +{ + return ge::GRAPH_SUCCESS; +} + +void SparseAttnSharedkvTiling::CalcUbBmm(SASTilingInfo *tilingInfo) +{ + uint32_t cubeMSize = tilingInfo->gSize * tilingInfo->s1Size; + uint32_t maxMSize = mBaseSize_; + if (cubeMSize > maxMSize) { + cubeMSize = maxMSize; + } + mmResUbSize_ = sInnerSizeAlign_ * Align(cubeMSize, 16U);// kernel按照16对齐写出，tiling按照这个原则分配内存 + bmm2ResUbSize_ = headDimAlign_ * Align(cubeMSize, 16U);// kernel按照16对齐写出，tiling按照这个原则分配内存 +} + +void SparseAttnSharedkvTiling::SplitBalanced(SASTilingInfo *tilingInfo) +{ + uint32_t s2Size = tilingInfo->s2Size; + sInnerSizeAlign_ = Align(sInnerSize_, BYTE_BLOCK); // 元素个数按照基本块大小对齐 + mBaseSize_ = tilingInfo->gSize; + headDimAlign_ = Align(tilingInfo->qHeadDim, BYTE_BLOCK); + CalcUbBmm(tilingInfo); + + tilingData_.baseParams.set_mBaseSize(mBaseSize_); + tilingData_.baseParams.set_s2BaseSize(sInnerSize_); + tilingData_.baseParams.set_mmResUbSize(mmResUbSize_); + tilingData_.baseParams.set_bmm2ResUbSize(bmm2ResUbSize_); +} + +// --------------------------SparseAttnSharedkvTiling类成员函数定义----------------------- +ge::graphStatus SparseAttnSharedkvTiling::DoOpTiling(SASTilingInfo *tilingInfo) +{ + // -------------set blockdim----------------- + auto ascendcPlatform = platform_ascendc::PlatformAscendC(tilingInfo->platformInfo); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + uint32_t blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum); + context_->SetBlockDim(blockDim); + OP_LOGI(tilingInfo->opName, "SAS block dim: %u aiv Num: %u aic Num: %u.", blockDim, aivNum, aicNum); + + SplitBalanced(tilingInfo); + // -------------set workspacesize----------------- + constexpr uint32_t MM1_RES_ELEM_SIZE = 4; // 4: fp32 + constexpr uint32_t VEC1_RES_ELEM_SIZE = 2; // 2: fp16/bf16 + constexpr uint32_t MM2_RES_ELEM_SIZE = 4; // 4: fp32 + constexpr uint32_t VEC2_RES_ELEM_SIZE = 4; // 4: fp32 + constexpr uint32_t PRELOAD_NUM = 2; // preload数量 + + uint32_t workspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize(); + // 主流程需Workspace大小 + workspaceSize += PRELOAD_NUM * mmResUbSize_ * MM1_RES_ELEM_SIZE * aicNum; + workspaceSize += PRELOAD_NUM * mmResUbSize_ * VEC1_RES_ELEM_SIZE * aicNum; + workspaceSize += PRELOAD_NUM * bmm2ResUbSize_ * MM2_RES_ELEM_SIZE * aicNum; + workspaceSize += PRELOAD_NUM * bmm2ResUbSize_ * VEC2_RES_ELEM_SIZE * aicNum; + if (tilingInfo->perfMode == SASTemplateMode::SCFA_TEMPLATE_MODE) { + workspaceSize += 4 * 512 * 512 * 2 * aicNum; // 4:bufNum 512:s2Size 512:D 2:sizeof(half) + workspaceSize += 4 * 128 * 4 * (2 * aicNum); // 4:缓存有效mte2 size长度 128:份数 4:512B对齐长度 2:aiv数量 + } + size_t *workSpaces = context_->GetWorkspaceSizes(1); + workSpaces[0] = workspaceSize; + + // -------------set tilingdata----------------- + tilingData_.baseParams.set_batchSize(tilingInfo->bSize); + tilingData_.baseParams.set_kvSeqSize(tilingInfo->s2Size); + tilingData_.baseParams.set_qSeqSize(tilingInfo->s1Size); + tilingData_.baseParams.set_nNumOfQInOneGroup(tilingInfo->gSize); + tilingData_.baseParams.set_paBlockSize(tilingInfo->blockSize); + tilingData_.baseParams.set_oriBlockSize(tilingInfo->oriBlockSize); + tilingData_.baseParams.set_cmpBlockSize(tilingInfo->cmpBlockSize); + tilingData_.baseParams.set_oriMaxBlockNumPerBatch(tilingInfo->oriMaxBlockNumPerBatch); + tilingData_.baseParams.set_actualLenDimsQ(tilingInfo->actualLenDimsQ); + tilingData_.baseParams.set_actualLenDimsKV(tilingInfo->actualLenDimsKV); + + tilingData_.baseParams.set_softmaxScale(tilingInfo->softmaxScale); + tilingData_.baseParams.set_outputLayout(static_cast(tilingInfo->outLayout)); + tilingData_.baseParams.set_oriMaskMode(tilingInfo->oriMaskMode); + tilingData_.baseParams.set_oriKvStride(tilingInfo->oriKvStride); + tilingData_.baseParams.set_oriWinLeft(tilingInfo->oriWinLeft); + tilingData_.baseParams.set_oriWinRight(tilingInfo->oriWinRight); + tilingData_.baseParams.set_sparseBlockSize(tilingInfo->sparseBlockSize); + tilingData_.baseParams.set_returnSoftmaxLse(tilingInfo->returnSoftmaxLse); + + tilingData_.cmpParams.set_cmpMaxBlockNumPerBatch(tilingInfo->cmpMaxBlockNumPerBatch); + tilingData_.cmpParams.set_sparseBlockCount(tilingInfo->sparseBlockCount); + tilingData_.cmpParams.set_cmpRatio(tilingInfo->cmpRatio); + tilingData_.cmpParams.set_cmpMaskMode(tilingInfo->cmpMaskMode); + tilingData_.cmpParams.set_cmpKvStride(tilingInfo->cmpKvStride); + + usedCoreNum_ = aicNum; + tilingData_.baseParams.set_usedCoreNum(usedCoreNum_); + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + + // -------------set tilingkey----------------- + // FLASH_DECODE, LAYOUT_T, KV_LAYOUT_T, TEMPLATE_MODE + uint32_t qLayout = static_cast(tilingInfo->qLayout); + uint32_t inputKvLayout = static_cast(tilingInfo->kvLayout); + + uint32_t tilingKey = + GET_TPL_TILING_KEY(0U, qLayout, inputKvLayout, static_cast(tilingInfo->perfMode)); + context_->SetScheduleMode(1); + context_->SetTilingKey(tilingKey); + + return ge::GRAPH_SUCCESS; +} + +// --------------------------Tiling函数定义--------------------------- +ge::graphStatus TilingSparseAttnSharedkv(gert::TilingContext *context) +{ + OP_CHECK_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("SparseAttnSharedkv", "Tiling context is null."), + return ge::GRAPH_FAILED); + SASTilingInfo sasInfo; + SASInfoParser sasInfoParser(context); + if (sasInfoParser.Parse(sasInfo) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + SASTilingCheck sasTilingChecker(sasInfo); + if (sasTilingChecker.Process() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + SparseAttnSharedkvTiling tiling(context); + return tiling.DoOpTiling(&sasInfo); +} +// --------------------------Tiling函数及TilingPrepare函数注册-------- +IMPL_OP_OPTILING(SparseAttnSharedkv) + .Tiling(TilingSparseAttnSharedkv) + .TilingParse(TilingPrepareForSparseAttnSharedkv); +} // namespace optiling diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_tiling.h b/csrc/ascend/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_tiling.h new file mode 100644 index 000000000..87cd1fad2 --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_host/sparse_attn_sharedkv_tiling.h @@ -0,0 +1,563 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_tiling.h + * \brief + */ +#ifndef SPARSE_ATTN_SHAREDKV_TILING_H +#define SPARSE_ATTN_SHAREDKV_TILING_H + +#include +#include +#include +#include "register/tilingdata_base.h" +#include "register/op_def_registry.h" +#include "tiling/tiling_api.h" +#include "log/log.h" +#include "log/error_code.h" +#include "err/ops_err.h" +#include "platform/platform_info.h" + +namespace optiling { +// ------------------公共定义-------------------------- +struct SASTilingRequiredParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; +}; + +struct SASTilingOptionalParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::Tensor *tensor; +}; + +enum class SASLayout : uint32_t { + BSND = 0, + TND = 1, + PA_ND = 2 +}; + +enum class SASAxis : uint32_t { + B = 0, + S = 1, + N = 2, + D = 3, + K = 3, // sparse_indices的K和key的D枚举值相同，表达相同位置, 最后一维 + T = 5, + Bn = 6, // block number + Bs = 7 // block size +}; + +enum class SASTemplateMode : uint32_t { + SWA_TEMPLATE_MODE = 0, + CFA_TEMPLATE_MODE = 1, + SCFA_TEMPLATE_MODE = 2 +}; + +enum class KvStorageMode : uint32_t { + BATCH_CONTINUOUS = 0, + TENSOR_LIST = 1, + PAGE_ATTENTION = 2 +}; + +// ------------------算子原型索引常量定义---------------- +// Inputs Index +constexpr uint32_t Q_INDEX = 0; +constexpr uint32_t ORI_KV_INDEX = 1; +constexpr uint32_t CMP_KV_INDEX = 2; +constexpr uint32_t ORI_SPARSE_INDICES_INDEX = 3; +constexpr uint32_t CMP_SPARSE_INDICES_INDEX = 4; +constexpr uint32_t ORI_BLOCK_TABLE_INDEX = 5; +constexpr uint32_t CMP_BLOCK_TABLE_INDEX = 6; +constexpr uint32_t CU_SEQLENS_Q_INDEX = 7; +constexpr uint32_t CU_SEQLENS_KV_INDEX = 8; +constexpr uint32_t CU_SEQLENS_CMP_KV_INDEX = 9; +constexpr uint32_t SEQUSED_Q_INDEX = 10; +constexpr uint32_t SEQUSED_KV_INDEX = 11; +constexpr uint32_t SINKS_INDEX = 12; +constexpr uint32_t METADATA_INDEX = 13; +// Outputs Index +constexpr uint32_t ATTN_OUT_INDEX = 0; + +// Attributes Index +constexpr uint32_t ATTR_SOFTMAX_SCALE_INDEX = 0; +constexpr uint32_t ATTR_CMP_RATIO_INDEX = 1; +constexpr uint32_t ATTR_ORI_MASK_MODE_INDEX = 2; +constexpr uint32_t ATTR_CMP_MASK_MODE_INDEX = 3; +constexpr uint32_t ATTR_ORI_KV_STRIDE_INDEX = 4; +constexpr uint32_t ATTR_CMP_KV_STRIDE_INDEX = 5; +constexpr uint32_t ATTR_ORI_WIN_LEFT_INDEX = 6; +constexpr uint32_t ATTR_ORI_WIN_RIGHT_INDEX = 7; +constexpr uint32_t ATTR_LAYOUT_Q_INDEX = 8; +constexpr uint32_t ATTR_LAYOUT_KV_INDEX = 9; +constexpr uint32_t ATTR_RETURN_SOFTMAX_LSE = 10; + +// Dim Index +constexpr uint32_t DIM_IDX_ONE = 1; +constexpr uint32_t DIM_IDX_TWO = 2; +constexpr uint32_t DIM_IDX_THREE = 3; +constexpr uint32_t DIM_IDX_FOUR = 4; + +// Dim Num +constexpr uint32_t DIM_NUM_ONE = 1; +constexpr uint32_t DIM_NUM_TWO = 2; +constexpr uint32_t DIM_NUM_THREE = 3; +constexpr uint32_t DIM_NUM_FOUR = 4; + +// 常量 +constexpr uint32_t MAX_BLOCK_SIZE = 1024; +constexpr uint32_t COPYND2NZ_SRC_STRIDE_LIMITATION = 65535; +constexpr uint32_t NUM_BYTES_FLOAT = 4; +constexpr uint32_t NUM_BYTES_FLOAT16 = 2; +constexpr uint32_t NUM_BYTES_BF16 = 2; +constexpr uint32_t BYTE_BLOCK = 32; + +// 入参限制常量 +constexpr uint32_t HEAD_DIM_LIMIT = 128; +constexpr uint32_t SPARSE_LIMIT = 2048; +constexpr uint32_t SPARSE_MODE_LOWER = 3; +constexpr uint32_t METADATA_LIMIT = 1024; +constexpr uint32_t DIM_LIMIT = 512; +constexpr uint32_t TOPK_LIMIT = 1024; +constexpr uint32_t BLOCK_SIZE_LIMIT = 1024; + +// -----------算子TilingData定义--------------- +BEGIN_TILING_DATA_DEF(SparseAttnSharedkvSwaParams) +TILING_DATA_FIELD_DEF(uint32_t, batchSize) +TILING_DATA_FIELD_DEF(uint32_t, qSeqSize) +TILING_DATA_FIELD_DEF(uint32_t, kvSeqSize) +TILING_DATA_FIELD_DEF(int64_t, paBlockSize) +TILING_DATA_FIELD_DEF(int64_t, oriBlockSize) +TILING_DATA_FIELD_DEF(int64_t, cmpBlockSize) +TILING_DATA_FIELD_DEF(uint32_t, oriMaxBlockNumPerBatch) +TILING_DATA_FIELD_DEF(uint32_t, nNumOfQInOneGroup) +TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsQ) +TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsKV) + +TILING_DATA_FIELD_DEF(float, softmaxScale) // 即 scaleValue +TILING_DATA_FIELD_DEF(uint32_t, outputLayout) +TILING_DATA_FIELD_DEF(uint64_t, oriMaskMode) +TILING_DATA_FIELD_DEF(int64_t, oriKvStride) +TILING_DATA_FIELD_DEF(int64_t, oriWinLeft) +TILING_DATA_FIELD_DEF(int64_t, oriWinRight) +TILING_DATA_FIELD_DEF(int64_t, sparseBlockSize) + +TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum); + +TILING_DATA_FIELD_DEF(uint32_t, mmResUbSize); +TILING_DATA_FIELD_DEF(uint32_t, bmm2ResUbSize); + +TILING_DATA_FIELD_DEF(uint32_t, mBaseSize) +TILING_DATA_FIELD_DEF(uint32_t, s2BaseSize) +TILING_DATA_FIELD_DEF(bool, returnSoftmaxLse) +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(SparseAttnSharedkvSwaParamsOp, SparseAttnSharedkvSwaParams) + +BEGIN_TILING_DATA_DEF(SparseAttnSharedkvCmpParams) + +TILING_DATA_FIELD_DEF(uint32_t, cmpMaxBlockNumPerBatch) +TILING_DATA_FIELD_DEF(uint32_t, sparseBlockCount) +TILING_DATA_FIELD_DEF(int64_t, cmpRatio) +TILING_DATA_FIELD_DEF(uint64_t, cmpMaskMode) +TILING_DATA_FIELD_DEF(int64_t, cmpKvStride) +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(SparseAttnSharedkvCmpParamsOp, SparseAttnSharedkvCmpParams) + +BEGIN_TILING_DATA_DEF(SparseAttnSharedkvTilingData) +TILING_DATA_FIELD_DEF_STRUCT(SparseAttnSharedkvSwaParams, baseParams); +TILING_DATA_FIELD_DEF_STRUCT(SparseAttnSharedkvCmpParams, cmpParams); +END_TILING_DATA_DEF + +REGISTER_TILING_DATA_CLASS(SparseAttnSharedkv, SparseAttnSharedkvTilingData) + +struct SASParaInfo { + SASTilingRequiredParaInfo q = {nullptr, nullptr}; + SASTilingOptionalParaInfo oriKv = {nullptr, nullptr}; + SASTilingOptionalParaInfo cmpKv = {nullptr, nullptr}; + SASTilingOptionalParaInfo oriSparseIndices = {nullptr, nullptr}; + SASTilingOptionalParaInfo cmpSparseIndices = {nullptr, nullptr}; + SASTilingOptionalParaInfo oriBlockTable = {nullptr, nullptr}; + SASTilingOptionalParaInfo cmpBlockTable = {nullptr, nullptr}; + SASTilingOptionalParaInfo cuSeqLensQ = {nullptr, nullptr}; + SASTilingOptionalParaInfo seqUsedQ = {nullptr, nullptr}; + SASTilingOptionalParaInfo cuSeqLensKv = {nullptr, nullptr}; + SASTilingOptionalParaInfo cuSeqLensCmpKv = {nullptr, nullptr}; + SASTilingOptionalParaInfo sequsedKv = {nullptr, nullptr}; + SASTilingOptionalParaInfo sinks = {nullptr, nullptr}; + SASTilingOptionalParaInfo metadata = {nullptr, nullptr}; + SASTilingRequiredParaInfo attnOut = {nullptr, nullptr}; + + const float *softmaxScale = nullptr; + const uint32_t *cmpRatio = nullptr; + const uint32_t *oriMaskMode = nullptr; + const uint32_t *cmpMaskMode = nullptr; + const uint32_t *oriKvStride = nullptr; + const uint32_t *cmpKvStride = nullptr; + const uint32_t *oriWinLeft = nullptr; + const uint32_t *oriWinRight = nullptr; + const char *layoutQ = nullptr; + const char *layoutKv = nullptr; + const bool *returnSoftmaxLse = nullptr; +}; + +static std::string SASDataTypeToSerialString(ge::DataType type); +std::string SASLayoutToSerialString(SASLayout layout); + +// -----------算子Tiling入参信息类--------------- +class SASTilingInfo { +public: + const char *opName = nullptr; + fe::PlatFormInfos *platformInfo = nullptr; + SASParaInfo opParamInfo; + + // Base Param + platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND910B; + uint32_t bSize = 0; + uint32_t n1Size = 0; + uint32_t n2Size = 0; + uint32_t s1Size = 0; + int64_t s2Size = 0; + uint32_t gSize = 0; + uint32_t qHeadDim = 0; + uint32_t oriKvHeadDim = 0; + uint32_t cmpKvHeadDim = 0; + uint32_t qTSize = 0; // 仅TND时生效 + + uint32_t actualLenDimsQ = 0; + uint32_t maxActualseq = 0; + bool actualSeqLenFlag = false; + bool isSameSeqAllKVTensor = true; + bool isSameActualseq = true; + uint32_t actualLenDimsKV = 0; + + float softmaxScale = 0; + int64_t cmpRatio = 1; + uint64_t oriMaskMode = 0; + uint64_t cmpMaskMode = 0; + uint64_t oriKvStride = 0; + uint64_t cmpKvStride = 0; + int64_t oriWinLeft = 0; + int64_t oriWinRight = 0; + int64_t sparseBlockSize = 0; + int64_t sparseBlockCount = 0; + // Mask + int32_t sparseMode = 0; + // Others Flag + uint32_t sparseCount = 0; + + bool returnSoftmaxLse = false; + // PageAttention + uint32_t blockTypeSize = 0; + uint32_t oriMaxBlockNumPerBatch = 0; + int32_t blockSize = 0; + int32_t oriBlockSize = 0; + int32_t cmpBlockSize = 0; + uint32_t cmpMaxBlockNumPerBatch = 0; + uint32_t totalBlockNum = 0; + + // DType + ge::DataType qType = ge::DT_FLOAT16; + ge::DataType oriKvType = ge::DT_FLOAT16; + ge::DataType cmpKvType = ge::DT_FLOAT16; + ge::DataType outputType = ge::DT_FLOAT16; + + // Layout + SASLayout qLayout = SASLayout::TND; + SASLayout cmpSparseIndicesLayout = SASLayout::TND; + SASLayout oriSparseIndicesLayout = SASLayout::TND; + SASLayout kvLayout = SASLayout::PA_ND; + SASLayout outLayout = SASLayout::BSND; + + // template mode + SASTemplateMode perfMode = SASTemplateMode::SWA_TEMPLATE_MODE; +}; + +// -----------算子Tiling入参信息解析及Check类--------------- +class SASTilingCheck { +public: + explicit SASTilingCheck(const SASTilingInfo &sasInfo) : sasInfo_(sasInfo) {}; + ~SASTilingCheck() = default; + virtual ge::graphStatus Process(); + +private: + void Init(); + + void LogErrorDtypeSupport(const std::vector &expectDtypeList, + const ge::DataType &actualDtype, const std::string &name) const; + ge::graphStatus CheckLayoutSupport(const SASLayout &actualLayout, const std::string &name) const; + template + void LogErrorDimNumSupport(const std::vector &expectNumberList, + const T &actualValue, const std::string &name) const; + template + void LogErrorNumberSupport(const std::vector &expectNumberList, + const T &actualValue, const std::string &name, const std::string subName) const; + ge::graphStatus CheckDimNumSupport(const gert::StorageShape *shape, + const std::vector &expectDimNumList, const std::string &name) const; + void LogErrorLayoutSupport(const std::vector &expectLayoutList, + const SASLayout &actualLayout, const std::string &name) const; + ge::graphStatus CheckDimNumInLayoutSupport(const SASLayout &layout, + const gert::StorageShape *shape, const std::string &name) const; + ge::graphStatus CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc, + const std::string &name) const; + ge::graphStatus CheckSinglePara() const; + ge::graphStatus CheckSingleParaQuery() const; + ge::graphStatus CheckSingleParaOriKv() const; + ge::graphStatus CheckSingleParaCmpKv() const; + ge::graphStatus CheckSingleParaNumHeads() const; + ge::graphStatus CheckSingleParaKvHeadNums() const; + ge::graphStatus CheckSingleParaOriSparseIndices() const; + ge::graphStatus CheckSingleParaCmpSparseIndices() const; + ge::graphStatus CheckSingleParaSinks() const; + ge::graphStatus CheckSingleParaMetadata() const; + ge::graphStatus CheckSingleParaCmpRatio() const; + ge::graphStatus CheckSingleParaOriMaskMode() const; + ge::graphStatus CheckSingleParaCmpMaskMode() const; + ge::graphStatus CheckSingleParaOriKvStride() const; + ge::graphStatus CheckSingleParaCmpKvStride() const; + ge::graphStatus CheckSingleParaOriWinLeft() const; + ge::graphStatus CheckSingleParaOriWinRight() const; + ge::graphStatus CheckSingleParaOriBlockTable() const; + ge::graphStatus CheckSingleParaCmpBlockTable() const; + + ge::graphStatus CheckParaExistence() const; + ge::graphStatus CheckExists(const void *pointer, const std::string &name) const; + ge::graphStatus CheckNotExists(const void *pointer, const std::string &name) const; + ge::graphStatus CheckExistsByMap(const std::map ¶mMap) const; + ge::graphStatus CheckNotExistsByMap(const std::map ¶mMap) const; + ge::graphStatus CheckExistenceByMap(std::map &existMap, + std::map ¬ExistMap) const; + + ge::graphStatus CheckFeature() const; + ge::graphStatus CheckFeatureShape() const; + ge::graphStatus CheckFeatureLayout() const; + ge::graphStatus CheckFeatureDtype() const; + ge::graphStatus CheckFeaturePa() const; + + ge::graphStatus CheckMultiParaConsistency(); + void SetSASShapeCompare(); + ge::graphStatus CheckDTypeConsistency(const ge::DataType &actualDtype, + const ge::DataType &expectDtype, const std::string &name) const; + ge::graphStatus CheckOriAndCmpKv() const; + ge::graphStatus CheckAttenOut() const; + ge::graphStatus CheckActualSeqLensQ() const; + ge::graphStatus CheckActualSeqLens() const; + ge::graphStatus CheckBlockTable() const; + + gert::Shape queryShapeCmp_{}; + gert::Shape oriKvShapeCmp_{}; + gert::Shape cmpKvShapeCmp_{}; + gert::Shape oriKvSparseIndicesCmp_{}; + gert::Shape cmpKvSparseIndicesCmp_{}; + gert::Shape attenOutShapeCmp_{}; + +private: + const char *opName_; + fe::PlatFormInfos *platformInfo_; + SASParaInfo opParamInfo_; + const SASTilingInfo &sasInfo_; + + uint32_t bSize_ = 0; + uint32_t n1Size_ = 0; + uint32_t n2Size_ = 0; + uint32_t gSize_ = 0; + uint32_t s1Size_ = 0; + int64_t s2Size_ = 0; + uint32_t qHeadDim_ = 0; + uint32_t oriKvHeadDim_ = 0; + uint32_t cmpKvHeadDim_ = 0; + + uint32_t qTSize_ = 0; // 仅TND时生效 + uint32_t kvTSize_ = 0; // 仅TND时生效 + int64_t cmpRatio_ = 1; + KvStorageMode kvStorageMode_ = KvStorageMode::BATCH_CONTINUOUS; + uint32_t sparseBlockCount_ = 0; + int64_t oriWinLeft_ = 0; + int64_t oriWinRight_ = 0; + SASLayout qLayout_ = SASLayout::TND; + SASLayout cmpSparseIndicesLayout_ = SASLayout::TND; + SASLayout oriSparseIndicesLayout_ = SASLayout::TND; + SASLayout outLayout_ = SASLayout::TND; + SASLayout kvLayout_ = SASLayout::PA_ND; + + uint32_t oriMaxBlockNumPerBatch_ = 0; + uint32_t cmpMaxBlockNumPerBatch_ = 0; + int64_t blockSize_ = 0; + int32_t oriBlockSize_ = 0; + int32_t cmpBlockSize_ = 0; + + uint32_t aicNum_ = 0; + uint32_t aivNum_ = 0; + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B; + uint64_t l2CacheSize_ = 0; + + bool isSameSeqAllKVTensor_ = true; + bool isSameActualseq_ = true; + uint32_t maxActualseq_ = 0; + + ge::DataType qType_ = ge::DT_FLOAT16; + ge::DataType oriKvType_ = ge::DT_FLOAT16; + ge::DataType cmpKvType_ = ge::DT_FLOAT16; + ge::DataType outputType_ = ge::DT_FLOAT16; +}; + +template inline T Align(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd) * (rnd))); +} + +class SASInfoParser { +public: + explicit SASInfoParser(gert::TilingContext *context) : context_(context) {} + ~SASInfoParser() = default; + + ge::graphStatus CheckRequiredInOutExistence() const; + ge::graphStatus CheckRequiredAttrExistence() const; + ge::graphStatus CheckRequiredParaExistence() const; + ge::graphStatus CheckUnrequiredParaExistence() const; + + ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + SASLayout &layout, const std::string &name) const; + ge::graphStatus GetActualSeqLenQSize(uint32_t &size); + ge::graphStatus GetOpName(); + ge::graphStatus GetNpuInfo(); + void GetOptionalInputParaInfo(); + void GetInputParaInfo(); + void GetOutputParaInfo(); + ge::graphStatus GetAttrParaInfo(); + ge::graphStatus GetKvCache(); + ge::graphStatus GetOpParaInfo(); + + ge::graphStatus GetInOutDataType(); + ge::graphStatus GetQueryAndOutLayout(); + ge::graphStatus GetKvLayout(); + ge::graphStatus GetSASTemplateMode(SASTilingInfo &sasInfo); + void SetSASShape(); + ge::graphStatus GetN1Size(); + ge::graphStatus GetN2Size(); + ge::graphStatus GetGSize(); + ge::graphStatus GetBatchSize(); + ge::graphStatus GetQTSize(); + ge::graphStatus GetKVTSize(); + ge::graphStatus GetS1Size(); + ge::graphStatus GetS2SizeForPageAttention(); + ge::graphStatus GetS2SizeForTND(); + ge::graphStatus GetS2Size(); + ge::graphStatus GetMaxBlockNumPerBatch(); + ge::graphStatus GetBlockSize(); + ge::graphStatus GetQHeadDim(); + ge::graphStatus GetValueHeadDim(); + ge::graphStatus GetSparseBlockCount(); + ge::graphStatus GetActualseqInfo(); + ge::graphStatus GetSinks(); + void GenerateInfo(SASTilingInfo &sasInfo); + ge::graphStatus Parse(SASTilingInfo &sasInfo); + +public: + gert::TilingContext *context_ = nullptr; + const char *opName_; + fe::PlatFormInfos *platformInfo_; + SASParaInfo opParamInfo_; + + bool HasAxis(const SASAxis &axis, const SASLayout &layout, const gert::Shape &shape) const; + size_t GetAxisIdx(const SASAxis &axis, const SASLayout &layout) const; + uint32_t GetAxisNum(const gert::Shape &shape, const SASAxis &axis,const SASLayout &layout) const; + static constexpr int64_t invalidDimValue_ = std::numeric_limits::min(); + + // BaseParams + uint32_t bSize_ = 0; + uint32_t n1Size_ = 0; + uint32_t n2Size_ = 0; + uint32_t gSize_ = 0; + uint32_t s1Size_ = 0; + int64_t s2Size_ = 0; + uint32_t headDim_ = 0; + uint32_t qTSize_ = 0; + uint32_t orikvTSize_ = 0; + uint32_t cmpkvTSize_ = 0; + uint32_t qHeadDim_ = 0; + uint32_t oriKvHeadDim_ = 0; + uint32_t cmpKvHeadDim_ = 0; + int64_t sparseBlockSize_ = 0; + int64_t sparseBlockCount_ = 0; + int64_t oriWinLeft_ = 0; + int64_t oriWinRight_ = 0; + uint32_t maxActualseq_ = 0; + bool isSameSeqAllKVTensor_ = true; + uint32_t actualLenDimsKV_ = 0; + uint32_t actualLenDimsQ_ = 0; + + uint32_t aicNum_ = 0; + uint32_t aivNum_ = 0; + // Layout + SASLayout qLayout_ = SASLayout::TND; + SASLayout cmpSparseIndicesLayout_ = SASLayout::TND; + SASLayout oriSparseIndicesLayout_ = SASLayout::TND; + SASLayout outLayout_ = SASLayout::BSND; + SASLayout kvLayout_ = SASLayout::PA_ND; + // PageAttention + uint32_t oriMaxBlockNumPerBatch_ = 0; + uint32_t cmpMaxBlockNumPerBatch_ = 0; + int32_t oriBlockSize_ = 0; + int32_t cmpBlockSize_ = 0; + + // template mode + SASTemplateMode perfMode_ = SASTemplateMode::SWA_TEMPLATE_MODE; + + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B; + ge::DataType qType_ = ge::DT_FLOAT16; + ge::DataType oriKvType_ = ge::DT_FLOAT16; + ge::DataType cmpKvType_ = ge::DT_FLOAT16; + ge::DataType cmpSparseIndicesType_ = ge::DT_INT32; + ge::DataType oriBlockTableType_ = ge::DT_INT32; + ge::DataType cmpBlockTableType_ = ge::DT_INT32; + ge::DataType cuSeqLensQType_ = ge::DT_INT32; + ge::DataType seqsedKvType_ = ge::DT_INT32; + ge::DataType sinksType_ = ge::DT_INT32; + ge::DataType metadataType_ = ge::DT_INT32; + ge::DataType outputType_ = ge::DT_FLOAT16; + + gert::Shape qShape_{}; + gert::Shape oriKvShape_{}; + gert::Shape cmpKvShape_{}; + gert::Shape oriSparseIndicesShape_{}; + gert::Shape cmpSparseIndicesShape_{}; +}; + +// ---------------算子Tiling类--------------- +class SparseAttnSharedkvTiling { +public: + explicit SparseAttnSharedkvTiling(gert::TilingContext *context) : context_(context){}; + ge::graphStatus DoOpTiling(SASTilingInfo *tilingInfo); + +private: + void SplitBalanced(SASTilingInfo *tilingInfo); + void CalcUbBmm(SASTilingInfo *tilingInfo); + gert::TilingContext *context_ = nullptr; + SASTemplateMode perfMode_ = SASTemplateMode::SWA_TEMPLATE_MODE; + SparseAttnSharedkvTilingData tilingData_; + uint32_t blockDim_{0}; + uint64_t workspaceSize_{0}; + uint64_t tilingKey_{0}; + + SASTilingInfo *sasInfo_ = nullptr; + + size_t mmResUbSize_ = 0; + size_t bmm2ResUbSize_ = 0; + uint32_t sInnerLoopTimes_ = 0; + uint32_t sInnerSize_ = 512; // s2固定切分512 + uint32_t sInnerSizeAlign_ = 0; + uint32_t usedCoreNum_ = 0; + + uint32_t headDimAlign_ = 0; + uint32_t mBaseSize_ = 64; +}; + +} +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_scfa_block_cube.h b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_scfa_block_cube.h new file mode 100644 index 000000000..f5ad02487 --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_scfa_block_cube.h @@ -0,0 +1,853 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_scfa_block_cube.h + * \brief use 7 buffer for matmul l1, better pipeline + */ +#ifndef SPARSE_ATTN_SHAREDKV_SCFA_BLOCK_CUBE_H +#define SPARSE_ATTN_SHAREDKV_SCFA_BLOCK_CUBE_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../sparse_attn_sharedkv_common.h" + +namespace SASKernel { +template +class SASCubeBlock { +public: + // 中间计算数据类型为float, 高精度模式 + using T = float; + using Q_T = typename SAST::queryType; + using KV_T = typename SAST::kvType; + using OUT_T = typename SAST::outputType; + using MM_OUT_T = T; + + __aicore__ inline SASCubeBlock(){}; + __aicore__ inline void InitParams(const ConstInfo &constInfo); + __aicore__ inline void InitMm1GlobalTensor(GlobalTensor queryGm, GlobalTensor oriKvGm, + GlobalTensor cmpKV, GlobalTensor mm1ResGm); + __aicore__ inline void InitMm2GlobalTensor(GlobalTensor vec1ResGm, GlobalTensor mm2ResGm, + GlobalTensor attentionOutGm); + __aicore__ inline void InitPageAttentionInfo(GlobalTensor oriKvGm, const GlobalTensor &kvMergeGm, + GlobalTensor oriBlockTableGm, + GlobalTensor cmpBlockTableGm); + __aicore__ inline void InitBuffers(TPipe *pipe); + + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void ComputeMm1(const RunInfo &info, const MSplitInfo mSplitInfo); + __aicore__ inline void ComputeMm2(const RunInfo &info, const MSplitInfo mSplitInfo); + +private: + static constexpr bool PAGE_ATTENTION = SAST::pageAttention; + static constexpr int TEMPLATE_MODE = SAST::templateMode; + static constexpr bool FLASH_DECODE = SAST::flashDecode; + static constexpr SAS_LAYOUT LAYOUT_T = SAST::layout; + static constexpr SAS_LAYOUT KV_LAYOUT_T = SAST::kvLayout; + + static constexpr uint32_t M_SPLIT_SIZE = 128; // m方向切分 + static constexpr uint32_t N_SPLIT_SIZE = 128; // n方向切分 + static constexpr uint32_t K_L0_SPLIT_SIZE = 128; // k方向L0切分 + static constexpr uint32_t K_L1_SPLIT_SIZE = 256; // k方向L1切分 + static constexpr uint32_t N_WORKSPACE_SIZE = 512; // n方向切分 + static constexpr uint32_t D_SPLIT_SIZE = 256; // d轴切分 + + + static constexpr uint32_t L1_BLOCK_SIZE = (64 * 512 * sizeof(Q_T)); + static constexpr uint32_t L1_BLOCK_OFFSET = 64 * 512; + + static constexpr uint32_t L0A_PP_SIZE = (32 * 1024); + static constexpr uint32_t L0B_PP_SIZE = (32 * 1024); + static constexpr uint32_t L0C_PP_SIZE = (64 * 1024); + + // mte2 <> mte1 EventID + // L1 3buf, 使用3个eventId + static constexpr uint32_t L1_EVENT0 = EVENT_ID2; + static constexpr uint32_t L1_EVENT1 = EVENT_ID3; + static constexpr uint32_t L1_EVENT2 = EVENT_ID4; + static constexpr uint32_t L1_EVENT3 = EVENT_ID5; + static constexpr uint32_t L1_EVENT4 = EVENT_ID6; + static constexpr uint32_t L1_EVENT5 = EVENT_ID7; + static constexpr uint32_t L1_EVENT6 = EVENT_ID1; + + // m <> mte1 EventID + static constexpr uint32_t L0AB_EVENT0 = EVENT_ID3; + static constexpr uint32_t L0AB_EVENT1 = EVENT_ID4; + + static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding + static constexpr uint32_t mte21QPIds[4] = {L1_EVENT0, L1_EVENT1, L1_EVENT2, L1_EVENT3}; // mte12复用 + static constexpr uint32_t mte21KVIds[3] = {L1_EVENT4, L1_EVENT5, L1_EVENT6}; + + ConstInfo constInfo{}; + + // L1分成3块buf, 用于记录 + uint32_t qpL1BufIter = 0; + uint32_t kvL1BufIter = -1; + uint32_t abL0BufIter = 0; + uint32_t cL0BufIter = 0; + + // mm1 + GlobalTensor queryGm; + GlobalTensor keyGm; + GlobalTensor mm1ResGm; + GlobalTensor oriKvGm; + GlobalTensor kvMergeGm_; + GlobalTensor cmpKvGm; + + // mm2 + GlobalTensor vec1ResGm; + GlobalTensor valueGm; + GlobalTensor mm2ResGm; + GlobalTensor attentionOutGm; + + // block_table + GlobalTensor oriBlockTableGm; + GlobalTensor cmpBlockTableGm; + + TBuf bufQPL1; + TBuf bufKVL1; + TBuf tmpBufL0A; + TBuf tmpBufL0B; + TBuf tmpBufL0C; + + LocalTensor l1QPTensor; + LocalTensor l1KVTensor; + LocalTensor aL0TensorPingPong; + LocalTensor bL0TensorPingPong; + LocalTensor cL0TensorPingPong; + + // L0AB m <> mte1 EventID + __aicore__ inline uint32_t Mte1MmABEventId(uint32_t idx) + { + return (L0AB_EVENT0 + idx); + } + + __aicore__ inline uint32_t GetQPL1RealIdx(uint32_t mIdx, uint32_t k1Idx) + { + uint32_t idxMap[] = {0, 2}; // 确保0块和1块连在一起, 2和3块连在一起, 来保证同一m块的地址相连 + return idxMap[mIdx % 2] + k1Idx; + } + + __aicore__ inline void CopyGmToL1(LocalTensor &l1Tensor, GlobalTensor &gmSrcTensor, uint32_t srcN, + uint32_t srcD, uint32_t srcDstride); + __aicore__ inline void CopyInMm1AToL1(LocalTensor &aL1Tensor, const RunInfo &info, uint32_t mSeqIdx, + uint32_t mSizeAct, uint32_t headSize, uint32_t headOffset); + + __aicore__ inline void CopyInMm2AToL1(LocalTensor &aL1Tensor, const RunInfo &info, uint32_t mSeqIdx, + uint32_t subMSizeAct, uint32_t nSize, uint32_t nOffset); + __aicore__ inline void LoadDataMm1A(LocalTensor &aL0Tensor, LocalTensor &aL1Tensor, uint32_t idx, + uint32_t kSplitSize, uint32_t mSize, uint32_t kSize); + __aicore__ inline void LoadDataMm1B(LocalTensor &bL0Tensor, LocalTensor &bL1Tensor, uint32_t idx, + uint32_t kSplitSize, uint32_t kSize, uint32_t nSize); +}; + +template +__aicore__ inline void SASCubeBlock::InitParams(const ConstInfo &constInfo) +{ + this->constInfo = constInfo; +} + +template +__aicore__ inline void SASCubeBlock::InitMm1GlobalTensor(GlobalTensor queryGm, GlobalTensor oriKvGm, + GlobalTensor cmpKvGm, + GlobalTensor mm1ResGm) +{ + // mm1 + this->queryGm = queryGm; + this->oriKvGm = oriKvGm; + this->cmpKvGm = cmpKvGm; + this->mm1ResGm = mm1ResGm; +} + +template +__aicore__ inline void SASCubeBlock::InitMm2GlobalTensor(GlobalTensor vec1ResGm, + GlobalTensor mm2ResGm, + GlobalTensor attentionOutGm) +{ + // mm2 + this->vec1ResGm = vec1ResGm; + this->mm2ResGm = mm2ResGm; + this->attentionOutGm = attentionOutGm; +} + +template +__aicore__ inline void +SASCubeBlock::InitPageAttentionInfo(GlobalTensor oriKvGm, const GlobalTensor &kvMergeGm, + GlobalTensor oriBlockTableGm, GlobalTensor cmpBlockTableGm) +{ + this->oriKvGm = oriKvGm; + this->kvMergeGm_ = kvMergeGm; + this->oriBlockTableGm = oriBlockTableGm; + this->cmpBlockTableGm = cmpBlockTableGm; +} + +template +__aicore__ inline void SASCubeBlock::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(bufQPL1, L1_BLOCK_SIZE * 4); + l1QPTensor = bufQPL1.Get(); + pipe->InitBuffer(bufKVL1, L1_BLOCK_SIZE * 3); + l1KVTensor = bufKVL1.Get(); + + // L0A + pipe->InitBuffer(tmpBufL0A, L0A_PP_SIZE * 2); // 64K + aL0TensorPingPong = tmpBufL0A.Get(); + // L0B + pipe->InitBuffer(tmpBufL0B, L0B_PP_SIZE * 2); // 64K + bL0TensorPingPong = tmpBufL0B.Get(); + // L0C + pipe->InitBuffer(tmpBufL0C, L0C_PP_SIZE * 2); // 128K + cL0TensorPingPong = tmpBufL0C.Get(); +} + +template +__aicore__ inline void SASCubeBlock::AllocEventID() +{ + SetFlag(L1_EVENT0); + SetFlag(L1_EVENT1); + SetFlag(L1_EVENT2); + SetFlag(L1_EVENT3); + SetFlag(L1_EVENT4); + SetFlag(L1_EVENT5); + SetFlag(L1_EVENT6); + SetFlag(L0AB_EVENT0); + SetFlag(L0AB_EVENT1); +} + +template +__aicore__ inline void SASCubeBlock::FreeEventID() +{ + WaitFlag(L1_EVENT0); + WaitFlag(L1_EVENT1); + WaitFlag(L1_EVENT2); + WaitFlag(L1_EVENT3); + WaitFlag(L1_EVENT4); + WaitFlag(L1_EVENT5); + WaitFlag(L1_EVENT6); + WaitFlag(L0AB_EVENT0); + WaitFlag(L0AB_EVENT1); +} + +template +__aicore__ inline void SASCubeBlock::CopyGmToL1(LocalTensor &l1Tensor, GlobalTensor &gmSrcTensor, + uint32_t srcN, uint32_t srcD, uint32_t srcDstride) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = srcN; // 行数 + nd2nzPara.dValue = srcD; + nd2nzPara.srcDValue = srcDstride; + nd2nzPara.dstNzC0Stride = (srcN + 15) / 16 * 16; // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(l1Tensor, gmSrcTensor, nd2nzPara); +} + +template +__aicore__ inline void SASCubeBlock::CopyInMm1AToL1(LocalTensor &l1Tensor, const RunInfo &info, + uint32_t mSeqIdx, uint32_t mSizeAct, uint32_t headSize, + uint32_t headOffset) +{ + auto srcGm = queryGm[info.tensorAOffset + mSeqIdx * constInfo.headDim + headOffset]; + CopyGmToL1(l1Tensor, srcGm, mSizeAct, headSize, constInfo.headDim); +} + +template +__aicore__ inline void SASCubeBlock::LoadDataMm1A(LocalTensor &aL0Tensor, LocalTensor &aL1Tensor, + uint32_t idx, uint32_t kSplitSize, uint32_t mSize, + uint32_t kSize) +{ + LocalTensor srcTensor = aL1Tensor[mSize * kSplitSize * idx]; + LoadData3DParamsV2 loadData3DParams; + // SetFmatrixParams + loadData3DParams.l1H = mSize / 16; // Hin=M1=8 + loadData3DParams.l1W = 16; // Win=M0 + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + // SetLoadToA0Params + loadData3DParams.mExtension = mSize; // M + loadData3DParams.kExtension = kSize; // K + loadData3DParams.mStartPt = 0; + loadData3DParams.kStartPt = 0; + loadData3DParams.strideW = 1; + loadData3DParams.strideH = 1; + loadData3DParams.filterW = 1; + loadData3DParams.filterSizeW = (1 >> 8) & 255; + loadData3DParams.filterH = 1; + loadData3DParams.filterSizeH = (1 >> 8) & 255; + loadData3DParams.dilationFilterW = 1; + loadData3DParams.dilationFilterH = 1; + loadData3DParams.enTranspose = 0; + loadData3DParams.fMatrixCtrl = 0; + loadData3DParams.channelSize = kSize; // Cin=K + LoadData(aL0Tensor, srcTensor, loadData3DParams); +} + +template +__aicore__ inline void SASCubeBlock::LoadDataMm1B(LocalTensor &l0Tensor, LocalTensor &l1Tensor, + uint32_t idx, uint32_t kSplitSize, uint32_t kSize, + uint32_t nSize) +{ + // N 方向全载 + LocalTensor srcTensor = l1Tensor[nSize * kSplitSize * idx]; + + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = (nSize + 15) / 16 * kSize / (32 / sizeof(KV_T)); + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = false; + LoadData(l0Tensor, srcTensor, loadData2DParams); +} + +template +__aicore__ inline void SASCubeBlock::CopyInMm2AToL1(LocalTensor &aL1Tensor, const RunInfo &info, + uint32_t mSeqIdx, uint32_t subMSizeAct, uint32_t nSize, + uint32_t nOffset) +{ + auto srcGm = vec1ResGm[(info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize + + mSeqIdx * info.actualSingleProcessSInnerSizeAlign + nOffset]; + CopyGmToL1(aL1Tensor, srcGm, subMSizeAct, nSize, info.actualSingleProcessSInnerSizeAlign); +} + +template +__aicore__ inline void SASCubeBlock::ComputeMm1(const RunInfo &info, const MSplitInfo mSplitInfo) +{ + uint32_t mSize = mSplitInfo.nBufferDealM; + uint32_t mL1Size = M_SPLIT_SIZE; + uint32_t mL1SizeAlign = SASAlign(M_SPLIT_SIZE, 16); + uint32_t mL1Loops = CeilDiv(mSize, M_SPLIT_SIZE); + + uint32_t nSize = info.actualSingleProcessSInnerSize; + uint32_t nL1Size = N_SPLIT_SIZE; + uint32_t nL1SizeAlign = SASAlign(N_SPLIT_SIZE, 16); + uint32_t nL1Loops = CeilDiv(nSize, N_SPLIT_SIZE); + + uint32_t kSize = 512; + uint32_t kL1Size = 256; + uint32_t kL1Loops = 2; + uint32_t kL0Size = 128; + uint32_t kL0Loops = CeilDiv(kL1Size, kL0Size); + + LocalTensor bL1Tensor; + uint32_t ka = 0, kb = 0; + + // L1 切n切k + for (uint32_t nL1 = 0; nL1 < nL1Loops; nL1++) { + if (nL1 == (nL1Loops - 1)) { + // 尾块重新计算size + nL1Size = nSize - (nL1Loops - 1) * N_SPLIT_SIZE; + nL1SizeAlign = SASAlign(nL1Size, 16); + } + + for (uint32_t kL1 = 0; kL1 < kL1Loops; kL1++) { + kvL1BufIter++; + uint32_t kb = kvL1BufIter % 3; + WaitFlag(mte21KVIds[kb]); + // 从k当中取当前的块 + bL1Tensor = l1KVTensor[kb * L1_BLOCK_OFFSET]; + uint32_t curSeqIdx = info.s2BatchOffset + nL1 * N_SPLIT_SIZE; + if (info.isOri) { + if constexpr (KV_LAYOUT_T == SAS_LAYOUT::PA_ND) { + uint32_t curS2Offset = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint; + uint32_t copyFinishRowCnt = 0; + LocalTensor kTensor; + uint32_t copyRowCnt = 0; + + while (copyFinishRowCnt < nL1Size) { + // 由于ori_left的存在， 即使第一块搬运也可能并非是pa_block的零点位 + copyRowCnt = constInfo.paOriBlockSize - curS2Offset % constInfo.paOriBlockSize; + if (copyFinishRowCnt + copyRowCnt > nL1Size) { + copyRowCnt = nL1Size - copyFinishRowCnt; + } + PAShape shape; + shape.blockSize = constInfo.paOriBlockSize; + shape.headNum = constInfo.kvHeadNum; + shape.headDim = constInfo.headDim; + shape.kvStride = constInfo.oriKvStride; + shape.actHeadDim = D_SPLIT_SIZE; + shape.maxblockNumPerBatch = constInfo.oriMaxBlockNumPerBatch; + shape.copyRowNum = copyRowCnt; + shape.copyRowNumAlign = nL1SizeAlign; + kTensor = bL1Tensor[copyFinishRowCnt * 16]; + + Position startPos; + startPos.bIdx = info.bIdx; + startPos.n2Idx = info.n2Idx; + startPos.s2Idx = curS2Offset; + startPos.dIdx = kL1 * D_SPLIT_SIZE; // mm1 右矩阵 bn2s2d, d为k轴不切; mm2 右矩阵, s2为k轴, d轴切分 + DataCopyPA(kTensor, oriKvGm, oriBlockTableGm, shape, startPos); + + // 更新循环变量 + copyFinishRowCnt += copyRowCnt; + curS2Offset += copyRowCnt; + } + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::BSND) { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nL1Size; // 行数 + nd2nzPara.dValue = D_SPLIT_SIZE; // 256 + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = nL1SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + + uint32_t headStride = constInfo.headDim; + uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim; + uint32_t batchStride = constInfo.kvSeqSize * seqStride; + + uint32_t curS2 = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint; + uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + (uint64_t)info.n2Idx * headStride + kL1 * D_SPLIT_SIZE; + DataCopy(bL1Tensor, oriKvGm[offset], nd2nzPara); + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::TND) { + uint32_t curS2Offset = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint; + if (kL1 == 0) { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nL1Size; + nd2nzPara.dValue = constInfo.headDim >> 1; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = nL1SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor, oriKvGm[info.tensorBOffset + curS2Offset * constInfo.headDim + + nL1 * N_SPLIT_SIZE * constInfo.headDim], nd2nzPara); + } else { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nL1Size; + nd2nzPara.dValue = constInfo.headDim >> 1; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = nL1SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor, + oriKvGm[info.tensorBOffset + curS2Offset * constInfo.headDim + (constInfo.headDim >> 1) + + nL1 * N_SPLIT_SIZE * constInfo.headDim], + nd2nzPara); + } + } + } else { + if (kL1 == 0) { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nL1Size; + nd2nzPara.dValue = constInfo.headDim >> 1; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = nL1SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor, + kvMergeGm_[info.cmpLoop % 4 * N_WORKSPACE_SIZE * kSize + + nL1 * N_SPLIT_SIZE * constInfo.headDim], + nd2nzPara); + } else { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nL1Size; + nd2nzPara.dValue = constInfo.headDim >> 1; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = nL1SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor, + kvMergeGm_[info.cmpLoop % 4 * N_WORKSPACE_SIZE * kSize + (constInfo.headDim >> 1) + + nL1 * N_SPLIT_SIZE * constInfo.headDim], + nd2nzPara); + } + } + SetFlag(mte21KVIds[kb]); + WaitFlag(mte21KVIds[kb]); + mL1Size = M_SPLIT_SIZE; + mL1SizeAlign = SASAlign(M_SPLIT_SIZE, 16U); + for (uint32_t mL1 = 0; mL1 < mL1Loops; mL1++) { + uint32_t aL1PaddingSize = 0; // 用于使左矩阵对齐到尾部, 以保证两块32K内存连续 + if (mL1 == (mL1Loops - 1)) { + mL1Size = mSize - (mL1Loops - 1) * M_SPLIT_SIZE; + mL1SizeAlign = SASAlign(mL1Size, 16U); + aL1PaddingSize = (M_SPLIT_SIZE - mL1SizeAlign) * 256; + } + uint32_t mIdx = qpL1BufIter + mL1; + ka = GetQPL1RealIdx(mIdx, kL1); + LocalTensor aL1Tensor = l1QPTensor[ka * L1_BLOCK_OFFSET + (1 - kL1) * aL1PaddingSize]; + if (nL1 == 0) { + if (kL1 == 0) { + WaitFlag(mte21QPIds[ka]); + WaitFlag(mte21QPIds[ka + 1]); + CopyInMm1AToL1(aL1Tensor, info, mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE, mL1Size, 256, 0); + } else { + LocalTensor qTmpTensor = aL1Tensor; + CopyInMm1AToL1(qTmpTensor, info, mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE, mL1Size, 256, + 256); + } + SetFlag(mte21QPIds[ka]); + WaitFlag(mte21QPIds[ka]); + } + // 使用unitflag同步 + LocalTensor cL0Tensor = + cL0TensorPingPong[(cL0BufIter % 2) * + (L0C_PP_SIZE / sizeof(MM_OUT_T))]; // 需要保证cL0BufIter和m步调一致 + for (uint32_t kL0 = 0; kL0 < kL0Loops; kL0++) { + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + LocalTensor aL0Tensor = aL0TensorPingPong[(abL0BufIter % 2) * (L0A_PP_SIZE / sizeof(KV_T))]; + LoadDataMm1A(aL0Tensor, aL1Tensor, kL0, kL0Size, mL1SizeAlign, kL0Size); + LocalTensor bL0Tensor = bL0TensorPingPong[(abL0BufIter % 2) * (L0B_PP_SIZE / sizeof(KV_T))]; + LoadDataMm1B(bL0Tensor, bL1Tensor, kL0, kL0Size, kL0Size, nL1SizeAlign); + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + + MmadParams mmadParams; + mmadParams.m = mL1SizeAlign; + mmadParams.n = nL1SizeAlign; + mmadParams.k = kL0Size; + mmadParams.cmatrixInitVal = (kL1 == 0 && kL0 == 0); + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = + (kL1 == 1 && kL0 == (kL0Loops - 1)) ? 0b11 : 0b10; // 累加最后一次翻转flag, 表示可以搬出 + Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + abL0BufIter++; + } + + if (nL1 == (nL1Loops - 1)) { + SetFlag(mte21QPIds[ka]); // 反向同步, 表示L1中的A已经被mte1消费完 + } + + if (kL1 == 1) { // 最后一轮kL1循环 + FixpipeParamsV220 fixParams; + fixParams.nSize = nL1SizeAlign; + fixParams.mSize = mL1SizeAlign; + fixParams.srcStride = mL1SizeAlign; + // 改成nSizeAlign + fixParams.dstStride = info.actualSingleProcessSInnerSizeAlign; // mm1ResGm两行之间的间隔 + fixParams.unitFlag = 0b11; + fixParams.ndNum = 1; // 输出ND + + Fixpipe(mm1ResGm[(info.loop % (constInfo.preLoadNum)) * constInfo.mmResUbSize + nL1 * N_SPLIT_SIZE + + (mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE) * + info.actualSingleProcessSInnerSizeAlign], + cL0Tensor, fixParams); + } + if (mL1Loops == 2) { + cL0BufIter++; + } + } + + SetFlag(mte21KVIds[kb]); // 反向同步, 表示L1已经被mte1消费完 + } + if (mL1Loops == 1) { + cL0BufIter++; + } + } + qpL1BufIter += mL1Loops; +} + + +template +__aicore__ inline void SASCubeBlock::ComputeMm2(const RunInfo &info, const MSplitInfo mSplitInfo) +{ + uint32_t mSize = mSplitInfo.nBufferDealM; + uint32_t mSizeAlign = (mSize + 16 - 1) / 16; + uint32_t mL1Loops = (mSize + M_SPLIT_SIZE - 1) / M_SPLIT_SIZE; + uint32_t mL1SizeAlign = M_SPLIT_SIZE; // 16对齐 + uint32_t mL1Size = M_SPLIT_SIZE; // m的实际大小 + + uint32_t nSize = BlockAlign(constInfo.headDim); + uint32_t nL1Loops = (nSize + N_SPLIT_SIZE - 1) / N_SPLIT_SIZE; + uint32_t nL1SizeAlign = N_SPLIT_SIZE; // 16对齐 + uint32_t nL1Size = N_SPLIT_SIZE; // n的实际大小 + + uint32_t kSize = info.actualSingleProcessSInnerSize; + uint32_t kL1Size = 256; + uint32_t kL1SizeAlign = SASAlign(kL1Size, 16U); + uint32_t kL1Loops = (kSize + kL1Size - 1) / kL1Size; + uint32_t kL0Size = 128; + uint32_t kL0Loops = (kL1Size + kL0Size - 1) / kL0Size; + uint32_t kL0SizeAlign = kL0Size; + LocalTensor bL1Tensor; + LocalTensor subvTensor; + + // ka表示左矩阵4buf选择哪一块buf, kb表示右矩阵3buf选择哪一块buf + uint32_t ka = 0, kb = 0; + uint32_t mBaseIdx = qpL1BufIter; + for (uint32_t nL1 = 0; nL1 < nL1Loops; nL1++) { // n切L1 -> D + if (nL1 == (nL1Loops - 1)) { + // 尾块 + nL1Size = nSize - (nL1Loops - 1) * N_SPLIT_SIZE; + nL1SizeAlign = SASAlign(nL1Size, 16U); + } + // k l1写成一个循环, 和mm1保持一致 + kL1Size = 256; + kL1SizeAlign = SASAlign(kL1Size, 16U); + uint32_t copyRowCnt = 0; + + for (uint32_t k1 = 0; k1 < kL1Loops; k1++) { // k切L1, 这里套了一层l0来操作 -> S2，每次256 + if (k1 == (kL1Loops - 1)) { + // 尾块 + kL1Size = kSize - (kL1Loops - 1) * 256; + kL1SizeAlign = SASAlign(kL1Size, 16U); + } + kvL1BufIter++; + uint32_t kb = kvL1BufIter % 3; + WaitFlag(mte21KVIds[kb]); + bL1Tensor = l1KVTensor[kb * L1_BLOCK_OFFSET]; + uint32_t kOffset = k1 * kL0Loops; + kL0Size = 128; + // 此处必须先初始化kL0Size, 再求kL0Loops, 否则由于循环会改变kL0Size大小, 导致kL0Loops错误 + kL0Loops = (kL1Size + kL0Size - 1) / kL0Size; + kL0SizeAlign = kL0Size; + for (uint32_t kL1 = kOffset; kL1 < kL0Loops + kOffset; kL1++) { // 128 循环搬pa，每次128 + if (kL1 == kOffset + kL0Loops - 1) { + // 尾块 + kL0Size = kL1Size - (kL0Loops - 1) * kL0Size; + kL0SizeAlign = SASAlign(kL0Size, 16U); + } + + uint32_t curSeqIdx = info.s2BatchOffset + (kL1 - kOffset) * 128 + k1 * 256; + if (info.isOri) { + if constexpr (KV_LAYOUT_T == SAS_LAYOUT::PA_ND) { + uint32_t copyFinishRowCnt = 0; + uint32_t curS2Offset = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint; + while (copyFinishRowCnt < kL0Size) { + copyRowCnt = constInfo.paOriBlockSize - curS2Offset % constInfo.paOriBlockSize; + if (copyFinishRowCnt + copyRowCnt > kL0Size) { + copyRowCnt = kL0Size - copyFinishRowCnt; + } + Position startPos; + startPos.bIdx = info.bIdx; + startPos.n2Idx = info.n2Idx; + startPos.s2Idx = curS2Offset; + startPos.dIdx = nL1 * N_SPLIT_SIZE; // mm1 右矩阵 bn2s2d, d为k轴不切; mm2 右矩阵, s2为k轴, d轴切分 + PAShape shape; + shape.blockSize = constInfo.paOriBlockSize; + shape.headNum = constInfo.kvHeadNum; + shape.headDim = constInfo.headDim; + shape.kvStride = constInfo.oriKvStride; + shape.actHeadDim = nL1Size; + shape.maxblockNumPerBatch = constInfo.oriMaxBlockNumPerBatch; + shape.copyRowNum = copyRowCnt; + shape.copyRowNumAlign = kL0SizeAlign; + subvTensor = bL1Tensor[(kL1 - kOffset) * 128 * N_SPLIT_SIZE + copyFinishRowCnt * 16]; + + DataCopyPA(subvTensor, oriKvGm, oriBlockTableGm, shape, startPos); + + // 更新循环变量 + copyFinishRowCnt += copyRowCnt; + curS2Offset += copyRowCnt; + } + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::BSND) { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = kL0Size; // 行数 + nd2nzPara.dValue = N_SPLIT_SIZE; // constInfo.headDim; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = kL0SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + + uint32_t headStride = constInfo.headDim; + uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim; + uint32_t batchStride = constInfo.kvSeqSize * seqStride; + + uint32_t curS2 = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint; + uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + (uint64_t)info.n2Idx * headStride + nL1 * N_SPLIT_SIZE; + subvTensor = bL1Tensor[(kL1 - kOffset) * 128 * N_SPLIT_SIZE]; + DataCopy(subvTensor, oriKvGm[offset], nd2nzPara); + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::TND) { + uint32_t curS2Offset = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = kL0Size; // 行数 + nd2nzPara.dValue = N_SPLIT_SIZE; // constInfo.headDim; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = kL0SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor[(kL1 - kOffset) * 128 * N_SPLIT_SIZE], + oriKvGm[info.tensorBOffset + curS2Offset * constInfo.headDim + kL1 * 128 * constInfo.headDim + + nL1 * N_SPLIT_SIZE], nd2nzPara); + } + } else { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = kL0Size; // 行数 + nd2nzPara.dValue = N_SPLIT_SIZE; // constInfo.headDim; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = kL0SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor[(kL1 - kOffset) * 128 * N_SPLIT_SIZE], + kvMergeGm_[info.cmpLoop % 4 * N_WORKSPACE_SIZE * 512 + kL1 * 128 * constInfo.headDim + + nL1 * N_SPLIT_SIZE], + nd2nzPara); + } + } + SetFlag(mte21KVIds[kb]); + WaitFlag(mte21KVIds[kb]); + mL1SizeAlign = M_SPLIT_SIZE; + mL1Size = M_SPLIT_SIZE; // m的实际大小 + for (uint32_t mL1 = 0; mL1 < mL1Loops; mL1++) { + if (mL1 == (mL1Loops - 1)) { + // 尾块 + mL1Size = mSize - (mL1Loops - 1) * M_SPLIT_SIZE; + mL1SizeAlign = SASAlign(mL1Size, 16U); + } + + uint32_t mIdx = mBaseIdx + mL1; + ka = GetQPL1RealIdx(mIdx, k1); + LocalTensor aL1Tensor = l1QPTensor[ka * L1_BLOCK_OFFSET]; + if (nL1 == 0) { + WaitFlag(mte21QPIds[ka]); + CopyInMm2AToL1(aL1Tensor, info, mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE, mL1Size, kL1Size, + 256 * k1); + SetFlag(mte21QPIds[ka]); + WaitFlag(mte21QPIds[ka]); + } + + LocalTensor cL0Tensor = + cL0TensorPingPong[(cL0BufIter % 2) * + (L0C_PP_SIZE / sizeof(MM_OUT_T))]; // 需要保证cL0BufIter和m步调一致 + uint32_t baseK = 128; + uint32_t baseN = 128; + kL0Size = 128; + kL0SizeAlign = kL0Size; + for (uint32_t kL0 = 0; kL0 < kL0Loops; kL0++) { + if (kL0 + 1 == kL0Loops) { + kL0Size = kL1Size - (kL0Loops - 1) * kL0Size; + kL0SizeAlign = SASAlign(kL0Size, 16U); + } + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + LocalTensor bL0Tensor = bL0TensorPingPong[(abL0BufIter % 2) * (L0B_PP_SIZE / sizeof(KV_T))]; + LoadData3DParamsV2 loadData3DParamsForB; + loadData3DParamsForB.l1H = kL0SizeAlign / 16; // 源操作数height + loadData3DParamsForB.l1W = 16; // 源操作数weight=16，目的height=l1H*L1W + loadData3DParamsForB.padList[0] = 0; + loadData3DParamsForB.padList[1] = 0; + loadData3DParamsForB.padList[2] = 0; + loadData3DParamsForB.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + loadData3DParamsForB.mExtension = kL0SizeAlign; // 在目的操作数height维度的传输长度 + loadData3DParamsForB.kExtension = nL1SizeAlign; // 在目的操作数width维度的传输长度 + loadData3DParamsForB.mStartPt = 0; // 卷积核在目的操作数width维度的起点 + loadData3DParamsForB.kStartPt = 0; // 卷积核在目的操作数height维度的起点 + loadData3DParamsForB.strideW = 1; + loadData3DParamsForB.strideH = 1; + loadData3DParamsForB.filterW = 1; + loadData3DParamsForB.filterSizeW = false; // 是否在filterW的基础上将卷积核width增加256个元素 + loadData3DParamsForB.filterH = 1; + loadData3DParamsForB.filterSizeH = false; // 是否在filterH的基础上将卷积核height增加256个元素 + loadData3DParamsForB.dilationFilterW = 1; // 卷积核width膨胀系数 + loadData3DParamsForB.dilationFilterH = 1; // 卷积核height膨胀系数 + loadData3DParamsForB.enTranspose = 1; // 是否启用转置功能 + loadData3DParamsForB.fMatrixCtrl = + 0; // 使用FMATRIX_LEFT还是使用FMATRIX_RIGHT，=0使用FMATRIX_LEFT，=1使用FMATRIX_RIGHT 1 + loadData3DParamsForB.channelSize = + nL1SizeAlign; // 源操作数的通道数。膨胀系数为1时，目的weight为filterW*filterH*channelSize + LoadData(bL0Tensor, bL1Tensor[kL0 * baseK * baseN], loadData3DParamsForB); + + LocalTensor aL0Tensor = aL0TensorPingPong[(abL0BufIter % 2) * (L0A_PP_SIZE / sizeof(KV_T))]; + LoadData3DParamsV2 loadData3DParamsForA; + loadData3DParamsForA.l1H = mL1SizeAlign / 16; // 源操作数height + loadData3DParamsForA.l1W = 16; // 源操作数weight + loadData3DParamsForA.padList[0] = 0; + loadData3DParamsForA.padList[1] = 0; + loadData3DParamsForA.padList[2] = 0; + loadData3DParamsForA.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + loadData3DParamsForA.mExtension = mL1SizeAlign; // 在目的操作数height维度的传输长度 + loadData3DParamsForA.kExtension = kL0SizeAlign; // 在目的操作数width维度的传输长度 + loadData3DParamsForA.mStartPt = 0; // 卷积核在目的操作数width维度的起点 + loadData3DParamsForA.kStartPt = 0; // 卷积核在目的操作数height维度的起点 + loadData3DParamsForA.strideW = 1; // 卷积核在源操作数width维度滑动的步长 + loadData3DParamsForA.strideH = 1; // 卷积核在源操作数height维度滑动的步长 + loadData3DParamsForA.filterW = 1; // 卷积核width + loadData3DParamsForA.filterSizeW = false; // 是否在filterW的基础上将卷积核width增加256个元素 + loadData3DParamsForA.filterH = 1; // 卷积核height + loadData3DParamsForA.filterSizeH = false; // 是否在filterH的基础上将卷积核height增加256个元素 + loadData3DParamsForA.dilationFilterW = 1; // 卷积核width膨胀系数 + loadData3DParamsForA.dilationFilterH = 1; // 卷积核height膨胀系数 + loadData3DParamsForA.enTranspose = 0; // 是否启用转置功能，对整个目标矩阵进行转置 + loadData3DParamsForA.fMatrixCtrl = 0; + loadData3DParamsForA.channelSize = + kL0SizeAlign; // 源操作数的通道数。膨胀系数为1时，目的weight为filterW*filterH*channelSize + LoadData(aL0Tensor, aL1Tensor[kL0 * baseK * mL1SizeAlign], + loadData3DParamsForA); + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + + MmadParams mmadParams; + mmadParams.m = mL1SizeAlign; + mmadParams.n = nL1SizeAlign; + mmadParams.k = kL0Size; + mmadParams.cmatrixInitVal = (kL0 == 0 && k1 == 0); + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = ((k1 == (kL1Loops - 1)) && (kL0 == (kL0Loops - 1))) ? 0b11 : 0b10; + + Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + abL0BufIter++; + } + + if (nL1 == (nL1Loops - 1)) { // nL1最后一轮, 需要将B驻留在L1中, 用于下一轮的计算？ + SetFlag(mte21QPIds[ka]); // 反向同步, 表示L1中的A已经被mte1消费完 + } + + if (k1 == (kL1Loops - 1)) { + // ND + FixpipeParamsV220 fixParams; + fixParams.nSize = nL1SizeAlign; + fixParams.mSize = mL1SizeAlign; + fixParams.srcStride = mL1SizeAlign; + fixParams.dstStride = nSize; // mm2ResGm两行之间的间隔 + fixParams.ndNum = 1; // 输出ND + fixParams.unitFlag = 0b11; + + uint64_t mm2Offset = (mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE) * nSize + nL1 * N_SPLIT_SIZE; + Fixpipe(mm2ResGm[(info.loop % (constInfo.preLoadNum)) * constInfo.bmm2ResUbSize + mm2Offset], + cL0Tensor, fixParams); + } + + if (mL1Loops == 2) { + cL0BufIter++; + } + } + SetFlag(mte21KVIds[kb]); // 反向同步, 表示L1已经被mte1消费完 + } + // cL0BufIter已经不在使用 + if (mL1Loops == 1) { + cL0BufIter++; + } + } + qpL1BufIter += mL1Loops; +} +} // namespace SASKernel +#endif // SPARSE_ATTN_SHAREDKV_SCFA_BLOCK_CUBE_H \ No newline at end of file diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_scfa_block_vector.h b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_scfa_block_vector.h new file mode 100644 index 000000000..bbce1450b --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_scfa_block_vector.h @@ -0,0 +1,1088 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_scfa_block_vector.h + * \brief + */ +#ifndef SPARSE_ATTN_SHAREDKV_SCFA_BLOCK_VECTOR_H +#define SPARSE_ATTN_SHAREDKV_SCFA_BLOCK_VECTOR_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../sparse_attn_sharedkv_common.h" + +namespace SASKernel { +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +template +class SASVectorBlock { +public: + // 中间计算数据类型为float，高精度模式 + using T = float; + using KV_T = typename SAST::kvType; + using OUT_T = typename SAST::outputType; + using UPDATE_T = T; + using SINKS_T = T; + using MM1_OUT_T = float; + using MM2_OUT_T = float; + + __aicore__ inline SASVectorBlock(){}; + __aicore__ inline void ProcessVec0L(const RunInfo &runInfo); + __aicore__ inline void ProcessVec1L(const RunInfo &info); + __aicore__ inline void ProcessVec2L(const RunInfo &info); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitParams(const struct ConstInfo &constInfo, + const SparseAttnSharedkvTilingData *__restrict tilingData); + __aicore__ inline void InitVec0GlobalTensor(const GlobalTensor &kvValidSizeGm, + const GlobalTensor &kvMergeGm, const GlobalTensor &oriKvGm, + const GlobalTensor &cmpKvGm, + const GlobalTensor &oriBlockTableGm, + const GlobalTensor &cmpBlockTableGm); + __aicore__ inline void InitVec1GlobalTensor(GlobalTensor mm1ResGm, GlobalTensor vec1ResGm, + GlobalTensor actualSeqLengthsQGm, + GlobalTensor actualSeqLengthsKVGm, + GlobalTensor topKGm, GlobalTensor sinksGm, GlobalTensor softmaxLseGm); + __aicore__ inline void InitVec2GlobalTensor(GlobalTensor accumOutGm, GlobalTensor vec2ResGm, + GlobalTensor mm2ResGm, GlobalTensor attentionOutGm); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void CopySinksIn(); + __aicore__ inline void SliceAndContactSinksValue(uint32_t nIdx, uint32_t dealRowCount); + __aicore__ inline void InitSoftmaxDefaultBuffer(); + // ================================Base Vector========================================== + __aicore__ inline void RowDivs(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline void RowMuls(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + // ================================Vector0========================================== + __aicore__ inline int64_t GetKeyGmOffset(int64_t realS2Idx, const RunInfo &runInfo, int64_t s2IdLimit); + __aicore__ inline void GetRealS2Idx(int64_t s2GmOffset, int64_t &realS2Idx, int64_t topkGmBaseOffset, + const RunInfo &runInfo); + __aicore__ inline void CopyInKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, int64_t realS2Idx1, + int64_t realS2Idx2, const RunInfo &runInfo); + __aicore__ inline void CopyOutMrgeResult(int64_t mte2Size, int64_t mte3Size, int64_t s2StartGmOffset, + int64_t mergeMte3Idx, const RunInfo &runInfo); + __aicore__ inline void CopyInSingleKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, int64_t realS2Idx, + int64_t keyBNBOffset, int64_t s2IdLimit, const RunInfo &runInfo); + // ================================Vector1========================================== + __aicore__ inline void ProcessVec1SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo); + __aicore__ inline void DealBmm1ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t startRow, + uint32_t dealRowCount, uint32_t columnCount, uint32_t loopId); + __aicore__ inline void SoftmaxFlashV2Compute(const RunInfo &info, const MSplitInfo &mSplitInfo, + LocalTensor &mmResUb, LocalTensor &softmaxTmpUb, + uint32_t startRow, uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + + __aicore__ inline void ElewiseCompute(const RunInfo &info, const LocalTensor &mmResUb, uint32_t dealRowCount, + uint32_t columnCount); + __aicore__ inline void ProcessLse(const RunInfo &info, const MSplitInfo &mSplitInfo); + // ================================Vecotr2========================================== + __aicore__ inline void ProcessVec2SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo); + __aicore__ inline void DealBmm2ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t startRow, + uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + __aicore__ inline void ProcessVec2Inner(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t mStartRow, + uint32_t mDealSize); + __aicore__ inline void Bmm2DataCopyOutTrans(const RunInfo &info, LocalTensor &attenOutUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + __aicore__ inline void Bmm2ResCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline void Bmm2CastAndCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline void Bmm2FDDataCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline uint64_t CalcAccumOffset(uint32_t bN2Idx, uint32_t gS1Idx); + + // BLOCK和REPEAT的字节数 + static constexpr uint64_t BYTE_BLOCK = 32UL; + static constexpr uint32_t REPEAT_BLOCK_BYTE = 256U; + // BLOCK和REPEAT的FP32元素数 + static constexpr uint32_t FP32_BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(float); + static constexpr uint32_t FP32_REPEAT_ELEMENT_NUM = REPEAT_BLOCK_BYTE / sizeof(float); + // repeat stride不能超过256 + static constexpr uint32_t REPEATE_STRIDE_UP_BOUND = 256; + +private: + static constexpr bool PAGE_ATTENTION = SAST::pageAttention; + static constexpr int TEMPLATE_MODE = SAST::templateMode; + static constexpr bool FLASH_DECODE = SAST::flashDecode; + static constexpr SAS_LAYOUT LAYOUT_T = SAST::layout; + static constexpr SAS_LAYOUT KV_LAYOUT_T = SAST::kvLayout; + + static constexpr uint64_t MERGE_CACHE_GM_BUF_NUM = 4; + static constexpr uint64_t SYNC_INPUT_BUF1_FLAG = 2; + static constexpr uint64_t SYNC_INPUT_BUF1_PONG_FLAG = 3; + static constexpr uint64_t SYNC_INPUT_BUF2_FLAG = 4; + static constexpr uint64_t SYNC_INPUT_BUF2_PONG_FLAG = 5; + static constexpr uint64_t SYNC_OUTPUT_BUF1_FLAG = 4; + static constexpr uint64_t SYNC_OUTPUT_BUF2_FLAG = 5; + static constexpr uint64_t SYNC_SINKS_BUF_FLAG = 6; + static constexpr uint64_t SYNC_INPUT_V0BUF_FLAG = 7; + static constexpr uint32_t INPUT1_BUFFER_OFFSET = ConstInfo::BUFFER_SIZE_BYTE_32K; + static constexpr uint32_t INPUT2_BUFFER_OFFSET = ConstInfo::BUFFER_SIZE_BYTE_16K; + static constexpr uint32_t SOFTMAX_TMP_BUFFER_OFFSET = ConstInfo::BUFFER_SIZE_BYTE_1K; + static constexpr uint32_t BASE_BLOCK_MAX_ELEMENT_NUM = ConstInfo::BUFFER_SIZE_BYTE_32K / sizeof(T); // 32768/4=8096 + static constexpr uint32_t BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(T); // 32/4=8 + static constexpr uint32_t MAX_N1_SIZE = 128U; + static constexpr T SOFTMAX_MIN_NUM = -2e38; + static constexpr SINKS_T R0 = 1.0f; + + const SparseAttnSharedkvTilingData *__restrict tilingData; + + uint32_t pingpongFlag = 0U; + ConstInfo constInfo = {}; + + GlobalTensor mm1ResGm; + GlobalTensor vec1ResGm; + GlobalTensor softmaxMaxGm; + GlobalTensor softmaxSumGm; + GlobalTensor sinksGm; + + GlobalTensor actualSeqLengthsQGm; + GlobalTensor actualSeqLengthsKVGm; + GlobalTensor vec2ResGm; + GlobalTensor mm2ResGm; + GlobalTensor accumOutGm; + GlobalTensor attentionOutGm; + GlobalTensor softmaxLseGm; + + GlobalTensor blkTableGm_; + GlobalTensor kvMergeGm_; + GlobalTensor keyGm_; + GlobalTensor topkGm_; + GlobalTensor kvValidSizeGm_; + GlobalTensor oriKvGm_; + GlobalTensor cmpKvGm_; + GlobalTensor oriBlockTableGm_; + GlobalTensor cmpBlockTableGm_; + + // ================================Local Buffer区==================================== + TBuf<> inputBuff1; // 32K + TBuf<> inputBuff2; // 16K + TBuf<> outputBuff1; // 32K + TBuf<> outputBuff2; // 32K + + TBuf<> tmpBuff1; // 32K + TBuf<> v0ValidSizeBuff; // 8K + + TBuf<> sinksBuff; // 1K + TBuf<> sinksBrcbBuff; // 12K + + TBuf<> softmaxMaxBuff; // PRE_LOAD_NUM * 2K + TBuf<> softmaxExpBuff; // PRE_LOAD_NUM * 2K + TBuf<> softmaxSumBuff; // PRE_LOAD_NUM * 2K + TBuf<> softmaxMaxDefaultBuff; // 2K + TBuf<> softmaxSumDefaultBuff; // 2K + + LocalTensor softmaxMaxDefaultUb; + LocalTensor softmaxSumDefaultUb; + + LocalTensor softmaxMaxUb; + LocalTensor softmaxSumUb; + LocalTensor softmaxExpUb; + LocalTensor kvMergUb_; + LocalTensor v0ValidSizeUb_; + LocalTensor sinksUb; + LocalTensor sinksBrcbUb; + + uint32_t mergeMte3Idx = 0; +}; + +template +__aicore__ inline void SASVectorBlock::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(inputBuff1, ConstInfo::BUFFER_SIZE_BYTE_32K * 2); // 2:pingpong + pipe->InitBuffer(inputBuff2, ConstInfo::BUFFER_SIZE_BYTE_16K * 2); // 2:pingpong + pipe->InitBuffer(outputBuff1, ConstInfo::BUFFER_SIZE_BYTE_32K); + if (constInfo.returnSoftmaxLse) { + pipe->InitBuffer(outputBuff2, ConstInfo::BUFFER_SIZE_BYTE_1K); + } + + pipe->InitBuffer(tmpBuff1, ConstInfo::BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(v0ValidSizeBuff, ConstInfo::BUFFER_SIZE_BYTE_8K); + + // M_MAX = 512/2vector = 256, 256 * sizeof(T) * N_Buffer + + pipe->InitBuffer(softmaxMaxBuff, ConstInfo::BUFFER_SIZE_BYTE_1K * constInfo.preLoadNum); + pipe->InitBuffer(softmaxExpBuff, ConstInfo::BUFFER_SIZE_BYTE_1K * constInfo.preLoadNum); + pipe->InitBuffer(softmaxSumBuff, ConstInfo::BUFFER_SIZE_BYTE_1K * constInfo.preLoadNum); + + pipe->InitBuffer(softmaxMaxDefaultBuff, ConstInfo::BUFFER_SIZE_BYTE_1K); + pipe->InitBuffer(softmaxSumDefaultBuff, ConstInfo::BUFFER_SIZE_BYTE_1K); + + pipe->InitBuffer(sinksBuff, MAX_N1_SIZE * sizeof(SINKS_T)); + // 分配256+N1大小内存，其中256是m轴VEC最大切块 + pipe->InitBuffer(sinksBrcbBuff, MAX_N1_SIZE * sizeof(SINKS_T) * BLOCK_ELEMENT_NUM * 3U); + + softmaxMaxUb = softmaxMaxBuff.Get(); + softmaxSumUb = softmaxSumBuff.Get(); + softmaxExpUb = softmaxExpBuff.Get(); + + softmaxMaxDefaultUb = softmaxMaxDefaultBuff.Get(); + softmaxSumDefaultUb = softmaxSumDefaultBuff.Get(); + + kvMergUb_ = inputBuff2.Get(); + + v0ValidSizeUb_ = v0ValidSizeBuff.Get(); + + sinksUb = sinksBuff.Get(); + sinksBrcbUb = sinksBrcbBuff.Get(); +} + +template +__aicore__ inline void SASVectorBlock::InitParams(const struct ConstInfo &constInfo, + const SparseAttnSharedkvTilingData *__restrict tilingData) +{ + this->constInfo = constInfo; + this->tilingData = tilingData; +} + +template +__aicore__ inline void SASVectorBlock::InitVec0GlobalTensor(const GlobalTensor &kvValidSizeGm, + const GlobalTensor &kvMergeGm, + const GlobalTensor &oriKvGm, + const GlobalTensor &cmpKvGm, + const GlobalTensor &oriBlockTableGm, + const GlobalTensor &cmpBlockTableGm) +{ + this->kvValidSizeGm_ = kvValidSizeGm; + this->kvMergeGm_ = kvMergeGm; + this->oriKvGm_ = oriKvGm; + this->cmpKvGm_ = cmpKvGm; + this->oriBlockTableGm_ = oriBlockTableGm; + this->cmpBlockTableGm_ = cmpBlockTableGm; +} + +template +__aicore__ inline void SASVectorBlock::InitVec1GlobalTensor( + GlobalTensor mm1ResGm, GlobalTensor vec1ResGm, + GlobalTensor actualSeqLengthsQGm, GlobalTensor actualSeqLengthsKVGm, + GlobalTensor topKGm, GlobalTensor sinksGm, GlobalTensor softmaxLseGm) +{ + this->mm1ResGm = mm1ResGm; + this->vec1ResGm = vec1ResGm; + this->actualSeqLengthsQGm = actualSeqLengthsQGm; + this->actualSeqLengthsKVGm = actualSeqLengthsKVGm; + this->topkGm_ = topKGm; + this->sinksGm = sinksGm; + this->softmaxLseGm = softmaxLseGm; +} + +template +__aicore__ inline void +SASVectorBlock::InitVec2GlobalTensor(GlobalTensor accumOutGm, GlobalTensor vec2ResGm, + GlobalTensor mm2ResGm, GlobalTensor attentionOutGm) +{ + this->accumOutGm = accumOutGm; + this->vec2ResGm = vec2ResGm; + this->mm2ResGm = mm2ResGm; + this->attentionOutGm = attentionOutGm; +} + +template +__aicore__ inline void SASVectorBlock::AllocEventID() +{ + SetFlag(SYNC_INPUT_BUF1_FLAG); + SetFlag(SYNC_INPUT_BUF1_PONG_FLAG); + SetFlag(SYNC_INPUT_BUF2_FLAG); + SetFlag(SYNC_INPUT_BUF2_PONG_FLAG); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); +} + +template +__aicore__ inline void SASVectorBlock::FreeEventID() +{ + WaitFlag(SYNC_INPUT_BUF1_FLAG); + WaitFlag(SYNC_INPUT_BUF1_PONG_FLAG); + WaitFlag(SYNC_INPUT_BUF2_FLAG); + WaitFlag(SYNC_INPUT_BUF2_PONG_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); +} + +template +__aicore__ inline void SASVectorBlock::CopySinksIn() +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = 1U; + dataCopyParams.blockLen = constInfo.qHeadNum * sizeof(T); + dataCopyParams.srcStride = 0U; + dataCopyParams.dstStride = 0U; + DataCopyPadExtParams padParams; + DataCopyPad(sinksUb, sinksGm, dataCopyParams, padParams); + SetFlag(SYNC_SINKS_BUF_FLAG); + WaitFlag(SYNC_SINKS_BUF_FLAG); + uint32_t repeatTimes = (constInfo.qHeadNum + BLOCK_ELEMENT_NUM - 1U) / BLOCK_ELEMENT_NUM; // 每次处理 8 datablocks + Brcb(sinksBrcbUb, sinksUb, repeatTimes, {1, BLOCK_ELEMENT_NUM}); + PipeBarrier(); + + DataCopyParams repeatParams; + repeatParams.blockCount = 1; // 搬到有一个块超过单个vec核减分核M轴大小即可，核间切分每个vec256 + repeatParams.blockLen = constInfo.qHeadNum; + repeatParams.srcStride = 0U; + repeatParams.dstStride = 0U; + for (uint32_t i = 1U; i <= 256U / constInfo.qHeadNum; i++) { + DataCopy(sinksBrcbUb[constInfo.qHeadNum * BLOCK_ELEMENT_NUM * i], sinksBrcbUb, repeatParams); + } + PipeBarrier(); +} + +template +__aicore__ inline void SASVectorBlock::SliceAndContactSinksValue(uint32_t nIdx, uint32_t dealRowCount) +{ + // WholeReduceMax接口中repeatTimes支持范围（0,255），因此需要分多次调用WholeReduceMax，每次repeatTime=128 + uint32_t repeatTimesOnce = 128; + uint32_t loopTimes = (dealRowCount + repeatTimesOnce - 1) / repeatTimesOnce; + uint32_t repeatTimes = repeatTimesOnce; + + for (uint32_t loop = 0; loop < loopTimes; ++loop) { + if (loop == loopTimes - 1) { + repeatTimes = dealRowCount - loop * repeatTimesOnce; + } + WholeReduceMax(softmaxMaxDefaultUb[loop * repeatTimesOnce], + sinksBrcbUb[(nIdx + loop * repeatTimesOnce) * BLOCK_ELEMENT_NUM], + BLOCK_ELEMENT_NUM * BLOCK_ELEMENT_NUM, repeatTimes, 1, 0, 1, ReduceOrder::ORDER_ONLY_VALUE); + PipeBarrier(); + } +} + +template +__aicore__ inline void SASVectorBlock::InitSoftmaxDefaultBuffer() +{ + CopySinksIn(); + Duplicate(softmaxMaxDefaultUb, SOFTMAX_MIN_NUM, SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)); + Duplicate(softmaxSumDefaultUb, R0, SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)); +} + +template +__aicore__ inline void SASVectorBlock::ElewiseCompute(const RunInfo &info, const LocalTensor &mmResUb, + uint32_t dealRowCount, uint32_t columnCount) +{ + Muls(mmResUb, mmResUb, static_cast(tilingData->baseParams.softmaxScale), dealRowCount * columnCount); +} + +template +__aicore__ inline void SASVectorBlock::ProcessLse(const RunInfo &info, const MSplitInfo &mSplitInfo) +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + uint64_t lseOffset; + if (constInfo.outputLayout == SAS_LAYOUT::TND) { + uint32_t tBase = actualSeqLengthsQGm.GetValue(info.bIdx); + lseOffset = (tBase + info.s1Idx) * constInfo.gSize + // T轴、s1轴偏移 + info.n2IdxReal * constInfo.qSeqSize * constInfo.gSize; // N2轴偏移 + } else if (constInfo.outputLayout == SAS_LAYOUT::BSND) { + lseOffset = info.bIdx * constInfo.qSeqSize * constInfo.kvHeadNum * constInfo.gSize + // B轴偏移 + info.n2IdxReal * constInfo.qSeqSize * constInfo.gSize + // N2轴偏移 + info.s1Idx * constInfo.gSize; // S1轴偏移 + } + lseOffset = lseOffset + mSplitInfo.nBufferStartM + mSplitInfo.vecStartM; + uint32_t baseOffset = mSplitInfo.nBufferStartM / 2; + uint32_t outIdx = info.loop % (constInfo.preLoadNum); + uint32_t softmaxOffset = outIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset; + auto sumTensor = softmaxSumUb[softmaxOffset]; + auto maxTensor = softmaxMaxUb[softmaxOffset]; + auto outLSETensor = outputBuff2.Get(); + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = 1; + dataCopyParams.blockLen = mSplitInfo.vecDealM * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = 0; + + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + PipeBarrier(); + Log(outLSETensor, sumTensor, mSplitInfo.vecDealM); + PipeBarrier(); + Add(outLSETensor, outLSETensor, maxTensor, mSplitInfo.vecDealM); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + + DataCopyPad(softmaxLseGm[lseOffset], outLSETensor, dataCopyParams); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); +} + +template +__aicore__ inline void +SASVectorBlock::SoftmaxFlashV2Compute(const RunInfo &info, const MSplitInfo &mSplitInfo, LocalTensor &mmResUb, + LocalTensor &softmaxTmpUb, uint32_t startRow, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount) +{ + LocalTensor inSumTensor; + LocalTensor inMaxTensor; + uint32_t baseOffset = mSplitInfo.nBufferStartM / 2 + startRow; + uint32_t outIdx = info.loop % (constInfo.preLoadNum); + uint32_t softmaxOutOffset = outIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset; + if (info.isFirstSInnerLoop) { + inMaxTensor = softmaxMaxDefaultUb[startRow]; + inSumTensor = softmaxSumDefaultUb; + } else { + uint32_t inIdx = (info.loop - 1) % (constInfo.preLoadNum); + inMaxTensor = softmaxMaxUb[inIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset]; + inSumTensor = softmaxSumUb[inIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset]; + } + if (actualColumnCount != 0) { + SoftMaxShapeInfo srcShape{dealRowCount, columnCount, dealRowCount, actualColumnCount}; + SoftMaxTiling newTiling = + SoftMaxFlashV2TilingFunc(srcShape, sizeof(T), sizeof(T), softmaxTmpUb.GetSize(), true, false); + SoftmaxFlashV2( + mmResUb, softmaxSumUb[softmaxOutOffset], softmaxMaxUb[softmaxOutOffset], mmResUb, + softmaxExpUb[softmaxOutOffset], inSumTensor, inMaxTensor, softmaxTmpUb, newTiling, srcShape); + } else { + uint32_t dealRowCountAlign = SASAlign(dealRowCount, FP32_BLOCK_ELEMENT_NUM); + DataCopy(softmaxSumUb[softmaxOutOffset], inSumTensor, dealRowCountAlign); + PipeBarrier(); + DataCopy(softmaxMaxUb[softmaxOutOffset], inMaxTensor, dealRowCountAlign); + } +} + +template +__aicore__ inline void SASVectorBlock::DealBmm1ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, + uint32_t startRow, uint32_t dealRowCount, + uint32_t columnCount, uint32_t loopId) +{ + uint32_t computeSize = dealRowCount * columnCount; + uint64_t inOutGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize + + (mSplitInfo.nBufferStartM + mSplitInfo.vecStartM + startRow) * columnCount; + LocalTensor mmResUb = inputBuff1.Get(); + mmResUb = mmResUb[pingpongFlag * INPUT1_BUFFER_OFFSET / sizeof(MM1_OUT_T)]; + WaitFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + + DataCopy(mmResUb, mm1ResGm[inOutGmOffset], computeSize); + SetFlag(SYNC_INPUT_BUF1_FLAG); + WaitFlag(SYNC_INPUT_BUF1_FLAG); + + ElewiseCompute(info, mmResUb, dealRowCount, columnCount); + + PipeBarrier(); + LocalTensor tmpAFloorUb = tmpBuff1.Get(); + LocalTensor softmaxTmpUb = tmpAFloorUb.template ReinterpretCast(); + + SoftmaxFlashV2Compute(info, mSplitInfo, mmResUb, softmaxTmpUb, startRow, dealRowCount, columnCount, + info.actualSingleProcessSInnerSize); + + PipeBarrier(); + LocalTensor tmpMMResCastTensor = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + + Cast(tmpMMResCastTensor, mmResUb, AscendC::RoundMode::CAST_ROUND, computeSize); + SetFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + pingpongFlag ^= 1; // pingpong 0 1 切换 + + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + DataCopy(vec1ResGm[inOutGmOffset], tmpMMResCastTensor, computeSize); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); +} + +template +__aicore__ inline void SASVectorBlock::ProcessVec1SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo) +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + uint32_t mSplitSize = info.actualSingleProcessSInnerSize == 0 ? + 16 : + BASE_BLOCK_MAX_ELEMENT_NUM / info.actualSingleProcessSInnerSizeAlign; + // 1. 向下8对齐是因为UB操作至少32B + // 2. info.actualSingleProcessSInnerSizeAlign最大512, mSplitSize可以确保最小为16 + mSplitSize = mSplitSize / 8 * 8; + + if (mSplitSize > mSplitInfo.vecDealM) { + mSplitSize = mSplitInfo.vecDealM; + } + uint32_t loopCount = (mSplitInfo.vecDealM + mSplitSize - 1) / mSplitSize; + uint32_t tailSplitSize = mSplitInfo.vecDealM - (loopCount - 1) * mSplitSize; + + SliceAndContactSinksValue((mSplitInfo.nBufferStartM + mSplitInfo.vecStartM) % constInfo.qHeadNum, + mSplitInfo.vecDealM); + + for (uint32_t i = 0, dealSize = mSplitSize; i < loopCount; i++) { + if (i == (loopCount - 1)) { + dealSize = tailSplitSize; + } + DealBmm1ResBaseBlock(info, mSplitInfo, i * mSplitSize, dealSize, info.actualSingleProcessSInnerSizeAlign, i); + } +} + +template +__aicore__ inline void SASVectorBlock::GetRealS2Idx(int64_t s2GmOffset, int64_t &realS2Idx, + int64_t topkGmBaseOffset, const RunInfo &runInfo) +{ + int64_t cmpS2Offset = s2GmOffset; + int64_t topkGmIdx = cmpS2Offset / constInfo.sparseBlockSize; + if (unlikely(topkGmIdx >= constInfo.sparseBlockCount || s2GmOffset >= runInfo.v0S2DealSize)) { + realS2Idx = -1; + return; + } + realS2Idx = topkGm_.GetValue(topkGmBaseOffset + topkGmIdx) * static_cast(constInfo.sparseBlockSize) + + static_cast(cmpS2Offset % constInfo.sparseBlockSize); +} + +template +__aicore__ inline int64_t SASVectorBlock::GetKeyGmOffset(int64_t realS2Idx, const RunInfo &runInfo, + int64_t s2IdLimit) +{ + if (realS2Idx < 0 || realS2Idx >= s2IdLimit) { + return -1; + } + int64_t realKeyGmOffset = 0; + if constexpr (KV_LAYOUT_T == SAS_LAYOUT::PA_ND) { + int64_t blkTableIdx = realS2Idx / constInfo.paCmpBlockSize; + int64_t blkTableOffset = realS2Idx % constInfo.paCmpBlockSize; + realKeyGmOffset = cmpBlockTableGm_.GetValue(runInfo.bIdx * constInfo.cmpMaxBlockNumPerBatch + blkTableIdx) * + static_cast(constInfo.cmpKvStride) + + blkTableOffset * static_cast(constInfo.kvHeadNum) * static_cast(constInfo.headDim); + + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::BSND) { + realKeyGmOffset = runInfo.bIdx * constInfo.kvSeqSize / constInfo.cmpRatio * constInfo.kvHeadNum + realS2Idx * constInfo.kvHeadNum; + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::TND) { + realKeyGmOffset = (runInfo.tensorCmpBOffset + realS2Idx * constInfo.kvHeadNum * constInfo.headDim) / + constInfo.headDim; + } + return realKeyGmOffset; +} + +template +__aicore__ inline void SASVectorBlock::CopyInSingleKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, + int64_t realS2Idx, int64_t keyBNBOffset, int64_t s2IdLimit, + const RunInfo &runInfo) +{ + if (keyBNBOffset < 0) { + return; + } + int64_t validS2Count = + (realS2Idx + constInfo.sparseBlockSize > s2IdLimit ? s2IdLimit - realS2Idx : constInfo.sparseBlockSize); + DataCopyExtParams intriParams; + intriParams.blockLen = validS2Count * constInfo.headDim * sizeof(KV_T); + intriParams.blockCount = 1; + intriParams.dstStride = 0; + intriParams.srcStride = 0; + DataCopyPadExtParams padParams; + if constexpr (KV_LAYOUT_T == SAS_LAYOUT::PA_ND) { + DataCopyPad( + kvMergUb_[mergeMte3Idx % 2 * INPUT2_BUFFER_OFFSET / sizeof(KV_T) + (mte2Size - mte3Size) * constInfo.headDim], + cmpKvGm_[keyBNBOffset], intriParams, padParams); + } else { + DataCopyPad( + kvMergUb_[mergeMte3Idx % 2 * INPUT2_BUFFER_OFFSET / sizeof(KV_T) + (mte2Size - mte3Size) * constInfo.headDim], + cmpKvGm_[keyBNBOffset * constInfo.headDim], intriParams, padParams); + } + mte2Size += validS2Count; +} + +template +__aicore__ inline void SASVectorBlock::CopyInKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, + int64_t realS2Idx1, int64_t realS2Idx2, const RunInfo &runInfo) +{ + int64_t s2IdLimit = runInfo.cmpS2IdLimit; + + int64_t keyOffset1 = GetKeyGmOffset(realS2Idx1, runInfo, s2IdLimit); + int64_t keyOffset2 = GetKeyGmOffset(realS2Idx2, runInfo, s2IdLimit); + if (unlikely(keyOffset1 < 0 && keyOffset2 < 0)) { + return; + } + + int64_t keySrcStride = 0; + if constexpr (KV_LAYOUT_T == SAS_LAYOUT::PA_ND) { + int64_t blkTableSrcStride = + ((keyOffset1 > keyOffset2 ? (keyOffset1 - keyOffset2) : + (keyOffset2 - keyOffset1)) - constInfo.sparseBlockSize * constInfo.headDim); + keySrcStride = blkTableSrcStride * sizeof(KV_T); + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::BSND) { + keySrcStride = ((keyOffset1 > keyOffset2 ? (keyOffset1 - keyOffset2) : + (keyOffset2 - keyOffset1)) - constInfo.sparseBlockSize) * constInfo.headDim * sizeof(KV_T); + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::TND) { + keySrcStride = ((keyOffset1 > keyOffset2 ? (keyOffset1 - keyOffset2) : + (keyOffset2 - keyOffset1)) - constInfo.sparseBlockSize) * constInfo.headDim * sizeof(KV_T); + } + if (unlikely(keySrcStride >= INT32_MAX || keySrcStride < 0 || + realS2Idx1 + constInfo.sparseBlockSize >= s2IdLimit || + realS2Idx2 + constInfo.sparseBlockSize >= s2IdLimit)) { + // stride溢出、stride为负数、s2超长等异常场景，还原成2条搬运指令 + // 因为需要拷贝两块 + CopyInSingleKv(mte2Size, mte3Size, mergeMte3Idx, realS2Idx1, keyOffset1, s2IdLimit, runInfo); + CopyInSingleKv(mte2Size, mte3Size, mergeMte3Idx, realS2Idx2, keyOffset2, s2IdLimit, runInfo); + } else { + DataCopyExtParams intriParams; + intriParams.blockLen = constInfo.sparseBlockSize * constInfo.headDim * sizeof(KV_T); + intriParams.blockCount = (keyOffset1 >= 0) + (keyOffset2 >= 0); + intriParams.dstStride = 0; + intriParams.srcStride = keySrcStride; + DataCopyPadExtParams padParams; + + int64_t startGmOffset = keyOffset1 > -1 ? keyOffset1 : keyOffset2; + if (keyOffset2 > -1 && keyOffset2 < keyOffset1) { + startGmOffset = keyOffset2; + } + if constexpr (KV_LAYOUT_T == SAS_LAYOUT::PA_ND) { + DataCopyPad(kvMergUb_[mergeMte3Idx % 2 * INPUT2_BUFFER_OFFSET / sizeof(KV_T) + + (mte2Size - mte3Size) * constInfo.headDim], + cmpKvGm_[startGmOffset], intriParams, padParams); + } else { + DataCopyPad(kvMergUb_[mergeMte3Idx % 2 * INPUT2_BUFFER_OFFSET / sizeof(KV_T) + + (mte2Size - mte3Size) * constInfo.headDim], + cmpKvGm_[startGmOffset * constInfo.headDim], intriParams, padParams); + } + mte2Size += ((keyOffset1 > -1) + (keyOffset2 > -1)) * constInfo.sparseBlockSize; + } +} + +template +__aicore__ inline void SASVectorBlock::CopyOutMrgeResult(int64_t mte2Size, int64_t mte3Size, + int64_t s2GmStartOffset, int64_t mergeMte3Idx, + const RunInfo &runInfo) +{ + if (mte2Size <= mte3Size) { + return; + } + SetFlag(mergeMte3Idx % 2 + SYNC_INPUT_BUF2_FLAG); + WaitFlag(mergeMte3Idx % 2 + SYNC_INPUT_BUF2_FLAG); + + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = mte2Size - mte3Size; + dataCopyParams.blockLen = constInfo.headDim * sizeof(KV_T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = 0; + + DataCopyPad(kvMergeGm_[runInfo.cmpLoop % 4 * 512 * 512 + + (s2GmStartOffset + mte3Size) * constInfo.headDim], + kvMergUb_[mergeMte3Idx % 2 * INPUT2_BUFFER_OFFSET / sizeof(KV_T)], dataCopyParams); +} + +// b s1 k +template +__aicore__ inline void SASVectorBlock::ProcessVec0L(const RunInfo &runInfo) +{ + int64_t s2ProcessSize = runInfo.v0S2DealSize; + int64_t s2Pair = CeilDiv(s2ProcessSize, 2 * constInfo.sparseBlockSize); + int64_t topkGmBaseOffset = runInfo.topKBaseOffset + runInfo.v0S2Start; + int64_t mte2Size = 0; + int64_t mte3Size = 0; + int64_t s2IdxArray0 = -1; + int64_t s2IdxArray1 = -1; + bool needWaitMte3ToMte2 = true; + int64_t s2SplitPoint = SASAlign(s2Pair, 2) * constInfo.sparseBlockSize; + int64_t s2GmStartOffset = GetSubBlockIdx() == 0 ? 0 : s2SplitPoint; + int64_t s2GmLimit = GetSubBlockIdx() == 0 ? s2SplitPoint : s2ProcessSize; + if (s2GmLimit > s2ProcessSize) { + s2GmLimit = s2ProcessSize; + } + // 处理两个基本块 + for (int64_t s2GmOffsetArray = s2GmStartOffset; s2GmOffsetArray < s2GmLimit; + s2GmOffsetArray += 2 * constInfo.sparseBlockSize) { + if (needWaitMte3ToMte2) { + WaitFlag(mergeMte3Idx % 2 + SYNC_INPUT_BUF2_FLAG); + needWaitMte3ToMte2 = false; + } + GetRealS2Idx(s2GmOffsetArray, s2IdxArray0, topkGmBaseOffset, runInfo); + if (unlikely(s2IdxArray0 < 0)) { + CopyOutMrgeResult(mte2Size, mte3Size, s2GmStartOffset, mergeMte3Idx, runInfo); + SetFlag(mergeMte3Idx % 2 + SYNC_INPUT_BUF2_FLAG); + mergeMte3Idx++; + break; + } + GetRealS2Idx(s2GmOffsetArray + constInfo.sparseBlockSize, s2IdxArray1, topkGmBaseOffset, runInfo); + CopyInKv(mte2Size, mte3Size, mergeMte3Idx, s2IdxArray0, s2IdxArray1, runInfo); + if ((mte2Size - mte3Size + 2 * constInfo.sparseBlockSize > 16) || + s2GmOffsetArray + 2 * constInfo.sparseBlockSize >= s2GmLimit) { + CopyOutMrgeResult(mte2Size, mte3Size, s2GmStartOffset, mergeMte3Idx, runInfo); + mte3Size = mte2Size; + SetFlag(mergeMte3Idx % 2 + SYNC_INPUT_BUF2_FLAG); + mergeMte3Idx++; + needWaitMte3ToMte2 = true; + } + } + return; +} + +template +__aicore__ inline void SASVectorBlock::ProcessVec1L(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferIdx = i; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + + mSplitInfo.vecDealM = (mSplitInfo.nBufferDealM <= 16) ? mSplitInfo.nBufferDealM : + (((mSplitInfo.nBufferDealM + 15) / 16 + 1) / 2 * 16); + mSplitInfo.vecStartM = 0; + if (GetBlockIdx() % 2 == 1) { + mSplitInfo.vecStartM = mSplitInfo.vecDealM; + mSplitInfo.vecDealM = mSplitInfo.nBufferDealM - mSplitInfo.vecDealM; + } + + CrossCoreWaitFlag(constInfo.syncC1V1); + // vec1 compute + ProcessVec1SingleBuf(info, mSplitInfo); + CrossCoreSetFlag(constInfo.syncV1C2); + + // move lse for flash decode or FA + if (constInfo.returnSoftmaxLse && info.s2Idx == info.curSInnerLoopTimes - 1) { + ProcessLse(info, mSplitInfo); + } + } +} + +template +__aicore__ inline uint64_t SASVectorBlock::CalcAccumOffset(uint32_t bN2Idx, uint32_t gS1Idx) +{ + return 0; +} + +template +__aicore__ inline void SASVectorBlock::ProcessVec2SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo) +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + + ProcessVec2Inner(info, mSplitInfo, 0, mSplitInfo.vecDealM); +} + +template +__aicore__ inline void SASVectorBlock::ProcessVec2L(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferIdx = i; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + + mSplitInfo.vecDealM = (mSplitInfo.nBufferDealM <= 16) ? mSplitInfo.nBufferDealM : + (((mSplitInfo.nBufferDealM + 15) / 16 + 1) / 2 * 16); + mSplitInfo.vecStartM = 0; + if (GetBlockIdx() % 2 == 1) { + mSplitInfo.vecStartM = mSplitInfo.vecDealM; + mSplitInfo.vecDealM = mSplitInfo.nBufferDealM - mSplitInfo.vecDealM; + } + CrossCoreWaitFlag(constInfo.syncC2V2); + ProcessVec2SingleBuf(info, mSplitInfo); + } +} + +template +__aicore__ inline void SASVectorBlock::ProcessVec2Inner(const RunInfo &info, const MSplitInfo &mSplitInfo, + uint32_t mStartRow, uint32_t mDealSize) +{ + uint32_t mSplitSize = BASE_BLOCK_MAX_ELEMENT_NUM / constInfo.headDim; + if (mSplitSize > mDealSize) { + mSplitSize = mDealSize; + } + + uint32_t loopCount = (mDealSize + mSplitSize - 1) / mSplitSize; + uint32_t tailSplitSize = mDealSize - (loopCount - 1) * mSplitSize; + for (uint32_t i = 0, dealSize = mSplitSize; i < loopCount; i++) { + if (i == (loopCount - 1)) { + dealSize = tailSplitSize; + } + DealBmm2ResBaseBlock(info, mSplitInfo, i * mSplitSize + mStartRow, dealSize, constInfo.headDim, + constInfo.headDim); + } +} + +template +__aicore__ inline void SASVectorBlock::Bmm2FDDataCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, + uint32_t wsMStart, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + LocalTensor tmp = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + DataCopy(tmp, bmm2ResUb, columnCount * dealRowCount); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + uint64_t accumTmpOutNum = CalcAccumOffset(info.bIdx, info.gS1Idx); + uint64_t offset = + accumTmpOutNum * constInfo.kvHeadNum * constInfo.mBaseSize * constInfo.headDim + // taskoffset + info.tndCoreStartKVSplitPos * constInfo.kvHeadNum * constInfo.mBaseSize * constInfo.headDim + // 份数offset + wsMStart * actualColumnCount; // m轴offset + GlobalTensor dst = accumOutGm[offset]; + if (info.actualSingleProcessSInnerSize == 0) { + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = dealRowCount; + dataCopyParams.blockLen = actualColumnCount * sizeof(T); + dataCopyParams.srcStride = (columnCount - actualColumnCount) / (BYTE_BLOCK / sizeof(T)); + dataCopyParams.dstStride = 0; + DataCopyPad(dst, tmp, dataCopyParams); + } else { + matmul::InitOutput(dst, dealRowCount * actualColumnCount, ConstInfo::FLOAT_ZERO); + } + SetFlag(SYNC_OUTPUT_BUF1_FLAG); +} + +template +__aicore__ inline void SASVectorBlock::Bmm2DataCopyOutTrans(const RunInfo &info, LocalTensor &attenOutUb, + uint32_t wsMStart, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = dealRowCount; + dataCopyParams.blockLen = actualColumnCount * sizeof(OUT_T); + dataCopyParams.srcStride = (columnCount - actualColumnCount) / (BYTE_BLOCK / sizeof(OUT_T)); + dataCopyParams.dstStride = 0; + DataCopyPad(attentionOutGm[info.attenOutOffset + wsMStart * actualColumnCount], attenOutUb, dataCopyParams); + return; +} + +template +__aicore__ inline void SASVectorBlock::Bmm2CastAndCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, + uint32_t wsMStart, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + LocalTensor tmpBmm2ResCastTensor = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + if constexpr (IsSameType::value) { // bf16 采取四舍六入五成双模式 + Cast(tmpBmm2ResCastTensor, bmm2ResUb, AscendC::RoundMode::CAST_RINT, dealRowCount * columnCount); + } else { + Cast(tmpBmm2ResCastTensor, bmm2ResUb, AscendC::RoundMode::CAST_ROUND, dealRowCount * columnCount); + } + + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + Bmm2DataCopyOutTrans(info, tmpBmm2ResCastTensor, wsMStart, dealRowCount, columnCount, actualColumnCount); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); +} + +template +__aicore__ inline void SASVectorBlock::Bmm2ResCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, + uint32_t wsMStart, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + if constexpr (FLASH_DECODE) { + if (info.tndIsS2SplitCore) { + Bmm2FDDataCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount); + } else { + Bmm2CastAndCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount); + } + } else { + Bmm2CastAndCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount); + } +} + +template +__aicore__ inline void SASVectorBlock::DealBmm2ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, + uint32_t startRow, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + uint32_t vec2ComputeSize = dealRowCount * columnCount; + uint32_t mStart = mSplitInfo.nBufferStartM + mSplitInfo.vecStartM + startRow; + uint64_t srcGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + mStart * columnCount; + LocalTensor tmpBmm2ResUb = inputBuff1.Get(); + tmpBmm2ResUb = tmpBmm2ResUb[pingpongFlag * INPUT1_BUFFER_OFFSET / sizeof(MM2_OUT_T)]; + WaitFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + DataCopy(tmpBmm2ResUb, mm2ResGm[srcGmOffset], vec2ComputeSize); + + SetFlag(SYNC_INPUT_BUF1_FLAG); + WaitFlag(SYNC_INPUT_BUF1_FLAG); + + LocalTensor bmm2ResUb = tmpBuff1.Get(); + bmm2ResUb.SetSize(vec2ComputeSize); + DataCopy(bmm2ResUb, tmpBmm2ResUb, vec2ComputeSize); + SetFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + pingpongFlag ^= 1; // pingpong 0 1切换 + + uint32_t inOutBaseOffset = mStart * columnCount; + uint32_t baseOffset = mSplitInfo.nBufferStartM / 2 + startRow; + + // 除第一个循环外，均需要更新中间计算结果 + if (!info.isFirstSInnerLoop) { + event_t eventIdMte2WaitMte3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventIdMte2WaitMte3); + WaitFlag(eventIdMte2WaitMte3); + + LocalTensor bmm2ResPreUb = inputBuff1.Get(); + bmm2ResPreUb = bmm2ResPreUb[pingpongFlag * INPUT1_BUFFER_OFFSET / sizeof(MM2_OUT_T)]; + WaitFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + + uint64_t vec2ResGmOffset = ((info.loop - 1) % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + inOutBaseOffset; + DataCopy(bmm2ResPreUb, vec2ResGm[vec2ResGmOffset], vec2ComputeSize); + + SetFlag(SYNC_INPUT_BUF1_FLAG); + WaitFlag(SYNC_INPUT_BUF1_FLAG); + + uint32_t idx = info.loop % (constInfo.preLoadNum); + LocalTensor expUb = v0ValidSizeBuff.Get()[384]; // sumUb用临时内存 16 * 32B = 512B + Brcb(expUb, softmaxExpUb[idx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset], (dealRowCount + 7) / 8, + {1, 8}); + PipeBarrier(); + + RowMuls(bmm2ResPreUb, bmm2ResPreUb, expUb, dealRowCount, columnCount, actualColumnCount); + AscendC::PipeBarrier(); + Add(bmm2ResUb, bmm2ResUb, bmm2ResPreUb, vec2ComputeSize); + AscendC::PipeBarrier(); + + SetFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + pingpongFlag ^= 1; // pingpong 0 1 切换 + } + + // 最后一次输出计算结果，否则将中间结果暂存至workspace + if (info.isLastS2Loop) { + uint32_t idx = info.loop % (constInfo.preLoadNum); + LocalTensor tmpSumUb = v0ValidSizeBuff.Get()[384]; // sumUb用临时内存 16 * 32B = 512B + Brcb(tmpSumUb, softmaxSumUb[idx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset], (dealRowCount + 7) / 8, + {1, 8}); + PipeBarrier(); + RowDivs(bmm2ResUb, bmm2ResUb, tmpSumUb, dealRowCount, columnCount, actualColumnCount); + PipeBarrier(); + Bmm2ResCopyOut(info, bmm2ResUb, mStart, dealRowCount, columnCount, actualColumnCount); + } else { + LocalTensor outUb = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + DataCopy(outUb, bmm2ResUb, dealRowCount * columnCount); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + uint64_t vec2ResGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + inOutBaseOffset; + DataCopy(vec2ResGm[vec2ResGmOffset], outUb, vec2ComputeSize); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + } +} + +template +__aicore__ inline void SASVectorBlock::RowDivs(LocalTensor dstUb, LocalTensor src0Ub, + LocalTensor src1Ub, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + // divs by row, 每行的元素除以相同的元素 + // dstUb[i, (j * 8) : (j * 8 + 7)] = src0Ub[i, (j * 8) : (j * 8 + 7)] / src1Ub[i, 0 : 7] + // src0Ub:[dealRowCount, columnCount], src1Ub:[dealRowCount, FP32_BLOCK_ELEMENT_NUM] dstUb:[dealRowCount, + // columnCount] + uint32_t dtypeMask = FP32_REPEAT_ELEMENT_NUM; + uint32_t dLoop = actualColumnCount / dtypeMask; + uint32_t dRemain = actualColumnCount % dtypeMask; + + BinaryRepeatParams repeatParamsDiv; + repeatParamsDiv.src0BlkStride = 1; + repeatParamsDiv.src1BlkStride = 0; + repeatParamsDiv.dstBlkStride = 1; + repeatParamsDiv.src0RepStride = columnCount / FP32_BLOCK_ELEMENT_NUM; + repeatParamsDiv.src1RepStride = 1; + repeatParamsDiv.dstRepStride = columnCount / FP32_BLOCK_ELEMENT_NUM; + uint32_t columnRepeatCount = dLoop; + if (columnRepeatCount <= dealRowCount) { + uint32_t offset = 0; + for (uint32_t i = 0; i < dLoop; i++) { + Div(dstUb[offset], src0Ub[offset], src1Ub, dtypeMask, dealRowCount, repeatParamsDiv); + offset += dtypeMask; + } + } else { + BinaryRepeatParams columnRepeatParams; + columnRepeatParams.src0BlkStride = 1; + columnRepeatParams.src1BlkStride = 0; + columnRepeatParams.dstBlkStride = 1; + columnRepeatParams.src0RepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + columnRepeatParams.src1RepStride = 0; + columnRepeatParams.dstRepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + uint32_t offset = 0; + for (uint32_t i = 0; i < dealRowCount; i++) { + Div(dstUb[offset], src0Ub[offset], src1Ub[i * FP32_BLOCK_ELEMENT_NUM], dtypeMask, columnRepeatCount, + columnRepeatParams); + offset += columnCount; + } + } + if (dRemain > 0) { + Div(dstUb[dLoop * dtypeMask], src0Ub[dLoop * dtypeMask], src1Ub, dRemain, dealRowCount, repeatParamsDiv); + } +} + +template +__aicore__ inline void SASVectorBlock::RowMuls(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount) +{ + // muls by row, 每行的元素乘以相同的元素 + // dstUb[i, (j * 8) : (j * 8 + 7)] = src0Ub[i, (j * 8) : (j * 8 + 7)] * src1Ub[i, 0 : 7] + // src0Ub:[dealRowCount, columnCount] src1Ub:[dealRowCount, FP32_BLOCK_ELEMENT_NUM] dstUb:[dealRowCount, + // columnCount] + // dealRowCount is repeat times, must be less 256 + uint32_t repeatElementNum = FP32_REPEAT_ELEMENT_NUM; + uint32_t blockElementNum = FP32_BLOCK_ELEMENT_NUM; + + if constexpr (std::is_same::value) { + // 此限制由于每个repeat至多连续读取256B数据 + repeatElementNum = FP32_REPEAT_ELEMENT_NUM * 2; // 256/4 * 2=128 + blockElementNum = FP32_BLOCK_ELEMENT_NUM * 2; // 32/4 * 2 = 16 + } + + // 每次只能连续读取256B的数据进行计算，故每次只能处理256B/sizeof(dType)= + // 列方向分dLoop次，每次处理8列数据 + uint32_t dLoop = actualColumnCount / repeatElementNum; + uint32_t dRemain = actualColumnCount % repeatElementNum; + // REPEATE_STRIDE_UP_BOUND=256， 此限制由于src0RepStride数据类型为uint8之多256个datablock间距 + if (columnCount < REPEATE_STRIDE_UP_BOUND * blockElementNum) { + BinaryRepeatParams repeatParams; + repeatParams.src0BlkStride = 1; + repeatParams.src1BlkStride = 0; + repeatParams.dstBlkStride = 1; + repeatParams.src0RepStride = columnCount / blockElementNum; + repeatParams.src1RepStride = 1; + repeatParams.dstRepStride = columnCount / blockElementNum; + + // 如果以列为repeat所处理的次数小于行处理次数，则以列方式处理。反之则以行进行repeat处理 + if (dLoop <= dealRowCount) { + uint32_t offset = 0; + for (uint32_t i = 0; i < dLoop; i++) { + Mul(dstUb[offset], src0Ub[offset], src1Ub, repeatElementNum, dealRowCount, repeatParams); + offset += repeatElementNum; + } + } else { + BinaryRepeatParams columnRepeatParams; + columnRepeatParams.src0BlkStride = 1; + columnRepeatParams.src1BlkStride = 0; + columnRepeatParams.dstBlkStride = 1; + columnRepeatParams.src0RepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + columnRepeatParams.src1RepStride = 0; + columnRepeatParams.dstRepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + for (uint32_t i = 0; i < dealRowCount; i++) { + Mul(dstUb[i * columnCount], src0Ub[i * columnCount], src1Ub[i * blockElementNum], repeatElementNum, + dLoop, columnRepeatParams); + } + } + + // 最后一次完成[dealRowCount, dRemain] * [dealRowCount, blockElementNum] 只计算有效部分 + if (dRemain > 0) { + Mul(dstUb[dLoop * repeatElementNum], src0Ub[dLoop * repeatElementNum], src1Ub, dRemain, dealRowCount, + repeatParams); + } + } else { + BinaryRepeatParams repeatParams; + repeatParams.src0RepStride = 8; // 每个repeat为256B数据，正好8个datablock + repeatParams.src0BlkStride = 1; + repeatParams.src1RepStride = 0; + repeatParams.src1BlkStride = 0; + repeatParams.dstRepStride = 8; + repeatParams.dstBlkStride = 1; + // 每次计算一行，共计算dealRowCount行 + for (uint32_t i = 0; i < dealRowCount; i++) { + // 计算一行中的dLoop个repeat, 每个repeat计算256/block_size 个data_block + Mul(dstUb[i * columnCount], src0Ub[i * columnCount], src1Ub[i * blockElementNum], repeatElementNum, dLoop, + repeatParams); + // 计算一行中的尾块 + if (dRemain > 0) { + Mul(dstUb[i * columnCount + dLoop * repeatElementNum], + src0Ub[i * columnCount + dLoop * repeatElementNum], src1Ub[i * blockElementNum], dRemain, 1, + repeatParams); + } + } + } +} +} // namespace SASKernel +#endif // SPARSE_ATTN_SHAREDKV_SCFA_BLOCK_VECTOR_H diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_scfa_kernel.h b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_scfa_kernel.h new file mode 100644 index 000000000..a2a3eb105 --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_scfa_kernel.h @@ -0,0 +1,830 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_scfa_kernel.h + * \brief + */ + +#ifndef SPARSE_ATTN_SHAREDKV_SCFA_KERNEL_H +#define SPARSE_ATTN_SHAREDKV_SCFA_KERNEL_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../sparse_attn_sharedkv_common.h" +#include "sparse_attn_sharedkv_scfa_block_cube.h" +#include "sparse_attn_sharedkv_scfa_block_vector.h" +#include "../sparse_attn_sharedkv_metadata.h" + +namespace SASKernel { +using namespace matmul; +using namespace optiling; +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +// 由于S2循环前，RunInfo还没有赋值，使用Bngs1Param临时存放B、N、S1轴相关的信息；同时减少重复计算 +struct TempLoopInfo { + uint32_t bn2IdxInCurCore = 0; + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint64_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小 + uint32_t s2LoopTimes = 0U; // S2方向循环的总次数，无论TND还是BXXD都是等于实际次数，不用减1 + + int32_t actS1Size = 0; // TND场景下当前Batch循环处理的S1轴的大小 + int32_t actOriS2Size = 0; + int32_t actCmpS2Size = 0; + + bool curActSeqLenIsZero = false; + + uint32_t tndCoreStartKVSplitPos = 0; + bool tndIsS2SplitCore = false; + uint32_t gS1Idx = 0U; + uint32_t s1StartIdx = 0; + uint32_t s1EndIdx = 0; + uint64_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小 + uint32_t cmpLoopTimes = 0; + uint32_t oriLoopTimes = 0; + uint32_t v0OriSize = 0; + uint32_t v0CmpSize = 0; + + // sparsemode = 4 + int32_t oriMaskRight = 0; + int32_t oriMaskLeft = 0; + + // sparsemode = 3 + int32_t cmpMaskRight = 0; + + uint64_t actualSeqQPrefixSum = 0; + uint64_t actualSeqKVPrefixSum = 0; + uint64_t actualSeqCmpKVPrefixSum = 0; +}; + +template +class SparseAttnSharedkvScfa { +public: + // 中间计算数据类型为float，高精度模式 + using T = float; + using Q_T = typename SAST::queryType; + using KV_T = typename SAST::kvType; + using OUT_T = typename SAST::outputType; + using SINKS_T = float; + using UPDATE_T = T; + using MM1_OUT_T = T; + using MM2_OUT_T = T; + + __aicore__ inline SparseAttnSharedkvScfa(){}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *oriKV, __gm__ uint8_t *cmpKV, + __gm__ uint8_t *cmpSparseIndices, __gm__ uint8_t *oriBlockTable, + __gm__ uint8_t *cmpBlockTable, __gm__ uint8_t *cuSeqlensQ, + __gm__ uint8_t* cuSeqlensKV, __gm__ uint8_t *cuSeqlensCmpKV, + __gm__ uint8_t *seqUsedQ, __gm__ uint8_t *seqUsedKV, __gm__ uint8_t *sinks, + __gm__ uint8_t *metadata, __gm__ uint8_t *attentionOut, + __gm__ uint8_t *softmaxLse, __gm__ uint8_t *workspace, + const SparseAttnSharedkvTilingData *__restrict tiling, __gm__ uint8_t *gmTiling, + TPipe *tPipe); + + __aicore__ inline void Process(); + +private: + static constexpr bool PAGE_ATTENTION = SAST::pageAttention; + static constexpr bool FLASH_DECODE = SAST::flashDecode; + static constexpr SAS_LAYOUT LAYOUT_T = SAST::layout; + static constexpr SAS_LAYOUT KV_LAYOUT_T = SAST::kvLayout; + + static constexpr uint32_t PRELOAD_NUM = 2; + static constexpr uint32_t N_BUFFER_M_BASIC_SIZE = 256; + static constexpr uint32_t SAS_PRELOAD_TASK_CACHE_SIZE = 3; + + static constexpr uint32_t SYNC_V0_C1_FLAG = 6; + static constexpr uint32_t SYNC_C1_V1_FLAG = 7; + static constexpr uint32_t SYNC_V1_C2_FLAG = 8; + static constexpr uint32_t SYNC_C2_V2_FLAG = 9; + + static constexpr uint64_t kvHeadNum = 1ULL; + static constexpr uint64_t headDim = 512ULL; + + static constexpr uint32_t dbWorkspaceRatio = PRELOAD_NUM; + + const SparseAttnSharedkvTilingData *__restrict tilingData = nullptr; + + TPipe *pipe = nullptr; + GlobalTensor metadataGm; + uint64_t mSizeVStart = 0ULL; + uint64_t topKBaseOffset = 0ULL; + uint64_t tensorACoreOffset = 0ULL; + uint64_t tensorBCoreOffset = 0ULL; + uint64_t tensorCmpBCoreOffset = 0ULL; + + uint32_t tmpBlockIdx = 0U; + uint32_t aiCoreIdx = 0U; + + ConstInfo constInfo{}; + TempLoopInfo tempLoopInfo{}; + + SASCubeBlock cubeBlock; + SASVectorBlock vectorBlock; + + GlobalTensor queryGm; + GlobalTensor oriKvGm; + GlobalTensor cmpKvGm; + GlobalTensor sinksGm; + + GlobalTensor attentionOutGm; + GlobalTensor softmaxLseGm; + + GlobalTensor oriBlockTableGm; + GlobalTensor cmpBlockTableGm; + GlobalTensor topKGm; + + GlobalTensor actualSeqLengthsQGm; + GlobalTensor actualSeqLengthsKVGm; + GlobalTensor actualSeqLengthsCmpKVGm; + + // workspace + GlobalTensor mm1ResGm; + GlobalTensor vec1ResGm; + GlobalTensor mm2ResGm; + GlobalTensor kvMergeGm_; + GlobalTensor kvValidSizeGm_; + + GlobalTensor vec2ResGm; + + GlobalTensor accumOutGm; + + // ================================Init functions================================== + __aicore__ inline void InitTilingData(); + __aicore__ inline void InitCalcParamsEach(); + __aicore__ inline void InitBuffers(); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsKV); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsKV, + __gm__ uint8_t *actualSeqLengthsCmpKV); + __aicore__ inline void InitOutputSingleCore(); + // ================================Process functions================================ + __aicore__ inline void ProcessBalance(); + __aicore__ inline void PreloadPipeline(uint32_t loop, uint32_t cmpLoop, uint64_t s2Start, uint64_t s2LoopIdx, + RunInfo extraInfo[SAS_PRELOAD_TASK_CACHE_SIZE]); + // ================================Offset Calc===================================== + __aicore__ inline void GetSparseActualSeqLen(); + __aicore__ inline void UpdateInnerLoopCond(); + __aicore__ inline void CalcParams(uint32_t loop, uint32_t cmpLoop, uint64_t s2Start, uint32_t s2LoopIdx, + RunInfo &info); + __aicore__ inline int32_t GetActualSeqLenQ(uint32_t bIdx); + __aicore__ inline int32_t GetActualSeqLenKV(uint32_t bIdx); + __aicore__ inline void GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx); + // ================================Mm1============================================== + __aicore__ inline void ComputeMm1(const RunInfo &info); + // ================================Mm2============================================== + __aicore__ inline void ComputeMm2(const RunInfo &info); + __aicore__ inline void InitAllZeroOutput(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx); +}; + +template +__aicore__ inline void SparseAttnSharedkvScfa::InitTilingData() +{ + // singleCoreParams + // singleCoreTensorSize + constInfo.mmResUbSize = tilingData->baseParams.mmResUbSize; + constInfo.bmm2ResUbSize = tilingData->baseParams.bmm2ResUbSize; + + // baseParams + constInfo.batchSize = tilingData->baseParams.batchSize; + constInfo.qHeadNum = constInfo.gSize = tilingData->baseParams.nNumOfQInOneGroup; + constInfo.kvSeqSize = tilingData->baseParams.kvSeqSize; + constInfo.qSeqSize = tilingData->baseParams.qSeqSize; + constInfo.oriMaxBlockNumPerBatch = tilingData->baseParams.oriMaxBlockNumPerBatch; + constInfo.cmpMaxBlockNumPerBatch = tilingData->cmpParams.cmpMaxBlockNumPerBatch; + constInfo.kvCacheBlockSize = tilingData->baseParams.paBlockSize; + constInfo.paOriBlockSize = tilingData->baseParams.oriBlockSize; + constInfo.paCmpBlockSize = tilingData->baseParams.cmpBlockSize; + constInfo.outputLayout = static_cast(tilingData->baseParams.outputLayout); + constInfo.kvHeadNum = kvHeadNum; + constInfo.headDim = headDim; + constInfo.oriMaskMode = tilingData->baseParams.oriMaskMode; + constInfo.oriKvStride = tilingData->baseParams.oriKvStride; + constInfo.oriWinLeft = tilingData->baseParams.oriWinLeft; + constInfo.oriWinRight = tilingData->baseParams.oriWinRight; + + constInfo.actualLenDimsQ = tilingData->baseParams.actualLenDimsQ; + constInfo.actualLenDimsKV = tilingData->baseParams.actualLenDimsKV; + constInfo.returnSoftmaxLse = tilingData->baseParams.returnSoftmaxLse; + // innerSplitParams + constInfo.mBaseSize = constInfo.gSize; + constInfo.s2BaseSize = tilingData->baseParams.s2BaseSize; + + constInfo.preLoadNum = PRELOAD_NUM; + constInfo.nBufferMBaseSize = N_BUFFER_M_BASIC_SIZE; + constInfo.syncV0C1 = SYNC_V0_C1_FLAG; + constInfo.syncC1V1 = SYNC_C1_V1_FLAG; + constInfo.syncV1C2 = SYNC_V1_C2_FLAG; + constInfo.syncC2V2 = SYNC_C2_V2_FLAG; + + // cmp + constInfo.cmpRatio = tilingData->cmpParams.cmpRatio; + constInfo.sparseBlockCount = tilingData->cmpParams.sparseBlockCount; + constInfo.sparseBlockSize = 1; // sparseBlockSize 固定为1 + constInfo.cmpMaskMode = tilingData->cmpParams.cmpMaskMode; + constInfo.cmpKvStride = tilingData->cmpParams.cmpKvStride; +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::InitBuffers() +{ + if ASCEND_IS_AIV { + vectorBlock.InitBuffers(pipe); + } else { + cubeBlock.InitBuffers(pipe); + } +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengthsKV) +{ + if (constInfo.actualLenDimsKV != 0) { + actualSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsKV, constInfo.actualLenDimsKV); + } + if (constInfo.actualLenDimsQ != 0) { + actualSeqLengthsQGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsQ, constInfo.actualLenDimsQ); + } +} + +template +__aicore__ inline void +SparseAttnSharedkvScfa::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsKV, + __gm__ uint8_t *actualSeqLengthsCmpKV) +{ + if (constInfo.actualLenDimsKV != 0) { + actualSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsKV, constInfo.actualLenDimsKV); + actualSeqLengthsCmpKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsCmpKV, constInfo.actualLenDimsKV); + } + if (constInfo.actualLenDimsQ != 0) { + actualSeqLengthsQGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsQ, constInfo.actualLenDimsQ); + } +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::InitAllZeroOutput(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx) +{ + if (constInfo.outputLayout == SAS_LAYOUT::TND) { + if (tempLoopInfo.actS1Size == 0) { + return; + } + uint32_t tBase = actualSeqLengthsQGm.GetValue(bIdx); + uint32_t s1Count = tempLoopInfo.actS1Size; + + uint64_t attenOutOffset = (tBase + s1Idx) * kvHeadNum * constInfo.gSize * headDim + // T轴、s1轴偏移 + n2Idx * constInfo.gSize * headDim; // N2轴偏移 + uint64_t lseOffset = (tBase + s1Idx) * constInfo.gSize + // T轴、s1轴偏移 + n2Idx * constInfo.qSeqSize * constInfo.gSize; // N2轴偏移 + matmul::InitOutput(attentionOutGm[attenOutOffset], constInfo.gSize * headDim, 0); + if (constInfo.returnSoftmaxLse) { + matmul::InitOutput(softmaxLseGm[lseOffset], constInfo.gSize, 0); + } + } else if (constInfo.outputLayout == SAS_LAYOUT::BSND) { + uint64_t attenOutOffset = bIdx * constInfo.qSeqSize * kvHeadNum * constInfo.gSize * headDim + + s1Idx * kvHeadNum * constInfo.gSize * headDim + // B轴、S1轴偏移 + n2Idx * constInfo.gSize * headDim; // N2轴偏移 + uint64_t lseOffset = bIdx * constInfo.qSeqSize * constInfo.kvHeadNum * constInfo.gSize + // B轴偏移 + n2Idx * constInfo.qSeqSize * constInfo.gSize + // N2轴偏移 + s1Idx * constInfo.gSize; // S1轴偏移 + matmul::InitOutput(attentionOutGm[attenOutOffset], constInfo.gSize * headDim, 0); + if (constInfo.returnSoftmaxLse) { + matmul::InitOutput(softmaxLseGm[lseOffset], constInfo.gSize, 0); + } + } +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::InitOutputSingleCore() +{ + uint32_t coreNum = GetBlockNum(); + if (coreNum != 0) { + uint64_t totalOutputSize = constInfo.batchSize * constInfo.qHeadNum * constInfo.qSeqSize * constInfo.headDim; + uint64_t singleCoreSize = (totalOutputSize + (2 * coreNum) - 1) / (2 * coreNum); // 2 means c:v = 1:2 + uint64_t tailSize = totalOutputSize - tmpBlockIdx * singleCoreSize; + uint64_t singleInitOutputSize = tailSize < singleCoreSize ? tailSize : singleCoreSize; + if (singleInitOutputSize > 0) { + matmul::InitOutput(attentionOutGm[tmpBlockIdx * singleCoreSize], singleInitOutputSize, 0); + } + SyncAll(); + } +} + +template +__aicore__ inline int32_t SparseAttnSharedkvScfa::GetActualSeqLenQ(uint32_t bIdx) +{ + if constexpr (LAYOUT_T == SAS_LAYOUT::TND) { + int32_t actualSeqQPrefixSum = actualSeqLengthsQGm.GetValue(bIdx); + int32_t actualSeqQNextSum = actualSeqLengthsQGm.GetValue(bIdx + 1); + tempLoopInfo.actualSeqQPrefixSum = static_cast(actualSeqQPrefixSum); + return actualSeqQNextSum - actualSeqQPrefixSum; + } else { + tempLoopInfo.actualSeqQPrefixSum = static_cast(bIdx * constInfo.qSeqSize); + if (constInfo.actualLenDimsQ == 0) { + return static_cast(constInfo.qSeqSize); + } else { + return actualSeqLengthsQGm.GetValue(bIdx); + } + } +} + +template +__aicore__ inline int32_t SparseAttnSharedkvScfa::GetActualSeqLenKV(uint32_t bIdx) +{ + if constexpr (KV_LAYOUT_T == SAS_LAYOUT::PA_ND) { + tempLoopInfo.actualSeqKVPrefixSum = static_cast(bIdx * constInfo.kvSeqSize); + if (constInfo.actualLenDimsKV == 0) { + return static_cast(constInfo.kvSeqSize); + } + return actualSeqLengthsKVGm.GetValue(bIdx); + } else if constexpr(KV_LAYOUT_T == SAS_LAYOUT::BSND) { + return static_cast(constInfo.kvSeqSize); + } else if constexpr(KV_LAYOUT_T == SAS_LAYOUT::TND) { + int32_t actualSeqKVPrefixSum = actualSeqLengthsKVGm.GetValue(bIdx); + int32_t actualSeqKVNextSum = actualSeqLengthsKVGm.GetValue(bIdx + 1); + tempLoopInfo.actualSeqCmpKVPrefixSum = actualSeqLengthsCmpKVGm.GetValue(bIdx); + tempLoopInfo.actualSeqKVPrefixSum = actualSeqKVPrefixSum; + return actualSeqKVNextSum - actualSeqKVPrefixSum; + } +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::GetSparseActualSeqLen() +{ + // 行无效通过ori部分判断, ori部分如果有行无效那么ori和cmp都有 + if (static_cast(tempLoopInfo.s1EndIdx) < -(tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size)) { + tempLoopInfo.actOriS2Size = 0; + tempLoopInfo.actCmpS2Size = 0; + return; + } + + // 对于cmp部分还有top k, tempLoopInfo.actS2Size只针对cmp + int32_t thresHold = (tempLoopInfo.cmpMaskRight + tempLoopInfo.s1EndIdx + 1) / constInfo.cmpRatio; + tempLoopInfo.actCmpS2Size = Min(constInfo.sparseBlockCount * constInfo.sparseBlockSize, thresHold); +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::UpdateInnerLoopCond() +{ + if ((tempLoopInfo.actCmpS2Size == 0 && tempLoopInfo.actOriS2Size == 0) || (tempLoopInfo.actS1Size == 0)) { + tempLoopInfo.curActSeqLenIsZero = true; + return; + } + tempLoopInfo.curActSeqLenIsZero = false; + tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize; + tempLoopInfo.mBasicSizeTail = + (tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail; +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::Init( + __gm__ uint8_t *query, __gm__ uint8_t *oriKV, __gm__ uint8_t *cmpKV, __gm__ uint8_t *cmpSparseIndices, + __gm__ uint8_t *oriBlockTable, __gm__ uint8_t *cmpBlockTable, __gm__ uint8_t *cuSeqlensQ, + __gm__ uint8_t* cuSeqlensKV, __gm__ uint8_t *cuSeqlensCmpKV, __gm__ uint8_t *seqUsedQ, + __gm__ uint8_t *seqUsedKV, __gm__ uint8_t *sinks, __gm__ uint8_t *metadata, __gm__ uint8_t *attentionOut, __gm__ uint8_t *softmaxLse, + __gm__ uint8_t *workspace, const SparseAttnSharedkvTilingData *__restrict tiling, __gm__ uint8_t *gmTiling, + TPipe *tPipe) +{ + if ASCEND_IS_AIV { + tmpBlockIdx = GetBlockIdx(); // vec:0-47 + aiCoreIdx = tmpBlockIdx / 2; + } else { + tmpBlockIdx = GetBlockIdx(); // cube:0-23 + aiCoreIdx = tmpBlockIdx; + } + + // init tiling data + tilingData = tiling; + InitTilingData(); + if (KV_LAYOUT_T == SAS_LAYOUT::TND && LAYOUT_T == SAS_LAYOUT::TND) { + InitActualSeqLen(cuSeqlensQ, cuSeqlensKV, cuSeqlensCmpKV); + } else if (KV_LAYOUT_T == SAS_LAYOUT::TND) { + InitActualSeqLen(seqUsedQ, cuSeqlensKV, cuSeqlensCmpKV); + } else if ((KV_LAYOUT_T == SAS_LAYOUT::PA_ND || KV_LAYOUT_T == SAS_LAYOUT::BSND) && LAYOUT_T == SAS_LAYOUT::TND) { + InitActualSeqLen(cuSeqlensQ, seqUsedKV); + } else if ((KV_LAYOUT_T == SAS_LAYOUT::PA_ND || KV_LAYOUT_T == SAS_LAYOUT::BSND)) { + InitActualSeqLen(seqUsedQ, seqUsedKV); + } + + metadataGm.SetGlobalBuffer((__gm__ uint32_t *)metadata); + InitCalcParamsEach(); + + pipe = tPipe; + // init global buffer + queryGm.SetGlobalBuffer((__gm__ Q_T *)query); + oriKvGm.SetGlobalBuffer((__gm__ KV_T *)oriKV); + cmpKvGm.SetGlobalBuffer((__gm__ KV_T *)cmpKV); + + if (sinks != nullptr) { + sinksGm.SetGlobalBuffer((__gm__ SINKS_T *)sinks); + } + + attentionOutGm.SetGlobalBuffer((__gm__ OUT_T *)attentionOut); + softmaxLseGm.SetGlobalBuffer((__gm__ T *)softmaxLse); + + if ASCEND_IS_AIV { + if (LAYOUT_T != SAS_LAYOUT::TND) { + if (constInfo.needInit) { + InitOutputSingleCore(); + } + } + } + + if constexpr (PAGE_ATTENTION) { + oriBlockTableGm.SetGlobalBuffer((__gm__ int32_t *)oriBlockTable); + cmpBlockTableGm.SetGlobalBuffer((__gm__ int32_t *)cmpBlockTable); + } + topKGm.SetGlobalBuffer((__gm__ int32_t *)cmpSparseIndices); + + // workspace 内存排布 + // |Q--|mm1ResGm|vec1ResGm|mm2ResGm|vec2ResGm + // |Core0_Q1-Core0_Q2-Core1_Q1-Core1_Q2....Core32_Q1-Core32_Q2|Core0_mmRes + uint64_t offset = 0; + mm1ResGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + offset + + aiCoreIdx * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(MM1_OUT_T))); + offset += GetBlockNum() * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(MM1_OUT_T); + + vec1ResGm.SetGlobalBuffer( + (__gm__ Q_T *)(workspace + offset + aiCoreIdx * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(KV_T))); + offset += GetBlockNum() * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(KV_T); + + mm2ResGm.SetGlobalBuffer( + (__gm__ MM2_OUT_T *)(workspace + offset + + aiCoreIdx * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(MM2_OUT_T))); + offset += GetBlockNum() * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(MM2_OUT_T); + + vec2ResGm.SetGlobalBuffer( + (__gm__ T *)(workspace + offset + aiCoreIdx * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(T))); + offset += GetBlockNum() * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(T); + + kvMergeGm_.SetGlobalBuffer((__gm__ KV_T *)(workspace + offset + aiCoreIdx * 512 * 512 * 4 * sizeof(KV_T))); + offset += GetBlockNum() * 512 * 512 * 4 * sizeof(KV_T); + + kvValidSizeGm_.SetGlobalBuffer( + (__gm__ int32_t *)(workspace + offset + (aiCoreIdx * 2) * 128 * 4 * sizeof(int32_t))); + + if ASCEND_IS_AIV { + vectorBlock.InitParams(constInfo, tilingData); + vectorBlock.InitVec0GlobalTensor(kvValidSizeGm_, kvMergeGm_, oriKvGm, cmpKvGm, oriBlockTableGm, + cmpBlockTableGm); + vectorBlock.InitVec1GlobalTensor(mm1ResGm, vec1ResGm, actualSeqLengthsQGm, actualSeqLengthsKVGm, topKGm, + sinksGm, softmaxLseGm); + vectorBlock.InitVec2GlobalTensor(accumOutGm, vec2ResGm, mm2ResGm, attentionOutGm); + } + + if ASCEND_IS_AIC { + cubeBlock.InitParams(constInfo); + cubeBlock.InitMm1GlobalTensor(queryGm, oriKvGm, cmpKvGm, mm1ResGm); + cubeBlock.InitMm2GlobalTensor(vec1ResGm, mm2ResGm, attentionOutGm); + cubeBlock.InitPageAttentionInfo(oriKvGm, kvMergeGm_, oriBlockTableGm, cmpBlockTableGm); + } + // 要在InitParams之后执行 + if (pipe != nullptr) { + InitBuffers(); + } +} + + +template +__aicore__ inline void SparseAttnSharedkvScfa::InitCalcParamsEach() +{ + if (aiCoreIdx != 0) { + constInfo.bN2Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_BN2_START_INDEX, false)); + constInfo.gS1Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_M_START_INDEX, false)); + constInfo.s2Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_S2_START_INDEX, false)); + } + constInfo.bN2End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_BN2_END_INDEX, false)); + constInfo.gS1End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_M_END_INDEX, false)); + constInfo.s2End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_S2_END_INDEX, false)); +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::CalcParams(uint32_t loop, uint32_t cmpLoop, uint64_t s2Start, + uint32_t s2LoopIdx, RunInfo &info) +{ + info.isValid = s2LoopIdx < tempLoopInfo.s2LoopTimes; + info.loop = loop; + info.cmpLoop = cmpLoop; + info.bIdx = tempLoopInfo.bIdx; + info.n2IdxReal = tempLoopInfo.n2Idx; + + info.gS1Idx = tempLoopInfo.gS1Idx; + info.s1Idx = tempLoopInfo.gS1Idx / constInfo.gSize; + info.s2Idx = s2LoopIdx; + info.curSInnerLoopTimes = tempLoopInfo.s2LoopTimes; + info.tndIsS2SplitCore = tempLoopInfo.tndIsS2SplitCore; + info.tndCoreStartKVSplitPos = tempLoopInfo.tndCoreStartKVSplitPos; + info.isBmm2Output = false; + info.actS1Size = tempLoopInfo.actS1Size; + + // M方向的尾块 + info.actMBaseSize = tempLoopInfo.mBasicSizeTail; + + if ASCEND_IS_AIV { + info.mSize = info.actMBaseSize; + info.mSizeV = (info.mSize <= 16) ? info.mSize : ((CeilDiv(info.mSize, 16) + 1) / 2 * 16); + info.mSizeVStart = 0; + if (tmpBlockIdx % 2 == 1) { + info.mSizeVStart = info.mSizeV; + info.mSizeV = info.mSize - info.mSizeV; + } + } + + info.isFirstSInnerLoop = s2LoopIdx == s2Start; + if (info.isFirstSInnerLoop) { + tempLoopInfo.bn2IdxInCurCore++; + } + info.isLastS2Loop = (s2LoopIdx == (tempLoopInfo.s2LoopTimes - 1)); + info.bn2IdxInCurCore = tempLoopInfo.bn2IdxInCurCore - 1; + + uint64_t tndBIdxOffsetForQ = tempLoopInfo.actualSeqQPrefixSum * constInfo.qHeadNum * constInfo.headDim; + uint64_t tndBIdxOffsetForKV = tempLoopInfo.actualSeqKVPrefixSum * constInfo.kvHeadNum * constInfo.headDim; + uint64_t tndBIdxOffsetForCmpKV = tempLoopInfo.actualSeqCmpKVPrefixSum * constInfo.kvHeadNum * constInfo.headDim; + + if (info.isFirstSInnerLoop) { + tensorACoreOffset = tndBIdxOffsetForQ + info.gS1Idx * constInfo.headDim; + tensorBCoreOffset = tndBIdxOffsetForKV + info.n2Idx * constInfo.headDim; // 当前为PA场景，该变量失效 + tensorCmpBCoreOffset = tndBIdxOffsetForCmpKV + info.n2Idx * constInfo.headDim; + if constexpr (LAYOUT_T == SAS_LAYOUT::BSND) { // B,S1,N2 K + topKBaseOffset = (info.bIdx * constInfo.qSeqSize + tempLoopInfo.s1StartIdx) * constInfo.kvHeadNum * + constInfo.sparseBlockCount + + info.n2Idx * constInfo.sparseBlockCount; + } else if (LAYOUT_T == SAS_LAYOUT::TND) { // T N2 K + topKBaseOffset = (tempLoopInfo.actualSeqQPrefixSum + tempLoopInfo.s1StartIdx) * constInfo.kvHeadNum * + constInfo.sparseBlockCount + + info.n2Idx * constInfo.sparseBlockCount; + } + } + info.tensorAOffset = tensorACoreOffset; + info.tensorBOffset = tensorBCoreOffset; + info.tensorCmpBOffset = tensorCmpBCoreOffset; + info.attenOutOffset = tensorACoreOffset; + info.topKBaseOffset = topKBaseOffset; + + if (s2LoopIdx < tempLoopInfo.oriLoopTimes) { + // S2首次循环只能在ori_kv + info.isOri = true; + info.relativeS2Idx = 0; + uint64_t s2Offset = info.s2Idx * constInfo.s2BaseSize; + if (s2LoopIdx + 1 == tempLoopInfo.oriLoopTimes) { + info.actualSingleProcessSInnerSize = (tempLoopInfo.oriMaskRight - tempLoopInfo.oriMaskLeft + 1) - s2Offset; + } else { + info.actualSingleProcessSInnerSize = constInfo.s2BaseSize; + } + info.s2StartPoint = tempLoopInfo.oriMaskLeft; + info.cmpS2IdLimit = (tempLoopInfo.cmpMaskRight + tempLoopInfo.s1EndIdx + 1) / constInfo.cmpRatio; + } else { + info.isOri = false; + info.relativeS2Idx = info.s2Idx - tempLoopInfo.oriLoopTimes; + uint64_t s2Offset = (info.s2Idx - tempLoopInfo.oriLoopTimes) * constInfo.s2BaseSize; + if (s2LoopIdx + 1 == tempLoopInfo.s2LoopTimes) { + info.actualSingleProcessSInnerSize = tempLoopInfo.actCmpS2Size - s2Offset; + } else { + info.actualSingleProcessSInnerSize = constInfo.s2BaseSize; + } + info.s2StartPoint = 0; + info.cmpS2IdLimit = (tempLoopInfo.cmpMaskRight + tempLoopInfo.s1EndIdx + 1) / constInfo.cmpRatio; + } + + info.actualSingleProcessSInnerSizeAlign = SASAlign(info.actualSingleProcessSInnerSize, SASVectorBlock::BYTE_BLOCK); + if (info.isOri) { + info.v0S2Start = 0; + info.v0S2DealSize = 0; + } else { + info.v0S2Start = 0; + if (s2LoopIdx + 1 == tempLoopInfo.s2LoopTimes && s2LoopIdx == 2) { // tail + info.v0S2Start = 512; + } + info.v0S2DealSize = 512; + } +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::ComputeMm1(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = CeilDiv(info.actMBaseSize, constInfo.nBufferMBaseSize); + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + cubeBlock.ComputeMm1(info, mSplitInfo); + CrossCoreSetFlag(constInfo.syncC1V1); + } +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::ComputeMm2(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + CrossCoreWaitFlag(constInfo.syncV1C2); + cubeBlock.ComputeMm2(info, mSplitInfo); + CrossCoreSetFlag(constInfo.syncC2V2); + } +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::Process() +{ + uint32_t hasLoad = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_CORE_ENABLE_INDEX, false)); + if (hasLoad == 0) { + return; + } + if ASCEND_IS_AIV { + vectorBlock.AllocEventID(); + vectorBlock.InitSoftmaxDefaultBuffer(); + } else { + cubeBlock.AllocEventID(); + } + ProcessBalance(); + if ASCEND_IS_AIV { + vectorBlock.FreeEventID(); + } else { + cubeBlock.FreeEventID(); + } +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx) +{ + bIdx = bN2Idx / kvHeadNum; + n2Idx = bN2Idx % kvHeadNum; +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::ProcessBalance() +{ + RunInfo extraInfo[SAS_PRELOAD_TASK_CACHE_SIZE]; + uint32_t gloop = 0; + uint32_t cmpLoop = 0; + uint32_t gS1LoopEnd = 0; + bool globalLoopStart = true; + + if ASCEND_IS_AIC { + CrossCoreSetFlag(3); + CrossCoreSetFlag(3); + CrossCoreSetFlag(3); + CrossCoreSetFlag(3); + } + + // 适配左闭右开 + if (constInfo.bN2Start == constInfo.bN2End) { + if (constInfo.gS1Start != constInfo.gS1End || constInfo.s2Start != constInfo.s2End) { + constInfo.bN2End += 1; + } + } else if ((constInfo.gS1End != 0) || (constInfo.s2End != 0)) { + constInfo.bN2End += 1; + } + + for (uint32_t bN2LoopIdx = constInfo.bN2Start; bN2LoopIdx < constInfo.bN2End; bN2LoopIdx++) { + GetBN2Idx(bN2LoopIdx, tempLoopInfo.bIdx, tempLoopInfo.n2Idx); + tempLoopInfo.actS1Size = GetActualSeqLenQ(tempLoopInfo.bIdx); // 获取actualSeqLength + bool isS1ZeroAndLastBatch = (tempLoopInfo.actS1Size == 0) && + ((constInfo.outputLayout == SAS_LAYOUT::BSND) || (bN2LoopIdx + 1 == constInfo.bN2End)); + uint32_t gS1SplitNum = CeilDiv(tempLoopInfo.actS1Size * constInfo.gSize, constInfo.mBaseSize); + + // 当处于最后一个BN2时, 且gS1End为0时, 说明当前BN2里的所有数据都在当前核处理 + gS1LoopEnd = (bN2LoopIdx + 1 == constInfo.bN2End && constInfo.gS1End != 0) ? constInfo.gS1End : gS1SplitNum; + // 当处于最后一个BN2且当前S1为0时，需要进入循环计算preload导致的未完成的部分 + gS1LoopEnd = isS1ZeroAndLastBatch ? gS1LoopEnd + 1 : gS1LoopEnd; + for (uint32_t gS1LoopIdx = constInfo.gS1Start; gS1LoopIdx < gS1LoopEnd; gS1LoopIdx++) { + tempLoopInfo.actOriS2Size = GetActualSeqLenKV(tempLoopInfo.bIdx); + // 计算需要的数据, 避免重复计算 + tempLoopInfo.gS1Idx = gS1LoopIdx * constInfo.mBaseSize; + tempLoopInfo.s1StartIdx = tempLoopInfo.gS1Idx / constInfo.gSize; + tempLoopInfo.s1EndIdx = + Min((tempLoopInfo.s1StartIdx + constInfo.mBaseSize / constInfo.gSize - 1), tempLoopInfo.actS1Size - 1); + + // 此处均为闭区间 + tempLoopInfo.oriMaskRight = tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size + + static_cast(tempLoopInfo.s1EndIdx) + constInfo.oriWinRight; + tempLoopInfo.oriMaskLeft = Max(tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size + + static_cast(tempLoopInfo.s1EndIdx) - constInfo.oriWinLeft, + 0); + tempLoopInfo.cmpMaskRight = tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size; + GetSparseActualSeqLen(); + UpdateInnerLoopCond(); + + uint32_t oriS2Size = tempLoopInfo.oriMaskRight - tempLoopInfo.oriMaskLeft + 1; + uint32_t oriSplitNum = 0; + uint32_t cmpSplitNum = 0; + uint32_t cmpS2Size = 0; + bool isEnd = (bN2LoopIdx + 1 == constInfo.bN2End) && (gS1LoopIdx + 1 == gS1LoopEnd); + if (tempLoopInfo.curActSeqLenIsZero) { + if ASCEND_IS_AIV { + InitAllZeroOutput(tempLoopInfo.bIdx, tempLoopInfo.s1StartIdx, tempLoopInfo.n2Idx); + } + if (!isEnd) { + continue; + } + } else { + oriSplitNum = CeilDiv(oriS2Size, constInfo.s2BaseSize); + cmpS2Size = tempLoopInfo.actCmpS2Size; + cmpSplitNum = CeilDiv(cmpS2Size, constInfo.s2BaseSize); + } + + uint32_t s2SplitNum = oriSplitNum + cmpSplitNum; + constexpr uint32_t V0_SPLIT = 32; // align to 32 + uint32_t v0OriSize = CeilDiv(oriS2Size * cmpS2Size, oriS2Size + cmpS2Size); + if (cmpS2Size > V0_SPLIT * oriSplitNum) { + v0OriSize = SASAlign(v0OriSize, V0_SPLIT * oriSplitNum); + } + uint32_t v0CmpSize = cmpS2Size - v0OriSize; + + tempLoopInfo.oriLoopTimes = oriSplitNum; + tempLoopInfo.cmpLoopTimes = cmpSplitNum; + tempLoopInfo.s2LoopTimes = s2SplitNum; + tempLoopInfo.v0OriSize = v0OriSize; + tempLoopInfo.v0CmpSize = v0CmpSize; + + uint32_t s2LoopEnd = (isEnd && constInfo.s2End != 0) ? constInfo.s2End : tempLoopInfo.s2LoopTimes; + tempLoopInfo.s2LoopTimes = s2LoopEnd; + // 分核修改后需要打开 + // 当前s2是否被切，决定了输出是否要写到attenOut上 + tempLoopInfo.tndIsS2SplitCore = ((constInfo.s2Start == 0) && (s2LoopEnd == s2SplitNum)) ? false : true; + tempLoopInfo.tndCoreStartKVSplitPos = globalLoopStart ? constInfo.coreStartKVSplitPos : 0; + uint32_t extraLoop = isEnd ? 2 : 0; + uint32_t curTopKIdx = 0; + for (uint32_t s2LoopIdx = constInfo.s2Start; s2LoopIdx < (s2LoopEnd + extraLoop); s2LoopIdx++) { + PreloadPipeline(gloop, cmpLoop, constInfo.s2Start, s2LoopIdx, extraInfo); + ++gloop; + if (s2LoopIdx >= tempLoopInfo.oriLoopTimes && s2LoopIdx < s2LoopEnd) { // 用于判断v0使用的循环GM的id + ++cmpLoop; + } + } + globalLoopStart = false; + constInfo.s2Start = 0; + } + constInfo.gS1Start = 0; + } + if ASCEND_IS_AIV { + CrossCoreWaitFlag(3); + CrossCoreWaitFlag(3); + CrossCoreWaitFlag(3); + CrossCoreWaitFlag(3); + } +} + +template +__aicore__ inline void SparseAttnSharedkvScfa::PreloadPipeline(uint32_t loop, uint32_t cmpLoop, uint64_t s2Start, + uint64_t s2LoopIdx, + RunInfo extraInfo[SAS_PRELOAD_TASK_CACHE_SIZE]) +{ + RunInfo &extraInfo0 = extraInfo[loop % SAS_PRELOAD_TASK_CACHE_SIZE]; // 本轮任务 + RunInfo &extraInfo2 = extraInfo[(loop + 2) % SAS_PRELOAD_TASK_CACHE_SIZE]; // 上一轮任务 + RunInfo &extraInfo1 = extraInfo[(loop + 1) % SAS_PRELOAD_TASK_CACHE_SIZE]; // 上两轮任务 + + CalcParams(loop, cmpLoop, s2Start, s2LoopIdx, extraInfo0); + if (extraInfo0.isValid) { + if ASCEND_IS_AIC { + if (!extraInfo0.isOri) { + CrossCoreWaitFlag(constInfo.syncV0C1); + } + ComputeMm1(extraInfo0); + } else { + if (extraInfo0.isFirstSInnerLoop) { + CrossCoreWaitFlag(3); + } + vectorBlock.ProcessVec0L(extraInfo0); + if (!extraInfo0.isOri) { + CrossCoreSetFlag(constInfo.syncV0C1); + } + } + } + if (extraInfo2.isValid) { + if ASCEND_IS_AIV { + vectorBlock.ProcessVec1L(extraInfo2); + } + if ASCEND_IS_AIC { + ComputeMm2(extraInfo2); + if (extraInfo2.isLastS2Loop) { + CrossCoreSetFlag(3); + } + } + } + if (extraInfo1.isValid) { + if ASCEND_IS_AIV { + vectorBlock.ProcessVec2L(extraInfo1); + } + extraInfo1.isValid = false; + } +} + +} // namespace SASKernel +#endif // SPARSE_ATTN_SHAREDKV_SCFA_KERNEL_H \ No newline at end of file diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_swa_block_cube.h b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_swa_block_cube.h new file mode 100644 index 000000000..b86d3d99f --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_swa_block_cube.h @@ -0,0 +1,953 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_swa_block_cube.h + * \brief use 7 buffer for matmul l1, better pipeline + */ +#ifndef SPARSE_ATTN_SHAREDKV_SWA_BLOCK_CUBE_H +#define SPARSE_ATTN_SHAREDKV_SWA_BLOCK_CUBE_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../sparse_attn_sharedkv_common.h" + +namespace SASKernel { +template +class SWACubeBlock { +public: + // 中间计算数据类型为float, 高精度模式 + using T = float; + using Q_T = typename SAST::queryType; + using KV_T = typename SAST::kvType; + using OUT_T = typename SAST::outputType; + using MM_OUT_T = T; + + __aicore__ inline SWACubeBlock(){}; + __aicore__ inline void InitParams(const ConstInfo &constInfo); + __aicore__ inline void InitMm1GlobalTensor(GlobalTensor queryGm, GlobalTensor oriKvGm, + GlobalTensor cmpKV, GlobalTensor mm1ResGm); + __aicore__ inline void InitMm2GlobalTensor(GlobalTensor vec1ResGm, GlobalTensor mm2ResGm, + GlobalTensor attentionOutGm); + __aicore__ inline void InitPageAttentionInfo(GlobalTensor oriKvGm, // const GlobalTensor& kvMergeGm, + GlobalTensor oriBlockTableGm, + GlobalTensor cmpBlockTableGm); + __aicore__ inline void InitBuffers(TPipe *pipe); + + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void ComputeMm1(const RunInfo &info, const MSplitInfo mSplitInfo); + __aicore__ inline void ComputeMm2(const RunInfo &info, const MSplitInfo mSplitInfo); + +private: + static constexpr bool PAGE_ATTENTION = SAST::pageAttention; + // static constexpr int TEMPLATE_MODE = SAST::templateMode; + static constexpr bool FLASH_DECODE = SAST::flashDecode; + static constexpr SAS_LAYOUT LAYOUT_T = SAST::layout; + static constexpr SAS_LAYOUT KV_LAYOUT_T = SAST::kvLayout; + + static constexpr uint32_t M_SPLIT_SIZE = 128; // m方向切分 + static constexpr uint32_t N_SPLIT_SIZE = 128; // n方向切分 + static constexpr uint32_t K_L0_SPLIT_SIZE = 128; // k方向L0切分 + static constexpr uint32_t K_L1_SPLIT_SIZE = 256; // k方向L1切分 + static constexpr uint32_t N_WORKSPACE_SIZE = 512; // n方向切分 + static constexpr uint32_t D_SPLIT_SIZE = 256; // d轴切分 + + static constexpr uint32_t L1_BLOCK_SIZE = (64 * 512 * sizeof(Q_T)); + static constexpr uint32_t L1_BLOCK_OFFSET = 64 * 512; + + static constexpr uint32_t L0A_PP_SIZE = (32 * 1024); + static constexpr uint32_t L0B_PP_SIZE = (32 * 1024); + static constexpr uint32_t L0C_PP_SIZE = (64 * 1024); + + // mte2 <> mte1 EventID + // L1 3buf, 使用3个eventId + static constexpr uint32_t L1_EVENT0 = EVENT_ID2; + static constexpr uint32_t L1_EVENT1 = EVENT_ID3; + static constexpr uint32_t L1_EVENT2 = EVENT_ID4; + static constexpr uint32_t L1_EVENT3 = EVENT_ID5; + static constexpr uint32_t L1_EVENT4 = EVENT_ID6; + static constexpr uint32_t L1_EVENT5 = EVENT_ID7; + static constexpr uint32_t L1_EVENT6 = EVENT_ID1; + + // m <> mte1 EventID + static constexpr uint32_t L0AB_EVENT0 = EVENT_ID3; + static constexpr uint32_t L0AB_EVENT1 = EVENT_ID4; + + static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding; + static constexpr uint32_t mte21QPIds[4] = {L1_EVENT0, L1_EVENT1, L1_EVENT2, L1_EVENT3}; // mte12复用 + static constexpr uint32_t mte21KVIds[3] = {L1_EVENT4, L1_EVENT5, L1_EVENT6}; + + ConstInfo constInfo{}; + + // L1分成3块buf, 用于记录 + uint32_t qpL1BufIter = 0; + uint32_t kvL1BufIter = -1; + uint32_t abL0BufIter = 0; + uint32_t cL0BufIter = 0; + + // mm1 + GlobalTensor queryGm; + GlobalTensor keyGm; + GlobalTensor mm1ResGm; + // GlobalTensor kvMergeGm_; + GlobalTensor oriKvGm; + GlobalTensor cmpKvGm; + + // mm2 + GlobalTensor vec1ResGm; + GlobalTensor valueGm; + GlobalTensor mm2ResGm; + GlobalTensor attentionOutGm; + + // block_table + GlobalTensor oriBlockTableGm; + GlobalTensor cmpBlockTableGm; + + TBuf bufQPL1; + TBuf bufKVL1; + TBuf tmpBufL0A; + TBuf tmpBufL0B; + TBuf tmpBufL0C; + + LocalTensor l1QPTensor; + LocalTensor l1KVTensor; + LocalTensor aL0TensorPingPong; + LocalTensor bL0TensorPingPong; + LocalTensor cL0TensorPingPong; + + // L0AB m <> mte1 EventID + __aicore__ inline uint32_t Mte1MmABEventId(uint32_t idx) + { + return (L0AB_EVENT0 + idx); + } + + __aicore__ inline uint32_t GetQPL1RealIdx(uint32_t mIdx, uint32_t k1Idx) + { + uint32_t idxMap[] = {0, 2}; // 确保0块和1块连在一起, 2和3块连在一起, 来保证同一m块的地址相连 + return idxMap[mIdx % 2] + k1Idx; + } + + __aicore__ inline void CopyGmToL1(LocalTensor &l1Tensor, GlobalTensor &gmSrcTensor, uint32_t srcN, + uint32_t srcD, uint32_t srcDstride); + __aicore__ inline void CopyInMm1AToL1(LocalTensor &aL1Tensor, const RunInfo &info, uint32_t mSeqIdx, + uint32_t mSizeAct, uint32_t headSize, uint32_t headOffset); + __aicore__ inline void CopyInMm2AToL1(LocalTensor &aL1Tensor, const RunInfo &info, uint32_t mSeqIdx, + uint32_t subMSizeAct, uint32_t nSize, uint32_t nOffset); + __aicore__ inline void LoadDataMm1A(LocalTensor &aL0Tensor, LocalTensor &aL1Tensor, uint32_t idx, + uint32_t kSplitSize, uint32_t mSize, uint32_t kSize); + __aicore__ inline void LoadDataMm1B(LocalTensor &bL0Tensor, LocalTensor &bL1Tensor, uint32_t idx, + uint32_t kSplitSize, uint32_t kSize, uint32_t nSize); +}; + + +template +__aicore__ inline void SWACubeBlock::InitParams(const ConstInfo &constInfo) +{ + this->constInfo = constInfo; +} + +template +__aicore__ inline void SWACubeBlock::InitMm1GlobalTensor(GlobalTensor queryGm, GlobalTensor oriKvGm, + GlobalTensor cmpKvGm, + GlobalTensor mm1ResGm) +{ + // mm1 + this->queryGm = queryGm; + this->oriKvGm = oriKvGm; + if (constInfo.templateMode == CFA_TEMPLATE) { + this->cmpKvGm = cmpKvGm; + } + this->mm1ResGm = mm1ResGm; +} + +template +__aicore__ inline void SWACubeBlock::InitMm2GlobalTensor(GlobalTensor vec1ResGm, + GlobalTensor mm2ResGm, + GlobalTensor attentionOutGm) +{ + // mm2 + this->vec1ResGm = vec1ResGm; + this->mm2ResGm = mm2ResGm; + this->attentionOutGm = attentionOutGm; +} + +template +__aicore__ inline void +SWACubeBlock::InitPageAttentionInfo(GlobalTensor oriKvGm, // const GlobalTensor& kvMergeGm, + GlobalTensor oriBlockTableGm, GlobalTensor cmpBlockTableGm) +{ + this->oriKvGm = oriKvGm; + this->oriBlockTableGm = oriBlockTableGm; + if (constInfo.templateMode == CFA_TEMPLATE) { + this->cmpBlockTableGm = cmpBlockTableGm; + } +} + +template +__aicore__ inline void SWACubeBlock::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(bufQPL1, L1_BLOCK_SIZE * 4); + l1QPTensor = bufQPL1.Get(); + pipe->InitBuffer(bufKVL1, L1_BLOCK_SIZE * 3); + l1KVTensor = bufKVL1.Get(); + // L0A + pipe->InitBuffer(tmpBufL0A, L0A_PP_SIZE * 2); // 64K + aL0TensorPingPong = tmpBufL0A.Get(); + // L0B + pipe->InitBuffer(tmpBufL0B, L0B_PP_SIZE * 2); // 64K + bL0TensorPingPong = tmpBufL0B.Get(); + // L0C + pipe->InitBuffer(tmpBufL0C, L0C_PP_SIZE * 2); // 128K + cL0TensorPingPong = tmpBufL0C.Get(); +} + +template +__aicore__ inline void SWACubeBlock::AllocEventID() +{ + SetFlag(L1_EVENT0); + SetFlag(L1_EVENT1); + SetFlag(L1_EVENT2); + SetFlag(L1_EVENT3); + SetFlag(L1_EVENT4); + SetFlag(L1_EVENT5); + SetFlag(L1_EVENT6); + SetFlag(L0AB_EVENT0); + SetFlag(L0AB_EVENT1); +} + +template +__aicore__ inline void SWACubeBlock::FreeEventID() +{ + WaitFlag(L1_EVENT0); + WaitFlag(L1_EVENT1); + WaitFlag(L1_EVENT2); + WaitFlag(L1_EVENT3); + WaitFlag(L1_EVENT4); + WaitFlag(L1_EVENT5); + WaitFlag(L1_EVENT6); + WaitFlag(L0AB_EVENT0); + WaitFlag(L0AB_EVENT1); +} + +template +__aicore__ inline void SWACubeBlock::CopyGmToL1(LocalTensor &l1Tensor, GlobalTensor &gmSrcTensor, + uint32_t srcN, uint32_t srcD, uint32_t srcDstride) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = srcN; // 行数 + nd2nzPara.dValue = srcD; + nd2nzPara.srcDValue = srcDstride; + nd2nzPara.dstNzC0Stride = (srcN + 15) / 16 * 16; // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(l1Tensor, gmSrcTensor, nd2nzPara); +} + +template +__aicore__ inline void SWACubeBlock::CopyInMm1AToL1(LocalTensor &l1Tensor, const RunInfo &info, + uint32_t mSeqIdx, uint32_t mSizeAct, uint32_t headSize, + uint32_t headOffset) +{ + auto srcGm = queryGm[info.tensorAOffset + mSeqIdx * constInfo.headDim + headOffset]; + CopyGmToL1(l1Tensor, srcGm, mSizeAct, headSize, constInfo.headDim); +} + +template +__aicore__ inline void SWACubeBlock::LoadDataMm1A(LocalTensor &aL0Tensor, LocalTensor &aL1Tensor, + uint32_t idx, uint32_t kSplitSize, uint32_t mSize, + uint32_t kSize) +{ + LocalTensor srcTensor = aL1Tensor[mSize * kSplitSize * idx]; + LoadData3DParamsV2 loadData3DParams; + // SetFmatrixParams + loadData3DParams.l1H = mSize / 16; // Hin=M1=8 + loadData3DParams.l1W = 16; // Win=M0 + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + // SetLoadToA0Params + loadData3DParams.mExtension = mSize; // M + loadData3DParams.kExtension = kSize; // K + loadData3DParams.mStartPt = 0; + loadData3DParams.kStartPt = 0; + loadData3DParams.strideW = 1; + loadData3DParams.strideH = 1; + loadData3DParams.filterW = 1; + loadData3DParams.filterSizeW = (1 >> 8) & 255; + loadData3DParams.filterH = 1; + loadData3DParams.filterSizeH = (1 >> 8) & 255; + loadData3DParams.dilationFilterW = 1; + loadData3DParams.dilationFilterH = 1; + loadData3DParams.enTranspose = 0; + loadData3DParams.fMatrixCtrl = 0; + loadData3DParams.channelSize = kSize; // Cin=K + LoadData(aL0Tensor, srcTensor, loadData3DParams); +} + +template +__aicore__ inline void SWACubeBlock::LoadDataMm1B(LocalTensor &l0Tensor, LocalTensor &l1Tensor, + uint32_t idx, uint32_t kSplitSize, uint32_t kSize, + uint32_t nSize) +{ + // N 方向全载 + LocalTensor srcTensor = l1Tensor[nSize * kSplitSize * idx]; + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = (nSize + 15) / 16 * kSize / (32 / sizeof(KV_T)); + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = false; + LoadData(l0Tensor, srcTensor, loadData2DParams); +} + +template +__aicore__ inline void SWACubeBlock::CopyInMm2AToL1(LocalTensor &aL1Tensor, const RunInfo &info, + uint32_t mSeqIdx, uint32_t subMSizeAct, uint32_t nSize, + uint32_t nOffset) +{ + auto srcGm = vec1ResGm[(info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize + + mSeqIdx * info.actualSingleProcessSInnerSizeAlign + nOffset]; + CopyGmToL1(aL1Tensor, srcGm, subMSizeAct, nSize, info.actualSingleProcessSInnerSizeAlign); +} + +template +__aicore__ inline void SWACubeBlock::ComputeMm1(const RunInfo &info, const MSplitInfo mSplitInfo) +{ + uint32_t mSize = mSplitInfo.nBufferDealM; + uint32_t mL1Size = M_SPLIT_SIZE; + uint32_t mL1SizeAlign = SASAlign(M_SPLIT_SIZE, 16); + uint32_t mL1Loops = CeilDiv(mSize, M_SPLIT_SIZE); + + uint32_t nSize = info.actualSingleProcessSInnerSize; + uint32_t nL1Size = N_SPLIT_SIZE; + uint32_t nL1SizeAlign = SASAlign(N_SPLIT_SIZE, 16); + uint32_t nL1Loops = CeilDiv(nSize, N_SPLIT_SIZE); + + uint32_t kSize = 512; + uint32_t kL1Size = 256; + uint32_t kL1Loops = 2; + uint32_t kL0Size = 128; + uint32_t kL0Loops = CeilDiv(kL1Size, kL0Size); + + LocalTensor bL1Tensor; + LocalTensor kTensor; + uint32_t ka = 0, kb = 0; + uint32_t copyRowCnt = 0; + uint32_t copyRowCntTmp = 0; + // L1 切n切k + for (uint32_t nL1 = 0; nL1 < nL1Loops; nL1++) { // L1切n, 512/128=4 + if (nL1 == (nL1Loops - 1)) { + // 尾块重新计算size + nL1Size = nSize - (nL1Loops - 1) * N_SPLIT_SIZE; + nL1SizeAlign = SASAlign(nL1Size, 16); + } + + for (uint32_t kL1 = 0; kL1 < kL1Loops; kL1++) { + kvL1BufIter++; + uint32_t kb = kvL1BufIter % 3; + WaitFlag(mte21KVIds[kb]); + // 从k当中取当前的块 + bL1Tensor = l1KVTensor[kb * L1_BLOCK_OFFSET]; + uint32_t curSeqIdx = info.s2BatchOffset + nL1 * N_SPLIT_SIZE; + uint32_t copyFinishRowCnt = 0; + + if (info.isOri) { + if constexpr (KV_LAYOUT_T == SAS_LAYOUT::PA_ND) { + uint32_t curS2Offset = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint + nL1 * N_SPLIT_SIZE; + uint32_t copyFinishRowCnt = 0; + LocalTensor kTensor; + uint32_t copyRowCnt = 0; + + while (copyFinishRowCnt < nL1Size) { + // 由于ori_left的存在， 即使第一块搬运也可能并非是pa_block的零点位 + copyRowCnt = constInfo.paOriBlockSize - curS2Offset % constInfo.paOriBlockSize; + if (copyFinishRowCnt + copyRowCnt > nL1Size) { + copyRowCnt = nL1Size - copyFinishRowCnt; + } + Position startPos; + startPos.bIdx = info.bIdx; + startPos.n2Idx = info.n2Idx; + startPos.s2Idx = curS2Offset; + // 256、32等待7buf命名更改 + startPos.dIdx = kL1 * 256; // mm1 右矩阵 bn2s2d, d为k轴不切; mm2 右矩阵, s2为k轴, d轴切分 + PAShape shape; + shape.blockSize = constInfo.paOriBlockSize; + shape.headNum = constInfo.kvHeadNum; + shape.headDim = constInfo.headDim; + shape.kvStride = constInfo.oriKvStride; + shape.actHeadDim = 256; + shape.maxblockNumPerBatch = constInfo.oriMaxBlockNumPerBatch; + shape.copyRowNum = copyRowCnt; + shape.copyRowNumAlign = nL1SizeAlign; + kTensor = bL1Tensor[copyFinishRowCnt * 16]; + DataCopyPA(kTensor, oriKvGm, oriBlockTableGm, shape, startPos); + // 更新循环变量 + copyFinishRowCnt += copyRowCnt; + curS2Offset += copyRowCnt; + } + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::BSND) { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nL1Size; // 行数 + nd2nzPara.dValue = D_SPLIT_SIZE; // 256 + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = nL1SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + + uint32_t headStride = constInfo.headDim; + uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim; + uint32_t batchStride = constInfo.kvSeqSize * seqStride; + + uint32_t curS2 = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint; + uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + (uint64_t)info.n2Idx * headStride + kL1 * D_SPLIT_SIZE; + DataCopy(bL1Tensor, oriKvGm[offset], nd2nzPara); + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::TND) { + uint32_t curS2Offset = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint + nL1 * N_SPLIT_SIZE; + if (kL1 == 0) { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nL1Size; + nd2nzPara.dValue = constInfo.headDim >> 1; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = nL1SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor, oriKvGm[info.tensorBOffset + curS2Offset * constInfo.headDim + + nL1 * N_SPLIT_SIZE * constInfo.headDim], nd2nzPara); + } else { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nL1Size; + nd2nzPara.dValue = constInfo.headDim >> 1; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = nL1SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor, + oriKvGm[info.tensorBOffset + curS2Offset * constInfo.headDim + (constInfo.headDim >> 1) + + nL1 * N_SPLIT_SIZE * constInfo.headDim], + nd2nzPara); + } + } + } else { + if constexpr (KV_LAYOUT_T == SAS_LAYOUT::PA_ND) { + uint32_t curS2Offset = info.relativeS2Idx * constInfo.s2BaseSize + nL1 * N_SPLIT_SIZE; + while (copyFinishRowCnt < nL1Size) { + // 由于ori_left的存在， 即使第一块搬运也可能并非是pa_block的零点位 + copyRowCnt = constInfo.paCmpBlockSize - curS2Offset % constInfo.paCmpBlockSize; + if (copyFinishRowCnt + copyRowCnt > nL1Size) { + copyRowCnt = nL1Size - copyFinishRowCnt; + } + + Position startPos; + startPos.bIdx = info.bIdx; + startPos.n2Idx = info.n2Idx; + startPos.s2Idx = curS2Offset; + // 256、32等待7buf命名更改 + startPos.dIdx = kL1 * 256; // mm1 右矩阵 bn2s2d, d为k轴不切; mm2 右矩阵, s2为k轴, d轴切分 + + PAShape shape; + shape.blockSize = constInfo.paCmpBlockSize; + shape.headNum = constInfo.kvHeadNum; + shape.headDim = constInfo.headDim; + shape.kvStride = constInfo.cmpKvStride; + shape.actHeadDim = 256; + shape.maxblockNumPerBatch = constInfo.cmpMaxBlockNumPerBatch; + shape.copyRowNum = copyRowCnt; + shape.copyRowNumAlign = nL1SizeAlign; + kTensor = bL1Tensor[copyFinishRowCnt * 16]; + DataCopyPA(kTensor, cmpKvGm, cmpBlockTableGm, shape, startPos); + // 更新循环变量 + copyFinishRowCnt += copyRowCnt; + curS2Offset += copyRowCnt; + } + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::BSND) { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nL1Size; // 行数 + nd2nzPara.dValue = D_SPLIT_SIZE; // 256 + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = nL1SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + + uint32_t headStride = constInfo.headDim; + uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim; + uint32_t batchStride = constInfo.kvSeqSize / constInfo.cmpRatio * seqStride; + + uint32_t curS2 = info.relativeS2Idx * constInfo.s2BaseSize + nL1 * N_SPLIT_SIZE; + uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + (uint64_t)info.n2Idx * headStride + kL1 * D_SPLIT_SIZE; + DataCopy(bL1Tensor, cmpKvGm[offset], nd2nzPara); + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::TND) { + uint32_t curS2Offset = info.relativeS2Idx * constInfo.s2BaseSize + nL1 * N_SPLIT_SIZE; + if (kL1 == 0) { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nL1Size; + nd2nzPara.dValue = constInfo.headDim >> 1; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = nL1SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor, cmpKvGm[info.tensorCmpBOffset + curS2Offset * constInfo.headDim], nd2nzPara); + } else { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nL1Size; + nd2nzPara.dValue = constInfo.headDim >> 1; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = nL1SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor, + cmpKvGm[info.tensorCmpBOffset + curS2Offset * constInfo.headDim + (constInfo.headDim >> 1)], + nd2nzPara); + } + } + } + + SetFlag(mte21KVIds[kb]); + WaitFlag(mte21KVIds[kb]); + mL1Size = M_SPLIT_SIZE; + mL1SizeAlign = SASAlign(M_SPLIT_SIZE, 16U); + for (uint32_t mL1 = 0; mL1 < mL1Loops; mL1++) { + uint32_t aL1PaddingSize = 0; // 用于使左矩阵对齐到尾部, 以保证两块32K内存连续 + if (mL1 == (mL1Loops - 1)) { + mL1Size = mSize - (mL1Loops - 1) * M_SPLIT_SIZE; + mL1SizeAlign = SASAlign(mL1Size, 16U); + aL1PaddingSize = (M_SPLIT_SIZE - mL1SizeAlign) * 256; + } + uint32_t mIdx = qpL1BufIter + mL1; + ka = GetQPL1RealIdx(mIdx, kL1); + LocalTensor aL1Tensor = l1QPTensor[ka * L1_BLOCK_OFFSET + (1 - kL1) * aL1PaddingSize]; + if (nL1 == 0) { + if (kL1 == 0) { + WaitFlag(mte21QPIds[ka]); + WaitFlag(mte21QPIds[ka + 1]); + CopyInMm1AToL1(aL1Tensor, info, mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE, mL1Size, 256, 0); + } else { + LocalTensor qTmpTensor = aL1Tensor; + CopyInMm1AToL1(qTmpTensor, info, mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE, mL1Size, 256, + 256); + } + SetFlag(mte21QPIds[ka]); + WaitFlag(mte21QPIds[ka]); + } + // 使用unitflag同步 + LocalTensor cL0Tensor = + cL0TensorPingPong[(cL0BufIter % 2) * + (L0C_PP_SIZE / sizeof(MM_OUT_T))]; // 需要保证cL0BufIter和m步调一致 + for (uint32_t kL0 = 0; kL0 < kL0Loops; kL0++) { + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + LocalTensor aL0Tensor = aL0TensorPingPong[(abL0BufIter % 2) * (L0A_PP_SIZE / sizeof(KV_T))]; + LoadDataMm1A(aL0Tensor, aL1Tensor, kL0, kL0Size, mL1SizeAlign, kL0Size); + LocalTensor bL0Tensor = bL0TensorPingPong[(abL0BufIter % 2) * (L0B_PP_SIZE / sizeof(KV_T))]; + LoadDataMm1B(bL0Tensor, bL1Tensor, kL0, kL0Size, kL0Size, nL1SizeAlign); + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + + MmadParams mmadParams; + mmadParams.m = mL1SizeAlign; + mmadParams.n = nL1SizeAlign; + mmadParams.k = kL0Size; + mmadParams.cmatrixInitVal = (kL1 == 0 && kL0 == 0); + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = + (kL1 == 1 && kL0 == (kL0Loops - 1)) ? 0b11 : 0b10; // 累加最后一次翻转flag, 表示可以搬出 + Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + abL0BufIter++; + } + + if (nL1 == (nL1Loops - 1)) { + SetFlag(mte21QPIds[ka]); // 反向同步, 表示L1中的A已经被mte1消费完 + } + + if (kL1 == 1) { // 最后一轮kL1循环 + FixpipeParamsV220 fixParams; + fixParams.nSize = nL1SizeAlign; + fixParams.mSize = mL1SizeAlign; + fixParams.srcStride = mL1SizeAlign; + // 改成nSizeAlign + fixParams.dstStride = info.actualSingleProcessSInnerSizeAlign; // mm1ResGm两行之间的间隔 + fixParams.unitFlag = 0b11; + fixParams.ndNum = 1; // 输出ND + + Fixpipe(mm1ResGm[(info.loop % (constInfo.preLoadNum)) * constInfo.mmResUbSize + nL1 * N_SPLIT_SIZE + + (mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE) * + info.actualSingleProcessSInnerSizeAlign], + cL0Tensor, fixParams); + } + if (mL1Loops == 2) { + cL0BufIter++; + } + } + + SetFlag(mte21KVIds[kb]); // 反向同步, 表示L1已经被mte1消费完 + } + if (mL1Loops == 1) { + cL0BufIter++; + } + } + qpL1BufIter += mL1Loops; +} + + +template +__aicore__ inline void SWACubeBlock::ComputeMm2(const RunInfo &info, const MSplitInfo mSplitInfo) +{ + uint32_t mSize = mSplitInfo.nBufferDealM; + uint32_t mSizeAlign = (mSize + 16 - 1) / 16; + uint32_t mL1Loops = (mSize + M_SPLIT_SIZE - 1) / M_SPLIT_SIZE; + uint32_t mL1SizeAlign = M_SPLIT_SIZE; // 16对齐 + uint32_t mL1Size = M_SPLIT_SIZE; // m的实际大小 + + uint32_t nSize = BlockAlign(constInfo.headDim); + uint32_t nL1Loops = (nSize + N_SPLIT_SIZE - 1) / N_SPLIT_SIZE; + uint32_t nL1SizeAlign = N_SPLIT_SIZE; // 16对齐 + uint32_t nL1Size = N_SPLIT_SIZE; // n的实际大小 + + uint32_t kSize = info.actualSingleProcessSInnerSize; + uint32_t kL1Size = 256; + uint32_t kL1SizeAlign = SASAlign(kL1Size, 16U); + uint32_t kL1Loops = (kSize + kL1Size - 1) / kL1Size; + uint32_t kL0Size = 128; + uint32_t kL0Loops = (kL1Size + kL0Size - 1) / kL0Size; + uint32_t kL0SizeAlign = kL0Size; + LocalTensor bL1Tensor; + LocalTensor subvTensor; + // ka表示左矩阵4buf选择哪一块buf, kb表示右矩阵3buf选择哪一块buf + uint32_t ka = 0, kb = 0; + uint32_t mBaseIdx = qpL1BufIter; + for (uint32_t nL1 = 0; nL1 < nL1Loops; nL1++) { // n切L1 + if (nL1 == (nL1Loops - 1)) { + // 尾块 + nL1Size = nSize - (nL1Loops - 1) * N_SPLIT_SIZE; + nL1SizeAlign = SASAlign(nL1Size, 16U); + } + // k l1写成一个循环, 和mm1保持一致 + kL1Size = 256; + kL1SizeAlign = SASAlign(kL1Size, 16U); + uint32_t copyRowCnt = 0; + for (uint32_t k1 = 0; k1 < kL1Loops; k1++) { // k切L1, 这里套了一层l0来操作 + if (k1 == (kL1Loops - 1)) { + // 尾块 + kL1Size = kSize - (kL1Loops - 1) * K_L1_SPLIT_SIZE; + kL1SizeAlign = SASAlign(kL1Size, 16U); + } + kvL1BufIter++; + uint32_t kb = kvL1BufIter % 3; + WaitFlag(mte21KVIds[kb]); + bL1Tensor = l1KVTensor[kb * L1_BLOCK_OFFSET]; + uint32_t kOffset = k1 * kL0Loops; + kL0Size = 128; + // 此处必须先初始化kL0Size, 再求kL0Loops, 否则由于循环会改变kL0Size大小, 导致kL0Loops错误 + kL0Loops = (kL1Size + kL0Size - 1) / kL0Size; + kL0SizeAlign = kL0Size; + for (uint32_t kL1 = kOffset; kL1 < kL0Loops + kOffset; kL1++) { // 128 循环搬pa + if (kL1 == kOffset + kL0Loops - 1) { + // 尾块 + kL0Size = kL1Size - (kL0Loops - 1) * kL0Size; + kL0SizeAlign = SASAlign(kL0Size, 16U); + } + + uint32_t copyFinishRowCnt = 0; + + if (info.isOri) { + if constexpr (KV_LAYOUT_T == SAS_LAYOUT::PA_ND) { + uint32_t curS2Offset = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint + kL1 * K_L0_SPLIT_SIZE; + while (copyFinishRowCnt < kL0Size) { + copyRowCnt = constInfo.paOriBlockSize - curS2Offset % constInfo.paOriBlockSize; + if (copyFinishRowCnt + copyRowCnt > kL0Size) { + copyRowCnt = kL0Size - copyFinishRowCnt; + } + Position startPos; + startPos.bIdx = info.bIdx; + startPos.n2Idx = info.n2Idx; + startPos.s2Idx = curS2Offset; + startPos.dIdx = nL1 * N_SPLIT_SIZE; // mm1 右矩阵 bn2s2d, d为k轴不切; mm2 右矩阵, s2为k轴, d轴切分 + PAShape shape; + shape.blockSize = constInfo.paOriBlockSize; + shape.headNum = constInfo.kvHeadNum; + shape.headDim = constInfo.headDim; + shape.kvStride = constInfo.oriKvStride; + shape.actHeadDim = nL1Size; + shape.maxblockNumPerBatch = constInfo.oriMaxBlockNumPerBatch; + shape.copyRowNum = copyRowCnt; + shape.copyRowNumAlign = kL0SizeAlign; + subvTensor = bL1Tensor[(kL1 - kOffset) * K_L0_SPLIT_SIZE * N_SPLIT_SIZE + copyFinishRowCnt * 16]; + DataCopyPA(subvTensor, oriKvGm, oriBlockTableGm, shape, startPos); + + // 更新循环变量 + copyFinishRowCnt += copyRowCnt; + curS2Offset += copyRowCnt; + } + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::BSND) { + subvTensor = bL1Tensor[(kL1 - kOffset) * K_L0_SPLIT_SIZE * N_SPLIT_SIZE]; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = kL0Size; // 行数 + nd2nzPara.dValue = nL1Size; // 256 + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = kL0SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + + uint32_t headStride = constInfo.headDim; + uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim; + uint32_t batchStride = constInfo.kvSeqSize * seqStride; + + uint32_t curS2 = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint + kL1 * K_L0_SPLIT_SIZE; + uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + (uint64_t)info.n2Idx * headStride + nL1 * N_SPLIT_SIZE; + DataCopy(subvTensor, oriKvGm[offset], nd2nzPara); + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::TND) { + uint32_t curS2Offset = info.s2Idx * constInfo.s2BaseSize + info.s2StartPoint + kL1 * K_L0_SPLIT_SIZE; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = kL0Size; // 行数 + nd2nzPara.dValue = N_SPLIT_SIZE; // constInfo.headDim; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = kL0SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor[(kL1 - kOffset) * K_L0_SPLIT_SIZE * N_SPLIT_SIZE], + oriKvGm[info.tensorBOffset + curS2Offset * constInfo.headDim + + nL1 * N_SPLIT_SIZE], nd2nzPara); + } + } else { + if constexpr (KV_LAYOUT_T == SAS_LAYOUT::PA_ND) { + uint32_t curS2Offset = info.relativeS2Idx * constInfo.s2BaseSize + K_L0_SPLIT_SIZE * kL1; + while (copyFinishRowCnt < kL0Size) { + copyRowCnt = constInfo.paCmpBlockSize - curS2Offset % constInfo.paCmpBlockSize; + if (copyFinishRowCnt + copyRowCnt > kL0Size) { + copyRowCnt = kL0Size - copyFinishRowCnt; + } + + Position startPos; + startPos.bIdx = info.bIdx; + startPos.n2Idx = info.n2Idx; + startPos.s2Idx = curS2Offset; + // 256、32等待7buf命名更改 + startPos.dIdx = nL1 * N_SPLIT_SIZE; // mm1 右矩阵 bn2s2d, d为k轴不切; mm2 右矩阵, s2为k轴, d轴切分 + + PAShape shape; + shape.blockSize = constInfo.paCmpBlockSize; + shape.headNum = constInfo.kvHeadNum; + shape.headDim = constInfo.headDim; + shape.kvStride = constInfo.cmpKvStride; + shape.actHeadDim = nL1Size; + shape.maxblockNumPerBatch = constInfo.cmpMaxBlockNumPerBatch; + shape.copyRowNum = copyRowCnt; + shape.copyRowNumAlign = kL0SizeAlign; + subvTensor = bL1Tensor[(kL1 - kOffset) * K_L0_SPLIT_SIZE * N_SPLIT_SIZE + copyFinishRowCnt * 16]; + DataCopyPA(subvTensor, cmpKvGm, cmpBlockTableGm, shape, startPos); + // 更新循环变量 + copyFinishRowCnt += copyRowCnt; + curS2Offset += copyRowCnt; + } + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::BSND) { + subvTensor = bL1Tensor[(kL1 - kOffset) * K_L0_SPLIT_SIZE * N_SPLIT_SIZE]; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = kL0Size; // 行数 + nd2nzPara.dValue = nL1Size; // 256 + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = kL0SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + + uint32_t headStride = constInfo.headDim; + uint32_t seqStride = constInfo.kvHeadNum * constInfo.headDim; + uint32_t batchStride = constInfo.kvSeqSize / constInfo.cmpRatio * seqStride; + + uint32_t curS2 = info.relativeS2Idx * constInfo.s2BaseSize + K_L0_SPLIT_SIZE * kL1; + uint64_t offset = (uint64_t)info.bIdx * batchStride + (uint64_t)curS2 * seqStride + (uint64_t)info.n2Idx * headStride + nL1 * N_SPLIT_SIZE; + DataCopy(subvTensor, cmpKvGm[offset], nd2nzPara); + } else if constexpr (KV_LAYOUT_T == SAS_LAYOUT::TND) { + uint32_t curS2Offset = info.relativeS2Idx * constInfo.s2BaseSize + K_L0_SPLIT_SIZE * kL1; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = kL0Size; // 行数 + nd2nzPara.dValue = N_SPLIT_SIZE; // constInfo.headDim; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = kL0SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor[(kL1 - kOffset) * K_L0_SPLIT_SIZE * N_SPLIT_SIZE], + cmpKvGm[info.tensorCmpBOffset + curS2Offset * constInfo.headDim + + nL1 * N_SPLIT_SIZE], nd2nzPara); + } + } + } + SetFlag(mte21KVIds[kb]); + WaitFlag(mte21KVIds[kb]); + mL1SizeAlign = M_SPLIT_SIZE; + mL1Size = M_SPLIT_SIZE; // m的实际大小 + for (uint32_t mL1 = 0; mL1 < mL1Loops; mL1++) { + if (mL1 == (mL1Loops - 1)) { + // 尾块 + mL1Size = mSize - (mL1Loops - 1) * M_SPLIT_SIZE; + mL1SizeAlign = SASAlign(mL1Size, 16U); + } + + uint32_t mIdx = mBaseIdx + mL1; + ka = GetQPL1RealIdx(mIdx, k1); + LocalTensor aL1Tensor = l1QPTensor[ka * L1_BLOCK_OFFSET]; + if (nL1 == 0) { + WaitFlag(mte21QPIds[ka]); + CopyInMm2AToL1(aL1Tensor, info, mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE, mL1Size, kL1Size, + 256 * k1); + SetFlag(mte21QPIds[ka]); + WaitFlag(mte21QPIds[ka]); + } + + LocalTensor cL0Tensor = + cL0TensorPingPong[(cL0BufIter % 2) * + (L0C_PP_SIZE / sizeof(MM_OUT_T))]; // 需要保证cL0BufIter和m步调一致 + uint32_t baseK = 128; + uint32_t baseN = 128; + kL0Size = 128; + kL0SizeAlign = kL0Size; + for (uint32_t kL0 = 0; kL0 < kL0Loops; kL0++) { + if (kL0 + 1 == kL0Loops) { + kL0Size = kL1Size - (kL0Loops - 1) * kL0Size; + kL0SizeAlign = SASAlign(kL0Size, 16U); + } + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + LocalTensor bL0Tensor = bL0TensorPingPong[(abL0BufIter % 2) * (L0B_PP_SIZE / sizeof(KV_T))]; + LoadData3DParamsV2 loadData3DParamsForB; + loadData3DParamsForB.l1H = kL0SizeAlign / 16; // 源操作数height + loadData3DParamsForB.l1W = 16; // 源操作数weight=16，目的height=l1H*L1W + loadData3DParamsForB.padList[0] = 0; + loadData3DParamsForB.padList[1] = 0; + loadData3DParamsForB.padList[2] = 0; + loadData3DParamsForB.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + loadData3DParamsForB.mExtension = kL0SizeAlign; // 在目的操作数height维度的传输长度 + loadData3DParamsForB.kExtension = nL1SizeAlign; // 在目的操作数width维度的传输长度 + loadData3DParamsForB.mStartPt = 0; // 卷积核在目的操作数width维度的起点 + loadData3DParamsForB.kStartPt = 0; // 卷积核在目的操作数height维度的起点 + loadData3DParamsForB.strideW = 1; + loadData3DParamsForB.strideH = 1; + loadData3DParamsForB.filterW = 1; + loadData3DParamsForB.filterSizeW = false; // 是否在filterW的基础上将卷积核width增加256个元素 + loadData3DParamsForB.filterH = 1; + loadData3DParamsForB.filterSizeH = false; // 是否在filterH的基础上将卷积核height增加256个元素 + loadData3DParamsForB.dilationFilterW = 1; // 卷积核width膨胀系数 + loadData3DParamsForB.dilationFilterH = 1; // 卷积核height膨胀系数 + loadData3DParamsForB.enTranspose = 1; // 是否启用转置功能 + loadData3DParamsForB.fMatrixCtrl = + 0; // 使用FMATRIX_LEFT还是使用FMATRIX_RIGHT，=0使用FMATRIX_LEFT，=1使用FMATRIX_RIGHT 1 + loadData3DParamsForB.channelSize = + nL1SizeAlign; // 源操作数的通道数。膨胀系数为1时，目的weight为filterW*filterH*channelSize + LoadData(bL0Tensor, bL1Tensor[kL0 * baseK * baseN], loadData3DParamsForB); + + LocalTensor aL0Tensor = aL0TensorPingPong[(abL0BufIter % 2) * (L0A_PP_SIZE / sizeof(KV_T))]; + LoadData3DParamsV2 loadData3DParamsForA; + loadData3DParamsForA.l1H = mL1SizeAlign / 16; // 源操作数height + loadData3DParamsForA.l1W = 16; // 源操作数weight + loadData3DParamsForA.padList[0] = 0; + loadData3DParamsForA.padList[1] = 0; + loadData3DParamsForA.padList[2] = 0; + loadData3DParamsForA.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + loadData3DParamsForA.mExtension = mL1SizeAlign; // 在目的操作数height维度的传输长度 + loadData3DParamsForA.kExtension = kL0SizeAlign; // 在目的操作数width维度的传输长度 + loadData3DParamsForA.mStartPt = 0; // 卷积核在目的操作数width维度的起点 + loadData3DParamsForA.kStartPt = 0; // 卷积核在目的操作数height维度的起点 + loadData3DParamsForA.strideW = 1; // 卷积核在源操作数width维度滑动的步长 + loadData3DParamsForA.strideH = 1; // 卷积核在源操作数height维度滑动的步长 + loadData3DParamsForA.filterW = 1; // 卷积核width + loadData3DParamsForA.filterSizeW = false; // 是否在filterW的基础上将卷积核width增加256个元素 + loadData3DParamsForA.filterH = 1; // 卷积核height + loadData3DParamsForA.filterSizeH = false; // 是否在filterH的基础上将卷积核height增加256个元素 + loadData3DParamsForA.dilationFilterW = 1; // 卷积核width膨胀系数 + loadData3DParamsForA.dilationFilterH = 1; // 卷积核height膨胀系数 + loadData3DParamsForA.enTranspose = 0; // 是否启用转置功能，对整个目标矩阵进行转置 + loadData3DParamsForA.fMatrixCtrl = 0; + loadData3DParamsForA.channelSize = + kL0SizeAlign; // 源操作数的通道数。膨胀系数为1时，目的weight为filterW*filterH*channelSize + LoadData(aL0Tensor, aL1Tensor[kL0 * baseK * mL1SizeAlign], + loadData3DParamsForA); + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + + MmadParams mmadParams; + mmadParams.m = mL1SizeAlign; + mmadParams.n = nL1SizeAlign; + mmadParams.k = kL0Size; + mmadParams.cmatrixInitVal = (kL0 == 0 && k1 == 0); + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = ((k1 == (kL1Loops - 1)) && (kL0 == (kL0Loops - 1))) ? 0b11 : 0b10; + + Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + abL0BufIter++; + } + + if (nL1 == (nL1Loops - 1)) { // nL1最后一轮, 需要将B驻留在L1中, 用于下一轮的计算？ + SetFlag(mte21QPIds[ka]); // 反向同步, 表示L1中的A已经被mte1消费完 + } + + if (k1 == (kL1Loops - 1)) { + // ND + FixpipeParamsV220 fixParams; + fixParams.nSize = nL1SizeAlign; + fixParams.mSize = mL1SizeAlign; + fixParams.srcStride = mL1SizeAlign; + fixParams.dstStride = nSize; // mm2ResGm两行之间的间隔 + fixParams.ndNum = 1; // 输出ND + fixParams.unitFlag = 0b11; + + uint64_t mm2Offset = (mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE) * nSize + nL1 * N_SPLIT_SIZE; + Fixpipe(mm2ResGm[(info.loop % (constInfo.preLoadNum)) * constInfo.bmm2ResUbSize + mm2Offset], + cL0Tensor, fixParams); + } + + if (mL1Loops == 2) { + cL0BufIter++; + } + } + SetFlag(mte21KVIds[kb]); // 反向同步, 表示L1已经被mte1消费完 + } + // cL0BufIter已经不在使用 + if (mL1Loops == 1) { + cL0BufIter++; + } + } + qpL1BufIter += mL1Loops; +} +} // namespace SASKernel +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_swa_block_vector.h b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_swa_block_vector.h new file mode 100644 index 000000000..4bbc16b6d --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_swa_block_vector.h @@ -0,0 +1,850 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_swa_block_vector.h + * \brief + */ +#ifndef SPARSE_ATTN_SHAREDKV_SWA_BLOCK_VECTOR_H +#define SPARSE_ATTN_SHAREDKV_SWA_BLOCK_VECTOR_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../sparse_attn_sharedkv_common.h" + +namespace SASKernel { +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +template +class SWAVectorBlock { +public: + // 中间计算数据类型为float，高精度模式 + using T = float; + using KV_T = typename SAST::kvType; + using OUT_T = typename SAST::outputType; + using UPDATE_T = T; + using SINKS_T = T; + using MM1_OUT_T = float; + using MM2_OUT_T = float; + + __aicore__ inline SWAVectorBlock(){}; + __aicore__ inline void ProcessVec1L(const RunInfo &info); + __aicore__ inline void ProcessVec2L(const RunInfo &info); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitParams(const struct ConstInfo &constInfo, + const SparseAttnSharedkvTilingData *__restrict tilingData); + __aicore__ inline void InitVec1GlobalTensor(GlobalTensor mm1ResGm, GlobalTensor vec1ResGm, + GlobalTensor actualSeqLengthsQGm, + GlobalTensor actualSeqLengthsKVGm, GlobalTensor sinksGm, GlobalTensor softmaxLseGm); + __aicore__ inline void InitVec2GlobalTensor(GlobalTensor accumOutGm, GlobalTensor vec2ResGm, + GlobalTensor mm2ResGm, GlobalTensor attentionOutGm); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void CopySinksIn(); + __aicore__ inline void SliceAndContactSinksValue(uint32_t nIdx, uint32_t dealRowCount); + __aicore__ inline void InitSoftmaxDefaultBuffer(); + // ================================Base Vector========================================== + __aicore__ inline void RowDivs(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline void RowMuls(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + // ================================Vector1========================================== + __aicore__ inline void ProcessVec1SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo); + __aicore__ inline void DealBmm1ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t startRow, + uint32_t dealRowCount, uint32_t columnCount, uint32_t loopId); + __aicore__ inline void SoftmaxFlashV2Compute(const RunInfo &info, const MSplitInfo &mSplitInfo, + LocalTensor &mmResUb, LocalTensor &softmaxTmpUb, + uint32_t startRow, uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + __aicore__ inline void ElewiseCompute(const RunInfo &info, const LocalTensor &mmResUb, uint32_t dealRowCount, + uint32_t columnCount); + __aicore__ inline void ProcessLse(const RunInfo &info, const MSplitInfo &mSplitInfo); + // ================================Vecotr2========================================== + __aicore__ inline void ProcessVec2SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo); + __aicore__ inline void DealBmm2ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t startRow, + uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + __aicore__ inline void ProcessVec2Inner(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t mStartRow, + uint32_t mDealSize); + __aicore__ inline void Bmm2DataCopyOutTrans(const RunInfo &info, LocalTensor &attenOutUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + __aicore__ inline void Bmm2ResCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline void Bmm2CastAndCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline void Bmm2FDDataCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline uint64_t CalcAccumOffset(uint32_t bN2Idx, uint32_t gS1Idx); + + // BLOCK和REPEAT的字节数 + static constexpr uint64_t BYTE_BLOCK = 32UL; + static constexpr uint32_t REPEAT_BLOCK_BYTE = 256U; + // BLOCK和REPEAT的FP32元素数 + static constexpr uint32_t FP32_BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(float); + static constexpr uint32_t FP32_REPEAT_ELEMENT_NUM = REPEAT_BLOCK_BYTE / sizeof(float); + // repeat stride不能超过256 + static constexpr uint32_t REPEATE_STRIDE_UP_BOUND = 256; + +private: + static constexpr bool PAGE_ATTENTION = SAST::pageAttention; + static constexpr bool FLASH_DECODE = SAST::flashDecode; + static constexpr SAS_LAYOUT LAYOUT_T = SAST::layout; + static constexpr SAS_LAYOUT KV_LAYOUT_T = SAST::kvLayout; + + static constexpr uint64_t SYNC_INPUT_BUF1_FLAG = 2; + static constexpr uint64_t SYNC_INPUT_BUF1_PONG_FLAG = 3; + static constexpr uint64_t SYNC_INPUT_BUF2_FLAG = 4; + static constexpr uint64_t SYNC_INPUT_BUF2_PONG_FLAG = 5; + static constexpr uint64_t SYNC_OUTPUT_BUF1_FLAG = 4; + static constexpr uint64_t SYNC_OUTPUT_BUF2_FLAG = 5; + static constexpr uint64_t SYNC_SINKS_BUF_FLAG = 6; + static constexpr uint32_t INPUT1_BUFFER_OFFSET = ConstInfo::BUFFER_SIZE_BYTE_32K; + static constexpr uint32_t SOFTMAX_TMP_BUFFER_OFFSET = ConstInfo::BUFFER_SIZE_BYTE_1K; + static constexpr uint32_t BASE_BLOCK_MAX_ELEMENT_NUM = ConstInfo::BUFFER_SIZE_BYTE_32K / sizeof(T); // 32768/4=8096 + static constexpr uint32_t BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(T); // 32/4=8 + static constexpr uint32_t MAX_N1_SIZE = 128U; + static constexpr T SOFTMAX_MIN_NUM = -2e38; + static constexpr SINKS_T R0 = 1.0f; + + const SparseAttnSharedkvTilingData *__restrict tilingData; + + uint32_t pingpongFlag = 0U; + ConstInfo constInfo = {}; + + GlobalTensor mm1ResGm; + GlobalTensor vec1ResGm; + GlobalTensor softmaxMaxGm; + GlobalTensor softmaxSumGm; + GlobalTensor sinksGm; + + GlobalTensor actualSeqLengthsQGm; + GlobalTensor actualSeqLengthsKVGm; + GlobalTensor vec2ResGm; + GlobalTensor mm2ResGm; + GlobalTensor accumOutGm; + GlobalTensor attentionOutGm; + GlobalTensor blkTableGm_; + GlobalTensor keyGm_; + GlobalTensor kvValidSizeGm_; + GlobalTensor oriKvGm_; + GlobalTensor cmpKvGm_; + GlobalTensor oriBlockTableGm_; + GlobalTensor cmpBlockTableGm_; + GlobalTensor softmaxLseGm; + + // ================================Local Buffer区==================================== + TBuf<> inputBuff1; // 32K + TBuf<> inputBuff2; // 16K + TBuf<> outputBuff1; // 32K + TBuf<> outputBuff2; // 4K + + TBuf<> tmpBuff1; // 32K + TBuf<> v0ValidSizeBuff; // 8K + + TBuf<> sinksBuff; // 1K + TBuf<> sinksBrcbBuff; // 12K + + TBuf<> softmaxMaxBuff; // PRE_LOAD_NUM * 2K + TBuf<> softmaxExpBuff; // PRE_LOAD_NUM * 2K + TBuf<> softmaxSumBuff; // PRE_LOAD_NUM * 2K + TBuf<> softmaxMaxDefaultBuff; // 2K + TBuf<> softmaxSumDefaultBuff; // 2K + + LocalTensor softmaxMaxDefaultUb; + LocalTensor softmaxSumDefaultUb; + + LocalTensor softmaxMaxUb; + LocalTensor softmaxSumUb; + LocalTensor softmaxExpUb; + LocalTensor sinksUb; + LocalTensor sinksBrcbUb; +}; + +// ============================== init ============================================== +template +__aicore__ inline void SWAVectorBlock::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(inputBuff1, ConstInfo::BUFFER_SIZE_BYTE_32K * 2); // 2:pingpong + pipe->InitBuffer(inputBuff2, ConstInfo::BUFFER_SIZE_BYTE_8K * 2); // 2:pingpong + pipe->InitBuffer(outputBuff1, ConstInfo::BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(outputBuff2, ConstInfo::BUFFER_SIZE_BYTE_4K); + + pipe->InitBuffer(tmpBuff1, ConstInfo::BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(v0ValidSizeBuff, ConstInfo::BUFFER_SIZE_BYTE_8K); + + // M_MAX = 512/2vector = 256, 256 * sizeof(T) * N_Buffer + pipe->InitBuffer(softmaxMaxBuff, ConstInfo::BUFFER_SIZE_BYTE_1K * constInfo.preLoadNum); + pipe->InitBuffer(softmaxExpBuff, ConstInfo::BUFFER_SIZE_BYTE_1K * constInfo.preLoadNum); + pipe->InitBuffer(softmaxSumBuff, ConstInfo::BUFFER_SIZE_BYTE_1K * constInfo.preLoadNum); + + pipe->InitBuffer(softmaxMaxDefaultBuff, ConstInfo::BUFFER_SIZE_BYTE_1K); + pipe->InitBuffer(softmaxSumDefaultBuff, ConstInfo::BUFFER_SIZE_BYTE_1K); + + pipe->InitBuffer(sinksBuff, MAX_N1_SIZE * sizeof(SINKS_T)); + // 分配256+N1大小内存，其中256是m轴VEC最大切块 + pipe->InitBuffer(sinksBrcbBuff, MAX_N1_SIZE * sizeof(SINKS_T) * BLOCK_ELEMENT_NUM * 3U); + + softmaxMaxUb = softmaxMaxBuff.Get(); + softmaxSumUb = softmaxSumBuff.Get(); + softmaxExpUb = softmaxExpBuff.Get(); + + softmaxMaxDefaultUb = softmaxMaxDefaultBuff.Get(); + softmaxSumDefaultUb = softmaxSumDefaultBuff.Get(); + + sinksUb = sinksBuff.Get(); + sinksBrcbUb = sinksBrcbBuff.Get(); +} + + +template +__aicore__ inline void SWAVectorBlock::InitParams(const struct ConstInfo &constInfo, + const SparseAttnSharedkvTilingData *__restrict tilingData) +{ + this->constInfo = constInfo; + this->tilingData = tilingData; +} + +template +__aicore__ inline void +SWAVectorBlock::InitVec1GlobalTensor(GlobalTensor mm1ResGm, GlobalTensor vec1ResGm, + GlobalTensor actualSeqLengthsQGm, + GlobalTensor actualSeqLengthsKVGm, + GlobalTensor sinksGm, GlobalTensor softmaxLseGm) +{ + this->mm1ResGm = mm1ResGm; + this->vec1ResGm = vec1ResGm; + this->actualSeqLengthsQGm = actualSeqLengthsQGm; + this->actualSeqLengthsKVGm = actualSeqLengthsKVGm; + this->sinksGm = sinksGm; + this->softmaxLseGm = softmaxLseGm; +} + +template +__aicore__ inline void +SWAVectorBlock::InitVec2GlobalTensor(GlobalTensor accumOutGm, GlobalTensor vec2ResGm, + GlobalTensor mm2ResGm, GlobalTensor attentionOutGm) +{ + this->accumOutGm = accumOutGm; + this->vec2ResGm = vec2ResGm; + this->mm2ResGm = mm2ResGm; + this->attentionOutGm = attentionOutGm; +} + +template +__aicore__ inline void SWAVectorBlock::AllocEventID() +{ + SetFlag(SYNC_INPUT_BUF1_FLAG); + SetFlag(SYNC_INPUT_BUF1_PONG_FLAG); + SetFlag(SYNC_INPUT_BUF2_FLAG); + SetFlag(SYNC_INPUT_BUF2_PONG_FLAG); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); +} + +template +__aicore__ inline void SWAVectorBlock::FreeEventID() +{ + WaitFlag(SYNC_INPUT_BUF1_FLAG); + WaitFlag(SYNC_INPUT_BUF1_PONG_FLAG); + WaitFlag(SYNC_INPUT_BUF2_FLAG); + WaitFlag(SYNC_INPUT_BUF2_PONG_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); +} + +template +__aicore__ inline void SWAVectorBlock::CopySinksIn() +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = 1U; + dataCopyParams.blockLen = constInfo.qHeadNum * sizeof(T); + dataCopyParams.srcStride = 0U; + dataCopyParams.dstStride = 0U; + DataCopyPadExtParams padParams; + DataCopyPad(sinksUb, sinksGm, dataCopyParams, padParams); + SetFlag(SYNC_SINKS_BUF_FLAG); + WaitFlag(SYNC_SINKS_BUF_FLAG); + uint32_t repeatTimes = (constInfo.qHeadNum + BLOCK_ELEMENT_NUM - 1U) / BLOCK_ELEMENT_NUM; // 每次处理 8 datablocks + Brcb(sinksBrcbUb, sinksUb, repeatTimes, {1, BLOCK_ELEMENT_NUM}); + PipeBarrier(); + + DataCopyParams repeatParams; + repeatParams.blockCount = 1; // 搬到有一个块超过单个vec核减分核M轴大小即可，核间切分每个vec256 + repeatParams.blockLen = constInfo.qHeadNum; + repeatParams.srcStride = 0U; + repeatParams.dstStride = 0U; + for (uint32_t i = 1U; i <= 256U / constInfo.qHeadNum; i++) { + DataCopy(sinksBrcbUb[constInfo.qHeadNum * BLOCK_ELEMENT_NUM * i], sinksBrcbUb, repeatParams); + } + PipeBarrier(); +} + +template +__aicore__ inline void SWAVectorBlock::SliceAndContactSinksValue(uint32_t nIdx, uint32_t dealRowCount) +{ + // 由于WholeReduceMax接口中repeatTimes支持范围（0,255），因此需要分多次调用WholeReduceMax，这里就使用每次repeatTime=128 + uint32_t repeatTimesOnce = 128; + uint32_t loopTimes = (dealRowCount + repeatTimesOnce - 1) / repeatTimesOnce; + uint32_t repeatTimes = repeatTimesOnce; + + for (uint32_t loop = 0; loop < loopTimes; ++loop) { + if (loop == loopTimes - 1) { + repeatTimes = dealRowCount - loop * repeatTimesOnce; + } + WholeReduceMax(softmaxMaxDefaultUb[loop * repeatTimesOnce], + sinksBrcbUb[(nIdx + loop * repeatTimesOnce) * BLOCK_ELEMENT_NUM], + BLOCK_ELEMENT_NUM * BLOCK_ELEMENT_NUM, repeatTimes, 1, 0, 1, ReduceOrder::ORDER_ONLY_VALUE); + PipeBarrier(); + } +} + +template +__aicore__ inline void SWAVectorBlock::InitSoftmaxDefaultBuffer() +{ + CopySinksIn(); + Duplicate(softmaxMaxDefaultUb, SOFTMAX_MIN_NUM, SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)); + Duplicate(softmaxSumDefaultUb, R0, SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)); +} + +template +__aicore__ inline void SWAVectorBlock::ElewiseCompute(const RunInfo &info, const LocalTensor &mmResUb, + uint32_t dealRowCount, uint32_t columnCount) +{ + Muls(mmResUb, mmResUb, static_cast(tilingData->baseParams.softmaxScale), dealRowCount * columnCount); +} + + +template +__aicore__ inline void +SWAVectorBlock::SoftmaxFlashV2Compute(const RunInfo &info, const MSplitInfo &mSplitInfo, LocalTensor &mmResUb, + LocalTensor &softmaxTmpUb, uint32_t startRow, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount) +{ + LocalTensor inSumTensor; + LocalTensor inMaxTensor; + uint32_t baseOffset = mSplitInfo.nBufferStartM / 2 + startRow; + uint32_t outIdx = info.loop % (constInfo.preLoadNum); + uint32_t softmaxOutOffset = outIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset; + if (info.isFirstSInnerLoop) { + inMaxTensor = softmaxMaxDefaultUb[startRow]; + inSumTensor = softmaxSumDefaultUb; + } else { + uint32_t inIdx = (info.loop - 1) % (constInfo.preLoadNum); + inMaxTensor = softmaxMaxUb[inIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset]; + inSumTensor = softmaxSumUb[inIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset]; + } + if (actualColumnCount != 0) { + SoftMaxShapeInfo srcShape{dealRowCount, columnCount, dealRowCount, actualColumnCount}; + SoftMaxTiling newTiling = + SoftMaxFlashV2TilingFunc(srcShape, sizeof(T), sizeof(T), softmaxTmpUb.GetSize(), true, false); + SoftmaxFlashV2( + mmResUb, softmaxSumUb[softmaxOutOffset], softmaxMaxUb[softmaxOutOffset], mmResUb, + softmaxExpUb[softmaxOutOffset], inSumTensor, inMaxTensor, softmaxTmpUb, newTiling, srcShape); + } else { + uint32_t dealRowCountAlign = SASAlign(dealRowCount, FP32_BLOCK_ELEMENT_NUM); + DataCopy(softmaxSumUb[softmaxOutOffset], inSumTensor, dealRowCountAlign); + PipeBarrier(); + DataCopy(softmaxMaxUb[softmaxOutOffset], inMaxTensor, dealRowCountAlign); + } +} + +template +__aicore__ inline void SWAVectorBlock::ProcessLse(const RunInfo &info, const MSplitInfo &mSplitInfo) +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + uint64_t lseOffset; + if (constInfo.outputLayout == SAS_LAYOUT::TND) { + uint32_t tBase = actualSeqLengthsQGm.GetValue(info.bIdx); + lseOffset = (tBase + info.s1Idx) * constInfo.gSize + // T轴、s1轴偏移 + info.n2IdxReal * constInfo.qSeqSize * constInfo.gSize; // N2轴偏移 + } else if (constInfo.outputLayout == SAS_LAYOUT::BSND) { + lseOffset = info.bIdx * constInfo.qSeqSize * constInfo.kvHeadNum * constInfo.gSize + // B轴偏移 + info.n2IdxReal * constInfo.qSeqSize * constInfo.gSize + // N2轴偏移 + info.s1Idx * constInfo.gSize; // S1轴偏移 + } + lseOffset = lseOffset + mSplitInfo.nBufferStartM + mSplitInfo.vecStartM; + uint32_t baseOffset = mSplitInfo.nBufferStartM / 2; + uint32_t outIdx = info.loop % (constInfo.preLoadNum); + uint32_t softmaxOffset = outIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset; + auto sumTensor = softmaxSumUb[softmaxOffset]; + auto maxTensor = softmaxMaxUb[softmaxOffset]; + auto outLSETensor = outputBuff2.Get(); + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = 1; + dataCopyParams.blockLen = mSplitInfo.vecDealM * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = 0; + + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + PipeBarrier(); + Log(outLSETensor, sumTensor, mSplitInfo.vecDealM); + PipeBarrier(); + Add(outLSETensor, outLSETensor, maxTensor, mSplitInfo.vecDealM); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + + DataCopyPad(softmaxLseGm[lseOffset], outLSETensor, dataCopyParams); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); +} + + +template +__aicore__ inline void SWAVectorBlock::DealBmm1ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, + uint32_t startRow, uint32_t dealRowCount, + uint32_t columnCount, uint32_t loopId) +{ + uint32_t computeSize = dealRowCount * columnCount; + uint64_t inOutGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize + + (mSplitInfo.nBufferStartM + mSplitInfo.vecStartM + startRow) * columnCount; + LocalTensor mmResUb = inputBuff1.Get(); + mmResUb = mmResUb[pingpongFlag * INPUT1_BUFFER_OFFSET / sizeof(MM1_OUT_T)]; + WaitFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + + DataCopy(mmResUb, mm1ResGm[inOutGmOffset], computeSize); + SetFlag(SYNC_INPUT_BUF1_FLAG); + WaitFlag(SYNC_INPUT_BUF1_FLAG); + + ElewiseCompute(info, mmResUb, dealRowCount, columnCount); + + PipeBarrier(); + LocalTensor tmpAFloorUb = tmpBuff1.Get(); + LocalTensor softmaxTmpUb = tmpAFloorUb.template ReinterpretCast(); + + SoftmaxFlashV2Compute(info, mSplitInfo, mmResUb, softmaxTmpUb, startRow, dealRowCount, columnCount, + info.actualSingleProcessSInnerSize); + PipeBarrier(); + LocalTensor tmpMMResCastTensor = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + + Cast(tmpMMResCastTensor, mmResUb, AscendC::RoundMode::CAST_ROUND, computeSize); + SetFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + DataCopy(vec1ResGm[inOutGmOffset], tmpMMResCastTensor, computeSize); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); +} + +template +__aicore__ inline void SWAVectorBlock::ProcessVec1SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo) +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + uint32_t mSplitSize = info.actualSingleProcessSInnerSize == 0 ? + 16 : + BASE_BLOCK_MAX_ELEMENT_NUM / info.actualSingleProcessSInnerSizeAlign; + // 1. 向下8对齐是因为UB操作至少32B + // 2. info.actualSingleProcessSInnerSizeAlign最大512, mSplitSize可以确保最小为16 + mSplitSize = mSplitSize / 8 * 8; + + if (mSplitSize > mSplitInfo.vecDealM) { + mSplitSize = mSplitInfo.vecDealM; + } + uint32_t loopCount = (mSplitInfo.vecDealM + mSplitSize - 1) / mSplitSize; + uint32_t tailSplitSize = mSplitInfo.vecDealM - (loopCount - 1) * mSplitSize; + + SliceAndContactSinksValue((mSplitInfo.nBufferStartM + mSplitInfo.vecStartM) % constInfo.qHeadNum, + mSplitInfo.vecDealM); + + for (uint32_t i = 0, dealSize = mSplitSize; i < loopCount; i++) { + if (i == (loopCount - 1)) { + dealSize = tailSplitSize; + } + DealBmm1ResBaseBlock(info, mSplitInfo, i * mSplitSize, dealSize, info.actualSingleProcessSInnerSizeAlign, i); + pingpongFlag ^= 1; // pingpong 0 1切换 + } +} + +// =======================vec1============================= + +template +__aicore__ inline void SWAVectorBlock::ProcessVec1L(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferIdx = i; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + + mSplitInfo.vecDealM = (mSplitInfo.nBufferDealM <= 16) ? mSplitInfo.nBufferDealM : + (((mSplitInfo.nBufferDealM + 15) / 16 + 1) / 2 * 16); + mSplitInfo.vecStartM = 0; + if (GetBlockIdx() % 2 == 1) { + mSplitInfo.vecStartM = mSplitInfo.vecDealM; + mSplitInfo.vecDealM = mSplitInfo.nBufferDealM - mSplitInfo.vecDealM; + } + + CrossCoreWaitFlag(constInfo.syncC1V1); + // vec1 compute + ProcessVec1SingleBuf(info, mSplitInfo); + CrossCoreSetFlag(constInfo.syncV1C2); + + // move lse for flash decode or FA + if (constInfo.returnSoftmaxLse && info.s2Idx == info.curSInnerLoopTimes - 1) { + ProcessLse(info, mSplitInfo); + } + } +} + +// =======================vec2============================= + +template +__aicore__ inline uint64_t SWAVectorBlock::CalcAccumOffset(uint32_t bN2Idx, uint32_t gS1Idx) +{ + return 0; +} + +template +__aicore__ inline void SWAVectorBlock::ProcessVec2SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo) +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + + ProcessVec2Inner(info, mSplitInfo, 0, mSplitInfo.vecDealM); +} + +template +__aicore__ inline void SWAVectorBlock::ProcessVec2L(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferIdx = i; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + + mSplitInfo.vecDealM = (mSplitInfo.nBufferDealM <= 16) ? mSplitInfo.nBufferDealM : + (((mSplitInfo.nBufferDealM + 15) / 16 + 1) / 2 * 16); + mSplitInfo.vecStartM = 0; + if (GetBlockIdx() % 2 == 1) { + mSplitInfo.vecStartM = mSplitInfo.vecDealM; + mSplitInfo.vecDealM = mSplitInfo.nBufferDealM - mSplitInfo.vecDealM; + } + CrossCoreWaitFlag(constInfo.syncC2V2); + ProcessVec2SingleBuf(info, mSplitInfo); + } +} + +template +__aicore__ inline void SWAVectorBlock::ProcessVec2Inner(const RunInfo &info, const MSplitInfo &mSplitInfo, + uint32_t mStartRow, uint32_t mDealSize) +{ + uint32_t mSplitSize = BASE_BLOCK_MAX_ELEMENT_NUM / constInfo.headDim; + if (mSplitSize > mDealSize) { + mSplitSize = mDealSize; + } + + uint32_t loopCount = (mDealSize + mSplitSize - 1) / mSplitSize; + uint32_t tailSplitSize = mDealSize - (loopCount - 1) * mSplitSize; + for (uint32_t i = 0, dealSize = mSplitSize; i < loopCount; i++) { + if (i == (loopCount - 1)) { + dealSize = tailSplitSize; + } + DealBmm2ResBaseBlock(info, mSplitInfo, i * mSplitSize + mStartRow, dealSize, constInfo.headDim, + constInfo.headDim); + pingpongFlag ^= 1; // pingpong 0 1切换 + } +} + +template +__aicore__ inline void SWAVectorBlock::Bmm2FDDataCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, + uint32_t wsMStart, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + LocalTensor tmp = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + DataCopy(tmp, bmm2ResUb, columnCount * dealRowCount); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + uint64_t accumTmpOutNum = CalcAccumOffset(info.bIdx, info.gS1Idx); + uint64_t offset = + accumTmpOutNum * constInfo.kvHeadNum * constInfo.mBaseSize * constInfo.headDim + // taskoffset + info.tndCoreStartKVSplitPos * constInfo.kvHeadNum * constInfo.mBaseSize * constInfo.headDim + // 份数offset + wsMStart * actualColumnCount; // m轴offset + GlobalTensor dst = accumOutGm[offset]; + if (info.actualSingleProcessSInnerSize == 0) { + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = dealRowCount; + dataCopyParams.blockLen = actualColumnCount * sizeof(T); + dataCopyParams.srcStride = (columnCount - actualColumnCount) / (BYTE_BLOCK / sizeof(T)); + dataCopyParams.dstStride = 0; + DataCopyPad(dst, tmp, dataCopyParams); + } else { + matmul::InitOutput(dst, dealRowCount * actualColumnCount, ConstInfo::FLOAT_ZERO); + } + SetFlag(SYNC_OUTPUT_BUF1_FLAG); +} + +template +__aicore__ inline void SWAVectorBlock::Bmm2DataCopyOutTrans(const RunInfo &info, LocalTensor &attenOutUb, + uint32_t wsMStart, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = dealRowCount; + dataCopyParams.blockLen = actualColumnCount * sizeof(OUT_T); + dataCopyParams.srcStride = (columnCount - actualColumnCount) / (BYTE_BLOCK / sizeof(OUT_T)); + dataCopyParams.dstStride = 0; + DataCopyPad(attentionOutGm[info.attenOutOffset + wsMStart * actualColumnCount], attenOutUb, dataCopyParams); + return; +} + +template +__aicore__ inline void SWAVectorBlock::Bmm2CastAndCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, + uint32_t wsMStart, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + LocalTensor tmpBmm2ResCastTensor = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + if constexpr (IsSameType::value) { // bf16 采取四舍六入五成双模式 + Cast(tmpBmm2ResCastTensor, bmm2ResUb, AscendC::RoundMode::CAST_RINT, dealRowCount * columnCount); + } else { + Cast(tmpBmm2ResCastTensor, bmm2ResUb, AscendC::RoundMode::CAST_ROUND, dealRowCount * columnCount); + } + + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + Bmm2DataCopyOutTrans(info, tmpBmm2ResCastTensor, wsMStart, dealRowCount, columnCount, actualColumnCount); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); +} + +template +__aicore__ inline void SWAVectorBlock::Bmm2ResCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, + uint32_t wsMStart, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + if constexpr (FLASH_DECODE) { + if (info.tndIsS2SplitCore) { + Bmm2FDDataCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount); + } else { + Bmm2CastAndCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount); + } + } else { + Bmm2CastAndCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount); + } +} + +template +__aicore__ inline void SWAVectorBlock::DealBmm2ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, + uint32_t startRow, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + uint32_t vec2ComputeSize = dealRowCount * columnCount; + uint32_t mStart = mSplitInfo.nBufferStartM + mSplitInfo.vecStartM + startRow; + uint64_t srcGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + mStart * columnCount; + LocalTensor tmpBmm2ResUb = inputBuff1.Get(); + tmpBmm2ResUb = tmpBmm2ResUb[pingpongFlag * INPUT1_BUFFER_OFFSET / sizeof(MM2_OUT_T)]; + WaitFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + DataCopy(tmpBmm2ResUb, mm2ResGm[srcGmOffset], vec2ComputeSize); + + SetFlag(SYNC_INPUT_BUF1_FLAG); + WaitFlag(SYNC_INPUT_BUF1_FLAG); + + LocalTensor bmm2ResUb = tmpBuff1.Get(); + bmm2ResUb.SetSize(vec2ComputeSize); + DataCopy(bmm2ResUb, tmpBmm2ResUb, vec2ComputeSize); + SetFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + + uint32_t inOutBaseOffset = mStart * columnCount; + uint32_t baseOffset = mSplitInfo.nBufferStartM / 2 + startRow; + + // 除第一个循环外，均需要更新中间计算结果 + if (!info.isFirstSInnerLoop) { + event_t eventIdMte2WaitMte3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventIdMte2WaitMte3); + WaitFlag(eventIdMte2WaitMte3); + + LocalTensor bmm2ResPreUb = inputBuff2.Get(); + WaitFlag(SYNC_INPUT_BUF2_FLAG); + + uint64_t vec2ResGmOffset = ((info.loop - 1) % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + inOutBaseOffset; + DataCopy(bmm2ResPreUb, vec2ResGm[vec2ResGmOffset], vec2ComputeSize); + + SetFlag(SYNC_INPUT_BUF2_FLAG); + WaitFlag(SYNC_INPUT_BUF2_FLAG); + + uint32_t idx = info.loop % (constInfo.preLoadNum); + LocalTensor expUb = v0ValidSizeBuff.Get()[384]; // sumUb用临时内存 16 * 32B = 512B + Brcb(expUb, softmaxExpUb[idx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset], (dealRowCount + 7) / 8, + {1, 8}); + PipeBarrier(); + + RowMuls(bmm2ResPreUb, bmm2ResPreUb, expUb, dealRowCount, columnCount, actualColumnCount); + AscendC::PipeBarrier(); + Add(bmm2ResUb, bmm2ResUb, bmm2ResPreUb, vec2ComputeSize); + AscendC::PipeBarrier(); + + SetFlag(SYNC_INPUT_BUF2_FLAG); + } + + // 最后一次输出计算结果，否则将中间结果暂存至workspace + if (info.isLastS2Loop) { + uint32_t idx = info.loop % (constInfo.preLoadNum); + LocalTensor tmpSumUb = v0ValidSizeBuff.Get()[384]; // sumUb用临时内存 16 * 32B = 512B + Brcb(tmpSumUb, softmaxSumUb[idx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset], (dealRowCount + 7) / 8, + {1, 8}); + PipeBarrier(); + RowDivs(bmm2ResUb, bmm2ResUb, tmpSumUb, dealRowCount, columnCount, actualColumnCount); + PipeBarrier(); + Bmm2ResCopyOut(info, bmm2ResUb, mStart, dealRowCount, columnCount, actualColumnCount); + } else { + LocalTensor outUb = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + DataCopy(outUb, bmm2ResUb, dealRowCount * columnCount); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + uint64_t vec2ResGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + inOutBaseOffset; + DataCopy(vec2ResGm[vec2ResGmOffset], outUb, vec2ComputeSize); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + } +} + +template +__aicore__ inline void SWAVectorBlock::RowDivs(LocalTensor dstUb, LocalTensor src0Ub, + LocalTensor src1Ub, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + // divs by row, 每行的元素除以相同的元素 + // dstUb[i, (j * 8) : (j * 8 + 7)] = src0Ub[i, (j * 8) : (j * 8 + 7)] / src1Ub[i, 0 : 7] + // src0Ub:[dealRowCount, columnCount], src1Ub:[dealRowCount, FP32_BLOCK_ELEMENT_NUM] dstUb:[dealRowCount, + // columnCount] + uint32_t dtypeMask = FP32_REPEAT_ELEMENT_NUM; + uint32_t dLoop = actualColumnCount / dtypeMask; + uint32_t dRemain = actualColumnCount % dtypeMask; + + BinaryRepeatParams repeatParamsDiv; + repeatParamsDiv.src0BlkStride = 1; + repeatParamsDiv.src1BlkStride = 0; + repeatParamsDiv.dstBlkStride = 1; + repeatParamsDiv.src0RepStride = columnCount / FP32_BLOCK_ELEMENT_NUM; + repeatParamsDiv.src1RepStride = 1; + repeatParamsDiv.dstRepStride = columnCount / FP32_BLOCK_ELEMENT_NUM; + uint32_t columnRepeatCount = dLoop; + if (columnRepeatCount <= dealRowCount) { + uint32_t offset = 0; + for (uint32_t i = 0; i < dLoop; i++) { + Div(dstUb[offset], src0Ub[offset], src1Ub, dtypeMask, dealRowCount, repeatParamsDiv); + offset += dtypeMask; + } + } else { + BinaryRepeatParams columnRepeatParams; + columnRepeatParams.src0BlkStride = 1; + columnRepeatParams.src1BlkStride = 0; + columnRepeatParams.dstBlkStride = 1; + columnRepeatParams.src0RepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + columnRepeatParams.src1RepStride = 0; + columnRepeatParams.dstRepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + uint32_t offset = 0; + for (uint32_t i = 0; i < dealRowCount; i++) { + Div(dstUb[offset], src0Ub[offset], src1Ub[i * FP32_BLOCK_ELEMENT_NUM], dtypeMask, columnRepeatCount, + columnRepeatParams); + offset += columnCount; + } + } + if (dRemain > 0) { + Div(dstUb[dLoop * dtypeMask], src0Ub[dLoop * dtypeMask], src1Ub, dRemain, dealRowCount, repeatParamsDiv); + } +} + +template +__aicore__ inline void SWAVectorBlock::RowMuls(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount) +{ + // muls by row, 每行的元素乘以相同的元素 + // dstUb[i, (j * 8) : (j * 8 + 7)] = src0Ub[i, (j * 8) : (j * 8 + 7)] * src1Ub[i, 0 : 7] + // src0Ub:[dealRowCount, columnCount] src1Ub:[dealRowCount, FP32_BLOCK_ELEMENT_NUM] dstUb:[dealRowCount, + // columnCount] + // dealRowCount is repeat times, must be less 256 + uint32_t repeatElementNum = FP32_REPEAT_ELEMENT_NUM; + uint32_t blockElementNum = FP32_BLOCK_ELEMENT_NUM; + + if constexpr (std::is_same::value) { + // 此限制由于每个repeat至多连续读取256B数据 + repeatElementNum = FP32_REPEAT_ELEMENT_NUM * 2; // 256/4 * 2=128 + blockElementNum = FP32_BLOCK_ELEMENT_NUM * 2; // 32/4 * 2 = 16 + } + + // 每次只能连续读取256B的数据进行计算，故每次只能处理256B/sizeof(dType)= + // 列方向分dLoop次，每次处理8列数据 + uint32_t dLoop = actualColumnCount / repeatElementNum; + uint32_t dRemain = actualColumnCount % repeatElementNum; + // REPEATE_STRIDE_UP_BOUND=256， 此限制由于src0RepStride数据类型为uint8之多256个datablock间距 + if (columnCount < REPEATE_STRIDE_UP_BOUND * blockElementNum) { + BinaryRepeatParams repeatParams; + repeatParams.src0BlkStride = 1; + repeatParams.src1BlkStride = 0; + repeatParams.dstBlkStride = 1; + repeatParams.src0RepStride = columnCount / blockElementNum; + repeatParams.src1RepStride = 1; + repeatParams.dstRepStride = columnCount / blockElementNum; + + // 如果以列为repeat所处理的次数小于行处理次数，则以列方式处理。反之则以行进行repeat处理 + if (dLoop <= dealRowCount) { + uint32_t offset = 0; + for (uint32_t i = 0; i < dLoop; i++) { + Mul(dstUb[offset], src0Ub[offset], src1Ub, repeatElementNum, dealRowCount, repeatParams); + offset += repeatElementNum; + } + } else { + BinaryRepeatParams columnRepeatParams; + columnRepeatParams.src0BlkStride = 1; + columnRepeatParams.src1BlkStride = 0; + columnRepeatParams.dstBlkStride = 1; + columnRepeatParams.src0RepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + columnRepeatParams.src1RepStride = 0; + columnRepeatParams.dstRepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + for (uint32_t i = 0; i < dealRowCount; i++) { + Mul(dstUb[i * columnCount], src0Ub[i * columnCount], src1Ub[i * blockElementNum], repeatElementNum, + dLoop, columnRepeatParams); + } + } + + // 最后一次完成[dealRowCount, dRemain] * [dealRowCount, blockElementNum] 只计算有效部分 + if (dRemain > 0) { + Mul(dstUb[dLoop * repeatElementNum], src0Ub[dLoop * repeatElementNum], src1Ub, dRemain, dealRowCount, + repeatParams); + } + } else { + BinaryRepeatParams repeatParams; + repeatParams.src0RepStride = 8; // 每个repeat为256B数据，正好8个datablock + repeatParams.src0BlkStride = 1; + repeatParams.src1RepStride = 0; + repeatParams.src1BlkStride = 0; + repeatParams.dstRepStride = 8; + repeatParams.dstBlkStride = 1; + // 每次计算一行，共计算dealRowCount行 + for (uint32_t i = 0; i < dealRowCount; i++) { + // 计算一行中的dLoop个repeat, 每个repeat计算256/block_size 个data_block + Mul(dstUb[i * columnCount], src0Ub[i * columnCount], src1Ub[i * blockElementNum], repeatElementNum, dLoop, + repeatParams); + // 计算一行中的尾块 + if (dRemain > 0) { + Mul(dstUb[i * columnCount + dLoop * repeatElementNum], + src0Ub[i * columnCount + dLoop * repeatElementNum], src1Ub[i * blockElementNum], dRemain, 1, + repeatParams); + } + } + } +} +} // namespace SASKernel +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_swa_kernel.h b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_swa_kernel.h new file mode 100644 index 000000000..864e4eafb --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/arch32/sparse_attn_sharedkv_swa_kernel.h @@ -0,0 +1,771 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_swa_kernel.h + * \brief + */ + +#ifndef SPARSE_ATTN_SHAREDKV_SWA_KERNEL_H +#define SPARSE_ATTN_SHAREDKV_SWA_KERNEL_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../sparse_attn_sharedkv_common.h" +#include "sparse_attn_sharedkv_swa_block_cube.h" +#include "sparse_attn_sharedkv_swa_block_vector.h" +#include "../sparse_attn_sharedkv_metadata.h" + +namespace SASKernel { +using namespace matmul; +using namespace optiling; +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +// 由于S2循环前，RunInfo还没有赋值，使用Bngs1Param临时存放B、N、S1轴相关的信息；同时减少重复计算 +struct SwaTempLoopInfo { + uint32_t bn2IdxInCurCore = 0; + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint64_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小 + uint32_t s2LoopTimes = 0U; // S2方向循环的总次数，无论TND还是BXXD都是等于实际次数，不用减1 + + int32_t actS1Size = 0; // TND场景下当前Batch循环处理的S1轴的大小 + int32_t actOriS2Size = 0; + int32_t actCmpS2Size = 0; + + bool curActSeqLenIsZero = false; + + uint32_t tndCoreStartKVSplitPos = 0; + bool tndIsS2SplitCore = false; + + uint32_t gS1Idx = 0U; + uint32_t s1StartIdx = 0; + uint32_t s1EndIdx = 0; + uint64_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小 + uint32_t cmpLoopTimes = 0; + uint32_t oriLoopTimes = 0; + + int32_t oriMaskRight = 0; + int32_t oriMaskLeft = 0; + int32_t cmpMaskRight = 0; + + uint64_t actualSeqQPrefixSum = 0; + uint64_t actualSeqKVPrefixSum = 0; + uint64_t actualSeqCmpKVPrefixSum = 0; +}; + +template +class SparseAttnSharedkvSwa { +public: + // 中间计算数据类型为float，高精度模式 + using T = float; + using Q_T = typename SAST::queryType; + using KV_T = typename SAST::kvType; + using OUT_T = typename SAST::outputType; + using SINKS_T = float; + using UPDATE_T = T; + using MM1_OUT_T = T; + using MM2_OUT_T = T; + + __aicore__ inline SparseAttnSharedkvSwa(){}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *oriKV, __gm__ uint8_t *cmpKV, + __gm__ uint8_t *cmpSparseIndices, __gm__ uint8_t *oriBlockTable, + __gm__ uint8_t *cmpBlockTable, __gm__ uint8_t *cuSeqlensQ, + __gm__ uint8_t* cuSeqlensKV, __gm__ uint8_t *cuSeqlensCmpKV, __gm__ uint8_t *seqUsedQ, + __gm__ uint8_t *seqUsedKV, __gm__ uint8_t *sinks, __gm__ uint8_t *metadata, + __gm__ uint8_t *attentionOut, __gm__ uint8_t *softmaxLse, __gm__ uint8_t *workspace, + const SparseAttnSharedkvTilingData *__restrict tiling, __gm__ uint8_t *gmTiling, + TPipe *tPipe); + + __aicore__ inline void Process(); + +private: + static constexpr bool PAGE_ATTENTION = SAST::pageAttention; + static constexpr int TEMPLATE_MODE = SAST::templateMode; + static constexpr bool FLASH_DECODE = SAST::flashDecode; + static constexpr SAS_LAYOUT LAYOUT_T = SAST::layout; + static constexpr SAS_LAYOUT KV_LAYOUT_T = SAST::kvLayout; + + static constexpr uint32_t PRELOAD_NUM = 2; + static constexpr uint32_t N_BUFFER_M_BASIC_SIZE = 256; + static constexpr uint32_t SAS_PRELOAD_TASK_CACHE_SIZE = 3; + + static constexpr uint32_t SYNC_V0_C1_FLAG = 6; + static constexpr uint32_t SYNC_C1_V1_FLAG = 7; + static constexpr uint32_t SYNC_V1_C2_FLAG = 8; + static constexpr uint32_t SYNC_C2_V2_FLAG = 9; + + static constexpr uint64_t SYNC_MM2RES_BUF1_FLAG = 10; + static constexpr uint64_t SYNC_MM2RES_BUF2_FLAG = 11; + static constexpr uint64_t SYNC_FDOUTPUT_BUF_FLAG = 12; + + static constexpr uint64_t kvHeadNum = 1ULL; + static constexpr uint64_t headDim = 512ULL; + static constexpr uint64_t headDimAlign = 512ULL; + static constexpr uint32_t msdIterNum = 2U; + + static constexpr uint32_t dbWorkspaceRatio = PRELOAD_NUM; + + const SparseAttnSharedkvTilingData *__restrict tilingData = nullptr; + + TPipe *pipe = nullptr; + GlobalTensor metadataGm; + + uint64_t mSizeVStart = 0ULL; + int64_t threshold = 0; + uint64_t s2BatchBaseOffset = 0; + uint64_t tensorACoreOffset = 0ULL; + uint64_t tensorBCoreOffset = 0ULL; + uint64_t tensorCmpBCoreOffset = 0ULL; + uint64_t attenOutOffset = 0ULL; + + uint32_t tmpBlockIdx = 0U; + uint32_t aiCoreIdx = 0U; + + ConstInfo constInfo{}; + SwaTempLoopInfo tempLoopInfo{}; + + SWACubeBlock cubeBlock; + SWAVectorBlock vectorBlock; + + GlobalTensor queryGm; + GlobalTensor oriKvGm; + GlobalTensor cmpKvGm; + GlobalTensor sinksGm; + + GlobalTensor attentionOutGm; + GlobalTensor softmaxLseGm; + + GlobalTensor oriBlockTableGm; + GlobalTensor cmpBlockTableGm; + + GlobalTensor actualSeqLengthsQGm; + GlobalTensor actualSeqLengthsKVGm; + GlobalTensor actualSeqLengthsCmpKVGm; + + // workspace + GlobalTensor mm1ResGm; + GlobalTensor vec1ResGm; + GlobalTensor mm2ResGm; + + GlobalTensor vec2ResGm; + + GlobalTensor accumOutGm; + // ================================Init functions================================== + __aicore__ inline void InitTilingData(); + __aicore__ inline void InitCalcParamsEach(); + __aicore__ inline void InitBuffers(); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsKv); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsKV, + __gm__ uint8_t *actualSeqLengthsCmpKV); + __aicore__ inline void InitOutputSingleCore(); + // ================================Process functions================================ + __aicore__ inline void ProcessBalance(); + __aicore__ inline void PreloadPipeline(uint32_t loop, uint32_t cmpLoop, uint64_t s2Start, uint64_t s2LoopIdx, + RunInfo extraInfo[SAS_PRELOAD_TASK_CACHE_SIZE]); + // ================================Offset Calc===================================== + __aicore__ inline void GetSparseActualSeqLen(); + __aicore__ inline void UpdateInnerLoopCond(); + __aicore__ inline void CalcParams(uint32_t loop, uint32_t cmpLoop, uint64_t s2Start, uint32_t s2LoopIdx, + RunInfo &info); + __aicore__ inline int32_t GetActualSeqLenQ(uint32_t bIdx); + __aicore__ inline int32_t GetActualSeqLenKV(uint32_t bIdx); + __aicore__ inline void GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx); + // ================================Mm1============================================== + __aicore__ inline void ComputeMm1(const RunInfo &info); + // ================================Mm2============================================== + __aicore__ inline void ComputeMm2(const RunInfo &info); + __aicore__ inline void InitAllZeroOutput(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx); +}; + +template +__aicore__ inline void SparseAttnSharedkvSwa::InitTilingData() +{ + // singleCoreParams + // singleCoreTensorSize + constInfo.mmResUbSize = tilingData->baseParams.mmResUbSize; + constInfo.bmm2ResUbSize = tilingData->baseParams.bmm2ResUbSize; + // baseParams + constInfo.batchSize = tilingData->baseParams.batchSize; + constInfo.qHeadNum = constInfo.gSize = tilingData->baseParams.nNumOfQInOneGroup; + constInfo.kvSeqSize = tilingData->baseParams.kvSeqSize; + constInfo.qSeqSize = tilingData->baseParams.qSeqSize; + constInfo.oriMaxBlockNumPerBatch = tilingData->baseParams.oriMaxBlockNumPerBatch; + constInfo.kvCacheBlockSize = tilingData->baseParams.paBlockSize; + + constInfo.paOriBlockSize = tilingData->baseParams.oriBlockSize; + constInfo.paCmpBlockSize = tilingData->baseParams.cmpBlockSize; + constInfo.outputLayout = static_cast(tilingData->baseParams.outputLayout); + constInfo.kvHeadNum = kvHeadNum; + constInfo.headDim = headDim; + constInfo.oriMaskMode = tilingData->baseParams.oriMaskMode; + constInfo.oriKvStride = tilingData->baseParams.oriKvStride; + constInfo.oriWinLeft = tilingData->baseParams.oriWinLeft; + constInfo.oriWinRight = tilingData->baseParams.oriWinRight; + constInfo.returnSoftmaxLse = tilingData->baseParams.returnSoftmaxLse; + + constInfo.actualLenDimsQ = tilingData->baseParams.actualLenDimsQ; + constInfo.actualLenDimsKV = tilingData->baseParams.actualLenDimsKV; + + // innerSplitParams + constInfo.mBaseSize = constInfo.gSize;; + constInfo.s2BaseSize = tilingData->baseParams.s2BaseSize; + // tilingData->baseParams.s2BaseSize + + constInfo.preLoadNum = PRELOAD_NUM; + constInfo.nBufferMBaseSize = N_BUFFER_M_BASIC_SIZE; + constInfo.syncV0C1 = SYNC_V0_C1_FLAG; + constInfo.syncC1V1 = SYNC_C1_V1_FLAG; + constInfo.syncV1C2 = SYNC_V1_C2_FLAG; + constInfo.syncC2V2 = SYNC_C2_V2_FLAG; + constInfo.templateMode = TEMPLATE_MODE; + + // cmp + if (constInfo.templateMode == CFA_TEMPLATE) { + constInfo.cmpRatio = tilingData->cmpParams.cmpRatio; + constInfo.cmpMaskMode = tilingData->cmpParams.cmpMaskMode; + constInfo.cmpKvStride = tilingData->cmpParams.cmpKvStride; + constInfo.cmpMaxBlockNumPerBatch = tilingData->cmpParams.cmpMaxBlockNumPerBatch; + } +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::InitBuffers() +{ + if ASCEND_IS_AIV { + vectorBlock.InitBuffers(pipe); + } else { + cubeBlock.InitBuffers(pipe); + } +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengthsKv) +{ + if (constInfo.actualLenDimsKV != 0) { + actualSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsKv, constInfo.actualLenDimsKV); + } + if (constInfo.actualLenDimsQ != 0) { + actualSeqLengthsQGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsQ, constInfo.actualLenDimsQ); + } +} + +template +__aicore__ inline void +SparseAttnSharedkvSwa::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsKV, + __gm__ uint8_t *actualSeqLengthsCmpKV) +{ + if (constInfo.actualLenDimsKV != 0) { + actualSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsKV, constInfo.actualLenDimsKV); + if (constInfo.templateMode == CFA_TEMPLATE) { + actualSeqLengthsCmpKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsCmpKV, constInfo.actualLenDimsKV); + } + } + if (constInfo.actualLenDimsQ != 0) { + actualSeqLengthsQGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsQ, constInfo.actualLenDimsQ); + } +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::InitAllZeroOutput(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx) +{ + if (constInfo.outputLayout == SAS_LAYOUT::TND) { + if (tempLoopInfo.actS1Size == 0) { + return; + } + uint32_t tBase = actualSeqLengthsQGm.GetValue(bIdx); + uint64_t attenOutOffset = (tBase + s1Idx) * kvHeadNum * constInfo.gSize * headDim + // T轴、s1轴偏移 + n2Idx * constInfo.gSize * headDim; // N2轴偏移 + uint64_t lseOffset = (tBase + s1Idx) * constInfo.gSize + // T轴、s1轴偏移 + n2Idx * constInfo.qSeqSize * constInfo.gSize; // N2轴偏移 + matmul::InitOutput(attentionOutGm[attenOutOffset], constInfo.gSize * headDim, 0); + if (constInfo.returnSoftmaxLse) { + matmul::InitOutput(softmaxLseGm[lseOffset], constInfo.gSize, 0); + } + } else if (constInfo.outputLayout == SAS_LAYOUT::BSND) { + uint64_t attenOutOffset = bIdx * constInfo.qSeqSize * kvHeadNum * constInfo.gSize * headDim + + s1Idx * kvHeadNum * constInfo.gSize * headDim + // B轴、S1轴偏移 + n2Idx * constInfo.gSize * headDim; // N2轴偏移 + uint64_t lseOffset = bIdx * constInfo.qSeqSize * constInfo.kvHeadNum * constInfo.gSize + // B轴偏移 + n2Idx * constInfo.qSeqSize * constInfo.gSize + // N2轴偏移 + s1Idx * constInfo.gSize; // S1轴偏移 + matmul::InitOutput(attentionOutGm[attenOutOffset], constInfo.gSize * headDim, 0); + if (constInfo.returnSoftmaxLse) { + matmul::InitOutput(softmaxLseGm[lseOffset], constInfo.gSize, 0); + } + } +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::InitOutputSingleCore() +{ + uint32_t coreNum = GetBlockNum(); + if (coreNum != 0) { + uint64_t totalOutputSize = constInfo.batchSize * constInfo.qHeadNum * constInfo.qSeqSize * constInfo.headDim; + uint64_t singleCoreSize = (totalOutputSize + (2 * coreNum) - 1) / (2 * coreNum); // 2 means c:v = 1:2 + uint64_t tailSize = totalOutputSize - tmpBlockIdx * singleCoreSize; + uint64_t singleInitOutputSize = tailSize < singleCoreSize ? tailSize : singleCoreSize; + if (singleInitOutputSize > 0) { + matmul::InitOutput(attentionOutGm[tmpBlockIdx * singleCoreSize], singleInitOutputSize, 0); + } + SyncAll(); + } +} + +template +__aicore__ inline int32_t SparseAttnSharedkvSwa::GetActualSeqLenQ(uint32_t bIdx) +{ + if constexpr (LAYOUT_T == SAS_LAYOUT::TND) { + int32_t actualSeqQPrefixSum = actualSeqLengthsQGm.GetValue(bIdx); + int32_t actualSeqQNextSum = actualSeqLengthsQGm.GetValue(bIdx + 1); + tempLoopInfo.actualSeqQPrefixSum = static_cast(actualSeqQPrefixSum); + return actualSeqQNextSum - actualSeqQPrefixSum; + } else { + tempLoopInfo.actualSeqQPrefixSum = static_cast(bIdx * constInfo.qSeqSize); + if (constInfo.actualLenDimsQ == 0) { + return static_cast(constInfo.qSeqSize); + } else { + return actualSeqLengthsQGm.GetValue(bIdx); + } + } +} + +template +__aicore__ inline int32_t SparseAttnSharedkvSwa::GetActualSeqLenKV(uint32_t bIdx) +{ + if constexpr (KV_LAYOUT_T == SAS_LAYOUT::PA_ND) { + tempLoopInfo.actualSeqKVPrefixSum = static_cast(bIdx * constInfo.kvSeqSize); + if (constInfo.actualLenDimsKV == 0) { + return static_cast(constInfo.kvSeqSize); + } + return actualSeqLengthsKVGm.GetValue(bIdx); + } else if constexpr(KV_LAYOUT_T == SAS_LAYOUT::BSND) { + return static_cast(constInfo.kvSeqSize); + } else if constexpr(KV_LAYOUT_T == SAS_LAYOUT::TND) { + int32_t actualSeqKVPrefixSum = actualSeqLengthsKVGm.GetValue(bIdx); + int32_t actualSeqKVNextSum = actualSeqLengthsKVGm.GetValue(bIdx + 1); + if (constInfo.templateMode == CFA_TEMPLATE) { + tempLoopInfo.actualSeqCmpKVPrefixSum = actualSeqLengthsCmpKVGm.GetValue(bIdx); + } + tempLoopInfo.actualSeqKVPrefixSum = actualSeqKVPrefixSum; + return actualSeqKVNextSum - actualSeqKVPrefixSum; + } +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::GetSparseActualSeqLen() +{ + // 行无效通过ori部分判断, ori部分如果有行无效那么ori和cmp都有 + if (static_cast(tempLoopInfo.s1EndIdx) < -(tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size)) { + tempLoopInfo.actOriS2Size = 0; + tempLoopInfo.actCmpS2Size = 0; + return; + } + + // 对于cmp部分还有top k, tempLoopInfo.actS2Size只针对cmp + if (constInfo.templateMode == CFA_TEMPLATE) { + int32_t thresHold = (tempLoopInfo.cmpMaskRight + tempLoopInfo.s1EndIdx + 1) / constInfo.cmpRatio; + tempLoopInfo.actCmpS2Size = thresHold; + } +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::UpdateInnerLoopCond() +{ + if ((tempLoopInfo.actCmpS2Size == 0 && tempLoopInfo.actOriS2Size == 0) || (tempLoopInfo.actS1Size == 0)) { + tempLoopInfo.curActSeqLenIsZero = true; + return; + } + tempLoopInfo.curActSeqLenIsZero = false; + tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize; + tempLoopInfo.mBasicSizeTail = + (tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail; +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::Init( + __gm__ uint8_t *query, __gm__ uint8_t *oriKV, __gm__ uint8_t *cmpKV, __gm__ uint8_t *cmpSparseIndices, + __gm__ uint8_t *oriBlockTable, __gm__ uint8_t *cmpBlockTable, __gm__ uint8_t *cuSeqlensQ, + __gm__ uint8_t *cuSeqlensKV, __gm__ uint8_t *cuSeqlensCmpKV, __gm__ uint8_t *seqUsedQ, + __gm__ uint8_t *seqUsedKV, __gm__ uint8_t *sinks, __gm__ uint8_t *metadata, __gm__ uint8_t *attentionOut, __gm__ uint8_t *softmaxLse, + __gm__ uint8_t *workspace, const SparseAttnSharedkvTilingData *__restrict tiling, __gm__ uint8_t *gmTiling, + TPipe *tPipe) +{ + if ASCEND_IS_AIV { + tmpBlockIdx = GetBlockIdx(); // vec:0-47 + aiCoreIdx = tmpBlockIdx / 2; + } else { + tmpBlockIdx = GetBlockIdx(); // cube:0-23 + aiCoreIdx = tmpBlockIdx; + } + + // init tiling data + tilingData = tiling; + + InitTilingData(); + if (KV_LAYOUT_T == SAS_LAYOUT::TND && LAYOUT_T == SAS_LAYOUT::TND) { + InitActualSeqLen(cuSeqlensQ, cuSeqlensKV, cuSeqlensCmpKV); + } else if (KV_LAYOUT_T == SAS_LAYOUT::TND) { + InitActualSeqLen(seqUsedQ, cuSeqlensKV, cuSeqlensCmpKV); + } else if ((KV_LAYOUT_T == SAS_LAYOUT::PA_ND || KV_LAYOUT_T == SAS_LAYOUT::BSND) && LAYOUT_T == SAS_LAYOUT::TND) { + InitActualSeqLen(cuSeqlensQ, seqUsedKV); + } else if ((KV_LAYOUT_T == SAS_LAYOUT::PA_ND || KV_LAYOUT_T == SAS_LAYOUT::BSND)) { + InitActualSeqLen(seqUsedQ, seqUsedKV); + } + metadataGm.SetGlobalBuffer((__gm__ uint32_t *)metadata); + InitCalcParamsEach(); + + pipe = tPipe; + // init global buffer + queryGm.SetGlobalBuffer((__gm__ Q_T *)query); + oriKvGm.SetGlobalBuffer((__gm__ KV_T *)oriKV); + if (constInfo.templateMode == CFA_TEMPLATE) { + cmpKvGm.SetGlobalBuffer((__gm__ KV_T *)cmpKV); + } + + if (sinks != nullptr) { + sinksGm.SetGlobalBuffer((__gm__ SINKS_T *)sinks); + } + + attentionOutGm.SetGlobalBuffer((__gm__ OUT_T *)attentionOut); + softmaxLseGm.SetGlobalBuffer((__gm__ T *)softmaxLse); + + if ASCEND_IS_AIV { + if (LAYOUT_T != SAS_LAYOUT::TND) { + if (constInfo.needInit) { + InitOutputSingleCore(); + } + } + } + + if constexpr (PAGE_ATTENTION) { + oriBlockTableGm.SetGlobalBuffer((__gm__ int32_t *)oriBlockTable); + if (constInfo.templateMode == CFA_TEMPLATE) { + cmpBlockTableGm.SetGlobalBuffer((__gm__ int32_t *)cmpBlockTable); + } + } + + // workspace 内存排布 + // |Q--|mm1ResGm|vec1ResGm|mm2ResGm|vec2ResGm + // |Core0_Q1-Core0_Q2-Core1_Q1-Core1_Q2....Core32_Q1-Core32_Q2|Core0_mmRes + uint64_t offset = 0; + mm1ResGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + offset + + aiCoreIdx * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(MM1_OUT_T))); + offset += GetBlockNum() * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(MM1_OUT_T); + + vec1ResGm.SetGlobalBuffer( + (__gm__ Q_T *)(workspace + offset + aiCoreIdx * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(KV_T))); + offset += GetBlockNum() * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(KV_T); + + mm2ResGm.SetGlobalBuffer( + (__gm__ MM2_OUT_T *)(workspace + offset + + aiCoreIdx * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(MM2_OUT_T))); + offset += GetBlockNum() * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(MM2_OUT_T); + + vec2ResGm.SetGlobalBuffer( + (__gm__ T *)(workspace + offset + aiCoreIdx * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(T))); + offset += GetBlockNum() * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(T); + + if ASCEND_IS_AIV { + vectorBlock.InitParams(constInfo, tilingData); + vectorBlock.InitVec1GlobalTensor(mm1ResGm, vec1ResGm, actualSeqLengthsQGm, actualSeqLengthsKVGm, sinksGm, softmaxLseGm); + vectorBlock.InitVec2GlobalTensor(accumOutGm, vec2ResGm, mm2ResGm, attentionOutGm); + } + + if ASCEND_IS_AIC { + cubeBlock.InitParams(constInfo); + cubeBlock.InitMm1GlobalTensor(queryGm, oriKvGm, cmpKvGm, mm1ResGm); + cubeBlock.InitMm2GlobalTensor(vec1ResGm, mm2ResGm, attentionOutGm); + cubeBlock.InitPageAttentionInfo(oriKvGm, oriBlockTableGm, cmpBlockTableGm); + } + // 要在InitParams之后执行 + if (pipe != nullptr) { + InitBuffers(); + } +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::InitCalcParamsEach() +{ + if (aiCoreIdx != 0) { + constInfo.bN2Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_BN2_START_INDEX, false)); + constInfo.gS1Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_M_START_INDEX, false)); + constInfo.s2Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_S2_START_INDEX, false)); + } + constInfo.bN2End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_BN2_END_INDEX, false)); + constInfo.gS1End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_M_END_INDEX, false)); + constInfo.s2End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_S2_END_INDEX, false)); +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::CalcParams(uint32_t loop, uint32_t cmpLoop, uint64_t s2Start, + uint32_t s2LoopIdx, RunInfo &info) +{ + info.isValid = s2LoopIdx < tempLoopInfo.s2LoopTimes; + info.loop = loop; + info.cmpLoop = cmpLoop; + info.bIdx = tempLoopInfo.bIdx; + info.gS1Idx = tempLoopInfo.gS1Idx; + info.s1Idx = tempLoopInfo.gS1Idx / constInfo.gSize; + info.s2Idx = s2LoopIdx; + info.n2IdxReal = tempLoopInfo.n2Idx; + + info.curSInnerLoopTimes = tempLoopInfo.s2LoopTimes; + info.tndIsS2SplitCore = tempLoopInfo.tndIsS2SplitCore; + info.tndCoreStartKVSplitPos = tempLoopInfo.tndCoreStartKVSplitPos; + info.isBmm2Output = false; + info.actS1Size = tempLoopInfo.actS1Size; + + // M方向的尾块 + info.actMBaseSize = tempLoopInfo.mBasicSizeTail; + + if ASCEND_IS_AIV { + info.mSize = info.actMBaseSize; + info.mSizeV = (info.mSize <= 16) ? info.mSize : ((CeilDiv(info.mSize, 16) + 1) / 2 * 16); + info.mSizeVStart = 0; + if (tmpBlockIdx % 2 == 1) { + info.mSizeVStart = info.mSizeV; + info.mSizeV = info.mSize - info.mSizeV; + } + } + + info.isFirstSInnerLoop = s2LoopIdx == s2Start; + if (info.isFirstSInnerLoop) { + tempLoopInfo.bn2IdxInCurCore++; + } + info.isLastS2Loop = (s2LoopIdx == (tempLoopInfo.s2LoopTimes - 1)); + info.bn2IdxInCurCore = tempLoopInfo.bn2IdxInCurCore - 1; + + uint64_t tndBIdxOffsetForQ = tempLoopInfo.actualSeqQPrefixSum * constInfo.qHeadNum * constInfo.headDim; + uint64_t tndBIdxOffsetForKV = tempLoopInfo.actualSeqKVPrefixSum * constInfo.kvHeadNum * constInfo.headDim; + uint64_t tndBIdxOffsetForCmpKV = tempLoopInfo.actualSeqCmpKVPrefixSum * constInfo.kvHeadNum * constInfo.headDim; + + if (info.isFirstSInnerLoop) { + tensorACoreOffset = tndBIdxOffsetForQ + info.gS1Idx * constInfo.headDim; + tensorBCoreOffset = tndBIdxOffsetForKV + info.n2Idx * constInfo.headDim; + tensorCmpBCoreOffset = tndBIdxOffsetForCmpKV + info.n2Idx * constInfo.headDim; + } + info.tensorAOffset = tensorACoreOffset; + info.tensorBOffset = tensorBCoreOffset; + info.tensorCmpBOffset = tensorCmpBCoreOffset; + info.attenOutOffset = tensorACoreOffset; + + if (s2LoopIdx < tempLoopInfo.oriLoopTimes) { + // S2首次循环只能在ori_kv + info.isOri = true; + info.relativeS2Idx = 0; + uint64_t s2Offset = info.s2Idx * constInfo.s2BaseSize; + if (s2LoopIdx + 1 == tempLoopInfo.oriLoopTimes) { + info.actualSingleProcessSInnerSize = (tempLoopInfo.oriMaskRight - tempLoopInfo.oriMaskLeft + 1) - s2Offset; + } else { + info.actualSingleProcessSInnerSize = constInfo.s2BaseSize; + } + info.s2StartPoint = tempLoopInfo.oriMaskLeft; + info.cmpS2IdLimit = 0; + } else { + if (constInfo.templateMode == CFA_TEMPLATE) { + info.isOri = false; + info.relativeS2Idx = info.s2Idx - tempLoopInfo.oriLoopTimes; + uint64_t s2Offset = (info.s2Idx - tempLoopInfo.oriLoopTimes) * constInfo.s2BaseSize; + if (s2LoopIdx + 1 == tempLoopInfo.s2LoopTimes) { + info.actualSingleProcessSInnerSize = tempLoopInfo.actCmpS2Size - s2Offset; + } else { + info.actualSingleProcessSInnerSize = constInfo.s2BaseSize; + } + info.s2StartPoint = 0; + info.cmpS2IdLimit = (tempLoopInfo.cmpMaskRight + tempLoopInfo.s1EndIdx + 1) / constInfo.cmpRatio; + } + } + + info.actualSingleProcessSInnerSizeAlign = + SASAlign(info.actualSingleProcessSInnerSize, SASVectorBlock::BYTE_BLOCK); +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::ComputeMm1(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = CeilDiv(info.actMBaseSize, constInfo.nBufferMBaseSize); + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + cubeBlock.ComputeMm1(info, mSplitInfo); + CrossCoreSetFlag(constInfo.syncC1V1); + } +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::ComputeMm2(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + CrossCoreWaitFlag(constInfo.syncV1C2); + cubeBlock.ComputeMm2(info, mSplitInfo); + CrossCoreSetFlag(constInfo.syncC2V2); + } +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::Process() +{ + uint32_t hasLoad = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, FA_CORE_ENABLE_INDEX, false)); + if (hasLoad == 0) { + return; + } + if ASCEND_IS_AIV { + vectorBlock.AllocEventID(); + vectorBlock.InitSoftmaxDefaultBuffer(); + } else { + cubeBlock.AllocEventID(); + } + ProcessBalance(); + if ASCEND_IS_AIV { + vectorBlock.FreeEventID(); + } else { + cubeBlock.FreeEventID(); + } +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx) +{ + bIdx = bN2Idx / kvHeadNum; + n2Idx = bN2Idx % kvHeadNum; +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::ProcessBalance() +{ + RunInfo extraInfo[SAS_PRELOAD_TASK_CACHE_SIZE]; + uint32_t gloop = 0; + uint32_t cmpLoop = 0; + uint32_t gS1LoopEnd = 0; + bool globalLoopStart = true; + // 适配左闭右开 + if (constInfo.bN2Start == constInfo.bN2End) { + if (constInfo.gS1Start != constInfo.gS1End || constInfo.s2Start != constInfo.s2End) { + constInfo.bN2End += 1; + } + } else if ((constInfo.gS1End != 0) || (constInfo.s2End != 0)) { + constInfo.bN2End += 1; + } + for (uint32_t bN2LoopIdx = constInfo.bN2Start; bN2LoopIdx < constInfo.bN2End; bN2LoopIdx++) { + GetBN2Idx(bN2LoopIdx, tempLoopInfo.bIdx, tempLoopInfo.n2Idx); + tempLoopInfo.actS1Size = GetActualSeqLenQ(tempLoopInfo.bIdx); // 获取actualSeqLength + bool isS1ZeroAndLastBatch = (tempLoopInfo.actS1Size == 0) && ((constInfo.outputLayout == SAS_LAYOUT::BSND) || + (bN2LoopIdx + 1 == constInfo.bN2End)); + uint32_t gS1SplitNum = CeilDiv(tempLoopInfo.actS1Size * constInfo.gSize, constInfo.mBaseSize); + + // 当处于最后一个BN2时, 且gS1End为0时, 说明当前BN2里的所有数据都在当前核处理 + gS1LoopEnd = (bN2LoopIdx == constInfo.bN2End - 1 && constInfo.gS1End != 0) ? constInfo.gS1End : gS1SplitNum; + // 当处于最后一个BN2且当前S1为0时，需要进入循环计算preload导致的未完成的部分 + gS1LoopEnd = isS1ZeroAndLastBatch ? gS1LoopEnd + 1 : gS1LoopEnd; + for (uint32_t gS1LoopIdx = constInfo.gS1Start; gS1LoopIdx < gS1LoopEnd; gS1LoopIdx++) { + tempLoopInfo.actOriS2Size = GetActualSeqLenKV(tempLoopInfo.bIdx); + // 对于各轴上的真实的idx, 采用左闭右闭的方案 + tempLoopInfo.gS1Idx = gS1LoopIdx * constInfo.mBaseSize; + tempLoopInfo.s1StartIdx = tempLoopInfo.gS1Idx / constInfo.gSize; + tempLoopInfo.s1EndIdx = + Min((tempLoopInfo.s1StartIdx + constInfo.mBaseSize / constInfo.gSize - 1), tempLoopInfo.actS1Size - 1); + // 此处均为闭区间 + tempLoopInfo.oriMaskRight = tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size + + static_cast(tempLoopInfo.s1EndIdx) + constInfo.oriWinRight; + tempLoopInfo.oriMaskLeft = Max(tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size + + static_cast(tempLoopInfo.s1EndIdx) - constInfo.oriWinLeft, + 0); + if (constInfo.templateMode == CFA_TEMPLATE) { + tempLoopInfo.cmpMaskRight = tempLoopInfo.actOriS2Size - tempLoopInfo.actS1Size; + } + GetSparseActualSeqLen(); + UpdateInnerLoopCond(); + uint32_t oriSplitNum = 0; + uint32_t s2SplitNum = 0; + bool isEnd = (bN2LoopIdx + 1 == constInfo.bN2End) && (gS1LoopIdx + 1 == gS1LoopEnd); + if (tempLoopInfo.curActSeqLenIsZero) { + if ASCEND_IS_AIV { + InitAllZeroOutput(tempLoopInfo.bIdx, tempLoopInfo.s1StartIdx, tempLoopInfo.n2Idx); + } + if (!isEnd) { + continue; + } + } else { + oriSplitNum = CeilDiv(tempLoopInfo.oriMaskRight - tempLoopInfo.oriMaskLeft + 1, constInfo.s2BaseSize); + s2SplitNum = oriSplitNum; + if (constInfo.templateMode == CFA_TEMPLATE) { + uint32_t cmpSplitNum = CeilDiv(tempLoopInfo.actCmpS2Size, constInfo.s2BaseSize); + s2SplitNum = oriSplitNum + cmpSplitNum; + tempLoopInfo.cmpLoopTimes = cmpSplitNum; + } + } + tempLoopInfo.s2LoopTimes = s2SplitNum; + tempLoopInfo.oriLoopTimes = oriSplitNum; + uint32_t s2LoopEnd = (isEnd && constInfo.s2End != 0) ? constInfo.s2End : tempLoopInfo.s2LoopTimes; + tempLoopInfo.s2LoopTimes = s2LoopEnd; + // 分核修改后需要打开 + // 当前s2是否被切，决定了输出是否要写到attenOut上 + tempLoopInfo.tndIsS2SplitCore = ((constInfo.s2Start == 0) && (s2LoopEnd == s2SplitNum)) ? false : true; + tempLoopInfo.tndCoreStartKVSplitPos = globalLoopStart ? constInfo.coreStartKVSplitPos : 0; + uint32_t extraLoop = isEnd ? PRELOAD_NUM : 0; + + for (uint32_t s2LoopIdx = constInfo.s2Start; s2LoopIdx < (s2LoopEnd + extraLoop); s2LoopIdx++) { + // PreloadPipeline loop初始值要求为 PRELOAD_NUM + PreloadPipeline(gloop, cmpLoop, constInfo.s2Start, s2LoopIdx, extraInfo); + ++gloop; + } + globalLoopStart = false; + constInfo.s2Start = 0; + } + constInfo.gS1Start = 0; + } +} + +template +__aicore__ inline void SparseAttnSharedkvSwa::PreloadPipeline(uint32_t loop, uint32_t cmpLoop, uint64_t s2Start, + uint64_t s2LoopIdx, + RunInfo extraInfo[SAS_PRELOAD_TASK_CACHE_SIZE]) +{ + RunInfo &extraInfo0 = extraInfo[loop % SAS_PRELOAD_TASK_CACHE_SIZE]; // 本轮任务 + RunInfo &extraInfo2 = extraInfo[(loop + 2) % SAS_PRELOAD_TASK_CACHE_SIZE]; // 上一轮任务 + RunInfo &extraInfo1 = extraInfo[(loop + 1) % SAS_PRELOAD_TASK_CACHE_SIZE]; // 上两轮任务 + + CalcParams(loop, cmpLoop, s2Start, s2LoopIdx, extraInfo0); + if (extraInfo0.isValid) { + if ASCEND_IS_AIC { + ComputeMm1(extraInfo0); + } + } + if (extraInfo2.isValid) { + if ASCEND_IS_AIV { + vectorBlock.ProcessVec1L(extraInfo2); + } + if ASCEND_IS_AIC { + ComputeMm2(extraInfo2); + } + } + if (extraInfo1.isValid) { + if ASCEND_IS_AIV { + vectorBlock.ProcessVec2L(extraInfo1); + } + extraInfo1.isValid = false; + } +} +} // namespace SASKernel +#endif // SPARSE_ATTN_SHAREDKV_SWA_KERNEL_H \ No newline at end of file diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/sparse_attn_sharedkv.cpp b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/sparse_attn_sharedkv.cpp new file mode 100644 index 000000000..b1b0a0dc5 --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/sparse_attn_sharedkv.cpp @@ -0,0 +1,72 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv.cpp + * \brief + */ + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" +#include "sparse_attn_sharedkv_template_tiling_key.h" +#include "arch32/sparse_attn_sharedkv_scfa_kernel.h" +#include "arch32/sparse_attn_sharedkv_swa_kernel.h" +#include "sparse_attn_sharedkv_metadata.h" + +using namespace AscendC; +using namespace optiling::detail; +using namespace SASKernel; + +#define SAS_OP_IMPL(templateClass, tilingdataClass, ...) \ + do { \ + templateClass> op; \ + GET_TILING_DATA_WITH_STRUCT(tilingdataClass, tiling_data_in, tiling); \ + const tilingdataClass *__restrict tiling_data = &tiling_data_in; \ + op.Init(query, oriKV, cmpKV, cmpSparseIndices, oriBlockTable, cmpBlockTable, cuSeqlensQ, \ + cuSeqlensOriKv, cuSeqlensCmpKv, seqUsedQ, seqUsedKV, \ + sinks, metadata, attentionOut, softmaxLse, user, tiling_data, tiling, &tPipe); \ + op.Process(); \ + } while (0) + + +template +__global__ __aicore__ void +sparse_attn_sharedkv(__gm__ uint8_t *query, __gm__ uint8_t *oriKV, __gm__ uint8_t *cmpKV, + __gm__ uint8_t *oriSparseIndices, __gm__ uint8_t *cmpSparseIndices, __gm__ uint8_t *oriBlockTable, + __gm__ uint8_t *cmpBlockTable, __gm__ uint8_t *cuSeqlensQ, __gm__ uint8_t *cuSeqlensOriKv, + __gm__ uint8_t *cuSeqlensCmpKv, __gm__ uint8_t *seqUsedQ, __gm__ uint8_t *seqUsedKV, + __gm__ uint8_t *sinks, __gm__ uint8_t *metadata, __gm__ uint8_t *attentionOut, + __gm__ uint8_t *softmaxLse, __gm__ uint8_t *workspace, __gm__ uint8_t *tiling) +{ + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + + TPipe tPipe; + __gm__ uint8_t *user = GetUserWorkspace(workspace); + if constexpr (ORIG_DTYPE_Q == DT_FLOAT16 && ORIG_DTYPE_ORI_KV == DT_FLOAT16 && ORIG_DTYPE_ATTN_OUT == DT_FLOAT16) { + if constexpr (TEMPLATE_MODE == SCFA_TEMPLATE) { + SAS_OP_IMPL(SparseAttnSharedkvScfa, SparseAttnSharedkvTilingData, half, half, half, FLASH_DECODE, + static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), TEMPLATE_MODE); + } else { + SAS_OP_IMPL(SparseAttnSharedkvSwa, SparseAttnSharedkvTilingData, half, half, half, FLASH_DECODE, + static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), TEMPLATE_MODE); + } + } + if constexpr (ORIG_DTYPE_Q == DT_BF16 && ORIG_DTYPE_ORI_KV == DT_BF16 && ORIG_DTYPE_ATTN_OUT == DT_BF16) { + if constexpr (TEMPLATE_MODE == SCFA_TEMPLATE) { + SAS_OP_IMPL(SparseAttnSharedkvScfa, SparseAttnSharedkvTilingData, bfloat16_t, bfloat16_t, bfloat16_t, + FLASH_DECODE, static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), + TEMPLATE_MODE); + } else { + SAS_OP_IMPL(SparseAttnSharedkvSwa, SparseAttnSharedkvTilingData, bfloat16_t, bfloat16_t, bfloat16_t, + FLASH_DECODE, static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), + TEMPLATE_MODE); + } + } +} \ No newline at end of file diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/sparse_attn_sharedkv_common.h b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/sparse_attn_sharedkv_common.h new file mode 100644 index 000000000..c353d7b13 --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/sparse_attn_sharedkv_common.h @@ -0,0 +1,325 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_common.h + * \brief + */ + +#ifndef SPARSE_ATTN_SHAREDKV_COMMON_H +#define SPARSE_ATTN_SHAREDKV_COMMON_H + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" + +namespace SASKernel { +using namespace AscendC; +// 将isCheckTiling设置为false, 输入输出的max&sum&exp的shape为(m, 1) +constexpr SoftmaxConfig SAS_SOFTMAX_FLASHV2_CFG_WITHOUT_BRC = {false, 0, 0, SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC}; + +enum class SAS_RUN_MODE { + SWA_MODE = 0, + SCFA_MODE = 1, + CFA_MODE = 2, +}; + +enum class SAS_LAYOUT { + BSND = 0, + TND = 1, + PA_ND = 2 +}; + +template +struct SASType { + using queryType = Q_T; + using kvType = KV_T; + using outputType = OUT_T; + static constexpr bool flashDecode = FLASH_DECODE; + static constexpr SAS_LAYOUT layout = LAYOUT_T; + static constexpr SAS_LAYOUT kvLayout = KV_LAYOUT_T; + static constexpr bool pageAttention = (KV_LAYOUT_T == SAS_LAYOUT::PA_ND); + static constexpr int templateMode = TEMPLATE_MODE; +}; + +// ================================Util functions================================== +template +__aicore__ inline T1 SASAlign(T1 num, T2 rnd) +{ + return (rnd == 0) ? 0 : ((num + rnd - 1) / rnd * rnd); +} + +template +__aicore__ inline T1 CeilDiv(T1 num, T2 rnd) +{ + return (rnd == 0) ? 0 : ((num + rnd - 1) / rnd); +} + +template +__aicore__ inline T1 Min(T1 a, T2 b) +{ + return (a > b) ? b : a; +} + +template +__aicore__ inline T1 Max(T1 a, T2 b) +{ + return (a > b) ? a : b; +} + +template +__aicore__ inline size_t BlockAlign(size_t s) +{ + if constexpr (IsSameType::value) { + return (s + 63) / 64 * 64; + } + size_t n = (32 / sizeof(T)); + return (s + n - 1) / n * n; +} + +struct PAShape { + uint32_t blockSize; + uint32_t headNum; // 一般为kv的head num，对应n2 + uint32_t headDim; // 512 对应d + uint32_t kvStride; + uint32_t maxblockNumPerBatch; // block table 每一行的最大个数 + uint32_t actHeadDim; // 实际拷贝col大小,考虑到N切块 s*d, 对应d + uint32_t copyRowNum; // 总共要拷贝的行数 + uint32_t copyRowNumAlign; +}; + +struct Position { + uint32_t bIdx; + uint32_t n2Idx; + uint32_t s2Idx; + uint32_t dIdx; + uint32_t s1Idx; +}; + +// 场景：query、key、value GM to L1 +// GM按ND格式存储 +// L1按NZ格式存储 +// GM的行、列、列的stride +template +__aicore__ inline void DataCopyGmNDToL1(LocalTensor &l1Tensor, GlobalTensor &gmTensor, uint32_t rowAct, + uint32_t rowAlign, + uint32_t col, // D + uint32_t colStride) // D or N*D +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = rowAct; // nd矩阵的行数 + // T为int4场景下，dValue = col / 2，srcDValue = colStride / 2 + nd2nzPara.dValue = col; // nd矩阵的列数 + nd2nzPara.srcDValue = colStride; // 同一nd矩阵相邻行起始地址间的偏移 + nd2nzPara.dstNzC0Stride = rowAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(l1Tensor, gmTensor, nd2nzPara); +} + +/* + 适用PA数据从GM拷贝到L1，支持ND、NZ数据； + PA的layout分 BNBD（blockNum,N,blockSize,D） BBH（blockNum,blockSize,N*D + BSH\BSND\TND 为BBH + shape.copyRowNumAlign 需要16字节对齐，如拷贝k矩阵，一次拷贝128*512，遇到尾块 10*512 需对齐到16*512 +*/ +template +__aicore__ inline void DataCopyPA(LocalTensor &dstTensor, //l1 + GlobalTensor &srcTensor, //gm + GlobalTensor &blockTableGm, + const PAShape &shape, // blockSize, headNum, headDim + const Position &startPos) // bacthIdx nIdx curSeqIdx +{ + uint32_t copyFinishRowCnt = 0; + uint64_t blockTableBaseOffset = startPos.bIdx * shape.maxblockNumPerBatch; + uint32_t curS2Idx = startPos.s2Idx; + uint32_t blockElementCnt = 32 / sizeof(T); + while (copyFinishRowCnt < shape.copyRowNum) { + uint64_t blockIdOffset = curS2Idx / shape.blockSize; // 获取block table上的索引 + uint64_t reaminRowCnt = curS2Idx % shape.blockSize; // 获取在单个块上超出的行数 + uint64_t idInBlockTable = + blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset); // 从block table上的获取编号 + uint32_t copyRowCnt = shape.blockSize - reaminRowCnt; // 一次只能处理一个Block + if (copyFinishRowCnt + copyRowCnt > shape.copyRowNum) { + copyRowCnt = shape.copyRowNum - copyFinishRowCnt; // 一个block未拷满 + } + // uint64_t offset = idInBlockTable * shape.blockSize * shape.headNum * shape.headDim; // PA的偏移 + uint64_t offset = idInBlockTable * shape.kvStride; // PA的偏移 + uint64_t dStride = shape.headDim; + offset += (uint64_t)(startPos.n2Idx * shape.headDim * shape.blockSize) + + reaminRowCnt * shape.headDim + startPos.dIdx; + + uint32_t dValue = shape.actHeadDim; + uint32_t srcDValue = dStride; + LocalTensor tmpDstTensor = dstTensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpSrcTensor = srcTensor[offset]; + DataCopyGmNDToL1(tmpDstTensor, tmpSrcTensor, copyRowCnt, shape.copyRowNumAlign, dValue, srcDValue); + copyFinishRowCnt += copyRowCnt; + curS2Idx += copyRowCnt; + } +} + +struct RunInfo { + uint32_t loop = 0; + uint32_t cmpLoop = 0; // 用于判断取 用于merge的4块GM 中的哪一块 + uint32_t bIdx = 0; + uint32_t gIdx = 0; + uint32_t s1Idx = 0; + uint32_t s2Idx = 0; + uint32_t n2IdxReal = 0; + uint32_t relativeS2Idx = 0; + uint32_t bn2IdxInCurCore = 0; + uint32_t curSInnerLoopTimes = 0; + uint64_t tndBIdxOffsetForQ = 0; + uint64_t tndBIdxOffsetForKV = 0; + uint64_t tensorCmpBOffset = 0; + uint64_t tensorAOffset = 0; + uint64_t tensorBOffset = 0; + uint64_t attenOutOffset = 0; + uint64_t attenMaskOffset = 0; + uint64_t topKBaseOffset = 0; + uint32_t actualSingleProcessSInnerSize = 0; + uint32_t actualSingleProcessSInnerSizeAlign = 0; + bool isFirstSInnerLoop = false; + uint32_t s2BatchOffset = 0; + uint32_t gSize = 0; + uint32_t s1Size = 0; + uint32_t s2Size = 0; + uint32_t mSize = 0; + uint32_t mSizeV = 0; + uint32_t mSizeVStart = 0; + uint32_t tndIsS2SplitCore = 0; + uint32_t tndCoreStartKVSplitPos = 0; + bool isBmm2Output = false; + bool isValid = false; + + static constexpr uint32_t n2Idx = 0; + uint64_t actS1Size = 1; + uint64_t actS2SizeOri = 0ULL; + uint32_t gS1Idx = 0; + uint64_t actS2Size = 1; + uint64_t actOriS2Size = 1; + uint32_t actMBaseSize = 0; + bool isLastS2Loop = 0; + int32_t nextTokensPerBatch = 0; + int64_t threshold = 0; + uint32_t curTopKIdx = 0; + uint64_t curOffsetInSparseBlock = 0; + bool isOri = true; // 判断当前块是在Ori部分还是Cmp部分 + uint64_t s2StartPoint = 0; + int64_t cmpS2IdLimit = 0; + int32_t v0S2DealSize = 0; + int32_t v0S2Start = 0; +}; + +struct ConstInfo { + // CUBE与VEC核间同步的模式 + static constexpr uint32_t SAS_SYNC_MODE2 = 2; + // BUFFER的字节数 + static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; + static constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64; + static constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256; + static constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512; + static constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024; + static constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048; + static constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096; + static constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192; + static constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384; + static constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768; + // FP32的0值和极大值 + static constexpr float FLOAT_ZERO = 0; + static constexpr float FLOAT_MAX = 3.402823466e+38F; + + // preLoad的总次数 + uint32_t preLoadNum = 0U; + uint32_t nBufferMBaseSize = 0U; + // CUBE和VEC的核间同步EventID + uint32_t syncV0C1 = 0U; + uint32_t syncC1V1 = 0U; + uint32_t syncV1C2 = 0U; + uint32_t syncC2V2 = 0U; + + uint32_t mmResUbSize = 0U; // Matmul1输出结果GM上的大小 + uint32_t vec1ResUbSize = 0U; // Vector1输出结果GM上的大小 + uint32_t bmm2ResUbSize = 0U; // Matmul2输出结果GM上的大小 + uint64_t batchSize = 0ULL; + uint64_t gSize = 0ULL; + uint64_t qHeadNum = 0ULL; + uint64_t kvHeadNum = 0; + uint64_t headDim = 0; + uint64_t kvSeqSize = 0ULL; // kv最大S长度 + uint64_t qSeqSize = 1ULL; // q最大S长度 + int64_t kvCacheBlockSize = 0; // PA场景的block size + uint64_t paCmpBlockSize = 0; + uint64_t paOriBlockSize = 0; + int64_t orikvCacheBlockSize = 0; + int64_t cmpkvCacheBlockSize = 0; + uint32_t oriMaxBlockNumPerBatch = 0; // PA场景的最大单batch block number + uint32_t cmpMaxBlockNumPerBatch = 0; + uint32_t splitKVNum = 0U; // S2核间切分的切分份数 + SAS_LAYOUT outputLayout; // 输出的Transpose格式 + uint32_t oriMaskMode = 0; + uint32_t cmpMaskMode = 0; + uint32_t oriKvStride = 0; + uint32_t cmpKvStride = 0; + bool needInit = false; + uint32_t templateMode = 0; + + // FlashDecoding + uint32_t actualCombineLoopSize = 0U; // FlashDecoding场景, S2在核间切分的最大份数 + uint64_t combineLseOffset = 0ULL; + uint64_t combineAccumOutOffset = 0ULL; + + uint32_t actualLenDimsQ = 0U; // query的actualSeqLength 的维度 + uint32_t actualLenDimsKV = 0U; // KV 的actualSeqLength 的维度 + + // TND + uint32_t s2Start = 0U; // TND场景下，S2的起始位置 + uint32_t s2End = 0U; // 单核TND场景下S2循环index上限 + + uint32_t bN2Start = 0U; + uint32_t bN2End = 0U; + uint32_t gS1Start = 0U; + uint32_t gS1End = 0U; + + uint32_t tndFDCoreArrLen = 0U; // TNDFlashDecoding相关分核信息array的长度 + uint32_t coreStartKVSplitPos = 0U; // TNDFlashDecoding kv起始位置 + + uint32_t mBaseSize = 1ULL; + uint32_t s2BaseSize = 1ULL; + + // sparse attr + int64_t sparseBlockSize = 0; + uint32_t sparseBlockCount = 0; + + // cmp attr + int64_t cmpRatio = 0; + + // win + int32_t oriWinRight = 0; + int32_t oriWinLeft = 128; + + // 是否返回SoftmaxLse + bool returnSoftmaxLse = false; +}; + +struct MSplitInfo { + uint32_t nBufferIdx = 0U; + uint32_t nBufferStartM = 0U; + uint32_t nBufferDealM = 0U; + uint32_t vecStartM = 0U; + uint32_t vecDealM = 0U; +}; +} // namespace SASKernel +#endif // SPARSE_ATTN_SHAREDKV_COMMON_H \ No newline at end of file diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/sparse_attn_sharedkv_metadata.h b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/sparse_attn_sharedkv_metadata.h new file mode 100644 index 000000000..3194141be --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/sparse_attn_sharedkv_metadata.h @@ -0,0 +1,79 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_metadata.h + * \brief + */ + +#ifndef SPARSE_ATTN_SHAREDKV_METADATA_H +#define SPARSE_ATTN_SHAREDKV_METADATA_H + +#include + +namespace optiling { + +// Constants +constexpr uint32_t AIC_CORE_NUM = 36; +constexpr uint32_t AIV_CORE_NUM = 72; +constexpr uint32_t SAS_META_SIZE = 1024; +using SAS_METADATA_T = int32_t; + +constexpr uint32_t FA_METADATA_SIZE = 8; +constexpr uint32_t FD_METADATA_SIZE = 8; + +// FA Metadata Index Definitions +constexpr uint32_t FA_CORE_ENABLE_INDEX = 0; +constexpr uint32_t FA_BN2_START_INDEX = 1; +constexpr uint32_t FA_M_START_INDEX = 2; +constexpr uint32_t FA_S2_START_INDEX = 3; +constexpr uint32_t FA_BN2_END_INDEX = 4; +constexpr uint32_t FA_M_END_INDEX = 5; +constexpr uint32_t FA_S2_END_INDEX = 6; +constexpr uint32_t FA_FIRST_FD_DATA_WORKSPACE_IDX_INDEX = 7; + +// FD Metadata Index Definitions +constexpr uint32_t FD_CORE_ENABLE_INDEX = 0; +constexpr uint32_t FD_BN2_IDX_INDEX = 1; +constexpr uint32_t FD_M_IDX_INDEX = 2; +constexpr uint32_t FD_WORKSPACE_IDX_INDEX = 3; +constexpr uint32_t FD_WORKSPACE_NUM_INDEX = 4; +constexpr uint32_t FD_M_START_INDEX = 5; +constexpr uint32_t FD_M_NUM_INDEX = 6; + +/** + * @brief 获取属性的绝对索引 + * @param coreIdx 核索引 + * @param metaIdx 元数据索引 + * @param isAIV 是否为AIV数据，默认为false + * @return 返回属性的绝对索引 + */ +#ifdef __CCE_AICORE__ +__aicore__ inline uint32_t GetAttrAbsIndex(uint32_t coreIdx, uint32_t metaIdx, bool isAIV = false) +{ + if (isAIV) { + return FA_METADATA_SIZE * AIC_CORE_NUM + FD_METADATA_SIZE * coreIdx + metaIdx; + } else { + return FA_METADATA_SIZE * coreIdx + metaIdx; + } +} +#endif + +namespace detail { +struct SasMetaData { + uint32_t faMetadata[AIC_CORE_NUM][FA_METADATA_SIZE]; + uint32_t fdMetadata[AIV_CORE_NUM][FD_METADATA_SIZE]; +}; +} // namespace detail + +static_assert(SAS_META_SIZE * sizeof(SAS_METADATA_T) >= sizeof(detail::SasMetaData)); +} // namespace optiling + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/sparse_attn_sharedkv_template_tiling_key.h b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/sparse_attn_sharedkv_template_tiling_key.h new file mode 100644 index 000000000..5119e89ac --- /dev/null +++ b/csrc/ascend/attention/sparse_attn_sharedkv/op_kernel/sparse_attn_sharedkv_template_tiling_key.h @@ -0,0 +1,59 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_template_tiling_key.h + * \brief + */ + +#ifndef SPARSE_ATTN_SHARED_TEMPLATE_TILING_KEY_H +#define SPARSE_ATTN_SHARED_TEMPLATE_TILING_KEY_H + +#include "ascendc/host_api/tiling/template_argument.h" + +#define SAS_LAYOUT_BSND 0 +#define SAS_LAYOUT_TND 1 +#define SAS_LAYOUT_PA_ND 2 + +#define ASCENDC_TPL_4_BW 4 + +#define SWA_TEMPLATE 0 +#define CFA_TEMPLATE 1 +#define SCFA_TEMPLATE 2 +// 模板参数支持的范围定义 +ASCENDC_TPL_ARGS_DECL(SparseAttnSharedkv, // 算子OpType + ASCENDC_TPL_BOOL_DECL(FLASH_DECODE, 0, 1), + ASCENDC_TPL_UINT_DECL(LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, SAS_LAYOUT_BSND, + SAS_LAYOUT_TND), + ASCENDC_TPL_UINT_DECL(KV_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, SAS_LAYOUT_PA_ND, SAS_LAYOUT_BSND, SAS_LAYOUT_TND), + ASCENDC_TPL_UINT_DECL(TEMPLATE_MODE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, SWA_TEMPLATE, + CFA_TEMPLATE, SCFA_TEMPLATE), ); + +// 支持的模板参数组合 +// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法 +ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, SAS_LAYOUT_BSND, SAS_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SAS_LAYOUT_PA_ND, SAS_LAYOUT_BSND, SAS_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, SWA_TEMPLATE), ), + + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, SAS_LAYOUT_BSND, SAS_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SAS_LAYOUT_PA_ND, SAS_LAYOUT_BSND, SAS_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, CFA_TEMPLATE), + ), + + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, SAS_LAYOUT_BSND, SAS_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SAS_LAYOUT_PA_ND, SAS_LAYOUT_BSND, SAS_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, SCFA_TEMPLATE), ), ); + +#endif // TEMPLATE_TILING_KEY \ No newline at end of file diff --git a/csrc/ascend/attention/sparse_flash_attention/CMakeLists.txt b/csrc/ascend/attention/sparse_flash_attention/CMakeLists.txt new file mode 100644 index 000000000..549b5c14c --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() diff --git a/csrc/ascend/attention/sparse_flash_attention/README.md b/csrc/ascend/attention/sparse_flash_attention/README.md new file mode 100644 index 000000000..e77aa38c8 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/README.md @@ -0,0 +1,232 @@ +# SparseFlashAttention + +## 产品支持情况 + +|产品 | 是否支持 | +|:----------------------------|:-----------:| +|Ascend 950PR/Ascend 950DT| √ | +|Atlas A3 训练系列产品/Atlas A3 推理系列产品| √ | +|Atlas A2 训练系列产品/Atlas A2 推理系列产品| √ | +|Atlas 200I/500 A2 推理产品| × | +|Atlas 推理系列产品| × | +|Atlas 训练系列产品| × | + +## 功能说明 + +- API功能：sparse_flash_attention（SFA）是针对大序列长度推理场景的高效注意力计算模块，该模块通过“只计算关键部分”大幅减少计算量，然而会引入大量的离散访存，造成数据搬运时间增加，进而影响整体性能。 + +- 计算公式： + + $$ + \text{softmax}(\frac{Q@\tilde{K}^T}{\sqrt{d_k}})@\tilde{V} + $$ + + 其中$\tilde{K},\tilde{V}$为基于某种选择算法（如`lightning_indexer`）得到的重要性较高的Key和Value，一般具有稀疏或分块稀疏的特征，$d_k$为$Q,\tilde{K}$每一个头的维度。 + 本次公布的`sparse_flash_attention`是面向Sparse Attention的全新算子，针对离散访存进行了指令缩减及搬运聚合的细致优化。 + +## 参数说明 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数名输入/输出/属性描述数据类型数据格式
query输入attention结构的Query输入，不支持空tensor和非连续。layout_query为BSND时，shape为(B,S1,N1,D)；layout_query为TND时，shape为(T1,N1,D)。FLOAT16、BFLOAT16ND
key输入attention结构的Key输入，不支持空tensor和非连续。layout_kv为PA_BSND时，shape为(block_num, block_size, KV_N, D)，其中block_num为PageAttention时block总数；layout_kv为BSND时，shape为(B, S2, KV_N, D)；layout_kv为TND时，shape为(T2, KV_N, D)。FLOAT16、BFLOAT16ND
value输入attention结构的Value输入，不支持空tensor和非连续，shape与key的shape一致。FLOAT16、BFLOAT16ND
sparseIndices输入离散取kvCache的索引，不支持空tensor和非连续。sparse_size为一次离散选取的block数，需要保证每行有效值均在前半部分、无效值均在后半部分，且sparse_size大于0。layout_query为BSND时，shape为(B, Q_S, KV_N, sparse_size)；layout_query为TND时，shape为(Q_T, KV_N, sparse_size)。INT32ND
blockTable输入表示PageAttention中kvCache存储使用的block映射表，不支持空tensor和非连续。第二维长度不小于所有batch中最大的S2对应的block数量，即S2_max / block_size向上取整；shape支持(B,S2/block_size)。INT32ND
actualSeqLengthsQuery输入表示不同Batch中query的有效token数，不支持空tensor和非连续。可传入None表示与query的S长度相同；支持长度为B的一维tensor，且每个Batch的有效token数不超过query中的维度S大小且不小于0。layout_query为TND时该入参必须传入，且以元素数量作为B值；每个元素表示当前batch与之前所有batch的token数总和。INT32ND
actualSeqLengthsKv输入表示不同Batch中key和value的有效token数，不支持空tensor和非连续。可传入None表示与key的S长度相同；支持长度为B的一维tensor，且每个Batch的有效token数不超过key/value中的维度S大小且不小于0。layout_kv为TND或PA_BSND时该入参必须传入；其中layout_kv为TND时，每个元素表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。INT32ND
queryRope输入表示MLA结构中的query的rope信息，不支持空tensor和非连续。layout_query为TND时，shape为(B,S1,N1,Dr)；layout_query为BSND时，shape为(T1,N1,Dr)。FLOAT16、BFLOAT16ND
keyRope输入表示MLA结构中的key的rope信息，不支持空tensor和非连续。layout_kv为TND时，shape为(B,S1,N1,Dr)；layout_kv为BSND时，shape为(T1,N1,Dr)；layout_kv为PA_BSND时，shape为(block_num,block_size,N2,Dr)。FLOAT16、BFLOAT16ND
scaleValue可选属性代表缩放系数。FLOAT16-
sparseBlockSize可选属性代表sparse阶段的block大小。sparse_block_size为1时，为Token-wise稀疏化场景；sparse_block_size大于1且小于等于128时，为Block-wise稀疏化场景，块内token共享相同的稀疏化决策。INT64-
layoutQuery可选属性标识输入query的数据排布格式，默认值为"BSND"，支持传入BSND和TND。STRING-
layoutKv可选属性标识输入key的数据排布格式，默认值为"BSND"，支持传入TND、BSND和PA_BSND，其中PA_BSND在使能PageAttention时使用。STRING-
sparseMode可选属性表示sparse的模式。sparse_mode为0时代表全部计算；sparse_mode为3时代表rightDownCausal模式的mask，对应以右下顶点往左上为划分线的下三角场景。INT64-
preTokens可选属性用于稀疏计算，表示attention需要和前几个Token计算关联，仅支持默认值2^63-1。INT64-
nextTokens可选属性用于稀疏计算，表示attention需要和后几个Token计算关联，仅支持默认值2^63-1。INT64-
attentionMode可选属性表示attention的模式，仅支持传入2，表示MLA-absorb模式。INT64-
returnSoftmaxLse可选属性用于表示是否返回softmax_max和softmax_sum。True表示返回，False表示不返回，默认值为False。该参数仅在训练且layout_kv不为PA_BSND场景支持。BOOL-
attentionOut输出公式中的输出，不支持空tensor和非连续。layout_query为BSND时，shape为(B,S1,N1,D)；layout_query为TND时，shape为(T1,N1,D)。FLOAT16、BFLOAT16ND
softmaxMaxOut输出Attention算法对query乘key的结果取max得到softmax_max，不支持空tensor和非连续。layout_query为BSND时，shape为(B,N2,S1,N1/N2)；layout_query为TND时，shape为(N2,T1,N1/N2)。FLOATND
softmaxSumOut输出Attention算法query乘key的结果减去softmax_max后取exp并求sum，得到softmax_sum，不支持空tensor和非连续。layout_query为BSND时，shape为(B,N2,S1,N1/N2)；layout_query为TND时，shape为(N2,T1,N1/N2)。FLOATND
+ +## 约束说明 + +- 该接口支持推理场景下使用。 +- 该接口支持图模式。 +- N1支持1/2/4/8/16/32/64/128。 +- block_size为一个block的token数，block_size取值为16的倍数，且最大支持1024。 +- 参数query中的D和key、value的D值相等为512，参数query_rope中的Dr和key_rope的Dr值相等为64。 +- 参数query、key、value的数据类型必须保持一致。 +- 当前只支持query_rope和key_rope传入，不支持rope为空。 +- 支持sparse_block_size整除block_size。 + - Ascend 950PR/Ascend 950DT： + - 只支持sparse_block_size为1。 + - Atlas A3 训练系列产品/Atlas A3 推理系列产品、Atlas A2 训练系列产品/Atlas A2 推理系列产品： + - 支持[1,128]，且要求是2的幂次方，在PageAttention场景下要求sparse_block_size整除block_size + +## 调用示例 + + + + + + + + + + + + + +
调用方式样例代码说明
aclnn接口 + test_aclnn_sparse_flash_attention + + + 通过 + aclnnSparseFlashAttention + + 接口方式调用算子 +
diff --git a/csrc/ascend/attention/sparse_flash_attention/docs/aclnnSparseFlashAttention.md b/csrc/ascend/attention/sparse_flash_attention/docs/aclnnSparseFlashAttention.md new file mode 100644 index 000000000..aa5fc60f9 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/docs/aclnnSparseFlashAttention.md @@ -0,0 +1,914 @@ +# aclnnSparseFlashAttention + +[📄 查看源码](https://gitcode.com/cann/ops-transformer/tree/master/attention/sparse_flash_attention) + +## 产品支持情况 + +| 产品 | 是否支持 | +| ------------------------------------------------------------ | :------: | +|Ascend 950PR/Ascend 950DT| √ | +|Atlas A2 推理系列产品 | √ | +|Atlas A3 推理系列产品 | √ | + +## 功能说明 + +- 接口功能：sparse_flash_attention（SFA）是针对大序列长度推理场景的高效注意力计算模块，该模块通过“只计算关键部分”大幅减少计算量，然而会引入大量的离散访存，造成数据搬运时间增加，进而影响整体性能。 + +- 计算公式： + +$$ +\text{softmax}(\frac{Q@\tilde{K}^T}{\sqrt{d_k}})@\tilde{V} +$$ + +其中$\tilde{K},\tilde{V}$为基于某种选择算法（如`lightning_indexer`）得到的重要性较高的Key和Value，一般具有稀疏或分块稀疏的特征，$d_k$为$Q,\tilde{K}$每一个头的维度。 + +## 函数原型 + +每个算子分为[两段式接口](../../../docs/zh/context/两段式接口.md)，必须先调用“aclnnSparseFlashAttentionGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnSparseFlashAttention”接口执行计算。 + +```Cpp +aclnnStatus aclnnSparseFlashAttentionGetWorkspaceSize( + const aclTensor *query, + const aclTensor *key, + const aclTensor *value, + const aclTensor *sparseIndices, + const aclTensor *blockTable, + const aclTensor *actualSeqLengthsQuery, + const aclTensor *actualSeqLengthsKv, + const aclTensor *queryRope, + const aclTensor *keyRope, + double scaleValue, + int64_t sparseBlockSize, + char *layoutQuery, + char *layoutKv, + int64_t sparseMode, + int64_t preTokens, + int64_t nextTokens, + int64_t attentionMode, + bool returnSoftmaxLse, + const aclTensor *attentionOutOut, + const aclTensor *softmaxMaxOut, + const aclTensor *softmaxSumOut, + uint64_t *workspaceSize, + aclOpExecutor **executor) +``` + +```Cpp +aclnnStatus aclnnSparseFlashAttention( + void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + const aclrtStream stream) +``` + +## aclnnSparseFlashAttentionGetWorkspaceSize + +- **参数说明：** + + > [!NOTE] + > + >- query、key、value参数维度含义：B（Batch Size）表示输入样本批量大小、S（Sequence Length）表示输入样本序列长度、H（Head Size）表示hidden层的大小、N（Head Num）表示多头数、D（Head Dim）表示hidden层最小的单元尺寸，且满足D=H/N、T表示所有Batch输入样本序列长度的累加和。 + >- Q\_S和S1表示query shape中的S，KV\_S和S2表示key shape中的S，Q\_N和N1表示num\_query\_heads，KV\_N和N2表示num\_key\_value\_heads，T1表示query shape中的T，T2表示key shape中的输入样本序列长度的累加和。 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensor
query（aclTensor）输入attention结构的Query输入。不支持空tensor。FLOAT16、BFLOAT16ND +
    +
  • layout_query为BSND时，shape为(B,S1,N1,D)。
    • +
  • layout_query为TND时，shape为(T1,N1,D)。
    • +
+ 		x
key（aclTensor）输入attention结构的Key输入 +
    +
  • 不支持空tensor。
    • +
  • block_num为PageAttention时block总数。
    • +
+ 		FLOAT16、BFLOAT16ND +
    +
  • layout_kv为PA_BSND时，shape为(block_num, block_size, KV_N, D)。
    • +
  • layout_kv为BSND时，shape为(B, S2, KV_N, D)。
    • +
  • layout_kv为TND时，shape为(T2, KV_N, D)。
    • +
+ 		x
value（aclTensor）输入attention结构的Value输入。不支持空tensor。FLOAT16、BFLOAT16NDshape与key的shape一致。x
sparseIndices（aclTensor）输入离散取kvCache的索引。 +
    +
  • 不支持空tensor。
    • +
  • sparse_size为一次离散选取的block数，需要保证每行有效值均在前半部分，无效值均在后半部分，且需要满足sparse_size大于0。
    • +
+ 		INT32ND +
    +
  • layout_query为BSND时，shape为(B, Q_S, KV_N, sparse_size)。
    • +
  • layout_query为TND时，shape为(Q_T, KV_N, sparse_size)。
    • +
+ 		x
blockTable（aclTensor）输入表示PageAttention中kvCache存储使用的block映射表。 +
    +
  • 不支持空tensor。
    • +
  • 第二维长度不小于所有batch中最大的S2对应的block数量，即S2_max / block_size向上取整。
    • +
+ 		INT32NDshape支持(B,S2/block_size)。x
actualSeqLengthsQuery（aclTensor）输入表示不同Batch中query的有效token数。 +
    +
  • 不支持空tensor。
    • +
  • 如果不指定seqlen可传入None，表示和query的shape的S长度相同。
    • +
  • 该入参中每个Batch的有效token数不超过query中的维度S大小且不小于0。支持长度为B的一维tensor。
    • +
  • layout_query为TND时，该入参必须传入，且以该入参元素的数量作为B值，该参数中每个元素的值表示当前batch与之前所有batch的token数总和。
    • +
+ 		INT32ND(B,)x
actualSeqLengthsKv（aclTensor）输入表示不同Batch中key和value的有效token数。 +
    +
  • 不支持空tensor。
    • +
  • 如果不指定seqlen可传入None，表示和key的shape的S长度相同。
    • +
  • 该参数中每个Batch的有效token数不超过key/value中的维度S大小且不小于0。支持长度为B的一维tensor。
    • +
  • 当layout_kv为TND或PA_BSND时，该入参必须传入。
    • +
  • layout_kv为TND，该参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。
    • +
+ 		INT32ND(B,)x
queryRope（aclTensor）输入表示MLA结构中的query的rope信息。不支持空tensor。FLOAT16、BFLOAT16ND +
    +
  • layout_query为TND时，shape为(B,S1,N1,Dr)。
    • +
  • layout_query为BSND时，shape为(T1,N1,Dr)。
    • +
+ 		x
keyRope（aclTensor）输入表示MLA结构中的key的rope信息。不支持空tensor。FLOAT16、BFLOAT16ND +
    +
  • layout_kv为TND时，shape为(B,S1,N1,Dr)。
    • +
  • layout_kv为BSND时，shape为(T1,N1,Dr)。
    • +
  • layout_kv为PA_BSND时，shape为(block_num,block_size,N2,Dr)。
    • +
+ 		x
scaleValue（double）输入代表缩放系数。-FLOAT16---
sparseBlockSize（int64_t）输入代表sparse阶段的block大小。 +
    +
  • sparse_block_size为1时，为Token-wise稀疏化场景，将每个token视为独立单元，在计算重要性分数时，评估每个查询token与每个键值token之间的独立关联程度。
    • +
  • sparse_block_size为大于1小于等于128时，为Block-wise稀疏化场景，将token序列划分为固定大小的连续块，以块为单位进行重要性评估，块内token共享相同的稀疏化决策。
    • +
+ 		INT64---
layoutQuery（char）输入标识输入query的数据排布格式。 +
    +
  • 用户不特意指定时可传入默认值"BSND"。
    • +
  • 支持传入BSND和TND。
    • +
+ 		STRING---
layoutKv（char）输入标识输入key的数据排布格式。 +
    +
  • 用户不特意指定时可传入默认值"BSND"。
    • +
  • 支持传入TND、BSND和PA_BSND，其中PA_BSND在使能PageAttention时使用。
    • +
+ 		STRING---
sparseMode（int64_t）输入表示sparse的模式。 +
    +
  • sparse_mode为0时，代表全部计算。
    • +
  • sparse_mode为3时，代表rightDownCausal模式的mask，对应以右下顶点往左上为划分线的下三角场景。
    • +
+ 		INT64---
preTokens（int64_t）输入用于稀疏计算，表示attention需要和前几个Token计算关联。仅支持默认值2^63-1。INT64---
nextTokens（int64_t）输入用于稀疏计算，表示attention需要和后几个Token计算关联。仅支持默认值2^63-1。INT64---
attentionMode（int64_t）输入-仅支持传入2，表示MLA-absorb模式。INT64---
returnSoftmaxLse（bool）输入用于表示是否返回softmax_max和softmax_sum。 +
    +
  • True表示返回，False表示不返回；默认值为False。
    • +
  • 该参数仅在训练且layout_kv不为PA_BSND场景支持。
    • +
+ 		BOOL---
attentionOut（aclTensor）输出公式中的输出。不支持空tensor。FLOAT16、BFLOAT16ND +
    +
  • layout_query为BSND时，shape为(B,S1,N1,D)。
    • +
  • layout_query为TND时shape为(T1,N1,D)。
    • +
+ 		x
softmaxMaxOut（aclTensor）输出Attention算法对query乘key的结果，取max得到softmax_max。不支持空tensor。FLOATND +
    +
  • layout_query为BSND时，shape为(B,N2,S1,N1/N2)。
    • +
  • layout_query为TND时shape为(N2,T1,N1/N2)。
    • +
+ 		x
softmaxSumOut（aclTensor）输出Attention算法query乘key的结果减去softmax_max, 再取exp，接着求sum，得到softmax_sum。不支持空tensor。FLOATND +
    +
  • layout_query为BSND时，shape为(B,N2,S1,N1/N2)。
    • +
  • layout_query为TND时shape为(N2,T1,N1/N2)。
    • +
+ 		x
workspaceSize（uint64_t*）输出返回需要在Device侧申请的workspace大小。-----
executor（aclOpExecutor）输出返回op执行器，包含了算子计算流程。-----
+ +- **返回值：** + + aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn返回码.md)。 + + 第一段接口会完成入参校验，出现以下场景时报错： + + + + + + + + + + + + + + + + + + + + + + + +
返回值错误码描述
ACLNN_ERR_PARAM_NULLPTR161001如果传入参数是必选输入，输出或者必选属性，且是空指针，则返回161001。
ACLNN_ERR_PARAM_INVALID161002query、key、value、sparseIndices、blockTable、actualSeqLengthsQuery、actualSeqLengthsKv、queryRope、keyRope、scaleValue、sparseBlockSize、layoutQuery、layoutKv、sparseMode、attentionMode、returnSoftmaxLse、attentionOut、softmaxMaxOut、softmaxSumOut的数据类型和数据格式不在支持的范围内。
+ +## aclnnSparseFlashAttention + +- **参数说明：** + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数名输入/输出描述
workspace输入在Device侧申请的workspace内存地址。
workspaceSize输入在Device侧申请的workspace大小，由第一段接口aclnnSparseFlashAttentionGetWorkspaceSize获取。
executor输入op执行器，包含了算子计算流程。
stream输入指定执行任务的Stream。
+ +- **返回值：** + + aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn返回码.md)。 + +## 约束说明 + +- 确定性计算：aclnnSparseFlashAttention默认确定性实现。 +- 该接口支持推理场景下使用。 +- N1支持1~64和128。 +- block_size为一个block的token数，block_size取值为16的倍数，且最大支持1024。 +- 参数query中的D和key、value的D值相等为512，参数query_rope中的Dr和key_rope的Dr值相等为64。 +- 参数query、key、value的数据类型必须保持一致。 +- 支持sparse_block_size整除block_size。 + - Ascend 950PR/Ascend 950DT： + - 只支持sparse_block_size为1。 + - Atlas A3 训练系列产品/Atlas A3 推理系列产品、Atlas A2 训练系列产品/Atlas A2 推理系列产品： + - 支持[1,128]，且要求是2的幂次方，在PageAttention场景下要求sparse_block_size整除block_size + +## 调用示例 + +示例代码如下，仅供参考，具体编译和执行过程请参考[编译与运行样例](../../../docs/zh/context/编译与运行样例.md)。 + +```Cpp +/** + * Copyright (c) 2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file test_incre_flash_attention_v4.cpp + * \brief + */ + +#include +#include +#include +#include +#include "securec.h" +#include "acl/acl.h" +#include "aclnnop/aclnn_sparse_flash_attention.h" + +using namespace std; + +namespace { + +#define CHECK_RET(cond) ((cond) ? true :(false)) + +#define LOG_PRINT(message, ...) \ + do { \ + (void)printf(message, ##__VA_ARGS__); \ + } while (0) + +int64_t GetShapeSize(const std::vector& shape) { + int64_t shapeSize = 1; + for (auto i : shape) { + shapeSize *= i; + } + return shapeSize; +} + +int Init(int32_t deviceId, aclrtStream* stream) { + auto ret = aclInit(nullptr); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclInit failed. ERROR: %d\n", ret); + return ret; + } + ret = aclrtSetDevice(deviceId); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); + return ret; + } + ret = aclrtCreateStream(stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); + return ret; + } + return 0; +} + +template +int CreateAclTensor(const std::vector& hostData, const std::vector& shape, void** deviceAddr, + aclDataType dataType, aclTensor** tensor) { + auto size = GetShapeSize(shape) * sizeof(T); + auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); + return ret; + } + + ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); + return ret; + } + + std::vector strides(shape.size(), 1); + for (int64_t i = shape.size() - 2; i >= 0; i--) { + strides[i] = shape[i + 1] * strides[i + 1]; + } + + *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, + shape.data(), shape.size(), *deviceAddr); + return 0; +} + +struct TensorResources { + void* queryDeviceAddr = nullptr; + void* keyDeviceAddr = nullptr; + void* valueDeviceAddr = nullptr; + void* sparseIndicesDeviceAddr = nullptr; + void* attentionOutDeviceAddr = nullptr; + void* softmaxMaxDeviceAddr = nullptr; + void* softmaxSumDeviceAddr = nullptr; + void* queryRopeDeviceAddr = nullptr; + void* keyRopeDeviceAddr = nullptr; + + aclTensor* queryTensor = nullptr; + aclTensor* keyTensor = nullptr; + aclTensor* valueTensor = nullptr; + aclTensor* sparseIndicesTensor = nullptr; + aclTensor* attentionOutTensor = nullptr; + aclTensor* softmaxMaxTensor = nullptr; + aclTensor* softmaxSumTensor = nullptr; + aclTensor* queryRopeTensor = nullptr; + aclTensor* keyRopeTensor = nullptr; +}; + +int InitializeTensors(TensorResources& resources) { + std::vector queryShape = {1, 2, 1, 512}; + std::vector keyShape = {1, 2, 1, 512}; + std::vector valueShape = {1, 2, 1, 512}; + std::vector sparseIndicesShape = {1, 2, 1, 2}; + std::vector attentionOutShape = {1, 2, 1, 512}; + std::vector softmaxMaxShape = {1, 2, 1, 16}; + std::vector softmaxSumShape = {1, 2, 1, 16}; + std::vector queryRopeShape = {1, 2, 1, 64}; + std::vector keyRopeShape = {1, 2, 1, 64}; + + int64_t queryShapeSize = GetShapeSize(queryShape); + int64_t keyShapeSize = GetShapeSize(keyShape); + int64_t valueShapeSize = GetShapeSize(valueShape); + int64_t sparseIndicesShapeSize = GetShapeSize(sparseIndicesShape); + int64_t attentionOutShapeSize = GetShapeSize(attentionOutShape); + int64_t softmaxMaxShapeSize = GetShapeSize(softmaxMaxShape); + int64_t softmaxSumShapeSize = GetShapeSize(softmaxSumShape); + int64_t queryRopeShapeSize = GetShapeSize(queryRopeShape); + int64_t keyRopeShapeSize = GetShapeSize(keyRopeShape); + + std::vector queryHostData(queryShapeSize, 1); + std::vector keyHostData(keyShapeSize, 1); + std::vector valueHostData(valueShapeSize, 1); + std::vector sparseIndicesHostData(sparseIndicesShapeSize, 1); + std::vector attentionOutHostData(attentionOutShapeSize, 1); + std::vector softmaxMaxHostData(softmaxMaxShapeSize, 1); + std::vector softmaxSumHostData(softmaxSumShapeSize, 1); + std::vector queryRopeHostData(queryRopeShapeSize, 1); + std::vector keyRopeHostData(keyRopeShapeSize, 1); + + // Create query aclTensor. + int ret = CreateAclTensor(queryHostData, queryShape, &resources.queryDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.queryTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create key aclTensor. + ret = CreateAclTensor(keyHostData, keyShape, &resources.keyDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.keyTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create value aclTensor. + ret = CreateAclTensor(valueHostData, valueShape, &resources.valueDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.valueTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create sparseIndices aclTensor. + ret = CreateAclTensor(sparseIndicesHostData, sparseIndicesShape, &resources.sparseIndicesDeviceAddr, + aclDataType::ACL_INT32, &resources.sparseIndicesTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create queryRope aclTensor. + ret = CreateAclTensor(queryRopeHostData, queryRopeShape, &resources.queryRopeDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.queryRopeTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create keyRope aclTensor. + ret = CreateAclTensor(keyRopeHostData, keyRopeShape, &resources.keyRopeDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.keyRopeTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create attention_out aclTensor. + ret = CreateAclTensor(attentionOutHostData, attentionOutShape, &resources.attentionOutDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.attentionOutTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create softmax_max aclTensor. + ret = CreateAclTensor(softmaxMaxHostData, softmaxMaxShape, &resources.softmaxMaxDeviceAddr, + aclDataType::ACL_FLOAT, &resources.softmaxMaxTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create softmax_sum aclTensor. + ret = CreateAclTensor(softmaxSumHostData, softmaxSumShape, &resources.softmaxSumDeviceAddr, + aclDataType::ACL_FLOAT, &resources.softmaxSumTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + return ACL_SUCCESS; +} + +int ExecuteSparseFlashAttention(TensorResources& resources, aclrtStream stream, + void** workspaceAddr, uint64_t* workspaceSize) { + int64_t d = 2; + double scaleValue = 1 / sqrt(d); + int64_t sparseBlockSize = 64; + constexpr const char layerOutStr[] = "BSND"; + constexpr size_t layerOutLen = sizeof(layerOutStr); + char layoutQuery[layerOutLen]; + char layoutKv[layerOutLen]; + errno_t memcpyRet = memcpy_s(layoutQuery, sizeof(layoutQuery), layerOutStr, layerOutLen); + if (memcpyRet != 0) { + LOG_PRINT("memcpy_s layoutQuery failed. ERROR: %d\n", memcpyRet); + return -1; + } + memcpyRet = memcpy_s(layoutKv, sizeof(layoutKv), layerOutStr, layerOutLen); + if (memcpyRet != 0) { + LOG_PRINT("memcpy_s layoutKv failed. ERROR: %d\n", memcpyRet); + return -1; + } + int64_t sparseMode = 3; + int64_t preTokens = 9223372036854775807; + int64_t nextTokens = 9223372036854775807; + int64_t attentionMode = 2; + bool returnSoftmaxLse = false; + aclOpExecutor* executor; + + int ret = aclnnSparseFlashAttentionGetWorkspaceSize(resources.queryTensor, resources.keyTensor, resources.valueTensor, resources.sparseIndicesTensor, nullptr, nullptr, nullptr, resources.queryRopeTensor, resources.keyRopeTensor, + scaleValue, sparseBlockSize, layoutQuery, layoutKv, sparseMode, preTokens, + nextTokens, attentionMode, returnSoftmaxLse, resources.attentionOutTensor, resources.softmaxMaxTensor, resources.softmaxSumTensor, workspaceSize, &executor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclnnSparseFlashAttentionGetWorkspaceSize failed. ERROR: %d\n", ret); + return ret; + } + + if (*workspaceSize > 0ULL) { + ret = aclrtMalloc(workspaceAddr, *workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); + return ret; + } + } + + ret = aclnnSparseFlashAttention(*workspaceAddr, *workspaceSize, executor, stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclnnSparseFlashAttention failed. ERROR: %d\n", ret); + return ret; + } + + return ACL_SUCCESS; +} + +int PrintOutResult(std::vector &shape, void** deviceAddr) { + auto size = GetShapeSize(shape); + std::vector resultData(size, 0); + auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), + *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); + return ret; + } + for (int64_t i = 0; i < size; i++) { + LOG_PRINT("mean result[%ld] is: %f\n", i, aclFloat16ToFloat(resultData[i])); + } + return ACL_SUCCESS; +} + +void CleanupResources(TensorResources& resources, void* workspaceAddr, + aclrtStream stream, int32_t deviceId) { + if (resources.queryTensor) { + aclDestroyTensor(resources.queryTensor); + } + if (resources.keyTensor) { + aclDestroyTensor(resources.keyTensor); + } + if (resources.valueTensor) { + aclDestroyTensor(resources.valueTensor); + } + if (resources.sparseIndicesTensor) { + aclDestroyTensor(resources.sparseIndicesTensor); + } + if (resources.attentionOutTensor) { + aclDestroyTensor(resources.attentionOutTensor); + } + if (resources.softmaxMaxTensor) { + aclDestroyTensor(resources.softmaxMaxTensor); + } + if (resources.softmaxSumTensor) { + aclDestroyTensor(resources.softmaxSumTensor); + } + if (resources.queryRopeTensor) { + aclDestroyTensor(resources.queryRopeTensor); + } + if (resources.keyRopeTensor) { + aclDestroyTensor(resources.keyRopeTensor); + } + + if (resources.queryDeviceAddr) { + aclrtFree(resources.queryDeviceAddr); + } + if (resources.keyDeviceAddr) { + aclrtFree(resources.keyDeviceAddr); + } + if (resources.valueDeviceAddr) { + aclrtFree(resources.valueDeviceAddr); + } + if (resources.sparseIndicesDeviceAddr) { + aclrtFree(resources.sparseIndicesDeviceAddr); + } + if (resources.attentionOutDeviceAddr) { + aclrtFree(resources.attentionOutDeviceAddr); + } + if (resources.softmaxMaxDeviceAddr) { + aclrtFree(resources.softmaxMaxDeviceAddr); + } + if (resources.softmaxSumDeviceAddr) { + aclrtFree(resources.softmaxSumDeviceAddr); + } + if (resources.queryRopeDeviceAddr) { + aclrtFree(resources.queryRopeDeviceAddr); + } + + if (resources.keyRopeDeviceAddr) { + aclrtFree(resources.keyRopeDeviceAddr); + } + + if (workspaceAddr) { + aclrtFree(workspaceAddr); + } + if (stream) { + aclrtDestroyStream(stream); + } + + aclrtResetDevice(deviceId); + aclFinalize(); +} + +} // namespace + +int main() { + + int32_t deviceId = 0; + aclrtStream stream = nullptr; + TensorResources resources = {}; + void* workspaceAddr = nullptr; + uint64_t workspaceSize = 0; + std::vector attentionOutShape = {1, 2, 1, 16}; + std::vector softmaxMaxShape = {1, 2, 1, 16}; + std::vector softmaxSumShape = {1, 2, 1, 16}; + int ret = ACL_SUCCESS; + + // 1. Initialize device and stream + ret = Init(deviceId, &stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("Init acl failed. ERROR: %d\n", ret); + return ret; + } + + + // 2. Initialize tensors + ret = InitializeTensors(resources); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + // 3. Execute the operation + ret = ExecuteSparseFlashAttention(resources, stream, &workspaceAddr, &workspaceSize); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + // 4. Synchronize stream + ret = aclrtSynchronizeStream(stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + // 5. Process results + printf("-----------attentionOut输出-----------\n"); + PrintOutResult(attentionOutShape, &resources.attentionOutDeviceAddr); + printf("-----------softmaxMax输出-----------\n"); + PrintOutResult(softmaxMaxShape, &resources.softmaxMaxDeviceAddr); + printf("-----------softmaxSum输出-----------\n"); + PrintOutResult(softmaxSumShape, &resources.softmaxSumDeviceAddr); + // 6. Cleanup resources + CleanupResources(resources, workspaceAddr, stream, deviceId); + return 0; +} +``` diff --git a/csrc/ascend/attention/sparse_flash_attention/examples/test_aclnn_sparse_flash_attention.cpp b/csrc/ascend/attention/sparse_flash_attention/examples/test_aclnn_sparse_flash_attention.cpp new file mode 100644 index 000000000..bc7d8561e --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/examples/test_aclnn_sparse_flash_attention.cpp @@ -0,0 +1,396 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file test_aclnn_sparse_flash_attention.cpp + * \brief + */ + +#include +#include +#include +#include +#include "securec.h" +#include "acl/acl.h" +#include "aclnnop/aclnn_sparse_flash_attention.h" + +using namespace std; + +namespace { + +#define CHECK_RET(cond) ((cond) ? true :(false)) + +#define LOG_PRINT(message, ...) \ + do { \ + (void)printf(message, ##__VA_ARGS__); \ + } while (0) + +int64_t GetShapeSize(const std::vector& shape) { + int64_t shapeSize = 1; + for (auto i : shape) { + shapeSize *= i; + } + return shapeSize; +} + +int32_t Init(int32_t deviceId, aclrtStream* stream) { + auto ret = aclInit(nullptr); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclInit failed. ERROR: %d\n", ret); + return ret; + } + ret = aclrtSetDevice(deviceId); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); + return ret; + } + ret = aclrtCreateStream(stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); + return ret; + } + return 0; +} + +template +int32_t CreateAclTensor(const std::vector& hostData, const std::vector& shape, void** deviceAddr, + aclDataType dataType, aclTensor** tensor) { + auto size = GetShapeSize(shape) * sizeof(T); + auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); + return ret; + } + + ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); + return ret; + } + + std::vector strides(shape.size(), 1); + for (int64_t i = shape.size() - 2; i >= 0; i--) { + strides[i] = shape[i + 1] * strides[i + 1]; + } + + *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, + shape.data(), shape.size(), *deviceAddr); + return 0; +} + +struct TensorResources { + void* queryDeviceAddr = nullptr; + void* keyDeviceAddr = nullptr; + void* valueDeviceAddr = nullptr; + void* sparseIndicesDeviceAddr = nullptr; + void* attentionOutDeviceAddr = nullptr; + void* softmaxMaxDeviceAddr = nullptr; + void* softmaxSumDeviceAddr = nullptr; + void* queryRopeDeviceAddr = nullptr; + void* keyRopeDeviceAddr = nullptr; + + aclTensor* queryTensor = nullptr; + aclTensor* keyTensor = nullptr; + aclTensor* valueTensor = nullptr; + aclTensor* sparseIndicesTensor = nullptr; + aclTensor* attentionOutTensor = nullptr; + aclTensor* softmaxMaxTensor = nullptr; + aclTensor* softmaxSumTensor = nullptr; + aclTensor* queryRopeTensor = nullptr; + aclTensor* keyRopeTensor = nullptr; +}; + +int32_t InitializeTensors(TensorResources& resources) { + std::vector queryShape = {1, 2, 1, 512}; + std::vector keyShape = {1, 2, 1, 512}; + std::vector valueShape = {1, 2, 1, 512}; + std::vector sparseIndicesShape = {1, 2, 1, 2}; + std::vector attentionOutShape = {1, 2, 1, 512}; + std::vector softmaxMaxShape = {1, 2, 1, 16}; + std::vector softmaxSumShape = {1, 2, 1, 16}; + std::vector queryRopeShape = {1, 2, 1, 64}; + std::vector keyRopeShape = {1, 2, 1, 64}; + + int64_t queryShapeSize = GetShapeSize(queryShape); + int64_t keyShapeSize = GetShapeSize(keyShape); + int64_t valueShapeSize = GetShapeSize(valueShape); + int64_t sparseIndicesShapeSize = GetShapeSize(sparseIndicesShape); + int64_t attentionOutShapeSize = GetShapeSize(attentionOutShape); + int64_t softmaxMaxShapeSize = GetShapeSize(softmaxMaxShape); + int64_t softmaxSumShapeSize = GetShapeSize(softmaxSumShape); + int64_t queryRopeShapeSize = GetShapeSize(queryRopeShape); + int64_t keyRopeShapeSize = GetShapeSize(keyRopeShape); + + std::vector queryHostData(queryShapeSize, 1); + std::vector keyHostData(keyShapeSize, 1); + std::vector valueHostData(valueShapeSize, 1); + std::vector sparseIndicesHostData(sparseIndicesShapeSize, 1); + std::vector attentionOutHostData(attentionOutShapeSize, 1); + std::vector softmaxMaxHostData(softmaxMaxShapeSize, 1); + std::vector softmaxSumHostData(softmaxSumShapeSize, 1); + std::vector queryRopeHostData(queryRopeShapeSize, 1); + std::vector keyRopeHostData(keyRopeShapeSize, 1); + + // Create query aclTensor. + int32_t ret = CreateAclTensor(queryHostData, queryShape, &resources.queryDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.queryTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create key aclTensor. + ret = CreateAclTensor(keyHostData, keyShape, &resources.keyDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.keyTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create value aclTensor. + ret = CreateAclTensor(valueHostData, valueShape, &resources.valueDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.valueTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create sparseIndices aclTensor. + ret = CreateAclTensor(sparseIndicesHostData, sparseIndicesShape, &resources.sparseIndicesDeviceAddr, + aclDataType::ACL_INT32, &resources.sparseIndicesTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create queryRope aclTensor. + ret = CreateAclTensor(queryRopeHostData, queryRopeShape, &resources.queryRopeDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.queryRopeTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create keyRope aclTensor. + ret = CreateAclTensor(keyRopeHostData, keyRopeShape, &resources.keyRopeDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.keyRopeTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create attention_out aclTensor. + ret = CreateAclTensor(attentionOutHostData, attentionOutShape, &resources.attentionOutDeviceAddr, + aclDataType::ACL_FLOAT16, &resources.attentionOutTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create softmax_max aclTensor. + ret = CreateAclTensor(softmaxMaxHostData, softmaxMaxShape, &resources.softmaxMaxDeviceAddr, + aclDataType::ACL_FLOAT, &resources.softmaxMaxTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + // Create softmax_sum aclTensor. + ret = CreateAclTensor(softmaxSumHostData, softmaxSumShape, &resources.softmaxSumDeviceAddr, + aclDataType::ACL_FLOAT, &resources.softmaxSumTensor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + return ret; + } + + return ACL_SUCCESS; +} + +int32_t ExecuteSparseFlashAttention(TensorResources& resources, aclrtStream stream, + void** workspaceAddr, uint64_t* workspaceSize) { + int64_t d = 2; + double scaleValue = 1 / sqrt(d); + int64_t sparseBlockSize = 64; + constexpr const char layerOutStr[] = "BSND"; + constexpr size_t layerOutLen = sizeof(layerOutStr); + char layoutQuery[layerOutLen]; + char layoutKv[layerOutLen]; + errno_t memcpyRet = memcpy_s(layoutQuery, sizeof(layoutQuery), layerOutStr, layerOutLen); + if (memcpyRet != 0) { + LOG_PRINT("memcpy_s layoutQuery failed. ERROR: %d\n", memcpyRet); + return -1; + } + memcpyRet = memcpy_s(layoutKv, sizeof(layoutKv), layerOutStr, layerOutLen); + if (memcpyRet != 0) { + LOG_PRINT("memcpy_s layoutKv failed. ERROR: %d\n", memcpyRet); + return -1; + } + int64_t sparseMode = 3; + int64_t preTokens = 9223372036854775807; + int64_t nextTokens = 9223372036854775807; + int64_t attentionMode = 2; + bool returnSoftmaxLse = false; + aclOpExecutor* executor; + + int32_t ret = aclnnSparseFlashAttentionGetWorkspaceSize(resources.queryTensor, resources.keyTensor, resources.valueTensor, resources.sparseIndicesTensor, nullptr, nullptr, nullptr, resources.queryRopeTensor, resources.keyRopeTensor, + scaleValue, sparseBlockSize, layoutQuery, layoutKv, sparseMode, preTokens, + nextTokens, attentionMode, returnSoftmaxLse, resources.attentionOutTensor, resources.softmaxMaxTensor, resources.softmaxSumTensor, workspaceSize, &executor); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclnnSparseFlashAttentionGetWorkspaceSize failed. ERROR: %d\n", ret); + return ret; + } + + if (*workspaceSize > 0ULL) { + ret = aclrtMalloc(workspaceAddr, *workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); + return ret; + } + } + + ret = aclnnSparseFlashAttention(*workspaceAddr, *workspaceSize, executor, stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclnnSparseFlashAttention failed. ERROR: %d\n", ret); + return ret; + } + + return ACL_SUCCESS; +} + +int32_t PrintOutResult(std::vector &shape, void** deviceAddr) { + auto size = GetShapeSize(shape); + std::vector resultData(size, 0); + auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), + *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); + return ret; + } + for (int64_t i = 0; i < size; i++) { + LOG_PRINT("mean result[%ld] is: %f\n", i, aclFloat16ToFloat(resultData[i])); + } + return ACL_SUCCESS; +} + +void CleanupResources(TensorResources& resources, void* workspaceAddr, + aclrtStream stream, int32_t deviceId) { + if (resources.queryTensor) { + aclDestroyTensor(resources.queryTensor); + } + if (resources.keyTensor) { + aclDestroyTensor(resources.keyTensor); + } + if (resources.valueTensor) { + aclDestroyTensor(resources.valueTensor); + } + if (resources.sparseIndicesTensor) { + aclDestroyTensor(resources.sparseIndicesTensor); + } + if (resources.attentionOutTensor) { + aclDestroyTensor(resources.attentionOutTensor); + } + if (resources.softmaxMaxTensor) { + aclDestroyTensor(resources.softmaxMaxTensor); + } + if (resources.softmaxSumTensor) { + aclDestroyTensor(resources.softmaxSumTensor); + } + if (resources.queryRopeTensor) { + aclDestroyTensor(resources.queryRopeTensor); + } + if (resources.keyRopeTensor) { + aclDestroyTensor(resources.keyRopeTensor); + } + + if (resources.queryDeviceAddr) { + aclrtFree(resources.queryDeviceAddr); + } + if (resources.keyDeviceAddr) { + aclrtFree(resources.keyDeviceAddr); + } + if (resources.valueDeviceAddr) { + aclrtFree(resources.valueDeviceAddr); + } + if (resources.sparseIndicesDeviceAddr) { + aclrtFree(resources.sparseIndicesDeviceAddr); + } + if (resources.attentionOutDeviceAddr) { + aclrtFree(resources.attentionOutDeviceAddr); + } + if (resources.softmaxMaxDeviceAddr) { + aclrtFree(resources.softmaxMaxDeviceAddr); + } + if (resources.softmaxSumDeviceAddr) { + aclrtFree(resources.softmaxSumDeviceAddr); + } + if (resources.queryRopeDeviceAddr) { + aclrtFree(resources.queryRopeDeviceAddr); + } + + if (resources.keyRopeDeviceAddr) { + aclrtFree(resources.keyRopeDeviceAddr); + } + + if (workspaceAddr) { + aclrtFree(workspaceAddr); + } + if (stream) { + aclrtDestroyStream(stream); + } + + aclrtResetDevice(deviceId); + aclFinalize(); +} + +} // namespace + +int32_t main() +{ + int32_t deviceId = 0; + aclrtStream stream = nullptr; + TensorResources resources = {}; + void* workspaceAddr = nullptr; + uint64_t workspaceSize = 0; + std::vector attentionOutShape = {1, 2, 1, 16}; + std::vector softmaxMaxShape = {1, 2, 1, 16}; + std::vector softmaxSumShape = {1, 2, 1, 16}; + int32_t ret = ACL_SUCCESS; + + // 1. Initialize device and stream + ret = Init(deviceId, &stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("Init acl failed. ERROR: %d\n", ret); + return ret; + } + + // 2. Initialize tensors + ret = InitializeTensors(resources); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + // 3. Execute the operation + ret = ExecuteSparseFlashAttention(resources, stream, &workspaceAddr, &workspaceSize); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + // 4. Synchronize stream + ret = aclrtSynchronizeStream(stream); + if (!CHECK_RET(ret == ACL_SUCCESS)) { + LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); + CleanupResources(resources, workspaceAddr, stream, deviceId); + return ret; + } + + // 5. Process results + printf("-----------attentionOut输出-----------\n"); + PrintOutResult(attentionOutShape, &resources.attentionOutDeviceAddr); + printf("-----------softmaxMax输出-----------\n"); + PrintOutResult(softmaxMaxShape, &resources.softmaxMaxDeviceAddr); + printf("-----------softmaxSum输出-----------\n"); + PrintOutResult(softmaxSumShape, &resources.softmaxSumDeviceAddr); + // 6. Cleanup resources + CleanupResources(resources, workspaceAddr, stream, deviceId); + return 0; +} diff --git a/csrc/ascend/attention/sparse_flash_attention/op_host/CMakeLists.txt b/csrc/ascend/attention/sparse_flash_attention/op_host/CMakeLists.txt new file mode 100644 index 000000000..e073395f6 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_host/CMakeLists.txt @@ -0,0 +1,30 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + set(sparse_flash_attention_depends attention/common CACHE INTERNAL "Dependencies for sparse_flash_attention") + target_sources(op_host_aclnnInner PRIVATE + sparse_flash_attention_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME SparseFlashAttention + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-vf-remove-membar=false + -mllvm -cce-aicore-hoist-movemask=false +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE sparse_flash_attention ACLNNTYPE aclnn_inner) +endif() diff --git a/csrc/ascend/attention/sparse_flash_attention/op_host/op_api/aclnn_sparse_flash_attention.cpp b/csrc/ascend/attention/sparse_flash_attention/op_host/op_api/aclnn_sparse_flash_attention.cpp new file mode 100644 index 000000000..6227f2bca --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_host/op_api/aclnn_sparse_flash_attention.cpp @@ -0,0 +1,145 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#include +#include "graph/types.h" +#include "aclnn_sparse_flash_attention.h" + +#include "opdev/make_op_executor.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/tensor_view_utils.h" +#include "opdev/op_def.h" +#include "opdev/op_log.h" +#include "opdev/shape_utils.h" +#include "opdev/common_types.h" +#include "opdev/data_type_utils.h" +#include "opdev/format_utils.h" + +using namespace op; + +#ifdef __cplusplus +extern "C" { +#endif + +namespace { + +extern aclnnStatus aclnnInnerSparseFlashAttentionGetWorkspaceSize( + const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *sparse_indices, + const aclTensor *blockTableOptional, const aclTensor *actualSeqLengthsQueryOptional, const aclTensor *actualSeqLengthsKvOptional, + const aclTensor *queryRopeOptional, const aclTensor *keyRopeOptional, double scaleValue, + int64_t sparseBlockSizeOptional, char *layoutQueryOptional, char *layoutKvOptional, + int64_t sparseMode, int64_t preTokens, int64_t nextTokens, int64_t attentionMode, + bool returnSoftmaxLse, const aclTensor *attentionOut, const aclTensor *softmaxMax, + const aclTensor *softmaxSum, uint64_t *workspaceSize, aclOpExecutor **executor); + +extern aclnnStatus aclnnInnerSparseFlashAttention(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, + const aclrtStream stream); + +class TensorHolder { +public: + TensorHolder(const aclTensor *&output, aclDataType dataType, std::string varName) { + inner_ = nullptr; + name_ = varName; + if (output == nullptr) { + std::vector shape = {0}; + int64_t addr = 0xff; + inner_ = aclCreateTensor(shape.data(), shape.size(), + dataType, shape.data(), 0, ACL_FORMAT_ND, + shape.data(), shape.size(), static_cast(&addr)); + output = inner_; + } + } + + ~TensorHolder() { + if (inner_) { + aclDestroyTensor(inner_); + inner_ = nullptr; + } + } + + void CheckTensorConditionalNotNull(bool conditional) const { + if (inner_ && conditional) { + OP_LOGW("Check %s != nullptr failed!", name_.c_str()); + } else if (!inner_ && !conditional) { + OP_LOGW("Check %s == nullptr failed!", name_.c_str()); + } + } + + bool IsTensorNotNull() const { + return inner_ == nullptr; + } + +private: + const aclTensor *inner_; + std::string name_; +}; + +aclnnStatus aclnnSparseFlashAttentionGetWorkspaceSize( + const aclTensor *query, + const aclTensor *key, + const aclTensor *value, + const aclTensor *sparseIndices, + const aclTensor *blockTableOptional, + const aclTensor *actualSeqLengthsQueryOptional, + const aclTensor *actualSeqLengthsKvOptional, + const aclTensor *queryRopeOptional, + const aclTensor *keyRopeOptional, + double scaleValue, + int64_t sparseBlockSizeOptional, + char *layoutQueryOptional, + char *layoutKvOptional, + int64_t sparseMode, + int64_t preTokens, + int64_t nextTokens, + int64_t attentionMode, + bool returnSoftmaxLse, + const aclTensor *attentionOut, + const aclTensor *softmaxMax, + const aclTensor *softmaxSum, + uint64_t *workspaceSize, + aclOpExecutor **executor) +{ + if (returnSoftmaxLse) { + if (softmaxMax == nullptr || softmaxSum == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "when returnSoftmaxLse is true, softmaxMax and softmaxSum cannot be nullptr."); + return ge::GRAPH_FAILED; + } + } else { + if (softmaxMax == nullptr && softmaxSum == nullptr) { + auto softmaxMaxHolder = TensorHolder(softmaxMax, aclDataType::ACL_FLOAT, std::string("softmaxMax")); + auto softmaxSumHolder = TensorHolder(softmaxSum, aclDataType::ACL_FLOAT, std::string("softmaxSum")); + if (softmaxMax == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Failed to create the holder of tensor softmaxMax!"); + return ge::GRAPH_FAILED; + } + if (softmaxSum == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Failed to create the holder of tensor softmaxSum!"); + return ge::GRAPH_FAILED; + } + } + } + return aclnnInnerSparseFlashAttentionGetWorkspaceSize( + query, key, value, sparseIndices, blockTableOptional, actualSeqLengthsQueryOptional, actualSeqLengthsKvOptional, queryRopeOptional, keyRopeOptional, + scaleValue, sparseBlockSizeOptional, layoutQueryOptional, layoutKvOptional, sparseMode, preTokens, + nextTokens, attentionMode, returnSoftmaxLse, attentionOut, + softmaxMax, softmaxSum, workspaceSize, executor); +} + +aclnnStatus aclnnSparseFlashAttention(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, + const aclrtStream stream) +{ + return aclnnInnerSparseFlashAttention(workspace, workspaceSize, executor, stream); +} + +} // namespace + +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/attention/sparse_flash_attention/op_host/op_api/aclnn_sparse_flash_attention.h b/csrc/ascend/attention/sparse_flash_attention/op_host/op_api/aclnn_sparse_flash_attention.h new file mode 100644 index 000000000..cc4e6f8e3 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_host/op_api/aclnn_sparse_flash_attention.h @@ -0,0 +1,63 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef ACLNN_SPARSE_FLASH_ATTENTION_H +#define ACLNN_SPARSE_FLASH_ATTENTION_H + +#include "aclnn/acl_meta.h" +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @brief The first interface of aclnnMlaPrologV2WeightNz calculates the workspace size based on the specific calculation process. + * @domain aclnn_ops_infer + */ +__attribute__((visibility("default"))) aclnnStatus aclnnSparseFlashAttentionGetWorkspaceSize( + const aclTensor *query, + const aclTensor *key, + const aclTensor *value, + const aclTensor *sparseIndices, + const aclTensor *blockTableOptional, + const aclTensor *actualSeqLengthsQueryOptional, + const aclTensor *actualSeqLengthsKvOptional, + const aclTensor *queryRopeOptional, + const aclTensor *keyRopeOptional, + double scaleValue, + int64_t sparseBlockSizeOptional, + char *layoutQueryOptional, + char *layoutKvOptional, + int64_t sparseMode, + int64_t preTokens, + int64_t nextTokens, + int64_t attentionMode, + bool returnSoftmaxLse, + const aclTensor *attentionOut, + const aclTensor *softmaxMax, + const aclTensor *softmaxSum, + uint64_t *workspaceSize, + aclOpExecutor **executor); + +/** + * @brief The second interface of ACLNN_SPARSE_FLASH_ATTENTION_H is used to perform calculations. + */ +__attribute__((visibility("default"))) aclnnStatus aclnnSparseFlashAttention(void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + const aclrtStream stream); + + +#ifdef __cplusplus +} +#endif + +#endif // ACLNN_SPARSE_FLASH_ATTENTION_H diff --git a/csrc/ascend/attention/sparse_flash_attention/op_host/sparse_flash_attention_def.cpp b/csrc/ascend/attention/sparse_flash_attention/op_host/sparse_flash_attention_def.cpp new file mode 100644 index 000000000..3a155320c --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_host/sparse_flash_attention_def.cpp @@ -0,0 +1,101 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_flash_attention_def.cpp + * \brief + */ + +#include "register/op_def_registry.h" + +namespace ops { +class SparseFlashAttention : public OpDef { +public: + explicit SparseFlashAttention(const char *name) : OpDef(name) + { + this->Input("query") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("key") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("value") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("sparse_indices") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("block_table") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("actual_seq_lengths_query") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("actual_seq_lengths_kv") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("query_rope") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("key_rope") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("attention_out") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("softmax_max") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("softmax_sum") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Attr("scale_value").AttrType(REQUIRED).Float(1.0); + this->Attr("sparse_block_size").AttrType(OPTIONAL).Int(1); + this->Attr("layout_query").AttrType(OPTIONAL).String("BSND"); + this->Attr("layout_kv").AttrType(OPTIONAL).String("BSND"); + this->Attr("sparse_mode").AttrType(OPTIONAL).Int(3); // 3:默认值，只计算下三角 + this->Attr("pre_tokens").AttrType(OPTIONAL).Int(INT64_MAX); + this->Attr("next_tokens").AttrType(OPTIONAL).Int(INT64_MAX); + this->Attr("attention_mode").AttrType(OPTIONAL).Int(2); + this->Attr("return_softmax_lse").AttrType(OPTIONAL).Bool(false); + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true); + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + } +}; +OP_ADD(SparseFlashAttention); +} // namespace ops diff --git a/csrc/ascend/attention/sparse_flash_attention/op_host/sparse_flash_attention_infershape.cpp b/csrc/ascend/attention/sparse_flash_attention/op_host/sparse_flash_attention_infershape.cpp new file mode 100644 index 000000000..9c1b0d532 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_host/sparse_flash_attention_infershape.cpp @@ -0,0 +1,130 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_flash_attention_proto.cpp + * \brief + */ + +#include +#include +#include "err/ops_err.h" + +using namespace ge; + +namespace ops { +constexpr size_t QUERY_INPUT_INDEX = 0; +constexpr size_t KEY_INPUT_INDEX = 1; + +constexpr uint32_t DIM_NUM_1 = 1; +constexpr uint32_t DIM_NUM_3 = 3; +constexpr uint32_t DIM_NUM_4 = 4; +constexpr uint32_t DIM_INDEX_0 = 0; +constexpr uint32_t DIM_INDEX_1 = 1; +constexpr uint32_t DIM_INDEX_2 = 2; +constexpr uint32_t DIM_INDEX_3 = 3; +constexpr uint32_t LAYOUT_KEY_ATTR_INDEX = 3; +constexpr uint32_t RETURN_SOFTMAX_LSE_INDEX = 8; + +constexpr uint32_t OUTPUT_INDEX_0 = 0; +constexpr uint32_t OUTPUT_INDEX_1 = 1; +constexpr uint32_t OUTPUT_INDEX_2 = 2; + +ge::graphStatus InferShapeSparseFlashAttention(gert::InferShapeContext *context) +{ + OP_CHECK_IF(context == nullptr, OP_LOGE("SparseFlashAttention", "InferShapeContext is nullptr"), + return ge::GRAPH_FAILED); + const gert::Shape *queryShape = context->GetInputShape(QUERY_INPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, queryShape); + + const gert::Shape *keyShape = context->GetInputShape(KEY_INPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, keyShape); + + gert::Shape *attentionOutShape = context->GetOutputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context, attentionOutShape); + *attentionOutShape = *queryShape; + + gert::Shape *softmaxMaxShape = context->GetOutputShape(1); + OP_CHECK_NULL_WITH_CONTEXT(context, softmaxMaxShape); + + gert::Shape *softmaxSumShape = context->GetOutputShape(2); + OP_CHECK_NULL_WITH_CONTEXT(context, softmaxSumShape); + + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const char *inputLayoutKeyPtr = attrs->GetAttrPointer(LAYOUT_KEY_ATTR_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, inputLayoutKeyPtr); + std::string inputLayoutKeyPtrStr = std::string(inputLayoutKeyPtr); + const bool *lse_flag = attrs->GetAttrPointer(RETURN_SOFTMAX_LSE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, lse_flag); + bool return_softmax_lse = (lse_flag != nullptr)? *lse_flag : false; + + if(return_softmax_lse){ + if(queryShape->GetDimNum() == DIM_NUM_3){ + if (inputLayoutKeyPtrStr == "PA_BSND") { + softmaxMaxShape->SetDimNum(DIM_NUM_3); + softmaxMaxShape->SetDim(DIM_INDEX_0, keyShape->GetDim(DIM_INDEX_2)); + softmaxMaxShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_0)); + softmaxMaxShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1) / keyShape->GetDim(DIM_INDEX_2)); + + softmaxSumShape->SetDimNum(DIM_NUM_3); + softmaxSumShape->SetDim(DIM_INDEX_0, keyShape->GetDim(DIM_INDEX_2)); + softmaxSumShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_0)); + softmaxSumShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1) / keyShape->GetDim(DIM_INDEX_2)); + } else { + softmaxMaxShape->SetDimNum(DIM_NUM_3); + softmaxMaxShape->SetDim(DIM_INDEX_0, keyShape->GetDim(DIM_INDEX_1)); + softmaxMaxShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_0)); + softmaxMaxShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1) / keyShape->GetDim(DIM_INDEX_1)); + + softmaxSumShape->SetDimNum(DIM_NUM_3); + softmaxSumShape->SetDim(DIM_INDEX_0, keyShape->GetDim(DIM_INDEX_1)); + softmaxSumShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_0)); + softmaxSumShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1) / keyShape->GetDim(DIM_INDEX_1)); + } + } else { + softmaxMaxShape->SetDimNum(DIM_NUM_4); + softmaxMaxShape->SetDim(DIM_INDEX_0, queryShape->GetDim(DIM_INDEX_0)); + softmaxMaxShape->SetDim(DIM_INDEX_1, keyShape->GetDim(DIM_INDEX_2)); + softmaxMaxShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1)); + softmaxMaxShape->SetDim(DIM_INDEX_3, queryShape->GetDim(DIM_INDEX_2) / keyShape->GetDim(DIM_INDEX_2)); + + softmaxSumShape->SetDimNum(DIM_NUM_4); + softmaxSumShape->SetDim(DIM_INDEX_0, queryShape->GetDim(DIM_INDEX_0)); + softmaxSumShape->SetDim(DIM_INDEX_1, keyShape->GetDim(DIM_INDEX_2)); + softmaxSumShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1)); + softmaxSumShape->SetDim(DIM_INDEX_3, queryShape->GetDim(DIM_INDEX_2) / keyShape->GetDim(DIM_INDEX_2)); + } + } else { + softmaxMaxShape->SetDimNum(DIM_NUM_1); + softmaxMaxShape->SetDim(DIM_INDEX_0, 0); + softmaxSumShape->SetDimNum(DIM_NUM_1); + softmaxSumShape->SetDim(DIM_INDEX_0, 0); + } + + return GRAPH_SUCCESS; +} + +ge::graphStatus InferDataTypeSparseFlashAttention(gert::InferDataTypeContext *context) +{ + OP_CHECK_IF(context == nullptr, OP_LOGE("SparseFlashAttention", "InferShapeContext is nullptr"), + return ge::GRAPH_FAILED); + const auto inputDataType = context->GetInputDataType(QUERY_INPUT_INDEX); + context->SetOutputDataType(OUTPUT_INDEX_0, inputDataType); + context->SetOutputDataType(OUTPUT_INDEX_1, ge::DT_FLOAT); + context->SetOutputDataType(OUTPUT_INDEX_2, ge::DT_FLOAT); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(SparseFlashAttention) + .InferShape(InferShapeSparseFlashAttention) + .InferDataType(InferDataTypeSparseFlashAttention); +} // namespace ops + diff --git a/csrc/ascend/attention/sparse_flash_attention/op_host/sparse_flash_attention_tiling.cpp b/csrc/ascend/attention/sparse_flash_attention/op_host/sparse_flash_attention_tiling.cpp new file mode 100644 index 000000000..923da2592 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_host/sparse_flash_attention_tiling.cpp @@ -0,0 +1,2073 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_flash_attention_tiling.cpp + * \brief + */ + +#include +#include +#include +#include +#include +#include "err/ops_err.h" +#include "register/op_def_registry.h" +#include "../op_kernel/sparse_flash_attention_template_tiling_key.h" +#include "sparse_flash_attention_tiling.h" + +using std::map; +using std::string; +using std::pair; + +using namespace ge; +using namespace AscendC; +namespace optiling { + +constexpr uint32_t PRE_LOAD_NUM = 2; +constexpr uint32_t BLOCK_TABLE_ELEM_BYTE = 4; +constexpr int32_t SPARSE_MODE_BAND = 4; + +static const std::string QUERY_NAME = "query"; +static const std::string KEY_NAME = "key"; +static const std::string VALUE_NAME = "value"; +static const std::string BLOCK_TABLE_NAME = "block_table"; +static const std::string SPARSE_INDICES_NAME = "sparse_indices"; +static const std::string QUERY_ROPE_NAME = "query_rope"; +static const std::string KEY_ROPE_NAME = "key_rope"; +static const std::string ATTEN_OUT_NAME = "attention_out"; +static const std::string SOFTMAX_MAX_NAME = "softmax_max"; +static const std::string SOFTMAX_SUM_NAME = "softmax_sum"; + +const std::map> DTYPE_SUPPORT_MAP = { + {QUERY_NAME, {ge::DT_FLOAT16, ge::DT_BF16}}, + {KEY_NAME, {ge::DT_FLOAT16, ge::DT_BF16}}, + {VALUE_NAME, {ge::DT_FLOAT16, ge::DT_BF16}}, + {QUERY_ROPE_NAME, {ge::DT_FLOAT16, ge::DT_BF16}}, + {KEY_ROPE_NAME, {ge::DT_FLOAT16, ge::DT_BF16}}, + {ATTEN_OUT_NAME, {ge::DT_FLOAT16, ge::DT_BF16}}, + {SOFTMAX_MAX_NAME, {ge::DT_FLOAT}}, + {SOFTMAX_SUM_NAME, {ge::DT_FLOAT}}, + {SPARSE_INDICES_NAME, {ge::DT_INT32}}, + {BLOCK_TABLE_NAME, {ge::DT_INT32}}, +}; + +const std::map> LAYOUT_SUPPORT_MAP = { + {QUERY_NAME, {SFALayout::BSND, SFALayout::TND}}, + {KEY_NAME, {SFALayout::BSND, SFALayout::TND, SFALayout::PA_BSND}}, + {VALUE_NAME, {SFALayout::BSND, SFALayout::TND, SFALayout::PA_BSND}}, + {ATTEN_OUT_NAME, {SFALayout::BSND, SFALayout::TND}}, + {SOFTMAX_MAX_NAME, {SFALayout::BNSG, SFALayout::NTG}}, + {SOFTMAX_SUM_NAME, {SFALayout::BNSG, SFALayout::NTG}}, +}; + +const std::map DATATYPE_TO_STRING_MAP = { + {ge::DT_UNDEFINED, "DT_UNDEFINED"}, // Used to indicate a DataType field has not been set. + {ge::DT_FLOAT, "DT_FLOAT"}, // float type + {ge::DT_FLOAT16, "DT_FLOAT16"}, // fp16 type + {ge::DT_INT8, "DT_INT8"}, // int8 type + {ge::DT_INT16, "DT_INT16"}, // int16 type + {ge::DT_UINT16, "DT_UINT16"}, // uint16 type + {ge::DT_UINT8, "DT_UINT8"}, // uint8 type + {ge::DT_INT32, "DT_INT32"}, // uint32 type + {ge::DT_INT64, "DT_INT64"}, // int64 type + {ge::DT_UINT32, "DT_UINT32"}, // unsigned int32 + {ge::DT_UINT64, "DT_UINT64"}, // unsigned int64 + {ge::DT_BOOL, "DT_BOOL"}, // bool type + {ge::DT_DOUBLE, "DT_DOUBLE"}, // double type + {ge::DT_DUAL, "DT_DUAL"}, // dual output type + {ge::DT_DUAL_SUB_INT8, "DT_DUAL_SUB_INT8"}, // dual output int8 type + {ge::DT_DUAL_SUB_UINT8, "DT_DUAL_SUB_UINT8"}, // dual output uint8 type + {ge::DT_COMPLEX32, "DT_COMPLEX32"}, // complex32 type + {ge::DT_COMPLEX64, "DT_COMPLEX64"}, // complex64 type + {ge::DT_COMPLEX128, "DT_COMPLEX128"}, // complex128 type + {ge::DT_QINT8, "DT_QINT8"}, // qint8 type + {ge::DT_QINT16, "DT_QINT16"}, // qint16 type + {ge::DT_QINT32, "DT_QINT32"}, // qint32 type + {ge::DT_QUINT8, "DT_QUINT8"}, // quint8 type + {ge::DT_QUINT16, "DT_QUINT16"}, // quint16 type + {ge::DT_RESOURCE, "DT_RESOURCE"}, // resource type + {ge::DT_STRING_REF, "DT_STRING_REF"}, // string ref type + {ge::DT_STRING, "DT_STRING"}, // string type + {ge::DT_VARIANT, "DT_VARIANT"}, // dt_variant type + {ge::DT_BF16, "DT_BFLOAT16"}, // dt_bfloat16 type + {ge::DT_INT4, "DT_INT4"}, // dt_variant type + {ge::DT_UINT1, "DT_UINT1"}, // dt_variant type + {ge::DT_INT2, "DT_INT2"}, // dt_variant type + {ge::DT_UINT2, "DT_UINT2"} // dt_variant type +}; + +struct SparseFlashAttentionCompileInfo { + int64_t core_num; +}; + +static const std::map> SFA_LAYOUT_AXIS_MAP = { + {SFALayout::BSND, {SFAAxis::B, SFAAxis::S, SFAAxis::N, SFAAxis::D}}, + {SFALayout::TND, {SFAAxis::T, SFAAxis::N, SFAAxis::D}}, + {SFALayout::PA_BSND, {SFAAxis::Bn, SFAAxis::Bs, SFAAxis::N, SFAAxis::D}}, + {SFALayout::BNSG, {SFAAxis::B, SFAAxis::N, SFAAxis::S, SFAAxis::G}}, + {SFALayout::NTG, {SFAAxis::N, SFAAxis::T, SFAAxis::G}}, +}; + +static const std::map SFA_LAYOUT_DIM_MAP = { + {SFALayout::BSND, DIM_NUM_FOUR}, + {SFALayout::TND, DIM_NUM_THREE}, + {SFALayout::PA_BSND, DIM_NUM_FOUR}, + {SFALayout::BNSG, DIM_NUM_FOUR}, + {SFALayout::NTG, DIM_NUM_THREE}, +}; + +static std::string GetShapeStr(gert::Shape shape) +{ + std::ostringstream oss; + oss << "["; + if (shape.GetDimNum() > 0) { + for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) { + oss << shape.GetDim(i) << ", "; + } + oss << shape.GetDim(shape.GetDimNum() - 1); + } + oss << "]"; + return oss.str(); +} + +static std::string SFADataTypeToSerialString(ge::DataType type) +{ + const auto it = DATATYPE_TO_STRING_MAP.find(type); + if (it != DATATYPE_TO_STRING_MAP.end()) { + return it->second; + } else { + OP_LOGE("SparseFlashAttention", "datatype %d not support", type); + return "UNDEFINED"; + } +} + +string SFATensorDesc2String(const gert::StorageShape *shape, const gert::CompileTimeTensorDesc *tensor) +{ + if (shape == nullptr || tensor == nullptr) { + return "nil "; + } + + std::ostringstream oss; + oss << "(dtype: " << ge::TypeUtils::DataTypeToAscendString(tensor->GetDataType()).GetString() << "),"; + oss << "(shape:" << SFAShape2String(shape->GetStorageShape()) << "),"; + oss << "(ori_shape:" << SFAShape2String(shape->GetOriginShape()) << "),"; + oss << "(format: " + << ge::TypeUtils::FormatToAscendString( + static_cast(ge::GetPrimaryFormat(tensor->GetStorageFormat()))) + .GetString() + << "),"; + oss << "(ori_format: " << ge::TypeUtils::FormatToAscendString(tensor->GetOriginFormat()).GetString() << ") "; + + return oss.str(); +} + +string SFADebugTilingContext(const gert::TilingContext *context) +{ + std::ostringstream oss; + for (size_t i = 0; i < context->GetComputeNodeInfo()->GetInputsNum(); ++i) { + oss << "input" << i << ": "; + oss << SFATensorDesc2String(context->GetInputShape(i), context->GetInputDesc(i)); + } + + for (size_t i = 0; i < context->GetComputeNodeInfo()->GetOutputsNum(); ++i) { + oss << "output" << i << ": "; + oss << SFATensorDesc2String(context->GetOutputShape(i), context->GetOutputDesc(i)); + } + return oss.str(); +} + +std::string SFALayoutToSerialString(SFALayout layout) +{ + switch (layout) { + case SFALayout::BSND: return "BSND"; + case SFALayout::TND: return "TND"; + case SFALayout::PA_BSND: return "PA_BSND"; + case SFALayout::BNSG: return "BNSG"; + case SFALayout::NTG: return "NTG"; + default: return "UNKNOWN"; + } +} + +ge::graphStatus SFAMlaTiling::SetBlockDim(uint32_t blockDim) const +{ + context_->SetBlockDim(blockDim); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAMlaTiling::SetTilingKey(uint64_t tilingKey) const +{ + context_->SetTilingKey(tilingKey); + context_->SetScheduleMode(1); // 1: batchmode模式 + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAMlaTiling::SetWorkspaceSize(uint64_t workspaceSize) const +{ + OP_CHECK_IF(context_->GetWorkspaceSizes(1) == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "workSpaceSize got from ge is nullptr"), + return ge::GRAPH_FAILED); + size_t *workSpaces = context_->GetWorkspaceSizes(1); + workSpaces[0] = workspaceSize; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAMlaTiling::SetTilingData(TilingDef &tilingData) const +{ + OP_CHECK_IF(context_->GetRawTilingData() == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "RawTilingData got from GE context is nullptr."), + return ge::GRAPH_FAILED); + + tilingData.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData.GetDataSize()); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAMlaTiling::GetPlatformInfo() +{ + OP_CHECK_IF(sfaInfo_->platformInfo == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(sfaInfo_->opName, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(sfaInfo_->platformInfo); + libapiSize_ = ascendcPlatform.GetLibApiWorkSpaceSize(); + aivNum_ = ascendcPlatform.GetCoreNumAiv(); + aicNum_ = ascendcPlatform.GetCoreNumAic(); + + OP_CHECK_IF(aicNum_ == 0 || aivNum_ == 0, + OPS_REPORT_VECTOR_INNER_ERR(sfaInfo_->opName, "num of core obtained is 0."), return GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void SFAMlaTiling::GenTilingKey() +{ + uint32_t inputQType = static_cast(sfaInfo_->inputQType); + uint32_t inputKvType = static_cast(sfaInfo_->inputKvType); + uint32_t outputType = static_cast(sfaInfo_->outputType); + uint32_t layoutQuery = static_cast(sfaInfo_->qLayout); + uint32_t layoutKV = static_cast(sfaInfo_->kvLayout); + uint32_t pageAttention = 0U; + if (sfaInfo_->kvLayout == SFALayout::PA_BSND) { + pageAttention = 1U; + } + + tilingKey_ = GET_TPL_TILING_KEY(0U, pageAttention, layoutQuery, layoutKV, + perfMode_ == SFAPerfMode::V_TEMPLATE_MODE, static_cast(sfaInfo_->gSize > 64)); // N1 > 128时核间切G + + OP_LOGI(sfaInfo_->opName, "SFA tilingKey_: %lu.", tilingKey_); +} + +void SFAMlaTiling::ZeroTensorProcess() const +{ + if (sfaInfo_->s2Size == 0) { + /* + * 1024，空tensor场景下，作为默认值完成后续计算 + * 避免matmal tiling softmax tiling异常 + * kernel计算使用真实的seqSize=0, 与actuseq_len流程归一 + */ + sfaInfo_->s2Size = 1024; + } +} + +void SFAMlaTiling::InitParams() +{ + if (sfaInfo_->s2Size != 0 && sfaInfo_->sparseBlockSize <= 4) { // 4:当前支持范围 + perfMode_ = SFAPerfMode::V_TEMPLATE_MODE; + } else { + perfMode_ = SFAPerfMode::C_TEMPLATE_MODE; + } + + coreNum_ = aicNum_; + + headDimAlign_ = Align(sfaInfo_->qkHeadDim, BYTE_BLOCK); // 元素个数按照基本块大小对齐 + ZeroTensorProcess(); +} + +void SFAMlaTiling::CalcUbBmm() +{ + uint32_t cubeMSize = sfaInfo_->gSize * sfaInfo_->s1Size; + uint32_t maxMSize = mBaseSize_; + if (cubeMSize > maxMSize) { + cubeMSize = maxMSize; + } + mmResUbSize_ = sInnerSizeAlign_ * Align(cubeMSize, 16U);// kernel按照16对齐写出，tiling按照这个原则分配内存 + bmm2ResUbSize_ = headDimAlign_ * Align(cubeMSize, 16U);// kernel按照16对齐写出，tiling按照这个原则分配内存 + + qPreSizeMla_ = sfaInfo_->gSize * (headDimAlign_ + 64U) * sfaInfo_->s1Size; +} + +void SFAMlaTiling::CheckUbSpace() +{ + CalcUbBmm(); +} + +void SFAMlaTiling::CalcInnerSize(uint32_t s2Size) +{ + sInnerSize_ = 512; // 512:s2默认切分大小 + // FlashDecode时，如果S2的计算量>=256(确保切分后不小于128)但又不足以分2次计算时，则修改sInnerSize_，均分为2份进行计算，确保Nbuffer=2 + if (splitKVFlag_ && sfaInfo_->qLayout != SFALayout::TND) { + if (s2Size == 256) { // 256:s2Size的阈值，判断sInnerSize_是否切分 + sInnerSize_ = 128; // 128:sInnerSize_值为s2Size的一半，均分为2份进行计算， + } else if (s2Size > 256 && s2Size <= sInnerSize_) { // 256:s2Size的阈值，判断sInnerSize_是否切分 + sInnerSize_ = (sInnerSize_ + 1) / 2; // 2:减半 + } + } + + sInnerLoopTimes_ = (s2Size + sInnerSize_ - 1) / sInnerSize_; + sInnerSizeTail_ = s2Size - (sInnerLoopTimes_ - 1) * sInnerSize_; + if (sInnerSize_ > s2Size) { + sInnerSize_ = s2Size; + } + sInnerSizeAlign_ = Align(sInnerSize_, BYTE_BLOCK); // 元素个数按照基本块大小对齐 + + CheckUbSpace(); +} + +void SFAMlaTiling::SplitBalanced() +{ + CalcInnerSize(sfaInfo_->s2Size); + + InnerSplitParams innerSplitParams; + innerSplitParams.s1GBaseSize = sfaInfo_->gSize; + innerSplitParams.s2BaseSize = sInnerSize_; + tilingData_.innerSplitParams.set_mBaseSize(innerSplitParams.s1GBaseSize); + tilingData_.innerSplitParams.set_s2BaseSize(innerSplitParams.s2BaseSize); + + usedCoreNum_ = aicNum_; +} + +void SFAMlaTiling::Split() +{ + SplitBalanced(); +} + +void SFAMlaTiling::FillTilingBaseParamsMla() +{ + tilingData_.baseParams.set_batchSize(sfaInfo_->bSize); + tilingData_.baseParams.set_seqSize(sfaInfo_->s2Size); + tilingData_.baseParams.set_qSeqSize(sfaInfo_->s1Size); + tilingData_.baseParams.set_blockSize(sfaInfo_->blockSize); + tilingData_.baseParams.set_maxBlockNumPerBatch(sfaInfo_->maxBlockNumPerBatch); + tilingData_.baseParams.set_scaleValue(sfaInfo_->scaleValue); + tilingData_.baseParams.set_nNumOfQInOneGroup(sfaInfo_->n1Size / sfaInfo_->n2Size); + tilingData_.baseParams.set_actualLenDimsQ(sfaInfo_->actualLenDimsQ); + tilingData_.baseParams.set_actualLenDimsKV(sfaInfo_->actualLenDimsKV); + tilingData_.baseParams.set_outputLayout(static_cast(sfaInfo_->outLayout)); + tilingData_.baseParams.set_sparseMode(sfaInfo_->sparseMode); + tilingData_.baseParams.set_preTokens(sfaInfo_->preTokens); + tilingData_.baseParams.set_nextTokens(sfaInfo_->nextTokens); + tilingData_.baseParams.set_sparseBlockSize(sfaInfo_->sparseBlockSize); + tilingData_.baseParams.set_sparseBlockCount(sfaInfo_->sparseBlockCount); + tilingData_.baseParams.set_attentionMode(sfaInfo_->attentionMode); + tilingData_.baseParams.set_returnSoftmaxLse(sfaInfo_->returnSoftmaxLse); + tilingData_.baseParams.set_isActualLenDimsNull(sfaInfo_->actualQSeqLenFlag ? 0U : 1U); + tilingData_.baseParams.set_isActualLenDimsKVNull(sfaInfo_->actualSeqLenFlag ? 0U : 1U); +} + +// for flash decode +void SFAMlaTiling::FillTilingSplitKVMla() +{ + tilingData_.splitKVParams.set_s2(kvSplitPart_); + + tilingData_.splitKVParams.set_accumOutSize(aicNum_ * 2 * sfaInfo_->n2Size * mBaseSize_ * headDimAlign_); // 2:每个核可能有头规约和尾规约，一共两份规约信息 + tilingData_.splitKVParams.set_logSumExpSize(2 * aicNum_ * 2 * sfaInfo_->n2Size * mBaseSize_ * // 2:每个核可能有头规约和尾规约，一共两份规约信息;sum + max + (BYTE_BLOCK / BLOCK_TABLE_ELEM_BYTE)); + + if (!splitKVFlag_) { + tilingData_.splitKVParams.set_s2(0); + } +} + +void SFAMlaTiling::FillTilingSingleCoreParamsMla() +{ + tilingData_.singleCoreParams.set_usedCoreNum(usedCoreNum_); +} + +void SFAMlaTiling::FillTilingSingleCoreTensorSizeMla() +{ + tilingData_.singleCoreTensorSize.set_mmResUbSize(mmResUbSize_); + tilingData_.singleCoreTensorSize.set_bmm2ResUbSize(bmm2ResUbSize_); +} + +void SFAMlaTiling::FillTiling() +{ + FillTilingBaseParamsMla(); + FillTilingSplitKVMla(); + FillTilingSingleCoreParamsMla(); + FillTilingSingleCoreTensorSizeMla(); +} + +uint32_t SFAMlaTiling::CalcBalanceFDParamNums(const uint32_t actCoreNum) const +{ + return actCoreNum * 2 * sfaInfo_->n2Size * mBaseSize_; // 2:每个核可能有头规约和尾规约，一共两份规约信息 +} + +void SFAMlaTiling::NormalCalcFDWorkSpace(const uint32_t actCoreNum) +{ + if (splitKVFlag_) { + uint32_t accumOutSize = 0; + uint32_t logSumExpSize = 0; + uint32_t FDParamNums = CalcBalanceFDParamNums(actCoreNum); //balanceModeFlag_ ? CalcBalanceFDParamNums(actCoreNum) : CalcUnbalanceFDParamNums(); + accumOutSize = FDParamNums * headDimAlign_; + logSumExpSize = 2 * FDParamNums * (BYTE_BLOCK / sfaInfo_->blockTypeSize); // log和sum的存储空间一致，共需要2份内存 + workspaceSize_ += (accumOutSize + logSumExpSize) * sfaInfo_->blockTypeSize; + } +} + +void SFAMlaTiling::CalcFDWorkSpace(const uint32_t actCoreNum) +{ + NormalCalcFDWorkSpace(actCoreNum); +} + +void SFAMlaTiling::GetWorkspaceSize() +{ + uint32_t actCoreNum = coreNum_; + if (sfaInfo_->isA5) { + workspaceSize_ = libapiSize_; + constexpr uint32_t TRIPLE_BUFFER_NUM = 3; + constexpr uint32_t S2_BASE_SIZE = 128; // S2轴基本块大小 + constexpr uint32_t D_SIZE = 576; + auto ascendcPlatform = platform_ascendc::PlatformAscendC(sfaInfo_->platformInfo); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + if (sfaInfo_->gSize > 64) { // N1大于64时切G，相邻两个cube核处理同一个s2Base + aicNum = aicNum >> 1; + } + workspaceSize_ += (S2_BASE_SIZE * D_SIZE * GetTypeSize(sfaInfo_->inputQType) \ + * TRIPLE_BUFFER_NUM * aicNum); + } else { + uint32_t mmResElemSize = 4; // 4:fp32 + uint32_t vec1ResElemSize = 2; // 2:fp16/bf16 + uint32_t bmm2ResElemSize = 4; // 4:fp32 + uint32_t qPreProcResElemSize = 0; // 普通场景不涉及Q预处理 + uint32_t nUpdateElemSize = 4; // 4:int32 + uint32_t softmaxSumElemSize = 4; // 4:int32 + float kvDtypeRatio = 1.0; + + workspaceSize_ = libapiSize_; + uint32_t preLoadNum = 1; + preLoadNum = PRE_LOAD_NUM; + + workspaceSize_ += preLoadNum * (mmResUbSize_ * actCoreNum * mmResElemSize); + workspaceSize_ += preLoadNum * static_cast(static_cast(mmResUbSize_ * \ + actCoreNum * vec1ResElemSize) * kvDtypeRatio); + workspaceSize_ += preLoadNum * bmm2ResUbSize_ * actCoreNum * bmm2ResElemSize; + workspaceSize_ += preLoadNum * static_cast(static_cast(qPreSizeMla_ * \ + actCoreNum * qPreProcResElemSize) * kvDtypeRatio); + workspaceSize_ += preLoadNum * mBaseSize_ * actCoreNum * nUpdateElemSize; + workspaceSize_ += preLoadNum * mBaseSize_ * actCoreNum * softmaxSumElemSize; + // topk BlkSize == 1场景, 需要额外空间缓存离散聚合的值 + // bufNum s2Base D dRope sizeOf(half) + // 4:bufNum 512:s2Base 512:D 64:dRope 2:sizeOf(half) + workspaceSize_ += 4 * 512 * (512 + 64) * 2 * actCoreNum; + // 缓存有效mte2 size的长度 份数 512B对齐的长度 sizeof(int32_t) aiv核数 + workspaceSize_ += 4 * 128 * 4 * (2 * actCoreNum); // 4:缓存有效mte2 size的长度 128:份数 4:512B对齐的长度 2:aiv核数 + } + + CalcFDWorkSpace(actCoreNum); +} + +void SFAMlaTiling::CalcBlockDim() +{ + auto ascendcPlatform = platform_ascendc::PlatformAscendC(sfaInfo_->platformInfo); + auto aicNum = usedCoreNum_; + auto aivNum = 2 * usedCoreNum_; + + blockDim_ = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum); + OP_LOGI(sfaInfo_->opName, "SFA block dim: %u aiv Num: %u aic Num: %u.", blockDim_, aivNum, aicNum); +} + +ge::graphStatus SFAMlaTiling::DoOpTiling(SFATilingInfo *sfaInfo) +{ + sfaInfo_ = sfaInfo; + if (GetPlatformInfo() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + InitParams(); + Split(); + FillTiling(); + CalcBlockDim(); + GetWorkspaceSize(); + GenTilingKey(); + + if ((SetBlockDim(blockDim_) != ge::GRAPH_SUCCESS) || + (SetTilingKey(tilingKey_) != ge::GRAPH_SUCCESS) || + (SetWorkspaceSize(workspaceSize_) != ge::GRAPH_SUCCESS) || + (SetTilingData(tilingData_) != ge::GRAPH_SUCCESS)) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingSparseFlashAttention(gert::TilingContext *context) +{ + SFATilingInfo sfaInfo; + SFAInfoParser sfaInfoParser(context); + if (sfaInfoParser.Parse(sfaInfo) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + SFATilingCheck tilingChecker(sfaInfo); + if (tilingChecker.Process() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + SFAMlaTiling tiling(context); + return tiling.DoOpTiling(&sfaInfo); +} + +ge::graphStatus TilingPrepareForSparseFlashAttention(gert::TilingParseContext* const context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::GetExpectedShape(gert::Shape &shapeExpected, + const SFATilingShapeCompareParam ¶m, const SFALayout &layout) const +{ + if (layout == SFALayout::BSND) { + shapeExpected = gert::Shape({param.B, param.S, param.N, param.D}); + } else if (layout == SFALayout::TND) { + shapeExpected = gert::Shape({param.T, param.N, param.D}); + } else if (layout == SFALayout::PA_BSND) { + shapeExpected = gert::Shape({param.Bn, param.Bs, param.N, param.D}); + } else if (layout == SFALayout::BNSG) { + shapeExpected = gert::Shape({param.B, param.N, param.S, param.G}); + } else if (layout == SFALayout::NTG) { + shapeExpected = gert::Shape({param.N, param.T, param.G}); + } else { + OP_LOGE(opName_, "layout %s is unsupported", SFALayoutToSerialString(layout).c_str()); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CompareShape(SFATilingShapeCompareParam ¶m, + const gert::Shape &shape, const SFALayout &layout, const std::string &name) const +{ + gert::Shape shapeExpected; + if (GetExpectedShape(shapeExpected, param, layout) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + if (shape.GetDimNum() != shapeExpected.GetDimNum()) { + OP_LOGE(opName_, + "%s dimension is %zu, expected dimension is %zu.", + name.c_str(), shape.GetDimNum(), shapeExpected.GetDimNum()); + return ge::GRAPH_FAILED; + } + + for (size_t i = 0; i < shape.GetDimNum(); i++) { + if (shape.GetDim(i) != shapeExpected.GetDim(i)) { + OP_LOGE(opName_, "%s layout is %s, shape is %s, expected shape is %s.", + name.c_str(), SFALayoutToSerialString(layout).c_str(), + GetShapeStr(shape).c_str(), GetShapeStr(shapeExpected).c_str()); + return ge::GRAPH_FAILED; + } + } + + return ge::GRAPH_SUCCESS; +} + +void SFATilingCheck::LogErrorDtypeSupport(const std::vector &expectDtypeList, + const ge::DataType &actualDtype, const std::string &name) const +{ + std::ostringstream oss; + for (size_t i = 0; i < expectDtypeList.size(); ++i) { + oss << SFADataTypeToSerialString(expectDtypeList[i]); + if (i < expectDtypeList.size() - 1) { + oss << ", "; + } + } + OP_LOGE(opName_, "Tensor %s only supports dtype %s, but got %s", + name.c_str(), oss.str().c_str(), SFADataTypeToSerialString(actualDtype).c_str()); +} + +ge::graphStatus SFATilingCheck::CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc, + const std::string &name) const +{ + if (desc != nullptr) { + const auto& it = DTYPE_SUPPORT_MAP.find(name); + OP_CHECK_IF(it == DTYPE_SUPPORT_MAP.end(), + OP_LOGE(opName_, "%s datatype support list should be specify in DTYPE_SUPPORT_MAP", name.c_str()), + return ge::GRAPH_FAILED); + auto &expectDtypeList = it->second; + OP_CHECK_IF(std::find( + expectDtypeList.begin(), expectDtypeList.end(), desc->GetDataType()) == expectDtypeList.end(), + LogErrorDtypeSupport(expectDtypeList, desc->GetDataType(), name), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +template +void SFATilingCheck::LogErrorNumberSupport(const std::vector &expectNumberList, + const T &actualValue, const std::string &name, const std::string subName) const +{ + std::ostringstream oss; + for (size_t i = 0; i < expectNumberList.size(); ++i) { + oss << std::to_string(expectNumberList[i]); + if (i < expectNumberList.size() - 1) { + oss << ", "; + } + } + + OP_LOGE(opName_, "%s %s only supports %s, but got %s", + name.c_str(), subName.c_str(), oss.str().c_str(), std::to_string(actualValue).c_str()); +} + +template +void SFATilingCheck::LogErrorDimNumSupport(const std::vector &expectNumberList, + const T &actualValue, const std::string &name) const +{ + LogErrorNumberSupport(expectNumberList, actualValue, name, "dimension"); +} + +ge::graphStatus SFATilingCheck::CheckDimNumInLayoutSupport(const SFALayout &layout, + const gert::StorageShape *shape, const std::string &name) const +{ + const auto& dimIt = SFA_LAYOUT_DIM_MAP.find(layout); + OP_CHECK_IF(shape->GetStorageShape().GetDimNum() != dimIt->second, + OP_LOGE(opName_, "When layout is %s, %s dimension should be %zu, but it's %zu", + SFALayoutToSerialString(layout).c_str(), name.c_str(), dimIt->second, + shape->GetStorageShape().GetDimNum()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckDimNumSupport(const gert::StorageShape *shape, + const std::vector &expectDimNumList, const std::string &name) const +{ + if (shape == nullptr) { + return ge::GRAPH_SUCCESS; + } + + if (std::find(expectDimNumList.begin(), expectDimNumList.end(), + shape->GetStorageShape().GetDimNum()) == expectDimNumList.end()) { + LogErrorDimNumSupport(expectDimNumList, shape->GetStorageShape().GetDimNum(), name); + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + + +void SFATilingCheck::LogErrorLayoutSupport(const std::vector &expectLayoutList, + const SFALayout &actualLayout, const std::string &name) const +{ + std::ostringstream oss; + for (size_t i = 0; i < expectLayoutList.size(); ++i) { + oss << SFALayoutToSerialString(expectLayoutList[i]); + if (i < expectLayoutList.size() - 1) { + oss << ", "; + } + } + OP_LOGE(opName_, "Tensor %s only supports layout %s, but got %s", + name.c_str(), oss.str().c_str(), SFALayoutToSerialString(actualLayout).c_str()); +} + +ge::graphStatus SFATilingCheck::CheckLayoutSupport(const SFALayout &actualLayout, const std::string &name) const +{ + const auto& it = LAYOUT_SUPPORT_MAP.find(name); + OP_CHECK_IF(it == LAYOUT_SUPPORT_MAP.end(), + OP_LOGE(opName_, "%s layout support list should be specify in LAYOUT_SUPPORT_MAP", name.c_str()), + return ge::GRAPH_FAILED); + auto &expectLayoutList = it->second; + OP_CHECK_IF(std::find( + expectLayoutList.begin(), expectLayoutList.end(), actualLayout) == expectLayoutList.end(), + LogErrorLayoutSupport(expectLayoutList, actualLayout, name), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckSingleParaQuery() const +{ + const std::vector queryDimNumList = {DIM_NUM_THREE, DIM_NUM_FOUR}; + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.query.desc, QUERY_NAME) || + ge::GRAPH_SUCCESS != CheckLayoutSupport(qLayout_, QUERY_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(opParamInfo_.query.shape, queryDimNumList, QUERY_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(qLayout_, opParamInfo_.query.shape, QUERY_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckSingleParaKey() const +{ + const std::vector keyDimNumList = {DIM_NUM_FOUR, DIM_NUM_THREE}; + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.key.desc, KEY_NAME) || + ge::GRAPH_SUCCESS != CheckLayoutSupport(kvLayout_, KEY_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumSupport(opParamInfo_.key.shape, keyDimNumList, KEY_NAME) || + ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(kvLayout_, opParamInfo_.key.shape, KEY_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckSingleParaNumHeads() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckSingleParaKvHeadNums() const +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckSingleParaSparseMode() const +{ + OP_CHECK_IF((*opParamInfo_.sparseMode != 3 && *opParamInfo_.sparseMode != 0), + OP_LOGE(opName_, "sparseMode must == 0/3, but got: %ld.", *opParamInfo_.sparseMode), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckSingleParaSparseBlockSize() const +{ + OP_CHECK_IF((*opParamInfo_.sparseBlockSize <= 0 || *opParamInfo_.sparseBlockSize > 128 || + (static_cast(*opParamInfo_.sparseBlockSize) & static_cast(*opParamInfo_.sparseBlockSize - 1L)) != 0UL), + OP_LOGE(opName_, "sparseBlockSize should be be in range [1, 128] and be a power of 2, but got: %ld.", + *opParamInfo_.sparseBlockSize), + return ge::GRAPH_FAILED); + + OP_CHECK_IF((npuArch_ == NpuArch::DAV_3510 && *opParamInfo_.sparseBlockSize != 1), + OP_LOGE(opName_, "when soc version is Ascend950, sparse_block_size only support 1, but now is %d.", + *opParamInfo_.sparseBlockSize), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckSingleParaSparseIndices() const +{ + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.sparseIndices.desc, SPARSE_INDICES_NAME)) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckSingleParaPreTokens() const +{ + OP_CHECK_IF((*opParamInfo_.preTokens != INT64_MAX), + OP_LOGE(opName_, "preTokens should be 9223372036854775807, but got: %ld.", *opParamInfo_.preTokens), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckSingleParaNextTokens() const +{ + OP_CHECK_IF((*opParamInfo_.nextTokens != INT64_MAX), + OP_LOGE(opName_, "nextTokens should be 9223372036854775807, but got: %ld.", *opParamInfo_.nextTokens), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckSinglePara() const +{ + if (ge::GRAPH_SUCCESS != CheckSingleParaQuery() || + ge::GRAPH_SUCCESS != CheckSingleParaKey() || + ge::GRAPH_SUCCESS != CheckSingleParaSparseIndices() || + ge::GRAPH_SUCCESS != CheckSingleParaNumHeads() || + ge::GRAPH_SUCCESS != CheckSingleParaKvHeadNums() || + ge::GRAPH_SUCCESS != CheckSingleParaSparseMode() || + ge::GRAPH_SUCCESS != CheckSingleParaSparseBlockSize() || + ge::GRAPH_SUCCESS != CheckSingleParaPreTokens() || + ge::GRAPH_SUCCESS != CheckSingleParaNextTokens()) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckRopeExistence() +{ + OP_CHECK_IF((opParamInfo_.queryRope.tensor != nullptr || opParamInfo_.keyRope.tensor != nullptr) + && *opParamInfo_.attentionMode == 0, + OP_LOGE(opName_, "In MHA/GQA situation(attentionMode=0), queryRope and keyRope should be null."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(*opParamInfo_.attentionMode != 2, + OP_LOGE(opName_, "attentionMode only support 2."), + return ge::GRAPH_FAILED); + OP_CHECK_IF((opParamInfo_.queryRope.tensor != nullptr && opParamInfo_.keyRope.tensor == nullptr), + OP_LOGE(opName_, "KeyRope is null, but queryRope exists, they should be both null or exist."), + return ge::GRAPH_FAILED); + OP_CHECK_IF((opParamInfo_.queryRope.tensor == nullptr && opParamInfo_.keyRope.tensor != nullptr), + OP_LOGE(opName_, "QueryRope is null, but keyRope exists, they should be both null or exist."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.keyRope.desc == nullptr || opParamInfo_.queryRope.desc == nullptr, + OP_LOGE(opName_, "In Mla situation, desc of keyRope and queryRope should not be null"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckExists(const void *pointer, const std::string &name) const +{ + OP_CHECK_IF(pointer == nullptr, + OP_LOGE(opName_, "%s should not be null", name.c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckNotExists(const void *pointer, const std::string &name) const +{ + OP_CHECK_IF(pointer != nullptr, + OP_LOGE(opName_, "%s should be null", name.c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckExistsByMap(const std::map ¶mMap) const +{ + for (const auto& kv : paramMap) { + if (CheckExists(kv.second, kv.first) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckNotExistsByMap(const std::map ¶mMap) const +{ + for (const auto& kv : paramMap) { + if (CheckNotExists(kv.second, kv.first) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckExistenceByMap(std::map &existMap, + std::map ¬ExistMap) const +{ + if (CheckExistsByMap(existMap) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + if (CheckNotExistsByMap(notExistMap) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +template +ge::graphStatus SFATilingCheck::CheckAttrValueByMap(std::map> &attrMap) const +{ + for (auto const &kv : attrMap) { + const std::string &name = kv.first; + const std::pair &pointerValuePair = kv.second; + if (pointerValuePair.first == nullptr) { + OP_LOGE(opName_, "Attr %s should not be nullptr", name.c_str()); + return ge::GRAPH_FAILED; + } + + if (*(pointerValuePair.first) != pointerValuePair.second) { + std::ostringstream ossExpect; + ossExpect << std::to_string(pointerValuePair.second); + std::ostringstream ossActual; + ossActual << std::to_string(*(pointerValuePair.first)); + OP_LOGE(opName_, + "%s value should be %s, but got %s", + name.c_str(), + ossExpect.str().c_str(), + ossActual.str().c_str()); + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckParaExistenceMlaNoquant() const +{ + if (kvStorageMode_ != KvStorageMode::PAGE_ATTENTION) { + return ge::GRAPH_SUCCESS; + } + std::map mlaNoquantParamExistMap = { + {"actualSeqLengths", opParamInfo_.actualSeqLengths.tensor}, + {"blockTable", opParamInfo_.blockTable.tensor}, + }; + std::map mlaNoquantParamNotExistMap = {}; + if (CheckExistenceByMap(mlaNoquantParamExistMap, mlaNoquantParamNotExistMap) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckParaExistenceMla() const +{ + return CheckParaExistenceMlaNoquant(); +} + +ge::graphStatus SFATilingCheck::CheckParaExistence() +{ + if (ge::GRAPH_SUCCESS != CheckRopeExistence()) { + return ge::GRAPH_FAILED; + } + + return CheckParaExistenceMla(); +} + +ge::graphStatus SFATilingCheck::GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + const SFALayout &layout, const std::string &name) const +{ + if (tensor == nullptr) { + OP_LOGE(opName_, "when layout of query is %s, %s must be provided.", + SFALayoutToSerialString(layout).c_str(), name.c_str()); + return ge::GRAPH_FAILED; + } + int64_t shapeSize = tensor->GetShapeSize(); + if (shapeSize <= 0) { + OP_LOGE(opName_, "the shape size of %s is %ld, it should be greater than 0.", + name.c_str(), shapeSize); + return ge::GRAPH_FAILED; + } + size = static_cast(shapeSize); + return ge::GRAPH_SUCCESS; +} + +void SFATilingCheck::SetSFAShapeCompare() +{ + queryShapeCmp_ = opParamInfo_.query.shape->GetStorageShape(); + topkShapeCmp_ = opParamInfo_.sparseIndices.shape->GetStorageShape(); + keyShapeCmp_ = opParamInfo_.key.shape->GetStorageShape(); + valueShapeCmp_ = opParamInfo_.value.shape->GetStorageShape(); + attenOutShapeCmp_ = opParamInfo_.attenOut.shape->GetStorageShape(); + queryRopeShapeCmp_ = opParamInfo_.queryRope.tensor->GetStorageShape(); + keyRopeShapeCmp_ = opParamInfo_.keyRope.tensor->GetStorageShape(); + softmaxMaxShapeCmp_ = opParamInfo_.softmaxMax.shape->GetStorageShape(); + softmaxSumShapeCmp_ = opParamInfo_.softmaxSum.shape->GetStorageShape(); +} + +ge::graphStatus SFATilingCheck::CheckBlockTable() const +{ + if (kvStorageMode_ != KvStorageMode::PAGE_ATTENTION) { + OP_CHECK_IF(opParamInfo_.blockTable.tensor != nullptr, + OP_LOGE(opName_, "when the layout_kv is %s, %s should be null", + SFALayoutToSerialString(kvLayout_).c_str(), BLOCK_TABLE_NAME.c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; + } + if (ge::GRAPH_SUCCESS != CheckDtypeSupport(opParamInfo_.blockTable.desc, BLOCK_TABLE_NAME)) { + return ge::GRAPH_FAILED; + } + uint32_t blockTableBatch = opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0); + OP_CHECK_IF(blockTableBatch != bSize_, + OP_LOGE(opName_, "%s's first dimension(%u) should be equal to batch size(%u)", + BLOCK_TABLE_NAME.c_str(), blockTableBatch, bSize_), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckDTypeConsistency(const ge::DataType &actualDtype, + const ge::DataType &expectDtype, const std::string &name) const +{ + if (actualDtype != expectDtype) { + OP_LOGE(opName_, "%s dtype should be %s, but it's %s.", name.c_str(), + SFADataTypeToSerialString(expectDtype).c_str(), + SFADataTypeToSerialString(actualDtype).c_str()); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckQRopeShape() +{ + SFATilingShapeCompareParam shapeParams; + shapeParams.B = bSize_; + shapeParams.N = n1Size_; + shapeParams.S = s1Size_; + shapeParams.D = ropeHeadDim_; + shapeParams.T = qTSize_; + return CompareShape(shapeParams, queryRopeShapeCmp_, qLayout_, QUERY_ROPE_NAME); +} + +ge::graphStatus SFATilingCheck::CheckTopkShape() +{ + SFATilingShapeCompareParam shapeParams; + shapeParams.B = bSize_; + shapeParams.N = n2Size_; + shapeParams.S = s1Size_; + shapeParams.D = sparseBlockCount_; + shapeParams.T = qTSize_; + return CompareShape(shapeParams, topkShapeCmp_, topkLayout_, SPARSE_INDICES_NAME); +} + +ge::graphStatus SFATilingCheck::CheckAttenOutShape() +{ + SFATilingShapeCompareParam shapeParams; + shapeParams.B = bSize_; + shapeParams.N = n1Size_; + shapeParams.S = s1Size_; + shapeParams.D = vHeadDim_; + shapeParams.T = qTSize_; + if (CompareShape(shapeParams, attenOutShapeCmp_, outLayout_, ATTEN_OUT_NAME) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckSoftmaxMaxShape() +{ + if (*opParamInfo_.returnSoftmaxLse) { + SFATilingShapeCompareParam shapeParams; + shapeParams.B = bSize_; + shapeParams.N = n2Size_; + shapeParams.S = s1Size_; + shapeParams.T = qTSize_; + shapeParams.G = n1Size_/n2Size_; + if (CompareShape(shapeParams, softmaxMaxShapeCmp_, softmaxMaxLayout_, SOFTMAX_MAX_NAME) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckSoftmaxSumShape() +{ + if (*opParamInfo_.returnSoftmaxLse) { + SFATilingShapeCompareParam shapeParams; + shapeParams.B = bSize_; + shapeParams.N = n2Size_; + shapeParams.S = s1Size_; + shapeParams.D = vHeadDim_; + shapeParams.T = qTSize_; + shapeParams.G = n1Size_/n2Size_; + if (CompareShape(shapeParams, softmaxSumShapeCmp_, softmaxSumLayout_, SOFTMAX_SUM_NAME) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckAttenOut() +{ + if (ge::GRAPH_SUCCESS != CheckDTypeConsistency(opParamInfo_.attenOut.desc->GetDataType(), + inputQType_, ATTEN_OUT_NAME) || + ge::GRAPH_SUCCESS != CheckAttenOutShape() || + ge::GRAPH_SUCCESS != CheckSoftmaxMaxShape() || + ge::GRAPH_SUCCESS != CheckSoftmaxSumShape()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} +ge::graphStatus SFATilingCheck::CheckSoftmaxMax() +{ + if (*opParamInfo_.returnSoftmaxLse) { + OP_CHECK_IF(opParamInfo_.softmaxMax.shape->GetStorageShape().GetShapeSize() == 0, + OP_LOGE(opName_, "When return_softmax_lse is true, SoftmaxMax tensor cannot be empty tensor."), + return ge::GRAPH_FAILED); + // type类型校验 + OP_CHECK_IF(opParamInfo_.softmaxMax.desc->GetDataType() != ge::DT_FLOAT, + OP_LOGE(opName_, "SoftmaxMax's dtype must be FLOAT."), + return ge::GRAPH_FAILED); + // shape和维度校验 + if (ge::GRAPH_SUCCESS != CheckAttenOutShape()) { + return ge::GRAPH_FAILED; + } + } else { + if (opParamInfo_.softmaxMax.shape->GetStorageShape().GetShapeSize() != 0) { + OP_LOGW(opName_, "When return_softmax_lse is false, SoftmaxMax tensor must be empty tensor."); + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckSoftmaxSum() +{ + if (*opParamInfo_.returnSoftmaxLse) { + OP_CHECK_IF(opParamInfo_.softmaxSum.shape->GetStorageShape().GetShapeSize() == 0, + OP_LOGE(opName_, "When return_softmax_lse is true, softmaxSum tensor cannot be empty tensor."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.softmaxSum.desc->GetDataType() != ge::DT_FLOAT, + OP_LOGE(opName_, "softmaxSum's dtype must be FLOAT."), + return ge::GRAPH_FAILED); + } else { + if (opParamInfo_.softmaxSum.shape->GetStorageShape().GetShapeSize() != 0) { + OP_LOGW(opName_, "When return_softmax_lse is false, softmaxSum tensor must be empty tensor."); + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckQRope() +{ + if (ge::GRAPH_SUCCESS != CheckDTypeConsistency(opParamInfo_.queryRope.desc->GetDataType(), + inputQType_, QUERY_ROPE_NAME) || + ge::GRAPH_SUCCESS != CheckQRopeShape()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckTopK() +{ + if (ge::GRAPH_SUCCESS != CheckTopkShape()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckVAndKRopeShapeForBatchContinuous() +{ + SFATilingShapeCompareParam shapeParams; + shapeParams.B = bSize_; + shapeParams.N = n2Size_; + shapeParams.S = s2Size_; + shapeParams.T = kvTSize_; + shapeParams.D = qkHeadDim_; + if (CompareShape(shapeParams, keyShapeCmp_, kvLayout_, KEY_NAME) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + shapeParams.D = vHeadDim_; + if (CompareShape(shapeParams, valueShapeCmp_, kvLayout_, VALUE_NAME) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + shapeParams.D = ropeHeadDim_; + if (CompareShape(shapeParams, keyRopeShapeCmp_, kvLayout_, KEY_ROPE_NAME) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +uint32_t SFAMlaTiling::GetTypeSize(ge::DataType dtype) const +{ + uint32_t typeSize = NUM_BYTES_FLOAT16; + switch (dtype) { + case ge::DT_FLOAT16: + typeSize = NUM_BYTES_FLOAT16; + break; + case ge::DT_BF16: + typeSize = NUM_BYTES_BF16; + break; + default: + typeSize = NUM_BYTES_FLOAT16; + } + return typeSize; +} + +ge::graphStatus SFATilingCheck::CheckVAndKRopeShapeForPageAttention() +{ + int64_t blockNum = keyShapeCmp_.GetDim(0); + OP_CHECK_IF(blockNum <= 0, + OP_LOGE(opName_, "The first dim(%ld) of key should be greater than 0", blockNum), + return ge::GRAPH_FAILED); + SFATilingShapeCompareParam shapeParams; + shapeParams.Bn = blockNum; + shapeParams.N = n2Size_; + shapeParams.Bs = blockSize_; + shapeParams.D = vHeadDim_; + shapeParams.T = kvTSize_; + if (CompareShape(shapeParams, valueShapeCmp_, kvLayout_, VALUE_NAME) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + shapeParams.D = ropeHeadDim_; + if (CompareShape(shapeParams, keyRopeShapeCmp_, kvLayout_, KEY_ROPE_NAME) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckVAndKRopeShape() +{ + if (kvStorageMode_ == KvStorageMode::BATCH_CONTINUOUS) { + return CheckVAndKRopeShapeForBatchContinuous(); + } + + if (kvStorageMode_ == KvStorageMode::PAGE_ATTENTION) { + return CheckVAndKRopeShapeForPageAttention(); + } + + OP_LOGE(opName_, "storage mode of key and value is %u, it is incorrect.", static_cast(kvStorageMode_)); + return ge::GRAPH_FAILED; +} + +ge::graphStatus SFATilingCheck::CheckVAndKRope() +{ + if (ge::GRAPH_SUCCESS != CheckDTypeConsistency(opParamInfo_.value.desc->GetDataType(), + inputKvType_, VALUE_NAME) || + ge::GRAPH_SUCCESS != CheckDTypeConsistency(opParamInfo_.keyRope.desc->GetDataType(), + inputKvType_, KEY_ROPE_NAME) || ge::GRAPH_SUCCESS != CheckVAndKRopeShape()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckActualSeqLensQ() +{ + if (ge::GRAPH_SUCCESS != CheckActualSeqLensQDType() || + ge::GRAPH_SUCCESS != CheckActualSeqLensQShape()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckActualSeqLensQDType() +{ + if (opParamInfo_.actualSeqLengthsQ.tensor == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (opParamInfo_.actualSeqLengthsQ.desc == nullptr) { + OP_LOGE(opName_, "actualSeqLengthsQ is not empty," + "but actualSeqLengthsQ's dtype is nullptr."); + return ge::GRAPH_FAILED; + } + if (opParamInfo_.actualSeqLengthsQ.desc->GetDataType() != ge::DT_INT32) { + OP_LOGE(opName_, "actualSeqLengthsQ's dtype is %s, it should be DT_INT32.", + SFADataTypeToSerialString(opParamInfo_.actualSeqLengthsQ.desc->GetDataType()).c_str()); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckActualSeqLensQShape() +{ + if (opParamInfo_.actualSeqLengthsQ.tensor == nullptr) { + return ge::GRAPH_SUCCESS; + } + uint32_t shapeSize = 0; + if (GetActualSeqLenSize(shapeSize, opParamInfo_.actualSeqLengthsQ.tensor, qLayout_, "actualSeqLengthsQ") != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + if (shapeSize != bSize_) { + OP_LOGE(opName_, "actualSeqLengthsQ shape size is %u, it should be equal to batch size[%u]", + shapeSize, bSize_); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckActualSeqLens() +{ + if (std::string(opParamInfo_.layoutKV) == "TND" && opParamInfo_.actualSeqLengths.tensor == nullptr) { + OP_LOGE(opName_, + "when the layout of key and value is TND, " + "the actualSeqLengths of key and value should not be empty."); + return ge::GRAPH_PARAM_INVALID; + } + if (ge::GRAPH_SUCCESS != CheckActualSeqLensDType() || + ge::GRAPH_SUCCESS != CheckActualSeqLensShape()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckActualSeqLensDType() +{ + if (opParamInfo_.actualSeqLengths.tensor == nullptr) { + return ge::GRAPH_SUCCESS; + } + if (opParamInfo_.actualSeqLengths.desc == nullptr) { + OP_LOGE(opName_, "actualSeqLengths is not empty," + "but actualSeqLengths's dtype is nullptr."); + return ge::GRAPH_FAILED; + } + if (opParamInfo_.actualSeqLengths.desc->GetDataType() != ge::DT_INT32) { + OP_LOGE(opName_, "actualSeqLengths's dtype is %s, it should be DT_INT32.", + SFADataTypeToSerialString(opParamInfo_.actualSeqLengths.desc->GetDataType()).c_str()); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckActualSeqLensShape() +{ + if (opParamInfo_.actualSeqLengths.tensor == nullptr) { + return ge::GRAPH_SUCCESS; + } + uint32_t shapeSize = 0; + if(GetActualSeqLenSize(shapeSize, opParamInfo_.actualSeqLengths.tensor, kvLayout_, "actualSeqLengths") != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + if (shapeSize != bSize_) { + OP_LOGE(opName_, "actualSeqLengths shape size is %u, it should be equal to batch size[%u].", + shapeSize, bSize_); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckMultiParaConsistency() +{ + SetSFAShapeCompare(); + if (ge::GRAPH_SUCCESS != CheckVAndKRope() || + ge::GRAPH_SUCCESS != CheckQRope() || + ge::GRAPH_SUCCESS != CheckTopK() || + ge::GRAPH_SUCCESS != CheckAttenOut() || + ge::GRAPH_SUCCESS != CheckSoftmaxMax() || + ge::GRAPH_SUCCESS != CheckSoftmaxSum() || + ge::GRAPH_SUCCESS != CheckActualSeqLensQ() || + ge::GRAPH_SUCCESS != CheckActualSeqLens() || + ge::GRAPH_SUCCESS != CheckBlockTable()) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckFeatureMlaNoQuantShape() const +{ + OP_CHECK_IF(bSize_ <= 0, + OP_LOGE(opName_, "batch_size should be greater than 0, but got %u", bSize_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(qTSize_ <= 0 && (qLayout_ == SFALayout::TND), + OP_LOGE(opName_, "T_size of query should be greater than 0, but got %u", qTSize_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(n1Size_ <= 0, + OP_LOGE(opName_, "q_head_num should be greater than 0, but got %u", n1Size_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(n2Size_ != 1, + OP_LOGE(opName_, "kv_head_num should be 1, but got %u", n2Size_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(n1Size_ % n2Size_ != 0, + OP_LOGE(opName_, "q_head_num(%u) must be divisible by kv_head_num(%u)", n1Size_, n2Size_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(gSize_ < 1 || (gSize_ > 64 && gSize_ != 128), + OP_LOGE(opName_, "group num should be in 1 ~ 64, 128, but got %u", gSize_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(qkHeadDim_ != 512, + OP_LOGE(opName_, "qk_head_dim only support 512, but got %u", qkHeadDim_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(qkHeadDim_ != vHeadDim_, + OP_LOGE(opName_, "qk_head_dim[%u] should be equal to v_head_dim[%u]", qkHeadDim_, vHeadDim_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(ropeHeadDim_ != 64, + OP_LOGE(opName_, "rope_head_dim should be 64, but got %u", ropeHeadDim_), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckFeatureMlaNoQuantLayout() const +{ + const std::vector layoutSupportList = { + "BSND", + "TND" + }; + std::string layoutQuery = opParamInfo_.layoutQuery; + OP_CHECK_IF(std::find(layoutSupportList.begin(), layoutSupportList.end(), layoutQuery) == layoutSupportList.end(), + OP_LOGE(opName_, "layoutQuery only supports BSND/TND, but got %s", layoutQuery.c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckFeatureMlaNoQuantDtype() const +{ + OP_CHECK_IF(inputQType_ != ge::DT_BF16 && inputQType_ != ge::DT_FLOAT16, + OP_LOGE(opName_, "query dtype only support %s and %s, but got %s", + SFADataTypeToSerialString(ge::DT_BF16).c_str(), SFADataTypeToSerialString(ge::DT_FLOAT16).c_str(), + SFADataTypeToSerialString(inputQType_).c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckFeatureMlaNoquantPa() const +{ + if (kvStorageMode_ != KvStorageMode::PAGE_ATTENTION) { + return ge::GRAPH_SUCCESS; + } + + OP_CHECK_IF(blockSize_ <= 0 || blockSize_ > static_cast(MAX_BLOCK_SIZE), + OP_LOGE(opName_, "when page attention is enabled, block_size(%d) should be in range (0, %u].", + blockSize_, MAX_BLOCK_SIZE), return ge::GRAPH_FAILED); + + OP_CHECK_IF(blockSize_ % 16 > 0, + OP_LOGE(opName_, "when page attention is enabled, block_size(%d) should be 16-aligned.", + blockSize_), return ge::GRAPH_FAILED); + + OP_CHECK_IF(blockSize_ % sparseBlockSize_ > 0, + OP_LOGE(opName_, "when page attention is enabled, block_size(%d) must be divided by sparse_block_size(%d), but now the remainder is %d.", + blockSize_, sparseBlockSize_, blockSize_ % sparseBlockSize_), return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckFeatureMlaNoquant() const +{ + if (ge::GRAPH_SUCCESS != CheckFeatureMlaNoQuantShape() || + ge::GRAPH_SUCCESS != CheckFeatureMlaNoQuantLayout() || + ge::GRAPH_SUCCESS != CheckFeatureMlaNoQuantDtype() || + ge::GRAPH_SUCCESS != CheckFeatureMlaNoquantPa()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFATilingCheck::CheckFeatureMla() const +{ + return CheckFeatureMlaNoquant(); +} + +ge::graphStatus SFATilingCheck::CheckFeature() const +{ + return CheckFeatureMla(); +} + +void SFATilingCheck::Init() +{ + opName_ = sfaInfo_.opName; + platformInfo_ = sfaInfo_.platformInfo; + opParamInfo_ = sfaInfo_.opParamInfo; + npuArch_ = sfaInfo_.npuArch; + isA5_ = sfaInfo_.isA5; + + bSize_ = sfaInfo_.bSize; + n1Size_ = sfaInfo_.n1Size; + n2Size_ = sfaInfo_.n2Size; + s1Size_ = sfaInfo_.s1Size; + s2Size_ = sfaInfo_.s2Size; + gSize_ = sfaInfo_.gSize; + qkHeadDim_ = sfaInfo_.qkHeadDim; + vHeadDim_ = sfaInfo_.vHeadDim; + ropeHeadDim_ = sfaInfo_.ropeHeadDim; + maxBlockNumPerBatch_ = sfaInfo_.maxBlockNumPerBatch; + qTSize_ = sfaInfo_.qTSize; + kvTSize_ = sfaInfo_.kvTSize; + blockSize_ = sfaInfo_.blockSize; + sparseBlockCount_ = sfaInfo_.sparseBlockCount; + sparseBlockSize_ = sfaInfo_.sparseBlockSize; + + inputQType_ = sfaInfo_.inputQType; + inputKvType_ = sfaInfo_.inputKvType; + inputQRopeType_ = sfaInfo_.inputQRopeType; + inputKRopeType_ = sfaInfo_.inputKRopeType; + outputType_ = sfaInfo_.outputType; + + qLayout_ = sfaInfo_.qLayout; + topkLayout_ = sfaInfo_.topkLayout; + kvLayout_ = sfaInfo_.kvLayout; + outLayout_ = sfaInfo_.outLayout; + softmaxMaxLayout_ = sfaInfo_.softmaxMaxLayout; + softmaxSumLayout_ = sfaInfo_.softmaxSumLayout; + + kvStorageMode_ = sfaInfo_.kvStorageMode; + l2CacheSize_ = sfaInfo_.l2CacheSize; +} + +ge::graphStatus SFATilingCheck::Process() +{ + Init(); + if (CheckSinglePara() != ge::GRAPH_SUCCESS || + CheckParaExistence() != ge::GRAPH_SUCCESS || + CheckFeature() != ge::GRAPH_SUCCESS || + CheckMultiParaConsistency() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +bool SFAInfoParser::HasAxis(const SFAAxis &axis, const SFALayout &layout, const gert::Shape &shape) const +{ + const auto& layoutIt = SFA_LAYOUT_AXIS_MAP.find(layout); + if (layoutIt == SFA_LAYOUT_AXIS_MAP.end()) { + return false; + } + + const std::vector& axes = layoutIt->second; + const auto& axisIt = std::find(axes.begin(), axes.end(), axis); + if (axisIt == axes.end()) { + return false; + } + const auto& dimIt = SFA_LAYOUT_DIM_MAP.find(layout); + if (dimIt == SFA_LAYOUT_DIM_MAP.end() || dimIt->second != shape.GetDimNum()) { + return false; + } + return true; +} + +size_t SFAInfoParser::GetAxisIdx(const SFAAxis &axis, const SFALayout &layout) const +{ + const std::vector& axes = SFA_LAYOUT_AXIS_MAP.find(layout)->second; + const auto& axisIt = std::find(axes.begin(), axes.end(), axis); + return std::distance(axes.begin(), axisIt); +} + +uint32_t SFAInfoParser::GetAxisNum(const gert::Shape &shape, const SFAAxis &axis,const SFALayout &layout) const +{ + return HasAxis(axis, layout, shape) ? shape.GetDim(GetAxisIdx(axis, layout)) : invalidDimValue_; +} + +ge::graphStatus SFAInfoParser::CheckRequiredInOutExistence() const +{ + OP_CHECK_IF(opParamInfo_.query.shape == nullptr, OP_LOGE(opName_, "Shape of tensor query is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.query.desc == nullptr, OP_LOGE(opName_, "Desc of tensor query is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key.shape == nullptr, OP_LOGE(opName_, "Shape of tensor k is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key.desc == nullptr, OP_LOGE(opName_, "Desc of tensor k is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.value.shape == nullptr, OP_LOGE(opName_, "Shape of tensor value is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.value.desc == nullptr, OP_LOGE(opName_, "Desc of tensor value is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.sparseIndices.shape == nullptr, OP_LOGE(opName_, "Shape of tensor sparseIndices is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.sparseIndices.desc == nullptr, OP_LOGE(opName_, "Desc of tensor sparseIndices is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.attenOut.shape == nullptr, OP_LOGE(opName_, "Shape of tensor output is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.attenOut.desc == nullptr, OP_LOGE(opName_, "Desc of tensor output is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.softmaxMax.shape == nullptr, OP_LOGE(opName_, "Shape of tensor softmaxMax is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.softmaxMax.desc == nullptr, OP_LOGE(opName_, "Desc of tensor softmaxMax is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.softmaxSum.shape == nullptr, OP_LOGE(opName_, "Shape of tensor softmaxSum is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.softmaxSum.desc == nullptr, OP_LOGE(opName_, "Desc of tensor softmaxSum is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.queryRope.tensor == nullptr, OP_LOGE(opName_, "Shape of queryRope is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.queryRope.desc == nullptr, OP_LOGE(opName_, "Desc of queryRope is nullptr"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::CheckRequiredAttrExistence() const +{ + OP_CHECK_IF(opParamInfo_.layoutQuery == nullptr, OP_LOGE(opName_, "attr layoutQuery is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.layoutKV == nullptr, OP_LOGE(opName_, "attr layoutKV is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.sparseBlockSize == nullptr, OP_LOGE(opName_, "attr sparseBlockSize is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.scaleValue == nullptr, OP_LOGE(opName_, "attr scaleValue is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.sparseMode == nullptr, OP_LOGE(opName_, "attr sparseMode is nullptr"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::CheckRequiredParaExistence() const +{ + if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS || + CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + SFALayout &layout, const std::string &name) const +{ + if ((tensor == nullptr)) { + OP_LOGE(opName_, "when layout of query is %s, %s must be provided.", + SFALayoutToSerialString(layout).c_str(), name.c_str()); + return ge::GRAPH_FAILED; + } + int64_t shapeSize = tensor->GetShapeSize(); + if (shapeSize <= 0) { + OP_LOGE(opName_, "the shape size of %s is %ld, it should be greater than 0.", + name.c_str(), shapeSize); + return ge::GRAPH_FAILED; + } + size = static_cast(shapeSize); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetActualSeqLenQSize(uint32_t &size) +{ + return GetActualSeqLenSize(size, opParamInfo_.actualSeqLengthsQ.tensor, qLayout_, "actualSeqLengthsQ"); +} + +ge::graphStatus SFAInfoParser::GetOpName() +{ + if (context_->GetNodeName() == nullptr) { + OP_LOGE("SparseFlashAttention", "opName got from TilingContext is nullptr"); + return ge::GRAPH_FAILED; + } + opName_ = context_->GetNodeName(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetNpuInfo() +{ + platformInfo_ = context_->GetPlatformInfo(); + OP_CHECK_IF(platformInfo_ == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(opName_, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo_); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + OP_CHECK_IF(aicNum == 0 || aivNum == 0, + OPS_REPORT_VECTOR_INNER_ERR(opName_, "num of core obtained is 0."), return GRAPH_FAILED); + + npuArch_ = ascendcPlatform.GetCurNpuArch(); + isA5_ = (npuArch_ == NpuArch::DAV_3510); + if (npuArch_ != NpuArch::DAV_2201 && npuArch_ != NpuArch::DAV_3510) { + OPS_REPORT_VECTOR_INNER_ERR(opName_, "Npu Arch Version[%d] is not support.", static_cast(npuArch_)); + return GRAPH_FAILED; + } + + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L2, l2CacheSize_); + + return ge::GRAPH_SUCCESS; +} + +void SFAInfoParser::GetOptionalInputParaInfo() +{ + opParamInfo_.blockTable.tensor = context_->GetOptionalInputTensor(BLOCK_TABLE_INPUT_INDEX); + opParamInfo_.blockTable.desc = context_->GetOptionalInputDesc(BLOCK_TABLE_INPUT_INDEX); + opParamInfo_.actualSeqLengthsQ.tensor = context_->GetOptionalInputTensor(ACT_SEQ_LEN_Q_INPUT_INDEX); + opParamInfo_.actualSeqLengthsQ.desc = context_->GetOptionalInputDesc(ACT_SEQ_LEN_Q_INPUT_INDEX); + opParamInfo_.actualSeqLengths.tensor = context_->GetOptionalInputTensor(ACT_SEQ_LEN_KV_INPUT_INDEX); + opParamInfo_.actualSeqLengths.desc = context_->GetOptionalInputDesc(ACT_SEQ_LEN_KV_INPUT_INDEX); + opParamInfo_.queryRope.tensor = context_->GetOptionalInputTensor(QUERY_ROPE_INPUT_INDEX); + opParamInfo_.queryRope.desc = context_->GetOptionalInputDesc(QUERY_ROPE_INPUT_INDEX); + opParamInfo_.keyRope.tensor = context_->GetOptionalInputTensor(KEY_ROPE_INPUT_INDEX); + opParamInfo_.keyRope.desc = context_->GetOptionalInputDesc(KEY_ROPE_INPUT_INDEX); +} + +void SFAInfoParser::GetInputParaInfo() +{ + opParamInfo_.query.desc = context_->GetInputDesc(QUERY_INPUT_INDEX); + opParamInfo_.query.shape = context_->GetInputShape(QUERY_INPUT_INDEX); + opParamInfo_.key.desc = context_->GetInputDesc(KEY_INPUT_INDEX); + opParamInfo_.key.shape = context_->GetInputShape(KEY_INPUT_INDEX); + opParamInfo_.value.desc = context_->GetInputDesc(VALUE_INPUT_INDEX); + opParamInfo_.value.shape = context_->GetInputShape(VALUE_INPUT_INDEX); + opParamInfo_.sparseIndices.desc = context_->GetInputDesc(SPARSE_INDICES_INPUT_INDEX); + opParamInfo_.sparseIndices.shape = context_->GetInputShape(SPARSE_INDICES_INPUT_INDEX); + GetOptionalInputParaInfo(); +} + +void SFAInfoParser::GetOutputParaInfo() +{ + opParamInfo_.attenOut.desc = context_->GetOutputDesc(OUTPUT_INDEX); + opParamInfo_.attenOut.shape = context_->GetOutputShape(OUTPUT_INDEX); + opParamInfo_.softmaxMax.desc = context_->GetOutputDesc(SOFTMAXMAX_INDEX); + opParamInfo_.softmaxMax.shape = context_->GetOutputShape(SOFTMAXMAX_INDEX); + opParamInfo_.softmaxSum.desc = context_->GetOutputDesc(SOFTMAXSUM_INDEX); + opParamInfo_.softmaxSum.shape = context_->GetOutputShape(SOFTMAXSUM_INDEX); +} + +ge::graphStatus SFAInfoParser::GetAttrParaInfo() +{ + auto attrs = context_->GetAttrs(); + OP_CHECK_IF(attrs == nullptr, OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "attrs got from ge is nullptr"), + return ge::GRAPH_FAILED); + + opParamInfo_.layoutQuery = attrs->GetStr(LAYOUT_QUERY_ATTR_INDEX); + opParamInfo_.layoutKV = attrs->GetStr(LAYOUT_KV_ATTR_INDEX); + opParamInfo_.sparseBlockSize = attrs->GetAttrPointer(SPARSE_BLOCK_SIZE_ATTR_INDEX); + opParamInfo_.scaleValue = attrs->GetAttrPointer(SCALE_VALUE_ATTR_INDEX); + opParamInfo_.sparseMode = attrs->GetAttrPointer(SPARSE_MODE_ATTR_INDEX); + opParamInfo_.preTokens = attrs->GetAttrPointer(PRE_TOKENS_ATTR_INDEX); + opParamInfo_.nextTokens = attrs->GetAttrPointer(NEXT_TOKENS_ATTR_INDEX); + opParamInfo_.attentionMode = attrs->GetAttrPointer(ATTENTION_MODE_ATTR_INDEX); + opParamInfo_.returnSoftmaxLse = attrs->GetAttrPointer(RETURN_SOFTMAX_LSE_ATTR_INDEX); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetOpParaInfo() +{ + GetInputParaInfo(); + GetOutputParaInfo(); + if (ge::GRAPH_SUCCESS != GetAttrParaInfo()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetInOutDataType() +{ + inputQType_ = opParamInfo_.query.desc->GetDataType(); + inputKvType_ = opParamInfo_.key.desc->GetDataType(); + outputType_ = opParamInfo_.attenOut.desc->GetDataType(); + if (opParamInfo_.queryRope.desc != nullptr) { + inputQRopeType_ = opParamInfo_.queryRope.desc->GetDataType(); + } + if (opParamInfo_.keyRope.desc != nullptr) { + inputKRopeType_ = opParamInfo_.keyRope.desc->GetDataType(); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetBatchSize() +{ + // 获取B基准值 + // 1、非TND时, 以query的batch_size维度为基准; + // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小 + if (qLayout_ == SFALayout::TND) { + return GetActualSeqLenQSize(bSize_); + } else { // BSND + bSize_ = GetAxisNum(queryShape_, SFAAxis::B, qLayout_); + return ge::GRAPH_SUCCESS; + } +} + +ge::graphStatus SFAInfoParser::GetQTSize() +{ + // 获取query的T基准值 + // 1、非TND时, 以query的batch_size维度为基准; + // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小 + qTSize_ = (qLayout_ == SFALayout::TND) ? GetAxisNum(queryShape_, SFAAxis::T, qLayout_) : 0; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetKVTSize() +{ + // 获取query的T基准值 + // 1、非TND时, 以key的batch_size维度为基准; + // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小 + kvTSize_ = (kvLayout_ == SFALayout::TND) ? GetAxisNum(keyShape_, SFAAxis::T, kvLayout_) : 0; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetQkHeadDim() +{ + // 获取qkHeadDim基准值 + // 以query的D维度为基准 + qkHeadDim_ = GetAxisNum(queryShape_, SFAAxis::D, qLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetS1Size() +{ + // 获取S1基准值 + // 1、非TND时, 以query的S维度为基准; + // 2、TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组中的最大值为基准 + if (qLayout_ == SFALayout::TND) { + s1Size_ = GetAxisNum(queryShape_, SFAAxis::T, qLayout_); + return ge::GRAPH_SUCCESS; + } else { // BSND + s1Size_ = GetAxisNum(queryShape_, SFAAxis::S, qLayout_); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetKvStorageMode() +{ + if (kvLayout_ == SFALayout::PA_BSND) { + kvStorageMode_ = KvStorageMode::PAGE_ATTENTION; + } else { + kvStorageMode_ = KvStorageMode::BATCH_CONTINUOUS; + } + // kv存储模式基准值 + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetKvLayout() +{ + const map layoutKVMap = { + {"BSND", SFALayout::BSND}, + {"PA_BSND", SFALayout::PA_BSND}, + {"TND", SFALayout::TND} + }; + + std::string layout(opParamInfo_.layoutKV); + auto it = layoutKVMap.find(layout); + if (it != layoutKVMap.end()) { + kvLayout_ = it->second; + } else { + OP_LOGE(opName_, "layoutKV is %s, it is unsupported.", layout.c_str()); + return ge::GRAPH_FAILED; + } + if (kvLayout_ != SFALayout::PA_BSND && qLayout_ != kvLayout_) { + OP_LOGE(opName_, "When layoutKV is not PA_BSND, layoutKV must be the same as layoutQ."); + return ge::GRAPH_FAILED; + } + uint32_t keyDimNum = opParamInfo_.key.shape->GetStorageShape().GetDimNum(); + if (kvLayout_ == SFALayout::PA_BSND && keyDimNum != 4U) { + OP_LOGE(opName_, "When layoutKV is PA_BSND, kvDimNum must be 4, but now is %d.", keyDimNum); + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetS2SizeForBatchContinuous() +{ + if (kvLayout_ == SFALayout::BSND) { // BSND + s2Size_ = GetAxisNum(keyShape_, SFAAxis::S, kvLayout_); + } else if (kvLayout_ == SFALayout::TND) { + s2Size_ = GetAxisNum(keyShape_, SFAAxis::T, kvLayout_); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetMaxBlockNumPerBatch() +{ + if (opParamInfo_.blockTable.tensor == nullptr) { + OP_LOGE(opName_, "the layout_kv is %s, blockTable must be provided.", SFALayoutToSerialString(kvLayout_).c_str()); + return ge::GRAPH_FAILED; + } + uint32_t dimNum = opParamInfo_.blockTable.tensor->GetStorageShape().GetDimNum(); + if (dimNum != DIM_NUM_TWO) { + OP_LOGE(opName_, "the dim num of block_table is %u, it should be %u.", dimNum, DIM_NUM_TWO); + return ge::GRAPH_FAILED; + } + if (opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1) <= 0) { + OP_LOGE(opName_, "%s's second dimension(%ld) should be greater than 0", + BLOCK_TABLE_NAME.c_str(), opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1)); + return ge::GRAPH_FAILED; + } + maxBlockNumPerBatch_ = opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetBlockSize() +{ + blockSize_ = GetAxisNum(keyShape_, SFAAxis::Bs, kvLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetSparseBlockCount() +{ + sparseBlockCount_ = GetAxisNum(sparseIndicesShape_, SFAAxis::K, qLayout_); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetS2SizeForPageAttention() +{ + if (GetMaxBlockNumPerBatch() != ge::GRAPH_SUCCESS || GetBlockSize() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + s2Size_ = maxBlockNumPerBatch_ * blockSize_; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetS2Size() +{ + // 获取S2基准值 + // 1、BATCH_CONTINUOUS时, 从key的S轴获取 + // 2、PAGE_ATTENTION时, S2 = block_table.dim1 * block_size + if (kvStorageMode_ == KvStorageMode::BATCH_CONTINUOUS) { + return GetS2SizeForBatchContinuous(); + } + return GetS2SizeForPageAttention(); +} + +ge::graphStatus SFAInfoParser::GetValueHeadDim() +{ + // 获取vHeadDim基准值 + // 以value的D维度为基准 + vHeadDim_ = GetAxisNum(valueShape_, SFAAxis::D, kvLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetRopeHeadDim() +{ + if (queryShape_.GetDimNum() != queryRopeShape_.GetDimNum()) { + OP_LOGE(opName_, "The dimensions of query and query_rope should be equal, but query has dimension %zu while query_rope has dimension %zu.", + queryShape_.GetDimNum(), queryRopeShape_.GetDimNum()); + return ge::GRAPH_PARAM_INVALID; + } + ropeHeadDim_ = GetAxisNum(queryRopeShape_, SFAAxis::D, qLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetQueryAndOutLayout() +{ + // 获取query和attentionOut的Layout基准值 + // layoutQuery: {qLayout, outLayout} + const map> layoutMap = { + {"BSND", {SFALayout::BSND, SFALayout::BSND}}, + {"TND", {SFALayout::TND, SFALayout::TND }}, + }; + + std::string layout(opParamInfo_.layoutQuery); + auto it = layoutMap.find(layout); + if (it != layoutMap.end()) { + qLayout_ = it->second.first; + outLayout_ = it->second.second; + } else { + OP_LOGE(opName_, "layoutQuery is %s, it is unsupported.", layout.c_str()); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetTopkLayout() +{ + topkLayout_ = qLayout_; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetSoftmaxMaxAndSumLayout() +{ + if (qLayout_ == SFALayout::BSND) { + softmaxMaxLayout_ = SFALayout::BNSG; + softmaxSumLayout_ = SFALayout::BNSG; + } else if (qLayout_ == SFALayout::TND) { + softmaxMaxLayout_ = SFALayout::NTG; + softmaxSumLayout_ = SFALayout::NTG; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetN1Size() +{ + n1Size_ = GetAxisNum(queryShape_, SFAAxis::N, qLayout_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetN2Size() +{ + n2Size_ = GetAxisNum(keyShape_, SFAAxis::N, kvLayout_); + return ge::GRAPH_SUCCESS; +} + +void SFAInfoParser::SetSFAShape() +{ + queryShape_ = opParamInfo_.query.shape->GetStorageShape(); + keyShape_ = opParamInfo_.key.shape->GetStorageShape(); + valueShape_ = opParamInfo_.value.shape->GetStorageShape(); + sparseIndicesShape_ = opParamInfo_.sparseIndices.shape->GetStorageShape(); + queryRopeShape_ = opParamInfo_.queryRope.tensor->GetStorageShape(); +} + +ge::graphStatus SFAInfoParser::GetGSize() +{ + if (n2Size_ != 0) { + gSize_ = n1Size_ / n2Size_; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus SFAInfoParser::GetActualseqInfo() +{ + maxActualseq_ = static_cast(s2Size_); + if (opParamInfo_.actualSeqLengths.tensor != nullptr) { + actualLenDimsKV_ = opParamInfo_.actualSeqLengths.tensor->GetShapeSize(); + } + if (opParamInfo_.actualSeqLengthsQ.tensor != nullptr) { + actualLenDimsQ_ = opParamInfo_.actualSeqLengthsQ.tensor->GetShapeSize(); + } + return ge::GRAPH_SUCCESS; +} + +void SFAInfoParser::GenerateInfo(SFATilingInfo &sfaInfo) +{ + sfaInfo.opName = opName_; + sfaInfo.platformInfo = platformInfo_; + sfaInfo.opParamInfo = opParamInfo_; + sfaInfo.npuArch = npuArch_; + sfaInfo.isA5 = isA5_; + + sfaInfo.bSize = bSize_; + sfaInfo.n1Size = n1Size_; + sfaInfo.n2Size = n2Size_; + sfaInfo.s1Size = s1Size_; + sfaInfo.s2Size = s2Size_; + sfaInfo.gSize = gSize_; + sfaInfo.qkHeadDim = qkHeadDim_; + sfaInfo.vHeadDim = vHeadDim_; + sfaInfo.ropeHeadDim = ropeHeadDim_; + sfaInfo.qTSize = qTSize_; + sfaInfo.kvTSize = kvTSize_; + sfaInfo.sparseBlockSize = *opParamInfo_.sparseBlockSize; + sfaInfo.sparseBlockCount = sparseBlockCount_; + + sfaInfo.inputQType = inputQType_; + sfaInfo.inputKvType = inputKvType_; + sfaInfo.inputQRopeType = inputQRopeType_; + sfaInfo.inputKRopeType = inputKRopeType_; + sfaInfo.outputType = outputType_; + + sfaInfo.kvStorageMode = kvStorageMode_; + sfaInfo.l2CacheSize = l2CacheSize_; + + sfaInfo.totalBlockNum = opParamInfo_.key.shape->GetStorageShape().GetDim(0); + sfaInfo.scaleValue = *opParamInfo_.scaleValue; + sfaInfo.pageAttentionFlag = (kvStorageMode_ == KvStorageMode::PAGE_ATTENTION); + sfaInfo.blockSize = blockSize_; + sfaInfo.blockTypeSize = sizeof(float); + sfaInfo.maxBlockNumPerBatch = maxBlockNumPerBatch_; + + sfaInfo.actualLenDimsQ = actualLenDimsQ_; + sfaInfo.actualLenDimsKV = actualLenDimsKV_; + sfaInfo.maxActualseq = maxActualseq_; + sfaInfo.isSameSeqAllKVTensor = isSameSeqAllKVTensor_; + sfaInfo.isSameActualseq = isSameActualseq_; + + sfaInfo.actualQSeqLenFlag = (opParamInfo_.actualSeqLengthsQ.tensor != nullptr); + sfaInfo.actualSeqLenFlag = (opParamInfo_.actualSeqLengths.tensor != nullptr); + + sfaInfo.sparseMode = *opParamInfo_.sparseMode; + sfaInfo.preTokens = *opParamInfo_.preTokens; + sfaInfo.nextTokens = *opParamInfo_.nextTokens; + sfaInfo.attentionMode = *opParamInfo_.attentionMode; + sfaInfo.returnSoftmaxLse = *opParamInfo_.returnSoftmaxLse; + + sfaInfo.qLayout = qLayout_; + sfaInfo.topkLayout = topkLayout_; + sfaInfo.kvLayout = kvLayout_; + sfaInfo.outLayout = outLayout_; + sfaInfo.softmaxMaxLayout = softmaxMaxLayout_; + sfaInfo.softmaxSumLayout = softmaxSumLayout_; +} + +ge::graphStatus SFAInfoParser::Parse(SFATilingInfo &sfaInfo) +{ + if (context_ == nullptr) { + OP_LOGE("SparseFlashAttention", "tiling context is nullptr!"); + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetOpName() || + ge::GRAPH_SUCCESS != GetNpuInfo() || + ge::GRAPH_SUCCESS != GetOpParaInfo() || + ge::GRAPH_SUCCESS != CheckRequiredParaExistence()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetInOutDataType() || + ge::GRAPH_SUCCESS != GetQueryAndOutLayout() || + ge::GRAPH_SUCCESS != GetTopkLayout() || + ge::GRAPH_SUCCESS != GetSoftmaxMaxAndSumLayout() || + ge::GRAPH_SUCCESS != GetKvLayout() || + ge::GRAPH_SUCCESS != GetKvStorageMode()) { + return ge::GRAPH_FAILED; + } + + SetSFAShape(); + if ( + ge::GRAPH_SUCCESS != GetN1Size() || + ge::GRAPH_SUCCESS != GetN2Size() || + ge::GRAPH_SUCCESS != GetGSize() || + ge::GRAPH_SUCCESS != GetBatchSize() || + ge::GRAPH_SUCCESS != GetQTSize() || + ge::GRAPH_SUCCESS != GetKVTSize() || + ge::GRAPH_SUCCESS != GetS1Size() || + ge::GRAPH_SUCCESS != GetQkHeadDim() || + ge::GRAPH_SUCCESS != GetS2Size() || + ge::GRAPH_SUCCESS != GetValueHeadDim() || + ge::GRAPH_SUCCESS != GetRopeHeadDim() || + ge::GRAPH_SUCCESS != GetSparseBlockCount()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetActualseqInfo()) { + return ge::GRAPH_FAILED; + } + + GenerateInfo(sfaInfo); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(SparseFlashAttention) + .Tiling(TilingSparseFlashAttention) + .TilingParse(TilingPrepareForSparseFlashAttention); +} // namespace optiling diff --git a/csrc/ascend/attention/sparse_flash_attention/op_host/sparse_flash_attention_tiling.h b/csrc/ascend/attention/sparse_flash_attention/op_host/sparse_flash_attention_tiling.h new file mode 100644 index 000000000..32ae1519e --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_host/sparse_flash_attention_tiling.h @@ -0,0 +1,639 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_flash_attention_tiling.h + * \brief + */ +#ifndef SPARSE_FLASH_ATTENTION_TILING_H +#define SPARSE_FLASH_ATTENTION_TILING_H + +#include +#include +#include +#include +#include "register/tilingdata_base.h" +#include "exe_graph/runtime/tiling_context.h" +#include "platform/soc_spec.h" + +namespace optiling { +// ------------------算子原型索引常量定义---------------- +// Inputs Index +constexpr uint32_t QUERY_INPUT_INDEX = 0; +constexpr uint32_t KEY_INPUT_INDEX = 1; +constexpr uint32_t VALUE_INPUT_INDEX = 2; +constexpr uint32_t SPARSE_INDICES_INPUT_INDEX = 3; +constexpr uint32_t BLOCK_TABLE_INPUT_INDEX = 4; +constexpr uint32_t ACT_SEQ_LEN_Q_INPUT_INDEX = 5; +constexpr uint32_t ACT_SEQ_LEN_KV_INPUT_INDEX = 6; +constexpr uint32_t QUERY_ROPE_INPUT_INDEX = 7; +constexpr uint32_t KEY_ROPE_INPUT_INDEX = 8; +// Outputs Index +constexpr uint32_t OUTPUT_INDEX = 0; +constexpr uint32_t SOFTMAXMAX_INDEX = 1; +constexpr uint32_t SOFTMAXSUM_INDEX = 2; + +// Attributes Index +constexpr uint32_t SCALE_VALUE_ATTR_INDEX = 0; +constexpr uint32_t SPARSE_BLOCK_SIZE_ATTR_INDEX = 1; +constexpr uint32_t LAYOUT_QUERY_ATTR_INDEX = 2; +constexpr uint32_t LAYOUT_KV_ATTR_INDEX = 3; +constexpr uint32_t SPARSE_MODE_ATTR_INDEX = 4; +constexpr uint32_t PRE_TOKENS_ATTR_INDEX = 5; +constexpr uint32_t NEXT_TOKENS_ATTR_INDEX = 6; +constexpr uint32_t ATTENTION_MODE_ATTR_INDEX = 7; +constexpr uint32_t RETURN_SOFTMAX_LSE_ATTR_INDEX = 8; +// Dim Num +constexpr size_t DIM_NUM_TWO = 2; +constexpr size_t DIM_NUM_THREE = 3; +constexpr size_t DIM_NUM_FOUR = 4; +// 常量 +constexpr uint32_t MAX_BLOCK_SIZE = 1024; +constexpr uint32_t COPYND2NZ_SRC_STRIDE_LIMITATION = 65535; +constexpr uint32_t NUM_BYTES_FLOAT = 4; +constexpr uint32_t NUM_BYTES_FLOAT16 = 2; +constexpr uint32_t NUM_BYTES_BF16 = 2; +constexpr uint32_t BYTE_BLOCK = 32; +const uint32_t SFA_MAX_AIC_CORE_NUM = 26; // 25 + 1 保证数组8字节对齐 + +// ------------------公共定义-------------------------- +enum class SFALayout : uint32_t { + BSND = 0, + TND = 1, + PA_BSND = 2, + BNSG = 3, + NTG = 4 +}; + +struct SFATilingShapeCompareParam { + int64_t B = 1; + int64_t S = 1; + int64_t N = 1; + int64_t D = 1; + int64_t T = 1; + int64_t G = 1; + // PA + int64_t Bs = 1; + int64_t Bn = 1; +}; + +enum class KvStorageMode : uint32_t { + BATCH_CONTINUOUS = 0, + PAGE_ATTENTION = 1 +}; + +enum class SFAPerfMode : uint32_t { + C_TEMPLATE_MODE = 0, + V_TEMPLATE_MODE +}; + +enum class SFAAxis : uint32_t { + B = 0, + S = 1, + N = 2, + D = 3, + K = 3, // sparse_indices的K和key的D枚举值相同，表达相同位置, 最后一维 + T = 5, + Bn = 6, // block number + Bs = 7, // block size + G = 8, +}; + +struct SFARequiredParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; +}; + +struct SFAOptionalParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::Tensor *tensor; +}; + +// -----------算子Tiling入参结构体定义--------------- +struct SFAParaInfo { + SFARequiredParaInfo query = {nullptr, nullptr}; + SFARequiredParaInfo key = {nullptr, nullptr}; + SFARequiredParaInfo value = {nullptr, nullptr}; + SFARequiredParaInfo sparseIndices = {nullptr, nullptr}; + SFAOptionalParaInfo blockTable = {nullptr, nullptr}; + SFAOptionalParaInfo actualSeqLengthsQ = {nullptr, nullptr}; + SFAOptionalParaInfo actualSeqLengths = {nullptr, nullptr}; + SFAOptionalParaInfo queryRope = {nullptr, nullptr}; + SFAOptionalParaInfo keyRope = {nullptr, nullptr}; + SFARequiredParaInfo attenOut = {nullptr, nullptr}; + SFARequiredParaInfo softmaxMax = {nullptr, nullptr}; + SFARequiredParaInfo softmaxSum = {nullptr, nullptr}; + + const char *layoutQuery = nullptr; + const char *layoutKV = nullptr; + const int64_t *sparseBlockSize = nullptr; + const float *scaleValue = nullptr; + const int64_t *sparseMode = nullptr; + const int64_t *preTokens = nullptr; + const int64_t *nextTokens = nullptr; + const int64_t *attentionMode = nullptr; + const bool *returnSoftmaxLse = nullptr; +}; + +struct InnerSplitParams { + uint32_t s1GBaseSize = 1; + uint32_t s2BaseSize = 1; +}; + +// -----------算子TilingData定义--------------- +BEGIN_TILING_DATA_DEF(SparseFlashAttentionBaseParamsMla) +TILING_DATA_FIELD_DEF(uint32_t, batchSize) +TILING_DATA_FIELD_DEF(uint32_t, seqSize) +TILING_DATA_FIELD_DEF(uint32_t, qSeqSize) +TILING_DATA_FIELD_DEF(int64_t, blockSize) +TILING_DATA_FIELD_DEF(uint32_t, maxBlockNumPerBatch) +TILING_DATA_FIELD_DEF(float, scaleValue) +TILING_DATA_FIELD_DEF(uint32_t, nNumOfQInOneGroup) +TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsQ) +TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsKV) +TILING_DATA_FIELD_DEF(uint32_t, outputLayout) +TILING_DATA_FIELD_DEF(uint32_t, sparseMode) +TILING_DATA_FIELD_DEF(int64_t, preTokens) +TILING_DATA_FIELD_DEF(int64_t, nextTokens) +TILING_DATA_FIELD_DEF(uint32_t, attentionMode) +TILING_DATA_FIELD_DEF(uint32_t, returnSoftmaxLse) +TILING_DATA_FIELD_DEF(int64_t, sparseBlockSize) +TILING_DATA_FIELD_DEF(uint32_t, sparseBlockCount) +TILING_DATA_FIELD_DEF(uint32_t, isActualLenDimsNull) +TILING_DATA_FIELD_DEF(uint32_t, isActualLenDimsKVNull) +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(SparseFlashAttentionBaseParamsMlaOp, SparseFlashAttentionBaseParamsMla) + +BEGIN_TILING_DATA_DEF(SparseFlashAttentionSingleCoreParamsMla) +TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum); +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(SparseFlashAttentionSingleCoreParamsMlaOp, SparseFlashAttentionSingleCoreParamsMla) + +BEGIN_TILING_DATA_DEF(SparseFlashAttentionSingleCoreTensorSizeMla) +TILING_DATA_FIELD_DEF(uint32_t, mmResUbSize); +TILING_DATA_FIELD_DEF(uint32_t, bmm2ResUbSize); +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(SparseFlashAttentionSingleCoreTensorSizeMlaOp, SparseFlashAttentionSingleCoreTensorSizeMla) + +BEGIN_TILING_DATA_DEF(SparseFlashAttentionSplitKVParamsMla) +TILING_DATA_FIELD_DEF(uint32_t, s2) // S2切分份数 +TILING_DATA_FIELD_DEF(uint32_t, accumOutSize) // FD workspace +TILING_DATA_FIELD_DEF(uint32_t, logSumExpSize) // FD workspace +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(SparseFlashAttentionSplitKVParamsMlaOp, SparseFlashAttentionSplitKVParamsMla) + +// 内切基本块参数 +BEGIN_TILING_DATA_DEF(SparseFlashAttentionInnerSplitParams) +TILING_DATA_FIELD_DEF(uint32_t, mBaseSize) +TILING_DATA_FIELD_DEF(uint32_t, s2BaseSize) +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(SparseFlashAttentionInnerSplitParamsOp, SparseFlashAttentionInnerSplitParams) + +BEGIN_TILING_DATA_DEF(SparseFlashAttentionTilingDataMla) +TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionBaseParamsMla, baseParams); +TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionSplitKVParamsMla, splitKVParams); +TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionSingleCoreParamsMla, singleCoreParams); +TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionSingleCoreTensorSizeMla, singleCoreTensorSize); +TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionInnerSplitParams, innerSplitParams); +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(SparseFlashAttention, SparseFlashAttentionTilingDataMla) + +template inline T Align(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd) * (rnd))); +} + +template +std::string SFAShape2String(const T &shape) +{ + std::ostringstream oss; + oss << "["; + if (shape.GetDimNum() > 0) { + for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) { + oss << shape.GetDim(i) << ", "; + } + oss << shape.GetDim(shape.GetDimNum() - 1); + } + oss << "]"; + return oss.str(); +} + +static std::string GetShapeStr(gert::Shape shape); +static std::string SFADataTypeToSerialString(ge::DataType type); +std::string SFATensorDesc2String(const gert::StorageShape *shape, const gert::CompileTimeTensorDesc *tensor); +std::string SFADebugTilingContext(const gert::TilingContext *context); +std::string SFALayoutToSerialString(SFALayout layout); + +// -----------算子Tiling入参信息类--------------- +struct SFATilingInfo { + const char *opName = nullptr; + fe::PlatFormInfos *platformInfo = nullptr; + SFAParaInfo opParamInfo; + + // Base Param + NpuArch npuArch = NpuArch::DAV_2201; + bool isA5 = false; + uint32_t bSize = 0; + uint32_t n1Size = 0; + uint32_t n2Size = 0; + uint32_t s1Size = 0; + int64_t s2Size = 0; + uint32_t qkHeadDim = 0; + uint32_t vHeadDim = 0; + uint32_t gSize = 0; + uint32_t ropeHeadDim = 0; + uint32_t qTSize = 0; // 仅TND时生效 + uint32_t kvTSize = 0; // 仅TND时生效 + float scaleValue = 0; + uint32_t innerPrecise = 0; + uint32_t l2CacheOffFlag = 0; + int64_t sparseBlockSize = 0; + int64_t sparseBlockCount = 0; + + bool pageAttentionFlag = false; + int64_t blockSize = 0; + uint32_t blockTypeSize = 0; + uint32_t maxBlockNumPerBatch = 0; + uint32_t totalBlockNum = 0; + + uint32_t actualLenDimsQ = 0; + uint32_t maxActualseq = 0; + + bool actualQSeqLenFlag = false; + bool actualSeqLenFlag = false; + bool isSameSeqAllKVTensor = true; + bool isSameActualseq = true; + uint32_t actualLenDimsKV = 0; + std::vector kvListSeqLens {}; + + uint32_t sparseMode = 0; + int64_t preTokens = INT64_MAX; + int64_t nextTokens = INT64_MAX; + uint32_t attentionMode = 2; + bool returnSoftmaxLse = false; + + ge::DataType inputQType = ge::DT_FLOAT16; + ge::DataType inputKvType = ge::DT_FLOAT16; + ge::DataType outputType = ge::DT_FLOAT16; + + KvStorageMode kvStorageMode = KvStorageMode::BATCH_CONTINUOUS; + + SFALayout qLayout = SFALayout::BSND; + SFALayout topkLayout = SFALayout::BSND; + SFALayout outLayout = SFALayout::BSND; + SFALayout kvLayout = SFALayout::BSND; + SFALayout softmaxMaxLayout = SFALayout::BNSG; + SFALayout softmaxSumLayout = SFALayout::BNSG; + + ge::DataType inputQRopeType = ge::DT_FLOAT16; + ge::DataType inputKRopeType = ge::DT_FLOAT16; + + uint64_t l2CacheSize = 0; +}; + +// ---------------算子Tiling类--------------- +class SFAMlaTiling { +public: + explicit SFAMlaTiling(gert::TilingContext *context) : context_(context) {} + ge::graphStatus DoOpTiling(SFATilingInfo *sfaInfo); + +private: + ge::graphStatus SetBlockDim(uint32_t blockDim) const; + ge::graphStatus SetTilingKey(uint64_t tilingKey) const; + ge::graphStatus SetWorkspaceSize(uint64_t workspaceSize) const; + ge::graphStatus SetTilingData(TilingDef &tilingData) const; + gert::TilingContext *context_ = nullptr; + ge::graphStatus GetPlatformInfo(); + void GenTilingKey(); + bool DealSameSeqEachBatch(); + + void ZeroTensorProcess() const; + void InitParams(); + + void Split(); + bool IsBalanceSplitCore(); + + void SplitBalanced(); + void CalcInnerSize(uint32_t s2Size); + + bool IsFlashDecode(uint32_t coreNum); + + void FillTilingBaseParamsMla(); + void FillTilingSplitKVMla(); + + void FillTilingSingleCoreParamsMla(); + void FillTilingSingleCoreTensorSizeMla(); + void FillTiling(); + + void CalcUbBmm(); + void CheckUbSpace(); + void NormalCalcFDWorkSpace(const uint32_t actCoreNum); + void CalcFDWorkSpace(const uint32_t actCoreNum); + void GetWorkspaceSize(); + + uint32_t CalcBalanceFDParamNums(const uint32_t actCoreNum) const; + + void CalcBlockDim(); + + uint32_t GetTypeSize(ge::DataType dtype) const; + + bool balanceModeFlag_ = false; + bool splitKVFlag_ = false; + + uint32_t coreNum_ = 0; + SFAPerfMode perfMode_ = SFAPerfMode::V_TEMPLATE_MODE; + uint32_t kvSplitPart_ = 1; + size_t mmResUbSize_ = 0; + size_t bmm2ResUbSize_ = 0; + size_t qPreSizeMla_= 0; + uint32_t sInnerLoopTimes_ = 0; + uint32_t sInnerSize_ = 0; + uint32_t sInnerSizeTail_ = 0; + uint32_t sInnerSizeAlign_ = 0; + uint32_t kvSplit_ = 0; + uint32_t usedCoreNum_ = 0; + uint32_t formerCoreNum_ = 0; + uint32_t blockSplitBn2Range_ = 0; + uint32_t tailSplitedBatchRange_ = 0; + + uint32_t aicNum_ = 0; + uint32_t aivNum_ = 0; + size_t libapiSize_ = 0; + + SparseFlashAttentionTilingDataMla tilingData_; + uint32_t blockDim_{0}; + uint64_t workspaceSize_{0}; + uint64_t tilingKey_{0}; + + uint32_t headDimAlign_ = 0; + uint32_t mBaseSize_ = 128; + uint32_t mFdBaseSize_ = 8; + + SFATilingInfo *sfaInfo_ = nullptr; +}; + +// -----------算子Tiling入参信息解析及Check类--------------- +class SFATilingCheck { +public: + explicit SFATilingCheck(const SFATilingInfo &sfaInfo) : sfaInfo_(sfaInfo) {}; + ~SFATilingCheck() = default; + virtual ge::graphStatus Process(); +private: + void Init(); + void LogErrorDtypeSupport(const std::vector &expectDtypeList, + const ge::DataType &actualDtype, const std::string &name) const; + ge::graphStatus CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc, + const std::string &name) const; + template void LogErrorNumberSupport(const std::vector &expectNumberList, + const T &actualValue, const std::string &name, const std::string subName) const; + template void LogErrorDimNumSupport(const std::vector &expectNumberList, + const T &actualValue, const std::string &name) const; + ge::graphStatus CheckDimNumSupport(const gert::StorageShape *shape, + const std::vector &expectDimNumList, const std::string &name) const; + ge::graphStatus CheckDimNumInLayoutSupport(const SFALayout &layout, + const gert::StorageShape *shape, const std::string &name) const; + void LogErrorLayoutSupport(const std::vector &expectLayoutList, + const SFALayout &actualLayout, const std::string &name) const; + ge::graphStatus GetExpectedShape(gert::Shape &shapeExpected, + const SFATilingShapeCompareParam ¶m, const SFALayout &layout) const; + ge::graphStatus CompareShape(SFATilingShapeCompareParam ¶m, + const gert::Shape &shape, const SFALayout &layout, const std::string &name) const; + ge::graphStatus CheckLayoutSupport(const SFALayout &actualLayout, const std::string &name) const; + ge::graphStatus CheckSingleParaQuery() const; + ge::graphStatus CheckSingleParaKey() const; + ge::graphStatus CheckSingleParaValue() const; + ge::graphStatus CheckSingleParaQueryRope() const; + ge::graphStatus CheckSingleParaKeyRope() const; + ge::graphStatus CheckSingleParaAttenOut() const; + ge::graphStatus CheckSingleParaNumHeads() const; + ge::graphStatus CheckSingleParaKvHeadNums() const; + ge::graphStatus CheckSingleParaLayout() const; + ge::graphStatus CheckSingleParaSparseMode() const; + ge::graphStatus CheckSingleParaSparseBlockSize() const; + ge::graphStatus CheckSingleParaSparseIndices() const; + ge::graphStatus CheckSinglePara() const; + ge::graphStatus CheckMultiParaConsistency() const; + ge::graphStatus CheckRopeExistence(); + ge::graphStatus CheckExists(const void *pointer, const std::string &name) const; + ge::graphStatus CheckNotExists(const void *pointer, const std::string &name) const; + ge::graphStatus CheckExistsByMap(const std::map ¶mMap) const; + ge::graphStatus CheckNotExistsByMap(const std::map ¶mMap) const; + ge::graphStatus CheckExistenceByMap(std::map &existMap, + std::map ¬ExistMap) const; + template ge::graphStatus CheckAttrValueByMap( + std::map> &attrMap) const; + ge::graphStatus CheckParaExistenceMlaNoquant() const; + ge::graphStatus CheckParaExistenceGqaNoquant() const; + ge::graphStatus CheckParaExistenceMla() const; + ge::graphStatus CheckParaExistence(); + ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + const SFALayout &layout, const std::string &name) const; + void SetSFAShapeCompare(); + ge::graphStatus CheckQRope(); + ge::graphStatus CheckQRopeShape(); + ge::graphStatus CheckVAndKRopeShapeForBatchContinuous(); + ge::graphStatus CheckVAndKRopeShapeForPageAttention(); + ge::graphStatus CheckVAndKRopeShape(); + ge::graphStatus CheckVAndKRope(); + ge::graphStatus CheckTopK(); + ge::graphStatus CheckTopkShape(); + ge::graphStatus CheckBlockTable() const; + ge::graphStatus CheckDTypeConsistency(const ge::DataType &actualDtype, + const ge::DataType &expectDtype, const std::string &name) const; + + ge::graphStatus CheckAttenOut(); + ge::graphStatus CheckAttenOutShape(); + ge::graphStatus CheckSoftmaxMax(); + ge::graphStatus CheckSoftmaxMaxShape(); + ge::graphStatus CheckSoftmaxSum(); + ge::graphStatus CheckSoftmaxSumShape(); + ge::graphStatus CheckActualSeqLensQ(); + ge::graphStatus CheckActualSeqLensQShape(); + ge::graphStatus CheckActualSeqLensQDType(); + ge::graphStatus CheckActualSeqLens(); + ge::graphStatus CheckActualSeqLensDType(); + ge::graphStatus CheckActualSeqLensShape(); + ge::graphStatus CheckMultiParaConsistency(); + + ge::graphStatus CheckFeatureMlaNoQuantShape() const; + ge::graphStatus CheckFeatureMlaNoQuantLayout() const; + ge::graphStatus CheckFeatureMlaNoQuantDtype() const; + ge::graphStatus CheckFeatureMlaNoquantPa() const; + ge::graphStatus CheckFeatureMlaNoquant() const; + ge::graphStatus CheckFeatureMla() const; + ge::graphStatus CheckFeature() const; + + ge::graphStatus CheckSingleParaPreTokens() const; + ge::graphStatus CheckSingleParaNextTokens() const; + +private: + const char *opName_; + fe::PlatFormInfos *platformInfo_; + SFAParaInfo opParamInfo_; + const SFATilingInfo &sfaInfo_; + + uint32_t bSize_ = 0; + uint32_t n1Size_ = 0; + uint32_t n2Size_ = 0; + uint32_t gSize_ = 0; + uint32_t s1Size_ = 0; + int64_t s2Size_ = 0; + uint32_t qkHeadDim_ = 0; + uint32_t vHeadDim_ = 0; + uint32_t ropeHeadDim_ = 0; + uint32_t qTSize_ = 0; // 仅TND时生效 + uint32_t kvTSize_ = 0; // 仅TND时生效 + KvStorageMode kvStorageMode_ = KvStorageMode::BATCH_CONTINUOUS; + uint32_t sparseBlockCount_ = 0; + int64_t sparseBlockSize_ = 0; + + SFALayout qLayout_ = SFALayout::BSND; + SFALayout topkLayout_ = SFALayout::BSND; + SFALayout outLayout_ = SFALayout::BSND; + SFALayout kvLayout_ = SFALayout::BSND; + SFALayout softmaxMaxLayout_ = SFALayout::BNSG; + SFALayout softmaxSumLayout_ = SFALayout::BNSG; + + uint32_t maxBlockNumPerBatch_ = 0; + int64_t blockSize_ = 0; + + uint32_t aicNum_ = 0; + uint32_t aivNum_ = 0; + NpuArch npuArch_ = NpuArch::DAV_2201; + bool isA5_ = false; + uint64_t l2CacheSize_ = 0; + + ge::DataType inputQType_ = ge::DT_FLOAT16; + ge::DataType inputKvType_ = ge::DT_FLOAT16; + ge::DataType outputType_ = ge::DT_FLOAT16; + ge::DataType inputQRopeType_ = ge::DT_FLOAT16; + ge::DataType inputKRopeType_ = ge::DT_FLOAT16; + + gert::Shape queryShapeCmp_{}; + gert::Shape keyShapeCmp_{}; + gert::Shape valueShapeCmp_{}; + gert::Shape topkShapeCmp_{}; + gert::Shape queryRopeShapeCmp_{}; + gert::Shape keyRopeShapeCmp_{}; + gert::Shape attenOutShapeCmp_{}; + gert::Shape softmaxMaxShapeCmp_{}; + gert::Shape softmaxSumShapeCmp_{}; +}; + +class SFAInfoParser { +public: + explicit SFAInfoParser(const gert::TilingContext *context) : context_(context) {} + ~SFAInfoParser() = default; + + ge::graphStatus CheckRequiredInOutExistence() const; + ge::graphStatus CheckRequiredAttrExistence() const; + ge::graphStatus CheckRequiredParaExistence() const; + + ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + SFALayout &layout, const std::string &name) const; + ge::graphStatus GetActualSeqLenQSize(uint32_t &size); + ge::graphStatus GetOpName(); + ge::graphStatus GetNpuInfo(); + void GetOptionalInputParaInfo(); + void GetInputParaInfo(); + void GetOutputParaInfo(); + ge::graphStatus GetAttrParaInfo(); + ge::graphStatus GetKvCache(); + ge::graphStatus GetOpParaInfo(); + + ge::graphStatus GetInOutDataType(); + ge::graphStatus GetBatchSize(); + ge::graphStatus GetQTSize(); + ge::graphStatus GetKVTSize(); + ge::graphStatus GetQkHeadDim(); + ge::graphStatus GetS1Size(); + ge::graphStatus GetKvStorageMode(); + ge::graphStatus GetKvLayout(); + void SetSFAShape(); + ge::graphStatus GetS2SizeForBatchContinuous(); + ge::graphStatus GetMaxBlockNumPerBatch(); + ge::graphStatus GetBlockSize(); + ge::graphStatus GetS2SizeForPageAttention(); + ge::graphStatus GetS2Size(); + ge::graphStatus GetValueHeadDim(); + ge::graphStatus GetRopeHeadDim(); + ge::graphStatus GetQueryAndOutLayout(); + ge::graphStatus GetTopkLayout(); + ge::graphStatus GetSoftmaxMaxAndSumLayout(); + ge::graphStatus GetN1Size(); + ge::graphStatus GetN2Size(); + ge::graphStatus GetGSize(); + ge::graphStatus GetSparseBlockCount(); + ge::graphStatus GetActualseqInfo(); + void GenerateInfo(SFATilingInfo &sfaInfo); + ge::graphStatus Parse(SFATilingInfo &sfaInfo); + +public: + bool HasAxis(const SFAAxis &axis, const SFALayout &layout, const gert::Shape &shape) const; + size_t GetAxisIdx(const SFAAxis &axis, const SFALayout &layout) const; + uint32_t GetAxisNum(const gert::Shape &shape, const SFAAxis &axis,const SFALayout &layout) const; + + const gert::TilingContext *context_ = nullptr; + + const char *opName_; + fe::PlatFormInfos *platformInfo_; + SFAParaInfo opParamInfo_; + static constexpr int64_t invalidDimValue_ = std::numeric_limits::min(); + + uint32_t bSize_ = 0; + uint32_t n1Size_ = 0; + uint32_t n2Size_ = 0; + uint32_t gSize_ = 0; + uint32_t s1Size_ = 0; + int64_t s2Size_ = 0; + uint32_t qkHeadDim_ = 0; + uint32_t vHeadDim_ = 0; + uint32_t ropeHeadDim_ = 0; + uint32_t qTSize_ = 0; // 仅TND时生效 + uint32_t kvTSize_ = 0; // 仅TND时生效 + KvStorageMode kvStorageMode_ = KvStorageMode::BATCH_CONTINUOUS; + uint32_t sparseBlockCount_ = 0; + + SFALayout qLayout_ = SFALayout::BSND; + SFALayout topkLayout_ = SFALayout::BSND; + SFALayout outLayout_ = SFALayout::BSND; + SFALayout kvLayout_ = SFALayout::BSND; + SFALayout softmaxMaxLayout_ = SFALayout::BNSG; + SFALayout softmaxSumLayout_ = SFALayout::BNSG; + uint32_t maxBlockNumPerBatch_ = 0; + uint32_t blockSize_ = 0; + + NpuArch npuArch_ = NpuArch::DAV_2201; + bool isA5_ = false; + + ge::DataType inputQType_ = ge::DT_FLOAT16; + ge::DataType inputKvType_ = ge::DT_FLOAT16; + ge::DataType outputType_ = ge::DT_FLOAT16; + ge::DataType inputQRopeType_ = ge::DT_FLOAT16; + ge::DataType inputKRopeType_ = ge::DT_FLOAT16; + + uint64_t l2CacheSize_ = 0; + + bool isSameSeqAllKVTensor_ = true; + bool isSameActualseq_ = true; + uint32_t maxActualseq_ = 0; + + uint32_t actualLenDimsQ_ = 0; + uint32_t actualLenDimsKV_ = 0; + + gert::Shape queryShape_{}; + gert::Shape keyShape_{}; + gert::Shape valueShape_{}; + gert::Shape sparseIndicesShape_{}; + gert::Shape queryRopeShape_{}; + gert::Shape keyRopeShape_{}; +}; +} // namespace optiling +#endif // SPARSE_FLASH_ATTENTION_TILING_H diff --git a/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch22/sparse_flash_attention_kernel_mla.h b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch22/sparse_flash_attention_kernel_mla.h new file mode 100644 index 000000000..75d34f875 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch22/sparse_flash_attention_kernel_mla.h @@ -0,0 +1,1019 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_flash_attention_kernel_mla.h + * \brief + */ + +#ifndef SPARSE_FLASH_ATTENTION_KERNEL_MLA_H +#define SPARSE_FLASH_ATTENTION_KERNEL_MLA_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../sparse_flash_attention_common.h" +#include "sparse_flash_attention_service_cube_mla.h" +#include "sparse_flash_attention_service_vector_mla.h" + +using namespace matmul; +using AscendC::CacheMode; +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +// 由于S2循环前，RunInfo还没有赋值，使用Bngs1Param临时存放B、N、S1轴相关的信息；同时减少重复计算 +struct TempLoopInfo { + uint32_t bn2IdxInCurCore = 0; + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint64_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小 + uint32_t s2LoopTimes = 0U; // S2方向循环的总次数，无论TND还是BXXD都是等于实际次数，不用减1 + uint64_t curActualSeqLen = 0ULL; + uint64_t curActualSeqLenOri = 0ULL; + bool curActSeqLenIsZero = false; + int32_t nextTokensPerBatch = 0; + + uint64_t actS1Size = 1ULL; // TND场景下当前Batch循环处理的S1轴的大小 + uint32_t tndCoreStartKVSplitPos; + bool tndIsS2SplitCore; + + uint32_t gS1Idx = 0U; + uint64_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小 +}; + +template class SparseFlashAttentionMla { +public: + // 中间计算数据类型为float，高精度模式 + using T = float; + using Q_T = typename SFAT::queryType; + using KV_T = typename SFAT::kvType; + using OUT_T = typename SFAT::outputType; + using Q_ROPE_T = Q_T; + using K_ROPE_T = KV_T; + using UPDATE_T = T; + using MM1_OUT_T = T; + using MM2_OUT_T = T; + + __aicore__ inline SparseFlashAttentionMla(){}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengths, __gm__ uint8_t *blockTable, + __gm__ uint8_t *queryRope, __gm__ uint8_t *keyRope, + __gm__ uint8_t *attentionOut, __gm__ uint8_t *softmaxMax, __gm__ uint8_t *softmaxSum, + __gm__ uint8_t *workspace, const SparseFlashAttentionTilingDataMla *__restrict tiling, + __gm__ uint8_t *gmTiling, TPipe *tPipe); + + __aicore__ inline void Process(); + +private: + static constexpr bool PAGE_ATTENTION = SFAT::pageAttention; + static constexpr int TEMPLATE_MODE = SFAT::templateMode; + static constexpr bool FLASH_DECODE = SFAT::flashDecode; + static constexpr SFA_LAYOUT LAYOUT_T = SFAT::layout; + static constexpr SFA_LAYOUT KV_LAYOUT_T = SFAT::kvLayout; + + static constexpr uint32_t PRELOAD_NUM = 2; + static constexpr uint32_t N_BUFFER_M_BASIC_SIZE = 256; + static constexpr uint32_t SFA_PRELOAD_TASK_CACHE_SIZE = 3; + + static constexpr uint32_t SYNC_V0_C1_FLAG = 6; + static constexpr uint32_t SYNC_C1_V1_FLAG = 7; + static constexpr uint32_t SYNC_V1_C2_FLAG = 8; + static constexpr uint32_t SYNC_C2_V2_FLAG = 9; + static constexpr uint32_t SYNC_C2_V1_FLAG = 4; + static constexpr uint32_t SYNC_V1_NUPDATE_C2_FLAG = 5; + + static constexpr uint64_t SYNC_MM2RES_BUF1_FLAG = 10; + static constexpr uint64_t SYNC_MM2RES_BUF2_FLAG = 11; + static constexpr uint64_t SYNC_FDOUTPUT_BUF_FLAG = 12; + + static constexpr uint32_t BLOCK_ELEMENT_NUM = SFAVectorService::BYTE_BLOCK / sizeof(T); + + static constexpr uint64_t kvHeadNum = 1ULL; + static constexpr uint64_t headDim = 512ULL; + static constexpr uint64_t headDimAlign = 512ULL; + static constexpr uint64_t headDimRope = 64ULL; + static constexpr uint32_t msdIterNum = 2U; + + static constexpr uint32_t dbWorkspaceRatio = PRELOAD_NUM; + + const SparseFlashAttentionTilingDataMla *__restrict tilingData = nullptr; + + TPipe *pipe = nullptr; + + uint64_t mSizeVStart = 0ULL; + int64_t threshold = 0; + uint64_t topKBaseOffset = 0ULL; + uint64_t s2BatchBaseOffset = 0; + uint64_t tensorACoreOffset = 0ULL; + uint64_t tensorBCoreOffset = 0ULL; + uint64_t tensorARopeCoreOffset = 0ULL; + uint64_t tensorBRopeCoreOffset = 0ULL; + uint64_t tensorBOffset = 0ULL; + uint64_t attenOutOffset = 0ULL; + + uint32_t tmpBlockIdx = 0U; + uint32_t aiCoreIdx = 0U; + uint32_t usedCoreNum = 0U; + + __gm__ uint8_t *keyPtr = nullptr; + __gm__ uint8_t *valuePtr = nullptr; + + ConstInfo constInfo{}; + TempLoopInfo tempLoopInfo{}; + + SFAMatmulService matmulService; + SFAVectorService vectorService; + + GlobalTensor queryGm; + GlobalTensor keyGm; + GlobalTensor valueGm; + GlobalTensor qRopeGm; + GlobalTensor kRopeGm; + + GlobalTensor attentionOutGm; + GlobalTensor softmaxMaxGm; + GlobalTensor softmaxSumGm; + GlobalTensor blockTableGm; + GlobalTensor topKGm; + + GlobalTensor actualSeqLengthsQGm; + GlobalTensor actualSeqLengthsKVGm; + + // workspace + GlobalTensor mm1ResGm; + GlobalTensor vec1ResGm; + GlobalTensor mm2ResGm; + GlobalTensor kvMergeGm_; + GlobalTensor kvValidSizeGm_; + + GlobalTensor mm2ResInt32Gm; + GlobalTensor vec2ResGm; + + GlobalTensor accumOutGm; + GlobalTensor lseSumFdGm; + GlobalTensor lseMaxFdGm; + + GlobalTensor lseSumFaGm; + GlobalTensor lseMaxFaGm; + + // ================================Init functions=================================== + __aicore__ inline void InitTilingData(); + __aicore__ inline void InitCalcParamsEach(); + __aicore__ inline void InitBuffers(); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths); + __aicore__ inline void InitOutputSingleCore(); + // ================================Process functions================================ + __aicore__ inline void ProcessBalance(); + __aicore__ inline void PreloadPipeline(uint32_t loop, uint64_t s2Start, uint64_t s2LoopIdx, + RunInfo extraInfo[SFA_PRELOAD_TASK_CACHE_SIZE], uint32_t &curTopKIdx, uint64_t &curOffsetInSparseBlock); + // ================================Offset Calc===================================== + __aicore__ inline void GetActualSeqLen(uint32_t bIdx, uint32_t s1Idx = 0); + __aicore__ inline void GetSparseActualSeqLen(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx); + __aicore__ inline void CalcSinnerTopKBegin(RunInfo &info, uint32_t &curTopKIdx, uint64_t &curOffsetInSparseBlock); + __aicore__ inline void UpdateInnerLoopCond(); + __aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx); + __aicore__ inline void CalcParams(uint32_t loop, uint64_t s2Start, uint32_t s2LoopIdx, RunInfo &info); + __aicore__ inline void GetAxisStartIdx(uint32_t bN2EndPrev, uint32_t gS1EndPrev, uint32_t s2EndPrev); + __aicore__ inline uint64_t GetBalanceActualSeqLengths(GlobalTensor &actualSeqLengths, uint32_t bIdx); + __aicore__ inline uint32_t GetActualSeqLenKV(uint32_t bIdx); + __aicore__ inline void GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, uint32_t &n2Idx); + __aicore__ inline void UpdateInner(uint32_t &s2End, uint32_t &curS2End, uint32_t s1Idx, bool isEnd); + __aicore__ inline void GetPreNextTokensLeftUp(); + // ================================Mm1============================================== + __aicore__ inline void ComputeMm1(const RunInfo &info); + // ================================Mm2============================================== + __aicore__ inline void ComputeMm2(const RunInfo &info); + __aicore__ inline void Bmm2DataCopyOut(uint64_t attenOutOffset, LocalTensor &attenOutUb, uint32_t startRow, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline void InitAllZeroOutput(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx); +}; + +template __aicore__ inline void SparseFlashAttentionMla::InitTilingData() +{ + usedCoreNum = tilingData->singleCoreParams.usedCoreNum; + constInfo.splitKVNum = tilingData->splitKVParams.s2; + constInfo.mmResUbSize = tilingData->singleCoreTensorSize.mmResUbSize; + constInfo.bmm2ResUbSize = tilingData->singleCoreTensorSize.bmm2ResUbSize; + constInfo.vec1ResUbSize = constInfo.mmResUbSize * msdIterNum; + + constInfo.batchSize = tilingData->baseParams.batchSize; + constInfo.qHeadNum = constInfo.gSize = tilingData->baseParams.nNumOfQInOneGroup; + constInfo.kvSeqSize = tilingData->baseParams.seqSize; + constInfo.qSeqSize = tilingData->baseParams.qSeqSize; + constInfo.maxBlockNumPerBatch = tilingData->baseParams.maxBlockNumPerBatch; + constInfo.kvCacheBlockSize = tilingData->baseParams.blockSize; + constInfo.outputLayout = static_cast(tilingData->baseParams.outputLayout); + constInfo.mBaseSize = tilingData->innerSplitParams.mBaseSize; + constInfo.s2BaseSize = tilingData->innerSplitParams.s2BaseSize; + constInfo.kvHeadNum = kvHeadNum; + constInfo.headDim = headDim; + constInfo.headDimRope = headDimRope; + constInfo.sparseBlockSize = tilingData->baseParams.sparseBlockSize; + constInfo.sparseBlockCount = tilingData->baseParams.sparseBlockCount; + constInfo.sparseMode = tilingData->baseParams.sparseMode; + constInfo.preTokens = tilingData->baseParams.preTokens; + constInfo.nextTokens = tilingData->baseParams.nextTokens; + constInfo.attentionMode = tilingData->baseParams.attentionMode; + constInfo.returnSoftmaxLse = tilingData->baseParams.returnSoftmaxLse; + + constInfo.preLoadNum = PRELOAD_NUM; + constInfo.nBufferMBaseSize = N_BUFFER_M_BASIC_SIZE; + constInfo.syncV0C1 = SYNC_V0_C1_FLAG; + constInfo.syncC1V1 = SYNC_C1_V1_FLAG; + constInfo.syncV1C2 = SYNC_V1_C2_FLAG; + constInfo.syncC2V2 = SYNC_C2_V2_FLAG; + constInfo.syncC2V1 = SYNC_C2_V1_FLAG; + constInfo.syncV1NupdateC2 = SYNC_V1_NUPDATE_C2_FLAG; +} + +template __aicore__ inline void SparseFlashAttentionMla::InitBuffers() +{ + if ASCEND_IS_AIV { + vectorService.InitBuffers(pipe); + } else { + matmulService.InitBuffers(pipe); + } +} + +template +__aicore__ inline void +SparseFlashAttentionMla::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengths) +{ + constInfo.actualLenDimsQ = tilingData->baseParams.actualLenDimsQ; + constInfo.actualLenDimsKV = tilingData->baseParams.actualLenDimsKV; + if (constInfo.actualLenDimsKV != 0) { + actualSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengths, constInfo.actualLenDimsKV); + } + if (constInfo.actualLenDimsQ != 0) { + actualSeqLengthsQGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsQ, constInfo.actualLenDimsQ); + } +} + +template +__aicore__ inline void SparseFlashAttentionMla::InitAllZeroOutput(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx) +{ + if (constInfo.outputLayout == SFA_LAYOUT::TND) { + uint32_t tBase = bIdx == 0 ? 0 : actualSeqLengthsQGm.GetValue(bIdx - 1); + uint32_t s1Count = tempLoopInfo.actS1Size; + + uint64_t attenOutOffset = (tBase + s1Idx) * kvHeadNum * constInfo.gSize * headDim + // T轴、s1轴偏移 + n2Idx * constInfo.gSize * headDim; // N2轴偏移 + matmul::InitOutput(attentionOutGm[attenOutOffset], constInfo.gSize * headDim, 0); + if (constInfo.returnSoftmaxLse) { + uint64_t softmaxSumOffset = n2Idx * actualSeqLengthsQGm.GetValue(constInfo.batchSize - 1) * constInfo.gSize + + (tBase + s1Idx) * constInfo.gSize; + uint64_t softmaxMaxOffset = softmaxSumOffset; + matmul::InitOutput(softmaxSumGm[softmaxSumOffset], constInfo.gSize, 0); + matmul::InitOutput(softmaxMaxGm[softmaxMaxOffset], constInfo.gSize, 0); + } + } else if (constInfo.outputLayout == SFA_LAYOUT::BSND) { + uint64_t attenOutOffset = bIdx * constInfo.qSeqSize * kvHeadNum * constInfo.gSize * headDim + + s1Idx * kvHeadNum * constInfo.gSize * headDim + // B轴、S1轴偏移 + n2Idx * constInfo.gSize * headDim; // N2轴偏移 + matmul::InitOutput(attentionOutGm[attenOutOffset], constInfo.gSize * headDim, 0); + if (constInfo.returnSoftmaxLse) { + uint64_t softmaxSumOffset = bIdx * kvHeadNum * constInfo.qSeqSize * constInfo.gSize + + n2Idx * constInfo.qSeqSize * constInfo.gSize + + s1Idx * constInfo.gSize; + uint64_t softmaxMaxOffset = softmaxSumOffset; + matmul::InitOutput(softmaxSumGm[softmaxSumOffset], constInfo.gSize, 0); + matmul::InitOutput(softmaxMaxGm[softmaxMaxOffset], constInfo.gSize, 0); + } + } +} + +template +__aicore__ inline void SparseFlashAttentionMla::InitOutputSingleCore() +{ + uint32_t coreNum = GetBlockNum(); + if (coreNum != 0) { + uint64_t totalOutputSize = constInfo.batchSize * constInfo.qHeadNum * constInfo.qSeqSize * constInfo.headDim; + uint64_t singleCoreSize = (totalOutputSize + (2 * coreNum) - 1) / (2 * coreNum); // 2 means c:v = 1:2 + uint64_t tailSize = totalOutputSize - tmpBlockIdx * singleCoreSize; + uint64_t singleInitOutputSize = tailSize < singleCoreSize ? tailSize : singleCoreSize; + if (tmpBlockIdx * singleCoreSize < totalOutputSize && singleInitOutputSize > 0) { + matmul::InitOutput(attentionOutGm[tmpBlockIdx * singleCoreSize], singleInitOutputSize, 0); + } + if (constInfo.returnSoftmaxLse) { + uint64_t totalReturnSoftmaxSize = constInfo.batchSize * constInfo.kvHeadNum * constInfo.qSeqSize * constInfo.gSize; + uint64_t singleCoreReturnSoftmaxSize = (totalReturnSoftmaxSize + (2 * coreNum) - 1) / (2 * coreNum); + uint64_t tailReturnSoftmaxSize = totalReturnSoftmaxSize - tmpBlockIdx * singleCoreReturnSoftmaxSize; + uint64_t singleInitReturnSoftmaxSize = tailReturnSoftmaxSize < singleCoreReturnSoftmaxSize ? + tailReturnSoftmaxSize : singleCoreReturnSoftmaxSize; + if (tmpBlockIdx * singleCoreReturnSoftmaxSize < totalReturnSoftmaxSize && singleInitReturnSoftmaxSize > 0) { + matmul::InitOutput(softmaxSumGm[tmpBlockIdx * singleCoreReturnSoftmaxSize], singleInitReturnSoftmaxSize, 0); + matmul::InitOutput(softmaxMaxGm[tmpBlockIdx * singleCoreReturnSoftmaxSize], singleInitReturnSoftmaxSize, 0); + } + } + SyncAll(); + } +} + +template +__aicore__ inline void SparseFlashAttentionMla::GetActualSeqLen(uint32_t bIdx, uint32_t s1Idx) +{ + tempLoopInfo.curActualSeqLenOri = GetActualSeqLenKV(bIdx); + tempLoopInfo.actS1Size = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bIdx); +} + +template +__aicore__ inline void SparseFlashAttentionMla::GetSparseActualSeqLen(uint32_t bIdx, uint32_t s1Idx, + uint32_t n2Idx) +{ + if (tempLoopInfo.nextTokensPerBatch < 0 && s1Idx < (-tempLoopInfo.nextTokensPerBatch)) { //存在行无效 + tempLoopInfo.curActualSeqLen = 0; + return; + } + int64_t threshold = tempLoopInfo.curActualSeqLenOri; + if (constInfo.sparseMode == 3) { + threshold = static_cast(tempLoopInfo.nextTokensPerBatch) + s1Idx + 1; + } + + tempLoopInfo.curActualSeqLen = (constInfo.sparseBlockCount * constInfo.sparseBlockSize > threshold) ? + threshold : + constInfo.sparseBlockCount * constInfo.sparseBlockSize; +} + +template +__aicore__ inline uint32_t SparseFlashAttentionMla::GetActualSeqLenKV(uint32_t bIdx) +{ + if constexpr (KV_LAYOUT_T == SFA_LAYOUT::TND) { + if (bIdx > 0) { + return actualSeqLengthsKVGm.GetValue(bIdx) - actualSeqLengthsKVGm.GetValue(bIdx - 1); + } else if (bIdx == 0) { + return actualSeqLengthsKVGm.GetValue(0); + } else { + return 0; + } + } else { + if (constInfo.actualLenDimsKV == 0) { + return constInfo.kvSeqSize; + } else if (constInfo.actualLenDimsKV == 1) { + return actualSeqLengthsKVGm.GetValue(0); + } else { + return actualSeqLengthsKVGm.GetValue(bIdx); + } + } +} + +template +__aicore__ inline void SparseFlashAttentionMla::DealActSeqLenIsZero(uint32_t bIdx, uint32_t s1Idx, uint32_t n2Idx) +{ + if ASCEND_IS_AIV { + InitAllZeroOutput(bIdx, s1Idx, n2Idx); + } +} + +template +__aicore__ inline void SparseFlashAttentionMla::GetPreNextTokensLeftUp() +{ + if (constInfo.sparseMode == 3) { + tempLoopInfo.nextTokensPerBatch = + static_cast(tempLoopInfo.curActualSeqLenOri) - static_cast(tempLoopInfo.actS1Size); + } +} + +template __aicore__ inline void SparseFlashAttentionMla::UpdateInnerLoopCond() +{ + if ((tempLoopInfo.curActualSeqLen == 0) || (tempLoopInfo.actS1Size == 0)) { + tempLoopInfo.curActSeqLenIsZero = true; + return; + } + tempLoopInfo.curActSeqLenIsZero = false; + tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize; + tempLoopInfo.mBasicSizeTail = + (tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail; + tempLoopInfo.s2LoopTimes = 0; +} + +template +__aicore__ inline void SparseFlashAttentionMla::UpdateInner(uint32_t &s2End, uint32_t &curS2End, + uint32_t s1Idx, bool isEnd) +{ + uint32_t s1BaseSize = 1; + int64_t s1Offset = s1BaseSize * s1Idx; + int64_t s2LastToken = Min(s1Offset + tempLoopInfo.nextTokensPerBatch + s1BaseSize,tempLoopInfo.curActualSeqLenOri); + s2LastToken = Min(constInfo.sparseBlockSize * constInfo.sparseBlockCount, s2LastToken); + curS2End = (s2LastToken + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + tempLoopInfo.s2LoopTimes = isEnd ? constInfo.s2End + 1 : curS2End; +} + +template +__aicore__ inline void SparseFlashAttentionMla::Init(__gm__ uint8_t *query, + __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengths, __gm__ uint8_t *blockTable, + __gm__ uint8_t *queryRope, __gm__ uint8_t *keyRope, + __gm__ uint8_t *attentionOut, __gm__ uint8_t *softmaxMax, __gm__ uint8_t *softmaxSum, + __gm__ uint8_t *workspace, const SparseFlashAttentionTilingDataMla *__restrict tiling, + __gm__ uint8_t *gmTiling, TPipe *tPipe) +{ + if ASCEND_IS_AIV { + tmpBlockIdx = GetBlockIdx(); // vec:0-47 + aiCoreIdx = tmpBlockIdx / 2; + } else { + tmpBlockIdx = GetBlockIdx(); // cube:0-23 + aiCoreIdx = tmpBlockIdx; + } + + // init tiling data + tilingData = tiling; + + InitTilingData(); + InitActualSeqLen(actualSeqLengthsQ, actualSeqLengths); + + // 初始化计算参数 + InitCalcParamsEach(); + pipe = tPipe; + keyPtr = key; + valuePtr = value; + + // init global buffer + queryGm.SetGlobalBuffer((__gm__ Q_T *)query); + keyGm.SetGlobalBuffer((__gm__ KV_T *)keyPtr); + valueGm.SetGlobalBuffer((__gm__ KV_T *)valuePtr); + qRopeGm.SetGlobalBuffer((__gm__ Q_ROPE_T *)queryRope); + kRopeGm.SetGlobalBuffer((__gm__ K_ROPE_T *)keyRope); + + attentionOutGm.SetGlobalBuffer((__gm__ OUT_T *)attentionOut); + softmaxMaxGm.SetGlobalBuffer((__gm__ T *)softmaxMax); + softmaxSumGm.SetGlobalBuffer((__gm__ T *)softmaxSum); + if ASCEND_IS_AIV { + if (constInfo.needInit && LAYOUT_T != SFA_LAYOUT::TND) { + InitOutputSingleCore(); + } + } + + if constexpr (PAGE_ATTENTION) { + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + } + topKGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices); + + // workspace 内存排布 + // |Q--|mm1ResGm(存S)|vec1ResGm(存A1,A2)|mm2ResGm(存O)|vec2ResGm + // |Core0_Q1-Core0_Q2-Core1_Q1-Core1_Q2....Core32_Q1-Core32_Q2|Core0_mmRes + uint64_t offset = 0; + mm1ResGm.SetGlobalBuffer( + (__gm__ MM1_OUT_T *)(workspace + offset + + aiCoreIdx * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(MM1_OUT_T))); + offset += GetBlockNum() * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(MM1_OUT_T); + + vec1ResGm.SetGlobalBuffer( + (__gm__ KV_T *)(workspace + offset + aiCoreIdx * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(KV_T))); + offset += GetBlockNum() * dbWorkspaceRatio * constInfo.mmResUbSize * sizeof(KV_T); + + mm2ResGm.SetGlobalBuffer( + (__gm__ MM2_OUT_T *)(workspace + offset + + aiCoreIdx * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(MM2_OUT_T))); + offset += GetBlockNum() * dbWorkspaceRatio * constInfo.bmm2ResUbSize * sizeof(MM2_OUT_T); + mm2ResInt32Gm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(mm2ResGm.GetPhyAddr(0))); + + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + // s2 d+rope bufNum + kvMergeGm_.SetGlobalBuffer((__gm__ KV_T *)(workspace + offset + aiCoreIdx * 512 * 576 * 4 * sizeof(KV_T))); + offset += GetBlockNum() * 512 * 576 * 4 * sizeof(KV_T); + + kvValidSizeGm_.SetGlobalBuffer( + (__gm__ int32_t *)(workspace + offset + (aiCoreIdx * 2) * 128 * 4 * sizeof(int32_t))); + } + + if constexpr (FLASH_DECODE) { + accumOutGm.SetGlobalBuffer((__gm__ float *)(workspace + offset)); + offset = offset + tilingData->splitKVParams.accumOutSize * sizeof(float); + lseSumFdGm.SetGlobalBuffer((__gm__ float *)(workspace + offset)); + lseMaxFdGm.SetGlobalBuffer((__gm__ float *)(workspace + offset) + tilingData->splitKVParams.logSumExpSize / 2); + offset = offset + tilingData->splitKVParams.logSumExpSize * sizeof(float); + } + + if ASCEND_IS_AIV { + vectorService.InitParams(constInfo, tilingData); + vectorService.InitMm2ResInt32GmGlobalTensor(mm2ResInt32Gm); + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + vectorService.InitVec0GlobalTensor(kvValidSizeGm_, kvMergeGm_, kRopeGm, keyGm, blockTableGm); + } + vectorService.InitVec1GlobalTensor(mm1ResGm, vec1ResGm, actualSeqLengthsQGm, + actualSeqLengthsKVGm, lseMaxFdGm, lseSumFdGm, topKGm, softmaxMaxGm, softmaxSumGm); + vectorService.InitVec2GlobalTensor(accumOutGm, vec2ResGm, mm2ResGm, attentionOutGm); + } + + if ASCEND_IS_AIC { + matmulService.InitParams(constInfo); + matmulService.InitMm1GlobalTensor(queryGm, qRopeGm, keyGm, kRopeGm, mm1ResGm); + matmulService.InitMm2GlobalTensor(vec1ResGm, valueGm, mm2ResGm, attentionOutGm); + matmulService.InitPageAttentionInfo(kvMergeGm_, blockTableGm, topKGm, + constInfo.kvCacheBlockSize, constInfo.maxBlockNumPerBatch); + } + // 要在InitParams之后执行 + if (pipe != nullptr) { + InitBuffers(); + } +} + +template __aicore__ inline void SparseFlashAttentionMla::InitCalcParamsEach() +{ + //计算总的基本块 + uint32_t totalBaseNum = 0; + uint32_t s1GBaseSize = constInfo.gSize; + uint32_t actBatchS2 = 1; + uint32_t coreNum = GetBlockNum(); + uint32_t currCoreIdx = aiCoreIdx; + uint32_t actBatchS1 = 1; + for (uint32_t bIdx = 0; bIdx < constInfo.batchSize; bIdx++) { + uint32_t actBatchS1 = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bIdx); + if (actBatchS1 < constInfo.qSeqSize) { + constInfo.needInit = true; + } + totalBaseNum += actBatchS1*actBatchS2 ; + } + uint32_t avgBaseNum = 1; + if (totalBaseNum > coreNum) { + avgBaseNum = (totalBaseNum + coreNum - 1) / coreNum; + }else { + usedCoreNum = totalBaseNum; + } + if(aiCoreIdx>=usedCoreNum){ + return; + } + //计算当前核的基本块 + uint32_t accumBaseNum = 0; // 当前累积的基本块数 + uint32_t targetBaseNum = 0; + uint32_t lastValidBIdx = 0; + uint32_t lastValidactBatchS1=0; + bool setStart=false; + targetBaseNum = (currCoreIdx + 1) * avgBaseNum; // 计算当前的目标权重 + uint32_t targetStartBaseNum = targetBaseNum-avgBaseNum; + for (uint32_t bN2Idx = 0; bN2Idx < constInfo.batchSize * constInfo.kvHeadNum; bN2Idx++) { + uint32_t bIdx = bN2Idx / constInfo.kvHeadNum; + actBatchS1 = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bIdx); + for (uint32_t s1GIdx = 0; s1GIdx < actBatchS1; s1GIdx++) { + accumBaseNum += 1; + if(!setStart && accumBaseNum >= targetStartBaseNum){ + constInfo.bN2Start = bN2Idx; + constInfo.gS1Start = s1GIdx; + setStart=true; + } + if (accumBaseNum >= targetBaseNum) { + // 更新当前核的End分核信息 + constInfo.bN2End = bN2Idx; + constInfo.gS1End = s1GIdx; + constInfo.s2End = 0; + constInfo.coreStartKVSplitPos = 0; + if (aiCoreIdx != 0) { + GetAxisStartIdx(constInfo.bN2Start, constInfo.gS1Start, 0); + } + return; + } + } + if ((actBatchS1 > 0) && (actBatchS2 > 0)) { + lastValidBIdx = bIdx; + lastValidactBatchS1 = actBatchS1; + } + } + if (!setStart){ + constInfo.bN2Start = lastValidBIdx; + constInfo.gS1Start = lastValidactBatchS1-1; + } + if (accumBaseNum < targetBaseNum) { + // 更新最后一个核的End分核信息 + constInfo.bN2End = lastValidBIdx; + constInfo.gS1End = lastValidactBatchS1-1; + constInfo.s2End = 0; + constInfo.coreStartKVSplitPos = 0; + if (aiCoreIdx != 0) { + GetAxisStartIdx(constInfo.bN2Start, constInfo.gS1Start, 0); + } + return; + } +} + +template +__aicore__ inline void +SparseFlashAttentionMla::Bmm2DataCopyOut(uint64_t attenOutOffset, LocalTensor &attenOutUb, + uint32_t startRow, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = dealRowCount; + dataCopyParams.blockLen = actualColumnCount * sizeof(OUT_T); + dataCopyParams.srcStride = (columnCount - actualColumnCount) / (SFAVectorService::BYTE_BLOCK / sizeof(OUT_T)); + dataCopyParams.dstStride = 0; + DataCopyPad(attentionOutGm[attenOutOffset + (mSizeVStart + startRow) * actualColumnCount], attenOutUb, + dataCopyParams); +} + + +template +__aicore__ inline void SparseFlashAttentionMla::CalcParams(uint32_t loop, uint64_t s2Start, + uint32_t s2LoopIdx, RunInfo &info) +{ + info.loop = loop; + info.bIdx = tempLoopInfo.bIdx; + info.gS1Idx = tempLoopInfo.gS1Idx; + info.s2Idx = s2LoopIdx; + info.curSInnerLoopTimes = tempLoopInfo.s2LoopTimes; + + info.tndIsS2SplitCore = tempLoopInfo.tndIsS2SplitCore; + info.tndCoreStartKVSplitPos = tempLoopInfo.tndCoreStartKVSplitPos; + info.isBmm2Output = false; + + info.actS1Size = tempLoopInfo.actS1Size; + + + info.actMBaseSize = constInfo.mBaseSize; + uint32_t remainedGS1Size = tempLoopInfo.actS1Size * constInfo.gSize - tempLoopInfo.gS1Idx; + if (remainedGS1Size <= constInfo.mBaseSize && remainedGS1Size > 0) { + info.actMBaseSize = tempLoopInfo.mBasicSizeTail; + } + + info.isValid = s2LoopIdx < tempLoopInfo.s2LoopTimes; + + if ASCEND_IS_AIV { + info.mSize = info.actMBaseSize; + info.mSizeV = (info.mSize <= 16) ? info.mSize : (((info.mSize + 15) / 16 + 1) / 2 * 16); + info.mSizeVStart = 0; + if (tmpBlockIdx % 2 == 1) { + info.mSizeVStart = info.mSizeV; + info.mSizeV = info.mSize - info.mSizeV; + } + } + + info.isChangeBatch = false; + + info.isFirstSInnerLoop = s2LoopIdx == s2Start; + if (info.isFirstSInnerLoop) { + tempLoopInfo.bn2IdxInCurCore++; + } + info.isLastS2Loop = s2LoopIdx == tempLoopInfo.s2LoopTimes - 1; + info.bn2IdxInCurCore = tempLoopInfo.bn2IdxInCurCore - 1; + uint64_t actualSeqQPrefixSum; + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + actualSeqQPrefixSum = (info.bIdx <= 0) ? 0 : actualSeqLengthsQGm.GetValue(info.bIdx - 1); + } else { + actualSeqQPrefixSum = (info.bIdx <= 0) ? 0 : info.bIdx * constInfo.qSeqSize; + } + info.tndBIdxOffsetForQ = actualSeqQPrefixSum * constInfo.qHeadNum * headDim; + + uint64_t actualSeqKVPrefixSum; + if constexpr (KV_LAYOUT_T == SFA_LAYOUT::TND) { + actualSeqKVPrefixSum = (info.bIdx <= 0) ? 0 : actualSeqLengthsKVGm.GetValue(info.bIdx - 1); + } else { + actualSeqKVPrefixSum = (info.bIdx <= 0) ? 0 : info.bIdx * constInfo.kvSeqSize; + } + info.tndBIdxOffsetForKV = actualSeqKVPrefixSum * constInfo.kvHeadNum * headDim; + + if (info.isFirstSInnerLoop) { + uint64_t tndBIdxRopeOffsetForQ = actualSeqQPrefixSum * constInfo.qHeadNum * headDimRope; + tensorACoreOffset = info.tndBIdxOffsetForQ + info.gS1Idx * headDim; + tensorARopeCoreOffset = tndBIdxRopeOffsetForQ + info.gS1Idx * headDimRope; + + uint64_t tndBIdxRopeOffsetForK = actualSeqKVPrefixSum * constInfo.kvHeadNum * headDimRope; + tensorBCoreOffset = info.tndBIdxOffsetForKV + info.n2Idx * headDim; + tensorBRopeCoreOffset = tndBIdxRopeOffsetForK + info.n2Idx * headDimRope; + if (constInfo.sparseMode == 3) { + threshold = static_cast(tempLoopInfo.nextTokensPerBatch) + info.gS1Idx / constInfo.gSize + 1; + } else { + threshold = tempLoopInfo.curActualSeqLenOri; + } + if constexpr(LAYOUT_T == SFA_LAYOUT::BSND) { // B,S1,N2 K + topKBaseOffset = info.bIdx * constInfo.qSeqSize * constInfo.kvHeadNum * constInfo.sparseBlockCount + + info.gS1Idx / constInfo.gSize * constInfo.kvHeadNum * constInfo.sparseBlockCount + + info.n2Idx * constInfo.sparseBlockCount; + } else if (LAYOUT_T == SFA_LAYOUT::TND) { // T N2 K + topKBaseOffset = info.tndBIdxOffsetForQ / constInfo.gSize / constInfo.headDim * constInfo.kvHeadNum * + constInfo.sparseBlockCount + info.n2Idx * constInfo.sparseBlockCount + + info.gS1Idx / constInfo.gSize * constInfo.kvHeadNum * constInfo.sparseBlockCount; + } else { // B N2 S1 K + topKBaseOffset = info.bIdx * constInfo.kvHeadNum * constInfo.qSeqSize * constInfo.sparseBlockCount + + info.n2Idx * constInfo.qSeqSize * constInfo.sparseBlockCount + + info.gS1Idx / constInfo.gSize * constInfo.sparseBlockCount; + } + } + info.topKBaseOffset = topKBaseOffset; + info.threshold = threshold; + info.tensorAOffset = tensorACoreOffset; + info.tensorARopeOffset = tensorARopeCoreOffset; + info.tensorBOffset = tensorBCoreOffset; + info.tensorBRopeOffset = tensorBRopeCoreOffset; + info.attenOutOffset = tensorACoreOffset; + + uint64_t sInnerOffsetDataSize = info.s2Idx * constInfo.s2BaseSize; + info.s2BatchOffset = s2BatchBaseOffset + sInnerOffsetDataSize; + + info.curActualSeqLenOri = tempLoopInfo.curActualSeqLenOri; + //计算实际基本块size + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + if (tempLoopInfo.curActualSeqLen > sInnerOffsetDataSize) { + info.actualSingleProcessSInnerSize = tempLoopInfo.curActualSeqLen - sInnerOffsetDataSize; + info.actualSingleProcessSInnerSize = info.actualSingleProcessSInnerSize > constInfo.s2BaseSize ? + constInfo.s2BaseSize : info.actualSingleProcessSInnerSize; + } else { + info.actualSingleProcessSInnerSize = 0; + } + info.actualSingleProcessSInnerSizeAlign = + SFAAlign((uint32_t)info.actualSingleProcessSInnerSize, (uint32_t)SFAVectorService::BYTE_BLOCK); + } +} + +template +__aicore__ inline void SparseFlashAttentionMla::ComputeMm1(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + matmulService.ComputeMm1(info, mSplitInfo); + CrossCoreSetFlag(constInfo.syncC1V1); + } +} + +template +__aicore__ inline void SparseFlashAttentionMla::ComputeMm2(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + CrossCoreWaitFlag(constInfo.syncV1C2); + matmulService.ComputeMm2(info, mSplitInfo); + CrossCoreSetFlag(constInfo.syncC2V2); + CrossCoreSetFlag(constInfo.syncC2V1); + } +} + +template __aicore__ inline void SparseFlashAttentionMla::Process() +{ + if (aiCoreIdx < usedCoreNum) { + if ASCEND_IS_AIV { + vectorService.AllocEventID(); + vectorService.InitSoftmaxDefaultBuffer(); + } else { + matmulService.AllocEventID(); + } + ProcessBalance(); + + if ASCEND_IS_AIV { + vectorService.FreeEventID(); + } else { + matmulService.FreeEventID(); + } + } +} + +template +__aicore__ inline void SparseFlashAttentionMla::GetBN2Idx(uint32_t bN2Idx, uint32_t &bIdx, + uint32_t &n2Idx) +{ + bIdx = bN2Idx / kvHeadNum; + n2Idx = bN2Idx % kvHeadNum; +} + +template __aicore__ inline void SparseFlashAttentionMla::ProcessBalance() +{ + RunInfo extraInfo[SFA_PRELOAD_TASK_CACHE_SIZE]; + uint32_t gloop = 0; + int gS1LoopEnd; + bool globalLoopStart = true; + if ASCEND_IS_AIC { + CrossCoreSetFlag(constInfo.syncC2V1); + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + CrossCoreSetFlag(3); + CrossCoreSetFlag(3); + CrossCoreSetFlag(3); + CrossCoreSetFlag(3); + } + } + for (uint32_t bN2LoopIdx = constInfo.bN2Start; bN2LoopIdx <= constInfo.bN2End; bN2LoopIdx++) { + GetBN2Idx(bN2LoopIdx, tempLoopInfo.bIdx, tempLoopInfo.n2Idx); + GetActualSeqLen(tempLoopInfo.bIdx); // 获取actualSeqLength及ActualSeqLengthKV + GetPreNextTokensLeftUp(); + if (tempLoopInfo.actS1Size == 0) { + continue; + } + int gS1SplitNum = (tempLoopInfo.actS1Size * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize; + gS1LoopEnd = (bN2LoopIdx == constInfo.bN2End) ? constInfo.gS1End : gS1SplitNum - 1; + for (uint32_t gS1LoopIdx = constInfo.gS1Start; gS1LoopIdx <= gS1LoopEnd; gS1LoopIdx++) { + tempLoopInfo.gS1Idx = gS1LoopIdx * constInfo.mBaseSize; + GetSparseActualSeqLen(tempLoopInfo.bIdx, gS1LoopIdx, tempLoopInfo.n2Idx); // TopK值sparse完后的ActualSeqLengthKV + UpdateInnerLoopCond(); + + if (tempLoopInfo.curActSeqLenIsZero) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, gS1LoopIdx, tempLoopInfo.n2Idx); + } + int s2SplitNum = + (tempLoopInfo.curActualSeqLen + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; // S2切分份数 + bool isEnd = (bN2LoopIdx == constInfo.bN2End) && (gS1LoopIdx == constInfo.gS1End); + tempLoopInfo.s2LoopTimes = s2SplitNum; + // 分核修改后需要打开 + // 当前s2是否被切，决定了输出是否要写到attenOut上 + tempLoopInfo.tndIsS2SplitCore = + ((constInfo.s2Start == 0) && (tempLoopInfo.s2LoopTimes == s2SplitNum)) ? false : true; + tempLoopInfo.tndCoreStartKVSplitPos = globalLoopStart ? constInfo.coreStartKVSplitPos : 0; + uint32_t extraLoop = isEnd ? 2 : 0; + + uint32_t curTopKIdx = 0; + uint64_t curOffsetInSparseBlock = 0; + for (int s2LoopIdx = constInfo.s2Start; s2LoopIdx < (tempLoopInfo.s2LoopTimes + extraLoop); s2LoopIdx++) { + // PreloadPipeline loop初始值要求为 PRELOAD_NUM + PreloadPipeline(gloop, constInfo.s2Start, s2LoopIdx, extraInfo, curTopKIdx, curOffsetInSparseBlock); + ++gloop; + } + globalLoopStart = false; + constInfo.s2Start = 0; + } + constInfo.gS1Start = 0; + } + if ASCEND_IS_AIV { + CrossCoreWaitFlag(constInfo.syncC2V1); + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + CrossCoreWaitFlag(3); + CrossCoreWaitFlag(3); + CrossCoreWaitFlag(3); + CrossCoreWaitFlag(3); + } + } +} + +template +__aicore__ inline void +SparseFlashAttentionMla::PreloadPipeline(uint32_t loop, uint64_t s2Start, uint64_t s2LoopIdx, + RunInfo extraInfo[SFA_PRELOAD_TASK_CACHE_SIZE], uint32_t &curTopKIdx, uint64_t &curOffsetInSparseBlock) +{ + RunInfo &extraInfo0 = extraInfo[loop % SFA_PRELOAD_TASK_CACHE_SIZE]; // 本轮任务 + RunInfo &extraInfo2 = extraInfo[(loop + 2) % SFA_PRELOAD_TASK_CACHE_SIZE]; // 上一轮任务 + RunInfo &extraInfo1 = extraInfo[(loop + 1) % SFA_PRELOAD_TASK_CACHE_SIZE]; // 上两轮任务 + + CalcParams(loop, s2Start, s2LoopIdx, extraInfo0); + CalcSinnerTopKBegin(extraInfo0, curTopKIdx, curOffsetInSparseBlock); + + if (extraInfo0.isValid) { + if ASCEND_IS_AIC { + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + CrossCoreWaitFlag(constInfo.syncV0C1); + } + ComputeMm1(extraInfo0); + } else { + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + CrossCoreWaitFlag(3); + vectorService.MergeKv(extraInfo0); + CrossCoreSetFlag(constInfo.syncV0C1); + } + } + } + if (extraInfo2.isValid) { + if ASCEND_IS_AIV { + vectorService.ProcessVec1L(extraInfo2); + } + if ASCEND_IS_AIC { + ComputeMm2(extraInfo2); + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + CrossCoreSetFlag(3); + } + } + } + if (extraInfo1.isValid) { + if ASCEND_IS_AIV { + vectorService.ProcessVec2L(extraInfo1); + } + extraInfo1.isValid = false; + } +} + +template +__aicore__ inline uint64_t +SparseFlashAttentionMla::GetBalanceActualSeqLengths(GlobalTensor &actualSeqLengths, + uint32_t bIdx) +{ + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + if (bIdx > 0) { + return actualSeqLengths.GetValue(bIdx) - actualSeqLengths.GetValue(bIdx - 1); + } else if (bIdx == 0) { + return actualSeqLengths.GetValue(0); + } else { + return 0; + } + } else { + if (constInfo.actualLenDimsQ == 0) { + return constInfo.qSeqSize; + } else if (constInfo.actualLenDimsQ == 1) { + return actualSeqLengths.GetValue(0); + } else { + return actualSeqLengths.GetValue(bIdx); + } + } +} + +template +__aicore__ inline void SparseFlashAttentionMla::GetAxisStartIdx(uint32_t bN2EndPrev, + uint32_t s1GEndPrev, + uint32_t s2EndPrev) +{ + uint32_t bEndPrev = bN2EndPrev / kvHeadNum; + uint32_t actualSeqQPrev = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bEndPrev); + uint32_t s1GPrevBaseNum = (actualSeqQPrev * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize; + constInfo.bN2Start = bN2EndPrev; + constInfo.gS1Start = s1GEndPrev; + + constInfo.s2Start = 0; + if (s1GEndPrev >= s1GPrevBaseNum - 1) { // 上个核把S1G处理完了 + constInfo.gS1Start = 0; + constInfo.bN2Start++; + } else { + constInfo.gS1Start++; + } +} + +template +__aicore__ inline void SparseFlashAttentionMla::CalcSinnerTopKBegin(RunInfo &info, uint32_t &curTopKIdx, uint64_t &curOffsetInSparseBlock) + +{ + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + return; + } + + uint64_t thresholdSparseCount = (info.threshold + constInfo.sparseBlockSize - 1) / constInfo.sparseBlockSize; + uint64_t validCount = (constInfo.sparseBlockCount > thresholdSparseCount) ? thresholdSparseCount : constInfo.sparseBlockCount; + + int32_t sparseIndices = topKGm.GetValue(info.topKBaseOffset + curTopKIdx); + if (sparseIndices == -1 || curTopKIdx == validCount) { + info.actualSingleProcessSInnerSize = 0; + info.actualSingleProcessSInnerSizeAlign = 0; + tempLoopInfo.s2BasicSizeTail = 0; + if (curTopKIdx == 0) { + DealActSeqLenIsZero(info.bIdx, info.gS1Idx / constInfo.gSize, tempLoopInfo.n2Idx); + } + return; + } + + uint32_t sparseLen = 0; + uint64_t blockBegin = sparseIndices * constInfo.sparseBlockSize; + uint64_t blockEnd = (blockBegin + constInfo.sparseBlockSize > info.threshold) ? info.threshold : blockBegin + constInfo.sparseBlockSize; + int32_t blockLen = blockEnd - blockBegin; + sparseLen += (blockLen > static_cast(curOffsetInSparseBlock)) ? blockLen - curOffsetInSparseBlock : 0; + + bool firstVaildFlag = false; + if (curTopKIdx > 0) { + info.curTopKIdx = curTopKIdx; + info.curOffsetInSparseBlock = curOffsetInSparseBlock; + } else if (curTopKIdx == 0 && sparseLen > 0) { + info.curTopKIdx = curTopKIdx; + info.curOffsetInSparseBlock = 0; + firstVaildFlag = true; + } + + for (uint64_t topkIdx = curTopKIdx + 1; topkIdx < validCount; topkIdx++) { + int32_t sparseIndices = topKGm.GetValue(info.topKBaseOffset + topkIdx); + if (sparseIndices == -1) { + curTopKIdx = topkIdx; + curOffsetInSparseBlock = 0; + break; + } + uint64_t blockBegin = sparseIndices * constInfo.sparseBlockSize; + if (blockBegin >= info.threshold) { + continue; + } + if (firstVaildFlag == false && curTopKIdx == 0) { + info.curTopKIdx = topkIdx; + info.curOffsetInSparseBlock = 0; + firstVaildFlag = true; + } + uint64_t blockEnd = (blockBegin + constInfo.sparseBlockSize > info.threshold) ? info.threshold : blockBegin + constInfo.sparseBlockSize; + uint64_t blockLen = blockEnd - blockBegin; + sparseLen += blockLen; + if (sparseLen >= constInfo.s2BaseSize) { + curTopKIdx = topkIdx; + curOffsetInSparseBlock = blockLen - (sparseLen - constInfo.s2BaseSize); + sparseLen = constInfo.s2BaseSize; + break; + } + + if (topkIdx == validCount - 1) { + curTopKIdx = validCount; + curOffsetInSparseBlock = 0; + } + } + + info.actualSingleProcessSInnerSize = sparseLen; + info.actualSingleProcessSInnerSizeAlign = SFAAlign((uint32_t)info.actualSingleProcessSInnerSize, (uint32_t)SFAVectorService::BYTE_BLOCK); + tempLoopInfo.s2BasicSizeTail = (sparseLen == constInfo.s2BaseSize) ? 0 : sparseLen; + if (curTopKIdx == 0 && sparseLen == 0) { + DealActSeqLenIsZero(info.bIdx, info.gS1Idx / constInfo.gSize, tempLoopInfo.n2Idx); + } +} + + +#endif // SPARSE_FLASH_ATTENTION_KERNEL_MLA_H diff --git a/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch22/sparse_flash_attention_service_cube_mla.h b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch22/sparse_flash_attention_service_cube_mla.h new file mode 100644 index 000000000..d8d8f88fb --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch22/sparse_flash_attention_service_cube_mla.h @@ -0,0 +1,1121 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_flash_attention_service_cube_mla.h + * \brief use 7 buffer for matmul l1, better pipeline + */ +#ifndef SPARSE_FLASH_ATTENTION_SERVICE_CUBE_MLA_H +#define SPARSE_FLASH_ATTENTION_SERVICE_CUBE_MLA_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../sparse_flash_attention_common.h" + +struct PAShape { + uint32_t blockSize; + uint32_t headNum; //一般为kv的head num，对应n2 + uint32_t headDim; //mla下rope为64，nope为512, 对应d + uint32_t maxblockNumPerBatch; //block table 每一行的最大个数 + uint32_t actHeadDim; //实际拷贝col大小,考虑到N切块 s*d, 对应d + uint32_t copyRowNum; //总共要拷贝的行数 + uint32_t copyRowNumAlign; +}; + +struct Position { + uint32_t bIdx; + uint32_t n2Idx; + uint32_t s2Idx; + uint32_t dIdx; +}; + +// 场景：query、queryRope、key、value GM to L1 +// GM按ND格式存储 +// L1按NZ格式存储 +// GM的行、列、列的stride +template +__aicore__ inline void DataCopyGmNDToL1(LocalTensor &l1Tensor, GlobalTensor &gmTensor, + uint32_t rowAct, + uint32_t rowAlign, + uint32_t col, // D + uint32_t colStride) // D or N*D +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = rowAct; //nd矩阵的行数 + // T为int4场景下，dValue = col / 2，srcDValue = colStride / 2 + nd2nzPara.dValue = col; //nd矩阵的列数 + nd2nzPara.srcDValue = colStride; //同一nd矩阵相邻行起始地址间的偏移 + nd2nzPara.dstNzC0Stride = rowAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(l1Tensor, gmTensor, nd2nzPara); +} + +/* + 适用PA数据从GM拷贝到L1，支持ND、NZ数据； + PA的layout分 BNBD（blockNum,N,blockSize,D） BBH（blockNum,blockSize,N*D + BSH\BSND\TND 为BBH + shape.copyRowNumAlign 需要16字节对齐，如拷贝k矩阵，一次拷贝128*512，遇到尾块 10*512 需对齐到16*512 +*/ +template +__aicore__ inline void DataCopyPA(LocalTensor &dstTensor, //l1 + GlobalTensor &srcTensor, //gm + GlobalTensor &blockTableGm, + const PAShape &shape, // blockSize, headNum, headDim + const Position &startPos) // bacthIdx nIdx curSeqIdx +{ + uint32_t copyFinishRowCnt = 0; + uint64_t blockTableBaseOffset = startPos.bIdx * shape.maxblockNumPerBatch; + uint32_t curS2Idx = startPos.s2Idx; + uint32_t blockElementCnt = 32 / sizeof(T); + while (copyFinishRowCnt < shape.copyRowNum) { + uint64_t blockIdOffset = curS2Idx / shape.blockSize; // 获取block table上的索引 + uint64_t reaminRowCnt = curS2Idx % shape.blockSize; // 获取在单个块上超出的行数 + uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset); // 从block table上的获取编号 + // 计算可以拷贝行数 + uint32_t copyRowCnt = shape.blockSize - reaminRowCnt; //一次只能处理一个Block + if (copyFinishRowCnt + copyRowCnt > shape.copyRowNum) { + copyRowCnt = shape.copyRowNum - copyFinishRowCnt; //一个block未拷满 + } + uint64_t offset = idInBlockTable * shape.blockSize * shape.headNum * shape.headDim ; //PA的偏移 + + uint64_t dStride = shape.headDim; + if constexpr (SRC_LAYOUT == SFA_LAYOUT::BSND || SRC_LAYOUT == SFA_LAYOUT::TND) { + offset += (uint64_t)(startPos.n2Idx * shape.headDim) + + reaminRowCnt * shape.headDim * shape.headNum + startPos.dIdx; + dStride = shape.headDim * shape.headNum; + } else { + offset += (uint64_t)(startPos.n2Idx * shape.headDim * shape.blockSize) + + reaminRowCnt * shape.headDim + startPos.dIdx; + } + + uint32_t dValue = shape.actHeadDim; + uint32_t srcDValue = dStride; + LocalTensor tmpDstTensor = dstTensor[copyFinishRowCnt * blockElementCnt]; + GlobalTensor tmpSrcTensor = srcTensor[offset]; + + DataCopyGmNDToL1(tmpDstTensor, tmpSrcTensor, copyRowCnt, shape.copyRowNumAlign, dValue, srcDValue); + copyFinishRowCnt += copyRowCnt; + curS2Idx += copyRowCnt; + } +} + +template class SFAMatmulService { +public: + // 中间计算数据类型为float, 高精度模式 + using T = float; + using Q_T = typename SFAT::queryType; + using KV_T = typename SFAT::kvType; + using OUT_T = typename SFAT::outputType; + using MM_OUT_T = T; + + __aicore__ inline SFAMatmulService(){}; + __aicore__ inline void InitParams(const ConstInfo &constInfo); + __aicore__ inline void InitMm1GlobalTensor(GlobalTensor queryGm, GlobalTensor qRopeGm, + GlobalTensor keyGm, GlobalTensor kRopeGm, + GlobalTensor mm1ResGm); + __aicore__ inline void InitMm2GlobalTensor(GlobalTensor vec1ResGm, GlobalTensor valueGm, + GlobalTensor mm2ResGm, GlobalTensor attentionOutGm); + __aicore__ inline void InitPageAttentionInfo(const GlobalTensor& kvMergeGm, + GlobalTensor blockTableGm, GlobalTensor topKGm, + uint32_t blockSize, uint32_t maxBlockNumPerBatch); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void UpdateKey(GlobalTensor keyGm); + __aicore__ inline void UpdateValue(GlobalTensor valueGm); + + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void CalcTopKBlockInfo(const RunInfo &info, uint32_t &curTopKIdx, + uint64_t &curOffsetInSparseBlock, uint32_t curSeqIdx, + uint32_t ©RowCnt, int64_t &idInTopK); + __aicore__ inline void ComputeMm1(const RunInfo &info, const MSplitInfo mSplitInfo); + __aicore__ inline void ComputeMm2(const RunInfo &info, const MSplitInfo mSplitInfo); + +private: + static constexpr bool PAGE_ATTENTION = SFAT::pageAttention; + static constexpr int TEMPLATE_MODE = SFAT::templateMode; + static constexpr bool FLASH_DECODE = SFAT::flashDecode; + static constexpr SFA_LAYOUT LAYOUT_T = SFAT::layout; + static constexpr SFA_LAYOUT KV_LAYOUT_T = SFAT::kvLayout; + + static constexpr uint32_t M_SPLIT_SIZE = 128; // m方向切分 + static constexpr uint32_t N_SPLIT_SIZE = 128; // n方向切分 + static constexpr uint32_t N_WORKSPACE_SIZE = 512; // n方向切分 + + static constexpr uint32_t L1_BLOCK_SIZE = (64 * (512 + 64) * sizeof(Q_T)); + static constexpr uint32_t L1_BLOCK_OFFSET = 64 * (512 + 64); // 72K的元素个数 + + static constexpr uint32_t L0A_PP_SIZE = (32 * 1024); + static constexpr uint32_t L0B_PP_SIZE = (32 * 1024); + static constexpr uint32_t L0C_PP_SIZE = (64 * 1024); + + // mte2 <> mte1 EventID + // L1 3buf, 使用3个eventId + static constexpr uint32_t L1_EVENT0 = EVENT_ID2; + static constexpr uint32_t L1_EVENT1 = EVENT_ID3; + static constexpr uint32_t L1_EVENT2 = EVENT_ID4; + static constexpr uint32_t L1_EVENT3 = EVENT_ID5; + static constexpr uint32_t L1_EVENT4 = EVENT_ID6; + static constexpr uint32_t L1_EVENT5 = EVENT_ID7; + static constexpr uint32_t L1_EVENT6 = EVENT_ID1; + + // m <> mte1 EventID + static constexpr uint32_t L0AB_EVENT0 = EVENT_ID3; + static constexpr uint32_t L0AB_EVENT1 = EVENT_ID4; + + static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding; + static constexpr uint32_t mte21QPIds[4] = {L1_EVENT0, L1_EVENT1, L1_EVENT2, L1_EVENT3}; // mte12复用 + static constexpr uint32_t mte21KVIds[3] = {L1_EVENT4, L1_EVENT5, L1_EVENT6}; + + uint32_t kvCacheBlockSize = 0; + uint32_t maxBlockNumPerBatch = 0; + ConstInfo constInfo{}; + + // L1分成3块buf, 用于记录 + uint32_t qpL1BufIter = 0; + uint32_t kvL1BufIter = -1; + uint32_t abL0BufIter = 0; + uint32_t cL0BufIter = 0; + + // mm1 + GlobalTensor queryGm; + GlobalTensor qRopeGm; + GlobalTensor keyGm; + GlobalTensor kRopeGm; + GlobalTensor mm1ResGm; + GlobalTensor kvMergeGm_; + + // mm2 + GlobalTensor vec1ResGm; + GlobalTensor valueGm; + GlobalTensor mm2ResGm; + GlobalTensor attentionOutGm; + + // block_table + GlobalTensor blockTableGm; + GlobalTensor topKGm; + + TBuf bufQPL1; + TBuf bufKVL1; + TBuf tmpBufL0A; + TBuf tmpBufL0B; + TBuf tmpBufL0C; + + LocalTensor l1QPTensor; + LocalTensor l1KVTensor; + LocalTensor aL0TensorPingPong; + LocalTensor bL0TensorPingPong; + LocalTensor cL0TensorPingPong; + + // L0AB m <> mte1 EventID + __aicore__ inline uint32_t Mte1MmABEventId(uint32_t idx) + { + return (L0AB_EVENT0 + idx); + } + + __aicore__ inline uint32_t GetQPL1RealIdx(uint32_t mIdx, uint32_t k1Idx) + { + uint32_t idxMap[] = {0, 2}; // 确保0块和1块连在一起, 2和3块连在一起, 来保证同一m块的地址相连 + return idxMap[mIdx % 2] + k1Idx; + } + + __aicore__ inline void CopyGmToL1(LocalTensor &l1Tensor, GlobalTensor &gmSrcTensor, uint32_t srcN, + uint32_t srcD, uint32_t srcDstride); + __aicore__ inline void CopyInMm1AToL1(LocalTensor &aL1Tensor, const RunInfo &info, uint32_t mSeqIdx, + uint32_t mSizeAct, uint32_t headSize, uint32_t headOffset); + __aicore__ inline void CopyInMm1ARopeToL1(LocalTensor &aL1Tensor, const RunInfo &info, uint32_t mSeqIdx, + uint32_t mSizeAct); + __aicore__ inline void CopyInMm1BToL1(LocalTensor &bL1Tensor, const uint64_t keyGmBaseOffset, + uint32_t copyTotalRowCntAlign, uint32_t copyStartRowCnt, + uint32_t nActCopyRowCount, uint32_t headSize); + __aicore__ inline void CopyInMm1BRopeToL1(LocalTensor &bL1Tensor, const uint64_t keyGmBaseOffset, + uint32_t copyTotalRowCntAlign, uint32_t copyStartRowCnt, + uint32_t nActCopyRowCount, uint32_t headSize); + __aicore__ inline void CopyInMm2AToL1(LocalTensor &aL1Tensor, const RunInfo &info, uint32_t mSeqIdx, + uint32_t subMSizeAct, uint32_t nSize, uint32_t nOffset); + __aicore__ inline void CopyInMm2BToL1(LocalTensor &bL1Tensor, const uint64_t valueGmBaseOffset, + uint32_t copyTotalRowCntAlign, uint32_t copyStartRowCnt, + uint32_t nActCopyRowCount, uint32_t copyStartColumnCount, + uint32_t copyColumnCount); + __aicore__ inline void LoadDataMm1A(LocalTensor &aL0Tensor, LocalTensor &aL1Tensor, uint32_t idx, + uint32_t kSplitSize, uint32_t mSize, uint32_t kSize); + __aicore__ inline void LoadDataMm1B(LocalTensor &bL0Tensor, LocalTensor &bL1Tensor, uint32_t idx, + uint32_t kSplitSize, uint32_t kSize, uint32_t nSize); +}; + +template __aicore__ inline void SFAMatmulService::InitParams(const ConstInfo &constInfo) +{ + this->constInfo = constInfo; +} + +template +__aicore__ inline void +SFAMatmulService::InitMm1GlobalTensor(GlobalTensor queryGm, GlobalTensor qRopeGm, + GlobalTensor keyGm, GlobalTensor kRopeGm, + GlobalTensor mm1ResGm) +{ + // mm1 + this->queryGm = queryGm; + this->qRopeGm = qRopeGm; + this->keyGm = keyGm; + this->kRopeGm = kRopeGm; + this->mm1ResGm = mm1ResGm; +} + +template +__aicore__ inline void +SFAMatmulService::InitMm2GlobalTensor(GlobalTensor vec1ResGm, GlobalTensor valueGm, + GlobalTensor mm2ResGm, GlobalTensor attentionOutGm) +{ + // mm2 + this->vec1ResGm = vec1ResGm; + this->valueGm = valueGm; + this->mm2ResGm = mm2ResGm; + this->attentionOutGm = attentionOutGm; +} + +template +__aicore__ inline void +SFAMatmulService::InitPageAttentionInfo(const GlobalTensor& kvMergeGm, GlobalTensor blockTableGm, + GlobalTensor topKGm, uint32_t blockSize, uint32_t maxBlockNumPerBatch) +{ + this->blockTableGm = blockTableGm; + this->topKGm = topKGm; + this->kvCacheBlockSize = blockSize; + this->maxBlockNumPerBatch = maxBlockNumPerBatch; + this->kvMergeGm_ = kvMergeGm; +} + +template __aicore__ inline void SFAMatmulService::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(bufQPL1, L1_BLOCK_SIZE * 4); // (64K + 8K) * 4 + l1QPTensor = bufQPL1.Get(); + pipe->InitBuffer(bufKVL1, L1_BLOCK_SIZE * 3); // (64K + 8K) * 3 + l1KVTensor = bufKVL1.Get(); + + // L0A + pipe->InitBuffer(tmpBufL0A, L0A_PP_SIZE * 2); // 64K + aL0TensorPingPong = tmpBufL0A.Get(); + // L0B + pipe->InitBuffer(tmpBufL0B, L0B_PP_SIZE * 2); // 64K + bL0TensorPingPong = tmpBufL0B.Get(); + // L0C + pipe->InitBuffer(tmpBufL0C, L0C_PP_SIZE * 2); // 128K + cL0TensorPingPong = tmpBufL0C.Get(); +} + +template __aicore__ inline void SFAMatmulService::UpdateKey(GlobalTensor keyGm) +{ + this->keyGm = keyGm; +} + +template __aicore__ inline void SFAMatmulService::UpdateValue(GlobalTensor valueGm) +{ + this->valueGm = valueGm; +} + +template __aicore__ inline void SFAMatmulService::AllocEventID() +{ + SetFlag(L1_EVENT0); + SetFlag(L1_EVENT1); + SetFlag(L1_EVENT2); + SetFlag(L1_EVENT3); + SetFlag(L1_EVENT4); + SetFlag(L1_EVENT5); + SetFlag(L1_EVENT6); + SetFlag(L0AB_EVENT0); + SetFlag(L0AB_EVENT1); +} + +template __aicore__ inline void SFAMatmulService::FreeEventID() +{ + WaitFlag(L1_EVENT0); + WaitFlag(L1_EVENT1); + WaitFlag(L1_EVENT2); + WaitFlag(L1_EVENT3); + WaitFlag(L1_EVENT4); + WaitFlag(L1_EVENT5); + WaitFlag(L1_EVENT6); + WaitFlag(L0AB_EVENT0); + WaitFlag(L0AB_EVENT1); +} + +template +__aicore__ inline void SFAMatmulService::CopyGmToL1(LocalTensor &l1Tensor, + GlobalTensor &gmSrcTensor, uint32_t srcN, + uint32_t srcD, uint32_t srcDstride) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = srcN; // 行数 + nd2nzPara.dValue = srcD; + nd2nzPara.srcDValue = srcDstride; + nd2nzPara.dstNzC0Stride = (srcN + 15) / 16 * 16; // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(l1Tensor, gmSrcTensor, nd2nzPara); +} + +template +__aicore__ inline void SFAMatmulService::CopyInMm1AToL1(LocalTensor &l1Tensor, const RunInfo &info, + uint32_t mSeqIdx, uint32_t mSizeAct, + uint32_t headSize, uint32_t headOffset) +{ + auto srcGm = queryGm[info.tensorAOffset + mSeqIdx * constInfo.headDim + headOffset]; + CopyGmToL1(l1Tensor, srcGm, mSizeAct, headSize, constInfo.headDim); +} + +template +__aicore__ inline void SFAMatmulService::CopyInMm1ARopeToL1(LocalTensor &l1Tensor, + const RunInfo &info, uint32_t mSeqIdx, + uint32_t mSizeAct) +{ + auto srcGm = qRopeGm[info.tensorARopeOffset + mSeqIdx * constInfo.headDimRope]; + CopyGmToL1(l1Tensor, srcGm, mSizeAct, constInfo.headDimRope, constInfo.headDimRope); +} + +template +__aicore__ inline void +SFAMatmulService::CopyInMm1BToL1(LocalTensor &bL1Tensor, const uint64_t keyGmBaseOffset, + uint32_t copyTotalRowCntAlign, uint32_t copyStartRowCnt, + uint32_t nActCopyRowCount, uint32_t headSize) +{ + uint64_t dStride = constInfo.headDim; + if constexpr (KV_LAYOUT_T == SFA_LAYOUT::BSND || KV_LAYOUT_T == SFA_LAYOUT::TND) { + dStride = constInfo.headDim * constInfo.kvHeadNum; + } + + uint32_t blockElementCnt = 32 / sizeof(KV_T); + + Nd2NzParams mm1Nd2NzParamsForB; + mm1Nd2NzParamsForB.ndNum = 1; + mm1Nd2NzParamsForB.nValue = nActCopyRowCount; + mm1Nd2NzParamsForB.dValue = headSize; + mm1Nd2NzParamsForB.srcDValue = dStride; + mm1Nd2NzParamsForB.dstNzC0Stride = copyTotalRowCntAlign; + mm1Nd2NzParamsForB.dstNzNStride = 1; + mm1Nd2NzParamsForB.srcNdMatrixStride = 0; + mm1Nd2NzParamsForB.dstNzMatrixStride = 0; + DataCopy(bL1Tensor[copyStartRowCnt * blockElementCnt], keyGm[keyGmBaseOffset], mm1Nd2NzParamsForB); +} + +template +__aicore__ inline void +SFAMatmulService::CopyInMm1BRopeToL1(LocalTensor &bL1Tensor, const uint64_t kRopeGmBaseOffset, + uint32_t copyTotalRowCntAlign, uint32_t copyStartRowCnt, + uint32_t nActCopyRowCount, uint32_t headSize) +{ + uint64_t dStride = constInfo.headDimRope; + if constexpr (KV_LAYOUT_T == SFA_LAYOUT::BSND || KV_LAYOUT_T == SFA_LAYOUT::TND) { + dStride = constInfo.headDimRope * constInfo.kvHeadNum; + } + + uint32_t blockElementCnt = 32 / sizeof(KV_T); + + Nd2NzParams mm1Nd2NzParamsForB; + mm1Nd2NzParamsForB.ndNum = 1; + mm1Nd2NzParamsForB.nValue = nActCopyRowCount; + mm1Nd2NzParamsForB.dValue = headSize; + mm1Nd2NzParamsForB.srcDValue = dStride; + mm1Nd2NzParamsForB.dstNzC0Stride = copyTotalRowCntAlign; + mm1Nd2NzParamsForB.dstNzNStride = 1; + mm1Nd2NzParamsForB.srcNdMatrixStride = 0; + mm1Nd2NzParamsForB.dstNzMatrixStride = 0; + DataCopy(bL1Tensor[copyStartRowCnt * blockElementCnt], kRopeGm[kRopeGmBaseOffset], mm1Nd2NzParamsForB); +} + +template +__aicore__ inline void SFAMatmulService::LoadDataMm1A(LocalTensor &aL0Tensor, + LocalTensor &aL1Tensor, uint32_t idx, + uint32_t kSplitSize, uint32_t mSize, uint32_t kSize) +{ + LocalTensor srcTensor = aL1Tensor[mSize * kSplitSize * idx]; + LoadData3DParamsV2 loadData3DParams; + // SetFmatrixParams + loadData3DParams.l1H = mSize / 16; // Hin=M1=8 + loadData3DParams.l1W = 16; // Win=M0 + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + // SetLoadToA0Params + loadData3DParams.mExtension = mSize; // M + loadData3DParams.kExtension = kSize; // K + loadData3DParams.mStartPt = 0; + loadData3DParams.kStartPt = 0; + loadData3DParams.strideW = 1; + loadData3DParams.strideH = 1; + loadData3DParams.filterW = 1; + loadData3DParams.filterSizeW = (1 >> 8) & 255; + loadData3DParams.filterH = 1; + loadData3DParams.filterSizeH = (1 >> 8) & 255; + loadData3DParams.dilationFilterW = 1; + loadData3DParams.dilationFilterH = 1; + loadData3DParams.enTranspose = 0; + loadData3DParams.fMatrixCtrl = 0; + loadData3DParams.channelSize = kSize; // Cin=K + LoadData(aL0Tensor, srcTensor, loadData3DParams); +} + +template +__aicore__ inline void SFAMatmulService::LoadDataMm1B(LocalTensor &l0Tensor, + LocalTensor &l1Tensor, uint32_t idx, + uint32_t kSplitSize, uint32_t kSize, uint32_t nSize) +{ + // N 方向全载 + LocalTensor srcTensor = l1Tensor[nSize * kSplitSize * idx]; + + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = (nSize + 15) / 16 * kSize / (32 / sizeof(KV_T)); + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = false; + LoadData(l0Tensor, srcTensor, loadData2DParams); +} + +template +__aicore__ inline void SFAMatmulService::CopyInMm2AToL1(LocalTensor &aL1Tensor, const RunInfo &info, + uint32_t mSeqIdx, uint32_t subMSizeAct, + uint32_t nSize, uint32_t nOffset) +{ + auto srcGm = vec1ResGm[(info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize + + mSeqIdx * info.actualSingleProcessSInnerSizeAlign + nOffset]; + CopyGmToL1(aL1Tensor, srcGm, subMSizeAct, nSize, info.actualSingleProcessSInnerSizeAlign); +} + +template +__aicore__ inline void SFAMatmulService::CopyInMm2BToL1( + LocalTensor &bL1Tensor, const uint64_t valueGmBaseOffset, uint32_t copyTotalRowCntAlign, + uint32_t copyStartRowCnt, uint32_t nActCopyRowCount, uint32_t copyStartColumnCount, uint32_t copyColumnCount) +{ + uint64_t step = constInfo.headDim; + if constexpr (KV_LAYOUT_T == SFA_LAYOUT::BSND || KV_LAYOUT_T == SFA_LAYOUT::TND) { + step = constInfo.headDim * constInfo.kvHeadNum; + } + + uint32_t blockElementCnt = 32 / sizeof(KV_T); + + Nd2NzParams mm1Nd2NzParamsForB; + mm1Nd2NzParamsForB.ndNum = 1; + mm1Nd2NzParamsForB.nValue = nActCopyRowCount; + mm1Nd2NzParamsForB.dValue = copyColumnCount; + mm1Nd2NzParamsForB.srcDValue = step; + mm1Nd2NzParamsForB.dstNzC0Stride = copyTotalRowCntAlign; + mm1Nd2NzParamsForB.dstNzNStride = 1; + mm1Nd2NzParamsForB.srcNdMatrixStride = 0; + mm1Nd2NzParamsForB.dstNzMatrixStride = 0; + DataCopy(bL1Tensor[copyStartRowCnt * blockElementCnt], valueGm[valueGmBaseOffset + copyStartColumnCount], + mm1Nd2NzParamsForB); +} + +template +__aicore__ inline void SFAMatmulService::CalcTopKBlockInfo( + const RunInfo &info, uint32_t &curTopKIdx, uint64_t &curOffsetInSparseBlock, uint32_t curSeqIdx, uint32_t ©RowCnt, int64_t &idInTopK) +{ + uint64_t blockBegin = idInTopK * constInfo.sparseBlockSize; + uint64_t blockEnd = (blockBegin + constInfo.sparseBlockSize > info.threshold) ? + info.threshold : blockBegin + constInfo.sparseBlockSize; + uint64_t blockLen = blockEnd - blockBegin; + if (curOffsetInSparseBlock + copyRowCnt < blockLen) { + curOffsetInSparseBlock += copyRowCnt; + copyRowCnt = blockLen - curOffsetInSparseBlock; + } else { + for (uint64_t topkidx = curTopKIdx + 1; topkidx < constInfo.sparseBlockCount; topkidx++) { + int64_t sparseIndices = topKGm.GetValue(info.topKBaseOffset + topkidx); + if (sparseIndices == -1) { + break; + } + + uint64_t blockBegin = sparseIndices * constInfo.sparseBlockSize; + if (blockBegin >= info.threshold) { + continue; + } + uint64_t blockEnd = (blockBegin + constInfo.sparseBlockSize > info.threshold) ? + info.threshold : blockBegin + constInfo.sparseBlockSize; + uint64_t blockLen = blockEnd - blockBegin; + curTopKIdx = topkidx; + idInTopK = sparseIndices; + curOffsetInSparseBlock = 0; + copyRowCnt = blockLen; + break; + } + } +} + +template +__aicore__ inline void SFAMatmulService::ComputeMm1(const RunInfo &info, const MSplitInfo mSplitInfo) +{ + // 最外层还需要一层m的循环 + uint32_t mSize = mSplitInfo.nBufferDealM; + uint32_t mL1Size = M_SPLIT_SIZE; + uint32_t mL1SizeAlign = SFAAlign(M_SPLIT_SIZE, 16U); + uint32_t mL1Loops = (mSize + M_SPLIT_SIZE - 1) / M_SPLIT_SIZE; + + uint32_t nSize = info.actualSingleProcessSInnerSize; + uint32_t nL1Size = N_SPLIT_SIZE; + uint32_t nL1SizeAlign = SFAAlign(N_SPLIT_SIZE, 16U); + uint32_t nL1Loops = (nSize + N_SPLIT_SIZE - 1) / N_SPLIT_SIZE; + + uint32_t kSize = 576; + uint32_t kL1Size = 288; + uint32_t kL1Loops = 2; // 2 : 576/288, mla专用 这里不考虑d泛化 + + uint32_t kL0Size = 96; + uint32_t kL0Loops = (kL1Size + kL0Size - 1) / kL0Size; // 288 / 96 = 3 kloops + + LocalTensor bL1Tensor; + LocalTensor kRopeTensor; + LocalTensor kTensor; + // ka表示左矩阵4buf选择哪一块buf, kb表示右矩阵3buf选择哪一块buf + uint32_t ka = 0, kb = 0; + + uint32_t curTopKIdx = info.curTopKIdx; + uint64_t curOffsetInSparseBlock = info.curOffsetInSparseBlock; //sparse Block块内偏移 + uint32_t copyRowCnt = 0; + int64_t idInTopK = topKGm.GetValue(info.topKBaseOffset + curTopKIdx); + + uint32_t curTopKIdxTmp = 0; + uint64_t curOffsetInSparseBlockTmp = 0; + uint32_t copyRowCntTmp = 0; + int64_t idInTopKTmp = 0; + + // L1 切n切k切m + for (uint32_t nL1 = 0; nL1 < nL1Loops; nL1++) { // L1切n, 512/128=4 + if (nL1 == (nL1Loops - 1)) { + // 尾块重新计算size + nL1Size = nSize - (nL1Loops - 1) * N_SPLIT_SIZE; + nL1SizeAlign = SFAAlign(nL1Size, 16U); + } + curTopKIdxTmp = curTopKIdx; + curOffsetInSparseBlockTmp = curOffsetInSparseBlock; + copyRowCntTmp = copyRowCnt; + idInTopKTmp = idInTopK; + + for (uint32_t kL1 = 0; kL1 < kL1Loops; kL1++) { // L1切k, 576/288, 这里不考虑d泛化 + kvL1BufIter++; + uint32_t kb = kvL1BufIter % 3; + WaitFlag(mte21KVIds[kb]); + // 从k当中取当前的块 + bL1Tensor = l1KVTensor[kb * L1_BLOCK_OFFSET]; + // mm1拷贝主流程 + + uint32_t curSeqIdx = info.s2BatchOffset + nL1 * N_SPLIT_SIZE; + uint32_t copyFinishRowCnt = 0; + curTopKIdx = curTopKIdxTmp; + curOffsetInSparseBlock = curOffsetInSparseBlockTmp; + copyRowCnt = copyRowCntTmp; + idInTopK = idInTopKTmp; + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + if (kL1 == 0) { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nL1Size; // 行数 + nd2nzPara.dValue = constInfo.headDim >> 1; // constInfo.headDim; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = nL1SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor, + kvMergeGm_[info.loop % 4 * N_WORKSPACE_SIZE * kSize + + nL1 * N_SPLIT_SIZE * constInfo.headDim], + nd2nzPara); + nd2nzPara.dValue = constInfo.headDimRope >> 1; + nd2nzPara.srcDValue = constInfo.headDimRope; + DataCopy( + bL1Tensor[nL1SizeAlign * (constInfo.headDim >> 1)], + kvMergeGm_[info.loop % 4 * N_WORKSPACE_SIZE * kSize + N_WORKSPACE_SIZE * constInfo.headDim + + nL1 * N_SPLIT_SIZE * constInfo.headDimRope], + nd2nzPara); + } else { + LocalTensor kTmpTensor = bL1Tensor[(constInfo.headDimRope >> 1) * nL1SizeAlign]; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = nL1Size; // 行数 + nd2nzPara.dValue = constInfo.headDim >> 1; // constInfo.headDim; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = nL1SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(kTmpTensor, + kvMergeGm_[info.loop % 4 * N_WORKSPACE_SIZE * kSize + (constInfo.headDim >> 1) + + nL1 * N_SPLIT_SIZE * constInfo.headDim], + nd2nzPara); + nd2nzPara.dValue = constInfo.headDimRope >> 1; + nd2nzPara.srcDValue = constInfo.headDimRope; + DataCopy( + bL1Tensor, + kvMergeGm_[info.loop % 4 * N_WORKSPACE_SIZE * kSize + N_WORKSPACE_SIZE * constInfo.headDim + + (constInfo.headDimRope >> 1) + nL1 * N_SPLIT_SIZE * constInfo.headDimRope], + nd2nzPara); + } + } else { + while (copyFinishRowCnt < nL1Size) { + CalcTopKBlockInfo(info, curTopKIdx, curOffsetInSparseBlock, curSeqIdx, copyRowCnt, idInTopK); + if (copyFinishRowCnt + copyRowCnt > nL1Size) { + copyRowCnt = nL1Size - copyFinishRowCnt; + } + + // BN2轴偏移 + if constexpr (PAGE_ATTENTION) { + Position startPos; + startPos.bIdx = info.bIdx; + startPos.n2Idx = info.n2Idx; + startPos.s2Idx = idInTopK * constInfo.sparseBlockSize + curOffsetInSparseBlock; + // 256、32等待7buf命名更改 + startPos.dIdx = kL1 * 256; // mm1 右矩阵 bn2s2d, d为k轴不切; mm2 右矩阵, s2为k轴, d轴切分 + Position ropeStartPos = startPos; + ropeStartPos.dIdx = kL1 * 32; + PAShape shape; + shape.blockSize = kvCacheBlockSize; + shape.headNum = constInfo.kvHeadNum; + shape.headDim = constInfo.headDim; + shape.actHeadDim = 256; + shape.maxblockNumPerBatch = maxBlockNumPerBatch; + shape.copyRowNum = copyRowCnt; + shape.copyRowNumAlign = nL1SizeAlign; + PAShape ropeShape = shape; + ropeShape.headDim = constInfo.headDimRope; + ropeShape.actHeadDim = 32; + if (kL1 == 0) { + kTensor = bL1Tensor[copyFinishRowCnt * 16]; + DataCopyPA(kTensor, keyGm, blockTableGm, shape, startPos); + kRopeTensor = bL1Tensor[(nL1SizeAlign * (BlockAlign(constInfo.headDim) >> 1)) + + copyFinishRowCnt * 16]; + DataCopyPA(kRopeTensor, kRopeGm, blockTableGm, ropeShape, + ropeStartPos); + } else { + kRopeTensor = bL1Tensor[copyFinishRowCnt * 16]; + DataCopyPA(kRopeTensor, kRopeGm, blockTableGm, ropeShape, + ropeStartPos); + LocalTensor kTmpTensor = bL1Tensor[32 * nL1SizeAlign + copyFinishRowCnt * 16]; + DataCopyPA(kTmpTensor, keyGm, blockTableGm, shape, startPos); + } + } else { + uint64_t keyOffset = info.tensorBOffset; + uint64_t kRopeOffset = info.tensorBRopeOffset; + if constexpr (KV_LAYOUT_T == SFA_LAYOUT::BSND || KV_LAYOUT_T == SFA_LAYOUT::TND) { + keyOffset += (idInTopK * constInfo.sparseBlockSize + curOffsetInSparseBlock) * + constInfo.kvHeadNum * constInfo.headDim; + kRopeOffset += (idInTopK * constInfo.sparseBlockSize + curOffsetInSparseBlock) * + constInfo.kvHeadNum * constInfo.headDimRope; + } else { + keyOffset += (idInTopK * constInfo.sparseBlockSize + curOffsetInSparseBlock) * + constInfo.headDim; + kRopeOffset += (idInTopK * constInfo.sparseBlockSize + curOffsetInSparseBlock) * + constInfo.headDimRope; + } + + if (kL1 == 0) { + CopyInMm1BToL1(bL1Tensor, keyOffset, nL1SizeAlign, copyFinishRowCnt, copyRowCnt, 256); + kRopeTensor = bL1Tensor[nL1SizeAlign * (BlockAlign(constInfo.headDim) >> 1)]; + CopyInMm1BRopeToL1(kRopeTensor, kRopeOffset, nL1SizeAlign, copyFinishRowCnt, copyRowCnt, + 32); + } else { + kRopeTensor = bL1Tensor; + CopyInMm1BRopeToL1(kRopeTensor, kRopeOffset + 32, nL1SizeAlign, copyFinishRowCnt, + copyRowCnt, 32); + LocalTensor kTmpTensor = bL1Tensor[nL1SizeAlign * 32]; + CopyInMm1BToL1(kTmpTensor, keyOffset + 256, nL1SizeAlign, copyFinishRowCnt, copyRowCnt, + 256); + } + } + + // 更新循环变量 + copyFinishRowCnt += copyRowCnt; + curSeqIdx += copyRowCnt; + } + } + + SetFlag(mte21KVIds[kb]); + WaitFlag(mte21KVIds[kb]); + mL1Size = M_SPLIT_SIZE; + mL1SizeAlign = SFAAlign(M_SPLIT_SIZE, 16U); + for (uint32_t mL1 = 0; mL1 < mL1Loops; mL1++) { + uint32_t aL1PaddingSize = 0; // 用于使左矩阵对齐到尾部, 以保证两块32K内存连续 + if (mL1 == (mL1Loops - 1)) { + // 尾块重新计算size + mL1Size = mSize - (mL1Loops - 1) * M_SPLIT_SIZE; + mL1SizeAlign = SFAAlign(mL1Size, 16U); + // mL1SizeAlign<128 kL1=0时需要偏移, 确保qRope能一半拷贝到当前tensor, 一半拷贝到下一个tensor + aL1PaddingSize = (M_SPLIT_SIZE - mL1SizeAlign) * 288; + } + + // 左矩阵L1选择12块还是34块的index, 由m l1 index决定 + // 左矩阵L1选择12块或34块的前一块还是后一块, 由k l1 index决定 + uint32_t mIdx = qpL1BufIter + mL1; + ka = GetQPL1RealIdx(mIdx, kL1); + LocalTensor aL1Tensor = + l1QPTensor[ka * L1_BLOCK_OFFSET + (1 - kL1) * aL1PaddingSize]; // kL1=0时需要偏移 + if (nL1 == 0) { // mL1=0, mL1=1两次 + if (kL1 == 0) { + WaitFlag(mte21QPIds[ka]); + WaitFlag(mte21QPIds[ka + 1]); + CopyInMm1AToL1(aL1Tensor, info, mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE, mL1Size, 256, 0); + // 由于L1里面是NZ, 这里q rope的偏移为整块q nope切k的后大小, 256为headDim的一半 + LocalTensor qRopeTensor = + aL1Tensor[mL1SizeAlign * + 256]; + CopyInMm1ARopeToL1(qRopeTensor, info, mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE, mL1Size); + } else { + // 32为rope headDim的一半 + LocalTensor qTmpTensor = aL1Tensor[mL1SizeAlign * 32]; + CopyInMm1AToL1(qTmpTensor, info, mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE, mL1Size, 256, + 256); + } + SetFlag(mte21QPIds[ka]); + WaitFlag(mte21QPIds[ka]); + } + + // 使用unitflag同步 + LocalTensor cL0Tensor = + cL0TensorPingPong[(cL0BufIter % 2) * + (L0C_PP_SIZE / sizeof(MM_OUT_T))]; // 需要保证cL0BufIter和m步调一致 + for (uint32_t kL0 = 0; kL0 < kL0Loops; kL0++) { + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + LocalTensor aL0Tensor = aL0TensorPingPong[(abL0BufIter % 2) * (L0A_PP_SIZE / sizeof(KV_T))]; + LoadDataMm1A(aL0Tensor, aL1Tensor, kL0, kL0Size, mL1SizeAlign, kL0Size); + LocalTensor bL0Tensor = bL0TensorPingPong[(abL0BufIter % 2) * (L0B_PP_SIZE / sizeof(KV_T))]; + LoadDataMm1B(bL0Tensor, bL1Tensor, kL0, kL0Size, kL0Size, nL1SizeAlign); + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + + // m == 1的时候需要特殊处理 + MmadParams mmadParams; + mmadParams.m = mL1SizeAlign; + mmadParams.n = nL1SizeAlign; + mmadParams.k = kL0Size; + mmadParams.cmatrixInitVal = (kL1 == 0 && kL0 == 0); + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = + (kL1 == 1 && kL0 == (kL0Loops - 1)) ? 0b11 : 0b10; // 累加最后一次翻转flag, 表示可以搬出 + Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams); + + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + abL0BufIter++; + } + + if (nL1 == (nL1Loops - 1)) { + SetFlag(mte21QPIds[ka]); // 反向同步, 表示L1中的A已经被mte1消费完 + } + + if (kL1 == 1) { // 最后一轮kL1循环 + FixpipeParamsV220 fixParams; + fixParams.nSize = nL1SizeAlign; + fixParams.mSize = mL1SizeAlign; + fixParams.srcStride = mL1SizeAlign; + // 改成nSizeAlign + fixParams.dstStride = info.actualSingleProcessSInnerSizeAlign; // mm1ResGm两行之间的间隔 + fixParams.unitFlag = 0b11; + fixParams.ndNum = 1; // 输出ND + + // 输出偏移info.loop % (constInfo.preLoadNum)) * mmResUbSize是否在matmul里计算 + Fixpipe(mm1ResGm[(info.loop % (constInfo.preLoadNum)) * constInfo.mmResUbSize + nL1 * N_SPLIT_SIZE + + (mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE) * + info.actualSingleProcessSInnerSizeAlign], + cL0Tensor, fixParams); + } + if (mL1Loops == 2) { + cL0BufIter++; + } + } + SetFlag(mte21KVIds[kb]); // 反向同步, 表示L1已经被mte1消费完 + } + if (mL1Loops == 1) { + cL0BufIter++; + } + } + qpL1BufIter += mL1Loops; +} + +template +__aicore__ inline void SFAMatmulService::ComputeMm2(const RunInfo &info, const MSplitInfo mSplitInfo) +{ + uint32_t mSize = mSplitInfo.nBufferDealM; + uint32_t mSizeAlign = (mSize + 16 - 1) / 16; + uint32_t mL1Loops = (mSize + M_SPLIT_SIZE - 1) / M_SPLIT_SIZE; + uint32_t mL1SizeAlign = M_SPLIT_SIZE; // 16对齐 + uint32_t mL1Size = M_SPLIT_SIZE; // m的实际大小 + + uint32_t nSize = BlockAlign(constInfo.headDim); + uint32_t nL1Loops = (nSize + N_SPLIT_SIZE - 1) / N_SPLIT_SIZE; + uint32_t nL1SizeAlign = N_SPLIT_SIZE; // 16对齐 + uint32_t nL1Size = N_SPLIT_SIZE; // n的实际大小 + + uint32_t kSize = info.actualSingleProcessSInnerSize; + uint32_t kL1Size = 256; + uint32_t kL1SizeAlign = SFAAlign(kL1Size, 16U); + uint32_t kL1Loops = (kSize + kL1Size - 1) / kL1Size; + uint32_t kL0Size = 128; + uint32_t kL0Loops = (kL1Size + kL0Size - 1) / kL0Size; + uint32_t kL0SizeAlign = kL0Size; + LocalTensor bL1Tensor; + LocalTensor subvTensor; + + // ka表示左矩阵4buf选择哪一块buf, kb表示右矩阵3buf选择哪一块buf + uint32_t ka = 0, kb = 0; + uint32_t mBaseIdx = qpL1BufIter; + for (uint32_t nL1 = 0; nL1 < nL1Loops; nL1++) { // n切L1 + if (nL1 == (nL1Loops - 1)) { + // 尾块 + nL1Size = nSize - (nL1Loops - 1) * N_SPLIT_SIZE; + nL1SizeAlign = SFAAlign(nL1Size, 16U); + } + + // k l1写成一个循环, 和mm1保持一致 + kL1Size = 256; + kL1SizeAlign = SFAAlign(kL1Size, 16U); + + uint32_t curTopKIdx = info.curTopKIdx; + uint64_t curOffsetInSparseBlock = info.curOffsetInSparseBlock; + uint32_t copyRowCnt = 0; + int64_t idInTopK = topKGm.GetValue(info.topKBaseOffset + curTopKIdx); + + for (uint32_t k1 = 0; k1 < kL1Loops; k1++) { // k切L1, 这里套了一层l0来操作 + if (k1 == (kL1Loops - 1)) { + // 尾块 + kL1Size = kSize - (kL1Loops - 1) * 256; + kL1SizeAlign = SFAAlign(kL1Size, 16U); + } + kvL1BufIter++; + uint32_t kb = kvL1BufIter % 3; + WaitFlag(mte21KVIds[kb]); + bL1Tensor = l1KVTensor[kb * L1_BLOCK_OFFSET]; + uint32_t kOffset = k1 * kL0Loops; + kL0Size = 128; + // 此处必须先初始化kL0Size, 再求kL0Loops, 否则由于循环会改变kL0Size大小, 导致kL0Loops错误 + kL0Loops = (kL1Size + kL0Size - 1) / kL0Size; + kL0SizeAlign = kL0Size; + for (uint32_t kL1 = kOffset; kL1 < kL0Loops + kOffset; kL1++) { // 128 循环搬pa + if (kL1 == kOffset + kL0Loops - 1) { + // 尾块 + kL0Size = kL1Size - (kL0Loops - 1) * kL0Size; + kL0SizeAlign = SFAAlign(kL0Size, 16U); + } + + uint32_t curSeqIdx = info.s2BatchOffset + (kL1 - kOffset) * 128 + k1 * 256; + uint32_t copyFinishRowCnt = 0; + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = kL0Size; // 行数 + nd2nzPara.dValue = N_SPLIT_SIZE; // constInfo.headDim; + nd2nzPara.srcDValue = constInfo.headDim; + nd2nzPara.dstNzC0Stride = kL0SizeAlign; + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(bL1Tensor[(kL1 - kOffset) * 128 * N_SPLIT_SIZE], + kvMergeGm_[info.loop % 4 * N_WORKSPACE_SIZE * 576 + kL1 * 128 * constInfo.headDim + + nL1 * N_SPLIT_SIZE], + nd2nzPara); + } else { + while (copyFinishRowCnt < kL0Size) { + CalcTopKBlockInfo(info, curTopKIdx, curOffsetInSparseBlock, curSeqIdx, copyRowCnt, idInTopK); + + if (copyFinishRowCnt + copyRowCnt > kL0Size) { + copyRowCnt = kL0Size - copyFinishRowCnt; + } + + if constexpr (PAGE_ATTENTION) { + Position startPos; + startPos.bIdx = info.bIdx; + startPos.n2Idx = info.n2Idx; + startPos.s2Idx = idInTopK * constInfo.sparseBlockSize + curOffsetInSparseBlock; + startPos.dIdx = + nL1 * N_SPLIT_SIZE; // mm1 右矩阵 bn2s2d, d为k轴不切; mm2 右矩阵, s2为k轴, d轴切分 + PAShape shape; + shape.blockSize = kvCacheBlockSize; + shape.headNum = constInfo.kvHeadNum; + shape.headDim = constInfo.headDim; + shape.actHeadDim = nL1Size; + shape.maxblockNumPerBatch = maxBlockNumPerBatch; + shape.copyRowNum = copyRowCnt; + shape.copyRowNumAlign = kL0SizeAlign; + subvTensor = bL1Tensor[(kL1 - kOffset) * 128 * N_SPLIT_SIZE + copyFinishRowCnt * 16]; + DataCopyPA(subvTensor, valueGm, blockTableGm, shape, startPos); + } else { + uint64_t valueOffset = info.tensorBOffset; + if constexpr (KV_LAYOUT_T == SFA_LAYOUT::BSND || KV_LAYOUT_T == SFA_LAYOUT::TND) { + valueOffset += (idInTopK * constInfo.sparseBlockSize + curOffsetInSparseBlock) * + constInfo.kvHeadNum * constInfo.headDim; + } else { + valueOffset += (idInTopK * constInfo.sparseBlockSize + curOffsetInSparseBlock) * + constInfo.headDim; + } + + subvTensor = bL1Tensor[(kL1 - kOffset) * 128 * N_SPLIT_SIZE]; + CopyInMm2BToL1(subvTensor, valueOffset, kL0SizeAlign, copyFinishRowCnt, copyRowCnt, + nL1 * N_SPLIT_SIZE, nL1Size); + } + // 更新循环变量 + copyFinishRowCnt += copyRowCnt; + curSeqIdx += copyRowCnt; + } + } + } + SetFlag(mte21KVIds[kb]); + WaitFlag(mte21KVIds[kb]); + mL1SizeAlign = M_SPLIT_SIZE; + mL1Size = M_SPLIT_SIZE; // m的实际大小 + for (uint32_t mL1 = 0; mL1 < mL1Loops; mL1++) { + if (mL1 == (mL1Loops - 1)) { + // 尾块 + mL1Size = mSize - (mL1Loops - 1) * M_SPLIT_SIZE; + mL1SizeAlign = SFAAlign(mL1Size, 16U); + } + + uint32_t mIdx = mBaseIdx + mL1; + ka = GetQPL1RealIdx(mIdx, k1); + LocalTensor aL1Tensor = l1QPTensor[ka * L1_BLOCK_OFFSET]; + if (nL1 == 0) { + WaitFlag(mte21QPIds[ka]); + CopyInMm2AToL1(aL1Tensor, info, mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE, mL1Size, kL1Size, + 256 * k1); + SetFlag(mte21QPIds[ka]); + WaitFlag(mte21QPIds[ka]); + } + + LocalTensor cL0Tensor = + cL0TensorPingPong[(cL0BufIter % 2) * + (L0C_PP_SIZE / sizeof(MM_OUT_T))]; // 需要保证cL0BufIter和m步调一致 + uint32_t baseK = 128; + uint32_t baseN = 128; + kL0Size = 128; + kL0SizeAlign = kL0Size; + for (uint32_t kL0 = 0; kL0 < kL0Loops; kL0++) { + if (kL0 + 1 == kL0Loops) { + kL0Size = kL1Size - (kL0Loops - 1) * kL0Size; + kL0SizeAlign = SFAAlign(kL0Size, 16U); + } + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + LocalTensor bL0Tensor = bL0TensorPingPong[(abL0BufIter % 2) * (L0B_PP_SIZE / sizeof(KV_T))]; + LoadData3DParamsV2 loadData3DParamsForB; + loadData3DParamsForB.l1H = kL0SizeAlign / 16; // 源操作数height + loadData3DParamsForB.l1W = 16; // 源操作数weight=16，目的height=l1H*L1W + loadData3DParamsForB.padList[0] = 0; + loadData3DParamsForB.padList[1] = 0; + loadData3DParamsForB.padList[2] = 0; + loadData3DParamsForB.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + loadData3DParamsForB.mExtension = kL0SizeAlign; // 在目的操作数height维度的传输长度 + loadData3DParamsForB.kExtension = nL1SizeAlign; // 在目的操作数width维度的传输长度 + loadData3DParamsForB.mStartPt = 0; // 卷积核在目的操作数width维度的起点 + loadData3DParamsForB.kStartPt = 0; // 卷积核在目的操作数height维度的起点 + loadData3DParamsForB.strideW = 1; + loadData3DParamsForB.strideH = 1; + loadData3DParamsForB.filterW = 1; + loadData3DParamsForB.filterSizeW = false; // 是否在filterW的基础上将卷积核width增加256个元素 + loadData3DParamsForB.filterH = 1; + loadData3DParamsForB.filterSizeH = false; // 是否在filterH的基础上将卷积核height增加256个元素 + loadData3DParamsForB.dilationFilterW = 1; // 卷积核width膨胀系数 + loadData3DParamsForB.dilationFilterH = 1; // 卷积核height膨胀系数 + loadData3DParamsForB.enTranspose = 1; // 是否启用转置功能 + loadData3DParamsForB.fMatrixCtrl = 0; // 使用FMATRIX_LEFT还是使用FMATRIX_RIGHT，=0使用FMATRIX_LEFT，=1使用FMATRIX_RIGHT 1 + loadData3DParamsForB.channelSize = nL1SizeAlign; // 源操作数的通道数。膨胀系数为1时，目的weight为filterW*filterH*channelSize + LoadData(bL0Tensor, bL1Tensor[kL0 * baseK * baseN], loadData3DParamsForB); + + LocalTensor aL0Tensor = aL0TensorPingPong[(abL0BufIter % 2) * (L0A_PP_SIZE / sizeof(KV_T))]; + LoadData3DParamsV2 loadData3DParamsForA; + loadData3DParamsForA.l1H = mL1SizeAlign / 16; // 源操作数height + loadData3DParamsForA.l1W = 16; // 源操作数weight + loadData3DParamsForA.padList[0] = 0; + loadData3DParamsForA.padList[1] = 0; + loadData3DParamsForA.padList[2] = 0; + loadData3DParamsForA.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + loadData3DParamsForA.mExtension = mL1SizeAlign; // 在目的操作数height维度的传输长度 + loadData3DParamsForA.kExtension = kL0SizeAlign; // 在目的操作数width维度的传输长度 + loadData3DParamsForA.mStartPt = 0; // 卷积核在目的操作数width维度的起点 + loadData3DParamsForA.kStartPt = 0; // 卷积核在目的操作数height维度的起点 + loadData3DParamsForA.strideW = 1; // 卷积核在源操作数width维度滑动的步长 + loadData3DParamsForA.strideH = 1; // 卷积核在源操作数height维度滑动的步长 + loadData3DParamsForA.filterW = 1; // 卷积核width + loadData3DParamsForA.filterSizeW = false; // 是否在filterW的基础上将卷积核width增加256个元素 + loadData3DParamsForA.filterH = 1; // 卷积核height + loadData3DParamsForA.filterSizeH = false; // 是否在filterH的基础上将卷积核height增加256个元素 + loadData3DParamsForA.dilationFilterW = 1; // 卷积核width膨胀系数 + loadData3DParamsForA.dilationFilterH = 1; // 卷积核height膨胀系数 + loadData3DParamsForA.enTranspose = 0; // 是否启用转置功能，对整个目标矩阵进行转置 + loadData3DParamsForA.fMatrixCtrl = 0; + loadData3DParamsForA.channelSize = kL0SizeAlign; // 源操作数的通道数。膨胀系数为1时，目的weight为filterW*filterH*channelSize + LoadData(aL0Tensor, aL1Tensor[kL0 * baseK * mL1SizeAlign], + loadData3DParamsForA); + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + WaitFlag(Mte1MmABEventId(abL0BufIter % 2)); + + MmadParams mmadParams; + mmadParams.m = mL1SizeAlign; + mmadParams.n = nL1SizeAlign; + mmadParams.k = kL0Size; + mmadParams.cmatrixInitVal = (kL0 == 0 && k1 == 0); + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = ((k1 == (kL1Loops - 1)) && (kL0 == (kL0Loops - 1))) ? 0b11 : 0b10; + + Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } + SetFlag(Mte1MmABEventId(abL0BufIter % 2)); + abL0BufIter++; + } + + if (nL1 == (nL1Loops - 1)) { // nL1最后一轮, 需要将B驻留在L1中, 用于下一轮的计算？ + SetFlag(mte21QPIds[ka]); // 反向同步, 表示L1中的A已经被mte1消费完 + } + + if (k1 == (kL1Loops - 1)) { + if (nL1 == 0 && mL1 == 0) { // 第一次Fixpipe前等待 + CrossCoreWaitFlag(constInfo.syncV1NupdateC2); + } + + SetAtomicAdd(); + // ND + FixpipeParamsV220 fixParams; + fixParams.nSize = nL1SizeAlign; + fixParams.mSize = mL1SizeAlign; + fixParams.srcStride = mL1SizeAlign; + fixParams.dstStride = nSize; // mm2ResGm两行之间的间隔 + fixParams.ndNum = 1; // 输出ND + fixParams.unitFlag = 0b11; + + uint64_t mm2Offset = (mSplitInfo.nBufferStartM + mL1 * M_SPLIT_SIZE) * nSize + nL1 * N_SPLIT_SIZE; + Fixpipe(mm2ResGm[(info.bn2IdxInCurCore % (constInfo.preLoadNum)) * + constInfo.bmm2ResUbSize + mm2Offset], cL0Tensor, fixParams); + SetAtomicNone(); + } + + if (mL1Loops == 2) { + cL0BufIter++; + } + } + SetFlag(mte21KVIds[kb]); // 反向同步, 表示L1已经被mte1消费完 + } + // cL0BufIter已经不在使用 + if (mL1Loops == 1) { + cL0BufIter++; + } + } + qpL1BufIter += mL1Loops; +} + +#endif // SPARSE_FLASH_ATTENTION_SERVICE_CUBE_MLA_H diff --git a/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch22/sparse_flash_attention_service_vector_mla.h b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch22/sparse_flash_attention_service_vector_mla.h new file mode 100644 index 000000000..a15add665 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch22/sparse_flash_attention_service_vector_mla.h @@ -0,0 +1,1464 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_flash_attention_service_vector_mla.h + * \brief + */ +#ifndef SPARSE_FLASH_ATTENTION_SERVICE_VECTOR_MLA_H +#define SPARSE_FLASH_ATTENTION_SERVICE_VECTOR_MLA_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "../sparse_flash_attention_common.h" + +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +template class SFAVectorService { +public: + // 中间计算数据类型为float，高精度模式 + using T = float; + using KV_T = typename SFAT::kvType; + using OUT_T = typename SFAT::outputType; + using UPDATE_T = T; + using MM1_OUT_T = float; + using MM2_OUT_T = float; + + __aicore__ inline SFAVectorService(){}; + __aicore__ inline void ProcessVec1L(const RunInfo &info); + __aicore__ inline void ProcessVec2L(const RunInfo &info); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitParams(const struct ConstInfo &constInfo, + const SparseFlashAttentionTilingDataMla *__restrict tilingData); + __aicore__ inline void InitMm2ResInt32GmGlobalTensor(GlobalTensor mm2ResInt32Gm); + __aicore__ inline void InitVec0GlobalTensor(const GlobalTensor &kvValidSizeGm, + const GlobalTensor &kvMergeGm, + const GlobalTensor &keyRopeGm, const GlobalTensor &keyGm, + const GlobalTensor &blkTableGm); + __aicore__ inline void InitVec1GlobalTensor(GlobalTensor mm1ResGm, GlobalTensor vec1ResGm, + GlobalTensor actualSeqLengthsQGm, + GlobalTensor actualSeqLengthsKVGm, GlobalTensor lseMaxFdGm, + GlobalTensor lseSumFdGm, GlobalTensor topKGm, + GlobalTensor softmaxMaxGm, GlobalTensor softmaxSumGm); + __aicore__ inline void InitVec2GlobalTensor(GlobalTensor accumOutGm, GlobalTensor vec2ResGm, + GlobalTensor mm2ResGm, GlobalTensor attentionOutGm); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void InitSoftmaxDefaultBuffer(); + // ================================Base Vector========================================== + __aicore__ inline void RowDivs(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline void RowMuls(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + // ================================Vector0========================================== + __aicore__ inline void MergeKv(const RunInfo &runInfo); + __aicore__ inline int64_t GetKeyGmOffset(int64_t realS2Idx, const RunInfo &runInfo, int64_t s2IdLimit); + __aicore__ inline int64_t GetKeyRopeGmOffset(int64_t realS2Idx, const RunInfo &runInfo, int64_t s2IdLimit); + __aicore__ inline void GetRealS2Idx(int64_t s2GmOffset, int64_t &realS2Idx, int64_t topkGmBaseOffset, + const RunInfo &runInfo); + __aicore__ inline void CopyInKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, int64_t realS2Idx1, + int64_t realS2Idx2, const RunInfo &runInfo); + __aicore__ inline void CopyOutMrgeResult(int64_t mte2Size, int64_t mte3Size, int64_t s2StartGmOffset, + int64_t mergeMte3Idx, const RunInfo &runInfo); + __aicore__ inline void SetInfInBlk(const LocalTensor &mmResUb, uint32_t dealRowCount, uint32_t columnCount, + uint64_t startId, uint64_t endId); + __aicore__ inline void SetMidInf(const LocalTensor &mmResUb, uint32_t dealRowCount, uint32_t columnCount, + uint64_t startId, uint64_t endId); + __aicore__ inline void CopyInSingleKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, int64_t realS2Idx, + int64_t keyBNBOffset,int64_t s2IdLimit, const RunInfo &runInfo); + // ================================Vector1========================================== + __aicore__ inline void ProcessVec1SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo); + __aicore__ inline void DealBmm1ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t startRow, + uint32_t dealRowCount, uint32_t columnCount, uint32_t loopId); + __aicore__ inline void SoftmaxFlashV2Compute(const RunInfo &info, const MSplitInfo &mSplitInfo, + LocalTensor &mmResUb, LocalTensor &softmaxTmpUb, + uint32_t startRow, uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + __aicore__ inline void AmlaVecCompute(const RunInfo &info, const MSplitInfo &mSplitInfo, LocalTensor &mmResUb, + LocalTensor &softmaxTmpUb, uint32_t startRow, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline void ElewiseCompute(const RunInfo &info, const LocalTensor &mmResUb, uint32_t dealRowCount, + uint32_t columnCount); + __aicore__ inline void ProcessAmlaNupdate(const RunInfo &info, const MSplitInfo &mSplitInfo); + __aicore__ inline void ComputeLogSumExpAndCopyToGm(const RunInfo &info, const MSplitInfo &mSplitInfo, + LocalTensor &softmaxSumUb, LocalTensor &softmaxMaxUb); + __aicore__ inline void CopyFALseToGm(const RunInfo &info, const MSplitInfo &mSplitInfo, + LocalTensor &softmaxSumUb, LocalTensor &softmaxMaxUb); + __aicore__ inline void SetBmm2FirstSInnerBias(const RunInfo &info, const MSplitInfo &mSplitInfo); + // ================================Vecotr2========================================== + __aicore__ inline void ProcessVec2SingleBuf(const RunInfo &info, const MSplitInfo &mSplitInfo); + __aicore__ inline void DealBmm2ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t startRow, + uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + __aicore__ inline void ProcessVec2Inner(const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t mStartRow, + uint32_t mDealSize); + __aicore__ inline void Bmm2DataCopyOutTrans(const RunInfo &info, LocalTensor &attenOutUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount); + __aicore__ inline void Bmm2ResCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline void Bmm2CastAndCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline void Bmm2FDDataCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount); + __aicore__ inline uint64_t CalcAccumOffset(uint32_t bN2Idx, uint32_t gS1Idx); + __aicore__ inline void GetConfusionTransposeTiling(int64_t numR, int64_t numC, const uint32_t stackBufferSize, + const uint32_t typeSize, ConfusionTransposeTiling &tiling); + + // BLOCK和REPEAT的字节数 + static constexpr uint64_t BYTE_BLOCK = 32UL; + static constexpr uint32_t REPEAT_BLOCK_BYTE = 256U; + // BLOCK和REPEAT的FP32元素数 + static constexpr uint32_t FP32_BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(float); + static constexpr uint32_t FP32_REPEAT_ELEMENT_NUM = REPEAT_BLOCK_BYTE / sizeof(float); + // repeat stride不能超过256 + static constexpr uint32_t REPEATE_STRIDE_UP_BOUND = 256; + +private: + static constexpr bool PAGE_ATTENTION = SFAT::pageAttention; + static constexpr int TEMPLATE_MODE = SFAT::templateMode; + static constexpr bool FLASH_DECODE = SFAT::flashDecode; + static constexpr SFA_LAYOUT LAYOUT_T = SFAT::layout; + static constexpr SFA_LAYOUT KV_LAYOUT_T = SFAT::kvLayout; + + static constexpr uint64_t MERGE_CACHE_GM_BUF_NUM = 4; + static constexpr uint64_t SYNC_INPUT_BUF1_FLAG = 2; + static constexpr uint64_t SYNC_INPUT_BUF1_PONG_FLAG = 3; + static constexpr uint64_t SYNC_INPUT_BUF2_FLAG = 4; + static constexpr uint64_t SYNC_INPUT_BUF2_PONG_FLAG = 5; + static constexpr uint64_t SYNC_OUTPUT_BUF1_FLAG = 4; + static constexpr uint64_t SYNC_OUTPUT_BUF2_FLAG = 5; + static constexpr uint64_t SYNC_INPUT_V0BUF_FLAG = 6; + static constexpr uint32_t INPUT1_BUFFER_OFFSET = ConstInfo::BUFFER_SIZE_BYTE_32K; + static constexpr uint32_t SOFTMAX_TMP_BUFFER_OFFSET = ConstInfo::BUFFER_SIZE_BYTE_1K; + static constexpr uint32_t BASE_BLOCK_MAX_ELEMENT_NUM = ConstInfo::BUFFER_SIZE_BYTE_32K / sizeof(T); // 32768/4=8096 + static constexpr uint32_t BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(T); // 32/4=8 + static constexpr T FLOAT_E_SCALAR = 8388608; + static constexpr T LN2 = 0.6931471805599453094172; + static constexpr T RECIP_OF_LN2 = 1 / LN2; + static constexpr T SOFTMAX_MIN_NUM = -2e38; + + const SparseFlashAttentionTilingDataMla *__restrict tilingData; + + uint32_t pingpongFlag = 0U; + ConstInfo constInfo = {}; + + GlobalTensor mm2ResInt32Gm; + GlobalTensor mm1ResGm; + GlobalTensor vec1ResGm; + GlobalTensor lseSumFdGm; + GlobalTensor lseMaxFdGm; + GlobalTensor softmaxMaxGm; + GlobalTensor softmaxSumGm; + + GlobalTensor actualSeqLengthsQGm; + GlobalTensor actualSeqLengthsKVGm; + GlobalTensor vec2ResGm; + GlobalTensor mm2ResGm; + GlobalTensor accumOutGm; + GlobalTensor attentionOutGm; + GlobalTensor blkTableGm_; + + GlobalTensor kvMergeGm_; + GlobalTensor keyRopeGm_; + GlobalTensor keyGm_; + GlobalTensor topkGm_; + GlobalTensor kvValidSizeGm_; + + // ================================Local Buffer区==================================== + TBuf<> inputBuff1; // 32K + TBuf<> inputBuff2; // 16K + TBuf<> outputBuff1; // 32K + TBuf<> outputBuff2; // 4K + + TBuf<> tmpBuff1; // 32K + TBuf<> v0ValidSizeBuff; // 8K + + TBuf<> nValueBuff; + TBuf<> cofValueBuff; + TBuf<> aMlaSumBuff; + TBuf<> softmaxMaxBuff; // PRE_LOAD_NUM * 2K + TBuf<> softmaxExpBuff; // PRE_LOAD_NUM * 2K + TBuf<> softmaxSumBuff; // PRE_LOAD_NUM * 2K + TBuf<> softmaxMaxDefaultBuff; // 2K + TBuf<> softmaxSumDefaultBuff; // 2K + + LocalTensor softmaxMaxDefaultUb; + LocalTensor softmaxSumDefaultUb; + + LocalTensor nValueUb; + LocalTensor cofValueUb; + LocalTensor aMlaSumUb; + LocalTensor softmaxMaxUb; + LocalTensor softmaxSumUb; + LocalTensor softmaxExpUb; + LocalTensor kvMergUb_; + LocalTensor ropeMergUb_; + LocalTensor v0ValidSizeUb_; +}; + +template __aicore__ inline void SFAVectorService::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(inputBuff1, ConstInfo::BUFFER_SIZE_BYTE_32K * 2); // 2:pingpong + pipe->InitBuffer(inputBuff2, ConstInfo::BUFFER_SIZE_BYTE_8K * 2); // 2:pingpong + pipe->InitBuffer(outputBuff1, ConstInfo::BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(outputBuff2, ConstInfo::BUFFER_SIZE_BYTE_4K); + + pipe->InitBuffer(tmpBuff1, ConstInfo::BUFFER_SIZE_BYTE_32K); + pipe->InitBuffer(v0ValidSizeBuff, ConstInfo::BUFFER_SIZE_BYTE_8K); + + // M_MAX = 512/2vector = 256, 256 * sizeof(T) * N_Buffer + pipe->InitBuffer(nValueBuff, ConstInfo::BUFFER_SIZE_BYTE_1K * constInfo.preLoadNum); + pipe->InitBuffer(cofValueBuff, ConstInfo::BUFFER_SIZE_BYTE_1K * constInfo.preLoadNum); + pipe->InitBuffer(aMlaSumBuff, ConstInfo::BUFFER_SIZE_BYTE_1K * constInfo.preLoadNum); + + pipe->InitBuffer(softmaxMaxBuff, ConstInfo::BUFFER_SIZE_BYTE_1K * constInfo.preLoadNum); + pipe->InitBuffer(softmaxExpBuff, ConstInfo::BUFFER_SIZE_BYTE_1K * constInfo.preLoadNum); + pipe->InitBuffer(softmaxSumBuff, ConstInfo::BUFFER_SIZE_BYTE_1K * constInfo.preLoadNum); + + pipe->InitBuffer(softmaxMaxDefaultBuff, ConstInfo::BUFFER_SIZE_BYTE_1K); + pipe->InitBuffer(softmaxSumDefaultBuff, ConstInfo::BUFFER_SIZE_BYTE_1K); + + nValueUb = nValueBuff.Get(); + cofValueUb = cofValueBuff.Get(); + aMlaSumUb = aMlaSumBuff.Get(); + + softmaxMaxUb = softmaxMaxBuff.Get(); + softmaxSumUb = softmaxSumBuff.Get(); + softmaxExpUb = softmaxExpBuff.Get(); + + softmaxMaxDefaultUb = softmaxMaxDefaultBuff.Get(); + softmaxSumDefaultUb = softmaxSumDefaultBuff.Get(); + + kvMergUb_ = inputBuff1.Get(); + ropeMergUb_ = inputBuff2.Get(); + + v0ValidSizeUb_ = v0ValidSizeBuff.Get(); +} + +template +__aicore__ inline void +SFAVectorService::InitParams(const struct ConstInfo &constInfo, + const SparseFlashAttentionTilingDataMla *__restrict tilingData) +{ + this->constInfo = constInfo; + this->tilingData = tilingData; +} + +template +__aicore__ inline void +SFAVectorService::InitMm2ResInt32GmGlobalTensor(GlobalTensor mm2ResInt32Gm) +{ + this->mm2ResInt32Gm = mm2ResInt32Gm; +} + +template +__aicore__ inline void SFAVectorService::InitVec0GlobalTensor( + const GlobalTensor &kvValidSizeGm, const GlobalTensor &kvMergeGm, + const GlobalTensor &keyRopeGm, const GlobalTensor &keyGm, const GlobalTensor &blkTableGm) +{ + this->kvMergeGm_ = kvMergeGm; + this->keyRopeGm_ = keyRopeGm; + this->keyGm_ = keyGm; + this->blkTableGm_ = blkTableGm; + this->kvValidSizeGm_ = kvValidSizeGm; +} + +template +__aicore__ inline void SFAVectorService::InitVec1GlobalTensor( + GlobalTensor mm1ResGm, GlobalTensor vec1ResGm, + GlobalTensor actualSeqLengthsQGm, GlobalTensor actualSeqLengthsKVGm, GlobalTensor lseMaxFdGm, + GlobalTensor lseSumFdGm, GlobalTensor topKGm, GlobalTensor softmaxMaxGm, GlobalTensor softmaxSumGm) +{ + this->mm1ResGm = mm1ResGm; + this->vec1ResGm = vec1ResGm; + this->actualSeqLengthsQGm = actualSeqLengthsQGm; + this->actualSeqLengthsKVGm = actualSeqLengthsKVGm; + this->lseMaxFdGm = lseMaxFdGm; + this->lseSumFdGm = lseSumFdGm; + this->topkGm_ = topKGm; + this->softmaxMaxGm = softmaxMaxGm; + this->softmaxSumGm = softmaxSumGm; +} + +template +__aicore__ inline void SFAVectorService::InitVec2GlobalTensor(GlobalTensor accumOutGm, + GlobalTensor vec2ResGm, + GlobalTensor mm2ResGm, + GlobalTensor attentionOutGm) +{ + this->accumOutGm = accumOutGm; + this->vec2ResGm = vec2ResGm; + this->mm2ResGm = mm2ResGm; + this->attentionOutGm = attentionOutGm; +} + +template __aicore__ inline void SFAVectorService::AllocEventID() +{ + SetFlag(SYNC_INPUT_BUF1_FLAG); + SetFlag(SYNC_INPUT_BUF1_PONG_FLAG); + SetFlag(SYNC_INPUT_BUF2_FLAG); + SetFlag(SYNC_INPUT_BUF2_PONG_FLAG); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); +} + +template __aicore__ inline void SFAVectorService::FreeEventID() +{ + WaitFlag(SYNC_INPUT_BUF1_FLAG); + WaitFlag(SYNC_INPUT_BUF1_PONG_FLAG); + WaitFlag(SYNC_INPUT_BUF2_FLAG); + WaitFlag(SYNC_INPUT_BUF2_PONG_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); +} + +template __aicore__ inline void SFAVectorService::InitSoftmaxDefaultBuffer() +{ + Duplicate(softmaxMaxDefaultUb, SOFTMAX_MIN_NUM, SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)); + Duplicate(softmaxSumDefaultUb, ConstInfo::FLOAT_ZERO, SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)); +} + +template +__aicore__ inline void SFAVectorService::CopyFALseToGm(const RunInfo &info, const MSplitInfo &mSplitInfo, + LocalTensor &softmaxSumUb, LocalTensor &softmaxMaxUb) + +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + uint64_t baseOffset = mSplitInfo.nBufferStartM / 2; + size_t size = mSplitInfo.vecDealM; + + int64_t offset = 0; + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { //lse layout为N2 T G + uint64_t actualSeqQTotal = (info.bIdx <= 0) ? 0 : actualSeqLengthsQGm.GetValue(constInfo.batchSize - 1); + uint64_t actualSeqQPrefixSum = (info.bIdx <= 0) ? 0 : actualSeqLengthsQGm.GetValue(info.bIdx - 1); + offset += info.n2Idx * actualSeqQTotal * constInfo.gSize + + (actualSeqQPrefixSum + info.gS1Idx / constInfo.gSize) * constInfo.gSize + + mSplitInfo.nBufferStartM + mSplitInfo.vecStartM; + } else { + offset += info.bIdx * constInfo.kvHeadNum * constInfo.qSeqSize * constInfo.gSize + + info.n2Idx * constInfo.qSeqSize * constInfo.gSize + + info.gS1Idx / constInfo.gSize * constInfo.gSize + + mSplitInfo.nBufferStartM + mSplitInfo.vecStartM; + } + + if (info.actualSingleProcessSInnerSize != 0) { + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = 1; + dataCopyParams.blockLen = sizeof(T) * size; + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = 0; + size_t alignedSize = (sizeof(T) * size + 31) / 32 * 32 / sizeof(T); + LocalTensor tmp = outputBuff2.Get(); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + DataCopy(tmp, softmaxMaxUb[baseOffset], alignedSize); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + DataCopyPad(softmaxMaxGm[offset], tmp, dataCopyParams); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + + tmp = outputBuff2.Get(); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + DataCopy(tmp, softmaxSumUb[baseOffset], alignedSize); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + DataCopyPad(softmaxSumGm[offset], tmp, dataCopyParams); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + } else { + matmul::InitOutput(softmaxSumGm[offset], size, ConstInfo::FLOAT_ZERO); + matmul::InitOutput(softmaxMaxGm[offset], size, SOFTMAX_MIN_NUM); + } +} + +template +__aicore__ inline void SFAVectorService::ComputeLogSumExpAndCopyToGm(const RunInfo &info, + const MSplitInfo &mSplitInfo, + LocalTensor &softmaxSumUb, + LocalTensor &softmaxMaxUb) +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + uint64_t baseOffset = mSplitInfo.nBufferStartM / 2; + size_t size = mSplitInfo.vecDealM * FP32_BLOCK_ELEMENT_NUM; + uint64_t accumTmpOutNum = CalcAccumOffset(info.bIdx, info.gS1Idx); + uint64_t offset = (accumTmpOutNum * constInfo.kvHeadNum * constInfo.mBaseSize + // taskoffset + info.tndCoreStartKVSplitPos * constInfo.kvHeadNum * constInfo.mBaseSize + // 份数offset + mSplitInfo.nBufferStartM + mSplitInfo.vecStartM) * + FP32_BLOCK_ELEMENT_NUM; // m轴offset + if (info.actualSingleProcessSInnerSize != 0) { + LocalTensor tmp = outputBuff2.Get(); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + Brcb(tmp, softmaxSumUb[baseOffset], (mSplitInfo.vecDealM + 7) / 8, {1, 8}); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + DataCopy(lseSumFdGm[offset], tmp, size); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + + tmp = outputBuff2.Get(); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + Brcb(tmp, softmaxMaxUb[baseOffset], (mSplitInfo.vecDealM + 7) / 8, {1, 8}); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + DataCopy(lseMaxFdGm[offset], tmp, size); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + } else { + matmul::InitOutput(lseSumFdGm[offset], size, ConstInfo::FLOAT_ZERO); + matmul::InitOutput(lseMaxFdGm[offset], size, SOFTMAX_MIN_NUM); + } +} + +template +__aicore__ inline void SFAVectorService::ElewiseCompute(const RunInfo &info, + const LocalTensor &mmResUb, + uint32_t dealRowCount, uint32_t columnCount) +{ + Muls(mmResUb, mmResUb, static_cast(tilingData->baseParams.scaleValue), dealRowCount * columnCount); + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + // v0的无效值判断 + uint64_t s2ValidSizeFirstPart = v0ValidSizeUb_.GetValue(128 + info.loop % MERGE_CACHE_GM_BUF_NUM); + uint64_t s2ValidSizeSecondPart = v0ValidSizeUb_.GetValue(256 + info.loop % MERGE_CACHE_GM_BUF_NUM); + + int64_t s2ProcessSize = info.actualSingleProcessSInnerSize; + int64_t s2Pair = CeilDiv(s2ProcessSize, 2L * constInfo.sparseBlockSize); + int64_t s2Mid = CeilDiv(s2Pair, 2L) * 2 * constInfo.sparseBlockSize; + if (s2Mid > s2ProcessSize) { + s2Mid = s2ProcessSize; + } + if (unlikely(s2ValidSizeFirstPart < s2Mid)) { + int64_t s2StartCeilAlign = CeilAlign(s2ValidSizeFirstPart, 8); + int64_t s2MidFloorAlign = s2Mid / 8 * 8; + // 场景一 s2Mid > s2ValidSizeFirstPart + oneBlk + // 可以推导出s2StartCeilAlign < s2Mid 第一阶段取到s2StartCeilAlign + // s2StartCeilAlign <= s2MidFloorAlign 第二阶段取到s2MidFloorAlign + // 场景二 s2Mid <= s2ValidSizeFirstPart + oneBlk + // 可以推导出 s2StartCeilAlign >= s2Mid 第一阶段取到mid + // s2StartCeilAlign > s2MidFloorAlign 第二阶段取到s2StartCeilAlign + SetInfInBlk(mmResUb, dealRowCount, columnCount, s2ValidSizeFirstPart, + s2StartCeilAlign >= s2Mid ? s2Mid : s2StartCeilAlign); + SetMidInf(mmResUb, dealRowCount, columnCount, s2StartCeilAlign, s2MidFloorAlign); + SetInfInBlk(mmResUb, dealRowCount, columnCount, + s2StartCeilAlign <= s2MidFloorAlign ? s2MidFloorAlign : s2StartCeilAlign, s2Mid); + } + if (unlikely(s2ValidSizeSecondPart < s2ProcessSize - s2Mid)) { + // 场景一 s2Mid + s2ValidSizeSecondPart > s2ProcessSize + oneBlk + // 可以推导出 s2StartCeilAlign < s2ProcessSize 第一阶段取到s2StartCeilAlign + // s2StartCeilAlign <= s2EndFloorAlign 第二阶段取到s2EndFloorAlign + // 场景二 s2Mid + s2ValidSizeSecondPart <= s2ProcessSize + oneBlk + // 可以推导出 s2StartCeilAlign >= s2ProcessSize 第一阶段取到s2ProcessSize + // s2StartCeilAlign > s2EndFloorAlign 第二阶段取到s2StartCeilAlign + int64_t s2StartCeilAlign = CeilAlign(s2Mid + s2ValidSizeSecondPart, 8); + int64_t s2EndFloorAlign = s2ProcessSize / 8 * 8; + SetInfInBlk(mmResUb, dealRowCount, columnCount, s2Mid + s2ValidSizeSecondPart, + s2StartCeilAlign >= s2ProcessSize ? s2ProcessSize : s2StartCeilAlign); + SetMidInf(mmResUb, dealRowCount, columnCount, s2StartCeilAlign, s2EndFloorAlign); + SetInfInBlk(mmResUb, dealRowCount, columnCount, + s2StartCeilAlign <= s2EndFloorAlign ? s2EndFloorAlign : s2StartCeilAlign, s2ProcessSize); + } + } +} + +template +__aicore__ inline void SFAVectorService::SetInfInBlk(const LocalTensor &mmResUb, + uint32_t dealRowCount, uint32_t columnCount, + uint64_t startId, uint64_t endId) +{ + // startId endId + // x x x 0 0 0 x x x + // 从startId到endId部分置-inf, endId、startId为endId一个blk内部的下标 + if (startId >= endId) { + return; + } + + uint64_t startFloorAlignSize = startId / BLOCK_ELEMENT_NUM * BLOCK_ELEMENT_NUM; + uint64_t notComputePreMaskOneBlk = (1 << (startId - startFloorAlignSize)) - 1; + uint64_t notComputePostMaskOneBlk = ~((1 << (endId - startFloorAlignSize)) - 1); + uint64_t notComputeMaskOneBlk = notComputePreMaskOneBlk ^ notComputePostMaskOneBlk; + + uint64_t maskOneBlk = ~notComputeMaskOneBlk; + uint64_t mask[1] = {maskOneBlk}; + for (int i = 1; i < 8; i++) { + mask[0] = mask[0] | (maskOneBlk << (i * 8)); + } + for (uint64_t rowId = 0; rowId < dealRowCount; rowId += 8) { + Duplicate(mmResUb[rowId * columnCount + startFloorAlignSize], SOFTMAX_MIN_NUM, mask, + 1, CeilDiv(columnCount, 8), 0); + } +} + +template +__aicore__ inline void SFAVectorService::SetMidInf(const LocalTensor &mmResUb, + uint32_t dealRowCount, uint32_t columnCount, + uint64_t startId, uint64_t endId) +{ + if (startId >= endId) { + return; + } + // startId endId + // 0 ... 0 + // 从startId到endId部分置-inf, startId、endId为32B对齐的下标 + for (uint64_t rowId = 0; rowId < dealRowCount; rowId++) { + Duplicate(mmResUb[rowId * columnCount + startId], SOFTMAX_MIN_NUM, endId - startId); + } +} + +template +__aicore__ inline void SFAVectorService::SoftmaxFlashV2Compute( + const RunInfo &info, const MSplitInfo &mSplitInfo, LocalTensor &mmResUb, LocalTensor &softmaxTmpUb, + uint32_t startRow, uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount) +{ + LocalTensor inSumTensor; + LocalTensor inMaxTensor; + uint32_t baseOffset = mSplitInfo.nBufferStartM / 2 + startRow; + uint32_t outIdx = info.loop % (constInfo.preLoadNum); + uint32_t softmaxOutOffset = outIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset; + if (info.isFirstSInnerLoop) { + inMaxTensor = softmaxMaxDefaultUb; + inSumTensor = softmaxSumDefaultUb; + } else { + uint32_t inIdx = (info.loop - 1) % (constInfo.preLoadNum); + inMaxTensor = softmaxMaxUb[inIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset]; + inSumTensor = softmaxSumUb[inIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset]; + } + if (actualColumnCount !=0) { + SoftMaxShapeInfo srcShape{dealRowCount, columnCount, dealRowCount, actualColumnCount}; + SoftMaxTiling newTiling = + SoftMaxFlashV2TilingFunc(srcShape, sizeof(T), sizeof(T), softmaxTmpUb.GetSize(), true, false); + SoftmaxFlashV2( + mmResUb, softmaxSumUb[softmaxOutOffset], softmaxMaxUb[softmaxOutOffset], mmResUb, + softmaxExpUb[softmaxOutOffset], inSumTensor, inMaxTensor, softmaxTmpUb, newTiling, srcShape); + } else { + uint32_t dealRowCountAlign = SFAAlign(dealRowCount, FP32_BLOCK_ELEMENT_NUM); + DataCopy(softmaxSumUb[softmaxOutOffset], inSumTensor, dealRowCountAlign); + PipeBarrier(); + DataCopy(softmaxMaxUb[softmaxOutOffset], inMaxTensor, dealRowCountAlign); + } +} + +template +__aicore__ inline void SFAVectorService::AmlaVecCompute( + const RunInfo &info, const MSplitInfo &mSplitInfo, LocalTensor &mmResUb, LocalTensor &softmaxTmpUb, + uint32_t startRow, uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount) +{ + uint32_t baseOffset = mSplitInfo.nBufferStartM / 2 + startRow; + uint32_t calCount = dealRowCount; + uint32_t outIdx = info.loop % (constInfo.preLoadNum); + uint32_t softmaxOutOffset = outIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset; + // compute n(i) + LocalTensor nTmp = softmaxTmpUb.template ReinterpretCast(); + LocalTensor nUpdateTmp = nTmp[SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)]; + Muls(nTmp, softmaxMaxUb[softmaxOutOffset], ((T)(-1.0)) * RECIP_OF_LN2, calCount); + + PipeBarrier(); + Cast(nTmp, nTmp, RoundMode::CAST_ROUND, calCount); + PipeBarrier(); + + uint32_t prOutIdx = (info.loop - 1) % (constInfo.preLoadNum); + uint32_t PreSoftmaxOutOffset = prOutIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset; + // n(i) - n(i-1) + if (info.isFirstSInnerLoop) { + Duplicate(nUpdateTmp, ConstInfo::FLOAT_ZERO, calCount); // n1=n0 + } else { + Sub(nUpdateTmp, nTmp, nValueUb[PreSoftmaxOutOffset], calCount); + } + PipeBarrier(); + // update n(i), DataCopy not support when calCount is not align 32B, so use Adds + Adds(nValueUb[softmaxOutOffset], nTmp, ConstInfo::FLOAT_ZERO, calCount); + PipeBarrier(); + + // update softmax res + LocalTensor nUpdateTmp2 = nTmp[2 * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)]; + LocalTensor nTmp_KvT = nTmp[3 * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)].template ReinterpretCast(); + LocalTensor tmpCofUb = nTmp[4 * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)]; + LocalTensor epsUb = nTmp[5 * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)]; + Muls(nUpdateTmp2, softmaxMaxUb[softmaxOutOffset], RECIP_OF_LN2, calCount); + PipeBarrier(); + Add(nTmp, nUpdateTmp2, nTmp, calCount); + PipeBarrier(); + Muls(nTmp, nTmp, LN2, calCount); + PipeBarrier(); + Exp(nTmp, nTmp, calCount); + PipeBarrier(); + Cast(nTmp_KvT, nTmp, RoundMode::CAST_ROUND, calCount); // fp32->fp16/bf16 + PipeBarrier(); + Cast(nUpdateTmp2, nTmp_KvT, RoundMode::CAST_NONE, calCount); // fp16/bf16->fp32 + PipeBarrier(); + if (info.s2Idx + 1 == info.curSInnerLoopTimes) { + Mul(aMlaSumUb[softmaxOutOffset], softmaxSumUb[softmaxOutOffset], nUpdateTmp2, calCount); + } + if (actualColumnCount == 0) { + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + return; + } + LocalTensor nTmp3 = nTmp[6 * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)]; + Brcb(nTmp3, nUpdateTmp2, (dealRowCount + 7) / 8, {1, 8}); + PipeBarrier(); + RowMuls(mmResUb, mmResUb, nTmp3, dealRowCount, columnCount, actualColumnCount); + + Div(tmpCofUb, nTmp, nUpdateTmp2, calCount); // cof(i)=tmpS32/tmpS16 + if (info.isFirstSInnerLoop) { + Duplicate(cofValueUb[softmaxOutOffset], (T)1.0, calCount); // cof_0=1 + PipeBarrier(); + Div(epsUb, cofValueUb[softmaxOutOffset], tmpCofUb, calCount); // 1 / cof(i) + } else { + PipeBarrier(); + Div(epsUb, cofValueUb[PreSoftmaxOutOffset], tmpCofUb, calCount); // cof(i - 1) / cof(i) + } + PipeBarrier(); + + Adds(cofValueUb[softmaxOutOffset], tmpCofUb, ConstInfo::FLOAT_ZERO, calCount); // store cof(i) + Adds(epsUb, epsUb, (T)(-1.0), calCount); // cof(i - 1) / cof(i) - 1 + PipeBarrier(); + Muls(epsUb, epsUb, (T)1.5, calCount); // (cof(i - 1) - cof(i)) / cof(i) * 1.5 + + Maxs(nUpdateTmp, nUpdateTmp, (T)(-30.0), calCount); // N = max(n(i) - n(i-1), -30) + PipeBarrier(); + Adds(epsUb, epsUb, (T)(0.000001), calCount); + PipeBarrier(); + Add(nUpdateTmp, nUpdateTmp, epsUb, calCount); + PipeBarrier(); + Muls(nUpdateTmp, nUpdateTmp, FLOAT_E_SCALAR, calCount); // N = N * pow(2, 23) + PipeBarrier(); + + // nUpdate int32 out + LocalTensor tmQue = outputBuff2.Get(); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + LocalTensor nInt32Out = tmQue[startRow]; // 缓存nUpdate + + Cast(nInt32Out, nUpdateTmp, RoundMode::CAST_ROUND, dealRowCount); + PipeBarrier(); + + SetFlag(SYNC_OUTPUT_BUF2_FLAG); +} + +template +__aicore__ inline void SFAVectorService::DealBmm1ResBaseBlock( + const RunInfo &info, const MSplitInfo &mSplitInfo, uint32_t startRow, uint32_t dealRowCount, + uint32_t columnCount, uint32_t loopId) +{ + uint32_t computeSize = dealRowCount * columnCount; + uint64_t inOutGmOffset = (info.loop % constInfo.preLoadNum) * constInfo.mmResUbSize + + (mSplitInfo.nBufferStartM + mSplitInfo.vecStartM + startRow) * columnCount; + LocalTensor mmResUb = inputBuff1.Get(); + mmResUb = mmResUb[pingpongFlag * INPUT1_BUFFER_OFFSET / sizeof(MM1_OUT_T)]; + WaitFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + + DataCopy(mmResUb, mm1ResGm[inOutGmOffset], computeSize); + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + if (loopId == 0) { + WaitFlag(0); + } + } + SetFlag(SYNC_INPUT_BUF1_FLAG); + WaitFlag(SYNC_INPUT_BUF1_FLAG); + + ElewiseCompute(info, mmResUb, dealRowCount, columnCount); + + PipeBarrier(); + LocalTensor tmpAFloorUb = tmpBuff1.Get(); + LocalTensor softmaxTmpUb = tmpAFloorUb.template ReinterpretCast(); + + SoftmaxFlashV2Compute(info, mSplitInfo, mmResUb, softmaxTmpUb, startRow, dealRowCount, columnCount, + info.actualSingleProcessSInnerSize); + + PipeBarrier(); + AmlaVecCompute(info, mSplitInfo, mmResUb, softmaxTmpUb, startRow, dealRowCount, columnCount, + info.actualSingleProcessSInnerSize); + + PipeBarrier(); + LocalTensor tmpMMResCastTensor = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + + Cast(tmpMMResCastTensor, mmResUb, AscendC::RoundMode::CAST_ROUND, computeSize); + SetFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + DataCopy(vec1ResGm[inOutGmOffset], tmpMMResCastTensor, computeSize); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); +} + +template +__aicore__ inline void SFAVectorService::SetBmm2FirstSInnerBias(const RunInfo &info, const MSplitInfo &mSplitInfo) +{ + uint32_t mSplitSize = 16U; + uint64_t baseoffset = (info.bn2IdxInCurCore % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + + (mSplitInfo.nBufferStartM + mSplitInfo.vecStartM) * constInfo.headDim; + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + LocalTensor tmpTensor = outputBuff1.Get(); + Duplicate(tmpTensor, static_cast(394264576), mSplitSize * constInfo.headDim); // 394264576 : fp32下2^(-80)的二进制表示对应的int32数值 + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + uint32_t loopCount = (mSplitInfo.vecDealM + mSplitSize - 1) / mSplitSize; + for (uint32_t loop = 0; loop < loopCount; loop++) { + DataCopy(mm2ResInt32Gm[baseoffset + loop * mSplitSize * constInfo.headDim], tmpTensor, mSplitSize * constInfo.headDim); + } + SetFlag(SYNC_OUTPUT_BUF1_FLAG); +} + +template +__aicore__ inline void SFAVectorService::ProcessAmlaNupdate(const RunInfo &info, const MSplitInfo &mSplitInfo) +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + if (info.isFirstSInnerLoop) { + SetBmm2FirstSInnerBias(info, mSplitInfo); + return; + } + + LocalTensor nUpdateTensor = outputBuff2.Get(); // shape:1/2*s1*g + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + SetFlag(SYNC_OUTPUT_BUF2_FLAG); + WaitFlag(SYNC_OUTPUT_BUF2_FLAG); + + constexpr uint32_t dGroupSize = 128U; + constexpr uint32_t mSplitSize = 64U; // tmpQue size 32KB，一次只能处理64个N，最大保存的数据大小：64*128*sizeof(int32) + constexpr uint32_t ONE_BLOCK_SIZE = 32U; // 32B + + uint32_t subMSize = SFAAlign(mSplitInfo.vecDealM, 16U); + uint16_t elementPerBlock = ONE_BLOCK_SIZE / sizeof(int32_t); // 单个datablock的元素数，int32_t类型的为32/4=8 + uint32_t loopCount = (subMSize + mSplitSize - 1) / mSplitSize; + uint32_t tailSplitSize = subMSize - (loopCount - 1) * mSplitSize; // 尾块 + + for (uint32_t loop = 0, processMSize = mSplitSize; loop < loopCount; loop++) { + if (loop == (loopCount - 1)) { + processMSize = tailSplitSize; + } + LocalTensor tmpQue = outputBuff1.Get(); + + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + // (m,1)单次brcb扩充成(m,8), 重复16次, 扩充为(m,128) + for (uint32_t i = 0; i < dGroupSize / elementPerBlock; i++) { + Brcb(tmpQue[i * elementPerBlock], + nUpdateTensor[loop * mSplitSize], + static_cast((processMSize + elementPerBlock - 1) / elementPerBlock), + {static_cast(dGroupSize / elementPerBlock), // 单次迭代内，目的操作数不同datablock间地址步长,单位为datablock + static_cast(dGroupSize)}); // 相邻迭代间，目的操作数相同datablock地址步长 + } + + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + + uint64_t baseoffset = (info.bn2IdxInCurCore % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + + (mSplitInfo.nBufferStartM + mSplitInfo.vecStartM + loop * mSplitSize) * constInfo.headDim; + + SetAtomicAdd(); + DataCopyParams dataCopyParams; + dataCopyParams.blockCount = static_cast(processMSize); + dataCopyParams.blockLen = dGroupSize * sizeof(int32_t) / ONE_BLOCK_SIZE; // 每个block是128个元素，单位为32B + dataCopyParams.srcStride = 0; // 前面一个数据块的尾与后面数据块的头的间隔 + dataCopyParams.dstStride = static_cast((constInfo.headDim - dGroupSize) * + sizeof(int32_t) / ONE_BLOCK_SIZE); // 单位为32B + for (uint32_t i = 0; i < constInfo.headDim / dGroupSize; i++) { // 4=512/128 + DataCopy(mm2ResInt32Gm[baseoffset + i * dGroupSize] ,tmpQue, dataCopyParams); + } + SetAtomicNone(); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + } + SetFlag(SYNC_OUTPUT_BUF2_FLAG); +} + +template +__aicore__ inline void SFAVectorService::ProcessVec1SingleBuf(const RunInfo &info, + const MSplitInfo &mSplitInfo) +{ + if (mSplitInfo.vecDealM == 0) { + return; + } + uint32_t mSplitSize = info.actualSingleProcessSInnerSize == 0 ? + 16 : BASE_BLOCK_MAX_ELEMENT_NUM / info.actualSingleProcessSInnerSizeAlign; + // 1. 向下8对齐是因为UB操作至少32B + // 2. info.actualSingleProcessSInnerSizeAlign最大512, mSplitSize可以确保最小为16 + mSplitSize = mSplitSize / 8 * 8; + + if (mSplitSize > mSplitInfo.vecDealM) { + mSplitSize = mSplitInfo.vecDealM; + } + uint32_t loopCount = (mSplitInfo.vecDealM + mSplitSize - 1) / mSplitSize; + uint32_t tailSplitSize = mSplitInfo.vecDealM - (loopCount - 1) * mSplitSize; + + if constexpr (TEMPLATE_MODE == V_TEMPLATE) { + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = 1; + dataCopyParams.blockLen = 256 * sizeof(int32_t); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = 0; + DataCopyPadExtParams padParams; + // 额外偏移128个元素，避免不同loop下v0和v1互相影响 + DataCopyPad(v0ValidSizeUb_[128], kvValidSizeGm_[info.loop % MERGE_CACHE_GM_BUF_NUM * (128 * 2)], + dataCopyParams, padParams); + SetFlag(0); + if (unlikely(loopCount == 0)) { + // scalar同步影响较大，挪到循环内部进行 + WaitFlag(0); + } + } + for (uint32_t i = 0, dealSize = mSplitSize; i < loopCount; i++) { + if (i == (loopCount - 1)) { + dealSize = tailSplitSize; + } + DealBmm1ResBaseBlock(info, mSplitInfo, i * mSplitSize, dealSize, info.actualSingleProcessSInnerSizeAlign, i); + pingpongFlag ^= 1; // pingpong 0 1切换 + } +} + +template +__aicore__ inline void SFAVectorService::GetRealS2Idx(int64_t s2GmOffset, int64_t &realS2Idx, + int64_t topkGmBaseOffset, const RunInfo &runInfo) +{ + int64_t topkGmIdx = (s2GmOffset + runInfo.s2Idx * constInfo.s2BaseSize) / constInfo.sparseBlockSize; + if (unlikely(topkGmIdx >= constInfo.sparseBlockCount)) { + realS2Idx = -1; + return; + } + realS2Idx = topkGm_.GetValue(topkGmBaseOffset + topkGmIdx) * static_cast(constInfo.sparseBlockSize) + + static_cast((s2GmOffset + runInfo.s2Idx * constInfo.s2BaseSize) % constInfo.sparseBlockSize); +} + +template +__aicore__ inline int64_t SFAVectorService::GetKeyGmOffset(int64_t realS2Idx, + const RunInfo &runInfo, int64_t s2IdLimit) +{ + if (realS2Idx < 0 || realS2Idx >= s2IdLimit) { + return -1; + } + int64_t realKeyGmOffset = 0; + if constexpr (PAGE_ATTENTION) { + int64_t blkTableIdx = realS2Idx / constInfo.kvCacheBlockSize; + int64_t blkTableOffset = realS2Idx % constInfo.kvCacheBlockSize; + realKeyGmOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo.maxBlockNumPerBatch + blkTableIdx) * + static_cast(constInfo.kvCacheBlockSize) * + static_cast(constInfo.kvHeadNum) + + blkTableOffset; + } else { + realKeyGmOffset = (runInfo.tensorBOffset + + realS2Idx * constInfo.kvHeadNum * constInfo.headDim) / + constInfo.headDim; + } + return realKeyGmOffset; +} + +template +__aicore__ inline int64_t SFAVectorService::GetKeyRopeGmOffset(int64_t realS2Idx, + const RunInfo &runInfo, int64_t s2IdLimit) +{ + if (realS2Idx < 0 || realS2Idx >= s2IdLimit) { + return -1; + } + int64_t realKeyRopeGmOffset = 0; + realKeyRopeGmOffset = (runInfo.tensorBRopeOffset + + realS2Idx * constInfo.kvHeadNum * constInfo.headDimRope) / + constInfo.headDimRope; + return realKeyRopeGmOffset; +} + +template +__aicore__ inline void +SFAVectorService::CopyInSingleKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, int64_t realS2Idx, + int64_t keyBNBOffset,int64_t s2IdLimit, const RunInfo &runInfo) +{ + if (keyBNBOffset < 0) { + return; + } + int64_t validS2Count = + (realS2Idx + constInfo.sparseBlockSize > s2IdLimit ? s2IdLimit - realS2Idx : constInfo.sparseBlockSize); + DataCopyExtParams intriParams; + intriParams.blockLen = validS2Count * constInfo.headDim * sizeof(KV_T); + intriParams.blockCount = 1; + intriParams.dstStride = 0; + intriParams.srcStride = 0; + DataCopyPadExtParams padParams; + DataCopyPad(kvMergUb_[mergeMte3Idx % 2 * 32 * 512 + (mte2Size - mte3Size) * constInfo.headDim], + keyGm_[keyBNBOffset * constInfo.headDim], intriParams, padParams); + intriParams.blockLen = validS2Count * constInfo.headDimRope * sizeof(KV_T); + + DataCopyPad(ropeMergUb_[mergeMte3Idx % 2 * 32 * 64 + (mte2Size - mte3Size) * constInfo.headDimRope], + keyRopeGm_[keyBNBOffset * constInfo.headDimRope], intriParams, padParams); + mte2Size += validS2Count; +} + +template +__aicore__ inline void SFAVectorService::CopyInKv(int64_t &mte2Size, int64_t mte3Size, int64_t mergeMte3Idx, + int64_t realS2Idx1, int64_t realS2Idx2, const RunInfo &runInfo) +{ + int64_t s2IdLimit = runInfo.curActualSeqLenOri; + if (constInfo.sparseMode == 3) { + s2IdLimit = runInfo.curActualSeqLenOri - runInfo.actS1Size + runInfo.gS1Idx / constInfo.gSize + 1; + } + + int64_t keyOffset1 = GetKeyGmOffset(realS2Idx1, runInfo, s2IdLimit); + int64_t keyOffset2 = GetKeyGmOffset(realS2Idx2, runInfo, s2IdLimit); + if (unlikely(keyOffset1 < 0 && keyOffset2 < 0)) { + return; + } + + int64_t keySrcStride = 0; + int64_t keyRopeSrcStride = 0; + if constexpr (PAGE_ATTENTION) { + int64_t blkTableSrcStride = + ((keyOffset1 > keyOffset2 ? (keyOffset1 - keyOffset2) : + (keyOffset2 - keyOffset1)) - constInfo.sparseBlockSize); + keySrcStride = blkTableSrcStride * constInfo.headDim * sizeof(KV_T); + keyRopeSrcStride = blkTableSrcStride * constInfo.headDimRope * sizeof(KV_T); + } else { + int64_t keyRopeOffset1 = GetKeyRopeGmOffset(realS2Idx1, runInfo, s2IdLimit); + int64_t keyRopeOffset2 = GetKeyRopeGmOffset(realS2Idx2, runInfo, s2IdLimit); + keySrcStride = ((keyOffset1 > keyOffset2 ? (keyOffset1 - keyOffset2) : + (keyOffset2 - keyOffset1)) - constInfo.sparseBlockSize) * constInfo.headDim * sizeof(KV_T); + keyRopeSrcStride = ((keyRopeOffset1 > keyRopeOffset2 ? (keyRopeOffset1 - keyRopeOffset2) : + (keyRopeOffset2 - keyRopeOffset1)) - constInfo.sparseBlockSize) * + constInfo.headDimRope * sizeof(KV_T); + } + + if (unlikely(keySrcStride >= INT32_MAX || keySrcStride < 0 || + (!PAGE_ATTENTION && (keyRopeSrcStride >= INT32_MAX || keyRopeSrcStride < 0)) || + realS2Idx1 + constInfo.sparseBlockSize >= s2IdLimit || + realS2Idx2 + constInfo.sparseBlockSize >= s2IdLimit)) { + // stride溢出、stride为负数、s2超长等异常场景，还原成2条搬运指令 + CopyInSingleKv(mte2Size, mte3Size, mergeMte3Idx, realS2Idx1, keyOffset1, s2IdLimit, runInfo); + CopyInSingleKv(mte2Size, mte3Size, mergeMte3Idx, realS2Idx2, keyOffset2, s2IdLimit, runInfo); + } else { + DataCopyExtParams intriParams; + intriParams.blockLen = constInfo.sparseBlockSize * constInfo.headDim * sizeof(KV_T); + intriParams.blockCount = (keyOffset1 >= 0) + (keyOffset2 >= 0); + intriParams.dstStride = 0; + intriParams.srcStride = keySrcStride; + DataCopyPadExtParams padParams; + + int64_t startGmOffset = keyOffset1 > -1 ? keyOffset1 : keyOffset2; + if (keyOffset2 > -1 && keyOffset2 < keyOffset1) { + startGmOffset = keyOffset2; + } + DataCopyPad(kvMergUb_[mergeMte3Idx % 2 * 32 * 512 + (mte2Size - mte3Size) * constInfo.headDim], + keyGm_[startGmOffset * constInfo.headDim], intriParams, padParams); + + intriParams.blockLen = constInfo.sparseBlockSize * constInfo.headDimRope * sizeof(KV_T); + intriParams.dstStride = 0; + intriParams.srcStride = keyRopeSrcStride; + DataCopyPad(ropeMergUb_[mergeMte3Idx % 2 * 32 * 64 + (mte2Size - mte3Size) * constInfo.headDimRope], + keyRopeGm_[startGmOffset * constInfo.headDimRope], intriParams, padParams); + mte2Size += ((keyOffset1 > -1) + (keyOffset2 > -1)) * constInfo.sparseBlockSize; + } +} + +template +__aicore__ inline void SFAVectorService::CopyOutMrgeResult(int64_t mte2Size, int64_t mte3Size, + int64_t s2GmStartOffset, int64_t mergeMte3Idx, + const RunInfo &runInfo) +{ + if (mte2Size <= mte3Size) { + return; + } + SetFlag(0); + WaitFlag(0); + + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = mte2Size - mte3Size; + dataCopyParams.blockLen = constInfo.headDim * sizeof(KV_T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = 0; + + DataCopyPad(kvMergeGm_[runInfo.loop % 4 * 512 * 576 + (s2GmStartOffset + mte3Size)*constInfo.headDim], + kvMergUb_[mergeMte3Idx % 2 * 32 * 512], dataCopyParams); + + dataCopyParams.blockLen = constInfo.headDimRope * sizeof(KV_T); + DataCopyPad(kvMergeGm_[runInfo.loop % 4 * 512 * 576 + 512 * 512 + (s2GmStartOffset + mte3Size) * + constInfo.headDimRope], ropeMergUb_[mergeMte3Idx % 2 * 32 * 64], dataCopyParams); +} + +// b s1 k +template +__aicore__ inline void SFAVectorService::MergeKv(const RunInfo &runInfo) +{ + int64_t s2ProcessSize = runInfo.actualSingleProcessSInnerSize; + int64_t s2Pair = CeilDiv(s2ProcessSize, 2L * constInfo.sparseBlockSize); + int64_t topkGmBaseOffset = 0; + + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + uint64_t actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsQGm.GetValue(runInfo.bIdx - 1); + topkGmBaseOffset += (actualSeqQPrefixSum + runInfo.gS1Idx / constInfo.gSize) * constInfo.kvHeadNum * + constInfo.sparseBlockCount + runInfo.n2Idx * constInfo.sparseBlockCount; + } else { + topkGmBaseOffset += runInfo.bIdx * constInfo.qSeqSize * constInfo.sparseBlockCount + + runInfo.gS1Idx / constInfo.gSize * constInfo.sparseBlockCount; + } + int64_t mergeMte3Idx = 0; + int64_t mte2Size = 0; + int64_t mte3Size = 0; + int64_t s2IdxArray0 = -1; + int64_t s2IdxArray1 = -1; + bool needWaitMte3ToMte2 = true; + SetFlag(0); + SetFlag(1); + int64_t s2GmStartOffset = GetSubBlockIdx() == 0 ? 0 : CeilDiv(s2Pair, 2L) * 2 * constInfo.sparseBlockSize; + int64_t s2GmLimit = GetSubBlockIdx() == 0 ? CeilDiv(s2Pair, 2L) * 2 * constInfo.sparseBlockSize: s2ProcessSize; + if (s2GmLimit > s2ProcessSize) { + s2GmLimit = s2ProcessSize; + } + for (int64_t s2GmOffsetArray = s2GmStartOffset; s2GmOffsetArray < s2GmLimit; s2GmOffsetArray += 2 * constInfo.sparseBlockSize) { + if (needWaitMte3ToMte2) { + WaitFlag(mergeMte3Idx % 2); + needWaitMte3ToMte2 = false; + } + GetRealS2Idx(s2GmOffsetArray, s2IdxArray0, topkGmBaseOffset, runInfo); + if (unlikely(s2IdxArray0 < 0)) { + CopyOutMrgeResult(mte2Size, mte3Size, s2GmStartOffset, mergeMte3Idx, runInfo); + SetFlag(mergeMte3Idx % 2); + mergeMte3Idx++; + break; + } + GetRealS2Idx(s2GmOffsetArray + constInfo.sparseBlockSize, s2IdxArray1, topkGmBaseOffset, runInfo); + CopyInKv(mte2Size, mte3Size, mergeMte3Idx, s2IdxArray0, s2IdxArray1, runInfo); + if ((mte2Size - mte3Size + 2 * constInfo.sparseBlockSize > 32) || + s2GmOffsetArray + 2 * constInfo.sparseBlockSize >= s2GmLimit) { + CopyOutMrgeResult(mte2Size, mte3Size, s2GmStartOffset, mergeMte3Idx, runInfo); + mte3Size = mte2Size; + SetFlag(mergeMte3Idx % 2); + mergeMte3Idx++; + needWaitMte3ToMte2 = true; + } + } + + if (unlikely(s2GmStartOffset + mte2Size < s2GmLimit)) { + SetFlag(0); + WaitFlag(0); + WaitFlag(mergeMte3Idx & 1); + Duplicate(kvMergUb_, static_cast(0.0), constInfo.headDim); + SetFlag(0); + WaitFlag(0); + + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = 1; + dataCopyParams.blockLen = constInfo.headDim * sizeof(KV_T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = 0; + for (int64_t s2GmOffset = s2GmStartOffset + mte2Size; s2GmOffset < s2GmLimit; s2GmOffset++) { + DataCopyPad(kvMergeGm_[runInfo.loop % MERGE_CACHE_GM_BUF_NUM * 512 * 576 + s2GmOffset * constInfo.headDim], + kvMergUb_, dataCopyParams); + } + dataCopyParams.blockLen = constInfo.headDimRope * sizeof(KV_T); + for (int64_t s2GmOffset = s2GmStartOffset + mte2Size; s2GmOffset < s2GmLimit; s2GmOffset++) { + DataCopyPad(kvMergeGm_[runInfo.loop % MERGE_CACHE_GM_BUF_NUM * 512 * 576 + 512 * constInfo.headDim + + s2GmOffset * constInfo.headDimRope], + kvMergUb_, dataCopyParams); + } + SetFlag(mergeMte3Idx & 1); + mergeMte3Idx++; + } + WaitFlag(0); + WaitFlag(1); + v0ValidSizeUb_.SetValue(runInfo.loop % MERGE_CACHE_GM_BUF_NUM, mte2Size); + SetFlag(1); + WaitFlag(1); + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = 1; + dataCopyParams.blockLen = 128 * sizeof(int32_t); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = 0; + DataCopyPad(kvValidSizeGm_[runInfo.loop % MERGE_CACHE_GM_BUF_NUM * (128 * 2) + GetSubBlockIdx() * 128], + v0ValidSizeUb_, dataCopyParams); + SetFlag(SYNC_INPUT_V0BUF_FLAG); + WaitFlag(SYNC_INPUT_V0BUF_FLAG); + return; +} + +template +__aicore__ inline void SFAVectorService::ProcessVec1L(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferIdx = i; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + + mSplitInfo.vecDealM = (mSplitInfo.nBufferDealM <= 16) ? mSplitInfo.nBufferDealM : + (((mSplitInfo.nBufferDealM + 15) / 16 + 1) / 2 * 16); + mSplitInfo.vecStartM = 0; + if (GetBlockIdx() % 2 == 1) { + mSplitInfo.vecStartM = mSplitInfo.vecDealM; + mSplitInfo.vecDealM = mSplitInfo.nBufferDealM - mSplitInfo.vecDealM; + } + + CrossCoreWaitFlag(constInfo.syncC1V1); + // vec1 compute + ProcessVec1SingleBuf(info, mSplitInfo); + CrossCoreSetFlag(constInfo.syncV1C2); + CrossCoreWaitFlag(constInfo.syncC2V1); + // add nUpdate to mm2ResGm + if (info.actualSingleProcessSInnerSize != 0) { + ProcessAmlaNupdate(info, mSplitInfo); + CrossCoreSetFlag(constInfo.syncV1NupdateC2); + } + // move lse for flash decode or FA + if (info.s2Idx == info.curSInnerLoopTimes - 1 && (constInfo.returnSoftmaxLse || info.tndIsS2SplitCore)) { + uint32_t outIdx = info.loop % (constInfo.preLoadNum); + auto sumTensor = softmaxSumUb[outIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)]; + auto maxTensor = softmaxMaxUb[outIdx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T)]; + if (constInfo.returnSoftmaxLse) { + CopyFALseToGm(info, mSplitInfo, sumTensor, maxTensor); + } + if (info.tndIsS2SplitCore) { + if constexpr (FLASH_DECODE) { + + ComputeLogSumExpAndCopyToGm(info, mSplitInfo, sumTensor, maxTensor); + } + } + } + } +} + +template +__aicore__ inline uint64_t SFAVectorService::CalcAccumOffset(uint32_t bN2Idx, uint32_t gS1Idx) +{ + return 0; +} + +template +__aicore__ inline void SFAVectorService::ProcessVec2SingleBuf(const RunInfo &info, + const MSplitInfo &mSplitInfo) +{ + if (info.s2Idx + 1 != info.curSInnerLoopTimes) { + return; + } + if (mSplitInfo.vecDealM == 0) { + return; + } + + ProcessVec2Inner(info, mSplitInfo, 0, mSplitInfo.vecDealM); +} + +template __aicore__ inline void SFAVectorService::ProcessVec2L(const RunInfo &info) +{ + uint32_t nBufferLoopTimes = (info.actMBaseSize + constInfo.nBufferMBaseSize - 1) / constInfo.nBufferMBaseSize; + uint32_t nBufferTail = info.actMBaseSize - (nBufferLoopTimes - 1) * constInfo.nBufferMBaseSize; + for (uint32_t i = 0; i < nBufferLoopTimes; i++) { + MSplitInfo mSplitInfo; + mSplitInfo.nBufferIdx = i; + mSplitInfo.nBufferStartM = i * constInfo.nBufferMBaseSize; + mSplitInfo.nBufferDealM = (i + 1 != nBufferLoopTimes) ? constInfo.nBufferMBaseSize : nBufferTail; + + mSplitInfo.vecDealM = (mSplitInfo.nBufferDealM <= 16) ? mSplitInfo.nBufferDealM : + (((mSplitInfo.nBufferDealM + 15) / 16 + 1) / 2 * 16); + mSplitInfo.vecStartM = 0; + if (GetBlockIdx() % 2 == 1) { + mSplitInfo.vecStartM = mSplitInfo.vecDealM; + mSplitInfo.vecDealM = mSplitInfo.nBufferDealM - mSplitInfo.vecDealM; + } + CrossCoreWaitFlag(constInfo.syncC2V2); + ProcessVec2SingleBuf(info, mSplitInfo); + } +} + +template +__aicore__ inline void SFAVectorService::ProcessVec2Inner(const RunInfo &info, + const MSplitInfo &mSplitInfo, + uint32_t mStartRow, uint32_t mDealSize) +{ + uint32_t mSplitSize = BASE_BLOCK_MAX_ELEMENT_NUM / constInfo.headDim; + if (mSplitSize > mDealSize) { + mSplitSize = mDealSize; + } + + uint32_t loopCount = (mDealSize + mSplitSize - 1) / mSplitSize; + uint32_t tailSplitSize = mDealSize - (loopCount - 1) * mSplitSize; + for (uint32_t i = 0, dealSize = mSplitSize; i < loopCount; i++) { + if (i == (loopCount - 1)) { + dealSize = tailSplitSize; + } + DealBmm2ResBaseBlock(info, mSplitInfo, i * mSplitSize + mStartRow, dealSize, + constInfo.headDim, constInfo.headDim); + pingpongFlag ^= 1; // pingpong 0 1切换 + } +} + + +template +__aicore__ inline void SFAVectorService::GetConfusionTransposeTiling( + int64_t numR, int64_t numC, const uint32_t stackBufferSize, const uint32_t typeSize, + ConfusionTransposeTiling &tiling) +{ + (void)stackBufferSize; + uint32_t blockSize = ONE_BLK_SIZE / typeSize; + uint32_t height = numC; + uint32_t width = numR; + uint32_t highBlock = height / BLOCK_CUBE; + uint32_t stride = height * blockSize * typeSize / ONE_BLK_SIZE; + uint32_t repeat = width / blockSize; + + tiling.param0 = blockSize; + tiling.param1 = height; + tiling.param2 = width; + tiling.param3 = highBlock; + tiling.param4 = stride; + tiling.param5 = repeat; +} + +template +__aicore__ inline void +SFAVectorService::Bmm2FDDataCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, + uint32_t wsMStart, uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount) +{ + LocalTensor tmp = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + DataCopy(tmp, bmm2ResUb, columnCount * dealRowCount); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + uint64_t accumTmpOutNum = CalcAccumOffset(info.bIdx, info.gS1Idx); + uint64_t offset = accumTmpOutNum * constInfo.kvHeadNum * constInfo.mBaseSize * constInfo.headDim + // taskoffset + info.tndCoreStartKVSplitPos * constInfo.kvHeadNum * constInfo.mBaseSize * constInfo.headDim + // 份数offset + wsMStart * actualColumnCount; // m轴offset + GlobalTensor dst = accumOutGm[offset]; + if (info.actualSingleProcessSInnerSize== 0) { + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = dealRowCount; + dataCopyParams.blockLen = actualColumnCount * sizeof(T); + dataCopyParams.srcStride = (columnCount - actualColumnCount) / (BYTE_BLOCK / sizeof(T)); + dataCopyParams.dstStride = 0; + DataCopyPad(dst, tmp, dataCopyParams); + } else { + matmul::InitOutput(dst, dealRowCount * actualColumnCount, ConstInfo::FLOAT_ZERO); + } + SetFlag(SYNC_OUTPUT_BUF1_FLAG); +} + +template +__aicore__ inline void +SFAVectorService::Bmm2DataCopyOutTrans(const RunInfo &info, LocalTensor &attenOutUb, + uint32_t wsMStart, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = dealRowCount; + dataCopyParams.blockLen = actualColumnCount * sizeof(OUT_T); + dataCopyParams.srcStride = (columnCount - actualColumnCount) / (BYTE_BLOCK / sizeof(OUT_T)); + dataCopyParams.dstStride = 0; + DataCopyPad(attentionOutGm[info.attenOutOffset + wsMStart * actualColumnCount], attenOutUb, dataCopyParams); + return; +} + +template +__aicore__ inline void +SFAVectorService::Bmm2CastAndCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, + uint32_t wsMStart, uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount) +{ + LocalTensor tmpBmm2ResCastTensor = outputBuff1.Get(); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + if constexpr (IsSameType::value) { // bf16 采取四舍六入五成双模式 + Cast(tmpBmm2ResCastTensor, bmm2ResUb, AscendC::RoundMode::CAST_RINT, dealRowCount * columnCount); + } else { + Cast(tmpBmm2ResCastTensor, bmm2ResUb, AscendC::RoundMode::CAST_ROUND, dealRowCount * columnCount); + } + + SetFlag(SYNC_OUTPUT_BUF1_FLAG); + WaitFlag(SYNC_OUTPUT_BUF1_FLAG); + Bmm2DataCopyOutTrans(info, tmpBmm2ResCastTensor, wsMStart, dealRowCount, columnCount, actualColumnCount); + SetFlag(SYNC_OUTPUT_BUF1_FLAG); +} + +template +__aicore__ inline void +SFAVectorService::Bmm2ResCopyOut(const RunInfo &info, LocalTensor &bmm2ResUb, uint32_t wsMStart, + uint32_t dealRowCount, uint32_t columnCount, + uint32_t actualColumnCount) +{ + if constexpr (FLASH_DECODE) { + if (info.tndIsS2SplitCore) { + Bmm2FDDataCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount); + } else { + Bmm2CastAndCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount); + } + } else { + Bmm2CastAndCopyOut(info, bmm2ResUb, wsMStart, dealRowCount, columnCount, actualColumnCount); + } +} + +template +__aicore__ inline void +SFAVectorService::DealBmm2ResBaseBlock(const RunInfo &info, const MSplitInfo &mSplitInfo, + uint32_t startRow, uint32_t dealRowCount, + uint32_t columnCount, uint32_t actualColumnCount) +{ + uint32_t vec2ComputeSize = dealRowCount * columnCount; + uint32_t mStart = mSplitInfo.nBufferStartM + mSplitInfo.vecStartM + startRow; + uint64_t srcGmOffset = (info.bn2IdxInCurCore % constInfo.preLoadNum) * constInfo.bmm2ResUbSize + + mStart * columnCount; + LocalTensor tmpBmm2ResUb = inputBuff1.Get(); + tmpBmm2ResUb = tmpBmm2ResUb[pingpongFlag * INPUT1_BUFFER_OFFSET / sizeof(MM2_OUT_T)]; + WaitFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + DataCopy(tmpBmm2ResUb, mm2ResGm[srcGmOffset], vec2ComputeSize); + + SetFlag(SYNC_INPUT_BUF1_FLAG); + WaitFlag(SYNC_INPUT_BUF1_FLAG); + + // 将绝对值大于1e10的数置为0 + LocalTensor bmm2ResUb = tmpBuff1.Get(); + bmm2ResUb.SetSize(vec2ComputeSize); + LocalTensor absBmm2ResUb = bmm2ResUb.template ReinterpretCast(); + Abs(absBmm2ResUb, tmpBmm2ResUb, vec2ComputeSize); + PipeBarrier(); + LocalTensor cmpMaskUb = absBmm2ResUb.template ReinterpretCast(); + CompareScalar(cmpMaskUb, absBmm2ResUb, (T)1e10, CMPMODE::LE, vec2ComputeSize); + PipeBarrier(); + Select(tmpBmm2ResUb, cmpMaskUb, tmpBmm2ResUb, ConstInfo::FLOAT_ZERO, + SELMODE::VSEL_TENSOR_SCALAR_MODE, vec2ComputeSize); + PipeBarrier(); + uint32_t baseOffset = mSplitInfo.nBufferStartM / 2 + startRow; + uint32_t idx = info.loop % (constInfo.preLoadNum); + LocalTensor tmpSumUb = v0ValidSizeBuff.Get()[384]; // sumUb用临时内存 16 * 32B = 512B + Brcb(tmpSumUb, aMlaSumUb[idx * SOFTMAX_TMP_BUFFER_OFFSET / sizeof(T) + baseOffset], (dealRowCount + 7) / 8, {1, 8}); + PipeBarrier(); + RowDivs(bmm2ResUb, tmpBmm2ResUb, tmpSumUb, dealRowCount, columnCount, actualColumnCount); + PipeBarrier(); + SetFlag(SYNC_INPUT_BUF1_FLAG + pingpongFlag); + Bmm2ResCopyOut(info, bmm2ResUb, mStart, dealRowCount, columnCount, actualColumnCount); +} + +template +__aicore__ inline void +SFAVectorService::RowDivs(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount) +{ + // divs by row, 每行的元素除以相同的元素 + // dstUb[i, (j * 8) : (j * 8 + 7)] = src0Ub[i, (j * 8) : (j * 8 + 7)] / src1Ub[i, 0 : 7] + // src0Ub:[dealRowCount, columnCount], src1Ub:[dealRowCount, FP32_BLOCK_ELEMENT_NUM] dstUb:[dealRowCount, + // columnCount] + uint32_t dtypeMask = FP32_REPEAT_ELEMENT_NUM; + uint32_t dLoop = actualColumnCount / dtypeMask; + uint32_t dRemain = actualColumnCount % dtypeMask; + + BinaryRepeatParams repeatParamsDiv; + repeatParamsDiv.src0BlkStride = 1; + repeatParamsDiv.src1BlkStride = 0; + repeatParamsDiv.dstBlkStride = 1; + repeatParamsDiv.src0RepStride = columnCount / FP32_BLOCK_ELEMENT_NUM; + repeatParamsDiv.src1RepStride = 1; + repeatParamsDiv.dstRepStride = columnCount / FP32_BLOCK_ELEMENT_NUM; + uint32_t columnRepeatCount = dLoop; + if (columnRepeatCount <= dealRowCount) { + uint32_t offset = 0; + for (uint32_t i = 0; i < dLoop; i++) { + Div(dstUb[offset], src0Ub[offset], src1Ub, dtypeMask, dealRowCount, repeatParamsDiv); + offset += dtypeMask; + } + } else { + BinaryRepeatParams columnRepeatParams; + columnRepeatParams.src0BlkStride = 1; + columnRepeatParams.src1BlkStride = 0; + columnRepeatParams.dstBlkStride = 1; + columnRepeatParams.src0RepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + columnRepeatParams.src1RepStride = 0; + columnRepeatParams.dstRepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + uint32_t offset = 0; + for (uint32_t i = 0; i < dealRowCount; i++) { + Div(dstUb[offset], src0Ub[offset], src1Ub[i * FP32_BLOCK_ELEMENT_NUM], dtypeMask, columnRepeatCount, + columnRepeatParams); + offset += columnCount; + } + } + if (dRemain > 0) { + Div(dstUb[dLoop * dtypeMask], src0Ub[dLoop * dtypeMask], src1Ub, dRemain, dealRowCount, repeatParamsDiv); + } +} + +template +__aicore__ inline void +SFAVectorService::RowMuls(LocalTensor dstUb, LocalTensor src0Ub, LocalTensor src1Ub, + uint32_t dealRowCount, uint32_t columnCount, uint32_t actualColumnCount) +{ + // muls by row, 每行的元素乘以相同的元素 + // dstUb[i, (j * 8) : (j * 8 + 7)] = src0Ub[i, (j * 8) : (j * 8 + 7)] * src1Ub[i, 0 : 7] + // src0Ub:[dealRowCount, columnCount] src1Ub:[dealRowCount, FP32_BLOCK_ELEMENT_NUM] dstUb:[dealRowCount, + // columnCount] + // dealRowCount is repeat times, must be less 256 + uint32_t repeatElementNum = FP32_REPEAT_ELEMENT_NUM; + uint32_t blockElementNum = FP32_BLOCK_ELEMENT_NUM; + + if constexpr (std::is_same::value) { + // 此限制由于每个repeat至多连续读取256B数据 + repeatElementNum = FP32_REPEAT_ELEMENT_NUM * 2; // 256/4 * 2=128 + blockElementNum = FP32_BLOCK_ELEMENT_NUM * 2; // 32/4 * 2 = 16 + } + + // 每次只能连续读取256B的数据进行计算，故每次只能处理256B/sizeof(dType)= + // 列方向分dLoop次，每次处理8列数据 + uint32_t dLoop = actualColumnCount / repeatElementNum; + uint32_t dRemain = actualColumnCount % repeatElementNum; + // REPEATE_STRIDE_UP_BOUND=256， 此限制由于src0RepStride数据类型为uint8之多256个datablock间距 + if (columnCount < REPEATE_STRIDE_UP_BOUND * blockElementNum) { + BinaryRepeatParams repeatParams; + repeatParams.src0BlkStride = 1; + repeatParams.src1BlkStride = 0; + repeatParams.dstBlkStride = 1; + repeatParams.src0RepStride = columnCount / blockElementNum; + repeatParams.src1RepStride = 1; + repeatParams.dstRepStride = columnCount / blockElementNum; + + // 如果以列为repeat所处理的次数小于行处理次数，则以列方式处理。反之则以行进行repeat处理 + if (dLoop <= dealRowCount) { + uint32_t offset = 0; + for (uint32_t i = 0; i < dLoop; i++) { + Mul(dstUb[offset], src0Ub[offset], src1Ub, repeatElementNum, dealRowCount, repeatParams); + offset += repeatElementNum; + } + } else { + BinaryRepeatParams columnRepeatParams; + columnRepeatParams.src0BlkStride = 1; + columnRepeatParams.src1BlkStride = 0; + columnRepeatParams.dstBlkStride = 1; + columnRepeatParams.src0RepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + columnRepeatParams.src1RepStride = 0; + columnRepeatParams.dstRepStride = 8; // 列方向上两次repeat起始地址间隔dtypeMask=64个元素，即8个block + for (uint32_t i = 0; i < dealRowCount; i++) { + Mul(dstUb[i * columnCount], src0Ub[i * columnCount], src1Ub[i * blockElementNum], repeatElementNum, + dLoop, columnRepeatParams); + } + } + + // 最后一次完成[dealRowCount, dRemain] * [dealRowCount, blockElementNum] 只计算有效部分 + if (dRemain > 0) { + Mul(dstUb[dLoop * repeatElementNum], src0Ub[dLoop * repeatElementNum], src1Ub, dRemain, dealRowCount, + repeatParams); + } + } else { + BinaryRepeatParams repeatParams; + repeatParams.src0RepStride = 8; // 每个repeat为256B数据，正好8个datablock + repeatParams.src0BlkStride = 1; + repeatParams.src1RepStride = 0; + repeatParams.src1BlkStride = 0; + repeatParams.dstRepStride = 8; + repeatParams.dstBlkStride = 1; + // 每次计算一行，共计算dealRowCount行 + for (uint32_t i = 0; i < dealRowCount; i++) { + // 计算一行中的dLoop个repeat, 每个repeat计算256/block_size 个data_block + Mul(dstUb[i * columnCount], src0Ub[i * columnCount], src1Ub[i * blockElementNum], repeatElementNum, dLoop, + repeatParams); + // 计算一行中的尾块 + if (dRemain > 0) { + Mul(dstUb[i * columnCount + dLoop * repeatElementNum], + src0Ub[i * columnCount + dLoop * repeatElementNum], src1Ub[i * blockElementNum], dRemain, 1, + repeatParams); + } + } + } +} + +#endif // SPARSE_FLASH_ATTENTION_SERVICE_VECTOR_MLA_H diff --git a/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_common_arch35.h b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_common_arch35.h new file mode 100644 index 000000000..8793dc30c --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_common_arch35.h @@ -0,0 +1,112 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_flash_attention_common_arch35.h + * \brief + */ +#ifndef SPARSE_FLASH_ATTENTION_COMMON_ARCH35_H +#define SPARSE_FLASH_ATTENTION_COMMON_ARCH35_H +#include +#include "kernel_tiling/kernel_tiling.h" + +constexpr uint64_t BLOCK_BYTE = 32; +constexpr uint32_t NEGATIVE_MIN_VALUE_FP32 = 0xFF7FFFFF; + +constexpr uint32_t L0AB_SHARED_SIZE_64K = 65536; // 65536表示64*1024 +constexpr uint32_t L0C_SHARED_SIZE_256K = 262144; // 262144表示256 * 1024 + +constexpr uint32_t BUFFER_SIZE_16K = 16384; // 16384表示16 * 1024 +constexpr uint32_t BUFFER_SIZE_32K = 32768; // 32768表示32 * 1024 +constexpr uint32_t BUFFER_SIZE_128K = 131072; // 131072表示128 * 1024 + +constexpr uint32_t CV_RATIO = 2; +constexpr uint64_t SYNC_MODE = 4; + +static constexpr uint32_t SFA_SYNC_MODE0 = 0; + +enum class SFA_LAYOUT { + BSND = 0, + TND = 1, + PA_BSND = 2, +}; + +enum class SFATemplateMode { + SWA_TEMPLATE_MODE = 0, + CFA_TEMPLATE_MODE = 1, + SCFA_TEMPLATE_MODE = 2 +}; + +namespace BaseApi { +__aicore__ constexpr uint64_t Align2Func(uint64_t data) +{ + return (data + 1UL) >> 1UL << 1UL; // 向上2对齐, +1移位2 +} + +__aicore__ constexpr uint64_t Align8Func(uint64_t data) +{ + return (data + 7UL) >> 3UL << 3UL; // 向上8对齐, +7移位3 +} + +__aicore__ constexpr uint64_t Align16Func(uint64_t data) +{ + return (data + 15UL) >> 4UL << 4UL; // 向上16对齐, +15移位4 +} + +__aicore__ constexpr uint64_t Align64Func(uint64_t data) +{ + return (data + 63UL) >> 6UL << 6UL; // 向上64对齐, +63移位6 +} +} + +#define TEMPLATE_INTF \ + template + +#define TEMPLATE_INTF_ARGS \ + Q_T, KV_T, T, OUTPUT_T, isFd, isPa, LAYOUT_T, KV_LAYOUT_T, TEMPLATE_MODE, IS_SPLIT_G + +#define CUBE_BLOCK_TRAITS_TYPE_FIELDS(X) \ + X(Q_T) \ + X(KV_T) \ + X(T) \ + X(OUTPUT_T) \ + +#define CUBE_BLOCK_TRAITS_CONST_FIELDS(X) \ + X(isFd, bool, false) \ + X(isPa, bool, true) \ + X(LAYOUT_T, SFA_LAYOUT, SFA_LAYOUT::BSND) \ + X(KV_LAYOUT_T, SFA_LAYOUT, SFA_LAYOUT::PA_BSND) \ + X(TEMPLATE_MODE, SFATemplateMode, SFATemplateMode::SCFA_TEMPLATE_MODE) \ + X(IS_SPLIT_G, bool, false) + + +/* 1. 生成带默认值的模版Template */ +#define GEN_TYPE_PARAM(name) typename name, +#define GEN_CONST_PARAM(name, type, default_val) type name = default_val, + +#define TEMPLATES_DEF \ +template + +/* 2. 生成不带默认值的模版Template */ +#define GEN_TEMPLATE_TYPE_NODEF(name) typename name, +#define GEN_TEMPLATE_CONST_NODEF(name, type, default_val) type name, +#define TEMPLATES_DEF_NO_DEFAULT \ +template + +/* 3. 生成有默认值的Args */ +#define GEN_ARG_NAME(name, ...) name, +#define TEMPLATE_ARGS \ + CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_ARG_NAME) \ + CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_ARG_NAME) end + +#endif // SPARSE_FLASH_ATTENTION_COMMON_ARCH35_H diff --git a/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_kernel_mla.h b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_kernel_mla.h new file mode 100644 index 000000000..51501c6cc --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_kernel_mla.h @@ -0,0 +1,717 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_flash_attention_kernel_mla.h + * \brief + */ + +#ifndef SPARSE_FLASH_ATTENTION_KERNEL_MLA_H +#define SPARSE_FLASH_ATTENTION_KERNEL_MLA_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "sparse_flash_attention_service_cube_mla.h" +#include "sparse_flash_attention_service_vector_mla.h" +#include "sparse_flash_attention_common_arch35.h" +#include "sparse_flash_attention_kvcache.h" + +#if __has_include("../../common/op_kernel/matmul.h") +#include "../../common/op_kernel/matmul.h" +#else +#include "../common/matmul.h" +#endif +#if __has_include("../../common/op_kernel/FixpipeOut.h") +#include "../../common/op_kernel/FixpipeOut.h" +#else +#include "../common/FixpipeOut.h" +#endif +#if __has_include("../../common/op_kernel/CopyInL1.h") +#include "../../common/op_kernel/CopyInL1.h" +#else +#include "../common/CopyInL1.h" +#endif + +using matmul::MatmulType; +using namespace AscendC; +using namespace AscendC::Impl::Detail; +using namespace regbaseutil; + +namespace BaseApi { +template class SparseFlashAttentionKernelMla { +public: + ARGS_TRAITS; + + __aicore__ inline SparseFlashAttentionKernelMla(){}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengths, __gm__ uint8_t *blockTable, + __gm__ uint8_t *queryRope, __gm__ uint8_t *keyRope, __gm__ uint8_t *attentionOut, + __gm__ uint8_t *softmaxMax, __gm__ uint8_t *softmaxSum, __gm__ uint8_t *workspace, + const SparseFlashAttentionTilingDataMla *__restrict tiling, + __gm__ uint8_t *gmTiling, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void ProcessMainLoop(); + __aicore__ inline void InitGlobalBuffer(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *queryRope, __gm__ uint8_t *keyRope, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *blockTable, __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths, + __gm__ uint8_t *softmaxMax, __gm__ uint8_t *softmaxSum, __gm__ uint8_t *workspace, + const SparseFlashAttentionTilingDataMla *__restrict tiling, + TPipe *tPipe); + __aicore__ inline void InitLocalBuffer(); + __aicore__ inline void InitMMResBuf(__gm__ uint8_t *workspace); + __aicore__ inline void ComputeConstexpr(); + __aicore__ inline void SetRunInfo(RunInfo &runInfo, RunParamStr &runParam, int64_t taskId, int64_t s2LoopCount, + int64_t s2LoopLimit, int64_t multiCoreInnerIdx); + __aicore__ inline void ComputeBmm1Tail(RunInfo &runInfo, RunParamStr &runParam); + __aicore__ inline void InitUniqueConstInfo(); + __aicore__ inline void ComputeAxisIdxByBnAndGs1(int64_t bnIndex, int64_t gS1Index, RunParamStr &runParam); + __aicore__ inline void InitUniqueRunInfo(const RunParamStr &runParam, RunInfo &runInfo); + + __aicore__ inline void InitCalcParamsEach(); + __aicore__ inline uint64_t GetBalanceActualSeqLengths(GlobalTensor &actualSeqLengths, uint32_t bIdx); + __aicore__ inline void GetAxisStartIdx(uint32_t bN2EndPrev, uint32_t s1GEndPrev, uint32_t s2EndPrev); + + TPipe *pipe; + + const SparseFlashAttentionTilingDataMla *__restrict tilingData; + static constexpr uint64_t SYNC_MODE = 4; + static constexpr uint32_t PRELOAD_NUM = 2; + /* 核间通道 */ + BufferManager gmBufferManager; + + BufferManager ubBufferManager; + BuffersPolicyDB bmm1Buffers; + BuffersPolicySingleBuffer bmm2Buffers; + + // mm2左矩阵P + BufferManager l1BufferManager; + BuffersPolicy3buff l1RightBuffers; + CVSharedParams sharedParams; + /* GM信息 */ + __gm__ int32_t *actualSeqKvlenAddr = nullptr; + __gm__ int32_t *actualSeqQlenAddr = nullptr; + + GlobalTensor actualSeqLengthsQGm; + uint32_t usedCoreNum = 0U; + + GlobalTensor oriTopkLengthGm; + bool hasOriTopkLength = false; + /* workspace 空间 */ + BuffersPolicy3buff v0ResGmBuffers; + /* 核Index信息 */ + int32_t aicIdx; + + /* 切G时最大s2Loop */ + int64_t maxS2LoopCnt; + + /* 初始化后不变的信息 */ + ConstInfo constInfo; + + /* 模板库Block */ + CubeBlockType cubeBlock; + VecBlockType vecBlock; + + uint32_t crossCoreSyncBufId = 0; +}; + +template + __aicore__ inline void SparseFlashAttentionKernelMla::Init( + __gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengths, __gm__ uint8_t *blockTable, + __gm__ uint8_t *queryRope, __gm__ uint8_t *keyRope, + __gm__ uint8_t *attentionOut, __gm__ uint8_t *softmaxMax, __gm__ uint8_t *softmaxSum, + __gm__ uint8_t *workspace, const SparseFlashAttentionTilingDataMla *__restrict tiling, + __gm__ uint8_t *gmTiling, TPipe *tPipe) +{ + fa_base_matmul::idCounterNum = 0; + constInfo.subBlockIdx = GetSubBlockIdx(); + if ASCEND_IS_AIC { + this->aicIdx = GetBlockIdx(); + constInfo.aivIdx = 0; + } else { + constInfo.aivIdx = GetBlockIdx(); + this->aicIdx = constInfo.aivIdx >> 1; + this->tilingData = tiling; + } + + constInfo.s1BaseSize = 64; + constInfo.s2BaseSize = 128; + + this->pipe = tPipe; + vecBlock.InitVecBlock(tPipe, this->tilingData, this->sharedParams, \ + this->aicIdx, constInfo.subBlockIdx, actualSeqLengthsQ, actualSeqLengths); + if ASCEND_IS_AIV { + constInfo.bSize = this->sharedParams.bSize; + constInfo.gSize = this->sharedParams.gSize; + constInfo.s1Size = this->sharedParams.s1Size; + constInfo.dSizeV = 512; + constInfo.needInit = this->sharedParams.needInit; + constInfo.returnSoftmaxLse = this->sharedParams.returnSoftmaxLse; + } + vecBlock.CleanOutput(attentionOut, softmaxMax, softmaxSum, constInfo); + /* cube侧不依赖sharedParams的scalar前置 */ + InitMMResBuf(workspace); + if ASCEND_IS_AIC { + cubeBlock.InitCubeBlock(pipe, l1BufferManager, query, queryRope); + /* wait kfc message */ + CrossCoreWaitFlag(15); + auto tempTilingSSbuf = reinterpret_cast<__ssbuf__ uint32_t*>(0); // 从ssbuf的0地址开始拷贝 + auto tempTiling = reinterpret_cast(&sharedParams); +#pragma unroll + for (int i = 0; i < sizeof(CVSharedParams) / sizeof(uint32_t); ++i, ++tempTilingSSbuf, ++tempTiling) { + *tempTiling = *tempTilingSSbuf; + } + } + this->ComputeConstexpr(); + this->InitGlobalBuffer(query, key, value, queryRope, keyRope, sparseIndices, \ + blockTable, actualSeqLengthsQ, actualSeqLengths, softmaxMax, softmaxSum, \ + workspace, tiling, tPipe); // gm设置 + this->InitCalcParamsEach(); + this->InitLocalBuffer(); +} + +template __aicore__ inline +void SparseFlashAttentionKernelMla::InitCalcParamsEach() +{ + // 计算总的基本块 + maxS2LoopCnt = 0; // 所有核中最大累计s2Loop + uint32_t totalBaseNum = 0; + uint32_t s1GBaseSize = constInfo.gSize; + uint32_t actBatchS2 = 1; + uint32_t coreNum = GetBlockNum(); // G128时相邻两个cube核处理一个s1，coreNum减半 + uint32_t currCoreIdx = aicIdx; + if constexpr (IS_SPLIT_G) { + coreNum = coreNum >> 1; + currCoreIdx = currCoreIdx >> 1; + } + uint32_t actBatchS1 = 1; + for (uint32_t bIdx = 0; bIdx < constInfo.bSize; bIdx++) { + actBatchS1 = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bIdx); // 不切S2，只关注S1 + if (actBatchS1 < constInfo.s1Size) { + constInfo.needInit = true; + } + totalBaseNum += actBatchS1 * actBatchS2; + } + uint32_t avgBaseNum = 1; + if (totalBaseNum > coreNum) { + avgBaseNum = (totalBaseNum + coreNum - 1) / coreNum; + if constexpr (IS_SPLIT_G) { + usedCoreNum = (totalBaseNum + avgBaseNum - 1) / avgBaseNum << 1; + } + } else { + if constexpr (IS_SPLIT_G) { + usedCoreNum = totalBaseNum << 1; + } else { + usedCoreNum = totalBaseNum; + } + } + + if constexpr (IS_SPLIT_G) { + maxS2LoopCnt = avgBaseNum * (constInfo.sparseBlockCount + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + } + + if (aicIdx >= usedCoreNum) { + return; + } + // 计算当前核的基本块 + uint32_t accumBaseNum = 0; // 当前累积的基本块数 + uint32_t targetBaseNum = 0; + uint32_t lastValidBIdx = 0; + uint32_t lastValidactBatchS1 = 0; + bool setStart = false; + targetBaseNum = (currCoreIdx + 1) * avgBaseNum; // 计算当前的目标权重 + uint32_t targetStartBaseNum = targetBaseNum - avgBaseNum; + for (uint32_t bN2Idx = 0; bN2Idx < constInfo.bSize * constInfo.n2Size; bN2Idx++) { + uint32_t bIdx = bN2Idx / constInfo.n2Size; + actBatchS1 = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bIdx); + for (uint32_t s1GIdx = 0; s1GIdx < actBatchS1; s1GIdx++) { + accumBaseNum += 1; + if (!setStart && accumBaseNum >= targetStartBaseNum) { + constInfo.bN2Start = bN2Idx; + constInfo.gS1Start = s1GIdx; + setStart = true; + } + if (accumBaseNum >= targetBaseNum) { + // 更新当前核的End分核信息 + constInfo.bN2End = bN2Idx; + constInfo.gS1End = s1GIdx; + constInfo.s2End = 0; + if (currCoreIdx != 0) { + GetAxisStartIdx(constInfo.bN2Start, constInfo.gS1Start, 0); + } + return; + } + } + if ((actBatchS1 > 0) && (actBatchS2 > 0)) { + lastValidBIdx = bIdx; + lastValidactBatchS1 = actBatchS1; + } + } + if (!setStart) { + constInfo.bN2Start = lastValidBIdx; + constInfo.gS1Start = lastValidactBatchS1 - 1; + } + if (accumBaseNum < targetBaseNum) { + // 更新最后一个核的End分核信息 + constInfo.bN2End = lastValidBIdx; + constInfo.gS1End = lastValidactBatchS1 - 1; + constInfo.s2End = 0; + if (currCoreIdx != 0) { + GetAxisStartIdx(constInfo.bN2Start, constInfo.gS1Start, 0); + } + return; + } +} + +template +__aicore__ inline uint64_t +SparseFlashAttentionKernelMla::GetBalanceActualSeqLengths( + GlobalTensor &actualSeqLengths, uint32_t bIdx) +{ + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + if (bIdx > 0) { + return actualSeqQlenAddr[bIdx] - actualSeqQlenAddr[bIdx - 1]; + } else if (bIdx == 0) { + return actualSeqQlenAddr[0]; + } else { + return 0; + } + } else { + if (constInfo.isActualLenDimsNull == 1) { + return constInfo.s1Size; + } else { + return actualSeqQlenAddr[bIdx]; + } + } +} + +template +__aicore__ inline void SparseFlashAttentionKernelMla::GetAxisStartIdx(uint32_t bN2EndPrev, + uint32_t s1GEndPrev, + uint32_t s2EndPrev) +{ + uint32_t bEndPrev = bN2EndPrev / constInfo.n2Size; + uint32_t actualSeqQPrev = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bEndPrev); + uint32_t s1GPrevBaseNum = actualSeqQPrev; + constInfo.bN2Start = bN2EndPrev; + constInfo.gS1Start = s1GEndPrev; + + constInfo.s2Start = 0; + if (s1GEndPrev >= s1GPrevBaseNum - 1) { // 上个核把S1G处理完了 + constInfo.gS1Start = 0; + constInfo.bN2Start++; + } else { + constInfo.gS1Start++; + } +} + +template __aicore__ inline +void SparseFlashAttentionKernelMla::InitGlobalBuffer( + __gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *queryRope, __gm__ uint8_t *keyRope, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *blockTable, __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths, + __gm__ uint8_t *softmaxMax, __gm__ uint8_t *softmaxSum, + __gm__ uint8_t *workspace, const SparseFlashAttentionTilingDataMla *__restrict tiling, TPipe *tPipe) +{ + if (actualSeqLengthsQ != nullptr) { + actualSeqQlenAddr = (__gm__ int32_t *)actualSeqLengthsQ; + } + if (actualSeqLengths != nullptr) { + actualSeqKvlenAddr = (__gm__ int32_t *)actualSeqLengths; + } + + vecBlock.InitGlobalBuffer(key, value, keyRope, sparseIndices, blockTable, softmaxMax, softmaxSum); + cubeBlock.InitCubeInput(key, keyRope, sparseIndices, blockTable, actualSeqLengthsQ, constInfo); +} + +template +__aicore__ inline void +SparseFlashAttentionKernelMla::InitMMResBuf(__gm__ uint8_t *workspace) +{ + uint32_t mm1ResultSize = constInfo.s1BaseSize / CV_RATIO * constInfo.s2BaseSize * sizeof(T); + uint32_t mm2ResultSize = constInfo.s1BaseSize / CV_RATIO * 512 * sizeof(T); + uint32_t mm2LeftSize = constInfo.s1BaseSize * constInfo.s2BaseSize * sizeof(Q_T); + uint32_t mm1RightSize = constInfo.s2BaseSize * 576 * sizeof(Q_T); + l1BufferManager.Init(pipe, 524288); // 512 * 1024 + // 保存p结果的L1内存必须放在第一个L1 policy上，保证和vec申请的地址相同 + l1RightBuffers.Init(l1BufferManager, mm1RightSize); + l1RightBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID); + crossCoreSyncBufId++; + l1RightBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID); + crossCoreSyncBufId++; + l1RightBuffers.Get().SetCrossCoreID(crossCoreSyncBufId, INVALID_CROSS_CORE_EVENT_ID); + crossCoreSyncBufId++; + if ASCEND_IS_AIC { + l1RightBuffers.Get().SetCrossCore(); + l1RightBuffers.Get().SetCrossCore(); + l1RightBuffers.Get().SetCrossCore(); + } + ubBufferManager.Init(pipe, mm1ResultSize * 2 + mm2ResultSize); + bmm2Buffers.Init(ubBufferManager, mm2ResultSize); + bmm2Buffers.Get().SetCrossCoreID(crossCoreSyncBufId, crossCoreSyncBufId); + crossCoreSyncBufId++; + if ASCEND_IS_AIV { + bmm2Buffers.Get().SetCrossCore(); + } + bmm1Buffers.Init(ubBufferManager, mm1ResultSize); + bmm1Buffers.Get().SetCrossCoreID(crossCoreSyncBufId, crossCoreSyncBufId); + crossCoreSyncBufId++; + bmm1Buffers.Get().SetCrossCoreID(crossCoreSyncBufId, crossCoreSyncBufId); + crossCoreSyncBufId++; + if ASCEND_IS_AIV { + bmm1Buffers.Get().SetCrossCore(); + bmm1Buffers.Get().SetCrossCore(); + } + + uint32_t v0ResSize = constInfo.s2BaseSize * 576U * sizeof(Q_T); + int64_t totalOffset; + if constexpr (IS_SPLIT_G) { + totalOffset = v0ResSize * 3 * (aicIdx >> 1U); + } else { + totalOffset = v0ResSize * 3 * aicIdx; + } + gmBufferManager.Init(workspace + totalOffset); + v0ResGmBuffers.Init(gmBufferManager, v0ResSize); + v0ResGmBuffers.Get().SetCrossCoreID(INVALID_CROSS_CORE_EVENT_ID, crossCoreSyncBufId); + crossCoreSyncBufId++; + v0ResGmBuffers.Get().SetCrossCoreID(INVALID_CROSS_CORE_EVENT_ID, crossCoreSyncBufId); + crossCoreSyncBufId++; + v0ResGmBuffers.Get().SetCrossCoreID(INVALID_CROSS_CORE_EVENT_ID, crossCoreSyncBufId); + crossCoreSyncBufId++; +} + +template +__aicore__ inline void SparseFlashAttentionKernelMla::InitLocalBuffer() +{ + vecBlock.InitLocalBuffer(pipe, constInfo); +} + +template +__aicore__ inline void SparseFlashAttentionKernelMla::ComputeConstexpr() +{ + // 计算轴的乘积 + usedCoreNum = sharedParams.usedCoreNum; + + if ASCEND_IS_AIC { + constInfo.bSize = this->sharedParams.bSize; + constInfo.gSize = this->sharedParams.gSize; + constInfo.s1Size = this->sharedParams.s1Size; + constInfo.dSizeV = 512; + constInfo.needInit = this->sharedParams.needInit; + } + constInfo.n2Size = sharedParams.n2Size; + constInfo.s2Size = sharedParams.s2Size; + constInfo.dSize = sharedParams.dSize; + constInfo.dSizeVInput = sharedParams.dSizeVInput; + constInfo.dSizeRope = 64; + constInfo.dSizeNope = 512; + constInfo.tileSize = sharedParams.tileSize; + constInfo.sparseBlockCount = sharedParams.sparseBlockCount; + constInfo.sparseBlockSize = 1; + constInfo.cmpRatio = sharedParams.cmpRatio; + constInfo.oriWinLeft = sharedParams.oriWinLeft; + constInfo.oriWinRight = sharedParams.oriWinRight; + constInfo.sparseMode = sharedParams.oriMaskMode; + constInfo.s1S2 = constInfo.s1Size * constInfo.s2Size; + constInfo.gS1 = constInfo.gSize * constInfo.s1Size; + constInfo.n2G = constInfo.n2Size * constInfo.gSize; + constInfo.gD = constInfo.gSize * constInfo.dSize; + constInfo.n2GD = constInfo.n2Size * constInfo.gD; + constInfo.s1Dv = constInfo.s1Size * constInfo.dSizeV; + constInfo.s2Dv = constInfo.s2Size * constInfo.dSizeV; + constInfo.n2Dv = constInfo.n2Size * constInfo.dSizeV; + constInfo.gDv = constInfo.gSize * constInfo.dSizeV; + constInfo.gS1Dv = constInfo.gSize * constInfo.s1Dv; + constInfo.isActualLenDimsNull = sharedParams.isActualSeqLengthsNull; + constInfo.isActualLenDimsKVNull = sharedParams.isActualSeqLengthsKVNull; + constInfo.n2S2Dv = constInfo.n2Size * constInfo.s2Dv; + constInfo.n2GDv = constInfo.n2Size * constInfo.gDv; + constInfo.s2BaseN2Dv = constInfo.s2BaseSize * constInfo.n2Dv; + constInfo.n2GS1Dv = constInfo.n2Size * constInfo.gS1Dv; + constInfo.layoutType = sharedParams.layoutType; + + constInfo.isActualLenDimsNull = sharedParams.isActualSeqLengthsNull; + constInfo.isActualLenDimsKVNull = sharedParams.isActualSeqLengthsKVNull; + + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + // (BS)ND + constInfo.s1BaseN2GDv = constInfo.s1BaseSize * constInfo.n2GDv; + constInfo.mm1Ka = constInfo.n2Size * constInfo.dSize; + if ASCEND_IS_AIV { + constInfo.attentionOutStride = (constInfo.n2G - constInfo.gSize) * constInfo.dSizeV * sizeof(OUTPUT_T); + } + } else if constexpr (LAYOUT_T == SFA_LAYOUT::BSND) { + // BSH/BSNGD + constInfo.s1BaseN2GDv = constInfo.s1BaseSize * constInfo.n2GDv; + constInfo.mm1Ka = constInfo.n2Size * constInfo.dSize; + if ASCEND_IS_AIV { + constInfo.attentionOutStride = (constInfo.n2G - constInfo.gSize) * constInfo.dSizeV * sizeof(OUTPUT_T); + } + } + if ASCEND_IS_AIV { + constInfo.softmaxScale = sharedParams.softmaxScale; + constInfo.oriBlockSize = sharedParams.oriBlockSize; + constInfo.oriMaxBlockNumPerBatch = sharedParams.oriMaxBlockNumPerBatch; + } + + InitUniqueConstInfo(); +} + +template +__aicore__ inline void SparseFlashAttentionKernelMla::InitUniqueConstInfo() +{ + // bsize + 1-> bsize + this->constInfo.actualSeqLenSize = this->sharedParams.bSize; + this->constInfo.actualSeqLenKVSize = this->sharedParams.bSize; +} + +template +__aicore__ inline void SparseFlashAttentionKernelMla::Process() +{ + // SyncAll Cube和Vector都需要调用 + if (this->sharedParams.needInit) { + SyncAll(); + } + + ProcessMainLoop(); +} + +template +__aicore__ inline void SparseFlashAttentionKernelMla::ProcessMainLoop() +{ + bool hasLoad = aicIdx < usedCoreNum; + if (!hasLoad) { + if ASCEND_IS_AIV { + if constexpr (IS_SPLIT_G) { + for (int64_t loopCnt = 0; loopCnt < maxS2LoopCnt; loopCnt++) { + CrossCoreSetFlag(15); + CrossCoreWaitFlag(15); + } + } + } + return; + } + + // 适配分核左闭右开 + uint32_t bIdx = constInfo.bN2End / constInfo.n2Size; + uint32_t actS1Size = GetBalanceActualSeqLengths(actualSeqLengthsQGm, bIdx); + uint32_t gS1max = actS1Size; + if (constInfo.gS1End + 1 < gS1max) { + /* constInfo.gS1End != gS1max时，gS1End需要往后加一格, bN2End不变 */ + constInfo.gS1End = constInfo.gS1End + 1; + } else { + /* constInfo.gS1End == gS1max，bN2End需要往后加一格，bN2End变为0，以代表末尾 */ + constInfo.bN2End = constInfo.bN2End + 1; + constInfo.gS1End = 0; + } + + // 分核信息 + uint32_t bN2StartIdx = constInfo.bN2Start; + uint32_t bN2EndIdx = constInfo.bN2End; + uint32_t gS1StartIdx = constInfo.gS1Start; + uint32_t nextGs1Idx = constInfo.gS1End; + uint32_t s2StartIdx = 0; + uint32_t s2EndIdx = 0; + uint32_t s2LoopLimit = 0; + + if (nextGs1Idx != 0) { + bN2EndIdx++; + } + + int64_t taskId = 0; + bool notLast = true; + RunInfo runInfo[3]; + RunParamStr runParam; + int64_t multiCoreInnerIdx = 1; + + for (int64_t bnIdx = bN2StartIdx; bnIdx < bN2EndIdx; bnIdx++) { + bool lastBN = (bnIdx == bN2EndIdx - 1); + runParam.boIdx = bnIdx; + runParam.n2oIdx = 0; + ComputeParamBatch(runParam, this->constInfo, + this->actualSeqQlenAddr, this->actualSeqKvlenAddr); + ComputeS1LoopInfo(runParam, this->constInfo, lastBN, nextGs1Idx, gS1StartIdx); + + int64_t gS1LoopEnd = lastBN ? (runParam.gs1LoopEndIdx + PRELOAD_NUM) : runParam.gs1LoopEndIdx; + for (int64_t gS1Index = runParam.gs1LoopStartIdx; gS1Index < gS1LoopEnd; gS1Index++) { + bool notLastTwoLoop = true; + if (lastBN) { + int32_t extraGS1 = gS1Index - runParam.gs1LoopEndIdx; + switch (extraGS1) { + case 0: + notLastTwoLoop = false; + break; + case 1: + notLast = false; + notLastTwoLoop = false; + break; + default: + break; + } + } + if (notLastTwoLoop) { + this->ComputeAxisIdxByBnAndGs1(bnIdx, gS1Index, runParam); + bool s1NoNeedCalc = ComputeParamS1( + runParam, this->constInfo, gS1Index, this->actualSeqQlenAddr); + // s1和s2有任意一个不需要算, 则continue, 如果是当前核最后一次循环，则补充计算taskIdx+2的部分 + bool s2NoNeedCalc = + ComputeS2LoopInfo(runParam, this->constInfo); + if (s1NoNeedCalc || s2NoNeedCalc) { + continue; + } + s2LoopLimit = runParam.s2LoopEndIdx - 1; + if constexpr (IS_SPLIT_G) { + maxS2LoopCnt -= (s2LoopLimit + 1); + } + } else { + s2LoopLimit = 0; + } + for (int64_t s2LoopCount = 0; s2LoopCount <= s2LoopLimit; ++s2LoopCount) { + if (notLastTwoLoop) { + RunInfo &runInfo1 = runInfo[taskId % 3]; + this->SetRunInfo(runInfo1, runParam, taskId, s2LoopCount, s2LoopLimit, multiCoreInnerIdx); + if ASCEND_IS_AIC { + this->cubeBlock.IterateBmm1(this->bmm1Buffers.Get(), + this->l1RightBuffers.Get(), v0ResGmBuffers.Get(), runInfo1, this->constInfo); + } else { + this->vecBlock.ProcessVec0(this->l1RightBuffers.Get(), v0ResGmBuffers.Get(), + runInfo1, this->constInfo, 0); + } + } else { + if ASCEND_IS_AIV { + if constexpr (IS_SPLIT_G) { + if (maxS2LoopCnt > 0) { + maxS2LoopCnt--; + CrossCoreSetFlag<0, PIPE_MTE3>(15); + CrossCoreWaitFlag<0, PIPE_MTE3>(15); + } + } + } + } + if (taskId > 0 && notLast) { + auto &runInfo2 = runInfo[(taskId + 2) % 3]; + if ASCEND_IS_AIV { + this->vecBlock.ProcessVec1(this->l1RightBuffers.GetReused(), + this->bmm1Buffers.Get(), runInfo2, this->constInfo); + } else { + RunInfo &runInfo2 = runInfo[(taskId + 2) % 3]; + this->cubeBlock.IterateBmm2(this->bmm2Buffers.Get(), this->l1RightBuffers, + this->l1RightBuffers.GetReused(), runInfo2, this->constInfo); + } + } + if (taskId > 1) { + if ASCEND_IS_AIV { + RunInfo &runInfo3 = runInfo[(taskId + 1) % 3]; + this->vecBlock.ProcessVec2(this->bmm2Buffers.Get(), runInfo3, this->constInfo); + } + } + ++taskId; + } + ++multiCoreInnerIdx; + } + gS1StartIdx = 0; + } + if ASCEND_IS_AIV { + if constexpr (IS_SPLIT_G) { + for (int64_t loopCnt = 0; loopCnt < maxS2LoopCnt; loopCnt++) { + CrossCoreSetFlag<0, PIPE_MTE3>(15); + CrossCoreWaitFlag<0, PIPE_MTE3>(15); + } + } + } +} + +template +__aicore__ inline void SparseFlashAttentionKernelMla::ComputeAxisIdxByBnAndGs1( + int64_t bnIndex, int64_t gS1Index, RunParamStr &runParam) +{ + // GS1合轴, 不切G, 只切S1 + runParam.s1oIdx = gS1Index * runParam.qSNumInOneBlock; + if constexpr (IS_SPLIT_G) { + runParam.goIdx = (aicIdx % 2 == 0) ? 0 : 64; // N1=128场景，相邻cube核处理一个s1，第一个cube核承担0-63行g，第二个cube核承担后64行g + } else { + runParam.goIdx = 0; + } +} + +template +__aicore__ inline void SparseFlashAttentionKernelMla::SetRunInfo( + RunInfo &runInfo, RunParamStr &runParam, int64_t taskId, int64_t s2LoopCount, + int64_t s2LoopLimit, int64_t multiCoreInnerIdx) +{ + if (s2LoopCount < runParam.kvLoopEndIdx) { + runInfo.s2StartIdx = runParam.s2LineStartIdx; + runInfo.s2EndIdx = runParam.s2LineEndIdx; + } + runInfo.s2LoopCount = s2LoopCount; + if (runInfo.multiCoreInnerIdx != multiCoreInnerIdx) { + runInfo.s1oIdx = runParam.s1oIdx; + runInfo.boIdx = runParam.boIdx; + runInfo.n2oIdx = runParam.n2oIdx; + runInfo.goIdx = runParam.goIdx; + runInfo.multiCoreInnerIdx = multiCoreInnerIdx; + runInfo.multiCoreIdxMod2 = multiCoreInnerIdx & 1; + runInfo.multiCoreIdxMod3 = multiCoreInnerIdx % 3; + } + + runInfo.taskId = taskId; + runInfo.taskIdMod2 = taskId & 1; + runInfo.taskIdMod3 = taskId % 3; + runInfo.s2LoopLimit = s2LoopLimit; + + runInfo.actualS1Size = runParam.actualS1Size; + runInfo.actualS2Size = runParam.actualS2Size; + runInfo.attentionOutOffset = runParam.attentionOutOffset; + runInfo.sOuterOffset = runParam.sOuterOffset; + runInfo.queryOffset = runParam.tensorQOffset; + this->ComputeBmm1Tail(runInfo, runParam); + InitUniqueRunInfo(runParam, runInfo); +} + +template +__aicore__ inline void SparseFlashAttentionKernelMla::InitUniqueRunInfo( + const RunParamStr &runParam, RunInfo &runInfo) +{ + InitTaskParamByRun(runParam, runInfo); +} + +template +__aicore__ inline void SparseFlashAttentionKernelMla::ComputeBmm1Tail( + RunInfo &runInfo, RunParamStr &runParam) +{ + // ------------------------S1 Base Related--------------------------- + runInfo.s1RealSize = runParam.s1RealSize; + runInfo.halfS1RealSize = runParam.halfS1RealSize; + runInfo.firstHalfS1RealSize = runParam.firstHalfS1RealSize; + runInfo.mRealSize = runParam.mRealSize; + runInfo.halfMRealSize = runParam.halfMRealSize; + runInfo.firstHalfMRealSize = runParam.firstHalfMRealSize; + + runInfo.vec2S1BaseSize = runInfo.halfS1RealSize; + runInfo.vec2MBaseSize = runInfo.halfMRealSize; + + // ------------------------S2 Base Related---------------------------- + runInfo.s2RealSize = constInfo.s2BaseSize; + runInfo.s2AlignedSize = runInfo.s2RealSize; + int64_t curS2LoopCnt = runInfo.s2LoopCount; + if (runInfo.s2StartIdx + (curS2LoopCnt + 1) * runInfo.s2RealSize > runInfo.s2EndIdx) { + runInfo.s2RealSize = runInfo.s2EndIdx - curS2LoopCnt * runInfo.s2RealSize - runInfo.s2StartIdx; + runInfo.s2AlignedSize = Align(runInfo.s2RealSize); + } +} +} +#endif // SPARSE_FLASH_ATTENTION_KERNEL_MLA_H diff --git a/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_kvcache.h b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_kvcache.h new file mode 100644 index 000000000..316906364 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_kvcache.h @@ -0,0 +1,300 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_attn_sharedkv_kvcache.h + * \brief + */ +#ifndef SPARSE_FLASH_ATTENTION_KVCACHE_H +#define SPARSE_FLASH_ATTENTION_KVCACHE_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "sparse_flash_attention_common_arch35.h" +#include "util_regbase.h" + +using namespace matmul; +using namespace regbaseutil; +using namespace AscendC; +using namespace AscendC::Impl::Detail; + +static constexpr uint32_t sparseModeZero = 0; +static constexpr uint32_t sparseModeThree = 3; + +TEMPLATE_INTF +__aicore__ inline void CalculateQueryOffset(RunParamStr& runParam, + const ConstInfo &constInfo, int32_t bIdx, + __gm__ int32_t* actualSeqQlenAddr) +{ + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + runParam.qBOffset = (bIdx == 0) ? 0 : actualSeqQlenAddr[bIdx - 1] * constInfo.gSize * 512; + runParam.qRopeBOffset = (bIdx == 0) ? 0 : actualSeqQlenAddr[bIdx - 1] * constInfo.gSize * 64; + } +} + +TEMPLATE_INTF +__aicore__ inline void GetSingleCoreParam(RunParamStr& runParam, const ConstInfo &constInfo, + __gm__ int32_t *actualSeqQlenAddr, __gm__ int32_t * actualSeqKvlenAddr) +{ + int32_t actualS1Size = 0; + int32_t actualS2Size = 0; + int32_t actualSeqMin = 1; + int32_t actualSeqKVMin = 1; + int32_t sIdx = runParam.boIdx; + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + // actual seq length first + if (actualSeqQlenAddr != nullptr) { + actualS1Size = (sIdx == 0) ? actualSeqQlenAddr[0] : + actualSeqQlenAddr[sIdx] - actualSeqQlenAddr[sIdx - 1]; + } else { + actualS1Size = constInfo.s1Size; + } + } else { + actualS1Size = (actualSeqQlenAddr == nullptr) ? constInfo.s1Size : + actualSeqQlenAddr[sIdx]; + } + + if (constInfo.isActualLenDimsKVNull) { + actualS2Size = constInfo.s2Size; + } else { + if constexpr (isPa) { + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + actualS2Size = actualSeqKvlenAddr[sIdx]; + } else { + actualS2Size = (constInfo.actualSeqLenKVSize == actualSeqKVMin) ? + actualSeqKvlenAddr[0] : actualSeqKvlenAddr[sIdx]; + } + } else { + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + actualS2Size = (sIdx == 0) ? actualSeqKvlenAddr[0] : + actualSeqKvlenAddr[sIdx] - actualSeqKvlenAddr[sIdx - 1]; + } else { + actualS2Size = (constInfo.actualSeqLenKVSize == actualSeqKVMin) ? + actualSeqKvlenAddr[0] : actualSeqKvlenAddr[sIdx]; + } + } + } + + runParam.actualS1Size = actualS1Size; + runParam.actualS2Size = actualS2Size; + if (constInfo.sparseMode == sparseModeZero) { + runParam.nextTokensPerBatch = MAX_PRE_NEXT_TOKENS; + } else { + runParam.nextTokensPerBatch = runParam.actualS2Size - runParam.actualS1Size; + } + runParam.preTokensPerBatch = runParam.actualS1Size; +} + +TEMPLATE_INTF +__aicore__ inline void ComputeParamBatch(RunParamStr& runParam, const ConstInfo &constInfo, + __gm__ int32_t *actualSeqQlenAddr, __gm__ int32_t *actualSeqKvlenAddr) +{ + GetSingleCoreParam(runParam, constInfo, actualSeqQlenAddr, actualSeqKvlenAddr); +} + +TEMPLATE_INTF +__aicore__ inline void ComputeS1LoopInfo(RunParamStr& runParam, const ConstInfo &constInfo, bool lastBN, + int64_t nextGs1Idx, int64_t gS1StartIdx) +{ + runParam.qSNumInOneBlock = 1; // 不切G轴, 计算每个基本块可以拷贝多少行s + runParam.gs1LoopStartIdx = gS1StartIdx; + if (runParam.nextTokensPerBatch < 0) { + int64_t gs1LoopStartIdx = runParam.nextTokensPerBatch * (-1) / runParam.qSNumInOneBlock + * runParam.qSNumInOneBlock; + if (gs1LoopStartIdx > gS1StartIdx) { + runParam.gs1LoopStartIdx = gs1LoopStartIdx; + } + } + + int32_t gs1LoopEndIdx = runParam.actualS1Size; // 不切G轴, 每次拷贝一行的topk，只算一行的qs + + // 不是最后一个bn, 赋值souterBlockNum + if (!lastBN) { + runParam.gs1LoopEndIdx = gs1LoopEndIdx; + } else { // 最后一个bn, 从数组下一个元素取值 + runParam.gs1LoopEndIdx = nextGs1Idx == 0 ? gs1LoopEndIdx : nextGs1Idx; + } + + if (runParam.gs1LoopStartIdx > runParam.gs1LoopEndIdx) { + runParam.gs1LoopStartIdx = runParam.gs1LoopEndIdx; + } +} + +TEMPLATE_INTF +__aicore__ inline void ComputeSouterParam(RunParamStr& runParam, const ConstInfo &constInfo, + uint32_t sOuterLoopIdx) +{ + int64_t cubeSOuterOffset = sOuterLoopIdx * runParam.qSNumInOneBlock; + if (runParam.actualS1Size == 0) { + runParam.s1RealSize = 0; + runParam.mRealSize = 0; + } else { + runParam.s1RealSize = Min(runParam.qSNumInOneBlock, runParam.actualS1Size - cubeSOuterOffset); + runParam.mRealSize = runParam.s1RealSize * constInfo.gSize; + if constexpr (IS_SPLIT_G) { + runParam.mRealSize = runParam.mRealSize >> 1; + } + } + + runParam.cubeMOuterOffset = cubeSOuterOffset * constInfo.gSize; + runParam.halfMRealSize = (runParam.mRealSize + 1) >> 1; + runParam.firstHalfMRealSize = runParam.halfMRealSize; + if (constInfo.subBlockIdx == 1) { + runParam.halfMRealSize = runParam.mRealSize - runParam.halfMRealSize; + runParam.mOuterOffset = runParam.cubeMOuterOffset + runParam.firstHalfMRealSize; + } else { + runParam.mOuterOffset = runParam.cubeMOuterOffset; + } + + runParam.halfS1RealSize = (runParam.s1RealSize + 1) >> 1; + runParam.firstHalfS1RealSize = runParam.halfS1RealSize; + if (constInfo.subBlockIdx == 1) { + runParam.halfS1RealSize = runParam.s1RealSize - runParam.halfS1RealSize; + runParam.sOuterOffset = cubeSOuterOffset + runParam.halfMRealSize / constInfo.gSize; + } else { + runParam.sOuterOffset = cubeSOuterOffset; + } + runParam.cubeSOuterOffset = cubeSOuterOffset; +} + +TEMPLATE_INTF +__aicore__ inline void LoopSOuterOffsetInit(RunParamStr& runParam, const ConstInfo &constInfo, + int32_t sIdx, __gm__ int32_t *cuSeqlensQAddr) +{ + if ASCEND_IS_AIV { + int64_t seqOffset = 0; + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + seqOffset = sIdx == 0 ? 0 : cuSeqlensQAddr[sIdx - 1]; + } else { + seqOffset = sIdx * constInfo.s1Size; + } + + int64_t attentionOutSeqOffset = seqOffset * constInfo.n2GDv; + if constexpr (LAYOUT_T == SFA_LAYOUT::BSND || LAYOUT_T == SFA_LAYOUT::TND) { + runParam.attentionOutOffset = attentionOutSeqOffset + + runParam.sOuterOffset * constInfo.n2GDv + runParam.n2oIdx * constInfo.gDv + + runParam.goIdx * constInfo.dSizeV; + } + if (constInfo.subBlockIdx == 1) { + runParam.attentionOutOffset += runParam.firstHalfMRealSize * constInfo.dSizeV; + } + if (constInfo.returnSoftmaxLse) { + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + // [N2, T, G] (TND) + runParam.softmaxLseOffset = runParam.n2oIdx * constInfo.s1Size * constInfo.gSize + + (seqOffset + runParam.sOuterOffset) * constInfo.gSize; + } else { + // [B, N2, S1, G] (BSND) + runParam.softmaxLseOffset = sIdx * constInfo.n2Size * constInfo.s1Size * constInfo.gSize + + runParam.n2oIdx * constInfo.s1Size * constInfo.gSize + + runParam.sOuterOffset * constInfo.gSize; + } + uint32_t aicIdx = constInfo.aivIdx >> 1U; + if (IS_SPLIT_G && aicIdx % 2U != 0) { + runParam.softmaxLseOffset += 64; // splitG时，需要偏移64 + } + if (constInfo.subBlockIdx == 1) { + runParam.softmaxLseOffset += runParam.firstHalfMRealSize; + } + } + } else { + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + runParam.tensorQOffset = runParam.qBOffset + runParam.cubeSOuterOffset * constInfo.n2GD + + runParam.n2oIdx * constInfo.gD + runParam.goIdx * constInfo.dSize; + runParam.tensorQRopeOffset = runParam.qRopeBOffset + runParam.cubeSOuterOffset * constInfo.n2GD + + runParam.n2oIdx * constInfo.gD + runParam.goIdx * constInfo.dSizeRope; + } else { + runParam.tensorQOffset = runParam.qBOffset + runParam.n2oIdx * constInfo.gS1D + + runParam.goIdx * constInfo.s1D + runParam.cubeSOuterOffset * constInfo.dSize; + runParam.tensorQRopeOffset = runParam.qRopeBOffset + runParam.n2oIdx * constInfo.gS1D + + runParam.goIdx * constInfo.s1D + runParam.cubeSOuterOffset * constInfo.dSizeRope; + } + } +} + +TEMPLATE_INTF +__aicore__ inline bool ComputeParamS1(RunParamStr& runParam, const ConstInfo &constInfo, + uint32_t sOuterLoopIdx, __gm__ int32_t *cuSeqlensQAddr) +{ + if (runParam.nextTokensPerBatch < 0) { + if (runParam.s1oIdx < (runParam.nextTokensPerBatch * (-1)) \ + / runParam.qSNumInOneBlock * runParam.qSNumInOneBlock) { + return true; + } + } + + ComputeSouterParam(runParam, constInfo, sOuterLoopIdx); + + LoopSOuterOffsetInit(runParam, constInfo, runParam.boIdx, cuSeqlensQAddr); + return false; +} + +TEMPLATE_INTF +__aicore__ inline bool ComputeLastBN(RunParamStr& runParam, __gm__ int32_t *cuSeqlensQAddr) +{ + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + // TND格式下 相邻Batch中当actualSeqQlen相等时则返回true + if (runParam.boIdx > 0 && ((runParam.boIdx == 0 && cuSeqlensQAddr[runParam.boIdx] == 0) || + (cuSeqlensQAddr[runParam.boIdx] - cuSeqlensQAddr[runParam.boIdx - 1] == 0))) { + return true; + } + } + return false; +} + +TEMPLATE_INTF +__aicore__ inline int64_t ClipSInnerTokenCube(int64_t sInnerToken, int64_t minValue, int64_t maxValue) +{ + sInnerToken = sInnerToken > minValue ? sInnerToken : minValue; + sInnerToken = sInnerToken < maxValue ? sInnerToken : maxValue; + return sInnerToken; +} + +TEMPLATE_INTF +__aicore__ inline bool ComputeS2LoopInfo(RunParamStr& runParam, const ConstInfo &constInfo) +{ + if (runParam.actualS2Size == 0) { + runParam.kvLoopEndIdx = 0; + runParam.s2LoopEndIdx = 0; + return true; + } + uint32_t s2BaseSize = constInfo.s2BaseSize; + + if (constInfo.sparseMode == sparseModeZero) { + runParam.s2LineStartIdx = 0; + runParam.s2LineEndIdx = Min(runParam.actualS2Size, constInfo.sparseBlockCount); + } else if (constInfo.sparseMode == sparseModeThree) { + runParam.s2LineStartIdx = ClipSInnerTokenCube( + runParam.cubeSOuterOffset - runParam.preTokensPerBatch, 0, runParam.actualS2Size); + runParam.s2LineEndIdx = ClipSInnerTokenCube( + runParam.cubeSOuterOffset + runParam.nextTokensPerBatch + + runParam.s1RealSize, 0, runParam.actualS2Size); + runParam.s2LineEndIdx = Min(runParam.s2LineEndIdx, constInfo.sparseBlockCount); // 当前LI输出的block size只可能是1 + } + + runParam.kvLoopEndIdx = (runParam.s2LineEndIdx + s2BaseSize - 1) / s2BaseSize; + runParam.s2LoopEndIdx = runParam.kvLoopEndIdx; + return false; +} + +TEMPLATE_INTF +__aicore__ inline void InitTaskParamByRun(const RunParamStr& runParam, RunInfo &runInfo) +{ + runInfo.boIdx = runParam.boIdx; + runInfo.preTokensPerBatch = runParam.preTokensPerBatch; + runInfo.nextTokensPerBatch = runParam.nextTokensPerBatch; + runInfo.actualS1Size = runParam.actualS1Size; + runInfo.actualS2Size = runParam.actualS2Size; + runInfo.softmaxLseOffset = runParam.softmaxLseOffset; + runInfo.qSNumInOneBlock = runParam.qSNumInOneBlock; + runInfo.kvLoopEndIdx = runParam.kvLoopEndIdx; +} + +#endif // SPARSE_FLASH_ATTENTION_KVCACHE_H diff --git a/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_service_cube_mla.h b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_service_cube_mla.h new file mode 100644 index 000000000..6ae227d08 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_service_cube_mla.h @@ -0,0 +1,387 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_flash_attention_service_cube_mla.h + * \brief + */ +#ifndef SPARSE_FLASH_ATTENTION_SERVICE_CUBE_MLA_H +#define SPARSE_FLASH_ATTENTION_SERVICE_CUBE_MLA_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "sparse_flash_attention_common_arch35.h" +#include "util_regbase.h" + +#if __has_include("../../common/op_kernel/offset_calculator.h") +#include "../../common/op_kernel/offset_calculator.h" +#else +#include "../common/offset_calculator.h" +#endif +#if __has_include("../../common/op_kernel/matmul.h") +#include "../../common/op_kernel/matmul.h" +#else +#include "../common/matmul.h" +#endif +#if __has_include("../../common/op_kernel/FixpipeOut.h") +#include "../../common/op_kernel/FixpipeOut.h" +#else +#include "../common/FixpipeOut.h" +#endif +#if __has_include("../../common/op_kernel/CopyInL1.h") +#include "../../common/op_kernel/CopyInL1.h" +#else +#include "../common/CopyInL1.h" +#endif + +using namespace AscendC; +using namespace AscendC::Impl::Detail; +using namespace regbaseutil; +using namespace fa_base_matmul; +namespace BaseApi { + +template +__aicore__ inline constexpr GmFormat GetQueryGmFormat() +{ + if constexpr (LAYOUT == SFA_LAYOUT::BSND) { + return GmFormat::BSNGD; + } else { + return GmFormat::TNGD; + } +} + +TEMPLATES_DEF +class SFAMatmulService { +public: + /* =================编译期常量的基本块信息================= */ + static constexpr uint32_t s1BaseSize = 64; + static constexpr uint32_t s2BaseSize = 128; + static constexpr uint32_t dBaseSize = 576; + static constexpr uint32_t dBaseMatmulSize = 128; + + __aicore__ inline SFAMatmulService() {}; + __aicore__ inline void InitCubeBlock(TPipe *pipe, BufferManager &l1BuffMgr, + __gm__ uint8_t *query, __gm__ uint8_t *queryRope); + __aicore__ inline void InitCubeInput(__gm__ uint8_t *key, __gm__ uint8_t *keyRope, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *blockTable, __gm__ uint8_t *actualSeqLengthsQ, const ConstInfo& constInfo); + __aicore__ inline void IterateBmm1(Buffer &output, + Buffer &inputRightBuf, + Buffer &v0ResGm, + RunInfo &runInfo, ConstInfo &constInfo); + + __aicore__ inline void IterateBmm2(Buffer &outputBuf, + BuffersPolicy3buff &inputLeftBuffers, + Buffer &inputRightBuf, RunInfo &runInfo, + ConstInfo &constInfo); + +private: + __aicore__ inline void InitLocalBuffer(BufferManager &l1BuffMgr); + __aicore__ inline void InitGmTensor(__gm__ uint8_t *cuSeqlensQ, const ConstInfo& constInfo); + + __aicore__ inline void IterateBmm1SFA(Buffer &outputBuf, + Buffer &inputRightBuf, + Buffer &v0ResGm, + RunInfo &runInfo, ConstInfo &constInfo); + + // --------------------Bmm2-------------------------- + __aicore__ inline void IterateBmm2SFA(Buffer &outputBuf, + BuffersPolicy3buff &inputLeftBuffers, + Buffer &inputRightBuf, RunInfo &runInfo, + ConstInfo &constInfo); + TPipe *tPipe; + /* =====================GM变量==================== */ + static constexpr GmFormat Q_FORMAT = GetQueryGmFormat(); + FaGmTensor queryGm; + FaGmTensor queryRopeGm; + + FaGmTensor keyGm; + GlobalTensor blockTableGm; + FaGmTensor curKvGm; + GlobalTensor cuSeqlensQGm; + + /* =====================运行时变量==================== */ + uint32_t kvCacheBlockSize = 0; + uint32_t maxBlockNumPerBatch = 0; + TEventID mte1ToMte2Id[3]; + TEventID mte2ToMte1Id[3]; + + /* =====================LocalBuffer变量==================== */ + BufferManager l0aBufferManager; + BufferManager l0bBufferManager; + BufferManager l0cBufferManager; + + // D小于等于256 mm1左矩阵Q，GS1循环内左矩阵复用, GS1循环间开pingpong；D大于256使用单块Buffer，S1循环间驻留；fp32场景单块不驻留 + BuffersPolicySingleBuffer l1QBuffers; + + // L0A + BuffersPolicyDB mmL0ABuffers; + // L0B + BuffersPolicyDB mmL0BBuffers; + // L0C + BuffersPolicyDB mmL0CBuffers; +}; + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void SFAMatmulService::InitCubeBlock( + TPipe *pipe, BufferManager &l1BuffMgr, __gm__ uint8_t *query, __gm__ uint8_t *queryRope) +{ + if ASCEND_IS_AIC { + tPipe = pipe; + this->queryGm.gmTensor.SetGlobalBuffer((__gm__ Q_T *)query); + this->queryRopeGm.gmTensor.SetGlobalBuffer((__gm__ Q_T *)queryRope); + InitLocalBuffer(l1BuffMgr); + } +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void SFAMatmulService::InitCubeInput(__gm__ uint8_t *key, __gm__ uint8_t *keyRope, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *blockTable, __gm__ uint8_t *actualSeqLengthsQ, + const ConstInfo& constInfo) +{ + if ASCEND_IS_AIC { + mte1ToMte2Id[0] = GetTPipePtr()->AllocEventID(); + mte1ToMte2Id[1] = GetTPipePtr()->AllocEventID(); + mte1ToMte2Id[2] = GetTPipePtr()->AllocEventID(); + mte2ToMte1Id[0] = GetTPipePtr()->AllocEventID(); + mte2ToMte1Id[1] = GetTPipePtr()->AllocEventID(); + mte2ToMte1Id[2] = GetTPipePtr()->AllocEventID(); + InitGmTensor(actualSeqLengthsQ, constInfo); + } +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void +SFAMatmulService::InitLocalBuffer(BufferManager &l1BuffMgr) +{ + constexpr uint32_t mm1LeftSize = s1BaseSize * dBaseSize * sizeof(Q_T); + l1QBuffers.Init(l1BuffMgr, mm1LeftSize); + + // L0A B C 当前写死，能否通过基础api获取 + l0aBufferManager.Init(tPipe, L0AB_SHARED_SIZE_64K); + l0bBufferManager.Init(tPipe, L0AB_SHARED_SIZE_64K); + l0cBufferManager.Init(tPipe, L0C_SHARED_SIZE_256K); + + mmL0ABuffers.Init(l0aBufferManager, BUFFER_SIZE_16K); // db类型，填入数值是总大小的一半 + mmL0BBuffers.Init(l0bBufferManager, BUFFER_SIZE_32K); + mmL0CBuffers.Init(l0cBufferManager, BUFFER_SIZE_128K); +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void +SFAMatmulService::InitGmTensor(__gm__ uint8_t *actualSeqLengthsQ, const ConstInfo& constInfo) +{ + if constexpr (LAYOUT_T == SFA_LAYOUT::BSND) { + this->queryGm.offsetCalculator.Init(constInfo.bSize, constInfo.n2Size, constInfo.gSize, + constInfo.s1Size, constInfo.dSize); + this->queryRopeGm.offsetCalculator.Init(constInfo.bSize, constInfo.n2Size, constInfo.gSize, + constInfo.s1Size, constInfo.dSizeRope); + } else { // SFA_LAYOUT::TND + GlobalTensor actualSeqQLen; + actualSeqQLen.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsQ); + this->queryGm.offsetCalculator.Init(constInfo.n2Size, constInfo.gSize, constInfo.dSize, + actualSeqQLen, constInfo.actualSeqLenSize); + this->queryRopeGm.offsetCalculator.Init(constInfo.n2Size, constInfo.gSize, constInfo.dSizeRope, + actualSeqQLen, constInfo.actualSeqLenSize); + } +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void SFAMatmulService::IterateBmm1( + Buffer &outputBuf, + Buffer &inputRightBuf, + Buffer &v0ResGm, RunInfo &runInfo, + ConstInfo &constInfo) +{ + IterateBmm1SFA(outputBuf, inputRightBuf, v0ResGm, runInfo, constInfo); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void SFAMatmulService::IterateBmm2( + Buffer &outputBuf, + BuffersPolicy3buff &inputLeftBuffers, + Buffer &inputRightBuf, RunInfo &runInfo, + ConstInfo &constInfo) +{ + IterateBmm2SFA(outputBuf, inputLeftBuffers, inputRightBuf, runInfo, constInfo); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void SFAMatmulService::IterateBmm1SFA( + Buffer &outputBuf, + Buffer &inputRightBuf, + Buffer &v0ResGm, RunInfo &runInfo, + ConstInfo &constInfo) +{ + Buffer inputLeftBuf; + // 左矩阵复用，S2的第一次循环加载左矩阵 + // 加载左矩阵到L1, 全载 + if (unlikely(runInfo.s2LoopCount == 0)) { // sOuter循环第一个基本块：搬运Q + inputLeftBuf = l1QBuffers.Get(); + inputLeftBuf.Wait(); // 占用L1A + LocalTensor inputLeftTensor = inputLeftBuf.GetTensor(); + uint32_t s1Coord = runInfo.s1oIdx * runInfo.qSNumInOneBlock; + uint64_t queryGmOffset = this->queryGm.offsetCalculator.GetOffset(runInfo.boIdx, runInfo.n2oIdx, + runInfo.goIdx, s1Coord, 0); + uint64_t queryRopeGmOffset = this->queryRopeGm.offsetCalculator.GetOffset(runInfo.boIdx, runInfo.n2oIdx, + runInfo.goIdx, s1Coord, 0); + CopyToL1Nd2Nz(inputLeftTensor, this->queryGm.gmTensor[queryGmOffset], + runInfo.mRealSize, 512, 512); // 64 constInfo.dSize constInfo.mm1Ka + CopyToL1Nd2Nz(inputLeftTensor[Align16Func(runInfo.mRealSize) * 512], + this->queryRopeGm.gmTensor[queryRopeGmOffset], runInfo.mRealSize, + 64, 64); // constInfo.dSize constInfo.mm1Ka + inputLeftBuf.Set(); // 通知 + } else { // 非S2的第一次循环直接复用Q + inputLeftBuf = l1QBuffers.GetPre(); + // 左矩阵复用时，sinner循环内不需要MTE2同步等待 + inputLeftBuf.Set(); // 通知 + } + + inputRightBuf.WaitCrossCore(); + SetFlag(mte2ToMte1Id[runInfo.taskIdMod3]); + WaitFlag(mte2ToMte1Id[runInfo.taskIdMod3]); + LocalTensor dst = inputRightBuf.GetTensor(); + v0ResGm.WaitCrossCore(); + GlobalTensor v0ResGmTensor = v0ResGm.template GetTensor(); + CopyToL1Nd2Nz(dst, v0ResGmTensor, runInfo.s2RealSize, 576, 576); + SetFlag(mte1ToMte2Id[runInfo.taskIdMod3]); + WaitFlag(mte1ToMte2Id[runInfo.taskIdMod3]); + + inputLeftBuf.Wait(); // 等待L1A + Buffer mm1ResL0C = mmL0CBuffers.Get(); + mm1ResL0C.Wait(); // 占用 + MMParam param = {static_cast(runInfo.mRealSize), // singleM + static_cast(runInfo.s2RealSize), // singleN + static_cast(constInfo.dSizeNope + constInfo.dSizeRope), // singleK + 0, // isLeftTranspose + 1 // isRightTranspose + }; + MatmulK( + inputLeftBuf.GetTensor(), inputRightBuf.GetTensor(), // mm1B直接用tensor的数据 + mmL0ABuffers, mmL0BBuffers, + mm1ResL0C.GetTensor(), + param); + if (unlikely(runInfo.s2LoopCount == runInfo.s2LoopLimit)) { + inputLeftBuf.Set(); // 释放L1A + } + + mm1ResL0C.Set(); // 通知 + mm1ResL0C.Wait(); // 等待L0C + + outputBuf.WaitCrossCore(); + FixpipeParamsC310 fixpipeParams; // L0C→UB + // L0C上的bmm1结果矩阵N方向的size大小; 同mmadParams.n; 为什么要8个元素对齐(32B对齐) // 128 + fixpipeParams.nSize = Align8Func(runInfo.s2RealSize); + // 有效数据不足16行，只需要输出部分行即可; L0C上的bmm1结果矩阵M方向的size大小(必须为偶数) // 128 + fixpipeParams.mSize = Align2Func(runInfo.mRealSize); + // L0C上bmm1结果相邻连续数据片段间隔(前面一个数据块的头与后面数据块的头的间隔), 单位为16*sizeof(T) + fixpipeParams.srcStride = Align16Func(fixpipeParams.mSize); + // mmResUb上两行之间的间隔，单位：element。 // 128:根据比对dump文件得到, ND方案(S1*S2)时脏数据用mask剔除 + fixpipeParams.dstStride = s2BaseSize; + fixpipeParams.dualDstCtl = 1; // 双目标模式，按M维度拆分，M / 2 * N写入每个UB, M必须为2的倍数 + fixpipeParams.params.ndNum = 1; + fixpipeParams.params.srcNdStride = 0; + fixpipeParams.params.dstNdStride = 0; + + Fixpipe(outputBuf.template GetTensor(), \ + mm1ResL0C.GetTensor(), fixpipeParams); // 将matmul结果从L0C搬运到UB + mm1ResL0C.Set(); // 释放L0C + outputBuf.SetCrossCore(); +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void SFAMatmulService::IterateBmm2SFA( + Buffer &outputBuf, + BuffersPolicy3buff &inputLeftBuffers, + Buffer &inputRightBuf, RunInfo &runInfo, + ConstInfo &constInfo) +{ + inputRightBuf.WaitCrossCore(); + + Buffer mm2ResL0C = mmL0CBuffers.Get(); + mm2ResL0C.Wait(); // 占用 + MMParam param = {static_cast(runInfo.mRealSize), // singleM + static_cast(constInfo.dSizeNope), // singleN + static_cast(runInfo.s2RealSize), // singleK + 0, // isLeftTranspose + 0 // isRightTranspose + }; + MatmulN( + inputRightBuf.GetTensor(s2BaseSize * constInfo.dSizeNope), // 左矩阵P 来自rope位置 + inputRightBuf.GetTensor(), // 右矩阵V nope + mmL0ABuffers, + mmL0BBuffers, + mm2ResL0C.GetTensor(), + param); + + inputRightBuf.SetCrossCore(); // bmm2才释放KV，在这里释放 + + mm2ResL0C.Set(); // 通知 + mm2ResL0C.Wait(); // 等待 + + outputBuf.WaitCrossCore(); + FixpipeParamsC310 fixpipeParams; // L0C→UB;FixpipeParamsM300:L0C→UB + fixpipeParams.nSize = Align8Func(constInfo.dSizeNope); // L0C上的bmm1结果矩阵N方向的size大小, 分档计算且vector2中通过mask筛选出实际有效值 + fixpipeParams.mSize = Align2Func(runInfo.mRealSize); // 有效数据不足16行，只需要输出部分行即可; L0C上的bmm1结果矩阵M方向的size大小; + fixpipeParams.srcStride = Align16Func(fixpipeParams.mSize); // L0C上bmm1结果相邻连续数据片段间隔（前面一个数据块的头与后面数据块的头的间隔） + fixpipeParams.dstStride = Align16Func(constInfo.dSizeNope); + fixpipeParams.dualDstCtl = 1; + fixpipeParams.params.ndNum = 1; + fixpipeParams.params.srcNdStride = 0; + fixpipeParams.params.dstNdStride = 0; + Fixpipe(outputBuf.template GetTensor(), + mm2ResL0C.GetTensor(), fixpipeParams); // 将matmul结果从L0C搬运到UB + mm2ResL0C.Set(); // 释放 + + outputBuf.SetCrossCore(); +} + +TEMPLATES_DEF +class SFAMatmulServiceDummy { +public: + __aicore__ inline SFAMatmulServiceDummy() {}; + __aicore__ inline void InitCubeBlock(TPipe *pipe, + BufferManager &l1BuffMgr, __gm__ uint8_t *query, __gm__ uint8_t *queryRope) {} + __aicore__ inline void InitCubeInput(__gm__ uint8_t *key, __gm__ uint8_t *keyRope, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *blockTable, + __gm__ uint8_t *actualSeqLengthsQ, const ConstInfo& constInfo) {} + __aicore__ inline void IterateBmm1(Buffer &outputBuf, + Buffer &inputRightBuf, + RunInfo &runInfo, ConstInfo &constInfo) {} + __aicore__ inline void IterateBmm2(Buffer &outputBuf, + BuffersPolicyDB &inputLeftBuffers, + Buffer &inputRightBuf, RunInfo &runInfo, + ConstInfo &constInfo) {} +}; + + +template +struct CubeBlockTraits; // 声明 + +/* 生成CubeBlockTraits */ +#define GEN_TRAIT_TYPE(name, ...) using name##_TRAITS = name; +#define GEN_TRAIT_CONST(name, type, ...) static constexpr type name##Traits = name; + +#define DEFINE_CUBE_BLOCK_TRAITS(CUBE_BLOCK_CLASS) \ + TEMPLATES_DEF_NO_DEFAULT \ + struct CubeBlockTraits> { \ + CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_TRAIT_TYPE) \ + CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_TRAIT_CONST) \ + } + +DEFINE_CUBE_BLOCK_TRAITS(SFAMatmulService); +DEFINE_CUBE_BLOCK_TRAITS(SFAMatmulServiceDummy); + +// /* 生成Arg Traits, kernel中只需要调用ARGS_TRAITS就可以获取所有CubeBlock中的模板参数 */ +#define GEN_ARGS_TYPE(name, ...) using name = typename CubeBlockTraits::name##_TRAITS; +#define GEN_ARGS_CONST(name, type, ...) static constexpr type name = CubeBlockTraits::name##Traits; +#define ARGS_TRAITS \ + CUBE_BLOCK_TRAITS_TYPE_FIELDS(GEN_ARGS_TYPE) \ + CUBE_BLOCK_TRAITS_CONST_FIELDS(GEN_ARGS_CONST) +} +#endif // SPARSE_FLASH_ATTENTION_SERVICE_CUBE_MLA_H diff --git a/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_service_vector_mla.h b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_service_vector_mla.h new file mode 100644 index 000000000..aed284adb --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/sparse_flash_attention_service_vector_mla.h @@ -0,0 +1,879 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_flash_attention_service_vector_mla.h + * \brief + */ +#ifndef SPARSE_FLASH_ATTENTION_SERVICE_VECTOR_MLA_H +#define SPARSE_FLASH_ATTENTION_SERVICE_VECTOR_MLA_H + +#include "util_regbase.h" +#include "sparse_flash_attention_common_arch35.h" +#include "kernel_operator_list_tensor_intf.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" + +#if __has_include("../../common/op_kernel/arch35/vf/vf_mul_sel_softmaxflashv2_cast_nz_sfa.h") +#include "../../common/op_kernel/arch35/vf/vf_mul_sel_softmaxflashv2_cast_nz_sfa.h" +#else +#include "../../common/arch35/vf/vf_mul_sel_softmaxflashv2_cast_nz_sfa.h" +#endif + +#if __has_include("../../common/op_kernel/arch35/vf/vf_flashupdate_new.h") +#include "../../common/op_kernel/arch35/vf/vf_flashupdate_new.h" +#else +#include "../../common/arch35/vf/vf_flashupdate_new.h" +#endif + +#if __has_include("../../common/op_kernel/buffers_policy.h") +#include "../../common/op_kernel/buffers_policy.h" +#else +#include "../../common/buffers_policy.h" +#endif +#if __has_include("../../common/op_kernel/buffer_manager.h") +#include "../../common/op_kernel/buffer_manager.h" +#else +#include "../../common/buffer_manager.h" +#endif +#if __has_include("../../common/op_kernel/buffer.h") +#include "../../common/op_kernel/buffer.h" +#else +#include "../../common/buffer.h" +#endif + +using namespace AscendC; +using namespace FaVectorApi; +using namespace AscendC::Impl::Detail; +using namespace regbaseutil; +using namespace matmul; + +namespace BaseApi { + +TEMPLATES_DEF +class SFAVectorService { +public: + // BUFFER的字节数 + static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; + /* =================编译期常量的基本块信息================= */ + static constexpr uint32_t s1BaseSize = 64; + static constexpr uint32_t s2BaseSize = 128; + static constexpr uint32_t vec1Srcstride = (s1BaseSize >> 1) + 1; + static constexpr uint32_t dVTemplateType = 512; + static constexpr uint32_t dTemplateAlign64 = Align64Func(dVTemplateType); + static constexpr uint32_t dVTemplateTypeInput = 576; + static constexpr float R0 = 1.0f; + static constexpr uint64_t SYNC_SINKS_BUF_FLAG = 6; + + // ==================== Functions ====================== + __aicore__ inline SFAVectorService() {}; + __aicore__ inline void InitVecBlock(TPipe *pipe, const SparseFlashAttentionTilingDataMla *__restrict tiling, + CVSharedParams &sharedParams, int32_t aicIdx, uint8_t subBlockIdx, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengths) + { + if ASCEND_IS_AIV { + tPipe = pipe; + tilingData = tiling; + if (actualSeqLengthsQ != nullptr) { + cuSeqlensQGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengthsQ); + } + if (actualSeqLengths != nullptr) { + actualSeqLengthsKVGm.SetGlobalBuffer((__gm__ int32_t *)actualSeqLengths); + } + this->InitCubeVecSharedParams(sharedParams, aicIdx, subBlockIdx); + this->GetExtremeValue(this->negativeFloatScalar); + } + } + + // 初始化LocalTensor + __aicore__ inline void InitLocalBuffer(TPipe *pipe, ConstInfo &constInfo); + // 初始化attentionOutGM + __aicore__ inline void CleanOutput(__gm__ uint8_t *attentionOut, __gm__ uint8_t *softmaxMax, + __gm__ uint8_t *softmaxSum, ConstInfo &constInfo); + __aicore__ inline void InitGlobalBuffer(__gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *keyRope, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *blockTable, __gm__ uint8_t *softmaxMax, + __gm__ uint8_t *softmaxSum); + __aicore__ inline void InitOutputSingleCore(ConstInfo &constInfo); + __aicore__ inline void ProcessVec0(Buffer &outputL1, + Buffer &v0ResGm, + const RunInfo &runInfo, ConstInfo &constInfo, int32_t startPos); + __aicore__ inline void ProcessVec1(Buffer &outputBuf, + Buffer &bmm1ResBuf, RunInfo &runInfo, + ConstInfo &constInfo); + using mm2ResPos = Buffer; + __aicore__ inline void ProcessVec2(mm2ResPos &bmm2ResBuf, RunInfo &runInfo, + ConstInfo &constInfo); + +private: + __aicore__ inline void ProcessSparseKv(Buffer &outputL1, + Buffer &v0ResGm, + const RunInfo &runInfo, ConstInfo &constInfo, int32_t startPos); + __aicore__ inline int64_t GetkeyOffset(int64_t s2Idx, const RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline void GetRealCmpS2Idx(int64_t &token0Idx, int64_t &token1Idx, int64_t s2IdxInBase, + const RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline void CopyInKvNotSparse(LocalTensor kvMergUb, int64_t v0Loop, int64_t dealRow, + int64_t s2StartIdx, const RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline uint32_t CopyInKvSparse(LocalTensor kvInUb, int64_t startRow, int64_t token0Idx, + int64_t token1Idx, const RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline void CalSparseCalSize(const RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline void CopyOutKvUb2Gm(Buffer &v0ResGm, + LocalTensor kvOutUb, int64_t dealRow, int64_t s2StartIdx, const RunInfo &runInfo, + ConstInfo &constInfo); + __aicore__ inline void CopyInSingleKv(LocalTensor kvInUb, + int64_t startRow, int64_t keyOffset, ConstInfo &constInfo); + /* VEC2_RES_T 表示bmm2ResUb当前的类型，VEC2_RES_T = Q_T那么不需要做Cast。另外，无效行场景当前默认需要做Cast */ + using VEC2_RES_T = T; + template + __aicore__ inline void Bmm2DataCopyOut(RunInfo &runInfo, ConstInfo &constInfo, + LocalTensor &vec2ResUb, int64_t vec2S1Idx, int64_t vec2CalcSize = 0); + template + __aicore__ inline void CopyOutAttentionOut(RunInfo &runInfo, + ConstInfo &constInfo, LocalTensor &vec2ResUb, + int64_t vec2S1Idx, int64_t vec2CalcSize); + __aicore__ inline void SoftmaxInitBuffer(); + __aicore__ inline void CopyFALseToGm(RunInfo &runInfo, ConstInfo &constInfo); + __aicore__ inline void InitCubeVecSharedParams(CVSharedParams &sharedParams, int32_t aicIdx, uint8_t subBlockIdx); + __aicore__ inline void GetExtremeValue(T &negativeScalar); + __aicore__ inline void InitSinksBuffer(ConstInfo &constInfo); + + TPipe *tPipe; + const SparseFlashAttentionTilingDataMla *__restrict tilingData; + + GlobalTensor attentionOutGm; + GlobalTensor keyGm; + GlobalTensor keyRopeGm; + GlobalTensor sparseIndicesGm; + GlobalTensor blockTableGm; + GlobalTensor cuSeqlensQGm; + GlobalTensor actualSeqLengthsKVGm; + GlobalTensor softmaxMaxGm; + GlobalTensor softmaxSumGm; + LocalTensor lseUb; + + TBuf<> commonTBuf; // common的复用空间 + TBuf<> sinksBuf; + TQue stage1OutQue[2]; // 2份表示可能存在pingpong + TBuf<> stage0OutBuf[2]; + TBuf<> stage2OutBuf; + TEventID mte3ToVAttnOutId; // 存放MTE3_V的eventId, 用于V2 attentionOut拷出阶段的同步 + TEventID vToMte3AttnOutId; // 存放V_MTE3的eventId, 用于V2 attentionOut拷出阶段的同步 + TEventID mte3ToVLseOutId; // 存放MTE3_V的eventId, 用于V1 LSE拷出阶段的同步 + TEventID vToMte3LseOutId; // 存放V_MTE3的eventId, 用于V1 LSE拷出阶段的同步 + TBuf<> softmaxMaxBuf[2]; + TBuf<> softmaxSumBuf[2]; + TBuf<> softmaxExpBuf[2]; + TBuf<> dequantScaleBuff; + TBuf<> lseBuf; + + TEventID mte2ToV; + TEventID mte2ToMte3[2]; + TEventID mte3ToMte2[2]; + + bool isSinks = false; + T negativeFloatScalar; + uint32_t maxBlockNumPerBatch; + uint32_t blockSize; + + int64_t sparseCalSize; + int64_t sparseS2Start; + int64_t sparseS2End; +}; + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void +SFAVectorService::GetRealCmpS2Idx(int64_t &token0Idx, int64_t &token1Idx, + int64_t s2IdxInBase, const RunInfo &runInfo, ConstInfo &constInfo) +{ + int64_t topkBS1Idx = 0; + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + uint64_t actualSeqQPrefixSum = runInfo.boIdx == 0 ? 0 : cuSeqlensQGm.GetValue(runInfo.boIdx - 1); + topkBS1Idx += (actualSeqQPrefixSum + runInfo.s1oIdx) * constInfo.sparseBlockCount; // T, N2(1), K + } else { + topkBS1Idx += runInfo.boIdx * constInfo.s1Size * constInfo.sparseBlockCount + + runInfo.s1oIdx * constInfo.sparseBlockCount; // B, S1, N2(1), K + } + int64_t cmpS2LoopCnt = runInfo.s2LoopCount; + int64_t topkKIdx = s2IdxInBase + cmpS2LoopCnt * constInfo.s2BaseSize; + if (unlikely(topkKIdx >= constInfo.sparseBlockCount)) { + token0Idx = -1; + } else { + token0Idx = sparseIndicesGm.GetValue(topkBS1Idx + topkKIdx) + runInfo.s2StartIdx; + } + topkKIdx += 1; + if (unlikely((topkKIdx >= constInfo.sparseBlockCount) || (s2IdxInBase + 1 >= sparseS2End))) { + token1Idx = -1; + } else { + token1Idx = sparseIndicesGm.GetValue(topkBS1Idx + topkKIdx) + runInfo.s2StartIdx; + } +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline +int64_t SFAVectorService::GetkeyOffset( + int64_t s2Idx, const RunInfo &runInfo, ConstInfo &constInfo) +{ + if (s2Idx < 0) { + return -1; + } + int64_t realkeyOffset = 0; + if constexpr (isPa) { + int64_t blkTableIdx = s2Idx / blockSize; + int64_t blkTableOffset = s2Idx % blockSize; + realkeyOffset = blockTableGm.GetValue(runInfo.boIdx * maxBlockNumPerBatch + blkTableIdx) * + static_cast(blockSize) + + blkTableOffset; // BlockNum, BlockSize, N(1), D + } else { + if constexpr (LAYOUT_T == SFA_LAYOUT::BSND) { + realkeyOffset = (runInfo.boIdx * constInfo.s2Size + s2Idx); // BSN(1)D + } else if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + int64_t batchKvStart = (runInfo.boIdx == 0) ? 0 : actualSeqLengthsKVGm.GetValue(runInfo.boIdx - 1); + realkeyOffset = (batchKvStart + s2Idx); + } + } + return realkeyOffset; +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void +SFAVectorService::CopyInSingleKv(LocalTensor kvInUb, int64_t startRow, + int64_t keyOffset, ConstInfo &constInfo) +{ + if (keyOffset < 0) { + return; + } + DataCopyExtParams intriParams; + + intriParams.blockCount = 1; + intriParams.dstStride = 0; + intriParams.srcStride = 0; + DataCopyPadExtParams padParams; + // 当前仅支持COMBINE模式 + uint32_t combineBytes = 512 * sizeof(KV_T); + intriParams.blockLen = combineBytes; + uint32_t combineDim = combineBytes / sizeof(KV_T); + uint32_t combineDimAlign = CeilAlign(combineBytes, BUFFER_SIZE_BYTE_32B) / sizeof(KV_T); + padParams.isPad = true; + padParams.leftPadding = 0; + padParams.rightPadding = combineDimAlign - combineDim; + padParams.paddingValue = 0; + DataCopyPad(kvInUb[startRow * 576], keyGm[keyOffset * 512], intriParams, padParams); + + intriParams.blockLen = constInfo.sparseBlockSize * constInfo.dSizeRope *sizeof(KV_T); + intriParams.dstStride = 512 / BUFFER_SIZE_BYTE_32B; + DataCopyPad(kvInUb[startRow * 576 + 512], keyRopeGm[keyOffset * 64], intriParams, padParams); // combineDimAlign +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline +uint32_t SFAVectorService::CopyInKvSparse( + LocalTensor kvInUb, int64_t startRow, int64_t token0Idx, + int64_t token1Idx, const RunInfo &runInfo, ConstInfo &constInfo) +{ + int64_t keyOffset0 = GetkeyOffset(token0Idx, runInfo, constInfo); + int64_t keyOffset1 = GetkeyOffset(token1Idx, runInfo, constInfo); + if (unlikely(keyOffset0 < 0 && keyOffset1 < 0)) { + return 0; + } + int64_t blkTableSrcStride = + ((keyOffset0 > keyOffset1 ? (keyOffset0 - keyOffset1) : + (keyOffset1 - keyOffset0)) - constInfo.sparseBlockSize); + int64_t keySrcStride = blkTableSrcStride * constInfo.dSizeNope * sizeof(KV_T); + int64_t keyRopeSrcStride = blkTableSrcStride * constInfo.dSizeRope * sizeof(KV_T); + if (unlikely(keyOffset1 < 0)) { + CopyInSingleKv(kvInUb, startRow, keyOffset0, constInfo); + } else if (keySrcStride >= INT32_MAX || keySrcStride < 0 || constInfo.sparseBlockSize > 1) { + // stride溢出、stride为负数、s2超长等异常场景，还原成2条搬运指令 + CopyInSingleKv(kvInUb, startRow, keyOffset0, constInfo); + CopyInSingleKv(kvInUb, startRow + 1, keyOffset1, constInfo); + } else { + DataCopyExtParams intriParams; + intriParams.blockCount = (keyOffset0 >= 0) + (keyOffset1 >= 0); + intriParams.blockLen = constInfo.sparseBlockSize * constInfo.dSizeNope *sizeof(KV_T); + intriParams.dstStride = constInfo.dSizeRope * sizeof(KV_T) / BUFFER_SIZE_BYTE_32B; + intriParams.srcStride = keySrcStride; + DataCopyPadExtParams padParams; + + int64_t keyOffset = keyOffset0 > -1 ? keyOffset0 : keyOffset1; + if (keyOffset1 > -1 && keyOffset1 < keyOffset0) { + keyOffset = keyOffset1; + } + DataCopyPad(kvInUb[startRow * 576], keyGm[keyOffset * constInfo.dSizeNope], + intriParams, padParams); // combineDimAlign + + intriParams.blockLen = constInfo.sparseBlockSize * constInfo.dSizeRope *sizeof(KV_T); + intriParams.dstStride = constInfo.dSizeNope * sizeof(KV_T) / BUFFER_SIZE_BYTE_32B; + intriParams.srcStride = keyRopeSrcStride; + DataCopyPad(kvInUb[startRow * 576 + 512], keyRopeGm[keyOffset * constInfo.dSizeRope], + intriParams, padParams); // combineDimAlign + } + return (keyOffset0 > -1) + (keyOffset1 > -1); +} + +// fp8->fp32 +static constexpr MicroAPI::CastTrait castTraitFp8_1 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; +// fp8->fp32 +static constexpr MicroAPI::CastTrait castTraitFp8_2 = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; +// fp32->fp16 +static constexpr MicroAPI::CastTrait castTraitFp8_3 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT}; +// fp32->fp16 +static constexpr MicroAPI::CastTrait castTraitFp8_4 = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT}; +template +__simd_vf__ void CastScaleImpl(__ubuf__ float* ubDstAddr, __ubuf__ int8_t* ubSrcAddr, uint32_t dealRowCount) +{ + MicroAPI::RegTensor vScale0; + MicroAPI::RegTensor vScale1; + MicroAPI::RegTensor vScalebf16Res0; + MicroAPI::RegTensor vScalebf16Res1; + MicroAPI::RegTensor vScalefp32Res0; + MicroAPI::RegTensor vScalefp32Res1; + __ubuf__ int8_t* ubScaleSrcAddrTemp = ubSrcAddr; + __ubuf__ float* ubDstAddrTmp = ubDstAddr; + MicroAPI::MaskReg bf16TypeMaskAll = MicroAPI::CreateMask(); + MicroAPI::MaskReg fp32MaskAll = MicroAPI::CreateMask(); + for (uint16_t i = 0; i < static_cast(dealRowCount); i++) { + // load scale + MicroAPI::LoadAlign( + (MicroAPI::RegTensor&)vScale0, ubScaleSrcAddrTemp, 640); + + MicroAPI::Cast(vScalebf16Res0, vScale0, bf16TypeMaskAll); + MicroAPI::Cast(vScalefp32Res0, vScalebf16Res0, fp32MaskAll); + + MicroAPI::StoreAlign( + ubDstAddrTmp, vScalefp32Res0, 64, bf16TypeMaskAll); + } +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void SFAVectorService::CopyOutKvUb2Gm( + Buffer &v0ResGm, LocalTensor kvOutUb, + int64_t dealRow, int64_t s2StartIdx, const RunInfo &runInfo, ConstInfo &constInfo) +{ + GlobalTensor v0ResGmTensor = v0ResGm.template GetTensor(); + DataCopy(v0ResGmTensor[s2StartIdx * 576], kvOutUb, dealRow * 576); +} + +TEMPLATES_DEF_NO_DEFAULT +__aicore__ inline void SFAVectorService::CalSparseCalSize(const RunInfo &runInfo, ConstInfo &constInfo) +{ + if constexpr (IS_SPLIT_G) { + uint32_t aicIdx = constInfo.aivIdx >> 1U; + uint32_t v0S2SizeFirstCore = CeilDiv(runInfo.s2RealSize, 2); + uint32_t v0S2SizeSecondCore = runInfo.s2RealSize - v0S2SizeFirstCore; + if (aicIdx % 2U == 0) { + if (GetSubBlockIdx() == 0) { + sparseCalSize = CeilDiv(v0S2SizeFirstCore, 2); + sparseS2Start = 0; + } else { + sparseCalSize = v0S2SizeFirstCore - CeilDiv(v0S2SizeFirstCore, 2); + sparseS2Start = CeilDiv(v0S2SizeFirstCore, 2); + } + } else { + if (GetSubBlockIdx() == 0) { + sparseCalSize = CeilDiv(v0S2SizeSecondCore, 2); + sparseS2Start = v0S2SizeFirstCore; + } else { + sparseCalSize = v0S2SizeSecondCore - CeilDiv(v0S2SizeSecondCore, 2); + sparseS2Start = v0S2SizeFirstCore + CeilDiv(v0S2SizeSecondCore, 2); + } + } + sparseS2End = sparseS2Start + sparseCalSize; + } else { + int64_t s2PerVecLoop = 2LL; + int64_t vecNum = 2LL; + int64_t s2Loops = CeilDiv(CeilDiv(runInfo.s2RealSize, vecNum), s2PerVecLoop); + sparseS2Start = GetSubBlockIdx() == 0 ? 0 : Min(s2Loops * s2PerVecLoop, runInfo.s2RealSize); + sparseS2End = GetSubBlockIdx() == 0 ? Min(s2Loops * s2PerVecLoop, runInfo.s2RealSize) : runInfo.s2RealSize; + sparseCalSize = sparseS2End - sparseS2Start; + } +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void SFAVectorService::ProcessVec0( + Buffer &outputL1, + Buffer &v0ResGm, + const RunInfo &runInfo, ConstInfo &constInfo, int32_t startPos) +{ + blockSize = constInfo.oriBlockSize; + maxBlockNumPerBatch = constInfo.oriMaxBlockNumPerBatch; + + CalSparseCalSize(runInfo, constInfo); + ProcessSparseKv(outputL1, v0ResGm, runInfo, constInfo, startPos); + if constexpr (IS_SPLIT_G) { + CrossCoreSetFlag<0, PIPE_MTE3>(15); + CrossCoreWaitFlag<0, PIPE_MTE3>(15); + } + outputL1.SetCrossCore(); + v0ResGm.SetCrossCore(); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void SFAVectorService::ProcessSparseKv( + Buffer &outputL1, + Buffer &v0ResGm, + const RunInfo &runInfo, ConstInfo &constInfo, int32_t startPos) +{ + if (sparseCalSize == 0) { + return; + } + bool meetEnd = false; + int64_t s2Start = sparseS2Start; + int64_t s2 = sparseS2Start; + int64_t token0Idx; + int64_t token1Idx; // 拷贝进入的两个token的index + // 处理一个s2的base块 + uint32_t pingPong = 0; + while ((s2 < sparseS2End) && !meetEnd) { // 拷贝到s2End或者遇到-1 + int64_t dealRow = 0; + // 1、copy kv in, gm ->ub + LocalTensor stage0OutUb = this->stage0OutBuf[pingPong].template Get(); + WaitFlag(mte3ToMte2[pingPong]); + while (dealRow < Min(16, sparseCalSize) && s2 < sparseS2End) { // 拷贝满16行或者遇到-1 + GetRealCmpS2Idx(token0Idx, token1Idx, s2, runInfo, constInfo); + s2 += 2; // 每次搬运2行 + if (token0Idx == -1 && token1Idx == -1) { + meetEnd = true; + break; + } + dealRow += CopyInKvSparse(stage0OutUb, dealRow, token0Idx, token1Idx, runInfo, constInfo); + if (token1Idx == -1) { + meetEnd = true; + break; + } + } + if (dealRow == 0) { + SetFlag(mte3ToMte2[pingPong]); + pingPong ^= 1; + return; + } + SetFlag(mte2ToMte3[pingPong]); + WaitFlag(mte2ToMte3[pingPong]); + // 2、copy kv out, ub -> l1 + CopyOutKvUb2Gm(v0ResGm, stage0OutUb, dealRow, s2Start, runInfo, constInfo); + SetFlag(mte3ToMte2[pingPong]); + s2Start += dealRow; + pingPong ^= 1; + } +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void SFAVectorService::ProcessVec1( + Buffer &outputBuf, + Buffer &bmm1ResBuf, RunInfo &runInfo, + ConstInfo &constInfo) +{ + bmm1ResBuf.WaitCrossCore(); + LocalTensor sumUb = this->softmaxSumBuf[runInfo.multiCoreIdxMod2].template Get(); + LocalTensor maxUb = this->softmaxMaxBuf[runInfo.multiCoreIdxMod2].template Get(); + LocalTensor expUb = this->softmaxExpBuf[runInfo.taskIdMod2].template Get(); + int64_t stage1Offset = runInfo.taskIdMod2; + auto stage1CastTensor = this->stage1OutQue[stage1Offset].template AllocTensor(); + + LocalTensor apiTmpBuffer = this->commonTBuf.template Get(); + LocalTensor mmRes = bmm1ResBuf.template GetTensor(); + + // loopCount = 0 但传入sinks时走update分支，maxUb通过sinks初始化，sumUb初始化为1.0 + if (runInfo.s2LoopCount == 0) { // sink 丢失首token信息，sink会增加首token信息，维度是n1 + if (likely(runInfo.s2RealSize == 128)) { // s2RealSize等于128分档, VF内常量化减少if判断 + ProcessVec1Vf( + stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize, runInfo.s2RealSize, + static_cast(constInfo.softmaxScale), negativeFloatScalar); + } else if (runInfo.s2RealSize <= 64) { // s2RealSize小于等于64分档, VF内常量化减少if判断 + ProcessVec1Vf( + stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize, + runInfo.s2RealSize, // 实际的计算有效元素， + static_cast(constInfo.softmaxScale), negativeFloatScalar); + } else if (runInfo.s2RealSize < 128) { // s2RealSize小于128分档, VF内常量化减少if判断 + ProcessVec1Vf( + stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize, runInfo.s2RealSize, + static_cast(constInfo.softmaxScale), negativeFloatScalar); + } + } else { + if (likely(runInfo.s2RealSize == 128)) { // s2RealSize等于128分档, VF内常量化减少if判断 + ProcessVec1Vf( + stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize, runInfo.s2RealSize, + static_cast(constInfo.softmaxScale), negativeFloatScalar); + } else if (runInfo.s2RealSize <= 64) { // s2RealSize小于等于64分档, VF内常量化减少if判断 + ProcessVec1Vf( + stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize, runInfo.s2RealSize, + static_cast(constInfo.softmaxScale), negativeFloatScalar); + } else if (runInfo.s2RealSize < 128) { // s2RealSize小于128分档, VF内常量化减少if判断 + ProcessVec1Vf( + stage1CastTensor, mmRes, sumUb, maxUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize, runInfo.s2RealSize, + static_cast(constInfo.softmaxScale), negativeFloatScalar); + } + } + bmm1ResBuf.SetCrossCore(); + + // ===================DataCopy to L1 ==================== + this->stage1OutQue[stage1Offset].template EnQue(stage1CastTensor); + this->stage1OutQue[stage1Offset].template DeQue(); + + LocalTensor mm2AL1Tensor = outputBuf.GetTensor(s2BaseSize * constInfo.dSizeV); + if (likely(runInfo.halfMRealSize != 0)) { + DataCopy(mm2AL1Tensor[constInfo.subBlockIdx * \ + (BLOCK_BYTE / sizeof(Q_T)) * (runInfo.mRealSize - runInfo.halfMRealSize)], + stage1CastTensor, {s2BaseSize / 16, (uint16_t)runInfo.halfMRealSize, + (uint16_t)(vec1Srcstride - runInfo.halfMRealSize), + (uint16_t)(Align16Func(runInfo.mRealSize) - runInfo.halfMRealSize)}); + } + + this->stage1OutQue[stage1Offset].template FreeTensor(stage1CastTensor); + + outputBuf.SetCrossCore(); + if (runInfo.s2LoopCount != 0) { + SFAUpdateExpSumAndExpMax(sumUb, maxUb, expUb, sumUb, maxUb, apiTmpBuffer, runInfo.halfMRealSize); + } + if (constInfo.returnSoftmaxLse && runInfo.s2LoopCount == runInfo.s2LoopLimit) { + CopyFALseToGm(runInfo, constInfo); + } +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void SFAVectorService::ProcessVec2( + Buffer &bmm2ResBuf, RunInfo &runInfo, + ConstInfo &constInfo) +{ + bmm2ResBuf.WaitCrossCore(); + if (unlikely(runInfo.vec2MBaseSize == 0)) { + bmm2ResBuf.SetCrossCore(); + return; + } + + runInfo.vec2S1RealSize = runInfo.vec2S1BaseSize; + runInfo.vec2MRealSize = runInfo.vec2MBaseSize; + int64_t vec2CalcSize = runInfo.vec2MRealSize * dTemplateAlign64; + + LocalTensor vec2ResUb = this->stage2OutBuf.template Get(); + LocalTensor mmRes = bmm2ResBuf.template GetTensor(); + + WaitFlag(mte3ToVAttnOutId); + if (unlikely(runInfo.s2LoopCount == 0)) { + DataCopy(vec2ResUb, mmRes, vec2CalcSize); + } else { + LocalTensor expUb = softmaxExpBuf[runInfo.taskIdMod2].template Get(); + if (runInfo.s2LoopCount < runInfo.s2LoopLimit) { + FlashUpdateNew( + vec2ResUb, mmRes, vec2ResUb, expUb, expUb, runInfo.vec2MRealSize, dTemplateAlign64, 1.0, 1.0); + } else { + LocalTensor sumUb = this->softmaxSumBuf[runInfo.multiCoreIdxMod2].template Get(); + FlashUpdateLastNew( + vec2ResUb, mmRes, vec2ResUb, expUb, expUb, sumUb, runInfo.vec2MRealSize, dTemplateAlign64, 1.0, 1.0); + } + } + + bmm2ResBuf.SetCrossCore(); + if (runInfo.s2LoopCount == runInfo.s2LoopLimit) { + if (unlikely(runInfo.s2LoopCount == 0)) { + LocalTensor sumUb = this->softmaxSumBuf[runInfo.multiCoreIdxMod2].template Get(); + LastDivNew( + vec2ResUb, vec2ResUb, sumUb, runInfo.vec2MRealSize, dTemplateAlign64, 1.0); + } + + this->CopyOutAttentionOut(runInfo, constInfo, vec2ResUb, 0, vec2CalcSize); + } + SetFlag(mte3ToVAttnOutId); +} + +TEMPLATES_DEF_NO_DEFAULT +template +__aicore__ inline void SFAVectorService::Bmm2DataCopyOut (RunInfo &runInfo, ConstInfo &constInfo, + LocalTensor &vec2ResUb, int64_t vec2S1Idx, int64_t vec2CalcSize) +{ + LocalTensor attenOut; + int64_t dSizeAligned64 = (int64_t)dTemplateAlign64; + + attenOut.SetAddr(vec2ResUb.address_); + Cast(attenOut, vec2ResUb, RoundMode::CAST_ROUND, vec2CalcSize); + SetFlag(vToMte3AttnOutId); + WaitFlag(vToMte3AttnOutId); + + DataCopyExtParams dataCopyParams; + dataCopyParams.blockLen = constInfo.dSizeV * sizeof(OUTPUT_T); + dataCopyParams.srcStride = (dSizeAligned64 - constInfo.dSizeV) >> 4; // 以32B为单位偏移，bf16类型即偏移16个数，右移4 + dataCopyParams.dstStride = constInfo.attentionOutStride; + dataCopyParams.blockCount = runInfo.vec2MRealSize; + + DataCopyPad(this->attentionOutGm[runInfo.attentionOutOffset], attenOut, dataCopyParams); +} + +TEMPLATES_DEF_NO_DEFAULT +template +__aicore__ inline void SFAVectorService::CopyOutAttentionOut( + RunInfo &runInfo, ConstInfo &constInfo, LocalTensor &vec2ResUb, int64_t vec2S1Idx, int64_t vec2CalcSize) +{ + this->Bmm2DataCopyOut(runInfo, constInfo, vec2ResUb, vec2S1Idx, vec2CalcSize); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline +void SFAVectorService::InitOutputSingleCore(ConstInfo &constInfo) +{ + uint32_t coreNum = GetBlockNum(); + uint64_t totalOutputSize = 0; + // n2 = 1, n1 = gn2 = gSize + if constexpr (LAYOUT_T == SFA_LAYOUT::BSND) { + totalOutputSize = constInfo.bSize * constInfo.gSize * constInfo.s1Size * constInfo.dSizeV; + } else if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + totalOutputSize = constInfo.s1Size * constInfo.gSize * constInfo.dSizeV; + } + + if (coreNum != 0) { + uint64_t singleCoreSize = (totalOutputSize + (CV_RATIO * coreNum) - 1) / (CV_RATIO * coreNum); + uint64_t tailSize = totalOutputSize - constInfo.aivIdx * singleCoreSize; + uint64_t singleInitOutputSize = tailSize < singleCoreSize ? tailSize : singleCoreSize; + if (constInfo.aivIdx * singleCoreSize < totalOutputSize && singleInitOutputSize > 0) { + matmul::InitOutput( + this->attentionOutGm[constInfo.aivIdx * singleCoreSize], singleInitOutputSize, 0); + } + } + + if (constInfo.returnSoftmaxLse) { + uint64_t totalReturnSoftmaxSize = 0; + if constexpr (LAYOUT_T == SFA_LAYOUT::BSND) { + totalReturnSoftmaxSize = constInfo.bSize * constInfo.n2Size * constInfo.s1Size * constInfo.gSize; + } else if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + totalReturnSoftmaxSize = constInfo.n2Size * constInfo.s1Size * constInfo.gSize; // (N2,T1,G) + } + if (coreNum != 0 && totalReturnSoftmaxSize > 0) { + uint64_t singleCoreSoftmaxSize = (totalReturnSoftmaxSize + (CV_RATIO * coreNum) - 1) / (CV_RATIO * coreNum); + uint64_t tailSoftmaxSize = totalReturnSoftmaxSize - constInfo.aivIdx * singleCoreSoftmaxSize; + uint64_t singleInitSoftmaxSize = tailSoftmaxSize < singleCoreSoftmaxSize ? + tailSoftmaxSize : singleCoreSoftmaxSize; + if (constInfo.aivIdx * singleCoreSoftmaxSize < totalReturnSoftmaxSize && singleInitSoftmaxSize > 0) { + matmul::InitOutput(this->softmaxSumGm[constInfo.aivIdx * singleCoreSoftmaxSize], singleInitSoftmaxSize, 0); + matmul::InitOutput(this->softmaxMaxGm[constInfo.aivIdx * singleCoreSoftmaxSize], singleInitSoftmaxSize, 0); + } + } + } + SyncAll(); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline +void SFAVectorService::CleanOutput(__gm__ uint8_t *attentionOut, __gm__ uint8_t *softmaxMax, + __gm__ uint8_t *softmaxSum, ConstInfo &constInfo) +{ + if ASCEND_IS_AIV { + this->attentionOutGm.SetGlobalBuffer((__gm__ OUTPUT_T *)attentionOut); + this->softmaxSumGm.SetGlobalBuffer((__gm__ float *)(softmaxSum)); + this->softmaxMaxGm.SetGlobalBuffer((__gm__ float *)(softmaxMax)); + if (constInfo.needInit == 1) { + InitOutputSingleCore(constInfo); + } + } +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline +void SFAVectorService::InitGlobalBuffer(__gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *keyRope, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *blockTable, __gm__ uint8_t *softmaxMax, + __gm__ uint8_t *softmaxSum) +{ + keyGm.SetGlobalBuffer((__gm__ KV_T *)(key)); + if constexpr (isPa) { + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable);; + } + sparseIndicesGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices); + keyRopeGm.SetGlobalBuffer((__gm__ KV_T *)(keyRope)); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void SFAVectorService::SoftmaxInitBuffer() +{ + constexpr uint32_t softmaxBufSize = 256; // VF单次操作256Byte + tPipe->InitBuffer(softmaxSumBuf[0], softmaxBufSize); + tPipe->InitBuffer(softmaxSumBuf[1], softmaxBufSize); + tPipe->InitBuffer(softmaxMaxBuf[0], softmaxBufSize); + tPipe->InitBuffer(softmaxMaxBuf[1], softmaxBufSize); + tPipe->InitBuffer(softmaxExpBuf[0], softmaxBufSize); + tPipe->InitBuffer(softmaxExpBuf[1], softmaxBufSize); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline +void SFAVectorService::CopyFALseToGm(RunInfo &runInfo, ConstInfo &constInfo) +{ + LocalTensor sumUb = this->softmaxSumBuf[runInfo.multiCoreIdxMod2].template Get(); + LocalTensor maxUb = this->softmaxMaxBuf[runInfo.multiCoreIdxMod2].template Get(); + + size_t alignedSize = (sizeof(float) * runInfo.halfMRealSize + 31) / 32 * 32 / sizeof(float); + + int64_t lseOffset = runInfo.softmaxLseOffset; + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = 1; + dataCopyParams.blockLen = sizeof(float) * runInfo.halfMRealSize; + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = 0; + + // 拷贝 softmaxMaxUb -> GM + WaitFlag(mte3ToVLseOutId); + DataCopy(lseUb, maxUb, alignedSize); + SetFlag(vToMte3LseOutId); + WaitFlag(vToMte3LseOutId); + DataCopyPad(this->softmaxMaxGm[lseOffset], lseUb, dataCopyParams); + SetFlag(mte3ToVLseOutId); + + // 拷贝 softmaxSumUb -> GM + WaitFlag(mte3ToVLseOutId); + DataCopy(lseUb, sumUb, alignedSize); + SetFlag(vToMte3LseOutId); + WaitFlag(vToMte3LseOutId); + DataCopyPad(this->softmaxSumGm[lseOffset], lseUb, dataCopyParams); + SetFlag(mte3ToVLseOutId); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void SFAVectorService::InitSinksBuffer(ConstInfo &constInfo) +{ + LocalTensor sinksUb = this->sinksBuf.template Get(); + const uint32_t maxN = constInfo.gSize; // N最大支持128, sink shape是[N] + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = 1U; + dataCopyParams.blockLen = maxN * sizeof(T); + dataCopyParams.srcStride = 0U; + dataCopyParams.dstStride = 0U; + DataCopyPadExtParams padParams; + DataCopyPad(sinksUb, this->sinksGm, dataCopyParams, padParams); + SetFlag(SYNC_SINKS_BUF_FLAG); + WaitFlag(SYNC_SINKS_BUF_FLAG); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline +void SFAVectorService::InitLocalBuffer(TPipe *pipe, ConstInfo &constInfo) +{ + SoftmaxInitBuffer(); + + tPipe->InitBuffer(commonTBuf, 512); // commonTBuf内存申请512B + tPipe->InitBuffer(sinksBuf, 512); // sinksBuf内存申请512B + tPipe->InitBuffer(lseBuf, 512); // lseBuf内存申请512B + lseUb = this->lseBuf.template Get(); + + tPipe->InitBuffer(stage0OutBuf[0], 576 * 16 * sizeof(KV_T)); + tPipe->InitBuffer(stage0OutBuf[1], 576 * 16 * sizeof(KV_T)); + + tPipe->InitBuffer(stage1OutQue[0], 1, vec1Srcstride * s2BaseSize * sizeof(Q_T)); + tPipe->InitBuffer(stage1OutQue[1], 1, vec1Srcstride * s2BaseSize * sizeof(Q_T)); + tPipe->InitBuffer(stage2OutBuf, (s1BaseSize / CV_RATIO) * dTemplateAlign64 * sizeof(T)); + + mte3ToVAttnOutId = GetTPipePtr()->AllocEventID(); + mte3ToVLseOutId = GetTPipePtr()->AllocEventID(); + SetFlag(mte3ToVAttnOutId); + SetFlag(mte3ToVLseOutId); + + vToMte3AttnOutId = GetTPipePtr()->AllocEventID(); + vToMte3LseOutId = GetTPipePtr()->AllocEventID(); + + mte2ToV = GetTPipePtr()->AllocEventID(); + mte3ToMte2[0] = GetTPipePtr()->AllocEventID(); + mte3ToMte2[1] = GetTPipePtr()->AllocEventID(); + SetFlag(mte3ToMte2[0]); + SetFlag(mte3ToMte2[1]); + mte2ToMte3[0] = GetTPipePtr()->AllocEventID(); + mte2ToMte3[1] = GetTPipePtr()->AllocEventID(); +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void SFAVectorService::InitCubeVecSharedParams( + CVSharedParams &sharedParams, int32_t aicIdx, uint8_t subBlockIdx) +{ + // TODO参数整改 + auto &sparseAttnSharedkvBaseParams = this->tilingData->baseParams; + sharedParams.bSize = sparseAttnSharedkvBaseParams.batchSize; + sharedParams.n2Size = 1; + sharedParams.gSize = sparseAttnSharedkvBaseParams.nNumOfQInOneGroup; + sharedParams.s1Size = sparseAttnSharedkvBaseParams.qSeqSize; + sharedParams.s2Size = sparseAttnSharedkvBaseParams.seqSize; + sharedParams.sparseBlockCount = sparseAttnSharedkvBaseParams.sparseBlockCount; + sharedParams.cmpRatio = 1; // 走sparse， 但不压缩 + sharedParams.oriMaskMode = sparseAttnSharedkvBaseParams.sparseMode; + sharedParams.oriWinLeft = -1; + sharedParams.oriWinRight = 0; + sharedParams.layoutType = sparseAttnSharedkvBaseParams.outputLayout; + sharedParams.dSizeRope = 64; + sharedParams.softmaxScale = sparseAttnSharedkvBaseParams.scaleValue; + sharedParams.dSize = 512; + sharedParams.dSizeVInput = 512; + sharedParams.usedCoreNum = this->tilingData->singleCoreParams.usedCoreNum; + + // pageAttention, rope在C侧搬运时使用 + if constexpr (isPa) { + sharedParams.oriBlockSize = sparseAttnSharedkvBaseParams.blockSize; + sharedParams.oriMaxBlockNumPerBatch = sparseAttnSharedkvBaseParams.maxBlockNumPerBatch; + } + + // actQ->TND, actKV pa场景任意layout均有 + sharedParams.isActualSeqLengthsNull = sparseAttnSharedkvBaseParams.isActualLenDimsNull; + sharedParams.isActualSeqLengthsKVNull = sparseAttnSharedkvBaseParams.isActualLenDimsKVNull; + sharedParams.returnSoftmaxLse = sparseAttnSharedkvBaseParams.returnSoftmaxLse; + sharedParams.needInit = 0; + for (uint32_t bIdx = 0; bIdx < sharedParams.bSize; bIdx++) { + int64_t s2Size; + if constexpr (KV_LAYOUT_T == SFA_LAYOUT::TND) { + s2Size = bIdx == 0 ? actualSeqLengthsKVGm.GetValue(bIdx) : \ + actualSeqLengthsKVGm.GetValue(bIdx) - actualSeqLengthsKVGm.GetValue(bIdx - 1); + } else { + if (sharedParams.isActualSeqLengthsKVNull) { + s2Size = sharedParams.s2Size; + } else { + s2Size = actualSeqLengthsKVGm.GetValue(bIdx); + } + } + int64_t s1Size; + if constexpr (LAYOUT_T == SFA_LAYOUT::TND) { + s1Size = bIdx == 0 ? cuSeqlensQGm.GetValue(bIdx) : \ + cuSeqlensQGm.GetValue(bIdx) - cuSeqlensQGm.GetValue(bIdx - 1); + } else { + if (sharedParams.isActualSeqLengthsNull) { + s1Size = sharedParams.s1Size; + } else { + s1Size = cuSeqlensQGm.GetValue(bIdx); + } + } + if (s1Size > s2Size || (LAYOUT_T == SFA_LAYOUT::BSND && s1Size < sharedParams.s1Size)) { + sharedParams.needInit = 1; + break; + } + } + + if ASCEND_IS_AIV { + if (subBlockIdx == 0) { + auto tempTilingSSbuf = reinterpret_cast<__ssbuf__ uint32_t*>(0); // 从ssbuf的0地址开始拷贝 + auto tempTiling = reinterpret_cast(&sharedParams); +#pragma unroll + for (int i = 0; i < sizeof(CVSharedParams) / sizeof(uint32_t); ++i, ++tempTilingSSbuf, ++tempTiling) { + *tempTilingSSbuf = *tempTiling; + } + CrossCoreSetFlag(15); + } + } +} + +TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void SFAVectorService::GetExtremeValue( + T &negativeScalar) +{ + uint32_t tmp1 = NEGATIVE_MIN_VALUE_FP32; + negativeScalar = *((float *)&tmp1); +} + + +TEMPLATES_DEF class SFAVectorServiceDummy { +public: + __aicore__ inline SFAVectorServiceDummy() {}; + __aicore__ inline void CleanOutput(__gm__ uint8_t *attentionOut, __gm__ uint8_t *softmaxMax, + __gm__ uint8_t *softmaxSum, ConstInfo &constInfo) {} + __aicore__ inline void InitGlobalBuffer(__gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *keyRope, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *blockTable, __gm__ uint8_t *softmaxMax, + __gm__ uint8_t *softmaxSum) {} + __aicore__ inline void InitVecBlock(TPipe *pipe, const SparseFlashAttentionTilingDataMla *__restrict tiling, + CVSharedParams &sharedParams, int32_t aicIdx, uint8_t subBlockIdx, __gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengths) {}; + __aicore__ inline void InitLocalBuffer(TPipe *pipe, ConstInfo &constInfo) {} + __aicore__ inline void ProcessVec1(Buffer &outputBuf, + Buffer &bmm1ResBuf, RunInfo &runInfo, + ConstInfo &constInfo) {} + + using mm2ResPos = Buffer; + __aicore__ inline void ProcessVec2(mm2ResPos &bmm2ResBuf, RunInfo &runInfo, + ConstInfo &constInfo) {} +}; +} +#endif // SPARSE_FLASH_ATTENTION_SERVICE_VECTOR_MLA_H diff --git a/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/util_regbase.h b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/util_regbase.h new file mode 100644 index 000000000..c9b934814 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_kernel/arch35/util_regbase.h @@ -0,0 +1,265 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file util_regbase.h + * \brief + */ + +#ifndef UTIL_REGBASE_H +#define UTIL_REGBASE_H + +#include "util.h" + +using AscendC::TQue; +using AscendC::QuePosition; + +namespace regbaseutil { +constexpr int64_t MAX_PRE_NEXT_TOKENS = 0x7FFFFFFF; + +#define COMMON_RUN_PARAM \ + int64_t boIdx; \ + int64_t s1oIdx; \ + int64_t n2oIdx; \ + int64_t goIdx; \ + int64_t s2LoopEndIdx; /* S2方向的循环控制信息 souter层确定 */ \ + int64_t s2LineStartIdx = 0; /* S2方向按行的起始位置 */ \ + int64_t s2LineEndIdx; /* S2方向按行的结束位置 */ \ + int64_t s2CmpLineEndIdx; \ + /* cube视角的sOuter，在SAMEAB场景中cubeSOuterSize为两倍的 halfS1RealSize souter层确定 */ \ + uint32_t s1RealSize; \ + uint32_t halfS1RealSize; \ + uint32_t firstHalfS1RealSize; \ + uint32_t mRealSize; \ + uint32_t halfMRealSize; \ + uint32_t firstHalfMRealSize; \ + int64_t attentionOutOffset; /* attentionOut的offset souter层确定 */ \ + int32_t actualS1Size; /* Q的actualSeqLength */ \ + int32_t actualS2Size; /* KV的actualSeqLength */ \ + int64_t tensorQOffset; \ + int64_t tensorQRopeOffset; \ + int64_t qBOffset; \ + int64_t qRopeBOffset; + +struct RunParamStr { // 分核与切块需要使用到参数 + COMMON_RUN_PARAM; + /* 推理新增 */ + int64_t gs1LoopStartIdx; + int64_t gs1LoopEndIdx; + // BN循环生产的数据 + int64_t preTokensPerBatch = MAX_PRE_NEXT_TOKENS; // 左上顶点的pretoken + int64_t nextTokensPerBatch = MAX_PRE_NEXT_TOKENS; // 左上顶点的nexttoken + + // NBS1循环生产的数据 + int64_t sOuterOffset; // 单个S内 souter的 souterIdx * halfS1RealSize souter层确定 + int64_t cubeSOuterOffset; // 单个S内 souter的 souterIdx * halfS1RealSize souter层确定 + int64_t mOuterOffset; + int64_t cubeMOuterOffset; + + // lse 输出offset + int64_t softmaxLseOffset; // souter层确定 + + int64_t qSNumInOneBlock; + int64_t kvLoopEndIdx; +}; + +#define COMMON_RUN_INFO \ + int64_t s2StartIdx; /* s2的起始位置，sparse场景下可能不是0 */ \ + int64_t s2EndIdx; \ + int64_t s2LoopCount; /* s2循环当前的循环index */ \ + int64_t s2LoopLimit; \ + int64_t s1oIdx = 0; /* s1轴的index */ \ + int64_t loop = 0; /* for v0 perload loop */ \ + int64_t boIdx = 0; /* b轴的index */ \ + int64_t n2oIdx = 0; /* n2轴的index */ \ + int64_t goIdx = 0; /* g轴的index */ \ + int32_t s1RealSize; \ + int32_t halfS1RealSize; /* vector侧实际的s1基本块大小，如果Cube基本块=128，那么halfS1RealSize=64 */ \ + int32_t firstHalfS1RealSize; /* 当s1RealSize不是2的整数倍时，v0比v1少计算一行，计算subblock偏移的时候需要使用v0的s1 size */ \ + int32_t mRealSize; \ + int32_t halfMRealSize; \ + int32_t firstHalfMRealSize; \ + int32_t s2RealSize; /* s2方向基本块的真实长度 */ \ + int64_t s2AlignedSize; /* s2方向基本块对齐到16之后的长度 */ \ + int32_t vec2S1BaseSize; /* vector2侧开循环之后，经过切分的S1大小，例如把64切分成两份32 */ \ + int32_t vec2S1RealSize; /* vector2侧开循环之后，经过切分的S1的尾块大小，例如把63切分成两份32和31，第二份的实际大小是31 */ \ + int32_t vec2MBaseSize; \ + int32_t vec2MRealSize; \ + int64_t taskId; \ + int64_t multiCoreInnerIdx = 0; \ + int64_t attentionOutOffset; \ + int32_t actualS1Size; /* 非TND场景=总s1Size, Tnd场景下当前batch对应的s1 */ \ + int32_t actualS2Size; /* 非TND场景=总s2Size, Tnd场景下当前batch对应的s2 */ \ + int64_t preTokensPerBatch; /* vector2 左上顶点的pretoken */ \ + int64_t nextTokensPerBatch; /* vector2 左上顶点的nexttoken */ \ + uint8_t taskIdMod2; \ + uint8_t taskIdMod3; \ + uint8_t multiCoreIdxMod2 = 0; \ + uint8_t multiCoreIdxMod3 = 0; \ + int64_t sOuterOffset; \ + int64_t mOuterOffset; \ + int64_t queryOffset; \ + int64_t queryRopeOffset + +struct RunInfo { + COMMON_RUN_INFO; + // 推理新增 + // lse 输出offset + int64_t softmaxLseOffset; + + int64_t qSNumInOneBlock; + int64_t kvLoopEndIdx; +}; + +#define COMMON_CONST_INFO \ + /* 全局的基本块信息 */ \ + uint32_t bSize; \ + uint32_t needInit; \ + uint32_t s1BaseSize; \ + uint32_t s2BaseSize; \ + int64_t dSize; /* query d 512 */ \ + int64_t dSizeV; /* key d 512 */ \ + int64_t dSizeVInput; /* key inpue d 656 = rope + nope + scale + pad */ \ + int64_t dSizeNope; /* key nope d 448 */ \ + int64_t dSizeRope; /* key rope d 64 */ \ + int64_t tileSize; /* 64 */ \ + int64_t sparseMode; \ + int64_t gSize; /* g轴的大小 */ \ + int64_t n2Size; \ + int64_t s1Size; /* s1总大小 */ \ + int64_t s2Size; /* s2总大小 */ \ + /* 轴的乘积 */ \ + int64_t s1D; \ + int64_t gS1D; \ + int64_t n2GS1D; \ + int64_t s2D; \ + int64_t n2S2D; \ + int64_t s1Dv; \ + int64_t gS1Dv; \ + int64_t n2GS1Dv; \ + int64_t s2Dv; \ + int64_t n2S2Dv; \ + int64_t s1S2; \ + int64_t gS1; \ + int64_t gD; \ + int64_t n2D; \ + int64_t bN2D; \ + int64_t gDv; \ + int64_t n2Dv; \ + int64_t bN2Dv; \ + int64_t n2G; \ + int64_t n2GD; \ + int64_t bN2GD; \ + int64_t n2GDv; \ + int64_t bN2GDv; \ + int64_t gS2; \ + int64_t s1Dr; \ + int64_t gS1Dr; \ + int64_t n2GS1Dr; \ + int64_t s2Dr; \ + int64_t n2S2Dr; \ + int64_t gDr; \ + int64_t n2Dr; \ + int64_t bN2Dr; \ + int64_t n2GDr; \ + int64_t bN2GDr; \ + int32_t s2BaseN2D; \ + int32_t s1BaseN2GD; \ + int64_t s2BaseBN2D; \ + int64_t s1BaseBN2GD; \ + int32_t s1BaseD; \ + int32_t s2BaseD; \ + int64_t s2BaseN2Dv; \ + int64_t s2BaseBN2Dv; \ + int64_t s1BaseN2GDv; \ + int64_t s1BaseBN2GDv; \ + int32_t s1BaseDv; \ + int32_t s2BaseDv; \ + bool returnSoftmaxLse; \ + /* matmul跳读参数 */ \ + int64_t mm1Ka; \ + /* dq 或者attentionOut的Stride */ \ + int64_t attentionOutStride; \ + uint32_t aivIdx; \ + uint8_t layoutType; \ + uint8_t subBlockIdx;\ + /* 分核相关 */ \ + uint32_t s2Start; \ + uint32_t s2End; \ + uint32_t bN2Start; \ + uint32_t bN2End; \ + uint32_t gS1Start; \ + uint32_t gS1End + +#define INFER_CONST_INFO \ + /* 推理 */ \ + bool isActualLenDimsNull; /* 判断是否有actualseq */ \ + bool isActualLenDimsKVNull; /* 判断是否有actualseq_kv */ \ + bool isSoftmaxLseEnable; \ + bool rsvd1; \ + uint32_t sparseBlockCount; \ + uint32_t actualSeqLenSize; /* 用户输入的actualseq的长度 */ \ + uint32_t actualSeqLenKVSize; /* 用户输入的actualseq_kv的长度 */ \ + /* service mm1 mm2 pageAttention */ \ + uint32_t oriBlockSize; \ + uint32_t cmpBlockSize; \ + uint32_t paLayoutType; \ + uint32_t oriMaxBlockNumPerBatch; \ + uint32_t cmpMaxBlockNumPerBatch; \ + int32_t oriWinLeft; \ + int32_t oriWinRight; \ + uint32_t sparseBlockSize; \ + uint32_t cmpRatio; \ + float softmaxScale + +#define CV_SHARED_PARAMS \ + /* base params */ \ + uint32_t s1BaseSize; \ + uint32_t s2BaseSize; \ + uint32_t bSize; \ + uint32_t n2Size; \ + uint32_t gSize; \ + uint32_t s1Size; \ + uint32_t s2Size; \ + uint32_t dSize : 10; \ + int64_t dSizeVInput : 12; \ + uint32_t needInit : 4; \ + uint32_t layoutType : 4; \ + uint32_t isActualSeqLengthsNull : 1; \ + uint32_t isActualSeqLengthsKVNull : 1; \ + uint32_t sparseBlockCount; \ + float softmaxScale; \ + uint32_t cmpRatio : 9; \ + uint32_t dSizeRope : 11; \ + uint32_t oriMaskMode : 6; \ + uint32_t cmpMaskMode : 6; \ + int32_t oriWinLeft; \ + int32_t oriWinRight; \ + uint32_t tileSize : 8; \ + /* pa params */ \ + uint32_t oriBlockSize : 12; \ + uint32_t cmpBlockSize : 12; \ + uint32_t oriMaxBlockNumPerBatch; \ + uint32_t cmpMaxBlockNumPerBatch; \ + uint32_t usedCoreNum; \ + bool returnSoftmaxLse + +struct ConstInfo { + COMMON_CONST_INFO; + INFER_CONST_INFO; +}; + +/* only support b32 or b64 */ +struct CVSharedParams { + CV_SHARED_PARAMS; +}; +} + +#endif // UTIL_REGBASE_H diff --git a/csrc/ascend/attention/sparse_flash_attention/op_kernel/sparse_flash_attention.cpp b/csrc/ascend/attention/sparse_flash_attention/op_kernel/sparse_flash_attention.cpp new file mode 100644 index 000000000..70479ceff --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_kernel/sparse_flash_attention.cpp @@ -0,0 +1,102 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + + /*! + * \file sparse_flash_attention.cpp + * \brief + */ + +#include "kernel_operator.h" +#include "sparse_flash_attention_template_tiling_key.h" +#if (__CCE_AICORE__ == 310) +#include "arch35/sparse_flash_attention_kernel_mla.h" +#else +#include "arch22/sparse_flash_attention_kernel_mla.h" +#endif + +using namespace AscendC; + +#if (__CCE_AICORE__ == 310) +#if defined(__DAV_C310_CUBE__) +#define SFA_OP_IMPL(templateClass, tilingdataClass, ...) \ + do { \ + using CubeBlockType = typename std::conditional, BaseApi::SFAMatmulServiceDummy<__VA_ARGS__>>::type; \ + using VecBlockType = typename std::conditional, BaseApi::SFAVectorService<__VA_ARGS__>>::type; \ + templateClass op; \ + GET_TILING_DATA_WITH_STRUCT(tilingdataClass, tiling_data_in, tiling); \ + op.Init(query, key, value, sparseIndices, actualSeqLengthsQuery, actualSeqLengthsKV, \ + blocktable, queryRope, keyRope, attentionOut, softmaxMax, softmaxSum, user, nullptr, tiling, &tPipe); \ + op.Process(); \ + } while (0) +#else +#define SFA_OP_IMPL(templateClass, tilingdataClass, ...) \ + do { \ + using CubeBlockType = typename std::conditional, BaseApi::SFAMatmulServiceDummy<__VA_ARGS__>>::type; \ + using VecBlockType = typename std::conditional, BaseApi::SFAVectorService<__VA_ARGS__>>::type; \ + templateClass op; \ + GET_TILING_DATA_WITH_STRUCT(tilingdataClass, tiling_data_in, tiling); \ + const tilingdataClass *__restrict tilingData = &tiling_data_in; \ + op.Init(query, key, value, sparseIndices, actualSeqLengthsQuery, actualSeqLengthsKV, \ + blocktable, queryRope, keyRope, attentionOut, softmaxMax, softmaxSum, user, tilingData, tiling, &tPipe); \ + op.Process(); \ + } while (0) +#endif +#else +#define SFA_OP_IMPL(templateClass, tilingdataClass, ...) \ + do { \ + templateClass> op; \ + GET_TILING_DATA_WITH_STRUCT(tilingdataClass, tiling_data_in, tiling); \ + const tilingdataClass *__restrict tiling_data = &tiling_data_in; \ + op.Init(query, key, value, sparseIndices, actualSeqLengthsQuery, actualSeqLengthsKV, \ + blocktable, queryRope, keyRope, attentionOut, softmaxMax, softmaxSum, user, tiling_data, tiling, &tPipe); \ + op.Process(); \ + } while (0) +#endif + +template + __global__ __aicore__ void +sparse_flash_attention(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *value, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *blocktable, + __gm__ uint8_t *actualSeqLengthsQuery, __gm__ uint8_t *actualSeqLengthsKV, + __gm__ uint8_t* queryRope, __gm__ uint8_t* keyRope, + __gm__ uint8_t *attentionOut, __gm__ uint8_t* softmaxMax, __gm__ uint8_t* softmaxSum, + __gm__ uint8_t *workspace, __gm__ uint8_t *tiling) +{ + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + + TPipe tPipe; + __gm__ uint8_t *user = GetUserWorkspace(workspace); + +#if (__CCE_AICORE__ == 310) + if constexpr (ORIG_DTYPE_QUERY == DT_BF16 && ORIG_DTYPE_KEY == DT_BF16 && + ORIG_DTYPE_ATTENTION_OUT == DT_BF16) { + SFA_OP_IMPL(BaseApi::SparseFlashAttentionKernelMla, SparseFlashAttentionTilingDataMla, bfloat16_t, bfloat16_t, + float, bfloat16_t, FLASH_DECODE, PAGE_ATTENTION, static_cast(LAYOUT_T), + static_cast(KV_LAYOUT_T), static_cast(TEMPLATE_MODE), IS_SPLIT_G); + } else { + SFA_OP_IMPL(BaseApi::SparseFlashAttentionKernelMla, SparseFlashAttentionTilingDataMla, half, half, + float, half, FLASH_DECODE, PAGE_ATTENTION, static_cast(LAYOUT_T), + static_cast(KV_LAYOUT_T), static_cast(TEMPLATE_MODE), IS_SPLIT_G); + } +#else + if constexpr (ORIG_DTYPE_QUERY == DT_FLOAT16 && ORIG_DTYPE_KEY == DT_FLOAT16 && + ORIG_DTYPE_ATTENTION_OUT == DT_FLOAT16) { + SFA_OP_IMPL(SparseFlashAttentionMla, SparseFlashAttentionTilingDataMla, half, half, half, + FLASH_DECODE, static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), TEMPLATE_MODE); + } else { // bf16 + SFA_OP_IMPL(SparseFlashAttentionMla, SparseFlashAttentionTilingDataMla, bfloat16_t, bfloat16_t, bfloat16_t, + FLASH_DECODE, static_cast(LAYOUT_T), static_cast(KV_LAYOUT_T), TEMPLATE_MODE); + } +#endif +} diff --git a/csrc/ascend/attention/sparse_flash_attention/op_kernel/sparse_flash_attention_common.h b/csrc/ascend/attention/sparse_flash_attention/op_kernel/sparse_flash_attention_common.h new file mode 100644 index 000000000..a250499e7 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_kernel/sparse_flash_attention_common.h @@ -0,0 +1,202 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_flash_attention_common.h + * \brief + */ + +#ifndef SPARSE_FLASH_ATTENTION_COMMON_H +#define SPARSE_FLASH_ATTENTION_COMMON_H + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" + +using namespace AscendC; +// 将isCheckTiling设置为false, 输入输出的max&sum&exp的shape为(m, 1) +constexpr SoftmaxConfig SFA_SOFTMAX_FLASHV2_CFG_WITHOUT_BRC = {false, 0, 0, SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC}; + +enum class SFA_LAYOUT +{ + BSND = 0, + TND = 1, + PA_BSND = 2, +}; + +template +struct SFAType { + using queryType = Q_T; + using kvType = KV_T; + using outputType = OUT_T; + static constexpr bool flashDecode = FLASH_DECODE; + static constexpr SFA_LAYOUT layout = LAYOUT_T; + static constexpr SFA_LAYOUT kvLayout = KV_LAYOUT_T; + static constexpr int templateMode = TEMPLATE_MODE; + static constexpr bool pageAttention = (KV_LAYOUT_T == SFA_LAYOUT::PA_BSND); +}; + +// ================================Util functions================================== +template __aicore__ inline T SFAAlign(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd) * (rnd))); +} + +template __aicore__ inline T1 Min(T1 a, T2 b) +{ + return (a > b) ? (b) : (a); +} + +template __aicore__ inline size_t BlockAlign(size_t s) +{ + if constexpr (IsSameType::value) { + return (s + 63) / 64 * 64; + } + size_t n = (32 / sizeof(T)); + return (s + n - 1) / n * n; +} + +struct RunInfo { + uint32_t loop; + uint32_t bIdx; + uint32_t gIdx; + uint32_t s1Idx; + uint32_t s2Idx; + uint32_t bn2IdxInCurCore; + uint32_t curSInnerLoopTimes; + uint64_t tndBIdxOffsetForQ; + uint64_t tndBIdxOffsetForKV; + uint64_t tensorAOffset; + uint64_t tensorBOffset; + uint64_t tensorARopeOffset; + uint64_t tensorBRopeOffset; + uint64_t attenOutOffset; + uint64_t attenMaskOffset; + uint64_t topKBaseOffset; + uint32_t actualSingleProcessSInnerSize; + uint32_t actualSingleProcessSInnerSizeAlign; + bool isFirstSInnerLoop; + bool isChangeBatch; + uint32_t s2BatchOffset; + uint32_t gSize; + uint32_t s1Size; + uint32_t s2Size; + uint32_t mSize; + uint32_t mSizeV; + uint32_t mSizeVStart; + uint32_t tndIsS2SplitCore; + uint32_t tndCoreStartKVSplitPos; + bool isBmm2Output; + bool isValid = false; + + static constexpr uint32_t n2Idx = 0; + uint64_t actS1Size = 1; + uint64_t curActualSeqLenOri = 0ULL; + + uint32_t gS1Idx; + uint64_t actS2Size = 1; + uint32_t actMBaseSize; + bool isLastS2Loop; + int32_t nextTokensPerBatch = 0; + int64_t threshold; + uint32_t curTopKIdx = 0; + uint64_t curOffsetInSparseBlock = 0; +}; + +struct ConstInfo { + // CUBE与VEC核间同步的模式 + static constexpr uint32_t SFA_SYNC_MODE2 = 2; + // BUFFER的字节数 + static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; + static constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64; + static constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256; + static constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512; + static constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024; + static constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048; + static constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096; + static constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192; + static constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384; + static constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768; + // FP32的0值和极大值 + static constexpr float FLOAT_ZERO = 0; + static constexpr float FLOAT_MAX = 3.402823466e+38F; + + // preLoad的总次数 + uint32_t preLoadNum = 0U; + uint32_t nBufferMBaseSize = 0U; + // CUBE和VEC的核间同步EventID + uint32_t syncV1NupdateC2 = 0U; + uint32_t syncV0C1 = 0U; + uint32_t syncC1V1 = 0U; + uint32_t syncV1C2 = 0U; + uint32_t syncC2V2 = 0U; + uint32_t syncC2V1 = 0U; + + uint32_t mmResUbSize = 0U; // Matmul1输出结果GM上的大小 + uint32_t vec1ResUbSize = 0U; // Vector1输出结果GM上的大小 + uint32_t bmm2ResUbSize = 0U; // Matmul2输出结果GM上的大小 + uint64_t batchSize = 0ULL; + uint64_t gSize = 0ULL; + uint64_t qHeadNum = 0ULL; + uint64_t kvHeadNum; + uint64_t headDim; + uint64_t headDimRope; + uint64_t kvSeqSize = 0ULL; // kv最大S长度 + uint64_t qSeqSize = 1ULL; // q最大S长度 + int64_t kvCacheBlockSize = 0; // PA场景的block size + uint32_t maxBlockNumPerBatch = 0; // PA场景的最大单batch block number + uint32_t splitKVNum = 0U; // S2核间切分的切分份数 + SFA_LAYOUT outputLayout; // 输出的Transpose格式 + uint32_t sparseMode = 0; + int64_t preTokens = INT64_MAX; + int64_t nextTokens = INT64_MAX; + uint32_t attentionMode = 2; + bool returnSoftmaxLse = false; + bool needInit = false; + + // FlashDecoding + uint32_t actualCombineLoopSize = 0U; // FlashDecoding场景, S2在核间切分的最大份数 + uint64_t combineLseOffset = 0ULL; + uint64_t combineAccumOutOffset = 0ULL; + + uint32_t actualLenDimsQ = 0U; // query的actualSeqLength 的维度 + uint32_t actualLenDimsKV = 0U; // KV 的actualSeqLength 的维度 + + // TND + uint32_t s2Start = 0U; // TND场景下，S2的起始位置 + uint32_t s2End = 0U; // 单核TND场景下S2循环index上限 + + uint32_t bN2Start = 0U; + uint32_t bN2End = 0U; + uint32_t gS1Start = 0U; + uint32_t gS1End = 0U; + + uint32_t tndFDCoreArrLen = 0U; // TNDFlashDecoding相关分核信息array的长度 + uint32_t coreStartKVSplitPos = 0U; // TNDFlashDecoding kv起始位置 + + uint32_t mBaseSize = 1ULL; + uint32_t s2BaseSize = 1ULL; + + // sparse attr + int64_t sparseBlockSize = 0; + uint32_t sparseBlockCount = 0; +}; + +struct MSplitInfo { + uint32_t nBufferIdx = 0U; + uint32_t nBufferStartM = 0U; + uint32_t nBufferDealM = 0U; + uint32_t vecStartM = 0U; + uint32_t vecDealM = 0U; +}; + +#endif // SPARSE_FLASH_ATTENTION_COMMON_H diff --git a/csrc/ascend/attention/sparse_flash_attention/op_kernel/sparse_flash_attention_template_tiling_key.h b/csrc/ascend/attention/sparse_flash_attention/op_kernel/sparse_flash_attention_template_tiling_key.h new file mode 100644 index 000000000..28524b0a7 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/op_kernel/sparse_flash_attention_template_tiling_key.h @@ -0,0 +1,117 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sparse_flash_attention_template_tiling_key.h + * \brief + */ + +#ifndef SPARSE_FLASH_ATTENTION_TEMPLATE_TILING_KEY_H +#define SPARSE_FLASH_ATTENTION_TEMPLATE_TILING_KEY_H + +#include "ascendc/host_api/tiling/template_argument.h" + +#define SFA_LAYOUT_BSND 0 +#define SFA_LAYOUT_TND 1 +#define SFA_LAYOUT_PA_BSND 2 + +#define ASCENDC_TPL_4_BW 4 + +#define C_TEMPLATE 0 +#define V_TEMPLATE 1 + +// 模板参数支持的范围定义 +ASCENDC_TPL_ARGS_DECL(SparseFlashAttention, // 算子OpType +ASCENDC_TPL_BOOL_DECL(FLASH_DECODE, 0, 1), +ASCENDC_TPL_BOOL_DECL(PAGE_ATTENTION, 0, 1), +ASCENDC_TPL_UINT_DECL(LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_BSND, SFA_LAYOUT_TND), +ASCENDC_TPL_UINT_DECL(KV_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_BSND, SFA_LAYOUT_TND, + SFA_LAYOUT_PA_BSND), +ASCENDC_TPL_UINT_DECL(TEMPLATE_MODE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, C_TEMPLATE, V_TEMPLATE), +ASCENDC_TPL_BOOL_DECL(IS_SPLIT_G, 0, 1), +); + +// 支持的模板参数组合 +// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法 +ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_BSND, SFA_LAYOUT_PA_BSND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, C_TEMPLATE), + ASCENDC_TPL_BOOL_SEL(IS_SPLIT_G, 0, 1), + ), + + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_TND, SFA_LAYOUT_PA_BSND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, C_TEMPLATE), + ASCENDC_TPL_BOOL_SEL(IS_SPLIT_G, 0, 1), + ), + + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_BSND, SFA_LAYOUT_PA_BSND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, V_TEMPLATE), + ASCENDC_TPL_BOOL_SEL(IS_SPLIT_G, 0, 1), + ), + + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_TND, SFA_LAYOUT_PA_BSND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, V_TEMPLATE), + ASCENDC_TPL_BOOL_SEL(IS_SPLIT_G, 0, 1), + ), + + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_BSND, SFA_LAYOUT_PA_BSND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, C_TEMPLATE), + ASCENDC_TPL_BOOL_SEL(IS_SPLIT_G, 0, 1), + ), + + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_TND, SFA_LAYOUT_PA_BSND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, C_TEMPLATE), + ASCENDC_TPL_BOOL_SEL(IS_SPLIT_G, 0, 1), + ), + + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_BSND, SFA_LAYOUT_PA_BSND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, V_TEMPLATE), + ASCENDC_TPL_BOOL_SEL(IS_SPLIT_G, 0, 1), + ), + + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_BOOL_SEL(FLASH_DECODE, 0), + ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(KV_LAYOUT_T, ASCENDC_TPL_UI_LIST, SFA_LAYOUT_TND, SFA_LAYOUT_PA_BSND), + ASCENDC_TPL_UINT_SEL(TEMPLATE_MODE, ASCENDC_TPL_UI_LIST, V_TEMPLATE), + ASCENDC_TPL_BOOL_SEL(IS_SPLIT_G, 0, 1), + ), +); + +#endif // TEMPLATE_TILING_KEY diff --git a/csrc/ascend/attention/sparse_flash_attention/sparse_flash_attention_torch_adpt.h b/csrc/ascend/attention/sparse_flash_attention/sparse_flash_attention_torch_adpt.h new file mode 100644 index 000000000..5aec8b3c5 --- /dev/null +++ b/csrc/ascend/attention/sparse_flash_attention/sparse_flash_attention_torch_adpt.h @@ -0,0 +1,152 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef SPARSE_FLASH_ATTENTION_TORCH_ADPT_H +#define SPARSE_FLASH_ATTENTION_TORCH_ADPT_H + +namespace vllm_fl { + +namespace { + +std::tuple construct_sparse_flash_attention_output_tensor( + const at::Tensor &query, const at::Tensor &key, + const std::string &layout_query_str, const std::string &layout_kv_str, bool return_softmax_lse) +{ + constexpr int64_t SIZE = 8; + constexpr int64_t DIM_0 = 0; + constexpr int64_t DIM_1 = 1; + constexpr int64_t DIM_2 = 2; + constexpr int64_t DIM_3 = 3; + constexpr int64_t DIM_4 = 4; + + TORCH_CHECK(layout_query_str == "BSND" || layout_query_str == "TND", + "The layout of query only support BSND and TND, but got ", + layout_query_str); + for (size_t i = 0; i < query.sizes().size(); i++) { + TORCH_CHECK(query.size(i) > 0, + "All values within query's shape should be greater " + "than 0, but shape[", + i, "] is ", query.size(i)); + } + + at::SmallVector output_size; + if (layout_query_str == "TND") { + TORCH_CHECK(query.dim() == DIM_3, + "When the layout of query is TND, the query dimension must be 3, but got ", + query.dim()); + output_size = {query.size(DIM_0), query.size(DIM_1), + query.size(DIM_2)}; + } else { + TORCH_CHECK(query.dim() == DIM_4, + "When the layout of query is BSND, the query dimension must be 4, but got ", + query.dim()); + output_size = {query.size(DIM_0), query.size(DIM_1), + query.size(DIM_2), query.size(DIM_3)}; + } + + at::Tensor attention_output = + at::empty(output_size, query.options().dtype(query.dtype())); + at::SmallVector softmax_size; + if (return_softmax_lse) { + if (query.dim() == DIM_3) { + const auto kv_head_num = + layout_kv_str == "PA_BSND" ? key.size(DIM_2) : key.size(DIM_1); + softmax_size = { + kv_head_num, + query.size(DIM_0), + query.size(DIM_1) / kv_head_num, + }; + } else { + softmax_size = { + query.size(DIM_0), + key.size(DIM_2), + query.size(DIM_1), + query.size(DIM_2) / key.size(DIM_2), + }; + } + } else { + softmax_size = {0}; + } + + at::Tensor softmax_max = + at::empty(softmax_size, query.options().dtype(at::kFloat)); + at::Tensor softmax_sum = + at::empty(softmax_size, query.options().dtype(at::kFloat)); + return std::tuple( + attention_output, softmax_max, softmax_sum); +} + +} // namespace + +std::tuple npu_sparse_flash_attention( + const at::Tensor &query, const at::Tensor &key, const at::Tensor &value, + const at::Tensor &sparse_indices, double scale_value, + const c10::optional &block_table, + const c10::optional &actual_seq_lengths_query, + const c10::optional &actual_seq_lengths_kv, + const c10::optional &query_rope, + const c10::optional &key_rope, int64_t sparse_block_size, + c10::string_view layout_query, c10::string_view layout_kv, + int64_t sparse_mode, int64_t pre_tokens, int64_t next_tokens, + int64_t attention_mode, bool return_softmax_lse) +{ + TORCH_CHECK(query.numel() > 0, "Tensor query is empty."); + TORCH_CHECK(key.numel() > 0, "Tensor key is empty."); + TORCH_CHECK(value.numel() > 0, "Tensor value is empty."); + TORCH_CHECK(sparse_indices.numel() > 0, "Tensor sparse_indices is empty."); + + std::string layout_query_str = std::string(layout_query); + std::string layout_kv_str = std::string(layout_kv); + + auto sparse_flash_attention_output = + construct_sparse_flash_attention_output_tensor( + query, key, layout_query_str, layout_kv_str, return_softmax_lse); + at::Tensor attention_output = std::get<0>(sparse_flash_attention_output); + at::Tensor softmax_max = std::get<1>(sparse_flash_attention_output); + at::Tensor softmax_sum = std::get<2>(sparse_flash_attention_output); + + // convert str + char *layout_query_ptr = const_cast(layout_query_str.c_str()); + char *layout_kv_ptr = const_cast(layout_kv_str.c_str()); + + EXEC_NPU_CMD( + aclnnSparseFlashAttention, + query, + key, + value, + sparse_indices, + block_table, + actual_seq_lengths_query, + actual_seq_lengths_kv, + query_rope, + key_rope, + scale_value, + sparse_block_size, + layout_query_ptr, + layout_kv_ptr, + sparse_mode, + pre_tokens, + next_tokens, + attention_mode, + return_softmax_lse, + attention_output, + softmax_max, + softmax_sum); + return std::tuple( + attention_output, softmax_max, softmax_sum); +} +} // namespace vllm_fl + +#endif // SPARSE_FLASH_ATTENTION_TORCH_ADPT_H diff --git a/csrc/ascend/attention/store_kv_block/CMakeLists.txt b/csrc/ascend/attention/store_kv_block/CMakeLists.txt new file mode 100644 index 000000000..86b308249 --- /dev/null +++ b/csrc/ascend/attention/store_kv_block/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/attention/store_kv_block/op_host/CMakeLists.txt b/csrc/ascend/attention/store_kv_block/op_host/CMakeLists.txt new file mode 100644 index 000000000..9ccd98be3 --- /dev/null +++ b/csrc/ascend/attention/store_kv_block/op_host/CMakeLists.txt @@ -0,0 +1,30 @@ +# This program is free software, you can redistribute it and/or modify it. +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This file is a part of the CANN Open Software. +# Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ====================================================================================================================== +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + store_kv_block_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME StoreKVBlock + OPTIONS + --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE store_kv_block ACLNNTYPE aclnn) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/attention/store_kv_block/op_host/store_kv_block_def.cpp b/csrc/ascend/attention/store_kv_block/op_host/store_kv_block_def.cpp new file mode 100644 index 000000000..1b81e0971 --- /dev/null +++ b/csrc/ascend/attention/store_kv_block/op_host/store_kv_block_def.cpp @@ -0,0 +1,64 @@ +/** + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/*! + * \file store_kv_block_def.cpp + * \brief Operator definition for StoreKVBlock + */ +#include "register/op_def_registry.h" + +namespace ops { +class StoreKVBlock : public OpDef { + public: + explicit StoreKVBlock(const char* name) : OpDef(name) { + this->Input("keyIn") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("keyCacheIn") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("groupLen") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32 , ge::DT_INT32 , ge::DT_INT32 }) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("groupKeyIdx") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32 , ge::DT_INT32 , ge::DT_INT32 }) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("groupKeyCacheIdx") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32 , ge::DT_INT32 , ge::DT_INT32 }) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Attr("blockSize").Int(); + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + } +}; + +OP_ADD(StoreKVBlock); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/attention/store_kv_block/op_host/store_kv_block_infershape.cpp b/csrc/ascend/attention/store_kv_block/op_host/store_kv_block_infershape.cpp new file mode 100644 index 000000000..07dd60236 --- /dev/null +++ b/csrc/ascend/attention/store_kv_block/op_host/store_kv_block_infershape.cpp @@ -0,0 +1,33 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file store_kv_block_infershape.cpp + * \brief InferShape implementation for StoreKVBlock + */ +#include +#include +#include "error/ops_error.h" + +using namespace ge; +namespace ops { + +static ge::graphStatus InferShape4StoreKVBlock(gert::InferShapeContext* context) +{ + return GRAPH_SUCCESS; +} + +static graphStatus InferDataType4StoreKVBlock(gert::InferDataTypeContext* context) +{ + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(StoreKVBlock).InferShape(InferShape4StoreKVBlock).InferDataType(InferDataType4StoreKVBlock); +} // namespace ops diff --git a/csrc/ascend/attention/store_kv_block/op_host/store_kv_block_tiling.cpp b/csrc/ascend/attention/store_kv_block/op_host/store_kv_block_tiling.cpp new file mode 100644 index 000000000..9eb931fe5 --- /dev/null +++ b/csrc/ascend/attention/store_kv_block/op_host/store_kv_block_tiling.cpp @@ -0,0 +1,156 @@ +/** + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "store_kv_block_tiling.h" +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" +#include "tiling_base/error_log.h" + +namespace optiling { + +constexpr uint32_t DIM_0 = 0; +constexpr uint32_t DIM_1 = 1; +constexpr uint32_t DIM_2 = 2; +constexpr int32_t MAX_UB_USE_SIZE = 180 * 1024; + +struct StoreKVBlockParams { + uint32_t numTokens{0}; + uint32_t numCache{0}; + uint32_t numHeads{1}; + uint32_t headSize[5]{1, 1, 1, 1, 1}; + uint32_t blockTableSize{0}; + uint32_t typeByte{0}; + uint32_t tokenSize{1}; + uint32_t tilingKey{1}; + uint64_t workspaceSize{0}; + uint32_t groupInfoLen{0}; + uint32_t corepernum{0}; + uint32_t coretail{0}; + uint64_t sysWorkspaceSize{0}; + uint32_t coreNum{0}; +}; + +static ge::graphStatus DoCommonTiling(gert::TilingContext* context, StoreKVBlockParams& params) { + auto kShape = context->GetInputShape(DIM_0); + auto kDimNum = kShape->GetStorageShape().GetDimNum(); + if (kDimNum < 2 || kDimNum > 7) { + OP_LOGE(context->GetNodeName(), "StoreKVBlock Input kDimNum dim < 2 || kDimNum>7"); + return ge::GRAPH_FAILED; + } + + for (int i = 0; i < kDimNum; i++) { + if (i == 0) params.numTokens = static_cast(kShape->GetStorageShape().GetDim(i)); + else if (i == 1) params.numHeads = static_cast(kShape->GetStorageShape().GetDim(i)); + else if (static_cast(kShape->GetStorageShape().GetDim(i)) != 0) + params.headSize[i - 2] = static_cast(kShape->GetStorageShape().GetDim(i)); + } + + auto kCacheShape = context->GetInputShape(DIM_1); + auto kCacheDimNum = kCacheShape->GetStorageShape().GetDimNum(); + if (kCacheDimNum < 2 || kCacheDimNum > 7) { + OP_LOGE(context->GetNodeName(), "StoreKVBlock Input kCacheDimNum < 2"); + return ge::GRAPH_FAILED; + } + params.numCache = kCacheShape->GetStorageShape().GetDim(0) * kCacheShape->GetStorageShape().GetDim(1); + + const int64_t* blockSizePtr = context->GetAttrs()->GetInt(0); + uint32_t blockSize = static_cast(*blockSizePtr); + params.tokenSize = params.numHeads * params.headSize[0] * params.headSize[1] * params.headSize[2] * params.headSize[3] * params.headSize[4]; + params.blockTableSize = blockSize; + + uint32_t typeByte = 0; + auto xDataType = context->GetInputDesc(DIM_0)->GetDataType(); + if (xDataType == ge::DataType::DT_INT8) { + typeByte = sizeof(int8_t); + params.tilingKey = 1; + } else if (xDataType == ge::DataType::DT_FLOAT16 || xDataType == ge::DataType::DT_BF16) { + typeByte = sizeof(uint16_t); + params.tilingKey = 2; + } else if (xDataType == ge::DataType::DT_INT32 || xDataType == ge::DataType::DT_UINT32) { + typeByte = sizeof(uint32_t); + params.tilingKey = 4; + } else { + OP_LOGE(context->GetNodeName(), "Unsupported type."); + return ge::GRAPH_FAILED; + } + + params.typeByte = typeByte; + + auto groupInfoShape = context->GetInputShape(DIM_2); + params.groupInfoLen = static_cast(groupInfoShape->GetStorageShape().GetDim(0)); + params.corepernum = params.groupInfoLen / params.coreNum; + params.coretail = params.groupInfoLen % params.coreNum; + + uint32_t pageBlockEleSize = params.blockTableSize * params.tokenSize; + if (pageBlockEleSize > MAX_UB_USE_SIZE) { + OP_LOGE(context->GetNodeName(), "pageBlockEleSize > MaxUBSize"); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus StoreKVBlockTilingFunc(gert::TilingContext* context) { + StoreKVBlockParams params; + + auto platformInfo = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + params.coreNum = ascendcPlatform.GetCoreNum(); + if (params.coreNum == 0) { + OP_LOGE(context->GetNodeName(), "Failed to get core num."); + return ge::GRAPH_FAILED; + } + params.sysWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize(); + + auto ret = DoCommonTiling(context, params); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + StoreKVBlockTilingData tilingData; + // if (params.blockTableSize > 0) tilingData.set_blockTableSize(params.blockTableSize); + // if (params.typeByte > 0) tilingData.set_typeByte(params.typeByte); + // if (params.tokenSize > 0) tilingData.set_tokenSize(params.tokenSize); + // if (params.corepernum > 0 || params.coretail != 0) tilingData.set_corePerNum(params.corepernum); + tilingData.set_blockTableSize(params.blockTableSize); + tilingData.set_typeByte(params.typeByte); + tilingData.set_tokenSize(params.tokenSize); + tilingData.set_corePerNum(params.corepernum); + if (params.coretail < params.coreNum) tilingData.set_coreTail(params.coretail); + if (params.numTokens > 0) tilingData.set_numTokens(params.numTokens); + if (params.numCache > 0) tilingData.set_numCache(params.numCache); + if (params.groupInfoLen > 0) tilingData.set_groupInfoLen(params.groupInfoLen); + + size_t* workspaceSize = context->GetWorkspaceSizes(1); + *workspaceSize = params.workspaceSize + params.sysWorkspaceSize; + context->SetTilingKey(params.tilingKey); + if (params.coreNum > 0) context->SetBlockDim(params.coreNum); + + tilingData.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tilingData.GetDataSize()); + + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus TilingParseForStoreKVBlock(gert::TilingParseContext* context) { + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(StoreKVBlock) + .Tiling(StoreKVBlockTilingFunc) + .TilingParse(TilingParseForStoreKVBlock); + +} // namespace optiling diff --git a/csrc/ascend/attention/store_kv_block/op_host/store_kv_block_tiling.h b/csrc/ascend/attention/store_kv_block/op_host/store_kv_block_tiling.h new file mode 100644 index 000000000..ff43c6055 --- /dev/null +++ b/csrc/ascend/attention/store_kv_block/op_host/store_kv_block_tiling.h @@ -0,0 +1,23 @@ +#include "register/tilingdata_base.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(StoreKVBlockTilingData) + TILING_DATA_FIELD_DEF(uint32_t, blockTableSize); + TILING_DATA_FIELD_DEF(uint32_t, typeByte); + TILING_DATA_FIELD_DEF(uint32_t, tokenSize); + TILING_DATA_FIELD_DEF(uint32_t, corePerNum); + TILING_DATA_FIELD_DEF(uint32_t, coreTail); + TILING_DATA_FIELD_DEF(uint32_t, numTokens); + TILING_DATA_FIELD_DEF(uint32_t, numCache); + TILING_DATA_FIELD_DEF(uint32_t, groupInfoLen); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(StoreKVBlock, StoreKVBlockTilingData) + +struct StoreKVBlockCompileInfo { + uint32_t coreNum; + uint64_t ubSizePlatForm; + uint32_t sysWorkspaceSize; +}; + +} // namespace optiling diff --git a/csrc/ascend/attention/store_kv_block/op_kernel/store_kv_block.cpp b/csrc/ascend/attention/store_kv_block/op_kernel/store_kv_block.cpp new file mode 100644 index 000000000..cb6c89112 --- /dev/null +++ b/csrc/ascend/attention/store_kv_block/op_kernel/store_kv_block.cpp @@ -0,0 +1,43 @@ +/** + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/*! + * \file store_kv_block.cpp + * \brief Kernel entry for StoreKVBlock operator + */ +#include "store_kv_block.h" + +extern "C" __global__ __aicore__ void store_kv_block( + GM_ADDR keyIn, GM_ADDR keyCacheIn, GM_ADDR groupLen, GM_ADDR groupKeyIdx, GM_ADDR groupKeyCacheIdx, GM_ADDR workspace, GM_ADDR tiling) +{ + AscendC::TPipe pipe; + REGISTER_TILING_DEFAULT(StoreKVBlock::StoreKVBlockTilingData); + GET_TILING_DATA(tilingData, tiling); + + if (TILING_KEY_IS(1)) { + StoreKVBlock::StoreKVBlockBase op; + op.Init( &pipe, &tilingData); + op.Process(keyIn,keyCacheIn, groupLen, groupKeyIdx, groupKeyCacheIdx); + } else if (TILING_KEY_IS(2)) { + StoreKVBlock::StoreKVBlockBase op; + op.Init( &pipe, &tilingData); + op.Process(keyIn,keyCacheIn, groupLen, groupKeyIdx, groupKeyCacheIdx); + } else if (TILING_KEY_IS(4)) { + StoreKVBlock::StoreKVBlockBase op; + op.Init( &pipe, &tilingData); + op.Process(keyIn,keyCacheIn, groupLen, groupKeyIdx, groupKeyCacheIdx); + } +} diff --git a/csrc/ascend/attention/store_kv_block/op_kernel/store_kv_block.h b/csrc/ascend/attention/store_kv_block/op_kernel/store_kv_block.h new file mode 100644 index 000000000..22a3dc367 --- /dev/null +++ b/csrc/ascend/attention/store_kv_block/op_kernel/store_kv_block.h @@ -0,0 +1,133 @@ +/** + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/*! + * \file store_kv_block.h + * \brief StoreKVBlock kernel operator + */ + +#ifndef ASCEND_STORE_KV_BLOCK_H +#define ASCEND_STORE_KV_BLOCK_H + +#include "kernel_operator.h" + +namespace StoreKVBlock { +using namespace AscendC; + + +#ifndef STORE_KV_BLOCK_TILING_DATA_H_ +#define STORE_KV_BLOCK_TILING_DATA_H_ +struct StoreKVBlockTilingData{ + uint32_t blockTableSize; + uint32_t typeByte; + uint32_t tokenSize; + uint32_t corePerNum; + uint32_t coreTail; + uint32_t numTokens; + uint32_t numCache; + uint32_t groupInfoLen; + +}; +#endif +template +class StoreKVBlockBase { +public: + + uint32_t tokenSize = 0; + uint32_t tokenByteSize = 0; + uint32_t blockTableSize = 0; + uint32_t typeByte = 0; + uint32_t numTokens = 0; + uint32_t numCache = 0; + uint32_t groupInfoLen = 0; + + uint32_t coreId = 0; + uint32_t coreTail = 0; + uint32_t corePerNum = 0; + uint32_t blockNum = 0; + AscendC::TPipe* pipeThis; + AscendC::LocalTensor tokenLocal; + AscendC::GlobalTensor keyInputGt; + AscendC::GlobalTensor keyCacheInputGt; + AscendC::GlobalTensor groupLenGt; + AscendC::GlobalTensor groupKeyIdxGt; + AscendC::GlobalTensor groupKeyCacheIdxGt; + AscendC::TBuf tokenBuf; + __aicore__ inline StoreKVBlockBase() {} + + __aicore__ inline uint32_t RoundUp(uint32_t x, uint32_t y = 16) + { + return y == 0 ? 0 : (x + y - 1) / y * y; + } + + __aicore__ inline void Init( AscendC::TPipe *pipe, StoreKVBlockTilingData *tilingData) + { + pipeThis = pipe; + typeByte = tilingData->typeByte; + tokenSize = tilingData->tokenSize; + tokenByteSize = tokenSize*typeByte; + blockTableSize = tilingData->blockTableSize; + numTokens = tilingData->numTokens; + numCache = tilingData->numCache; + groupInfoLen = tilingData->groupInfoLen; + + coreId = AscendC::GetBlockIdx(); + coreTail = tilingData->coreTail; + blockNum = AscendC::GetBlockNum(); + if (coreId < coreTail){ + // Not all cores have corePerNum+1 items; only coreTail cores get one extra. + // If corePerNum is 0, cores beyond coreTail have no work and will not access any address. + corePerNum = tilingData->corePerNum+1; + }else { + corePerNum = tilingData->corePerNum; + } + } + __aicore__ inline void Process(GM_ADDR keyIn, GM_ADDR keyCacheIn, GM_ADDR groupLen, GM_ADDR groupKeyIdx, GM_ADDR groupKeyCacheIdx) + { + + keyInputGt.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(keyIn)); + keyCacheInputGt.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(keyCacheIn)); + groupLenGt.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(groupLen)); + groupKeyIdxGt.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(groupKeyIdx)); + groupKeyCacheIdxGt.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(groupKeyCacheIdx)); + + pipeThis->InitBuffer(tokenBuf, blockTableSize*tokenByteSize); + tokenLocal = tokenBuf.Get(); + + AscendC::DataCopyExtParams copyParams{1, 0, 0, 0, 0}; // todo: full block length + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + for (int32_t i = 0; i < corePerNum; i++) { + int32_t idx = (coreId+i*blockNum); + + if( groupLenGt.GetValue(idx)<= 0 || groupKeyIdxGt.GetValue(idx)<0 || groupKeyCacheIdxGt.GetValue(idx)<0){ + continue; + } + + copyParams.blockLen = groupLenGt.GetValue(idx)*tokenByteSize; // in bytes + DataCopyPad(tokenLocal, keyInputGt[ groupKeyIdxGt.GetValue(idx)*tokenSize], copyParams, padParams); // note: offset order + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + DataCopyPad(keyCacheInputGt[groupKeyCacheIdxGt.GetValue(idx)*tokenSize], tokenLocal, copyParams); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + } + + } + +}; +} + +#endif diff --git a/csrc/ascend/attention/store_kv_block/store_kv_block_torch_adpt.h b/csrc/ascend/attention/store_kv_block/store_kv_block_torch_adpt.h new file mode 100644 index 000000000..5a1b92ce4 --- /dev/null +++ b/csrc/ascend/attention/store_kv_block/store_kv_block_torch_adpt.h @@ -0,0 +1,37 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// #include "../aclnn_torch_adapter/op_api_common.h" + +#ifndef STORE_KV_BLOCK_TORCH_ADPT_H +#define STORE_KV_BLOCK_TORCH_ADPT_H +#include +namespace vllm_fl { + +void store_kv_block( + const at::Tensor &key_in, + const at::Tensor &key_cache_in, + const at::Tensor &group_len, + const at::Tensor &group_key_idx, + const at::Tensor &group_key_cache_idx, + int64_t block_size) +{ + + EXEC_NPU_CMD(aclnnStoreKVBlock, key_in, key_cache_in,group_len, group_key_idx, group_key_cache_idx, block_size); + +} + +} +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/CMakeLists.txt b/csrc/ascend/attention/vllm_quant_lightning_indexer/CMakeLists.txt new file mode 100644 index 000000000..a3ab34a93 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/README.md b/csrc/ascend/attention/vllm_quant_lightning_indexer/README.md new file mode 100644 index 000000000..47645c1f6 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/README.md @@ -0,0 +1,252 @@ +# VllmQuantLightningIndexer + +## 产品支持情况 + +| 产品 | 是否支持 | +| ------------------------------------------------------------ | :------: | +|Ascend 950PR/Ascend 950DT| √ | +|Atlas A3 训练系列产品/Atlas A3 推理系列产品| √ | +|Atlas A2 训练系列产品/Atlas A2 推理系列产品| √ | +|Atlas 200I/500 A2 推理产品| × | +|Atlas 推理系列加速卡产品| × | +|Atlas 训练系列产品| × | + +## 功能说明 + +- API功能：VllmQuantLightningIndexer是推理场景下，稀疏attention前处理的计算，选出关键的稀疏token，并对输入query和key进行量化实现存8算8，获取最大收益。 + +- 计算公式： + $$out = \text{Top-}k\left\{[1]_{1\times g}@\left[(W@[1]_{1\times S_{k}})\odot\text{ReLU}\left(\left(Scale_Q@Scale_K^T\right)\odot\left(Q_{index}^{Quant}@{\left(K_{index}^{Quant}\right)}^T\right)\right)\right]\right\}$$ + 主要计算过程为： + 1. 将某个token对应的输入参数`query`（$Q_{index}^{Quant}\in\R^{g\times d}$）乘以给定上下文`key`（$K_{index}^{Quant}\in\R^{S_{k}\times d}$），得到相关性。 + 2. 相关性结果与`query`和`key`对应的反量化系数`query_dequant_scale`（$Scale_Q$）和`key_dequant_scale`（$Scale_K^T$）相乘，通过激活函数$ReLU$过滤无效负相关信号后，得到当前Token与所有前序Token的相关性分数向量。 + 3. 将其与权重系数`weights`（$W$）相乘后，沿g的方向，选取前$Top-k$个索引值得到输出$out$，作为Attention的输入。 + +## 参数说明 + +| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 | +|----------------------------|-----------|----------------------------------------------------------------------|----------------|------------| +| query | 输入 | 公式中的$Q_{index}\in\R^{g\times d}，表示输入Index Query$ | INT8、FLOAT8_e4m3fn | ND | +| key | 输入 | 公式的$K_{index}\in\R^{S_{k}\times d}，表示压缩后的输入Index Key$ | INT8、FLOAT8_e4m3fn | ND | +| weights | 输入 | 公式中的$W$，表示权重系数，不支持非连续。| FLOAT16、FLOAT32 | ND | +| query_dequant_scale | 输入 | 公式中的$Scale_Q$，表示Index Query的反量化系数，不支持非连续 | FLOAT16、FLOAT32 | ND | +| key_dequant_scale | 输入 | 公式中的$Scale_Q$，表示Index Key的反量化系数，不支持非连续 | FLOAT16、FLOAT32 | ND | +| actual_seq_lengths_query | 可选输入 | 表示不同Batch中`query`的有效token数 | INT32 | ND | +| actual_seq_lengths_key | 可选输入 | 表示不同Batch中`key`的有效token数 | INT32 | ND | +| block_table | 可选输入 | 表示PageAttention中KV存储使用的block映射表 | INT32 | ND | +| metadata | 可选输入 | VllmQuantLightningIndexerMetadata算子传入的分核信息，包含使用核数、分块大小以及每个核处理数据的起始点等内容，shape大小为[1024]，当前不支持传空 | INT32 | ND | +| query_quant_mode | 可选属性 | 用于标识输入`query`的量化模式，当前支持Per-Token-Head量化模式，当前仅支持传入0 | INT32 | - | +| key_quant_mode | 可选属性| 用于标识输入`key`的量化模式，当前支持Per-Token-Head量化模式，当前仅支持传入0 | INT32 | - | +| layout_query | 可选属性| 用于标识输入`query`的数据排布格式，当前支持BSND、TND，默认值"BSND" | STRING | - | +| layout_key | 可选属性 | 用于标识输入`key`的数据排布格式，当前仅支持传入PA_BSND | STRING | - | +| sparse_count | 可选属性 | 代表topK阶段需要保留的block数量，Atlas A3 推理系列产品支持[1, 2048]，Ascend 950PR/Ascend 950DT支持512 | INT32 | - | +| sparse_mode | 可选属性 | 表示sparse的模式，支持0/3，数据类型支持`int32`。 sparse_mode为0时，代表defaultMask模式。sparse_mode为3时，代表rightDownCausal模式的mask，对应以右顶点为划分的下三角场景。 | INT32 | - | +| pre_tokens | 可选属性 | 预留参数，表示attention需要和前几个Token计算关联，仅支持默认值2^63-1 | INT64 | - | +| next_tokens | 可选属性 | 预留参数，表示attention需要和前几个Token计算关联，仅支持默认值2^63-1 | INT64 | - | +| cmp_ratio | 可选属性 | 用于稀疏计算，表示key的压缩倍数。数据类型支持`int32`。Atlas A3 推理系列产品支持1/2/4/8/16/32/64/128，Ascend 950PR/Ascend 950DT支持1/4/128。 | INT32 | - | +| return_value | 可选属性 | 表示是否输出`sparse_values`。True表示输出，False表示不输出；仅支持默认值False | BOOL | - | +| stride | 可选属性 | 表示key的首轴的stride | INT32 | - | +| sparse_indices | 输出 | 公式中的输出Out，参与稀疏attention计算的token索引值 | INT32 | ND | +| sparse_values | 输出 | 公式中的Indices输出对应的value值，**目前暂不支持返回sparse_values。** | FLOAT32 | ND | + +- Ascend 950PR/Ascend 950DT：query、key不支持INT8；weights、query_dequant_scale和key_dequant_scale不支持FLOAT16。 +- Atlas A3 训练系列产品/Atlas A3 推理系列产品、Atlas A2 训练系列产品/Atlas A2 推理系列产品：query、key不支持FLOAT8_e4m3fn；weights、query_dequant_scale和key_dequant_scale不支持FLOAT32。 + +## 约束说明 + +- 该接口支持图模式。 +- 该接口要求$W \odot Scale_Q$的结果在`float16`(Atlas A3)/`float32`(Ascend 950PR/Ascend 950DT)的表示范围内。 +- 该接口的TopK过程对NAN排序是未定义行为。 +- 参数query中的D轴和参数key中的D轴值相等为128。 +- 参数query和key中的N轴分别仅支持64和1。 +- 当`layout_query`为TND时，`actual_seq_lengths_query`必须传入，且以该入参元素的数量作为B值，该入参中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。不能出现负值。 +- 当`layout_key`为PA_BSND时，`actual_seq_lengths_key`该入参必须传入。 +- PageAttention场景下，`block_table`必须为二维，第一维长度需要等于B，第二维长度不能小于maxBlockNumPerSeq(maxBlockNumPerSeq为每个batch中最大`actual_seq_lengths_key`对应的block数量)，支持block_size取值为16的整数倍，最大支持到1024。 +- query、key、weights、query_dequant_scale、key_dequant_scale数据排布格式支持从多种维度解读，其中B（Batch Size）表示输入样本批量大小、S（Sequence Length）表示输入样本序列长度、H（Head Size）表示hidden层的大小、N（Head Num）表示多头数、D（Head Dim）表示hidden层最小的单元尺寸，且满足D=H/N、T表示所有Batch输入样本序列长度的累加和。 + +## Atlas A3 推理系列产品 调用说明 + +- 单算子模式调用 + ```python + import torch + import torch_npu + import numpy as np + import torch.nn as nn + import math + import custom_ops + + n1 = 64 + n2 = 1 + d = 128 + block_size = 128 + layout_key = "PA_BSND" + layout_query = "BSND" + query_quant_mode = 0 + key_quant_mode = 0 + np.random.seed(0) + # ------------- + b = 24 + t = None + s1 = 4 + s2 = 512 + act_seq_q = None + act_seq_k = None + sparse_mode = 0 + sparse_count = 512 + cmp_ratio = 1 + max_block_table_num = (s2 + block_size - 1) // block_size + block_table = torch.tensor([range(b * max_block_table_num)], dtype = torch.int32).reshape(b, -1) + key = torch.tensor(np.random.uniform(-128, 127, (b * max_block_table_num, block_size, n2, d))).to(torch.int8) + key_dequant_scale = torch.tensor(np.random.uniform(0, 10, (b * max_block_table_num, block_size, n2))) + key_dequant_scale = key_dequant_scale.to(torch.float16) + query = torch.tensor(np.random.uniform(-128, 127, (b, s1, n1, d))).to(torch.int8) + query_dequant_scale = torch.tensor(np.random.uniform(0, 10, (b, s1, n1))).to(torch.float16) + weights = torch.tensor(np.random.uniform(0, 0.01, (b, s1, n1))).to(torch.float16) + actual_seq_lengths_query = torch.tensor(np.random.uniform(s1, s1, (b))).to(torch.int32) \ + if act_seq_q is None else torch.tensor(act_seq_q).to(torch.int32) + actual_seq_lengths_key = torch.tensor(np.random.uniform(s2, s2, (b))).to(torch.int32) \ + if act_seq_k is None else torch.tensor(act_seq_k).to(torch.int32) + max_seqlen_q = actual_seq_lengths_query.max().item() + max_seqlen_k = actual_seq_lengths_key.max().item() + metadata = torch.ops.custom.npu_vllm_quant_lightning_indexer_metadata ( + actual_seq_lengths_query = actual_seq_lengths_query.npu(), + actual_seq_lengths_key = actual_seq_lengths_key.npu(), + num_heads_q = n1, + num_heads_k = n2, + head_dim = d, + query_quant_mode = query_quant_mode, + key_quant_mode = key_quant_mode, + batch_size = b, + max_seqlen_q = max_seqlen_q, + max_seqlen_k = max_seqlen_k, + layout_query = layout_query, + layout_key = layout_key, + sparse_count = sparse_count, + sparse_mode = sparse_mode, + pre_tokens = (1<<63)-1, + next_tokens = (1<<63)-1, + cmp_ratio = cmp_ratio, + device = 'npu:0') + + sparse_indices, sparse_values = torch.ops.custom.npu_vllm_quant_lightning_indexer(query.npu(), key.npu(), weights.npu(), query_dequant_scale.npu(), + key_dequant_scale.npu(), + actual_seq_lengths_query=actual_seq_lengths_query.npu(), + actual_seq_lengths_key=actual_seq_lengths_key.npu(), + block_table=block_table.npu(), + metadata = metadata, + query_quant_mode=query_quant_mode, + key_quant_mode=key_quant_mode, + layout_query=layout_query, + layout_key=layout_key, sparse_count=sparse_count, + sparse_mode=sparse_mode, pre_tokens=(1<<63)-1, + next_tokens=(1<<63)-1, cmp_ratio=cmp_ratio) + ``` +- aclgarph调用 + + ```python + import torch + import torch_npu + import numpy as np + import torch.nn as nn + import math + import torchair + import custom_ops + from torchair.configs.compiler_config import CompilerConfig + + n1 = 64 + n2 = 1 + d = 128 + block_size = 128 + layout_key = "PA_BSND" + layout_query = "BSND" + query_quant_mode = 0 + key_quant_mode = 0 + np.random.seed(0) + # ------------- + b = 24 + t = None + s1 = 4 + s2 = 512 + act_seq_q = None + act_seq_k = None + sparse_mode = 3 + sparse_count = 512 + pre_tokens=(1<<63)-1 + next_tokens=(1<<63)-1 + cmp_ratio = 4 + max_block_table_num = (s2 + block_size - 1) // block_size + block_table = torch.tensor([range(b * max_block_table_num)], dtype = torch.int32).reshape(b, -1).npu() + key = torch.tensor(np.random.uniform(-128, 127, (b * max_block_table_num, block_size, n2, d))).to(torch.int8).npu() + key_dequant_scale = torch.tensor(np.random.uniform(0, 10, (b * max_block_table_num, block_size, n2))).npu() + key_dequant_scale = key_dequant_scale.to(torch.float16).npu() + query = torch.tensor(np.random.uniform(-128, 127, (b, s1, n1, d))).to(torch.int8).npu() + query_dequant_scale = torch.tensor(np.random.uniform(0, 10, (b, s1, n1))).to(torch.float16).npu() + weights = torch.tensor(np.random.uniform(0, 0.01, (b, s1, n1))).to(torch.float16).npu() + actual_seq_lengths_query = torch.tensor(np.random.uniform(s1, s1, (b))).to(torch.int32).npu() \ + if act_seq_q is None else torch.tensor(act_seq_q).to(torch.int32).npu() + actual_seq_lengths_key = torch.tensor(np.random.uniform(s2, s2, (b))).to(torch.int32).npu() \ + if act_seq_k is None else torch.tensor(act_seq_k).to(torch.int32).npu() + max_seqlen_q = actual_seq_lengths_query.max().item() + max_seqlen_k = actual_seq_lengths_key.max().item() + + class QLINetwork(nn.Module): + def __init__(self): + super(QLINetwork, self).__init__() + + def forward(self, query, key, weights, q_scale, k_scale, query_quant_mode, key_quant_mode, + batch_size, num_heads_q, num_heads_k, head_dim, + actual_seq_lengths_query=None, actual_seq_lengths_key=None, + block_table=None, layout_query='BSND', layout_key='BSND', + sparse_count=512, sparse_mode=3, pre_tokens=(1<<63)-1, + next_tokens=(1<<63)-1, cmp_ratio=cmp_ratio, return_value=False): + metadata = torch.ops.custom.npu_vllm_quant_lightning_indexer_metadata( + actual_seq_lengths_query = actual_seq_lengths_query, + actual_seq_lengths_key = actual_seq_lengths_key, + num_heads_q = num_heads_q, + num_heads_k = num_heads_k, + head_dim = head_dim, + query_quant_mode = query_quant_mode, + key_quant_mode = key_quant_mode, + batch_size = batch_size, + max_seqlen_q = max_seqlen_q, + max_seqlen_k = max_seqlen_k, + layout_query = layout_query, + layout_key = layout_key, + sparse_count = sparse_count, + sparse_mode = sparse_mode, + pre_tokens = (1<<63)-1, + next_tokens = (1<<63)-1, + cmp_ratio = cmp_ratio, + device = 'npu:0') + + sparse_indices, sparse_values = torch.ops.custom.npu_vllm_quant_lightning_indexer(query, key, weights, + q_scale, k_scale, + actual_seq_lengths_query=actual_seq_lengths_query, + actual_seq_lengths_key=actual_seq_lengths_key, + block_table=block_table, metadata=metadata, + query_quant_mode=query_quant_mode, + key_quant_mode=key_quant_mode, + layout_query=layout_query, + layout_key=layout_key, sparse_count=sparse_count, + sparse_mode=sparse_mode,pre_tokens=pre_tokens, + next_tokens=next_tokens, cmp_ratio=cmp_ratio, + return_value=return_value) + return sparse_indices + + + config = CompilerConfig() + npu_backend = torchair.get_npu_backend(compiler_config=config) + torch._dynamo.reset() + npu_mode = torch.compile(QLINetwork().npu(), fullgraph=True, backend=npu_backend, dynamic=False) + sparse_indices = npu_mode( query, key, weights, query_dequant_scale, key_dequant_scale, + query_quant_mode, key_quant_mode, b, n1, n2, d, + actual_seq_lengths_query=actual_seq_lengths_query, + actual_seq_lengths_key=actual_seq_lengths_key, + block_table=block_table, + layout_query=layout_query, layout_key=layout_key, + sparse_count=sparse_count, sparse_mode=sparse_mode, + pre_tokens=pre_tokens, next_tokens=next_tokens, + cmp_ratio=cmp_ratio, return_value=False) + ``` + +更多使用示例见[pytest示例](./tests/pytest/README.md)。 diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/CMakeLists.txt b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/CMakeLists.txt new file mode 100644 index 000000000..28fee16f4 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/CMakeLists.txt @@ -0,0 +1,29 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + vllm_quant_lightning_indexer_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME VllmQuantLightningIndexer + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-vf-remove-membar=false + -mllvm -cce-aicore-hoist-movemask=false +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE vllm_quant_lightning_indexer ACLNNTYPE aclnn) +endif() diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/vllm_quant_lightning_indexer_def.cpp b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/vllm_quant_lightning_indexer_def.cpp new file mode 100644 index 000000000..5e7e24623 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/vllm_quant_lightning_indexer_def.cpp @@ -0,0 +1,152 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vllm_quant_lightning_indexer_def.cpp + * \brief + */ + +#include "register/op_def_registry.h" + +namespace ops { +class VllmQuantLightningIndexer : public OpDef { +public: + explicit VllmQuantLightningIndexer(const char *name) : OpDef(name) + { + this->Input("query") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("key") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8}) + .Format({ge::FORMAT_ND}) + .IgnoreContiguous(); + this->Input("weights") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("query_dequant_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("key_dequant_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND}) + .IgnoreContiguous(); + this->Input("actual_seq_lengths_query") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("actual_seq_lengths_key") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("block_table") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("metadata") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("sparse_indices").ParamType(REQUIRED).DataType({ge::DT_INT32}).Format({ge::FORMAT_ND}); + this->Output("sparse_values").ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}); + this->Attr("query_quant_mode").AttrType(REQUIRED).Int(0); // 0: 默认值，per-token-head + this->Attr("key_quant_mode").AttrType(REQUIRED).Int(0); // 0: 默认值，per-token-head + this->Attr("layout_query").AttrType(OPTIONAL).String("BSND"); + this->Attr("layout_key").AttrType(OPTIONAL).String("BSND"); + this->Attr("sparse_count").AttrType(OPTIONAL).Int(2048); // 2048: 默认值，筛选前2048 + this->Attr("sparse_mode").AttrType(OPTIONAL).Int(3); // 3: 默认值，只计算下三角 + this->Attr("pre_tokens").AttrType(OPTIONAL).Int(9223372036854775807); // 9223372036854775807: 默认值，int64的最大值 + this->Attr("next_tokens").AttrType(OPTIONAL).Int(9223372036854775807); // 9223372036854775807: 默认值，int64的最大值 + this->Attr("cmp_ratio").AttrType(OPTIONAL).Int(1); // 1: 压缩率 + this->Attr("return_values").AttrType(OPTIONAL).Bool(false); // 是否返回sparse_values + this->Attr("stride").AttrType(OPTIONAL).Int(1); // stride参数 + this->Attr("scale_stride").AttrType(OPTIONAL).Int(1); // scaleStride参数 + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true); + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + + OpAICoreConfig aicore_config_950; + aicore_config_950.Input("query") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT8_E4M3FN}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_950.Input("key") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT8_E4M3FN}) + .Format({ge::FORMAT_ND}) + .IgnoreContiguous(); + aicore_config_950.Input("weights") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_950.Input("query_dequant_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_950.Input("key_dequant_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .IgnoreContiguous(); + aicore_config_950.Input("actual_seq_lengths_query") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_950.Input("actual_seq_lengths_key") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_950.Input("block_table") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_950.Input("metadata") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .AutoContiguous(); + aicore_config_950.Output("sparse_indices").ParamType(REQUIRED).DataType({ge::DT_INT32}).Format({ge::FORMAT_ND}); + aicore_config_950.Output("sparse_values").ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}); + aicore_config_950.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("aclnnSupport.value", "support_aclnn") + .ExtendCfgInfo("opFile.value", "vllm_quant_lightning_indexer") + .ExtendCfgInfo("jitCompile.flag", "static_false,dynamic_false"); + } +}; +OP_ADD(VllmQuantLightningIndexer); +} // namespace ops diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/vllm_quant_lightning_indexer_infershape.cpp b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/vllm_quant_lightning_indexer_infershape.cpp new file mode 100644 index 000000000..501cddfde --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/vllm_quant_lightning_indexer_infershape.cpp @@ -0,0 +1,106 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vllm_quant_lightning_indexer_infershape.cpp + * \brief + */ +#include +#include + +#include "err/ops_err.h" +#include "log/log.h" + +using namespace ge; + +namespace ops { +constexpr uint32_t QUERY_INDEX = 0; +constexpr uint32_t KEY_INDEX = 1; +constexpr uint32_t ATTR_QUERY_LAYOUT_INDEX = 2; +constexpr uint32_t ATTR_KV_LAYOUT_INDEX = 3; +constexpr uint32_t ATTR_SPARSE_COUNT_INDEX = 4; +constexpr uint32_t ATTR_RETURN_VALUE_INDEX = 9; +constexpr uint32_t DIM_NUM_3 = 3; +constexpr uint32_t DIM_NUM_4 = 4; + +static ge::graphStatus InferShapeVllmQuantLightningIndexer(gert::InferShapeContext *context) +{ + if (context == nullptr) { + OP_LOGE("VllmQuantLightningIndexer", "context is nullptr!"); + return ge::GRAPH_FAILED; + } + const gert::Shape *queryShape = context->GetInputShape(QUERY_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, queryShape); + const gert::Shape *keyShape = context->GetInputShape(KEY_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, keyShape); + gert::Shape *sparseIndicesShape = context->GetOutputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context, sparseIndicesShape); + gert::Shape *sparseValuesShape = context->GetOutputShape(1); + + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const char *inputLayoutQueryPtr = attrs->GetAttrPointer(ATTR_QUERY_LAYOUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, inputLayoutQueryPtr); + const char *inputLayoutKeyPtr = attrs->GetAttrPointer(ATTR_KV_LAYOUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, inputLayoutKeyPtr); + const int64_t *sparse_count = attrs->GetInt(ATTR_SPARSE_COUNT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, sparse_count); + + std::string inputLayoutQueryPtrStr = std::string(inputLayoutQueryPtr); + std::string inputLayoutKeyPtrStr = std::string(inputLayoutKeyPtr); + if (inputLayoutQueryPtrStr != "TND" && inputLayoutQueryPtrStr != "BSND") { + OP_LOGE(context, "The input layout query should be TND or BSND, but got %s.", inputLayoutQueryPtrStr.c_str()); + return GRAPH_FAILED; + } + + int64_t keyHeadNum = (inputLayoutKeyPtrStr == "TND") ? keyShape->GetDim(1) : keyShape->GetDim(2); + if (inputLayoutQueryPtrStr == "BSND") { + sparseIndicesShape->SetDimNum(DIM_NUM_4); + sparseIndicesShape->SetDim(0, queryShape->GetDim(0)); // 0:Dim B + sparseIndicesShape->SetDim(1, queryShape->GetDim(1)); // 1:Dim S + sparseIndicesShape->SetDim(2, keyHeadNum); // 2:Dim N + sparseIndicesShape->SetDim(3, *sparse_count); // 3:Dim K + } else { + sparseIndicesShape->SetDimNum(DIM_NUM_3); + sparseIndicesShape->SetDim(0, queryShape->GetDim(0)); // 0:Dim T + sparseIndicesShape->SetDim(1, keyHeadNum); // 1:output shape's N Dim, 2: key shape's N Dim + sparseIndicesShape->SetDim(2, *sparse_count); // 2:Dim K + } + const bool *return_value = attrs->GetAttrPointer(ATTR_RETURN_VALUE_INDEX); + bool returnValueFlag = (return_value != nullptr) ? *return_value : false; + if (returnValueFlag) { + *sparseValuesShape = *sparseIndicesShape; + } else { + sparseValuesShape->SetDimNum(1); + sparseValuesShape->SetDim(0, 0); + } + + OP_LOGD(context->GetNodeName(), "VllmQuantLightningIndexer InferShape end."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataTypeVllmQuantLightningIndexer(gert::InferDataTypeContext *context) +{ + if (context == nullptr) { + OP_LOGE("VllmQuantLightningIndexer", "InferDataTypeContext context is nullptr!"); + return ge::GRAPH_FAILED; + } + OP_LOGD(context->GetNodeName(), "Enter VllmQuantLightningIndexer InferDataType impl."); + // default index data type is int32 + ge::DataType outputType = ge::DT_INT32; + context->SetOutputDataType(0, outputType); + OP_LOGD(context->GetNodeName(), "VllmQuantLightningIndexer InferDataType end."); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(VllmQuantLightningIndexer) + .InferShape(InferShapeVllmQuantLightningIndexer) + .InferDataType(InferDataTypeVllmQuantLightningIndexer); +} // namespace ops diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/vllm_quant_lightning_indexer_tiling.cpp b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/vllm_quant_lightning_indexer_tiling.cpp new file mode 100644 index 000000000..513df707a --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/vllm_quant_lightning_indexer_tiling.cpp @@ -0,0 +1,948 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vllm_quant_lightning_indexer_tiling.cpp + * \brief + */ + +#include "vllm_quant_lightning_indexer_tiling.h" + +#include "../op_kernel/vllm_quant_lightning_indexer_template_tiling_key.h" + +using namespace ge; +using namespace AscendC; +using std::map; +using std::string; +namespace optiling { +// --------------------------QLIInfoParser类成员函数定义------------------------------------- +ge::graphStatus QLIInfoParser::CheckRequiredInOutExistence() const +{ + OP_CHECK_IF(opParamInfo_.query.shape == nullptr, OP_LOGE(opName_, "Shape of tensor query is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.query.desc == nullptr, OP_LOGE(opName_, "Desc of tensor query is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key.shape == nullptr, OP_LOGE(opName_, "Shape of tensor key is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key.desc == nullptr, OP_LOGE(opName_, "Desc of tensor key is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.weights.shape == nullptr, OP_LOGE(opName_, "Shape of tensor weights is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.weights.desc == nullptr, OP_LOGE(opName_, "Desc of tensor weights is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.query_dequant_scale.shape == nullptr, + OP_LOGE(opName_, "Shape of tensor query_dequant_scale is nullptr"), return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.query_dequant_scale.desc == nullptr, + OP_LOGE(opName_, "Desc of tensor query_dequant_scale is nullptr"), return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key_dequant_scale.shape == nullptr, + OP_LOGE(opName_, "Shape of tensor key_dequant_scale is nullptr"), return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.key_dequant_scale.desc == nullptr, + OP_LOGE(opName_, "Desc of tensor key_dequant_scale is nullptr"), return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.attenOut.shape == nullptr, OP_LOGE(opName_, "Shape of tensor output is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.attenOut.desc == nullptr, OP_LOGE(opName_, "Desc of tensor output is nullptr"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::CheckRequiredAttrExistence() const +{ + OP_CHECK_IF(opParamInfo_.layOutQuery == nullptr, OP_LOGE(opName_, "attr layout_query is nullptr"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(opParamInfo_.layOutKey == nullptr, OP_LOGE(opName_, "attr layout_key is nullptr"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(opParamInfo_.sparseCount == nullptr, OP_LOGE(opName_, "attr sparse_count is nullptr"), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(opParamInfo_.sparseMode == nullptr, OP_LOGE(opName_, "attr sparse_mode is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.queryQuantMode == nullptr, OP_LOGE(opName_, "query_quant_mode is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.keyQuantMode == nullptr, OP_LOGE(opName_, "key_quant_mode is nullptr"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::CheckRequiredParaExistence() const +{ + if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS || CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetOpName() +{ + if (context_->GetNodeName() == nullptr) { + OP_LOGE("VllmQuantLightningIndexer", "opName got from TilingContext is nullptr"); + return ge::GRAPH_FAILED; + } + opName_ = context_->GetNodeName(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetNpuInfo() +{ + platformInfo_ = context_->GetPlatformInfo(); + OP_CHECK_IF(platformInfo_ == nullptr, OP_LOGE(opName_, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo_); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + OP_CHECK_IF(aicNum == 0 || aivNum == 0, OP_LOGE(opName_, "num of core obtained is 0."), return GRAPH_FAILED); + + socVersion_ = ascendcPlatform.GetSocVersion(); + if ((socVersion_ != platform_ascendc::SocVersion::ASCEND910B) && + (socVersion_ != platform_ascendc::SocVersion::ASCEND910_93) && + (socVersion_ != platform_ascendc::SocVersion::ASCEND950)) { + OP_LOGE(opName_, "SOC Version[%d] is not support.", static_cast(socVersion_)); + return GRAPH_FAILED; + } + OP_CHECK_IF(context_->GetWorkspaceSizes(1) == nullptr, OP_LOGE(opName_, "workSpaceSize got from ge is nullptr"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(context_->GetRawTilingData() == nullptr, + OP_LOGE(context_->GetNodeName(), "RawTilingData got from GE context is nullptr."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void QLIInfoParser::GetOptionalInputParaInfo() +{ + opParamInfo_.actualSeqLengthsQ.tensor = context_->GetOptionalInputTensor(ACTUAL_SEQ_Q_INDEX); + opParamInfo_.actualSeqLengthsQ.desc = context_->GetOptionalInputDesc(ACTUAL_SEQ_Q_INDEX); + opParamInfo_.actualSeqLengthsK.tensor = context_->GetOptionalInputTensor(ACTUAL_SEQ_K_INDEX); + opParamInfo_.actualSeqLengthsK.desc = context_->GetOptionalInputDesc(ACTUAL_SEQ_K_INDEX); + opParamInfo_.blockTable.tensor = context_->GetOptionalInputTensor(BLOCK_TABLE_INDEX); + opParamInfo_.blockTable.desc = context_->GetOptionalInputDesc(BLOCK_TABLE_INDEX); + opParamInfo_.metadata.tensor = context_->GetOptionalInputTensor(METADATA_INDEX); + opParamInfo_.metadata.desc = context_->GetOptionalInputDesc(METADATA_INDEX); +} + +void QLIInfoParser::GetInputParaInfo() +{ + opParamInfo_.query.desc = context_->GetInputDesc(QUERY_INDEX); + opParamInfo_.query.shape = context_->GetInputShape(QUERY_INDEX); + opParamInfo_.key.desc = context_->GetInputDesc(KEY_INDEX); + opParamInfo_.key.shape = context_->GetInputShape(KEY_INDEX); + opParamInfo_.weights.desc = context_->GetInputDesc(WEIGTHS_INDEX); + opParamInfo_.weights.shape = context_->GetInputShape(WEIGTHS_INDEX); + opParamInfo_.query_dequant_scale.desc = context_->GetInputDesc(QUERY_DEQUANT_SCALE_INDEX); + opParamInfo_.query_dequant_scale.shape = context_->GetInputShape(QUERY_DEQUANT_SCALE_INDEX); + opParamInfo_.key_dequant_scale.desc = context_->GetInputDesc(KEY_DEQUANT_SCALE_INDEX); + opParamInfo_.key_dequant_scale.shape = context_->GetInputShape(KEY_DEQUANT_SCALE_INDEX); + GetOptionalInputParaInfo(); +} + +void QLIInfoParser::GetOutputParaInfo() +{ + opParamInfo_.attenOut.desc = context_->GetOutputDesc(vllm_quant_lightning_indexer); + opParamInfo_.attenOut.shape = context_->GetOutputShape(vllm_quant_lightning_indexer); +} + +ge::graphStatus QLIInfoParser::GetAttrParaInfo() +{ + auto attrs = context_->GetAttrs(); + OP_CHECK_IF(attrs == nullptr, OP_LOGE(context_->GetNodeName(), "attrs got from ge is nullptr"), + return ge::GRAPH_FAILED); + + OP_LOGI(context_->GetNodeName(), "GetAttrParaInfo start"); + opParamInfo_.layOutQuery = attrs->GetStr(ATTR_QUERY_LAYOUT_INDEX); + opParamInfo_.layOutKey = attrs->GetStr(ATTR_KEY_LAYOUT_INDEX); + + opParamInfo_.queryQuantMode = attrs->GetAttrPointer(ATTR_QUERY_QUANT_MODE_INDEX); + opParamInfo_.keyQuantMode = attrs->GetAttrPointer(ATTR_KEY_QUANT_MODE_INDEX); + opParamInfo_.layOutQuery = attrs->GetStr(ATTR_QUERY_LAYOUT_INDEX); + opParamInfo_.layOutKey = attrs->GetStr(ATTR_KEY_LAYOUT_INDEX); + opParamInfo_.sparseCount = attrs->GetAttrPointer(ATTR_SPARSE_COUNT_INDEX); + opParamInfo_.sparseMode = attrs->GetAttrPointer(ATTR_SPARSE_MODE_INDEX); + opParamInfo_.preTokens = attrs->GetAttrPointer(ATTR_PRE_TOKENS_INDEX); + opParamInfo_.nextTokens = attrs->GetAttrPointer(ATTR_NEXT_TOKENS_INDEX); + opParamInfo_.cmpRatio = attrs->GetAttrPointer(ATTR_CMP_RATIO_INDEX); + opParamInfo_.returnValues = attrs->GetAttrPointer(ATTR_RETURN_VALUES_INDEX); + opParamInfo_.stride = attrs->GetAttrPointer(ATTR_STRIDE_INDEX); + opParamInfo_.scaleStride = attrs->GetAttrPointer(ATTR_SCALE_STRIDE_INDEX); + + if (opParamInfo_.layOutQuery != nullptr) { + OP_LOGI(context_->GetNodeName(), "layout_query is:%s", opParamInfo_.layOutQuery); + } + if (opParamInfo_.layOutKey != nullptr) { + OP_LOGI(context_->GetNodeName(), "layout_key is:%s", opParamInfo_.layOutKey); + } + if (opParamInfo_.sparseCount != nullptr) { + OP_LOGI(context_->GetNodeName(), "selscted count is:%d", *opParamInfo_.sparseCount); + } + if (opParamInfo_.sparseMode != nullptr) { + OP_LOGI(context_->GetNodeName(), "sparse mode is:%d", *opParamInfo_.sparseMode); + } + if (opParamInfo_.preTokens != nullptr) { + OP_LOGI(context_->GetNodeName(), "preTokens is:%d", *opParamInfo_.preTokens); + } + if (opParamInfo_.nextTokens != nullptr) { + OP_LOGI(context_->GetNodeName(), "nextTokens is:%d", *opParamInfo_.nextTokens); + } + if (opParamInfo_.cmpRatio != nullptr) { + OP_LOGI(context_->GetNodeName(), "cmpRatio is:%d", *opParamInfo_.cmpRatio); + } + if (opParamInfo_.returnValues != nullptr) { + OP_LOGI(context_->GetNodeName(), "returnValues is:%s", *opParamInfo_.returnValues ? "true" : "false"); + } + if (opParamInfo_.queryQuantMode != nullptr) { + OP_LOGI(context_->GetNodeName(), "query_quant_mode mode is:%d", *opParamInfo_.queryQuantMode); + } + if (opParamInfo_.keyQuantMode != nullptr) { + OP_LOGI(context_->GetNodeName(), "key_quant_mode mode is:%d", *opParamInfo_.keyQuantMode); + } + OP_LOGI(context_->GetNodeName(), "GetAttrParaInfo end"); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::CheckAttrParaInfo() +{ + std::string layout_key(opParamInfo_.layOutKey); + std::string layout_query(opParamInfo_.layOutQuery); + + OP_CHECK_IF( + ((std::string(opParamInfo_.layOutKey) != "PA_BSND")), + OP_LOGE(opName_, "input attr layout_key only supported PA_BSND," + "but now layout_key is %s.", layout_key.c_str()), + return ge::GRAPH_FAILED); + + if ((socVersion_ == platform_ascendc::SocVersion::ASCEND910B) || + (socVersion_ == platform_ascendc::SocVersion::ASCEND910_93)) { + OP_CHECK_IF(!((*opParamInfo_.sparseCount > 0) && (*opParamInfo_.sparseCount <= SPARSE_LIMIT)), + OP_LOGE(opName_, "input attr sparse_count must > 0 and <= %d, but now sparse_count is %d", + SPARSE_LIMIT, *opParamInfo_.sparseCount),return ge::GRAPH_FAILED); + OP_CHECK_IF((*opParamInfo_.cmpRatio <= 0) || (*opParamInfo_.cmpRatio > 128) || + ((*opParamInfo_.cmpRatio & (*opParamInfo_.cmpRatio - 1)) != 0), + OP_LOGE(opName_, "input attr cmpRatio must > 0 and <= 128 and should be powers of 2, but now cmpRatio is %ld.", + *opParamInfo_.cmpRatio), return ge::GRAPH_FAILED); + } else if (socVersion_ == platform_ascendc::SocVersion::ASCEND950) { + OP_CHECK_IF(!((*opParamInfo_.sparseCount > 0) && (*opParamInfo_.sparseCount <= SPARSE_LIMIT)), + OP_LOGE(opName_, "input attr sparse_count must > 0 and <= %d, but now sparse_count is %d", + SPARSE_LIMIT, *opParamInfo_.sparseCount),return ge::GRAPH_FAILED); + OP_CHECK_IF((*opParamInfo_.cmpRatio != 1) && (*opParamInfo_.cmpRatio != 4) && (*opParamInfo_.cmpRatio != 128), + OP_LOGE(opName_, "input attr cmpRatio must be 1、4 or 128, but now cmpRatio is %ld.", + *opParamInfo_.cmpRatio), return ge::GRAPH_FAILED); + } + + OP_CHECK_IF(((std::string(opParamInfo_.layOutQuery) != "BSND") && (std::string(opParamInfo_.layOutQuery) != "TND")), + OP_LOGE(opName_, "input attr layout_query only supported BSND or TND."), return ge::GRAPH_FAILED); + OP_CHECK_IF( + ((std::string(opParamInfo_.layOutKey) != "PA_BSND") && + (std::string(opParamInfo_.layOutQuery)) != (std::string(opParamInfo_.layOutKey))), + OP_LOGE(opName_, "outside of PA, input attr layout_query and input attr layout_key must be the same," + "but now layout_key is %s, layout_query is %s.", + layout_key.c_str(), layout_query.c_str()), return ge::GRAPH_FAILED); + OP_CHECK_IF(!((*opParamInfo_.sparseMode == 0) || (*opParamInfo_.sparseMode == SPARSE_MODE_LOWER)), + OP_LOGE(opName_, "input attr sparse_mode only supported 0 or 3, but now sparseMode is %d.", + *opParamInfo_.sparseMode), return ge::GRAPH_FAILED); + OP_CHECK_IF(*opParamInfo_.preTokens != 9223372036854775807, + OP_LOGE(opName_, "input attr preTokens only supported 9223372036854775807, but now preTokens is %ld.", + *opParamInfo_.preTokens), return ge::GRAPH_FAILED); + OP_CHECK_IF(*opParamInfo_.nextTokens != 9223372036854775807, + OP_LOGE(opName_, "input attr nextTokens only supported 9223372036854775807, but now nextTokens is %ld.", + *opParamInfo_.nextTokens), return ge::GRAPH_FAILED); + + OP_CHECK_IF(*opParamInfo_.queryQuantMode != 0, OP_LOGE(opName_, "input attr query_quant_mode only supported 0."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(*opParamInfo_.keyQuantMode != 0, OP_LOGE(opName_, "input attr key_quant_mode only supported 0."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(*opParamInfo_.returnValues, OP_LOGE(opName_, "input attr returnValues only supported False."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetOpParaInfo() +{ + GetInputParaInfo(); + GetOutputParaInfo(); + if (ge::GRAPH_SUCCESS != GetAttrParaInfo()) { + return ge::GRAPH_FAILED; + } + if (ge::GRAPH_SUCCESS != CheckAttrParaInfo()) { + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetAndCheckInOutDataType() +{ + inputQType_ = opParamInfo_.query.desc->GetDataType(); + inputKType_ = opParamInfo_.key.desc->GetDataType(); + weightsType_ = opParamInfo_.weights.desc->GetDataType(); + inputQueryScaleType_ = opParamInfo_.query_dequant_scale.desc->GetDataType(); + inputKeyScaleType_ = opParamInfo_.key_dequant_scale.desc->GetDataType(); + outputType_ = opParamInfo_.attenOut.desc->GetDataType(); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo_); + socVersion_ = ascendcPlatform.GetSocVersion(); + + OP_CHECK_IF(!(inputQType_ == inputKType_), + OP_LOGE(opName_, "The data types of the input query and key must be the same."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF( + !(inputQueryScaleType_ == inputKeyScaleType_), + OP_LOGE(opName_, "The data types of the input query_dequant_scale and key_dequant_scale must be the same."), + return ge::GRAPH_FAILED); + if ((socVersion_ == platform_ascendc::SocVersion::ASCEND910B) || + (socVersion_ == platform_ascendc::SocVersion::ASCEND910_93)) { + OP_CHECK_IF(inputQType_ != ge::DT_INT8, + OP_LOGE(opName_, "The data types of the input query and key must be int8."), return ge::GRAPH_FAILED); + OP_CHECK_IF( + inputQueryScaleType_ != ge::DT_FLOAT16, + OP_LOGE(opName_, "The data types of the input query_dequant_scale and key_dequant_scale must be float16."), + return ge::GRAPH_FAILED); + } else if (socVersion_ == platform_ascendc::SocVersion::ASCEND950) { + OP_CHECK_IF(inputQType_ != ge::DT_FLOAT8_E4M3FN, + OP_LOGE(opName_, "The data types of the input query and key must be float8_e4m3."), return ge::GRAPH_FAILED); + OP_CHECK_IF( + inputQueryScaleType_ != ge::DT_FLOAT, + OP_LOGE(opName_, "The data types of the input query_dequant_scale and key_dequant_scale must be float."), + return ge::GRAPH_FAILED); + } + + if ((socVersion_ == platform_ascendc::SocVersion::ASCEND910B) || + (socVersion_ == platform_ascendc::SocVersion::ASCEND910_93)) { + OP_CHECK_IF(weightsType_ != ge::DT_FLOAT16, + OP_LOGE(opName_, "The data types of the input weights must be float16."), return ge::GRAPH_FAILED); + } else if (socVersion_ == platform_ascendc::SocVersion::ASCEND950) { + OP_CHECK_IF(weightsType_ != ge::DT_FLOAT, + OP_LOGE(opName_, "The data types of the input weights must be float."), return ge::GRAPH_FAILED); + } + + OP_CHECK_IF(outputType_ != ge::DT_INT32, + OP_LOGE(opName_, "The data types of the output sparse_indices must be int32."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetQueryKeyAndOutLayout() +{ + // 获取query,key的Layout基准值 + const map layoutQueryMap = {{"BSND", DataLayout::BSND}, {"TND", DataLayout::TND}}; + + std::string layout_query(opParamInfo_.layOutQuery); + auto QLayout_ = layoutQueryMap.find(layout_query); + if (QLayout_ != layoutQueryMap.end()) { + qLayout_ = QLayout_->second; + } + + const map layoutKeyMap = { + {"BSND", DataLayout::BSND}, {"TND", DataLayout::TND}, + {"PA_BSND", DataLayout::PA_BSND}, {"PA_BBND", DataLayout::PA_BSND}}; + std::string layout_key(opParamInfo_.layOutKey); + auto KLayout = layoutKeyMap.find(layout_key); + if (KLayout != layoutKeyMap.end()) { + kLayout_ = KLayout->second; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetAndCheckOptionalInput() +{ + if (kLayout_ == DataLayout::PA_BSND) { + OP_CHECK_IF(opParamInfo_.blockTable.tensor == nullptr, + OP_LOGE(opName_, "key layout only supported PA_BSND, input block_table must not be null"), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + opParamInfo_.actualSeqLengthsK.tensor == nullptr, + OP_LOGE(opName_, "key layout only supported PA_BSND, input actual_seq_lengths_key must not be null"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.blockTable.desc->GetDataType() != ge::DT_INT32, + OP_LOGE(opName_, "input block_table data type only support int32"), return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF(opParamInfo_.blockTable.tensor != nullptr, + OP_LOGE(opName_, "key layout is not PA_BSND, input block_table must be null"), + return ge::GRAPH_FAILED); + } + + if (kLayout_ == DataLayout::TND) { + OP_CHECK_IF(opParamInfo_.actualSeqLengthsK.tensor == nullptr, + OP_LOGE(opName_, "when layout_key is TND, input actual_seq_lengths_key must not be null"), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(opParamInfo_.actualSeqLengthsK.tensor != nullptr && + opParamInfo_.actualSeqLengthsK.desc->GetDataType() != ge::DT_INT32, + OP_LOGE(opName_, "input actual_seq_lengths_key data type only support int32"), + return ge::GRAPH_FAILED); + if (qLayout_ == DataLayout::TND) { + OP_CHECK_IF(opParamInfo_.actualSeqLengthsQ.tensor == nullptr, + OP_LOGE(opName_, "when layout_query is TND, input actual_seq_lengths_query must not be null"), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(opParamInfo_.actualSeqLengthsQ.tensor != nullptr && + opParamInfo_.actualSeqLengthsQ.desc->GetDataType() != ge::DT_INT32, + OP_LOGE(opName_, "input actual_seq_lengths_query data type only support int32"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(opParamInfo_.metadata.tensor == nullptr, + OP_LOGE(opName_, "input metadata must not be null"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::CheckShapeDim() +{ + OP_CHECK_IF((opParamInfo_.blockTable.tensor != nullptr) && + (opParamInfo_.blockTable.tensor->GetStorageShape().GetDimNum() != DIM_NUM_TWO), + OP_LOGE(opName_, "the dim num of block_table's shape should be 2, but now is %u", + opParamInfo_.blockTable.tensor->GetStorageShape().GetDimNum()), return ge::GRAPH_FAILED); + OP_CHECK_IF( + ((kLayout_ == DataLayout::PA_BSND)||(kLayout_ == DataLayout::BSND)) && + (opParamInfo_.key.shape->GetStorageShape().GetDimNum() != DIM_NUM_FOUR), + OP_LOGE(opName_, "the dim num of key's shape should be 4, but now is %u", + opParamInfo_.key.shape->GetStorageShape().GetDimNum()), return ge::GRAPH_FAILED); + OP_CHECK_IF( + (kLayout_ == DataLayout::TND) && (opParamInfo_.key.shape->GetStorageShape().GetDimNum() != DIM_NUM_THREE), + OP_LOGE(opName_, "the dim num of key's shape should be 3, but now is %u", + opParamInfo_.key.shape->GetStorageShape().GetDimNum()), return ge::GRAPH_FAILED); + + uint32_t qShapeDim = opParamInfo_.query.shape->GetStorageShape().GetDimNum(); + uint32_t weightsShapeDim = opParamInfo_.weights.shape->GetStorageShape().GetDimNum(); + uint32_t outShapeDim = opParamInfo_.attenOut.shape->GetStorageShape().GetDimNum(); + uint32_t expectShapeDim = DIM_NUM_FOUR; + if (qLayout_ == DataLayout::TND) { + expectShapeDim = DIM_NUM_THREE; + } + OP_CHECK_IF( + qShapeDim != expectShapeDim, + OP_LOGE(opName_, "the dim num of query's shape should be %u, but now is %u", expectShapeDim, qShapeDim), + return ge::GRAPH_FAILED); + OP_CHECK_IF(outShapeDim != expectShapeDim, + OP_LOGE(opName_, "the dim num of sparse_indices's shape should be %u, but now is %u", expectShapeDim, + outShapeDim), + return ge::GRAPH_FAILED); + OP_CHECK_IF(!(weightsShapeDim == expectShapeDim - 1), + OP_LOGE(opName_, "the dim num of weights's shape should be %u, but now is %u", expectShapeDim - 1, + weightsShapeDim), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetN1Size() +{ + if (qLayout_ == DataLayout::BSND) { + n1Size_ = static_cast(opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_TWO)); + } else { + // TND + n1Size_ = static_cast(opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_ONE)); + } + OP_LOGI(context_->GetNodeName(), "n1Size is %d", n1Size_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + const std::string &actualSeqLenName) const +{ + size = static_cast(tensor->GetShapeSize()); + if (size <= 0) { + OP_LOGE(opName_, "%s's shape size is %u, it should be greater than 0.", actualSeqLenName.c_str(), size); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetAndCheckN2Size() +{ + // PA_BSND + if (kLayout_ == DataLayout::TND) { + n2Size_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_ONE)); + } else { + n2Size_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_TWO)); + } + OP_LOGI(context_->GetNodeName(), "N2 is %d", n2Size_); + OP_CHECK_IF(n2Size_ != 1, OP_LOGE(opName_, "key shape[2] is numhead, only support 1."), return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetGSize() +{ + if (n1Size_ % n2Size_ != 0) { + OP_LOGE(opName_, "input query's head_num %u can not be a multiple of key's head_num %u.", n1Size_, n2Size_); + return ge::GRAPH_FAILED; + } + gSize_ = n1Size_ / n2Size_; + OP_CHECK_IF(gSize_ != G_SIZE_LIMIT, + OP_LOGE(opName_, "N1 is %u, N2 is %u, N1 divided by N2 must equal 64.", n1Size_, n2Size_), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetBatchSize() +{ + // 获取B基准值 + // 1、非TND时, 以query的batch_size维度为基准; + // 2、Q和K都为TND时, actual_seq_lens_q必须传入, 以actual_seq_lens_q数组的长度为B轴大小 + // 3、Q为TND，K为PA_BSND时，以actual_seq_lens_k数组的长度为B轴大小 + if (qLayout_ == DataLayout::BSND) { + bSize_ = opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_ZERO); + OP_LOGI(context_->GetNodeName(), "b: %d, s: %d, n: %d,d :%d", + opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_ZERO), + opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_ONE), + opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_TWO), + opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_THREE)); + return ge::GRAPH_SUCCESS; + } else { // TND + uint32_t bSizeQuery; + uint32_t bSizeKey; + GetActualSeqLenSize(bSizeQuery, opParamInfo_.actualSeqLengthsQ.tensor, "input actual_seq_lengths_query"); + GetActualSeqLenSize(bSizeKey, opParamInfo_.actualSeqLengthsK.tensor, "input actual_seq_lengths_key"); + if (kLayout_ == DataLayout::TND) { + OP_CHECK_IF(bSizeQuery != bSizeKey, + OP_LOGE(opName_, "the lengths of actual_seq_lengths_query and actual_seq_lengths_key is %u, %u respectively, they must be same.", + bSizeQuery, bSizeKey), + return ge::GRAPH_FAILED); + bSize_ = bSizeQuery; + } else { + if (bSizeQuery == bSizeKey + 1) { + batchSupperFlag_ = true; + } + OP_CHECK_IF((bSizeQuery != bSizeKey) && !batchSupperFlag_, + OP_LOGE(opName_, "the lengths of actual_seq_lengths_query and actual_seq_lengths_key is %u, %u respectively, they must be same.", + bSizeQuery, bSizeKey), + return ge::GRAPH_FAILED); + bSize_ = bSizeKey; // Q为TND，batch从Key中获取 + } + return ge::GRAPH_SUCCESS; + } +} + +ge::graphStatus QLIInfoParser::GetHeadDim() +{ + // 以query的D维度为基准 + uint32_t dIndex = DIM_IDX_TWO; + // 根据layout确定D维度在shape中的位置 + switch (qLayout_) { + case DataLayout::TND: + // TND格式: [Total, N, D] -> D是第2维(索引2) + dIndex = DIM_IDX_TWO; + break; + case DataLayout::BSND: + // BSND格式: [Batch, SeqLen, N, D] -> D是第3维(索引3) + dIndex = DIM_IDX_THREE; + break; + default: + OP_LOGE(opName_, "unsupported layout for getting head dim."); + return ge::GRAPH_FAILED; + } + headDim_ = opParamInfo_.query.shape->GetStorageShape().GetDim(dIndex); + OP_CHECK_IF(headDim_ != HEAD_DIM_LIMIT, OP_LOGE(opName_, "input query's last dim head_dim only support 128, but now is %u.", headDim_), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetS1Size() +{ + if (qLayout_ == DataLayout::BSND) { + s1Size_ = opParamInfo_.query.shape->GetStorageShape().GetDim(1); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetAndCheckBlockSize() +{ + blockSize_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(1)); + OP_LOGI(context_->GetNodeName(), "blockSize_ is %d", blockSize_); + + OP_CHECK_IF( + ((blockSize_ % BLOCK_SIZE_FACTOR != 0) || (blockSize_ == 0) || (blockSize_ > BLOCK_SIZE_LIMIT)), + OP_LOGE(opName_, "input key's block_size must be a multiple of 16 and belong to (0, 1024], but now is %d.", blockSize_), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetS2SizeForPageAttention() +{ + if (GetAndCheckBlockSize() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + int32_t blockCount_ = static_cast(opParamInfo_.key.shape->GetStorageShape().GetDim(0)); + OP_CHECK_IF((blockCount_ == 0), OP_LOGE(opName_, "input key's block_count cannot be 0."), return ge::GRAPH_FAILED); + + maxBlockNumPerBatch_ = opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1); + s2Size_ = maxBlockNumPerBatch_ * blockSize_; + OP_LOGI(context_->GetNodeName(), "maxBlockNumPerBatch_ is %d, blockSize_ is %d, s2Size_ is %d", + maxBlockNumPerBatch_, blockSize_, s2Size_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetS2SizeForBatchContinuous() +{ + std::string layout_key(opParamInfo_.layOutKey); + if (kLayout_ == DataLayout::BSND) { + s2Size_ = opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_ONE); + } else if (kLayout_ == DataLayout::TND) { + s2Size_ = opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_ZERO); + } + OP_CHECK_IF((kLayout_ != DataLayout::BSND) && (kLayout_ != DataLayout::TND), + OP_LOGE(opName_, "the layout of key is %s, it is unsupported.", layout_key.c_str()), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::GetS2Size() +{ + // 获取S2基准值 + // 1、BATCH_CONTINUOUS时, 从key的S轴获取 + // 3、PAGE_ATTENTION时, S2 = block_table.dim1 * block_size + if (kLayout_ == DataLayout::PA_BSND) { + return GetS2SizeForPageAttention(); + } + return GetS2SizeForBatchContinuous(); +} + +ge::graphStatus QLIInfoParser::ValidateInputShapesMatch() +{ + /* + TND: + query [T,N1,D], + key [BlockNum,BlockSize,N2,D], + weight [T,N1], + block_table [BatchSize, BatchMaxBlockNum], + act_seq_k [BatchSize] + act_seq_q [BatchSize], + out [T,N2,topk] + ---------------------- + BSND: + query [BatchSize,S1,N1,D], + key [BlockNum,BlockSize,N2,D], + weight [BatchSize,S1,N1], + block_table [BatchSize, BatchMaxBlockNum], + act_seq_k [BatchSize] + act_seq_q [BatchSize] 可选 + out [BatchSize,S1,N2,topk] + */ + uint32_t queryWeightsN1Dim = 1; + uint32_t outN2Dim = 1; + + if (qLayout_ == DataLayout::TND) { + // -----------------------check BatchSize------------------- + // bSize_ 来源于act_seq_q + OP_CHECK_IF((kLayout_ == DataLayout::PA_BSND) && + ((opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize() != bSize_) || + (opParamInfo_.blockTable.tensor != nullptr && + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) != bSize_)), + OP_LOGE( + opName_, + "TND case input actual_seq_lengths_query, actual_seq_lengths_key, block_table dim 0 are %u, %u, %u " + "respectively, they must be same.", + bSize_, opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize(), + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + OP_CHECK_IF((kLayout_ != DataLayout::PA_BSND) && + (opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize() != bSize_), + OP_LOGE( + opName_, + "TND case input actual_seq_lengths_query, actual_seq_lengths_key, are %u, %u " + "respectively, they must be same.", + bSize_, opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize()), + return ge::GRAPH_FAILED); + // -----------------------check T------------------- + uint32_t qTsize = opParamInfo_.query.shape->GetStorageShape().GetDim(0); + OP_CHECK_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != qTsize) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != qTsize), + OP_LOGE(opName_, + "TND case input query, weights, sparse_indices dim 0 are %u, %u, %u " + "respectively, they must be same.", + qTsize, opParamInfo_.weights.shape->GetStorageShape().GetDim(0), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + } else { + // -----------------------check BatchSize------------------- + // bSize_ 来源于query + OP_CHECK_IF((kLayout_ == DataLayout::PA_BSND) && + ((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != bSize_) || + (opParamInfo_.blockTable.tensor != nullptr && + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) != bSize_) || + (opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize() != bSize_) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != bSize_)), + OP_LOGE(opName_, + "BSND case input query, weight, actual_seq_lengths_key, block_table, sparse_indices dim 0 are %u, %u, %u, %u, %u respectively, they must be same.", + bSize_, opParamInfo_.weights.shape->GetStorageShape().GetDim(0), + opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize(), + opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + OP_CHECK_IF((kLayout_ != DataLayout::PA_BSND) && + ((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != bSize_) || + (opParamInfo_.actualSeqLengthsK.tensor != nullptr && + opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize() != bSize_) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != bSize_)), + OP_LOGE(opName_, + "BSND case input query, weight, actual_seq_lengths_key, sparse_indices dim 0 are %u, %u, %u, %u respectively, they must be same.", + bSize_, opParamInfo_.weights.shape->GetStorageShape().GetDim(0), + opParamInfo_.actualSeqLengthsK.tensor->GetShapeSize(), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + (opParamInfo_.actualSeqLengthsQ.tensor != nullptr) && + (opParamInfo_.actualSeqLengthsQ.tensor->GetShapeSize() != bSize_), + OP_LOGE( + opName_, + "BSND case input query, actual_seq_lengths_query dim 0 are %u, %ld respectively, they must be same", + bSize_, opParamInfo_.actualSeqLengthsQ.tensor->GetShapeSize()), + return ge::GRAPH_FAILED); + // -----------------------check S1------------------- + OP_CHECK_IF( + (opParamInfo_.weights.shape->GetStorageShape().GetDim(1) != s1Size_) || + (opParamInfo_.attenOut.shape->GetStorageShape().GetDim(1) != s1Size_), + OP_LOGE(opName_, "BSND case input query, weight, sparse_indices dim 1 are %u, %u, %u, they must be same.", + s1Size_, opParamInfo_.weights.shape->GetStorageShape().GetDim(1), + opParamInfo_.attenOut.shape->GetStorageShape().GetDim(1)), + return ge::GRAPH_FAILED); + queryWeightsN1Dim = DIM_IDX_TWO; + outN2Dim = DIM_IDX_TWO; + } + // -----------------------check N1------------------- + OP_CHECK_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(queryWeightsN1Dim) != n1Size_), + OP_LOGE(opName_, "input query, weight shape dim N1 must be same, but now are %u, %u respectively, they must be same.", + opParamInfo_.weights.shape->GetStorageShape().GetDim(queryWeightsN1Dim), n1Size_), + return ge::GRAPH_FAILED); + // -----------------------check D------------------- + OP_CHECK_IF( + ((kLayout_ != DataLayout::TND && opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_THREE) != headDim_) + || (kLayout_ == DataLayout::TND && opParamInfo_.key.shape->GetStorageShape().GetDim(DIM_IDX_TWO) != headDim_)), + OP_LOGE(opName_, "input query, key shape last dim must be same."), return ge::GRAPH_FAILED); + // -----------------------check N2------------------- + OP_CHECK_IF((opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim) != n2Size_), + OP_LOGE(opName_, "input query and output sparse_indices shape n2 dim must be same."), + return ge::GRAPH_FAILED); + // -----------------------check sparse_count------------------- + OP_CHECK_IF((opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim + 1) != *opParamInfo_.sparseCount), + OP_LOGE(opName_, "output sparse_indices shape last dim must be same as attr sparse_count."), + return ge::GRAPH_FAILED); + // -----------------------check metadata------------------- + OP_CHECK_IF((opParamInfo_.metadata.tensor->GetShapeSize() != METADATA_LIMIT), + OP_LOGE(opName_, "input metadata dim 0 must be %u.", METADATA_LIMIT), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus QLIInfoParser::CheckScaleShape() +{ + uint32_t qShapeDim = opParamInfo_.query.shape->GetStorageShape().GetDimNum(); + uint32_t kShapeDim = opParamInfo_.key.shape->GetStorageShape().GetDimNum(); + uint32_t qDequantScaleShapeDim = opParamInfo_.query_dequant_scale.shape->GetStorageShape().GetDimNum(); + uint32_t kDequantScaleShapeDim = opParamInfo_.key_dequant_scale.shape->GetStorageShape().GetDimNum(); + OP_CHECK_IF(qDequantScaleShapeDim != (qShapeDim - 1), + OP_LOGE(opName_, "the dim num of query_dequant_scale's shape should be %u, but now is %u", + qShapeDim - 1, qDequantScaleShapeDim), + return ge::GRAPH_FAILED); + OP_CHECK_IF(kDequantScaleShapeDim != (kShapeDim - 1), + OP_LOGE(opName_, "the dim num of key_dequant_scale's shape should be %u, but now is %u", kShapeDim - 1, + kDequantScaleShapeDim), + return ge::GRAPH_FAILED); + // check q scale + for (uint32_t i = 0; i < (qShapeDim - 1); i++) { + uint32_t dimValueQueryScale = opParamInfo_.query_dequant_scale.shape->GetStorageShape().GetDim(i); + uint32_t dimValueQuery = opParamInfo_.query.shape->GetStorageShape().GetDim(i); + OP_CHECK_IF(dimValueQueryScale != dimValueQuery, + OP_LOGE(opName_, "query_dequant_scale's shape[%u] %u and query's shape[%u] %u is not same", i, + dimValueQueryScale, i, dimValueQuery), + return ge::GRAPH_FAILED); + } + // check k scale + for (uint32_t i = 0; i < (kShapeDim - 1); i++) { + uint32_t dimValueKeyScale = opParamInfo_.key_dequant_scale.shape->GetStorageShape().GetDim(i); + uint32_t dimValueKey = opParamInfo_.key.shape->GetStorageShape().GetDim(i); + OP_CHECK_IF(dimValueKeyScale != dimValueKey, + OP_LOGE(opName_, "key_dequant_scale's shape[%u] %u and key's shape[%u] %u is not same", i, + dimValueKeyScale, i, dimValueKey), + return ge::GRAPH_FAILED); + } + + return ge::GRAPH_SUCCESS; +} + +void QLIInfoParser::GenerateInfo(QLITilingInfo &QLIInfo) +{ + QLIInfo.opName = opName_; + QLIInfo.platformInfo = platformInfo_; + QLIInfo.opParamInfo = opParamInfo_; + QLIInfo.socVersion = socVersion_; + + QLIInfo.bSize = bSize_; + QLIInfo.n1Size = n1Size_; + QLIInfo.n2Size = n2Size_; + QLIInfo.s1Size = s1Size_; + QLIInfo.s2Size = s2Size_; + QLIInfo.gSize = gSize_; + + QLIInfo.inputQType = inputQType_; + QLIInfo.inputKType = inputKType_; + QLIInfo.outputType = outputType_; + + QLIInfo.blockSize = blockSize_; + QLIInfo.maxBlockNumPerBatch = maxBlockNumPerBatch_; + + QLIInfo.pageAttentionFlag = (kLayout_ == DataLayout::PA_BSND); + QLIInfo.batchSupperFlag = batchSupperFlag_; + QLIInfo.sparseMode = *opParamInfo_.sparseMode; + QLIInfo.sparseCount = *opParamInfo_.sparseCount; + QLIInfo.preTokens = *opParamInfo_.preTokens; + QLIInfo.nextTokens = *opParamInfo_.nextTokens; + QLIInfo.cmpRatio = *opParamInfo_.cmpRatio; + QLIInfo.returnValues = *opParamInfo_.returnValues; + QLIInfo.stride = *opParamInfo_.stride; + QLIInfo.scaleStride = *opParamInfo_.scaleStride; + + QLIInfo.inputQLayout = qLayout_; + QLIInfo.inputKLayout = kLayout_; +} + +ge::graphStatus QLIInfoParser::ParseAndCheck(QLITilingInfo &QLIInfo) +{ + if (ge::GRAPH_SUCCESS != GetOpName() || ge::GRAPH_SUCCESS != GetNpuInfo() || ge::GRAPH_SUCCESS != GetOpParaInfo() || + ge::GRAPH_SUCCESS != CheckRequiredParaExistence()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetAndCheckInOutDataType() || ge::GRAPH_SUCCESS != GetQueryKeyAndOutLayout() || + ge::GRAPH_SUCCESS != GetAndCheckOptionalInput()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != CheckShapeDim() || ge::GRAPH_SUCCESS != GetN1Size() || + ge::GRAPH_SUCCESS != GetAndCheckN2Size() || ge::GRAPH_SUCCESS != GetGSize()) { + return ge::GRAPH_FAILED; + } + + if (ge::GRAPH_SUCCESS != GetBatchSize() || ge::GRAPH_SUCCESS != GetS1Size() || ge::GRAPH_SUCCESS != GetHeadDim() || + ge::GRAPH_SUCCESS != GetS2Size()) { + return ge::GRAPH_FAILED; + } + if (ge::GRAPH_SUCCESS != ValidateInputShapesMatch() || ge::GRAPH_SUCCESS != CheckScaleShape()) { + return ge::GRAPH_FAILED; + } + + GenerateInfo(QLIInfo); + + return ge::GRAPH_SUCCESS; +} + +// --------------------------TilingPrepare函数定义------------------------------------- +static ge::graphStatus TilingPrepareForVllmQuantLightningIndexer(gert::TilingParseContext * /* context */) +{ + return ge::GRAPH_SUCCESS; +} + +// --------------------------VllmQuantLightningIndexerTiling类成员函数定义----------------------- +ge::graphStatus VllmQuantLightningIndexerTiling::DoTiling(QLITilingInfo *tilingInfo) +{ + // -------------set blockdim----------------- + auto ascendcPlatform = platform_ascendc::PlatformAscendC(tilingInfo->platformInfo); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint32_t aicNum = ascendcPlatform.GetCoreNumAic(); + uint32_t blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum); + context_->SetBlockDim(blockDim); + + // -------------set workspacesize----------------- + constexpr uint32_t MM1_RES_ELEM_SIZE = 4; // 4: fp32 + constexpr uint32_t DOUBLE_BUFFER = 2; // 双Buffer + constexpr uint32_t M_BASE_SIZE = 512; // m轴基本块大小 + constexpr uint32_t S2_BASE_SIZE = 512; // S2轴基本块大小 + constexpr uint32_t V1_RES_ELEM_SIZE = 4; // 4: int32 + constexpr uint32_t V1_RES_ELEM_TYPE = 2; // 保留Index和Value 2种数据 + constexpr uint32_t V1_DECODE_PARAM_ELEM_SIZE = 8; // 8: int64 + constexpr uint32_t V1_DECODE_PARAM_NUM = 16; // Decode参数个数 + constexpr uint32_t V1_DECODE_DATA_NUM = 2; // Decode每个核需要存储头和尾部两块数据 + constexpr uint32_t S1_BASE_SIZE = 8; // S1轴基本块的大小 + constexpr uint32_t TOPK_MAX_SIZE = 2048; // TopK选取个数 + constexpr uint32_t ASCEND950_S1_BASE_SIZE = 4; // Ascend 950 S1轴基本块的大小 + constexpr uint32_t ASCEND950_S2_BASE_SIZE = 128; // Ascend 950 S2轴基本块的大小 + uint32_t workspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize(); + // 主流程需Workspace大小 + platform_ascendc::SocVersion socVersion_ = ascendcPlatform.GetSocVersion(); + if (socVersion_ == platform_ascendc::SocVersion::ASCEND950) { + constexpr uint32_t s1Base = ASCEND950_S1_BASE_SIZE; + constexpr uint32_t s2Base = ASCEND950_S2_BASE_SIZE; + workspaceSize += s1Base * ((tilingInfo->s2Size + s2Base - 1) / s2Base) * s2Base * sizeof(uint32_t) * aicNum; + } else { + uint32_t mm1ResSize = M_BASE_SIZE * S2_BASE_SIZE; + workspaceSize += mm1ResSize * MM1_RES_ELEM_SIZE * DOUBLE_BUFFER * aicNum; + // Decode流程(LD)需要Workspace大小 + // 临时存储Decode中间结果大小: 2(头/尾)*8(s1Base)*2(idx/value)*2048(K)*sizeof(int32)*24=6M + workspaceSize += V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_RES_ELEM_TYPE * TOPK_MAX_SIZE * V1_RES_ELEM_SIZE * aicNum; + // 临时存储Decode中间参数信息大小: 2(头/尾)*8(s1Base)*16(paramNum)*sizeof(int64_t)*24=48k + workspaceSize += V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_DECODE_PARAM_NUM * V1_DECODE_PARAM_ELEM_SIZE * aicNum; + } + size_t *workSpaces = context_->GetWorkspaceSizes(1); + workSpaces[0] = workspaceSize; + + // -------------set tilingdata----------------- + tilingData_.set_bSize(tilingInfo->bSize); + tilingData_.set_s2Size(tilingInfo->s2Size); + tilingData_.set_s1Size(tilingInfo->s1Size); + tilingData_.set_sparseCount(tilingInfo->sparseCount); + tilingData_.set_gSize(tilingInfo->gSize); + tilingData_.set_blockSize(tilingInfo->blockSize); + tilingData_.set_maxBlockNumPerBatch(tilingInfo->maxBlockNumPerBatch); + tilingData_.set_sparseMode(tilingInfo->sparseMode); + tilingData_.set_cmpRatio(tilingInfo->cmpRatio); + tilingData_.set_returnValues(tilingInfo->returnValues); + tilingData_.set_usedCoreNum(blockDim); + tilingData_.set_batchSupperFlag(tilingInfo->batchSupperFlag); + tilingData_.set_stride(tilingInfo->stride); + tilingData_.set_scaleStride(tilingInfo->scaleStride); + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + + // -------------set tilingkey----------------- + // DT_Q, DT_KV, DT_OUT, PAGE_ATTENTION, FLASH_DECODE, LAYOUT_T, KV_LAYOUT_T + uint32_t inputQType = static_cast(tilingInfo->inputQType); + uint32_t inputKType = static_cast(tilingInfo->inputKType); + uint32_t outputType = static_cast(tilingInfo->outputType); + uint32_t pageAttentionFlag = static_cast(tilingInfo->pageAttentionFlag); + uint32_t inputQLayout = static_cast(tilingInfo->inputQLayout); + uint32_t inputKLayout = static_cast(tilingInfo->inputKLayout); + uint32_t tilingKey = + GET_TPL_TILING_KEY(inputQType, inputKType, outputType, pageAttentionFlag, inputQLayout, inputKLayout); + context_->SetTilingKey(tilingKey); + context_->SetScheduleMode(1); + + return ge::GRAPH_SUCCESS; +} + +// --------------------------Tiling函数定义--------------------------- +ge::graphStatus TilingForVllmQuantLightningIndexer(gert::TilingContext *context) +{ + OP_CHECK_IF(context == nullptr, OP_LOGE("VllmQuantLightningIndexer", "Tiling context is null."), + return ge::GRAPH_FAILED); + QLITilingInfo QLIInfo; + QLIInfoParser QLIInfoParser(context); + if (QLIInfoParser.ParseAndCheck(QLIInfo) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + VllmQuantLightningIndexerTiling QLITiling(context); + return QLITiling.DoTiling(&QLIInfo); +} + +// --------------------------Tiling及函数TilingPrepare函数注册-------- +IMPL_OP_OPTILING(VllmQuantLightningIndexer) + .Tiling(TilingForVllmQuantLightningIndexer) + .TilingParse(TilingPrepareForVllmQuantLightningIndexer); + +} // namespace optiling diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/vllm_quant_lightning_indexer_tiling.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/vllm_quant_lightning_indexer_tiling.h new file mode 100644 index 000000000..ae169be59 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_host/vllm_quant_lightning_indexer_tiling.h @@ -0,0 +1,263 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vllm_quant_lightning_indexer_tiling.h + * \brief + */ + +#ifndef QUANT_LIGHTNING_INDEXER_TILING_H +#define QUANT_LIGHTNING_INDEXER_TILING_H + +#include "err/ops_err.h" +#include "exe_graph/runtime/tiling_context.h" +#include "platform/platform_info.h" +#include "register/op_def_registry.h" +#include "register/tilingdata_base.h" +#include "tiling/platform/platform_ascendc.h" +#include "tiling/tiling_api.h" + +namespace optiling { +// ------------------公共定义-------------------------- +struct TilingRequiredParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; +}; + +struct TilingOptionalParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::Tensor *tensor; +}; + +enum class DataLayout : uint32_t { + BSND = 0, + TND = 1, + PA_BSND = 2 +}; + +// ------------------算子原型索引常量定义---------------- +// Inputs Index +constexpr uint32_t QUERY_INDEX = 0; +constexpr uint32_t KEY_INDEX = 1; +constexpr uint32_t WEIGTHS_INDEX = 2; +constexpr uint32_t QUERY_DEQUANT_SCALE_INDEX = 3; +constexpr uint32_t KEY_DEQUANT_SCALE_INDEX = 4; +constexpr uint32_t ACTUAL_SEQ_Q_INDEX = 5; +constexpr uint32_t ACTUAL_SEQ_K_INDEX = 6; +constexpr uint32_t BLOCK_TABLE_INDEX = 7; +constexpr uint32_t METADATA_INDEX = 8; +constexpr uint32_t vllm_quant_lightning_indexer = 0; +// Attributes Index +constexpr uint32_t ATTR_QUERY_QUANT_MODE_INDEX = 0; +constexpr uint32_t ATTR_KEY_QUANT_MODE_INDEX = 1; +constexpr uint32_t ATTR_QUERY_LAYOUT_INDEX = 2; +constexpr uint32_t ATTR_KEY_LAYOUT_INDEX = 3; +constexpr uint32_t ATTR_SPARSE_COUNT_INDEX = 4; +constexpr uint32_t ATTR_SPARSE_MODE_INDEX = 5; +constexpr uint32_t ATTR_PRE_TOKENS_INDEX = 6; +constexpr uint32_t ATTR_NEXT_TOKENS_INDEX = 7; +constexpr uint32_t ATTR_CMP_RATIO_INDEX = 8; +constexpr uint32_t ATTR_RETURN_VALUES_INDEX = 9; +constexpr uint32_t ATTR_STRIDE_INDEX = 10; +constexpr uint32_t ATTR_SCALE_STRIDE_INDEX = 11; +// Dim Index +constexpr uint32_t DIM_IDX_ZERO = 0; +constexpr uint32_t DIM_IDX_ONE = 1; +constexpr uint32_t DIM_IDX_TWO = 2; +constexpr uint32_t DIM_IDX_THREE = 3; +// Dim Num +constexpr uint32_t DIM_NUM_TWO = 2; +constexpr uint32_t DIM_NUM_THREE = 3; +constexpr uint32_t DIM_NUM_FOUR = 4; +// 入参限制常量 +constexpr uint32_t HEAD_DIM_LIMIT = 128; +constexpr uint32_t SPARSE_LIMIT = 2048; +constexpr uint32_t G_SIZE_LIMIT = 64; +constexpr uint32_t BLOCK_SIZE_LIMIT = 1024; +constexpr uint32_t BLOCK_SIZE_FACTOR = 16; +constexpr uint32_t SPARSE_MODE_LOWER = 3; +constexpr uint32_t METADATA_LIMIT = 1024; + +// -----------算子TilingData定义--------------- +BEGIN_TILING_DATA_DEF(QLITilingData) +TILING_DATA_FIELD_DEF(uint32_t, bSize) +TILING_DATA_FIELD_DEF(uint32_t, n2Size) +TILING_DATA_FIELD_DEF(uint32_t, gSize) +TILING_DATA_FIELD_DEF(uint32_t, s1Size) +TILING_DATA_FIELD_DEF(uint32_t, s2Size) +TILING_DATA_FIELD_DEF(uint32_t, sparseCount) +TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum) +TILING_DATA_FIELD_DEF(uint32_t, blockSize) +TILING_DATA_FIELD_DEF(uint32_t, maxBlockNumPerBatch) +TILING_DATA_FIELD_DEF(uint32_t, sparseMode) +TILING_DATA_FIELD_DEF(uint32_t, cmpRatio) +TILING_DATA_FIELD_DEF(uint32_t, batchSupperFlag) +TILING_DATA_FIELD_DEF(uint32_t, returnValues) +TILING_DATA_FIELD_DEF(int64_t, stride) +TILING_DATA_FIELD_DEF(int64_t, scaleStride) +END_TILING_DATA_DEF +REGISTER_TILING_DATA_CLASS(VllmQuantLightningIndexer, QLITilingData) + +// -----------算子CompileInfo定义------------------- +struct QLICompileInfo {}; + +// -----------算子Tiling入参结构体定义--------------- +struct QLIParaInfo { + TilingRequiredParaInfo query = {nullptr, nullptr}; + TilingRequiredParaInfo key = {nullptr, nullptr}; + TilingRequiredParaInfo weights = {nullptr, nullptr}; + TilingRequiredParaInfo query_dequant_scale = {nullptr, nullptr}; + TilingRequiredParaInfo key_dequant_scale = {nullptr, nullptr}; + TilingOptionalParaInfo actualSeqLengthsQ = {nullptr, nullptr}; + TilingOptionalParaInfo actualSeqLengthsK = {nullptr, nullptr}; + TilingOptionalParaInfo blockTable = {nullptr, nullptr}; + TilingOptionalParaInfo metadata = {nullptr, nullptr}; + TilingRequiredParaInfo attenOut = {nullptr, nullptr}; + + const int64_t *queryQuantMode = nullptr; + const int64_t *keyQuantMode = nullptr; + const char *layOutQuery = nullptr; + const char *layOutKey = nullptr; + const int64_t *blockSize = nullptr; + const int64_t *sparseMode = nullptr; + const int64_t *sparseCount = nullptr; + const int64_t *preTokens = nullptr; + const int64_t *nextTokens = nullptr; + const int64_t *cmpRatio = nullptr; + const int32_t *batchSupperFlag = nullptr; + const bool *returnValues = nullptr; + const int64_t *stride = nullptr; + const int64_t *scaleStride = nullptr; +}; + +// -----------算子Tiling入参信息类--------------- +class QLITilingInfo { +public: + const char *opName = nullptr; + fe::PlatFormInfos *platformInfo = nullptr; + QLIParaInfo opParamInfo; + // Base Param + platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND910B; + uint32_t bSize = 0; + uint32_t n1Size = 0; + uint32_t n2Size = 0; + uint32_t s1Size = 0; + int64_t s2Size = 0; + uint32_t qkHeadDim = 0; + uint32_t gSize = 0; + // PageAttention + bool pageAttentionFlag = false; + int32_t blockSize = 0; + uint32_t maxBlockNumPerBatch = 0; + // Mask + int32_t sparseMode = 0; + // Others Flag + uint32_t sparseCount = 0; + int64_t preTokens = 0; + int64_t nextTokens = 0; + uint32_t cmpRatio = 1; + bool batchSupperFlag = false; + bool returnValues = false; + int64_t stride = 1; + int64_t scaleStride = 1; + // DType + ge::DataType inputQType = ge::DT_FLOAT16; + ge::DataType inputKType = ge::DT_FLOAT16; + ge::DataType outputType = ge::DT_INT32; + // Layout + DataLayout inputQLayout = DataLayout::BSND; + DataLayout inputKLayout = DataLayout::PA_BSND; +}; + +// -----------算子Tiling入参信息解析及Check类--------------- +class QLIInfoParser { +public: + explicit QLIInfoParser(gert::TilingContext *context) : context_(context) {} + ~QLIInfoParser() = default; + + ge::graphStatus CheckRequiredInOutExistence() const; + ge::graphStatus CheckRequiredAttrExistence() const; + ge::graphStatus CheckRequiredParaExistence() const; + ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor, + const std::string &actualSeqLenName) const; + ge::graphStatus GetOpName(); + ge::graphStatus GetNpuInfo(); + void GetOptionalInputParaInfo(); + void GetInputParaInfo(); + void GetOutputParaInfo(); + ge::graphStatus GetAttrParaInfo(); + ge::graphStatus CheckAttrParaInfo(); + ge::graphStatus GetOpParaInfo(); + ge::graphStatus ValidateInputShapesMatch(); + ge::graphStatus CheckScaleShape(); + ge::graphStatus GetAndCheckInOutDataType(); + ge::graphStatus GetBatchSize(); + ge::graphStatus GetHeadDim(); + ge::graphStatus GetS1Size(); + ge::graphStatus GetAndCheckOptionalInput(); + ge::graphStatus CheckShapeDim(); + ge::graphStatus GetAndCheckBlockSize(); + ge::graphStatus GetS2SizeForPageAttention(); + ge::graphStatus GetS2SizeForBatchContinuous(); + ge::graphStatus GetS2Size(); + ge::graphStatus GetQueryKeyAndOutLayout(); + ge::graphStatus GetN1Size(); + ge::graphStatus GetAndCheckN2Size(); + ge::graphStatus GetGSize(); + ge::graphStatus GetAttenMaskInfo(); + ge::graphStatus GetActualSeqInfo(); + void GenerateInfo(QLITilingInfo &QLIInfo); + ge::graphStatus ParseAndCheck(QLITilingInfo &QLIInfo); + +public: + gert::TilingContext *context_ = nullptr; + const char *opName_; + fe::PlatFormInfos *platformInfo_; + QLIParaInfo opParamInfo_; + + // BaseParams + uint32_t bSize_ = 0; + uint32_t n1Size_ = 0; + uint32_t n2Size_ = 0; + uint32_t gSize_ = 0; + uint32_t s1Size_ = 0; + int64_t s2Size_ = 0; + uint32_t headDim_ = 0; + bool batchSupperFlag_ = false; + // Layout + DataLayout qLayout_ = DataLayout::BSND; + DataLayout kLayout_ = DataLayout::PA_BSND; + // PageAttention + uint32_t maxBlockNumPerBatch_ = 0; + int32_t blockSize_ = 0; + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B; + ge::DataType inputQType_ = ge::DT_FLOAT16; + ge::DataType inputKType_ = ge::DT_FLOAT16; + ge::DataType weightsType_ = ge::DT_FLOAT16; + ge::DataType inputQueryScaleType_ = ge::DT_FLOAT16; + ge::DataType inputKeyScaleType_ = ge::DT_FLOAT16; + ge::DataType blockTableType_ = ge::DT_FLOAT16; + ge::DataType inputKRopeType_ = ge::DT_FLOAT16; + ge::DataType outputType_ = ge::DT_FLOAT16; +}; + +// ---------------算子Tiling类--------------- +class VllmQuantLightningIndexerTiling { +public: + explicit VllmQuantLightningIndexerTiling(gert::TilingContext *context) : context_(context) {}; + ge::graphStatus DoTiling(QLITilingInfo *tilingInfo); + +private: + gert::TilingContext *context_ = nullptr; + QLITilingData tilingData_; +}; + +} // namespace optiling +#endif // QUANT_LIGHTNING_INDEXER_TILING_H \ No newline at end of file diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_common.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_common.h new file mode 100644 index 000000000..3ee29eb27 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_common.h @@ -0,0 +1,148 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_common.h + * \brief + */ +#ifndef QUANT_LIGHTNING_INDEXER_COMMON_H +#define QUANT_LIGHTNING_INDEXER_COMMON_H + +namespace QLICommon { + +// 与tiling的layout保持一致 +enum class LI_LAYOUT : uint32_t { + BSND = 0, + TND = 1, + PA_BSND = 2 +}; + +template +struct QLIType { + using queryType = Q_T; + using keyType = K_T; + using outputType = OUT_T; + static constexpr bool pageAttention = PAGE_ATTENTION; + static constexpr LI_LAYOUT layout = Q_LAYOUT_T; + static constexpr LI_LAYOUT keyLayout = K_LAYOUT_T; +}; + +struct RunInfo { + uint32_t loop; + uint32_t bN2Idx; + uint32_t bIdx; + uint32_t n2Idx = 0; + uint32_t gS1Idx; + uint32_t s2Idx; + + uint32_t actS1Size = 1; + uint32_t actS2Size = 1; + uint32_t actS2SizeOrig = 1; + uint32_t actMBaseSize; + uint32_t actualSingleProcessSInnerSize; + uint32_t actualSingleProcessSInnerSizeAlign; + + uint64_t tensorQueryOffset; + uint64_t tensorKeyOffset; + uint64_t tensorKeyScaleOffset; + uint64_t tensorWeightsOffset; + uint64_t indiceOutOffset; + + bool isFirstS2InnerLoop; + bool isLastS2InnerLoop; + bool isValid = false; +}; + +struct ConstInfo { + // CUBE与VEC核间同步的模式 + static constexpr uint32_t FIA_SYNC_MODE2 = 2; + // BUFFER的字节数 + static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; + static constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64; + static constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256; + static constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512; + static constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024; + static constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048; + static constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096; + static constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192; + static constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384; + static constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768; + // 无效索引 + static constexpr int INVALID_IDX = -1; + + // CUBE和VEC的核间同步EventID + uint32_t syncC1V1 = 0U; + uint32_t syncC1V0 = 2U; + uint32_t syncV1C1 = 0U; + uint32_t syncV0C1 = 1U; + + // 基本块大小 + uint32_t mBaseSize = 1ULL; + uint32_t s1BaseSize = 1ULL; + uint32_t s2BaseSize = 1ULL; + + uint64_t batchSize = 0ULL; + uint64_t gSize = 0ULL; + uint64_t qHeadNum = 0ULL; + uint64_t kHeadNum; + uint64_t headDim; + uint64_t sparseCount; // topK选取大小 + uint64_t kSeqSize = 0ULL; // kv最大S长度 + uint64_t qSeqSize = 1ULL; // q最大S长度 + uint32_t kCacheBlockSize = 0; // PA场景的block size + uint32_t maxBlockNumPerBatch = 0; // PA场景的最大单batch block number + LI_LAYOUT outputLayout; // 输出的格式 + bool attenMaskFlag = false; + uint32_t cmpRatio = 1; // 压缩率 + bool batchSupperFlag = false; // Qactual_se长度是否为B+1 + int64_t stride = 1; + int64_t scaleStride = 1; + + uint32_t actualLenQDims = 0U; // query的actualSeqLength 的维度 + uint32_t actualLenDims = 0U; // KV 的actualSeqLength 的维度 + bool isAccumSeqS1 = false; // 是否累加模式 + bool isAccumSeqS2 = false; // 是否累加模式 + + uint32_t s2Start = 0U; + uint32_t s2End = 0U; + uint32_t bN2Start = 0U; + uint32_t bN2End = 0U; + uint32_t gS1Start = 0U; + uint32_t gS1End = 0U; + uint32_t coreEnable = 0U; +}; + +template +__aicore__ inline T1 Align(T1 num, T2 rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd) * (rnd))); +} + +template +__aicore__ inline T1 Min(T1 a, T2 b) +{ + return (a > b) ? (b) : (a); +} + +template +__aicore__ inline T1 Max(T1 a, T2 b) +{ + return (a > b) ? (a) : (b); +} + +template +__aicore__ inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd))); +} +} // namespace QLICommon + +#endif // QUANT_LIGHTNING_INDEXER_COMMON_H \ No newline at end of file diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_kernel.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_kernel.h new file mode 100644 index 000000000..180580c44 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_kernel.h @@ -0,0 +1,667 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_kernel.h + * \brief + */ + +#ifndef QUANT_LIGHTNING_INDEXER_KERNEL_H +#define QUANT_LIGHTNING_INDEXER_KERNEL_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "quant_lightning_indexer_common.h" +#include "quant_lightning_indexer_service_vector.h" +#include "quant_lightning_indexer_service_cube.h" +#include "../vllm_quant_lightning_indexer_metadata.h" + +namespace QLIKernel { +using namespace QLICommon; +using namespace QLIServiceVec; +using namespace matmul; +using namespace optiling::detail; +using namespace optiling; +using AscendC::CacheMode; +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +// 由于S2循环前，RunInfo还没有赋值，使用TempLoopInfo临时存放B、N、S1轴相关的信息；同时减少重复计算 +struct TempLoopInfo { + uint32_t bN2Idx = 0; + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint32_t gS1Idx = 0U; + uint32_t gS1LoopEnd = 0U; // gS1方向循环的结束Idx + uint32_t s2LoopEnd = 0U; // S2方向循环的结束Idx + uint32_t actS1Size = 1ULL; // 当前Batch循环处理的S1轴的实际大小 + uint32_t actS2Size = 0ULL; + uint32_t actS2SizeOrig = 0ULL; + bool curActSeqLenIsZero = false; + bool needDealActS1LessThanS1 = false; // S1的实际长度小于shape的S1长度时，是否需要清理输出 + uint32_t actMBaseSize = 0U; // m轴(gS1)方向实际大小 + uint32_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小 + uint32_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小 + uint32_t validS2Len = 0U; +}; + +template +class QLIPreload { +public: + __aicore__ inline QLIPreload(){}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale, __gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengthsK, __gm__ uint8_t *blockTable, __gm__ uint8_t *metadata, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *workspace, + const QLITilingData *__restrict tiling, TPipe *tPipe); + __aicore__ inline void Process(); + + // =================================类型定义区================================= + using Q_T = typename QLIT::queryType; + using K_T = typename QLIT::keyType; + using OUT_T = typename QLIT::outputType; + static constexpr bool PAGE_ATTENTION = QLIT::pageAttention; + static constexpr LI_LAYOUT Q_LAYOUT_T = QLIT::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = QLIT::keyLayout; + + using MM1_OUT_T = float; + + QLIMatmul matmulService; + QLIVector vectorService; + + // =================================常量区================================= + static constexpr uint32_t SYNC_C1_V1_FLAG = 4; + static constexpr uint32_t SYNC_V1_C1_FLAG = 5; + + static constexpr uint32_t M_BASE_SIZE = 256; + static constexpr uint32_t S2_BASE_SIZE = 2048; + static constexpr uint32_t HEAD_DIM = 128; + static constexpr uint32_t K_HEAD_NUM = 1; + static constexpr uint32_t GM_ALIGN_BYTES = 512; + static constexpr uint32_t LI_QUANT_PRELOAD_TASK_CACHE_SIZE = 2; + + // for workspace double + static constexpr uint32_t WS_DOBULE = 2; + static constexpr uint32_t ELE_NUM_PER_BLOCK = 16; + +protected: + TPipe *pipe = nullptr; + + // offset + uint64_t queryCoreOffset = 0ULL; + uint64_t keyCoreOffset = 0ULL; + uint64_t keyScaleCoreOffset = 0ULL; + uint64_t weightsCoreOffset = 0ULL; + uint64_t indiceOutCoreOffset = 0ULL; + uint32_t coreZeroEnable = 1U; + + // ================================Global Buffer区================================= + GlobalTensor queryGm; + GlobalTensor keyGm; + GlobalTensor weightsGm; + GlobalTensor metadataGm; + GlobalTensor indiceOutGm; + GlobalTensor blockTableGm; + + GlobalTensor actualSeqLengthsGmQ; + GlobalTensor actualSeqLengthsGm; + + // ================================类成员变量==================================== + // aic、aiv核信息 + uint32_t tmpBlockIdx = 0U; + uint32_t aiCoreIdx = 0U; + + QLICommon::ConstInfo constInfo{}; + TempLoopInfo tempLoopInfo{}; + + // ================================Init functions================================== + __aicore__ inline void InitTilingData(const QLITilingData *__restrict tilingData); + __aicore__ inline void InitBuffers(); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK); + // ================================Split Core================================ + __aicore__ inline void SplitCore(); + __aicore__ inline uint32_t GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, uint32_t actS2SizeOrig, + uint32_t &validS2Len); + __aicore__ inline uint32_t GetTotalBaseBlockNum(); + // ================================Process functions================================ + __aicore__ inline void ProcessMain(); + __aicore__ inline void ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, + QLICommon::RunInfo runInfo[LI_QUANT_PRELOAD_TASK_CACHE_SIZE]); + __aicore__ inline void ProcessInvalid(); + // ================================Params Calc===================================== + __aicore__ inline void CalcGS1LoopParams(uint32_t bN2Idx); + __aicore__ inline void GetBN2Idx(uint32_t bN2Idx); + __aicore__ inline uint32_t GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, uint32_t defaultSeqLen); + __aicore__ inline uint32_t GetActualSeqLenKey(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, uint32_t defaultSeqLen, uint32_t cmpRatio); + __aicore__ inline void GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, uint32_t &actS2Size, uint32_t &actS2SizeOrig); + __aicore__ inline void CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx); + __aicore__ inline void CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, QLICommon::RunInfo &runInfo); + __aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start); +}; + +template +__aicore__ inline void QLIPreload::InitTilingData(const QLITilingData *__restrict tilingData) +{ + constInfo.batchSize = tilingData->bSize; + constInfo.qHeadNum = constInfo.gSize = tilingData->gSize; + constInfo.kSeqSize = tilingData->s2Size; + constInfo.qSeqSize = tilingData->s1Size; + constInfo.attenMaskFlag = (tilingData->sparseMode == 3); + constInfo.kCacheBlockSize = tilingData->blockSize; + constInfo.maxBlockNumPerBatch = tilingData->maxBlockNumPerBatch; + constInfo.sparseCount = tilingData->sparseCount; + constInfo.cmpRatio = tilingData->cmpRatio; + constInfo.batchSupperFlag = tilingData->batchSupperFlag; + constInfo.stride = tilingData->stride; + constInfo.scaleStride = tilingData->scaleStride; + + constInfo.outputLayout = Q_LAYOUT_T; // 输出和输入形状一致 + if (Q_LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS1 = true; + } + if (K_LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS2 = true; + } + + constInfo.kHeadNum = K_HEAD_NUM; + constInfo.headDim = HEAD_DIM; + + constInfo.mBaseSize = M_BASE_SIZE; + constInfo.s2BaseSize = S2_BASE_SIZE; + constInfo.s1BaseSize = (constInfo.mBaseSize + constInfo.gSize - 1) / constInfo.gSize; +} + +template +__aicore__ inline void QLIPreload::InitBuffers() +{ + if ASCEND_IS_AIV { + vectorService.InitBuffers(pipe); + } else { + matmulService.InitBuffers(pipe); + } +} + +template +__aicore__ inline void QLIPreload::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengthsK) +{ + if (actualSeqLengthsQ == nullptr) { + constInfo.actualLenQDims = 0; + } else { + constInfo.actualLenQDims = (constInfo.batchSupperFlag) ? constInfo.batchSize + 1 : constInfo.batchSize; + actualSeqLengthsGmQ.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsQ, constInfo.actualLenQDims); + } + if (actualSeqLengthsK == nullptr) { + constInfo.actualLenDims = 0; + } else { + constInfo.actualLenDims = constInfo.batchSize; + actualSeqLengthsGm.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsK, constInfo.actualLenDims); + } +} + +template +__aicore__ inline uint32_t QLIPreload::GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, + uint32_t defaultSeqLen) +{ + if (actualLenDims == 0) { + return defaultSeqLen; + } else if (constInfo.batchSupperFlag) { + return actualSeqLengthsGm.GetValue(bIdx + 1) - actualSeqLengthsGm.GetValue(bIdx); + } else if (isAccumSeq && bIdx > 0) { + return actualSeqLengthsGm.GetValue(bIdx) - actualSeqLengthsGm.GetValue(bIdx - 1); + } else { + return actualSeqLengthsGm.GetValue(bIdx); + } +} + +template +__aicore__ inline uint32_t QLIPreload::GetActualSeqLenKey(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, + uint32_t defaultSeqLen, uint32_t cmpRatio) +{ + if (actualLenDims == 0) { + return defaultSeqLen * cmpRatio; + } else if (isAccumSeq && bIdx > 0) { + return actualSeqLengthsGm.GetValue(bIdx) - actualSeqLengthsGm.GetValue(bIdx - 1); + } else { + return actualSeqLengthsGm.GetValue(bIdx); + } +} + +template +__aicore__ inline void QLIPreload::GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, uint32_t &actS2Size, uint32_t &actS2SizeOrig) +{ + actS1Size = GetActualSeqLen(bIdx, constInfo.actualLenQDims, constInfo.isAccumSeqS1, actualSeqLengthsGmQ, + constInfo.qSeqSize); + actS2SizeOrig = + GetActualSeqLenKey(bIdx, constInfo.actualLenDims, constInfo.isAccumSeqS2, actualSeqLengthsGm, constInfo.kSeqSize, constInfo.cmpRatio); // 压缩前的actS2Size + actS2Size = actS2SizeOrig / constInfo.cmpRatio; // 真实使用的压缩后S2长度 +} + +template +__aicore__ inline uint32_t QLIPreload::GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, + uint32_t actS2SizeOrig, uint32_t &validS2Len) +{ + if (actS2SizeOrig / constInfo.cmpRatio == 0) { + validS2Len = 0; + return 0; + } + uint32_t s1Offset = constInfo.s1BaseSize * s1gIdx; + int32_t validS2LenBase = static_cast(actS2SizeOrig) - static_cast(actS1Size); // 压缩前的validS2LenBase + validS2Len = (static_cast(s1Offset) + validS2LenBase + static_cast(constInfo.s1BaseSize)) / static_cast(constInfo.cmpRatio); + validS2Len = Min(validS2Len, static_cast(actS2SizeOrig) / constInfo.cmpRatio); + validS2Len = Max(validS2Len, 1); + return (validS2Len + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; +} + +template +__aicore__ inline uint32_t QLIPreload::GetTotalBaseBlockNum() +{ + uint32_t totalBlockNum = 0; + uint32_t actS1Size, actS2Size, actS2SizeOrig; + uint32_t s1GBaseNum, s2BaseNum; + uint32_t validS2Len = 0; + for (uint32_t bIdx = 0; bIdx < constInfo.batchSize; bIdx++) { + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size, actS2SizeOrig); + s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize); + if (!constInfo.attenMaskFlag) { + s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize); + totalBlockNum += s1GBaseNum * s2BaseNum * constInfo.kHeadNum; + continue; + } + for (uint32_t s1gIdx = 0; s1gIdx < s1GBaseNum; s1gIdx++) { + s2BaseNum = GetS2BaseBlockNumOnMask(s1gIdx, actS1Size, actS2SizeOrig, validS2Len); + totalBlockNum += s2BaseNum * constInfo.kHeadNum; + } + } + return totalBlockNum; +} + +// 多核版本，双闭区间。基本原则：计算每个核最少处理的块数, 剩余的部分前面的核每个核多处理一块 +template +__aicore__ void inline QLIPreload::SplitCore() +{ + constInfo.coreEnable = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_CORE_ENABLE_INDEX, false)); + if (aiCoreIdx != 0) { + constInfo.bN2Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_BN2_START_INDEX, false)); + constInfo.gS1Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_M_START_INDEX, false)); + constInfo.s2Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_S2_START_INDEX, false)); + } + constInfo.bN2End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_BN2_END_INDEX, false)); + constInfo.gS1End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_M_END_INDEX, false)); + constInfo.s2End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_S2_END_INDEX, false)); + + // 如果0核都没有启动，说明所有核都没启动 + coreZeroEnable = metadataGm.GetValue(GetAttrAbsIndex(0, LI_CORE_ENABLE_INDEX, false)); +} + +template +__aicore__ inline void QLIPreload::DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start) +{ + if ASCEND_IS_AIV { + if (constInfo.outputLayout == LI_LAYOUT::TND) { + uint32_t tSizeIdx = (constInfo.batchSupperFlag) ? constInfo.batchSize : constInfo.batchSize - 1; + uint32_t tBaseIdx = (constInfo.batchSupperFlag) ? bIdx : bIdx - 1; + uint32_t tSize = actualSeqLengthsGmQ.GetValue(constInfo.batchSize - 1); + uint32_t tBase = bIdx == 0 ? 0 : actualSeqLengthsGmQ.GetValue(tBaseIdx); + uint32_t s1Count = tempLoopInfo.actS1Size; + + for (uint32_t s1Idx = s1Start; s1Idx < s1Count; s1Idx++) { + uint64_t indiceOutOffset = + (tBase + s1Idx) * constInfo.kHeadNum * constInfo.sparseCount + // T轴、s1轴偏移 + n2Idx * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } else if (constInfo.outputLayout == LI_LAYOUT::BSND) { + for (uint32_t s1Idx = s1Start; s1Idx < constInfo.qSeqSize; s1Idx++) { + // B,S1,N2,K + uint64_t indiceOutOffset = bIdx * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount + + s1Idx * constInfo.kHeadNum * constInfo.sparseCount + // B轴、S1轴偏移 + n2Idx * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } + } +} + +template +__aicore__ inline void QLIPreload::Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK, + __gm__ uint8_t *blockTable, __gm__ uint8_t *metadata, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *workspace, + const QLITilingData *__restrict tiling, TPipe *tPipe) +{ + if ASCEND_IS_AIV { + tmpBlockIdx = GetBlockIdx(); // vec:0-47 + aiCoreIdx = tmpBlockIdx / 2; + } else { + tmpBlockIdx = GetBlockIdx(); // cube:0-23 + aiCoreIdx = tmpBlockIdx; + } + + InitTilingData(tiling); + InitActualSeqLen(actualSeqLengthsQ, actualSeqLengthsK); + + if (metadata != nullptr) { + metadataGm.SetGlobalBuffer((__gm__ uint32_t *)metadata); + // 计算分核 + SplitCore(); + } + + pipe = tPipe; + // workspace 内存排布 + // |mm1ResGm(存S) + uint64_t offset = 0; + + // mm1开DoubleBuffer + GlobalTensor mm1ResGm; // 存放S + uint64_t singleCoreMm1ResSize = WS_DOBULE * constInfo.s1BaseSize * constInfo.s2BaseSize * sizeof(MM1_OUT_T); + mm1ResGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + aiCoreIdx * singleCoreMm1ResSize)); + offset += GetBlockNum() * singleCoreMm1ResSize; + + GlobalTensor weightWorkspaceGm; // v1阶段处理w*scale后的结果 + uint64_t weightMemSize = BLOCK_CUBE * constInfo.mBaseSize * WS_DOBULE * sizeof(half); + weightWorkspaceGm.SetGlobalBuffer((__gm__ half *)(workspace + offset + aiCoreIdx * weightMemSize)); + offset += GetBlockNum() * weightMemSize; + + GlobalTensor qScaleGm; + GlobalTensor kScaleGm; + if ASCEND_IS_AIV { + vectorService.InitParams(constInfo, tiling); + indiceOutGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices); + weightsGm.SetGlobalBuffer((__gm__ half *)weights); + qScaleGm.SetGlobalBuffer((__gm__ half *)queryScale); + kScaleGm.SetGlobalBuffer((__gm__ half *)keyScale); + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + vectorService.InitVecInputTensor(weightsGm, qScaleGm, kScaleGm, indiceOutGm, blockTableGm); + vectorService.InitVecWorkspaceTensor(weightWorkspaceGm, mm1ResGm); + } else { + matmulService.InitParams(constInfo); + queryGm.SetGlobalBuffer((__gm__ Q_T *)query); + if constexpr (PAGE_ATTENTION) { + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + } + keyGm.SetGlobalBuffer((__gm__ K_T *)key); + matmulService.InitMm1GlobalTensor(blockTableGm, keyGm, queryGm, mm1ResGm, weightWorkspaceGm); + } + InitBuffers(); +} + +template +__aicore__ inline void QLIPreload::GetBN2Idx(uint32_t bN2Idx) +{ + tempLoopInfo.bN2Idx = bN2Idx; + tempLoopInfo.bIdx = bN2Idx / constInfo.kHeadNum; + tempLoopInfo.n2Idx = bN2Idx % constInfo.kHeadNum; +} + +template +__aicore__ inline void QLIPreload::CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx) +{ + tempLoopInfo.gS1Idx = gS1LoopIdx; + tempLoopInfo.actMBaseSize = constInfo.mBaseSize; + uint32_t remainedGS1Size = tempLoopInfo.actS1Size * constInfo.gSize - tempLoopInfo.gS1Idx * constInfo.mBaseSize; + if (remainedGS1Size <= constInfo.mBaseSize && remainedGS1Size > 0) { + tempLoopInfo.actMBaseSize = tempLoopInfo.mBasicSizeTail; + } + + bool isEnd = (bN2LoopIdx + 1 == constInfo.bN2End) && (gS1LoopIdx + 1 == tempLoopInfo.gS1LoopEnd); + uint32_t s2BlockNum; + uint32_t validS2Len = 0; + if (constInfo.attenMaskFlag) { + s2BlockNum = GetS2BaseBlockNumOnMask(gS1LoopIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2SizeOrig, + tempLoopInfo.validS2Len); + } else { + s2BlockNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + tempLoopInfo.validS2Len = tempLoopInfo.actS2Size; + } + tempLoopInfo.s2LoopEnd = (isEnd && constInfo.s2End != 0) ? constInfo.s2End : s2BlockNum; + tempLoopInfo.s2BasicSizeTail = tempLoopInfo.validS2Len % constInfo.s2BaseSize; + tempLoopInfo.s2BasicSizeTail = (tempLoopInfo.s2BasicSizeTail == 0) ? + constInfo.s2BaseSize : tempLoopInfo.s2BasicSizeTail; +} + +template +__aicore__ inline void QLIPreload::CalcGS1LoopParams(uint32_t bN2LoopIdx) +{ + GetBN2Idx(bN2LoopIdx); + GetS1S2ActualSeqLen(tempLoopInfo.bIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2Size, tempLoopInfo.actS2SizeOrig); + if ((tempLoopInfo.actS2Size == 0) || (tempLoopInfo.actS1Size == 0)) { + tempLoopInfo.curActSeqLenIsZero = true; + return; + } + tempLoopInfo.curActSeqLenIsZero = false; + tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize; + tempLoopInfo.mBasicSizeTail = + (tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail; + + uint32_t gS1SplitNum = (tempLoopInfo.actS1Size * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize; + tempLoopInfo.gS1LoopEnd = (bN2LoopIdx + 1 == constInfo.bN2End && constInfo.gS1End != 0) ? constInfo.gS1End : gS1SplitNum; + if constexpr (Q_LAYOUT_T == LI_LAYOUT::BSND) { + if (tempLoopInfo.gS1LoopEnd == gS1SplitNum && constInfo.qSeqSize > tempLoopInfo.actS1Size) { + tempLoopInfo.needDealActS1LessThanS1 = true; + } + } +} + +template +__aicore__ inline void QLIPreload::CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, QLICommon::RunInfo &runInfo) +{ + runInfo.loop = loop; + runInfo.bIdx = tempLoopInfo.bIdx; + runInfo.gS1Idx = tempLoopInfo.gS1Idx; + runInfo.s2Idx = s2LoopIdx; + runInfo.bN2Idx = tempLoopInfo.bN2Idx; + runInfo.isValid = s2LoopIdx < tempLoopInfo.s2LoopEnd; + + if (!runInfo.isValid) { + return; // 需要验证， v1 时候需要runInfo + } + + runInfo.actS1Size = tempLoopInfo.actS1Size; + runInfo.actS2Size = tempLoopInfo.actS2Size; + runInfo.actS2SizeOrig = tempLoopInfo.actS2SizeOrig; + // 计算实际基本块size + runInfo.actMBaseSize = tempLoopInfo.actMBaseSize; + runInfo.actualSingleProcessSInnerSize = constInfo.s2BaseSize; + uint32_t s2SplitNum = (tempLoopInfo.validS2Len + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + if (runInfo.s2Idx == s2SplitNum - 1) { + runInfo.actualSingleProcessSInnerSize = tempLoopInfo.s2BasicSizeTail; + } + runInfo.actualSingleProcessSInnerSizeAlign = + QLICommon::Align((uint32_t)runInfo.actualSingleProcessSInnerSize, QLICommon::ConstInfo::BUFFER_SIZE_BYTE_32B); + + runInfo.isFirstS2InnerLoop = s2LoopIdx == constInfo.s2Start; + runInfo.isLastS2InnerLoop = (s2LoopIdx + 1 == tempLoopInfo.s2LoopEnd); + + if (runInfo.isFirstS2InnerLoop) { + uint64_t actualSeqQPrefixSum; + if constexpr (Q_LAYOUT_T == LI_LAYOUT::TND) { + uint32_t actualSeqLengthsGmQIdx = (constInfo.batchSupperFlag) ? runInfo.bIdx : runInfo.bIdx - 1; + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGmQ.GetValue(actualSeqLengthsGmQIdx); + } else { // BSND + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.qSeqSize; + } + uint64_t tndBIdxOffset = actualSeqQPrefixSum * constInfo.qHeadNum * constInfo.headDim; + // B,S1,N1(N2,G),D + queryCoreOffset = tndBIdxOffset + runInfo.gS1Idx * constInfo.mBaseSize * constInfo.headDim; + // B,S1,N1(N2,G)/T,N1(N2,G) + weightsCoreOffset = actualSeqQPrefixSum * constInfo.qHeadNum + runInfo.n2Idx * constInfo.gSize; + // B,S1,N2,k/T,N2,k + indiceOutCoreOffset = + actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + runInfo.n2Idx * constInfo.sparseCount; + } + uint64_t actualSeqKPrefixSum; + if constexpr (K_LAYOUT_T == LI_LAYOUT::TND) { // T N2 D + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGm.GetValue(runInfo.bIdx - 1); + } else { + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.kSeqSize; + } + uint64_t tndBIdxOffsetForK = actualSeqKPrefixSum * constInfo.kHeadNum * constInfo.headDim; + keyCoreOffset = tndBIdxOffsetForK + runInfo.s2Idx * constInfo.s2BaseSize * constInfo.kHeadNum * constInfo.headDim; + keyScaleCoreOffset = (actualSeqKPrefixSum + runInfo.s2Idx * constInfo.s2BaseSize) * constInfo.kHeadNum; + runInfo.tensorQueryOffset = queryCoreOffset; + runInfo.tensorKeyOffset = keyCoreOffset; + runInfo.tensorKeyScaleOffset = keyScaleCoreOffset; + runInfo.tensorWeightsOffset = weightsCoreOffset; + runInfo.indiceOutOffset = indiceOutCoreOffset; +} + +template +__aicore__ inline void QLIPreload::Process() +{ + // 没有计算任务，直接清理输出 + if (coreZeroEnable == 0) { + ProcessInvalid(); + return; + } + ProcessMain(); +} + +template +__aicore__ inline void QLIPreload::ProcessInvalid() +{ + if ASCEND_IS_AIV { + uint32_t aivCoreNum = GetBlockNum() * 2; // 2 means c:v = 1:2 + uint64_t totalOutputSize = + constInfo.batchSize * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount; + uint64_t singleCoreSize = + QLICommon::Align((totalOutputSize + aivCoreNum - 1) / aivCoreNum, GM_ALIGN_BYTES / sizeof(OUT_T)); + uint64_t baseSize = tmpBlockIdx * singleCoreSize; + if (baseSize < totalOutputSize) { + uint64_t dealSize = + (baseSize + singleCoreSize <= totalOutputSize) ? singleCoreSize : totalOutputSize - baseSize; + GlobalTensor output = indiceOutGm[baseSize]; + AscendC::InitGlobalMemory(output, dealSize, constInfo.INVALID_IDX); + } + } +} + +template +__aicore__ inline void QLIPreload::ProcessMain() +{ + // 无任务核直接返回 + if (constInfo.coreEnable == 0) { + return; + } + + if ASCEND_IS_AIV { + vectorService.AllocEventID(); + CrossCoreSetFlag(constInfo.syncV1C1); + CrossCoreSetFlag(constInfo.syncV1C1); + } else { + matmulService.AllocEventID(); + CrossCoreSetFlag(constInfo.syncC1V0); + CrossCoreSetFlag(constInfo.syncC1V0); + } + + QLICommon::RunInfo runInfo[LI_QUANT_PRELOAD_TASK_CACHE_SIZE]; + + // 适配左闭右开 + if (constInfo.bN2Start == constInfo.bN2End) { + if (constInfo.gS1Start != constInfo.gS1End || constInfo.s2Start != constInfo.s2End) { + constInfo.bN2End += 1; + } + } else if ((constInfo.gS1End != 0) || (constInfo.s2End != 0)){ + constInfo.bN2End += 1; + } + + uint32_t gloop = 0; + for (uint32_t bN2LoopIdx = constInfo.bN2Start; bN2LoopIdx < constInfo.bN2End; bN2LoopIdx++) { + CalcGS1LoopParams(bN2LoopIdx); + if (tempLoopInfo.curActSeqLenIsZero) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, 0U); + + if ASCEND_IS_AIV { + if (bN2LoopIdx + 1 == constInfo.bN2End && gloop > 0) { + CrossCoreWaitFlag(constInfo.syncC1V1); + vectorService.ProcessVec1(runInfo[1 - gloop % LI_QUANT_PRELOAD_TASK_CACHE_SIZE]); + CrossCoreSetFlag( + constInfo.syncV1C1); // 反向同步 1 + } + } + continue; + } + for (uint32_t gS1LoopIdx = constInfo.gS1Start; gS1LoopIdx < tempLoopInfo.gS1LoopEnd; gS1LoopIdx++) { + CalcS2LoopParams(bN2LoopIdx, gS1LoopIdx); + bool isEnd = (bN2LoopIdx + 1 == constInfo.bN2End) && (gS1LoopIdx + 1 == tempLoopInfo.gS1LoopEnd); + uint32_t extraLoop = isEnd ? LI_QUANT_PRELOAD_TASK_CACHE_SIZE - 1 : 0; // 只preload一轮 + + for (uint32_t s2LoopIdx = constInfo.s2Start; s2LoopIdx < (tempLoopInfo.s2LoopEnd + extraLoop); s2LoopIdx++) { + ProcessBaseBlock(gloop, s2LoopIdx, runInfo); + ++gloop; + } + constInfo.s2Start = 0; + } + if (tempLoopInfo.needDealActS1LessThanS1) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, tempLoopInfo.actS1Size); + } + constInfo.gS1Start = 0; + } + + if ASCEND_IS_AIV { + vectorService.FreeEventID(); + CrossCoreWaitFlag(constInfo.syncC1V0); + CrossCoreWaitFlag(constInfo.syncC1V0); + } else { + matmulService.FreeEventID(); + CrossCoreWaitFlag(constInfo.syncV1C1); + CrossCoreWaitFlag(constInfo.syncV1C1); + } +} + +template +__aicore__ inline void QLIPreload::ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, + QLICommon::RunInfo runInfo[LI_QUANT_PRELOAD_TASK_CACHE_SIZE]) +{ + int32_t curTaskId = loop % LI_QUANT_PRELOAD_TASK_CACHE_SIZE; + QLICommon::RunInfo &curRunInfo = runInfo[curTaskId]; + QLICommon::RunInfo &lastRunInfo = runInfo[1 - curTaskId]; + + CalcRunInfo(loop, s2LoopIdx, curRunInfo); + + if (curRunInfo.isValid) { + if ASCEND_IS_AIC { + if (curRunInfo.isFirstS2InnerLoop) { + CrossCoreWaitFlag(constInfo.syncV0C1); + } + CrossCoreWaitFlag(constInfo.syncV1C1); // 反向同步 1 + matmulService.ComputeMm1(curRunInfo); + CrossCoreSetFlag(constInfo.syncC1V1); + if (curRunInfo.isLastS2InnerLoop) { + CrossCoreSetFlag(constInfo.syncC1V0); // 反向同步 0 + } + } else { + if (curRunInfo.isFirstS2InnerLoop) { + CrossCoreWaitFlag(constInfo.syncC1V0); // 反向同步 0 + vectorService.ProcessVec0(curRunInfo); + CrossCoreSetFlag(constInfo.syncV0C1); + } + } + } + + if (lastRunInfo.isValid) { + if ASCEND_IS_AIV { + CrossCoreWaitFlag(constInfo.syncC1V1); + vectorService.ProcessVec1(lastRunInfo); + CrossCoreSetFlag(constInfo.syncV1C1); // 反向同步 1 + } + lastRunInfo.isValid = false; + } +} +} // namespace QLIKernel +#endif // QUANT_LIGHTNING_INDEXER_KERNEL_H diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_service_cube.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_service_cube.h new file mode 100644 index 000000000..daa58e150 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_service_cube.h @@ -0,0 +1,613 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_service_cube.h + * \brief use 5 buffer for matmul l1, better pipeline + */ +#ifndef QUANT_LIGHTNING_INDEXER_SERVICE_CUBE_H +#define QUANT_LIGHTNING_INDEXER_SERVICE_CUBE_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "quant_lightning_indexer_common.h" + +namespace QLIKernel { +using namespace QLICommon; +struct MmInfo { + int64_t s2L0LoopId; + int64_t s1gL0LoopId; + int64_t s2L0RealSize; + int64_t s2GmOffset; +}; + +template +class QLIMatmul { +public: + using Q_T = typename QLIT::queryType; + using K_T = typename QLIT::keyType; + + __aicore__ inline QLIMatmul(){}; + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitMm1GlobalTensor(const GlobalTensor &blkTableGm, const GlobalTensor &keyGm, + const GlobalTensor &queryGm, const GlobalTensor &mm1ResGm, + const GlobalTensor &weightWorkspaceGm); + __aicore__ inline void InitParams(const ConstInfo &constInfo); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void ComputeMm1(const QLICommon::RunInfo &runInfo); + + static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding; + static constexpr uint64_t DOUBLE_BUF_NUM = 2; + static constexpr uint64_t L0AB_BUF_NUM = 4; + + static constexpr uint32_t KEY_MTE1_MTE2_EVENT = EVENT_ID2; + static constexpr uint32_t QW_MTE1_MTE2_EVENT = EVENT_ID5; // KEY_MTE1_MTE2_EVENT + DOUBLE_BUF_NUM; + static constexpr uint32_t M_MTE1_EVENT = EVENT_ID3; + static constexpr uint32_t M_FIX_EVENT = EVENT_ID0; + static constexpr uint32_t FIX_M_EVENT = EVENT_ID2; + static constexpr uint32_t FIX_MTE1_EVENT = EVENT_ID4; + + static constexpr uint64_t S8_BLOCK_CUBE = 32; + + static constexpr uint32_t MTE2_MTE1_EVENT = EVENT_ID2; + static constexpr uint32_t MTE1_M_EVENT = EVENT_ID2; + + static constexpr uint64_t D_BASIC_BLOCK = 128; + static constexpr uint64_t S1G_BASIC_BLOCK_L1 = 256; + + static constexpr uint64_t S1G_BASIC_BLOCK_L0 = 128; + static constexpr uint64_t S2_BASIC_BLOCK_L0 = 128; + + static constexpr uint64_t QUERY_BUFFER_OFFSET = S1G_BASIC_BLOCK_L1 * D_BASIC_BLOCK; + static constexpr uint64_t SL1_BUFFER_OFFSET = S1G_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0; + static constexpr uint64_t KEY_BUFFER_OFFSET = S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK; + static constexpr uint64_t WEIGHT_BUFFER_OFFSET = S1G_BASIC_BLOCK_L1 * BLOCK_CUBE; + static constexpr uint64_t L0AB_BUFFER_OFFSET_S8_16K = 16 * 1024; + static constexpr uint64_t L0AB_BUFFER_OFFSET_FP16_16K = 16 * 512; + static constexpr uint64_t L0C_BUFFER_OFFSET = 64 * 256; + +private: + __aicore__ inline void WeightDmaCopy(uint64_t s1gL1RealSize, const QLICommon::RunInfo &runInfo); + __aicore__ inline void LoadKeyToL0b(uint64_t s2L0RealSize); + __aicore__ inline void LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL1RealSize, uint64_t s1gL0RealSize); + __aicore__ inline void QueryNd2Nz(uint64_t s1gL1RealSize, const QLICommon::RunInfo &runInfo); + __aicore__ inline void KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, const QLICommon::RunInfo &runInfo); + __aicore__ inline void KeyNd2Nz(uint64_t s2L1RealSize, const MmInfo &mmInfo, const QLICommon::RunInfo &runInfo); + __aicore__ inline void FixpSToL1(uint64_t s1gL0RealSize, uint64_t s2L0RealSize); + __aicore__ inline void LoadSToL0b(uint64_t s1gL1RealSize, uint64_t s2L0RealSize, uint64_t sL1BufIdx, + int64_t mStartPt); + __aicore__ inline void LoadWeightToL0a(uint64_t s1gL1Offset); + __aicore__ inline void ComputeWs(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, int64_t s1gOffset); + __aicore__ inline void FixpResToGm(uint64_t s1L0RealCount, uint64_t s2L0RealSize, uint64_t s1GmOffset, + uint64_t s2GmOffset, const QLICommon::RunInfo &runInfo); + __aicore__ inline void ComputeQk(uint64_t s1gL0RealSize, uint64_t s2L0RealSize); + __aicore__ inline void ProcessWs(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t sL1BufIdx, + const MmInfo &mmInfo, const QLICommon::RunInfo &runInfo); + __aicore__ inline void ProcessQk(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t s1L0LoopCnt, + const MmInfo &mmInfo, const QLICommon::RunInfo &runInfo); + __aicore__ inline void CalcMmInfo(MmInfo &mmInfo, uint64_t loopIdx, uint64_t s1L0LoopCnt, const MmInfo &lastMmInfo, + const QLICommon::RunInfo &runInfo); + static constexpr LI_LAYOUT Q_LAYOUT_T = QLIT::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = QLIT::keyLayout; + GlobalTensor blkTableGm_; + GlobalTensor keyGm_; + GlobalTensor queryGm_; + GlobalTensor weightGm_; + GlobalTensor mm1ResGm_; + + TBuf bufQL1_; + LocalTensor queryL1_; + TBuf bufKeyL1_; + LocalTensor keyL1_; + TBuf bufWeightL1_; + LocalTensor weightL1_; + TBuf bufSL1_; + LocalTensor sL1_; + + TBuf bufL0A_; + LocalTensor l0a_; + TBuf bufL0B_; + LocalTensor l0b_; + + TBuf bufL0C_; + LocalTensor cL0_; + + uint64_t keyL1BufIdx_ = 0; + uint64_t qwL1Mte2BufIdx_ = 0; + uint64_t sL1BufIdx_ = 0; + uint64_t l0BufIdx_ = 0; + uint64_t l0cBufIdx_ = 0; + + ConstInfo constInfo_; +}; + +template +__aicore__ inline void QLIMatmul::InitParams(const ConstInfo &constInfo) +{ + constInfo_ = constInfo; +} + +template +__aicore__ inline void QLIMatmul::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(bufQL1_, DOUBLE_BUF_NUM * S1G_BASIC_BLOCK_L1 * D_BASIC_BLOCK * sizeof(Q_T)); + queryL1_ = bufQL1_.Get(); + pipe->InitBuffer(bufKeyL1_, DOUBLE_BUF_NUM * S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK * sizeof(K_T)); + keyL1_ = bufKeyL1_.Get(); + + pipe->InitBuffer(bufWeightL1_, DOUBLE_BUF_NUM * S1G_BASIC_BLOCK_L1 * BLOCK_CUBE * sizeof(half)); + weightL1_ = bufWeightL1_.Get(); + pipe->InitBuffer(bufSL1_, DOUBLE_BUF_NUM * S2_BASIC_BLOCK_L0 * S1G_BASIC_BLOCK_L0 * sizeof(half)); + sL1_ = bufSL1_.Get(); + + pipe->InitBuffer(bufL0A_, 64 * 1024); + l0a_ = bufL0A_.Get(); + pipe->InitBuffer(bufL0B_, 64 * 1024); + l0b_ = bufL0B_.Get(); + + pipe->InitBuffer(bufL0C_, 128 * 1024); + cL0_ = bufL0C_.Get(); +} + +template +__aicore__ inline void QLIMatmul::InitMm1GlobalTensor(const GlobalTensor &blkTableGm, + const GlobalTensor &keyGm, + const GlobalTensor &queryGm, + const GlobalTensor &mm1ResGm, + const GlobalTensor &weightWorkspaceGm) +{ + blkTableGm_ = blkTableGm; + keyGm_ = keyGm; + queryGm_ = queryGm; + mm1ResGm_ = mm1ResGm; + weightGm_ = weightWorkspaceGm; +} + +template +__aicore__ inline void QLIMatmul::ProcessWs(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t sL1BufIdx, + const MmInfo &mmInfo, const QLICommon::RunInfo &runInfo) +{ + WaitFlag(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM); + for (int64_t s1gOffset = 0; s1gOffset < s1gL0RealSize; s1gOffset += constInfo_.gSize) { + WaitFlag(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM); + LoadSToL0b(s1gL0RealSize, mmInfo.s2L0RealSize, sL1BufIdx, s1gOffset); + LoadWeightToL0a(s1gOffset + s1gL1Offset); + + ComputeWs(s1gL0RealSize, mmInfo.s2L0RealSize, s1gOffset); + + SetFlag(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM); + l0BufIdx_++; + } + + FixpResToGm(s1gL0RealSize / constInfo_.gSize, mmInfo.s2L0RealSize, s1gL1Offset / constInfo_.gSize, + mmInfo.s2L0LoopId * S2_BASIC_BLOCK_L0, runInfo); + SetFlag(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM); + l0cBufIdx_++; +} + +template +__aicore__ inline void QLIMatmul::ProcessQk(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t s1L0LoopCnt, + const MmInfo &mmInfo, const QLICommon::RunInfo &runInfo) +{ + if (mmInfo.s1gL0LoopId == 0) { + WaitFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % DOUBLE_BUF_NUM); + if constexpr (K_LAYOUT_T == LI_LAYOUT::PA_BSND) { + KeyNd2NzForPA(mmInfo.s2L0RealSize, runInfo.s2Idx * constInfo_.s2BaseSize + mmInfo.s2GmOffset, runInfo); + } else { + KeyNd2Nz(mmInfo.s2L0RealSize, mmInfo, runInfo); + } + + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + } + + WaitFlag(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM); + LoadQueryToL0a(s1gL1Offset, runInfo.actMBaseSize, s1gL0RealSize); + LoadKeyToL0b(mmInfo.s2L0RealSize); + + if (mmInfo.s1gL0LoopId + 1 >= s1L0LoopCnt) { + SetFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % DOUBLE_BUF_NUM); + keyL1BufIdx_++; + } + + WaitFlag(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM); + ComputeQk(s1gL0RealSize, mmInfo.s2L0RealSize); + SetFlag(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM); + + FixpSToL1(s1gL0RealSize, mmInfo.s2L0RealSize); + SetFlag(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM); + l0BufIdx_++; + l0cBufIdx_++; +} + +template +__aicore__ inline void QLIMatmul::CalcMmInfo(MmInfo &mmInfo, uint64_t loopIdx, uint64_t s1L0LoopCnt, + const MmInfo &lastMmInfo, const QLICommon::RunInfo &runInfo) +{ + mmInfo.s2L0LoopId = loopIdx / s1L0LoopCnt; + mmInfo.s1gL0LoopId = loopIdx % s1L0LoopCnt; + + if (mmInfo.s1gL0LoopId == 0) { + mmInfo.s2GmOffset = mmInfo.s2L0LoopId * S2_BASIC_BLOCK_L0; + mmInfo.s2L0RealSize = mmInfo.s2GmOffset + S2_BASIC_BLOCK_L0 > runInfo.actualSingleProcessSInnerSize + ? runInfo.actualSingleProcessSInnerSize - mmInfo.s2GmOffset + : S2_BASIC_BLOCK_L0; + } else { + mmInfo.s2L0RealSize = lastMmInfo.s2L0RealSize; + } +} + +template +__aicore__ inline void QLIMatmul::ComputeMm1(const QLICommon::RunInfo &runInfo) +{ + if (runInfo.isFirstS2InnerLoop) { + WaitFlag(QW_MTE1_MTE2_EVENT + qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM); + QueryNd2Nz(runInfo.actMBaseSize, runInfo); // 256 * 128 // L1BasicBlock + WeightDmaCopy(runInfo.actMBaseSize, runInfo); + } + int64_t loopIdx = 0; + int64_t s2L0LoopCnt = CeilDiv(runInfo.actualSingleProcessSInnerSize, S2_BASIC_BLOCK_L0); // 2048取128 + int64_t s1L0LoopCnt = CeilDiv(runInfo.actMBaseSize, S1G_BASIC_BLOCK_L0); // 256取128 + int64_t s1gL1Offset[2] = {0, static_cast(S1G_BASIC_BLOCK_L0)}; + int64_t s1gL0RealSize[2] = {s1L0LoopCnt > 1 ? static_cast(S1G_BASIC_BLOCK_L0) : runInfo.actMBaseSize, + runInfo.actMBaseSize - s1gL1Offset[1]}; + MmInfo mmInfo[2]; + CalcMmInfo(mmInfo[loopIdx & 1], loopIdx, s1L0LoopCnt, mmInfo[(loopIdx + 1) & 1], runInfo); + + ProcessQk(s1gL0RealSize[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], + s1gL1Offset[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], s1L0LoopCnt, mmInfo[loopIdx & 1], + runInfo); + + SetFlag(FIX_MTE1_EVENT + sL1BufIdx_ % DOUBLE_BUF_NUM); + sL1BufIdx_++; + loopIdx++; + + while (loopIdx < s2L0LoopCnt * s1L0LoopCnt) { + CalcMmInfo(mmInfo[loopIdx & 1], loopIdx, s1L0LoopCnt, mmInfo[(loopIdx + 1) & 1], runInfo); + + ProcessQk(s1gL0RealSize[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], + s1gL1Offset[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], s1L0LoopCnt, mmInfo[loopIdx & 1], + runInfo); + + SetFlag(FIX_MTE1_EVENT + sL1BufIdx_ % DOUBLE_BUF_NUM); + sL1BufIdx_++; + + WaitFlag(FIX_MTE1_EVENT + sL1BufIdx_ % DOUBLE_BUF_NUM); + + ProcessWs(s1gL0RealSize[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], + s1gL1Offset[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], sL1BufIdx_, + mmInfo[(loopIdx + 1) & 1], runInfo); + loopIdx++; + } + + WaitFlag(FIX_MTE1_EVENT + (sL1BufIdx_ + 1) % DOUBLE_BUF_NUM); + + ProcessWs(s1gL0RealSize[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], + s1gL1Offset[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], sL1BufIdx_ - 1, + mmInfo[(loopIdx + 1) & 1], runInfo); + + if (runInfo.isLastS2InnerLoop) { + SetFlag(QW_MTE1_MTE2_EVENT + qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM); + qwL1Mte2BufIdx_++; + } +} + +// blkNum, blkSize, N2, D +template +__aicore__ inline void QLIMatmul::KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const QLICommon::RunInfo &runInfo) +{ + uint64_t s2L1Offset = 0; + while (s2L1Offset < s2L1RealSize) { + uint64_t s2BlkId = (s2L1Offset + s2GmOffset) / constInfo_.kCacheBlockSize; + uint64_t s2BlkOffset = (s2L1Offset + s2GmOffset) % constInfo_.kCacheBlockSize; + uint64_t keyGmOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId) * + constInfo_.stride + + s2BlkOffset * constInfo_.headDim; + uint64_t s2Mte2Size = s2L1RealSize - s2L1Offset; + s2Mte2Size = s2BlkOffset + s2Mte2Size >= constInfo_.kCacheBlockSize ? constInfo_.kCacheBlockSize - s2BlkOffset + : s2Mte2Size; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2Mte2Size; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(keyL1_[(keyL1BufIdx_ % DOUBLE_BUF_NUM) * KEY_BUFFER_OFFSET + s2L1Offset * S8_BLOCK_CUBE], + keyGm_[keyGmOffset], nd2nzPara); + + s2L1Offset += s2Mte2Size; + } +} + +template +__aicore__ inline void QLIMatmul::KeyNd2Nz(uint64_t s2L1RealSize, const MmInfo &mmInfo, + const QLICommon::RunInfo &runInfo) +{ + uint64_t dStride = constInfo_.headDim; + if constexpr (K_LAYOUT_T == LI_LAYOUT::BSND || K_LAYOUT_T == LI_LAYOUT::TND) { + dStride = constInfo_.headDim * constInfo_.kHeadNum; // constInfo_.kHeadNum + } + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2L1RealSize; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = dStride; + nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + // 默认一块buf最多放两份 + DataCopy(keyL1_[(keyL1BufIdx_ % DOUBLE_BUF_NUM) * KEY_BUFFER_OFFSET], + keyGm_[runInfo.tensorKeyOffset + mmInfo.s2GmOffset * constInfo_.headDim], nd2nzPara); +} + +// batch, s1, g, 1 +template +__aicore__ inline void QLIMatmul::WeightDmaCopy(uint64_t s1gL1RealSize, const QLICommon::RunInfo &runInfo) +{ + DataCopyParams copyInParams; + copyInParams.blockCount = 1; + copyInParams.blockLen = s1gL1RealSize; + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + DataCopy(weightL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * WEIGHT_BUFFER_OFFSET], + weightGm_[runInfo.loop % DOUBLE_BUF_NUM * BLOCK_CUBE * constInfo_.mBaseSize], copyInParams); +} + +// batch, s1, n2, g, d +template +__aicore__ inline void QLIMatmul::QueryNd2Nz(uint64_t s1gL1RealSize, const QLICommon::RunInfo &runInfo) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s1gL1RealSize; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(s1gL1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + // 默认一块buf最多放两份 + DataCopy(queryL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * QUERY_BUFFER_OFFSET], queryGm_[runInfo.tensorQueryOffset], + nd2nzPara); +} + +// s1g, d +template +__aicore__ inline void QLIMatmul::LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL1RealSize, + uint64_t s1gL0RealSize) +{ + LoadData3DParamsV2 loadData3DParams; + // SetFmatrixParams + loadData3DParams.l1H = CeilDiv(s1gL1RealSize, BLOCK_CUBE); // Hin=M1=8 + loadData3DParams.l1W = BLOCK_CUBE; // Win=M0 + loadData3DParams.channelSize = constInfo_.headDim; // Cin=K + + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + // SetLoadToA0Params + loadData3DParams.mExtension = CeilAlign(s1gL0RealSize, BLOCK_CUBE); // M height维度目的 + loadData3DParams.kExtension = constInfo_.headDim; // K width维度目的 + loadData3DParams.mStartPt = s1gL1Offset; + loadData3DParams.kStartPt = 0; + loadData3DParams.strideW = 1; + loadData3DParams.strideH = 1; + loadData3DParams.filterW = 1; + loadData3DParams.filterSizeW = (1 >> 8) & 255; + loadData3DParams.filterH = 1; + loadData3DParams.filterSizeH = (1 >> 8) & 255; + loadData3DParams.dilationFilterW = 1; + loadData3DParams.dilationFilterH = 1; + loadData3DParams.enTranspose = 0; + loadData3DParams.fMatrixCtrl = 0; + + LoadData(l0a_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K], + queryL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * QUERY_BUFFER_OFFSET], + loadData3DParams); +} + +// s1, g, s2 --> 2 * 64* 128 +template +__aicore__ inline void QLIMatmul::LoadSToL0b(uint64_t s1gL1RealSize, uint64_t s2L0RealSize, uint64_t sL1BufIdx, + int64_t mStartPt) +{ + LoadData3DParamsV2 loadData3DParams; + // SetFmatrixParams + loadData3DParams.l1H = S1G_BASIC_BLOCK_L0 / BLOCK_CUBE; // Hin=M1=8 + loadData3DParams.l1W = BLOCK_CUBE; // Win=M0 + loadData3DParams.channelSize = CeilAlign(s2L0RealSize, BLOCK_CUBE); // Cin=K + + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + // SetLoadToA0Params + loadData3DParams.mExtension = constInfo_.gSize; // M height维度目的 + loadData3DParams.kExtension = CeilAlign(s2L0RealSize, BLOCK_CUBE); // K width维度目的 + loadData3DParams.kStartPt = 0; + loadData3DParams.strideW = 1; + loadData3DParams.strideH = 1; + loadData3DParams.filterW = 1; + loadData3DParams.filterSizeW = (1 >> 8) & 255; + loadData3DParams.filterH = 1; + loadData3DParams.filterSizeH = (1 >> 8) & 255; + loadData3DParams.dilationFilterW = 1; + loadData3DParams.dilationFilterH = 1; + loadData3DParams.enTranspose = 1; + loadData3DParams.fMatrixCtrl = 0; + + loadData3DParams.mStartPt = mStartPt; + LoadData( + l0b_.template ReinterpretCast()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K], + sL1_[(sL1BufIdx % DOUBLE_BUF_NUM) * SL1_BUFFER_OFFSET], loadData3DParams); +} + +// s1,g,1(16), 2,64,16 +template +__aicore__ inline void QLIMatmul::LoadWeightToL0a(uint64_t s1gL1Offset) +{ + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = CeilDiv(constInfo_.gSize, BLOCK_CUBE); + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = true; + LoadData(l0a_.template ReinterpretCast()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K], + weightL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * WEIGHT_BUFFER_OFFSET + s1gL1Offset* BLOCK_CUBE], + loadData2DParams); +} + +// s2, d -> 128,128 +template +__aicore__ inline void QLIMatmul::LoadKeyToL0b(uint64_t s2L0RealSize) +{ + LoadData2DParams loadData2DParams; + loadData2DParams.startIndex = 0; + loadData2DParams.repeatTimes = CeilDiv(s2L0RealSize, BLOCK_CUBE) * CeilDiv(constInfo_.headDim, S8_BLOCK_CUBE); + loadData2DParams.srcStride = 1; + loadData2DParams.dstGap = 0; + loadData2DParams.ifTranspose = false; + LoadData(l0b_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K], + keyL1_[(keyL1BufIdx_ % DOUBLE_BUF_NUM) * KEY_BUFFER_OFFSET], loadData2DParams); +} + +// A: s1,g,1(16) B: s1,g,s2 C: s1, 1(16), s2 +template +__aicore__ inline void QLIMatmul::ComputeWs(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, int64_t s1gOffset) +{ + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + MmadParams mmadParams; + mmadParams.m = BLOCK_CUBE; + mmadParams.n = s2L0RealSize; + mmadParams.k = constInfo_.gSize; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + Mmad(cL0_.template ReinterpretCast()[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET + + s1gOffset * S2_BASIC_BLOCK_L0], + l0a_.template ReinterpretCast()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K], + l0b_.template ReinterpretCast()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K], + mmadParams); +} + +template +__aicore__ inline void QLIMatmul::ComputeQk(uint64_t s1gL0RealSize, uint64_t s2L0RealSize) +{ + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + + MmadParams mmadParams; + mmadParams.m = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + mmadParams.n = s2L0RealSize; + mmadParams.k = constInfo_.headDim; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + Mmad(cL0_[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET], + l0a_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K], + l0b_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K], mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } +} + +template +__aicore__ inline void QLIMatmul::FixpSToL1(uint64_t s1gL0RealSize, uint64_t s2L0RealSize) +{ + SetFlag(M_FIX_EVENT); + WaitFlag(M_FIX_EVENT); + DataCopyCO12DstParams params; + params.mSize = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + params.nSize = CeilAlign(s2L0RealSize, BLOCK_CUBE); + params.dstStride = S1G_BASIC_BLOCK_L0; + params.srcStride = params.mSize; + params.quantPre = QuantMode_t::DEQF16; + params.reluPre = 1; + params.channelSplit = 0; + params.nz2ndEn = 0; + SetFixpipePreQuantFlag(0x3a800000); + DataCopy(sL1_[(sL1BufIdx_ % DOUBLE_BUF_NUM) * SL1_BUFFER_OFFSET], + cL0_[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET], params); +} + +template +__aicore__ inline void QLIMatmul::FixpResToGm(uint64_t s1L0RealCount, uint64_t s2L0RealSize, uint64_t s1GmOffset, + uint64_t s2GmOffset, const QLICommon::RunInfo &runInfo) +{ + SetFlag(M_FIX_EVENT); + WaitFlag(M_FIX_EVENT); + + AscendC::DataCopyCO12DstParams intriParams; + intriParams.mSize = 1; + intriParams.nSize = s2L0RealSize; + intriParams.dstStride = constInfo_.s2BaseSize; + intriParams.srcStride = 16; + // set mode according to dtype + intriParams.quantPre = QuantMode_t::NoQuant; + intriParams.nz2ndEn = true; + intriParams.reluPre = 0; + AscendC::SetFixpipeNz2ndFlag(s1L0RealCount, CeilDiv(constInfo_.gSize, BLOCK_CUBE) * S2_BASIC_BLOCK_L0 / BLOCK_CUBE, + 2048); + AscendC::DataCopy(mm1ResGm_[(runInfo.loop % 2) * constInfo_.mBaseSize / constInfo_.gSize * constInfo_.s2BaseSize + + s1GmOffset * intriParams.dstStride + s2GmOffset], + cL0_.template ReinterpretCast()[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET], + intriParams); +} + +template +__aicore__ inline void QLIMatmul::AllocEventID() +{ + SetFlag(KEY_MTE1_MTE2_EVENT + 0); + SetFlag(KEY_MTE1_MTE2_EVENT + 1); + SetFlag(KEY_MTE1_MTE2_EVENT + 2); + + SetFlag(QW_MTE1_MTE2_EVENT + 0); + SetFlag(QW_MTE1_MTE2_EVENT + 1); + + SetFlag(M_MTE1_EVENT + 0); + SetFlag(M_MTE1_EVENT + 1); + SetFlag(M_MTE1_EVENT + 2); + SetFlag(M_MTE1_EVENT + 3); + + SetFlag(FIX_M_EVENT + 0); + SetFlag(FIX_M_EVENT + 1); +} + +template +__aicore__ inline void QLIMatmul::FreeEventID() +{ + WaitFlag(KEY_MTE1_MTE2_EVENT + 0); + WaitFlag(KEY_MTE1_MTE2_EVENT + 1); + WaitFlag(KEY_MTE1_MTE2_EVENT + 2); + + WaitFlag(QW_MTE1_MTE2_EVENT + 0); + WaitFlag(QW_MTE1_MTE2_EVENT + 1); + + WaitFlag(M_MTE1_EVENT + 0); + WaitFlag(M_MTE1_EVENT + 1); + WaitFlag(M_MTE1_EVENT + 2); + WaitFlag(M_MTE1_EVENT + 3); + + WaitFlag(FIX_M_EVENT + 0); + WaitFlag(FIX_M_EVENT + 1); +} +} // namespace QLIKernel +#endif // QUANT_LIGHTNING_INDEXER_SERVICE_CUBE_H \ No newline at end of file diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_service_vector.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_service_vector.h new file mode 100644 index 000000000..192022ac4 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_service_vector.h @@ -0,0 +1,437 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_service_vector.h + * \brief + */ +#ifndef QUANT_LIGHTNING_INDEXER_SERVICE_VECTOR_H +#define QUANT_LIGHTNING_INDEXER_SERVICE_VECTOR_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "quant_lightning_indexer_common.h" +#include "quant_lightning_indexer_vector.h" + +namespace QLIKernel { +using namespace QLICommon; +using namespace QLIServiceVec; +constexpr uint32_t BASE_TOPK = 2048; +constexpr uint32_t BASE_TOPK_VALUE_IDX_SIZE = 4096; +constexpr uint32_t ELE_NUM_32 = 32; +constexpr uint32_t ELE_NUM_128 = 128; +constexpr uint32_t ELE_NUM_512 = 512; + +template +class QLIVector { +public: + // =================================类型定义区================================= + static constexpr LI_LAYOUT Q_LAYOUT_T = QLIT::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = QLIT::keyLayout; + static constexpr bool PAGE_ATTENTION = QLIT::pageAttention; + // MM输出数据类型, 当前只支持float + using MM1_OUT_T = float; + + __aicore__ inline QLIVector(){}; + __aicore__ inline void ProcessVec0(const QLICommon::RunInfo &info); + __aicore__ inline void ProcessVec1(const QLICommon::RunInfo &info); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitParams(const struct QLICommon::ConstInfo &constInfo, + const QLITilingData *__restrict tilingData); + __aicore__ inline void InitVecWorkspaceTensor(GlobalTensor vec0OutGm, GlobalTensor mm1ResGm); + __aicore__ inline void InitVecInputTensor(GlobalTensor weightsGm, GlobalTensor qScaleGm, + GlobalTensor kScaleGm, GlobalTensor indiceOutGm, + GlobalTensor blockTableGm); + __aicore__ inline void CleanInvalidOutput(int64_t invalidS1offset); + __aicore__ inline int32_t AlignS2(int32_t cuS2Len); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + +protected: + GlobalTensor mm1ResGm; + GlobalTensor weightsGm; + GlobalTensor qScaleGm; + GlobalTensor kScaleGm; + GlobalTensor vec0OutGm; + GlobalTensor indiceOutGm; + GlobalTensor blockTableGm; + // =================================常量区================================= + +private: + __aicore__ inline void GetKeyScale(const QLICommon::RunInfo &runInfo, const LocalTensor &resUb, + int64_t batchId, int64_t startS2, int64_t getLen); + // ================================Local Buffer区==================================== + // queue + TQue inQueue_; + TQue outQueue_; + + // tmp buff for vector + TBuf sortOutBuf_; + TBuf indexBuf_; + TBuf tmpBuf_; + + LocalTensor globalTopkIndice_; + LocalTensor globalTopkUb_; + + int32_t blockId_ = -1; + // para for vector + int32_t groupInner_ = 0; + int32_t globalTopkNum_ = 0; + int64_t blockS2StartIdx_ = 0; + int32_t gSize_ = 0; + int32_t kSeqSize_ = 0; + int32_t kHeadNum_ = 0; + int32_t qHeadNum_ = 0; + int32_t s1BaseSize_ = 0; + int32_t s2BaseSize_ = 0; + int32_t kCacheBlockSize_ = 0; + int32_t maxBlockNumPerBatch_ = 0; + + struct QLICommon::ConstInfo constInfo_; +}; + +template +__aicore__ inline void QLIVector::GetKeyScale(const QLICommon::RunInfo &runInfo, const LocalTensor &resUb, + int64_t batchId, int64_t startS2, int64_t getLen) +{ + // startS2一定能整除kCacheBlockSize_ + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams copyInParams; + if constexpr (PAGE_ATTENTION) { + int32_t startBlockTableIdx = startS2 / kCacheBlockSize_; + int32_t startBlockTableOffset = startS2 % kCacheBlockSize_; + int32_t blockTableBatchOffset = batchId * maxBlockNumPerBatch_; + copyInParams.blockCount = 1; + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + int32_t resUbBaseOffset = 0; + if (startBlockTableOffset > 0) { + int32_t firstPartLen = + kCacheBlockSize_ - startBlockTableOffset > getLen ? getLen : kCacheBlockSize_ - startBlockTableOffset; + copyInParams.blockLen = firstPartLen * sizeof(half); + int32_t blockId = blockTableGm.GetValue(blockTableBatchOffset + startBlockTableIdx); + SetWaitFlag(HardEvent::S_MTE2); + AscendC::DataCopyPad(resUb, kScaleGm[blockId * constInfo_.scaleStride + startBlockTableOffset], + copyInParams, padParams); + startBlockTableIdx++; + getLen = getLen - firstPartLen; + resUbBaseOffset = firstPartLen; + } + int32_t getLoopNum = CeilDiv(getLen, kCacheBlockSize_); + copyInParams.blockLen = kCacheBlockSize_ * sizeof(half); + for (int32_t i = 0; i < getLoopNum; i++) { + if (i == getLoopNum - 1) { + copyInParams.blockLen = (getLen - i * kCacheBlockSize_) * sizeof(half); + } + int32_t blockId = blockTableGm.GetValue(blockTableBatchOffset + startBlockTableIdx + i); + SetWaitFlag(HardEvent::S_MTE2); + AscendC::DataCopyPad(resUb[resUbBaseOffset + i * kCacheBlockSize_], kScaleGm[blockId * constInfo_.scaleStride], + copyInParams, padParams); + } + } else { + copyInParams.blockCount = 1; + copyInParams.blockLen = getLen * sizeof(half); + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + AscendC::DataCopyPad(resUb, kScaleGm[runInfo.tensorKeyScaleOffset], copyInParams, padParams); + } +} + +template +__aicore__ inline void QLIVector::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(inQueue_, 2, s2BaseSize_ * sizeof(float) * 2); // 32KB + pipe->InitBuffer(outQueue_, 1, BASE_TOPK * sizeof(float)); // 8 KB + pipe->InitBuffer(indexBuf_, s2BaseSize_ * sizeof(int32_t)); // 8 KB + pipe->InitBuffer(tmpBuf_, 64 * 1024); // 64KB + pipe->InitBuffer(sortOutBuf_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK_VALUE_IDX_SIZE * sizeof(float)); // 32KB + + globalTopkIndice_ = indexBuf_.Get(); + globalTopkUb_ = sortOutBuf_.Get(); + globalTopkNum_ = 0; + + // 基本块执行前初始化UB和GM + // step1. 初始化一个有序索引 0 - s2BaseSize_ + ArithProgression(globalTopkIndice_, 0, 1, s2BaseSize_); + // step2. globalTopkUb_ [CeilDiv(s1BaseSize_, 2), BASE_TOPK, 2] -inf,-1 + InitSortOutBuf(globalTopkUb_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK_VALUE_IDX_SIZE); +} + +template +__aicore__ inline void QLIVector::InitParams(const struct QLICommon::ConstInfo &constInfo, + const QLITilingData *__restrict tilingData) +{ + this->constInfo_ = constInfo; + blockS2StartIdx_ = 0; + gSize_ = constInfo.gSize; + kSeqSize_ = constInfo.kSeqSize; + // define N2 para + kHeadNum_ = constInfo.kHeadNum; + qHeadNum_ = constInfo.qHeadNum; + // define MMBase para + s1BaseSize_ = constInfo.s1BaseSize; // 4 + s2BaseSize_ = constInfo.s2BaseSize; // 2048 + kCacheBlockSize_ = constInfo.kCacheBlockSize; + maxBlockNumPerBatch_ = constInfo.maxBlockNumPerBatch; + blockId_ = GetBlockIdx(); +} + +template +__aicore__ inline void QLIVector::InitVecInputTensor(GlobalTensor weightsGm, GlobalTensor qScaleGm, + GlobalTensor kScaleGm, + GlobalTensor indiceOutGm, + GlobalTensor blockTableGm) +{ + this->weightsGm = weightsGm; + this->qScaleGm = qScaleGm; + this->kScaleGm = kScaleGm; + this->indiceOutGm = indiceOutGm; + this->blockTableGm = blockTableGm; +} + +template +__aicore__ inline void QLIVector::InitVecWorkspaceTensor(GlobalTensor vec0OutGm, + GlobalTensor mm1ResGm) +{ + this->mm1ResGm = mm1ResGm; + this->vec0OutGm = vec0OutGm; +} + +template +__aicore__ inline void QLIVector::AllocEventID() +{ +} + +template +__aicore__ inline void QLIVector::FreeEventID() +{ +} + +template +__aicore__ inline void QLIVector::CleanInvalidOutput(int64_t invalidS1offset) +{ + // init -1 and copy to output + LocalTensor valueULocal = outQueue_.AllocTensor(); + LocalTensor idxULocal1 = valueULocal.template ReinterpretCast(); + Duplicate(idxULocal1, constInfo_.INVALID_IDX, constInfo_.sparseCount); + outQueue_.EnQue(valueULocal); + valueULocal = outQueue_.DeQue(); + QLIServiceVec::CopyOut(indiceOutGm[invalidS1offset], idxULocal1, constInfo_.sparseCount); + outQueue_.FreeTensor(valueULocal); +} + +template +__aicore__ inline void QLIVector::ProcessVec0(const QLICommon::RunInfo &info) +{ + // 只需要一个v核做 + if (blockId_ % 2 != 0) { + return; + } + int32_t cuBaseS1Idx = info.gS1Idx * s1BaseSize_; + // 计算输出w基地址偏移 偶数循环 -> 0 + aic_offset 奇数循环 -> 4*64 + aic_offset + int64_t vec0OutGmOffset = (info.loop % 2) * ((s1BaseSize_ * gSize_ * BLOCK_CUBE)); + // 计算输入weight的地址偏移，qScale的地址偏移与weight相同 + int64_t weightGmOffset = info.tensorWeightsOffset + cuBaseS1Idx * qHeadNum_; + // 当前需要计算的S1行数，处理尾块场景 + int32_t cuS1ProcNum = cuBaseS1Idx + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + int32_t cuProcEleNum = cuS1ProcNum * gSize_; + + LocalTensor inWeightsUb = inQueue_.AllocTensor(); + LocalTensor inQScaleUb = inWeightsUb[cuProcEleNum]; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams copyInParams; + copyInParams.blockCount = 1; + copyInParams.blockLen = cuProcEleNum * sizeof(half); + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + AscendC::DataCopyPad(inWeightsUb, weightsGm[weightGmOffset], copyInParams, padParams); + AscendC::DataCopyPad(inQScaleUb, qScaleGm[weightGmOffset], copyInParams, padParams); + + inQueue_.EnQue(inWeightsUb); + inWeightsUb = inQueue_.DeQue(); + AscendC::Mul(inWeightsUb, inWeightsUb, inQScaleUb, cuProcEleNum); + PipeBarrier(); + LocalTensor resUb = outQueue_.AllocTensor(); + AscendC::Brcb(resUb, inWeightsUb, static_cast(cuProcEleNum / 8), {1, 8}); + inQueue_.FreeTensor(inWeightsUb); + + outQueue_.EnQue(resUb); + resUb = outQueue_.DeQue(); + AscendC::DataCopyParams copyOutParams; + copyOutParams.blockCount = 1; + copyOutParams.blockLen = cuProcEleNum * BLOCK_CUBE * sizeof(half); + copyOutParams.srcStride = 0; + copyOutParams.dstStride = 0; + AscendC::DataCopyPad(vec0OutGm[vec0OutGmOffset], resUb, copyOutParams); + outQueue_.FreeTensor(resUb); +} + +template +__aicore__ inline int32_t QLIVector::AlignS2(int32_t cuS2Len) +{ + // 限制：当前cuS2Len最大为2048，暂不考虑更长 + // 该函数目的是将cuS2Len对齐到形如 32*(4^n)*m 的形式 (m ∈ [1, 3])，方便后续sort/merge + if (cuS2Len <= ELE_NUM_128) { + return Align(cuS2Len, ELE_NUM_32); + } else if (cuS2Len <= ELE_NUM_512) { + return Align(cuS2Len, ELE_NUM_128); + } else { + return Align(cuS2Len, ELE_NUM_512); + } +} + +template +__aicore__ inline void QLIVector::ProcessVec1(const QLICommon::RunInfo &info) +{ + int32_t cuBaseS1Idx = info.gS1Idx * s1BaseSize_; + int32_t cuBaseS2Idx = info.s2Idx * s2BaseSize_; + + // 计算基本块基地址偏移 偶数循环 -> 0 + aic_offset 奇数循环 -> 4*2048 + aic_offset + int64_t mmGmOffset = (info.loop % 2) * (s1BaseSize_ * s2BaseSize_); + + // cuS1BeginIdxPerAiv: 每个AIV的S1起始偏移 + int32_t cuS1BeginIdxPerAiv = cuBaseS1Idx; + int32_t cuS1ProcNum = + cuS1BeginIdxPerAiv + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + // cuS1ProcNumPerAiv: 每个AIv的S1计算量 + int32_t cuS1ProcNumPerAiv = blockId_ % 2 == 0 ? CeilDiv(cuS1ProcNum, 2) : (cuS1ProcNum / 2); + cuS1BeginIdxPerAiv += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2); + // 基本块基地址偏移奇数核加一个S1地址偏移 + mmGmOffset += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2) * s2BaseSize_; + // 非首个基本块, M(S1)轴发生切换需要初始化 + if (info.loop != 0 && info.s2Idx == 0) { + // globalTopkUb_ value,index=-inf,-1 + InitSortOutBuf(globalTopkUb_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK_VALUE_IDX_SIZE); + blockS2StartIdx_ = 0; + } else if (info.loop == 0) { + blockS2StartIdx_ = info.s2Idx; + } + // cuRealAcSeq: 当前基本块S1对应的AcSeq + int32_t cuRealAcSeq = info.actS2Size; + int32_t cuRealAcSeqCount = 0; + if (constInfo_.attenMaskFlag) { + // attenMask true场景 + cuRealAcSeq = info.actS2SizeOrig - info.actS1Size + cuS1BeginIdxPerAiv; + } + int32_t cuRealAcSeqIni = cuRealAcSeq; + + + // LD输出S1方向偏移，保证2个Vector输出的内容连续 + uint32_t ldS1Offset = (blockId_ % 2 == 0) ? s1BaseSize_ / 2 - cuS1ProcNumPerAiv : 0; + for (int innerS1Idx = 0; innerS1Idx < cuS1ProcNumPerAiv; innerS1Idx++) { + if (constInfo_.attenMaskFlag) { + cuRealAcSeqCount += 1; + cuRealAcSeq = (cuRealAcSeqCount + cuRealAcSeqIni) / static_cast(constInfo_.cmpRatio); + } + int32_t cuS2Len = cuBaseS2Idx + s2BaseSize_ >= cuRealAcSeq ? cuRealAcSeq - cuBaseS2Idx : s2BaseSize_; + int32_t cuS1Idx = cuS1BeginIdxPerAiv + innerS1Idx; + if (cuRealAcSeq > 0 && cuS2Len > 0) { + int32_t cuS2LenVecAlign = AlignS2(cuS2Len); + LocalTensor mmInUb = inQueue_.AllocTensor(); + LocalTensor kScaleUb = mmInUb[cuS2LenVecAlign]; + LocalTensor kScaleTUb = kScaleUb.template ReinterpretCast()[cuS2LenVecAlign]; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPadExtParams padTParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams copyInParams; + copyInParams.blockCount = 1; + copyInParams.blockLen = cuS2Len * sizeof(float); + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + AscendC::DataCopyPad(mmInUb, mm1ResGm[mmGmOffset + innerS1Idx * s2BaseSize_], copyInParams, padParams); + GetKeyScale(info, kScaleTUb, info.bIdx, cuBaseS2Idx, cuS2Len); + inQueue_.EnQue(mmInUb); + mmInUb = inQueue_.DeQue(); + AscendC::Cast(kScaleUb, kScaleTUb, RoundMode::CAST_NONE, cuS2Len); + PipeBarrier(); + AscendC::Mul(mmInUb, mmInUb, kScaleUb, cuS2Len); + PipeBarrier(); + LocalTensor sortBuff = tmpBuf_.Get(); + LocalTensor sortScoreUb = sortBuff; + LocalTensor sortIndiceUb = sortBuff[cuS2LenVecAlign]; + PipeBarrier(); + Duplicate(sortScoreUb.template ReinterpretCast(), QLIServiceVec::NEG_INF, cuS2LenVecAlign); + PipeBarrier(); + Adds(sortScoreUb, mmInUb, 0.0f, cuS2Len); + PipeBarrier(); + inQueue_.FreeTensor(mmInUb); + LocalTensor sortIndiceUbInt = sortIndiceUb.template ReinterpretCast(); + // 无效数据索引填充为-1 + if (cuS2LenVecAlign != cuS2Len) { + Duplicate(sortIndiceUbInt, -1, cuS2LenVecAlign); + PipeBarrier(); + } + Adds(sortIndiceUbInt, globalTopkIndice_, static_cast(cuBaseS2Idx), cuS2Len); + PipeBarrier(); + LocalTensor tmpSortBuf = sortBuff[2 * cuS2LenVecAlign]; + QLIServiceVec::SortAll(sortBuff, tmpSortBuf, cuS2LenVecAlign); + PipeBarrier(); + QLIServiceVec::MergeSort(globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], BASE_TOPK, sortBuff, + cuS2LenVecAlign, tmpSortBuf); + PipeBarrier(); + bool isS2End = cuBaseS2Idx + s2BaseSize_ >= cuRealAcSeq; + bool needCopyOutGm = blockS2StartIdx_ == 0 && isS2End; + if (needCopyOutGm) { + LocalTensor idxULocal = outQueue_.AllocTensor(); + ExtractIndex(idxULocal, + globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE].template ReinterpretCast(), + BASE_TOPK); + PipeBarrier(); + InitSortOutBuf(globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], BASE_TOPK_VALUE_IDX_SIZE); + outQueue_.EnQue(idxULocal); + idxULocal = outQueue_.DeQue(); + QLIServiceVec::CopyOut(indiceOutGm[info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount], + idxULocal.template ReinterpretCast(), constInfo_.sparseCount); + outQueue_.FreeTensor(idxULocal); + } + } else if (cuRealAcSeq <= 0) { + CleanInvalidOutput(info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount); + } + } + + // BNSD场景无效S1 输出-1 + if (Q_LAYOUT_T == LI_LAYOUT::BSND) { + // 最后一个S1的基本块, 需要 >= info.actS1Size + bool isS1LoopEnd = (cuBaseS1Idx + s1BaseSize_) >= info.actS1Size; + int32_t invalidS1Num = constInfo_.qSeqSize - info.actS1Size; + // blockS2StartIdx_ == 0 控制S2从开始的核去做冗余清理 + if (invalidS1Num > 0 && isS1LoopEnd && blockS2StartIdx_ == 0) { + int32_t s1NumPerAiv = blockId_ % 2 == 0 ? CeilDiv(invalidS1Num, 2) : (invalidS1Num / 2); + int32_t s1OffsetPerAiv = info.actS1Size + (blockId_ % 2) * CeilDiv(invalidS1Num, 2); + for (int innerS1Idx = 0; innerS1Idx < s1NumPerAiv; innerS1Idx++) { + CleanInvalidOutput(info.indiceOutOffset + (s1OffsetPerAiv + innerS1Idx) * constInfo_.sparseCount); + } + } + + int32_t invalidS1Num2 = info.actS1Size - info.actS2SizeOrig; + if (invalidS1Num2 > 0 && isS1LoopEnd && blockS2StartIdx_ == 0 && constInfo_.attenMaskFlag) { + int32_t s1NumPerAiv = blockId_ % 2 == 0 ? CeilDiv(invalidS1Num2, 2) : (invalidS1Num2 / 2); + int32_t s1OffsetPerAiv = (blockId_ % 2) * CeilDiv(invalidS1Num2, 2); + for (int innerS1Idx = 0; innerS1Idx < s1NumPerAiv; innerS1Idx++) { + CleanInvalidOutput((info.bN2Idx * constInfo_.qSeqSize + s1OffsetPerAiv + innerS1Idx) * + constInfo_.sparseCount); + } + } + } + + if (info.isLastS2InnerLoop) { + // S2最后一个Loop后, 下一个基本块初始从0开始 + blockS2StartIdx_ = 0; + } +} + +} // namespace QLIKernel +#endif // QUANT_LIGHTNING_INDEXER_SERVICE_VECTOR_H \ No newline at end of file diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_vector.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_vector.h new file mode 100644 index 000000000..ed9348d4b --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_vector.h @@ -0,0 +1,193 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_vector.h + * \brief + */ +#ifndef QUANT_LIGHTNING_INDEXER_VECTOR_H +#define QUANT_LIGHTNING_INDEXER_VECTOR_H + +#include "kernel_operator.h" +#include "quant_lightning_indexer_vector.h" + +namespace QLIServiceVec { +using namespace AscendC; + +constexpr int32_t NEG_INF = 0xFF800000; +constexpr int32_t INVALID_INDEX = -1; +constexpr uint8_t VEC_REPEAT_MAX = 255; +constexpr uint8_t B32_VEC_ELM_NUM = 64; +constexpr uint8_t B32_BLOCK_ALIGN_NUM = 8; +constexpr uint8_t B32_VEC_REPEAT_STRIDE = 8; +constexpr uint64_t VEC_REPEAT_BYTES = 256; +constexpr int32_t CONST_TWO = 2; +constexpr int64_t VALUE_AND_INDEX_NUM = 2; +constexpr int64_t BLOCK_BYTES = 32; +constexpr int64_t MRG_QUE_0 = 0; +constexpr int64_t MRG_QUE_1 = 1; +constexpr int64_t MRG_QUE_2 = 2; +constexpr int64_t MRG_QUE_3 = 3; +constexpr int64_t MRG_BLOCK_2 = 2; +constexpr int64_t MRG_BLOCK_3 = 3; +constexpr int64_t MRG_BLOCK_4 = 4; + +template +__aicore__ inline void CopyOut(const GlobalTensor &dstGm, const LocalTensor &srcUb, int64_t copyCount) +{ + AscendC::DataCopyParams dataCopyOutyParams; + dataCopyOutyParams.blockCount = 1; + dataCopyOutyParams.blockLen = copyCount * sizeof(T); + dataCopyOutyParams.srcStride = 0; + dataCopyOutyParams.dstStride = 0; + AscendC::DataCopyPad(dstGm, srcUb, dataCopyOutyParams); +} + +/** + src: 传入的初始化空间 + eleNum: 需要初始化的元素个数需为64整数倍，元素将被初始化为交错排布的-inf，-1 + */ +__aicore__ inline void InitSortOutBuf(const LocalTensor &src, int64_t eleNum) +{ + uint64_t mask1[2] = {0x5555555555555555, 0}; + uint64_t mask0[2] = {0xaaaaaaaaaaaaaaaa, 0}; + int64_t repeatNum = eleNum / B32_VEC_ELM_NUM; + int64_t forLoop = repeatNum / VEC_REPEAT_MAX; + int64_t forRemain = repeatNum % VEC_REPEAT_MAX; + for (int i = 0; i < forLoop; i++) { + AscendC::Duplicate(src.template ReinterpretCast(), NEG_INF, mask1, VEC_REPEAT_MAX, 1, + B32_VEC_REPEAT_STRIDE); + AscendC::Duplicate(src.template ReinterpretCast(), INVALID_INDEX, mask0, VEC_REPEAT_MAX, 1, + B32_VEC_REPEAT_STRIDE); + } + if (forRemain > 0) { + AscendC::Duplicate(src.template ReinterpretCast()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM], NEG_INF, + mask1, forRemain, 1, B32_VEC_REPEAT_STRIDE); + AscendC::Duplicate(src.template ReinterpretCast()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM], + INVALID_INDEX, mask0, forRemain, 1, B32_VEC_REPEAT_STRIDE); + } + AscendC::PipeBarrier(); +} + +/** + src: logits和索引，前logitsNum为logits，后logitsNum为索引 + tmp: 计算使用到的临时空间，大小与src一致 + logitsNum: 排序的元素个数, 暂只支持[128,256,384,512,1024,2048] + */ +__aicore__ inline void SortAll(LocalTensor &src, LocalTensor &tmp, int64_t logitsNum) +{ + int64_t sort32Repeats = logitsNum / BLOCK_BYTES; + AscendC::Sort32(tmp, src, src[logitsNum].ReinterpretCast(), sort32Repeats); + AscendC::PipeBarrier(); + + int64_t mrgGroups = sort32Repeats; + int64_t mrgElements = BLOCK_BYTES; + int64_t i = 0; + AscendC::LocalTensor srcTensor; + AscendC::LocalTensor dstTensor; + while (true) { + if (i % CONST_TWO == 0) { + srcTensor = tmp; + dstTensor = src; + } else { + srcTensor = src; + dstTensor = tmp; + } + AscendC::MrgSort4Info params; + params.elementLengths[0] = mrgElements; + params.elementLengths[MRG_QUE_1] = mrgElements; + params.elementLengths[MRG_QUE_2] = mrgElements; + params.elementLengths[MRG_QUE_3] = mrgElements; + params.ifExhaustedSuspension = false; + params.validBit = 0b1111; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = srcTensor[0]; + srcList.src2 = srcTensor[MRG_QUE_1 * VALUE_AND_INDEX_NUM * mrgElements]; + srcList.src3 = srcTensor[MRG_QUE_2 * VALUE_AND_INDEX_NUM * mrgElements]; + srcList.src4 = srcTensor[MRG_QUE_3 * VALUE_AND_INDEX_NUM * mrgElements]; + if (mrgGroups <= MRG_BLOCK_4) { + params.repeatTimes = 1; + if (mrgGroups == 1) { + break; + } else if (mrgGroups == MRG_BLOCK_2) { + params.validBit = 0b0011; + } else if (mrgGroups == MRG_BLOCK_3) { + params.validBit = 0b0111; + } else if (mrgGroups == MRG_BLOCK_4) { + params.validBit = 0b1111; + } + AscendC::MrgSort(dstTensor, srcList, params); + i += 1; + break; + } else { + params.repeatTimes = mrgGroups / MRG_BLOCK_4; + AscendC::MrgSort(dstTensor, srcList, params); + i += 1; + mrgElements = mrgElements * MRG_BLOCK_4; + mrgGroups = mrgGroups / MRG_BLOCK_4; + } + AscendC::PipeBarrier(); + } + if (i % CONST_TWO == 0) { + AscendC::DataCopy(src, tmp, logitsNum * VALUE_AND_INDEX_NUM); + AscendC::PipeBarrier(); + } +} + +/** + mrgDst: 合并进的Tensor + mrgSrc: 待合并的Tensor + tmpTensor：空间为mrgDst+mrgSrc + */ +__aicore__ inline void MergeSort(const LocalTensor &mrgDst, int32_t mrgDstNum, LocalTensor &mrgSrc, + int32_t mrgSrcNum, LocalTensor &tmpTensor) +{ + AscendC::MrgSort4Info params; + params.elementLengths[0] = mrgSrcNum; + params.elementLengths[1] = mrgDstNum; + params.ifExhaustedSuspension = false; + params.validBit = 0b0011; + params.repeatTimes = 1; + + AscendC::MrgSortSrcList srcList; + srcList.src1 = mrgSrc; + srcList.src2 = mrgDst; + + AscendC::MrgSort(tmpTensor, srcList, params); + AscendC::PipeBarrier(); + AscendC::DataCopy(mrgDst, tmpTensor, mrgDstNum * VALUE_AND_INDEX_NUM); + AscendC::PipeBarrier(); +} + +__aicore__ inline void ExtractIndex(const LocalTensor &idxULocal, const LocalTensor &sortLocal, + int64_t extractNum) +{ + AscendC::GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = Ceil(extractNum * sizeof(float) * VALUE_AND_INDEX_NUM, VEC_REPEAT_BYTES); + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = B32_VEC_REPEAT_STRIDE; + gatherMaskParams.src1RepeatStride = 0; + uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数 + uint8_t src1Pattern = 2; // 固定模式2,表示筛选出奇数索引的数 + AscendC::GatherMask(idxULocal, sortLocal, src1Pattern, false, static_cast(0), gatherMaskParams, rsvdCnt); + AscendC::PipeBarrier(); +} + +template +__aicore__ inline void SetWaitFlag(HardEvent evt) +{ + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(evt)); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); +} + +} // namespace QLIServiceVec +#endif // QUANT_LIGHTNING_INDEXER_VECTOR_H \ No newline at end of file diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/quant_lightning_indexer_common.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/quant_lightning_indexer_common.h new file mode 100644 index 000000000..1e4b19d48 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/quant_lightning_indexer_common.h @@ -0,0 +1,179 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_common.h + * \brief + */ +#ifndef quant_lightning_indexer_COMMON_H +#define quant_lightning_indexer_COMMON_H +using namespace AscendC; +namespace QLICommon { + +// 与tiling的layout保持一致 +enum class LI_LAYOUT : uint32_t { + BSND = 0, + TND = 1, + PA_BSND = 2 +}; + +template +struct QLIType { + static_assert( + (std::is_same_v && + (std::is_same_v || std::is_same_v)) || + (std::is_same_v && + std::is_same_v), + "Invalid combination of QK_T and SCORE_T" + ); + using queryType = Q_T; + using keyType = K_T; + using queryKeyType = QK_T; + using scoreType = SCORE_T; + using outputType = OUT_T; + + static constexpr bool pageAttention = PAGE_ATTENTION; + static constexpr LI_LAYOUT layout = Q_LAYOUT_T; + static constexpr LI_LAYOUT keyLayout = K_LAYOUT_T; +}; + +struct RunInfo { + uint32_t loop; + uint32_t bN2Idx; + uint32_t bIdx; + uint32_t n2Idx = 0; + uint32_t gS1Idx; + uint32_t s2Idx; + + uint32_t actS1Size = 1; + uint32_t actS2Size = 1; + uint32_t actS2SizeOrig = 1; + uint32_t actMBaseSize; + uint32_t actualSingleProcessSInnerSize; + uint32_t actualSingleProcessSInnerSizeAlign; + + uint64_t tensorQueryOffset; + uint64_t tensorKeyOffset; + uint64_t tensorKeyScaleOffset; + uint64_t tensorWeightsOffset; + uint64_t indiceOutOffset; + + bool isFirstS2InnerLoop; + bool isLastS2InnerLoop; + bool isAllLoopEnd = false; + bool isValid = false; +}; + +struct ConstInfo { + // CUBE与VEC核间同步的模式 + static constexpr uint32_t QLI_SYNC_MODE4 = 4; + static constexpr uint32_t AIV0_AIV1_OFFSET = 16; + static constexpr uint32_t CROSS_VC_EVENT = 0; + static constexpr uint32_t CROSS_CV_EVENT = 2; + // BUFFER的字节数 + static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32; + static constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64; + static constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256; + static constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512; + static constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024; + static constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048; + static constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096; + static constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192; + static constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384; + static constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768; + // 无效索引 + static constexpr int INVALID_IDX = -1; + + // CUBE和VEC的核间同步EventID + uint32_t syncC1V1 = 0U; + uint32_t syncC1V0 = 2U; + uint32_t syncV1C1 = 0U; + uint32_t syncV0C1 = 1U; + + // 基本块大小 + uint32_t mBaseSize = 1ULL; + uint32_t s1BaseSize = 1ULL; + uint32_t s2BaseSize = 1ULL; + + uint64_t batchSize = 0ULL; + uint64_t gSize = 0ULL; + uint64_t qHeadNum = 0ULL; + uint64_t kHeadNum; + uint64_t headDim; + uint64_t sparseCount; // topK选取大小 + uint64_t kSeqSize = 0ULL; // kv最大S长度 + uint64_t qSeqSize = 1ULL; // q最大S长度 + uint32_t kCacheBlockSize = 0; // PA场景的block size + uint32_t maxBlockNumPerBatch = 0; // PA场景的最大单batch block number + LI_LAYOUT outputLayout; // 输出的格式 + bool attenMaskFlag = false; + uint32_t cmpRatio = 1; + bool batchSupperFlag = false; // Qactual_se长度是否为B+1 + int64_t stride = 1; + int64_t scaleStride = 1; + + uint32_t actualLenQDims = 0U; // query的actualSeqLength 的维度 + uint32_t actualLenDims = 0U; // KV 的actualSeqLength 的维度 + bool isAccumSeqS1 = false; // 是否累加模式 + bool isAccumSeqS2 = false; // 是否累加模式 + bool isLDOpen = false; +}; + +struct SplitCoreInfo { + uint32_t s2Start = 0U; // S2的起始位置 + uint32_t s2End = 0U; // S2循环index上限 + uint32_t bN2Start = 0U; + uint32_t bN2End = 0U; + uint32_t gS1Start = 0U; + uint32_t gS1End = 0U; + bool isLD = false; // 当前核是否需要进行Decode归约任务 + bool isCoreEnable = false; +}; + +template +__aicore__ inline T Align(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd) * (rnd))); +} + +template +__aicore__ inline T1 Min(T1 a, T2 b) +{ + return (a > b) ? (b) : (a); +} + +template +__aicore__ inline T1 Max(T1 a, T2 b) +{ + return (a > b) ? (a) : (b); +} + +template +__aicore__ inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd))); +} +} // namespace QLICommon + +// bank冲突优化 +// david 256KB bank layout +// shape ( bank_depth ( banks bank_groups block)) (512 ( 2 8 32)) +// stride (banks*bank_groups*block (bank_groups*block block 1)) (512 (256 32 1)) +#define UB_BLOCK 32 // 32B +#define UB_BANK_GROUPS 8 +#define UB_BANKS 2 +#define UB_BANK_DEPTH 512 + +#define UB_BANK_GROUP_STRIDE UB_BLOCK // 32B +#define UB_BANK_STRIDE (UB_BANK_GROUPS * UB_BLOCK) // 256B +#define UB_BANK_DEPTH_STRIDE (UB_BANKS * UB_BANK_GROUPS * UB_BLOCK) // 512B + +#endif // quant_lightning_indexer_COMMON_H diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/quant_lightning_indexer_kernel.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/quant_lightning_indexer_kernel.h new file mode 100644 index 000000000..ec2b7f8b6 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/quant_lightning_indexer_kernel.h @@ -0,0 +1,640 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_kernel.h + * \brief + */ + +#ifndef quant_lightning_indexer_KERNEL_H +#define quant_lightning_indexer_KERNEL_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "quant_lightning_indexer_common.h" +#include "quant_lightning_indexer_service_vector.h" +#include "quant_lightning_indexer_service_cube.h" +#include "../vllm_quant_lightning_indexer_metadata.h" + +namespace QLIKernel { +using namespace QLICommon; +using namespace matmul; +using namespace optiling; +using namespace optiling::detail; +using AscendC::CacheMode; +using AscendC::CrossCoreSetFlag; +using AscendC::CrossCoreWaitFlag; + +// 由于S2循环前，RunInfo还没有赋值，使用TempLoopInfo临时存放B、N、S1轴相关的信息；同时减少重复计算 +struct TempLoopInfo { + uint32_t bN2Idx = 0; + uint32_t bIdx = 0U; + uint32_t n2Idx = 0U; + uint32_t gS1Idx = 0U; + uint32_t gS1LoopEnd = 0U; // gS1方向循环的结束Idx + uint32_t s2LoopEnd = 0U; // S2方向循环的结束Idx + uint32_t actS1Size = 1ULL; // 当前Batch循环处理的S1轴的实际大小 + uint32_t actS2Size = 0ULL; + uint32_t actS2SizeOrig = 0ULL;//压缩前s2 + bool curActSeqLenIsZero = false; + bool needDealActS1LessThanS1 = false; // S1的实际长度小于shape的S1长度时，是否需要清理输出 + uint32_t actMBaseSize = 0U; // m轴(gS1)方向实际大小 + uint32_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小 + uint32_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小 +}; + +template +class QLIPreload { +public: + __aicore__ inline QLIPreload(){}; + __aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale, __gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengthsK, __gm__ uint8_t *blockTable, + __gm__ uint8_t *metadata, __gm__ uint8_t *sparseIndices, + __gm__ uint8_t *workspace, const QLITilingData *__restrict tiling, TPipe *tPipe); + __aicore__ inline void Process(); + + // =================================类型定义区================================= + using Q_T = typename QLIT::queryType; + using K_T = typename QLIT::keyType; + using OUT_T = typename QLIT::outputType; + static constexpr bool PAGE_ATTENTION = QLIT::pageAttention; + static constexpr LI_LAYOUT Q_LAYOUT_T = QLIT::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = QLIT::keyLayout; + + using SCORE_T = typename QLIT::scoreType; + + QLIMatmul matmulService; + QLIVector vectorService; + + // =================================常量区================================= + static constexpr uint32_t SYNC_C1_V1_FLAG = 4; + static constexpr uint32_t SYNC_V1_C1_FLAG = 5; + + static constexpr uint32_t M_BASE_SIZE = 256; + static constexpr uint32_t S2_BASE_SIZE = 128; + static constexpr uint32_t HEAD_DIM = 128; + static constexpr uint32_t K_HEAD_NUM = 1; + static constexpr uint32_t GM_ALIGN_BYTES = 512; + + static constexpr int64_t LD_PREFETCH_LEN = 2; + // for workspace double + static constexpr uint32_t WS_DOBULE = 2; + +protected: + TPipe *pipe = nullptr; + + // offset + uint64_t queryCoreOffset = 0ULL; + uint64_t keyCoreOffset = 0ULL; + uint64_t keyScaleCoreOffset = 0ULL; + uint64_t weightsCoreOffset = 0ULL; + uint64_t indiceOutCoreOffset = 0ULL; + bool isUsedCoreEqZero = false; + // ================================Global Buffer区================================= + GlobalTensor queryGm; + GlobalTensor keyGm; + GlobalTensor weightsGm; + GlobalTensor qScaleGm; + GlobalTensor kScaleGm; + GlobalTensor metadataGm; + + GlobalTensor indiceOutGm; + GlobalTensor blockTableGm; + + GlobalTensor actualSeqLengthsGmQ; + GlobalTensor actualSeqLengthsGm; + + // ================================类成员变量==================================== + // aic、aiv核信息 + uint32_t tmpBlockIdx = 0U; + uint32_t aiCoreIdx = 0U; + uint32_t usedCoreNum = 0U; + + QLICommon::ConstInfo constInfo{}; + TempLoopInfo tempLoopInfo{}; + QLICommon::SplitCoreInfo splitCoreInfo{}; + + // ================================Init functions================================== + __aicore__ inline void InitTilingData(const QLITilingData *__restrict tilingData); + __aicore__ inline void InitBuffers(); + __aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK); + // ================================Split Core================================ + __aicore__ inline void SplitCoreByAICPU(uint32_t curCoreIdx, GlobalTensor &metadataGm); + __aicore__ inline uint32_t GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, uint32_t actS2SizeOrig); + // ================================Process functions================================ + __aicore__ inline void ProcessMain(); + __aicore__ inline void ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, + QLICommon::RunInfo runInfo); + __aicore__ inline void ProcessInvalid(); + // ================================Params Calc===================================== + __aicore__ inline void CalcGS1LoopParams(uint32_t bN2Idx); + __aicore__ inline void GetBN2Idx(uint32_t bN2Idx); + __aicore__ inline uint32_t GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, uint32_t defaultSeqLen); + __aicore__ inline uint32_t GetActualSeqLenKey(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, uint32_t defaultSeqLen, uint32_t cmpRatio); + __aicore__ inline void GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, uint32_t &actS2Size, uint32_t &actS2SizeOrig); + __aicore__ inline void CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx); + __aicore__ inline void CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, QLICommon::RunInfo &runInfo); + __aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start); +}; + +template +__aicore__ inline void QLIPreload::InitTilingData(const QLITilingData *__restrict tilingData) +{ + usedCoreNum = tilingData->usedCoreNum; + constInfo.batchSize = tilingData->bSize; + constInfo.qHeadNum = constInfo.gSize = tilingData->gSize; + constInfo.kSeqSize = tilingData->s2Size; + constInfo.qSeqSize = tilingData->s1Size; + constInfo.attenMaskFlag = (tilingData->sparseMode == 3); + constInfo.kCacheBlockSize = tilingData->blockSize; + constInfo.maxBlockNumPerBatch = tilingData->maxBlockNumPerBatch; + constInfo.sparseCount = tilingData->sparseCount; + constInfo.cmpRatio = tilingData->cmpRatio; + constInfo.batchSupperFlag = tilingData->batchSupperFlag; + constInfo.stride = tilingData->stride; + constInfo.scaleStride = tilingData->scaleStride; + constInfo.outputLayout = Q_LAYOUT_T; // 输出和输入形状一致 + if (Q_LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS1 = true; + } + if (K_LAYOUT_T == LI_LAYOUT::TND) { + constInfo.isAccumSeqS2 = true; + } + + constInfo.kHeadNum = K_HEAD_NUM; + constInfo.headDim = HEAD_DIM; + + constInfo.mBaseSize = M_BASE_SIZE; + constInfo.s2BaseSize = S2_BASE_SIZE; + constInfo.s1BaseSize = (constInfo.mBaseSize + constInfo.gSize - 1) / constInfo.gSize; +} + +template +__aicore__ inline void QLIPreload::InitBuffers() +{ + if ASCEND_IS_AIV { + vectorService.InitBuffers(pipe); + } else { + matmulService.InitBuffers(pipe); + } +} + +template +__aicore__ inline void QLIPreload::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, + __gm__ uint8_t *actualSeqLengthsK) +{ + if (actualSeqLengthsQ == nullptr) { + constInfo.actualLenQDims = 0; + } else { + constInfo.actualLenQDims = (constInfo.batchSupperFlag) ? constInfo.batchSize + 1 : constInfo.batchSize; + actualSeqLengthsGmQ.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsQ, constInfo.actualLenQDims); + } + if (actualSeqLengthsK == nullptr) { + constInfo.actualLenDims = 0; + } else { + constInfo.actualLenDims = constInfo.batchSize; + actualSeqLengthsGm.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsK, constInfo.actualLenDims); + } +} + +template +__aicore__ inline uint32_t QLIPreload::GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, + uint32_t defaultSeqLen) +{ + bIdx = (constInfo.batchSupperFlag)? bIdx + 1 : bIdx; // 如果为B+1情况，则向后移动一位 + if (actualLenDims == 0) { + return defaultSeqLen; + } else if (isAccumSeq && bIdx > 0) { + return actualSeqLengthsGm.GetValue(bIdx) - actualSeqLengthsGm.GetValue(bIdx - 1); + } else { + return actualSeqLengthsGm.GetValue(bIdx); + } +} + +template +__aicore__ inline uint32_t QLIPreload::GetActualSeqLenKey(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq, + GlobalTensor &actualSeqLengthsGm, + uint32_t defaultSeqLen, uint32_t cmpRatio) +{ + if (actualLenDims == 0) { + return defaultSeqLen * cmpRatio; + } else if (isAccumSeq && bIdx > 0) { + return actualSeqLengthsGm.GetValue(bIdx) - actualSeqLengthsGm.GetValue(bIdx - 1); + } else { + return actualSeqLengthsGm.GetValue(bIdx); + } +} + +template +__aicore__ inline void QLIPreload::GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, uint32_t &actS2Size, uint32_t &actS2SizeOrig) +{ + actS1Size = GetActualSeqLen(bIdx, constInfo.actualLenQDims, constInfo.isAccumSeqS1, actualSeqLengthsGmQ, + constInfo.qSeqSize); + actS2SizeOrig = + GetActualSeqLenKey(bIdx, constInfo.actualLenDims, constInfo.isAccumSeqS2, actualSeqLengthsGm, constInfo.kSeqSize, constInfo.cmpRatio); // 压缩前的actS2Size + actS2Size = actS2SizeOrig / constInfo.cmpRatio; // 真实使用的压缩后S2长度 +} + +template +__aicore__ inline uint32_t QLIPreload::GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, + uint32_t actS2SizeOrig) +{ + if (actS2SizeOrig / constInfo.cmpRatio == 0) { + return 0; + } + uint32_t s1Offset = constInfo.s1BaseSize * s1gIdx; + int32_t validS2LenBase = static_cast(actS2SizeOrig) - static_cast(actS1Size); // 压缩前的validS2LenBase + int32_t validS2Len = (static_cast(s1Offset) + validS2LenBase + static_cast(constInfo.s1BaseSize)) / static_cast(constInfo.cmpRatio); + validS2Len = Min(validS2Len, static_cast(actS2SizeOrig) / constInfo.cmpRatio); + validS2Len = Max(validS2Len, 1); + return (validS2Len + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; +} + +template +__aicore__ inline void QLIPreload::SplitCoreByAICPU(uint32_t curCoreIdx, GlobalTensor &metadataGm) +{ + uint32_t liCoreEnableIndex = GetAttrAbsIndex(curCoreIdx, LI_CORE_ENABLE_INDEX); + uint32_t bN2StartIndex = GetAttrAbsIndex(curCoreIdx, LI_BN2_START_INDEX); + uint32_t mStartIndex = GetAttrAbsIndex(curCoreIdx, LI_M_START_INDEX); + uint32_t s2StartIndex = GetAttrAbsIndex(curCoreIdx, LI_S2_START_INDEX); + uint32_t bN2EndIndex = GetAttrAbsIndex(curCoreIdx, LI_BN2_END_INDEX); + uint32_t mEndIndex = GetAttrAbsIndex(curCoreIdx, LI_M_END_INDEX); + uint32_t s2EndIndex = GetAttrAbsIndex(curCoreIdx, LI_S2_END_INDEX); + + uint32_t liZeroCoreEnableIndex = GetAttrAbsIndex(0, LI_CORE_ENABLE_INDEX); + if (metadataGm.GetValue(liZeroCoreEnableIndex) == 0) { + isUsedCoreEqZero = true; + } + if (metadataGm.GetValue(liCoreEnableIndex) == 0) { + splitCoreInfo.isCoreEnable = false; + return; + } else { + splitCoreInfo.isCoreEnable = true; + } + + splitCoreInfo.bN2Start = metadataGm.GetValue(bN2StartIndex); + splitCoreInfo.gS1Start = metadataGm.GetValue(mStartIndex); + splitCoreInfo.s2Start = metadataGm.GetValue(s2StartIndex); + splitCoreInfo.bN2End = metadataGm.GetValue(bN2EndIndex); + splitCoreInfo.gS1End = metadataGm.GetValue(mEndIndex); + splitCoreInfo.s2End = metadataGm.GetValue(s2EndIndex); + + if (splitCoreInfo.s2End != 0) { + // 此时只需要s2End往前退一格，bN2End和gS1End都不变 + splitCoreInfo.s2End = splitCoreInfo.s2End - 1; + } else { + if (splitCoreInfo.gS1End != 0) { + // splitCoreInfo.gS1End != 0 splitCoreInfo.s2End == 0 时，gS1End需要往前退一格, bN2End不变 + // 此时需要使用bIdx获取实际Actal S2来计算出 s2End + splitCoreInfo.gS1End = splitCoreInfo.gS1End - 1; + // 需要获取当前的Actaul S2 + uint32_t bIdx = splitCoreInfo.bN2End / constInfo.kHeadNum; + uint32_t actS1Size, actS2Size, actS2SizeOrig; + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size, actS2SizeOrig); + // s2的切块数量 + uint32_t s2BaseNum; + if (constInfo.attenMaskFlag) { + s2BaseNum = GetS2BaseBlockNumOnMask(splitCoreInfo.gS1End, actS1Size, actS2SizeOrig); + } else { + s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize); + } + splitCoreInfo.s2End = s2BaseNum - 1; + } else { + // splitCoreInfo.gS1End == 0 splitCoreInfo.s2End == 0 时，bN2End需要往前退一格 + // 此时需要使用bIdx获取实际Actal S1和S2来计算出 gS1End 和 s2End + splitCoreInfo.bN2End = splitCoreInfo.bN2End - 1; + + // 需要获取当前的Actaul S1 S2 + uint32_t bIdx = splitCoreInfo.bN2End / constInfo.kHeadNum; + uint32_t actS1Size, actS2Size, actS2SizeOrig; + GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size, actS2SizeOrig); + + // s1的切块数量 + uint32_t s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize); + splitCoreInfo.gS1End = s1GBaseNum - 1; + + // s2的切块数量 + uint32_t s2BaseNum; + if (constInfo.attenMaskFlag) { + s2BaseNum = GetS2BaseBlockNumOnMask(splitCoreInfo.gS1End, actS1Size, actS2SizeOrig); + } else { + s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize); + } + splitCoreInfo.s2End = s2BaseNum - 1; + } + } + + splitCoreInfo.isLD = false; +} + +template +__aicore__ inline void QLIPreload::DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start) +{ + if ASCEND_IS_AIV { + if (constInfo.outputLayout == LI_LAYOUT::TND) { + uint32_t tSizeIdx = (constInfo.batchSupperFlag) ? constInfo.batchSize : constInfo.batchSize - 1; + uint32_t tBaseIdx = (constInfo.batchSupperFlag) ? bIdx : bIdx - 1; + uint32_t tSize = actualSeqLengthsGmQ.GetValue(tSizeIdx); + uint32_t tBase = bIdx == 0 ? 0 : actualSeqLengthsGmQ.GetValue(tBaseIdx); + uint32_t s1Count = tempLoopInfo.actS1Size; + + for (uint32_t s1Idx = s1Start; s1Idx < s1Count; s1Idx++) { + uint64_t indiceOutOffset = + (tBase + s1Idx) * constInfo.kHeadNum * constInfo.sparseCount + // T轴、s1轴偏移 + n2Idx * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } else if (constInfo.outputLayout == LI_LAYOUT::BSND) { + for (uint32_t s1Idx = s1Start; s1Idx < constInfo.qSeqSize; s1Idx++) { + // B,S1,N2,K + uint64_t indiceOutOffset = bIdx * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount + + s1Idx * constInfo.kHeadNum * constInfo.sparseCount + // B轴、S1轴偏移 + n2Idx * constInfo.sparseCount; // N2轴偏移 + vectorService.CleanInvalidOutput(indiceOutOffset); + } + } + } +} + +template +__aicore__ inline void QLIPreload::Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK, + __gm__ uint8_t *blockTable, __gm__ uint8_t *metadata, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *workspace, + const QLITilingData *__restrict tiling, TPipe *tPipe) +{ + if ASCEND_IS_AIV { + tmpBlockIdx = GetBlockIdx(); // vec:0-47 + aiCoreIdx = tmpBlockIdx / 2; + } else { + tmpBlockIdx = GetBlockIdx(); // cube:0-23 + aiCoreIdx = tmpBlockIdx; + } + + InitTilingData(tiling); + InitActualSeqLen(actualSeqLengthsQ, actualSeqLengthsK); + + // 获取分核信息 + metadataGm.SetGlobalBuffer((__gm__ uint32_t *)metadata); + SplitCoreByAICPU(aiCoreIdx, metadataGm); + + pipe = tPipe; + + uint64_t offset = 0; + //vec 把整个s2的score存储在GM，大小为s1BaseSize * 16K * 4 + GlobalTensor scoreGm; //存放vec核写出的score + if ASCEND_IS_AIV { + uint64_t singleCoreScoreSize = constInfo.s1BaseSize * QLICommon::Align((uint64_t)constInfo.kSeqSize, (uint64_t)constInfo.s2BaseSize) * sizeof(SCORE_T); + scoreGm.SetGlobalBuffer((__gm__ SCORE_T *)(workspace + aiCoreIdx * singleCoreScoreSize)); + offset += GetBlockNum() * singleCoreScoreSize; + } + + if ASCEND_IS_AIV { + vectorService.InitParams(constInfo, tiling); + indiceOutGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices); + weightsGm.SetGlobalBuffer((__gm__ float *)weights); + qScaleGm.SetGlobalBuffer((__gm__ float *)queryScale); + kScaleGm.SetGlobalBuffer((__gm__ float *)keyScale); + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + vectorService.InitVecInputTensor(weightsGm, qScaleGm, kScaleGm, indiceOutGm, blockTableGm); + vectorService.InitVecWorkspaceTensor(scoreGm); + } else { + matmulService.InitParams(constInfo); + queryGm.SetGlobalBuffer((__gm__ Q_T *)query); + if constexpr (PAGE_ATTENTION) { + blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable); + } + keyGm.SetGlobalBuffer((__gm__ K_T *)key); + matmulService.InitMm1GlobalTensor(blockTableGm, keyGm, queryGm); + } + InitBuffers(); +} + +template +__aicore__ inline void QLIPreload::GetBN2Idx(uint32_t bN2Idx) +{ + tempLoopInfo.bN2Idx = bN2Idx; + tempLoopInfo.bIdx = bN2Idx / constInfo.kHeadNum; + tempLoopInfo.n2Idx = bN2Idx % constInfo.kHeadNum; +} + +template +__aicore__ inline void QLIPreload::CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx) +{ + tempLoopInfo.gS1Idx = gS1LoopIdx; + tempLoopInfo.actMBaseSize = constInfo.mBaseSize; + uint32_t remainedGS1Size = tempLoopInfo.actS1Size * constInfo.gSize - tempLoopInfo.gS1Idx * constInfo.mBaseSize; + if (remainedGS1Size <= constInfo.mBaseSize && remainedGS1Size > 0) { + tempLoopInfo.actMBaseSize = tempLoopInfo.mBasicSizeTail; + } + + bool isEnd = (bN2LoopIdx == splitCoreInfo.bN2End) && (gS1LoopIdx == splitCoreInfo.gS1End); + uint32_t s2BlockNum; + if (constInfo.attenMaskFlag) { + s2BlockNum = GetS2BaseBlockNumOnMask(gS1LoopIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2SizeOrig); + } else { + s2BlockNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + } + tempLoopInfo.s2LoopEnd = isEnd ? splitCoreInfo.s2End : s2BlockNum - 1; +} + +template +__aicore__ inline void QLIPreload::CalcGS1LoopParams(uint32_t bN2LoopIdx) +{ + GetBN2Idx(bN2LoopIdx); + GetS1S2ActualSeqLen(tempLoopInfo.bIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2Size, tempLoopInfo.actS2SizeOrig); + if ((tempLoopInfo.actS2Size == 0) || (tempLoopInfo.actS1Size == 0)) { + tempLoopInfo.curActSeqLenIsZero = true; + return; + } + tempLoopInfo.curActSeqLenIsZero = false; + tempLoopInfo.s2BasicSizeTail = tempLoopInfo.actS2Size % constInfo.s2BaseSize; + tempLoopInfo.s2BasicSizeTail = + (tempLoopInfo.s2BasicSizeTail == 0) ? constInfo.s2BaseSize : tempLoopInfo.s2BasicSizeTail; + tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize; + tempLoopInfo.mBasicSizeTail = + (tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail; + + uint32_t gS1SplitNum = (tempLoopInfo.actS1Size * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize; + tempLoopInfo.gS1LoopEnd = (bN2LoopIdx == splitCoreInfo.bN2End) ? splitCoreInfo.gS1End : gS1SplitNum - 1; + if constexpr (Q_LAYOUT_T == LI_LAYOUT::BSND) { + if (tempLoopInfo.gS1LoopEnd == gS1SplitNum - 1 && constInfo.qSeqSize > tempLoopInfo.actS1Size) { + tempLoopInfo.needDealActS1LessThanS1 = true; + } + } +} + +template +__aicore__ inline void QLIPreload::CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, QLICommon::RunInfo &runInfo) +{ + runInfo.loop = loop; + runInfo.bIdx = tempLoopInfo.bIdx; + runInfo.gS1Idx = tempLoopInfo.gS1Idx; + runInfo.s2Idx = s2LoopIdx; + runInfo.bN2Idx = tempLoopInfo.bN2Idx; + runInfo.isValid = s2LoopIdx <= tempLoopInfo.s2LoopEnd; + + if (!runInfo.isValid) { + return; // 需要验证， v1 时候需要runInfo + } + + runInfo.actS1Size = tempLoopInfo.actS1Size; + runInfo.actS2Size = tempLoopInfo.actS2Size; + runInfo.actS2SizeOrig = tempLoopInfo.actS2SizeOrig; + // 计算实际基本块size + runInfo.actMBaseSize = tempLoopInfo.actMBaseSize; + runInfo.actualSingleProcessSInnerSize = constInfo.s2BaseSize; + uint32_t s2SplitNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize; + if (runInfo.s2Idx == s2SplitNum - 1) { + runInfo.actualSingleProcessSInnerSize = tempLoopInfo.s2BasicSizeTail; + } + runInfo.actualSingleProcessSInnerSizeAlign = + QLICommon::Align((uint32_t)runInfo.actualSingleProcessSInnerSize, QLICommon::ConstInfo::BUFFER_SIZE_BYTE_32B); + + runInfo.isFirstS2InnerLoop = s2LoopIdx == splitCoreInfo.s2Start; + runInfo.isLastS2InnerLoop = s2LoopIdx == tempLoopInfo.s2LoopEnd; + runInfo.isAllLoopEnd = (runInfo.bN2Idx == splitCoreInfo.bN2End) && (runInfo.gS1Idx == splitCoreInfo.gS1End) && + (runInfo.s2Idx == splitCoreInfo.s2End); + + if (runInfo.isFirstS2InnerLoop) { + uint64_t actualSeqQPrefixSum; + if constexpr (Q_LAYOUT_T == LI_LAYOUT::TND) { + uint32_t actualSeqLengthsGmQIdx = (constInfo.batchSupperFlag) ? runInfo.bIdx : runInfo.bIdx - 1; + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGmQ.GetValue(actualSeqLengthsGmQIdx); + } else { // BSND + actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.qSeqSize; + } + uint64_t tndBIdxOffset = actualSeqQPrefixSum * constInfo.qHeadNum * constInfo.headDim; + // B,S1,N1(N2,G),D + queryCoreOffset = tndBIdxOffset + runInfo.gS1Idx * constInfo.mBaseSize * constInfo.headDim; + // B,S1,N1(N2,G)/T,N1(N2,G) + weightsCoreOffset = actualSeqQPrefixSum * constInfo.qHeadNum + runInfo.n2Idx * constInfo.gSize; + // B,S1,N2,k/T,N2,k + indiceOutCoreOffset = + actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + runInfo.n2Idx * constInfo.sparseCount; + } + uint64_t actualSeqKPrefixSum; + if constexpr (K_LAYOUT_T == LI_LAYOUT::TND) { // T N2 D + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGm.GetValue(runInfo.bIdx - 1); + actualSeqKPrefixSum = actualSeqKPrefixSum / constInfo.cmpRatio; + } else { + actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.kSeqSize; + } + uint64_t tndBIdxOffsetForK = actualSeqKPrefixSum * constInfo.kHeadNum * constInfo.headDim; + keyCoreOffset = tndBIdxOffsetForK + runInfo.s2Idx * constInfo.s2BaseSize * constInfo.kHeadNum * constInfo.headDim; + keyScaleCoreOffset = (actualSeqKPrefixSum + runInfo.s2Idx * constInfo.s2BaseSize) * constInfo.kHeadNum; + runInfo.tensorQueryOffset = queryCoreOffset; + runInfo.tensorKeyOffset = keyCoreOffset; + runInfo.tensorKeyScaleOffset = keyScaleCoreOffset; + runInfo.tensorWeightsOffset = weightsCoreOffset; + runInfo.indiceOutOffset = indiceOutCoreOffset; +} + +template +__aicore__ inline void QLIPreload::Process() +{ + if (isUsedCoreEqZero) { + // 没有计算任务，直接清理输出 + ProcessInvalid(); + return; + } + + ProcessMain(); +} + +template +__aicore__ inline void QLIPreload::ProcessInvalid() +{ + if ASCEND_IS_AIV { + uint32_t aivCoreNum = GetBlockNum() * 2; // 2 means c:v = 1:2 + uint64_t totalOutputSize = + constInfo.batchSize * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount; + uint64_t singleCoreSize = + QLICommon::Align((totalOutputSize + aivCoreNum - 1) / aivCoreNum, GM_ALIGN_BYTES / sizeof(OUT_T)); + uint64_t baseSize = tmpBlockIdx * singleCoreSize; + if (baseSize < totalOutputSize) { + uint64_t dealSize = + (baseSize + singleCoreSize <= totalOutputSize) ? singleCoreSize : totalOutputSize - baseSize; + GlobalTensor output = indiceOutGm[baseSize]; + AscendC::InitGlobalMemory(output, dealSize, constInfo.INVALID_IDX); + } + } +} + +template +__aicore__ inline void QLIPreload::ProcessMain() +{ + if(!splitCoreInfo.isCoreEnable){ + return; + } + + if ASCEND_IS_AIV { + vectorService.AllocEventID(); + CrossCoreSetFlag(QLICommon::ConstInfo::CROSS_VC_EVENT + 0); + CrossCoreSetFlag(QLICommon::ConstInfo::CROSS_VC_EVENT + 1); + } else { + matmulService.AllocEventID(); + } + + QLICommon::RunInfo runInfo; + uint32_t gloop = 0; + for (uint32_t bN2LoopIdx = splitCoreInfo.bN2Start; bN2LoopIdx <= splitCoreInfo.bN2End; bN2LoopIdx++) { + CalcGS1LoopParams(bN2LoopIdx); + if (tempLoopInfo.curActSeqLenIsZero) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, 0U); + continue; + } + for (uint32_t gS1LoopIdx = splitCoreInfo.gS1Start; gS1LoopIdx <= tempLoopInfo.gS1LoopEnd; gS1LoopIdx++) { + CalcS2LoopParams(bN2LoopIdx, gS1LoopIdx); + for (int s2LoopIdx = splitCoreInfo.s2Start; s2LoopIdx <= tempLoopInfo.s2LoopEnd; s2LoopIdx++) { + ProcessBaseBlock(gloop, s2LoopIdx, runInfo); + ++gloop; + } + splitCoreInfo.s2Start = 0; + } + if (tempLoopInfo.needDealActS1LessThanS1) { + DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, tempLoopInfo.actS1Size); + } + splitCoreInfo.gS1Start = 0; + } + + if ASCEND_IS_AIV { + vectorService.FreeEventID(); + } else { + matmulService.FreeEventID(); + CrossCoreWaitFlag(QLICommon::ConstInfo::CROSS_VC_EVENT + 0); + CrossCoreWaitFlag(QLICommon::ConstInfo::CROSS_VC_EVENT + 1); + } +} + +template +__aicore__ inline void QLIPreload::ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx, QLICommon::RunInfo runInfo) +{ + CalcRunInfo(loop, s2LoopIdx, runInfo); + if ASCEND_IS_AIC { + matmulService.ComputeMm1(runInfo); + } else { + vectorService.ProcessVec1(runInfo); + if (runInfo.isLastS2InnerLoop) { //本核s2last + vectorService.ProcessTopK(runInfo); + } + } +} + +} // namespace QLIKernel +#endif // quant_lightning_indexer_KERNEL_H diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/quant_lightning_indexer_service_cube.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/quant_lightning_indexer_service_cube.h new file mode 100644 index 000000000..548126f4b --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/quant_lightning_indexer_service_cube.h @@ -0,0 +1,438 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file lightning_indexer_service_cube.h + * \brief use 5 buffer for matmul l1, better pipeline + */ +#ifndef quant_lightning_indexer_SERVICE_CUBE_H +#define quant_lightning_indexer_SERVICE_CUBE_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "quant_lightning_indexer_common.h" + +namespace QLIKernel { +using namespace QLICommon; +template +class QLIMatmul { +public: + using Q_T = typename QLIT::queryType; + using K_T = typename QLIT::keyType; + using QK_T = typename QLIT::queryKeyType; + + __aicore__ inline QLIMatmul(){}; + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitMm1GlobalTensor(const GlobalTensor &blkTableGm, const GlobalTensor &keyGm, + const GlobalTensor &queryGm); + __aicore__ inline void InitParams(const ConstInfo &constInfo); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + __aicore__ inline void ComputeMm1(const QLICommon::RunInfo &runInfo); + + static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding; + static constexpr uint64_t KEY_BUF_NUM = 3; + static constexpr uint64_t QUERY_BUF_NUM = 2; + static constexpr uint64_t L0_BUF_NUM = 2; + + static constexpr uint32_t KEY_MTE1_MTE2_EVENT = EVENT_ID2; + static constexpr uint32_t QUERY_MTE1_MTE2_EVENT = EVENT_ID5; // KEY_MTE1_MTE2_EVENT + KEY_BUF_NUM; + static constexpr uint32_t M_MTE1_EVENT = EVENT_ID3; + + static constexpr uint32_t MTE2_MTE1_EVENT = EVENT_ID2; + static constexpr uint32_t MTE1_M_EVENT = EVENT_ID2; + static constexpr uint32_t FIX_M_EVENT = EVENT_ID2; + static constexpr uint32_t M_FIX_EVENT = EVENT_ID3; + + static constexpr uint64_t M_BASIC_BLOCK = 256; + static constexpr uint64_t D_BASIC_BLOCK = 128; + static constexpr uint64_t S2_BASIC_BLOCK = 128; + + static constexpr uint64_t M_BASIC_BLOCK_L0 = 256; + static constexpr uint64_t D_BASIC_BLOCK_L0 = 128; + static constexpr uint64_t S2_BASIC_BLOCK_L0 = 128; + + static constexpr uint64_t FP8_BLOCK_CUBE = 32; + static constexpr FixpipeConfig QLI_CFG_ROW_MAJOR_UB = {CO2Layout::ROW_MAJOR, true}; // ROW_MAJOR: 使能NZ2ND，输出数据格式为ND格式; true: 用于用户指定目的地址的位置是否是UB + + static constexpr uint64_t QUERY_BUFFER_OFFSET = M_BASIC_BLOCK * D_BASIC_BLOCK; + static constexpr uint64_t KEY_BUFFER_OFFSET = S2_BASIC_BLOCK * D_BASIC_BLOCK; + static constexpr uint64_t L0AB_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0; + static constexpr uint64_t L0C_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0; + +protected: + __aicore__ inline void Fixp(uint64_t s1gGmOffset, uint64_t s2GmOffset, uint64_t s1gL0RealSize, + uint64_t s2L0RealSize, const QLICommon::RunInfo &runInfo); + __aicore__ inline void ComuteL0c(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, const QLICommon::RunInfo &runInfo); + __aicore__ inline void LoadKeyToL0b(uint64_t s2L0Offset, uint64_t s2L1RealSize, uint64_t s2L0RealSize, + const QLICommon::RunInfo &runInfo); + __aicore__ inline void LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL0Offset, uint64_t s1gL1RealSize, + uint64_t s1gL0RealSize, const QLICommon::RunInfo &runInfo); + __aicore__ inline void QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gL1Offset, const QLICommon::RunInfo &runInfo); + __aicore__ inline void KeyNd2Nz(uint64_t s2L1RealSize, uint64_t s2GmOffset, const QLICommon::RunInfo &runInfo); + __aicore__ inline void KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, const QLICommon::RunInfo &runInfo); + GlobalTensor blkTableGm_; + GlobalTensor keyGm_; + GlobalTensor queryGm_; + + TBuf bufQL1_; + LocalTensor queryL1_; + TBuf bufKeyL1_; + LocalTensor keyL1_; + + TBuf bufQL0_; + LocalTensor queryL0_; + TBuf bufKeyL0_; + LocalTensor keyL0_; + + TBuf bufL0C_; + LocalTensor cL0_; + + TBuf bufUB_; + LocalTensor mm1ResUB_; + + uint64_t keyL1BufIdx_ = 0; + uint64_t queryL1Mte2BufIdx_ = 0; + uint64_t queryL1Mte1BufIdx_ = 0; + uint64_t l0BufIdx_ = 0; + + ConstInfo constInfo_; + +private: + static constexpr bool PAGE_ATTENTION = QLIT::pageAttention; +}; + +template +__aicore__ inline void QLIMatmul::InitParams(const ConstInfo &constInfo) +{ + constInfo_ = constInfo; +} + +template +__aicore__ inline void QLIMatmul::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(bufUB_, 2 * CeilDiv(constInfo_.mBaseSize, 2) * constInfo_.s2BaseSize * sizeof(QK_T)); //大小：2(开dB) * 2 * 64 * 128 * 4 = 128KB + mm1ResUB_ = bufUB_.Get(); + pipe->InitBuffer(bufQL1_, QUERY_BUF_NUM * M_BASIC_BLOCK * D_BASIC_BLOCK * sizeof(Q_T)); + queryL1_ = bufQL1_.Get(); + pipe->InitBuffer(bufKeyL1_, KEY_BUF_NUM * S2_BASIC_BLOCK * D_BASIC_BLOCK * sizeof(K_T)); + keyL1_ = bufKeyL1_.Get(); + + pipe->InitBuffer(bufQL0_, L0_BUF_NUM * M_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0 * sizeof(Q_T)); + queryL0_ = bufQL0_.Get(); + pipe->InitBuffer(bufKeyL0_, L0_BUF_NUM * D_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0 * sizeof(K_T)); + keyL0_ = bufKeyL0_.Get(); + + pipe->InitBuffer(bufL0C_, L0_BUF_NUM * M_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0 * sizeof(float)); + cL0_ = bufL0C_.Get(); +} + +template +__aicore__ inline void +QLIMatmul::InitMm1GlobalTensor(const GlobalTensor &blkTableGm, const GlobalTensor &keyGm, + const GlobalTensor &queryGm) +{ + blkTableGm_ = blkTableGm; + keyGm_ = keyGm; + queryGm_ = queryGm; +} + +template +__aicore__ inline void QLIMatmul::ComputeMm1(const QLICommon::RunInfo &runInfo) +{ + CrossCoreWaitFlag(QLICommon::ConstInfo::CROSS_VC_EVENT + runInfo.loop % 2); + CrossCoreWaitFlag(QLICommon::ConstInfo::CROSS_VC_EVENT + runInfo.loop % 2 + QLICommon::ConstInfo::AIV0_AIV1_OFFSET); + uint64_t s2GmBaseOffset = runInfo.s2Idx * constInfo_.s2BaseSize; + uint64_t s1gProcessSize = runInfo.actMBaseSize; + uint64_t s2ProcessSize = runInfo.actualSingleProcessSInnerSize; + for (uint64_t s2GmOffset = 0; s2GmOffset < s2ProcessSize; s2GmOffset += S2_BASIC_BLOCK) { + WaitFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM); + uint64_t s2L1RealSize = + s2GmOffset + S2_BASIC_BLOCK > s2ProcessSize ? s2ProcessSize - s2GmOffset : S2_BASIC_BLOCK; + if (PAGE_ATTENTION) { + KeyNd2NzForPA(s2L1RealSize, s2GmBaseOffset + s2GmOffset, runInfo); + }else { + KeyNd2Nz(s2L1RealSize, s2GmOffset, runInfo); + } + + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + // s1gProcessSize当前必定不会超过2倍的s1g basic block + for (uint64_t s1gGmOffset = 0; s1gGmOffset < s1gProcessSize; s1gGmOffset += M_BASIC_BLOCK) { + uint64_t s1gL1RealSize = + s1gGmOffset + M_BASIC_BLOCK > s1gProcessSize ? s1gProcessSize - s1gGmOffset : M_BASIC_BLOCK; + uint64_t s1gL1SizeAlign2G = CeilAlign(s1gL1RealSize, 2 * constInfo_.gSize); + if (runInfo.isFirstS2InnerLoop && s2GmOffset == 0) { + queryL1Mte2BufIdx_++; + queryL1Mte1BufIdx_ = queryL1Mte2BufIdx_; + WaitFlag(QUERY_MTE1_MTE2_EVENT + queryL1Mte2BufIdx_ % QUERY_BUF_NUM); + QueryNd2Nz(s1gL1SizeAlign2G, s1gGmOffset, runInfo); + SetFlag(MTE2_MTE1_EVENT); + WaitFlag(MTE2_MTE1_EVENT); + } else { + queryL1Mte1BufIdx_ = + queryL1Mte2BufIdx_ - (CeilDiv(s1gProcessSize, M_BASIC_BLOCK) - 1 - (s1gGmOffset > 0)); + } + for (uint64_t s2L1Offset = 0; s2L1Offset < s2L1RealSize; s2L1Offset += S2_BASIC_BLOCK_L0) { + uint64_t s2L0RealSize = + s2L1Offset + S2_BASIC_BLOCK_L0 > s2L1RealSize ? s2L1RealSize - s2L1Offset : S2_BASIC_BLOCK_L0; + for (uint64_t s1gL1Offset = 0; s1gL1Offset < s1gL1SizeAlign2G; s1gL1Offset += M_BASIC_BLOCK_L0) { + WaitFlag(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM); + uint64_t s1gL0RealSize = + s1gL1Offset + M_BASIC_BLOCK_L0 > s1gL1SizeAlign2G ? s1gL1SizeAlign2G - s1gL1Offset : M_BASIC_BLOCK_L0; + LoadQueryToL0a(s1gGmOffset, s1gL1Offset, s1gL1SizeAlign2G, s1gL0RealSize, runInfo); + LoadKeyToL0b(s2L1Offset, s2L1RealSize, s2L0RealSize, runInfo); + + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + + WaitFlag(FIX_M_EVENT + l0BufIdx_ % L0_BUF_NUM); + ComuteL0c(s1gL0RealSize, s2L0RealSize, runInfo); + + SetFlag(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM); + + Fixp(s1gGmOffset + s1gL1Offset, s2GmOffset + s2L1Offset, s1gL0RealSize, s2L0RealSize, runInfo); + SetFlag(FIX_M_EVENT + l0BufIdx_ % L0_BUF_NUM); + l0BufIdx_++; + } + } + if (s2GmOffset + S2_BASIC_BLOCK >= s2ProcessSize && runInfo.isLastS2InnerLoop) { + SetFlag(QUERY_MTE1_MTE2_EVENT + queryL1Mte1BufIdx_ % QUERY_BUF_NUM); + } + } + SetFlag(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM); + keyL1BufIdx_++; + } + CrossCoreSetFlag(QLICommon::ConstInfo::CROSS_CV_EVENT + runInfo.loop % 2); + CrossCoreSetFlag(QLICommon::ConstInfo::CROSS_CV_EVENT + runInfo.loop % 2 + QLICommon::ConstInfo::AIV0_AIV1_OFFSET); +} + +template +__aicore__ inline void QLIMatmul::KeyNd2Nz(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const QLICommon::RunInfo &runInfo) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2L1RealSize; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + // 默认一块buf最多放两份 + DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET], + keyGm_[runInfo.tensorKeyOffset + s2GmOffset * constInfo_.headDim], nd2nzPara); +} + +// blkNum, blkSize, N2, D +template +__aicore__ inline void QLIMatmul::KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, + const QLICommon::RunInfo &runInfo) +{ + uint64_t s2L1Offset = 0; + while (s2L1Offset < s2L1RealSize) { + uint64_t s2BlkId = (s2L1Offset + s2GmOffset) / constInfo_.kCacheBlockSize; + uint64_t s2BlkOffset = (s2L1Offset + s2GmOffset) % constInfo_.kCacheBlockSize; + uint64_t keyGmOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId) * + constInfo_.stride + + s2BlkOffset * constInfo_.headDim; + + uint64_t s2Mte2Size = s2L1RealSize - s2L1Offset; + s2Mte2Size = s2BlkOffset + s2Mte2Size >= constInfo_.kCacheBlockSize ? constInfo_.kCacheBlockSize - s2BlkOffset + : s2Mte2Size; + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s2Mte2Size; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET + s2L1Offset * FP8_BLOCK_CUBE], + keyGm_[keyGmOffset], nd2nzPara); + + s2L1Offset += s2Mte2Size; + } +} + +// batch, s1, n2, g, d +template +__aicore__ inline void QLIMatmul::QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gGmOffset, + const QLICommon::RunInfo &runInfo) +{ + Nd2NzParams nd2nzPara; + nd2nzPara.ndNum = 1; + nd2nzPara.nValue = s1gL1RealSize; // 行数 + nd2nzPara.dValue = constInfo_.headDim; + nd2nzPara.srcDValue = constInfo_.headDim; + nd2nzPara.dstNzC0Stride = CeilAlign(s1gL1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block + nd2nzPara.dstNzNStride = 1; + nd2nzPara.srcNdMatrixStride = 0; + nd2nzPara.dstNzMatrixStride = 0; + // 默认一块buf最多放两份 + DataCopy(queryL1_[(queryL1Mte2BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET], + queryGm_[runInfo.tensorQueryOffset + s1gGmOffset * constInfo_.headDim], nd2nzPara); +} + +template +__aicore__ inline void QLIMatmul::LoadQueryToL0a(uint64_t s1gGmOffset, uint64_t s1gL1Offset, uint64_t s1gL1RealSize, + uint64_t s1gL0RealSize, const QLICommon::RunInfo &runInfo) +{ + LoadData2DParamsV2 loadData2DParamsV2; + loadData2DParamsV2.mStartPosition = CeilDiv(s1gL1Offset, BLOCK_CUBE); + loadData2DParamsV2.kStartPosition = 0; + loadData2DParamsV2.mStep = CeilDiv(s1gL0RealSize, BLOCK_CUBE); + loadData2DParamsV2.kStep = CeilDiv(constInfo_.headDim, FP8_BLOCK_CUBE); + loadData2DParamsV2.srcStride = CeilDiv(s1gL1RealSize, BLOCK_CUBE); + loadData2DParamsV2.dstStride = CeilDiv(s1gL0RealSize, BLOCK_CUBE); + loadData2DParamsV2.ifTranspose = false; + + LoadData(queryL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], + queryL1_[(queryL1Mte1BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET], loadData2DParamsV2); +} + +template +__aicore__ inline void QLIMatmul::LoadKeyToL0b(uint64_t s2L1Offset, uint64_t s2L1RealSize, uint64_t s2L0RealSize, + const QLICommon::RunInfo &runInfo) +{ + LoadData2DParamsV2 loadData2DParamsV2; + loadData2DParamsV2.mStartPosition = CeilDiv(s2L1Offset, BLOCK_CUBE); + loadData2DParamsV2.kStartPosition = 0; + loadData2DParamsV2.mStep = CeilDiv(s2L0RealSize, BLOCK_CUBE); + loadData2DParamsV2.kStep = CeilDiv(constInfo_.headDim, FP8_BLOCK_CUBE); + loadData2DParamsV2.srcStride = CeilDiv(s2L1RealSize, BLOCK_CUBE); + loadData2DParamsV2.dstStride = CeilDiv(s2L0RealSize, BLOCK_CUBE); + loadData2DParamsV2.ifTranspose = false; + + LoadData(keyL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], + keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET], loadData2DParamsV2); +} + +template +__aicore__ inline void QLIMatmul::ComuteL0c(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, + const QLICommon::RunInfo &runInfo) +{ + MmadParams mmadParams; + mmadParams.m = CeilAlign(s1gL0RealSize, BLOCK_CUBE); + mmadParams.n = s2L0RealSize; + mmadParams.k = constInfo_.headDim; + mmadParams.cmatrixInitVal = true; + mmadParams.cmatrixSource = false; + Mmad(cL0_[(l0BufIdx_ % L0_BUF_NUM) * L0C_BUFFER_OFFSET], queryL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], + keyL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], mmadParams); + if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) { + PipeBarrier(); + } +} + +template +__aicore__ inline void QLIMatmul::Fixp(uint64_t s1gGmOffset, uint64_t s2GmOffset, uint64_t s1gL0RealSize, + uint64_t s2L0RealSize, const QLICommon::RunInfo &runInfo) +{ + SetFlag(M_FIX_EVENT + l0BufIdx_ % L0_BUF_NUM); + WaitFlag(M_FIX_EVENT + l0BufIdx_ % L0_BUF_NUM); + + static_assert(S2_BASIC_BLOCK == S2_BASIC_BLOCK_L0 && S2_BASIC_BLOCK_L0 == 128); + if constexpr (std::is_same_v) { + // s1gL0RealSize：2*gSize(128)对齐, 最大256 + // s2L0RealSize <= S2_BASIC_BLOCK_L0, 未约束 + uint32_t nSize = (s2L0RealSize + 7) >> 3 << 3; // 32B对齐 + uint32_t mSize = (s1gL0RealSize + 1) >> 1 << 1; + FixpipeParamsC310 fixpipeParams; + // 固定参数 + fixpipeParams.mSize = mSize; + fixpipeParams.srcStride = mSize; // 已16对齐 + fixpipeParams.dstStride = UB_BANK_DEPTH_STRIDE / sizeof(QK_T); // 落到同一个bank + fixpipeParams.dualDstCtl = 1; // 双目标模式，按M维度拆分， M / 2 * N写入每个UB，M必须为2的倍数 + + // nSize已保证N方向32B对齐 + if (nSize <= (256 / sizeof(float))) { + // N方向小于一个bank(256B), 只需搬一个ND块, 且不用补齐 + fixpipeParams.nSize = nSize; + fixpipeParams.params.ndNum = 1; + fixpipeParams.params.srcNdStride = 0; + fixpipeParams.params.dstNdStride = 0; + } else { + // N方向在(256B, 512B]范围， 直接按512B搬, 注意此时不能开unitflag + fixpipeParams.nSize = S2_BASIC_BLOCK_L0 / 2; // 分2个ND搬, S2_BASIC_BLOCK_L0不为128会有问题 + fixpipeParams.params.ndNum = 2; + fixpipeParams.params.srcNdStride = ((fixpipeParams.mSize + 15) / 16) * fixpipeParams.nSize; + fixpipeParams.params.dstNdStride = constInfo_.s2BaseSize * constInfo_.mBaseSize / 2; // S2_BASIC_BLOCK * M_BASE_SIZE / 2 + } + Fixpipe(mm1ResUB_[(runInfo.loop % 2) * constInfo_.s2BaseSize / 2], // 未考虑s1gGmOffset和s2GmOffset + cL0_[(l0BufIdx_ % L0_BUF_NUM) * L0C_BUFFER_OFFSET], fixpipeParams); // 将matmul结果从L0C搬运到UB + } else { + // nSize * sizeof(QT) <= 256B, 小于一个UB bank大小(VL) + uint32_t nSize = (s2L0RealSize + 7) >> 3 << 3; // 8个元素（32B)对齐 + uint32_t mSize = (s1gL0RealSize + 1) >> 1 << 1; // 有效数据不足16行，只需输出部分行即可;L0C上的bmm1结果矩阵M方向的size大小必须是偶数 + uint32_t srcStride = ((mSize + 15) / 16) * 16; // L0C上matmul结果相邻连续数据片断间隔（前面一个数据块的头与后面数据块的头的间隔），单位为16 *sizeof(T) //源NZ矩阵中相邻Z排布的起始地址偏移 + FixpipeParamsC310 fixpipeParams; // L0C->UB + fixpipeParams.nSize = nSize; // N方向全部输出 + fixpipeParams.mSize = mSize / 2; // M方向每个AIV一半 + fixpipeParams.srcStride = srcStride; + fixpipeParams.dstStride = UB_BANK_DEPTH_STRIDE / sizeof(QK_T); // 落到同一个bank + fixpipeParams.params.ndNum = 1; + fixpipeParams.params.srcNdStride = 0; + fixpipeParams.params.dstNdStride = 0; + fixpipeParams.dualDstCtl = 0; + fixpipeParams.quantPre = F322BF16; + fixpipeParams.reluEn = true; // ReLU激活 + fixpipeParams.subBlockId = 0; + Fixpipe(mm1ResUB_[(runInfo.loop % 2) * (UB_BANK_STRIDE / sizeof(QK_T))], // 未考虑s1gGmOffset和s2GmOffset + cL0_[(l0BufIdx_ % L0_BUF_NUM) * L0C_BUFFER_OFFSET], fixpipeParams); // 将matmul结果从L0C搬运到UB + + fixpipeParams.subBlockId = 1; + Fixpipe(mm1ResUB_[(runInfo.loop % 2) * (UB_BANK_STRIDE / sizeof(QK_T))], // 未考虑s1gGmOffset和s2GmOffset + cL0_[(l0BufIdx_ % L0_BUF_NUM) * L0C_BUFFER_OFFSET + mSize / 2 * 16], fixpipeParams); // 将matmul结果从L0C搬运到UB + } +} + +template +__aicore__ inline void QLIMatmul::AllocEventID() +{ + SetMMLayoutTransform(true); + SetFlag(KEY_MTE1_MTE2_EVENT + 0); + SetFlag(KEY_MTE1_MTE2_EVENT + 1); + SetFlag(KEY_MTE1_MTE2_EVENT + 2); + + SetFlag(QUERY_MTE1_MTE2_EVENT + 0); + SetFlag(QUERY_MTE1_MTE2_EVENT + 1); + + SetFlag(M_MTE1_EVENT + 0); + SetFlag(M_MTE1_EVENT + 1); + + SetFlag(FIX_M_EVENT + 0); + SetFlag(FIX_M_EVENT + 1); +} + +template +__aicore__ inline void QLIMatmul::FreeEventID() +{ + SetMMLayoutTransform(false); + WaitFlag(KEY_MTE1_MTE2_EVENT + 0); + WaitFlag(KEY_MTE1_MTE2_EVENT + 1); + WaitFlag(KEY_MTE1_MTE2_EVENT + 2); + + WaitFlag(QUERY_MTE1_MTE2_EVENT + 0); + WaitFlag(QUERY_MTE1_MTE2_EVENT + 1); + + WaitFlag(M_MTE1_EVENT + 0); + WaitFlag(M_MTE1_EVENT + 1); + + WaitFlag(FIX_M_EVENT + 0); + WaitFlag(FIX_M_EVENT + 1); +} +} // namespace QLIKernel +#endif diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/quant_lightning_indexer_service_vector.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/quant_lightning_indexer_service_vector.h new file mode 100644 index 000000000..5e83eca2d --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/quant_lightning_indexer_service_vector.h @@ -0,0 +1,519 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_service_vector.h + * \brief + */ +#ifndef quant_lightning_indexer_SERVICE_VECTOR_H +#define quant_lightning_indexer_SERVICE_VECTOR_H + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "lib/matrix/matmul/tiling.h" +#include "quant_lightning_indexer_common.h" +#include "../arch35/vf/quant_lightning_indexer_vector1.h" +#include "../arch35/vf/quant_lightning_indexer_topk.h" + +namespace QLIKernel { +using namespace QLICommon; +constexpr uint32_t TRUNK_LEN_16K = 16384; +template +class QLIVector { +public: + // =================================类型定义区================================= + static constexpr LI_LAYOUT Q_LAYOUT_T = QLIT::layout; + static constexpr LI_LAYOUT K_LAYOUT_T = QLIT::keyLayout; + static constexpr bool PAGE_ATTENTION = QLIT::pageAttention; + + using QK_T = typename QLIT::queryKeyType; + using SCORE_T = typename QLIT::scoreType; + + __aicore__ inline QLIVector(){}; + __aicore__ inline void ProcessVec1(const QLICommon::RunInfo &info); + __aicore__ inline void ProcessTopK(const QLICommon::RunInfo &info); + __aicore__ inline void InitBuffers(TPipe *pipe); + __aicore__ inline void InitParams(const struct QLICommon::ConstInfo &constInfo, + const QLITilingData *__restrict tilingData); + __aicore__ inline void InitVecWorkspaceTensor(GlobalTensor scoreGm); + __aicore__ inline void InitVecInputTensor(GlobalTensor weightsGm, GlobalTensor qScaleGm, + GlobalTensor kScaleGm, GlobalTensor indiceOutGm, + GlobalTensor blockTableGm); + __aicore__ inline void CleanInvalidOutput(int64_t invalidS1offset); + __aicore__ inline void AllocEventID(); + __aicore__ inline void FreeEventID(); + +protected: + GlobalTensor scoreGm; + GlobalTensor weightsGm; + GlobalTensor qScaleGm; + GlobalTensor kScaleGm; + GlobalTensor indiceOutGm; + GlobalTensor blockTableGm; + // =================================常量区================================= + static constexpr uint32_t VEC1_V_MTE2_EVENT = EVENT_ID0; + static constexpr uint32_t VEC1_MTE2_V_EVENT = EVENT_ID1; + static constexpr uint32_t VEC1_V_MTE3_EVENT = EVENT_ID2; + static constexpr uint32_t VEC1_MTE3_V_EVENT = EVENT_ID3; + + static constexpr uint32_t TOPK_V_MTE2_EVENT = EVENT_ID4; + static constexpr uint32_t TOPK_MTE2_V_EVENT = EVENT_ID5; + static constexpr uint32_t TOPK_V_MTE3_EVENT = EVENT_ID6; + static constexpr uint32_t TOPK_MTE3_V_EVENT = EVENT_ID7; + + static constexpr uint32_t KSCALE_S_MTE2_EVENT = EVENT_ID7; + static constexpr uint32_t MTE3_MTE2_EVENT = EVENT_ID0; + static constexpr uint32_t V_MTE2_EVENT = EVENT_ID7; + static constexpr uint32_t V_MTE2_EVENT1 = EVENT_ID2; + static constexpr uint32_t V_MTE2_EVENT2 = EVENT_ID3; + static constexpr uint32_t V_MTE2_EVENT3 = EVENT_ID5; + +private: + __aicore__ inline void GetKeyScale(const QLICommon::RunInfo &runInfo, LocalTensor &kScaleUB, + int64_t batchId, int64_t startS2, int64_t getLen); + // ================================Local Buffer区==================================== + + // tmp buff for vector + TBuf resMm1Buf_; + LocalTensor resMm1UB_; + //tmp buff for weight + TBuf weightBuf_; + LocalTensor weightUB_; + //tmp buff for kScale + TBuf kScaleBuf_; + LocalTensor kScaleUB_; + //tmp buff for qScale + TBuf qScaleBuf_; + LocalTensor qScaleUB_; + //tmp buff for out + TBuf outBuf_; + LocalTensor vec1OutUB_; + // tmp buff for LD + + // tmp buff for topk + TBuf mrgValueBuf_; + LocalTensor mrgValueLocal_; + + TBuf indicesOutBuf_; + LocalTensor indicesOutLocal_; + + TBuf scoreOutBuf_; + LocalTensor scoreOutLocal_; + + TBuf topkSharedTmpBuf_; + LocalTensor topkSharedTmpLocal_; + + TBuf outInvalidBuf_; + LocalTensor outInvalidLocal_; + + int32_t blockId_ = -1; + // para for vector + int32_t groupInner_ = 0; + int32_t globalTopkNum_ = 0; + int64_t blockS2StartIdx_ = 0; + int32_t gSize_ = 0; + int32_t kSeqSize_ = 0; + int32_t kHeadNum_ = 0; + int32_t qHeadNum_ = 0; + int32_t s1BaseSize_ = 0; + int32_t s2BaseSize_ = 0; + int32_t kCacheBlockSize_ = 0; + int32_t maxBlockNumPerBatch_ = 0; + uint32_t topkCount_ = 0; + uint32_t topkCountAlign256_ = 0; // topkCount对齐到256(直方图需要)，支持topk泛化 + uint32_t trunkLen_ = 0; + + struct QLICommon::ConstInfo constInfo_; + topk::LITopk topkOp_; +}; + +template +__aicore__ inline void QLIVector::InitBuffers(TPipe *pipe) +{ + pipe->InitBuffer(resMm1Buf_, 2 * CeilDiv(constInfo_.mBaseSize, 2) * s2BaseSize_ * sizeof(QK_T)); //大小：2(开dB) * 2 * 64 * 128 * 4 = 128KB + resMm1UB_ = resMm1Buf_.Get();//qk + pipe->InitBuffer(weightBuf_, 2 * CeilDiv(s1BaseSize_, 2) * gSize_* sizeof(float)); // 大小：2(开dB) * 2 * 64 * 2 = 0.5KB + weightUB_ = weightBuf_.Get();//weight + pipe->InitBuffer(kScaleBuf_, 2 * s2BaseSize_ * sizeof(float)); // 大小：2(开dB) * 128 * 4 = 1KB + kScaleUB_ = kScaleBuf_.Get();//kScale + pipe->InitBuffer(qScaleBuf_, 2 * CeilDiv(s1BaseSize_, 2) * gSize_* sizeof(float)); // 大小：2(开dB) * 2 * 64 * 4 = 1KB + qScaleUB_ = qScaleBuf_.Get();//qScale + pipe->InitBuffer(outBuf_, 2 * CeilDiv(s1BaseSize_, 2) * s2BaseSize_ * sizeof(SCORE_T)); // 大小：2(开dB) * 2 * 128 * 4 = 2KB + vec1OutUB_ = outBuf_.Get();//out + + // Topk + pipe->InitBuffer(mrgValueBuf_, (topkCountAlign256_ + trunkLen_) * sizeof(SCORE_T)); // 大小：(topkCountAlign256_ + 每次排序长度) * sizeof(SCORE_T) + mrgValueLocal_ = mrgValueBuf_.Get(); + + pipe->InitBuffer(indicesOutBuf_, (topkCountAlign256_ + 64) * sizeof(uint32_t)); // 大小：(topkCountAlign256_ + 64) * 4 64:duplicate刷-1需要额外空间 + indicesOutLocal_ = indicesOutBuf_.Get(); + + pipe->InitBuffer(scoreOutBuf_, topkCountAlign256_ * sizeof(SCORE_T)); // 大小：topkCountAlign256_ * sizeof(SCORE_T) + scoreOutLocal_ = scoreOutBuf_.Get(); + + uint64_t topkSharedTmpSize = topkOp_.GetSharedTmpBufferSize(); + pipe->InitBuffer(topkSharedTmpBuf_, topkSharedTmpSize); + topkSharedTmpLocal_ = topkSharedTmpBuf_.Get(); + topkOp_.InitBuffers(topkSharedTmpLocal_); + + //刷-1 + pipe->InitBuffer(outInvalidBuf_, topkCount_ * sizeof(int32_t)); + outInvalidLocal_ = outInvalidBuf_.Get(); + Duplicate(kScaleUB_, float(0), 2 * s2BaseSize_); +} + +template +__aicore__ inline void QLIVector::InitParams(const struct QLICommon::ConstInfo &constInfo, + const QLITilingData *__restrict tilingData) +{ + this->constInfo_ = constInfo; + blockS2StartIdx_ = 0; + gSize_ = constInfo.gSize; + kSeqSize_ = constInfo.kSeqSize; + // define N2 para + kHeadNum_ = constInfo.kHeadNum; + qHeadNum_ = constInfo.qHeadNum; + // define MMBase para + s1BaseSize_ = constInfo.s1BaseSize; // 4 + s2BaseSize_ = constInfo.s2BaseSize; // 128 + kCacheBlockSize_ = constInfo.kCacheBlockSize; + maxBlockNumPerBatch_ = constInfo.maxBlockNumPerBatch; + blockId_ = GetBlockIdx(); + trunkLen_ = TRUNK_LEN_16K; + topkCount_ = constInfo.sparseCount; + topkCountAlign256_ = QLICommon::Align(constInfo.sparseCount, (uint64_t)256); // topkCount对齐到256 + topkOp_.Init(topkCount_, trunkLen_); +} + +template +__aicore__ inline void QLIVector::InitVecInputTensor(GlobalTensor weightsGm, GlobalTensor qScaleGm, + GlobalTensor kScaleGm, + GlobalTensor indiceOutGm, + GlobalTensor blockTableGm) +{ + this->weightsGm = weightsGm; + this->qScaleGm = qScaleGm; + this->kScaleGm = kScaleGm; + this->indiceOutGm = indiceOutGm; + this->blockTableGm = blockTableGm; +} + +template +__aicore__ inline void QLIVector::InitVecWorkspaceTensor(GlobalTensor scoreGm) +{ + this->scoreGm = scoreGm;//resucesum*k +} + +template +__aicore__ inline void QLIVector::AllocEventID() +{ + SetFlag(VEC1_V_MTE2_EVENT + 0); + SetFlag(VEC1_V_MTE2_EVENT + 1); + SetFlag(VEC1_MTE3_V_EVENT + 0); + SetFlag(VEC1_MTE3_V_EVENT + 1); + + SetFlag(TOPK_V_MTE2_EVENT); + SetFlag(TOPK_MTE3_V_EVENT); + SetFlag(V_MTE2_EVENT1); +} + +template +__aicore__ inline void QLIVector::FreeEventID() +{ + WaitFlag(VEC1_V_MTE2_EVENT + 0); + WaitFlag(VEC1_V_MTE2_EVENT + 1); + WaitFlag(VEC1_MTE3_V_EVENT + 0); + WaitFlag(VEC1_MTE3_V_EVENT + 1); + + WaitFlag(TOPK_V_MTE2_EVENT); + WaitFlag(TOPK_MTE3_V_EVENT); + WaitFlag(V_MTE2_EVENT1); +} + +template +__aicore__ inline void QLIVector::CleanInvalidOutput(int64_t invalidS1Offset) +{ + // init -1 and copy to output + Duplicate(outInvalidLocal_, constInfo_.INVALID_IDX, constInfo_.sparseCount); + + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + + AscendC::DataCopyParams dataCopyOutParams; + dataCopyOutParams.blockCount = 1; + dataCopyOutParams.blockLen = constInfo_.sparseCount * sizeof(int32_t); + dataCopyOutParams.srcStride = 0; + dataCopyOutParams.dstStride = 0; + AscendC::DataCopyPad(indiceOutGm[invalidS1Offset], outInvalidLocal_, dataCopyOutParams); +} + +template +__aicore__ inline void QLIVector::GetKeyScale(const QLICommon::RunInfo &runInfo, LocalTensor &kScaleUB, + int64_t batchId, int64_t startS2, int64_t getLen) +{ + // startS2一定能整除kCacheBlockSize_ + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams copyInParams; + if constexpr (PAGE_ATTENTION) { + int32_t startBlockTableIdx = startS2 / kCacheBlockSize_; + int32_t startBlockTableOffset = startS2 % kCacheBlockSize_; + int32_t blockTableBatchOffset = batchId * maxBlockNumPerBatch_; + copyInParams.blockCount = 1; + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + int32_t resUbBaseOffset = 0; + if (startBlockTableOffset > 0) { + int32_t firstPartLen = + kCacheBlockSize_ - startBlockTableOffset > getLen ? getLen : kCacheBlockSize_ - startBlockTableOffset; + copyInParams.blockLen = firstPartLen * sizeof(float); + int32_t blockId = blockTableGm.GetValue(blockTableBatchOffset + startBlockTableIdx); + SetFlag(KSCALE_S_MTE2_EVENT); + WaitFlag(KSCALE_S_MTE2_EVENT); + AscendC::DataCopyPad(kScaleUB[(runInfo.loop % 2) * s2BaseSize_], + kScaleGm[blockId * constInfo_.scaleStride + startBlockTableOffset], + copyInParams, padParams); + startBlockTableIdx++; + getLen = getLen - firstPartLen; + resUbBaseOffset = firstPartLen; + } + int32_t getLoopNum = CeilDiv(getLen, kCacheBlockSize_); + copyInParams.blockLen = kCacheBlockSize_ * sizeof(float); + for (int32_t i = 0; i < getLoopNum; i++) { + if (i == getLoopNum - 1) { + copyInParams.blockLen = (getLen - i * kCacheBlockSize_) * sizeof(float); + } + int32_t blockId = blockTableGm.GetValue(blockTableBatchOffset + startBlockTableIdx + i); + SetFlag(KSCALE_S_MTE2_EVENT); + WaitFlag(KSCALE_S_MTE2_EVENT); + AscendC::DataCopyPad(kScaleUB[(runInfo.loop % 2) * s2BaseSize_ + resUbBaseOffset + i * kCacheBlockSize_], + kScaleGm[blockId * constInfo_.scaleStride], + copyInParams, padParams); + } + } else { + copyInParams.blockCount = 1; + copyInParams.blockLen = getLen * sizeof(float); + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + AscendC::DataCopyPad(kScaleUB[(runInfo.loop % 2) * s2BaseSize_], kScaleGm[runInfo.tensorKeyScaleOffset], copyInParams, padParams); + } +} + +template +__aicore__ inline void QLIVector::ProcessVec1(const QLICommon::RunInfo &info) +{ + auto pingpong = (info.loop % 2); + auto s1BaseSizePerAIV = CeilDiv(s1BaseSize_, 2); + int64_t curS1Idx = info.gS1Idx * s1BaseSize_; + int64_t curS2Idx = info.s2Idx * s2BaseSize_; + int64_t curS1ProcNum = curS1Idx + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + int64_t curAivS1Idx = curS1Idx + (blockId_ % 2) * CeilDiv(curS1ProcNum, 2); + int64_t curAivS1ProcNum = (blockId_ % 2 == 0) ? CeilDiv(curS1ProcNum, 2) : curS1ProcNum / 2; + + if (curAivS1ProcNum == 0) { + CrossCoreWaitFlag(QLICommon::ConstInfo::CROSS_CV_EVENT + pingpong); // V核等C核计算完mm1，mm1Res已搬运到UB + CrossCoreSetFlag(QLICommon::ConstInfo::CROSS_VC_EVENT + pingpong); // V核处理完，通知C核可以把mm1Res搬运到UB + return; + } + WaitFlag(VEC1_V_MTE2_EVENT + pingpong); + //weightsGm --> weightUB_ + int64_t weightGmOffset = info.tensorWeightsOffset + curAivS1Idx * kHeadNum_ * gSize_; + DataCopyPadExtParams padWeightsParams{false, 0, 0, 0}; + DataCopyExtParams qwDataCopyExtParams; + qwDataCopyExtParams.blockCount = curAivS1ProcNum; + qwDataCopyExtParams.blockLen = gSize_ * sizeof(float); + qwDataCopyExtParams.srcStride = 0; + qwDataCopyExtParams.dstStride = (UB_BANK_DEPTH_STRIDE - UB_BANK_STRIDE) / 32; + DataCopyPad(weightUB_[pingpong * (UB_BANK_STRIDE / sizeof(float))], + weightsGm[weightGmOffset], qwDataCopyExtParams, padWeightsParams); + + //qScaleGm --> qScaleUB_ + DataCopyPadExtParams padQScaleParams{false, 0, 0, 0}; + DataCopyPad(qScaleUB_[pingpong * (UB_BANK_STRIDE / sizeof(float))], + qScaleGm[weightGmOffset], qwDataCopyExtParams, padQScaleParams); + + //kScaleGm --> kScaleUB_ + GetKeyScale(info, kScaleUB_, info.bIdx, curS2Idx, info.actualSingleProcessSInnerSize); + SetFlag(VEC1_MTE2_V_EVENT + pingpong); + WaitFlag(VEC1_MTE2_V_EVENT + pingpong); + WaitFlag(VEC1_MTE3_V_EVENT + pingpong); + + //CV同步 + CrossCoreWaitFlag(QLICommon::ConstInfo::CROSS_CV_EVENT + info.loop % 2); //V核等C核计算完mm1，mm1Res已搬运到UB + + static_assert(std::is_same_v); + auto outBase = vec1OutUB_[pingpong * (UB_BANK_STRIDE / sizeof(SCORE_T))]; + auto weightBase = weightUB_[pingpong * (UB_BANK_STRIDE / sizeof(float))]; + auto qScaleBase = qScaleUB_[pingpong * (UB_BANK_STRIDE / sizeof(float))]; + auto kScaleBase = kScaleUB_[pingpong * s2BaseSize_]; + + auto qkBase = resMm1UB_[pingpong * (UB_BANK_STRIDE / sizeof(QK_T))]; + auto qkVLstride = (UB_BANK_DEPTH_STRIDE / sizeof(QK_T)) / 2 * constInfo_.mBaseSize; + vector1::BatchMulWeightAndReduceSum(outBase, UB_BANK_DEPTH_STRIDE / sizeof(SCORE_T), + qkBase, qkVLstride, (uint32_t)(gSize_ * UB_BANK_DEPTH_STRIDE / sizeof(QK_T)), + weightBase, UB_BANK_DEPTH_STRIDE / sizeof(float), + kScaleBase, (uint32_t)0, + qScaleBase, UB_BANK_DEPTH_STRIDE / sizeof(float), + gSize_, curAivS1ProcNum); + SetFlag(VEC1_V_MTE2_EVENT + pingpong); + SetFlag(VEC1_V_MTE3_EVENT + pingpong); + WaitFlag(VEC1_V_MTE3_EVENT + pingpong); + //outUB_ ---> scoreGm + int64_t vec1OutGmOffset = blockId_ % 2 == 0 ? curS2Idx : + s1BaseSizePerAIV * QLICommon::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_) + curS2Idx; + DataCopyExtParams copyOutParams; + copyOutParams.blockCount = curAivS1ProcNum; + copyOutParams.blockLen = s2BaseSize_ * sizeof(SCORE_T); + copyOutParams.srcStride = (UB_BANK_DEPTH_STRIDE - UB_BANK_STRIDE) / 32; + copyOutParams.dstStride = (QLICommon::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_) - s2BaseSize_) * sizeof(SCORE_T); + DataCopyPad(scoreGm[vec1OutGmOffset], outBase, copyOutParams); + SetFlag(VEC1_MTE3_V_EVENT + pingpong); + CrossCoreSetFlag(QLICommon::ConstInfo::CROSS_VC_EVENT + pingpong); //V核处理完，通知C核可以把mm1Res搬运到UB +} + +template +__aicore__ inline void QLIVector::ProcessTopK(const QLICommon::RunInfo &info) +{ + SetFlag(MTE3_MTE2_EVENT); + WaitFlag(MTE3_MTE2_EVENT); + + int64_t curS1Idx = info.gS1Idx * s1BaseSize_; + int64_t curS2Idx = info.s2Idx * s2BaseSize_; + int64_t curS1ProcNum = curS1Idx + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_; + int64_t curAivS1Idx = curS1Idx + (blockId_ % 2) * CeilDiv(curS1ProcNum, 2); + int64_t curAivS1ProcNum = (blockId_ % 2 == 0) ? CeilDiv(curS1ProcNum, 2) : curS1ProcNum / 2; + + AscendC::DataCopyExtParams copyInParams; + copyInParams.blockCount = 1; + copyInParams.srcStride = 0; + copyInParams.dstStride = 0; + copyInParams.rsv = 0; + + AscendC::DataCopyParams copyOutParams; + copyOutParams.blockCount = 1; + copyOutParams.blockLen = topkCount_ * sizeof(uint32_t); // bytes + copyOutParams.srcStride = 0; + copyOutParams.dstStride = 0; + + int32_t cuRealAcSeq = info.actS2Size; + if (constInfo_.attenMaskFlag) { + cuRealAcSeq = info.actS2SizeOrig - info.actS1Size + curAivS1Idx + 1; + } + + int32_t validS2Len = cuRealAcSeq; + for (uint32_t i = 0; i < curAivS1ProcNum; i++) { + uint32_t rowIdx = blockId_ % 2 * CeilDiv(curS1ProcNum, 2) + i; + uint32_t vecOffset = blockId_ % 2 * CeilDiv(s1BaseSize_, 2) + i; + + SCORE_T zero = 0; + int32_t neg = -1; + if (constInfo_.attenMaskFlag) { + validS2Len = ((int32_t)i + cuRealAcSeq) / static_cast(constInfo_.cmpRatio); + } + if (validS2Len <= 0) { + WaitFlag(TOPK_MTE3_V_EVENT); + Duplicate(indicesOutLocal_.ReinterpretCast(), neg, topkCount_); + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + AscendC::DataCopyPad(indiceOutGm[info.indiceOutOffset + (curS1Idx + rowIdx) * topkCount_], indicesOutLocal_.ReinterpretCast(), copyOutParams); + SetFlag(TOPK_MTE3_V_EVENT); + continue; + } + + WaitFlag(TOPK_V_MTE2_EVENT); + WaitFlag(TOPK_MTE3_V_EVENT); + + AscendC::DataCopyPadExtParams padParams{true, 0, 0, 0}; + if (validS2Len >= topkCount_) { + uint32_t s2LoopNum = (validS2Len + trunkLen_ - 1) / trunkLen_; + if (s2LoopNum == 1) { + uint32_t validS2LenAlign = QLICommon::Align(validS2Len, (int32_t)256); + Duplicate(mrgValueLocal_[validS2Len / 256 * 256], zero, validS2LenAlign - validS2Len / 256 * 256); + SetFlag(V_MTE2_EVENT); + WaitFlag(V_MTE2_EVENT); + copyInParams.blockLen = validS2Len * sizeof(SCORE_T); // byte + AscendC::DataCopyPadExtParams padParams{true, 0, 0, 0}; + AscendC::DataCopyPad(mrgValueLocal_, scoreGm[vecOffset * QLICommon::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_)], copyInParams, padParams); + SetFlag(TOPK_MTE2_V_EVENT); + WaitFlag(TOPK_MTE2_V_EVENT); + topkOp_(mrgValueLocal_, indicesOutLocal_, scoreOutLocal_, validS2LenAlign, 0, 1); + } else { + for (uint32_t loopIdx = 0; loopIdx < s2LoopNum; loopIdx++) { + if (loopIdx == 0) { + copyInParams.blockLen = trunkLen_ * sizeof(SCORE_T); // byte + AscendC::DataCopyPad(mrgValueLocal_, scoreGm[vecOffset * QLICommon::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_)], copyInParams, padParams); + SetFlag(TOPK_MTE2_V_EVENT); + WaitFlag(TOPK_MTE2_V_EVENT); + topkOp_(mrgValueLocal_, indicesOutLocal_, scoreOutLocal_, trunkLen_, loopIdx, s2LoopNum); + continue; + } + SetFlag(V_MTE2_EVENT2); + WaitFlag(V_MTE2_EVENT2); + uint32_t validTrunkLen = (loopIdx * trunkLen_ + trunkLen_) > validS2Len ? validS2Len % trunkLen_ : trunkLen_; + uint32_t offset = vecOffset * QLICommon::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_) + loopIdx * trunkLen_; + AscendC::DataCopy(mrgValueLocal_, scoreOutLocal_, topkCountAlign256_); + // topk如果没有对齐到256，则把topkCountAlign256_ - topkCount_部分刷0 + if (topkCountAlign256_ != topkCount_) { + uint64_t mask[1]; + mask[0] = ~0; + mask[0] = mask[0] << (topkCount_ % 64); + PipeBarrier(); + // 把topkCount_对齐到64刷0，此处由于duplicate的限制mask[0]刷64个数 + Duplicate(mrgValueLocal_[topkCount_ / 64 * 64], zero, mask, 1, 1, 0); + PipeBarrier(); + // 把topk剩余对齐到256的部分刷0 + Duplicate(mrgValueLocal_[topkCount_ / 64 * 64 + 64], zero, topkCountAlign256_ - (topkCount_ / 64 * 64 + 64)); + SetFlag(V_MTE2_EVENT3); + WaitFlag(V_MTE2_EVENT3); + } + copyInParams.blockLen = validTrunkLen * sizeof(SCORE_T); // byte + // TOPK 直方图一次必须计算256，输入处理数据需要和256对齐 + if ((topkCountAlign256_ + validTrunkLen) % 256 != 0) { + Duplicate(mrgValueLocal_[topkCountAlign256_ + validTrunkLen / 256 * 256], zero, QLICommon::Align(validTrunkLen, (uint32_t)256) - validTrunkLen / 256 * 256); + SetFlag(V_MTE2_EVENT); + WaitFlag(V_MTE2_EVENT); + } + WaitFlag(V_MTE2_EVENT1); + AscendC::DataCopyPad(mrgValueLocal_[topkCountAlign256_], scoreGm[offset], copyInParams, padParams); + SetFlag(TOPK_MTE2_V_EVENT); + WaitFlag(TOPK_MTE2_V_EVENT); + topkOp_(mrgValueLocal_, indicesOutLocal_, scoreOutLocal_, QLICommon::Align(topkCountAlign256_ + validTrunkLen, (uint32_t)256), loopIdx, s2LoopNum); + SetFlag(V_MTE2_EVENT1); + } + } + } else { + AscendC::CreateVecIndex(indicesOutLocal_.ReinterpretCast(), (int32_t)zero, validS2Len); + } + + if (validS2Len < topkCount_) { + uint64_t mask[1]; + mask[0] = ~0; + mask[0] = mask[0] << (validS2Len % 8); + PipeBarrier(); + Duplicate(indicesOutLocal_.ReinterpretCast()[validS2Len / 8 * 8], neg, mask, 1, 1, 0); + } + + if (validS2Len / 8 * 8 + 64 < topkCount_) { + PipeBarrier(); + Duplicate(indicesOutLocal_.ReinterpretCast()[validS2Len / 8 * 8 + 64], neg, topkCount_ - (validS2Len / 8 * 8 + 64)); + } + + SetFlag(TOPK_V_MTE2_EVENT); + SetFlag(TOPK_V_MTE3_EVENT); + WaitFlag(TOPK_V_MTE3_EVENT); + AscendC::DataCopyPad(indiceOutGm[info.indiceOutOffset + (curS1Idx + rowIdx) * topkCount_], indicesOutLocal_.ReinterpretCast(), copyOutParams); + SetFlag(TOPK_MTE3_V_EVENT); + } +} +} // namespace QLIKernel +#endif diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/vf/quant_lightning_indexer_topk.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/vf/quant_lightning_indexer_topk.h new file mode 100644 index 000000000..e698b871f --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/vf/quant_lightning_indexer_topk.h @@ -0,0 +1,165 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_topk.h + * \brief + */ +#ifndef quant_lightning_indexer_TOPK_H +#define quant_lightning_indexer_TOPK_H + +#include "kernel_operator.h" +#include "vf_topk.h" +#include "vf_topk_16_gather.h" + +namespace topk { +template +class LITopk +{ +public: + __aicore__ inline void operator()(LocalTensor& outputIdxLocal, + LocalTensor& inputLocal, + uint32_t s2SeqLen) + { + } +}; + +template<> +class LITopk { +public: + static __aicore__ inline uint32_t GetSharedTmpBufferSize(uint32_t topK) + { + return 2 * topK * sizeof(uint32_t) + 5 * 256 * sizeof(uint32_t) + 64 * sizeof(uint32_t) + + (topK + 64) * sizeof(uint32_t); // for output value tensor + } + + static __aicore__ inline uint32_t GetIndexBufferSize(uint32_t topK) + { + return (topK + 64) * sizeof(uint32_t); + } + + __aicore__ inline void Init(uint32_t topK) + { + this->topK = topK; + } + + __aicore__ inline void InitBuffers(LocalTensor& sharedTmpBuffer) + { + tmpIdxLocal = sharedTmpBuffer[0]; + tmpValueLocal = tmpIdxLocal[topK]; + histogramsLocal = tmpValueLocal[topK]; + idx0Local = histogramsLocal[256]; + idx1Local = idx0Local[256]; + idx2Local = idx1Local[256]; + idx3Local = idx2Local[256]; + nkValueLocal = idx3Local[256]; + outputValueLocal = nkValueLocal[64]; + } + + __aicore__ inline void operator()(LocalTensor& outputIdxLocal, + LocalTensor& inputLocal, + uint32_t s2SeqLen) + { + topkb32::LiTopKVF(outputIdxLocal, // filter阶段使用输出value Buf topK * 4B + outputValueLocal, // filter阶段使用输出 Idx Buf topK * 4B + inputLocal, // 输入 s2SeqLen * 4B + tmpIdxLocal, // filter阶段使用暂存index Buf topK * 4B + tmpValueLocal, // filter阶段使用暂存value Buf topK * 4B + histogramsLocal, // 直方图的临时Buf 256 * 4B + idx0Local, // 输入数据第1个8位Buf 256 * 4B + idx1Local, // 输入数据第2个8位Buf 256 * 4B + idx2Local, // 输入数据第3个8位Buf 256 * 4B + idx3Local, // 输入数据第4个8位Buf 256 * 4B + nkValueLocal, // next_k 暂存Buf 64 * 4B + topK, // topk数量 + s2SeqLen); // 输入元素总数 + } +private: + LocalTensor tmpIdxLocal; // filter阶段使用暂存index Buf topK * 4B + LocalTensor tmpValueLocal; // filter阶段使用暂存value Buf topK * 4B + LocalTensor histogramsLocal; // 直方图的临时Buf 256 * 4B + LocalTensor idx0Local; // 输入数据第1个8位Buf 256 * 4B + LocalTensor idx1Local; // 输入数据第2个8位Buf 256 * 4B + LocalTensor idx2Local; // 输入数据第3个8位Buf 256 * 4B + LocalTensor idx3Local; // 输入数据第4个8位Buf 256 * 4B + LocalTensor nkValueLocal; // next_k 暂存Buf 64 * 4B + LocalTensor outputValueLocal; // 输出value tensor + uint32_t topK; +}; + +template<> +class LITopk { +public: + __aicore__ inline uint32_t GetSharedTmpBufferSize() + { + // 2 * QLICommon::Align(topK, (uint32_t)256):两块hisIndexLocal；3 * 256：histogramsLocal idxHighLocal idxLowLocal；64：nkValueLocal + uint64_t bufferSize1 = (2 * QLICommon::Align(topK, (uint32_t)256) + 3 * 256 + 64) * sizeof(uint32_t); + // QLICommon::Align(topK, (uint32_t)256) + trunkLen：tmpIndexLocal + uint64_t bufferSize2 = (QLICommon::Align(topK, (uint32_t)256) + trunkLen) * sizeof(uint16_t); + return bufferSize1 + bufferSize2; + } + + __aicore__ inline void Init(uint32_t topK, uint32_t trunkLen) + { + this->topK = topK; + this->trunkLen = trunkLen; + } + + __aicore__ inline void InitBuffers(LocalTensor& sharedTmpBuffer) + { + LocalTensor hisIndexLocal1 = sharedTmpBuffer[0]; + LocalTensor hisIndexLocal2 = hisIndexLocal1[QLICommon::Align(topK, (uint32_t)256)]; + hisIndexLocal[0] = hisIndexLocal1; + hisIndexLocal[1] = hisIndexLocal2; + histogramsLocal = hisIndexLocal2[QLICommon::Align(topK, (uint32_t)256)]; + idxHighLocal = histogramsLocal[256]; + idxLowLocal = idxHighLocal[256]; + nkValueLocal = idxLowLocal[256]; + LocalTensor tmpIndexLocalTmp = nkValueLocal[64]; + tmpIndexLocal = tmpIndexLocalTmp.template ReinterpretCast(); + } + + __aicore__ inline void operator()(LocalTensor& mrgValueLocal, LocalTensor& indicesOutLocal, + LocalTensor& hisValueLocal, uint32_t s2SeqLen, uint32_t loopIdx, uint32_t s2LoopNum) + { + if (s2LoopNum == 1) { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, mrgValueLocal, histogramsLocal, idxHighLocal, idxLowLocal, nkValueLocal, topK, s2SeqLen); + PipeBarrier(); + Cast(indicesOutLocal, tmpIndexLocal, RoundMode::CAST_NONE, topK); + return; + } + + if (loopIdx == 0) { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, mrgValueLocal, histogramsLocal, idxHighLocal, idxLowLocal, nkValueLocal, topK, s2SeqLen); + PipeBarrier(); + Cast(hisIndexLocal[(loopIdx + 1) % 2], tmpIndexLocal, RoundMode::CAST_NONE, topK); + } else { + topkb16gather::LiTopKVF(tmpIndexLocal, hisValueLocal, mrgValueLocal, histogramsLocal, idxHighLocal, idxLowLocal, nkValueLocal, topK, s2SeqLen); + PipeBarrier(); + topkb16gather::LiTopKGatherVF(hisIndexLocal[(loopIdx + 1) % 2], hisValueLocal, mrgValueLocal, tmpIndexLocal, hisIndexLocal[loopIdx % 2], + topK, loopIdx * trunkLen - QLICommon::Align(topK, (uint32_t)256), s2SeqLen); + if (loopIdx == s2LoopNum - 1) { + PipeBarrier(); + AscendC::DataCopy(indicesOutLocal, hisIndexLocal[(loopIdx + 1) % 2], QLICommon::Align(topK, (uint32_t)256)); + } + } + } +private: + LocalTensor hisIndexLocal[2]; // 每trunkLen长度的s2选出的topK个索引 + LocalTensor histogramsLocal; // 直方图的临时Buf 256 * 4B + LocalTensor idxHighLocal; // 输入数据高8位Buf 256 * 4B + LocalTensor idxLowLocal; // 输入数据低8位Buf 256 * 4B + LocalTensor nkValueLocal; // next_k 暂存Buf 64 * 4B + LocalTensor tmpIndexLocal; // 每trunkLen + topK的临时index + uint32_t topK = 512; + uint32_t trunkLen = 16384; +}; +} +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/vf/quant_lightning_indexer_vector1.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/vf/quant_lightning_indexer_vector1.h new file mode 100644 index 000000000..b03edde38 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/vf/quant_lightning_indexer_vector1.h @@ -0,0 +1,614 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file quant_lightning_indexer_vector1.h + * \brief + */ +#ifndef quant_lightning_indexer_VECTOR1_H +#define quant_lightning_indexer_VECTOR1_H + +#include "kernel_operator.h" + +namespace vector1 { + +template +struct FloatSortTraits; + +// fp32 +template <> +struct FloatSortTraits { + using UInt = uint32_t; + static constexpr UInt ZERO = 0x00000000; + static constexpr UInt SIGN_MASK = 0x80000000; + static constexpr UInt NAN_MASK = 0x7FC00000; + static constexpr UInt ALL_ONE = 0xFFFFFFFF; +}; + +// bf16 +template <> +struct FloatSortTraits { + using UInt = uint16_t; + static constexpr UInt ZERO = 0x0000; + static constexpr UInt SIGN_MASK = 0x8000; + static constexpr UInt NAN_MASK = 0x7FC0; + static constexpr UInt ALL_ONE = 0xFFFF; +}; + + +template +struct FloatSortConstCtx { + using Traits = FloatSortTraits; + using UInt = typename Traits::UInt; + AscendC::MicroAPI::RegTensor zeros; + AscendC::MicroAPI::RegTensor allOnes; + AscendC::MicroAPI::RegTensor signMask; + AscendC::MicroAPI::RegTensor nan; +}; + + +template +__simd_callee__ inline void InitFloatSortConstCtx(FloatSortConstCtx& ctx, AscendC::MicroAPI::MaskReg& maskAll) +{ + using Traits = FloatSortTraits; + AscendC::MicroAPI::Duplicate(ctx.zeros, Traits::ZERO, maskAll); + AscendC::MicroAPI::Duplicate(ctx.allOnes, Traits::ALL_ONE, maskAll); + AscendC::MicroAPI::Duplicate(ctx.signMask, Traits::SIGN_MASK, maskAll); + AscendC::MicroAPI::Duplicate(ctx.nan, Traits::NAN_MASK, maskAll); +} + + +template +__simd_callee__ inline void FloatToSortableKey(AscendC::MicroAPI::RegTensor::UInt>& outKey, + AscendC::MicroAPI::RegTensor& inVal, + FloatSortConstCtx& ctx, + AscendC::MicroAPI::MaskReg& maskAll) +{ + using Traits = FloatSortTraits; + using UInt = typename Traits::UInt; + + AscendC::MicroAPI::RegTensor regTemp; + AscendC::MicroAPI::RegTensor regMask; + AscendC::MicroAPI::MaskReg regSelectNan; + AscendC::MicroAPI::MaskReg regSelectSign; + + auto& inBits = (AscendC::MicroAPI::RegTensor&)inVal; + + // 1. NaN check + AscendC::MicroAPI::Compare(regSelectNan, inBits, ctx.nan, maskAll); + + // 2. NaN -> ALL_ONE + AscendC::MicroAPI::Select(outKey, ctx.allOnes, inBits, regSelectNan); + + // 3. sign bit + AscendC::MicroAPI::And(regTemp, outKey, ctx.signMask, maskAll); + + AscendC::MicroAPI::Compare(regSelectSign, regTemp, ctx.zeros, maskAll); + + // 4. xor mask + AscendC::MicroAPI::Select(regMask, ctx.allOnes, ctx.signMask, regSelectSign); + AscendC::MicroAPI::Xor(outKey, outKey, regMask, maskAll); +} + +template +__simd_callee__ inline void FloatX2ToSortableKey(AscendC::MicroAPI::RegTensor::UInt>& outKey0, + AscendC::MicroAPI::RegTensor::UInt>& outKey1, + AscendC::MicroAPI::RegTensor& inVal0, + AscendC::MicroAPI::RegTensor& inVal1, + FloatSortConstCtx& ctx, + AscendC::MicroAPI::MaskReg& maskAll) +{ + using Traits = FloatSortTraits; + using UInt = typename Traits::UInt; + + AscendC::MicroAPI::RegTensor regTemp[2]; + AscendC::MicroAPI::RegTensor regMask[2]; + AscendC::MicroAPI::MaskReg regSelectNan[2]; + AscendC::MicroAPI::MaskReg regSelectSign[2]; + + auto& inBits0 = (AscendC::MicroAPI::RegTensor&)inVal0; + auto& inBits1 = (AscendC::MicroAPI::RegTensor&)inVal1; + + // 1. NaN check + AscendC::MicroAPI::Compare(regSelectNan[0], inBits0, ctx.nan, maskAll); + AscendC::MicroAPI::Compare(regSelectNan[1], inBits1, ctx.nan, maskAll); + + // 2. NaN -> ALL_ONE + AscendC::MicroAPI::Select(outKey0, ctx.allOnes, inBits0, regSelectNan[0]); + AscendC::MicroAPI::Select(outKey1, ctx.allOnes, inBits1, regSelectNan[1]); + + // 3. sign bit + AscendC::MicroAPI::And(regTemp[0], outKey0, ctx.signMask, maskAll); + AscendC::MicroAPI::And(regTemp[1], outKey1, ctx.signMask, maskAll); + + AscendC::MicroAPI::Compare(regSelectSign[0], regTemp[0], ctx.zeros, maskAll); + AscendC::MicroAPI::Compare(regSelectSign[1], regTemp[1], ctx.zeros, maskAll); + + // 4. xor mask + AscendC::MicroAPI::Select(regMask[0], ctx.allOnes, ctx.signMask, regSelectSign[0]); + AscendC::MicroAPI::Select(regMask[1], ctx.allOnes, ctx.signMask, regSelectSign[1]); + AscendC::MicroAPI::Xor(outKey0, outKey0, regMask[0], maskAll); + AscendC::MicroAPI::Xor(outKey1, outKey1, regMask[1], maskAll); +} + + +template +__simd_callee__ inline void DuplicateZero(AscendC::MicroAPI::RegTensor (®Array)[N], + AscendC::MicroAPI::MaskReg& mask) +{ + static_assert(N <= 4, "N must be <= 4"); + // 不能用循环, 会导致fatal error: error in backend: Unsupported Inst must be hoisted. + if constexpr (N >= 1) { + AscendC::MicroAPI::Duplicate(regArray[0], static_cast(0), mask); + } + if constexpr (N >= 2) { + AscendC::MicroAPI::Duplicate(regArray[1], static_cast(0), mask); + } + if constexpr (N >= 3) { + AscendC::MicroAPI::Duplicate(regArray[2], static_cast(0), mask); + } + if constexpr (N >= 4) { + AscendC::MicroAPI::Duplicate(regArray[3], static_cast(0), mask); + } +} + + +template +__simd_callee__ inline void WeightedAccum(AscendC::MicroAPI::RegTensor (&accum)[N], + AscendC::MicroAPI::RegTensor (&input)[N], + AscendC::MicroAPI::RegTensor& weight, + AscendC::MicroAPI::MaskReg& mask) +{ + static_assert(N <= 2, "N must be <= 2"); + // ---- Relu block ---- + if constexpr (ApplyRelu) { + if constexpr (N >= 1) { + AscendC::MicroAPI::Relu(input[0], input[0], mask); + } + if constexpr (N >= 2) { + AscendC::MicroAPI::Relu(input[1], input[1], mask); + } + } + // ---- MulAdd block ---- + if constexpr (N >= 1) { + AscendC::MicroAPI::MulAddDst(accum[0], input[0], weight, mask); + } + if constexpr (N >= 2) { + AscendC::MicroAPI::MulAddDst(accum[1], input[1], weight, mask); + } +} + + +__simd_callee__ inline void BroadcastLane(AscendC::MicroAPI::RegTensor& dst, + AscendC::MicroAPI::RegTensor& src, + uint16_t laneIdx) +{ + AscendC::MicroAPI::RegTensor brcGatherIndex; + AscendC::MicroAPI::Duplicate(brcGatherIndex, laneIdx); + AscendC::MicroAPI::Gather(dst, src, brcGatherIndex); +} + +__simd_callee__ inline void BroadcastLane(AscendC::MicroAPI::RegTensor& dst, + __local_mem__ float* src, + uint16_t laneIdx) +{ + AscendC::MicroAPI::LoadAlign(dst, src + laneIdx); +} + +// float in uint16 out +__aicore__ inline void MulWeightAndReduceSum(const LocalTensor &out_, // out [S2Base] [128 ] + const LocalTensor &qk_, // q*k^t [G, S2Base] [64 128] + const uint32_t qkVLStride, + const LocalTensor &weight_, // w [G] [64 ] + const LocalTensor &kScale_, // kScale [S2Base] [128 ] + const LocalTensor &qScale_, // qScale [G] [64 ] + const int gSize) // G 64 +{ + auto weight = (__local_mem__ float*)weight_.GetPhyAddr(); + auto qScale = (__local_mem__ float*)qScale_.GetPhyAddr(); + auto kScale = (__local_mem__ float*)kScale_.GetPhyAddr(); + auto qk = (__local_mem__ float*)qk_.GetPhyAddr(); + auto out = (__local_mem__ uint16_t*)out_.GetPhyAddr(); + + __VEC_SCOPE__ + { + AscendC::MicroAPI::RegTensor regwBrc; + AscendC::MicroAPI::RegTensor regQK[2]; + AscendC::MicroAPI::RegTensor regW; + + AscendC::MicroAPI::RegTensor regQScale; + AscendC::MicroAPI::RegTensor regKScale[2]; + AscendC::MicroAPI::RegTensor regSum0[2]; + AscendC::MicroAPI::RegTensor regSum1[2]; + AscendC::MicroAPI::MaskReg maskAllB32 = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg maskAllB16 = AscendC::MicroAPI::CreateMask(); + + FloatSortConstCtx bf16Ctx; + InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_ROUND}; + + AscendC::MicroAPI::LoadAlign(regW, weight); + AscendC::MicroAPI::LoadAlign(regQScale, qScale); + AscendC::MicroAPI::Mul(regW, regW, regQScale, maskAllB32); + + DuplicateZero(regSum0, maskAllB32); + DuplicateZero(regSum1, maskAllB32); + + MicroAPI::LoadAlign(regKScale[0], kScale); + MicroAPI::LoadAlign(regKScale[1], kScale + 64); + + // unroll2 + for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i += 2) { + MicroAPI::LoadAlign(regQK[0], qk + 128 * i); // RowStride是128, 行都落在一个bank上 + MicroAPI::LoadAlign(regQK[1], qk + 128 * i + qkVLStride); + BroadcastLane(regwBrc, regW, i); + WeightedAccum(regSum0, regQK, regwBrc, maskAllB32); + + MicroAPI::LoadAlign(regQK[0], qk + 128 * i + 128); + MicroAPI::LoadAlign(regQK[1], qk + 128 * i + 128 + qkVLStride); + BroadcastLane(regwBrc, regW, i + 1); + WeightedAccum(regSum1, regQK, regwBrc, maskAllB32); + } + + AscendC::MicroAPI::Add(regSum0[0], regSum0[0], regSum1[0], maskAllB32); + AscendC::MicroAPI::Add(regSum0[1], regSum0[1], regSum1[1], maskAllB32); + + AscendC::MicroAPI::Mul(regSum0[0], regSum0[0], regKScale[0], maskAllB32); + AscendC::MicroAPI::Mul(regSum0[1], regSum0[1], regKScale[1], maskAllB32); + + AscendC::MicroAPI::RegTensor regSumBF16; + // interleave cast ==> regSum[1] high regSum[0] low + AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + AscendC::MicroAPI::Cast(regSumBF16, regSum0[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16, regSum0[0], maskAllB32); + + AscendC::MicroAPI::RegTensor regOut; + FloatToSortableKey(regOut, regSumBF16, bf16Ctx, maskAllB16); + // normal store + AscendC::MicroAPI::StoreAlign(out, regOut, maskAllB16); + } +} + + +// bfloat16_t in uint16 out +__aicore__ inline void MulWeightAndReduceSum(const LocalTensor &out_, // out [S2Base] [128 ] + const LocalTensor &qk_, // q*k^t [G, S2Base] [64 128] + const uint32_t qkVLStride, // unused for bfloat16 + const LocalTensor &weight_, // w [G] [64 ] + const LocalTensor &kScale_, // kScale [S2Base] [128 ] + const LocalTensor &qScale_, // qScale [G] [64 ] + const int gSize) // G 64 +{ + auto weight = (__local_mem__ float*)weight_.GetPhyAddr(); + auto qScale = (__local_mem__ float*)qScale_.GetPhyAddr(); + auto qk = (__local_mem__ bfloat16_t*)qk_.GetPhyAddr(); + auto kScale = (__local_mem__ float*)kScale_.GetPhyAddr(); + auto out = (__local_mem__ uint16_t*)out_.GetPhyAddr(); + + __VEC_SCOPE__ + { + AscendC::MicroAPI::RegTensor regQK[4]; + AscendC::MicroAPI::RegTensor regQKB16[2]; + AscendC::MicroAPI::RegTensor regW; + AscendC::MicroAPI::RegTensor regwBrc[2]; + AscendC::MicroAPI::RegTensor regQScale; + AscendC::MicroAPI::RegTensor regKScale[2]; + AscendC::MicroAPI::RegTensor regSum[2]; + + AscendC::MicroAPI::MaskReg maskAllB32 = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg maskAllB16 = AscendC::MicroAPI::CreateMask(); + + AscendC::MicroAPI::RegTensor regSumBF16; + + FloatSortConstCtx bf16Ctx; + InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + + using CastTrait = AscendC::MicroAPI::CastTrait; + static constexpr CastTrait castTraitB162B32_EVEN = {AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN, + AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + static constexpr CastTrait castTraitB162B32_ODD = {AscendC::MicroAPI::RegLayout::ONE, AscendC::MicroAPI::SatMode::UNKNOWN, + AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + constexpr static CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, RoundMode::CAST_ROUND}; + constexpr static CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_ROUND}; + + AscendC::MicroAPI::LoadAlign(regW, weight); + AscendC::MicroAPI::LoadAlign(regQScale, qScale); + AscendC::MicroAPI::Mul(regW, regW, regQScale, maskAllB32); + AscendC::MicroAPI::StoreAlign(weight, regW, maskAllB32); + AscendC::MicroAPI::LocalMemBar(); + + DuplicateZero(regSum, maskAllB32); + + // interleave load + MicroAPI::LoadAlign(regKScale[0], regKScale[1], kScale); + + // Duplicate + Gather方法劣化 + // Relu在cube随路做 + for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i++) { + AscendC::MicroAPI::LoadAlign(regQKB16[0], qk + 256 * i); // RowStride是256, 行都落在一个bank上 + AscendC::MicroAPI::LoadAlign(regwBrc[0], weight + i); + // interleave cast + AscendC::MicroAPI::Cast(regQK[0], regQKB16[0], maskAllB16); + AscendC::MicroAPI::Cast(regQK[1], regQKB16[0], maskAllB16); + AscendC::MicroAPI::MulAddDst(regSum[0], regQK[0], regwBrc[0], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum[1], regQK[1], regwBrc[0], maskAllB32); + } + + AscendC::MicroAPI::Mul(regSum[0], regSum[0], regKScale[0], maskAllB32); + AscendC::MicroAPI::Mul(regSum[1], regSum[1], regKScale[1], maskAllB32); + // interleave cast back + AscendC::MicroAPI::Cast(regSumBF16, regSum[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16, regSum[0], maskAllB32); + + AscendC::MicroAPI::RegTensor regOut; + FloatToSortableKey(regOut, regSumBF16, bf16Ctx, maskAllB16); + // norm load + AscendC::MicroAPI::StoreAlign(out, regOut, maskAllB16); + } +} + + +// 计算S1=2 +// float in uint16 out +__aicore__ inline void MulWeightAndReduceSum2(const LocalTensor &out_, // out [2, S2Base] [128 ] + uint32_t outStride, + const LocalTensor &qk_, // q*k^t [2, G, S2Base] [64 128] + uint32_t qkVLStride, + uint32_t qkStride, + const LocalTensor &weight_, // w [2, G] [64 ] + uint32_t weightStride, + const LocalTensor &kScale_, // kScale [S2Base] [128 ] + uint32_t kScaleStride, + const LocalTensor &qScale_, // qScale [2, G] [64 ] + uint32_t qScaleStride, + const int gSize) // G 64 +{ + auto weight0 = (__local_mem__ float*)weight_.GetPhyAddr(); + auto qScale0 = (__local_mem__ float*)qScale_.GetPhyAddr(); + auto kScale0 = (__local_mem__ float*)kScale_.GetPhyAddr(); + auto qk0 = (__local_mem__ float*)qk_.GetPhyAddr(); + auto out0 = (__local_mem__ uint16_t*)out_.GetPhyAddr(); + + auto weight1 = weight0 + weightStride; + auto qScale1 = qScale0 + qScaleStride; + auto qk1 = qk0 + qkStride; + // kScaleStride is zero + auto out1 = out0 + outStride; + + __VEC_SCOPE__ + { + AscendC::MicroAPI::RegTensor regwBrc[2]; + AscendC::MicroAPI::RegTensor regQK0[2]; + AscendC::MicroAPI::RegTensor regQK1[2]; + AscendC::MicroAPI::RegTensor regW[2]; + + AscendC::MicroAPI::RegTensor regQScale[2]; + AscendC::MicroAPI::RegTensor regKScale[2]; + AscendC::MicroAPI::RegTensor regSum0[2]; + AscendC::MicroAPI::RegTensor regSum1[2]; + AscendC::MicroAPI::MaskReg maskAllB32 = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg maskAllB16 = AscendC::MicroAPI::CreateMask(); + + FloatSortConstCtx bf16Ctx; + InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_ROUND}; + + AscendC::MicroAPI::LoadAlign(regW[0], weight0); + AscendC::MicroAPI::LoadAlign(regW[1], weight1); + AscendC::MicroAPI::LoadAlign(regQScale[0], qScale0); + AscendC::MicroAPI::LoadAlign(regQScale[1], qScale1); + AscendC::MicroAPI::Mul(regW[0], regW[0], regQScale[0], maskAllB32); + AscendC::MicroAPI::Mul(regW[1], regW[1], regQScale[1], maskAllB32); + // regW[0]与weight1混合使用 + AscendC::MicroAPI::StoreAlign(weight1, regW[1], maskAllB32); + AscendC::MicroAPI::LocalMemBar(); + DuplicateZero(regSum0, maskAllB32); + DuplicateZero(regSum1, maskAllB32); + + MicroAPI::LoadAlign(regKScale[0], kScale0); + MicroAPI::LoadAlign(regKScale[1], kScale0 + 64); + + for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i++) { + MicroAPI::LoadAlign(regQK0[0], qk0 + 128 * i); + MicroAPI::LoadAlign(regQK0[1], qk0 + 128 * i + qkVLStride); + MicroAPI::LoadAlign(regQK1[0], qk1 + 128 * i); + MicroAPI::LoadAlign(regQK1[1], qk1 + 128 * i + qkVLStride); + // 混合使用对整体性能更好 + BroadcastLane(regwBrc[0], regW[0], i); + // Weight无bank冲突，用LoadAlign来提取weight标量 + BroadcastLane(regwBrc[1], weight1, i); + AscendC::MicroAPI::Relu(regQK0[0], regQK0[0], maskAllB32); + AscendC::MicroAPI::Relu(regQK0[1], regQK0[1], maskAllB32); + AscendC::MicroAPI::Relu(regQK1[0], regQK1[0], maskAllB32); + AscendC::MicroAPI::Relu(regQK1[1], regQK1[1], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum0[0], regQK0[0], regwBrc[0], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum0[1], regQK0[1], regwBrc[0], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum1[0], regQK1[0], regwBrc[1], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum1[1], regQK1[1], regwBrc[1], maskAllB32); + } + + // Apply kScale scaling + AscendC::MicroAPI::Mul(regSum0[0], regSum0[0], regKScale[0], maskAllB32); + AscendC::MicroAPI::Mul(regSum0[1], regSum0[1], regKScale[1], maskAllB32); + AscendC::MicroAPI::Mul(regSum1[0], regSum1[0], regKScale[0], maskAllB32); + AscendC::MicroAPI::Mul(regSum1[1], regSum1[1], regKScale[1], maskAllB32); + + + // Convert to bfloat16 and store output channel + AscendC::MicroAPI::RegTensor regSumBF16[2]; + AscendC::MicroAPI::RegTensor regOut[2]; + AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]); + AscendC::MicroAPI::DeInterleave(regSum1[0], regSum1[1], regSum1[0], regSum1[1]); + AscendC::MicroAPI::Cast(regSumBF16[0], regSum0[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[1], regSum1[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[0], regSum0[0], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[1], regSum1[0], maskAllB32); + + FloatX2ToSortableKey(regOut[0], regOut[1], regSumBF16[0], regSumBF16[1], bf16Ctx, maskAllB16); + AscendC::MicroAPI::StoreAlign(out0, regOut[0], maskAllB16); + AscendC::MicroAPI::StoreAlign(out1, regOut[1], maskAllB16); + } +} + + +// 计算S1=2 +// bfloat16 in uint16 out +__aicore__ inline void MulWeightAndReduceSum2(const LocalTensor &out_, // out [2, S2Base] [128 ] + uint32_t outStride, + const LocalTensor &qk_, // q*k^t [2, G, S2Base] [64 128] + uint32_t qkVLStride, + uint32_t qkStride, // gSize * 256 + const LocalTensor &weight_, // w [2, G] [64 ] + uint32_t weightStride, + const LocalTensor &kScale_, // kScale [S2Base] [128 ] + uint32_t kScaleStride, + const LocalTensor &qScale_, // qScale [2, G] [64 ] + uint32_t qScaleStride, + const int gSize) // G 64 +{ + auto weight0 = (__local_mem__ float*)weight_.GetPhyAddr(); + auto qScale0 = (__local_mem__ float*)qScale_.GetPhyAddr(); + auto kScale0 = (__local_mem__ float*)kScale_.GetPhyAddr(); + auto qk0 = (__local_mem__ bfloat16_t*)qk_.GetPhyAddr(); + auto out0 = (__local_mem__ uint16_t*)out_.GetPhyAddr(); + + auto weight1 = weight0 + weightStride; + auto qScale1 = qScale0 + qScaleStride; + auto qk1 = qk0 + qkStride; + // kScaleStride is zero + auto out1 = out0 + outStride; + + __VEC_SCOPE__ + { + AscendC::MicroAPI::RegTensor regwBrc[2]; + AscendC::MicroAPI::RegTensor regQK0[2]; + AscendC::MicroAPI::RegTensor regQK1[2]; + AscendC::MicroAPI::RegTensor regW[2]; + AscendC::MicroAPI::RegTensor regQKB16[2]; + + AscendC::MicroAPI::RegTensor regQScale[2]; + AscendC::MicroAPI::RegTensor regKScale[2]; + AscendC::MicroAPI::RegTensor regSum0[2]; + AscendC::MicroAPI::RegTensor regSum1[2]; + AscendC::MicroAPI::MaskReg maskAllB32 = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg maskAllB16 = AscendC::MicroAPI::CreateMask(); + + FloatSortConstCtx bf16Ctx; + InitFloatSortConstCtx(bf16Ctx, maskAllB16); + + using CastTrait = AscendC::MicroAPI::CastTrait; + static constexpr CastTrait castTraitB162B32_EVEN = {AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN, + AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + static constexpr CastTrait castTraitB162B32_ODD = {AscendC::MicroAPI::RegLayout::ONE, AscendC::MicroAPI::SatMode::UNKNOWN, + AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::MERGING, RoundMode::CAST_ROUND}; + constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_ROUND}; + + AscendC::MicroAPI::LoadAlign(regW[0], weight0); + AscendC::MicroAPI::LoadAlign(regW[1], weight1); + AscendC::MicroAPI::LoadAlign(regQScale[0], qScale0); + AscendC::MicroAPI::LoadAlign(regQScale[1], qScale1); + AscendC::MicroAPI::Mul(regW[0], regW[0], regQScale[0], maskAllB32); + AscendC::MicroAPI::Mul(regW[1], regW[1], regQScale[1], maskAllB32); + // 读写依赖，寄存器可以保序 + AscendC::MicroAPI::StoreAlign(weight0, regW[0], maskAllB32); + AscendC::MicroAPI::StoreAlign(weight1, regW[1], maskAllB32); + DuplicateZero(regSum0, maskAllB32); + DuplicateZero(regSum1, maskAllB32); + + // interleave load + MicroAPI::LoadAlign(regKScale[0], regKScale[1], kScale0); + + for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i++) { + AscendC::MicroAPI::LoadAlign(regQKB16[0], qk0 + 256 * i); // RowStride是256, 行都落在一个bank上 + AscendC::MicroAPI::LoadAlign(regQKB16[1], qk1 + 256 * i); // RowStride是256, 行都落在一个bank上 + AscendC::MicroAPI::LoadAlign(regwBrc[0], weight0 + i); + AscendC::MicroAPI::LoadAlign(regwBrc[1], weight1 + i); + // interleave cast + AscendC::MicroAPI::Cast(regQK0[0], regQKB16[0], maskAllB32); + AscendC::MicroAPI::Cast(regQK0[1], regQKB16[0], maskAllB32); + AscendC::MicroAPI::Cast(regQK1[0], regQKB16[1], maskAllB32); + AscendC::MicroAPI::Cast(regQK1[1], regQKB16[1], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum0[0], regQK0[0], regwBrc[0], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum0[1], regQK0[1], regwBrc[0], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum1[0], regQK1[0], regwBrc[1], maskAllB32); + AscendC::MicroAPI::MulAddDst(regSum1[1], regQK1[1], regwBrc[1], maskAllB32); + } + + // Apply kScale scaling + AscendC::MicroAPI::Mul(regSum0[0], regSum0[0], regKScale[0], maskAllB32); + AscendC::MicroAPI::Mul(regSum0[1], regSum0[1], regKScale[1], maskAllB32); + AscendC::MicroAPI::Mul(regSum1[0], regSum1[0], regKScale[0], maskAllB32); + AscendC::MicroAPI::Mul(regSum1[1], regSum1[1], regKScale[1], maskAllB32); + + // Convert to bfloat16 and store output channel + AscendC::MicroAPI::RegTensor regSumBF16[2]; + AscendC::MicroAPI::RegTensor regOut[2]; + AscendC::MicroAPI::Cast(regSumBF16[0], regSum0[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[1], regSum1[1], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[0], regSum0[0], maskAllB32); + AscendC::MicroAPI::Cast(regSumBF16[1], regSum1[0], maskAllB32); + + FloatX2ToSortableKey(regOut[0], regOut[1], regSumBF16[0], regSumBF16[1], bf16Ctx, maskAllB16); + AscendC::MicroAPI::StoreAlign(out0, regOut[0], maskAllB16); + AscendC::MicroAPI::StoreAlign(out1, regOut[1], maskAllB16); + } +} + + +template +__aicore__ inline void BatchMulWeightAndReduceSum(const LocalTensor &out_, // out [S2Base] [128 ] + uint32_t outStride, + const LocalTensor &qk_, // q*k^t [G, S2Base] [64 128] + uint32_t qkVLStride, + uint32_t qkStride, + const LocalTensor &weight_, // w [G] [64 ] + uint32_t weightStride, + const LocalTensor &kScale_, // kScale [S2Base] [128 ] + uint32_t kScaleStride, + const LocalTensor &qScale_, // qScale [G] [64 ] + uint32_t qScaleStride, + const int gSize, // G 64 + const int batch) +{ + // 暂只支持这两种情况, 后续改成循环 + if (batch != 2 && batch != 1) { + return; + } + if (batch == 2) { + MulWeightAndReduceSum2(out_, outStride, + qk_, qkVLStride, qkStride, + weight_, weightStride, + kScale_, kScaleStride, + qScale_, qScaleStride, + gSize); + } else { + MulWeightAndReduceSum(out_, qk_, qkVLStride, weight_, kScale_, qScale_, gSize); + } +} + +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/vf/vf_topk.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/vf/vf_topk.h new file mode 100644 index 000000000..7834ee922 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/vf/vf_topk.h @@ -0,0 +1,678 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file vf_top_k.h +* \brief +*/ + +#ifndef VF_TOP_K_H +#define VF_TOP_K_H + +namespace topkb32 { +template +__simd_vf__ void HistogramsFirstVFImpl(__ubuf__ uint32_t* histogramsBuf, __ubuf__ uint32_t* inputBuf, uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图cout0 0-127 cout1 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + // 32bit 高16bit + MicroAPI::RegTensor vreg0U16; + // 32bit 低16bit + MicroAPI::RegTensor vreg1U16; + MicroAPI::RegTensor vreg2U16; + MicroAPI::RegTensor vreg3U16; + + MicroAPI::RegTensor vreg0; + MicroAPI::RegTensor vreg1; + MicroAPI::RegTensor vreg2; + MicroAPI::RegTensor vreg3; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vreg1U16, vreg0U16, inputBuf + i * 256); + MicroAPI::LoadAlign(vreg3U16, vreg2U16, inputBuf + (i * 256) + 128); + + MicroAPI::DeInterleave(vreg1, vreg0, (MicroAPI::RegTensor&)vreg0U16, (MicroAPI::RegTensor&)vreg2U16); + + MicroAPI::Histograms(cout0, vreg0, pregB8); + MicroAPI::Histograms(cout1, vreg0, pregB8); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, cout0U32Even, cout0U32Odd, pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, cout1U32Even, cout1U32Odd, pregB32); +} + +__simd_vf__ void FindFirstTargetBinVFImpl(__ubuf__ uint32_t* idx0Buf, __ubuf__ uint32_t* nkValueBuf, __ubuf__ uint32_t* histogramsBuf, uint32_t bottomK) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx0; + + MicroAPI::RegTensor btmK; + MicroAPI::Duplicate(btmK, bottomK); + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdx0; + + MicroAPI::MaskReg pregGE = MicroAPI::CreateMask(); + + MicroAPI::Arange(idxC, i * 64); + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + MicroAPI::Compare(pregGE, cout, btmK, pregB32); + MicroAPI::Squeeze(sqzIdx0, (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idx0Buf, sqzIdx0, alignIdx0); + } + MicroAPI::StoreUnAlignPost(idx0Buf, alignIdx0); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idx0; + MicroAPI::LoadAlign(idx0, idx0Buf); + + MicroAPI::RegTensor idxAll1; + MicroAPI::RegTensor idxPrev0; + MicroAPI::RegTensor prevBinValue; + MicroAPI::Duplicate(idxAll1, 1); + + MicroAPI::RegTensor zeroAll; + MicroAPI::Duplicate(zeroAll, 0); + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + MicroAPI::Compare(preg0, idx0, zeroAll, pregB32); + MicroAPI::Sub(idxPrev0, idx0, (MicroAPI::RegTensor&)idxAll1, pregB32); + MicroAPI::ShiftRights(idxPrev0, idxPrev0, (int16_t)24, pregB32); + + MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev0, pregB32); + MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, preg0); + + MicroAPI::RegTensor nextK; + MicroAPI::Sub(nextK, btmK, prevBinValue, pregB32); + MicroAPI::StoreAlign(nkValueBuf, nextK, pregB32); +} + +template +__simd_vf__ void HistogramsSecondVFImpl(__ubuf__ uint32_t* histogramsBuf, __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* idx0Buf, uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图0-127 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + MicroAPI::RegTensor idx0; + // 0x000000fc -> 0xfcfcfcfc + MicroAPI::LoadAlign(idx0, idx0Buf); + + MicroAPI::RegTensor vreg0U16; + MicroAPI::RegTensor vreg1U16; + MicroAPI::RegTensor vreg2U16; + MicroAPI::RegTensor vreg3U16; + + MicroAPI::RegTensor vreg0; + MicroAPI::RegTensor vreg1; + MicroAPI::RegTensor vreg2; + MicroAPI::RegTensor vreg3; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vreg1U16, vreg0U16, inputBuf + i * 256); + MicroAPI::LoadAlign(vreg3U16, vreg2U16, inputBuf + (i * 256) + 128); + + MicroAPI::DeInterleave(vreg1, vreg0, (MicroAPI::RegTensor&)vreg0U16, (MicroAPI::RegTensor&)vreg2U16); + + MicroAPI::MaskReg pregEQ = MicroAPI::CreateMask(); + MicroAPI::Compare(pregEQ, vreg0, (MicroAPI::RegTensor&)idx0, pregB8); + + MicroAPI::Histograms(cout0, vreg1, pregEQ); + MicroAPI::Histograms(cout1, vreg1, pregEQ); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, cout0U32Even, cout0U32Odd, pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, cout1U32Even, cout1U32Odd, pregB32); +} + +// kValue新的bottomK +__simd_vf__ void FindSecondTargetBinVFImpl(__ubuf__ uint32_t* idx1Buf, __ubuf__ uint32_t* nkValueBuf, __ubuf__ uint32_t* kValue, __ubuf__ uint32_t* histogramsBuf) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx1; + + MicroAPI::RegTensor btmK1; + MicroAPI::LoadAlign(btmK1, kValue); + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdx1; + + MicroAPI::MaskReg pregGE = MicroAPI::CreateMask(); + + MicroAPI::Arange(idxC, i * 64); + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + MicroAPI::Compare(pregGE, cout, btmK1, pregB32); + MicroAPI::Squeeze(sqzIdx1, (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idx1Buf, sqzIdx1, alignIdx1); + } + MicroAPI::StoreUnAlignPost(idx1Buf, alignIdx1); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idx1; + MicroAPI::LoadAlign(idx1, idx1Buf); + + MicroAPI::RegTensor idxAll1; + MicroAPI::RegTensor idxPrev1; + MicroAPI::RegTensor prevBinValue; + MicroAPI::Duplicate(idxAll1, 1); + + MicroAPI::RegTensor zeroAll; + MicroAPI::Duplicate(zeroAll, 0); + + MicroAPI::MaskReg preg1 = MicroAPI::CreateMask(); + MicroAPI::Compare(preg1, idx1, zeroAll, pregB32); + MicroAPI::Sub(idxPrev1, idx1, (MicroAPI::RegTensor&)idxAll1, pregB32); + MicroAPI::ShiftRights(idxPrev1, idxPrev1, (int16_t)24, pregB32); + + MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev1, pregB32); + MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, preg1); + + MicroAPI::RegTensor nextK; + MicroAPI::Sub(nextK, btmK1, prevBinValue, pregB32); + MicroAPI::StoreAlign(nkValueBuf, nextK, pregB32); +} + +template +__simd_vf__ void HistogramsThirdVFImpl(__ubuf__ uint32_t* histogramsBuf, __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* idx0Buf, __ubuf__ uint32_t* idx1Buf, uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图0-127 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + MicroAPI::RegTensor idx0; + MicroAPI::RegTensor idx1; + // 0x000000fc -> 0xfcfcfcfc + MicroAPI::LoadAlign(idx0, idx0Buf); + MicroAPI::LoadAlign(idx1, idx1Buf); + + MicroAPI::RegTensor vreg0U16; + MicroAPI::RegTensor vreg1U16; + MicroAPI::RegTensor vreg2U16; + MicroAPI::RegTensor vreg3U16; + + MicroAPI::RegTensor vreg0; + MicroAPI::RegTensor vreg1; + MicroAPI::RegTensor vreg2; + MicroAPI::RegTensor vreg3; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vreg1U16, vreg0U16, inputBuf + i * 256); + MicroAPI::LoadAlign(vreg3U16, vreg2U16, inputBuf + (i * 256) + 128); + + MicroAPI::DeInterleave(vreg1, vreg0, (MicroAPI::RegTensor&)vreg0U16, (MicroAPI::RegTensor&)vreg2U16); + MicroAPI::DeInterleave(vreg3, vreg2, (MicroAPI::RegTensor&)vreg1U16, (MicroAPI::RegTensor&)vreg3U16); + + MicroAPI::MaskReg pregEQ0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregEQ1 = MicroAPI::CreateMask(); + MicroAPI::Compare(pregEQ0, vreg0, (MicroAPI::RegTensor&)idx0, pregB8); + MicroAPI::Compare(pregEQ1, vreg1, (MicroAPI::RegTensor&)idx1, pregB8); + + MicroAPI::MaskReg pregEQ = MicroAPI::CreateMask(); + MicroAPI::And(pregEQ, pregEQ0, pregEQ1, pregB8); + + MicroAPI::Histograms(cout0, vreg2, pregEQ); + MicroAPI::Histograms(cout1, vreg2, pregEQ); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, cout0U32Even, cout0U32Odd, pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, cout1U32Even, cout1U32Odd, pregB32); +} + +__simd_vf__ void FindThirdTargetBinVFImpl(__ubuf__ uint32_t* idx2Buf, __ubuf__ uint32_t* nkValueBuf, __ubuf__ uint32_t* kValue, __ubuf__ uint32_t* histogramsBuf) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx2; + + MicroAPI::RegTensor btmK2; + MicroAPI::LoadAlign(btmK2, kValue); + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdx2; + + MicroAPI::MaskReg pregGE = MicroAPI::CreateMask(); + + MicroAPI::Arange(idxC, i * 64); + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + MicroAPI::Compare(pregGE, cout, btmK2, pregB32); + MicroAPI::Squeeze(sqzIdx2, (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idx2Buf, sqzIdx2, alignIdx2); + } + MicroAPI::StoreUnAlignPost(idx2Buf, alignIdx2); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idx2; + MicroAPI::LoadAlign(idx2, idx2Buf); + + MicroAPI::RegTensor idxAll1; + MicroAPI::RegTensor idxPrev2; + MicroAPI::RegTensor prevBinValue; + MicroAPI::Duplicate(idxAll1, 1); + + MicroAPI::RegTensor zeroAll; + MicroAPI::Duplicate(zeroAll, 0); + + MicroAPI::MaskReg preg2 = MicroAPI::CreateMask(); + MicroAPI::Compare(preg2, idx2, zeroAll, pregB32); + MicroAPI::Sub(idxPrev2, idx2, (MicroAPI::RegTensor&)idxAll1, pregB32); + MicroAPI::ShiftRights(idxPrev2, idxPrev2, (int16_t)24, pregB32); + + MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev2, pregB32); + MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, preg2); + + MicroAPI::RegTensor nextK; + MicroAPI::Sub(nextK, btmK2, prevBinValue, pregB32); + MicroAPI::StoreAlign(nkValueBuf, nextK, pregB32); +} + +template +__simd_vf__ void HistogramsLastVFImpl(__ubuf__ uint32_t* histogramsBuf, __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* idx0Buf, __ubuf__ uint32_t* idx1Buf, __ubuf__ uint32_t* idx2Buf, uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图0-127 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + MicroAPI::RegTensor idx0; + MicroAPI::RegTensor idx1; + MicroAPI::RegTensor idx2; + // 0x000000fc -> 0xfcfcfcfc + MicroAPI::LoadAlign(idx0, idx0Buf); + MicroAPI::LoadAlign(idx1, idx1Buf); + MicroAPI::LoadAlign(idx2, idx2Buf); + + MicroAPI::RegTensor vreg0U16; + MicroAPI::RegTensor vreg1U16; + MicroAPI::RegTensor vreg2U16; + MicroAPI::RegTensor vreg3U16; + + MicroAPI::RegTensor vreg0; + MicroAPI::RegTensor vreg1; + MicroAPI::RegTensor vreg2; + MicroAPI::RegTensor vreg3; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vreg1U16, vreg0U16, inputBuf + i * 256); + MicroAPI::LoadAlign(vreg3U16, vreg2U16, inputBuf + (i * 256) + 128); + + MicroAPI::DeInterleave(vreg1, vreg0, (MicroAPI::RegTensor&)vreg0U16, (MicroAPI::RegTensor&)vreg2U16); + MicroAPI::DeInterleave(vreg3, vreg2, (MicroAPI::RegTensor&)vreg1U16, (MicroAPI::RegTensor&)vreg3U16); + + MicroAPI::MaskReg pregEQ0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregEQ1 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregEQ2 = MicroAPI::CreateMask(); + MicroAPI::Compare(pregEQ0, vreg0, (MicroAPI::RegTensor&)idx0, pregB8); + MicroAPI::Compare(pregEQ1, vreg1, (MicroAPI::RegTensor&)idx1, pregB8); + MicroAPI::Compare(pregEQ2, vreg2, (MicroAPI::RegTensor&)idx2, pregB8); + + MicroAPI::MaskReg pregEQ0And1 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregEQAll = MicroAPI::CreateMask(); + MicroAPI::And(pregEQ0And1, pregEQ0, pregEQ1, pregB8); + MicroAPI::And(pregEQAll, pregEQ0And1, pregEQ2, pregB8); + + MicroAPI::Histograms(cout0, vreg3, pregEQAll); + MicroAPI::Histograms(cout1, vreg3, pregEQAll); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, cout0U32Even, cout0U32Odd, pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, cout1U32Even, cout1U32Odd, pregB32); +} + +__simd_vf__ void FindKthVFImpl(__ubuf__ uint32_t* kValue, __ubuf__ uint32_t* histogramsBuf, __ubuf__ uint32_t* idx0Buf, __ubuf__ uint32_t* idx1Buf, __ubuf__ uint32_t* idx2Buf, __ubuf__ uint32_t* idx3Buf) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx3; + + MicroAPI::RegTensor btmK3; + MicroAPI::LoadAlign(btmK3, kValue); + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdx3; + + MicroAPI::MaskReg pregGE = MicroAPI::CreateMask(); + + MicroAPI::Arange(idxC, i * 64); + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + MicroAPI::Compare(pregGE, cout, btmK3, pregB32); + MicroAPI::Squeeze(sqzIdx3, (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idx3Buf, sqzIdx3, alignIdx3); + } + MicroAPI::StoreUnAlignPost(idx3Buf, alignIdx3); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idx0; + MicroAPI::RegTensor idx1; + MicroAPI::RegTensor idx2; + MicroAPI::RegTensor idx3; + MicroAPI::LoadAlign(idx0, idx0Buf); + MicroAPI::LoadAlign(idx1, idx1Buf); + MicroAPI::LoadAlign(idx2, idx2Buf); + MicroAPI::LoadAlign(idx3, idx3Buf); + + MicroAPI::ShiftLefts(idx0, idx0, (int16_t)24, pregB32); + MicroAPI::ShiftLefts(idx1, idx1, (int16_t)16, pregB32); + MicroAPI::ShiftLefts(idx2, idx2, (int16_t)8, pregB32); + + // ADD + MicroAPI::Add(idx0, idx0, idx1, pregB32); + MicroAPI::Add(idx0, idx0, idx2, pregB32); + MicroAPI::Add(idx0, idx0, idx3, pregB32); + + MicroAPI::StoreAlign(kValue, idx0, pregB32); +} + +__simd_vf__ void FindIdxGTOutputVFImpl(__ubuf__ uint32_t* outputIdxBuf, __ubuf__ uint32_t* inputBuf, uint32_t beginIdx, __ubuf__ uint32_t* kValue, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::RegTensor idxC; + MicroAPI::Arange(idxC, beginIdx + i * 64); + + MicroAPI::LoadAlign(vregInput, inputBuf + i * 64); + + MicroAPI::MaskReg poutGT = MicroAPI::CreateMask(); + + MicroAPI::RegTensor sqzIdxOut; + MicroAPI::Compare(poutGT, vregInput, kthValue, pregB32); + + MicroAPI::Squeeze(sqzIdxOut, (MicroAPI::RegTensor&)idxC, poutGT); + MicroAPI::StoreUnAlign(outputIdxBuf, sqzIdxOut, alignIdx); + } + MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx); +} + +__simd_vf__ void FindIdxEQOutputVFImpl(__ubuf__ uint32_t* outputIdxBuf, __ubuf__ uint32_t* inputBuf, uint32_t beginIdx, __ubuf__ uint32_t* kValue) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::UnalignRegForStore alignIdx; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + + MicroAPI::RegTensor idxC; + MicroAPI::Arange(idxC, beginIdx); + + MicroAPI::LoadAlign(vregInput, inputBuf); + + MicroAPI::MaskReg poutEQ = MicroAPI::CreateMask(); + + MicroAPI::RegTensor sqzIdxOut; + MicroAPI::Compare(poutEQ, vregInput, kthValue, pregB32); + + MicroAPI::Squeeze(sqzIdxOut, (MicroAPI::RegTensor&)idxC, poutEQ); + MicroAPI::StoreUnAlign(outputIdxBuf, sqzIdxOut, alignIdx); + MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx); +} + +__simd_vf__ void FindValueGTOutputVFImpl(__ubuf__ uint32_t* outputValueBuf, __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* kValue, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignValue; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::LoadAlign(vregInput, inputBuf + i * 64); + + MicroAPI::MaskReg poutGT = MicroAPI::CreateMask(); + + MicroAPI::RegTensor sqzValueOut; + MicroAPI::Compare(poutGT, vregInput, kthValue, pregB32); + + MicroAPI::Squeeze(sqzValueOut, vregInput, poutGT); + MicroAPI::StoreUnAlign(outputValueBuf, sqzValueOut, alignValue); + } + MicroAPI::StoreUnAlignPost(outputValueBuf, alignValue); +} + +__simd_vf__ void FindValueEQOutputVFImpl(__ubuf__ uint32_t* outputValueBuf, __ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* kValue) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::UnalignRegForStore alignValue; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + + MicroAPI::LoadAlign(vregInput, inputBuf); + + MicroAPI::MaskReg poutEQ = MicroAPI::CreateMask(); + + MicroAPI::RegTensor sqzValueOut; + MicroAPI::Compare(poutEQ, vregInput, kthValue, pregB32); + + MicroAPI::Squeeze(sqzValueOut, vregInput, poutEQ); + MicroAPI::StoreUnAlign(outputValueBuf, sqzValueOut, alignValue); + MicroAPI::StoreUnAlignPost(outputValueBuf, alignValue); +} + +__aicore__ inline void LiTopKVF(const LocalTensor& outputIdxLocal, + const LocalTensor& outputValueLocal, + const LocalTensor& inputLocal, + const LocalTensor& tmpIdxLocal, + const LocalTensor& tmpValueLocal, + const LocalTensor& histogramsLocal, + const LocalTensor& idx0Local, + const LocalTensor& idx1Local, + const LocalTensor& idx2Local, + const LocalTensor& idx3Local, + const LocalTensor& nkValueLocal, + uint32_t topK, + uint32_t s2SeqLen) +{ + __ubuf__ uint32_t* outputIdxBuf = (__ubuf__ uint32_t*)outputIdxLocal.GetPhyAddr(); + __ubuf__ uint32_t* outputValueBuf = (__ubuf__ uint32_t*)outputValueLocal.GetPhyAddr(); + __ubuf__ uint32_t* inputBuf = (__ubuf__ uint32_t*)inputLocal.GetPhyAddr(); + __ubuf__ uint32_t* tmpIdxBuf = (__ubuf__ uint32_t*)tmpIdxLocal.GetPhyAddr(); + __ubuf__ uint32_t* tmpValueBuf = (__ubuf__ uint32_t*)tmpValueLocal.GetPhyAddr(); + __ubuf__ uint32_t* histogramsBuf = (__ubuf__ uint32_t*)histogramsLocal.GetPhyAddr(); + __ubuf__ uint32_t* idx0Buf = (__ubuf__ uint32_t*)idx0Local.GetPhyAddr(); + __ubuf__ uint32_t* idx1Buf = (__ubuf__ uint32_t*)idx1Local.GetPhyAddr(); + __ubuf__ uint32_t* idx2Buf = (__ubuf__ uint32_t*)idx2Local.GetPhyAddr(); + __ubuf__ uint32_t* idx3Buf = (__ubuf__ uint32_t*)idx3Local.GetPhyAddr(); + __ubuf__ uint32_t* nkValueBuf = (__ubuf__ uint32_t*)nkValueLocal.GetPhyAddr(); + + uint32_t bottomK = s2SeqLen - topK + 1; + uint32_t beginIdx = 0; + bool flag = true; + + const uint16_t repeatSize8 = 256; + const uint16_t repeatSize32 = 64; + + uint16_t histogramsLoopNum = (s2SeqLen + repeatSize8 - 1) / repeatSize8; + uint16_t inputLoopNum = (s2SeqLen + repeatSize32 - 1) / repeatSize32; + uint16_t topkLoopNum = (topK + 64 - 1) / 64; + + // find kth-value + HistogramsFirstVFImpl(histogramsBuf, inputBuf, histogramsLoopNum, flag); + FindFirstTargetBinVFImpl(idx0Buf, nkValueBuf, histogramsBuf, bottomK); + HistogramsSecondVFImpl(histogramsBuf, inputBuf, idx0Buf, histogramsLoopNum, flag); + FindSecondTargetBinVFImpl(idx1Buf, nkValueBuf, nkValueBuf, histogramsBuf); + HistogramsThirdVFImpl(histogramsBuf, inputBuf, idx0Buf, idx1Buf, histogramsLoopNum, flag); + FindThirdTargetBinVFImpl(idx2Buf, nkValueBuf, nkValueBuf, histogramsBuf); + HistogramsLastVFImpl(histogramsBuf, inputBuf, idx0Buf, idx1Buf, idx2Buf, histogramsLoopNum, flag); + FindKthVFImpl(nkValueBuf, histogramsBuf, idx0Buf, idx1Buf, idx2Buf, idx3Buf); + + // filter + // 输出大于k-value的值value + FindValueGTOutputVFImpl(outputValueBuf, inputBuf, nkValueBuf, inputLoopNum); + // value-当前偏移大于k-value的值在AR特殊寄存器中的有效字节数 + int64_t arValueNum = AscendC::GetSpr(); + // value-剩余需要输出等于k-value的数量 + int64_t remainValueNum = topK - (arValueNum / sizeof(uint32_t)); + for(uint16_t i = 0; i < inputLoopNum; ++i) { + int64_t arValueNumPerLoop = AscendC::GetSpr(); + if (((arValueNumPerLoop - arValueNum) / sizeof(uint32_t)) < remainValueNum) { + // 调用一次查找等于k-value情况的过程 + FindValueEQOutputVFImpl(outputValueBuf, inputBuf + i * 64, nkValueBuf); + } else { + break; + } + } + + // 输出大于k-value的值idx + FindIdxGTOutputVFImpl(outputIdxBuf, inputBuf, (uint32_t)(0), nkValueBuf, inputLoopNum); + // idx-当前偏移大于k-value的值在AR特殊寄存器中的有效字节数 + int64_t arIdxNum = AscendC::GetSpr(); + int64_t remainIdxNum = topK - (arIdxNum / sizeof(uint32_t)); + for(uint16_t i = 0; i < inputLoopNum; ++i) { + int64_t arIdxNumPerLoop = AscendC::GetSpr(); + if (((arIdxNumPerLoop - arIdxNum) / sizeof(uint32_t)) < remainIdxNum) { + // 调用一次查找等于k-value情况的过程 + beginIdx = i * 64; + FindIdxEQOutputVFImpl(outputIdxBuf, inputBuf + i * 64, beginIdx, nkValueBuf); + } else { + break; + } + } +} +} +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/vf/vf_topk_16_gather.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/vf/vf_topk_16_gather.h new file mode 100644 index 000000000..0978251a3 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/arch35/vf/vf_topk_16_gather.h @@ -0,0 +1,430 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file vf_top_k_16_gather.h +* \brief +*/ + +#ifndef VF_TOP_K_16_GATHER_H +#define VF_TOP_K_16_GATHER_H + +namespace topkb16gather { + +template +__simd_vf__ void HistogramsHighVFImpl(__ubuf__ uint32_t* histogramsBuf, __ubuf__ uint16_t* inputBuf, uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + // 计算直方图cout0 0-127 cout1 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + MicroAPI::RegTensor vregHigh; + MicroAPI::RegTensor vregLow; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vregLow, vregHigh, inputBuf + i * 256); + + MicroAPI::Histograms(cout0, (MicroAPI::RegTensor&)vregHigh, pregB8); + MicroAPI::Histograms(cout1, (MicroAPI::RegTensor&)vregHigh, pregB8); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, cout0U32Even, cout0U32Odd, pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, cout1U32Even, cout1U32Odd, pregB32); +} + +__simd_vf__ void FindHighTargetBinVFImpl(__ubuf__ uint32_t* idxHighBuf, __ubuf__ uint32_t* nkValueBuf, __ubuf__ uint32_t* histogramsBuf, uint32_t bottomK) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg pregGE; + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdxHigh; + + MicroAPI::RegTensor btmK; + MicroAPI::Duplicate(btmK, bottomK); + + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdxHigh; + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::Arange(idxC, i * 64); + + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + + MicroAPI::Compare(pregGE, cout, btmK, pregB32); + + MicroAPI::Squeeze(sqzIdxHigh, (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idxHighBuf, sqzIdxHigh, alignIdxHigh); + } + MicroAPI::StoreUnAlignPost(idxHighBuf, alignIdxHigh); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idxHigh; + MicroAPI::LoadAlign(idxHigh, idxHighBuf); + + MicroAPI::RegTensor idxAll1; + MicroAPI::RegTensor idxPrev0; + MicroAPI::RegTensor prevBinValue; + MicroAPI::Duplicate(idxAll1, 1); + + MicroAPI::RegTensor zeroAll; + MicroAPI::Duplicate(zeroAll, 0); + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + MicroAPI::Compare(preg0, idxHigh, zeroAll, pregB32); + MicroAPI::Sub(idxPrev0, idxHigh, (MicroAPI::RegTensor&)idxAll1, pregB32); + MicroAPI::ShiftRights(idxPrev0, idxPrev0, (int16_t)24, pregB32); + + MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev0, pregB32); + MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, preg0); + + MicroAPI::RegTensor nextK; + MicroAPI::Sub(nextK, btmK, prevBinValue, pregB32); + MicroAPI::StoreAlign(nkValueBuf, nextK, pregB32); +} + +template +__simd_vf__ void HistogramsLowVFImpl(__ubuf__ uint32_t* histogramsBuf, __ubuf__ uint16_t* inputBuf, __ubuf__ uint32_t* idxHighBuf, uint16_t vfLoop, bool init) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg pregEQ; + + // 计算直方图0-127 128-255 + MicroAPI::RegTensor cout0; + MicroAPI::RegTensor cout1; + MicroAPI::Duplicate(cout0, 0); + MicroAPI::Duplicate(cout1, 0); + + MicroAPI::RegTensor cout0U32Even; + MicroAPI::RegTensor cout0U32Odd; + MicroAPI::RegTensor cout1U32Even; + MicroAPI::RegTensor cout1U32Odd; + + MicroAPI::RegTensor idxHigh; + MicroAPI::LoadAlign(idxHigh, idxHighBuf); + + MicroAPI::RegTensor vregHigh; + MicroAPI::RegTensor vregLow; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t i = 0; i < vfLoop; ++i) { + MicroAPI::LoadAlign(vregLow, vregHigh, inputBuf + i * 256); + + MicroAPI::Compare(pregEQ, (MicroAPI::RegTensor&)vregHigh, (MicroAPI::RegTensor&)idxHigh, pregB8); + + MicroAPI::Histograms(cout0, (MicroAPI::RegTensor&)vregLow, pregEQ); + MicroAPI::Histograms(cout1, (MicroAPI::RegTensor&)vregLow, pregEQ); + } + + MicroAPI::Cast(cout0U32Even, cout0, pregB16); + MicroAPI::Cast(cout0U32Odd, cout0, pregB16); + MicroAPI::Cast(cout1U32Even, cout1, pregB16); + MicroAPI::Cast(cout1U32Odd, cout1, pregB16); + + MicroAPI::StoreAlign(histogramsBuf, cout0U32Even, cout0U32Odd, pregB32); + MicroAPI::StoreAlign(histogramsBuf + 128, cout1U32Even, cout1U32Odd, pregB32); +} + +__simd_vf__ void FindKthVFImpl(__ubuf__ uint32_t* kValue, __ubuf__ uint32_t* histogramsBuf, __ubuf__ uint32_t* idxHighBuf, __ubuf__ uint32_t* idxLowBuf) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg pregGE; + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdxLow; + + MicroAPI::RegTensor btmK; + MicroAPI::LoadAlign(btmK, kValue); + + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor cout; + MicroAPI::RegTensor sqzIdxLow; + + for (uint16_t i = 0; i < (uint16_t)(4); ++i) { + MicroAPI::Arange(idxC, i * 64); + + MicroAPI::LoadAlign(cout, histogramsBuf + i * 64); + + MicroAPI::Compare(pregGE, cout, btmK, pregB32); + + MicroAPI::Squeeze(sqzIdxLow, (MicroAPI::RegTensor&)idxC, pregGE); + MicroAPI::StoreUnAlign(idxLowBuf, sqzIdxLow, alignIdxLow); + } + MicroAPI::StoreUnAlignPost(idxLowBuf, alignIdxLow); + + MicroAPI::LocalMemBar(); + + MicroAPI::RegTensor idxHigh; + MicroAPI::RegTensor idxLow; + MicroAPI::LoadAlign(idxHigh, idxHighBuf); + MicroAPI::LoadAlign(idxLow, idxLowBuf); + + MicroAPI::RegTensor idxTmp; + MicroAPI::Duplicate(idxTmp, 0xff00); + + MicroAPI::And(idxHigh, idxHigh, (MicroAPI::RegTensor&)idxTmp, pregB32); + + MicroAPI::RegTensor idxK; + MicroAPI::Add(idxK, idxHigh, idxLow, pregB16); + + MicroAPI::StoreAlign(kValue, idxK, pregB32); +} + +/** + 输出所有大于的kth-value的Index + */ +__simd_vf__ void FindIdxGTOutputVFImpl(__ubuf__ uint16_t* outputIdxBuf, __ubuf__ uint16_t* inputValueBuf, uint16_t beginIdx, __ubuf__ uint32_t* kValue, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg poutGT; + + MicroAPI::ClearSpr(); + + MicroAPI::UnalignRegForStore alignIdx; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor sqzIdxOut; + + for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::Arange(idxC, beginIdx + i * 128); + + MicroAPI::LoadAlign(vregInput, inputValueBuf + i * 128); + + MicroAPI::Compare(poutGT, vregInput, (MicroAPI::RegTensor&)kthValue, pregB16); + + MicroAPI::Squeeze(sqzIdxOut, (MicroAPI::RegTensor&)idxC, poutGT); + MicroAPI::StoreUnAlign(outputIdxBuf, sqzIdxOut, alignIdx); + } + MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx); +} + +/** + 输出所有等于的kth-value的Index + */ +__simd_vf__ void FindIdxEQOutputVFImpl(__ubuf__ uint16_t* outputIdxBuf, __ubuf__ uint16_t* inputValueBuf, uint16_t beginIdx, __ubuf__ uint32_t* kValue, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg poutEQ; + + MicroAPI::UnalignRegForStore alignIdx; + + MicroAPI::RegTensor kthValue; + MicroAPI::LoadAlign(kthValue, kValue); + + MicroAPI::RegTensor vregInput; + MicroAPI::RegTensor idxC; + MicroAPI::RegTensor sqzIdxOut; + + for(uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::Arange(idxC, beginIdx + i * 128); + + MicroAPI::LoadAlign(vregInput, inputValueBuf + i * 128); + + MicroAPI::Compare(poutEQ, vregInput, (MicroAPI::RegTensor&)kthValue, pregB16); + + MicroAPI::Squeeze(sqzIdxOut, (MicroAPI::RegTensor&)idxC, poutEQ); + MicroAPI::StoreUnAlign(outputIdxBuf, sqzIdxOut, alignIdx); + } + MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx); +} + +/** + 输出最终的Value + */ +__simd_vf__ void FindValueOutputVFImpl(__ubuf__ uint16_t* outputValueBuf, __ubuf__ uint16_t* inputValueBuf, __ubuf__ uint16_t* tmpIdxBuf, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask(); + + MicroAPI::RegTensor tmpIdx; + MicroAPI::RegTensor outputValue; + + for(uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::LoadAlign(tmpIdx, tmpIdxBuf + i * 128); + + MicroAPI::Gather(outputValue, inputValueBuf, tmpIdx, pregB16); + + MicroAPI::StoreAlign(outputValueBuf + i * 128, outputValue, pregB16); + } +} + +/** + 输出最终的Idx + */ +__simd_vf__ void FindRealIndexVFImpl(__ubuf__ uint32_t* outputIdxBuf, __ubuf__ uint16_t* tmpIdxBuf, __ubuf__ uint32_t* hisIdxBuf, uint32_t topK, uint32_t loopIndex, uint16_t vfLoop) +{ + MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask(); + + MicroAPI::MaskReg pregNow; + MicroAPI::MaskReg pregHis; + + MicroAPI::RegTensor tmpIdx; + MicroAPI::RegTensor outputGatherIdx; + MicroAPI::RegTensor outputAddsIdx; + + for(uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) { + MicroAPI::LoadAlign(tmpIdx, tmpIdxBuf + i * 64); + + MicroAPI::Compares(pregNow, (MicroAPI::RegTensor&)tmpIdx, topK - 1, pregB32); + MicroAPI::Xor(pregHis, pregNow, pregB32, pregB32); + + MicroAPI::Gather(outputGatherIdx, hisIdxBuf, (MicroAPI::RegTensor&)tmpIdx, pregHis); + MicroAPI::Adds(outputAddsIdx, (MicroAPI::RegTensor&)tmpIdx, loopIndex, pregNow); + + MicroAPI::Add(outputGatherIdx, outputGatherIdx, outputAddsIdx, pregB32); + + MicroAPI::StoreAlign(outputIdxBuf + i * 64, outputGatherIdx, pregB32); + } +} + +/** + * @brief LiTopKVF 对一个validLen的输入进行topk算法，输出idx_tmp + * @param tmpIdxLocal Temp阶段输出的TopKIndex;如果s2SeqLen < 16K作为最终输出 validLen * 2B + * @param outputValueLocal 如果s2SeqLen > 16K并且是首轮输出Value topK * 2B + * @param inputValueLocal 输入Value validLen * 2B + * @param histogramsLocal 直方图 256 * 4B + * @param idxHighLocal 目标桶高八位 256 * 4B + * @param idxLowLocal 目标桶低八位 256 * 4B + * @param nkValueLocal 存储next_k的值 64 * 4B + * @param topK topK元素 + * @param validLen 有效元素个数:QLICommon::Align(topkCountAlign256_ + validTrunkLen, (uint32_t)256) + */ +template // 是否输出VALUE +__aicore__ inline void LiTopKVF(const LocalTensor& tmpIdxLocal, + const LocalTensor& outputValueLocal, + const LocalTensor& inputValueLocal, + const LocalTensor& histogramsLocal, + const LocalTensor& idxHighLocal, + const LocalTensor& idxLowLocal, + const LocalTensor& nkValueLocal, + uint32_t topK, + uint32_t validLen) +{ + __ubuf__ uint16_t* tmpIdxBuf = (__ubuf__ uint16_t*)tmpIdxLocal.GetPhyAddr(); + __ubuf__ uint16_t* outputValueBuf = (__ubuf__ uint16_t*)outputValueLocal.GetPhyAddr(); + __ubuf__ uint16_t* inputValueBuf = (__ubuf__ uint16_t*)inputValueLocal.GetPhyAddr(); + __ubuf__ uint32_t* histogramsBuf = (__ubuf__ uint32_t*)histogramsLocal.GetPhyAddr(); + __ubuf__ uint32_t* idxHighBuf = (__ubuf__ uint32_t*)idxHighLocal.GetPhyAddr(); + __ubuf__ uint32_t* idxLowBuf = (__ubuf__ uint32_t*)idxLowLocal.GetPhyAddr(); + __ubuf__ uint32_t* nkValueBuf = (__ubuf__ uint32_t*)nkValueLocal.GetPhyAddr(); + + uint32_t bottomK = validLen - topK + 1; + uint32_t beginIdx = 0; + bool flag = true; + + const uint16_t repeatSize8 = 256; + const uint16_t repeatSize16 = 128; + const uint16_t repeatSize32 = 64; + + uint16_t histogramsLoopNum = (validLen + repeatSize8 - 1) / repeatSize8; + uint16_t inputLoopNum = (validLen + repeatSize16 - 1) / repeatSize16; + uint16_t topkLoopNum = (topK + repeatSize32 - 1) / repeatSize32; + uint16_t topkLoopNum16 = (topK + repeatSize16 - 1) / repeatSize16; + + // find kth-value + HistogramsHighVFImpl(histogramsBuf, inputValueBuf, histogramsLoopNum, flag); + FindHighTargetBinVFImpl(idxHighBuf, nkValueBuf, histogramsBuf, bottomK); + + HistogramsLowVFImpl(histogramsBuf, inputValueBuf, idxHighBuf, histogramsLoopNum, flag); + FindKthVFImpl(nkValueBuf, histogramsBuf, idxHighBuf, idxLowBuf); + + // filter + // 输出大于k-value的值idx + FindIdxGTOutputVFImpl(tmpIdxBuf, inputValueBuf, (uint32_t)(0), nkValueBuf, inputLoopNum); + // 输出等于k-value的值idx + FindIdxEQOutputVFImpl(tmpIdxBuf, inputValueBuf, (uint32_t)(0), nkValueBuf, inputLoopNum); + + // 是否输出Value + if constexpr (ISOUTVALUE) { + FindValueOutputVFImpl(outputValueBuf, inputValueBuf, tmpIdxBuf, topkLoopNum16); + } +} + +/** + * @brief 通过idx_tmp gather出实际的TopKIndex，s2SeqLen > 16K才会执行 + * @param outputIdxLocal 输出Idx 有效:topK * 2B + * @param outputValueLocal 输出Value topK * 2B(以后需要输出实际value使用) + * @param inputValueLocal 输入Value validLen * 2B + * @param tmpIdxLocal 本轮tmpIdx输入 validLen * 2B (0 ~ validLen - 1) + * @param hisIdxLocal 上一轮实际Idx输入 有效:topK * 4B + * @param topK topK元素个数 + * @param loopBasicIdx 当前循环需要加上得基准Index + * @param validLen 有效元素个数 + */ +__aicore__ inline void LiTopKGatherVF(const LocalTensor& outputIdxLocal, + const LocalTensor& outputValueLocal, + const LocalTensor& inputValueLocal, + const LocalTensor& tmpIdxLocal, + const LocalTensor& hisIdxLocal, + uint32_t topK, + uint32_t loopBasicIdx, + uint32_t validLen) +{ + __ubuf__ uint32_t* outputIdxBuf = (__ubuf__ uint32_t*)outputIdxLocal.GetPhyAddr(); + __ubuf__ uint16_t* outputValueBuf = (__ubuf__ uint16_t*)outputValueLocal.GetPhyAddr(); + __ubuf__ uint16_t* inputValueBuf = (__ubuf__ uint16_t*)inputValueLocal.GetPhyAddr(); + __ubuf__ uint16_t* tmpIdxBuf = (__ubuf__ uint16_t*)tmpIdxLocal.GetPhyAddr(); + __ubuf__ uint32_t* hisIdxBuf = (__ubuf__ uint32_t*)hisIdxLocal.GetPhyAddr(); + + const uint16_t repeatSize32 = 64; + const uint16_t repeatSize16 = 128; + uint16_t topkLoopNum16 = (topK + repeatSize16 - 1) / repeatSize16; + uint16_t topkLoopNum32 = (topK + repeatSize32 - 1) / repeatSize32; + + FindRealIndexVFImpl(outputIdxBuf, tmpIdxBuf, hisIdxBuf, topK, loopBasicIdx, topkLoopNum32); +} +} +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/vllm_quant_lightning_indexer.cpp b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/vllm_quant_lightning_indexer.cpp new file mode 100644 index 000000000..76622ba0c --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/vllm_quant_lightning_indexer.cpp @@ -0,0 +1,55 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vllm_quant_lightning_indexer.cpp + * \brief + */ + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" +#if (__CCE_AICORE__ == 310) + #include "arch35/quant_lightning_indexer_kernel.h" +#else + #include "arch32/quant_lightning_indexer_kernel.h" +#endif +#include "vllm_quant_lightning_indexer_template_tiling_key.h" +using namespace QLIKernel; +using namespace optiling::detail; + +#define INVOKE_LI_NO_KFC_OP_IMPL(templateClass, ...) \ + do { \ + templateClass> op; \ + GET_TILING_DATA_WITH_STRUCT(QLITilingData, tiling_data_in, tiling); \ + const QLITilingData *__restrict tiling_data = &tiling_data_in; \ + op.Init(query, key, weights, queryScale, keyScale, actualSeqLengthsQ, actualSeqLengthsK, blocktable, \ + metadata, sparseIndices, user, tiling_data, &tPipe); \ + op.Process(); \ + } while (0) + +template +__global__ __aicore__ void vllm_quant_lightning_indexer(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights, + __gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale, + __gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK, + __gm__ uint8_t *blocktable, __gm__ uint8_t *metadata, + __gm__ uint8_t *sparseIndices, __gm__ uint8_t *sparseValues, + __gm__ uint8_t *workspace, __gm__ uint8_t *tiling) +{ + TPipe tPipe; + __gm__ uint8_t *user = GetUserWorkspace(workspace); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + #if (__CCE_AICORE__ == 310) + INVOKE_LI_NO_KFC_OP_IMPL(QLIPreload, fp8_e4m3fn_t, fp8_e4m3fn_t, float, uint16_t, int32_t, + PAGE_ATTENTION, LI_LAYOUT(Q_LAYOUT_T), LI_LAYOUT(K_LAYOUT_T)); + #else + INVOKE_LI_NO_KFC_OP_IMPL(QLIPreload, int8_t, int8_t, int32_t, + PAGE_ATTENTION, LI_LAYOUT(Q_LAYOUT_T), LI_LAYOUT(K_LAYOUT_T)); + #endif +} diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/vllm_quant_lightning_indexer_metadata.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/vllm_quant_lightning_indexer_metadata.h new file mode 100644 index 000000000..4c0f4a94c --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/vllm_quant_lightning_indexer_metadata.h @@ -0,0 +1,79 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file vllm_quant_lightning_indexer_metadata.h +* \brief +*/ + +#ifndef QUANT_LIGHTNING_INDEXER_METADATA_H +#define QUANT_LIGHTNING_INDEXER_METADATA_H + +#include + +namespace optiling { + +// Constants +inline constexpr uint32_t AIC_CORE_NUM = 36; +inline constexpr uint32_t AIV_CORE_NUM = 72; +constexpr uint32_t QLI_META_SIZE = 1024; +using QLI_METADATA_T = int32_t; + +inline constexpr uint32_t LI_METADATA_SIZE = 8; +inline constexpr uint32_t LD_METADATA_SIZE = 8; + +// LI Metadata Index Definitions +inline constexpr uint32_t LI_CORE_ENABLE_INDEX = 0; +inline constexpr uint32_t LI_BN2_START_INDEX = 1; +inline constexpr uint32_t LI_M_START_INDEX = 2; +inline constexpr uint32_t LI_S2_START_INDEX = 3; +inline constexpr uint32_t LI_BN2_END_INDEX = 4; +inline constexpr uint32_t LI_M_END_INDEX = 5; +inline constexpr uint32_t LI_S2_END_INDEX = 6; +inline constexpr uint32_t LI_FIRST_LD_DATA_WORKSPACE_IDX_INDEX = 7; + +// LD Metadata Index Definitions +inline constexpr uint32_t LD_CORE_ENABLE_INDEX = 0; +inline constexpr uint32_t LD_BN2_IDX_INDEX = 1; +inline constexpr uint32_t LD_M_IDX_INDEX = 2; +inline constexpr uint32_t LD_WORKSPACE_IDX_INDEX = 3; +inline constexpr uint32_t LD_WORKSPACE_NUM_INDEX = 4; +inline constexpr uint32_t LD_M_START_INDEX = 5; +inline constexpr uint32_t LD_M_NUM_INDEX = 6; + + /** + * @brief 获取属性的绝对索引 + * @param coreIdx 核索引 + * @param metaIdx 元数据索引 + * @param isAIV 是否为AIV数据，默认为false + * @return 返回属性的绝对索引 + */ +#ifdef __CCE_AICORE__ +__aicore__ inline uint32_t GetAttrAbsIndex(uint32_t coreIdx, uint32_t metaIdx, bool isAIV=false) +{ + if (isAIV) { + return LI_METADATA_SIZE * AIC_CORE_NUM + LD_METADATA_SIZE * coreIdx + metaIdx; + } else { + return LI_METADATA_SIZE * coreIdx + metaIdx; + } +} +#endif + +namespace detail { + struct QliMetaData { + uint32_t LIMetadata[AIC_CORE_NUM][LI_METADATA_SIZE]; + uint32_t LDMetadata[AIV_CORE_NUM][LD_METADATA_SIZE]; + }; +}; + +static_assert(QLI_META_SIZE * sizeof(QLI_METADATA_T) >= sizeof(detail::QliMetaData)); +}; + +#endif \ No newline at end of file diff --git a/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/vllm_quant_lightning_indexer_template_tiling_key.h b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/vllm_quant_lightning_indexer_template_tiling_key.h new file mode 100644 index 000000000..176e658c3 --- /dev/null +++ b/csrc/ascend/attention/vllm_quant_lightning_indexer/op_kernel/vllm_quant_lightning_indexer_template_tiling_key.h @@ -0,0 +1,79 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file vllm_quant_lightning_indexer_template_tiling_key.h + * \brief + */ + +#ifndef QUANT_LIGHTNING_INDEXER_TEMPLATE_TILING_KEY_H +#define QUANT_LIGHTNING_INDEXER_TEMPLATE_TILING_KEY_H + +#include "ascendc/host_api/tiling/template_argument.h" + +#define QLI_TPL_INT8 2 +#define QLI_TPL_INT32 3 +#define QLI_TPL_FLOAT32_E4M3FN 36 +#define QLI_LAYOUT_BSND 0 +#define QLI_LAYOUT_TND 1 +#define QLI_LAYOUT_PA_BSND 2 + +#define ASCENDC_TPL_4_BW 4 + +// 模板参数支持的范围定义 +#if (__CCE_AICORE__ == 310) + ASCENDC_TPL_ARGS_DECL(VllmQuantLightningIndexer, // 算子OpType + ASCENDC_TPL_DTYPE_DECL(DT_Q, QLI_TPL_FLOAT32_E4M3FN), ASCENDC_TPL_DTYPE_DECL(DT_K, QLI_TPL_FLOAT32_E4M3FN), + ASCENDC_TPL_DTYPE_DECL(DT_OUT, QLI_TPL_INT32), ASCENDC_TPL_BOOL_DECL(PAGE_ATTENTION, 1, 0), + ASCENDC_TPL_UINT_DECL(Q_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_BSND, + QLI_LAYOUT_TND), + ASCENDC_TPL_UINT_DECL(K_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, + QLI_LAYOUT_BSND, QLI_LAYOUT_TND, QLI_LAYOUT_PA_BSND), ); + // 支持的模板参数组合 + // 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法 + ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLI_TPL_FLOAT32_E4M3FN), ASCENDC_TPL_DTYPE_SEL(DT_K, QLI_TPL_FLOAT32_E4M3FN), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLI_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_BSND, QLI_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_PA_BSND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLI_TPL_FLOAT32_E4M3FN), ASCENDC_TPL_DTYPE_SEL(DT_K, QLI_TPL_FLOAT32_E4M3FN), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLI_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_BSND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLI_TPL_FLOAT32_E4M3FN), ASCENDC_TPL_DTYPE_SEL(DT_K, QLI_TPL_FLOAT32_E4M3FN), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLI_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_TND), ), ); +#else + ASCENDC_TPL_ARGS_DECL(VllmQuantLightningIndexer, // 算子OpType + ASCENDC_TPL_DTYPE_DECL(DT_Q, QLI_TPL_INT8), ASCENDC_TPL_DTYPE_DECL(DT_K, QLI_TPL_INT8), + ASCENDC_TPL_DTYPE_DECL(DT_OUT, QLI_TPL_INT32), ASCENDC_TPL_BOOL_DECL(PAGE_ATTENTION, 1, 0), + ASCENDC_TPL_UINT_DECL(Q_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_BSND, + QLI_LAYOUT_TND), + ASCENDC_TPL_UINT_DECL(K_LAYOUT_T, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, + QLI_LAYOUT_BSND, QLI_LAYOUT_TND, QLI_LAYOUT_PA_BSND), ); + // 支持的模板参数组合 + // 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法 + ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLI_TPL_INT8), ASCENDC_TPL_DTYPE_SEL(DT_K, QLI_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLI_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 1), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_BSND, QLI_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_PA_BSND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLI_TPL_INT8), ASCENDC_TPL_DTYPE_SEL(DT_K, QLI_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLI_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_BSND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_BSND), ), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_DTYPE_SEL(DT_Q, QLI_TPL_INT8), ASCENDC_TPL_DTYPE_SEL(DT_K, QLI_TPL_INT8), + ASCENDC_TPL_DTYPE_SEL(DT_OUT, QLI_TPL_INT32), ASCENDC_TPL_BOOL_SEL(PAGE_ATTENTION, 0), + ASCENDC_TPL_UINT_SEL(Q_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_TND), + ASCENDC_TPL_UINT_SEL(K_LAYOUT_T, ASCENDC_TPL_UI_LIST, QLI_LAYOUT_TND), ), ); +#endif + +#endif \ No newline at end of file diff --git a/csrc/ascend/build.sh b/csrc/ascend/build.sh new file mode 100755 index 000000000..6f01a004c --- /dev/null +++ b/csrc/ascend/build.sh @@ -0,0 +1,225 @@ +#!/bin/bash +# Copyright (c) 2024 Huawei Technologies Co., Ltd. +# This file is a part of the CANN Open Software. +# Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ====================================================================================================================== +# +# Low-level CANN custom-op build script for vllm-plugin-FL framework operators. +# This is the equivalent of vllm-ascend/csrc/build.sh, adapted to the +# vllm-plugin-FL directory layout where the CANN CMakeLists.txt lives under +# csrc/ascend/. + +set -e + +CURRENT_DIR=$(dirname $(readlink -f ${BASH_SOURCE[0]})) +BUILD_DIR=${CURRENT_DIR}/build +OUTPUT_DIR=${CURRENT_DIR}/output +USER_ID=$(id -u) +PARENT_JOB="false" +CHECK_COMPATIBLE="true" +VERBOSE="false" + +if [ "${USER_ID}" != "0" ]; then + DEFAULT_TOOLKIT_INSTALL_DIR="${HOME}/Ascend/ascend-toolkit/latest" + DEFAULT_INSTALL_DIR="${HOME}/Ascend/latest" +else + DEFAULT_TOOLKIT_INSTALL_DIR="/usr/local/Ascend/ascend-toolkit/latest" + DEFAULT_INSTALL_DIR="/usr/local/Ascend/latest" +fi + +ENABLE_BUILD_PKG="OFF" +ENABLE_BUILT_IN="OFF" + +BASE_CUSTOM_OPTION="-DBUILD_OPEN_PROJECT=ON -DBUILD_TYPE=Release -DENABLE_OPS_HOST=ON -DENABLE_OPS_KERNEL=ON" +CUSTOM_OPTION="${BASE_CUSTOM_OPTION}" + +function help_info() { + echo "Usage: $0 [options]" + echo "Options:" + echo + echo "-h|--help Displays help message." + echo + echo "-n|--op-name Specifies the compiled operator. If there are multiple values, separate them with semicolons and use quotation marks. The default is all." + echo " For example: -n \"flash_attention_score\" or -n \"flash_attention_score;flash_attention_score_grad\"" + echo + echo "-c|--compute-unit Specifies the chip type. If there are multiple values, separate them with semicolons and use quotation marks. The default is ascend910b." + echo " For example: -c \"ascend910b\" or -c \"ascend910b;ascend310p\"" + echo + echo "--pkg Build a self-extracting .run package." + echo + echo "--ops=OPS Same as -n (upstream style)." + echo + echo "--soc=SOC Same as -c (upstream style)." + echo + echo "--verbose Displays more compilation information." + echo +} + +function log() { + local current_time=$(date +"%Y-%m-%d %H:%M:%S") + echo "[${current_time}] $1" +} + +function set_env() +{ + source ${ASCEND_CANN_PACKAGE_PATH}/bin/setenv.bash || echo "0" + + export BISHENG_REAL_PATH=$(which bisheng || true) + + if [ -z "${BISHENG_REAL_PATH}" ];then + log "Error: bisheng compilation tool not found, Please check whether the cann package or environment variables are set." + exit 1 + fi +} + +function clean() +{ + if [ -n "${BUILD_DIR}" ];then + rm -rf ${BUILD_DIR} + fi + if [ -n "${OUTPUT_DIR}" ];then + rm -rf ${OUTPUT_DIR} + fi + mkdir -p ${BUILD_DIR} ${OUTPUT_DIR} +} + +function cmake_config() +{ + local extra_option="$1" + log "Info: cmake config ${CUSTOM_OPTION} ${extra_option} ." + cmake -S ${CURRENT_DIR} -B ${BUILD_DIR} ${CUSTOM_OPTION} ${extra_option} +} + +function build() +{ + local target="$1" + if [ "${VERBOSE}" == "true" ];then + local option="--verbose" + fi + cmake --build ${BUILD_DIR} --target ${target} ${JOB_NUM} ${option} +} + +function gen_bisheng(){ + local ccache_program=$1 + local gen_bisheng_dir=${BUILD_DIR}/gen_bisheng_dir + + if [ ! -d "${gen_bisheng_dir}" ];then + mkdir -p ${gen_bisheng_dir} + fi + + pushd ${gen_bisheng_dir} > /dev/null + cat > bisheng <> bisheng + fi + cat >> bisheng <<'EOF' +eval "${ccache_args} ${args}" +EOF + chmod +x bisheng + + export PATH=${gen_bisheng_dir}:$PATH + popd > /dev/null +} + +function build_package(){ + build package +} + +while [[ $# -gt 0 ]]; do + case $1 in + -h|--help) + help_info + exit + ;; + -n|--op-name) + ascend_op_name="$2" + shift 2 + ;; + -c|--compute-unit) + ascend_compute_unit="$2" + shift 2 + ;; + --pkg) + ENABLE_BUILD_PKG="ON" + ENABLE_BUILT_IN="OFF" + shift + ;; + --ops) + ascend_op_name="$2" + shift 2 + ;; + --soc) + ascend_compute_unit="$2" + shift 2 + ;; + --ops=*) + ascend_op_name="${1#*=}" + shift + ;; + --soc=*) + ascend_compute_unit="${1#*=}" + shift + ;; + --verbose) + VERBOSE="true" + shift + ;; + *) + help_info + exit 1 + ;; + esac +done + +CUSTOM_OPTION="${BASE_CUSTOM_OPTION} -DENABLE_BUILD_PKG=${ENABLE_BUILD_PKG} -DENABLE_BUILT_IN=${ENABLE_BUILT_IN}" + +if [ -n "${ascend_compute_unit}" ];then + CUSTOM_OPTION="${CUSTOM_OPTION} -DASCEND_COMPUTE_UNIT=${ascend_compute_unit}" +fi + +if [ -n "${ascend_op_name}" ];then + CUSTOM_OPTION="${CUSTOM_OPTION} -DASCEND_OP_NAME=${ascend_op_name}" +fi + +if [ -n "${ASCEND_HOME_PATH}" ];then + ASCEND_CANN_PACKAGE_PATH=${ASCEND_HOME_PATH} +elif [ -n "${ASCEND_OPP_PATH}" ];then + ASCEND_CANN_PACKAGE_PATH=$(dirname ${ASCEND_OPP_PATH}) +elif [ -d "${DEFAULT_TOOLKIT_INSTALL_DIR}" ];then + ASCEND_CANN_PACKAGE_PATH=${DEFAULT_TOOLKIT_INSTALL_DIR} +elif [ -d "${DEFAULT_INSTALL_DIR}" ];then + ASCEND_CANN_PACKAGE_PATH=${DEFAULT_INSTALL_DIR} +else + log "Error: Please set the toolkit package installation directory through ASCEND_HOME_PATH or ASCEND_OPP_PATH." + exit 1 +fi + +if [ "${PARENT_JOB}" == "false" ];then + CPU_NUM=$(($(grep -c "^processor" /proc/cpuinfo)*2)) + JOB_NUM="-j${CPU_NUM}" +fi + +CUSTOM_OPTION="${CUSTOM_OPTION} -DCUSTOM_ASCEND_CANN_PACKAGE_PATH=${ASCEND_CANN_PACKAGE_PATH} -DCHECK_COMPATIBLE=${CHECK_COMPATIBLE} -DCANN_3RD_LIB_PATH=${CURRENT_DIR}/third_party" + +set_env +clean + +ccache_system=$(which ccache || true) +if [ -n "${ccache_system}" ];then + CUSTOM_OPTION="${CUSTOM_OPTION} -DENABLE_CCACHE=ON -DCUSTOM_CCACHE=${ccache_system}" + gen_bisheng ${ccache_system} +fi + +cmake_config +build_package + +log "Info: CANN framework operator package built at ${BUILD_DIR}." +log "Info: Install with: bash ${BUILD_DIR}/cann-ops-transformer-*.run --install-path=\$YOUR_INSTALL_DIR" diff --git a/csrc/ascend/build_aclnn.sh b/csrc/ascend/build_aclnn.sh new file mode 100755 index 000000000..e7dcecb08 --- /dev/null +++ b/csrc/ascend/build_aclnn.sh @@ -0,0 +1,132 @@ +#!/bin/bash +# Copyright (c) 2026 BAAI. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Top-level CANN custom-op build entry for vllm-plugin-FL framework operators. +# Equivalent to vllm-ascend/csrc/build_aclnn.sh: handles SOC detection, catlass +# setup, operator selection and final installation. + +set -e + +ROOT_DIR=$(cd "$(dirname $(readlink -f ${BASH_SOURCE[0]}))/../../" && pwd) +SOC_VERSION="ascend910b" +CLEAN_THIRD_PARTY=0 + +# Parse arguments. The SOC_VERSION positional argument can appear anywhere; +# --clean-third-party is the only supported flag. +for arg in "$@"; do + case "$arg" in + --clean-third-party) + CLEAN_THIRD_PARTY=1 + ;; + --*) + echo "Unknown option: $arg" + exit 1 + ;; + *) + SOC_VERSION="$arg" + ;; + esac +done + +if [[ "$SOC_VERSION" =~ ^ascend310 ]]; then + echo "No custom aclnn ops for ASCEND310 series." + exit 0 +elif [[ "$SOC_VERSION" =~ ^ascend910b ]]; then + # ASCEND910B (A2) series + git config --global --add safe.directory "$ROOT_DIR" || true + CATLASS_PATH=${ROOT_DIR}/csrc/ascend/third_party/catlass/include + if [[ ! -d "${CATLASS_PATH}" ]]; then + echo "dependency catlass is missing, try to fetch it..." + if ! git submodule update --init --recursive; then + echo "fetch failed" + exit 1 + fi + fi + ABSOLUTE_CATLASS_PATH=$(cd "${CATLASS_PATH}" && pwd) + export CPATH=${ABSOLUTE_CATLASS_PATH}:${CPATH} + + CUSTOM_OPS="moe_grouped_matmul;grouped_matmul_swiglu_quant;grouped_matmul_swiglu_quant_v2;grouped_matmul_swiglu_quant_weight_nz_tensor_list;lightning_indexer;lightning_indexer_quant;sparse_flash_attention;kv_quant_sparse_flash_attention;sparse_attn_sharedkv;vllm_quant_lightning_indexer;compressor;compressor_metadata;fused_gdn_gating;inplace_partial_rotary_mul;recurrent_gated_delta_rule;rms_norm_dynamic_quant;store_kv_block;reshape_and_cache_bnsd;ngram_spec_decode;matmul_allreduce_add_rmsnorm;moe_init_routing_custom;moe_gating_top_k;moe_gating_top_k_hash;add_rms_norm_bias;apply_top_k_top_p_custom;causal_conv1d;chunk_fwd_o;chunk_gated_delta_rule_fwd_h;copy_and_expand_eagle_inputs;dequant_swiglu_quant;hamming_dist_top_k;hc_post;hc_pre;hc_pre_inv_rms;hc_pre_sinkhorn;scatter_nd_update_v2;transpose_kv_cache_by_block;" + SOC_ARG="ascend910b" +elif [[ "$SOC_VERSION" =~ ^ascend910_93 ]]; then + # ASCEND910C (A3) series + git config --global --add safe.directory "$ROOT_DIR" || true + CATLASS_PATH=${ROOT_DIR}/csrc/ascend/third_party/catlass/include + if [[ ! -d "${CATLASS_PATH}" ]]; then + echo "dependency catlass is missing, try to fetch it..." + if ! git submodule update --init --recursive; then + echo "fetch failed" + exit 1 + fi + fi + ABSOLUTE_CATLASS_PATH=$(cd "${CATLASS_PATH}" && pwd) + export CPATH=${ABSOLUTE_CATLASS_PATH}:${CPATH} + + if [[ -n "${HCCL_STRUCT_FILE_PATH}" ]]; then + yes | cp "${HCCL_STRUCT_FILE_PATH}" "${ROOT_DIR}/csrc/ascend/utils/inc/kernel" + fi + + CUSTOM_OPS_ARRAY=( + "grouped_matmul_swiglu_quant_weight_nz_tensor_list" + "lightning_indexer_vllm" + "sparse_flash_attention" + "dispatch_ffn_combine" + "dispatch_ffn_combine_bf16" + "dispatch_gmm_combine_decode" + "moe_combine_normal" + "moe_dispatch_normal" + "dispatch_layout" + "notify_dispatch" + "moe_init_routing_custom" + "moe_gating_top_k" + "add_rms_norm_bias" + "apply_top_k_top_p_custom" + "transpose_kv_cache_by_block" + "copy_and_expand_eagle_inputs" + "causal_conv1d" + "moe_grouped_matmul" + "lightning_indexer_quant" + ) + CUSTOM_OPS=$(IFS=';'; echo "${CUSTOM_OPS_ARRAY[*]}") + SOC_ARG="ascend910_93" +else + echo "Unsupported SOC version: $SOC_VERSION" + exit 1 +fi + +# Build custom ops +cd ${ROOT_DIR}/csrc/ascend +rm -rf build output build_out + +echo "building custom ops $CUSTOM_OPS for $SOC_VERSION" +bash build.sh --pkg --ops="$CUSTOM_OPS" --soc="$SOC_ARG" + +# Install custom ops to vllm_fl/_cann_ops_custom (isolated from system CANN). +INSTALL_DIR=${ROOT_DIR}/vllm_fl/_cann_ops_custom +RUN_PACKAGE=$(ls build/cann-ops-transformer*.run 2> /dev/null | head -n1) +if [[ -z "${RUN_PACKAGE}" ]]; then + RUN_PACKAGE=$(ls build_out/cann-ops-transformer*.run 2> /dev/null | head -n1) +fi +if [[ -z "${RUN_PACKAGE}" ]]; then + echo "Error: no .run package found under build/ or build_out/" + exit 1 +fi + +echo "installing ${RUN_PACKAGE} to ${INSTALL_DIR}" +bash "${RUN_PACKAGE}" --install-path="${INSTALL_DIR}" + +# Clean downloaded third-party build artifacts only when explicitly requested. +# catlass and pto-isa are source submodules and must be kept. +clean_third_party_artifacts() { + local third_party_dir="${ROOT_DIR}/csrc/ascend/third_party" + echo "[build_aclnn] cleaning downloaded third-party build artifacts ..." + rm -rf "${third_party_dir}/abseil-cpp" + rm -rf "${third_party_dir}/ascend_protobuf" + rm -rf "${third_party_dir}/json" + rm -rf "${third_party_dir}/pkg" +} +if [[ "${CLEAN_THIRD_PARTY}" == "1" ]]; then + clean_third_party_artifacts +fi + +echo "CANN framework operators built and installed for $SOC_VERSION." diff --git a/csrc/ascend/build_batch_invariant_ops.sh b/csrc/ascend/build_batch_invariant_ops.sh new file mode 100755 index 000000000..988bb0a0f --- /dev/null +++ b/csrc/ascend/build_batch_invariant_ops.sh @@ -0,0 +1,100 @@ +#!/bin/bash + +set -euo pipefail + +# arguments: +# $1: SOC_ARG (ascend910b, ascend910_93, ascend950) + +SOC_ARG="${1:-}" + +log() { + echo "[install_batch_invariant] $*" +} + +# validate arguments +if [[ -z "${SOC_ARG}" ]]; then + log "ERROR: SOC_ARG is required as first argument" + exit 1 +fi + +log "Starting batch_invariant installation..." +log "SOC_ARG=${SOC_ARG}" + +# determine device type from SOC_ARG +case "${SOC_ARG}" in + ascend910b) + BATCH_INVARIANT_DEVICE="910b" + ;; + ascend910_93) + BATCH_INVARIANT_DEVICE="A3" + ;; + *) + log "Warning: batch_invariant not available for SOC_ARG=${SOC_ARG}; skipping" + exit 0 + ;; +esac + +# detect system architecture +ARCH_INFO=$(uname -m) +case "${ARCH_INFO}" in + aarch64) + ARCH_SUFFIX="aarch64" + ;; + x86_64) + ARCH_SUFFIX="x86_64" + ;; + *) + log "Warning: unknown architecture ${ARCH_INFO}; cannot determine batch_invariant package" + exit 0 + ;; +esac + +# download and install run package +BATCH_INVARIANT_RUN_URL="https://vllm-ascend.obs.cn-north-4.myhuaweicloud.com/vllm-ascend/cann-ops-batch_invariant-${BATCH_INVARIANT_DEVICE}-1.0.0-linux.${ARCH_SUFFIX}.run" +BATCH_INVARIANT_RUN_FILE="cann-ops-batch_invariant-${BATCH_INVARIANT_DEVICE}-1.0.0-linux.${ARCH_SUFFIX}.run" + +log "Downloading batch_invariant run package..." +unset ASCEND_CUSTOM_OPP_PATH +if curl --max-time 60 -sS -k -O "${BATCH_INVARIANT_RUN_URL}" && [[ -f "${BATCH_INVARIANT_RUN_FILE}" ]]; then + chmod +x "${BATCH_INVARIANT_RUN_FILE}" + log "Running installer: ${BATCH_INVARIANT_RUN_FILE}" + if "./${BATCH_INVARIANT_RUN_FILE}"; then + log "batch_invariant run package installed successfully" + else + log "Failed to install batch_invariant run package" + fi +else + log "Failed to download batch_invariant run package: ${BATCH_INVARIANT_RUN_URL}" +fi +# clean up downloaded run file (always clean, regardless of success/failure) +rm -f "${BATCH_INVARIANT_RUN_FILE}" + +# download and install whl package +BATCH_INVARIANT_WHL_URL="https://vllm-ascend.obs.cn-north-4.myhuaweicloud.com/vllm-ascend/batch_invariant-torch_ops_extension-1.0.0.zip" +BATCH_INVARIANT_WHL_FILE="batch_invariant-torch_ops_extension-1.0.0.zip" + +log "Downloading batch_invariant whl package..." +if curl --max-time 3 -sS -k -O "${BATCH_INVARIANT_WHL_URL}" >/dev/null 2>&1 && [[ -f "${BATCH_INVARIANT_WHL_FILE}" ]]; then + if unzip -o "${BATCH_INVARIANT_WHL_FILE}" >/dev/null 2>&1; then + if [[ -d "torch_ops_extension/batch_invariant_ops" ]]; then + cd torch_ops_extension/batch_invariant_ops + log "Building and installing batch_invariant whl package..." + if bash build_and_install.sh; then + log "batch_invariant whl package installed successfully" + else + log "Failed to build and install batch_invariant whl package" + fi + cd - + else + log "batch_invariant_ops directory not found in zip" + fi + else + log "Failed to unzip batch_invariant whl package" + fi +else + log "Failed to download batch_invariant whl package: ${BATCH_INVARIANT_WHL_URL}" +fi +# clean up downloaded files (always clean, regardless of success/failure) +rm -rf "${BATCH_INVARIANT_WHL_FILE}" torch_ops_extension + +log "batch_invariant_ops build completed" diff --git a/csrc/ascend/camem_allocator.cpp b/csrc/ascend/camem_allocator.cpp new file mode 100644 index 000000000..04e896256 --- /dev/null +++ b/csrc/ascend/camem_allocator.cpp @@ -0,0 +1,347 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include + +extern "C" { + +#define PY_SSIZE_T_CLEAN +#include + +#include +#include "acl/acl.h" + +// Global references to Python callables +// NOTE: this is borrowed reference, so we don't need to DECREF them. +// This brings the limitation that the allocator needs to be singleton. +static PyObject* g_python_malloc_callback = nullptr; +static PyObject* g_python_free_callback = nullptr; + + +// --------------------------------------------------------------------------- +// Helper functions: + +void ensure_context(unsigned long long device) { + aclrtContext pctx; + aclrtGetCurrentContext(&pctx); + if (!pctx) { + // Ensure device context. + aclrtCreateContext(&pctx, device); + aclrtSetCurrentContext(pctx); + } +} + +void create_and_map(unsigned long long device, ssize_t size, void* d_mem, + aclrtDrvMemHandle* p_memHandle) { + ensure_context(device); + // Define memory allocation properties + aclrtPhysicalMemProp prop = {}; + prop.handleType = ACL_MEM_HANDLE_TYPE_NONE; + prop.allocationType = ACL_MEM_ALLOCATION_TYPE_PINNED; + prop.memAttr = ACL_HBM_MEM_HUGE; + prop.location.id = device; + prop.location.type = ACL_MEM_LOCATION_TYPE_DEVICE; + prop.reserve = 0; + + // Allocate memory using aclrtMallocPhysical + aclError error_code = aclrtMallocPhysical(p_memHandle, size, &prop, 0); + if (error_code != 0) { + if (error_code == ACL_ERROR_RT_MEMORY_ALLOCATION) { + throw std::runtime_error("aclrtMallocPhysical failed with acl error code: " + + std::to_string(error_code) + "(OOM: Out of Memory, allocation failed) " + + __FILE__ + ":" + std::to_string(__LINE__)); + } else { + throw std::runtime_error("aclrtMallocPhysical failed with acl error code: " + + std::to_string(error_code) + " " + __FILE__ + ":" + std::to_string(__LINE__)); + } + } + + // Map memory + error_code = aclrtMapMem(d_mem, size, 0, *p_memHandle, 0); + if (error_code != 0) { + throw std::runtime_error("aclrtMapMem failed with acl error code: " + + std::to_string(error_code) + " " + __FILE__ + ":" + std::to_string(__LINE__)); + } +} + +void unmap_and_release(unsigned long long device, ssize_t size, + void* d_mem, + aclrtDrvMemHandle* p_memHandle) { + // std::cout << "unmap_and_release: device=" << device << ", size=" << size << + // ", d_mem=" << d_mem << ", p_memHandle=" << p_memHandle << std::endl; + ensure_context(device); + aclError error_code = aclrtUnmapMem(d_mem); + if (error_code != 0) { + throw std::runtime_error("aclrtUnmapMem failed with acl error code: " + + std::to_string(error_code) + " " + __FILE__ + ":" + std::to_string(__LINE__)); + } + error_code = aclrtFreePhysical(*p_memHandle); + if (error_code != 0) { + throw std::runtime_error("aclrtFreePhysical failed with acl error code: " + + std::to_string(error_code) + " " + __FILE__ + ":" + std::to_string(__LINE__)); + } +} + +PyObject* create_tuple_from_c_integers(unsigned long long a, + unsigned long long b, + unsigned long long c, + unsigned long long d) { + // Create a new tuple of size 4 + PyObject* tuple = PyTuple_New(4); + if (!tuple) { + return NULL; // Return NULL on failure + } + + // Convert integers to Python objects and set them in the tuple + PyTuple_SetItem( + tuple, 0, + PyLong_FromUnsignedLongLong(a)); // Steals reference to the PyLong + PyTuple_SetItem(tuple, 1, PyLong_FromUnsignedLongLong(b)); + PyTuple_SetItem(tuple, 2, PyLong_FromUnsignedLongLong(c)); + PyTuple_SetItem(tuple, 3, PyLong_FromUnsignedLongLong(d)); + + // Note: PyTuple_SetItem "steals" a reference to each object, + // so we do not need to Py_DECREF the PyLong objects explicitly. + + return tuple; // Return the created tuple +} + +// --------------------------------------------------------------------------- +// Our exported C functions that call Python: + +__attribute__ ((visibility("default"))) void* my_malloc(ssize_t size, int device, aclrtStream stream) { + ensure_context(device); + + // first allocation, align the size, and reserve an address, and also allocate + // a aclrtDrvMemHandle + + // Define memory allocation properties + aclrtPhysicalMemProp prop = {}; + prop.handleType = ACL_MEM_HANDLE_TYPE_NONE ; + prop.allocationType = ACL_MEM_ALLOCATION_TYPE_PINNED; + prop.memAttr = ACL_HBM_MEM_HUGE; + prop.location.id = device; + prop.location.type = ACL_MEM_LOCATION_TYPE_DEVICE; + prop.reserve = 0; + + // Check if the allocation is supported + size_t granularity; + aclError error_code = aclrtMemGetAllocationGranularity(&prop, + ACL_RT_MEM_ALLOC_GRANULARITY_MINIMUM, + &granularity); + if (error_code != 0) { + throw std::runtime_error("aclrtMemGetAllocationGranularity failed with acl error code: " + + std::to_string(error_code) + " " + __FILE__ + ":" + std::to_string(__LINE__)); + } + size_t alignedSize = ((size + granularity - 1) / granularity) * granularity; + void *d_mem; + error_code = aclrtReserveMemAddress(&d_mem, alignedSize, 0, nullptr, 0); + if (error_code != 0) { + if (error_code == ACL_ERROR_RT_MEMORY_ALLOCATION) { + throw std::runtime_error("aclrtReserveMemAddress failed with acl error code: " + + std::to_string(error_code) + "(OOM: Out of Memory, allocation failed) " + + __FILE__ + ":" + std::to_string(__LINE__)); + } else { + throw std::runtime_error("aclrtReserveMemAddress failed with acl error code: " + + std::to_string(error_code) + " " + __FILE__ + ":" + std::to_string(__LINE__)); + } + } + // allocate the aclrtDrvMemHandle + aclrtDrvMemHandle* p_memHandle = + (aclrtDrvMemHandle*)malloc(sizeof(aclrtDrvMemHandle)); + + if (!g_python_malloc_callback) { + throw std::runtime_error("my_malloc ERROR: g_python_malloc_callback not set." + + std::string(" ") + __FILE__ + ":" + std::to_string(__LINE__)); + } + + // Acquire GIL (not in stable ABI officially, but often works) + PyGILState_STATE gstate = PyGILState_Ensure(); + + PyObject* arg_tuple = create_tuple_from_c_integers( + (unsigned long long)device, (unsigned long long)alignedSize, + (unsigned long long)d_mem, (unsigned long long)p_memHandle); + + // Call g_python_malloc_callback + PyObject* py_result = + PyObject_CallFunctionObjArgs(g_python_malloc_callback, arg_tuple, NULL); + Py_DECREF(arg_tuple); + + if (!py_result) { + PyErr_Print(); + PyGILState_Release(gstate); + return nullptr; + } + + PyGILState_Release(gstate); + + // do the final mapping + create_and_map(device, alignedSize, d_mem, p_memHandle); + + return (void*)d_mem; +} + +__attribute__ ((visibility("default"))) void my_free(void* ptr, ssize_t size, int device, aclrtStream stream) { + // get memory handle from the pointer + if (!g_python_free_callback) { + throw std::runtime_error("aclrtDrvMemHandle ERROR: g_python_malloc_callback not set." + + std::string(" ") + __FILE__ + ":" + std::to_string(__LINE__)); + } + + // Acquire GIL (not in stable ABI officially, but often works) + PyGILState_STATE gstate = PyGILState_Ensure(); + + PyObject* py_ptr = + PyLong_FromUnsignedLongLong(reinterpret_cast(ptr)); + + PyObject* py_result = + PyObject_CallFunctionObjArgs(g_python_free_callback, py_ptr, NULL); + + if (!py_result || !PyTuple_Check(py_result) || PyTuple_Size(py_result) != 4) { + PyErr_SetString(PyExc_TypeError, "Expected a tuple of size 4"); + return; + } + + unsigned long long recv_device, recv_size; + unsigned long long recv_d_mem, recv_p_memHandle; + // Unpack the tuple into four C integers + if (!PyArg_ParseTuple(py_result, "KKKK", &recv_device, &recv_size, + &recv_d_mem, &recv_p_memHandle)) { + // PyArg_ParseTuple sets an error if it fails + return; + } + + PyGILState_Release(gstate); + + // recv_size == size + // recv_device == device + + // Free memory + + void *d_mem = (void*)recv_d_mem; + // allocate the aclrtDrvMemHandle + aclrtDrvMemHandle* p_memHandle = + (aclrtDrvMemHandle*)recv_p_memHandle; + unmap_and_release(device, size, d_mem, p_memHandle); + + // free address and the handle + aclError error_code = aclrtReleaseMemAddress(d_mem); + if (error_code != 0) { + throw std::runtime_error("aclrtReleaseMemAddress failed with acl error code: " + + std::to_string(error_code) + " " + __FILE__ + ":" + std::to_string(__LINE__)); + } + free(p_memHandle); +} + +// --------------------------------------------------------------------------- +// Python extension boilerplate: + +// Python-exposed function: init_module(python_malloc, python_free) +static PyObject* py_init_module(PyObject* self, PyObject* args) { + PyObject* malloc_callback = nullptr; + PyObject* free_callback = nullptr; + + if (!PyArg_ParseTuple(args, "OO", &malloc_callback, &free_callback)) { + return nullptr; + } + + if (!PyCallable_Check(malloc_callback) || !PyCallable_Check(free_callback)) { + PyErr_SetString(PyExc_TypeError, "Both arguments must be callables"); + return nullptr; + } + + // Save the Python callables + // This module does not handle GC of these objects, so they must be kept alive + // outside of this module. + g_python_malloc_callback = malloc_callback; + g_python_free_callback = free_callback; + + Py_RETURN_NONE; +} + +static PyObject* python_unmap_and_release(PyObject* self, PyObject* args) { + if (!args || !PyTuple_Check(args) || PyTuple_Size(args) != 4) { + PyErr_SetString(PyExc_TypeError, "Expected a tuple of size 4"); + return nullptr; + } + + unsigned long long recv_device, recv_size; + unsigned long long recv_d_mem, recv_p_memHandle; + // Unpack the tuple into four C integers + if (!PyArg_ParseTuple(args, "KKKK", &recv_device, &recv_size, &recv_d_mem, + &recv_p_memHandle)) { + // PyArg_ParseTuple sets an error if it fails + return nullptr; + } + + void *d_mem_ptr = (void*)recv_d_mem; + aclrtDrvMemHandle* p_memHandle = + (aclrtDrvMemHandle*)recv_p_memHandle; + + unmap_and_release(recv_device, recv_size, d_mem_ptr, p_memHandle); + + Py_RETURN_NONE; +} + +static PyObject* python_create_and_map(PyObject* self, PyObject* args) { + if (!args || !PyTuple_Check(args) || PyTuple_Size(args) != 4) { + PyErr_SetString(PyExc_TypeError, "Expected a tuple of size 4"); + return nullptr; + } + + unsigned long long recv_device, recv_size; + unsigned long long recv_d_mem, recv_p_memHandle; + // Unpack the tuple into four C integers + if (!PyArg_ParseTuple(args, "KKKK", &recv_device, &recv_size, &recv_d_mem, + &recv_p_memHandle)) { + // PyArg_ParseTuple sets an error if it fails + return nullptr; + } + + void *d_mem_ptr = (void*)recv_d_mem; + aclrtDrvMemHandle* p_memHandle = + (aclrtDrvMemHandle*)recv_p_memHandle; + + create_and_map(recv_device, recv_size, d_mem_ptr, p_memHandle); + + Py_RETURN_NONE; +} + +static PyMethodDef module_methods[] = { + {"init_module", (PyCFunction)py_init_module, METH_VARARGS, + "Initialize module with python_malloc and python_free callables."}, + {"python_create_and_map", (PyCFunction)python_create_and_map, METH_VARARGS, + "Create and map memory on the device."}, + {"python_unmap_and_release", (PyCFunction)python_unmap_and_release, + METH_VARARGS, "Unmap and release memory on the device."}, + {NULL, NULL, 0, NULL} // sentinel +}; + +static struct PyModuleDef camem_allocator_module = { + PyModuleDef_HEAD_INIT, "camem_allocator", + "CANN-mem-based allocator for NPUPluggableAllocator", -1, module_methods}; + +PyMODINIT_FUNC PyInit__C_ascend(void) { + // Initialize the module + PyObject* module = PyModule_Create(&camem_allocator_module); + if (!module) { + return NULL; + } + return module; +} +} // extern "C" diff --git a/csrc/ascend/cmake/Third_Party_Open_Source_Software_List.yaml b/csrc/ascend/cmake/Third_Party_Open_Source_Software_List.yaml new file mode 100644 index 000000000..09d009074 --- /dev/null +++ b/csrc/ascend/cmake/Third_Party_Open_Source_Software_List.yaml @@ -0,0 +1,26 @@ +observability: + googletest: + cpeName: googletest + groupId: + url: https://github.com/google/googletest.git + version: v1.14.0 + makeself: + cpeName: makeself + groupId: + url: https://github.com/megastep/makeself.git + version: release-2.5.0-patch1 + protobuf: + cpeName: protobuf + groupId: + url: https://github.com/protocolbuffers/protobuf.git + version: v25.1 + json: + cpeName: json + groupId: + url: https://github.com/nlohmann/json.git + version: v3.11.3 + libboundscheck: + cpeName: libboundscheck + groupId: + url: https://github.com/openeuler/libboundscheck.git + version: v1.1.16 diff --git a/csrc/ascend/cmake/aclnn_ops_transformer.h.in b/csrc/ascend/cmake/aclnn_ops_transformer.h.in new file mode 100644 index 000000000..4c99dea0d --- /dev/null +++ b/csrc/ascend/cmake/aclnn_ops_transformer.h.in @@ -0,0 +1,15 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef @aclnn_master_header_name@_H_ +#define @aclnn_master_header_name@_H_ + +@aclnn_all_header_include_content@ +#endif // @aclnn_master_header_name@_H_ \ No newline at end of file diff --git a/csrc/ascend/cmake/build_empty_package.cmake b/csrc/ascend/cmake/build_empty_package.cmake new file mode 100644 index 000000000..3cfd4079e --- /dev/null +++ b/csrc/ascend/cmake/build_empty_package.cmake @@ -0,0 +1,81 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +function(cpack_empty_package) + include(cmake/third_party/makeself-fetch.cmake) + if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64") + message(STATUS "Detected architecture: x86_64") + set(ARCH x86_64) + elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64|arm") + message(STATUS "Detected architecture: ARM64") + set(ARCH aarch64) + else () + message(WARNING "Unknown architecture: ${CMAKE_SYSTEM_PROCESSOR}") + endif () + + # CPack config + install(FILES ${CMAKE_SOURCE_DIR}/version.info + DESTINATION share/info/ops_transformer + ) + install(FILES ${CMAKE_SOURCE_DIR}/scripts/package/ops_transformer/scripts/help.info + DESTINATION share/info/ops_transformer/script + ) + install(FILES ${CMAKE_SOURCE_DIR}/scripts/package/ops_transformer/scripts/empty_package_scripts/install.sh + DESTINATION share/info/ops_transformer/script + PERMISSIONS OWNER_EXECUTE OWNER_READ OWNER_WRITE GROUP_READ GROUP_EXECUTE WORLD_READ WORLD_EXECUTE + ) + install(FILES ${CMAKE_SOURCE_DIR}/scripts/package/ops_transformer/scripts/empty_package_scripts/cleanup.sh + DESTINATION share/info/ops_transformer/script + PERMISSIONS OWNER_EXECUTE OWNER_READ OWNER_WRITE GROUP_READ GROUP_EXECUTE WORLD_READ WORLD_EXECUTE + ) + string(FIND "${ASCEND_COMPUTE_UNIT}" ";" SEMICOLON_INDEX) + if (SEMICOLON_INDEX GREATER -1) + # 截取分号前的字串 + math(EXPR SUBSTRING_LENGTH "${SEMICOLON_INDEX}") + string(SUBSTRING "${ASCEND_COMPUTE_UNIT}" 0 "${SUBSTRING_LENGTH}" compute_unit) + else() + # 没有分号取全部内容 + set(compute_unit "${ASCEND_COMPUTE_UNIT}") + endif() + set(CMAKE_INSTALL_PREFIX ${CMAKE_SOURCE_DIR}/build_out) + string(REGEX REPLACE "^.*[Aa]scend" "" soc_version_temp "${ASCEND_COMPUTE_UNIT}") + # 检查是否成功提取 + if("${soc_version_temp}" STREQUAL "${ASCEND_COMPUTE_UNIT}") + set(soc_version "unknown") + else() + set(soc_version "${soc_version_temp}") + endif() + + if("${VERSION}" STREQUAL "") + set(CPACK_PACKAGE_FILE_NAME "cann-${soc_version}-ops-transformer_linux-${ARCH}.run") + else() + set(CPACK_PACKAGE_FILE_NAME "cann-${soc_version}-ops-transformer_${VERSION}_linux-${ARCH}.run") + endif() + + set(CPACK_INSTALL_PREFIX "/") + + set(CPACK_CMAKE_SOURCE_DIR "${CMAKE_SOURCE_DIR}") + set(CPACK_CMAKE_BINARY_DIR "${CMAKE_BINARY_DIR}") + set(CPACK_CMAKE_INSTALL_PREFIX "${CMAKE_INSTALL_PREFIX}") + set(CPACK_CMAKE_CURRENT_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}") + set(CPACK_MAKESELF_PATH "${MAKESELF_PATH}") + set(CPACK_SOC "${compute_unit}") + set(CPACK_ARCH "${ARCH}") + set(CPACK_SET_DESTDIR ON) + set(CPACK_GENERATOR External) + if (ENABLE_BUILT_IN) + set(CPACK_EXTERNAL_PACKAGE_SCRIPT "${CMAKE_SOURCE_DIR}/cmake/makeself_built_in.cmake") + endif() + set(CPACK_EXTERNAL_ENABLE_STAGING true) + set(CPACK_PACKAGE_DIRECTORY "${CMAKE_INSTALL_PREFIX}") + + message(STATUS "CMAKE_INSTALL_PREFIX = ${CMAKE_INSTALL_PREFIX}") + include(CPack) +endfunction() \ No newline at end of file diff --git a/csrc/ascend/cmake/config.cmake b/csrc/ascend/cmake/config.cmake new file mode 100644 index 000000000..12ddb70ca --- /dev/null +++ b/csrc/ascend/cmake/config.cmake @@ -0,0 +1,287 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +######################################################################################################################## +# 环境检查 +######################################################################################################################## + +# Python3 +find_package(Python3) +if ((NOT Python3_FOUND) OR (${Python3_EXECUTABLE} STREQUAL "")) + message(FATAL_ERROR "Can't find python3.") +endif () +set(HI_PYTHON "${Python3_EXECUTABLE}" CACHE STRING "python executor") + +# 获取基础 CANN 路径 +if (CUSTOM_ASCEND_CANN_PACKAGE_PATH) + set(ASCEND_CANN_PACKAGE_PATH ${CUSTOM_ASCEND_CANN_PACKAGE_PATH}) +elseif (DEFINED ENV{ASCEND_HOME_PATH}) + set(ASCEND_CANN_PACKAGE_PATH $ENV{ASCEND_HOME_PATH}) +elseif (DEFINED ENV{ASCEND_OPP_PATH}) + get_filename_component(ASCEND_CANN_PACKAGE_PATH "$ENV{ASCEND_OPP_PATH}/.." ABSOLUTE) +else() + set(ASCEND_CANN_PACKAGE_PATH "/usr/local/Ascend/latest") +endif () +message(STATUS "ASCEND_CANN_PACKAGE_PATH=${ASCEND_CANN_PACKAGE_PATH}") + +######################################################################################################################## +# 公共配置 +######################################################################################################################## + +# 开关类 +option(PREPARE_BUILD "Prepare build." OFF) +option(ENABLE_OPS_HOST "Build ops host." ON) +option(ENABLE_OPS_KERNEL "Build ops kernel." ON) +if (TESTS_EXAMPLE_OPS_TEST OR TESTS_UT_OPS_TEST) + set(ENABLE_OPS_KERNEL OFF) +endif () +set(OP_DEBUG_CONFIG "false" CACHE STRING "op debug config") + +# 路径配置 +# 源码树相关路径 +get_filename_component(OPS_ADV_DIR "${CMAKE_CURRENT_SOURCE_DIR}" REALPATH) +get_filename_component(OPS_ADV_CMAKE_DIR "${OPS_ADV_DIR}/cmake" REALPATH) +get_filename_component(OPS_ADV_UTILS_KERNEL_INC "${OPS_ADV_DIR}/common/include/kernel" REALPATH) + + +# 构建树相关路径 +set(ASCEND_IMPL_OUT_DIR ${CMAKE_CURRENT_BINARY_DIR}/impl CACHE STRING "ascend impl output directories") +set(ASCEND_BINARY_OUT_DIR ${CMAKE_CURRENT_BINARY_DIR}/binary CACHE STRING "ascend binary output directories") +set(ASCEND_AUTOGEN_DIR ${CMAKE_CURRENT_BINARY_DIR}/autogen CACHE STRING "Auto generate file directories") +set(ASCEND_CUSTOM_OPTIONS ${ASCEND_AUTOGEN_DIR}/custom_compile_options.ini) +set(ASCEND_CUSTOM_TILING_KEYS ${ASCEND_AUTOGEN_DIR}/custom_tiling_keys.ini) +set(ASCEND_CUSTOM_OPC_OPTIONS ${ASCEND_AUTOGEN_DIR}/custom_opc_options.ini) +set(OP_BUILD_TOOL ${ASCEND_CANN_PACKAGE_PATH}/tools/opbuild/op_build CACHE STRING "op_build tool") +file(MAKE_DIRECTORY ${ASCEND_AUTOGEN_DIR}) +file(REMOVE ${ASCEND_CUSTOM_OPTIONS}) +file(TOUCH ${ASCEND_CUSTOM_OPTIONS}) +file(REMOVE ${ASCEND_CUSTOM_TILING_KEYS}) +file(TOUCH ${ASCEND_CUSTOM_TILING_KEYS}) +file(REMOVE ${ASCEND_CUSTOM_OPC_OPTIONS}) +file(TOUCH ${ASCEND_CUSTOM_OPC_OPTIONS}) +if (BUILD_OPEN_PROJECT) + if(EXISTS ${ASCEND_CANN_PACKAGE_PATH}/${SYSTEM_PREFIX}/tikcpp/ascendc_kernel_cmake) + set(ASCEND_PROJECT_DIR ${ASCEND_CANN_PACKAGE_PATH}/${SYSTEM_PREFIX}/tikcpp/ascendc_kernel_cmake) + endif() + set(ASCEND_CMAKE_DIR ${ASCEND_PROJECT_DIR}/cmake CACHE STRING "ascend project cmake") + set(IMPL_INSTALL_DIR packages/vendors/${VENDOR_NAME}_transformer/op_impl/ai_core/tbe/${VENDOR_NAME}_transformer_impl) + set(IMPL_DYNAMIC_INSTALL_DIR packages/vendors/${VENDOR_NAME}_transformer/op_impl/ai_core/tbe/${VENDOR_NAME}_transformer_impl/dynamic) + set(ACLNN_INC_INSTALL_DIR packages/vendors/${VENDOR_NAME}_transformer/op_api/include/aclnnop) + set(ACLNN_INC_LEVEL2_INSTALL_DIR packages/vendors/${VENDOR_NAME}_transformer/op_api/include/aclnnop/level2) +else() + set(ASCEND_CMAKE_DIR ${TOP_DIR}/asl/ops/cann/ops/built-in/ascendc/samples/customize/cmake CACHE STRING "ascend project cmake") + set(IMPL_INSTALL_DIR lib/ascendc/impl) + set(IMPL_DYNAMIC_INSTALL_DIR lib/ascendc/impl/dynamic) + set(ACLNN_INC_INSTALL_DIR lib/include) + set(OPS_STATIC_TYPES infer train) + set(OPS_STATIC_SCRIPT ${TOP_DIR}/asl/ops/cann/ops/built-in/kernel/binary_script/build_opp_kernel_static.py) +endif () +if (EXISTS ${OPS_ADV_CMAKE_DIR}/scripts/util) + set(ASCENDC_CMAKE_UTIL_DIR ${OPS_ADV_CMAKE_DIR}/scripts/util) +else() + set(ASCENDC_CMAKE_UTIL_DIR ${ASCEND_CMAKE_DIR}/util) +endif() +set(CUSTOM_DIR ${CMAKE_BINARY_DIR}/custom) +set(TILING_CUSTOM_DIR ${CUSTOM_DIR}/op_impl/ai_core/tbe/op_tiling) +set(TILING_CUSTOM_FILE ${TILING_CUSTOM_DIR}/liboptiling.so) + +# 兼容ascendc变更临时适配，待切换新版本ascendc新版本后删除 +if(EXISTS ${ASCENDC_CMAKE_UTIL_DIR}/ascendc_gen_options.py) + set(ADD_OPS_COMPILE_OPTION_V2 ON) +else() + set(ADD_OPS_COMPILE_OPTION_V2 OFF) +endif() + +######################################################################################################################## +# CMake 选项, 缺省参数设置 +# 按 CMake 构建过程对 CMake 选项, CMake 缺省参数进行配置 +# CMake 构建过程: 1) 配置阶段(Configure); 2) 构建阶段(Build); 3) 安装阶段(Install); +######################################################################################################################## +if (BUILD_OPEN_PROJECT) + # 构建阶段(Build) + # 构建类型 + # CMake中的Generator(生成器)是用于生成本地/本机构建系统的工具。一般分为两种: + # 1. 单配置生成器(Single-configuration generator): + # 在配置(Configuration)阶段，仅允许指定一种构建类型，通过变量 CMAKE_BUILD_TYPE 指定; + # 在构建阶段(Build)无法更改构建类型，仅允许使用配置(Configuration)阶段通过变量 CMAKE_BUILD_TYPE 指定的构建类型; + # 常见的此类型生成器有: Ninja, Unix Makefiles + # 2. 多配置生成器(Multi-configuration generator) : + # 在配置(Configuration)阶段，仅指定构建阶段(Build)可用的构建类型列表，通过变量 CMAKE_CONFIGURATION_TYPES 指定; + # 在构建阶段(Build)通过 ”--config“ 参数，指定构建阶段具体的构建类型; + # 常见的此类型生成器有: Xcode, Visual Studio + # 所以: + # 1. 单配置生成器(Single-configuration generator)场景下，如果构建类型(CMAKE_BUILD_TYPE)未指定，则默认为 Debug ; + # 2. 多配置生成器(Multi-configuration generator)场景下，如果构建阶段可选的构建类型(CMAKE_CONFIGURATION_TYPES)未指定， + # 则默认将其指定为CMake允许的构建类型全集 [Debug;Release;MinSizeRel;RelWithDebInfo] + if (NOT BUILD_OPS_RTY_KERNEL) + if (ENABLE_TEST) + set(DEFAULT_BUILD_TYPE "Debug") + else() + set(DEFAULT_BUILD_TYPE "Release") + endif() + if (NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) + set(CMAKE_BUILD_TYPE "${DEFAULT_BUILD_TYPE}" CACHE STRING "Choose the build type: Release/Debug" FORCE) + endif() + endif() + get_property(GENERATOR_IS_MULTI_CONFIG GLOBAL PROPERTY GENERATOR_IS_MULTI_CONFIG) + if (GENERATOR_IS_MULTI_CONFIG) + if (NOT CMAKE_CONFIGURATION_TYPES) + set(CMAKE_CONFIGURATION_TYPES "Debug;Release;MinSizeRel;RelWithDebInfo" CACHE STRING "Configuration Build type" FORCE) + endif () + else () + if (NOT CMAKE_BUILD_TYPE) + set(CMAKE_BUILD_TYPE "Debug" CACHE STRING "Build type(default Debug)" FORCE) + endif () + endif () + + # 构建阶段(Build) + # 可执行文件运行时库文件搜索路径 RPATH + # 在 UTest 及 Example 场景不略去 RPATH + if (TESTS_UT_OPS_TEST OR TESTS_EXAMPLE_OPS_TEST) + set(CMAKE_SKIP_RPATH FALSE) + else () + set(CMAKE_SKIP_RPATH TRUE) + endif () + + # 构建阶段(Build) + # CCACHE 配置 + if (ENABLE_CCACHE) + if (CUSTOM_CCACHE) + set(CCACHE_PROGRAM ${CUSTOM_CCACHE}) + else() + find_program(CCACHE_PROGRAM ccache) + endif () + if (CCACHE_PROGRAM) + set(CMAKE_C_COMPILER_LAUNCHER ${CCACHE_PROGRAM} CACHE PATH "C cache Compiler") + set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE_PROGRAM} CACHE PATH "CXX cache Compiler") + endif () + endif () + + # 安装阶段(Install) + # 安装路径 + # 未显示设置 CMAKE_INSTALL_PREFIX (即 CMAKE_INSTALL_PREFIX 取缺省值)时, + # 修正其取值与构建树根目录 CMAKE_CURRENT_BINARY_DIR 平级 + if (CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT) + get_filename_component(_Install_Path_Prefix "${CMAKE_CURRENT_BINARY_DIR}/../output" REALPATH) + set(CMAKE_INSTALL_PREFIX "${_Install_Path_Prefix}" CACHE STRING "Install path" FORCE) + endif () +endif () + +######################################################################################################################## +# 公开编译参数 +######################################################################################################################## +list(TRANSFORM ASCEND_COMPUTE_UNIT TOLOWER) +if (BUILD_OPEN_PROJECT) + message(STATUS "ENABLE_CCACHE=${ENABLE_CCACHE}, CUSTOM_CCACHE=${CUSTOM_CCACHE}") + message(STATUS "CCACHE_PROGRAM=${CCACHE_PROGRAM}") + message(STATUS "ASCEND_COMPUTE_UNIT=${ASCEND_COMPUTE_UNIT}") + message(STATUS "ASCEND_OP_NAME=${ASCEND_OP_NAME}") + message(STATUS "TILING_KEY=${TILING_KEY}") + message(STATUS "TESTS_UT_OPS_TEST=${TESTS_UT_OPS_TEST}") + message(STATUS "TESTS_EXAMPLE_OPS_TEST=${TESTS_EXAMPLE_OPS_TEST}") + message(STATUS "CMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}") + message(STATUS "VERSION=${VERSION}") +endif () + +######################################################################################################################## +# 预处理 +######################################################################################################################## +if (BUILD_OPEN_PROJECT) + # 与基础 CANN 配套关系检查 + option(CHECK_COMPATIBLE "check compatibility" ON) + set(CHECK_COMPATIBLE OFF) + if (CHECK_COMPATIBLE) + set(_param + "--cann_path=${ASCEND_CANN_PACKAGE_PATH}" + "--cann_package_name=toolkit" + "check_code_compatible" + "--code_version_info_file=${CMAKE_CURRENT_SOURCE_DIR}/version.info" + ) + execute_process( + COMMAND ${HI_PYTHON} ${CMAKE_CURRENT_SOURCE_DIR}/cmake/scripts/check_version_compatible.py ${_param} + RESULT_VARIABLE result + OUTPUT_STRIP_TRAILING_WHITESPACE + OUTPUT_VARIABLE CANN_VERSION + ) + if (result) + message(FATAL_ERROR "Check version compatibility failed.") + else() + string(TOLOWER ${CANN_VERSION} CANN_VERSION) + endif () + endif () + + string(REPLACE "," ";" ASCEND_OP_NAME "${ASCEND_OP_NAME}") + + if (NOT PREPARE_BUILD AND ENABLE_OPS_KERNEL) + if (TILING_KEY) + string(REPLACE ";" "::" EP_TILING_KEY "${TILING_KEY}") + else() + set(EP_TILING_KEY FALSE) + endif () + + if (OPS_COMPILE_OPTIONS) + string(REPLACE ";" "::" EP_OPS_COMPILE_OPTIONS "${OPS_COMPILE_OPTIONS}") + else() + set(EP_OPS_COMPILE_OPTIONS FALSE) + endif () + + string(REPLACE ";" "::" EP_ASCEND_COMPUTE_UNIT "${ASCEND_COMPUTE_UNIT}") + + string(REPLACE ";" "::" EP_ASCEND_OP_NAME "${ASCEND_OP_NAME}") + + execute_process(COMMAND bash ${CMAKE_CURRENT_SOURCE_DIR}/cmake/scripts/prepare.sh + -s ${CMAKE_CURRENT_SOURCE_DIR} + -b ${CMAKE_CURRENT_BINARY_DIR}/prepare_build + -p ${ASCEND_CANN_PACKAGE_PATH} + --autogen-dir ${ASCEND_AUTOGEN_DIR} + --build-open-project ${BUILD_OPEN_PROJECT} + --binary-out-dir ${ASCEND_BINARY_OUT_DIR} + --impl-out-dir ${ASCEND_IMPL_OUT_DIR} + --op-build-tool ${OP_BUILD_TOOL} + --ascend-cmake-dir ${ASCEND_CMAKE_DIR} + --tiling-key ${EP_TILING_KEY} + --ops-compile-options ${EP_OPS_COMPILE_OPTIONS} + --check-compatible ${CHECK_COMPATIBLE} + --ascend-compute_unit ${EP_ASCEND_COMPUTE_UNIT} + --ascend-op_name ${EP_ASCEND_OP_NAME} + --op_debug_config ${OP_DEBUG_CONFIG} + --build_ops_rty_kernel ${BUILD_OPS_RTY_KERNEL} + --enable_built_in ${ENABLE_BUILT_IN} + --enable_static ${ENABLE_STATIC} + --enable_experimental ${ENABLE_EXPERIMENTAL} + --enable_ccache ${ENABLE_CCACHE} + --cann_3rd_lib_path ${CANN_3RD_LIB_PATH} + --build_type ${BUILD_TYPE} + --version ${VERSION} + --enable_oom ${ENABLE_OOM} + RESULT_VARIABLE result + OUTPUT_STRIP_TRAILING_WHITESPACE + OUTPUT_VARIABLE PREPARE_BUILD_OUTPUT_VARIABLE) + if (result) + message(FATAL_ERROR "Error: ops prepare build failed.") + endif () + + file(REMOVE ${ASCEND_CUSTOM_OPTIONS}) + file(TOUCH ${ASCEND_CUSTOM_OPTIONS}) + file(REMOVE ${ASCEND_CUSTOM_TILING_KEYS}) + file(TOUCH ${ASCEND_CUSTOM_TILING_KEYS}) + file(REMOVE ${ASCEND_CUSTOM_OPC_OPTIONS}) + file(TOUCH ${ASCEND_CUSTOM_OPC_OPTIONS}) + endif () +endif () + +######################################################################################################################## +# 其他配置 +######################################################################################################################## +if (BUILD_OPEN_PROJECT) + if (TESTS_UT_OPS_TEST) + include(${OPS_ADV_CMAKE_DIR}/config_utest.cmake) + endif () +endif () diff --git a/csrc/ascend/cmake/config_utest.cmake b/csrc/ascend/cmake/config_utest.cmake new file mode 100644 index 000000000..122a5bb47 --- /dev/null +++ b/csrc/ascend/cmake/config_utest.cmake @@ -0,0 +1,88 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +######################################################################################################################## +# 预定义变量 +######################################################################################################################## + +# 所使用的产品类型 +set(OPS_ADV_UTEST_OPS_TEST_ASCEND_PRODUCT_TYPE ascend910B1) + + +######################################################################################################################## +# 环境检查 +######################################################################################################################## + +if (EXISTS ${ASCEND_CANN_PACKAGE_PATH}/tools/tikicpulib/lib/cmake) + list(APPEND CMAKE_PREFIX_PATH ${ASCEND_CANN_PACKAGE_PATH}/tools/tikicpulib/lib/cmake) +else() + list(APPEND CMAKE_PREFIX_PATH ${ASCEND_CANN_PACKAGE_PATH}/toolkit/tools/tikicpulib/lib/cmake) +endif() +find_package(tikicpulib REQUIRED) + +# ASAN / UBSAN 场景随编译执行用例场景下, 将相关检查在编译前执行, 避免出现编译完成后又无法执行的情况, 影响使用体验. +# 仅 GNU 编译器需要设置 LD_PRELOAD +if ((ENABLE_ASAN OR ENABLE_UBSAN) AND "${CMAKE_C_COMPILER_ID}" STREQUAL "GNU") + message(STATUS "CMAKE_CXX_COMPILER_ID=${CMAKE_CXX_COMPILER_ID}") + if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU") + set(SAN_LD_PRELOAD "LD_PRELOAD=") + if (ENABLE_ASAN) + # libasan.so + execute_process(COMMAND ${CMAKE_C_COMPILER} --print-file-name=libasan.so + RESULT_VARIABLE _RST + OUTPUT_VARIABLE ASAN_SHARED_PATH) + if (_RST) + message(FATAL_ERROR "Can't get libasan.so path with ${CMAKE_C_COMPILER}") + endif () + get_filename_component(ASAN_SHARED_PATH "${ASAN_SHARED_PATH}" DIRECTORY) + get_filename_component(ASAN_SHARED_PATH "${ASAN_SHARED_PATH}/libasan.so" REALPATH) + if (NOT EXISTS ${ASAN_SHARED_PATH}) + message(FATAL_ERROR "ASAN_SHARED_PATH=${ASAN_SHARED_PATH} not exist.") + endif () + set(SAN_LD_PRELOAD "${SAN_LD_PRELOAD}:${ASAN_SHARED_PATH}") + endif () + if (ENABLE_UBSAN) + # libubsan.so + execute_process(COMMAND ${CMAKE_C_COMPILER} --print-file-name=libubsan.so + RESULT_VARIABLE _RST + OUTPUT_VARIABLE UBSAN_SHARED_PATH) + if (_RST) + message(FATAL_ERROR "Can't get libubsan.so path with ${CMAKE_C_COMPILER}") + endif () + get_filename_component(UBSAN_SHARED_PATH "${UBSAN_SHARED_PATH}" DIRECTORY) + get_filename_component(UBSAN_SHARED_PATH "${UBSAN_SHARED_PATH}/libubsan.so" REALPATH) + if (NOT EXISTS ${UBSAN_SHARED_PATH}) + message(FATAL_ERROR "UBSAN_SHARED_PATH=${UBSAN_SHARED_PATH} not exist.") + endif () + set(SAN_LD_PRELOAD "${SAN_LD_PRELOAD}:${UBSAN_SHARED_PATH}") + endif () + # libstdc++.so + execute_process(COMMAND ${CMAKE_C_COMPILER} --print-file-name=libstdc++.so + RESULT_VARIABLE _RST + OUTPUT_VARIABLE STDC_SHARED_PATH) + if (_RST) + message(FATAL_ERROR "Can't get libstdc++.so path with ${CMAKE_C_COMPILER}") + endif () + get_filename_component(STDC_SHARED_PATH "${STDC_SHARED_PATH}" DIRECTORY) + get_filename_component(STDC_SHARED_PATH "${STDC_SHARED_PATH}/libstdc++.so" REALPATH) + if (NOT EXISTS ${STDC_SHARED_PATH}) + message(FATAL_ERROR "STDC_SHARED_PATH=${STDC_SHARED_PATH} not exist.") + endif () + set(SAN_LD_PRELOAD "${SAN_LD_PRELOAD}:${STDC_SHARED_PATH}") + endif () +endif () + + +######################################################################################################################## +# 公共配置 +######################################################################################################################## + +# 开关类 +option(TESTS_UT_OPS_TEST_CI_PR "Build UTest in push request scene." OFF) diff --git a/csrc/ascend/cmake/custom_build.cmake b/csrc/ascend/cmake/custom_build.cmake new file mode 100644 index 000000000..e0bdb391f --- /dev/null +++ b/csrc/ascend/cmake/custom_build.cmake @@ -0,0 +1,885 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if (BUILD_OPEN_PROJECT) + set(CMAKE_MODULE_PATH + ${CMAKE_MODULE_PATH} + ${CMAKE_CURRENT_LIST_DIR}/cmake/modules + ) + + set(CMAKE_PREFIX_PATH + ${CMAKE_PREFIX_PATH} + ${ASCEND_CANN_PACKAGE_PATH} + ) + + set(_op_host_aclnn_link + $ + exe_graph + register + c_sec + ) + + find_package(alog MODULE) + + find_package(unified_dlog MODULE) + + if(NOT ${alog_FOUND}) + add_definitions(-DALOG_NOT_FOUND) + endif() + + add_library(op_host_aclnn SHARED EXCLUDE_FROM_ALL) + target_link_libraries(op_host_aclnn PRIVATE + ${_op_host_aclnn_link} + ) + target_compile_options(op_host_aclnn PRIVATE + $<$:-std=gnu++1z> + ) + + add_library(op_host_aclnnInner SHARED EXCLUDE_FROM_ALL) + target_link_libraries(op_host_aclnnInner PRIVATE + ${_op_host_aclnn_link} + ) + target_compile_options(op_host_aclnnInner PRIVATE + $<$:-std=gnu++1z> + ) + + add_library(op_host_aclnnExc SHARED EXCLUDE_FROM_ALL) + target_link_libraries(op_host_aclnnExc PRIVATE + ${_op_host_aclnn_link} + ) + target_compile_options(op_host_aclnnExc PRIVATE + $<$:-std=gnu++1z> + ) + + # op api + add_library(cust_opapi SHARED) + # When compiling a specified operator, there is an operator without aclnn src. + add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/cust_opapi_stub.cpp + COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/cust_opapi_stub.cpp + ) + target_sources(cust_opapi PRIVATE + ${CMAKE_CURRENT_BINARY_DIR}/cust_opapi_stub.cpp + ) + target_compile_options(cust_opapi PRIVATE + $<$:-std=gnu++1z> + ) + target_include_directories(cust_opapi PRIVATE + $ + $ + ) + target_compile_options(cust_opapi PRIVATE + -Werror=format + ) + target_compile_definitions(cust_opapi PRIVATE + -DACLNN_LOG_FMT_CHECK + ) + if(BUILD_WITH_3_8_PACKAGE) # 3~8 package links opapi_math and does not depend on opsbase. + target_link_libraries(cust_opapi PRIVATE + $ + -Wl,--whole-archive + ops_aclnn + -Wl,--no-whole-archive + # Avoid exporting CANN built-in ACLNN symbols through libcust_opapi.so. + $<$:$> + nnopbase + profapi + ge_common_base + ascend_dump + ascendalog + dl + ) + else() + target_link_libraries(cust_opapi PRIVATE + $ + -Wl,--whole-archive + ops_aclnn + -Wl,--no-whole-archive + # -lopapi + nnopbase + profapi + ge_common_base + ascend_dump + ascendalog + dl + ) + endif() + set_target_properties(cust_opapi PROPERTIES OUTPUT_NAME + cust_opapi + ) + if (NOT ENABLE_BUILT_IN) + install(TARGETS cust_opapi + LIBRARY DESTINATION packages/vendors/${VENDOR_NAME}_transformer/op_api/lib + ) + endif() + + # op proto + add_library(cust_proto SHARED) + target_compile_options(cust_proto PRIVATE + $<$:-std=c++11> + -fvisibility=hidden + ) + target_compile_definitions(cust_proto PRIVATE + LOG_CPP + PROCESS_LOG + ) + target_link_libraries(cust_proto PRIVATE + $ + $ + $<$:$> + $<$:$> + -Wl,--whole-archive + rt2_registry + -Wl,--no-whole-archive + -Wl,--no-as-needed + exe_graph + graph + graph_base + register + ascendalog + error_manager + platform + -Wl,--as-needed + c_sec + ) + set_target_properties(cust_proto PROPERTIES OUTPUT_NAME + cust_opsproto_rt2.0 + ) + if (NOT ENABLE_BUILT_IN) + install(TARGETS cust_proto + LIBRARY DESTINATION packages/vendors/${VENDOR_NAME}_transformer/op_proto/lib/linux/${CMAKE_SYSTEM_PROCESSOR} + ) + endif() + + # op tiling + add_library(cust_opmaster SHARED) + target_include_directories(cust_opmaster PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/mc2/common/inc + $<$:$> + ) + target_compile_options(cust_opmaster PRIVATE + $<$:-std=c++11> + -fvisibility=hidden + ) + target_compile_definitions(cust_opmaster PRIVATE + LOG_CPP + PROCESS_LOG + ) + target_link_libraries(cust_opmaster PRIVATE + $ + $ + $<$:$> + $<$:$> + -Wl,--whole-archive + rt2_registry + -Wl,--no-whole-archive + -Wl,--no-as-needed + graph + graph_base + exe_graph + platform + register + error_manager + ascendalog + unified_dlog + -Wl,--as-needed + -Wl,--whole-archive + tiling_api + -Wl,--no-whole-archive + c_sec + ) + set_target_properties(cust_opmaster PROPERTIES OUTPUT_NAME + cust_opmaster_rt2.0 + ) + add_custom_command(TARGET cust_opmaster + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E make_directory ${TILING_CUSTOM_DIR} + COMMAND ln -sf $ ${TILING_CUSTOM_FILE} + ) + if (NOT ENABLE_BUILT_IN) + install(TARGETS cust_opmaster + LIBRARY DESTINATION packages/vendors/${VENDOR_NAME}_transformer/op_impl/ai_core/tbe/op_tiling/lib/linux/${CMAKE_SYSTEM_PROCESSOR} + ) + endif() + + # optiling compat + set(compat_optiling_dir ${CMAKE_CURRENT_BINARY_DIR}/compat) + set(compat_optiling_file ${compat_optiling_dir}/liboptiling.so) + add_custom_target(optiling_compat ALL + DEPENDS ${compat_optiling_file} + ) + + add_custom_command( + OUTPUT ${compat_optiling_file} + COMMAND ${CMAKE_COMMAND} -E make_directory ${compat_optiling_dir} + COMMAND ln -sf lib/linux/${CMAKE_SYSTEM_PROCESSOR}/$ ${compat_optiling_file} + ) + + if (NOT ENABLE_BUILT_IN) + install(FILES ${compat_optiling_file} + DESTINATION packages/vendors/${VENDOR_NAME}_transformer/op_impl/ai_core/tbe/op_tiling + ) + endif() + + add_ops_tiling_keys( + OP_NAME "ALL" + TILING_KEYS ${TILING_KEY} + ) + + add_opc_config( + OP_NAME "ALL" + CONFIG ${OP_DEBUG_CONFIG} + ) + + if(ADD_OPS_COMPILE_OPTION_V2) + add_ops_compile_options( + OP_NAME "ALL" + OPTIONS ${OPS_COMPILE_OPTIONS} + ) + endif() +endif () + +add_subdirectory(common) + +set(OP_LIST) +set(OP_DIR_LIST) +op_add_subdirectory(OP_LIST OP_DIR_LIST) + +if (BUILD_OPEN_PROJECT) + if (ENABLE_TEST) + set(OP_UT_LIST) + set(OP_UT_DIR_LIST) + op_add_ut_subdirectory(OP_UT_LIST OP_UT_DIR_LIST) + foreach (OP_UT_LIST ${OP_UT_DIR_LIST}) + # 仅通过op_add_subdirectory添加的算子目录，需要在这里add tests + if(OP_UT_LIST IN_LIST OP_DIR_LIST) + add_subdirectory(${OP_UT_LIST}/tests) + endif() + endforeach () + + if (TESTS_UT_OPS_TEST) + add_subdirectory(tests/ut/framework_special) + add_definitions(-Wno-builtin-macro-redefined) + endif() + endif () + if (TESTS_EXAMPLE_OPS_TEST) + add_subdirectory(examples) + endif () +endif () + + +foreach (OP_DIR ${OP_DIR_LIST}) + if (EXISTS "${OP_DIR}/op_host") + add_subdirectory(${OP_DIR}/op_host) + if(EXISTS "${OP_DIR}/op_graph/CMakeLists.txt") + add_subdirectory(${OP_DIR}/op_graph) + endif() + else() + add_subdirectory(${OP_DIR}) + endif() +endforeach () + +if(ENABLE_EXPERIMENTAL) + # genop新增experimental算子分类 + # add_subdirectory(${op_class}) + add_subdirectory(experimental/attention) +else() + # genop新增非experimental算子分类 + # add_subdirectory(${op_class}) + add_subdirectory(attention) +endif() + + +if (UT_TEST_ALL OR OP_HOST_UT OR OP_API_UT OR OP_KERNEL_UT OR OP_GRAPH_UT) + add_subdirectory(tests/ut/framework_normal) +endif() + +if("${ASCEND_OP_NAME}" STREQUAL "add_example") + add_subdirectory(examples) + list(APPEND OP_DIR_LIST ${CMAKE_CURRENT_SOURCE_DIR}/examples/${ASCEND_OP_NAME}) +endif() + +if("${ASCEND_OP_NAME}" STREQUAL "all_gather_add") + add_subdirectory(examples/mc2) + list(APPEND OP_DIR_LIST ${CMAKE_CURRENT_SOURCE_DIR}/examples/mc2/${ASCEND_OP_NAME}) +endif() + +list(APPEND OP_LIST ${COMPILED_OPS}) +list(APPEND OP_DIR_LIST ${COMPILED_OP_DIRS}) + +if(ENABLE_TEST) + foreach (OP_DIR ${OP_DIR_LIST}) + file(READ "${OP_DIR}/tests/CMakeLists.txt" CML_CONTENT) + if (CML_CONTENT MATCHES "OpsTest_Level2_AddOp") + set(UTEST_FRAMEWORK_OLD TRUE CACHE BOOL "UTEST_FRAMEWORK_OLD" FORCE) + else() + set(UTEST_FRAMEWORK_NEW TRUE CACHE BOOL "UTEST_FRAMEWORK_NEW" FORCE) + endif() + endforeach() + if(TESTS_UT_OPS_TEST) + OpsTest_AddLaunch() + endif() +endif() + + +if (DEFINED MC2_OPT AND EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/mc2/common/CMakeLists.txt AND EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/mc2/3rd/CMakeLists.txt) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/mc2/common) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/mc2/3rd) +endif() + +set(OP_DEPEND_DIR_LIST) +op_add_depend_directory( + OP_LIST ${OP_LIST} + OP_DIR_LIST OP_DEPEND_DIR_LIST +) +# 仅针对被依赖的算子重新add_subdirectory +foreach (OP_DEPEND_DIR ${OP_DEPEND_DIR_LIST}) + get_filename_component(SUB_DIR ${OP_DEPEND_DIR} NAME) + if ("${ASCEND_OP_NAME}" STREQUAL "all" OR "${ASCEND_OP_NAME}" STREQUAL "ALL") + if ( "${OP_DEPEND_DIR}" MATCHES ".*attention.*") + continue() + endif() + endif() + if (NOT ${SUB_DIR} IN_LIST ASCEND_OP_NAME) + list(APPEND ASCEND_OP_NAME ${SUB_DIR}) + if (EXISTS "${OP_DEPEND_DIR}/op_host") + add_subdirectory(${OP_DEPEND_DIR}/op_host) + else() + add_subdirectory(${OP_DEPEND_DIR}) + endif() + endif () + if ( "${OP_DEPEND_DIR}" MATCHES ".*moe_inplace_index_add_with_sorted.*") + list(APPEND OP_DIR_LIST ${OPS_TRANSFORMER_DIR}/moe/3rd/moe_inplace_index_add_with_sorted) + endif() +endforeach () + +# ------------------------------------------------ aclnn ------------------------------------------------ +get_target_property(base_aclnn_srcs op_host_aclnn SOURCES) +get_target_property(base_aclnn_inner_srcs op_host_aclnnInner SOURCES) +get_target_property(base_aclnn_exclude_srcs op_host_aclnnExc SOURCES) + +if (BUILD_OPEN_PROJECT) + set(base_aclnn_binary_dir ${ASCEND_AUTOGEN_DIR}) +else() + get_target_property(base_aclnn_binary_dir op_host_aclnn BINARY_DIR) +endif () + +set(generate_aclnn_srcs) +set(generate_aclnn_inner_srcs) +set(generate_aclnn_headers) +set(generate_proto_dir ${base_aclnn_binary_dir}) +set(generate_exclude_proto_srcs) +set(generate_proto_srcs) +set(generate_proto_headers) + +if (base_aclnn_srcs) + foreach (_src ${base_aclnn_srcs}) + string(REGEX MATCH "^${CMAKE_CURRENT_SOURCE_DIR}" is_match "${_src}") + if (is_match) + get_filename_component(name_without_ext ${_src} NAME_WE) + + string(REGEX REPLACE "_def$" "" _op_name ${name_without_ext}) + list(APPEND generate_aclnn_srcs ${base_aclnn_binary_dir}/aclnn_${_op_name}.cpp) + list(APPEND generate_aclnn_headers ${base_aclnn_binary_dir}/aclnn_${_op_name}.h) + append_versioned_aclnn_outputs("${_src}" "aclnn" "${_op_name}" "${base_aclnn_binary_dir}" + generate_aclnn_srcs generate_aclnn_headers) + list(APPEND generate_proto_srcs ${generate_proto_dir}/${_op_name}_proto.cpp) + list(APPEND generate_proto_headers ${generate_proto_dir}/${_op_name}_proto.h) + endif () + endforeach () +else () + add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_stub.cpp + COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_stub.cpp + ) + + target_sources(op_host_aclnn PRIVATE + ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_stub.cpp + ) +endif () + +if (base_aclnn_inner_srcs) + foreach (_src ${base_aclnn_inner_srcs}) + string(REGEX MATCH "^${CMAKE_CURRENT_SOURCE_DIR}" is_match "${_src}") + if (is_match) + get_filename_component(name_without_ext ${_src} NAME_WE) + string(REGEX REPLACE "_def$" "" _op_name ${name_without_ext}) + list(APPEND generate_aclnn_inner_srcs ${base_aclnn_binary_dir}/inner/aclnnInner_${_op_name}.cpp) + list(APPEND generate_proto_srcs ${generate_proto_dir}/inner/${_op_name}_proto.cpp) + list(APPEND generate_proto_headers ${generate_proto_dir}/inner/${_op_name}_proto.h) + endif () + endforeach () +else () + add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_inner_stub.cpp + COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_inner_stub.cpp + ) + + target_sources(op_host_aclnnInner PRIVATE + ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_inner_stub.cpp + ) +endif () + +if (base_aclnn_exclude_srcs) + foreach (_src ${base_aclnn_exclude_srcs}) + string(REGEX MATCH "^${CMAKE_CURRENT_SOURCE_DIR}" is_match "${_src}") + if (is_match) + get_filename_component(name_without_ext ${_src} NAME_WE) + string(REGEX REPLACE "_def$" "" _op_name ${name_without_ext}) + list(APPEND generate_exclude_proto_srcs ${generate_proto_dir}/exc/${_op_name}_proto.cpp) + list(APPEND generate_proto_srcs ${generate_proto_dir}/exc/${_op_name}_proto.cpp) + list(APPEND generate_proto_headers ${generate_proto_dir}/exc/${_op_name}_proto.h) + endif () + endforeach () +else() + add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_exc_stub.cpp + COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_exc_stub.cpp + ) + + target_sources(op_host_aclnnExc PRIVATE + ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_exc_stub.cpp + ) +endif () + +if (BUILD_OPEN_PROJECT) + if (generate_aclnn_srcs OR generate_aclnn_inner_srcs) + set(ops_aclnn_src ${generate_aclnn_srcs} ${generate_aclnn_inner_srcs}) + else () + set(ops_aclnn_src ${CMAKE_CURRENT_BINARY_DIR}/ops_aclnn_src_stub.cpp) + + add_custom_command(OUTPUT ${ops_aclnn_src} + COMMAND touch ${ops_aclnn_src} + ) + endif () + + set_source_files_properties(${ops_aclnn_src} + PROPERTIES GENERATED TRUE + ) + add_library(ops_aclnn STATIC + ${ops_aclnn_src} + ) + target_include_directories(ops_aclnn PRIVATE + ${PROJECT_SOURCE_DIR}/common/include/common + ${PROJECT_SOURCE_DIR}/common/include/static + ) + target_compile_options(ops_aclnn PRIVATE + $<$:-std=gnu++1z> + ) + target_link_libraries(ops_aclnn PRIVATE + $ + ) + if (ENABLE_STATIC) + add_custom_target(opbuild_gen_aclnn_static + COMMAND python3 ${PROJECT_SOURCE_DIR}/scripts/util/modify_gen_aclnn.py ${CMAKE_BINARY_DIR} + DEPENDS opbuild_gen_default opbuild_gen_inner opbuild_gen_exc + ) + add_dependencies(ops_aclnn opbuild_gen_default opbuild_gen_inner opbuild_gen_aclnn_static) + else() + add_dependencies(ops_aclnn opbuild_gen_default opbuild_gen_inner) + endif() + + set_source_files_properties(${generate_proto_srcs} + PROPERTIES GENERATED TRUE + ) + target_sources(cust_proto PRIVATE + ${generate_proto_srcs} + ) + add_dependencies(cust_proto ops_transformer_proto_headers) + + if (NOT ENABLE_BUILT_IN) + install(FILES ${generate_proto_headers} + DESTINATION packages/vendors/${VENDOR_NAME}_transformer/op_proto/inc OPTIONAL + ) + endif() + + merge_graph_headers( + TARGET merge_ops_proto ALL + OUT_DIR ${ASCEND_GRAPH_CONF_DST} + ) + + add_dependencies(cust_proto merge_ops_proto) + target_sources(cust_proto PRIVATE + ${ASCEND_GRAPH_CONF_DST}/ops_proto_transformer.cpp + ) + + redefine_file_macro( + TARGET_NAME + op_host_aclnn + op_host_aclnnInner + op_host_aclnnExc + cust_opapi + cust_proto + cust_opmaster + ops_aclnn + ) +else() + if (generate_aclnn_srcs OR generate_aclnn_inner_srcs) + set_source_files_properties(${generate_aclnn_srcs} ${generate_aclnn_inner_srcs} + TARGET_DIRECTORY acl_op_builtin + PROPERTIES GENERATED TRUE + ) + + target_sources(acl_op_builtin PRIVATE + ${generate_aclnn_srcs} + ${generate_aclnn_inner_srcs} + ) + endif () + + if (generate_proto_srcs) + set_source_files_properties(${generate_proto_srcs} + TARGET_DIRECTORY cust_proto opsproto_rt2.0 + PROPERTIES GENERATED TRUE + ) + target_sources(cust_proto PRIVATE + ${generate_proto_srcs} + ) + add_dependencies(cust_proto ops_transformer_proto_headers) + + target_sources(opsproto_rt2.0 PRIVATE + ${generate_proto_srcs} + ) + add_dependencies(opsproto_rt2.0 ops_transformer_proto_headers) + endif () + + add_target_source( + TARGET_NAME opmaster_rt2.0 opmaster_static_rt2.0 + BASE_TARGET cust_opmaster + SRC_DIR ${CMAKE_CURRENT_SOURCE_DIR} + ) + + add_target_source( + TARGET_NAME opsproto_rt2.0 opsproto_static_rt2.0 + BASE_TARGET cust_proto + SRC_DIR ${CMAKE_CURRENT_SOURCE_DIR} + ) + + add_static_ops( + ACLNN_SRC ${generate_aclnn_srcs} + ACLNN_INNER_SRC ${generate_aclnn_inner_srcs} + SRC_DIR ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif () +target_sources(cust_opapi PRIVATE + $<$:$>) +if(NOT BUILD_WITH_3_8_PACKAGE) +target_link_libraries( + cust_opapi + PRIVATE $<$:$> + $<$:opsbase> +) +endif() + +if(BUILD_WITH_3_8_PACKAGE) +target_link_libraries( + cust_opmaster + PUBLIC ${OPHOST_NAME}_tiling_obj + PUBLIC $<$:$> + PUBLIC $<$:$> + PRIVATE $<$:$> +) +else () +target_link_libraries( + cust_opmaster + PUBLIC ${OPHOST_NAME}_tiling_obj + PUBLIC $<$:$> + PUBLIC $<$:$> + PRIVATE $<$:$> + $<$:opsbase> +) +endif() + +if(TARGET ${COMMON_NAME}_obj) + add_dependencies(cust_opmaster ${COMMON_NAME}_obj) +else() + message(WARNING "Target ${COMMON_NAME}_obj not found, dependency not added!") +endif() + +if(BUILD_WITH_3_8_PACKAGE) +target_link_libraries( + cust_proto + PUBLIC ${OPHOST_NAME}_infer_obj +) +else() +target_link_libraries( + cust_proto + PUBLIC ${OPHOST_NAME}_infer_obj + PRIVATE $<$:opsbase> +) +endif() +if (generate_aclnn_headers) + install(FILES ${generate_aclnn_headers} + DESTINATION ${ACLNN_INC_INSTALL_DIR} OPTIONAL + ) +endif () + +add_library(ops_transformer_proto_headers INTERFACE) + +target_include_directories(ops_transformer_proto_headers INTERFACE + $ + $ + $ + $ +) + +if ((NOT BUILD_OPEN_PROJECT) AND ("${PRODUCT_SIDE}" STREQUAL "device")) + ExternalProject_Add(extern_opbuild_gen_transformer + SOURCE_DIR ${TOP_DIR}/cmake/superbuild + CONFIGURE_COMMAND ${CMAKE_COMMAND} + -G ${CMAKE_GENERATOR} + -DHOST_PACKAGE=opp + -DBUILD_MOD=ops + -DUSE_CCACHE=${USE_CCACHE} + -DCMAKE_INSTALL_PREFIX=${CMAKE_CURRENT_BINARY_DIR}/opbuild_output + -DFEATURE_LIST=custom_opbuild_out_dir=${generate_proto_dir} + + BUILD_COMMAND TARGETS=opbuild_gen_all $(MAKE) + INSTALL_COMMAND "" + LIST_SEPARATOR :: + EXCLUDE_FROM_ALL TRUE + ) + add_dependencies(ops_transformer_proto_headers extern_opbuild_gen_transformer) +else() + add_dependencies(ops_transformer_proto_headers opbuild_gen_default opbuild_gen_inner opbuild_gen_exc) +endif () + +if (NOT BUILD_OPEN_PROJECT) + if (generate_proto_srcs) + install_package( + PACKAGE ops_adv + TARGETS ops_transformer_proto_headers + FILES ${generate_proto_headers} + DESTINATION include/ops_adv/proto + ) + endif () +endif () + +# ------------------------------------------------ opbuild ------------------------------------------------ +if (BUILD_OPEN_PROJECT) + string(REPLACE ";" "\;" OPS_PRODUCT_NAME "${ASCEND_COMPUTE_UNIT}") + if (generate_aclnn_srcs) + add_custom_command(OUTPUT ${generate_aclnn_srcs} ${generate_aclnn_headers} + COMMAND mkdir -p ${base_aclnn_binary_dir} + COMMAND OPS_PROTO_SEPARATE=1 + OPS_ACLNN_GEN=1 + OPS_PROJECT_NAME=aclnn + OPS_PRODUCT_NAME=\"${OPS_PRODUCT_NAME}\" + ${OP_BUILD_TOOL} + $ + ${base_aclnn_binary_dir} + ) + endif () + + add_custom_target(opbuild_gen_default + DEPENDS ${generate_aclnn_srcs} ${generate_aclnn_headers} op_host_aclnn + ) + + if (generate_aclnn_inner_srcs) + add_custom_command(OUTPUT ${generate_aclnn_inner_srcs} + COMMAND mkdir -p ${base_aclnn_binary_dir}/inner + COMMAND OPS_PROTO_SEPARATE=1 + OPS_ACLNN_GEN=1 + OPS_PROJECT_NAME=aclnnInner + OPS_PRODUCT_NAME=\"${OPS_PRODUCT_NAME}\" + ${OP_BUILD_TOOL} + $ + ${base_aclnn_binary_dir}/inner + ) + endif () + + add_custom_target(opbuild_gen_inner + DEPENDS ${generate_aclnn_inner_srcs} op_host_aclnnInner + ) + + if (generate_exclude_proto_srcs) + add_custom_command(OUTPUT ${generate_exclude_proto_srcs} + COMMAND mkdir -p ${base_aclnn_binary_dir}/exc + COMMAND OPS_PROTO_SEPARATE=1 + OPS_ACLNN_GEN=0 + OPS_PROJECT_NAME=aclnnExc + OPS_PRODUCT_NAME=\"${OPS_PRODUCT_NAME}\" + ${OP_BUILD_TOOL} + $ + ${base_aclnn_binary_dir}/exc + ) + endif () + + add_custom_target(opbuild_gen_exc + DEPENDS ${generate_exclude_proto_srcs} op_host_aclnnExc + ) +endif () + +# ------------------------------------------------ generate adapt py ------------------------------------------------ +add_custom_target(generate_transformer_adapt_py + COMMAND ${HI_PYTHON} ${CMAKE_CURRENT_SOURCE_DIR}/cmake/scripts/util/ascendc_impl_build.py + \"\" + \"\" + \"\" + \"\" + ${ASCEND_IMPL_OUT_DIR} + ${ASCEND_AUTOGEN_DIR} + --opsinfo-dir ${base_aclnn_binary_dir} ${base_aclnn_binary_dir}/inner ${base_aclnn_binary_dir}/exc +) + +add_dependencies(generate_transformer_adapt_py opbuild_gen_default opbuild_gen_inner opbuild_gen_exc) + +foreach (_op_name ${OP_LIST}) + install(FILES ${ASCEND_IMPL_OUT_DIR}/dynamic/${_op_name}.py + DESTINATION ${IMPL_DYNAMIC_INSTALL_DIR} + OPTIONAL + ) + install(FILES ${ASCEND_IMPL_OUT_DIR}/dynamic/${_op_name}_apt.py + DESTINATION ${IMPL_DYNAMIC_INSTALL_DIR} + OPTIONAL + ) +endforeach () + +install(DIRECTORY ${OPS_ADV_UTILS_KERNEL_INC}/ + DESTINATION ${IMPL_INSTALL_DIR}/ascendc/common +) + +# install(DIRECTORY ${OPS_ADV_DIR}/mc2/common/inc/kernel +# DESTINATION ${IMPL_INSTALL_DIR}/ascendc/common/inc +# ) + +# install(DIRECTORY ${OPS_ADV_DIR}/mc2/3rd/ +# DESTINATION ${IMPL_INSTALL_DIR}/ascendc/3rd +# ) + +foreach (op_dir ${OP_DIR_LIST}) + get_filename_component(_op_name "${op_dir}" NAME) + set(CURRENT_KERNEL_DIR "${op_dir}/op_kernel") + file(GLOB KERNEL_SUB_DIRS RELATIVE "${CURRENT_KERNEL_DIR}" "${CURRENT_KERNEL_DIR}/*") + filter_copy_files(SELECTED_FILES SELECTED_DIRS) + install(FILES ${SELECTED_FILES} + DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_name} + OPTIONAL + ) + install(DIRECTORY ${SELECTED_DIRS} + DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_name} + OPTIONAL + ) + + foreach (op_depend_dir ${${_op_name}_depends}) + set(CURRENT_KERNEL_DIR "${OPS_TRANSFORMER_DIR}/${op_depend_dir}/op_kernel") + file(GLOB KERNEL_SUB_DIRS RELATIVE "${CURRENT_KERNEL_DIR}" "${CURRENT_KERNEL_DIR}/*") + get_filename_component(_op_depened_name "${op_depend_dir}" NAME) + filter_copy_files(SELECTED_DEPEND_FILES SELECTED_DEPEND_DIRS) + install(FILES ${SELECTED_DEPEND_FILES} + DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_depened_name} + OPTIONAL + ) + install(DIRECTORY ${SELECTED_DEPEND_DIRS} + DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_depened_name} + OPTIONAL + ) + endforeach () +endforeach () + +# ------------------------------------------------ generate compile cmd ------------------------------------------------ +if (BUILD_OPEN_PROJECT) + add_custom_target(prepare_build ALL) + add_custom_target(generate_compile_cmd ALL) + add_custom_target(generate_ops_info ALL) + add_dependencies(prepare_build generate_transformer_adapt_py generate_compile_cmd) + + foreach (compute_unit ${ASCEND_COMPUTE_UNIT}) + add_compile_cmd_target( + COMPUTE_UNIT ${compute_unit} + ) + + add_ops_info_target( + COMPUTE_UNIT ${compute_unit} + ) + endforeach () +else() + add_dependencies(tbe_ops_json_info generate_transformer_adapt_py) +endif () + +# ------------------------------------------------ opp kernel ------------------------------------------------ +if (ENABLE_OPS_KERNEL) + add_custom_target(ops_transformer_kernel ALL) + add_custom_target(ops_transformer_config ALL) + add_dependencies(ops_transformer_kernel ops_transformer_config) + + foreach (compute_unit ${ASCEND_COMPUTE_UNIT}) + add_bin_compile_target( + COMPUTE_UNIT + ${compute_unit} + OP_INFO + ${OP_DIR_LIST} + ) + endforeach () +endif () + +if (NOT ENABLE_BUILT_IN AND BUILD_OPEN_PROJECT) + add_custom_target(modify_vendor ALL + DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/scripts/install.sh ${CMAKE_CURRENT_BINARY_DIR}/scripts/upgrade.sh + ) + + # modify VENDOR_NAME in install.sh and upgrade.sh + if (EXISTS ${ASCEND_PROJECT_DIR}/fwk_modules/scripts) + set(ASCEND_PROJECT_DIR_SCRIPTS_PATH ${ASCEND_PROJECT_DIR}/fwk_modules/scripts) + else() + set(ASCEND_PROJECT_DIR_SCRIPTS_PATH ${CMAKE_SOURCE_DIR}/cmake/scripts/custom) + endif() + add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/scripts/install.sh ${CMAKE_CURRENT_BINARY_DIR}/scripts/upgrade.sh + COMMAND mkdir -p ${CMAKE_CURRENT_BINARY_DIR}/scripts + COMMAND cp -r ${ASCEND_PROJECT_DIR_SCRIPTS_PATH}/* ${CMAKE_CURRENT_BINARY_DIR}/scripts/ + COMMAND chmod +w ${CMAKE_CURRENT_BINARY_DIR}/scripts/* + COMMAND sed -i "s/vendor_name=customize/vendor_name=${VENDOR_NAME}_transformer/g" ${CMAKE_CURRENT_BINARY_DIR}/scripts/* + ) + + install(DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/scripts/ + DESTINATION . FILE_PERMISSIONS OWNER_EXECUTE OWNER_READ GROUP_READ + ) + + # gen version.info + set(version_info_dir ${CMAKE_CURRENT_BINARY_DIR}) + set(version_info_file ${version_info_dir}/version.info) + add_custom_target(gen_version_info ALL + DEPENDS ${version_info_file} + ) + + add_custom_command(OUTPUT ${version_info_file} + COMMAND bash ${ASCENDC_CMAKE_UTIL_DIR}/gen_version_info.sh ${ASCEND_CANN_PACKAGE_PATH} ${version_info_dir} + ) + + install(FILES ${version_info_file} + DESTINATION packages/vendors/${VENDOR_NAME}_transformer/ + ) + + if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64") + message(STATUS "Detected architecture: x86_64") + set(ARCH x86_64) + elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64|arm") + message(STATUS "Detected architecture: ARM64") + set(ARCH aarch64) + else () + message(WARNING "Unknown architecture: ${CMAKE_SYSTEM_PROCESSOR}") + endif () + + # CPack config + set(CPACK_PACKAGE_NAME ${CMAKE_PROJECT_NAME}) + set(CPACK_PACKAGE_VERSION ${CMAKE_PROJECT_VERSION}) + set(CPACK_PACKAGE_DESCRIPTION "CPack ops project") + set(CPACK_PACKAGE_DESCRIPTION_SUMMARY "CPack ops project") + set(CPACK_PACKAGE_DIRECTORY ${CMAKE_BINARY_DIR}) + set(CPACK_PACKAGE_FILE_NAME "cann-ops-transformer-${VENDOR_NAME}_linux-${ARCH}.run") + set(CPACK_GENERATOR External) + set(CPACK_CMAKE_GENERATOR "${CMAKE_GENERATOR}") + set(CPACK_EXTERNAL_ENABLE_STAGING TRUE) + if (ENABLE_BUILD_PKG) + if (EXISTS ${ASCEND_CMAKE_DIR}/makeself.cmake) + set(CPACK_EXTERNAL_PACKAGE_SCRIPT ${ASCEND_CMAKE_DIR}/makeself.cmake) + else() + set(CPACK_MAKESELF_PATH ${OPS_TRANSFORMER_DIR}/third_party/makeself) + set(CPACK_EXTERNAL_PACKAGE_SCRIPT ${CMAKE_SOURCE_DIR}/cmake/makeself_custom.cmake) + endif() + endif() + set(CPACK_EXTERNAL_BUILT_PACKAGES ${CPACK_PACKAGE_DIRECTORY}/_CPack_Packages/Linux/External/${CPACK_PACKAGE_FILE_NAME}/${CPACK_PACKAGE_FILE_NAME}) + include(CPack) +endif () diff --git a/csrc/ascend/cmake/dependencies.cmake b/csrc/ascend/cmake/dependencies.cmake new file mode 100644 index 000000000..51241a780 --- /dev/null +++ b/csrc/ascend/cmake/dependencies.cmake @@ -0,0 +1,64 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +# Ascend mode +if(DEFINED ENV{ASCEND_HOME_PATH}) + set(ASCEND_DIR $ENV{ASCEND_HOME_PATH}) +else() + if ("$ENV{USER}" STREQUAL "root") + if(EXISTS /usr/local/Ascend/ascend-toolkit/latest) + set(ASCEND_DIR /usr/local/Ascend/ascend-toolkit/latest) + else() + set(ASCEND_DIR /usr/local/Ascend/latest) + endif() + else() + if(EXISTS $ENV{HOME}/Ascend/ascend-toolkit/latest) + set(ASCEND_DIR $ENV{HOME}/Ascend/ascend-toolkit/latest) + else() + set(ASCEND_DIR $ENV{HOME}/Ascend/latest) + endif() + endif() +endif() +message(STATUS "Search libs under install path ${ASCEND_DIR}") + +set(CMAKE_PREFIX_PATH ${ASCEND_DIR}/) + +set(CMAKE_MODULE_PATH + ${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules + ${CMAKE_MODULE_PATH} +) +message(STATUS "CMAKE_MODULE_PATH :${CMAKE_MODULE_PATH}") + +set(OPS_TRANSFORMER_CXX_FLAGS) +string(APPEND OPS_TRANSFORMER_CXX_FLAGS " ${COMPILE_OP_MODE}") +string(APPEND OPS_TRANSFORMER_CXX_FLAGS " -Wall") +string(APPEND OPS_TRANSFORMER_CXX_FLAGS " -Wextra") +string(APPEND OPS_TRANSFORMER_CXX_FLAGS " -Wshadow") +string(APPEND OPS_TRANSFORMER_CXX_FLAGS " -Wformat=2") +string(APPEND OPS_TRANSFORMER_CXX_FLAGS " -fno-common") +string(APPEND OPS_TRANSFORMER_CXX_FLAGS " -fPIC") +if(NOT "${CMAKE_CXX_COMPILER_ID}" STREQUAL "Clang") +# string(APPEND OPS_TRANSFORMER_CXX_FLAGS " -Werror") # TODO: add -Werror when fix all compile warnings + string(APPEND OPS_TRANSFORMER_CXX_FLAGS " -Wformat-signedness") + string(APPEND OPS_TRANSFORMER_CXX_FLAGS " -Wno-missing-include-dirs") + string(APPEND OPS_TRANSFORMER_CXX_FLAGS " -Wno-write-strings") +endif() +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OPS_TRANSFORMER_CXX_FLAGS}") +message(STATUS "compile option:${CMAKE_CXX_FLAGS}") + +find_package(dlog MODULE REQUIRED) +find_package(securec MODULE) +find_package(OPBASE MODULE REQUIRED) +find_package(platform MODULE REQUIRED) +find_package(metadef MODULE REQUIRED) +find_package(runtime MODULE REQUIRED) +find_package(nnopbase MODULE REQUIRED) +find_package(tilingapi MODULE REQUIRED) +find_package(aicpu MODULE REQUIRED) \ No newline at end of file diff --git a/csrc/ascend/cmake/func.cmake b/csrc/ascend/cmake/func.cmake new file mode 100644 index 000000000..46664fe79 --- /dev/null +++ b/csrc/ascend/cmake/func.cmake @@ -0,0 +1,904 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +function(filter_copy_files SELECTED_FILES SELECTED_DIRS) + set(_selected_files "") + set(_selected_dirs "") + foreach(item ${KERNEL_SUB_DIRS}) + set(path "${CURRENT_KERNEL_DIR}/${item}") + if(IS_DIRECTORY "${path}") + if(item MATCHES "^arch") + list(FIND ARCH_DIRECTORY "${item}" idx) + if(idx EQUAL -1) + continue() + endif() + endif() + list(APPEND _selected_dirs "${path}") + else() + list(APPEND _selected_files "${path}") + endif() + endforeach() + + set(${SELECTED_FILES} "${_selected_files}" PARENT_SCOPE) + set(${SELECTED_DIRS} "${_selected_dirs}" PARENT_SCOPE) +endfunction() + +function(add_target_source) + cmake_parse_arguments(ADD "" "BASE_TARGET;SRC_DIR" "TARGET_NAME" ${ARGN}) + + get_target_property(all_srcs ${ADD_BASE_TARGET} SOURCES) + set(add_srcs) + foreach(_src ${all_srcs}) + string(REGEX MATCH "^${ADD_SRC_DIR}" is_match "${_src}") + if (is_match) + list(APPEND add_srcs ${_src}) + endif () + endforeach() + + get_target_property(all_includes ${ADD_BASE_TARGET} INCLUDE_DIRECTORIES) + set(add_includes) + foreach(_include ${all_includes}) + string(REGEX MATCH "^${ADD_SRC_DIR}" is_match "${_include}") + if (is_match) + list(APPEND add_includes ${_include}) + endif () + endforeach() + + foreach(_target_name ${ADD_TARGET_NAME}) + target_sources(${_target_name} PRIVATE + ${add_srcs} + ) + + target_include_directories(${_target_name} PRIVATE + ${add_includes} + ) + endforeach() +endfunction() + +function(op_add_subdirectory OP_LIST OP_DIR_LIST) + set(_OP_LIST) + set(_OP_DIR_LIST) + + if(ENABLE_EXPERIMENTAL) + message(STATUS "Build experimental module") + file(GLOB OP_HOST_CMAKE_FILES + "${CMAKE_CURRENT_SOURCE_DIR}/experimental/ffn/**/op_host/CMakeLists.txt" + "${CMAKE_CURRENT_SOURCE_DIR}/experimental/gmm/**/op_host/CMakeLists.txt" + "${CMAKE_CURRENT_SOURCE_DIR}/experimental/mc2/**/op_host/CMakeLists.txt" + "${CMAKE_CURRENT_SOURCE_DIR}/experimental/moe/**/op_host/CMakeLists.txt" + "${CMAKE_CURRENT_SOURCE_DIR}/experimental/posembedding/**/op_host/CMakeLists.txt" + ) + else() + file(GLOB OP_HOST_CMAKE_FILES + "${CMAKE_CURRENT_SOURCE_DIR}/gmm/**/op_host/CMakeLists.txt" + "${CMAKE_CURRENT_SOURCE_DIR}/gmm/**/CMakeLists.txt" + ) + if(BUILD_OPEN_PROJECT AND (NOT BUILD_OPS_RTY_KERNEL)) + file(GLOB CANNDEV_OPS_HOST_CMAKE_FILES + "${CMAKE_CURRENT_SOURCE_DIR}/posembedding/**/op_host/CMakeLists.txt" + "${CMAKE_CURRENT_SOURCE_DIR}/moe/**/op_host/CMakeLists.txt" + "${CMAKE_CURRENT_SOURCE_DIR}/ffn/**/op_host/CMakeLists.txt" + "${CMAKE_CURRENT_SOURCE_DIR}/mc2/**/op_host/CMakeLists.txt" + "${CMAKE_CURRENT_SOURCE_DIR}/posembedding/**/framework/CMakeLists.txt" + "${CMAKE_CURRENT_SOURCE_DIR}/moe/**/framework/CMakeLists.txt" + "${CMAKE_CURRENT_SOURCE_DIR}/ffn/**/framework/CMakeLists.txt" + "${CMAKE_CURRENT_SOURCE_DIR}/mc2/**/framework/CMakeLists.txt" + ) + List(APPEND OP_HOST_CMAKE_FILES ${CANNDEV_OPS_HOST_CMAKE_FILES}) + endif() + endif() + + foreach(OP_CMAKE_FILE ${OP_HOST_CMAKE_FILES}) + if ("${OP_CMAKE_FILE}" MATCHES "op_host") + get_filename_component(OP_HOST_DIR "${OP_CMAKE_FILE}" DIRECTORY) + get_filename_component(OP_DIR "${OP_HOST_DIR}" DIRECTORY) + else() + get_filename_component(OP_DIR "${OP_CMAKE_FILE}" DIRECTORY) + endif() + get_filename_component(OP_NAME "${OP_DIR}" NAME) + + if (NOT BUILD_OPEN_PROJECT) + if (EXISTS ${TOP_DIR}/asl/ops/cann/ops/built-in/tbe/impl/ascendc/${OP_NAME}) + continue() + endif () + endif () + + if (DEFINED ASCEND_OP_NAME AND NOT "${ASCEND_OP_NAME}" STREQUAL "") + if (NOT "${ASCEND_OP_NAME}" STREQUAL "all" AND NOT "${ASCEND_OP_NAME}" STREQUAL "ALL") + if (NOT ${OP_NAME} IN_LIST ASCEND_OP_NAME) + continue() + endif () + endif () + endif () + + if (DEFINED ASCEND950_SKIP_OPS AND ${OP_NAME} IN_LIST ASCEND950_SKIP_OPS) + message(STATUS "Skipping ${OP_NAME} in op_add_subdirectory because ASCEND950 is not available.") + continue() + endif () + + if (ENABLE_TEST) + file(READ "${OP_DIR}/tests/CMakeLists.txt" CML_CONTENT) + if (CML_CONTENT MATCHES "OpsTest_Level2_AddOp") + set(UTEST_FRAMEWORK_OLD TRUE CACHE BOOL "UTEST_FRAMEWORK_OLD" FORCE) + else() + set(UTEST_FRAMEWORK_NEW TRUE CACHE BOOL "UTEST_FRAMEWORK_NEW" FORCE) + endif() + endif() + + list(APPEND _OP_LIST ${OP_NAME}) + list(APPEND _OP_DIR_LIST ${OP_DIR}) + endforeach() + + list(REMOVE_DUPLICATES _OP_LIST) + list(REMOVE_DUPLICATES _OP_DIR_LIST) + list(SORT _OP_LIST) + list(SORT _OP_DIR_LIST) + set(${OP_LIST} ${_OP_LIST} PARENT_SCOPE) + set(${OP_DIR_LIST} ${_OP_DIR_LIST} PARENT_SCOPE) +endfunction() + +macro(add_op_to_compiled_list) + get_filename_component(PARENT_DIR ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) + get_filename_component(OP_NAME ${PARENT_DIR} NAME) + # 记录全局的COMPILED_OPS和COMPILED_OP_DIRS，其中COMPILED_OP_DIRS只记录到算子名，例如moe/moe_token_permute_with_routing_map_grad + set(COMPILED_OPS ${COMPILED_OPS} ${OP_NAME} CACHE STRING "Compiled Ops" FORCE) + set(COMPILED_OP_DIRS ${COMPILED_OP_DIRS} ${PARENT_DIR} CACHE STRING "Compiled Ops Dirs" FORCE) +endmacro() + + +function(op_add_depend_directory) + cmake_parse_arguments(DEP "" "OP_DIR_LIST" "OP_LIST" ${ARGN}) + set(_OP_DEPEND_DIR_LIST) + foreach(op_name ${DEP_OP_LIST}) + if (DEFINED ${op_name}_depends) + foreach(depend_info ${${op_name}_depends}) + if (NOT EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${depend_info}/op_host/CMakeLists.txt AND NOT EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/src/${depend_info}/CMakeLists.txt) + continue() + endif () + + get_filename_component(_depend_op_name "${depend_info}" NAME) + if (NOT BUILD_OPEN_PROJECT) + if (EXISTS ${TOP_DIR}/asl/ops/cann/ops/built-in/tbe/impl/ascendc/${_depend_op_name}) + continue() + endif () + endif () + + if (NOT ${_depend_op_name} IN_LIST DEP_OP_LIST) + list(APPEND _OP_DEPEND_DIR_LIST ${CMAKE_CURRENT_SOURCE_DIR}/${depend_info}) + endif () + endforeach() + endif() + endforeach() + + list(SORT _OP_DEPEND_DIR_LIST) + list(REMOVE_DUPLICATES _OP_DEPEND_DIR_LIST) + set(${DEP_OP_DIR_LIST} ${_OP_DEPEND_DIR_LIST} PARENT_SCOPE) +endfunction() + +function(add_compile_cmd_target) + cmake_parse_arguments(CMD "" "COMPUTE_UNIT" "" ${ARGN}) + + if(ADD_OPS_COMPILE_OPTION_V2) + set(OP_DEBUG_CONFIG_OPTION --opc-config-file ${ASCEND_CUSTOM_OPC_OPTIONS}) + else() + if(OP_DEBUG_CONFIG) + set(OP_DEBUG_CONFIG_OPTION --op-debug-config ${OP_DEBUG_CONFIG}) + endif() + set(OP_TILING_KEY_OPTION --tiling-keys ${ASCEND_CUSTOM_TILING_KEYS}) + endif() + + set(_OUT_DIR ${ASCEND_BINARY_OUT_DIR}/${CMD_COMPUTE_UNIT}) + set(GEN_OUT_DIR ${_OUT_DIR}/gen) + set(COMPILE_CMD_TARGET generate_compile_cmd_${CMD_COMPUTE_UNIT}) + + set(SED_SCRIPT ${CMAKE_CURRENT_SOURCE_DIR}/cmake/scripts/fix_format.sh) + + add_custom_target(${COMPILE_CMD_TARGET} ALL + COMMAND ${CMAKE_COMMAND} -E make_directory ${GEN_OUT_DIR} + COMMAND ${HI_PYTHON} ${ASCENDC_CMAKE_UTIL_DIR}/ascendc_bin_param_build.py + ${base_aclnn_binary_dir}/aic-${CMD_COMPUTE_UNIT}-ops-info.ini + ${GEN_OUT_DIR} + ${CMD_COMPUTE_UNIT} + ${OP_TILING_KEY_OPTION} + ${OP_DEBUG_CONFIG_OPTION} + COMMAND ${HI_PYTHON} ${ASCENDC_CMAKE_UTIL_DIR}/ascendc_bin_param_build.py + ${base_aclnn_binary_dir}/inner/aic-${CMD_COMPUTE_UNIT}-ops-info.ini + ${GEN_OUT_DIR} + ${CMD_COMPUTE_UNIT} + ${OP_TILING_KEY_OPTION} + ${OP_DEBUG_CONFIG_OPTION} + COMMAND ${HI_PYTHON} ${ASCENDC_CMAKE_UTIL_DIR}/ascendc_bin_param_build.py + ${base_aclnn_binary_dir}/exc/aic-${CMD_COMPUTE_UNIT}-ops-info.ini + ${GEN_OUT_DIR} + ${CMD_COMPUTE_UNIT} + ${OP_TILING_KEY_OPTION} + ${OP_DEBUG_CONFIG_OPTION} + COMMAND bash ${SED_SCRIPT} ${GEN_OUT_DIR} + ) + + add_dependencies(${COMPILE_CMD_TARGET} opbuild_gen_default opbuild_gen_inner opbuild_gen_exc) + add_dependencies(generate_compile_cmd ${COMPILE_CMD_TARGET}) +endfunction() + +function(add_ops_info_target) + cmake_parse_arguments(OPINFO "" "COMPUTE_UNIT" "" ${ARGN}) + + set(OPS_INFO_TARGET generate_ops_info_${OPINFO_COMPUTE_UNIT}) + if (ENABLE_BUILT_IN) + set(OPS_INFO_JSON ${ASCEND_AUTOGEN_DIR}/aic-${OPINFO_COMPUTE_UNIT}-ops-info-transformer.json) + else() + set(OPS_INFO_JSON ${ASCEND_AUTOGEN_DIR}/aic-${OPINFO_COMPUTE_UNIT}-ops-info.json) + endif() + set(CUSTOM_OPS_INFO_DIR ${CUSTOM_DIR}/op_impl/ai_core/tbe/config/${OPINFO_COMPUTE_UNIT}) + + set(OPS_INFO_INI ${base_aclnn_binary_dir}/aic-${OPINFO_COMPUTE_UNIT}-ops-info.ini) + set(OPS_INFO_INNER_INI ${base_aclnn_binary_dir}/inner/aic-${OPINFO_COMPUTE_UNIT}-ops-info.ini) + set(OPS_INFO_EXCLUDE_INI ${base_aclnn_binary_dir}/exc/aic-${OPINFO_COMPUTE_UNIT}-ops-info.ini) + + add_custom_command(OUTPUT ${OPS_INFO_JSON} + COMMAND ${HI_PYTHON} ${ASCENDC_CMAKE_UTIL_DIR}/parse_ini_to_json.py + ${OPS_INFO_INI} + ${OPS_INFO_INNER_INI} + ${OPS_INFO_EXCLUDE_INI} + ${OPS_INFO_JSON} + COMMAND mkdir -p ${CUSTOM_OPS_INFO_DIR} + COMMAND cp -f ${OPS_INFO_JSON} ${CUSTOM_OPS_INFO_DIR} + ) + + add_custom_target(${OPS_INFO_TARGET} ALL + DEPENDS ${OPS_INFO_JSON} + ) + + add_dependencies(${OPS_INFO_TARGET} opbuild_gen_default opbuild_gen_inner opbuild_gen_exc) + add_dependencies(generate_ops_info ${OPS_INFO_TARGET}) + + if (ENABLE_BUILT_IN) + install(FILES ${OPS_INFO_JSON} + DESTINATION ops_transformer/built-in/op_impl/ai_core/tbe/config/${OPINFO_COMPUTE_UNIT} OPTIONAL + ) + else() + install(FILES ${OPS_INFO_JSON} + DESTINATION packages/vendors/${VENDOR_NAME}_transformer/op_impl/ai_core/tbe/config/${OPINFO_COMPUTE_UNIT} OPTIONAL + ) + endif() +endfunction() + +function(add_ops_compile_options) + cmake_parse_arguments(OP_COMPILE "" "OP_NAME" "COMPUTE_UNIT;OPTIONS" ${ARGN}) + + if(NOT OP_COMPILE_OPTIONS) + return() + endif() + + if(ADD_OPS_COMPILE_OPTION_V2) + execute_process(COMMAND ${HI_PYTHON} ${ASCENDC_CMAKE_UTIL_DIR}/ascendc_gen_options.py + ${ASCEND_CUSTOM_OPTIONS} ${OP_COMPILE_OP_NAME} + ${OP_COMPILE_COMPUTE_UNIT} ${OP_COMPILE_OPTIONS} + RESULT_VARIABLE EXEC_RESULT + OUTPUT_VARIABLE EXEC_INFO + ERROR_VARIABLE EXEC_ERROR) + if (EXEC_RESULT) + message("add ops compile options info: ${EXEC_INFO}") + message("add ops compile options error: ${EXEC_ERROR}") + message(FATAL_ERROR "Error: add ops compile options failed!") + endif () + else() + file(APPEND ${ASCEND_CUSTOM_OPTIONS} + "${OP_COMPILE_OP_NAME},${OP_COMPILE_COMPUTE_UNIT},${OP_COMPILE_OPTIONS}\n" + ) + endif() +endfunction() + +function(add_ops_tiling_keys) + cmake_parse_arguments(OP_COMPILE "" "OP_NAME" "COMPUTE_UNIT;TILING_KEYS" ${ARGN}) + + if(NOT OP_COMPILE_TILING_KEYS) + return() + endif() + + if(ADD_OPS_COMPILE_OPTION_V2) + list(JOIN OP_COMPILE_TILING_KEYS "," STRING_TILING_KEYS) + add_ops_compile_options( + OP_NAME ${OP_COMPILE_OP_NAME} + OPTIONS --tiling_key=${STRING_TILING_KEYS} + ) + else() + file(APPEND ${ASCEND_CUSTOM_TILING_KEYS} + "${OP_COMPILE_OP_NAME},${OP_COMPILE_COMPUTE_UNIT},${OP_COMPILE_TILING_KEYS}\n" + ) + endif() +endfunction() + +function(add_opc_config) + cmake_parse_arguments(OP_COMPILE "" "OP_NAME" "COMPUTE_UNIT;CONFIG" ${ARGN}) + + if(NOT ADD_OPS_COMPILE_OPTION_V2) + return() + endif() + + set(_OPC_CONFIG) + + if(NOT OP_COMPILE_CONFIG) + list(APPEND _OPC_CONFIG "-DNOT_DYNAMIC_COMPILE") + else() + string(REPLACE "," ";" OP_COMPILE_CONFIG_LIST "${OP_COMPILE_CONFIG}") + list(APPEND _OPC_CONFIG "-DNOT_DYNAMIC_COMPILE") + + foreach(_option ${OP_COMPILE_CONFIG_LIST}) + if("${_option}" STREQUAL "ccec_g") + list(APPEND _OPC_CONFIG "-g") + elseif("${_option}" STREQUAL "ccec_O0") + list(APPEND _OPC_CONFIG "-O0") + elseif("${_option}" STREQUAL "sanitizer") + list(APPEND _OPC_CONFIG "-sanitizer") + elseif("${_option}" STREQUAL "dump_cce") + list(APPEND _OPC_CONFIG "--save-temp-files") + endif() + endforeach() + endif() + + if(ENABLE_OOM) + list(APPEND _OPC_CONFIG "--oom") + list(APPEND _OPC_CONFIG "-ffunction-sections -fdata-sections") + endif() + + if(_OPC_CONFIG) + add_ops_compile_options( + OP_NAME ${OP_COMPILE_OP_NAME} + OPTIONS ${_OPC_CONFIG} + ) + endif() +endfunction() + +function(add_ops_src_copy) + cmake_parse_arguments(SRC_COPY "" "TARGET_NAME;SRC;DST;BE_RELIED;COMPUTE_UNIT" "" ${ARGN}) + + set(OPS_UTILS_INC_KERNEL_TARGET ops_utils_inc_kernel_${SRC_COPY_COMPUTE_UNIT}) + if (EXISTS ${OPS_ADV_UTILS_KERNEL_INC}) + if (NOT TARGET ${OPS_UTILS_INC_KERNEL_TARGET}) + get_filename_component(_ROOT_OPS_SRC_DIR "${SRC_COPY_DST}" DIRECTORY) + set(OPS_UTILS_INC_KERNEL_DIR ${_ROOT_OPS_SRC_DIR}/ascendc/common) + add_custom_command(OUTPUT ${OPS_UTILS_INC_KERNEL_DIR} + COMMAND mkdir -p ${OPS_UTILS_INC_KERNEL_DIR}/regbase + COMMAND cp -rf ${OPS_ADV_UTILS_KERNEL_INC}/*.* ${OPS_UTILS_INC_KERNEL_DIR} + ) + + add_custom_target(${OPS_UTILS_INC_KERNEL_TARGET} + DEPENDS ${OPS_UTILS_INC_KERNEL_DIR} + ) + endif () + endif () + + # set(MC2_OPS_LIST "matmul_reduce_scatter;" + # "grouped_mat_mul_allto_allv;" + # "grouped_mat_mul_all_reduce;" + # "batch_mat_mul_reduce_scatter_allto_all;" + # "allto_allv_grouped_mat_mul;" + # "allto_all_all_gather_batch_mat_mul;" + # "distribute_barrier;" + # "moe_distribute_combine_add_rms_norm;" + # "moe_distribute_dispatch;" + # "moe_distribute_combine;" + # "moe_distribute_dispatch_v2;" + # "moe_distribute_combine_v2;" + # "moe_update_expert;" + # "all_gather_matmul;" + # "matmul_all_reduce;" + # "matmul_all_reduce_add_rms_norm;" + # "inplace_matmul_all_reduce_add_rms_norm;" + # "attention_to_ffn;" + # "ffn_to_attention;" + # ) # mc2算子列表 + set(MC2_OPS_LIST "" + ) # mc2算子列表 + + get_filename_component(FOLDER_NAME "${SRC_COPY_DST}" NAME_WE) + list(FIND MC2_OPS_LIST "${FOLDER_NAME}" INDEX) + if(NOT INDEX EQUAL -1) + set(BELONG_MC2_OPS TRUE) + endif() + + if(NOT BUILD_OPS_RTY_KERNEL AND BELONG_MC2_OPS) + file(GLOB SRC_FILES ${SRC_COPY_SRC}/* ${SRC_COPY_SRC}/op_kernel/*) + else() + file(GLOB SRC_FILES ${SRC_COPY_SRC}/*) + endif() + list(FILTER SRC_FILES EXCLUDE REGEX "op_host") + + get_filename_component(PARENT_PTH "${SRC_COPY_SRC}" DIRECTORY) + get_filename_component(CUR_NAME "${SRC_COPY_SRC}" NAME) + get_filename_component(PARENT_NAME "${PARENT_PTH}" NAME) + + set(DOING_TARGET_NAME ${SRC_COPY_TARGET_NAME}) + if(${CUR_NAME} STREQUAL "common") + set(DOING_TARGET_NAME ${PARENT_NAME}_${DOING_TARGET_NAME}) + endif() + + if (NOT TARGET ${DOING_TARGET_NAME}) + set(_BUILD_FLAG ${SRC_COPY_DST}/${DOING_TARGET_NAME}.done) + if (NOT BUILD_OPS_RTY_KERNEL AND BELONG_MC2_OPS) + add_custom_command(OUTPUT ${_BUILD_FLAG} + COMMAND mkdir -p ${SRC_COPY_DST} + COMMAND cp -rf ${SRC_FILES} ${SRC_COPY_DST} + COMMAND rm -rf ${SRC_COPY_DST}/op_kernel/ + COMMAND touch ${_BUILD_FLAG} + ) + else() + add_custom_command(OUTPUT ${_BUILD_FLAG} + COMMAND mkdir -p ${SRC_COPY_DST} + COMMAND cp -rf ${SRC_FILES} ${SRC_COPY_DST} + COMMAND touch ${_BUILD_FLAG} + ) + endif() + + add_custom_target(${DOING_TARGET_NAME} + DEPENDS ${_BUILD_FLAG} + ) + endif () + + if (TARGET ${OPS_UTILS_INC_KERNEL_TARGET}) + add_dependencies(${DOING_TARGET_NAME} ${OPS_UTILS_INC_KERNEL_TARGET}) + endif () + + if (DEFINED SRC_COPY_BE_RELIED) + add_dependencies(${SRC_COPY_BE_RELIED} ${DOING_TARGET_NAME}) + endif () + +endfunction() + +function(add_bin_compile_target) + cmake_parse_arguments(BINARY "" "COMPUTE_UNIT" "OP_INFO" ${ARGN}) + + if (ENABLE_BUILT_IN) + set(_INSTALL_DIR ops_transformer/built-in/op_impl/ai_core/tbe/kernel) + else() + set(_INSTALL_DIR packages/vendors/${VENDOR_NAME}_transformer/op_impl/ai_core/tbe/kernel) + endif() + set(_OUT_DIR ${ASCEND_BINARY_OUT_DIR}/${BINARY_COMPUTE_UNIT}) + + set(BIN_OUT_DIR ${_OUT_DIR}/bin) + set(GEN_OUT_DIR ${_OUT_DIR}/gen) + set(SRC_OUT_DIR ${_OUT_DIR}/src) + file(MAKE_DIRECTORY ${BIN_OUT_DIR}) + + foreach(_op_info ${BINARY_OP_INFO}) + get_filename_component(_op_name "${_op_info}" NAME) + set(${_op_name}_dir ${_op_info}) + set(${_op_name}_apt_dir ${_op_info}) + endforeach() + + set(_ops_target_list) + set(compile_scripts) + file(GLOB scripts_list ${GEN_OUT_DIR}/*.sh) + list(APPEND compile_scripts ${scripts_list}) + + foreach(bin_script ${compile_scripts}) + get_filename_component(bin_file ${bin_script} NAME_WE) + string(REPLACE "-" ";" bin_sep ${bin_file}) + list(GET bin_sep 0 op_type) + list(GET bin_sep 1 op_file) + list(GET bin_sep 2 op_index) + + if (NOT DEFINED ${op_file}_dir) + continue() + endif () + + if (NOT TARGET ${op_file}) + add_custom_target(${op_file}) + add_dependencies(ops_transformer_kernel ${op_file}) + endif () + + set(OP_TARGET_NAME ${op_file}_${BINARY_COMPUTE_UNIT}) + + if (NOT TARGET ${OP_TARGET_NAME}) + add_custom_target(${OP_TARGET_NAME}) + add_dependencies(${op_file} ${OP_TARGET_NAME}) + list(APPEND _ops_target_list ${OP_TARGET_NAME}) + + set(OP_SRC_OUT_DIR ${SRC_OUT_DIR}/${op_file}) + set(OP_BIN_OUT_DIR ${BIN_OUT_DIR}/${op_file}) + file(MAKE_DIRECTORY ${OP_SRC_OUT_DIR}) + + add_ops_src_copy( + TARGET_NAME + ${OP_TARGET_NAME}_src_copy + SRC + ${${op_file}_dir} + DST + ${OP_SRC_OUT_DIR} + COMPUTE_UNIT + ${BINARY_COMPUTE_UNIT} + ) + + if (DEFINED ${op_file}_depends) + foreach(depend_info ${${op_file}_depends}) + get_filename_component(_depend_op_name "${depend_info}" NAME) + set(_depend_op_target ${_depend_op_name}_${BINARY_COMPUTE_UNIT}_src_copy) + add_ops_src_copy( + TARGET_NAME + ${_depend_op_target} + SRC + ${CMAKE_SOURCE_DIR}/${depend_info} + DST + ${SRC_OUT_DIR}/${_depend_op_name} + COMPUTE_UNIT + ${BINARY_COMPUTE_UNIT} + BE_RELIED + ${OP_TARGET_NAME}_src_copy + ) + endforeach() + endif () + + set(DYNAMIC_PY_FILE ${OP_SRC_OUT_DIR}/${op_type}.py) + add_custom_command(OUTPUT ${DYNAMIC_PY_FILE} + COMMAND cp -rf ${ASCEND_IMPL_OUT_DIR}/dynamic/${op_file}.py ${DYNAMIC_PY_FILE} + # COMMAND bash ${CMAKE_CURRENT_SOURCE_DIR}/cmake/scripts/update_get_kernel_source.sh ${DYNAMIC_PY_FILE} + ) + + add_custom_target(${OP_TARGET_NAME}_py_copy + DEPENDS ${DYNAMIC_PY_FILE} + ) + + add_custom_command(OUTPUT ${OP_BIN_OUT_DIR} + COMMAND mkdir -p ${OP_BIN_OUT_DIR} + ) + + add_custom_target(${OP_TARGET_NAME}_mkdir + DEPENDS ${OP_BIN_OUT_DIR} + ) + + if (ENABLE_BUILT_IN) + install(DIRECTORY ${OP_BIN_OUT_DIR} + DESTINATION ${_INSTALL_DIR}/${BINARY_COMPUTE_UNIT}/ops_transformer OPTIONAL + ) + install(FILES ${BIN_OUT_DIR}/${op_file}.json + DESTINATION ${_INSTALL_DIR}/config/${BINARY_COMPUTE_UNIT}/ops_transformer OPTIONAL + ) + else() + install(DIRECTORY ${OP_BIN_OUT_DIR} + DESTINATION ${_INSTALL_DIR}/${BINARY_COMPUTE_UNIT} OPTIONAL + ) + install(FILES ${BIN_OUT_DIR}/${op_file}.json + DESTINATION ${_INSTALL_DIR}/config/${BINARY_COMPUTE_UNIT} OPTIONAL + ) + endif() + endif () + + set(_group "1-0") + if (DEFINED ASCEND_OP_NAME AND NOT "${ASCEND_OP_NAME}" STREQUAL "") + if (NOT "${ASCEND_OP_NAME}" STREQUAL "all" AND NOT "${ASCEND_OP_NAME}" STREQUAL "ALL") + if (${op_file} IN_LIST ASCEND_OP_NAME) + list(LENGTH ASCEND_OP_NAME _len) + list(FIND ASCEND_OP_NAME ${op_file} _index) + math(EXPR _next_index "${_index} + 1") + if (${_next_index} LESS ${_len}) + list(GET ASCEND_OP_NAME ${_next_index} _group_str) + set(_regex "^[0-9]+-[0-9]+$") + string(REGEX MATCH "${_regex}" match "${_group_str}") + if (match) + set(_group ${_group_str}) + endif () + endif () + endif () + endif () + endif () + + string(REPLACE "-" ";" _group_sep ${_group}) + + list(GET _group_sep 1 start_index) + set(end_index ${op_index}) + list(GET _group_sep 0 step) + + set(_compile_flag false) + if (${start_index} LESS ${end_index}) + foreach(i RANGE ${start_index} ${end_index} ${step}) + if (${i} EQUAL ${end_index}) + set(_compile_flag true) + break() + endif () + endforeach() + elseif (${start_index} EQUAL ${end_index}) + set(_compile_flag true) + else() + set(_compile_flag false) + endif () + + if (_compile_flag) + set(_BUILD_COMMAND) + set(_BUILD_FLAG ${GEN_OUT_DIR}/${OP_TARGET_NAME}_${op_index}.done) + if (ENABLE_OPS_HOST OR ENABLE_HOST_TILING) + list(APPEND _BUILD_COMMAND export ASCEND_CUSTOM_OPP_PATH=${CUSTOM_DIR} &&) + endif () + list(APPEND _BUILD_COMMAND export HI_PYTHON="python3" &&) + list(APPEND _BUILD_COMMAND export TILINGKEY_PAR_COMPILE=1 &&) + list(APPEND _BUILD_COMMAND export BIN_FILENAME_HASHED=1 &&) + list(APPEND _BUILD_COMMAND bash ${bin_script} ${OP_SRC_OUT_DIR}/${op_type}.py ${OP_BIN_OUT_DIR}) + if(CMAKE_GENERATOR MATCHES "Unix Makefiles") + list(APPEND _BUILD_COMMAND && echo $(MAKE)) + endif() + + add_custom_command(OUTPUT ${_BUILD_FLAG} + COMMAND ${_BUILD_COMMAND} + COMMAND touch ${_BUILD_FLAG} + WORKING_DIRECTORY ${GEN_OUT_DIR} + ) + + add_custom_target(${OP_TARGET_NAME}_${op_index} + DEPENDS ${_BUILD_FLAG} + ) + + if (ENABLE_OPS_HOST OR ENABLE_HOST_TILING) + add_dependencies(${OP_TARGET_NAME}_${op_index} optiling_compat generate_ops_info) + endif () + add_dependencies(${OP_TARGET_NAME}_${op_index} ${OP_TARGET_NAME}_src_copy ${OP_TARGET_NAME}_py_copy ${OP_TARGET_NAME}_mkdir) + add_dependencies(${OP_TARGET_NAME} ${OP_TARGET_NAME}_${op_index}) + endif () + endforeach() + + if (_ops_target_list) + set(OPS_CONFIG_TARGET ops_config_${BINARY_COMPUTE_UNIT}) + set(BINARY_INFO_CONFIG_FILE ${BIN_OUT_DIR}/binary_info_config.json) + set(RELOCATABLE_KERNEL_INFO_CONFIG_FILE ${BIN_OUT_DIR}/relocatable_kernel_info_config.json) + + add_custom_command(OUTPUT ${BINARY_INFO_CONFIG_FILE} + COMMAND ${HI_PYTHON} ${ASCENDC_CMAKE_UTIL_DIR}/ascendc_ops_config.py -p ${BIN_OUT_DIR} -s ${BINARY_COMPUTE_UNIT} + ) + + add_custom_target(${OPS_CONFIG_TARGET} + DEPENDS ${BINARY_INFO_CONFIG_FILE} + ) + + add_dependencies(ops_transformer_config ${OPS_CONFIG_TARGET}) + + foreach(_op_target ${_ops_target_list}) + add_dependencies(${OPS_CONFIG_TARGET} ${_op_target}) + endforeach() + + if (ENABLE_BUILT_IN) + install(FILES ${BINARY_INFO_CONFIG_FILE} + DESTINATION ${_INSTALL_DIR}/config/${BINARY_COMPUTE_UNIT}/ops_transformer OPTIONAL + ) + install(FILES ${RELOCATABLE_KERNEL_INFO_CONFIG_FILE} + DESTINATION ${_INSTALL_DIR}/config/${BINARY_COMPUTE_UNIT}/ops_transformer OPTIONAL + ) + else() + install(FILES ${RELOCATABLE_KERNEL_INFO_CONFIG_FILE} + DESTINATION ${_INSTALL_DIR}/config/${BINARY_COMPUTE_UNIT} OPTIONAL + ) + install(FILES ${BINARY_INFO_CONFIG_FILE} + DESTINATION ${_INSTALL_DIR}/config/${BINARY_COMPUTE_UNIT} OPTIONAL + ) + endif() + endif () +endfunction() + +function(redefine_file_macro) + cmake_parse_arguments(_FILE "" "" "TARGET_NAME" ${ARGN}) + + foreach(_target_name ${_FILE_TARGET_NAME}) + target_compile_options(${_target_name} PRIVATE + -Wno-builtin-macro-redefined + ) + + get_target_property(_srcs ${_target_name} SOURCES) + + foreach(_src ${_srcs}) + get_filename_component(_src_name "${_src}" NAME) + set_source_files_properties(${_src} + PROPERTIES COMPILE_DEFINITIONS __FILE__="${_src_name}" + ) + endforeach() + endforeach() +endfunction() + +function(add_static_ops) + cmake_parse_arguments(STATIC "" "SRC_DIR" "ACLNN_SRC;ACLNN_INNER_SRC" ${ARGN}) + set(prepare_ops_adv_static_target prepare_ops_transformer_static) + set(static_src_temp_dir ${CMAKE_CURRENT_BINARY_DIR}/static_src_temp_dir) + set(modified_files) + foreach(ops_type ${OPS_STATIC_TYPES}) + get_target_property(all_srcs aclnn_ops_${ops_type} SOURCES) + set(add_srcs) + set(generate_aclnn_srcs) + foreach(_src ${all_srcs}) + string(REGEX MATCH "^${STATIC_SRC_DIR}" is_match "${_src}") + if (is_match) + list(APPEND add_srcs ${_src}) + endif () + endforeach() + + foreach(_src ${add_srcs}) + get_filename_component(name_without_ext ${_src} NAME_WE) + string(REGEX REPLACE "^aclnn_" "" _op_name ${name_without_ext}) + + foreach(_aclnn_src ${STATIC_ACLNN_SRC}) + get_filename_component(aclnn_name ${_aclnn_src} NAME_WE) + if("aclnn_${_op_name}" STREQUAL "${aclnn_name}") + list(APPEND generate_aclnn_srcs ${_aclnn_src}) + break() + endif() + endforeach() + + foreach(_aclnn_inner_src ${STATIC_ACLNN_INNER_SRC}) + get_filename_component(aclnn_inner_name ${_aclnn_inner_src} NAME_WE) + if("aclnnInner_${_op_name}" STREQUAL "${aclnn_inner_name}") + list(APPEND generate_aclnn_srcs ${_aclnn_inner_src}) + break() + endif() + endforeach() + endforeach() + + if(add_srcs) + list(TRANSFORM add_srcs REPLACE "${STATIC_SRC_DIR}" "${static_src_temp_dir}" OUTPUT_VARIABLE add_static_srcs) + list(APPEND modified_files ${add_static_srcs}) + set(aclnn_ops_static_target aclnn_ops_${ops_type}_static) + set_source_files_properties(${add_static_srcs} + TARGET_DIRECTORY ${aclnn_ops_static_target} + PROPERTIES GENERATED TRUE + ) + + target_sources(${aclnn_ops_static_target} PRIVATE + ${add_static_srcs} + ) + add_dependencies(${aclnn_ops_static_target} ${prepare_ops_adv_static_target}) + endif() + + if(generate_aclnn_srcs) + list(REMOVE_DUPLICATES generate_aclnn_srcs) + set(aclnn_op_target acl_op_${ops_type}_builtin) + set_source_files_properties(${generate_aclnn_srcs} + TARGET_DIRECTORY ${aclnn_op_target} + PROPERTIES GENERATED TRUE + ) + + target_sources(${aclnn_op_target} PRIVATE + ${generate_aclnn_srcs} + ) + endif() + endforeach() + + if(NOT TARGET ${prepare_ops_adv_static_target}) + list(REMOVE_DUPLICATES modified_files) + add_custom_command(OUTPUT ${static_src_temp_dir} + COMMAND mkdir -p ${static_src_temp_dir} + COMMAND cp -rf ${STATIC_SRC_DIR}/gmm ${STATIC_SRC_DIR}/mc2 ${STATIC_SRC_DIR}/attention ${static_src_temp_dir} || true + COMMAND ${HI_PYTHON} -B ${OPS_STATIC_SCRIPT} InsertIni -p ${static_src_temp_dir} -f ${modified_files} + ) + + add_custom_target(${prepare_ops_adv_static_target} + DEPENDS ${static_src_temp_dir} + ) + endif() +endfunction() +function(add_aicpu_kernel_modules) + if(NOT TARGET ${OPHOST_NAME}_aicpu_obj) + add_library(${OPHOST_NAME}_aicpu_obj OBJECT) + target_include_directories(${OPHOST_NAME}_aicpu_obj PRIVATE ${AICPU_INCLUDE}) + target_compile_definitions( + ${OPHOST_NAME}_aicpu_obj PRIVATE _FORTIFY_SOURCE=2 google=ascend_private + $<$:ASCEND_AICPU_UT> + ) + target_compile_options( + ${OPHOST_NAME}_aicpu_obj PRIVATE $<$>:-DDISABLE_COMPILE_V1> -Dgoogle=ascend_private + -fvisibility=hidden ${AICPU_DEFINITIONS} + ) + target_link_libraries( + ${OPHOST_NAME}_aicpu_obj + PRIVATE $,intf_llt_pub_asan_cxx17,intf_pub_cxx17>> + $ + ) + endif() +endfunction() + +function(add_aicpu_cust_kernel_modules op_name aicpu_sources aicpu_jsons) + set(target_name ${op_name}_obj) + if(NOT TARGET ${target_name}) + add_library(${target_name} OBJECT) + target_include_directories(${target_name} PRIVATE ${AICPU_INCLUDE}) + target_compile_definitions( + ${target_name} PRIVATE + _FORTIFY_SOURCE=2 _GLIBCXX_USE_CXX11_ABI=1 + google=ascend_private + $<$:ASCEND_AICPU_UT> + ) + target_compile_options( + ${target_name} PRIVATE + $<$>:-DDISABLE_COMPILE_V1> -Dgoogle=ascend_private + -fvisibility=hidden ${AICPU_DEFINITIONS} + ) + target_link_libraries( + ${target_name} + PRIVATE $,intf_llt_pub_asan_cxx17,intf_pub_cxx17>> + $ + -Wl,--no-whole-archive + ) + if (NOT (UT_TEST_ALL OR OP_KERNEL_AICPU_UT)) + set_property(TARGET ${target_name} PROPERTY + CXX_COMPILER_LAUNCHER ${ASCEND_DIR}/toolkit/toolchain/hcc/bin/aarch64-target-linux-gnu-g++) + endif() + target_sources(${target_name} PRIVATE ${aicpu_sources}) + set_property(GLOBAL APPEND PROPERTY AICPU_JSON_FILES ${aicpu_jsons}) + if (NOT ${target_name} IN_LIST AICPU_CUST_OBJ_TARGETS) + set(AICPU_CUST_OBJ_TARGETS ${AICPU_CUST_OBJ_TARGETS} ${target_name} CACHE INTERNAL "All aicpu cust obj targets") + endif() + endif() +endfunction() + +# 添加待编译算子 +function(add_need_compile_ops op_name) + if(NOT ASCEND_OP_NAME) + # 为空则不需要更新 + return() + endif() + + set(NEW_OP_NAMES ${ASCEND_OP_NAME} ${op_name}) + list(REMOVE_DUPLICATES NEW_OP_NAMES) + set(ASCEND_OP_NAME + ${NEW_OP_NAMES} + CACHE STRING "Ascend op names to compile" FORCE + ) +endfunction() + +function(add_dependent_ops dependent_ops) + foreach(dep_op ${dependent_ops}) + # 查询依赖算子所在目录 + set(dep_op_path_list "") + if(ENABLE_EXPERIMENTAL) + foreach(ops_category ${OPS_CATEGORY_LIST}) + file(GLOB dep_op_path "${PROJECT_SOURCE_DIR}/experimental/${ops_category}/${dep_op}") + list(APPEND dep_op_path_list ${dep_op_path}) + endforeach() + endif() + set(outside_experimental FALSE) + # 如果非 experimental，或 experimental 下没找到，则去常规目录查找 + if(NOT dep_op_path_list) + foreach(ops_category ${OPS_CATEGORY_LIST}) + file(GLOB dep_op_path "${PROJECT_SOURCE_DIR}/${ops_category}/${dep_op}") + list(APPEND dep_op_path_list ${dep_op_path}) + endforeach() + set(outside_experimental ${ENABLE_EXPERIMENTAL}) + endif() + # 检查依赖存在 + if(NOT dep_op_path_list) + message(FATAL_ERROR "dependent operator(${dep_op}) not exists") + endif() + list(LENGTH ${dep_op_path_list} find_dep_ops_count) + if(find_dep_ops_count GREATER 1) + message(FATAL_ERROR "dependent operator(${dep_op}) is not unique, the found operators:${dep_op_path_list}") + endif() + + # ASCEND_OP_NAME 为空表示全部编译，则不需要特意添加目录；但指定experimental时未扫描常规算子，如果依赖常规算子，需要添加目录 + # 已在待编译列表，则不需要加入 + # 已在编译列表，则不需要重复加入 + + if((ASCEND_OP_NAME OR outside_experimental) + AND (NOT (dep_op IN_LIST ASCEND_OP_NAME)) + AND (NOT (dep_op IN_LIST COMPILED_OPS)) + ) + # 加入依赖并去重 + add_need_compile_ops("${dep_op}") + # 添加目录 + get_filename_component(dep_op_path "${dep_op_path_list}" ABSOLUTE) + get_filename_component(dep_op_parent "${dep_op_path}" DIRECTORY) + get_filename_component(parent_path "${dep_op_parent}" NAME) + message(STATUS "add dependent operator: ${parent_path}/${dep_op}, path: ${dep_op_path}") + add_subdirectory("${dep_op_path}" "${CMAKE_BINARY_DIR}/dependent-ops/${parent_path}/${dep_op}") + endif() + endforeach() +endfunction() + +if (BUILD_OPEN_PROJECT) + if (TESTS_UT_OPS_TEST) + include(${OPS_ADV_CMAKE_DIR}/func_utest.cmake) + endif () + if (TESTS_EXAMPLE_OPS_TEST) + include(${OPS_ADV_CMAKE_DIR}/func_examples.cmake) + endif () +endif () diff --git a/csrc/ascend/cmake/func_examples.cmake b/csrc/ascend/cmake/func_examples.cmake new file mode 100644 index 000000000..1caa8d8d9 --- /dev/null +++ b/csrc/ascend/cmake/func_examples.cmake @@ -0,0 +1,69 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +function(add_execute_example) + cmake_parse_arguments(EXAMPLE "" "TARGET_NAME;SCRIPT;TEST_CASE;ACLNN_FUNC" "" ${ARGN}) + + # 获取需使用的 opapi 动态库绝对路径 + if (NOT EXAMPLE_ACLNN_FUNC) + message(FATAL_ERROR "Example(${EXAMPLE_TEST_CASE}) not give aclnn func name") + endif () + set(_get_py ${OPS_ADV_DIR}/cmake/scripts/examples/get_opapi_abs_path.py) + execute_process( + COMMAND ${HI_PYTHON} ${_get_py} "-f=${EXAMPLE_ACLNN_FUNC}" + OUTPUT_VARIABLE OPAPI_SHARED_REL_PATH + ) + message(STATUS "Example(${EXAMPLE_TEST_CASE}) Func(${EXAMPLE_ACLNN_FUNC}) use ${OPAPI_SHARED_REL_PATH}") + if (NOT OPAPI_SHARED_REL_PATH) + message(FATAL_ERROR "Example(${EXAMPLE_TEST_CASE}) can't get opapi path") + endif () + get_filename_component(OPAPI_SHARED_REL_DIR ${OPAPI_SHARED_REL_PATH} DIRECTORY) + + target_link_libraries(${EXAMPLE_TARGET_NAME} + PRIVATE + $ + ${OPAPI_SHARED_REL_PATH} + -lascendcl + -lnnopbase + -lc_sec + ) + target_compile_options(${EXAMPLE_TARGET_NAME} + PRIVATE + $<$:-std=gnu++1z> + ) + set_target_properties(${EXAMPLE_TARGET_NAME} + PROPERTIES + RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/${EXAMPLE_TEST_CASE} + BUILD_RPATH "${OPAPI_SHARED_REL_DIR}" + ) + + set(_execute_flag OFF) + if ("${TESTS_EXAMPLE_OPS_TEST}" STREQUAL "") + return() + elseif ("ALL" IN_LIST TESTS_EXAMPLE_OPS_TEST OR "all" IN_LIST TESTS_EXAMPLE_OPS_TEST) + set(_execute_flag ON) + elseif ("${EXAMPLE_TARGET_NAME}" IN_LIST TESTS_EXAMPLE_OPS_TEST) + set(_execute_flag ON) + endif () + + if (_execute_flag) + add_custom_command( + TARGET ${EXAMPLE_TARGET_NAME} POST_BUILD + COMMAND bash ${EXAMPLE_SCRIPT} ${EXAMPLE_TARGET_NAME} $ ${EXAMPLE_TEST_CASE} + WORKING_DIRECTORY $ + COMMENT "Run ${EXAMPLE_TARGET_NAME}" + ) + endif () + + if (NOT TARGET ops_test_example) + add_custom_target(ops_test_example) + endif () + add_dependencies(ops_test_example ${EXAMPLE_TARGET_NAME}) +endfunction() diff --git a/csrc/ascend/cmake/func_utest.cmake b/csrc/ascend/cmake/func_utest.cmake new file mode 100644 index 000000000..80c5b84a6 --- /dev/null +++ b/csrc/ascend/cmake/func_utest.cmake @@ -0,0 +1,1198 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +######################################################################################################################## +# 预定义变量 +######################################################################################################################## + +# 缓存所有算子 UTest 场景 OpApi 动态库相关信息 +set(_OpsTestUt_OpApiSources "" CACHE INTERNAL "" FORCE) # Sources +set(_OpsTestUt_OpApiPrivateIncludesExt "" CACHE INTERNAL "" FORCE) # PrivateIncludesExt +set(_OpsTestUt_OpApiLinkLibrariesExt "" CACHE INTERNAL "" FORCE) # LinkLibrariesExt +add_library(_OpsTestUt_OpApi_Wno INTERFACE) +target_compile_options(_OpsTestUt_OpApi_Wno + INTERFACE + $<$:-Wno-mismatched-tags> + $<$:-Wno-format-signedness> + -Wno-extra + -Wno-redundant-decls +) + +add_definitions(-DNOT_DYNAMIC_COMPILE) + +# 缓存所有算子 UTest 场景 OpProto 动态库相关信息 +set(_OpsTestUt_OpProtoSources "" CACHE INTERNAL "" FORCE) # Sources +set(_OpsTestUt_OpProtoPrivateIncludesExt "" CACHE INTERNAL "" FORCE) # PrivateIncludesExt +set(_OpsTestUt_OpProtoLinkLibrariesExt "" CACHE INTERNAL "" FORCE) # LinkLibrariesExt + +# 缓存所有算子 UTest 场景 OpTiling 动态库相关信息 +set(_OpsTestUt_OpTilingSources "" CACHE INTERNAL "" FORCE) # Sources +set(_OpsTestUt_OpTilingPrivateIncludesExt "" CACHE INTERNAL "" FORCE) # PrivateIncludesExt +set(_OpsTestUt_OpTilingLinkLibrariesExt "" CACHE INTERNAL "" FORCE) # LinkLibrariesExt +add_library(_OpsTestUt_OpTiling_Wno INTERFACE) +target_compile_options(_OpsTestUt_OpTiling_Wno + INTERFACE + $<$:-Wno-unused-private-field> + $<$:-Wno-redundant-move> + -Wno-shadow + $<$:-Wno-format-signedness> + -Wno-extra +) + +# 缓存当前算子 UTest 场景 OpKernel 目标 +set(_OpsTestUt_OpKernelLibraries "" CACHE INTERNAL "" FORCE) # LinkLibrariesExt +add_library(_OpsTestUt_OpKernel_Wno INTERFACE) +target_compile_options(_OpsTestUt_OpKernel_Wno + INTERFACE + -Wno-undef + -Wno-cast-qual + -Wno-shadow + -Wno-sign-compare + -Wno-unused-macros + -Wno-unused-variable + $<$:-Wno-unused-but-set-variable> + $<$:$<$:-Wno-unused-but-set-variable>> + $<$:-Wno-duplicated-branches> + -Wno-extra + -Wno-float-conversion + -Wno-parentheses +) +# 缓存当前算子 UTest 场景 UTestCommon 目标 +set(_OpsTestUt_UTestCommonLibrary "" CACHE INTERNAL "" FORCE) + +# 缓存所有算子 UTest 场景 UTest用例 动态库相关信息 +set(_OpsTestUt_UTestCaseLibraries "" CACHE INTERNAL "" FORCE) +set(_OpsTestUt_UTestAclnnCaseLibraries "" CACHE INTERNAL "" FORCE) +add_library(_OpsTestUt_UTestCaseStatic_Wno INTERFACE) +target_compile_options(_OpsTestUt_UTestCaseStatic_Wno + INTERFACE + $<$:-Wno-vla> +) + +# GTest 版本兼容性保证 +add_library(_OpsTestUt_GTest_Wno INTERFACE) +target_compile_options(_OpsTestUt_GTest_Wno + INTERFACE + -Wno-undef +) + +######################################################################################################################## +# 编译方法 +######################################################################################################################## + +# Level1, 添加算子 OpApi 动态库 +#[[ +调用参数: + one_value_keywords: + SUB_SYSTEM : 必选参数, 用于指定算子所属子系统, 如 transformer + BRIEF : 必选参数, 用于指定算子缩略名(建议以大驼峰命名, 与算子实际名称无强制对应关系), 如 Fag/Fas/Fa + SNAKE : 必选参数, 用于指定算子全名, 如 flash_attention_score_grad + multi_value_keywords: + SOURCES_EXT : 可选参数, 额外源文件 + PRIVATE_INCLUDES_EXT : 可选参数, 额外头文件搜索路径 + PRIVATE_LINK_LIBRARIES_EXT : 可选参数, 额外链接库 +备注说明: + 本函数提供编译算子对应 opapi.so 的功能. 下面介绍在调用本函数时所需了解的一些背景知识和注意事项. + 1. 本函数假设算子对应 OpApi 的源文件位于 'src/${SUB_SYSTEM}/${SNAKE}/op_host' 路径, 并自动添加以下内容: + a) 源文件: ${SNAKE}.cpp, aclnn_${SNAKE}.cpp + b) 头文件搜索路径: 'src/${SUB_SYSTEM}/${SNAKE}/op_host' + 2. 若算子 '需要额外的源文件’ 或 ‘源文件不满足上述约定的默认路径', 则可通过 SOURCES_EXT 参数指定 '额外的源文件'; + 3. 若算子 '需要额外头文件搜索路径’ 或 ‘本函数现有实现的头文件搜索路径设置不满足', + 则可通过 PRIVATE_INCLUDES_EXT 参数指定 '额外的头文件搜索路径'; 参数 PRIVATE_LINK_LIBRARIES_EXT 设置逻辑同理; +]] +function(OpsTest_Level1_AddOpApiShared) + cmake_parse_arguments( + TMP + "" + "SUB_SYSTEM;BRIEF;SNAKE" + "SOURCES_EXT;PRIVATE_INCLUDES_EXT;PRIVATE_LINK_LIBRARIES_EXT" + "" + ${ARGN} + ) + get_filename_component(_L0_Src "${OPS_ADV_DIR}/${TMP_SUB_SYSTEM}/${TMP_SNAKE}/op_host/${TMP_SNAKE}.cpp" REALPATH) + if (NOT EXISTS "${_L0_Src}") + set(_L0_Src) + endif () + get_filename_component(_L2_Src "${OPS_ADV_DIR}/${TMP_SUB_SYSTEM}/${TMP_SNAKE}/op_host/aclnn_${TMP_SNAKE}.cpp" REALPATH) + if (NOT EXISTS "${_L2_Src}") + set(_L2_Src) + endif () + get_filename_component(_Inc "${OPS_ADV_DIR}/${TMP_SUB_SYSTEM}/${TMP_SNAKE}/op_host/" REALPATH) + if (NOT EXISTS "${_Inc}") + set(_Inc) + endif () + + set(_Sources ${TMP_SOURCES_EXT} ${_L0_Src} ${_L2_Src}) + list(REMOVE_DUPLICATES _Sources) + if (_Sources) + set(_PrivateIncludeDirs ${TMP_PRIVATE_INCLUDES_EXT} ${_Inc}) + list(REMOVE_DUPLICATES _PrivateIncludeDirs) + set(_OpsTestUt_OpApiSources ${_OpsTestUt_OpApiSources} ${_Sources} CACHE INTERNAL "" FORCE) + set(_OpsTestUt_OpApiPrivateIncludesExt ${_OpsTestUt_OpApiPrivateIncludesExt} ${_PrivateIncludeDirs} CACHE INTERNAL "" FORCE) + set(_OpsTestUt_OpApiLinkLibrariesExt ${_OpsTestUt_OpApiLinkLibrariesExt} ${TMP_PRIVATE_LINK_LIBRARIES_EXT} CACHE INTERNAL "" FORCE) + endif () +endfunction() + +# 私有函数, 外部不可直接调用, 用于实际生成 OpApi 动态库 +#[[ +]] +function(OpsTest_AddOpApiShared) + list(REMOVE_DUPLICATES _OpsTestUt_OpApiSources) + list(REMOVE_DUPLICATES _OpsTestUt_OpApiPrivateIncludesExt) + list(REMOVE_DUPLICATES _OpsTestUt_OpApiLinkLibrariesExt) + if ("${_OpsTestUt_OpApiSources}" STREQUAL "") + return() + endif () + set(_Target ${UTest_NamePrefix}_OpApi) + add_library(${_Target} SHARED) + target_sources(${_Target} + PRIVATE + ${_OpsTestUt_OpApiSources} + ) + target_include_directories(${_Target} + PRIVATE + ${OPAPI_INCLUDE} + ${_OpsTestUt_OpApiPrivateIncludesExt} + ${ASCEND_CANN_PACKAGE_PATH}/runtime/pkg_inc + ${ASCEND_CANN_PACKAGE_PATH}/runtime/pkg_inc/profiling + ) + target_compile_options(${_Target} + PRIVATE + $<$:-std=gnu++1z> + ) + target_compile_definitions(${_Target} + PRIVATE + ACLNN_LOG_FMT_CHECK + LOG_CPP + PROCESS_LOG + ) + target_link_libraries(${_Target} + PRIVATE + -Wl,--whole-archive + ${_OpsTestUt_OpApiLinkLibrariesExt} + $ + $ + -Wl,--no-whole-archive + PRIVATE $<$:$> + nnopbase + profapi + ge_common_base + ascend_dump + ascendalog + dl + ) + set(_UTest_OpApiLibrary ${_Target} PARENT_SCOPE) +endfunction() + +# Level1, 添加算子 OpProto 动态库 +#[[ +调用参数: + one_value_keywords: + SUB_SYSTEM : 必选参数, 用于指定算子所属子系统, 如 transformer + BRIEF : 必选参数, 用于指定算子缩略名(建议以大驼峰命名, 与算子实际名称无强制对应关系), 如 Fag/Fas/Fa + SNAKE : 必选参数, 用于指定算子全名, 如 flash_attention_score_grad + multi_value_keywords: + SOURCES_EXT : 可选参数, 额外源文件 + PRIVATE_INCLUDES_EXT : 可选参数, 额外头文件搜索路径 + PRIVATE_LINK_LIBRARIES_EXT : 可选参数, 额外链接库 +备注说明: + 本函数提供编译算子对应 OpProto.so 的功能. 下面介绍在调用本函数时所需了解的一些背景知识和注意事项. + 1. 本函数假设算子对应 OpProto 的源文件位于 'src/${SUB_SYSTEM}/${SNAKE}/op_host' 路径, 并自动添加以下内容: + a) 源文件: ${SNAKE}_proto.cpp + b) 头文件搜索路径: 'src/${SUB_SYSTEM}/${SNAKE}/op_host' + 2. 若算子 '需要额外的源文件’ 或 ‘源文件不满足上述约定的默认路径', 则可通过 SOURCES_EXT 参数指定 '额外的源文件'; + 3. 若算子 '需要额外头文件搜索路径’ 或 ‘本函数现有实现的头文件搜索路径设置不满足', + 则可通过 PRIVATE_INCLUDES_EXT 参数指定 '额外的头文件搜索路径'; 参数 PRIVATE_LINK_LIBRARIES_EXT 设置逻辑同理; +]] +function(OpsTest_Level1_AddOpProtoShared) + cmake_parse_arguments( + TMP + "" + "SUB_SYSTEM;BRIEF;SNAKE" + "SOURCES_EXT;PRIVATE_INCLUDES_EXT;PRIVATE_LINK_LIBRARIES_EXT" + "" + ${ARGN} + ) + get_filename_component(_Src "${OPS_ADV_DIR}/${TMP_SUB_SYSTEM}/${TMP_SNAKE}/op_host/${TMP_SNAKE}_proto.cpp" REALPATH) + if (NOT EXISTS "${_Src}") + set(_Src) + endif () + get_filename_component(_Inc "${OPS_ADV_DIR}/${TMP_SUB_SYSTEM}/${TMP_SNAKE}/op_host/" REALPATH) + if (NOT EXISTS "${_Inc}") + set(_Inc) + endif () + + set(_Sources ${TMP_SOURCES_EXT} ${_Src}) + list(REMOVE_DUPLICATES _Sources) + if (_Sources) + set(_PrivateIncludeDirs ${TMP_PRIVATE_INCLUDES_EXT} ${_Inc}) + list(REMOVE_DUPLICATES _PrivateIncludeDirs) + set(_OpsTestUt_OpProtoSources ${_OpsTestUt_OpProtoSources} ${_Sources} CACHE INTERNAL "" FORCE) + set(_OpsTestUt_OpProtoPrivateIncludesExt ${_OpsTestUt_OpProtoPrivateIncludesExt} ${_PrivateIncludeDirs} CACHE INTERNAL "" FORCE) + set(_OpsTestUt_OpProtoLinkLibrariesExt ${_OpsTestUt_OpProtoLinkLibrariesExt} ${TMP_PRIVATE_LINK_LIBRARIES_EXT} CACHE INTERNAL "" FORCE) + endif () +endfunction() + +# 私有函数, 外部不可直接调用, 用于实际生成 OpProto 动态库 +#[[ +]] +function(OpsTest_AddOpProtoShared) + list(REMOVE_DUPLICATES _OpsTestUt_OpProtoSources) + list(REMOVE_DUPLICATES _OpsTestUt_OpProtoPrivateIncludesExt) + list(REMOVE_DUPLICATES _OpsTestUt_OpProtoLinkLibrariesExt) + if ("${_OpsTestUt_OpProtoSources}" STREQUAL "") + return() + endif () + set(_Target ${UTest_NamePrefix}_OpProto) + add_library(${_Target} SHARED) + target_sources(${_Target} + PRIVATE + ${_OpsTestUt_OpProtoSources} + ) + target_include_directories(${_Target} + PRIVATE + ${OPBASE_INC_DIRS} + ${_OpsTestUt_OpProtoPrivateIncludesExt} + ) + target_compile_options(${_Target} + PRIVATE + $<$:-std=gnu++1z> + ) + target_compile_definitions(${_Target} + PRIVATE + LOG_CPP + PROCESS_LOG + ) + target_link_libraries(${_Target} + PRIVATE + -Wl,--whole-archive + ${_OpsTestUt_OpProtoLinkLibrariesExt} + $ + $ + $<$:$> + $<$:$> + -Wl,--no-whole-archive + ascendalog + ) + set(_UTest_OpProtoLibrary ${_Target} PARENT_SCOPE) +endfunction() + +# Level1, 添加算子 OpTiling 动态库 +#[[ +调用参数: + one_value_keywords: + SUB_SYSTEM : 必选参数, 用于指定算子所属子系统, 如 transformer + BRIEF : 必选参数, 用于指定算子缩略名(建议以大驼峰命名, 与算子实际名称无强制对应关系), 如 Fag/Fas/Fa + SNAKE : 必选参数, 用于指定算子全名, 如 flash_attention_score_grad + multi_value_keywords: + SOURCES_EXT : 可选参数, 额外源文件 + PRIVATE_INCLUDES_EXT : 可选参数, 额外头文件搜索路径 + PRIVATE_LINK_LIBRARIES_EXT : 可选参数, 额外链接库 +备注说明: + 本函数提供编译算子对应 OpTiling.so 的功能. 下面介绍在调用本函数时所需了解的一些背景知识和注意事项. + 1. 本函数假设算子对应 OpTiling 的源文件位于 'src/${SUB_SYSTEM}/${SNAKE}/op_host/${SNAKE}_tiling.cpp/.cc' 路径; + 2. 若算子 '需要额外的源文件’ 或 ‘源文件不满足上述约定的默认路径', 则可通过 SOURCES_EXT 参数指定 '额外的源文件'; + 3. 若算子 '需要额外头文件搜索路径’ 或 ‘本函数现有实现的头文件搜索路径设置不满足', + 则可通过 PRIVATE_INCLUDES_EXT 参数指定 '额外的头文件搜索路径'; 参数 PRIVATE_LINK_LIBRARIES_EXT 设置逻辑同理; +]] +function(OpsTest_Level1_AddOpTilingShared) + cmake_parse_arguments( + TMP + "" + "SUB_SYSTEM;BRIEF;SNAKE" + "SOURCES_EXT;PRIVATE_INCLUDES_EXT;PRIVATE_LINK_LIBRARIES_EXT" + "" + ${ARGN} + ) + file(GLOB _Src1 "${OPS_ADV_DIR}/${TMP_SUB_SYSTEM}/${TMP_SNAKE}/op_host/${TMP_SNAKE}_tiling.cc") + file(GLOB _Src2 "${OPS_ADV_DIR}/${TMP_SUB_SYSTEM}/${TMP_SNAKE}/op_host/${TMP_SNAKE}_tiling.cpp") + list(APPEND _Sources ${TMP_SOURCES_EXT} ${_Src1} ${_Src2}) + list(REMOVE_DUPLICATES _Sources) + get_filename_component(_Inc "${OPS_ADV_DIR}/${TMP_SUB_SYSTEM}/${TMP_SNAKE}/op_host/" REALPATH) + if (NOT EXISTS "${_Inc}") + set(_Inc) + endif () + if (_Sources) + set(_PrivateIncludeDirs ${TMP_PRIVATE_INCLUDES_EXT} ${_Inc}) + list(REMOVE_DUPLICATES _PrivateIncludeDirs) + set(_OpsTestUt_OpTilingSources ${_OpsTestUt_OpTilingSources} ${_Sources} CACHE INTERNAL "" FORCE) + set(_OpsTestUt_OpTilingPrivateIncludesExt ${_OpsTestUt_OpTilingPrivateIncludesExt} ${_PrivateIncludeDirs} CACHE INTERNAL "" FORCE) + set(_OpsTestUt_OpTilingLinkLibrariesExt ${_OpsTestUt_OpTilingLinkLibrariesExt} ${TMP_PRIVATE_LINK_LIBRARIES_EXT} CACHE INTERNAL "" FORCE) + endif () +endfunction() + +# 私有函数, 外部不可直接调用, 用于实际生成 OpTiling 动态库 +#[[ +]] +function(OpsTest_AddOpTilingShared) + list(REMOVE_DUPLICATES _OpsTestUt_OpTilingSources) + list(REMOVE_DUPLICATES _OpsTestUt_OpTilingPrivateIncludesExt) + list(REMOVE_DUPLICATES _OpsTestUt_OpTilingLinkLibrariesExt) + set(_Target ${UTest_NamePrefix}_OpTiling) + add_library(${_Target} SHARED) + target_sources(${_Target} + PRIVATE + ${_OpsTestUt_OpTilingSources} + ${OPS_ADV_DIR}/tests/ut/framework_special/stubs/tiling/tiling_templates_registry.cpp + ) + target_include_directories(${_Target} + PRIVATE + ${OPBASE_INC_DIRS} + ${_OpsTestUt_OpTilingPrivateIncludesExt} + ${OPS_TRANSFORMER_DIR}/common/include + ) + target_compile_definitions(${_Target} + PRIVATE + OP_TILING_LIB + LOG_CPP + PROCESS_LOG + ) + target_compile_options(${_Target} + PRIVATE + $<$:-std=c++11> + ) + target_link_libraries(${_Target} + PRIVATE + -Wl,--as-needed + -Wl,--no-whole-archive + ${_OpsTestUt_OpTilingLinkLibrariesExt} + $ + $ + $ + $<$:$> + $<$:$> + graph + graph_base + exe_graph + platform + register + ascendalog + tiling_api + c_sec + ) + set(_UTest_OpTilingLibrary ${_Target} PARENT_SCOPE) +endfunction() + +# Level1, 添加算子 Kernel 静态库 +#[[ +调用参数: + one_value_keywords: + SUB_SYSTEM : 必选参数, 用于指定算子所属子系统, 如 transformer + BRIEF : 必选参数, 用于指定算子缩略名(建议以大驼峰命名, 与算子实际名称无强制对应关系), 如 Fag/Fas/Fa + SNAKE : 必选参数, 用于指定算子全名, 如 flash_attention_score_grad + multi_value_keywords: + SOURCES_EXT : 可选参数, 额外源文件 + TILING_DATA_DEF_H : 可选参数, 算子 Kernel 所需 TilingData 定义头文件 + PRIVATE_INCLUDES_EXT : 可选参数, 额外头文件搜索路径 + PRIVATE_COMPILE_DEFINITIONS_EXT : 可选参数, 额外编译宏 +备注说明: + 本函数提供编译算子对应 kernel.a 的功能. 下面介绍在调用本函数时所需了解的一些背景知识和注意事项. + 1. 本函数假设算子对应 Kernel 的源文件位于 'src/${SUB_SYSTEM}/${SNAKE}/${SNAKE}.cpp' 路径; + 2. 若算子 '需要额外的源文件’ 或 ‘源文件不满足上述约定的默认路径', 则可通过 SOURCES_EXT 参数指定 '额外的源文件'; + 3. 当前 Ascend C 融合算子 Kernel 一般需要多个(>2) TilingData 定义头文件, + 如 flash_attention_score 算子需要 data_copy_transpose_tiling_def.h 和 flash_attention_score_tiling.h; + - 在 NPU 编译时, 编译框架使用 ccec 编译器并使用 -include 编译选项注入所需的多个 TilingData 定义头文件 至 Kernel 源文件; + - 但 CPU 编译时, CMake 早期版本在处理 -include 选项时有 Bug(只第一个 -include 指定的头文件生效). + 故本框架通过新增一个 {BRIEF}_tiling_data.h, 再在该文件中 include 所需的多个 TilingData 定义头文件, + 再通过 -include 编译选项注入 {BRIEF}_tiling_data.h 的方式规避 CPU 编译时 CMake 不能处理多个 -include 选项的 Bug; + 4. PRIVATE_COMPILE_DEFINITIONS_EXT 设置格式如下: + optional{KernelCtrlParam func suffix} optional{OtherCompileDefinitions}} + 其中 KernelCtrlParam 用于指定 Kernel 编译控制参数, func 为标识 Kernel 原始入口函数名, + suffix 为标识当前 Kernel 二进制及 Kernel 入口函数后缀; 其设置逻辑如下: + 4.0 KernelCtrlParam 内 func, suffix 必需按序设置; + 4.1 当不设置 KernelCtrlParam 时, 仅会编译一个 Kernel.a 并以 OtherCompileDefinitions 设置编译 Definitions(如设置); + 4.2 若当设置 KernelCtrlParam 时, 本函数会按 KernelCtrlParam 个数逐个编译 Kernel.a, + 并以 KernelCtrlParam 间 OtherCompileDefinitions 设置编译宏定义; 并将 func 重命名为 {func}_{suffix}; + 4.3 KernelCtrlParam 内 func 支持指定多个, 指定多个时 func 之间用','间隔; +]] +function(OpsTest_Level1_AddOpKernelStatic) + cmake_parse_arguments( + TMP + "" + "SUB_SYSTEM;BRIEF;SNAKE" + "SOURCES_EXT;TILING_DATA_DEF_H;PRIVATE_INCLUDES_EXT;PRIVATE_COMPILE_DEFINITIONS_EXT" + "" + ${ARGN} + ) + # 生成 Kernel 所需的结构体表示的对应 tiling.h + string(TOLOWER ${TMP_BRIEF} tmp_brief) + set(_tmp_files ${TMP_TILING_DATA_DEF_H}) + list(REMOVE_DUPLICATES _tmp_files) + set(_ori_files) + set(_define_py ${OPS_ADV_DIR}/cmake/scripts/utest/gen_tiling_data_stub.py) + foreach (_tmp ${_tmp_files}) + if (EXISTS ${_tmp}) + list(APPEND _ori_files "-s=${_tmp}") + else () + message(FATAL_ERROR "${_tmp} not exist.") + endif () + endforeach () + # 生成目标根目录 + get_filename_component(_OpsTest_GenDir "${CMAKE_CURRENT_BINARY_DIR}/gen" REALPATH) + get_filename_component(_OpsTest_GenDirInc "${_OpsTest_GenDir}/inc" REALPATH) + execute_process( + COMMAND ${HI_PYTHON} ${_define_py} "-o=${tmp_brief}" ${_ori_files} "-d=${_OpsTest_GenDirInc}" + ) + set(_Target ${UTest_NamePrefix}_${TMP_BRIEF}_OpTilingDataDef) + add_library(${_Target} INTERFACE) + target_include_directories(${_Target} INTERFACE ${_OpsTest_GenDirInc} ${OPS_ADV_UTILS_KERNEL_INC}) + + # 编译变量处理 + set(_TargetPrefix ${UTest_NamePrefix}_${TMP_BRIEF}_OpKernel) + aux_source_directory(${OPS_ADV_DIR}/${TMP_SUB_SYSTEM}/${TMP_SNAKE} _Sources) + list(APPEND _Sources ${TMP_SOURCES_EXT}) + set(_PrivateIncludeDirectories + ${_OpsTest_GenDirInc} + ${TMP_PRIVATE_INCLUDES_EXT} + /${SYSTEM_PREFIX}/ascendc/include/basic_api + /${SYSTEM_PREFIX}/ascendc/include/basic_api/impl + /${SYSTEM_PREFIX}/ascendc/include/basic_api/interface + /tools/tikicpulib/lib/include + /include/ascendc + /x86_64-linux/include/ascendc/highlevel_api + /pkg_inc/runtime/runtime + /runtime/pkg_inc + ) + get_filename_component(_Inc "${OPS_ADV_DIR}/${TMP_SUB_SYSTEM}/${TMP_SNAKE}" REALPATH) + if (EXISTS "${_Inc}") + list(APPEND _PrivateIncludeDirectories ${_Inc}) + endif () + set(_PrivateCompileOptions + -include ${_OpsTest_GenDirInc}/tiling/${tmp_brief}/tiling_stub.h + ) + set(_PrivateLinkLibraries + -Wl,--as-needed + -Wl,--no-whole-archive + c_sec + $ + $ + ) + + # 多 Kernel 处理 + set(_OpKernelLibraries) + list(FIND TMP_PRIVATE_COMPILE_DEFINITIONS_EXT KernelCtrlParam _GrpIdx) + if ("${_GrpIdx}" STREQUAL "-1") + # 不存在多 Kernel 配置时, 不添加后缀, 编译一个 Kernel.a + set(_Target ${_TargetPrefix}) + add_library(${_Target} STATIC) + target_sources(${_Target} PRIVATE ${_Sources}) + target_include_directories(${_Target} PRIVATE ${_PrivateIncludeDirectories}) + target_compile_definitions(${_Target} PRIVATE ${TMP_PRIVATE_COMPILE_DEFINITIONS_EXT}) + target_compile_options(${_Target} PRIVATE ${_PrivateCompileOptions}) + target_link_libraries(${_Target} + PUBLIC + tikicpulib::${OPS_ADV_UTEST_OPS_TEST_ASCEND_PRODUCT_TYPE} + PRIVATE + ${_PrivateLinkLibraries} + ) + list(APPEND _OpKernelLibraries ${_Target}) + else () + # 存在 1/n 多 Kernel 配置时, 添加后缀, 编译多 Kernel.a + while (NOT "${_GrpIdx}" STREQUAL "-1") + # 获取当前 Func, Suffix, CompileDefinitions + math(EXPR _FuncIdx "${_GrpIdx} + 1") + math(EXPR _SuffixIdx "${_GrpIdx} + 2") + math(EXPR _SubLstBgnIdx "${_GrpIdx} + 3") + list(GET TMP_PRIVATE_COMPILE_DEFINITIONS_EXT ${_FuncIdx} _FuncOriValList) + list(GET TMP_PRIVATE_COMPILE_DEFINITIONS_EXT ${_SuffixIdx} _SuffixVal) + list(SUBLIST TMP_PRIVATE_COMPILE_DEFINITIONS_EXT ${_SubLstBgnIdx} -1 TMP_PRIVATE_COMPILE_DEFINITIONS_EXT) + list(FIND TMP_PRIVATE_COMPILE_DEFINITIONS_EXT KernelCtrlParam _GrpIdx) # _GrpIdx 移动 + list(SUBLIST TMP_PRIVATE_COMPILE_DEFINITIONS_EXT 0 ${_GrpIdx} _SubCompileDefinitions) + string(REPLACE "," ";" _FuncOriValList "${_FuncOriValList}") + set(_FuncValDef) + foreach (_f ${_FuncOriValList}) + list(APPEND _FuncValDef "-D${_f}=${_f}_${_SuffixVal}") + endforeach () + # 编译目标 + set(_Target ${_TargetPrefix}_${_SuffixVal}) + add_library(${_Target} STATIC) + target_sources(${_Target} PRIVATE ${_Sources}) + target_include_directories(${_Target} PRIVATE ${_PrivateIncludeDirectories}) + target_compile_definitions(${_Target} PRIVATE ${_FuncValDef} ${_SubCompileDefinitions}) + target_compile_options(${_Target} PRIVATE ${_PrivateCompileOptions}) + target_link_libraries(${_Target} + PUBLIC + tikicpulib::${OPS_ADV_UTEST_OPS_TEST_ASCEND_PRODUCT_TYPE} + PRIVATE + ${_PrivateLinkLibraries} + ) + list(APPEND _OpKernelLibraries ${_Target}) + endwhile () + endif () + set(_OpsTestUt_OpKernelLibraries ${_OpKernelLibraries} CACHE INTERNAL "" FORCE) +endfunction() + +# Level1, 添加算子 UTest 用例 Common 静态库 +#[[ +调用参数: + one_value_keywords: + BRIEF : 必选参数, 用于指定算子缩略名(建议以大驼峰命名, 与算子实际名称无强制对应关系), 如 Fag/Fas/Fa + SNAKE : 必选参数, 用于指定算子全名, 如 flash_attention_score_grad + TILING_DATA_DEF_BRIEF: 可选参数, 用于指定链接的 TilingData 头文件目标 + multi_value_keywords: + SOURCES_EXT : 可选参数, 额外源文件 + SOURCES_EXT_FORCE : 可选参数, 强制使用 SOURCES_EXT 指定的额外源文件 (Common), 不再自动寻找 comm 下源文件 + PUBLIC_INCLUDES_EXT : 可选参数, 额外头文件搜索路径(public) + PRIVATE_INCLUDES_EXT : 可选参数, 额外头文件搜索路径 + PRIVATE_LINK_LIBRARIES_EXT : 可选参数, 额外链接库 + PRIVATE_COMPILE_DEFINITIONS : 可选参数, 额外的编译宏定义 +]] +function(OpsTest_Level1_AddUTestCommonStatic) + cmake_parse_arguments( + TMP + "" + "BRIEF;SNAKE;TILING_DATA_DEF_BRIEF" + "SOURCES_EXT;SOURCES_EXT_FORCE;PUBLIC_INCLUDES_EXT;PRIVATE_INCLUDES_EXT;PRIVATE_LINK_LIBRARIES_EXT;PRIVATE_COMPILE_DEFINITIONS" + "" + ${ARGN} + ) + if (NOT TMP_TILING_DATA_DEF_BRIEF) + set(TMP_TILING_DATA_DEF_BRIEF ${TMP_BRIEF}) + endif () + + # 强制指定源码处理 + set(_Src1) + set(_Src2) + set(_Common_Sources ${TMP_SOURCES_EXT}) + if (NOT TMP_SOURCES_EXT_FORCE) + if (EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/comm) + if (EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/comm/inc) + aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/comm/src _Src1) + else () + file(GLOB_RECURSE _Src1 "${CMAKE_CURRENT_SOURCE_DIR}/comm/*.cc") + file(GLOB_RECURSE _Src2 "${CMAKE_CURRENT_SOURCE_DIR}/comm/*.cpp") + endif () + endif () + endif () + list(APPEND _Common_Sources ${_Src1} ${_Src2}) + list(REMOVE_DUPLICATES _Common_Sources) + + set(_Target_UTest_Common) + if (_Common_Sources) + set(_Target_UTest_Common ${UTest_NamePrefix}_${TMP_BRIEF}_UTest_Common) + add_library(${_Target_UTest_Common} STATIC) + target_sources(${_Target_UTest_Common} PRIVATE ${_Common_Sources}) + if (EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/comm/inc) + target_include_directories(${_Target_UTest_Common} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/comm/inc) + elseif (EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/comm) + target_include_directories(${_Target_UTest_Common} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/comm) + endif () + target_include_directories(${_Target_UTest_Common} + PUBLIC + ${TMP_PUBLIC_INCLUDES_EXT} + PRIVATE + ${OPBASE_INC_DIRS} + ${OPS_TRANSFORMER_DIR}/common/include + ${ASCEND_CANN_PACKAGE_PATH}/include + ${TMP_PRIVATE_INCLUDES_EXT} + ) + target_compile_definitions(${_Target_UTest_Common} + PRIVATE + ${TMP_PRIVATE_COMPILE_DEFINITIONS} + ) + target_link_libraries(${_Target_UTest_Common} + PRIVATE + -Wl,--as-needed + -Wl,--no-whole-archive + platform + register + opp_registry + error_manager + metadef + tiling_api + $ + $ + $<$:$> + $<$:$> + tikicpulib::${OPS_ADV_UTEST_OPS_TEST_ASCEND_PRODUCT_TYPE} + ${UTest_NamePrefix}_${TMP_TILING_DATA_DEF_BRIEF}_OpTilingDataDef + ${UTest_NamePrefix}_Utils + ${TMP_PRIVATE_LINK_LIBRARIES_EXT} + ${_OpsTestUt_OpKernelLibraries} # 当前各算子 common 内以 extern 方式声明 kernel 入口, 故添加依赖 + ) + endif () + set(_OpsTestUt_UTestCommonLibrary ${_Target_UTest_Common} CACHE INTERNAL "" FORCE) +endfunction() + +# Level1, 添加算子 UTest 用例静态库 +#[[ +调用参数: + one_value_keywords: + BRIEF : 必选参数, 用于指定算子缩略名(建议以大驼峰命名, 与算子实际名称无强制对应关系), 如 Fag/Fas/Fa + SNAKE : 必选参数, 用于指定算子全名, 如 flash_attention_score_grad + TILING_DATA_DEF_BRIEF: 可选参数, 用于指定链接的 TilingData 头文件目标 + multi_value_keywords: + SOURCES_EXT : 可选参数, 额外源文件 + SOURCES_EXT_FORCE : 可选参数, 强制使用 SOURCES_EXT 指定的额外源文件 (UTest), 不再自动寻找 utest 下源文件 + PRIVATE_INCLUDES_EXT : 可选参数, 额外头文件搜索路径 + PRIVATE_LINK_LIBRARIES_EXT : 可选参数, 额外链接库 + PRIVATE_COMPILE_DEFINITIONS : 可选参数, 额外的编译宏定义 +]] +function(OpsTest_Level1_AddUTestCaseStatic) + cmake_parse_arguments( + TMP + "" + "BRIEF;SNAKE;TILING_DATA_DEF_BRIEF" + "SOURCES_EXT;SOURCES_EXT_FORCE;PRIVATE_INCLUDES_EXT;PRIVATE_LINK_LIBRARIES_EXT;PRIVATE_COMPILE_DEFINITIONS" + "" + ${ARGN} + ) + if (NOT TMP_TILING_DATA_DEF_BRIEF) + set(TMP_TILING_DATA_DEF_BRIEF ${TMP_BRIEF}) + endif () + + # 强制指定源码处理 + set(_UTest_Sources ${TMP_SOURCES_EXT}) + if (NOT TMP_SOURCES_EXT_FORCE) + file(GLOB_RECURSE _Src1 "${CMAKE_CURRENT_SOURCE_DIR}/utest/*.cc") + file(GLOB_RECURSE _Src2 "${CMAKE_CURRENT_SOURCE_DIR}/utest/*.cpp") + list(APPEND _UTest_Sources ${_Src1} ${_Src2}) + endif () + list(REMOVE_DUPLICATES _UTest_Sources) + if (_UTest_Sources) + set(_Target_UTest ${UTest_NamePrefix}_${TMP_BRIEF}_UTest_Case) + add_library(${_Target_UTest} STATIC) + target_sources(${_Target_UTest} PRIVATE ${_UTest_Sources}) + target_include_directories(${_Target_UTest} + PRIVATE + ${ASCEND_CANN_PACKAGE_PATH}/include + ${TMP_PRIVATE_INCLUDES_EXT} + ${CMAKE_CURRENT_SOURCE_DIR}/utest + ) + target_compile_definitions(${_Target_UTest} + PRIVATE + ${TMP_PRIVATE_COMPILE_DEFINITIONS} + ) + target_link_libraries(${_Target_UTest} + PRIVATE + -Wl,--as-needed + -Wl,--no-whole-archive + gtest + $ + $ + $<$:$> + $<$:$> + tikicpulib::${OPS_ADV_UTEST_OPS_TEST_ASCEND_PRODUCT_TYPE} + ${UTest_NamePrefix}_${TMP_TILING_DATA_DEF_BRIEF}_OpTilingDataDef + ${UTest_NamePrefix}_Utils + $ + ${UTest_NamePrefix}_Utest + ${TMP_PRIVATE_LINK_LIBRARIES_EXT} + ${_OpsTestUt_UTestCommonLibrary} + ) + set(_OpsTestUt_UTestCaseLibraries ${_OpsTestUt_UTestCaseLibraries} ${_Target_UTest} CACHE INTERNAL "" FORCE) + endif () +endfunction() + +# Level1, 添加算子 UTest 用例(Aclnn)静态库 +#[[ +调用参数: + one_value_keywords: + BRIEF : 必选参数, 用于指定算子缩略名(建议以大驼峰命名, 与算子实际名称无强制对应关系), 如 Fag/Fas/Fa + SNAKE : 必选参数, 用于指定算子全名, 如 flash_attention_score_grad + TILING_DATA_DEF_BRIEF: 可选参数, 用于指定链接的 TilingData 头文件目标 + multi_value_keywords: + SOURCES_EXT : 可选参数, 额外源文件 + SOURCES_EXT_FORCE : 可选参数, 强制使用 SOURCES_EXT 指定的额外源文件 (UTest_Aclnn), 不再自动寻找 utest_aclnn 下源文件 + PRIVATE_INCLUDES_EXT : 可选参数, 额外头文件搜索路径 + PRIVATE_COMPILE_DEFINITIONS : 可选参数, 额外的编译宏定义 + PRIVATE_LINK_LIBRARIES_EXT : 可选参数, 额外链接库 +]] +function(OpsTest_Level1_AddUTestAclnnCaseStatic) + cmake_parse_arguments( + TMP + "" + "BRIEF;SNAKE;TILING_DATA_DEF_BRIEF" + "SOURCES_EXT;SOURCES_EXT_FORCE;PRIVATE_INCLUDES_EXT;PRIVATE_COMPILE_DEFINITIONS;PRIVATE_LINK_LIBRARIES_EXT" + "" + ${ARGN} + ) + if (NOT TMP_TILING_DATA_DEF_BRIEF) + set(TMP_TILING_DATA_DEF_BRIEF ${TMP_BRIEF}) + endif () + + # 强制指定源码处理 + set(_UTest_Aclnn_Sources ${TMP_SOURCES_EXT}) + if (NOT TMP_SOURCES_EXT_FORCE) + file(GLOB_RECURSE _Src1 "${CMAKE_CURRENT_SOURCE_DIR}/utest_aclnn/*.cc") + file(GLOB_RECURSE _Src2 "${CMAKE_CURRENT_SOURCE_DIR}/utest_aclnn/*.cpp") + list(APPEND _UTest_Aclnn_Sources ${_Src1} ${_Src2}) + endif () + list(REMOVE_DUPLICATES _UTest_Aclnn_Sources) + if (_UTest_Aclnn_Sources) + set(_Target_UTest_Aclnn ${UTest_NamePrefix}_${TMP_BRIEF}_UTest_Case_Aclnn) + add_library(${_Target_UTest_Aclnn} STATIC) + target_sources(${_Target_UTest_Aclnn} PRIVATE ${_UTest_Aclnn_Sources}) + target_include_directories(${_Target_UTest_Aclnn} + PRIVATE + ${ASCEND_CANN_PACKAGE_PATH}/include + ${TMP_PRIVATE_INCLUDES_EXT} + ${CMAKE_CURRENT_SOURCE_DIR}/utest_aclnn + ) + target_compile_definitions(${_Target_UTest_Aclnn} + PRIVATE + ${TMP_PRIVATE_COMPILE_DEFINITIONS} + ) + target_link_libraries(${_Target_UTest_Aclnn} + PRIVATE + -Wl,--as-needed + -Wl,--no-whole-archive + gtest + $ + $ + $<$:$> + $<$:$> + tikicpulib::${OPS_ADV_UTEST_OPS_TEST_ASCEND_PRODUCT_TYPE} + ${UTest_NamePrefix}_${TMP_TILING_DATA_DEF_BRIEF}_OpTilingDataDef + ${UTest_NamePrefix}_Utils + $ + ${UTest_NamePrefix}_Utest + ${TMP_PRIVATE_LINK_LIBRARIES_EXT} + ${_OpsTestUt_UTestCommonLibrary} + ) + set(_OpsTestUt_UTestAclnnCaseLibraries ${_OpsTestUt_UTestAclnnCaseLibraries} ${_Target_UTest_Aclnn} CACHE INTERNAL "" FORCE) + endif () +endfunction() + +# Level2, 添加算子 UTest 用例动态库 及其所需全部库(OpTiling.so, OpKernel.a 等) +#[[ +调用参数: + one_value_keywords: + SUB_SYSTEM : 必选参数, 用于指定算子所属子系统, 如 transformer + BRIEF : 必选参数, 用于指定算子缩略名(建议以大驼峰命名, 与算子实际名称无强制对应关系), 如 Fag/Fas/Fa + SNAKE : 必选参数, 用于指定算子全名, 如 flash_attention_score_grad + TILING_DATA_DEF_BRIEF: 可选参数, 用于指定链接的 TilingData 头文件目标 + multi_value_keywords: + OPAPI_SOURCES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddOpApiShared 函数说明 + OPAPI_PRIVATE_INCLUDES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddOpApiShared 函数说明 + OPAPI_PRIVATE_LINK_LIBRARIES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddOpApiShared 函数说明 + PROTO_SOURCES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddOpProtoShared 函数说明 + PROTO_PRIVATE_INCLUDES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddOpProtoShared 函数说明 + PROTO_PRIVATE_LINK_LIBRARIES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddOpProtoShared 函数说明 + TILING_SOURCES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddOpTilingShared 函数说明 + TILING_PRIVATE_INCLUDES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddOpTilingShared 函数说明 + TILING_PRIVATE_LINK_LIBRARIES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddOpTilingShared 函数说明 + KERNEL_SOURCES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddOpKernelStatic 函数说明 + KERNEL_TILING_DATA_DEF_H : 透传参数, 详情参见 OpsTest_Level1_AddOpKernelStatic 函数说明 + KERNEL_PRIVATE_INCLUDES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddOpKernelStatic 函数说明 + KERNEL_PRIVATE_COMPILE_DEFINITIONS_EXT : 透传参数, 详情参见 OpsTest_Level1_AddOpKernelStatic 函数说明 + UTEST_COMMON_SOURCES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddUTestCommonStatic 函数说明 + UTEST_COMMON_SOURCES_EXT_FORCE : 透传参数, 详情参见 OpsTest_Level1_AddUTestCommonStatic 函数说明 + UTEST_COMMON_PUBLIC_INCLUDES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddUTestCommonStatic 函数说明 + UTEST_COMMON_PRIVATE_INCLUDES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddUTestCommonStatic 函数说明 + UTEST_COMMON_PRIVATE_COMPILE_DEFINITIONS : 透传参数, 详情参见 OpsTest_Level1_AddUTestCommonStatic 函数说明 + UTEST_COMMON_PRIVATE_LINK_LIBRARIES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddUTestCommonStatic 函数说明 + UTEST_SOURCES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddUTestCaseStatic 函数说明 + UTEST_SOURCES_EXT_FORCE : 透传参数, 详情参见 OpsTest_Level1_AddUTestCaseStatic 函数说明 + UTEST_PRIVATE_INCLUDES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddUTestCaseStatic 函数说明 + UTEST_PRIVATE_COMPILE_DEFINITIONS : 透传参数, 详情参见 OpsTest_Level1_AddUTestCaseStatic 函数说明 + UTEST_PRIVATE_LINK_LIBRARIES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddUTestCaseStatic 函数说明 + UTEST_ACLNN_SOURCES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddUTestAclnnCaseStatic 函数说明 + UTEST_ACLNN_SOURCES_EXT_FORCE : 透传参数, 详情参见 OpsTest_Level1_AddUTestAclnnCaseStatic 函数说明 + UTEST_ACLNN_PRIVATE_INCLUDES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddUTestAclnnCaseStatic 函数说明 + UTEST_ACLNN_PRIVATE_COMPILE_DEFINITIONS : 透传参数, 详情参见 OpsTest_Level1_AddUTestAclnnCaseStatic 函数说明 + UTEST_ACLNN_PRIVATE_LINK_LIBRARIES_EXT : 透传参数, 详情参见 OpsTest_Level1_AddUTestAclnnCaseStatic 函数说明 +]] +function(OpsTest_Level2_AddOp) + cmake_parse_arguments( + TMP + "" + "SUB_SYSTEM;BRIEF;SNAKE;TILING_DATA_DEF_BRIEF" + "OPAPI_SOURCES_EXT;OPAPI_PRIVATE_INCLUDES_EXT;OPAPI_PRIVATE_LINK_LIBRARIES_EXT;PROTO_SOURCES_EXT;PROTO_PRIVATE_INCLUDES_EXT;PROTO_PRIVATE_LINK_LIBRARIES_EXT;TILING_SOURCES_EXT;TILING_PRIVATE_INCLUDES_EXT;TILING_PRIVATE_LINK_LIBRARIES_EXT;KERNEL_SOURCES_EXT;KERNEL_TILING_DATA_DEF_H;KERNEL_PRIVATE_INCLUDES_EXT;KERNEL_PRIVATE_COMPILE_DEFINITIONS_EXT;UTEST_COMMON_SOURCES_EXT;UTEST_COMMON_SOURCES_EXT_FORCE;UTEST_COMMON_PUBLIC_INCLUDES_EXT;UTEST_COMMON_PRIVATE_INCLUDES_EXT;UTEST_COMMON_PRIVATE_COMPILE_DEFINITIONS;UTEST_COMMON_PRIVATE_LINK_LIBRARIES_EXT;UTEST_SOURCES_EXT;UTEST_SOURCES_EXT_FORCE;UTEST_PRIVATE_INCLUDES_EXT;UTEST_PRIVATE_COMPILE_DEFINITIONS;UTEST_PRIVATE_LINK_LIBRARIES_EXT;UTEST_ACLNN_SOURCES_EXT;UTEST_ACLNN_SOURCES_EXT_FORCE;UTEST_ACLNN_PRIVATE_INCLUDES_EXT;UTEST_ACLNN_PRIVATE_COMPILE_DEFINITIONS;UTEST_ACLNN_PRIVATE_LINK_LIBRARIES_EXT" + "" + ${ARGN} + ) + + OpsTest_Level1_AddOpApiShared( + SUB_SYSTEM ${TMP_SUB_SYSTEM} + BRIEF ${TMP_BRIEF} + SNAKE ${TMP_SNAKE} + SOURCES_EXT ${TMP_OPAPI_SOURCES_EXT} + PRIVATE_INCLUDES_EXT ${TMP_OPAPI_PRIVATE_INCLUDES_EXT} + PRIVATE_LINK_LIBRARIES_EXT ${TMP_OPAPI_PRIVATE_LINK_LIBRARIES_EXT} + ) + OpsTest_Level1_AddOpProtoShared( + SUB_SYSTEM ${TMP_SUB_SYSTEM} + BRIEF ${TMP_BRIEF} + SNAKE ${TMP_SNAKE} + SOURCES_EXT ${TMP_PROTO_SOURCES_EXT} + PRIVATE_INCLUDES_EXT ${TMP_PROTO_PRIVATE_INCLUDES_EXT} + PRIVATE_LINK_LIBRARIES_EXT ${TMP_PROTO_PRIVATE_LINK_LIBRARIES_EXT} + ) + OpsTest_Level1_AddOpTilingShared( + SUB_SYSTEM ${TMP_SUB_SYSTEM} + BRIEF ${TMP_BRIEF} + SNAKE ${TMP_SNAKE} + SOURCES_EXT ${TMP_TILING_SOURCES_EXT} + PRIVATE_INCLUDES_EXT ${TMP_TILING_PRIVATE_INCLUDES_EXT} + PRIVATE_LINK_LIBRARIES_EXT ${TMP_TILING_PRIVATE_LINK_LIBRARIES_EXT} + ) + OpsTest_Level1_AddOpKernelStatic( + SUB_SYSTEM ${TMP_SUB_SYSTEM} + BRIEF ${TMP_BRIEF} + SNAKE ${TMP_SNAKE} + SOURCES_EXT ${TMP_KERNEL_SOURCES_EXT} + TILING_DATA_DEF_H ${TMP_KERNEL_TILING_DATA_DEF_H} + PRIVATE_INCLUDES_EXT ${TMP_KERNEL_PRIVATE_INCLUDES_EXT} + PRIVATE_COMPILE_DEFINITIONS_EXT ${TMP_KERNEL_PRIVATE_COMPILE_DEFINITIONS_EXT} + ) + OpsTest_Level1_AddUTestCommonStatic( + BRIEF ${TMP_BRIEF} + SNAKE ${TMP_SNAKE} + SOURCES_EXT ${TMP_UTEST_COMMON_SOURCES_EXT} + SOURCES_EXT_FORCE ${TMP_UTEST_COMMON_SOURCES_EXT_FORCE} + TILING_DATA_DEF_BRIEF ${TMP_TILING_DATA_DEF_BRIEF} + PUBLIC_INCLUDES_EXT ${TMP_UTEST_COMMON_PUBLIC_INCLUDES_EXT} + PRIVATE_INCLUDES_EXT ${TMP_UTEST_COMMON_PRIVATE_INCLUDES_EXT} + PRIVATE_COMPILE_DEFINITIONS ${TMP_UTEST_COMMON_PRIVATE_COMPILE_DEFINITIONS} + PRIVATE_LINK_LIBRARIES_EXT ${TMP_UTEST_COMMON_PRIVATE_LINK_LIBRARIES_EXT} + ) + OpsTest_Level1_AddUTestCaseStatic( + BRIEF ${TMP_BRIEF} + SNAKE ${TMP_SNAKE} + TILING_DATA_DEF_BRIEF ${TMP_TILING_DATA_DEF_BRIEF} + SOURCES_EXT ${TMP_UTEST_SOURCES_EXT} + SOURCES_EXT_FORCE ${TMP_UTEST_SOURCES_EXT_FORCE} + PRIVATE_INCLUDES_EXT ${TMP_UTEST_PRIVATE_INCLUDES_EXT} + PRIVATE_COMPILE_DEFINITIONS ${TMP_UTEST_PRIVATE_COMPILE_DEFINITIONS} + PRIVATE_LINK_LIBRARIES_EXT ${TMP_UTEST_PRIVATE_LINK_LIBRARIES_EXT} + ) + OpsTest_Level1_AddUTestAclnnCaseStatic( + BRIEF ${TMP_BRIEF} + SNAKE ${TMP_SNAKE} + TILING_DATA_DEF_BRIEF ${TMP_TILING_DATA_DEF_BRIEF} + SOURCES_EXT ${TMP_UTEST_ACLNN_SOURCES_EXT} + SOURCES_EXT_FORCE ${TMP_UTEST_ACLNN_SOURCES_EXT_FORCE} + PRIVATE_INCLUDES_EXT ${TMP_UTEST_ACLNN_PRIVATE_INCLUDES_EXT} + PRIVATE_COMPILE_DEFINITIONS ${TMP_UTEST_ACLNN_PRIVATE_COMPILE_DEFINITIONS} + PRIVATE_LINK_LIBRARIES_EXT ${TMP_UTEST_ACLNN_PRIVATE_LINK_LIBRARIES_EXT} + ) +endfunction() + +function(op_add_ut_subdirectory OP_UT_LIST OP_UT_DIR_LIST) + set(_OP_UT_LIST) + set(_OP_UT_DIR_LIST) + + file(GLOB OP_HOST_CMAKE_FILES "${CMAKE_CURRENT_SOURCE_DIR}/**/**/tests/CMakeLists.txt") + + foreach(OP_CMAKE_FILE ${OP_HOST_CMAKE_FILES}) + if ("${OP_CMAKE_FILE}" MATCHES "tests") + get_filename_component(OP_HOST_DIR "${OP_CMAKE_FILE}" DIRECTORY) + get_filename_component(OP_DIR "${OP_HOST_DIR}" DIRECTORY) + else() + get_filename_component(OP_DIR "${OP_CMAKE_FILE}" DIRECTORY) + endif() + get_filename_component(OP_NAME "${OP_DIR}" NAME) + + if (DEFINED ASCEND_OP_NAME AND NOT "${ASCEND_OP_NAME}" STREQUAL "") + if (NOT "${ASCEND_OP_NAME}" STREQUAL "all" AND NOT "${ASCEND_OP_NAME}" STREQUAL "ALL") + if (NOT ${OP_NAME} IN_LIST ASCEND_OP_NAME) + continue() + endif () + endif () + endif () + + list(APPEND _OP_UT_LIST ${OP_NAME}) + list(APPEND _OP_UT_DIR_LIST ${OP_DIR}) + endforeach() + + list(REMOVE_DUPLICATES _OP_UT_LIST) + list(REMOVE_DUPLICATES _OP_UT_DIR_LIST) + list(SORT _OP_UT_LIST) + list(SORT _OP_UT_DIR_LIST) + set(${OP_UT_LIST} ${_OP_UT_LIST} PARENT_SCOPE) + set(${OP_UT_DIR_LIST} ${_OP_UT_DIR_LIST} PARENT_SCOPE) +endfunction() + +# 私有函数, 外部不可直接调用, 用于执行包含多个算子的 UT 可执行程序 +#[[ +]] +function(OpsTest_RunLaunch) + cmake_parse_arguments( + TMP + "" + "TARGET;EXECUTABLE" + "" + "" + ${ARGN} + ) + if (ENABLE_UT_EXEC) + set(LD_LIBRARY_PATH_ "LD_LIBRARY_PATH=$ENV{LD_LIBRARY_PATH}:${ASCEND_CANN_PACKAGE_PATH}/compiler/lib64:${ASCEND_CANN_PACKAGE_PATH}/toolkit/tools/simulator/Ascend910B1/lib") + if (ENABLE_ASAN OR ENABLE_UBSAN) + if (ENABLE_ASAN) + # 谨慎修改 ASAN_OPTIONS_ 取值, 当前出现告警会使 UT 失败. + # halt_on_error=1, 出现告警时停止运行进而触发构建失败, 避免主进程或 CPU孪生调试 fork 出的子进程出现错误无法发现的情况 + # detect_stack_use_after_return=1, 栈空间返回后使用检测 + # check_initialization_order, 尝试捕获初始化顺序问题 + # strict_init_order, 动态初始化器永远不能访问来自其他模块的全局变量, 及时或者已经初始化 + # strict_string_checks, 检查字符串参数是否正确以 null 终止 + # detect_leaks=1, 内存泄漏检测 + set(ASAN_OPTIONS_ "ASAN_OPTIONS=halt_on_error=1,detect_stack_use_after_return=1,check_initialization_order=1,strict_init_order=1,strict_string_checks=1,detect_leaks=1") + message(STATUS "${TMP_EXECUTABLE}: ${ASAN_OPTIONS_}") + endif () + if (ENABLE_UBSAN) + # 谨慎修改 UBSAN_OPTIONS_ 取值, 当前出现告警会使 UT 失败. + # halt_on_error=1, 出现告警时停止运行进而触发构建失败, 避免主进程或 CPU孪生调试 fork 出的子进程出现错误无法发现的情况 + # print_stacktrace=1, 出错时打印调用栈 + set(UBSAN_OPTIONS_ "UBSAN_OPTIONS=halt_on_error=1,print_stacktrace=1") + message(STATUS "${TMP_EXECUTABLE}: ${UBSAN_OPTIONS_}") + endif () + # 用例执行 + # 采用链接 ASAN/MSAN 动态库方式执行, 受 SAN 机制限制, 若采用静态库, 则会导致如 OpTiling/OpProto 内 SAN 功能失效. + if (NOT "${SAN_LD_PRELOAD}" STREQUAL "") + message(STATUS "${SAN_LD_PRELOAD}") + endif () + if ("${CMAKE_CXX_COMPILER_ID}" STREQUAL "Clang") + add_custom_command( + TARGET ${TMP_TARGET} POST_BUILD + COMMAND export ${LD_LIBRARY_PATH_} && ulimit -s 32768 && ${ASAN_OPTIONS_} ${UBSAN_OPTIONS_} ./${TMP_EXECUTABLE} + COMMENT "Run ${TMP_EXECUTABLE} with Sanitizer, ASAN(${ENABLE_ASAN}), UBSAN(${ENABLE_UBSAN})" + ) + else () + add_custom_command( + TARGET ${TMP_TARGET} POST_BUILD + COMMAND export ${LD_LIBRARY_PATH_} && export ${SAN_LD_PRELOAD} && ulimit -s 32768 && ${ASAN_OPTIONS_} ${UBSAN_OPTIONS_} ./${TMP_EXECUTABLE} + COMMENT "Run ${TMP_EXECUTABLE} with Sanitizer, ASAN(${ENABLE_ASAN}), UBSAN(${ENABLE_UBSAN})" + ) + endif () + else() + # 用例执行 + add_custom_command( + TARGET ${TMP_TARGET} POST_BUILD + COMMAND export ${LD_LIBRARY_PATH_} && ./${TMP_EXECUTABLE} + COMMENT "Run ${TMP_EXECUTABLE}" + ) + endif () + endif () +endfunction() + +# 生成覆盖率 +#[[ +调用参数: + one_value_keywords: + TARGET : 必选参数, 用于指定覆盖率生成所依赖的目标(POST_BUILD) + multi_value_keywords: + FILTER_DIRECTORIES : 可选参数, 用于指定覆盖率结果过滤目录 +]] +function(OpsTest_GenerateCoverage) + cmake_parse_arguments( + TMP + "" + "TARGET" + "FILTER_DIRECTORIES" + "" + ${ARGN} + ) + if (ENABLE_UT_EXEC) + if (ENABLE_GCOV) + find_program(LCOV lcov REQUIRED) + get_filename_component(GEN_COV_PY ${OPS_ADV_CMAKE_DIR}/scripts/utest/gen_coverage.py REALPATH) + get_filename_component(ASCEND_CANN_PACKAGE_PATH_PARENT "${ASCEND_CANN_PACKAGE_PATH}/../" REALPATH) + get_filename_component(GEM_COV_DATA_DIR "${CMAKE_CURRENT_BINARY_DIR}" REALPATH) + # 获取 gcc 默认头文件搜索路径 + execute_process( + COMMAND ${CMAKE_C_COMPILER} --print-sysroot-headers-suffix + RESULT_VARIABLE _RST + OUTPUT_VARIABLE _SUFFIX + ERROR_QUIET + ) + if (_RST) + get_filename_component(SYS_ROOT "/usr/include" REALPATH) + else () + get_filename_component(SYS_ROOT "${_SUFFIX}/usr/include" REALPATH) + endif () + list(REMOVE_DUPLICATES _FilterCmds) + add_custom_command( + TARGET ${TMP_TARGET} POST_BUILD + COMMAND ${HI_PYTHON} ${GEN_COV_PY} + "-s=${OPS_ADV_DIR}" + "-c=${GEM_COV_DATA_DIR}" + "-f=/tmp/*" + "-f=/usr/include/*" + "-f=${ASCEND_CANN_PACKAGE_PATH_PARENT}/*" + "-y=${OPS_ADV_DIR}/tests/test_config.yaml" + COMMENT "Generate coverage for ${TMP_TARGET}" + ) + endif () + endif () +endfunction() + +# 生成包含多个算子的 UT 可执行程序 +#[[ +]] +function(OpsTest_AddLaunch) + add_custom_target(ops_test_utest) + + if (NOT _OpsTestUt_UTestCaseLibraries AND NOT _OpsTestUt_UTestAclnnCaseLibraries) + # 当 _OpsTestUt_UTestCaseLibraries 与 _OpsTestUt_UTestAclnnCaseLibraries 都为空时 + # 说明区分领域的 TESTS_UT_OPS_TEST 选项被错误设置, 需排查对应触发 ut 的 python 脚本逻辑及 build.sh. + message(STATUS "_OpsTestUt_UTestCaseLibraries = ${_OpsTestUt_UTestCaseLibraries}.") + message(STATUS "_OpsTestUt_UTestAclnnCaseLibraries = ${_OpsTestUt_UTestAclnnCaseLibraries}") + message(STATUS "_OpsTestUt_UTestCaseLibraries and _OpsTestUt_UTestAclnnCaseLibraries both empty.") + return () + endif () + + set(_UTest_OpApiLibrary) + set(_UTest_OpProtoLibrary) + set(_UTest_OpTilingLibrary) + set(_UTest_Main) # UTest 用例可执行程序(常规) + set(_UTest_Main_Aclnn) # UTest 用例可执行程序(Aclnn) + set(_Utest_ExecutableList) # 为提高PR场景执行效率, PR场景各可执行文件并行执行, 否则串行执行 + + # OpApi 动态库(可选) + OpsTest_AddOpApiShared() + + # OpProto 动态库(可选) + OpsTest_AddOpProtoShared() + + # OpTiling 动态库(必选) + OpsTest_AddOpTilingShared() + + # UTest 用例可执行程序(常规) + if (_OpsTestUt_UTestCaseLibraries) + # 支持按算子分离可执行文件, 便于彼此冲突的 Kernel 侧编译宏设置 + foreach (_OpsTestUt_UTestCaseLibrary ${_OpsTestUt_UTestCaseLibraries}) + # 获取算子工程名 + string(REPLACE "_UTest_Case" "" _OpBrief ${_OpsTestUt_UTestCaseLibrary}) + string(REPLACE "UTest_" "" _OpBrief ${_OpBrief}) + + set(_UTest_Main ${UTest_NamePrefix}_Main_Normal_${_OpBrief}) + add_executable(${_UTest_Main}) + target_sources(${_UTest_Main} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tests/ut/framework_special/main.cpp) + target_compile_options(${_UTest_Main} PRIVATE -fPIC) + target_link_libraries(${_UTest_Main} + PRIVATE + -Wl,--no-as-needed + -Wl,--whole-archive + $ + gtest + $ + ${_OpsTestUt_UTestCaseLibrary} + ${UTest_NamePrefix}_Stubs + ${UTest_NamePrefix}_Utils + ${_UTest_OpApiLibrary} # 若算子在 UTest_{Op}_Common 内实现 Aclnn 相关执行逻辑, 则需要连接 + -Wl,--as-needed + -Wl,--no-whole-archive + c_sec + ) + add_dependencies(${_UTest_Main} ${_UTest_OpTilingLibrary}) + add_dependencies(ops_test_utest ${_UTest_Main}) + if (TESTS_UT_OPS_TEST_CI_PR) + OpsTest_RunLaunch(TARGET ${_UTest_Main} EXECUTABLE ${_UTest_Main}) + else () + list(APPEND _Utest_ExecutableList ${_UTest_Main}) + endif () + endforeach () + endif () + + # UTest 用例可执行程序(Aclnn) + set(_param + "--cann_path=${ASCEND_CANN_PACKAGE_PATH}" + "--cann_package_name=opp" + "get_package_version" + ) + + if (_OpsTestUt_UTestAclnnCaseLibraries) + # 支持按算子分离可执行文件, 便于彼此冲突的 Kernel 侧编译宏设置 + foreach (_OpsTestUt_UTestAclnnCaseLibrary ${_OpsTestUt_UTestAclnnCaseLibraries}) + # 获取算子工程名 + string(REPLACE "_UTest_Case_Aclnn" "" _OpBrief ${_OpsTestUt_UTestAclnnCaseLibrary}) + string(REPLACE "UTest_" "" _OpBrief ${_OpBrief}) + set(_UTest_Main_Aclnn ${UTest_NamePrefix}_Main_Aclnn_${_OpBrief}) + add_executable(${_UTest_Main_Aclnn}) + target_sources(${_UTest_Main_Aclnn} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tests/ut/framework_special/main_aclnn.cpp) + target_compile_options(${_UTest_Main_Aclnn} PRIVATE -fPIC) + target_link_libraries(${_UTest_Main_Aclnn} + PRIVATE + -Wl,--no-as-needed + -Wl,--whole-archive + $ + gtest + $ + $<$:$> + $<$:$> + ${_OpsTestUt_UTestAclnnCaseLibrary} + ${UTest_NamePrefix}_Stubs + ${UTest_NamePrefix}_Utils + ${_UTest_OpApiLibrary} + -Wl,--as-needed + -Wl,--no-whole-archive + c_sec + ) + add_dependencies(${_UTest_Main_Aclnn} ${_UTest_OpProtoLibrary} ${_UTest_OpTilingLibrary}) + add_dependencies(ops_test_utest ${_UTest_Main_Aclnn}) + if (TESTS_UT_OPS_TEST_CI_PR) + OpsTest_RunLaunch(TARGET ${_UTest_Main_Aclnn} EXECUTABLE ${_UTest_Main_Aclnn}) + else () + list(APPEND _Utest_ExecutableList ${_UTest_Main_Aclnn}) + endif () + endforeach () + endif () + + # 为提高PR场景执行效率, PR场景各可执行文件并行执行, 否则串行执行 + if (NOT TESTS_UT_OPS_TEST_CI_PR) + foreach (_Exe ${_Utest_ExecutableList}) + OpsTest_RunLaunch(TARGET ops_test_utest EXECUTABLE ${_Exe}) + endforeach () + endif () + + # 生成覆盖率 + set(_FilterDirectories) + list(REMOVE_DUPLICATES _FilterDirectories) + OpsTest_GenerateCoverage(TARGET ops_test_utest FILTER_DIRECTORIES ${_FilterDirectories}) +endfunction() + +# 添加算子UT路径 +#[[ +]] +function(OpsTestUt_AddSubdirectory) + cmake_parse_arguments( + TMP + "" + "" + "CLANG_NOT_SUPPORTED_OP_LIST" + "" + ${ARGN} + ) + + if ("${TESTS_UT_OPS_TEST}" STREQUAL "") + return() + elseif ("ALL" IN_LIST TESTS_UT_OPS_TEST OR "all" IN_LIST TESTS_UT_OPS_TEST) + file(GLOB sub_dirs ${CMAKE_CURRENT_SOURCE_DIR}/*) + foreach (_dir ${sub_dirs}) + if (IS_DIRECTORY ${_dir}) + get_filename_component(_op_type ${_dir} NAME) + if ("${CMAKE_C_COMPILER_ID}" STREQUAL "Clang" AND "${_op_type}" IN_LIST TMP_CLANG_NOT_SUPPORTED_OP_LIST) + message(STATUS "The UTest of ${_op_type} don't supported clang compiler yet.") + continue() + endif () + add_subdirectory(${_dir}) + endif () + endforeach () + else () + set(_added_op_type_list) + foreach (_op_type ${TESTS_UT_OPS_TEST}) + if (DEFINED ${_op_type}_alias) + set(_op_type ${${_op_type}_alias}) + endif () + if (NOT "${_op_type}" IN_LIST _added_op_type_list) + set(_dir ${CMAKE_CURRENT_SOURCE_DIR}/${_op_type}) + if (IS_DIRECTORY ${_dir}) + if ("${CMAKE_C_COMPILER_ID}" STREQUAL "Clang" AND "${_op_type}" IN_LIST TMP_CLANG_NOT_SUPPORTED_OP_LIST) + message(STATUS "The UTest of ${_op_type} don't supported clang compiler yet.") + continue() + endif () + add_subdirectory(${_dir}) + list(APPEND _added_op_type_list ${_op_type}) + endif () + endif () + endforeach () + endif () +endfunction() diff --git a/csrc/ascend/cmake/gen_ops_info.cmake b/csrc/ascend/cmake/gen_ops_info.cmake new file mode 100644 index 000000000..833e41feb --- /dev/null +++ b/csrc/ascend/cmake/gen_ops_info.cmake @@ -0,0 +1,368 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +################################################################################################### +# copy kernel src to tbe/ascendc path +################################################################################################### +function(kernel_src_copy) + set(oneValueArgs TARGET DST_DIR) + set(multiValueArgs IMPL_DIR) + cmake_parse_arguments(KNCPY "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + add_custom_target(${KNCPY_TARGET}) + foreach(OP_DIR ${KNCPY_IMPL_DIR}) + get_filename_component(OP_NAME ${OP_DIR} NAME) + message(STATUS "start copy kernel file: ${OP_NAME} to ${KNCPY_DST_DIR}") + if(NOT TARGET ${OP_NAME}_src_copy) + set(SRC_DIR ${OP_DIR}/op_kernel) + if(NOT EXISTS ${SRC_DIR}) + continue() + endif() + add_custom_target(${OP_NAME}_src_copy + COMMAND ${CMAKE_COMMAND} -E make_directory ${KNCPY_DST_DIR}/${OP_NAME} + COMMAND bash -c "find ${SRC_DIR} -mindepth 1 -maxdepth 1 -exec cp -r {} ${KNCPY_DST_DIR}/${OP_NAME} \\;" + VERBATIM + ) + add_dependencies(${KNCPY_TARGET} ${OP_NAME}_src_copy) + if(ENABLE_PACKAGE) + install( + DIRECTORY ${SRC_DIR}/ + DESTINATION ${IMPL_INSTALL_DIR}/${OP_NAME} + ) + endif() + endif() + endforeach() +endfunction() + +################################################################################################### +# generate operator dynamic python script for compile, generenate out path ${CMAKE_BINARY_DIR}/tbe, +# and install to packages/vendors/${VENDOR_NAME}_transformer/op_impl/ai_core/tbe/${VENDOR_NAME}_impl/dynamic +################################################################################################### +function(add_ops_impl_target) + set(oneValueArgs TARGET OPS_INFO_DIR IMPL_DIR OUT_DIR INSTALL_DIR) + cmake_parse_arguments(OPIMPL "" "${oneValueArgs}" "OPS_BATCH;OPS_ITERATE" ${ARGN}) + + add_custom_command(OUTPUT ${OPIMPL_OUT_DIR}/.impl_timestamp + COMMAND mkdir -m 700 -p ${OPIMPL_OUT_DIR}/dynamic + COMMAND ${ASCEND_PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/util/ascendc_impl_build.py + \"\" \"${OPIMPL_OPS_BATCH}\" \"${OPIMPL_OPS_ITERATE}\" + ${OPIMPL_IMPL_DIR} ${OPIMPL_OUT_DIR}/dynamic ${ASCEND_AUTOGEN_PATH} + --opsinfo-dir ${OPIMPL_OPS_INFO_DIR} ${OPIMPL_OPS_INFO_DIR}/inner ${OPIMPL_OPS_INFO_DIR}/exc + COMMAND rm -rf ${OPIMPL_OUT_DIR}/.impl_timestamp + COMMAND touch ${OPIMPL_OUT_DIR}/.impl_timestamp + DEPENDS ${CMAKE_SOURCE_DIR}/scripts/util/ascendc_impl_build.py + ) + add_custom_target(${OPIMPL_TARGET} ALL + DEPENDS ${OPIMPL_OUT_DIR}/.impl_timestamp + ) + + file(GLOB dynamic_impl ${OPIMPL_OUT_DIR}/dynamic/*.py) + if(ENABLE_PACKAGE) + install( + FILES ${dynamic_impl} + DESTINATION ${OPIMPL_INSTALL_DIR} + OPTIONAL + ) + endif() +endfunction() + +################################################################################################### +# generate aic-${compute_unit}-ops-info.json from aic-${compute_unit}-ops-info.ini +# generate outpath: ${CMAKE_BINARY_DIR}/tbe/op_info_cfg/ai_core/${compute_unit}/ +# install path: packages/vendors/${VENDOR_NAME}_transformer/op_impl/ai_core/tbe/config/${compute_unit} +################################################################################################### +function(add_ops_info_target) + set(oneValueArgs TARGET OPS_INFO_DIR COMPUTE_UNIT OUTPUT INSTALL_DIR) + cmake_parse_arguments(OPINFO "" "${oneValueArgs}" "" ${ARGN}) + get_filename_component(opinfo_file_path "${OPINFO_OUTPUT}" DIRECTORY) + add_custom_command(OUTPUT ${OPINFO_OUTPUT} + COMMAND mkdir -p ${opinfo_file_path} + COMMAND ${ASCEND_PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/util/parse_ini_to_json.py + ${OPINFO_OPS_INFO_DIR}/aic-${OPINFO_COMPUTE_UNIT}-ops-info.ini + ${OPINFO_OPS_INFO_DIR}/inner/aic-${OPINFO_COMPUTE_UNIT}-ops-info.ini + ${OPINFO_OPS_INFO_DIR}/exc/aic-${OPINFO_COMPUTE_UNIT}-ops-info.ini + ${OPINFO_OUTPUT} + ) + add_custom_target(${OPINFO_TARGET} ALL + DEPENDS ${OPINFO_OUTPUT} + ) + + if(ENABLE_PACKAGE) + install(FILES ${OPINFO_OUTPUT} + DESTINATION ${OPINFO_INSTALL_DIR} + ) + endif() +endfunction() + +################################################################################################### +# merge ops info ini in aclnn/aclnn_inner/aclnn_exc to a total ini file +# srcpath: ${ASCEND_AUTOGEN_PATH} +# generate outpath: ${CMAKE_BINARY_DIR}/tbe/config +################################################################################################### +function(merge_ini_files) + set(oneValueArgs TARGET OPS_INFO_DIR COMPUTE_UNIT) + cmake_parse_arguments(MGINI "" "${oneValueArgs}" "" ${ARGN}) + add_custom_command(OUTPUT ${ASCEND_KERNEL_CONF_DST}/aic-${MGINI_COMPUTE_UNIT}-ops-info.ini + COMMAND touch ${MGINI_OPS_INFO_DIR}/aic-merged-${MGINI_COMPUTE_UNIT}-ops-info.ini + COMMAND ${ASCEND_PYTHON_EXECUTABLE} ${OPS_KERNEL_BINARY_SCRIPT}/merge_ini_files.py + ${MGINI_OPS_INFO_DIR}/aic-${MGINI_COMPUTE_UNIT}-ops-info.ini + ${MGINI_OPS_INFO_DIR}/inner/aic-${MGINI_COMPUTE_UNIT}-ops-info.ini + ${MGINI_OPS_INFO_DIR}/exc/aic-${MGINI_COMPUTE_UNIT}-ops-info.ini + --output-file ${ASCEND_KERNEL_CONF_DST}/aic-${MGINI_COMPUTE_UNIT}-ops-info.ini + ) + add_custom_target(${MGINI_TARGET} ALL + DEPENDS ${ASCEND_KERNEL_CONF_DST}/aic-${MGINI_COMPUTE_UNIT}-ops-info.ini + ) +endfunction() + +# ################################################################################################## +# merge ops proto headers in aclnn/aclnn_inner/aclnn_exc to a total proto file +# srcpath: ${ASCEND_AUTOGEN_PATH} +# generate outpath: ${CMAKE_BINARY_DIR}/tbe/graph +# ################################################################################################## +function(merge_graph_headers) + set(oneValueArgs TARGET OUT_DIR) + cmake_parse_arguments(MGPROTO "" "${oneValueArgs}" "" ${ARGN}) + get_target_property(proto_headers ${GRAPH_PLUGIN_NAME}_proto_headers INTERFACE_SOURCES) + add_custom_command(OUTPUT ${MGPROTO_OUT_DIR}/ops_proto_math.h + COMMAND ${ASCEND_PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/util/merge_proto.py + ${proto_headers} + --output-file ${MGPROTO_OUT_DIR}/ops_proto_math.h + ) + add_custom_target(${MGPROTO_TARGET} ALL + DEPENDS ${MGPROTO_OUT_DIR}/ops_proto_math.h + ) +endfunction() + +################################################################################################### +# generate binary compile shell script and binary json +# srcpath: ${ASCEND_AUTOGEN_PATH} +# outpath: ${CMAKE_BINARY_DIR}/binary/${compute_unit} +################################################################################################### +function(generate_bin_scripts) + set(oneValueArgs TARGET OP_NAME OPS_INFO_DIR COMPUTE_UNIT OUT_DIR) + cmake_parse_arguments(GENBIN "" "${oneValueArgs}" "" ${ARGN}) + file(MAKE_DIRECTORY ${GENBIN_OUT_DIR}/gen) + file(MAKE_DIRECTORY ${GENBIN_OUT_DIR}/gen/${GENBIN_OP_NAME}) + message(STATUS "start generate_bin_scripts for op: ${GENBIN_OP_NAME}") + add_custom_target(generate_bin_scripts_${GENBIN_COMPUTE_UNIT}_${GENBIN_OP_NAME} + COMMAND ${ASCEND_PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/util/ascendc_bin_param_build.py + ${GENBIN_OPS_INFO_DIR}/aic-${GENBIN_COMPUTE_UNIT}-ops-info.ini + ${GENBIN_OUT_DIR}/gen/${GENBIN_OP_NAME} ${GENBIN_COMPUTE_UNIT} + --opc-config-file ${ASCEND_AUTOGEN_PATH}/${CUSTOM_OPC_OPTIONS} + --ops ${GENBIN_OP_NAME} + COMMAND ${ASCEND_PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/util/ascendc_bin_param_build.py + ${GENBIN_OPS_INFO_DIR}/inner/aic-${GENBIN_COMPUTE_UNIT}-ops-info.ini + ${GENBIN_OUT_DIR}/gen/${GENBIN_OP_NAME} ${GENBIN_COMPUTE_UNIT} + --opc-config-file ${ASCEND_AUTOGEN_PATH}/${CUSTOM_OPC_OPTIONS} + --ops ${GENBIN_OP_NAME} + COMMAND ${ASCEND_PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/util/ascendc_bin_param_build.py + ${GENBIN_OPS_INFO_DIR}/exc/aic-${GENBIN_COMPUTE_UNIT}-ops-info.ini + ${GENBIN_OUT_DIR}/gen/${GENBIN_OP_NAME} ${GENBIN_COMPUTE_UNIT} + --opc-config-file ${ASCEND_AUTOGEN_PATH}/${CUSTOM_OPC_OPTIONS} + --ops ${GENBIN_OP_NAME} + ) + if(NOT TARGET ${GENBIN_TARGET}) + add_custom_target(${GENBIN_TARGET}) + endif() + add_dependencies(${GENBIN_TARGET} generate_bin_scripts_${GENBIN_COMPUTE_UNIT}_${GENBIN_OP_NAME} + ) +endfunction() + +################################################################################################### +# copy binary config from op_host/config to tbe/config path +################################################################################################### +function(binary_config_copy) + set(oneValueArgs TARGET OP_NAME CONF_DIR DST_DIR COMPUTE_UNIT) + cmake_parse_arguments(CNFCPY "" "${oneValueArgs}" "" ${ARGN}) + file(MAKE_DIRECTORY ${CNFCPY_DST_DIR}/${CNFCPY_COMPUTE_UNIT}/${CNFCPY_OP_NAME}) + add_custom_target(${CNFCPY_TARGET} + COMMAND rm -rf ${CNFCPY_DST_DIR}/${CNFCPY_COMPUTE_UNIT}/${CNFCPY_OP_NAME}/* + COMMAND cp -r ${CNFCPY_CONF_DIR}/${CNFCPY_COMPUTE_UNIT}/* ${CNFCPY_DST_DIR}/${CNFCPY_COMPUTE_UNIT}/${CNFCPY_OP_NAME} + ) +endfunction() + +################################################################################################### +# compile binary from op_host/config binary json files +# generate outpath: ${CMAKE_BINARY_DIR}/binary/${compute_unit}/bin +# install path: ${BIN_KERNEL_INSTALL_DIR}/${compute_unit} +################################################################################################### +function(compile_from_config) + set(oneValueArgs TARGET OP_NAME OPS_INFO_DIR IMPL_DIR CONFIG_DIR OP_PYTHON_DIR OUT_DIR INSTALL_DIR COMPUTE_UNIT) + cmake_parse_arguments(CONFCMP "" "${oneValueArgs}" "" ${ARGN}) + file(MAKE_DIRECTORY ${CONFCMP_OUT_DIR}/src) + file(MAKE_DIRECTORY ${CONFCMP_OUT_DIR}/bin) + file(MAKE_DIRECTORY ${CONFCMP_OUT_DIR}/gen) + snake_to_camel("${CONFCMP_OP_NAME}" OP_TYPE) + message(STATUS "start to compile op: ${CONFCMP_OP_NAME}, op_type: ${OP_TYPE}") + # add Environment Variable Configurations of python & ccache + set(_ASCENDC_ENV_VAR) + list(APPEND _ASCENDC_ENV_VAR export HI_PYTHON=${ASCEND_PYTHON_EXECUTABLE} &&) + # whether need judging CMAKE_C_COMPILER_LAUNCHER + if(${CMAKE_CXX_COMPILER_LAUNCHER} MATCHES "ccache$") + list(APPEND _ASCENDC_ENV_VAR export ASCENDC_CCACHE_EXECUTABLE=${CMAKE_CXX_COMPILER_LAUNCHER} &&) + endif() + # copy binary config file to tbe/config + binary_config_copy( + TARGET bin_conf_${CONFCMP_OP_NAME}_${CONFCMP_COMPUTE_UNIT}_copy + OP_NAME ${CONFCMP_OP_NAME} + CONF_DIR ${CONFCMP_CONFIG_DIR} + DST_DIR ${ASCEND_KERNEL_CONF_DST} + COMPUTE_UNIT ${CONFCMP_COMPUTE_UNIT} + ) + + add_custom_target(config_compile_${CONFCMP_COMPUTE_UNIT}_${CONFCMP_OP_NAME} + COMMAND ${_ASCENDC_ENV_VAR} bash ${OPS_KERNEL_BINARY_SCRIPT}/build_binary_single_op.sh + ${OP_TYPE} + ${CONFCMP_COMPUTE_UNIT} + ${CONFCMP_OUT_DIR}/bin + WORKING_DIRECTORY ${OPS_KERNEL_BINARY_SCRIPT} + DEPENDS ${ASCEND_KERNEL_CONF_DST}/aic-${CONFCMP_COMPUTE_UNIT}-ops-info.ini + ascendc_kernel_src_copy + bin_conf_${CONFCMP_OP_NAME}_${CONFCMP_COMPUTE_UNIT}_copy + ) + + if(NOT TARGET binary) + add_custom_target(binary) + endif() + add_custom_target(${CONFCMP_TARGET} + COMMAND cp -r ${CONFCMP_IMPL_DIR}/*.* ${CONFCMP_OUT_DIR}/src + COMMAND cp ${CONFCMP_OP_PYTHON_DIR}/${CONFCMP_OP_NAME}.py ${CONFCMP_OUT_DIR}/src + ) + add_dependencies(binary config_compile_${CONFCMP_COMPUTE_UNIT}_${CONFCMP_OP_NAME} ${CONFCMP_TARGET}) + + if(ENABLE_PACKAGE) + install(DIRECTORY ${CONFCMP_OUT_DIR}/bin/${CONFCMP_COMPUTE_UNIT}/${CONFCMP_OP_NAME} + DESTINATION ${BIN_KERNEL_INSTALL_DIR}/${CONFCMP_COMPUTE_UNIT} OPTIONAL + ) + install(FILES ${CONFCMP_OUT_DIR}/bin/config/${CONFCMP_COMPUTE_UNIT}/${CONFCMP_OP_NAME}.json + DESTINATION ${BIN_KERNEL_CONFIG_INSTALL_DIR}/${CONFCMP_COMPUTE_UNIT} OPTIONAL + ) + endif() +endfunction() + +################################################################################################### +# generate binary_info_config.json +# generate outpath: ${CMAKE_BINARY_DIR}/binary/${compute_unit}/bin/config +# install path: packages/vendors/${VENDOR_NAME}_transformer/op_impl/ai_core/tbe/kernel/config +################################################################################################### +function(gen_binary_info_config_json) + set(oneValueArgs TARGET BIN_DIR COMPUTE_UNIT) + cmake_parse_arguments(GENBIN_INFOCFG "" "${oneValueArgs}" "" ${ARGN}) + + add_custom_command(OUTPUT ${GENBIN_INFOCFG_BIN_DIR}/bin/config/${GENBIN_INFOCFG_COMPUTE_UNIT}/binary_info_config.json + COMMAND ${ASCEND_PYTHON_EXECUTABLE} ${OPS_KERNEL_BINARY_SCRIPT}/gen_binary_info_config.py + ${GENBIN_INFOCFG_BIN_DIR}/bin + ${GENBIN_INFOCFG_COMPUTE_UNIT} + DEPENDS ${GENBIN_INFOCFG_BIN_DIR}/bin/config + ) + add_custom_target(${GENBIN_INFOCFG_TARGET} + DEPENDS ${GENBIN_INFOCFG_BIN_DIR}/bin/config/${GENBIN_INFOCFG_COMPUTE_UNIT}/binary_info_config.json + ) + + if(NOT TARGET gen_bin_info_config) + add_custom_target(gen_bin_info_config) + endif() + add_dependencies(gen_bin_info_config ${GENBIN_INFOCFG_TARGET}) + + if(ENABLE_PACKAGE) + install( + FILES ${GENBIN_INFOCFG_BIN_DIR}/bin/config/${GENBIN_INFOCFG_COMPUTE_UNIT}/binary_info_config.json + DESTINATION ${BIN_KERNEL_CONFIG_INSTALL_DIR}/${GENBIN_INFOCFG_COMPUTE_UNIT} OPTIONAL + ) + endif() +endfunction() + +# binary compile +function(gen_ops_info_and_python) + gen_aclnn_with_opdef() + if(NOT TARGET opbuild_custom_gen_aclnn_all) + message(STATUS "no need build binary, for all the ops do not have any operator def") + return() + endif() + + kernel_src_copy( + TARGET ascendc_kernel_src_copy + OP_LIST ${COMPILED_OPS} + IMPL_DIR ${COMPILED_OP_DIRS} + DST_DIR ${ASCEND_KERNEL_SRC_DST} + ) + + add_ops_impl_target( + TARGET ascendc_impl_gen + OPS_INFO_DIR ${ASCEND_AUTOGEN_PATH} + IMPL_DIR ${ASCEND_KERNEL_SRC_DST} + OUT_DIR ${CMAKE_BINARY_DIR}/tbe + INSTALL_DIR ${IMPL_DYNAMIC_INSTALL_DIR} + ) + + merge_graph_headers( + TARGET merge_ops_proto ALL + OUT_DIR ${ASCEND_GRAPH_CONF_DST} + ) + + set(ascendc_impl_gen_depends ascendc_kernel_src_copy opbuild_custom_gen_aclnn_all) + foreach(compute_unit ${ASCEND_COMPUTE_UNIT}) + # generate aic-${compute_unit}-ops-info.json, operator infos + add_ops_info_target( + TARGET ops_info_gen_${compute_unit} + OUTPUT ${CMAKE_BINARY_DIR}/tbe/op_info_cfg/ai_core/${compute_unit}/aic-${compute_unit}-ops-info.json + OPS_INFO_DIR ${ASCEND_AUTOGEN_PATH} + COMPUTE_UNIT ${compute_unit} + INSTALL_DIR ${OPS_INFO_INSTALL_DIR} + ) + + # merge ops info ini files + merge_ini_files(TARGET merge_ini_${compute_unit} + OPS_INFO_DIR ${ASCEND_AUTOGEN_PATH} + COMPUTE_UNIT ${compute_unit} + ) + list(APPEND ascendc_impl_gen_depends ops_info_gen_${compute_unit}) + endforeach() + add_dependencies(ascendc_impl_gen ${ascendc_impl_gen_depends}) + + if(ENABLE_BINARY OR ENABLE_CUSTOM) + foreach(compute_unit ${ASCEND_COMPUTE_UNIT}) + foreach(OP_DIR ${COMPILED_OP_DIRS}) + get_filename_component(op_name ${OP_DIR} NAME) + # generate opc shell scripts for autogen binary config ops + generate_bin_scripts( + TARGET gen_bin_scripts + OP_NAME ${op_name} + OPS_INFO_DIR ${ASCEND_AUTOGEN_PATH} + COMPUTE_UNIT ${compute_unit} + OUT_DIR ${CMAKE_BINARY_DIR}/binary/${compute_unit} + ) + if(EXISTS ${OP_DIR}/op_host/config/${compute_unit}/${op_name}_binary.json) + # binary compile from binary json config + message(STATUS "[INFO] On [${compute_unit}], [${op_name}] compile binary with self config.") + compile_from_config( + TARGET ascendc_bin_${compute_unit}_${op_name} + OP_NAME ${op_name} + OPS_INFO_DIR ${ASCEND_AUTOGEN_PATH} + IMPL_DIR ${OP_DIR}/op_kernel + CONFIG_DIR ${OP_DIR}/op_host/config + OP_PYTHON_DIR ${CMAKE_BINARY_DIR}/tbe/dynamic + OUT_DIR ${CMAKE_BINARY_DIR}/binary/${compute_unit} + INSTALL_DIR ${BIN_KERNEL_INSTALL_DIR} + COMPUTE_UNIT ${compute_unit} + ) + add_dependencies(ascendc_bin_${compute_unit}_${op_name} merge_ini_${compute_unit} ascendc_impl_gen) + endif() + endforeach() + + # generate binary_info_config.json + gen_binary_info_config_json( + TARGET gen_bin_info_config_${compute_unit} + BIN_DIR ${CMAKE_BINARY_DIR}/binary/${compute_unit} + COMPUTE_UNIT ${compute_unit} + ) + endforeach() + endif() +endfunction() diff --git a/csrc/ascend/cmake/intf.cmake b/csrc/ascend/cmake/intf.cmake new file mode 100644 index 000000000..5b4240f19 --- /dev/null +++ b/csrc/ascend/cmake/intf.cmake @@ -0,0 +1,21 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if (BUILD_OPEN_PROJECT) + include(${OPS_ADV_CMAKE_DIR}/intf_pub.cmake) + + if (TESTS_UT_OPS_TEST) + include(${OPS_ADV_CMAKE_DIR}/intf_pub_utest.cmake) + endif () + + if (TESTS_EXAMPLE_OPS_TEST) + include(${OPS_ADV_CMAKE_DIR}/intf_pub_examples.cmake) + endif () +endif () diff --git a/csrc/ascend/cmake/intf_pub.cmake b/csrc/ascend/cmake/intf_pub.cmake new file mode 100644 index 000000000..8ab03dc74 --- /dev/null +++ b/csrc/ascend/cmake/intf_pub.cmake @@ -0,0 +1,167 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +# Custom 包场景, Host 侧各 Target 公共编译配置 +# 注意: 为保证与 built-in 包编译流程兼容, intf_pub 名称不可变更 +add_library(intf_pub INTERFACE) +target_include_directories(intf_pub + INTERFACE + ${ASCEND_CANN_PACKAGE_PATH}/include + ${ASCEND_CANN_PACKAGE_PATH}/include/external + ${ASCEND_CANN_PACKAGE_PATH}/include/platform + ${ASCEND_CANN_PACKAGE_PATH}/include/experiment/runtime + ${ASCEND_CANN_PACKAGE_PATH}/include/experiment/msprof +) +target_link_directories(intf_pub + INTERFACE + ${ASCEND_CANN_PACKAGE_PATH}/lib64 +) +target_compile_options(intf_pub + INTERFACE + -fPIC + -O2 + -Wall -Wundef -Wcast-qual -Wpointer-arith -Wdate-time + -Wfloat-equal -Wformat=2 -Wshadow + -Wsign-compare -Wunused-macros -Wvla -Wdisabled-optimization -Wempty-body -Wignored-qualifiers + -Wimplicit-fallthrough=3 -Wtype-limits -Wshift-negative-value -Wswitch-default + -Wframe-larger-than=98304 -Woverloaded-virtual + -Wnon-virtual-dtor -Wshift-overflow=2 -Wshift-count-overflow + -Wwrite-strings -Wmissing-format-attribute -Wformat-nonliteral + -Wdelete-non-virtual-dtor -Wduplicated-cond + -Wtrampolines -Wsized-deallocation -Wlogical-op -Wsuggest-attribute=format + -Wduplicated-branches + -Wformat-signedness + -Wreturn-local-addr -Wextra + -Wredundant-decls -Wfloat-conversion + -Wno-write-strings -Wall -Wno-dangling-else -Wno-comment -Wno-conversion-null -Wno-return-type + -Wno-unknown-pragmas -Wno-sign-compare + -Wno-error=undef + -Wno-error=comment + -Wno-error=conversion-null + -Wno-error=dangling-else + -Wno-error=return-type + -Wno-error=shadow + -Wno-error=sign-compare + -Wno-error=unknown-pragmas + -Wno-error=unused-parameter + -Wno-error=cast-qual + -Wno-error=format= + -Wno-error=maybe-uninitialized + -Wno-error=missing-field-initializers + -Wno-error=redundant-decls + -Wno-error=unused-variable + $<$:-Wnested-externs> + $<$:-g> + $,-fstack-protector-strong,-fstack-protector-all> + $<$:-fprofile-arcs -ftest-coverage> +) +target_compile_definitions(intf_pub + INTERFACE + $<$:_GLIBCXX_USE_CXX11_ABI=0> # 必须设置, 以保证与 CANN 包内其他依赖库兼容 + $<$:_FORTIFY_SOURCE=2> +) +target_link_options(intf_pub + INTERFACE + $<$,EXECUTABLE>:-pie> + $<$:-s> + -Wl,-z,relro + -Wl,-z,now + -Wl,-z,noexecstack + $<$:-fprofile-arcs -ftest-coverage> + +) + +# intf_pub_cxx14 for c++14 +add_library(intf_pub_cxx14 INTERFACE) +target_compile_options(intf_pub_cxx14 INTERFACE + -Wall + -fPIC + $,-fstack-protector-strong,-fstack-protector-all> + $<$:-g> + $<$:-std=c++14> + $<$:-fprofile-arcs -ftest-coverage> +) +target_compile_definitions(intf_pub_cxx14 INTERFACE + _GLIBCXX_USE_CXX11_ABI=0 + $<$:CFG_BUILD_NDEBUG> + $<$:CFG_BUILD_DEBUG> + WIN64=1 + LINUX=0 +) +target_link_options(intf_pub_cxx14 INTERFACE + -Wl,-z,relro + -Wl,-z,now + -Wl,-z,noexecstack + $<$:-s> + $<$:-Wl,--build-id=none> + $<$:-fprofile-arcs -ftest-coverage> +) +target_link_directories(intf_pub_cxx14 INTERFACE) +target_link_libraries(intf_pub_cxx14 INTERFACE + -lpthread +) + +# intf_pub_cxx17 for c++17 +add_library(intf_pub_cxx17 INTERFACE) +target_compile_options(intf_pub_cxx17 INTERFACE + -Wall + -fPIC + $,-fstack-protector-strong,-fstack-protector-all> + $<$:-g> + $<$:-std=c++17> + $<$:-fprofile-arcs -ftest-coverage> + ) +target_compile_definitions(intf_pub_cxx17 INTERFACE + _GLIBCXX_USE_CXX11_ABI=0 + $<$:CFG_BUILD_NDEBUG> + $<$:CFG_BUILD_DEBUG> + WIN64=1 + LINUX=0) +target_link_options(intf_pub_cxx17 INTERFACE + -Wl,-z,relro + -Wl,-z,now + -Wl,-z,noexecstack + $<$:-s> + $<$:-Wl,--build-id=none> + $<$:-fprofile-arcs -ftest-coverage> + ) +target_link_directories(intf_pub_cxx17 INTERFACE) +target_link_libraries(intf_pub_cxx17 INTERFACE + -lpthread) + +#########intf_pub_aicpu######### +add_library(intf_pub_aicpu INTERFACE) +target_compile_options(intf_pub_aicpu INTERFACE + -Wall + -fPIC + $,-fstack-protector-strong,-fstack-protector-all> + $<$:-g> + $<$:-std=c++17> + $<$:-fprofile-arcs -ftest-coverage> +) +target_compile_definitions(intf_pub_aicpu INTERFACE + $<$>:_GLIBCXX_USE_CXX11_ABI=0> + $<$:_GLIBCXX_USE_CXX11_ABI=1> + $<$:CFG_BUILD_NDEBUG> + $<$:CFG_BUILD_DEBUG> + WIN64=1 + LINUX=0 +) +target_link_options(intf_pub_aicpu INTERFACE + -Wl,-z,relro + -Wl,-z,now + -Wl,-z,noexecstack + $<$:-Wl,--build-id=none> + $<$:-fprofile-arcs -ftest-coverage> +) +target_link_directories(intf_pub_aicpu INTERFACE) +target_link_libraries(intf_pub_aicpu INTERFACE + -lpthread +) diff --git a/csrc/ascend/cmake/intf_pub_examples.cmake b/csrc/ascend/cmake/intf_pub_examples.cmake new file mode 100644 index 000000000..33f7d44fa --- /dev/null +++ b/csrc/ascend/cmake/intf_pub_examples.cmake @@ -0,0 +1,9 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- diff --git a/csrc/ascend/cmake/intf_pub_linux.cmake b/csrc/ascend/cmake/intf_pub_linux.cmake new file mode 100644 index 000000000..18de76ce5 --- /dev/null +++ b/csrc/ascend/cmake/intf_pub_linux.cmake @@ -0,0 +1,114 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +if(TARGET intf_pub) + message(STATUS "intf_pub has been found, no need add library") + return() +endif() + +# intf_pub for c++11 +add_library(intf_pub INTERFACE) +target_compile_options(intf_pub INTERFACE + -Wall + -fPIC + $,-fstack-protector-strong,-fstack-protector-all> + $<$:-std=c++14> +) +target_compile_definitions(intf_pub INTERFACE + _GLIBCXX_USE_CXX11_ABI=0 + $<$:CFG_BUILD_NDEBUG> + $<$:CFG_BUILD_DEBUG> + WIN64=1 + LINUX=0 +) +target_link_options(intf_pub INTERFACE + -Wl,-z,relro + -Wl,-z,now + -Wl,-z,noexecstack + $<$:-Wl,--build-id=none> +) +target_link_directories(intf_pub INTERFACE) +target_link_libraries(intf_pub INTERFACE + -lpthread +) + +# intf_pub_cxx14 for c++14 +add_library(intf_pub_cxx14 INTERFACE) +target_compile_options(intf_pub_cxx14 INTERFACE + -Wall + -fPIC + $,-fstack-protector-strong,-fstack-protector-all> + $<$:-std=c++14> +) +target_compile_definitions(intf_pub_cxx14 INTERFACE + _GLIBCXX_USE_CXX11_ABI=0 + $<$:CFG_BUILD_NDEBUG> + $<$:CFG_BUILD_DEBUG> + WIN64=1 + LINUX=0 +) +target_link_options(intf_pub_cxx14 INTERFACE + -Wl,-z,relro + -Wl,-z,now + -Wl,-z,noexecstack + $<$:-Wl,--build-id=none> +) +target_link_directories(intf_pub_cxx14 INTERFACE) +target_link_libraries(intf_pub_cxx14 INTERFACE + -lpthread +) + +# intf_pub_cxx14 for c++17 +add_library(intf_pub_cxx17 INTERFACE) +target_compile_options(intf_pub_cxx17 INTERFACE + -Wall + -fPIC + $,-fstack-protector-strong,-fstack-protector-all> + $<$:-std=c++17>) +target_compile_definitions(intf_pub_cxx17 INTERFACE + _GLIBCXX_USE_CXX11_ABI=0 + $<$:CFG_BUILD_NDEBUG> + $<$:CFG_BUILD_DEBUG> + WIN64=1 + LINUX=0) +target_link_options(intf_pub_cxx17 INTERFACE + -Wl,-z,relro + -Wl,-z,now + -Wl,-z,noexecstack + $<$:-Wl,--build-id=none>) +target_link_directories(intf_pub_cxx17 INTERFACE) +target_link_libraries(intf_pub_cxx17 INTERFACE + -lpthread) + +#########intf_pub_aicpu######### +add_library(intf_pub_aicpu INTERFACE) +target_compile_options(intf_pub_aicpu INTERFACE + -Wall + -fPIC + $,-fstack-protector-strong,-fstack-protector-all> + $<$:-std=c++11> +) +target_compile_definitions(intf_pub_aicpu INTERFACE + $<$>:_GLIBCXX_USE_CXX11_ABI=0> + $<$:_GLIBCXX_USE_CXX11_ABI=1> + $<$:CFG_BUILD_NDEBUG> + $<$:CFG_BUILD_DEBUG> + WIN64=1 + LINUX=0 +) +target_link_options(intf_pub_aicpu INTERFACE + -Wl,-z,relro + -Wl,-z,now + -Wl,-z,noexecstack + $<$:-Wl,--build-id=none> +) +target_link_directories(intf_pub_aicpu INTERFACE) +target_link_libraries(intf_pub_aicpu INTERFACE + -lpthread +) \ No newline at end of file diff --git a/csrc/ascend/cmake/intf_pub_llt_gccnative.cmake b/csrc/ascend/cmake/intf_pub_llt_gccnative.cmake new file mode 100644 index 000000000..6cdb8f05c --- /dev/null +++ b/csrc/ascend/cmake/intf_pub_llt_gccnative.cmake @@ -0,0 +1,134 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +if(TARGET intf_llt_pub) + message(STATUS "intf_llt_pub has been found, no need add library") + return() +endif() +add_library(intf_llt_pub INTERFACE) +target_include_directories(intf_llt_pub INTERFACE + ${GTEST_INCLUDE} + ${GMOCK_INCLUDE} +) +target_compile_definitions(intf_llt_pub INTERFACE + _GLIBCXX_USE_CXX11_ABI=0 + CFG_BUILD_DEBUG +) +target_compile_options(intf_llt_pub INTERFACE + -g + --coverage + -fprofile-arcs + -ftest-coverage + -w + $<$:-std=c++11> + -fPIC +) +target_link_options(intf_llt_pub INTERFACE + -fprofile-arcs -ftest-coverage +) +target_link_libraries(intf_llt_pub INTERFACE + gcov + pthread +) + +if(TARGET intf_llt_pub_asan) + message(STATUS "intf_llt_pub_asan has been found, no need add library") + return() +endif() +add_library(intf_llt_pub_asan INTERFACE) +target_include_directories(intf_llt_pub_asan INTERFACE + ${GTEST_INCLUDE} + ${GMOCK_INCLUDE} +) +target_compile_definitions(intf_llt_pub_asan INTERFACE + _GLIBCXX_USE_CXX11_ABI=0 + CFG_BUILD_DEBUG +) +target_compile_options(intf_llt_pub_asan INTERFACE + -g + --coverage + -fprofile-arcs + -ftest-coverage + -w + $<$:-std=c++11> + $<$:-fsanitize=address -fsanitize-recover=address,all -fno-omit-frame-pointer -g> + -fPIC +) +target_link_options(intf_llt_pub_asan INTERFACE + -fprofile-arcs -ftest-coverage + $<$:-fsanitize=address> +) +target_link_libraries(intf_llt_pub_asan INTERFACE + gcov + pthread +) + +if(TARGET intf_llt_pub_asan_cxx14) + message(STATUS "intf_llt_pub_asan_cxx14 has been found, no need add library") + return() +endif() +add_library(intf_llt_pub_asan_cxx14 INTERFACE) +target_include_directories(intf_llt_pub_asan_cxx14 INTERFACE + ${GTEST_INCLUDE} + ${GMOCK_INCLUDE} +) +target_compile_definitions(intf_llt_pub_asan_cxx14 INTERFACE + _GLIBCXX_USE_CXX11_ABI=0 + CFG_BUILD_DEBUG +) +target_compile_options(intf_llt_pub_asan_cxx14 INTERFACE + -g + --coverage + -fprofile-arcs + -ftest-coverage + -w + $<$:-std=c++14> + $<$:-fsanitize=address -fsanitize-recover=address,all -fno-omit-frame-pointer -g> + -fPIC +) +target_link_options(intf_llt_pub_asan_cxx14 INTERFACE + -fprofile-arcs -ftest-coverage + $<$:-fsanitize=address> +) +target_link_libraries(intf_llt_pub_asan_cxx14 INTERFACE + gcov + pthread +) + +if(TARGET intf_llt_pub_asan_cxx17) + message(STATUS "intf_llt_pub_asan_cxx17 has been found, no need add library") + return() +endif() +add_library(intf_llt_pub_asan_cxx17 INTERFACE) +target_include_directories(intf_llt_pub_asan_cxx17 INTERFACE + ${GTEST_INCLUDE} + ${GMOCK_INCLUDE} +) +target_compile_definitions(intf_llt_pub_asan_cxx17 INTERFACE + _GLIBCXX_USE_CXX11_ABI=0 + CFG_BUILD_DEBUG +) +target_compile_options(intf_llt_pub_asan_cxx17 INTERFACE + -g + --coverage + -fprofile-arcs + -ftest-coverage + -w + $<$:-std=c++17> + $<$:-fsanitize=address -fsanitize-recover=address,all -fno-omit-frame-pointer -g> + -fPIC +) +target_link_options(intf_llt_pub_asan_cxx17 INTERFACE + -fprofile-arcs -ftest-coverage + $<$:-fsanitize=address> +) +target_link_libraries(intf_llt_pub_asan_cxx17 INTERFACE + gcov + pthread +) diff --git a/csrc/ascend/cmake/intf_pub_utest.cmake b/csrc/ascend/cmake/intf_pub_utest.cmake new file mode 100644 index 000000000..aeff595f7 --- /dev/null +++ b/csrc/ascend/cmake/intf_pub_utest.cmake @@ -0,0 +1,94 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +# UTest 场景, 编译 Target 名称公共前缀 +set(UTest_NamePrefix UTest) + +# UTest 场景, 公共配置 +add_library(intf_pub_utest INTERFACE) +target_compile_definitions(intf_pub_utest + INTERFACE + $<$:_GLIBCXX_USE_CXX11_ABI=0> # 必须设置, 以保证与 CANN 包内其他依赖库兼容 + ASCENDC_OP_TEST + ASCENDC_OP_TEST_UT + $<$:CFG_BUILD_NDEBUG> + $<$:CFG_BUILD_DEBUG> +) +target_compile_options(intf_pub_utest + INTERFACE + -fPIC + $<$:$,-fstack-protector-strong,-fstack-protector-all>> + $<$:$,-fstack-protector-strong,-fstack-protector-all>> + -g + $<$:$<$:--coverage -fprofile-arcs -ftest-coverage>> + $<$:-fsanitize=address -fsanitize-address-use-after-scope -fsanitize=leak> + # 在 Clang 编译器场景下 使能 -fsanitize=undefined 会默认开启基本所有的 UBSAN 检查项, 只有以下检查项不会开启 + # float-divide-by-zero, unsigned-integer-overflow, implicit-conversion, local-bounds 及 nullability-* 类检查. + # 故在 Clang 编译器使能 UBSAN 场景下, 需开启 -fsanitize=undefined 使能时仍未开启的对应检查项 + # 在 GNU 编译器场景下, 官方文档并未对使能 -fsanitize=undefined 时开启的默认检查项范围进行说明, 故手工开启常用基本检查项, 避免能力遗漏 + $<$:-fsanitize=undefined -fsanitize=float-divide-by-zero -fno-sanitize=alignment> + # $<$:$<$:-fsanitize=unsigned-integer-overflow>> # GNU 不支持这些检查项 + # $<$:$<$:$<$:-fsanitize=implicit-conversion>>> # GNU 不支持这些检查项, Clang高版本才支持这些检查项 + $<$:$<$:-fsanitize=shift -fsanitize=integer-divide-by-zero -fsanitize=signed-integer-overflow -fsanitize=float-divide-by-zero -fsanitize=float-cast-overflow -fsanitize=bool -fsanitize=enum -fsanitize=vptr>> + $<$:-fno-omit-frame-pointer -fsanitize-recover=all> + -Wall -fno-common -fno-strict-aliasing + -Wundef -Wcast-qual -Wpointer-arith -Wdate-time + -Wfloat-equal -Wformat=2 -Wshadow + -Wsign-compare -Wunused-macros -Wvla -Wdisabled-optimization -Wempty-body -Wignored-qualifiers + $<$:-Wimplicit-fallthrough=3> -Wtype-limits -Wshift-negative-value -Wswitch-default + -Wframe-larger-than=67108864 # 67108864=65536 * 1024, 兼容 ASAN场景对栈的额外消耗 + -Woverloaded-virtual + -Wnon-virtual-dtor $<$:-Wshift-overflow=2> -Wshift-count-overflow + -Wwrite-strings -Wmissing-format-attribute -Wformat-nonliteral + -Wdelete-non-virtual-dtor $<$:-Wduplicated-cond> + $<$:-Wtrampolines> + $<$:-Wsized-deallocation> + $<$:-Wlogical-op> + $<$:-Wsuggest-attribute=format> + $<$:-Wnested-externs> + $<$:-Wduplicated-branches> + $<$:-Wformat-signedness> + $<$:-Wreturn-local-addr> -Wextra + -Wredundant-decls -Wfloat-conversion + $<$:-Wno-tautological-unsigned-enum-zero-compare> + $<$:-fprofile-arcs -ftest-coverage> +) +target_include_directories(intf_pub_utest + INTERFACE + ${ASCEND_CANN_PACKAGE_PATH}/include + ${ASCEND_CANN_PACKAGE_PATH}/include/external + ${ASCEND_CANN_PACKAGE_PATH}/include/experiment + ${ASCEND_CANN_PACKAGE_PATH}/include/experiment/metadef + ${ASCEND_CANN_PACKAGE_PATH}/include/experiment/runtime + ${ASCEND_CANN_PACKAGE_PATH}/include/experiment/msprof +) +if (TESTS_UT_OPS_TEST_CI_PR) + target_compile_definitions(intf_pub_utest + INTERFACE + TESTS_UT_OPS_TEST_CI_PR + ) +endif () +target_link_directories(intf_pub_utest + INTERFACE + ${ASCEND_CANN_PACKAGE_PATH}/lib64 + ${ASCEND_CANN_PACKAGE_PATH}/runtime/lib64 + ${ASCEND_CANN_PACKAGE_PATH}/runtime/lib64/stub +) +target_link_libraries(intf_pub_utest + INTERFACE + $<$:$<$:gcov>> + pthread +) +target_link_options(intf_pub_utest + INTERFACE + $<$:$<$:-fprofile-arcs -ftest-coverage>> + $<$:-fsanitize=address> + $<$:-fsanitize=undefined> +) diff --git a/csrc/ascend/cmake/makeself_built_in.cmake b/csrc/ascend/cmake/makeself_built_in.cmake new file mode 100644 index 000000000..52c5ffdad --- /dev/null +++ b/csrc/ascend/cmake/makeself_built_in.cmake @@ -0,0 +1,108 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +# makeself.cmake - 自定义 makeself 打包脚本 + +# 设置 makeself 路径 +set(MAKESELF_EXE ${CPACK_MAKESELF_PATH}/makeself.sh) +set(MAKESELF_HEADER_EXE ${CPACK_MAKESELF_PATH}/makeself-header.sh) +if(NOT MAKESELF_EXE) + message(FATAL_ERROR "makeself not found!") +endif() + +# 创建临时安装目录 +set(STAGING_DIR "${CPACK_CMAKE_BINARY_DIR}/_CPack_Packages/makeself_staging") +# 先删除再创建 +file(REMOVE_RECURSE "${STAGING_DIR}") +file(MAKE_DIRECTORY "${STAGING_DIR}") + +# 执行安装到临时目录 +execute_process( + COMMAND "${CMAKE_COMMAND}" --install "${CPACK_CMAKE_BINARY_DIR}" --prefix "${STAGING_DIR}" + RESULT_VARIABLE INSTALL_RESULT +) + +if(NOT INSTALL_RESULT EQUAL 0) + message(FATAL_ERROR "Installation to staging directory failed: ${INSTALL_RESULT}") +endif() + +# 生成安装配置文件 +set(CSV_OUTPUT ${CPACK_CMAKE_BINARY_DIR}/filelist.csv) +execute_process( + COMMAND python3 ${CPACK_CMAKE_SOURCE_DIR}/scripts/package/package.py --pkg_name ops_transformer --chip_name ${CPACK_SOC} --os_arch linux-${CPACK_ARCH} --version_dir ${CPACK_VERSION} + WORKING_DIRECTORY ${CPACK_CMAKE_BINARY_DIR} + OUTPUT_VARIABLE result + ERROR_VARIABLE error + RESULT_VARIABLE code + OUTPUT_STRIP_TRAILING_WHITESPACE +) +message(STATUS "package.py result: ${code}") +if (NOT code EQUAL 0) + message(FATAL_ERROR "Filelist generation failed: ${result}") +else () + message(STATUS "Filelist generated successfully: ${result}") + + if (NOT EXISTS ${CSV_OUTPUT}) + message(FATAL_ERROR "Output file not created: ${CSV_OUTPUT}") + endif () +endif () +set(SCENE_OUT_PUT + ${CPACK_CMAKE_BINARY_DIR}/scene.info +) +set(TRANSFORMER_VERSION_OUT_PUT + ${CPACK_CMAKE_BINARY_DIR}/ops_transformer_version.h +) + +configure_file( + ${SCENE_OUT_PUT} + ${STAGING_DIR}/share/info/ops_transformer/ + COPYONLY +) +configure_file( + ${CSV_OUTPUT} + ${STAGING_DIR}/share/info/ops_transformer/script/ + COPYONLY +) +configure_file( + ${TRANSFORMER_VERSION_OUT_PUT} + ${STAGING_DIR}/share/info/ops_transformer/ + COPYONLY +) +# makeself打包 +file(STRINGS ${CPACK_CMAKE_BINARY_DIR}/makeself.txt script_output) +string(REPLACE " " ";" makeself_param_string "${script_output}") +string(REGEX MATCH "cann.*\\.run" package_name "${makeself_param_string}") + +list(LENGTH makeself_param_string LIST_LENGTH) +math(EXPR INSERT_INDEX "${LIST_LENGTH} - 2") +list(INSERT makeself_param_string ${INSERT_INDEX} "${STAGING_DIR}") + +message(STATUS "script output: ${script_output}") +message(STATUS "makeself: ${makeself_param_string}") +message(STATUS "package: ${package_name}") + +execute_process(COMMAND bash ${MAKESELF_EXE} + --header ${MAKESELF_HEADER_EXE} + --help-header share/info/ops_transformer/script/help.info + ${makeself_param_string} share/info/ops_transformer/script/install.sh + WORKING_DIRECTORY ${STAGING_DIR} + RESULT_VARIABLE EXEC_RESULT + ERROR_VARIABLE EXEC_ERROR +) + +if(NOT EXEC_RESULT EQUAL 0) + message(FATAL_ERROR "makeself packaging failed: ${EXEC_ERROR}") +endif() + +execute_process( + COMMAND mkdir -p ${CPACK_PACKAGE_DIRECTORY} + COMMAND mv ${STAGING_DIR}/${package_name} ${CPACK_PACKAGE_DIRECTORY}/ + COMMAND echo "build pkg success: ${CPACK_PACKAGE_DIRECTORY}/${package_name}" + WORKING_DIRECTORY ${STAGING_DIR} +) diff --git a/csrc/ascend/cmake/makeself_custom.cmake b/csrc/ascend/cmake/makeself_custom.cmake new file mode 100644 index 000000000..c694b1950 --- /dev/null +++ b/csrc/ascend/cmake/makeself_custom.cmake @@ -0,0 +1,51 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +# makeself_custom.cmake - 自定义 makeself 打包脚本 + +# 设置 makeself 路径 +set(MAKESELF_EXE ${CPACK_MAKESELF_PATH}/makeself.sh) +set(MAKESELF_HEADER_EXE ${CPACK_MAKESELF_PATH}/makeself-header.sh) +if(NOT MAKESELF_EXE) + message(FATAL_ERROR "makeself not found!") +endif() + +execute_process(COMMAND bash ${MAKESELF_EXE} + --header ${MAKESELF_HEADER_EXE} + --help-header ./help.info --tar-format posix + --gzip --complevel 4 --nomd5 --sha256 + ./ ${CPACK_PACKAGE_FILE_NAME} "version:1.0" ./install.sh + WORKING_DIRECTORY ${CPACK_TEMPORARY_DIRECTORY} + RESULT_VARIABLE EXEC_RESULT + ERROR_VARIABLE EXEC_ERROR +) + +if (NOT "${EXEC_RESULT}x" STREQUAL "0x") + message(FATAL_ERROR "CPack Command error: ${EXEC_RESULT}\n${EXEC_ERROR}") +endif() + +execute_process(COMMAND cp ${CPACK_EXTERNAL_BUILT_PACKAGES} ${CPACK_PACKAGE_DIRECTORY}/ + COMMAND echo "Copy ${CPACK_EXTERNAL_BUILT_PACKAGES} to ${CPACK_PACKAGE_DIRECTORY}/" + WORKING_DIRECTORY ${CPACK_TEMPORARY_DIRECTORY} + ) + +if (NOT "${CPACK_PACKAGE_DIRECTORY}x" STREQUAL "${CPACK_INSTALL_PREFIX}x") + execute_process( + COMMAND ${CMAKE_COMMAND} -E make_directory ${CPACK_INSTALL_PREFIX} + WORKING_DIRECTORY ${CPACK_TEMPORARY_DIRECTORY} + ) + + execute_process( + COMMAND cp ${CPACK_EXTERNAL_BUILT_PACKAGES} ${CPACK_INSTALL_PREFIX}/ + COMMAND echo "Copy ${CPACK_EXTERNAL_BUILT_PACKAGES} to ${CPACK_INSTALL_PREFIX}/" + WORKING_DIRECTORY ${CPACK_TEMPORARY_DIRECTORY} + ) +endif() + + diff --git a/csrc/ascend/cmake/modules/FindOPBASE.cmake b/csrc/ascend/cmake/modules/FindOPBASE.cmake new file mode 100644 index 000000000..3003db5a2 --- /dev/null +++ b/csrc/ascend/cmake/modules/FindOPBASE.cmake @@ -0,0 +1,68 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if(OPBASE_FOUND) + message(STATUS "OpBase has been found") + return() +endif() + +include(FindPackageHandleStandardArgs) + +set(OPBASE_HEAD_SEARCH_PATHS + ${ASCEND_DIR}/${SYSTEM_PREFIX}/pkg_inc + ${TOP_DIR}/ops-base/pkg_inc # compile with ci +) + +set(OPBASE_LIB_SEARCH_PATHS ${ASCEND_DIR}/${SYSTEM_PREFIX}) + +find_path(OPBASE_INC_DIR + NAMES op_common/op_host/util/opbase_export.h + PATHS ${OPBASE_HEAD_SEARCH_PATHS} + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) + +find_library(OPBASE_LIB_DIR + NAME ops_base + PATHS ${OPBASE_LIB_SEARCH_PATHS} + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) + +find_package_handle_standard_args(OPBASE + REQUIRED_VARS OPBASE_INC_DIR) + +get_filename_component(OPBASE_INC_DIR ${OPBASE_INC_DIR} REALPATH) +if(OPBASE_LIB_DIR) + get_filename_component(OPBASE_LIB_DIR ${OPBASE_LIB_DIR} REALPATH) + add_library(opsbase SHARED IMPORTED) + set_target_properties(opsbase PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES ${OPBASE_INC_DIR} + IMPORTED_LOCATION ${OPBASE_LIB_DIR} + ) +else() + if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + message(STATUS "Cannot find library ops_base") + endif() +endif() + +if(OPBASE_FOUND) + if(NOT OPBASE_FIND_QUIETLY) + message(STATUS "Found OPABSE include:${OPBASE_INC_DIR}") + message(STATUS "Found OPABSE lib:${OPBASE_LIB_DIR}") + endif() + set(OPBASE_INC_DIRS + ${OPBASE_INC_DIR} + ${OPBASE_INC_DIR}/op_common + ${OPBASE_INC_DIR}/op_common/op_host + ${OPBASE_INC_DIR}/op_common/atvoss + ) +endif() \ No newline at end of file diff --git a/csrc/ascend/cmake/modules/FindPython.cmake b/csrc/ascend/cmake/modules/FindPython.cmake new file mode 100644 index 000000000..8c21bd419 --- /dev/null +++ b/csrc/ascend/cmake/modules/FindPython.cmake @@ -0,0 +1,60 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +include_guard(GLOBAL) + +find_package(Python3 3.7 COMPONENTS Interpreter) +set(HI_PYTHON ${Python3_EXECUTABLE}) +message(STATUS "HI_PYTHON = ${Python3_EXECUTABLE}") + +macro(get_python) + execute_process(COMMAND ${HI_PYTHON} -c "import sys; print(sys.executable)" + RESULT_VARIABLE result + OUTPUT_VARIABLE _TEMP_PYTHON_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE) + if (result) + set(FATAL_ERROR "Please install python3 software first.") + endif() + + execute_process(COMMAND ${HI_PYTHON} -c "import sysconfig; print(sysconfig.get_config_var('INCLUDEPY'))" + RESULT_VARIABLE result + OUTPUT_VARIABLE HI_PYTHON_INC_TEMP + OUTPUT_STRIP_TRAILING_WHITESPACE) + if (result) + set(HI_PYTHON_INC_TEMP "unknown") + endif() + + execute_process(COMMAND ${HI_PYTHON} -c "import sysconfig; print(sysconfig.get_config_var('LIBDIR'))" + RESULT_VARIABLE result + OUTPUT_VARIABLE HI_PYTHON_LIBDIR + OUTPUT_STRIP_TRAILING_WHITESPACE) + if (result) + set(HI_PYTHON_LIBDIR ".") + endif() + + execute_process(COMMAND ${HI_PYTHON} -c "import sysconfig; print(sysconfig.get_config_var('BLDLIBRARY'))" + RESULT_VARIABLE result + OUTPUT_VARIABLE HI_PYTHON_LIB_TEMP + OUTPUT_STRIP_TRAILING_WHITESPACE) + if (result) + set(HI_PYTHON_LIB_TEMP "") + endif() + if(${HI_PYTHON_LIB_TEMP} MATCHES "^lib.*\\.a$") + set(HI_PYTHON_LIB_TEMP "${HI_PYTHON_LIBDIR}/${HI_PYTHON_LIB_TEMP}") + else() + STRING(REPLACE "-L." "-L${HI_PYTHON_LIBDIR}" HI_PYTHON_LIB_TEMP ${HI_PYTHON_LIB_TEMP}) + endif() +endmacro(get_python) +get_python() +message(STATUS "HI_PYTHON_INC_TEMP = ${HI_PYTHON_INC_TEMP}") +message(STATUS "HI_PYTHON_LIB_TEMP = ${HI_PYTHON_LIB_TEMP}") + +set(HI_PYTHON_INC ${HI_PYTHON_INC_TEMP} CACHE STRING "python include path" FORCE) +set(HI_PYTHON_LIB ${HI_PYTHON_LIB_TEMP} CACHE STRING "python library libpython3.7m.a" FORCE) diff --git a/csrc/ascend/cmake/modules/Findaicpu.cmake b/csrc/ascend/cmake/modules/Findaicpu.cmake new file mode 100644 index 000000000..850c2d608 --- /dev/null +++ b/csrc/ascend/cmake/modules/Findaicpu.cmake @@ -0,0 +1,52 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if(aicpu_FOUND) + message(STATUS "aicpu has been found") + return() +endif() + +include(FindPackageHandleStandardArgs) + +if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + set(AICPU_INC_DIRS + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include/aicpu + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include/experiment + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include/experiment/cce + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include/experiment/msprof + ${ASCEND_DIR}/${SYSTEM_PREFIX}/pkg_inc + ${ASCEND_DIR}/${SYSTEM_PREFIX}/pkg_inc/aicpu + ${ASCEND_DIR}/${SYSTEM_PREFIX}/pkg_inc/aicpu_common/context + ${ASCEND_DIR}/${SYSTEM_PREFIX}/pkg_inc/aicpu_common/context/common + ${ASCEND_DIR}/${SYSTEM_PREFIX}/pkg_inc/aicpu_common/context/cpu_proto + ${ASCEND_DIR}/${SYSTEM_PREFIX}/pkg_inc/aicpu_common/context/utils + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include/aicpu_common/context + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include/aicpu_common/context/common + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include/aicpu_common/context/cpu_proto + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include/experiment/datagw/aicpu/common + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include/aicpu_common/context/utils + ) +else() + set(AICPU_INC_DIRS + ${TOP_DIR}/abl/msprof/inc # codespell:ignore abl + ${TOP_DIR}/ace/comop/inc + ${TOP_DIR}/inc/aicpu/cpu_kernels + ${TOP_DIR}/inc/external/aicpu + ${TOP_DIR}/asl/ops/cann/ops/built-in/aicpu/context/inc + ${TOP_DIR}/asl/ops/cann/ops/built-in/aicpu/impl/utils + ${TOP_DIR}/asl/ops/cann/ops/built-in/aicpu/impl + ${TOP_DIR}/ops-base/pkg_inc/aicpu_common/context/common + ${TOP_DIR}/ops-base/include/aicpu_common/context/common + ${TOP_DIR}/open_source/eigen + ${TOP_DIR}/runtime/pkg_inc/aicpu_sched/common + ) +endif() + +message(STATUS "Using AICPU include dirs: ${AICPU_INC_DIRS}") diff --git a/csrc/ascend/cmake/modules/Findalog.cmake b/csrc/ascend/cmake/modules/Findalog.cmake new file mode 100644 index 000000000..30c95c03c --- /dev/null +++ b/csrc/ascend/cmake/modules/Findalog.cmake @@ -0,0 +1,119 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if (alog_FOUND) + message(STATUS "Package alog has been found.") + return() +endif() + +set(_cmake_targets_defined "") +set(_cmake_targets_not_defined "") +set(_cmake_expected_targets "") +foreach(_cmake_expected_target IN ITEMS slog_a alog_a alog_headers) + list(APPEND _cmake_expected_targets "${_cmake_expected_target}") + if(TARGET "${_cmake_expected_target}") + list(APPEND _cmake_targets_defined "${_cmake_expected_target}") + else() + list(APPEND _cmake_targets_not_defined "${_cmake_expected_target}") + endif() +endforeach() +unset(_cmake_expected_target) + +if(_cmake_targets_defined STREQUAL _cmake_expected_targets) + unset(_cmake_targets_defined) + unset(_cmake_targets_not_defined) + unset(_cmake_expected_targets) + unset(CMAKE_IMPORT_FILE_VERSION) + cmake_policy(POP) + return() +endif() + +if(NOT _cmake_targets_defined STREQUAL "") + string(REPLACE ";" ", " _cmake_targets_defined_text "${_cmake_targets_defined}") + string(REPLACE ";" ", " _cmake_targets_not_defined_text "${_cmake_targets_not_defined}") + message(FATAL_ERROR "Some (but not all) targets in this export set were already defined.\nTargets Defined: ${_cmake_targets_defined_text}\nTargets not yet defined: ${_cmake_targets_not_defined_text}\n") +endif() +unset(_cmake_targets_defined) +unset(_cmake_targets_not_defined) +unset(_cmake_expected_targets) + +set(ALOG_HEAD_SEARCH_PATHS + ${ASCEND_DIR}/pkg_inc +) + +find_path(_INCLUDE_DIR + NAMES base/alog_pub.h + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH + PATHS ${ALOG_HEAD_SEARCH_PATHS}) + +find_library(slog_a_SHARED_LIBRARY + NAMES libascendalog.so + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH) + +find_library(alog_a_SHARED_LIBRARY + NAMES libascendalog.so + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(alog + FOUND_VAR + alog_FOUND + REQUIRED_VARS + _INCLUDE_DIR + slog_a_SHARED_LIBRARY + alog_a_SHARED_LIBRARY +) + +if(alog_FOUND) + set(alog_a_INCLUDE_DIR "${_INCLUDE_DIR}") + include(CMakePrintHelpers) + message(STATUS "Variables in alog module:") + cmake_print_variables(alog_a_INCLUDE_DIR) + cmake_print_variables(slog_a_SHARED_LIBRARY) + cmake_print_variables(alog_a_SHARED_LIBRARY) + + add_library(slog_a SHARED IMPORTED) + set_target_properties(slog_a PROPERTIES + INTERFACE_COMPILE_DEFINITIONS "LOG_CPP;PROCESS_LOG" + INTERFACE_LINK_LIBRARIES "alog_headers" + IMPORTED_LOCATION "${slog_a_SHARED_LIBRARY}" + ) + + add_library(alog_a SHARED IMPORTED) + set_target_properties(alog_a PROPERTIES + INTERFACE_COMPILE_DEFINITIONS "LOG_CPP;PROCESS_LOG" + INTERFACE_LINK_LIBRARIES "alog_headers" + IMPORTED_LOCATION "${alog_a_SHARED_LIBRARY}" + ) + + add_library(alog_headers INTERFACE IMPORTED) + set_target_properties(alog_headers PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${alog_a_INCLUDE_DIR};${alog_a_INCLUDE_DIR}/base" + ) + + include(CMakePrintHelpers) + cmake_print_properties(TARGETS slog_a + PROPERTIES INTERFACE_COMPILE_DEFINITIONS INTERFACE_LINK_LIBRARIES IMPORTED_LOCATION + ) + cmake_print_properties(TARGETS alog_a + PROPERTIES INTERFACE_COMPILE_DEFINITIONS INTERFACE_LINK_LIBRARIES IMPORTED_LOCATION + ) + cmake_print_properties(TARGETS alog_headers + PROPERTIES INTERFACE_INCLUDE_DIRECTORIES + ) +endif() + +# Cleanup temporary variables. +set(_INCLUDE_DIR) diff --git a/csrc/ascend/cmake/modules/Finddlog.cmake b/csrc/ascend/cmake/modules/Finddlog.cmake new file mode 100644 index 000000000..79a2ae98e --- /dev/null +++ b/csrc/ascend/cmake/modules/Finddlog.cmake @@ -0,0 +1,105 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if (dlog_FOUND) + message(STATUS "Package dlog has been found.") + return() +endif() + +set(_cmake_targets_defined "") +set(_cmake_targets_not_defined "") +set(_cmake_expected_targets "") +foreach(_cmake_expected_target IN ITEMS dlog_a dlog_headers) + list(APPEND _cmake_expected_targets "${_cmake_expected_target}") + if(TARGET "${_cmake_expected_target}") + list(APPEND _cmake_targets_defined "${_cmake_expected_target}") + else() + list(APPEND _cmake_targets_not_defined "${_cmake_expected_target}") + endif() +endforeach() +unset(_cmake_expected_target) + +if(_cmake_targets_defined STREQUAL _cmake_expected_targets) + unset(_cmake_targets_defined) + unset(_cmake_targets_not_defined) + unset(_cmake_expected_targets) + unset(CMAKE_IMPORT_FILE_VERSION) + cmake_policy(POP) + return() +endif() + +if(NOT _cmake_targets_defined STREQUAL "") + string(REPLACE ";" ", " _cmake_targets_defined_text "${_cmake_targets_defined}") + string(REPLACE ";" ", " _cmake_targets_not_defined_text "${_cmake_targets_not_defined}") + message(FATAL_ERROR "Some (but not all) targets in this export set were already defined.\nTargets Defined: ${_cmake_targets_defined_text}\nTargets not yet defined: ${_cmake_targets_not_defined_text}\n") +endif() +unset(_cmake_targets_defined) +unset(_cmake_targets_not_defined) +unset(_cmake_expected_targets) + +set(DLOG_HEAD_SEARCH_PATHS + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include/toolchain + ${ASCEND_DIR}/pkg_inc/base # new slog directory structure + ${TOP_DIR}/abl/slog/inc/toolchain # compile with CI # codespell:ignore abl +) + +find_path(dlog_TRANSFORMER_INCLUDE_DIR + NAMES dlog_pub.h + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH + PATHS ${DLOG_HEAD_SEARCH_PATHS}) + +find_library(dlog_SHARED_LIBRARY + NAMES libascendalog.so + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(dlog + FOUND_VAR + dlog_FOUND + REQUIRED_VARS + dlog_TRANSFORMER_INCLUDE_DIR +) + +if(dlog_FOUND) + include(CMakePrintHelpers) + message(STATUS "Transformer Variables in dlog module:") + cmake_print_variables(dlog_TRANSFORMER_INCLUDE_DIR) + cmake_print_variables(dlog_SHARED_LIBRARY) + + if(dlog_SHARED_LIBRARY) + add_library(dlog SHARED IMPORTED) + set_target_properties(dlog PROPERTIES + INTERFACE_COMPILE_DEFINITIONS "LOG_CPP;PROCESS_LOG" + INTERFACE_LINK_LIBRARIES "dlog_headers" + IMPORTED_LOCATION "${dlog_SHARED_LIBRARY}" + ) + endif() + + set(TRANSFORMER_INTERFACE_INCLUDE "${dlog_TRANSFORMER_INCLUDE_DIR}") + string(FIND "${dlog_TRANSFORMER_INCLUDE_DIR}" "toolchain" IDX) + if(NOT IDX EQUAL -1) + set(TRANSFORMER_INTERFACE_INCLUDE "${TRANSFORMER_INTERFACE_INCLUDE};${dlog_TRANSFORMER_INCLUDE_DIR}/..") + endif() + add_library(dlog_headers INTERFACE IMPORTED) + set_target_properties(dlog_headers PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${TRANSFORMER_INTERFACE_INCLUDE}" + ) + + include(CMakePrintHelpers) + cmake_print_properties(TARGETS dlog + PROPERTIES INTERFACE_COMPILE_DEFINITIONS INTERFACE_LINK_LIBRARIES IMPORTED_LOCATION + ) + cmake_print_properties(TARGETS dlog_headers + PROPERTIES INTERFACE_INCLUDE_DIRECTORIES + ) +endif() diff --git a/csrc/ascend/cmake/modules/Findjson.cmake b/csrc/ascend/cmake/modules/Findjson.cmake new file mode 100644 index 000000000..fc4bfbc83 --- /dev/null +++ b/csrc/ascend/cmake/modules/Findjson.cmake @@ -0,0 +1,36 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if (json_FOUND) + message(STATUS "Package json has been found.") + return() +endif() + +find_path(JSON_INCLUDE + NAMES nlohmann/json.hpp + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(json + FOUND_VAR + json_FOUND + REQUIRED_VARS + JSON_INCLUDE + ) + +if(json_FOUND) + set(JSON_INCLUDE_DIR ${JSON_INCLUDE_DIR}) + + add_library(json INTERFACE IMPORTED) + set_target_properties(json PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${JSON_INCLUDE_DIR}") + target_compile_definitions(json INTERFACE nlohmann=ascend_nlohmann) +endif() diff --git a/csrc/ascend/cmake/modules/Findmetadef.cmake b/csrc/ascend/cmake/modules/Findmetadef.cmake new file mode 100644 index 000000000..f1e2b9c7d --- /dev/null +++ b/csrc/ascend/cmake/modules/Findmetadef.cmake @@ -0,0 +1,226 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if(metadef_FOUND) + message(STATUS "metadef has been found") + return() +endif() + +include(FindPackageHandleStandardArgs) + +set(METADEF_HEAD_SEARCH_PATHS + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include + ${TOP_DIR}/metadef/inc/external/ # compile with ci +) + +set(METADEF_LIB_SEARCH_PATHS + ${ASCEND_DIR}/${SYSTEM_PREFIX} +) + +find_path(METADEF_INC_DIR + NAMES register/register.h + PATHS ${METADEF_HEAD_SEARCH_PATHS} + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) + +find_library(REGISTER_LIB_DIR + NAME register + PATHS ${METADEF_LIB_SEARCH_PATHS} + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) + +if(REGISTER_LIB_DIR) + get_filename_component(REGISTER_LIB_DIR ${REGISTER_LIB_DIR} REALPATH) + add_library(register SHARED IMPORTED) + set_target_properties(register PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES ${METADEF_INC_DIR} + IMPORTED_LOCATION ${REGISTER_LIB_DIR} + ) + message(STATUS "Found register library:${REGISTER_LIB_DIR}") +else() + if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + message(STATUS "Cannot find library register") + endif() +endif() + +find_library(OPP_REGISTER_LIB_DIR + NAME opp_registry + PATHS ${METADEF_LIB_SEARCH_PATHS} + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) + +if(OPP_REGISTER_LIB_DIR) + get_filename_component(OPP_REGISTER_LIB_DIR ${OPP_REGISTER_LIB_DIR} REALPATH) + add_library(opp_registry SHARED IMPORTED) + set_target_properties(opp_registry PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES ${METADEF_INC_DIR} + IMPORTED_LOCATION ${OPP_REGISTER_LIB_DIR} + ) + message(STATUS "Found opp_registry library:${OPP_REGISTER_LIB_DIR}") +else() + if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + message(STATUS "Cannot find library opp_registry") + endif() +endif() + +find_library(EXEGRAPH_LIB_DIR + NAME exe_graph + PATHS ${METADEF_LIB_SEARCH_PATHS} + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) + +if(EXEGRAPH_LIB_DIR) + get_filename_component(EXEGRAPH_LIB_DIR ${EXEGRAPH_LIB_DIR} REALPATH) + add_library(exe_graph SHARED IMPORTED) + set_target_properties(exe_graph PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES ${METADEF_INC_DIR}/exe_graph + IMPORTED_LOCATION ${EXEGRAPH_LIB_DIR} + ) + message(STATUS "Found exe_graph library:${EXEGRAPH_LIB_DIR}") +else() + if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + message(STATUS "Cannot find library exe_graph") + endif() +endif() + +find_library(GRAPH_BASE_LIB_DIR + NAME graph_base + PATHS ${METADEF_LIB_SEARCH_PATHS} + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH + ) + +if(GRAPH_BASE_LIB_DIR) + get_filename_component(GRAPH_BASE_LIB_DIR ${GRAPH_BASE_LIB_DIR} REALPATH) + add_library(graph_base SHARED IMPORTED) + set_target_properties(graph_base PROPERTIES + IMPORTED_LOCATION ${GRAPH_BASE_LIB_DIR} + ) + message(STATUS "Found graph_base library:${GRAPH_BASE_LIB_DIR}") +else() + if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + message(STATUS "Cannot find library graph_base") + endif() +endif() + +find_library(ERROR_MANAGER_LIB_DIR + NAME error_manager + PATHS ${METADEF_LIB_SEARCH_PATHS} + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH + ) + +if(ERROR_MANAGER_LIB_DIR) + get_filename_component(ERROR_MANAGER_LIB_DIR ${ERROR_MANAGER_LIB_DIR} REALPATH) + add_library(error_manager SHARED IMPORTED) + set_target_properties(error_manager PROPERTIES + IMPORTED_LOCATION ${ERROR_MANAGER_LIB_DIR} + ) + message(STATUS "Found error_manager library:${ERROR_MANAGER_LIB_DIR}") +else() + if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + message(STATUS "Cannot find library error_manager") + endif() +endif() + +find_library(METADEF_LIB_DIR + NAME metadef + PATHS ${METADEF_LIB_SEARCH_PATHS} + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH + ) + +if(METADEF_LIB_DIR) + get_filename_component(METADEF_LIB_DIR ${METADEF_LIB_DIR} REALPATH) + add_library(metadef SHARED IMPORTED) + set_target_properties(metadef PROPERTIES + IMPORTED_LOCATION ${METADEF_LIB_DIR} + ) + message(STATUS "Found metadef library:${METADEF_LIB_DIR}") +else() + if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + message(STATUS "Cannot find library metadef") + endif() +endif() + +find_library(REGISTER_STATIC_LIB_DIR + NAME librt2_registry.a + PATHS ${METADEF_LIB_SEARCH_PATHS} + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) + +if(REGISTER_STATIC_LIB_DIR) + get_filename_component(REGISTER_STATIC_LIB_DIR ${REGISTER_STATIC_LIB_DIR} REALPATH) + add_library(rt2_registry_static STATIC IMPORTED) + set_target_properties(rt2_registry_static PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES ${METADEF_INC_DIR} + IMPORTED_LOCATION ${REGISTER_STATIC_LIB_DIR} + ) + message(STATUS "Found rt2_registry library:${REGISTER_STATIC_LIB_DIR}") +else() + if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + message(STATUS "Cannot find library rt2_registry") + endif() +endif() + +find_package_handle_standard_args(metadef + REQUIRED_VARS METADEF_INC_DIR) + +get_filename_component(METADEF_INC_DIR ${METADEF_INC_DIR} REALPATH) +if(metadef_FOUND) + set(METADEF_INCLUDE_DIRS + ${METADEF_INC_DIR}/ + ${METADEF_INC_DIR}/exe_graph + ) + + if(NOT BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + set(METADEF_INCLUDE_DIRS ${METADEF_INC_DIR}/../ ${METADEF_INCLUDE_DIRS}) + endif() + message(STATUS "Found source metadef include dir: ${METADEF_INCLUDE_DIRS}") +endif() + + + +set(METADEF_LIB_SEARCH_PATHS + ${ASCEND_DIR}/${SYSTEM_PREFIX} +) + +find_library(GRAPH_LIB_DIR + NAME graph + PATHS ${METADEF_LIB_SEARCH_PATHS} + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) + +if(GRAPH_LIB_DIR) + get_filename_component(GRAPH_LIB_DIR ${GRAPH_LIB_DIR} REALPATH) + add_library(graph SHARED IMPORTED) + set_target_properties(graph PROPERTIES + IMPORTED_LOCATION ${GRAPH_LIB_DIR} + ) + message(STATUS "Found graph library:${GRAPH_LIB_DIR}") +else() + if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + message(STATUS "Cannot find library graph") + endif() +endif() \ No newline at end of file diff --git a/csrc/ascend/cmake/modules/Findnnopbase.cmake b/csrc/ascend/cmake/modules/Findnnopbase.cmake new file mode 100644 index 000000000..d66dbb251 --- /dev/null +++ b/csrc/ascend/cmake/modules/Findnnopbase.cmake @@ -0,0 +1,84 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if(nnopbase_FOUND) + message(STATUS "nnopbase has been found") + return() +endif() + +set(nnopbase_FOUND ON) +include(FindPackageHandleStandardArgs) + +set(NNOPBASE_ACLNN_HEAD_SEARCH_PATHS + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include + ${TOP_DIR}/ace/npuruntime/inc/external/ # compile with ci +) + +set(NNOPBASE_OPDEV_HEAD_SEARCH_PATHS + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include/aclnn + ${TOP_DIR}/ace/npuruntime/inc/nnopbase/ # compile with ci +) + +set(NNOPBASE_LIB_SEARCH_PATHS + ${ASCEND_DIR}/${SYSTEM_PREFIX} +) + +find_path(NNOPBASE_ACLNN_INC_DIR + NAMES aclnn/aclnn_base.h + PATHS ${NNOPBASE_ACLNN_HEAD_SEARCH_PATHS} + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) +if(NOT NNOPBASE_ACLNN_INC_DIR) + set(nnopbase_FOUND OFF) +endif() + +find_path(NNOPBASE_OPDEV_INC_DIR + NAMES opdev/op_errno.h + PATHS ${NNOPBASE_OPDEV_HEAD_SEARCH_PATHS} + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) +if(NOT NNOPBASE_OPDEV_INC_DIR) + set(nnopbase_FOUND OFF) +endif() + +find_library(NNOPBASE_LIB_DIR + NAME nnopbase + PATHS ${NNOPBASE_LIB_SEARCH_PATHS} + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) + +get_filename_component(NNOPBASE_ACLNN_INC_DIR ${NNOPBASE_ACLNN_INC_DIR} REALPATH) +get_filename_component(NNOPBASE_OPDEV_INC_DIR ${NNOPBASE_OPDEV_INC_DIR} REALPATH) + +if(NNOPBASE_LIB_DIR) + get_filename_component(NNOPBASE_LIB_DIR ${NNOPBASE_LIB_DIR} REALPATH) + add_library(nnopbase SHARED IMPORTED) + set_target_properties(nnopbase PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES ${NNOPBASE_OPDEV_INC_DIR} + IMPORTED_LOCATION ${NNOPBASE_LIB_DIR} + ) +else() + if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + message(STATUS "Cannot find library nnopbase") + endif() +endif() + +if(nnopbase_FOUND) + set(NNOPBASE_INCLUDE_DIRS + ${NNOPBASE_ACLNN_INC_DIR} + ${NNOPBASE_OPDEV_INC_DIR} + ) + message(STATUS "Found aclnn include dir: ${NNOPBASE_ACLNN_INC_DIR}") + message(STATUS "Found opdev include dir: ${NNOPBASE_OPDEV_INC_DIR}") +endif() \ No newline at end of file diff --git a/csrc/ascend/cmake/modules/Findoptiling.cmake b/csrc/ascend/cmake/modules/Findoptiling.cmake new file mode 100644 index 000000000..0cf3f69b7 --- /dev/null +++ b/csrc/ascend/cmake/modules/Findoptiling.cmake @@ -0,0 +1,26 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +unset(OPTILING_SO_LIB_DIR) + +file(GLOB_RECURSE LIB_OP_TILING_SO + ${ASCEND_DIR}/../liboptiling.so +) + +list(FILTER LIB_OP_TILING_SO INCLUDE REGEX "lib/linux/${CMAKE_SYSTEM_PROCESSOR}") + +if(LIB_OP_TILING_SO) + list(GET LIB_OP_TILING_SO 0 LIB_OP_TILING_SO_PATH) + get_filename_component(LIB_OP_TILING_SO_PATH ${LIB_OP_TILING_SO_PATH} DIRECTORY) + message(STATUS "Found optiling so lib:${LIB_OP_TILING_SO_PATH}") +else() + message(STATUS "Cannot find library optiling so") +endif() + diff --git a/csrc/ascend/cmake/modules/Findplatform.cmake b/csrc/ascend/cmake/modules/Findplatform.cmake new file mode 100644 index 000000000..32474d5f7 --- /dev/null +++ b/csrc/ascend/cmake/modules/Findplatform.cmake @@ -0,0 +1,64 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if(platform_FOUND) + message(STATUS "platform has been found") + return() +endif() + +include(FindPackageHandleStandardArgs) + +set(PLATFORM_HEAD_SEARCH_PATHS + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include + ${TOP_DIR}/metadef/inc/external # compile with ci + ${TOP_DIR}/ace/npuruntime/runtime/platform/inc # compile with ci + ${TOP_DIR}/runtime/pkg_inc # compile with ci +) + +set(PLATFORM_LIB_SEARCH_PATHS ${ASCEND_DIR}/${SYSTEM_PREFIX}) + +find_path(PLATFORM_INC_DIR + NAMES platform/platform_info.h + PATHS ${PLATFORM_HEAD_SEARCH_PATHS} + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) + +find_library(PLATFORM_LIB_DIR + NAME platform + PATHS ${PLATFORM_LIB_SEARCH_PATHS} + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) + +find_package_handle_standard_args(platform + REQUIRED_VARS PLATFORM_INC_DIR) + +get_filename_component(PLATFORM_INC_DIR ${PLATFORM_INC_DIR} REALPATH) +if(PLATFORM_LIB_DIR) + get_filename_component(PLATFORM_LIB_DIR ${PLATFORM_LIB_DIR} REALPATH) + add_library(platform SHARED IMPORTED) + set_target_properties(platform PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES ${PLATFORM_INC_DIR} + IMPORTED_LOCATION ${PLATFORM_LIB_DIR} + ) +else() + if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + message(STATUS "Cannot find library platform") + endif() +endif() + +if(platform_FOUND) + if(NOT platform_FIND_QUIETLY) + message(STATUS "Found platform include:${PLATFORM_INC_DIR}") + endif() + set(PLATFORM_INC_DIRS ${PLATFORM_INC_DIR}) +endif() \ No newline at end of file diff --git a/csrc/ascend/cmake/modules/Findruntime.cmake b/csrc/ascend/cmake/modules/Findruntime.cmake new file mode 100644 index 000000000..ee1437b25 --- /dev/null +++ b/csrc/ascend/cmake/modules/Findruntime.cmake @@ -0,0 +1,58 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +include(FindPackageHandleStandardArgs) +set(runtime_FOUND ON) +#search acl.h +set(ACL_HEAD_SEARCH_PATHS + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include + ${TOP_DIR}/ace/npuruntime/acl/inc/external # compile with ci +) +find_path(ACL_INC_DIR + NAMES acl/acl.h + PATHS ${ACL_HEAD_SEARCH_PATHS} + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) +if(NOT ACL_INC_DIR) + set(runtime_FOUND OFF) + message(FATAL_ERROR "no source acl include dir found") +endif() +get_filename_component(ACL_INC_DIR ${ACL_INC_DIR} REALPATH) +message(STATUS "Found source acl include dir: ${ACL_INC_DIR}") + +#search rt.h +set(RUNTIME_SEARCH_PATH + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include/experiment/runtime + ${TOP_DIR}/ace/npuruntime/inc # compile with ci + ${ASCEND_DIR}/${SYSTEM_PREFIX}/pkg_inc/runtime +) +find_path(RUNTIME_INC_DIR + NAMES runtime/rt.h rt_external.h + PATHS ${RUNTIME_SEARCH_PATH} + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) +if(NOT RUNTIME_INC_DIR) + set(runtime_FOUND OFF) + message(FATAL_ERROR "no source runtime include dir found") +endif() +get_filename_component(RUNTIME_INC_DIR ${RUNTIME_INC_DIR} REALPATH) + +if(runtime_FOUND) + if(NOT runtime_FIND_QUIETLY) + message(STATUS "Found source npuruntime include dir: ${RUNTIME_INC_DIR}") + endif() + set(NPURUNTIME_INCLUDE_DIRS + ${ACL_INC_DIR} + ${RUNTIME_INC_DIR} + ${RUNTIME_INC_DIR}/runtime + ) +endif() \ No newline at end of file diff --git a/csrc/ascend/cmake/modules/Findsecurec.cmake b/csrc/ascend/cmake/modules/Findsecurec.cmake new file mode 100644 index 000000000..783e14224 --- /dev/null +++ b/csrc/ascend/cmake/modules/Findsecurec.cmake @@ -0,0 +1,89 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if (securec_FOUND) + message(STATUS "Package securec has been found.") + return() +endif() + +set(C_SEC_HEAD_SEARCH_PATHS + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include + ${TOP_DIR}/abl/libc_sec/include # compile with CI # codespell:ignore abl +) + +set(C_SEC_LIB_SEARCH_PATHS ${ASCEND_DIR}/${SYSTEM_PREFIX}) + +find_path(C_SEC_INCLUDE NAMES securec.h PATHS ${C_SEC_HEAD_SEARCH_PATHS}) +if(WIN32) + find_library(C_SEC_STATIC_LIBRARY NAMES libc_sec_static.lib) + find_file(C_SEC_SHARED_DLL NAMES libc_sec.dll PATH_SUFFIXES lib) + find_library(C_SEC_IMPLIB NAMES libc_sec.lib) +else() + find_library(C_SEC_STATIC_LIBRARY NAMES libc_sec.a PATHS ${C_SEC_LIB_SEARCH_PATHS} PATH_SUFFIXES lib64) + find_library(C_SEC_SHARED_LIBRARY NAMES libc_sec.so PATHS ${C_SEC_LIB_SEARCH_PATHS} PATH_SUFFIXES lib64) +endif() + +include(FindPackageHandleStandardArgs) +if(WIN32) + find_package_handle_standard_args(securec + FOUND_VAR + securec_FOUND + REQUIRED_VARS + C_SEC_INCLUDE + C_SEC_STATIC_LIBRARY + C_SEC_SHARED_DLL + C_SEC_IMPLIB + ) +elseif("${ENABLE_SECUREC_SHARED}" STREQUAL "OFF") + find_package_handle_standard_args(securec + FOUND_VAR + securec_FOUND + REQUIRED_VARS + C_SEC_INCLUDE + C_SEC_STATIC_LIBRARY + ) +else() + find_package_handle_standard_args(securec + FOUND_VAR + securec_FOUND + REQUIRED_VARS + C_SEC_INCLUDE + C_SEC_STATIC_LIBRARY + C_SEC_SHARED_LIBRARY + ) +endif() + +if(securec_FOUND) + set(C_SEC_INCLUDE_DIR ${C_SEC_INCLUDE}) + get_filename_component(C_SEC_LIBRARY_DIR ${C_SEC_STATIC_LIBRARY} DIRECTORY) + + add_library(c_sec_headers INTERFACE IMPORTED) + target_include_directories(c_sec_headers INTERFACE ${C_SEC_INCLUDE}) + + add_library(c_sec_static STATIC IMPORTED) + set_target_properties(c_sec_static PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${C_SEC_INCLUDE}" + IMPORTED_LOCATION "${C_SEC_STATIC_LIBRARY}" + ) + + add_library(c_sec SHARED IMPORTED) + if(WIN32) + set_target_properties(c_sec PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${C_SEC_INCLUDE}" + IMPORTED_IMPLIB "${C_SEC_IMPLIB}" + IMPORTED_LOCATION "${C_SEC_SHARED_DLL}" + ) + else() + set_target_properties(c_sec PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${C_SEC_INCLUDE}" + IMPORTED_LOCATION "${C_SEC_SHARED_LIBRARY}" + ) + endif() +endif() diff --git a/csrc/ascend/cmake/modules/Findtilingapi.cmake b/csrc/ascend/cmake/modules/Findtilingapi.cmake new file mode 100644 index 000000000..0a546a797 --- /dev/null +++ b/csrc/ascend/cmake/modules/Findtilingapi.cmake @@ -0,0 +1,71 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if(tilingapi_FOUND) + message(STATUS "tilingapi has been found") + return() +endif() + +include(FindPackageHandleStandardArgs) + +set(TILINGAPI_HEAD_SEARCH_PATHS + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include +) + +set(TILINGAPI_LIB_SEARCH_PATHS + ${ASCEND_DIR}/${SYSTEM_PREFIX} +) + +find_path(TILINGAPI_INC_DIR + NAMES tiling/tiling_api.h + PATHS ${TILINGAPI_SEARCH_PATHS} + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) + +find_library(TILINGAPI_LIB_DIR + NAME libtiling_api.a + PATHS ${TILINGAPI_LIB_SEARCH_PATHS} + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH +) + +if(TILINGAPI_INC_DIR) + get_filename_component(TILINGAPI_INC_DIR ${TILINGAPI_INC_DIR} REALPATH) + message(STATUS "Found tilingapi include:${TILINGAPI_INC_DIR}") +endif() + +get_filename_component(TILINGAPI_INC_PREFIX ${TILINGAPI_INC_DIR} DIRECTORY) +set(TILINGAPI_INC_DIRS + ${TILINGAPI_INC_PREFIX}/include + ${TILINGAPI_INC_PREFIX}/ascendc + ${TILINGAPI_INC_PREFIX}/ascendc/include + ${TILINGAPI_INC_PREFIX}/ascendc/include/highlevel_api + ${TILINGAPI_INC_PREFIX}/ascendc/include/highlevel_api/lib + ${TILINGAPI_INC_PREFIX}/ascendc/include/highlevel_api/impl + ${TILINGAPI_INC_PREFIX}/ascendc/include/highlevel_api/tiling +) +message(STATUS "TILINGAPI_INC_DIRS: ${TILINGAPI_INC_DIRS}") + +if(TILINGAPI_LIB_DIR) + get_filename_component(TILINGAPI_LIB_DIR ${TILINGAPI_LIB_DIR} REALPATH) + message(STATUS "Found tilingapi lib:${TILINGAPI_LIB_DIR}") + if(NOT TARGET tiling_api) + add_library(tiling_api STATIC IMPORTED) + set_target_properties(tiling_api PROPERTIES + IMPORTED_LOCATION ${TILINGAPI_LIB_DIR} + ) + endif() +else() + if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + message(STATUS "Cannot find library tiling_api") + endif() +endif() \ No newline at end of file diff --git a/csrc/ascend/cmake/modules/Findunified_dlog.cmake b/csrc/ascend/cmake/modules/Findunified_dlog.cmake new file mode 100644 index 000000000..d1b840706 --- /dev/null +++ b/csrc/ascend/cmake/modules/Findunified_dlog.cmake @@ -0,0 +1,96 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if (unified_dlog_FOUND) + message(STATUS "Package unified_dlog has been found.") + return() +endif() + +set(_cmake_targets_defined "") +set(_cmake_targets_not_defined "") +set(_cmake_expected_targets "") +foreach(_cmake_expected_target IN ITEMS unified_dlog unified_dlog_headers) + list(APPEND _cmake_expected_targets "${_cmake_expected_target}") + if(TARGET "${_cmake_expected_target}") + list(APPEND _cmake_targets_defined "${_cmake_expected_target}") + else() + list(APPEND _cmake_targets_not_defined "${_cmake_expected_target}") + endif() +endforeach() +unset(_cmake_expected_target) + +if(_cmake_targets_defined STREQUAL _cmake_expected_targets) + unset(_cmake_targets_defined) + unset(_cmake_targets_not_defined) + unset(_cmake_expected_targets) + unset(CMAKE_IMPORT_FILE_VERSION) + cmake_policy(POP) + return() +endif() + +if(NOT _cmake_targets_defined STREQUAL "") + string(REPLACE ";" ", " _cmake_targets_defined_text "${_cmake_targets_defined}") + string(REPLACE ";" ", " _cmake_targets_not_defined_text "${_cmake_targets_not_defined}") + message(FATAL_ERROR "Some (but not all) targets in this export set were already defined.\nTargets Defined: ${_cmake_targets_defined_text}\nTargets not yet defined: ${_cmake_targets_not_defined_text}\n") +endif() +unset(_cmake_targets_defined) +unset(_cmake_targets_not_defined) +unset(_cmake_expected_targets) + +find_path(_INCLUDE_ROOT_DIR + NAMES toolchain/dlog_pub.h + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH) + +find_library(unified_dlog_SHARED_LIBRARY + NAMES libunified_dlog.so + PATH_SUFFIXES lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(unified_dlog + FOUND_VAR + unified_dlog_FOUND + REQUIRED_VARS + _INCLUDE_ROOT_DIR + unified_dlog_SHARED_LIBRARY +) + +if(unified_dlog_FOUND) + set(unified_dlog_INCLUDE_DIR "${_INCLUDE_ROOT_DIR}") + include(CMakePrintHelpers) + message(STATUS "Variables in unified_dlog module:") + cmake_print_variables(unified_dlog_INCLUDE_DIR) + cmake_print_variables(unified_dlog_SHARED_LIBRARY) + + add_library(unified_dlog SHARED IMPORTED) + set_target_properties(unified_dlog PROPERTIES + INTERFACE_COMPILE_DEFINITIONS "LOG_CPP;PROCESS_LOG" + INTERFACE_LINK_LIBRARIES "unified_dlog_headers" + IMPORTED_LOCATION "${unified_dlog_SHARED_LIBRARY}" + ) + + add_library(unified_dlog_headers INTERFACE IMPORTED) + set_target_properties(unified_dlog_headers PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${unified_dlog_INCLUDE_DIR};${unified_dlog_INCLUDE_DIR}/toolchain" + ) + + include(CMakePrintHelpers) + cmake_print_properties(TARGETS unified_dlog + PROPERTIES INTERFACE_COMPILE_DEFINITIONS INTERFACE_LINK_LIBRARIES IMPORTED_LOCATION + ) + cmake_print_properties(TARGETS unified_dlog_headers + PROPERTIES INTERFACE_INCLUDE_DIRECTORIES + ) +endif() + +# Cleanup temporary variables. +set(_INCLUDE_ROOT_DIR) diff --git a/csrc/ascend/cmake/obj_func.cmake b/csrc/ascend/cmake/obj_func.cmake new file mode 100644 index 000000000..e77610d7d --- /dev/null +++ b/csrc/ascend/cmake/obj_func.cmake @@ -0,0 +1,840 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +# usage: add_modules_sources(DIR OPTYPE ACLNNTYPE) +# ACLNNTYPE 支持类型aclnn/aclnn_inner/aclnn_exclude +# OPTYPE 和 ACLNNTYPE 需一一对应 + +# 用于custom自定算子包host侧obj生成 +macro(add_modules_sources) + set(oneValueArgs OP_API_INDEPENDENT OP_API_DIR) + set(multiValueArgs OPTYPE ACLNNTYPE) + + cmake_parse_arguments(MODULE "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + set(SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) + + # 该段代码作用为兼容op_api新旧目录结构(旧： 嵌套于op_host下； 新： 与op_host同级) + if (NOT DEFINED MODULE_OP_API_INDEPENDENT) + set(MODULE_OP_API_INDEPENDENT OFF) + endif() + if(MODULE_OP_API_INDEPENDENT) + # 新结构：op_api与op_host同级，需要指定有效路径 + if (NOT DEFINED MODULE_OP_API_DIR OR NOT EXISTS "${MODULE_OP_API_DIR}") + message(FATAL_ERROR "OP_API_INDEPENDENT=ON时，必须传递有效的OP_API_DIR路径") + endif() + set(OP_API_SRC_DIR "${MODULE_OP_API_DIR}") + else() + # 旧结构：op_api嵌套在op_host目录下 + set(OP_API_SRC_DIR "${SOURCE_DIR}/op_api") + endif() + + # opapi 默认全部编译 + file(GLOB OPAPI_SRCS ${OP_API_SRC_DIR}/*.cpp) + if (OPAPI_SRCS) + # aclnn + add_opapi_modules() + target_sources(${OPHOST_NAME}_opapi_obj PRIVATE ${OPAPI_SRCS}) + else() + if (NOT TARGET ${OPHOST_NAME}_opapi_obj) + add_opapi_modules() + add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/opapi_stub.cpp + COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/opapi_stub.cpp + ) + target_sources(${OPHOST_NAME}_opapi_obj PRIVATE + ${CMAKE_CURRENT_BINARY_DIR}/opapi_stub.cpp + ) + endif() + endif() + file(GLOB OPAPI_HEADERS ${OP_API_SRC_DIR}/aclnn_*.h) + if (OPAPI_HEADERS) + target_sources(${OPHOST_NAME}_aclnn_exclude_headers INTERFACE ${OPAPI_HEADERS}) + endif() + + # 是否编译该算子已经由op_add_subdirectory和每个二级目录判断完毕，默认走到这里全编 + file(GLOB OPINFER_SRCS ${SOURCE_DIR}/*_infershape*.cpp) + add_infer_modules() + set(PROTO_STUB_FILE ${CMAKE_CURRENT_BINARY_DIR}/proto_stub.cpp) + + if(NOT EXISTS ${PROTO_STUB_FILE}) + file(WRITE ${PROTO_STUB_FILE} "// Auto-generated stub file\n") + endif() + + # 标记为生成的文件 + set_source_files_properties( + ${PROTO_STUB_FILE} + PROPERTIES GENERATED TRUE + ) + + if (OPINFER_SRCS) + target_sources(${OPHOST_NAME}_infer_obj PRIVATE ${OPINFER_SRCS}) + else() + target_sources(${OPHOST_NAME}_infer_obj PRIVATE ${PROTO_STUB_FILE}) + endif() + + file(GLOB_RECURSE SUB_OPTILING_SRC ${SOURCE_DIR}/op_tiling/*.cpp) + file(GLOB OPTILING_SRCS + ${SOURCE_DIR}/*fallback*.cpp + ${SOURCE_DIR}/*_tiling*.cpp + ${SOURCE_DIR}/op_tiling/arch35/*.cpp + ${SOURCE_DIR}/../op_graph/fallback_*.cpp + ${SOURCE_DIR}/../graph_plugin/fallback_*.cpp) + if (OPTILING_SRCS OR SUB_OPTILING_SRC) + # tiling + add_tiling_modules() + target_sources(${OPHOST_NAME}_tiling_obj PRIVATE ${OPTILING_SRCS} ${SUB_OPTILING_SRC}) + # target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE ${SOURCE_DIR}/../../ ${SOURCE_DIR}) + endif() + + file(GLOB AICPU_SRCS ${SOURCE_DIR}/*_aicpu*.cpp) + if(AICPU_SRCS) + add_aicpu_kernel_modules() + target_sources(${OPHOST_NAME}_aicpu_obj PRIVATE ${AICPU_SRCS}) + endif() + + if (MODULE_OPTYPE) + list(LENGTH MODULE_OPTYPE OpTypeLen) + list(LENGTH MODULE_ACLNNTYPE AclnnTypeLen) + if(NOT ${OpTypeLen} EQUAL ${AclnnTypeLen}) + message(FATAL_ERROR "OPTYPE AND ACLNNTYPE Should be One-to-One") + endif() + math(EXPR index "${OpTypeLen} - 1") + foreach(i RANGE ${index}) + list(GET MODULE_OPTYPE ${i} OpType) + list(GET MODULE_ACLNNTYPE ${i} AclnnType) + if (${AclnnType} STREQUAL "aclnn" OR ${AclnnType} STREQUAL "aclnn_inner" OR ${AclnnType} STREQUAL "aclnn_exclude") + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/${OpType}_def*.cpp) + + if (OPDEF_SRCS) + target_sources(${OPHOST_NAME}_opdef_${AclnnType}_obj INTERFACE ${OPDEF_SRCS}) + endif() + elseif(${AclnnType} STREQUAL "no_need_aclnn") + message(STATUS "aicpu or host aicpu no need aclnn.") + else() + message(FATAL_ERROR "ACLNN TYPE UNSUPPORTED, ONLY SUPPORT aclnn/aclnn_inner/aclnn_exclude") + endif() + endforeach() + else() + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/*_def*.cpp) + if(OPDEF_SRCS) + message(FATAL_ERROR + "Should Manually specify aclnn/aclnn_inner/aclnn_exclude\n" + "usage: add_modules_sources(OPTYPE optypes ACLNNTYPE aclnntypes)\n" + "example: add_modules_sources(OPTYPE add ACLNNTYPE aclnn_exclude)" + ) + endif() + endif() +endmacro() + +macro(add_modules_sources_with_soc) + set(oneValueArgs OP_API_INDEPENDENT OP_API_DIR) + set(multiValueArgs OPTYPE ACLNNTYPE) + + cmake_parse_arguments(MODULE "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + set(SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) + + # 该段代码作用为兼容op_api新旧目录结构(旧： 嵌套于op_host下； 新： 与op_host同级) + if (NOT DEFINED MODULE_OP_API_INDEPENDENT) + set(MODULE_OP_API_INDEPENDENT OFF) + endif() + if(MODULE_OP_API_INDEPENDENT) + # 新结构：op_api与op_host同级，需要指定有效路径 + if (NOT DEFINED MODULE_OP_API_DIR OR NOT EXISTS "${MODULE_OP_API_DIR}") + message(FATAL_ERROR "OP_API_INDEPENDENT=ON时，必须传递有效的OP_API_DIR路径") + endif() + set(OP_API_SRC_DIR "${MODULE_OP_API_DIR}") + else() + # 旧结构：op_api嵌套在op_host目录下 + set(OP_API_SRC_DIR "${SOURCE_DIR}/op_api") + endif() + + # opapi 默认全部编译 + file(GLOB OPAPI_SRCS ${OP_API_SRC_DIR}/*.cpp) + if (OPAPI_SRCS) + # aclnn + add_opapi_modules() + target_sources(${OPHOST_NAME}_opapi_obj PRIVATE ${OPAPI_SRCS}) + else() + if (NOT TARGET ${OPHOST_NAME}_opapi_obj) + add_opapi_modules() + add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/opapi_stub.cpp + COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/opapi_stub.cpp + ) + target_sources(${OPHOST_NAME}_opapi_obj PRIVATE + ${CMAKE_CURRENT_BINARY_DIR}/opapi_stub.cpp + ) + endif() + endif() + file(GLOB OPAPI_HEADERS ${OP_API_SRC_DIR}/aclnn_*.h) + if (OPAPI_HEADERS) + target_sources(${OPHOST_NAME}_aclnn_exclude_headers INTERFACE ${OPAPI_HEADERS}) + endif() + + # 是否编译该算子已经由op_add_subdirectory和每个二级目录判断完毕，默认走到这里全编 + + file(GLOB OPINFER_SRCS ${SOURCE_DIR}/*_infershape*.cpp) + foreach(ARCH ${ARCH_DIRECTORY}) + file(GLOB_RECURSE files ${SOURCE_DIR}/${ARCH}/*_infershape*.cpp) + list(APPEND OPINFER_SRCS ${files}) + endforeach() + + if (OPINFER_SRCS) + # proto + add_infer_modules() + target_sources(${OPHOST_NAME}_infer_obj PRIVATE ${OPINFER_SRCS}) + else() + if (NOT TARGET ${OPHOST_NAME}_infer_obj) + add_library(${OPHOST_NAME}_infer_obj OBJECT) + add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/proto_stub.cpp + COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/proto_stub.cpp + ) + target_sources(${OPHOST_NAME}_infer_obj PRIVATE + ${CMAKE_CURRENT_BINARY_DIR}/proto_stub.cpp + ) + endif() + endif() + + foreach(ARCH ${ARCH_DIRECTORY}) + file(GLOB_RECURSE files ${SOURCE_DIR}/${ARCH}/*_tiling*.cpp) + list(APPEND SUB_OPTILING_SRC ${files}) + endforeach() + file(GLOB OPTILING_SRCS + ${SOURCE_DIR}/*fallback*.cpp + ${SOURCE_DIR}/*_tiling*.cpp + ${SOURCE_DIR}/../op_graph/fallback_*.cpp + ${SOURCE_DIR}/../graph_plugin/fallback_*.cpp) + if (OPTILING_SRCS OR SUB_OPTILING_SRC) + # tiling + add_tiling_modules() + target_sources(${OPHOST_NAME}_tiling_obj PRIVATE ${OPTILING_SRCS} ${SUB_OPTILING_SRC}) + # target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE ${SOURCE_DIR}/../../ ${SOURCE_DIR}) + endif() + + file(GLOB AICPU_SRCS ${SOURCE_DIR}/*_aicpu*.cpp) + if(AICPU_SRCS) + add_aicpu_kernel_modules() + target_sources(${OPHOST_NAME}_aicpu_obj PRIVATE ${AICPU_SRCS}) + endif() + + if (MODULE_OPTYPE) + list(LENGTH MODULE_OPTYPE OpTypeLen) + list(LENGTH MODULE_ACLNNTYPE AclnnTypeLen) + if(NOT ${OpTypeLen} EQUAL ${AclnnTypeLen}) + message(FATAL_ERROR "OPTYPE AND ACLNNTYPE Should be One-to-One") + endif() + math(EXPR index "${OpTypeLen} - 1") + foreach(i RANGE ${index}) + list(GET MODULE_OPTYPE ${i} OpType) + list(GET MODULE_ACLNNTYPE ${i} AclnnType) + if (${AclnnType} STREQUAL "aclnn" OR ${AclnnType} STREQUAL "aclnn_inner" OR ${AclnnType} STREQUAL "aclnn_exclude") + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/${OpType}_def*.cpp) + + if (OPDEF_SRCS) + target_sources(${OPHOST_NAME}_opdef_${AclnnType}_obj INTERFACE ${OPDEF_SRCS}) + endif() + elseif(${AclnnType} STREQUAL "no_need_aclnn") + message(STATUS "aicpu or host aicpu no need aclnn.") + else() + message(FATAL_ERROR "ACLNN TYPE UNSUPPORTED, ONLY SUPPORT aclnn/aclnn_inner/aclnn_exclude") + endif() + endforeach() + else() + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/*_def*.cpp) + if(OPDEF_SRCS) + message(FATAL_ERROR + "Should Manually specify aclnn/aclnn_inner/aclnn_exclude\n" + "usage: add_modules_sources(OPTYPE optypes ACLNNTYPE aclnntypes)\n" + "example: add_modules_sources(OPTYPE add ACLNNTYPE aclnn_exclude)" + ) + endif() + endif() +endmacro() + +macro(add_mc2_modules_sources) + set(multiValueArgs OPTYPE ACLNNTYPE) + + cmake_parse_arguments(MODULE "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + set(SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) + + #opapi 默认全部编译 + file(GLOB OPAPI_SRCS ${SOURCE_DIR}/../op_api/*.cpp) + if (OPAPI_SRCS) + # aclnn + add_opapi_modules() + target_sources(${OPHOST_NAME}_opapi_obj PRIVATE ${OPAPI_SRCS}) + endif() + + # 获取算子层级目录名称，判断是否编译该算子 + get_filename_component(PARENT_DIR ${SOURCE_DIR} DIRECTORY) + get_filename_component(OP_NAME ${PARENT_DIR} NAME) + list(FIND ASCEND_OP_NAME ${OP_NAME} INDEX) + # 记录全局的COMPILED_OPS和COMPILED_OP_DIRS，其中COMPILED_OP_DIRS只记录到算子名，例如moe/moe_token_permute_with_routing_map_grad + set(COMPILED_OPS ${COMPILED_OPS} ${OP_NAME} CACHE STRING "Compiled Ops" FORCE) + set(COMPILED_OP_DIRS ${COMPILED_OP_DIRS} ${PARENT_DIR} CACHE STRING "Compiled Ops Dirs" FORCE) + + file(GLOB OPINFER_SRCS ${SOURCE_DIR}/*_infershape*.cpp) + if (OPINFER_SRCS) + # proto + add_infer_modules() + target_sources(${OPHOST_NAME}_infer_obj PRIVATE ${OPINFER_SRCS}) + endif() + + file(GLOB_RECURSE OPTILING_SRCS + ${SOURCE_DIR}/op_tiling/*.cpp + ${SOURCE_DIR}/op_tiling/arch35/*.cpp + ${SOURCE_DIR}/op_tiling/common/*.cpp + ${SOURCE_DIR}/../op_graph/fallback*.cpp + ) + if (OPTILING_SRCS) + # tiling + add_tiling_modules() + target_sources(${OPHOST_NAME}_tiling_obj PRIVATE + ${OPTILING_SRCS} + ${OPS_TRANSFORMER_DIR}/mc2/common/src/matmul_formulaic_tiling.cpp + ${OPS_TRANSFORMER_DIR}/mc2/common/src/mc2_tiling_utils.cpp + ${OPS_TRANSFORMER_DIR}/mc2/common/src/mc2_matmul_tiling_cfg.cpp + ${OPS_TRANSFORMER_DIR}/mc2/common/src/mc2_log.cpp + ${OPS_TRANSFORMER_DIR}/mc2/3rd/ops_legacy/op_tiling/op_cache_tiling.cpp + ${OPS_TRANSFORMER_DIR}/mc2/3rd/ops_legacy/op_tiling/runtime_kb_api.cpp + ) + endif() + + file(GLOB GENTASK_SRCS + ${SOURCE_DIR}/../op_graph/*_gen_task*.cpp + ) + if(GENTASK_SRCS) + add_opmaster_ct_gentask_modules() + target_sources(${OPHOST_NAME}_opmaster_ct_gentask_obj PRIVATE ${GENTASK_SRCS}) + endif() + + file(GLOB AICPU_SRCS ${SOURCE_DIR}/*_aicpu*.cpp) + if(AICPU_SRCS) + add_aicpu_kernel_modules() + target_sources(${OPHOST_NAME}_aicpu_obj PRIVATE ${AICPU_SRCS}) + endif() + + if (MODULE_OPTYPE) + list(LENGTH MODULE_OPTYPE OpTypeLen) + list(LENGTH MODULE_ACLNNTYPE AclnnTypeLen) + if(NOT ${OpTypeLen} EQUAL ${AclnnTypeLen}) + message(FATAL_ERROR "OPTYPE AND ACLNNTYPE Should be One-to-One") + endif() + math(EXPR index "${OpTypeLen} - 1") + foreach(i RANGE ${index}) + list(GET MODULE_OPTYPE ${i} OpType) + list(GET MODULE_ACLNNTYPE ${i} AclnnType) + if (${AclnnType} STREQUAL "aclnn" OR ${AclnnType} STREQUAL "aclnn_inner" OR ${AclnnType} STREQUAL "aclnn_exclude") + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/${OpType}_def*.cpp) + + if (OPDEF_SRCS) + target_sources(${OPHOST_NAME}_opdef_${AclnnType}_obj INTERFACE ${OPDEF_SRCS}) + endif() + elseif(${AclnnType} STREQUAL "no_need_aclnn") + message(STATUS "aicpu or host aicpu no need aclnn.") + else() + message(FATAL_ERROR "ACLNN TYPE UNSUPPORTED, ONLY SUPPORT aclnn/aclnn_inner/aclnn_exclude") + endif() + endforeach() + else() + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/*_def*.cpp) + if(OPDEF_SRCS) + message(FATAL_ERROR + "Should Manually specify aclnn/aclnn_inner/aclnn_exclude\n" + "usage: add_modules_sources(OPTYPE optypes ACLNNTYPE aclnntypes)\n" + "example: add_modules_sources(OPTYPE add ACLNNTYPE aclnn_exclude)" + ) + endif() + endif() +endmacro() + + +# usage: add_modules_sources_aicpu(OPTYPE ACLNNTYPE DEPENDENCIES COMPUTE_UNIT TILING_DIR DISABLE_IN_OPP) +# ACLNNTYPE 支持类型aclnn/aclnn_inner/aclnn_exclude +# OPTYPE 和 ACLNNTYPE 需一一对应 +# DEPENDENCIES 指定依赖的算子名称列表，如果开启 experimental，则会优先加载 experimental 下的算子 +# COMPUTE_UNIT 设置支持芯片版本号，必须与TILING_DIR一一对应，示例：ascend910b ascend950 +# TILING_DIR 设置所支持芯片类型对应的tiling文件目录，必须与COMPUTE_UNIT一一对应，示例：arch32 arch35 +# DISABLE_IN_OPP 设置是否在opp包中编译tiling文件，布尔类型：TRUE，FALSE +macro(add_modules_sources_aicpu) + set(oneValueArgs DISABLE_IN_OPP) + set(multiValueArgs OPTYPE ACLNNTYPE DEPENDENCIES COMPUTE_UNIT TILING_DIR) + + cmake_parse_arguments(MODULE "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + set(SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) + get_filename_component(OP_NAME ${SOURCE_DIR} NAME) + + add_dependent_ops("${MODULE_DEPENDENCIES}") + + + # opapi l0 默认全部编译 + file(GLOB OPAPI_L0_SRCS ${SOURCE_DIR}/op_api/*.cpp) + list(FILTER OPAPI_L0_SRCS EXCLUDE REGEX "aclnn_") + if(OPAPI_L0_SRCS) + add_opapi_modules() + target_sources(${OPHOST_NAME}_opapi_obj PRIVATE ${OPAPI_L0_SRCS}) + endif() + + file(GLOB OPAPI_HEADERS ${SOURCE_DIR}/op_api/aclnn_*.h) + if(OPAPI_HEADERS) + target_sources(${OPHOST_NAME}_aclnn_exclude_headers INTERFACE ${OPAPI_HEADERS}) + endif() + + file(GLOB OPAPI_L2_SRCS ${SOURCE_DIR}/op_api/aclnn_*.cpp) + if(OPAPI_L2_SRCS) + add_opapi_modules() + target_sources(${OPHOST_NAME}_opapi_obj PRIVATE ${OPAPI_L2_SRCS}) + endif() + + file(GLOB OPINFER_SRCS ${SOURCE_DIR}/op_host/*_infershape*.cpp) + if(OPINFER_SRCS) + add_infer_modules() + target_sources(${OPHOST_NAME}_infer_obj PRIVATE ${OPINFER_SRCS}) + endif() + + file(GLOB OPTILING_SRCS ${SOURCE_DIR}/*_tiling*.cpp) + if(OPTILING_SRCS) + add_tiling_modules() + target_sources(${OPHOST_NAME}_tiling_obj PRIVATE ${OPTILING_SRCS}) + else() + if (NOT TARGET ${OPHOST_NAME}_tiling_obj) + add_tiling_modules() + add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/optiling_stub.cpp + COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/optiling_stub.cpp + ) + target_sources(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_BINARY_DIR}/optiling_stub.cpp + ) + endif() + endif() + + + file(GLOB AICPU_SRCS ${SOURCE_DIR}/op_kernel_aicpu/*_aicpu*.cpp) + + if(AICPU_SRCS) + if(NOT BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + add_aicpu_kernel_modules() + target_sources(${OPHOST_NAME}_aicpu_obj PRIVATE ${AICPU_SRCS}) + else() + file(GLOB AICPU_JSON_FILE ${SOURCE_DIR}/op_kernel_aicpu/*.json) + add_aicpu_cust_kernel_modules(${OP_NAME} ${AICPU_SRCS} ${AICPU_JSON_FILE}) + endif() + endif() + + if(MODULE_OPTYPE) + list(LENGTH MODULE_OPTYPE OpTypeLen) + list(LENGTH MODULE_ACLNNTYPE AclnnTypeLen) + if(NOT ${OpTypeLen} EQUAL ${AclnnTypeLen}) + message(FATAL_ERROR "OPTYPE AND ACLNNTYPE Should be One-to-One") + endif() + math(EXPR index "${OpTypeLen} - 1") + foreach(i RANGE ${index}) + list(GET MODULE_OPTYPE ${i} OpType) + list(GET MODULE_ACLNNTYPE ${i} AclnnType) + if(${AclnnType} STREQUAL "aclnn" + OR ${AclnnType} STREQUAL "aclnn_inner" + OR ${AclnnType} STREQUAL "aclnn_exclude" + ) + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/op_host/${OpType}_def*.cpp) + if(OPDEF_SRCS) + target_sources(${OPHOST_NAME}_opdef_${AclnnType}_obj INTERFACE ${OPDEF_SRCS}) + endif() + else() + message(FATAL_ERROR "ACLNN TYPE UNSUPPORTED, ONLY SUPPORT aclnn/aclnn_inner/aclnn_exclude") + endif() + endforeach() + else() + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/op_host/*_def*.cpp) + if(OPDEF_SRCS) + message( + FATAL_ERROR + "Should Manually specify aclnn/aclnn_inner/aclnn_exclude\n" + "usage: add_modules_sources_aicpu(OPTYPE optypes ACLNNTYPE aclnntypes)\n" + "example: add_modules_sources_aicpu(OPTYPE add ACLNNTYPE aclnn_exclude)" + ) + endif() + endif() + + file(GLOB OP_GRAPH_PROTO_HEADERS ${SOURCE_DIR}/op_graph/*_proto*.h) + if(OP_GRAPH_PROTO_HEADERS) + target_sources(${GRAPH_PLUGIN_NAME}_proto_headers INTERFACE ${OP_GRAPH_PROTO_HEADERS}) + endif() + + set(ENABLE_AICPU ON CACHE BOOL "enable aicpu kernel" FORCE) + +endmacro() + +# 添加opapi object +function(add_opapi_modules) + if (NOT TARGET ${OPHOST_NAME}_opapi_obj) + add_library(${OPHOST_NAME}_opapi_obj OBJECT) + unset(OPAPI_UT_DEPEND_INC) + if(UT_TEST_ALL OR OP_API_UT) + set(OPAPI_UT_DEPEND_INC ${UT_PATH}/op_api/stub) + endif() + target_include_directories(${OPHOST_NAME}_opapi_obj + PRIVATE + ${OPAPI_INCLUDE} + ${OPAPI_UT_DEPEND_INC} + $ + $ + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $<$:$> + ${OPS_TRANSFORMER_DIR}/mc2/common/inc + ${OPS_TRANSFORMER_DIR}/mc2/3rd + ) + target_compile_definitions(${OPHOST_NAME}_opapi_obj PRIVATE + _GLIBCXX_USE_CXX11_ABI=0 + BUILD_OPEN_PROJECT_API=1 + ) + target_compile_options(${OPHOST_NAME}_opapi_obj + PRIVATE + -Dgoogle=ascend_private + -DACLNN_LOG_FMT_CHECK + ) + target_link_libraries(${OPHOST_NAME}_opapi_obj + PUBLIC + $ + $,intf_llt_pub_asan_cxx17,intf_pub_cxx17>> + -Wl,--whole-archive + ops_aclnn + -Wl,--no-whole-archive + $<$:$> + nnopbase + profapi + ge_common_base + ascend_dump + ascendalog + dl + ) + endif() +endfunction() + +set(INFER_OBJ_INCLUDE + ${OP_PROTO_INCLUDE} + $ + $ + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $<$:$> + ${OPS_TRANSFORMER_DIR}/mc2/common/inc + ${OPS_TRANSFORMER_DIR}/mc2/3rd +) + +# 添加infer object +function(add_infer_modules) + if (NOT TARGET ${OPHOST_NAME}_infer_obj) + add_library(${OPHOST_NAME}_infer_obj OBJECT) + target_include_directories(${OPHOST_NAME}_infer_obj + PRIVATE ${INFER_OBJ_INCLUDE} + ) + target_compile_definitions(${OPHOST_NAME}_infer_obj + PRIVATE + LOG_CPP + OPS_UTILS_LOG_SUB_MOD_NAME="OP_PROTO" + $<$:ASCEND_OPSPROTO_UT> + ) + target_compile_options(${OPHOST_NAME}_infer_obj + PRIVATE + $<$>:-DDISABLE_COMPILE_V1> + -Dgoogle=ascend_private + -fvisibility=hidden + ) + target_link_libraries(${OPHOST_NAME}_infer_obj + PRIVATE + $,intf_llt_pub_asan_cxx17,intf_pub_cxx17>> + $ + $ + $<$:$> + $<$:$> + -Wl,--whole-archive + rt2_registry_static + -Wl,--no-whole-archive + -Wl,--no-as-needed + exe_graph + graph + graph_base + register + ascendalog + error_manager + platform + -Wl,--as-needed + c_sec + ) + endif() +endfunction() + +# 添加tiling object +function(add_tiling_modules) + if (NOT TARGET ${OPHOST_NAME}_tiling_obj) + add_library(${OPHOST_NAME}_tiling_obj OBJECT) + add_dependencies(${OPHOST_NAME}_tiling_obj json) + target_include_directories(${OPHOST_NAME}_tiling_obj + PRIVATE ${OP_TILING_INCLUDE} + $ + $ + + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $<$:$> + ${OPS_TRANSFORMER_DIR}/mc2/common/inc + ${OPS_TRANSFORMER_DIR}/mc2/3rd + ) + target_compile_definitions(${OPHOST_NAME}_tiling_obj + PRIVATE + LOG_CPP + OPS_UTILS_LOG_SUB_MOD_NAME="OP_TILING" + $<$:ASCEND_OPTILING_UT> + ) + target_compile_options(${OPHOST_NAME}_tiling_obj + PRIVATE + $<$>:-DDISABLE_COMPILE_V1> + -Dgoogle=ascend_private + -fvisibility=hidden + -fno-strict-aliasing + ) + target_link_libraries(${OPHOST_NAME}_tiling_obj + PRIVATE + $,intf_llt_pub_asan_cxx17,intf_pub_cxx17>> + $ + $ + $<$:$> + $<$:$> + -Wl,--whole-archive + rt2_registry_static + -Wl,--no-whole-archive + -Wl,--no-as-needed + graph + graph_base + exe_graph + platform + register + # ascendalog + error_manager + -Wl,--as-needed + -Wl,--whole-archive + tiling_api + -Wl,--no-whole-archive + # mmpa + c_sec + ) + endif() +endfunction() + +function(add_graph_plugin_modules) + if(NOT TARGET ${GRAPH_PLUGIN_NAME}_obj) + add_library(${GRAPH_PLUGIN_NAME}_obj OBJECT) + target_include_directories(${GRAPH_PLUGIN_NAME}_obj PRIVATE + ${OP_PROTO_INCLUDE} + ) + target_compile_definitions(${GRAPH_PLUGIN_NAME}_obj PRIVATE OPS_UTILS_LOG_SUB_MOD_NAME="GRAPH_PLUGIN" LOG_CPP) + target_compile_options( + ${GRAPH_PLUGIN_NAME}_obj PRIVATE $<$>:-DDISABLE_COMPILE_V1> -Dgoogle=ascend_private + -fvisibility=hidden + ) + target_link_libraries( + ${GRAPH_PLUGIN_NAME}_obj + PRIVATE $,intf_llt_pub_asan_cxx17,intf_pub_cxx17>> + $ + $<$:$> + $<$:$> + ) + endif() +endfunction() + +# 添加gentask object +function(add_opmaster_ct_gentask_modules) + message(STATUS "add_opmaster_ct_gentask_modules start") + if (NOT TARGET ${OPHOST_NAME}_opmaster_ct_gentask_obj) + add_library(${OPHOST_NAME}_opmaster_ct_gentask_obj OBJECT) + add_dependencies(${OPHOST_NAME}_opmaster_ct_gentask_obj json) + + target_include_directories(${OPHOST_NAME}_opmaster_ct_gentask_obj + PRIVATE ${OP_TILING_INCLUDE} + $<$:$> + $<$:$> + ) + target_compile_definitions(${OPHOST_NAME}_opmaster_ct_gentask_obj + PRIVATE + LOG_CPP + ) + target_compile_options(${OPHOST_NAME}_opmaster_ct_gentask_obj + PRIVATE + $<$>:-DDISABLE_COMPILE_V1> + -Dgoogle=ascend_private + -fvisibility=hidden + -fno-strict-aliasing + ) + message(STATUS "xxxx compile add_opmaster_ct_gentask_modules") + target_link_libraries(${OPHOST_NAME}_opmaster_ct_gentask_obj + PRIVATE + $ + $<$:$> + $<$:$> + # $<$:$> + # $<$>:$> + ) + endif() +endfunction() + + +# usage: add_graph_plugin_sources() +macro(add_graph_plugin_sources) + set(SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) + + # 获取算子层级目录名称，判断是否编译该算子 + get_filename_component(PARENT_DIR ${SOURCE_DIR} DIRECTORY) + get_filename_component(OP_NAME ${PARENT_DIR} NAME) + if(DEFINED ASCEND_OP_NAME + AND NOT "${ASCEND_OP_NAME}" STREQUAL "" + AND NOT "${ASCEND_OP_NAME}" STREQUAL "all" + AND NOT "${ASCEND_OP_NAME}" STREQUAL "ALL" + ) + if(NOT ${OP_NAME} IN_LIST ASCEND_OP_NAME) + return() + endif() + endif() + + file(GLOB GRAPH_PLUGIN_SRCS + ${SOURCE_DIR}/*_graph_plugin*.cpp + ) + if(GRAPH_PLUGIN_SRCS) + add_graph_plugin_modules() + target_sources(${GRAPH_PLUGIN_NAME}_obj PRIVATE ${GRAPH_PLUGIN_SRCS}) + endif() + + file(GLOB GRAPH_PLUGIN_PROTO_HEADERS ${SOURCE_DIR}/*_proto*.h) + if(GRAPH_PLUGIN_PROTO_HEADERS) + target_sources(${GRAPH_PLUGIN_NAME}_proto_headers INTERFACE ${GRAPH_PLUGIN_PROTO_HEADERS}) + endif() +endmacro() + +function(protobuf_generate_external comp c_var h_var) + if (NOT ARGN) + message(SEND_ERROR "Error: protobuf_generate_external() called without any proto files") + return() + endif() + + set(${c_var}) + set(${h_var}) + set(_add_target FALSE) + + set(extra_option "") + foreach(arg ${ARGN}) + if ("${arg}" MATCHES "--proto_path") + set(extra_option ${arg}) + endif() + endforeach() + + foreach(file ${ARGN}) + if ("${file}" STREQUAL "TARGET") + set(_add_target TRUE) + continue() + endif() + + if ("${file}" MATCHES "--proto_path") + continue() + endif() + + get_filename_component(abs_file ${file} ABSOLUTE) + get_filename_component(file_name ${file} NAME_WE) + get_filename_component(file_dir ${abs_file} PATH) + get_filename_component(parent_subdir ${file_dir} NAME) + + if ("${parent_subdir}" STREQUAL "proto") + set(proto_output_path ${CMAKE_BINARY_DIR}/proto/${comp}/proto) + else() + set(proto_output_path ${CMAKE_BINARY_DIR}/proto/${comp}/proto/${parent_subdir}) + endif() + list(APPEND ${c_var} "${proto_output_path}/${file_name}.pb.cc") + list(APPEND ${h_var} "${proto_output_path}/${file_name}.pb.h") + + add_custom_command( + OUTPUT "${proto_output_path}/${file_name}.pb.cc" "${proto_output_path}/${file_name}.pb.h" + WORKING_DIRECTORY ${PROJECT_SOURCE_DIR} + COMMAND ${CMAKE_COMMAND} -E make_directory "${proto_output_path}" + COMMAND ${CMAKE_COMMAND} -E echo "generate proto cpp_out ${comp} by ${abs_file}" + COMMAND ${Protobuf_PROTOC_EXECUTABLE} -I${file_dir} ${extra_option} --cpp_out=${proto_output_path} ${abs_file} + DEPENDS ${abs_file} ascend_protobuf_build_transformer + COMMENT "Running C++ protocol buffer compiler on ${file}" VERBATIM) + + endforeach() + + if (_add_target) + add_custom_target( + ${comp} DEPENDS ${${c_var}} ${${h_var}}) + endif() + + set_source_files_properties(${${c_var}} ${${h_var}} PROPERTIES GENERATED TRUE) + set(${c_var} ${${c_var}} PARENT_SCOPE) + set(${h_var} ${${h_var}} PARENT_SCOPE) + +endfunction() + +function(add_onnx_plugin_modules) + if (NOT TARGET ${ONNX_PLUGIN_NAME}_obj) + set(ge_onnx_proto_srcs + ${ASCEND_DIR}/include/proto/ge_onnx.proto) + + protobuf_generate_external(onnx ge_onnx_proto_cc ge_onnx_proto_h ${ge_onnx_proto_srcs}) + + add_library(${ONNX_PLUGIN_NAME}_obj OBJECT ${ge_onnx_proto_h}) + # 为特定目标设置C++14标准 + set_target_properties(${ONNX_PLUGIN_NAME}_obj PROPERTIES + CXX_STANDARD 14 + CXX_STANDARD_REQUIRED ON + CXX_EXTENSIONS OFF + ) + target_include_directories(${ONNX_PLUGIN_NAME}_obj PRIVATE ${OP_PROTO_INCLUDE} ${Protobuf_INCLUDE} ${Protobuf_PATH} ${CMAKE_BINARY_DIR}/proto ${ONNX_PLUGIN_COMMON_INCLUDE} ${JSON_INCLUDE_DIR} ${ABSL_SOURCE_DIR}) + target_compile_definitions(${ONNX_PLUGIN_NAME}_obj PRIVATE OPS_UTILS_LOG_SUB_MOD_NAME="ONNX_PLUGIN" LOG_CPP) + + if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + target_compile_options( + ${ONNX_PLUGIN_NAME}_obj PRIVATE -Dgoogle=ascend_private -fvisibility=hidden -Wno-shadow -Wno-unused-parameter + ) + else() + target_compile_options( + ${ONNX_PLUGIN_NAME}_obj PRIVATE $<$>:-DDISABLE_COMPILE_V1> -Dgoogle=ascend_private + -fvisibility=hidden -Wno-shadow -Wno-unused-parameter + ) + endif() + + target_link_libraries( + ${ONNX_PLUGIN_NAME}_obj + PRIVATE $,intf_llt_pub_asan_cxx14,intf_pub_cxx14>> + $ + $<$:$> + $<$:$> + json + ) + endif() + +endfunction() + +macro(add_onnx_plugin_sources) + set(SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) + + file(GLOB ONNX_PLUGIN_SRCS ${SOURCE_DIR}/*_onnx_plugin.cpp) + if(ONNX_PLUGIN_SRCS) + add_onnx_plugin_modules() + target_sources(${ONNX_PLUGIN_NAME}_obj PRIVATE ${ONNX_PLUGIN_SRCS}) + else() + message(WARNING "No onnx plugin source files found in ${SOURCE_DIR}") + endif() +endmacro() diff --git a/csrc/ascend/cmake/opbuild.cmake b/csrc/ascend/cmake/opbuild.cmake new file mode 100644 index 000000000..5bbd03023 --- /dev/null +++ b/csrc/ascend/cmake/opbuild.cmake @@ -0,0 +1,307 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +# ###################################################################################################################### +# 调用opbuild工具，生成aclnn/aclnnInner/.ini的算子信息库 等文件 +# generate outpath: ${ASCEND_AUTOGEN_PATH}/${sub_dir} +# ###################################################################################################################### +function(gen_opbuild_target) + set(oneValueArgs TARGET PREFIX GENACLNN OUT_DIR OUT_SUB_DIR) + set(multiValueArgs IN_SRCS OUT_SRCS OUT_HEADERS) + cmake_parse_arguments(OPBUILD "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + if(NOT OPBUILD_IN_SRCS) + message(STATUS "No ${OPBUILD_PREFIX} srcs, skip ${OPBUILD_TARGET}") + return() + endif() + + add_library(gen_op_host_${OPBUILD_PREFIX} SHARED ${OPBUILD_IN_SRCS}) + target_link_libraries(gen_op_host_${OPBUILD_PREFIX} PRIVATE + $ + exe_graph + register + c_sec + ) + target_compile_options(gen_op_host_${OPBUILD_PREFIX} PRIVATE + -fno-common + ) + + add_custom_command(OUTPUT ${OPBUILD_OUT_SRCS} ${OPBUILD_OUT_HEADERS} + COMMAND OPS_PROTO_SEPARATE=1 + OPS_PROJECT_NAME=${OPBUILD_PREFIX} + OPS_ACLNN_GEN=${OPBUILD_GENACLNN} + OPS_PRODUCT_NAME=\"${ASCEND_COMPUTE_UNIT}\" + ${OP_BUILD_TOOL} + $ + ${OPBUILD_OUT_DIR}/${OPBUILD_OUT_SUB_DIR} + ) + + add_custom_target(${OPBUILD_TARGET} + DEPENDS ${OPBUILD_OUT_SRCS} ${OPBUILD_OUT_HEADERS} + ) + add_dependencies(${OPBUILD_TARGET} gen_op_host_${OPBUILD_PREFIX}) + if(TARGET op_build) + add_dependencies(${OPBUILD_TARGET} op_build) + endif() +endfunction() + +function(append_versioned_aclnn_outputs op_def_src file_prefix op_name out_dir out_srcs_var out_headers_var) + if(NOT EXISTS "${op_def_src}") + return() + endif() + + file(READ "${op_def_src}" op_def_content) + string(REGEX MATCHALL "\\.Version\\([ \t]*[A-Za-z0-9_]+[ \t]*\\)" version_exprs "${op_def_content}") + if(NOT version_exprs) + return() + endif() + + set(version_values) + foreach(version_expr ${version_exprs}) + string(REGEX REPLACE ".*\\.Version\\([ \t]*([A-Za-z0-9_]+)[ \t]*\\).*" "\\1" version_token "${version_expr}") + set(version_value "") + if(version_token MATCHES "^[0-9]+$") + set(version_value "${version_token}") + else() + string(REGEX MATCH "[A-Za-z_][A-Za-z0-9_]*[ \t]+${version_token}[ \t]*=[ \t]*[0-9]+" version_decl "${op_def_content}") + if(version_decl) + string(REGEX REPLACE ".*=[ \t]*([0-9]+).*" "\\1" version_value "${version_decl}") + endif() + endif() + + if(version_value AND version_value GREATER 1) + list(APPEND version_values "${version_value}") + endif() + endforeach() + + if(NOT version_values) + return() + endif() + + list(REMOVE_DUPLICATES version_values) + foreach(version_value ${version_values}) + list(APPEND ${out_srcs_var} ${out_dir}/${file_prefix}_${op_name}_v${version_value}.cpp) + list(APPEND ${out_headers_var} ${out_dir}/${file_prefix}_${op_name}_v${version_value}.h) + endforeach() + + set(${out_srcs_var} "${${out_srcs_var}}" PARENT_SCOPE) + set(${out_headers_var} "${${out_headers_var}}" PARENT_SCOPE) +endfunction() + +function(gen_aclnn_classify host_obj prefix ori_out_srcs ori_out_headers opbuild_out_srcs opbuild_out_headers) + get_target_property(module_sources ${host_obj} INTERFACE_SOURCES) + set(sub_dir) + # aclnn\aclnnExc以aclnn开头，aclnnInner以aclnnInner开头 + if("${prefix}" STREQUAL "aclnn") + set(file_prefix "aclnn") + set(need_gen_aclnn 1) + elseif("${prefix}" STREQUAL "aclnnInner") + set(sub_dir inner) + set(file_prefix "aclnnInner") + set(need_gen_aclnn 1) + elseif("${prefix}" STREQUAL "aclnnExc") + set(sub_dir exc) + set(file_prefix "aclnn") + set(need_gen_aclnn 0) + else() + message(FATAL_ERROR "UnSupported aclnn prefix type, must be in aclnn/aclnnInner/aclnnExc") + endif() + + set(out_src_path ${ASCEND_AUTOGEN_PATH}/${sub_dir}) + file(MAKE_DIRECTORY ${out_src_path}) + get_filename_component(out_src_path ${out_src_path} REALPATH) + set(in_srcs) + set(out_srcs) + set(out_headers) + if(module_sources) + foreach(file ${module_sources}) + get_filename_component(name_without_ext ${file} NAME_WE) + string(REGEX REPLACE "_def$" "" _op_name ${name_without_ext}) + list(APPEND in_srcs ${file}) + list(APPEND out_srcs ${out_src_path}/${file_prefix}_${_op_name}.cpp) + list(APPEND out_headers ${out_src_path}/${file_prefix}_${_op_name}.h) + if(need_gen_aclnn) + append_versioned_aclnn_outputs("${file}" "${file_prefix}" "${_op_name}" "${out_src_path}" out_srcs out_headers) + endif() + endforeach() + endif() + # opbuild_gen_aclnn/opbuild_gen_aclnnInner/opbuild_gen_aclnnExc + if("${prefix}" STREQUAL "aclnnExc") + get_target_property(exclude_headers ${OPHOST_NAME}_aclnn_exclude_headers INTERFACE_SOURCES) + if(exclude_headers) + set(${opbuild_out_headers} ${ori_out_headers} ${exclude_headers} PARENT_SCOPE) + endif() + else() + set(${opbuild_out_srcs} ${ori_out_srcs} ${out_srcs} PARENT_SCOPE) + set(${opbuild_out_headers} ${ori_out_headers} ${out_headers} PARENT_SCOPE) + endif() +endfunction() + +function(gen_aclnn_master_header aclnn_master_header_name aclnn_master_header opbuild_out_headers) + # 规范化，防止生成的代码编译失败 + string(REGEX REPLACE "[^a-zA-Z0-9_]" "_" aclnn_master_header_name "${aclnn_master_header_name}") + string(TOUPPER ${aclnn_master_header_name} aclnn_master_header_name) + + # 生成include内容 + set(aclnn_all_header_include_content "") + foreach(header_file ${opbuild_out_headers}) + get_filename_component(header_name ${header_file} NAME) + set(aclnn_all_header_include_content "${aclnn_all_header_include_content}#include \"${header_name}\"\n") + endforeach() + + # 根据模板生成头文件 + message(STATUS "create aclnn master header file: ${aclnn_master_header}") + configure_file( + "${CMAKE_CURRENT_SOURCE_DIR}/cmake/aclnn_ops_transformer.h.in" + "${aclnn_master_header}" + @ONLY + ) +endfunction() + +function(gen_aclnn_with_opdef) + set(opbuild_out_srcs) + set(opbuild_out_headers) + gen_aclnn_classify(${OPHOST_NAME}_opdef_aclnn_obj aclnn "${opbuild_out_srcs}" "${opbuild_out_headers}" + opbuild_out_srcs opbuild_out_headers) + gen_aclnn_classify(${OPHOST_NAME}_opdef_aclnn_inner_obj aclnnInner "${opbuild_out_srcs}" "${opbuild_out_headers}" + opbuild_out_srcs opbuild_out_headers) + gen_aclnn_classify(${OPHOST_NAME}_opdef_aclnn_exclude_obj aclnnExc "${opbuild_out_srcs}" "${opbuild_out_headers}" + opbuild_out_srcs opbuild_out_headers) + + # 创建汇总头文件 + if(NOT ENABLE_BUILT_IN) + set(aclnn_master_header_name "aclnn_ops_transformer_${VENDOR_NAME}") + else() + set(aclnn_master_header_name "aclnn_ops_transformer") + endif() + set(aclnn_master_header "${CMAKE_CURRENT_BINARY_DIR}/${aclnn_master_header_name}.h") + gen_aclnn_master_header(${aclnn_master_header_name} "${aclnn_master_header}" "${opbuild_out_headers}") + + set(mc2_op_aclnn_name + "all_gather_matmul" + "all_to_all_all_gather_batch_matmul" + "allto_allv_grouped_mat_mul" + "batch_matmul_reduce_scatter_all_to_all" + "distribute_barrier" + "distribute_barrier_v2" + "grouped_mat_mul_allto_allv" + "inplace_matmul_all_reduce_add_rms_norm" + "inplace_quant_matmul_all_reduce_add_rms_norm" + "inplace_weight_quant_matmul_all_reduce_add_rms_norm" + "matmul_all_reduce" + "matmul_all_reduce_add_rms_norm" + "matmul_all_reduce_v2" + "matmul_reduce_scatter" + "moe_distribute_combine" + "moe_distribute_combine_add_rms_norm" + "moe_distribute_combine_add_rms_norm_v2" + "moe_distribute_combine_v2" + "moe_distribute_combine_v3" + "moe_distribute_dispatch" + "moe_distribute_dispatch_v2" + "moe_distribute_dispatch_v3" + "moe_update_expert" + "weight_quant_matmul_all_reduce" + "weight_quant_matmul_all_reduce_add_rms_norm" + ) + set(mc2_aclnn_master_headers "") + foreach(op_aclnn_name ${mc2_op_aclnn_name}) + if (NOT ENABLE_BUILT_IN AND NOT ("${ASCEND_OP_NAME}" STREQUAL "ALL")) + foreach(op_name IN LISTS ASCEND_OP_NAME) + file(GLOB matching_file "${OPS_TRANSFORMER_DIR}/mc2/${op_name}/op_api/aclnn_${op_aclnn_name}.h") + list(APPEND mc2_aclnn_master_headers ${matching_file}) + endforeach() + else() + file(GLOB matching_file "${OPS_TRANSFORMER_DIR}/mc2/*/op_api/aclnn_${op_aclnn_name}.h") + list(APPEND mc2_aclnn_master_headers ${matching_file}) + endif() + endforeach() + + # 将头文件安装到packages/vendors/vendor_name/op_api/include + if (NOT ENABLE_BUILT_IN) + install(FILES ${opbuild_out_headers} DESTINATION ${ACLNN_INC_INSTALL_DIR} OPTIONAL) + install(FILES ${aclnn_master_header} DESTINATION ${ACLNN_INC_INSTALL_DIR} OPTIONAL) + if (BUILD_OPEN_PROJECT AND mc2_aclnn_master_headers) + install(FILES ${mc2_aclnn_master_headers} DESTINATION ${ACLNN_INC_INSTALL_DIR} OPTIONAL) + endif() + else() + install(FILES ${opbuild_out_headers} DESTINATION ${ACLNN_INC_INSTALL_DIR} OPTIONAL) + install(FILES ${aclnn_master_header} DESTINATION ${ACLNN_INC_INSTALL_DIR} OPTIONAL) + install(FILES ${opbuild_out_headers} DESTINATION ${ACLNN_INC_LEVEL2_INSTALL_DIR} OPTIONAL) + install(FILES ${aclnn_master_header} DESTINATION ${ACLNN_INC_LEVEL2_INSTALL_DIR} OPTIONAL) + if (BUILD_OPEN_PROJECT) + install(FILES ${mc2_aclnn_master_headers} DESTINATION ${ACLNN_INC_INSTALL_DIR} OPTIONAL) + install(FILES ${mc2_aclnn_master_headers} DESTINATION ${ACLNN_INC_LEVEL2_INSTALL_DIR} OPTIONAL) + endif() + endif() + + if (ENABLE_STATIC) + install(FILES ${opbuild_out_headers} DESTINATION ${CMAKE_BINARY_DIR}/static_library_files/include/aclnnop OPTIONAL) + install(FILES ${aclnn_master_header} DESTINATION ${CMAKE_BINARY_DIR}/static_library_files/include/aclnnop OPTIONAL) + install(FILES ${opbuild_out_headers} DESTINATION ${CMAKE_BINARY_DIR}/static_library_files/include/aclnnop/level2 OPTIONAL) + install(FILES ${aclnn_master_header} DESTINATION ${CMAKE_BINARY_DIR}/static_library_files/include/aclnnop/level2 OPTIONAL) + if (BUILD_OPEN_PROJECT) + install(FILES ${mc2_aclnn_master_headers} DESTINATION ${CMAKE_BINARY_DIR}/static_library_files/include/aclnnop OPTIONAL) + install(FILES ${mc2_aclnn_master_headers} DESTINATION ${CMAKE_BINARY_DIR}/static_library_files/include/aclnnop/level2 OPTIONAL) + endif() + endif() + + # ascendc_impl_gen depends opbuild_custom_gen_aclnn_all, for opbuild will generate .ini + set(dependency_list) + if(TARGET opbuild_gen_aclnn) + list(APPEND dependency_list opbuild_gen_aclnn) + endif() + if(TARGET opbuild_gen_aclnnInner) + list(APPEND dependency_list opbuild_gen_aclnnInner) + endif() + if(TARGET opbuild_gen_aclnnExc) + list(APPEND dependency_list opbuild_gen_aclnnExc) + endif() + if(NOT dependency_list) + message(STATUS "no operator info to generate") + return() + endif() + add_custom_target(opbuild_custom_gen_aclnn_all) + add_dependencies(opbuild_custom_gen_aclnn_all ${dependency_list}) + if(opbuild_out_srcs) + set_source_files_properties(${opbuild_out_srcs} PROPERTIES GENERATED TRUE) + add_library(opbuild_gen_aclnn_all OBJECT ${opbuild_out_srcs}) + add_dependencies( + opbuild_gen_aclnn_all + opbuild_custom_gen_aclnn_all + ) + target_include_directories(opbuild_gen_aclnn_all + PRIVATE + ${OPAPI_INCLUDE} + ) + endif() +endfunction() + +function(merge_graph_headers) + set(oneValueArgs TARGET OUT_DIR) + cmake_parse_arguments(MGPROTO "" "${oneValueArgs}" "" ${ARGN}) + get_target_property(proto_headers ${GRAPH_PLUGIN_NAME}_proto_headers INTERFACE_SOURCES) + + add_custom_command(OUTPUT ${MGPROTO_OUT_DIR}/ops_proto_transformer.h + COMMAND ${ASCEND_PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/scripts/util/merge_proto.py + ${proto_headers} + --output-file ${MGPROTO_OUT_DIR}/ops_proto_transformer.h + ) + + add_custom_command( + OUTPUT ${MGPROTO_OUT_DIR}/ops_proto_transformer.cpp + COMMAND ${CMAKE_COMMAND} -E copy + ${MGPROTO_OUT_DIR}/ops_proto_transformer.h + ${MGPROTO_OUT_DIR}/ops_proto_transformer.cpp + DEPENDS ${MGPROTO_OUT_DIR}/ops_proto_transformer.h + ) + + add_custom_target(${MGPROTO_TARGET} ALL + DEPENDS ${MGPROTO_OUT_DIR}/ops_proto_transformer.h ${MGPROTO_OUT_DIR}/ops_proto_transformer.cpp + ) +endfunction() diff --git a/csrc/ascend/cmake/package.cmake b/csrc/ascend/cmake/package.cmake new file mode 100644 index 000000000..be869b011 --- /dev/null +++ b/csrc/ascend/cmake/package.cmake @@ -0,0 +1,216 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +#### CPACK to package run ##### + +# download makeself package +include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/third_party/makeself-fetch.cmake) + +# mc2_matmul KB install +include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/runtimeKB.cmake) + +function(pack_custom) + message(STATUS "System processor: ${CMAKE_SYSTEM_PROCESSOR}") + if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64") + message(STATUS "Detected architecture: x86_64") + set(ARCH x86_64) + elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64|arm") + message(STATUS "Detected architecture: ARM64") + set(ARCH aarch64) + else () + message(WARNING "Unknown architecture: ${CMAKE_SYSTEM_PROCESSOR}") + endif () + set(PACK_CUSTOM_NAME "cann-ops-transformer-${VENDOR_NAME}-linux-${ARCH}") + npu_op_package(${PACK_CUSTOM_NAME} + TYPE RUN + CONFIG + ENABLE_SOURCE_PACKAGE True + ENABLE_BINARY_PACKAGE True + INSTALL_PATH ${CMAKE_INSTALL_PREFIX}/ + ENABLE_DEFAULT_PACKAGE_NAME_RULE False + ) + + npu_op_package_add(${PACK_CUSTOM_NAME} + LIBRARY + cust_opapi + ) + if (TARGET cust_proto) + npu_op_package_add(${PACK_CUSTOM_NAME} + LIBRARY + cust_proto + ) + endif() + if (TARGET cust_opmaster) + npu_op_package_add(${PACK_CUSTOM_NAME} + LIBRARY + cust_opmaster + ) + endif() +endfunction() + +function(pack_tiling_sink) + ExternalProject_Get_Property(tiling_sink_task BINARY_DIR) + + if(ENABLE_BUILT_IN) + set(TRANSFORMER_OPMASTER_SO ${BINARY_DIR}/libtiling_device_transformer.so) + set(INSTALL_DIR "ops_transformer/built-in/op_impl/ai_core/tbe/op_tiling_device/lib") + else() + set(TRANSFORMER_OPMASTER_SO ${BINARY_DIR}/libcust_opmaster.so) + set(INSTALL_DIR "packages/vendors/${VENDOR_NAME}_transformer/op_impl/ai_core/tbe/op_master_device/lib") + endif() + install(CODE " + if(EXISTS \"${TRANSFORMER_OPMASTER_SO}\") + file( + INSTALL DESTINATION \"\${CMAKE_INSTALL_PREFIX}/${INSTALL_DIR}\" + TYPE FILE FILES \"${TRANSFORMER_OPMASTER_SO}\") + endif() + ") +endfunction() + +function(pack_built_in) + #### built-in package #### + message(STATUS "System processor: ${CMAKE_SYSTEM_PROCESSOR}") + if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64") + message(STATUS "Detected architecture: x86_64") + set(ARCH x86_64) + elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64|arm") + message(STATUS "Detected architecture: ARM64") + set(ARCH aarch64) + else () + message(WARNING "Unknown architecture: ${CMAKE_SYSTEM_PROCESSOR}") + endif () + + set(script_prefix ${CMAKE_SOURCE_DIR}/scripts/package/ops_transformer/scripts) + install(DIRECTORY ${script_prefix}/ + DESTINATION share/info/ops_transformer/script + FILE_PERMISSIONS + OWNER_READ OWNER_WRITE OWNER_EXECUTE # 文件权限 + GROUP_READ GROUP_EXECUTE + WORLD_READ WORLD_EXECUTE + DIRECTORY_PERMISSIONS + OWNER_READ OWNER_WRITE OWNER_EXECUTE # 目录权限 + GROUP_READ GROUP_EXECUTE + WORLD_READ WORLD_EXECUTE + REGEX "(setenv|prereq_check)\\.(bash|fish|csh)" EXCLUDE + ) + + set(SCRIPTS_FILES + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/check_version_required.awk + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/common_func.inc + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/common_interface.sh + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/common_interface.csh + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/common_interface.fish + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/version_compatiable.inc + ${CMAKE_SOURCE_DIR}/scripts/package/common/py/merge_binary_info_config.py + ) + + install(FILES ${SCRIPTS_FILES} + DESTINATION share/info/ops_transformer/script + ) + set(COMMON_FILES + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/install_common_parser.sh + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/common_func_v2.inc + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/common_installer.inc + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/script_operator.inc + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/version_cfg.inc + ) + + set(PACKAGE_FILES + ${COMMON_FILES} + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/multi_version.inc + ) + set(LATEST_MANGER_FILES + ${COMMON_FILES} + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/common_func.inc + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/version_compatiable.inc + ${CMAKE_SOURCE_DIR}/scripts/package/common/sh/check_version_required.awk + ) + set(CONF_FILES + ${CMAKE_SOURCE_DIR}/scripts/package/common/cfg/path.cfg + ) + install(FILES ${CMAKE_SOURCE_DIR}/version.info + DESTINATION share/info/ops_transformer + ) + install(FILES ${CONF_FILES} + DESTINATION ops_transformer/conf + ) + install(FILES ${PACKAGE_FILES} + DESTINATION share/info/ops_transformer/script + ) + install(FILES ${LATEST_MANGER_FILES} + DESTINATION latest_manager + ) + install(DIRECTORY ${CMAKE_SOURCE_DIR}/scripts/package/latest_manager/scripts/ + DESTINATION latest_manager + ) + + string(FIND "${ASCEND_COMPUTE_UNIT}" ";" SEMICOLON_INDEX) + if (SEMICOLON_INDEX GREATER -1) + # 截取分号前的字串 + math(EXPR SUBSTRING_LENGTH "${SEMICOLON_INDEX}") + string(SUBSTRING "${ASCEND_COMPUTE_UNIT}" 0 "${SUBSTRING_LENGTH}" compute_unit) + else() + # 没有分号取全部内容 + set(compute_unit "${ASCEND_COMPUTE_UNIT}") + endif() + + message(STATUS "current compute_unit is: ${compute_unit}") + pack_tiling_sink() + + # ============= CPack ============= + set(CPACK_PACKAGE_NAME "${PROJECT_NAME}") + set(CPACK_PACKAGE_VERSION "${PROJECT_VERSION}") + string(REGEX REPLACE "^.*[Aa]scend" "" soc_version_temp "${ASCEND_COMPUTE_UNIT}") + # 检查是否成功提取 + if("${soc_version_temp}" STREQUAL "${ASCEND_COMPUTE_UNIT}") + set(soc_version "unknown") + else() + set(soc_version "${soc_version_temp}") + endif() + + if(NOT ENABLE_OPS_KERNEL) + set(CPACK_PACKAGE_FILE_NAME "CANN--${CPACK_PACKAGE_NAME}.run") + else() + if("${VERSION}" STREQUAL "") + if("${soc_version}" STREQUAL "910_93") + set(CPACK_PACKAGE_FILE_NAME "cann-A3-ops-transformer_linux-${ARCH}.run") + else() + set(CPACK_PACKAGE_FILE_NAME "cann-${soc_version}-ops-transformer_linux-${ARCH}.run") + endif() + else() + if("${soc_version}" STREQUAL "910_93") + set(CPACK_PACKAGE_FILE_NAME "cann-A3-ops-transformer_${VERSION}_linux-${ARCH}.run") + else() + set(CPACK_PACKAGE_FILE_NAME "cann-${soc_version}-ops-transformer_${VERSION}_linux-${ARCH}.run") + endif() + endif() + + endif() + + set(CPACK_INSTALL_PREFIX "/") + + set(CPACK_CMAKE_SOURCE_DIR "${CMAKE_SOURCE_DIR}") + set(CPACK_CMAKE_BINARY_DIR "${CMAKE_BINARY_DIR}") + set(CPACK_CMAKE_INSTALL_PREFIX "${CMAKE_INSTALL_PREFIX}") + set(CPACK_CMAKE_CURRENT_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}") + set(CPACK_MAKESELF_PATH "${MAKESELF_PATH}") + set(CPACK_SOC "${compute_unit}") + set(CPACK_ARCH "${ARCH}") + set(CPACK_SET_DESTDIR ON) + set(CPACK_VERSION "${VERSION}") + set(CPACK_GENERATOR External) + if (ENABLE_BUILT_IN) + set(CPACK_EXTERNAL_PACKAGE_SCRIPT "${CMAKE_SOURCE_DIR}/cmake/makeself_built_in.cmake") + endif() + set(CPACK_EXTERNAL_ENABLE_STAGING true) + set(CPACK_PACKAGE_DIRECTORY "${CMAKE_INSTALL_PREFIX}") + + message(STATUS "CMAKE_INSTALL_PREFIX = ${CMAKE_INSTALL_PREFIX}") + include(CPack) +endfunction() \ No newline at end of file diff --git a/csrc/ascend/cmake/rty_obj_func.cmake b/csrc/ascend/cmake/rty_obj_func.cmake new file mode 100644 index 000000000..4c39647ec --- /dev/null +++ b/csrc/ascend/cmake/rty_obj_func.cmake @@ -0,0 +1,592 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +# usage: add_modules_sources(DIR OPTYPE ACLNNTYPE) +# ACLNNTYPE 支持类型aclnn/aclnn_inner/aclnn_exclude +# OPTYPE 和 ACLNNTYPE 需一一对应 + +# 添加infer object +function(add_infer_modules) + if (NOT TARGET ${OPHOST_NAME}_infer_obj) + add_library(${OPHOST_NAME}_infer_obj OBJECT) + target_include_directories(${OPHOST_NAME}_infer_obj + PRIVATE ${OP_PROTO_INCLUDE} + ) + target_compile_definitions(${OPHOST_NAME}_infer_obj + PRIVATE + OPS_UTILS_LOG_SUB_MOD_NAME="OP_PROTO" + $<$:ASCEND_OPSPROTO_UT> + LOG_CPP + ) + target_compile_options(${OPHOST_NAME}_infer_obj + PRIVATE + $<$>:-DDISABLE_COMPILE_V1> + -Dgoogle=ascend_private + -fvisibility=hidden + ) + target_link_libraries(${OPHOST_NAME}_infer_obj + PRIVATE + $,intf_llt_pub_asan_cxx17,intf_pub_cxx17>> + $ + $<$:$> + $<$:$> + + $<$:$> + $<$:$> + tiling_api + ) + endif() +endfunction() + +# 添加tiling object +function(add_tiling_modules) + if (NOT TARGET ${OPHOST_NAME}_tiling_obj) + add_library(${OPHOST_NAME}_tiling_obj OBJECT) + target_include_directories(${OPHOST_NAME}_tiling_obj + PRIVATE ${OP_TILING_INCLUDE} + ) + target_compile_definitions(${OPHOST_NAME}_tiling_obj + PRIVATE + OPS_UTILS_LOG_SUB_MOD_NAME="OP_TILING" + $<$:ASCEND_OPTILING_UT> + LOG_CPP + ) + target_compile_options(${OPHOST_NAME}_tiling_obj + PRIVATE + $<$>:-DDISABLE_COMPILE_V1> + -Dgoogle=ascend_private + -fvisibility=hidden + -fno-strict-aliasing + ) + target_link_libraries(${OPHOST_NAME}_tiling_obj + PRIVATE + $,intf_llt_pub_asan_cxx17,intf_pub_cxx17>> + $ + $<$:$> + $<$:$> + $<$:$> + tiling_api + ) + endif() +endfunction() + +# 添加opapi object +function(add_opapi_modules) + if (NOT TARGET ${OPHOST_NAME}_opapi_obj) + add_library(${OPHOST_NAME}_opapi_obj OBJECT) + target_include_directories(${OPHOST_NAME}_opapi_obj + PRIVATE + ${OPAPI_INCLUDE} + ) + target_compile_options(${OPHOST_NAME}_opapi_obj + PRIVATE + -Dgoogle=ascend_private + -DACLNN_LOG_FMT_CHECK + ) + target_compile_definitions(${OPHOST_NAME}_opapi_obj + PRIVATE + LOG_CPP + ) + target_link_libraries(${OPHOST_NAME}_opapi_obj + PUBLIC + $,intf_llt_pub_asan_cxx17,intf_pub_cxx17>> + PRIVATE + $ + $) + endif() +endfunction() + +# 添加gentask object +function(add_opmaster_ct_gentask_modules) + message(STATUS "add_opmaster_ct_gentask_modules start") + if (NOT TARGET ${OPHOST_NAME}_opmaster_ct_gentask_obj) + add_library(${OPHOST_NAME}_opmaster_ct_gentask_obj OBJECT) + + #如果protobuf还没生成的话，要生成.h + if(NOT TARGET ops_proto_gen) + set(_op_proto_utils_protolist + "${TOP_DIR}/metadef/proto/task.proto" + "${TOP_DIR}/metadef/proto/ge_ir.proto" + ) + protobuf_generate(ops_proto_gen _proto_cc _proto_h ${_op_proto_utils_protolist} TARGET) + message("task.pb.h generate location: ${_proto_h}") + endif() + add_dependencies(${OPHOST_NAME}_opmaster_ct_gentask_obj ops_proto_gen) + + list(GET _proto_h 0 first_proto_header) + get_filename_component(proto_gen_dir "${first_proto_header}" DIRECTORY) + get_filename_component(task_pb_dir "${proto_gen_dir}" DIRECTORY) + + target_include_directories(${OPHOST_NAME}_opmaster_ct_gentask_obj + PRIVATE ${OP_TILING_INCLUDE} + ${task_pb_dir} + ) + target_compile_definitions(${OPHOST_NAME}_opmaster_ct_gentask_obj + PRIVATE + OP_TILING_LIB + ) + target_compile_options(${OPHOST_NAME}_opmaster_ct_gentask_obj + PRIVATE + $<$>:-DDISABLE_COMPILE_V1> + -Dgoogle=ascend_private + -fvisibility=hidden + -fno-strict-aliasing + ) + set(_op_master_ct_gen_task_link_libs + -Wl,--no-as-needed + graph + graph_base + exe_graph + platform + register + alog + error_manager + ops_utils_tiling + -Wl,--as-needed + c_sec + json + platform + mmpa + ascend_protobuf + ) + target_link_libraries(${OPHOST_NAME}_opmaster_ct_gentask_obj + PRIVATE + $ + $<$:$> + $<$>:$> + ${_op_master_ct_gen_task_link_libs} + ) + endif() +endfunction() + +# usage: add_aicpu_kernel_modules() +# 添加aicpu kernel object +function(add_aicpu_kernel_modules) + message(STATUS "add_aicpu_kernel_modules") + if(NOT TARGET ${OPHOST_NAME}_aicpu_obj) + add_library(${OPHOST_NAME}_aicpu_obj OBJECT) + target_include_directories(${OPHOST_NAME}_aicpu_obj PRIVATE ${AICPU_INCLUDE}) + target_compile_definitions( + ${OPHOST_NAME}_aicpu_obj PRIVATE _FORTIFY_SOURCE=2 google=ascend_private + $<$:ASCEND_AICPU_UT> + ) + target_compile_options( + ${OPHOST_NAME}_aicpu_obj PRIVATE $<$>:-DDISABLE_COMPILE_V1> -Dgoogle=ascend_private + -fvisibility=hidden ${AICPU_DEFINITIONS} + ) + target_link_libraries( + ${OPHOST_NAME}_aicpu_obj + PRIVATE $,intf_llt_pub_asan_cxx17,intf_pub_cxx17>> + $ + ) + endif() +endfunction() + +# usage: add_aicpu_cust_kernel_modules(target_name) +# 添加aicpu cust kernel object target +function(add_aicpu_cust_kernel_modules target_name) + message(STATUS "add_aicpu_cust_kernel_modules for ${target_name}") + if(NOT TARGET ${target_name}) + add_library(${target_name} OBJECT) + target_include_directories(${target_name} PRIVATE ${AICPU_INCLUDE}) + target_compile_definitions( + ${target_name} PRIVATE + _FORTIFY_SOURCE=2 _GLIBCXX_USE_CXX11_ABI=1 + google=ascend_private + $<$:ASCEND_AICPU_UT> + ) + target_compile_options( + ${target_name} PRIVATE + $<$>:-DDISABLE_COMPILE_V1> -Dgoogle=ascend_private + -fvisibility=hidden ${AICPU_DEFINITIONS} + ) + target_link_libraries( + ${target_name} + PRIVATE $,intf_llt_pub_asan_cxx17,intf_pub_cxx17>> + $ + -Wl,--no-whole-archive + Eigen3::EigenCv + ) + if (NOT ${target_name} IN_LIST AICPU_CUST_OBJ_TARGETS) + set(AICPU_CUST_OBJ_TARGETS ${AICPU_CUST_OBJ_TARGETS} ${target_name} CACHE INTERNAL "All aicpu cust obj targets") + endif() + endif() +endfunction() + +# usage: add_modules_sources(DIR OPTYPE ACLNNTYPE) +# ACLNNTYPE 支持类型aclnn/aclnn_inner/aclnn_exclude +# OPTYPE 和 ACLNNTYPE 需一一对应 +macro(add_modules_sources) + set(oneValueArgs OP_API_INDEPENDENT OP_API_DIR) + set(multiValueArgs OPTYPE ACLNNTYPE) + + cmake_parse_arguments(MODULE "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + set(SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) + + # 该段代码作用为兼容op_api新旧目录结构(旧： 嵌套于op_host下； 新： 与op_host同级) + if (NOT DEFINED MODULE_OP_API_INDEPENDENT) + set(MODULE_OP_API_INDEPENDENT OFF) + endif() + if(MODULE_OP_API_INDEPENDENT) + # 新结构：op_api与op_host同级，需要指定有效路径 + if (NOT DEFINED MODULE_OP_API_DIR OR NOT EXISTS "${MODULE_OP_API_DIR}") + message(FATAL_ERROR "OP_API_INDEPENDENT=ON时，必须传递有效的OP_API_DIR路径") + endif() + set(OP_API_SRC_DIR "${MODULE_OP_API_DIR}") + else() + # 旧结构：op_api嵌套在op_host目录下 + set(OP_API_SRC_DIR "${SOURCE_DIR}/op_api") + endif() + + # opapi 默认全部编译 + file(GLOB OPAPI_SRCS ${OP_API_SRC_DIR}/*.cpp) + if (OPAPI_SRCS) + add_opapi_modules() + target_sources(${OPHOST_NAME}_opapi_obj PRIVATE ${OPAPI_SRCS}) + endif() + + file(GLOB OPAPI_HEADERS ${OP_API_SRC_DIR}/aclnn_*.h) + if (OPAPI_HEADERS) + target_sources(${OPHOST_NAME}_aclnn_exclude_headers INTERFACE ${OPAPI_HEADERS}) + endif() + + # 获取算子层级目录名称，判断是否编译该算子 + get_filename_component(PARENT_DIR ${SOURCE_DIR} DIRECTORY) + get_filename_component(OP_NAME ${PARENT_DIR} NAME) + list(FIND ASCEND_OP_NAME ${OP_NAME} INDEX) + if(NOT "${ASCEND_OP_NAME}" STREQUAL "" AND INDEX EQUAL -1) + #ASCEND_OP_NAME 为空表示全部编译 + return() + endif() + # 记录全局的COMPILED_OPS和COMPILED_OP_DIRS，其中COMPILED_OP_DIRS只记录到算子名，例如transformer/abs + set(COMPILED_OPS ${COMPILED_OPS} ${OP_NAME} CACHE STRING "Compiled Ops" FORCE) + set(COMPILED_OP_DIRS ${COMPILED_OP_DIRS} ${PARENT_DIR} CACHE STRING "Compiled Ops Dirs" FORCE) + + file(GLOB OPINFER_SRCS ${SOURCE_DIR}/*_infershape*.cpp) + if (OPINFER_SRCS) + add_infer_modules() + target_sources(${OPHOST_NAME}_infer_obj PRIVATE ${OPINFER_SRCS}) + endif() + + file(GLOB_RECURSE SUB_OPTILING_SRC ${SOURCE_DIR}/op_tiling/*.cpp) + file(GLOB OPTILING_SRCS + ${SOURCE_DIR}/*_tiling*.cpp + ${SOURCE_DIR}/*fallback*.cpp + ${SOURCE_DIR}/op_tiling/arch35/*.cpp + ${SOURCE_DIR}/../graph_plugin/fallback_*.cpp + ) + if (OPTILING_SRCS OR SUB_OPTILING_SRC) + add_tiling_modules() + target_sources(${OPHOST_NAME}_tiling_obj PRIVATE ${OPTILING_SRCS} ${SUB_OPTILING_SRC}) + # target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE ${SOURCE_DIR}/../../ ${SOURCE_DIR}) + endif() + + file(GLOB AICPU_SRCS ${MODULE_DIR}/*_aicpu*.cpp) + if (AICPU_SRCS) + add_aicpu_kernel_modules() + target_sources(${OPHOST_NAME}_aicpu_obj PRIVATE ${AICPU_SRCS}) + endif() + + if (MODULE_OPTYPE) + list(LENGTH MODULE_OPTYPE OpTypeLen) + list(LENGTH MODULE_ACLNNTYPE AclnnTypeLen) + if(NOT ${OpTypeLen} EQUAL ${AclnnTypeLen}) + message(FATAL_ERROR "OPTYPE AND ACLNNTYPE Should be One-to-One") + endif() + math(EXPR index "${OpTypeLen} - 1") + foreach(i RANGE ${index}) + list(GET MODULE_OPTYPE ${i} OpType) + list(GET MODULE_ACLNNTYPE ${i} AclnnType) + if (${AclnnType} STREQUAL "aclnn" OR ${AclnnType} STREQUAL "aclnn_inner" OR ${AclnnType} STREQUAL "aclnn_exclude") + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/${OpType}_def*.cpp) + if (OPDEF_SRCS) + target_sources(${OPHOST_NAME}_opdef_${AclnnType}_obj INTERFACE ${OPDEF_SRCS}) + endif() + elseif(${AclnnType} STREQUAL "no_need_aclnn") + message(STATUS "aicpu or host aicpu no need aclnn.") + else() + message(FATAL_ERROR "ACLNN TYPE UNSUPPORTED, ONLY SUPPORT aclnn/aclnn_inner/aclnn_exclude") + endif() + endforeach() + else() + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/*_def*.cpp) + if(OPDEF_SRCS) + message(FATAL_ERROR + "Should Manually specify aclnn/aclnn_inner/aclnn_exclude\n" + "usage: add_modules_sources(OPTYPE optypes ACLNNTYPE aclnntypes)\n" + "example: add_modules_sources(OPTYPE add ACLNNTYPE aclnn_exclude)" + ) + endif() + endif() +endmacro() + +macro(add_modules_sources_with_soc) + set(oneValueArgs OP_API_INDEPENDENT OP_API_DIR) + set(multiValueArgs OPTYPE ACLNNTYPE) + + cmake_parse_arguments(MODULE "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + set(SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) + + # 该段代码作用为兼容op_api新旧目录结构(旧： 嵌套于op_host下； 新： 与op_host同级) + if (NOT DEFINED MODULE_OP_API_INDEPENDENT) + set(MODULE_OP_API_INDEPENDENT OFF) + endif() + if(MODULE_OP_API_INDEPENDENT) + # 新结构：op_api与op_host同级，需要指定有效路径 + if (NOT DEFINED MODULE_OP_API_DIR OR NOT EXISTS "${MODULE_OP_API_DIR}") + message(FATAL_ERROR "OP_API_INDEPENDENT=ON时，必须传递有效的OP_API_DIR路径") + endif() + set(OP_API_SRC_DIR "${MODULE_OP_API_DIR}") + else() + # 旧结构：op_api嵌套在op_host目录下 + set(OP_API_SRC_DIR "${SOURCE_DIR}/op_api") + endif() + + # opapi 默认全部编译 + file(GLOB OPAPI_SRCS ${OP_API_SRC_DIR}/*.cpp) + if (OPAPI_SRCS) + add_opapi_modules() + target_sources(${OPHOST_NAME}_opapi_obj PRIVATE ${OPAPI_SRCS}) + endif() + + file(GLOB OPAPI_HEADERS ${OP_API_SRC_DIR}/aclnn_*.h) + if (OPAPI_HEADERS) + target_sources(${OPHOST_NAME}_aclnn_exclude_headers INTERFACE ${OPAPI_HEADERS}) + endif() + + # 获取算子层级目录名称，判断是否编译该算子 + get_filename_component(PARENT_DIR ${SOURCE_DIR} DIRECTORY) + get_filename_component(OP_NAME ${PARENT_DIR} NAME) + list(FIND ASCEND_OP_NAME ${OP_NAME} INDEX) + if(NOT "${ASCEND_OP_NAME}" STREQUAL "" AND INDEX EQUAL -1) + #ASCEND_OP_NAME 为空表示全部编译 + return() + endif() + # 记录全局的COMPILED_OPS和COMPILED_OP_DIRS，其中COMPILED_OP_DIRS只记录到算子名，例如transformer/abs + set(COMPILED_OPS ${COMPILED_OPS} ${OP_NAME} CACHE STRING "Compiled Ops" FORCE) + set(COMPILED_OP_DIRS ${COMPILED_OP_DIRS} ${PARENT_DIR} CACHE STRING "Compiled Ops Dirs" FORCE) + + file(GLOB OPINFER_SRCS ${SOURCE_DIR}/*_infershape*.cpp) + if (OPINFER_SRCS) + add_infer_modules() + target_sources(${OPHOST_NAME}_infer_obj PRIVATE ${OPINFER_SRCS}) + endif() + + file(GLOB_RECURSE SUB_OPTILING_SRC ${SOURCE_DIR}/*_tiling*.cpp) + file(GLOB OPTILING_SRCS + ${SOURCE_DIR}/*_tiling*.cpp + ${SOURCE_DIR}/*fallback*.cpp + ${SOURCE_DIR}/../op_graph/fallback_*.cpp + ${SOURCE_DIR}/../graph_plugin/fallback_*.cpp + ) + if (OPTILING_SRCS OR SUB_OPTILING_SRC) + add_tiling_modules() + target_sources(${OPHOST_NAME}_tiling_obj PRIVATE ${OPTILING_SRCS} ${SUB_OPTILING_SRC}) + # target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE ${SOURCE_DIR}/../../ ${SOURCE_DIR}) + endif() + + file(GLOB AICPU_SRCS ${MODULE_DIR}/*_aicpu*.cpp) + if (AICPU_SRCS) + add_aicpu_kernel_modules() + target_sources(${OPHOST_NAME}_aicpu_obj PRIVATE ${AICPU_SRCS}) + endif() + + if (MODULE_OPTYPE) + list(LENGTH MODULE_OPTYPE OpTypeLen) + list(LENGTH MODULE_ACLNNTYPE AclnnTypeLen) + if(NOT ${OpTypeLen} EQUAL ${AclnnTypeLen}) + message(FATAL_ERROR "OPTYPE AND ACLNNTYPE Should be One-to-One") + endif() + math(EXPR index "${OpTypeLen} - 1") + foreach(i RANGE ${index}) + list(GET MODULE_OPTYPE ${i} OpType) + list(GET MODULE_ACLNNTYPE ${i} AclnnType) + if (${AclnnType} STREQUAL "aclnn" OR ${AclnnType} STREQUAL "aclnn_inner" OR ${AclnnType} STREQUAL "aclnn_exclude") + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/${OpType}_def*.cpp) + if (OPDEF_SRCS) + target_sources(${OPHOST_NAME}_opdef_${AclnnType}_obj INTERFACE ${OPDEF_SRCS}) + endif() + elseif(${AclnnType} STREQUAL "no_need_aclnn") + message(STATUS "aicpu or host aicpu no need aclnn.") + else() + message(FATAL_ERROR "ACLNN TYPE UNSUPPORTED, ONLY SUPPORT aclnn/aclnn_inner/aclnn_exclude") + endif() + endforeach() + else() + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/*_def*.cpp) + if(OPDEF_SRCS) + message(FATAL_ERROR + "Should Manually specify aclnn/aclnn_inner/aclnn_exclude\n" + "usage: add_modules_sources(OPTYPE optypes ACLNNTYPE aclnntypes)\n" + "example: add_modules_sources(OPTYPE add ACLNNTYPE aclnn_exclude)" + ) + endif() + endif() +endmacro() + +# mc2算子回黄编译框架 +macro(add_mc2_modules_sources) + set(multiValueArgs OPTYPE ACLNNTYPE) + + cmake_parse_arguments(MODULE "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + set(SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) + + # 获取父目录和祖父目录路径 + get_filename_component(CMAKE_PARENT_DIR ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) + get_filename_component(CMAKE_GRANDPARENT_DIR ${CMAKE_PARENT_DIR} DIRECTORY) + get_filename_component(ASCEND_PARENT_DIR ${ASCEND_CANN_PACKAGE_PATH} DIRECTORY) + + # opapi 默认全部编译 + file(GLOB OPAPI_SRCS ${SOURCE_DIR}/../op_api/*.cpp) + if (OPAPI_SRCS) + add_opapi_modules() + target_sources(${OPHOST_NAME}_opapi_obj PRIVATE ${OPAPI_SRCS}) + endif() + + # file(GLOB OPAPI_HEADERS ${SOURCE_DIR}/op_api/aclnn_*.h) + # if (OPAPI_HEADERS) + # target_sources(${OPHOST_NAME}_aclnn_inner_headers INTERFACE ${OPAPI_HEADERS}) + # endif() + + # 获取算子层级目录名称，判断是否编译该算子 + get_filename_component(PARENT_DIR ${SOURCE_DIR} DIRECTORY) + get_filename_component(OP_NAME ${PARENT_DIR} NAME) + list(FIND ASCEND_OP_NAME ${OP_NAME} INDEX) + if(NOT "${ASCEND_OP_NAME}" STREQUAL "" AND INDEX EQUAL -1) + #ASCEND_OP_NAME 为空表示全部编译 + return() + endif() + # 记录全局的COMPILED_OPS和COMPILED_OP_DIRS，其中COMPILED_OP_DIRS只记录到算子名，例如transformer/abs + set(COMPILED_OPS ${COMPILED_OPS} ${OP_NAME} CACHE STRING "Compiled Ops" FORCE) + set(COMPILED_OP_DIRS ${COMPILED_OP_DIRS} ${PARENT_DIR} CACHE STRING "Compiled Ops Dirs" FORCE) + + file(GLOB OPINFER_SRCS ${SOURCE_DIR}/*_infershape*.cpp) + if (OPINFER_SRCS) + add_infer_modules() + target_sources(${OPHOST_NAME}_infer_obj PRIVATE ${OPINFER_SRCS}) + endif() + + file(GLOB_RECURSE OPTILING_SRCS + ${SOURCE_DIR}/op_tiling/*.cpp + ${SOURCE_DIR}/../op_graph/fallback_*.cpp + ${SOURCE_DIR}/../graph_plugin/fallback_*.cpp + ${SOURCE_DIR}/../../common/src/matmul_formulaic_tiling.cpp + ${SOURCE_DIR}/../../common/src/mc2_hcom_topo_info.cpp + ${SOURCE_DIR}/../../common/src/mc2_matmul_tiling_cfg.cpp) + if (OPTILING_SRCS) + add_tiling_modules() + target_sources(${OPHOST_NAME}_tiling_obj PRIVATE ${OPTILING_SRCS}) + endif() + + file(GLOB AICPU_SRCS ${MODULE_DIR}/*_aicpu*.cpp) + if (AICPU_SRCS) + add_aicpu_kernel_modules() + target_sources(${OPHOST_NAME}_aicpu_obj PRIVATE ${AICPU_SRCS}) + endif() + + file(GLOB GENTASK_SRCS + ${SOURCE_DIR}/../op_graph/*_gen_task*.cpp + ${SOURCE_DIR}/../../common/src/mc2_a5_gen_task_utils.cpp + ) + if(GENTASK_SRCS) + add_opmaster_ct_gentask_modules() + target_sources(${OPHOST_NAME}_opmaster_ct_gentask_obj PRIVATE ${GENTASK_SRCS}) + endif() + + if (MODULE_OPTYPE) + list(LENGTH MODULE_OPTYPE OpTypeLen) + list(LENGTH MODULE_ACLNNTYPE AclnnTypeLen) + if(NOT ${OpTypeLen} EQUAL ${AclnnTypeLen}) + message(FATAL_ERROR "OPTYPE AND ACLNNTYPE Should be One-to-One") + endif() + math(EXPR index "${OpTypeLen} - 1") + foreach(i RANGE ${index}) + list(GET MODULE_OPTYPE ${i} OpType) + list(GET MODULE_ACLNNTYPE ${i} AclnnType) + if (${AclnnType} STREQUAL "aclnn" OR ${AclnnType} STREQUAL "aclnn_inner" OR ${AclnnType} STREQUAL "aclnn_exclude") + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/${OpType}_def*.cpp) + if (OPDEF_SRCS) + target_sources(${OPHOST_NAME}_opdef_${AclnnType}_obj INTERFACE ${OPDEF_SRCS}) + endif() + elseif(${AclnnType} STREQUAL "no_need_aclnn") + message(STATUS "aicpu or host aicpu no need aclnn.") + else() + message(FATAL_ERROR "ACLNN TYPE UNSUPPORTED, ONLY SUPPORT aclnn/aclnn_inner/aclnn_exclude") + endif() + endforeach() + else() + file(GLOB OPDEF_SRCS ${SOURCE_DIR}/*_def*.cpp) + if(OPDEF_SRCS) + message(FATAL_ERROR + "Should Manually specify aclnn/aclnn_inner/aclnn_exclude\n" + "usage: add_modules_sources(OPTYPE optypes ACLNNTYPE aclnntypes)\n" + "example: add_modules_sources(OPTYPE add ACLNNTYPE aclnn_exclude)" + ) + endif() + endif() +endmacro() + +# usage: add_graph_plugin_sources() +macro(add_graph_plugin_sources) + set(SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) + + # 获取算子层级目录名称，判断是否编译该算子 + get_filename_component(PARENT_DIR ${SOURCE_DIR} DIRECTORY) + get_filename_component(OP_NAME ${PARENT_DIR} NAME) + if(DEFINED ASCEND_OP_NAME + AND NOT "${ASCEND_OP_NAME}" STREQUAL "" + AND NOT "${ASCEND_OP_NAME}" STREQUAL "all" + AND NOT "${ASCEND_OP_NAME}" STREQUAL "ALL" + ) + if(NOT ${OP_NAME} IN_LIST ASCEND_OP_NAME) + return() + endif() + endif() + + file(GLOB GRAPH_PLUGIN_SRCS + ${SOURCE_DIR}/*_graph_plugin*.cpp + ) + if(GRAPH_PLUGIN_SRCS) + add_graph_plugin_modules() + target_sources(${GRAPH_PLUGIN_NAME}_obj PRIVATE ${GRAPH_PLUGIN_SRCS}) + endif() + + file(GLOB GRAPH_PLUGIN_PROTO_HEADERS ${SOURCE_DIR}/*_proto*.h) + if(GRAPH_PLUGIN_PROTO_HEADERS) + target_sources(${GRAPH_PLUGIN_NAME}_proto_headers INTERFACE ${GRAPH_PLUGIN_PROTO_HEADERS}) + endif() +endmacro() + +set(_op_tiling_link_libs + -Wl,--no-as-needed + graph + graph_base + exe_graph + platform + register + alog + error_manager + ops_utils_tiling + -Wl,--as-needed + -Wl,--whole-archive + tiling_api + -Wl,--no-whole-archive + c_sec + json + platform + mmpa + ascend_protobuf +) + +if(NOT ${CMAKE_BUILD_MODE} STREQUAL "FALSE") + set(compile_opt_mode ${CMAKE_BUILD_MODE}) + else() + set(compile_opt_mode -O2) +endif() diff --git a/csrc/ascend/cmake/runtimeKB.cmake b/csrc/ascend/cmake/runtimeKB.cmake new file mode 100644 index 000000000..0aac3b906 --- /dev/null +++ b/csrc/ascend/cmake/runtimeKB.cmake @@ -0,0 +1,34 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +#### CPACK runtime kb to package run ##### + +set(OPS_RTKB_DIR ops_transformer/built-in/data/op) +set(KB_CATEGORY_LIST "mc2") +set(RTKB_FILE_ALL) +foreach(OP_RTKB_CATEGORY ${KB_CATEGORY_LIST}) + if (IS_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/${OP_RTKB_CATEGORY}) + if (IS_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/${OP_RTKB_CATEGORY}/3rd) + file(GLOB_RECURSE RTKB_FILE ${CMAKE_CURRENT_SOURCE_DIR}/${OP_RTKB_CATEGORY}/3rd/*/op_host/config/${compute_unit}/*_runtime_kb.json) + list(APPEND RTKB_FILE_ALL ${RTKB_FILE}) + endif() + endif() +endforeach() + +foreach(RTKB_DATA_FILE ${RTKB_FILE_ALL}) + string(FIND "${RTKB_DATA_FILE}" "/" last_slash_pos REVERSE) + string(SUBSTRING "${RTKB_DATA_FILE}" ${last_slash_pos} -1 RTKB_DATA_FILE_NAME) + string(FIND "${RTKB_DATA_FILE_NAME}" "_" first_underscore_pos) + string(SUBSTRING "${RTKB_DATA_FILE_NAME}" 0 ${first_underscore_pos} fileSocPath) + + install(FILES ${RTKB_DATA_FILE} + DESTINATION ${OPS_RTKB_DIR}${fileSocPath}/unified_bank + OPTIONAL + ) +endforeach() \ No newline at end of file diff --git a/csrc/ascend/cmake/scripts/check_version_compatible.py b/csrc/ascend/cmake/scripts/check_version_compatible.py new file mode 100644 index 000000000..0e36f99a5 --- /dev/null +++ b/csrc/ascend/cmake/scripts/check_version_compatible.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +""" +版本兼容性检查 + +检查当前代码仓与基础 CANN 包间的兼容性. +""" + +import argparse +import logging +from pathlib import Path +from typing import NoReturn + + +class VersionChecker: + @classmethod + def main(cls) -> NoReturn: + parser = argparse.ArgumentParser(description="Check Version Compatible", epilog="Best Regards!") + sub_parser = parser.add_subparsers(help="Sub-Command") + # 参数注册 + parser.add_argument("--cann_path", required=True, nargs=1, type=str, help="CANN install path") + parser.add_argument("--cann_package_name", required=True, nargs=1, type=str, help="CANN package name") + # 子命令行(Check) + p_chk = sub_parser.add_parser("check_code_compatible", help="Check Version Compatible.") + p_chk.add_argument( + "--code_version_info_file", required=True, nargs=1, type=str, help="Code version info file path" + ) + p_chk.set_defaults(func=VersionChecker._check_compatible) + # 子命令行(Get) + p_get = sub_parser.add_parser("get_package_version", help="Get Package Version") + p_get.set_defaults(func=VersionChecker._get_package_version) + # 参数处理 + args = parser.parse_args() + # 基本合法性检查, 版本号获取 + cann_version_info_file = Path(args.cann_path[0], args.cann_package_name[0], "version.info").absolute() + if not cann_version_info_file.exists(): + raise ValueError(f"CANN version info file({cann_version_info_file}) not exist.") + ret, cann_version = cls._get_version_str(file=cann_version_info_file) + if not ret: + raise ValueError(f"Can't get version from CANN version info file({cann_version_info_file}).") + rst = args.func(cann_version, args) + return rst + + @classmethod + def _check_compatible(cls, cann_version: str, args) -> str: + code_version_info_file = Path(args.code_version_info_file[0]).absolute() + if not code_version_info_file.exists(): + raise ValueError(f"Code version info file({code_version_info_file}) not exist.") + ret, code_version = cls._get_version_str(file=code_version_info_file) + if not ret: + raise ValueError(f"Can't get version from Code version info file({code_version_info_file}).") + # 兼容性检查 + cann_sub_version = cann_version.rsplit(".", 1)[0] + code_sub_version = code_version.rsplit(".", 1)[0] + if cann_sub_version != code_sub_version: + raise ValueError( + f"The version number of the current code is {code_sub_version}, " + f"and the version number of the cann package used is {cann_sub_version}. " + f"Please install version {code_sub_version} of the cann package." + ) + return cann_sub_version + + @classmethod + def _get_package_version(cls, cann_version: str, args) -> str: + cann_sub_version = cann_version.rsplit(".", 1)[0] + return cann_sub_version + + @classmethod + def _get_version_str(cls, file: Path): + with open(file) as fh: + lines = fh.readlines() + for line in lines: + if not line.startswith("Version="): + continue + version = line[8:].replace("\r", "").replace("\n", "") + return True, version + return False, "" + + +if __name__ == "__main__": + logging.basicConfig(format="%(filename)s:%(lineno)d [%(levelname)s] %(message)s", level=logging.INFO) + try: + print(VersionChecker.main()) + except Exception as e: + logging.error(e) + exit(1) diff --git a/csrc/ascend/cmake/scripts/convert_yaml.py b/csrc/ascend/cmake/scripts/convert_yaml.py new file mode 100644 index 000000000..9494cd979 --- /dev/null +++ b/csrc/ascend/cmake/scripts/convert_yaml.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +""" +test_config.yaml 格式转换 + +转换成 ops-nn 仓的格式, 方便ci读取 +""" + +import logging +import os +import sys + +import yaml + + +def load_test_config(test_config_path: str): + """读取并解析test_config.yaml文件""" + try: + with open(test_config_path, encoding="utf-8") as file: + return yaml.safe_load(file) + + except Exception as e: + logging.error("Failed to read test_config.yaml file: %s", e) + return None + + +def extract_src_and_exclude(data): + """从数据中提取所有有options的算子的src和exclude路径""" + src_paths = set() + exclude_paths = set() + + def extract_from_dict(obj, current_key=None): + if isinstance(obj, dict): + if "src" in obj and isinstance(obj["src"], list): + for path in obj["src"]: + src_paths.add(path) + if "exclude" in obj and isinstance(obj["exclude"], list): + for path in obj["exclude"]: + exclude_paths.add(path) + if "ut_cov_exclude" in obj and isinstance(obj["ut_cov_exclude"], list): + for path in obj["ut_cov_exclude"]: + exclude_paths.add(f'"{path}"' if path.startswith("*") else path) + + # 递归处理所有值 + for key, value in obj.items(): + extract_from_dict(value, key) + elif isinstance(obj, list): + for item in obj: + extract_from_dict(item, current_key) + + extract_from_dict(data) + + return sorted(src_paths), sorted(exclude_paths) + + +def write_new_format(new_file_path: str, src_paths: list, exclude_paths: list): + """以新格式写入文件""" + try: + with open(new_file_path, "w", encoding="utf-8") as file: + file.write("ops-transformer:\n") + file.write(" src:\n") + + file.write(" release:\n") + for path in src_paths: + file.write(f" - {path}\n") + + file.write(" unrelease:\n") + for path in exclude_paths: + file.write(f" - {path}\n") + + return True + + except Exception as e: + logging.error("Failed to write file: %s", e) + return False + + +def main(test_config_path: str, output_path: str): + """主函数""" + # 检查文件是否存在 + if not os.path.exists(test_config_path): + logging.error("File does not exist: %s", test_config_path) + return + + # 读取test_config文件 + data = load_test_config(test_config_path) + if data is None: + return + + # 提取所有有options的算子的src和exclude路径 + src_paths, exclude_paths = extract_src_and_exclude(data) + + logging.info("Found %s src paths", len(src_paths)) + logging.info("Found %s exclude paths", len(exclude_paths)) + + # 以新格式写回 + if write_new_format(output_path, src_paths, exclude_paths): + logging.info("File conversion completed") + else: + logging.error("File conversion failed") + + +if __name__ == "__main__": + logging.basicConfig(format="[%(asctime)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO) + + if len(sys.argv) == 1: + main("test_config.yaml", "test_config.yaml") + elif len(sys.argv) == 2: + main(sys.argv[1], sys.argv[1]) + elif len(sys.argv) == 3: + main(sys.argv[1], sys.argv[2]) + else: + logging.error("usage: convert_yaml.py test_config_path [output_path]") + exit(1) diff --git a/csrc/ascend/cmake/scripts/custom/help.info b/csrc/ascend/cmake/scripts/custom/help.info new file mode 100644 index 000000000..337c286ff --- /dev/null +++ b/csrc/ascend/cmake/scripts/custom/help.info @@ -0,0 +1,2 @@ + --install-path Install operator package to specific dir path + --install-for-all Allow other users to use the operator package \ No newline at end of file diff --git a/csrc/ascend/cmake/scripts/custom/install.sh b/csrc/ascend/cmake/scripts/custom/install.sh new file mode 100644 index 000000000..e1cca7533 --- /dev/null +++ b/csrc/ascend/cmake/scripts/custom/install.sh @@ -0,0 +1,344 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +vendor_name=customize +targetdir=/usr/local/Ascend/opp +target_custom=0 + +sourcedir=$PWD/packages +vendordir=vendors/$vendor_name + +QUIET="y" +INSTALL_FOR_ALL="n" + + +while true +do + case $1 in + --quiet) + QUIET="y" + shift + ;; + --install-path=*) + INSTALL_PATH=$(echo $1 | cut -d"=" -f2-) + INSTALL_PATH=${INSTALL_PATH%*/} + shift + ;; + --install-for-all) + INSTALL_FOR_ALL="y" + shift + ;; + --*) + shift + ;; + *) + break + ;; + esac +done + +log() { + cur_date=`date +"%Y-%m-%d %H:%M:%S"` + echo "[ops_custom] [$cur_date] "$1 +} + +if [ -n "${INSTALL_PATH}" ]; then + if [[ ! "${INSTALL_PATH}" = /* ]]; then + log "[ERROR] use absolute path for --install-path argument" + exit 1 + fi + if [ ! -d ${INSTALL_PATH} ]; then + mkdir ${INSTALL_PATH} >> /dev/null 2>&1 + if [ $? -ne 0 ]; then + log "[ERROR] create ${INSTALL_PATH} failed" + exit 1 + fi + fi + targetdir=${INSTALL_PATH} +elif [ -n "${ASCEND_CUSTOM_OPP_PATH}" ]; then + if [[ "${ASCEND_CUSTOM_OPP_PATH}" == *:* ]]; then + log "[ERROR] environment variable ASCEND_CUSTOM_OPP_PATH=${ASCEND_CUSTOM_OPP_PATH} is set and \ + has multiple path in it (colon inside), which will cause the custom op installed incorrectly. \ + Please use the --install-path option to specify an installation path instead." + exit 1 + fi + if [ ! -d ${ASCEND_CUSTOM_OPP_PATH} ]; then + mkdir -p ${ASCEND_CUSTOM_OPP_PATH} >> /dev/null 2>&1 + if [ $? -ne 0 ]; then + log "[ERROR] create ${ASCEND_CUSTOM_OPP_PATH} failed" + fi + fi + targetdir=${ASCEND_CUSTOM_OPP_PATH} +else + if [ "x${ASCEND_OPP_PATH}" == "x" ]; then + log "[ERROR] env ASCEND_OPP_PATH no exist" + exit 1 + fi + targetdir="${ASCEND_OPP_PATH}" +fi + +if [ ! -d $targetdir ];then + log "[ERROR] $targetdir no exist" + exit 1 +fi + +if [ ! -x $targetdir ] || [ ! -w $targetdir ] || [ ! -r $targetdir ];then + log "[WARNING] The directory $targetdir does not have sufficient permissions. \ + Please check and modify the folder permissions (e.g., using chmod), \ + or use the --install-path option to specify an installation path and \ + change the environment variable ASCEND_CUSTOM_OPP_PATH to the specified path." +fi + +upgrade() +{ + if [ ! -d ${sourcedir}/$vendordir/$1 ]; then + log "[INFO] no need to upgrade ops $1 files" + return 0 + fi + + if [ ! -d ${targetdir}/$vendordir/$1 ];then + log "[INFO] create ${targetdir}/$vendordir/$1." + mkdir -p ${targetdir}/$vendordir/$1 + if [ $? -ne 0 ];then + log "[ERROR] create ${targetdir}/$vendordir/$1 failed" + return 1 + fi + else + has_same_file=-1 + for file_a in ${sourcedir}/$vendordir/$1/*; do + file_b=${file_a##*/}; + if [ "ls ${targetdir}/$vendordir/$1" = "" ]; then + log "[INFO] ${targetdir}/$vendordir/$1 is empty !!" + return 1 + fi + grep -q $file_b <<<`ls ${targetdir}/$vendordir/$1`; + if [[ $? -eq 0 ]]; then + echo -n "${file_b} " + has_same_file=0 + fi + done + if [ 0 -eq $has_same_file ]; then + echo + if test $QUIET = "n"; then + echo "[INFO]: has old version in ${targetdir}/$vendordir/$1, \ + you want to Overlay Installation , please enter:[o]; \ + or replace directory installation , please enter: [r]; \ + or not install , please enter:[n]." + + while true + do + read orn + if [ "$orn" = n ]; then + return 0 + elif [ "$orn" = o ]; then + break; + elif [ "$orn" = r ]; then + [ -d "${targetdir}/$vendordir/$1/" ] && rm -rf "${targetdir}/$vendordir/$1"/* + break; + else + log "[ERROR] input error, please input again!" + fi + done + else + [ -d "${targetdir}/$vendordir/$1/" ] && rm -rf "${targetdir}/$vendordir/$1"/* + fi + fi + log "[INFO] replace or merge old ops $1 files ......" + fi + + log "[INFO] copy new ops $1 files ......" + if [ -d ${targetdir}/$vendordir/$1/ ]; then + chmod -R +w "$targetdir/$vendordir/$1/" >/dev/null 2>&1 + fi + cp -rf ${sourcedir}/$vendordir/$1/* $targetdir/$vendordir/$1/ + if [ $? -ne 0 ];then + log "[ERROR] copy new $1 files failed" + return 1 + fi + + return 0 +} +upgrade_proto() +{ + if [ ! -f ${sourcedir}/$vendordir/custom.proto ]; then + log "[INFO] no need to upgrade custom.proto files" + return 0 + fi + if [ ! -d ${targetdir}/$vendordir/framework/caffe ];then + log "[INFO] create ${targetdir}/$vendordir/framework/caffe." + mkdir -p ${targetdir}/$vendordir/framework/caffe + if [ $? -ne 0 ];then + log "[ERROR] create ${targetdir}/$vendordir/framework/caffe failed" + return 1 + fi + else + if [ -f ${targetdir}/$vendordir/framework/caffe/custom.proto ]; then + # 有老版本,判断是否要覆盖式安装 + if test $QUIET = "n"; then + echo "[INFO] ${targetdir}/$vendordir/framework/caffe has old version"\ + "custom.proto file. Do you want to replace? [y/n] " + + while true + do + read yn + if [ "$yn" = n ]; then + return 0 + elif [ "$yn" = y ]; then + break; + else + log "[ERROR] input error, please input again!" + fi + done + fi + fi + log "[INFO] replace old caffe.proto files ......" + fi + chmod -R +w "$targetdir/$vendordir/framework/caffe/" >/dev/null 2>&1 + cp -rf ${sourcedir}/$vendordir/custom.proto ${targetdir}/$vendordir/framework/caffe/ + if [ $? -ne 0 ];then + log "[ERROR] copy new custom.proto failed" + return 1 + fi + log "[INFO] copy custom.proto success" + + return 0 +} + +upgrade_file() +{ + if [ ! -e ${sourcedir}/$vendordir/$1 ]; then + log "[INFO] no need to upgrade ops $1 file" + return 0 + fi + + log "[INFO] copy new $1 files ......" + cp -f ${sourcedir}/$vendordir/$1 $targetdir/$vendordir/$1 + if [ $? -ne 0 ];then + log "[ERROR] copy new $1 file failed" + return 1 + fi + + return 0 +} + +delete_optiling_file() +{ + if [ ! -d ${targetdir}/vendors ];then + log "[INFO] $1 not exist, no need to uninstall" + return 0 + fi + sys_info=$(uname -m) + if [ ! -d ${sourcedir}/$vendordir/$1/ai_core/tbe/op_tiling/lib/linux/${sys_info} ];then + rm -rf ${sourcedir}/$vendordir/$1/ai_core/tbe/op_tiling/liboptiling.so + fi + return 0 +} + +log "[INFO] copy uninstall sh success" + +if [ ! -d ${targetdir}/vendors ];then + log "[INFO] create ${targetdir}/vendors." + mkdir -p ${targetdir}/vendors + if [ $? -ne 0 ];then + log "[ERROR] create ${targetdir}/vendors failed" + exit 1 + fi +fi +chmod u+w ${targetdir}/vendors + +log "[INFO] upgrade framework" +upgrade framework +if [ $? -ne 0 ];then + exit 1 +fi + +log "[INFO] upgrade op proto" +upgrade op_proto +if [ $? -ne 0 ];then + exit 1 +fi + +log "[INFO] upgrade op impl" +delete_optiling_file op_impl +upgrade op_impl +if [ $? -ne 0 ];then + exit 1 +fi + +log "[INFO] upgrade op api" +upgrade op_api +if [ $? -ne 0 ];then + exit 1 +fi + +log "[INFO] upgrade version.info" +upgrade_file version.info +if [ $? -ne 0 ];then + exit 1 +fi + +upgrade_proto +if [ $? -ne 0 ];then + exit 1 +fi + +# set the set_env.bash +if [ -n "${INSTALL_PATH}" ] && [ -d ${INSTALL_PATH} ]; then + _ASCEND_CUSTOM_OPP_PATH=${targetdir}/${vendordir} + bin_path="${_ASCEND_CUSTOM_OPP_PATH}/bin" + set_env_variable="#!/bin/bash\nexport ASCEND_CUSTOM_OPP_PATH=${_ASCEND_CUSTOM_OPP_PATH}:\${ASCEND_CUSTOM_OPP_PATH}\nexport LD_LIBRARY_PATH=${_ASCEND_CUSTOM_OPP_PATH}/op_api/lib/:\${LD_LIBRARY_PATH}" + if [ ! -d ${bin_path} ]; then + mkdir -p ${bin_path} >> /dev/null 2>&1 + if [ $? -ne 0 ]; then + log "[ERROR] create ${bin_path} failed" + exit 1 + fi + fi + echo -e ${set_env_variable} > ${bin_path}/set_env.bash + if [ $? -ne 0 ]; then + log "[ERROR] write ASCEND_CUSTOM_OPP_PATH to set_env.bash failed" + exit 1 + else + log "[INFO] using requirements: when custom module install finished or before you run the custom module, \ + execute the command [ source ${bin_path}/set_env.bash ] to set the environment path" + fi +else + _ASCEND_CUSTOM_OPP_PATH=${targetdir}/${vendordir} + config_file=${targetdir}/vendors/config.ini + if [ ! -f ${config_file} ]; then + touch ${config_file} + chmod 640 ${config_file} + echo "load_priority=$vendor_name" > ${config_file} + if [ $? -ne 0 ];then + log "[ERROR] echo load_priority failed" + exit 1 + fi + else + found_vendors="$(grep -w "load_priority" "$config_file" | cut --only-delimited -d"=" -f2-)" + found_vendor=$(echo $found_vendors | sed "s/\<$vendor_name\>//g" | tr ',' ' ') + vendor=$(echo $found_vendor | tr -s ' ' ',') + if [ "$vendor" != "" ]; then + sed -i "/load_priority=$found_vendors/s@load_priority=$found_vendors@load_priority=$vendor_name,$vendor@g" "$config_file" + fi + fi + if test $INSTALL_FOR_ALL = "y"; then + chmod 755 ${config_file} + fi + log "[INFO] using requirements: when custom module install finished or before you run the custom module, \ + execute the command [ export LD_LIBRARY_PATH=${_ASCEND_CUSTOM_OPP_PATH}/op_api/lib/:\${LD_LIBRARY_PATH} ] to set the environment path" +fi + +if [ -d ${targetdir}/$vendordir/op_impl/cpu/aicpu_kernel/impl/ ]; then + chmod -R 440 ${targetdir}/$vendordir/op_impl/cpu/aicpu_kernel/impl/* >/dev/null 2>&1 +fi + +echo "SUCCESS" +exit 0 diff --git a/csrc/ascend/cmake/scripts/custom/upgrade.sh b/csrc/ascend/cmake/scripts/custom/upgrade.sh new file mode 100644 index 000000000..282789644 --- /dev/null +++ b/csrc/ascend/cmake/scripts/custom/upgrade.sh @@ -0,0 +1,153 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +vendor_name=customize +targetdir=/usr/local/Ascend/opp +target_custom=0 + +sourcedir=$PWD/packages +vendordir=vendors/$vendor_name + +log() { + cur_date=`date +"%Y-%m-%d %H:%M:%S"` + echo "[ops_custom] [$cur_date] "$1 +} + +if [[ "x${ASCEND_OPP_PATH}" == "x" ]];then + log "[ERROR] env ASCEND_OPP_PATH no exist" + exit 1 +fi + +targetdir=${ASCEND_OPP_PATH} + +if [ ! -d $targetdir ];then + log "[ERROR] $targetdir no exist" + exit 1 +fi + +if [ ! -x $targetdir ] || [ ! -w $targetdir ] || [ ! -r $targetdir ];then + log "[WARNING] The directory $targetdir does not have sufficient permissions. \ + Please check and modify the folder permissions (e.g., using chmod), \ + or use the --install-path option to specify an installation path and \ + change the environment variable ASCEND_CUSTOM_OPP_PATH to the specified path." +fi + +upgrade() +{ + if [ ! -d ${sourcedir}/$vendordir/$1 ]; then + log "[INFO] no need to upgrade ops $1 files" + return 0 + fi + + if [ ! -d ${targetdir}/$vendordir/$1 ];then + log "[INFO] create ${targetdir}/$vendordir/$1." + mkdir -p ${targetdir}/$vendordir/$1 + if [ $? -ne 0 ];then + log "[ERROR] create ${targetdir}/$vendordir/$1 failed" + return 1 + fi + else + vendor_installed_dir=$(ls "$targetdir/vendors" 2> /dev/null) + for i in $vendor_installed_dir;do + vendor_installed_file=$(ls "$vendor_installed_dir/$vendor_name/$i" 2> /dev/null) + if [ "$i" = "$vendor_name" ] && [ "$vendor_installed_file" != "" ]; then + echo "[INFO]: $vendor_name custom opp package has been installed on the path $vendor_installed_dir, \ + you want to Overlay Installation , please enter:[o]; \ + or replace directory installation , please enter: [r]; \ + or not install , please enter:[n]." + fi + while true + do + read mrn + if [ "$mrn" = o ]; then + break + elif [ "$mrn" = r ]; then + [ -n "$vendor_installed_file" ] && rm -rf "$vendor_installed_file" + break + elif [ "$mrn" = n ]; then + return 0 + else + log "[WARNING]: Input error, please input m or r or n to choose!" + fi + done + done + log "[INFO] replace old ops $1 files ......" + fi + + log "copy new ops $1 files ......" + cp -rf ${sourcedir}/$vendordir/$1/* $targetdir/$vendordir/$1/ + if [ $? -ne 0 ];then + log "[ERROR] copy new $1 files failed" + return 1 + fi + + return 0 +} + +upgrade_file() +{ + if [ ! -e ${sourcedir}/$vendordir/$1 ]; then + log "[INFO] no need to upgrade ops $1 file" + return 0 + fi + + log "copy new $1 files ......" + cp -f ${sourcedir}/$vendordir/$1 $targetdir/$vendordir/$1 + if [ $? -ne 0 ];then + log "[ERROR] copy new $1 file failed" + return 1 + fi + + return 0 +} + +log "[INFO] copy uninstall sh success" + +log "[INFO] upgrade framework" +upgrade framework +if [ $? -ne 0 ];then + exit 1 +fi + +log "[INFO] upgrade op proto" +upgrade op_proto +if [ $? -ne 0 ];then + exit 1 +fi + +log "[INFO] upgrade op impl" +upgrade op_impl +if [ $? -ne 0 ];then + exit 1 +fi + +log "[INFO] upgrade op api" +upgrade op_api +if [ $? -ne 0 ];then + exit 1 +fi + +log "[INFO] upgrade version.info" +upgrade_file version.info +if [ $? -ne 0 ];then + exit 1 +fi + +config_file=${targetdir}/vendors/config.ini +found_vendors="$(grep -w "load_priority" "$config_file" | cut --only-delimited -d"=" -f2-)" +found_vendor=$(echo $found_vendors | sed "s/\<$vendor_name\>//g" | tr ',' ' ') +vendor=$(echo $found_vendor | tr -s ' ' ',') +if [ "$vendor" != "" ]; then + sed -i "/load_priority=$found_vendors/s@load_priority=$found_vendors@load_priority=$vendor_name,$vendor@g" "$config_file" +fi + +echo "SUCCESS" +exit 0 diff --git a/csrc/ascend/cmake/scripts/examples/get_opapi_abs_path.py b/csrc/ascend/cmake/scripts/examples/get_opapi_abs_path.py new file mode 100644 index 000000000..12407fdce --- /dev/null +++ b/csrc/ascend/cmake/scripts/examples/get_opapi_abs_path.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +""" +获取 opapi 二进制绝对路径 + +Examples 场景下, 用于 built-in 包与 custom 包共存场景下获取正确的 opapi 动态库绝对路径. +""" + +import argparse +import logging +import os +import subprocess +from pathlib import Path + + +class OpApiMgr: + @staticmethod + def has_symbol(file_path: Path, sym: str) -> bool: + cmd = f"nm -D {file_path}".split() + ret = subprocess.run(cmd, capture_output=True, check=True, encoding="utf-8") + ret.check_returncode() + return sym in ret.stdout + + @staticmethod + def get_environ_custom_lib_paths() -> list[Path]: + paths: list[Path] = [] + env = os.getenv("ASCEND_CUSTOM_OPP_PATH") + if env is None: + logging.debug("ASCEND_CUSTOM_OPP_PATH is none.") + return paths + str_paths = str(env).split(sep=":") + if len(str_paths) == 0: + return paths + for s in str_paths: + if len(s) == 0: + continue + p = Path(s, "op_api/lib/libcust_opapi.so").resolve(strict=False) + if not p.exists(): + logging.warning("Skip not exist path(%s)", p) + continue + paths.append(p) + return paths + + @staticmethod + def get_default_custom_lib_paths() -> list[Path]: + paths: list[Path] = [] + env = os.getenv("ASCEND_OPP_PATH") + if env is None: + logging.warning("ASCEND_OPP_PATH is none.") + return paths + path_env = Path(env).resolve(strict=False) + if not path_env.exists(): + logging.warning("ASCEND_CUSTOM_OPP_PATH(%s) not exist.", path_env) + return paths + cfg_file = Path(path_env, "vendors/config.ini").resolve(strict=False) + if not cfg_file.exists(): + logging.debug("Config file(%s) not exist.", cfg_file) + return paths + # 手工解析 ini 文件 + with open(cfg_file) as fh: + lines = fh.readlines() + for line in lines: + if not line.startswith("load_priority="): + continue + sub_str = line[14:] + sub_str = sub_str.split(sep="#")[0] + sub_str = sub_str.replace("\r", "").replace("\n", "").replace(" ", "") + if len(sub_str) == 0: + continue + vendors = sub_str.split(sep=",") + for v in vendors: + if len(v) == 0: + continue + p = Path(path_env, "vendors", v, "op_api/lib/libcust_opapi.so") + if not p.exists(): + logging.warning("Skip not exist path(%s)", p) + continue + paths.append(p) + return paths + + @staticmethod + def get_default_builtin_lib_paths() -> list[Path]: + paths: list[Path] = [] + env = os.getenv("ASCEND_OPP_PATH") + if env is None: + logging.warning("ASCEND_OPP_PATH is none.") + return paths + path_env = Path(env).resolve(strict=False) + if not path_env.exists(): + logging.warning("ASCEND_CUSTOM_OPP_PATH(%s) not exist.", path_env) + return paths + shared = Path(path_env, "lib64/libopapi.so").resolve(strict=False) + if not shared.exists(): + logging.error("Can't get built-in libopapi.so(%s)", shared) + return paths + paths.append(shared) + return paths + + @staticmethod + def judge_lib_path(sym: str) -> Path | None: + path = None + environ_custom_lib_paths = OpApiMgr.get_environ_custom_lib_paths() + default_custom_lib_paths = OpApiMgr.get_default_custom_lib_paths() + default_builtin_lib_paths = OpApiMgr.get_default_builtin_lib_paths() + path_list = environ_custom_lib_paths + default_custom_lib_paths + default_builtin_lib_paths + for p in path_list: + if OpApiMgr.has_symbol(file_path=p, sym=sym): + path = p + break + return path + + @staticmethod + def main() -> str: + ps = argparse.ArgumentParser(description="Get opapi path", epilog="Best Regards!") + ps.add_argument("-f", "--func", required=True, nargs=1, type=str, help="Func name") + args = ps.parse_args() + sym = args.func[0] + if sym is None or len(sym) == 0: + return "" + lib = OpApiMgr.judge_lib_path(sym=sym) + if lib is None: + return "" + else: + return str(lib) + + +if __name__ == "__main__": + logging.basicConfig(format="%(filename)s:%(lineno)d [%(levelname)s] %(message)s", level=logging.INFO) + print(OpApiMgr.main(), end="") diff --git a/csrc/ascend/cmake/scripts/examples/get_soc_info.py b/csrc/ascend/cmake/scripts/examples/get_soc_info.py new file mode 100644 index 000000000..31c9fbb0d --- /dev/null +++ b/csrc/ascend/cmake/scripts/examples/get_soc_info.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +""" +获取 Soc 相关信息 + +Examples 场景下, 用于获取 Soc 相关信息. +""" + +import argparse +import ctypes +import logging + + +class SocInfoMgr: + @staticmethod + def get_soc_name() -> str: + acl_lib = ctypes.cdll.LoadLibrary("libascendcl.so") + acl_lib.aclrtGetSocName.restype = ctypes.c_char_p + rst = acl_lib.aclrtGetSocName() + if rst: + rst = str(rst, encoding="utf-8") + else: + rst = "" + return rst + + @staticmethod + def main() -> str: + ps = argparse.ArgumentParser(description="Get soc info", epilog="Best Regards!") + ps.add_argument("-i", "--info", required=True, type=str, help="SocInfo") + args = ps.parse_args() + rst = "" + if args.info == "soc_name": + rst = SocInfoMgr.get_soc_name() + else: + logging.error("Unknown SocInfo name %s", args.info) + return rst + + +if __name__ == "__main__": + logging.basicConfig(format="%(filename)s:%(lineno)d [%(levelname)s] %(message)s", level=logging.INFO) + g_rst = "" + try: + g_rst = SocInfoMgr.main() + except Exception as e: + logging.error(e) + print(g_rst, end="") diff --git a/csrc/ascend/cmake/scripts/fix_format.sh b/csrc/ascend/cmake/scripts/fix_format.sh new file mode 100644 index 000000000..5d843babb --- /dev/null +++ b/csrc/ascend/cmake/scripts/fix_format.sh @@ -0,0 +1,35 @@ +#!/bin/bash\n" +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +set -e +if find $1 -type f -name "GroupedMatmul_*json" -print0 | grep -q .; then + echo "Found GroupedMatmul_*json files, performing sed operation..." + find $1 -type f -name "GroupedMatmul_*json" -print0 | xargs -0 sed -i 's/FormatAgnostic/FormatDefault/g' +fi + +if find $1 -type f -name "MlaProlog_*json" -print0 | grep -q .; then + echo "Found MlaProlog_*json files, performing sed operation..." + find $1 -type f -name "MlaProlog_*json" -print0 | xargs -0 sed -i 's/FormatAgnostic/FormatDefault/g' +fi + +if find $1 -type f -name "MlaPrologV2_*json" -print0 | grep -q .; then + echo "Found MlaPrologV2_*json files, performing sed operation..." + find $1 -type f -name "MlaPrologV2_*json" -print0 | xargs -0 sed -i 's/FormatAgnostic/FormatDefault/g' +fi + +if find $1 -type f -name "MlaPrologV3_*json" -print0 | grep -q .; then + echo "Found MlaPrologV3_*json files, performing sed operation..." + find $1 -type f -name "MlaPrologV3_*json" -print0 | xargs -0 sed -i 's/FormatAgnostic/FormatDefault/g' +fi + +if find $1 -type f -name "GroupedMatmulSwigluQuant_*json" -print0 | grep -q .; then + echo "Found GroupedMatmulSwigluQuant_*json files, performing sed operation..." + find $1 -type f -name "GroupedMatmulSwigluQuant_*json" -print0 | xargs -0 sed -i 's/FormatAgnostic/FormatDefault/g' +fi \ No newline at end of file diff --git a/csrc/ascend/cmake/scripts/parse_changed_files.py b/csrc/ascend/cmake/scripts/parse_changed_files.py new file mode 100644 index 000000000..d19495c0b --- /dev/null +++ b/csrc/ascend/cmake/scripts/parse_changed_files.py @@ -0,0 +1,282 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +""" +获取修改文件应触发的测试范围. + +当前仅支持对应触发的 UTest 用例进行分析, 切仅支持 ops_test 这个 UTest 目标. +""" + +import argparse +import logging +from pathlib import Path +from typing import Any + +import yaml + + +class Module: + def __init__(self, name): + self.name: str = name + self.src_files: list[Path] = [] + self.src_exclude_files: list[Path] = [] + self.tests_ut_ops_test_src_files: list[Path] = [] + self.tests_ut_ops_test_src_exclude_files: list[Path] = [] + self.tests_ut_ops_test_options: list[str] = [] + self.options: list[str] = [] + self.test_excludes: list[str] = [] + + @staticmethod + def _add_str_cfg(src, dst: list[str]): + if isinstance(src, str): + src = [src] + for s in src: + if s not in dst: + dst.append(s) + return True + + @staticmethod + def _add_test_excludes(test_options, dst: list[str]): + if isinstance(test_options, dict): + if "examples" in test_options and not test_options["examples"]: + dst.append("examples") + if "ut" in test_options and not test_options["ut"]: + dst.append("ut") + return True + + def update_classify_cfg(self, desc: dict[str, Any]) -> bool: + if not self._update_src(desc=desc): + return False + if not self._update_exclude_src(desc=desc): + return False + if not self._update_test_excludes(desc=desc): + return False + return self._update_options(desc=desc) + + def get_test_options(self, f: Path) -> list[str]: + def is_excluded(e_f: Path): + for e in self.src_exclude_files: + try: + e_f.relative_to(e) + return True + except ValueError: + continue + return False + + related_options: list[str] = [] + for s in self.src_files: + if is_excluded(e_f=f): + continue + try: + if f.relative_to(s): + # 当同一个修改文件需要触发多个 Options 时, 需要把这些 Options 全部添加 + related_options.extend(self.options) + except ValueError: + continue + # 关联 Options 去重 + related_options = list(set(related_options)) + return related_options + + def get_test_example_ops_test_options(self, f: Path) -> list[str]: + return self.get_test_options(f) + + def print_details(self): + dbg_str = ( + f"Name={self.name} SrcLen={len(self.src_files)} " + f"TestUtOpsTestSrcLen={len(self.tests_ut_ops_test_src_files)} " + f"TestUtOpsTestOptions={self.options} " + f"TestUtOpsTestOptions={self.tests_ut_ops_test_options}" + ) + logging.debug(dbg_str) + + def _add_rel_path(self, src, dst: list[Path]): + if isinstance(src, (str, Path)): + src = [src] + for p in src: + p = Path(p) + if p.is_absolute(): + logging.error("[%s]'s Path[%s] is absolute path.", self.name, p) + return False + if p not in dst: + dst.append(p) + return True + + def _update_src(self, desc: dict[str, Any]) -> bool: + src_paths = desc.get("src", []) + return self._add_rel_path(src=src_paths, dst=self.src_files) + + def _update_exclude_src(self, desc: dict[str, Any]) -> bool: + src_paths = desc.get("exclude", []) + return self._add_rel_path(src=src_paths, dst=self.src_exclude_files) + + def _update_test_excludes(self, desc: dict[str, Any]) -> bool: + test_options = desc.get("test", []) + return self._add_test_excludes(test_options=test_options, dst=self.test_excludes) + + def _update_options(self, desc: dict[str, Any]) -> bool: + options = desc.get("options", []) + return self._add_str_cfg(src=options, dst=self.options) + + +class Parser: + """ + 规则文件、修改文件列表文件解析. + """ + + _Modules: list[Module] = [] # 保存规则文件(tests/test_config.yaml)内设置的模块列表 + _ChangedPaths: list[Path] = [] # 修改文件列表文件(changed_file)内设置的修改文件列表 + _UTExcludes: list[str] = [] + _ExamplesExcludes: list[str] = [] + + @classmethod + def print_details(cls): + for m in cls._Modules: + m.print_details() + for p in cls._ChangedPaths: + logging.debug(p) + + @classmethod + def parse_classify_file(cls, file: Path) -> bool: + file = Path(file).resolve() + if not file.exists(): + logging.error("Classify file(%s) not exist.", file) + return False + with open(file, encoding="utf-8") as f: + desc: dict[str, Any] = yaml.load(f, Loader=yaml.SafeLoader) + + def extract_from_dict(obj, current_key="root") -> bool: + # 只看 dict 类型 + if not isinstance(obj, dict): + return True + + # 递归到 module 时说明到达最后一层 + if "module" in obj: + return cls._parse_classify_item(current_key, desc) + + # 递归处理其他值 + return all(extract_from_dict(value, key) for key, value in obj.items()) + + return extract_from_dict(desc) + + @classmethod + def parse_changed_file(cls, file: Path) -> bool: + file = Path(file).resolve() + if not file.exists(): + logging.error("Change files desc file(%s) not exist.", file) + return False + with open(file) as fh: + lines = fh.readlines() + for cur_line in lines: + cur_line = cur_line.strip() + f = Path(cur_line) + if f.is_absolute(): + logging.error("%s is absolute path.", f) + return False + cls._ChangedPaths.append(f) + return True + + @classmethod + def get_related_ut(cls): + ops_test_option_lst: list[str] = [] + for p in cls._ChangedPaths: + for m in cls._Modules: + new_options = m.get_test_options(f=p) + for opt in new_options: + if opt not in ops_test_option_lst: + ops_test_option_lst.append(opt) + if len(ops_test_option_lst) == 0: + logging.info("Don't trigger any UT.") + return "" + ops_test_ut_str: str = "" + if "all" in ops_test_option_lst: + ops_test_ut_str = "all" + else: + for opt in ops_test_option_lst: + if opt not in cls._UTExcludes: + ops_test_ut_str += f"{opt};" + ops_test_ut_str = f"{ops_test_ut_str}" + logging.info("Trigger UT: %s", ops_test_ut_str) + return ops_test_ut_str + + @classmethod + def get_ops_test_option_lst(cls) -> list[str]: + ops_test_option_lst: list[str] = [] + for p in cls._ChangedPaths: + for m in cls._Modules: + new_options = m.get_test_example_ops_test_options(f=p) + for opt in new_options: + if opt not in ops_test_option_lst: + ops_test_option_lst.append(opt) + return ops_test_option_lst + + @classmethod + def get_related_examples(cls) -> str: + ops_test_option_lst = cls.get_ops_test_option_lst() + if len(ops_test_option_lst) == 0: + logging.info("Don't trigger any examples.") + return "" + ops_test_examples_str: str = "" + if "all" in ops_test_option_lst: + ops_test_examples_str = "all" + else: + for opt in ops_test_option_lst: + if opt not in cls._ExamplesExcludes: + ops_test_examples_str += f"{opt};" + ops_test_examples_str = f"{ops_test_examples_str}" + logging.info("Trigger examples: %s", ops_test_examples_str) + return ops_test_examples_str + + @classmethod + def _parse_classify_item(cls, name: str, desc: dict[str, Any] | None = None) -> bool: + if desc is None: + logging.error("[%s]'s desc is None.", name) + return False + if desc.get("module", False): + mod = Module(name=name) + rst = mod.update_classify_cfg(desc=desc) + if rst: + cls._Modules.append(mod) + short_name = name.split("/")[-1] + if "examples" in mod.test_excludes: + cls._ExamplesExcludes.append(short_name) + if "ut" in mod.test_excludes: + cls._UTExcludes.append(short_name) + return rst + return all(cls._parse_classify_item(name=name + "/" + k, desc=sub_desc) for k, sub_desc in desc.items()) + + @staticmethod + def main() -> str: + # 参数注册 + ps = argparse.ArgumentParser(description="Parse changed files", epilog="Best Regards!") + ps.add_argument("-c", "--classify", required=True, nargs=1, type=Path, help="tests/test_config.yaml") + ps.add_argument("-f", "--file", required=True, nargs=1, type=Path, help="changed files desc file.") + # 子命令行 + sub_ps = ps.add_subparsers(help="Sub-Command") + p_ut = sub_ps.add_parser("get_related_ut", help="Get related ut.") + p_ut.set_defaults(func=Parser.get_related_ut) + p_examples = sub_ps.add_parser("get_related_examples", help="Get related examples.") + p_examples.set_defaults(func=Parser.get_related_examples) + # 处理 + args = ps.parse_args() + logging.debug(args) + if not Parser.parse_classify_file(file=Path(args.classify[0])): + return "" + if not Parser.parse_changed_file(file=Path(args.file[0])): + return "" + Parser.print_details() + rst = args.func() + return rst + + +if __name__ == "__main__": + logging.basicConfig( + format="[%(asctime)s][%(filename)s:%(lineno)d] %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO + ) + print(Parser.main()) diff --git a/csrc/ascend/cmake/scripts/prepare.sh b/csrc/ascend/cmake/scripts/prepare.sh new file mode 100644 index 000000000..6cda89707 --- /dev/null +++ b/csrc/ascend/cmake/scripts/prepare.sh @@ -0,0 +1,178 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +CPU_NUM=$(($(cat /proc/cpuinfo | grep "^processor" | wc -l)*2)) +JOB_NUM="-j${CPU_NUM}" + +while [[ $# -gt 0 ]]; do + case $1 in + -s) + PATH_TO_SOURCE="$2" + shift 2 + ;; + -b) + PATH_TO_BUILD="$2" + shift 2 + ;; + -p) + ASCEND_CANN_PACKAGE_PATH="$2" + shift 2 + ;; + --autogen-dir) + ASCEND_AUTOGEN_DIR="$2" + shift 2 + ;; + --build-open-project) + BUILD_OPEN_PROJECT="$2" + shift 2 + ;; + --binary-out-dir) + ASCEND_BINARY_OUT_DIR="$2" + shift 2 + ;; + --impl-out-dir) + ASCEND_IMPL_OUT_DIR="$2" + shift 2 + ;; + --op-build-tool) + OP_BUILD_TOOL="$2" + shift 2 + ;; + --ascend-cmake-dir) + ASCEND_CMAKE_DIR="$2" + shift 2 + ;; + --tiling-key) + TILING_KEY="$2" + shift 2 + ;; + --ops-compile-options) + OPS_COMPILE_OPTIONS="$2" + shift 2 + ;; + --check-compatible) + CHECK_COMPATIBLE="$2" + shift 2 + ;; + --ascend-compute_unit) + ASCEND_COMPUTE_UNIT="$2" + shift 2 + ;; + --ascend-op_name) + ASCEND_OP_NAME="$2" + shift 2 + ;; + --op_debug_config) + OP_DEBUG_CONFIG="$2" + shift 2 + ;; + --build_type) + BUILD_TYPE="$2" + shift 2 + ;; + --version) + VERSION="$2" + shift 2 + ;; + --build_ops_rty_kernel) + BUILD_OPS_RTY_KERNEL="$2" + shift 2 + ;; + --enable_built_in) + ENABLE_BUILT_IN="$2" + shift 2 + ;; + --enable_static) + ENABLE_STATIC="$2" + shift 2 + ;; + --enable_experimental) + ENABLE_EXPERIMENTAL="$2" + shift 2 + ;; + --enable_ccache) + ENABLE_CCACHE="$2" + shift 2 + ;; + --enable_oom) + ENABLE_OOM="$2" + shift 2 + ;; + --cann_3rd_lib_path) + CANN_3RD_LIB_PATH="$(realpath $2)" + shift 2 + ;; + *) + break + ;; + esac +done + +function clean() { + if [ -n "${PATH_TO_BUILD}" ];then + rm -rf ${PATH_TO_BUILD} + mkdir -p ${PATH_TO_BUILD} + fi +} + +function convert_string() { + local _input=$1 + _output=$(echo $_input | sed 's/::/;/g') + echo "${_output}" +} + +function set_env() { + CONVERT_TILING_KEY="$(convert_string ${TILING_KEY})" + + CONVERT_OPS_COMPILE_OPTIONS="$(convert_string ${OPS_COMPILE_OPTIONS})" + + CONVERT_ASCEND_COMPUTE_UNIT="$(convert_string ${ASCEND_COMPUTE_UNIT})" + + CONVERT_ASCEND_OP_NAME="$(convert_string ${ASCEND_OP_NAME})" +} + +function build() { + cd ${PATH_TO_BUILD} + cmake ${PATH_TO_SOURCE} \ + -DBUILD_OPEN_PROJECT=${BUILD_OPEN_PROJECT} \ + -DPREPARE_BUILD=ON \ + -DCUSTOM_ASCEND_CANN_PACKAGE_PATH=${ASCEND_CANN_PACKAGE_PATH} \ + -DASCEND_AUTOGEN_DIR=${ASCEND_AUTOGEN_DIR} \ + -DASCEND_BINARY_OUT_DIR=${ASCEND_BINARY_OUT_DIR} \ + -DASCEND_IMPL_OUT_DIR=${ASCEND_IMPL_OUT_DIR} \ + -DOP_BUILD_TOOL=${OP_BUILD_TOOL} \ + -DASCEND_CMAKE_DIR=${ASCEND_CMAKE_DIR} \ + -DCHECK_COMPATIBLE=${CHECK_COMPATIBLE} \ + -DTILING_KEY="${CONVERT_TILING_KEY}" \ + -DOPS_COMPILE_OPTIONS="${CONVERT_OPS_COMPILE_OPTIONS}" \ + -DASCEND_COMPUTE_UNIT=${CONVERT_ASCEND_COMPUTE_UNIT} \ + -DASCEND_OP_NAME=${CONVERT_ASCEND_OP_NAME} \ + -DENABLE_CCACHE=${ENABLE_CCACHE} \ + -DBUILD_OPS_RTY_KERNEL=${BUILD_OPS_RTY_KERNEL} \ + -DENABLE_BUILT_IN=${ENABLE_BUILT_IN} \ + -DENABLE_STATIC=${ENABLE_STATIC} \ + -DENABLE_EXPERIMENTAL=${ENABLE_EXPERIMENTAL} \ + -DOP_DEBUG_CONFIG=${OP_DEBUG_CONFIG} \ + -DCANN_3RD_LIB_PATH=${CANN_3RD_LIB_PATH} \ + -DCMAKE_BUILD_TYPE=${BUILD_TYPE} \ + -DVERSION=${VERSION} \ + -DENABLE_OOM=${ENABLE_OOM} + + make ${JOB_NUM} prepare_build +} + +function main() { + clean + set_env + build +} + +main diff --git a/csrc/ascend/cmake/scripts/utest/gen_coverage.py b/csrc/ascend/cmake/scripts/utest/gen_coverage.py new file mode 100644 index 000000000..6ae51bf47 --- /dev/null +++ b/csrc/ascend/cmake/scripts/utest/gen_coverage.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +""" +生成覆盖率 +""" + +import argparse +import dataclasses +import logging +import os +import subprocess +from pathlib import Path + +import yaml + + +class GenCoverage: + @dataclasses.dataclass + class Param: + source_dir: Path | None = None + data_dir: Path | None = None + info_file: Path | None = None + info_file_filtered: Path | None = None + html_report_dir: Path | None = None + filter_str: str = "" + + @staticmethod + def get_exclude_paths_from_yaml(yaml_path: Path): + with open(yaml_path, encoding="utf-8") as file: + data = yaml.safe_load(file) + + exclude_paths = set() + + def extract_from_dict(obj, current_key=None): + if isinstance(obj, dict): + if "exclude" in obj and isinstance(obj["exclude"], list): + for path in obj["exclude"]: + exclude_paths.add(path) + if "ut_cov_exclude" in obj and isinstance(obj["ut_cov_exclude"], list): + for path in obj["ut_cov_exclude"]: + exclude_paths.add(path) + + # 递归处理所有值 + for key, value in obj.items(): + extract_from_dict(value, key) + elif isinstance(obj, list): + for item in obj: + extract_from_dict(item, current_key) + + extract_from_dict(data) + + return exclude_paths + + def init_filter_str(self, fs: list[list[str]] | None): + if not fs: + return + for fl in fs: + self.filter_str += f"{fl[0]} " + + def init_filter_str_from_yaml(self, source_dir: Path, yaml_path: Path): + exclude_paths = self.get_exclude_paths_from_yaml(yaml_path) + + for path in sorted(exclude_paths): + if path.startswith("*"): + lcov_path = path + else: + full_path = source_dir / Path(path) + if full_path.is_dir(): + lcov_path = f"{full_path}/*" + else: + lcov_path = f"{full_path}" + + self.filter_str += f"{lcov_path} " + + @classmethod + def main(cls): + # 参数注册 + parser = argparse.ArgumentParser(description="Generate Coverage", epilog="Best Regards!") + parser.add_argument( + "-s", + "--source_base_dir", + required=True, + nargs=1, + type=Path, + help="Explicitly specify the source base directory.", + ) + parser.add_argument( + "-c", + "--coverage_data_dir", + required=True, + nargs=1, + type=Path, + help="Explicitly specify the *.da's base directory.", + ) + parser.add_argument( + "-i", "--info_file", required=False, nargs=1, type=Path, help="Explicitly specify coverage info file path." + ) + # 考虑最低支持 Python 版本为 3.7, 此处用 append 而非 extend + parser.add_argument( + "-f", + "--filter", + required=False, + action="append", + nargs="*", + type=str, + help="Explicitly specify filter file/dir in coverage info.", + ) + parser.add_argument( + "-y", + "--yaml", + required=False, + nargs=1, + type=Path, + help="Explicitly specify filter file/dir from tests/test_config.yaml.", + ) + parser.add_argument( + "--html_report", required=False, nargs=1, type=Path, help="Explicitly specify coverage html report dir." + ) + # 参数解析, 默认值处理 + p = cls.Param() + args = parser.parse_args() + p.source_dir = Path(args.source_base_dir[0]).absolute() + p.data_dir = Path(args.coverage_data_dir[0]).absolute() + if args.info_file: + p.info_file = Path(args.info_file[0]).absolute() + p.info_file_filtered = Path(p.info_file.parent, f"{p.info_file.stem}_filtered{p.info_file.suffix}") + else: + p.info_file = Path(p.data_dir, "cov_result/coverage.info") + p.info_file_filtered = p.info_file + p.html_report_dir = args.html_report[0] if args.html_report else Path(p.info_file.parent, "html_report") + p.html_report_dir = Path(p.html_report_dir).absolute() + p.init_filter_str(fs=args.filter) + p.init_filter_str_from_yaml(source_dir=p.source_dir, yaml_path=args.yaml[0]) if args.yaml else None + logging.debug("[DEBUG] filter_str=%s", p.filter_str) + # 参数检查 + if not p.data_dir.exists(): + logging.error("[ERROR] The dir(%s) required to find the .da files not exist.", p.data_dir) + exit(1) + if not p.info_file.exists(): + p.info_file.parent.mkdir(parents=True, exist_ok=True) + if not p.html_report_dir.exists(): + p.html_report_dir.mkdir(parents=True, exist_ok=True) + # 环境检查 + if not cls._chk_env(): + exit(1) + # 生成覆盖率数据 + cls._gen_cov(param=p) + + @classmethod + def _chk_env(cls): + try: + ret = subprocess.run(["lcov", "--version"], capture_output=True, check=True, encoding="utf-8") + ret.check_returncode() + except FileNotFoundError: + logging.error("[ERROR] lcov is required to generate coverage data, please install.") + return False + try: + ret = subprocess.run(["genhtml", "--version"], capture_output=True, check=True, encoding="utf-8") + ret.check_returncode() + except FileNotFoundError: + logging.error("[ERROR] genhtml is required to generate coverage html report, please install.") + return False + return True + + @classmethod + def _gen_cov(cls, param: Param): + """ + 使用 lcov 生成覆盖率 + """ + # 当 log 等级小于 INFO 时，lcov 不带 -q 标签 + lcov_log_tag = "" if logging.getLogger().level <= logging.INFO else "-q" + logging.critical("================================================================================") + logging.critical("Coverage Report") + logging.critical("================================================================================") + + # 生成覆盖率 + cmd = f"lcov -c -d {param.data_dir} -o {param.info_file} {lcov_log_tag}" + logging.debug("[DEBUG] Generate origin coverage file, cmd=`%s`", cmd) + ret = subprocess.run(cmd.split(), capture_output=False, check=True, encoding="utf-8") + ret.check_returncode() + if param.info_file.stat().st_size == 0: + logging.critical("No file found in origin coverage file.") + return + logging.debug("[DEBUG] Generated origin coverage file %s", param.info_file) + # 滤掉某些文件/路径的覆盖率信息 + cmd = f"lcov --remove {param.info_file} {param.filter_str} -o {param.info_file_filtered} {lcov_log_tag}" + logging.debug("[DEBUG] Generate filtered coverage file, cmd=`%s`", cmd) + ret = subprocess.run(cmd.split(), capture_output=False, check=True, encoding="utf-8") + ret.check_returncode() + logging.debug("[DEBUG] Generated filtered coverage file %s", param.info_file_filtered) + logging.info("[INFO] Generated coverage result in %s", os.path.dirname(param.info_file)) + + if param.info_file_filtered.stat().st_size == 0: + logging.critical("No file found in filtered coverage file.") + return + # 生成 html 报告 + sub_cmd_prefix = f"-p {param.source_dir}" if param.source_dir else "" + cmd = f"genhtml {param.info_file_filtered} {sub_cmd_prefix} -o {param.html_report_dir} {lcov_log_tag}" + logging.debug("[DEBUG] Generate filtered coverage html report, cmd=`%s`", cmd) + ret = subprocess.run(cmd.split(), capture_output=False, check=True, encoding="utf-8") + ret.check_returncode() + logging.info("[INFO] Generated filtered coverage html report. %s", param.html_report_dir) + # 输出覆盖率数据到终端 + cmd = f"lcov --list {param.info_file_filtered}" + ret = subprocess.run(cmd.split(), capture_output=False, check=True, encoding="utf-8") + logging.critical("================================================================================") + ret.check_returncode() + + +if __name__ == "__main__": + # 将环境变量中的 ASCEND_GLOBAL_LOG_LEVEL 换算成 python 的 log 等级 + log_level = (int(os.getenv("ASCEND_GLOBAL_LOG_LEVEL", "3")) + 1) * 10 + logging.basicConfig(format="[%(asctime)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=log_level) + GenCoverage.main() diff --git a/csrc/ascend/cmake/scripts/utest/gen_tiling_data_stub.py b/csrc/ascend/cmake/scripts/utest/gen_tiling_data_stub.py new file mode 100644 index 000000000..2cc5e217d --- /dev/null +++ b/csrc/ascend/cmake/scripts/utest/gen_tiling_data_stub.py @@ -0,0 +1,432 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +""" +生成 TilingData 桩 + +用于 UTest 场景下, 生成 Struct 表示的 TilingData 相关头文件. +""" + +import argparse +import datetime +import logging +import os +import stat +from pathlib import Path + +import regex as re + + +def check_if_new_tiling_file_path_existed(ori_file: Path) -> Path: + current_path = ori_file + target_dir = ori_file / "op_kernel" + while True: + if target_dir.is_dir(): + break + parent_path = current_path.parent + if parent_path == current_path: + return ori_file, False + current_path = parent_path + target_dir = current_path / "op_kernel" + tiling_files = list(target_dir.glob("*tiling_data.h")) + if not tiling_files: + return ori_file, False + new_file = tiling_files[0] + new_path = target_dir / new_file.name + return new_path, True + + +def process_class_fields(fields_str, class_name): + result = [] + seen = set() + arr_pattern = re.compile( + r"^\s*" + r"(?P(?:[\w:<>]+\s+)*[\w:<>* &]+?)\s+" # 类型（含修饰符/指针/引用） + r"(?P\w+)\s*" # 变量名 + r"\[(?P\d+)\]" # 数组长度 + r"(?:\s*=\s*\{\s*[^\}]*\s*})?" # 可选初始化，支持 {} 或 {0, ...} + r"\s*;\s*" # 以 ; 结尾 + r"(?:[ \t]*(?://[^\n]*)?)?$", # 行尾可有 // 注释 + re.MULTILINE, + ) + var_pattern = re.compile(r"^\s*(?P[\w:<>]+)\s+(?P\w+)\s*(?:=\s*[^;]*)?;", re.MULTILINE) + # 逐行处理，保持顺序 + for line in fields_str.splitlines(): + # 跳空行或纯注释行 + if not line.strip() or line.strip().startswith("//"): + continue + + m = arr_pattern.match(line) + if m: + t, n, ln = m.group("type"), m.group("name"), m.group("len") + result.append(("array", t, n, ln)) + seen.add(n) + continue + + m = var_pattern.match(line) + if m: + t, n = m.group("type"), m.group("name") + if n not in seen: + result.append(("normal", t, n)) + seen.add(n) + continue + return result + + +def find_classes(content): + out = [] + class_re = re.compile(r"\bclass\s+(\w+)\s*{") + for m in class_re.finditer(content): + class_name = m.group(1) + start = m.end() + idx = start + braces = 1 + while idx < len(content): + c = content[idx] + if c == "{": + braces += 1 + elif c == "}": + braces -= 1 + if braces == 0: + out.append((class_name, content[start:idx].strip())) + break + idx += 1 + return out + + +def convert_template_tilingkey(ori_file: Path): + with open(ori_file) as f: + content = f.read() + + classes = find_classes(content) + output = [] + for class_name, fields_str in classes: + fields = process_class_fields(fields_str, class_name) + output.append(f"BEGIN_TILING_DATA_DEF({class_name})") + for entry in fields: + if entry[0] == "normal": + _, field_type, field_name = entry + if field_type in [ + "uint32_t", + "int32_t", + "uint8_t", + "uint16_t", + "float", + "uint64_t", + "int64_t", + "double", + ]: + output.append(f"TILING_DATA_FIELD_DEF({field_type}, {field_name});") + else: + output.append(f"TILING_DATA_FIELD_DEF_STRUCT({field_type}, {field_name});") + elif entry[0] == "array": + _, field_type, field_name, field_len = entry + output.append(f"TILING_DATA_FIELD_DEF_ARR({field_type}, {field_len}, {field_name});") + output.append("END_TILING_DATA_DEF;") + output.append(f"REGISTER_TILING_DATA_CLASS({class_name}Op, {class_name})\n") + result_code = "\n".join(output) + + return result_code + + +def process_fields(fields_str, struct_name): + field_pattern = re.compile(r"(\w+)\s+(\w+)(?:\s*=\d+)?;") + fields = field_pattern.findall(fields_str) + return fields + + +def convert_to_old_tiling_struct_style(redirected_file_path): + with open(redirected_file_path) as f: + content = f.read() + struct_pattern = re.compile(r"struct (\w+) {([^}]*)}", re.DOTALL) + structs = struct_pattern.findall(content) + output = [] + for struct_name, fields_str in structs: + fields = process_fields(fields_str, struct_name) + output.append(f"BEGIN_TILING_DATA_DEF({struct_name})") + for field_type, field_name in fields: + if field_type in ["uint32_t", "uint8_t", "uint16_t"]: + output.append(f"TILING_DATA_FIELD_DEF({field_type}, {field_name});") + else: + output.append(f"TILING_DATA_FIELD_DEF_STRUCT({field_type}, {field_name});") + output.append("END_TILING_DATA_DEF;") + output.append(f"REGISTER_TILING_DATA_CLASS({struct_name}Op, {struct_name})\n") + result_code = "\n".join(output) + return result_code + + +class Process: + _WRITE_FLAGS = os.O_WRONLY | os.O_CREAT | os.O_TRUNC + _WRITE_MODES = stat.S_IWUSR | stat.S_IRUSR + + @classmethod + def _write_file(cls, file: Path, src: str): + with os.fdopen(os.open(file, cls._WRITE_FLAGS, cls._WRITE_MODES), "w") as fh: + fh.write(src) + + @classmethod + def _get_begin_source(cls, ori_file: Path, gen_file: Path) -> str: + bgn_src: str = ( + "/**\n" + " * This program is free software, you can redistribute it and/or modify.\n" + " * Copyright (c) {year} Huawei Technologies Co., Ltd.\n" + " * This file is a part of the CANN Open Software.\n" + ' * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").\n' + " * Please refer to the License for details. " + "You may not use this file except in compliance with the License.\n" + ' * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, ' + "WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, " + "INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.\n" + " * See LICENSE in the root of the software repository for the full text of the License.\n" + " */\n" + ).format(year=datetime.datetime.today().year) + bgn_src += "\n" + bgn_src += ("/*!\n * \\file {gen_file_name}\n * \\brief Generate {ori_file_name}\n */\n").format( + gen_file_name=gen_file.name, ori_file_name=ori_file.name + ) + bgn_src += "\n" + bgn_src += "#pragma once\n" + bgn_src += "\n" + return bgn_src + + @classmethod + def _get_tiling_source(cls, ori_file: Path, isTemplateTilingKey: bool = False) -> str: + """ + 获取 TilingData 定义源码 + + :param ori_file: 原始文件 + :return: 生成文件内容 + """ + rst_source = ( + "#include \n#include \n#include \n#include \n\n" + ) + pattern = re.compile(r"[(](.*)[)]", re.S) + if isTemplateTilingKey: + lines = convert_template_tilingkey(ori_file) + lines = lines.splitlines() + else: + ori_file, existed_flag = check_if_new_tiling_file_path_existed(ori_file) + if existed_flag: + lines = convert_to_old_tiling_struct_style(ori_file) + lines = lines.splitlines() + else: + with open(ori_file) as fd: + lines = fd.readlines() + for line in lines: + line = line.strip() + struct_src = "" + if line.startswith("BEGIN_TILING_DATA_DEF"): + struct_name = re.findall(pattern, line)[0] + struct_src += ("#pragma pack(1)\nstruct {}\n").format(struct_name) + struct_src += "{\n" + struct_offset = 0 + elif line.startswith("TILING_DATA_FIELD_DEF_ARR"): + field_params = re.findall(pattern, line)[0] + fds = field_params.split(",") + fds_dtype = fds[0].strip() + fds_num = int(fds[1].strip()) + fds_name = fds[2].strip() + tmp_src, tmp_offset = cls._get_tmp_src( + offset=struct_offset, dtype=fds_dtype, name=fds_name, num=fds_num + ) + struct_src += tmp_src + struct_offset += tmp_offset + elif line.startswith("TILING_DATA_FIELD_DEF_STRUCT"): + field_params = re.findall(pattern, line)[0] + fds = field_params.split(",") + struct_src += " {} {};\n".format(fds[0].strip(), fds[1].strip()) + elif line.startswith("TILING_DATA_FIELD_DEF"): + field_params = re.findall(pattern, line)[0] + fds = field_params.split(",") + fds_dtype = fds[0].strip() + fds_num = 1 + fds_name = fds[1].strip() + tmp_src, tmp_offset = cls._get_tmp_src( + offset=struct_offset, dtype=fds_dtype, name=fds_name, num=fds_num + ) + struct_src += tmp_src + struct_offset += tmp_offset + elif line.startswith("END_TILING_DATA_DEF"): + # 要求结构体满足 8 字节对齐 + if struct_offset % 8 != 0: + pad_num = 8 - (struct_offset % 8) + struct_src += " uint8_t {}_PH[{}] = {{}};\n".format(struct_name, pad_num) + struct_offset += pad_num + struct_src += "};" + struct_src += "\n" + struct_src += "#pragma pack()\n" + struct_src += "\n" + struct_src += "inline void Init{struct_name}(uint8_t* tiling, {struct_name}* const_data)\n".format( + struct_name=struct_name + ) + struct_src += "{\n" + struct_src += ( + " (void)memcpy_s(const_data, sizeof({struct_name}), tiling, sizeof({struct_name}));\n".format( + struct_name=struct_name + ) + ) + struct_src += "}\n" + struct_src += "\n" + rst_source += struct_src + rst_source += ( + "" + "#undef GET_TILING_DATA\n" + "#define GET_TILING_DATA(tiling_data, tiling_arg) \\\n" + "{struct_name} tiling_data; \\\n" + "Init{struct_name}(tiling_arg, &tiling_data)\n" + "\n" + ).format(struct_name=struct_name) + return rst_source + + @classmethod + def _get_tiling_whole(cls, ori_file: Path, isTemplateTilingKey: bool = False) -> str: + with open(ori_file) as f: + content = f.read() + return content + + @classmethod + def _gen_tiling_h(cls, ori_file: Path, gen_dir: Path): + gen_file = Path(gen_dir, "_gen_" + ori_file.name) + flag = "op_kernel" in [part for part in ori_file.parts] + if not gen_file.exists(): + if not flag: + bgn_src = cls._get_begin_source(ori_file=ori_file, gen_file=gen_file) + def_src = cls._get_tiling_source(ori_file=ori_file, isTemplateTilingKey=flag) + source = bgn_src + def_src + else: + source = "\n" + cls._write_file(file=gen_file, src=source) + logging.info("Generate TilingDefFile: %s", gen_file) + return gen_file + + @classmethod + def _get_type_size(cls, dtype: str): + mp = { + "int8_t": 1, + "int16_t": 2, + "int32_t": 4, + "int64_t": 8, + "uint8_t": 1, + "uint16_t": 2, + "uint32_t": 4, + "uint64_t": 8, + "float": 4, + } + d_len = mp.get(dtype) + if d_len is None: + raise ValueError(f"Unknown dtype({dtype})") + return d_len + + @classmethod + def _get_tmp_src(cls, offset: int, dtype: str, name: str, num: int): + source = "" + result = 0 + dtype_size = cls._get_type_size(dtype=dtype) + + if offset % dtype_size != 0: + pad_num = dtype_size - (offset % dtype_size) + source += " uint8_t {}_PH[{}] = {{}};\n".format(name, pad_num) + result += pad_num + + if num == 1: + source += " {} {} = 0;\n".format(dtype, name) + else: + source += " {} {}[{}] = {{}};\n".format(dtype, name, num) + result += cls._get_type_size(dtype=dtype) * num + return source, result + + @classmethod + def gen_tiling_h(cls, ori_files: list[Path], gen_dir: Path): + gen_files: list[Path] = [] + gen_dir.mkdir(parents=True, exist_ok=True) + for ori_file in ori_files: + if not ori_file.exists(): + raise ValueError(f"Origin file({ori_file}) not exist.") + gen_file = cls._gen_tiling_h(ori_file=ori_file, gen_dir=gen_dir) + gen_files.append(gen_file) + return gen_files + + @classmethod + def gen_tiling_data_h(cls, op: str, gen_files: list[Path], data_file: Path): + if not data_file.exists(): + bgn_src = cls._get_begin_source(ori_file=data_file, gen_file=data_file) + def_src = "" + for gen_f in gen_files: + def_src += '#include "tiling/{op}/{file_name}"\n'.format(op=op, file_name=gen_f.name) + source = bgn_src + def_src + cls._write_file(file=data_file, src=source) + logging.info("Generate TilingDataFile: %s", data_file) + return data_file + + @classmethod + def gen_tiling_stub_h(cls, data_file: Path, stub_file: Path): + if not stub_file.exists(): + bgn_src = cls._get_begin_source(ori_file=stub_file, gen_file=stub_file) + def_src = "" + def_src += '#include "{}"\n'.format(data_file.name) + def_src += ( + "\n" + "#undef GET_TILING_DATA_WITH_STRUCT\n" + "#define GET_TILING_DATA_WITH_STRUCT(tiling_struct, tiling_data, tiling_arg) \\\n" + "tiling_struct tiling_data; \\\n" + "(void)memcpy_s(&tiling_data, sizeof(tiling_struct), tiling_arg, sizeof(tiling_struct));\n" + "\n" + ) + def_src += ( + "\n" + "#undef GET_TILING_DATA_MEMBER\n" + "#define GET_TILING_DATA_MEMBER(tiling_type, member, var, tiling) \\\n" + "decltype(tiling_type::member) var; \\\n" + "size_t offset##var = (size_t)(&((tiling_type *)0)->member); \\\n" + "(void)memcpy_s(&var, sizeof(decltype(var)), tiling + offset##var, sizeof(decltype(var))); \n" + ) + source = bgn_src + def_src + cls._write_file(file=stub_file, src=source) + logging.info("Generate TilingStubFile: %s", stub_file) + return stub_file + + @classmethod + def main(cls): + # 参数注册 + parser = argparse.ArgumentParser(description="TilingData Generator", epilog="Best Regards!") + parser.add_argument("-o", "--operator", required=True, nargs=1, type=str, help="Target operator.") + parser.add_argument( + "-s", + "--srcs", + required=True, + action="append", + nargs="+", + type=Path, + help="Origin tiling data define files(.h).", + ) + parser.add_argument("-d", "--dest", required=True, nargs=1, type=Path, help="Generate directory.") + # 参数解析 + result = parser.parse_args() + op = result.operator[0].lower() + ori_files: list[Path] = [] + for file in result.srcs: + ori_files.append(file[0].absolute()) + gen_dir = Path(result.dest[0], "tiling/{}".format(op)).absolute() + data_file = Path(gen_dir, "tiling_data.h") + stub_file = Path(gen_dir, "tiling_stub.h") + + # 流程处理 + gen_files = cls.gen_tiling_h(ori_files=ori_files, gen_dir=gen_dir) + cls.gen_tiling_data_h(op=op, gen_files=gen_files, data_file=data_file) + cls.gen_tiling_stub_h(data_file=data_file, stub_file=stub_file) + + +if __name__ == "__main__": + logging.basicConfig(format="%(filename)s:%(lineno)d [%(levelname)s] %(message)s", level=logging.DEBUG) + try: + Process.main() + except Exception as e: + logging.error(e) + raise e diff --git a/csrc/ascend/cmake/scripts/util/ascendc_bin_param_build.py b/csrc/ascend/cmake/scripts/util/ascendc_bin_param_build.py new file mode 100644 index 000000000..373f7a950 --- /dev/null +++ b/csrc/ascend/cmake/scripts/util/ascendc_bin_param_build.py @@ -0,0 +1,534 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import argparse +import copy +import hashlib +import json +import os +import sys +from collections import defaultdict +from typing import NamedTuple + +import const_var +import opdesc_parser +import regex as re + +PYF_PATH = os.path.dirname(os.path.realpath(__file__)) + + +class ParamInfo(NamedTuple): + dtype_list: list + format_list: list + dtype_for_bin_list: dict + format_for_bin_list: dict + + +class BinParamBuilder(opdesc_parser.OpDesc): + def __init__(self: any, op_type: str): + super().__init__(op_type) + self.soc = "" + self.out_path = "" + self.tiling_keys = set() + self.op_debug_config = "" + self.op_super_config = [] + + def set_soc_version(self: any, soc: str): + self.soc = soc + + def set_out_path(self: any, out_path: str): + self.out_path = out_path + + def set_tiling_key(self: any, tiling_key_info: set): + if tiling_key_info: + self.tiling_keys.update(tiling_key_info) + + def set_op_debug_config(self: any, op_debug_config: str): + if op_debug_config: + self.op_debug_config = op_debug_config + + def set_op_super_config(self: any, op_super_config: str): + if op_super_config: + self.op_super_config = op_super_config + + def get_full_list(self: any): + dtype_list = [] + for dtype_in in self.input_dtype: + dtype_list.append(dtype_in.split(",")) + for dtype_out in self.output_dtype: + dtype_list.append(dtype_out.split(",")) + + format_list = [] + for fmt_in in self.input_fmt: + format_list.append(fmt_in.split(",")) + for fmt_out in self.output_fmt: + format_list.append(fmt_out.split(",")) + + dtype_for_bin_list = [[] for _ in range(len(self.input_dtype) + len(self.output_dtype))] + format_for_bin_list = copy.deepcopy(dtype_for_bin_list) + + for key, value in self.input_dtype_for_bin.items(): + dtype_for_bin_list[key] = value.split(",") + for key, value in self.output_dtype_for_bin.items(): + dtype_for_bin_list[key + len(self.input_dtype)] = value.split(",") + for key, value in self.input_fmt_for_bin.items(): + format_for_bin_list[key] = value.split(",") + for key, value in self.output_fmt_for_bin.items(): + format_for_bin_list[key + len(self.input_dtype)] = value.split(",") + + return ParamInfo(dtype_list, format_list, dtype_for_bin_list, format_for_bin_list) + + def gen_bin_cprs_list(self: any, param_info: ParamInfo): + combine_dict = {} + origin_combine_dict = {} + for cob_idx in range(0, len(self.input_dtype[0].split(","))): + origin_combine = "" + combine = "" + for param_idx in range(0, len(self.input_dtype) + len(self.output_dtype)): + if param_info.dtype_for_bin_list[param_idx]: + combine += param_info.dtype_for_bin_list[param_idx][cob_idx] + else: + combine += param_info.dtype_list[param_idx][cob_idx] + origin_combine += param_info.dtype_list[param_idx][cob_idx] + if param_info.format_for_bin_list[param_idx]: + combine += param_info.format_for_bin_list[param_idx][cob_idx] + else: + combine += param_info.format_list[param_idx][cob_idx] + origin_combine += param_info.format_list[param_idx][cob_idx] + if combine not in combine_dict: + combine_dict[combine] = [] + combine_dict[combine].append(cob_idx) + origin_combine_dict[origin_combine] = cob_idx + for key, value in combine_dict.items(): + if key not in origin_combine_dict: + print(f"WARNING: ForBinQuery {key} not in origin combine") + self.bin_save_list += value + continue + if len(value) == 1 and value[0] == origin_combine_dict[key]: + self.bin_save_list += value + continue + self.bin_cprs_head.append(origin_combine_dict[key]) + self.bin_cprs_list.append(value) + for index, sub_list in enumerate(self.bin_cprs_list): + if self.bin_cprs_head[index] not in self.bin_save_list: + continue + sub_list.append(self.bin_cprs_head[index]) + self.bin_save_list += self.bin_cprs_head + + def gen_for_bin_list(self: any, param_info: ParamInfo): + combine_size = len(self.input_dtype[0].split(",")) + input_size = len(self.input_dtype) + output_size = len(self.output_dtype) + + self.input_dtype_for_bin_list = [[] for _ in range(input_size)] + self.output_dtype_for_bin_list = [[] for _ in range(output_size)] + for i in range(0, input_size): + self.input_dtype_for_bin_list[i] = [[] for _ in range(combine_size)] + for i in range(0, output_size): + self.output_dtype_for_bin_list[i] = [[] for _ in range(combine_size)] + self.input_fmt_for_bin_list = copy.deepcopy(self.input_dtype_for_bin_list) + self.output_fmt_for_bin_list = copy.deepcopy(self.output_dtype_for_bin_list) + + for index, sub_list in enumerate(self.bin_cprs_list): + head_idx = self.bin_cprs_head[index] + for cmb_idx in sub_list: + for i in range(0, input_size): + self.input_dtype_for_bin_list[i][head_idx].append(param_info.dtype_list[i][cmb_idx]) + self.input_fmt_for_bin_list[i][head_idx].append(param_info.format_list[i][cmb_idx]) + for i in range(0, output_size): + self.output_dtype_for_bin_list[i][head_idx].append(param_info.dtype_list[i + input_size][cmb_idx]) + self.output_fmt_for_bin_list[i][head_idx].append(param_info.format_list[i + input_size][cmb_idx]) + + def rm_cprs_cmb(self: any, dtype_list, format_list, input_size, output_size): + for i in range(0, input_size): + self.input_dtype_for_bin_list[i] = [ + element for index, element in enumerate(self.input_dtype_for_bin_list[i]) if index in self.bin_save_list + ] + self.input_fmt_for_bin_list[i] = [ + element for index, element in enumerate(self.input_fmt_for_bin_list[i]) if index in self.bin_save_list + ] + new_dtype_list = [element for index, element in enumerate(dtype_list[i]) if index in self.bin_save_list] + new_dtype_str = "" + for dtype in new_dtype_list: + new_dtype_str += f"{dtype}," + self.input_dtype[i] = new_dtype_str[:-1] + new_format_list = [element for index, element in enumerate(format_list[i]) if index in self.bin_save_list] + new_format_str = "" + for fmt in new_format_list: + new_format_str += f"{fmt}," + self.input_fmt[i] = new_format_str[:-1] + for i in range(0, output_size): + self.output_dtype_for_bin_list[i] = [ + element + for index, element in enumerate(self.output_dtype_for_bin_list[i]) + if index in self.bin_save_list + ] + self.output_fmt_for_bin_list[i] = [ + element for index, element in enumerate(self.output_fmt_for_bin_list[i]) if index in self.bin_save_list + ] + new_dtype_list = [ + element for index, element in enumerate(dtype_list[i + input_size]) if index in self.bin_save_list + ] + new_dtype_str = "" + for dtype in new_dtype_list: + new_dtype_str += f"{dtype}," + self.output_dtype[i] = new_dtype_str[:-1] + new_format_list = [ + element for index, element in enumerate(format_list[i + input_size]) if index in self.bin_save_list + ] + new_format_str = "" + for fmt in new_format_list: + new_format_str += f"{fmt}," + self.output_fmt[i] = new_format_str[:-1] + + def is_set_for_bin_query(self: any): + return any( + [ + self.input_dtype_for_bin, + self.output_dtype_for_bin, + self.input_fmt_for_bin, + self.output_fmt_for_bin, + ] + ) + + def for_bin_list_match(self: any): + if not self.is_set_for_bin_query(): + return + input_size = len(self.input_dtype) + output_size = len(self.output_dtype) + param_info = self.get_full_list() + self.gen_bin_cprs_list(param_info) + self.gen_for_bin_list(param_info) + if len(self.bin_save_list) == len(self.input_dtype[0].split(",")): + print("WARNING: ForBinQuery can not compress number of bin file with this set, please check!!.") + return + self.rm_cprs_cmb(param_info.dtype_list, param_info.format_list, input_size, output_size) + + def gen_input_json(self: any, auto_gen_path: str): + key_map = {} + self.for_bin_list_match() + if len(self.input_dtype) == 0: + count = len(self.output_dtype[0].split(",")) + else: + count = len(self.input_dtype[0].split(",")) + if count == 0: + raise RuntimeError(f"Op {self.op_type} must have at least one input or output") + required_parameters = set() + index_value = -1 + + for i in range(0, count): + inputs = [] + outputs = [] + attrs = [] + required_parameter = [] + op_node = {} + + for idx in range(0, len(self.input_name)): + idtypes = self.input_dtype[idx].split(",") + ifmts = self.input_fmt[idx].split(",") + itype = self.input_type[idx] + para = {} + para["name"] = self.input_name[idx][:-5] + para["index"] = idx + para["dtype"] = idtypes[i] + if self.is_set_for_bin_query() and self.input_dtype_for_bin_list[idx][i]: + para["dtypeForBinQuery"] = self.input_dtype_for_bin_list[idx][i] + para["format"] = ifmts[i] + if self.is_set_for_bin_query() and self.input_fmt_for_bin_list[idx][i]: + para["formatForBinQuery"] = self.input_fmt_for_bin_list[idx][i] + para["paramType"] = itype + para["shape"] = [-2] + para["format_match_mode"] = "FormatAgnostic" + + input_parameter_key = (idtypes[i], ifmts[i]) + if itype == "dynamic": + inputs.append([para]) + required_parameter.append(input_parameter_key) + elif itype == "required": + inputs.append(para) + required_parameter.append(input_parameter_key) + else: + inputs.append(para) + + for idx in range(0, len(self.output_name)): + odtypes = self.output_dtype[idx].split(",") + ofmts = self.output_fmt[idx].split(",") + otype = self.output_type[idx] + para = {} + para["name"] = self.output_name[idx][:-5] + para["index"] = idx + para["dtype"] = odtypes[i] + if self.is_set_for_bin_query() and self.output_dtype_for_bin_list[idx][i]: + para["dtypeForBinQuery"] = self.output_dtype_for_bin_list[idx][i] + para["format"] = ofmts[i] + if self.is_set_for_bin_query() and self.output_fmt_for_bin_list[idx][i]: + para["formatForBinQuery"] = self.output_fmt_for_bin_list[idx][i] + para["paramType"] = otype + para["shape"] = [-2] + para["format_match_mode"] = "FormatAgnostic" + output_parameter_key = (odtypes[i], ofmts[i]) + if otype == "dynamic": + outputs.append([para]) + required_parameter.append(output_parameter_key) + elif otype == "required": + outputs.append(para) + required_parameter.append(output_parameter_key) + else: + outputs.append(para) + + for attr in self.attr_list: + att = {} + att["name"] = attr + atype = self.attr_val.get(attr).get("type").lower() + att["dtype"] = atype + att["value"] = const_var.ATTR_DEF_VAL.get(atype) + attrs.append(att) + + required_parameter_tuple = tuple(required_parameter) + if required_parameter_tuple in required_parameters: + continue + else: + required_parameters.add(required_parameter_tuple) + index_value += 1 + + op_node["bin_filename"] = "" + op_node["inputs"] = inputs + op_node["outputs"] = outputs + if len(attrs) > 0: + op_node["attrs"] = attrs + + param = {} + param["op_type"] = self.op_type + param["op_list"] = [op_node] + objstr = json.dumps(param, indent=" ") + md5sum = hashlib.md5(objstr.encode("utf-8")).hexdigest() + while key_map.get(md5sum) is not None: + objstr += "1" + md5sum = hashlib.md5(objstr.encode("utf-8")).hexdigest() + key_map[md5sum] = md5sum + bin_file = self.op_type + "_" + md5sum + op_node["bin_filename"] = bin_file + param_file = os.path.join(self.out_path, bin_file + "_param.json") + param_file = os.path.realpath(param_file) + + self._write_build_json(param_file, param) + self._write_build_cmd(param_file, bin_file, index_value, auto_gen_path) + if self.op_super_config: + bin_file += "_relocatable" + op_node["bin_filename"] = bin_file + param_file = os.path.join(self.out_path, bin_file + "_param.json") + param_file = os.path.realpath(param_file) + self._write_build_json(param_file, param) + index_value += 1 + self._write_build_cmd(param_file, bin_file, index_value, auto_gen_path, True) + + def _write_build_json(self: any, param_file: str, param): + with os.fdopen(os.open(param_file, const_var.WFLAGS, const_var.WMODES), "w") as fd: + json.dump(param, fd, indent=" ") + + def _generate_check_result(self: any, enable_tiling_keys: bool, bin_file: str): + check_result = "" + if enable_tiling_keys is False: + check_result += 'echo "${res}"\n' + check_result += const_var.CHK_CMD.format(res_file=bin_file + ".json") + check_result += const_var.CHK_CMD.format(res_file=bin_file + ".o") + else: + check_result += "if [ $? -eq 1 ]; then\n" + check_result += ' if echo "${res}" | \ +grep -q "None of the given tiling keys are in the supported list"; then\n' + check_result += ' echo "${res}"\n' + check_result += " else\n" + check_result += ' echo "${res}"\n' + check_result += " exit 1\n" + check_result += " fi\n" + check_result += "else\n" + check_result += 'echo "${res}"\n' + check_result += const_var.CHK_CMD.format(res_file=bin_file + ".json") + check_result += const_var.CHK_CMD.format(res_file=bin_file + ".o") + check_result += "fi\n" + return check_result + + def _write_build_cmd(self: any, param_file: str, bin_file: str, index: int, auto_gen_path: str, super_mode=False): + hard_soc = const_var.conv_soc_ver(self.soc) + if not hard_soc: + hard_soc = self.soc.capitalize() + name_com = [self.op_type, self.op_file, str(index)] + compile_file = os.path.join(self.out_path, "-".join(name_com) + ".sh") + compile_file = os.path.realpath(compile_file) + + bin_cmd_str = "res=$(opc $1 --main_func={fun} --input_param={param} --soc_version={soc} \ + --output=$2 --impl_mode={impl} --simplified_key_mode=0 --op_mode=dynamic " + + build_cmd_var = "#!/bin/bash\n" + build_cmd_var += f'echo "[{self.soc}] Generating {bin_file} ..."\n' + plog_level = os.environ.get("ASCEND_GLOBAL_LOG_LEVEL") + plog_stdout = os.environ.get("ASCEND_SLOG_PRINT_TO_STDOUT") + if plog_level is None: + build_cmd_var += const_var.SET_PLOG_LEVEL_ERROR + if plog_stdout is None: + build_cmd_var += const_var.SET_PLOG_STDOUT + build_cmd_var += const_var.SRC_ENV + if hard_soc == "Ascend610Lite": + build_cmd_var += f"export ASCEND_CUSTOM_OPP_PATH={auto_gen_path}:$ASCEND_CUSTOM_OPP_PATH \n" + build_cmd_var += bin_cmd_str.format( + fun=self.op_intf, soc=hard_soc, param=param_file, impl="high_performance,optional" + ) + enable_tiling_keys = False + if self.tiling_keys: + tiling_keys_list = sorted(list(self.tiling_keys)) + tiling_key_str = ",".join([str(_key) for _key in tiling_keys_list]) + build_cmd_var += f' --tiling_key="{tiling_key_str}"' + enable_tiling_keys = True + + if self.op_debug_config: + op_debug_str = ",".join([str(_key) for _key in list(self.op_debug_config)]) + build_cmd_var += f" --op_debug_config={op_debug_str}" + + if super_mode and self.op_super_config: + op_super_config_str = " ".join([str(_key) for _key in list(self.op_super_config)]) + build_cmd_var += f" {op_super_config_str}" + + build_cmd_var += ")\n" + build_cmd_var += "\n" + + check_result = self._generate_check_result(enable_tiling_keys, bin_file) + build_cmd_var += check_result + build_cmd_var += f'echo "[{self.soc}] Generating {bin_file} Done"\n' + + with os.fdopen(os.open(compile_file, const_var.WFLAGS, const_var.WMODES), "w") as fd: + fd.write(build_cmd_var) + + +def get_tiling_keys(tiling_keys: str) -> set: + all_tiling_keys = set() + if not tiling_keys: + return all_tiling_keys + + tiling_key_list = tiling_keys.split(";") + for tiling_key_value in tiling_key_list: + pattern = r"(? int(end): + continue + for i in range(int(start), int(end) + 1): + all_tiling_keys.add(i) + elif tiling_key_value.isdigit(): + all_tiling_keys.add(int(tiling_key_value)) + return all_tiling_keys + + +def trans_soc_verion(soc_ver: str): + low_soc_ver = soc_ver.lower() + if low_soc_ver not in opdesc_parser.SOC_TO_SHORT_SOC_MAP: + return low_soc_ver + return opdesc_parser.SOC_TO_SHORT_SOC_MAP[low_soc_ver] + + +def parse_op_debug_confg(opc_config_file: str, soc: str) -> dict: + tiling_key_info = defaultdict(set) + op_debug_config = defaultdict(set) + if not opc_config_file: + return tiling_key_info, op_debug_config + + if not os.path.exists(opc_config_file): + return tiling_key_info, op_debug_config + + with open(opc_config_file) as file: + contents = file.readlines() + + for _content in contents: + content = _content.strip() + opc_configs = content.split("@") + if len(opc_configs) < 3: + continue + + op_type = opc_configs[0] + if not op_type: + continue + + compute_unit = opc_configs[1] + if compute_unit: + compute_unit_list = compute_unit.split(";") + soc_lists = [] + for soc_ver in compute_unit_list: + short_soc_ver = trans_soc_verion(soc_ver) + soc_lists.append(short_soc_ver) + if soc not in soc_lists: + continue + + for options in opc_configs[2:]: + if "--tiling_key" in options: + format_tiling_keys = get_tiling_keys(options.split("=")[1]) + if format_tiling_keys: + tiling_key_info[op_type].update(format_tiling_keys) + if "--op_debug_config" in options: + first_index = options.find("=") + if first_index != -1: + debug_config = options[first_index + 1 :] + else: + debug_config = "" + + format_debug_config = set(debug_config.split(";")) + for _config in format_debug_config: + op_debug_config[op_type].add(_config) + return tiling_key_info, op_debug_config + + +def gen_bin_param_file(cfgfile: str, out_dir: str, soc: str, opc_config_file: str = "", ops: list = None): + if not os.path.exists(cfgfile): + print(f"INFO: {cfgfile} does not exists in this project, skip generating compile commands.") + return + + debug_config = defaultdict(set) + super_config = defaultdict(set) + + op_descs = opdesc_parser.get_op_desc(cfgfile, [], [], BinParamBuilder, ops) + tiling_key_info, op_debug_config = parse_op_debug_confg(opc_config_file, soc) + for _op_type, _op_option in op_debug_config.items(): + for _option in _op_option: + if _option.startswith("--op_relocatable_kernel_binary") or _option.startswith("--op_super_kernel_options"): + super_config[_op_type].add(_option) + else: + debug_config[_op_type].add(_option) + + auto_gen_path_dir = os.path.dirname(cfgfile) + all_soc_key = "ALL" + for op_desc in op_descs: + op_desc.set_soc_version(soc) + op_desc.set_out_path(out_dir) + if op_desc.op_type in debug_config: + op_desc.set_op_debug_config(debug_config[op_desc.op_type]) + if all_soc_key in debug_config: + op_desc.set_op_debug_config(debug_config[all_soc_key]) + if op_desc.op_type in super_config: + op_desc.set_op_super_config(super_config[op_desc.op_type]) + if op_desc.op_type in tiling_key_info: + op_desc.set_tiling_key(tiling_key_info[op_desc.op_type]) + if all_soc_key in tiling_key_info: + op_desc.set_tiling_key(tiling_key_info[all_soc_key]) + op_desc.gen_input_json(auto_gen_path_dir) + + +def parse_args(argv): + """Command line parameter parsing""" + parser = argparse.ArgumentParser() + parser.add_argument("argv", nargs="+") + parser.add_argument("--opc-config-file", nargs="?", const="", default="") + return parser.parse_args(argv) + + +if __name__ == "__main__": + args = parse_args(sys.argv) + if len(args.argv) <= 3: + raise RuntimeError("arguments must greater than 3") + gen_bin_param_file(args.argv[1], args.argv[2], args.argv[3], opc_config_file=args.opc_config_file) diff --git a/csrc/ascend/cmake/scripts/util/ascendc_gen_options.py b/csrc/ascend/cmake/scripts/util/ascendc_gen_options.py new file mode 100644 index 000000000..9d6fd87d2 --- /dev/null +++ b/csrc/ascend/cmake/scripts/util/ascendc_gen_options.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import os +import stat +import sys + + +def write_options_to_file(file_name: str, options_str: str, op_type: str, compute_unit: str, split_char: str): + flags = os.O_WRONLY | os.O_CREAT + modes = stat.S_IWUSR | stat.S_IRUSR + try: + with os.fdopen(os.open(file_name, flags, modes), "a") as fd: + fd.write(op_type + split_char + compute_unit + split_char + options_str + "\n") + except Exception as err: + print("write compile options config file failed") + raise (err) + + +def gen_compile_options(compile_options_file: str, op_type: str, compute_unit: str, compile_options: list): + base_dir = os.path.dirname(compile_options_file) + opc_config_file = os.path.join(base_dir, "custom_opc_options.ini") + compile_opt = [] + opc_debug_config = [] + opc_tiling_keys = "" + for opts in compile_options: + if "oom" in opts: + if opts == "--oom": + opc_debug_config.append("oom") + else: + raise RuntimeError(f"Unknown oom option format {opts}") + elif "--save-temp-files" in opts: + opc_debug_config.append("dump_cce") + elif opts.startswith("--op_relocatable_kernel_binary") or opts.startswith("--op_super_kernel_options"): + opc_debug_config.append(opts) + elif "--tiling_key" in opts: + keys = opts.strip().split("=")[1].split(",") + keys_str = ";".join([key for key in keys]) + opc_tiling_keys = keys_str + else: + compile_opt.append(opts) + if len(compile_opt) > 0: + options_str = ";".join([opt for opt in compile_opt]) + write_options_to_file(compile_options_file, options_str, op_type, compute_unit, ",") + opc_config_str = "" + if opc_debug_config: + opc_config_str = "--op_debug_config=" + ";".join([opt for opt in opc_debug_config]) + if len(opc_tiling_keys) > 0: + if opc_config_str != "": + opc_config_str += "@" + opc_config_str += "--tiling_key=" + opc_tiling_keys + + if opc_config_str != "": + write_options_to_file(opc_config_file, opc_config_str, op_type, compute_unit, "@") + + +if __name__ == "__main__": + if len(sys.argv) < 4: + raise RuntimeError("arguments must greater than 4") + compute_soc = "" + comp_options = [] + for i in range(len(sys.argv) - 3): + if sys.argv[i + 3].upper().startswith("ASCEND"): + compute_soc += sys.argv[i + 3] + ";" + else: + comp_options.append(sys.argv[i + 3]) + if compute_soc != "": + compute_soc = compute_soc[0:-1] + gen_compile_options(sys.argv[1], sys.argv[2], compute_soc, comp_options) diff --git a/csrc/ascend/cmake/scripts/util/ascendc_impl_build.py b/csrc/ascend/cmake/scripts/util/ascendc_impl_build.py new file mode 100644 index 000000000..5e847f624 --- /dev/null +++ b/csrc/ascend/cmake/scripts/util/ascendc_impl_build.py @@ -0,0 +1,782 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import argparse +import datetime +import glob +import json +import os +import sys + +import const_var +import opdesc_parser +import regex as re + +PYF_PATH = os.path.dirname(os.path.realpath(__file__)) + +IMPL_HEAD = '''#!/usr/bin/env python +# -*- coding: UTF-8 -*- +""" +Copyright (c) Huawei Technologies Co., Ltd. {}-{}. All rights reserved. +""" + +import regex as re +import os, sys +import ctypes +import json +import shutil +from tbe.common.platform import get_soc_spec +from tbe.common.utils import para_check +from tbe.tikcpp import compile_op, replay_op, check_op_cap, generalize_op_params, get_code_channel, OpInfo +from tbe.tikcpp.compile_op import CommonUtility, AscendCLogLevel +from tbe.common.buildcfg import get_default_build_config +from tbe.common.buildcfg import get_current_build_config +import tbe.common.register as tbe_register +PYF_PATH = os.path.dirname(os.path.realpath(__file__)) + +DTYPE_MAP = {{"float32": ["DT_FLOAT", "float"], + "float16": ["DT_FLOAT16", "half"], + "int8": ["DT_INT8", "int8_t"], + "int16": ["DT_INT16", "int16_t"], + "int32": ["DT_INT32", "int32_t"], + "int64": ["DT_INT64", "int64_t"], + "uint1": ["DT_UINT1", "uint1b_t"], + "uint8": ["DT_UINT8", "uint8_t"], + "uint16": ["DT_UINT16", "uint16_t"], + "uint32": ["DT_UINT32", "uint32_t"], + "uint64": ["DT_UINT64", "uint64_t"], + "bool": ["DT_BOOL", "bool"], + "double": ["DT_DOUBLE", "double"], + "dual": ["DT_DUAL", "unknown"], + "dual_sub_int8": ["DT_DUAL_SUB_INT8", "unknown"], + "dual_sub_uint8": ["DT_DUAL_SUB_UINT8", "unknown"], + "string": ["DT_STRING", "unknown"], + "complex32": ["DT_COMPLEX32", "complex32"], + "complex64": ["DT_COMPLEX64", "complex64"], + "complex128": ["DT_COMPLEX128", "unknown"], + "qint8": ["DT_QINT8", "unknown"], + "qint16": ["DT_QINT16", "unknown"], + "qint32": ["DT_QINT32", "unknown"], + "quint8": ["DT_QUINT8", "unknown"], + "quint16": ["DT_QUINT16", "unknown"], + "resource": ["DT_RESOURCE", "unknown"], + "string_ref": ["DT_STRING_REF", "unknown"], + "int4": ["DT_INT4", "int4b_t"], + "bfloat16": ["DT_BF16", "bfloat16_t"], + "float8_e5m2": ["DT_FLOAT8_E5M2", "fp8_e5m2_t"], + "float8_e4m3fn": ["DT_FLOAT8_E4M3FN", "fp8_e4m3fn_t"], + "hifloat8":["DT_HIFLOAT8", "hifloat8_t"], + "float8_e8m0":["DT_FLOAT8_E8M0", "fp8_e8m0_t"], + "float4_e2m1":["DT_FLOAT4_E2M1", "fp4x2_e2m1_t"], + "float4_e1m2":["DT_FLOAT4_E1M2", "fp4x2_e1m2_t"], + "int2": ["DT_INT2", "int2b_t"]}} + +def add_dtype_fmt_option_single(x, x_n, is_ref: bool = False): + options = [] + x_fmt = x.get("format") + x_dtype = x.get("dtype") + x_n_in_kernel = x_n + '_REF' if is_ref else x_n + options.append("-DDTYPE_{{n}}={{t}}".format(n=x_n_in_kernel, t=DTYPE_MAP.get(x_dtype)[1])) + options.append("-DORIG_DTYPE_{{n}}={{orig_t}}".format(n=x_n_in_kernel, orig_t=DTYPE_MAP.get(x_dtype)[0])) + options.append("-DFORMAT_{{n}}=FORMAT_{{f}}".format(n=x_n_in_kernel, f=x_fmt)) + return options + +def get_dtype_fmt_options(__inputs__, __outputs__): + options = [] + input_names = {} + output_names = {} + unique_param_name_set = set() + for idx, x in enumerate(__inputs__): + if x is None: + continue + x_n = input_names[idx].upper() + unique_param_name_set.add(x_n) + options += add_dtype_fmt_option_single(x, x_n) + + for idx, x in enumerate(__outputs__): + if x is None: + continue + x_n = output_names[idx].upper() + if x_n in unique_param_name_set: + options += add_dtype_fmt_option_single(x, x_n, True) + else: + options += add_dtype_fmt_option_single(x, x_n) + return options + +def load_dso(so_path): + try: + ctypes.CDLL(so_path) + except OSError as error : + CommonUtility.print_compile_log("", error, AscendCLogLevel.LOG_ERROR) + raise RuntimeError("cannot open %s" %(so_path)) + else: + msg = "load so succ " + so_path + CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO) + +def get_shortsoc_compile_option(compile_option_list: list, shortsoc:str): + compile_options = [] + if shortsoc in compile_option_list: + compile_options.extend(compile_option_list[shortsoc]) + if '__ALLSOC__' in compile_option_list: + compile_options.extend(compile_option_list['__ALLSOC__']) + return compile_options + +def get_kernel_source(src_file, dir_snake, dir_ex): + src = os.path.join(PYF_PATH, "op_kernel", src_file) + if os.path.exists(src): + return src + src = os.path.join(PYF_PATH, "..", "ascendc", dir_snake, "op_kernel", src_file) + if os.path.exists(src): + return src + src_ex = os.path.join(PYF_PATH, "..", "ascendc", dir_ex, "op_kernel", src_file) + if os.path.exists(src_ex): + return src_ex + src_ex = os.path.join(PYF_PATH, "..", "ascendc", dir_ex, src_file) + if os.path.exists(src_ex): + return src_ex + src = os.environ.get('BUILD_KERNEL_SRC') + if src and os.path.exists(src): + return src + src = os.path.join(PYF_PATH, "..", "ascendc", dir_snake, src_file) + if os.path.exists(src): + return src + src = os.path.join(PYF_PATH, src_file) + if os.path.exists(src): + return src + src = os.path.join(PYF_PATH, "..", "ascendc", dir_snake, dir_snake + ".cpp") + if os.path.exists(src): + return src + src = os.path.join(PYF_PATH, "..", "ascendc", dir_ex, dir_ex + ".cpp") + if os.path.exists(src): + return src + src = os.path.join(PYF_PATH, "..", "ascendc", os.path.splitext(src_file)[0], src_file) + if os.path.exists(src): + return src + return src_ex + +''' + +IMPL_API = """ +@tbe_register.register_operator("{}", trans_bool_to_s8=False) +@para_check.check_op_params({}) +def {}({}, kernel_name="{}"{}): +{} + if get_current_build_config("enable_op_prebuild"): + return + __inputs__, __outputs__, __attrs__ = _build_args({}) + options = get_dtype_fmt_options(__inputs__, __outputs__) + options += ["-x", "cce"] + bisheng = os.environ.get('BISHENG_REAL_PATH') + if bisheng is None: + bisheng = shutil.which("bisheng") + if bisheng != None: + bisheng_path = os.path.dirname(bisheng) + tikcpp_path = os.path.realpath(os.path.join(bisheng_path, "..", "..", "tikcpp")) + else: + toolkit_path = os.environ.get('ASCEND_HOME_PATH', '/usr/local/Ascend/latest') + tikcpp_path = os.path.realpath(os.path.join(toolkit_path, "compiler", "tikcpp")) + options.append("-I" + tikcpp_path) + options.append("-I" + os.path.join(tikcpp_path, "..", "..", "include")) + options.append("-I" + os.path.join(tikcpp_path, "tikcfw")) + options.append("-I" + os.path.join(tikcpp_path, "tikcfw", "impl")) + options.append("-I" + os.path.join(tikcpp_path, "tikcfw", "interface")) + options.append("-I" + os.path.join(tikcpp_path, "..", "ascendc", "act")) + options.append("-I" + os.path.join(PYF_PATH, "..", "ascendc", "common")) + toolkit_path = os.environ.get('ASCEND_HOME_PATH') + if toolkit_path is None: + toolkit_path = os.path.realpath("/usr/local/Ascend/latest/") + options.append("-I" + toolkit_path + os.path.join("/", os.uname().machine +"-linux", "asc", "atcos")) + op_common_path = os.path.realpath(toolkit_path + "/pkg_inc/op_common/") + options.append("-I" + op_common_path) + if "impl_mode" in locals(): + if impl_mode == "high_performance": + options.append("-DHIGH_PERFORMANCE=1") + elif impl_mode == "high_precision": + options.append("-DHIGH_PRECISION=1") + elif "high_precision" in impl_mode and "high_performance" in impl_mode: + options.append("-DHIGH_PRECISION=1 -DHIGH_PERFORMANCE=1") + if get_current_build_config("enable_deterministic_mode") == 1: + options.append("-DDETERMINISTIC_MODE=1") + else: + options.append("-DDETERMINISTIC_MODE=0") + ascendc_api_version_header_path = os.path.join(tikcpp_path, "tikcfw/lib/ascendc_api_version.h") + if os.path.exists(ascendc_api_version_header_path): + with open(ascendc_api_version_header_path, "r") as ascendc_api_version_file: + ascendc_api_version = re.findall(r"#define ASCENDC_API_VERSION (\d+)", ascendc_api_version_file.read()) + if ascendc_api_version: + options.append(f"-DASCENDC_API_VERSION={{ascendc_api_version[0]}}") + custom_compile_options = {}, + custom_all_compile_options = {}, + soc_version = get_soc_spec("SOC_VERSION") + soc_short = get_soc_spec("SHORT_SOC_VERSION").lower() + custom_compile_options_soc = get_shortsoc_compile_option(custom_compile_options[0], soc_short) + custom_all_compile_options_soc = get_shortsoc_compile_option(custom_all_compile_options[0], soc_short) + options += custom_all_compile_options_soc + options += custom_compile_options_soc + + origin_func_name = "{}" + ascendc_src_dir_ex = "{}" + ascendc_src_dir = "{}" + ascendc_src_file = "{}" + src = get_kernel_source(ascendc_src_file, ascendc_src_dir, ascendc_src_dir_ex) +""" + +REPLAY_OP_API = """ + msg = "start replay Ascend C Operator {}, kernel name is {}" + CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO) + tikreplay_codegen_path = tikcpp_path + "/tikreplaylib/lib" + tikreplay_stub_path = tikcpp_path + "/tikreplaylib/lib/" + soc_version + msg = "start load libtikreplaylib_codegen.so and libtikreplaylib_stub.so" + CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO) + codegen_so_path = tikreplay_codegen_path + "/libtikreplaylib_codegen.so" + replaystub_so_path = tikreplay_stub_path + "/libtikreplaylib_stub.so" + if PYF_PATH.endswith("dynamic"): + op_replay_path = os.path.join(PYF_PATH, "..", "..", "op_replay") + else: + op_replay_path = os.path.join(PYF_PATH, "..", "op_replay") + replayapi_so_path = os.path.join(op_replay_path, "libreplay_{}_" + soc_short + ".so") + load_dso(codegen_so_path) + load_dso(replaystub_so_path) + load_dso(replayapi_so_path) + op_type = "{}" + entry_obj = os.path.join(op_replay_path, "{}_entry_" + soc_short + ".o") + code_channel = get_code_channel(src, kernel_name, op_type, options) + op_info = OpInfo(kernel_name = kernel_name, op_type = op_type, inputs = __inputs__, outputs = __outputs__,\\ + attrs = __attrs__, impl_mode = impl_mode, param_type_dynamic = {}) + res, msg = replay_op(op_info, entry_obj, code_channel, src, options) + if not res: + print("call replay op failed for %s and get into call compile op" %(msg)) + compile_op(src, origin_func_name, op_info, options, code_channel, '{}') +""" + +COMPILE_OP_API = """ + msg = "start compile Ascend C Operator {}, kernel name is " + kernel_name + CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO) + op_type = "{}" + code_channel = get_code_channel(src, kernel_name, op_type, options) + op_info = OpInfo(kernel_name = kernel_name, op_type = op_type, inputs = __inputs__, outputs = __outputs__,\\ + attrs = __attrs__ {}, origin_inputs=[{}], origin_outputs = [{}],\\ + param_type_dynamic = {}, mc2_ctx = {}, param_type_list = {}, init_value_list = {},\\ + output_shape_depend_on_compute = {}) + compile_op(src, origin_func_name, op_info, options, code_channel, '{}', {}) +""" +COMPILE_OP_API_BUILT_IN = """ + msg = "start compile Ascend C Operator {}, kernel name is " + kernel_name + CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO) + op_type = "{}" + code_channel = get_code_channel(src, kernel_name, op_type, options) + op_info = OpInfo(kernel_name = kernel_name, op_type = op_type, inputs = __inputs__, outputs = __outputs__,\\ + attrs = __attrs__ {}, origin_inputs=[{}], origin_outputs = [{}],\\ + param_type_dynamic = {}, mc2_ctx = {}, param_type_list = {}, init_value_list = {},\\ + output_shape_depend_on_compute = {}) + + op_compile_option = '{}' + opp_path = os.environ.get('ASCEND_OPP_PATH') + dat_path = os.path.realpath(os.path.join(opp_path, "built-in", "op_impl", "ai_core", "tbe", "ascendc_impl.dat")) + if opp_path and os.path.exists(dat_path): + # dat file exists: built in hidden src file online compiling process. append vfs compile option in compile_op + abs_rel_kernel_src_path = "{}" + extend_options = {} + extend_options['opp_kernel_hidden_dat_path'] = dat_path + compile_op(abs_rel_kernel_src_path, origin_func_name, op_info, options, code_channel, op_compile_option,\\ + extend_options) + else: + raise RuntimeError("built-in opp compile, ascendc_impl.dat file path does not exist: %s" %(dat_path)) +""" +SUP_API = """ +def {}({}{}): + __inputs__, __outputs__, __attrs__ = _build_args({}) + ret_str = check_op_cap("{}", "{}", __inputs__, __outputs__, __attrs__) + ret_dict = json.loads(ret_str) + err_code = ret_dict.get("ret_code") + sup = "Unknown" + reason = "Unknown reason" + if err_code is not None: + if err_code == 0: + sup = "True" + reason = "" + elif err_code == 1: + sup = "False" + reason = ret_dict.get("reason") + else: + sup = "Unknown" + reason = ret_dict.get("reason") + return sup, reason +""" +CAP_API = """ +def {}({}{}): + __inputs__, __outputs__, __attrs__ = _build_args({}) + result = check_op_cap("{}", "{}", __inputs__, __outputs__, __attrs__) + return result.decode("utf-8") +""" +GLZ_API = """ +@tbe_register.register_param_generalization("{}") +def {}_generalization({}, generalize_config=None): + __inputs__, __outputs__, __attrs__ = _build_args({}) + ret_str = generalize_op_params("{}", __inputs__, __outputs__, __attrs__, generalize_config) + return [json.loads(ret_str)] +""" + +ATTR_DEFAULT = { + "bool": "False", + "int": "0", + "float": "0.0", + "list_int": "[]", + "list_float": "[]", + "list_bool": "[]", + "list_list_int": "[[]]", + "str": "", +} + + +def optype_snake(origin_str): + temp_str = origin_str[0].lower() + origin_str[1:] + new_str = re.sub(r"([A-Z])", r"_\1", temp_str).lower() + return new_str + + +def optype_snake_ex(s): + snake_case = "" + for i, c in enumerate(s): + if i == 0: + snake_case += c.lower() + elif c.isupper(): + if s[i - 1] != "_": + if not s[i - 1].isupper() or s[i - 1].isupper() and (i + 1) < len(s) and s[i + 1].islower(): + snake_case += "_" + snake_case += c.lower() + else: + snake_case += c + return snake_case + + +class AdpBuilder(opdesc_parser.OpDesc): + def __init__(self: any, op_type: str): + self.argsdefv = [] + self.op_compile_option: str = "{}" + super().__init__(op_type) + + def write_adapt(self: any, impl_path, path: str, op_compile_option_all: list = None): + self._build_paradefault() + if os.environ.get("BUILD_BUILTIN_OPP") != "1" and impl_path != "": + src_file = os.path.join(impl_path, self.op_file + ".cpp") + if not os.path.exists(src_file): + print(f"[ERROR]: operator: {self.op_file} source file: {src_file} does not found, please check.") + return + out_path = os.path.abspath(path) + if self.dynamic_shape and not out_path.endswith("dynamic"): + out_path = os.path.join(path, "dynamic") + os.makedirs(out_path, exist_ok=True) + adpfile = os.path.join(out_path, self.op_file + ".py") + self._gen_op_compile_option(op_compile_option_all) + with os.fdopen(os.open(adpfile, const_var.WFLAGS, const_var.WMODES), "w") as fd: + self._write_head(fd) + self._write_argparse(fd) + self._get_impl_mode() + self._write_impl(fd, impl_path) + if self.op_chk_support: + self._write_cap("check_supported", fd) + self._write_cap("get_op_support_info", fd) + if self.op_fmt_sel: + self._write_cap("op_select_format", fd) + self._write_cap("get_op_specific_info", fd) + if self.op_range_limit == "limited" or self.op_range_limit == "dynamic": + self._write_glz(fd) + + def _gen_op_compile_option(self: any, op_compile_option_all: list = None): + if op_compile_option_all is not None: + if self.op_type in op_compile_option_all: + self.op_compile_option = op_compile_option_all[self.op_type] + elif "__all__" in op_compile_option_all: + self.op_compile_option = op_compile_option_all["__all__"] + + def _ip_argpack(self: any, default: bool = True) -> list: + args = [] + for i in range(len(self.input_name)): + arg = self.input_name[i] + if default and self.argsdefv[i] is not None: + arg += "=" + self.argsdefv[i] + args.append(arg) + return args + + def _op_argpack(self: any, default: bool = True) -> list: + args = [] + argidx = len(self.input_name) + for i in range(len(self.output_name)): + arg = self.output_name[i] + if default and self.argsdefv[i + argidx] is not None: + arg += "=" + self.argsdefv[i + argidx] + args.append(arg) + return args + + def _attr_argpack(self: any, default: bool = True) -> list: + args = [] + argidx = len(self.input_name) + len(self.output_name) + for i in range(len(self.attr_list)): + att = self.attr_list[i] + arg = att + if default and self.argsdefv[i + argidx] is not None: + if self.attr_val.get(att).get("type") == "str": + arg += '="' + self.argsdefv[i + argidx] + '"' + elif self.attr_val.get(att).get("type") == "bool": + arg += "=" + self.argsdefv[i + argidx].capitalize() + elif self.attr_val.get(att).get("type") == "list_bool": + arg += ( + "=" + + "[" + + ", ".join( + word.strip().capitalize() for word in self.argsdefv[i + argidx].strip("[]").split(",") + ) + + "]" + ) + else: + arg += "=" + self.argsdefv[i + argidx] + args.append(arg) + return args + + def _build_paralist(self: any, default: bool = True) -> str: + args = [] + args.extend(self._ip_argpack(default)) + args.extend(self._op_argpack(default)) + args.extend(self._attr_argpack(default)) + return ", ".join(args) + + def _io_parachk(self: any, types: list, type_name: str) -> list: + chk = [] + for iot in types: + if iot == "optional": + ptype = "OPTION" + else: + ptype = iot.upper() + chk.append("para_check.{}_{}".format(ptype, type_name)) + return chk + + def _attr_parachk(self: any) -> list: + chk = [] + for att in self.attr_list: + att_type = self.attr_val.get(att).get("type").upper() + chk.append("para_check.{}_ATTR_{}".format("OPTION", att_type)) + return chk + + def _build_parachk(self: any) -> str: + chk = [] + chk.extend(self._io_parachk(self.input_type, "INPUT")) + chk.extend(self._io_parachk(self.output_type, "OUTPUT")) + chk.extend(self._attr_parachk()) + chk.append("para_check.KERNEL_NAME") + return ", ".join(chk) + + def _build_virtual(self: any) -> str: + virt_exp = [] + for index in range(len(self.input_name)): + if self.input_virt.get(index) is None: + continue + val = [] + val.append('"param_name":"{}"'.format(self.input_name[index])) + val.append('"index":{}'.format(index)) + val.append('"dtype":"{}"'.format(self.input_dtype[index].split(",")[0])) + val.append('"format":"{}"'.format(self.input_fmt[index].split(",")[0])) + val.append('"ori_format":"{}"'.format(self.input_fmt[index].split(",")[0])) + val.append('"paramType":"optional"') + val.append('"shape":[1]') + val.append('"ori_shape":[1]') + virt_exp.append(" " + self.input_name[index] + " = {" + ",".join(val) + "}") + if len(virt_exp) > 0: + return "\n".join(virt_exp) + else: + return " # do ascendc build step" + + def _build_mc2_ctx(self: any): + if len(self.mc2_ctx) != 0: + return '["' + '", "'.join(self.mc2_ctx) + '"]' + return "[]" + + def _build_paradefault(self: any): + optional = False + argtypes = [] + argtypes.extend(self.input_type) + argtypes.extend(self.output_type) + for atype in argtypes: + if atype == "optional": + optional = True + if optional: + self.argsdefv.append("None") + else: + self.argsdefv.append(None) + for attr in self.attr_list: + atype = self.attr_val.get(attr).get("paramType") + if atype == "optional": + optional = True + attrval = self.attr_val.get(attr).get("defaultValue") + if attrval is not None: + optional = True + if atype == "bool": + attrval = attrval.capitalize() + elif atype == "str": + attrval = '"' + attrval + '"' + self.argsdefv.append(attrval) + continue + if optional: + self.argsdefv.append(ATTR_DEFAULT.get(self.attr_val.get(attr).get("type"))) + else: + self.argsdefv.append(None) + + def _write_head(self: any, fd: object): + now = datetime.datetime.now() + curr_year = now.year + former_year = curr_year - 1 + fd.write(IMPL_HEAD.format(former_year, curr_year, self.input_ori_name, self.output_ori_name)) + + def _write_argparse(self: any, fd: object): + args = self._build_paralist(False) + fd.write("def _build_args({}):\n".format(args)) + fd.write(" __inputs__ = []\n") + fd.write(" for arg in [{}]:\n".format(", ".join(self.input_name))) + fd.write(" if arg != None:\n") + fd.write(" if isinstance(arg, (list, tuple)):\n") + fd.write(" if len(arg) == 0:\n") + fd.write(" continue\n") + fd.write(" __inputs__.append(arg[0])\n") + fd.write(" else:\n") + fd.write(" __inputs__.append(arg)\n") + fd.write(" else:\n") + fd.write(" __inputs__.append(arg)\n") + fd.write(" __outputs__ = []\n") + fd.write(" for arg in [{}]:\n".format(", ".join(self.output_name))) + fd.write(" if arg != None:\n") + fd.write(" if isinstance(arg, (list, tuple)):\n") + fd.write(" if len(arg) == 0:\n") + fd.write(" continue\n") + fd.write(" __outputs__.append(arg[0])\n") + fd.write(" else:\n") + fd.write(" __outputs__.append(arg)\n") + fd.write(" else:\n") + fd.write(" __outputs__.append(arg)\n") + fd.write(" __attrs__ = []\n") + for attr in self.attr_list: + fd.write(" if {} != None:\n".format(attr)) + fd.write(" attr = {}\n") + fd.write(' attr["name"] = "{}"\n'.format(attr)) + fd.write(' attr["dtype"] = "{}"\n'.format(self.attr_val.get(attr).get("type"))) + fd.write(' attr["value"] = {}\n'.format(attr)) + fd.write(" __attrs__.append(attr)\n") + fd.write(" return __inputs__, __outputs__, __attrs__\n") + + def _get_kernel_source(self: any, kernel_src_dir, src_file, dir_snake, dir_ex): + src = os.path.join(kernel_src_dir, "op_kernel", src_file) + if os.path.exists(src): + return src + src = os.path.join(kernel_src_dir, "..", "ascendc", dir_snake, "op_kernel", src_file) + if os.path.exists(src): + return src + src_ex = os.path.join(kernel_src_dir, "..", "ascendc", dir_ex, "op_kernel", src_file) + if os.path.exists(src_ex): + return src_ex + src_ex = os.path.join(kernel_src_dir, dir_ex, src_file) + if os.path.exists(src_ex): + return src_ex + src = os.environ.get("BUILD_KERNEL_SRC") + if src and os.path.exists(src): + return src + src = os.path.join(kernel_src_dir, dir_snake, src_file) + if os.path.exists(src): + return src + src = os.path.join(kernel_src_dir, src_file) + if os.path.exists(src): + return src + src = os.path.join(kernel_src_dir, dir_snake, dir_snake + ".cpp") + if os.path.exists(src): + return src + src = os.path.join(kernel_src_dir, dir_ex, dir_ex + ".cpp") + if os.path.exists(src): + return src + src = os.path.join(kernel_src_dir, os.path.splitext(src_file)[0], src_file) + if os.path.exists(src): + return src + return src_ex + + def _get_impl_mode(self: any): + op_compile_options = json.loads(self.op_compile_option) + if "impl_mode" in op_compile_options: + if op_compile_options["impl_mode"] == "": + self.impl_mode = "" + self.impl_mode_op_info = "" + del op_compile_options["impl_mode"] + self.op_compile_option = json.dumps(op_compile_options) + else: + self.impl_mode = ", impl_mode ='" + op_compile_options["impl_mode"] + "'" + self.impl_mode_op_info = ", impl_mode ='" + op_compile_options["impl_mode"] + "'" + else: + self.impl_mode = ', impl_mode = ""' + self.impl_mode_op_info = ", impl_mode = impl_mode" + + def _write_impl(self: any, fd: object, impl_path: str = ""): + argsdef = self._build_paralist() + argsval = self._build_paralist(False) + pchk = self._build_parachk() + if len(self.kern_name) > 0: + kern_name = self.kern_name + else: + kern_name = self.op_intf + src = self.op_file + ".cpp" + virt_exprs = self._build_virtual() + fd.write( + IMPL_API.format( + self.op_type, + pchk, + self.op_intf, + argsdef, + kern_name, + self.impl_mode, + virt_exprs, + argsval, + self.custom_compile_options, + self.custom_all_compile_options, + self.op_intf, + optype_snake_ex(self.op_type), + optype_snake(self.op_type), + src, + ) + ) + if self.op_replay_flag: + fd.write( + REPLAY_OP_API.format( + self.op_type, + kern_name, + self.op_file, + self.op_type, + self.op_file, + self.param_type_dynamic, + self.op_compile_option, + ) + ) + else: + value_depend_obj = {key: value for key, value in self.input_value_depend.items()} + extend_opt = {"valueDepend": value_depend_obj} + if os.environ.get("BUILD_BUILTIN_OPP") == "1": + relative_kernel_src_path = os.path.realpath( + self._get_kernel_source(impl_path, src, optype_snake(self.op_type), optype_snake_ex(self.op_type)) + ) + # to match src path in .dat file system, turn relative path into absolute path + abs_rel_kernel_src_path = os.path.join("/", os.path.relpath(relative_kernel_src_path, impl_path)) + + # compiling hidden src file requires src path before packaging .dat file, + # hard code such src path to .py + fd.write( + COMPILE_OP_API_BUILT_IN.format( + self.op_type, + self.op_type, + self.impl_mode_op_info, + ", ".join(self.input_name), + ", ".join(self.output_name), + self.param_type_dynamic, + self._build_mc2_ctx(), + self.input_type + self.output_type, + self.output_init_value, + self.output_shape_depend_on_compute, + self.op_compile_option, + abs_rel_kernel_src_path, + repr(extend_opt), + ) + ) + else: + fd.write( + COMPILE_OP_API.format( + self.op_type, + self.op_type, + self.impl_mode_op_info, + ", ".join(self.input_name), + ", ".join(self.output_name), + self.param_type_dynamic, + self._build_mc2_ctx(), + self.input_type + self.output_type, + self.output_init_value, + self.output_shape_depend_on_compute, + self.op_compile_option, + repr(extend_opt), + ) + ) + + def _write_cap(self: any, cap_name: str, fd: object): + argsdef = self._build_paralist() + argsval = self._build_paralist(False) + if cap_name == "check_supported": + fd.write(SUP_API.format(cap_name, argsdef, self.impl_mode, argsval, cap_name, self.op_type)) + else: + fd.write(CAP_API.format(cap_name, argsdef, self.impl_mode, argsval, cap_name, self.op_type)) + + def _write_glz(self: any, fd: object): + argsdef = self._build_paralist() + argsval = self._build_paralist(False) + fd.write(GLZ_API.format(self.op_type, self.op_intf, argsdef, argsval, self.op_type)) + + +def write_scripts(cfgfile: str, cfgs: dict, dirs: dict, ops: list = None, op_compile_option: list = None): + batch_lists = cfgs.get(const_var.REPLAY_BATCH).split(";") + iterator_lists = cfgs.get(const_var.REPLAY_ITERATE).split(";") + file_map = {} + op_descs = opdesc_parser.get_op_desc( + cfgfile, batch_lists, iterator_lists, AdpBuilder, ops, dirs.get(const_var.AUTO_GEN_DIR) + ) + for op_desc in op_descs: + op_desc.write_adapt(dirs.get(const_var.CFG_IMPL_DIR), dirs.get(const_var.CFG_OUT_DIR), op_compile_option) + file_map[op_desc.op_type] = op_desc.op_file + return file_map + + +class OpFileNotExistsError(Exception): + """File does not exist error.""" + + def __str__(self) -> str: + return f"File aic-*-ops-info.ini does not exist in directory {super().__str__()}" + + +def get_ops_info_files(opsinfo_dir: list[str]) -> list[str]: + """Get all ops info files.""" + ops_info_files = [] + for _dir in opsinfo_dir: + ops_info_files.extend(glob.glob(f"{_dir}/aic-*-ops-info.ini")) + return sorted(ops_info_files) + + +def parse_args(argv): + """Command line parameter parsing""" + parser = argparse.ArgumentParser() + parser.add_argument("argv", nargs="+") + parser.add_argument("--opsinfo-dir", nargs="*", default=None) + return parser.parse_args(argv) + + +if __name__ == "__main__": + args = parse_args(sys.argv) + + if len(args.argv) <= 6: + raise RuntimeError("arguments must greater equal than 6") + + rep_cfg = {} + rep_cfg[const_var.REPLAY_BATCH] = args.argv[2] + rep_cfg[const_var.REPLAY_ITERATE] = args.argv[3] + + cfg_dir = {} + cfg_dir[const_var.CFG_IMPL_DIR] = args.argv[4] + cfg_dir[const_var.CFG_OUT_DIR] = args.argv[5] + cfg_dir[const_var.AUTO_GEN_DIR] = args.argv[6] + + ops_infos = [] + if args.opsinfo_dir: + ops_infos.extend(get_ops_info_files(args.opsinfo_dir)) + if not ops_infos: + raise OpFileNotExistsError(args.opsinfo_dir) + else: + ops_infos.append(args.argv[1]) + + for ops_info in ops_infos: + write_scripts(cfgfile=ops_info, cfgs=rep_cfg, dirs=cfg_dir) diff --git a/csrc/ascend/cmake/scripts/util/ascendc_ops_config.py b/csrc/ascend/cmake/scripts/util/ascendc_ops_config.py new file mode 100644 index 000000000..e97425c1f --- /dev/null +++ b/csrc/ascend/cmake/scripts/util/ascendc_ops_config.py @@ -0,0 +1,360 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import argparse +import glob +import json +import os +import sys +from typing import NamedTuple + +import const_var + + +class OpConfig(NamedTuple): + op_type: str + support_info: dict + core_type: str + task_ration: str + obj_file: str + + +def load_json(json_file: str): + with open(json_file, encoding="utf-8") as file: + json_content = json.load(file) + return json_content + + +def get_specified_suffix_file(root_dir, suffix): + specified_suffix = os.path.join(root_dir, "**/*{}".format(suffix)) + all_suffix_files = glob.glob(specified_suffix, recursive=True) + return sorted(all_suffix_files) + + +def add_dict_key(dict_to_add, key, value): + if value is None: + return + dict_to_add[key] = value + + +def correct_format_mode(format_mode): + if format_mode == "FormatDefault": + return "nd_agnostic" + if format_mode == "FormatAgnostic": + return "static_nd_agnostic" + if format_mode == "FormatFixed": + return "normal" + return format_mode + + +def get_input_or_output_config(in_or_out): + param_dict = {} + name = in_or_out.get("name") + index = in_or_out.get("index") + param_type = in_or_out.get("paramType") + + format_match_mode = in_or_out.get("format_match_mode") + format_mode = correct_format_mode(format_match_mode) + + dtype_mode = in_or_out.get("dtype_match_mode") + if dtype_mode == "DtypeByte": + dtype_mode = "bit" + + add_dict_key(param_dict, "name", name) + add_dict_key(param_dict, "index", index) + add_dict_key(param_dict, "paramType", param_type) + add_dict_key(param_dict, "dtypeMode", dtype_mode) + add_dict_key(param_dict, "formatMode", format_mode) + return param_dict + + +def get_inputs_or_outputs_config(inputs_or_outputs): + if inputs_or_outputs is None: + return None + inputs_or_outputs_list = [] + + for in_or_out in inputs_or_outputs: + if isinstance(in_or_out, dict): + dict_param_config = get_input_or_output_config(in_or_out) + inputs_or_outputs_list.append(dict_param_config) + elif isinstance(in_or_out, list): + param_info = in_or_out[0] + list_param_config = get_input_or_output_config(param_info) + tmp_list = [list_param_config] + inputs_or_outputs_list.append(tmp_list) + return inputs_or_outputs_list + + +def gen_attrs_config(attrs): + attrs_list = [] + for attr in attrs: + attrs_dict = {} + name = attr.get("name") + mode = attr.get("mode") + add_dict_key(attrs_dict, "name", name) + add_dict_key(attrs_dict, "mode", mode) + attrs_list.append(attrs_dict) + return attrs_list + + +def get_params_config(support_info): + params_dict = {} + + inputs = support_info.get("inputs") + inputs_list = get_inputs_or_outputs_config(inputs) + params_dict["inputs"] = inputs_list + + outputs = support_info.get("outputs") + outputs_list = get_inputs_or_outputs_config(outputs) + params_dict["outputs"] = outputs_list + + attrs = support_info.get("attrs") + if attrs is not None: + attrs_list = gen_attrs_config(attrs) + params_dict["attrs"] = attrs_list + + return params_dict + + +def add_simplified_config(op_info, binary_info_config, config): + simplified_key = op_info.support_info.get("simplifiedKey") + + json_path = op_info.obj_file.split(".")[0] + ".json" + + simple_cfg = config.get(binary_info_config) + op_cfg = simple_cfg.get(op_info.op_type) + if not op_cfg: + op_cfg = {"dynamicRankSupport": True} + + simplified_key_mode = op_info.support_info.get("simplifiedKeyMode") + add_dict_key(op_cfg, "simplifiedKeyMode", simplified_key_mode) + + optional_input_mode = op_info.support_info.get("optionalInputMode") + optional_output_mode = op_info.support_info.get("optionalOutputMode") + add_dict_key(op_cfg, "optionalInputMode", optional_input_mode) + if optional_output_mode is not None: + add_dict_key(op_cfg, "optionalOutputMode", optional_output_mode) + + params_info = get_params_config(op_info.support_info) + op_cfg["params"] = params_info + op_cfg["binaryList"] = [] + simple_cfg[op_info.op_type] = op_cfg + + bin_list = op_cfg.get("binaryList") + if op_info.core_type == 0 and op_info.task_ration == "tilingKey": + bin_list.append( + { + "coreType": op_info.core_type, + "simplifiedKey": simplified_key, + "multiKernelType": 1, + "binPath": op_info.obj_file, + "jsonPath": json_path, + } + ) + else: + bin_list.append( + { + "coreType": op_info.core_type, + "simplifiedKey": simplified_key, + "binPath": op_info.obj_file, + "jsonPath": json_path, + } + ) + + +def add_op_config(op_file, bin_info, config): + op_cfg = config.get(op_file) + if not op_cfg: + op_cfg = {"binList": []} + config[op_file] = op_cfg + op_cfg.get("binList").append(bin_info) + + +def gen_ops_config(json_file, soc, binary_info_config, config): + core_type_map = {"MIX": 0, "AiCore": 1, "VectorCore": 2, "MIX_AICORE": 3, "MIX_VECTOR_CORE": 4, "MIX_AIV": 4} + contents = load_json(json_file) + if ("binFileName" not in contents) or ("supportInfo" not in contents): + return + json_base_name = os.path.basename(json_file) + op_dir = os.path.basename(os.path.dirname(json_file)) + + support_info = contents.get("supportInfo") + bin_name = contents.get("binFileName") + bin_suffix = contents.get("binFileSuffix") + core_type = contents.get("coreType") + task_ration = contents.get("taskRation") + core_type = core_type_map.get(core_type, -1) + if core_type == -1 and soc != "ascend310b": + raise Exception("[ERROR]: must set coreType in json when soc version is {soc}.") + + bin_file_name = bin_name + bin_suffix + op_type = bin_name.split("_")[0] + op_file = op_dir + ".json" + bin_info = {} + + add_dict_key(bin_info, "implMode", support_info.get("implMode")) + add_dict_key(bin_info, "int64Mode", support_info.get("int64Mode")) + add_dict_key(bin_info, "simplifiedKeyMode", support_info.get("simplifiedKeyMode")) + + simplified_key = support_info.get("simplifiedKey") + if simplified_key is not None: + bin_info["simplifiedKey"] = simplified_key + obj_file = os.path.join(soc, op_dir, bin_file_name) + op_info = OpConfig( + op_type=op_type, + support_info=support_info, + core_type=core_type, + task_ration=task_ration, + obj_file=obj_file, + ) + add_simplified_config(op_info, binary_info_config, config) + + add_dict_key(bin_info, "dynamicParamMode", support_info.get("dynamicParamMode")) + bin_info["staticKey"] = support_info.get("staticKey") + bin_info["inputs"] = support_info.get("inputs") + bin_info["outputs"] = support_info.get("outputs") + if support_info.get("attrs"): + bin_info["attrs"] = support_info.get("attrs") + + add_dict_key(bin_info, "opMode", support_info.get("opMode")) + add_dict_key(bin_info, "optionalInputMode", support_info.get("optionalInputMode")) + add_dict_key(bin_info, "deterministic", support_info.get("deterministic")) + if support_info.get("optionalOutputMode") is not None: + add_dict_key(bin_info, "optionalOutputMode", support_info.get("optionalOutputMode")) + + bin_info["binInfo"] = {"jsonFilePath": os.path.join(soc, op_dir, json_base_name)} + add_op_config(op_file, bin_info, config) + + +def check_single_op_is_void(root_dir): + for root, dirs, _ in os.walk(root_dir): + for sub_dir in dirs: + dir_path = os.path.join(root, sub_dir) + if len(os.listdir(dir_path)) == 0: + print(f"[ERROR] op {sub_dir}: not any obj compile success") + sys.exit(1) + + +def write_jsons(out_dir, file_list, config): + for json_name in file_list: + json_file = os.path.join(out_dir, json_name) + with os.fdopen(os.open(json_file, const_var.WFLAGS, const_var.WMODES), "w") as fd: + json.dump(config.get(json_name), fd, indent=" ") + + +def generate_operator_cfg_file(json_files, binary_info_config, soc, out_dir, gen_json_status): + if not json_files: + return + + if gen_json_status == "not_generated": + return + + json_files.sort() + config = {binary_info_config: {}} + for _json in json_files: + gen_ops_config(_json, soc, binary_info_config, config) + + if gen_json_status == "single_json": + file_list = [json_file for json_file in config if json_file != binary_info_config] + elif gen_json_status == "summary_json": + file_list = [binary_info_config] + else: + file_list = config.keys() + + write_jsons(out_dir, file_list, config) + + +def gen_all_config(root_dir, soc, out_dir, skip_binary_info_config, op_range="all"): + if op_range != "relocatable": + check_single_op_is_void(root_dir) + all_json_files = get_specified_suffix_file(root_dir, ".json") + relocatable_json_files = get_specified_suffix_file(root_dir, "_relocatable.json") + normal_json_files = list(set(all_json_files) - set(relocatable_json_files)) + os.makedirs(out_dir, exist_ok=True) + + if op_range != "relocatable": + for _json in all_json_files: + file_path = soc + _json.split(soc, maxsplit=1)[1] + with open(_json, "r+") as f: + data = json.load(f) + data["filePath"] = file_path + f.seek(0) + json.dump(data, f, indent=" ") + f.truncate() + + if skip_binary_info_config: + gen_normale_json = "single_json" + gen_relocatable_json = "not_generated" + else: + gen_normale_json = "all_json" + gen_relocatable_json = "summary_json" + + # normal kernel + if op_range == "all" or op_range == "normal": + binary_info_config = "binary_info_config.json" + generate_operator_cfg_file(normal_json_files, binary_info_config, soc, out_dir, gen_normale_json) + + # relocatable kernel + if op_range == "all" or op_range == "relocatable": + binary_info_config = "relocatable_kernel_info_config.json" + generate_operator_cfg_file(relocatable_json_files, binary_info_config, soc, out_dir, gen_relocatable_json) + + +# Parse multiple soc_versions ops in single path. +def gen_all_soc_config(all_path): + soc_roots = glob.glob(os.path.join(all_path, "ascend*")) + + for soc_root in soc_roots: + soc = os.path.basename(soc_root) + gen_all_config(soc_root, soc, soc_root, True) + cfg_files = glob.glob(os.path.join(soc_root, "*.json")) + cfg_path = os.path.join(all_path, "config", soc) + os.makedirs(cfg_path, exist_ok=True) + for cfg_file in cfg_files: + new_file = os.path.join(cfg_path, os.path.basename(cfg_file)) + os.rename(cfg_file, new_file) + + +def args_prase(): + parser = argparse.ArgumentParser() + parser.add_argument("-p", "--path", nargs="?", required=True, help="Parse the path of the json file.") + + parser.add_argument("-s", "--soc", nargs="?", required=True, help="Parse the soc_version of ops.") + + parser.add_argument("-o", "--out", nargs="?", help="Output directory.") + + parser.add_argument( + "--skip-binary-info-config", action="store_true", help="binary_info_config.json file is not parsed." + ) + + parser.add_argument( + "--op-range", + type=str, + choices=["all", "normal", "relocatable"], + default="all", + help="all operators/normal operators/relocatable operators.", + ) + + return parser.parse_args() + + +def main(): + args = args_prase() + if args.out is None: + out_dir = args.path + else: + out_dir = args.out + + gen_all_config(args.path, args.soc, out_dir, args.skip_binary_info_config, args.op_range) + + +if __name__ == "__main__": + main() diff --git a/csrc/ascend/cmake/scripts/util/const_var.py b/csrc/ascend/cmake/scripts/util/const_var.py new file mode 100644 index 000000000..d4afd9718 --- /dev/null +++ b/csrc/ascend/cmake/scripts/util/const_var.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import os +import stat + +REPLAY_BATCH = "batch" +REPLAY_ITERATE = "iterate" +CFG_IMPL_DIR = "impl_dir" +CFG_OUT_DIR = "out_dir" +AUTO_GEN_DIR = "auto_gen_dir" +WFLAGS = os.O_WRONLY | os.O_CREAT | os.O_TRUNC +WMODES = stat.S_IWUSR | stat.S_IRUSR +SOC_MAP_EXT = { + "ascend310p": "Ascend310P3", + "ascend310b": "Ascend310B1", + "ascend910": "Ascend910A", + "ascend910b": "Ascend910B1", + "ascend910_93": "Ascend910_9391", + "ascend610lite": "Ascend610Lite", + "ascend950": "Ascend950PR_9599", + "kirinx90": "KirinX90", +} +BIN_CMD = "opc $1 --main_func={fun} --input_param={param} --soc_version={soc} \ +--output=$2 --impl_mode={impl} --simplified_key_mode=0 --op_mode=dynamic\n" +SET_PLOG_LEVEL_ERROR = "export ASCEND_GLOBAL_LOG_LEVEL=3\n" +SET_PLOG_STDOUT = "export ASCEND_SLOG_PRINT_TO_STDOUT=1\n" +SRC_ENV = """ +while true; do + case "$1" in + --kernel-src=*) + export BUILD_KERNEL_SRC=$(echo "$1" | cut -d"=" -f2-) + shift + ;; + -*) + shift + ;; + *) + break + ;; + esac +done +""" +CHK_CMD = """ +if ! test -f $2/{res_file} ; then + echo "$2/{res_file} not generated!" + exit 1 +fi +""" +ATTR_DEF_VAL = { + "str": "", + "int": 0, + "float": 0.0, + "bool": False, + "list_bool": [], + "list_int": [], + "list_float": [], + "list_list_int": [[]], +} + + +def conv_soc_ver(ver: str): + return SOC_MAP_EXT.get(ver) diff --git a/csrc/ascend/cmake/scripts/util/gen_version_info.sh b/csrc/ascend/cmake/scripts/util/gen_version_info.sh new file mode 100644 index 000000000..b483e487f --- /dev/null +++ b/csrc/ascend/cmake/scripts/util/gen_version_info.sh @@ -0,0 +1,18 @@ +#!/bin/bash\n" +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + + +ascend_install_dir=$1 +gen_file_dir=$2 + +# create version.info +compiler_version=$(grep "Version" -w ${ascend_install_dir}/compiler/version.info | awk -F = '{print $2}') +echo "custom_opp_compiler_version=${compiler_version}" > ${gen_file_dir}/version.info \ No newline at end of file diff --git a/csrc/ascend/cmake/scripts/util/opdesc_parser.py b/csrc/ascend/cmake/scripts/util/opdesc_parser.py new file mode 100644 index 000000000..ea3aec6d5 --- /dev/null +++ b/csrc/ascend/cmake/scripts/util/opdesc_parser.py @@ -0,0 +1,409 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import os + +OP_ALL = "__ALLOP__" +SOC_ALL = "__ALLSOC__" +SOC_TO_SHORT_SOC_MAP = { + "ascend910a": "ascend910", + "ascend910proa": "ascend910", + "ascend910b": "ascend910", + "ascend910prob": "ascend910", + "ascend910premiuma": "ascend910", + "ascend910b1": "ascend910b", + "ascend910b2": "ascend910b", + "ascend910b2c": "ascend910b", + "ascend910b3": "ascend910b", + "ascend910b4": "ascend910b", + "ascend910b4-1": "ascend910b", + "ascend910_9391": "ascend910_93", + "ascend910_9381": "ascend910_93", + "ascend910_9372": "ascend910_93", + "ascend910_9392": "ascend910_93", + "ascend910_9382": "ascend910_93", + "ascend910_9362": "ascend910_93", + "ascend310p1": "ascend310p", + "ascend310p3": "ascend310p", + "ascend310p5": "ascend310p", + "ascend310p7": "ascend310p", + "ascend310p3vir01": "ascend310p", + "ascend310p3vir02": "ascend310p", + "ascend310p3vir04": "ascend310p", + "ascend310p3vir08": "ascend310p", + "ascend310b1": "ascend310b", + "bs9sx1aa": "bs9sx1a", + "ascend610lite": "ascend610lite", + "ascend950": "ascend950", +} +CONFLICT_KEYWORDS = { + "and", + "as", + "assert", + "break", + "class", + "continue", + "def", + "del", + "elif", + "else", + "except", + "finally", + "for", + "from", + "global", + "if", + "import", + "in", + "is", + "lambda", + "not", + "or", + "pass", + "raise", + "return", + "try", + "while", + "with", + "yield", + "False", + "None", + "True", + "nonlocal", + "arg", + "__inputs__", + "__outputs__", + "options", + "bisheng", + "bisheng_path", + "tikcpp_path", + "impl_mode", + "custom_compile_options", + "custom_all_compile_options", + "soc_version", + "soc_short", + "custom_compile_options_soc", + "custom_all_compile_options_soc", + "origin_func_name", + "ascendc_src_dir_ex", + "ascendc_src_dir", + "ascendc_src_file", + "src", + "op_type", + "code_channel", + "op_info", + "compile_op", + "get_code_channel", + "result", + "__attrs__", + "isinstance", + "attr", + "get_current_build_config", + "_build_args", + "get_dtype_fmt_options", + "shutil", + "os", + "get_kernel_source", +} + + +class OpDesc: + def __init__(self: any, op_type: str): + self.op_type = op_type + self.attr_list = [] + self.attr_val = {} + self.input_name = [] + self.input_ori_name = [] + self.input_type = [] + self.input_dtype = [] + self.input_dtype_for_bin_list = [] + self.input_dtype_for_bin = {} + self.input_fmt = [] + self.input_fmt_for_bin_list = [] + self.input_fmt_for_bin = {} + self.input_virt = {} + self.input_value_depend = {} + self.output_name = [] + self.output_ori_name = [] + self.output_type = [] + self.output_dtype = [] + self.output_dtype_for_bin_list = [] + self.output_dtype_for_bin = {} + self.output_fmt = [] + self.output_fmt_for_bin_list = [] + self.output_fmt_for_bin = {} + self.output_init_value = [] + self.output_shape_depend_on_compute = [] + self.op_fmt_sel = False + self.op_chk_support = False + self.op_intf = "" + self.kern_name = "" + self.op_file = "" + self.op_replay_flag = False + self.op_replay_batch = False + self.input_idx = -1 + self.output_idx = -1 + self.max_block_dim = 32 + self.max_shape_size = 268435456 + self.dynamic_shape = False + self.op_range_limit = "" + self.custom_compile_options = {} + self.custom_all_compile_options = {} + self.param_type_dynamic = False + self.mc2_ctx = [] + self.bin_cprs_list = [] + self.bin_cprs_head = [] + self.bin_save_list = [] + + @staticmethod + def _parse_digit(conf: str) -> int: + return int(conf.split("=")[1]) + + @staticmethod + def _parse_flag(conf: str) -> bool: + return conf.split("=")[1] == "true" + + @staticmethod + def _parse_str(conf: str) -> str: + return conf.split("=")[1] + + @staticmethod + def _parse_list(conf: str) -> list: + return conf.split("=")[1].split(",") + + def parse_input(self: any, conf: str): + if conf.startswith("input{}.name".format(int(self.input_idx) + 1)): + self.input_idx += 1 + self.input_ori_name.append(self._parse_str(conf)) + self.input_name.append(self.input_ori_name[-1] + "_in__") + elif conf.startswith("input{}.paramType".format(int(self.input_idx))): + param_type = self._parse_str(conf) + self.input_type.append(param_type) + if param_type == "dynamic": + self.param_type_dynamic = True + elif conf.startswith("input{}.dtype".format(int(self.input_idx))): + self.input_dtype.append(self._parse_str(conf)) + elif conf.startswith("input{}.for_bin_dtype".format(int(self.input_idx))): + self.input_dtype_for_bin.update({self.input_idx: self._parse_str(conf)}) + elif conf.startswith("input{}.format".format(int(self.input_idx))): + self.input_fmt.append(self._parse_str(conf)) + elif conf.startswith("input{}.for_bin_format".format(int(self.input_idx))): + self.input_fmt_for_bin.update({self.input_idx: self._parse_str(conf)}) + elif conf.startswith("input{}.virtual".format(int(self.input_idx))): + self.input_virt[self.input_idx] = self._parse_str(conf) + elif conf.startswith("input{}.valueDepend".format(int(self.input_idx))): + self.input_value_depend[self.input_idx] = self._parse_str(conf) + elif conf.startswith("input{}.initValue".format(int(self.input_idx))): + raise Exception( + f"[ERROR]: Op: {{'{self.op_type}'}} input {self.input_ori_name[int(self.input_idx)]}\ + has InitValue, which is not support!" + ) + else: + return + + def parse_output(self: any, conf: str): + if conf.startswith("output{}.name".format(int(self.output_idx) + 1)): + self.output_idx += 1 + self.output_ori_name.append(self._parse_str(conf)) + self.output_name.append(self.output_ori_name[-1] + "_out_") + self.output_init_value.append(None) + elif conf.startswith("output{}.paramType".format(int(self.output_idx))): + param_type = self._parse_str(conf) + self.output_type.append(param_type) + if param_type == "dynamic": + self.param_type_dynamic = True + elif conf.startswith("output{}.dtype".format(int(self.output_idx))): + self.output_dtype.append(self._parse_str(conf)) + elif conf.startswith("output{}.for_bin_dtype".format(int(self.output_idx))): + self.output_dtype_for_bin.update({self.output_idx: self._parse_str(conf)}) + elif conf.startswith("output{}.format".format(int(self.output_idx))): + self.output_fmt.append(self._parse_str(conf)) + elif conf.startswith("output{}.for_bin_format".format(int(self.output_idx))): + self.output_fmt_for_bin.update({self.output_idx: self._parse_str(conf)}) + elif conf.startswith("output{}.initValue".format(int(self.output_idx))): + self.output_init_value[int(self.output_idx)] = self._parse_str(conf) + elif conf.startswith("output{}.outputShapeDependOnCompute=true".format(int(self.output_idx))): + self.output_shape_depend_on_compute.append(int(self.output_idx)) + else: + return + + def parse_op_format(self: any, conf: str): + self.op_fmt_sel = self._parse_flag(conf) + + def parse_check_support(self: any, conf: str): + self.op_chk_support = self._parse_flag(conf) + + def parse_range_limit(self: any, conf: str): + self.op_range_limit = self._parse_str(conf) + + def parse_kern_name(self: any, conf: str): + self.kern_name = self._parse_str(conf) + + def parse_op_intf(self: any, conf: str): + self.op_intf = self._parse_str(conf) + + def parse_op_file(self: any, conf: str): + self.op_file = self._parse_str(conf) + + def parse_dynamic_shape(self: any, conf: str): + self.dynamic_shape = self._parse_flag(conf) + + def parse_attr_list(self: any, conf: str): + self.attr_list = self._parse_list(conf) + intersection_element = set(self.attr_list) & CONFLICT_KEYWORDS + if intersection_element: + raise Exception( + f"[ERROR]: The attribute name: {intersection_element} in op: {{'{self.op_type}'}} \ +conflicts with the built-in variable name. Use a complex name or prefix the operator name." + ) + + def parse_mc2_ctx(self: any, conf: str): + self.mc2_ctx = self._parse_list(conf) + + @staticmethod + def _camel_to_snake(camel_case_str: str): + snake_case_str = "" + for i, c in enumerate(camel_case_str): + if i == 0: + snake_case_str += c.lower() + elif c.isupper(): + snake_case_str += "_" + c.lower() + else: + snake_case_str += c + return snake_case_str + + def parse_attr_val(self: any, conf: str): + for attr in self.attr_list: + if self.attr_val.get(attr) is None: + self.attr_val[attr] = {} + if conf.startswith("attr_{}.type".format(attr)): + self.attr_val.get(attr)["type"] = self._camel_to_snake(self._parse_str(conf)) + elif conf.startswith("attr_{}.paramType".format(attr)): + self.attr_val.get(attr)["paramType"] = self._parse_str(conf) + elif conf.startswith("attr_{}.defaultValue".format(attr)): + self.attr_val.get(attr)["defaultValue"] = self._parse_str(conf) + + def parse_replay_val(self: any, batch_list: list, iterator_list: list): + if self.op_type in batch_list: + self.op_replay_flag = True + self.op_replay_batch = True + elif self.op_type in iterator_list: + self.op_replay_flag = True + self.op_replay_batch = False + + +def _is_op_type_in_opdesc(op_descs: list, op_type: str): + return any(op_type == op.op_type for op in op_descs) + + +def _set_all_options_to_opdescs(op_descs, soc_ver_compile_options): + for op in op_descs: + op.custom_all_compile_options = soc_ver_compile_options + + +def _set_options_to_opdesc(op_descs, op_type, soc_ver_compile_options): + for op in op_descs: + if op.op_type != op_type: + continue + op.custom_compile_options.update(soc_ver_compile_options) + + +def _trans_soc_ver_to_short(soc_ver: str): + low_soc_ver = soc_ver.lower() + if low_soc_ver not in SOC_TO_SHORT_SOC_MAP: + print(f"WARNING: caution: {soc_ver} will trans into ascend910, if not your intention,use ascend910b1~4 instead") + return SOC_TO_SHORT_SOC_MAP[low_soc_ver] + + +def _get_op_custom_options(op_descs: list, auto_gen_dir: str): + if auto_gen_dir is None: + return {} + file = os.path.join(auto_gen_dir, "custom_compile_options.ini") + if not os.path.exists(file): + print(f"WARNING: cannot find {auto_gen_dir}/custom_compile_options.ini") + return {} + with open(file) as fd: + lines = fd.readlines() + for line in lines: + param_list = str.split(line.rstrip("\n"), ",") + if len(param_list) != 3: + raise Exception(f"ERROR: custom compile option {param_list} len is not 3") + op_type = param_list[0] + if op_type.upper() == "ALL": + op_type = OP_ALL + if op_type != OP_ALL and not _is_op_type_in_opdesc(op_descs, op_type): + continue + soc_ver_compile_options = {} + soc_ver = param_list[1] + options_str = param_list[2] + options = str.split(options_str, ";") + if soc_ver == "": + soc_ver_compile_options[SOC_ALL] = options + else: + soc_ver_list = str.split(soc_ver, ";") + for ver in soc_ver_list: + short_ver = _trans_soc_ver_to_short(ver) + soc_ver_compile_options[short_ver] = options + if op_type == OP_ALL: + _set_all_options_to_opdescs(op_descs, soc_ver_compile_options) + else: + _set_options_to_opdesc(op_descs, op_type, soc_ver_compile_options) + + +def get_op_desc( + file: str, batch_list: list, iterator_list: list, builder: any, op_type: list, auto_gen_dir: str = None +) -> list: + op_descs = [] + op_match = False + with open(file) as fd: + lines = fd.readlines() + for line in lines: + line = line.strip() + if line.startswith("["): + name = line[1:-1] + if op_type is None or name in op_type: + op_match = True + op_desc = builder(name) + op_desc.parse_replay_val(batch_list, iterator_list) + op_descs.append(op_desc) + else: + op_match = False + if op_type is not None and len(op_descs) == len(op_type): + break + continue + if not op_match: + continue + if line.startswith("input"): + op_desc.parse_input(line) + elif line.startswith("output"): + op_desc.parse_output(line) + elif line.startswith("dynamicFormat.flag"): + op_desc.parse_op_format(line) + elif line.startswith("needCheckSupport.flag"): + op_desc.parse_check_support(line) + elif line.startswith("rangeLimit.value"): + op_desc.parse_range_limit(line) + elif line.startswith("opInterface.value"): + op_desc.parse_op_intf(line) + elif line.startswith("kernel.name"): + op_desc.parse_kern_name(line) + elif line.startswith("opFile.value"): + op_desc.parse_op_file(line) + elif line.startswith("dynamicShapeSupport.flag"): + op_desc.parse_dynamic_shape(line) + elif line.startswith("mc2.ctx"): + op_desc.parse_mc2_ctx(line) + elif line.startswith("attr.list"): + op_desc.parse_attr_list(line) + elif line.startswith("attr_"): + op_desc.parse_attr_val(line) + _get_op_custom_options(op_descs, auto_gen_dir) + return op_descs diff --git a/csrc/ascend/cmake/scripts/util/parse_ini_to_json.py b/csrc/ascend/cmake/scripts/util/parse_ini_to_json.py new file mode 100644 index 000000000..27e1186f4 --- /dev/null +++ b/csrc/ascend/cmake/scripts/util/parse_ini_to_json.py @@ -0,0 +1,429 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import json +import os +import stat +import sys + +ATTR_TYPE_LIST = [ + "int", + "float", + "bool", + "str", + "listInt", + "listFloat", + "listBool", + "listStr", + "listListInt", + "type", + "listType", + "tensor", + "listTensor", +] +ATTR_PARAMTYPE_LIST = ["optional", "required"] +BOOL_FLAG_KEY = [ + "dynamicFormat", + "dynamicShapeSupport", + "dynamicRankSupport", + "precision_reduce", + "heavyOp", + "needCheckSupport", + "enableVectorCore", +] +BOOL_LIST = ["true", "false"] +DTYPE_LIST = [ + "float16", + "float", + "float32", + "int8", + "int16", + "int32", + "uint8", + "uint16", + "uint32", + "bool", + "int64", + "uint64", + "qint8", + "qint16", + "qint32", + "quint8", + "quint16", + "double", + "complex32", + "complex64", + "complex128", + "string", + "resource", + "dual", + "dual_sub_int8", + "dual_sub_uint8", + "string_ref", + "int4", + "bfloat16", + "uint1", + "hifloat8", + "float8_e4m3fn", + "float8_e5m2", + "float8_e8m0", + "float4_e2m1", + "float4_e1m2", + "int2", +] +FORMAT_LIST = [ + "NCHW", + "NHWC", + "ND", + "NC1HWC0", + "FRACTAL_Z", + "NC1C0HWPAD", + "NHWC1C0", + "FSR_NCHW", + "FRACTAL_DECONV", + "C1HWNC0", + "FRACTAL_DECONV_TRANSPOSE", + "FRACTAL_DECONV_SP_STRIDE_TRANS", + "NC1HWC0_C04", + "FRACTAL_Z_C04", + "CHWN", + "FRACTAL_DECONV_SP_STRIDE8_TRANS", + "HWCN", + "NC1KHKWHWC0", + "BN_WEIGHT", + "FILTER_HWCK", + "HASHTABLE_LOOKUP_LOOKUPS", + "HASHTABLE_LOOKUP_KEYS", + "HASHTABLE_LOOKUP_VALUE", + "HASHTABLE_LOOKUP_OUTPUT", + "HASHTABLE_LOOKUP_HITS", + "C1HWNCoC0", + "MD", + "NDHWC", + "FRACTAL_ZZ", + "FRACTAL_NZ", + "NCDHW", + "DHWCN", + "NDC1HWC0", + "FRACTAL_Z_3D", + "CN", + "NC", + "DHWNC", + "FRACTAL_Z_3D_TRANSPOSE", + "FRACTAL_ZN_LSTM", + "FRACTAL_ZN_RNN", + "FRACTAL_Z_G", + "NULL", + "FRACTAL_NZ_C0_2", + "FRACTAL_NZ_C0_4", + "FRACTAL_NZ_C0_16", + "FRACTAL_NZ_C0_32", +] + + +def parse_ini_files(ini_files): + """ + parse ini files to json + Parameters: + ---------------- + ini_files:input file list + return:ops_info + ---------------- + """ + tbe_ops_info = {} + for ini_file in ini_files: + check_file_size(ini_file) + parse_ini_to_obj(ini_file, tbe_ops_info) + return tbe_ops_info + + +def check_file_size(input_file): + try: + file_size = os.path.getsize(input_file) + except OSError as os_error: + print(f'[ERROR] Failed to open "{input_file}". {os_error}') + raise OSError from os_error + if file_size > 10 * 1024 * 1024: + print(f"[WARN] The size of {input_file} exceeds 10MB, it may take more time to run, please wait.") + + +def parse_ini_to_obj(ini_file, tbe_ops_info): + """ + parse ini file to json obj + Parameters: + ---------------- + ini_file:ini file path + tbe_ops_info:ops_info + ---------------- + """ + with open(ini_file) as ini_file: + lines = ini_file.readlines() + op_dict = {} + op_name = "" + find_op_type = False + for line in lines: + line = line.rstrip() + if line == "": + continue + if line.startswith("["): + if line.endswith("]"): + op_name = line[1:-1] + op_dict = {} + tbe_ops_info[op_name] = op_dict + find_op_type = True + elif "=" in line: + key1 = line[: line.index("=")] + key2 = line[line.index("=") + 1 :] + key1_0, key1_1 = key1.split(".") + if key1_0 not in op_dict: + op_dict[key1_0] = {} + if key1_1 in op_dict.get(key1_0): + raise RuntimeError("Op:" + op_name + " " + key1_0 + " " + key1_1 + " is repeated!") + dic_key = op_dict.get(key1_0) + dic_key[key1_1] = key2 + else: + continue + if not find_op_type: + raise RuntimeError("Not find OpType in .ini file.") + + +def check_output_exist(op_dict, is_valid): + """ + Function Description: + Check output is exist + Parameter: op_dict + Parameter: is_valid + """ + if "output0" in op_dict: + output0_dict = op_dict.get("output0") + if output0_dict.get("name", None) is None: + is_valid = False + print("output0.name is required in .ini file!") + else: + is_valid = False + print("output0 is required in .ini file!") + return is_valid + + +def check_attr_dict(attr_dict, is_valid, attr): + """ + Function Description: + Check attr_dict + Parameter: attr_dict + Parameter: is_valid + Parameter: attr + """ + attr_type = attr_dict.get("type") + value = attr_dict.get("value") + param_type = attr_dict.get("paramType") + if attr_type is None or value is None: + is_valid = False + print(f"If attr.list is exist, {attr}.type and {attr}.value is required") + if param_type and param_type not in ATTR_PARAMTYPE_LIST: + is_valid = False + print(f"{attr}.paramType only support {ATTR_PARAMTYPE_LIST}.") + if attr_type and attr_type not in ATTR_TYPE_LIST: + is_valid = False + print(f"{attr}.type only support {ATTR_TYPE_LIST}.") + return is_valid + + +def check_attr(op_dict, is_valid): + """ + Function Description: + Check attr + Parameter: op_dict + Parameter: is_valid + """ + if "attr" in op_dict: + attr_dict = op_dict.get("attr") + attr_list_str = attr_dict.get("list", None) + if attr_list_str is None: + is_valid = False + print("attr.list is required in .ini file!") + else: + attr_list = attr_list_str.split(",") + for attr_name in attr_list: + attr = "attr_" + attr_name.strip() + attr_dict = op_dict.get(attr) + if attr_dict: + is_valid = check_attr_dict(attr_dict, is_valid, attr) + else: + is_valid = False + print(f"{attr} is required in .ini file, when attr.list is {attr_list_str}!") + return is_valid + + +def check_bool_flag(op_dict, is_valid): + """ + Function Description: + check_bool_flag + Parameter: op_dict + Parameter: is_valid + """ + for key in BOOL_FLAG_KEY: + if key in op_dict: + op_bool_key = op_dict.get(key) + if op_bool_key.get("flag").strip() not in BOOL_LIST: + is_valid = False + print(f"{key}.flag only support {BOOL_LIST}.") + return is_valid + + +def check_type_format(op_info, is_valid, op_info_key): + """ + Function Description: + Check type and format + Parameter: op_info + Parameter: is_valid + Parameter: op_info_key + """ + op_info_dtype_str = op_info.get("dtype") + op_info_dtype_num = 0 + op_info_format_num = 0 + if op_info_dtype_str: + op_info_dtype = op_info_dtype_str.split(",") + op_info_dtype_num = len(op_info_dtype) + for dtype in op_info_dtype: + if dtype.strip() not in DTYPE_LIST: + is_valid = False + print(f"{op_info_key}.dtype not support {dtype}.") + op_info_format_str = op_info.get("format") + if op_info_format_str: + op_info_format = op_info_format_str.split(",") + op_info_format_num = len(op_info_format) + for op_format in op_info_format: + if op_format.strip() not in FORMAT_LIST: + is_valid = False + print(f"{op_info_key}.format not support {op_format}.") + if op_info_dtype_num > 0 and op_info_format_num > 0: + if op_info_dtype_num != op_info_format_num: + is_valid = False + print("The number of {0}.dtype not match the number of {0}.format.".format(op_info_key)) + return is_valid + + +def check_op_info(tbe_ops): + """ + Function Description: + Check info. + Parameter: tbe_ops + Return Value: is_valid + """ + print("\n\n==============check valid for ops info start==============") + required_op_input_info_keys = ["paramType", "name"] + required_op_output_info_keys = ["paramType", "name"] + param_type_valid_value = ["dynamic", "optional", "required"] + is_valid = True + for op_key in tbe_ops: + op_dict = tbe_ops[op_key] + for op_info_key in op_dict: + if op_info_key.startswith("input"): + op_input_info = op_dict[op_info_key] + missing_keys = [] + for required_op_input_info_key in required_op_input_info_keys: + if required_op_input_info_key not in op_input_info: + missing_keys.append(required_op_input_info_key) + if len(missing_keys) > 0: + print("op: " + op_key + " " + op_info_key + " missing: " + ",".join(missing_keys)) + is_valid = False + else: + if op_input_info["paramType"] not in param_type_valid_value: + print( + "op: " + op_key + " " + op_info_key + " paramType not valid, valid key:[dynamic, " + "optional, required]" + ) + is_valid = False + is_valid = check_type_format(op_input_info, is_valid, op_info_key) + if op_info_key.startswith("output"): + op_input_info = op_dict[op_info_key] + missing_keys = [] + for required_op_input_info_key in required_op_output_info_keys: + if required_op_input_info_key not in op_input_info: + missing_keys.append(required_op_input_info_key) + if len(missing_keys) > 0: + print("op: " + op_key + " " + op_info_key + " missing: " + ",".join(missing_keys)) + is_valid = False + else: + if op_input_info["paramType"] not in param_type_valid_value: + print( + "op: " + op_key + " " + op_info_key + " paramType not valid, valid key:[dynamic, " + "optional, required]" + ) + is_valid = False + is_valid = check_type_format(op_input_info, is_valid, op_info_key) + is_valid = check_attr(op_dict, is_valid) + is_valid = check_bool_flag(op_dict, is_valid) + print("==============check valid for ops info end================\n\n") + return is_valid + + +def write_json_file(tbe_ops_info, json_file_path): + """ + Save info to json file + Parameters: + ---------------- + tbe_ops_info: ops_info + json_file_path: json file path + ---------------- + """ + json_file_real_path = os.path.realpath(json_file_path) + wr_flag = os.O_WRONLY | os.O_CREAT + wr_mode = stat.S_IWUSR | stat.S_IRUSR + with os.fdopen(os.open(json_file_real_path, wr_flag, wr_mode), "w") as file_path: + # The owner have all rights£¬group only have read rights + os.chmod(json_file_real_path, stat.S_IWUSR + stat.S_IRGRP + stat.S_IRUSR) + json.dump(tbe_ops_info, file_path, sort_keys=True, indent=4, separators=(",", ":")) + print("Compile op info cfg successfully.") + + +def parse_ini_to_json(ini_file_paths, outfile_path): + """ + parse ini files to json file + Parameters: + ---------------- + ini_file_paths: list of ini file path + outfile_path: output file path + ---------------- + """ + tbe_ops_info = parse_ini_files(ini_file_paths) + if not check_op_info(tbe_ops_info): + print("Compile op info cfg failed.") + return False + write_json_file(tbe_ops_info, outfile_path) + return True + + +if __name__ == "__main__": + args = sys.argv + + OUTPUT_FILE_PATH = "tbe_ops_info.json" + ini_file_path_list = [] + parse_ini_list = [] + + for arg in args: + if arg.endswith("ini"): + ini_file_path_list.append(arg) + OUTPUT_FILE_PATH = arg.replace(".ini", ".json") + if arg.endswith("json"): + OUTPUT_FILE_PATH = arg + + if not ini_file_path_list: + ini_file_path_list.append("tbe_ops_info.ini") + + for ini_file in ini_file_path_list: + if os.path.exists(ini_file): + parse_ini_list.append(ini_file) + + if parse_ini_list: + if not parse_ini_to_json(parse_ini_list, OUTPUT_FILE_PATH): + sys.exit(1) + sys.exit(0) diff --git a/csrc/ascend/cmake/static.cmake b/csrc/ascend/cmake/static.cmake new file mode 100644 index 000000000..181a06a5d --- /dev/null +++ b/csrc/ascend/cmake/static.cmake @@ -0,0 +1,98 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +if (TARGET ${OPHOST_NAME}_infer_obj OR TARGET ${OPHOST_NAME}_tiling_obj OR TARGET ${OPHOST_NAME}_aicpu_objs) + add_library( + ${OPHOST_NAME}_static STATIC + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $<$:$> + ) + add_custom_command(TARGET ${OPHOST_NAME}_static + POST_BUILD + COMMAND python3 ${PROJECT_SOURCE_DIR}/scripts/util/build_opp_kernel_static.py + GenerateSymbol -l ${PROJECT_SOURCE_DIR}/build/lib${OPHOST_NAME}_static.a + -s ${PROJECT_SOURCE_DIR}/build/${OPHOST_NAME}.txt + WORKING_DIRECTORY ${CMAKE_BINARY_DIR}/bin_tmp) + target_link_libraries( + ${OPHOST_NAME}_static + PRIVATE $ + c_sec + -Wl,--no-as-needed + register + $<$:$> + $<$:opsbase> + -Wl,--as-needed + -Wl,--whole-archive + rt2_registry_static + tiling_api + -Wl,--no-whole-archive + ) +endif() + + +if (TARGET ${OPHOST_NAME}_opapi_obj OR TARGET opbuild_gen_aclnn_all) + if (TARGET ops_aclnn) + target_compile_definitions(ops_aclnn PUBLIC ACLNN_WITH_BINARY) + endif() + if (TARGET ${OPHOST_NAME}_opapi_obj) + target_compile_definitions(${OPHOST_NAME}_opapi_obj PUBLIC ACLNN_WITH_BINARY) + endif() + if (TARGET opbuild_gen_aclnn_all) + target_compile_definitions(opbuild_gen_aclnn_all PUBLIC ACLNN_WITH_BINARY) + endif() + + add_library( + ${OPAPI_NAME}_static STATIC + $<$:$> + $<$:$> + ) + add_dependencies(${OPAPI_NAME}_static ${OPHOST_NAME}_static) + add_custom_command(TARGET ${OPAPI_NAME}_static + POST_BUILD + COMMAND ${CMAKE_AR} x ${PROJECT_SOURCE_DIR}/build/libops_aclnn.a + COMMAND ${CMAKE_AR} x ${PROJECT_SOURCE_DIR}/build/lib${OPAPI_NAME}_static.a + COMMAND ${CMAKE_AR} qcs lib${OPAPI_NAME}_static.a *.o + COMMAND rm *.o + COMMAND python3 ${PROJECT_SOURCE_DIR}/scripts/util/build_opp_kernel_static.py + GenerateSymbol -l ${PROJECT_SOURCE_DIR}/build/bin_tmp/lib${OPAPI_NAME}_static.a + -s ${PROJECT_SOURCE_DIR}/build/${OPAPI_NAME}.txt + WORKING_DIRECTORY ${CMAKE_BINARY_DIR}/bin_tmp) +endif() + +add_custom_target(${OPSTATIC_NAME}) +foreach(compute_unit ${ASCEND_COMPUTE_UNIT}) + set(RESOURCE_PATH ${CMAKE_BINARY_DIR}/autogen/${compute_unit}/aclnnop_resource) + file(GLOB RESOURCE_CPP ${RESOURCE_PATH}/*.cpp) + set_source_files_properties(${RESOURCE_CPP} PROPERTIES GENERATED TRUE) + add_library(resource_${compute_unit}_static STATIC + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $<$:$>) + target_sources(resource_${compute_unit}_static PRIVATE ${RESOURCE_CPP}) + target_include_directories(resource_${compute_unit}_static PRIVATE + ${OPAPI_INCLUDE}) + set_target_properties(resource_${compute_unit}_static PROPERTIES + ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin_tmp/${compute_unit} + OUTPUT_NAME ${OPSTATIC_NAME}) + add_dependencies(${OPSTATIC_NAME} resource_${compute_unit}_static ${OPHOST_NAME}_static ${OPAPI_NAME}_static) + add_custom_command(TARGET resource_${compute_unit}_static + POST_BUILD + COMMAND ${CMAKE_AR} x ${CMAKE_BINARY_DIR}/libops_aclnn.a + COMMAND ${CMAKE_AR} qcs lib${OPSTATIC_NAME}.a *.o + WORKING_DIRECTORY ${CMAKE_BINARY_DIR}/bin_tmp/${compute_unit}) + install(FILES ${CMAKE_BINARY_DIR}/bin_tmp/${compute_unit}/lib${OPSTATIC_NAME}.a + DESTINATION ${CMAKE_BINARY_DIR}/static_library_files/lib64 + OPTIONAL) +endforeach() diff --git a/csrc/ascend/cmake/symbol.cmake b/csrc/ascend/cmake/symbol.cmake new file mode 100644 index 000000000..50d6dc07d --- /dev/null +++ b/csrc/ascend/cmake/symbol.cmake @@ -0,0 +1,330 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +function(gen_common_symbol) + install(DIRECTORY ${OPS_TRANSFORMER_COMMON_INC_HEADERS} + DESTINATION ${COMMON_INC_INSTALL_DIR} + ) +endfunction() + +# ophost shared +function(gen_ophost_symbol) + add_library(${OPHOST_NAME} SHARED + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $<$:$> + ) + + target_link_libraries( + ${OPHOST_NAME} + PRIVATE $ + c_sec + -Wl,--no-as-needed + register + $<$:opsbase> + -Wl,--as-needed + -Wl,--whole-archive + rt2_registry_static + -Wl,--no-whole-archive + tiling_api + -Wl,-Bsymbolic + ) + + target_link_directories(${OPHOST_NAME} + PRIVATE + ${ASCEND_DIR}/${SYSTEM_PREFIX}/lib64 + ) + + install(TARGETS ${OPHOST_NAME} + LIBRARY DESTINATION ${OPHOST_LIB_INSTALL_PATH} + ) +endfunction() + +# graph_plugin shared +function(gen_opgraph_symbol) + add_library(${OPGRAPH_NAME} SHARED + $<$:$> + ) + + target_link_libraries( + ${OPGRAPH_NAME} + PRIVATE $ + c_sec + -Wl,--no-as-needed + register + $<$:opsbase> + -Wl,--as-needed + -Wl,--whole-archive + rt2_registry_static + -Wl,--no-whole-archive + -Wl,-Bsymbolic + ) + + target_link_directories(${OPGRAPH_NAME} + PRIVATE + ${ASCEND_DIR}/${SYSTEM_PREFIX}/lib64 + ) + + set_target_properties(${OPGRAPH_NAME} PROPERTIES OUTPUT_NAME "opgraph_transformer") + install(TARGETS ${OPGRAPH_NAME} + LIBRARY DESTINATION ${OPGRAPH_LIB_INSTALL_DIR} + ) + install(FILES ${ASCEND_GRAPH_CONF_DST}/ops_proto_transformer.h + DESTINATION ${OPGRAPH_INC_INSTALL_DIR} OPTIONAL + ) +endfunction() + +function(gen_opapi_symbol) + # opapi shared + add_library(${OPAPI_NAME} SHARED + $<$:$> + $<$:$> + ) + + target_link_libraries(${OPAPI_NAME} + PUBLIC + $ + -Wl,--whole-archive + ops_aclnn + -Wl,--no-whole-archive + nnopbase + profapi + ge_common_base + ascend_dump + ascendalog + opapi_math + dl + ) +endfunction() + +function(gen_built_in_opapi_symbol) + install(TARGETS ${OPAPI_NAME} + LIBRARY DESTINATION ${ACLNN_LIB_INSTALL_DIR} + ) +endfunction() + +function(gen_cust_opapi_symbol) + #op_api + set_target_properties(${OPAPI_NAME} PROPERTIES OUTPUT_NAME "cust_opapi") + + install(TARGETS ${OPAPI_NAME} + LIBRARY DESTINATION ${ACLNN_LIB_INSTALL_DIR} + ) +endfunction() + +function(gen_cust_optiling_symbol) + # op_tiling + if(NOT TARGET ${OPHOST_NAME}_tiling_obj) + return() + endif() + add_library(cust_opmaster SHARED + $<$:$ + $<$:$>> + ) + target_link_libraries(cust_opmaster + PRIVATE + c_sec + tiling_api + -Wl,--no-as-needed + register + $<$:opsbase> + -Wl,--as-needed + -Wl,--whole-archive + rt2_registry_static + -Wl,--no-whole-archive + ) + target_link_directories(cust_opmaster + PRIVATE + ${ASCEND_DIR}/${SYSTEM_PREFIX}/lib64 + ) + set_target_properties(cust_opmaster PROPERTIES OUTPUT_NAME "cust_opmaster_rt2.0") + + install(TARGETS cust_opmaster + LIBRARY DESTINATION ${OPTILING_LIB_INSTALL_DIR} + ) + add_custom_target(optiling_compat ALL + COMMAND ln -sf lib/linux/${CMAKE_SYSTEM_PROCESSOR}/$ + ${CMAKE_BINARY_DIR}/liboptiling.so + ) + install(FILES ${CMAKE_BINARY_DIR}/liboptiling.so + DESTINATION ${OPTILING_INSTALL_DIR} + ) +endfunction() + +function(gen_cust_proto_symbol) + # op_proto + if(NOT TARGET ${OPHOST_NAME}_infer_obj) + return() + endif() + add_library(cust_proto SHARED + $<$:$> + $<$:$> + ) + target_link_libraries(cust_proto + PRIVATE + c_sec + -Wl,--no-as-needed + register + $<$:opsbase> + -Wl,--as-needed + -Wl,--whole-archive + rt2_registry_static + -Wl,--no-whole-archive + ) + target_link_directories(cust_proto + PRIVATE + ${ASCEND_DIR}/${SYSTEM_PREFIX}/lib64 + ) + set_target_properties(cust_proto PROPERTIES OUTPUT_NAME "cust_opsproto_rt2.0") + + install(TARGETS cust_proto + LIBRARY DESTINATION ${OPPROTO_LIB_INSTALL_DIR} + ) + file(GLOB_RECURSE proto_headers + ${ASCEND_AUTOGEN_PATH}/*_proto.h + ) + install(FILES ${proto_headers} + DESTINATION ${OPPROTO_INC_INSTALL_DIR} OPTIONAL + ) +endfunction() + +function(gen_cust_aicpu_json_symbol) + get_property(ALL_AICPU_JSON_FILES GLOBAL PROPERTY AICPU_JSON_FILES) + if(NOT ALL_AICPU_JSON_FILES) + message(STATUS "No aicpu json files to merge, skipping.") + return() + endif() + + set(MERGED_JSON ${CMAKE_BINARY_DIR}/cust_aicpu_kernel.json) + add_custom_command( + OUTPUT ${MERGED_JSON} + COMMAND bash ${CMAKE_SOURCE_DIR}/scripts/util/merge_aicpu_info_json.sh ${CMAKE_SOURCE_DIR} ${MERGED_JSON} ${ALL_AICPU_JSON_FILES} + DEPENDS ${ALL_AICPU_JSON_FILES} + COMMENT "Merging Json files into ${MERGED_JSON}" + VERBATIM + ) + add_custom_target(merge_aicpu_json ALL DEPENDS ${MERGED_JSON}) + install( + FILES ${MERGED_JSON} + DESTINATION packages/vendors/${VENDOR_NAME}_transformer/op_impl/cpu/config + OPTIONAL + ) +endfunction() + +function(gen_cust_aicpu_kernel_symbol) + if(NOT AICPU_CUST_OBJ_TARGETS) + message(STATUS "No aicpu cust obj targets found, skipping.") + return() + endif() + + set(ARM_CXX_COMPILER ${ASCEND_DIR}/toolkit/toolchain/hcc/bin/aarch64-target-linux-gnu-g++) + set(ARM_SO_OUTPUT ${CMAKE_BINARY_DIR}/libtransformer_aicpu_kernels.so) + + set(ALL_OBJECTS "") + foreach(tgt IN LISTS AICPU_CUST_OBJ_TARGETS) + list(APPEND ALL_OBJECTS $) + endforeach() + + message(STATUS "Linking cust_aicpu_kernels with ARM toolchain: ${ARM_CXX_COMPILER}") + message(STATUS "Objects: ${ALL_OBJECTS}") + message(STATUS "Output: ${ARM_SO_OUTPUT}") + + if(EXISTS ${ASCEND_DIR}/ops_base/lib64/libaicpu_context.a) + set(LIBAICPU_CONTEXT_PATH ${ASCEND_DIR}/ops_base/lib64/libaicpu_context.a) + else() + set(LIBAICPU_CONTEXT_PATH ${ASCEND_DIR}/lib64/libaicpu_context.a) + endif() + + if(EXISTS ${ASCEND_DIR}/ops_base/lib64/libbase_ascend_protobuf.a) + set(LIBBASE_ASCEND_PROTOBUF_PATH ${ASCEND_DIR}/ops_base/lib64/libbase_ascend_protobuf.a) + else() + set(LIBBASE_ASCEND_PROTOBUF_PATH ${ASCEND_DIR}/lib64/libbase_ascend_protobuf.a) + endif() + + add_custom_command( + OUTPUT ${ARM_SO_OUTPUT} + COMMAND ${ARM_CXX_COMPILER} -shared ${ALL_OBJECTS} + -Wl,--whole-archive + ${LIBAICPU_CONTEXT_PATH} + ${LIBBASE_ASCEND_PROTOBUF_PATH} + -Wl,--no-whole-archive + -Wl,-Bsymbolic + -Wl,--exclude-libs=libbase_ascend_protobuf.a + -s + -o ${ARM_SO_OUTPUT} + DEPENDS ${AICPU_CUST_OBJ_TARGETS} + COMMENT "Linking cust_aicpu_kernels.so using ARM toolchain" + ) + add_custom_target(cust_aicpu_kernels ALL DEPENDS ${ARM_SO_OUTPUT}) + + install( + FILES ${ARM_SO_OUTPUT} + DESTINATION packages/vendors/${VENDOR_NAME}_transformer/op_impl/cpu/aicpu_kernel/impl + OPTIONAL + ) +endfunction() + +function(gen_onnx_plugin_symbol) + add_library( + ${ONNX_PLUGIN_NAME} SHARED + $<$:$> + ) + + target_link_libraries( + ${ONNX_PLUGIN_NAME} + PRIVATE $ + c_sec + -Wl,--no-as-needed + register + $<$:opsbase> + -Wl,--as-needed + -Wl,--whole-archive + rt2_registry_static + -Wl,--no-whole-archive + ) + + target_link_directories(${ONNX_PLUGIN_NAME} PRIVATE ${ASCEND_DIR}/${SYSTEM_PREFIX}/lib64) + + install( + TARGETS ${ONNX_PLUGIN_NAME} + LIBRARY DESTINATION ${ONNX_PLUGIN_LIB_INSTALL_DIR} + ) + +endfunction() + +function(gen_norm_symbol) + gen_common_symbol() + + if (ENABLE_OPS_HOST) + gen_ophost_symbol() + gen_opapi_symbol() + gen_built_in_opapi_symbol() + endif() + + gen_opgraph_symbol() + gen_onnx_plugin_symbol() + +endfunction() + +function(gen_cust_symbol) + gen_opapi_symbol() + + gen_cust_opapi_symbol() + + gen_cust_optiling_symbol() + + gen_cust_proto_symbol() + + gen_cust_aicpu_json_symbol() + + gen_cust_aicpu_kernel_symbol() +endfunction() \ No newline at end of file diff --git a/csrc/ascend/cmake/third_party/.gitkeep b/csrc/ascend/cmake/third_party/.gitkeep new file mode 100644 index 000000000..e69de29bb diff --git a/csrc/ascend/cmake/third_party/abseil-cpp.cmake b/csrc/ascend/cmake/third_party/abseil-cpp.cmake new file mode 100644 index 000000000..6e4dd0ec9 --- /dev/null +++ b/csrc/ascend/cmake/third_party/abseil-cpp.cmake @@ -0,0 +1,59 @@ +# ---------------------------------------------------------------------------- +# This program is free software, you can redistribute it and/or modify it. +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This file is a part of the CANN Open Software. +# Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ---------------------------------------------------------------------------- + +include(ExternalProject) +set(ABSEIL_VERSION_PKG abseil-cpp-20230802.1.tar.gz) + +unset(abseil-cpp_FOUND CACHE) +unset(ABSL_SOURCE_DIR CACHE) + +find_path(ABSL_SOURCE_DIR + NAMES absl/log/absl_log.h + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH + PATHS ${PROJECT_SOURCE_DIR}/third_party/abseil-cpp) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(abseil-cpp + FOUND_VAR + abseil-cpp_FOUND + REQUIRED_VARS + ABSL_SOURCE_DIR) + +set(ABSEIL_SOURCE_DIR ${PROJECT_SOURCE_DIR}/third_party/abseil-cpp) +if(abseil-cpp_FOUND) + message(STATUS "Found abseil-cpp in ${CANN_3RD_LIB_PATH}/abseil-cpp") +else() + # 初始化可选参数列表 + if(EXISTS "${CANN_3RD_LIB_PATH}/abseil-cpp/${ABSEIL_VERSION_PKG}") + set(REQ_URL "file://${CANN_3RD_LIB_PATH}/abseil-cpp/${ABSEIL_VERSION_PKG}") + message(STATUS "[ThirdPartyLib][abseil-cpp] found in ${REQ_URL}.") + elseif(EXISTS "${CANN_3RD_LIB_PATH}/pkg/${ABSEIL_VERSION_PKG}") + set(REQ_URL "file://${CANN_3RD_LIB_PATH}/pkg/${ABSEIL_VERSION_PKG}") + message(STATUS "[ThirdPartyLib][abseil-cpp] found in ${REQ_URL}.") + else() + set(REQ_URL "https://gitcode.com/cann-src-third-party/abseil-cpp/releases/download/20230802.1/abseil-cpp-20230802.1.tar.gz") + message(STATUS "[ThirdPartyLib][abseil-cpp] ${REQ_URL} not found, need download.") + endif() + + ExternalProject_Add(abseil_build_transformer + URL ${REQ_URL} + DOWNLOAD_DIR ${CANN_3RD_LIB_PATH}/pkg + PATCH_COMMAND patch -p1 < ${CMAKE_CURRENT_LIST_DIR}/build/modules/patch/protobuf-hide_absl_symbols.patch + SOURCE_DIR ${ABSEIL_SOURCE_DIR} + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + INSTALL_COMMAND "" + EXCLUDE_FROM_ALL TRUE + ) + + ExternalProject_Get_Property(abseil_build_transformer SOURCE_DIR) + set(ABSL_SOURCE_DIR ${SOURCE_DIR}) +endif() \ No newline at end of file diff --git a/csrc/ascend/cmake/third_party/ascend_protobuf.cmake b/csrc/ascend/cmake/third_party/ascend_protobuf.cmake new file mode 100644 index 000000000..374223592 --- /dev/null +++ b/csrc/ascend/cmake/third_party/ascend_protobuf.cmake @@ -0,0 +1,89 @@ +# ---------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ---------------------------------------------------------------------------------------------------------- +include(ExternalProject) +set(PROTOBUF_VERSION_PKG protobuf-25.1.tar.gz) +set(ASCEND_PROTOBUF_DIR ${CANN_3RD_LIB_PATH}/ascend_protobuf) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(ascend_protobuf_build_transformer + FOUND_VAR + ascend_protobuf_build_transformer_FOUND + REQUIRED_VARS + ASCEND_PROTOBUF_SHARED_INCLUDE +) + +set(ASCEND_PROTOBUF_SOURCE_DIR ${PROJECT_SOURCE_DIR}/third_party/ascend_protobuf) +if(ascend_protobuf_build_transformer_FOUND AND NOT FORCE_REBUILD_CANN_3RD) + message(STATUS "[ThirdPartyLib][ascend protobuf] ascend_protobuf_shared found, skip compile.") + cmake_print_variables(ASCEND_PROTOBUF_SHARED_INCLUDE) + cmake_print_variables(ASCEND_PROTOC) + set(Protobuf_INCLUDE ${ASCEND_PROTOBUF_SHARED_INCLUDE}) + set(Protobuf_PATH ${ASCEND_PROTOC}) + set(Protobuf_PROTOC_EXECUTABLE ${Protobuf_PATH}/protoc) + add_library(ascend_protobuf_build_transformer INTERFACE) +else() + message(STATUS "[ThirdPartyLib][ascend protobuf] ascend protobuf shared not found, finding binary file.") + if(EXISTS "${CANN_3RD_LIB_PATH}/protobuf/protobuf-all-25.1.tar.gz") + set(REQ_URL "file://${CANN_3RD_LIB_PATH}/protobuf/protobuf-all-25.1.tar.gz") + message(STATUS "[ThirdPartyLib][ascend protobuf] found in ${REQ_URL}.") + elseif(EXISTS "${CANN_3RD_LIB_PATH}/pkg/${PROTOBUF_VERSION_PKG}") + set(REQ_URL "file://${CANN_3RD_LIB_PATH}/pkg/${PROTOBUF_VERSION_PKG}") + message(STATUS "[ThirdPartyLib][ascend protobuf] found in ${REQ_URL}.") + else() + set(REQ_URL "https://gitcode.com/cann-src-third-party/protobuf/releases/download/v25.1/protobuf-25.1.tar.gz") + message(STATUS "[ThirdPartyLib][ascend protobuf] ${REQ_URL} not found, need download.") + endif() + + set(protobuf_CXXFLAGS "-Wno-maybe-uninitialized -Wno-unused-parameter -fPIC -fstack-protector-all -D_FORTIFY_SOURCE=2 -D_GLIBCXX_USE_CXX11_ABI=0 -O2 -Dgoogle=ascend_private") + set(protobuf_LDFLAGS "-Wl,-z,relro,-z,now,-z,noexecstack") + + ExternalProject_Add(ascend_protobuf_build_transformer + URL ${REQ_URL} + DOWNLOAD_DIR ${CANN_3RD_LIB_PATH}/pkg + PATCH_COMMAND patch -p1 < ${CMAKE_CURRENT_LIST_DIR}/build/modules/patch/protobuf_25.1_change_version.patch + CONFIGURE_COMMAND ${CMAKE_COMMAND} + -DCMAKE_MESSAGE_LOG_LEVEL=ERROR + -DCMAKE_INSTALL_LIBDIR=lib + -Dprotobuf_WITH_ZLIB=OFF + -DLIB_PREFIX=ascend_ + -DCMAKE_SKIP_RPATH=TRUE + -Dprotobuf_BUILD_TESTS=OFF + -DBUILD_SHARED_LIBS=OFF + -DCMAKE_CXX_STANDARD=14 + -DCMAKE_CXX_FLAGS=${protobuf_CXXFLAGS} + -DCMAKE_CXX_LDFLAGS=${protobuf_LDFLAGS} + -DCMAKE_C_COMPILER_LAUNCHER=${CMAKE_C_COMPILER_LAUNCHER} + -DCMAKE_CXX_COMPILER_LAUNCHER=${CMAKE_CXX_COMPILER_LAUNCHER} + -DCMAKE_INSTALL_PREFIX=${ASCEND_PROTOBUF_DIR} + -Dprotobuf_BUILD_PROTOC_BINARIES=ON + -Dprotobuf_ABSL_PROVIDER=module + -DABSL_ROOT_DIR=${ABSL_SOURCE_DIR} + + SOURCE_DIR ${ASCEND_PROTOBUF_SOURCE_DIR} + BUILD_COMMAND ${CMAKE_COMMAND} --build . + INSTALL_COMMAND "" + EXCLUDE_FROM_ALL TRUE + ) + if(TARGET abseil_build_transformer) + add_dependencies(ascend_protobuf_build_transformer abseil_build_transformer) + endif() + + ExternalProject_Get_Property(ascend_protobuf_build_transformer SOURCE_DIR) + ExternalProject_Get_Property(ascend_protobuf_build_transformer BINARY_DIR) + + set(Protobuf_INCLUDE ${SOURCE_DIR}/src) + set(Protobuf_PATH ${BINARY_DIR}) + set(Protobuf_PROTOC_EXECUTABLE ${Protobuf_PATH}/protoc) + + add_custom_command( + OUTPUT ${Protobuf_PROTOC_EXECUTABLE} + DEPENDS ascend_protobuf_build_transformer + ) +endif() \ No newline at end of file diff --git a/csrc/ascend/cmake/third_party/build/modules/patch/protobuf-hide_absl_symbols.patch b/csrc/ascend/cmake/third_party/build/modules/patch/protobuf-hide_absl_symbols.patch new file mode 100644 index 000000000..2af21cf64 --- /dev/null +++ b/csrc/ascend/cmake/third_party/build/modules/patch/protobuf-hide_absl_symbols.patch @@ -0,0 +1,61 @@ +diff --git a/CMakeLists.txt b/CMakeLists.txt +index eef6626..c054264 100644 +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -18,6 +18,9 @@ + # As of 2022-09-06, CMake 3.10 is the minimum supported version. + cmake_minimum_required(VERSION 3.10) + ++# NOTICE: force using static libraries ++set(BUILD_SHARED_LIBS OFF) ++ + # Compiler id for Apple Clang is now AppleClang. + if (POLICY CMP0025) + cmake_policy(SET CMP0025 NEW) +@@ -219,6 +222,7 @@ if(ABSL_ENABLE_INSTALL) + PATTERN "*.h" + PATTERN "copts" EXCLUDE + PATTERN "testdata" EXCLUDE ++ PATTERN "options.h" EXCLUDE + ) + + file(READ "absl/base/options.h" ABSL_INTERNAL_OPTIONS_H_CONTENTS) +diff --git a/absl/base/options.h b/absl/base/options.h +index f308e1b..09617e0 100644 +--- a/absl/base/options.h ++++ b/absl/base/options.h +@@ -200,7 +200,7 @@ + // allowed. + + #define ABSL_OPTION_USE_INLINE_NAMESPACE 1 +-#define ABSL_OPTION_INLINE_NAMESPACE_NAME lts_20230802 ++#define ABSL_OPTION_INLINE_NAMESPACE_NAME lts_ascend_private + + // ABSL_OPTION_HARDENED + // +diff --git a/absl/hash/internal/hash.cc b/absl/hash/internal/hash.cc +index 11451e5..75886d3 100644 +--- a/absl/hash/internal/hash.cc ++++ b/absl/hash/internal/hash.cc +@@ -44,7 +44,7 @@ uint64_t MixingHashState::CombineLargeContiguousImpl64( + std::integral_constant{}); + } + +-ABSL_CONST_INIT const void* const MixingHashState::kSeed = &kSeed; ++ABSL_CONST_INIT const void* const MixingHashState::kSeed = 0; + + // The salt array used by LowLevelHash. This array is NOT the mechanism used to + // make absl::Hash non-deterministic between program invocations. See `Seed()` +diff --git a/absl/hash/internal/hash.h b/absl/hash/internal/hash.h +index ef3f366..7714cdb 100644 +--- a/absl/hash/internal/hash.h ++++ b/absl/hash/internal/hash.h +@@ -1159,7 +1159,7 @@ class ABSL_DLL MixingHashState : public HashStateBase { + #if (!defined(__clang__) || __clang_major__ > 11) && \ + (!defined(__apple_build_version__) || \ + __apple_build_version__ >= 19558921) // Xcode 12 +- return static_cast(reinterpret_cast(&kSeed)); ++ return static_cast(reinterpret_cast(kSeed)); + #else + // Workaround the absence of + // https://github.com/llvm/llvm-project/commit/bc15bf66dcca76cc06fe71fca35b74dc4d521021. \ No newline at end of file diff --git a/csrc/ascend/cmake/third_party/build/modules/patch/protobuf_25.1_change_version.patch b/csrc/ascend/cmake/third_party/build/modules/patch/protobuf_25.1_change_version.patch new file mode 100644 index 000000000..2754e0d7e --- /dev/null +++ b/csrc/ascend/cmake/third_party/build/modules/patch/protobuf_25.1_change_version.patch @@ -0,0 +1,357 @@ +diff --git a/CMakeLists.txt b/CMakeLists.txt +index 4137ce2..1d9e1ca 100644 +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -79,7 +79,7 @@ if (protobuf_BUILD_SHARED_LIBS) + endif () + + # Version metadata +-set(protobuf_VERSION_STRING "4.25.1") ++set(protobuf_VERSION_STRING "3.13.0") + set(protobuf_DESCRIPTION "Protocol Buffers") + set(protobuf_CONTACT "protobuf@googlegroups.com") + +@@ -105,7 +105,7 @@ message(STATUS "${protobuf_VERSION_PRERELEASE}") + + # Package version + set(protobuf_VERSION +- "${protobuf_VERSION_MINOR}.${protobuf_VERSION_PATCH}") ++ "${protobuf_VERSION_MAJOR}.${protobuf_VERSION_MINOR}.${protobuf_VERSION_PATCH}") + + if(protobuf_VERSION_PRERELEASE) + set(protobuf_VERSION "${protobuf_VERSION}.${protobuf_VERSION_PRERELEASE}") +@@ -122,30 +122,6 @@ if(protobuf_VERBOSE) + message(STATUS "]") + endif() + +-file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/cmaketest.map +-"{ +- global: +- main; +- local: +- *; +-};") +-# CheckLinkerFlag module available in CMake >=3.18. +-if(${CMAKE_VERSION} VERSION_GREATER 3.18 OR ${CMAKE_VERSION} VERSION_EQUAL 3.18) +- include(CheckLinkerFlag) +- check_linker_flag(CXX -Wl,--version-script=${CMAKE_CURRENT_BINARY_DIR}/cmaketest.map protobuf_HAVE_LD_VERSION_SCRIPT) +-else() +- include(CheckCXXSourceCompiles) +- set(OLD_CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS}) +- set(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_BINARY_DIR}/cmaketest.map) +- check_cxx_source_compiles(" +- int main() { +- return 0; +- } +- " protobuf_HAVE_LD_VERSION_SCRIPT) +- set(CMAKE_REQUIRED_FLAGS ${OLD_CMAKE_REQUIRED_FLAGS}) +-endif() +-file(REMOVE ${CMAKE_CURRENT_BINARY_DIR}/cmaketest.map) +- + find_package(Threads REQUIRED) + + # We can install dependencies from submodules if we're running +diff --git a/src/google/protobuf/arenastring.cc b/src/google/protobuf/arenastring.cc +index bf9d0bb..0e4f450 100644 +--- a/src/google/protobuf/arenastring.cc ++++ b/src/google/protobuf/arenastring.cc +@@ -11,6 +11,7 @@ + + #include "absl/log/absl_check.h" + #include "absl/strings/string_view.h" ++#include "absl/strings/escaping.h" + #include "absl/synchronization/mutex.h" + #include "google/protobuf/io/coded_stream.h" + #include "google/protobuf/message_lite.h" +@@ -116,7 +117,7 @@ void ArenaStringPtr::Set(absl::string_view value, Arena* arena) { + old->assign("garbagedata"); + } + #else // PROTOBUF_FORCE_COPY_DEFAULT_STRING +- UnsafeMutablePointer()->assign(value.data(), value.length()); ++ UnsafeMutablePointer()->assign(std::string(value.data(), value.length())); + #endif // PROTOBUF_FORCE_COPY_DEFAULT_STRING + } + } +@@ -263,6 +264,30 @@ void ArenaStringPtr::ClearToDefault(const LazyString& default_value, + } + } + ++void ArenaStringPtr::Set(EmptyDefault, const std::string& value, Arena* arena) { ++ Set(value, arena); ++} ++ ++void ArenaStringPtr::Set(EmptyDefault, absl::string_view value, Arena* arena) { ++ Set(value, arena); ++} ++ ++void ArenaStringPtr::Set(EmptyDefault, std::string&& value, Arena* arena) { ++ Set(value, arena); ++} ++ ++void ArenaStringPtr::Set(EmptyDefault, const char* s, Arena* arena) { ++ Set(s, arena); ++} ++ ++void ArenaStringPtr::Set(EmptyDefault, const char* s, size_t n, Arena* arena) { ++ Set(s, n, arena); ++} ++ ++void ArenaStringPtr::Set(EmptyDefault, const void* p, size_t n, Arena* arena) { ++ Set(absl::string_view{static_cast(p), n}, arena); ++} ++ + const char* EpsCopyInputStream::ReadArenaString(const char* ptr, + ArenaStringPtr* s, + Arena* arena) { +@@ -277,8 +302,11 @@ const char* EpsCopyInputStream::ReadArenaString(const char* ptr, + GOOGLE_PROTOBUF_PARSER_ASSERT(ptr); + return ptr; + } +- + } // namespace internal ++ ++std::string CEscape(const std::string &src) { ++ return absl::CEscape(absl::string_view(src)); ++} + } // namespace protobuf + } // namespace google + +diff --git a/src/google/protobuf/arenastring.h b/src/google/protobuf/arenastring.h +index dea2aec..ffe2343 100644 +--- a/src/google/protobuf/arenastring.h ++++ b/src/google/protobuf/arenastring.h +@@ -19,6 +19,8 @@ + #include "google/protobuf/arena.h" + #include "google/protobuf/explicitly_constructed.h" + #include "google/protobuf/port.h" ++#include "absl/strings/string_view.h" ++#include "absl/strings/escaping.h" + + // must be last: + #include "google/protobuf/port_def.inc" +@@ -298,7 +300,15 @@ struct PROTOBUF_EXPORT ArenaStringPtr { + void Set(const std::string& value, Arena* arena); + void Set(const char* s, Arena* arena); + void Set(const char* s, size_t n, Arena* arena); +- ++ // Compatibility with undefined symbol modifications in opp_kernel packages compiled based on protobuf 3.14 ++ struct EmptyDefault {}; ++ void Set(EmptyDefault, const std::string& value, Arena* arena); ++ void Set(EmptyDefault, absl::string_view value, Arena* arena); ++ void Set(EmptyDefault, std::string&& value, Arena* arena); ++ void Set(EmptyDefault, const char* s, Arena* arena); ++ void Set(EmptyDefault, const char* s, size_t n, Arena* arena); ++ void Set(EmptyDefault, const void* p, size_t n, Arena* arena); ++ + void SetBytes(absl::string_view value, Arena* arena); + void SetBytes(std::string&& value, Arena* arena); + template +@@ -528,9 +538,9 @@ inline std::string* ArenaStringPtr::UnsafeMutablePointer() { + ABSL_DCHECK(tagged_ptr_.Get() != nullptr); + return tagged_ptr_.Get(); + } +- +- + } // namespace internal ++ ++PROTOBUF_EXPORT std::string CEscape(const std::string& src); + } // namespace protobuf + } // namespace google + +diff --git a/src/google/protobuf/compiler/cpp/field_generators/string_field.cc b/src/google/protobuf/compiler/cpp/field_generators/string_field.cc +index 748a8d6..1834c39 100644 +--- a/src/google/protobuf/compiler/cpp/field_generators/string_field.cc ++++ b/src/google/protobuf/compiler/cpp/field_generators/string_field.cc +@@ -347,7 +347,7 @@ void SingularString::ReleaseImpl(io::Printer* p) const { + p->Emit(R"cc( + auto* released = $field_$.Release(); + #ifdef PROTOBUF_FORCE_COPY_DEFAULT_STRING +- $field_$.Set("", $set_args$); ++ $field_$.Set(::google::protobuf::internal::ArenaStringPtr::EmptyDefault{}, "", $set_args$); + #endif // PROTOBUF_FORCE_COPY_DEFAULT_STRING + return released; + )cc"); +@@ -393,7 +393,7 @@ void SingularString::SetAllocatedImpl(io::Printer* p) const { + p->Emit(R"cc( + #ifdef PROTOBUF_FORCE_COPY_DEFAULT_STRING + if ($field_$.IsDefault()) { +- $field_$.Set("", $set_args$); ++ $field_$.Set(::google::protobuf::internal::ArenaStringPtr::EmptyDefault{}, "", $set_args$); + } + #endif // PROTOBUF_FORCE_COPY_DEFAULT_STRING + )cc"); +@@ -442,7 +442,7 @@ void SingularString::GenerateInlineAccessorDefinitions(io::Printer* p) const { + $TsanDetectConcurrentMutation$; + $PrepareSplitMessageForWrite$; + $update_hasbit$; +- $field_$.$Set$(static_cast(arg), args..., $set_args$); ++ $field_$.Set(::google::protobuf::internal::ArenaStringPtr::EmptyDefault{}, static_cast(arg), args..., $set_args$); + $annotate_set$; + // @@protoc_insertion_point(field_set:$pkg.Msg.field$) + } +@@ -463,7 +463,7 @@ void SingularString::GenerateInlineAccessorDefinitions(io::Printer* p) const { + $update_hasbit$; + //~ Don't use $Set$ here; we always want the std::string variant + //~ regardless of whether this is a `bytes` field. +- $field_$.Set(value, $set_args$); ++ $field_$.Set(::google::protobuf::internal::ArenaStringPtr::EmptyDefault{}, value, $set_args$); + } + inline std::string* $Msg$::_internal_mutable_$name$() { + $TsanDetectConcurrentMutation$; +@@ -599,7 +599,7 @@ void SingularString::GenerateConstructorCode(io::Printer* p) const { + if (IsString(field_, *opts_) && EmptyDefault()) { + p->Emit(R"cc( + #ifdef PROTOBUF_FORCE_COPY_DEFAULT_STRING +- $field_$.Set("", GetArena()); ++ $field_$.Set(::google::protobuf::internal::ArenaStringPtr::EmptyDefault{}, "", GetArena()); + #endif // PROTOBUF_FORCE_COPY_DEFAULT_STRING + )cc"); + } +@@ -636,7 +636,7 @@ void SingularString::GenerateCopyConstructorCode(io::Printer* p) const { + }}}, + R"cc( + if ($hazzer$) { +- _this->$field_$.Set(from._internal_$name$(), $set_args$); ++ _this->$field_$.Set(::google::protobuf::internal::ArenaStringPtr::EmptyDefault{}, from._internal_$name$(), $set_args$); + } + )cc"); + } +diff --git a/src/google/protobuf/compiler/cpp/helpers.cc b/src/google/protobuf/compiler/cpp/helpers.cc +index 2c8d275..d0abac8 100644 +--- a/src/google/protobuf/compiler/cpp/helpers.cc ++++ b/src/google/protobuf/compiler/cpp/helpers.cc +@@ -1706,8 +1706,7 @@ bool IsFileDescriptorProto(const FileDescriptor* file, const Options& options) { + } + + bool ShouldGenerateClass(const Descriptor* descriptor, const Options& options) { +- return !IsMapEntryMessage(descriptor) || +- HasDescriptorMethods(descriptor->file(), options); ++ return true; + } + + } // namespace cpp +diff --git a/src/google/protobuf/compiler/cpp/helpers.h b/src/google/protobuf/compiler/cpp/helpers.h +index 9b74835..a1b6f58 100644 +--- a/src/google/protobuf/compiler/cpp/helpers.h ++++ b/src/google/protobuf/compiler/cpp/helpers.h +@@ -1002,7 +1002,7 @@ inline bool ShouldGenerateExternSpecializations(const Options& options) { + // Some compilers can't handle that much input in a single translation unit. + // These specializations are just a link size optimization and do not affect + // correctness or performance, so it is ok to omit them. +- return !options.opensource_runtime; ++ return true; + } + + struct OneOfRangeImpl { +diff --git a/src/google/protobuf/compiler/csharp/csharp_helpers.cc b/src/google/protobuf/compiler/csharp/csharp_helpers.cc +index af5c0b2..254e77c 100644 +--- a/src/google/protobuf/compiler/csharp/csharp_helpers.cc ++++ b/src/google/protobuf/compiler/csharp/csharp_helpers.cc +@@ -254,7 +254,7 @@ std::string GetFieldConstantName(const FieldDescriptor* field) { + + std::string GetPropertyName(const FieldDescriptor* descriptor) { + // Names of members declared or overridden in the message. +- static const auto& reserved_member_names = *new absl::flat_hash_set({ ++ static const auto& reserved_member_names = *new std::set({ + "Types", + "Descriptor", + "Equals", +diff --git a/src/google/protobuf/compiler/java/names.cc b/src/google/protobuf/compiler/java/names.cc +index d934973..caa3c47 100644 +--- a/src/google/protobuf/compiler/java/names.cc ++++ b/src/google/protobuf/compiler/java/names.cc +@@ -37,7 +37,7 @@ const char* DefaultPackage(Options options) { + + bool IsReservedName(absl::string_view name) { + static const auto& kReservedNames = +- *new absl::flat_hash_set({ ++ *new std::set({ + "abstract", "assert", "boolean", "break", "byte", + "case", "catch", "char", "class", "const", + "continue", "default", "do", "double", "else", +@@ -49,7 +49,7 @@ bool IsReservedName(absl::string_view name) { + "switch", "synchronized", "this", "throw", "throws", + "transient", "try", "void", "volatile", "while", + }); +- return kReservedNames.contains(name); ++ return kReservedNames.count(name) != 0U; + } + + bool IsForbidden(absl::string_view field_name) { +@@ -59,7 +59,7 @@ bool IsForbidden(absl::string_view field_name) { + // Keep this list in sync with specialFieldNames in + // java/core/src/main/java/com/google/protobuf/DescriptorMessageInfoFactory.java + static const auto& kForbiddenNames = +- *new absl::flat_hash_set({ ++ *new std::set({ + // java.lang.Object: + "Class", + // com.google.protobuf.MessageLiteOrBuilder: +@@ -75,7 +75,7 @@ bool IsForbidden(absl::string_view field_name) { + // obsolete. kept for backwards compatibility of generated code + "CachedSize", + }); +- return kForbiddenNames.contains(UnderscoresToCamelCase(field_name, true)); ++ return kForbiddenNames.count(UnderscoresToCamelCase(field_name, true)) != 0U; + } + + std::string FieldName(const FieldDescriptor* field) { +diff --git a/src/google/protobuf/extension_set.h b/src/google/protobuf/extension_set.h +index f45a1c8..e2191b7 100644 +--- a/src/google/protobuf/extension_set.h ++++ b/src/google/protobuf/extension_set.h +@@ -166,7 +166,7 @@ class PROTOBUF_EXPORT GeneratedExtensionFinder { + // off to the ExtensionSet for parsing. Etc. + class PROTOBUF_EXPORT ExtensionSet { + public: +- constexpr ExtensionSet() : ExtensionSet(nullptr) {} ++ constexpr ExtensionSet(); + ExtensionSet(const ExtensionSet& rhs) = delete; + + // Arena enabled constructors: for internal use only. +@@ -912,6 +912,9 @@ class PROTOBUF_EXPORT ExtensionSet { + static void DeleteFlatMap(const KeyValue* flat, uint16_t flat_capacity); + }; + ++constexpr ExtensionSet::ExtensionSet() ++ : arena_(nullptr), flat_capacity_(0), flat_size_(0), map_{nullptr} {} ++ + constexpr ExtensionSet::ExtensionSet(Arena* arena) + : arena_(arena), flat_capacity_(0), flat_size_(0), map_{nullptr} {} + +diff --git a/src/google/protobuf/message_lite.h b/src/google/protobuf/message_lite.h +index de8bd19..60bc923 100644 +--- a/src/google/protobuf/message_lite.h ++++ b/src/google/protobuf/message_lite.h +@@ -505,12 +505,10 @@ class PROTOBUF_EXPORT MessageLite { + return Arena::CreateMaybeMessage(arena); + } + +-#ifdef PROTOBUF_EXPLICIT_CONSTRUCTORS + template + static T* CreateMaybeMessage(Arena* arena, const T& from) { + return Arena::CreateMaybeMessage(arena, from); + } +-#endif // PROTOBUF_EXPLICIT_CONSTRUCTORS + + inline explicit MessageLite(Arena* arena) : _internal_metadata_(arena) {} + +diff --git a/src/google/protobuf/port_def.inc b/src/google/protobuf/port_def.inc +index 25abc30..f17d861 100644 +--- a/src/google/protobuf/port_def.inc ++++ b/src/google/protobuf/port_def.inc +@@ -572,14 +572,6 @@ static_assert(PROTOBUF_ABSL_MIN(20230125, 3), + #error PROTOBUF_FORCE_COPY_DEFAULT_STRING was previously defined + #endif + +-// Set (and leave set) PROTOBUF_EXPLICIT_CONSTRUCTORS +-// This flag changes the internal implementation class and its uses from using +-// aggregate initialization to using explicit constructor based initialization. +-// TODO: remove this flag and make the code permanent. +-#ifndef PROTOBUF_EXPLICIT_CONSTRUCTORS +-#define PROTOBUF_EXPLICIT_CONSTRUCTORS +-#endif +- + #ifdef PROTOBUF_FORCE_ALLOCATION_ON_CONSTRUCTION + #error PROTOBUF_FORCE_ALLOCATION_ON_CONSTRUCTION was previously defined + #endif \ No newline at end of file diff --git a/csrc/ascend/cmake/third_party/gtest.cmake b/csrc/ascend/cmake/third_party/gtest.cmake new file mode 100644 index 000000000..517d6f866 --- /dev/null +++ b/csrc/ascend/cmake/third_party/gtest.cmake @@ -0,0 +1,115 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +include_guard(GLOBAL) + +unset(gtest_FOUND CACHE) +unset(GTEST_INCLUDE CACHE) +unset(GTEST_STATIC_LIBRARY CACHE) +unset(GTEST_MAIN_STATIC_LIBRARY CACHE) +unset(GMOCK_STATIC_LIBRARY CACHE) +unset(GMOCK_MAIN_STATIC_LIBRARY CACHE) + +set(GTEST_INSTALL_PATH ${CANN_3RD_LIB_PATH}/gtest) +message("GTEST_INSTALL_PATH=${GTEST_INSTALL_PATH}") +find_path(GTEST_INCLUDE + NAMES gtest/gtest.h + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH + PATHS ${GTEST_INSTALL_PATH}/include) +find_library(GTEST_STATIC_LIBRARY + NAMES libgtest.a + PATH_SUFFIXES lib lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH + PATHS ${GTEST_INSTALL_PATH}) +find_library(GTEST_MAIN_STATIC_LIBRARY + NAMES libgtest_main.a + PATH_SUFFIXES lib lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH + PATHS ${GTEST_INSTALL_PATH}) +find_library(GMOCK_STATIC_LIBRARY + NAMES libgmock.a + PATH_SUFFIXES lib lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH + PATHS ${GTEST_INSTALL_PATH}) +find_library(GMOCK_MAIN_STATIC_LIBRARY + NAMES libgmock_main.a + PATH_SUFFIXES lib lib64 + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH + PATHS ${GTEST_INSTALL_PATH}) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(gtest + FOUND_VAR + gtest_FOUND + REQUIRED_VARS + GTEST_INCLUDE + GTEST_STATIC_LIBRARY + GTEST_MAIN_STATIC_LIBRARY + GMOCK_STATIC_LIBRARY + GMOCK_MAIN_STATIC_LIBRARY + ) +message("gtest found:${gtest_FOUND}") + +if(gtest_FOUND AND NOT FORCE_REBUILD_CANN_3RD) + message("gtest found in ${GTEST_INSTALL_PATH}, and not force rebuild cann third_party") +else() + set(REQ_URL "https://gitcode.com/cann-src-third-party/googletest/releases/download/v1.14.0/googletest-1.14.0.tar.gz") + set (gtest_CXXFLAGS "-D_GLIBCXX_USE_CXX11_ABI=0 -O2 -D_FORTIFY_SOURCE=2 -fPIC -fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack") + set (gtest_CFLAGS "-D_GLIBCXX_USE_CXX11_ABI=0 -O2 -D_FORTIFY_SOURCE=2 -fPIC -fstack-protector-all -Wl,-z,relro,-z,now,-z,noexecstack") + + include(ExternalProject) + ExternalProject_Add(third_party_gtest + URL ${REQ_URL} + TLS_VERIFY OFF + DOWNLOAD_DIR ${CANN_3RD_PKG_PATH} + CONFIGURE_COMMAND ${CMAKE_COMMAND} + -DCMAKE_CXX_FLAGS=${gtest_CXXFLAGS} + -DCMAKE_C_FLAGS=${gtest_CFLAGS} + -DCMAKE_INSTALL_PREFIX=${GTEST_INSTALL_PATH} + -DCMAKE_INSTALL_LIBDIR=lib + -DBUILD_SHARED_LIBS=OFF + + BUILD_COMMAND $(MAKE) + INSTALL_COMMAND $(MAKE) install + EXCLUDE_FROM_ALL TRUE + ) +endif() + +set(GTEST_INCLUDE ${GTEST_INSTALL_PATH}/include) + +add_library(gtest STATIC IMPORTED) +add_dependencies(gtest third_party_gtest) + +add_library(gmock STATIC IMPORTED) +add_dependencies(gmock third_party_gtest) + +add_library(gtest_main STATIC IMPORTED) +add_dependencies(gtest_main third_party_gtest) + +if (NOT EXISTS ${GTEST_INSTALL_PATH}/include) + file(MAKE_DIRECTORY "${GTEST_INSTALL_PATH}/include") +endif () + +set_target_properties(gtest PROPERTIES + IMPORTED_LOCATION ${GTEST_INSTALL_PATH}/lib/libgtest.a + INTERFACE_INCLUDE_DIRECTORIES ${GTEST_INSTALL_PATH}/include) + +set_target_properties(gmock PROPERTIES + IMPORTED_LOCATION ${GTEST_INSTALL_PATH}/lib/libgmock.a + INTERFACE_INCLUDE_DIRECTORIES ${GTEST_INSTALL_PATH}/include) + +set_target_properties(gtest_main PROPERTIES + IMPORTED_LOCATION ${GTEST_INSTALL_PATH}/lib/libgtest_main.a + INTERFACE_INCLUDE_DIRECTORIES ${GTEST_INSTALL_PATH}/include) diff --git a/csrc/ascend/cmake/third_party/json.cmake b/csrc/ascend/cmake/third_party/json.cmake new file mode 100644 index 000000000..da073c08b --- /dev/null +++ b/csrc/ascend/cmake/third_party/json.cmake @@ -0,0 +1,69 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +include_guard(GLOBAL) + +if(json_FOUND) + return() +endif() + +unset(json_FOUND CACHE) +unset(JSON_INCLUDE CACHE) + +if(NOT CANN_3RD_PKG_PATH) + set(CANN_3RD_PKG_PATH ${PROJECT_SOURCE_DIR}/third_party/pkg) +endif() + +set(JSON_DOWNLOAD_PATH ${CANN_3RD_LIB_PATH}/pkg) +set(JSON_INSTALL_PATH ${CANN_3RD_LIB_PATH}/json) + +find_path(JSON_INCLUDE + NAMES nlohmann/json.hpp + NO_CMAKE_SYSTEM_PATH + NO_CMAKE_FIND_ROOT_PATH + PATHS ${JSON_INSTALL_PATH}/include) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(json + FOUND_VAR + json_FOUND + REQUIRED_VARS + JSON_INCLUDE + ) + +if(json_FOUND AND NOT FORCE_REBUILD_CANN_3RD) + message("json found in ${JSON_INSTALL_PATH}, and not force rebuild cann third_party") + set(JSON_INCLUDE_DIR ${JSON_INSTALL_PATH}/include) + add_library(json INTERFACE IMPORTED) +else() + set(REQ_URL "https://gitcode.com/cann-src-third-party/json/releases/download/v3.11.3/include.zip") + + include(ExternalProject) + ExternalProject_Add(third_party_json + URL ${REQ_URL} + TLS_VERIFY OFF + DOWNLOAD_DIR ${JSON_DOWNLOAD_PATH} + DOWNLOAD_NO_EXTRACT TRUE + SOURCE_DIR ${JSON_INSTALL_PATH} + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + INSTALL_COMMAND + ${CMAKE_COMMAND} -E make_directory ${JSON_INSTALL_PATH} && + ${CMAKE_COMMAND} -E chdir ${JSON_INSTALL_PATH} ${CMAKE_COMMAND} -E tar xf "${JSON_DOWNLOAD_PATH}/include.zip" --format=zip + UPDATE_COMMAND "" + ) + + ExternalProject_Get_Property(third_party_json SOURCE_DIR) + ExternalProject_Get_Property(third_party_json BINARY_DIR) + set(JSON_INCLUDE_DIR ${SOURCE_DIR}/include) + add_library(json INTERFACE) + target_include_directories(json INTERFACE ${JSON_INCLUDE_DIR}) + add_dependencies(json third_party_json) +endif() diff --git a/csrc/ascend/cmake/third_party/makeself-fetch.cmake b/csrc/ascend/cmake/third_party/makeself-fetch.cmake new file mode 100644 index 000000000..9759b6a9d --- /dev/null +++ b/csrc/ascend/cmake/third_party/makeself-fetch.cmake @@ -0,0 +1,35 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +set(MAKESELF_NAME "makeself") +set(MAKESELF_PATH ${CANN_3RD_LIB_PATH}/makeself) + +# 默认配置的makeself还是不存在则下载 +if (NOT EXISTS "${MAKESELF_PATH}/makeself-header.sh" OR NOT EXISTS "${MAKESELF_PATH}/makeself.sh") + set(MAKESELF_URL "https://gitcode.com/cann-src-third-party/makeself/releases/download/release-2.5.0-patch1.0/makeself-release-2.5.0-patch1.tar.gz") + message(STATUS "Downloading ${MAKESELF_NAME} from ${MAKESELF_URL}") + + include(FetchContent) + FetchContent_Declare( + ${MAKESELF_NAME} + URL ${MAKESELF_URL} + URL_HASH SHA256=bfa730a5763cdb267904a130e02b2e48e464986909c0733ff1c96495f620369a + SOURCE_DIR "${MAKESELF_PATH}" # 直接解压到此目录 + ) + FetchContent_MakeAvailable(${MAKESELF_NAME}) + execute_process( + COMMAND chmod 700 "${CMAKE_BINARY_DIR}/makeself/makeself.sh" + COMMAND chmod 700 "${CMAKE_BINARY_DIR}/makeself/makeself-header.sh" + -E env + CMAKE_TLS_VERIFY=0 + RESULT_VARIABLE CHMOD_RESULT + ERROR_VARIABLE CHMOD_ERROR + ) +endif() \ No newline at end of file diff --git a/csrc/ascend/cmake/third_party/protobuf.cmake b/csrc/ascend/cmake/third_party/protobuf.cmake new file mode 100644 index 000000000..19d095663 --- /dev/null +++ b/csrc/ascend/cmake/third_party/protobuf.cmake @@ -0,0 +1,43 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +set(_protobuf_url "") +if(CANN_PKG_SERVER) + set(_protobuf_url "${CANN_PKG_SERVER}/libs/protobuf/v25.1.tar.gz") +endif() +include(ExternalProject) +ExternalProject_Add(external_protobuf + URL ${_protobuf_url} + https://gitcode.com/cann-src-third-party/protobuf/releases/download/v25.1/protobuf-25.1.tar.gz + DOWNLOAD_DIR download/protobuf + PREFIX third_party + SOURCE_SUBDIR cmake + CMAKE_CACHE_ARGS + -DProtobuf_BUILD_TYPE:STRING=${CMAKE_BUILD_TYPE} + -Dprotobuf_BUILD_TESTS:BOOL=OFF + -Dprotobuf_BUILD_EXAMPLES:BOOL=OFF + -Dprotobuf_BUILD_SHARED_LIBS:BOOL=OFF + -DProtobuf_CXX_COMPILER:STRING=${CMAKE_CXX_COMPILER} + INSTALL_COMMAND "" +) + +ExternalProject_Get_Property(external_protobuf SOURCE_DIR) +ExternalProject_Get_Property(external_protobuf BINARY_DIR) + +set(Protobuf_INCLUDE ${SOURCE_DIR}/src) +set(Protobuf_PATH ${BINARY_DIR}) +set(Protobuf_PROTOC_EXECUTABLE ${Protobuf_PATH}/protoc) + +add_custom_command( + OUTPUT ${Protobuf_PROTOC_EXECUTABLE} + DEPENDS external_protobuf +) +add_custom_target( + protoc ALL DEPENDS ${Protobuf_PROTOC_EXECUTABLE} +) diff --git a/csrc/ascend/cmake/third_party/secure_c.cmake b/csrc/ascend/cmake/third_party/secure_c.cmake new file mode 100644 index 000000000..39c65f069 --- /dev/null +++ b/csrc/ascend/cmake/third_party/secure_c.cmake @@ -0,0 +1,108 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +set(_secure_c_url "") +if(CANN_PKG_SERVER) + set(_secure_c_url "${CANN_PKG_SERVER}/libs/securec/v1.1.10.tar.gz") +endif() +include(ExternalProject) +ExternalProject_Add(secure_c + URL ${_secure_c_url} + https://gitee.com/openeuler/libboundscheck/repository/archive/v1.1.16.tar.gz + URL_MD5 ae4865cec1bfb52f7dca03f5c05ac98a + DOWNLOAD_DIR download/secure_c + PREFIX third_party + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + INSTALL_COMMAND "" +) + +ExternalProject_Get_Property(secure_c SOURCE_DIR) +ExternalProject_Get_Property(secure_c BINARY_DIR) + +set(SEC_C_SRCS + ${SOURCE_DIR}/src/fscanf_s.c + ${SOURCE_DIR}/src/fwscanf_s.c + ${SOURCE_DIR}/src/gets_s.c + ${SOURCE_DIR}/src/input.inl + ${SOURCE_DIR}/src/memcpy_s.c + ${SOURCE_DIR}/src/memmove_s.c + ${SOURCE_DIR}/src/memset_s.c + ${SOURCE_DIR}/src/output.inl + ${SOURCE_DIR}/src/scanf_s.c + ${SOURCE_DIR}/src/secinput.h + ${SOURCE_DIR}/src/securecutil.c + ${SOURCE_DIR}/src/securecutil.h + ${SOURCE_DIR}/src/secureinput_a.c + ${SOURCE_DIR}/src/secureinput_w.c + ${SOURCE_DIR}/src/secureprintoutput_a.c + ${SOURCE_DIR}/src/secureprintoutput.h + ${SOURCE_DIR}/src/secureprintoutput_w.c + ${SOURCE_DIR}/src/snprintf_s.c + ${SOURCE_DIR}/src/sprintf_s.c + ${SOURCE_DIR}/src/sscanf_s.c + ${SOURCE_DIR}/src/strcat_s.c + ${SOURCE_DIR}/src/strcpy_s.c + ${SOURCE_DIR}/src/strncat_s.c + ${SOURCE_DIR}/src/strncpy_s.c + ${SOURCE_DIR}/src/strtok_s.c + ${SOURCE_DIR}/src/swprintf_s.c + ${SOURCE_DIR}/src/swscanf_s.c + ${SOURCE_DIR}/src/vfscanf_s.c + ${SOURCE_DIR}/src/vfwscanf_s.c + ${SOURCE_DIR}/src/vscanf_s.c + ${SOURCE_DIR}/src/vsnprintf_s.c + ${SOURCE_DIR}/src/vsprintf_s.c + ${SOURCE_DIR}/src/vsscanf_s.c + ${SOURCE_DIR}/src/vswprintf_s.c + ${SOURCE_DIR}/src/vswscanf_s.c + ${SOURCE_DIR}/src/vwscanf_s.c + ${SOURCE_DIR}/src/wcscat_s.c + ${SOURCE_DIR}/src/wcscpy_s.c + ${SOURCE_DIR}/src/wcsncat_s.c + ${SOURCE_DIR}/src/wcsncpy_s.c + ${SOURCE_DIR}/src/wcstok_s.c + ${SOURCE_DIR}/src/wmemcpy_s.c + ${SOURCE_DIR}/src/wmemmove_s.c + ${SOURCE_DIR}/src/wscanf_s.c +) + +add_library(c_sec SHARED ${SEC_C_SRCS}) + +add_dependencies(c_sec secure_c) + +set_source_files_properties( + ${SEC_C_SRCS} + PROPERTIES + GENERATED TRUE +) + +target_include_directories(c_sec + PUBLIC + "${SOURCE_DIR}/include" +) + +set(C_SEC_INCLUDE ${SOURCE_DIR}/include) + +target_compile_options(c_sec + PRIVATE + -fstack-protector-strong -fPIC -Wall -D_FORTIFY_SOURCE=2 -O2 +) + +set_target_properties(c_sec + PROPERTIES + LINKER_LANGUAGE C + LIBRARY_OUTPUT_DIRECTORY ${BINARY_DIR} +) + +cann_install( + TARGET c_sec + FILES $ + DESTINATION "${CANN_ROOT}/lib" +) diff --git a/csrc/ascend/cmake/ut.cmake b/csrc/ascend/cmake/ut.cmake new file mode 100644 index 000000000..6d42c2e0b --- /dev/null +++ b/csrc/ascend/cmake/ut.cmake @@ -0,0 +1,554 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +include_guard(GLOBAL) + +function(register_op_name op_name) + get_property(op_list GLOBAL PROPERTY ALL_OP_LIST) + if(NOT op_list) + set(op_list "") + endif() + + list(FIND op_list "${op_name}" idx) + if(idx EQUAL -1) + list(APPEND op_list "${op_name}") + set_property(GLOBAL PROPERTY ALL_OP_LIST "${op_list}") + endif() +endfunction() + +if(UT_TEST_ALL OR OP_HOST_UT) + set(OP_TILING_MODULE_NAME + ${PKG_NAME}_op_tiling_ut + CACHE STRING "op_tiling ut module name" FORCE + ) + set(OP_INFERSHAPE_MODULE_NAME + ${PKG_NAME}_op_infershape_ut + CACHE STRING "op_infershape ut module name" FORCE + ) + function(add_optiling_ut_modules OP_TILING_MODULE_NAME) + # add optiling ut common object: transformer_op_tiling_ut_common_obj + add_library(${OP_TILING_MODULE_NAME}_common_obj OBJECT) + file(GLOB OP_TILING_UT_COMMON_SRC ${UT_COMMON_INC}/tiling_context_faker.cpp + ${UT_COMMON_INC}/tiling_case_executor.cpp + ) + target_sources(${OP_TILING_MODULE_NAME}_common_obj PRIVATE ${OP_TILING_UT_COMMON_SRC}) + target_include_directories( + ${OP_TILING_MODULE_NAME}_common_obj PRIVATE ${JSON_INCLUDE_DIR} ${GTEST_INCLUDE} + ${ASCEND_DIR}/include/base/context_builder ${ASCEND_DIR}/pkg_inc + ${ASCEND_DIR}/include/ascendc/basic_api + ) + target_link_libraries( + ${OP_TILING_MODULE_NAME}_common_obj PRIVATE $ json gtest c_sec + ) + + # add optiling ut cases object: transformer_op_tiling_ut_cases_obj + if(NOT TARGET ${OP_TILING_MODULE_NAME}_cases_obj) + add_library(${OP_TILING_MODULE_NAME}_cases_obj OBJECT ${UT_PATH}/empty.cpp) + endif() + target_include_directories( + ${OP_TILING_MODULE_NAME}_cases_obj PRIVATE ${UT_COMMON_INC} ${GTEST_INCLUDE} ${ASCEND_DIR}/include + ${ASCEND_DIR}/include/base/context_builder ${OP_TILING_INCLUDE} + $<$:$> + ${ASCEND_DIR}/include/ascendc/basic_api + ) + target_compile_definitions(${OP_TILING_MODULE_NAME}_cases_obj PRIVATE + LOG_CPP + ) + target_link_libraries(${OP_TILING_MODULE_NAME}_cases_obj + PRIVATE $ + $<$:$> + gtest) + + # add op tiling ut cases static lib: libtransformer_op_tiling_ut_cases.a + add_library(${OP_TILING_MODULE_NAME}_cases STATIC) + target_link_libraries( + ${OP_TILING_MODULE_NAME}_cases PRIVATE ${OP_TILING_MODULE_NAME}_common_obj ${OP_TILING_MODULE_NAME}_cases_obj + ) + endfunction() + + function(add_infershape_ut_modules OP_INFERSHAPE_MODULE_NAME) + # add opinfershape ut common object: transformer_op_infershape_ut_common_obj + add_library(${OP_INFERSHAPE_MODULE_NAME}_common_obj OBJECT) + file(GLOB OP_INFERSHAPE_UT_COMMON_SRC + ${UT_COMMON_INC}/infer_shape_context_faker.cpp + ${UT_COMMON_INC}/infer_shape_case_executor.cpp + ${UT_COMMON_INC}/infer_datatype_context_faker.cpp + ${UT_COMMON_INC}/infer_shaperange_context_faker.cpp + ) + target_sources(${OP_INFERSHAPE_MODULE_NAME}_common_obj PRIVATE ${OP_INFERSHAPE_UT_COMMON_SRC}) + target_include_directories( + ${OP_INFERSHAPE_MODULE_NAME}_common_obj PRIVATE ${ASCEND_DIR}/include/base/context_builder + ${ASCEND_DIR}/pkg_inc + ${ASCEND_DIR}/include/ascendc/basic_api + ) + target_link_libraries( + ${OP_INFERSHAPE_MODULE_NAME}_common_obj PRIVATE $ json gtest c_sec + ) + + # add opinfershape ut cases object: transformer_op_infershape_ut_cases_obj + if(NOT TARGET ${OP_INFERSHAPE_MODULE_NAME}_cases_obj) + add_library(${OP_INFERSHAPE_MODULE_NAME}_cases_obj OBJECT ${UT_PATH}/empty.cpp) + endif() + target_include_directories( + ${OP_INFERSHAPE_MODULE_NAME}_cases_obj PRIVATE ${UT_COMMON_INC} ${GTEST_INCLUDE} ${ASCEND_DIR}/include + ${ASCEND_DIR}/pkg_inc ${ASCEND_DIR}/include/base/context_builder + ${OPBASE_INC_DIRS} + ${ASCEND_DIR}/include/ascendc/basic_api + ) + target_link_libraries( + ${OP_INFERSHAPE_MODULE_NAME}_cases_obj + PRIVATE $ + $<$:$> + gtest + ) + + # add op infershape ut cases static lib: libtransformer_op_infershape_ut_cases.a + add_library(${OP_INFERSHAPE_MODULE_NAME}_cases STATIC) + target_link_libraries( + ${OP_INFERSHAPE_MODULE_NAME}_cases PRIVATE ${OP_INFERSHAPE_MODULE_NAME}_common_obj + ${OP_INFERSHAPE_MODULE_NAME}_cases_obj + ) + endfunction() +endif() + +if(UT_TEST_ALL OR OP_API_UT) + set(OP_API_MODULE_NAME + ${PKG_NAME}_op_api_ut + CACHE STRING "op_api ut module name" FORCE + ) + function(add_opapi_ut_modules OP_API_MODULE_NAME) + # add opapi ut L2 obj + if(NOT TARGET ${OP_API_MODULE_NAME}_cases_obj) + add_library(${OP_API_MODULE_NAME}_cases_obj OBJECT) + endif() + target_sources(${OP_API_MODULE_NAME}_cases_obj PRIVATE + ${UT_PATH}/op_api/stub/opdev/platform.cpp + ${UT_PATH}/op_api/stub/opdev/nnopbase.cpp) + target_include_directories( + ${OP_API_MODULE_NAME}_cases_obj + PRIVATE ${JSON_INCLUDE_DIR} ${HI_PYTHON_INC_TEMP} ${UT_PATH}/op_api/stub ${OP_API_UT_COMMON_INC} + ${ASCEND_DIR}/include ${ASCEND_DIR}/include/aclnn ${ASCEND_DIR}/include/aclnnop + ${OPAPI_INCLUDE} + ${ASCEND_DIR}/pkg_inc + ${ASCEND_DIR}/include/ascendc/basic_api + ${ASCEND_CANN_PACKAGE_PATH}/runtime/pkg_inc + ) + target_link_libraries(${OP_API_MODULE_NAME}_cases_obj + PRIVATE $ + $<$:$> + gtest) + endfunction() +endif() + +if(UT_TEST_ALL OR OP_KERNEL_UT) + set(OP_KERNEL_MODULE_NAME + ${PKG_NAME}_op_kernel_ut + CACHE STRING "op_kernel ut module name" FORCE + ) + function(add_opkernel_ut_modules OP_KERNEL_MODULE_NAME) + # add opkernel ut common object: transformer_op_kernel_ut_common_obj + add_library(${OP_KERNEL_MODULE_NAME}_common_obj OBJECT) + file(GLOB OP_KERNEL_UT_COMMON_SRC ${UT_COMMON_INC}/tiling_context_faker.cpp + ${UT_COMMON_INC}/tiling_case_executor.cpp ${PROJECT_SOURCE_DIR}/tests/ut/framework_normal/op_kernel/data_utils.cpp + ) + target_sources(${OP_KERNEL_MODULE_NAME}_common_obj PRIVATE ${OP_KERNEL_UT_COMMON_SRC}) + target_include_directories( + ${OP_KERNEL_MODULE_NAME}_common_obj PRIVATE ${JSON_INCLUDE_DIR} ${GTEST_INCLUDE} + ${OPBASE_INC_DIRS} ${ASCEND_DIR}/include/base + ${ASCEND_DIR}/include/base/context_builder + ${ASCEND_DIR}/include/experiment + ${ASCEND_DIR}/include/experiment/metadef/common/util + ) + target_link_libraries( + ${OP_KERNEL_MODULE_NAME}_common_obj PRIVATE $ json gtest c_sec + ) + + foreach(socVersion ${fastOpTestSocVersions}) + # add op kernel ut cases obj: transformer_op_tiling_ut_${socVersion}_cases + if(NOT TARGET ${OP_KERNEL_MODULE_NAME}_${socVersion}_cases_obj) + add_library(${OP_KERNEL_MODULE_NAME}_${socVersion}_cases_obj OBJECT) + endif() + target_link_libraries(${OP_KERNEL_MODULE_NAME}_${socVersion}_cases_obj PRIVATE gcov) + + # add op kernel ut cases dynamic lib: libtransformer_op_tiling_ut_${socVersion}_cases.so + add_library( + ${OP_KERNEL_MODULE_NAME}_${socVersion}_cases SHARED + $ + $ + ) + target_link_libraries( + ${OP_KERNEL_MODULE_NAME}_${socVersion}_cases + PRIVATE $ ${OP_KERNEL_MODULE_NAME}_common_obj + ${OP_KERNEL_MODULE_NAME}_${socVersion}_cases_obj + ) + endforeach() + endfunction() +endif() + +if(UT_TEST_ALL + OR OP_HOST_UT + OR OP_API_UT + ) + function(add_modules_ut_sources) + set(options OPTION_RESERVED) + set(oneValueArgs UT_NAME MODE DIR) + set(multiValueArgs MULIT_RESERVED) + cmake_parse_arguments(MODULE "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + + get_filename_component(ARCH_NAME ${MODULE_DIR} NAME) + if(${ARCH_NAME} STREQUAL "op_host") + # tiling.cpp, infershape.cpp under op_host/ + if("${MODULE_UT_NAME}" STREQUAL "${OP_TILING_MODULE_NAME}") + get_filename_component(UT_DIR ${MODULE_DIR} DIRECTORY) + get_filename_component(TESTS_DIR ${UT_DIR} DIRECTORY) + get_filename_component(OP_NAME_DIR ${TESTS_DIR} DIRECTORY) + get_filename_component(OP_NAME ${OP_NAME_DIR} NAME) + list(FIND ASCEND_OP_NAME ${OP_NAME} INDEX) + # if "--ops" is not NULL, opName not include, jump over. if "--ops" is NULL, include all. + if(NOT "${ASCEND_OP_NAME}" STREQUAL "ALL" AND INDEX EQUAL -1) + return() + endif() + + # add op_tiling ut common object: transformer_op_tiling_ut_cases_obj + if(NOT TARGET ${MODULE_UT_NAME}_cases_obj) + add_library(${MODULE_UT_NAME}_cases_obj OBJECT) + endif() + file(GLOB OPHOST_TILING_CASES_SRC ${MODULE_DIR}/test_*_tiling.cpp) + target_sources(${MODULE_UT_NAME}_cases_obj ${MODULE_MODE} ${OPHOST_TILING_CASES_SRC}) + endif() + + if("${MODULE_UT_NAME}" STREQUAL "${OP_INFERSHAPE_MODULE_NAME}") + get_filename_component(UT_DIR ${MODULE_DIR} DIRECTORY) + get_filename_component(TESTS_DIR ${UT_DIR} DIRECTORY) + get_filename_component(OP_NAME_DIR ${TESTS_DIR} DIRECTORY) + get_filename_component(OP_NAME ${OP_NAME_DIR} NAME) + list(FIND ASCEND_OP_NAME ${OP_NAME} INDEX) + # if "--ops" is not NULL, opName not include, jump over. if "--ops" is NULL, include all. + if(NOT "${ASCEND_OP_NAME}" STREQUAL "ALL" AND INDEX EQUAL -1) + return() + endif() + + # add op_infershape ut common object: transformer_op_infershape_ut_cases_obj + if(NOT TARGET ${MODULE_UT_NAME}_cases_obj) + add_library(${MODULE_UT_NAME}_cases_obj OBJECT) + endif() + file(GLOB OPHOST_INFERSHAPE_CASES_SRC ${MODULE_DIR}/test_*_infershape.cpp) + target_sources(${MODULE_UT_NAME}_cases_obj ${MODULE_MODE} ${OPHOST_INFERSHAPE_CASES_SRC}) + endif() + endif() + + # tiling.cpp, infershape.cpp under arch*/ + if(${ARCH_NAME} MATCHES "^arch") + # if "--soc" not include current soc, jump over + list(FIND ARCH_DIRECTORY ${ARCH_NAME} INDEX) + if(INDEX EQUAL -1) + return() + endif() + + if("${MODULE_UT_NAME}" STREQUAL "${OP_TILING_MODULE_NAME}") + get_filename_component(UT_TYPE_DIR ${MODULE_DIR} DIRECTORY) + get_filename_component(UT_DIR ${UT_TYPE_DIR} DIRECTORY) + get_filename_component(TESTS_DIR ${UT_DIR} DIRECTORY) + get_filename_component(OP_NAME_DIR ${TESTS_DIR} DIRECTORY) + get_filename_component(OP_NAME ${OP_NAME_DIR} NAME) + list(FIND ASCEND_OP_NAME ${OP_NAME} INDEX) + # if "--ops" is not NULL, opName not include, jump over. if "--ops" is NULL, include all. + if(NOT "${ASCEND_OP_NAME}" STREQUAL "ALL" AND INDEX EQUAL -1) + return() + endif() + + # add op_tiling ut common object: transformer_op_tiling_ut_cases_obj + if(NOT TARGET ${MODULE_UT_NAME}_cases_obj) + add_library(${MODULE_UT_NAME}_cases_obj OBJECT) + endif() + file(GLOB OPHOST_TILING_CASES_SRC ${MODULE_DIR}/test_*_tiling.cpp) + target_sources(${MODULE_UT_NAME}_cases_obj ${MODULE_MODE} ${OPHOST_TILING_CASES_SRC}) + endif() + + if("${MODULE_UT_NAME}" STREQUAL "${OP_INFERSHAPE_MODULE_NAME}") + get_filename_component(UT_TYPE_DIR ${MODULE_DIR} DIRECTORY) + get_filename_component(UT_DIR ${UT_TYPE_DIR} DIRECTORY) + get_filename_component(TESTS_DIR ${UT_DIR} DIRECTORY) + get_filename_component(OP_NAME_DIR ${TESTS_DIR} DIRECTORY) + get_filename_component(OP_NAME ${OP_NAME_DIR} NAME) + + list(FIND ASCEND_OP_NAME ${OP_NAME} INDEX) + # if "--ops" is not NULL, opName not include, jump over. if "--ops" is NULL, include all. + if(NOT "${ASCEND_OP_NAME}" STREQUAL "ALL" AND INDEX EQUAL -1) + return() + endif() + + # add op_infershape ut common object: transformer_op_infershape_ut_cases_obj + if(NOT TARGET ${MODULE_UT_NAME}_cases_obj) + add_library(${MODULE_UT_NAME}_cases_obj OBJECT) + endif() + file(GLOB OPHOST_INFERSHAPE_CASES_SRC ${MODULE_DIR}/test_*_infershape.cpp) + target_sources(${MODULE_UT_NAME}_cases_obj ${MODULE_MODE} ${OPHOST_INFERSHAPE_CASES_SRC}) + endif() + endif() + + # op_api ut + if("${MODULE_UT_NAME}" STREQUAL "${OP_API_MODULE_NAME}") + get_filename_component(OP_HOST_DIR ${MODULE_DIR} DIRECTORY) + get_filename_component(OP_HOST_NAME ${OP_HOST_DIR} NAME) + if("${OP_HOST_NAME}" STREQUAL "op_host") + get_filename_component(UT_DIR ${OP_HOST_DIR} DIRECTORY) + else() + get_filename_component(UT_DIR ${MODULE_DIR} DIRECTORY) + endif() + get_filename_component(TESTS_DIR ${UT_DIR} DIRECTORY) + get_filename_component(OP_NAME_DIR ${TESTS_DIR} DIRECTORY) + get_filename_component(OP_NAME ${OP_NAME_DIR} NAME) + list(FIND ASCEND_OP_NAME ${OP_NAME} INDEX) + # if "--ops" is not NULL, opName not include, jump over. if "--ops" is NULL, include all. + if(NOT "${ASCEND_OP_NAME}" STREQUAL "ALL" AND INDEX EQUAL -1) + return() + endif() + + # add op_api ut common object: transformer_op_api_ut_cases_obj + register_op_name(${OP_NAME}) + if(NOT TARGET ${MODULE_UT_NAME}_cases_obj) + add_library(${MODULE_UT_NAME}_cases_obj OBJECT) + endif() + file(GLOB OPAPI_CASES_SRC ${MODULE_DIR}/test_aclnn_*.cpp) + target_sources(${MODULE_UT_NAME}_cases_obj ${MODULE_MODE} ${OPAPI_CASES_SRC}) + endif() + endfunction() +endif() + +if(UT_TEST_ALL OR OP_KERNEL_UT) + set(fastOpTestSocVersions + "" + CACHE STRING "fastOp Test SocVersions" + ) + function(AddOpTestCase opName supportedSocVersion otherCompileOptions tilingSrcFiles) + get_filename_component(ARCH_NAME ${CMAKE_CURRENT_SOURCE_DIR} NAME) + if(${ARCH_NAME} STREQUAL "op_kernel") + get_filename_component(UT_DIR ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) + elseif(${ARCH_NAME} MATCHES "^arch") + get_filename_component(UT_TYPE_DIR ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) + get_filename_component(UT_DIR ${UT_TYPE_DIR} DIRECTORY) + endif() + get_filename_component(TESTS_DIR ${UT_DIR} DIRECTORY) + get_filename_component(OP_NAME_DIR ${TESTS_DIR} DIRECTORY) + get_filename_component(OP_NAME ${OP_NAME_DIR} NAME) + list(FIND ASCEND_OP_NAME ${OP_NAME} INDEX) + # if "--ops" is not NULL, opName not include, jump over. if "--ops" is NULL, include all. + if(NOT "${ASCEND_OP_NAME}" STREQUAL "ALL" AND INDEX EQUAL -1) + return() + endif() + + # find kernel file + if(ARGN) + list(GET ARGN 0 KernelFile) + else() + file(GLOB KernelFile "${PROJECT_SOURCE_DIR}/*/${opName}/op_kernel/${opName}.cpp") + endif() + + # find case file + get_filename_component(ARCH_NAME ${CMAKE_CURRENT_SOURCE_DIR} NAME) + # arch35 + if(${ARCH_NAME} STREQUAL "arch35") + list(FIND ARCH_DIRECTORY ${ARCH_NAME} INDEX) + if(NOT INDEX EQUAL -1) + file(GLOB OPKERNEL_CASES_SRC ${CMAKE_CURRENT_SOURCE_DIR}/test_${opName}*.cpp) + else() + return() + endif() + # op_kernel/arch20/arch22/arch32/arch38 + else() + list(FIND ARCH_DIRECTORY "arch35" INDEX) + if(INDEX EQUAL -1) + file(GLOB OPKERNEL_CASES_SRC ${CMAKE_CURRENT_SOURCE_DIR}/test_${opName}*.cpp) + else() + return() + endif() + endif() + + # standardize opType + set(opType "") + string(REPLACE "_" ";" opTypeTemp "${opName}") + foreach(word IN LISTS opTypeTemp) + string(SUBSTRING "${word}" 0 1 firstLetter) + string(SUBSTRING "${word}" 1 -1 restOfWord) + string(TOUPPER "${firstLetter}" firstLetter) + string(TOLOWER "${restOfWord}" restOfWord) + set(opType "${opType}${firstLetter}${restOfWord}") + endforeach() + + # standardize tiling files + string(REPLACE "," ";" tilingSrc "${tilingSrcFiles}") + + foreach(oriSocVersion ${supportedSocVersion}) + # standardize socVersion + string(REPLACE "ascend" "Ascend" socVersion "${oriSocVersion}") + + # add tiling tmp so: ${opName}_${socVersion}_tiling_tmp.so + add_library(${opName}_${socVersion}_tiling_tmp SHARED ${tilingSrc} $) + target_include_directories( + ${opName}_${socVersion}_tiling_tmp + PRIVATE ${OPBASE_INC_DIRS} ${ASCEND_DIR}/include/base + ${PROJECT_SOURCE_DIR}/common/include + ${ASCEND_DIR}/include/experiment + ${ASCEND_DIR}/include/experiment/metadef/common/util + ${ASCEND_DIR}/include + ) + target_compile_definitions(${opName}_${socVersion}_tiling_tmp PRIVATE LOG_CPP _GLIBCXX_USE_CXX11_ABI=0) + target_link_libraries( + ${opName}_${socVersion}_tiling_tmp + PRIVATE -Wl,--no-as-needed $<$:opsbase> -Wl,--as-needed -Wl,--whole-archive tiling_api + -Wl,--no-whole-archive gcov + $<$:$> + ) + + # gen ascendc tiling head files + set(tilingFile ${CMAKE_CURRENT_BINARY_DIR}/${opName}_tiling_data.h) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${opName}_tiling.h") + set(compileOptions -include "${CMAKE_CURRENT_SOURCE_DIR}/${opName}_tiling.h") + else() + set(compileOptions -include ${tilingFile}) + endif() + set(CUSTOM_TILING_DATA_KEYS "") + string(REGEX MATCH "-DUT_CUSTOM_TILING_DATA_KEYS=([^ ]+)" matchedPart "${otherCompileOptions}") + if(CMAKE_MATCH_1) + set(CUSTOM_TILING_DATA_KEYS ${CMAKE_MATCH_1}) + string(REGEX REPLACE "-DUT_CUSTOM_TILING_DATA_KEYS=[^ ]+" "" modifiedString ${otherCompileOptions}) + set(otherCompileOptions ${modifiedString}) + endif() + string(REPLACE " " ";" options "${otherCompileOptions}") + foreach(option IN LISTS options) + set(compileOptions ${compileOptions} ${option}) + endforeach() + message("compileOptions: ${compileOptions}") + set(gen_tiling_head_file ${OPS_TRANSFORMER_DIR}/tests/ut/framework_normal/op_kernel/scripts/gen_tiling_head_file.sh) + set(gen_tiling_so_path ${CMAKE_CURRENT_BINARY_DIR}/lib${opName}_${socVersion}_tiling_tmp.so) + set(gen_tiling_head_tag ${opName}_${socVersion}_gen_head) + set(gen_cmd "bash ${gen_tiling_head_file} ${opType} ${opName} ${gen_tiling_so_path} ${CUSTOM_TILING_DATA_KEYS}") + message("gen tiling head file to ${tilingFile}, command:") + message("${gen_cmd}") + add_custom_command( + OUTPUT ${tilingFile} + COMMAND rm -f ${tilingFile} + COMMAND bash -c ${gen_cmd} + WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} + DEPENDS ${opName}_${socVersion}_tiling_tmp + ) + add_custom_target(${gen_tiling_head_tag} ALL DEPENDS ${tilingFile}) + + # add object: ${opName}_${socVersion}_cases_obj + add_library(${opName}_${socVersion}_cases_obj OBJECT ${KernelFile} ${OPKERNEL_CASES_SRC}) + add_dependencies(${opName}_${socVersion}_cases_obj ${gen_tiling_head_tag}) + target_compile_options( + ${opName}_${socVersion}_cases_obj PRIVATE -g ${compileOptions} -DUT_SOC_VERSION="${socVersion}" + ) + target_include_directories( + ${opName}_${socVersion}_cases_obj + PRIVATE ${ASCEND_DIR}/include/base/context_builder ${PROJECT_SOURCE_DIR}/tests/ut/framework_normal/op_kernel + ${PROJECT_SOURCE_DIR}/tests/ut/framework_normal/common + ${ASCEND_DIR}/${SYSTEM_PREFIX}/asc/impl/basic_api + ${ASCEND_DIR}/${SYSTEM_PREFIX}/asc + ${ASCEND_DIR}/${SYSTEM_PREFIX}/asc/include + ${ASCEND_DIR}/${SYSTEM_PREFIX}/asc/include/basic_api + ${ASCEND_DIR}/${SYSTEM_PREFIX}/asc/include/adv_api + ${ASCEND_DIR}/${SYSTEM_PREFIX}/include/ascendc/highlevel_api + ) + target_link_libraries( + ${opName}_${socVersion}_cases_obj PRIVATE $ tikicpulib::${socVersion} + gtest + ) + + # add object: transformer_op_kernel_ut_${oriSocVersion}_cases_obj + if(NOT TARGET ${OP_KERNEL_MODULE_NAME}_${oriSocVersion}_cases_obj) + add_library( + ${OP_KERNEL_MODULE_NAME}_${oriSocVersion}_cases_obj OBJECT + $ + ) + endif() + target_link_libraries( + ${OP_KERNEL_MODULE_NAME}_${oriSocVersion}_cases_obj PRIVATE $ + $ + ) + + list(FIND fastOpTestSocVersions "${oriSocVersion}" index) + if(index EQUAL -1) + set(fastOpTestSocVersions + ${fastOpTestSocVersions} ${oriSocVersion} + CACHE STRING "fastOp Test SocVersions" FORCE + ) + endif() + endforeach() + endfunction() +endif() + +if(UT_TEST_ALL OR OP_KERNEL_AICPU_UT) + include(${PROJECT_SOURCE_DIR}/cmake/third_party/gtest.cmake) + function(AddAicpuOpTestCase opName) + get_filename_component(UT_DIR ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY) + get_filename_component(OP_NAME ${UT_DIR} NAME) + list(FIND ASCEND_OP_NAME ${OP_NAME} INDEX) + # if "--ops" is not NULL, opName not include, jump over. if "--ops" is NULL, include all. + if(NOT "${ASCEND_OP_NAME}" STREQUAL "" AND INDEX EQUAL -1) + return() + endif() + + ## find kernel file + file(GLOB KernelFile "${PROJECT_SOURCE_DIR}/*/${opName}/op_kernel_aicpu/${opName}_aicpu.cpp") + + ## add object: ${opName}_cases_obj + file(GLOB OPKERNEL_CASES_SRC ${UT_DIR}/tests/ut/op_kernel_aicpu/test_${opName}*.cpp) + + message(STATUS "aicpu kernel info: ${opName}, ${KernelFile}, ${OPKERNEL_CASES_SRC}") + + add_library(${opName}_cases_obj OBJECT + ${KernelFile} + ${OPKERNEL_CASES_SRC} + ) + target_compile_options(${opName}_cases_obj PRIVATE + -g + ) + + ## add op_kernel_aicpu test header file search path, so that header files can be referenced based on relative path + target_include_directories(${opName}_cases_obj PRIVATE + ${AICPU_INCLUDE} + ${OPBASE_INC_DIRS} + ${AICPU_INC_DIRS} + ${PROJECT_SOURCE_DIR}/tests/ut/op_kernel_aicpu + ) + target_link_libraries(${opName}_cases_obj PRIVATE + $ + -ldl + gtest + c_sec + Eigen3::EigenMath + $<$:opsbase> + ) + + ## add object: math_op_kernel_ut_cases_obj + if(NOT TARGET ${AICPU_OP_KERNEL_MODULE_NAME}_cases_obj) + add_library( + ${AICPU_OP_KERNEL_MODULE_NAME}_cases_obj OBJECT + $ + ) + else() + target_sources(${AICPU_OP_KERNEL_MODULE_NAME}_cases_obj PRIVATE $) + endif() + + target_link_libraries(${AICPU_OP_KERNEL_MODULE_NAME}_cases_obj PRIVATE + $ + $ + -ldl + $ + gtest + c_sec + Eigen3::EigenMath + $<$:opsbase> + ) + endfunction() +endif() \ No newline at end of file diff --git a/csrc/ascend/cmake/variables.cmake b/csrc/ascend/cmake/variables.cmake new file mode 100644 index 000000000..a8e7d2ba4 --- /dev/null +++ b/csrc/ascend/cmake/variables.cmake @@ -0,0 +1,319 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +set(COMMON_NAME common_${PKG_NAME}) +set(OPHOST_NAME ophost_${PKG_NAME}) +set(OPSTATIC_NAME cann_${PKG_NAME}_static) +set(OPAPI_NAME opapi_${PKG_NAME}) +set(OPGRAPH_NAME opgraph_${PKG_NAME}) +set(GRAPH_PLUGIN_NAME graph_plugin_${PKG_NAME}) +set(ONNX_PLUGIN_NAME op_${PKG_NAME}_onnx_plugin) +if(NOT CANN_3RD_LIB_PATH) + set(CANN_3RD_LIB_PATH ${PROJECT_SOURCE_DIR}/third_party) +endif() +if(NOT CANN_3RD_PKG_PATH) + set(CANN_3RD_PKG_PATH ${PROJECT_SOURCE_DIR}/third_party/pkg) +endif() +# interface, 用于收集aclnn/aclnn_inner/aclnn_exclude的def文件 +add_library(${OPHOST_NAME}_opdef_aclnn_obj INTERFACE) +add_library(${OPHOST_NAME}_opdef_aclnn_inner_obj INTERFACE) +add_library(${OPHOST_NAME}_opdef_aclnn_exclude_obj INTERFACE) +add_library(${OPHOST_NAME}_aclnn_exclude_headers INTERFACE) +# interface, 用于收集ops proto头文件 +add_library(${GRAPH_PLUGIN_NAME}_proto_headers INTERFACE) + +# global variables +set(COMPILED_OPS CACHE STRING "Compiled Ops" FORCE) +set(COMPILED_OP_DIRS CACHE STRING "Compiled Ops Dirs" FORCE) + +# src path +get_filename_component(OPS_TRANSFORMER_CMAKE_DIR "${OPS_TRANSFORMER_DIR}/cmake" REALPATH) +get_filename_component(OPS_TRANSFORMER_COMMON_INC "${OPS_TRANSFORMER_DIR}/common/include" REALPATH) +get_filename_component(OPS_TRANSFORMER_COMMON_INC_COMMON "${OPS_TRANSFORMER_COMMON_INC}/common" REALPATH) +get_filename_component(OPS_TRANSFORMER_COMMON_INC_EXTERNAL "${OPS_TRANSFORMER_COMMON_INC}/external" REALPATH) +get_filename_component(OPS_TRANSFORMER_COMMON_INC_HEADERS "${OPS_TRANSFORMER_COMMON_INC_EXTERNAL}/aclnn_kernels" REALPATH) +get_filename_component(OPS_KERNEL_BINARY_SCRIPT "${OPS_TRANSFORMER_DIR}/scripts/kernel/binary_script" REALPATH) +get_filename_component(OPS_KERNEL_BINARY_CONFIG "${OPS_TRANSFORMER_DIR}/scripts/kernel/binary_config" REALPATH) + +# python +if(NOT DEFINED ASCEND_PYTHON_EXECUTABLE) + set(ASCEND_PYTHON_EXECUTABLE python3 CACHE STRING "") +endif() + +set(BIN_STATIC_INSTALL_DIR packages/vendors/${VENDOR_NAME}_transformer/op_impl/ai_core/tbe/static) +if (ENABLE_BUILT_IN) + set(ACLNN_INC_INSTALL_DIR ops_transformer/built-in/op_impl/ai_core/tbe/op_api/include/aclnnop) + set(ACLNN_INC_LEVEL2_INSTALL_DIR ops_transformer/built-in/op_impl/ai_core/tbe/op_api/include/aclnnop/level2) + set(ACLNN_LIB_INSTALL_DIR ops_transformer/built-in/op_impl/ai_core/tbe/op_api/lib/linux/${CMAKE_SYSTEM_PROCESSOR}) + set(OPS_INFO_INSTALL_DIR ops_transformer/built-in/op_impl/ai_core/tbe/config) + set(IMPL_INSTALL_DIR ops_transformer/built-in/op_impl/ai_core/tbe/impl/ops_transformer/ascendc) + set(IMPL_DYNAMIC_INSTALL_DIR ops_transformer/built-in/op_impl/ai_core/tbe/impl/ops_transformer/dynamic) + set(BIN_KERNEL_INSTALL_DIR ops_transformer/built-in/op_impl/ai_core/tbe/kernel) + set(BIN_KERNEL_CONFIG_INSTALL_DIR ops_transformer/built-in/op_impl/ai_core/tbe/kernel/config) + set(BIN_STATIC_INSTALL_DIR ops_transformer/built-in/op_impl/ai_core/tbe/static) + set(OPHOST_INC_INSTALL_PATH ops_transformer/built-in/op_impl/ai_core/tbe/op_host/include) + set(OPHOST_LIB_INSTALL_PATH ops_transformer/built-in/op_impl/ai_core/tbe/op_host/lib/linux/${CMAKE_SYSTEM_PROCESSOR}) + set(OPTILING_LIB_INSTALL_DIR ${OPHOST_LIB_INSTALL_PATH}) + set(OPGRAPH_INC_INSTALL_DIR ops_transformer/built-in/op_graph/inc) + set(OPGRAPH_LIB_INSTALL_DIR ops_transformer/built-in/op_graph/lib/linux/${CMAKE_SYSTEM_PROCESSOR}) + set(ONNX_PLUGIN_LIB_INSTALL_DIR ops_transformer/built-in/framework/onnx) + set(COMMON_INC_INSTALL_DIR ops_transformer/include) + set(COMMON_LIB_INSTALL_DIR ops_transformer/lib) + set(VERSION_INFO_INSTALL_DIR ops_transformer) + set(IMPL_INSTALL_DIR ops_transformer/built-in/op_impl/ai_core/tbe/impl/ops_transformer) +endif() + +if (ENABLE_TEST) + set(UTEST_FRAMEWORK_OLD FALSE CACHE BOOL "UTEST_FRAMEWORK_OLD") + set(UTEST_FRAMEWORK_NEW FALSE CACHE BOOL "UTEST_FRAMEWORK_NEW") +endif() + +# util path +set(ASCEND_TENSOR_COMPILER_PATH ${ASCEND_DIR}/compiler) +set(ASCEND_CCEC_COMPILER_PATH ${ASCEND_TENSOR_COMPILER_PATH}/ccec_compiler/bin) +set(OP_BUILD_TOOL ${ASCEND_DIR}/tools/opbuild/op_build) +set(UT_PATH ${CMAKE_CURRENT_SOURCE_DIR}/tests/ut/framework_normal) + +# output path +set(ASCEND_AUTOGEN_PATH ${CMAKE_BINARY_DIR}/autogen) +set(ASCEND_KERNEL_SRC_DST ${CMAKE_BINARY_DIR}/tbe/ascendc) +set(ASCEND_KERNEL_CONF_DST ${CMAKE_BINARY_DIR}/tbe/config) +set(ASCEND_GRAPH_CONF_DST ${CMAKE_BINARY_DIR}/tbe/graph) +file(MAKE_DIRECTORY ${ASCEND_AUTOGEN_PATH}) +file(MAKE_DIRECTORY ${ASCEND_KERNEL_SRC_DST}) +file(MAKE_DIRECTORY ${ASCEND_KERNEL_CONF_DST}) +file(MAKE_DIRECTORY ${ASCEND_GRAPH_CONF_DST}) +set(CUSTOM_COMPILE_OPTIONS "custom_compile_options.ini") +set(CUSTOM_OPC_OPTIONS "custom_opc_options.ini") +execute_process( + COMMAND rm -rf ${ASCEND_AUTOGEN_PATH}/${CUSTOM_COMPILE_OPTIONS} + COMMAND rm -rf ${ASCEND_AUTOGEN_PATH}/${CUSTOM_OPC_OPTIONS} + COMMAND touch ${ASCEND_AUTOGEN_PATH}/${CUSTOM_COMPILE_OPTIONS} + COMMAND touch ${ASCEND_AUTOGEN_PATH}/${CUSTOM_OPC_OPTIONS} +) + +# pack path +set(CMAKE_INSTALL_PREFIX ${CMAKE_SOURCE_DIR}/build_out) + +set(OPAPI_INCLUDE + ${C_SEC_INCLUDE} + ${PLATFORM_INC_DIRS} + ${METADEF_INCLUDE_DIRS} + ${NNOPBASE_INCLUDE_DIRS} + ${NPURUNTIME_INCLUDE_DIRS} + ${AICPU_INC_DIRS} + ${OPS_TRANSFORMER_DIR}/ + ${OPS_TRANSFORMER_DIR}/common/include + ${OPS_TRANSFORMER_DIR}/common/include/external + ${OPS_TRANSFORMER_DIR}/common/include/common + ${OPS_TRANSFORMER_DIR}/common/include/static + ${OPS_TRANSFORMER_DIR}/common/stub/op_api + $<$>:${TOP_DIR}/output/${PRODUCT}/aclnnop_resource> + + ${OPS_TRANSFORMER_DIR}/mc2/common/inc + ${OPS_TRANSFORMER_DIR}/mc2/common/new_mc2_mm + ${OPS_TRANSFORMER_DIR}/mc2/3rd + ${OPS_TRANSFORMER_DIR}/mc2 +) + +if (NOT BUILD_OPEN_PROJECT) + list(APPEND OPAPI_INCLUDE + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_proto/runtime + ${TOP_DIR}/ace/comop/inc/external + ${TOP_DIR}/ops/ops-nn/matmul/common/op_host/op_api + ${TOP_DIR}/asl/ops/cann/ops/utils/inc/log/inner + ${TOP_DIR}/asl/ops/cann/ops/utils/inc/error + ${TOP_DIR}/ace/comop/inc/external + ${TOP_DIR}/ace/npuruntime/inc/external + ${TOP_DIR}/ace/npuruntime/inc/nnopbase + ${TOP_DIR}/asl/ops/cann/ops/mc2/communication_and_computation + ${TOP_DIR}/ace/npuruntime/acl/inc/external/acl/error_codes + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_tiling/runtime + ${TOP_DIR}/asl/ops/cann/ops/built-in + ${TOP_DIR}/ops-base/pkg_inc/op_common/op_host + ${TOP_DIR}/ops-base/pkg_inc + ${TOP_DIR}/ops-base/include + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_fallback + ) +else() + list(APPEND OPAPI_INCLUDE ${OPBASE_INC_DIRS}) +endif() + +set(OP_TILING_INCLUDE + ${C_SEC_INCLUDE} + ${PLATFORM_INC_DIRS} + ${METADEF_INCLUDE_DIRS} + ${TILINGAPI_INC_DIRS} + ${NPURUNTIME_INCLUDE_DIRS} + ${OPBASE_INC_DIRS} + ${NNOPBASE_INCLUDE_DIRS} + ${AICPU_INC_DIRS} + ${OPS_TRANSFORMER_DIR} + ${JSON_INCLUDE_DIR} + ${OPS_TRANSFORMER_DIR}/common/include + ${OPS_TRANSFORMER_DIR}/common/include/ + ${OPS_TRANSFORMER_DIR}/common/stub/op_tiling + + ${OPS_TRANSFORMER_DIR}/mc2/common + ${OPS_TRANSFORMER_DIR}/mc2/common/inc + ${OPS_TRANSFORMER_DIR}/mc2/common/new_mc2_mm + ${OPS_TRANSFORMER_DIR}/mc2/3rd + ${OPS_TRANSFORMER_DIR}/mc2 + ${NNOPBASE_INCLUDE_DIRS} + ${AICPU_INC_DIRS} +) + +if (NOT BUILD_OPEN_PROJECT) + list(APPEND OP_TILING_INCLUDE + ${TOP_DIR}/abl/msprof/inc # codespell:ignore abl + ${METADEF_INC_DIR}/../common/util + ${TOP_DIR}/asl/ops/cann/ops/utils/inc + ${TOP_DIR}/ace/comop/inc + ${TOP_DIR}/ace/comop/hccl/open_source/inc + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_tiling/cube + ${TOP_DIR}/ace/npuruntime/inc/external + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_tiling + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_api/inc + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_proto/runtime + ${TOP_DIR}/asl/ops/cann/ops/common/inc + ${TOP_DIR}/asl/ops/cann/ops/ops-nn/inner + ${TOP_DIR}/asl/ops/cann/ops/matmul + ${TOP_DIR}/ace/npuruntime/acl/inc/external/acl/error_codes + ${TOP_DIR}/asl/ops/cann/ops/mc2/common/inc + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_tiling/runtime + ${TOP_DIR}/asl/ops/cann/ops/built-in + + ${TOP_DIR}/asl/ops/cann/ops/mc2/communication_and_computation + ${TOP_DIR}/ops-base/pkg_inc/op_common/op_host + ${TOP_DIR}/ops-base/pkg_inc + ${TOP_DIR}/ops-base/include + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_fallback + + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_proto/runtime + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_tiling/device/error + ${METADEF_DIR} + ${RUNTIME_INC_DIR}/runtime/platform/inc + ${METADEF_DIR}/inc/external/ge + ) +endif() + +set(OP_PROTO_INCLUDE + ${C_SEC_INCLUDE} + ${PLATFORM_INC_DIRS} + ${METADEF_INCLUDE_DIRS} + ${OPBASE_INC_DIRS} + ${NPURUNTIME_INCLUDE_DIRS} + ${OPS_TRANSFORMER_DIR}/common/include/ + ${OPS_TRANSFORMER_DIR} + ${OPS_TRANSFORMER_DIR}/mc2/common + + ${OPS_TRANSFORMER_DIR}/common/include + + ${OPS_TRANSFORMER_DIR}/mc2/common/inc + ${OPS_TRANSFORMER_DIR}/mc2/common/new_mc2_mm + ${OPS_TRANSFORMER_DIR}/mc2/3rd + ${OPS_TRANSFORMER_DIR}/mc2 +) + +if (NOT BUILD_OPEN_PROJECT) + list(APPEND OP_PROTO_INCLUDE + ${TOP_DIR}/abl/msprof/inc # codespell:ignore abl + ${METADEF_INC_DIR}/../common/util + ${TOP_DIR}/ace/comop/inc + ${TOP_DIR}/ace/comop/hccl/open_source/inc + ${TOP_DIR}/asl/ops/cann/ops/utils/inc + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_tiling/cube + ${TOP_DIR}/asl/ops/cann/ops/utils/inc/log/inner + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_tiling + ${TOP_DIR}/ace/npuruntime/inc/external + ${TOP_DIR}/asl/ops/cann/ops/common/inc + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_proto/runtime + ${TOP_DIR}/tmp/host-prefix/src/host-build/atc/opcompiler/ascendc_compiler/api/kernel_tiling + ${TOP_DIR}/asl/ops/cann/ops/ops-nn/inner + ${TOP_DIR}/asl/ops/cann/ops/matmul + ${TOP_DIR}/ace/npuruntime/acl/inc/external/acl/error_codes + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_tiling/runtime + ${TOP_DIR}/asl/ops/cann/ops/built-in + ${TOP_DIR}/ops-base/pkg_inc/op_common/op_host + ${TOP_DIR}/ops-base/pkg_inc + ${TOP_DIR}/ops-base/include + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_fallback + ) +endif() + +if (NOT BUILD_OPEN_PROJECT) + list(APPEND AICPU_INCLUDE + ${TOP_DIR}/asl/ops/cann/ops/built-in/aicpu/impl + ${TOP_DIR}/asl/ops/cann/ops/built-in/aicpu/impl/utils + ${TOP_DIR}/asl/ops/cann/ops/built-in/aicpu/impl/kernels/host/runtime/utils + ${TOP_DIR}/asl/ops/cann/ops/built-in/aicpu/impl/kernels/normalized/random + ${TOP_DIR}/asl/ops/cann/ops/built-in/aicpu/context/inc + ${TOP_DIR}/asl/ops/cann/ops/built-in/aicpu/context/common/*.h + ${TOP_DIR}/ace/comop/inc + ${TOP_DIR}/ops-base/include + ${TOP_DIR}/ops-base/pkg_inc + ${TOP_DIR}/asl/ops/cann/ops/built-in/op_fallback + ) +endif() + +set(AICPU_INCLUDE + ${OPBASE_INC_DIRS} + ${AICPU_INC_DIRS} + ${C_SEC_INCLUDE} + ${NNOPBASE_INCLUDE_DIRS} + ${HCCL_EXTERNAL_INCLUDE} + ${OPS_TRANSFORMER_DIR}/common/inc/common + ${METADEF_INCLUDE_DIRS} +) + +set(ONNX_PLUGIN_COMMON_INCLUDE + ${OPS_TRANSFORMER_DIR} + ${OPS_TRANSFORMER_DIR}/common/include/framework + ${OPS_TRANSFORMER_DIR}/common/include/op_graph + ${OPS_TRANSFORMER_DIR}/common/stub/inc/framework +) + +set(AICPU_DEFINITIONS + -O2 + -std=c++17 + -fstack-protector-all + -fvisibility-inlines-hidden + -fvisibility=hidden + -frename-registers + -fpeel-loops + -DEIGEN_NO_DEBUG + -DEIGEN_MPL2_ONLY + -DNDEBUG + -DEIGEN_HAS_CXX11_MATH + -DEIGEN_OS_GNULINUX + -DEigen=ascend_Eigen + -fno-common + -fPIC +) + +if(EXISTS ${TOP_DIR}/build/product/onetrack/sys_version/sys_version.conf) + execute_process(COMMAND grep -Po "^\\d+\\.\\d+" ${TOP_DIR}/build/product/onetrack/sys_version/sys_version.conf + OUTPUT_VARIABLE SYS_VERSION + OUTPUT_STRIP_TRAILING_WHITESPACE + ) +else() + execute_process(COMMAND grep -Po "(?<=Version=)[0-9]+\.[0-9]+" ${OPS_TRANSFORMER_DIR}/version.info + OUTPUT_VARIABLE SYS_VERSION + OUTPUT_STRIP_TRAILING_WHITESPACE + ) +endif() + +include_directories(${PROJECT_SOURCE_DIR}/utils/inc) + +set(OPS_CATEGORY_LIST + "attention" + "moe" + "mc2" +) diff --git a/csrc/ascend/common/CMakeLists.txt b/csrc/ascend/common/CMakeLists.txt new file mode 100644 index 000000000..efbaa8fb9 --- /dev/null +++ b/csrc/ascend/common/CMakeLists.txt @@ -0,0 +1,131 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) + foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() + endforeach() +endif() + +if (BUILD_OPEN_PROJECT) + add_library(ops_transformer_utils_tiling_headers INTERFACE) + + target_include_directories(ops_transformer_utils_tiling_headers INTERFACE + $ + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $ + ) + + target_compile_definitions(ops_transformer_utils_tiling_headers INTERFACE + OPS_UTILS_LOG_SUB_MOD_NAME="OP_TILING" + OPS_UTILS_LOG_PACKAGE_TYPE=$,"[Custom]",""> + ) + + add_library(ops_transformer_utils_proto_headers INTERFACE) + + target_include_directories(ops_transformer_utils_proto_headers INTERFACE + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $ + ) + + target_compile_definitions(ops_transformer_utils_proto_headers INTERFACE + OPS_UTILS_LOG_SUB_MOD_NAME="OP_PROTO" + OPS_UTILS_LOG_PACKAGE_TYPE=$,"[Custom]",""> + ) + + add_library(${COMMON_NAME}_obj OBJECT) + + file(GLOB CPP_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/src/tiling_base/*.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/src/*.cpp") + + target_sources(${COMMON_NAME}_obj PRIVATE ${CPP_SOURCES}) + + target_include_directories(${COMMON_NAME}_obj + PRIVATE + ${C_SEC_INCLUDE} + ${TILINGAPI_INC_DIRS} + ${METADEF_INCLUDE_DIRS} + ${OPS_TRANSFORMER_DIR}/common/include + ${ASCEND_CANN_PACKAGE_PATH}/include/experiment/runtime + ${ASCEND_CANN_PACKAGE_PATH}/include/experiment/msprof + ${ASCEND_CANN_PACKAGE_PATH}/include/op_common + ${ASCEND_CANN_PACKAGE_PATH}/pkg_inc + ${ASCEND_CANN_PACKAGE_PATH}/pkg_inc/op_common + ${ASCEND_CANN_PACKAGE_PATH}/${SYSTEM_PREFIX}/pkg_inc/runtime + ${ASCEND_CANN_PACKAGE_PATH}/${SYSTEM_PREFIX}/pkg_inc + ${ASCEND_CANN_PACKAGE_PATH}/${SYSTEM_PREFIX}/pkg_inc/op_common + ${ASCEND_CANN_PACKAGE_PATH}/pkg_inc/profiling + ) + + target_link_libraries(${COMMON_NAME}_obj + PRIVATE + $ + tiling_api + dlog + ) +else() + add_library(${COMMON_NAME}_obj OBJECT) + + file(GLOB CPP_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/src/tiling_base/*.cpp") + + target_sources(${COMMON_NAME}_obj PRIVATE ${CPP_SOURCES}) + + target_include_directories(${COMMON_NAME}_obj + PRIVATE + ${C_SEC_INCLUDE} + ${TILINGAPI_INC_DIRS} + ${METADEF_INCLUDE_DIRS} + ${OPS_TRANSFORMER_DIR}/common/include + ${OPS_TRANSFORMER_DIR}/common/inc + ) + + target_link_libraries(${COMMON_NAME}_obj + PRIVATE + $ + tiling_api + ) +endif() + +if (NOT BUILD_OPS_RTY_KERNEL) + if (BUILD_OPEN_PROJECT) + ExternalProject_Add(tiling_sink_task + SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/src/tiling_sink + CONFIGURE_COMMAND ${CMAKE_COMMAND} + -G ${CMAKE_GENERATOR} + -DBUILD_OPEN_PROJECT=${BUILD_OPEN_PROJECT} + -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} + -DASCEND_CANN_PACKAGE_PATH=${ASCEND_CANN_PACKAGE_PATH} + -DCMAKE_C_COMPILER=${ASCEND_CANN_PACKAGE_PATH}/toolkit/toolchain/hcc/bin/aarch64-target-linux-gnu-gcc + -DCMAKE_CXX_COMPILER=${ASCEND_CANN_PACKAGE_PATH}/toolkit/toolchain/hcc/bin/aarch64-target-linux-gnu-g++ + -DCMAKE_C_COMPILER_LAUNCHER=${CMAKE_C_COMPILER_LAUNCHER} + -DCMAKE_CXX_COMPILER_LAUNCHER=${CMAKE_CXX_COMPILER_LAUNCHER} + -DSYS_VERSION=${SYS_VERSION} + -DOPS_ADV_CMAKE_DIR=${OPS_ADV_CMAKE_DIR} + -DPRODUCT_SIDE=device + -DENABLE_BUILT_IN=${ENABLE_BUILT_IN} + + INSTALL_COMMAND "" + BUILD_ALWAYS TRUE + BUILD_BYPRODUCTS + ) + add_subdirectory(src/framework) + else() + add_subdirectory(src/tiling_sink) + endif() +endif() diff --git a/csrc/ascend/common/aicpu/cpu_context_util.h b/csrc/ascend/common/aicpu/cpu_context_util.h new file mode 100644 index 000000000..e269c1413 --- /dev/null +++ b/csrc/ascend/common/aicpu/cpu_context_util.h @@ -0,0 +1,91 @@ +/** + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file cou_context_util.h + * \brief + */ + +#ifndef CPU_CONTEXT_UTIL_H +#define CPU_CONTEXT_UTIL_H + +#include "cpu_context.h" +#include "cpu_kernel.h" +#include "cpu_tensor.h" +#include "log.h" +#include +#include + +#define KERNEL_STATUS_OK 0 +#define KERNEL_STATUS_PARAM_INVALID 1 + +namespace aicpu { +template +inline typename std::enable_if, bool>::type +GetAttrValue(CpuKernelContext &ctx, const std::string &name, T &value) { + auto attr = ctx.GetAttr(name); + if (!attr) { + KERNEL_LOG_ERROR("attr is null: %s", name.c_str()); + return false; + } + value = static_cast(attr->GetInt()); + return true; +} + +inline bool GetAttrValue(CpuKernelContext &ctx, const std::string &name, + std::string &value) { + auto attr = ctx.GetAttr(name); + if (!attr) { + KERNEL_LOG_ERROR("attr is null: %s", name.c_str()); + return false; + } + value = attr->GetString(); + return true; +} + +inline bool GetAttrValue(CpuKernelContext &ctx, const std::string &name, + bool &value) { + auto attr = ctx.GetAttr(name); + if (!attr) { + KERNEL_LOG_ERROR("attr is null: %s", name.c_str()); + return false; + } + value = attr->GetBool(); + return true; +} + +template +inline typename std::enable_if, void>::type +GetAttrValueOpt(CpuKernelContext &ctx, const std::string &name, T &value) { + auto attr = ctx.GetAttr(name); + if (attr != nullptr) { + value = static_cast(attr->GetInt()); + } +} + +inline void GetAttrValueOpt(CpuKernelContext &ctx, const std::string &name, + std::string &value) { + auto attr = ctx.GetAttr(name); + if (attr != nullptr) { + value = attr->GetString(); + } +} + +inline void GetAttrValueOpt(CpuKernelContext &ctx, const std::string &name, + bool &value) { + auto attr = ctx.GetAttr(name); + if (attr != nullptr) { + value = attr->GetBool(); + } +} + +} // namespace aicpu + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/include/cann_compat.h b/csrc/ascend/common/include/cann_compat.h new file mode 100644 index 000000000..c04c64d6a --- /dev/null +++ b/csrc/ascend/common/include/cann_compat.h @@ -0,0 +1,30 @@ +#pragma once + +#include + +// B080 op_common/log/log.h stopped exposing the unqualified OP module id used +// by inherited ops-transformer tiling/error headers. Include the CANN log type +// header early so OP still comes from the active CANN version. +#if defined(__has_include) +#if __has_include("base/log_types.h") +#include "base/log_types.h" +#elif __has_include("toolchain/log_types.h") +#include "toolchain/log_types.h" +#endif +#endif + +#if !defined(LOG_TYPES_H_) && !defined(OP) +#define OP 63 +#endif + +#if defined(LOG_CPP) && !defined(DLOG_PUB_H_) +#ifdef __cplusplus +extern "C" { +#endif +int32_t CheckLogLevel(int32_t moduleId, int32_t logLevel); +void DlogRecord(int32_t moduleId, int32_t level, const char *fmt, ...); +#ifdef __cplusplus +} +#endif +#define DLOG_PUB_H_ +#endif diff --git a/csrc/ascend/common/include/common/op_api_def.h b/csrc/ascend/common/include/common/op_api_def.h new file mode 100644 index 000000000..ca83c7bf6 --- /dev/null +++ b/csrc/ascend/common/include/common/op_api_def.h @@ -0,0 +1,29 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file op_api_def.h + * \brief + */ + +#ifndef Transformer_COMMON_OP_API_DEF_H +#define Transformer_COMMON_OP_API_DEF_H + +namespace op { + constexpr size_t MAX_SUPPORT_DIMS_NUMS = 8; + constexpr size_t BN_MIN_SUPPORT_DIMS_NUMS = 2; + constexpr int8_t FP16FP32_KEEP_DTYPE = -1; + constexpr int8_t KEEP_DTYPE = 0; + constexpr int8_t ALLOW_FP32_DOWN_PRECISION = 1; + constexpr int8_t USE_FP16 = 2; + constexpr int8_t USE_HF32 = 3; + constexpr size_t MAX_MASK_LEN64 = 64; +} // namespace op +#endif // Transformer_COMMON_OP_API_DEF_H \ No newline at end of file diff --git a/csrc/ascend/common/include/common/tensor_util.cpp b/csrc/ascend/common/include/common/tensor_util.cpp new file mode 100644 index 000000000..3527339ab --- /dev/null +++ b/csrc/ascend/common/include/common/tensor_util.cpp @@ -0,0 +1,230 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "tensor_util.h" +#include "aclnn_kernels/transdata.h" +#include "aclnn_kernels/transpose.h" +#include "aclnn_kernels/reshape.h" +#include "aclnn_kernels/cast.h" +#include "aclnn_kernels/contiguous.h" +#include "level0/unsqueeze.h" +#include "level0/squeeze.h" +#include "level0/fill.h" +#include "aclnn/aclnn_base.h" + +namespace op { +const aclIntArray* getAllDims(const aclTensor* self, aclOpExecutor* executor) { + auto input_shape = self->GetViewShape(); + const size_t input_dim_num = input_shape.GetDimNum(); + std::vector dims(input_dim_num); + for (size_t idx = 0; idx < input_dim_num; idx++) { + dims[idx] = idx; + } + return executor->AllocIntArray(dims.data(), input_dim_num); +} + +constexpr size_t MAX_DIM_CNT = 5; +const aclTensor* ResizeFrom1D(const aclTensor* cdim, const aclTensor* input, bool isSupportNcdhw, aclOpExecutor* executor) { + auto cdimContiguous = l0op::Contiguous(cdim, executor); + if (cdimContiguous == nullptr) { + return cdimContiguous; + } + + auto cdimCast = l0op::Cast(cdimContiguous, DataType::DT_FLOAT, executor); + if (cdimCast == nullptr) { + return cdimCast; + } + + size_t inputDim = input->GetViewShape().GetDimNum(); + + const int64_t appendDim[] = {0, 2, 3}; + aclIntArray* newShape = executor->AllocIntArray(appendDim, sizeof(appendDim) / sizeof(int64_t)); + if (inputDim == MAX_DIM_CNT) { + const int64_t value[] = {0, 2, 3, 4}; + newShape = executor->AllocIntArray(value, sizeof(value) / sizeof(int64_t)); + } + auto cdimUnsqueeze = l0op::UnsqueezeNd(cdimCast, newShape, executor); + if (cdimUnsqueeze == nullptr) { + return cdimUnsqueeze; + } + + op::Format format = inputDim == MAX_DIM_CNT ? Format::FORMAT_NCDHW : Format::FORMAT_NCHW; + auto cdimFormat = l0op::ReFormat(cdimUnsqueeze, format); + if (cdimFormat == nullptr) { + return cdimFormat; + } + + if ((inputDim == MAX_DIM_CNT) && !isSupportNcdhw) { + return l0op::TransDataSpecial(cdimFormat, Format::FORMAT_NDC1HWC0, 0, executor); + } + + return cdimFormat; +} + +const aclTensor* ResizeTo1D(const aclTensor* result, const aclTensor* output, bool isSupportNcdhw, aclOpExecutor* executor) { + auto resultTransdata = result; + size_t resultDim = result->GetViewShape().GetDimNum(); + if (resultDim >= MAX_DIM_CNT && !isSupportNcdhw) { + resultTransdata = l0op::TransDataSpecial(result, Format::FORMAT_NCDHW, 0, executor); + if (resultTransdata == nullptr) { + return resultTransdata; + } + } + + const int64_t appendDim[] = {0, 2, 3}; + aclIntArray* newShape = executor->AllocIntArray(appendDim, sizeof(appendDim) / sizeof(int64_t)); + if (resultTransdata->GetViewShape().GetDimNum() == MAX_DIM_CNT) { + const int64_t value[] = {0, 2, 3, 4}; + newShape = executor->AllocIntArray(value, sizeof(value) / sizeof(int64_t)); + } + auto resultNchw = l0op::SqueezeNd(resultTransdata, newShape, executor); + if (resultNchw == nullptr) { + return resultNchw; + } + + auto resultNd = l0op::ReFormat(resultNchw, Format::FORMAT_ND); + if (resultNd == nullptr) { + return resultNd; + } + + auto resultCast = l0op::Cast(resultNd, output->GetDataType(), executor); + if (resultCast == nullptr) { + return resultCast; + } + + return l0op::ViewCopy(resultCast, output, executor); +} + +const aclTensor* ResizeFromND(const aclTensor* input, aclOpExecutor* executor) { + const int nchw_dims = 4; + auto inputShape = input->GetViewShape(); + int64_t nchwShape[nchw_dims]; + for (size_t i = 0; i < nchw_dims; i++) { + nchwShape[i] = i < inputShape.GetDimNum() ? inputShape[i] : 1; + } + aclIntArray* nchwArray = executor->AllocIntArray(nchwShape, nchw_dims); + + auto inputReshape = l0op::Reshape(input, nchwArray, executor); + if (inputReshape == nullptr) { + return inputReshape; + } + + return l0op::ReFormat(inputReshape, Format::FORMAT_NCHW); +} + +const aclTensor* ResizeToND(const aclTensor* output, const aclTensor* input, aclOpExecutor* executor) { + auto inputShape = input->GetViewShape(); + size_t dimNum = inputShape.GetDimNum(); + + int64_t ndShape[dimNum]; + for (size_t i = 0; i < inputShape.GetDimNum(); i++) { + ndShape[i] = inputShape[i]; + } + aclIntArray* ndArray = executor->AllocIntArray(ndShape, dimNum); + + auto outputReshape = l0op::Reshape(output, ndArray, executor); + if (outputReshape == nullptr) { + return outputReshape; + } + + return l0op::ReFormat(outputReshape, input->GetViewFormat()); +} + +const aclTensor* ResizeFrom5D(const aclTensor* input, aclOpExecutor* executor) { + auto inputShape = input->GetViewShape(); + // NCDHW -> NDCHW + const int64_t value[] = {0, 2, 1, 3, 4}; + aclIntArray* ndchwShape = executor->AllocIntArray(value, MAX_DIM_CNT); + auto inputTranspose = l0op::Transpose(input, ndchwShape, executor); + if (inputTranspose == nullptr) { + return inputTranspose; + } + + // NDCHW -> NCHW + const int64_t nchwShape[] = {inputShape[0] * inputShape[2], inputShape[1], inputShape[3], inputShape[4]}; + aclIntArray* nchwArray = executor->AllocIntArray(nchwShape, sizeof(nchwShape) / sizeof(int64_t)); + auto inputReshape = l0op::Reshape(inputTranspose, nchwArray, executor); + if (inputReshape == nullptr) { + return inputReshape; + } + + return l0op::ReFormat(inputReshape, Format::FORMAT_NCHW); +} + +const aclTensor* ResizeTo5D(const aclTensor* output, const aclTensor* input, aclOpExecutor* executor) { + auto inputShape = input->GetViewShape(); + // nchw -> ndchw + const int64_t ndchwShape[] = {inputShape[0], inputShape[2], inputShape[1], inputShape[3], inputShape[4]}; + aclIntArray* ndchwArray = executor->AllocIntArray(ndchwShape, MAX_DIM_CNT); + auto outputReshape = l0op::Reshape(output, ndchwArray, executor); + if (outputReshape == nullptr) { + return outputReshape; + } + + auto outputFormat = l0op::ReFormat(outputReshape, Format::FORMAT_NCDHW); + if (outputFormat == nullptr) { + return outputFormat; + } + // ndchw -> ncdhw + const int64_t ncdhwShape[] = {0, 2, 1, 3, 4}; + aclIntArray* ncdhwArray = executor->AllocIntArray(ncdhwShape, MAX_DIM_CNT); + return l0op::Transpose(outputFormat, ncdhwArray, executor); +} + +aclTensor* FillScalar(int64_t dim, int value, aclOpExecutor* executor) { + const aclScalar* dimScalar = executor->AllocScalar(dim); + const aclTensor* dimTensor = executor->ConvertToTensor(dimScalar, op::DataType::DT_INT32); + aclIntArray* outShape = executor->AllocIntArray(&dim, 1); + + const aclScalar* valueScalar = executor->AllocScalar(value); + const aclTensor* valueTensor = executor->ConvertToTensor(valueScalar, op::DataType::DT_FLOAT); + + auto fillTensor = l0op::Fill(dimTensor, valueTensor, outShape, executor); + if (fillTensor == nullptr) { + return nullptr; + } + + return const_cast(fillTensor); +} + +aclTensor* FillVector(const op::Shape dstShape, const aclTensor* src, float value, aclOpExecutor* executor) { + op::FVector fillDims = op::ToShapeVector(dstShape); + auto shapes = executor->AllocIntArray(fillDims.data(), src->GetViewShape().GetDimNum()); + const aclTensor* dimTensor = executor->ConvertToTensor(shapes, op::DataType::DT_INT32); + const aclScalar* valueScalar = executor->AllocScalar(value); + const aclTensor* valueTensor = executor->ConvertToTensor(valueScalar, src->GetDataType()); + auto fillTensor = l0op::Fill(dimTensor, valueTensor, shapes, executor); + if (fillTensor == nullptr) { + return nullptr; + } + fillTensor = l0op::ReFormat(fillTensor, op::Format::FORMAT_ND); + return const_cast(fillTensor); +} + +aclnnStatus ProcessEmptyTensorWithValue(aclTensor* src, float initValue, aclOpExecutor* executor) { + auto srcShape = src->GetViewShape(); + auto dst = FillVector(srcShape, src, initValue, executor); + auto dstCopyResult = l0op::ViewCopy(dst, src, executor); + CHECK_RET(dstCopyResult != nullptr, ACLNN_ERR_INNER_NULLPTR); + return ACLNN_SUCCESS; +} + +op::DataType CombineCategories(op::DataType higher, op::DataType lower) { + if (IsFloatingType(higher)) { + return higher; + } + + if (IsFloatingType(lower) || higher == op::DataType::DT_BOOL) { + return op::PromoteType(higher, lower); + } + + return (higher != op::DataType::DT_UNDEFINED) ? higher : lower; +} +} // namespace op \ No newline at end of file diff --git a/csrc/ascend/common/include/common/tensor_util.h b/csrc/ascend/common/include/common/tensor_util.h new file mode 100644 index 000000000..8120e73a3 --- /dev/null +++ b/csrc/ascend/common/include/common/tensor_util.h @@ -0,0 +1,47 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#include "aclnn/aclnn_base.h" +#include "opdev/common_types.h" + +namespace op { +const aclIntArray* getAllDims(const aclTensor* self, aclOpExecutor* executor); + +const aclTensor* ResizeFrom1D(const aclTensor* cdim, const aclTensor* input, bool isSupportNcdhw, + aclOpExecutor* executor); + +const aclTensor* ResizeTo1D(const aclTensor* result, const aclTensor* output, bool isSupportNcdhw, + aclOpExecutor* executor); + +const aclTensor* ResizeFromND(const aclTensor* input, aclOpExecutor* executor); + +const aclTensor* ResizeToND(const aclTensor* output, const aclTensor* input, aclOpExecutor* executor); + +const aclTensor* ResizeFrom5D(const aclTensor* input, aclOpExecutor* executor); + +const aclTensor* ResizeTo5D(const aclTensor* output, const aclTensor* input, aclOpExecutor* executor); + +aclTensor* FillScalar(int64_t dim, int value, aclOpExecutor* executor); + +aclnnStatus ProcessEmptyTensorWithValue(aclTensor* src, float initValue, aclOpExecutor* executor); + +op::DataType CombineCategories(op::DataType higher, op::DataType lower); +} // namespace op + +#ifdef __cplusplus +extern "C" { +#endif + +aclnnStatus BatchNorm(const aclTensor* input, const aclTensor* weight, const aclTensor* bias, aclTensor* runningMean, + aclTensor* runningVar, bool training, float momentum, float eps, aclTensor** output, + aclTensor* saveMean, aclTensor* saveInvstd, aclOpExecutor* executor); + +#ifdef __cplusplus +} +#endif \ No newline at end of file diff --git a/csrc/ascend/common/include/err/ops_err.h b/csrc/ascend/common/include/err/ops_err.h new file mode 100644 index 000000000..e83b89843 --- /dev/null +++ b/csrc/ascend/common/include/err/ops_err.h @@ -0,0 +1,32 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file ops_err.h + * \brief + */ + +#ifndef Transformer_COMMON_OPS_ERR_H +#define Transformer_COMMON_OPS_ERR_H + +#include "log/log.h" + +#define OPS_INNER_ERR_STUB(ERR_CODE_STR, OPS_DESC, FMT, ...) \ + do { \ + OpLogSub(OP, DLOG_ERROR, OPS_DESC, FMT, ##__VA_ARGS__); \ + REPORT_INNER_ERR_MSG(ERR_CODE_STR, FMT, ##__VA_ARGS__); \ + } while (0) + + +/* 基础报错 */ +#define OPS_REPORT_VECTOR_INNER_ERR(OPS_DESC, ...) OPS_INNER_ERR_STUB("E89999", OPS_DESC, __VA_ARGS__) +#define OPS_REPORT_CUBE_INNER_ERR(OPS_DESC, ...) OPS_INNER_ERR_STUB("E69999", OPS_DESC, __VA_ARGS__) + +#endif // Transformer_COMMON_OPS_ERR_H \ No newline at end of file diff --git a/csrc/ascend/common/include/external/aclnn_kernels/cast.h b/csrc/ascend/common/include/external/aclnn_kernels/cast.h new file mode 100644 index 000000000..25f5c91f7 --- /dev/null +++ b/csrc/ascend/common/include/external/aclnn_kernels/cast.h @@ -0,0 +1,24 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_CAST_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_CAST_H + +#include "opdev/op_executor.h" +#include "opdev/make_op_executor.h" + +namespace l0op { +const aclTensor* Cast(const aclTensor* self, op::DataType dstDtype, aclOpExecutor* executor); + +// 专攻卷积反向定制 +const aclTensor* CastOnlyForConvBackward(const aclTensor* self, op::DataType dstDtype, aclOpExecutor* executor); +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_CAST_H diff --git a/csrc/ascend/common/include/external/aclnn_kernels/common/op_error_check.h b/csrc/ascend/common/include/external/aclnn_kernels/common/op_error_check.h new file mode 100644 index 000000000..48f41354a --- /dev/null +++ b/csrc/ascend/common/include/external/aclnn_kernels/common/op_error_check.h @@ -0,0 +1,251 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_ERROR_CHECK_H__ +#define OP_ERROR_CHECK_H__ + +#include "opdev/op_log.h" +#include "opdev/common_types.h" +#include "opdev/data_type_utils.h" +#include "opdev/shape_utils.h" + +const int32_t NCHW_N_DIM = 0; +const int32_t NCHW_C_DIM = 1; +const int32_t NHWC_N_DIM = 0; +const int32_t NHWC_C_DIM = 3; + +static inline bool IsNullptr(const aclTensor *tensor, const char *name) { + if (tensor == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a proper Tensor but got null for argument %s.", name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclTensorList *tensorList, const char *name) { + if (tensorList == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a proper TensorList but got null for argument %s.", name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclScalar *scalar, const char *name) { + if (scalar == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a value of type number for argument %s but instead found type null.", + name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclIntArray *intArr, const char *name) { + if (intArr == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a value of type List[int] for argument %s but instead found type null.", + name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclBoolArray *boolArr, const char *name) { + if (boolArr == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a value of type List[bool] for argument %s but instead found type null.", + name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclFloatArray *floatArr, const char *name) { + if (floatArr == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a value of type List[float] for argument %s but instead found type \ + null.", name); + return true; + } + return false; +} + +static inline bool CheckDims(const aclTensor *tensor) { + const auto& xShape = tensor->GetViewShape(); + for(size_t i = 0; i < xShape.GetDimNum(); i++) { + if (xShape.GetDim(i) > INT32_MAX) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The tensor's shape cannot be larger than %d.", INT32_MAX); + return false; + } + } + return true; +} + +static inline bool CheckReduceOutShape(const aclTensor *inferOut, const aclTensor *out) +{ + auto const &xShape = inferOut->GetViewShape(); + auto const &yShape = out->GetViewShape(); + if (xShape != yShape) { + if (!(xShape.GetShapeSize() == 1 && yShape.GetShapeSize() == 1)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The out tensor's shape[%s] is not equal with inferOut shape[%s].", + op::ToString(out->GetViewShape()).GetString(), op::ToString(inferOut->GetViewShape()).GetString()); + return false; + } + } + return true; +} + +static inline bool CheckNCDimValid(const aclTensor *self, const aclTensor *out) { + auto format = self->GetStorageFormat(); + int64_t selfDimN = 0; + int64_t selfDimC = 0; + int64_t outDimN = 0; + int64_t outDimC = 0; + if (format == op::Format::FORMAT_NCHW) { + selfDimN = self->GetViewShape().GetDim(NCHW_N_DIM); + selfDimC = self->GetViewShape().GetDim(NCHW_C_DIM); + outDimN = out->GetViewShape().GetDim(NCHW_N_DIM); + outDimC = out->GetViewShape().GetDim(NCHW_C_DIM); + } else if (format == op::Format::FORMAT_NHWC) { + selfDimN = self->GetViewShape().GetDim(NHWC_N_DIM); + selfDimC = self->GetViewShape().GetDim(NHWC_C_DIM); + outDimN = out->GetViewShape().GetDim(NHWC_N_DIM); + outDimC = out->GetViewShape().GetDim(NHWC_C_DIM); + } else { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "Input and output format only support [NCHW, NHWC] format ."); + return false; + } + if ((selfDimN != outDimN) || (selfDimC != outDimC)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "The selfDimN[%ld]/outDimN[%ld] or selfDimC[%ld]/outDimC[%ld] not equal .", + selfDimN, outDimN, selfDimC, outDimC); + return false; + } + return true; +} + + +#define OP_CHECK_NULL(param, retExpr) \ + if (IsNullptr(param, #param)) { \ + retExpr; \ + } + +#define OP_CHECK_DTYPE_NOT_SUPPORT(tensor, supportList, retExpr) \ + if (!CheckType(tensor->GetDataType(), supportList)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Tensor %s not implemented for %s, should be in dtype support list %s.", \ + #tensor, op::ToString(tensor->GetDataType()).GetString(), op::ToString(supportList).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_DTYPE_NOT_MATCH(tensor, expectedDtype, retExpr) \ + if (tensor->GetDataType() != expectedDtype) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Tensor %s expected dtype is %s but found %s.", \ + #tensor, op::ToString(expectedDtype).GetString(), op::ToString(tensor->GetDataType()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_DTYPE_NOT_SAME(tensor1, tensor2, retExpr) \ + if (tensor1->GetDataType() != tensor2->GetDataType()) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected both tensors to have same dtype, but found %s %s and %s %s.", \ + #tensor1, op::ToString(tensor1->GetDataType()).GetString(), \ + #tensor2, op::ToString(tensor2->GetDataType()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_RESULT_DTYPE_CAST_FAILED(dtype, desiredDtype, retExpr); \ + if (!CanCast(dtype, desiredDtype)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Result type %s can't be cast to the desired output type %s.", \ + op::ToString(dtype).GetString(), op::ToString(desiredDtype).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_BROADCAST(tensor1, tensor2, retExpr) \ + if (!CheckBroadcastShape(tensor1->GetViewShape(), tensor2->GetViewShape())) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The size of tensor %s %s must match the size of tensor %s %s.", \ + #tensor1, op::ToString(tensor1->GetViewShape()).GetString(), \ + #tensor2, op::ToString(tensor2->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_BROADCAST_WITH_SHAPE(tensor, shape, retExpr) \ + if (!CheckBroadcastShape(tensor->GetViewShape(), shape)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The size of tensor %s %s must match the size %s.", \ + #tensor, op::ToString(tensor->GetViewShape()).GetString(), op::ToString(shape).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_BROADCAST_AND_INFER_SHAPE(tensor1, tensor2, retShape, retExpr) \ + if (!BroadcastInferShape(tensor1->GetViewShape(), tensor2->GetViewShape(), retShape)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The size of tensor %s %s must match the size of tensor %s %s.", \ + #tensor1, op::ToString(tensor1->GetViewShape()).GetString(), \ + #tensor2, op::ToString(tensor2->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_SHAPE_NOT_EQUAL(tensor1, tensor2, retExpr) \ + if (tensor1->GetViewShape() != tensor2->GetViewShape()) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected tensor for %s to have same size as tensor for %s, but %s does not " \ + "equal %s.", #tensor1, #tensor2, op::ToString(tensor1->GetViewShape()).GetString(), \ + op::ToString(tensor2->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(tensor, shape, retExpr) \ + if (tensor->GetViewShape() != shape) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected tensor for %s to have same size as %s, but got %s.", \ + #tensor, op::ToString(shape).GetString(), op::ToString(tensor->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_WRONG_DIMENSION(tensor, expectedDimNum, retExpr) \ + if (tensor->GetViewShape().GetDimNum() != expectedDimNum) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected %zu dimension input, but got %s with sizes %s.", \ + static_cast(expectedDimNum), #tensor, op::ToString(tensor->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_MAX_DIM(tensor, maxDim, retExpr) \ + if (tensor->GetViewShape().GetDimNum() > static_cast(maxDim)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The %s tensor cannot be larger than %zu dimensions.", \ + #tensor, static_cast(maxDim)); \ + retExpr; \ + } + +#define OP_CHECK_MIN_DIM(tensor, minDim, retExpr) \ + if (tensor->GetViewShape().GetDimNum() < static_cast(minDim)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The %s tensor must have at least %zu dimensions.", \ + #tensor, static_cast(minDim)); \ + retExpr; \ + } + +#define OP_CHECK_COMM_INPUT(workspaceSize, executor) \ + if (workspaceSize == nullptr || executor == nullptr) { \ + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "The workspaceSize or executor is nullptr."); \ + return ACLNN_ERR_PARAM_NULLPTR; \ + } + +#define OP_CHECK_ADD_TO_LAUNCHER_LIST_AICORE(cond, retExpr, errMsg, ...) \ + if (cond) { \ + OP_LOGE(ACLNN_ERR_INNER_STATIC_WORKSPACE_INVALID, errMsg, ##__VA_ARGS__); \ + retExpr; \ + } + +#define OP_CHECK_INFERSHAPE(cond, retExpr, errMsg, ...) \ + if (cond) { \ + OP_LOGE(ACLNN_ERR_INNER_INFERSHAPE_ERROR, errMsg, ##__VA_ARGS__); \ + retExpr; \ + } + +#define OP_CHECK_TENSORLIST_SIZE_EQUAL(tensorlist1, tensorlist2, retExpr) \ + if ((tensorlist1)->Size() != (tensorlist2)->Size()) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, \ + "The %s tensorlist and %s tensorlist must have the same number of tensors, but got %ld and %ld.", \ + #tensorlist1, #tensorlist2, (tensorlist1)->Size(), (tensorlist2)->Size()); \ + retExpr; \ + } + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/include/external/aclnn_kernels/contiguous.h b/csrc/ascend/common/include/external/aclnn_kernels/contiguous.h new file mode 100644 index 000000000..0c49bfdce --- /dev/null +++ b/csrc/ascend/common/include/external/aclnn_kernels/contiguous.h @@ -0,0 +1,89 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_CONTIGUOUS_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_CONTIGUOUS_H + +#include "opdev/op_def.h" +#include "opdev/common_types.h" + +namespace l0op { + +typedef struct { + // 每个op::Shape 18ns + int64_t viewOffset; + + // Transpose + op::Shape transposeSrcShape; + op::Shape transposeDstShape; + op::FVector perm; + + // broadcast to + op::Shape broadcastSrcShape; + op::Shape broadcastDstShape; + op::FVector shape; + + // slice + op::Shape sliceSrcShape; + op::Shape sliceDstShape; + op::FVector offset; + op::FVector size; + + // strided slice + op::Shape stridedsliceSrcShape; + op::Shape stridedsliceDstShape; + op::FVector begin; + op::FVector end; + op::FVector strides; + + // optimizer + bool mayBroadcast; + bool mayTranspose; + bool maySlice; + bool mayStridedslice; +} ContiguousParam; + +/** + * @brief 将非连续Tensor转换为连续Tensor + * @param x + * @param executor + * @return aclTensor 转换后的tensor + */ +const aclTensor* Contiguous(const aclTensor* x, aclOpExecutor* executor); + +/** + * @brief 将连续tensor拷贝到非连续的tensor上 + * @param x + * @param y + * @param executor + * @return aclTensor 转换后的tensor + */ +const aclTensor* ViewCopy(const aclTensor* x, const aclTensor* y, aclOpExecutor* executor); + +/** + * @brief 对Tensor创建一个View，要求Tensor满足PickView的条件 + * @param x 输入Tensor，可以是一整块的非连续Tensor + * @param executor + * @return 输出Shape是一个连续Tensor + */ +const aclTensor* PickViewAsContiguous(const aclTensor* x, aclOpExecutor* executor); + +const aclTensor* ReViewToOut(const aclTensor* x, const aclTensor* y, aclOpExecutor* executor); + +// ============内部接口============= +bool CanOptimizeContiguous( + const op::Shape& viewShape, const op::Strides& strides, int64_t offset, int64_t storageSize, + ContiguousParam& param); + +bool CanOptimizeView(const op::Shape& viewShape, const op::Strides& strides, int64_t offset, ContiguousParam& param); +// ============内部接口============= +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_CONTIGUOUS_H diff --git a/csrc/ascend/common/include/external/aclnn_kernels/pad.h b/csrc/ascend/common/include/external/aclnn_kernels/pad.h new file mode 100644 index 000000000..ab740a435 --- /dev/null +++ b/csrc/ascend/common/include/external/aclnn_kernels/pad.h @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_PAD_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_PAD_H + +#include "opdev/op_executor.h" +#include "opdev/make_op_executor.h" + +namespace l0op { +const aclTensor* Pad(const aclTensor* self, const aclTensor* paddings, aclOpExecutor* executor); +} +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_PAD_H diff --git a/csrc/ascend/common/include/external/aclnn_kernels/reshape.h b/csrc/ascend/common/include/external/aclnn_kernels/reshape.h new file mode 100644 index 000000000..6a9a62d2c --- /dev/null +++ b/csrc/ascend/common/include/external/aclnn_kernels/reshape.h @@ -0,0 +1,37 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_RESHAPE_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_RESHAPE_H + +#include "opdev/shape_utils.h" +#include "opdev/op_def.h" + +namespace l0op { +/** + * @brief Modify input tensor's shape. + * @param x Input Tensor. Should be contiguous. + * @param shape Target Shape. Only one dimension can be -1. + * @param executor aclOpExecutor.ldd + * @return *aclTensor Output tensor. + */ +const aclTensor* Reshape(const aclTensor* x, const op::Shape& shape, aclOpExecutor* executor); + +/** + * @brief Modify input tensor's shape. + * @param x Input Tensor. Should be contiguous. + * @param shape Target Shape. Only one dimension can be -1. + * @param executor aclOpExecutor. + * @return *aclTensor Output tensor. + */ +const aclTensor* Reshape(const aclTensor* x, const aclIntArray* shape, aclOpExecutor* executor); +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_RESHAPE_H diff --git a/csrc/ascend/common/include/external/aclnn_kernels/slice.h b/csrc/ascend/common/include/external/aclnn_kernels/slice.h new file mode 100644 index 000000000..3594f478c --- /dev/null +++ b/csrc/ascend/common/include/external/aclnn_kernels/slice.h @@ -0,0 +1,25 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_SLICE_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_SLICE_H + +#include "opdev/op_def.h" + +namespace l0op { + +const aclTensor* Slice( + const aclTensor* x, const aclTensor* y, const aclTensor* offset, const aclTensor* size, aclOpExecutor* executor); + +const aclTensor* Slice( + const aclTensor* x, const aclIntArray* offsets, const aclIntArray* size, aclOpExecutor* executor); +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_SLICE_H diff --git a/csrc/ascend/common/include/external/aclnn_kernels/transdata.h b/csrc/ascend/common/include/external/aclnn_kernels/transdata.h new file mode 100644 index 000000000..33e7bd7df --- /dev/null +++ b/csrc/ascend/common/include/external/aclnn_kernels/transdata.h @@ -0,0 +1,56 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSDATA_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSDATA_H + +#include "opdev/op_executor.h" + +namespace l0op { + +const aclTensor* ReFormat(const aclTensor* x, const op::Format& format, aclOpExecutor* executor = nullptr); + +/** + * TransData + * Formal Transdata. Set the c0 size strictly based on the data type and chip block size. + * support data type as follows: fp16,fp32,int32,uint32,int8,uint8 + * fp16: block_size/2 + * fp32/int32/uint32: block_size/4 (this is different from `TransDataSpecial`) + * int8/uint8: block_size/1 + * + * @param x : aclTensor need to transpose + * @param dstPrimaryFormat: dstPrimaryFormat like NC1HWC0 + * @param groups: groups + * @param executor: executor should not be null + * @return trans format tensor + */ +const aclTensor* TransData(const aclTensor* x, op::Format dstPrimaryFormat, int64_t groups, aclOpExecutor* executor); +/** + * Special Transdata. Set the c0 size strictly based on the data type and chip block size. + * this transdata c0 size rule: + * fp16: block_size/2 + * fp32/int32/uint32: block_size/2 + * int8/uint8: block_size/1 + * bool not supported, should do: + * (NCHW, bool)-> cast -> (NCHW, fp16) -> TransDataSpecial -> (5HD, fp16) -> cast -> (5HD, bool) + * (5HD, bool)-> cast -> (5HD, fp16) -> TransDataSpecial -> (NCHW, fp16) -> cast -> (NCHW, bool) + * + * @param x : aclTensor need to transpose + * @param dstPrimaryFormat: dstPrimaryFormat like NC1HWC0 + * @param groups: groups + * @param executor: executor should not be null + * @return trans format tensor + */ +const aclTensor* TransDataSpecial( + const aclTensor* x, op::Format dstPrimaryFormat, int64_t groups, aclOpExecutor* executor); + +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSDATA_H diff --git a/csrc/ascend/common/include/external/aclnn_kernels/transpose.h b/csrc/ascend/common/include/external/aclnn_kernels/transpose.h new file mode 100644 index 000000000..b67fa20a9 --- /dev/null +++ b/csrc/ascend/common/include/external/aclnn_kernels/transpose.h @@ -0,0 +1,22 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSPOSE_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSPOSE_H + +#include "opdev/op_def.h" + +namespace l0op { + +const aclTensor* Transpose(const aclTensor* x, const aclTensor* y, const aclTensor* perm, aclOpExecutor* executor); +const aclTensor* Transpose(const aclTensor* x, const aclIntArray* perm, aclOpExecutor* executor); +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSPOSE_H diff --git a/csrc/ascend/common/include/external/aclnn_util.h b/csrc/ascend/common/include/external/aclnn_util.h new file mode 100644 index 000000000..1b46d2df2 --- /dev/null +++ b/csrc/ascend/common/include/external/aclnn_util.h @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file aclnn_util.h + * \brief + */ +#ifndef Transformer_COMMON_ACLNN_UTIL_H +#define Transformer_COMMON_ACLNN_UTIL_H + +#define ACLNN_API __attribute__((visibility("default"))) + +#endif // Transformer_COMMON_ACLNN_UTIL_H \ No newline at end of file diff --git a/csrc/ascend/common/include/fallback/fallback.h b/csrc/ascend/common/include/fallback/fallback.h new file mode 100644 index 000000000..4fb9fc690 --- /dev/null +++ b/csrc/ascend/common/include/fallback/fallback.h @@ -0,0 +1,499 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file fallback.h + * \brief + */ + +#ifndef ACLNNFALLBACK_OPAPI_H_ +#define ACLNNFALLBACK_OPAPI_H_ + +#include + +#include +#include +#include +#include + +#include "aclnn/aclnn_base.h" +#include "fallback/fallback_comm.h" +#include "mc2_log.h" +#include "runtime/base.h" +#include "log/log.h" + +namespace fallback { +using namespace std; +using namespace gert; +using namespace ge; +using namespace std; + +namespace std_utils { + template + struct index_sequence {}; + + template + struct make_index_sequence_helper : make_index_sequence_helper {}; + + template + struct make_index_sequence_helper<0, Is...> { + using type = index_sequence; + }; + + template + using make_index_sequence = typename make_index_sequence_helper::type; +} + +using aclOpExecutor = struct aclOpExecutor; +using aclTensor = struct aclTensor; +using aclScalar = struct aclScalar; +using aclIntArray = struct aclIntArray; +using aclFloatArray = struct aclFloatArray; +using aclBoolArray = struct aclBoolArray; +using aclTensorList = struct aclTensorList; + +using _aclCreateTensor = aclTensor* (*)(const int64_t* view_dims, uint64_t view_dims_num, aclDataType data_type, + const int64_t* stride, int64_t offset, aclFormat format, + const int64_t* storage_dims, uint64_t storage_dims_num, void* tensor_data); + +using _aclCreateScalar = aclScalar* (*)(void* value, aclDataType data_type); +using _aclCreateIntArray = aclIntArray* (*)(const int64_t* value, uint64_t size); +using _aclCreateFloatArray = aclFloatArray* (*)(const float* value, uint64_t size); +using _aclCreateBoolArray = aclBoolArray* (*)(const bool* value, uint64_t size); +using _aclCreateTensorList = aclTensorList* (*)(const aclTensor* const *value, uint64_t size); + +using _aclDestroyTensor = int (*)(const aclTensor* tensor); +using _aclDestroyScalar = int (*)(const aclScalar* scalar); +using _aclDestroyIntArray = int (*)(const aclIntArray* array); +using _aclDestroyFloatArray = int (*)(const aclFloatArray* array); +using _aclDestroyBoolArray = int (*)(const aclBoolArray* array); +using _aclDestroyTensorList = int (*)(const aclTensorList* array); + +#define GET_OP_API_FUNC(apiName) reinterpret_cast<_##apiName>(GetOpApiFuncAddr(#apiName)) + +inline const char* GetOpApiLibName(void) { + return "libopapi.so"; +} + +inline const char* GetCustOpApiLibName(void) { + return "libcust_opapi.so"; +} + +inline void* GetOpApiFuncAddrInLib(void* handler, const char* libName, const char* apiName) { + auto funcAddr = dlsym(handler, apiName); + if (funcAddr == nullptr) { + OP_LOGW("aclnnfallback", "dlsym %s from %s failed, error:%s.", apiName, libName, dlerror()); + } + return funcAddr; +} + +inline void* GetOpApiLibHandler(const char* libName) { + auto handler = dlopen(libName, RTLD_LAZY); + if (handler == nullptr) { + OP_LOGW("aclnnfallback", "dlopen %s failed, error:%s.", libName, dlerror()); + } + return handler; +} + +inline void* GetAclnnArrdByApiName(const char *apiName) { + vector libs = {"libaclnn_ops_infer.so", "libaclnn_ops_train.so", "libaclnn_math.so", + "libaclnn_rand.so", "libaclnn_sparse.so", "libaclnn_fft.so"}; + for (const auto &libName : libs) { + static auto libHandler = GetOpApiLibHandler(libName.c_str()); + if (libHandler != nullptr) { + auto funcAddr = GetOpApiFuncAddrInLib(libHandler, libName.c_str(), apiName); + if (funcAddr != nullptr) { + return funcAddr; + } + } + } + OP_LOGE("aclnnfallback", "api %s can't find in any aclnn lib.", apiName); + return nullptr; +} + +inline void* GetOpApiFuncAddr(const char* apiName) { + static auto custOpApiHandler = GetOpApiLibHandler(GetCustOpApiLibName()); + if (custOpApiHandler != nullptr) { + auto funcAddr = GetOpApiFuncAddrInLib(custOpApiHandler, GetCustOpApiLibName(), apiName); + if (funcAddr != nullptr) { + return funcAddr; + } + } + + static auto opApiHandler = GetOpApiLibHandler(GetOpApiLibName()); + if (opApiHandler != nullptr) { + auto funcAddr = GetOpApiFuncAddrInLib(opApiHandler, GetOpApiLibName(), apiName); + if (funcAddr != nullptr) { + return funcAddr; + } + } + OP_LOGD("aclnnfallback", "opapi lib is not exist,will use aclnn lib."); + return GetAclnnArrdByApiName(apiName); +} + +inline aclTensor* ConvertType(aclTensor* ge_tensor) { + return ge_tensor; +} + +inline aclIntArray* ConvertType(const std::vector &arr) { + if (arr.empty()) { + return nullptr; + } + static const auto aclCreateIntArray = GET_OP_API_FUNC(aclCreateIntArray); + auto array = aclCreateIntArray(arr.data(), arr.size()); + return array; +} + +inline aclDataType GetConvertType(const gert::Tensor* ge_tensor) { + // convert data type + auto dataType_ge = ge_tensor->GetDataType(); + auto dataType = aclDataType::ACL_FLOAT16; + if (dataType_ge == DT_FLOAT) { + dataType = aclDataType::ACL_FLOAT; + } else if (dataType_ge == DT_BF16) { + dataType = aclDataType::ACL_BF16; + } else if (dataType_ge == DT_BOOL) { + dataType = aclDataType::ACL_BOOL; + } else if (dataType_ge == DT_INT64) { + dataType = aclDataType::ACL_INT64; + } else if (dataType_ge == DT_INT32) { + dataType = aclDataType::ACL_INT32; + } else if (dataType_ge == DT_UINT64) { + dataType = aclDataType::ACL_UINT64; + } else if (dataType_ge == DT_UINT32) { + dataType = aclDataType::ACL_UINT32; + } else if (dataType_ge == DT_INT8) { + dataType = aclDataType::ACL_INT8; + } else if (dataType_ge == DT_UINT8) { + dataType = aclDataType::ACL_UINT8; + } else if (dataType_ge == DT_INT4) { + dataType = aclDataType::ACL_INT4; + } else if (dataType_ge == DT_FLOAT8_E4M3FN) { + dataType = aclDataType::ACL_FLOAT8_E4M3FN; + } else { + dataType = aclDataType::ACL_FLOAT16; + } + + return dataType; +} + +inline aclTensor* ConvertType(const gert::Tensor* ge_tensor) { + if (ge_tensor == nullptr) { + return nullptr; + } + + static const auto aclCreateTensor = GET_OP_API_FUNC(aclCreateTensor); + OP_CHECK_IF(aclCreateTensor == nullptr, OP_LOGE("aclnnfallback", "aclCreateTensor nullptr"), return nullptr); + + void* device_addr = nullptr; + device_addr = const_cast(ge_tensor->GetAddr()); + + auto dataType = GetConvertType(ge_tensor); + + OP_LOGD("aclnnfallback", "aclCreateTensor: tensor type is %d", dataType); + + // convert shape + auto gert_shape = ge_tensor->GetStorageShape(); + std::vector shape; + for (size_t i = 0; i < gert_shape.GetDimNum(); ++i) { + shape.push_back(gert_shape.GetDim(i)); + } + + // 计算连续tensor的strides + std::vector strides(shape.size(), 1); + for (int64_t i = shape.size() - 2; i >= 0; i--) { + strides[i] = shape[i + 1] * strides[i + 1]; + } + + aclTensor* out = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), + 0, aclFormat::ACL_FORMAT_ND, + shape.data(), shape.size(), device_addr); + + OP_CHECK_IF(out == nullptr, + OP_LOGE("aclnnfallback", "out nullptr"), return nullptr); + + return out; +} + +inline aclTensorList* ConvertType(std::vector& ge_tenserList) { + OP_CHECK_IF(ge_tenserList.size() == 0, + OP_LOGE("aclnnfallback", "ge_tenserList size 0"), return nullptr); + + static const auto aclCreateTensorList = GET_OP_API_FUNC(aclCreateTensorList); + OP_CHECK_IF(aclCreateTensorList == nullptr, + OP_LOGE("aclnnfallback", "ge_tenserList size 0"), return nullptr); + + std::vector tmp; + for (size_t i = 0; i < ge_tenserList.size(); i++) { + auto t_acl = ConvertType(ge_tenserList[i]); + tmp.push_back(t_acl); + } + + aclTensorList* tensorList = aclCreateTensorList(tmp.data(), tmp.size()); + return tensorList; +} + +template +inline aclScalar* ConvertScalarType(T value) { + static const auto aclCreateScalar = GET_OP_API_FUNC(aclCreateScalar); + OP_CHECK_IF(aclCreateScalar == nullptr, + OP_LOGE("aclnnfallback", "aclCreateScalar nullptr"), return nullptr); + if (typeid(value) == typeid(float)) { + return aclCreateScalar(&value, aclDataType::ACL_FLOAT); + } + return nullptr; +} + +template +T ConvertType(T value) { + return value; +} + +inline aclTensor* ConvertMmType(const gert::Tensor* ge_tensor, bool transpose, bool enable_NZ=false) { + if (ge_tensor == nullptr) { + return nullptr; + } + auto gert_shape = ge_tensor->GetStorageShape(); + if (gert_shape.GetDimNum() <= 1) { + return ConvertType(ge_tensor); + } + + static const auto aclCreateTensor = GET_OP_API_FUNC(aclCreateTensor); + OP_CHECK_IF(aclCreateTensor == nullptr, OP_LOGE("aclnnfallback", "aclCreateTensor nullptr"), return nullptr); + + void* device_addr = const_cast(ge_tensor->GetAddr()); + // convert data type + auto dataType_ge = ge_tensor->GetDataType(); + auto dataType = ToAclDataType(dataType_ge); + // convert shape + std::vector shape; + for (size_t i = 0; i < gert_shape.GetDimNum(); ++i) { + shape.push_back(gert_shape.GetDim(i)); + } + // 计算连续tensor的strides + std::vector strides(shape.size(), 1); + for (int64_t i = shape.size() - 2; i >= 0; i--) { + strides[i] = shape[i + 1] * strides[i + 1]; + } + + auto viewShape = shape; + // 对于transpose后的tensor对后两维度进行strides, viewShape转换 + if (transpose) { + // dimM 为倒数第二维， dimN 为倒数第一维度 + auto dimM = shape.size() - 2; + auto dimN = shape.size() - 1; + auto swap = strides[dimN]; + strides[dimN] = strides[dimM]; + strides[dimM] = swap; + // 修改viewShape + viewShape[dimN] = shape[dimM]; + viewShape[dimM] = shape[dimN]; + } + auto acl_format = aclFormat::ACL_FORMAT_ND; + if (enable_NZ && GetPrimaryFormat(ge_tensor->GetStorageFormat()) == ge::Format::FORMAT_FRACTAL_NZ) { + acl_format = aclFormat::ACL_FORMAT_FRACTAL_NZ; + } + aclTensor* out = aclCreateTensor(viewShape.data(), shape.size(), dataType, strides.data(), + 0, acl_format, shape.data(), shape.size(), device_addr); + OP_CHECK_IF(out == nullptr, OP_LOGE("aclnnfallback", "out nullptr"), return nullptr); + + return out; +} + +inline void Release(aclTensor* p) { + static const auto aclDestroyTensor = GET_OP_API_FUNC(aclDestroyTensor); + OP_CHECK_IF(aclDestroyTensor == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyTensor is null"), return); + aclDestroyTensor(p); +} + +inline void Release(aclScalar* p) { + static const auto aclDestroyScalar = GET_OP_API_FUNC(aclDestroyScalar); + OP_CHECK_IF(aclDestroyScalar == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyScalar is null"), return); + aclDestroyScalar(p); +} + +inline void Release(aclIntArray* p) { + static const auto aclDestroyIntArray = GET_OP_API_FUNC(aclDestroyIntArray); + OP_CHECK_IF(aclDestroyIntArray == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyIntArray is null"), return); + aclDestroyIntArray(p); +} + +inline void Release(aclBoolArray* p) { + static const auto aclDestroyBoolArray = GET_OP_API_FUNC(aclDestroyBoolArray); + OP_CHECK_IF(aclDestroyBoolArray == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyBoolArray is null"), return); + aclDestroyBoolArray(p); +} + +inline void Release(aclTensorList* p) { + static const auto aclDestroyTensorList = GET_OP_API_FUNC(aclDestroyTensorList); + OP_CHECK_IF(aclDestroyTensorList == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyTensorList is null"), return); + aclDestroyTensorList(p); +} + +template +void Release(T value) { + (void)value; +} + +template +void CallRelease(Tuple t, std_utils::index_sequence) { + (void)std::initializer_list{(Release(std::get(t)), 0)...}; +} + +template +void ReleaseConvertTypes(Tuple& t) { + static constexpr auto size = std::tuple_size::value; + CallRelease(t, std_utils::make_index_sequence{}); +} + +template +auto ConvertTypes(Ts&... args) -> decltype(std::make_tuple(ConvertType(args)...)) { + auto tp = std::make_tuple(ConvertType(args)...); + return tp; +} + +template +auto call(Function f, Tuple t, std_utils::index_sequence) -> int { + return f(std::get(t)...); +} + +template +auto call(Function f, Tuple t) -> int { + static constexpr auto size = std::tuple_size::value; + return call(f, t, std_utils::make_index_sequence{}); +} + +template +auto ConvertToOpApiFunc(const Tuple& params, void* opApiAddr, std_utils::index_sequence) + -> int (*)(typename std::decay(params))>::type...) { + using LocalOpApiFunc = int (*)(typename std::decay(params))>::type...); + auto func = reinterpret_cast(opApiAddr); + return func; +} + +template +auto ConvertToOpApiFunc(const Tuple& params, void* opApiAddr) + -> typename std::enable_if::value != 0, + decltype(ConvertToOpApiFunc(params, opApiAddr, std_utils::make_index_sequence::value>{}))>::type { + static constexpr auto size = std::tuple_size::value; + return ConvertToOpApiFunc(params, opApiAddr, std_utils::make_index_sequence{}); +} + +template +class ConvertedParams { + public: + ConvertedParams(Tuple&& convertedParams) : convertedParams_(std::move(convertedParams)){}; + ConvertedParams(ConvertedParams&& other) : convertedParams_(std::move(other.convertedParams_)) { + other.validParams_ = false; + }; + ConvertedParams& operator=(ConvertedParams&& other) { + if (this == &other) { + return *this; + } + + convertedParams_ = std::move(other.convertedParams_); + validParams_ = true; + other.validParams_ = false; + return *this; + } + + ConvertedParams() = delete; + ConvertedParams(const ConvertedParams& other) = delete; + ConvertedParams& operator=(const ConvertedParams& other) = delete; + + ~ConvertedParams() { + if (validParams_) { + ReleaseConvertTypes(convertedParams_); + } + } + + const Tuple& GetConvertedParams() const { + return convertedParams_; + } + + private: + Tuple convertedParams_; + bool validParams_{true}; +}; + +using InitHugeMemThreadLocal = int (*)(void*, bool); +using UnInitHugeMemThreadLocal = void (*)(void*, bool); +using ReleaseHugeMem = void (*)(void*, bool); +using PTAGetExecCache = aclOpExecutor* (*)(uint64_t, uint64_t*); +using InitPTACacheThreadLocal = void (*)(); +using SetPTAHashKey = void (*)(uint64_t); +using CanUsePTACache = bool (*)(const char*); + +using ResetCacheThreadLocal = void (*)(); + +#define EXEC_OPAPI_CMD(aclnn_api, ...) \ + ({ \ + static auto ret = GRAPH_SUCCESS; \ + do { \ + static const auto ResetCacheThreadLocalAddr = GetOpApiFuncAddr("ResetCacheThreadLocal"); \ + static const auto getWorkspaceSizeFuncAddr = GetOpApiFuncAddr(#aclnn_api "GetWorkspaceSize"); \ + static const auto opApiFuncAddr = GetOpApiFuncAddr(#aclnn_api); \ + if (getWorkspaceSizeFuncAddr == nullptr || opApiFuncAddr == nullptr || ResetCacheThreadLocalAddr == nullptr) { \ + OP_LOGE("aclnnfallback", "%s or %s not in %s or %s or ResetCacheThreadLocal not found.", \ + #aclnn_api "GetWorkspaceSize", #aclnn_api, GetOpApiLibName(), GetOpApiLibName()); \ + ret = GRAPH_FAILED; \ + break; \ + } \ + auto ResetCacheThreadLocalFunc = reinterpret_cast(ResetCacheThreadLocalAddr); \ + ResetCacheThreadLocalFunc(); \ + uint64_t workspace_size = 0; \ + uint64_t* workspace_size_addr = &workspace_size; \ + aclOpExecutor* executor = nullptr; \ + aclOpExecutor** executor_addr = &executor; \ + auto converted_params = ConvertTypes(__VA_ARGS__, workspace_size_addr, executor_addr); \ + static auto getWorkspaceSizeFunc = ConvertToOpApiFunc(converted_params, getWorkspaceSizeFuncAddr); \ + auto workspace_status = call(getWorkspaceSizeFunc, converted_params); \ + if (workspace_status != 0) { \ + OP_LOGE("aclnnfallback", "call %s failed:", #aclnn_api); \ + ret = GRAPH_FAILED; \ + break; \ + } \ + void* workspace_addr = nullptr; \ + if (workspace_size > 0) { \ + workspace_addr = host_api_ctx->MallocWorkspace(workspace_size); \ + if (workspace_addr == nullptr) { \ + OP_LOGE("aclnnfallback", "call %s allocate workspace failed", #aclnn_api); \ + ret = GRAPH_FAILED; \ + break; \ + } \ + } \ + auto acl_stream = host_api_ctx->GetStream(); \ + auto acl_call = [converted_params, workspace_addr, workspace_size, host_api_ctx, acl_stream, \ + executor]() -> int { \ + using OpApiFunc = int (*)(void*, uint64_t, aclOpExecutor*, const aclrtStream); \ + OpApiFunc opApiFunc = reinterpret_cast(opApiFuncAddr); \ + auto api_ret_inner = opApiFunc(workspace_addr, workspace_size, executor, acl_stream); \ + ReleaseConvertTypes(converted_params); \ + host_api_ctx->FreeWorkspace(); \ + if (api_ret_inner != 0) { \ + OP_LOGE("aclnnfallback", "call %s allocate workspace failed api_ret_inner: %d", #aclnn_api, api_ret_inner); \ + return GRAPH_FAILED; \ + } \ + return api_ret_inner; \ + }; \ + \ + ret = acl_call(); \ + } while (false); \ + (ret); \ + }) + +} // namespace fallback + +#endif // ACLNNFALLBACK_OPAPI_H_ diff --git a/csrc/ascend/common/include/fallback/fallback_2stages.h b/csrc/ascend/common/include/fallback/fallback_2stages.h new file mode 100644 index 000000000..8564403ad --- /dev/null +++ b/csrc/ascend/common/include/fallback/fallback_2stages.h @@ -0,0 +1,126 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef ACLNNFALLBACK_OPAPI_TWOSTAGES_H_ +#define ACLNNFALLBACK_OPAPI_TWOSTAGES_H_ + +#include + +#include +#include +#include +#include +#include + +#include "aclnn/aclnn_base.h" +#include "fallback.h" +#include "fallback_comm.h" +#include "fallback_comm_2stages.h" +#include "log/log.h" +#include "mc2_log.h" + +namespace fallback { +using namespace std; +using namespace gert; +using namespace ge; + +inline void Collect(aclTensor *p, std::vector ¶ms) { + static const auto aclDestroyTensor = GET_OP_API_FUNC(aclDestroyTensor); + OPS_ERR_IF(aclDestroyTensor == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyTensor is null"), return); + params.emplace_back(OpApiAnyValue{p, [](void *param) {aclDestroyTensor(static_cast(param));}}); +} + +inline void Collect(aclScalar *p, std::vector ¶ms) { + static const auto aclDestroyScalar = GET_OP_API_FUNC(aclDestroyScalar); + OPS_ERR_IF(aclDestroyScalar == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyScalar is null"), return); + params.emplace_back(OpApiAnyValue{p, [](void *param) {aclDestroyScalar(static_cast(param));}}); +} + +inline void Collect(aclIntArray *p, std::vector ¶ms) { + static const auto aclDestroyIntArray = GET_OP_API_FUNC(aclDestroyIntArray); + OPS_ERR_IF(aclDestroyIntArray == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyIntArray is null"), return); + params.emplace_back(OpApiAnyValue{p, [](void *param) {aclDestroyIntArray(static_cast(param));}}); +} + +inline void Collect(aclBoolArray *p, std::vector ¶ms) { + static const auto aclDestroyBoolArray = GET_OP_API_FUNC(aclDestroyBoolArray); + OPS_ERR_IF(aclDestroyBoolArray == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyBoolArray is null"), return); + params.emplace_back(OpApiAnyValue{p, [](void *param) {aclDestroyBoolArray(static_cast(param));}}); +} + +inline void Collect(aclTensorList *p, std::vector ¶ms) { + static const auto aclDestroyTensorList = GET_OP_API_FUNC(aclDestroyTensorList); + OPS_ERR_IF(aclDestroyTensorList == nullptr, + OP_LOGE("aclnnfallback", "aclDestroyTensorList is null"), return); + params.emplace_back(OpApiAnyValue{p, [](void *param) {aclDestroyTensorList(static_cast(param));}}); +} + +template +void Collect(T value, std::vector ¶ms) { + (void)value; + params.emplace_back(OpApiAnyValue{nullptr, nullptr}); +} + +template +void CallCollect(Tuple t, std_utils::index_sequence, std::vector ¶ms) { + (void)std::initializer_list{(Collect(std::get(t), params), 0)...}; +} + +template +void CollectConvertedTypes(Tuple &t, std::vector ¶ms) { + static constexpr auto size = std::tuple_size::value; + CallCollect(t, std_utils::make_index_sequence{}, params); +} + +#define EXEC_OPAPI_PREPARE_CMD(aclnn_api, ...) \ + ({ \ + static auto ret = GRAPH_SUCCESS; \ + do { \ + static const auto ResetCacheThreadLocalAddr = GetOpApiFuncAddr("ResetCacheThreadLocal"); \ + static const auto getWorkspaceSizeFuncAddr = GetOpApiFuncAddr(#aclnn_api "GetWorkspaceSize"); \ + static const auto opApiFuncAddr = GetOpApiFuncAddr(#aclnn_api); \ + if (getWorkspaceSizeFuncAddr == nullptr || opApiFuncAddr == nullptr || ResetCacheThreadLocalAddr == nullptr) { \ + OP_LOGE("aclnnfallback", "%s or %s not in %s or %s or ResetCacheThreadLocal not found.", \ + #aclnn_api "GetWorkspaceSize", #aclnn_api, GetOpApiLibName(), GetOpApiLibName()); \ + ret = GRAPH_FAILED; \ + break; \ + } \ + auto *op_api_params = new (std::nothrow) OpApiParams(); \ + auto ResetCacheThreadLocalFunc = reinterpret_cast(ResetCacheThreadLocalAddr); \ + ResetCacheThreadLocalFunc(); \ + op_api_params->op_api_func = reinterpret_cast(opApiFuncAddr); \ + uint64_t workspace_size = 0; \ + uint64_t* workspace_size_addr = &workspace_size; \ + aclOpExecutor** executor_addr = &op_api_params->executor; \ + auto converted_params = ConvertTypes(__VA_ARGS__, workspace_size_addr, executor_addr); \ + using TupleT = decltype(converted_params); \ + constexpr size_t tuple_size = std::tuple_size::value; \ + op_api_params->converted_params.reserve(tuple_size); \ + CollectConvertedTypes(converted_params, op_api_params->converted_params); \ + host_api_ctx->SetOpApiParamsWithDefaultDeleter(op_api_params); \ + static auto getWorkspaceSizeFunc = ConvertToOpApiFunc(converted_params, getWorkspaceSizeFuncAddr); \ + auto workspace_status = call(getWorkspaceSizeFunc, converted_params); \ + if (workspace_status != 0) { \ + OP_LOGE("aclnnfallback", "call %s failed:", #aclnn_api); \ + ret = GRAPH_FAILED; \ + break; \ + } \ + ret = host_api_ctx->SetWorkspaceSizes({workspace_size}); \ + } while (false); \ + (ret); \ + }) + +} // namespace fallback + +#endif // ACLNNFALLBACK_OPAPI_TWOSTAGES_H_ \ No newline at end of file diff --git a/csrc/ascend/common/include/fallback/fallback_comm.h b/csrc/ascend/common/include/fallback/fallback_comm.h new file mode 100644 index 000000000..778c568c6 --- /dev/null +++ b/csrc/ascend/common/include/fallback/fallback_comm.h @@ -0,0 +1,42 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file fallback_comm.h + * \brief + */ + +#ifndef INC_EXTERNAL_GRAPH_FALLBACK_COMMON_H_ +#define INC_EXTERNAL_GRAPH_FALLBACK_COMMON_H_ + +#include "aclnn/aclnn_base.h" +#include "exe_graph/runtime/op_execute_context.h" +#include "exe_graph/runtime/tensor.h" +#include "register/op_impl_registry.h" +#if __has_include("runtime/base.h") +#include "runtime/base.h" +#else +#include "runtime/rt_external_base.h" +#endif + +#ifdef __cplusplus +extern "C" { +#endif + +namespace fallback { + +aclDataType ToAclDataType(ge::DataType dtype); +} // namespace fallback + +#ifdef __cplusplus +} +#endif + +#endif // INC_EXTERNAL_GRAPH_FALLBACK_COMMON_H_ diff --git a/csrc/ascend/common/include/fallback/fallback_comm_2stages.h b/csrc/ascend/common/include/fallback/fallback_comm_2stages.h new file mode 100644 index 000000000..dbd741ecf --- /dev/null +++ b/csrc/ascend/common/include/fallback/fallback_comm_2stages.h @@ -0,0 +1,56 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef INC_EXTERNAL_GRAPH_FALLBACK_COMMON_TWOSTAGES_H_ +#define INC_EXTERNAL_GRAPH_FALLBACK_COMMON_TWOSTAGES_H_ + +#include "aclnn/aclnn_base.h" +#include "aclnn/acl_meta.h" +#include "exe_graph/runtime/op_execute_context.h" +#include "exe_graph/runtime/op_execute_prepare_context.h" +#include "exe_graph/runtime/op_execute_launch_context.h" +#include "exe_graph/runtime/tensor.h" +#include "register/op_impl_kernel_registry.h" +#include "register/op_impl_registry.h" +#if __has_include("runtime/base.h") +#include "runtime/base.h" +#else +#include "runtime/rt_external_base.h" +#endif + +#ifdef __cplusplus +extern "C" { +#endif + +namespace fallback { + +using OpApiAnyValueDeleter = void (*)(void *); +typedef struct { + void *pointer; + OpApiAnyValueDeleter deleter; +} OpApiAnyValue; + +// aclnn算子params结构体，用于传递算子一阶段到二阶段的参数，定义在算子仓，由算子感知，GE框架不感知 +using OpApiFunc = int (*)(void *, uint64_t, aclOpExecutor *, const aclrtStream); +struct OpApiParams { + std::vector converted_params; // 算子下发依赖的参数 + aclOpExecutor *executor = nullptr; // aclOpExecutor指针 + OpApiFunc op_api_func = nullptr; // aclnnxx函数指针，实现算子launch下发 +}; + +// aclnn算子注册的二阶段launch func，函数实现可以与算子类型无关，所有算子使用同一个二阶段注册接口 +ge::graphStatus ExecuteOpLaunch(gert::OpExecuteLaunchContext *context); +} // namespace fallback + +#ifdef __cplusplus +} +#endif + +#endif // INC_EXTERNAL_GRAPH_FALLBACK_COMMON_TWOSTAGES_H_ \ No newline at end of file diff --git a/csrc/ascend/common/include/framework/onnx_common.h b/csrc/ascend/common/include/framework/onnx_common.h new file mode 100644 index 000000000..f88455e83 --- /dev/null +++ b/csrc/ascend/common/include/framework/onnx_common.h @@ -0,0 +1,85 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file onnx_common.h + * \brief + */ + +#ifndef MATH_COMMON_ONNX_COMMON_H +#define MATH_COMMON_ONNX_COMMON_H + +#include +#include +#include + +#include "stub_ops.h" +#include "register/register.h" +#include "graph/operator.h" +#include "graph/graph.h" +#include "base/err_msg.h" +#include "log/log.h" +#include "onnx/proto/ge_onnx.pb.h" + +namespace domi { +template +inline std::string GetOpName(const T& op) +{ + ge::AscendString op_ascend_name; + ge::graphStatus ret = op.GetName(op_ascend_name); + if (ret != ge::GRAPH_SUCCESS) { + std::string op_name = "None"; + return op_name; + } + return op_ascend_name.GetString(); +} + +template +inline ge::Tensor Vec2Tensor(vector& vals, const vector& dims, ge::DataType dtype, ge::Format format = ge::FORMAT_ND) { + ge::Shape shape(dims); + ge::TensorDesc desc(shape, format, dtype); + ge::Tensor tensor(desc, reinterpret_cast(vals.data()), vals.size() * sizeof(T)); + return tensor; +} + +template +inline ge::Tensor CreateScalar(T val, ge::DataType dtype, ge::Format format = ge::FORMAT_ND) { + vector dims_scalar = {}; + ge::Shape shape(dims_scalar); + ge::TensorDesc desc(shape, format, dtype); + ge::Tensor tensor(desc, reinterpret_cast(&val), sizeof(T)); + return tensor; +} + +inline Status ChangeFormatFromOnnx(ge::Operator& op, const int idx, ge::Format format, bool is_input) { + if (is_input) { + ge::TensorDesc org_tensor = op.GetInputDesc(idx); + org_tensor.SetOriginFormat(format); + org_tensor.SetFormat(format); + auto ret = op.UpdateInputDesc(idx, org_tensor); + if (ret != ge::GRAPH_SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "change input format failed."); + return FAILED; + } + } else { + ge::TensorDesc org_tensor_y = op.GetOutputDesc(idx); + org_tensor_y.SetOriginFormat(format); + org_tensor_y.SetFormat(format); + auto ret_y = op.UpdateOutputDesc(idx, org_tensor_y); + if (ret_y != ge::GRAPH_SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "change output format failed."); + return FAILED; + } + } + return SUCCESS; +} +} // namespace domi + +#endif // MATH_COMMON_ONNX_COMMON_H \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/common.h b/csrc/ascend/common/include/kernel/common.h new file mode 100644 index 000000000..0b3af9296 --- /dev/null +++ b/csrc/ascend/common/include/kernel/common.h @@ -0,0 +1,29 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file common.h + * \brief + */ + +#ifndef INCLUDE_COMMON_H +#define INCLUDE_COMMON_H + +#define CONST_2 2 + +#define SET_FLAG(trigger, waiter, e) AscendC::SetFlag((e)) +#define WAIT_FLAG(trigger, waiter, e) AscendC::WaitFlag((e)) +#define PIPE_BARRIER(pipe) AscendC::PipeBarrier() + +#ifndef FORCE_INLINE +#define FORCE_INLINE inline __attribute__((always_inline)) +#endif + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/common_func.h b/csrc/ascend/common/include/kernel/common_func.h new file mode 100644 index 000000000..4ff536a4a --- /dev/null +++ b/csrc/ascend/common/include/kernel/common_func.h @@ -0,0 +1,117 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file common_func.h + * \brief + */ + + #ifndef INCLUDE_COMMON_FUNC_H + #define INCLUDE_COMMON_FUNC_H + + #include + #include + + #ifdef __CCE_KT_TEST__ + #include "stub_def.h" + #include "stub_fun.h" + #else + #include "kernel_macros.h" + #endif + + template + inline __aicore__ T RoundUp(const T val) + { + static_assert(ALIGN != 0, "align must not be zero"); + static_assert(std::is_arithmetic::value, "T must be an arithmetic type"); + T align = ALIGN; + if (val + align - 1 < val) { + return val; + } + return (val + align - 1) / align * align; + } + + template + inline __aicore__ T RoundUp(const T val, const T align) + { + static_assert(std::is_arithmetic::value, "T must be an arithmetic type"); + if (align == 0 || val + align - 1 < val) { + return val; + } + return (val + align - 1) / align * align; + } + + template + inline __aicore__ T CeilDiv(const T dividend) + { + static_assert(DIVISOR != 0, "align must not be zero"); + static_assert(std::is_arithmetic::value, "T must be an arithmetic type"); + T divisor = DIVISOR; + if (dividend + divisor - 1 < dividend) { + return dividend; + } + return (dividend + divisor - 1) / divisor; + } + + template + constexpr T T_MAX = std::numeric_limits::max(); + + template + inline __aicore__ T CeilDiv(const T dividend, const T divisor) + { + static_assert(std::is_arithmetic::value, "T must be an arithmetic type"); + if (divisor == 0 || dividend + divisor - 1 < dividend) { + return T_MAX; + } + return (dividend + divisor - 1) / divisor; + } + + template + __aicore__ inline T Min(const T lhs, const T rhs) + { + return lhs < rhs ? lhs : rhs; + } + + template __aicore__ __attribute__((always_inline)) inline uint32_t BlockSize() + { + return 32 / sizeof(Dtype); + } + + template __aicore__ __attribute__((always_inline)) inline uint32_t MatrixSize() + { + return 512 / sizeof(Dtype); + } + + template __aicore__ __attribute__((always_inline)) inline uint64_t BlockSizeRoundUp(uint64_t num) + { + return (num + BlockSize() - 1) / BlockSize() * BlockSize(); + } + + template __aicore__ __attribute__((always_inline)) inline uint64_t NumBlocksRoundUp(uint64_t num) + { + return (num + BlockSize() - 1) / BlockSize(); + } + + template __aicore__ __attribute__((always_inline)) inline uint64_t MatrixSizeRoundUp(uint64_t num) + { + return (num + MatrixSize() - 1) / MatrixSize() * MatrixSize(); + } + + template __aicore__ __attribute__((always_inline)) inline uint64_t NumMatrixsRoundUp(uint64_t num) + { + return (num + MatrixSize() - 1) / MatrixSize(); + } + + template __aicore__ __attribute__((always_inline)) inline uint64_t L0HalfSize() + { + return 32 * 1024 / sizeof(Dtype); + } + + #endif \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/dropmask.h b/csrc/ascend/common/include/kernel/dropmask.h new file mode 100644 index 000000000..b57b0126d --- /dev/null +++ b/csrc/ascend/common/include/kernel/dropmask.h @@ -0,0 +1,121 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dropmask.h + * \brief + */ + +#ifndef DROPMASK_H +#define DROPMASK_H + +#include "util.h" + +using AscendC::DROPOUT_MODE_BIT_MISALIGN; +using AscendC::DropOutShapeInfo; +using AscendC::DropOut; + +struct DropMaskInfo { + // for compute dropout mask offset + // 参数按B N G S1 S2全部切分设置进行偏移计算，没有切分的轴对应的参数设置为合适的0或者原始值 + int64_t n2G; // n2 * g + int64_t gSize; // g + int64_t s1Size; // s1 + int64_t s2Size; // s2 + int64_t gOutIdx; // g out index + int64_t bSSOffset; // boidx * s1 * s2 ===bSSOffset + int64_t n2OutIdx; // n out index + int64_t s1OutIdx; // s1 out index ===s1oIdx + int64_t s1InnerIdx; // s1 inner index, 配比 ===loopIdx + int64_t s1BaseSize; // S1基本块大小 + int64_t splitS1BaseSize; // s1 split size ===vec1S1BaseSize + int64_t s2StartIdx; // s2 start index + int64_t s2Idx; // s2 index =====s2LoopCount + int64_t s2BaseNratioSize; // s2的配比长度: s2BaseSize(S2基本块大小) * nRatio + + // for copy in dropout mask + uint32_t s1CopySize; + uint32_t s2CopySize; + int64_t s2TotalSize; + + // for compute dropout mask + uint32_t firstAxis; + uint32_t lstAxis; + uint32_t maskLstAxis; + int64_t vecCoreOffset = 0; + float keepProb; + + bool boolMode; +}; + +template +__aicore__ inline int64_t ComputeDropOffset(DropMaskInfo &dropMaskInfo) +{ + if constexpr (hasDrop == true) { + // boidx * n2 * g* s1 * s2 + int64_t bOffset = dropMaskInfo.bSSOffset * dropMaskInfo.n2G; + // n2oIdx * g * s1 *s2 + int64_t n2Offset = dropMaskInfo.n2OutIdx * dropMaskInfo.gSize * dropMaskInfo.s1Size * dropMaskInfo.s2Size; + // goIdx * s1 * s2 + int64_t gOffset = dropMaskInfo.gOutIdx * dropMaskInfo.s1Size * dropMaskInfo.s2Size; + // s1oIdx * s1BaseSize * s2Size + s1innerindex * vec1S1BaseSize * s2Size + int64_t s1Offset = (dropMaskInfo.s1OutIdx * dropMaskInfo.s1BaseSize + dropMaskInfo.vecCoreOffset + + dropMaskInfo.s1InnerIdx * dropMaskInfo.splitS1BaseSize) * dropMaskInfo.s2Size; + // s2StartIdx + s2index * s2BaseNratioSize + int64_t s2Offset = dropMaskInfo.s2StartIdx + dropMaskInfo.s2Idx * dropMaskInfo.s2BaseNratioSize; + return bOffset + n2Offset + gOffset + s1Offset + s2Offset; + } else { + return 0; + } +} + +template +__aicore__ inline void CopyInDropMask(LocalTensor&dstTensor, GlobalTensor& srcBoolTensor, + GlobalTensor& srcByteTensor, DropMaskInfo &dropMaskInfo, int64_t alignedSize = blockBytes) +{ + if constexpr (hasDrop == true) { + int64_t dropMaskOffset = ComputeDropOffset(dropMaskInfo); + if (unlikely(dropMaskInfo.boolMode)) { + BoolCopyIn(dstTensor, srcBoolTensor, dropMaskOffset, + dropMaskInfo.s1CopySize, dropMaskInfo.s2CopySize, dropMaskInfo.s2TotalSize, alignedSize); + } else { + Bit2Int8CopyIn(dstTensor, srcByteTensor, dropMaskOffset, 1, + dropMaskInfo.s1CopySize, dropMaskInfo.s2CopySize, dropMaskInfo.s2TotalSize, alignedSize); + } + return; + } +} + +template +__aicore__ inline void ComputeDropMask(LocalTensor& dstTensor, LocalTensor& srcTensor, + LocalTensor& dropoutBuffer, LocalTensor& tmpDropBuffer, DropMaskInfo &dropMaskInfo) +{ + if constexpr (hasDrop == true) { + DropOutShapeInfo dropOutShapeInfo; + dropOutShapeInfo.firstAxis = dropMaskInfo.firstAxis; + dropOutShapeInfo.srcLastAxis = dropMaskInfo.lstAxis; + + if (unlikely(dropMaskInfo.boolMode)) { + dropOutShapeInfo.maskLastAxis = CeilDiv(dropMaskInfo.maskLstAxis, blockBytes) * blockBytes; + DropOut(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, dropMaskInfo.keepProb, dropOutShapeInfo); + } else { + dropOutShapeInfo.maskLastAxis = CeilDiv(dropMaskInfo.maskLstAxis / byteBitRatio, blockBytes) * blockBytes; + if (likely(dropMaskInfo.lstAxis / byteBitRatio % blockBytes == 0)) { + DropOut(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, dropMaskInfo.keepProb, dropOutShapeInfo); + } else { + DropOut(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, + dropMaskInfo.keepProb, dropOutShapeInfo); + } + } + return; + } +} + +#endif // DROPMASK_H diff --git a/csrc/ascend/common/include/kernel/gm_to_l1_iterator.h b/csrc/ascend/common/include/kernel/gm_to_l1_iterator.h new file mode 100644 index 000000000..401a7a93d --- /dev/null +++ b/csrc/ascend/common/include/kernel/gm_to_l1_iterator.h @@ -0,0 +1,169 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file gm_to_l1_iterator.h + * \brief + */ + +#ifndef GM_TO_L1_ITERATOR_H +#define GM_TO_L1_ITERATOR_H + +#include "iterator.h" + +constexpr uint32_t STRIDE_LIMIT_H = 65536; + +// Partial specialization for V220, ND_in, ND_out +template +struct gm_to_l1 { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + + __aicore__ gm_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::GlobalTensor gmTensor, + uint32_t nTileActual, + uint32_t nTileCeil, + uint32_t nVal, + uint32_t dTileActual, + uint32_t dTileCeil, + uint32_t dVal) + { + AscendC::DataCopy(l1Tensor, + gmTensor, + AscendC::DataCopyParams(1, // nBurst + CeilDiv(nTileActual * dTileActual), // lenBurst + 0, // srcGap + 0)); // dstGap + }; +}; + +// Partial specialization for NZ_in, NZ_out +template +struct gm_to_l1 { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + + __aicore__ gm_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::GlobalTensor gmTensor, + uint32_t nTileActual, + uint32_t nTileCeil, + uint32_t nVal, + uint32_t dTileActual, + uint32_t dTileCeil, + uint32_t dVal) + { + uint64_t srcStride = nTileCeil - nTileActual; + if (srcStride < STRIDE_LIMIT_H) { + AscendC::DataCopy(l1Tensor, gmTensor, + AscendC::DataCopyParams(dTileActual / BLOCK_SIZE, // nBurst + nTileActual, // lenBurst + nTileCeil - nTileActual, // srcGap + 0)); // dstGap + } else { + for (uint64_t i = 0; i < dTileActual / BLOCK_SIZE; i++) { + uint64_t dstOffset = i * nTileActual * BLOCK_SIZE; + uint64_t srcOffset = i * nTileCeil * BLOCK_SIZE; + AscendC::DataCopy(l1Tensor[dstOffset], gmTensor[srcOffset], + AscendC::DataCopyParams(1, // nBurst + nTileActual, // lenBurst + 0, // srcGap + 0)); // dstGap + } + } + }; +}; + +// Partial specialization for V220, ND_in, ND_out +template +struct gm_to_l1 { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + + __aicore__ gm_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::GlobalTensor gmTensor, + uint32_t nTileActual, + uint32_t nTileCeil, + uint32_t nVal, + uint32_t dTileActual, + uint32_t dTileCeil, + uint32_t dVal) + { + if (dVal < STRIDE_LIMIT_H) { + AscendC::DataCopy(l1Tensor, + gmTensor, + AscendC::Nd2NzParams(1, // ndNum + nTileActual, // nValue + dTileActual, // dValue + 0, // srcNdMatrixStride, unused + dVal, // srcDValue + nTileCeil, // dstNzC0Stride + 1, // dstNzNStride + 0)); // dstNzMatrixStride, unused + } else { + for (uint32_t i = 0; i < nTileActual; i++) { + AscendC::DataCopy(l1Tensor[i * BLOCK_SIZE], + gmTensor[i * dVal], + AscendC::Nd2NzParams(1, // ndNum + 1, // nValue + dTileActual, // dValue + 0, // srcNdMatrixStride, unused + 0, // srcDValue + nTileCeil, // dstNzC0Stride + 0, // dstNzNStride + 0)); // dstNzMatrixStride, unused + } + } + }; +}; + +// Partial specialization for V220, ND_in, NZ_out +template +struct gm_to_l1 { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + + __aicore__ gm_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::GlobalTensor gmTensor, + uint32_t nTileActual, + uint32_t nTileCeil, + uint32_t nVal, + uint32_t dTileActual, + uint32_t dTileCeil, + uint32_t dVal) + { + if (dVal < STRIDE_LIMIT_H) { + AscendC::DataCopy(l1Tensor, + gmTensor, + AscendC::Nd2NzParams(1, // ndNum + nTileActual, // nValue + dTileActual, // dValue + 0, // srcNdMatrixStride, unused + dVal, // srcDValue + nTileCeil, // dstNzC0Stride + 1, // dstNzNStride + 0)); // dstNzMatrixStride, unused + } else { + for (uint32_t i = 0; i < nTileActual; ++i) { + AscendC::DataCopy(l1Tensor, + gmTensor, + AscendC::Nd2NzParams(1, // ndNum + 1, // nValue + dTileActual, // dValue + 0, // srcNdMatrixStride, unused + 0, // srcDValue + nTileCeil, // dstNzC0Stride + 0, // dstNzNStride + 0)); // dstNzMatrixStride, unused + } + } + }; +}; + +#endif // GM_TO_L1_ITERATOR_H \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/gm_to_ub_iterator.h b/csrc/ascend/common/include/kernel/gm_to_ub_iterator.h new file mode 100644 index 000000000..a906bc005 --- /dev/null +++ b/csrc/ascend/common/include/kernel/gm_to_ub_iterator.h @@ -0,0 +1,97 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file gm_to_ub_iterator.h + * \brief + */ + +#ifndef GM_TO_UB_ITERATOR_H +#define GM_TO_UB_ITERATOR_H + +#include "iterator.h" + +constexpr uint32_t STRIDE_LIMIT_I = 65536; + +template struct gm_to_ub { + __aicore__ inline gm_to_ub(AscendC::LocalTensor dstTensor, AscendC::GlobalTensor srcTensor, + uint8_t sid, uint16_t nBurst, uint16_t lenBurst, uint16_t srcStride, uint16_t dstStride) + { + AscendC::DataCopy(dstTensor, srcTensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride)); + }; +}; + +template struct gm_to_ub_align { + __aicore__ inline gm_to_ub_align(AscendC::LocalTensor dstTensor, AscendC::GlobalTensor srcTensor, + uint8_t sid, uint16_t nBurst, uint32_t lenBurst, uint8_t leftPaddingNum, + uint8_t rightPaddingNum, uint32_t srcGap, uint32_t dstGap) + { + AscendC::DataCopyPad(dstTensor, srcTensor, AscendC::DataCopyExtParams(nBurst, lenBurst, srcGap, dstGap, 0), + AscendC::DataCopyPadExtParams(false, leftPaddingNum, rightPaddingNum, 0)); + }; +}; + +template struct ub_to_ub { + __aicore__ inline ub_to_ub(AscendC::LocalTensor dstTensor, AscendC::LocalTensor srcTensor, + uint8_t sid, uint16_t nBurst, uint16_t lenBurst, uint16_t srcStride, uint16_t dstStride) + { + AscendC::DataCopy(dstTensor, srcTensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride)); + }; +}; + +template +struct ub_to_gm { + __aicore__ inline ub_to_gm(AscendC::GlobalTensor dstTensor, AscendC::LocalTensor srcTensor, + uint8_t sid, uint16_t nBurst, uint16_t lenBurst, uint16_t srcStride, uint16_t dstStride) + { + AscendC::DataCopy(dstTensor, srcTensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride)); + }; +}; + +template struct ub_to_gm { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + + __aicore__ ub_to_gm(AscendC::GlobalTensor gmTensor, AscendC::LocalTensor l1Tensor, + uint32_t nTileActual, uint32_t nTileCeil, uint32_t nVal, uint32_t dTileActual, + uint32_t dTileCeil, uint32_t dVal) + { + uint64_t dstStride = nTileCeil - nTileActual; + if (dstStride < STRIDE_LIMIT_I) { + AscendC::DataCopy(gmTensor, l1Tensor, + AscendC::DataCopyParams(dTileActual / BLOCK_SIZE, // nBurst + nTileActual, // lenBurst + 0, // srcGap + dstStride)); // dstGap + } else { + for (uint64_t i = 0; i < dTileActual / BLOCK_SIZE; i++) { + uint64_t srcOffset = i * nTileActual * BLOCK_SIZE; + uint64_t dstOffset = i * nTileCeil * BLOCK_SIZE; + AscendC::DataCopy(gmTensor[dstOffset], l1Tensor[srcOffset], + AscendC::DataCopyParams(1, // nBurst + nTileActual, // lenBurst + 0, // srcGap + 0)); // dstGap + } + } + }; +}; + +template struct ub_to_gm_align { + __aicore__ inline ub_to_gm_align(AscendC::GlobalTensor dstTensor, AscendC::LocalTensor srcTensor, + uint8_t sid, uint16_t nBurst, uint32_t lenBurst, uint8_t leftPaddingNum, + uint8_t rightPaddingNum, uint32_t srcGap, uint32_t dstGap) + { + AscendC::DataCopyPad(dstTensor, srcTensor, AscendC::DataCopyExtParams(nBurst, lenBurst, srcGap, dstGap, 0)); + }; +}; + +#endif // GM_TO_UB_ITERATOR_H \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/hardware.h b/csrc/ascend/common/include/kernel/hardware.h new file mode 100644 index 000000000..e4c91c705 --- /dev/null +++ b/csrc/ascend/common/include/kernel/hardware.h @@ -0,0 +1,40 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hardware.h + * \brief + */ + +#ifndef INCLUDE_HARDWARE_H +#define INCLUDE_HARDWARE_H + +enum class ArchType { ASCEND_V220, ASCEND_V200, ASCEND_M200 }; + +template +struct HardwareInfo { + static uint32_t const l2BW = 5; + static uint32_t const hbmBW = 1; + static uint32_t const supportMix = 0; + static uint32_t const l1Size = 512 * 1024; + static uint32_t const l0ASize = 64 * 1024; + static uint32_t const l0BSize = 64 * 1024; + static uint32_t const l0CSize = 128 * 1024; + static uint32_t const l2Size = 192 * 1024 * 1024; + static uint32_t const biasSize = 1024; + static uint32_t const fixBufSize = 7 * 1024; + static uint32_t const ubSize = 192 * 1024; + static uint32_t const fractalSize = 512; + static uint32_t const l1l0BlockSize = 32; + static uint32_t const btBlockSize = 64; + static uint32_t const fbBlockSize = 128; +}; + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/iterator.h b/csrc/ascend/common/include/kernel/iterator.h new file mode 100644 index 000000000..95180180f --- /dev/null +++ b/csrc/ascend/common/include/kernel/iterator.h @@ -0,0 +1,123 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file iterator.h + * \brief + */ + +#ifndef INCLUDE_ITERTOR_H +#define INCLUDE_ITERTOR_H + +#include "common_func.h" +#include "hardware.h" +#include "kernel_operator.h" +#include "layout.h" +#include "mem.h" + +///////////////////////////////////////////////////// +// gm_to_l1 +///////////////////////////////////////////////////// +template +struct gm_to_l1 { + __aicore__ gm_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::GlobalTensor gmTensor, + uint32_t nTileActual, + uint32_t nTileCeil, + uint32_t nVal, + uint32_t dTileActual, + uint32_t dTileCeil, + uint32_t dVal) {}; +}; + +///////////////////////////////////////////////////// +// l1_to_l0_a +///////////////////////////////////////////////////// +template +struct l1_to_l0_a { + __aicore__ l1_to_l0_a(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t mTileCeil, + uint32_t kPartCeil, + uint32_t mSrcStride, + uint32_t kSrcStride, + uint32_t mDstStride, + uint32_t kDstStride) {}; +}; + +///////////////////////////////////////////////////// +// l1_to_l0_b +///////////////////////////////////////////////////// +template +struct l1_to_l0_b { + __aicore__ l1_to_l0_b(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t nTileCeil, + uint32_t kPartCeil, + uint32_t nSrcStride, + uint32_t kSrcStride, + uint32_t nDstStride, + uint32_t kDstStride) {}; +}; + +// l1_to_l0_a +///////////////////////////////////////////////////// +template +struct l1_to_l0_a_v1 { + __aicore__ l1_to_l0_a_v1(AscendC::LocalTensor l0_tensor, + AscendC::LocalTensor l1_tensor, + uint32_t m_tile_ceil, + uint32_t k_tile_ceil, + uint32_t k_part, + uint32_t k_part_ceil, + uint32_t k_part_idx) {}; +}; + +///////////////////////////////////////////////////// +// l1_to_l0_b +///////////////////////////////////////////////////// +template +struct l1_to_l0_b_v1 { + __aicore__ l1_to_l0_b_v1(AscendC::LocalTensor l0_tensor, + AscendC::LocalTensor l1_tensor, + int32_t n_tile_ceil, + int32_t k_tile_ceil, + int32_t k_part_ceil, + int32_t k_part_idx) {}; +}; + +///////////////////////////////////////////////////// +// l0c_to_gm +///////////////////////////////////////////////////// +template +struct l0c_to_gm { + __aicore__ l0c_to_gm(AscendC::GlobalTensor gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t mTileCeil, + uint32_t nActual) {}; +}; + +///////////////////////////////////////////////////// +// l0c_to_l1 +///////////////////////////////////////////////////// +template +struct l0c_to_l1 { + __aicore__ l0c_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::LocalTensor l0cTensor, + AscendC::LocalTensor deqTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t mTileCeil, + uint32_t nActual) {}; +}; + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/l0c_to_gm_iterator.h b/csrc/ascend/common/include/kernel/l0c_to_gm_iterator.h new file mode 100644 index 000000000..9770dd3ac --- /dev/null +++ b/csrc/ascend/common/include/kernel/l0c_to_gm_iterator.h @@ -0,0 +1,213 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l0c_to_gm_iterator.h + * \brief + */ + +#ifndef L0C_TO_GM_ITERATOR_H +#define L0C_TO_GM_ITERATOR_H + +#ifdef __CCE_KT_TEST__ +#define __bf16 bfloat16_t +#endif + +#include "iterator.h" +constexpr uint32_t BLOCK_NUM = 16; +constexpr uint32_t BLOCK_SIZE_INT8 = 32; + +template <> +struct l0c_to_gm { + /** + * @brief Copy data from L0C buffer to global memory, partial specialized for + * + * @param gmTensor the destination tensor on global memory, which is stored in ND format. + * @param l0cTensor the source tensor on L0C buffer, which is stored in FRACTAL_NZ format. + * @param mTileActual the m-direction size of the matrix in L0C buffer. + * @param nTileActual the n-direction size of the matrix in L0C buffer. + * @param srcStride the source stride between the adjacent fractal matrices along n-direction in unit of C0_SIZE. + * @param dstStride the leading dimension of the destination matrix in unit of element. + */ + __aicore__ l0c_to_gm(AscendC::GlobalTensor gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t srcStride, + uint32_t dstStride) + { +#ifdef __DAV_C220_CUBE__ + auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize + mTileActual, // mSize + srcStride, // srcStride + dstStride, // dstStride + false); // enRelu + + intriParams.quantPre = QuantMode_t::F322F16; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#else + AscendC::FixpipeParams intriParams( + (nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE, + static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8), + 0, + dstStride); + intriParams.nz2ndParams = {true, 1, 0, 0, static_cast(nTileActual)}; + intriParams.quantParams = {QuantMode_t::F322F16}; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#endif + }; +}; + +template <> +struct l0c_to_gm { + __aicore__ l0c_to_gm(AscendC::GlobalTensor gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t srcStride, + uint32_t dstStride) + { +#ifdef __DAV_C220_CUBE__ + auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize + mTileActual, // mSize + srcStride, // srcStride + dstStride, // dstStride + false); // enRelu + + intriParams.quantPre = QuantMode_t::VDEQF16; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#else + AscendC::FixpipeParams intriParams( + (nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE, + static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8), + 0, + dstStride); + intriParams.nz2ndParams = {true, 1, 0, 0, static_cast(nTileActual)}; + intriParams.quantParams = {QuantMode_t::VDEQF16}; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#endif + }; +}; + +template <> +struct l0c_to_gm { + __aicore__ l0c_to_gm(AscendC::GlobalTensor<__bf16> gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t srcStride, + uint32_t dstStride) + { +#ifdef __DAV_C220_CUBE__ + auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize + mTileActual, // mSize + srcStride, // srcStride + dstStride, // dstStride + false); // enRelu + + intriParams.quantPre = QuantMode_t::F322BF16; + AscendC::Fixpipe<__bf16, float, AscendC::CFG_ROW_MAJOR>(gmTensor, l0cTensor, intriParams); +#else + AscendC::FixpipeParams intriParams( + (nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE, + static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8), + 0, + dstStride); + intriParams.nz2ndParams = {true, 1, 0, 0, static_cast(nTileActual)}; + intriParams.quantParams = {QuantMode_t::F322BF16}; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#endif + }; +}; + +// Partial specialization ND, float +template <> +struct l0c_to_gm { + __aicore__ l0c_to_gm(AscendC::GlobalTensor gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t srcStride, + uint32_t dstStride) + { +#ifdef __DAV_C220_CUBE__ + auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize + mTileActual, // mSize + srcStride, // srcStride + dstStride, // dstStride + false); // enRelu + + intriParams.quantPre = QuantMode_t::NoQuant; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#else + AscendC::FixpipeParams intriParams( + (nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE, + static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8), + 0, + dstStride); + intriParams.nz2ndParams = {true, 1, 0, 0, static_cast(nTileActual)}; + intriParams.quantParams = {QuantMode_t::NoQuant}; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#endif + }; +}; + +template <> +struct l0c_to_gm { + __aicore__ l0c_to_gm(AscendC::GlobalTensor gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t srcStride, + uint32_t dstStride) + { +#ifdef __DAV_C220_CUBE__ + auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize + mTileActual, // mSize + srcStride, // srcStride + dstStride, // dstStride + false); // enRelu + + intriParams.quantPre = QuantMode_t::F322F16; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#else + AscendC::FixpipeParams intriParams( + (nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE, + static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8), + 0, + dstStride - (nTileActual * sizeof(half) / sizeof(float))); + intriParams.quantParams = {QuantMode_t::F322F16}; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#endif + }; +}; + +template <> +struct l0c_to_gm { + __aicore__ l0c_to_gm(AscendC::GlobalTensor gmTensor, + AscendC::LocalTensor l0cTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t srcStride, + uint32_t dstStride){ +#ifdef __DAV_C220_CUBE__ + auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize + mTileActual, // mSize + srcStride, // srcStride + dstStride, // dstStride + false); // enRelu + + intriParams.quantPre = QuantMode_t::NoQuant; + AscendC::Fixpipe(gmTensor, l0cTensor, intriParams); +#endif +}; +}; + +#endif // L0C_TO_GM_ITERATOR_H diff --git a/csrc/ascend/common/include/kernel/l0c_to_l1_iterator.h b/csrc/ascend/common/include/kernel/l0c_to_l1_iterator.h new file mode 100644 index 000000000..9f3fc5c38 --- /dev/null +++ b/csrc/ascend/common/include/kernel/l0c_to_l1_iterator.h @@ -0,0 +1,51 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l0c_to_l1_iterator.h + * \brief + */ + +#ifndef L0C_TO_L1_ITERATOR_H +#define L0C_TO_L1_ITERATOR_H + +#include "iterator.h" +///////////////////////////////////////////////////// +// l0c_to_l1 +///////////////////////////////////////////////////// + +// Partial specialization ZN, half, int32_t +template +struct l0c_to_l1 { + using ElementOut = half; + using ElementIn = int32_t; + __aicore__ l0c_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::LocalTensor l0cTensor, + AscendC::LocalTensor deqTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t mTileCeil, + uint32_t nActual) + { + constexpr uint32_t BLOCK_NUM = 16; + constexpr uint32_t BLOCK_SIZE = 32; + AscendC::FixpipeParams intriParams( + (nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE, + static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE), + 0, + mTileCeil - static_cast(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE) * + sizeof(ElementOut) / sizeof(ElementIn)); + intriParams.nz2ndParams = {false, 1, 0, 0, static_cast(nTileActual)}; + intriParams.quantParams = {QuantMode_t::VDEQF16}; + AscendC::Fixpipe(l1Tensor, l0cTensor, deqTensor, intriParams); + }; +}; + +#endif // L0C_TO_L1_ITERATOR_H \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/l0c_to_ub_iterator.h b/csrc/ascend/common/include/kernel/l0c_to_ub_iterator.h new file mode 100644 index 000000000..5d2ea8e71 --- /dev/null +++ b/csrc/ascend/common/include/kernel/l0c_to_ub_iterator.h @@ -0,0 +1,72 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l0c_to_ub_iterator.h + * \brief + */ +#ifndef L0C_TO_UB_ITERATOR_H +#define L0C_TO_UB_ITERATOR_H + +#include "iterator.h" + +///////////////////////////////////////////////////// +// l0c_to_ub +///////////////////////////////////////////////////// + +// Partial specialization ZN, half, int32_t +template struct l0c_to_ub { + __aicore__ l0c_to_ub(AscendC::LocalTensor ubTensor, AscendC::LocalTensor l0cTensor, + uint16_t nBurst, uint16_t lenBurst, uint16_t srcStride, uint16_t dstStride) + { + constexpr auto mode = + MatrixMode ? AscendC::BlockMode::BLOCK_MODE_MATRIX : AscendC::BlockMode::BLOCK_MODE_VECTOR; + AscendC::DataCopy(ubTensor, l0cTensor, + AscendC::DataCopyParams(nBurst, // count + lenBurst, // len + srcStride, // srcStrideIn + dstStride), // dstStrideIn + AscendC::DataCopyEnhancedParams(mode, // blockModeIn + AscendC::DeqScale::DEQ_NONE, // deqScaleIn + 0, // deqValueIn + 0, // sidStoreModeIn + false, // isReluIn + pad_t::PAD_NONE, // padModeIn + 0) // padValueIn + ); + }; +}; + +template +struct l0c_to_ub { + __aicore__ l0c_to_ub(AscendC::LocalTensor ubTensor, + AscendC::LocalTensor l0cTensor, + uint16_t nBurst, + uint16_t lenBurst, + uint16_t srcStride, + uint16_t dstStride) + { + AscendC::DataCopy(ubTensor, l0cTensor, + AscendC::DataCopyParams(nBurst, // count + lenBurst, // len + srcStride, // srcStrideIn + dstStride), // dstStrideIn + AscendC::DataCopyEnhancedParams(AscendC::BlockMode::BLOCK_MODE_MATRIX, // blockModeIn + AscendC::DeqScale::VDEQ16, // deqScaleIn + 0, // deqValueIn + 0, // sidStoreModeIn + false, // isReluIn + pad_t::PAD_NONE, // padModeIn + 0) // padValueIn + ); + }; +}; + +#endif // L0C_TO_UB_ITERATOR_H \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/l1_to_bt_iterator.h b/csrc/ascend/common/include/kernel/l1_to_bt_iterator.h new file mode 100644 index 000000000..a2db67a33 --- /dev/null +++ b/csrc/ascend/common/include/kernel/l1_to_bt_iterator.h @@ -0,0 +1,39 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l1_to_bt_iterator.h + * \brief + */ +#ifndef L1_TO_BT_ITERATOR_H +#define L1_TO_BT_ITERATOR_H + +#include "iterator.h" + +///////////////////////////////////////////////////// +// l1_to_bt +///////////////////////////////////////////////////// + +// Partial specialization for V220 +template +struct l1_to_bt { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::btBlockSize / sizeof(DataType); + + __aicore__ l1_to_bt(AscendC::LocalTensor biasTableTensor, + AscendC::LocalTensor biasL1Tensor, + uint32_t ntileActual) + { + AscendC::DataCopy( + biasTableTensor, biasL1Tensor, {1, static_cast(CeilDiv(ntileActual)), 0, 0}); + }; +}; + +#endif // L1_TO_BT_ITERATOR_H \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/l1_to_fb_iterator.h b/csrc/ascend/common/include/kernel/l1_to_fb_iterator.h new file mode 100644 index 000000000..8f75074ac --- /dev/null +++ b/csrc/ascend/common/include/kernel/l1_to_fb_iterator.h @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l1_to_fb_iterator.h + * \brief + */ + +#ifndef L1_TO_FB_ITERATOR_H +#define L1_TO_FB_ITERATOR_H + +#include "iterator.h" + +///////////////////////////////////////////////////// +// l1_to_fb +///////////////////////////////////////////////////// + +// Partial specialization for V220 +template +struct l1_to_fb { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::fbBlockSize / sizeof(DataType); + + __aicore__ + l1_to_fb(AscendC::LocalTensor fbTensor, AscendC::LocalTensor l1Tensor, uint32_t ntileActual) + { + copy_cbuf_to_fbuf((__fbuf__ DataType *)fbTensor.GetPhyAddr(), + (__cbuf__ DataType *)l1Tensor.GetPhyAddr(), + 1, + CeilDiv(ntileActual), + 0, + 0); + }; +}; + +#endif // L1_TO_FB_ITERATOR_H \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/l1_to_l0_iterator.h b/csrc/ascend/common/include/kernel/l1_to_l0_iterator.h new file mode 100644 index 000000000..72e58ddc6 --- /dev/null +++ b/csrc/ascend/common/include/kernel/l1_to_l0_iterator.h @@ -0,0 +1,259 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l1_to_l0_iterator.h + * \brief + */ + +#ifndef L1_TO_L0_ITERATOR_H +#define L1_TO_L0_ITERATOR_H + +#include "iterator.h" + +///////////////////////////////////////////////////// +// l1_to_l0_a +///////////////////////////////////////////////////// + +// Partial specialization for vector +template +struct l1_to_l0_a { + using HardwareParams = HardwareInfo; + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + + __aicore__ l1_to_l0_a(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t mTileCeil, + uint32_t kPartCeil, + uint32_t mSrcStride, + uint32_t kSrcStride, + uint32_t mDstStride, + uint32_t kDstStride) + { + AscendC::LoadData(l0Tensor, + l1Tensor, + AscendC::LoadData2dParams(0, // baseIdx + kPartCeil, // repeat + kSrcStride, // srcStride + 0, // sid + kDstStride, // dstStride + IsTransPose, // transpose + 0)); // addrCalMode + }; +}; + +// Partial specialization for no transpose, not vector +template +struct l1_to_l0_a { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize; + + __aicore__ l1_to_l0_a(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t mTileCeil, + uint32_t kPartCeil, + uint32_t mSrcStride, + uint32_t kSrcStride, + uint32_t mDstStride, + uint32_t kDstStride) + { + for (uint32_t i = 0; i < mTileCeil / BLOCK_NUM_PER_FRACTAL; i++) { + AscendC::LoadData(l0Tensor[i * mDstStride * FRACTAL_SIZE], + l1Tensor[i * mSrcStride * FRACTAL_SIZE], + AscendC::LoadData2dParams(0, // baseIdx + static_cast(kPartCeil / BLOCK_SIZE), // repeat + kSrcStride, // srcStride + 0, // sid + kDstStride - 1, // dstStride + false, // transpose + 0)); // addrCalMode + } + }; +}; + +// Partial specialization for transpose, not vector +template +struct l1_to_l0_a { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize; + + __aicore__ l1_to_l0_a(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t mTileCeil, + uint32_t kPartCeil, + uint32_t mSrcStride, + uint32_t kSrcStride, + uint32_t mDstStride, + uint32_t kDstStride) + { + for (uint32_t i = 0; i < mTileCeil / BLOCK_SIZE; i++) { + AscendC::LoadData(l0Tensor[i * mDstStride * FRACTAL_SIZE], + l1Tensor[i * mSrcStride * FRACTAL_SIZE], + AscendC::LoadData2dParams(0, + static_cast(kPartCeil / BLOCK_NUM_PER_FRACTAL), + kSrcStride, + 0, + kDstStride - 1, + true, + 0)); + } + }; +}; + +template +struct l1_to_l0_a { + using HardwareParams = HardwareInfo; + // 16 * 32 + static constexpr uint32_t ROW_BLOCK_SIZE = 16; + static constexpr uint32_t COL_BLOCK_SIZE = 32 / sizeof(DataType); + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize; + + __aicore__ l1_to_l0_a(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t mTileCeil, + uint32_t kPartCeil, + uint32_t mSrcStride, + uint32_t kSrcStride, + uint32_t mDstStride, + uint32_t kDstStride) + { + for (uint32_t i = 0; i < mTileCeil / ROW_BLOCK_SIZE; i++) { + AscendC::LoadData(l0Tensor[i * ROW_BLOCK_SIZE * kPartCeil], + l1Tensor[i * FRACTAL_SIZE], + AscendC::LoadData2dParams(0, + static_cast(kPartCeil / COL_BLOCK_SIZE), + mTileCeil / ROW_BLOCK_SIZE, + 0, + 0, + false, + 0)); + } + }; +}; + +///////////////////////////////////////////////////// +// l1_to_l0_b +///////////////////////////////////////////////////// + +// Partial specialization for vector +template +struct l1_to_l0_b { + using HardwareParams = HardwareInfo; + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + + __aicore__ l1_to_l0_b(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t nTileCeil, + uint32_t kPartCeil, + uint32_t nSrcStride, + uint32_t kSrcStride, + uint32_t nDstStride, + uint32_t kDstStride) + { + AscendC::LoadData( + l0Tensor, l1Tensor, AscendC::LoadData2dParams(0, kPartCeil, kSrcStride, 0, kDstStride, IsTransPose, 0)); + }; +}; + +template +struct l1_to_l0_b { + using HardwareParams = HardwareInfo; + using DataType = int8_t; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + + __aicore__ l1_to_l0_b(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t nTileCeil, + uint32_t kPartCeil, + uint32_t nSrcStride, + uint32_t kSrcStride, + uint32_t nDstStride, + uint32_t kDstStride) + { + for (uint32_t i = 0; i < nTileCeil / BLOCK_SIZE; i++) { + AscendC::LoadDataWithTranspose(l0Tensor[i * kPartCeil * BLOCK_SIZE], + l1Tensor[i * BLOCK_SIZE * BLOCK_SIZE], + AscendC::LoadData2dTransposeParams(0, // startIndexIn + kPartCeil / BLOCK_SIZE, // repeatTimesIn + nTileCeil / BLOCK_SIZE, // srcStrideIn + 1, // dstGapIn + 0, // dstfracGapIn + 0) // addrModeIn + ); + } + }; +}; + +// Partial specialization for no transpose, not vector +template +struct l1_to_l0_b { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize; + + __aicore__ l1_to_l0_b(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t nTileCeil, + uint32_t kPartCeil, + uint32_t nSrcStride, + uint32_t kSrcStride, + uint32_t nDstStride, + uint32_t kDstStride) + { + for (uint32_t i = 0; i < kPartCeil / BLOCK_NUM_PER_FRACTAL; i++) { + AscendC::LoadData(l0Tensor[i * kDstStride * FRACTAL_SIZE], + l1Tensor[i * kSrcStride * FRACTAL_SIZE], + AscendC::LoadData2dParams(0, // baseIdx + static_cast(nTileCeil / BLOCK_SIZE), // repeat + nSrcStride, // srcStride + 0, // sid + nDstStride - 1, // dstStride + true, // transpose + 0)); // addrCalMode + } + }; +}; + +// Partial specialization for transpose, not vector +template +struct l1_to_l0_b { + using HardwareParams = HardwareInfo; + static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType); + static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType); + static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize; + __aicore__ l1_to_l0_b(AscendC::LocalTensor l0Tensor, + AscendC::LocalTensor l1Tensor, + uint32_t nTileCeil, + uint32_t kPartCeil, + uint32_t nSrcStride, + uint32_t kSrcStride, + uint32_t nDstStride, + uint32_t kDstStride) + { + AscendC::LoadData( + l0Tensor, + l1Tensor, + AscendC::LoadData2dParams(0, // baseIdx + static_cast(kPartCeil * nTileCeil / FRACTAL_SIZE), // repeat + 1, // srcStride + 0, // sid + 0, // dstStride + false, // transpose + 0)); // addr_cal_mode_t + }; +}; + +#endif // L1_TO_L0_ITERATOR_H \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/l1_to_ub_iterator.h b/csrc/ascend/common/include/kernel/l1_to_ub_iterator.h new file mode 100644 index 000000000..2446d110c --- /dev/null +++ b/csrc/ascend/common/include/kernel/l1_to_ub_iterator.h @@ -0,0 +1,52 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file l1_to_ub_iterator.h + * \brief + */ + +#ifndef L1_TO_UB_ITERATOR_H +#define L1_TO_UB_ITERATOR_H + +#include "iterator.h" + +///////////////////////////////////////////////////// +// l1_to_ub +///////////////////////////////////////////////////// +template +struct l1_to_ub { + __aicore__ l1_to_ub(AscendC::LocalTensor ubTensor, + AscendC::LocalTensor l1Tensor, + uint16_t nBurst, + uint16_t lenBurst, + uint16_t srcStride, + uint16_t dstStride) + { + AscendC::DataCopy(ubTensor, l1Tensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride)); + }; +}; + +///////////////////////////////////////////////////// +// ub_to_l1 +///////////////////////////////////////////////////// +template +struct ub_to_l1 { + __aicore__ ub_to_l1(AscendC::LocalTensor l1Tensor, + AscendC::LocalTensor ubTensor, + uint16_t nBurst, + uint16_t lenBurst, + uint16_t srcStride, + uint16_t dstStride) + { + AscendC::DataCopy(l1Tensor, ubTensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride)); + }; +}; +#endif // L1_TO_UB_ITERATOR_H \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/layout.h b/csrc/ascend/common/include/kernel/layout.h new file mode 100644 index 000000000..dec4c49ae --- /dev/null +++ b/csrc/ascend/common/include/kernel/layout.h @@ -0,0 +1,28 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file layout.h + * \brief + */ + +#ifndef INCLUDE_LAYOUT_H +#define INCLUDE_LAYOUT_H + +enum class DataFormatT { + ND = 0, + NZ, + ZN, + ZZ, + NN, + VECTOR +}; + +#endif // INCLUDE_LAYOUT_H \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/mem.h b/csrc/ascend/common/include/kernel/mem.h new file mode 100644 index 000000000..6b41f2db9 --- /dev/null +++ b/csrc/ascend/common/include/kernel/mem.h @@ -0,0 +1,79 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file mem.h + * \brief + */ + +#ifndef INCLUDE_MEM_H +#define INCLUDE_MEM_H + +#include "hardware.h" +#include "kernel_event.h" +#include "kernel_tensor.h" + +enum class BufferType { ASCEND_UB, ASCEND_CB, ASCEND_L0A, ASCEND_L0B, ASCEND_L0C, ASCEND_MAX }; + +template +__aicore__ constexpr AscendC::TPosition GetPosition() +{ + if constexpr (BufferType_ == BufferType::ASCEND_UB) { + return AscendC::TPosition::VECIN; + } else if constexpr (BufferType_ == BufferType::ASCEND_CB) { + return AscendC::TPosition::A1; + } else if constexpr (BufferType_ == BufferType::ASCEND_L0A) { + return AscendC::TPosition::A2; + } else if constexpr (BufferType_ == BufferType::ASCEND_L0B) { + return AscendC::TPosition::B2; + } else if constexpr (BufferType_ == BufferType::ASCEND_L0C) { + return AscendC::TPosition::CO1; + } + return AscendC::TPosition::GM; +} + +template +struct AsdopsBuffer { +public: + __aicore__ AsdopsBuffer() + { + constexpr uint32_t bufferSize[(uint32_t)BufferType::ASCEND_MAX] = {HardwareInfo::ubSize, + HardwareInfo::l1Size, + HardwareInfo::l0ASize, + HardwareInfo::l0BSize, + HardwareInfo::l0CSize}; +#ifdef __DAV_C220_VEC__ + tensor[(uint32_t)BufferType::ASCEND_UB] = AscendC::LocalTensor(AscendC::TPosition::VECIN, 0, bufferSize[(uint32_t)BufferType::ASCEND_UB]); +#elif __DAV_C220_CUBE__ + tensor[(uint32_t)BufferType::ASCEND_CB] = AscendC::LocalTensor(AscendC::TPosition::A1, 0, bufferSize[(uint32_t)BufferType::ASCEND_CB]); + tensor[(uint32_t)BufferType::ASCEND_L0A] = AscendC::LocalTensor(AscendC::TPosition::A2, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0A]); + tensor[(uint32_t)BufferType::ASCEND_L0B] = AscendC::LocalTensor(AscendC::TPosition::B2, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0B]); + tensor[(uint32_t)BufferType::ASCEND_L0C] = AscendC::LocalTensor(AscendC::TPosition::CO1, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0C]); +#else +#ifndef __clang__ + tensor[(uint32_t)BufferType::ASCEND_UB] = AscendC::LocalTensor(AscendC::TPosition::VECIN, 0, bufferSize[(uint32_t)BufferType::ASCEND_UB]); + tensor[(uint32_t)BufferType::ASCEND_CB] = AscendC::LocalTensor(AscendC::TPosition::A1, 0, bufferSize[(uint32_t)BufferType::ASCEND_CB]); + tensor[(uint32_t)BufferType::ASCEND_L0A] = AscendC::LocalTensor(AscendC::TPosition::A2, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0A]); + tensor[(uint32_t)BufferType::ASCEND_L0B] = AscendC::LocalTensor(AscendC::TPosition::B2, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0B]); + tensor[(uint32_t)BufferType::ASCEND_L0C] = AscendC::LocalTensor(AscendC::TPosition::CO1, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0C]); +#endif +#endif + }; + + template + __aicore__ AscendC::LocalTensor GetBuffer(const uint32_t offset) const + { + return tensor[(uint32_t)BufferType_][offset].template ReinterpretCast(); + } + +public: + AscendC::LocalTensor tensor[(uint32_t)BufferType::ASCEND_MAX]; +}; +#endif \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/mma.h b/csrc/ascend/common/include/kernel/mma.h new file mode 100644 index 000000000..72ab9fd36 --- /dev/null +++ b/csrc/ascend/common/include/kernel/mma.h @@ -0,0 +1,79 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file mma.h + * \brief + */ + +#ifndef INCLUDE_MMA_H +#define INCLUDE_MMA_H + +#include "hardware.h" +#include "kernel_tensor.h" + +template +struct mmad { + __aicore__ mmad(AscendC::LocalTensor l0cTensor, + AscendC::LocalTensor l0aTensor, + AscendC::LocalTensor l0bTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t kPartActual, + bool initC) {}; + + __aicore__ mmad(AscendC::LocalTensor l0cTensor, + AscendC::LocalTensor l0aTensor, + AscendC::LocalTensor l0bTensor, + uint64_t biasBt, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t kPartActual, + bool initC) {}; +}; + +// Partial specialization for V220, int8_t, not_vector_A, not TransposeA +template +struct mmad { + __aicore__ mmad(AscendC::LocalTensor l0cTensor, + AscendC::LocalTensor l0aTensor, + AscendC::LocalTensor l0bTensor, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t kPartActual, + bool initC) + { + AscendC::Mmad(l0cTensor, + l0aTensor, + l0bTensor, + AscendC::MmadParams(mTileActual, nTileActual, kPartActual, 0, false, initC)); + }; + + __aicore__ mmad(AscendC::LocalTensor l0cTensor, + AscendC::LocalTensor l0aTensor, + AscendC::LocalTensor l0bTensor, + uint64_t biasBt, + uint32_t mTileActual, + uint32_t nTileActual, + uint32_t kPartActual, + bool initC) + { + AscendC::LocalTensor biasTensor; + biasTensor.InitBuffer(biasBt, mTileActual); + biasTensor.address_.logicPos = static_cast(AscendC::TPosition::C2); + AscendC::Mmad(l0cTensor, + l0aTensor, + l0bTensor, + biasTensor, + AscendC::MmadParams(mTileActual, nTileActual, kPartActual, 0, false, initC)); + }; +}; + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/pse.h b/csrc/ascend/common/include/kernel/pse.h new file mode 100644 index 000000000..185a7ea38 --- /dev/null +++ b/csrc/ascend/common/include/kernel/pse.h @@ -0,0 +1,483 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file pse.h + * \brief + */ + +#ifndef FLASH_ATTENTION_SCORE_PSE_H +#define FLASH_ATTENTION_SCORE_PSE_H + +#include "kernel_operator.h" +#include "util.h" + +constexpr static int64_t pseS1S2 = 0; +constexpr static int64_t pse1S2 = 1; +constexpr static int64_t pseSlopeBn = 2; +constexpr static int64_t pseSlopeN = 3; + +constexpr static uint8_t pseEncodeALibiS2Full = 0x11; + +enum class PseTypeEnum { + PSE_OUTER_MUL_ADD_TYPE = 0, // default + PSE_OUTER_ADD_MUL_TYPE, + PSE_INNER_MUL_ADD_TYPE, + PSE_INNER_MUL_ADD_SQRT_TYPE, + PSE_INVALID_TYPE +}; + +struct PseInfo { + int64_t blockCount; + int64_t bSSOffset; // boidx * s1 * s2 + int64_t boIdx; + int64_t gSize; + int64_t goIdx; + int64_t loopIdx; + int64_t n2G; + int64_t n2oIdx; + int64_t pseBSize; + int64_t pseS1Size; // for alibi + int64_t pseS2ComputeSize; // for alibi, do not need assignment + int64_t pseS2Size; // for alibi + uint32_t pseShapeType; + int64_t readS2Size; // for alibi, do not need assignment + int64_t s1BaseSize; + int64_t s1Size; + int64_t s1oIdx; + int64_t s2AlignedSize; + int64_t s2BaseNratioSize; + int64_t s2LoopCount; + int64_t s2RealSize; + int64_t s2Size; + int64_t s2SizeAcc; // accumulated sum of s2 size + int64_t s2StartIdx; + int64_t vec1S1BaseSize; + int64_t vec1S1RealSize; + uint32_t pseEncodeType; // for distinguish alibi + uint32_t pseType; // 0: outer, mul-add 1:outer, add-mul 2:inner, mul-add 3:inner, mul-add-sqrt + int64_t pseAlibiBaseS1; + int64_t pseAlibiBaseS2; + int64_t qStartIdx; + int64_t kvStartIdx; + int64_t vecCoreOffset = 0; + bool needCast; + bool align8 = false; + bool pseEndogenous = false; +}; + +template +__aicore__ inline void DataCopyInCommon(LocalTensor &dstTensor, GlobalTensor &srcTensor, int64_t offset, + int64_t s1Size, int64_t s2Size, int64_t actualS2Len, int32_t dtypeSize, + int32_t alignedS2Size) +{ + if constexpr (hasPse == true) { + uint32_t shapeArray[] = {static_cast(s1Size), static_cast(alignedS2Size)}; + dstTensor.SetShapeInfo(ShapeInfo(2, shapeArray, DataFormat::ND)); + dstTensor.SetSize(s1Size * alignedS2Size); + DataCopyParams dataCopyParams; + dataCopyParams.blockCount = s1Size; + dataCopyParams.blockLen = CeilDiv(s2Size * dtypeSize, blockBytes); // 单位32B + dataCopyParams.dstStride = alignedS2Size * dtypeSize / blockBytes - dataCopyParams.blockLen; // gap + if (actualS2Len * dtypeSize % blockBytes == 0) { + dataCopyParams.srcStride = + (actualS2Len * dtypeSize - dataCopyParams.blockLen * blockBytes) / blockBytes; // srcGap + DataCopy(dstTensor, srcTensor[offset], dataCopyParams); + } else { + dataCopyParams.blockLen = s2Size * dtypeSize; // 单位Byte + dataCopyParams.srcStride = (actualS2Len * dtypeSize - dataCopyParams.blockLen); + dataCopyParams.dstStride = (alignedS2Size - s2Size) * dtypeSize / blockBytes; + DataCopyPadParams dataCopyPadParams; + dataCopyPadParams.isPad = false; + DataCopyPad(dstTensor, srcTensor[offset], dataCopyParams, dataCopyPadParams); + } + } +} + +template +__aicore__ inline void DataCopyIn(LocalTensor &dstTensor, GlobalTensor &srcTensor, int64_t offset, + int64_t s1Size, int64_t s2Size, int64_t actualS2Len, int64_t alignedSize = 16) +{ + if constexpr (hasPse == true) { + int32_t dtypeSize = sizeof(INPUT_T); + int32_t alignedS2Size = CeilDiv(s2Size, alignedSize) * alignedSize; + DataCopyInCommon(dstTensor, srcTensor, offset, s1Size, s2Size, + actualS2Len, dtypeSize, alignedS2Size); + } +} + +template +__aicore__ inline void DataCopyInAlign8(LocalTensor &dstTensor, GlobalTensor &srcTensor, int64_t offset, + int64_t s1Size, int64_t s2Size, int64_t actualS2Len) +{ + if constexpr (hasPse == true) { + int32_t dtypeSize = sizeof(INPUT_T); + if (dtypeSize == 0){ + return; + } + int32_t alignedS2Size = CeilDiv(s2Size, 32 / dtypeSize) * (32 / dtypeSize); + DataCopyInCommon(dstTensor, srcTensor, offset, s1Size, s2Size, + actualS2Len, dtypeSize, alignedS2Size); + } +} + +/* +dst = BroadcastAdd(src0, src1) +src0 shape: (s1, s2) +src1 shape: (1, s2) +dst shape: (s1, s2) +*/ +template +__aicore__ inline void BroadcastAdd(const LocalTensor &src0Tensor, const LocalTensor &src1Tensor, + int64_t src0Offset, int32_t src1Size, int32_t repeatTimes) +{ + if constexpr (hasPse == true) { + /* Total data number of single step should be smaller than 256bytes. + * If larger, we need to do add multiple times. */ + int32_t innerLoop = src1Size / repeatMaxSize; // s2轴整块计算次数 + int32_t innerRemain = src1Size % repeatMaxSize; // s2轴尾块计算量 + BinaryRepeatParams binaryRepeatParams; + binaryRepeatParams.src0BlkStride = 1; + binaryRepeatParams.src0RepStride = src1Size / blockSize; + binaryRepeatParams.src1BlkStride = 1; + binaryRepeatParams.src1RepStride = 0; + binaryRepeatParams.dstRepStride = binaryRepeatParams.src0RepStride; + binaryRepeatParams.blockNumber = binaryRepeatParams.src0RepStride; + + for (int32_t j = 0; j < innerLoop; j++) { + auto innerOffset = j * repeatMaxSize; + auto ubOffset = src0Offset + innerOffset; + Add(src0Tensor[ubOffset], src0Tensor[ubOffset], src1Tensor[innerOffset], repeatMaxSize, repeatTimes, + binaryRepeatParams); + } + if (innerRemain > 0) { + auto innerOffset = innerLoop * repeatMaxSize; + auto ubOffset = src0Offset + innerOffset; + Add(src0Tensor[ubOffset], src0Tensor[ubOffset], src1Tensor[innerOffset], innerRemain, repeatTimes, + binaryRepeatParams); + } + } +} + +template +__aicore__ inline void PseBroadcastAdd(int32_t s1Size, int32_t s2Size, int32_t computeSize, const LocalTensor &pseUb, + const LocalTensor &dstTensor, uint32_t pseShapeType) +{ + if constexpr (hasPse == true) { + if (pseShapeType == pseS1S2 || pseShapeType == pseSlopeBn || pseShapeType == pseSlopeN) { + Add(dstTensor, dstTensor, pseUb, computeSize); + } else { + /* Total repeated times should be <= repeatMaxTimes. If larger, + * we need to do multiple inner loops. */ + int32_t s1OuterLoop = s1Size / repeatMaxTimes; + int32_t s1OuterRemain = s1Size % repeatMaxTimes; + for (int32_t s1OuterIdx = 0; s1OuterIdx < s1OuterLoop; s1OuterIdx++) { + int32_t s1OuterOffset = s1OuterIdx * repeatMaxTimes * s2Size; + BroadcastAdd(dstTensor, pseUb, s1OuterOffset, s2Size, repeatMaxTimes); + } + if (s1OuterRemain > 0) { + int32_t s1OuterOffset = s1OuterLoop * repeatMaxTimes * s2Size; + BroadcastAdd(dstTensor, pseUb, s1OuterOffset, s2Size, s1OuterRemain); + } + } + } +} +template __aicore__ inline int64_t PseComputeOffset(PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + int64_t bOffset = 0; + int64_t n2Offset = 0; + int64_t s1Offset = 0; + int64_t s2Offset = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize; + int64_t gOffset = 0; + if (pseInfo.pseShapeType == pseS1S2) { + // b, n2, g, s1, s2 + bOffset = pseInfo.bSSOffset * pseInfo.n2G; + n2Offset = pseInfo.n2oIdx * pseInfo.gSize * pseInfo.s1Size * pseInfo.s2Size; + gOffset = pseInfo.goIdx * pseInfo.s1Size * pseInfo.s2Size; + s1Offset = (pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + pseInfo.loopIdx * pseInfo.vec1S1BaseSize) * pseInfo.s2Size; + } else if (pseInfo.pseShapeType == pse1S2) { + // b, n2, g, 1, s2 + bOffset = pseInfo.s2SizeAcc * pseInfo.n2G; + n2Offset = pseInfo.n2oIdx * pseInfo.gSize * pseInfo.s2Size; + gOffset = pseInfo.goIdx * pseInfo.s2Size; + } + if (pseInfo.pseBSize == 1) { + bOffset = 0; + } + return bOffset + n2Offset + gOffset + s1Offset + s2Offset; + } else { + return 0; + } +} + +template __aicore__ inline int64_t PseAlibiComputeOffset(PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + int64_t bOffset = (pseInfo.boIdx % pseInfo.pseBSize) * pseInfo.n2G * pseInfo.pseS2Size * pseInfo.pseS1Size; + int64_t n2Offset = pseInfo.n2oIdx * pseInfo.gSize * pseInfo.pseS2Size * pseInfo.pseS1Size; + int64_t gOffset = pseInfo.goIdx * pseInfo.pseS2Size * pseInfo.pseS1Size; + int64_t row = pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + pseInfo.loopIdx * pseInfo.vec1S1BaseSize; + int64_t column = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize; + int64_t m = 0; + int64_t k = 0; + if constexpr (layOutType != LayOutTypeEnum::LAYOUT_TND) { + int64_t threshold = pseInfo.s1Size - pseInfo.pseS1Size; + if (row >= threshold) { + m = row - threshold; + k = column; + } else { + m = row % pseInfo.pseS1Size; + k = pseInfo.pseS2Size - (row - column) - (pseInfo.pseS1Size - m); + } + } else { + int64_t threshold = pseInfo.pseS2Size - pseInfo.pseS1Size; + int64_t posVal = row - column - threshold; + if (threshold >= 0) { + if (posVal >= 0) { + m = posVal; + k = 0; + } else { + m = 0; + k = -posVal; + } + } else { + m = posVal; + k = 0; + } + } + int64_t s1Offset = m * pseInfo.pseS2Size; + int64_t s2Offset = k; + pseInfo.readS2Size = Min(pseInfo.s2AlignedSize, pseInfo.pseS2Size - k); + pseInfo.pseS2ComputeSize = Align(pseInfo.readS2Size); + + return bOffset + n2Offset + gOffset + s1Offset + s2Offset; + } else { + return 0; + } +} + +template __aicore__ inline bool NeedPseAlibiCompute(PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + // Alibi编码只计算下三角 + if (pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + (pseInfo.loopIdx + 1) * pseInfo.vec1S1BaseSize <= + pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize) { + return false; + } + return true; + } else { + return false; + } +} + +template +__aicore__ inline void PseAlibiCopyIn(LocalTensor &dstTensor, LocalTensor &tmpTensor, + GlobalTensor &srcTensor, PseInfo &pseInfo, int64_t alignedSize = 16) +{ + if constexpr (hasPse == true) { + if (!NeedPseAlibiCompute(pseInfo)) { + return; + } + int64_t offset = PseAlibiComputeOffset(pseInfo); + if constexpr (IsSameType::value) { + if (!pseInfo.align8){ + DataCopyIn(dstTensor, srcTensor, offset, pseInfo.vec1S1RealSize, pseInfo.readS2Size, + pseInfo.pseS2Size, alignedSize); + } else { + DataCopyInAlign8(dstTensor, srcTensor, offset, pseInfo.vec1S1RealSize, + pseInfo.readS2Size, pseInfo.pseS2Size); + } + return; + } + + DataCopyIn(tmpTensor, srcTensor, offset, pseInfo.vec1S1RealSize, pseInfo.readS2Size, + pseInfo.pseS2Size, alignedSize); + if (pseInfo.needCast) { + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(dstTensor, tmpTensor, RoundMode::CAST_NONE, pseInfo.vec1S1RealSize * pseInfo.pseS2ComputeSize); + } + return; + } +} + +template +__aicore__ inline void PseSlopeCopyIn(LocalTensor &dstTensor, LocalTensor &helpTensor, + __gm__ uint8_t *pseSlope, GlobalTensor &alibiGm, PseInfo &pseInfo, + int64_t alignedSize = 16) { + if constexpr (hasPse == true) { + int64_t bOffset = 0; + int64_t n2Offset = pseInfo.n2oIdx * pseInfo.gSize; + int64_t gOffset = pseInfo.goIdx; + + if (pseInfo.pseShapeType == pseSlopeBn) { + bOffset = pseInfo.boIdx * pseInfo.n2G; + } + int64_t offset = bOffset + n2Offset + gOffset; + + DataCopyIn(helpTensor, alibiGm, 0, pseInfo.vec1S1RealSize, + pseInfo.s2RealSize, pseInfo.pseAlibiBaseS2, alignedSize); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + + if (pseInfo.needCast) { + int64_t computeSize = pseInfo.vec1S1RealSize * pseInfo.s2AlignedSize; + Cast(dstTensor, helpTensor, RoundMode::CAST_NONE, computeSize); + AscendC::PipeBarrier(); + + int64_t s1Offset = pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + pseInfo.loopIdx * pseInfo.vec1S1BaseSize; + int64_t s2Offset = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize; + + float posShift = float(s2Offset + pseInfo.kvStartIdx - s1Offset - pseInfo.qStartIdx); + + Adds(dstTensor, dstTensor, posShift, computeSize); + AscendC::PipeBarrier(); + Abs(dstTensor, dstTensor, computeSize); + AscendC::PipeBarrier(); + float slopes = ((__gm__ T *)pseSlope)[offset] * -1; + if (pseInfo.pseType == (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_SQRT_TYPE) { + Sqrt(dstTensor, dstTensor, computeSize); + AscendC::PipeBarrier(); + } + Muls(dstTensor, dstTensor, slopes, computeSize); + AscendC::PipeBarrier(); + } + } +} + +template +__aicore__ inline void PseSlopeCast(LocalTensor &dstTensor, LocalTensor &helpTensor, + __gm__ uint8_t *pseSlope, PseInfo &pseInfo) { + if constexpr (hasPse == true) { + int64_t bOffset = 0; + int64_t n2Offset = pseInfo.n2oIdx * pseInfo.gSize; + int64_t gOffset = pseInfo.goIdx; + + if (pseInfo.pseShapeType == pseSlopeBn) { + bOffset = pseInfo.boIdx * pseInfo.n2G; + } + int64_t offset = bOffset + n2Offset + gOffset; + int64_t computeSize = pseInfo.vec1S1RealSize * pseInfo.s2AlignedSize; + Cast(dstTensor, helpTensor, RoundMode::CAST_NONE, computeSize); + AscendC::PipeBarrier(); + + int64_t s1Offset = pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + pseInfo.loopIdx * pseInfo.vec1S1BaseSize; + int64_t s2Offset = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize; + + float posShift = float(s2Offset + pseInfo.kvStartIdx - s1Offset - pseInfo.qStartIdx); + + Adds(dstTensor, dstTensor, posShift, computeSize); + AscendC::PipeBarrier(); + Abs(dstTensor, dstTensor, computeSize); + AscendC::PipeBarrier(); + float slopes = ((__gm__ T *)pseSlope)[offset] * -1; + if (pseInfo.pseType == (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_SQRT_TYPE) { + Sqrt(dstTensor, dstTensor, computeSize); + AscendC::PipeBarrier(); + } + Muls(dstTensor, dstTensor, slopes, computeSize); + AscendC::PipeBarrier(); + } +} + +template +__aicore__ inline void PseCopyIn(LocalTensor &dstTensor, LocalTensor &tmpTensor, + GlobalTensor &srcTensor, PseInfo &pseInfo, int64_t alignedSize = 16) +{ + if constexpr (hasPse == true) { + if (pseInfo.pseEncodeType == pseEncodeALibiS2Full) { + return PseAlibiCopyIn(dstTensor, tmpTensor, srcTensor, pseInfo, alignedSize); + } + int64_t offset = PseComputeOffset(pseInfo); + int64_t s1Size = pseInfo.pseShapeType == pse1S2 ? (pseInfo.blockCount == 0 ? 1 : pseInfo.blockCount) : + pseInfo.vec1S1RealSize; + + if constexpr (IsSameType::value) { + if (!pseInfo.align8){ + DataCopyIn(dstTensor, srcTensor, offset, s1Size, pseInfo.s2RealSize, + pseInfo.s2Size, alignedSize); + } else { + DataCopyInAlign8(dstTensor, srcTensor, offset, s1Size, pseInfo.s2RealSize, pseInfo.s2Size); + } + return; + } + DataCopyIn(tmpTensor, srcTensor, offset, s1Size, pseInfo.s2RealSize, pseInfo.s2Size, + alignedSize); + if (pseInfo.needCast) { + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(dstTensor, tmpTensor, RoundMode::CAST_NONE, s1Size * pseInfo.s2AlignedSize); + } + return; + } +} + +template +__aicore__ inline void PseAlibiCompute(LocalTensor &dstTensor, LocalTensor &pseTensor, PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + if (!NeedPseAlibiCompute(pseInfo)) { + return; + } + Add(dstTensor, dstTensor, pseTensor, pseInfo.vec1S1RealSize * pseInfo.pseS2ComputeSize); + return; + } +} + +template +__aicore__ inline void PseCompute(LocalTensor &dstTensor, LocalTensor &pseTensor, PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + if (pseInfo.pseEncodeType == pseEncodeALibiS2Full) { + return PseAlibiCompute(dstTensor, pseTensor, pseInfo); + } + int64_t computeSize = (pseInfo.pseShapeType == pseS1S2 || pseInfo.pseShapeType == pseSlopeBn || + pseInfo.pseShapeType == pseSlopeN) + ? pseInfo.vec1S1RealSize * pseInfo.s2AlignedSize + : pseInfo.s2AlignedSize; + PseBroadcastAdd(pseInfo.vec1S1RealSize, pseInfo.s2AlignedSize, computeSize, pseTensor, + dstTensor, pseInfo.pseShapeType); + return; + } +} + +template +__aicore__ inline void PseInnerAlibiCreate(GlobalTensor &dstTensor, LocalTensor &helpTensor, PseInfo &pseInfo) { + if constexpr (hasPse == true) { + if (pseInfo.pseType != (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_TYPE && pseInfo.pseType != (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_SQRT_TYPE) { + return; + } + event_t eventIdMte3ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V)); + event_t eventIdMte3ToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S)); + event_t eventIdVToMte3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + float tmpValue = -1.0; + + for (int64_t i = 0; i < pseInfo.pseAlibiBaseS1; i++) { + CreateVecIndex(helpTensor, (half)(i * tmpValue), pseInfo.pseAlibiBaseS2); + SetFlag(eventIdVToMte3); + WaitFlag(eventIdVToMte3); + DataCopy(dstTensor[i * pseInfo.pseAlibiBaseS2], helpTensor, pseInfo.pseAlibiBaseS2); + SetFlag(eventIdMte3ToV); + WaitFlag(eventIdMte3ToV); + SetFlag(eventIdMte3ToS); + WaitFlag(eventIdMte3ToS); + } + } +} +#endif diff --git a/csrc/ascend/common/include/kernel/simd.h b/csrc/ascend/common/include/kernel/simd.h new file mode 100644 index 000000000..de67a52a7 --- /dev/null +++ b/csrc/ascend/common/include/kernel/simd.h @@ -0,0 +1,433 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file simd.h + * \brief + */ + +#ifndef INCLUDE_SIMD_H +#define INCLUDE_SIMD_H + +#ifdef __CCE_KT_TEST__ +#define __bf16 bfloat16_t +#endif + +#include "hardware.h" +#include "kernel_operator.h" + +///////////////////////////////////////////////////// +// vadd +///////////////////////////////////////////////////// +template +__aicore__ inline void add_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint8_t repeat, + uint8_t dstBlockStride, + uint8_t src0BlockStride, + uint8_t src1BlockStride, + uint8_t dstRepeatStride, + uint8_t src0RepeatStride, + uint8_t src1RepeatStride) +{ + AscendC::Add( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::BinaryRepeatParams( + dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride)); +} + +///////////////////////////////////////////////////// +// vadds +///////////////////////////////////////////////////// +template +__aicore__ inline void adds_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + DType scalarValue, + uint8_t repeat, + uint8_t dstBlockStride, + uint8_t srcBlockStride, + uint8_t dstRepeatStride, + uint8_t srcRepeatStride) +{ + AscendC::Adds( + dst, + src, + scalarValue, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vcadd +///////////////////////////////////////////////////// +template +__aicore__ inline void cadd_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstRepeatStride, + uint16_t srcBlockStride, + uint16_t srcRepeatStride) +{ + AscendC::RepeatReduceSum(dst, src, repeat, 0, 0, srcBlockStride, dstRepeatStride, srcRepeatStride); +} +///////////////////////////////////////////////////// +// vbrcb +///////////////////////////////////////////////////// +template +__aicore__ inline void brcb_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint16_t dstBlockStride, + uint16_t dstRepeatStride, + uint8_t repeat) +{ + AscendC::Brcb(dst, src, repeat, AscendC::BrcbRepeatParams(dstBlockStride, dstRepeatStride)); +} + +///////////////////////////////////////////////////// +// vcmax +///////////////////////////////////////////////////// +template +__aicore__ inline void cmax_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstRepeatStride, + uint16_t srcBlockStride, + uint16_t srcRepeatStride) +{ +#if defined(__DAV_C220_VEC__) + AscendC::WholeReduceMax( + dst, src, (int32_t)0, repeat, dstRepeatStride, srcBlockStride, srcRepeatStride, OrderType); +#else + AscendC::WholeReduceMax( + dst, src, (int32_t)0, repeat, dstRepeatStride, srcBlockStride, srcRepeatStride); +#endif +} + +///////////////////////////////////////////////////// +// vconv +///////////////////////////////////////////////////// +template +__aicore__ inline void conv_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + if constexpr (std::is_same::value && std::is_same::value) { + AscendC::Cast( + dst, + src, + AscendC::RoundMode::CAST_RINT, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); + } else { + AscendC::Cast( + dst, + src, + AscendC::RoundMode::CAST_NONE, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); + } +} + +///////////////////////////////////////////////////// +// vconv_f322bf16r +///////////////////////////////////////////////////// +template +__aicore__ inline void convr_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Cast( + dst, + src, + AscendC::RoundMode::CAST_RINT, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vdiv +///////////////////////////////////////////////////// +template +__aicore__ inline void div_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint8_t repeat, + uint8_t dstBlockStride, + uint8_t src0BlockStride, + uint8_t src1BlockStride, + uint8_t dstRepeatStride, + uint8_t src0RepeatStride, + uint8_t src1RepeatStride) +{ + AscendC::Div( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::BinaryRepeatParams( + dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride)); +} + +///////////////////////////////////////////////////// +// vexp +///////////////////////////////////////////////////// +template +__aicore__ inline void exp_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Exp( + dst, + src, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vmax +///////////////////////////////////////////////////// +template +__aicore__ inline void max_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint8_t repeat, + uint8_t dstBlockStride, + uint8_t src0BlockStride, + uint8_t src1BlockStride, + uint8_t dstRepeatStride, + uint8_t src0RepeatStride, + uint8_t src1RepeatStride) +{ + AscendC::Max( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::BinaryRepeatParams( + dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride)); +} + +///////////////////////////////////////////////////// +// vmul +///////////////////////////////////////////////////// +template +__aicore__ inline void mul_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint8_t repeat, + uint8_t dstBlockStride, + uint8_t src0BlockStride, + uint8_t src1BlockStride, + uint8_t dstRepeatStride, + uint8_t src0RepeatStride, + uint8_t src1RepeatStride) +{ + AscendC::Mul( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::BinaryRepeatParams( + dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride)); +} + +///////////////////////////////////////////////////// +// vmuls +///////////////////////////////////////////////////// +template +__aicore__ inline void muls_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + DType src1, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Muls( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vsub +///////////////////////////////////////////////////// +template +__aicore__ inline void sub_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + AscendC::LocalTensor src1, + uint8_t repeat, + uint8_t dstBlockStride, + uint8_t src0BlockStride, + uint8_t src1BlockStride, + uint8_t dstRepeatStride, + uint8_t src0RepeatStride, + uint8_t src1RepeatStride) +{ + AscendC::Sub( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::BinaryRepeatParams( + dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride)); +} + +///////////////////////////////////////////////////// +// vmaxs +///////////////////////////////////////////////////// +template +__aicore__ inline void maxs_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + DType src1, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Maxs( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vmins +///////////////////////////////////////////////////// +template +__aicore__ inline void mins_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src0, + DType src1, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Mins( + dst, + src0, + src1, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vsqrt +///////////////////////////////////////////////////// +template +__aicore__ inline void sqrt_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Sqrt( + dst, + src, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vln +///////////////////////////////////////////////////// +template +__aicore__ inline void ln_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + uint8_t repeat, + uint16_t dstBlockStride, + uint16_t srcBlockStride, + uint16_t dstRepeatStride, + uint16_t srcRepeatStride) +{ + AscendC::Ln( + dst, + src, + (uint64_t)0, + repeat, + AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride)); +} + +///////////////////////////////////////////////////// +// vtranspose +///////////////////////////////////////////////////// +template +__aicore__ inline void tranpose_v(AscendC::LocalTensor dst, AscendC::LocalTensor src) +{ + AscendC::Transpose(dst, src); +} + +///////////////////////////////////////////////////// +// vcgmax +///////////////////////////////////////////////////// +template +__aicore__ inline void cgmax_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + const int32_t repeat, + const int32_t dstRepStride, + const int32_t srcBlkStride, + const int32_t srcRepStride) +{ + AscendC::BlockReduceMax(dst, src, repeat, 0, dstRepStride, srcBlkStride, srcRepStride); +} + +///////////////////////////////////////////////////// +// vcgadd +///////////////////////////////////////////////////// +template +__aicore__ inline void cgadd_v(AscendC::LocalTensor dst, + AscendC::LocalTensor src, + const int32_t repeat, + const int32_t dstRepStride, + const int32_t srcBlkStride, + const int32_t srcRepStride) +{ + AscendC::BlockReduceSum(dst, src, repeat, 0, dstRepStride, srcBlkStride, srcRepStride); +} +#endif \ No newline at end of file diff --git a/csrc/ascend/common/include/kernel/util.h b/csrc/ascend/common/include/kernel/util.h new file mode 100644 index 000000000..71ac733bb --- /dev/null +++ b/csrc/ascend/common/include/kernel/util.h @@ -0,0 +1,159 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file util.h + * \brief + */ + +#ifndef FLASH_ATTENTION_UTIL_H +#define FLASH_ATTENTION_UTIL_H + +constexpr int32_t blockBytes = 32; +constexpr int32_t byteBitRatio = 8; +constexpr int64_t prefixAttenMaskDownHeight = 1024; +constexpr static int32_t blockSize = blockBytes / 4; // 4 means sizeof(T) +constexpr static int32_t repeatMaxBytes = 256; +constexpr static int32_t repeatMaxTimes = 255; +constexpr static int32_t repeatMaxSize = repeatMaxBytes / 4; // 4 means sizeof(T) + +using AscendC::LocalTensor; +using AscendC::GlobalTensor; +using AscendC::DataFormat; +using AscendC::ShapeInfo; +using AscendC::DataCopyParams; +using AscendC::DataCopyExtParams; +using AscendC::DataCopyPadParams; +using AscendC::DataCopyPadExtParams; +using AscendC::BinaryRepeatParams; +using AscendC::IsSameType; +using AscendC::HardEvent; +using AscendC::SetFlag; +using AscendC::WaitFlag; + +enum class LayOutTypeEnum { None = 0, LAYOUT_BSH = 1, LAYOUT_SBH = 2, LAYOUT_BNSD = 3, LAYOUT_TND = 4, LAYOUT_NTD_TND = 5}; + +namespace math { +template __aicore__ inline T Ceil(T a, T b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +template __aicore__ inline T Align(T a, T b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b * b; +} +} + +template +__aicore__ inline T1 CeilDiv(T1 a, T2 b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +template +__aicore__ inline T1 Max(T1 a, T2 b) +{ + return (a > b) ? (a) : (b); +} + +template +__aicore__ inline T1 Min(T1 a, T2 b) +{ + return (a > b) ? (b) : (a); +} + +__aicore__ inline void BoolCopyIn(LocalTensor &dstTensor, GlobalTensor &srcTensor, + int64_t srcOffset, uint32_t s1Size, uint32_t s2Size, int64_t totalS2Size, int64_t alignedSize = blockBytes) +{ + uint32_t alignedS2Size = CeilDiv(s2Size, alignedSize) * alignedSize; + uint32_t shapeArray[] = {s1Size, alignedS2Size}; + dstTensor.SetShapeInfo(ShapeInfo(2, shapeArray, DataFormat::ND)); + dstTensor.SetSize(s1Size * alignedS2Size); + DataCopyParams dataCopyParams; + dataCopyParams.blockCount = s1Size; + dataCopyParams.dstStride = 0; + if (totalS2Size == blockBytes && alignedSize == 64) { // totalS2Size < 64 && totalS2Size % blockBytes == 0 + dataCopyParams.dstStride = 1; + alignedSize = blockBytes; + alignedS2Size = CeilDiv(s2Size, blockBytes) * blockBytes; + } + if (likely(totalS2Size - s2Size <= UINT16_MAX)) { + if (totalS2Size % alignedSize == 0) { + dataCopyParams.blockLen = alignedS2Size / blockBytes; + dataCopyParams.srcStride = (totalS2Size - alignedS2Size) / blockBytes; + DataCopy(dstTensor, srcTensor[srcOffset], dataCopyParams); + } else { + dataCopyParams.blockLen = s2Size; + dataCopyParams.srcStride = totalS2Size - s2Size; + DataCopyPadParams dataCopyPadParams; + dataCopyPadParams.isPad = true; + dataCopyPadParams.rightPadding = Min(alignedS2Size - s2Size, blockBytes); + dataCopyPadParams.paddingValue = 1; + DataCopyPad(dstTensor, srcTensor[srcOffset], dataCopyParams, dataCopyPadParams); + } + } else { + DataCopyExtParams extParams; + extParams.blockCount = s1Size; + extParams.dstStride = 0; + extParams.blockLen = s2Size; + extParams.srcStride = totalS2Size - s2Size; + DataCopyPadExtParams dataCopyPadParams; + dataCopyPadParams.isPad = true; + dataCopyPadParams.rightPadding = Min(alignedS2Size - s2Size, blockBytes); + dataCopyPadParams.paddingValue = 1; + DataCopyPad(dstTensor, srcTensor[srcOffset], extParams, dataCopyPadParams); + } +} + +__aicore__ inline void Bit2Int8CopyIn(LocalTensor &dstTensor, GlobalTensor &srcTensor, + int64_t srcOffset, uint32_t batchSize, uint32_t s1BaseSize, uint32_t s2BaseSize, int64_t s2TotalSize, + int64_t alignedSize = blockBytes) +{ + uint32_t alignedS2Size = CeilDiv(s2BaseSize / byteBitRatio, alignedSize) * alignedSize; + uint32_t shapeArray[] = {batchSize * s1BaseSize, alignedS2Size}; + dstTensor.SetShapeInfo(ShapeInfo(2, shapeArray, DataFormat::ND)); + dstTensor.SetSize(batchSize * s1BaseSize * alignedS2Size); + DataCopyParams dataCopyParams; + dataCopyParams.blockCount = batchSize * s1BaseSize; + dataCopyParams.blockLen = CeilDiv(s2BaseSize / byteBitRatio, blockBytes); + dataCopyParams.dstStride = 0; + if (s2TotalSize / byteBitRatio % alignedSize == 0 && s2BaseSize / byteBitRatio % alignedSize == 0) { + dataCopyParams.srcStride = + (s2TotalSize / byteBitRatio - dataCopyParams.blockLen * blockBytes) / blockBytes; + DataCopy(dstTensor, srcTensor[srcOffset / byteBitRatio], dataCopyParams); + } else { + dataCopyParams.blockLen = CeilDiv(s2BaseSize , byteBitRatio); + dataCopyParams.srcStride = (s2TotalSize - s2BaseSize) / byteBitRatio; + DataCopyPadParams dataCopyPadParams; + dataCopyPadParams.isPad = true; + dataCopyPadParams.rightPadding = 0; + dataCopyPadParams.paddingValue = 0; + DataCopyPad(dstTensor, srcTensor[srcOffset / byteBitRatio], dataCopyParams, dataCopyPadParams); + } +} + +__aicore__ inline int32_t Align(int32_t shape) +{ + int32_t alignFactor = 16; + int32_t alignedSize = CeilDiv(shape, alignFactor) * alignFactor; + return alignedSize; +} + +#endif // FLASH_ATTENTION_UTIL_H diff --git a/csrc/ascend/common/include/op_graph/op_transformer_proto_extend.h b/csrc/ascend/common/include/op_graph/op_transformer_proto_extend.h new file mode 100644 index 000000000..457bc4c2b --- /dev/null +++ b/csrc/ascend/common/include/op_graph/op_transformer_proto_extend.h @@ -0,0 +1,100 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file op_transformer_proto_extend.h + * \brief + */ +#ifndef OPS_OP_MATH_PROTO_EXTEND_H_ +#define OPS_OP_MATH_PROTO_EXTEND_H_ + +#include "graph/operator_reg.h" + +namespace ge { +/** +* @brief swin_transformer model specific structure.Operator only supports swin_transformer. + +* @par Inputs: +* Three inputs, including: +* @li x: An ND Tensor. Must be one of the following types: float16, float, bfloat16, + the shape should be (B*W, N, S1, S2) or (B, W, N, S1, S2). +* @li atten_mask: An ND Tensor. Must be one of the following types: float16, float, bfloat16, + the shape should be (W, S1, S2) or (W, 1, S1, S2) or (1, W, 1, S1, S2) +* @li relative_pos_bias: An ND Tensor. Must be one of the following types: float16, float, bfloat16. + the shape sholud be (N, S1, S2) or (1, N, S1, S2) or (1, 1, N, S1, S2) + +* @par Attributes: +* @li scale_value: A optional attribute, the type is float. Defaults to 1.0. +* @li inner_precision_mode: A optional attribute, the type is int. Defaults to 0, reserved field. + +* @par Outputs: +* One output, including: +* @li y: An ND Tensor. Must be one of the following types: float16, float, bfloat16, + the shape should be same with x. +*/ +REG_OP(MaskedSoftmaxWithRelPosBias) + .INPUT(x, TensorType({DT_FLOAT16, DT_BFLOAT16, DT_FLOAT})) + .OPTIONAL_INPUT(atten_mask, TensorType({DT_FLOAT16, DT_BFLOAT16, DT_FLOAT})) + .INPUT(relative_pos_bias, TensorType({DT_FLOAT16, DT_BFLOAT16, DT_FLOAT})) + .OUTPUT(y, TensorType({DT_FLOAT16, DT_BFLOAT16, DT_FLOAT})) + .ATTR(scale_value, Float, 1.0) + .ATTR(inner_precision_mode, Int, 0) + .OP_END_FACTORY_REG(MaskedSoftmaxWithRelPosBias) + +/** +* @brief AttentionScore's forward calculation. + +* @par Inputs: +* six inputs, including: +* @li query: A matrix Tensor. The type only support float16. Enter a 4D Tensor. +* @li key: A matrix Tensor. The type only support float16. Enter a 4D Tensor. +* @li value: A matrix Tensor. The type only support float16. Enter a 4D Tensor. +* @li padding_mask: A matrix Tensor. The type only support float16. Enter a 4D Tensor. +* @li scale: A scalar. The type only support float16. Enter a 4D Tensor. +* @li drop_mask: A matrix Tensor. An optional input parameter. The type only support uint8. Enter a 4D Tensor. + +* @par Attributes: +* @li keep_prob: A float. The keep probability of dropout. Default: 1.0. +* @li query_transpose: A bool. If True, changes the shape of "query" from [B, N, S, D] to [B, N, D, S]. +* Default: false. +* @li key_transpose: A bool. If True, changes the shape of "key" from [B, N, S, D] to [B, N, D, S]. +* Default: false. +* @li bmm_score_transpose_a: A bool. If True, changes the shape of "mid_data" from [B, N, S, D] to [B, N, D, S]. +* Default: false. +* @li bmm_score_transpose_b: A bool. If True, changes the shape of "value" from [B, N, S, D] to [B, N, D, S]. +* Default: false. +* @li softmax_axes: A list of int. The dimension softmax would be performed on. Defaults to "[-1]". + +* @par Outputs: +* attention_score: The result matrix Tensor. The type only support float16. The output shape is the same as query. +* softmax_output: The result matrix Tensor. The type only support float16. The output shape is the same as query. + +* @par Restrictions: +* Warning: THIS FUNCTION IS EXPERIMENTAL. Please do not use. +*/ +REG_OP(AttentionScore) + .INPUT(query, TensorType({DT_FLOAT16})) + .INPUT(key, TensorType({DT_FLOAT16})) + .INPUT(value, TensorType({DT_FLOAT16})) + .INPUT(padding_mask, TensorType({DT_FLOAT16})) + .INPUT(scale, TensorType({DT_FLOAT16})) + .OPTIONAL_INPUT(drop_mask, TensorType({DT_INT8})) + .OUTPUT(attention_score, TensorType({DT_FLOAT16})) + .OUTPUT(softmax_output, TensorType({DT_FLOAT16})) + .ATTR(keep_prob, Float, 1.0) + .ATTR(query_transpose, Bool, false) + .ATTR(key_transpose, Bool, false) + .ATTR(bmm_score_transpose_a, Bool, false) + .ATTR(bmm_score_transpose_b, Bool, false) + .ATTR(softmax_axes, ListInt, {-1}) + .OP_END_FACTORY_REG(AttentionScore) +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/include/static/op_resource.h b/csrc/ascend/common/include/static/op_resource.h new file mode 100644 index 000000000..2ad2ed1ae --- /dev/null +++ b/csrc/ascend/common/include/static/op_resource.h @@ -0,0 +1,32 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file op_resource.h + * \brief + */ +#ifndef COMMON_NN_OP_RESOURCE_H +#define COMMON_NN_OP_RESOURCE_H + +#define EXTERN_OP_RESOURCE(kernelName) \ +namespace l0op { \ + extern void * kernelName##TilingRegisterResource(); \ + extern void * kernelName##InferShapeRegisterResource(); \ + extern void * kernelName##TuningRegisterResource(); \ + extern const OP_BINARY_RES& kernelName##KernelResource(); \ + extern const OP_RUNTIME_KB_RES& kernelName##TuningResource(); \ + [[maybe_unused]] uint32_t kernelName##_kernelName_Be_Defined_Multi_Times___; \ +} + +#define AUTO_GEN_OP_RESOURCE(kernelName) {{ #kernelName, \ + {{l0op::kernelName##TilingRegisterResource(), l0op::kernelName##InferShapeRegisterResource(), l0op::kernelName##TuningRegisterResource()}, \ + l0op::kernelName##KernelResource(), l0op::kernelName##TuningResource()}}} \ + +#endif // COMMON_NN_OP_RESOURCE_H \ No newline at end of file diff --git a/csrc/ascend/common/include/static/static_space.h b/csrc/ascend/common/include/static/static_space.h new file mode 100644 index 000000000..3491adc5b --- /dev/null +++ b/csrc/ascend/common/include/static/static_space.h @@ -0,0 +1,34 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file static_space.h + * \brief + */ +#ifndef CANN_OPS_STATIC_SPACE_H_ +#define CANN_OPS_STATIC_SPACE_H_ +#include "base/registry/op_impl_space_registry_v2.h" + +class StaticSpaceInitializer { +public: + static StaticSpaceInitializer& GetInstance() { + static StaticSpaceInitializer instance; + return instance; + } +private: + StaticSpaceInitializer () { + auto space_registry = gert::DefaultOpImplSpaceRegistryV2::GetInstance().GetSpaceRegistry(); + if (space_registry == nullptr) { + space_registry = std::make_shared(); + gert::DefaultOpImplSpaceRegistryV2::GetInstance().SetSpaceRegistry(space_registry); + } + } +}; +#endif \ No newline at end of file diff --git a/csrc/ascend/common/include/tiling_base/data_copy_transpose_tiling.h b/csrc/ascend/common/include/tiling_base/data_copy_transpose_tiling.h new file mode 100644 index 000000000..2f2c3e4ae --- /dev/null +++ b/csrc/ascend/common/include/tiling_base/data_copy_transpose_tiling.h @@ -0,0 +1,51 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file data_copy_transpose_tiling.h + * \brief + */ + +#pragma once + +#include +#include +#include "data_copy_transpose_tiling_def.h" + +namespace optiling { + +inline void GetDataCopyTransposeTiling(const ge::Shape &dstShape, const ge::Shape &srcShape, const uint32_t typeSize, + optiling::CopyTransposeTiling &tiling) +{ + constexpr int64_t B_INDEX = 0; + constexpr int64_t N_INDEX = 1; + constexpr int64_t S_INDEX = 2; + constexpr int64_t H_INDEX = 3; + std::vector dstShapeInfo = dstShape.GetDims(); + std::vector srcShapeInfo = srcShape.GetDims(); + + tiling.set_dstShapeB(dstShapeInfo[B_INDEX]); + tiling.set_dstShapeN(dstShapeInfo[N_INDEX]); + tiling.set_dstShapeS(dstShapeInfo[S_INDEX]); + tiling.set_dstShapeH(dstShapeInfo[H_INDEX]); + tiling.set_dstShapeHN(tiling.get_dstShapeH() / tiling.get_dstShapeN()); + + tiling.set_srcShapeB(srcShapeInfo[B_INDEX]); + tiling.set_srcShapeN(srcShapeInfo[N_INDEX]); + tiling.set_srcShapeS(srcShapeInfo[S_INDEX]); + tiling.set_srcShapeHN(srcShapeInfo[H_INDEX]); + tiling.set_originalShapeNLen(tiling.get_srcShapeHN() * typeSize); + tiling.set_shapeSHValue(tiling.get_dstShapeS() * tiling.get_dstShapeH()); + tiling.set_shapeNsValue(tiling.get_dstShapeN() * tiling.get_dstShapeS()); + tiling.set_shapeNsnValue(tiling.get_dstShapeN() * tiling.get_srcShapeS() * tiling.get_srcShapeN()); + tiling.set_shapeBHValue(tiling.get_dstShapeB() * tiling.get_dstShapeH()); +} + +} // namespace optiling diff --git a/csrc/ascend/common/include/tiling_base/data_copy_transpose_tiling_def.h b/csrc/ascend/common/include/tiling_base/data_copy_transpose_tiling_def.h new file mode 100644 index 000000000..891a566a0 --- /dev/null +++ b/csrc/ascend/common/include/tiling_base/data_copy_transpose_tiling_def.h @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file data_copy_transpose_tiling_def.h + * \brief + */ + +#pragma once + +#include +#include + +namespace optiling { + +BEGIN_TILING_DATA_DEF(CopyTransposeTiling) +TILING_DATA_FIELD_DEF(uint32_t, dstShapeB); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeN); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeS); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeHN); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeH); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeB); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeN); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeS); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeHN); +TILING_DATA_FIELD_DEF(uint32_t, originalShapeNLen); +TILING_DATA_FIELD_DEF(uint32_t, shapeSHValue); +TILING_DATA_FIELD_DEF(uint32_t, shapeNsValue); +TILING_DATA_FIELD_DEF(uint32_t, shapeNsnValue); +TILING_DATA_FIELD_DEF(uint32_t, invalidParamCopyTransposeTiling); +TILING_DATA_FIELD_DEF(uint32_t, shapeBHValue); +TILING_DATA_FIELD_DEF(uint32_t, paramsAlign); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(CopyTransposeTilingOp, CopyTransposeTiling) + +} // namespace optiling diff --git a/csrc/ascend/common/include/tiling_base/error_log.h b/csrc/ascend/common/include/tiling_base/error_log.h new file mode 100644 index 000000000..ca2499252 --- /dev/null +++ b/csrc/ascend/common/include/tiling_base/error_log.h @@ -0,0 +1,73 @@ +#ifndef OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_ +#define OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_ + +#include +#include + +#include "tiling_base/tiling_util.h" +#include "toolchain/slog.h" + +#define OP_LOGI(opname, ...) +#define OP_LOGD(opname, ...) + +#define OP_LOGW(opname, ...) \ + do { \ + (void)(opname); \ + std::printf("[WARN] "); \ + std::printf(__VA_ARGS__); \ + std::printf("\n"); \ + } while (0) + +#define OP_LOGE_WITHOUT_REPORT(opname, ...) \ + do { \ + (void)(opname); \ + std::printf("[ERRORx] "); \ + std::printf(__VA_ARGS__); \ + std::printf("\n"); \ + } while (0) + +#define OP_LOGE(opname, ...) \ + do { \ + (void)(opname); \ + std::printf("[ERROR] "); \ + std::printf(__VA_ARGS__); \ + std::printf("\n"); \ + } while (0) + +namespace optiling { + +#define VECTOR_INNER_ERR_REPORT_TILIING(op_name, err_msg, ...) \ + do { \ + OP_LOGE_WITHOUT_REPORT(op_name, err_msg, ##__VA_ARGS__); \ + } while (0) + +#define OP_CHECK_IF(cond, log_func, expr) \ + do { \ + if (cond) { \ + log_func; \ + expr; \ + } \ + } while (0) + +#define OP_TILING_CHECK(cond, log_func, expr) \ + do { \ + if (cond) { \ + log_func; \ + expr; \ + } \ + } while (0) + +#define OP_CHECK_NULL_WITH_CONTEXT(context, ptr) \ + do { \ + if ((ptr) == nullptr) { \ + OP_LOGE(context->GetNodeType(), "%s is null", #ptr); \ + return ge::GRAPH_FAILED; \ + } \ + } while (0) + +} // namespace optiling + +using Ops::Transformer::CeilAlign; +using Ops::Transformer::CeilDiv; + +#endif // OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_ diff --git a/csrc/ascend/common/include/tiling_base/tiling_base.h b/csrc/ascend/common/include/tiling_base/tiling_base.h new file mode 100644 index 000000000..36be50a78 --- /dev/null +++ b/csrc/ascend/common/include/tiling_base/tiling_base.h @@ -0,0 +1,256 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_base.h + * \brief + */ + +#pragma once + +#include +#include +#include +#include "tiling/platform/platform_ascendc.h" +#include "tiling_base/error_log.h" + +#ifdef ASCENDC_OP_TEST +#define ASCENDC_EXTERN_C extern "C" +#else +#define ASCENDC_EXTERN_C +#endif + +namespace Ops { +namespace Transformer { +namespace OpTiling { + +struct AiCoreParams { + uint64_t ubSize = 0; + uint64_t blockDim = 0; + uint64_t aicNum = 0; + uint64_t l1Size = 0; + uint64_t l0aSize = 0; + uint64_t l0bSize = 0; + uint64_t l0cSize = 0; +}; + +struct CompileInfoCommon { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; + int32_t socVersion; + uint32_t rsvd; +}; + +struct FlashAttentionScoreGradCompileInfo { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; + platform_ascendc::SocVersion socVersion; +}; + +struct FACompileInfoCommon { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; + int32_t socVersion; + uint32_t rsvd; +}; + +class TilingBaseClass { +public: + explicit TilingBaseClass(gert::TilingContext* context) : context_(context) + {} + + virtual ~TilingBaseClass() = default; + + // Tiling执行框架 + // 1、GRAPH_SUCCESS: 成功，并且不需要继续执行后续Tiling类的实现 + // 2、GRAPH_FAILED: 失败，中止整个Tiling流程 + // 3、GRAPH_PARAM_INVALID: 本类不支持，需要继续往下执行其他Tiling类的实现 + ge::graphStatus DoTiling() + { + auto ret = GetShapeAttrsInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetPlatformInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + if (!IsCapable()) { + return ge::GRAPH_PARAM_INVALID; + } + ret = DoOpTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = DoLibApiTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetWorkspaceSize(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = PostTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + context_->SetTilingKey(GetTilingKey()); + DumpTilingInfo(); + return ge::GRAPH_SUCCESS; + } + + // 更新 context + virtual void Reset(gert::TilingContext* context) + { + context_ = context; + } + +protected: + virtual bool IsCapable() = 0; + // 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小 + virtual ge::graphStatus GetPlatformInfo() = 0; + // 2、获取INPUT/OUTPUT/ATTR信息 + virtual ge::graphStatus GetShapeAttrsInfo() = 0; + // 3、计算数据切分TilingData + virtual ge::graphStatus DoOpTiling() = 0; + // 4、计算高阶API的TilingData + virtual ge::graphStatus DoLibApiTiling() = 0; + // 5、计算TilingKey + [[nodiscard]] virtual uint64_t GetTilingKey() const = 0; + // 6、计算Workspace 大小 + virtual ge::graphStatus GetWorkspaceSize() = 0; + // 7、保存Tiling数据 + virtual ge::graphStatus PostTiling() = 0; + // 8、Dump Tiling数据 + virtual void DumpTilingInfo() + { + int32_t enable = CheckLogLevel(static_cast(OP), DLOG_DEBUG); + if (enable != 1) { + return; + } + auto buf = (uint32_t*)context_->GetRawTilingData()->GetData(); + auto bufLen = context_->GetRawTilingData()->GetDataSize(); + std::ostringstream oss; + oss << "Start to dump tiling info. tilingkey:" << context_->GetTilingKey() << ", tiling data size:" << bufLen + << ", content:"; + for (size_t i = 0; i < bufLen / sizeof(uint32_t); i++) { + oss << *(buf + i) << ","; + if (oss.str().length() > 640) { // Split according to 640 to avoid truncation + OP_LOGD(context_, "%s", oss.str().c_str()); + oss.str(""); + } + } + OP_LOGD(context_, "%s", oss.str().c_str()); + } + + static uint32_t CalcTschBlockDim(uint32_t sliceNum, uint32_t aicCoreNum, uint32_t aivCoreNum) + { + uint32_t ration; + if (aicCoreNum == 0 || aivCoreNum == 0 || aicCoreNum > aivCoreNum) { + return sliceNum; + } + ration = aivCoreNum / aicCoreNum; + return (sliceNum + (ration - 1)) / ration; + } + + template + [[nodiscard]] std::string GetShapeDebugStr(const T& shape) const + { + std::ostringstream oss; + oss << "["; + if (shape.GetDimNum() > 0) { + for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) { + oss << shape.GetDim(i) << ", "; + } + oss << shape.GetDim(shape.GetDimNum() - 1); + } + oss << "]"; + return oss.str(); + } + + [[nodiscard]] std::string GetTensorDebugStr( + const gert::StorageShape* shape, const gert::CompileTimeTensorDesc* tensor) + { + if (shape == nullptr || tensor == nullptr) { + return "nil "; + } + std::ostringstream oss; + oss << "(dtype: " << ge::TypeUtils::DataTypeToSerialString(tensor->GetDataType()) << "),"; + oss << "(shape:" << GetShapeDebugStr(shape->GetStorageShape()) << "),"; + oss << "(ori_shape:" << GetShapeDebugStr(shape->GetOriginShape()) << "),"; + oss << "(format: " + << ge::TypeUtils::FormatToSerialString( + static_cast(ge::GetPrimaryFormat(tensor->GetStorageFormat()))) + << "),"; + oss << "(ori_format: " << ge::TypeUtils::FormatToSerialString(tensor->GetOriginFormat()) << ") "; + return oss.str(); + } + + [[nodiscard]] std::string GetTilingContextDebugStr() + { + std::ostringstream oss; + for (size_t i = 0; i < context_->GetComputeNodeInfo()->GetInputsNum(); ++i) { + oss << "input" << i << ": "; + oss << GetTensorDebugStr(context_->GetInputShape(i), context_->GetInputDesc(i)); + } + + for (size_t i = 0; i < context_->GetComputeNodeInfo()->GetOutputsNum(); ++i) { + oss << "output" << i << ": "; + oss << GetTensorDebugStr(context_->GetOutputShape(i), context_->GetOutputDesc(i)); + } + return oss.str(); + } + + [[nodiscard]] std::string GetTilingDataDebugStr() const + { + auto rawTilingData = context_->GetRawTilingData(); + auto rawTilingDataSize = rawTilingData->GetDataSize(); + auto data = reinterpret_cast(rawTilingData->GetData()); + size_t len = rawTilingDataSize / sizeof(int32_t); + std::ostringstream oss; + for (size_t i = 0; i < len; i++) { + oss << data[i] << ", "; + } + return oss.str(); + } + +protected: + gert::TilingContext* context_ = nullptr; + std::unique_ptr ascendcPlatform_{nullptr}; + uint32_t blockDim_{0}; + uint64_t workspaceSize_{0}; + uint64_t tilingKey_{0}; + AiCoreParams aicoreParams_; +}; + +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops diff --git a/csrc/ascend/common/include/tiling_base/tiling_key.h b/csrc/ascend/common/include/tiling_base/tiling_key.h new file mode 100644 index 000000000..411bf96a0 --- /dev/null +++ b/csrc/ascend/common/include/tiling_base/tiling_key.h @@ -0,0 +1,63 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_key.h + * \brief + */ + +#pragma once + +#include + +namespace Ops { +namespace Transformer { +namespace OpTiling { +constexpr uint64_t RecursiveSum() +{ + return 0; +} + +constexpr uint64_t kBase = 10; // 10进制进位基数 +template constexpr uint64_t RecursiveSum(T templateId, Args... templateIds) +{ + return static_cast(templateId) + kBase * RecursiveSum(templateIds...); +} + +// TilingKey 的生成规则： +// FlashAttentionScore/FlashAttentionScoreGrad 十进制位组装tiling key，包含以下关键参数，从低位到高位依次是：Ub0, Ub1, +// Block, DataType, Format, Sparse, 特化模板 Ub0、Ub1: +// 表示Ub核内切分的轴，使用枚举AxisEnum表示，因为我们允许最多切分两根轴，所以存在UB0和UB1，如果没有UB核内切分， +// 那么填AXIS_NONE。UB0和UB1各占一个十进制位; +// Block: 表示UB用来分核的轴，使用枚举AxisEnum表示，占一个十进制位; +// DataType: 表示当前tiling key支持的输入输出的数据类型，使用枚举SupportedDtype来表示，占一个十进制位 +// Format: 表示当前tiling key支持的Format, 使用枚举InputLayout表示，占一个十进制位 +// Sparse: 表示当前tiling key是否支持Sparse，使用枚举SparseCapability表示，占一个十进制位 +// 其余特化场景，定义自己的位域和值 +// usage: get tilingKey from inputted types +// uint64_t tilingKey = GET_FLASHATTENTION_TILINGKEY(AxisEnum::AXIS_S1, AxisEnum::AXIS_S2, AxisEnum::AXIS_N2, +// SupportedDtype::FLOAT32, InputLayout::BSH, SparseCapability::SUPPORT_ALL) + +constexpr uint64_t TILINGKEYOFFSET = uint64_t(10000000000000000000UL); // 10^19 +template constexpr uint64_t GET_TILINGKEY(Args... templateIds) +{ + return TILINGKEYOFFSET + RecursiveSum(templateIds...); +} + +// usage: get tilingKey from inputted types +// uint64_t tilingKey = TILINGKEY(S2, S1, N2, FLOAT32, BSND, ALL) + +#define TILINGKEY(ub2, ub1, block, dtype, layout, sparse) \ + (GET_TILINGKEY(AxisEnum::ub2, AxisEnum::ub1, AxisEnum::block, DtypeEnum::dtype, LayoutEnum::layout, \ + SparseEnum::sparse)) + +} // namespace Optiling +} // namespace Transformer +} // namespace Ops diff --git a/csrc/ascend/common/include/tiling_base/tiling_templates_registry.h b/csrc/ascend/common/include/tiling_base/tiling_templates_registry.h new file mode 100644 index 000000000..c0a923763 --- /dev/null +++ b/csrc/ascend/common/include/tiling_base/tiling_templates_registry.h @@ -0,0 +1,350 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_templates_registry.h + * \brief + */ + +#pragma once + +#include +#include +#include +#include "exe_graph/runtime/tiling_context.h" +#include "tiling_base/tiling_base.h" +#include "tiling_base/error_log.h" + +namespace Ops { +namespace Transformer { +namespace OpTiling { + +template +std::unique_ptr TILING_CLASS(gert::TilingContext* context) +{ + return std::unique_ptr(new (std::nothrow) T(context)); +} + +using TilingClassCase = std::unique_ptr (*)(gert::TilingContext*); + +class TilingCases { +public: + explicit TilingCases(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + void AddTiling(int32_t priority) + { + OP_CHECK_IF( + cases_.find(priority) != cases_.end(), OP_LOGE(op_type_, "There are duplicate registrations."), return); + cases_[priority] = TILING_CLASS; + OP_CHECK_IF( + cases_[priority] == nullptr, + OP_LOGE(op_type_, "Register op tiling func failed, please check the class name."), return); + } + + const std::map& GetTilingCases() + { + return cases_; + } + +private: + std::map cases_; + const std::string op_type_; +}; + +// --------------------------------Interfacce with soc version -------------------------------- +class TilingRegistryNew { +public: + TilingRegistryNew() = default; + +#ifdef ASCENDC_OP_TEST + static TilingRegistryNew& GetInstance(); +#else + static TilingRegistryNew& GetInstance() + { + static TilingRegistryNew registry_impl_; + return registry_impl_; + } +#endif + + std::shared_ptr RegisterOp(const std::string& op_type, int32_t soc_version) + { + auto soc_iter = registry_map_.find(soc_version); + if (soc_iter == registry_map_.end()) { + std::map> op_type_map; + op_type_map[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + registry_map_[soc_version] = op_type_map; + } else { + if (soc_iter->second.find(op_type) == soc_iter->second.end()) { + soc_iter->second[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + } + } + + OP_CHECK_IF( + registry_map_[soc_version][op_type] == nullptr, + OP_LOGE(op_type, "Register tiling func failed, please check the class name."), return nullptr); + return registry_map_[soc_version][op_type]; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context) + { + int32_t soc_version = (int32_t)platform_ascendc::SocVersion::RESERVED_VERSION; + const char* op_type = context->GetNodeType(); + fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + auto compileInfoPtr = static_cast(context->GetCompileInfo()); + OP_CHECK_IF( + compileInfoPtr == nullptr, OP_LOGE(op_type, "compileInfoPtr is null."), return ge::GRAPH_FAILED); + soc_version = compileInfoPtr->socVersion; + OP_LOGD(context, "soc version in compileInfo is %d", soc_version); + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + soc_version = static_cast(ascendcPlatform.GetSocVersion()); + OP_LOGD(context, "soc version is %d", soc_version); + if (soc_version == (int32_t)platform_ascendc::SocVersion::RESERVED_VERSION) { + OP_LOGE(op_type, "Do op tiling failed, cannot find soc version."); + return ge::GRAPH_FAILED; + } + } + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type, soc_version); + for (auto it = tilingTemplateRegistryMap.begin(); it != tilingTemplateRegistryMap.end(); ++it) { + auto tilingTemplate = it->second(context); + if (tilingTemplate != nullptr) { + ge::graphStatus status = tilingTemplate->DoTiling(); + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do general op tiling success priority=%d", it->first); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", it->first); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context, const std::vector& priorities) + { + int32_t soc_version; + const char* op_type = context->GetNodeType(); + auto platformInfoPtr = context->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + auto compileInfoPtr = reinterpret_cast(context->GetCompileInfo()); + OP_CHECK_IF( + compileInfoPtr == nullptr, OP_LOGE(op_type, "compileInfoPtr is null."), return ge::GRAPH_FAILED); + soc_version = compileInfoPtr->socVersion; + OP_LOGD(context, "soc version in compileInfo is %d", soc_version); + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + soc_version = static_cast(ascendcPlatform.GetSocVersion()); + OP_LOGD(context, "soc version is %d", soc_version); + } + + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type, soc_version); + for (auto priority_id : priorities) { + auto tilingCaseIter = tilingTemplateRegistryMap.find(priority_id); + if (tilingCaseIter != tilingTemplateRegistryMap.end()) { + auto templateFunc = tilingCaseIter->second(context); + if (templateFunc != nullptr) { + ge::graphStatus status = templateFunc->DoTiling(); + if (status == ge::GRAPH_SUCCESS) { + OP_LOGD(context, "Do general op tiling success priority=%d", priority_id); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", priority_id); + } + } + } + return ge::GRAPH_FAILED; + } + + const std::map& GetTilingTemplates(const std::string& op_type, int32_t soc_version) + { + auto soc_iter = registry_map_.find(soc_version); + OP_CHECK_IF( + soc_iter == registry_map_.end(), + OP_LOGE(op_type, "Get op tiling func failed, please check the soc version %d", soc_version), + return empty_tiling_case_); + auto op_iter = soc_iter->second.find(op_type); + OP_CHECK_IF( + op_iter == soc_iter->second.end(), OP_LOGE(op_type, "Get op tiling func failed, please check the op name."), + return empty_tiling_case_); + return op_iter->second->GetTilingCases(); + } + +private: + std::map>> registry_map_; // key is socversion + const std::map empty_tiling_case_{}; +}; + +class RegisterNew { +public: + explicit RegisterNew(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + RegisterNew& tiling(int32_t priority, int32_t soc_version) + { + auto tilingCases = TilingRegistryNew::GetInstance().RegisterOp(op_type_, soc_version); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), return *this); + tilingCases->AddTiling(priority); + return *this; + } + + template + RegisterNew& tiling(int32_t priority, const std::vector& soc_versions) + { + for (int32_t soc_version : soc_versions) { + auto tilingCases = TilingRegistryNew::GetInstance().RegisterOp(op_type_, soc_version); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), + return *this); + tilingCases->AddTiling(priority); + } + return *this; + } + +private: + const std::string op_type_; +}; + +// --------------------------------Interfacce without soc version -------------------------------- +class TilingRegistry { +public: + TilingRegistry() = default; + +#ifdef ASCENDC_OP_TEST + static TilingRegistry& GetInstance(); +#else + static TilingRegistry& GetInstance() + { + static TilingRegistry registry_impl_; + return registry_impl_; + } +#endif + + std::shared_ptr RegisterOp(const std::string& op_type) + { + if (registry_map_.find(op_type) == registry_map_.end()) { + registry_map_[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + } + OP_CHECK_IF( + registry_map_[op_type] == nullptr, + OP_LOGE(op_type, "Register tiling func failed, please check the class name."), return nullptr); + return registry_map_[op_type]; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context) + { + const char* op_type = context->GetNodeType(); + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type); + for (auto it = tilingTemplateRegistryMap.begin(); it != tilingTemplateRegistryMap.end(); ++it) { + auto tilingTemplate = it->second(context); + if (tilingTemplate != nullptr) { + ge::graphStatus status = tilingTemplate->DoTiling(); + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do general op tiling success priority=%d", it->first); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", it->first); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context, const std::vector& priorities) + { + const char* op_type = context->GetNodeType(); + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type); + for (auto priorityId : priorities) { + auto templateFunc = tilingTemplateRegistryMap[priorityId](context); + if (templateFunc != nullptr) { + ge::graphStatus status = templateFunc->DoTiling(); + if (status == ge::GRAPH_SUCCESS) { + OP_LOGD(context, "Do general op tiling success priority=%d", priorityId); + return status; + } + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do op tiling failed"); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", priorityId); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + const std::map& GetTilingTemplates(const std::string& op_type) + { + OP_CHECK_IF( + registry_map_.find(op_type) == registry_map_.end(), + OP_LOGE(op_type, "Get op tiling func failed, please check the op name."), return empty_tiling_case_); + return registry_map_[op_type]->GetTilingCases(); + } + +private: + std::map> registry_map_; + const std::map empty_tiling_case_; +}; + +class Register { +public: + explicit Register(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + Register& tiling(int32_t priority) + { + auto tilingCases = TilingRegistry::GetInstance().RegisterOp(op_type_); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), return *this); + tilingCases->AddTiling(priority); + return *this; + } + +private: + const std::string op_type_; +}; +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops + +// op_type: 算子名称， class_name: 注册的 tiling 类, soc_version：芯片版本号 +// priority: tiling 类的优先级, 越小表示优先级越高, 即会优先选择这个tiling类 +#define REGISTER_TILING_TEMPLATE_WITH_SOCVERSION(op_type, class_name, soc_versions, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::RegisterNew VAR_UNUSED##op_type##class_name##priority_register = \ + Ops::Transformer::OpTiling::RegisterNew(#op_type).tiling(priority, soc_versions) + +// op_type: 算子名称， class_name: 注册的 tiling 类, +// priority: tiling 类的优先级, 越小表示优先级越高, 即被选中的概率越大 +#define REGISTER_TILING_TEMPLATE(op_type, class_name, priority) \ + static Ops::Transformer::OpTiling::Register VAR_UNUSED##op_type_##class_name##priority_register = \ + Ops::Transformer::OpTiling::Register(op_type).tiling(priority) + +// op_type: 算子名称， class_name: 注册的 tiling 类, +// soc_version: soc版本，用于区分不同的soc +// priority: tiling 类的优先级, 越小表示优先级越高, 即会优先选择这个tiling类 +#define REGISTER_TILING_TEMPLATE_NEW(op_type, class_name, soc_version, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::RegisterNew VAR_UNUSED##op_type##class_name##priority_register = \ + Ops::Transformer::OpTiling::RegisterNew(#op_type).tiling(priority, soc_version) + +// op_type: 算子名称， class_name: 注册的 tiling 类, +// priority: tiling 类的优先级, 越小表示优先级越高, 即被选中的概率越大 +// 取代 REGISTER_TILING_TEMPLATE , 传入的op_type如果是字符串常量，需要去掉引号 +#define REGISTER_OPS_TILING_TEMPLATE(op_type, class_name, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::Register \ + __attribute__((unused)) tiling_##op_type##_##class_name##_##priority##_register = \ + Ops::Transformer::OpTiling::Register(#op_type).tiling(priority) diff --git a/csrc/ascend/common/include/tiling_base/tiling_type.h b/csrc/ascend/common/include/tiling_base/tiling_type.h new file mode 100644 index 000000000..bf1904173 --- /dev/null +++ b/csrc/ascend/common/include/tiling_base/tiling_type.h @@ -0,0 +1,139 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_type.h + * \brief + */ + +#pragma once + +#include + +namespace optiling { + +enum class AxisEnum { + B = 0, + N2 = 1, + G = 2, + S1 = 3, + S2 = 4, + D = 5, + NONE = 9, +}; + +enum class DtypeEnum { + FLOAT16 = 0, + FLOAT32 = 1, + BFLOAT16 = 2, + FLOAT16_PRECISION = 3, +}; + +enum class PerformanceOrientedEnum { + BIG_BUFFER = 1, + BIG_DOUBLE_BUFFER = 2, +}; + +enum class MatmulConfig { + NULL_CONFIG = 0, + NORMAL_CONFIG = 1, + MDL_CONFIG = 2 +}; + +enum class PseConfig { + NO_PSE = 0, + EXIST_PSE = 1 +}; + +enum class AttenMaskConfig { + NO_ATTEN_MASK = 0, + EXIST_ATTEN_MASK = 1 +}; + +enum class DropOutConfig { + NO_DROP_OUT = 0, + EXIST_DROP_OUT = 1 +}; + +enum class CubeFormatEnum { + ND = 0, + NZ = 1 +}; +enum class LayoutEnum { + BSND = 0, + SBND = 1, + BNSD = 2, + TND = 3, + NTD_TND = 4 +}; + +enum class CubeInputSourceEnum { + GM = 0, + L1 = 1 +}; + +enum class OptionEnum { + DISABLE = 0, + ENABLE = 1 +}; + +enum class SparseEnum { + ALL = 0, + NONE = 1, + ANY = 2, + CAUSAL = 3, + BAND = 4, + PREFIX = 5, + BAND_COMPRESS = 6, + RIGHT_DOWN_CAUSAL = 7, + RIGHT_DOWN_CAUSAL_BAND = 8, + BAND_LEFT_UP_CAUSAL = 9 +}; + +constexpr uint64_t RecursiveSum() +{ + return 0; +} + +constexpr int64_t base10Multiplier = 10; + +template constexpr uint64_t RecursiveSum(T templateId, Args... templateIds) +{ + return static_cast(templateId) + base10Multiplier * RecursiveSum(templateIds...); +} + +// TilingKey 的生成规则： +// FlashAttentionScore/FlashAttentionScoreGrad 十进制位组装tiling key，包含以下关键参数，从低位到高位依次是：Ub0, Ub1, +// Block, DataType, Format, Sparse, 特化模板 Ub0、Ub1: +// 表示Ub核内切分的轴，使用枚举AxisEnum表示，因为我们允许最多切分两根轴，所以存在UB0和UB1，如果没有UB核内切分， +// 那么填AXIS_NONE。UB0和UB1各占一个十进制位; +// Block: 表示UB用来分核的轴，使用枚举AxisEnum表示，占一个十进制位; +// DataType: 表示当前tiling key支持的输入输出的数据类型，使用枚举SupportedDtype来表示，占一个十进制位 +// Format: 表示当前tiling key支持的Format, 使用枚举InputLayout表示，占一个十进制位 +// Sparse: 表示当前tiling key是否支持Sparse，使用枚举SparseCapability表示，占一个十进制位 +// 其余特化场景，定义自己的位域和值 +// usage: get tilingKey from inputted types +// uint64_t tilingKey = GET_FLASHATTENTION_TILINGKEY(AxisEnum::AXIS_S1, AxisEnum::AXIS_S2, AxisEnum::AXIS_N2, +// SupportedDtype::FLOAT32, InputLayout::BSH, SparseCapability::SUPPORT_ALL) + +constexpr uint64_t TILINGKEYOFFSET = uint64_t(10000000000000000000UL); // 10^19 +template constexpr uint64_t GET_TILINGKEY(Args... templateIds) +{ + return TILINGKEYOFFSET + RecursiveSum(templateIds...); +} + +// usage: get tilingKey from inputted types +// uint64_t tilingKey = TILINGKEY(S2, S1, N2, FLOAT32, BSND, ALL) + +#define TILINGKEY(ub2, ub1, block, dtype, layout, sparse) \ + (GET_TILINGKEY(AxisEnum::ub2, AxisEnum::ub1, AxisEnum::block, DtypeEnum::dtype, LayoutEnum::layout, \ + SparseEnum::sparse)) + +} // namespace optiling diff --git a/csrc/ascend/common/include/tiling_base/tiling_util.h b/csrc/ascend/common/include/tiling_base/tiling_util.h new file mode 100644 index 000000000..85845e79f --- /dev/null +++ b/csrc/ascend/common/include/tiling_base/tiling_util.h @@ -0,0 +1,46 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_util.h + * \brief + */ + +#pragma once + +#include "register/op_impl_registry.h" + +namespace Ops { +namespace Transformer { + +template +T CeilAlign(T a, T b) +{ + return (a + b - 1) / b * b; +} + +template +T CeilDiv(T a, T b) +{ + if (b == 0) { + return a; + } + return (a + b - 1) / b; +} + +namespace OpTiling { +bool IsRegbaseSocVersion(const gert::TilingParseContext* context); + +bool IsRegbaseSocVersion(const gert::TilingContext* context); + +const gert::Shape& EnsureNotScalar(const gert::Shape& inShape); +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops diff --git a/csrc/ascend/common/include/tiling_sink/device_op_impl_registry_impl.h b/csrc/ascend/common/include/tiling_sink/device_op_impl_registry_impl.h new file mode 100644 index 000000000..994d4a817 --- /dev/null +++ b/csrc/ascend/common/include/tiling_sink/device_op_impl_registry_impl.h @@ -0,0 +1,49 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file device_op_impl_registry_impl.h + * \brief + */ + +#ifndef OP_TILING_DEVICE_OP_IMPL_REGISTRY_IMPL_H +#define OP_TILING_DEVICE_OP_IMPL_REGISTRY_IMPL_H + +#include +#include +#include "register/device_op_impl_registry.h" + +namespace optiling { +class DeviceOpImplRegistry { + public: + static DeviceOpImplRegistry& GetSingleton(); + void RegisterSinkTiling(std::string &opType, SinkTilingFunc& func); + SinkTilingFunc GetSinkTilingFunc(std::string &opType); + + private: + DeviceOpImplRegistry() = default; + ~DeviceOpImplRegistry() = default; + + private: + std::map sinkTilingFuncsMap_; +}; + +class DeviceOpImplRegisterImpl { + public: + DeviceOpImplRegisterImpl() = default; + ~DeviceOpImplRegisterImpl(); + std::string& GetOpType(); + + private: + std::string opType_ = ""; +}; +} // namespace optiling + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/include/tiling_sink/tiling_aicpu_task.h b/csrc/ascend/common/include/tiling_sink/tiling_aicpu_task.h new file mode 100644 index 000000000..24e013327 --- /dev/null +++ b/csrc/ascend/common/include/tiling_sink/tiling_aicpu_task.h @@ -0,0 +1,30 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_aicpu_task.h + * \brief + */ + +#ifndef TILING_SINK_TILING_AICPU_TASK_H_ +#define TILING_SINK_TILING_AICPU_TASK_H_ +#include "exe_graph/runtime/tiling_context.h" + +namespace tilingsink { +struct TilingAicpuTask { + gert::TilingContext *tilingContext; + const char *opType; + uint64_t notifyAddr; + uint64_t workspaceAddr; + uint64_t workspaceSize; +}; +} // namespace optiling + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/src/fallback_comm.cpp b/csrc/ascend/common/src/fallback_comm.cpp new file mode 100644 index 000000000..4385d9cff --- /dev/null +++ b/csrc/ascend/common/src/fallback_comm.cpp @@ -0,0 +1,53 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file fallback_comm.cpp + * \brief + */ + +#include "fallback/fallback_comm.h" + +#include +#include +#include +#include + +#include "aclnn/aclnn_base.h" +#include "runtime/base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +namespace fallback { +using namespace std; +using namespace gert; +using namespace ge; + +aclDataType ToAclDataType(ge::DataType dtype) { + static const std::vector CANN_CONVERT_TO_ACL_DataType_LIST = { + ge::DataType::DT_FLOAT, ge::DataType::DT_FLOAT16, ge::DataType::DT_INT8, ge::DataType::DT_INT32, + ge::DataType::DT_UINT8, ge::DataType::DT_INT16, ge::DataType::DT_UINT16, ge::DataType::DT_UINT32, + ge::DataType::DT_INT64, ge::DataType::DT_DOUBLE, ge::DataType::DT_BOOL, ge::DataType::DT_STRING, + ge::DataType::DT_COMPLEX64, ge::DataType::DT_COMPLEX128, ge::DataType::DT_BF16, ge::DataType::DT_UINT64, + ge::DataType::DT_INT4}; + auto iter = std::find(CANN_CONVERT_TO_ACL_DataType_LIST.begin(), CANN_CONVERT_TO_ACL_DataType_LIST.end(), dtype); + if (iter == CANN_CONVERT_TO_ACL_DataType_LIST.end()) { + return aclDataType::ACL_DT_UNDEFINED; + } + return static_cast(dtype); +} + +} // namespace fallback + +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/common/src/fallback_comm_2stages.cpp b/csrc/ascend/common/src/fallback_comm_2stages.cpp new file mode 100644 index 000000000..70de472ed --- /dev/null +++ b/csrc/ascend/common/src/fallback_comm_2stages.cpp @@ -0,0 +1,67 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file fallback_comm_2stages.cpp + * \brief + */ + +#include "fallback/fallback_comm_2stages.h" + +#include +#include +#include +#include + +#include "aclnn/aclnn_base.h" +#include "runtime/base.h" +#include "log/log.h" + +#ifdef __cplusplus +extern "C" { +#endif + +namespace fallback { +using namespace std; +using namespace gert; +using namespace ge; + +ge::graphStatus ExecuteOpLaunch(gert::OpExecuteLaunchContext *context) { + auto params = reinterpret_cast(context->GetOpApiParams()); + auto workspace_sizes = context->GetWorkspaceSizes(); + auto workspace_addrs = context->GetWorkspaceAddrs(); + OP_CHECK_IF((workspace_sizes->GetSize() == 0) || (workspace_addrs->GetSize() == 0), + OP_LOGE("aclnnfallback", "no workspace addrs"), return ge::GRAPH_FAILED); + auto workspace_size = workspace_sizes->GetData()[0]; + auto workspace_addr = workspace_addrs->GetData()[0]->GetAddr(); + + auto acl_stream = context->GetStream(); + auto opApiFunc = params->op_api_func; + OP_CHECK_IF(opApiFunc == nullptr, + OP_LOGE("aclnnfallback", "opApiFunc nullptr"), return ge::GRAPH_FAILED); + auto op_api_ret = opApiFunc(workspace_addr, workspace_size, params->executor, acl_stream); + for (auto &av : params->converted_params) { + if (av.deleter != nullptr) { + av.deleter(av.pointer); + } + } + params->converted_params.clear(); + if (op_api_ret != 0) { + OP_LOGE("aclnnfallback", "call %s allocate workspace failed op_api_ret: %d", context->GetNodeName(), op_api_ret); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +} // namespace fallback + +#ifdef __cplusplus +} +#endif \ No newline at end of file diff --git a/csrc/ascend/common/src/framework/CMakeLists.txt b/csrc/ascend/common/src/framework/CMakeLists.txt new file mode 100644 index 000000000..54a2c43c6 --- /dev/null +++ b/csrc/ascend/common/src/framework/CMakeLists.txt @@ -0,0 +1,14 @@ +# ---------------------------------------------------------------------------- +# This program is free software, you can redistribute it and/or modify. +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This file is a part of the CANN Open Software. +# Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING +# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of +# the software repository for the full text of the License. +# ---------------------------------------------------------------------------- + +if (BUILD_OPEN_PROJECT AND NOT BUILD_OPS_RTY_KERNEL) + add_onnx_plugin_sources() +endif() \ No newline at end of file diff --git a/csrc/ascend/common/src/framework/embedding_bag_onnx_plugin.cpp b/csrc/ascend/common/src/framework/embedding_bag_onnx_plugin.cpp new file mode 100644 index 000000000..ef2153155 --- /dev/null +++ b/csrc/ascend/common/src/framework/embedding_bag_onnx_plugin.cpp @@ -0,0 +1,78 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file embedding_bag_onnx_plugin.cpp + * \brief + */ + +#include "onnx_common.h" + +namespace domi { +using NodeProto = ge::onnx::NodeProto; + +static Status ParseParamsEmbeddingBag(const Message *op_src, ge::Operator &op_dest) { + const NodeProto *node = dynamic_cast(op_src); + if (node == nullptr) { + OP_LOGE(GetOpName(op_dest), "Dynamic cast op_src to NodeProto failed."); + return FAILED; + } + // set attr mode_value + std::string mode_value; + for (const auto &attr : node->attribute()) { + if (attr.name() == "mode" && attr.type() == ge::onnx::AttributeProto::STRING) { + mode_value = attr.s(); + op_dest.SetAttr("mode", mode_value); + } + } + // set attr scale_grad_by_freq + bool scale_grad_by_freq = false; + for (const auto &attr : node->attribute()) { + if (attr.name() == "scale_grad_by_freq" && attr.i() != 0) { + scale_grad_by_freq = true; + break; + } + } + op_dest.SetAttr("scale_grad_by_freq", scale_grad_by_freq); + // set attr sparse + bool sparse = false; + for (const auto &attr : node->attribute()) { + if (attr.name() == "sparse" && attr.i() != 0) { + sparse = true; + break; + } + } + op_dest.SetAttr("sparse", sparse); + // set attr include_last_offset + bool include_last_offset = false; + for (const auto &attr : node->attribute()) { + if (attr.name() == "include_last_offset" && attr.i() != 0) { + include_last_offset = true; + break; + } + } + op_dest.SetAttr("include_last_offset", include_last_offset); + return SUCCESS; +} + +REGISTER_CUSTOM_OP("EmbeddingBag") + .FrameworkType(ONNX) + .OriginOpType({ge::AscendString("ai.onnx::8::EmbeddingBag"), + ge::AscendString("ai.onnx::9::EmbeddingBag"), + ge::AscendString("ai.onnx::10::EmbeddingBag"), + ge::AscendString("ai.onnx::11::EmbeddingBag"), + ge::AscendString("ai.onnx::12::EmbeddingBag"), + ge::AscendString("ai.onnx::13::EmbeddingBag"), + ge::AscendString("ai.onnx::14::EmbeddingBag"), + ge::AscendString("ai.onnx::15::EmbeddingBag"), + ge::AscendString("ai.onnx::16::EmbeddingBag")}) + .ParseParamsFn(ParseParamsEmbeddingBag) + .ImplyType(ImplyType::TVM); +} // namespace domi diff --git a/csrc/ascend/common/src/framework/fillwindowcache_onnx_plugin.cpp b/csrc/ascend/common/src/framework/fillwindowcache_onnx_plugin.cpp new file mode 100644 index 000000000..75d5014f1 --- /dev/null +++ b/csrc/ascend/common/src/framework/fillwindowcache_onnx_plugin.cpp @@ -0,0 +1,58 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file fillwindowcache_onnx_plugin.cpp + * \brief + */ + +#include "onnx_common.h" + +namespace domi { +static Status parseParamsFillWindowCache(const Message* op_src, ge::Operator& op_dest) +{ + const ge::onnx::NodeProto *node = reinterpret_cast(op_src); + if (node == nullptr) { + OP_LOGE(GetOpName(op_dest), "Dynamic fillwindowcache op_src to NodeProto failed."); + return FAILED; + } + + int axis = 0; + int cache_depth = 0; + for (const auto &attr : node->attribute()) { + if (attr.name() == "axis") { + axis = attr.i(); + } + if (attr.name() == "cache_depth") { + cache_depth = attr.i(); + } + } + op_dest.SetAttr("axis", axis); + op_dest.SetAttr("cache_depth", cache_depth); + return SUCCESS; +} + +// register FillWindowCache op info to GE +REGISTER_CUSTOM_OP("FillWindowCache") + .FrameworkType(ONNX) + .OriginOpType({ge::AscendString("ai.onnx::8::FillWindowCache"), + ge::AscendString("ai.onnx::9::FillWindowCache"), + ge::AscendString("ai.onnx::10::FillWindowCache"), + ge::AscendString("ai.onnx::11::FillWindowCache"), + ge::AscendString("ai.onnx::12::FillWindowCache"), + ge::AscendString("ai.onnx::13::FillWindowCache"), + ge::AscendString("ai.onnx::14::FillWindowCache"), + ge::AscendString("ai.onnx::15::FillWindowCache"), + ge::AscendString("ai.onnx::16::FillWindowCache"), + ge::AscendString("ai.onnx::17::FillWindowCache"), + ge::AscendString("ai.onnx::18::FillWindowCache")}) + .ParseParamsFn(parseParamsFillWindowCache) + .ImplyType(ImplyType::TVM); +} // namespace domi diff --git a/csrc/ascend/common/src/framework/multi_head_attention_onnx_plugin.cpp b/csrc/ascend/common/src/framework/multi_head_attention_onnx_plugin.cpp new file mode 100644 index 000000000..42b40f24d --- /dev/null +++ b/csrc/ascend/common/src/framework/multi_head_attention_onnx_plugin.cpp @@ -0,0 +1,84 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file multi_head_attention_onnx_plugin.cpp + * \brief + */ + +#include "onnx_common.h" + +namespace domi { +using NodeProto = ge::onnx::NodeProto; +static const int REQUIRED_ATTRS_NUM = 6; +static Status ParseParamsMultiHeadAttention(const Message* op_src, ge::Operator& op_dest) { + const NodeProto* node = dynamic_cast(op_src); + if (node == nullptr) { + OP_LOGE("MultiHeadAttention", "Dynamic cast op_src to NodeProto failed."); + return FAILED; + } + int attn_head_num = 0; + int attn_dim_per_head = 0; + int src_len = 0; + int tgt_len = 0; + float dropout_prob = 0.0f; + int softmax_use_float = 0; + int attr_num = 0; + for (const auto& attr : node->attribute()) { + if (attr.name() == "attn_head_num" && attr.type() == ge::onnx::AttributeProto::INT) { + attn_head_num = attr.i(); + ++attr_num; + } else if (attr.name() == "attn_dim_per_head" && attr.type() == ge::onnx::AttributeProto::INT) { + attn_dim_per_head = attr.i(); + ++attr_num; + } else if (attr.name() == "src_len" && attr.type() == ge::onnx::AttributeProto::INT) { + src_len = attr.i(); + ++attr_num; + } else if (attr.name() == "tgt_len" && attr.type() == ge::onnx::AttributeProto::INT) { + tgt_len = attr.i(); + ++attr_num; + } else if (attr.name() == "dropout_prob" && attr.type() == ge::onnx::AttributeProto::FLOAT) { + dropout_prob = attr.f(); + ++attr_num; + } else if (attr.name() == "softmax_use_float" && attr.type() == ge::onnx::AttributeProto::INT) { + softmax_use_float = attr.i(); + ++attr_num; + } + } + + if (attr_num != REQUIRED_ATTRS_NUM) { + OP_LOGE(GetOpName(op_dest).c_str(), "Node must have attrs attn_head_num/attn_dim_per_head/" + "src_len/tgt_len/dropout_prob/softmax_use_float"); + return FAILED; + } + op_dest.SetAttr("attn_head_num", attn_head_num); + op_dest.SetAttr("attn_dim_per_head", attn_dim_per_head); + op_dest.SetAttr("src_len", src_len); + op_dest.SetAttr("tgt_len", tgt_len); + op_dest.SetAttr("keep_prob", static_cast(1 - dropout_prob)); + op_dest.SetAttr("softmax_use_float", static_cast(softmax_use_float)); + return SUCCESS; +} + +// register Yolo op info to GE +REGISTER_CUSTOM_OP("MultiHeadAttention") + .FrameworkType(ONNX) + .OriginOpType({ge::AscendString("ai.onnx::11::NPUMultiHeadAttention"), + ge::AscendString("ai.onnx::12::NPUMultiHeadAttention"), + ge::AscendString("ai.onnx::13::NPUMultiHeadAttention"), + ge::AscendString("ai.onnx::14::NPUMultiHeadAttention"), + ge::AscendString("ai.onnx::15::NPUMultiHeadAttention"), + ge::AscendString("ai.onnx::16::NPUMultiHeadAttention"), + ge::AscendString("ai.onnx::17::NPUMultiHeadAttention"), + ge::AscendString("ai.onnx::18::NPUMultiHeadAttention"), + ge::AscendString("npu::1::NPUMultiHeadAttention")}) + .ParseParamsFn(ParseParamsMultiHeadAttention) + .ImplyType(ImplyType::TVM); +} // namespace domi diff --git a/csrc/ascend/common/src/framework/npu_fused_attention_score_fwd_onnx_plugin.cpp b/csrc/ascend/common/src/framework/npu_fused_attention_score_fwd_onnx_plugin.cpp new file mode 100644 index 000000000..039e1592f --- /dev/null +++ b/csrc/ascend/common/src/framework/npu_fused_attention_score_fwd_onnx_plugin.cpp @@ -0,0 +1,176 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "onnx_common.h" +#include "op_transformer_proto_extend.h" + +namespace domi { +using NodeProto = ge::onnx::NodeProto; +constexpr int REQUIRED_ATTR = 1; +constexpr int ONE = 1; +constexpr int INDEX_TWO = 2; +constexpr int INDEX_THREE = 3; +constexpr int ACL_FLOAT16 = 1; + +static Status ParseParamsNpuFusedAttentionScoreFwd(const Message *op_src, ge::Operator &op_dest) { + const NodeProto *node = dynamic_cast(op_src); + if (node == nullptr) { + OP_LOGE(GetOpName(op_dest), "Dynamic cast op_src to NodeProto failed."); + return FAILED; + } + + int input_size = node->input_size(); + int output_size = node->output_size(); + op_dest.DynamicInputRegister("x", input_size); + op_dest.DynamicOutputRegister("y", output_size); + + int required_attr_num = 0; + float scale = 0; + float keep_prob = 1.; + bool query_transpose = false; + bool key_transpose = false; + bool bmm_score_transpose_a = false; + bool bmm_score_transpose_b = false; + + for (const auto &attr : node->attribute()) { + if (attr.name() == "scale" && attr.type() == ge::onnx::AttributeProto::FLOAT) { + scale = attr.f(); + required_attr_num++; + } else if (attr.name() == "keep_prob" && attr.type() == ge::onnx::AttributeProto::FLOAT) { + keep_prob = attr.f(); + } else if (attr.name() == "query_transpose" && attr.type() == ge::onnx::AttributeProto::INT) { + query_transpose = (attr.i() == 1); + } else if (attr.name() == "key_transpose" && attr.type() == ge::onnx::AttributeProto::INT) { + key_transpose = (attr.i() == 1); + } else if (attr.name() == "bmm_score_transpose_a" && attr.type() == ge::onnx::AttributeProto::INT) { + bmm_score_transpose_a = (attr.i() == 1); + } else if (attr.name() == "bmm_score_transpose_b" && attr.type() == ge::onnx::AttributeProto::INT) { + bmm_score_transpose_b = (attr.i() == 1); + } + } + + if (required_attr_num != REQUIRED_ATTR) { + OP_LOGE(GetOpName(op_dest), "attr scale is required."); + return FAILED; + } + + op_dest.SetAttr("name", node->name()); + op_dest.SetAttr("scale", scale); + op_dest.SetAttr("keep_prob", keep_prob); + op_dest.SetAttr("query_transpose", query_transpose); + op_dest.SetAttr("key_transpose", key_transpose); + op_dest.SetAttr("bmm_score_transpose_a", bmm_score_transpose_a); + op_dest.SetAttr("bmm_score_transpose_b", bmm_score_transpose_b); + op_dest.SetAttr("original_type", "npu::1::NPUFusedAttentionScoreFwd"); + return SUCCESS; +} + +namespace{ +static Status GetAttrFromOperator(const ge::Operator& op, float& scale, float& keep_prob, bool& query_transpose, + bool& key_transpose, bool& bmm_score_transpose_a, bool& bmm_score_transpose_b) { + if (op.GetAttr("scale", scale) != SUCCESS) { + OP_LOGE(GetOpName(op), "get scale from op failed"); + return FAILED; + } + if (op.GetAttr("keep_prob", keep_prob) != SUCCESS) { + OP_LOGE(GetOpName(op), "get keep_prob from op failed"); + return FAILED; + } + if (op.GetAttr("query_transpose", query_transpose) != SUCCESS) { + OP_LOGE(GetOpName(op), "get query_transpose from op failed"); + return FAILED; + } + if (op.GetAttr("key_transpose", key_transpose) != SUCCESS) { + OP_LOGE(GetOpName(op), "get key_transpose from op failed"); + return FAILED; + } + if (op.GetAttr("bmm_score_transpose_a", bmm_score_transpose_a) != SUCCESS) { + OP_LOGE(GetOpName(op), "get bmm_score_transpose_a from op failed"); + return FAILED; + } + if (op.GetAttr("bmm_score_transpose_b", bmm_score_transpose_b) != SUCCESS) { + OP_LOGE(GetOpName(op), "get bmm_score_transpose_b from op failed"); + return FAILED; + } + return SUCCESS; + } +} + +static Status ParseOpToGraphNpuFusedAttentionScoreFwd(const ge::Operator& op, ge::Graph& graph) { + std::string ori_name; + if (op.GetAttr("name", ori_name) != SUCCESS) { + OP_LOGE(GetOpName(op), "get name from op failed."); + return FAILED; + } + + auto data0 = ge::op::Data((ori_name + "_data0").c_str()).set_attr_index(0); + auto data1 = ge::op::Data((ori_name + "_data1").c_str()).set_attr_index(1); + auto data2 = ge::op::Data((ori_name + "_data2").c_str()).set_attr_index(2); + auto data3 = ge::op::Data((ori_name + "_data3").c_str()).set_attr_index(3); + + float scale = 0; + float keep_prob = 0; + bool query_transpose = false; + bool key_transpose = false; + bool bmm_score_transpose_a = false; + bool bmm_score_transpose_b = false; + Status ret = GetAttrFromOperator( + op, scale, keep_prob, query_transpose, key_transpose, bmm_score_transpose_a, bmm_score_transpose_b); + if (ret != SUCCESS) { + return FAILED; + } + // create const input tensor "drop_mask" which is filled with the scalar value 1 for inferencing + // deop_mask.size = {query_size[0], query_size[1], query_size[2], query_size[2]} + ge::Tensor saclar_one = CreateScalar(ONE, ge::DT_UINT8); + auto const_one = ge::op::Const((ori_name + "_Const_one").c_str()).set_attr_value(saclar_one); + std::vector dims = op.GetInputDesc(0).GetShape().GetDims(); + dims[INDEX_THREE] = dims[INDEX_TWO]; + auto tensor_dims = Vec2Tensor(dims, {4}, ge::DT_INT64); + auto const_dims = ge::op::Const((ori_name + "_Const_dims").c_str()).set_attr_value(tensor_dims); + auto drop_mask = ge::op::Fill((ori_name + "_Fill_ones").c_str()).set_input_dims(const_dims) + .set_input_value(const_one); + + ge::Tensor tensor_scale = CreateScalar(scale, ge::DT_FLOAT); + auto const_scale = ge::op::Const((ori_name + "_Const_scale").c_str()).set_attr_value(tensor_scale); + auto cast_const_scale = ge::op::Cast((ori_name + "_Cast_const_scale").c_str()).set_input_x(const_scale) + .set_attr_dst_type(ACL_FLOAT16); + + auto AttentionScore = ge::op::AttentionScore((ori_name + "_AttentionScore").c_str()).set_input_query(data0) + .set_input_key(data1).set_input_value(data2).set_input_padding_mask(data3) + .set_input_scale(cast_const_scale).set_input_drop_mask(drop_mask) + .set_attr_keep_prob(keep_prob).set_attr_query_transpose(query_transpose) + .set_attr_key_transpose(key_transpose).set_attr_bmm_score_transpose_a(bmm_score_transpose_a) + .set_attr_bmm_score_transpose_b(bmm_score_transpose_b).set_attr_softmax_axes({-1}); + + std::vector inputs{data0, data1, data2, data3}; + std::vector>> outputs; + outputs.emplace_back(AttentionScore, std::vector{0}); + outputs.emplace_back(AttentionScore, std::vector{1}); + outputs.emplace_back(drop_mask, std::vector{0}); + graph.SetInputs(inputs).SetOutputs(outputs); + return SUCCESS; +} + +// register npu_fused_attention_score_fwd op info to GE +REGISTER_CUSTOM_OP("PartitionedCall") + .FrameworkType(ONNX) + .OriginOpType({ge::AscendString("npu::1::NPUFusedAttentionScoreFwd"), + ge::AscendString("ai.onnx::11::NPUFusedAttentionScoreFwd"), + ge::AscendString("ai.onnx::12::NPUFusedAttentionScoreFwd"), + ge::AscendString("ai.onnx::13::NPUFusedAttentionScoreFwd"), + ge::AscendString("ai.onnx::14::NPUFusedAttentionScoreFwd"), + ge::AscendString("ai.onnx::15::NPUFusedAttentionScoreFwd"), + ge::AscendString("ai.onnx::16::NPUFusedAttentionScoreFwd"), + ge::AscendString("ai.onnx::17::NPUFusedAttentionScoreFwd"), + ge::AscendString("ai.onnx::18::NPUFusedAttentionScoreFwd")}) + .ParseParamsFn(ParseParamsNpuFusedAttentionScoreFwd) + .ParseOpToGraphFn(ParseOpToGraphNpuFusedAttentionScoreFwd) + .ImplyType(ImplyType::TVM); +} // namespace domi diff --git a/csrc/ascend/common/src/framework/npu_fused_attention_score_onnx_plugin.cpp b/csrc/ascend/common/src/framework/npu_fused_attention_score_onnx_plugin.cpp new file mode 100644 index 000000000..c98612c24 --- /dev/null +++ b/csrc/ascend/common/src/framework/npu_fused_attention_score_onnx_plugin.cpp @@ -0,0 +1,185 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "onnx_common.h" +#include "op_transformer_proto_extend.h" + +namespace domi { +using NodeProto = ge::onnx::NodeProto; +constexpr int REQUIRED_ATTR = 1; +constexpr int ONE = 1; +constexpr int INDEX_TWO = 2; +constexpr int INDEX_THREE = 3; +constexpr int ACL_FLOAT16 = 1; + +static Status ParseParamsNpuFusedAttentionScore(const Message *op_src, ge::Operator &op_dest) { + const NodeProto *node = dynamic_cast(op_src); + if (node == nullptr) { + OP_LOGE(GetOpName(op_dest).c_str(), "Dynamic cast op_src to NodeProto failed."); + return FAILED; + } + + int input_size = node->input_size(); + int output_size = node->output_size(); + op_dest.DynamicInputRegister("x", input_size); + op_dest.DynamicOutputRegister("y", output_size); + + int required_attr_num = 0; + float scale = 0; + float keep_prob = 1.; + bool query_transpose = false; + bool key_transpose = false; + bool bmm_score_transpose_a = false; + bool bmm_score_transpose_b = false; + + for (const auto &attr : node->attribute()) { + if (attr.name() == "scale" && attr.type() == ge::onnx::AttributeProto::FLOAT) { + scale = attr.f(); + required_attr_num++; + } else if (attr.name() == "keep_prob" && attr.type() == ge::onnx::AttributeProto::FLOAT) { + keep_prob = attr.f(); + } else if (attr.name() == "query_transpose" && attr.type() == ge::onnx::AttributeProto::INT) { + query_transpose = (attr.i() == 1); + } else if (attr.name() == "key_transpose" && attr.type() == ge::onnx::AttributeProto::INT) { + key_transpose = (attr.i() == 1); + } else if (attr.name() == "bmm_score_transpose_a" && attr.type() == ge::onnx::AttributeProto::INT) { + bmm_score_transpose_a = (attr.i() == 1); + } else if (attr.name() == "bmm_score_transpose_b" && attr.type() == ge::onnx::AttributeProto::INT) { + bmm_score_transpose_b = (attr.i() == 1); + } + } + + if (required_attr_num != REQUIRED_ATTR) { + OP_LOGE(GetOpName(op_dest).c_str(), "attr scale is required."); + return FAILED; + } + + op_dest.SetAttr("name", node->name()); + op_dest.SetAttr("scale", scale); + op_dest.SetAttr("keep_prob", keep_prob); + op_dest.SetAttr("query_transpose", query_transpose); + op_dest.SetAttr("key_transpose", key_transpose); + op_dest.SetAttr("bmm_score_transpose_a", bmm_score_transpose_a); + op_dest.SetAttr("bmm_score_transpose_b", bmm_score_transpose_b); + op_dest.SetAttr("original_type", "npu::1::NPUFusedAttentionScore"); + return SUCCESS; +} + +namespace { +static Status GetAttrFromPre3(const ge::Operator& op, float& scale, float& keep_prob, bool& query_transpose) { + if (op.GetAttr("scale", scale) != SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "get scale from op failed"); + return FAILED; + } + if (op.GetAttr("keep_prob", keep_prob) != SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "get keep_prob from op failed"); + return FAILED; + } + if (op.GetAttr("query_transpose", query_transpose) != SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "get query_transpose from op failed"); + return FAILED; + } + return SUCCESS; +} + +static Status GetAttrFromLast3( + const ge::Operator& op, bool& key_transpose, bool& bmm_score_transpose_a, bool& bmm_score_transpose_b) { + if (op.GetAttr("key_transpose", key_transpose) != SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "get key_transpose from op failed"); + return FAILED; + } + if (op.GetAttr("bmm_score_transpose_a", bmm_score_transpose_a) != SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "get bmm_score_transpose_a from op failed"); + return FAILED; + } + if (op.GetAttr("bmm_score_transpose_b", bmm_score_transpose_b) != SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "get bmm_score_transpose_b from op failed"); + return FAILED; + } + return SUCCESS; +} +} // namespace + +static Status ParseOpToGraphNpuFusedAttentionScore(const ge::Operator& op, ge::Graph& graph) { + std::string ori_name; + if (op.GetAttr("name", ori_name) != SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "get name from op failed."); + return FAILED; + } + + auto data0 = ge::op::Data((ori_name + "_data0").c_str()).set_attr_index(0); + auto data1 = ge::op::Data((ori_name + "_data1").c_str()).set_attr_index(1); + auto data2 = ge::op::Data((ori_name + "_data2").c_str()).set_attr_index(2); + auto data3 = ge::op::Data((ori_name + "_data3").c_str()).set_attr_index(3); + + float scale = 0; + float keep_prob = 0; + bool query_transpose = false; + if (GetAttrFromPre3(op, scale, keep_prob, query_transpose) != SUCCESS) { + return FAILED; + } + bool key_transpose = false; + bool bmm_score_transpose_a = false; + bool bmm_score_transpose_b = false; + if (GetAttrFromLast3(op, key_transpose, bmm_score_transpose_a, bmm_score_transpose_b) != SUCCESS) { + return FAILED; + } + + // create const input tensor "drop_mask" which is filled with the scalar value 1 for inferencing + // deop_mask.size = {query_size[0], query_size[1], query_size[2], query_size[2]} + ge::Tensor saclar_one = CreateScalar(ONE, ge::DT_UINT8); + auto const_one = ge::op::Const((ori_name + "_Const_one").c_str()).set_attr_value(saclar_one); + std::vector dims = op.GetInputDesc(0).GetShape().GetDims(); + dims[INDEX_THREE] = dims[INDEX_TWO]; + auto tensor_dims = Vec2Tensor(dims, {4}, ge::DT_INT64); + auto const_dims = ge::op::Const((ori_name + "_Const_dims").c_str()).set_attr_value(tensor_dims); + auto drop_mask = ge::op::Fill((ori_name + "_Fill_ones").c_str()).set_input_dims(const_dims) + .set_input_value(const_one); + // set drop_mask["dtype"] to fp16 for inferencing + auto cast_drop_mask = ge::op::Cast((ori_name + "_Cast_drop_mask").c_str()).set_input_x(drop_mask) + .set_attr_dst_type(ACL_FLOAT16); + + ge::Tensor tensor_scale = CreateScalar(scale, ge::DT_FLOAT); + auto const_scale = ge::op::Const((ori_name + "_Const_scale").c_str()).set_attr_value(tensor_scale); + auto cast_const_scale = ge::op::Cast((ori_name + "_Cast_const_scale").c_str()).set_input_x(const_scale) + .set_attr_dst_type(ACL_FLOAT16); + + auto AttentionScore = ge::op::AttentionScore((ori_name + "_AttentionScore").c_str()) + .set_input_query(data0).set_input_key(data1).set_input_value(data2) + .set_input_padding_mask(data3).set_input_scale(cast_const_scale) + .set_input_drop_mask(cast_drop_mask).set_attr_keep_prob(keep_prob) + .set_attr_query_transpose(query_transpose).set_attr_key_transpose(key_transpose) + .set_attr_bmm_score_transpose_a(bmm_score_transpose_a) + .set_attr_bmm_score_transpose_b(bmm_score_transpose_b) + .set_attr_softmax_axes({-1}); + + std::vector inputs{data0, data1, data2, data3}; + std::vector>> outputs; + outputs.emplace_back(AttentionScore, std::vector{0}); + graph.SetInputs(inputs).SetOutputs(outputs); + return SUCCESS; +} + +// register npu_fused_attention_score op info to GE +REGISTER_CUSTOM_OP("PartitionedCall") + .FrameworkType(ONNX) + .OriginOpType({ge::AscendString("npu::1::NPUFusedAttentionScore"), + ge::AscendString("ai.onnx::11::NPUFusedAttentionScore"), + ge::AscendString("ai.onnx::12::NPUFusedAttentionScore"), + ge::AscendString("ai.onnx::13::NPUFusedAttentionScore"), + ge::AscendString("ai.onnx::14::NPUFusedAttentionScore"), + ge::AscendString("ai.onnx::15::NPUFusedAttentionScore"), + ge::AscendString("ai.onnx::16::NPUFusedAttentionScore"), + ge::AscendString("ai.onnx::17::NPUFusedAttentionScore"), + ge::AscendString("ai.onnx::18::NPUFusedAttentionScore")}) + .ParseParamsFn(ParseParamsNpuFusedAttentionScore) + .ParseOpToGraphFn(ParseOpToGraphNpuFusedAttentionScore) + .ImplyType(ImplyType::TVM); +} // namespace domi diff --git a/csrc/ascend/common/src/framework/npu_masked_softmax_with_relposbias_onnx_plugin.cpp b/csrc/ascend/common/src/framework/npu_masked_softmax_with_relposbias_onnx_plugin.cpp new file mode 100644 index 000000000..0ae50f532 --- /dev/null +++ b/csrc/ascend/common/src/framework/npu_masked_softmax_with_relposbias_onnx_plugin.cpp @@ -0,0 +1,102 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file npu_masked_softmax_with_relposbias_onnx_plugin.cpp + * \brief + */ + +#include "onnx_common.h" +#include "op_transformer_proto_extend.h" + +namespace domi { +using NodeProto = ge::onnx::NodeProto; +constexpr int OUTPUT_INDEX = 1; +// x, atten_mask, relative_pos_bias, scale_value, inner_precision_mode float scale_value=1.0, int inner_precision_mode=0 +static Status ParseParamsNpuMaskedSoftmaxWithRelPosBias(const Message* op_src, ge::Operator& op_dest) { + const NodeProto* node = dynamic_cast(op_src); + if (node == nullptr) { + OP_LOGE("MaskedSoftmaxWithRelPosBias", "Dynamic cast op_src to NodeProto failed."); + return FAILED; + } + + int input_size = node->input_size(); + int output_size = node->output_size(); + op_dest.DynamicInputRegister("x", input_size); + op_dest.DynamicOutputRegister("y", output_size); + + float scale_value = 0; + int inner_precision_mode = 1.0; + for (const auto& attr : node->attribute()) { + if (attr.name() == "inner_precision_mode" && attr.type() == ge::onnx::AttributeProto::INT) { + inner_precision_mode = attr.i(); + } else if (attr.name() == "scale_value" && attr.type() == ge::onnx::AttributeProto::FLOAT) { + scale_value = attr.f(); + } + } + op_dest.SetAttr("name", node->name()); + op_dest.SetAttr("inner_precision_mode", inner_precision_mode); + op_dest.SetAttr("scale_value", scale_value); + op_dest.SetAttr("original_type", "npu::1::NPUMaskedSoftmaxWithRelPosBias"); + return SUCCESS; +} + +static Status ParseOpToGraphNpuMaskedSoftmaxWithRelPosBias(const ge::Operator& op, ge::Graph& graph) { + std::string ori_name; + if (op.GetAttr("name", ori_name) != SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "get name from op failed."); + return FAILED; + } + + auto data0 = ge::op::Data((ori_name + "_data0").c_str()).set_attr_index(0); + auto data1 = ge::op::Data((ori_name + "_data1").c_str()).set_attr_index(1); + auto data2 = ge::op::Data((ori_name + "_data2").c_str()).set_attr_index(2); + + int inner_precision_mode = 0; + if (op.GetAttr("inner_precision_mode", inner_precision_mode) != SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "get inner_precision_mode from op failed"); + return FAILED; + } + float scale_value = 1.0f; + if (op.GetAttr("scale_value", scale_value) != SUCCESS) { + OP_LOGE(GetOpName(op).c_str(), "get scale_value from op failed"); + return FAILED; + } + auto masked_softmax_with_relposbias = ge::op::MaskedSoftmaxWithRelPosBias((ori_name + "_MaskedSoftmaxWithRelPosBias").c_str()) + .set_input_x(data0) + .set_input_atten_mask(data1) + .set_input_relative_pos_bias(data2) + .set_attr_scale_value(scale_value) + .set_attr_inner_precision_mode(inner_precision_mode); + + std::vector inputs{ data0, data1, data2 }; + std::vector>> outputs; + outputs.emplace_back(masked_softmax_with_relposbias, std::vector{OUTPUT_INDEX}); + graph.SetInputs(inputs).SetOutputs(outputs); + return SUCCESS; +} + +// register npu_masked_softmax_with_rel_pos_bias op info to GE +REGISTER_CUSTOM_OP("PartitionedCall") + .FrameworkType(ONNX) + .OriginOpType({ge::AscendString("npu::1::NPUMaskedSoftmaxWithRelPosBias"), + ge::AscendString("ai.onnx::11::NPUMaskedSoftmaxWithRelPosBias"), + ge::AscendString("ai.onnx::12::NPUMaskedSoftmaxWithRelPosBias"), + ge::AscendString("ai.onnx::13::NPUMaskedSoftmaxWithRelPosBias"), + ge::AscendString("ai.onnx::14::NPUMaskedSoftmaxWithRelPosBias"), + ge::AscendString("ai.onnx::15::NPUMaskedSoftmaxWithRelPosBias"), + ge::AscendString("ai.onnx::16::NPUMaskedSoftmaxWithRelPosBias"), + ge::AscendString("ai.onnx::17::NPUMaskedSoftmaxWithRelPosBias"), + ge::AscendString("ai.onnx::18::NPUMaskedSoftmaxWithRelPosBias"), + ge::AscendString("ai.onnx::19::NPUMaskedSoftmaxWithRelPosBias")}) + .ParseParamsFn(ParseParamsNpuMaskedSoftmaxWithRelPosBias) + .ParseOpToGraphFn(ParseOpToGraphNpuMaskedSoftmaxWithRelPosBias) + .ImplyType(ImplyType::TVM); +} // namespace domi \ No newline at end of file diff --git a/csrc/ascend/common/src/framework/npu_scaled_masked_softmax_onnx_plugin.cpp b/csrc/ascend/common/src/framework/npu_scaled_masked_softmax_onnx_plugin.cpp new file mode 100644 index 000000000..e5b8f2bca --- /dev/null +++ b/csrc/ascend/common/src/framework/npu_scaled_masked_softmax_onnx_plugin.cpp @@ -0,0 +1,57 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file npu_scaled_masked_softmax_onnx_plugin.cpp + * \brief onnx plugin for npu custom operator npu_scaled_masked_softmax + */ + +#include "onnx_common.h" + +namespace domi { +using NodeProto = ge::onnx::NodeProto; + +static Status ParseParamsNPUScaledMaskedSoftmax(const Message* op_src, ge::Operator& op_dest) { + const NodeProto* node = dynamic_cast(op_src); + if (node == nullptr) { + OP_LOGE(GetOpName(op_dest).c_str(), "Dynamic cast op_src to NodeProto failed!"); + return FAILED; + } + + float scale = 1.0; + bool fixed_triu_mask = false; + for (const auto& attr : node->attribute()) { + if (attr.name() == "scale" && attr.type() == ge::onnx::AttributeProto::FLOAT) { + scale = attr.f(); + } else if (attr.name() == "fixed_triu_mask" && attr.type() == ge::onnx::AttributeProto::INT && attr.i() == 1) { + fixed_triu_mask = true; + } + } + op_dest.SetAttr("scale", scale); + op_dest.SetAttr("fixed_triu_mask", fixed_triu_mask); + + return SUCCESS; +} + +// register npu_scaled_masked_softmax op info to GE +REGISTER_CUSTOM_OP("ScaledMaskedSoftmax") + .FrameworkType(ONNX) + .OriginOpType({ge::AscendString("npu::1::NPUScaledMaskedSoftmax"), + ge::AscendString("ai.onnx::11::NPUScaledMaskedSoftmax"), + ge::AscendString("ai.onnx::12::NPUScaledMaskedSoftmax"), + ge::AscendString("ai.onnx::13::NPUScaledMaskedSoftmax"), + ge::AscendString("ai.onnx::14::NPUScaledMaskedSoftmax"), + ge::AscendString("ai.onnx::15::NPUScaledMaskedSoftmax"), + ge::AscendString("ai.onnx::16::NPUScaledMaskedSoftmax"), + ge::AscendString("ai.onnx::17::NPUScaledMaskedSoftmax"), + ge::AscendString("ai.onnx::18::NPUScaledMaskedSoftmax")}) + .ParseParamsFn(ParseParamsNPUScaledMaskedSoftmax) + .ImplyType(ImplyType::TVM); +} // namespace domi \ No newline at end of file diff --git a/csrc/ascend/common/src/framework/tfidf_vectorizer_onnx_plugin.cpp b/csrc/ascend/common/src/framework/tfidf_vectorizer_onnx_plugin.cpp new file mode 100644 index 000000000..2691b0443 --- /dev/null +++ b/csrc/ascend/common/src/framework/tfidf_vectorizer_onnx_plugin.cpp @@ -0,0 +1,62 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "onnx_common.h" + +namespace domi { + +static Status ParseParamsTfIdfVectorizer(const Message* op_src, ge::Operator& op_dest) { + const ge::onnx::NodeProto* node = dynamic_cast(op_src); + if (node == nullptr) { + std::string reportErrorCode = "E50058"; + std::vector errKeys = {"op_name", "description"}; + std::vector errValues = {"TfIdfVectorizer", "Dynamic cast op_src to NodeProto failed!"}; + ge::ReportPredefinedErrMsg(reportErrorCode.c_str(), errKeys, errValues); + return FAILED; + } + + for (const auto& attr : node->attribute()) { + const std::string& name = attr.name(); + if (attr.type() == ge::onnx::AttributeProto::INT) { + if (name == "max_gram_length" || name == "max_skip_count" || name == "min_gram_length") { + op_dest.SetAttr(name.c_str(), attr.i()); + } + } else if (name == "mode" && attr.type() == ge::onnx::AttributeProto::STRING) { + op_dest.SetAttr(name.c_str(), attr.s()); + } else if (name == "weights" && attr.type() == ge::onnx::AttributeProto::FLOATS) { + op_dest.SetAttr(name.c_str(), std::vector{attr.floats().begin(), attr.floats().end()}); + } else if (attr.type() == ge::onnx::AttributeProto::INTS) { + if (name == "ngram_counts" || name == "ngram_indexes" || name == "pool_int64s") { + op_dest.SetAttr(name.c_str(), std::vector{attr.ints().begin(), attr.ints().end()}); + } + } else if (name == "pool_strings" && attr.type() == ge::onnx::AttributeProto::STRINGS) { + op_dest.SetAttr(name.c_str(), std::vector{attr.strings().begin(), attr.strings().end()}); + } + } + + return SUCCESS; +} + +// register op info to GE +REGISTER_CUSTOM_OP("TfIdfVectorizer") + .FrameworkType(ONNX) + .OriginOpType({ge::AscendString("ai.onnx::9::TfIdfVectorizer"), + ge::AscendString("ai.onnx::10::TfIdfVectorizer"), + ge::AscendString("ai.onnx::11::TfIdfVectorizer"), + ge::AscendString("ai.onnx::12::TfIdfVectorizer"), + ge::AscendString("ai.onnx::13::TfIdfVectorizer"), + ge::AscendString("ai.onnx::14::TfIdfVectorizer"), + ge::AscendString("ai.onnx::15::TfIdfVectorizer"), + ge::AscendString("ai.onnx::16::TfIdfVectorizer"), + ge::AscendString("ai.onnx::17::TfIdfVectorizer"), + ge::AscendString("ai.onnx::18::TfIdfVectorizer")}) + .ParseParamsFn(ParseParamsTfIdfVectorizer) + .ImplyType(ImplyType::TVM); +} // namespace domi \ No newline at end of file diff --git a/csrc/ascend/common/src/tiling_base/tiling_util.cpp b/csrc/ascend/common/src/tiling_base/tiling_util.cpp new file mode 100644 index 000000000..c85de69f0 --- /dev/null +++ b/csrc/ascend/common/src/tiling_base/tiling_util.cpp @@ -0,0 +1,54 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_util.cpp + * \brief + */ + +#include "tiling_base/tiling_util.h" +#include "platform/platform_ascendc.h" + +namespace Ops { +namespace Transformer { +namespace OpTiling { +static const gert::Shape g_vec_1_shape = {1}; + +static bool IsRegbaseSocVersion(platform_ascendc::SocVersion version) +{ + const static std::set regbaseSocVersions = { + platform_ascendc::SocVersion::ASCEND950}; + + return regbaseSocVersions.find(version) != regbaseSocVersions.end(); +} + +bool IsRegbaseSocVersion(const gert::TilingParseContext* context) +{ + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + auto socVersion = ascendcPlatform.GetSocVersion(); + return IsRegbaseSocVersion(socVersion); +} + +bool IsRegbaseSocVersion(const gert::TilingContext* context) +{ + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + auto socVersion = ascendcPlatform.GetSocVersion(); + return IsRegbaseSocVersion(socVersion); +} + +const gert::Shape &EnsureNotScalar(const gert::Shape &inShape) { + if (inShape.IsScalar()) { + return g_vec_1_shape; + } + return inShape; +} +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops diff --git a/csrc/ascend/common/src/tiling_sink/CMakeLists.txt b/csrc/ascend/common/src/tiling_sink/CMakeLists.txt new file mode 100644 index 000000000..cefe9c45c --- /dev/null +++ b/csrc/ascend/common/src/tiling_sink/CMakeLists.txt @@ -0,0 +1,197 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +cmake_minimum_required(VERSION 3.16) + +project(tiling_sink_project) + +set(CMAKE_SKIP_RPATH TRUE) + +include(${CMAKE_CURRENT_SOURCE_DIR}/../../../cmake/intf.cmake) + +set(ATTENTION_SRC ${CMAKE_CURRENT_SOURCE_DIR}/../../../attention) +set(src_files + ${ATTENTION_SRC}/common/op_host/fia_tiling_info.cpp + ${ATTENTION_SRC}/common/op_host/arch32/fia_tiling_nonquant_mla.cpp + ${ATTENTION_SRC}/common/op_host/arch32/fia_tiling_nonquant.cpp + ${ATTENTION_SRC}/common/op_host/arch32/fia_tiling_empty_tensor.cpp + ${ATTENTION_SRC}/common/op_host/fia_tiling_shape.cpp + ${ATTENTION_SRC}/common/op_host/split_core.cpp + ${ATTENTION_SRC}/fused_infer_attention_score/op_host/fused_infer_attention_score_tiling.cpp + ${ATTENTION_SRC}/fused_infer_attention_score/op_host/arch32/fused_infer_attention_score_tiling_check.cpp + ${ATTENTION_SRC}/fused_infer_attention_score/op_host/arch32/fused_infer_attention_score_tiling_check_single_para.cpp + ${ATTENTION_SRC}/fused_infer_attention_score/op_host/arch32/fused_infer_attention_score_tiling_check_existence.cpp + ${ATTENTION_SRC}/fused_infer_attention_score/op_host/arch32/fused_infer_attention_score_tiling_check_feature.cpp + ${ATTENTION_SRC}/fused_infer_attention_score/op_host/arch32/fused_infer_attention_score_tiling_check_consistency.cpp + ${ATTENTION_SRC}/fused_infer_attention_score/op_host/arch32/fused_infer_attention_score_tiling_info_parser.cpp + ${ATTENTION_SRC}/fused_infer_attention_score/op_host/arch32/fused_infer_attention_score_tiling_v3.cpp + ${ATTENTION_SRC}/fused_infer_attention_score/op_host/arch35/fused_infer_attention_score_tiling_v2.cpp + ${ATTENTION_SRC}/prompt_flash_attention/op_host/prompt_flash_attention_tiling.cpp + ${ATTENTION_SRC}/prompt_flash_attention/op_host/prompt_flash_attention_tiling_v2.cpp + ${ATTENTION_SRC}/incre_flash_attention/op_host/incre_flash_attention_tiling.cpp + ${ATTENTION_SRC}/incre_flash_attention/op_host/incre_flash_attention_tiling_v2.cpp + ${ATTENTION_SRC}/incre_flash_attention/op_host/incre_flash_attention_tiling_check.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/tiling_sink_registry.cpp +) + +foreach(f ${src_files}) + if(NOT EXISTS ${f}) + message(WARNING "File not found: ${f}") + return() + endif() +endforeach() + +add_library(opmaster SHARED + ${src_files} +) + +if (ENABLE_BUILT_IN) + set(OPMASTER_NAME "tiling_device_transformer") +else() + set(OPMASTER_NAME "cust_opmaster") +endif() + +set_target_properties(opmaster PROPERTIES OUTPUT_NAME ${OPMASTER_NAME}) + +if (BUILD_OPEN_PROJECT) + add_custom_command( + TARGET opmaster + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy + $ + ${CMAKE_BINARY_DIR}/lib${OPMASTER_NAME}.so + ) +endif () + +target_include_directories(opmaster PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ${ATTENTION_SRC}/common/op_host + ${ATTENTION_SRC}/common/op_kernel + ${ATTENTION_SRC}/fused_infer_attention_score/op_host + ${ATTENTION_SRC}/prompt_flash_attention/op_host + ${ATTENTION_SRC}/incre_flash_attention/op_host + ${CMAKE_CURRENT_SOURCE_DIR}/../../include + ${CMAKE_CURRENT_SOURCE_DIR}/../../include/tiling_sink +) + +if (BUILD_OPEN_PROJECT) + target_include_directories(opmaster PRIVATE + ${ASCEND_CANN_PACKAGE_PATH}/include + ${ASCEND_CANN_PACKAGE_PATH}/include/base + ${ASCEND_CANN_PACKAGE_PATH}/pkg_inc + ${ASCEND_CANN_PACKAGE_PATH}/pkg_inc/op_common + ${ASCEND_CANN_PACKAGE_PATH}/include/aicpu + ${ASCEND_CANN_PACKAGE_PATH}/pkg_inc/aicpu_common/context/common + ${ASCEND_CANN_PACKAGE_PATH}/pkg_inc/aicpu_common/context/utils + ${ASCEND_CANN_PACKAGE_PATH}/pkg_inc/aicpu_common/context/cpu_proto + ${ASCEND_CANN_PACKAGE_PATH}/include/toolchain + ${ASCEND_CANN_PACKAGE_PATH}/pkg_inc/base + ) +else () + target_include_directories(opmaster PRIVATE + ${TOP_DIR}/ops-base/include + ${TOP_DIR}/ops-base/pkg_inc + ${TOP_DIR}/ops-base/pkg_inc/op_common + ${TOP_DIR}/ops-base/include/base + ${TOP_DIR}/ops-base/pkg_inc/aicpu_common/context/common + ${TOP_DIR}/ops-base/pkg_inc/aicpu_common/context/utils + ${TOP_DIR}/ops-base/pkg_inc/aicpu_common/context/cpu_proto + ${TOP_DIR}/metadef/inc/external + ${TOP_DIR}/metadef/inc + ${TOP_DIR}/abl/slog/inc/toolchain # codespell:ignore abl + ${TOP_DIR}/metadef/inc/external/base + ${TOP_DIR}/abl/slog/inc # codespell:ignore abl + ) +endif () + +if (BUILD_OPEN_PROJECT) + target_link_directories(opmaster PRIVATE + ${ASCEND_CANN_PACKAGE_PATH}/lib64/device/lib64 + ${ASCEND_CANN_PACKAGE_PATH}/devlib/device + ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/lib64 + ) +endif() + +target_link_libraries(opmaster PRIVATE + $,intf_pub_aicpu,intf_pub_cxx17>> + -Wl,--whole-archive + -Wl,--no-whole-archive + mmpa + tiling_api + exe_meta_device + platform_static + ascend_protobuf + -Wl,--no-as-needed + c_sec + -ldl + -Wl,--as-needed +) + +find_library(DEVICE_REGISTER_EXIST device_register + ${ASCEND_CANN_PACKAGE_PATH}/lib64/device/lib64 + ${ASCEND_CANN_PACKAGE_PATH}/devlib/device + ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/lib64) +if (DEVICE_REGISTER_EXIST STREQUAL "DEVICE_REGISTER_EXIST-NOTFOUND") + message(WARNING "device_register library not found!") +else() + target_link_libraries(opmaster PRIVATE + -Wl,--whole-archive + device_register + -Wl,--no-whole-archive + ) +endif() + +if (NOT BUILD_OPEN_PROJECT) + target_link_libraries(opmaster PRIVATE + $ + -Wl,--no-as-needed + slog + -Wl,--as-needed + ) +endif () + +target_compile_definitions(opmaster PRIVATE + _FORTIFY_SOURCE=2 + google=ascend_private + DEVICE_OP_TILING_LIB + BUILT_IN_TILING_SINK +) + +target_compile_definitions(opmaster PRIVATE + OPS_UTILS_LOG_SUB_MOD_NAME="OP_TILING" + OPS_UTILS_LOG_PACKAGE_TYPE=$,"[Custom]",""> +) + +target_compile_options(opmaster PRIVATE + -O2 + -fstack-protector-all + -fvisibility-inlines-hidden + -fvisibility=hidden + -frename-registers + -fpeel-loops + -DEIGEN_NO_DEBUG + -DEIGEN_MPL2_ONLY + -DNDEBUG + -DEIGEN_HAS_CXX11_MATH + -DEIGEN_OS_GNULINUX + -DEigen=ascend_Eigen +) + +target_link_options(opmaster PRIVATE + -Wl,-z,relro,-z,now + $<$>:-s> + -Wl,-Bsymbolic +) + +if (NOT BUILD_OPEN_PROJECT) + install ( + TARGETS opmaster OPTIONAL + LIBRARY DESTINATION lib + ) +endif () diff --git a/csrc/ascend/common/src/tiling_sink/tiling_sink_registry.cpp b/csrc/ascend/common/src/tiling_sink/tiling_sink_registry.cpp new file mode 100644 index 000000000..d7a77f54a --- /dev/null +++ b/csrc/ascend/common/src/tiling_sink/tiling_sink_registry.cpp @@ -0,0 +1,22 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_sink_registry.cpp + * \brief + */ + + #include "register/device_op_impl_registry.h" + #include "fused_infer_attention_score_tiling.h" + + #ifndef ASCEND_OPTILING_UT + DEVICE_IMPL_OP_OPTILING(FusedInferAttentionScore).Tiling(optiling::DeviceDoOpTilingFusedInferAttentionScore); + DEVICE_IMPL_OP_OPTILING(IncreFlashAttention).Tiling(optiling::DeviceDoOpTilingIncreFlashAttention); + #endif \ No newline at end of file diff --git a/csrc/ascend/common/stub/CMakeLists.txt b/csrc/ascend/common/stub/CMakeLists.txt new file mode 100644 index 000000000..c882e1975 --- /dev/null +++ b/csrc/ascend/common/stub/CMakeLists.txt @@ -0,0 +1,18 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG) + file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) + foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() + endforeach() +endif() \ No newline at end of file diff --git a/csrc/ascend/common/stub/inc/framework/stub_ops.h b/csrc/ascend/common/stub/inc/framework/stub_ops.h new file mode 100644 index 000000000..67f6e19a1 --- /dev/null +++ b/csrc/ascend/common/stub/inc/framework/stub_ops.h @@ -0,0 +1,140 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file stub_ops.h + * \brief + */ + +#ifndef MATH_COMMON_STUB_OPS_H +#define MATH_COMMON_STUB_OPS_H + +#include "graph/operator_reg.h" +#include "graph/operator.h" + +namespace ge { +/** +*@brief Input data for other operators. \n + +*@par Inputs: +*x: A tensor. \n + +*@par Attributes: +*index: Index of the input tensor.The data type must be int32 or int64. +Assume that net has three data nodes, one should be set 0, another should +be set 1, and the left should be set 2. \n + +*@par Outputs: +*y: A tensor. \n + +*@par Third-party framework compatibility +*Compatible with the Caffe operator Data. +*/ +REG_OP(Data) + .INPUT(x, TensorType::ALL()) + .OUTPUT(y, TensorType::ALL()) + .ATTR(index, Int, 0) + .OP_END_FACTORY_REG(Data) + +/** +*@brief Creates a constant tensor from a tensor-like object. This operator is used for inference. +Operator Const has the same definition as operator Constant. \n + +*@par Attributes: +*value: Required. The value and type of the resulting tensor, and no restrictions on type. \n + +*@par Outputs: +*y: A constant tensor. \n + +*@par Third-party framework compatibility +*Compatible with the TensorFlow operator Const. +*/ +REG_OP(Const) + .OUTPUT(y, TensorType::ALL()) + .ATTR(value, Tensor, Tensor()) + .OP_END_FACTORY_REG(Const) + +/** +*@brief Cast a tensor from src data type to dst data type. + +*@par Inputs: +*One input: +* x:An ND or 5HD tensor. Support 1D~8D. Must be one of the following types: bool, float16, float, int8, int32, uint32, uint8, bfloat16, uint1, + int64, uint64, int16, uint16, double, complex32, complex64, complex128, qint8, quint8, qint16, quint16, qint32, + hifloat8, float8_e5m2, float8_e4m3fn, float4_e1m2, float4_e2m1. + +*@par Attributes: +*dst_type: A required attribute of type int32, specifying the dst data type. + +*@par Outputs: +*y:An ND Tensor with same shape as x, and data type is specified by dst_type. + +*@attention Constraints: +* @li In the scenario where the data type is converted from float16 to int16: \n +* If the input data contains inf, inf is converted into the maximum value of int16. \n +* If the input data contains -inf, -inf is converted into the minimum value of int16. \n +* @li In the scenarios where the data type is converted from INT32 to INT8: \n +* It can only guarantee that the input data has no precision errors within the range of (-2048, 1920). +* @li Atlas Inference Series Product in the scenarios where the data type is converted from FLOAT32 to INT8: \n +* It can only guarantee that the input data has no precision errors within the range of (-2048, 1920). +* @li Atlas Inference Series Product in the scenarios where the data type is converted from FLOAT32 to INT64 and from FLOAT32 to UINT8: \n +* It can only guarantee that the input data has no precision errors within the range of (-2147483648, 2147483583). +* @li Atlas Inference Series Product in the scenarios where the data type is converted from INT64 to FLOAT32: \n +* It can only guarantee that the input data has no precision errors within the range of (-2147483648, 2147483647). +*/ +REG_OP(Cast) + .INPUT(x, TensorType({DT_BOOL, DT_FLOAT16, DT_FLOAT, DT_INT8, DT_INT32, DT_UINT32, DT_UINT8, + DT_INT64, DT_UINT64, DT_INT16, DT_UINT16, DT_DOUBLE, DT_COMPLEX64, + DT_COMPLEX128, DT_QINT8, DT_QUINT8, DT_QINT16, DT_QUINT16, DT_QINT32, DT_BF16, DT_UINT1, + DT_COMPLEX32, DT_HIFLOAT8, DT_FLOAT8_E5M2, DT_FLOAT8_E4M3FN, + DT_FLOAT4_E1M2, DT_FLOAT4_E2M1})) + .OUTPUT(y, TensorType({DT_BOOL, DT_FLOAT16, DT_FLOAT, DT_INT8, DT_INT32, DT_UINT32, DT_UINT8, + DT_INT64, DT_UINT64, DT_INT16, DT_UINT16, DT_DOUBLE, DT_COMPLEX64, + DT_COMPLEX128, DT_QINT8, DT_QUINT8, DT_QINT16, DT_QUINT16, DT_QINT32, + DT_BF16, DT_COMPLEX32, DT_HIFLOAT8, DT_FLOAT8_E5M2, DT_FLOAT8_E4M3FN, + DT_FLOAT4_E1M2, DT_FLOAT4_E2M1})) + .REQUIRED_ATTR(dst_type, Int) + .OP_END_FACTORY_REG(Cast) + +/** +* @brief Creates a tensor filled with a scalar value. +* This operation creates a tensor of shape "dims" and fills it with "value". +* +* @par Inputs: +* @li dims: A 1D tensor of types int32 or int64. Represents the shape of the output tensor . + The size of each dimension must be less than or equal to 8. \n + +* @li value: A 0D scalar. Specifies the value to fill the returned tensor. +* Must be one of the following types: +* bfloat16, float16, float32, double, int32, uint8, int16, int8, complex64, int64, bool, +* qint8, quint8, qint32, qint16, quint16, uint16, complex128, uint32, uint64, string. +* +* @par Outputs: +* y: A tensor. Has the same type as "value". +* +* @par Third-party framework compatibility +* @li Compatible with the TensorFlow operator Fill. +* @li Compatible with the Caffe operator Filler. +* +*/ +REG_OP(Fill) + .INPUT(dims, TensorType::IndexNumberType()) + .INPUT(value, "T") + .OUTPUT(y, "T") + .DATATYPE(T, TensorType({DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT16, + DT_INT8, DT_COMPLEX64, DT_INT64, DT_BOOL, DT_QINT8, + DT_QUINT8, DT_QINT32, DT_QINT16, DT_QUINT16, DT_UINT16, + DT_COMPLEX128, DT_FLOAT16, DT_BF16, DT_UINT32, DT_UINT64, DT_STRING})) + .OP_END_FACTORY_REG(Fill) + + +} // namespace ge + +#endif // MATH_COMMON_STUB_OPS_H diff --git a/csrc/ascend/common/stub/op_api/CMakeLists.txt b/csrc/ascend/common/stub/op_api/CMakeLists.txt new file mode 100644 index 000000000..d8518de5e --- /dev/null +++ b/csrc/ascend/common/stub/op_api/CMakeLists.txt @@ -0,0 +1,40 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}/build) + +file(GLOB_RECURSE OP_API_FILES "*.cpp") + +add_library(opapi_math SHARED ${OP_API_FILES}) + +target_compile_definitions(opapi_math PRIVATE + _GLIBCXX_USE_CXX11_ABI=0 + LOG_CPP +) + +target_include_directories(opapi_math PRIVATE + ${OPAPI_INCLUDE} + ${ASCEND_CANN_PACKAGE_PATH}/pkg_inc +) + +if(ENABLE_TEST) + add_library(opapi_stub SHARED ${OP_API_FILES}) + + target_compile_definitions(opapi_stub PRIVATE + _GLIBCXX_USE_CXX11_ABI=0 + LOG_CPP + ) + + target_include_directories(opapi_stub PRIVATE + ${OPAPI_INCLUDE} + ${ASCEND_CANN_PACKAGE_PATH}/pkg_inc + ${ASCEND_CANN_PACKAGE_PATH}/include/ascendc/basic_api + ) +endif() \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_api/aclnn_kernels/cast.h b/csrc/ascend/common/stub/op_api/aclnn_kernels/cast.h new file mode 100644 index 000000000..25f5c91f7 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/aclnn_kernels/cast.h @@ -0,0 +1,24 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_CAST_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_CAST_H + +#include "opdev/op_executor.h" +#include "opdev/make_op_executor.h" + +namespace l0op { +const aclTensor* Cast(const aclTensor* self, op::DataType dstDtype, aclOpExecutor* executor); + +// 专攻卷积反向定制 +const aclTensor* CastOnlyForConvBackward(const aclTensor* self, op::DataType dstDtype, aclOpExecutor* executor); +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_CAST_H diff --git a/csrc/ascend/common/stub/op_api/aclnn_kernels/common/op_error_check.h b/csrc/ascend/common/stub/op_api/aclnn_kernels/common/op_error_check.h new file mode 100644 index 000000000..48f41354a --- /dev/null +++ b/csrc/ascend/common/stub/op_api/aclnn_kernels/common/op_error_check.h @@ -0,0 +1,251 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_ERROR_CHECK_H__ +#define OP_ERROR_CHECK_H__ + +#include "opdev/op_log.h" +#include "opdev/common_types.h" +#include "opdev/data_type_utils.h" +#include "opdev/shape_utils.h" + +const int32_t NCHW_N_DIM = 0; +const int32_t NCHW_C_DIM = 1; +const int32_t NHWC_N_DIM = 0; +const int32_t NHWC_C_DIM = 3; + +static inline bool IsNullptr(const aclTensor *tensor, const char *name) { + if (tensor == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a proper Tensor but got null for argument %s.", name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclTensorList *tensorList, const char *name) { + if (tensorList == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a proper TensorList but got null for argument %s.", name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclScalar *scalar, const char *name) { + if (scalar == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a value of type number for argument %s but instead found type null.", + name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclIntArray *intArr, const char *name) { + if (intArr == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a value of type List[int] for argument %s but instead found type null.", + name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclBoolArray *boolArr, const char *name) { + if (boolArr == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a value of type List[bool] for argument %s but instead found type null.", + name); + return true; + } + return false; +} + +static inline bool IsNullptr(const aclFloatArray *floatArr, const char *name) { + if (floatArr == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a value of type List[float] for argument %s but instead found type \ + null.", name); + return true; + } + return false; +} + +static inline bool CheckDims(const aclTensor *tensor) { + const auto& xShape = tensor->GetViewShape(); + for(size_t i = 0; i < xShape.GetDimNum(); i++) { + if (xShape.GetDim(i) > INT32_MAX) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The tensor's shape cannot be larger than %d.", INT32_MAX); + return false; + } + } + return true; +} + +static inline bool CheckReduceOutShape(const aclTensor *inferOut, const aclTensor *out) +{ + auto const &xShape = inferOut->GetViewShape(); + auto const &yShape = out->GetViewShape(); + if (xShape != yShape) { + if (!(xShape.GetShapeSize() == 1 && yShape.GetShapeSize() == 1)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The out tensor's shape[%s] is not equal with inferOut shape[%s].", + op::ToString(out->GetViewShape()).GetString(), op::ToString(inferOut->GetViewShape()).GetString()); + return false; + } + } + return true; +} + +static inline bool CheckNCDimValid(const aclTensor *self, const aclTensor *out) { + auto format = self->GetStorageFormat(); + int64_t selfDimN = 0; + int64_t selfDimC = 0; + int64_t outDimN = 0; + int64_t outDimC = 0; + if (format == op::Format::FORMAT_NCHW) { + selfDimN = self->GetViewShape().GetDim(NCHW_N_DIM); + selfDimC = self->GetViewShape().GetDim(NCHW_C_DIM); + outDimN = out->GetViewShape().GetDim(NCHW_N_DIM); + outDimC = out->GetViewShape().GetDim(NCHW_C_DIM); + } else if (format == op::Format::FORMAT_NHWC) { + selfDimN = self->GetViewShape().GetDim(NHWC_N_DIM); + selfDimC = self->GetViewShape().GetDim(NHWC_C_DIM); + outDimN = out->GetViewShape().GetDim(NHWC_N_DIM); + outDimC = out->GetViewShape().GetDim(NHWC_C_DIM); + } else { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "Input and output format only support [NCHW, NHWC] format ."); + return false; + } + if ((selfDimN != outDimN) || (selfDimC != outDimC)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "The selfDimN[%ld]/outDimN[%ld] or selfDimC[%ld]/outDimC[%ld] not equal .", + selfDimN, outDimN, selfDimC, outDimC); + return false; + } + return true; +} + + +#define OP_CHECK_NULL(param, retExpr) \ + if (IsNullptr(param, #param)) { \ + retExpr; \ + } + +#define OP_CHECK_DTYPE_NOT_SUPPORT(tensor, supportList, retExpr) \ + if (!CheckType(tensor->GetDataType(), supportList)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Tensor %s not implemented for %s, should be in dtype support list %s.", \ + #tensor, op::ToString(tensor->GetDataType()).GetString(), op::ToString(supportList).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_DTYPE_NOT_MATCH(tensor, expectedDtype, retExpr) \ + if (tensor->GetDataType() != expectedDtype) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Tensor %s expected dtype is %s but found %s.", \ + #tensor, op::ToString(expectedDtype).GetString(), op::ToString(tensor->GetDataType()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_DTYPE_NOT_SAME(tensor1, tensor2, retExpr) \ + if (tensor1->GetDataType() != tensor2->GetDataType()) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected both tensors to have same dtype, but found %s %s and %s %s.", \ + #tensor1, op::ToString(tensor1->GetDataType()).GetString(), \ + #tensor2, op::ToString(tensor2->GetDataType()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_RESULT_DTYPE_CAST_FAILED(dtype, desiredDtype, retExpr); \ + if (!CanCast(dtype, desiredDtype)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Result type %s can't be cast to the desired output type %s.", \ + op::ToString(dtype).GetString(), op::ToString(desiredDtype).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_BROADCAST(tensor1, tensor2, retExpr) \ + if (!CheckBroadcastShape(tensor1->GetViewShape(), tensor2->GetViewShape())) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The size of tensor %s %s must match the size of tensor %s %s.", \ + #tensor1, op::ToString(tensor1->GetViewShape()).GetString(), \ + #tensor2, op::ToString(tensor2->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_BROADCAST_WITH_SHAPE(tensor, shape, retExpr) \ + if (!CheckBroadcastShape(tensor->GetViewShape(), shape)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The size of tensor %s %s must match the size %s.", \ + #tensor, op::ToString(tensor->GetViewShape()).GetString(), op::ToString(shape).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_BROADCAST_AND_INFER_SHAPE(tensor1, tensor2, retShape, retExpr) \ + if (!BroadcastInferShape(tensor1->GetViewShape(), tensor2->GetViewShape(), retShape)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The size of tensor %s %s must match the size of tensor %s %s.", \ + #tensor1, op::ToString(tensor1->GetViewShape()).GetString(), \ + #tensor2, op::ToString(tensor2->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_SHAPE_NOT_EQUAL(tensor1, tensor2, retExpr) \ + if (tensor1->GetViewShape() != tensor2->GetViewShape()) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected tensor for %s to have same size as tensor for %s, but %s does not " \ + "equal %s.", #tensor1, #tensor2, op::ToString(tensor1->GetViewShape()).GetString(), \ + op::ToString(tensor2->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(tensor, shape, retExpr) \ + if (tensor->GetViewShape() != shape) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected tensor for %s to have same size as %s, but got %s.", \ + #tensor, op::ToString(shape).GetString(), op::ToString(tensor->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_WRONG_DIMENSION(tensor, expectedDimNum, retExpr) \ + if (tensor->GetViewShape().GetDimNum() != expectedDimNum) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected %zu dimension input, but got %s with sizes %s.", \ + static_cast(expectedDimNum), #tensor, op::ToString(tensor->GetViewShape()).GetString()); \ + retExpr; \ + } + +#define OP_CHECK_MAX_DIM(tensor, maxDim, retExpr) \ + if (tensor->GetViewShape().GetDimNum() > static_cast(maxDim)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The %s tensor cannot be larger than %zu dimensions.", \ + #tensor, static_cast(maxDim)); \ + retExpr; \ + } + +#define OP_CHECK_MIN_DIM(tensor, minDim, retExpr) \ + if (tensor->GetViewShape().GetDimNum() < static_cast(minDim)) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The %s tensor must have at least %zu dimensions.", \ + #tensor, static_cast(minDim)); \ + retExpr; \ + } + +#define OP_CHECK_COMM_INPUT(workspaceSize, executor) \ + if (workspaceSize == nullptr || executor == nullptr) { \ + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "The workspaceSize or executor is nullptr."); \ + return ACLNN_ERR_PARAM_NULLPTR; \ + } + +#define OP_CHECK_ADD_TO_LAUNCHER_LIST_AICORE(cond, retExpr, errMsg, ...) \ + if (cond) { \ + OP_LOGE(ACLNN_ERR_INNER_STATIC_WORKSPACE_INVALID, errMsg, ##__VA_ARGS__); \ + retExpr; \ + } + +#define OP_CHECK_INFERSHAPE(cond, retExpr, errMsg, ...) \ + if (cond) { \ + OP_LOGE(ACLNN_ERR_INNER_INFERSHAPE_ERROR, errMsg, ##__VA_ARGS__); \ + retExpr; \ + } + +#define OP_CHECK_TENSORLIST_SIZE_EQUAL(tensorlist1, tensorlist2, retExpr) \ + if ((tensorlist1)->Size() != (tensorlist2)->Size()) { \ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, \ + "The %s tensorlist and %s tensorlist must have the same number of tensors, but got %ld and %ld.", \ + #tensorlist1, #tensorlist2, (tensorlist1)->Size(), (tensorlist2)->Size()); \ + retExpr; \ + } + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_api/aclnn_kernels/contiguous.h b/csrc/ascend/common/stub/op_api/aclnn_kernels/contiguous.h new file mode 100644 index 000000000..0c49bfdce --- /dev/null +++ b/csrc/ascend/common/stub/op_api/aclnn_kernels/contiguous.h @@ -0,0 +1,89 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_CONTIGUOUS_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_CONTIGUOUS_H + +#include "opdev/op_def.h" +#include "opdev/common_types.h" + +namespace l0op { + +typedef struct { + // 每个op::Shape 18ns + int64_t viewOffset; + + // Transpose + op::Shape transposeSrcShape; + op::Shape transposeDstShape; + op::FVector perm; + + // broadcast to + op::Shape broadcastSrcShape; + op::Shape broadcastDstShape; + op::FVector shape; + + // slice + op::Shape sliceSrcShape; + op::Shape sliceDstShape; + op::FVector offset; + op::FVector size; + + // strided slice + op::Shape stridedsliceSrcShape; + op::Shape stridedsliceDstShape; + op::FVector begin; + op::FVector end; + op::FVector strides; + + // optimizer + bool mayBroadcast; + bool mayTranspose; + bool maySlice; + bool mayStridedslice; +} ContiguousParam; + +/** + * @brief 将非连续Tensor转换为连续Tensor + * @param x + * @param executor + * @return aclTensor 转换后的tensor + */ +const aclTensor* Contiguous(const aclTensor* x, aclOpExecutor* executor); + +/** + * @brief 将连续tensor拷贝到非连续的tensor上 + * @param x + * @param y + * @param executor + * @return aclTensor 转换后的tensor + */ +const aclTensor* ViewCopy(const aclTensor* x, const aclTensor* y, aclOpExecutor* executor); + +/** + * @brief 对Tensor创建一个View，要求Tensor满足PickView的条件 + * @param x 输入Tensor，可以是一整块的非连续Tensor + * @param executor + * @return 输出Shape是一个连续Tensor + */ +const aclTensor* PickViewAsContiguous(const aclTensor* x, aclOpExecutor* executor); + +const aclTensor* ReViewToOut(const aclTensor* x, const aclTensor* y, aclOpExecutor* executor); + +// ============内部接口============= +bool CanOptimizeContiguous( + const op::Shape& viewShape, const op::Strides& strides, int64_t offset, int64_t storageSize, + ContiguousParam& param); + +bool CanOptimizeView(const op::Shape& viewShape, const op::Strides& strides, int64_t offset, ContiguousParam& param); +// ============内部接口============= +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_CONTIGUOUS_H diff --git a/csrc/ascend/common/stub/op_api/aclnn_kernels/pad.h b/csrc/ascend/common/stub/op_api/aclnn_kernels/pad.h new file mode 100644 index 000000000..ab740a435 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/aclnn_kernels/pad.h @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_PAD_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_PAD_H + +#include "opdev/op_executor.h" +#include "opdev/make_op_executor.h" + +namespace l0op { +const aclTensor* Pad(const aclTensor* self, const aclTensor* paddings, aclOpExecutor* executor); +} +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_PAD_H diff --git a/csrc/ascend/common/stub/op_api/aclnn_kernels/reshape.h b/csrc/ascend/common/stub/op_api/aclnn_kernels/reshape.h new file mode 100644 index 000000000..6a9a62d2c --- /dev/null +++ b/csrc/ascend/common/stub/op_api/aclnn_kernels/reshape.h @@ -0,0 +1,37 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_RESHAPE_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_RESHAPE_H + +#include "opdev/shape_utils.h" +#include "opdev/op_def.h" + +namespace l0op { +/** + * @brief Modify input tensor's shape. + * @param x Input Tensor. Should be contiguous. + * @param shape Target Shape. Only one dimension can be -1. + * @param executor aclOpExecutor.ldd + * @return *aclTensor Output tensor. + */ +const aclTensor* Reshape(const aclTensor* x, const op::Shape& shape, aclOpExecutor* executor); + +/** + * @brief Modify input tensor's shape. + * @param x Input Tensor. Should be contiguous. + * @param shape Target Shape. Only one dimension can be -1. + * @param executor aclOpExecutor. + * @return *aclTensor Output tensor. + */ +const aclTensor* Reshape(const aclTensor* x, const aclIntArray* shape, aclOpExecutor* executor); +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_RESHAPE_H diff --git a/csrc/ascend/common/stub/op_api/aclnn_kernels/slice.h b/csrc/ascend/common/stub/op_api/aclnn_kernels/slice.h new file mode 100644 index 000000000..3594f478c --- /dev/null +++ b/csrc/ascend/common/stub/op_api/aclnn_kernels/slice.h @@ -0,0 +1,25 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_SLICE_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_SLICE_H + +#include "opdev/op_def.h" + +namespace l0op { + +const aclTensor* Slice( + const aclTensor* x, const aclTensor* y, const aclTensor* offset, const aclTensor* size, aclOpExecutor* executor); + +const aclTensor* Slice( + const aclTensor* x, const aclIntArray* offsets, const aclIntArray* size, aclOpExecutor* executor); +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_SLICE_H diff --git a/csrc/ascend/common/stub/op_api/aclnn_kernels/transdata.h b/csrc/ascend/common/stub/op_api/aclnn_kernels/transdata.h new file mode 100644 index 000000000..33e7bd7df --- /dev/null +++ b/csrc/ascend/common/stub/op_api/aclnn_kernels/transdata.h @@ -0,0 +1,56 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSDATA_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSDATA_H + +#include "opdev/op_executor.h" + +namespace l0op { + +const aclTensor* ReFormat(const aclTensor* x, const op::Format& format, aclOpExecutor* executor = nullptr); + +/** + * TransData + * Formal Transdata. Set the c0 size strictly based on the data type and chip block size. + * support data type as follows: fp16,fp32,int32,uint32,int8,uint8 + * fp16: block_size/2 + * fp32/int32/uint32: block_size/4 (this is different from `TransDataSpecial`) + * int8/uint8: block_size/1 + * + * @param x : aclTensor need to transpose + * @param dstPrimaryFormat: dstPrimaryFormat like NC1HWC0 + * @param groups: groups + * @param executor: executor should not be null + * @return trans format tensor + */ +const aclTensor* TransData(const aclTensor* x, op::Format dstPrimaryFormat, int64_t groups, aclOpExecutor* executor); +/** + * Special Transdata. Set the c0 size strictly based on the data type and chip block size. + * this transdata c0 size rule: + * fp16: block_size/2 + * fp32/int32/uint32: block_size/2 + * int8/uint8: block_size/1 + * bool not supported, should do: + * (NCHW, bool)-> cast -> (NCHW, fp16) -> TransDataSpecial -> (5HD, fp16) -> cast -> (5HD, bool) + * (5HD, bool)-> cast -> (5HD, fp16) -> TransDataSpecial -> (NCHW, fp16) -> cast -> (NCHW, bool) + * + * @param x : aclTensor need to transpose + * @param dstPrimaryFormat: dstPrimaryFormat like NC1HWC0 + * @param groups: groups + * @param executor: executor should not be null + * @return trans format tensor + */ +const aclTensor* TransDataSpecial( + const aclTensor* x, op::Format dstPrimaryFormat, int64_t groups, aclOpExecutor* executor); + +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSDATA_H diff --git a/csrc/ascend/common/stub/op_api/aclnn_kernels/transpose.h b/csrc/ascend/common/stub/op_api/aclnn_kernels/transpose.h new file mode 100644 index 000000000..b67fa20a9 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/aclnn_kernels/transpose.h @@ -0,0 +1,22 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSPOSE_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSPOSE_H + +#include "opdev/op_def.h" + +namespace l0op { + +const aclTensor* Transpose(const aclTensor* x, const aclTensor* y, const aclTensor* perm, aclOpExecutor* executor); +const aclTensor* Transpose(const aclTensor* x, const aclIntArray* perm, aclOpExecutor* executor); +} // namespace l0op + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_TRANSPOSE_H diff --git a/csrc/ascend/common/stub/op_api/level0/add.h b/csrc/ascend/common/stub/op_api/level0/add.h new file mode 100644 index 000000000..44cdff899 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/add.h @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add.h + * \brief + */ + +#ifndef OP_API_INC_LEVEL0_OP_ADD_OP_H_ +#define OP_API_INC_LEVEL0_OP_ADD_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor* Add(const aclTensor* self, const aclTensor* other, aclOpExecutor* executor); + +} + +#endif // OP_API_INC_LEVEL0_OP_ADD_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/arange.h b/csrc/ascend/common/stub/op_api/level0/arange.h new file mode 100644 index 000000000..30d8e7e8e --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/arange.h @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file arange.h + * \brief + */ + +#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_ARANGE_OP_H_ +#define PTA_NPU_OP_API_INC_LEVEL0_OP_ARANGE_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor* Arange(const aclScalar* start, const aclScalar* end, const aclScalar* step, const aclTensor* out, + const bool isClosed, aclOpExecutor* executor); +} + +#endif diff --git a/csrc/ascend/common/stub/op_api/level0/axpy.h b/csrc/ascend/common/stub/op_api/level0/axpy.h new file mode 100644 index 000000000..51530bf15 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/axpy.h @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file axpy.h +* \brief +*/ + +#ifndef OP_API_INC_LEVEL0_OP_AXPY_OP_H_ +#define OP_API_INC_LEVEL0_OP_AXPY_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor* Axpy(const aclTensor* self, const aclTensor* other, float alpha, aclOpExecutor* executor); + +} + +#endif // OP_API_INC_LEVEL0_OP_AXPY_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/batch_norm_backward.h b/csrc/ascend/common/stub/op_api/level0/batch_norm_backward.h new file mode 100644 index 000000000..fe57ed67e --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/batch_norm_backward.h @@ -0,0 +1,45 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_BATCH_NORM_GRAD_OP_H_ +#define PTA_NPU_OP_API_INC_LEVEL0_OP_BATCH_NORM_GRAD_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const std::array BNTrainingUpdateGrad(const aclTensor* gradOut, const aclTensor* x, + const aclTensor* saveMean, const aclTensor* saveInvstd, float eps, + aclOpExecutor* executor); +const std::array BN3DTrainingUpdateGrad(const aclTensor* gradOut, const aclTensor* x, + const aclTensor* saveMean, const aclTensor* saveInvstd, + float eps, aclOpExecutor* executor); + +const aclTensor* BNTrainingReduceGrad(const aclTensor* gradOut, const aclTensor* x, const aclTensor* gradWeight, + const aclTensor* gradBias, const aclTensor* weight, const aclTensor* saveMean, + const aclTensor* saveInvstd, float eps, aclOpExecutor* executor); +const aclTensor* BN3DTrainingReduceGrad(const aclTensor* gradOut, const aclTensor* x, const aclTensor* gradWeight, + const aclTensor* gradBias, const aclTensor* weight, const aclTensor* saveMean, + const aclTensor* saveInvstd, float eps, aclOpExecutor* executor); + +const aclTensor* BNInferGrad(const aclTensor* gradOut, const aclTensor* weight, const aclTensor* runningVar, float eps, + aclOpExecutor* executor); + +constexpr size_t BN_GRAD_V3_OUTPUT_NUM = 3; +const std::array BatchNormGradV3(const aclTensor* gradOut, + const aclTensor* input, + const aclTensor* weight, + const aclTensor* runningMean, + const aclTensor* runningVar, + const aclTensor* saveMean, + const aclTensor* saveInvstd, + bool training, float eps, + aclOpExecutor* executor); +} // namespace l0op + +#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_BATCH_NORM_GRAD_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/broadcast_to.h b/csrc/ascend/common/stub/op_api/level0/broadcast_to.h new file mode 100644 index 000000000..069a6916b --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/broadcast_to.h @@ -0,0 +1,28 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file broadcast_to.h +* \brief +*/ + +#ifndef OP_API_INC_LEVEL0_BROADCAST_TO_H +#define OP_API_INC_LEVEL0_BROADCAST_TO_H + +# include "opdev/op_def.h" + +namespace l0op { + +const aclTensor *BroadcastTo(const aclTensor *x, const aclTensor *y, const aclTensor *shape, aclOpExecutor *executor); +const aclTensor *BroadcastTo(const aclTensor *x, const aclIntArray *shape, aclOpExecutor *executor); + +} // l0op + +#endif // OP_API_INC_LEVEL0_BROADCAST_TO_H diff --git a/csrc/ascend/common/stub/op_api/level0/dilation.h b/csrc/ascend/common/stub/op_api/level0/dilation.h new file mode 100644 index 000000000..035a1dfec --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/dilation.h @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_API_OP_API_COMMON_INC_LEVEL0_OP_DILATION_OP_H_ +#define OP_API_OP_API_COMMON_INC_LEVEL0_OP_DILATION_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +// Dilation +const aclTensor *Dilation(const aclTensor *x, const aclIntArray *dilations, const aclIntArray *pads, + float paddingValue, aclOpExecutor *executor); +} // namespace l0op + +#endif // OP_API_OP_API_COMMON_INC_LEVEL0_OP_DILATION_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/div.h b/csrc/ascend/common/stub/op_api/level0/div.h new file mode 100644 index 000000000..6325d98a3 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/div.h @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_DIV_OP_H_ +#define PTA_NPU_OP_API_INC_LEVEL0_OP_DIV_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor *Div(const aclTensor *self, const aclTensor *other, aclOpExecutor *executor); + +} + +#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_ADD_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/dot.h b/csrc/ascend/common/stub/op_api/level0/dot.h new file mode 100644 index 000000000..a31098361 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/dot.h @@ -0,0 +1,27 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dot.h + * \brief + */ + +#ifndef OP_API_INC_LEVEL0_DOT_H_ +#define OP_API_INC_LEVEL0_DOT_H_ + +#include "opdev/op_executor.h" + +namespace l0op { + +const aclTensor *Dot(const aclTensor *self, const aclTensor *tensor, aclOpExecutor *executor); + +} + +#endif // OP_API_INC_LEVEL0_DOT_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/expand.h b/csrc/ascend/common/stub/op_api/level0/expand.h new file mode 100644 index 000000000..e9d11347f --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/expand.h @@ -0,0 +1,25 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file expand.h + * \brief + */ +#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_EXPAND_OP_H_ +#define PTA_NPU_OP_API_INC_LEVEL0_OP_EXPAND_OP_H_ + +#include "opdev/op_executor.h" +#include "opdev/make_op_executor.h" + +namespace l0op { +const aclTensor *Expand(const aclTensor *self, const aclIntArray *shape, aclOpExecutor *executor); +} + +#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_EXPAND_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/fault_injection.h b/csrc/ascend/common/stub/op_api/level0/fault_injection.h new file mode 100644 index 000000000..48c547b42 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/fault_injection.h @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_API_INC_LEVEL0_FAULT_INJECTION_H_ +#define OP_API_INC_LEVEL0_FAULT_INJECTION_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor* FaultInjection(const aclTensor *injectObj, const aclTensor *injectPara, aclTensor *out, aclOpExecutor *executor); +} // namespace l0op + +#endif // OP_API_INC_LEVEL0_FAULT_INJECTION_H_ \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_api/level0/fill.h b/csrc/ascend/common/stub/op_api/level0/fill.h new file mode 100644 index 000000000..2b72604e4 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/fill.h @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file fill.h + * \brief + */ + +#ifndef OP_API_INC_LEVEL0_OP_FILL_OP_H_ +#define OP_API_INC_LEVEL0_OP_FILL_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor *Fill(const aclTensor *dims, const aclTensor *value, const aclIntArray *outShape, + aclOpExecutor *executor); +} + +#endif diff --git a/csrc/ascend/common/stub/op_api/level0/gather_elements.h b/csrc/ascend/common/stub/op_api/level0/gather_elements.h new file mode 100644 index 000000000..0e725d5bf --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/gather_elements.h @@ -0,0 +1,27 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file gather_elements.h + * \brief + */ +#ifndef OP_API_INC_LEVEL0_OP_GATHER_ELEMENTS_H_ +#define OP_API_INC_LEVEL0_OP_GATHER_ELEMENTS_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor *GatherElements(const aclTensor *self, + const int64_t dim, + const aclTensor *index, + aclOpExecutor *executor); +} // l0op + +#endif // OP_API_INC_LEVEL0_OP_GATHER_ELEMENTS_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/gather_v2.h b/csrc/ascend/common/stub/op_api/level0/gather_v2.h new file mode 100644 index 000000000..5b02bb6d5 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/gather_v2.h @@ -0,0 +1,25 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_API_INC_LEVEL0_OP_GATHER_V2_H_ +#define OP_API_INC_LEVEL0_OP_GATHER_V2_H_ + +#include "opdev/op_executor.h" + +namespace l0op { + +const aclTensor *GatherV2(const aclTensor *self, int64_t axis, const aclTensor *indices, aclOpExecutor *executor, + int batchDims = 0, bool negativeIndexSupport = false); + +const aclTensor *GatherV2WithImplMode(const aclTensor *self, int64_t axis, const aclTensor *indices, int64_t implMode, + aclOpExecutor *executor, int batchDims = 0, bool negativeIndexSupport = false); + +} + +#endif // OP_API_INC_LEVEL0_OP_GATHER_V2_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/gather_v3.h b/csrc/ascend/common/stub/op_api/level0/gather_v3.h new file mode 100644 index 000000000..271831cdb --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/gather_v3.h @@ -0,0 +1,23 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_INC_EXTERNAL_ACLNN_KERNELS_GATHER_V3_H +#define COMMON_INC_EXTERNAL_ACLNN_KERNELS_GATHER_V3_H + +#include "opdev/op_executor.h" + +namespace l0op { + +const aclTensor *GatherV3(const aclTensor *self, int64_t axis, const aclTensor *indices, aclOpExecutor *executor, + int batchDims = 0, bool negativeIndexSupport = false); + +} + +#endif // COMMON_INC_EXTERNAL_ACLNN_KERNELS_GATHER_V3_H diff --git a/csrc/ascend/common/stub/op_api/level0/inplace_index_add.h b/csrc/ascend/common/stub/op_api/level0/inplace_index_add.h new file mode 100644 index 000000000..fc319a440 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/inplace_index_add.h @@ -0,0 +1,34 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file inplace_index_add.h + * \brief + */ +#ifndef OP_API_INC_LEVEL0_OP_INDEX_ADD_H_ +#define OP_API_INC_LEVEL0_OP_INDEX_ADD_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor *InplaceIndexAddAiCore(const aclTensor *self, const int64_t dim, const aclTensor *index, + const aclTensor *source, const aclTensor *alphaTensor, + aclOpExecutor *executor); + +const aclTensor *InplaceIndexAddAiCpu(const aclTensor *self, const int64_t dim, const aclTensor *index, + const aclTensor *source, const aclTensor *alphaTensor, + aclOpExecutor *executor); + +const aclTensor *InplaceIndexAddWithSorted(const aclTensor *self, const int64_t dim, const aclTensor *sortedIndices, + const aclTensor *pos, const aclTensor *value, const aclTensor *alphaTensor, + aclOpExecutor *executor); +} + +#endif diff --git a/csrc/ascend/common/stub/op_api/level0/masked_scatter.h b/csrc/ascend/common/stub/op_api/level0/masked_scatter.h new file mode 100644 index 000000000..dc1e5d5d1 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/masked_scatter.h @@ -0,0 +1,25 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file masked_scatter.h + * \brief + */ +#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_MASKED_SCATTER_OP_H_ +#define PTA_NPU_OP_API_INC_LEVEL0_OP_MASKED_SCATTER_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor* MaskedScatter(const aclTensor* self, const aclTensor* mask, const aclTensor* source, + aclOpExecutor* executor); +} + +#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_MASKED_SCATTER_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/matmul_v2tov3.h b/csrc/ascend/common/stub/op_api/level0/matmul_v2tov3.h new file mode 100644 index 000000000..0c776ff15 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/matmul_v2tov3.h @@ -0,0 +1,22 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_MATMUL_V2TOV3_H_ +#define PTA_NPU_OP_API_INC_LEVEL0_OP_MATMUL_V2TOV3_H_ + +#include "opdev/op_executor.h" + +namespace l0op { + +bool MmCheckHitV3Shape(const aclTensor* x1, const aclTensor* x2, const aclTensor* bias, + const bool transposeX1, const bool transposeX2, op::Format mat2_format, bool supportSplitK); +} + +#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_MATMUL_V2TOV3_H_ \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_api/level0/maximum.h b/csrc/ascend/common/stub/op_api/level0/maximum.h new file mode 100644 index 000000000..36b80fa2a --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/maximum.h @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file maximum.h +* \brief +*/ + +#ifndef OP_API_INC_LEVEL0_OP_MAXIMUM_OP_H_ +#define OP_API_INC_LEVEL0_OP_MAXIMUM_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor* Maximum(const aclTensor* self, const aclTensor* other, aclOpExecutor* executor); + +} + +#endif // OP_API_INC_LEVEL0_OP_MAXIMUM_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/minimum.h b/csrc/ascend/common/stub/op_api/level0/minimum.h new file mode 100644 index 000000000..a1d6effca --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/minimum.h @@ -0,0 +1,25 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file minimum.h +* \brief +*/ + +#ifndef OP_API_INC_LEVEL0_OP_MINIMUM_OP_H_ +#define OP_API_INC_LEVEL0_OP_MINIMUM_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor* Minimum(const aclTensor* self, const aclTensor* other, aclOpExecutor* executor); +} + +#endif // OP_API_INC_LEVEL0_OP_MINIMUM_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/mul.h b/csrc/ascend/common/stub/op_api/level0/mul.h new file mode 100644 index 000000000..d91da526f --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/mul.h @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file mul.h +* \brief +*/ + +#ifndef OP_API_INC_LEVEL0_MUL_H_ +#define OP_API_INC_LEVEL0_MUL_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor *Mul(const aclTensor *self, const aclTensor *other, aclOpExecutor *executor); + +} + +#endif // OP_API_INC_LEVEL0_MUL_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/muls.h b/csrc/ascend/common/stub/op_api/level0/muls.h new file mode 100644 index 000000000..9a677c0b6 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/muls.h @@ -0,0 +1,28 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + + /*! +* \file muls.h +* \brief +*/ + +#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_MULS_OP_H_ +#define PTA_NPU_OP_API_INC_LEVEL0_OP_MULS_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor *Muls(const aclTensor *self, + float alpha, + aclOpExecutor *executor); + +} + +#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_MULS_OP_H_ \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_api/level0/ones_like.h b/csrc/ascend/common/stub/op_api/level0/ones_like.h new file mode 100644 index 000000000..8e4664264 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/ones_like.h @@ -0,0 +1,27 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_ONES_LIKE_OP_H_ +#define PTA_NPU_OP_API_INC_LEVEL0_OP_ONES_LIKE_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +static const std::initializer_list AICPU_DTYPE_SUPPORT_LIST = { + op::DataType::DT_BOOL, op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_INT8, + op::DataType::DT_INT16, op::DataType::DT_UINT16, op::DataType::DT_UINT8, op::DataType::DT_INT32, + op::DataType::DT_INT64, op::DataType::DT_DOUBLE, op::DataType::DT_COMPLEX64, op::DataType::DT_COMPLEX128, + op::DataType::DT_BF16}; +const aclTensor *OnesLike(const aclTensor *self, aclOpExecutor *executor); +inline static bool IsAiCpuSupport(const aclTensor *self) { + return op::CheckType(self->GetDataType(), AICPU_DTYPE_SUPPORT_LIST); +} +} // namespace l0op + +#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_ONES_LIKE_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/padv3.h b/csrc/ascend/common/stub/op_api/level0/padv3.h new file mode 100644 index 000000000..b381e8467 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/padv3.h @@ -0,0 +1,25 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file padv3.h +* \brief +*/ +#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_PADV3_OP_H_ +#define PTA_NPU_OP_API_INC_LEVEL0_OP_PADV3_OP_H_ + +#include "opdev/op_executor.h" +#include "opdev/make_op_executor.h" + +namespace l0op { +const aclTensor *PadV3(const aclTensor *self, const aclTensor *paddings, const aclTensor *constant_values, + const std::string& mode, const bool paddingsContiguous, aclOpExecutor *executor); +} +#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_PADV3_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/reduce_mean.h b/csrc/ascend/common/stub/op_api/level0/reduce_mean.h new file mode 100644 index 000000000..af13b697f --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/reduce_mean.h @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_API_INC_LEVEL0_REDUCE_MEAN_H +#define OP_API_INC_LEVEL0_REDUCE_MEAN_H +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor* ReduceMean(const aclTensor* self, const aclIntArray* dim, bool keepDim, aclOpExecutor* executor); +const aclTensor* ReduceMean( + const aclTensor* self, const aclIntArray* dim, bool keepDim, bool noopWithEmptyAxes, aclOpExecutor* executor); +} // namespace l0op + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_api/level0/reduce_sum_op.h b/csrc/ascend/common/stub/op_api/level0/reduce_sum_op.h new file mode 100644 index 000000000..44b754cf0 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/reduce_sum_op.h @@ -0,0 +1,27 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file padv3.h +* \brief +*/ + +#ifndef OP_API_OP_API_COMMON_INC_LEVEL0_OP_REDUCE_SUM_OP_H_ +#define OP_API_OP_API_COMMON_INC_LEVEL0_OP_REDUCE_SUM_OP_H_ + +#include "opdev/op_executor.h" +#include "opdev/make_op_executor.h" + +namespace l0op { +const aclTensor *ReduceSumOp(const aclTensor *x, const aclIntArray *axes, + bool keep_dims, aclOpExecutor *executor); +} + +#endif // OP_API_OP_API_COMMON_INC_LEVEL0_OP_REDUCE_SUM_OP_H_ \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_api/level0/shape_op.h b/csrc/ascend/common/stub/op_api/level0/shape_op.h new file mode 100644 index 000000000..29169d75e --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/shape_op.h @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_API_INC_LEVEL0_SHAPE_OP_H +#define OP_API_INC_LEVEL0_SHAPE_OP_H +#include "opdev/op_executor.h" +namespace l0op { +const aclTensor *Shape_op(const aclTensor *x, aclOpExecutor *executor); +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_api/level0/sort.h b/csrc/ascend/common/stub/op_api/level0/sort.h new file mode 100644 index 000000000..d07d636fe --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/sort.h @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sort.h + * \brief + */ +#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_SORT_OP_H_ +#define PTA_NPU_OP_API_INC_LEVEL0_OP_SORT_OP_H_ + +#include "opdev/op_executor.h" +#include "opdev/fast_vector.h" + +namespace l0op { +const std::tuple Sort(const aclTensor* self, int64_t dim, bool descending, bool stable, + op::DataType indicesType, aclOpExecutor* executor); +} + +#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_SORT_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/squeeze.h b/csrc/ascend/common/stub/op_api/level0/squeeze.h new file mode 100644 index 000000000..9b1c9a526 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/squeeze.h @@ -0,0 +1,24 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_API_INC_LEVEL0_SQUEEZE_ND_H +#define OP_API_INC_LEVEL0_SQUEEZE_ND_H + +# include "opdev/op_def.h" + +namespace l0op { + +const aclTensor *SqueezeNd(const aclTensor *x, const aclIntArray* dim, aclOpExecutor *executor); + +const aclTensor *SqueezeNd(const aclTensor *x, int64_t dim, aclOpExecutor *executor); + +} // l0op + +#endif // OP_API_INC_LEVEL0_SQUEEZE_ND_H \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_api/level0/sub.h b/csrc/ascend/common/stub/op_api/level0/sub.h new file mode 100644 index 000000000..2d8acfe5b --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/sub.h @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_SUB_OP_H_ +#define PTA_NPU_OP_API_INC_LEVEL0_OP_SUB_OP_H_ + +#include "opdev/op_executor.h" +#include "opdev/make_op_executor.h" + +namespace l0op { +const aclTensor *Sub(const aclTensor *self, const aclTensor *other, aclOpExecutor *executor); + +} + +#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_SUB_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/level0/tensor_move.h b/csrc/ascend/common/stub/op_api/level0/tensor_move.h new file mode 100644 index 000000000..cba9b62e8 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/tensor_move.h @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tensor_move.h + * \brief + */ + +#ifndef OP_API_INC_LEVEL0_TENSOR_MOVE_H +#define OP_API_INC_LEVEL0_TENSOR_MOVE_H + +# include "opdev/op_def.h" + +namespace l0op { + +const aclTensor *TensorMove(const aclTensor *x, const aclTensor *y, aclOpExecutor *executor); +} // l0op + +#endif // OP_API_INC_LEVEL0_TENSOR_MOVE_H diff --git a/csrc/ascend/common/stub/op_api/level0/unsqueeze.h b/csrc/ascend/common/stub/op_api/level0/unsqueeze.h new file mode 100644 index 000000000..393e4348e --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/unsqueeze.h @@ -0,0 +1,29 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! +* \file unsqueeze.h +* \brief +*/ + +#ifndef OP_API_INC_LEVEL0_UNSQUEEZE_ND_H +#define OP_API_INC_LEVEL0_UNSQUEEZE_ND_H + +# include "opdev/op_def.h" + +namespace l0op { + +const aclTensor *UnsqueezeNd(const aclTensor *x, const aclIntArray* dim, aclOpExecutor *executor); + +const aclTensor *UnsqueezeNd(const aclTensor *x, int64_t dim, aclOpExecutor *executor); + +} // l0op + +#endif // OP_API_INC_LEVEL0_UNSQUEEZE_ND_H diff --git a/csrc/ascend/common/stub/op_api/level0/zero_op.h b/csrc/ascend/common/stub/op_api/level0/zero_op.h new file mode 100644 index 000000000..7dbddb69a --- /dev/null +++ b/csrc/ascend/common/stub/op_api/level0/zero_op.h @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_ZERO_OP_H_ +#define PTA_NPU_OP_API_INC_LEVEL0_OP_ZERO_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor *ZerosLike(const aclTensor *self, aclOpExecutor *executor); + +} + +#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_ZERO_OP_H_ diff --git a/csrc/ascend/common/stub/op_api/opapi_stub.cpp b/csrc/ascend/common/stub/op_api/opapi_stub.cpp new file mode 100644 index 000000000..b62a8f796 --- /dev/null +++ b/csrc/ascend/common/stub/op_api/opapi_stub.cpp @@ -0,0 +1,303 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file op_legacy_api.cpp + * \brief + */ + +#include "opdev/op_executor.h" +#include "aclnn_kernels/cast.h" +#include "aclnn_kernels/contiguous.h" +#include "aclnn_kernels/pad.h" +#include "aclnn_kernels/reshape.h" +#include "aclnn_kernels/slice.h" +#include "aclnn_kernels/transdata.h" +#include "aclnn_kernels/transpose.h" + +#include "level0/add.h" +#include "level0/axpy.h" +#include "level0/broadcast_to.h" +#include "level0/dot.h" +#include "level0/fill.h" +#include "level0/mul.h" +#include "level0/muls.h" +#include "level0/reduce_mean.h" +#include "level0/padv3.h" +#include "level0/sort.h" +#include "level0/dilation.h" +#include "level0/zero_op.h" +#include "level0/squeeze.h" +#include "level0/unsqueeze.h" + +namespace l0op { +const aclTensor *TensorMove(const aclTensor *x, const aclTensor * /*y*/, aclOpExecutor * /*executor*/) +{ + return x; +} +const aclTensor *ZerosLike(const aclTensor *self, aclOpExecutor * /*executor*/) +{ + return self; +} +const aclTensor *Maximum(const aclTensor *self, const aclTensor * /*other*/, aclOpExecutor * /*executor*/) +{ + return self; +} +const aclTensor *GatherV2(const aclTensor *self, int64_t /*axis*/, const aclTensor * /*indices*/, + aclOpExecutor * /*executor*/, int /*batchDims = 0*/, bool /*negativeIndexSupport = false*/) +{ + return self; +} + +const aclTensor *GatherV2WithImplMode(const aclTensor *self, int64_t /*axis*/, const aclTensor * /*indices*/, + int64_t /*implMode*/, aclOpExecutor * /*executor*/, int /*batchDims = 0*/, + bool /*negativeIndexSupport = false*/) +{ + return self; +} +const aclTensor *GatherElements(const aclTensor *self, const int64_t /*dim*/, const aclTensor * /*index*/, + aclOpExecutor * /*executor*/) +{ + return self; +} +const aclTensor *Minimum(const aclTensor *self, const aclTensor * /*other*/, aclOpExecutor * /*executor*/) +{ + return self; +} +const aclTensor *Cast(const aclTensor *self, op::DataType /*dstDtype*/, aclOpExecutor * /*executor*/) +{ + return self; +} + +const aclTensor *CastOnlyForConvBackward(const aclTensor *self, op::DataType /*dstDtype*/, aclOpExecutor * /*executor*/) +{ + return self; +} + +const aclTensor *Contiguous(const aclTensor *x, aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *ViewCopy(const aclTensor *x, const aclTensor * /*y*/, aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *PickViewAsContiguous(const aclTensor *x, aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *ReViewToOut(const aclTensor *x, const aclTensor * /*y*/, aclOpExecutor * /*executor*/) +{ + return x; +} + +const std::tuple Sort(const aclTensor * self, int64_t /*dim*/, bool /*descending*/, + bool /*stable*/, op::DataType /*indicesType*/, + aclOpExecutor * /*executor*/) +{ + return std::tuple(const_cast(self), const_cast(self)); +} + +bool CanOptimizeContiguous(const op::Shape & /*viewShape*/, const op::Strides & /*strides*/, int64_t /*offset*/, + int64_t /*storageSize*/, ContiguousParam & /*param*/) +{ + return true; +} + +bool CanOptimizeView(const op::Shape & /*viewShape*/, const op::Strides & /*strides*/, int64_t /*offset*/, + ContiguousParam & /*param*/) +{ + return true; +} + +const aclTensor *Pad(const aclTensor *self, const aclTensor * /*paddings*/, aclOpExecutor * /*executor*/) +{ + return self; +} + +const aclTensor *Reshape(const aclTensor *x, const op::Shape & /*shape*/, aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *Reshape(const aclTensor *x, const aclIntArray * /*shape*/, aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *Slice(const aclTensor *x, const aclTensor * /*y*/, const aclTensor * /*offset*/, + const aclTensor * /*size*/, aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *Slice(const aclTensor *x, const aclIntArray * /*offsets*/, const aclIntArray * /*size*/, + aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *ReFormat(const aclTensor *x, const op::Format & /*format*/, aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *TransData(const aclTensor *x, op::Format /*dstPrimaryFormat*/, int64_t /*groups*/, + aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *TransDataSpecial(const aclTensor *x, op::Format /*dstPrimaryFormat*/, int64_t /*groups*/, + aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *Transpose(const aclTensor *x, const aclTensor * /*y*/, const aclTensor * /*perm*/, + aclOpExecutor * /*executor*/) +{ + return x; +} +const aclTensor *Transpose(const aclTensor *x, const aclIntArray * /*perm*/, aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *Add(const aclTensor *self, const aclTensor * /*other*/, aclOpExecutor * /*executor*/) +{ + return self; +} +const aclTensor *Axpy(const aclTensor *self, const aclTensor * /*other*/, float /*alpha*/, aclOpExecutor * /*executor*/) +{ + return self; +} + +const aclTensor *BroadcastTo(const aclTensor *x, const aclTensor * /*y*/, const aclTensor * /*shape*/, + aclOpExecutor * /*executor*/) +{ + return x; +} +const aclTensor *BroadcastTo(const aclTensor *x, const aclIntArray * /*shape*/, aclOpExecutor * /*executor*/) +{ + return x; +} +const aclTensor *Dot(const aclTensor *self, const aclTensor * /*tensor*/, aclOpExecutor * /*executor*/) +{ + return self; +} +const aclTensor *Fill(const aclTensor * /*dims*/, const aclTensor *value, const aclIntArray * /*outShape*/, + aclOpExecutor * /*executor*/) +{ + return value; +} +const aclTensor *Mul(const aclTensor *self, const aclTensor * /*other*/, aclOpExecutor * /*executor*/) +{ + return self; +} +const aclTensor *Muls(const aclTensor *self, float /*alpha*/, aclOpExecutor * /*executor*/) +{ + return self; +} +const aclTensor *ReduceMean(const aclTensor *self, const aclIntArray * /*dim*/, bool /*keepDim*/, + aclOpExecutor * /*executor*/) +{ + return self; +} +const aclTensor *ReduceMean(const aclTensor *self, const aclIntArray * /*dim*/, bool /*keepDim*/, + bool /*noopWithEmptyAxes*/, aclOpExecutor * /*executor*/) +{ + return self; +} +const aclTensor *Dilation(const aclTensor *x, const aclIntArray * /*dilations*/, const aclIntArray * /*pads*/, + float /*paddingValue*/, aclOpExecutor * /*executor*/) +{ + return x; +} +const aclTensor *Shape_op(const aclTensor *x, aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *SqueezeNd(const aclTensor *x, const aclIntArray * /*dim*/, aclOpExecutor * /*executor*/) +{ + return x; +} +const aclTensor *SqueezeNd(const aclTensor *x, int64_t /*dim*/, aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *PadV3(const aclTensor *self, const aclTensor * /*paddings*/, const aclTensor * /*constant_values*/, + const std::string & /*mode*/, const bool /*paddingsContiguous*/, aclOpExecutor * /*executor*/) +{ + return self; +} + +const aclTensor *UnsqueezeNd(const aclTensor *x, const aclIntArray * /*dim*/, aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *UnsqueezeNd(const aclTensor *x, int64_t /*dim*/, aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *ReduceSumOp(const aclTensor *x, const aclIntArray * /*axes*/, bool /*keep_dims*/, + aclOpExecutor * /*executor*/) +{ + return x; +} + +const aclTensor *MaskedScatter(const aclTensor * self, const aclTensor * /*mask*/, const aclTensor * /*source*/, + aclOpExecutor * /*executor*/) +{ + return self; +} + +const aclTensor *InplaceIndexAddAiCore(const aclTensor * self, const int64_t /*dim*/, const aclTensor * /*index*/, + const aclTensor * /*source*/, const aclTensor * /*alphaTensor*/, + aclOpExecutor * /*executor*/) +{ + return self; +} + +const aclTensor *InplaceIndexAddAiCpu(const aclTensor * self, const int64_t /*dim*/, const aclTensor * /*index*/, + const aclTensor * /*source*/, const aclTensor * /*alphaTensor*/, + aclOpExecutor * /*executor*/) +{ + return self; +} + +const aclTensor *InplaceIndexAddWithSorted(const aclTensor * self, const int64_t /*dim*/, + const aclTensor * /*sortedIndices*/, const aclTensor * /*pos*/, + const aclTensor * /*value*/, const aclTensor * /*alphaTensor*/, + aclOpExecutor * /*executor*/) +{ + return self; +} + +const aclTensor *GatherV3(const aclTensor *self, int64_t axis, const aclTensor *indices, aclOpExecutor *executor, + int batchDims = 0, bool negativeIndexSupport = false) +{ + (void)self; + (void)axis; + (void)indices; + (void)executor; + (void)batchDims; + (void)negativeIndexSupport; + return self; +} + +} // namespace l0op diff --git a/csrc/ascend/common/stub/op_tiling/CMakeLists.txt b/csrc/ascend/common/stub/op_tiling/CMakeLists.txt new file mode 100644 index 000000000..cb0c97e52 --- /dev/null +++ b/csrc/ascend/common/stub/op_tiling/CMakeLists.txt @@ -0,0 +1,30 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}/build) + +file(GLOB_RECURSE OP_TILING_FILES "*.cpp") + +add_library(optiling SHARED ${OP_TILING_FILES}) + +add_dependencies(optiling json) + +target_compile_definitions(optiling PRIVATE + _GLIBCXX_USE_CXX11_ABI=0 + LOG_CPP +) + +target_include_directories(optiling PRIVATE + ${OP_TILING_INCLUDE} +) + +target_link_libraries(optiling PRIVATE + $ +) \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_tiling/op_cache_def_tiling.h b/csrc/ascend/common/stub/op_tiling/op_cache_def_tiling.h new file mode 100644 index 000000000..8905b8f73 --- /dev/null +++ b/csrc/ascend/common/stub/op_tiling/op_cache_def_tiling.h @@ -0,0 +1,198 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file op_cache_tiling.h + * \brief + */ + +#ifndef OPS_BUILT_IN_OP_TILING_OP_CACHE_DEF_TILING_H +#define OPS_BUILT_IN_OP_TILING_OP_CACHE_DEF_TILING_H + +#include +#include "exe_graph/runtime/tiling_context.h" + +namespace optiling { +struct BatchmatmulCompileParas { + bool binary_mode_flag = false; + bool bias_flag = false; + bool at_l1_flag = true; + bool split_k_flag = false; + bool pattern_flag = false; + bool zero_flag = false; + bool sparse_4to2_flag = false; + bool binary_constant_flag = false; + bool vector_pre_conv_mode = false; + float fused_double_operand_num = 0; + float aub_double_num = 0; + float bub_double_num = 0; + int64_t quant_scale = 0; + int64_t eltwise_src = 0; + int8_t enable_pad = 0; + bool enable_nz_fusion = false; + bool enable_rt_bank_cache = false; +}; + +struct BatchmatmulRunParas { + bool nd_flag = false; + bool use_pre_ub = false; + bool trans_a_flag = false; + bool trans_b_flag = false; + bool format_a_nd = false; + bool format_b_nd = false; + bool format_out_nd = false; + ge::Format format_a = ge::FORMAT_ND; + ge::Format format_b = ge::FORMAT_ND; + ge::Format format_out = ge::FORMAT_ND; + bool reserved_bool = false; + bool b_have_batch = false; // dim num > 2 + bool is_batch_matmul_mode = false; // dynamic_mode == "dynamic_mknb" + bool is_batch_matmul_op = false; // BatchMatMulV2 or BatchMatMul + bool used_aligned_pattern = false; + bool non_factor_k = false; + bool non_factor_bmn = false; + bool bias_flag = false; + bool pattern_flag = false; + bool do_not_multi_batch = false; + bool performance_flag = false; + bool unaligned_flag = false; + bool zero_flag = false; + bool is_compress_quant = false; + bool is_bmm_fixp = false; + bool enable_nz_fusion = false; + bool weight_nz_flag = false; + int8_t enable_pad = 0; + int8_t hf32_flag = 1; + int8_t pad_flag = 0; + int8_t nz_fusion_flag = 0; + int32_t dtype_a = 0; + int32_t dtype_b = 0; + int32_t dtype_out = 0; + int32_t dtype_bias = 0; + int64_t m_mapped = 1; + int64_t k_mapped = 1; + int64_t n_mapped = 1; + int64_t batch_mapped = 1; + int64_t m = 1; + int64_t k = 1; + int64_t n = 1; + int64_t batch = 1; + int64_t ori_shape_m = 1; + int64_t ori_shape_k = 1; + int64_t ori_shape_n = 1; + int64_t m_pad = 0; + int64_t k_pad = 0; + int64_t n_pad = 0; + int64_t nl0 = 1; + int64_t kl0 = 1; + int64_t dim0_a = 0; + int64_t dim1_a = 0; + int64_t dim2_a = 0; + int64_t dim0_b = 0; + int64_t dim1_b = 0; + int64_t dim2_b = 0; + int64_t batch_a1 = 1; + int64_t batch_a2 = 1; + int64_t batch_a3 = 1; + int64_t batch_a4 = 1; + int64_t batch_b1 = 1; + int64_t batch_b2 = 1; + int64_t batch_b3 = 1; + int64_t batch_b4 = 1; + int64_t batch_c1 = 1; + int64_t batch_c2 = 1; + int64_t batch_c3 = 1; + int64_t batch_c4 = 1; + int32_t offset_x = 0; + int32_t index_size = 0; + bool m_quant_check = false; + bool n_quant_check = false; + bool is_weight_quant_bmm = false; + bool vector_pre_conv_mode = false; + bool is_quant_batch_matmul_v3 = false; + bool is_weight_quant_batch_matmul_v2 = false; + bool is_pertoken = false; + // 3 is perm_a dim + std::array perm_a = {0, 0, 0}; + // 3 is perm_b dim + std::array perm_b = {0, 0, 0}; + ge::DataType bias_dtype = ge::DT_FLOAT16; +}; + +class CacheTilingData +{ +public: + uint64_t tiling_id; + int64_t n_cub = 1; + int64_t db_cub = 1; + int64_t m_l0 = 1; + int64_t k_l0 = 1; + int64_t n_l0 = 1; + int64_t batch_dim = 1; + int64_t n_dim = 1; + int64_t m_dim = 1; + int64_t k_dim = 1; + int64_t kal1_16 = 1; + int64_t kbl1_16 = 1; + int64_t kal1_factor = 1; + int64_t kbl1_factor = 1; + int64_t m_al1 = 1; + int64_t n_bl1 = 1; + int64_t db_al1 = 1; + int64_t db_bl1 = 1; + int64_t k_aub = 1; + int64_t m_aub = 1; + int64_t db_aub = 1; + int64_t k_bub = 1; + int64_t n_bub = 1; + int64_t db_bub = 1; + int64_t aub_dim = 1; + int64_t bub_dim = 1; + int64_t m1_aub = 1; + int64_t n1_bub = 1; + int64_t k1_aub = 1; + int64_t k1_bub = 1; + int64_t m_aub_dim = 1; + int64_t n_bub_dim = 1; + int64_t k_aub_dim = 1; + int64_t k_bub_dim = 1; + int64_t k_org_dim = 1; + int64_t db_l0c = 1; + int64_t batch_l0 = 1; + int64_t batch_aub = 1; + int64_t batch_bub = 1; + int64_t batch_cub = 1; + int32_t out_branch_flag = 1; + int32_t bias_flag = 0; + int32_t aub_multi_flag = 0; + int32_t bub_multi_flag = 0; + int64_t a_align_value = 1; + int64_t b_align_value = 1; + int64_t aub_align_bound = 0; + int64_t bub_align_bound = 0; + int64_t min_kl1_cmp_kl0 = 0; + int32_t al1_attach_flag = 0; + int32_t bl1_attach_flag = 0; + int32_t abkl1_attach_flag = 0; + int32_t l0c_multi_batch = 0; + int64_t m_single_core = 1; + int64_t n_single_core = 1; + bool flag_cub_solving_bank_conflict = false; + bool al1_full_load = false; + bool bl1_full_load = false; + int8_t hf32_flag = 1; + int32_t zero_flag = 0; + bool datatype_bf16 = false; + uint64_t deq_scale_var = 0x3F800000; + uint32_t l2_cache_flag = 0; +}; +} // namespace optiling + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_tiling/op_cache_tiling.cpp b/csrc/ascend/common/stub/op_tiling/op_cache_tiling.cpp new file mode 100644 index 000000000..beb8101cd --- /dev/null +++ b/csrc/ascend/common/stub/op_tiling/op_cache_tiling.cpp @@ -0,0 +1,35 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +/*! + * \file op_cache_tiling.cpp + * \brief + */ + +#include "op_cache_tiling.h" + +namespace optiling { +bool TilingPrepareForOpCache(gert::TilingContext* /*context*/) +{ + return true; +} + +bool TilingPrepareForOpCache(gert::TilingParseContext* /*context*/) +{ + return true; +} + +bool GenTiling( + const std::string& /*op_type*/, const BatchmatmulCompileParas& /*compile_params*/, + BatchmatmulRunParas& /*run_params*/, CacheTilingData& /*tiling*/, gert::TilingContext* /*context*/) +{ + return true; +} + +} // namespace optiling diff --git a/csrc/ascend/common/stub/op_tiling/op_cache_tiling.h b/csrc/ascend/common/stub/op_tiling/op_cache_tiling.h new file mode 100644 index 000000000..836391bea --- /dev/null +++ b/csrc/ascend/common/stub/op_tiling/op_cache_tiling.h @@ -0,0 +1,34 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file cop_ache_tiling.h + * \brief + */ + +#ifndef OPS_BUILT_IN_OP_TILING_OP_CACHE_TILING_H +#define OPS_BUILT_IN_OP_TILING_OP_CACHE_TILING_H + +#include +#include "exe_graph/runtime/tiling_context.h" +#include "exe_graph/runtime/tiling_parse_context.h" +#include "op_cache_def_tiling.h" + +namespace optiling { + +bool TilingPrepareForOpCache(gert::TilingContext* context); +bool TilingPrepareForOpCache(gert::TilingParseContext* context); + +bool GenTiling( + const std::string& op_type, const BatchmatmulCompileParas& compile_params, BatchmatmulRunParas& run_params, + CacheTilingData& tiling, gert::TilingContext* context); +} // namespace optiling + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_tiling/register/tuning_bank_key_registry.h b/csrc/ascend/common/stub/op_tiling/register/tuning_bank_key_registry.h new file mode 100644 index 000000000..c54030ee7 --- /dev/null +++ b/csrc/ascend/common/stub/op_tiling/register/tuning_bank_key_registry.h @@ -0,0 +1,209 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef __INC_REGISTER_TUNING_BANK_KEY_REGISTRY_HEADER__ +#define __INC_REGISTER_TUNING_BANK_KEY_REGISTRY_HEADER__ +#include +#include +#include +#include +#include "graph/ascend_string.h" +#include "register/register_types.h" +#include "exe_graph/runtime/tiling_context.h" + +// v1 stub +#define REGISTER_OP_BANK_KEY_CONVERT_FUN(op, opfunc) REGISTER_OP_BANK_KEY_CONVERT_FUN_UNIQ_HELPER(op, (opfunc)) + +#define REGISTER_OP_BANK_KEY_CONVERT_FUN_UNIQ_HELPER(optype, opfunc) REGISTER_OP_BANK_KEY_UNIQ(optype, (opfunc)) + +#define REGISTER_OP_BANK_KEY_UNIQ(optype, opfunc) \ + static tuningtiling::OpBankKeyFuncRegistry g_##optype##BankKeyRegistryInterf(#optype, (opfunc)) + +#define REGISTER_OP_BANK_KEY_PARSE_FUN(op, parse_func, load_func) \ + REGISTER_OP_BANK_KEY_PARSE_FUN_UNIQ_HELPER(op, (parse_func), (load_func)) + +#define REGISTER_OP_BANK_KEY_PARSE_FUN_UNIQ_HELPER(optype, parse_func, load_func) \ + REGISTER_OP_BANK_KEY_PARSE_UNIQ(optype, (parse_func), (load_func)) + +#define REGISTER_OP_BANK_KEY_PARSE_UNIQ(optype, parse_func, load_func) \ + static tuningtiling::OpBankKeyFuncRegistry g_##optype##BankParseInterf(#optype, (parse_func), (load_func)) + +// v2 +#define REGISTER_OP_BANK_KEY_CONVERT_FUN_V2(op, opfunc) REGISTER_OP_BANK_KEY_CONVERT_FUN_UNIQ_HELPER_V2(op, (opfunc)) + +#define REGISTER_OP_BANK_KEY_CONVERT_FUN_UNIQ_HELPER_V2(optype, opfunc) REGISTER_OP_BANK_KEY_UNIQ_V2(optype, (opfunc)) + +#define REGISTER_OP_BANK_KEY_UNIQ_V2(optype, opfunc) \ + static tuningtiling::OpBankKeyFuncRegistryV2 g_##optype##BankKeyRegistryInterf(#optype, (opfunc)) + +#define REGISTER_OP_BANK_KEY_PARSE_FUN_V2(op, parse_func, load_func) \ + REGISTER_OP_BANK_KEY_PARSE_FUN_UNIQ_HELPER_V2(op, (parse_func), (load_func)) + +#define REGISTER_OP_BANK_KEY_PARSE_FUN_UNIQ_HELPER_V2(optype, parse_func, load_func) \ + REGISTER_OP_BANK_KEY_PARSE_UNIQ_V2(optype, (parse_func), (load_func)) + +#define REGISTER_OP_BANK_KEY_PARSE_UNIQ_V2(optype, parse_func, load_func) \ + static tuningtiling::OpBankKeyFuncRegistryV2 g_##optype##BankParseInterf(#optype, (parse_func), (load_func)) + +#define TUNING_TILING_MAKE_SHARED(exec_expr0, exec_expr1) \ + do { \ + try { \ + exec_expr0; \ + } catch (...) { \ + exec_expr1; \ + } \ + } while (0) + +// v1 stub +#define DECLARE_STRUCT_RELATE_WITH_OP(op, bank_key, ...) \ + do { \ + NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(bank_key, __VA_ARGS__); \ + static bool ParseFunc##op##bank_key( \ + const std::shared_ptr& in_args, size_t len, ge::AscendString& bank_key_str) \ + { \ + if (sizeof(bank_key_str) != len || in_args == nullptr) { \ + return false; \ + } \ + return false; \ + } \ + static bool LoadFunc##op##bank_key( \ + std::shared_ptr& in_args, size_t& len, const ge::AscendString& bank_key_str) \ + { \ + len = sizeof(bank_key_str); \ + TUNING_TILING_MAKE_SHARED(in_args = std::make_shared(), return false); \ + auto op_ky = std::static_pointer_cast(in_args); \ + return false; \ + } \ + REGISTER_OP_BANK_KEY_PARSE_FUN(op, ParseFunc##op##bank_key, LoadFunc##op##bank_key) \ + } while (0) + + +// v2 +#define DECLARE_STRUCT_RELATE_WITH_OP_V2(op, bank_key, ...) \ + do { \ + NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(bank_key, __VA_ARGS__); \ + static bool ParseFuncV2##op##bank_key( \ + const std::shared_ptr& in_args, size_t len, ge::AscendString& bank_key_json_str) \ + { \ + if (sizeof(bank_key) != len || in_args == nullptr) { \ + return false; \ + } \ + nlohmann::json bank_key_json; \ + bank_key_json = *(std::static_pointer_cast(in_args)); \ + try { \ + std::string json_dump_str = bank_key_json.dump(); \ + bank_key_json_str = ge::AscendString(json_dump_str.c_str()); \ + } catch (std::exception & e) { \ + return false; \ + } \ + return true; \ + } \ + static bool LoadFuncV2##op##bank_key( \ + std::shared_ptr& in_args, size_t& len, const ge::AscendString& bank_key_json_str) \ + { \ + len = sizeof(bank_key); \ + TUNING_TILING_MAKE_SHARED(in_args = std::make_shared(), return false); \ + nlohmann::json bank_key_json; \ + try { \ + bank_key_json = nlohmann::json::parse(bank_key_json_str.GetString()); \ + auto op_ky = std::static_pointer_cast(in_args); \ + *op_ky = bank_key_json.get(); \ + } catch (std::exception & e) { \ + return false; \ + } \ + return true; \ + } \ + REGISTER_OP_BANK_KEY_PARSE_FUN_V2(op, ParseFuncV2##op##bank_key, LoadFuncV2##op##bank_key) \ + } while (0) + + +namespace tuningtiling { +// v1兼容老版本om +using OpBankKeyConvertFun = std::function&, size_t&)>; +using OpBankParseFun = std::function&, size_t, ge::AscendString&)>; +using OpBankLoadFun = std::function&, size_t&, const ge::AscendString&)>; + +// v2 +using OpBankKeyConvertFunV2 = std::function&, size_t&)>; +using OpBankParseFunV2 = std::function&, size_t, ge::AscendString&)>; +using OpBankLoadFunV2 = std::function&, size_t&, const ge::AscendString&)>; +// v1兼容老版本om +class FMK_FUNC_HOST_VISIBILITY OpBankKeyFuncInfo +{ +public: + explicit OpBankKeyFuncInfo(const ge::AscendString& optype); + OpBankKeyFuncInfo() = default; + ~OpBankKeyFuncInfo() = default; + void SetOpConvertFunc(const OpBankKeyConvertFun& convert_func); + void SetOpParseFunc(const OpBankParseFun& parse_func); + void SetOpLoadFunc(const OpBankLoadFun& load_func); + const OpBankKeyConvertFun& GetBankKeyConvertFunc() const; + const OpBankParseFun& GetBankKeyParseFunc() const; + const OpBankLoadFun& GetBankKeyLoadFunc() const; + const ge::AscendString& GetOpType() const + { + return optype_; + } + +private: + ge::AscendString optype_; + OpBankKeyConvertFun convert_func_; + OpBankParseFun parse_func_; + OpBankLoadFun load_func_; +}; + +// v2 +class FMK_FUNC_HOST_VISIBILITY OpBankKeyFuncInfoV2 +{ +public: + explicit OpBankKeyFuncInfoV2(const ge::AscendString& optypeV2); + OpBankKeyFuncInfoV2() = default; + ~OpBankKeyFuncInfoV2() = default; + void SetOpConvertFuncV2(const OpBankKeyConvertFunV2& convert_funcV2); + void SetOpParseFuncV2(const OpBankParseFunV2& parse_funcV2); + void SetOpLoadFuncV2(const OpBankLoadFunV2& load_funcV2); + const OpBankKeyConvertFunV2& GetBankKeyConvertFuncV2() const; + const OpBankParseFunV2& GetBankKeyParseFuncV2() const; + const OpBankLoadFunV2& GetBankKeyLoadFuncV2() const; + const ge::AscendString& GetOpTypeV2() const + { + return optypeV2_; + } + +private: + ge::AscendString optypeV2_; + OpBankKeyConvertFunV2 convert_funcV2_; + OpBankParseFunV2 parse_funcV2_; + OpBankLoadFunV2 load_funcV2_; +}; + +// v1兼容老版本om +class FMK_FUNC_HOST_VISIBILITY OpBankKeyFuncRegistry +{ +public: + OpBankKeyFuncRegistry(const ge::AscendString& optype, const OpBankKeyConvertFun& convert_func); + OpBankKeyFuncRegistry( + const ge::AscendString& optype, const OpBankParseFun& parse_func, const OpBankLoadFun& load_func); + ~OpBankKeyFuncRegistry() = default; + static std::unordered_map& RegisteredOpFuncInfo(); +}; + +// v2 +class FMK_FUNC_HOST_VISIBILITY OpBankKeyFuncRegistryV2 +{ +public: + OpBankKeyFuncRegistryV2(const ge::AscendString& optype, const OpBankKeyConvertFunV2& convert_funcV2); + OpBankKeyFuncRegistryV2( + const ge::AscendString& optype, const OpBankParseFunV2& parse_funcV2, const OpBankLoadFunV2& load_funcV2); + ~OpBankKeyFuncRegistryV2() = default; + static std::unordered_map& RegisteredOpFuncInfoV2(); +}; +} // namespace tuningtiling +#endif \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_tiling/register/tuning_tiling_reflection_utils.h b/csrc/ascend/common/stub/op_tiling/register/tuning_tiling_reflection_utils.h new file mode 100644 index 000000000..3ce90436d --- /dev/null +++ b/csrc/ascend/common/stub/op_tiling/register/tuning_tiling_reflection_utils.h @@ -0,0 +1,191 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef __INC_REGISTER_TUNING_TILING_REFLECTION_UTILS_HEADER__ +#define __INC_REGISTER_TUNING_TILING_REFLECTION_UTILS_HEADER__ +#include +#include +#include +#include + +namespace tuningtiling { +// implement for std c++11 +template +using decay_t = typename std::decay::type; + +template +using enable_if_t = typename std::enable_if::type; + +template +struct integer_sequence { + using value_type = T; + static constexpr std::size_t size() + { + return sizeof...(Ints); + } +}; + +template +using index_sequence = integer_sequence; + +template +struct make_integer_sequence : make_integer_sequence { +}; + +template +struct make_integer_sequence : integer_sequence { +}; + +template +using make_index_sequence = make_integer_sequence; + +template +struct StructInfo { + static std::tuple<> Info() + { + return std::make_tuple(); + } +}; + +#define DECLARE_SCHEMA(Struct, ...) \ + template <> \ + struct StructInfo { \ + static decltype(std::make_tuple(__VA_ARGS__)) Info() \ + { \ + return std::make_tuple(__VA_ARGS__); \ + } \ + }; + +#define FIELD(class, FieldName) std::make_tuple(#FieldName, &class ::FieldName) + +template +void ForEachTuple(Tuple&& tuple, Field&& fields, Fn&& fn, index_sequence) +{ + (void)std::initializer_list{ + (fn(std::get<0>(std::get(fields)), tuple.*std::get<1>(std::get(fields))), Is)...}; +} + +template +void ForEachTuple(Tuple&& tuple, Fn&& fn) +{ + const auto fields = StructInfo>::Info(); + ForEachTuple( + std::forward(tuple), fields, std::forward(fn), + make_index_sequence::value>{}); +} + +template +struct is_optional : std::false_type { +}; + +template +struct is_optional> : std::true_type { +}; + +template +bool is_optional_v() +{ + return is_optional>::value; +} + +template +decltype(std::begin(T()), std::true_type{}) containable(size_t); + +template +std::false_type containable(...); + +template +using is_containable = decltype(containable(0U)); + +template +constexpr bool IsSerializeType() +{ + return ((!std::is_class>::value) || is_containable>()); +} + +template +void ForEachField(T&& value, Fn&& fn) +{ + ForEachTuple(std::forward(value), std::forward(fn)); +} + +template +struct DumpFunctor; + +template ()>* = nullptr> +void DumpObj(T&& obj, const std::string& field_name, Js& j) +{ + if (field_name.empty()) { + ForEachField(std::forward(obj), DumpFunctor(j)); + return; + } + ForEachField(std::forward(obj), DumpFunctor(j[field_name])); +} + +template ()>* = nullptr> +void DumpObj(T&& obj, const std::string& field_name, Js& j) +{ + if (field_name.empty()) { + return; + } + j[field_name] = std::forward(obj); +} + +template +struct DumpFunctor { + explicit DumpFunctor(T& j) : js(j) + {} + template + void operator()(Name&& name, Field&& field) const + { + DumpObj(std::forward(field), std::forward(name), js); + } + T& js; +}; + +template +struct FromJsonFunctor; + +template ()>* = nullptr> +void FromJsonImpl(T&& obj, const std::string& field_name, const Js& j) +{ + if (field_name.empty()) { + ForEachField(std::forward(obj), FromJsonFunctor(j)); + return; + } + if (j.find(field_name) == j.cend()) { + return; + } + ForEachField(std::forward(obj), FromJsonFunctor(j[field_name])); +} + +template ()>* = nullptr> +void FromJsonImpl(T&& obj, const std::string& field_name, const Js& j) +{ + // ignore missing field of optional + if ((tuningtiling::is_optional_v()) || (j.find(field_name) == j.cend())) { + return; + } + j.at(field_name).get_to(std::forward(obj)); +} + +template +struct FromJsonFunctor { + explicit FromJsonFunctor(const Js& j) : js(j) + {} + template + void operator()(Name&& name, Field&& field) const + { + FromJsonImpl(std::forward(field), std::forward(name), js); + } + const Js& js; +}; +} // namespace tuningtiling +#endif \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_tiling/register/tuning_tiling_registry.h b/csrc/ascend/common/stub/op_tiling/register/tuning_tiling_registry.h new file mode 100644 index 000000000..dd55282be --- /dev/null +++ b/csrc/ascend/common/stub/op_tiling/register/tuning_tiling_registry.h @@ -0,0 +1,111 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef __INC_REGISTER_TUNING_TILING_REGISTRY_HEADER__ +#define __INC_REGISTER_TUNING_TILING_REGISTRY_HEADER__ +#include +#include +#include +#include +#include "graph/ascend_string.h" +#include "register/tuning_tiling_reflection_utils.h" +namespace tuningtiling { +struct TilingItem { + ge::AscendString dtype_; + ge::AscendString name_; +}; + +class TuningTilingDef +{ +public: + virtual void FromJson(const nlohmann::json& j) = 0; + virtual void ToJson(nlohmann::json& j) = 0; + ge::AscendString GetClassName() const; + virtual std::vector GetItemInfo() const = 0; + +protected: + TuningTilingDef() = default; + virtual ~TuningTilingDef() = default; + // dtype , name + std::vector field_info_; + ge::AscendString class_name_; +}; + +#define BEGIN_TUNING_TILING_DEF(class_name) \ + class class_name : public TuningTilingDef \ + { \ + public: \ + virtual void FromJson(const nlohmann::json& j) \ + { \ + FromJsonImpl(*this, "", j); \ + } \ + \ + virtual void ToJson(nlohmann::json& j) \ + { \ + DumpObj(*this, "", j); \ + } \ + \ + std::vector GetItemInfo() const \ + { \ + return field_info_; \ + } \ + \ + class FieldHandler \ + { \ + public: \ + FieldHandler(class_name* pinstance, const ge::AscendString& dtype, const ge::AscendString& name) \ + { \ + pinstance->field_info_.push_back({dtype, name}); \ + } \ + }; \ + friend class FieldHandler; \ + \ + public: \ + class_name() \ + { \ + class_name_ = #class_name; \ + }; + +#define TUNING_TILING_DATA_FIELD_DEF(data_type, field_name) \ +public: \ + data_type field_name; \ + FieldHandler field_name##_handler_ = FieldHandler(this, #data_type, #field_name); + +#define END_TUNING_TILING_DEF \ + } \ + ; + +using TuningTilingDefConstructor = std::shared_ptr (*)(); +class TuningTilingClassFactory +{ +public: + static std::map& RegisterInfo(); + static void RegisterTilingData(const ge::AscendString& optype, TuningTilingDefConstructor const constructor); + static std::shared_ptr CreateTilingDataInstance(const ge::AscendString& optype); +}; + +#define REGISTER_TUNING_TILING_CLASS(optype, class_name) \ + class optype##Helper \ + { \ + public: \ + optype##Helper() \ + { \ + TuningTilingClassFactory::RegisterTilingData(#optype, optype##Helper::CreateTilingDataInstance); \ + } \ + static std::shared_ptr CreateTilingDataInstance() \ + { \ + return std::make_shared(); \ + } \ + }; \ + optype##Helper g_tuning_tiling_##optype##Helper; +using TuningTilingDefPtr = std::shared_ptr; +} // namespace tuningtiling + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_tiling/runtime_kb_api.cpp b/csrc/ascend/common/stub/op_tiling/runtime_kb_api.cpp new file mode 100644 index 000000000..4d02aef5e --- /dev/null +++ b/csrc/ascend/common/stub/op_tiling/runtime_kb_api.cpp @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "runtime_kb_api.h" + +namespace RuntimeKb { +uint32_t QueryBank( + const void* /*src*/, size_t /*src_len*/, const std::string& /*op_type*/, const std::string& /*soc_version*/, + uint32_t /*core_num*/, tuningtiling::TuningTilingDefPtr& /*tiling*/) +{ + return 0; +} +} // namespace RuntimeKb \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_tiling/runtime_kb_api.h b/csrc/ascend/common/stub/op_tiling/runtime_kb_api.h new file mode 100644 index 000000000..c01a5007c --- /dev/null +++ b/csrc/ascend/common/stub/op_tiling/runtime_kb_api.h @@ -0,0 +1,23 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef RUNTIME_KB_RUNTIME_KB_API_H +#define RUNTIME_KB_RUNTIME_KB_API_H + +#include +#include "exe_graph/runtime/tiling_context.h" +#include "register/tuning_tiling_registry.h" + +namespace RuntimeKb { +uint32_t QueryBank(const void *src, size_t src_len, const std::string &op_type, const std::string &soc_version, + uint32_t core_num, tuningtiling::TuningTilingDefPtr &tiling); +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_tiling/tbe_tiling_api.cpp b/csrc/ascend/common/stub/op_tiling/tbe_tiling_api.cpp new file mode 100644 index 000000000..ef93a414e --- /dev/null +++ b/csrc/ascend/common/stub/op_tiling/tbe_tiling_api.cpp @@ -0,0 +1,39 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#include "tbe_tiling_api.h" + +using namespace optiling; + +namespace optiling { +bool GetTbeTiling(const gert::TilingContext* context, Conv3dBpFilterV2RunInfo& runInfoForV2, Conv3dBackpropV2TBETilingData& tbeTilingForV2) +{ + (void)context; + (void)runInfoForV2; + (void)tbeTilingForV2; + return true; +} +bool GetTbeTiling(gert::TilingContext* context, Conv3dBpInputV2RunInfo& runInfoV2, + Conv3dBackpropV2TBETilingData& tbeTilingForV2, const optiling::OpTypeV2 opType) +{ + (void)context; + (void)runInfoV2; + (void)tbeTilingForV2; + (void)opType; + return true; +} + +bool GetTbeTiling(gert::TilingContext* context, Conv3dBackpropV2TBETilingData& tbeTilingForV2, const optiling::OpTypeV2 opType) +{ + (void)context; + (void)tbeTilingForV2; + (void)opType; + return true; +} +} \ No newline at end of file diff --git a/csrc/ascend/common/stub/op_tiling/tbe_tiling_api.h b/csrc/ascend/common/stub/op_tiling/tbe_tiling_api.h new file mode 100644 index 000000000..0a0f54fbb --- /dev/null +++ b/csrc/ascend/common/stub/op_tiling/tbe_tiling_api.h @@ -0,0 +1,193 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tbe_tiling_api.h + * \brief + */ +#ifndef TBE_TILING_API_H +#define TBE_TILING_API_H + +#include +#include +#include +#include "graph/utils/type_utils.h" +#include "platform/platform_infos_def.h" + +namespace optiling { +struct Conv3dBackpropV2TBETilingData { + // L0 tiling parameters + int32_t m_l0; // Base M dimension at L0 + int32_t k_l0; // Base K dimension at L0 + int32_t n_l0; // Base N dimension at L0 + + // L1 tiling parameters + int32_t m_al1; // Step M dimension at L1 + int32_t n_bl1; // Step N dimension at L1 + int32_t k_al1; // Step K dimension A at L1 + int32_t k_bl1; // Step K dimension B at L1 + + // Buffer parameters + int32_t db_l0c; // L0C buffer size + int32_t db_al1; // AL1 buffer size + int32_t db_bl1; // BL1 buffer size + + // Dimension parameters + int32_t batch_dim; // Batch dimension + int32_t d_dim; // Depth dimension + int32_t group_dim; // Group dimension + int32_t m_dim; // M dimension + int32_t n_dim; // N dimension + int32_t k_dim; // K dimension +}; + +struct Conv3dBpFilterV2RunInfo { + int32_t batch; + int32_t co; // output channels + int32_t ci; // input channels + int32_t cout1_g; // output channels per group + int32_t cin1_g; // input channels per group + int32_t dout; // output depth // codespell:ignore dout + int32_t wo; // output width + int32_t ho; // output height + int32_t wi; // input width + int32_t hi; // input height + int32_t di; // input depth + int32_t kw; // kernel width + int32_t kh; // kernel height + int32_t kd; // kernel depth + int32_t real_g; // actual groups + int32_t stride_w; + int32_t stride_h; + int32_t stride_d; + int32_t pad_l; // left padding + int32_t pad_r; // right padding + int32_t pad_u; // up padding + int32_t pad_d; // down padding + int32_t pad_f; // front padding + int32_t pad_b; // back padding + int32_t dilation_w; + int32_t dilation_h; + int32_t dilation_d; + int32_t ci1; // another input channels parameter + uint64_t bl1_bound; // buffer limit 1 bound, tiling结果的衍生参数，不建议放在这里 + int32_t batch_dout_single_core; // batch*dout per core, tiling结果的衍生参数，不建议放在这里 // codespell:ignore dout + + // Tiling parameters + uint32_t k0; + uint32_t m0; + uint32_t n0; + uint32_t hf32Flag; + + ge::DataType a_dtype = ge::DT_FLOAT16; + ge::DataType b_dtype = ge::DT_FLOAT16; + ge::DataType c_dtype = ge::DT_FLOAT16; + int32_t a_dtype_bytes = 2; + int32_t b_dtype_bytes = 2; + int32_t c_dtype_bytes = 2; + uint32_t core_num; +}; + +struct Conv3dBpInputV2RunInfo { + // Batch and group related + int32_t batch_n; // Batch size + int32_t real_g; // Number of groups + + // Input dimensions (dedx) + int32_t dedx_d; // Input depth + int32_t dedx_cin; // Input channels per group + int32_t dedx_cin1; // Input channels per group + int32_t dedx_cin1_g; // Input channels per group (grouped) + int32_t dedx_h; // Input height + int32_t dedx_w; // Input width + + // Output dimensions (dedy) + int32_t dedy_d; // Output depth + int32_t dedy_cout; // Output channels per group + int32_t dedy_cout1; // Output channels per group + int32_t dedy_cout1_g; // Output channels per group (grouped) + int32_t dedy_h; // Output height + int32_t dedy_w; // Output width + + // Kernel dimensions + int32_t kernel_d; // Kernel depth + int32_t kernel_h; // Kernel height + int32_t kernel_w; // Kernel width + + // Strides + int32_t stride_d; // Stride depth + int32_t stride_h; // Stride height + int32_t stride_w; // Stride width + + // Padding + int32_t pad_h; // Padding height + int32_t pad_t; // Padding top + int32_t pad_u; // Padding up + int32_t pad_d; // Padding down + int32_t pad_l; // Padding left + int32_t pad_r; // Padding right + + // Dilation + int32_t dilation_d; // Dilation depth + int32_t dilation_h; // Dilation height + int32_t dilation_w; // Dilation width + + // Backprop padding + int32_t backprop_pad_h; // Backprop padding height + int32_t backprop_pad_t; // Backprop padding top + int32_t backprop_pad_u; // Backprop padding up + int32_t backprop_pad_d; // Backprop padding down + int32_t backprop_pad_l; // Backprop padding left + int32_t backprop_pad_r; // Backprop padding right + + // Other flags + int32_t hf32_flag; // Flag for FP32 handling + int32_t a_dtype_bytes = 2; + int32_t b_dtype_bytes = 2; + int32_t c_dtype_bytes = 2; + int32_t initOutputFlag = 0; +}; + +struct Conv3DBackpropV2CompileInfo { + std::string soc_version = ""; + platform_ascendc::SocVersion shortSocVersion = platform_ascendc::SocVersion::ASCEND910B; + + uint32_t core_num = 0; + uint64_t ub_size = 0; + uint64_t l1_size = 0; + uint64_t l2_size = 0; + uint64_t l0a_size = 0; + uint64_t l0b_size = 0; + uint64_t l0c_size = 0; + uint64_t bt_size = 0; + int32_t cube_freq = 0; + bool load3d_constraints = true; + bool intrinsic_data_move_l12ub = true; + bool intrinsic_matmul_ub_to_ub = false; + bool intrinsic_conv_ub_to_ub = false; + bool intrinsic_data_move_l0c2ub = true; + bool intrinsic_fix_pipe_l0c2out = false; + bool intrinsic_fix_pipe_l0c2ub = false; + bool intrinsic_data_move_out2l1_nd2nz = false; + bool intrinsic_data_move_l12bt_bf16 = false; +}; + +enum OpTypeV2 : size_t { + kConv3DBackpropFilterV2, + kConv3DBackpropInputV2, + kConv3DTransposeV2, +}; + +bool GetTbeTiling(const gert::TilingContext* context, Conv3dBpFilterV2RunInfo& runInfoForV2, Conv3dBackpropV2TBETilingData& tbeTilingForV2); +bool GetTbeTiling(gert::TilingContext* context, Conv3dBpInputV2RunInfo& runInfoV2, + Conv3dBackpropV2TBETilingData& tbeTilingForV2, const optiling::OpTypeV2 opType); +bool GetTbeTiling(gert::TilingContext* context, Conv3dBackpropV2TBETilingData& tbeTilingForV2, const optiling::OpTypeV2 opType); +} +#endif // TBE_TILING_API_H diff --git a/csrc/ascend/gmm/CMakeLists.txt b/csrc/ascend/gmm/CMakeLists.txt new file mode 100644 index 000000000..22a6bd86b --- /dev/null +++ b/csrc/ascend/gmm/CMakeLists.txt @@ -0,0 +1,21 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB SUBDIRECTORIES LIST_DIRECTORIES true RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +foreach(SUBDIR ${SUBDIRECTORIES}) + if(DEFINED ASCEND950_SKIP_OPS AND ${SUBDIR} IN_LIST ASCEND950_SKIP_OPS) + message(STATUS "Skipping ${SUBDIR} because ASCEND950 is not available.") + continue() + endif() + + if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${SUBDIR}/CMakeLists.txt) + add_subdirectory(${SUBDIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/CMakeLists.txt b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/CMakeLists.txt new file mode 100644 index 000000000..2dbca2086 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/CMakeLists.txt @@ -0,0 +1,18 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB SUBDIRECTORIES LIST_DIRECTORIES true RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +# 遍历子目录 +foreach(SUBDIR ${SUBDIRECTORIES}) + # 检查子目录中是否存在 CMakeLists.txt + if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${SUBDIR}/CMakeLists.txt) + add_subdirectory(${SUBDIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/README.md b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/README.md new file mode 100644 index 000000000..d4ed1109e --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/README.md @@ -0,0 +1,185 @@ +# GroupedMatmulSwigluQuant + +## 产品支持情况 + +| 产品 | 是否支持 | +| :----------------------------------------------------------- | :------: | +| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | √ | +| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | √ | +| Kirin X90 处理器系列产品 | √ | +| Kirin 9030 处理器系列产品 | √ | + +## 功能说明 + +- 接口功能：融合GroupedMatmul 、dquant、swiglu和quant，详细解释见计算公式。 +- 计算公式： + + - **定义**： + + - **⋅** 表示矩阵乘法。 + - **⊙** 表示逐元素乘法。 + - $\left \lfloor x\right \rceil$ 表示将x四舍五入到最近的整数。 + - $\mathbb{Z_8} = \{ x \in \mathbb{Z} | −128≤x≤127 \}$ + - $\mathbb{Z_{32}} = \{ x \in \mathbb{Z} | -2147483648≤x≤2147483647 \}$ + - **输入**： + + - $X∈\mathbb{Z_8}^{M \times K}$：输入矩阵（左矩阵），M是总token 数，K是特征维度。 + - $W∈\mathbb{Z_8}^{E \times K \times N}$：分组权重矩阵（右矩阵），E是专家个数，K是特征维度，N是输出维度。 + - $bias∈\mathbb{Z_{32}}^{E \times N}$：矩阵乘计算的偏移值，E是专家个数，N是输出维度。 + - $offset∈\mathbb{R}^{E \times N}$：per-channel非对称反量化的偏移，E是专家个数，N是输出维度。 + - $w\_scale∈\mathbb{R}^{E \times N}$：分组权重矩阵（右矩阵）的逐通道缩放因子，E是专家个数，N是输出维度。 + - $x\_scale∈\mathbb{R}^{M}$：输入矩阵（左矩阵）的逐 token缩放因子，M是总token 数。 + - $groupList∈\mathbb{N}^{E}$：前缀和的分组索引列表。 + - **输出**： + + - $Q∈\mathbb{Z_8}^{M \times N / 2}$：量化后的输出矩阵。 + - $Q\_scale∈\mathbb{R}^{M}$：量化缩放因子。 + - $Q\_offset∈\mathbb{R}^{M}$：量化偏移因子。 + - **计算过程** + + - 1.根据groupList[i]确定当前分组的 token ，$i \in [0,Len(groupList)]$。 + + >例子：假设groupList=[3,4,4,6]，从0开始计数。 + > + >第0个右矩阵`W[0,:,:]`，对应索引位置[0,3)的token`x[0:3]`（共3-0=3个token），对应`x_scale[0:3]`、`w_scale[0]`、`bias[0]`、`offset[0]`、`Q[0:3]`、`Q_scale[0:3]`、`Q_offset[0:3]`； + > + >第1个右矩阵`W[1,:,:]`，对应索引位置[3,4)的token`x[3:4]`（共4-3=1个token），对应`x_scale[3:4]`、`w_scale[1]`、`bias[1]`、`offset[1]`、`Q[3:4]`、`Q_scale[3:4]`、`Q_offset[3:4]`； + > + >第2个右矩阵`W[2,:,:]`，对应索引位置[4,4)的token`x[4:4]`（共4-4=0个token），对应`x_scale[4:4]`、`w_scale[2]`、`bias[2]`、`offset[2]`、`Q[4:4]`、`Q_scale[4:4]`、`Q_offset[4:4]`； + > + >第3个右矩阵`W[3,:,:]`，对应索引位置[4,6)的token`x[4:6]`（共6-4=2个token），对应`x_scale[4:6]`、`w_scale[3]`、`bias[3]`、`offset[3]`、`Q[4:6]`、`Q_scale[4:6]`、`Q_offset[4:6]`； + > + >请注意：groupList中未指定的部分将不会参与更新。 + >例如groupList=[12,14,18]，X的shape为[30，:]。 + > + >则第一个输出Q的shape为[30，:]，其中Q[18:，：]的部分不会进行更新和初始化，其中数据为显存空间申请时的原数据。 + > + >同理，第二个输出Q的shape为[30]，其中Q\_scale[18:]的部分不会进行更新或初始化，其中数据为显存空间申请时的原数据。 + > + >即输出的Q[:groupList[-1],:]和Q\_scale[:groupList[-1]]为有效数据部分。 + + - 2.根据分组确定的入参进行如下计算： + + $C_{i} = (X_{i}\cdot W_{i} )\odot x\_scale_{i\ BroadCast} \odot w\_scale_{i\ BroadCast}$ + + $C_{i,act}, gate_{i} = split(C_{i})$ + + $S_{i}=Swish(C_{i,act})\odot gate_{i}$   其中$Swish(x)=\frac{x}{1+e^{-x}}$ + + >注：当前版本不支持$bias_{i}$、$offset_{i}$，未来版本将支持的计算公式如下： + >$C_{i} =(X_{i}\cdot W_{i} + bias_{i\ BroadCast})\odot x\_scale_{i\ BroadCast} \odot w\_scale_{i\ BroadCast}+offset_{i\ BroadCast}$ + + - 3.确定量化方式 + + - 当量化方式为对称量化时： + + $Q\_scale_{i} = \frac{max(|S_{i}|)}{127}$ + + $Q_{i} = \left \lfloor \frac{S_{i}}{Q\_scale_{i}}\right \rceil $ + + - 当量化方式为非对称量化时：(暂不支持) + + $Q\_scale_{i} = \frac{max(S_{i})-min(S_{i})}{255}$ + + $Q\_offset_{i} = -128 - \left \lfloor \frac{min(S_{i})}{Q\_scale_{i}}\right \rceil$ + + $Q_{i} = \left \lfloor \frac{S_{i}}{ Q\_scale_{i} } + Q\_offset_{i}\right \rceil $ + +## 参数说明 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
参数名输入/输出/属性描述数据类型数据格式
x输入左矩阵，公式中的X。INT8ND
weight输入权重矩阵，公式中的W。INT8ND / NZ
bias输入矩阵乘计算的偏移值，公式中的bias。INT32ND
offset输入per-channel非对称反量化的偏移，公式中的offset。FLOAT32ND
weightScale输入右矩阵的量化因子，公式中的w_scale。FLOAT、FLOAT16、BFLOAT16ND
xScale输入左矩阵的量化因子，公式中的x_scale。FLOAT32ND
groupList输入指示每个分组参与计算的Token个数，公式中的groupList。INT64ND
output输出输出的量化因子，公式中的Q。FLOATND
outputScale输出输出的量化因子，公式中的Q_scale。FLOATND
outputOffset输出输出的非对称量化的偏移，公式中的Q_offset。FLOATND
+ +- Kirin X90/Kirin 9030 处理器系列产品: 不支持BFLOAT16。 + +## 约束说明 + +- N轴长度不能超过10240。 +- K轴长度不能超过65536。 + +## 调用说明 + +| 调用方式 | 调用样例 | 说明 | +|--------------|-------------------------|--------------------------------------------------------------| +| aclnn调用 | [test_aclnn_grouped_matmul_swiglu_quant](examples/test_aclnn_grouped_matmul_swiglu_quant.cpp) | 通过接口方式调用[GroupedMatmulSwigluQuant](docs/aclnnGroupedMatmulSwigluQuant.md)算子。 | diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/CMakeLists.txt b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/CMakeLists.txt new file mode 100644 index 000000000..a72dfb574 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/CMakeLists.txt @@ -0,0 +1,28 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +add_op_to_compiled_list() +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnExc PRIVATE + grouped_matmul_swiglu_quant_def.cpp + ) + add_ops_compile_options( + OP_NAME GroupedMatmulSwigluQuant + OPTIONS --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror + ) +endif() + +if(NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE grouped_matmul_swiglu_quant ACLNNTYPE aclnn_exclude) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/grouped_matmul_swiglu_quant_def.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/grouped_matmul_swiglu_quant_def.cpp new file mode 100644 index 000000000..d1585e071 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/grouped_matmul_swiglu_quant_def.cpp @@ -0,0 +1,122 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_def.cpp + * \brief + */ + +#include "register/op_def_registry.h" +namespace ops { +class GroupedMatmulSwigluQuant : public OpDef { +public: + explicit GroupedMatmulSwigluQuant(const char *name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("weight") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT4, ge::DT_INT4}) + .Format({ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_ND, + ge::FORMAT_FRACTAL_NZ}); + this->Input("weight_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_UINT64, ge::DT_UINT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("x_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("weight_assistance_matrix") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("group_list") + .ParamType(REQUIRED) + .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("y_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Attr("is_enable_weight_assistance_matrix").AttrType(OPTIONAL).Bool(true); + this->Attr("dequant_mode").AttrType(OPTIONAL).Int(0); + this->Attr("limited").AttrType(OPTIONAL).Float(0.0f); + + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true); + + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + + OpAICoreConfig config_kirin = GetKirinCoreConfig(); + this->AICore().AddConfig("kirinx90", config_kirin); + this->AICore().AddConfig("kirin9030", config_kirin); + } + +private: + OpAICoreConfig GetKirinCoreConfig() const + { + OpAICoreConfig config_kirin; + config_kirin.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true); + config_kirin.Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_INT8}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Input("weight") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_INT8}) + .Format({ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ}); + config_kirin.Input("weight_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Input("x_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Input("weight_assistance_matrix") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Input("group_list") + .ParamType(REQUIRED) + .DataType({ge::DT_INT64, ge::DT_INT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_INT8}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Output("y_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}); + return config_kirin; + } +}; + +OP_ADD(GroupedMatmulSwigluQuant); +} // namespace ops diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/grouped_matmul_swiglu_quant_infershape.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/grouped_matmul_swiglu_quant_infershape.cpp new file mode 100644 index 000000000..bd74ff2da --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/grouped_matmul_swiglu_quant_infershape.cpp @@ -0,0 +1,53 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_proto.cpp + * \brief + */ +#include "register/op_impl_registry.h" +#include "log/log.h" +#include "platform/platform_info.h" + +using namespace ge; +namespace ops { +const int64_t X_INDEX = 0; +const int64_t WEIGHTSCALE_INDEX = 2; +const int64_t M_DIM_INDEX = 0; +const int64_t DIM_LEN = 2; +const int64_t SPLIT_RATIO = 2; +static ge::graphStatus InferShape4GroupedMatmulSwigluQuant(gert::InferShapeContext *context) +{ + const gert::Shape *xShape = context->GetInputShape(X_INDEX); + const gert::Shape *weightScaleShape = context->GetInputShape(WEIGHTSCALE_INDEX); + int64_t m = xShape->GetDim(M_DIM_INDEX); + int64_t N_DIM_INDEX = weightScaleShape->GetDimNum() - 1; + int64_t n = static_cast(weightScaleShape->GetDim(N_DIM_INDEX) / SPLIT_RATIO); + auto outShape = context->GetOutputShape(0); + outShape->SetDimNum(DIM_LEN); + outShape->SetDim(0, m); + outShape->SetDim(1, n); + auto outScaleShape = context->GetOutputShape(1); + outScaleShape->SetDimNum(1); + outScaleShape->SetDim(0, m); + return GRAPH_SUCCESS; +} + +static graphStatus InferDataType4GroupedMatmulSwigluQuant(gert::InferDataTypeContext *context) +{ + context->SetOutputDataType(0, DataType::DT_INT8); + context->SetOutputDataType(1, DataType::DT_FLOAT); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(GroupedMatmulSwigluQuant) + .InferShape(InferShape4GroupedMatmulSwigluQuant) + .InferDataType(InferDataType4GroupedMatmulSwigluQuant); +} // namespace ops diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/grouped_matmul_swiglu_quant_tiling.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/grouped_matmul_swiglu_quant_tiling.cpp new file mode 100644 index 000000000..9c1243f0e --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/grouped_matmul_swiglu_quant_tiling.cpp @@ -0,0 +1,288 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_tiling.cpp + * \brief + */ +#include +#include +#include "register/op_impl_registry.h" +#include "log/log.h" +#include "err/ops_err.h" +#include "tiling_base/tiling_base.h" +#include "grouped_matmul_swiglu_quant_tiling.h" +using namespace ge; +using namespace AscendC; +using namespace GroupedMatmulSwigluQuantTiling; +using namespace Ops::Transformer::OpTiling; +namespace { +template +static inline auto AlignUp(T a, T base) -> T +{ + if (base == 0) { + return 0; + } + return (a + base - 1) / base * base; +} +} // namespace + +namespace optiling { + +struct GMMSwigluCompileInfo { + uint64_t ubSize_ = 0; + uint32_t aicNum_ = 0; + uint32_t baseM_ = 128; + uint32_t baseN_ = 256; +}; + +static int64_t CalMaxRowInUb_A8W4(const gert::TilingContext *context, const uint64_t ubSize, const uint64_t n) +{ + const uint64_t ALIGNMENT = 8; + const float WEIGHT_FACTOR = 8.5; + const uint64_t ALIGNMENT_TERM_FACTOR = 4; + const uint64_t LINEAR_TERM_FACTOR = 6; + const uint64_t CONSTANT_TERM = 64; + const uint64_t MIN_ROW_THRESHOLD = 1; + + // 表达式：8.5 * row * n + 4 * alignUp(row, 8) + 6n + 64 <= ubSize + + // 忽略对齐项的初始估计 + int64_t maxRowEstimate = + (ubSize - CONSTANT_TERM - LINEAR_TERM_FACTOR * n) / static_cast(WEIGHT_FACTOR * n); + + // 考虑对齐影响 + uint64_t alignedRow = (maxRowEstimate + ALIGNMENT - 1) / ALIGNMENT * ALIGNMENT; + uint64_t totalSize = static_cast(WEIGHT_FACTOR * maxRowEstimate * n) + + ALIGNMENT_TERM_FACTOR * alignedRow + LINEAR_TERM_FACTOR * n + CONSTANT_TERM; + + // 如果超过UB大小，逐步减少row直到满足条件 + while (totalSize > ubSize && maxRowEstimate > 0) { + maxRowEstimate--; + alignedRow = (maxRowEstimate + ALIGNMENT - 1) / ALIGNMENT * ALIGNMENT; + totalSize = static_cast(WEIGHT_FACTOR * maxRowEstimate * n) + ALIGNMENT_TERM_FACTOR * alignedRow + + LINEAR_TERM_FACTOR * n + CONSTANT_TERM; + } + + if (maxRowEstimate < MIN_ROW_THRESHOLD) { + OP_LOGE(context->GetNodeName(), "GMM_SWIGLU_QUANT TILING: No valid row found for n = %lu, ubSize = %lu\n", n, + ubSize); + return 0; + } + return maxRowEstimate; +} + +static int64_t CalMaxRowInUb(const gert::TilingContext *context, const uint64_t ubSize, const uint64_t n) +{ + uint64_t tmpBufSize = (n / SWIGLU_REDUCE_FACTOR) * FP32_DTYPE_SIZE; + uint64_t perchannleBufSize = n * FP32_DTYPE_SIZE * DOUBLE_BUFFER; + uint64_t reduceMaxResBufSize = BLOCK_BYTE; + uint64_t reduceMaxTmpBufSize = BLOCK_BYTE; + const uint64_t CONSTANT_TERM = 64; + int64_t remainUbSize = ubSize - tmpBufSize - perchannleBufSize - reduceMaxResBufSize - reduceMaxTmpBufSize; + int64_t maxRowInUb = + remainUbSize / (n * INT32_DTYPE_SIZE + n / SWIGLU_REDUCE_FACTOR + FP32_DTYPE_SIZE) / DOUBLE_BUFFER; + int64_t curUb = DOUBLE_BUFFER * (maxRowInUb * (INT32_DTYPE_SIZE * n + n / SWIGLU_REDUCE_FACTOR) + + AlignUp(maxRowInUb, FP32_BLOCK_SIZE) * FP32_DTYPE_SIZE); + if (curUb > remainUbSize) { + // 64 : make sure ub does not excceed maxUbSize after align up to 8 + maxRowInUb = (remainUbSize - CONSTANT_TERM) / + (n * INT32_DTYPE_SIZE + n / SWIGLU_REDUCE_FACTOR + FP32_DTYPE_SIZE) / DOUBLE_BUFFER; + } + if (maxRowInUb < 1) { + // when n > (ubSize - 72) / 19 = 10330, maxRowInUb < 1 + OP_LOGE(context->GetNodeName(), "GMM_SWIGLU_QUANT TILING: n should not be greater than 10240, now is %lu\n", n); + } + return maxRowInUb; +} + +static void SetTilingKey(gert::TilingContext *context, bool isSplitWorkSpace, bool isA8W4MSD) +{ + if (isA8W4MSD) { // A8W4 MSD tiling_key使用4 + context->SetTilingKey(A8W4_MSD_TILING_KEY_MODE); + context->SetScheduleMode(BATCH_MODE_SCHEDULE); + } else if (isSplitWorkSpace) { + context->SetTilingKey(SPLITWORKSPACE_TILING_KEY_MODE); + context->SetScheduleMode(BATCH_MODE_SCHEDULE); + } else { + context->SetTilingKey(COMMON_TILING_KEY_MODE); + context->SetScheduleMode(BATCH_MODE_SCHEDULE); + } +} + +ASCENDC_EXTERN_C graphStatus TilingGMMSwigluQuant(gert::TilingContext *context) +{ + // set info + OP_LOGD(context->GetNodeName(), "Begin Run GMM Swiglu Tiling ."); + auto xDesc = context->GetInputDesc(X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, xDesc); + auto weightDesc = context->GetInputDesc(WEIGHT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, weightDesc); + ge::DataType xDType = xDesc->GetDataType(); + ge::DataType weightDType = weightDesc->GetDataType(); + + bool isA8W4MSD = (xDType == ge::DataType::DT_INT8 && weightDType == ge::DataType::DT_INT4); + auto compileInfoPtr = context->GetCompileInfo(); + auto xTensor = context->GetInputTensor(X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, xTensor); + const int64_t m = xTensor->GetStorageShape().GetDim(0); + const int64_t k = xTensor->GetStorageShape().GetDim(1); + auto wTensor = context->GetInputTensor(WEIGHT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, wTensor); + // swiglu limit 0 means clamp is disabled. + auto attrs = context->GetAttrs(); + float limited = 0.0f; + if (attrs != nullptr) { + if (const double *limitedPtr = attrs->GetAttrPointer(ATTR_INDEX_LIMITED)) { + limited = static_cast(*limitedPtr); + } + } + OP_CHECK_IF(!(limited >= 0.0f), + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "limited should be non-negative"), + return GRAPH_FAILED); + int64_t n = 0; + if (wTensor->GetStorageShape().GetDimNum() == ND_WEIGHT_DIM_LIMIT) { // ND + n = wTensor->GetStorageShape().GetDim(DIM_2); + } else if (wTensor->GetStorageShape().GetDimNum() == NZ_WEIGHT_DIM_LIMIT) { // NZ + n = wTensor->GetStorageShape().GetDim(DIM_1) * wTensor->GetStorageShape().GetDim(DIM_4); + } + auto wScaleTensor = context->GetInputTensor(WEIGHT_SCALE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, wScaleTensor); + int64_t quantGroupNum = 0; + if (wScaleTensor->GetStorageShape().GetDimNum() == PERCHANNEL_WSCALE_DIM_LIMIT) { // perChannel + quantGroupNum = 1; + } else if (wScaleTensor->GetStorageShape().GetDimNum() == PERGROUP_WSCALE_DIM_LIMIT) { // perGroup + quantGroupNum = wScaleTensor->GetStorageShape().GetDim(1); + } + auto groupListTensor = context->GetDynamicInputTensor(GROUPLIST_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, groupListTensor); + const int64_t groupNum = groupListTensor->GetStorageShape().GetDim(0); + GMMSwigluQuantTilingData tilingData; + int64_t row = 0; + if (isA8W4MSD) { + row = CalMaxRowInUb_A8W4(context, compileInfoPtr->ubSize_, n); + } else { + row = CalMaxRowInUb(context, compileInfoPtr->ubSize_, n); + } + + tilingData.gmmSwigluBaseParams.set_groupNum(groupNum); + tilingData.gmmSwigluBaseParams.set_coreNum(compileInfoPtr->aicNum_); + tilingData.gmmSwigluBaseParams.set_K(k); + tilingData.gmmSwigluBaseParams.set_N(n); + tilingData.gmmSwigluBaseParams.set_M(m); + tilingData.gmmSwigluBaseParams.set_baseM(A8W4_BASEM); + tilingData.gmmSwigluBaseParams.set_baseN(A8W4_BASEN); + tilingData.gmmSwigluBaseParams.set_limited(limited); + tilingData.gmmSwiglu.set_maxProcessRowNum(row); + tilingData.gmmSwiglu.set_groupListLen(groupNum); + tilingData.gmmSwiglu.set_tokenLen(n); + + tilingData.gmmSwigluBaseParams.set_quantGroupNum(quantGroupNum); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + using namespace matmul_tiling; + + MatmulApiTiling tiling(ascendcPlatform); + tiling.SetAType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_INT4); + tiling.SetBType(TPosition::GM, CubeFormat::NZ, matmul_tiling::DataType::DT_INT4); + tiling.SetCType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16); + tiling.SetBias(false); + tiling.SetShape(A8W4_BASEM, A8W4_BASEN, k); + tiling.SetFixSplit(A8W4_BASEM, A8W4_BASEN, A8W4_BASEK); + tiling.SetOrgShape(m, n, k); + tiling.SetBufferSpace(-1, -1, -1); + OP_CHECK_IF( + tiling.GetTiling(tilingData.mmTilingData) == -1, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "grouped_matmul_swiglu_quant_tiling, get tiling failed"), + return GRAPH_FAILED); + if (isA8W4MSD) { + tilingData.mmTilingData.set_baseM(A8W4_BASEM); + tilingData.mmTilingData.set_baseN(A8W4_BASEN); + tilingData.mmTilingData.set_baseK(A8W4_BASEK); + tilingData.mmTilingData.set_dbL0B(DOUBLE_BUFFER); + tilingData.mmTilingData.set_stepKa(NUM_FOUR); + tilingData.mmTilingData.set_stepKb(NUM_FOUR); + tilingData.mmTilingData.set_depthA1(NUM_EIGHT); + tilingData.mmTilingData.set_depthB1(NUM_EIGHT); + tilingData.mmTilingData.set_stepM(1); + tilingData.mmTilingData.set_stepN(1); + } + auto workspaceSizes = context->GetWorkspaceSizes(1); + int64_t usrWorkspaceLimit = USER_WORKSPACE_LIMIT; + int64_t mLimit = 0; + if (isA8W4MSD) { + mLimit = ((usrWorkspaceLimit / DOUBLE_WORKSPACE_SPLIT) / (k * sizeof(int8_t) + DOUBLE_ROW * n * sizeof(half))); + } else { + mLimit = ((usrWorkspaceLimit / DOUBLE_WORKSPACE_SPLIT) / INT32_DTYPE_SIZE) / n; + } + OP_CHECK_IF(mLimit <= 0, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "mLimit is %ld must over then 0.", mLimit), + return GRAPH_FAILED); + tilingData.gmmSwigluBaseParams.set_mLimit(mLimit); + if (isA8W4MSD) { + int workSpaceMTemp = mLimit * DOUBLE_WORKSPACE_SPLIT; + tilingData.gmmSwigluBaseParams.set_workSpaceOffset1(workSpaceMTemp * k * sizeof(int8_t)); + tilingData.gmmSwigluBaseParams.set_workSpaceOffset2(2 * workSpaceMTemp * n * sizeof(half)); + workspaceSizes[0] = + SYS_WORKSPACE_SIZE + // 系统预留16MB + (workSpaceMTemp * k * sizeof(int8_t)) + // 第一阶段 预处理左矩阵 (mLimit, K) * int8 * 2(double WorkSpace) + (DOUBLE_ROW * workSpaceMTemp * n * + sizeof(half)); // 第二阶段 矩阵乘结果 (2 * mLimit, N) * fp16 * 2(double WorkSpace) + } else { + int workSpaceMTemp = (mLimit * DOUBLE_WORKSPACE_SPLIT > m ? m : mLimit * DOUBLE_WORKSPACE_SPLIT); + tilingData.gmmSwigluBaseParams.set_workSpaceOffset1(0); + tilingData.gmmSwigluBaseParams.set_workSpaceOffset2(0); + workspaceSizes[0] = SYS_WORKSPACE_SIZE + (workSpaceMTemp * n * sizeof(int32_t)); + } + bool isSplitWorkSpace = m > mLimit * DOUBLE_WORKSPACE_SPLIT; + OP_LOGD(context->GetNodeName(), "grouped_matmul_swiglu_quant_tiling."); + OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.groupNum: %ld", groupNum); + OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.coreNum: %u ", compileInfoPtr->aicNum_); + OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.M: %ld", m); + OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.K: %ld", k); + OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.N: %ld", n); + OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.baseM: %ld", A8W4_BASEM); + OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.baseN: %ld", A8W4_BASEN); + OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.mLimit: %ld", mLimit); + OP_LOGD(context->GetNodeName(), "gmmSwigluBaseParams.quantGroupNum: %ld", quantGroupNum); + OP_LOGD(context->GetNodeName(), "gmmSwiglu.maxProcessRowNum: %ld", row); + OP_LOGD(context->GetNodeName(), "gmmSwiglu.groupListLen: %ld", groupNum); + OP_LOGD(context->GetNodeName(), "gmmSwiglu.tokenLen: %ld", n); + OP_LOGD(context->GetNodeName(), "USER_WORKSPACE_LIMIT: %ld", usrWorkspaceLimit); + OP_LOGD(context->GetNodeName(), "workspaceSizes: %lu", workspaceSizes[0]); + OP_LOGD(context->GetNodeName(), "isSplitWorkSpace: %s", isSplitWorkSpace ? "true" : "false"); + OP_LOGD(context->GetNodeName(), "GMMSWIGLUQUANT_TILING: baseM is %u, baseK is %u, baseN is %u.", A8W4_BASEM, A8W4_BASEK, A8W4_BASEN); + SetTilingKey(context, isSplitWorkSpace, isA8W4MSD); + tilingData.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); + context->SetBlockDim(compileInfoPtr->aicNum_); // block dim is the number of aicube + context->GetRawTilingData()->SetDataSize(tilingData.GetDataSize()); + OP_LOGD(context->GetNodeName(), "End Run GMM Swiglu Tiling."); + return GRAPH_SUCCESS; +} + +ASCENDC_EXTERN_C graphStatus TilingPrepareForGMMSwigluQuant(gert::TilingParseContext *context) +{ + // get info + fe::PlatFormInfos *platformInfoPtr = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfoPtr); + auto compileInfoPtr = context->GetCompiledInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + compileInfoPtr->aicNum_ = ascendcPlatform.GetCoreNumAic(); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, compileInfoPtr->ubSize_); + OP_LOGD(context->GetNodeName(), "ubSize is %lu, aicNum is %u.", compileInfoPtr->ubSize_, compileInfoPtr->aicNum_); + return GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(GroupedMatmulSwigluQuant) + .Tiling(TilingGMMSwigluQuant) + .TilingParse(TilingPrepareForGMMSwigluQuant); +} // namespace optiling diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/grouped_matmul_swiglu_quant_tiling.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/grouped_matmul_swiglu_quant_tiling.h new file mode 100644 index 000000000..621b4e6a2 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/grouped_matmul_swiglu_quant_tiling.h @@ -0,0 +1,95 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_tiling.h + * \brief + */ +#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_GROUPED_MATMUL_SWIGLU_QUANT_H +#define AIR_CXX_RUNTIME_V2_OP_IMPL_GROUPED_MATMUL_SWIGLU_QUANT_H + +#include +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" + +namespace optiling { +// GMM 基本信息 +BEGIN_TILING_DATA_DEF(GMMSwigluBaseParams) +TILING_DATA_FIELD_DEF(uint32_t, groupNum); +TILING_DATA_FIELD_DEF(uint32_t, coreNum); +TILING_DATA_FIELD_DEF(uint32_t, K); +TILING_DATA_FIELD_DEF(uint32_t, N); +TILING_DATA_FIELD_DEF(uint32_t, M); +TILING_DATA_FIELD_DEF(uint32_t, baseM); +TILING_DATA_FIELD_DEF(uint32_t, baseN); +TILING_DATA_FIELD_DEF(uint32_t, mLimit); +TILING_DATA_FIELD_DEF(uint32_t, workSpaceOffset1); +TILING_DATA_FIELD_DEF(uint32_t, workSpaceOffset2); +TILING_DATA_FIELD_DEF(uint32_t, quantGroupNum); +TILING_DATA_FIELD_DEF(float, limited); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(GMMSwigluBaseParamsOp, GMMSwigluBaseParams) + +// SwigluQuant部分tiling 基本信息 +BEGIN_TILING_DATA_DEF(GMMSwiglu) +TILING_DATA_FIELD_DEF(uint32_t, maxProcessRowNum); +TILING_DATA_FIELD_DEF(uint32_t, groupListLen); +TILING_DATA_FIELD_DEF(uint32_t, tokenLen); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(GMMSwigluOp, GMMSwiglu) + +// 结构体集合 +BEGIN_TILING_DATA_DEF(GMMSwigluQuantTilingData) +TILING_DATA_FIELD_DEF_STRUCT(GMMSwigluBaseParams, gmmSwigluBaseParams); +TILING_DATA_FIELD_DEF_STRUCT(GMMSwiglu, gmmSwiglu); +TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, mmTilingData); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(GroupedMatmulSwigluQuant, GMMSwigluQuantTilingData) +} // namespace optiling + +namespace GroupedMatmulSwigluQuantTiling { +constexpr uint32_t X_INDEX = 0; +constexpr uint32_t WEIGHT_INDEX = 1; +constexpr uint32_t WEIGHT_SCALE_INDEX = 2; +constexpr uint32_t GROUPLIST_INDEX = 5; +constexpr uint32_t BATCH_MODE_SCHEDULE = 1; +constexpr uint32_t DIM_0 = 0; +constexpr uint32_t DIM_1 = 1; +constexpr uint32_t DIM_2 = 2; +constexpr uint32_t DIM_3 = 3; +constexpr uint32_t DIM_4 = 4; +constexpr uint32_t NUM_FOUR = 4; +constexpr uint32_t NUM_EIGHT = 8; +constexpr uint32_t SYS_WORKSPACE_SIZE = 16 * 1024 * 1024; +constexpr int64_t USER_WORKSPACE_LIMIT = 64 * 1024 * 1024; +constexpr int64_t DOUBLE_WORKSPACE_SPLIT = 2; +constexpr uint32_t ATTR_INDEX_LIMITED = 2; +constexpr int64_t INT32_DTYPE_SIZE = 4; +constexpr int64_t FP32_DTYPE_SIZE = 4; +constexpr int64_t FP32_BLOCK_SIZE = 8; +constexpr int64_t BLOCK_BYTE = 32; +constexpr int64_t SWIGLU_REDUCE_FACTOR = 2; +constexpr int64_t DOUBLE_BUFFER = 2; +constexpr int64_t ND_WEIGHT_DIM_LIMIT = 3; +constexpr int64_t NZ_WEIGHT_DIM_LIMIT = 5; +constexpr int64_t DOUBLE_ROW = 2; +constexpr int64_t PERCHANNEL_WSCALE_DIM_LIMIT = 2; +constexpr int64_t PERGROUP_WSCALE_DIM_LIMIT = 3; +constexpr int64_t A8W4_MSD_TILING_KEY_MODE = 2; +constexpr int64_t SPLITWORKSPACE_TILING_KEY_MODE = 1; +constexpr int64_t COMMON_TILING_KEY_MODE = 0; +constexpr int64_t A8W4_TOKEN_THRESHOLD = 32; +constexpr int64_t A8W4_BASEM = 128; +constexpr int64_t A8W4_BASEK = 256; +constexpr int64_t A8W4_BASEN = 256; +} // namespace GroupedMatmulSwigluQuantTiling + +#endif // AIR_CXX_RUNTIME_V2_OP_IMPL_GROUPED_MATMUL_SWIGLU_QUANT_H \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/aclnn_grouped_matmul_swiglu_quant.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/aclnn_grouped_matmul_swiglu_quant.cpp new file mode 100644 index 000000000..37f4197e9 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/aclnn_grouped_matmul_swiglu_quant.cpp @@ -0,0 +1,540 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#include +#include +#include "aclnn_kernels/contiguous.h" +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "opdev/common_types.h" +#include "opdev/data_type_utils.h" +#include "opdev/format_utils.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/platform.h" +#include "opdev/shape_utils.h" +#include "opdev/tensor_view_utils.h" +#include "opdev/make_op_executor.h" +#include "grouped_matmul_swiglu_quant.h" +#include "aclnn_grouped_matmul_swiglu_quant_weight_nz.h" +#include "aclnn_grouped_matmul_swiglu_quant.h" + +using namespace op; + +#ifdef __cplusplus +extern "C" { +#endif + +static constexpr int64_t SPLIT = 2L; +static constexpr int64_t K_LIMIT_A8W8 = 65536L; +static constexpr int64_t K_LIMIT_A8W4 = 20000L; +static constexpr int64_t N_LIMIT = 10240L; +static constexpr int64_t NZ_DIM_4_INT8 = 32L; +static constexpr int64_t NZ_DIM_4_INT4 = 64L; +static constexpr int64_t NZ_DIM_3 = 16L; +static constexpr int64_t OUTPUT_IDX_0 = 0L; +static constexpr int64_t OUTPUT_IDX_1 = 1L; +static constexpr int64_t DIM_IDX_0 = 0L; +static constexpr int64_t DIM_IDX_1 = 1L; +static constexpr int64_t DIM_IDX_2 = 2L; +static constexpr int64_t DIM_IDX_3 = 4L; +static constexpr size_t X_DIM_LIMIT = 2UL; +static constexpr size_t WEIGHT_ND_DIM_LIMIT = 3UL; +static constexpr size_t WEIGHT_NZ_DIM_LIMIT = 5UL; +static constexpr size_t WEIGHT_SCALE_DIM_LIMIT = 2UL; +static constexpr size_t WEIGHT_SCALE_PERGROUP_DIM_LIMIT = 3UL; +static constexpr size_t WEIGHT_SCALE_PERCHANNEL_DIM_LIMIT = 2UL; +static constexpr size_t TOKEN_SCALE_DIM_LIMIT = 1UL; +static constexpr size_t BIAS_DIM_LIMIT = 2UL; +static constexpr size_t GROUP_LIST_DIM_LIMIT = 1UL; +static constexpr size_t QUANTOUT_DIM_LIMIT = 2UL; +static constexpr size_t QUANTSCALEOUT_DIM_LIMIT = 1UL; +static constexpr size_t INT4_PER_INT32 = 8UL; +bool isEnableWeightAssistanceMatrix = false; +int dequantMode = 0; + +static const std::initializer_list X_DTYPE_SUPPORT_LIST = {DataType::DT_INT8}; +static const std::initializer_list WEIGHT_DTYPE_SUPPORT_LIST = {DataType::DT_INT8, DataType::DT_INT4}; +static const std::initializer_list WEIGHT_SCALE_DTYPE_SUPPORT_LIST = { + DataType::DT_FLOAT, DataType::DT_FLOAT16, DataType::DT_BF16}; +static const std::initializer_list WEIGHT_SCALE_A8W4_DTYPE_SUPPORT_LIST = {DataType::DT_UINT64}; +static const std::initializer_list X_SCALE_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT, DataType::DT_FLOAT16, + DataType::DT_BF16}; +static const std::initializer_list GROUP_LIST_DTYPE_SUPPORT_LIST = {DataType::DT_INT64}; +static const std::initializer_list QUANTOUT_DTYPE_SUPPORT_LIST = {DataType::DT_INT8}; +static const std::initializer_list QUANTSCALEOUT_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT}; +static const std::initializer_list BIAS_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT}; +static bool CheckNotNull(const aclTensor *x, const aclTensor *weight, const aclTensor *bias, const aclTensor *offset, + const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList, + const aclTensor *output, const aclTensor *outputScale, const aclTensor *outputOffset) +{ + OP_CHECK_NULL(x, return false); + OP_CHECK_NULL(weight, return false); + OP_CHECK_NULL(weightScale, return false); + OP_CHECK_NULL(xScale, return false); + OP_CHECK_NULL(groupList, return false); + OP_CHECK_NULL(output, return false); + OP_CHECK_NULL(outputScale, return false); + if (x->GetDataType() == DataType::DT_INT8 && weight->GetDataType() == DataType::DT_INT8 && bias != nullptr) { + OP_LOGW("aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario that bias is not 0. " + "Features and accuracy are not guaranteed if inputting bias with values other than 0."); + } else if (x->GetDataType() == DataType::DT_INT8 && weight->GetDataType() == DataType::DT_INT4 && bias == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario that without bias. " + "When x is Int8 and weight is int4, bias serves as an auxiliary matrix to weight, and this parameter " + "cannot be nullptr."); + return false; + } + if (offset != nullptr) { + OP_LOGW( + "aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario where offset is not 0. " + "Features and accuracy are not guaranteed if inputting bias with values other than 0s."); + } + if (outputOffset != nullptr) { + OP_LOGW("aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario where outputOffset " + "is not 0. Features and accuracy are not guaranteed if inputting bias with values other than 0s."); + } + return true; +} + +static bool CheckInputOutDims_A8W8(const aclTensor *x, const aclTensor *weight, const aclTensor *weightScale, + const aclTensor *xScale, const aclTensor *groupList, const aclTensor *output, + const aclTensor *outputScale) +{ + OP_CHECK_WRONG_DIMENSION(x, X_DIM_LIMIT, return false); + op::Format weightViewFormat = weight->GetViewFormat(); + if (IsPrivateFormat(weightViewFormat)) { + OP_CHECK_WRONG_DIMENSION(weight, WEIGHT_NZ_DIM_LIMIT, return false); + } else { + OP_CHECK_WRONG_DIMENSION(weight, WEIGHT_ND_DIM_LIMIT, return false); + } + OP_CHECK_WRONG_DIMENSION(weightScale, WEIGHT_SCALE_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(xScale, TOKEN_SCALE_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(groupList, GROUP_LIST_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(output, QUANTOUT_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(outputScale, QUANTSCALEOUT_DIM_LIMIT, return false); + return true; +} + +static bool CheckInputOutDims_A8W4(const aclTensor *x, const aclTensor *weight, const aclTensor *bias, + const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList, + const aclTensor *output, const aclTensor *outputScale) +{ + OP_CHECK_WRONG_DIMENSION(x, X_DIM_LIMIT, return false); + op::Format weightViewFormat = weight->GetViewFormat(); + if (IsPrivateFormat(weightViewFormat)) { + OP_CHECK_WRONG_DIMENSION(weight, WEIGHT_NZ_DIM_LIMIT, return false); + } else { + OP_CHECK_WRONG_DIMENSION(weight, WEIGHT_ND_DIM_LIMIT, return false); + } + // 支持pergroup、perchannel量化weightScale分别为2维和3维 + OP_CHECK_MAX_DIM(weightScale, WEIGHT_SCALE_PERGROUP_DIM_LIMIT, return false); + OP_CHECK_MIN_DIM(weightScale, WEIGHT_SCALE_PERCHANNEL_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(bias, BIAS_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(xScale, TOKEN_SCALE_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(groupList, GROUP_LIST_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(output, QUANTOUT_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(outputScale, QUANTSCALEOUT_DIM_LIMIT, return false); + return true; +} + +static bool CheckInputOutShape_A8W8(const aclTensor *x, const aclTensor *weight, const aclTensor *weightScale, + const aclTensor *xScale, const aclTensor *groupList, const aclTensor *output, + const aclTensor *outputScale) +{ + int64_t m = x->GetViewShape().GetDim(0); + int64_t k = x->GetViewShape().GetDim(1); + int64_t n = weightScale->GetViewShape().GetDim(1); + int64_t e = weight->GetViewShape().GetDim(0); + if (n % SPLIT != 0) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "aclnnGroupedMatmulSwiGluQuant, N is %ld , not an even number.", n); + return false; + } + int64_t nAfterHalve = static_cast(n / SPLIT); + // x的shape期望为[M, K] + op::Shape xExpectShape = {m, k}; + // weight的NDshape期望为[E, K, N] + op::Shape weightNDExpectShape = {e, k, n}; + // weight的NZshape期望为[E, N // 32, K // 16, 16, 32] + op::Shape weightNZExpectShape = {e, static_cast(n / NZ_DIM_4_INT8), static_cast(k / NZ_DIM_3), + NZ_DIM_3, NZ_DIM_4_INT8}; + // weightScale的shape期望为[E, N] + op::Shape weightScaleExpectShape = {e, n}; + // xScale的shape期望为[E, N] + op::Shape xScaleExpectShape = {m}; + // output的shape期望为[M, N] + op::Shape outputExpectShape = {m, nAfterHalve}; + // outputScale的shape期望为[M] + op::Shape outputScaleExpectShape = {m}; + op::Format weightViewFormat = weight->GetViewFormat(); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(x, xExpectShape, return false); + if (IsPrivateFormat(weightViewFormat)) { + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weight, weightNZExpectShape, return false); + } else { + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weight, weightNDExpectShape, return false); + } + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weightScale, weightScaleExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(xScale, xScaleExpectShape, return false); + + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(output, outputExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(outputScale, outputScaleExpectShape, return false); + // groupList的长度应小于等于weight的专家数 + int64_t groupListLen = groupList->GetViewShape().GetDim(0); + if (groupListLen > e) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwiGluQuant A8W8, Length of 'groupList' out of range" + " (expected to be in range of [1, %ld], but got %ld)", + e, groupListLen); + return false; + } + if (n > N_LIMIT) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwiGluQuant A8W8: The current version does not support the scenario that " + "N(%ld) is greater than %ld.", + n, N_LIMIT); + return false; + } + if (k >= K_LIMIT_A8W8) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwiGluQuant A8W8, The current version does not support the scenario." + "The tail axis dimension of input0(x) is %ld, which need lower than %ld.", + k, K_LIMIT_A8W8); + return false; + } + return true; +} + +static bool CheckInputOutShape_A8W4(const aclTensor *x, const aclTensor *weight, const aclTensor *bias, + const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList, + const aclTensor *output, const aclTensor *outputScale) +{ + int64_t e = weight->GetViewShape().GetDim(0); + int64_t m = x->GetViewShape().GetDim(0); + int64_t k = x->GetViewShape().GetDim(1); + int64_t n = 1; + int64_t KGroupCount = 1; // K轴的组数，perchannel场景相当于pergroup场景中的组数为1 + int64_t KGroupSize = k; // K轴每组的元素个数 + op::Shape weightScaleExpectShape; + // 通过weightScale的维度判断是否为perchannel 或 pergroup量化模式 + if (weightScale->GetViewShape().GetDimNum() == WEIGHT_SCALE_PERCHANNEL_DIM_LIMIT) { + // weightScale入参在perchannel场景期望shape [E, N] + n = weightScale->GetViewShape().GetDim(DIM_IDX_1); + weightScaleExpectShape = {e, n}; + } else if (weightScale->GetViewShape().GetDimNum() == WEIGHT_SCALE_PERGROUP_DIM_LIMIT) { + // weightScale入参在pergroup场景期望shape [E, KGroupCount, N] + n = weightScale->GetViewShape().GetDim(DIM_IDX_2); + KGroupCount = weightScale->GetViewShape().GetDim(DIM_IDX_1); + KGroupSize = KGroupCount > 0 ? k / KGroupCount : k; + weightScaleExpectShape = {e, KGroupCount, n}; + } + if (KGroupCount == 0 || k % KGroupCount != 0) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwiGluQuant, " + "The number of groups along the k-axis is %ld, and the length of the k-axis is %ld, which is illegal. " + "The number of groups must be greater than 0, and k-axis length %% number of groups == 0 must be true.", + KGroupCount, k); + return false; + } + if (n % SPLIT != 0) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "aclnnGroupedMatmulSwiGluQuant, N is %ld , which must even number.", n); + return false; + } + int64_t nAfterHalve = static_cast(n / SPLIT); + // x的shape期望为[M, K] + op::Shape xExpectShape = {m, k}; + // weight的NDshape期望为[E, K, N] + op::Shape weightNDExpectShape = {e, k, n}; + op::Shape biasExpectShape = {e, n}; + // weight的NZshape期望为[E, N // 64, K // 16, 16, 64] + op::Shape weightNZExpectShape = {e, static_cast(n / NZ_DIM_4_INT4), static_cast(k / NZ_DIM_3), + NZ_DIM_3, NZ_DIM_4_INT4}; + // xScale的shape期望为[E, N] + op::Shape xScaleExpectShape = {m}; + // output的shape期望为[M, N] + op::Shape outputExpectShape = {m, nAfterHalve}; + // outputScale的shape期望为[M] + op::Shape outputScaleExpectShape = {m}; + op::Format weightViewFormat = weight->GetViewFormat(); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(x, xExpectShape, return false); + if (IsPrivateFormat(weightViewFormat)) { + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weight, weightNZExpectShape, return false); + } else { + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weight, weightNDExpectShape, return false); + } + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(bias, biasExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weightScale, weightScaleExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(xScale, xScaleExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(output, outputExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(outputScale, outputScaleExpectShape, return false); + // groupList的长度应小于等于weight的专家数 + int64_t groupListLen = groupList->GetViewShape().GetDim(0); + if (groupListLen > e) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwiGluQuant A8W4, Length of 'groupList' out of range" + " (expected to be in range of [1, %ld], but got %ld)", + e, groupListLen); + return false; + } + if (n > N_LIMIT) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwiGluQuant A8W4, The current version does not support the scenario." + "where N after halve is %ld greater than %ld.", + n, N_LIMIT); + return false; + } + if (k >= K_LIMIT_A8W4) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwiGluQuant A8W4, The current version does not support the scenario." + "The tail axis dimension of input0(x) is %ld, which need lower than %ld.", + k, K_LIMIT_A8W4); + return false; + } + (void)KGroupSize; + return true; +} + +static bool CheckDtypeValid(const aclTensor *x, const aclTensor *weight, const aclTensor *bias, + const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList, + const aclTensor *output, const aclTensor *outputScale) +{ + OP_CHECK_DTYPE_NOT_SUPPORT(x, X_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(weight, WEIGHT_DTYPE_SUPPORT_LIST, return false); + if (weight->GetDataType() == DataType::DT_INT4) { + OP_CHECK_DTYPE_NOT_SUPPORT(bias, BIAS_DTYPE_SUPPORT_LIST, return false); + } + if (weight->GetDataType() == DataType::DT_INT4) { + OP_CHECK_DTYPE_NOT_SUPPORT(weightScale, WEIGHT_SCALE_A8W4_DTYPE_SUPPORT_LIST, return false); + } else if (weight->GetDataType() == DataType::DT_INT8) { + OP_CHECK_DTYPE_NOT_SUPPORT(weightScale, WEIGHT_SCALE_DTYPE_SUPPORT_LIST, return false); + } + OP_CHECK_DTYPE_NOT_SUPPORT(xScale, X_SCALE_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(groupList, GROUP_LIST_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(output, QUANTOUT_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(outputScale, QUANTSCALEOUT_DTYPE_SUPPORT_LIST, return false); + + return true; +} + +static bool CheckFormat(const aclTensor *x, const aclTensor *weight, const aclTensor *output) +{ + bool isNZ = weight->GetStorageFormat() == op::Format::FORMAT_FRACTAL_NZ; + if ((x->GetDataType() == DataType::DT_INT8 && weight->GetDataType() == DataType::DT_INT8) && !isNZ) { + // fp16 in fp32 out that is split k template, not precision-advanced now + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario." + "weight Format expect is FRACTAL_NZ, but got [%s].", + op::ToString(weight->GetStorageFormat()).GetString()); + return false; + } + if (IsPrivateFormat(x->GetStorageFormat())) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario." + "x Format Not support Private Format."); + return false; + } + if (IsPrivateFormat(output->GetStorageFormat())) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwiGluQuant, The current version does not support the scenario." + "output Format Not support Private Format."); + return false; + } + return true; +} + +static void UnpackInt32ToInt4(const aclTensor *&tensorS32, const std::string &tensorType) +{ + OP_LOGD("Unpack %s from int32 to int4 start.", tensorType.c_str()); + auto tensorS4 = const_cast(tensorS32); + op::Shape tensorShape = tensorS4->GetViewShape(); + auto viewShapeDim = tensorShape.GetDimNum(); + tensorShape[viewShapeDim - 1] = tensorShape[viewShapeDim - 1] * INT4_PER_INT32; + tensorS4->SetViewShape(tensorShape); + tensorS4->SetStorageShape(tensorShape); + tensorS4->SetDataType(DataType::DT_INT4); + OP_LOGD("Unpack %s from int32 to int4 finished.", tensorType.c_str()); +} + +static aclnnStatus CheckParams(const aclTensor *x, const aclTensor *weight, const aclTensor *bias, + const aclTensor *offset, const aclTensor *weightScale, const aclTensor *xScale, + const aclTensor *groupList, const aclTensor *output, const aclTensor *outputScale, + const aclTensor *outputOffset) +{ + // 1. 检查参数是否为空指针 + CHECK_RET(CheckNotNull(x, weight, bias, offset, weightScale, xScale, groupList, output, outputScale, outputOffset), + ACLNN_ERR_PARAM_NULLPTR); + // A8W8场景 + if (x->GetDataType() == DataType::DT_INT8 && weight->GetDataType() == DataType::DT_INT8) { + // 2. 校验输入、输出参数维度 + CHECK_RET(CheckInputOutDims_A8W8(x, weight, weightScale, xScale, groupList, output, outputScale), + ACLNN_ERR_PARAM_INVALID); + + // 3. 校验输入、输出shape参数 + CHECK_RET(CheckInputOutShape_A8W8(x, weight, weightScale, xScale, groupList, output, outputScale), + ACLNN_ERR_PARAM_INVALID); + } + // A8W4场景 INT32为兼容torch_npu考虑，实际计算时，1个INT32数据会被视为8个INT4数据 + if ((x->GetDataType() == DataType::DT_INT8 && weight->GetDataType() == DataType::DT_INT4) || + (x->GetDataType() == DataType::DT_INT8 && weight->GetDataType() == DataType::DT_INT32)) { + // 将INT32视为8个Int4数据，调整viewShape和dtype便于后续统一校验 + if (weight->GetDataType() == DataType::DT_INT32) { + UnpackInt32ToInt4(weight, "weight"); + } + if (weightScale->GetDataType() == DataType::DT_INT64) { + auto weightScale_fix = const_cast(weightScale); + weightScale_fix->SetDataType(DataType::DT_UINT64); + } + // 2. 校验输入、输出参数维度 + CHECK_RET(CheckInputOutDims_A8W4(x, weight, bias, weightScale, xScale, groupList, output, outputScale), + ACLNN_ERR_PARAM_INVALID); + + // 3. 校验输入、输出shape参数 + CHECK_RET(CheckInputOutShape_A8W4(x, weight, bias, weightScale, xScale, groupList, output, outputScale), + ACLNN_ERR_PARAM_INVALID); + } + // 4. 检查输入的数据类型是否在支持的数据类型范围之内 + CHECK_RET(CheckDtypeValid(x, weight, bias, weightScale, xScale, groupList, output, outputScale), + ACLNN_ERR_PARAM_INVALID); + + // 5. 检查数据形状是否支持 + CHECK_RET(CheckFormat(x, weight, output), ACLNN_ERR_PARAM_INVALID); + return ACLNN_SUCCESS; +} + +static aclnnStatus aclnnGroupedMatmulSwigluQuantGetWorkspaceSizeCommon( + const aclTensor *x, const aclTensor *weight, const aclTensor *bias, const aclTensor *offset, + const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList, double limited, aclTensor *output, + aclTensor *outputScale, aclTensor *outputOffset, uint64_t *workspaceSize, aclOpExecutor **executor) +{ + // 固定写法，创建OpExecutor + auto uniqueExecutor = CREATE_EXECUTOR(); + CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + // 固定写法，参数检查 + + auto ret = CheckParams(x, weight, bias, offset, weightScale, xScale, groupList, output, outputScale, outputOffset); + + CHECK_RET(ret == ACLNN_SUCCESS, ret); + // 空Tensor场景 + if (output->IsEmpty() || groupList->IsEmpty() || outputScale->IsEmpty()) { + *workspaceSize = 0; + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; + } + // 转连续 + x = l0op::Contiguous(x, uniqueExecutor.get()); + CHECK_RET(x != nullptr, ACLNN_ERR_INNER_NULLPTR); + // 若weight为私有格式，则不应该做连续性转换 （l0op::Contiguous接口会把viewShape赋值给storageShape） + if (IsPrivateFormat(weight->GetStorageFormat())) { + weight->SetOriginalShape(weight->GetViewShape()); + } else { + weight = l0op::Contiguous(weight, uniqueExecutor.get()); + } + CHECK_RET(weight != nullptr, ACLNN_ERR_INNER_NULLPTR); + weightScale = l0op::Contiguous(weightScale, uniqueExecutor.get()); + CHECK_RET(weightScale != nullptr, ACLNN_ERR_INNER_NULLPTR); + xScale = l0op::Contiguous(xScale, uniqueExecutor.get()); + CHECK_RET(xScale != nullptr, ACLNN_ERR_INNER_NULLPTR); + groupList = l0op::Contiguous(groupList, uniqueExecutor.get()); + CHECK_RET(groupList != nullptr, ACLNN_ERR_INNER_NULLPTR); + // 调用L0算子能力 + if (bias != nullptr) { + isEnableWeightAssistanceMatrix = true; + bias = l0op::Contiguous(bias, uniqueExecutor.get()); + CHECK_RET(bias != nullptr, ACLNN_ERR_INNER_NULLPTR); + } + if (isEnableWeightAssistanceMatrix && weightScale->GetViewShape().GetDimNum() == WEIGHT_SCALE_PERGROUP_DIM_LIMIT) { + dequantMode = 1; + } + auto ret_0 = l0op::GroupedMatmulSwigluQuant(x, weight, weightScale, xScale, groupList, limited, bias, + isEnableWeightAssistanceMatrix, dequantMode, uniqueExecutor.get()); + CHECK_RET(ret_0 != std::tuple(nullptr, nullptr), ACLNN_ERR_INNER_NULLPTR); + auto out0 = std::get(ret_0); + auto ret_1 = l0op::ViewCopy(out0, output, uniqueExecutor.get()); + CHECK_RET(ret_1 != nullptr, ACLNN_ERR_INNER_NULLPTR); + auto out1 = std::get(ret_0); + auto ret_2 = l0op::ViewCopy(out1, outputScale, uniqueExecutor.get()); + CHECK_RET(ret_2 != nullptr, ACLNN_ERR_INNER_NULLPTR); + *workspaceSize = uniqueExecutor->GetWorkspaceSize(); + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; +} + +aclnnStatus aclnnGroupedMatmulSwigluQuantGetWorkspaceSize(const aclTensor *x, const aclTensor *weight, + const aclTensor *bias, const aclTensor *offset, + const aclTensor *weightScale, const aclTensor *xScale, + const aclTensor *groupList, double limited, aclTensor *output, + aclTensor *outputScale, aclTensor *outputOffset, + uint64_t *workspaceSize, aclOpExecutor **executor) +{ + OP_CHECK_COMM_INPUT(workspaceSize, executor); + L2_DFX_PHASE_1(aclnnGroupedMatmulSwigluQuant, DFX_IN(x, weight, bias, offset, weightScale, xScale, groupList, limited), + DFX_OUT(output, outputScale, outputOffset)); + // 固定写法，创建OpExecutor + return aclnnGroupedMatmulSwigluQuantGetWorkspaceSizeCommon(x, weight, bias, offset, weightScale, xScale, groupList, limited, + output, outputScale, outputOffset, workspaceSize, + executor); +} + +aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNZGetWorkspaceSize(const aclTensor *x, const aclTensor *weight, + const aclTensor *bias, const aclTensor *offset, + const aclTensor *weightScale, const aclTensor *xScale, + const aclTensor *groupList, double limited, aclTensor *output, + aclTensor *outputScale, aclTensor *outputOffset, + uint64_t *workspaceSize, aclOpExecutor **executor) +{ + OP_CHECK_COMM_INPUT(workspaceSize, executor); + L2_DFX_PHASE_1(aclnnGroupedMatmulSwigluQuantWeightNZ, + DFX_IN(x, weight, bias, offset, weightScale, xScale, groupList), + DFX_OUT(output, outputScale, outputOffset)); + // weight在该场景下强制绑定StorageFormat 和 ViewFormat 为NZ + CHECK_RET(weight != nullptr, ACLNN_ERR_PARAM_NULLPTR); + auto storgeShape = weight->GetStorageShape(); + auto viewShape = weight->GetViewShape(); + aclTensor *weightNZ = const_cast(weight); + CHECK_COND((storgeShape.GetDimNum() == WEIGHT_NZ_DIM_LIMIT), ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwigluQuantWeightNZ, The dimnum of storageShape for second input (weight)" + "must be 5. \n But StorageShape got %s , and dimNum is %lu.", + op::ToString(storgeShape).GetString(), storgeShape.GetDimNum()); + // weight的StorageFormat无条件视为NZ + weightNZ->SetStorageFormat(op::Format::FORMAT_FRACTAL_NZ); + if (viewShape.GetDimNum() == WEIGHT_NZ_DIM_LIMIT) { + // 若weight的viewShape为5维则视为NZ + weightNZ->SetViewFormat(op::Format::FORMAT_FRACTAL_NZ); + } else if (viewShape.GetDimNum() == WEIGHT_ND_DIM_LIMIT) { + // 若weight的viewShape为3维则视为ND + weightNZ->SetViewFormat(op::Format::FORMAT_ND); + } + // 调用公共接口 + return aclnnGroupedMatmulSwigluQuantGetWorkspaceSizeCommon(x, weight, bias, offset, weightScale, xScale, groupList, limited, + output, outputScale, outputOffset, workspaceSize, + executor); +} + +aclnnStatus aclnnGroupedMatmulSwigluQuant(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, + aclrtStream stream) +{ + L2_DFX_PHASE_2(aclnnGroupedMatmulSwigluQuant); + CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER, + "This is an error in GroupedMatmulSwigluQuant launch aicore"); + return ACLNN_SUCCESS; +} + +aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNZ(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, + aclrtStream stream) +{ + L2_DFX_PHASE_2(aclnnGroupedMatmulSwigluQuantWeightNZ); + CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER, + "This is an error in GroupedMatmulSwigluQuantWeightNZ launch aicore"); + return ACLNN_SUCCESS; +} + +#ifdef __cplusplus +} +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/aclnn_grouped_matmul_swiglu_quant.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/aclnn_grouped_matmul_swiglu_quant.h new file mode 100644 index 000000000..6b829a394 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/aclnn_grouped_matmul_swiglu_quant.h @@ -0,0 +1,58 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_API_INC_GROUPED_MATMUL_SWIGLU_QUANT_H +#define OP_API_INC_GROUPED_MATMUL_SWIGLU_QUANT_H +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @brief aclnnGroupedMatmulSwigluQuant的第一段接口，根据具体的计算流程，计算workspace大小。 + * @domain aclnn_ops_infer + * + * @param [in] x: 表示公式中的x，数据类型支持INT8数据类型，数据格式支持ND。 + * @param [in] weight: + * 表示公式中的weight，数据类型支持INT8数据类型，数据格式支持NZ。 + * @param [in] weightScale: + * 表示量化参数，数据类型支持FLOAT16、BFLOAT16、FLOAT32数据类型，数据格式支持ND，支持的最大长度为128个。 表示per + * Channel参数，数据类型支持FLOAT16，BFLOAT16数据类型，数据格式支持ND。 + * @param [in] xScale: + * 表示per Token量化参数，数据类型支持FLOAT32数据类型，数据格式支持ND。 + * @param [in] groupList: 必选参数，代表输入和输出分组轴上的索引情况，数据类型支持INT64。 + * @param [out] quantOutput: 表示公式中的out，数据类型支持INT8数据类型，数据格式支持ND。 + * @param [out] quantScaleOutput: 表示公式中的outQuantScale，数据类型支持Float32数据类型。 + * @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。 + * @param [out] executor: 返回op执行器，包含算子计算流程。 + * @return aclnnStatus: 返回状态码。 + */ +__attribute__((visibility("default"))) aclnnStatus aclnnGroupedMatmulSwigluQuantGetWorkspaceSize( + const aclTensor *x, const aclTensor *weight, const aclTensor *bias, const aclTensor *offset, + const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList, double limited, aclTensor *output, + aclTensor *outputScale, aclTensor *outputOffset, uint64_t *workspaceSize, aclOpExecutor **executor); + +/** + * @brief aclnnGroupedMatmulSwigluQuant的第二段接口，用于执行计算。 + * @param [in] workspace: 在npu device侧申请的workspace内存起址。 + * @param [in] workspaceSize: 在npu + * device侧申请的workspace大小，由第一段接口aclnnGroupedMatmulSwigluQuantGetWorkspaceSize获取。 + * @param [in] stream: acl stream流。 + * @param [in] executor: op执行器，包含了算子计算流程。 + * @return aclnnStatus: 返回状态码。 + */ +__attribute__((visibility("default"))) aclnnStatus +aclnnGroupedMatmulSwigluQuant(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_weight_nz.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_weight_nz.h new file mode 100644 index 000000000..110157385 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_weight_nz.h @@ -0,0 +1,60 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_API_INC_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_H +#define OP_API_INC_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_H +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @brief aclnnGroupedMatmulSwigluQuantWeightNZ的第一段接口，根据具体的计算流程，计算workspace大小。 + * @domain aclnn_ops_infer + * + * @param [in] x: 表示公式中的x，数据类型支持INT8数据类型，数据格式支持ND。 + * @param [in] weight: + * 表示公式中的weight，数据类型支持INT8数据类型，数据格式支持NZ。 + * @param [in] weightScale: + * 表示量化参数，数据类型支持FLOAT16、BFLOAT16、FLOAT32数据类型，数据格式支持ND，支持的最大长度为128个。 表示per + * Channel参数，数据类型支持FLOAT16，BFLOAT16数据类型，数据格式支持ND。 + * @param [in] xScale: + * 表示per Token量化参数，数据类型支持FLOAT32数据类型，数据格式支持ND。 + * @param [in] groupList: 必选参数，代表输入和输出分组轴上的索引情况，数据类型支持INT64。 + * @param [out] quantOutput: 表示公式中的out，数据类型支持INT8数据类型，数据格式支持ND。 + * @param [out] quantScaleOutput: 表示公式中的outQuantScale，数据类型支持Float32数据类型。 + * @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。 + * @param [out] executor: 返回op执行器，包含算子计算流程。 + * @return aclnnStatus: 返回状态码。 + */ +__attribute__((visibility("default"))) aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNZGetWorkspaceSize( + const aclTensor *x, const aclTensor *weight, const aclTensor *bias, const aclTensor *offset, + const aclTensor *weightScale, const aclTensor *xScale, const aclTensor *groupList, double limited, aclTensor *output, + aclTensor *outputScale, aclTensor *outputOffset, uint64_t *workspaceSize, aclOpExecutor **executor); + +/** + * @brief aclnnGroupedMatmulSwigluQuantWeightNZ的第二段接口，用于执行计算。 + * @param [in] workspace: 在npu device侧申请的workspace内存起址。 + * @param [in] workspaceSize: 在npu + * device侧申请的workspace大小，由第一段接口aclnnGroupedMatmulSwigluQuantWeightNZGetWorkspaceSize获取。 + * @param [in] stream: acl stream流。 + * @param [in] executor: op执行器，包含了算子计算流程。 + * @return aclnnStatus: 返回状态码。 + */ +__attribute__((visibility("default"))) aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNZ(void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/grouped_matmul_swiglu_quant.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/grouped_matmul_swiglu_quant.cpp new file mode 100644 index 000000000..ff39a2fab --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/grouped_matmul_swiglu_quant.cpp @@ -0,0 +1,58 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "opdev/op_log.h" +#include "opdev/op_dfx.h" +#include "opdev/make_op_executor.h" +#include "grouped_matmul_swiglu_quant.h" + +using namespace op; + +namespace l0op { +OP_TYPE_REGISTER(GroupedMatmulSwigluQuant); + +const std::tuple +GroupedMatmulSwigluQuant(const aclTensor *x, const aclTensor *weight, const aclTensor *perChannelScale, + const aclTensor *perTokenScale, const aclTensor *groupList, double limited, + const aclTensor *weightAssistanceMatrix, bool isEnableWeightAssistanceMatrix, int dequantMode, + aclOpExecutor *executor) +{ + L0_DFX(GroupedMatmulSwigluQuant, x, weight, perChannelScale, perTokenScale, weightAssistanceMatrix, groupList, limited, + isEnableWeightAssistanceMatrix, dequantMode); + if (x == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "x is nullptr."); + return std::tuple(nullptr, nullptr); + } + int64_t m = perTokenScale->GetViewShape().GetDim(0); + int64_t n = perChannelScale->GetViewShape().GetDim(1); + int64_t nAfterHalve = static_cast(n / 2); + gert::Shape outShape({m, nAfterHalve}); + gert::Shape scaleOutShape({m}); + auto out = executor->AllocTensor(outShape, DataType::DT_INT8, ge::FORMAT_ND); + auto scaleOut = executor->AllocTensor(scaleOutShape, DataType::DT_FLOAT, ge::FORMAT_ND); + auto ret = INFER_SHAPE(GroupedMatmulSwigluQuant, + OP_INPUT(x, weight, perChannelScale, perTokenScale, weightAssistanceMatrix, groupList), + OP_OUTPUT(out, scaleOut), OP_ATTR(isEnableWeightAssistanceMatrix, dequantMode, limited)); + if (ret != ACLNN_SUCCESS) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "InferShape failed."); + return std::tuple(nullptr, nullptr); + } + ret = ADD_TO_LAUNCHER_LIST_AICORE( + GroupedMatmulSwigluQuant, + OP_INPUT(x, weight, perChannelScale, perTokenScale, weightAssistanceMatrix, groupList), + OP_OUTPUT(out, scaleOut), OP_ATTR(isEnableWeightAssistanceMatrix, dequantMode, limited)); + if (ret != ACLNN_SUCCESS) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "ADD_TO_LAUNCHER_LIST_AICORE failed."); + return std::tuple(nullptr, nullptr); + } + return std::tie(out, scaleOut); +} + +} // namespace l0op \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/grouped_matmul_swiglu_quant.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/grouped_matmul_swiglu_quant.h new file mode 100644 index 000000000..97d9b9d44 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_host/op_api/grouped_matmul_swiglu_quant.h @@ -0,0 +1,23 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_API_INC_LEVEL0_OP_GROUPED_MATMUL_SWIGLU_QUANT_OP_H +#define OP_API_INC_LEVEL0_OP_GROUPED_MATMUL_SWIGLU_QUANT_OP_H + +#include "opdev/op_executor.h" + +namespace l0op { +const std::tuple +GroupedMatmulSwigluQuant(const aclTensor *x, const aclTensor *weight, const aclTensor *perChannelScale, + const aclTensor *perTokenScale, const aclTensor *groupList, double limited, + const aclTensor *weightAssistanceMatrix, bool isEnableWeightAssistanceMatrix, int dequantMode, + aclOpExecutor *executor); +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_pipeline.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_pipeline.h new file mode 100644 index 000000000..2b0268814 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_pipeline.h @@ -0,0 +1,197 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant.cpp + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_PIPELINE_H +#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_PIPELINE_H +#include "grouped_matmul_swiglu_quant.h" +#include +#include "grouped_matmul_swiglu_quant_a8w4_msd_pre.h" +#include "grouped_matmul_swiglu_quant_a8w4_msd_mid.h" +#include "grouped_matmul_swiglu_quant_a8w4_msd_post.h" +#include "grouped_matmul_swiglu_quant_utils.h" +using namespace AscendC; +using namespace matmul; +#ifdef GMM_SWIGLU_QUANT_A8W4_MSD + +namespace GROUPED_MATMUL_SWIGLU_QUANT { + +template +class GMMSwigluQuantPipelineSchedule { +private: + typename mmType::MT &mm; + TPipe *pipe; + const GMMSwigluBaseParams *__restrict gmmBaseParams; + const GMMSwiglu *__restrict gmmSwiglu; + // WorkSpaceSplitConfig控制Workspace切割方式的结构体; + WorkSpaceSplitConfig workspaceSplitConfig; + WorkSpaceSplitConfig tempWorkspaceSplitConfig; + // 记录GM_ADDR的结构体 + GMAddrParams gmAddrParams; + // 前处理GMMA8W4PreProcess类 + GMMA8W4PreProcess preProcess; + // 中间处理GMMA8W4MidProcess类 + GMMA8W4MidProcess midProcess; + // 后处理GMMA8W4PostProcess类 + GMMA8W4PostProcess postProcess; + GlobalTensor groupListGM; + __aicore__ inline void InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig, + int32_t workspaceSplitLoopIdx); + +public: + __aicore__ inline GMMSwigluQuantPipelineSchedule(typename mmType::MT &mm_, + const GMMSwigluBaseParams *__restrict gmmBaseParamsIN, + const GMMSwiglu *__restrict gmmSwigluIN, TPipe *tPipeIN) + : mm(mm_), midProcess(mm), gmmBaseParams(gmmBaseParamsIN), gmmSwiglu(gmmSwigluIN), pipe(tPipeIN) + { + } + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, GM_ADDR xScale, + GM_ADDR weightAssistanceMatrix, GM_ADDR groupList, GM_ADDR y, GM_ADDR yScale, + GM_ADDR workspace); + __aicore__ inline void Process(); +}; + +template +__aicore__ inline void GMMSwigluQuantPipelineSchedule::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, + GM_ADDR xScale, GM_ADDR weightAssistanceMatrix, + GM_ADDR groupList, GM_ADDR y, GM_ADDR yScale, + GM_ADDR workspace) +{ + gmAddrParams.xGM = x; + gmAddrParams.weightGM = weight; + gmAddrParams.weightScaleGM = weightScale; + gmAddrParams.xScaleGM = xScale; + gmAddrParams.weightAuxiliaryMatrixGM = weightAssistanceMatrix; + gmAddrParams.groupListGM = groupList; + gmAddrParams.yGM = y; + gmAddrParams.yScaleGM = yScale; + gmAddrParams.workSpaceGM = workspace; + gmAddrParams.workSpaceOffset1 = gmmBaseParams->workSpaceOffset1 / 2; + gmAddrParams.workSpaceOffset2 = gmmBaseParams->workSpaceOffset1; + gmAddrParams.workSpaceOffset3 = gmmBaseParams->workSpaceOffset1 + gmmBaseParams->workSpaceOffset2 / 2; + groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM); + InitWorkSpaceSplitConfig(workspaceSplitConfig); +} + +template +__aicore__ inline void GMMSwigluQuantPipelineSchedule::Process() +{ + // 1.对每次workspace切分做大循环。 + preProcess.Init(gmAddrParams, gmmBaseParams); + midProcess.Init(gmAddrParams, gmmBaseParams); + postProcess.Init(gmAddrParams, gmmBaseParams, gmmSwiglu); + + // 1.前处理提前下发一次 + preProcess.Process(workspaceSplitConfig, 0, pipe); + for (int64_t workspaceSplitLoopIdx = 0; workspaceSplitLoopIdx < workspaceSplitConfig.loopCount; + workspaceSplitLoopIdx++) { + // 更新workspaceSplitConfig + UpdateWorkSpaceSplitConfig(workspaceSplitConfig, workspaceSplitLoopIdx); + if ASCEND_IS_AIV { + pipe->Reset(); + } + + SyncAll(); + // 2.第n次中处理 && 第n+1次前处理 && 第n-1次后处理 并行 + midProcess.Process(workspaceSplitConfig, workspaceSplitLoopIdx); + + preProcess.Process(workspaceSplitConfig, workspaceSplitLoopIdx + 1, pipe); + if ASCEND_IS_AIV { + pipe->Reset(); + SyncAll(); + } + postProcess.Process(tempWorkspaceSplitConfig, workspaceSplitLoopIdx - 1, pipe); + // 3.第n-1次后处理需要保留第n次的切分数据 + tempWorkspaceSplitConfig = workspaceSplitConfig; + // reset + if ASCEND_IS_AIV { + pipe->Reset(); + } + SyncAll(); + // 3.前一次后处理 && 后一次MM 并行 + } + // reset + if ASCEND_IS_AIV { + pipe->Reset(); + } + SyncAll(); + // // 4.最后一次后处理 + postProcess.Process(workspaceSplitConfig, workspaceSplitConfig.loopCount - 1, pipe); + if ASCEND_IS_AIV { + pipe->Destroy(); + } +} + +template +__aicore__ inline void +GMMSwigluQuantPipelineSchedule::InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig) +{ + workspaceSplitConfig.M = groupListGM.GetValue(gmmSwiglu->groupListLen - 1); + workspaceSplitConfig.loopCount = Ceil(workspaceSplitConfig.M, gmmBaseParams->mLimit); + workspaceSplitConfig.notLastTaskSize = gmmBaseParams->mLimit; + workspaceSplitConfig.lastLoopTaskSize = + workspaceSplitConfig.M - (workspaceSplitConfig.loopCount - 1) * gmmBaseParams->mLimit; + workspaceSplitConfig.leftMatrixStartIndex = 0; + workspaceSplitConfig.rightMatrixExpertStartIndex = 0; + workspaceSplitConfig.rightMatrixExpertNextStartIndex = 0; + workspaceSplitConfig.isLastLoop = false; +} + +template +__aicore__ inline void +GMMSwigluQuantPipelineSchedule::UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig, + int32_t workspaceSplitLoopIdx) +{ + if (workspaceSplitLoopIdx < 0) + return; + workspaceSplitConfig.leftMatrixStartIndex = workspaceSplitLoopIdx * gmmBaseParams->mLimit; + workspaceSplitConfig.rightMatrixExpertStartIndex = workspaceSplitConfig.rightMatrixExpertNextStartIndex; + workspaceSplitConfig.rightMatrixExpertEndIndex = workspaceSplitConfig.rightMatrixExpertStartIndex; + // 计算右专家矩阵的终止索引(rightMatrixExpertEndIndex) 和下一次的起始索引(rightMatrixExpertNextStartIndex) + int32_t curTaskNum = 0; + int32_t nextTaskNum = 0; + while (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwiglu->groupListLen) { + curTaskNum = groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex) - + workspaceSplitConfig.leftMatrixStartIndex; + int32_t nextTaskIdx = workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwiglu->groupListLen - 1 ? + gmmSwiglu->groupListLen - 1 : + workspaceSplitConfig.rightMatrixExpertEndIndex + 1; + nextTaskNum = groupListGM.GetValue(nextTaskIdx) - workspaceSplitConfig.leftMatrixStartIndex; + if (curTaskNum > gmmBaseParams->mLimit) { + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex; + break; + } else if (curTaskNum == gmmBaseParams->mLimit && nextTaskNum > gmmBaseParams->mLimit) { + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex + 1; + break; + } else if (nextTaskNum > gmmBaseParams->mLimit) { + workspaceSplitConfig.rightMatrixExpertEndIndex++; + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex; + break; + } + workspaceSplitConfig.rightMatrixExpertEndIndex++; + } + workspaceSplitConfig.isLastLoop = workspaceSplitLoopIdx == workspaceSplitConfig.loopCount - 1 ? true : false; + + if (workspaceSplitConfig.isLastLoop) { + workspaceSplitConfig.rightMatrixExpertEndIndex = + workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwiglu->groupListLen ? + gmmSwiglu->groupListLen - 1 : + workspaceSplitConfig.rightMatrixExpertEndIndex; + } +} + +} // namespace GROUPED_MATMUL_SWIGLU_QUANT +#endif // GMM_SWIGLU_QUANT_A8W4_MSD +#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_PIPELINE_H diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant.cpp new file mode 100644 index 000000000..2d562f4ab --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant.cpp @@ -0,0 +1,102 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant.cpp + * \brief + */ +#include "grouped_matmul_swiglu_quant.h" +#include "grouped_matmul_swiglu_pipeline.h" +#include "grouped_matmul_swiglu_quant_utils.h" +#include +#include "grouped_matmul_swiglu_quant_split_ws.h" +using namespace AscendC; +using namespace matmul; +using namespace GROUPED_MATMUL_SWIGLU_QUANT; + +#define GMM_CV_SPLIT_IMP(computeClass, dtypeWeightScale, transA, transB, sync) \ + do { \ + using xType = MatmulType; \ + using weightType = MatmulType; \ + using yType = MatmulType; \ + using matmulType = MMImplTypeStatic; \ + matmulType::MT mm; \ + GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, gmmSwigluBaseParams, gmmSwigluBaseParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, gmmSwiglu, gmmSwiglu_, tiling); \ + if ASCEND_IS_AIC { \ + mm.SetSubBlockIdx(0); \ + mm.Init(&mmTilingData_, &tPipe); \ + } \ + computeClass computeOp(mm); \ + computeOp.Init(x, weight, weightScale, xScale, groupList, y, yScale, user1, &gmmSwigluBaseParams_, \ + &mmTilingData_, &gmmSwiglu_, &tPipe); \ + computeOp.Process(); \ + } while (0) + +#define GMM_CV_SPLIT_IMP_A8W4_MSD(computeClass, dtypeWeightScale, transA, transB, sync) \ + do { \ + GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, gmmSwigluBaseParams, gmmSwigluBaseParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, gmmSwiglu, gmmSwiglu_, tiling); \ + using xType = MatmulType; \ + using weightType = MatmulType; \ + using yType = MatmulType; \ + using matmulType = MMImplType; \ + matmulType::MT mm; \ + if ASCEND_IS_AIC { \ + mm.SetSubBlockIdx(0); \ + mm.Init(&mmTilingData_); \ + } \ + computeClass op(mm, &gmmSwigluBaseParams_, &gmmSwiglu_, &tPipe); \ + op.Init(x, weight, weightScale, xScale, weightAssistanceMatrix, groupList, y, yScale, user1); \ + \ + op.Process(); \ + } while (0) + +extern "C" __global__ __aicore__ void grouped_matmul_swiglu_quant(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, + GM_ADDR xScale, GM_ADDR weightAssistanceMatrix, + GM_ADDR groupList, GM_ADDR y, GM_ADDR yScale, + GM_ADDR workspace, GM_ADDR tiling) +{ + TPipe tPipe; + AscendCUtils::SetOverflow(1); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + GM_ADDR user1 = GetUserWorkspace(workspace); +#if defined(GMM_SWIGLU_QUANT_A8W8) + if (TILING_KEY_IS(0)) { // antiquant msd + KERNEL_TASK_TYPE(0, KERNEL_TYPE_MIX_AIC_1_2); + GMM_CV_SPLIT_IMP(GMMSwigluCompute, // computeClass + DTYPE_WEIGHT_SCALE, + false, // transA + false, // transB + false // sync + ); + } else if (TILING_KEY_IS(1)) { + KERNEL_TASK_TYPE(1, KERNEL_TYPE_MIX_AIC_1_2); + GMM_CV_SPLIT_IMP(GMMSwigluSplitWorkSpaceCompute, // computeClass + DTYPE_WEIGHT_SCALE, + false, // transA + false, // transB + false // sync + ); + } +#elif defined(GMM_SWIGLU_QUANT_A8W4_MSD) + if (TILING_KEY_IS(2)) { + KERNEL_TASK_TYPE(2, KERNEL_TYPE_MIX_AIC_1_2); + GMM_CV_SPLIT_IMP_A8W4_MSD(GMMSwigluQuantPipelineSchedule, // computeClass + DTYPE_WEIGHT_SCALE, + false, // transA + false, // transB + false // sync + ); + } +#endif +} diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant.h new file mode 100644 index 000000000..27563c79a --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant.h @@ -0,0 +1,547 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant.h + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_H +#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_H + +#include "grouped_matmul_swiglu_quant_utils.h" +namespace GROUPED_MATMUL_SWIGLU_QUANT { +/** @brief internal computation class + */ +template +class GMMSwigluCompute { +public: + using AT = typename mmType::AT::T; + using BT = typename mmType::BT::T; + using B = typename mmType::BT; + using CT = typename mmType::CT::T; + using BiasT = typename mmType::BiasT::T; + using WT = int8_t; + constexpr static bool transposeX = mmType::AT::isTrans; + constexpr static bool transposeW = mmType::BT::isTrans; + static constexpr float FLOAT_INF = 3e+99; + /** @brief constructor */ + __aicore__ inline GMMSwigluCompute(typename mmType::MT &mm_) : mm(mm_) + { + } + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale, + GM_ADDR groupList, GM_ADDR quantOutput, GM_ADDR quantScaleOutput, GM_ADDR workspace, + const GMMSwigluBaseParams *__restrict gmmBaseParamsIN, + const TCubeTiling *__restrict mmTilingDataIN, const GMMSwiglu *__restrict gmmSwigluIN, + TPipe *tPipeIN); + __aicore__ inline void Process(); + +private: + __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig &mnConfig, uint32_t coreIdx); + + __aicore__ inline void UpdateMnConfig(MNConfig &mnConfig); + + __aicore__ inline void SetMNConfig(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig); + + __aicore__ inline void SetMKN(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig); + + __aicore__ inline uint64_t GetWOffset(uint32_t tailN, uint32_t k); + + __aicore__ inline void CubeProcess(MNConfig &mnConfig); + __aicore__ inline void VecProcess(VecConfig &vecConfig); + __aicore__ inline void MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock, const uint32_t count, + const uint32_t thresholdM_dimN); + template + __aicore__ inline void UpdateChannelScale(uint32_t loopidx, VecConfig &vecConfig); + __aicore__ inline void VectorCompute(uint32_t loopidx, VecConfig &vecConfig); + template + __aicore__ inline void PreLoadTokenAndChannel(LocalTensor &channelScaleLocal, VecConfig &vecConfig); + __aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig); + __aicore__ inline void customDataCopyIn(uint32_t outLoopIdx, VecConfig &vecConfig); + __aicore__ inline void customDataCopyOut(VecConfig &vecConfig); + __aicore__ inline void Dequant(uint32_t loopidx, VecConfig &vecConfig); + __aicore__ inline void Quant(uint32_t loopidx); + __aicore__ inline void Swiglu(uint32_t loopidx); + +private: + typename mmType::MT &mm; + const GMMSwigluBaseParams *__restrict gmmBaseParams; + const GMMSwiglu *__restrict gmmSwiglu; + const TCubeTiling *__restrict mmTilingData; + TPipe *pipe; + GlobalTensor xGM; + GlobalTensor weightGM; + GlobalTensor perChannelScaleGM; + GlobalTensor perTokenScaleGM; + GlobalTensor groupListGM; + GlobalTensor quantOutputGM; + GlobalTensor quantScaleOutputGM; + GlobalTensor mmOutGM; + // define the que + TQue mmOutQueue; + TQue perChannelScaleInQueue; + TQue quantOutQueue; + TQue quantScaleOutQueue; + TBuf reduceWorkspace; + uint32_t blockIdx = 0; + int32_t preOffset = 0; + int64_t aicCoreNum = 0; + int64_t aivCoreNum = 0; + float limited = FLOAT_INF; + GM_ADDR xTensorPtr; + GM_ADDR weightTensorPtr; +}; + +template +__aicore__ inline void GMMSwigluCompute::Init( + GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale, GM_ADDR groupList, GM_ADDR quantOutput, + GM_ADDR quantScaleOutput, GM_ADDR workspace, const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIn, + const TCubeTiling *__restrict mmTilingDataIN, const GMMSwiglu *__restrict gmmSwigluIN, TPipe *tPipeIN) +{ + aicCoreNum = GetBlockNum(); + aivCoreNum = aicCoreNum * 2; + blockIdx = GetBlockIdx(); + mmTilingData = mmTilingDataIN; + gmmBaseParams = gmmSwigluBaseParamsIn; + gmmSwiglu = gmmSwigluIN; + pipe = tPipeIN; + xTensorPtr = x; + limited = gmmBaseParams->limited; + weightTensorPtr = weight; + groupListGM.SetGlobalBuffer((__gm__ int64_t *)groupList, gmmSwiglu->groupListLen); + mmOutGM.SetGlobalBuffer((__gm__ int32_t *)workspace, gmmBaseParams->M * gmmSwiglu->tokenLen); + if ASCEND_IS_AIV { + perChannelScaleGM.SetGlobalBuffer((__gm__ CHANNELDTYPE *)perChannelScale, + gmmSwiglu->groupListLen * gmmSwiglu->tokenLen); + perTokenScaleGM.SetGlobalBuffer((__gm__ float *)perTokenScale, gmmBaseParams->M); + quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)quantOutput, + gmmBaseParams->M * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR); + quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)quantScaleOutput, gmmBaseParams->M); + } +} + +template +__aicore__ inline void GMMSwigluCompute::Process() +{ + MNConfig mnConfig; + VecConfig vecConfig; + CubeProcess(mnConfig); + VecProcess(vecConfig); +} + +template +template +__aicore__ inline void +GMMSwigluCompute::PreLoadTokenAndChannel(LocalTensor &channelScaleLocal, + VecConfig &vecConfig) +{ + DataCopyExtParams copyChannelParams{1, static_cast(gmmSwiglu->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + if constexpr (!IsSameType::value) { + LocalTensor dstLocalT = channelScaleLocal.template ReinterpretCast(); + DataCopyPad(dstLocalT[gmmSwiglu->tokenLen], perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], + copyChannelParams, padParams); + PipeBarrier(); + Cast(channelScaleLocal, dstLocalT[gmmSwiglu->tokenLen], RoundMode::CAST_NONE, gmmSwiglu->tokenLen); + } else { + DataCopyPad(channelScaleLocal, perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], + copyChannelParams, padParams); + } + perChannelScaleInQueue.EnQue(channelScaleLocal); +} + +template +__aicore__ inline void GMMSwigluCompute::MMCompute(uint32_t groupIdx, MNConfig &mnConfig, + uint32_t coreIdx) +{ + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN; + uint32_t curSingleM = + mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM : mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + uint64_t xOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.k; + if constexpr (transposeX) { + xOffset = mnConfig.mIdx * mnConfig.singleM; + } + uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN; + xGM.SetGlobalBuffer((__gm__ int8_t *)xTensorPtr + mnConfig.xBaseOffset); + weightGM.SetGlobalBuffer((__gm__ int8_t *)weightTensorPtr + mnConfig.wBaseOffset + GetWOffset(tailN, mnConfig.k)); + if (mnConfig.blockDimM == 1) { + weightGM.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } + mnConfig.workSpaceOffset = outOffset + mnConfig.yBaseOffset; + mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k); + mm.SetSingleShape(curSingleM, curSingleN, mnConfig.k); + mm.SetTensorA(xGM[xOffset], transposeX); + mm.SetTensorB(weightGM, transposeW); + mm.template IterateAll(mmOutGM[mnConfig.workSpaceOffset], 0); +} + +template +__aicore__ inline void GMMSwigluCompute::UpdateMnConfig(MNConfig &mnConfig) +{ + if constexpr (B::format == CubeFormat::NZ) { + mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<32>(mnConfig.n); // 16: nz format last two dim size + } else { + mnConfig.wBaseOffset += mnConfig.k * mnConfig.n; + } + mnConfig.nAxisBaseOffset += mnConfig.n; + mnConfig.mAxisBaseOffset += mnConfig.m; + mnConfig.xBaseOffset += mnConfig.m * mnConfig.k; + mnConfig.yBaseOffset += mnConfig.m * mnConfig.n; +} + +template +__aicore__ inline void GMMSwigluCompute::SetMNConfig(const int32_t splitValue, + const uint32_t groupIdx, + MNConfig &mnConfig) +{ + SetMKN(splitValue, groupIdx, mnConfig); + mnConfig.baseM = BASIC_M; + mnConfig.baseN = BASIC_N; + mnConfig.singleM = SINGLE_CORE_M; + mnConfig.singleN = SINGLE_CORE_N; +} + +template +__aicore__ inline void GMMSwigluCompute::SetMKN(const int32_t splitValue, + const uint32_t groupIdx, MNConfig &mnConfig) +{ + mnConfig.m = static_cast(splitValue); + mnConfig.k = gmmBaseParams->K; // tilingData + mnConfig.n = gmmBaseParams->N; // tilingData +} + +template +__aicore__ inline uint64_t GMMSwigluCompute::GetWOffset(uint32_t tailN, uint32_t k) +{ + uint64_t wOffset = 0; + if constexpr (mmType::BT::format == CubeFormat::NZ) { + wOffset = tailN * AlignUp<16>(k); // 16: nz format last two dim size + } else { + wOffset = tailN; + } + return wOffset; +} + +template +__aicore__ inline void GMMSwigluCompute::CubeProcess(MNConfig &mnConfig) +{ + if ASCEND_IS_AIC { + preOffset = 0; + int32_t prevSplitValue = 0; + for (uint32_t groupIdx = 0, count = 0; groupIdx < gmmSwiglu->groupListLen; ++groupIdx) { + UpdateMnConfig(mnConfig); + int32_t currSplitValue = static_cast(groupListGM.GetValue(groupIdx)); + int32_t splitValue = currSplitValue - prevSplitValue; + prevSplitValue = currSplitValue; + SetMNConfig(splitValue, groupIdx, mnConfig); + if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) { + continue; + } + mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM); + mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN); + + uint32_t curCount = count + mnConfig.blockDimM * mnConfig.blockDimN; + uint32_t curBlock = blockIdx >= count ? blockIdx : blockIdx + gmmBaseParams->coreNum; + uint32_t thresholdM_dimN = THRESHOLD_BLOCK_NUM * mnConfig.blockDimN; + + while (curBlock < curCount) { + MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN); + MMCompute(groupIdx, mnConfig, blockIdx); + curBlock += aicCoreNum; + } + count = curCount % gmmBaseParams->coreNum; + } + SyncAll(); + } +} + +template +__aicore__ inline void GMMSwigluCompute::VecProcess(VecConfig &vecConfig) +{ + if ASCEND_IS_AIV { + UpdateVecConfig(blockIdx, vecConfig); + if (blockIdx < vecConfig.usedCoreNum) { + LocalTensor channelScaleLocal = perChannelScaleInQueue.AllocTensor(); + LocalTensor mmLocal = mmOutQueue.AllocTensor(); + LocalTensor quantLocal = quantOutQueue.AllocTensor(); + LocalTensor quantScaleLocal = quantScaleOutQueue.AllocTensor(); + mmOutQueue.EnQue(mmLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); + quantOutQueue.EnQue(quantLocal); + PreLoadTokenAndChannel(channelScaleLocal, vecConfig); + } + SyncAll(); + if (blockIdx < vecConfig.usedCoreNum) { + for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) { + vecConfig.innerLoopNum = + outLoopIdx == (vecConfig.outLoopNum - 1) ? vecConfig.tailLoopNum : gmmSwiglu->maxProcessRowNum; + customDataCopyIn(outLoopIdx, vecConfig); + for (uint32_t innerLoopIdx = 0; innerLoopIdx < vecConfig.innerLoopNum; innerLoopIdx++) { + UpdateChannelScale(innerLoopIdx, vecConfig); + VectorCompute(innerLoopIdx, vecConfig); + } + customDataCopyOut(vecConfig); + } + + LocalTensor channelScaleLocal = perChannelScaleInQueue.DeQue(); + LocalTensor mmLocal = mmOutQueue.DeQue(); + LocalTensor quantLocal = quantOutQueue.DeQue(); + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + perChannelScaleInQueue.FreeTensor(channelScaleLocal); + mmOutQueue.FreeTensor(mmLocal); + quantScaleOutQueue.FreeTensor(quantScaleLocal); + quantOutQueue.FreeTensor(quantLocal); + } else { + return; + } + } +} + +template +__aicore__ inline void +GMMSwigluCompute::MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock, + const uint32_t count, const uint32_t thresholdM_dimN) +{ + mnConfig.mIdx = (curBlock - count) / mnConfig.blockDimN; + mnConfig.nIdx = (curBlock - count) % mnConfig.blockDimN; +} + +template +__aicore__ inline void GMMSwigluCompute::UpdateVecConfig(uint32_t blockIdx, + VecConfig &vecConfig) +{ + // 第一步 读取grouplist reduceSum 计算总数据个数 + int64_t prevM = 0; + for (uint32_t groupIdx = 0; groupIdx < gmmSwiglu->groupListLen; groupIdx++) { + int64_t currM = groupListGM.GetValue(groupIdx); + int64_t tempM = currM - prevM; + prevM = currM; + vecConfig.M += tempM; + } + // 第二步 计算分核 + uint32_t eachCoreTaskNum = (vecConfig.M + aivCoreNum - 1) / aivCoreNum; + vecConfig.usedCoreNum = vecConfig.M >= aivCoreNum ? aivCoreNum : vecConfig.M; + uint32_t tailCoreIdx = vecConfig.M - (eachCoreTaskNum - 1) * vecConfig.usedCoreNum; + vecConfig.taskNum = blockIdx < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1; + vecConfig.startIdx = + blockIdx < tailCoreIdx ? eachCoreTaskNum * blockIdx : ((eachCoreTaskNum - 1) * blockIdx + tailCoreIdx); + vecConfig.curIdx = vecConfig.startIdx; + vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen; + vecConfig.curOffset = vecConfig.startOffset; + int64_t curStartIdx = vecConfig.startIdx; + prevM = 0; + for (uint32_t groupIdx = 0; groupIdx < gmmSwiglu->groupListLen; groupIdx++) { + int64_t currM = groupListGM.GetValue(groupIdx); + int64_t tempM = currM - prevM; + prevM = currM; + if (curStartIdx >= 0 && curStartIdx - tempM < 0) { + vecConfig.curGroupIdx = groupIdx; + vecConfig.nextUpadteInterVal = tempM - curStartIdx; + } + curStartIdx -= tempM; + } + // 第三步 计算总数据量 + vecConfig.outLoopNum = (vecConfig.taskNum + gmmSwiglu->maxProcessRowNum - 1) / gmmSwiglu->maxProcessRowNum; + vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwiglu->maxProcessRowNum ? + vecConfig.taskNum % gmmSwiglu->maxProcessRowNum : + gmmSwiglu->maxProcessRowNum; + pipe->Reset(); + // 第四步 申请空间 + pipe->InitBuffer(mmOutQueue, DOUBLE_BUFFER, gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen * sizeof(int32_t)); + pipe->InitBuffer(perChannelScaleInQueue, DOUBLE_BUFFER, gmmSwiglu->tokenLen * sizeof(float)); + pipe->InitBuffer(quantOutQueue, DOUBLE_BUFFER, + gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)); + pipe->InitBuffer(quantScaleOutQueue, DOUBLE_BUFFER, + AlignUp(gmmSwiglu->maxProcessRowNum, ALIGN_8_ELE) * sizeof(float)); + // two 32 byte buffer for reduceMax calculation in Quant. + pipe->InitBuffer(reduceWorkspace, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(float) + UB_BLOCK_UNIT_SIZE + + UB_BLOCK_UNIT_SIZE); +} + +template +__aicore__ inline void GMMSwigluCompute::customDataCopyIn(uint32_t outLoopIdx, + VecConfig &vecConfig) +{ + LocalTensor _inMMLocal_0 = mmOutQueue.DeQue(); + DataCopyExtParams copyParams_0{ + 1, static_cast(vecConfig.innerLoopNum * gmmSwiglu->tokenLen * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams padParams_0{false, 0, 0, 0}; + DataCopyPad(_inMMLocal_0, mmOutGM[vecConfig.curOffset], copyParams_0, padParams_0); + + mmOutQueue.EnQue(_inMMLocal_0); + + LocalTensor _inMMLocal_1 = mmOutQueue.DeQue(); + + Cast(_inMMLocal_1.ReinterpretCast(), _inMMLocal_1, RoundMode::CAST_NONE, + vecConfig.innerLoopNum * gmmSwiglu->tokenLen); + + mmOutQueue.EnQue(_inMMLocal_1); + LocalTensor _inMMLocal_2 = mmOutQueue.DeQue(); + SetFlag(EVENT_ID0); + for (uint32_t i = 0; i < vecConfig.innerLoopNum; i++) { + WaitFlag(EVENT_ID0); + float scale = perTokenScaleGM.GetValue(vecConfig.curIdx); + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + Muls(_inMMLocal_2[i * gmmSwiglu->tokenLen], _inMMLocal_2[i * gmmSwiglu->tokenLen], scale, gmmSwiglu->tokenLen); + SetFlag(EVENT_ID0); + vecConfig.curIdx++; + } + WaitFlag(EVENT_ID0); + vecConfig.curOffset = vecConfig.curIdx * gmmSwiglu->tokenLen; + mmOutQueue.EnQue(_inMMLocal_2); +} + +template +template +__aicore__ inline void GMMSwigluCompute::UpdateChannelScale(uint32_t loopIdx, + VecConfig &vecConfig) +{ + // 更新perChannel + if (unlikely(vecConfig.nextUpadteInterVal == 0)) { + int64_t loop = gmmSwiglu->groupListLen - vecConfig.curGroupIdx; + while (loop--) { + int64_t curTemp = groupListGM.GetValue(vecConfig.curGroupIdx); + vecConfig.curGroupIdx++; + int64_t nextTemp = groupListGM.GetValue(vecConfig.curGroupIdx); + if (nextTemp != curTemp) { + vecConfig.nextUpadteInterVal = nextTemp - curTemp; + break; + } + } + LocalTensor _inChannel = perChannelScaleInQueue.DeQue(); + DataCopyExtParams copyParams{1, static_cast(gmmSwiglu->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + if constexpr (!IsSameType::value) { + LocalTensor dstLocalT = _inChannel.template ReinterpretCast(); + DataCopyPad(dstLocalT[gmmSwiglu->tokenLen], perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], + copyParams, padParams); + PipeBarrier(); + Cast(_inChannel, dstLocalT[gmmSwiglu->tokenLen], RoundMode::CAST_NONE, gmmSwiglu->tokenLen); + } else { + DataCopyPad(_inChannel, perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], copyParams, + padParams); + } + PipeBarrier(); + perChannelScaleInQueue.EnQue(_inChannel); + } +} + +template +__aicore__ inline void GMMSwigluCompute::VectorCompute(uint32_t loopIdx, + VecConfig &vecConfig) +{ + Dequant(loopIdx, vecConfig); + Swiglu(loopIdx); + Quant(loopIdx); +} + +template +__aicore__ inline void GMMSwigluCompute::Dequant(uint32_t loopIdx, VecConfig &vecConfig) +{ + // perChanelScale * perTokenScale + LocalTensor mmLocal = mmOutQueue.DeQue(); + LocalTensor perChannelLocal = perChannelScaleInQueue.DeQue(); + Mul(mmLocal[loopIdx * gmmSwiglu->tokenLen], mmLocal[loopIdx * gmmSwiglu->tokenLen], perChannelLocal, + gmmSwiglu->tokenLen); + vecConfig.nextUpadteInterVal--; + mmOutQueue.EnQue(mmLocal); + perChannelScaleInQueue.EnQue(perChannelLocal); +} + +template +__aicore__ inline void GMMSwigluCompute::Swiglu(uint32_t loopIdx) +{ + // 高阶API swiglu + LocalTensor _inMMLocal = mmOutQueue.DeQue(); + float beta = 1.0f; + LocalTensor workspaceLocal = reduceWorkspace.Get(); + LocalTensor src0Local = + _inMMLocal[loopIdx * gmmSwiglu->tokenLen + gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR]; + LocalTensor src1Local = _inMMLocal[loopIdx * gmmSwiglu->tokenLen]; + if (limited > 0.0f) { + Mins(src0Local, src0Local, limited, gmmSwiglu->tokenLen / 2); + PipeBarrier(); + Maxs(src0Local, src0Local, (-1.0f * limited), gmmSwiglu->tokenLen / 2); + PipeBarrier(); + Mins(src1Local, src1Local, limited, gmmSwiglu->tokenLen / 2); + PipeBarrier(); + } + SwiGLU(workspaceLocal, src0Local, src1Local, beta, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR); + PipeBarrier(); + DataCopyParams repeatParams{1, static_cast((gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR) / ALIGN_8_ELE), 0, + 0}; + DataCopy(_inMMLocal[loopIdx * gmmSwiglu->tokenLen], workspaceLocal, repeatParams); + mmOutQueue.EnQue(_inMMLocal); +} + +template +__aicore__ inline void GMMSwigluCompute::Quant(uint32_t loopIdx) +{ + LocalTensor _inMMLocal = mmOutQueue.DeQue(); + uint64_t preOffset = loopIdx * gmmSwiglu->tokenLen; + uint64_t halfTokenLen = gmmSwiglu->tokenLen / BISECT; + Abs(_inMMLocal[preOffset + gmmSwiglu->tokenLen / BISECT], _inMMLocal[preOffset], halfTokenLen); + PipeBarrier(); + // reduceMax + LocalTensor workLocal = reduceWorkspace.Get(halfTokenLen); + LocalTensor reduceResLocal = + reduceWorkspace.GetWithOffset(FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float)); + LocalTensor reduceTmpLocal = reduceWorkspace.GetWithOffset( + FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float) + UB_BLOCK_UNIT_SIZE); + ReduceMaxTemplate(reduceResLocal, workLocal, _inMMLocal[preOffset + gmmSwiglu->tokenLen / BISECT], reduceTmpLocal, + static_cast(halfTokenLen)); + + int32_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float quantScale = reduceResLocal.GetValue(0) / QUANT_SCALE_INT8; + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + quantScaleLocal.SetValue(loopIdx, quantScale); + quantScale = 1 / quantScale; + int32_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Muls(_inMMLocal[preOffset], _inMMLocal[preOffset], quantScale, halfTokenLen); + PipeBarrier(); + LocalTensor quantLocal = quantOutQueue.DeQue(); + int32_t dstTempOffset = static_cast(preOffset / BISECT); + int32_t srcTempOffset = static_cast(preOffset); + int32_t tempCount = static_cast(halfTokenLen); + LocalTensor castSpace = reduceWorkspace.Get(UB_BLOCK_UNIT_SIZE); + CastFp32ToInt8Template(quantLocal, _inMMLocal, castSpace, dstTempOffset, srcTempOffset, tempCount); + mmOutQueue.EnQue(_inMMLocal); + quantOutQueue.EnQue(quantLocal); +} + +template +__aicore__ inline void GMMSwigluCompute::customDataCopyOut(VecConfig &vecConfig) +{ + // perChanelScale * perTokenScale + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + DataCopyParams copyParams_0{1, (uint16_t)(vecConfig.innerLoopNum * sizeof(float)), 0, 0}; + PipeBarrier(); + DataCopyPad(quantScaleOutputGM[vecConfig.startIdx], quantScaleLocal, copyParams_0); + LocalTensor quantLocal = quantOutQueue.DeQue(); + DataCopyParams copyParams_1{ + 1, (uint16_t)(vecConfig.innerLoopNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)), 0, 0}; + PipeBarrier(); + DataCopyPad(quantOutputGM[vecConfig.startIdx * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR], quantLocal, + copyParams_1); + PipeBarrier(); + vecConfig.startIdx += vecConfig.innerLoopNum; + vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen; + quantOutQueue.EnQue(quantLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); +} + +} // namespace GROUPED_MATMUL_SWIGLU_QUANT +#endif // ASCENDC_GROUPED_MATMUL_QUANT_MIXCORE_H diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_a8w4_msd_mid.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_a8w4_msd_mid.h new file mode 100644 index 000000000..b5b36e2d7 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_a8w4_msd_mid.h @@ -0,0 +1,251 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_antiquant_a8w4_msd.h + * \brief + */ + +#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_MID_H +#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_MID_H + +#include "grouped_matmul_swiglu_quant_utils.h" +#include "grouped_matmul_swiglu_quant.h" + +#ifdef GMM_SWIGLU_QUANT_A8W4_MSD +namespace GROUPED_MATMUL_SWIGLU_QUANT { +using namespace matmul; +using namespace AscendC; + +constexpr uint32_t BUFFER_NUM = 1; + +template +__aicore__ inline void DataCopyPad2DA8W4(const LocalTensor dst, const GlobalTensor src, uint32_t dim1, + uint32_t dim0, uint32_t srcDim0) +{ + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + params.srcStride = (srcDim0 - dim0) * sizeof(T); + // 32: int32 -> float16, 为防止跨行数据进入同一32B block，提前每行按偶数block对齐 + params.dstStride = Ceil(dim0 * sizeof(T), 32) % 2; + + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(dst, src, params, padParams); +} + +template +__aicore__ inline void DataCopyPad2DA8W4ND(const LocalTensor dst, const GlobalTensor src, uint32_t dim1, + uint32_t dim0, uint32_t srcDim0) +{ + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + params.srcStride = (srcDim0 - dim0) * sizeof(T); + params.dstStride = 0; + + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(dst, src, params, padParams); + return; +} + +template +__aicore__ inline void DataCopyPad2DA8W4(const GlobalTensor dst, const LocalTensor src, uint32_t dim1, + uint32_t dim0, uint32_t srcDim0, uint32_t dstDim0) +{ + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + // 32: ub访问粒度为32B + params.srcStride = (srcDim0 - dim0) * sizeof(T) / 32; + params.dstStride = (dstDim0 - dim0) * sizeof(T); + DataCopyPad(dst, src, params); +} + +template +class GMMA8W4MidProcess { +public: + using bT = typename mmType::BT; + +public: + __aicore__ inline GMMA8W4MidProcess(typename mmType::MT &matmul) : mm(matmul) + { + } + __aicore__ inline void Init(const GMAddrParams gmAddrParams, + const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN); + __aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx); + +private: + __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig &mnConfig, WorkSpaceSplitConfig &workspaceSplitConfig); + __aicore__ inline void SetMNConfig(const int32_t splitValue, MNConfig &mnConfig); + __aicore__ inline void UpdateMnConfig(MNConfig &mnConfig); + +private: + typename mmType::MT &mm; + const uint32_t HALF_ALIGN = 16; + GlobalTensor xGM; + GlobalTensor xGM1; + GlobalTensor xGM2; + GlobalTensor weightGM; + + GlobalTensor mmOutGM; + GlobalTensor mmOutGM1; + GlobalTensor mmOutGM2; + GlobalTensor groupListGM; + GlobalTensor weightScaleGM; + // define the que + uint32_t subBlockIdx = 0; + uint32_t coreIdx = 0; + uint32_t quantGroupSize = 0; + uint32_t vecCount = 0; + uint32_t xRowSumCount = 0; + const GMMSwigluBaseParams *__restrict gmmBaseParams; +}; + +template +__aicore__ inline void GMMA8W4MidProcess::Init(const GMAddrParams gmAddrParams, + const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN) +{ + if ASCEND_IS_AIC { + gmmBaseParams = gmmSwigluBaseParamsIN; + xRowSumCount = gmmBaseParams->M; + xGM1.SetGlobalBuffer((__gm__ int4b_t *)gmAddrParams.workSpaceGM); // 从前处理中获得的结果 + xGM2.SetGlobalBuffer( + (__gm__ int4b_t *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset1)); + weightGM.SetGlobalBuffer((__gm__ int4b_t *)gmAddrParams.weightGM); + weightScaleGM.SetGlobalBuffer((__gm__ uint64_t *)gmAddrParams.weightScaleGM); + groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM); + mmOutGM1.SetGlobalBuffer( + (__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset2)); + mmOutGM2.SetGlobalBuffer( + (__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset3)); + quantGroupSize = gmmBaseParams->K / gmmBaseParams->quantGroupNum; // 约束为整除关系 + subBlockIdx = GetSubBlockIdx(); + coreIdx = GetBlockIdx(); + } +} + +template +__aicore__ inline void GMMA8W4MidProcess::UpdateMnConfig(MNConfig &mnConfig) +{ + if constexpr (bT::format == CubeFormat::NZ) { + mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<32>(mnConfig.n); // 16: nz format last two dim size + } else { + mnConfig.wBaseOffset += mnConfig.k * mnConfig.n; + } + mnConfig.nAxisBaseOffset += mnConfig.n; + mnConfig.mAxisBaseOffset += mnConfig.m; + mnConfig.xBaseOffset += mnConfig.m * mnConfig.k; + mnConfig.yBaseOffset += mnConfig.m * mnConfig.n; +} + +template +__aicore__ inline void GMMA8W4MidProcess::SetMNConfig(const int32_t splitValue, MNConfig &mnConfig) +{ + mnConfig.m = static_cast(splitValue); + mnConfig.baseM = gmmBaseParams->baseM; + mnConfig.baseN = gmmBaseParams->baseN; + mnConfig.singleM = gmmBaseParams->baseM; + mnConfig.singleN = gmmBaseParams->baseN; +} + +template +__aicore__ inline void GMMA8W4MidProcess::Process(WorkSpaceSplitConfig &workspaceSplitConfig, + int64_t workspaceSplitLoopIdx) +{ + if ASCEND_IS_AIC { + if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount || workspaceSplitLoopIdx < 0) { + return; + } + xGM = (workspaceSplitLoopIdx % 2 == 0 ? xGM1 : xGM2); + mmOutGM = (workspaceSplitLoopIdx % 2 == 0 ? mmOutGM1 : mmOutGM2); + MNConfig mnConfig; + mnConfig.baseM = gmmBaseParams->baseM; + mnConfig.baseN = gmmBaseParams->baseN; + mnConfig.singleM = gmmBaseParams->baseM; + mnConfig.singleN = gmmBaseParams->baseN; + mnConfig.k = gmmBaseParams->K; // tilingData + mnConfig.n = gmmBaseParams->N; // tilingData + mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN); + int32_t prevSplitValue = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize; + for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex, preCount = 0; + groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; ++groupIdx) { + UpdateMnConfig(mnConfig); + int32_t currSplitValue = static_cast(groupListGM.GetValue(groupIdx)); + currSplitValue = currSplitValue > (workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit ? + (workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit : + currSplitValue; + + int32_t splitValue = (currSplitValue - prevSplitValue) * 2; // 2: int8 has been split in 2 int4 + prevSplitValue = currSplitValue; + + SetMNConfig(splitValue, mnConfig); + if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) { + continue; + } + mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM); + mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k); + uint32_t curCount = preCount + mnConfig.blockDimN * mnConfig.blockDimM; + uint32_t curBlock = coreIdx >= preCount ? coreIdx : coreIdx + gmmBaseParams->coreNum; + while (curBlock < curCount) { + mnConfig.mIdx = (curBlock - preCount) / mnConfig.blockDimN; + mnConfig.nIdx = (curBlock - preCount) % mnConfig.blockDimN; + MMCompute(groupIdx, mnConfig, workspaceSplitConfig); + curBlock += gmmBaseParams->coreNum; + } + preCount = curCount % gmmBaseParams->coreNum; + } + } +} + +template +__aicore__ inline void GMMA8W4MidProcess::MMCompute(uint32_t groupIdx, MNConfig &mnConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint32_t curSingleN = mnConfig.singleN; + if (unlikely(mnConfig.nIdx == mnConfig.blockDimN - 1)) { + curSingleN = gmmBaseParams->N - tailN; + } + uint32_t curSingleM = mnConfig.singleM; + if (unlikely(mnConfig.mIdx == mnConfig.blockDimM - 1)) { + curSingleM = mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + } + uint64_t weightOffset = 0; + if constexpr (mmType::BT::format == CubeFormat::NZ) { + weightOffset = static_cast(groupIdx) * gmmBaseParams->N * gmmBaseParams->K + tailN * gmmBaseParams->K; + } else { + weightOffset = static_cast(groupIdx) * gmmBaseParams->N * gmmBaseParams->K + tailN; + } + mm.SetSingleShape(curSingleM, curSingleN, quantGroupSize); // 8, 256, 512 --> 514us + GlobalTensor weightSlice; + uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN; + mnConfig.workSpaceOffset = outOffset + mnConfig.yBaseOffset; + for (uint32_t loopK = 0; loopK < gmmBaseParams->quantGroupNum; loopK++) { + mm.SetTensorA( + xGM[mnConfig.xBaseOffset + mnConfig.mIdx * mnConfig.k * mnConfig.singleM + loopK * quantGroupSize]); + if constexpr (mmType::BT::format == CubeFormat::NZ) { + weightSlice = weightGM[weightOffset + loopK * quantGroupSize * 64]; + } else { + weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmBaseParams->N]; + } + if (mnConfig.blockDimM == 1) { + weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } + mm.SetTensorB(weightSlice); + mm.SetQuantVector(weightScaleGM[groupIdx * gmmBaseParams->N * gmmBaseParams->quantGroupNum + + loopK * gmmBaseParams->N + tailN]); + mm.Iterate(); + mm.GetTensorC(mmOutGM[mnConfig.workSpaceOffset], loopK == 0 ? 0 : 1); + } +} +} // namespace GROUPED_MATMUL_SWIGLU_QUANT +#endif // GMM_SWIGLU_QUANT_A8W4_MSD +#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_MID_H \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_a8w4_msd_post.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_a8w4_msd_post.h new file mode 100644 index 000000000..365e36d52 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_a8w4_msd_post.h @@ -0,0 +1,385 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_a8w4_msd_post.h + * \brief + */ + +#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_POST_H +#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_POST_H +#include "grouped_matmul_swiglu_quant_utils.h" +#include "kernel_operator.h" +#ifdef GMM_SWIGLU_QUANT_A8W4_MSD +namespace GROUPED_MATMUL_SWIGLU_QUANT { +using namespace AscendC; +#define DOUBLE_BUFFER 2 +constexpr float DEFAULT_MUL_SCALE = 16.0f; +class GMMA8W4PostProcess { +public: + __aicore__ inline GMMA8W4PostProcess(){}; + __aicore__ inline void Init(const GMAddrParams gmAddrParams, + const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN, + const GMMSwiglu *__restrict gmmSwigluIN); + + __aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx, + TPipe *pipe); + +private: + __aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx, + TPipe *pipe); + + __aicore__ inline void UpdateAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig); + + __aicore__ inline void VectorCompute(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void customDataCopyIn(uint32_t outLoopIdx, GlobalTensor &mmOutGM, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void customDataCopyOut(VecConfig &vecConfig, WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void PreLoadAuxiliaryMatrix(VecConfig &vecConfig); + + __aicore__ inline void Quant(uint32_t loopIdx, VecConfig &vecConfig); + + __aicore__ inline void Swiglu(uint32_t loopIdx, VecConfig &vecConfig); + + __aicore__ inline void MergeAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig); + + __aicore__ inline void MulPertokenScale(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig); + const GMMSwiglu *__restrict gmmSwiglu; + const GMMSwigluBaseParams *__restrict gmmBaseParams; + GlobalTensor perTokenScaleGM; + GlobalTensor groupListGM; + GlobalTensor quantOutputGM; + GlobalTensor weightAuxiliaryMatrixGM; + GlobalTensor quantScaleOutputGM; + GlobalTensor mmOutGM1; + GlobalTensor mmOutGM2; + GlobalTensor mmOutGM; + LocalTensor mmLocal_fp32; + LocalTensor mmLocal_fp16; + TQue weightAuxiliaryMatrixInQueue; + TQue mmOutQueue; + TQue quantOutQueue; + TQue quantScaleOutQueue; + TBuf reduceWorkspace; + uint32_t blockIdx = 0; + int64_t aicCoreNum = 0; + int64_t aivCoreNum = 0; +}; + +__aicore__ inline void GMMA8W4PostProcess::Init(const GMAddrParams gmAddrParams, + const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN, + const GMMSwiglu *__restrict gmmSwigluIN) +{ + if ASCEND_IS_AIV { + aicCoreNum = GetBlockNum(); + aivCoreNum = aicCoreNum * 2; + blockIdx = GetBlockIdx(); + gmmBaseParams = gmmSwigluBaseParamsIN; + gmmSwiglu = gmmSwigluIN; + weightAuxiliaryMatrixGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.weightAuxiliaryMatrixGM); // E, N + groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM, gmmSwiglu->groupListLen); + mmOutGM1.SetGlobalBuffer( + (__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset2)); + mmOutGM2.SetGlobalBuffer( + (__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset3)); + perTokenScaleGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.xScaleGM, gmmBaseParams->M); + quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.yGM, + gmmBaseParams->M * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR); + quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.yScaleGM, gmmBaseParams->M); + } +} + +__aicore__ inline void GMMA8W4PostProcess::customDataCopyIn(uint32_t outLoopIdx, GlobalTensor &mmOutGM, + VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + mmLocal_fp16 = mmOutQueue.DeQue(); + mmLocal_fp32 = mmLocal_fp16.ReinterpretCast(); + const int64_t processNum = 2 * vecConfig.innerLoopNum * gmmSwiglu->tokenLen; + DataCopyExtParams copyParams_0{1, static_cast(processNum * sizeof(half)), 0, 0, 0}; + DataCopyPadExtParams padParams_0{false, 0, 0, 0}; + DataCopyPad(mmLocal_fp16[processNum], mmOutGM[vecConfig.curOffset * DOUBLE_ROW], copyParams_0, padParams_0); + + mmOutQueue.EnQue(mmLocal_fp16); + mmLocal_fp16 = mmOutQueue.DeQue(); + // 1. fp16 -> fp32 + Cast(mmLocal_fp32, mmLocal_fp16[processNum], RoundMode::CAST_NONE, processNum); + PipeBarrier(); + int32_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + // 2. high_4bit * 16 + low_4bit + for (uint32_t i = 0; i < vecConfig.innerLoopNum; i++) { + Muls(mmLocal_fp32[(DOUBLE_ROW * i) * gmmSwiglu->tokenLen], mmLocal_fp32[(DOUBLE_ROW * i) * gmmSwiglu->tokenLen], + DEFAULT_MUL_SCALE, gmmSwiglu->tokenLen); + PipeBarrier(); + Add(mmLocal_fp32[i * gmmSwiglu->tokenLen], mmLocal_fp32[(DOUBLE_ROW * i) * gmmSwiglu->tokenLen], + mmLocal_fp32[(DOUBLE_ROW * i + 1) * gmmSwiglu->tokenLen], gmmSwiglu->tokenLen); + PipeBarrier(); + vecConfig.curIdx++; + } + vecConfig.curOffset = vecConfig.curIdx * gmmSwiglu->tokenLen; + PipeBarrier(); +} + +__aicore__ inline void GMMA8W4PostProcess::VectorCompute(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + // 1.辅助矩阵加回 + MergeAuxiliaryMatrix(loopIdx, vecConfig); + // 2.perToken反量化 + MulPertokenScale(loopIdx, vecConfig, workspaceSplitConfig); + // 3.Swiglu + Swiglu(loopIdx, vecConfig); + // 4.Quant + Quant(loopIdx, vecConfig); +} + +__aicore__ inline void GMMA8W4PostProcess::MergeAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig) +{ + // perChanelScale * perTokenScale + mmLocal_fp32 = mmOutQueue.DeQue(); + LocalTensor weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue(); + Add(mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], weightAuxiliaryMatrixLocal, + gmmSwiglu->tokenLen); + vecConfig.nextUpadteInterVal--; + PipeBarrier(); + weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal); +} + +__aicore__ inline void GMMA8W4PostProcess::MulPertokenScale(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + float scale = perTokenScaleGM.GetValue(loopIdx + workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx); + int32_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Muls(mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], scale, gmmSwiglu->tokenLen); + PipeBarrier(); +} + +__aicore__ inline void GMMA8W4PostProcess::Swiglu(uint32_t loopIdx, VecConfig &vecConfig) +{ + // 高阶API swiglu + float beta = 1.0f; + LocalTensor workspaceLocal = reduceWorkspace.Get(); + LocalTensor src0Local = + mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen + gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR]; + LocalTensor src1Local = mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen]; + + SwiGLU(workspaceLocal, src0Local, src1Local, beta, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR); + PipeBarrier(); + DataCopyParams repeatParams{1, static_cast((gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR) / ALIGN_8_ELE), 0, + 0}; + DataCopy(mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], workspaceLocal, repeatParams); + + PipeBarrier(); +} + +__aicore__ inline void GMMA8W4PostProcess::Quant(uint32_t loopIdx, VecConfig &vecConfig) +{ + uint64_t preOffset = loopIdx * gmmSwiglu->tokenLen; + uint64_t halfTokenLen = gmmSwiglu->tokenLen / BISECT; + Abs(mmLocal_fp32[preOffset + gmmSwiglu->tokenLen / BISECT], mmLocal_fp32[preOffset], halfTokenLen); + PipeBarrier(); + // reduceMax + LocalTensor workLocal = reduceWorkspace.Get(halfTokenLen); + LocalTensor reduceResLocal = + reduceWorkspace.GetWithOffset(FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float)); + LocalTensor reduceTmpLocal = reduceWorkspace.GetWithOffset( + FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float) + UB_BLOCK_UNIT_SIZE); + ReduceMaxTemplate(reduceResLocal, workLocal, mmLocal_fp32[preOffset + gmmSwiglu->tokenLen / BISECT], reduceTmpLocal, + static_cast(halfTokenLen)); + int32_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float quantScale = reduceResLocal.GetValue(0) / QUANT_SCALE_INT8; + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + quantScaleLocal.SetValue(loopIdx, quantScale); + quantScale = 1 / quantScale; + int32_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Muls(mmLocal_fp32[preOffset], mmLocal_fp32[preOffset], quantScale, halfTokenLen); + PipeBarrier(); + LocalTensor quantLocal = quantOutQueue.DeQue(); + int32_t dstTempOffset = static_cast(preOffset / BISECT); + int32_t srcTempOffset = static_cast(preOffset); + int32_t tempCount = static_cast(halfTokenLen); + LocalTensor castSpace = reduceWorkspace.Get(UB_BLOCK_UNIT_SIZE); + CastFp32ToInt8Template(quantLocal, mmLocal_fp32, castSpace, dstTempOffset, srcTempOffset, tempCount); + mmOutQueue.EnQue(mmLocal_fp32); + quantOutQueue.EnQue(quantLocal); +} + +__aicore__ inline void GMMA8W4PostProcess::UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig, + int64_t workspaceSplitLoopIdx, TPipe *pipe) +{ + // 第一步 读取grouplist reduceSum 计算总数据个数 + vecConfig.M = workspaceSplitLoopIdx < workspaceSplitConfig.loopCount - 1 ? workspaceSplitConfig.notLastTaskSize : + workspaceSplitConfig.lastLoopTaskSize; + // 第二步 计算分核 + uint32_t eachCoreTaskNum = (vecConfig.M + aivCoreNum - 1) / aivCoreNum; + vecConfig.usedCoreNum = vecConfig.M >= aivCoreNum ? aivCoreNum : vecConfig.M; + uint32_t tailCoreIdx = vecConfig.M - (eachCoreTaskNum - 1) * vecConfig.usedCoreNum; + vecConfig.taskNum = blockIdx < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1; + vecConfig.startIdx = + blockIdx < tailCoreIdx ? eachCoreTaskNum * blockIdx : ((eachCoreTaskNum - 1) * blockIdx + tailCoreIdx); + vecConfig.curIdx = vecConfig.startIdx; + vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen; + vecConfig.curOffset = vecConfig.startOffset; + int64_t curStartIdx = vecConfig.startIdx; + int64_t prevM = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize; + for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex; + groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; groupIdx++) { + int64_t currM = groupListGM.GetValue(groupIdx); + int64_t tempM = currM - prevM; + prevM = currM; + if (curStartIdx >= 0 && curStartIdx - tempM < 0) { + vecConfig.curGroupIdx = groupIdx; + vecConfig.nextUpadteInterVal = tempM - curStartIdx; + } + curStartIdx -= tempM; + } + // 第三步 计算总数据量 + vecConfig.outLoopNum = (vecConfig.taskNum + gmmSwiglu->maxProcessRowNum - 1) / gmmSwiglu->maxProcessRowNum; + vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwiglu->maxProcessRowNum ? + vecConfig.taskNum % gmmSwiglu->maxProcessRowNum : + gmmSwiglu->maxProcessRowNum; + + // 第四步 申请空间 + // 2 * row * n * sizeof(float) + row * n / 2 * sizeof(int8) + alignUp * sizeof(float) + n * sizeof(float) + + // n / 2 *sizeof(float) + 64 < 191 * 1024 + pipe->InitBuffer(mmOutQueue, 1, 2 * gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen * sizeof(float)); + pipe->InitBuffer(quantOutQueue, 1, + gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)); + pipe->InitBuffer(quantScaleOutQueue, 1, AlignUp(gmmSwiglu->maxProcessRowNum, ALIGN_8_ELE) * sizeof(float)); + pipe->InitBuffer(weightAuxiliaryMatrixInQueue, 1, gmmSwiglu->tokenLen * sizeof(float)); + // two 32 byte buffer for reduceMax calculation in Quant. + pipe->InitBuffer(reduceWorkspace, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(float) + UB_BLOCK_UNIT_SIZE + + UB_BLOCK_UNIT_SIZE); +} + +__aicore__ inline void GMMA8W4PostProcess::PreLoadAuxiliaryMatrix(VecConfig &vecConfig) +{ + LocalTensor weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue(); + DataCopyExtParams copyAuxiliaryMatrixParams{1, static_cast(gmmSwiglu->tokenLen * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(weightAuxiliaryMatrixLocal, weightAuxiliaryMatrixGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], + copyAuxiliaryMatrixParams, padParams); + weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal); +} + +__aicore__ inline void GMMA8W4PostProcess::UpdateAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig) +{ + // 更新weightAuxiliaryMatrix + if (unlikely(vecConfig.nextUpadteInterVal == 0)) { + int64_t loop = gmmSwiglu->groupListLen - vecConfig.curGroupIdx; + while (loop--) { + int64_t curTemp = groupListGM.GetValue(vecConfig.curGroupIdx); + vecConfig.curGroupIdx++; + int64_t nextTemp = groupListGM.GetValue(vecConfig.curGroupIdx); + if (nextTemp != curTemp) { + vecConfig.nextUpadteInterVal = nextTemp - curTemp; + break; + } + } + LocalTensor weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue(); + DataCopyExtParams copyParams{1, static_cast(gmmSwiglu->tokenLen * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(weightAuxiliaryMatrixLocal, weightAuxiliaryMatrixGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], + copyParams, padParams); + weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal); + } +} + +__aicore__ inline void GMMA8W4PostProcess::Process(WorkSpaceSplitConfig &workspaceSplitConfig, + int64_t workspaceSplitLoopIdx, TPipe *pipe) +{ + if ASCEND_IS_AIV { + if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount || workspaceSplitLoopIdx < 0) { + return; + } + VecConfig vecConfig; + UpdateVecConfig(blockIdx, vecConfig, workspaceSplitConfig, workspaceSplitLoopIdx, pipe); + + if (blockIdx < vecConfig.usedCoreNum) { + mmOutGM = (workspaceSplitLoopIdx % 2 == 0 ? mmOutGM1 : mmOutGM2); + LocalTensor weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.AllocTensor(); + LocalTensor mmLocal_fp32 = mmOutQueue.AllocTensor(); + LocalTensor quantScaleLocal = quantScaleOutQueue.AllocTensor(); + LocalTensor quantLocal = quantOutQueue.AllocTensor(); + + mmOutQueue.EnQue(mmLocal_fp32); + quantScaleOutQueue.EnQue(quantScaleLocal); + quantOutQueue.EnQue(quantLocal); + weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal); + PreLoadAuxiliaryMatrix(vecConfig); + for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) { + vecConfig.innerLoopNum = + outLoopIdx == (vecConfig.outLoopNum - 1) ? vecConfig.tailLoopNum : gmmSwiglu->maxProcessRowNum; + int32_t eventIdMTE3ToMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventIdMTE3ToMTE2); + WaitFlag(eventIdMTE3ToMTE2); + // 1.matmul中间结果搬入 + 高四位与低四位合并 + customDataCopyIn(outLoopIdx, mmOutGM, vecConfig, workspaceSplitConfig); + + for (uint32_t innerLoopIdx = 0; innerLoopIdx < vecConfig.innerLoopNum; innerLoopIdx++) { + // 2.如果涉及group切换，更新辅助矩阵 + UpdateAuxiliaryMatrix(innerLoopIdx, vecConfig); + // 3. 四步vector计算（辅助矩阵加回、perToken反量化、Swiglu、Quant） + VectorCompute(innerLoopIdx, vecConfig, workspaceSplitConfig); + } + int32_t eventIdVToMTE3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventIdVToMTE3); + WaitFlag(eventIdVToMTE3); + customDataCopyOut(vecConfig, workspaceSplitConfig); + } + weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue(); + mmLocal_fp32 = mmOutQueue.DeQue(); + quantScaleLocal = quantScaleOutQueue.DeQue(); + quantLocal = quantOutQueue.DeQue(); + + weightAuxiliaryMatrixInQueue.FreeTensor(weightAuxiliaryMatrixLocal); + mmOutQueue.FreeTensor(mmLocal_fp32); + quantScaleOutQueue.FreeTensor(quantScaleLocal); + quantOutQueue.FreeTensor(quantLocal); + } + } +} + +__aicore__ inline void GMMA8W4PostProcess::customDataCopyOut(VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + DataCopyParams copyParams_0{1, (uint16_t)(vecConfig.innerLoopNum * sizeof(float)), 0, 0}; + DataCopyPad(quantScaleOutputGM[workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx], quantScaleLocal, + copyParams_0); + LocalTensor quantLocal = quantOutQueue.DeQue(); + DataCopyParams copyParams_1{ + 1, (uint16_t)(vecConfig.innerLoopNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)), 0, 0}; + DataCopyPad(quantOutputGM[(workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx) * gmmSwiglu->tokenLen / + SWIGLU_REDUCE_FACTOR], + quantLocal, copyParams_1); + + vecConfig.startIdx += vecConfig.innerLoopNum; + vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen; + quantOutQueue.EnQue(quantLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); +} + +} // namespace GROUPED_MATMUL_SWIGLU_QUANT +#endif // GMM_SWIGLU_QUANT_A8W4_MSD +#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_AFTER_H \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_a8w4_msd_pre.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_a8w4_msd_pre.h new file mode 100644 index 000000000..bfe9f2a49 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_a8w4_msd_pre.h @@ -0,0 +1,219 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_a8w4_msd_pre.h + * \brief + */ + +#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_PRE_H +#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_PRE_H +#include "grouped_matmul_swiglu_quant_utils.h" +#include "kernel_operator.h" +#ifdef GMM_SWIGLU_QUANT_A8W4_MSD +namespace GROUPED_MATMUL_SWIGLU_QUANT { +using namespace AscendC; +#define BUFFER_NUM_A8W4_PRE 1 +constexpr int TWO = 2; +constexpr int EIGHT = 8; +constexpr size_t LEN_128 = 128; // 16bit operator +constexpr int DATA_BLOCK_SIZE_32 = 32; +class GMMA8W4PreProcess { +public: + __aicore__ inline GMMA8W4PreProcess(){}; + __aicore__ inline void Init(const GMAddrParams gmAddrParams, + const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN); + __aicore__ inline void CalculateTaskInfoEachCore(uint32_t &curCoreTaskNum_, uint32_t &curCoreStartOffset_); + __aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx, + TPipe *pipe); + __aicore__ inline void CustomInitBuffer(TPipe *pipe); + +private: + TQue vecInQueueX, vecInQueueXBak; + TQue vecOutQueueA1; + TQue vecOutQueueA2; + TQue vecOutQueueA3; + TQue vecOutQueue0F; + TQue vecOutQueueRowSum; + TBuf tempBuff; + const GMMSwigluBaseParams *__restrict gmmSwigluBaseParams; + LocalTensor xTensor; + LocalTensor xHighHalfTensor; + LocalTensor xHighFloatTensor; + LocalTensor xLowHalfTensor; + LocalTensor xLowHalfTensor2; + LocalTensor xHighI4Tensor; + LocalTensor xLowI4Tensor; + LocalTensor xLowI16Tensor; + LocalTensor groupListTensor; + LocalTensor groupListFTensor; + LocalTensor xRowSumTensor; + + GlobalTensor xGM; + GlobalTensor yGm; + GlobalTensor yGm1; + GlobalTensor yGm2; + GlobalTensor groupListGM; + + uint32_t vK{0}; + uint32_t vKAlign{0}; + uint32_t totalM{0}; + uint32_t blockDim{0}; + uint32_t curCoreId{0}; + uint32_t curCoreTaskNum{0}; + uint32_t curCoreStartOffset{0}; + uint32_t curCoreOuterLoopNum{0}; + uint32_t curCoreInnerTailLoopNum{0}; + uint32_t groupNum{0}; +}; + +__aicore__ inline void GMMA8W4PreProcess::Init(const GMAddrParams gmAddrParams, + const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN) +{ + if ASCEND_IS_AIV { + xGM.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.xGM); + yGm1.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.workSpaceGM); + yGm2.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset1); + groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM); + gmmSwigluBaseParams = gmmSwigluBaseParamsIN; + vK = gmmSwigluBaseParams->K; + groupNum = static_cast(gmmSwigluBaseParams->groupNum); + // M * K * 7B (1B + 0.5B + 0.5B + 2B + 4B) <= UBsize - 256B + blockDim = GetBlockNum() * GetTaskRation(); + } +} + +__aicore__ inline void GMMA8W4PreProcess::CustomInitBuffer(TPipe *pipe) +{ + pipe->InitBuffer(vecInQueueX, BUFFER_NUM_A8W4_PRE, vK * sizeof(int8_t)); // K * 1B + pipe->InitBuffer(vecOutQueueA1, BUFFER_NUM_A8W4_PRE, vK * sizeof(int4b_t)); // K * 0.5B + pipe->InitBuffer(vecOutQueueA2, BUFFER_NUM_A8W4_PRE, vK * sizeof(int4b_t)); // K * 0.5B + pipe->InitBuffer(vecOutQueueA3, BUFFER_NUM_A8W4_PRE, vK * sizeof(half)); // K * 2B + // xLowHalfTensor, xLowHalfTensor2 and xHighFloatTensor share the same buffer + pipe->InitBuffer(tempBuff, vK * sizeof(float)); // K * 4B + constexpr int BUFFER_SIZE_256B = 128 * sizeof(int16_t); + pipe->InitBuffer(vecOutQueue0F, BUFFER_NUM_A8W4_PRE, BUFFER_SIZE_256B); // 256B +} + + +__aicore__ inline void GMMA8W4PreProcess::CalculateTaskInfoEachCore(uint32_t &curCoreTaskNum_, + uint32_t &curCoreStartOffset_) +{ + // 均分任务数 + int64_t eachCoreTaskNum = (totalM + blockDim - 1) / blockDim; // 每个核处理的数据量 + // 尾核任务数 + int64_t taskNumPertailCore = eachCoreTaskNum - 1; + // 实际使用核数 + int64_t usedCoreNum = totalM >= blockDim ? blockDim : totalM; + // 尾核起始索引 + uint32_t tailCoreIdx = totalM - (eachCoreTaskNum - 1) * usedCoreNum; + curCoreId = GetBlockIdx(); + // 每个核处理的任务数量 = 是否为尾核 ？均分任务数 ：(均分任务数 - 1) + curCoreTaskNum_ = curCoreId < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1; + // 每个核处理的起始偏移地址 = 是否为尾核 ？均分任务数 * blockId : (均分任务数 - 1) * blockId + 尾核起始索引 + curCoreStartOffset_ = + curCoreId < tailCoreIdx ? eachCoreTaskNum * curCoreId : ((eachCoreTaskNum - 1) * curCoreId + tailCoreIdx); +} + +__aicore__ inline void GMMA8W4PreProcess::Process(WorkSpaceSplitConfig &workspaceSplitConfig, + int64_t workspaceSplitLoopIdx, TPipe *pipe) +{ + if ASCEND_IS_AIV { + if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount) { + return; + } + yGm = (workspaceSplitLoopIdx % 2 == 0 ? yGm1 : yGm2); + CustomInitBuffer(pipe); + constexpr int32_t MASK = 128; + xTensor = vecInQueueX.AllocTensor(); + xHighI4Tensor = vecOutQueueA1.AllocTensor(); + xLowI4Tensor = vecOutQueueA2.AllocTensor(); + xHighHalfTensor = vecOutQueueA3.AllocTensor(); + const uint32_t xLowHalfOffset = vK * sizeof(half); + xLowHalfTensor = tempBuff.GetWithOffset(xLowHalfOffset, 0); + xLowHalfTensor2 = tempBuff.GetWithOffset(xLowHalfOffset, xLowHalfOffset); + xLowI16Tensor = vecOutQueue0F.AllocTensor(); + + Duplicate(xLowI16Tensor, static_cast(0x0F0F), MASK); // get rid of high 4 bits in every int8 + PipeBarrier(); + const size_t LEN_VK = (vK / 2) / 128; + const size_t LAST_LEN_VK = (vK % 256) / 2; + const half ONE_SIXTEENTH = static_cast(0.0625f); + // groupList仅支持count + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + totalM = workspaceSplitLoopIdx < workspaceSplitConfig.loopCount - 1 ? workspaceSplitConfig.notLastTaskSize : + workspaceSplitConfig.lastLoopTaskSize; + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + CalculateTaskInfoEachCore(curCoreTaskNum, curCoreStartOffset); + SetFlag(EVENT_ID0); // 0 + SetFlag(EVENT_ID0); // 1 + SetFlag(EVENT_ID1); // 2 + + for (uint32_t xloop = 0; xloop < curCoreTaskNum; xloop++) { + uint64_t relStartAddr = (xloop + curCoreStartOffset) * vK; + uint64_t absStartAddr = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize * vK + relStartAddr; + // 高四位处理开始 + WaitFlag(EVENT_ID0); // 0 + DataCopy(xTensor, xGM[absStartAddr], vK); + SetFlag(EVENT_ID0); // 3 + WaitFlag(EVENT_ID0); // 3 + Cast(xHighHalfTensor, xTensor, AscendC::RoundMode::CAST_NONE, vK); + PipeBarrier(); + Muls(xHighHalfTensor, xHighHalfTensor, ONE_SIXTEENTH, vK); + PipeBarrier(); + WaitFlag(EVENT_ID1); // 2 + Cast(xHighI4Tensor, xHighHalfTensor, AscendC::RoundMode::CAST_FLOOR, vK); + SetFlag(EVENT_ID0); // 4 + WaitFlag(EVENT_ID0); // 4 + DataCopy(yGm[relStartAddr], xHighI4Tensor.ReinterpretCast(), vK / 2); + // 高四位处理结束 + + // 低四位处理开始 + SetFlag(EVENT_ID1); // 2 + And(xLowHalfTensor.ReinterpretCast(), xTensor.ReinterpretCast(), xLowI16Tensor, LEN_128, + LEN_VK, {1, 1, 1, 8, 8, 0}); + if (LAST_LEN_VK > 0) { + And(xLowHalfTensor[LEN_VK * LEN_128].ReinterpretCast(), + xTensor[LEN_VK * LEN_128 * TWO].ReinterpretCast(), xLowI16Tensor, LAST_LEN_VK, 1, + {1, 1, 1, 8, 8, 0}); + } + PipeBarrier(); + SetFlag(EVENT_ID0); // 0 + Cast(xLowHalfTensor2.ReinterpretCast(), xLowHalfTensor.ReinterpretCast(), + AscendC::RoundMode::CAST_NONE, vK); + PipeBarrier(); + const half MINUS_EIGHT = static_cast(-8); + Adds(xHighHalfTensor, xLowHalfTensor2, MINUS_EIGHT, vK); + PipeBarrier(); + WaitFlag(EVENT_ID0); // 1 + Cast(xLowI4Tensor, xHighHalfTensor.ReinterpretCast(), AscendC::RoundMode::CAST_NONE, vK); + SetFlag(EVENT_ID1); // 5 + WaitFlag(EVENT_ID1); // 5 + DataCopy(yGm[relStartAddr + vK / TWO], xLowI4Tensor.ReinterpretCast(), vK / TWO); + SetFlag(EVENT_ID0); // 1 + // 低四位处理结束 + } + + WaitFlag(EVENT_ID0); // 0 + WaitFlag(EVENT_ID0); // 1 + WaitFlag(EVENT_ID1); // 2 + vecInQueueX.FreeTensor(xTensor); + vecOutQueueA1.FreeTensor(xHighI4Tensor); + vecOutQueueA2.FreeTensor(xLowI4Tensor); + vecOutQueueA3.FreeTensor(xHighHalfTensor); + vecOutQueue0F.FreeTensor(xLowI16Tensor); + } +} + +} // namespace GROUPED_MATMUL_SWIGLU_QUANT +#endif // GMM_SWIGLU_QUANT_A8W4_MSD +#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_PRE_H \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_split_ws.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_split_ws.h new file mode 100644 index 000000000..46d5ca61a --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_split_ws.h @@ -0,0 +1,635 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_split_ws.h + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_SPLIT_WS_H +#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_SPLIT_WS_H + +#include "grouped_matmul_swiglu_quant_utils.h" +namespace GROUPED_MATMUL_SWIGLU_QUANT { +/** @brief internal computation class + */ + +template +class GMMSwigluSplitWorkSpaceCompute { +public: + using AT = typename mmType::AT::T; + using BT = typename mmType::BT::T; + using B = typename mmType::BT; + using CT = typename mmType::CT::T; + using BiasT = typename mmType::BiasT::T; + using WT = int8_t; + constexpr static bool transposeX = mmType::AT::isTrans; + constexpr static bool transposeW = mmType::BT::isTrans; + static constexpr float FLOAT_INF = 3e+99; + /** @brief constructor */ + __aicore__ inline GMMSwigluSplitWorkSpaceCompute(typename mmType::MT &mm_) : mm(mm_) + { + } + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale, + GM_ADDR groupList, GM_ADDR quantOutput, GM_ADDR quantScaleOutput, GM_ADDR workspace, + const GMMSwigluBaseParams *__restrict gmmBaseParamsIN, + const TCubeTiling *__restrict mmTilingDataIN, const GMMSwiglu *__restrict gmmSwigluIN, + TPipe *tPipeIN); + __aicore__ inline void Process(); + +private: + __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig &mnConfig, uint32_t coreIdx, + GlobalTensor &mmOutGM); + + __aicore__ inline void UpdateMnConfig(MNConfig &mnConfig); + + __aicore__ inline void SetMNConfig(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig); + + __aicore__ inline void SetMKN(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig); + + __aicore__ inline uint64_t GetWOffset(uint32_t tailN, uint32_t k); + + __aicore__ inline void MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock, const uint32_t count, + const uint32_t thresholdM_dimN); + + template + __aicore__ inline void UpdateChannelScale(uint32_t loopidx, VecConfig &vecConfig); + + __aicore__ inline void VectorCompute(uint32_t loopidx, VecConfig &vecConfig); + + template + __aicore__ inline void PreLoadTokenAndChannel(LocalTensor &channelScaleLocal, VecConfig &vecConfig); + + __aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig); + + __aicore__ inline void UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig, + int32_t workspaceSplitLoopIdx); + + __aicore__ inline void InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void customDataCopyIn(uint32_t outLoopIdx, GlobalTensor &mmOutGM, VecConfig &vecConfig); + + __aicore__ inline void customDataCopyOut(VecConfig &vecConfig); + + __aicore__ inline void Dequant(uint32_t loopidx, VecConfig &vecConfig); + + __aicore__ inline void Quant(uint32_t loopidx, VecConfig &vecConfig); + + __aicore__ inline void Swiglu(uint32_t loopidx, VecConfig &vecConfig); + +private: + typename mmType::MT &mm; + const GMMSwigluBaseParams *__restrict gmmBaseParams; + const GMMSwiglu *__restrict gmmSwiglu; + const TCubeTiling *__restrict mmTilingData; + WorkSpaceSplitConfig workspaceSplitConfig; + TPipe *pipe; + GlobalTensor xGM; + GlobalTensor weightGM; + GlobalTensor perChannelScaleGM; + GlobalTensor perTokenScaleGM; + GlobalTensor groupListGM; + GlobalTensor quantOutputGM; + GlobalTensor quantScaleOutputGM; + GlobalTensor mmOutGM1; + GlobalTensor mmOutGM2; + // define the que + TQue mmOutQueue; + TQue perChannelScaleInQueue; + TQue quantOutQueue; + TQue quantScaleOutQueue; + TBuf reduceWorkspace; + uint32_t blockIdx = 0; + int64_t aicCoreNum = 0; + int64_t aivCoreNum = 0; + GM_ADDR xTensorPtr; + GM_ADDR weightTensorPtr; + float limited = FLOAT_INF; +}; + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::Init( + GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale, GM_ADDR groupList, GM_ADDR quantOutput, + GM_ADDR quantScaleOutput, GM_ADDR workspace, const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIn, + const TCubeTiling *__restrict mmTilingDataIN, const GMMSwiglu *__restrict gmmSwigluIN, TPipe *tPipeIN) +{ + aicCoreNum = GetBlockNum(); + aivCoreNum = aicCoreNum * 2; + blockIdx = GetBlockIdx(); + pipe = tPipeIN; + xTensorPtr = x; + weightTensorPtr = weight; + mmTilingData = mmTilingDataIN; + gmmBaseParams = gmmSwigluBaseParamsIn; + gmmSwiglu = gmmSwigluIN; + limited = gmmBaseParams->limited; + groupListGM.SetGlobalBuffer((__gm__ int64_t *)groupList, gmmSwiglu->groupListLen); + mmOutGM1.SetGlobalBuffer((__gm__ int32_t *)workspace, gmmBaseParams->mLimit * gmmSwiglu->tokenLen); + mmOutGM2.SetGlobalBuffer((__gm__ int32_t *)workspace + gmmBaseParams->mLimit * gmmSwiglu->tokenLen, + gmmBaseParams->mLimit * gmmSwiglu->tokenLen); + if ASCEND_IS_AIV { + perChannelScaleGM.SetGlobalBuffer((__gm__ CHANNELDTYPE *)perChannelScale, + gmmSwiglu->groupListLen * gmmSwiglu->tokenLen); + perTokenScaleGM.SetGlobalBuffer((__gm__ float *)perTokenScale, gmmBaseParams->M); + quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)quantOutput, + gmmBaseParams->M * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR); + quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)quantScaleOutput, gmmBaseParams->M); + } +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::InitWorkSpaceSplitConfig( + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + workspaceSplitConfig.M = groupListGM.GetValue(gmmSwiglu->groupListLen - 1); + workspaceSplitConfig.loopCount = Ceil(workspaceSplitConfig.M, gmmBaseParams->mLimit); + workspaceSplitConfig.notLastTaskSize = gmmBaseParams->mLimit; + workspaceSplitConfig.lastLoopTaskSize = + workspaceSplitConfig.M - (workspaceSplitConfig.loopCount - 1) * gmmBaseParams->mLimit; + workspaceSplitConfig.leftMatrixStartIndex = 0; + workspaceSplitConfig.rightMatrixExpertStartIndex = 0; + workspaceSplitConfig.rightMatrixExpertNextStartIndex = 0; + workspaceSplitConfig.isLastLoop = false; +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::UpdateWorkSpaceSplitConfig( + WorkSpaceSplitConfig &workspaceSplitConfig, int32_t workspaceSplitLoopIdx) +{ + workspaceSplitConfig.leftMatrixStartIndex = workspaceSplitLoopIdx * gmmBaseParams->mLimit; + workspaceSplitConfig.rightMatrixExpertStartIndex = workspaceSplitConfig.rightMatrixExpertNextStartIndex; + workspaceSplitConfig.rightMatrixExpertEndIndex = workspaceSplitConfig.rightMatrixExpertStartIndex; + // 计算右专家矩阵的终止索引(rightMatrixExpertEndIndex) 和下一次的起始索引(rightMatrixExpertNextStartIndex) + int32_t curTaskNum = 0; + int32_t nextTaskNum = 0; + while (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwiglu->groupListLen) { + curTaskNum = groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex) - + workspaceSplitConfig.leftMatrixStartIndex; + int32_t nextTaskIdx = workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwiglu->groupListLen - 1 ? + gmmSwiglu->groupListLen - 1 : + workspaceSplitConfig.rightMatrixExpertEndIndex + 1; + nextTaskNum = groupListGM.GetValue(nextTaskIdx) - workspaceSplitConfig.leftMatrixStartIndex; + if (curTaskNum > gmmBaseParams->mLimit) { + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex; + break; + } else if (curTaskNum == gmmBaseParams->mLimit && nextTaskNum > gmmBaseParams->mLimit) { + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex + 1; + break; + } else if (nextTaskNum > gmmBaseParams->mLimit) { + workspaceSplitConfig.rightMatrixExpertEndIndex++; + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex; + break; + } + workspaceSplitConfig.rightMatrixExpertEndIndex++; + } + workspaceSplitConfig.isLastLoop = workspaceSplitLoopIdx == workspaceSplitConfig.loopCount - 1 ? true : false; + + if (workspaceSplitConfig.isLastLoop) { + workspaceSplitConfig.rightMatrixExpertEndIndex = + workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwiglu->groupListLen ? + gmmSwiglu->groupListLen - 1 : + workspaceSplitConfig.rightMatrixExpertEndIndex; + } +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::Process() +{ + InitWorkSpaceSplitConfig(workspaceSplitConfig); + int32_t parallelNum = 2; // 2: double workspace buffer + for (int32_t workspaceSplitLoopIdx = 0; workspaceSplitLoopIdx < workspaceSplitConfig.loopCount; + workspaceSplitLoopIdx++) { + UpdateWorkSpaceSplitConfig(workspaceSplitConfig, workspaceSplitLoopIdx); + GlobalTensor mmOutGM = (workspaceSplitLoopIdx % parallelNum == 0) ? mmOutGM1 : mmOutGM2; + + if ASCEND_IS_AIC { + if (workspaceSplitLoopIdx >= parallelNum) { // first parallelNum core no need to wait + SyncAll(); + } + MNConfig mnConfig; + int32_t prevSplitValue = workspaceSplitConfig.leftMatrixStartIndex; + for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex, count = 0; + groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; ++groupIdx) { + UpdateMnConfig(mnConfig); + int32_t currSplitValue = static_cast(groupListGM.GetValue(groupIdx)); + currSplitValue = currSplitValue > (workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit ? + (workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit : + currSplitValue; + int32_t splitValue = currSplitValue - prevSplitValue; + prevSplitValue = currSplitValue; + SetMNConfig(splitValue, groupIdx, mnConfig); + if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) { + continue; + } + mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM); + mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN); + + uint32_t curCount = count + mnConfig.blockDimM * mnConfig.blockDimN; + uint32_t curBlock = blockIdx >= count ? blockIdx : blockIdx + gmmBaseParams->coreNum; + uint32_t thresholdM_dimN = THRESHOLD_BLOCK_NUM * mnConfig.blockDimN; + + while (curBlock < curCount) { + MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN); + MMCompute(groupIdx, mnConfig, blockIdx, mmOutGM); + curBlock += aicCoreNum; + } + count = curCount % gmmBaseParams->coreNum; + } + SyncAll(); + } + + if ASCEND_IS_AIV { + VecConfig vecConfig; + UpdateVecConfig(blockIdx, vecConfig); + if (blockIdx < vecConfig.usedCoreNum) { + LocalTensor channelScaleLocal = perChannelScaleInQueue.AllocTensor(); + LocalTensor mmLocal = mmOutQueue.AllocTensor(); + LocalTensor quantLocal = quantOutQueue.AllocTensor(); + LocalTensor quantScaleLocal = quantScaleOutQueue.AllocTensor(); + mmOutQueue.EnQue(mmLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); + quantOutQueue.EnQue(quantLocal); + PreLoadTokenAndChannel(channelScaleLocal, vecConfig); + } + SyncAll(); + if (blockIdx < vecConfig.usedCoreNum) { + for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) { + vecConfig.innerLoopNum = + outLoopIdx == (vecConfig.outLoopNum - 1) ? vecConfig.tailLoopNum : gmmSwiglu->maxProcessRowNum; + int32_t eventIdMTE3ToMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventIdMTE3ToMTE2); + WaitFlag(eventIdMTE3ToMTE2); + customDataCopyIn(outLoopIdx, mmOutGM, vecConfig); + for (uint32_t innerLoopIdx = 0; innerLoopIdx < vecConfig.innerLoopNum; innerLoopIdx++) { + UpdateChannelScale(innerLoopIdx, vecConfig); + VectorCompute(innerLoopIdx, vecConfig); + } + int32_t eventIdVToMTE3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventIdVToMTE3); + WaitFlag(eventIdVToMTE3); + customDataCopyOut(vecConfig); + } + + LocalTensor channelScaleLocal = perChannelScaleInQueue.DeQue(); + LocalTensor mmLocal = mmOutQueue.DeQue(); + LocalTensor quantLocal = quantOutQueue.DeQue(); + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + perChannelScaleInQueue.FreeTensor(channelScaleLocal); + mmOutQueue.FreeTensor(mmLocal); + quantScaleOutQueue.FreeTensor(quantScaleLocal); + quantOutQueue.FreeTensor(quantLocal); + } + if (workspaceSplitLoopIdx < workspaceSplitConfig.loopCount - parallelNum) { + SyncAll(); + } + } + } +} + +template +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::PreLoadTokenAndChannel( + LocalTensor &channelScaleLocal, VecConfig &vecConfig) +{ + DataCopyExtParams copyChannelParams{1, static_cast(gmmSwiglu->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + if constexpr (!IsSameType::value) { + LocalTensor dstLocalT = channelScaleLocal.template ReinterpretCast(); + DataCopyPad(dstLocalT[gmmSwiglu->tokenLen], perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], + copyChannelParams, padParams); + int32_t eventIdMTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMTE2ToV); + WaitFlag(eventIdMTE2ToV); + Cast(channelScaleLocal, dstLocalT[gmmSwiglu->tokenLen], RoundMode::CAST_NONE, gmmSwiglu->tokenLen); + } else { + DataCopyPad(channelScaleLocal, perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], + copyChannelParams, padParams); + } + perChannelScaleInQueue.EnQue(channelScaleLocal); +} + +template +__aicore__ inline void +GMMSwigluSplitWorkSpaceCompute::MMCompute(uint32_t groupIdx, MNConfig &mnConfig, + uint32_t coreIdx, GlobalTensor &mmOutGM) +{ + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN; + uint32_t curSingleM = + mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM : mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + uint64_t xOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.k; + if constexpr (transposeX) { + xOffset = mnConfig.mIdx * mnConfig.singleM; + } + uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN; + xGM.SetGlobalBuffer((__gm__ int8_t *)xTensorPtr + mnConfig.xBaseOffset + + workspaceSplitConfig.leftMatrixStartIndex * mnConfig.k); + weightGM.SetGlobalBuffer((__gm__ int8_t *)weightTensorPtr + groupIdx * mnConfig.k * mnConfig.n + + GetWOffset(tailN, mnConfig.k)); + if (mnConfig.blockDimM == 1) { + weightGM.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } else { + weightGM.SetL2CacheHint(CacheMode::CACHE_MODE_NORMAL); + } + mnConfig.workSpaceOffset = outOffset + mnConfig.yBaseOffset; + mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k); + mm.SetSingleShape(curSingleM, curSingleN, mnConfig.k); + mm.SetTensorA(xGM[xOffset], transposeX); + mm.SetTensorB(weightGM, transposeW); + mm.template IterateAll(mmOutGM[mnConfig.workSpaceOffset], 0); +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::UpdateMnConfig(MNConfig &mnConfig) +{ + if constexpr (B::format == CubeFormat::NZ) { + mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<32>(mnConfig.n); // 16: nz format last two dim size + } else { + mnConfig.wBaseOffset += mnConfig.k * mnConfig.n; + } + mnConfig.nAxisBaseOffset += mnConfig.n; + mnConfig.mAxisBaseOffset += mnConfig.m; + mnConfig.xBaseOffset += mnConfig.m * mnConfig.k; + mnConfig.yBaseOffset += mnConfig.m * mnConfig.n; +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::SetMNConfig(const int32_t splitValue, + const uint32_t groupIdx, + MNConfig &mnConfig) +{ + SetMKN(splitValue, groupIdx, mnConfig); + mnConfig.baseM = BASIC_M; + mnConfig.baseN = BASIC_N; + mnConfig.singleM = SINGLE_CORE_M; + mnConfig.singleN = SINGLE_CORE_N; +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::SetMKN(const int32_t splitValue, + const uint32_t groupIdx, + MNConfig &mnConfig) +{ + mnConfig.m = static_cast(splitValue); + mnConfig.k = gmmBaseParams->K; // tilingData + mnConfig.n = gmmBaseParams->N; // tilingData +} + +template +__aicore__ inline uint64_t GMMSwigluSplitWorkSpaceCompute::GetWOffset(uint32_t tailN, + uint32_t k) +{ + uint64_t wOffset = 0; + if constexpr (mmType::BT::format == CubeFormat::NZ) { + wOffset = tailN * AlignUp<16>(k); // 16: nz format last two dim size + } else { + wOffset = tailN; + } + return wOffset; +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::MNBlockIdxCompute( + MNConfig &mnConfig, const uint32_t curBlock, const uint32_t count, const uint32_t thresholdM_dimN) +{ + mnConfig.mIdx = (curBlock - count) / mnConfig.blockDimN; + mnConfig.nIdx = (curBlock - count) % mnConfig.blockDimN; +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::UpdateVecConfig(uint32_t blockIdx, + VecConfig &vecConfig) +{ + // 第一步 读取grouplist reduceSum 计算总数据个数 + vecConfig.M = + workspaceSplitConfig.isLastLoop ? workspaceSplitConfig.lastLoopTaskSize : workspaceSplitConfig.notLastTaskSize; + // 第二步 计算分核 + uint32_t eachCoreTaskNum = (vecConfig.M + aivCoreNum - 1) / aivCoreNum; + vecConfig.usedCoreNum = vecConfig.M >= aivCoreNum ? aivCoreNum : vecConfig.M; + uint32_t tailCoreIdx = vecConfig.M - (eachCoreTaskNum - 1) * vecConfig.usedCoreNum; + vecConfig.taskNum = blockIdx < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1; + vecConfig.startIdx = + blockIdx < tailCoreIdx ? eachCoreTaskNum * blockIdx : ((eachCoreTaskNum - 1) * blockIdx + tailCoreIdx); + vecConfig.curIdx = vecConfig.startIdx; + vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen; + vecConfig.curOffset = vecConfig.startOffset; + int64_t curStartIdx = vecConfig.startIdx; + int64_t prevM = workspaceSplitConfig.leftMatrixStartIndex; + for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex; + groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; groupIdx++) { + int64_t currM = groupListGM.GetValue(groupIdx); + int64_t tempM = currM - prevM; + prevM = currM; + if (curStartIdx >= 0 && curStartIdx - tempM < 0) { + vecConfig.curGroupIdx = groupIdx; + vecConfig.nextUpadteInterVal = tempM - curStartIdx; + } + curStartIdx -= tempM; + } + // 第三步 计算总数据量 + vecConfig.outLoopNum = (vecConfig.taskNum + gmmSwiglu->maxProcessRowNum - 1) / gmmSwiglu->maxProcessRowNum; + vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwiglu->maxProcessRowNum ? + vecConfig.taskNum % gmmSwiglu->maxProcessRowNum : + gmmSwiglu->maxProcessRowNum; + pipe->Reset(); + // 第四步 申请空间 + pipe->InitBuffer(mmOutQueue, DOUBLE_BUFFER, gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen * sizeof(int32_t)); + pipe->InitBuffer(perChannelScaleInQueue, DOUBLE_BUFFER, gmmSwiglu->tokenLen * sizeof(float)); + pipe->InitBuffer(quantOutQueue, DOUBLE_BUFFER, + gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)); + pipe->InitBuffer(quantScaleOutQueue, DOUBLE_BUFFER, + AlignUp(gmmSwiglu->maxProcessRowNum, ALIGN_8_ELE) * sizeof(float)); + // two 32 byte buffer for reduceMax calculation in Quant. + pipe->InitBuffer(reduceWorkspace, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(float) + UB_BLOCK_UNIT_SIZE + + UB_BLOCK_UNIT_SIZE); +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::customDataCopyIn( + uint32_t outLoopIdx, GlobalTensor &mmOutGM, VecConfig &vecConfig) +{ + LocalTensor _inMMLocal_0 = mmOutQueue.DeQue(); + DataCopyExtParams copyParams_0{ + 1, static_cast(vecConfig.innerLoopNum * gmmSwiglu->tokenLen * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams padParams_0{false, 0, 0, 0}; + DataCopyPad(_inMMLocal_0, mmOutGM[vecConfig.curOffset], copyParams_0, padParams_0); + mmOutQueue.EnQue(_inMMLocal_0); + + LocalTensor _inMMLocal_1 = mmOutQueue.DeQue(); + + Cast(_inMMLocal_1.ReinterpretCast(), _inMMLocal_1, RoundMode::CAST_NONE, + vecConfig.innerLoopNum * gmmSwiglu->tokenLen); + + mmOutQueue.EnQue(_inMMLocal_1); + LocalTensor _inMMLocal_2 = mmOutQueue.DeQue(); + int32_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + for (uint32_t i = 0; i < vecConfig.innerLoopNum; i++) { + WaitFlag(eventIdSToV); + float scale = perTokenScaleGM.GetValue(vecConfig.curIdx + workspaceSplitConfig.leftMatrixStartIndex); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Muls(_inMMLocal_2[i * gmmSwiglu->tokenLen], _inMMLocal_2[i * gmmSwiglu->tokenLen], scale, gmmSwiglu->tokenLen); + SetFlag(eventIdSToV); + vecConfig.curIdx++; + } + WaitFlag(eventIdSToV); + vecConfig.curOffset = vecConfig.curIdx * gmmSwiglu->tokenLen; + mmOutQueue.EnQue(_inMMLocal_2); +} + +template +template +__aicore__ inline void +GMMSwigluSplitWorkSpaceCompute::UpdateChannelScale(uint32_t loopIdx, VecConfig &vecConfig) +{ + // 更新perChannel + if (unlikely(vecConfig.nextUpadteInterVal == 0)) { + int64_t loop = gmmSwiglu->groupListLen - vecConfig.curGroupIdx; + while (loop--) { + int64_t curTemp = groupListGM.GetValue(vecConfig.curGroupIdx); + vecConfig.curGroupIdx++; + int64_t nextTemp = groupListGM.GetValue(vecConfig.curGroupIdx); + if (nextTemp != curTemp) { + vecConfig.nextUpadteInterVal = nextTemp - curTemp; + break; + } + } + LocalTensor _inChannel = perChannelScaleInQueue.DeQue(); + DataCopyExtParams copyParams{1, static_cast(gmmSwiglu->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + if constexpr (!IsSameType::value) { + LocalTensor dstLocalT = _inChannel.template ReinterpretCast(); + DataCopyPad(dstLocalT[gmmSwiglu->tokenLen], perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], + copyParams, padParams); + int32_t eventIdMTE2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMTE2ToV); + WaitFlag(eventIdMTE2ToV); + Cast(_inChannel, dstLocalT[gmmSwiglu->tokenLen], RoundMode::CAST_NONE, gmmSwiglu->tokenLen); + } else { + DataCopyPad(_inChannel, perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], copyParams, + padParams); + } + perChannelScaleInQueue.EnQue(_inChannel); + } +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::VectorCompute(uint32_t loopIdx, + VecConfig &vecConfig) +{ + Dequant(loopIdx, vecConfig); + Swiglu(loopIdx, vecConfig); + Quant(loopIdx, vecConfig); +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::Dequant(uint32_t loopIdx, + VecConfig &vecConfig) +{ + // perChanelScale * perTokenScale + LocalTensor mmLocal = mmOutQueue.DeQue(); + LocalTensor perChannelLocal = perChannelScaleInQueue.DeQue(); + Mul(mmLocal[loopIdx * gmmSwiglu->tokenLen], mmLocal[loopIdx * gmmSwiglu->tokenLen], perChannelLocal, + gmmSwiglu->tokenLen); + vecConfig.nextUpadteInterVal--; + mmOutQueue.EnQue(mmLocal); + perChannelScaleInQueue.EnQue(perChannelLocal); +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::Swiglu(uint32_t loopIdx, + VecConfig &vecConfig) +{ + // 高阶API swiglu + LocalTensor _inMMLocal = mmOutQueue.DeQue(); + float beta = 1.0f; + LocalTensor workspaceLocal = reduceWorkspace.Get(); + LocalTensor src0Local = + _inMMLocal[loopIdx * gmmSwiglu->tokenLen + gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR]; + LocalTensor src1Local = _inMMLocal[loopIdx * gmmSwiglu->tokenLen]; + if (limited > 0.0f) { + Mins(src0Local, src0Local, limited, gmmSwiglu->tokenLen / 2); + PipeBarrier(); + Maxs(src0Local, src0Local, (-1.0f * limited), gmmSwiglu->tokenLen / 2); + PipeBarrier(); + Mins(src1Local, src1Local, limited, gmmSwiglu->tokenLen / 2); + PipeBarrier(); + } + SwiGLU(workspaceLocal, src0Local, src1Local, beta, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR); + PipeBarrier(); + DataCopyParams repeatParams{1, static_cast((gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR) / ALIGN_8_ELE), 0, + 0}; + DataCopy(_inMMLocal[loopIdx * gmmSwiglu->tokenLen], workspaceLocal, repeatParams); + mmOutQueue.EnQue(_inMMLocal); +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::Quant(uint32_t loopIdx, + VecConfig &vecConfig) +{ + LocalTensor _inMMLocal = mmOutQueue.DeQue(); + uint64_t preOffset = loopIdx * gmmSwiglu->tokenLen; + uint64_t halfTokenLen = gmmSwiglu->tokenLen / BISECT; + Abs(_inMMLocal[preOffset + gmmSwiglu->tokenLen / BISECT], _inMMLocal[preOffset], halfTokenLen); + PipeBarrier(); + // reduceMax + LocalTensor workLocal = reduceWorkspace.Get(halfTokenLen); + LocalTensor reduceResLocal = + reduceWorkspace.GetWithOffset(FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float)); + LocalTensor reduceTmpLocal = reduceWorkspace.GetWithOffset( + FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float) + UB_BLOCK_UNIT_SIZE); + ReduceMaxTemplate(reduceResLocal, workLocal, _inMMLocal[preOffset + gmmSwiglu->tokenLen / BISECT], reduceTmpLocal, + static_cast(halfTokenLen)); + + int32_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float quantScale = reduceResLocal.GetValue(0) / QUANT_SCALE_INT8; + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + quantScaleLocal.SetValue(loopIdx, quantScale); + quantScale = 1 / quantScale; + int32_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Muls(_inMMLocal[preOffset], _inMMLocal[preOffset], quantScale, halfTokenLen); + PipeBarrier(); + LocalTensor quantLocal = quantOutQueue.DeQue(); + int32_t dstTempOffset = static_cast(preOffset / BISECT); + int32_t srcTempOffset = static_cast(preOffset); + int32_t tempCount = static_cast(halfTokenLen); + LocalTensor castSpace = reduceWorkspace.Get(UB_BLOCK_UNIT_SIZE); + CastFp32ToInt8Template(quantLocal, _inMMLocal, castSpace, dstTempOffset, srcTempOffset, tempCount); + mmOutQueue.EnQue(_inMMLocal); + quantOutQueue.EnQue(quantLocal); +} + +template +__aicore__ inline void +GMMSwigluSplitWorkSpaceCompute::customDataCopyOut(VecConfig &vecConfig) +{ + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + DataCopyParams copyParams_0{1, (uint16_t)(vecConfig.innerLoopNum * sizeof(float)), 0, 0}; + DataCopyPad(quantScaleOutputGM[workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx], quantScaleLocal, + copyParams_0); + LocalTensor quantLocal = quantOutQueue.DeQue(); + DataCopyParams copyParams_1{ + 1, (uint16_t)(vecConfig.innerLoopNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)), 0, 0}; + DataCopyPad(quantOutputGM[(workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx) * gmmSwiglu->tokenLen / + SWIGLU_REDUCE_FACTOR], + quantLocal, copyParams_1); + vecConfig.startIdx += vecConfig.innerLoopNum; + vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen; + quantOutQueue.EnQue(quantLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); +} + +} // namespace GROUPED_MATMUL_SWIGLU_QUANT +#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_SPLIT_WS_H diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_utils.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_utils.h new file mode 100644 index 000000000..80cf1a7c3 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant/op_kernel/grouped_matmul_swiglu_quant_utils.h @@ -0,0 +1,332 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_utils.h + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_UTILS_H +#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_UTILS_H + +#include "kernel_tiling/kernel_tiling.h" +#include "kernel_operator.h" +#include "lib/matmul_intf.h" + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 +// A8W4 MSD场景 +#if defined(ORIG_DTYPE_X) && defined(DT_INT8) && ORIG_DTYPE_X == DT_INT8 && defined(ORIG_DTYPE_WEIGHT) && \ + defined(DT_INT4) && ORIG_DTYPE_WEIGHT == DT_INT4 + #define GMM_SWIGLU_QUANT_A8W4_MSD + using DTYPE_X_A8W4_MSD = AscendC::int4b_t; +// A8W8 场景 +#elif defined(ORIG_DTYPE_X) && defined(DT_INT8) && ORIG_DTYPE_X == DT_INT8 && defined(ORIG_DTYPE_WEIGHT) && \ + defined(DT_INT8) && ORIG_DTYPE_WEIGHT == DT_INT8 + #define GMM_SWIGLU_QUANT_A8W8 +#endif // 场景分类 + +#if defined(FORMAT_WEIGHT) && FORMAT_WEIGHT == FORMAT_FRACTAL_NZ + constexpr CubeFormat wFormat = CubeFormat::NZ; +#elif defined(FORMAT_WEIGHT) && FORMAT_WEIGHT == FORMAT_ND + constexpr CubeFormat wFormat = CubeFormat::ND; +#endif // weight格式分类 + +#endif // 芯片型号分类 + +namespace GROUPED_MATMUL_SWIGLU_QUANT { +using namespace AscendC; +constexpr uint32_t INT8_BITS = 8; // a int8 number has 8 bits +constexpr uint32_t UB_BLOCK_UNIT_SIZE = 32; // 32: a block has 32 bytes data +constexpr uint32_t THRESHOLD_BLOCK_NUM = 8; +constexpr uint32_t UB_BLOCK_DOUBLE_UNIT_SIZE = 64; // 64: a block has 64 bytes data +constexpr uint32_t HALF_UB_BLOCK_UNIT_SIZE = UB_BLOCK_UNIT_SIZE / 2; // 2: a float16 data has two bytes +constexpr uint32_t FLOAT_UB_BLOCK_UNIT_SIZE = 8; // 2: a float16 data has two bytes +constexpr uint32_t SINGLE_CORE_M = 128; +constexpr uint32_t SINGLE_CORE_N = 256; +constexpr uint32_t SINGLE_CORE_K = 7168; +constexpr uint32_t BASIC_M = 128; +constexpr uint32_t BASIC_N = 256; +constexpr uint32_t BASIC_K = 128; +constexpr uint32_t STEP_M = 1; +constexpr uint32_t STEP_N = 1; +constexpr uint32_t STEP_Ka = 4; +constexpr uint32_t STEP_Kb = 4; +constexpr uint32_t DEPTH_A1 = 8; +constexpr uint32_t DEPTH_B1 = 8; +constexpr uint32_t VEC_LEN_ONCE_REPEAT_ELE = 64; +constexpr uint32_t VEC_LEN_ONCE_REPEAT_BLOCK = 8; +constexpr uint32_t FP32_LEN_64_REPEAT = 4096; +constexpr uint32_t REPEAT_64 = 64; +constexpr uint32_t REPEAT_8 = 8; +constexpr uint32_t BISECT = 2; +constexpr uint32_t MOD_32_MASK = 0x1F; +constexpr uint32_t MOD_16_MASK = 0x0F; +constexpr uint32_t ALIGN_8_ELE = 8; +constexpr uint32_t ALIGN_16_ELE = 16; +constexpr float QUANT_SCALE_INT8 = 127.0f; +constexpr int64_t SWIGLU_REDUCE_FACTOR = 2; +constexpr int64_t DOUBLE_BUFFER = 2; +constexpr uint8_t NUM_8 = 8; +constexpr bool NO_BIAS = false; +constexpr int64_t DOUBLE_ROW = 2; +constexpr MatmulConfig CUSTOM_CFG_MDL = GetMDLConfig(false, false, 0, true, false, false, true); +constexpr MatmulConfig GetMMStaticCFG() +{ + MatmulConfig MM_CFG = CUSTOM_CFG_MDL; + MM_CFG.singleCoreM = SINGLE_CORE_M; + MM_CFG.singleCoreN = SINGLE_CORE_N; + MM_CFG.singleCoreK = SINGLE_CORE_K; + MM_CFG.basicM = BASIC_M; + MM_CFG.basicN = BASIC_N; + MM_CFG.basicK = BASIC_K; + return MM_CFG; +} + +constexpr static MatmulApiStaticTiling GetMMTiling(const MatmulApiStaticTiling &mmTiling) +{ + MatmulApiStaticTiling tiling = mmTiling; + tiling.stepM = STEP_M; + tiling.stepN = STEP_N; + tiling.stepKa = STEP_Ka; + tiling.stepKb = STEP_Kb; + tiling.depthA1 = DEPTH_A1; + tiling.depthB1 = DEPTH_B1; + tiling.isBias = NO_BIAS; + return tiling; +} + +template +struct MMImplTypeStatic { + using AT = AT_; + using BT = BT_; + using CT = CT_; + // bias未被使用但高阶模板参数需要传入 + using BiasT = MatmulType; + static constexpr MatmulConfig cfg = GetMMStaticCFG(); + static constexpr MatmulApiStaticTiling mdl = GetMMTiling(GetMatmulApiTiling(cfg)); + using MT = matmul::MatmulImpl; +}; + +template +struct MMImplType { + using AT = AT_; + using BT = BT_; + using CT = CT_; + // bias未被使用但高阶模板参数需要传入 + using BiasT = MatmulType; + using MT = matmul::MatmulImpl; +}; + +struct MNConfig { + int64_t m = 0; + int64_t k = 0; + int64_t n = 0; + int64_t baseM = 0; + int64_t baseN = 0; + int64_t mIdx = 0; + int64_t nIdx = 0; + int64_t blockDimM = 0; + int64_t blockDimN = 0; + int64_t singleM = 0; + int64_t singleN = 0; + int64_t wBaseOffset = 0; + int64_t nAxisBaseOffset = 0; + int64_t mAxisBaseOffset = 0; + int64_t xBaseOffset = 0; + int64_t yBaseOffset = 0; + int64_t wOutOffset = 0; + int64_t workSpaceOffset = 0; +}; + +struct VecConfig { + int64_t M = 0; + int64_t usedCoreNum = 0; + int64_t startOffset = 0; + int64_t curOffset = 0; + int64_t startIdx = 0; + int64_t curIdx = 0; + int64_t taskNum = 0; + int64_t curGroupIdx = 0; + int64_t outLoopNum = 0; + int64_t innerLoopNum = 0; + int64_t tailLoopNum = 0; + int64_t nextUpadteInterVal = 0; +}; + +struct WorkSpaceSplitConfig { + int64_t M = 0; + int64_t loopCount = 0; + int64_t leftMatrixStartIndex = 0; + int64_t rightMatrixExpertStartIndex = 0; + int64_t rightMatrixExpertNextStartIndex = 0; + int64_t rightMatrixExpertEndIndex = 0; + int64_t notLastTaskSize = 0; + int64_t lastLoopTaskSize = 0; + bool isLastLoop = false; +}; + +struct GMAddrParams { + // 输入 GM Tensor + GM_ADDR xGM; // 左矩阵 + GM_ADDR weightGM; // 右矩阵 + GM_ADDR weightScaleGM; // 权重scale + GM_ADDR xScaleGM; // 激活scale + GM_ADDR weightAuxiliaryMatrixGM; // 权重辅助矩阵 + GM_ADDR groupListGM; // 分组矩阵 + // 输出 GM Tensor + GM_ADDR yGM; // 输出量化矩阵 + GM_ADDR yScaleGM; // 输出scale矩阵 + // workspace GM Tensor + GM_ADDR workSpaceGM; // 左矩阵前处理结果矩阵 (double workspace) + 中间处理结果矩阵 (double workspace) + int64_t workSpaceOffset1; + int64_t workSpaceOffset2; + int64_t workSpaceOffset3; +}; + +template +__aicore__ inline auto AlignUp(T a) -> T +{ + if (unlikely(base == 0)) { + return a; + } + return (a + base - 1) / base * base; +} + +template +__aicore__ inline auto AlignUp(T a, T base) -> T +{ + if (unlikely(base == 0)) { + return a; + } + return (a + base - 1) / base * base; +} + +template +__aicore__ inline auto AlignDown(T a, T base) -> T +{ + if (unlikely(base == 0)) { + return a; + } + return a / base * base; +} + +template <> +__aicore__ inline uint32_t AlignUp<4, uint32_t>(uint32_t a) +{ + // to be Multiple of 4, result should be in a format of b(xxxx,x100). + // This means last two bits should be zero, requiring that + // result = num & b(1111,1100) = num & (~3). + // &(~3) operator may reduces num into the range [num, num - 3]. + // As the result should be no less than a (result >= a), it means num - 3 >= a in the worst case. + // In this case, num >= a+3. On the other hand, num should also be less then a+4, otherwise, + // the result will not be least multiple of 4 for 3. In other cases like [num, num - 2], + // num = a + 3 also satisfies the goal condition. + return (a + 3) & ~3; // & ~3: set last two bits of (a+3) to be zero +} + +template <> +__aicore__ inline uint32_t AlignUp<8, uint32_t>(uint32_t a) +{ + // In general, if we want to get the least multiple of b (b is the power of 2) for a, + // it comes to a conclusion from the above comment: result = (a + (b - 1)) & (~b) + return (a + 7) & ~7; // & ~7: set last four bits of (a+7) to be zero +} + +template <> +__aicore__ inline uint32_t AlignUp<16, uint32_t>(uint32_t a) +{ + // In general, if we want to get the least multiple of b (b is the power of 2) for a, + // it comes to a conclusion from the above comment: result = (a + (b - 1)) & (~b) + return (a + 15) & ~15; // & ~15: set last four bits of (a+15) to be zero +} + +template <> +__aicore__ inline uint32_t AlignUp<32, uint32_t>(uint32_t a) +{ + // refer to the above comments. + return (a + 31) & ~31; // & ~31: set last five bits of (a+31) to be zero} +} + +__aicore__ inline void ReduceMaxSmall(const LocalTensor &dstLocal, const LocalTensor &workLocal, + const LocalTensor &srcLocal, uint32_t count) +{ + /** + * @brief ReduceMaxSmall 此函数仅支持入参count小于4096。 + */ + uint32_t repeat = count / VEC_LEN_ONCE_REPEAT_ELE; + uint32_t tailNum = count % VEC_LEN_ONCE_REPEAT_ELE; + if (likely(repeat > 0)) { + WholeReduceMax(workLocal, srcLocal, VEC_LEN_ONCE_REPEAT_ELE, repeat, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK, + ReduceOrder::ORDER_ONLY_VALUE); + PipeBarrier(); + } + if (unlikely(tailNum != 0)) { + WholeReduceMax(workLocal[repeat], srcLocal[count - tailNum], tailNum, 1, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK, + ReduceOrder::ORDER_ONLY_VALUE); + PipeBarrier(); + repeat += 1; + } + WholeReduceMax(dstLocal, workLocal, repeat, 1, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK, ReduceOrder::ORDER_ONLY_VALUE); +} + +__aicore__ inline void ReduceMaxTemplate(const LocalTensor &dstLocal, const LocalTensor &workLocal, + const LocalTensor &srcLocal, const LocalTensor &resTmpLocal, + uint32_t count) +{ + /** + * @brief 当前算子仅支持[32, 10240]长度的词向量维度N，对应此函数count入参范围在[16, 5120]。 + * @param [in] count: 本函数支持count范围为[1,8192]。 + */ + if (count <= FP32_LEN_64_REPEAT) { + ReduceMaxSmall(dstLocal, workLocal, srcLocal, count); + PipeBarrier(); + } else { + BlockReduceMax(workLocal, srcLocal, REPEAT_64, VEC_LEN_ONCE_REPEAT_ELE, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK); + PipeBarrier(); + + BlockReduceMax(workLocal, workLocal, REPEAT_8, VEC_LEN_ONCE_REPEAT_ELE, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK); + PipeBarrier(); + + WholeReduceMax(resTmpLocal, workLocal, VEC_LEN_ONCE_REPEAT_ELE, 1, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK, + ReduceOrder::ORDER_ONLY_VALUE); + PipeBarrier(); + + ReduceMaxSmall(dstLocal, workLocal, srcLocal[FP32_LEN_64_REPEAT], count - FP32_LEN_64_REPEAT); + PipeBarrier(); + + const BinaryRepeatParams repeatParams = {1, 1, 1, NUM_8, NUM_8, NUM_8}; + Max(dstLocal, dstLocal, resTmpLocal, 1, 1, repeatParams); + } +} + +__aicore__ inline void CastFp32ToInt8Template(LocalTensor &dstLocal, LocalTensor &srcLocal, + LocalTensor &oneBlockWorkspace, int32_t dstOffset, + int32_t srcOffset, int32_t count) +{ + Cast(srcLocal[srcOffset].ReinterpretCast(), srcLocal[srcOffset], RoundMode::CAST_RINT, count); + PipeBarrier(); + if ((dstOffset & MOD_32_MASK) == 0) { + Cast(dstLocal[dstOffset], srcLocal[srcOffset].ReinterpretCast(), RoundMode::CAST_RINT, count); + } else if ((dstOffset & MOD_16_MASK) == 0) { + Cast(dstLocal[dstOffset + ALIGN_16_ELE], srcLocal[srcOffset + ALIGN_8_ELE].ReinterpretCast(), + RoundMode::CAST_RINT, count - ALIGN_16_ELE); + PipeBarrier(); + Cast(oneBlockWorkspace, srcLocal[srcOffset].ReinterpretCast(), RoundMode::CAST_RINT, ALIGN_16_ELE); + PipeBarrier(); + for (int32_t i = 0; i < ALIGN_16_ELE; i++) { + int8_t temp = oneBlockWorkspace.GetValue(i); + dstLocal.SetValue(dstOffset + i, temp); + } + PipeBarrier(); + } +} + +} // namespace GROUPED_MATMUL_SWIGLU_QUANT + +#endif // ASCENDC_GROUPED_MATMUL_UTILS_H diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/CMakeLists.txt b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/CMakeLists.txt new file mode 100644 index 000000000..86b308249 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/grouped_matmul_swiglu_quant_v2_torch_adpt.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/grouped_matmul_swiglu_quant_v2_torch_adpt.h new file mode 100644 index 000000000..e8d24e650 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/grouped_matmul_swiglu_quant_v2_torch_adpt.h @@ -0,0 +1,75 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef GROUPED_MATMUL_SWIGLU_QUANT_V2_TORCH_ADPT_H +#define GROUPED_MATMUL_SWIGLU_QUANT_V2_TORCH_ADPT_H +namespace vllm_fl { + +std::tuple grouped_matmul_swiglu_quant_v2( + const at::Tensor & x, + const at::TensorList &weight, + const at::TensorList &weight_scale, + const at::Tensor & x_scale, + const at::Tensor & group_list, + const c10::optional & smooth_scale, + const c10::optional weight_assist_matrix, + const c10::optional & bias, + c10::optional dequant_mode, + c10::optional dequant_dtype, + c10::optional quant_mode, + c10::optional quant_dtype, + bool transpose_weight, + int64_t group_list_type, + at::IntArrayRef tuning_config, + double swiglu_limit) +{ + + auto x_size = x.sizes(); + int n = weight_scale[0].sizes().back(); + int m = x_size[0]; + int k = x_size[1]; + + at::Tensor output = at::empty({m, n/2}, x.options().dtype(at::kChar)); + at::Tensor output_scale = at::empty({m}, x.options().dtype(at::kFloat)); + int64_t dequant_mode_real = dequant_mode.value_or(0); + int64_t dequant_dtype_real = dequant_dtype.value_or(0); + int64_t quant_mode_real = quant_mode.value_or(0); + auto bias_real = bias.value_or(at::Tensor()); + auto smooth_scale_real = smooth_scale.value_or(at::Tensor()); + double swiglu_limit_f = static_cast(swiglu_limit); + auto ws=weight[0].sizes(); + EXEC_NPU_CMD( + aclnnGroupedMatmulSwigluQuantWeightNzV2, + x, + weight, + weight_scale, + weight_assist_matrix, + bias_real, + x_scale, + smooth_scale_real, + group_list, + dequant_mode_real, + dequant_dtype_real, + quant_mode_real, + group_list_type, + tuning_config, + swiglu_limit_f, + output, + output_scale); + return std::tuple(output, output_scale); +} + +} +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/CMakeLists.txt b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/CMakeLists.txt new file mode 100644 index 000000000..375c46404 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/CMakeLists.txt @@ -0,0 +1,31 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnExc PRIVATE + grouped_matmul_swiglu_quant_v2_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME GroupedMatmulSwigluQuantV2 + OPTIONS + --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE grouped_matmul_swiglu_quant_v2 ACLNNTYPE aclnn_exclude) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_base_tiling.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_base_tiling.cpp new file mode 100644 index 000000000..378f0bac2 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_base_tiling.cpp @@ -0,0 +1,507 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file grouped_matmul_swiglu_quant_v2_base_tiling.cpp + * \brief + */ +#include "grouped_matmul_swiglu_quant_v2_base_tiling.h" +#include "util/math_util.h" +#include "err/ops_err.h" + +using namespace matmul_tiling; + +namespace optiling { +namespace GroupedMatmulSwigluQuantV2Tiling { + +constexpr int64_t ND_WEIGHT_MULTI_TENSOR_DIM = 2; +constexpr int64_t NZ_WEIGHT_MULTI_TENSOR_DIM = 4; +constexpr float EFFECTIVE_TASK_RATIO = 0.95f; +constexpr int32_t MIN_BASE_M = 16; + +template +static inline auto AlignUp(T a, T base) -> T +{ + if (base == 0) { + return 0; + } + return (a + base - 1) / base * base; +} + +template +auto CeilDiv(T1 a, T2 b) -> T1 +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + + +static inline uint32_t SixteenAlign(uint32_t a, bool up = false) +{ + if (up) { + a += 15U; + } + return a & ~15U; +} + +int64_t GroupedMatmulSwigluQuantV2BaseTiling::CalMaxRowInUbA8W4(const uint64_t ubSize, const uint64_t n) const +{ + const uint64_t ALIGNMENT = 8; + const float WEIGHT_FACTOR = isA4W4_ ? 4.5f : 8.5f; + const uint64_t ALIGNMENT_TERM_FACTOR = 4; + const uint64_t LINEAR_TERM_FACTOR = 6; + const uint64_t CONSTANT_TERM = 64; + const int64_t MIN_ROW_THRESHOLD = 1; + + // A8W4 表达式：8.5 * row * n + 4 * alignUp(row, 8) + 6n + 64 <= ubSize + // A4W4 表达式：4.5 * row * n + 4 * alignUp(row, 8) + 6n + 64 <= ubSize + + // 忽略对齐项的初始估计 + int64_t maxRowEstimate = + (ubSize - CONSTANT_TERM - LINEAR_TERM_FACTOR * n) / static_cast(WEIGHT_FACTOR * n); + + // 考虑对齐影响 + uint64_t alignedRow = (maxRowEstimate + ALIGNMENT - 1) / ALIGNMENT * ALIGNMENT; + uint64_t totalSize = static_cast(WEIGHT_FACTOR * maxRowEstimate * n) + + ALIGNMENT_TERM_FACTOR * alignedRow + LINEAR_TERM_FACTOR * n + CONSTANT_TERM; + + // 如果超过UB大小，逐步减少row直到满足条件 + while (totalSize > ubSize && maxRowEstimate > 0) { + maxRowEstimate--; + alignedRow = (maxRowEstimate + ALIGNMENT - 1) / ALIGNMENT * ALIGNMENT; + totalSize = static_cast(WEIGHT_FACTOR * maxRowEstimate * n) + ALIGNMENT_TERM_FACTOR * alignedRow + + LINEAR_TERM_FACTOR * n + CONSTANT_TERM; + } + + if (maxRowEstimate < MIN_ROW_THRESHOLD) { + OP_LOGE(context_->GetNodeName(), "GMM_SWIGLU_QUANT TILING: No valid row found for n = %lu, ubSize = %lu\n", n, + ubSize); + return 0; + } + return maxRowEstimate; +} + +int64_t GroupedMatmulSwigluQuantV2BaseTiling::CalMaxRowInUb(const uint64_t ubSize, const uint64_t n) const +{ + uint64_t tmpBufSize = (n / SWIGLU_REDUCE_FACTOR) * FP32_DTYPE_SIZE; + uint64_t perchannleBufSize = n * FP32_DTYPE_SIZE * DOUBLE_BUFFER; + uint64_t reduceMaxResBufSize = BLOCK_BYTE; + uint64_t reduceMaxTmpBufSize = BLOCK_BYTE; + const uint64_t CONSTANT_TERM = 64; + int64_t remainUbSize = ubSize - tmpBufSize - perchannleBufSize - reduceMaxResBufSize - reduceMaxTmpBufSize; + int64_t maxRowInUb = + remainUbSize / (n * INT32_DTYPE_SIZE + n / SWIGLU_REDUCE_FACTOR + FP32_DTYPE_SIZE) / DOUBLE_BUFFER; + int64_t curUb = DOUBLE_BUFFER * (maxRowInUb * (INT32_DTYPE_SIZE * n + n / SWIGLU_REDUCE_FACTOR) + + AlignUp(maxRowInUb, FP32_BLOCK_SIZE) * FP32_DTYPE_SIZE); + if (curUb > remainUbSize) { + // 64 : make sure ub does not excceed maxUbSize after align up to 8 + maxRowInUb = (remainUbSize - CONSTANT_TERM) / + (n * INT32_DTYPE_SIZE + n / SWIGLU_REDUCE_FACTOR + FP32_DTYPE_SIZE) / DOUBLE_BUFFER; + } + if (maxRowInUb < 1) { + // when n > (ubSize - 72) / 19 = 10330, maxRowInUb < 1 + OP_LOGE(context_->GetNodeName(), "GMM_SWIGLU_QUANT TILING: n should not be greater than 10240, now is %lu\n", + n); + } + return maxRowInUb; +} + +bool GroupedMatmulSwigluQuantV2BaseTiling::IsCapable() +{ + auto weightDesc = context_->GetInputDesc(WEIGHT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, weightDesc); + ge::DataType weightDType = weightDesc->GetDataType(); + if (weightDType != ge::DataType::DT_INT4) { + return false; + } + + auto wTensor = context_->GetDynamicInputTensor(WEIGHT_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context_, wTensor); + if (!(wTensor->GetStorageShape().GetDimNum() == ND_WEIGHT_DIM_LIMIT || + wTensor->GetStorageShape().GetDimNum() == ND_WEIGHT_MULTI_TENSOR_DIM || + wTensor->GetStorageShape().GetDimNum() == NZ_WEIGHT_DIM_LIMIT || + wTensor->GetStorageShape().GetDimNum() == NZ_WEIGHT_MULTI_TENSOR_DIM)) { + return false; + } + + return true; +} + +ge::graphStatus GroupedMatmulSwigluQuantV2BaseTiling::ParseInputAndAttr() +{ + auto xDesc = context_->GetInputDesc(X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, xDesc); + auto weightDesc = context_->GetInputDesc(WEIGHT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, weightDesc); + auto wTensor = context_->GetDynamicInputTensor(WEIGHT_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context_, wTensor); + auto xTensor = context_->GetDynamicInputTensor(X_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context_, xTensor); + auto wScaleTensor = context_->GetDynamicInputTensor(WEIGHT_SCALE_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context_, wScaleTensor); + auto groupListTensor = context_->GetDynamicInputTensor(GROUPLIST_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context_, groupListTensor); + + auto wDimNum = wTensor->GetStorageShape().GetDimNum(); + if (wDimNum == ND_WEIGHT_DIM_LIMIT || wDimNum == NZ_WEIGHT_DIM_LIMIT) { + isSingleTensor_ = 1; + } else { + isSingleTensor_ = 0; + } + + auto attr = context_->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context_, attr); // check attr is not null + const int64_t *dequantModePtr = attr->GetAttrPointer(ATTR_INDEX_DEQUANT_MODE); + auto dequantMode = dequantModePtr != nullptr ? *dequantModePtr : 0; + OP_CHECK_IF(!(dequantMode == 0 || dequantMode == 1), + OP_LOGE(context_->GetNodeName(), "dequantMode must be 0 or 1, but actual value is %ld.", dequantMode), + return ge::GRAPH_FAILED); + + const auto swigluLimtPtr = attr->GetAttrPointer(ATTR_INDEX_SWIGLU_LIMIT); + double swigluLimt_ = swigluLimtPtr != nullptr ? *swigluLimtPtr : 0.0f; + OP_CHECK_IF(!(swigluLimt_ >= 0.0), + OP_LOGE(context_->GetNodeName(), "swigluLimit must be non-negative, but actual value is %f.", + swigluLimt_), + return ge::GRAPH_FAILED); + tilingData_.gmmSwigluQuantV2BaseParams.set_swigluLimit(swigluLimt_); + const int64_t *groupListTypePtr = attr->GetAttrPointer(ATTR_INDEX_GROUPLIST_TYPE); + groupListType_ = groupListTypePtr != nullptr ? *groupListTypePtr : 0; + OP_CHECK_IF( + !(groupListType_ == 0 || groupListType_ == 1), + OP_LOGE(context_->GetNodeName(), "GroupListType must be 0 or 1, but actual value is %ld.", groupListType_), + return ge::GRAPH_FAILED); + + ge::DataType xDType = xDesc->GetDataType(); + ge::DataType weightDType = weightDesc->GetDataType(); + + isA8W4MSD_ = (xDType == ge::DataType::DT_INT8 && weightDType == ge::DataType::DT_INT4); + isA4W4_ = (xDType == ge::DataType::DT_INT4 && weightDType == ge::DataType::DT_INT4); + if (isA4W4_) { + auto smoothScaleTensor = context_->GetDynamicInputTensor(SMOOTH_SCALE_INDEX, 0); + if (smoothScaleTensor == nullptr) { + smoothScaleDimNum_ = 0; + } else { + smoothScaleDimNum_ = smoothScaleTensor->GetStorageShape().GetDimNum(); + } + } + + auto compileInfoPtr = context_->GetCompileInfo(); + OP_CHECK_IF(compileInfoPtr == nullptr, OP_LOGE(context_->GetNodeName(), "CompileInfo is nullptr"), + return ge::GRAPH_FAILED); + + m_ = xTensor->GetStorageShape().GetDim(0); + k_ = xTensor->GetStorageShape().GetDim(1); + auto wScaleDimNum = wScaleTensor->GetStorageShape().GetDimNum(); + isWeightTrans_ = false; + if (wTensor->GetStorageShape().GetDimNum() == NZ_WEIGHT_DIM_LIMIT || wTensor->GetStorageShape().GetDimNum() == NZ_WEIGHT_MULTI_TENSOR_DIM) { + isNz_ = true; + } + const auto tuningConfigPtr = attr->GetAttrPointer(ATTR_INDEX_TUNING_CONFIG); + tuningConfig_ = tuningConfigPtr != nullptr && tuningConfigPtr->GetSize() > 1? + (reinterpret_cast(tuningConfigPtr->GetData()))[0] : 0; + + if (isA4W4_) { + n_ = wScaleTensor->GetStorageShape().GetDim(wScaleDimNum - DIM_1); + } else { + if (wTensor->GetStorageShape().GetDimNum() == ND_WEIGHT_DIM_LIMIT) { + // ND SingleTensor [E, K, N] + n_ = wTensor->GetStorageShape().GetDim(DIM_2); + } else if (wTensor->GetStorageShape().GetDimNum() == NZ_WEIGHT_DIM_LIMIT) { + // NZ SingleTensor [E, N // 64, K // 16, 16, 64] + n_ = wTensor->GetStorageShape().GetDim(DIM_1) * wTensor->GetStorageShape().GetDim(DIM_4); + } else if (wTensor->GetStorageShape().GetDimNum() == ND_WEIGHT_MULTI_TENSOR_DIM) { + // ND MultiTensor [K, N] + n_ = wTensor->GetStorageShape().GetDim(DIM_1); + } else if (wTensor->GetStorageShape().GetDimNum() == NZ_WEIGHT_MULTI_TENSOR_DIM) { + // NZ MultiTensor [N // 64, K // 16, 16, 64] + n_ = wTensor->GetStorageShape().GetDim(DIM_0) * wTensor->GetStorageShape().GetDim(DIM_3); + } + } + + isWeightTrans_ = *attr->GetAttrPointer(ATTR_INDEX_TRANSPOSE_WEIGHT); + + if (dequantMode == 1) { // perGroup量化模式：单tensor场景[E, KGroupCount, N]，多tensor场景[KGroupCount, N] + quantGroupNum_ = wScaleTensor->GetStorageShape().GetDim(wScaleDimNum - DIM_2); + } else { // perChannel量化模式 + quantGroupNum_ = 1; + } + + groupNum_ = groupListTensor->GetStorageShape().GetDim(0); + + if (isA8W4MSD_ || isA4W4_) { + maxProcessRowNum_ = CalMaxRowInUbA8W4(compileInfoPtr->ubSize_, n_); + } else { + maxProcessRowNum_ = CalMaxRowInUb(compileInfoPtr->ubSize_, n_); + } + + blockDim_ = compileInfoPtr->aicNum_; + return ge::GRAPH_SUCCESS; +} + +int32_t GroupedMatmulSwigluQuantV2BaseTiling::FindBestSingleN(const uint32_t &aicNum, int64_t baseM, int64_t baseN) const +{ + uint64_t quantGroupNum = quantGroupNum_; + if (n_ < baseN || tuningConfig_ <= 0 || !(quantGroupNum == 1)) { + return baseN; + } + int32_t mDim = CeilDiv(tuningConfig_, baseM); + int32_t nDim = CeilDiv(n_, baseN); + int32_t taskNum = mDim * nDim * static_cast(groupNum_); + int32_t taskNumPerCore = CeilDiv(taskNum, aicNum); + // 每个核只需要做1个基本块的时候，任务量太少，无需处理 + if (taskNumPerCore <= 1) { + return baseN; + } + int32_t curNDim = 0; + int32_t curTaskNum = 0; + int32_t bestSingleN = baseN; + float ratio = 0; + for (uint32_t i = 1; i <= aicNum; ++i) { + if (isNz_) { + bestSingleN = CeilDiv(static_cast(n_), i); + if (bestSingleN != n_ && bestSingleN % baseN != 0) { + continue; + } + } else { + // 暂时只NZ格式开启动态分块 + return baseN; + } + curNDim = CeilDiv(n_, bestSingleN); + curTaskNum = mDim * curNDim * static_cast(groupNum_); + ratio = static_cast(curTaskNum) / AlignUp(static_cast(curTaskNum), aicNum); + if (ratio >= EFFECTIVE_TASK_RATIO) { + return bestSingleN; + } + } + return baseN; +} + +bool GroupedMatmulSwigluQuantV2BaseTiling::TryFullLoadA(int32_t baseM, int64_t baseN, int64_t baseK, uint64_t l1Size) +{ + // 暂时只支持A4W4 + float sizeofweightDtype = 0.5f; + float sizeofxDtype = 0.5f; + auto matBl1Size = static_cast(tilingData_.mmTilingData.get_depthB1() * baseN * baseK * sizeofweightDtype); + auto remainL1Size = l1Size - matBl1Size - 8 * baseN; + int32_t newDepthA1 = CeilDiv(k_, baseK); + if (static_cast(newDepthA1 * baseM * baseK * sizeofxDtype) < static_cast(remainL1Size)) { + tilingData_.mmTilingData.set_stepKa(newDepthA1); + tilingData_.mmTilingData.set_depthA1(newDepthA1); + return true; + } + return false; +} + + +ge::graphStatus GroupedMatmulSwigluQuantV2BaseTiling::DynamicTilingSingleN(gert::TilingContext *context, const uint32_t &aicNum, + int64_t baseM, int64_t baseN, int64_t baseK) +{ + //get info + auto platformInfoPtr = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfoPtr); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + uint64_t l1Size = 0; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L1, l1Size); + tilingData_.gmmSwigluQuantV2BaseParams.set_singleN(0); + + if (n_ < baseN || tuningConfig_ <= 0 || !isA4W4_) { + return ge::GRAPH_SUCCESS; + } + int32_t bestSingleN = FindBestSingleN(aicNum, baseM, baseN); + if (bestSingleN == baseN) { // 没找到更优的singleN + return ge::GRAPH_SUCCESS; + } + tilingData_.gmmSwigluQuantV2BaseParams.set_singleN(bestSingleN); + // 先不改看看baseM能否全载左矩阵 + if (TryFullLoadA(baseM, baseN, baseK, l1Size)) { + return ge::GRAPH_SUCCESS; + } + // 可以尝试减小baseM来全载左矩阵 + int32_t newBaseM = static_cast(SixteenAlign(tuningConfig_, true)); + // 防止不均匀情况 + newBaseM += MIN_BASE_M; + // 再看看能否全载左矩阵 + if (newBaseM < baseM && TryFullLoadA(newBaseM, baseN, baseK, l1Size)) { + tilingData_.mmTilingData.set_baseM(newBaseM); + return ge::GRAPH_SUCCESS; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus GroupedMatmulSwigluQuantV2BaseTiling::DoOpTiling() +{ + OP_LOGD(context_->GetNodeName(), "Begin Run GMM Swiglu Tiling ."); + + if (ParseInputAndAttr() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo()); + MatmulApiTiling tiling(ascendcPlatform); + tiling.SetAType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_INT4); + tiling.SetBType(TPosition::GM, CubeFormat::NZ, matmul_tiling::DataType::DT_INT4); + tiling.SetCType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16); + tiling.SetBias(false); + tiling.SetShape(A8W4_BASEM, A8W4_BASEN, k_); + tiling.SetFixSplit(A8W4_BASEM, A8W4_BASEN, A8W4_BASEK); + tiling.SetOrgShape(m_, n_, k_); + tiling.SetBufferSpace(-1, -1, -1); + OP_CHECK_IF(tiling.GetTiling(tilingData_.mmTilingData) == -1, + OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), + "grouped_matmul_swiglu_quant_base_tiling, get tiling failed"), + return ge::GRAPH_FAILED); + if (isA8W4MSD_ || isA4W4_) { + tilingData_.mmTilingData.set_baseM(A8W4_BASEM); + tilingData_.mmTilingData.set_baseN(A8W4_BASEN); + tilingData_.mmTilingData.set_baseK(A8W4_BASEK); + tilingData_.mmTilingData.set_dbL0B(DOUBLE_BUFFER); + tilingData_.mmTilingData.set_stepKa(NUM_FOUR); + tilingData_.mmTilingData.set_stepKb(NUM_FOUR); + tilingData_.mmTilingData.set_depthA1(NUM_EIGHT); + tilingData_.mmTilingData.set_depthB1(NUM_EIGHT); + tilingData_.mmTilingData.set_stepM(1); + tilingData_.mmTilingData.set_stepN(1); + + } + + usrWorkspaceLimit_ = USER_WORKSPACE_LIMIT; + mLimit_ = 0; + if (isA8W4MSD_) { + mLimit_ = + ((usrWorkspaceLimit_ / DOUBLE_WORKSPACE_SPLIT) / (k_ * sizeof(int8_t) + DOUBLE_ROW * n_ * SIZE_OF_HALF_2)); + } else if (isA4W4_) { + mLimit_ = ((usrWorkspaceLimit_ / DOUBLE_WORKSPACE_SPLIT) / (n_ * SIZE_OF_HALF_2)); + } else { + mLimit_ = ((usrWorkspaceLimit_ / DOUBLE_WORKSPACE_SPLIT) / INT32_DTYPE_SIZE) / n_; + } + + OP_CHECK_IF(mLimit_ <= 0, + OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "mLimit_ is %ld must over then 0.", mLimit_), + return ge::GRAPH_FAILED); + tilingData_.gmmSwigluQuantV2BaseParams.set_mLimit(mLimit_); + + DynamicTilingSingleN(context_, blockDim_, A8W4_BASEM, A8W4_BASEN, A8W4_BASEK); + + if (isA8W4MSD_) { + int workSpaceMTemp = mLimit_ * DOUBLE_WORKSPACE_SPLIT; + tilingData_.gmmSwigluQuantV2BaseParams.set_workSpaceOffset1(workSpaceMTemp * k_ * sizeof(int8_t)); + tilingData_.gmmSwigluQuantV2BaseParams.set_workSpaceOffset2(DOUBLE_ROW * workSpaceMTemp * n_ * SIZE_OF_HALF_2); + workspaceSize_ = + SYS_WORKSPACE_SIZE + // 系统预留16MB + (workSpaceMTemp * k_ * sizeof(int8_t)) + // 第一阶段 预处理左矩阵 (mLimit_, K) * int8 * 2(double WorkSpace) + (DOUBLE_ROW * workSpaceMTemp * n_ * + SIZE_OF_HALF_2); // 第二阶段 矩阵乘结果 (2 * mLimit_, N) * fp16 * 2(double WorkSpace) + } else if (isA4W4_) { + int workSpaceMTemp = mLimit_ * DOUBLE_WORKSPACE_SPLIT; + tilingData_.gmmSwigluQuantV2BaseParams.set_workSpaceOffset1(mLimit_ * n_ * SIZE_OF_HALF_2); + tilingData_.gmmSwigluQuantV2BaseParams.set_workSpaceOffset2(0); + workspaceSize_ = SYS_WORKSPACE_SIZE + (workSpaceMTemp * n_ * SIZE_OF_HALF_2); + } else { + int workSpaceMTemp = (mLimit_ * DOUBLE_WORKSPACE_SPLIT > m_ ? m_ : mLimit_ * DOUBLE_WORKSPACE_SPLIT); + tilingData_.gmmSwigluQuantV2BaseParams.set_workSpaceOffset1(0); + tilingData_.gmmSwigluQuantV2BaseParams.set_workSpaceOffset2(0); + workspaceSize_ = SYS_WORKSPACE_SIZE + (workSpaceMTemp * n_ * sizeof(int32_t)); + } + + isSplitWorkSpace_ = m_ > mLimit_ * DOUBLE_WORKSPACE_SPLIT; + SetTilingKeyAndScheMode(); + FillTilingData(); + PrintTilingData(); + OP_LOGD(context_->GetNodeName(), "End Run GMM Swiglu Tiling."); + return ge::GRAPH_SUCCESS; +} + +uint64_t GroupedMatmulSwigluQuantV2BaseTiling::GetTilingKey() const +{ + return tilingKey_; +} + +void GroupedMatmulSwigluQuantV2BaseTiling::FillTilingData() +{ + tilingData_.gmmSwigluQuantV2BaseParams.set_groupNum(groupNum_); + tilingData_.gmmSwigluQuantV2BaseParams.set_coreNum(blockDim_); + tilingData_.gmmSwigluQuantV2BaseParams.set_K(k_); + tilingData_.gmmSwigluQuantV2BaseParams.set_N(n_); + tilingData_.gmmSwigluQuantV2BaseParams.set_M(m_); + tilingData_.gmmSwigluQuantV2BaseParams.set_baseM(A8W4_BASEM); + tilingData_.gmmSwigluQuantV2BaseParams.set_baseN(A8W4_BASEN); + tilingData_.gmmSwigluQuantV2BaseParams.set_quantGroupNum(quantGroupNum_); + tilingData_.gmmSwigluQuantV2BaseParams.set_isSingleTensor(isSingleTensor_); + tilingData_.gmmSwigluQuantV2BaseParams.set_groupListType(groupListType_); + tilingData_.gmmSwigluQuantV2BaseParams.set_smoothScaleDimNum(smoothScaleDimNum_); + tilingData_.gmmSwigluQuantV2.set_maxProcessRowNum(maxProcessRowNum_); + tilingData_.gmmSwigluQuantV2.set_groupListLen(groupNum_); + tilingData_.gmmSwigluQuantV2.set_tokenLen(n_); +} + +void GroupedMatmulSwigluQuantV2BaseTiling::PrintTilingData() +{ + OP_LOGD(context_->GetNodeName(), "grouped_matmul_swiglu_quant_base_tiling."); + OP_LOGD(context_->GetNodeName(), "groupNum: %ld", tilingData_.gmmSwigluQuantV2BaseParams.get_groupNum()); + OP_LOGD(context_->GetNodeName(), "coreNum: %ld", tilingData_.gmmSwigluQuantV2BaseParams.get_coreNum()); + OP_LOGD(context_->GetNodeName(), "M: %ld", tilingData_.gmmSwigluQuantV2BaseParams.get_M()); + OP_LOGD(context_->GetNodeName(), "K: %ld", tilingData_.gmmSwigluQuantV2BaseParams.get_K()); + OP_LOGD(context_->GetNodeName(), "N: %ld", tilingData_.gmmSwigluQuantV2BaseParams.get_N()); + OP_LOGD(context_->GetNodeName(), "baseM: %ld", tilingData_.gmmSwigluQuantV2BaseParams.get_baseM()); + OP_LOGD(context_->GetNodeName(), "baseN: %ld", tilingData_.gmmSwigluQuantV2BaseParams.get_baseN()); + OP_LOGD(context_->GetNodeName(), "mLimit: %ld", tilingData_.gmmSwigluQuantV2BaseParams.get_mLimit()); + OP_LOGD(context_->GetNodeName(), "quantGroupNum: %ld", tilingData_.gmmSwigluQuantV2BaseParams.get_quantGroupNum()); + OP_LOGD(context_->GetNodeName(), "isSingleTensor:%ld", tilingData_.gmmSwigluQuantV2BaseParams.get_isSingleTensor()); + OP_LOGD(context_->GetNodeName(), "groupListType: %ld", tilingData_.gmmSwigluQuantV2BaseParams.get_groupListType()); + OP_LOGD(context_->GetNodeName(), "get_swigluLimit: %ld", tilingData_.gmmSwigluQuantV2BaseParams.get_swigluLimit()); + OP_LOGD(context_->GetNodeName(), "smoothScaleDimNum: %ld", + tilingData_.gmmSwigluQuantV2BaseParams.get_smoothScaleDimNum()); + OP_LOGD(context_->GetNodeName(), "maxProcessRowNum: %ld", tilingData_.gmmSwigluQuantV2.get_maxProcessRowNum()); + OP_LOGD(context_->GetNodeName(), "groupListLen: %ld", tilingData_.gmmSwigluQuantV2.get_groupListLen()); + OP_LOGD(context_->GetNodeName(), "tokenLen: %ld", tilingData_.gmmSwigluQuantV2.get_tokenLen()); + OP_LOGD(context_->GetNodeName(), "USER_WORKSPACE_LIMIT: %ld", usrWorkspaceLimit_); + OP_LOGD(context_->GetNodeName(), "workspaceSizes: %lu", workspaceSize_); + OP_LOGD(context_->GetNodeName(), "isSplitWorkSpace: %s", isSplitWorkSpace_ ? "true" : "false"); +} + +void GroupedMatmulSwigluQuantV2BaseTiling::SetTilingKeyAndScheMode() +{ + if (isA8W4MSD_) { // A8W4 MSD tiling_key + tilingKey_ = A8W4_MSD_TILING_KEY_MODE; + context_->SetScheduleMode(BATCH_MODE_SCHEDULE); + } else if (isA4W4_ && !isWeightTrans_) { + tilingKey_ = A4W4_WEIGHT_NOTRANS_TILING_KEY_MODE; + context_->SetScheduleMode(BATCH_MODE_SCHEDULE); + } else if (isA4W4_ && isWeightTrans_) { + tilingKey_ = A4W4_WEIGHT_TRANS_TILING_KEY_MODE; + context_->SetScheduleMode(BATCH_MODE_SCHEDULE); + } else if (isSplitWorkSpace_) { + tilingKey_ = SPLITWORKSPACE_TILING_KEY_MODE; + context_->SetScheduleMode(BATCH_MODE_SCHEDULE); + } else { + tilingKey_ = COMMON_TILING_KEY_MODE; + context_->SetScheduleMode(BATCH_MODE_SCHEDULE); + } +} + +ge::graphStatus GroupedMatmulSwigluQuantV2BaseTiling::PostTiling() +{ + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetRawTilingData()); + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + context_->SetBlockDim(blockDim_); + + size_t *workspaces = context_->GetWorkspaceSizes(1); // set workspace + OP_CHECK_IF(workspaces == nullptr, OPS_REPORT_CUBE_INNER_ERR(context_->GetNodeName(), "workspaces is null"), + return ge::GRAPH_FAILED); + workspaces[0] = workspaceSize_; + + return ge::GRAPH_SUCCESS; +} + +} // namespace GroupedMatmulSwigluQuantV2Tiling +} // namespace optiling diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_base_tiling.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_base_tiling.h new file mode 100644 index 000000000..6784d9bc5 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_base_tiling.h @@ -0,0 +1,77 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_base_tiling.h + * \brief + */ +#ifndef __OP_HOST_OP_TILING_GROUPED_MATMUL_SWIGLU_QUANT_V2_BASE_TILING_H__ +#define __OP_HOST_OP_TILING_GROUPED_MATMUL_SWIGLU_QUANT_V2_BASE_TILING_H__ + +#include "grouped_matmul_swiglu_quant_v2_tiling.h" +#include "tiling_base/tiling_base.h" +#include "err/ops_err.h" + +namespace optiling { +namespace GroupedMatmulSwigluQuantV2Tiling { + +class GroupedMatmulSwigluQuantV2BaseTiling : public GroupedMatmulSwigluQuantV2Tiling { +public: + explicit GroupedMatmulSwigluQuantV2BaseTiling(gert::TilingContext* context) : GroupedMatmulSwigluQuantV2Tiling(context) {}; + + ~GroupedMatmulSwigluQuantV2BaseTiling() override = default; + +protected: + bool IsCapable() override; + + ge::graphStatus DoOpTiling() override; + + uint64_t GetTilingKey() const override; + + ge::graphStatus PostTiling() override; + + void FillTilingData() override; + void PrintTilingData() override; + void SetTilingKeyAndScheMode(void); + ge::graphStatus ParseInputAndAttr(); + int64_t CalMaxRowInUbA8W4(const uint64_t ubSize, const uint64_t n) const; + int64_t CalMaxRowInUb(const uint64_t ubSize, const uint64_t n) const; + int32_t FindBestSingleN(const uint32_t &aicNum, int64_t baseM, int64_t baseN) const; + bool TryFullLoadA(int32_t baseM, int64_t baseN, int64_t baseK, uint64_t l1Size); + ge::graphStatus DynamicTilingSingleN(gert::TilingContext *context, const uint32_t &aicNum, + int64_t baseM, int64_t baseN, int64_t baseK); + +private: + GMMSwigluQuantV2TilingData tilingData_; + int64_t k_ = 0; + int64_t m_ = 0; + int64_t n_ = 0; + int64_t quantGroupNum_ = 0; + int64_t mLimit_ = 0; + int64_t blockDim_ = 0; + int64_t maxProcessRowNum_ = 0; + int64_t groupNum_ = 0; + int64_t isSingleTensor_ = 1; + int64_t groupListType_ = 0; + int64_t smoothScaleDimNum_ = 0; + int64_t usrWorkspaceLimit_ = 0; + uint64_t workspaceSize_ = 0; + int64_t tuningConfig_ = 0; + float swigluLimtPtr_ = 0.0f; + bool isA8W4MSD_ = false; + bool isA4W4_ = false; + bool isNz_ = false; + bool isWeightTrans_ = false; + bool isSplitWorkSpace_ = false; +}; + +} +} +#endif // __OP_HOST_OP_TILING_GROUPED_MATMUL_SWIGLU_QUANT_V2_BASE_TILING_H__ diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_def.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_def.cpp new file mode 100644 index 000000000..d0a6b88a4 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_def.cpp @@ -0,0 +1,260 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_def.cpp + * \brief + */ + +#include "register/op_def_registry.h" + +namespace ops { +class GroupedMatmulSwigluQuantV2 : public OpDef { +public: + explicit GroupedMatmulSwigluQuantV2(const char *name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT4, ge::DT_INT4}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("x_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("group_list") + .ParamType(REQUIRED) + .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("weight") + .ParamType(DYNAMIC) + .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT4, ge::DT_INT4, ge::DT_INT4, ge::DT_INT4}) + .Format({ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_ND, + ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_ND}); + this->Input("weight_scale") + .ParamType(DYNAMIC) + .DataType({ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_UINT64, ge::DT_UINT64, ge::DT_UINT64, ge::DT_UINT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("weight_assist_matrix") + .ParamType(DYNAMIC) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("bias") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("smooth_scale") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("y_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + this->Attr("dequant_mode").AttrType(OPTIONAL).Int(0); + this->Attr("dequant_dtype").AttrType(OPTIONAL).Int(0); + this->Attr("quant_mode").AttrType(OPTIONAL).Int(0); + this->Attr("quant_dtype").AttrType(OPTIONAL).Int(0); + this->Attr("transpose_weight").AttrType(OPTIONAL).Bool(0); + this->Attr("group_list_type").AttrType(OPTIONAL).Int(0); + this->Attr("tuning_config").AttrType(OPTIONAL).ListInt({0}); + this->Attr("swiglu_limit").AttrType(OPTIONAL).Float(0.0f); + + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true); + + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + + OpAICoreConfig config950; + config950.Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, + ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, + ge::DT_FLOAT4_E2M1, + ge::DT_FLOAT4_E2M1, + ge::DT_FLOAT4_E2M1, + ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_HIFLOAT8, + ge::DT_HIFLOAT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, + ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E5M2, + ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, + ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E5M2, + ge::DT_FLOAT8_E5M2}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}); + config950.Input("x_scale") + .ParamType(REQUIRED) + .DataType( + {ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, + ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, + ge::DT_FLOAT8_E8M0, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}); + config950.Input("group_list") + .ParamType(REQUIRED) + .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}); + config950.Input("weight") + .ParamType(DYNAMIC) + .DataType({ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, + ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, + ge::DT_FLOAT4_E2M1, + ge::DT_FLOAT4_E2M1, + ge::DT_FLOAT4_E2M1, + ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_HIFLOAT8, + ge::DT_HIFLOAT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, + ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, + ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, + ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, + ge::DT_FLOAT8_E5M2}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}); + config950.Input("weight_scale") + .ParamType(DYNAMIC) + .DataType( + {ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, + ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, + ge::DT_FLOAT8_E8M0, ge::DT_BF16, + ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_BF16, + ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, + ge::DT_BF16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + config950.Input("weight_assist_matrix") + .ParamType(DYNAMIC) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}); + config950.Input("bias") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}); + config950.Input("smooth_scale") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}); + + config950.Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E5M2, + ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN, + ge::DT_FLOAT8_E5M2, + ge::DT_FLOAT8_E4M3FN, + ge::DT_FLOAT4_E2M1, + ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_HIFLOAT8, + ge::DT_HIFLOAT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, + ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, + ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, + ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, + ge::DT_FLOAT8_E5M2}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}); + config950.Output("y_scale") + .ParamType(REQUIRED) + .DataType( + {ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, + ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, ge::DT_FLOAT8_E8M0, + ge::DT_FLOAT8_E8M0, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}); + + config950.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("prebuildPattern.value", "Opaque") + .ExtendCfgInfo("coreType.value", "AiCore") + .ExtendCfgInfo("aclnnSupport.value", "support_aclnn") + .ExtendCfgInfo("opFile.value","grouped_matmul_swiglu_quant_v2_apt"); + } +}; + +OP_ADD(GroupedMatmulSwigluQuantV2); +} // namespace ops diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_fusion_tiling.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_fusion_tiling.cpp new file mode 100644 index 000000000..c3d8cfb41 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_fusion_tiling.cpp @@ -0,0 +1,204 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file grouped_matmul_swiglu_quant_v2_fusion_tiling.cpp + * \brief + */ +#include "grouped_matmul_swiglu_quant_v2_fusion_tiling.h" +#include "util/math_util.h" +#include "err/ops_err.h" + +namespace optiling { +namespace GroupedMatmulSwigluQuantV2Tiling { + +constexpr int64_t BASE_M = 128; +constexpr int64_t BASE_K = 128; +constexpr int64_t BASE_N = 256; +constexpr int64_t UB_Y_FACTOR = 2; +constexpr int64_t EXTEND_WORKSPACE_SIZE = (20 * 1024 * 1024); +constexpr int64_t NZ_WEIGHT_SINGLE_TENSOR_DIM = 5; // single: [E, N/32, K/16, 16, 32] +constexpr int64_t NZ_WEIGHT_MULTI_TENSOR_DIM = 4; // multi: each [N/32, K/16, 16, 32] +constexpr int64_t MIN_UB_FACTOR_DIM_X_N = 4600; +constexpr int64_t MID_UB_FACTOR_DIM_X_N = 8192; + +using namespace matmul_tiling; + +bool GroupedMatmulSwigluQuantV2FusionTiling::IsCapable() +{ + auto weightDesc = context_->GetInputDesc(WEIGHT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, weightDesc); + ge::DataType weightDType = weightDesc->GetDataType(); + if (weightDType != ge::DataType::DT_INT8) { + return false; + } + + auto wTensor = context_->GetDynamicInputTensor(WEIGHT_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context_, wTensor); + if (!(wTensor->GetStorageShape().GetDimNum() == NZ_WEIGHT_DIM_LIMIT || + wTensor->GetStorageShape().GetDimNum() == NZ_WEIGHT_MULTI_TENSOR_DIM)) { + return false; + } + return true; +} + +ge::graphStatus GroupedMatmulSwigluQuantV2FusionTiling::ParseInputAndAttr() +{ + auto xTensor = context_->GetDynamicInputTensor(X_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context_, xTensor); + auto wTensor = context_->GetDynamicInputTensor(WEIGHT_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context_, wTensor); + auto groupListTensor = context_->GetDynamicInputTensor(GROUPLIST_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context_, groupListTensor); + groupNum_ = groupListTensor->GetStorageShape().GetDim(0); + auto wDimNum = wTensor->GetStorageShape().GetDimNum(); + if (wDimNum == NZ_WEIGHT_DIM_LIMIT) { + isSingleTensor_ = 1; + } else { + isSingleTensor_ = 0; // multi tensor: 4D per weight [N/32, K/16, 16, 32] + } + auto attr = context_->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context_, attr); // check attr is not null + const int64_t *groupListTypePtr = attr->GetAttrPointer(ATTR_INDEX_GROUPLIST_TYPE); + groupListType_ = groupListTypePtr != nullptr ? *groupListTypePtr : 0; + OP_CHECK_IF(!(groupListType_ == 0 || groupListType_ == 1), + OP_LOGE(context_->GetNodeName(), "GroupListType must be 0 or 1, but actual value is %ld.", groupListType_), + return ge::GRAPH_FAILED); + + const auto swigluLimtPtr = attr->GetAttrPointer(ATTR_INDEX_SWIGLU_LIMIT); + double swigluLimt_ = swigluLimtPtr != nullptr ? *swigluLimtPtr : 0.0f; + OP_CHECK_IF(!(swigluLimt_ >= 0.0), + OP_LOGE(context_->GetNodeName(), "swigluLimit must be non-negative, but actual value is %f.", + swigluLimt_), + return ge::GRAPH_FAILED); + tilingData_.set_swigluLimit(swigluLimt_); + m_ = xTensor->GetStorageShape().GetDim(0); + k_ = xTensor->GetStorageShape().GetDim(1); + if (wDimNum == NZ_WEIGHT_DIM_LIMIT) { + n_ = wTensor->GetStorageShape().GetDim(DIM_1) * wTensor->GetStorageShape().GetDim(DIM_4); + } else { + // 4D multi tensor: [N/32, K/16, 16, 32] -> N = dim0 * dim3 + n_ = wTensor->GetStorageShape().GetDim(0) * wTensor->GetStorageShape().GetDim(3); + } + if (n_ < MIN_UB_FACTOR_DIM_X_N) { + ubFactorDimx_ = 0x4; + } else if (n_ >= MIN_UB_FACTOR_DIM_X_N && n_ < MID_UB_FACTOR_DIM_X_N) { + ubFactorDimx_ = 0x2; + } else { + ubFactorDimx_ = 1; + } + + auto platformInfo = context_->GetPlatformInfo(); + if (platformInfo == nullptr) { + auto compileInfoPtr = context_->GetCompileInfo(); + OP_CHECK_IF(compileInfoPtr == nullptr, OP_LOGE(context_->GetNodeName(), "CompileInfo is nullptr"), + return ge::GRAPH_FAILED); + aicCoreNum_ = compileInfoPtr->aicNum_; + aivCoreNum_ = compileInfoPtr->aivNum_; + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + aicCoreNum_ = ascendcPlatform.GetCoreNumAic(); + aivCoreNum_ = ascendcPlatform.GetCoreNumAiv(); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus GroupedMatmulSwigluQuantV2FusionTiling::DoOpTiling() +{ + OP_LOGD(context_->GetNodeName(), "Begin Run GMM Swiglu Fusion Tiling."); + + if (ParseInputAndAttr() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo()); + MatmulApiTiling tiling(ascendcPlatform); + tiling.SetAType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_INT8); + tiling.SetBType(TPosition::GM, CubeFormat::NZ, matmul_tiling::DataType::DT_INT8); + tiling.SetCType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_INT32); + tiling.SetBias(false); + tiling.SetShape(m_, BASE_N, k_); + tiling.SetOrgShape(m_, n_, k_); + tiling.SetBufferSpace(-1, -1, -1); + OP_CHECK_IF( + tiling.GetTiling(tilingData_.matmulTiling) == -1, + OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "grouped_matmul_swiglu_quant_tiling, get tiling failed"), + return ge::GRAPH_FAILED); + + workspaceSize_ = static_cast(m_) * static_cast(n_) * sizeof(int32_t) + EXTEND_WORKSPACE_SIZE; + tilingKey_ = A8W8_FUSION_KEY_MODE; + FillTilingData(); + PrintTilingData(); + return ge::GRAPH_SUCCESS; +} + +uint64_t GroupedMatmulSwigluQuantV2FusionTiling::GetTilingKey() const +{ + return tilingKey_; +} + +void GroupedMatmulSwigluQuantV2FusionTiling::PrintTilingData() +{ + OP_LOGD(context_->GetNodeName(), "cubeBlockDim: %d", tilingData_.get_cubeBlockDim()); + OP_LOGD(context_->GetNodeName(), "vectorBlockDim: %d", tilingData_.get_vectorBlockDim()); + OP_LOGD(context_->GetNodeName(), "K: %d", tilingData_.get_K()); + OP_LOGD(context_->GetNodeName(), "M: %d", tilingData_.get_M()); + OP_LOGD(context_->GetNodeName(), "N: %d", tilingData_.get_N()); + OP_LOGD(context_->GetNodeName(), "ubFactorDimx: %d", tilingData_.get_ubFactorDimx()); + OP_LOGD(context_->GetNodeName(), "ubFactorDimy: %d", tilingData_.get_ubFactorDimy()); + OP_LOGD(context_->GetNodeName(), "groupListType: %ld", tilingData_.get_groupListType()); + OP_LOGD(context_->GetNodeName(), "isSingleTensor: %d", tilingData_.get_isSingleTensor()); +} + +void GroupedMatmulSwigluQuantV2FusionTiling::FillTilingData() +{ + tilingData_.set_cubeBlockDim(aicCoreNum_); + tilingData_.set_vectorBlockDim(aivCoreNum_); + tilingData_.set_groupNum(groupNum_); + tilingData_.set_K(k_); + tilingData_.set_N(n_); + tilingData_.set_M(m_); + tilingData_.set_ubFactorDimx(ubFactorDimx_); + tilingData_.set_ubFactorDimy(n_ / UB_Y_FACTOR); + tilingData_.set_groupListType(groupListType_); + tilingData_.set_isSingleTensor(isSingleTensor_); + + blockDim_ = aicCoreNum_; + tilingData_.matmulTiling.set_usedCoreNum(aicCoreNum_); + tilingData_.matmulTiling.set_shareMode(0); + tilingData_.matmulTiling.set_dbL0C(1); + tilingData_.matmulTiling.set_baseM(BASE_M); + tilingData_.matmulTiling.set_baseN(BASE_N); + tilingData_.matmulTiling.set_baseK(BASE_K); + tilingData_.matmulTiling.set_stepKa(0x4); // 4: L1中左矩阵单次搬运基于baseK的4倍数据 + tilingData_.matmulTiling.set_stepKb(0x4); // 4: L1中右矩阵单次搬运基于baseK的4倍数据 + tilingData_.matmulTiling.set_depthA1(0x8); // 8: stepKa的两倍，开启double buffer + tilingData_.matmulTiling.set_depthB1(0x8); // 8: stepKb的两倍，开启double buffer + tilingData_.matmulTiling.set_stepM(1); + tilingData_.matmulTiling.set_stepN(1); +} + +ge::graphStatus GroupedMatmulSwigluQuantV2FusionTiling::PostTiling() +{ + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetRawTilingData()); + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->SetBlockDim(blockDim_); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + + size_t *workspaces = context_->GetWorkspaceSizes(1); // set workspace + OP_CHECK_IF(workspaces == nullptr, + OPS_REPORT_CUBE_INNER_ERR(context_->GetNodeName(), "fusion tiling workspaces is null"), + return ge::GRAPH_FAILED); + workspaces[0] = workspaceSize_; + + return ge::GRAPH_SUCCESS; +} +} // namespace GroupedMatmulSwigluQuantV2Tiling +} // namespace optiling diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_fusion_tiling.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_fusion_tiling.h new file mode 100644 index 000000000..ef3ec09a4 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_fusion_tiling.h @@ -0,0 +1,59 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_fusion_tiling.h + * \brief + */ +#ifndef __OP_HOST_OP_TILING_GROUPED_MATMUL_SWIGLU_QUANT_V2_FUSION_TILING_H__ +#define __OP_HOST_OP_TILING_GROUPED_MATMUL_SWIGLU_QUANT_V2_FUSION_TILING_H__ + +#include "grouped_matmul_swiglu_quant_v2_tiling.h" +#include "tiling_base/tiling_base.h" +#include "err/ops_err.h" + +namespace optiling { +namespace GroupedMatmulSwigluQuantV2Tiling { + +class GroupedMatmulSwigluQuantV2FusionTiling : public GroupedMatmulSwigluQuantV2Tiling { +public: + explicit GroupedMatmulSwigluQuantV2FusionTiling(gert::TilingContext* context) : GroupedMatmulSwigluQuantV2Tiling(context) {}; + + ~GroupedMatmulSwigluQuantV2FusionTiling() override = default; + +protected: + bool IsCapable() override; + + ge::graphStatus DoOpTiling() override; + + uint64_t GetTilingKey() const override; + + ge::graphStatus PostTiling() override; + ge::graphStatus ParseInputAndAttr(); + void FillTilingData() override; + void PrintTilingData() override; +private: + GMMSwigluQuantV2TilingFusionData tilingData_; + uint64_t workspaceSize_; + uint32_t blockDim_; + int64_t k_; + int64_t m_; + int64_t n_; + int32_t groupNum_; + int32_t aicCoreNum_; + int32_t aivCoreNum_; + int64_t ubFactorDimx_; + int64_t groupListType_ = 0; + int8_t isSingleTensor_; +}; + +} +} +#endif // __OP_HOST_OP_TILING_GROUPED_MATMUL_SWIGLU_QUANT_V2_FUSION_TILING_H__ diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_host_utils.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_host_utils.h new file mode 100644 index 000000000..09decd885 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_host_utils.h @@ -0,0 +1,46 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_host_utils.h + * \brief + */ + +#ifndef OP_HOST_GROUPED_MATMUL_SWIGLU_QUANT_V2_HOST_UTILS_H +#define OP_HOST_GROUPED_MATMUL_SWIGLU_QUANT_V2_HOST_UTILS_H + +#include + +namespace GroupedMatmulSwigluQuantParamsV2 { +constexpr uint32_t X_INDEX = 0UL; +constexpr uint32_t PER_TOKEN_SCALE_INDEX = 1UL; +constexpr uint32_t GROUPLIST_INDEX = 2UL; +constexpr uint32_t WEIGHT_INDEX = 3UL; +constexpr uint32_t SCALE_INDEX = 4UL; +constexpr uint32_t Y_DATA_INDEX = 0UL; +constexpr uint32_t Y_SCALE_INDEX = 1UL; +constexpr uint64_t TILING_KEY = 0UL; +constexpr uint64_t ATTR_INDEX_DEQUANT_MODE = 0UL; +constexpr uint32_t ATTR_INDEX_DEQUANT_DTYPE = 1UL; +constexpr uint64_t ATTR_INDEX_QUANT_MODE = 2UL; +constexpr uint32_t ATTR_INDEX_QUANT_DTYPE = 3UL; +constexpr uint64_t ATTR_INDEX_TRANS_W = 4UL; +constexpr uint32_t ATTR_INDEX_GROUP_LIST_TYPE = 5UL; +constexpr size_t PRECHANNEL_WEIGHT_SCALE_DIM = 2UL; +constexpr size_t PERTOKEN_X_SCALE_DIM = 1UL; +constexpr size_t MX_WEIGHT_SCALE_DIM = 4UL; +constexpr size_t MX_X_SCALE_DIM = 3UL; +constexpr size_t MXQuantMode = 2UL; +constexpr uint64_t B4_DATACOPY_MIN_NUM = 2; +constexpr int32_t SPLIT_M = 0; +constexpr uint64_t MXFP4_K_MIN_VALUE = 2UL; +constexpr uint64_t MXFP4_N_MIN_VALUE = 4UL; +} // namespace GroupedMatmulSwigluQuantParamsV2 +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_infershape.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_infershape.cpp new file mode 100644 index 000000000..03b03ec76 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_infershape.cpp @@ -0,0 +1,142 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_proto.cpp + * \brief + */ +#include "register/op_impl_registry.h" +#include "log/log.h" +#include "platform/platform_info.h" +#include "util/math_util.h" +#include "graph/utils/type_utils.h" + +using namespace ge; +namespace ops { +const int64_t X_INDEX = 0; +const int64_t WEIGHT_INDEX = 3; +const int64_t WEIGHTSCALE_DIM_PERTOKEN = 2; +const int64_t WEIGHTSCALE_INDEX = 4; +const int64_t M_DIM_INDEX = 0; +const int64_t DIM_LEN = 2; +const int64_t SPLIT_RATIO = 2; +const int64_t OUT_DIM_LEN = 3; +const int64_t N_SPLIT_RATIO = 128; +constexpr size_t GMMSQ_INDEX_ATTR_QUANT_DTYPE = 3UL; +constexpr size_t GMMSQ_INDEX_ATTR_QUANT_MODE = 2UL; +constexpr size_t QUANT_MODE_MX_TYPE = 2; +constexpr size_t QUANT_MODE_PERTOKEN_TYPE = 0; +constexpr int64_t DYNAMIC_GRAPH_FIRST_INFERSHAPE_DIM_VALUE = -1; + +static std::set GmmDavidSupportSoc = {"Ascend950"}; +static const std::unordered_set DavidSupportedInputDtypes = { + ge::DataType::DT_FLOAT8_E5M2, ge::DataType::DT_FLOAT8_E4M3FN, + ge::DataType::DT_FLOAT4_E2M1, ge::DataType::DT_INT8, ge::DataType::DT_HIFLOAT8}; +bool isSupportedInputDtypeForDavid(ge::DataType dtype) +{ + return DavidSupportedInputDtypes.find(dtype) != DavidSupportedInputDtypes.end(); +} + +static ge::graphStatus InferShape4GroupedMatmulSwigluQuantV2(gert::InferShapeContext *context) +{ + const gert::Shape *xShape = context->GetDynamicInputShape(X_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, xShape); + const gert::Shape *weightScaleShape = context->GetDynamicInputShape(WEIGHTSCALE_INDEX, 0); + OP_CHECK_NULL_WITH_CONTEXT(context, weightScaleShape); + int64_t m = xShape->GetDim(M_DIM_INDEX); + int64_t nDimIndex = weightScaleShape->GetDimNum() - 1; + auto outScaleShape = context->GetOutputShape(1); + OP_CHECK_NULL_WITH_CONTEXT(context, outScaleShape); + if (nDimIndex == OUT_DIM_LEN) { + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const bool *transposeWeightPtr = attrs->GetBool(WEIGHTSCALE_INDEX); + const bool transposeWeight = (transposeWeightPtr != nullptr ? *transposeWeightPtr : false); + nDimIndex = transposeWeight ? weightScaleShape->GetDimNum() - OUT_DIM_LEN : + weightScaleShape->GetDimNum() - DIM_LEN; + int64_t dimValue = static_cast(weightScaleShape->GetDim(nDimIndex)); + int64_t n = 0; + if (dimValue == DYNAMIC_GRAPH_FIRST_INFERSHAPE_DIM_VALUE) { + n = dimValue; + } else { + n = static_cast(Ops::Base::CeilDiv(weightScaleShape->GetDim(nDimIndex), N_SPLIT_RATIO)); + } + outScaleShape->SetDimNum(OUT_DIM_LEN); + outScaleShape->SetDim(0, m); + outScaleShape->SetDim(1, n); + outScaleShape->SetDim(2, SPLIT_RATIO); // 设置outScaleShape的第2维度 + } else { + outScaleShape->SetDimNum(1); + outScaleShape->SetDim(0, m); + } + + int64_t dimValue = static_cast(weightScaleShape->GetDim(nDimIndex)); + int64_t n = 0; + if (dimValue == DYNAMIC_GRAPH_FIRST_INFERSHAPE_DIM_VALUE) { + n = dimValue; + } else { + n = static_cast(weightScaleShape->GetDim(nDimIndex) / SPLIT_RATIO); + if (weightScaleShape->GetDimNum() == WEIGHTSCALE_DIM_PERTOKEN) { + n = static_cast(weightScaleShape->GetDim(1) / SPLIT_RATIO); + } + } + auto outShape = context->GetOutputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context, outShape); + outShape->SetDimNum(DIM_LEN); + outShape->SetDim(0, m); + outShape->SetDim(1, n); + return GRAPH_SUCCESS; +} + +static graphStatus InferDataType4GroupedMatmulSwigluQuantV2(gert::InferDataTypeContext *context) +{ + OP_CHECK_NULL_WITH_CONTEXT(context, context); + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const int64_t* outDtype = attrs->GetInt(GMMSQ_INDEX_ATTR_QUANT_DTYPE); + OP_CHECK_NULL_WITH_CONTEXT(context, outDtype); + const int64_t* quantMode = attrs->GetInt(GMMSQ_INDEX_ATTR_QUANT_MODE); + OP_CHECK_NULL_WITH_CONTEXT(context, quantMode); + + fe::PlatformInfo platformInfo; + fe::OptionalInfo optionalInfo; + auto ret = fe::PlatformInfoManager::Instance().GetPlatformInfoWithOutSocVersion(platformInfo, optionalInfo); + if (ret == GRAPH_SUCCESS && GmmDavidSupportSoc.count(platformInfo.str_info.short_soc_version) > 0) { + auto xDtype = context->GetInputDataType(X_INDEX); + auto weightDtype = context->GetDynamicInputDataType(WEIGHT_INDEX, 0); + OP_CHECK_IF(!isSupportedInputDtypeForDavid(xDtype) || !isSupportedInputDtypeForDavid(weightDtype), + OP_LOGE(context->GetNodeName(), "Invalid Input on this platform, expected FLOAT8_E4M3," + "FLOAT8_E5M2, FLOAT4_E2M1, INT_8, HIFLOAT8, but actual value of x is %s, weight is %s.", + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str(), + ge::TypeUtils::DataTypeToSerialString(weightDtype).c_str()), return GRAPH_FAILED); + + OP_CHECK_IF(*quantMode != QUANT_MODE_MX_TYPE && *quantMode != QUANT_MODE_PERTOKEN_TYPE, + OP_LOGE(context->GetNodeName(), "On this platform, quantMode should be 0(Pertoken) or 2(MX)," + " but actual value is %ld.", *quantMode), return GRAPH_FAILED); + } + auto weightScaleDtype = context->GetDynamicInputDataType(WEIGHTSCALE_INDEX, 0); + if (*quantMode == QUANT_MODE_MX_TYPE) { + if (weightScaleDtype == ge::DataType::DT_FLOAT8_E8M0) { + context->SetOutputDataType(1, DataType::DT_FLOAT8_E8M0); + } else { + OP_LOGE(context->GetNodeName(), "In mx quant mode, quantMode should be 2, but actual value is %ld.", *quantMode); + return GRAPH_FAILED; + } + } else if (*quantMode == QUANT_MODE_PERTOKEN_TYPE) { + context->SetOutputDataType(1, DataType::DT_FLOAT); + } + context->SetOutputDataType(0, static_cast(*outDtype)); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(GroupedMatmulSwigluQuantV2) + .InferShape(InferShape4GroupedMatmulSwigluQuantV2) + .InferDataType(InferDataType4GroupedMatmulSwigluQuantV2); +} // namespace ops diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_tiling.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_tiling.cpp new file mode 100644 index 000000000..7c861af49 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_tiling.cpp @@ -0,0 +1,78 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_tiling.cpp + * \brief + */ + +#include "grouped_matmul_swiglu_quant_v2_tiling.h" +#include +#include +#include "register/op_impl_registry.h" +#include "log/log.h" +#include "err/ops_err.h" +#include "tiling_base/tiling_base.h" +#include "register/op_def_registry.h" +#include "tiling_base/tiling_templates_registry.h" +#include "grouped_matmul_swiglu_quant_v2_fusion_tiling.h" +#include "grouped_matmul_swiglu_quant_v2_base_tiling.h" +#include "platform/platform_infos_def.h" + +using namespace ge; +using namespace AscendC; +using namespace optiling::GroupedMatmulSwigluQuantV2Tiling; +using namespace Ops::Transformer::OpTiling; + +namespace optiling { + +REGISTER_OPS_TILING_TEMPLATE(GroupedMatmulSwigluQuantV2, GroupedMatmulSwigluQuantV2FusionTiling, 0); +REGISTER_OPS_TILING_TEMPLATE(GroupedMatmulSwigluQuantV2, GroupedMatmulSwigluQuantV2BaseTiling, 1); + +static ge::graphStatus GroupedMatmulSwigluQuantV2TilingFunc(gert::TilingContext *context) +{ + OP_CHECK_IF(context == nullptr, + OPS_REPORT_CUBE_INNER_ERR("GroupedMatmulSwigluQuantV2TilingFunc", "Tilingcontext is null"), + return ge::GRAPH_FAILED); + auto compileInfoPtr = context->GetCompileInfo(); + if (compileInfoPtr->supportL12BtBf16) { + std::vector registerList = {2}; + OP_LOGD("GroupedMatmulSwigluQuantV2TilingFunc", "Using the tiling strategy in the mxfp8"); + return TilingRegistry::GetInstance().DoTilingImpl(context, registerList); + }else { + std::vector registerList = {0,1}; + OP_LOGD("GroupedMatmulSwigluQuantV2TilingFunc", "Using the tiling strategy in the int8"); + return TilingRegistry::GetInstance().DoTilingImpl(context, registerList); + } +} + +ASCENDC_EXTERN_C graphStatus TilingPrepareForGMMSwigluQuantV2(gert::TilingParseContext *context) +{ + // get info + auto platformInfoPtr = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfoPtr); + auto compileInfoPtr = context->GetCompiledInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + compileInfoPtr->aicNum_ = ascendcPlatform.GetCoreNumAic(); + compileInfoPtr->aivNum_ = ascendcPlatform.GetCoreNumAiv(); + std::string platformRes; + platformInfoPtr->GetPlatformRes("AICoreintrinsicDtypeMap", "Intrinsic_data_move_l12bt", platformRes); + compileInfoPtr->supportL12BtBf16 = (platformRes.find("bf16") != std::string::npos); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, compileInfoPtr->ubSize_); + OP_LOGD(context->GetNodeName(), "ubSize is %lu, aicNum is %u.", compileInfoPtr->ubSize_, compileInfoPtr->aicNum_); + return GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(GroupedMatmulSwigluQuantV2) + .Tiling(GroupedMatmulSwigluQuantV2TilingFunc) + .TilingParse(TilingPrepareForGMMSwigluQuantV2); +} // namespace optiling diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_tiling.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_tiling.h new file mode 100644 index 000000000..440f51e8c --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/grouped_matmul_swiglu_quant_v2_tiling.h @@ -0,0 +1,174 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_tiling.h + * \brief + */ +#ifndef __OP_HOST_OP_TILING_GROUPED_MATMUL_SWIGLU_QUANT_V2_TILING_H__ +#define __OP_HOST_OP_TILING_GROUPED_MATMUL_SWIGLU_QUANT_V2_TILING_H__ + +#include +#include "tiling_base/tiling_base.h" +#include "tiling/tiling_api.h" + +namespace optiling { + +// GMM 基本信息 +BEGIN_TILING_DATA_DEF(GMMSwigluQuantV2BaseParams) +TILING_DATA_FIELD_DEF(int64_t, groupNum); +TILING_DATA_FIELD_DEF(int64_t, coreNum); +TILING_DATA_FIELD_DEF(int64_t, K); +TILING_DATA_FIELD_DEF(int64_t, N); +TILING_DATA_FIELD_DEF(int64_t, M); +TILING_DATA_FIELD_DEF(int64_t, baseM); +TILING_DATA_FIELD_DEF(int64_t, baseN); +TILING_DATA_FIELD_DEF(int64_t, mLimit); +TILING_DATA_FIELD_DEF(int64_t, workSpaceOffset1); +TILING_DATA_FIELD_DEF(int64_t, workSpaceOffset2); +TILING_DATA_FIELD_DEF(int64_t, quantGroupNum); +TILING_DATA_FIELD_DEF(int64_t, isSingleTensor); +TILING_DATA_FIELD_DEF(int64_t, groupListType); +TILING_DATA_FIELD_DEF(int64_t, smoothScaleDimNum); +TILING_DATA_FIELD_DEF(int64_t, singleN); +TILING_DATA_FIELD_DEF(float, swigluLimit); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(GMMSwigluQuantV2BaseParamsOp, GMMSwigluQuantV2BaseParams) + +// SwigluQuant部分tiling 基本信息 +BEGIN_TILING_DATA_DEF(GMMSwigluQuantV2) +TILING_DATA_FIELD_DEF(int64_t, maxProcessRowNum); +TILING_DATA_FIELD_DEF(int64_t, groupListLen); +TILING_DATA_FIELD_DEF(int64_t, tokenLen); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(GMMSwigluQuantV2Op, GMMSwigluQuantV2) + +// 结构体集合 +BEGIN_TILING_DATA_DEF(GMMSwigluQuantV2TilingData) +TILING_DATA_FIELD_DEF_STRUCT(GMMSwigluQuantV2BaseParams, gmmSwigluQuantV2BaseParams); +TILING_DATA_FIELD_DEF_STRUCT(GMMSwigluQuantV2, gmmSwigluQuantV2); +TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, mmTilingData); +END_TILING_DATA_DEF; + +BEGIN_TILING_DATA_DEF(GMMSwigluQuantV2TilingFusionData) +TILING_DATA_FIELD_DEF(int64_t, cubeBlockDim); +TILING_DATA_FIELD_DEF(int64_t, vectorBlockDim); +TILING_DATA_FIELD_DEF(int64_t, groupNum); +TILING_DATA_FIELD_DEF(int64_t, K); +TILING_DATA_FIELD_DEF(int64_t, N); +TILING_DATA_FIELD_DEF(int64_t, M); +// vector +TILING_DATA_FIELD_DEF(int64_t, ubFactorDimx); +TILING_DATA_FIELD_DEF(int64_t, ubFactorDimy); +TILING_DATA_FIELD_DEF(int64_t, actRight); +TILING_DATA_FIELD_DEF(int64_t, groupListType); +TILING_DATA_FIELD_DEF(int8_t, isSingleTensor); +TILING_DATA_FIELD_DEF(float, swigluLimit); +TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, matmulTiling); +END_TILING_DATA_DEF; + +BEGIN_TILING_DATA_DEF(GMMSwigluQuantParams) +TILING_DATA_FIELD_DEF(uint32_t, groupNum); +TILING_DATA_FIELD_DEF(uint8_t, groupListType); +TILING_DATA_FIELD_DEF(uint8_t, quantDtype); +TILING_DATA_FIELD_DEF(uint8_t, reserved1); +TILING_DATA_FIELD_DEF(uint8_t, dequantDtype); +TILING_DATA_FIELD_DEF(uint32_t, rowLen); +TILING_DATA_FIELD_DEF(uint32_t, ubAvail); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(GMMSwigluQuantParamsOp, GMMSwigluQuantParams) + +BEGIN_TILING_DATA_DEF(GMMSwigluQuantTilingDataParams) +TILING_DATA_FIELD_DEF_STRUCT(GMMSwigluQuantParams, gmmSwigluQuantParams); +TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, mmTilingData); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(GroupedMatmulSwigluQuantV2_0, GMMSwigluQuantTilingDataParams) +REGISTER_TILING_DATA_CLASS(GroupedMatmulSwigluQuantV2_1, GMMSwigluQuantTilingDataParams) + +REGISTER_TILING_DATA_CLASS(GroupedMatmulSwigluQuantV2, GMMSwigluQuantV2TilingData) +REGISTER_TILING_DATA_CLASS(GroupedMatmulSwigluQuantV2_3, GMMSwigluQuantV2TilingFusionData) + +struct GMMSwigluV2CompileInfo { + uint64_t ubSize_ = 0; + uint32_t aicNum_ = 0; + uint32_t aivNum_ = 0; + uint32_t baseM_ = 128; + uint32_t baseN_ = 256; + bool supportL12BtBf16; +}; + +namespace GroupedMatmulSwigluQuantV2Tiling { +constexpr uint32_t X_INDEX = 0; +constexpr uint32_t WEIGHT_INDEX = 3; +constexpr uint32_t WEIGHT_SCALE_INDEX = 4; +constexpr uint32_t GROUPLIST_INDEX = 2; +constexpr uint32_t SMOOTH_SCALE_INDEX = 7; +constexpr uint32_t BATCH_MODE_SCHEDULE = 1; +constexpr uint32_t ATTR_INDEX_DEQUANT_MODE = 0; +constexpr uint32_t ATTR_INDEX_GROUPLIST_TYPE = 5; +constexpr uint32_t ATTR_INDEX_TUNING_CONFIG = 6; +constexpr uint32_t ATTR_INDEX_SWIGLU_LIMIT = 7; +constexpr uint32_t ATTR_INDEX_TRANSPOSE_WEIGHT = 4; +constexpr uint32_t DIM_0 = 0; +constexpr uint32_t DIM_1 = 1; +constexpr uint32_t DIM_2 = 2; +constexpr uint32_t DIM_3 = 3; +constexpr uint32_t DIM_4 = 4; +constexpr uint32_t NUM_FOUR = 4; +constexpr uint32_t NUM_EIGHT = 8; +constexpr uint32_t SYS_WORKSPACE_SIZE = static_cast(16 * 1024 * 1024); +constexpr int64_t USER_WORKSPACE_LIMIT = static_cast(64 * 1024 * 1024); +constexpr int64_t DOUBLE_WORKSPACE_SPLIT = 2; +constexpr int64_t INT32_DTYPE_SIZE = 4; +constexpr int64_t FP32_DTYPE_SIZE = 4; +constexpr int64_t FP32_BLOCK_SIZE = 8; +constexpr int64_t BLOCK_BYTE = 32; +constexpr int64_t SWIGLU_REDUCE_FACTOR = 2; +constexpr int64_t DOUBLE_BUFFER = 2; +constexpr int64_t ND_WEIGHT_DIM_LIMIT = 3; +constexpr int64_t NZ_WEIGHT_DIM_LIMIT = 5; +constexpr int64_t DOUBLE_ROW = 2; +constexpr int64_t PERCHANNEL_WSCALE_DIM_LIMIT = 2; +constexpr int64_t PERGROUP_WSCALE_DIM_LIMIT = 3; +constexpr int64_t A4W4_WEIGHT_NOTRANS_TILING_KEY_MODE = 4; +constexpr int64_t A4W4_WEIGHT_TRANS_TILING_KEY_MODE = 5; +constexpr int64_t A8W8_FUSION_KEY_MODE = 3; +constexpr int64_t A8W4_MSD_TILING_KEY_MODE = 2; +constexpr int64_t SPLITWORKSPACE_TILING_KEY_MODE = 1; +constexpr int64_t COMMON_TILING_KEY_MODE = 0; +constexpr int64_t A8W4_BASEM = 128; +constexpr int64_t A8W4_BASEK = 256; +constexpr int64_t A8W4_BASEN = 256; +constexpr int64_t SIZE_OF_HALF_2 = 2; + +class GroupedMatmulSwigluQuantV2Tiling : public Ops::Transformer::OpTiling::TilingBaseClass { +public: + explicit GroupedMatmulSwigluQuantV2Tiling(gert::TilingContext* context) : Ops::Transformer::OpTiling::TilingBaseClass(context) {}; + + ~GroupedMatmulSwigluQuantV2Tiling() override = default; + +protected: + ge::graphStatus GetPlatformInfo() override {return ge::GRAPH_SUCCESS;}; + + ge::graphStatus GetShapeAttrsInfo() override {return ge::GRAPH_SUCCESS;}; + + ge::graphStatus DoLibApiTiling() override {return ge::GRAPH_SUCCESS;}; + + ge::graphStatus GetWorkspaceSize() override {return ge::GRAPH_SUCCESS;}; + + virtual void FillTilingData() = 0; + virtual void PrintTilingData() = 0; +}; + +} // namespace GroupedMatmulSwigluQuantV2Tiling +} // namespace optiling + +#endif // __OP_HOST_OP_TILING_GROUPED_MATMUL_SWIGLU_QUANT_V2_TILING_H__ \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_v2.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_v2.cpp new file mode 100644 index 000000000..4ed455e76 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_v2.cpp @@ -0,0 +1,208 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include +#include +#include +#include +#include "gmm_dsq_base.h" +#include "grouped_matmul_swiglu_quant_v2_utils.h" +#include "grouped_matmul_swiglu_quant_v2.h" +#include "aclnn_grouped_matmul_swiglu_quant_weight_nz_v2.h" +#include "aclnn_grouped_matmul_swiglu_quant_v2.h" + +using namespace op; +using namespace gmm_dsq; +using namespace gmm_dsq_base; + +class GmmDsqHandlerFactory { +private: + std::unordered_map> handlers_; + +public: + void registerHandler(NpuArch npuArch, std::unique_ptr handler) + { + handlers_[npuArch] = std::move(handler); + } + + GroupedMatmulSwigluQuantHandler *getHandler(NpuArch npuArch) + { + auto it = handlers_.find(npuArch); + return it != handlers_.end() ? it->second.get() : nullptr; + } +}; + +static aclnnStatus aclnnGroupedMatmulSwigluQuantGetWorkspaceSizeCommon(const char* interfaceName, + GroupedMatmulSwigluQuantParamsBase ¶ms, uint64_t *workspaceSize, aclOpExecutor **executor) +{ + GmmDsqHandlerFactory factory; + auto npuArch = op::GetCurrentPlatformInfo().GetCurNpuArch(); + factory.registerHandler(NpuArch::DAV_2201, + std::make_unique()); + factory.registerHandler(NpuArch::DAV_3510, + std::make_unique()); + + if (auto *handler = factory.getHandler(npuArch)) { + handler->Initialize(interfaceName, params, workspaceSize, executor); + return handler->Process(); + } else { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "interfaceName failed: the soc version is not support"); + } + + return ACLNN_ERR_PARAM_INVALID; +} + +#ifdef __cplusplus +extern "C" { +#endif + +aclnnStatus aclnnGroupedMatmulSwigluQuantV2GetWorkspaceSize(const aclTensor *x, + const aclTensorList *weight, const aclTensorList *weightScale, + const aclTensorList *weightAssistMatrix, const aclTensor *bias, + const aclTensor *xScale, const aclTensor *smoothScale, + const aclTensor *groupList, int64_t dequantMode, + int64_t dequantDtype, int64_t quantMode, + int64_t groupListType, const aclIntArray *tuningConfigOptional, double swigluLimit, + aclTensor *output, aclTensor *outputScale, + uint64_t *workspaceSize, aclOpExecutor **executor) +{ + OP_CHECK_COMM_INPUT(workspaceSize, executor); + L2_DFX_PHASE_1(aclnnGroupedMatmulSwigluQuantV2, + DFX_IN(x, weight, weightScale, xScale, groupList), + DFX_OUT(output, outputScale)); + CHECK_COND((output != nullptr), ACLNN_ERR_PARAM_INVALID, + "Expected a proper Tensor but got null for argument output."); + + GroupedMatmulSwigluQuantParamsBase params = + GroupedMatmulSwigluQuantParamsBuilder::Create(x, weight, weightScale, output, outputScale) + .SetXScale(xScale).SetSmoothScale(smoothScale) + .SetGroupList(groupList).SetGroupListType(groupListType) + .SetWeightAssistMatrix(weightAssistMatrix) + .SetDequantAttr(dequantMode, dequantDtype) + .SetQuantAttr(quantMode, static_cast (output->GetDataType())) + .SetTransposeAttr(false).SetBias(bias) + .SetLimitAttr(swigluLimit) + .SetScenario() + .SetTuningConfig(tuningConfigOptional).Build(); + // 调用公共接口 + return aclnnGroupedMatmulSwigluQuantGetWorkspaceSizeCommon(__FUNCTION__, params, workspaceSize, executor); +} + +aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNzV2GetWorkspaceSize(const aclTensor *x, + const aclTensorList *weight, const aclTensorList *weightScale, + const aclTensorList *weightAssistMatrix, const aclTensor *bias, + const aclTensor *xScale, const aclTensor *smoothScale, + const aclTensor *groupList, int64_t dequantMode, + int64_t dequantDtype, int64_t quantMode, + int64_t groupListType, const aclIntArray *tuningConfigOptional, double swigluLimit, + aclTensor *output, aclTensor *outputScale, + uint64_t *workspaceSize, aclOpExecutor **executor) +{ + OP_CHECK_COMM_INPUT(workspaceSize, executor); + L2_DFX_PHASE_1(aclnnGroupedMatmulSwigluQuantWeightNzV2, + DFX_IN(x, weight, weightScale, xScale, groupList), + DFX_OUT(output, outputScale)); + // weight在该场景下强制绑定StorageFormat 和 ViewFormat 为NZ + CHECK_RET(weight != nullptr, ACLNN_ERR_PARAM_NULLPTR); + size_t wLength = weight->Size(); + if (wLength == 1) { + // 单Tensor场景 + auto w = (*weight)[0]; + auto storgeShape = w->GetStorageShape(); + auto viewShape = w->GetViewShape(); + aclTensor *weightNZ = const_cast(w); + auto storageShape = w->GetStorageShape(); + auto groupListViewShape = groupList->GetViewShape(); + auto expertNum = groupListViewShape[0]; + auto weightScale0 = (*weightScale)[0]; + auto weightScaleStorageShape = weightScale0->GetViewShape(); + auto n = weightScaleStorageShape[1]; + auto xViewShape = x->GetViewShape(); + auto k = xViewShape[1]; + storageShape = {expertNum, n / 64, k / 16, 16, 8}; + w->SetStorageShape(storageShape); + CHECK_COND((storgeShape.GetDimNum() == WEIGHT_NZ_DIM_LIMIT), ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwigluQuantWeightNzV2, The dimnum of storageShape for second input (weight)" + "must be 5. \n But StorageShape got %s , and dimNum is %lu.", + op::ToString(storgeShape).GetString(), storgeShape.GetDimNum()); + // weight的StorageFormat无条件视为NZ + weightNZ->SetStorageFormat(op::Format::FORMAT_FRACTAL_NZ); + if (viewShape.GetDimNum() == WEIGHT_NZ_DIM_LIMIT) { + // 若weight的viewShape为5维则视为NZ + weightNZ->SetViewFormat(op::Format::FORMAT_FRACTAL_NZ); + } else if (viewShape.GetDimNum() == WEIGHT_ND_DIM_LIMIT) { + // 若weight的viewShape为3维则视为ND + weightNZ->SetViewFormat(op::Format::FORMAT_ND); + } + } else { + // 多Tensor场景 + for (size_t i = 0; i < wLength; i++) { + auto w = (*weight)[i]; + auto storgeShape = w->GetStorageShape(); + auto viewShape = w->GetViewShape(); + aclTensor *weightNZ = const_cast(w); + auto storageShape = w->GetStorageShape(); + auto groupListViewShape = groupList->GetViewShape(); + auto weightScale0 = (*weightScale)[i]; + auto weightScaleStorageShape = weightScale0->GetViewShape(); + auto n = weightScaleStorageShape[0]; + auto xViewShape = x->GetViewShape(); + auto k = xViewShape[1]; + storageShape = {n / 64, k / 16, 16, 8}; + w->SetStorageShape(storageShape); + CHECK_COND((storgeShape.GetDimNum() == MULTI_WEIGHT_NZ_DIM_LIMIT), ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwigluQuantWeightNzV2, The dimnum of storageShape for second input (weight)" + "must be 4. \n But StorageShape got %s , and dimNum is %lu.", + op::ToString(storgeShape).GetString(), storgeShape.GetDimNum()); + // weight的StorageFormat无条件视为NZ + weightNZ->SetStorageFormat(op::Format::FORMAT_FRACTAL_NZ); + if (viewShape.GetDimNum() == MULTI_WEIGHT_NZ_DIM_LIMIT) { + // 若weight的viewShape为4维则视为NZ + weightNZ->SetViewFormat(op::Format::FORMAT_FRACTAL_NZ); + } else if (viewShape.GetDimNum() == MULTI_WEIGHT_ND_DIM_LIMIT) { + // 若weight的viewShape为2维则视为ND + weightNZ->SetViewFormat(op::Format::FORMAT_ND); + } + } + } + GroupedMatmulSwigluQuantParamsBase params = + GroupedMatmulSwigluQuantParamsBuilder::Create(x, weight, weightScale, output, outputScale) + .SetXScale(xScale).SetSmoothScale(smoothScale) + .SetGroupList(groupList).SetGroupListType(groupListType) + .SetWeightAssistMatrix(weightAssistMatrix) + .SetDequantAttr(dequantMode, dequantDtype) + .SetLimitAttr(swigluLimit) + .SetScenario() + .SetTuningConfig(tuningConfigOptional).Build(); + // 调用公共接口 + return aclnnGroupedMatmulSwigluQuantGetWorkspaceSizeCommon(__FUNCTION__, params, workspaceSize, executor); +} + +aclnnStatus aclnnGroupedMatmulSwigluQuantV2(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, + aclrtStream stream) +{ + L2_DFX_PHASE_2(aclnnGroupedMatmulSwigluQuantV2); + CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER, + "This is an error in GroupedMatmulSwigluQuantV2 launch aicore"); + return ACLNN_SUCCESS; +} + +aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNzV2(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, + aclrtStream stream) +{ + L2_DFX_PHASE_2(aclnnGroupedMatmulSwigluQuantWeightNzV2); + CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER, + "This is an error in GroupedMatmulSwigluQuantWeightNzV2 launch aicore"); + return ACLNN_SUCCESS; +} + +#ifdef __cplusplus +} +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_v2.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_v2.h new file mode 100644 index 000000000..f45c46521 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_v2.h @@ -0,0 +1,76 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_HOST_OP_API_ACLNN_GROUPED_MATMUL_SWIGLU_QUANT_V2_H +#define OP_HOST_OP_API_ACLNN_GROUPED_MATMUL_SWIGLU_QUANT_V2_H +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @brief aclnnGroupedMatmulSwigluQuantV2 的第一段接口，根据具体的计算流程，计算workspace大小。 + * @domain aclnn_ops_infer + * + * @param [in] x: 表示公式中的x，数据类型支持INT8、FLOAT4_E2M1、FLOAT8_E4M3FN、FLOAT8_E5M2、HIFLOAT8数据类型，数据格式支持ND。 + * @param [in] weight: + * 表示公式中的weight，数据类型支持INT4、FLOAT4_E2M1、FLOAT8_E4M3FN、FLOAT8_E5M2、INT8、HIFLOAT8数据类型，数据格式支持ND。 + * @param [in] weightScale: + * 表示量化参数，数据类型支持UINT64、FLOAT32、FLOAT8_E8M0、BF16、FLOAT16数据类型，数据格式支持ND。 + * @param [in] weightAssistMatrix: + * 表示weight辅助矩阵，数据类型支持FLOAT32数据类型。 + * @param [in] bias: + * 表示偏移，数据类型支持FLOAT32数据类型，数据格式支持ND。 + * @param [in] xScale: + * 表示perToken量化参数，数据类型支持FLOAT8_E8M0、FLOAT32数据类型，数据格式支持ND。 + * @param [in] smoothScale: + * 左矩阵的的量化因子，数据类型支持FLOAT32数据类型，数据格式支持ND。 + * @param [in] groupList: 必选参数，表示每个分组参与计算的Token个数，数据类型支持INT64。 + * @param [in] dequantMode: 表示反量化计算类型，用于确定激活矩阵与权重矩阵的反量化方式。 + * @param [in] dequantDtype: 表示中间GroupedMatmul的结果数据类型。 + * @param [in] quantMode: 表示量化计算类型，用于确定swiglu结果的量化模式。 + * @param [in] groupListType: 表示指定分组的解释方式，用于确定groupList的语义。 + * @param [in] tuningConfig: 用于算子预估m/e的大小，走不同的算子模板，以适配不不同场景性能要求。 + * @param [in] swigluLimit: clamp。 + * @param [out] quantOutput: 表示公式中的out，数据类型支持INT8、FLOAT4_E2M1、FLOAT8_E4M3FN、FLOAT8_E5M2、HIFLOAT8数据类型，数据格式支持ND。 + * @param [out] quantScaleOutput: 表示公式中的outQuantScale，数据类型支持FLOAT32、FLOAT8_E8M0数据类型。 + * @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。 + * @param [out] executor: 返回op执行器，包含算子计算流程。 + * @return aclnnStatus: 返回状态码。 + */ +aclnnStatus aclnnGroupedMatmulSwigluQuantV2GetWorkspaceSize(const aclTensor *x, + const aclTensorList *weight, const aclTensorList *weightScale, + const aclTensorList *weightAssistMatrix, const aclTensor *bias, + const aclTensor *xScale, const aclTensor *smoothScale, + const aclTensor *groupList, int64_t dequantMode, + int64_t dequantDtype, int64_t quantMode, int64_t groupListType, + const aclIntArray *tuningConfigOptional, double swigluLimit, + aclTensor *output, aclTensor *outputScale, + uint64_t *workspaceSize, aclOpExecutor **executor); + +/** + * @brief aclnnGroupedMatmulSwigluQuantV2的第二段接口，用于执行计算。 + * @param [in] workspace: 在npu device侧申请的workspace内存起址。 + * @param [in] workspaceSize: 在npu + * device侧申请的workspace大小，由第一段接口aclnnGroupedMatmulSwigluQuantV2GetWorkspaceSize获取。 + * @param [in] stream: acl stream流。 + * @param [in] executor: op执行器，包含了算子计算流程。 + * @return aclnnStatus: 返回状态码。 + */ +__attribute__((visibility("default"))) aclnnStatus aclnnGroupedMatmulSwigluQuantV2(void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_weight_nz_v2.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_weight_nz_v2.h new file mode 100644 index 000000000..9efcc913e --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_weight_nz_v2.h @@ -0,0 +1,75 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP__HOST_OP_API_ACLNN_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_V2_H +#define OP__HOST_OP_API_ACLNN_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_V2_H +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @brief aclnnGroupedMatmulSwigluQuantWeightNzV2 的第一段接口，根据具体的计算流程，计算workspace大小。 + * @domain aclnn_ops_infer + * + * @param [in] x: 表示公式中的x，数据类型支持INT8数据类型，数据格式支持ND。 + * @param [in] weight: + * 表示公式中的weight，数据类型支持INT8、INT4数据类型，数据格式支持NZ。 + * @param [in] weightScale: + * 表示量化参数，数据类型支持FLOAT32、UINT64数据类型，数据格式支持ND。 + * @param [in] weightAssistMatrix: + * 表示weight辅助矩阵，数据类型支持FLOAT32数据类型。 + * @param [in] bias: + * 表示偏移，数据类型支持FLOAT32数据类型，数据格式支持ND。 + * @param [in] xScale: + * 表示perToken量化参数，数据类型支持FLOAT8_E8M0数据类型，数据格式支持ND。 + * @param [in] smoothScale: + * 左矩阵的的量化因子，数据类型支持FLOAT32数据类型，数据格式支持ND。 + * @param [in] groupList: 必选参数，代表输入和输出分组轴上的索引情况，数据类型支持INT64。 + * @param [in] dequantMode: 表示反量化计算类型，用于确定激活矩阵与权重矩阵的反量化方式。 + * @param [in] dequantDtype: 表示中间GroupedMatmul的结果数据类型。 + * @param [in] quantMode: 表示量化计算类型，用于确定swiglu结果的量化模式。 + * @param [in] groupListType: 表示指定分组的解释方式，用于确定groupList的语义。 + * @param [in] tuningConfig: 用于算子预估m/e的大小，走不同的算子模板，以适配不不同场景性能要求。 + * @param [out] quantOutput: 表示公式中的out，数据类型支持INT8、FLOAT8_E4M3FN、FLOAT8_E5M2数据类型，数据格式支持ND。 + * @param [out] quantScaleOutput: 表示公式中的outQuantScale，数据类型支持FLOAT32、FLOAT8_E8M0数据类型。 + * @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。 + * @param [out] executor: 返回op执行器，包含算子计算流程。 + * @return aclnnStatus: 返回状态码。 + */ +aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNzV2GetWorkspaceSize(const aclTensor *x, + const aclTensorList *weight, const aclTensorList *weightScale, + const aclTensorList *weightAssistMatrix, const aclTensor *bias, + const aclTensor *xScale, const aclTensor *smoothScale, + const aclTensor *groupList, int64_t dequantMode, + int64_t dequantDtype, int64_t quantMode, int64_t groupListType, + const aclIntArray *tuningConfigOptional, double swigluLimit, + aclTensor *output, aclTensor *outputScale, + uint64_t *workspaceSize, aclOpExecutor **executor); + +/** + * @brief aclnnGroupedMatmulSwigluQuantWeightNzV2的第二段接口，用于执行计算。 + * @param [in] workspace: 在npu device侧申请的workspace内存起址。 + * @param [in] workspaceSize: 在npu + * device侧申请的workspace大小，由第一段接口aclnnGroupedMatmulSwigluQuantWeightNzV2GetWorkspaceSize获取。 + * @param [in] stream: acl stream流。 + * @param [in] executor: op执行器，包含了算子计算流程。 + * @return aclnnStatus: 返回状态码。 + */ +__attribute__((visibility("default"))) aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNzV2(void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/gmm_dsq_base.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/gmm_dsq_base.h new file mode 100644 index 000000000..b81adf02c --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/gmm_dsq_base.h @@ -0,0 +1,682 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_HOST_OP_API_ACLNN_GMM_DSQ_BASE_H +#define OP_HOST_OP_API_ACLNN_GMM_DSQ_BASE_H + +#include "grouped_matmul_swiglu_quant_utils.h" + +namespace gmm_dsq_base { + +using namespace gmm_dsq; + +constexpr int64_t SPLIT = 2L; +constexpr int64_t K_LIMIT_A8W8 = 65536L; +constexpr int64_t K_LIMIT_A8W4 = 20000L; +constexpr int64_t N_LIMIT = 10240L; +constexpr int64_t NZ_DIM_4_INT8 = 32L; +constexpr int64_t NZ_DIM_4_INT4 = 64L; +constexpr int64_t NZ_DIM_3 = 16L; +constexpr int64_t OUTPUT_IDX_0 = 0L; +constexpr int64_t OUTPUT_IDX_1 = 1L; +constexpr int64_t DIM_IDX_0 = 0L; +constexpr int64_t DIM_IDX_1 = 1L; +constexpr int64_t DIM_IDX_2 = 2L; +constexpr int64_t DIM_IDX_3 = 4L; +constexpr size_t X_DIM_LIMIT = 2UL; +constexpr size_t MULTI_WEIGHT_NZ_DIM_LIMIT = 4UL; +constexpr size_t MULTI_WEIGHT_ND_DIM_LIMIT = 2UL; +constexpr size_t WEIGHT_SCALE_DIM_LIMIT = 2UL; +constexpr size_t SINGLE_WEIGHT_SCALE_PERGROUP_DIM_LIMIT = 3UL; +constexpr size_t SINGLE_WEIGHT_SCALE_PERCHANNEL_DIM_LIMIT = 2UL; +constexpr size_t MULTI_WEIGHT_SCALE_PERGROUP_DIM_LIMIT = 2UL; +constexpr size_t MULTI_WEIGHT_SCALE_PERCHANNEL_DIM_LIMIT = 1UL; +constexpr size_t TOKEN_SCALE_DIM_LIMIT = 1UL; +constexpr size_t SINGLE_WEIGHT_ASSIST_MATRIX_DIM_LIMIT = 2UL; +constexpr size_t MULTI_WEIGHT_ASSIST_MATRIX_DIM_LIMIT = 1UL; +constexpr size_t GROUP_LIST_DIM_LIMIT = 1UL; +constexpr size_t QUANTOUT_DIM_LIMIT = 2UL; +constexpr size_t QUANTSCALEOUT_DIM_LIMIT = 1UL; +constexpr size_t INT4_PER_INT32 = 8UL; +constexpr size_t NZ_ALIGN_K = 16UL; +constexpr size_t NZ_ALIGN_N = 32UL; +constexpr size_t SMOOTH_SCALE_1D_DIM_LIMIT = 1UL; +constexpr size_t SMOOTH_SCALE_2D_DIM_LIMIT = 2UL; + +const std::initializer_list X_DTYPE_SUPPORT_LIST = {DataType::DT_INT8, DataType::DT_INT4}; +const std::initializer_list WEIGHT_DTYPE_SUPPORT_LIST = {DataType::DT_INT8, DataType::DT_INT4}; +const std::initializer_list WEIGHT_SCALE_DTYPE_SUPPORT_LIST = { + DataType::DT_FLOAT, DataType::DT_FLOAT16, DataType::DT_BF16}; +const std::initializer_list WEIGHT_SCALE_A8W4_DTYPE_SUPPORT_LIST = {DataType::DT_UINT64}; +const std::initializer_list X_SCALE_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT}; +const std::initializer_list GROUP_LIST_DTYPE_SUPPORT_LIST = {DataType::DT_INT64}; +const std::initializer_list QUANTOUT_DTYPE_SUPPORT_LIST = {DataType::DT_INT8}; +const std::initializer_list QUANTSCALEOUT_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT}; +const std::initializer_list WEIGHT_ASSIST_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT}; +const std::initializer_list SMOOTH_SCALE_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT}; + +class GroupedMatmulSwigluQuantBaseHandler : public GroupedMatmulSwigluQuantHandler { +protected: + bool CheckInputOutDimsA8W8() + { + OP_CHECK_WRONG_DIMENSION(gmmDsqParams_.x, X_DIM_LIMIT, return false); + size_t wLength = gmmDsqParams_.weight->Size(); + for (size_t i = 0; i < wLength; i++) { + const aclTensor* w = (*gmmDsqParams_.weight)[i]; + const aclTensor* wScale = (*gmmDsqParams_.weightScale)[i]; + op::Format wFormat = w->GetViewFormat(); + if (wLength == static_cast(1)) { // 单Tensor场景 + if (IsPrivateFormat(wFormat)) { + OP_CHECK_WRONG_DIMENSION(w, WEIGHT_NZ_DIM_LIMIT, return false); + } else { + OP_CHECK_WRONG_DIMENSION(w, WEIGHT_ND_DIM_LIMIT, return false); + } + OP_CHECK_WRONG_DIMENSION(wScale, WEIGHT_SCALE_DIM_LIMIT, return false); + } else { // 多Tensor场景 + if (IsPrivateFormat(wFormat)) { + OP_CHECK_WRONG_DIMENSION(w, MULTI_WEIGHT_NZ_DIM_LIMIT, return false); + } else { + OP_CHECK_WRONG_DIMENSION(w, MULTI_WEIGHT_ND_DIM_LIMIT, return false); + } + OP_CHECK_WRONG_DIMENSION(wScale, MULTI_WEIGHT_SCALE_PERCHANNEL_DIM_LIMIT, return false); + } + } + + OP_CHECK_WRONG_DIMENSION(gmmDsqParams_.xScale, TOKEN_SCALE_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(gmmDsqParams_.groupList, GROUP_LIST_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(gmmDsqParams_.output, QUANTOUT_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(gmmDsqParams_.outputScale, QUANTSCALEOUT_DIM_LIMIT, return false); + return true; + } + + bool CheckInputOutDimsA4W4orA8W4() + { + OP_CHECK_WRONG_DIMENSION(gmmDsqParams_.x, X_DIM_LIMIT, return false); + if (gmmDsqParams_.isA4W4 && gmmDsqParams_.weightAssistMatrix != nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "In the A4W4 scenario, the weightAssistMatrix input must be nullptr."); + return false; + } else if (gmmDsqParams_.isA8W4 && gmmDsqParams_.weightAssistMatrix == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "In the A8W4 scenario, the weightAssistMatrix input must not be nullptr."); + return false; + } + size_t wLength = gmmDsqParams_.weight->Size(); + for (size_t i = 0; i < wLength; i++) { + const aclTensor* w = (*gmmDsqParams_.weight)[i]; + const aclTensor* wScale = (*gmmDsqParams_.weightScale)[i]; + op::Format weightViewFormat = w->GetViewFormat(); + bool isSingle = (wLength == 1); + // 检查权重维度 + OP_CHECK_WRONG_DIMENSION(w, + (isSingle ? + (IsPrivateFormat(weightViewFormat) ? WEIGHT_NZ_DIM_LIMIT : WEIGHT_ND_DIM_LIMIT) : + (IsPrivateFormat(weightViewFormat) ? MULTI_WEIGHT_NZ_DIM_LIMIT : MULTI_WEIGHT_ND_DIM_LIMIT)), + return false); + // 检查权重Scale维度 + OP_CHECK_WRONG_DIMENSION(wScale, + (isSingle ? + (gmmDsqParams_.dequantMode == 0 ? SINGLE_WEIGHT_SCALE_PERCHANNEL_DIM_LIMIT : SINGLE_WEIGHT_SCALE_PERGROUP_DIM_LIMIT) : + (gmmDsqParams_.dequantMode == 0 ? MULTI_WEIGHT_SCALE_PERCHANNEL_DIM_LIMIT : MULTI_WEIGHT_SCALE_PERGROUP_DIM_LIMIT)), + return false); + // 检查辅助矩阵（A8W4模式） + if (gmmDsqParams_.isA8W4) { + const aclTensor* weightAssistMatrix = (*gmmDsqParams_.weightAssistMatrix)[i]; + OP_CHECK_WRONG_DIMENSION(weightAssistMatrix, + (isSingle ? SINGLE_WEIGHT_ASSIST_MATRIX_DIM_LIMIT : MULTI_WEIGHT_ASSIST_MATRIX_DIM_LIMIT), + return false); + } + } + OP_CHECK_WRONG_DIMENSION(gmmDsqParams_.xScale, TOKEN_SCALE_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(gmmDsqParams_.groupList, GROUP_LIST_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(gmmDsqParams_.output, QUANTOUT_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(gmmDsqParams_.outputScale, QUANTSCALEOUT_DIM_LIMIT, return false); + return true; + } + + bool CheckSingleTensorListTypeA8W8(int64_t e, int64_t k, int64_t n) + { + // weight的NDshape期望为[E, K, N] + op::Shape weightNDExpectShape1 = {e, k, n}; + // 单tesnsor weight的NZshape期望为[E, N // 32, K // 16, 16, 32] + op::Shape weightNZExpectShape1 = {e, static_cast(n / NZ_DIM_4_INT8), static_cast(k / NZ_DIM_3), + NZ_DIM_3, NZ_DIM_4_INT8}; + // weight的NDshape期望为[K, N] + op::Shape weightNDExpectShape2 = {k, n}; + // weight的NZshape期望为[N // 32, K // 16, 16, 32] + op::Shape weightNZExpectShape2 = {static_cast(n / NZ_DIM_4_INT8), static_cast(k / NZ_DIM_3), + NZ_DIM_3, NZ_DIM_4_INT8}; + + // weightScale的shape期望为[E, N] + op::Shape weightScaleExpectShape1 = {e, n}; + op::Shape weightScaleExpectShape2 = {n}; + + const aclTensor* w = (*gmmDsqParams_.weight)[0]; + const aclTensor* wScale = (*gmmDsqParams_.weightScale)[0]; + + op::Format wFormat = w->GetViewFormat(); + op::Format storageFormat = w->GetStorageFormat(); + if (IsPrivateFormat(wFormat)) { + if (!(w->GetViewShape() == weightNZExpectShape1 || w->GetViewShape() == weightNZExpectShape2)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected tensor for weight to have same size as %s or %s, but got %s.", + op::ToString(weightNZExpectShape1).GetString(), + op::ToString(weightNZExpectShape2).GetString(), + op::ToString(w->GetViewShape()).GetString()); + return false; + } + } else { + if (!(w->GetViewShape() == weightNDExpectShape1 || w->GetViewShape() == weightNDExpectShape2)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected tensor for weight to have same size as %s or %s, but got %s.", + op::ToString(weightNDExpectShape1).GetString(), + op::ToString(weightNDExpectShape2).GetString(), + op::ToString(w->GetViewShape()).GetString()); + return false; + } + + if (IsPrivateFormat(storageFormat) && (k % NZ_ALIGN_K != 0 || n % NZ_ALIGN_N != 0)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "In W8a8 Nz mode, k should align to 16, n align to 32"); + return false; + } + } + + if (!(wScale->GetViewShape() == weightScaleExpectShape1 || wScale->GetViewShape() == weightScaleExpectShape2)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected tensor for weight_scale to have same size as %s or %s, but got %s.", + op::ToString(weightScaleExpectShape1).GetString(), + op::ToString(weightScaleExpectShape2).GetString(), + op::ToString(wScale->GetViewShape()).GetString()); + return false; + } + return true; + } + + bool CheckMultiTensorTypeA8W8(int64_t k, int64_t n) + { + // weight的NDshape期望为[K, N] + op::Shape weightNDExpectShape = {k, n}; + // weight的NZshape期望为[N // 32, K // 16, 16, 32] + op::Shape weightNZExpectShape = {static_cast(n / NZ_DIM_4_INT8), static_cast(k / NZ_DIM_3), + NZ_DIM_3, NZ_DIM_4_INT8}; + + // weightScale的shape期望为[N] + op::Shape weightScaleExpectShape = {n}; + size_t wLength = gmmDsqParams_.weight->Size(); + + for (size_t i = 0; i < wLength; i++) { + const aclTensor* w = (*gmmDsqParams_.weight)[0]; + const aclTensor* wScale = (*gmmDsqParams_.weightScale)[0]; + op::Format wFormat = w->GetViewFormat(); + op::Format storageFormat = w->GetStorageFormat(); + if (IsPrivateFormat(wFormat)) { + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(w, weightNZExpectShape, return false); + } else { + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(w, weightNDExpectShape, return false); + if (IsPrivateFormat(storageFormat) && (k % NZ_ALIGN_K != 0 || n % NZ_ALIGN_N != 0)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "In W8a8 Nz mode, k should align to 16, n align to 32"); + return false; + } + } + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(wScale, weightScaleExpectShape, return false); + } + + return true; + } + + bool CheckTensorListShapeA8W8(int64_t e,int64_t k, int64_t n) + { + size_t wLength = gmmDsqParams_.weight->Size(); + if (wLength == static_cast(1)) { + return CheckSingleTensorListTypeA8W8(e, k, n); + } + + return CheckMultiTensorTypeA8W8(k, n); + } + + bool CheckInputOutShapeA8W8() + { + int64_t m = gmmDsqParams_.x->GetViewShape().GetDim(0); + int64_t k = gmmDsqParams_.x->GetViewShape().GetDim(1); + auto n_index = ((*gmmDsqParams_.weightScale)[0])->GetViewShape().GetDimNum() - 1; + int64_t n = ((*gmmDsqParams_.weightScale)[0])->GetViewShape().GetDim(n_index); + size_t wLength = gmmDsqParams_.weight->Size(); + int64_t e = wLength; + if (wLength == static_cast(1)) { + e = ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(0); + } + if (n % SPLIT != 0) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "%s, N is %ld , not an even number.", interfaceName_.c_str(), n); + return false; + } + int64_t nAfterHalve = static_cast(n / SPLIT); + // x的shape期望为[M, K] + op::Shape xExpectShape = {m, k}; + // xScale的shape期望为[E, N] + op::Shape xScaleExpectShape = {m}; + // output的shape期望为[M, N / 2] + op::Shape outputExpectShape = {m, nAfterHalve}; + // outputScale的shape期望为[M] + op::Shape outputScaleExpectShape = {m}; + + auto ret = CheckTensorListShapeA8W8(e, k, n); + if (!ret) { + return false; + } + + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(gmmDsqParams_.x, xExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(gmmDsqParams_.xScale, xScaleExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(gmmDsqParams_.output, outputExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(gmmDsqParams_.outputScale, outputScaleExpectShape, return false); + // groupList的长度应小于等于weight的专家数 + int64_t groupListLen = gmmDsqParams_.groupList->GetViewShape().GetDim(0); + if (groupListLen > e) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "%s A8W8, Length of 'groupList' out of range (expected to be in range of [1, " + "%ld], but got %ld)", interfaceName_.c_str(), + e, groupListLen); + return false; + } + if (n > N_LIMIT) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "%s A8W8: The current version does not support the scenario that " + "N(%ld) is greater than %ld.", interfaceName_.c_str(), + n, N_LIMIT); + return false; + } + if (k >= K_LIMIT_A8W8) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "%s A8W8, The current version does not support the scenario." + "The tail axis dimension of input0(x) is %ld, which need lower than %ld.", + interfaceName_.c_str(), k, K_LIMIT_A8W8); + return false; + } + if (gmmDsqParams_.smoothScale != nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "%s, smoothScale must be nullptr in A8W8 scenario.", interfaceName_.c_str()); + return false; + } + return true; + } + + bool CheckSingleTensorListTypeA8W4orA4W4(int64_t e, int64_t k, int64_t n) + { + // weight的NDshape期望为[E, K, N] + op::Shape weightNDExpectShape = {e, k, n}; + // 单tesnsor weight的NZshape期望为[E, N // 64, K // 16, 16, 64] + op::Shape weightNZExpectShape = {e, static_cast(n / NZ_DIM_4_INT4), static_cast(k / NZ_DIM_3), + NZ_DIM_3, NZ_DIM_4_INT4}; + // 单tensor NZ转置 + op::Shape weightNZTransposeExpectShape1 = {e, static_cast(k / NZ_DIM_4_INT4), static_cast(n / NZ_DIM_3), + NZ_DIM_4_INT4, NZ_DIM_3}; + op::Shape weightNZTransposeExpectShape2 = {e, static_cast(k / NZ_DIM_4_INT4), static_cast(n / NZ_DIM_3), + NZ_DIM_3, NZ_DIM_4_INT4}; + + // 辅助矩阵的shape期望为[E, N] + op::Shape weightAssistMatrixExpectShape = {e, n}; + + const aclTensor* w = (*gmmDsqParams_.weight)[0]; + const aclTensor* weightAssistMatrix = nullptr; + if (gmmDsqParams_.weightAssistMatrix != nullptr && (*gmmDsqParams_.weightAssistMatrix)[0] != nullptr) { + weightAssistMatrix = (*gmmDsqParams_.weightAssistMatrix)[0]; + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weightAssistMatrix, weightAssistMatrixExpectShape, return false); + } + op::Format weightViewFormat = w->GetViewFormat(); + if (IsPrivateFormat(weightViewFormat)) { + if (!(w->GetViewShape() == weightNZExpectShape || w->GetViewShape() == weightNZTransposeExpectShape1 || w->GetViewShape() == weightNZTransposeExpectShape2)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Expected tensor for weight to have same size as %s %s or %s, but got %s.", + op::ToString(weightNZExpectShape).GetString(), + op::ToString(weightNZTransposeExpectShape1).GetString(), + op::ToString(weightNZTransposeExpectShape2).GetString(), + op::ToString(w->GetViewShape()).GetString()); + return false; + } + } else { + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(w, weightNDExpectShape, return false); + } + return true; + } + + bool CheckMultiTensorTypeA8W4orA4W4(int64_t k, int64_t n) + { + // weight的NDshape期望为[K, N] + op::Shape weightNDExpectShape = {k, n}; + // weight的NZshape期望为[N // 64, K // 16, 16, 64] + op::Shape weightNZExpectShape = {static_cast(n / NZ_DIM_4_INT4), static_cast(k / NZ_DIM_3), + NZ_DIM_3, NZ_DIM_4_INT4}; + + op::Shape weightAssistMatrixExpectShape = {n}; + size_t wLength = gmmDsqParams_.weight->Size(); + + for (size_t i = 0; i < wLength; i++) { + const aclTensor* w = (*gmmDsqParams_.weight)[i]; + const aclTensor* weightAssistMatrix = nullptr; + if (gmmDsqParams_.weightAssistMatrix != nullptr && (*gmmDsqParams_.weightAssistMatrix)[i] != nullptr) { + weightAssistMatrix = (*gmmDsqParams_.weightAssistMatrix)[i]; + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weightAssistMatrix, weightAssistMatrixExpectShape, return false); + } + op::Format weightViewFormat = w->GetViewFormat(); + if (IsPrivateFormat(weightViewFormat)) { + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(w, weightNZExpectShape, return false); + } else { + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(w, weightNDExpectShape, return false); + } + } + + return true; + } + + bool CheckSmoothScaleA4W4(int64_t e, int64_t nAfterHalve) + { + if (gmmDsqParams_.smoothScale == nullptr) { + return true; + } + OP_CHECK_DTYPE_NOT_SUPPORT(gmmDsqParams_.smoothScale, SMOOTH_SCALE_DTYPE_SUPPORT_LIST, return false); + size_t dimNum = gmmDsqParams_.smoothScale->GetViewShape().GetDimNum(); + if (dimNum == SMOOTH_SCALE_1D_DIM_LIMIT) { + op::Shape expectShape = {e}; + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(gmmDsqParams_.smoothScale, expectShape, return false); + } else if (dimNum == SMOOTH_SCALE_2D_DIM_LIMIT) { + op::Shape expectShape = {e, nAfterHalve}; + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(gmmDsqParams_.smoothScale, expectShape, return false); + } else { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "%s, smoothScale dimNum should be 1 or 2 in A4W4 scenario, but got %lu.", + interfaceName_.c_str(), dimNum); + return false; + } + return true; + } + + bool CheckTensorListShapeA8W4orA4W4(int64_t e, int64_t k, int64_t n) + { + size_t wLength = gmmDsqParams_.weight->Size(); + if (wLength == static_cast(1)) { + return CheckSingleTensorListTypeA8W4orA4W4(e, k, n); + } + + return CheckMultiTensorTypeA8W4orA4W4(k, n); + } + + bool CheckInputOutShapeA8W4orA4W4() + { + int64_t m = gmmDsqParams_.x->GetViewShape().GetDim(0); + int64_t k = gmmDsqParams_.x->GetViewShape().GetDim(1); + int64_t e = 1; + int64_t n = 1; + int64_t KGroupCount = 1; // K轴的组数，perchannel场景相当于pergroup场景中的组数为1 + int64_t KGroupSize = k; // K轴每组的元素个数 + op::Shape weightScaleExpectShape; + size_t wLength = gmmDsqParams_.weight->Size(); + if (gmmDsqParams_.dequantMode == 0 && wLength == static_cast(1)) { + e = ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(0); + // weightScale入参在perchannel单tensor场景期望shape [E, N] + n = ((*gmmDsqParams_.weightScale)[0])->GetViewShape().GetDim(DIM_IDX_1); + weightScaleExpectShape = {e, n}; // 单 + } else if (gmmDsqParams_.dequantMode == 0 && wLength != static_cast(1)) { + e = wLength; + // weightScale入参在perchannel多tensor场景期望shape [N] + n = ((*gmmDsqParams_.weightScale)[0])->GetViewShape().GetDim(DIM_IDX_0); + weightScaleExpectShape = {n}; // 多 + } else if (gmmDsqParams_.dequantMode == 1 && wLength == static_cast(1)) { + e = ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(0); + // weightScale入参在pergroup单tensor场景期望shape [E, KGroupCount, N] + n = ((*gmmDsqParams_.weightScale)[0])->GetViewShape().GetDim(DIM_IDX_2); + KGroupCount = ((*gmmDsqParams_.weightScale)[0])->GetViewShape().GetDim(DIM_IDX_1); + KGroupSize = KGroupCount > 0 ? k / KGroupCount : k; + weightScaleExpectShape = {e, KGroupCount, n}; // 单 + } else if (gmmDsqParams_.dequantMode == 1 && wLength != static_cast(1)) { + e = wLength; + // weightScale入参在pergroup多tensor场景期望shape [KGroupCount, N] + n = ((*gmmDsqParams_.weightScale)[0])->GetViewShape().GetDim(DIM_IDX_1); + KGroupCount = ((*gmmDsqParams_.weightScale)[0])->GetViewShape().GetDim(DIM_IDX_0); + KGroupSize = KGroupCount > 0 ? k / KGroupCount : k; + weightScaleExpectShape = {KGroupCount, n}; // 多 + } + if (KGroupCount == 0 || k % KGroupCount != 0) { + OP_LOGE( + ACLNN_ERR_PARAM_INVALID, + "%s, " + "The number of groups along the k-axis is %ld, and the length of the k-axis is %ld, which is illegal. " + "The number of groups must be greater than 0, and k-axis length %% number of groups == 0 must be true.", + interfaceName_.c_str(), KGroupCount, k); + return false; + } + if (n % SPLIT != 0) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "%s, N is %ld , not an even number.", interfaceName_.c_str(), n); + return false; + } + int64_t nAfterHalve = static_cast(n / SPLIT); + // x的shape期望为[M, K] + op::Shape xExpectShape = {m, k}; + // xScale的shape期望为[E, N] + op::Shape xScaleExpectShape = {m}; + // output的shape期望为[M, N / 2] + op::Shape outputExpectShape = {m, nAfterHalve}; + // outputScale的shape期望为[M] + op::Shape outputScaleExpectShape = {m}; + auto ret = CheckTensorListShapeA8W4orA4W4(e, k, n); + if (!ret) { + return false; + } + + for (size_t i = 0; i < wLength; i++) { + const aclTensor* wScale = (*gmmDsqParams_.weightScale)[i]; + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(wScale, weightScaleExpectShape, return false); + } + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(gmmDsqParams_.x, xExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(gmmDsqParams_.xScale, xScaleExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(gmmDsqParams_.output, outputExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(gmmDsqParams_.outputScale, outputScaleExpectShape, return false); + // groupList的长度应小于等于weight的专家数 + int64_t groupListLen = gmmDsqParams_.groupList->GetViewShape().GetDim(0); + if (groupListLen > e) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "%s A8W4 or A4W4, Length of 'groupList' out of range (expected to be in range of [1, " + "%ld], but got %ld)", interfaceName_.c_str(), + e, groupListLen); + return false; + } + if (n > N_LIMIT) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "%s A8W4 or A4W4: The current version does not support the scenario that " + "N(%ld) is greater than %ld.", interfaceName_.c_str(), + n, N_LIMIT); + return false; + } + if (k >= K_LIMIT_A8W4) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "%s A8W4 or A4W4, The current version does not support the scenario." + "The tail axis dimension of input0(x) is %ld, which need lower than %ld.", + interfaceName_.c_str(), k, K_LIMIT_A8W4); + return false; + } + if (gmmDsqParams_.isA4W4) { + if (!CheckSmoothScaleA4W4(e, nAfterHalve)) { + return false; + } + } else if (gmmDsqParams_.isA8W4 && gmmDsqParams_.smoothScale != nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "%s, smoothScale must be nullptr in A8W4 scenario.", interfaceName_.c_str()); + return false; + } + (void)KGroupSize; + return true; + } + + bool IsTransposeLastTwoDims(const aclTensor *tensor) + { + auto shape = tensor->GetViewShape(); + int64_t dim1 = shape.GetDimNum() - 1; + int64_t dim2 = shape.GetDimNum() - 2; + auto strides = tensor->GetViewStrides(); + if (strides[dim2] == 1 && strides[dim1] == shape.GetDim(dim2)) { + int64_t tmpNxD = shape.GetDim(dim1) * shape.GetDim(dim2); + for (int64_t batchDim = shape.GetDimNum() - 3; batchDim >= 0; batchDim--) { + if (strides[batchDim] != tmpNxD) { + return false; + } + tmpNxD *= shape.GetDim(batchDim); + } + return true; + } + return false; + } + + void UnpackInt32ToInt4(const aclTensor *&tensorS32, const std::string &tensorType) + { + OP_LOGD("Unpack %s from int32 to int4 start.", tensorType.c_str()); + auto tensorS4 = const_cast(tensorS32); + op::Shape tensorShape = tensorS4->GetViewShape(); + auto viewShapeDim = tensorShape.GetDimNum(); + op::Strides newStride = tensorS4->GetViewStrides(); + bool transposeTensor = false; + auto changeDimIdx = viewShapeDim - 1; + // 轴大于等于2才判断是否转置 + if (viewShapeDim >= DIM_IDX_2 && IsTransposeLastTwoDims(tensorS4)) { + transposeTensor = true; + changeDimIdx = viewShapeDim - DIM_IDX_2; + } + tensorShape[changeDimIdx] = tensorShape.GetDim(changeDimIdx) * INT4_PER_INT32; + bool isNz = tensorS4->GetStorageFormat() == op::Format::FORMAT_FRACTAL_NZ; + tensorS4->SetViewShape(tensorShape); + tensorS4->SetDataType(DataType::DT_INT4); + if (isNz){ + OP_LOGD("Reset %s storageShape because tensor is NZ format.", tensorType.c_str()); + auto storageShape = tensorS4->GetStorageShape(); + auto storageShapeDim = storageShape.GetDimNum(); + storageShape[storageShapeDim - 1] *= INT4_PER_INT32; + tensorS4->SetStorageShape(storageShape); + } + if (transposeTensor) { + OP_LOGD("Reset %s stride because tensor is transposed.", tensorType.c_str()); + auto strideSize = newStride.size(); + // 转置场景，B32承载B4时Strides缩小了8倍，需要调整回来 + newStride[strideSize - 1] *= INT4_PER_INT32; + for(int64_t batchDim = strideSize - 3; batchDim >= 0; batchDim--) { + newStride[batchDim] *= INT4_PER_INT32; + } + tensorS4->SetViewStrides(newStride); + } + OP_LOGD("Unpack %s from int32 to int4 finished.", tensorType.c_str()); + } + + bool CheckInputOutDims() override + { + if (gmmDsqParams_.x->GetDataType() == DataType::DT_INT8 + && ((*gmmDsqParams_.weight)[0])->GetDataType() == DataType::DT_INT8) { + return CheckInputOutDimsA8W8(); + } + // A8W4或者A4W4场景 INT32为兼容torch_npu考虑，实际计算时，1个INT32数据会被视为8个INT4数据 + if (gmmDsqParams_.isA8W4 || gmmDsqParams_.isA4W4) { + bool transposeWeight = IsTransposeLastTwoDims((*gmmDsqParams_.weight)[0]); + gmmDsqParams_.transposeWeight = transposeWeight; + // 将INT32视为8个Int4数据，调整viewShape和dtype便于后续统一校验 + if (gmmDsqParams_.x->GetDataType() == DataType::DT_INT32) { + UnpackInt32ToInt4(gmmDsqParams_.x, "x"); + } + if (((*gmmDsqParams_.weight)[0])->GetDataType() == DataType::DT_INT32) { + size_t wLength = gmmDsqParams_.weight->Size(); + for (size_t i = 0; i < wLength; i++) { + const aclTensor *w = (*gmmDsqParams_.weight)[i]; + UnpackInt32ToInt4(w, "weight"); + } + } + + if (transposeWeight == true){ + const aclTensor* w = (*gmmDsqParams_.weight)[0]; + bool isNZ = w->GetStorageFormat() == op::Format::FORMAT_FRACTAL_NZ; + if (!isNZ) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "In weight Transpose scenario.weight Format expect is FRACTAL_NZ when weight is transposed, but got [%s].", + op::ToString(w->GetStorageFormat()).GetString()); + return false; + } + if (!gmmDsqParams_.isA4W4) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "In weight Transpose scenario, only A4W4 is supported."); + return false; + } + } + if (((*gmmDsqParams_.weightScale)[0])->GetDataType() == DataType::DT_INT64) { + size_t weightScaleLength = gmmDsqParams_.weightScale->Size(); + for (size_t i = 0; i < weightScaleLength; i++) { + auto weightScale_fix = const_cast((*gmmDsqParams_.weightScale)[i]); + weightScale_fix->SetDataType(DataType::DT_UINT64); + } + } + return CheckInputOutDimsA4W4orA8W4(); + } + return false; + } + + bool CheckInputOutShape() override + { + if (gmmDsqParams_.x->GetDataType() == DataType::DT_INT8 + && ((*gmmDsqParams_.weight)[0])->GetDataType() == DataType::DT_INT8) { + return CheckInputOutShapeA8W8(); + } + // A8W4场景或A4W4场景 + if (gmmDsqParams_.isA8W4 || gmmDsqParams_.isA4W4) { + return CheckInputOutShapeA8W4orA4W4(); + } + return false; + } + + bool CheckDtypeValid() override + { + size_t wLength = gmmDsqParams_.weight->Size(); + for (size_t i = 0; i < wLength; i++) { + const aclTensor* wScale = (*gmmDsqParams_.weightScale)[i]; + const aclTensor* w = (*gmmDsqParams_.weight)[i]; + + OP_CHECK_DTYPE_NOT_SUPPORT(w, WEIGHT_DTYPE_SUPPORT_LIST, return false); + + if (w->GetDataType() == DataType::DT_INT4) { + OP_CHECK_DTYPE_NOT_SUPPORT(wScale, WEIGHT_SCALE_A8W4_DTYPE_SUPPORT_LIST, return false); + if (gmmDsqParams_.weightAssistMatrix != nullptr && (*gmmDsqParams_.weightAssistMatrix)[i] != nullptr) { + const aclTensor* weightAssistMatrix = (*gmmDsqParams_.weightAssistMatrix)[i]; + OP_CHECK_DTYPE_NOT_SUPPORT(weightAssistMatrix, WEIGHT_ASSIST_DTYPE_SUPPORT_LIST, return false); + } + } else if (w->GetDataType() == DataType::DT_INT8) { + OP_CHECK_DTYPE_NOT_SUPPORT(wScale, WEIGHT_SCALE_DTYPE_SUPPORT_LIST, return false); + } + } + OP_CHECK_DTYPE_NOT_SUPPORT(gmmDsqParams_.x, X_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(gmmDsqParams_.xScale, X_SCALE_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(gmmDsqParams_.groupList, GROUP_LIST_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(gmmDsqParams_.output, QUANTOUT_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(gmmDsqParams_.outputScale, QUANTSCALEOUT_DTYPE_SUPPORT_LIST, return false); + + return true; + } + + bool CheckFormat() override + { + const aclTensor* w = (*gmmDsqParams_.weight)[0]; + bool isNZ = w->GetStorageFormat() == op::Format::FORMAT_FRACTAL_NZ; + if ((gmmDsqParams_.x->GetDataType() == DataType::DT_INT8 && w->GetDataType() == DataType::DT_INT8) && !isNZ) { + // fp16 in fp32 out that is split k template, not precision-advanced now + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "%s, The current version does not support the scenario." + "weight Format expect is FRACTAL_NZ, but got [%s].", interfaceName_.c_str(), + op::ToString(w->GetStorageFormat()).GetString()); + return false; + } + if (IsPrivateFormat(gmmDsqParams_.x->GetStorageFormat())) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "%s, The current version does not support the scenario." + "x Format Not support Private Format.", interfaceName_.c_str()); + return false; + } + if (IsPrivateFormat(gmmDsqParams_.output->GetStorageFormat())) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "%s, The current version does not support the scenario." + "output Format Not support Private Format.", interfaceName_.c_str()); + return false; + } + return true; + } +}; +} +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/grouped_matmul_swiglu_quant_utils.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/grouped_matmul_swiglu_quant_utils.h new file mode 100644 index 000000000..4bb3810bb --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/grouped_matmul_swiglu_quant_utils.h @@ -0,0 +1,409 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_HOST_OP_API_GROUPED_MATMUL_SWIGLU_QUANT_UTILS_H +#define OP_HOST_OP_API_GROUPED_MATMUL_SWIGLU_QUANT_UTILS_H + +#include "aclnn_kernels/contiguous.h" +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "opdev/common_types.h" +#include "opdev/data_type_utils.h" +#include "opdev/format_utils.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/platform.h" +#include "opdev/shape_utils.h" +#include "opdev/tensor_view_utils.h" +#include "opdev/make_op_executor.h" +#include "grouped_matmul_swiglu_quant_v2.h" + +namespace gmm_dsq { +using namespace op; +constexpr int64_t OUTPUT_IDX_0 = 0L; +constexpr int64_t OUTPUT_IDX_1 = 1L; +constexpr size_t WEIGHT_NZ_DIM_LIMIT = 5UL; +constexpr size_t WEIGHT_ND_DIM_LIMIT = 3UL; + +struct GroupedMatmulSwigluQuantParamsBase { + const aclTensor *x = nullptr; + const aclTensorList *weight = nullptr; + const aclTensorList *weightScale = nullptr; + const aclTensorList *weightAssistMatrix = nullptr; + const aclTensor *xScale = nullptr; + const aclTensor *bias = nullptr; + const aclTensor *smoothScale = nullptr; + const aclTensor *groupList = nullptr; + const aclTensor *output = nullptr; + const aclTensor *outputScale = nullptr; + const aclIntArray *tuningConfig = nullptr; + int64_t dequantMode = 0; + int64_t dequantDtype = 0; + int64_t quantMode = 0; + int64_t quantDtype = 0; + int64_t groupListType = 0; + bool transposeWeight = false; + double swigluLimit=0; + bool isA8W4 = false; + bool isA4W4 = false; +}; + +class GroupedMatmulSwigluQuantParamsBuilder { +public: + static GroupedMatmulSwigluQuantParamsBuilder Create(const aclTensor *x, const aclTensorList *weight, + const aclTensorList *weightScale, const aclTensor *output, const aclTensor *outputScale) + { + GroupedMatmulSwigluQuantParamsBuilder b; + b.p_.x = x; + b.p_.weight = weight; + b.p_.weightScale = weightScale; + b.p_.output = output; + b.p_.outputScale = outputScale; + return b; + } + + GroupedMatmulSwigluQuantParamsBuilder &SetWeightAssistMatrix(const aclTensorList *weightAssistMatrix) + { + p_.weightAssistMatrix = weightAssistMatrix; + return *this; + } + + GroupedMatmulSwigluQuantParamsBuilder &SetXScale(const aclTensor *xScale) + { + p_.xScale = xScale; + return *this; + } + + GroupedMatmulSwigluQuantParamsBuilder &SetSmoothScale(const aclTensor *smoothScale) + { + p_.smoothScale = smoothScale; + return *this; + } + + GroupedMatmulSwigluQuantParamsBuilder &SetBias(const aclTensor *bias) + { + p_.bias = bias; + return *this; + } + + GroupedMatmulSwigluQuantParamsBuilder &SetGroupList(const aclTensor *groupList) + { + p_.groupList = groupList; + return *this; + } + + GroupedMatmulSwigluQuantParamsBuilder &SetGroupListType(const int64_t groupListType) + { + p_.groupListType = groupListType; + return *this; + } + + GroupedMatmulSwigluQuantParamsBuilder &SetTuningConfig(const aclIntArray *tuningConfig) + { + p_.tuningConfig = tuningConfig; + return *this; + } + + GroupedMatmulSwigluQuantParamsBuilder &SetDequantAttr(int64_t dequantMode, int64_t dequantDtype) + { + p_.dequantMode = dequantMode; + p_.dequantDtype = dequantDtype; + return *this; + } + + GroupedMatmulSwigluQuantParamsBuilder &SetQuantAttr(int64_t quantMode, int64_t quantDtype) + { + p_.quantMode = quantMode; + p_.quantDtype = quantDtype; + return *this; + } + + GroupedMatmulSwigluQuantParamsBuilder &SetTransposeAttr(bool transposeWeight) + { + p_.transposeWeight = transposeWeight; + return *this; + } + GroupedMatmulSwigluQuantParamsBuilder &SetLimitAttr(double swigluLimit) + { + p_.swigluLimit = swigluLimit; + return *this; + } + GroupedMatmulSwigluQuantParamsBuilder &SetScenario() + { + p_.isA8W4 = ((this->p_.x->GetDataType() == DataType::DT_INT8 && + ((*this->p_.weight)[0])->GetDataType() == DataType::DT_INT4) || + (this->p_.x->GetDataType() == DataType::DT_INT8 && + ((*this->p_.weight)[0])->GetDataType() == DataType::DT_INT32)); + p_.isA4W4 = ((this->p_.x->GetDataType() == DataType::DT_INT4 && + ((*this->p_.weight)[0])->GetDataType() == DataType::DT_INT4) || + (this->p_.x->GetDataType() == DataType::DT_INT4 && + ((*this->p_.weight)[0])->GetDataType() == DataType::DT_INT32) || + (this->p_.x->GetDataType() == DataType::DT_INT32 && + ((*this->p_.weight)[0])->GetDataType() == DataType::DT_INT4) || + (this->p_.x->GetDataType() == DataType::DT_INT32 && + ((*this->p_.weight)[0])->GetDataType() == DataType::DT_INT32)); + return *this; + } + + GroupedMatmulSwigluQuantParamsBase Build() const + { + return p_; + } + +private: + GroupedMatmulSwigluQuantParamsBase p_; +}; + +class GroupedMatmulSwigluQuantHandler { +public: + virtual ~GroupedMatmulSwigluQuantHandler() = default; + +protected: + bool CheckTensorListNull(const aclTensorList *&tensors) const + { + OP_CHECK_NULL(tensors, return false); + if (tensors->Size() == 0) { + return true; + } else if ((tensors->Size() == 1) && ((*tensors)[0] == nullptr)) { + return true; + } + + return false; + } + + virtual bool CheckNotNull(void) + { + OP_CHECK_NULL(gmmDsqParams_.x, return false); + OP_CHECK_NULL(gmmDsqParams_.weight, return false); + OP_CHECK_NULL(gmmDsqParams_.weightScale, return false); + OP_CHECK_NULL(gmmDsqParams_.xScale, return false); + OP_CHECK_NULL(gmmDsqParams_.groupList, return false); + OP_CHECK_NULL(gmmDsqParams_.output, return false); + OP_CHECK_NULL(gmmDsqParams_.outputScale, return false); + + auto ret = CheckTensorListNull(gmmDsqParams_.weight); + if (ret) { + return false; + } + + ret = CheckTensorListNull(gmmDsqParams_.weightScale); + if (ret) { + return false; + } + + if (!gmmDsqParams_.weight || !gmmDsqParams_.weightScale) { + OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, + "The weight or weightScale is nullptr."); + return false; + } + return true; + } + + virtual bool CheckEmptyTensor(void) + { + if ((*gmmDsqParams_.weight)[0]->IsEmpty() || (*gmmDsqParams_.weightScale)[0]->IsEmpty()) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The weight or weightScale is an empty container."); + return false; + } + return true; + } + + virtual bool CheckInputOutDims() = 0; + virtual bool CheckInputOutShape() = 0; + virtual bool CheckDtypeValid() = 0; + virtual bool CheckFormat() = 0; + + virtual aclnnStatus CheckParams() + { + // 1. 检查参数是否为空指针、空tensor + CHECK_RET(CheckNotNull(), ACLNN_ERR_PARAM_NULLPTR); + CHECK_RET(CheckEmptyTensor(), ACLNN_ERR_PARAM_INVALID); + + // 2. 校验输入、输出参数维度 + CHECK_RET(CheckInputOutDims(), ACLNN_ERR_PARAM_INVALID); + + // 3. 校验输入、输出shape参数 + CHECK_RET(CheckInputOutShape(), ACLNN_ERR_PARAM_INVALID); + + // 4. 检查输入的数据类型是否在支持的数据类型范围之内 + CHECK_RET(CheckDtypeValid(), ACLNN_ERR_PARAM_INVALID); + + // 5. 检查数据形状是否支持 + CHECK_RET(CheckFormat(), ACLNN_ERR_PARAM_INVALID); + + return ACLNN_SUCCESS; + } + + void CheckOptionalTensorListEmpty(const aclTensorList *&tensorList) const + { + if (tensorList == nullptr) { + return; + } + + if (tensorList->Size() == 0) { + tensorList = nullptr; + } else if (tensorList->Size() == 1) { + op::Shape shape = (*tensorList)[0]->GetViewShape(); + if (shape.GetDimNum() == 1 && shape.GetDim(0) == 0) { + tensorList = nullptr; + } + } + } + + void CreateEmptyTensor(const aclDataType dataType, const aclTensorList *&tensorList, + aclTensorList *&emptyTensorList) const + { + if (tensorList != nullptr) { + return; + } + + FVector emptyTensors; + aclTensor *emptyTensor = l0Executor_->AllocTensor({0}, static_cast(dataType)); + emptyTensors.emplace_back(emptyTensor); + emptyTensorList = l0Executor_->AllocTensorList(emptyTensors.data(), emptyTensors.size()); + tensorList = emptyTensorList; + } + + aclnnStatus DataContiguous(const aclTensorList *&tensors) const + { + std::vector tensorsVec; + const aclTensor *contiguousTensor = nullptr; + for (size_t i = 0; i < tensors->Size(); ++i) { + const aclTensor *tensor = (*tensors)[i]; + contiguousTensor = l0op::Contiguous(tensor, l0Executor_); + CHECK_RET(contiguousTensor != nullptr, ACLNN_ERR_INNER_NULLPTR); + tensorsVec.push_back(contiguousTensor); + } + tensors = l0Executor_->AllocTensorList(tensorsVec.data(), tensorsVec.size()); + return ACLNN_SUCCESS; + } + + aclnnStatus DataContiguousWeight(const aclTensorList *&tensors) const + { + std::vector tensorsVec; + const aclTensor *contiguousTensor = nullptr; + for (size_t i = 0; i < tensors->Size(); ++i) { + const aclTensor *tensor = (*tensors)[i]; + if (!IsPrivateFormat(tensor->GetStorageFormat())) { + contiguousTensor = l0op::Contiguous(tensor, l0Executor_); + CHECK_RET(contiguousTensor != nullptr, ACLNN_ERR_INNER_NULLPTR); + tensorsVec.push_back(contiguousTensor); + } else { + tensorsVec.push_back(tensor); + } + } + tensors = l0Executor_->AllocTensorList(tensorsVec.data(), tensorsVec.size()); + return ACLNN_SUCCESS; + } + + virtual aclnnStatus CovertDataContiguous() + { + aclTensorList *emptyWeightAssistMatrixList = nullptr; + CreateEmptyTensor(aclDataType::ACL_FLOAT, gmmDsqParams_.weightAssistMatrix, + emptyWeightAssistMatrixList); + + CHECK_COND(DataContiguousWeight(gmmDsqParams_.weight) == ACLNN_SUCCESS, ACLNN_ERR_INNER_NULLPTR, + "Contiguous weight failed."); + CHECK_COND(DataContiguous(gmmDsqParams_.weightScale) == ACLNN_SUCCESS, ACLNN_ERR_INNER_NULLPTR, + "Contiguous weightScale failed."); + if (gmmDsqParams_.weightAssistMatrix != nullptr && gmmDsqParams_.weightAssistMatrix->Size() != 0) { + CHECK_COND(DataContiguous(gmmDsqParams_.weightAssistMatrix) == ACLNN_SUCCESS, ACLNN_ERR_INNER_NULLPTR, + "Contiguous weightAssistMatrix failed."); + } + + gmmDsqParams_.x = l0op::Contiguous(gmmDsqParams_.x, l0Executor_); + CHECK_COND(gmmDsqParams_.x != nullptr, ACLNN_ERR_INNER_NULLPTR, "Contiguous groupList failed."); + gmmDsqParams_.xScale = l0op::Contiguous(gmmDsqParams_.xScale, l0Executor_); + CHECK_COND(gmmDsqParams_.xScale != nullptr, ACLNN_ERR_INNER_NULLPTR, "Contiguous xScale failed."); + gmmDsqParams_.groupList = l0op::Contiguous(gmmDsqParams_.groupList, l0Executor_); + CHECK_COND(gmmDsqParams_.groupList != nullptr, ACLNN_ERR_INNER_NULLPTR, "Contiguous groupList failed."); + + return ACLNN_SUCCESS; + } + +public: + void Initialize(const char *interfaceName, GroupedMatmulSwigluQuantParamsBase ¶ms, uint64_t *workspaceSize, aclOpExecutor **executor) + { + interfaceName_ = interfaceName; + gmmDsqParams_ = params; + workspaceSize_ = workspaceSize; + executor_ = executor; + } + + aclnnStatus Process() + { + // 固定写法，创建OpExecutor + auto uniqueExecutor = CREATE_EXECUTOR(); + CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + l0Executor_ = uniqueExecutor.get(); + + auto ret = CheckParams(); + CHECK_RET(ret == ACLNN_SUCCESS, ret); + + if (op::GetCurrentPlatformInfo().GetCurNpuArch() == NpuArch::DAV_3510) { + auto x1MDim = gmmDsqParams_.x->GetViewShape().GetDim(0); + auto x2NIndex = (*gmmDsqParams_.weight)[0]->GetViewShape().GetDimNum() - 1; + auto x2NDim = (*gmmDsqParams_.weight)[0]->GetViewShape().GetDim(x2NIndex); + if (x1MDim == 0 || x2NDim == 0) { + *workspaceSize_ = 0ULL; + uniqueExecutor.ReleaseTo(executor_); + return ACLNN_SUCCESS; + } + } + for (size_t i = 0; i < gmmDsqParams_.weight->Size(); i++) { + auto *w = (*gmmDsqParams_.weight)[i]; + if (IsPrivateFormat(w->GetStorageFormat())) { + w->SetOriginalShape(w->GetViewShape()); + } + } + // 空Tensor场景 + if (gmmDsqParams_.output->IsEmpty() || gmmDsqParams_.groupList->IsEmpty() || gmmDsqParams_.outputScale->IsEmpty()) { + *workspaceSize_ = 0ULL; + uniqueExecutor.ReleaseTo(executor_); + return ACLNN_SUCCESS; + } + + ret = CovertDataContiguous(); + CHECK_RET(ret == ACLNN_SUCCESS, ret); + auto ret0 = l0op::GroupedMatmulSwigluQuantV2(gmmDsqParams_.x, gmmDsqParams_.weight, gmmDsqParams_.weightScale, + gmmDsqParams_.xScale, gmmDsqParams_.weightAssistMatrix, + gmmDsqParams_.bias, + gmmDsqParams_.smoothScale, gmmDsqParams_.groupList, + gmmDsqParams_.dequantMode, gmmDsqParams_.dequantDtype, + gmmDsqParams_.quantMode, gmmDsqParams_.quantDtype, + gmmDsqParams_.transposeWeight, gmmDsqParams_.groupListType, + gmmDsqParams_.tuningConfig,gmmDsqParams_.swigluLimit, uniqueExecutor.get()); + CHECK_RET(ret0 != std::tuple(nullptr, nullptr), ACLNN_ERR_INNER_NULLPTR); + + auto out0 = std::get(ret0); + auto ret1 = l0op::ViewCopy(out0, gmmDsqParams_.output, uniqueExecutor.get()); + CHECK_RET(ret1 != nullptr, ACLNN_ERR_INNER_NULLPTR); + + auto out1 = std::get(ret0); + auto ret2 = l0op::ViewCopy(out1, gmmDsqParams_.outputScale, uniqueExecutor.get()); + CHECK_RET(ret2 != nullptr, ACLNN_ERR_INNER_NULLPTR); + + *workspaceSize_ = uniqueExecutor->GetWorkspaceSize(); + uniqueExecutor.ReleaseTo(executor_); + return ACLNN_SUCCESS; + } + +protected: + string interfaceName_; + GroupedMatmulSwigluQuantParamsBase gmmDsqParams_; + uint64_t *workspaceSize_; + aclOpExecutor **executor_; + aclOpExecutor *l0Executor_; +}; + +} // namespace gmm_dsq +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/grouped_matmul_swiglu_quant_v2.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/grouped_matmul_swiglu_quant_v2.cpp new file mode 100644 index 000000000..a52b0e139 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/grouped_matmul_swiglu_quant_v2.cpp @@ -0,0 +1,87 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "opdev/op_log.h" +#include "opdev/op_dfx.h" +#include "opdev/make_op_executor.h" +#include "util/math_util.h" +#include "grouped_matmul_swiglu_quant_utils.h" +#include "grouped_matmul_swiglu_quant_v2.h" + +using namespace op; +using namespace gmm_dsq; + +namespace l0op { +OP_TYPE_REGISTER(GroupedMatmulSwigluQuantV2); + +constexpr int64_t SWIGLU_SPLIT_SIZE = 64L; + +const std::tuple GroupedMatmulSwigluQuantV2(const aclTensor *x, const aclTensorList *weight, + const aclTensorList *weightScale, + const aclTensor *xScale, const aclTensorList *weightAssistanceMatrix, + const aclTensor *bias, const aclTensor *smoothScale, + const aclTensor *groupList, int64_t dequantMode, int64_t dequantDtype, + int64_t quantMode, int64_t quantDtype, bool transposeWeight, int64_t groupListType, + const aclIntArray *tuningConfigOptional, double swigluLimit,aclOpExecutor *executor) +{ + L0_DFX(GroupedMatmulSwigluQuantV2, x, weight, weightScale, xScale, weightAssistanceMatrix, smoothScale, + groupList, dequantMode, dequantDtype, quantMode, quantDtype, transposeWeight, tuningConfigOptional, swigluLimit); + if (x == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "x is nullptr."); + return std::tuple(nullptr, nullptr); + } + int64_t m = xScale->GetViewShape().GetDim(0); + int64_t n = (*weightScale)[0]->GetViewShape().GetDim(1); + int64_t nAfterHalve = static_cast(n / 2); + gert::Shape outShape({m, nAfterHalve}); + gert::Shape scaleOutShape({m}); + auto out = executor->AllocTensor(outShape, DataType::DT_INT8, ge::FORMAT_ND); + auto scaleOut = executor->AllocTensor(scaleOutShape, DataType::DT_FLOAT, ge::FORMAT_ND); + if (op::GetCurrentPlatformInfo().GetCurNpuArch() == NpuArch::DAV_3510) { + n = transposeWeight ? (*weightScale)[0]->GetViewShape().GetDim(1) : // 转置情况下weightScale的第1维是n + (*weightScale)[0]->GetViewShape().GetDim(2); // 非转置情况下weightScale的第2维是n + nAfterHalve = static_cast(n / 2); // outShape需要为[M, N / 2] + gert::Shape outShapeV2({m, nAfterHalve}); + gert::Shape scaleOutShapeV2; + // 当quantMode等于2时，out_scale 的形状为三维 + if (quantMode == 2) { + int64_t nAfterSplit = static_cast(Ops::Base::CeilDiv(nAfterHalve, SWIGLU_SPLIT_SIZE)); + scaleOutShapeV2 = gert::Shape({m, nAfterSplit, 2}); + } else { + scaleOutShapeV2 = gert::Shape({m}); + } + out = executor->AllocTensor(outShapeV2, static_cast(quantDtype), ge::FORMAT_ND); + // 当quantMode等于2时，outScale的DataType为FLOAT8_E8M0 + scaleOut = quantMode == 2 ? executor->AllocTensor(scaleOutShapeV2, DataType::DT_FLOAT8_E8M0, ge::FORMAT_ND) : + executor->AllocTensor(scaleOutShapeV2, DataType::DT_FLOAT, ge::FORMAT_ND); + } + auto ret = INFER_SHAPE(GroupedMatmulSwigluQuantV2, + OP_INPUT(x, xScale, groupList, weight, weightScale, weightAssistanceMatrix, bias, smoothScale), + OP_OUTPUT(out, scaleOut), OP_ATTR(dequantMode, dequantDtype, quantMode, quantDtype, transposeWeight, + groupListType, tuningConfigOptional, swigluLimit)); + if (ret != ACLNN_SUCCESS) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "InferShape failed."); + return std::tuple(nullptr, nullptr); + } + + ret = ADD_TO_LAUNCHER_LIST_AICORE( + GroupedMatmulSwigluQuantV2, + OP_INPUT(x, xScale, groupList, weight, weightScale, weightAssistanceMatrix, bias, smoothScale), + OP_OUTPUT(out, scaleOut), OP_ATTR(dequantMode, dequantDtype, quantMode, quantDtype, transposeWeight, + groupListType, tuningConfigOptional, swigluLimit)); + if (ret != ACLNN_SUCCESS) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "ADD_TO_LAUNCHER_LIST_AICORE failed."); + return std::tuple(nullptr, nullptr); + } + + return std::tie(out, scaleOut); +} + +} // namespace l0op \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/grouped_matmul_swiglu_quant_v2.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/grouped_matmul_swiglu_quant_v2.h new file mode 100644 index 000000000..f7c926f83 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/grouped_matmul_swiglu_quant_v2.h @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_HOST_OP_API_GROUPED_MATMUL_SWIGLU_QUANT_V2_H +#define OP_HOST_OP_API_GROUPED_MATMUL_SWIGLU_QUANT_V2_H + +#include "opdev/op_executor.h" + +namespace l0op { + +const std::tuple GroupedMatmulSwigluQuantV2(const aclTensor *x, const aclTensorList *weight, + const aclTensorList *weightScale, + const aclTensor *xScale, const aclTensorList *weightAssistanceMatrix, + const aclTensor *bias, const aclTensor *smoothScale, + const aclTensor *groupList, int64_t dequantMode, int64_t dequantDtype, + int64_t quantMode, int64_t quantDtype, bool transposeWeight, int64_t groupListType, + const aclIntArray *tuningConfigOptional, double swigluLimit, aclOpExecutor *executor); +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/grouped_matmul_swiglu_quant_v2_utils.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/grouped_matmul_swiglu_quant_v2_utils.h new file mode 100644 index 000000000..163ea467b --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_host/op_api/grouped_matmul_swiglu_quant_v2_utils.h @@ -0,0 +1,814 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_HOST_OP_API_GROUPED_MATMUL_SWIGLU_QUANT_V2_UTILS_H +#define OP_HOST_OP_API_GROUPED_MATMUL_SWIGLU_QUANT_V2_UTILS_H + +#include "grouped_matmul_swiglu_quant_utils.h" +#include "util/math_util.h" + +namespace gmmSwigluQuantV2 { + +using namespace gmm_dsq; + +constexpr int64_t OUTPUT_IDX_0 = 0L; +constexpr int64_t OUTPUT_IDX_1 = 1L; +constexpr size_t MX_SPLIT_K_PER_TOKEN_SCALE_DIM = 3UL; +constexpr size_t LAST_SECOND_DIM_INDEX = 2; +constexpr size_t LAST_THIRD_DIM_INDEX = 3; +constexpr int64_t MXFP_MULTI_BASE_SIZE = 2L; +constexpr size_t MX_SPLIT_M_SCALE_DIM = 4UL; +constexpr size_t MX_X_DIM = 2UL; +constexpr size_t MX_X_SCALE_DIM = 3UL; +constexpr size_t MX_WEIGHT_DIM = 3UL; +constexpr size_t MX_WEIGHT_SCALE_DIM = 4UL; +constexpr size_t MX_OUTPUT_DIM = 2UL; +constexpr size_t MX_OUTPUT_SCALE_DIM = 3UL; +constexpr size_t PERTOKEN_X_DIM = 2; +constexpr size_t PERTOKEN_X_SCALE_DIM = 1; +constexpr size_t PERTOKEN_WEIGHT_DIM = 3; +constexpr size_t PERTOKEN_WEIGHT_SCALE_DIM = 2; +constexpr size_t PERTOKEN_OUTPUT_DIM = 2; +constexpr size_t PERTOKEN_OUTPUT_SCALE_DIM = 1; +constexpr int64_t SWIGLU_SPLIT_FACTOR = 2L; +constexpr int64_t SWIGLU_SPLIT_SIZE = 64L; +constexpr int64_t MXFP4_K_CONSTRAINT = 2L; +constexpr int64_t SWIGLU_N_CONSTRAINT = 2L; +constexpr int64_t MXFP4_N_CONSTRAINT = 4L; +constexpr size_t SINGLE_TENSOR_SIZE = 1; +constexpr int64_t MAX_GROUP_LIST_SIZE = 1024L; +constexpr int64_t QUNAT_MODE_MX = 2; +constexpr int64_t QUNAT_MODE_PERTOKEN = 0; + +const std::initializer_list X_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT8_E4M3FN, DataType::DT_FLOAT8_E5M2}; +const std::initializer_list X_DTYPE_SUPPORT_LIST_MXFP4 = {DataType::DT_FLOAT4_E2M1}; +const std::initializer_list XW_DTYPE_SUPPORT_LIST_PERTOKEN = { + DataType::DT_INT8, DataType::DT_FLOAT8_E4M3FN, DataType::DT_FLOAT8_E5M2, DataType::DT_HIFLOAT8}; +const std::initializer_list WEIGHT_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT8_E4M3FN, + DataType::DT_FLOAT8_E5M2}; +const std::initializer_list WEIGHT_DTYPE_SUPPORT_LIST_MXFP4 = {DataType::DT_FLOAT4_E2M1}; +const std::initializer_list WEIGHT_SCALE_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT8_E8M0}; +const std::initializer_list WEIGHT_SCALE_DTYPE_SUPPORT_LIST_PERTOKEN_XINT8 = { + DataType::DT_FLOAT16, DataType::DT_BF16, DataType::DT_FLOAT}; +const std::initializer_list WEIGHT_SCALE_DTYPE_SUPPORT_LIST_PERTOKEN_XFP8HIF8 = {DataType::DT_BF16, + DataType::DT_FLOAT}; +const std::initializer_list X_SCALE_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT8_E8M0}; +const std::initializer_list X_SCALE_DTYPE_SUPPORT_LIST_PERTOKEN = {DataType::DT_FLOAT}; +const std::initializer_list GROUP_LIST_DTYPE_SUPPORT_LIST = {DataType::DT_INT64}; +const std::initializer_list QUANTOUT_DTYPE_SUPPORT_LIST_MXFP4 = { + DataType::DT_FLOAT8_E4M3FN, DataType::DT_FLOAT8_E5M2, DataType::DT_FLOAT4_E2M1}; +const std::initializer_list QUANTOUT_DTYPE_SUPPORT_LIST_PERTOKEN = { + DataType::DT_INT8, DataType::DT_FLOAT8_E4M3FN, DataType::DT_FLOAT8_E5M2, DataType::DT_HIFLOAT8}; +const std::initializer_list QUANTSCALEOUT_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT8_E8M0}; +const std::initializer_list QUANTSCALEOUT_DTYPE_SUPPORT_LIST_PERTOKEN = {DataType::DT_FLOAT}; + +class GroupedMatmulSwigluQuantBaseHandler : public GroupedMatmulSwigluQuantHandler { +protected: + bool IsTransposeForMxShape(const aclTensor *tensor) const + { + auto shape = tensor->GetViewShape(); + if (shape.GetDimNum() < MX_SPLIT_K_PER_TOKEN_SCALE_DIM) { + return false; + } + int64_t firstLastDim = shape.GetDimNum() - 1; + int64_t secondLastDim = shape.GetDimNum() - LAST_SECOND_DIM_INDEX; + int64_t thirdLastDim = shape.GetDimNum() - LAST_THIRD_DIM_INDEX; + auto strides = tensor->GetViewStrides(); + if (strides[firstLastDim] == 1 && strides[thirdLastDim] == MXFP_MULTI_BASE_SIZE && + strides[secondLastDim] == shape.GetDim(thirdLastDim) * MXFP_MULTI_BASE_SIZE) { + return true; + } + return false; + } + + bool IsTransposeLastTwoDims(const aclTensor *tensor) const + { + auto shape = tensor->GetViewShape(); + int64_t dim1 = shape.GetDimNum() - 1; + int64_t dim2 = shape.GetDimNum() - 2; + auto strides = tensor->GetViewStrides(); + if (strides[dim2] == 1 && strides[dim1] == shape.GetDim(dim2)) { + int64_t tmpNxD = shape.GetDim(dim1) * shape.GetDim(dim2); + for (int64_t batchDim = shape.GetDimNum() - 3; batchDim >= 0; batchDim--) { + if (strides[batchDim] != tmpNxD) { + return false; + } + tmpNxD *= shape.GetDim(batchDim); + } + return true; + } + return false; + } + + void CreateContiguousTensorListForMXTypeMScale(const aclTensorList *tensorList, + std::vector &newTensorList, + aclOpExecutor *executor) const + { + op::Shape shape; + for (uint64_t idx = 0; idx < (*tensorList).Size(); idx++) { + const aclTensor *inputTensor = (*tensorList)[idx]; + op::Shape viewShape = inputTensor->GetViewShape(); + shape.SetScalar(); + if (viewShape.GetDimNum() < MX_SPLIT_M_SCALE_DIM) { + continue; + } + shape.AppendDim(viewShape.GetDim(0)); + shape.AppendDim(viewShape.GetDim(viewShape.GetDimNum() - LAST_SECOND_DIM_INDEX)); + shape.AppendDim(viewShape.GetDim(viewShape.GetDimNum() - LAST_THIRD_DIM_INDEX)); + shape.AppendDim(viewShape.GetDim(viewShape.GetDimNum() - 1)); + aclTensor *tensor = + executor->CreateView(inputTensor, shape, inputTensor->GetViewOffset()); // use executor to create tensor + tensor->SetStorageFormat(inputTensor->GetStorageFormat()); + newTensorList.emplace_back(tensor); + } + } + + void CreateContiguousTensorList(const aclTensorList *tensorList, std::vector &newTensorList, + aclOpExecutor *executor) const + { + op::Shape shape; + for (uint64_t idx = 0; idx < (*tensorList).Size(); idx++) { + const aclTensor *inputTensor = (*tensorList)[idx]; + op::Shape viewShape = inputTensor->GetViewShape(); + uint32_t viewShapeDimsNum = viewShape.GetDimNum(); + shape.SetScalar(); + // 2: the second last dimension; in for-loops, it indicates dimensions before the second last remain unchanged. + for (uint32_t i = 0; i < viewShapeDimsNum - 2; ++i) { + shape.AppendDim(viewShape.GetDim(i)); + } + // viewShapeDimsNum - 1, the dim value of the last dim. viewShapeDimsNum - 2, the dim value of the second + // last dim. + shape.AppendDim(viewShape.GetDim(viewShapeDimsNum - 1)); + shape.AppendDim(viewShape.GetDim(viewShapeDimsNum - 2)); // 2:the second last dim. + aclTensor *tensor = + executor->CreateView(inputTensor, shape, inputTensor->GetViewOffset()); // use executor to create tensor + tensor->SetStorageFormat(inputTensor->GetStorageFormat()); + newTensorList.emplace_back(tensor); + } + } + + static void CheckOptionalTensorListEmpty(const aclTensorList *&tensorList) + { + if (tensorList != nullptr) { + if (tensorList->Size() == 0) { + tensorList = nullptr; + } else if ((*tensorList)[0] == nullptr) { + tensorList = nullptr; + } else if (tensorList->Size() == 1) { + op::Shape shape = (*tensorList)[0]->GetViewShape(); + if (shape.GetDimNum() == 1 && shape.GetDim(0) == 0) { + tensorList = nullptr; + } + } + } + } + + bool CheckAttrs() + { + CheckOptionalTensorListEmpty(gmmDsqParams_.weightAssistMatrix); + if (gmmDsqParams_.tuningConfig != nullptr && gmmDsqParams_.tuningConfig->Size() == 0) { + gmmDsqParams_.tuningConfig = nullptr; + } + if (gmmDsqParams_.weightAssistMatrix != nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "The current version does not support weightAssistMatrix, it should be nullptr."); + return false; + } + if (gmmDsqParams_.bias != nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The current version does not support bias, it should be nullptr."); + return false; + } + if (gmmDsqParams_.smoothScale != nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The current version does not support smoothScale, it should be nullptr."); + return false; + } + if (gmmDsqParams_.tuningConfig != nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "The current version does not support tuningConfig, it should be nullptr."); + return false; + } + if ((gmmDsqParams_.dequantMode != QUNAT_MODE_MX && gmmDsqParams_.dequantMode != QUNAT_MODE_PERTOKEN) || + (gmmDsqParams_.quantMode != QUNAT_MODE_MX && gmmDsqParams_.quantMode != QUNAT_MODE_PERTOKEN) || + (gmmDsqParams_.dequantMode != gmmDsqParams_.quantMode)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "Both dequantMode and quantMode must be 0 (pertoken) or 2 (mx), and they must be equal. Actual " + "value: dequantMode=%lu, dequantMode=%lu.", + gmmDsqParams_.dequantMode, gmmDsqParams_.quantMode); + return false; + } + ge::DataType dequantDtype = static_cast(gmmDsqParams_.dequantDtype); + if (gmmDsqParams_.quantMode == QUNAT_MODE_MX && dequantDtype != ge::DT_FLOAT) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "In mx quant mode, dequantDtype should be 0, but actual value is %lu.", + gmmDsqParams_.dequantDtype); + return false; + } + if (gmmDsqParams_.quantMode == QUNAT_MODE_PERTOKEN && dequantDtype != ge::DT_FLOAT && dequantDtype != ge::DT_BF16 && + dequantDtype != ge::DT_FLOAT16) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "In pertoken quant mode, dequantDtype should be 0, 1, 27, but actual value is %lu.", + gmmDsqParams_.dequantDtype); + return false; + } + return true; + } + + bool CheckMXTranspose() + { + // 判断weight和weightScale是否转置，是则对两者进行转置动作 + bool transposeWeightScale = IsTransposeForMxShape((*gmmDsqParams_.weightScale)[0]); + bool transposeWeight = IsTransposeLastTwoDims((*gmmDsqParams_.weight)[0]); + bool transposeX = IsTransposeLastTwoDims(gmmDsqParams_.x); + bool transposeXScale = IsTransposeForMxShape(gmmDsqParams_.xScale); + + if (transposeWeightScale != transposeWeight) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "The transposition of weightScale/weight should be equal, but actual transpositions are %s/%s.", + transposeWeightScale ? "true" : "false", transposeWeight ? "true" : "false"); + return false; + } + + if (transposeWeightScale && transposeWeight) { + gmmDsqParams_.transposeWeight = true; + auto uniqueExecutor = CREATE_EXECUTOR(); + CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + aclOpExecutor *executorPtr = uniqueExecutor.get(); + CHECK_RET(executorPtr != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + std::vector scaleTensorList; + std::vector weightTensorList; + CreateContiguousTensorListForMXTypeMScale(gmmDsqParams_.weightScale, scaleTensorList, executorPtr); + gmmDsqParams_.weightScale = executorPtr->AllocTensorList(scaleTensorList.data(), scaleTensorList.size()); + CreateContiguousTensorList(gmmDsqParams_.weight, weightTensorList, executorPtr); + gmmDsqParams_.weight = executorPtr->AllocTensorList(weightTensorList.data(), weightTensorList.size()); + uniqueExecutor.ReleaseTo(executor_); + } + + if ((gmmDsqParams_.x->GetViewShape().GetDim(0) == 1 && gmmDsqParams_.x->GetViewShape().GetDim(1) == 1) || + (gmmDsqParams_.xScale->GetViewShape().GetDim(0) == 1 && + gmmDsqParams_.xScale->GetViewShape().GetDim(1) == 1)) { + return true; + } + if (transposeX || transposeXScale) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "The transposition of x/xScale should be false, but actual transposition are %s/%s.", + transposeX ? "true" : "false", transposeXScale ? "true" : "false"); + return false; + } + return true; + } + + bool CheckPertokenTranspose() + { + bool transposeWeight = IsTransposeLastTwoDims((*gmmDsqParams_.weight)[0]); + bool transposeX = IsTransposeLastTwoDims(gmmDsqParams_.x); + + if (transposeWeight) { + gmmDsqParams_.transposeWeight = true; + auto uniqueExecutor = CREATE_EXECUTOR(); + CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + aclOpExecutor *executorPtr = uniqueExecutor.get(); + CHECK_RET(executorPtr != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + std::vector weightTensorList; + CreateContiguousTensorList(gmmDsqParams_.weight, weightTensorList, executorPtr); + gmmDsqParams_.weight = executorPtr->AllocTensorList(weightTensorList.data(), weightTensorList.size()); + uniqueExecutor.ReleaseTo(executor_); + } + if ((gmmDsqParams_.x->GetViewShape().GetDim(0) == 1 && gmmDsqParams_.x->GetViewShape().GetDim(1) == 1) || + (gmmDsqParams_.xScale->GetViewShape().GetDim(0) == 1 && + gmmDsqParams_.xScale->GetViewShape().GetDim(1) == 1)) { + return true; + } + if (transposeX) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The transposition of x should be false, but actual transposition are %s.", + transposeX ? "true" : "false"); + return false; + } + return true; + } + + bool CheckMXShape() + { + int64_t m = gmmDsqParams_.x->GetViewShape().GetDim(0); // 从x的第0维获取m + int64_t k = gmmDsqParams_.x->GetViewShape().GetDim(1); // 从x的第1维获取k + // 转置情况下从weight的第1维获取n，非转置情况下从weight的第2维获取n + int64_t n = gmmDsqParams_.transposeWeight ? ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(1) : + ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(2); + int64_t e = ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(0); // 从weight的第0维获取e + + // x的shape期望为[M, K] + op::Shape xExpectShape = {m, k}; + // xScale的shape期望为[M, CeilDiv(K, 64), 2] + op::Shape xScaleExpectShape = {m, Ops::Base::CeilDiv(k, SWIGLU_SPLIT_SIZE), SWIGLU_SPLIT_FACTOR}; + // weight的shape期望为[E, K, N] + op::Shape weightExpectShape = {e, k, n}; + // weightScale的shape期望为[E, CeilDiv(K, 64), N, 2] + op::Shape weightScaleExpectShape = {e, Ops::Base::CeilDiv(k, SWIGLU_SPLIT_SIZE), n, SWIGLU_SPLIT_FACTOR}; + // weight转置的shape期望为[E, N, K] + op::Shape weightTransExpectShape = {e, n, k}; + // weightScale转置的shape期望为[E, N, CeilDiv(K, 64), 2] + op::Shape weightScaleTransExpectShape = {e, n, Ops::Base::CeilDiv(k, SWIGLU_SPLIT_SIZE), SWIGLU_SPLIT_FACTOR}; + int64_t nAfterHalve = static_cast(n / SWIGLU_SPLIT_FACTOR); + // output的shape期望为[M, N / 2] + op::Shape outputExpectShape = {m, nAfterHalve}; + // outputScale的shape期望为[M, CeilDiv(N / 2, 64), 2] + op::Shape outputScaleExpectShape = {m, Ops::Base::CeilDiv(nAfterHalve, SWIGLU_SPLIT_SIZE), SWIGLU_SPLIT_FACTOR}; + const aclTensor *x = gmmDsqParams_.x; + const aclTensor *xScale = gmmDsqParams_.xScale; + const aclTensor *output = gmmDsqParams_.output; + const aclTensor *outputScale = gmmDsqParams_.outputScale; + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(x, xExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(xScale, xScaleExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(output, outputExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(outputScale, outputScaleExpectShape, return false); + + const aclTensor *weightScale = (*gmmDsqParams_.weightScale)[0]; + const aclTensor *weight = (*gmmDsqParams_.weight)[0]; + if (gmmDsqParams_.transposeWeight) { + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weightScale, weightScaleTransExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weight, weightTransExpectShape, return false); + } else { + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weightScale, weightScaleExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weight, weightExpectShape, return false); + } + // 进行swiglu操作需满足n为偶数 + if (n % SWIGLU_N_CONSTRAINT != 0) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Swiglu operation requires n to be even , but n actual value is %lu.", n); + return false; + } + // groupList的长度应等于weight的专家数 + int64_t groupListLen = gmmDsqParams_.groupList->GetViewShape().GetDim(0); + if (groupListLen != e) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "Length of 'groupList' should be equal to the number of experts in weight."); + return false; + } + return true; + } + + bool CheckPertokenShape() + { + int64_t m = gmmDsqParams_.x->GetViewShape().GetDim(0); // 从x的第0维获取m + int64_t k = gmmDsqParams_.x->GetViewShape().GetDim(1); // 从x的第1维获取k + // 转置情况下从weight的第1维获取n，非转置情况下从weight的第2维获取n + int64_t n = gmmDsqParams_.transposeWeight ? ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(1) : + ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(2); + int64_t e = ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(0); // 从weight的第0维获取e + + // x的shape期望为[M, K] + op::Shape xExpectShape = {m, k}; + // xScale的shape期望为[M] + op::Shape xScaleExpectShape = {m}; + // weight的shape期望为根据转置的情况来具体确认[E, K, N] 或者[E, N, K] + op::Shape weightExpectShape = gmmDsqParams_.transposeWeight ? op::Shape{e, n, k} : op::Shape{e, k, n}; + // weightScale的shape期望为[E, N] + op::Shape weightScaleExpectShape = {e, n}; + int64_t nAfterHalve = static_cast(n / SWIGLU_SPLIT_FACTOR); + // output的shape期望为[M, N / 2] + op::Shape outputExpectShape = {m, nAfterHalve}; + // outputScale的shape期望为[M] + op::Shape outputScaleExpectShape = {m}; + const aclTensor *x = gmmDsqParams_.x; + const aclTensor *xScale = gmmDsqParams_.xScale; + const aclTensor *weight = (*gmmDsqParams_.weight)[0]; + const aclTensor *weightScale = (*gmmDsqParams_.weightScale)[0]; + const aclTensor *output = gmmDsqParams_.output; + const aclTensor *outputScale = gmmDsqParams_.outputScale; + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(x, xExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(xScale, xScaleExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weight, weightExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(weightScale, weightScaleExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(output, outputExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(outputScale, outputScaleExpectShape, return false); + return true; + } + + bool CheckFp8DtypeValid(const aclTensor *x, const aclTensor *xScale, const aclTensor *groupList, + const aclTensor *output, const aclTensor *outputScale) + { + size_t weightLength = gmmDsqParams_.weight->Size(); + for (size_t i = 0; i < weightLength; i++) { + const aclTensor *weightScale = (*gmmDsqParams_.weightScale)[i]; + const aclTensor *weight = (*gmmDsqParams_.weight)[i]; + OP_CHECK_DTYPE_NOT_SUPPORT(weight, WEIGHT_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(weightScale, WEIGHT_SCALE_DTYPE_SUPPORT_LIST, return false); + } + OP_CHECK_DTYPE_NOT_SUPPORT(x, X_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(xScale, X_SCALE_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(groupList, GROUP_LIST_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(outputScale, QUANTSCALEOUT_DTYPE_SUPPORT_LIST, return false); + DataType outputDtype = gmmDsqParams_.output->GetDataType(); + if (outputDtype != DataType::DT_FLOAT8_E4M3FN && outputDtype != DataType::DT_FLOAT8_E5M2) { + OP_LOGE( + ACLNN_ERR_PARAM_INVALID, + "When the dtypes of x and weight inputs are DT_FLOAT8_E4M3FN or " + "DT_FLOAT8_E5M2, the dtypes of output should be DT_FLOAT8_E4M3FN or DT_FLOAT8_E5M2, but actual value " + "is %s.", + op::ToString(outputDtype).GetString()); + return false; + } + return true; + } + + bool CheckFp4DtypeValid(const aclTensor *x, const aclTensor *xScale, const aclTensor *groupList, + const aclTensor *output, const aclTensor *outputScale) + { + size_t weightLength = gmmDsqParams_.weight->Size(); + for (size_t i = 0; i < weightLength; i++) { + const aclTensor *weightScale = (*gmmDsqParams_.weightScale)[i]; + const aclTensor *weight = (*gmmDsqParams_.weight)[i]; + OP_CHECK_DTYPE_NOT_SUPPORT(weight, WEIGHT_DTYPE_SUPPORT_LIST_MXFP4, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(weightScale, WEIGHT_SCALE_DTYPE_SUPPORT_LIST, return false); + } + OP_CHECK_DTYPE_NOT_SUPPORT(x, X_DTYPE_SUPPORT_LIST_MXFP4, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(xScale, X_SCALE_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(groupList, GROUP_LIST_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(output, QUANTOUT_DTYPE_SUPPORT_LIST_MXFP4, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(outputScale, QUANTSCALEOUT_DTYPE_SUPPORT_LIST, return false); + return true; + } + + bool CheckPertokenDtypeValid(const aclTensor *x, const aclTensor *xScale, const aclTensor *groupList, + const aclTensor *output, const aclTensor *outputScale) + { + OP_CHECK_DTYPE_NOT_SUPPORT(x, XW_DTYPE_SUPPORT_LIST_PERTOKEN, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(xScale, X_SCALE_DTYPE_SUPPORT_LIST_PERTOKEN, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(groupList, GROUP_LIST_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(output, QUANTOUT_DTYPE_SUPPORT_LIST_PERTOKEN, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(outputScale, QUANTSCALEOUT_DTYPE_SUPPORT_LIST_PERTOKEN, return false); + size_t weightLength = gmmDsqParams_.weight->Size(); + for (size_t i = 0; i < weightLength; i++) { + const aclTensor *weight = (*gmmDsqParams_.weight)[i]; + const aclTensor *weightScale = (*gmmDsqParams_.weightScale)[i]; + OP_CHECK_DTYPE_NOT_SUPPORT(weight, XW_DTYPE_SUPPORT_LIST_PERTOKEN, return false); + DataType xDtype = gmmDsqParams_.x->GetDataType(); + if (xDtype == DataType::DT_INT8) { + OP_CHECK_DTYPE_NOT_SUPPORT(weightScale, WEIGHT_SCALE_DTYPE_SUPPORT_LIST_PERTOKEN_XINT8, return false); + } else { + OP_CHECK_DTYPE_NOT_SUPPORT(weightScale, WEIGHT_SCALE_DTYPE_SUPPORT_LIST_PERTOKEN_XFP8HIF8, + return false); + } + } + DataType xDtype = gmmDsqParams_.x->GetDataType(); + return IsDtypeCompatiblePertoken(xDtype, ((*gmmDsqParams_.weight)[0])->GetDataType()); + } + + bool IsDtypeCompatiblePertoken(const DataType a, const DataType b) const + { + if ((a == DataType::DT_FLOAT8_E4M3FN || a == DataType::DT_FLOAT8_E5M2) && + (b == DataType::DT_FLOAT8_E4M3FN || b == DataType::DT_FLOAT8_E5M2)) { + return true; + } + return a == b; + } + + bool checkMxfp4InputShape() + { + int64_t kValue = gmmDsqParams_.x->GetViewShape().GetDim(1); + // 转置情况下从weight的第1维获取n，非转置情况下从weight的第2维获取n + int64_t nValue = gmmDsqParams_.transposeWeight ? ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(1) : + ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(2); + // mxfp4场景不支持k=2 + if (kValue == MXFP4_K_CONSTRAINT) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "When the dtypes of x and weight inputs are DT_FLOAT4_E2M1, the K value \ +should be greater than 2, but actual value is %lu.", + kValue); + return false; + } + + // 1：检查K是否为偶数 + int64_t kModValue = kValue % MXFP4_K_CONSTRAINT; + // 2：检查N是否为偶数 + int64_t nModValue = nValue % MXFP4_N_CONSTRAINT; + if (kModValue != 0) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "When the dtypes of x and weight inputs are DT_FLOAT4_E2M1, the K value \ +should be even, but actual value is %lu.", + kValue); + return false; + } + + // mxfp4场景下，当输出类型为fp4时，N需要满足为大于等于4的偶数 + DataType outputDtype = gmmDsqParams_.output->GetDataType(); + if (outputDtype == DataType::DT_FLOAT4_E2M1) { + if (!(nValue >= MXFP4_N_CONSTRAINT && nModValue == 0)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "When the output dtype is DT_FLOAT4_E2M1, the N value should be even \ +and greater or equal to 4, but actual value is %lu.", + nValue); + return false; + } + } + + return true; + } + + bool CheckEmptyTensor() override + { + if (gmmDsqParams_.x->GetViewShape().GetDim(1) <= 0) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "When the M value is not 0, the K value in x should be positive, but actual value is %ld", + gmmDsqParams_.x->GetViewShape().GetDim(1)); + return false; + } + auto weightKIndex = ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDimNum() - LAST_SECOND_DIM_INDEX; + if (((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(weightKIndex) <= 0) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "When the N value is not 0, the K value in weight should be positive, but actual value is %ld", + ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(weightKIndex)); + return false; + } + return true; + } + + bool CheckInputOutDims() override + { + if (!CheckAttrs()) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "CheckAttrs failed."); + return false; + } + + if (gmmDsqParams_.quantMode == QUNAT_MODE_MX) { + return CheckInputOutDimsForMX(); + } else if (gmmDsqParams_.quantMode == QUNAT_MODE_PERTOKEN) { + return CheckInputOutDimsForPertoken(); + } else { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "Quant mode %d is not supported. Supported modes are 0 (pertoken) and 2 (MX).", + gmmDsqParams_.quantMode); + return false; + } + return true; + } + + bool CheckInputOutDimsForMX() + { + auto xDimNumber = gmmDsqParams_.x->GetViewShape().GetDimNum(); + auto xScaleDimNumber = gmmDsqParams_.xScale->GetViewShape().GetDimNum(); + auto outputDimNumber = gmmDsqParams_.output->GetViewShape().GetDimNum(); + auto outputScaleDimNumber = gmmDsqParams_.outputScale->GetViewShape().GetDimNum(); + if (xDimNumber != MX_X_DIM) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The dim num of x should be equal 2, current dim is %lu.", xDimNumber); + return false; + } + if (xScaleDimNumber != MX_X_SCALE_DIM) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The dim num of xScale should be equal 3, current dim is %lu.", + xScaleDimNumber); + return false; + } + if (gmmDsqParams_.weight->Size() != SINGLE_TENSOR_SIZE) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The size of weight should be 1, current size is %lu.", + gmmDsqParams_.weight->Size()); + return false; + } + if (gmmDsqParams_.weightScale->Size() != SINGLE_TENSOR_SIZE) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The size of weightScale should be 1, current size is %lu.", + gmmDsqParams_.weightScale->Size()); + return false; + } + if (outputDimNumber != MX_OUTPUT_DIM) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The dim num of output should be equal 2, current dim is %lu.", + outputDimNumber); + return false; + } + if (outputScaleDimNumber != MX_OUTPUT_SCALE_DIM) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The dim num of outputScale should be equal 3, current dim is %lu.", + outputScaleDimNumber); + return false; + } + auto weightDimNumber = ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDimNum(); + auto weightScaleDimNumber = ((*gmmDsqParams_.weightScale)[0])->GetViewShape().GetDimNum(); + if (weightScaleDimNumber != MX_WEIGHT_SCALE_DIM) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The dim num of weightScale should be equal 2, current dim is %lu.", + weightScaleDimNumber); + return false; + } + if (weightDimNumber != MX_WEIGHT_DIM) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The dim num of weight should be equal 3, current dim is %lu.", + weightDimNumber); + return false; + } + return true; + } + bool CheckInputOutDimsForPertoken() + { + auto xDimNumber = gmmDsqParams_.x->GetViewShape().GetDimNum(); + auto xScaleDimNumber = gmmDsqParams_.xScale->GetViewShape().GetDimNum(); + auto outputDimNumber = gmmDsqParams_.output->GetViewShape().GetDimNum(); + auto outputScaleDimNumber = gmmDsqParams_.outputScale->GetViewShape().GetDimNum(); + if (xDimNumber != PERTOKEN_X_DIM) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The dim num of x should be equal 2, current dim is %lu.", xDimNumber); + return false; + } + if (xScaleDimNumber != PERTOKEN_X_SCALE_DIM) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The dim num of xScale should be equal 1, current dim is %lu.", + xScaleDimNumber); + return false; + } + if (outputDimNumber != PERTOKEN_OUTPUT_DIM) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The dim num of output should be equal 2, current dim is %lu.", + outputDimNumber); + return false; + } + if (outputScaleDimNumber != PERTOKEN_OUTPUT_SCALE_DIM) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The dim num of outputScale should be equal 1, current dim is %lu.", + outputScaleDimNumber); + return false; + } + if (gmmDsqParams_.weight->Size() != SINGLE_TENSOR_SIZE) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The size of weight should be 1, current size is %lu.", + gmmDsqParams_.weight->Size()); + return false; + } + if (gmmDsqParams_.weightScale->Size() != SINGLE_TENSOR_SIZE) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The size of weightScale should be 1, current size is %lu.", + gmmDsqParams_.weightScale->Size()); + return false; + } + auto weightDimNumber = ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDimNum(); + auto weightScaleDimNumber = ((*gmmDsqParams_.weightScale)[0])->GetViewShape().GetDimNum(); + if (weightScaleDimNumber != PERTOKEN_WEIGHT_SCALE_DIM) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The dim num of weightScale should be equal 2, current dim is %lu.", + weightScaleDimNumber); + return false; + } + if (weightDimNumber != PERTOKEN_WEIGHT_DIM) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The dim num of weight should be equal 3, current dim is %lu.", + weightDimNumber); + return false; + } + return true; + } + + bool CheckInputOutShape() override + { + int64_t groupListLen = gmmDsqParams_.groupList->GetViewShape().GetDim(0); + if (groupListLen > MAX_GROUP_LIST_SIZE) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "The length of groupList should not be greater than 1024, but actual is %ld.", groupListLen); + return false; + } + // 从x的第1维获取k + int64_t kInX = gmmDsqParams_.x->GetViewShape().GetDim(1); + // 根据是否转置从weight中读取维度k + int64_t kInWeight = gmmDsqParams_.transposeWeight ? ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(2) : + ((*gmmDsqParams_.weight)[0])->GetViewShape().GetDim(1); + if (kInX != kInWeight) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "Expected input tensor x and weight tensor to have consistent k-dimension, but k=%ld in x, while " + "k=%ld in weight.", + kInX, kInWeight); + return false; + } + if (gmmDsqParams_.quantMode == QUNAT_MODE_MX) { + return CheckInputOutShapeForMX(); + } else if (gmmDsqParams_.quantMode == QUNAT_MODE_PERTOKEN) { + return CheckInputOutShapeForPertoken(); + } else { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "Quant mode %d is not supported. Supported modes are 0 (pertoken) and 2 (MX).", + gmmDsqParams_.quantMode); + return false; + } + + return true; + } + + bool CheckInputOutShapeForMX() + { + if (!CheckMXTranspose()) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "CheckMXTranspose failed."); + return false; + } + if (!CheckMXShape()) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "CheckMXShape failed."); + return false; + } + DataType xDtype = gmmDsqParams_.x->GetDataType(); + DataType weightDtype = ((*gmmDsqParams_.weight)[0])->GetDataType(); + if (xDtype == DataType::DT_FLOAT4_E2M1 && weightDtype == DataType::DT_FLOAT4_E2M1) { + return checkMxfp4InputShape(); + } + return true; + } + + bool CheckInputOutShapeForPertoken() + { + if (!CheckPertokenTranspose()) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "CheckPertokenTranspose failed."); + return false; + } + if (!CheckPertokenShape()) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "CheckPerTokenShape failed."); + return false; + } + return true; + } + + bool CheckDtypeValid() override + { + DataType xDtype = gmmDsqParams_.x->GetDataType(); + DataType weightDtype = ((*gmmDsqParams_.weight)[0])->GetDataType(); + DataType xScaleDtype = gmmDsqParams_.xScale->GetDataType(); + DataType weightScaleDtype = ((*gmmDsqParams_.weightScale)[0])->GetDataType(); + const aclTensor *x = gmmDsqParams_.x; + const aclTensor *xScale = gmmDsqParams_.xScale; + const aclTensor *groupList = gmmDsqParams_.groupList; + const aclTensor *output = gmmDsqParams_.output; + const aclTensor *outputScale = gmmDsqParams_.outputScale; + if(std::find(X_DTYPE_SUPPORT_LIST.begin(), X_DTYPE_SUPPORT_LIST.end(), xDtype) == X_DTYPE_SUPPORT_LIST.end() && + std::find(X_DTYPE_SUPPORT_LIST_MXFP4.begin(), X_DTYPE_SUPPORT_LIST_MXFP4.end(), xDtype) == X_DTYPE_SUPPORT_LIST_MXFP4.end() && + std::find(XW_DTYPE_SUPPORT_LIST_PERTOKEN.begin(), XW_DTYPE_SUPPORT_LIST_PERTOKEN.end(), xDtype) == XW_DTYPE_SUPPORT_LIST_PERTOKEN.end()){ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Quant case with x dtype %s is not supported; supported types are: INT8, FLOAT8_E4M3FN, " + "FLOAT8_E5M2, HIFLOAT8, and FLOAT4_E2M1.", op::ToString(xDtype).GetString()); + return false; + } + if(std::find(WEIGHT_DTYPE_SUPPORT_LIST.begin(), WEIGHT_DTYPE_SUPPORT_LIST.end(), weightDtype) == WEIGHT_DTYPE_SUPPORT_LIST.end() && + std::find(WEIGHT_DTYPE_SUPPORT_LIST_MXFP4.begin(), WEIGHT_DTYPE_SUPPORT_LIST_MXFP4.end(), weightDtype) == WEIGHT_DTYPE_SUPPORT_LIST_MXFP4.end() && + std::find(XW_DTYPE_SUPPORT_LIST_PERTOKEN.begin(), XW_DTYPE_SUPPORT_LIST_PERTOKEN.end(), weightDtype) == XW_DTYPE_SUPPORT_LIST_PERTOKEN.end()){ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Quant case with weight dtype %s is not supported; supported types are: INT8, FLOAT8_E4M3FN, " + "FLOAT8_E5M2, HIFLOAT8, and FLOAT4_E2M1.", op::ToString(weightDtype).GetString()); + return false; + } + if (gmmDsqParams_.quantMode == QUNAT_MODE_MX && + (xDtype == DataType::DT_FLOAT8_E4M3FN || xDtype == DataType::DT_FLOAT8_E5M2) && + (weightDtype == DataType::DT_FLOAT8_E4M3FN || weightDtype == DataType::DT_FLOAT8_E5M2)) { + return CheckFp8DtypeValid(x, xScale, groupList, output, outputScale); + } else if (gmmDsqParams_.quantMode == QUNAT_MODE_MX && xDtype == DataType::DT_FLOAT4_E2M1 && + weightDtype == DataType::DT_FLOAT4_E2M1) { + return CheckFp4DtypeValid(x, xScale, groupList, output, outputScale); + } else if (gmmDsqParams_.quantMode == QUNAT_MODE_PERTOKEN && + std::find(XW_DTYPE_SUPPORT_LIST_PERTOKEN.begin(), XW_DTYPE_SUPPORT_LIST_PERTOKEN.end(), xDtype) != + XW_DTYPE_SUPPORT_LIST_PERTOKEN.end() && + std::find(XW_DTYPE_SUPPORT_LIST_PERTOKEN.begin(), XW_DTYPE_SUPPORT_LIST_PERTOKEN.end(), + weightDtype) != XW_DTYPE_SUPPORT_LIST_PERTOKEN.end()) { + return CheckPertokenDtypeValid(x, xScale, groupList, output, outputScale); + } else { + OP_LOGE( + ACLNN_ERR_PARAM_INVALID, + "In quantization mode %d, the combination of x dtype %s, weight dtype %s is not supported. " + "Supported combinations are: " + "Quantmode 0 (pertoken): (x=int8, weight=int8) or (x=float8_e4m3fn/float8_e5m2, " + "weight=float8_e4m3fn/float8_e5m2) or (x=hifloat8, weight=hifloat8); " + "Quantmode 2 (mx): (x=float8_e4m3fn/float8_e5m2, weight=float8_e4m3fn/float8_e5m2) or (x=float4_e2m1, weight=float4_e2m1).", + gmmDsqParams_.quantMode, op::ToString(xDtype).GetString(), op::ToString(weightDtype).GetString()); + return false; + } + return true; + } + + bool CheckFormat() override + { + size_t wLength = gmmDsqParams_.weight->Size(); + for (size_t i = 0; i < wLength; i++) { + const aclTensor *weightScale = (*gmmDsqParams_.weightScale)[i]; + const aclTensor *weight = (*gmmDsqParams_.weight)[i]; + if (op::IsPrivateFormat(weight->GetStorageFormat())) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Format of weight should be ND, current format is format is %s.", + op::ToString(weight->GetStorageFormat()).GetString()); + return false; + } + if (op::IsPrivateFormat(weightScale->GetStorageFormat())) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Format of weightScale should be ND, current format is format is %s.", + op::ToString(weightScale->GetStorageFormat()).GetString()); + return false; + } + } + + if (op::IsPrivateFormat(gmmDsqParams_.x->GetStorageFormat())) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Format of x should be ND, current format is format is %s.", + op::ToString(gmmDsqParams_.x->GetStorageFormat()).GetString()); + return false; + } + if (op::IsPrivateFormat(gmmDsqParams_.xScale->GetStorageFormat())) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Format of xScale should be ND, current format is format is %s.", + op::ToString(gmmDsqParams_.xScale->GetStorageFormat()).GetString()); + return false; + } + if (op::IsPrivateFormat(gmmDsqParams_.groupList->GetStorageFormat())) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Format of groupList should be ND, current format is format is %s.", + op::ToString(gmmDsqParams_.groupList->GetStorageFormat()).GetString()); + return false; + } + if (op::IsPrivateFormat(gmmDsqParams_.output->GetStorageFormat())) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Format of output should be ND, current format is format is %s.", + op::ToString(gmmDsqParams_.output->GetStorageFormat()).GetString()); + return false; + } + if (op::IsPrivateFormat(gmmDsqParams_.outputScale->GetStorageFormat())) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Format of outputScale should be ND, current format is format is %s.", + op::ToString(gmmDsqParams_.outputScale->GetStorageFormat()).GetString()); + return false; + } + return true; + } +}; +} // namespace gmmSwigluQuantV2 +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/arch35/grouped_matmul_swiglu_quant_v2_mxquant.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/arch35/grouped_matmul_swiglu_quant_v2_mxquant.h new file mode 100644 index 000000000..e248fba9d --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/arch35/grouped_matmul_swiglu_quant_v2_mxquant.h @@ -0,0 +1,76 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_mxquant.h + * \brief + */ + +#ifndef GROUPED_MATMUL_SWIGLU_QUANT_V2_MXQUANT_H +#define GROUPED_MATMUL_SWIGLU_QUANT_V2_MXQUANT_H + +#include "cgmct/kernel/kernel_gmm_swiglu_mxquant.h" +#include "cgmct/block/block_mx_mm_aic_to_aiv_builder.h" +#include "cgmct/block/block_scheduler_gmm_aswt_with_tail_split.h" + +using namespace Cgmct::Gemm; +using namespace Cgmct::Gemm::Kernel; + +template +__aicore__ inline void GmmSwigluAswt(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, GM_ADDR xScale, + GM_ADDR weightAssistanceMatrix, GM_ADDR smoothScale, GM_ADDR groupList, + GM_ADDR y, GM_ADDR yScale, GM_ADDR workspace, GM_ADDR tiling) +{ + GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingDataParams, gmmSwigluQuantParams, gmmSwigluQuantParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingDataParams, mmTilingData, mmTilingData_, tiling); \ + // 定义L1和L0的TileShape + using L1TileShape = AscendC::Shape<_0, _0, _0>; + using L0TileShape = AscendC::Shape<_0, _0, _0>; + // 定义矩阵的类型和布局 + using AType = DTYPE_X; + using BType = DTYPE_WEIGHT; + using CType = DTYPE_Y; + using LayoutA = layoutA; + using LayoutB = layoutB; + using LayoutC = layout::RowMajorAlign; + using weightscaleType = AscendC::fp8_e8m0_t; + using BiasType = float; + // 定义scheduler类型 + using BlockScheduler = GroupedMatmulAswtWithTailSplitScheduler; + // 定义MMAD类型 + using C1Type = float; + // 定义BlockEpilogue类型 + using BlockEpilogue = Block::BlockEpilogueSwigluQuant; + // 定义shape的形状，tuple保存 m n k batch + using ProblemShape = MatmulShape; + using BlockMmad = Block::BlockMxMmAicToAivBuilder, + Tile::TileCopy>; + using QGmmKernel = + Kernel::KernelGmmSwiGluMixOnlineDynamic; + using Params = typename QGmmKernel::Params; + using GMMTiling = typename QGmmKernel::GMMTiling; + GMMTiling gmmParams{gmmSwigluQuantParams_.groupNum, gmmSwigluQuantParams_.groupListType, mmTilingData_.baseM, + mmTilingData_.baseN, mmTilingData_.baseK}; + gmmParams.matmulTiling = &mmTilingData_; + Params params = {// template shape, gmm shape can not get now + {1, 1, 1, 1}, + // mmad args + {x, weight, weightScale, xScale, y, groupList}, + {y, yScale, nullptr, nullptr, nullptr, static_cast(mmTilingData_.baseM), + static_cast(mmTilingData_.baseN)}, + // gmm tiling data + gmmParams}; + QGmmKernel op; + op(params); +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/arch35/grouped_matmul_swiglu_quant_v2_pertoken_quant.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/arch35/grouped_matmul_swiglu_quant_v2_pertoken_quant.h new file mode 100644 index 000000000..9a76c26ad --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/arch35/grouped_matmul_swiglu_quant_v2_pertoken_quant.h @@ -0,0 +1,115 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_pertoken_quant.h + * \brief + */ + +#ifndef GROUPED_MATMUL_SWIGLU_QUANT_V2_PERTOKEN_QUANT_H +#define GROUPED_MATMUL_SWIGLU_QUANT_V2_PERTOKEN_QUANT_H + +#include "cgmct/kernel/kernel_gmm_swiglu_pertoken_quant.h" +#include "cgmct/block/block_mmad_builder.h" +#include "cgmct/block/block_scheduler_gmm_aswt_with_tail_split.h" + +using namespace Cgmct::Gemm; +using namespace Cgmct::Gemm::Kernel; + +static constexpr uint8_t BF16_VALUE = 27; + +template +__aicore__ inline void GmmSwigluAswtPertokenKernel(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, GM_ADDR xScale, + GM_ADDR weightAssistanceMatrix, GM_ADDR smoothScale, + GM_ADDR groupList, GM_ADDR y, GM_ADDR yScale, GM_ADDR workspace, + GM_ADDR tiling, TPipe *pipe) +{ + /* 1. 取 tiling 数据 */ + GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingDataParams, gmmSwigluQuantParams, gmmSwigluQuantParams_, tiling); + GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingDataParams, mmTilingData, mmTilingData_, tiling); + + /* 2. 编译期常量决定 DequantType / C1Type */ + using DequantType = + std::conditional_t>; + using AType = DTYPE_X; + using BType = DTYPE_WEIGHT; + using CType = DTYPE_Y; // y dtype + using C1Type = std::conditional_t, int32_t, float>; // matmul output dtype + + /* 3. 其余别名 */ + using L0TileShape = AscendC::Shape<_0, _0, _0>; + using L1TileShape = AscendC::Shape<_0, _0, _0>; + using LayoutA = layoutA; + using LayoutB = layoutB; + using LayoutC = layout::RowMajorAlign; + using weightscaleType = DTYPE_WEIGHT_SCALE; + using xscaleType = float; + using BiasType = float; + using BlockScheduler = GroupedMatmulAswtWithTailSplitScheduler; + using BlockEpilogueDequantAndSwiglu = + Block::BlockEpilogueDequantSwiglu; + using BlockEpiloguePertokenQuant = Block::BlockEpiloguePertokenQuant; + using ProblemShape = MatmulShape; + using BlockMmad = + Block::BlockMmadBuilder, + Tile::TileCopy>; + using QGmmKernel = + Kernel::KernelGmmSwiGluPertokenQuant; + + /* 4. 拼参数、launch */ + using Params = typename QGmmKernel::Params; + using GMMTiling = typename QGmmKernel::GMMTiling; + GMMTiling gmmParams{gmmSwigluQuantParams_.groupNum, gmmSwigluQuantParams_.groupListType, mmTilingData_.baseM, + mmTilingData_.baseN, mmTilingData_.baseK}; + gmmParams.matmulTiling = &mmTilingData_; + Params params = { + {1, 1, 1, 1}, + // mmad args + {x, weight, y, nullptr, groupList}, + {workspace, weightScale, xScale, static_cast(mmTilingData_.baseM), + static_cast(mmTilingData_.baseN)}, + {workspace, smoothScale, y, yScale, gmmSwigluQuantParams_.rowLen, gmmSwigluQuantParams_.ubAvail, false}, + // gmm tiling data + gmmParams}; + QGmmKernel op(pipe); + op(params); +} + +/* ---------------------------------------------------------- + * 5. 最外层入口：只做 switch，把运行期值 → 编译期常量 + * ---------------------------------------------------------- */ +template +__aicore__ inline void GmmSwigluAswtPertoken(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, GM_ADDR xScale, + GM_ADDR weightAssistanceMatrix, GM_ADDR smoothScale, GM_ADDR groupList, + GM_ADDR y, GM_ADDR yScale, GM_ADDR workspace, GM_ADDR tiling, TPipe *pipe) +{ + GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingDataParams, gmmSwigluQuantParams, gmmSwigluQuantParams_, tiling); + + switch (gmmSwigluQuantParams_.dequantDtype) { + case 1: + GmmSwigluAswtPertokenKernel<1, layoutA, layoutB>(x, weight, weightScale, xScale, weightAssistanceMatrix, + smoothScale, groupList, y, yScale, workspace, tiling, + pipe); + break; + case BF16_VALUE: + GmmSwigluAswtPertokenKernel(x, weight, weightScale, xScale, + weightAssistanceMatrix, smoothScale, groupList, y, + yScale, workspace, tiling, pipe); + break; + default: + GmmSwigluAswtPertokenKernel<0, layoutA, layoutB>(x, weight, weightScale, xScale, weightAssistanceMatrix, + smoothScale, groupList, y, yScale, workspace, tiling, + pipe); + break; + } +} +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/arch35/grouped_matmul_swiglu_quant_v2_tiling_key.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/arch35/grouped_matmul_swiglu_quant_v2_tiling_key.h new file mode 100644 index 000000000..ee7523a0f --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/arch35/grouped_matmul_swiglu_quant_v2_tiling_key.h @@ -0,0 +1,40 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_tiling_key.h + * \brief + */ + +#ifndef __OP_KERNEL_GMM_SWIGLU_QUANT_V2_TILING_KEY_H__ +#define __OP_KERNEL_GMM_SWIGLU_QUANT_V2_TILING_KEY_H__ + +#include "ascendc/host_api/tiling/template_argument.h" + +#define GMM_SWIGLU_QUANT_NO_TRANS 0 +#define GMM_SWIGLU_QUANT_TRANS 1 + +// 模板参数 +ASCENDC_TPL_ARGS_DECL(GroupedMatmulSwigluQuantV2, // 算子OpType + ASCENDC_TPL_UINT_DECL(QUANT_B_TRANS, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, + GMM_SWIGLU_QUANT_NO_TRANS, GMM_SWIGLU_QUANT_TRANS), + ASCENDC_TPL_UINT_DECL(QUANT_A_TRANS, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, + GMM_SWIGLU_QUANT_NO_TRANS, GMM_SWIGLU_QUANT_TRANS)); + +// 模板参数组合 +// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法 +ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_SWIGLU_QUANT_NO_TRANS), + ASCENDC_TPL_UINT_SEL(QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_SWIGLU_QUANT_NO_TRANS)), + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2), + ASCENDC_TPL_UINT_SEL(QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_SWIGLU_QUANT_TRANS), + ASCENDC_TPL_UINT_SEL(QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_SWIGLU_QUANT_NO_TRANS))); +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_spilit_fusion.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_spilit_fusion.h new file mode 100644 index 000000000..8ea7b6590 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_spilit_fusion.h @@ -0,0 +1,591 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grrouped_matmul_swiglu_quant_spilit_fusion.h + * \brief + */ + +#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_SPLIT_FUSION_H +#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_SPLIT_FUSION_H + +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#include "grouped_matmul_swiglu_quant_v2_utils.h" + +namespace GroupedMatmulDequantSwigluQuant { +using namespace AscendC; +constexpr int64_t BLOCK_SIZE = 32; +constexpr int64_t BLOCK_ELEM = BLOCK_SIZE / sizeof(float); +constexpr int64_t SWI_FACTOR = 2; +constexpr float DYNAMIC_QUANT_FACTOR = 1.0 / static_cast(127.0); +constexpr uint64_t MAX_CALC_NUM = 64; +constexpr uint64_t REDUCEMAX_CALC_NUM = 64; +constexpr uint64_t SPILI_NUM = 2; +constexpr uint64_t VC_SYNC_MAX_TIMES = 14; +constexpr uint64_t RESRERVE_MEM_SIZE = 192; + +class GroupedMatmulDequantSwigluQuantFusion { +public: + using aType = MatmulType; + using bType = MatmulType; + using cType = MatmulType; + using biasType = MatmulType; + using matmulType = MMImplType; + matmulType::MT mm; + + __aicore__ inline GroupedMatmulDequantSwigluQuantFusion( + TPipe* pipe, const GMMSwigluQuantV2TilingFusionData* __restrict tiling, + const TCubeTiling* __restrict matmulTilingData) + : pipe_(pipe), tilingData_(tiling), matmulTilingData_(matmulTilingData) { + } + + __aicore__ inline int CeilDiv(int a, int b) { + return (a + b - 1) / b; + } + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR weight_scale, GM_ADDR activation_scale, + GM_ADDR weightAssistanceMatrix, GM_ADDR group_list, + GM_ADDR y, GM_ADDR scale, GM_ADDR workspace) { + xGm_.SetGlobalBuffer((__gm__ int8_t*)x); + groupListGm_.SetGlobalBuffer((__gm__ int64_t*)group_list); + weightGm_.SetGlobalBuffer(GetTensorAddr(0, weight)); + weightScaleGm_.SetGlobalBuffer(GetTensorAddr(0, weight_scale)); + workspaceGm_.SetGlobalBuffer((__gm__ int32_t*)workspace); + activateScaleGm_.SetGlobalBuffer((__gm__ float*)activation_scale); + scaleGm_.SetGlobalBuffer((__gm__ float*)scale); + yGm_.SetGlobalBuffer((__gm__ int8_t*)y); + weightScaleTensorPtr_ = weight_scale; + weightTensorPtr_ = weight; + + nBasicsBlocks = CeilDiv(tilingData_->N, matmulTilingData_->baseN); + totalBasicBlocks = 0; + for (int groupId = 0; groupId < tilingData_->groupNum; groupId++) { + int tokens = groupListGm_.GetValue(groupId); + if (tilingData_->groupListType == 0 && groupId > 0) { + tokens = groupListGm_.GetValue(groupId) - groupListGm_.GetValue(groupId - 1); + } + int mBasicBlocks = CeilDiv(tokens, matmulTilingData_->baseM); + totalBasicBlocks += mBasicBlocks * nBasicsBlocks; + } + + totalSyncTimes = CeilDiv(totalBasicBlocks, tilingData_->cubeBlockDim); + if ASCEND_IS_AIV { + pipe_->InitBuffer(xActQueue_, 1, (tilingData_->ubFactorDimx * (tilingData_->N / SPILI_NUM) * SWI_FACTOR + tilingData_->ubFactorDimx * BLOCK_ELEM) * sizeof(int32_t)); + pipe_->InitBuffer(inScaleQueue_, 1, ((tilingData_->N / SPILI_NUM) * SWI_FACTOR + (tilingData_->N / SPILI_NUM)) * sizeof(float)); + pipe_->InitBuffer(outQueue_, 1, tilingData_->ubFactorDimx * (tilingData_->N / SPILI_NUM) * sizeof(int8_t) + tilingData_->ubFactorDimx * sizeof(float) + RESRERVE_MEM_SIZE); + pipe_->InitBuffer(tmpBuf1_, tilingData_->ubFactorDimx * (tilingData_->N / SPILI_NUM) * SWI_FACTOR * sizeof(float)); + } + } + + __aicore__ inline void FindCurrentGroup(uint32_t basicBlockIdxInGlobal, uint32_t& currentGroupId, + uint32_t& globalMOffset, uint32_t& processedBasicBlock) { + for (int groupId = currentGroupId; groupId < tilingData_->groupNum; groupId++) { + int tokens = groupListGm_.GetValue(groupId); + if (tilingData_->groupListType == 0 && groupId > 0) { + tokens = groupListGm_.GetValue(groupId) - groupListGm_.GetValue(groupId - 1); + } + int mBasicBlocks = CeilDiv(tokens, matmulTilingData_->baseM); + if (processedBasicBlock + mBasicBlocks * nBasicsBlocks > basicBlockIdxInGlobal) { + currentGroupId = groupId; + break; + } else { + globalMOffset += tokens; + processedBasicBlock += mBasicBlocks * nBasicsBlocks; + } + } + } + + __aicore__ inline void CalculateBlockSizes(int tokens, int currentBasicBlockMId, int currentBasicBlockNId, + int& realMSize, int& realNSize) { + realMSize = matmulTilingData_->baseM; + if (currentBasicBlockMId * matmulTilingData_->baseM + realMSize > tokens) { + realMSize = tokens - currentBasicBlockMId * matmulTilingData_->baseM; + } + realNSize = matmulTilingData_->baseN; + if (currentBasicBlockNId * matmulTilingData_->baseN + realNSize > tilingData_->N) { + realNSize = tilingData_->N - currentBasicBlockNId * matmulTilingData_->baseN; + } + } + + __aicore__ inline void SetupMatmulShape(int tokens, int realMSize, int realNSize) { + mm.SetOrgShape(tokens, tilingData_->N, tilingData_->K); + mm.SetSingleShape(realMSize, realNSize, tilingData_->K); + } + + __aicore__ inline void SetupMatmulWeight(int currentGroupId, int currentBasicBlockNId) { + if (tilingData_->isSingleTensor == 0) { + weightGm_.SetGlobalBuffer(GetTensorAddr(currentGroupId, weightTensorPtr_)); + mm.SetTensorB(weightGm_[0x8 * currentBasicBlockNId * tilingData_->K * 0x20]); + } else { + int64_t tensorBOffset = currentGroupId * tilingData_->K * tilingData_->N + 0x8 * currentBasicBlockNId * tilingData_->K * 0x20; + mm.SetTensorB(weightGm_[tensorBOffset]); + } + } + + __aicore__ inline void ProcessCubeBlock(uint32_t basicBlockIdxInGlobal, uint32_t& currentGroupId, + uint32_t& globalMOffset, uint32_t& processedBasicBlock) { + FindCurrentGroup(basicBlockIdxInGlobal, currentGroupId, globalMOffset, processedBasicBlock); + int tokens = groupListGm_.GetValue(currentGroupId); + if (tilingData_->groupListType == 0 && currentGroupId > 0) { + tokens = groupListGm_.GetValue(currentGroupId) - groupListGm_.GetValue(currentGroupId - 1); + } + int basicBlockIdxInCurrentGroup = basicBlockIdxInGlobal - processedBasicBlock; + int mBasicBlocks = CeilDiv(tokens, matmulTilingData_->baseM); + int currentBasicBlockMId = basicBlockIdxInCurrentGroup / nBasicsBlocks; + int currentBasicBlockNId = basicBlockIdxInCurrentGroup % nBasicsBlocks; + int realMSize = 0; + int realNSize = 0; + CalculateBlockSizes(tokens, currentBasicBlockMId, currentBasicBlockNId, realMSize, realNSize); + SetupMatmulShape(tokens, realMSize, realNSize); + int64_t tensorAOffset = currentBasicBlockMId * matmulTilingData_->baseM * tilingData_->K + globalMOffset * tilingData_->K; + mm.SetTensorA(xGm_[tensorAOffset]); + SetupMatmulWeight(currentGroupId, currentBasicBlockNId); + int64_t workspaceOffset = globalMOffset * tilingData_->N + currentBasicBlockMId * matmulTilingData_->baseM * tilingData_->N + + currentBasicBlockNId * matmulTilingData_->baseN; + mm.template IterateAll(workspaceGm_[workspaceOffset]); + } + + __aicore__ inline void FinalizeCubeSync(uint32_t& syncId) { + while (syncId < totalSyncTimes) { + AscendC::CrossCoreSetFlag<0x2, PIPE_FIX>(0x8); + syncId += 1; + } + } + + __aicore__ inline void CubeProcess() { + if ASCEND_IS_AIC { + uint32_t currentBlockId = GetBlockIdx(); + uint32_t rsvBlockNum = 0; + uint32_t calcBlockNum = 0; + uint32_t cvTimes = 0; + uint32_t syncId = 0; + uint32_t globalMOffset = 0; + uint32_t processedBasicBlock = 0; + uint32_t currentGroupId = 0; + uint32_t realSyncId = 0; + while (currentBlockId < totalBasicBlocks) { + cvTimes = CeilDiv(nBasicsBlocks - rsvBlockNum, tilingData_->cubeBlockDim); + calcBlockNum += cvTimes * tilingData_->cubeBlockDim; + rsvBlockNum = calcBlockNum % nBasicsBlocks; + + for (uint32_t cvId = 0; cvId < cvTimes; cvId++) { + uint32_t basicBlockIdxInGlobal = currentBlockId; + if (basicBlockIdxInGlobal >= totalBasicBlocks) { + break; + } + ProcessCubeBlock(basicBlockIdxInGlobal, currentGroupId, globalMOffset, processedBasicBlock); + currentBlockId += tilingData_->cubeBlockDim; + syncId += 1; + } + AscendC::CrossCoreSetFlag<0x2, PIPE_FIX>(0x8); + + realSyncId += 1; + if (realSyncId > 0 && realSyncId % VC_SYNC_MAX_TIMES == 0) { + AscendC::CrossCoreWaitFlag(0x9); + } + } + FinalizeCubeSync(syncId); + } + } + + __aicore__ inline void CalculateEndGroupInfo(int endBasicBlockId, int endGroupId, int& endGroupMOffset, + int& basicBlockCountBeforeEndGroup) { + endGroupMOffset = 0; + basicBlockCountBeforeEndGroup = 0; + for (int gId = 0; gId < endGroupId; gId++) { + int tokens = groupListGm_.GetValue(gId); + if (tilingData_->groupListType == 0 && gId > 0) { + tokens = groupListGm_.GetValue(gId) - groupListGm_.GetValue(gId - 1); + } + int mBasicBlocks = CeilDiv(tokens, matmulTilingData_->baseM); + basicBlockCountBeforeEndGroup += mBasicBlocks * nBasicsBlocks; + endGroupMOffset += tokens; + } + int basicBlockIdxInCurrentGroup = endBasicBlockId - basicBlockCountBeforeEndGroup; + int currentBasicBlockMId = basicBlockIdxInCurrentGroup / nBasicsBlocks; + endGroupMOffset += currentBasicBlockMId * matmulTilingData_->baseM; + } + + __aicore__ inline void ProcessGroupRange(int startGroupId, int endGroupId, int endGroupMOffset, + uint32_t& globalMOffset, bool &isSyncAll) { + int currentGroupMOffset = 0; + for (int gId = 0; gId < startGroupId; gId++) { + currentGroupMOffset += groupListGm_.GetValue(gId); + } + for (int groupId = startGroupId; groupId <= endGroupId; groupId++) { + if (tilingData_->groupListType == 0 && groupId > 0) { + currentGroupMOffset = groupListGm_.GetValue(groupId); + } else { + currentGroupMOffset += groupListGm_.GetValue(groupId); + } + int calcCount = 0; + if (currentGroupMOffset <= endGroupMOffset) { + calcCount = currentGroupMOffset - globalMOffset; + } else { + calcCount = endGroupMOffset - globalMOffset; + } + ProcessDSQ(groupId, globalMOffset, calcCount, isSyncAll); + globalMOffset += calcCount; + } + } + + __aicore__ inline void ProcessVectorBlock(uint32_t syncId, bool& isSyncAll, uint32_t& globalMOffset) { + int startBasicBlockId = syncId * tilingData_->cubeBlockDim; + int endBasicBlockId = startBasicBlockId + tilingData_->cubeBlockDim; + if (totalBasicBlocks < endBasicBlockId) { + endBasicBlockId = totalBasicBlocks; + } + int startGroupId = GetGroupId(startBasicBlockId); + int endGroupId = GetGroupId(endBasicBlockId); + int endGroupMOffset = 0; + int basicBlockCountBeforeEndGroup = 0; + CalculateEndGroupInfo(endBasicBlockId, endGroupId, endGroupMOffset, basicBlockCountBeforeEndGroup); + ProcessGroupRange(startGroupId, endGroupId, endGroupMOffset, globalMOffset, isSyncAll); + } + + __aicore__ inline void VectorProcess() { + if ASCEND_IS_AIV { + weightCacheGroupId_ = -1; + uint32_t currentBlockId = GetBlockIdx() / 2; + uint32_t rsvBlockNum = 0; + uint32_t calcBlockNum = 0; + uint32_t cvTimes = 0; + uint32_t syncId = 0; + uint32_t globalMOffset = 0; + uint32_t processedBasicBlock = 0; + uint32_t currentGroupId = 0; + uint32_t realSyncId = 0; + bool isSyncAll = false; + while (syncId < totalSyncTimes) { + cvTimes = CeilDiv(nBasicsBlocks - rsvBlockNum, tilingData_->cubeBlockDim); + calcBlockNum += cvTimes * tilingData_->cubeBlockDim; + rsvBlockNum = calcBlockNum % nBasicsBlocks; + isSyncAll = true; + + for (uint32_t cvId = 0; cvId < cvTimes; cvId++) { + ProcessVectorBlock(syncId, isSyncAll, globalMOffset); + currentBlockId += tilingData_->cubeBlockDim; + syncId += 1; + } + + realSyncId += 1; + if (realSyncId > 0 && (realSyncId % VC_SYNC_MAX_TIMES == 0)) { + AscendC::CrossCoreSetFlag<0x2, PIPE_MTE2>(0x9); + } + } + } + } + + __aicore__ inline void Process() { + CubeProcess(); + VectorProcess(); + } + + __aicore__ inline int GetGroupId(int basicBlockId) { + int processedBasicBlock = 0; + int currentGroupId = 0; + int globalMOffset = 0; + for (int groupId = 0; groupId < tilingData_->groupNum; groupId++) { + int tokens = groupListGm_.GetValue(groupId); + if (tilingData_->groupListType == 0 && groupId > 0) { + tokens = groupListGm_.GetValue(groupId) - groupListGm_.GetValue(groupId - 1); + } + int mBasicBlocks = CeilDiv(tokens, matmulTilingData_->baseM); + if (processedBasicBlock + mBasicBlocks * nBasicsBlocks >= basicBlockId) { + return groupId; + } else { + processedBasicBlock += mBasicBlocks * nBasicsBlocks; + } + } + return tilingData_->groupNum - 1; + } + + __aicore__ inline void ComputeReduceMax(const LocalTensor& tempRes, int32_t calcCount) { + uint32_t vectorCycles = calcCount / MAX_CALC_NUM; + uint32_t remainElements = calcCount % MAX_CALC_NUM; + + BinaryRepeatParams repeatParams; + repeatParams.dstBlkStride = 1; + repeatParams.src0BlkStride = 1; + repeatParams.src1BlkStride = 1; + repeatParams.dstRepStride = 0; + repeatParams.src0RepStride = 0x8; + repeatParams.src1RepStride = 0; + + if (vectorCycles > 0 && remainElements > 0) { + Max(tempRes, tempRes, tempRes[vectorCycles * MAX_CALC_NUM], remainElements, 1, repeatParams); + PipeBarrier(); + } + + if (vectorCycles > 1) { + Max(tempRes, tempRes[MAX_CALC_NUM], tempRes, MAX_CALC_NUM, vectorCycles - 1, repeatParams); + PipeBarrier(); + } + } + + __aicore__ inline void ProcessDSQ(int groupId, int globalOffset, int calcCount, bool &isSyncAll) { + int32_t blockDimxFactor = (calcCount + tilingData_->vectorBlockDim - 1) / tilingData_->vectorBlockDim; + int32_t realCoreDim = calcCount == 0 ? 0 : (calcCount + blockDimxFactor - 1) / blockDimxFactor; + + if (GetBlockIdx() >= realCoreDim) { + if (isSyncAll) { + AscendC::CrossCoreWaitFlag(0x8); + SyncAll(); + isSyncAll = false; + } + return; + } + + DataCopyPadParams padParams{false, 0, 0, 0}; + LocalTensor inScaleLocal = inScaleQueue_.AllocTensor(); + + if (weightCacheGroupId_ != groupId) { + DataCopyParams dataCopyWeightScaleParams; + dataCopyWeightScaleParams.blockCount = 1; + dataCopyWeightScaleParams.blockLen = tilingData_->N * sizeof(float); + dataCopyWeightScaleParams.srcStride = 0; + dataCopyWeightScaleParams.dstStride = 0; + if (tilingData_->isSingleTensor == 0) { + weightScaleGm_.SetGlobalBuffer(GetTensorAddr(groupId, weightScaleTensorPtr_)); + DataCopyPad(inScaleLocal, weightScaleGm_, dataCopyWeightScaleParams, padParams); + } else { + DataCopyPad(inScaleLocal, weightScaleGm_[groupId * tilingData_->N], dataCopyWeightScaleParams, padParams); + } + DataCopyParams dataCopyQuantScaleParams; + dataCopyQuantScaleParams.blockCount = 1; + dataCopyQuantScaleParams.blockLen = (tilingData_->N / SPILI_NUM) * sizeof(float); + dataCopyQuantScaleParams.srcStride = 0; + dataCopyQuantScaleParams.dstStride = 0; + weightCacheGroupId_ = groupId; + } + + inScaleQueue_.EnQue(inScaleLocal); + inScaleLocal = inScaleQueue_.DeQue(); + + int32_t blockDimxTailFactor = calcCount - blockDimxFactor * (realCoreDim - 1); + int32_t DimxCore = GetBlockIdx() == (realCoreDim - 1) ? blockDimxTailFactor : blockDimxFactor; + + int32_t ubDimxLoop = (DimxCore + tilingData_->ubFactorDimx - 1) / tilingData_->ubFactorDimx; + int32_t ubDimxTailFactor = DimxCore - tilingData_->ubFactorDimx * (ubDimxLoop - 1); + + int64_t coreDimxOffset = blockDimxFactor * GetBlockIdx(); + int32_t actOffset = tilingData_->actRight * tilingData_->ubFactorDimy; + int32_t gateOffset = tilingData_->ubFactorDimy - actOffset; + + LocalTensor weightScaleLocal = inScaleLocal; + LocalTensor quantScaleLocal = inScaleLocal[tilingData_->N]; + + for (uint32_t loopIdx = 0; loopIdx < ubDimxLoop; loopIdx++) { + int64_t xDimxOffset = (coreDimxOffset + loopIdx * tilingData_->ubFactorDimx) + globalOffset; + int32_t proDimsx = loopIdx == (ubDimxLoop - 1) ? ubDimxTailFactor : tilingData_->ubFactorDimx; + LocalTensor tmpUbF32 = tmpBuf1_.AllocTensor(); + SetMaskCount(); + SetVectorMask(tilingData_->ubFactorDimy * SWI_FACTOR); + Copy(tmpUbF32, weightScaleLocal, MASK_PLACEHOLDER, proDimsx, + {1, 1, static_cast((tilingData_->ubFactorDimy * SWI_FACTOR) / BLOCK_ELEM), 0}); + SetMaskNorm(); + ResetMask(); + + LocalTensor xActLocal = xActQueue_.AllocTensor(); + DataCopyParams dataCopyActScaleParams; + dataCopyActScaleParams.blockCount = proDimsx; + dataCopyActScaleParams.blockLen = sizeof(float); + dataCopyActScaleParams.srcStride = 0; + dataCopyActScaleParams.dstStride = 0; + LocalTensor xActLocalF32 = xActLocal.template ReinterpretCast(); + DataCopyPad(xActLocalF32[tilingData_->ubFactorDimx * tilingData_->N], activateScaleGm_[xDimxOffset], + dataCopyActScaleParams, padParams); + + if (isSyncAll) { + AscendC::CrossCoreWaitFlag(0x8); + SyncAll(); + isSyncAll = false; + } + + DataCopyParams dataCopyXParams; + dataCopyXParams.blockCount = proDimsx; + dataCopyXParams.blockLen = tilingData_->N * sizeof(int32_t); + dataCopyXParams.srcStride = 0; + dataCopyXParams.dstStride = 0; + DataCopyPad(xActLocal, workspaceGm_[xDimxOffset * tilingData_->N], dataCopyXParams, padParams); + xActQueue_.EnQue(xActLocal); + xActLocal = xActQueue_.DeQue(); + + LocalTensor xLocal = xActLocal; + xActLocalF32 = xActLocal.template ReinterpretCast(); + LocalTensor xLocalF32 = xActLocalF32; + LocalTensor activationScaleLocal = xActLocalF32[tilingData_->ubFactorDimx * tilingData_->N]; + + Cast(xLocalF32, xLocal, RoundMode::CAST_NONE, SWI_FACTOR * proDimsx * tilingData_->ubFactorDimy); + PipeBarrier(); + + Mul(xLocalF32, tmpUbF32, xLocalF32, tilingData_->ubFactorDimy * SWI_FACTOR * proDimsx); + PipeBarrier(); + + SetMaskCount(); + SetVectorMask(tilingData_->ubFactorDimy * SWI_FACTOR); + Copy(tmpUbF32, activationScaleLocal, AscendC::MASK_PLACEHOLDER, proDimsx, + {1, 0, static_cast((tilingData_->ubFactorDimy * SWI_FACTOR) / BLOCK_ELEM), 1}); + SetMaskNorm(); + ResetMask(); + PipeBarrier(); + + Mul(xLocalF32, tmpUbF32, xLocalF32, tilingData_->ubFactorDimy * SWI_FACTOR * proDimsx); + PipeBarrier(); + + LocalTensor tmpUbF32Act = tmpUbF32; + LocalTensor tmpUbF32Gate = tmpUbF32[tilingData_->ubFactorDimy * proDimsx]; + SetMaskCount(); + SetVectorMask(tilingData_->ubFactorDimy); + Copy(tmpUbF32Act, xLocalF32[actOffset], AscendC::MASK_PLACEHOLDER, proDimsx, + {1, 1, static_cast(tilingData_->ubFactorDimy / BLOCK_ELEM), + static_cast(tilingData_->ubFactorDimy / BLOCK_ELEM * SWI_FACTOR)}); + Copy(tmpUbF32Gate, xLocalF32[gateOffset], AscendC::MASK_PLACEHOLDER, proDimsx, + {1, 1, static_cast(tilingData_->ubFactorDimy / BLOCK_ELEM), + static_cast(tilingData_->ubFactorDimy / BLOCK_ELEM * SWI_FACTOR)}); + SetMaskNorm(); + ResetMask(); + PipeBarrier(); + limited=tilingData_->swigluLimit; + + if (limited > 0.0f) { + Mins(tmpUbF32Gate, tmpUbF32Gate, limited, tilingData_->ubFactorDimy * proDimsx); + PipeBarrier(); + Maxs(tmpUbF32Gate, tmpUbF32Gate, (-1.0f * limited), tilingData_->ubFactorDimy * proDimsx); + PipeBarrier(); + Mins(tmpUbF32Act, tmpUbF32Act, limited, tilingData_->ubFactorDimy * proDimsx); + PipeBarrier(); + } + + Muls(xLocalF32, tmpUbF32Act, static_cast(-1.0), tilingData_->ubFactorDimy * proDimsx); + PipeBarrier(); + Exp(xLocalF32, xLocalF32, tilingData_->ubFactorDimy * proDimsx); + PipeBarrier(); + Adds(xLocalF32, xLocalF32, static_cast(1.0), tilingData_->ubFactorDimy * proDimsx); + PipeBarrier(); + Div(tmpUbF32Act, tmpUbF32Act, xLocalF32, tilingData_->ubFactorDimy * proDimsx); + PipeBarrier(); + + xActQueue_.FreeTensor(xActLocal); + Mul(tmpUbF32Act, tmpUbF32Gate, tmpUbF32Act, tilingData_->ubFactorDimy * proDimsx); + PipeBarrier(); + + Abs(tmpUbF32Gate, tmpUbF32Act, tilingData_->ubFactorDimy * proDimsx); + + LocalTensor outLocal = outQueue_.AllocTensor(); + + uint64_t scaleOutOffset = tilingData_->ubFactorDimx * (tilingData_->N / SPILI_NUM) * sizeof(int8_t) / sizeof(float); + uint64_t alignScaleOutOffset = Ceil(scaleOutOffset, uint32_t(8)) * 8; // 8: num int32_t in 32B ub block + LocalTensor scaleOut = outLocal[alignScaleOutOffset]; + LocalTensor yOut = outLocal.template ReinterpretCast(); + PipeBarrier(); + + for (uint32_t i = 0; i < proDimsx; i++) { + ComputeReduceMax(tmpUbF32Gate[i * tilingData_->ubFactorDimy], tilingData_->ubFactorDimy); + } + + uint64_t realReduceMaxCalcNum = REDUCEMAX_CALC_NUM; + if (tilingData_->ubFactorDimy < REDUCEMAX_CALC_NUM) { + realReduceMaxCalcNum = tilingData_->ubFactorDimy; + } + + WholeReduceMax(tmpUbF32Gate, tmpUbF32Gate, realReduceMaxCalcNum, proDimsx, 1, 1, + tilingData_->ubFactorDimy / BLOCK_ELEM, ReduceOrder::ORDER_ONLY_VALUE); + PipeBarrier(); + + Muls(scaleOut, tmpUbF32Gate, DYNAMIC_QUANT_FACTOR, proDimsx); + PipeBarrier(); + + int64_t blockCount = (proDimsx + BLOCK_ELEM - 1) / BLOCK_ELEM; + Brcb(outLocal, scaleOut, blockCount, {1, 8}); + PipeBarrier(); + + SetMaskCount(); + SetVectorMask(tilingData_->ubFactorDimy); + Copy(tmpUbF32Gate, outLocal, AscendC::MASK_PLACEHOLDER, proDimsx, + {1, 0, static_cast(tilingData_->ubFactorDimy / BLOCK_ELEM), 1}); + SetMaskNorm(); + ResetMask(); + PipeBarrier(); + + Div(tmpUbF32Act, tmpUbF32Act, tmpUbF32Gate, tilingData_->ubFactorDimy * proDimsx); + PipeBarrier(); + + LocalTensor tmpUbF32ActI32 = tmpUbF32Act.ReinterpretCast(); + Cast(tmpUbF32ActI32, tmpUbF32Act, RoundMode::CAST_RINT, tilingData_->ubFactorDimy * proDimsx); + SetDeqScale((half)1.000000e+00f); + + LocalTensor tmpUbF32Gate16 = tmpUbF32Gate.template ReinterpretCast(); + Cast(tmpUbF32Gate16, tmpUbF32ActI32, RoundMode::CAST_ROUND, tilingData_->ubFactorDimy * proDimsx); + PipeBarrier(); + + Cast(yOut, tmpUbF32Gate16, RoundMode::CAST_TRUNC, tilingData_->ubFactorDimy * proDimsx); + PipeBarrier(); + + tmpBuf1_.FreeTensor(tmpUbF32); + outQueue_.EnQue(outLocal); + outLocal = outQueue_.DeQue(); + scaleOut = outLocal[alignScaleOutOffset]; + yOut = outLocal.template ReinterpretCast(); + + DataCopyParams dataCopyOutScaleParams; + dataCopyOutScaleParams.blockCount = 1; + dataCopyOutScaleParams.blockLen = proDimsx * sizeof(float); + dataCopyOutScaleParams.srcStride = 0; + dataCopyOutScaleParams.dstStride = 0; + DataCopyPad(scaleGm_[xDimxOffset], scaleOut, dataCopyOutScaleParams); + + DataCopyParams dataCopyOutyParams; + dataCopyOutyParams.blockCount = 1; + dataCopyOutyParams.blockLen = proDimsx * (tilingData_->N / SPILI_NUM) * sizeof(int8_t); + dataCopyOutyParams.srcStride = 0; + dataCopyOutyParams.dstStride = 0; + DataCopyPad(yGm_[xDimxOffset * (tilingData_->N / SPILI_NUM)], yOut, dataCopyOutyParams); + outQueue_.FreeTensor(outLocal); + } + inScaleQueue_.FreeTensor(inScaleLocal); + } + +private: + TPipe *pipe_ = nullptr; + const GMMSwigluQuantV2TilingFusionData* __restrict tilingData_; + const TCubeTiling* __restrict matmulTilingData_; + static constexpr float FLOAT_INF = 3e+99; + GlobalTensor xGm_; + GlobalTensor weightGm_; + GlobalTensor yGm_; + GlobalTensor workspaceGm_; + GlobalTensor weightScaleGm_; + GlobalTensor activateScaleGm_; + GlobalTensor scaleGm_; + GlobalTensor groupListGm_; + int nBasicsBlocks = 0; + int totalBasicBlocks = 0; + int totalSyncTimes = 0; + int32_t weightCacheGroupId_ = -1; + float limited = FLOAT_INF; + + TQue inQue_; + TQue xQue_; + TBuf tmpBuf_; + TQue scaleOutQue_; + TQue yOutQue_; + TQue xActQueue_; + TQue outQueue_; + TQue inScaleQueue_; + TBuf tmpBuf1_; + + GM_ADDR weightTensorPtr_; + GM_ADDR weightScaleTensorPtr_; +}; +} + +#endif diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2.cpp new file mode 100644 index 000000000..5a89e3e59 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2.cpp @@ -0,0 +1,109 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2.cpp + * \brief + */ + +#include "kernel_tiling/kernel_tiling.h" +#include "kernel_operator.h" +#include "lib/matmul_intf.h" +#include "grouped_matmul_swiglu_quant_spilit_fusion.h" +#include "grouped_matmul_swiglu_quant_v2_a8w4_msd_pipeline.h" +#include "grouped_matmul_swiglu_quant_v2_a4w4_pipeline.h" +#include "grouped_matmul_swiglu_quant_v2_utils.h" +using namespace AscendC; +using namespace matmul; +using namespace GroupedMatmulDequantSwigluQuant; +extern "C" __global__ __aicore__ void grouped_matmul_swiglu_quant_v2(GM_ADDR x, GM_ADDR xScale, GM_ADDR groupList, + GM_ADDR weight, GM_ADDR weightScale, + GM_ADDR weightAssistanceMatrix, GM_ADDR bias, + GM_ADDR smoothScale, GM_ADDR y, GM_ADDR yScale, + GM_ADDR workspace, GM_ADDR tiling) +{ + TPipe tPipe; + GM_ADDR userWorkspace = GetUserWorkspace(workspace); + +#if defined(GMM_SWIGLU_QUANT_V2_A8W4_MSD) + if (TILING_KEY_IS(2)) { + KERNEL_TASK_TYPE(2, KERNEL_TYPE_MIX_AIC_1_2); + GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, gmmSwigluQuantV2BaseParams, gmmSwigluQuantV2BaseParams_, + tiling); + GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, mmTilingData, mmTilingData_, tiling); + GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, gmmSwigluQuantV2, gmmSwiglu_, tiling); + using xType = MatmulType; + using weightType = MatmulType; + using yType = MatmulType; + using matmulType = MMImplTypeCustom; + matmulType::MT mm; + if ASCEND_IS_AIC { + mm.SetSubBlockIdx(0); + mm.Init(&mmTilingData_); + } + GMMSwigluQuantPipelineSchedule op(mm, &gmmSwigluQuantV2BaseParams_, &gmmSwiglu_, &tPipe); + op.Init(x, weight, weightScale, xScale, weightAssistanceMatrix, groupList, y, yScale, userWorkspace); + op.Process(); + } +#endif +#if defined(GMM_SWIGLU_QUANT_V2_A4W4) + if (TILING_KEY_IS(4)) { + KERNEL_TASK_TYPE(4, KERNEL_TYPE_MIX_AIC_1_2); + GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, gmmSwigluQuantV2BaseParams, gmmSwigluQuantV2BaseParams_, + tiling); + GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, mmTilingData, mmTilingData_, tiling); + GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, gmmSwigluQuantV2, gmmSwiglu_, tiling); + using xType = MatmulType; + using weightType = MatmulType; + using yType = MatmulType; + using matmulType = MMImplTypeCustom; + matmulType::MT mm; + if ASCEND_IS_AIC { + mm.SetSubBlockIdx(0); + mm.Init(&mmTilingData_); + } + GMMSwigluQuantPipelineSchedule op(mm, &gmmSwigluQuantV2BaseParams_, &gmmSwiglu_, &tPipe); + op.Init(x, weight, weightScale, xScale, weightAssistanceMatrix, groupList, smoothScale, y, yScale, userWorkspace); + op.Process(); + } else if (TILING_KEY_IS(5)) { + KERNEL_TASK_TYPE(5, KERNEL_TYPE_MIX_AIC_1_2); + GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, gmmSwigluQuantV2BaseParams, gmmSwigluQuantV2BaseParams_, + tiling); + + GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, gmmSwigluQuantV2, gmmSwiglu_, tiling); + GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, mmTilingData, mmTilingData_, tiling); + using xType = MatmulType; + using weightType = MatmulType; + using yType = MatmulType; + using matmulType = MMImplTypeCustom; + matmulType::MT mm; + if ASCEND_IS_AIC { + mm.SetSubBlockIdx(0); + mm.Init(&mmTilingData_); + } + GMMSwigluQuantPipelineSchedule op(mm, &gmmSwigluQuantV2BaseParams_, &gmmSwiglu_, &tPipe); + op.Init(x, weight, weightScale, xScale, weightAssistanceMatrix, groupList, smoothScale, y, yScale, userWorkspace); + op.Process(); + } +#endif + if (TILING_KEY_IS(3)) { + KERNEL_TASK_TYPE(3, KERNEL_TYPE_MIX_AIC_1_2); + GET_TILING_DATA_WITH_STRUCT(GMMSwigluQuantV2TilingFusionData, tilingData, tiling); + GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingFusionData, matmulTiling, matmulTilingData, tiling); + GroupedMatmulDequantSwigluQuantFusion op(&tPipe, &tilingData, &matmulTilingData); + if ASCEND_IS_AIC { + op.mm.SetSubBlockIdx(0); + op.mm.Init(&matmulTilingData, &tPipe); + } + + op.Init(x, weight, weightScale, xScale, weightAssistanceMatrix, groupList, y, yScale, userWorkspace); + op.Process(); + } +} diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a4w4_mid.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a4w4_mid.h new file mode 100644 index 000000000..7ea407fc2 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a4w4_mid.h @@ -0,0 +1,255 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_a4w4_mid.h + * \brief + */ + +#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_MID_H +#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_MID_H + +#include "grouped_matmul_swiglu_quant_v2_utils.h" + +#ifdef GMM_SWIGLU_QUANT_V2_A4W4 + +namespace GroupedMatmulDequantSwigluQuant { +using namespace matmul; +using namespace AscendC; + +constexpr uint32_t BUFFER_NUM = 1; + +template +class GMMA4W4MidProcess { +public: + using bT = typename mmType::BT; + +public: + __aicore__ inline GMMA4W4MidProcess(typename mmType::MT &matmul) : mm(matmul) + { + } + __aicore__ inline void Init(const GMAddrParams gmAddrParams, + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN); + __aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx); + +private: + __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig &mnConfig, WorkSpaceSplitConfig &workspaceSplitConfig); + __aicore__ inline void SetMNConfig(const int32_t splitValue, MNConfig &mnConfig); + __aicore__ inline void UpdateMnConfig(MNConfig &mnConfig, bool resetOutputOffset); + +private: + typename mmType::MT &mm; + const uint32_t HALF_ALIGN = 16; + GlobalTensor xGM; + GlobalTensor weightGM; + + GlobalTensor mmOutGM; + GlobalTensor mmOutGM1; + GlobalTensor mmOutGM2; + GlobalTensor groupListGM; + GlobalTensor weightScaleGM; + + GM_ADDR weightTensorPtr; + GM_ADDR weightScaleTensorPtr; + + MNConfig mnConfig; + + // define the que + uint32_t subBlockIdx = 0; + uint32_t coreIdx = 0; + uint32_t quantGroupSize = 0; + uint32_t vecCount = 0; + uint32_t xRowSumCount = 0; + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams = nullptr; +}; + +template +__aicore__ inline void +GMMA4W4MidProcess::Init(const GMAddrParams gmAddrParams, + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN) +{ + if ASCEND_IS_AIC { + gmmSwigluQuantV2BaseParams = gmmSwigluQuantV2BaseParamsIN; + xRowSumCount = gmmSwigluQuantV2BaseParams->M; + xGM.SetGlobalBuffer((__gm__ int4b_t *)gmAddrParams.xGM); + weightGM.SetGlobalBuffer(GetTensorAddr(0, gmAddrParams.weightGM)); + weightScaleGM.SetGlobalBuffer(GetTensorAddr(0, gmAddrParams.weightScaleGM)); + groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM); + mmOutGM1.SetGlobalBuffer((__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM)); + mmOutGM2.SetGlobalBuffer( + (__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset1)); + quantGroupSize = gmmSwigluQuantV2BaseParams->K / gmmSwigluQuantV2BaseParams->quantGroupNum; // 约束为整除关系 + subBlockIdx = GetSubBlockIdx(); + coreIdx = GetBlockIdx(); + weightTensorPtr = gmAddrParams.weightGM; + weightScaleTensorPtr = gmAddrParams.weightScaleGM; + } +} + +template +__aicore__ inline void GMMA4W4MidProcess::UpdateMnConfig(MNConfig &mnConfig, bool resetOutputOffset) +{ + if constexpr (bT::format == CubeFormat::NZ) { + mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<32>(mnConfig.n); // 16: nz format last two dim size + } else { + mnConfig.wBaseOffset += mnConfig.k * mnConfig.n; + } + mnConfig.nAxisBaseOffset += mnConfig.n; + mnConfig.mAxisBaseOffset += mnConfig.m; + mnConfig.xBaseOffset += mnConfig.m * mnConfig.k; + if (resetOutputOffset) { + mnConfig.yBaseOffset = 0; + } else { + mnConfig.yBaseOffset += mnConfig.m * mnConfig.n; + } +} + +template +__aicore__ inline void GMMA4W4MidProcess::SetMNConfig(const int32_t splitValue, MNConfig &mnConfig) +{ + mnConfig.m = static_cast(splitValue); + mnConfig.baseM = gmmSwigluQuantV2BaseParams->baseM; + mnConfig.baseN = gmmSwigluQuantV2BaseParams->baseN; + mnConfig.singleM = gmmSwigluQuantV2BaseParams->baseM; + mnConfig.singleN = gmmSwigluQuantV2BaseParams->singleN != 0 && gmmSwigluQuantV2BaseParams->quantGroupNum == 1? + gmmSwigluQuantV2BaseParams->singleN : gmmSwigluQuantV2BaseParams->baseN; +} + +template +__aicore__ inline void GMMA4W4MidProcess::Process(WorkSpaceSplitConfig &workspaceSplitConfig, + int64_t workspaceSplitLoopIdx) +{ + if ASCEND_IS_AIC { + if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount || workspaceSplitLoopIdx < 0) { + return; + } + mmOutGM = (workspaceSplitLoopIdx % NUM_2 == 0 ? mmOutGM1 : mmOutGM2); + mnConfig.baseM = gmmSwigluQuantV2BaseParams->baseM; + mnConfig.baseN = gmmSwigluQuantV2BaseParams->baseN; + mnConfig.singleM = gmmSwigluQuantV2BaseParams->baseM; + mnConfig.singleN = gmmSwigluQuantV2BaseParams->singleN != 0 && gmmSwigluQuantV2BaseParams->quantGroupNum == 1? + gmmSwigluQuantV2BaseParams->singleN : gmmSwigluQuantV2BaseParams->baseN; + mnConfig.k = gmmSwigluQuantV2BaseParams->K; // tilingData + mnConfig.n = gmmSwigluQuantV2BaseParams->N; // tilingData + mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN); + int32_t prevSplitValue = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize; + int32_t totalTmp = 0; + if (gmmSwigluQuantV2BaseParams->groupListType == 1) { + for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertStartIndex; i++) { + totalTmp += groupListGM.GetValue(i); + } + } + // 当workspace切换时,需要将输出的地址偏移初始化为0,使用resetOutputOffset控制 + bool resetOutputOffset = true; + for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex, preCount = 0; + groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; ++groupIdx) { + UpdateMnConfig(mnConfig, resetOutputOffset); + resetOutputOffset = false; + int32_t currSplitValue = 0; + if (gmmSwigluQuantV2BaseParams->groupListType == 0) { + currSplitValue = static_cast(groupListGM.GetValue(groupIdx)); + } else { + totalTmp += static_cast(groupListGM.GetValue(groupIdx)); + currSplitValue = totalTmp; + } + currSplitValue = currSplitValue > (workspaceSplitLoopIdx + 1) * gmmSwigluQuantV2BaseParams->mLimit ? + (workspaceSplitLoopIdx + 1) * gmmSwigluQuantV2BaseParams->mLimit : + currSplitValue; + + int32_t splitValue = (currSplitValue - prevSplitValue); + prevSplitValue = currSplitValue; + + SetMNConfig(splitValue, mnConfig); + if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) { + continue; + } + mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM); + mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k); + uint32_t curCount = preCount + mnConfig.blockDimN * mnConfig.blockDimM; + uint32_t curBlock = coreIdx >= preCount ? coreIdx : coreIdx + gmmSwigluQuantV2BaseParams->coreNum; + while (curBlock < curCount) { + mnConfig.mIdx = (curBlock - preCount) / mnConfig.blockDimN; + mnConfig.nIdx = (curBlock - preCount) % mnConfig.blockDimN; + MMCompute(groupIdx, mnConfig, workspaceSplitConfig); + curBlock += gmmSwigluQuantV2BaseParams->coreNum; + } + preCount = curCount % gmmSwigluQuantV2BaseParams->coreNum; + } + } +} + +template +__aicore__ inline void GMMA4W4MidProcess::MMCompute(uint32_t groupIdx, MNConfig &mnConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint32_t curSingleN = mnConfig.singleN; + if (unlikely(mnConfig.nIdx == mnConfig.blockDimN - 1)) { + curSingleN = gmmSwigluQuantV2BaseParams->N - tailN; + } + uint32_t curSingleM = mnConfig.singleM; + if (unlikely(mnConfig.mIdx == mnConfig.blockDimM - 1)) { + curSingleM = mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + } + uint64_t weightOffset = 0; + mm.SetSingleShape(curSingleM, curSingleN, quantGroupSize); + GlobalTensor weightSlice; + uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN; + mnConfig.workspaceOffset = outOffset + mnConfig.yBaseOffset; + for (uint32_t loopK = 0; loopK < gmmSwigluQuantV2BaseParams->quantGroupNum; loopK++) { + mm.SetTensorA( + xGM[mnConfig.xBaseOffset + mnConfig.mIdx * mnConfig.k * mnConfig.singleM + loopK * quantGroupSize]); + if (gmmSwigluQuantV2BaseParams->isSingleTensor == 0) { + weightGM.SetGlobalBuffer(GetTensorAddr(groupIdx, weightTensorPtr)); + if constexpr (mmType::BT::format == CubeFormat::NZ && mmType::BT::isTrans == true) { + weightOffset = tailN * 64; + weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmSwigluQuantV2BaseParams->N]; + } else if constexpr (mmType::BT::format == CubeFormat::NZ && mmType::BT::isTrans == false) { + weightOffset = tailN * gmmSwigluQuantV2BaseParams->K; + weightSlice = weightGM[weightOffset + loopK * quantGroupSize * 64]; + } else { + weightOffset = tailN; + weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmSwigluQuantV2BaseParams->N]; + } + } else { + if constexpr (mmType::BT::format == CubeFormat::NZ && mmType::BT::isTrans == true) { + weightOffset = static_cast(groupIdx) * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->K + + tailN * 64; + weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmSwigluQuantV2BaseParams->N]; + } else if constexpr (mmType::BT::format == CubeFormat::NZ && mmType::BT::isTrans == false) { + weightOffset = + static_cast(groupIdx) * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->K + + tailN * gmmSwigluQuantV2BaseParams->K; + weightSlice = weightGM[weightOffset + loopK * quantGroupSize * 64]; + } else { + weightOffset = + static_cast(groupIdx) * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->K + + tailN; + weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmSwigluQuantV2BaseParams->N]; + } + } + if (mnConfig.blockDimM == 1) { + weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } + mm.SetTensorB(weightSlice, mmType::BT::isTrans); + if (gmmSwigluQuantV2BaseParams->isSingleTensor == 0) { + weightScaleGM.SetGlobalBuffer(GetTensorAddr(groupIdx, weightScaleTensorPtr)); + mm.SetQuantVector(weightScaleGM[loopK * gmmSwigluQuantV2BaseParams->N + tailN]); + } else { + mm.SetQuantVector( + weightScaleGM[groupIdx * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->quantGroupNum + + loopK * gmmSwigluQuantV2BaseParams->N + tailN]); + } + mm.IterateAll(mmOutGM[mnConfig.workspaceOffset], loopK == 0 ? 0 : 1); + } +} +} // namespace GroupedMatmulDequantSwigluQuant +#endif // GMM_SWIGLU_QUANT_V2_A4W4 +#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_MID_H \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a4w4_pipeline.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a4w4_pipeline.h new file mode 100644 index 000000000..031cba8df --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a4w4_pipeline.h @@ -0,0 +1,229 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_a4w4_pipeline.h + * \brief + */ +#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_PIPELINE_H +#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_PIPELINE_H + +#include +#include "grouped_matmul_swiglu_quant_v2_a4w4_mid.h" +#include "grouped_matmul_swiglu_quant_v2_a4w4_post.h" +#include "grouped_matmul_swiglu_quant_v2_utils.h" + +using namespace AscendC; +using namespace matmul; + +#ifdef GMM_SWIGLU_QUANT_V2_A4W4 + +namespace GroupedMatmulDequantSwigluQuant { + +template +class GMMSwigluQuantPipelineSchedule { +private: + typename mmType::MT &mm; + TPipe *pipe; + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams; + const GMMSwigluQuantV2 *__restrict gmmSwigluQuantV2; + // WorkSpaceSplitConfig控制Workspace切割方式的结构体; + WorkSpaceSplitConfig workspaceSplitConfig; + WorkSpaceSplitConfig tempWorkspaceSplitConfig; + // 记录GM_ADDR的结构体 + GMAddrParams gmAddrParams; + // 中间处理GMMA4W4MidProcess类 + GMMA4W4MidProcess midProcess; + // 后处理GMMA4W4PostProcess类 + GMMA4W4PostProcess postProcess; + GlobalTensor groupListGM; + __aicore__ inline void InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig, + int32_t workspaceSplitLoopIdx); + +public: + __aicore__ inline GMMSwigluQuantPipelineSchedule( + typename mmType::MT &mm_, const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN, + const GMMSwigluQuantV2 *__restrict gmmSwigluIN, TPipe *tPipeIN) + : mm(mm_), midProcess(mm), gmmSwigluQuantV2BaseParams(gmmSwigluQuantV2BaseParamsIN), + gmmSwigluQuantV2(gmmSwigluIN), pipe(tPipeIN) + { + } + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, GM_ADDR xScale, + GM_ADDR weightAssistanceMatrix, GM_ADDR groupList, GM_ADDR smoothScale, GM_ADDR y, GM_ADDR yScale, + GM_ADDR workspace); + __aicore__ inline void Process(); +}; + +template +__aicore__ inline void GMMSwigluQuantPipelineSchedule::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, + GM_ADDR xScale, GM_ADDR weightAssistanceMatrix, + GM_ADDR groupList, GM_ADDR smoothScale, GM_ADDR y, GM_ADDR yScale, + GM_ADDR workspace) +{ + gmAddrParams.xGM = x; + gmAddrParams.weightGM = weight; + gmAddrParams.weightScaleGM = weightScale; + gmAddrParams.xScaleGM = xScale; + gmAddrParams.weightAuxiliaryMatrixGM = weightAssistanceMatrix; + gmAddrParams.groupListGM = groupList; + gmAddrParams.smoothScaleGM = smoothScale; + gmAddrParams.yGM = y; + gmAddrParams.yScaleGM = yScale; + gmAddrParams.workSpaceGM = workspace; + gmAddrParams.workSpaceOffset1 = gmmSwigluQuantV2BaseParams->workSpaceOffset1; + gmAddrParams.workSpaceOffset2 = 0; + gmAddrParams.workSpaceOffset3 = 0; + groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM); + InitWorkSpaceSplitConfig(workspaceSplitConfig); +} + +template +__aicore__ inline void GMMSwigluQuantPipelineSchedule::Process() +{ + // 1.对每次workspace切分做大循环。 + midProcess.Init(gmAddrParams, gmmSwigluQuantV2BaseParams); + postProcess.Init(gmAddrParams, gmmSwigluQuantV2BaseParams, gmmSwigluQuantV2); + + for (int64_t workspaceSplitLoopIdx = 0; workspaceSplitLoopIdx < workspaceSplitConfig.loopCount; + workspaceSplitLoopIdx++) { + // 更新workspaceSplitConfig + UpdateWorkSpaceSplitConfig(workspaceSplitConfig, workspaceSplitLoopIdx); + if ASCEND_IS_AIV { + pipe->Reset(); + } + + SyncAll(); + // 2.第n次中处理 && 第n-1次后处理 并行 + midProcess.Process(workspaceSplitConfig, workspaceSplitLoopIdx); + + if ASCEND_IS_AIV { + pipe->Reset(); + SyncAll(); + } + postProcess.Process(tempWorkspaceSplitConfig, workspaceSplitLoopIdx - 1, pipe); + // 3.第n-1次后处理需要保留第n次的切分数据 + tempWorkspaceSplitConfig = workspaceSplitConfig; + // reset + if ASCEND_IS_AIV { + pipe->Reset(); + } + SyncAll(); + // 3.前一次后处理 && 后一次MM 并行 + } + // reset + if ASCEND_IS_AIV { + pipe->Reset(); + } + SyncAll(); + // // 4.最后一次后处理 + postProcess.Process(workspaceSplitConfig, workspaceSplitConfig.loopCount - 1, pipe); + if ASCEND_IS_AIV { + pipe->Destroy(); + } +} + +template +__aicore__ inline void +GMMSwigluQuantPipelineSchedule::InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig) +{ + if (gmmSwigluQuantV2BaseParams->groupListType == 0) { + workspaceSplitConfig.M = groupListGM.GetValue(gmmSwigluQuantV2->groupListLen - 1); + } else { + int64_t totalTmp = 0; + for (uint32_t i = 0; i < gmmSwigluQuantV2->groupListLen; i++) { + totalTmp += groupListGM.GetValue(i); + } + workspaceSplitConfig.M = totalTmp; + } + workspaceSplitConfig.loopCount = Ceil(workspaceSplitConfig.M, gmmSwigluQuantV2BaseParams->mLimit); + workspaceSplitConfig.notLastTaskSize = gmmSwigluQuantV2BaseParams->mLimit; + workspaceSplitConfig.lastLoopTaskSize = + workspaceSplitConfig.M - (workspaceSplitConfig.loopCount - 1) * gmmSwigluQuantV2BaseParams->mLimit; + workspaceSplitConfig.leftMatrixStartIndex = 0; + workspaceSplitConfig.rightMatrixExpertStartIndex = 0; + workspaceSplitConfig.rightMatrixExpertNextStartIndex = 0; + workspaceSplitConfig.isLastLoop = false; +} + +template +__aicore__ inline void +GMMSwigluQuantPipelineSchedule::UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig, + int32_t workspaceSplitLoopIdx) +{ + if (workspaceSplitLoopIdx < 0) + return; + workspaceSplitConfig.leftMatrixStartIndex = workspaceSplitLoopIdx * gmmSwigluQuantV2BaseParams->mLimit; + workspaceSplitConfig.rightMatrixExpertStartIndex = workspaceSplitConfig.rightMatrixExpertNextStartIndex; + workspaceSplitConfig.rightMatrixExpertEndIndex = workspaceSplitConfig.rightMatrixExpertStartIndex; + // 计算右专家矩阵的终止索引(rightMatrixExpertEndIndex) 和下一次的起始索引(rightMatrixExpertNextStartIndex) + int32_t curTaskNum = 0; + int32_t nextTaskNum = 0; + int32_t curTaskNumTmp = 0; + int32_t nextTaskNumTmp = 0; + if (gmmSwigluQuantV2BaseParams->groupListType == 1) { + for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertEndIndex; i++) { + curTaskNumTmp += groupListGM.GetValue(i); + } + if (workspaceSplitConfig.rightMatrixExpertEndIndex == 0) { + nextTaskNumTmp = groupListGM.GetValue(0); + } else { + for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertEndIndex; i++) { + nextTaskNumTmp += groupListGM.GetValue(i); + } + } + } + while (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwigluQuantV2->groupListLen) { + if (gmmSwigluQuantV2BaseParams->groupListType == 0) { + curTaskNum = groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex) - + workspaceSplitConfig.leftMatrixStartIndex; + } else { + curTaskNumTmp += groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex); + curTaskNum = curTaskNumTmp - workspaceSplitConfig.leftMatrixStartIndex; + } + int32_t nextTaskIdx = workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwigluQuantV2->groupListLen - 1 ? + gmmSwigluQuantV2->groupListLen - 1 : + workspaceSplitConfig.rightMatrixExpertEndIndex + 1; + if (gmmSwigluQuantV2BaseParams->groupListType == 0) { + nextTaskNum = groupListGM.GetValue(nextTaskIdx) - workspaceSplitConfig.leftMatrixStartIndex; + } else { + if (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwigluQuantV2->groupListLen - 1) { + nextTaskNumTmp += groupListGM.GetValue(nextTaskIdx); + } + nextTaskNum = nextTaskNumTmp - workspaceSplitConfig.leftMatrixStartIndex; + } + if (curTaskNum > gmmSwigluQuantV2BaseParams->mLimit) { + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex; + break; + } else if (curTaskNum == gmmSwigluQuantV2BaseParams->mLimit && + nextTaskNum > gmmSwigluQuantV2BaseParams->mLimit) { + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex + 1; + break; + } else if (nextTaskNum > gmmSwigluQuantV2BaseParams->mLimit) { + workspaceSplitConfig.rightMatrixExpertEndIndex++; + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex; + break; + } + workspaceSplitConfig.rightMatrixExpertEndIndex++; + } + workspaceSplitConfig.isLastLoop = workspaceSplitLoopIdx == workspaceSplitConfig.loopCount - 1 ? true : false; + + if (workspaceSplitConfig.isLastLoop) { + workspaceSplitConfig.rightMatrixExpertEndIndex = + workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwigluQuantV2->groupListLen ? + gmmSwigluQuantV2->groupListLen - 1 : + workspaceSplitConfig.rightMatrixExpertEndIndex; + } +} + +} // namespace GroupedMatmulDequantSwigluQuant +#endif // GMM_SWIGLU_QUANT_V2_A4W4 +#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_PIPELINE_H \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a4w4_post.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a4w4_post.h new file mode 100644 index 000000000..06815a28c --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a4w4_post.h @@ -0,0 +1,392 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_a4w4_post.h + * \brief + */ + +#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_POST_H +#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_POST_H + +#include "grouped_matmul_swiglu_quant_v2_utils.h" +#include "kernel_operator.h" + +#ifdef GMM_SWIGLU_QUANT_V2_A4W4 + +namespace GroupedMatmulDequantSwigluQuant { +using namespace AscendC; +#define DOUBLE_BUFFER 2 +constexpr float DEFAULT_MUL_SCALE = 16.0f; +class GMMA4W4PostProcess { +public: + __aicore__ inline GMMA4W4PostProcess(){}; + __aicore__ inline void Init(const GMAddrParams gmAddrParams, + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN, + const GMMSwigluQuantV2 *__restrict gmmSwigluIN); + + __aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx, + TPipe *pipe); + static constexpr float FLOAT_INF = 3e+99; + +private: + __aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx, + TPipe *pipe); + + __aicore__ inline void VectorCompute(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void customDataCopyIn(uint32_t outLoopIdx, GlobalTensor &mmOutGM, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void customDataCopyOut(VecConfig &vecConfig, WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void Quant(uint32_t loopIdx, VecConfig &vecConfig); + + __aicore__ inline void Swiglu(uint32_t loopIdx, VecConfig &vecConfig); + + __aicore__ inline void MulPertokenScale(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void ApplySmoothScale(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig); + + const GMMSwigluQuantV2 *__restrict gmmSwigluQuantV2; + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams; + GlobalTensor perTokenScaleGM; + GlobalTensor groupListGM; + GlobalTensor smoothScaleGM; + GlobalTensor quantOutputGM; + GlobalTensor quantScaleOutputGM; + GlobalTensor mmOutGM1; + GlobalTensor mmOutGM2; + GlobalTensor mmOutGM; + LocalTensor mmLocal_fp32; + LocalTensor mmLocal_fp16; + TQue mmOutQueue; + TQue quantOutQueue; + TQue quantScaleOutQueue; + TBuf reduceWorkspace; + uint32_t blockIdx = 0; + int64_t aicCoreNum = 0; + int64_t aivCoreNum = 0; + float limited = FLOAT_INF; +}; + +__aicore__ inline void GMMA4W4PostProcess::Init(const GMAddrParams gmAddrParams, + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN, + const GMMSwigluQuantV2 *__restrict gmmSwigluIN) +{ + if ASCEND_IS_AIV { + aicCoreNum = GetBlockNum(); + aivCoreNum = aicCoreNum * NUM_2; + blockIdx = GetBlockIdx(); + gmmSwigluQuantV2BaseParams = gmmSwigluQuantV2BaseParamsIN; + gmmSwigluQuantV2 = gmmSwigluIN; + groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM, gmmSwigluQuantV2->groupListLen); + mmOutGM1.SetGlobalBuffer((__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM)); + mmOutGM2.SetGlobalBuffer( + (__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset1)); + perTokenScaleGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.xScaleGM, gmmSwigluQuantV2BaseParams->M); + smoothScaleGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.smoothScaleGM); + quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.yGM, gmmSwigluQuantV2BaseParams->M * + gmmSwigluQuantV2->tokenLen / + SWIGLU_REDUCE_FACTOR); + quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.yScaleGM, gmmSwigluQuantV2BaseParams->M); + limited = gmmSwigluQuantV2BaseParams->swigluLimit; + } +} + +__aicore__ inline void GMMA4W4PostProcess::customDataCopyIn(uint32_t outLoopIdx, GlobalTensor &mmOutGM, + VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + mmLocal_fp16 = mmOutQueue.DeQue(); + mmLocal_fp32 = mmLocal_fp16.ReinterpretCast(); + const int64_t processNum = vecConfig.innerLoopNum * gmmSwigluQuantV2->tokenLen; + DataCopyExtParams copyParams_0{1, static_cast(processNum * SIZE_OF_HALF_2), 0, 0, 0}; + DataCopyPadExtParams padParams_0{false, 0, 0, 0}; + DataCopyPad(mmLocal_fp16[processNum], mmOutGM[vecConfig.curOffset], copyParams_0, padParams_0); + + mmOutQueue.EnQue(mmLocal_fp16); + mmLocal_fp16 = mmOutQueue.DeQue(); + + // 1. fp16 -> fp32 + Cast(mmLocal_fp32, mmLocal_fp16[processNum], RoundMode::CAST_NONE, processNum); + PipeBarrier(); + + vecConfig.curIdx += vecConfig.innerLoopNum; + vecConfig.curOffset = vecConfig.curIdx * gmmSwigluQuantV2->tokenLen; +} + +__aicore__ inline void GMMA4W4PostProcess::VectorCompute(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + // 1.perToken反量化 + MulPertokenScale(loopIdx, vecConfig, workspaceSplitConfig); + // 2.Swiglu + Swiglu(loopIdx, vecConfig); + // 3.ApplySmoothScale（smoothScaleDimNum为0时跳过，表示smoothScale为空指针） + if (gmmSwigluQuantV2BaseParams->smoothScaleDimNum != 0) { + ApplySmoothScale(loopIdx, vecConfig, workspaceSplitConfig); + } + // 4.Quant + Quant(loopIdx, vecConfig); +} + +__aicore__ inline void GMMA4W4PostProcess::MulPertokenScale(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + if (loopIdx != 0) { + mmLocal_fp32 = mmOutQueue.DeQue(); + } + float scale = perTokenScaleGM.GetValue(loopIdx + workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx); + PipeBarrier(); + Muls(mmLocal_fp32[loopIdx * gmmSwigluQuantV2->tokenLen], mmLocal_fp32[loopIdx * gmmSwigluQuantV2->tokenLen], scale, + gmmSwigluQuantV2->tokenLen); +} + +__aicore__ inline void GMMA4W4PostProcess::Swiglu(uint32_t loopIdx, VecConfig &vecConfig) +{ + // 高阶API swiglu + float beta = 1.0f; + LocalTensor workspaceLocal = reduceWorkspace.Get(); + LocalTensor src0Local = + mmLocal_fp32[loopIdx * gmmSwigluQuantV2->tokenLen + gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR]; + LocalTensor src1Local = mmLocal_fp32[loopIdx * gmmSwigluQuantV2->tokenLen]; + if (limited > 0.0f) { + Mins(src0Local, src0Local, limited, gmmSwigluQuantV2->tokenLen / 2); + PipeBarrier(); + Maxs(src0Local, src0Local, (-1.0f * limited), gmmSwigluQuantV2->tokenLen / 2); + PipeBarrier(); + Mins(src1Local, src1Local, limited, gmmSwigluQuantV2->tokenLen / 2); + PipeBarrier(); + } + SwiGLU(workspaceLocal, src0Local, src1Local, beta, gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR); + PipeBarrier(); + DataCopyParams repeatParams{ + 1, static_cast((gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR) / ALIGN_8_ELE), 0, 0}; + DataCopy(mmLocal_fp32[loopIdx * gmmSwigluQuantV2->tokenLen], workspaceLocal, repeatParams); +} + +__aicore__ inline void GMMA4W4PostProcess::ApplySmoothScale(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + int64_t smoothScaleDimNum = gmmSwigluQuantV2BaseParams->smoothScaleDimNum; + int64_t halfTokenLen = gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR; + int64_t currentTokenIdx = workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx + loopIdx; + + // 找到当前token所属的group + uint32_t groupIdx = 0; + int64_t prevM = 0; + int64_t totalTmp = 0; + if (gmmSwigluQuantV2BaseParams->groupListType == 1) { + for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertStartIndex; i++) { + totalTmp += groupListGM.GetValue(i); + } + } + for (uint32_t i = workspaceSplitConfig.rightMatrixExpertStartIndex; + i <= workspaceSplitConfig.rightMatrixExpertEndIndex; i++) { + int64_t currM = 0; + if (gmmSwigluQuantV2BaseParams->groupListType == 0) { + currM = groupListGM.GetValue(i); + } else { + totalTmp += groupListGM.GetValue(i); + currM = totalTmp; + } + if (currentTokenIdx < currM) { + groupIdx = i; + break; + } + prevM = currM; + } + + uint64_t preOffset = loopIdx * gmmSwigluQuantV2->tokenLen; + + if (smoothScaleDimNum == NUM_2) { + // smoothScale形状为 (E, N/2)，只需要当前group的那一行 + for (uint32_t j = 0; j < halfTokenLen; j++) { + float scale = smoothScaleGM.GetValue(groupIdx * halfTokenLen + j); + float val = mmLocal_fp32.GetValue(preOffset + j); + mmLocal_fp32.SetValue(preOffset + j, val * scale); + } + } else if (smoothScaleDimNum == 1) { + // smoothScale形状为 (E,)，需要广播到 (N/2) + float scale = smoothScaleGM.GetValue(groupIdx); + PipeBarrier(); + Muls(mmLocal_fp32[preOffset], mmLocal_fp32[preOffset], scale, halfTokenLen); + } +} + +__aicore__ inline void GMMA4W4PostProcess::Quant(uint32_t loopIdx, VecConfig &vecConfig) +{ + uint64_t preOffset = loopIdx * gmmSwigluQuantV2->tokenLen; + uint64_t halfTokenLen = gmmSwigluQuantV2->tokenLen / BISECT; + PipeBarrier(); + Abs(mmLocal_fp32[preOffset + gmmSwigluQuantV2->tokenLen / BISECT], mmLocal_fp32[preOffset], halfTokenLen); + PipeBarrier(); + // reduceMax + LocalTensor workLocal = reduceWorkspace.Get(halfTokenLen); + LocalTensor reduceResLocal = + reduceWorkspace.GetWithOffset(FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float)); + LocalTensor reduceTmpLocal = reduceWorkspace.GetWithOffset( + FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float) + UB_BLOCK_UNIT_SIZE); + ReduceMaxTemplate(reduceResLocal, workLocal, mmLocal_fp32[preOffset + gmmSwigluQuantV2->tokenLen / BISECT], + reduceTmpLocal, static_cast(halfTokenLen)); + float quantScale = reduceResLocal.GetValue(0) / QUANT_SCALE_INT8; + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + quantScaleLocal.SetValue(loopIdx, quantScale); + quantScale = QUANT_SCALE_INT8 / reduceResLocal.GetValue(0); + Muls(mmLocal_fp32[preOffset], mmLocal_fp32[preOffset], quantScale, halfTokenLen); + PipeBarrier(); + LocalTensor quantLocal = quantOutQueue.DeQue(); + int32_t dstTempOffset = static_cast(preOffset / BISECT); + int32_t srcTempOffset = static_cast(preOffset); + int32_t tempCount = static_cast(halfTokenLen); + LocalTensor castSpace = reduceWorkspace.Get(UB_BLOCK_UNIT_SIZE); + CastFp32ToInt8Template(quantLocal, mmLocal_fp32, castSpace, dstTempOffset, srcTempOffset, tempCount); + mmOutQueue.EnQue(mmLocal_fp32); + quantOutQueue.EnQue(quantLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); +} + +__aicore__ inline void GMMA4W4PostProcess::UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig, + int64_t workspaceSplitLoopIdx, TPipe *pipe) +{ + // 第一步 读取grouplist reduceSum 计算总数据个数 + vecConfig.M = workspaceSplitLoopIdx < workspaceSplitConfig.loopCount - 1 ? workspaceSplitConfig.notLastTaskSize : + workspaceSplitConfig.lastLoopTaskSize; + // 第二步 计算分核 + uint32_t eachCoreTaskNum = (vecConfig.M + aivCoreNum - 1) / aivCoreNum; + vecConfig.usedCoreNum = vecConfig.M >= aivCoreNum ? aivCoreNum : vecConfig.M; + uint32_t tailCoreIdx = vecConfig.M - (eachCoreTaskNum - 1) * vecConfig.usedCoreNum; + vecConfig.taskNum = blockIdx < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1; + vecConfig.startIdx = + blockIdx < tailCoreIdx ? eachCoreTaskNum * blockIdx : ((eachCoreTaskNum - 1) * blockIdx + tailCoreIdx); + vecConfig.curIdx = vecConfig.startIdx; + vecConfig.startOffset = vecConfig.startIdx * gmmSwigluQuantV2->tokenLen; + vecConfig.curOffset = vecConfig.startOffset; + int64_t curStartIdx = vecConfig.startIdx; + int64_t prevM = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize; + int64_t totalTmp = 0; + if (gmmSwigluQuantV2BaseParams->groupListType == 1) { + for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertStartIndex; i++) { + totalTmp += groupListGM.GetValue(i); + } + } + for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex; + groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; groupIdx++) { + int64_t currM = 0; + if (gmmSwigluQuantV2BaseParams->groupListType == 0) { + currM = groupListGM.GetValue(groupIdx); + } else { + totalTmp += groupListGM.GetValue(groupIdx); + currM = totalTmp; + } + int64_t tempM = currM - prevM; + prevM = currM; + curStartIdx -= tempM; + } + // 第三步 计算总数据量 + vecConfig.outLoopNum = + (vecConfig.taskNum + gmmSwigluQuantV2->maxProcessRowNum - 1) / gmmSwigluQuantV2->maxProcessRowNum; + vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwigluQuantV2->maxProcessRowNum ? + vecConfig.taskNum % gmmSwigluQuantV2->maxProcessRowNum : + gmmSwigluQuantV2->maxProcessRowNum; + + // 第四步 申请空间 + // 2 * row * n * sizeof(float) + row * n / 2 * sizeof(int8) + alignUp * sizeof(float) + n * sizeof(float) + + // n / 2 *sizeof(float) + 64 < 191 * 1024 + pipe->InitBuffer(mmOutQueue, 1, + gmmSwigluQuantV2->maxProcessRowNum * gmmSwigluQuantV2->tokenLen * sizeof(float)); + pipe->InitBuffer(quantOutQueue, 1, + gmmSwigluQuantV2->maxProcessRowNum * gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR * + sizeof(int8_t)); + pipe->InitBuffer(quantScaleOutQueue, 1, + AlignUp(gmmSwigluQuantV2->maxProcessRowNum, ALIGN_8_ELE) * sizeof(float)); + // two 32 byte buffer for reduceMax calculation in Quant. + pipe->InitBuffer(reduceWorkspace, gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(float) + + UB_BLOCK_UNIT_SIZE + UB_BLOCK_UNIT_SIZE); +} + +__aicore__ inline void GMMA4W4PostProcess::Process(WorkSpaceSplitConfig &workspaceSplitConfig, + int64_t workspaceSplitLoopIdx, TPipe *pipe) +{ + if ASCEND_IS_AIV { + if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount || workspaceSplitLoopIdx < 0) { + return; + } + VecConfig vecConfig; + UpdateVecConfig(blockIdx, vecConfig, workspaceSplitConfig, workspaceSplitLoopIdx, pipe); + + if (blockIdx < vecConfig.usedCoreNum) { + mmOutGM = (workspaceSplitLoopIdx % NUM_2 == 0 ? mmOutGM1 : mmOutGM2); + LocalTensor mmLocal = mmOutQueue.AllocTensor(); + LocalTensor quantScaleLocal = quantScaleOutQueue.AllocTensor(); + LocalTensor quantLocal = quantOutQueue.AllocTensor(); + + mmOutQueue.EnQue(mmLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); + quantOutQueue.EnQue(quantLocal); + for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) { + vecConfig.innerLoopNum = outLoopIdx == (vecConfig.outLoopNum - 1) ? vecConfig.tailLoopNum : + gmmSwigluQuantV2->maxProcessRowNum; + int32_t eventIdMTE3ToMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventIdMTE3ToMTE2); + WaitFlag(eventIdMTE3ToMTE2); + // 1.matmul中间结果搬入 + customDataCopyIn(outLoopIdx, mmOutGM, vecConfig, workspaceSplitConfig); + + for (uint32_t innerLoopIdx = 0; innerLoopIdx < vecConfig.innerLoopNum; innerLoopIdx++) { + // 2. 四步vector计算（perToken反量化、Swiglu、SmoothScale、Quant） + VectorCompute(innerLoopIdx, vecConfig, workspaceSplitConfig); + } + int32_t eventIdVToMTE3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventIdVToMTE3); + WaitFlag(eventIdVToMTE3); + customDataCopyOut(vecConfig, workspaceSplitConfig); + } + mmLocal = mmOutQueue.DeQue(); + quantScaleLocal = quantScaleOutQueue.DeQue(); + quantLocal = quantOutQueue.DeQue(); + + mmOutQueue.FreeTensor(mmLocal); + quantScaleOutQueue.FreeTensor(quantScaleLocal); + quantOutQueue.FreeTensor(quantLocal); + } + } +} + +__aicore__ inline void GMMA4W4PostProcess::customDataCopyOut(VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + DataCopyParams copyParams_0{1, (uint16_t)(vecConfig.innerLoopNum * sizeof(float)), 0, 0}; + DataCopyPad(quantScaleOutputGM[workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx], quantScaleLocal, + copyParams_0); + LocalTensor quantLocal = quantOutQueue.DeQue(); + DataCopyParams copyParams_1{ + 1, (uint16_t)(vecConfig.innerLoopNum * gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)), 0, + 0}; + DataCopyPad(quantOutputGM[(workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx) * + gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR], + quantLocal, copyParams_1); + + vecConfig.startIdx += vecConfig.innerLoopNum; + vecConfig.startOffset = vecConfig.startIdx * gmmSwigluQuantV2->tokenLen; + quantOutQueue.EnQue(quantLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); +} + +} // namespace GroupedMatmulDequantSwigluQuant +#endif // GMM_SWIGLU_QUANT_V2_A4W4 +#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_POST_H diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a8w4_msd_mid.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a8w4_msd_mid.h new file mode 100644 index 000000000..27ba3780d --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a8w4_msd_mid.h @@ -0,0 +1,288 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_a8w4_msd_mid.h + * \brief + */ + +#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_MID_H +#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_MID_H + +#include "grouped_matmul_swiglu_quant_v2_utils.h" + +#ifdef GMM_SWIGLU_QUANT_V2_A8W4_MSD + +namespace GroupedMatmulDequantSwigluQuant { +using namespace matmul; +using namespace AscendC; + +constexpr uint32_t BUFFER_NUM = 1; + +template +__aicore__ inline void DataCopyPad2DA8W4(const LocalTensor dst, const GlobalTensor src, uint32_t dim1, + uint32_t dim0, uint32_t srcDim0) +{ + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + params.srcStride = (srcDim0 - dim0) * sizeof(T); + // 32: int32 -> float16, 为防止跨行数据进入同一32B block，提前每行按偶数block对齐 + params.dstStride = Ceil(dim0 * sizeof(T), 32) % 2; + + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(dst, src, params, padParams); +} + +template +__aicore__ inline void DataCopyPad2DA8W4ND(const LocalTensor dst, const GlobalTensor src, uint32_t dim1, + uint32_t dim0, uint32_t srcDim0) +{ + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + params.srcStride = (srcDim0 - dim0) * sizeof(T); + params.dstStride = 0; + + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(dst, src, params, padParams); + return; +} + +template +__aicore__ inline void DataCopyPad2DA8W4(const GlobalTensor dst, const LocalTensor src, uint32_t dim1, + uint32_t dim0, uint32_t srcDim0, uint32_t dstDim0) +{ + DataCopyExtParams params; + params.blockCount = dim1; + params.blockLen = dim0 * sizeof(T); + // 32: ub访问粒度为32B + params.srcStride = (srcDim0 - dim0) * sizeof(T) / 32; + params.dstStride = (dstDim0 - dim0) * sizeof(T); + DataCopyPad(dst, src, params); +} + +template +class GMMA8W4MidProcess { +public: + using bT = typename mmType::BT; + +public: + __aicore__ inline GMMA8W4MidProcess(typename mmType::MT &matmul) : mm(matmul) + { + } + __aicore__ inline void Init(const GMAddrParams gmAddrParams, + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN); + __aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx); + +private: + __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig &mnConfig, WorkSpaceSplitConfig &workspaceSplitConfig); + __aicore__ inline void SetMNConfig(const int32_t splitValue, MNConfig &mnConfig); + __aicore__ inline void UpdateMnConfig(MNConfig &mnConfig); + +private: + typename mmType::MT &mm; + const uint32_t HALF_ALIGN = 16; + GlobalTensor xGM; + GlobalTensor xGM1; + GlobalTensor xGM2; + GlobalTensor weightGM; + + GlobalTensor mmOutGM; + GlobalTensor mmOutGM1; + GlobalTensor mmOutGM2; + GlobalTensor groupListGM; + GlobalTensor weightScaleGM; + + GM_ADDR weightTensorPtr; + GM_ADDR weightScaleTensorPtr; + + // define the que + uint32_t subBlockIdx = 0; + uint32_t coreIdx = 0; + uint32_t quantGroupSize = 0; + uint32_t vecCount = 0; + uint32_t xRowSumCount = 0; + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams; +}; + +template +__aicore__ inline void +GMMA8W4MidProcess::Init(const GMAddrParams gmAddrParams, + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN) +{ + if ASCEND_IS_AIC { + gmmSwigluQuantV2BaseParams = gmmSwigluQuantV2BaseParamsIN; + xRowSumCount = gmmSwigluQuantV2BaseParams->M; + xGM1.SetGlobalBuffer((__gm__ int4b_t *)gmAddrParams.workSpaceGM); // 从前处理中获得的结果 + xGM2.SetGlobalBuffer( + (__gm__ int4b_t *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset1)); + weightGM.SetGlobalBuffer(GetTensorAddr(0, gmAddrParams.weightGM)); + weightScaleGM.SetGlobalBuffer(GetTensorAddr(0, gmAddrParams.weightScaleGM)); + groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM); + mmOutGM1.SetGlobalBuffer( + (__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset2)); + mmOutGM2.SetGlobalBuffer( + (__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset3)); + quantGroupSize = gmmSwigluQuantV2BaseParams->K / gmmSwigluQuantV2BaseParams->quantGroupNum; // 约束为整除关系 + subBlockIdx = GetSubBlockIdx(); + coreIdx = GetBlockIdx(); + weightTensorPtr = gmAddrParams.weightGM; + weightScaleTensorPtr = gmAddrParams.weightScaleGM; + } +} + +template +__aicore__ inline void GMMA8W4MidProcess::UpdateMnConfig(MNConfig &mnConfig) +{ + if constexpr (bT::format == CubeFormat::NZ) { + mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<32>(mnConfig.n); // 16: nz format last two dim size + } else { + mnConfig.wBaseOffset += mnConfig.k * mnConfig.n; + } + mnConfig.nAxisBaseOffset += mnConfig.n; + mnConfig.mAxisBaseOffset += mnConfig.m; + mnConfig.xBaseOffset += mnConfig.m * mnConfig.k; + mnConfig.yBaseOffset += mnConfig.m * mnConfig.n; +} + +template +__aicore__ inline void GMMA8W4MidProcess::SetMNConfig(const int32_t splitValue, MNConfig &mnConfig) +{ + mnConfig.m = static_cast(splitValue); + mnConfig.baseM = gmmSwigluQuantV2BaseParams->baseM; + mnConfig.baseN = gmmSwigluQuantV2BaseParams->baseN; + mnConfig.singleM = gmmSwigluQuantV2BaseParams->baseM; + mnConfig.singleN = gmmSwigluQuantV2BaseParams->baseN; +} + +template +__aicore__ inline void GMMA8W4MidProcess::Process(WorkSpaceSplitConfig &workspaceSplitConfig, + int64_t workspaceSplitLoopIdx) +{ + if ASCEND_IS_AIC { + if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount || workspaceSplitLoopIdx < 0) { + return; + } + xGM = (workspaceSplitLoopIdx % 2 == 0 ? xGM1 : xGM2); + mmOutGM = (workspaceSplitLoopIdx % 2 == 0 ? mmOutGM1 : mmOutGM2); + MNConfig mnConfig; + mnConfig.baseM = gmmSwigluQuantV2BaseParams->baseM; + mnConfig.baseN = gmmSwigluQuantV2BaseParams->baseN; + mnConfig.singleM = gmmSwigluQuantV2BaseParams->baseM; + mnConfig.singleN = gmmSwigluQuantV2BaseParams->baseN; + mnConfig.k = gmmSwigluQuantV2BaseParams->K; // tilingData + mnConfig.n = gmmSwigluQuantV2BaseParams->N; // tilingData + mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN); + int32_t prevSplitValue = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize; + int32_t totalTmp = 0; + if (gmmSwigluQuantV2BaseParams->groupListType == 1) { + for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertStartIndex; i++) { + totalTmp += groupListGM.GetValue(i); + } + } + for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex, preCount = 0; + groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; ++groupIdx) { + UpdateMnConfig(mnConfig); + int32_t currSplitValue = 0; + if (gmmSwigluQuantV2BaseParams->groupListType == 0) { + currSplitValue = static_cast(groupListGM.GetValue(groupIdx)); + } else { + totalTmp += static_cast(groupListGM.GetValue(groupIdx)); + currSplitValue = totalTmp; + } + currSplitValue = currSplitValue > (workspaceSplitLoopIdx + 1) * gmmSwigluQuantV2BaseParams->mLimit ? + (workspaceSplitLoopIdx + 1) * gmmSwigluQuantV2BaseParams->mLimit : + currSplitValue; + + int32_t splitValue = (currSplitValue - prevSplitValue) * 2; // 2: int8 has been split in 2 int4 + prevSplitValue = currSplitValue; + + SetMNConfig(splitValue, mnConfig); + if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) { + continue; + } + mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM); + mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k); + uint32_t curCount = preCount + mnConfig.blockDimN * mnConfig.blockDimM; + uint32_t curBlock = coreIdx >= preCount ? coreIdx : coreIdx + gmmSwigluQuantV2BaseParams->coreNum; + while (curBlock < curCount) { + mnConfig.mIdx = (curBlock - preCount) / mnConfig.blockDimN; + mnConfig.nIdx = (curBlock - preCount) % mnConfig.blockDimN; + MMCompute(groupIdx, mnConfig, workspaceSplitConfig); + curBlock += gmmSwigluQuantV2BaseParams->coreNum; + } + preCount = curCount % gmmSwigluQuantV2BaseParams->coreNum; + } + } +} + +template +__aicore__ inline void GMMA8W4MidProcess::MMCompute(uint32_t groupIdx, MNConfig &mnConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint32_t curSingleN = mnConfig.singleN; + if (unlikely(mnConfig.nIdx == mnConfig.blockDimN - 1)) { + curSingleN = gmmSwigluQuantV2BaseParams->N - tailN; + } + uint32_t curSingleM = mnConfig.singleM; + if (unlikely(mnConfig.mIdx == mnConfig.blockDimM - 1)) { + curSingleM = mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + } + uint64_t weightOffset = 0; + mm.SetSingleShape(curSingleM, curSingleN, quantGroupSize); // 8, 256, 512 --> 514us + GlobalTensor weightSlice; + uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN; + mnConfig.workspaceOffset = outOffset + mnConfig.yBaseOffset; + for (uint32_t loopK = 0; loopK < gmmSwigluQuantV2BaseParams->quantGroupNum; loopK++) { + mm.SetTensorA( + xGM[mnConfig.xBaseOffset + mnConfig.mIdx * mnConfig.k * mnConfig.singleM + loopK * quantGroupSize]); + if (gmmSwigluQuantV2BaseParams->isSingleTensor == 0) { + weightGM.SetGlobalBuffer(GetTensorAddr(groupIdx, weightTensorPtr)); + if constexpr (mmType::BT::format == CubeFormat::NZ) { + weightOffset = tailN * gmmSwigluQuantV2BaseParams->K; + weightSlice = weightGM[weightOffset + loopK * quantGroupSize * 64]; + } else { + weightOffset = tailN; + weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmSwigluQuantV2BaseParams->N]; + } + } else { + if constexpr (mmType::BT::format == CubeFormat::NZ) { + weightOffset = + static_cast(groupIdx) * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->K + + tailN * gmmSwigluQuantV2BaseParams->K; + weightSlice = weightGM[weightOffset + loopK * quantGroupSize * 64]; + } else { + weightOffset = + static_cast(groupIdx) * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->K + + tailN; + weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmSwigluQuantV2BaseParams->N]; + } + } + if (mnConfig.blockDimM == 1) { + weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } + mm.SetTensorB(weightSlice); + if (gmmSwigluQuantV2BaseParams->isSingleTensor == 0) { + weightScaleGM.SetGlobalBuffer(GetTensorAddr(groupIdx, weightScaleTensorPtr)); + mm.SetQuantVector(weightScaleGM[loopK * gmmSwigluQuantV2BaseParams->N + tailN]); + } else { + mm.SetQuantVector( + weightScaleGM[groupIdx * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->quantGroupNum + + loopK * gmmSwigluQuantV2BaseParams->N + tailN]); + } + mm.Iterate(); + mm.GetTensorC(mmOutGM[mnConfig.workspaceOffset], loopK == 0 ? 0 : 1); + } +} +} // namespace GroupedMatmulDequantSwigluQuant +#endif // GMM_SWIGLU_QUANT_V2_A8W4_MSD +#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_MID_H \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a8w4_msd_pipeline.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a8w4_msd_pipeline.h new file mode 100644 index 000000000..60e788dd8 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a8w4_msd_pipeline.h @@ -0,0 +1,236 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_a8w4_msd_pipeline.h + * \brief + */ +#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_PIPELINE_H +#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_PIPELINE_H + +#include +#include "grouped_matmul_swiglu_quant_v2_a8w4_msd_pre.h" +#include "grouped_matmul_swiglu_quant_v2_a8w4_msd_mid.h" +#include "grouped_matmul_swiglu_quant_v2_a8w4_msd_post.h" +#include "grouped_matmul_swiglu_quant_v2_utils.h" + +using namespace AscendC; +using namespace matmul; + +#ifdef GMM_SWIGLU_QUANT_V2_A8W4_MSD + +namespace GroupedMatmulDequantSwigluQuant { + +template +class GMMSwigluQuantPipelineSchedule { +private: + typename mmType::MT &mm; + TPipe *pipe; + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams; + const GMMSwigluQuantV2 *__restrict gmmSwigluQuantV2; + // WorkSpaceSplitConfig控制Workspace切割方式的结构体; + WorkSpaceSplitConfig workspaceSplitConfig; + WorkSpaceSplitConfig tempWorkspaceSplitConfig; + // 记录GM_ADDR的结构体 + GMAddrParams gmAddrParams; + // 前处理GMMA8W4PreProcess类 + GMMA8W4PreProcess preProcess; + // 中间处理GMMA8W4MidProcess类 + GMMA8W4MidProcess midProcess; + // 后处理GMMA8W4PostProcess类 + GMMA8W4PostProcess postProcess; + GlobalTensor groupListGM; + __aicore__ inline void InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig, + int32_t workspaceSplitLoopIdx); + +public: + __aicore__ inline GMMSwigluQuantPipelineSchedule( + typename mmType::MT &mm_, const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN, + const GMMSwigluQuantV2 *__restrict gmmSwigluIN, TPipe *tPipeIN) + : mm(mm_), midProcess(mm), gmmSwigluQuantV2BaseParams(gmmSwigluQuantV2BaseParamsIN), + gmmSwigluQuantV2(gmmSwigluIN), pipe(tPipeIN) + { + } + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, GM_ADDR xScale, + GM_ADDR weightAssistanceMatrix, GM_ADDR groupList, GM_ADDR y, GM_ADDR yScale, + GM_ADDR workspace); + __aicore__ inline void Process(); +}; + +template +__aicore__ inline void GMMSwigluQuantPipelineSchedule::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, + GM_ADDR xScale, GM_ADDR weightAssistanceMatrix, + GM_ADDR groupList, GM_ADDR y, GM_ADDR yScale, + GM_ADDR workspace) +{ + gmAddrParams.xGM = x; + gmAddrParams.weightGM = weight; + gmAddrParams.weightScaleGM = weightScale; + gmAddrParams.xScaleGM = xScale; + gmAddrParams.weightAuxiliaryMatrixGM = weightAssistanceMatrix; + gmAddrParams.groupListGM = groupList; + gmAddrParams.yGM = y; + gmAddrParams.yScaleGM = yScale; + gmAddrParams.workSpaceGM = workspace; + gmAddrParams.workSpaceOffset1 = gmmSwigluQuantV2BaseParams->workSpaceOffset1 / 2; + gmAddrParams.workSpaceOffset2 = gmmSwigluQuantV2BaseParams->workSpaceOffset1; + gmAddrParams.workSpaceOffset3 = + gmmSwigluQuantV2BaseParams->workSpaceOffset1 + gmmSwigluQuantV2BaseParams->workSpaceOffset2 / 2; + groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM); + InitWorkSpaceSplitConfig(workspaceSplitConfig); +} + +template +__aicore__ inline void GMMSwigluQuantPipelineSchedule::Process() +{ + // 1.对每次workspace切分做大循环。 + preProcess.Init(gmAddrParams, gmmSwigluQuantV2BaseParams); + midProcess.Init(gmAddrParams, gmmSwigluQuantV2BaseParams); + postProcess.Init(gmAddrParams, gmmSwigluQuantV2BaseParams, gmmSwigluQuantV2); + + // 1.前处理提前下发一次 + preProcess.Process(workspaceSplitConfig, 0, pipe); + for (int64_t workspaceSplitLoopIdx = 0; workspaceSplitLoopIdx < workspaceSplitConfig.loopCount; + workspaceSplitLoopIdx++) { + // 更新workspaceSplitConfig + UpdateWorkSpaceSplitConfig(workspaceSplitConfig, workspaceSplitLoopIdx); + if ASCEND_IS_AIV { + pipe->Reset(); + } + + SyncAll(); + // 2.第n次中处理 && 第n+1次前处理 && 第n-1次后处理 并行 + midProcess.Process(workspaceSplitConfig, workspaceSplitLoopIdx); + + preProcess.Process(workspaceSplitConfig, workspaceSplitLoopIdx + 1, pipe); + if ASCEND_IS_AIV { + pipe->Reset(); + SyncAll(); + } + postProcess.Process(tempWorkspaceSplitConfig, workspaceSplitLoopIdx - 1, pipe); + // 3.第n-1次后处理需要保留第n次的切分数据 + tempWorkspaceSplitConfig = workspaceSplitConfig; + // reset + if ASCEND_IS_AIV { + pipe->Reset(); + } + SyncAll(); + // 3.前一次后处理 && 后一次MM 并行 + } + // reset + if ASCEND_IS_AIV { + pipe->Reset(); + } + SyncAll(); + // // 4.最后一次后处理 + postProcess.Process(workspaceSplitConfig, workspaceSplitConfig.loopCount - 1, pipe); + if ASCEND_IS_AIV { + pipe->Destroy(); + } +} + +template +__aicore__ inline void +GMMSwigluQuantPipelineSchedule::InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig) +{ + if (gmmSwigluQuantV2BaseParams->groupListType == 0) { + workspaceSplitConfig.M = groupListGM.GetValue(gmmSwigluQuantV2->groupListLen - 1); + } else { + int64_t totalTmp = 0; + for (uint32_t i = 0; i < gmmSwigluQuantV2->groupListLen; i++) { + totalTmp += groupListGM.GetValue(i); + } + workspaceSplitConfig.M = totalTmp; + } + workspaceSplitConfig.loopCount = Ceil(workspaceSplitConfig.M, gmmSwigluQuantV2BaseParams->mLimit); + workspaceSplitConfig.notLastTaskSize = gmmSwigluQuantV2BaseParams->mLimit; + workspaceSplitConfig.lastLoopTaskSize = + workspaceSplitConfig.M - (workspaceSplitConfig.loopCount - 1) * gmmSwigluQuantV2BaseParams->mLimit; + workspaceSplitConfig.leftMatrixStartIndex = 0; + workspaceSplitConfig.rightMatrixExpertStartIndex = 0; + workspaceSplitConfig.rightMatrixExpertNextStartIndex = 0; + workspaceSplitConfig.isLastLoop = false; +} + +template +__aicore__ inline void +GMMSwigluQuantPipelineSchedule::UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig, + int32_t workspaceSplitLoopIdx) +{ + if (workspaceSplitLoopIdx < 0) + return; + workspaceSplitConfig.leftMatrixStartIndex = workspaceSplitLoopIdx * gmmSwigluQuantV2BaseParams->mLimit; + workspaceSplitConfig.rightMatrixExpertStartIndex = workspaceSplitConfig.rightMatrixExpertNextStartIndex; + workspaceSplitConfig.rightMatrixExpertEndIndex = workspaceSplitConfig.rightMatrixExpertStartIndex; + // 计算右专家矩阵的终止索引(rightMatrixExpertEndIndex) 和下一次的起始索引(rightMatrixExpertNextStartIndex) + int32_t curTaskNum = 0; + int32_t nextTaskNum = 0; + int32_t curTaskNumTmp = 0; + int32_t nextTaskNumTmp = 0; + if (gmmSwigluQuantV2BaseParams->groupListType == 1) { + for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertEndIndex; i++) { + curTaskNumTmp += groupListGM.GetValue(i); + } + if (workspaceSplitConfig.rightMatrixExpertEndIndex == 0) { + nextTaskNumTmp = groupListGM.GetValue(0); + } else { + for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertEndIndex; i++) { + nextTaskNumTmp += groupListGM.GetValue(i); + } + } + } + while (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwigluQuantV2->groupListLen) { + if (gmmSwigluQuantV2BaseParams->groupListType == 0) { + curTaskNum = groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex) - + workspaceSplitConfig.leftMatrixStartIndex; + } else { + curTaskNumTmp += groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex); + curTaskNum = curTaskNumTmp - workspaceSplitConfig.leftMatrixStartIndex; + } + int32_t nextTaskIdx = workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwigluQuantV2->groupListLen - 1 ? + gmmSwigluQuantV2->groupListLen - 1 : + workspaceSplitConfig.rightMatrixExpertEndIndex + 1; + if (gmmSwigluQuantV2BaseParams->groupListType == 0) { + nextTaskNum = groupListGM.GetValue(nextTaskIdx) - workspaceSplitConfig.leftMatrixStartIndex; + } else { + if (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwigluQuantV2->groupListLen - 1) { + nextTaskNumTmp += groupListGM.GetValue(nextTaskIdx); + } + nextTaskNum = nextTaskNumTmp - workspaceSplitConfig.leftMatrixStartIndex; + } + if (curTaskNum > gmmSwigluQuantV2BaseParams->mLimit) { + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex; + break; + } else if (curTaskNum == gmmSwigluQuantV2BaseParams->mLimit && + nextTaskNum > gmmSwigluQuantV2BaseParams->mLimit) { + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex + 1; + break; + } else if (nextTaskNum > gmmSwigluQuantV2BaseParams->mLimit) { + workspaceSplitConfig.rightMatrixExpertEndIndex++; + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex; + break; + } + workspaceSplitConfig.rightMatrixExpertEndIndex++; + } + workspaceSplitConfig.isLastLoop = workspaceSplitLoopIdx == workspaceSplitConfig.loopCount - 1 ? true : false; + + if (workspaceSplitConfig.isLastLoop) { + workspaceSplitConfig.rightMatrixExpertEndIndex = + workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwigluQuantV2->groupListLen ? + gmmSwigluQuantV2->groupListLen - 1 : + workspaceSplitConfig.rightMatrixExpertEndIndex; + } +} + +} // namespace GroupedMatmulDequantSwigluQuant +#endif // GMM_SWIGLU_QUANT_V2_A8W4_MSD +#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_PIPELINE_H \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a8w4_msd_post.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a8w4_msd_post.h new file mode 100644 index 000000000..741b22817 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a8w4_msd_post.h @@ -0,0 +1,444 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_a8w4_msd_post.h + * \brief + */ + +#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_POST_H +#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_POST_H + +#include "grouped_matmul_swiglu_quant_v2_utils.h" +#include "kernel_operator.h" + +#ifdef GMM_SWIGLU_QUANT_V2_A8W4_MSD + +namespace GroupedMatmulDequantSwigluQuant { +using namespace AscendC; +#define DOUBLE_BUFFER 2 +constexpr float DEFAULT_MUL_SCALE = 16.0f; +class GMMA8W4PostProcess { +public: + __aicore__ inline GMMA8W4PostProcess(){}; + __aicore__ inline void Init(const GMAddrParams gmAddrParams, + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN, + const GMMSwigluQuantV2 *__restrict gmmSwigluIN); + + __aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx, + TPipe *pipe); + static constexpr float FLOAT_INF = 3e+99; +private: + __aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx, + TPipe *pipe); + + __aicore__ inline void UpdateAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig); + + __aicore__ inline void VectorCompute(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void customDataCopyIn(uint32_t outLoopIdx, GlobalTensor &mmOutGM, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void customDataCopyOut(VecConfig &vecConfig, WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void PreLoadAuxiliaryMatrix(VecConfig &vecConfig); + + __aicore__ inline void Quant(uint32_t loopIdx, VecConfig &vecConfig); + + __aicore__ inline void Swiglu(uint32_t loopIdx, VecConfig &vecConfig); + + __aicore__ inline void MergeAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig); + + __aicore__ inline void MulPertokenScale(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig); + const GMMSwigluQuantV2 *__restrict gmmSwigluQuantV2; + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams; + GlobalTensor perTokenScaleGM; + GlobalTensor groupListGM; + GlobalTensor quantOutputGM; + GlobalTensor weightAuxiliaryMatrixGM; + GlobalTensor quantScaleOutputGM; + GlobalTensor mmOutGM1; + GlobalTensor mmOutGM2; + GlobalTensor mmOutGM; + TQue weightAuxiliaryMatrixInQueue; + TQue mmOutQueue; + TQue quantOutQueue; + TQue quantScaleOutQueue; + TBuf reduceWorkspace; + uint32_t blockIdx = 0; + int64_t aicCoreNum = 0; + int64_t aivCoreNum = 0; + GM_ADDR weightAuxiliaryMatrixTensorPtr; + float limited = FLOAT_INF; +}; + +__aicore__ inline void +GMMA8W4PostProcess::Init(const GMAddrParams gmAddrParams, + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN, + const GMMSwigluQuantV2 *__restrict gmmSwigluIN) +{ + if ASCEND_IS_AIV { + aicCoreNum = GetBlockNum(); + aivCoreNum = aicCoreNum * 2; + blockIdx = GetBlockIdx(); + gmmSwigluQuantV2BaseParams = gmmSwigluQuantV2BaseParamsIN; + gmmSwigluQuantV2 = gmmSwigluIN; + weightAuxiliaryMatrixGM.SetGlobalBuffer(GetTensorAddr(0, gmAddrParams.weightAuxiliaryMatrixGM)); + groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM, gmmSwigluQuantV2->groupListLen); + mmOutGM1.SetGlobalBuffer( + (__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset2)); + mmOutGM2.SetGlobalBuffer( + (__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset3)); + perTokenScaleGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.xScaleGM, gmmSwigluQuantV2BaseParams->M); + quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.yGM, gmmSwigluQuantV2BaseParams->M * + gmmSwigluQuantV2->tokenLen / + SWIGLU_REDUCE_FACTOR); + quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.yScaleGM, gmmSwigluQuantV2BaseParams->M); + weightAuxiliaryMatrixTensorPtr = gmAddrParams.weightAuxiliaryMatrixGM; + limited = gmmSwigluQuantV2BaseParams->swigluLimit; + } +} + +__aicore__ inline void GMMA8W4PostProcess::customDataCopyIn(uint32_t outLoopIdx, GlobalTensor &mmOutGM, + VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + LocalTensor _inMMLocal_0 = mmOutQueue.DeQue(); + const int64_t processNum = 2 * vecConfig.innerLoopNum * gmmSwigluQuantV2->tokenLen; + DataCopyExtParams copyParams_0{1, static_cast(processNum * SIZE_OF_HALF_2), 0, 0, 0}; + DataCopyPadExtParams padParams_0{false, 0, 0, 0}; + DataCopyPad(_inMMLocal_0[processNum], mmOutGM[vecConfig.curOffset * DOUBLE_ROW], copyParams_0, padParams_0); + + mmOutQueue.EnQue(_inMMLocal_0); + + LocalTensor _inMMLocal_1 = mmOutQueue.DeQue(); + // 1. fp16 -> fp32 + Cast(_inMMLocal_1.ReinterpretCast(), _inMMLocal_1[processNum], RoundMode::CAST_NONE, processNum); + + mmOutQueue.EnQue(_inMMLocal_1); + LocalTensor _inMMLocal_2 = mmOutQueue.DeQue(); + int32_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + // 2. high_4bit * 16 + low_4bit + for (uint32_t i = 0; i < vecConfig.innerLoopNum; i++) { + Muls(_inMMLocal_2[(DOUBLE_ROW * i) * gmmSwigluQuantV2->tokenLen], + _inMMLocal_2[(DOUBLE_ROW * i) * gmmSwigluQuantV2->tokenLen], DEFAULT_MUL_SCALE, + gmmSwigluQuantV2->tokenLen); + PipeBarrier(); + Add(_inMMLocal_2[i * gmmSwigluQuantV2->tokenLen], _inMMLocal_2[(DOUBLE_ROW * i) * gmmSwigluQuantV2->tokenLen], + _inMMLocal_2[(DOUBLE_ROW * i + 1) * gmmSwigluQuantV2->tokenLen], gmmSwigluQuantV2->tokenLen); + PipeBarrier(); + vecConfig.curIdx++; + } + vecConfig.curOffset = vecConfig.curIdx * gmmSwigluQuantV2->tokenLen; + mmOutQueue.EnQue(_inMMLocal_2); +} + +__aicore__ inline void GMMA8W4PostProcess::VectorCompute(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + // 1.辅助矩阵加回 + MergeAuxiliaryMatrix(loopIdx, vecConfig); + // 2.perToken反量化 + MulPertokenScale(loopIdx, vecConfig, workspaceSplitConfig); + // 3.Swiglu + Swiglu(loopIdx, vecConfig); + // 4.Quant + Quant(loopIdx, vecConfig); +} + +__aicore__ inline void GMMA8W4PostProcess::MergeAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig) +{ + // perChanelScale * perTokenScale + LocalTensor mmLocal = mmOutQueue.DeQue(); + LocalTensor weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue(); + Add(mmLocal[loopIdx * gmmSwigluQuantV2->tokenLen], mmLocal[loopIdx * gmmSwigluQuantV2->tokenLen], + weightAuxiliaryMatrixLocal, gmmSwigluQuantV2->tokenLen); + vecConfig.nextUpdateInterVal--; + mmOutQueue.EnQue(mmLocal); + weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal); +} + +__aicore__ inline void GMMA8W4PostProcess::MulPertokenScale(uint32_t loopIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + LocalTensor mmLocal = mmOutQueue.DeQue(); + int32_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + float scale = perTokenScaleGM.GetValue(loopIdx + workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Muls(mmLocal[loopIdx * gmmSwigluQuantV2->tokenLen], mmLocal[loopIdx * gmmSwigluQuantV2->tokenLen], scale, + gmmSwigluQuantV2->tokenLen); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); +} + +__aicore__ inline void GMMA8W4PostProcess::Swiglu(uint32_t loopIdx, VecConfig &vecConfig) +{ + // 高阶API swiglu + LocalTensor _inMMLocal = mmOutQueue.DeQue(); + float beta = 1.0f; + LocalTensor workspaceLocal = reduceWorkspace.Get(); + LocalTensor src0Local = + _inMMLocal[loopIdx * gmmSwigluQuantV2->tokenLen + gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR]; + LocalTensor src1Local = _inMMLocal[loopIdx * gmmSwigluQuantV2->tokenLen]; + + if (limited > 0.0f) { + Mins(src0Local, src0Local, limited, gmmSwigluQuantV2->tokenLen / 2); + PipeBarrier(); + Maxs(src0Local, src0Local, (-1.0f * limited), gmmSwigluQuantV2->tokenLen / 2); + PipeBarrier(); + Mins(src1Local, src1Local, limited, gmmSwigluQuantV2->tokenLen / 2); + PipeBarrier(); + } + SwiGLU(workspaceLocal, src0Local, src1Local, beta, gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR); + PipeBarrier(); + DataCopyParams repeatParams{ + 1, static_cast((gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR) / ALIGN_8_ELE), 0, 0}; + DataCopy(_inMMLocal[loopIdx * gmmSwigluQuantV2->tokenLen], workspaceLocal, repeatParams); + + mmOutQueue.EnQue(_inMMLocal); +} + +__aicore__ inline void GMMA8W4PostProcess::Quant(uint32_t loopIdx, VecConfig &vecConfig) +{ + LocalTensor _inMMLocal = mmOutQueue.DeQue(); + uint64_t preOffset = loopIdx * gmmSwigluQuantV2->tokenLen; + uint64_t halfTokenLen = gmmSwigluQuantV2->tokenLen / BISECT; + Abs(_inMMLocal[preOffset + gmmSwigluQuantV2->tokenLen / BISECT], _inMMLocal[preOffset], halfTokenLen); + PipeBarrier(); + // reduceMax + LocalTensor workLocal = reduceWorkspace.Get(halfTokenLen); + LocalTensor reduceResLocal = + reduceWorkspace.GetWithOffset(FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float)); + LocalTensor reduceTmpLocal = reduceWorkspace.GetWithOffset( + FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float) + UB_BLOCK_UNIT_SIZE); + ReduceMaxTemplate(reduceResLocal, workLocal, _inMMLocal[preOffset + gmmSwigluQuantV2->tokenLen / BISECT], + reduceTmpLocal, static_cast(halfTokenLen)); + int32_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float quantScale = reduceResLocal.GetValue(0) / QUANT_SCALE_INT8; + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + quantScaleLocal.SetValue(loopIdx, quantScale); + quantScale = 1 / quantScale; + int32_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Muls(_inMMLocal[preOffset], _inMMLocal[preOffset], quantScale, halfTokenLen); + PipeBarrier(); + LocalTensor quantLocal = quantOutQueue.DeQue(); + int32_t dstTempOffset = static_cast(preOffset / BISECT); + int32_t srcTempOffset = static_cast(preOffset); + int32_t tempCount = static_cast(halfTokenLen); + LocalTensor castSpace = reduceWorkspace.Get(UB_BLOCK_UNIT_SIZE); + CastFp32ToInt8Template(quantLocal, _inMMLocal, castSpace, dstTempOffset, srcTempOffset, tempCount); + mmOutQueue.EnQue(_inMMLocal); + quantOutQueue.EnQue(quantLocal); +} + +__aicore__ inline void GMMA8W4PostProcess::UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig, + int64_t workspaceSplitLoopIdx, TPipe *pipe) +{ + // 第一步 读取grouplist reduceSum 计算总数据个数 + vecConfig.M = workspaceSplitLoopIdx < workspaceSplitConfig.loopCount - 1 ? workspaceSplitConfig.notLastTaskSize : + workspaceSplitConfig.lastLoopTaskSize; + // 第二步 计算分核 + uint32_t eachCoreTaskNum = (vecConfig.M + aivCoreNum - 1) / aivCoreNum; + vecConfig.usedCoreNum = vecConfig.M >= aivCoreNum ? aivCoreNum : vecConfig.M; + uint32_t tailCoreIdx = vecConfig.M - (eachCoreTaskNum - 1) * vecConfig.usedCoreNum; + vecConfig.taskNum = blockIdx < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1; + vecConfig.startIdx = + blockIdx < tailCoreIdx ? eachCoreTaskNum * blockIdx : ((eachCoreTaskNum - 1) * blockIdx + tailCoreIdx); + vecConfig.curIdx = vecConfig.startIdx; + vecConfig.startOffset = vecConfig.startIdx * gmmSwigluQuantV2->tokenLen; + vecConfig.curOffset = vecConfig.startOffset; + int64_t curStartIdx = vecConfig.startIdx; + int64_t prevM = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize; + int64_t totalTmp = 0; + if (gmmSwigluQuantV2BaseParams->groupListType == 1) { + for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertStartIndex; i++) { + totalTmp += groupListGM.GetValue(i); + } + } + for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex; + groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; groupIdx++) { + int64_t currM = 0; + if (gmmSwigluQuantV2BaseParams->groupListType == 0) { + currM = groupListGM.GetValue(groupIdx); + } else { + totalTmp += groupListGM.GetValue(groupIdx); + currM = totalTmp; + } + int64_t tempM = currM - prevM; + prevM = currM; + if (curStartIdx >= 0 && curStartIdx - tempM < 0) { + vecConfig.curGroupIdx = groupIdx; + vecConfig.nextUpdateInterVal = tempM - curStartIdx; + } + curStartIdx -= tempM; + } + // 第三步 计算总数据量 + vecConfig.outLoopNum = + (vecConfig.taskNum + gmmSwigluQuantV2->maxProcessRowNum - 1) / gmmSwigluQuantV2->maxProcessRowNum; + vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwigluQuantV2->maxProcessRowNum ? + vecConfig.taskNum % gmmSwigluQuantV2->maxProcessRowNum : + gmmSwigluQuantV2->maxProcessRowNum; + + // 第四步 申请空间 + // 2 * row * n * sizeof(float) + row * n / 2 * sizeof(int8) + alignUp * sizeof(float) + n * sizeof(float) + + // n / 2 *sizeof(float) + 64 < 191 * 1024 + pipe->InitBuffer(mmOutQueue, 1, + 2 * gmmSwigluQuantV2->maxProcessRowNum * gmmSwigluQuantV2->tokenLen * sizeof(float)); + pipe->InitBuffer(quantOutQueue, 1, + gmmSwigluQuantV2->maxProcessRowNum * gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR * + sizeof(int8_t)); + pipe->InitBuffer(quantScaleOutQueue, 1, + AlignUp(gmmSwigluQuantV2->maxProcessRowNum, ALIGN_8_ELE) * sizeof(float)); + pipe->InitBuffer(weightAuxiliaryMatrixInQueue, 1, gmmSwigluQuantV2->tokenLen * sizeof(float)); + // two 32 byte buffer for reduceMax calculation in Quant. + pipe->InitBuffer(reduceWorkspace, gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(float) + + UB_BLOCK_UNIT_SIZE + UB_BLOCK_UNIT_SIZE); +} + +__aicore__ inline void GMMA8W4PostProcess::PreLoadAuxiliaryMatrix(VecConfig &vecConfig) +{ + LocalTensor weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue(); + DataCopyExtParams copyAuxiliaryMatrixParams{1, static_cast(gmmSwigluQuantV2->tokenLen * sizeof(float)), 0, + 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + if (gmmSwigluQuantV2BaseParams->isSingleTensor == 0) { + weightAuxiliaryMatrixGM.SetGlobalBuffer( + GetTensorAddr(vecConfig.curGroupIdx, weightAuxiliaryMatrixTensorPtr)); + DataCopyPad(weightAuxiliaryMatrixLocal, weightAuxiliaryMatrixGM, copyAuxiliaryMatrixParams, padParams); + } else { + DataCopyPad(weightAuxiliaryMatrixLocal, + weightAuxiliaryMatrixGM[vecConfig.curGroupIdx * gmmSwigluQuantV2->tokenLen], + copyAuxiliaryMatrixParams, padParams); + } + weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal); +} + +__aicore__ inline void GMMA8W4PostProcess::UpdateAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig) +{ + // 更新weightAuxiliaryMatrix + if (unlikely(vecConfig.nextUpdateInterVal == 0)) { + int64_t loop = gmmSwigluQuantV2->groupListLen - vecConfig.curGroupIdx; + while (loop--) { + if (gmmSwigluQuantV2BaseParams->groupListType == 0) { + int64_t curTemp = groupListGM.GetValue(vecConfig.curGroupIdx); + vecConfig.curGroupIdx++; + int64_t nextTemp = groupListGM.GetValue(vecConfig.curGroupIdx); + if (nextTemp != curTemp) { + vecConfig.nextUpdateInterVal = nextTemp - curTemp; + break; + } + } else { + vecConfig.curGroupIdx++; + int64_t nextUpdateInterValTmp = groupListGM.GetValue(vecConfig.curGroupIdx); + if (nextUpdateInterValTmp != 0) { + vecConfig.nextUpdateInterVal = nextUpdateInterValTmp; + break; + } + } + } + LocalTensor weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue(); + DataCopyExtParams copyParams{1, static_cast(gmmSwigluQuantV2->tokenLen * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(weightAuxiliaryMatrixLocal, + weightAuxiliaryMatrixGM[vecConfig.curGroupIdx * gmmSwigluQuantV2->tokenLen], copyParams, padParams); + weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal); + } +} + +__aicore__ inline void GMMA8W4PostProcess::Process(WorkSpaceSplitConfig &workspaceSplitConfig, + int64_t workspaceSplitLoopIdx, TPipe *pipe) +{ + if ASCEND_IS_AIV { + if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount || workspaceSplitLoopIdx < 0) { + return; + } + VecConfig vecConfig; + UpdateVecConfig(blockIdx, vecConfig, workspaceSplitConfig, workspaceSplitLoopIdx, pipe); + + if (blockIdx < vecConfig.usedCoreNum) { + mmOutGM = (workspaceSplitLoopIdx % 2 == 0 ? mmOutGM1 : mmOutGM2); + LocalTensor weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.AllocTensor(); + LocalTensor mmLocal = mmOutQueue.AllocTensor(); + LocalTensor quantScaleLocal = quantScaleOutQueue.AllocTensor(); + LocalTensor quantLocal = quantOutQueue.AllocTensor(); + + mmOutQueue.EnQue(mmLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); + quantOutQueue.EnQue(quantLocal); + weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal); + PreLoadAuxiliaryMatrix(vecConfig); + for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) { + vecConfig.innerLoopNum = outLoopIdx == (vecConfig.outLoopNum - 1) ? vecConfig.tailLoopNum : + gmmSwigluQuantV2->maxProcessRowNum; + int32_t eventIdMTE3ToMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventIdMTE3ToMTE2); + WaitFlag(eventIdMTE3ToMTE2); + // 1.matmul中间结果搬入 + 高四位与低四位合并 + customDataCopyIn(outLoopIdx, mmOutGM, vecConfig, workspaceSplitConfig); + + for (uint32_t innerLoopIdx = 0; innerLoopIdx < vecConfig.innerLoopNum; innerLoopIdx++) { + // 2.如果涉及group切换，更新辅助矩阵 + UpdateAuxiliaryMatrix(innerLoopIdx, vecConfig); + // 3. 四步vector计算（辅助矩阵加回、perToken反量化、Swiglu、Quant） + VectorCompute(innerLoopIdx, vecConfig, workspaceSplitConfig); + } + int32_t eventIdVToMTE3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventIdVToMTE3); + WaitFlag(eventIdVToMTE3); + customDataCopyOut(vecConfig, workspaceSplitConfig); + } + weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue(); + mmLocal = mmOutQueue.DeQue(); + quantScaleLocal = quantScaleOutQueue.DeQue(); + quantLocal = quantOutQueue.DeQue(); + + weightAuxiliaryMatrixInQueue.FreeTensor(weightAuxiliaryMatrixLocal); + mmOutQueue.FreeTensor(mmLocal); + quantScaleOutQueue.FreeTensor(quantScaleLocal); + quantOutQueue.FreeTensor(quantLocal); + } + } +} + +__aicore__ inline void GMMA8W4PostProcess::customDataCopyOut(VecConfig &vecConfig, + WorkSpaceSplitConfig &workspaceSplitConfig) +{ + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + DataCopyParams copyParams_0{1, (uint16_t)(vecConfig.innerLoopNum * sizeof(float)), 0, 0}; + DataCopyPad(quantScaleOutputGM[workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx], quantScaleLocal, + copyParams_0); + LocalTensor quantLocal = quantOutQueue.DeQue(); + DataCopyParams copyParams_1{ + 1, (uint16_t)(vecConfig.innerLoopNum * gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)), 0, + 0}; + DataCopyPad(quantOutputGM[(workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx) * + gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR], + quantLocal, copyParams_1); + + vecConfig.startIdx += vecConfig.innerLoopNum; + vecConfig.startOffset = vecConfig.startIdx * gmmSwigluQuantV2->tokenLen; + quantOutQueue.EnQue(quantLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); +} + +} // namespace GroupedMatmulDequantSwigluQuant +#endif // GMM_SWIGLU_QUANT_V2_A8W4_MSD +#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_POST_H diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a8w4_msd_pre.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a8w4_msd_pre.h new file mode 100644 index 000000000..4037b1756 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_a8w4_msd_pre.h @@ -0,0 +1,219 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_a8w4_msd_pre.h + * \brief + */ + +#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_PRE_H +#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_PRE_H + +#include "grouped_matmul_swiglu_quant_v2_utils.h" +#include "kernel_operator.h" + +#ifdef GMM_SWIGLU_QUANT_V2_A8W4_MSD + +namespace GroupedMatmulDequantSwigluQuant { +using namespace AscendC; +#define BUFFER_NUM_A8W4_PRE 1 +constexpr int TWO = 2; +constexpr int EIGHT = 8; +constexpr size_t LEN_128 = 128; // 16bit operator +constexpr int DATA_BLOCK_SIZE_32 = 32; +class GMMA8W4PreProcess { +public: + __aicore__ inline GMMA8W4PreProcess(){}; + __aicore__ inline void Init(const GMAddrParams gmAddrParams, + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN); + __aicore__ inline void CalculateTaskInfoEachCore(uint32_t &curCoreTaskNum_, uint32_t &curCoreStartOffset_); + __aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx, + TPipe *pipe); + __aicore__ inline void CustomInitBuffer(TPipe *pipe); + +private: + TQue vecInQueueX, vecInQueueXBak; + TQue vecOutQueueA1; + TQue vecOutQueueA2; + TQue vecOutQueueA3; + TQue vecOutQueue0F; + TQue vecOutQueueRowSum; + TBuf tempBuff; + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams; + LocalTensor xTensor; + LocalTensor xHighHalfTensor; + LocalTensor xHighFloatTensor; + LocalTensor xLowHalfTensor; + LocalTensor xLowHalfTensor2; + LocalTensor xHighI4Tensor; + LocalTensor xLowI4Tensor; + LocalTensor xLowI16Tensor; + LocalTensor xRowSumTensor; + + GlobalTensor xGM; + GlobalTensor yGm; + GlobalTensor yGm1; + GlobalTensor yGm2; + + uint32_t vK{0}; + uint32_t vKAlign{0}; + uint32_t totalM{0}; + uint32_t blockDim{0}; + uint32_t curCoreId{0}; + uint32_t curCoreTaskNum{0}; + uint32_t curCoreStartOffset{0}; + uint32_t curCoreOuterLoopNum{0}; + uint32_t curCoreInnerTailLoopNum{0}; + uint32_t groupNum{0}; +}; + +__aicore__ inline void +GMMA8W4PreProcess::Init(const GMAddrParams gmAddrParams, + const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN) +{ + if ASCEND_IS_AIV { + xGM.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.xGM); + yGm1.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.workSpaceGM); + yGm2.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset1); + gmmSwigluQuantV2BaseParams = gmmSwigluQuantV2BaseParamsIN; + vK = gmmSwigluQuantV2BaseParams->K; + groupNum = static_cast(gmmSwigluQuantV2BaseParams->groupNum); + // M * K * 7B (1B + 0.5B + 0.5B + 2B + 4B) <= UBsize - 256B + blockDim = GetBlockNum() * GetTaskRation(); + } +} + +__aicore__ inline void GMMA8W4PreProcess::CustomInitBuffer(TPipe *pipe) +{ + pipe->InitBuffer(vecInQueueX, BUFFER_NUM_A8W4_PRE, vK * sizeof(int8_t)); // K * 1B + pipe->InitBuffer(vecOutQueueA1, BUFFER_NUM_A8W4_PRE, vK * sizeof(int4b_t)); // K * 0.5B + pipe->InitBuffer(vecOutQueueA2, BUFFER_NUM_A8W4_PRE, vK * sizeof(int4b_t)); // K * 0.5B + pipe->InitBuffer(vecOutQueueA3, BUFFER_NUM_A8W4_PRE, vK * SIZE_OF_HALF_2); // K * 2B + // xLowHalfTensor, xLowHalfTensor2 and xHighFloatTensor share the same buffer + pipe->InitBuffer(tempBuff, vK * sizeof(float)); // K * 4B + constexpr int BUFFER_SIZE_256B = 128 * sizeof(int16_t); + pipe->InitBuffer(vecOutQueue0F, BUFFER_NUM_A8W4_PRE, BUFFER_SIZE_256B); // 256B +} + + +__aicore__ inline void GMMA8W4PreProcess::CalculateTaskInfoEachCore(uint32_t &curCoreTaskNum_, + uint32_t &curCoreStartOffset_) +{ + // 均分任务数 + int64_t eachCoreTaskNum = (totalM + blockDim - 1) / blockDim; // 每个核处理的数据量 + // 尾核任务数 + int64_t taskNumPertailCore = eachCoreTaskNum - 1; + // 实际使用核数 + int64_t usedCoreNum = totalM >= blockDim ? blockDim : totalM; + // 尾核起始索引 + uint32_t tailCoreIdx = totalM - (eachCoreTaskNum - 1) * usedCoreNum; + curCoreId = GetBlockIdx(); + // 每个核处理的任务数量 = 是否为尾核 ？均分任务数 ：(均分任务数 - 1) + curCoreTaskNum_ = curCoreId < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1; + // 每个核处理的起始偏移地址 = 是否为尾核 ？均分任务数 * blockId : (均分任务数 - 1) * blockId + 尾核起始索引 + curCoreStartOffset_ = + curCoreId < tailCoreIdx ? eachCoreTaskNum * curCoreId : ((eachCoreTaskNum - 1) * curCoreId + tailCoreIdx); +} + +__aicore__ inline void GMMA8W4PreProcess::Process(WorkSpaceSplitConfig &workspaceSplitConfig, + int64_t workspaceSplitLoopIdx, TPipe *pipe) +{ + if ASCEND_IS_AIV { + if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount) { + return; + } + yGm = (workspaceSplitLoopIdx % 2 == 0 ? yGm1 : yGm2); + CustomInitBuffer(pipe); + constexpr int32_t MASK = 128; + xTensor = vecInQueueX.AllocTensor(); + xHighI4Tensor = vecOutQueueA1.AllocTensor(); + xLowI4Tensor = vecOutQueueA2.AllocTensor(); + xHighHalfTensor = vecOutQueueA3.AllocTensor(); + const uint32_t xLowHalfOffset = vK * SIZE_OF_HALF_2; + xLowHalfTensor = tempBuff.GetWithOffset(xLowHalfOffset, 0); + xLowHalfTensor2 = tempBuff.GetWithOffset(xLowHalfOffset, xLowHalfOffset); + xLowI16Tensor = vecOutQueue0F.AllocTensor(); + + Duplicate(xLowI16Tensor, static_cast(0x0F0F), MASK); // get rid of high 4 bits in every int8 + PipeBarrier(); + const size_t LEN_VK = (vK / 2) / 128; + const size_t LAST_LEN_VK = (vK % 256) / 2; + const half ONE_SIXTEENTH = static_cast(0.0625f); + // groupList仅支持count + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + totalM = workspaceSplitLoopIdx < workspaceSplitConfig.loopCount - 1 ? workspaceSplitConfig.notLastTaskSize : + workspaceSplitConfig.lastLoopTaskSize; + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + CalculateTaskInfoEachCore(curCoreTaskNum, curCoreStartOffset); + SetFlag(EVENT_ID0); // 0 + SetFlag(EVENT_ID0); // 1 + SetFlag(EVENT_ID1); // 2 + + for (uint32_t xloop = 0; xloop < curCoreTaskNum; xloop++) { + uint64_t relStartAddr = (xloop + curCoreStartOffset) * vK; + uint64_t absStartAddr = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize * vK + relStartAddr; + // 高四位处理开始 + WaitFlag(EVENT_ID0); // 0 + DataCopy(xTensor, xGM[absStartAddr], vK); + SetFlag(EVENT_ID0); // 3 + WaitFlag(EVENT_ID0); // 3 + Cast(xHighHalfTensor, xTensor, AscendC::RoundMode::CAST_NONE, vK); + PipeBarrier(); + Muls(xHighHalfTensor, xHighHalfTensor, ONE_SIXTEENTH, vK); + PipeBarrier(); + WaitFlag(EVENT_ID1); // 2 + Cast(xHighI4Tensor, xHighHalfTensor, AscendC::RoundMode::CAST_FLOOR, vK); + SetFlag(EVENT_ID0); // 4 + WaitFlag(EVENT_ID0); // 4 + DataCopy(yGm[relStartAddr], xHighI4Tensor.ReinterpretCast(), vK / 2); + // 高四位处理结束 + + // 低四位处理开始 + SetFlag(EVENT_ID1); // 2 + And(xLowHalfTensor.ReinterpretCast(), xTensor.ReinterpretCast(), xLowI16Tensor, LEN_128, + LEN_VK, {1, 1, 1, 8, 8, 0}); + if (LAST_LEN_VK > 0) { + And(xLowHalfTensor[LEN_VK * LEN_128].ReinterpretCast(), + xTensor[LEN_VK * LEN_128 * TWO].ReinterpretCast(), xLowI16Tensor, LAST_LEN_VK, 1, + {1, 1, 1, 8, 8, 0}); + } + PipeBarrier(); + SetFlag(EVENT_ID0); // 0 + Cast(xLowHalfTensor2.ReinterpretCast(), xLowHalfTensor.ReinterpretCast(), + AscendC::RoundMode::CAST_NONE, vK); + PipeBarrier(); + const half MINUS_EIGHT = static_cast(-8); + Adds(xHighHalfTensor, xLowHalfTensor2, MINUS_EIGHT, vK); + PipeBarrier(); + WaitFlag(EVENT_ID0); // 1 + Cast(xLowI4Tensor, xHighHalfTensor.ReinterpretCast(), AscendC::RoundMode::CAST_NONE, vK); + SetFlag(EVENT_ID1); // 5 + WaitFlag(EVENT_ID1); // 5 + DataCopy(yGm[relStartAddr + vK / TWO], xLowI4Tensor.ReinterpretCast(), vK / TWO); + SetFlag(EVENT_ID0); // 1 + // 低四位处理结束 + } + + WaitFlag(EVENT_ID0); // 0 + WaitFlag(EVENT_ID0); // 1 + WaitFlag(EVENT_ID1); // 2 + vecInQueueX.FreeTensor(xTensor); + vecOutQueueA1.FreeTensor(xHighI4Tensor); + vecOutQueueA2.FreeTensor(xLowI4Tensor); + vecOutQueueA3.FreeTensor(xHighHalfTensor); + vecOutQueue0F.FreeTensor(xLowI16Tensor); + } +} + +} // namespace GroupedMatmulDequantSwigluQuant +#endif // GMM_SWIGLU_QUANT_V2_A8W4_MSD +#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_PRE_H \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_apt.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_apt.cpp new file mode 100644 index 000000000..20fe9ebc3 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_apt.cpp @@ -0,0 +1,66 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_apt.cpp + * \brief + */ + +#include "kernel_operator.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" +#if ORIG_DTYPE_X_SCALE == DT_FLOAT8_E8M0 + #include "arch35/grouped_matmul_swiglu_quant_v2_mxquant.h" +#elif ORIG_DTYPE_X_SCALE == DT_FLOAT + #include "arch35/grouped_matmul_swiglu_quant_v2_pertoken_quant.h" +#endif +#include "arch35/grouped_matmul_swiglu_quant_v2_tiling_key.h" + +#define FLOAT_OVERFLOW_MODE_CTRL 60 + +using namespace AscendC; +using namespace matmul; + +template +__global__ __aicore__ void grouped_matmul_swiglu_quant_v2(GM_ADDR x, GM_ADDR xScale, GM_ADDR groupList, GM_ADDR weight, + GM_ADDR weightScale, GM_ADDR weightAssistanceMatrix, + GM_ADDR bias, GM_ADDR smoothScale, GM_ADDR y, GM_ADDR yScale, + GM_ADDR workspace, GM_ADDR tiling) +{ + TPipe tPipe; + GM_ADDR userWorkspace = GetUserWorkspace(workspace); + int64_t oriOverflowMode = AscendC::GetCtrlSpr(); + // enable overflow mode to avoid nan/inf value + AscendC::SetCtrlSpr(0); +#if ORIG_DTYPE_X_SCALE == DT_FLOAT8_E8M0 + if (QUANT_B_TRANS == GMM_SWIGLU_QUANT_NO_TRANS && QUANT_A_TRANS == GMM_SWIGLU_QUANT_NO_TRANS) { // transX = false, transW = false + GmmSwigluAswt( + x, weight, weightScale, xScale, weightAssistanceMatrix, smoothScale, groupList, y, yScale, workspace, + tiling); + } else if (QUANT_B_TRANS == GMM_SWIGLU_QUANT_TRANS && QUANT_A_TRANS == GMM_SWIGLU_QUANT_NO_TRANS) { // transX = false, transW = true + GmmSwigluAswt( + x, weight, weightScale, xScale, weightAssistanceMatrix, smoothScale, groupList, y, yScale, workspace, + tiling); + } +#elif ORIG_DTYPE_X_SCALE == DT_FLOAT + if (QUANT_B_TRANS == GMM_SWIGLU_QUANT_NO_TRANS && + QUANT_A_TRANS == GMM_SWIGLU_QUANT_NO_TRANS) { // transX = false, transW = false + GmmSwigluAswtPertoken( + x, weight, weightScale, xScale, weightAssistanceMatrix, smoothScale, groupList, y, yScale, workspace, + tiling, &tPipe); + } else if (QUANT_B_TRANS == GMM_SWIGLU_QUANT_TRANS && + QUANT_A_TRANS == GMM_SWIGLU_QUANT_NO_TRANS) { // transX = false, transW = true + GmmSwigluAswtPertoken( + x, weight, weightScale, xScale, weightAssistanceMatrix, smoothScale, groupList, y, yScale, workspace, + tiling, &tPipe); + } +#endif + AscendC::SetCtrlSpr(oriOverflowMode); +} diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_utils.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_utils.h new file mode 100644 index 000000000..42ca81977 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_v2/op_kernel/grouped_matmul_swiglu_quant_v2_utils.h @@ -0,0 +1,297 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_v2_utils.h + * \brief + */ + +#ifndef OP_KERNEL_GROUPED_MATMUL_DEQUANT_SWIGLU_QUANT_V2_UTILS_H +#define OP_KERNEL_GROUPED_MATMUL_DEQUANT_SWIGLU_QUANT_V2_UTILS_H + +// A8W4 MSD场景 +#if defined(ORIG_DTYPE_X) && defined(DT_INT8) && ORIG_DTYPE_X == DT_INT8 && defined(ORIG_DTYPE_WEIGHT) && \ + defined(DT_INT4) && ORIG_DTYPE_WEIGHT == DT_INT4 + #define GMM_SWIGLU_QUANT_V2_A8W4_MSD + using DTYPE_X_A8W4_MSD = AscendC::int4b_t; +// A4W4 场景 +#elif defined(ORIG_DTYPE_X) && defined(DT_INT4) && ORIG_DTYPE_X == DT_INT4 && defined(ORIG_DTYPE_WEIGHT) && \ + defined(DT_INT4) && ORIG_DTYPE_WEIGHT == DT_INT4 + #define GMM_SWIGLU_QUANT_V2_A4W4 +// A8W8 场景 +#elif defined(ORIG_DTYPE_X) && defined(DT_INT8) && ORIG_DTYPE_X == DT_INT8 && defined(ORIG_DTYPE_WEIGHT) && \ + defined(DT_INT8) && ORIG_DTYPE_WEIGHT == DT_INT8 + #define GMM_SWIGLU_QUANT_V2_A8W8 +#endif // 场景分类 + +#if defined(FORMAT_WEIGHT) && FORMAT_WEIGHT == FORMAT_FRACTAL_NZ + constexpr CubeFormat wFormat = CubeFormat::NZ; +#elif defined(FORMAT_WEIGHT) && FORMAT_WEIGHT == FORMAT_ND + constexpr CubeFormat wFormat = CubeFormat::ND; +#endif // weight格式分类 + + +namespace GroupedMatmulDequantSwigluQuant { +using namespace AscendC; + + +constexpr uint32_t UB_BLOCK_UNIT_SIZE = 32; +constexpr uint32_t FLOAT_UB_BLOCK_UNIT_SIZE = 8; +constexpr uint32_t VEC_LEN_ONCE_REPEAT_ELE = 64; +constexpr uint32_t VEC_LEN_ONCE_REPEAT_BLOCK = 8; +constexpr uint32_t FP32_LEN_64_REPEAT = 4096; +constexpr uint32_t REPEAT_64 = 64; +constexpr uint32_t REPEAT_8 = 8; +constexpr uint32_t BISECT = 2; +constexpr uint32_t MOD_32_MASK = 0x1F; +constexpr uint32_t MOD_16_MASK = 0x0F; +constexpr uint32_t ALIGN_8_ELE = 8; +constexpr uint32_t ALIGN_16_ELE = 16; +constexpr uint32_t NUM_2 = 2; +constexpr int64_t SWIGLU_REDUCE_FACTOR = 2; +constexpr int64_t DOUBLE_BUFFER = 2; +constexpr int64_t DOUBLE_ROW = 2; +constexpr int64_t SIZE_OF_HALF_2 = 2; +constexpr uint8_t NUM_8 = 8; +constexpr float QUANT_SCALE_INT8 = 127.0f; + +constexpr MatmulConfig matmulCFGUnitFlag{false, false, true, 0, 0, 0, false, false, false, false, false, 0, 0, 0, + 0, 0, 0, 0, false}; +constexpr MatmulConfig NZ_CFG_MDL = GetMDLConfig(false, false, 0, true, false, false, false); +constexpr MatmulConfig CUSTOM_CFG_MDL = GetMDLConfig(false, false, 0, true, false, false, true); + +template +struct MMImplType { + using AT = AT_; + using BT = BT_; + using CT = CT_; + using BiasT = BiasT_; + using MT = matmul::MatmulImpl; +}; + +template +struct MMImplTypeCustom { + using AT = AT_; + using BT = BT_; + using CT = CT_; + // bias未被使用但高阶模板参数需要传入 + using BiasT = MatmulType; + using MT = matmul::MatmulImpl; +}; + +struct MNConfig { + uint32_t m = 0; + uint32_t k = 0; + uint32_t n = 0; + uint32_t baseM = 0; + uint32_t baseN = 0; + uint32_t baseK = 0; + uint32_t mIdx = 0; + uint32_t nIdx = 0; + uint32_t blockDimM = 0; + uint32_t blockDimN = 0; + uint32_t singleM = 0; + uint32_t singleN = 0; + uint64_t wBaseOffset = 0; + uint64_t mAxisBaseOffset = 0; + uint64_t nAxisBaseOffset = 0; + uint64_t xBaseOffset = 0; + uint64_t yBaseOffset = 0; + uint64_t wOutOffset = 0; + uint64_t workspaceOffset = 0; +}; + +struct VecConfig { + int64_t M = 0; + int64_t usedCoreNum = 0; + int64_t startOffset = 0; + int64_t curOffset = 0; + int64_t startIdx = 0; + int64_t curIdx = 0; + int64_t taskNum = 0; + int64_t curGroupIdx = 0; + int64_t outLoopNum = 0; + int64_t innerLoopNum = 0; + int64_t tailLoopNum = 0; + int64_t nextUpdateInterVal = 0; +}; + +struct WorkSpaceSplitConfig { + int64_t M = 0; + int64_t loopCount = 0; + int64_t leftMatrixStartIndex = 0; + int64_t rightMatrixExpertStartIndex = 0; + int64_t rightMatrixExpertNextStartIndex = 0; + int64_t rightMatrixExpertEndIndex = 0; + int64_t notLastTaskSize = 0; + int64_t lastLoopTaskSize = 0; + bool isLastLoop = false; +}; + +struct GMAddrParams { + // 输入 GM Tensor + GM_ADDR xGM; // 左矩阵 + GM_ADDR weightGM; // 右矩阵 + GM_ADDR weightScaleGM; // 权重scale + GM_ADDR xScaleGM; // 激活scale + GM_ADDR weightAuxiliaryMatrixGM; // 权重辅助矩阵 + GM_ADDR groupListGM; // 分组矩阵 + GM_ADDR smoothScaleGM; // 平滑缩放因子 + // 输出 GM Tensor + GM_ADDR yGM; // 输出量化矩阵 + GM_ADDR yScaleGM; // 输出scale矩阵 + // workspace GM Tensor + GM_ADDR workSpaceGM; // 左矩阵前处理结果矩阵 (double workspace) + 中间处理结果矩阵 (double workspace) + int64_t workSpaceOffset1; + int64_t workSpaceOffset2; + int64_t workSpaceOffset3; +}; + +template +__aicore__ inline auto AlignUp(T a) -> T +{ + if (unlikely(base == 0)) { + return a; + } + return (a + base - 1) / base * base; +} + +template +__aicore__ inline auto AlignUp(T a, T base) -> T +{ + if (unlikely(base == 0)) { + return a; + } + return (a + base - 1) / base * base; +} + +template <> +__aicore__ inline uint32_t AlignUp<4, uint32_t>(uint32_t a) +{ + // to be Multiple of 4, result should be in a format of b(xxxx,x100). + // This means last two bits should be zero, requiring that + // result = num & b(1111,1100) = num & (~3). + // &(~3) operator may reduces num into the range [num, num - 3]. + // As the result should be no less than a (result >= a), it means num - 3 >= a in the worst case. + // In this case, num >= a+3. On the other hand, num should also be less then a+4, otherwise, + // the result will not be least multiple of 4 for 3. In other cases like [num, num - 2], + // num = a + 3 also satisfies the goal condition. + return (a + 3) & ~3; // & ~3: set last two bits of (a+3) to be zero +} + +template <> +__aicore__ inline uint32_t AlignUp<8, uint32_t>(uint32_t a) +{ + // In general, if we want to get the least multiple of b (b is the power of 2) for a, + // it comes to a conclusion from the above comment: result = (a + (b - 1)) & (~b) + return (a + 7) & ~7; // & ~7: set last four bits of (a+7) to be zero +} + +template <> +__aicore__ inline uint32_t AlignUp<16, uint32_t>(uint32_t a) +{ + // In general, if we want to get the least multiple of b (b is the power of 2) for a, + // it comes to a conclusion from the above comment: result = (a + (b - 1)) & (~b) + return (a + 15) & ~15; // & ~15: set last four bits of (a+15) to be zero +} + +template <> +__aicore__ inline uint32_t AlignUp<32, uint32_t>(uint32_t a) +{ + // refer to the above comments. + return (a + 31) & ~31; // & ~31: set last five bits of (a+31) to be zero} +} + +__aicore__ inline void ReduceMaxSmall(const LocalTensor &dstLocal, const LocalTensor &workLocal, + const LocalTensor &srcLocal, uint32_t count) +{ + /** + * @brief ReduceMaxSmall 此函数仅支持入参count小于4096。 + */ + uint32_t repeat = count / VEC_LEN_ONCE_REPEAT_ELE; + uint32_t tailNum = count % VEC_LEN_ONCE_REPEAT_ELE; + if (likely(repeat > 0)) { + WholeReduceMax(workLocal, srcLocal, VEC_LEN_ONCE_REPEAT_ELE, repeat, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK, + ReduceOrder::ORDER_ONLY_VALUE); + PipeBarrier(); + } + if (unlikely(tailNum != 0)) { + WholeReduceMax(workLocal[repeat], srcLocal[count - tailNum], tailNum, 1, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK, + ReduceOrder::ORDER_ONLY_VALUE); + PipeBarrier(); + repeat += 1; + } + WholeReduceMax(dstLocal, workLocal, repeat, 1, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK, ReduceOrder::ORDER_ONLY_VALUE); +} + +__aicore__ inline void ReduceMaxTemplate(const LocalTensor &dstLocal, const LocalTensor &workLocal, + const LocalTensor &srcLocal, const LocalTensor &resTmpLocal, + uint32_t count) +{ + /** + * @brief 当前算子仅支持[32, 10240]长度的词向量维度N，对应此函数count入参范围在[16, 5120]。 + * @param [in] count: 本函数支持count范围为[1,8192]。 + */ + if (count <= FP32_LEN_64_REPEAT) { + ReduceMaxSmall(dstLocal, workLocal, srcLocal, count); + PipeBarrier(); + } else { + BlockReduceMax(workLocal, srcLocal, REPEAT_64, VEC_LEN_ONCE_REPEAT_ELE, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK); + PipeBarrier(); + + BlockReduceMax(workLocal, workLocal, REPEAT_8, VEC_LEN_ONCE_REPEAT_ELE, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK); + PipeBarrier(); + + WholeReduceMax(resTmpLocal, workLocal, VEC_LEN_ONCE_REPEAT_ELE, 1, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK, + ReduceOrder::ORDER_ONLY_VALUE); + PipeBarrier(); + + ReduceMaxSmall(dstLocal, workLocal, srcLocal[FP32_LEN_64_REPEAT], count - FP32_LEN_64_REPEAT); + PipeBarrier(); + + const BinaryRepeatParams repeatParams = {1, 1, 1, NUM_8, NUM_8, NUM_8}; + Max(dstLocal, dstLocal, resTmpLocal, 1, 1, repeatParams); + } +} + +__aicore__ inline void CastFp32ToInt8Template(LocalTensor &dstLocal, LocalTensor &srcLocal, + LocalTensor &oneBlockWorkspace, int32_t dstOffset, + int32_t srcOffset, int32_t count) +{ + Cast(srcLocal[srcOffset].ReinterpretCast(), srcLocal[srcOffset], RoundMode::CAST_RINT, count); + PipeBarrier(); + if ((dstOffset & MOD_32_MASK) == 0) { + Cast(dstLocal[dstOffset], srcLocal[srcOffset].ReinterpretCast(), RoundMode::CAST_RINT, count); + } else if ((dstOffset & MOD_16_MASK) == 0) { + Cast(dstLocal[dstOffset + ALIGN_16_ELE], srcLocal[srcOffset + ALIGN_8_ELE].ReinterpretCast(), + RoundMode::CAST_RINT, count - ALIGN_16_ELE); + PipeBarrier(); + Cast(oneBlockWorkspace, srcLocal[srcOffset].ReinterpretCast(), RoundMode::CAST_RINT, ALIGN_16_ELE); + PipeBarrier(); + for (int32_t i = 0; i < ALIGN_16_ELE; i++) { + int8_t temp = oneBlockWorkspace.GetValue(i); + dstLocal.SetValue(dstOffset + i, temp); + } + PipeBarrier(); + } +} + +template +__aicore__ inline __gm__ T* GetTensorAddr(uint16_t index, GM_ADDR tensorPtr) +{ + __gm__ uint64_t* dataAddr = reinterpret_cast<__gm__ uint64_t*>(tensorPtr); + uint64_t tensorPtrOffset = *dataAddr; + + __gm__ uint64_t* retPtr = dataAddr + (tensorPtrOffset >> 3); + return reinterpret_cast<__gm__ T*>(*(retPtr + index)); +} +} // namespace GroupedMatmulDequantSwigluQuant + +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/CMakeLists.txt b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/CMakeLists.txt new file mode 100644 index 000000000..298ed5677 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/grouped_matmul_swiglu_quant_torch_adpt.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/grouped_matmul_swiglu_quant_torch_adpt.h new file mode 100644 index 000000000..5ca311b29 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/grouped_matmul_swiglu_quant_torch_adpt.h @@ -0,0 +1,99 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef GROUPED_MATMUL_SWIGLU_QUANT_TORCH_ADPT_H +#define GROUPED_MATMUL_SWIGLU_QUANT_TORCH_ADPT_H +namespace vllm_fl { +const int64_t INT4_NUMS_IN_INT32 = 8; +std::tuple grouped_matmul_swiglu_quant_weight_nz( + const at::Tensor &x, const at::Tensor &weight, const at::Tensor &weight_scale, const at::Tensor &x_scale, + const at::Tensor &group_list, const c10::optional &bias, const c10::optional &offset, + double swiglu_limit) +{ + int m = x.sizes()[0]; + int n = weight.sizes()[2]; + bool is_a8w4 = x.dtype() == at::kChar && weight.dtype() == at::kInt; + if (is_a8w4) { + n *= INT4_NUMS_IN_INT32; + } + + at::Tensor output = at::empty({m, n/2}, x.options().dtype(c10::ScalarType::Char)); + at::Tensor output_scale = at::empty({m}, x.options().dtype(c10::ScalarType::Float)); + at::Tensor output_offset = at::empty({}, x.options().dtype(c10::ScalarType::Float)); + double swiglu_limit_f = static_cast(swiglu_limit); + + EXEC_NPU_CMD( + aclnnGroupedMatmulSwigluQuantWeightNZ, + x, + weight, + bias, + offset, + weight_scale, + x_scale, + group_list, + swiglu_limit_f, + output, + output_scale, + output_offset); + return std::tuple(output, output_scale, output_offset); +} + +std::tuple grouped_matmul_swiglu_quant( + const at::Tensor &x, const at::Tensor &weight, const at::Tensor &weight_scale, const at::Tensor &x_scale, + const at::Tensor &group_list, const c10::optional &bias, const c10::optional &offset, + double swiglu_limit) +{ + return grouped_matmul_swiglu_quant_weight_nz( + x, weight, weight_scale, x_scale, group_list, bias, offset, swiglu_limit); +} + +std::tuple grouped_matmul_swiglu_quant_weight_nz_tensor_list( + const at::Tensor & x, + const at::TensorList & weight, + const at::TensorList & weight_scale, + const at::Tensor & x_scale, + const at::Tensor & group_list, + const c10::optional & bias, + const c10::optional & offset, + double swiglu_limit) +{ + auto x_size = x.sizes(); + int n = weight[0].sizes()[1]; + int m = x_size[0]; + int k = x_size[1]; + + at::Tensor output = at::empty({m, n/2}, x.options().dtype(at::kChar)); + at::Tensor output_scale = at::empty({m}, x.options().dtype(at::kFloat)); + at::Tensor output_offset = at::empty({m}, x.options().dtype(at::kFloat)); + float swiglu_limit_f = static_cast(swiglu_limit); + + EXEC_NPU_CMD( + aclnnGroupedMatmulSwigluQuantWeightNzTensorList, + x, + weight, + bias, + offset, + weight_scale, + x_scale, + group_list, + swiglu_limit_f, + output, + output_scale, + output_offset); + + return std::tuple(output, output_scale, output_offset); +} +} +#endif \ No newline at end of file diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/CMakeLists.txt b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/CMakeLists.txt new file mode 100644 index 000000000..3df51a55e --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/CMakeLists.txt @@ -0,0 +1,22 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnExc PRIVATE + grouped_matmul_swiglu_quant_weight_nz_tensor_list_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME GroupedMatmulSwigluQuantWeightNzTensorList + OPTIONS + --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE grouped_matmul_swiglu_quant_weight_nz_tensor_list ACLNNTYPE aclnn_exclude) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list.h new file mode 100644 index 000000000..bfb3ab2f1 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list.h @@ -0,0 +1,25 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_API_INC_LEVEL0_OP_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_OP_H +#define OP_API_INC_LEVEL0_OP_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_OP_H + +#include "opdev/op_executor.h" + +namespace l0op { +const std::tuple GroupedMatmulSwigluQuantWeightNzTensorList(const aclTensor *x, + const aclTensorList *weight, + const aclTensorList *perChannelScale, + const aclTensor *perTokenScale, + const aclTensor *groupList, + float swigluLimit, + aclOpExecutor *executor); +} + +#endif diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_def.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_def.cpp new file mode 100644 index 000000000..b028e1718 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_def.cpp @@ -0,0 +1,66 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_weight_nz_tensor_list_def.cpp + * \brief + */ + +#include +#include "register/op_def_registry.h" +namespace ops { +class GroupedMatmulSwigluQuantWeightNzTensorList : public OpDef { +public: + explicit GroupedMatmulSwigluQuantWeightNzTensorList(const char* name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8,ge::DT_INT8,ge::DT_INT8}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("weight") + .ParamType(DYNAMIC) + .DataType({ge::DT_INT8,ge::DT_INT8,ge::DT_INT8}) + .Format({ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ}); + this->Input("weight_scale") + .ParamType(DYNAMIC) + .DataType({ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("x_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT,ge::DT_FLOAT,ge::DT_FLOAT}) + .Format({ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND}); + this->Input("group_list") + .ParamType(REQUIRED) + .DataType({ge::DT_INT64,ge::DT_INT64,ge::DT_INT64}) + .Format({ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND}); + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8,ge::DT_INT8,ge::DT_INT8}) + .Format({ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND}); + this->Output("y_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT,ge::DT_FLOAT,ge::DT_FLOAT}) + .Format({ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND}); + this->Attr("swiglu_limit").AttrType(REQUIRED).Float(0.0f); + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true); + + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + } +}; + +OP_ADD(GroupedMatmulSwigluQuantWeightNzTensorList); +} diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_infershape.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_infershape.cpp new file mode 100644 index 000000000..eff3cad37 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_infershape.cpp @@ -0,0 +1,51 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_weight_nz_tensor_list_infershape.cpp + * \brief + */ +#include "register/op_impl_registry.h" +#include "log/ops_log.h" +#include "platform/platform_info.h" + +using namespace ge; + +namespace ops { +const int64_t X_INDEX = 0; +const int64_t WEIGHTSCALE_INDEX = 2; +const int64_t M_DIM_INDEX = 0; +const int64_t N_DIM_INDEX = 0; + +static ge::graphStatus InferShape4GroupedMatmulSwigluQuantWeightNzTensorList(gert::InferShapeContext* context) { + const gert::Shape* xShape = context->GetInputShape(X_INDEX); + const gert::Shape* weightScaleShape = context->GetDynamicInputShape(WEIGHTSCALE_INDEX, 0); + int64_t m = xShape->GetDim(M_DIM_INDEX); + int64_t n = static_cast(weightScaleShape->GetDim(N_DIM_INDEX) / 2); + auto outShape = context->GetOutputShape(0); + outShape->SetDimNum(2); + outShape->SetDim(0, m); + outShape->SetDim(1, n); + auto outScaleShape = context->GetOutputShape(1); + outScaleShape->SetDimNum(1); + outScaleShape->SetDim(0, m); + return GRAPH_SUCCESS; +} + +static graphStatus InferDataType4GroupedMatmulSwigluQuantWeightNzTensorList(gert::InferDataTypeContext* context) { + context->SetOutputDataType(0, DataType::DT_INT8); + context->SetOutputDataType(1, DataType::DT_FLOAT); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(GroupedMatmulSwigluQuantWeightNzTensorList) + .InferShape(InferShape4GroupedMatmulSwigluQuantWeightNzTensorList) + .InferDataType(InferDataType4GroupedMatmulSwigluQuantWeightNzTensorList); +} // namespace ops diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_proto.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_proto.cpp new file mode 100644 index 000000000..5e3d44320 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_proto.cpp @@ -0,0 +1,49 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_weight_nz_tensor_list_proto.cpp + * \brief + */ +#include "register/op_impl_registry.h" +#include "log/ops_log.h" +#include "platform/platform_info.h" + +using namespace ge; +namespace ops { +const int64_t X_INDEX = 0; +const int64_t WEIGHTSCALE_INDEX = 2; +const int64_t M_DIM_INDEX = 0; +const int64_t N_DIM_INDEX = 0; +static ge::graphStatus InferShape4GroupedMatmulSwigluQuantWeightNzTensorList(gert::InferShapeContext* context) { + const gert::Shape* xShape = context->GetInputShape(X_INDEX); + const gert::Shape* weightScaleShape = context->GetDynamicInputShape(WEIGHTSCALE_INDEX, 0); + int64_t m = xShape->GetDim(M_DIM_INDEX); + int64_t n = static_cast(weightScaleShape->GetDim(N_DIM_INDEX) / 2); + auto outShape = context->GetOutputShape(0); + outShape->SetDimNum(2); + outShape->SetDim(0, m); + outShape->SetDim(1, n); + auto outScaleShape = context->GetOutputShape(1); + outScaleShape->SetDimNum(1); + outScaleShape->SetDim(0, m); + return GRAPH_SUCCESS; +} + +static graphStatus InferDataType4GroupedMatmulSwigluQuantWeightNzTensorList(gert::InferDataTypeContext* context) { + context->SetOutputDataType(0, DataType::DT_INT8); + context->SetOutputDataType(1, DataType::DT_FLOAT); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(GroupedMatmulSwigluQuantWeightNzTensorList) + .InferShape(InferShape4GroupedMatmulSwigluQuantWeightNzTensorList) + .InferDataType(InferDataType4GroupedMatmulSwigluQuantWeightNzTensorList); +} // namespace ops diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_tiling.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_tiling.cpp new file mode 100644 index 000000000..496822c83 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_tiling.cpp @@ -0,0 +1,196 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_weight_nz_tensor_list_tiling.cpp + * \brief + */ +#include +#include +#include "register/op_impl_registry.h" +#include "log/ops_log.h" +#include "error/ops_error.h" +#include "tiling/tiling_base.h" +#include "grouped_matmul_swiglu_quant_weight_nz_tensor_list_tiling.h" +using namespace ge; +using namespace AscendC; +using namespace GroupedMatmulSwigluQuantWeightNzTensorListTiling; + +template +static T1 CeilDiv(T1 a, T2 b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +namespace optiling { + +struct GMMSwigluCompileInfo { + uint64_t ubSize_ = 0; + uint32_t aicNum_ = 0; + uint32_t baseM_ = 128; + uint32_t baseN_ = 256; +}; + +static uint64_t CalcMaxTmpSize(const uint32_t row, const uint64_t n) { + std::vector shape_vec = {static_cast(row * n)}; + Shape shape(shape_vec); + uint32_t max; + uint32_t min; + GetSwiGLUMaxMinTmpSize(shape, 4, max, min, false); + uint32_t averageTmp = (max + min) >> 1; + GetAscendQuantMaxMinTmpSize(shape, 4, max, min); + uint32_t average = (max + min) >> 1; + average = average > averageTmp ? average : averageTmp; + GetAscendDequantMaxMinTmpSize(shape, 4, max, min); + averageTmp = (max + min) >> 1; + return average > averageTmp ? average : averageTmp; +} + +static uint64_t CalRows(const uint64_t ubSize, const uint64_t n) { + uint64_t tokenSize = n << 2; + uint64_t expectSize = ubSize - tokenSize; + uint64_t rows = expectSize / (8 + tokenSize); + uint64_t realSize = (8 + tokenSize) * rows + CalcMaxTmpSize(rows, n); + while (expectSize < realSize) { + rows -= CeilDiv(realSize - expectSize, (8 + tokenSize) << 2); + realSize = (8 + tokenSize) * rows + CalcMaxTmpSize(rows, n); + } + return rows; +} + +static void SetTilingKey(gert::TilingContext* context, bool isSplitWorkSpace) { + if(isSplitWorkSpace){ + context->SetTilingKey(1); + context->SetScheduleMode(BATCH_MODE_SCHEDULE); + } else { + context->SetTilingKey(0); + context->SetScheduleMode(BATCH_MODE_SCHEDULE); + } +} + +static bool IsPreFill(GMMSwigluQuantTensorListTilingData &tilingData) { + int64_t k = tilingData.gmmSwigluTensorListBaseParams.get_K(); + int64_t n = tilingData.gmmSwigluTensorListBaseParams.get_N(); + int64_t m = tilingData.gmmSwigluTensorListBaseParams.get_M(); + int64_t groupNum = tilingData.gmmSwigluTensorListBaseParams.get_groupNum(); + if (groupNum == 128 && m >= PREFILL_M_MIN_SIZE) { // 128:prefiling groupNum + std::array kNList = {k, n}; // 2: kNList size + if (PREFILL_WHITE_LIST.count(kNList)) { + return true; + } + } + return false; +} + +ASCENDC_EXTERN_C graphStatus TilingGMMSwigluQuantTensorList(gert::TilingContext* context) { + // set info + OPS_LOG_I(context->GetNodeName(), "Begin Run GMM Swiglu Tiling ."); + + auto compileInfoPtr = context->GetCompileInfo(); + auto xTensor = context->GetInputTensor(X_INDEX); + OPS_LOG_E_IF_NULL(context, xTensor, return GRAPH_FAILED); + const int64_t m = xTensor->GetStorageShape().GetDim(0); + const int64_t k = xTensor->GetStorageShape().GetDim(1); + auto wTensor = context->GetDynamicInputTensor(WEIGHT_INDEX, 0); + OPS_LOG_E_IF_NULL(context, wTensor, return GRAPH_FAILED); + const int64_t n = wTensor->GetStorageShape().GetDim(0) * wTensor->GetStorageShape().GetDim(3); + auto groupListTensor = context->GetDynamicInputTensor(GROUPLIST_INDEX, 0); + OPS_LOG_E_IF_NULL(context, groupListTensor, return GRAPH_FAILED); + const int64_t groupNum = groupListTensor->GetStorageShape().GetDim(0); + GMMSwigluQuantTensorListTilingData tilingData; + const int64_t row = CalRows(compileInfoPtr->ubSize_, n); + auto attrs = context->GetAttrs(); + float swiglu_limit = *attrs->GetFloat(0); + OPS_ERR_IF(!(swiglu_limit >= 0.0f), + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "swiglu_limit should be non-negative"), + return GRAPH_FAILED); + + tilingData.gmmSwigluTensorListBaseParams.set_groupNum(groupNum); + tilingData.gmmSwigluTensorListBaseParams.set_coreNum(compileInfoPtr->aicNum_); + tilingData.gmmSwigluTensorListBaseParams.set_K(k); + tilingData.gmmSwigluTensorListBaseParams.set_N(n); + tilingData.gmmSwigluTensorListBaseParams.set_M(m); + tilingData.gmmSwigluTensorList.set_maxProcessRowNum(row); + tilingData.gmmSwigluTensorList.set_groupListLen(groupNum); + tilingData.gmmSwigluTensorList.set_tokenLen(n); + tilingData.gmmSwigluTensorList.set_swigluLimit(swiglu_limit); + + OPS_LOG_D(context->GetNodeName(),"grouped_matmul_swiglu_quant_weight_nz_tensor_list_tiling."); + OPS_LOG_D(context->GetNodeName(),"gmmSwigluTensorListBaseParams.groupNum: %ld", groupNum); + OPS_LOG_D(context->GetNodeName(),"gmmSwigluTensorListBaseParams.coreNum: %u ", compileInfoPtr->aicNum_); + OPS_LOG_D(context->GetNodeName(),"gmmSwigluTensorListBaseParams.M: %ld", m); + OPS_LOG_D(context->GetNodeName(),"gmmSwigluTensorListBaseParams.K: %ld", k); + OPS_LOG_D(context->GetNodeName(),"gmmSwigluTensorListBaseParams.N: %ld", n); + OPS_LOG_D(context->GetNodeName(),"gmmSwigluTensorList.maxProcessRowNum: %ld", row); + OPS_LOG_D(context->GetNodeName(),"gmmSwigluTensorList.groupListLen: %ld", groupNum); + OPS_LOG_D(context->GetNodeName(),"gmmSwigluTensorList.tokenLen: %ld", n); + OPS_LOG_D(context->GetNodeName(),"gmmSwigluTensorList.swigluLimit: %f", swiglu_limit); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + using namespace matmul_tiling; + MatmulApiTiling tiling(ascendcPlatform); + tiling.SetAType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_INT8); + tiling.SetBType(TPosition::GM, CubeFormat::NZ, matmul_tiling::DataType::DT_INT8); + tiling.SetCType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_INT32); + tiling.SetBias(false); + tiling.SetShape(compileInfoPtr->baseM_, compileInfoPtr->baseN_, k); + tiling.SetOrgShape(m, n, k); + tiling.SetBufferSpace(-1, -1, -1); + OPS_ERR_IF(tiling.GetTiling(tilingData.mmTilingData) == -1, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "grouped_matmul_swiglu_quant_weight_nz_tensor_list_tiling, get tiling failed"), + return GRAPH_FAILED); + auto workspaceSizes = context->GetWorkspaceSizes(1); + bool isPreFill = IsPreFill(tilingData); + tilingData.gmmSwigluTensorListBaseParams.set_isPreFill(isPreFill); + int64_t usrWorkspaceLimut = isPreFill ? PREFILL_USER_WORKSPACE_LIMIT : USER_WORKSPACE_LIMIT; + int64_t mLimit = ((usrWorkspaceLimut / DOUBLE_WORKSPACE_SPLIT) / INT32_DTYPE_SIZE) / n; + OPS_ERR_IF(mLimit <= 0, + OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(),"mLimit is %ld must over then 0.", mLimit), + return GRAPH_FAILED); + tilingData.gmmSwigluTensorListBaseParams.set_mLimit(mLimit); + workspaceSizes[0] = SYS_WORKSPACE_SIZE + ((mLimit * DOUBLE_WORKSPACE_SPLIT > m \ + ? m \ + : mLimit * DOUBLE_WORKSPACE_SPLIT) * n * sizeof(int32_t)); + bool isSplitWorkSpace = m > mLimit * DOUBLE_WORKSPACE_SPLIT; + OPS_LOG_D(context->GetNodeName(), "USER_WORKSPACE_LIMIT: %ld", usrWorkspaceLimut); + OPS_LOG_D(context->GetNodeName(), "mLimit: %ld", mLimit); + OPS_LOG_D(context->GetNodeName(), "workspaceSizes: %lu", workspaceSizes[0]); + OPS_LOG_D(context->GetNodeName(), "isSplitWorkSpace: %s", isSplitWorkSpace ? "true" : "false"); + OPS_LOG_D(context->GetNodeName(), "isPreFill: %s", isPreFill ? "true" : "false"); + SetTilingKey(context, isSplitWorkSpace); + tilingData.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); + context->SetBlockDim(compileInfoPtr->aicNum_); // block dim is the number of aicube + context->GetRawTilingData()->SetDataSize(tilingData.GetDataSize()); + + OPS_LOG_D(context->GetNodeName(), "End Run GMM Swiglu Tiling."); + return GRAPH_SUCCESS; +} + +ASCENDC_EXTERN_C graphStatus TilingPrepareForGMMSwigluQuantTensorList(gert::TilingParseContext* context) { + // get info + fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo(); + OPS_LOG_E_IF_NULL(context, platformInfoPtr, return GRAPH_FAILED); + auto compileInfoPtr = context->GetCompiledInfo(); + OPS_LOG_E_IF_NULL(context, compileInfoPtr, return GRAPH_FAILED); + + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + compileInfoPtr->aicNum_ = ascendcPlatform.GetCoreNumAic(); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, compileInfoPtr->ubSize_); + OPS_LOG_D(context->GetNodeName(), "ubSize is %lu, aicNum is %u.", compileInfoPtr->ubSize_, compileInfoPtr->aicNum_); + return GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(GroupedMatmulSwigluQuantWeightNzTensorList) +.Tiling(TilingGMMSwigluQuantTensorList) +.TilingParse(TilingPrepareForGMMSwigluQuantTensorList); +} // namespace optiling diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_tiling.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_tiling.h new file mode 100644 index 000000000..7331c679e --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/grouped_matmul_swiglu_quant_weight_nz_tensor_list_tiling.h @@ -0,0 +1,69 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_weight_nz_tensor_list_tiling.h + * \brief + */ +#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_H +#define AIR_CXX_RUNTIME_V2_OP_IMPL_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_H + +#include +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(GMMSwigluTensorListBaseParams) + TILING_DATA_FIELD_DEF(uint32_t, groupNum); + TILING_DATA_FIELD_DEF(uint32_t, coreNum); + TILING_DATA_FIELD_DEF(uint32_t, K); + TILING_DATA_FIELD_DEF(uint32_t, N); + TILING_DATA_FIELD_DEF(uint32_t, M); + TILING_DATA_FIELD_DEF(uint32_t, mLimit); + TILING_DATA_FIELD_DEF(uint64_t, isPreFill); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(GMMSwigluTensorListBaseParamsOp, GMMSwigluTensorListBaseParams) + +BEGIN_TILING_DATA_DEF(GMMSwigluTensorList) + TILING_DATA_FIELD_DEF(uint32_t, maxProcessRowNum); + TILING_DATA_FIELD_DEF(uint32_t, groupListLen); + TILING_DATA_FIELD_DEF(uint32_t, tokenLen); + TILING_DATA_FIELD_DEF(float, swigluLimit); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(GMMSwigluTensorListOp, GMMSwigluTensorList) + +BEGIN_TILING_DATA_DEF(GMMSwigluQuantTensorListTilingData) + TILING_DATA_FIELD_DEF_STRUCT(GMMSwigluTensorListBaseParams, gmmSwigluTensorListBaseParams); + TILING_DATA_FIELD_DEF_STRUCT(GMMSwigluTensorList, gmmSwigluTensorList); + TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, mmTilingData); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(GroupedMatmulSwigluQuantWeightNzTensorList, GMMSwigluQuantTensorListTilingData) +} + +namespace GroupedMatmulSwigluQuantWeightNzTensorListTiling { +constexpr uint32_t X_INDEX = 0; +constexpr uint32_t WEIGHT_INDEX = 1; +constexpr uint32_t GROUPLIST_INDEX = 4; +constexpr uint32_t BATCH_MODE_SCHEDULE = 1; +constexpr uint32_t SYS_WORKSPACE_SIZE = 16 * 1024 * 1024; +constexpr int64_t USER_WORKSPACE_LIMIT = 256 * 1024 * 1024; +constexpr int64_t PREFILL_USER_WORKSPACE_LIMIT = 64 * 1024 * 1024; +constexpr int64_t DOUBLE_WORKSPACE_SPLIT = 2; +constexpr int64_t INT32_DTYPE_SIZE = 4; +constexpr uint32_t PREFILL_M_MIN_SIZE = 16 * 1024; + +const std::set> PREFILL_WHITE_LIST = { // used for preFill case + {{2048, 1536}}, + {{4096, 3072}} +}; +} // namespace GroupedMatmulSwigluQuantWeightNzTensorListTiling + +#endif // AIR_CXX_RUNTIME_V2_OP_IMPL_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_H diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_weight_nz_tensor_list.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_weight_nz_tensor_list.cpp new file mode 100644 index 000000000..3c7a9bf13 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_weight_nz_tensor_list.cpp @@ -0,0 +1,329 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#include +#include +#include "aclnn_kernels/contiguous.h" +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "opdev/common_types.h" +#include "opdev/data_type_utils.h" +#include "opdev/format_utils.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/platform.h" +#include "opdev/shape_utils.h" +#include "opdev/tensor_view_utils.h" +#include "opdev/make_op_executor.h" +#include "../grouped_matmul_swiglu_quant_weight_nz_tensor_list.h" +#include "aclnn_grouped_matmul_swiglu_quant_weight_nz_tensor_list.h" + +using namespace op; + +#ifdef __cplusplus +extern "C" { +#endif + +static constexpr int64_t SPLIT = 2; +static constexpr int64_t K_LIMIT = 65536; +static constexpr int64_t N_LIMIT = 4096; +static constexpr int64_t NZ_DIM_3 = 32; +static constexpr int64_t NZ_DIM_2 = 16; +static constexpr int64_t OUTPUT_IDX_0 = 0; +static constexpr int64_t OUTPUT_IDX_1 = 1; +static constexpr size_t X_DIM_LIMIT = 2; +static constexpr size_t WEIGHT_ND_DIM_LIMIT = 2; +static constexpr size_t WEIGHT_NZ_DIM_LIMIT = 4; +static constexpr size_t WEIGHT_SCALE_DIM_LIMIT = 1; +static constexpr size_t TOKEN_SCALE_DIM_LIMIT = 1; +static constexpr size_t GROUP_LIST_DIM_LIMIT = 1; +static constexpr size_t QUANTOUT_DIM_LIMIT = 2; +static constexpr size_t QUANTSCALEOUT_DIM_LIMIT = 1; + +static const std::initializer_list X_DTYPE_SUPPORT_LIST = {DataType::DT_INT8}; +static const std::initializer_list WEIGHT_DTYPE_SUPPORT_LIST = {DataType::DT_INT8}; +static const std::initializer_list WEIGHT_SCALE_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT, DataType::DT_FLOAT16, DataType::DT_BF16}; +static const std::initializer_list X_SCALE_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT, DataType::DT_FLOAT16, DataType::DT_BF16}; +static const std::initializer_list GROUP_LIST_DTYPE_SUPPORT_LIST = {DataType::DT_INT64}; +static const std::initializer_list QUANTOUT_DTYPE_SUPPORT_LIST = {DataType::DT_INT8}; +static const std::initializer_list QUANTSCALEOUT_DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT}; + +static bool CheckNotNull(const aclTensor* x, const aclTensorList* weight, const aclTensor* bias, const aclTensor* offset, + const aclTensorList* weightScale, const aclTensor* xScale, const aclTensor* groupList, + const aclTensor* output, const aclTensor* outputScale, const aclTensor* outputOffset) +{ + OP_CHECK_NULL(x, return false); + OP_CHECK_NULL(weight, return false); + OP_CHECK_NULL(weightScale, return false); + OP_CHECK_NULL(xScale, return false); + OP_CHECK_NULL(groupList, return false); + OP_CHECK_NULL(output, return false); + OP_CHECK_NULL(outputScale, return false); + if (bias != nullptr) { + OP_LOGW("aclnnGroupedMatmulSwigluQuantWeightNzTensorList, The current version does not support the scenario where bias is not 0. " + "Features and accuracy are not guaranteed if inputting bias with values other than 0s."); + } + if (offset != nullptr) { + OP_LOGW("aclnnGroupedMatmulSwigluQuantWeightNzTensorList, The current version does not support the scenario where offset is not 0. " + "Features and accuracy are not guaranteed if inputting bias with values other than 0s."); + } + if (outputOffset != nullptr) { + OP_LOGW("aclnnGroupedMatmulSwigluQuantWeightNzTensorList, The current version does not support the scenario where outputOffset is not 0. " + "Features and accuracy are not guaranteed if inputting bias with values other than 0s."); + } + return true; +} + +static bool CheckInputOutDims(const aclTensor* x, const aclTensorList* weight, const aclTensorList* weightScale, + const aclTensor* xScale, const aclTensor* groupList, + const aclTensor* output, const aclTensor* outputScale) +{ + OP_CHECK_WRONG_DIMENSION(x, X_DIM_LIMIT, return false); + op::Format weightViewFormat = (*weight)[0]->GetViewFormat(); + if (IsPrivateFormat(weightViewFormat)){ + OP_CHECK_WRONG_DIMENSION((*weight)[0], WEIGHT_NZ_DIM_LIMIT, return false); + } else { + OP_CHECK_WRONG_DIMENSION((*weight)[0], WEIGHT_ND_DIM_LIMIT, return false); + } + OP_CHECK_WRONG_DIMENSION((*weightScale)[0], WEIGHT_SCALE_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(xScale, TOKEN_SCALE_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(groupList, GROUP_LIST_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(output, QUANTOUT_DIM_LIMIT, return false); + OP_CHECK_WRONG_DIMENSION(outputScale, QUANTSCALEOUT_DIM_LIMIT, return false); + return true; +} + +static bool CheckInputOutShape(const aclTensor* x, const aclTensorList* weight, const aclTensorList* weightScale, + const aclTensor* xScale, const aclTensor* groupList, + const aclTensor* output, const aclTensor* outputScale) +{ + int64_t m = x->GetViewShape().GetDim(0); + int64_t k = x->GetViewShape().GetDim(1); + int64_t n = (*weightScale)[0]->GetViewShape().GetDim(0); + int64_t e = weight->Size(); + if (n % SPLIT != 0){ + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwigluQuantWeightNzTensorList, N is %ld , not an even number.", n); + return false; + } + int64_t nAfterHalve = static_cast(n / SPLIT); + // x shape is expected to be [M, K] + op::Shape xExpectShape = {m, k}; + // The ND shape of each weight in TensorList is expected to be [K, N] + op::Shape weightNDExpectShape = {k, n}; + // The NZ shape of each weight in TensorList is expected to be [N // 32, K // 16, 16, 32] + op::Shape weightNZExpectShape = {static_cast(n / NZ_DIM_3), + static_cast(k / NZ_DIM_2), + NZ_DIM_2, NZ_DIM_3}; + // weightScale shape is expected to be [N] + op::Shape weightScaleExpectShape = {n}; + // xScale shape is expected to be [E, N] + op::Shape xScaleExpectShape = {m}; + // output shape is expected to be [M, N] + op::Shape outputExpectShape = {m, nAfterHalve}; + // outputScale shape is expected to be [M] + op::Shape outputScaleExpectShape = {m}; + for (size_t i = 0; i < weight->Size(); ++i) { + op::Format weightViewFormat = (*weight)[i]->GetViewFormat(); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(x, xExpectShape, return false); + if (IsPrivateFormat(weightViewFormat)){ + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE((*weight)[i], weightNZExpectShape, return false); + } else { + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE((*weight)[i], weightNDExpectShape, return false); + } + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE((*weightScale)[i], weightScaleExpectShape, return false); + } + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(xScale, xScaleExpectShape, return false); + + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(output, outputExpectShape, return false); + OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(outputScale, outputScaleExpectShape, return false); + // The length of groupList should be less than or equal to the number of experts in weight + int64_t groupListLen = groupList->GetViewShape().GetDim(0); + if(groupListLen > e) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwigluQuantWeightNzTensorList, Length of 'groupList' out of range (expected to be in range of [1, %ld], but got %ld)", + e, groupListLen); + return false; + } + if(nAfterHalve > N_LIMIT) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwigluQuantWeightNzTensorList, The current version does not support the scenario.\ + where N after halve is %ld greater than %ld.", + nAfterHalve, N_LIMIT); + return false; + } + if(k >= K_LIMIT) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwigluQuantWeightNzTensorList, The current version does not support the scenario.\ + The tail axis dimension of input0(x) is %ld, which need lower than %ld.", + k, K_LIMIT); + return false; + } + return true; +} + +static bool CheckDtypeValid(const aclTensor* x, const aclTensorList* weight, const aclTensorList* weightScale, + const aclTensor* xScale, const aclTensor* groupList, + const aclTensor* output, const aclTensor* outputScale) +{ + OP_CHECK_DTYPE_NOT_SUPPORT(x, X_DTYPE_SUPPORT_LIST, return false); + for (size_t i = 0; i < weight->Size(); ++i) { + OP_CHECK_DTYPE_NOT_SUPPORT((*weight)[i], WEIGHT_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT((*weightScale)[i], WEIGHT_SCALE_DTYPE_SUPPORT_LIST, return false); + } + OP_CHECK_DTYPE_NOT_SUPPORT(xScale, X_SCALE_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(groupList, GROUP_LIST_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(output, QUANTOUT_DTYPE_SUPPORT_LIST, return false); + OP_CHECK_DTYPE_NOT_SUPPORT(outputScale, QUANTSCALEOUT_DTYPE_SUPPORT_LIST, return false); + return true; +} + +static bool CheckFormat(const aclTensor* x, const aclTensorList* weight, const aclTensor* output) +{ + bool isNZ = (*weight)[0]->GetStorageFormat() == op::Format::FORMAT_FRACTAL_NZ; + if (!isNZ) { + // fp16 in fp32 out that is split k template, not precision-advanced now + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "aclnnGroupedMatmulSwigluQuantWeightNzTensorList, The current version does not support the scenario.\ + weight Format expect is FRACTAL_NZ, but got [%s].", op::ToString((*weight)[0]->GetStorageFormat()).GetString()); + return false; + } + if (IsPrivateFormat(x->GetStorageFormat())) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "aclnnGroupedMatmulSwigluQuantWeightNzTensorList, The current version does not support the scenario.\ + x Format Not support Private Format."); + return false; + } + if (IsPrivateFormat(output->GetStorageFormat())) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "aclnnGroupedMatmulSwigluQuantWeightNzTensorList, The current version does not support the scenario.\ + output Format Not support Private Format."); + return false; + } + return true; +} + +static aclnnStatus CheckParams(const aclTensor* x, const aclTensorList* weight, const aclTensor* bias, const aclTensor* offset, + const aclTensorList* weightScale, const aclTensor* xScale, const aclTensor* groupList, + const aclTensor* output, const aclTensor* outputScale, const aclTensor* outputOffset) { + // 1. Check if parameters are null pointers + CHECK_RET(CheckNotNull(x, weight, bias, offset, weightScale, xScale, + groupList, output, outputScale, outputOffset), ACLNN_ERR_PARAM_NULLPTR); + + // 2. Verify input and output parameter dimensions + CHECK_RET(CheckInputOutDims(x, weight, weightScale, xScale, + groupList, output, outputScale), ACLNN_ERR_PARAM_INVALID); + + // 3. Verify input and output shape parameters + CHECK_RET(CheckInputOutShape(x, weight, weightScale, xScale, + groupList, output, outputScale), ACLNN_ERR_PARAM_INVALID); + + // 4. Check if the input data types are within the supported data type range + CHECK_RET(CheckDtypeValid(x, weight, weightScale, xScale, + groupList, output, outputScale), ACLNN_ERR_PARAM_INVALID); + + // 5. Check if data format is supported + CHECK_RET(CheckFormat(x, weight, output), ACLNN_ERR_PARAM_INVALID); + + return ACLNN_SUCCESS; +} + +static aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNzTensorListGetWorkspaceSizeCommon(const aclTensor *x, const aclTensorList *weight, + const aclTensor *bias, const aclTensor *offset, + const aclTensorList *weightScale, const aclTensor *xScale, + const aclTensor *groupList, float swigluLimit, + aclTensor *output, aclTensor *outputScale, + aclTensor *outputOffset, uint64_t *workspaceSize, + aclOpExecutor **executor){ + // Fixed pattern, create OpExecutor + auto uniqueExecutor = CREATE_EXECUTOR(); + CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + // Fixed pattern, parameter check + auto ret = CheckParams(x, weight, bias, offset, weightScale, xScale, + groupList, output, outputScale, outputOffset); + CHECK_RET(ret == ACLNN_SUCCESS, ret); + // Empty tensor scenario + if (output->IsEmpty() || groupList->IsEmpty() || outputScale->IsEmpty()) { + *workspaceSize = 0; + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; + } + // Convert to contiguous + x = l0op::Contiguous(x, uniqueExecutor.get()); + CHECK_RET(x != nullptr, ACLNN_ERR_INNER_NULLPTR); + for (size_t i = 0; i < weight->Size(); ++i) { + (*weight)[i]->SetOriginalShape((*weight)[i]->GetViewShape()); + } + xScale = l0op::Contiguous(xScale, uniqueExecutor.get()); + CHECK_RET(xScale != nullptr, ACLNN_ERR_INNER_NULLPTR); + groupList = l0op::Contiguous(groupList, uniqueExecutor.get()); + CHECK_RET(groupList != nullptr, ACLNN_ERR_INNER_NULLPTR); + // Call L0 operator capability + auto ret_0 = l0op::GroupedMatmulSwigluQuantWeightNzTensorList(x, weight, weightScale, xScale, groupList, swigluLimit, uniqueExecutor.get()); + CHECK_RET(ret_0 != std::tuple(nullptr, nullptr), ACLNN_ERR_INNER_NULLPTR); + auto out0 = std::get(ret_0); + auto ret_1 = l0op::ViewCopy(out0, output, uniqueExecutor.get()); + CHECK_RET(ret_1 != nullptr, ACLNN_ERR_INNER_NULLPTR); + auto out1 = std::get(ret_0); + auto ret_2 = l0op::ViewCopy(out1, outputScale, uniqueExecutor.get()); + CHECK_RET(ret_2 != nullptr, ACLNN_ERR_INNER_NULLPTR); + *workspaceSize = uniqueExecutor->GetWorkspaceSize(); + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; +} + +aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNzTensorListGetWorkspaceSize(const aclTensor *x, const aclTensorList *weight, + const aclTensor *bias, const aclTensor *offset, + const aclTensorList *weightScale, const aclTensor *xScale, + const aclTensor *groupList, float swigluLimit, + aclTensor *output, aclTensor *outputScale, + aclTensor *outputOffset, uint64_t *workspaceSize, + aclOpExecutor **executor) { + OP_CHECK_COMM_INPUT(workspaceSize, executor); + L2_DFX_PHASE_1(aclnnGroupedMatmulSwigluQuantWeightNzTensorList, + DFX_IN(x, weight, bias, offset, weightScale, xScale, groupList,swigluLimit), + DFX_OUT(output, outputScale, outputOffset)); + // weight is forcibly bound to StorageFormat and ViewFormat as NZ in this scenario + CHECK_RET(weight != nullptr, ACLNN_ERR_PARAM_NULLPTR); + for (size_t i = 0; i < weight->Size(); ++i) { + auto storgeShape = (*weight)[i]->GetStorageShape(); + auto viewShape = (*weight)[i]->GetViewShape(); + aclTensor* weightNZ = const_cast((*weight)[i]); + CHECK_COND((storgeShape.GetDimNum() == WEIGHT_NZ_DIM_LIMIT), + ACLNN_ERR_PARAM_INVALID, + "aclnnGroupedMatmulSwigluQuantWeightNZTensorList, The dimnum of storageShape for second input (weight) \ + must be 4. \n But StorageShape got %s , and dimNum is %lu.", + op::ToString(storgeShape).GetString(), storgeShape.GetDimNum()); + // The StorageFormat of weight is unconditionally regarded as NZ + weightNZ->SetStorageFormat(op::Format::FORMAT_FRACTAL_NZ); + if (viewShape.GetDimNum() == WEIGHT_NZ_DIM_LIMIT){ + // If the viewShape of weight is 4-dimensional, it is regarded as NZ + weightNZ->SetViewFormat(op::Format::FORMAT_FRACTAL_NZ); + } else if (viewShape.GetDimNum() == WEIGHT_ND_DIM_LIMIT){ + // If the viewShape of weight is 2-dimensional, it is regarded as ND + weightNZ->SetViewFormat(op::Format::FORMAT_ND); + } + } + // Call the common interface + return aclnnGroupedMatmulSwigluQuantWeightNzTensorListGetWorkspaceSizeCommon(x, weight, bias, offset, weightScale, xScale, groupList, swigluLimit, + output, outputScale, outputOffset, workspaceSize, executor); +} + +aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNzTensorList(void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream) { + L2_DFX_PHASE_2(aclnnGroupedMatmulSwigluQuantWeightNzTensorList); + CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER, + "This is an error in GroupedMatmulSwigluQuantWeightNzTensorList launch aicore"); + return ACLNN_SUCCESS; +} + +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_weight_nz_tensor_list.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_weight_nz_tensor_list.h new file mode 100644 index 000000000..30980428c --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/op_api/aclnn_grouped_matmul_swiglu_quant_weight_nz_tensor_list.h @@ -0,0 +1,57 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_API_INC_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_H +#define OP_API_INC_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_H +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @brief The first interface of aclnnGroupedMatmulSwigluQuantWeightNzTensorList, which calculates the workspace size according to the specific calculation process. + * @domain aclnn_ops_infer + * + * @param [in] x: Represents x in the formula. The data type supports INT8, and the data format supports ND. + * @param [in] weight: + * Represents weight in the formula. The data type supports INT8, and the data format supports NZ. + * @param [in] weightScale: Represents quantization parameters. The data type supports FLOAT16, BFLOAT16, and FLOAT32. The data format supports ND, with a maximum length of 128. + * Represents per Channel parameters. The data type supports FLOAT16 and BFLOAT16. The data format supports ND. + * @param [in] xScale: + * Represents per Token quantization parameters. The data type supports FLOAT32, and the data format supports ND. + * @param [in] groupList: Required parameter, representing the index situation on the input and output grouping axes. The data type supports INT64. + * @param [in] swigluLimit: Required parameter. + * @param [out] quantOutput: Represents out in the formula. The data type supports INT8, and the data format supports ND. + * @param [out] quantScaleOutput: Represents outQuantScale in the formula. The data type supports Float32. + * @param [out] workspaceSize: Returns the workspace size that users need to apply for on the npu device side. + * @param [out] executor: Returns the op executor, containing the operator calculation process. + * @return aclnnStatus: Returns the status code. + */ +__attribute__((visibility("default"))) aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNzTensorListGetWorkspaceSize( + const aclTensor *x, const aclTensorList *weight, const aclTensor *bias, const aclTensor *offset, + const aclTensorList *weightScale, const aclTensor *xScale, const aclTensor *groupList, float swigluLimit, + aclTensor *output, aclTensor *outputScale, aclTensor *outputOffset, uint64_t *workspaceSize, aclOpExecutor **executor); + +/** + * @brief The second interface of aclnnGroupedMatmulSwigluQuantWeightNzTensorList, used to execute calculations. + * @param [in] workspace: The starting address of the workspace memory applied for on the npu device side. + * @param [in] workspaceSize: The workspace size applied for on the npu device side, obtained from the first interface aclnnGroupedMatmulSwigluQuantWeightNzTensorListGetWorkspaceSize. + * @param [in] stream: acl stream. + * @param [in] executor: op executor, containing the operator calculation process. + * @return aclnnStatus: Returns the status code. + */ +__attribute__((visibility("default"))) aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNzTensorList(void* workspace, + uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/op_api/grouped_matmul_swiglu_quant_weight_nz_tensor_list.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/op_api/grouped_matmul_swiglu_quant_weight_nz_tensor_list.cpp new file mode 100644 index 000000000..6cc3595d1 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_host/op_api/grouped_matmul_swiglu_quant_weight_nz_tensor_list.cpp @@ -0,0 +1,59 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "opdev/op_log.h" +#include "opdev/op_dfx.h" +#include "opdev/make_op_executor.h" +#include "../grouped_matmul_swiglu_quant_weight_nz_tensor_list.h" + +using namespace op; + +namespace l0op { +OP_TYPE_REGISTER(GroupedMatmulSwigluQuantWeightNzTensorList); + +const std::tuple GroupedMatmulSwigluQuantWeightNzTensorList(const aclTensor *x, + const aclTensorList *weight, + const aclTensorList *perChannelScale, + const aclTensor *perTokenScale, + const aclTensor *groupList, + float swigluLimit, + aclOpExecutor *executor) { + L0_DFX(GroupedMatmulSwigluQuantWeightNzTensorList, x, weight, perChannelScale, perTokenScale, groupList); + if (x == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "x is nullptr."); + return std::tuple(nullptr, nullptr); + } + int64_t m = perTokenScale->GetViewShape().GetDim(0); + int64_t n = (*perChannelScale)[0]->GetViewShape().GetDim(0); + int64_t nAfterHalve = static_cast(n / 2); + gert::Shape outShape({m, nAfterHalve}); + gert::Shape scaleOutShape({m}); + auto out = executor->AllocTensor(outShape, DataType::DT_INT8, ge::FORMAT_ND); + auto scaleOut = executor->AllocTensor(scaleOutShape, DataType::DT_FLOAT, ge::FORMAT_ND); + auto ret = INFER_SHAPE(GroupedMatmulSwigluQuantWeightNzTensorList, + OP_INPUT(x, weight, perChannelScale, perTokenScale, groupList), + OP_OUTPUT(out, scaleOut), + OP_ATTR(swigluLimit)); + if (ret != ACLNN_SUCCESS) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "InferShape failed."); + return std::tuple(nullptr, nullptr); + } + ret = ADD_TO_LAUNCHER_LIST_AICORE(GroupedMatmulSwigluQuantWeightNzTensorList, + OP_INPUT(x, weight, perChannelScale, perTokenScale, groupList), + OP_OUTPUT(out, scaleOut), + OP_ATTR(swigluLimit)); + if (ret != ACLNN_SUCCESS) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "ADD_TO_LAUNCHER_LIST_AICORE failed."); + return std::tuple(nullptr, nullptr); + } + return std::tie(out, scaleOut); +} + +} // namespace l0op diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_kernel/grouped_matmul_swiglu_quant_weight_nz_tensor_list.cpp b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_kernel/grouped_matmul_swiglu_quant_weight_nz_tensor_list.cpp new file mode 100644 index 000000000..713a8cb41 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_kernel/grouped_matmul_swiglu_quant_weight_nz_tensor_list.cpp @@ -0,0 +1,80 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_weight_nz_tensor_list.cpp + * \brief + */ +#include "grouped_matmul_swiglu_quant_weight_nz_tensor_list.h" +#include +#include "grouped_matmul_swiglu_quant_weight_nz_tensor_list_split_ws.h" +using namespace AscendC; +using namespace matmul; +using namespace GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST; +using MM_DTYPE_Y = int32_t; + +template +using xType = MatmulType; + +template +using weightType = MatmulType; + +using yType = MatmulType; + +#define GMM_CV_SPLIT_IMP(computeClass, dtypeC, transA, transB, sync, cfg, aType, bType, cType) \ + do { \ + using matmulType = MMImplType, bType, cType, cType, cfg>; \ + matmulType::MT mm; \ + GET_TILING_DATA_MEMBER(GMMSwigluQuantTensorListTilingData, gmmSwigluTensorListBaseParams, gmmSwigluTensorListBaseParams_, tiling); \ + GET_TILING_DATA_MEMBER(GMMSwigluQuantTensorListTilingData, mmTilingData, mmTilingData_, tiling); \ + GET_TILING_DATA_MEMBER(GMMSwigluQuantTensorListTilingData, gmmSwigluTensorList, gmmSwiglu_, tiling); \ + if ASCEND_IS_AIC { \ + mm.SetSubBlockIdx(0); \ + mm.Init(&mmTilingData_, &tPipe); \ + } \ + computeClass computeOp(mm); \ + computeOp.Init(x, weight, perChannelScale, perTokenScale, groupList, quantOutput, quantScaleOutput, \ + user1, &gmmSwigluTensorListBaseParams_, &mmTilingData_, &gmmSwiglu_, &tPipe); \ + computeOp.Process(); \ + } while (0) + +extern "C" __global__ __aicore__ void grouped_matmul_swiglu_quant_weight_nz_tensor_list(GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale, + GM_ADDR groupList, GM_ADDR quantOutput, GM_ADDR quantScaleOutput, + GM_ADDR workspace, GM_ADDR tiling) { + TPipe tPipe; + AscendCUtils::SetOverflow(1); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + GM_ADDR user1 = GetUserWorkspace(workspace); + if (TILING_KEY_IS(0)) { // antiquant msd + KERNEL_TASK_TYPE(0, KERNEL_TYPE_MIX_AIC_1_2); + GMM_CV_SPLIT_IMP( + GMMSwigluCompute, // computeClass + DTYPE_WEIGHT_SCALE, + false, // transA + false, // transB + false, // sync + NZ_CFG_MDL, // cfg + xType, // aType + weightType, // bType + yType); // cType + } else if(TILING_KEY_IS(1)){ + KERNEL_TASK_TYPE(1, KERNEL_TYPE_MIX_AIC_1_2); + GMM_CV_SPLIT_IMP( + GMMSwigluSplitWorkSpaceCompute, // computeClass + DTYPE_WEIGHT_SCALE, + false, // transA + false, // transB + false, // sync + NZ_CFG_MDL, // cfg + xType, // aType + weightType, // bType + yType); // cType + } +} diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_kernel/grouped_matmul_swiglu_quant_weight_nz_tensor_list.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_kernel/grouped_matmul_swiglu_quant_weight_nz_tensor_list.h new file mode 100644 index 000000000..e1ea355f8 --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_kernel/grouped_matmul_swiglu_quant_weight_nz_tensor_list.h @@ -0,0 +1,507 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_weight_nz_tensor_list.h + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_H +#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_H + +#include "grouped_matmul_swiglu_quant_weight_nz_tensor_list_utils.h" +namespace GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST { +/** @brief internal computation class +*/ +template +class GMMSwigluCompute{ + public: + using AT = typename mmType::AT::T; + using BT = typename mmType::BT::T; + using B = typename mmType::BT; + using CT = typename mmType::CT::T; + using BiasT = typename mmType::BiasT::T; + using WT = int8_t; + constexpr static bool transposeX = mmType::AT::isTrans; + constexpr static bool transposeW = mmType::BT::isTrans; + + /** @brief constructor */ + __aicore__ inline GMMSwigluCompute(typename mmType::MT& mm_): mm(mm_) {} + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale, + GM_ADDR groupList, GM_ADDR quantOutput, GM_ADDR quantScaleOutput, + GM_ADDR workspace, + const GMMSwigluTensorListBaseParams* __restrict gmmBaseParamsIN, + const TCubeTiling* __restrict mmTilingDataIN, + const GMMSwigluTensorList* __restrict gmmSwigluIN, TPipe* tPipeIN); + __aicore__ inline void Process(); + private: + __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig, uint32_t coreIdx); + + __aicore__ inline void UpdateMnConfig(MNConfig &mnConfig); + + __aicore__ inline void SetMNConfig(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig); + + __aicore__ inline void SetMKN(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig); + + __aicore__ inline uint64_t GetWOffset(uint32_t tailN, uint32_t k); + + __aicore__ inline void MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock, + const uint32_t count, const uint32_t thresholdM_dimN); + template + __aicore__ inline void UpdateChannelScale(uint32_t loopidx); + __aicore__ inline void VectorCompute(uint32_t loopidx); + template + __aicore__ inline void PreLoadTokenAndChannel(LocalTensor& channelScaleLocal); + __aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig& vecConfig); + __aicore__ inline void customDataCopyIn(uint32_t outLoopIdx); + __aicore__ inline void customDataCopyOut(); + __aicore__ inline void Dequant(uint32_t loopidx); + __aicore__ inline void Quant(uint32_t loopidx); + __aicore__ inline void Swiglu(uint32_t loopidx); + private: + typename mmType::MT& mm; + const GMMSwigluTensorListBaseParams* __restrict gmmBaseParams; + const GMMSwigluTensorList* __restrict gmmSwigluTensorList; + const TCubeTiling* __restrict mmTilingData; + uint32_t blockIdx; + VecConfig vecConfig; + TPipe* pipe; + GlobalTensor xGM, weightGM; + GlobalTensor perChannelScaleGM; + GlobalTensor perTokenScaleGM; + GlobalTensor groupListGM; + GlobalTensor quantOutputGM; + GlobalTensor quantScaleOutputGM; + GlobalTensor mmOutGM; + // define the que + TQue mmOutQueue; + TQue perChannelScaleInQueue; + TQue quantOutQueue; + TQue quantScaleOutQueue; + TBuf reduceWorkspace; + TBuf castWorkspace; + bool sequentialWrite = true; + uint32_t cubeNum; // Matmul completions on the kernel + uint32_t groupNum; // Matmul completions on the kernel + int32_t preOffset; + int64_t aicCoreNum; + int64_t aivCoreNum; + GM_ADDR xTensorPtr; + GM_ADDR weightTensorPtr; + GM_ADDR perChannelScalePtr; +}; + +template +__aicore__ inline void GMMSwigluCompute::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale, + GM_ADDR groupList, GM_ADDR quantOutput, GM_ADDR quantScaleOutput, + GM_ADDR workspace, + const GMMSwigluTensorListBaseParams* __restrict gmmSwigluTensorListBaseParamsIn, + const TCubeTiling* __restrict mmTilingDataIN, + const GMMSwigluTensorList* __restrict gmmSwigluIN, TPipe* tPipeIN) +{ + aicCoreNum = GetBlockNum(); + aivCoreNum = aicCoreNum * 2; + blockIdx = GetBlockIdx(); + mmTilingData = mmTilingDataIN; + gmmBaseParams = gmmSwigluTensorListBaseParamsIn; + gmmSwigluTensorList = gmmSwigluIN; + pipe = tPipeIN; + xTensorPtr = x; + weightTensorPtr = weight; + perChannelScalePtr = perChannelScale; + groupNum = gmmSwigluTensorList->groupListLen; + if ASCEND_IS_AIC { + groupListGM.SetGlobalBuffer((__gm__ int64_t *)groupList, gmmSwigluTensorList->groupListLen); + mmOutGM.SetGlobalBuffer((__gm__ int32_t *)workspace, gmmBaseParams->M * gmmSwigluTensorList->tokenLen); + } + if ASCEND_IS_AIV { + mmOutGM.SetGlobalBuffer((__gm__ int32_t *)workspace, gmmBaseParams->M * gmmSwigluTensorList->tokenLen); + perChannelScaleGM.SetGlobalBuffer((__gm__ CHANNELDTYPE *)perChannelScale, gmmSwigluTensorList->groupListLen * gmmSwigluTensorList->tokenLen); + perTokenScaleGM.SetGlobalBuffer((__gm__ float *)perTokenScale, gmmSwigluTensorList->maxProcessRowNum); + groupListGM.SetGlobalBuffer((__gm__ int64_t *)groupList, gmmSwigluTensorList->groupListLen); + quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)quantOutput, gmmBaseParams->M * gmmSwigluTensorList->tokenLen / 2); + quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)quantScaleOutput, gmmSwigluTensorList->maxProcessRowNum); + } +} + +template +__aicore__ inline void GMMSwigluCompute::Process() { + MNConfig mnConfig; + if ASCEND_IS_AIC { + preOffset = 0; + int32_t prevSplitValue = 0; + for (uint32_t groupIdx = 0, count = 0; groupIdx < gmmSwigluTensorList->groupListLen; ++groupIdx) { + UpdateMnConfig(mnConfig); + int32_t currSplitValue = static_cast(groupListGM.GetValue(groupIdx)); + int32_t splitValue = currSplitValue - prevSplitValue; + prevSplitValue = currSplitValue; + SetMNConfig(splitValue, groupIdx, mnConfig); + if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) { + continue; + } + mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM); + mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN); + + uint32_t curCount = count + mnConfig.blockDimM * mnConfig.blockDimN; + uint32_t curBlock = blockIdx >= count ? blockIdx : blockIdx + gmmBaseParams->coreNum; + uint32_t thresholdM_dimN = THRESHOLD_BLOCK_NUM * mnConfig.blockDimN; + + while (curBlock < curCount) { + MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN); + MMCompute(groupIdx, mnConfig, blockIdx); + curBlock += aicCoreNum; + } + count = curCount % gmmBaseParams->coreNum; + } + SyncAll(); + } + + if ASCEND_IS_AIV { + UpdateVecConfig(blockIdx, vecConfig); + if (blockIdx < vecConfig.usedCoreNum) { + LocalTensor channelScaleLocal = perChannelScaleInQueue.AllocTensor(); + LocalTensor mmLocal = mmOutQueue.AllocTensor(); + LocalTensor quantLocal = quantOutQueue.AllocTensor(); + LocalTensor quantScaleLocal = quantScaleOutQueue.AllocTensor(); + mmOutQueue.EnQue(mmLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); + quantOutQueue.EnQue(quantLocal); + PreLoadTokenAndChannel(channelScaleLocal); + } + SyncAll(); + if (blockIdx < vecConfig.usedCoreNum) { + for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) { + vecConfig.innerLoopNum = outLoopIdx == (vecConfig.outLoopNum - 1) + ? vecConfig.tailLoopNum + : gmmSwigluTensorList->maxProcessRowNum; + customDataCopyIn(outLoopIdx); + for (uint32_t innerLoopIdx = 0; innerLoopIdx < vecConfig.innerLoopNum; innerLoopIdx++) { + UpdateChannelScale(innerLoopIdx); + VectorCompute(innerLoopIdx); + } + customDataCopyOut(); + } + + LocalTensor channelScaleLocal = perChannelScaleInQueue.DeQue(); + LocalTensor mmLocal = mmOutQueue.DeQue(); + LocalTensor quantLocal = quantOutQueue.DeQue(); + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + perChannelScaleInQueue.FreeTensor(channelScaleLocal); + mmOutQueue.FreeTensor(mmLocal); + quantScaleOutQueue.FreeTensor(quantScaleLocal); + quantOutQueue.FreeTensor(quantLocal); + } else { + return; + } + } +} + +template +template +__aicore__ inline void GMMSwigluCompute::PreLoadTokenAndChannel(LocalTensor& channelScaleLocal) +{ + GlobalTensor perChannelScaleTensor; + perChannelScaleTensor.SetGlobalBuffer(GetTensorAddr(vecConfig.curGroupIdx, perChannelScalePtr)); + + DataCopyExtParams copyChannelParams{1, static_cast(gmmSwigluTensorList->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0 ,0, 0}; + if constexpr(!IsSameType::value) { + LocalTensor dstLocalT = channelScaleLocal.template ReinterpretCast(); + DataCopyPad(dstLocalT[gmmSwigluTensorList->tokenLen], perChannelScaleTensor, copyChannelParams, padParams); + PipeBarrier(); + Cast(channelScaleLocal, dstLocalT[gmmSwigluTensorList->tokenLen], RoundMode::CAST_NONE, gmmSwigluTensorList->tokenLen); + } else { + DataCopyPad(channelScaleLocal, perChannelScaleTensor, copyChannelParams, padParams); + } + perChannelScaleInQueue.EnQue(channelScaleLocal); +} + +template +__aicore__ inline void GMMSwigluCompute::MMCompute(uint32_t groupIdx, MNConfig& mnConfig, uint32_t coreIdx) +{ + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN; + uint32_t curSingleM = mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM + : mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + uint64_t xOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.k; + if constexpr (transposeX) { + xOffset = mnConfig.mIdx * mnConfig.singleM; + } + uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN; + xGM.SetGlobalBuffer((__gm__ int8_t *)xTensorPtr + mnConfig.xBaseOffset); + weightGM.SetGlobalBuffer(GetTensorAddr(groupIdx, weightTensorPtr) + GetWOffset(tailN, mnConfig.k)); + if (mnConfig.blockDimM == 1){ + weightGM.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } + mnConfig.workSpaceOffset = outOffset + mnConfig.yBaseOffset; + mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k); + mm.SetSingleShape(curSingleM, curSingleN, mnConfig.k); + mm.SetTensorA(xGM[xOffset], transposeX); + mm.SetTensorB(weightGM, transposeW); + mm.template IterateAll(mmOutGM[mnConfig.workSpaceOffset], 0); +} + +template +__aicore__ inline void GMMSwigluCompute::UpdateMnConfig(MNConfig &mnConfig) { + if constexpr (B::format == CubeFormat::NZ) { + mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<32>(mnConfig.n); // 16: nz format last two dim size + } else { + mnConfig.wBaseOffset += mnConfig.k * mnConfig.n; + } + mnConfig.nAxisBaseOffset += mnConfig.n; + mnConfig.mAxisBaseOffset += mnConfig.m; + mnConfig.xBaseOffset += mnConfig.m * mnConfig.k; + mnConfig.yBaseOffset += mnConfig.m * mnConfig.n; +} + +template +__aicore__ inline void GMMSwigluCompute::SetMNConfig(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig) { + SetMKN(splitValue, groupIdx, mnConfig); + mnConfig.baseM = BASIC_M; + mnConfig.baseN = BASIC_N; + mnConfig.singleM = SINGLE_CORE_M; + mnConfig.singleN = SINGLE_CORE_N; +} + +template +__aicore__ inline void GMMSwigluCompute::SetMKN(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig) +{ + mnConfig.m = static_cast(splitValue); + mnConfig.k = gmmBaseParams->K; // tilingData + mnConfig.n = gmmBaseParams->N; // tilingData +} + +template +__aicore__ inline uint64_t GMMSwigluCompute::GetWOffset(uint32_t tailN, uint32_t k) { + uint64_t wOffset = 0; + if constexpr (mmType::BT::format == CubeFormat::NZ) { + wOffset = tailN * AlignUp<16>(k); // 16: nz format last two dim size + } else { + wOffset = tailN; + } + return wOffset; +} + +template +__aicore__ inline void GMMSwigluCompute::MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock, + const uint32_t count, const uint32_t thresholdM_dimN) { + mnConfig.mIdx = (curBlock - count) / mnConfig.blockDimN; + mnConfig.nIdx = (curBlock - count) % mnConfig.blockDimN; +} + +template +__aicore__ inline void GMMSwigluCompute::UpdateVecConfig(uint32_t blockIdx, VecConfig& vecConfig) +{ + // Step 1: Read grouplist reduceSum to calculate total data count + int64_t prevM = 0; + for (uint32_t groupIdx = 0; groupIdx < gmmSwigluTensorList->groupListLen; groupIdx++){ + int64_t currM = groupListGM.GetValue(groupIdx); + int64_t tempM = currM - prevM; + prevM = currM; + vecConfig.M += tempM; + } + // Step 2: Calculate core allocation + uint32_t eachCoreTaskNum = (vecConfig.M + aivCoreNum - 1) / aivCoreNum; + vecConfig.usedCoreNum = vecConfig.M >= aivCoreNum ? aivCoreNum : vecConfig.M; + uint32_t tailCoreIdx = vecConfig.M - (eachCoreTaskNum - 1) * vecConfig.usedCoreNum; + vecConfig.taskNum = blockIdx < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1; + vecConfig.startIdx = blockIdx < tailCoreIdx + ? eachCoreTaskNum * blockIdx + :((eachCoreTaskNum - 1) * blockIdx + tailCoreIdx); + vecConfig.curIdx = vecConfig.startIdx; + vecConfig.startOffset = vecConfig.startIdx * gmmSwigluTensorList->tokenLen; + vecConfig.curOffset = vecConfig.startOffset; + int64_t curStartIdx = vecConfig.startIdx; + prevM = 0; + for (uint32_t groupIdx = 0; groupIdx < gmmSwigluTensorList->groupListLen; groupIdx++){ + int64_t currM = groupListGM.GetValue(groupIdx); + int64_t tempM = currM - prevM; + prevM = currM; + if (curStartIdx >= 0 && curStartIdx - tempM < 0) { + vecConfig.curGroupIdx = groupIdx; + vecConfig.nextUpadteInterVal = tempM - curStartIdx; + } + curStartIdx -= tempM; + } + // Step 3: Calculate total data volume + vecConfig.outLoopNum = (vecConfig.taskNum + gmmSwigluTensorList->maxProcessRowNum - 1) / gmmSwigluTensorList->maxProcessRowNum; + vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwigluTensorList->maxProcessRowNum + ? vecConfig.taskNum % gmmSwigluTensorList->maxProcessRowNum + : gmmSwigluTensorList->maxProcessRowNum; + pipe->Reset(); + // Step 4: Allocate space + pipe->InitBuffer(mmOutQueue, 1, gmmSwigluTensorList->maxProcessRowNum * gmmSwigluTensorList->tokenLen * sizeof(int32_t)); + pipe->InitBuffer(perChannelScaleInQueue, 1, gmmSwigluTensorList->tokenLen * sizeof(float)); + pipe->InitBuffer(quantOutQueue, 1, gmmSwigluTensorList->maxProcessRowNum * gmmSwigluTensorList->tokenLen / 2 * sizeof(int8_t)); + pipe->InitBuffer(quantScaleOutQueue, 1, AlignUp(gmmSwigluTensorList->maxProcessRowNum, 8) * sizeof(float)); + pipe->InitBuffer(reduceWorkspace, 1024 * sizeof(float)); + pipe->InitBuffer(castWorkspace, 32 * sizeof(int8_t)); +} + +template +__aicore__ inline void GMMSwigluCompute::customDataCopyIn(uint32_t outLoopIdx) +{ + LocalTensor _inMMLocal_0 = mmOutQueue.DeQue(); + DataCopyExtParams copyParams_0{1, static_cast(vecConfig.innerLoopNum * gmmSwigluTensorList->tokenLen * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams padParams_0{false, 0 ,0, 0}; + DataCopyPad(_inMMLocal_0, mmOutGM[vecConfig.curOffset], copyParams_0, padParams_0); + + mmOutQueue.EnQue(_inMMLocal_0); + + LocalTensor _inMMLocal_1 = mmOutQueue.DeQue(); + + Cast(_inMMLocal_1.ReinterpretCast(), _inMMLocal_1, RoundMode::CAST_NONE, vecConfig.innerLoopNum * gmmSwigluTensorList->tokenLen); + + mmOutQueue.EnQue(_inMMLocal_1); + LocalTensor _inMMLocal_2 = mmOutQueue.DeQue(); + set_flag(PIPE_S, PIPE_V, EVENT_ID0); + for (uint32_t i = 0; i < vecConfig.innerLoopNum; i++){ + wait_flag(PIPE_S, PIPE_V, EVENT_ID0); + float scale = perTokenScaleGM.GetValue(vecConfig.curIdx); + set_flag(PIPE_S, PIPE_V, EVENT_ID0); + wait_flag(PIPE_S, PIPE_V, EVENT_ID0); + Muls(_inMMLocal_2[i * gmmSwigluTensorList->tokenLen], _inMMLocal_2[i * gmmSwigluTensorList->tokenLen], scale, gmmSwigluTensorList->tokenLen); + set_flag(PIPE_S, PIPE_V, EVENT_ID0); + vecConfig.curIdx++; + } + wait_flag(PIPE_S, PIPE_V, EVENT_ID0); + vecConfig.curOffset = vecConfig.curIdx * gmmSwigluTensorList->tokenLen; + mmOutQueue.EnQue(_inMMLocal_2); +} + +template +template +__aicore__ inline void GMMSwigluCompute::UpdateChannelScale(uint32_t loopIdx){ + // Update perChannel + if (unlikely(vecConfig.nextUpadteInterVal == 0)) { + int64_t loop = gmmSwigluTensorList->groupListLen - vecConfig.curGroupIdx; + while (loop--) { + int64_t curTemp = groupListGM.GetValue(vecConfig.curGroupIdx); + vecConfig.curGroupIdx++; + int64_t nextTemp = groupListGM.GetValue(vecConfig.curGroupIdx); + if(nextTemp != curTemp){ + vecConfig.nextUpadteInterVal = nextTemp - curTemp; + break; + } + } + LocalTensor _inChannel = perChannelScaleInQueue.DeQue(); + DataCopyExtParams copyParams{1, static_cast(gmmSwigluTensorList->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0 ,0, 0}; + + GlobalTensor perChannelScaleTensor; + perChannelScaleTensor.SetGlobalBuffer(GetTensorAddr(vecConfig.curGroupIdx, perChannelScalePtr)); + + if constexpr(!IsSameType::value) { + LocalTensor dstLocalT = _inChannel.template ReinterpretCast(); + DataCopyPad(dstLocalT[gmmSwigluTensorList->tokenLen], perChannelScaleTensor, copyParams, padParams); + PipeBarrier(); + Cast(_inChannel, dstLocalT[gmmSwigluTensorList->tokenLen], RoundMode::CAST_NONE, gmmSwigluTensorList->tokenLen); + } else { + DataCopyPad(_inChannel, perChannelScaleTensor, copyParams, padParams); + } + PipeBarrier(); + perChannelScaleInQueue.EnQue(_inChannel); + } +} + +template +__aicore__ inline void GMMSwigluCompute::VectorCompute(uint32_t loopIdx) { + Dequant(loopIdx); + Swiglu(loopIdx); + Quant(loopIdx); +} + +template +__aicore__ inline void GMMSwigluCompute::Dequant(uint32_t loopIdx) { + // perChanelScale * perTokenScale + LocalTensor mmLocal = mmOutQueue.DeQue(); + LocalTensor perChannelLocal = perChannelScaleInQueue.DeQue(); + Mul(mmLocal[loopIdx * gmmSwigluTensorList->tokenLen], mmLocal[loopIdx * gmmSwigluTensorList->tokenLen], perChannelLocal, gmmSwigluTensorList->tokenLen); + vecConfig.nextUpadteInterVal--; + mmOutQueue.EnQue(mmLocal); + perChannelScaleInQueue.EnQue(perChannelLocal); +} + +template +__aicore__ inline void GMMSwigluCompute::Swiglu(uint32_t loopIdx) { + // High-level API swiglu + LocalTensor _inMMLocal = mmOutQueue.DeQue(); + float beta = 1.0f; + LocalTensor workspaceLocal= reduceWorkspace.Get(); + LocalTensor src0Local = _inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen + gmmSwigluTensorList->tokenLen / 2]; + LocalTensor src1Local = _inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen]; + if(gmmSwigluTensorList->swigluLimit>1e-8f){ + Mins(src0Local, src0Local, gmmSwigluTensorList->swigluLimit, gmmSwigluTensorList->tokenLen / 2); + PipeBarrier(); + Maxs(src0Local, src0Local, (-1.0f * gmmSwigluTensorList->swigluLimit), gmmSwigluTensorList->tokenLen / 2); + PipeBarrier(); + + Mins(src1Local, src1Local, gmmSwigluTensorList->swigluLimit, gmmSwigluTensorList->tokenLen / 2); + PipeBarrier(); + } + SwiGLU(workspaceLocal, src0Local, src1Local, beta, gmmSwigluTensorList->tokenLen / 2); + PipeBarrier(); + DataCopyParams repeatParams{1, static_cast((gmmSwigluTensorList->tokenLen / 2) / 8), 0, 0}; + DataCopy(_inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen], workspaceLocal, repeatParams); + mmOutQueue.EnQue(_inMMLocal); +} + +template +__aicore__ inline void GMMSwigluCompute::Quant(uint32_t loopIdx) { + LocalTensor _inMMLocal = mmOutQueue.DeQue(); + Abs(_inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen + gmmSwigluTensorList->tokenLen / BISECT], + _inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen], + gmmSwigluTensorList->tokenLen / BISECT); + LocalTensor workspaceLocal= reduceWorkspace.Get(); + PipeBarrier(); + ReduceMaxTemplate(workspaceLocal, + _inMMLocal, loopIdx * gmmSwigluTensorList->tokenLen + gmmSwigluTensorList->tokenLen / BISECT, gmmSwigluTensorList->tokenLen / BISECT); + PipeBarrier(); + float quantScale = workspaceLocal.GetValue(0) / QUANT_SCALE_INT8; + PipeBarrier(); + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + PipeBarrier(); + quantScaleLocal.SetValue(loopIdx, quantScale); + PipeBarrier(); + quantScale = 1 / quantScale; + PipeBarrier(); + Muls(_inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen], _inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen], + quantScale, gmmSwigluTensorList->tokenLen / BISECT); + PipeBarrier(); + LocalTensor quantLocal = quantOutQueue.DeQue(); + int32_t dstTempOffset = static_cast(loopIdx * gmmSwigluTensorList->tokenLen / BISECT); + int32_t srcTempOffset = static_cast(loopIdx * gmmSwigluTensorList->tokenLen); + int32_t tempCount = static_cast(gmmSwigluTensorList->tokenLen / BISECT); + LocalTensor castSpace = castWorkspace.Get(); + CastFp32ToInt8Template(quantLocal, _inMMLocal, castSpace, dstTempOffset, srcTempOffset, tempCount); + mmOutQueue.EnQue(_inMMLocal); + quantOutQueue.EnQue(quantLocal); +} + +template +__aicore__ inline void GMMSwigluCompute::customDataCopyOut() { + // perChanelScale * perTokenScale + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + DataCopyParams copyParams_0{1, (uint16_t)(vecConfig.innerLoopNum * sizeof(float)), 0, 0}; + PipeBarrier(); + DataCopyPad(quantScaleOutputGM[vecConfig.startIdx], quantScaleLocal, copyParams_0); + LocalTensor quantLocal = quantOutQueue.DeQue(); + DataCopyParams copyParams_1{1, (uint16_t)(vecConfig.innerLoopNum * gmmSwigluTensorList->tokenLen / 2 * sizeof(int8_t)), 0, 0}; + PipeBarrier(); + DataCopyPad(quantOutputGM[vecConfig.startIdx * gmmSwigluTensorList->tokenLen / 2], quantLocal, copyParams_1); + PipeBarrier(); + vecConfig.startIdx += vecConfig.innerLoopNum; + vecConfig.startOffset = vecConfig.startIdx * gmmSwigluTensorList->tokenLen; + quantOutQueue.EnQue(quantLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); +} + +} // namespace GROUPED_MATMUL +#endif // ASCENDC_GROUPED_MATMUL_QUANT_MIXCORE_H diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_kernel/grouped_matmul_swiglu_quant_weight_nz_tensor_list_split_ws.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_kernel/grouped_matmul_swiglu_quant_weight_nz_tensor_list_split_ws.h new file mode 100644 index 000000000..d2aad686c --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_kernel/grouped_matmul_swiglu_quant_weight_nz_tensor_list_split_ws.h @@ -0,0 +1,596 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_weight_nz_tensor_list_split_ws.h + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_SPLIT_WS_H +#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_SPLIT_WS_H + +#include "grouped_matmul_swiglu_quant_weight_nz_tensor_list_utils.h" +namespace GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST { +/** @brief internal computation class +*/ + +template +class GMMSwigluSplitWorkSpaceCompute{ + public: + using AT = typename mmType::AT::T; + using BT = typename mmType::BT::T; + using B = typename mmType::BT; + using CT = typename mmType::CT::T; + using BiasT = typename mmType::BiasT::T; + using WT = int8_t; + constexpr static bool transposeX = mmType::AT::isTrans; + constexpr static bool transposeW = mmType::BT::isTrans; + + /** @brief constructor */ + __aicore__ inline GMMSwigluSplitWorkSpaceCompute(typename mmType::MT& mm_): mm(mm_) {} + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale, + GM_ADDR groupList, GM_ADDR quantOutput, GM_ADDR quantScaleOutput, + GM_ADDR workspace, + const GMMSwigluTensorListBaseParams* __restrict gmmBaseParamsIN, + const TCubeTiling* __restrict mmTilingDataIN, + const GMMSwigluTensorList* __restrict gmmSwigluIN, TPipe* tPipeIN); + __aicore__ inline void Process(); + + private: + __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig, uint32_t coreIdx, GlobalTensor &mmOutGM); + + __aicore__ inline void UpdateMnConfig(MNConfig &mnConfig); + + __aicore__ inline void SetMNConfig(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig); + + __aicore__ inline void SetMKN(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig); + + __aicore__ inline uint64_t GetWOffset(uint32_t tailN, uint32_t k); + + __aicore__ inline void MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock, + const uint32_t count, const uint32_t thresholdM_dimN); + + template + __aicore__ inline void UpdateChannelScale(uint32_t loopidx, VecConfig& vecConfig); + + __aicore__ inline void VectorCompute(uint32_t loopidx, VecConfig& vecConfig); + + template + __aicore__ inline void PreLoadTokenAndChannel(LocalTensor& channelScaleLocal, VecConfig& vecConfig); + + __aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig& vecConfig); + + __aicore__ inline void UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig, int32_t workspaceSplitLoopIdx); + + __aicore__ inline void InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig); + + __aicore__ inline void customDataCopyIn(uint32_t outLoopIdx, GlobalTensor &mmOutGM, VecConfig& vecConfig); + + __aicore__ inline void customDataCopyOut(VecConfig& vecConfig); + + __aicore__ inline void Dequant(uint32_t loopidx, VecConfig& vecConfig); + + __aicore__ inline void Quant(uint32_t loopidx, VecConfig& vecConfig); + + __aicore__ inline void Swiglu(uint32_t loopidx, VecConfig& vecConfig); + + private: + typename mmType::MT& mm; + const GMMSwigluTensorListBaseParams* __restrict gmmBaseParams; + const GMMSwigluTensorList* __restrict gmmSwigluTensorList; + const TCubeTiling* __restrict mmTilingData; + uint32_t blockIdx; + WorkSpaceSplitConfig workspaceSplitConfig; + TPipe* pipe; + GlobalTensor xGM; + GlobalTensor weightGM; + GlobalTensor perChannelScaleGM; + GlobalTensor perTokenScaleGM; + GlobalTensor groupListGM; + GlobalTensor quantOutputGM; + GlobalTensor quantScaleOutputGM; + GlobalTensor mmOutGM1; + GlobalTensor mmOutGM2; + // define the que + TQue mmOutQueue; + TQue perChannelScaleInQueue; + TQue quantOutQueue; + TQue quantScaleOutQueue; + TBuf reduceWorkspace; + TBuf castWorkspace; + bool sequentialWrite = true; + uint32_t cubeNum; // Matmul completions on the kernel + uint32_t groupNum; // Matmul completions on the kernel + int64_t aicCoreNum; + int64_t aivCoreNum; + GM_ADDR xTensorPtr; + GM_ADDR weightTensorPtr; + GM_ADDR perChannelScalePtr; +}; + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::Init(GM_ADDR x, GM_ADDR weight, + GM_ADDR perChannelScale, GM_ADDR perTokenScale, + GM_ADDR groupList, GM_ADDR quantOutput, + GM_ADDR quantScaleOutput, GM_ADDR workspace, + const GMMSwigluTensorListBaseParams* __restrict gmmSwigluTensorListBaseParamsIn, + const TCubeTiling* __restrict mmTilingDataIN, + const GMMSwigluTensorList* __restrict gmmSwigluIN, TPipe* tPipeIN) +{ + aicCoreNum = GetBlockNum(); + aivCoreNum = aicCoreNum * 2; + blockIdx = GetBlockIdx(); + pipe = tPipeIN; + xTensorPtr = x; + weightTensorPtr = weight; + perChannelScalePtr = perChannelScale; + mmTilingData = mmTilingDataIN; + gmmBaseParams = gmmSwigluTensorListBaseParamsIn; + gmmSwigluTensorList = gmmSwigluIN; + groupNum = gmmSwigluTensorList->groupListLen; + if ASCEND_IS_AIC { + groupListGM.SetGlobalBuffer((__gm__ int64_t *)groupList, gmmSwigluTensorList->groupListLen); + mmOutGM1.SetGlobalBuffer((__gm__ int32_t *)workspace, gmmBaseParams->mLimit * gmmSwigluTensorList->tokenLen); + mmOutGM2.SetGlobalBuffer((__gm__ int32_t *)workspace + gmmBaseParams->mLimit * gmmSwigluTensorList->tokenLen, + gmmBaseParams->mLimit * gmmSwigluTensorList->tokenLen); + } + if ASCEND_IS_AIV { + mmOutGM1.SetGlobalBuffer((__gm__ int32_t *)workspace, gmmBaseParams->mLimit * gmmSwigluTensorList->tokenLen); + mmOutGM2.SetGlobalBuffer((__gm__ int32_t *)workspace + gmmBaseParams->mLimit * gmmSwigluTensorList->tokenLen, + gmmBaseParams->mLimit * gmmSwigluTensorList->tokenLen); + perChannelScaleGM.SetGlobalBuffer((__gm__ CHANNELDTYPE *)perChannelScale, + gmmSwigluTensorList->groupListLen * gmmSwigluTensorList->tokenLen); + perTokenScaleGM.SetGlobalBuffer((__gm__ float *)perTokenScale, gmmBaseParams->M); + groupListGM.SetGlobalBuffer((__gm__ int64_t *)groupList, gmmSwigluTensorList->groupListLen); + quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)quantOutput, gmmBaseParams->M * gmmSwigluTensorList->tokenLen / 2); + quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)quantScaleOutput, gmmBaseParams->M); + } +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig) +{ + workspaceSplitConfig.M = groupListGM.GetValue(gmmSwigluTensorList->groupListLen - 1); + workspaceSplitConfig.loopCount = Ceil(workspaceSplitConfig.M, gmmBaseParams->mLimit); + workspaceSplitConfig.notLastTaskSize = gmmBaseParams->mLimit; + workspaceSplitConfig.lastLoopTaskSize = workspaceSplitConfig.M - (workspaceSplitConfig.loopCount - 1) * gmmBaseParams->mLimit; + workspaceSplitConfig.leftMatrixStartIndex = 0; + workspaceSplitConfig.rightMatrixExpertStartIndex = 0; + workspaceSplitConfig.rightMatrixExpertNextStartIndex = 0; + workspaceSplitConfig.isLastLoop = false; +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig, int32_t workspaceSplitLoopIdx) +{ + workspaceSplitConfig.leftMatrixStartIndex = workspaceSplitLoopIdx * gmmBaseParams->mLimit; + workspaceSplitConfig.rightMatrixExpertStartIndex = workspaceSplitConfig.rightMatrixExpertNextStartIndex; + workspaceSplitConfig.rightMatrixExpertEndIndex = workspaceSplitConfig.rightMatrixExpertStartIndex; + // Calculate the right expert matrix end index (rightMatrixExpertEndIndex) and the next start index (rightMatrixExpertNextStartIndex) + int32_t curTaskNum = 0; + int32_t nextTaskNum = 0; + while(workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwigluTensorList->groupListLen) + { + curTaskNum = groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex) - workspaceSplitConfig.leftMatrixStartIndex; + int32_t nextTaskIdx = workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwigluTensorList->groupListLen - 1 \ + ? gmmSwigluTensorList->groupListLen - 1 \ + : workspaceSplitConfig.rightMatrixExpertEndIndex + 1; + nextTaskNum = groupListGM.GetValue(nextTaskIdx) - workspaceSplitConfig.leftMatrixStartIndex; + if (curTaskNum > gmmBaseParams->mLimit){ + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex; + break; + } else if (curTaskNum == gmmBaseParams->mLimit && nextTaskNum > gmmBaseParams->mLimit){ + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex + 1; + break; + } else if (nextTaskNum > gmmBaseParams->mLimit){ + workspaceSplitConfig.rightMatrixExpertEndIndex++; + workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex; + break; + } + workspaceSplitConfig.rightMatrixExpertEndIndex++; + } + workspaceSplitConfig.isLastLoop = workspaceSplitLoopIdx == workspaceSplitConfig.loopCount - 1 ? true : false; + + if (workspaceSplitConfig.isLastLoop) { + workspaceSplitConfig.rightMatrixExpertEndIndex = workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwigluTensorList->groupListLen \ + ? gmmSwigluTensorList->groupListLen - 1 \ + : workspaceSplitConfig.rightMatrixExpertEndIndex; + } +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::Process() { + InitWorkSpaceSplitConfig(workspaceSplitConfig); + int32_t parallelNum = gmmBaseParams->isPreFill ? 2 : 1; // 2: double workspace buffer + for (int32_t workspaceSplitLoopIdx = 0; workspaceSplitLoopIdx < workspaceSplitConfig.loopCount; workspaceSplitLoopIdx++) { + UpdateWorkSpaceSplitConfig(workspaceSplitConfig, workspaceSplitLoopIdx); + GlobalTensor mmOutGM = (workspaceSplitLoopIdx % 2 == 0 ) ? mmOutGM1 : mmOutGM2; + + if ASCEND_IS_AIC { + if (workspaceSplitLoopIdx >= parallelNum){ // first parallelNum core no need to wait + SyncAll(); + } + MNConfig mnConfig; + int32_t prevSplitValue = workspaceSplitConfig.leftMatrixStartIndex; + for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex, count = 0; groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; ++groupIdx) { + UpdateMnConfig(mnConfig); + int32_t currSplitValue = static_cast(groupListGM.GetValue(groupIdx)); + currSplitValue = currSplitValue > (workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit \ + ? (workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit \ + : currSplitValue; + int32_t splitValue = currSplitValue - prevSplitValue; + prevSplitValue = currSplitValue; + SetMNConfig(splitValue, groupIdx, mnConfig); + if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) { + continue; + } + mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM); + mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN); + + uint32_t curCount = count + mnConfig.blockDimM * mnConfig.blockDimN; + uint32_t curBlock = blockIdx >= count ? blockIdx : blockIdx + gmmBaseParams->coreNum; + uint32_t thresholdM_dimN = THRESHOLD_BLOCK_NUM * mnConfig.blockDimN; + + while (curBlock < curCount) { + MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN); + MMCompute(groupIdx, mnConfig, blockIdx, mmOutGM); + curBlock += aicCoreNum; + } + count = curCount % gmmBaseParams->coreNum; + } + SyncAll(); + } + + if ASCEND_IS_AIV { + VecConfig vecConfig; + UpdateVecConfig(blockIdx, vecConfig); + if (blockIdx < vecConfig.usedCoreNum) { + LocalTensor channelScaleLocal = perChannelScaleInQueue.AllocTensor(); + LocalTensor mmLocal = mmOutQueue.AllocTensor(); + LocalTensor quantLocal = quantOutQueue.AllocTensor(); + LocalTensor quantScaleLocal = quantScaleOutQueue.AllocTensor(); + mmOutQueue.EnQue(mmLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); + quantOutQueue.EnQue(quantLocal); + PreLoadTokenAndChannel(channelScaleLocal, vecConfig); + } + SyncAll(); + if (blockIdx < vecConfig.usedCoreNum) { + for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) { + vecConfig.innerLoopNum = outLoopIdx == (vecConfig.outLoopNum - 1) + ? vecConfig.tailLoopNum + : gmmSwigluTensorList->maxProcessRowNum; + PipeBarrier(); + customDataCopyIn(outLoopIdx, mmOutGM, vecConfig); + PipeBarrier(); + for (uint32_t innerLoopIdx = 0; innerLoopIdx < vecConfig.innerLoopNum; innerLoopIdx++) { + UpdateChannelScale(innerLoopIdx, vecConfig); + VectorCompute(innerLoopIdx, vecConfig); + } + PipeBarrier(); + customDataCopyOut(vecConfig); + PipeBarrier(); + } + + LocalTensor channelScaleLocal = perChannelScaleInQueue.DeQue(); + LocalTensor mmLocal = mmOutQueue.DeQue(); + LocalTensor quantLocal = quantOutQueue.DeQue(); + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + perChannelScaleInQueue.FreeTensor(channelScaleLocal); + mmOutQueue.FreeTensor(mmLocal); + quantScaleOutQueue.FreeTensor(quantScaleLocal); + quantOutQueue.FreeTensor(quantLocal); + } + if (workspaceSplitLoopIdx < workspaceSplitConfig.loopCount - parallelNum){ + SyncAll(); + } + } + } +} + +template +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::PreLoadTokenAndChannel(LocalTensor& channelScaleLocal, VecConfig& vecConfig) +{ + GlobalTensor perChannelScaleTensor; + perChannelScaleTensor.SetGlobalBuffer(GetTensorAddr(vecConfig.curGroupIdx, perChannelScalePtr)); + + DataCopyExtParams copyChannelParams{1, static_cast(gmmSwigluTensorList->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0 ,0, 0}; + if constexpr(!IsSameType::value) { + LocalTensor dstLocalT = channelScaleLocal.template ReinterpretCast(); + DataCopyPad(dstLocalT[gmmSwigluTensorList->tokenLen], perChannelScaleTensor, copyChannelParams, padParams); + PipeBarrier(); + Cast(channelScaleLocal, dstLocalT[gmmSwigluTensorList->tokenLen], RoundMode::CAST_NONE, gmmSwigluTensorList->tokenLen); + } else { + DataCopyPad(channelScaleLocal, perChannelScaleTensor, copyChannelParams, padParams); + } + perChannelScaleInQueue.EnQue(channelScaleLocal); +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::MMCompute(uint32_t groupIdx, MNConfig& mnConfig, uint32_t coreIdx, GlobalTensor &mmOutGM) +{ + uint32_t tailN = mnConfig.nIdx * mnConfig.singleN; + uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN; + uint32_t curSingleM = mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM + : mnConfig.m - mnConfig.mIdx * mnConfig.singleM; + uint64_t xOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.k; + if constexpr (transposeX) { + xOffset = mnConfig.mIdx * mnConfig.singleM; + } + uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN; + xGM.SetGlobalBuffer((__gm__ int8_t *)xTensorPtr + mnConfig.xBaseOffset + workspaceSplitConfig.leftMatrixStartIndex * mnConfig.k); + weightGM.SetGlobalBuffer(GetTensorAddr(groupIdx, weightTensorPtr) + GetWOffset(tailN, mnConfig.k)); + if (mnConfig.blockDimM == 1){ + weightGM.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } else { + weightGM.SetL2CacheHint(CacheMode::CACHE_MODE_NORMAL); + } + mnConfig.workSpaceOffset = outOffset + mnConfig.yBaseOffset; + mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k); + mm.SetSingleShape(curSingleM, curSingleN, mnConfig.k); + mm.SetTensorA(xGM[xOffset], transposeX); + mm.SetTensorB(weightGM, transposeW); + mm.template IterateAll(mmOutGM[mnConfig.workSpaceOffset], 0); +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::UpdateMnConfig(MNConfig &mnConfig) { + if constexpr (B::format == CubeFormat::NZ) { + mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<32>(mnConfig.n); // 16: nz format last two dim size + } else { + mnConfig.wBaseOffset += mnConfig.k * mnConfig.n; + } + mnConfig.nAxisBaseOffset += mnConfig.n; + mnConfig.mAxisBaseOffset += mnConfig.m; + mnConfig.xBaseOffset += mnConfig.m * mnConfig.k; + mnConfig.yBaseOffset += mnConfig.m * mnConfig.n; +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::SetMNConfig(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig) { + SetMKN(splitValue, groupIdx, mnConfig); + mnConfig.baseM = BASIC_M; + mnConfig.baseN = BASIC_N; + mnConfig.singleM = SINGLE_CORE_M; + mnConfig.singleN = SINGLE_CORE_N; +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::SetMKN(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig) +{ + mnConfig.m = static_cast(splitValue); + mnConfig.k = gmmBaseParams->K; // tilingData + mnConfig.n = gmmBaseParams->N; // tilingData +} + +template +__aicore__ inline uint64_t GMMSwigluSplitWorkSpaceCompute::GetWOffset(uint32_t tailN, uint32_t k) { + uint64_t wOffset = 0; + if constexpr (mmType::BT::format == CubeFormat::NZ) { + wOffset = tailN * AlignUp<16>(k); // 16: nz format last two dim size + } else { + wOffset = tailN; + } + return wOffset; +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock, + const uint32_t count, const uint32_t thresholdM_dimN) { + mnConfig.mIdx = (curBlock - count) / mnConfig.blockDimN; + mnConfig.nIdx = (curBlock - count) % mnConfig.blockDimN; +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::UpdateVecConfig(uint32_t blockIdx, VecConfig& vecConfig) +{ + // Step 1: Read grouplist reduceSum to calculate total data count + vecConfig.M = workspaceSplitConfig.isLastLoop \ + ? workspaceSplitConfig.lastLoopTaskSize\ + : workspaceSplitConfig.notLastTaskSize; + // Step 2: Calculate core allocation + uint32_t eachCoreTaskNum = (vecConfig.M + aivCoreNum - 1) / aivCoreNum; + vecConfig.usedCoreNum = vecConfig.M >= aivCoreNum ? aivCoreNum : vecConfig.M; + uint32_t tailCoreIdx = vecConfig.M - (eachCoreTaskNum - 1) * vecConfig.usedCoreNum; + vecConfig.taskNum = blockIdx < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1; + vecConfig.startIdx = blockIdx < tailCoreIdx + ? eachCoreTaskNum * blockIdx + :((eachCoreTaskNum - 1) * blockIdx + tailCoreIdx); + vecConfig.curIdx = vecConfig.startIdx; + vecConfig.startOffset = vecConfig.startIdx * gmmSwigluTensorList->tokenLen; + vecConfig.curOffset = vecConfig.startOffset; + int64_t curStartIdx = vecConfig.startIdx; + int64_t prevM = workspaceSplitConfig.leftMatrixStartIndex; + for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex; groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; groupIdx++){ + int64_t currM = groupListGM.GetValue(groupIdx); + int64_t tempM = currM - prevM; + prevM = currM; + if (curStartIdx >= 0 && curStartIdx - tempM < 0) { + vecConfig.curGroupIdx = groupIdx; + vecConfig.nextUpadteInterVal = tempM - curStartIdx; + } + curStartIdx -= tempM; + } + // Step 3: Calculate total data volume + vecConfig.outLoopNum = (vecConfig.taskNum + gmmSwigluTensorList->maxProcessRowNum - 1) / gmmSwigluTensorList->maxProcessRowNum; + vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwigluTensorList->maxProcessRowNum + ? vecConfig.taskNum % gmmSwigluTensorList->maxProcessRowNum + : gmmSwigluTensorList->maxProcessRowNum; + pipe->Reset(); + // Step 4: Allocate space + pipe->InitBuffer(mmOutQueue, 1, gmmSwigluTensorList->maxProcessRowNum * gmmSwigluTensorList->tokenLen * sizeof(int32_t)); + pipe->InitBuffer(perChannelScaleInQueue, 1, gmmSwigluTensorList->tokenLen * sizeof(float)); + pipe->InitBuffer(quantOutQueue, 1, gmmSwigluTensorList->maxProcessRowNum * gmmSwigluTensorList->tokenLen / 2 * sizeof(int8_t)); + pipe->InitBuffer(quantScaleOutQueue, 1, AlignUp(gmmSwigluTensorList->maxProcessRowNum, 8) * sizeof(float)); + pipe->InitBuffer(reduceWorkspace, 1024 * sizeof(float)); + pipe->InitBuffer(castWorkspace, 32 * sizeof(int8_t)); +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::customDataCopyIn(uint32_t outLoopIdx, GlobalTensor &mmOutGM, VecConfig& vecConfig) +{ + LocalTensor _inMMLocal_0 = mmOutQueue.DeQue(); + DataCopyExtParams copyParams_0{1, static_cast(vecConfig.innerLoopNum * gmmSwigluTensorList->tokenLen * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams padParams_0{false, 0 ,0, 0}; + PipeBarrier(); + DataCopyPad(_inMMLocal_0, mmOutGM[vecConfig.curOffset], copyParams_0, padParams_0); + mmOutQueue.EnQue(_inMMLocal_0); + + LocalTensor _inMMLocal_1 = mmOutQueue.DeQue(); + + Cast(_inMMLocal_1.ReinterpretCast(), _inMMLocal_1, RoundMode::CAST_NONE, vecConfig.innerLoopNum * gmmSwigluTensorList->tokenLen); + + mmOutQueue.EnQue(_inMMLocal_1); + LocalTensor _inMMLocal_2 = mmOutQueue.DeQue(); + set_flag(PIPE_S, PIPE_V, EVENT_ID0); + for (uint32_t i = 0; i < vecConfig.innerLoopNum; i++){ + wait_flag(PIPE_S, PIPE_V, EVENT_ID0); + float scale = perTokenScaleGM.GetValue(vecConfig.curIdx + workspaceSplitConfig.leftMatrixStartIndex); + set_flag(PIPE_S, PIPE_V, EVENT_ID0); + wait_flag(PIPE_S, PIPE_V, EVENT_ID0); + Muls(_inMMLocal_2[i * gmmSwigluTensorList->tokenLen], _inMMLocal_2[i * gmmSwigluTensorList->tokenLen], scale, gmmSwigluTensorList->tokenLen); + set_flag(PIPE_S, PIPE_V, EVENT_ID0); + vecConfig.curIdx++; + } + wait_flag(PIPE_S, PIPE_V, EVENT_ID0); + vecConfig.curOffset = vecConfig.curIdx * gmmSwigluTensorList->tokenLen; + mmOutQueue.EnQue(_inMMLocal_2); +} + +template +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::UpdateChannelScale(uint32_t loopIdx, VecConfig& vecConfig){ + // Update perChannel + if (unlikely(vecConfig.nextUpadteInterVal == 0)) { + int64_t loop = gmmSwigluTensorList->groupListLen - vecConfig.curGroupIdx; + while (loop--) { + int64_t curTemp = groupListGM.GetValue(vecConfig.curGroupIdx); + vecConfig.curGroupIdx++; + int64_t nextTemp = groupListGM.GetValue(vecConfig.curGroupIdx); + if(nextTemp != curTemp){ + vecConfig.nextUpadteInterVal = nextTemp - curTemp; + break; + } + } + LocalTensor _inChannel = perChannelScaleInQueue.DeQue(); + DataCopyExtParams copyParams{1, static_cast(gmmSwigluTensorList->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0 ,0, 0}; + + GlobalTensor perChannelScaleTensor; + perChannelScaleTensor.SetGlobalBuffer(GetTensorAddr(vecConfig.curGroupIdx, perChannelScalePtr)); + + if constexpr(!IsSameType::value) { + LocalTensor dstLocalT = _inChannel.template ReinterpretCast(); + DataCopyPad(dstLocalT[gmmSwigluTensorList->tokenLen], perChannelScaleTensor, copyParams, padParams); + PipeBarrier(); + Cast(_inChannel, dstLocalT[gmmSwigluTensorList->tokenLen], RoundMode::CAST_NONE, gmmSwigluTensorList->tokenLen); + } else { + DataCopyPad(_inChannel, perChannelScaleTensor, copyParams, padParams); + } + PipeBarrier(); + + perChannelScaleInQueue.EnQue(_inChannel); + } +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::VectorCompute(uint32_t loopIdx, VecConfig& vecConfig) { + Dequant(loopIdx, vecConfig); + Swiglu(loopIdx, vecConfig); + Quant(loopIdx, vecConfig); +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::Dequant(uint32_t loopIdx, VecConfig& vecConfig) { + // perChanelScale * perTokenScale + LocalTensor mmLocal = mmOutQueue.DeQue(); + LocalTensor perChannelLocal = perChannelScaleInQueue.DeQue(); + Mul(mmLocal[loopIdx * gmmSwigluTensorList->tokenLen], mmLocal[loopIdx * gmmSwigluTensorList->tokenLen], perChannelLocal, gmmSwigluTensorList->tokenLen); + vecConfig.nextUpadteInterVal--; + mmOutQueue.EnQue(mmLocal); + perChannelScaleInQueue.EnQue(perChannelLocal); +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::Swiglu(uint32_t loopIdx, VecConfig& vecConfig) { + // High-level API swiglu + LocalTensor _inMMLocal = mmOutQueue.DeQue(); + float beta = 1.0f; + LocalTensor workspaceLocal= reduceWorkspace.Get(); + LocalTensor src0Local = _inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen + gmmSwigluTensorList->tokenLen / 2]; + LocalTensor src1Local = _inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen]; + if(gmmSwigluTensorList->swigluLimit>1e-8f){ + Mins(src0Local, src0Local, gmmSwigluTensorList->swigluLimit, gmmSwigluTensorList->tokenLen / 2); + PipeBarrier(); + Maxs(src0Local, src0Local, (-1.0f * gmmSwigluTensorList->swigluLimit), gmmSwigluTensorList->tokenLen / 2); + PipeBarrier(); + Mins(src1Local, src1Local, gmmSwigluTensorList->swigluLimit, gmmSwigluTensorList->tokenLen / 2); + PipeBarrier(); + } + SwiGLU(workspaceLocal, src0Local, src1Local, beta, gmmSwigluTensorList->tokenLen / 2); + PipeBarrier(); + DataCopyParams repeatParams{1, static_cast((gmmSwigluTensorList->tokenLen / 2) / 8), 0, 0}; + DataCopy(_inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen], workspaceLocal, repeatParams); + mmOutQueue.EnQue(_inMMLocal); +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::Quant(uint32_t loopIdx, VecConfig& vecConfig) { + LocalTensor _inMMLocal = mmOutQueue.DeQue(); + Abs(_inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen + gmmSwigluTensorList->tokenLen / BISECT], + _inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen], + gmmSwigluTensorList->tokenLen / BISECT); + LocalTensor workspaceLocal= reduceWorkspace.Get(); + PipeBarrier(); + ReduceMaxTemplate(workspaceLocal, + _inMMLocal, loopIdx * gmmSwigluTensorList->tokenLen + gmmSwigluTensorList->tokenLen / BISECT, gmmSwigluTensorList->tokenLen / BISECT); + PipeBarrier(); + float quantScale = workspaceLocal.GetValue(0) / QUANT_SCALE_INT8; + PipeBarrier(); + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + PipeBarrier(); + quantScaleLocal.SetValue(loopIdx, quantScale); + PipeBarrier(); + quantScale = 1 / quantScale; + PipeBarrier(); + Muls(_inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen], _inMMLocal[loopIdx * gmmSwigluTensorList->tokenLen], + quantScale, gmmSwigluTensorList->tokenLen / BISECT); + PipeBarrier(); + LocalTensor quantLocal = quantOutQueue.DeQue(); + int32_t dstTempOffset = static_cast(loopIdx * gmmSwigluTensorList->tokenLen / BISECT); + int32_t srcTempOffset = static_cast(loopIdx * gmmSwigluTensorList->tokenLen); + int32_t tempCount = static_cast(gmmSwigluTensorList->tokenLen / BISECT); + LocalTensor castSpace = castWorkspace.Get(); + CastFp32ToInt8Template(quantLocal, _inMMLocal, castSpace, dstTempOffset, srcTempOffset, tempCount); + mmOutQueue.EnQue(_inMMLocal); + quantOutQueue.EnQue(quantLocal); +} + +template +__aicore__ inline void GMMSwigluSplitWorkSpaceCompute::customDataCopyOut(VecConfig& vecConfig) { + LocalTensor quantScaleLocal = quantScaleOutQueue.DeQue(); + DataCopyParams copyParams_0{1, (uint16_t)(vecConfig.innerLoopNum * sizeof(float)), 0, 0}; + PipeBarrier(); + DataCopyPad(quantScaleOutputGM[workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx], quantScaleLocal, copyParams_0); + LocalTensor quantLocal = quantOutQueue.DeQue(); + DataCopyParams copyParams_1{1, (uint16_t)(vecConfig.innerLoopNum * gmmSwigluTensorList->tokenLen / 2 * sizeof(int8_t)), 0, 0}; + PipeBarrier(); + DataCopyPad(quantOutputGM[(workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx) * gmmSwigluTensorList->tokenLen / 2], quantLocal, copyParams_1); + PipeBarrier(); + vecConfig.startIdx += vecConfig.innerLoopNum; + vecConfig.startOffset = vecConfig.startIdx * gmmSwigluTensorList->tokenLen; + quantOutQueue.EnQue(quantLocal); + quantScaleOutQueue.EnQue(quantScaleLocal); +} + +} // namespace GROUPED_MATMUL +#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_SPLIT_WS_H diff --git a/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_kernel/grouped_matmul_swiglu_quant_weight_nz_tensor_list_utils.h b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_kernel/grouped_matmul_swiglu_quant_weight_nz_tensor_list_utils.h new file mode 100644 index 000000000..37ddd845d --- /dev/null +++ b/csrc/ascend/gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/op_kernel/grouped_matmul_swiglu_quant_weight_nz_tensor_list_utils.h @@ -0,0 +1,240 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file grouped_matmul_swiglu_quant_weight_nz_tensor_list_utils.h + * \brief + */ +#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_UTILS_H +#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST_UTILS_H + +#include "kernel_tiling/kernel_tiling.h" +#include "kernel_operator.h" +#include "lib/matmul_intf.h" + +namespace GROUPED_MATMUL_SWIGLU_QUANT_WEIGHT_NZ_TENSOR_LIST { +using namespace AscendC; +constexpr uint32_t INT8_BITS = 8; // a int8 number has 8 bits +constexpr uint32_t UB_BLOCK_UNIT_SIZE = 32; // 32: a block has 32 bytes data +constexpr uint32_t THRESHOLD_BLOCK_NUM = 8; +constexpr uint32_t UB_BLOCK_DOUBLE_UNIT_SIZE = 64; // 64: a block has 64 bytes data +constexpr uint32_t HALF_UB_BLOCK_UNIT_SIZE = UB_BLOCK_UNIT_SIZE / 2; // 2: a float16 data has two bytes +constexpr uint32_t SINGLE_CORE_M = 128; +constexpr uint32_t SINGLE_CORE_N = 512; +constexpr uint32_t SINGLE_CORE_K = 7168; +constexpr uint32_t BASIC_M = 128; +constexpr uint32_t BASIC_N = 256; +constexpr uint32_t BASIC_K = 128; +constexpr uint32_t STEP_M = 1; +constexpr uint32_t STEP_N = 1; +constexpr uint32_t STEP_Ka = 4; +constexpr uint32_t STEP_Kb = 4; +constexpr uint32_t DEPTH_A1 = 8; +constexpr uint32_t DEPTH_B1 = 8; +constexpr uint32_t VEC_LEN_ONCE_REPEAT_ELE = 64; +constexpr uint32_t VEC_LEN_ONCE_REPEAT_BLOCK = 8; +constexpr uint32_t BISECT = 2; +constexpr uint32_t MOD_32_MASK = 0x1F; +constexpr uint32_t MOD_16_MASK = 0x0F; +constexpr uint32_t ALIGN_8_ELE = 8; +constexpr uint32_t ALIGN_16_ELE = 16; +constexpr float QUANT_SCALE_INT8 = 127.0f; +constexpr MatmulConfig NZ_CFG_MDL = GetMDLConfig(false, false, 0, true, false, false, true); +constexpr MatmulConfig GetMMCFG() { + MatmulConfig MM_CFG = NZ_CFG_MDL; + MM_CFG.singleCoreM = SINGLE_CORE_M; + MM_CFG.singleCoreN= SINGLE_CORE_N; + MM_CFG.singleCoreK= SINGLE_CORE_K; + MM_CFG.basicM= BASIC_M; + MM_CFG.basicN= BASIC_N; + MM_CFG.basicK= BASIC_K; + return MM_CFG; +} + +constexpr static MatmulApiStaticTiling GetMMTiling(const MatmulApiStaticTiling& mmTiling) +{ + MatmulApiStaticTiling tiling = mmTiling; + tiling.stepM = STEP_M; + tiling.stepN = STEP_N; + tiling.stepKa = STEP_Ka; + tiling.stepKb = STEP_Kb; + tiling.depthA1 = DEPTH_A1; + tiling.depthB1 = DEPTH_B1; + return tiling; +} +template +struct MMImplType { + using AT = AT_; + using BT = BT_; + using CT = CT_; + using BiasT = BiasT_; + static constexpr MatmulConfig cfg = GetMMCFG(); + static constexpr MatmulApiStaticTiling mdl = GetMMTiling(GetMatmulApiTiling(cfg)); + using MT = matmul::MatmulImpl; +}; + +struct MNConfig { + int64_t m = 0; + int64_t k = 0; + int64_t n = 0; + int64_t baseM = 0; + int64_t baseN = 0; + int64_t mIdx = 0; + int64_t nIdx = 0; + int64_t blockDimM = 0; + int64_t blockDimN = 0; + int64_t singleM = 0; + int64_t singleN = 0; + int64_t wBaseOffset = 0; + int64_t nAxisBaseOffset = 0; + int64_t mAxisBaseOffset = 0; + int64_t xBaseOffset = 0; + int64_t yBaseOffset = 0; + int64_t wOutOffset = 0; + int64_t workSpaceOffset = 0; +}; + +struct VecConfig { + int64_t M = 0; + int64_t usedCoreNum = 0; + int64_t startOffset = 0; + int64_t curOffset = 0; + int64_t startIdx = 0; + int64_t curIdx = 0; + int64_t taskNum = 0; + int64_t curGroupIdx = 0; + int64_t outLoopNum = 0; + int64_t innerLoopNum = 0; + int64_t tailLoopNum = 0; + int64_t nextUpadteInterVal = 0; +}; + +struct WorkSpaceSplitConfig { + int64_t M = 0; + int64_t loopCount = 0; + int64_t leftMatrixStartIndex = 0; + int64_t rightMatrixExpertStartIndex = 0; + int64_t rightMatrixExpertNextStartIndex = 0; + int64_t rightMatrixExpertEndIndex = 0; + int64_t notLastTaskSize = 0; + int64_t lastLoopTaskSize = 0; + bool isLastLoop = false; +}; + +template +__aicore__ inline T AlignUp(T a) { + return (a + base - 1) / base * base; +} + +template +__aicore__ inline T AlignUp(T a, T base) { + return (a + base - 1) / base * base; +} + +template +__aicore__ inline T AlignDown(T a, T base) { + if (unlikely(base == 0)) { + return a; + } + return a / base * base; +} + +template <> +__aicore__ inline uint32_t AlignUp<4, uint32_t>(uint32_t a) { + // to be Multiple of 4, result should be in a format of b(xxxx,x100). + // This means last two bits should be zero, requiring that + // result = num & b(1111,1100) = num & (~3). + // &(~3) operator may reduces num into the range [num, num - 3]. + // As the result should be no less than a (result >= a), it means num - 3 >= a in the worst case. + // In this case, num >= a+3. On the other hand, num should also be less then a+4, otherwise, + // the result will not be least multiple of 4 for 3. In other cases like [num, num - 2], + // num = a + 3 also satisfies the goal condition. + return (a + 3) & ~3; // & ~3: set last two bits of (a+3) to be zero +} + +template <> +__aicore__ inline uint32_t AlignUp<8, uint32_t>(uint32_t a) { + // In general, if we want to get the least multiple of b (b is the power of 2) for a, + // it comes to a conclusion from the above comment: result = (a + (b - 1)) & (~b) + return (a + 7) & ~7; // & ~7: set last four bits of (a+7) to be zero +} + +template <> +__aicore__ inline uint32_t AlignUp<16, uint32_t>(uint32_t a) { + // In general, if we want to get the least multiple of b (b is the power of 2) for a, + // it comes to a conclusion from the above comment: result = (a + (b - 1)) & (~b) + return (a + 15) & ~15; // & ~15: set last four bits of (a+15) to be zero +} + +template <> +__aicore__ inline uint32_t AlignUp<32, uint32_t>(uint32_t a) { + // refer to the above comments. + return (a + 31) & ~31; // & ~31: set last five bits of (a+31) to be zero} +} + +__aicore__ inline void ReduceMaxTemplate(LocalTensor& dstLocal, LocalTensor& srcLocal, + uint32_t srcOffset, uint32_t count) +{ + if (likely(count > VEC_LEN_ONCE_REPEAT_ELE && count % VEC_LEN_ONCE_REPEAT_ELE == 0)){ + WholeReduceMax(dstLocal, + srcLocal[srcOffset], VEC_LEN_ONCE_REPEAT_ELE, + count / VEC_LEN_ONCE_REPEAT_ELE, 1, 1, + VEC_LEN_ONCE_REPEAT_BLOCK, ReduceOrder::ORDER_ONLY_VALUE); + PipeBarrier(); + WholeReduceMax(dstLocal, dstLocal, + count / VEC_LEN_ONCE_REPEAT_ELE, 1, 1, 1, + VEC_LEN_ONCE_REPEAT_BLOCK, ReduceOrder::ORDER_ONLY_VALUE); + } else if (count <= VEC_LEN_ONCE_REPEAT_ELE) { + WholeReduceMax(dstLocal, + srcLocal[srcOffset], + count, 1, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK, ReduceOrder::ORDER_ONLY_VALUE); + } else { + ReduceMax(dstLocal, srcLocal[srcOffset], dstLocal, count, false); + } +} + +__aicore__ inline void CastFp32ToInt8Template(LocalTensor& dstLocal, LocalTensor& srcLocal, + LocalTensor& oneBlockWorkspace, + int32_t dstOffset, int32_t srcOffset, int32_t count) +{ + Cast(srcLocal[srcOffset].ReinterpretCast(), srcLocal[srcOffset], RoundMode::CAST_RINT, count); + PipeBarrier(); + if ((dstOffset & MOD_32_MASK) == 0) { + Cast(dstLocal[dstOffset], + srcLocal[srcOffset].ReinterpretCast(), + RoundMode::CAST_RINT, count); + } else if ((dstOffset & MOD_16_MASK) == 0) { + Cast(dstLocal[dstOffset + ALIGN_16_ELE], + srcLocal[srcOffset + ALIGN_8_ELE].ReinterpretCast(), + RoundMode::CAST_RINT, count - ALIGN_16_ELE); + PipeBarrier(); + Cast(oneBlockWorkspace, srcLocal[srcOffset].ReinterpretCast(), + RoundMode::CAST_RINT, ALIGN_16_ELE); + PipeBarrier(); + for (int32_t i = 0; i < ALIGN_16_ELE; i++) { + int8_t temp = oneBlockWorkspace.GetValue(i); + dstLocal.SetValue(dstOffset + i, temp); + } + PipeBarrier(); + } +} + +template +__aicore__ inline __gm__ T* GetTensorAddr(uint16_t index, GM_ADDR tensorPtr) { + __gm__ uint64_t* dataAddr = reinterpret_cast<__gm__ uint64_t*>(tensorPtr); + uint64_t tensorPtrOffset = *dataAddr; // The offset of the data address from the first address. + // Moving 3 bits to the right means dividing by sizeof(uint64 t). + __gm__ uint64_t* retPtr = dataAddr + (tensorPtrOffset >> 3); + return reinterpret_cast<__gm__ T*>(*(retPtr + index)); +} + +} // namespace GROUPED_MATMUL + +#endif // ASCENDC_GROUPED_MATMUL_UTILS_H diff --git a/csrc/ascend/kernels/bgmv_expand.cpp b/csrc/ascend/kernels/bgmv_expand.cpp new file mode 100644 index 000000000..28df2d758 --- /dev/null +++ b/csrc/ascend/kernels/bgmv_expand.cpp @@ -0,0 +1,369 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kernel_operator.h" +#include "types.h" + +template +class BGMVExpand { +public: + using X_T = float; + using W_T = scalar_t; + using Y_T = scalar_t; + + static constexpr uint64_t LORA_RANK_8 = 8; + static constexpr uint64_t LORA_RANK_16 = 16; + static constexpr uint64_t LORA_RANK_32 = 32; + static constexpr uint64_t LORA_RANK_64 = 64; + static constexpr uint64_t SUPPORTED_RANKS[] = {LORA_RANK_8, LORA_RANK_16, LORA_RANK_32, LORA_RANK_64}; + static constexpr int32_t BUFFER_NUM = 2; + + // The vector unit reads 8 blocks (32 bytes each and 256 bytes in total) of contiguous data each time. + static constexpr int32_t NUM_BYTES_PER_REPEAT = 256; + static constexpr int32_t NUM_BLOCKS_PER_REPEAT = 8; + // The maximum number of elements in a single iteration is 256 / sizeof(intermediate data type). + static constexpr int32_t NUM_ELEMENTS_PER_REPEAT = NUM_BYTES_PER_REPEAT / sizeof(float); + // Mask is used to control the elements that participate in computation in each iteration. + static constexpr int32_t MASK_COUNT = NUM_BYTES_PER_REPEAT / sizeof(float); + // Refer to numOutputElementsPerInputTile_ initialization for the constraints on the following constants. + static constexpr int32_t W_IN_TILE_NUM_ELEMENTS = 8192; + static constexpr int32_t Y_OUT_TILE_NUM_ELEMENTS = 4096; + static constexpr int32_t BLOCK_REDUCE_NUM_REPEATS = W_IN_TILE_NUM_ELEMENTS / NUM_ELEMENTS_PER_REPEAT; + // BlockReduceSum would generate(BLOCK_REDUCE_NUM_REPEATS * NUM_BLOCKS_PER_REPEAT)floats. + // So need to read them all and apply PairReduceSum + static constexpr int32_t PAIR_REDUCE_NUM_REPEATS_16 = + (BLOCK_REDUCE_NUM_REPEATS * NUM_BLOCKS_PER_REPEAT + NUM_ELEMENTS_PER_REPEAT - 1) / NUM_ELEMENTS_PER_REPEAT; + // The second PairReduceSum for rank=32, needs half of the repetition that happened for rank=16. + // Same for rank=64, we do not support ranks greater than 64. + static constexpr int32_t PAIR_REDUCE_NUM_REPEATS_32 = (PAIR_REDUCE_NUM_REPEATS_16 + 1) / 2; + +public: + __aicore__ inline BGMVExpand(AscendC::TPipe* pipe) : pipe_(pipe) {} + + __aicore__ inline void Init(__gm__ void* x, __gm__ void* weight, __gm__ void* indices, + uint32_t indicesSize, __gm__ void* yIn, __gm__ void* yOut, + uint32_t batchSize, uint32_t numTokensPerCore, uint32_t maxLoRARank, + uint32_t outputHiddenDim, uint32_t sliceOffset, uint32_t outputFullDim) + { + batchSize_ = batchSize; + numTokensPerCore_ = numTokensPerCore; + maxLoRARank_ = maxLoRARank; + outputHiddenDim_ = outputHiddenDim; + sliceOffset_ = sliceOffset; + outputFullDim_ = outputFullDim; + singleLoRAWeightLen_ = maxLoRARank_ * outputHiddenDim_; + + xGm_.SetGlobalBuffer((__gm__ X_T *)x); + wGm_.SetGlobalBuffer((__gm__ W_T *)weight); + yInGm_.SetGlobalBuffer((__gm__ Y_T *)yIn); + yOutGm_.SetGlobalBuffer((__gm__ Y_T *)yOut); + indicesGm_.SetGlobalBuffer((__gm__ int64_t *)indices, indicesSize); + + pipe_->InitBuffer(inQueueX_, 1, NUM_ELEMENTS_PER_REPEAT * sizeof(X_T)); + pipe_->InitBuffer(inQueueW_, BUFFER_NUM, W_IN_TILE_NUM_ELEMENTS * sizeof(W_T)); + pipe_->InitBuffer(inQueueY_, BUFFER_NUM, Y_OUT_TILE_NUM_ELEMENTS * sizeof(Y_T)); + pipe_->InitBuffer(outQueueY_, BUFFER_NUM, Y_OUT_TILE_NUM_ELEMENTS * sizeof(Y_T)); + + pipe_->InitBuffer(dupBufferX_, NUM_ELEMENTS_PER_REPEAT * sizeof(float)); + pipe_->InitBuffer(tmpBufferW_, W_IN_TILE_NUM_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(inBufferY_, Y_OUT_TILE_NUM_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(tmpBufferY_, Y_OUT_TILE_NUM_ELEMENTS * sizeof(float)); + + // Each compute iteration would generate not one, but several output elements. + // Therefore, the following variable would determine how many output elements are calculated in each iteration. + numOutputElementsPerInputTile_ = BLOCK_REDUCE_NUM_REPEATS * (NUM_ELEMENTS_PER_REPEAT / maxLoRARank_); + numStreamInPerOutputTile_ = Y_OUT_TILE_NUM_ELEMENTS / numOutputElementsPerInputTile_; + + } + + __aicore__ inline void Process() + { + int64_t blockIdx = AscendC::GetBlockIdx(); + int64_t startIdx = blockIdx * numTokensPerCore_; + int64_t endIdx = startIdx + numTokensPerCore_; + if (endIdx > batchSize_) { + endIdx = batchSize_; + } + for (int64_t idx = startIdx; idx < endIdx; idx++) { + yOffset_ = outputFullDim_ * idx + sliceOffset_; + + // Set up LoRA index + CopyInIndex(idx); + if (reqLoRAIndex_ < 0) { + continue; + } + reqLoRAWeightOffset_ = reqLoRAIndex_ * singleLoRAWeightLen_; + + CopyInX(idx); + int32_t numStreamOut = outputHiddenDim_ / Y_OUT_TILE_NUM_ELEMENTS; + for (int32_t i = 0; i < numStreamOut; i++) { + CopyInY(i); + for (int32_t j = 0; j < numStreamInPerOutputTile_; j++) { + CopyInW(i * numStreamInPerOutputTile_ + j); + Compute(j * numOutputElementsPerInputTile_); + } + ScaleOutput(); + CopyOut(i); + } + ComputeLastIteration(); + } + } + +private: + __aicore__ inline void CopyInIndex(const int64_t idx) + { + // Look up the LoRA index + reqLoRAIndex_ = indicesGm_.GetValue(idx); + } + + __aicore__ inline void ComputeLastIteration() + { + int32_t remainingY = outputHiddenDim_ % Y_OUT_TILE_NUM_ELEMENTS; + if (remainingY == 0) { + return; + } + int32_t numStreamOut = outputHiddenDim_ / Y_OUT_TILE_NUM_ELEMENTS; + int32_t remainingW = remainingY * maxLoRARank_; + int32_t numCompleteWTileInForLastIteration = remainingW / W_IN_TILE_NUM_ELEMENTS; + int32_t remainingWForLastRepeat = remainingW % W_IN_TILE_NUM_ELEMENTS; + + CopyInY(numStreamOut, remainingY); + + int32_t outputIdx = 0; + for (outputIdx = 0; outputIdx < numCompleteWTileInForLastIteration; outputIdx++) { + CopyInW(numStreamOut * numStreamInPerOutputTile_ + outputIdx); + Compute(outputIdx * numOutputElementsPerInputTile_); + } + + if (remainingWForLastRepeat != 0) { + CopyInW(numStreamOut * numStreamInPerOutputTile_ + numCompleteWTileInForLastIteration, + remainingWForLastRepeat); + int32_t lastRepeatCount = remainingWForLastRepeat / NUM_ELEMENTS_PER_REPEAT; + int32_t pairReduceRepeat16 = + (lastRepeatCount * NUM_BLOCKS_PER_REPEAT + NUM_ELEMENTS_PER_REPEAT - 1) / NUM_ELEMENTS_PER_REPEAT; + int32_t pairReduceRepeat32 = (pairReduceRepeat16 + 1) / 2; + int32_t lastComputeOutputElement = outputIdx * numOutputElementsPerInputTile_; + Compute(lastComputeOutputElement, lastRepeatCount, pairReduceRepeat16, pairReduceRepeat32); + } + + ScaleOutput(remainingY); + CopyOut(numStreamOut, remainingY); + } + + __aicore__ inline void CopyInX(const int64_t idx) + { + AscendC::LocalTensor xLocal = inQueueX_.AllocTensor(); + if constexpr (std::is_same_v) { + DataCopy(xLocal, xGm_[maxLoRARank_ * idx], maxLoRARank_); + } else { + uint16_t blockLen = static_cast(maxLoRARank_ * sizeof(X_T)); + DataCopyPad(xLocal, xGm_[maxLoRARank_ * idx], {1, blockLen, 0, 0}, {}); + } + inQueueX_.EnQue(xLocal); + xLocal = inQueueX_.DeQue(); + AscendC::LocalTensor xDup = dupBufferX_.Get(); + + // As we are generating multiple output elements with one API invocation, + // we need to duplicate the X vector multiple times to fill one NUM_BYTES_PER_REPEAT + if constexpr (std::is_same_v) { + for (int32_t i = 0; i < NUM_ELEMENTS_PER_REPEAT; i += maxLoRARank_) { + for (int32_t j = 0; j < maxLoRARank_; j++) { + float entry = xLocal.GetValue(j); + xDup.SetValue(i + j, entry); + } + } + } else { + Cast(xDup, xLocal, AscendC::RoundMode::CAST_NONE, maxLoRARank_); + AscendC::PipeBarrier(); + + for (int32_t i = maxLoRARank_; i < NUM_ELEMENTS_PER_REPEAT; i += maxLoRARank_) { + for (int32_t j = 0; j < maxLoRARank_; j++) { + float entry = xDup.GetValue(j); + xDup.SetValue(i + j, entry); + } + } + } + inQueueX_.FreeTensor(xLocal); + } + + __aicore__ inline void CopyInY(int32_t progress, int32_t numElements = Y_OUT_TILE_NUM_ELEMENTS) + { + AscendC::LocalTensor yInLocal = inQueueY_.AllocTensor(); + DataCopy(yInLocal, yInGm_[yOffset_ + progress * Y_OUT_TILE_NUM_ELEMENTS], numElements); + inQueueY_.EnQue(yInLocal); + } + + __aicore__ inline void CopyInW(int32_t progress, int32_t numElements = W_IN_TILE_NUM_ELEMENTS) + { + AscendC::LocalTensor wLocal = inQueueW_.AllocTensor(); + DataCopy(wLocal, wGm_[reqLoRAWeightOffset_ + progress * W_IN_TILE_NUM_ELEMENTS], numElements); + inQueueW_.EnQue(wLocal); + } + + __aicore__ inline void ScaleOutput(int32_t numElements = Y_OUT_TILE_NUM_ELEMENTS) + { + AscendC::LocalTensor yLocal = tmpBufferY_.Get(); + AscendC::LocalTensor yInLocal = inQueueY_.DeQue(); + AscendC::LocalTensor yInLocalFP32 = inBufferY_.Get(); + Cast(yInLocalFP32, yInLocal, AscendC::RoundMode::CAST_NONE, numElements); + AscendC::PipeBarrier(); + inQueueY_.FreeTensor(yInLocal); + + Add(yLocal, yLocal, yInLocalFP32, numElements); + AscendC::PipeBarrier(); + + AscendC::LocalTensor yOutLocal = outQueueY_.AllocTensor(); + Cast(yOutLocal, yLocal, AscendC::RoundMode::CAST_RINT, numElements); + AscendC::PipeBarrier(); + + outQueueY_.EnQue(yOutLocal); + } + + __aicore__ inline void Compute(int32_t progress, + int32_t blockReduceRepeatCount=BLOCK_REDUCE_NUM_REPEATS, + int32_t pairReduceRepeat16=PAIR_REDUCE_NUM_REPEATS_16, + int32_t pairReduceRepeat32=PAIR_REDUCE_NUM_REPEATS_32) + { + AscendC::LocalTensor yLocal = tmpBufferY_.Get(); + AscendC::LocalTensor xDup = dupBufferX_.Get(); + AscendC::LocalTensor wLocal = inQueueW_.DeQue(); + AscendC::LocalTensor wTmpTensor = tmpBufferW_.Get(); + + Cast(wTmpTensor, wLocal, AscendC::RoundMode::CAST_NONE, MASK_COUNT, blockReduceRepeatCount, castParams_); + AscendC::PipeBarrier(); + inQueueW_.FreeTensor(wLocal); + + Mul(wTmpTensor, xDup, wTmpTensor, MASK_COUNT, blockReduceRepeatCount, dotProductParams_); + AscendC::PipeBarrier(); + + if (maxLoRARank_ == LORA_RANK_8) { + BlockReduceSum(yLocal[progress], wTmpTensor, blockReduceRepeatCount, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + } else if (maxLoRARank_ == LORA_RANK_16) { + BlockReduceSum(wTmpTensor, wTmpTensor, blockReduceRepeatCount, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + PairReduceSum(yLocal[progress], wTmpTensor, pairReduceRepeat16, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + } else if (maxLoRARank_ == LORA_RANK_32) { + BlockReduceSum(wTmpTensor, wTmpTensor, blockReduceRepeatCount, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + PairReduceSum(wTmpTensor, wTmpTensor, pairReduceRepeat16, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + PairReduceSum(yLocal[progress], wTmpTensor, pairReduceRepeat32, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + } else if (maxLoRARank_ == LORA_RANK_64) { + BlockReduceSum(wTmpTensor, wTmpTensor, blockReduceRepeatCount, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + BlockReduceSum(yLocal[progress], wTmpTensor, pairReduceRepeat16, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + } + } + + __aicore__ inline void CopyOut(int32_t progress, int32_t numElements = Y_OUT_TILE_NUM_ELEMENTS) + { + AscendC::LocalTensor yOutLocal = outQueueY_.DeQue(); + DataCopy(yOutGm_[yOffset_ + progress * Y_OUT_TILE_NUM_ELEMENTS], yOutLocal, numElements); + outQueueY_.FreeTensor(yOutLocal); + } + +private: + AscendC::TPipe* pipe_; + AscendC::TQue inQueueY_, inQueueW_; + AscendC::TQue inQueueX_; + AscendC::TQue outQueueY_; + AscendC::TBuf tmpBufferW_, dupBufferX_, inBufferY_, tmpBufferY_; + AscendC::GlobalTensor xGm_; + AscendC::GlobalTensor wGm_; + AscendC::GlobalTensor yInGm_; + AscendC::GlobalTensor yOutGm_; + AscendC::GlobalTensor indicesGm_; + uint32_t batchSize_; + uint32_t numTokensPerCore_; + uint32_t maxLoRARank_; + uint32_t outputHiddenDim_; + uint32_t sliceOffset_; + uint32_t outputFullDim_; + uint32_t singleLoRAWeightLen_; + int64_t reqLoRAIndex_; + uint64_t reqLoRAWeightOffset_; + uint32_t numOutputElementsPerInputTile_; + uint32_t numStreamInPerOutputTile_; + uint64_t yOffset_; + + // The block stride is set to 1, and 8 blocks in the same repeat are processed continuously. + // The repeat stride is 8, so the vector unit reads 8 consecutive blocks in the first repeat, + // reads next 8 consecutive blocks in the second repeat. + AscendC::UnaryRepeatParams castParams_ = {1, 1, 8, 4}; + + // For each repeat in BlockReduceSum and PairReduceSum we should move forward only one block, + // so we set dstRepStride = 1 + AscendC::UnaryRepeatParams reduceSumParams_ = {1, 1, 1, 8}; + + // When the repeat stride is 0, the vector unit repeatedly reads and computes the first 8 consecutive blocks. + // For xDup we repeatedly use it, so we set src0RepStride = 0 + AscendC::BinaryRepeatParams dotProductParams_ = {1, 1, 1, 8, 0, 8}; + +}; + +#define BGMV_EXPAND_TYPE_DECLARE(TYPE) \ + extern "C" __global__ __aicore__ void bgmv_expand_##TYPE(__gm__ void* x, __gm__ void* weight, __gm__ void* indices,\ + uint32_t indicesSize, __gm__ void* yIn, __gm__ void* yOut,\ + uint32_t batchSize, uint32_t numTokensPerCore, \ + uint32_t maxLoRARank, uint32_t outputHiddenDim, \ + uint32_t sliceOffset, uint32_t outputFullDim) \ + { \ + AscendC::TPipe pipe; \ + BGMVExpand op(&pipe); \ + op.Init(x, weight, indices, indicesSize, yIn, yOut, batchSize, numTokensPerCore, maxLoRARank, \ + outputHiddenDim, sliceOffset, outputFullDim); \ + op.Process(); \ + } + +// declare all dtype kernel +BGMV_EXPAND_TYPE_DECLARE(half) +#if !defined(__CCE_AICORE__) || (__CCE_AICORE__ >= 220) + BGMV_EXPAND_TYPE_DECLARE(bfloat16_t) +#endif + +namespace vllm_fl { +extern void bgmv_expand_impl(AscendType type, void* stream, void* x, void* weight, void* indices, uint32_t indicesSize, + void* yIn, void* yOut, uint32_t batchSize, uint32_t numTokensPerCore, uint32_t maxLoRARank, + uint32_t outputHiddenDim, uint32_t sliceOffset, uint32_t outputFullDim) +{ + uint32_t blockDim = (batchSize + numTokensPerCore - 1) / numTokensPerCore; + if (type == AscendType::FP16) { + bgmv_expand_half<<>>(x, weight, indices, indicesSize, yIn, yOut, batchSize, numTokensPerCore, + maxLoRARank, outputHiddenDim, sliceOffset, outputFullDim); + } else if (type == AscendType::BF16) { + #if !defined(__CCE_AICORE__) || (__CCE_AICORE__ >= 220) + bgmv_expand_bfloat16_t<<>>(x, weight, indices, indicesSize, yIn, yOut, batchSize, + numTokensPerCore, maxLoRARank, outputHiddenDim, + sliceOffset, outputFullDim); + #endif + } else { + return; + } +} + +} // namespace vllm_fl \ No newline at end of file diff --git a/csrc/ascend/kernels/bgmv_shrink.cpp b/csrc/ascend/kernels/bgmv_shrink.cpp new file mode 100644 index 000000000..f7513fd10 --- /dev/null +++ b/csrc/ascend/kernels/bgmv_shrink.cpp @@ -0,0 +1,252 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kernel_operator.h" +#include "types.h" + +template +class BGMVShrink { +public: + using X_T = scalar_t; + using W_T = scalar_t; + using Y_T = float; + + static constexpr uint64_t BUFFER_NUM = 1; + static constexpr uint64_t TILE_LENGTH = 11776; // optimal performance tile length + +public: + __aicore__ inline BGMVShrink(AscendC::TPipe *pipe) : pipe_(pipe) {} + __aicore__ inline void Init(__gm__ void *x, __gm__ void *weight, __gm__ void *indices, uint32_t indicesSize, __gm__ void *y, + uint32_t batchSize, uint32_t numTokensPerCore, uint32_t inputHiddenDim, + uint32_t maxLoRARank, float scale) + { + batchSize_ = batchSize; + numTokensPerCore_ = numTokensPerCore; + inputHiddenDim_ = inputHiddenDim; + maxLoRARank_ = maxLoRARank; + scale_ = scale; + singleLoRAWeightLen_ = inputHiddenDim_ * maxLoRARank_; + incremental_ = inputHiddenDim_ > TILE_LENGTH; + + xGm_.SetGlobalBuffer((__gm__ X_T *)x); + yOutGm_.SetGlobalBuffer((__gm__ Y_T *)y); + wGm_.SetGlobalBuffer((__gm__ W_T *)weight); + indicesGm_.SetGlobalBuffer((__gm__ int64_t *)indices, indicesSize); + + pipe_->InitBuffer(inQueueX_, BUFFER_NUM, TILE_LENGTH * sizeof(X_T)); + pipe_->InitBuffer(inQueueW_, BUFFER_NUM, TILE_LENGTH * sizeof(W_T)); + pipe_->InitBuffer(tmpBufferX_, TILE_LENGTH * sizeof(float)); + pipe_->InitBuffer(tmpBufferW_, TILE_LENGTH * sizeof(float)); + + pipe_->InitBuffer(outQueueY_, 1, maxLoRARank_ * sizeof(Y_T)); + pipe_->InitBuffer(outBufferY_, maxLoRARank_ * sizeof(float)); + } + + __aicore__ inline void Process() + { + int64_t blockIdx = AscendC::GetBlockIdx(); + int64_t startIdx = blockIdx * numTokensPerCore_; + int64_t endIdx = startIdx + numTokensPerCore_; + if (endIdx > batchSize_) { + endIdx = batchSize_; + } + for (int64_t idx = startIdx; idx < endIdx; idx++) { + // set up LoRA index + CopyInIndex(idx); + if (reqLoRAIndex_ < 0) { + continue; + } + reqLoRAWeightOffset_ = reqLoRAIndex_ * singleLoRAWeightLen_; + + if (incremental_) { + ProcessImpl(idx); + } else { + ProcessImpl(idx); + } + + ScaleOutput(); + CopyOut(idx); + } + } + +private: + template + __aicore__ inline void ProcessImpl(const int64_t idx) + { + AscendC::LocalTensor yOutLocal = outBufferY_.Get(); + if constexpr (!INCREMENTAL_MODE) { + CopyInX(idx, 0, inputHiddenDim_); + AscendC::LocalTensor xTmpTensor = tmpBufferX_.Get(); + AscendC::LocalTensor xLocal = inQueueX_.DeQue(); + Cast(xTmpTensor, xLocal, AscendC::RoundMode::CAST_NONE, inputHiddenDim_); + AscendC::PipeBarrier(); + inQueueX_.FreeTensor(xLocal); + } + + for (int i = 0; i < maxLoRARank_; i++) { + float acc(0); + for (int32_t j = 0; j < inputHiddenDim_ / TILE_LENGTH; j++) { + if constexpr (INCREMENTAL_MODE) { + CopyInX(idx, j); + } + CopyInW(i, j); + Compute(acc); + } + CopyAndComputeLastIteration(idx, i, acc); + yOutLocal.SetValue(i, acc); + } + } + + __aicore__ inline void CopyInIndex(const int64_t idx) + { + // look up the LoRA index + reqLoRAIndex_ = indicesGm_.GetValue(idx); + } + + __aicore__ inline void CopyInX(const int64_t idx, int32_t colIdx, int32_t numElements = TILE_LENGTH) + { + AscendC::LocalTensor xLocal = inQueueX_.AllocTensor(); + DataCopy(xLocal, xGm_[inputHiddenDim_ * idx + colIdx * TILE_LENGTH], numElements); + inQueueX_.EnQue(xLocal); + } + + __aicore__ inline void CopyInW(int32_t rowIdx, int32_t colIdx, int32_t numElements = TILE_LENGTH) + { + AscendC::LocalTensor wLocal = inQueueW_.AllocTensor(); + DataCopy(wLocal, wGm_[reqLoRAWeightOffset_ + rowIdx * inputHiddenDim_ + colIdx * TILE_LENGTH], numElements); + inQueueW_.EnQue(wLocal); + } + + template + __aicore__ inline void Compute(float &acc, int32_t numElements = TILE_LENGTH) + { + AscendC::LocalTensor wLocal = inQueueW_.DeQue(); + AscendC::LocalTensor xTmpTensor = tmpBufferX_.Get(); + AscendC::LocalTensor wTmpTensor = tmpBufferW_.Get(); + + if constexpr (INCREMENTAL_MODE) { + AscendC::LocalTensor xLocal = inQueueX_.DeQue(); + Cast(xTmpTensor, xLocal, AscendC::RoundMode::CAST_NONE, numElements); + Cast(wTmpTensor, wLocal, AscendC::RoundMode::CAST_NONE, numElements); + AscendC::PipeBarrier(); + inQueueX_.FreeTensor(xLocal); + inQueueW_.FreeTensor(wLocal); + } else { + Cast(wTmpTensor, wLocal, AscendC::RoundMode::CAST_NONE, numElements); + AscendC::PipeBarrier(); + inQueueW_.FreeTensor(wLocal); + } + // dot product of the one tile of X and W + Mul(wTmpTensor, xTmpTensor, wTmpTensor, numElements); + AscendC::PipeBarrier(); + // reduce sum generate one number, which is the summation of all the dot product + ReduceSum(wTmpTensor, wTmpTensor, wTmpTensor, numElements); + AscendC::PipeBarrier(); + + acc += wTmpTensor.GetValue(0); + } + + template + __aicore__ inline void CopyAndComputeLastIteration(const int64_t idx, int32_t rowIdx, float &acc) + { + int32_t colIdx = inputHiddenDim_ / TILE_LENGTH; + int32_t remaining = inputHiddenDim_ % TILE_LENGTH; + if (remaining == 0) { + return; + } + if constexpr (INCREMENTAL_MODE) { + CopyInX(idx, colIdx, remaining); + } + CopyInW(rowIdx, colIdx, remaining); + Compute(acc, remaining); + } + + __aicore__ inline void ScaleOutput() + { + AscendC::LocalTensor yLocal = outBufferY_.Get(); + AscendC::LocalTensor yOutLocal = outQueueY_.AllocTensor(); + + Muls(yOutLocal, yLocal, scale_, maxLoRARank_); + AscendC::PipeBarrier(); + + outQueueY_.EnQue(yOutLocal); + } + + __aicore__ inline void CopyOut(const int64_t idx) + { + AscendC::LocalTensor yOutLocal = outQueueY_.DeQue(); + DataCopy(yOutGm_[maxLoRARank_ * idx], yOutLocal, maxLoRARank_); + outQueueY_.FreeTensor(yOutLocal); + } + +private: + AscendC::TPipe *pipe_; + AscendC::TQue inQueueX_, inQueueW_; + AscendC::TQue outQueueY_; + AscendC::TBuf tmpBufferX_, tmpBufferW_, outBufferY_; + AscendC::GlobalTensor xGm_; + AscendC::GlobalTensor wGm_; + AscendC::GlobalTensor indicesGm_; + AscendC::GlobalTensor yOutGm_; + uint32_t batchSize_; + uint32_t numTokensPerCore_; + uint32_t inputHiddenDim_; + uint32_t maxLoRARank_; + float scale_; + uint32_t singleLoRAWeightLen_; + int64_t reqLoRAIndex_; + uint64_t reqLoRAWeightOffset_; + bool incremental_; +}; + +#define BGMV_SHRINK_TYPE_DECLARE(TYPE) \ + extern "C" __global__ __aicore__ void bgmv_shrink_##TYPE(__gm__ void* x, __gm__ void* weight, __gm__ void* indices,\ + uint32_t indicesSize, __gm__ void* y, uint32_t batchSize, \ + uint32_t numTokensPerCore, uint32_t inputHiddenDim, \ + uint32_t maxLoRARank, float scale) \ + { \ + AscendC::TPipe pipe; \ + BGMVShrink op(&pipe); \ + op.Init(x, weight, indices, indicesSize, y, batchSize, numTokensPerCore, inputHiddenDim, maxLoRARank, scale); \ + op.Process(); \ + } + +// declare all dtype kernel +BGMV_SHRINK_TYPE_DECLARE(half) +#if !defined(__CCE_AICORE__) || (__CCE_AICORE__ >= 220) + BGMV_SHRINK_TYPE_DECLARE(bfloat16_t) +#endif + +namespace vllm_fl { +extern void bgmv_shrink_impl(AscendType type, void* stream, void* x, void* weight, void* indices, uint32_t indicesSize, + void* y, uint32_t batchSize, uint32_t numTokensPerCore, uint32_t inputHiddenDim, + uint32_t maxLoRARank, float scale) +{ + uint32_t blockDim = (batchSize + numTokensPerCore - 1) / numTokensPerCore; + if (type == AscendType::FP16) { + bgmv_shrink_half<<>>(x, weight, indices, indicesSize, y, batchSize, numTokensPerCore, + inputHiddenDim, maxLoRARank, scale); + } else if (type == AscendType::BF16) { + #if !defined(__CCE_AICORE__) || (__CCE_AICORE__ >= 220) + bgmv_shrink_bfloat16_t<<>>(x, weight, indices, indicesSize, y, batchSize, numTokensPerCore, + inputHiddenDim, maxLoRARank, scale); + #endif + } else { + return; + } +} + +} // namespace vllm_fl \ No newline at end of file diff --git a/csrc/ascend/kernels/math_utils.h b/csrc/ascend/kernels/math_utils.h new file mode 100644 index 000000000..62b46921c --- /dev/null +++ b/csrc/ascend/kernels/math_utils.h @@ -0,0 +1,15 @@ +#ifndef KERNEL_MATH_UTILS_H +#define KERNEL_MATH_UTILS_H +#include + +namespace device_utils { + +template +__aicore__ __force_inline__ T RoundUp(const T &val) +{ + return (val + roundVal - 1) / roundVal * roundVal; +} + +}; // namespace device_utils + +#endif diff --git a/csrc/ascend/kernels/sgmv_expand.cpp b/csrc/ascend/kernels/sgmv_expand.cpp new file mode 100644 index 000000000..a61ef74f0 --- /dev/null +++ b/csrc/ascend/kernels/sgmv_expand.cpp @@ -0,0 +1,389 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kernel_operator.h" +#include "types.h" + +template +class SGMVExpand { +public: + using X_T = float; + using W_T = scalar_t; + using Y_T = scalar_t; + + static constexpr uint64_t LORA_RANK_8 = 8; + static constexpr uint64_t LORA_RANK_16 = 16; + static constexpr uint64_t LORA_RANK_32 = 32; + static constexpr uint64_t LORA_RANK_64 = 64; + static constexpr uint64_t SUPPORTED_RANKS[] = {LORA_RANK_8, LORA_RANK_16, LORA_RANK_32, LORA_RANK_64}; + static constexpr int32_t BUFFER_NUM = 2; + + // The vector unit reads 8 blocks (32 bytes each and 256 bytes in total) of contiguous data each time. + static constexpr int32_t NUM_BYTES_PER_REPEAT = 256; + static constexpr int32_t NUM_BLOCKS_PER_REPEAT = 8; + // The maximum number of elements in a single iteration is 256 / sizeof(intermediate data type). + static constexpr int32_t NUM_ELEMENTS_PER_REPEAT = NUM_BYTES_PER_REPEAT / sizeof(float); + // Mask is used to control the elements that participate in computation in each iteration. + static constexpr int32_t MASK_COUNT = NUM_BYTES_PER_REPEAT / sizeof(float); + // Refer to numOutputElementsPerInputTile_ initialization for the constraints on the following constants. + static constexpr int32_t W_IN_TILE_NUM_ELEMENTS = 8192; + static constexpr int32_t Y_OUT_TILE_NUM_ELEMENTS = 4096; + static constexpr int32_t BLOCK_REDUCE_NUM_REPEATS = W_IN_TILE_NUM_ELEMENTS / NUM_ELEMENTS_PER_REPEAT; + // BlockReduceSum would generate(BLOCK_REDUCE_NUM_REPEATS * NUM_BLOCKS_PER_REPEAT)floats. + // So need to read them all and apply PairReduceSum + static constexpr int32_t PAIR_REDUCE_NUM_REPEATS_16 = + (BLOCK_REDUCE_NUM_REPEATS * NUM_BLOCKS_PER_REPEAT + NUM_ELEMENTS_PER_REPEAT - 1) / NUM_ELEMENTS_PER_REPEAT; + // The second PairReduceSum for rank=32, needs half of the repetition that happened for rank=16. + // Same for rank=64, we do not support ranks greater than 64. + static constexpr int32_t PAIR_REDUCE_NUM_REPEATS_32 = (PAIR_REDUCE_NUM_REPEATS_16 + 1) / 2; + +public: + __aicore__ inline SGMVExpand(AscendC::TPipe* pipe) : pipe_(pipe) {} + + __aicore__ inline void Init(__gm__ void* x, __gm__ void* weight, __gm__ void* loraIndices, uint32_t loraIndicesSize, + __gm__ void* seqLen, uint32_t seqLenSize, __gm__ void* yIn, __gm__ void* yOut, + uint32_t batchSize, uint32_t numTokensPerCore, uint32_t maxLoRARank, + uint32_t outputHiddenDim, uint32_t sliceOffset, uint32_t outputFullDim) + { + batchSize_ = batchSize; + numTokensPerCore_ = numTokensPerCore; + maxLoRARank_ = maxLoRARank; + outputHiddenDim_ = outputHiddenDim; + sliceOffset_ = sliceOffset; + outputFullDim_ = outputFullDim; + singleLoRAWeightLen_ = maxLoRARank_ * outputHiddenDim_; + + xGm_.SetGlobalBuffer((__gm__ X_T *)x); + wGm_.SetGlobalBuffer((__gm__ W_T *)weight); + yInGm_.SetGlobalBuffer((__gm__ Y_T *)yIn); + yOutGm_.SetGlobalBuffer((__gm__ Y_T *)yOut); + loraIndicesGm_.SetGlobalBuffer((__gm__ int64_t *)loraIndices, loraIndicesSize); + seqLenGm_.SetGlobalBuffer((__gm__ int64_t *)seqLen, seqLenSize); + + pipe_->InitBuffer(inQueueX_, 1, NUM_ELEMENTS_PER_REPEAT * sizeof(X_T)); + pipe_->InitBuffer(inQueueW_, BUFFER_NUM, W_IN_TILE_NUM_ELEMENTS * sizeof(W_T)); + pipe_->InitBuffer(inQueueY_, BUFFER_NUM, Y_OUT_TILE_NUM_ELEMENTS * sizeof(Y_T)); + pipe_->InitBuffer(outQueueY_, BUFFER_NUM, Y_OUT_TILE_NUM_ELEMENTS * sizeof(Y_T)); + + pipe_->InitBuffer(dupBufferX_, NUM_ELEMENTS_PER_REPEAT * sizeof(float)); + pipe_->InitBuffer(tmpBufferW_, W_IN_TILE_NUM_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(inBufferY_, Y_OUT_TILE_NUM_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(tmpBufferY_, Y_OUT_TILE_NUM_ELEMENTS * sizeof(float)); + + // Each compute iteration would generate not one, but several output elements. + // Therefore, the following variable would determine how many output elements are calculated in each iteration. + numOutputElementsPerInputTile_ = BLOCK_REDUCE_NUM_REPEATS * (NUM_ELEMENTS_PER_REPEAT / maxLoRARank_); + numStreamInPerOutputTile_ = Y_OUT_TILE_NUM_ELEMENTS / numOutputElementsPerInputTile_; + + } + + __aicore__ inline void Process() + { + int64_t blockIdx = AscendC::GetBlockIdx(); + int64_t startIdx = blockIdx * numTokensPerCore_; + int64_t endIdx = startIdx + numTokensPerCore_; + if (endIdx > batchSize_) { + endIdx = batchSize_; + } + for (int64_t idx = startIdx; idx < endIdx; idx++) { + yOffset_ = outputFullDim_ * idx + sliceOffset_; + + // Set up LoRA index + CopyInIndex(idx); + if (reqLoRAIndex_ < 0) { + continue; + } + reqLoRAWeightOffset_ = reqLoRAIndex_ * singleLoRAWeightLen_; + + CopyInX(idx); + int32_t numStreamOut = outputHiddenDim_ / Y_OUT_TILE_NUM_ELEMENTS; + for (int32_t i = 0; i < numStreamOut; i++) { + CopyInY(i); + for (int32_t j = 0; j < numStreamInPerOutputTile_; j++) { + CopyInW(i * numStreamInPerOutputTile_ + j); + Compute(j * numOutputElementsPerInputTile_); + } + ScaleOutput(); + CopyOut(i); + } + ComputeLastIteration(); + } + } + +private: + __aicore__ inline void CopyInIndex(const int64_t idx) + { + // Look up the LoRA index + int64_t weightIdx = idx; + uint64_t i = 0; + for (; i < seqLenGm_.GetSize(); i++) { + int64_t repeatValue = seqLenGm_.GetValue(i); + if (weightIdx >= repeatValue) { + weightIdx -= repeatValue; + continue; + } + break; + } + reqLoRAIndex_ = (i < seqLenGm_.GetSize()) ? loraIndicesGm_.GetValue(i) : -1; + } + + __aicore__ inline void ComputeLastIteration() + { + int32_t remainingY = outputHiddenDim_ % Y_OUT_TILE_NUM_ELEMENTS; + if (remainingY == 0) { + return; + } + int32_t numStreamOut = outputHiddenDim_ / Y_OUT_TILE_NUM_ELEMENTS; + int32_t remainingW = remainingY * maxLoRARank_; + int32_t numCompleteWTileInForLastIteration = remainingW / W_IN_TILE_NUM_ELEMENTS; + int32_t remainingWForLastRepeat = remainingW % W_IN_TILE_NUM_ELEMENTS; + + CopyInY(numStreamOut, remainingY); + + int32_t outputIdx = 0; + for (outputIdx = 0; outputIdx < numCompleteWTileInForLastIteration; outputIdx++) { + CopyInW(numStreamOut * numStreamInPerOutputTile_ + outputIdx); + Compute(outputIdx * numOutputElementsPerInputTile_); + } + + if (remainingWForLastRepeat != 0) { + CopyInW(numStreamOut * numStreamInPerOutputTile_ + numCompleteWTileInForLastIteration, + remainingWForLastRepeat); + int32_t lastRepeatCount = remainingWForLastRepeat / NUM_ELEMENTS_PER_REPEAT; + int32_t pairReduceRepeat16 = + (lastRepeatCount * NUM_BLOCKS_PER_REPEAT + NUM_ELEMENTS_PER_REPEAT - 1) / NUM_ELEMENTS_PER_REPEAT; + int32_t pairReduceRepeat32 = (pairReduceRepeat16 + 1) / 2; + int32_t lastComputeOutputElement = outputIdx * numOutputElementsPerInputTile_; + Compute(lastComputeOutputElement, lastRepeatCount, pairReduceRepeat16, pairReduceRepeat32); + } + + ScaleOutput(remainingY); + CopyOut(numStreamOut, remainingY); + } + + __aicore__ inline void CopyInX(const int64_t idx) + { + AscendC::LocalTensor xLocal = inQueueX_.AllocTensor(); + if constexpr (std::is_same_v) { + DataCopy(xLocal, xGm_[maxLoRARank_ * idx], maxLoRARank_); + } else { + uint16_t blockLen = static_cast(maxLoRARank_ * sizeof(X_T)); + DataCopyPad(xLocal, xGm_[maxLoRARank_ * idx], {1, blockLen, 0, 0}, {}); + } + inQueueX_.EnQue(xLocal); + xLocal = inQueueX_.DeQue(); + AscendC::LocalTensor xDup = dupBufferX_.Get(); + + // As we are generating multiple output elements with one API invocation, + // we need to duplicate the X vector multiple times to fill one NUM_BYTES_PER_REPEAT + if constexpr (std::is_same_v) { + for (int32_t i = 0; i < NUM_ELEMENTS_PER_REPEAT; i += maxLoRARank_) { + for (int32_t j = 0; j < maxLoRARank_; j++) { + float entry = xLocal.GetValue(j); + xDup.SetValue(i + j, entry); + } + } + } else { + Cast(xDup, xLocal, AscendC::RoundMode::CAST_NONE, maxLoRARank_); + AscendC::PipeBarrier(); + + for (int32_t i = maxLoRARank_; i < NUM_ELEMENTS_PER_REPEAT; i += maxLoRARank_) { + for (int32_t j = 0; j < maxLoRARank_; j++) { + float entry = xDup.GetValue(j); + xDup.SetValue(i + j, entry); + } + } + } + inQueueX_.FreeTensor(xLocal); + } + + __aicore__ inline void CopyInY(int32_t progress, int32_t numElements = Y_OUT_TILE_NUM_ELEMENTS) + { + AscendC::LocalTensor yInLocal = inQueueY_.AllocTensor(); + DataCopy(yInLocal, yInGm_[yOffset_ + progress * Y_OUT_TILE_NUM_ELEMENTS], numElements); + inQueueY_.EnQue(yInLocal); + } + + __aicore__ inline void CopyInW(int32_t progress, int32_t numElements = W_IN_TILE_NUM_ELEMENTS) + { + AscendC::LocalTensor wLocal = inQueueW_.AllocTensor(); + DataCopy(wLocal, wGm_[reqLoRAWeightOffset_ + progress * W_IN_TILE_NUM_ELEMENTS], numElements); + inQueueW_.EnQue(wLocal); + } + + __aicore__ inline void ScaleOutput(int32_t numElements = Y_OUT_TILE_NUM_ELEMENTS) + { + AscendC::LocalTensor yLocal = tmpBufferY_.Get(); + AscendC::LocalTensor yInLocal = inQueueY_.DeQue(); + AscendC::LocalTensor yInLocalFP32 = inBufferY_.Get(); + Cast(yInLocalFP32, yInLocal, AscendC::RoundMode::CAST_NONE, numElements); + AscendC::PipeBarrier(); + inQueueY_.FreeTensor(yInLocal); + + Add(yLocal, yLocal, yInLocalFP32, numElements); + AscendC::PipeBarrier(); + + AscendC::LocalTensor yOutLocal = outQueueY_.AllocTensor(); + Cast(yOutLocal, yLocal, AscendC::RoundMode::CAST_RINT, numElements); + AscendC::PipeBarrier(); + + outQueueY_.EnQue(yOutLocal); + } + + __aicore__ inline void Compute(int32_t progress, + int32_t blockReduceRepeatCount=BLOCK_REDUCE_NUM_REPEATS, + int32_t pairReduceRepeat16=PAIR_REDUCE_NUM_REPEATS_16, + int32_t pairReduceRepeat32=PAIR_REDUCE_NUM_REPEATS_32) + { + AscendC::LocalTensor yLocal = tmpBufferY_.Get(); + AscendC::LocalTensor xDup = dupBufferX_.Get(); + AscendC::LocalTensor wLocal = inQueueW_.DeQue(); + AscendC::LocalTensor wTmpTensor = tmpBufferW_.Get(); + + Cast(wTmpTensor, wLocal, AscendC::RoundMode::CAST_NONE, MASK_COUNT, blockReduceRepeatCount, castParams_); + AscendC::PipeBarrier(); + inQueueW_.FreeTensor(wLocal); + + Mul(wTmpTensor, xDup, wTmpTensor, MASK_COUNT, blockReduceRepeatCount, dotProductParams_); + AscendC::PipeBarrier(); + + if (maxLoRARank_ == LORA_RANK_8) { + BlockReduceSum(yLocal[progress], wTmpTensor, blockReduceRepeatCount, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + } else if (maxLoRARank_ == LORA_RANK_16) { + BlockReduceSum(wTmpTensor, wTmpTensor, blockReduceRepeatCount, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + PairReduceSum(yLocal[progress], wTmpTensor, pairReduceRepeat16, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + } else if (maxLoRARank_ == LORA_RANK_32) { + BlockReduceSum(wTmpTensor, wTmpTensor, blockReduceRepeatCount, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + PairReduceSum(wTmpTensor, wTmpTensor, pairReduceRepeat16, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + PairReduceSum(yLocal[progress], wTmpTensor, pairReduceRepeat32, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + } else if (maxLoRARank_ == LORA_RANK_64) { + BlockReduceSum(wTmpTensor, wTmpTensor, blockReduceRepeatCount, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + BlockReduceSum(yLocal[progress], wTmpTensor, pairReduceRepeat16, MASK_COUNT, + reduceSumParams_.dstRepStride, reduceSumParams_.srcBlkStride, reduceSumParams_.srcRepStride); + AscendC::PipeBarrier(); + } + } + + __aicore__ inline void CopyOut(int32_t progress, int32_t numElements = Y_OUT_TILE_NUM_ELEMENTS) + { + AscendC::LocalTensor yOutLocal = outQueueY_.DeQue(); + DataCopy(yOutGm_[yOffset_ + progress * Y_OUT_TILE_NUM_ELEMENTS], yOutLocal, numElements); + outQueueY_.FreeTensor(yOutLocal); + } + +private: + AscendC::TPipe* pipe_; + AscendC::TQue inQueueY_, inQueueW_; + AscendC::TQue inQueueX_; + AscendC::TQue outQueueY_; + AscendC::TBuf tmpBufferW_, dupBufferX_, inBufferY_, tmpBufferY_; + AscendC::GlobalTensor xGm_; + AscendC::GlobalTensor wGm_; + AscendC::GlobalTensor yInGm_; + AscendC::GlobalTensor yOutGm_; + AscendC::GlobalTensor loraIndicesGm_; + AscendC::GlobalTensor seqLenGm_; + uint32_t batchSize_; + uint32_t numTokensPerCore_; + uint32_t maxLoRARank_; + uint32_t outputHiddenDim_; + uint32_t sliceOffset_; + uint32_t outputFullDim_; + uint32_t singleLoRAWeightLen_; + int64_t reqLoRAIndex_; + uint64_t reqLoRAWeightOffset_; + uint32_t numOutputElementsPerInputTile_; + uint32_t numStreamInPerOutputTile_; + uint64_t yOffset_; + + // The block stride is set to 1, and 8 blocks in the same repeat are processed continuously. + // The repeat stride is 8, so the vector unit reads 8 consecutive blocks in the first repeat, + // reads next 8 consecutive blocks in the second repeat. + AscendC::UnaryRepeatParams castParams_ = {1, 1, 8, 4}; + + // For each repeat in BlockReduceSum and PairReduceSum we should move forward only one block, + // so we set dstRepStride = 1 + AscendC::UnaryRepeatParams reduceSumParams_ = {1, 1, 1, 8}; + + // When the repeat stride is 0, the vector unit repeatedly reads and computes the first 8 consecutive blocks. + // For xDup we repeatedly use it, so we set src0RepStride = 0 + AscendC::BinaryRepeatParams dotProductParams_ = {1, 1, 1, 8, 0, 8}; + +}; + +#define SGMV_EXPAND_TYPE_DECLARE(TYPE) \ + extern "C" __global__ __aicore__ void sgmv_expand_##TYPE(__gm__ void* x, __gm__ void* weight, \ + __gm__ void* loraIndices, uint32_t loraIndicesSize, \ + __gm__ void* seqLen, uint32_t seqLenSize, \ + __gm__ void* yIn, __gm__ void* yOut, \ + uint32_t batchSize, uint32_t numTokensPerCore, \ + uint32_t maxLoRARank, uint32_t outputHiddenDim, \ + uint32_t sliceOffset, uint32_t outputFullDim) \ + { \ + AscendC::TPipe pipe; \ + SGMVExpand op(&pipe); \ + op.Init(x, weight, loraIndices, loraIndicesSize, seqLen, seqLenSize, \ + yIn, yOut, batchSize, numTokensPerCore, maxLoRARank, \ + outputHiddenDim, sliceOffset, outputFullDim); \ + op.Process(); \ + } + +// declare all dtype kernel +SGMV_EXPAND_TYPE_DECLARE(half) +#if !defined(__CCE_AICORE__) || (__CCE_AICORE__ >= 220) + SGMV_EXPAND_TYPE_DECLARE(bfloat16_t) +#endif + +namespace vllm_fl { +extern void sgmv_expand_impl(AscendType type, void* stream, void* x, void* weight, + void* loraIndices, uint32_t loraIndicesSize, + void* seqLen, uint32_t seqLenSize, + void* yIn, void* yOut, uint32_t batchSize, uint32_t numTokensPerCore, uint32_t maxLoRARank, + uint32_t outputHiddenDim, uint32_t sliceOffset, uint32_t outputFullDim) +{ + uint32_t blockDim = (batchSize + numTokensPerCore - 1) / numTokensPerCore; + if (type == AscendType::FP16) { + sgmv_expand_half<<>>(x, weight, loraIndices, loraIndicesSize, seqLen, seqLenSize, + yIn, yOut, batchSize, + numTokensPerCore, maxLoRARank, outputHiddenDim, sliceOffset, + outputFullDim); + } else if (type == AscendType::BF16) { + #if !defined(__CCE_AICORE__) || (__CCE_AICORE__ >= 220) + sgmv_expand_bfloat16_t<<>>(x, weight, loraIndices, loraIndicesSize, + seqLen, seqLenSize, yIn, yOut, batchSize, + numTokensPerCore, maxLoRARank, outputHiddenDim, + sliceOffset, outputFullDim); + #endif + } else { + return; + } +} + +} // namespace vllm_fl \ No newline at end of file diff --git a/csrc/ascend/kernels/sgmv_shrink.cpp b/csrc/ascend/kernels/sgmv_shrink.cpp new file mode 100644 index 000000000..77d3db78d --- /dev/null +++ b/csrc/ascend/kernels/sgmv_shrink.cpp @@ -0,0 +1,275 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kernel_operator.h" +#include "types.h" + +template +class SGMVShrink { +public: + using X_T = scalar_t; + using W_T = scalar_t; + using Y_T = float; + + static constexpr uint64_t BUFFER_NUM = 1; + static constexpr uint64_t TILE_LENGTH = 11776; // optimal performance tile length + +public: + __aicore__ inline SGMVShrink(AscendC::TPipe *pipe) : pipe_(pipe) {} + __aicore__ inline void Init(__gm__ void *x, __gm__ void *weight, __gm__ void *loraIndices, uint32_t loraIndicesSize, + __gm__ void *seqLen, uint32_t seqLenSize, + __gm__ void *y, uint32_t batchSize, uint32_t numTokensPerCore, uint32_t inputHiddenDim, + uint32_t maxLoRARank, float scale) + { + batchSize_ = batchSize; + numTokensPerCore_ = numTokensPerCore; + inputHiddenDim_ = inputHiddenDim; + maxLoRARank_ = maxLoRARank; + scale_ = scale; + singleLoRAWeightLen_ = inputHiddenDim_ * maxLoRARank_; + incremental_ = inputHiddenDim_ > TILE_LENGTH; + + xGm_.SetGlobalBuffer((__gm__ X_T *)x); + yOutGm_.SetGlobalBuffer((__gm__ Y_T *)y); + wGm_.SetGlobalBuffer((__gm__ W_T *)weight); + loraIndicesGm_.SetGlobalBuffer((__gm__ int64_t *)loraIndices, loraIndicesSize); + seqLenGm_.SetGlobalBuffer((__gm__ int64_t *)seqLen, seqLenSize); + + pipe_->InitBuffer(inQueueX_, BUFFER_NUM, TILE_LENGTH * sizeof(X_T)); + pipe_->InitBuffer(inQueueW_, BUFFER_NUM, TILE_LENGTH * sizeof(W_T)); + pipe_->InitBuffer(tmpBufferX_, TILE_LENGTH * sizeof(float)); + pipe_->InitBuffer(tmpBufferW_, TILE_LENGTH * sizeof(float)); + + pipe_->InitBuffer(outQueueY_, 1, maxLoRARank_ * sizeof(Y_T)); + pipe_->InitBuffer(outBufferY_, maxLoRARank_ * sizeof(float)); + } + + __aicore__ inline void Process() + { + int64_t blockIdx = AscendC::GetBlockIdx(); + int64_t startIdx = blockIdx * numTokensPerCore_; + int64_t endIdx = startIdx + numTokensPerCore_; + if (endIdx > batchSize_) { + endIdx = batchSize_; + } + for (int64_t idx = startIdx; idx < endIdx; idx++) { + // set up LoRA index + CopyInIndex(idx); + if (reqLoRAIndex_ < 0) { + continue; + } + reqLoRAWeightOffset_ = reqLoRAIndex_ * singleLoRAWeightLen_; + + if (incremental_) { + ProcessImpl(idx); + } else { + ProcessImpl(idx); + } + + ScaleOutput(); + CopyOut(idx); + } + } + +private: + template + __aicore__ inline void ProcessImpl(const int64_t idx) + { + AscendC::LocalTensor yOutLocal = outBufferY_.Get(); + if constexpr (!INCREMENTAL_MODE) { + CopyInX(idx, 0, inputHiddenDim_); + AscendC::LocalTensor xTmpTensor = tmpBufferX_.Get(); + AscendC::LocalTensor xLocal = inQueueX_.DeQue(); + Cast(xTmpTensor, xLocal, AscendC::RoundMode::CAST_NONE, inputHiddenDim_); + AscendC::PipeBarrier(); + inQueueX_.FreeTensor(xLocal); + } + + for (int i = 0; i < maxLoRARank_; i++) { + float acc(0); + for (int32_t j = 0; j < inputHiddenDim_ / TILE_LENGTH; j++) { + if constexpr (INCREMENTAL_MODE) { + CopyInX(idx, j); + } + CopyInW(i, j); + Compute(acc); + } + CopyAndComputeLastIteration(idx, i, acc); + yOutLocal.SetValue(i, acc); + } + } + + __aicore__ inline void CopyInIndex(const int64_t idx) + { + // look up the LoRA index + int64_t weightIdx = idx; + uint64_t i = 0; + for (; i < seqLenGm_.GetSize(); i++) { + int64_t repeatValue = seqLenGm_.GetValue(i); + if (weightIdx >= repeatValue) { + weightIdx -= repeatValue; + continue; + } + break; + } + reqLoRAIndex_ = (i < seqLenGm_.GetSize()) ? loraIndicesGm_.GetValue(i) : -1; + } + + __aicore__ inline void CopyInX(const int64_t idx, int32_t colIdx, int32_t numElements = TILE_LENGTH) + { + AscendC::LocalTensor xLocal = inQueueX_.AllocTensor(); + DataCopy(xLocal, xGm_[inputHiddenDim_ * idx + colIdx * TILE_LENGTH], numElements); + inQueueX_.EnQue(xLocal); + } + + __aicore__ inline void CopyInW(int32_t rowIdx, int32_t colIdx, int32_t numElements = TILE_LENGTH) + { + AscendC::LocalTensor wLocal = inQueueW_.AllocTensor(); + DataCopy(wLocal, wGm_[reqLoRAWeightOffset_ + rowIdx * inputHiddenDim_ + colIdx * TILE_LENGTH], numElements); + inQueueW_.EnQue(wLocal); + } + + template + __aicore__ inline void Compute(float &acc, int32_t numElements = TILE_LENGTH) + { + AscendC::LocalTensor wLocal = inQueueW_.DeQue(); + AscendC::LocalTensor xTmpTensor = tmpBufferX_.Get(); + AscendC::LocalTensor wTmpTensor = tmpBufferW_.Get(); + + if constexpr (INCREMENTAL_MODE) { + AscendC::LocalTensor xLocal = inQueueX_.DeQue(); + Cast(xTmpTensor, xLocal, AscendC::RoundMode::CAST_NONE, numElements); + Cast(wTmpTensor, wLocal, AscendC::RoundMode::CAST_NONE, numElements); + AscendC::PipeBarrier(); + inQueueX_.FreeTensor(xLocal); + inQueueW_.FreeTensor(wLocal); + } else { + Cast(wTmpTensor, wLocal, AscendC::RoundMode::CAST_NONE, numElements); + AscendC::PipeBarrier(); + inQueueW_.FreeTensor(wLocal); + } + // dot product of the one tile of X and W + Mul(wTmpTensor, xTmpTensor, wTmpTensor, numElements); + AscendC::PipeBarrier(); + // reduce sum generate one number, which is the summation of all the dot product + ReduceSum(wTmpTensor, wTmpTensor, wTmpTensor, numElements); + AscendC::PipeBarrier(); + + acc += wTmpTensor.GetValue(0); + } + + template + __aicore__ inline void CopyAndComputeLastIteration(const int64_t idx, int32_t rowIdx, float &acc) + { + int32_t colIdx = inputHiddenDim_ / TILE_LENGTH; + int32_t remaining = inputHiddenDim_ % TILE_LENGTH; + if (remaining == 0) { + return; + } + if constexpr (INCREMENTAL_MODE) { + CopyInX(idx, colIdx, remaining); + } + CopyInW(rowIdx, colIdx, remaining); + Compute(acc, remaining); + } + + __aicore__ inline void ScaleOutput() + { + AscendC::LocalTensor yLocal = outBufferY_.Get(); + AscendC::LocalTensor yOutLocal = outQueueY_.AllocTensor(); + + Muls(yOutLocal, yLocal, scale_, maxLoRARank_); + AscendC::PipeBarrier(); + + outQueueY_.EnQue(yOutLocal); + } + + __aicore__ inline void CopyOut(const int64_t idx) + { + AscendC::LocalTensor yOutLocal = outQueueY_.DeQue(); + DataCopy(yOutGm_[maxLoRARank_ * idx], yOutLocal, maxLoRARank_); + outQueueY_.FreeTensor(yOutLocal); + } + +private: + AscendC::TPipe *pipe_; + AscendC::TQue inQueueX_, inQueueW_; + AscendC::TQue outQueueY_; + AscendC::TBuf tmpBufferX_, tmpBufferW_, outBufferY_; + AscendC::GlobalTensor xGm_; + AscendC::GlobalTensor wGm_; + AscendC::GlobalTensor loraIndicesGm_; + AscendC::GlobalTensor seqLenGm_; + AscendC::GlobalTensor yOutGm_; + uint32_t batchSize_; + uint32_t numTokensPerCore_; + uint32_t inputHiddenDim_; + uint32_t maxLoRARank_; + float scale_; + uint32_t singleLoRAWeightLen_; + int64_t reqLoRAIndex_; + uint64_t reqLoRAWeightOffset_; + bool incremental_; +}; + +#define SGMV_SHRINK_TYPE_DECLARE(TYPE) \ + extern "C" __global__ __aicore__ void sgmv_shrink_##TYPE(__gm__ void* x, __gm__ void* weight, \ + __gm__ void* loraIndices, uint32_t loraIndicesSize, \ + __gm__ void* seqLen, uint32_t seqLenSize, \ + __gm__ void* y, uint32_t batchSize, \ + uint32_t numTokensPerCore, uint32_t inputHiddenDim, \ + uint32_t maxLoRARank, float scale) \ + { \ + AscendC::TPipe pipe; \ + SGMVShrink op(&pipe); \ + op.Init(x, weight, loraIndices, loraIndicesSize, seqLen, seqLenSize, \ + y, batchSize, numTokensPerCore, inputHiddenDim, maxLoRARank, scale); \ + op.Process(); \ + } + +// declare all dtype kernel +SGMV_SHRINK_TYPE_DECLARE(half) +#if !defined(__CCE_AICORE__) || (__CCE_AICORE__ >= 220) + SGMV_SHRINK_TYPE_DECLARE(bfloat16_t) +#endif + +namespace vllm_fl { +extern void sgmv_shrink_impl(AscendType type, void* stream, void* x, void* weight, + void* loraIndices, uint32_t loraIndicesSize, + void* seqLen, uint32_t seqLenSize, + void* y, uint32_t batchSize, uint32_t numTokensPerCore, uint32_t inputHiddenDim, + uint32_t maxLoRARank, float scale) +{ + uint32_t blockDim = (batchSize + numTokensPerCore - 1) / numTokensPerCore; + if (type == AscendType::FP16) { + sgmv_shrink_half<<>>(x, weight, loraIndices, loraIndicesSize, seqLen, seqLenSize, + y, batchSize, + numTokensPerCore, inputHiddenDim, maxLoRARank, + scale); + } else if (type == AscendType::BF16) { + #if !defined(__CCE_AICORE__) || (__CCE_AICORE__ >= 220) + sgmv_shrink_bfloat16_t<<>>(x, weight, loraIndices, loraIndicesSize, + seqLen, seqLenSize, + y, batchSize, + numTokensPerCore, inputHiddenDim, maxLoRARank, + scale); + #endif + } else { + return; + } +} + +} // namespace vllm_fl \ No newline at end of file diff --git a/csrc/ascend/kernels/types.h b/csrc/ascend/kernels/types.h new file mode 100644 index 000000000..a5b05e6ee --- /dev/null +++ b/csrc/ascend/kernels/types.h @@ -0,0 +1,25 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +namespace vllm_fl { +enum struct AscendType { + FP16 = 0, + BF16 = 1, + FP32 = 2, +}; +} \ No newline at end of file diff --git a/csrc/ascend/kernels/utils.h b/csrc/ascend/kernels/utils.h new file mode 100644 index 000000000..b64795520 --- /dev/null +++ b/csrc/ascend/kernels/utils.h @@ -0,0 +1,51 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once +#include "kernel_type.h" +namespace vllm_fl { + +template struct AccType; + +#if (__CCE_AICORE__ >= 220) +template <> struct AccType { + using type = float; +}; +#endif + +template <> struct AccType { + using type = half; +}; + +template <> struct AccType { + using type = float; +}; + +template <> struct AccType { + using type = int; +}; + +template +__aicore__ inline void local_mem_copy(AscendC::LocalTensor dst, AscendC::LocalTensor src, int size) +{ + constexpr int loadSize = 256 / sizeof(scalar_t); + int loopCnt = size / loadSize; + int tailSize = size % loadSize; + if (loopCnt) + AscendC::Copy(dst, src, loadSize, loopCnt, {1, 1, 8, 8}); + AscendC::Copy(dst[loopCnt * loadSize], src[loopCnt * loadSize], tailSize, 1, {1, 1, 8, 8}); +} +} // namespace vllm_fl \ No newline at end of file diff --git a/csrc/ascend/mc2/CMakeLists.txt b/csrc/ascend/mc2/CMakeLists.txt new file mode 100644 index 000000000..a07a71c6e --- /dev/null +++ b/csrc/ascend/mc2/CMakeLists.txt @@ -0,0 +1,17 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB SUBDIRECTORIES LIST_DIRECTORIES true RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +foreach(SUBDIR ${SUBDIRECTORIES}) + # 检查子目录中是否存在 CMakeLists.txt + if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${SUBDIR}/CMakeLists.txt) + add_subdirectory(${SUBDIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/CMakeLists.txt b/csrc/ascend/mc2/dispatch_ffn_combine/CMakeLists.txt new file mode 100644 index 000000000..298ed5677 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/dispatch_ffn_combine_torch_adpt.h b/csrc/ascend/mc2/dispatch_ffn_combine/dispatch_ffn_combine_torch_adpt.h new file mode 100644 index 000000000..4e22d8010 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/dispatch_ffn_combine_torch_adpt.h @@ -0,0 +1,89 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OF CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef DISPATCH_FFN_COMBINE_TORCH_ADPT_H +#define DISPATCH_FFN_COMBINE_TORCH_ADPT_H + +namespace vllm_fl { +std::tuple dispatch_ffn_combine( + const at::Tensor& x, + const at::TensorList& weight1, + const at::TensorList& weight2, + const at::Tensor& expert_idx, + const at::TensorList& scale1, + const at::TensorList& scale2, + const at::TensorList& bias1, + const at::TensorList& bias2, + const at::Tensor& probs, + c10::string_view group, + int64_t max_output_size, + at::Tensor& out, + at::Tensor& expert_token_nums, + const c10::optional& x_active_mask, + double swiglu_limit +) { + char *group_ep_ptr = const_cast(group.data()); + bool is_int8 = weight1[0].dtype() == at::kChar; + bool is_int4 = weight1[0].dtype() == at::kInt; + if (is_int8) { + EXEC_NPU_CMD(aclnnDispatchFFNCombine, + x, + weight1, + weight2, + expert_idx, + scale1, + scale2, + probs, + x_active_mask.has_value() ? x_active_mask.value() : at::Tensor(), + group_ep_ptr, + max_output_size, + swiglu_limit, + out, + expert_token_nums); + } else if (is_int4){ + EXEC_NPU_CMD(aclnnDispatchFFNCombineW4A8, + x, + weight1, + weight2, + expert_idx, + scale1, + scale2, + bias1, + bias2, + probs, + x_active_mask.has_value() ? x_active_mask.value() : at::Tensor(), + group_ep_ptr, + max_output_size, + swiglu_limit, + out, + expert_token_nums); + } else { + EXEC_NPU_CMD(aclnnDispatchFFNCombineBF16, + x, + weight1, + weight2, + expert_idx, + scale1, + scale2, + probs, + group_ep_ptr, + max_output_size, + out, + expert_token_nums); + } + return {out, expert_token_nums}; +} +} +#endif \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_host/CMakeLists.txt b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/CMakeLists.txt new file mode 100644 index 000000000..863122e91 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/CMakeLists.txt @@ -0,0 +1,38 @@ +set(_DISPATCH_FFN_INC_OPTS) +if (EXISTS ${ASCEND_CANN_PACKAGE_PATH}/aarch64-linux/ascendc/include) + list(APPEND _DISPATCH_FFN_INC_OPTS -I${ASCEND_CANN_PACKAGE_PATH}/aarch64-linux/ascendc/include) +elseif (EXISTS ${ASCEND_CANN_PACKAGE_PATH}/arm64-linux/ascendc/include) + list(APPEND _DISPATCH_FFN_INC_OPTS -I${ASCEND_CANN_PACKAGE_PATH}/arm64-linux/ascendc/include) +elseif (EXISTS ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/ascendc/include) + list(APPEND _DISPATCH_FFN_INC_OPTS -I${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/ascendc/include) +endif() +if (EXISTS ${CMAKE_SOURCE_DIR}/third_party/catlass/include) + list(APPEND _DISPATCH_FFN_INC_OPTS -I${CMAKE_SOURCE_DIR}/third_party/catlass/include) +endif() + +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnInner PRIVATE + dispatch_ffn_combine_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME DispatchFFNCombine + OPTIONS + --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror + -DHCCL_COMM + -DCATLASS_ARCH=2201 + ${_DISPATCH_FFN_INC_OPTS} +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE dispatch_ffn_combine ACLNNTYPE aclnn_inner) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ${CMAKE_CURRENT_SOURCE_DIR}/../op_kernel + ) +endif() diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_host/dispatch_ffn_combine_def.cpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/dispatch_ffn_combine_def.cpp new file mode 100644 index 000000000..2c9ca5352 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/dispatch_ffn_combine_def.cpp @@ -0,0 +1,99 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dispatch_ffn_combine_def.cpp + * \brief + */ +#include "register/op_def_registry.h" + +namespace ops { +class DispatchFFNCombine : public OpDef { + public: + explicit DispatchFFNCombine(const char *name) : OpDef(name) { + this->Input("a") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("w1") + .ParamType(DYNAMIC) + .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ}) + .IgnoreContiguous(); + this->Input("w2") + .ParamType(DYNAMIC) + .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ}) + .IgnoreContiguous(); + this->Input("expertIdx") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("scale1") + .ParamType(DYNAMIC) + .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("scale2") + .ParamType(DYNAMIC) + .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("probs") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("xActiveMaskOptional") + .ParamType(OPTIONAL) + .DataType({ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + // Output + this->Output("out") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND,ge::FORMAT_ND}); + this->Output("expert_token_nums") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + this->Attr("group").AttrType(REQUIRED).String(); + this->Attr("M").AttrType(OPTIONAL).Int(); + this->Attr("transB").AttrType(OPTIONAL).Bool(false); + this->Attr("weightNz").AttrType(OPTIONAL).Bool(false); + this->Attr("swigluLimit").AttrType(OPTIONAL).Float(0.0f); + + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("aclnnSupport.value", "support_aclnn") + .ExtendCfgInfo("jitCompile.flag", "static_false") + .ExtendCfgInfo("multiKernelSupportDynamicGraph.value", "multi_kernel"); + this->AICore().AddConfig("ascend910_93", aicore_config); + this->AICore().AddConfig("ascend910b", aicore_config); + this->MC2().HcclGroup("group"); + } +}; + +OP_ADD(DispatchFFNCombine); +} // namespace ops diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_host/dispatch_ffn_combine_proto.cpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/dispatch_ffn_combine_proto.cpp new file mode 100644 index 000000000..9401e7f6d --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/dispatch_ffn_combine_proto.cpp @@ -0,0 +1,42 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dispatch_ffn_proto.cpp + * \brief + */ +#include +#include +// #include "../../common/ophost/op_util.h" +// #include "../../common/ophost/hcom_topo_info.h" +// #include "log/ops_log.h" + +using namespace ge; +namespace ops { +const size_t ATTR_GROUP = 0; +const size_t ATTR_RANK_SIZE = 1; +const size_t SUPPORT_DIM_SIZE = 2; + +static ge::graphStatus InferShapeDispatchFFNCombine(gert::InferShapeContext* context) { + (void) context; + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataTypeDispatchFFNCombine(gert::InferDataTypeContext* context) { + (void) context; + // auto d_type = context->GetInputDataType(0); + // context->SetOutputDataType(0, d_type); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(DispatchFFNCombine) + .InferShape(InferShapeDispatchFFNCombine) + .InferDataType(InferDataTypeDispatchFFNCombine); +} // namespace ops diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_host/dispatch_ffn_combine_tiling.cpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/dispatch_ffn_combine_tiling.cpp new file mode 100644 index 000000000..85a9be129 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/dispatch_ffn_combine_tiling.cpp @@ -0,0 +1,339 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +/*! + * \file dispatch_ffn_tiling.cpp + * \brief + */ +#include "vector" +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "tiling_base/error_log.h" +#include "hcom_topo_info.h" +#include "register/op_def_registry.h" +#include "../op_kernel/dispatch_ffn_combine_tiling.h" +#include +#include +#include +#include "../op_kernel/moe_init_routing_quant_v2/moe_init_routing_quant_v2_tiling.h" + +using namespace AscendC; +using namespace ge; + +#define HCCL_BUFFSIZE "HCCL_BUFFSIZE" + +namespace { + // 1. Constant definitions + const char *K_INNER_DEBUG = "DispatchFFNCombine Tiling Debug"; + constexpr uint32_t ATTR_GROUP_INDEX = 0; + constexpr uint32_t ATTR_MAX_OUTPUT_SIZE_INDEX = 1; + constexpr uint32_t ATTR_IS_TRANS_B = 2; + constexpr uint32_t ATTR_WEIGHT_NZ = 3; + constexpr uint32_t ATTR_SWIGLU_LIMIT = 4; + constexpr uint64_t INIT_TILINGKEY = 1000000; + constexpr uint64_t TILINGKEY_TRANS_B = 1U; + constexpr uint64_t TILINGKEY_WEIGHT_NZ = 10; + constexpr uint32_t X_INDEX = 0; + constexpr uint32_t WEIGHT_INDEX = 1; + constexpr uint32_t WEIGHT2_INDEX = 2; + constexpr uint32_t EXPERTID_INDEX = 3; + constexpr uint32_t X_ACTIVE_MASK_INDEX = 7; + constexpr uint32_t BLOCK_NUM = 20; + constexpr uint32_t SYSTEM_NEED_WORKSPACE = 16 * 1024 * 1024; + constexpr uint64_t MB_SIZE = 1024 * 1024UL; +} + +namespace optiling { + +static int32_t CeilDev(int32_t num, int32_t div) +{ + if (div == 0) { + return 0; + } + return (num + div - 1) / div; +} + +static uint64_t GetMaxWindowSize() +{ + uint16_t defaultWindowSize = 200; + const char* hccl_buffsize_env = getenv(HCCL_BUFFSIZE); + if (hccl_buffsize_env != nullptr) { + try { + std::string envStr(hccl_buffsize_env); + unsigned long val = std::stoul(envStr); + if (val <= std::numeric_limits::max()) { + defaultWindowSize = static_cast(val); + } else { + OP_LOGW(K_INNER_DEBUG, "HCCL_BUFFSIZE value %lu is out of range, using default.", val); + } + } catch (const std::exception& e) { + OP_LOGE(K_INNER_DEBUG, "Exception encountered when parsing env HCCL_BUFFSIZE: %s", e.what()); + } + } else { + OP_LOGD(K_INNER_DEBUG, "Env HCCL_BUFFSIZE not set"); + } + const uint64_t maxWindowSize = static_cast(defaultWindowSize) * MB_SIZE; + OP_LOGD(K_INNER_DEBUG, "Get maxWindowSize is %lu", maxWindowSize); + return maxWindowSize; +} + +// Parse and validate rankId, group, worldSize, and isTransB attributes +static ge::graphStatus DispatchFFNCombineCheckAttrAndSetTiling(gert::TilingContext *context, DispatchFFNCombineInfo& info) +{ + auto attrs = context->GetAttrs(); + OP_TILING_CHECK(attrs == nullptr, OP_LOGE(K_INNER_DEBUG, "attrs is null."), return ge::GRAPH_FAILED); + + // TODO: set, validate, and print tiling data related to attributes + auto groupPtr = attrs->GetAttrPointer(static_cast(ATTR_GROUP_INDEX)); + auto maxOutputSizePtr = attrs->GetAttrPointer(ATTR_MAX_OUTPUT_SIZE_INDEX); + auto is_trans_b = attrs->GetAttrPointer(ATTR_IS_TRANS_B); + auto weight_nz = attrs->GetAttrPointer(ATTR_WEIGHT_NZ); + auto swiglu_limit = attrs->GetAttrPointer(ATTR_SWIGLU_LIMIT); + OP_TILING_CHECK(groupPtr == nullptr || strlen(groupPtr) == 0, + OP_LOGE(K_INNER_DEBUG, "group is invalid."), return GRAPH_FAILED); + + OP_TILING_CHECK(is_trans_b == nullptr, + OP_LOGE(K_INNER_DEBUG, "is_trans_b is invalid."), return GRAPH_FAILED); + OP_TILING_CHECK(weight_nz == nullptr, + OP_LOGE(K_INNER_DEBUG, "weight_nz is invalid."), return GRAPH_FAILED); + + info.maxOutputSize = *maxOutputSizePtr; + info.isTransposeB = *is_trans_b; + info.isWeightNz = *weight_nz; + info.swigluLimit = swiglu_limit != nullptr ? *swiglu_limit : 0.0f; + + int64_t rankSize; + (void)ge::HcomTopoInfo::Instance().GetGroupRankSize(groupPtr, rankSize); + info.worldSize = rankSize; + + OP_LOGD(K_INNER_DEBUG, "maxOutputSize=%d ", info.maxOutputSize); + OP_LOGD(K_INNER_DEBUG, "rankSize=%d ", info.worldSize); + + return ge::GRAPH_SUCCESS; +} + +// Extract shapes of input tensors A and B to compute M, K, N +static ge::graphStatus DispatchFFNCombineCheckShapeAndSetTiling(gert::TilingContext *context, DispatchFFNCombineInfo &info) +{ + const char *nodeName = context->GetNodeName(); + + const gert::StorageShape *aStorageShape = context->GetInputShape(X_INDEX); + auto expertIdxTensor = context->GetDynamicInputTensor(EXPERTID_INDEX, 0); + uint32_t M = aStorageShape->GetStorageShape().GetDim(0); + uint32_t K = aStorageShape->GetStorageShape().GetDim(1); + + auto wTensor = context->GetDynamicInputTensor(WEIGHT_INDEX, 0); + uint32_t wTensorDims = wTensor->GetOriginShape().GetDimNum(); + uint32_t N = wTensor->GetStorageShape().GetDim(wTensorDims - 1); + + uint32_t topK = expertIdxTensor->GetStorageShape().GetDim(1); + uint32_t listLen = 0; + while (true) { + auto wTensorT = context->GetDynamicInputTensor(WEIGHT_INDEX, ++listLen); + if (wTensorT == nullptr) {break;} + } + + uint32_t expertPerRank; + if (listLen == 1) { + expertPerRank = wTensor->GetStorageShape().GetDim(0); + } else { + expertPerRank = listLen; + } + + info.M = M; + info.N = N; + info.K = K; + info.expertPerRank = expertPerRank; + info.topK = topK; + info.listLen = listLen; + OP_LOGD(K_INNER_DEBUG, "M=%d ", info.M); + OP_LOGD(K_INNER_DEBUG, "K=%d ", info.K); + OP_LOGD(K_INNER_DEBUG, "N=%d ", info.N); + OP_LOGD(K_INNER_DEBUG, "expertPerRank=%d ", info.expertPerRank); + OP_LOGD(K_INNER_DEBUG, "topK=%d ", info.topK); + OP_LOGD(K_INNER_DEBUG, "listLen=%d ", info.listLen); + + return ge::GRAPH_SUCCESS; +} + +// Get hardware info such as AI Core count and UB capacity for the current chip platform. +static ge::graphStatus DispatchFFNCombineGetPlatformInfoAndSetTiling(gert::TilingContext *context, DispatchFFNCombineInfo& info) +{ + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + uint64_t ubSize = 0U; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); + info.aivNum = aivNum; + info.totalUbSize = ubSize; + + OP_LOGD(K_INNER_DEBUG, "aivNum=%d", info.aivNum); + OP_LOGD(K_INNER_DEBUG, "ubSize=%lu", info.totalUbSize); + + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus CheckXActiveMaskShape(gert::TilingContext *context, const char *nodeName, DispatchFFNCombineInfo &info) +{ + const gert::StorageShape* xActiveMaskStorageShape = context->GetOptionalInputShape(X_ACTIVE_MASK_INDEX); + if (xActiveMaskStorageShape != nullptr) { + OP_TILING_CHECK(xActiveMaskStorageShape->GetStorageShape().GetDimNum() != 1, + OP_LOGE(nodeName, "xActiveMask shape dims must be 1, but current dim num is %lu.", + xActiveMaskStorageShape->GetStorageShape().GetDimNum()), return ge::GRAPH_FAILED); + const int64_t xActiveMaskDim0 = xActiveMaskStorageShape->GetStorageShape().GetDim(0); + OP_TILING_CHECK(xActiveMaskDim0 != static_cast(info.M), + OP_LOGE(nodeName, "xActiveMask Dim0 must be M(%u), but current dim is %lu.", info.M, xActiveMaskDim0), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +void SetTilingData(CoCTiling &cocTilingData, DispatchFFNCombineInfo &info) +{ + cocTilingData.m0 = 128; + cocTilingData.k0 = 256; + cocTilingData.n0 = 256; + cocTilingData.swizzleDirect = 1; + cocTilingData.swizzleOffset = 7; + cocTilingData.ubMoveNum = 16 * 1024; + cocTilingData.pValue = 1; + cocTilingData.commNpuSplit = info.worldSize; + cocTilingData.commDataSplit = 1; + cocTilingData.lenPerLoop = cocTilingData.m0 * cocTilingData.n0 / 2; +} + +// Main scheduling function: +// Get tilingData ➝ check Attr ➝ check Shape ➝ get platform info +// ➝ call SetTilingData (based on rank count) ➝ set blockDim ➝ set tilingKey ➝ set workspace ➝ configure communication parameters + +static ge::graphStatus DispatchFFNCombineTilingFuncImpl(gert::TilingContext *context) +{ + const char *nodeName = context->GetNodeName(); + OP_LOGI(nodeName, "Enter DispatchFFNCombine tiling func."); + + // 1. tilingData + DispatchFFNCombineTilingData *tilingData = context->GetTilingData(); + OP_TILING_CHECK(tilingData == nullptr, OP_LOGE(nodeName, "tilingData is nullptr."), + return ge::GRAPH_FAILED); + OP_LOGI(nodeName, "DispatchFFNCombine get tilingData."); + DispatchFFNCombineInfo& info = tilingData->dispatchFFNCombineInfo; + OP_LOGI(nodeName, "DispatchFFNCombine get tilingData info."); + + OP_TILING_CHECK(DispatchFFNCombineCheckAttrAndSetTiling(context, info) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "DispatchFFNCombine CheckAttrAndSetTiling Failed"), + return ge::GRAPH_FAILED); + OP_TILING_CHECK(DispatchFFNCombineCheckShapeAndSetTiling(context, info) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "DispatchFFNCombine CheckShapeAndSetTiling Failed"), + return ge::GRAPH_FAILED); + OP_TILING_CHECK(DispatchFFNCombineGetPlatformInfoAndSetTiling(context, info) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "DispatchFFNCombine GetPlatformInfoAndSetTiling Failed"), + return ge::GRAPH_FAILED); + OP_TILING_CHECK(CheckXActiveMaskShape(context, nodeName, info) != ge::GRAPH_SUCCESS, + OP_LOGE(context->GetNodeName(), "DispatchFFNCombine CheckXActiveMaskShape Failed"), + return ge::GRAPH_FAILED); + + SetTilingData(tilingData->cocTiling, info); + + // 2. set blockDim + uint32_t blockDim = 1U; + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + auto aicNum = ascendcPlatform.GetCoreNumAic(); + auto aivNum = ascendcPlatform.GetCoreNumAiv(); + blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum); + context->SetBlockDim(blockDim); + + // 3. set tilingKey + uint64_t tilingKey = INIT_TILINGKEY; + tilingKey += info.isTransposeB ? TILINGKEY_TRANS_B : 0; + tilingKey += info.isWeightNz ? TILINGKEY_WEIGHT_NZ : 0; + context->SetTilingKey(tilingKey); + + OP_LOGD(K_INNER_DEBUG, "tilingKey=%d", tilingKey); + + optiling::MoeInitRoutingQuantV2TilingBase moeInitRoutingQuantV2TilingBase; + int64_t inuptXDtypeSize = sizeof(int16_t); + int64_t scaleDim0 = 0; + int64_t ubSize = 196352; + int64_t expertCapacity = 0; + int64_t expertNum = info.expertPerRank * info.worldSize + 1; // enable expertId == expertNum + int64_t activeNum = 0; + int64_t dropPadMode = 0; + int64_t expertTokensCountOrCumsumFlag = 2; + bool expertTokensBeforeCapacityFlag = false; + int64_t quantMode = 1; + uint32_t aivNumInitRouting = 2 * BLOCK_NUM; + moeInitRoutingQuantV2TilingBase.DoTiling(info.M, info.K, info.topK, expertCapacity, expertNum, activeNum, dropPadMode, + expertTokensCountOrCumsumFlag, expertTokensBeforeCapacityFlag, inuptXDtypeSize, quantMode, scaleDim0, aivNumInitRouting, ubSize); + uint64_t initRoutingQuantTilingKey = moeInitRoutingQuantV2TilingBase.tilingKey_; + size_t initRoutingWorkspace = moeInitRoutingQuantV2TilingBase.workspaceSize_; + + tilingData->cocTiling.moeInitRoutingQuantV2TilingData = moeInitRoutingQuantV2TilingBase.quantTilingData; + tilingData->cocTiling.moeInitRoutingQuantV2TilingData.vbsComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.vbsComputeParamsOp; + tilingData->cocTiling.moeInitRoutingQuantV2TilingData.vmsMiddleComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.vmsMiddleComputeParamsOp; + tilingData->cocTiling.moeInitRoutingQuantV2TilingData.sortOutComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.sortOutComputeParamsOp; + tilingData->cocTiling.moeInitRoutingQuantV2TilingData.srcToDstComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.srcToDstComputeParamsOp; + tilingData->cocTiling.moeInitRoutingQuantV2TilingData.srcToDstCapacityComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.srcToDstCapacityComputeParamsOp; + tilingData->cocTiling.moeInitRoutingQuantV2TilingData.gatherOutComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.gatherOutComputeParamsOp; + tilingData->cocTiling.initRoutingQuantTilingKey = initRoutingQuantTilingKey; + + uint64_t maxWindowSize = GetMaxWindowSize(); + uint64_t actualSize = static_cast(info.M) * info.topK * info.K * sizeof(int8_t) * 3 + 10 * MB_SIZE ; + OP_TILING_CHECK((actualSize > maxWindowSize), + OP_LOGE(nodeName, "HCCL_BUFFSIZE is too SMALL, m = %lu, k = %lu, topK = %lu" + " expected HCCL_BUFFSIZE is ((m * k * topK * sizeof(int8_t)) * 3 + 3MB)= %luMB, HCCL_BUFFSIZE=%luMB.", + info.M, info.K, info.topK, (actualSize + MB_SIZE - 1) / MB_SIZE, maxWindowSize / MB_SIZE), + return ge::GRAPH_FAILED); + + // 4. workspace + size_t *workSpaces = context->GetWorkspaceSizes(1); + OP_TILING_CHECK(workSpaces == nullptr, OP_LOGE(nodeName, "workSpaces is nullptr."), + return ge::GRAPH_FAILED); + + uint32_t n2 = info.K; + uint32_t k2 = info.N / 2; + + uint64_t cocWorkspace = (info.M + 256 - 1) / 256 * 256 * info.topK *sizeof(int32_t) + + info.worldSize * info.worldSize * info.expertPerRank * sizeof(int32_t) * 2 + + info.maxOutputSize * sizeof(float) * 2 + + info.maxOutputSize * n2 * sizeof(int16_t) + + info.maxOutputSize * info.K * sizeof(int8_t) + + info.worldSize * sizeof(int32_t) * 16 + + (info.expertPerRank + info.worldSize) * sizeof(int32_t) * 16; + + workSpaces[0] = SYSTEM_NEED_WORKSPACE + std::max(cocWorkspace, initRoutingWorkspace); + + + // 5. communication + auto attrs = context->GetAttrs(); + auto group = attrs->GetAttrPointer(static_cast(ATTR_GROUP_INDEX)); + uint32_t opType = 8U; + std::string algConfig = "AlltoAll=level0:fullmesh;level1:pairwise"; + AscendC::Mc2CcTilingConfig mc2CcTilingConfig(group, opType, algConfig); + mc2CcTilingConfig.GetTiling(tilingData->mc2InitTiling); + mc2CcTilingConfig.GetTiling(tilingData->mc2CcTiling); + + OP_LOGI(nodeName, "Leave DispatchFFNCombine tiling func."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus DispatchFFNCombineTilingFunc(gert::TilingContext* context) +{ + return DispatchFFNCombineTilingFuncImpl(context); +} + +struct DispatchFFNCombineCompileInfo {}; +ge::graphStatus TilingParseForDispatchFFNCombine(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(DispatchFFNCombine) + .Tiling(DispatchFFNCombineTilingFunc) + .TilingParse(TilingParseForDispatchFFNCombine); +} // namespace optiling diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_host/hcom_topo_info.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/hcom_topo_info.h new file mode 100644 index 000000000..827d4c5b6 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/hcom_topo_info.h @@ -0,0 +1,72 @@ +/* Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ===================================================================================================================*/ + +#ifndef METADEF_CXX_INC_EXTERNAL_HCOM_HCOM_TOPO_INFO_H_ +#define METADEF_CXX_INC_EXTERNAL_HCOM_HCOM_TOPO_INFO_H_ + +#include +#include + +using Status = int32_t; + +namespace ge { +static constexpr uint32_t COMM_MESH = 0b1U; +static constexpr uint32_t COMM_SWITCH = (COMM_MESH << 1U); +static constexpr uint32_t COMM_RING = (COMM_MESH << 2U); +static constexpr uint32_t COMM_PAIRWISE = (COMM_MESH << 3U); +class HcomTopoInfo { + public: + enum class TopoLevel { + L0 = 0, + L1, + MAX, + }; + struct TopoLevelDesc { + uint32_t comm_sets; + uint32_t rank_size; + }; + using TopoDescs = TopoLevelDesc[static_cast(TopoLevel::MAX)]; + struct TopoInfo { + int64_t rank_size; + void *notify_handle; + TopoDescs topo_level_descs; + }; + static HcomTopoInfo &Instance(); + bool TopoInfoHasBeenSet(const char_t *group); + bool TryGetGroupTopoInfo(const char_t *group, TopoInfo &info); + Status SetGroupTopoInfo(const char_t *group, const TopoInfo &info); + Status GetGroupRankSize(const char_t *group, int64_t &rank_size); + TopoDescs *GetGroupTopoDesc(const char_t *group); + Status GetGroupNotifyHandle(const char_t *group, void *¬ify_handle); + void UnsetGroupTopoInfo(const char_t *group) { + const std::lock_guard lock(mutex_); + (void) rank_info_.erase(group); + } + + Status SetGroupOrderedStream(const char_t *group, void *stream); + Status GetGroupOrderedStream(const char_t *group, void *&stream); + void UnsetGroupOrderedStream(const char_t *group) { + const std::lock_guard lock(mutex_); + (void) group_to_ordered_stream_.erase(group); + }; + + Status SetGroupOrderedStream(const int32_t device_id, const char_t *group, void *stream); + Status GetGroupOrderedStream(const int32_t device_id, const char_t *group, void *&stream); + void UnsetGroupOrderedStream(const int32_t device_id, const char_t *group); + private: + HcomTopoInfo() = default; + ~HcomTopoInfo() = default; + std::unordered_map rank_info_; + std::mutex mutex_; + std::unordered_map group_to_ordered_stream_; // Ordered stream for the communication domain + std::unordered_map> device_id_to_group_to_ordered_stream_; // Ordered stream for the communication domain +}; +} + +#endif // METADEF_CXX_INC_EXTERNAL_HCOM_HCOM_TOPO_INFO_H_ diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_host/op_api/aclnn_dispatch_ffn_combine.cpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/op_api/aclnn_dispatch_ffn_combine.cpp new file mode 100644 index 000000000..01aff26d1 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/op_api/aclnn_dispatch_ffn_combine.cpp @@ -0,0 +1,83 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#include "aclnn_dispatch_ffn_combine.h" +#include +// #include "aclnn_kernels/common/op_error_check.h" +// #include "opdev/op_log.h" +// #include "opdev/common_types.h" +// #include "opdev/platform.h" +// #include "ophost/matmul_util.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +// using namespace op; + +// using namespace op; + +#ifdef __cplusplus +extern "C" { +#endif + +static constexpr size_t TWO_DIMS = 2; +static constexpr int64_t KVALUE_MIN = 256; +static constexpr int64_t KVALUE_MAX = 65535; +static constexpr size_t HCCL_GROUP_NAME_MAX = 128U; +enum NnopbaseHcclServerType { + NNOPBASE_HCCL_SERVER_TYPE_AICPU = 0, + NNOPBASE_HCCL_SERVER_TYPE_MTE, + NNOPBASE_HCCL_SERVER_TYPE_END +}; + +extern aclnnStatus aclnnInnerDispatchFFNCombineGetWorkspaceSize(const aclTensor* x, const aclTensorList* weight1, const aclTensorList* weight2, + const aclTensor* expertId, const aclTensorList* scale1, const aclTensorList* scale2, + const aclTensor* probs, const aclTensor* xActiveMask, + const char* group, int64_t maxOutputSize, + bool transB, bool weightNz, double swigluLimit, + const aclTensor* out, const aclTensor* expertTokenNums, + uint64_t* workspaceSize, aclOpExecutor** executor); +extern aclnnStatus aclnnInnerDispatchFFNCombine(void *workspace, uint64_t workspaceSize, + aclOpExecutor *executor, aclrtStream stream); +extern "C" void __attribute__((weak)) NnopbaseSetHcclServerType(void *executor, NnopbaseHcclServerType sType); + + + +aclnnStatus aclnnDispatchFFNCombineGetWorkspaceSize(const aclTensor* x, const aclTensorList* weight1, const aclTensorList* weight2, + const aclTensor* expertId, const aclTensorList* scale1, const aclTensorList* scale2, + const aclTensor* probs, const aclTensor* xActiveMask, + const char* group, int64_t maxOutputSize, double swigluLimit, + const aclTensor* out, const aclTensor* expertTokenNums, + uint64_t* workspaceSize, aclOpExecutor** executor) +{ + bool transB = false; + bool weightNz = true; + + aclnnStatus ret = aclnnInnerDispatchFFNCombineGetWorkspaceSize(x, weight1, weight2, expertId, scale1, scale2, probs, xActiveMask, group, + maxOutputSize, transB, weightNz, swigluLimit, + out, expertTokenNums, workspaceSize, executor); + return ret; +} + +aclnnStatus aclnnDispatchFFNCombine(void* workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream) +{ + if (NnopbaseSetHcclServerType) { + NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE); + } + aclnnStatus ret = aclnnInnerDispatchFFNCombine(workspace, workspaceSize, executor, stream); + return ret; +} +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_host/op_api/aclnn_dispatch_ffn_combine.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/op_api/aclnn_dispatch_ffn_combine.h new file mode 100644 index 000000000..6a2f4a177 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/op_api/aclnn_dispatch_ffn_combine.h @@ -0,0 +1,64 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_API_INC_DISPATCH_FFN_COMBINE_ +#define OP_API_INC_DISPATCH_FFN_COMBINE_ + +#include + +#include "aclnn/aclnn_base.h" +#include "hccl/hccl.h" +#include "hccl/hccl_types.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * Operator function: fuse all distributed MoE ops from InitRouting through Unpermute. + * @brief First-stage interface of aclnnDispatchFFNCombine that calculates workspace size based on the specific compute flow. + * @domain aclnn_ops_infer + * @param [in] x: The input tensor. + * @param [in] weight1: The first weight tensor. + * @param [in] weight2: The second weight tensor. + * @param [in] expertId: The expert ID tensor. + * @param [in] scale1: The first scale tensor. + * @param [in] scale2: The second scale tensor. + * @param [in] probs: The probabilities tensor. + * @param [in] group: string identifying the communication domain name. + * @param [in] maxOutputSize: The maximum output size. + * @param [out] out: result of computation + communication; same dtype as input. + * @param [out] workspaceSize: workspace size to allocate on the NPU device side. + * @param [out] executor: op executor containing the operator compute flow. + * @return aclnnStatus: status code. + */ +__attribute__((visibility("default"))) aclnnStatus aclnnDispatchFFNCombineGetWorkspaceSize(const aclTensor* x, const aclTensorList* weight1, const aclTensorList* weight2, + const aclTensor* expertId, const aclTensorList* scale1, const aclTensorList* scale2, + const aclTensor* probs, const aclTensor* xActiveMask, + const char* group, int64_t maxOutputSize, double swigluLimit, + const aclTensor* out, const aclTensor* expertTokenNums, + uint64_t* workspaceSize, aclOpExecutor** executor); + +/** + * @brief Second-stage interface of aclnnDispatchFFNCombine to execute computation. + * @param [in] workspace: workspace memory address allocated on the NPU device side. + * @param [in] workspace_size: workspace size allocated on the NPU device side, obtained from aclnnDispatchFFNCombineGetWorkspaceSize. + * @param [in] executor: op executor containing the operator compute flow. + * @param [in] stream: acl stream. + * @return aclnnStatus: status code. + */ +__attribute__((visibility("default"))) aclnnStatus aclnnDispatchFFNCombine(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, + aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif // OP_API_INC_DISPATCH_FFN_COMBINE_ diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_host/tiling_args.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/tiling_args.h new file mode 100644 index 000000000..950cbe904 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_host/tiling_args.h @@ -0,0 +1,9 @@ +#ifndef TILING_ARGS_H +#define TILING_ARGS_H +#include + +namespace Moe { +constexpr uint64_t COMBINE_STATE_WIN_OFFSET = 3U * 1024UL * 1024UL; +constexpr uint64_t NOTIFY_DISPATCH_WIN_OFFSET = 204U * 1024UL * 1024UL; +} // namespace Moe +#endif // TILING_ARGS_H diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine.cpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine.cpp new file mode 100644 index 000000000..815c13669 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine.cpp @@ -0,0 +1,33 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file dispatch_ffn_combine.cpp + * \brief + */ +#include "kernel_operator.h" +#include "lib/matmul_intf.h" +#include "dispatch_ffn_combine_tiling.h" +#include "dispatch_ffn_combine.h" + +using namespace AscendC; +using namespace DispatchFFNCombineImpl; +extern "C" __global__ __aicore__ void dispatch_ffn_combine(GM_ADDR x, GM_ADDR w1, GM_ADDR w2, GM_ADDR expertId, GM_ADDR scale1, GM_ADDR scale2, GM_ADDR probs, + GM_ADDR xActiveMask, GM_ADDR c, GM_ADDR expertTokenNums, GM_ADDR workspaceGM, GM_ADDR tilingGM) +{ + REGISTER_TILING_DEFAULT(DispatchFFNCombineTilingData); + if (TILING_KEY_IS(1000010)) { + KERNEL_TASK_TYPE(1000010, KERNEL_TYPE_MIX_AIC_1_2); + GET_TILING_DATA_WITH_STRUCT(DispatchFFNCombineTilingData, tilingData, tilingGM); + DispatchFFNCombine op; + op.Init(x, w1, w2, expertId, scale1, scale2, probs, xActiveMask, c, expertTokenNums, workspaceGM, tilingGM); + op.Process(); + } +} \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine.h new file mode 100644 index 000000000..9b759fd3e --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine.h @@ -0,0 +1,291 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dispatch_ffn_combine.h + * \brief + */ + +#ifndef DISPATCH_FFN_COMBINE_H +#define DISPATCH_FFN_COMBINE_H + +using namespace AscendC; + +#include "kernel_operator.h" + +#include "utils/moe_distribute_base.h" + +#include "dispatch_ffn_combine_tiling.h" + +#include "catlass/catlass.hpp" +#include "catlass/arch/arch.hpp" +#include "catlass/epilogue/dispatch_policy.hpp" +#include "catlass/epilogue/block/block_epilogue.hpp" +#include "catlass/epilogue/tile/tile_copy.hpp" +#include "catlass/epilogue/tile/tile_elemwise_add.hpp" +#include "catlass/epilogue/tile/tile_elemwise_muls.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "catlass/gemm/block/block_swizzle.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/kernel/matmul_epilogue.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" + +#include "utils/select_helper.hpp" +#include "utils/const_args.hpp" +#include "dispatch_ffn_combine_kernel.hpp" +#include "moe_init_routing_quant_v2/moe_init_routing_quant_v2_tiling.h" + +using namespace Catlass; + +namespace DispatchFFNCombineImpl { +#define TemplateMMA2AClass typename AType_, typename BType_, typename CType_, bool TB_, bool Nz_ +#define TemplateMMA2ACFunc AType_, BType_, CType_, TB_, Nz_ + +using namespace AscendC; +template +class DispatchFFNCombine { +public: + __aicore__ inline DispatchFFNCombine() {}; + __aicore__ inline void Init(GM_ADDR xGM, GM_ADDR weight1GM, GM_ADDR weight2GM, GM_ADDR expertIdGM, GM_ADDR scale1GM, GM_ADDR scale2GM, + GM_ADDR probs, GM_ADDR xActiveMaskGM, GM_ADDR outGM, GM_ADDR expertTokenNums, GM_ADDR workspaceGM, GM_ADDR tilingGM); + __aicore__ inline void Process(); + + +private: + GM_ADDR xGM_; + GM_ADDR weight1GM_; + GM_ADDR weight2GM_; + GM_ADDR expertIdGM_; + GM_ADDR scale1GM_; + GM_ADDR scale2GM_; + GM_ADDR probs_; + GM_ADDR xActiveMaskGM_; + GM_ADDR outGM_; + GM_ADDR gmExpertTokenNums_; + GM_ADDR workspaceGM_; + + GM_ADDR moeInitRoutingQuantV2Scale = nullptr; + GM_ADDR moeInitRoutingQuantV2Offset = nullptr; + GM_ADDR expertTokensBeforeCapacity = nullptr; + + + TBuf uBuf_; + + int32_t rank; + int32_t rankSize; + int32_t aivNum; + + int32_t m0; + int32_t k0; + int32_t n0; + int32_t swizzlOffset; + int32_t swizzlDirect; + int32_t ubMoveNum; + int32_t pValue; + + int32_t commNpuSplit; + int32_t commDataSplit; + int32_t lenPerLoop; + + int32_t m; + int32_t k; + int32_t n; + int32_t topK; + int32_t expertPerRank; + int32_t maxOutputSize; + int32_t EP; + int32_t listLen; + float swigluLimit; + + optiling::MoeInitRoutingQuantV2TilingData moeInitRoutingQuantV2TilingData; + uint64_t initRoutingQuantTilingKey; + + // Hccl hccl_; + +}; + + +template +__aicore__ inline void DispatchFFNCombine::Init(GM_ADDR xGM, GM_ADDR weight1GM, GM_ADDR weight2GM, GM_ADDR expertIdGM, GM_ADDR scale1GM, GM_ADDR scale2GM, + GM_ADDR probs, GM_ADDR xActiveMaskGM, GM_ADDR outGM, GM_ADDR expertTokenNums, GM_ADDR workspaceGM, GM_ADDR tilingGM) +{ + REGISTER_TILING_DEFAULT(DispatchFFNCombineTilingData); + auto tiling = (__gm__ DispatchFFNCombineTilingData*)tilingGM; + GET_TILING_DATA(tilingData, tilingGM); + + xGM_ = xGM; + weight1GM_ = weight1GM; + weight2GM_ = weight2GM; + expertIdGM_ = expertIdGM; + scale1GM_ = scale1GM; + scale2GM_ = scale2GM; + probs_ = probs; + xActiveMaskGM_ = xActiveMaskGM; + + outGM_ = outGM; + gmExpertTokenNums_ = expertTokenNums; + + workspaceGM_ = workspaceGM; + + aivNum = tilingData.dispatchFFNCombineInfo.aivNum; + + m = tilingData.dispatchFFNCombineInfo.M; + k = tilingData.dispatchFFNCombineInfo.K; + n = tilingData.dispatchFFNCombineInfo.N; + EP = tilingData.dispatchFFNCombineInfo.worldSize; + topK = tilingData.dispatchFFNCombineInfo.topK; + expertPerRank = tilingData.dispatchFFNCombineInfo.expertPerRank; + maxOutputSize = tilingData.dispatchFFNCombineInfo.maxOutputSize; + listLen = tilingData.dispatchFFNCombineInfo.listLen; + swigluLimit = tilingData.dispatchFFNCombineInfo.swigluLimit; + + m0 = tilingData.cocTiling.m0; + k0 = tilingData.cocTiling.k0; + n0 = tilingData.cocTiling.n0; + swizzlDirect = tilingData.cocTiling.swizzleDirect; + swizzlOffset = tilingData.cocTiling.swizzleOffset; + ubMoveNum = tilingData.cocTiling.ubMoveNum; + pValue = tilingData.cocTiling.pValue; + commNpuSplit = tilingData.cocTiling.commNpuSplit; + commDataSplit = tilingData.cocTiling.commDataSplit; + lenPerLoop = tilingData.cocTiling.lenPerLoop; + moeInitRoutingQuantV2TilingData = tilingData.cocTiling.moeInitRoutingQuantV2TilingData; + initRoutingQuantTilingKey = tilingData.cocTiling.initRoutingQuantTilingKey; + + auto contextGM0 = AscendC::GetHcclContext(); + __gm__ HcclOpResParamCustom *WinContext_{nullptr}; + WinContext_ = (__gm__ HcclOpResParamCustom *)contextGM0; + + rank = WinContext_->localUsrRankId; + rankSize = WinContext_->rankSize; +} + +template +__aicore__ inline void DispatchFFNCombine::Process() +{ + // Define ArchTag + using ArchTag = Arch::AtlasA2; + constexpr bool enableUnitFlag = false; + constexpr bool enableShuffleK = true; + + uint32_t k2 = n/2; + uint32_t n2 = k; + + int64_t activeNum = 0; + int64_t expertCapacity = 0; + int64_t expertNum = expertPerRank * EP; + int64_t dropPadMode = 0; + int64_t expertTokensCountOrCumsumFlag = 2; + bool expertTokensBeforeCapacityFlag = false; + int64_t quantMode = 1; + + using LayoutA = layout::RowMajor; + using LayoutB = typename std::conditional< + Nz_, + layout::zN, + typename std::conditional::type + >::type; + + LayoutB layoutB1 = LayoutBInitializer::create(k, n); + LayoutB layoutB2 = LayoutBInitializer::create(k2, n2); + using LayoutC = layout::RowMajor; + using L1TileShape = GemmShape<128, 256, 512>; // M, N, K + + constexpr uint32_t workspaceStages = 2; + constexpr uint32_t preloadStages = 1; + constexpr uint32_t l1Stages = 2; + constexpr uint32_t l0AStages = 2; + constexpr uint32_t l0BStages = 2; + constexpr uint32_t l0CStages = 1; + + using DispatchPolicy = Gemm::MmadAtlasA2PreloadAsyncFixpipe< + preloadStages, + l1Stages, l0AStages, l0BStages, l0CStages, + enableUnitFlag, enableShuffleK + >; + + using L0TileShape = GemmShape<128, 256, 128>; + using AType = Gemm::GemmType; + using BType = Gemm::GemmType; + using CType = Gemm::GemmType; + using D1Type = Gemm::GemmType; + + using D2Type = typename std::conditional< + std::is_same_v, + Gemm::GemmType, + Gemm::GemmType + >::type; + + using BlockMmad = Gemm::Block::BlockMmad; + constexpr uint32_t ubStages = 2; + + using EpilogueDispatchPolicy1 = Epilogue::EpilogueAtlasA2PerTokenDequantSwigluQuant; + + using ScaleType = Gemm::GemmType; + using PerTokenScaleType = Gemm::GemmType; + using ElementMulType = Gemm::GemmType; + using TileElemWiseMuls = Epilogue::Tile::TileElemWiseMuls; + + using TileCopy1 = Epilogue::Tile::TileCopy; + using BlockEpilogue1 = Epilogue::Block::BlockEpilogue; + + using EpilogueDispatchPolicy2 = Epilogue::EpilogueAtlasA2PerTokenDequant; + using EpilogueDispatchPolicy3 = Epilogue::EpilogueAtlasA2PerTokenDequantV2; + + using TileCopy2 = Epilogue::Tile::TileCopy; + using BlockEpilogue2 = Epilogue::Block::BlockEpilogue; + using BlockEpilogue3 = Epilogue::Block::BlockEpilogue; + + + using BlockScheduler = typename Gemm::Block::GemmIdentityBlockSwizzle<9, 1>; + using ElementGroupList = int64_t; + using MatmulKernel = Gemm::Kernel::DispatchFFNCombineKernel; + + LayoutA layoutA1{static_cast(m), static_cast(k)}; + LayoutA layoutA2{static_cast(m), static_cast(k2)}; + layout::VectorLayout layoutScale1{static_cast(n)}; + layout::VectorLayout layoutScale2{static_cast(n2)}; + layout::RowMajor layoutD1{static_cast(maxOutputSize), static_cast(k2)}; + layout::RowMajor layoutD2{static_cast(m*topK), static_cast(n2)}; + // Prepare params + + GemmCoord problemShape{static_cast(m), static_cast(n), static_cast(k)}; + + uint32_t epilogueCoreNum = aivNum; + uint32_t epilogueGranularity = expertPerRank - 3; + if (expertPerRank <= 4) { + epilogueGranularity = expertPerRank - 1; + } + typename MatmulKernel::Params params{ + problemShape, static_cast(EP), static_cast(listLen), static_cast(expertPerRank), static_cast(maxOutputSize), + static_cast(rank), static_cast(rankSize), + static_cast(topK), initRoutingQuantTilingKey, + epilogueCoreNum, epilogueGranularity, + xGM_, layoutA1, layoutA2, + weight1GM_, layoutB1, + weight2GM_, layoutB2, + scale1GM_, layoutScale1, + scale2GM_, layoutScale2, + outGM_, layoutD1, layoutD2, + expertIdGM_, moeInitRoutingQuantV2Scale, moeInitRoutingQuantV2Offset, + expertTokensBeforeCapacity, probs_, + workspaceGM_, gmExpertTokenNums_, ubMoveNum, xActiveMaskGM_, moeInitRoutingQuantV2TilingData, swigluLimit}; + //Call kernel + MatmulKernel kernel(params); + kernel(params); +} + +} // DispatchFFNCombineImpl +#endif // DISPATCH_FFN_COMBINE_H \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine_kernel.hpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine_kernel.hpp new file mode 100644 index 000000000..ba95cc5d7 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine_kernel.hpp @@ -0,0 +1,1190 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef DISPATCH_FFN_COMBINE_KERNEL_HPP +#define DISPATCH_FFN_COMBINE_KERNEL_HPP + +#include "kernel_operator.h" + +#include "catlass/catlass.hpp" +#include "catlass/arch/cross_core_sync.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/coord.hpp" +#include "catlass/detail/callback.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/matrix_coord.hpp" +#include "catlass/epilogue/tile/tile_copy.hpp" + +#ifndef HCCL_COMM + #include "block_mmad_preload_async_fixpipe_quant.hpp" + #include "copy_gm_to_l1_custom.hpp" + #include "block_epilogue_pertoken_row.hpp" + #include "block_epilogue_pertoken_v2.hpp" + #include "block_epilogue_pertoken_swiglu.hpp" + #include "hccl_shmem.hpp" + #include "const_args.hpp" + #include "layout3d.hpp" + #include "tiling/moe_init_routing_quant_v2_tiling.h" + #include "moe_init_routing_quant_v2/moe_init_routing_quant_v2.cpp" + #include "moe_init_routing_quant_v2/moe_v2_fullload_dynamic_quant.h" + #include "moe_token_unpermute.h" + #include "get_tensor_addr.hpp" + inline __gm__ struct OpSystemRunCfg g_opSystemRunCfg{Catlass::L2_OFFSET}; +#else + #include "utils/block_mmad_preload_async_fixpipe_quant.hpp" + #include "utils/copy_gm_to_l1_custom.hpp" + #include "utils/block_epilogue_pertoken_row.hpp" + #include "utils/block_epilogue_pertoken_v2.hpp" + #include "utils/block_epilogue_pertoken_swiglu.hpp" + #include "utils/hccl_shmem.hpp" + #include "utils/const_args.hpp" + #include "utils/layout3d.hpp" + #include "moe_init_routing_quant_v2/moe_init_routing_quant_v2_tiling.h" + #include "moe_init_routing_quant_v2/moe_init_routing_quant_v2.cpp" + #include "moe_init_routing_quant_v2/moe_v2_fullload_dynamic_quant.h" + #include "unpermute/moe_token_unpermute.h" + #include "utils/get_tensor_addr.hpp" +#endif + +using namespace AscendC; + +namespace Catlass::Gemm::Kernel { + +constexpr uint16_t SYNCFLAGC2V = 9; +constexpr uint16_t SYNCFLAGV2C = 10; + +template < + class BlockMmad_, + class BlockScheduler_, + class ElementGroupList_, + class BlockEpilogue1_, + class BlockEpilogue2_, + class BlockEpilogue3_ +> +class DispatchFFNCombineKernel { +public: + using BlockMmad = BlockMmad_; + using ArchTag = typename BlockMmad::ArchTag; + using L1TileShape = typename BlockMmad::L1TileShape; + using ElementA = typename BlockMmad::ElementA; + using LayoutA = typename BlockMmad::LayoutA; + using ElementB = typename BlockMmad::ElementB; + using LayoutB = typename BlockMmad::LayoutB; + using ElementC = typename BlockMmad::ElementC; + using LayoutC = typename BlockMmad::LayoutC; + using ElementAccumulator = typename BlockMmad::ElementAccumulator; + using ElementScale = uint64_t; + using LayoutScale = typename layout::VectorLayout; + using ElementPerTokenScale = float; + using LayoutPerTokenScale = typename layout::VectorLayout; + using BlockScheduler = BlockScheduler_; + + using BlockEpilogue1 = BlockEpilogue1_; + using BlockEpilogue2 = BlockEpilogue2_; + using BlockEpilogue3 = BlockEpilogue3_; + + using ElementD1 = typename BlockEpilogue1::ElementD; + using LayoutD1 = typename BlockEpilogue1::LayoutD; + using ElementD2 = typename BlockEpilogue2::ElementD; + using LayoutD2 = typename BlockEpilogue2::LayoutD; + + /// Parameters structure + struct Params { + // Data members + GemmCoord problemShape; + __gm__ ElementA *ptrA; + LayoutA layoutA; + LayoutA layoutA2; + GM_ADDR ptrB1; + LayoutB layoutB1; + GM_ADDR ptrB2; + LayoutB layoutB2; + GM_ADDR ptrScale1; + LayoutScale layoutScale1; + GM_ADDR ptrScale2; + LayoutScale layoutScale2; + __gm__ ElementD2 *ptrOutput; + LayoutD1 layoutD1; + LayoutD2 layoutD2; + GM_ADDR ptrWorkspace; + GM_ADDR ptrExpertTokenNums; + int32_t EP; + int32_t listLen; + int32_t expertPerRank; + uint32_t maxOutputSize; + uint32_t rank; + uint32_t rankSize; + int32_t ubMoveNum; + GM_ADDR symmetricPtr; + //-------------- + GM_ADDR expertIdx; + GM_ADDR moeInitRoutingQuantV2Scale; + GM_ADDR moeInitRoutingQuantV2Offset; + GM_ADDR expandedX; + GM_ADDR expandedRowIdx; + GM_ADDR expertTokensCountOrCumsum; + GM_ADDR expertTokensBeforeCapacity; + GM_ADDR dynamicQuantScale; + GM_ADDR probs; + GM_ADDR ptrXActiveMask; + int64_t topK; + uint64_t initRoutingQuantTilingKey; + uint32_t epilogueCoreNum; + uint32_t epilogueGranularity; + optiling::MoeInitRoutingQuantV2TilingData moeInitRoutingQuantV2TilingData; + float swigluLimit; + //-------------- + + // Methods + CATLASS_HOST_DEVICE + Params() {} + + CATLASS_HOST_DEVICE + Params( + GemmCoord problemShape_, + uint32_t EP_, uint32_t listLen_, uint32_t expertPerRank_, uint32_t maxOutputSize_, + uint32_t rank_, uint32_t rankSize_, int64_t topK_, + uint64_t initRoutingQuantTilingKey_, uint32_t epilogueCoreNum_, uint32_t epilogueGranularity_, + GM_ADDR ptrA_, LayoutA layoutA_, LayoutA layoutA2_, + GM_ADDR ptrB1_, LayoutB layoutB1_, + GM_ADDR ptrB2_, LayoutB layoutB2_, + GM_ADDR ptrScale1_, LayoutScale layoutScale1_, + GM_ADDR ptrScale2_, LayoutScale layoutScale2_, + GM_ADDR ptrOutput_, LayoutD2 layoutD1_, LayoutD2 layoutD2_, + GM_ADDR expertIdx_, GM_ADDR moeInitRoutingQuantV2Scale_, + GM_ADDR moeInitRoutingQuantV2Offset_, + GM_ADDR expertTokensBeforeCapacity_, GM_ADDR probs_, + GM_ADDR ptrWorkspace_, GM_ADDR gmExpertTokenNums_, int32_t ubMoveNum_, + GM_ADDR ptrXActiveMask_, + optiling::MoeInitRoutingQuantV2TilingData moeInitRoutingQuantV2TilingData_, + float swigluLimit_ + ) : problemShape(problemShape_), + EP(EP_), listLen(listLen_), expertPerRank(expertPerRank_), maxOutputSize(maxOutputSize_), + rank(rank_), rankSize(rankSize_), topK(topK_), + initRoutingQuantTilingKey(initRoutingQuantTilingKey_), + epilogueCoreNum(epilogueCoreNum_), epilogueGranularity(epilogueGranularity_), + ptrA(reinterpret_cast<__gm__ ElementA *>(ptrA_)), layoutA(layoutA_), layoutA2(layoutA2_), + ptrB1(ptrB1_), layoutB1(layoutB1_), + ptrB2(ptrB2_), layoutB2(layoutB2_), + ptrScale1(ptrScale1_), layoutScale1(layoutScale1_), + ptrScale2(ptrScale2_), layoutScale2(layoutScale2_), + ptrOutput(reinterpret_cast<__gm__ ElementD2 *>(ptrOutput_)), layoutD1(layoutD1_), layoutD2(layoutD2_), + expertIdx(expertIdx_), moeInitRoutingQuantV2Scale(moeInitRoutingQuantV2Scale_), + moeInitRoutingQuantV2Offset(moeInitRoutingQuantV2Offset_), + expertTokensBeforeCapacity(expertTokensBeforeCapacity_), probs(probs_), + ptrWorkspace(ptrWorkspace_), ptrExpertTokenNums(gmExpertTokenNums_), ubMoveNum(ubMoveNum_), + ptrXActiveMask(ptrXActiveMask_), + moeInitRoutingQuantV2TilingData(moeInitRoutingQuantV2TilingData_), + swigluLimit(swigluLimit_) + { + } + }; + + // Methods + CATLASS_DEVICE + DispatchFFNCombineKernel(Params const ¶ms) + { + if ASCEND_IS_AIC { + coreIdx = AscendC::GetBlockIdx(); + coreNum = AscendC::GetBlockNum(); + } + + if ASCEND_IS_AIV { + coreIdx = get_block_idx() + get_subblockid() * get_block_num(); + coreNum = get_block_num() * get_subblockdim(); + } + + initBuffer(params); + } + + CATLASS_DEVICE + ~DispatchFFNCombineKernel() + { + } + + template + CATLASS_DEVICE + void operator()(Params const ¶ms); + + template <> + CATLASS_DEVICE + void operator()(Params const ¶ms) + { + GMM1(params); + AscendC::CrossCoreWaitFlag<0x2>(SYNCFLAGV2C); + GMM2(params); + } + + + template <> + CATLASS_DEVICE + void operator()(Params const ¶ms) + { + DispatchAndCombine(params); + } + +private: + CATLASS_DEVICE void initBuffer(Params const ¶ms) { + #ifndef HCCL_COMM + shmem.initShmem(params.symmetricPtr, params.rank, params.rankSize); + #endif + workspaceInfo = WorkspaceInfo(params); + peermemInfo = PeermemInfo(params, shmem); + cumsumMM.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(workspaceInfo.ptrcumsumMM)); + gmA.SetGlobalBuffer(reinterpret_cast<__gm__ ElementA *>(workspaceInfo.ptrA)); + gmC.SetGlobalBuffer(reinterpret_cast<__gm__ ElementC *>(workspaceInfo.ptrC)); + gmPermutedToken.SetGlobalBuffer(reinterpret_cast<__gm__ ElementD1 *>(workspaceInfo.ptrPermutedToken)); + gmC2.SetGlobalBuffer(reinterpret_cast<__gm__ ElementC *>(workspaceInfo.ptrC2)); + gmPerTokenScale1.SetGlobalBuffer(reinterpret_cast<__gm__ ElementPerTokenScale *>(workspaceInfo.ptrPerTokenScale)); + gmPerTokenScale2.SetGlobalBuffer(reinterpret_cast<__gm__ ElementPerTokenScale *>(workspaceInfo.ptrPerTokenScale2)); + tokenPerExpert.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(shmem() + peermemInfo.offsetPeerTokenPerExpert)); + paddedExpertNumAligned = AlignUp(params.EP * params.expertPerRank + 1, ALIGN_128); + tokenPerExpertLayout = Layout3D(paddedExpertNumAligned, params.expertPerRank); + preSumBeforeRank.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(workspaceInfo.ptrSumBeforeRank)); + gmXActiveMask.SetGlobalBuffer(reinterpret_cast<__gm__ bool*>(params.ptrXActiveMask)); + + isCombineV1 = true; + if (params.problemShape.m() * params.topK <= 4096) { + isCombineV1 = false; + } + } + + template + CATLASS_DEVICE void CopyGMToGM( + AscendC::GlobalTensor dst, + AscendC::GlobalTensor src, + int32_t elemNum, + int32_t ubMoveNum + ) + { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + + using TType = Gemm::GemmType; + using CopyGmToUb = Epilogue::Tile::CopyGm2Ub; + using CopyUbToGm = Epilogue::Tile::CopyUb2Gm; + CopyGmToUb copyGmToUb; + CopyUbToGm copyUbToGm; + constexpr int32_t BufferNum = 2; + int tmpBufferSize = 32 * 1024 / sizeof(T); // 32 KB + AscendC::LocalTensor tmpBuffer1 = resource.ubBuf.template GetBufferByByte(0); + tmpBuffer1.SetSize(tmpBufferSize); + int tmpBufferOffset = 96 * 1024; // half of UB + AscendC::LocalTensor tmpBuffer2 = resource.ubBuf.template GetBufferByByte(tmpBufferOffset); + tmpBuffer2.SetSize(tmpBufferSize); + + // [ReduceScatter] 2. Pre Interface Sync + int pingpongId = 0; + auto processCount = CeilDiv(elemNum, ubMoveNum); + for (uint32_t processIndex = 0; processIndex < processCount; ++processIndex) { + uint32_t curProcessNum = (processIndex == processCount - 1) ? elemNum - ubMoveNum * (processCount - 1) : ubMoveNum; + AscendC::TEventID EVENT_ID = pingpongId == 0 ? EVENT_ID0 : EVENT_ID1; + AscendC::LocalTensor buf = pingpongId == 0 ? tmpBuffer1 : tmpBuffer2; + auto processOffset = processIndex * ubMoveNum; + + auto inputOffset = processOffset; + auto outputOffset = processOffset; + // [ReduceScatter] 2. Pre Interface Sync + AscendC::WaitFlag(EVENT_ID); + // [ReduceScatter] 3. Start shmem_mte_get_mem_nbi + copyGmToUb(buf, src[inputOffset], layout::RowMajor{ 1, curProcessNum}, layout::RowMajor{1, curProcessNum}); + AscendC::SetFlag(EVENT_ID); + AscendC::WaitFlag(EVENT_ID); + copyUbToGm(dst[outputOffset], buf, layout::RowMajor{ 1, curProcessNum}, layout::RowMajor{1, curProcessNum}); + + // [ReduceScatter] 4. Post Interface Sync + AscendC::SetFlag(EVENT_ID); + pingpongId = (pingpongId + 1) % BufferNum; + } + // [ReduceScatter] 4. Post Interface Sync + + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + } + + template + CATLASS_DEVICE void CopyGMToGMPerToken( + AscendC::GlobalTensor dst, + AscendC::GlobalTensor dstScale, + AscendC::GlobalTensor src, + int32_t rows, + int32_t hiddenSize, + int32_t ubMoveNum, + int32_t& pingpongId + ) { + + constexpr int32_t BufferNum = 2; + AscendC::LocalTensor tmpBuffer1 = resource.ubBuf.template GetBufferByByte(0); + constexpr int tmpBufferOffset = 96 * 1024; // half of UB + AscendC::LocalTensor tmpBuffer2 = resource.ubBuf.template GetBufferByByte(tmpBufferOffset); + uint32_t copyInNum = hiddenSize + UB_ALIGN; + auto processCount = CeilDiv(rows, ubMoveNum); + for (uint32_t processIndex = 0; processIndex < processCount; ++processIndex) { + pingpongId = (pingpongId + 1) % BufferNum; + AscendC::TEventID EVENT_ID = pingpongId == 0 ? EVENT_ID0 : EVENT_ID1; + AscendC::LocalTensor buf = pingpongId == 0 ? tmpBuffer1 : tmpBuffer2; + AscendC::LocalTensor bufScale = buf[hiddenSize].template ReinterpretCast(); + auto inputOffset = processIndex * ubMoveNum * copyInNum; + + int32_t rowNum = ubMoveNum; + if (processIndex == processCount - 1) { + rowNum = rows - processIndex * ubMoveNum; + } + + AscendC::WaitFlag(EVENT_ID); + int64_t dataLen = rowNum * copyInNum; + AscendC::DataCopy(buf, src[inputOffset], dataLen); + + AscendC::SetFlag(EVENT_ID); + AscendC::WaitFlag(EVENT_ID); + auto outputOffset = processIndex * ubMoveNum * hiddenSize; + #define U16(x) static_cast(x) + AscendC::DataCopyPad(dst[outputOffset], + buf, {U16(rowNum), U16(hiddenSize), 1, 0, 0}); + AscendC::DataCopyPad(dstScale[processIndex * ubMoveNum], + bufScale, {U16(rowNum), U16(sizeof(float)), static_cast(hiddenSize / 32), 0, 0}); + AscendC::SetFlag(EVENT_ID); + + } + } + + + CATLASS_DEVICE + void ApplyXActiveMask(Params const ¶ms) { + if (params.ptrXActiveMask == nullptr) { + return; + } + int32_t m = params.problemShape.m(); + int32_t topK = params.topK; + int32_t expertNum = params.expertPerRank * params.EP; + AscendC::GlobalTensor expertIdxGm; + expertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(params.expertIdx)); + + int32_t totalElements = m * topK; + int32_t base = totalElements / coreNum; + int32_t rem = totalElements % coreNum; + + int32_t startIdx = coreIdx * base + min(coreIdx, rem); + int32_t endIdx = (coreIdx + 1) * base + min(coreIdx + 1, rem); + + AscendC::LocalTensor tmpExpertIdx = resource.ubBuf.template GetBufferByByte(0); + int32_t copySize = endIdx - startIdx; + + AscendC::DataCopyPad(tmpExpertIdx[0], expertIdxGm[startIdx], + {1, static_cast(copySize * sizeof(int32_t)), 0, 0}, {} + ); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + for (int32_t i = 0; i < copySize; ++i) { + int32_t tokenIdx = (startIdx + i) / topK; + bool isActive = gmXActiveMask(tokenIdx); + if (!isActive) { + tmpExpertIdx.SetValue(i, expertNum); + } + } + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopyPad(expertIdxGm[startIdx], tmpExpertIdx[0], {1, static_cast(copySize * sizeof(int32_t)), 0, 0, 0}); + AscendC::SyncAll(); + } + + CATLASS_DEVICE + void GetCumsumForMMAIV(AscendC::GlobalTensor & tokenPerExpert, AscendC::GlobalTensor & result, uint32_t expertPerRank, uint32_t rankId, uint32_t EP) + { + int32_t expertPerRankAligned = (expertPerRank + 8 - 1) / 8 * 8; + AscendC::LocalTensor tmpBuffer1 = resource.ubBuf.template GetBufferByByte(0); + AscendC::LocalTensor tmpResult = resource.ubBuf.template GetBufferByByte(EP * expertPerRank * sizeof(int32_t)); + #define U16(x) static_cast(x) + + AscendC::DataCopyPad( + tmpBuffer1, + tokenPerExpert[rankId * expertPerRank], + {U16(EP), U16(expertPerRank * sizeof(int32_t)), U16((paddedExpertNumAligned - expertPerRank) * sizeof(int32_t)), 0}, + {} + ); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + for (uint32_t i = 1; i < EP; ++i) { + AscendC::Add(tmpBuffer1[i * expertPerRankAligned], tmpBuffer1[i * expertPerRankAligned], tmpBuffer1[(i - 1) * expertPerRankAligned], expertPerRank); + AscendC::PipeBarrier(); + } + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::DataCopyPad( + result, + tmpBuffer1, + {U16(EP), U16((expertPerRank) * sizeof(int32_t)), 0, 0} + ); + } + + CATLASS_DEVICE + void GMM1(Params const ¶ms){ + icache_preload(8); + BlockScheduler blockScheduler; + BlockMmad blockMmad(resource); + float aivFinishGroups = 0.0f; + __gm__ float* aivFinishPtr = workspaceInfo.ptrSoftFlagBase + params.EP * FLAGSTRIDE; + + int64_t gmGroupOffsetA = 0; + int64_t gmGroupOffsetB = 0; + int64_t gmGroupOffsetC = 0; + uint32_t startCoreIdx = 0; + uint32_t syncGroupIdx = 0; + int64_t preCurrentmSum = 0; + int32_t syncLoopIdx = -1; + + uint16_t syncgmmIdx = 0; + AscendC::CrossCoreWaitFlag<0x2>(syncgmmIdx / CROSS_CORE_FLAG_MAX_SET_COUNT); // Wait for AIV to finish cumsum for matmul + syncgmmIdx++; + + for (uint32_t groupIdx = 0; groupIdx < params.expertPerRank; ++groupIdx) { + uint32_t currentM = cumsumMM((params.EP - 1) * params.expertPerRank + groupIdx); + if (preCurrentmSum >= params.maxOutputSize) { + currentM = 0; + } else if (preCurrentmSum + currentM >= params.maxOutputSize) { + currentM = params.maxOutputSize - preCurrentmSum; + } + AscendC::GlobalTensor gmB1; + AscendC::GlobalTensor gmS; + int32_t arrayGroupIdx = params.listLen == 1 ? 0 : groupIdx; + gmB1.SetGlobalBuffer(reinterpret_cast<__gm__ ElementB *>(GetTensorAddr(arrayGroupIdx, params.ptrB1))); + gmS.SetGlobalBuffer(reinterpret_cast<__gm__ ElementScale *>(GetTensorAddr(arrayGroupIdx, params.ptrScale1))); + if (currentM <= L1TileShape::M) { + gmB1.SetL2CacheHint(AscendC::CacheMode::CACHE_MODE_DISABLE); + } + GemmCoord inGroupProblemShape{currentM, params.problemShape.n(), params.problemShape.k()}; + LayoutA layoutA = params.layoutA.GetTileLayout(inGroupProblemShape.GetCoordMK()); + LayoutB layoutB1 = params.layoutB1; + LayoutScale layoutScale = params.layoutScale1; + LayoutC layoutC = LayoutC(inGroupProblemShape.m(), inGroupProblemShape.n(), params.problemShape.k()); + blockScheduler.Update(inGroupProblemShape, MakeCoord(L1TileShape::M, L1TileShape::N)); + uint32_t coreLoops = blockScheduler.GetCoreLoops(); + // Determine the starting loopIdx of the current core under the current groupIdx + uint32_t startLoopIdx = ((coreIdx < startCoreIdx) ? (coreIdx + coreNum) : coreIdx) - startCoreIdx; + // Loop through the matmul of each groupIdx + + for (uint32_t loopIdx = startLoopIdx; loopIdx < coreLoops; loopIdx += coreNum) { + for(;syncGroupIdx <= groupIdx; syncGroupIdx++) { + AscendC::CrossCoreWaitFlag<0x2>(syncgmmIdx / CROSS_CORE_FLAG_MAX_SET_COUNT); + syncgmmIdx ++; + } + + // Compute block location + GemmCoord blockCoord = blockScheduler.GetBlockCoord(loopIdx); + GemmCoord actualBlockShape = blockScheduler.GetActualBlockShape(blockCoord); + // Compute initial location in logical coordinates + MatrixCoord offsetA{blockCoord.m() * L1TileShape::M, blockCoord.k() * L1TileShape::K}; + MatrixCoord offsetB{blockCoord.k() * L1TileShape::K, blockCoord.n() * L1TileShape::N}; + MatrixCoord offsetC{blockCoord.m() * L1TileShape::M, blockCoord.n() * L1TileShape::N}; + int64_t gmOffsetA = layoutA.GetOffset(offsetA); + int64_t gmOffsetB = layoutB1.GetOffset(offsetB); + int64_t gmOffsetC = layoutC.GetOffset(offsetC); + int64_t gmOffsetS = blockCoord.n() * L1TileShape::N + (params.listLen == 1 ? groupIdx * params.problemShape.n() : 0); + if (currentM > 0) { + blockMmad( + gmA[gmGroupOffsetA + gmOffsetA], layoutA, + gmB1[gmGroupOffsetB + gmOffsetB], layoutB1, + gmC[gmGroupOffsetC + gmOffsetC], layoutC, + gmS[gmOffsetS], layoutScale, + actualBlockShape + ); + } + } + + if ((groupIdx + 1) == params.epilogueGranularity && (groupIdx < params.expertPerRank - 1)) { + syncLoopIdx ++; + if constexpr (BlockMmad::DispatchPolicy::ASYNC) { + blockMmad.SynchronizeBlock(); + } + // Synchronization signal: GMM1 notifies SwiGLU [1] + blockMmad.Finalize(syncLoopIdx, SYNCFLAGC2V); + } + + preCurrentmSum += currentM; + gmGroupOffsetA += inGroupProblemShape.m() * inGroupProblemShape.k(); + if (params.listLen == 1) { + gmGroupOffsetB += inGroupProblemShape.k() * inGroupProblemShape.n(); + } + gmGroupOffsetC += inGroupProblemShape.m() * inGroupProblemShape.k(); + startCoreIdx = (startCoreIdx + coreLoops) % coreNum; + } + + for(;syncGroupIdx < params.expertPerRank; syncGroupIdx++) { + AscendC::CrossCoreWaitFlag<0x2>(syncgmmIdx / CROSS_CORE_FLAG_MAX_SET_COUNT); + syncgmmIdx ++; + } + + if constexpr (BlockMmad::DispatchPolicy::ASYNC) { + blockMmad.SynchronizeBlock(); + } + // Synchronization signal: GMM1 notifies SwiGLU [2] + blockMmad.Finalize(syncLoopIdx + 1, SYNCFLAGC2V); + } + + CATLASS_DEVICE + void GMM2(Params const ¶ms) { + icache_preload(8); + BlockScheduler blockScheduler; + BlockMmad blockMmad(resource); + + uint32_t n2 = params.problemShape.k(); + uint32_t k2 = params.problemShape.n() / 2; + + int64_t gmGroupOffsetA = 0; + int64_t gmGroupOffsetB = 0; + int64_t gmGroupOffsetC = 0; + + uint32_t startCoreIdx = 0; + + int64_t preCurrentmSum = 0; + int32_t syncLoopIdx = -1; + uint32_t lastDequantExpertNum = params.expertPerRank; + + if (params.epilogueGranularity < params.expertPerRank) { + lastDequantExpertNum = params.expertPerRank - params.epilogueGranularity; + } + + for (uint32_t groupIdx = 0; groupIdx < params.expertPerRank; ++groupIdx) { + uint32_t currentM = cumsumMM((params.EP - 1) * params.expertPerRank + groupIdx); + if (preCurrentmSum >= params.maxOutputSize) { + currentM = 0; + } else if (preCurrentmSum + currentM > params.maxOutputSize) { + currentM = params.maxOutputSize - preCurrentmSum; + } + AscendC::GlobalTensor gmB2; + AscendC::GlobalTensor gmS2; + int32_t arrayGroupIdx = params.listLen == 1 ? 0 : groupIdx; + gmB2.SetGlobalBuffer(reinterpret_cast<__gm__ ElementB *>(GetTensorAddr(arrayGroupIdx, params.ptrB2))); + gmS2.SetGlobalBuffer(reinterpret_cast<__gm__ ElementScale *>(GetTensorAddr(arrayGroupIdx, params.ptrScale2))); + if (currentM <= L1TileShape::M) { + gmB2.SetL2CacheHint(AscendC::CacheMode::CACHE_MODE_DISABLE); + } + GemmCoord inGroupProblemShape{currentM, n2, k2}; // M N K + + LayoutA layoutA = params.layoutA2.GetTileLayout(inGroupProblemShape.GetCoordMK()); + LayoutB layoutB2 = params.layoutB2; + LayoutScale layoutScale = params.layoutScale2; + LayoutC layoutC = LayoutC(inGroupProblemShape.m(), inGroupProblemShape.n()); + + blockScheduler.Update(inGroupProblemShape, MakeCoord(L1TileShape::M, L1TileShape::N)); + uint32_t coreLoops = blockScheduler.GetCoreLoops(); + + // Determine the starting loopIdx of the current core under the current groupIdx + uint32_t startLoopIdx = ((coreIdx < startCoreIdx) ? (coreIdx + coreNum) : coreIdx) - startCoreIdx; + // Loop through the matmul of each groupIdx + if (params.expertPerRank > lastDequantExpertNum && groupIdx + 1 == params.expertPerRank - lastDequantExpertNum) { + AscendC::CrossCoreWaitFlag<0x2>(SYNCFLAGV2C); + } + + for (uint32_t loopIdx = startLoopIdx; loopIdx < coreLoops; loopIdx += coreNum) { + if (loopIdx + coreNum >= coreLoops) { + syncLoopIdx = groupIdx; + } + + // Compute block location + GemmCoord blockCoord = blockScheduler.GetBlockCoord(loopIdx); + GemmCoord actualBlockShape = blockScheduler.GetActualBlockShape(blockCoord); + + // Compute initial location in logical coordinates + MatrixCoord offsetA{blockCoord.m() * L1TileShape::M, blockCoord.k() * L1TileShape::K}; + MatrixCoord offsetB{blockCoord.k() * L1TileShape::K, blockCoord.n() * L1TileShape::N}; + MatrixCoord offsetC{blockCoord.m() * L1TileShape::M, blockCoord.n() * L1TileShape::N}; + + int64_t gmOffsetA = layoutA.GetOffset(offsetA); + int64_t gmOffsetB = layoutB2.GetOffset(offsetB); + int64_t gmOffsetC = layoutC.GetOffset(offsetC); + int64_t gmOffsetS = blockCoord.n() * L1TileShape::N + (params.listLen == 1 ? groupIdx * n2 : 0); // One scale group per expert + if (currentM > 0) { + blockMmad( + gmPermutedToken[gmGroupOffsetA + gmOffsetA], layoutA, + gmB2[gmGroupOffsetB + gmOffsetB], layoutB2, + gmC2[gmGroupOffsetC + gmOffsetC], layoutC, + gmS2[gmOffsetS], layoutScale, + actualBlockShape, syncLoopIdx, 0 + ); + } + } + preCurrentmSum += currentM; + gmGroupOffsetA += inGroupProblemShape.m() * inGroupProblemShape.k(); + if (params.listLen == 1) { + gmGroupOffsetB += inGroupProblemShape.k() * inGroupProblemShape.n(); + } + gmGroupOffsetC += inGroupProblemShape.m() * inGroupProblemShape.n(); + + startCoreIdx = (startCoreIdx + coreLoops) % coreNum; + } + if constexpr (BlockMmad::DispatchPolicy::ASYNC) { + blockMmad.SynchronizeBlock(); + } + if (isCombineV1) { + blockMmad.Finalize(params.expertPerRank - 1, 0); + } + } + + + CATLASS_DEVICE + void InitArithProgress(Params const ¶ms) { + AscendC::LocalTensor tmpBuffer1 = resource.ubBuf.template GetBufferByByte(0); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::Duplicate(tmpBuffer1, 0.0f, (params.EP + 1) * FLAGSTRIDE); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor flagGlobalBase; + flagGlobalBase.SetGlobalBuffer(workspaceInfo.ptrSoftFlagBase); + AscendC::DataCopy(flagGlobalBase, tmpBuffer1, (params.EP + 1) * FLAGSTRIDE); + } + + + CATLASS_DEVICE + void CrossRankSyncAndlocalTokenPerExpertAllGatherAndGetSumPreRankV2(Params const ¶ms, int64_t localTokenPerExpertOffset){ + uint32_t numPerCore = paddedExpertNumAligned; + AscendC::LocalTensor tmpBuffer = resource.ubBuf.template GetBufferByByte(0); + AscendC::LocalTensor prevSumBuf = tmpBuffer[numPerCore]; + + for(int32_t dstEpIdx = coreIdx; dstEpIdx < params.EP; dstEpIdx += coreNum) { + if (dstEpIdx == params.rank) { + continue; + } + AscendC::GlobalTensor srcAddress; + srcAddress.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(shmem() + localTokenPerExpertOffset)); + AscendC::GlobalTensor dstAddress; + __gm__ void* dstPeermemPtr = shmem(localTokenPerExpertOffset, coreIdx); + dstAddress.SetGlobalBuffer((__gm__ int32_t * )dstPeermemPtr); + + AscendC::SetFlag(EVENT_ID0); + using TType = Gemm::GemmType; + using CopyGmToUb = Epilogue::Tile::CopyGm2Ub; + using CopyUbToGm = Epilogue::Tile::CopyUb2Gm; + CopyGmToUb copyGmToUb; + CopyUbToGm copyUbToGm; + + AscendC::WaitFlag(EVENT_ID0); + + copyGmToUb(tmpBuffer, srcAddress[0], + layout::RowMajor{ 1, numPerCore}, + layout::RowMajor{1, numPerCore}); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::Adds(tmpBuffer, tmpBuffer, 0x800000, numPerCore); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + copyUbToGm(dstAddress[0], tmpBuffer, + layout::RowMajor{ 1, numPerCore}, + layout::RowMajor{1, numPerCore}); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + for(int32_t dstEpIdx = coreIdx; dstEpIdx < params.EP; dstEpIdx += coreNum) { + if (dstEpIdx != params.rank) { + int32_t intPer512 = CACHE_LINE / sizeof(int); + for(int32_t checkIdx = 0; checkIdx < paddedExpertNumAligned; checkIdx += intPer512) { + __gm__ int32_t* sync_check = reinterpret_cast<__gm__ int32_t*>(shmem() + peermemInfo.offsetPeerTokenPerExpert) + tokenPerExpertLayout(dstEpIdx, 0, checkIdx); + gm_signal_wait_until_ne(sync_check, 0); + } + AscendC::DataCopy(tmpBuffer, tokenPerExpert[tokenPerExpertLayout(dstEpIdx, 0, 0)], numPerCore); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::Adds(tmpBuffer, tmpBuffer, -0x800000, numPerCore); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopy(tokenPerExpert[tokenPerExpertLayout(dstEpIdx, 0, 0)], tmpBuffer, numPerCore); + } else { + AscendC::DataCopy(tmpBuffer, tokenPerExpert[tokenPerExpertLayout(dstEpIdx, 0, 0)], numPerCore); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::PipeBarrier(); + int32_t prevSum = 0; + int32_t j = 0; + for (int32_t i = 0; i < (params.rank + 1) * params.expertPerRank; i++) { + if (i >= params.rank * params.expertPerRank) { + prevSumBuf(j) = prevSum; + j++; + } + prevSum += tmpBuffer(i); + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopyPad(preSumBeforeRank[dstEpIdx * params.expertPerRank], prevSumBuf, + AscendC::DataCopyParams{1, static_cast(params.expertPerRank * sizeof(int32_t)), 0, 0}); + } + + AscendC::SyncAll(); + } + + CATLASS_DEVICE + void ResetTokenPerExpert(int32_t num) + { + if (coreIdx != coreNum - 1) { + return; + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::LocalTensor tmp = resource.ubBuf.template GetBufferByByte(0); + AscendC::Duplicate(tmp, 0, num); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopy(tokenPerExpert, tmp, num); + } + + CATLASS_DEVICE + void UpdateAicFlags(const Params ¶ms) + { + float flagBase = 1.0f * params.expertPerRank; + __gm__ float* aicFinishPtr = workspaceInfo.ptrSoftFlagBase + params.EP * FLAGSTRIDE; + float flag = 0.0f; + float lastflag = -1.0f; + AscendC::LocalTensor tmpBuffer1 = resource.ubBuf.template GetBufferByByte(0); + __gm__ float* flagPtr = workspaceInfo.ptrSoftFlagBase; + AscendC::GlobalTensor flagGM; + flagGM.SetGlobalBuffer(flagPtr); + int32_t flagBufferSize = max(4, params.EP) * FLAGSTRIDE; + AscendC::LocalTensor dstValueBuffer = resource.ubBuf.template GetBufferByByte(flagBufferSize); + AscendC::LocalTensor sharedTmpBuffer = resource.ubBuf.template GetBufferByByte((flagBufferSize + 64)); + uint64_t mask[1] = {0}; + uint32_t repeatNum = (flagBufferSize / (4 * FLAGSTRIDE)); + for (int32_t i = 0; i < 4; i ++) { + if (i < params.EP) { + mask[0] |= 1ull * (1ull << (i * 16)); + } + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + while (flag < flagBase) { + flag = flagBase; + AscendC::DataCopy(tmpBuffer1, flagGM, params.EP * FLAGSTRIDE); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::ReduceMin(dstValueBuffer, tmpBuffer1, sharedTmpBuffer, mask, repeatNum, 8, false); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + flag = min(flag, dstValueBuffer.GetValue(0)); + + if (flag > lastflag) { + *aicFinishPtr = flag; + gm_dcci(aicFinishPtr); + lastflag = flag; + } + } + } + + + CATLASS_DEVICE + void DispatchAndCombine(Params const ¶ms) { + icache_preload(8); + int64_t localTokenPerExpertOffset = peermemInfo.offsetPeerTokenPerExpert + tokenPerExpertLayout(params.rank, 0, 0) * sizeof(int32_t); + GM_ADDR localTokenPerExpert = shmem() + localTokenPerExpertOffset; // Place the entire communication matrix in peermem + uint32_t expandedRowIdxOffset = AlignUp(params.problemShape.m(), 256) * params.topK * sizeof(int32_t); + + ApplyXActiveMask(params); + + //---initRouting------ + moe_init_routing_quant_v2(reinterpret_cast (params.ptrA), params.expertIdx, + params.moeInitRoutingQuantV2Scale, params.moeInitRoutingQuantV2Offset, shmem() + peermemInfo.offsetA, + workspaceInfo.expandedRowIdx, localTokenPerExpert, params.expertTokensBeforeCapacity, + shmem() + peermemInfo.offsetPeerPerTokenScale, + params.ptrWorkspace + expandedRowIdxOffset, + ¶ms.moeInitRoutingQuantV2TilingData, params.initRoutingQuantTilingKey); + + AscendC::SyncAll(); + + CrossRankSyncAndlocalTokenPerExpertAllGatherAndGetSumPreRankV2(params, localTokenPerExpertOffset); + + if (coreIdx == 0) { + GetCumsumForMMAIV(tokenPerExpert, cumsumMM, params.expertPerRank, params.rank, params.EP); + } + + uint32_t curGroupOffset = 0; + int32_t prevSumBeforeRank = 0; + int32_t prevSum = 0; + if (coreIdx < params.EP) { + prevSum = preSumBeforeRank(coreIdx * params.expertPerRank); + } + AscendC::SyncAll(); + + AscendC::GlobalTensor ExpertTokenNums; + ExpertTokenNums.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(params.ptrExpertTokenNums)); + if(coreIdx == 0) + { + CopyGMToGM(ExpertTokenNums, cumsumMM[(params.EP - 1) * params.expertPerRank], params.expertPerRank, params.ubMoveNum); + } + uint16_t syncgmm1Idx = 0; + AscendC::CrossCoreSetFlag<0x2, PIPE_MTE3>(syncgmm1Idx / CROSS_CORE_FLAG_MAX_SET_COUNT); + syncgmm1Idx++; + + uint32_t prevGroupSum1 = 0, dequantSum1 = 0, dequantSum2 = 0; + uint32_t dequantSum = 0; + + icache_preload(8); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + int32_t pingpongIdx = 0; + for (int32_t groupIdx = 0; groupIdx < params.expertPerRank; ++groupIdx) { + // The ith core reads data from the ith rank's peermem + uint32_t currentM = cumsumMM((params.EP - 1) * params.expertPerRank + groupIdx); + for(int32_t dstEpIdx = coreIdx; dstEpIdx < params.EP; dstEpIdx += coreNum) { + uint32_t rowStart = (dstEpIdx == 0 ? 0 : cumsumMM((dstEpIdx - 1) * params.expertPerRank + groupIdx)) + prevGroupSum1; + if (rowStart < params.maxOutputSize) { + uint32_t rows = tokenPerExpert(tokenPerExpertLayout(dstEpIdx, params.rank, groupIdx)); + if (rowStart + rows > params.maxOutputSize) { + rows = params.maxOutputSize - rowStart; + } + uint32_t rowSrc = prevSum; + prevSum += rows; + GM_ADDR otherRankPtr = shmem(0, dstEpIdx); + AscendC::GlobalTensor gmRemoteA; + gmRemoteA.SetGlobalBuffer(reinterpret_cast<__gm__ ElementA*>(otherRankPtr + peermemInfo.offsetA)); + MatrixCoord offsetA{rowStart, 0}; + int64_t gmOffsetA = params.layoutA.GetOffset(offsetA); + int64_t gmOffsetPeer = rowSrc * (params.problemShape.k() + UB_ALIGN); + int32_t ubMoveNum = 2; + CopyGMToGMPerToken(gmA[gmOffsetA], gmPerTokenScale1[rowStart], gmRemoteA[gmOffsetPeer], rows, params.problemShape.k(), ubMoveNum, pingpongIdx); + } + + } + AscendC::SyncAll(); + AscendC::CrossCoreSetFlag<0x2, PIPE_MTE3>(syncgmm1Idx / CROSS_CORE_FLAG_MAX_SET_COUNT); + syncgmm1Idx ++; + + prevGroupSum1 += currentM; + + // Token count and truncation logic for the first SwiGLU operation + if (groupIdx + 1 <= params.epilogueGranularity) { + if (dequantSum1 + currentM <= params.maxOutputSize) { + dequantSum1 += currentM; + } else if (dequantSum1 < params.maxOutputSize) { + dequantSum1 = params.maxOutputSize; + } + } + + // Token count and truncation logic for the second SwiGLU operation + if (groupIdx + 1 > params.epilogueGranularity && dequantSum1 < params.maxOutputSize) { + if (dequantSum1 + dequantSum2 + currentM <= params.maxOutputSize) { + dequantSum2 += currentM; + } else if (dequantSum1 + dequantSum2 < params.maxOutputSize) { + dequantSum2 += params.maxOutputSize - dequantSum1 - dequantSum2; + } + } + } + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + + uint32_t n2 = params.problemShape.k(); + + + typename BlockEpilogue2::Params epilogueParams2{ + static_cast(params.EP), + static_cast(params.expertPerRank), + reinterpret_cast<__gm__ int32_t *>(shmem() + peermemInfo.offsetPeerTokenPerExpert), + static_cast(n2) + }; + + typename BlockEpilogue3::Params epilogueParams3{ + static_cast(params.EP), + static_cast(params.expertPerRank), + static_cast(params.rank), + reinterpret_cast<__gm__ int32_t *>(shmem() + peermemInfo.offsetPeerTokenPerExpert), + params.layoutD2, + static_cast(n2), + static_cast(L1TileShape::N), + shmem, + static_cast(peermemInfo.offsetD), + tokenPerExpertLayout + }; + + uint32_t n = params.problemShape.n(); + BlockEpilogue2 blockEpilogue2(resource, epilogueParams2); + BlockEpilogue3 blockEpilogue3(resource, epilogueParams3); + BlockEpilogue1 blockEpilogue1(resource, n); + + // Synchronous wait: SwiGLU waits for GMM1 [1] + AscendC::CrossCoreWaitFlag<0x2>(SYNCFLAGC2V); + AscendC::SyncAll(); + if (dequantSum1 > 0) { + uint32_t rowStartThisCore = 0; + MatrixCoord offsetC{0U, 0}; + MatrixCoord shapeC{dequantSum1, params.problemShape.n()}; + LayoutC layoutC{dequantSum1, params.problemShape.n()}; + int64_t gmOffsetC = layoutC.GetOffset(offsetC); + int64_t gmOffsetD = params.layoutD1.GetOffset(offsetC); + blockEpilogue1(gmC[gmOffsetC], shapeC, gmPerTokenScale1[rowStartThisCore], gmPermutedToken[gmOffsetD], + gmPerTokenScale2[rowStartThisCore], resource, params.epilogueCoreNum, params.swigluLimit, params.problemShape.k()); + } + AscendC::SyncAll(); + // Synchronization signal: SwiGLU notifies GMM2 [1] + AscendC::CrossCoreSetFlag<0x2, PIPE_MTE3>(SYNCFLAGV2C); + + if ((params.epilogueGranularity < params.expertPerRank && params.epilogueGranularity > 0)) { + // Synchronous wait: SwiGLU waits for GMM1 [2] + AscendC::CrossCoreWaitFlag<0x2>(SYNCFLAGC2V); + AscendC::SyncAll(); + if (dequantSum2 > 0) { + uint32_t rowStartThisCore = dequantSum1; + MatrixCoord offsetC{rowStartThisCore, 0}; + uint32_t dequantLen = dequantSum2; + MatrixCoord shapeC{dequantLen, params.problemShape.n()}; + LayoutC layoutC{dequantLen, params.problemShape.k()}; + int64_t gmOffsetC = layoutC.GetOffset(offsetC); + int64_t gmOffsetD = params.layoutD1.GetOffset(offsetC); + blockEpilogue1(gmC[gmOffsetC], shapeC, gmPerTokenScale1[rowStartThisCore], gmPermutedToken[gmOffsetD], + gmPerTokenScale2[rowStartThisCore], resource, coreNum, params.swigluLimit, params.problemShape.k()); + } + AscendC::SyncAll(); + // Synchronization signal: SwiGLU notifies GMM2 [2] + AscendC::CrossCoreSetFlag<0x2, PIPE_MTE3>(SYNCFLAGV2C); + } + + blockEpilogue1.Finalize(); + if (isCombineV1) { + blockEpilogue2.SetFlag(); + CombineV1(params, blockEpilogue2); + } else { + blockEpilogue3.SetFlag(); + CombineV2(params, blockEpilogue3); + } + + + + AscendC::SyncAll(); + ResetTokenPerExpert(params.EP * paddedExpertNumAligned); + + shmem.CrossRankSync(); + + MoeTokenUnpermuteTilingData tilingData; + MoeTokenUnpermuteTiling(params.problemShape.m() * params.topK, n2, params.topK, tilingData, coreNum); + KernelMoeTokenUnpermute kernelMoeTokenUnpermuteOp; + kernelMoeTokenUnpermuteOp.Init(shmem() + peermemInfo.offsetD, workspaceInfo.expandedRowIdx, params.probs, reinterpret_cast(params.ptrOutput), &tilingData); + kernelMoeTokenUnpermuteOp.Process(); + } + + CATLASS_DEVICE + void CombineV1(Params const ¶ms, BlockEpilogue2 & blockEpilogue) { + uint32_t n2 = params.problemShape.k(); + int32_t prevGroupSum2 = 0; + + icache_preload(8); + for (uint32_t t_groupIdx = 0; t_groupIdx < params.expertPerRank; ++t_groupIdx) { + int32_t flagId = t_groupIdx / CROSS_CORE_FLAG_MAX_SET_COUNT; + AscendC::CrossCoreWaitFlag<0x2>(flagId); + AscendC::SyncAll(); + + uint32_t groupIdx = t_groupIdx; + + for(int32_t dstEpIdx = coreIdx; dstEpIdx < params.EP; dstEpIdx += coreNum) { + __gm__ void* dstPeermemPtr = shmem(peermemInfo.offsetD, dstEpIdx); + AscendC::GlobalTensor gmRemotePeer; + gmRemotePeer.SetGlobalBuffer(reinterpret_cast<__gm__ ElementD2*>(dstPeermemPtr)); + uint32_t srcRowOffset = (dstEpIdx == 0 ? 0 : cumsumMM((dstEpIdx - 1) * params.expertPerRank + groupIdx)) + prevGroupSum2; + if (srcRowOffset < params.maxOutputSize) { + uint32_t dataRows = tokenPerExpert(tokenPerExpertLayout(dstEpIdx, params.rank, groupIdx)); + if (srcRowOffset + dataRows > params.maxOutputSize) { + dataRows = params.maxOutputSize - srcRowOffset; + } + //uint32_t dstRowOffset = preSumBeforeRank(2 * dstEpIdx * FLAGSTRIDE + groupIdx); + int32_t tmpBlock = AlignUp(params.expertPerRank, FLAGSTRIDE); + //uint32_t dstRowOffset = preSumBeforeRank(dstEpIdx * tmpBlock + groupIdx); + uint32_t dstRowOffset = preSumBeforeRank(dstEpIdx * params.expertPerRank + groupIdx); + MatrixCoord offsetC{srcRowOffset, 0}; + MatrixCoord offsetPeer{dstRowOffset, 0}; + MatrixCoord shapeC{dataRows, n2}; + int64_t gmOffsetC = params.layoutD2.GetOffset(offsetC); + int64_t gmOffsetPeer = params.layoutD2.GetOffset(offsetPeer); + if constexpr (std::is_same_v) { + blockEpilogue(gmC2[gmOffsetC], shapeC, gmPerTokenScale2[srcRowOffset], gmRemotePeer[gmOffsetPeer]); + } else { + blockEpilogue(gmC2[gmOffsetC], shapeC, gmRemotePeer[gmOffsetPeer]); + } + } + } + prevGroupSum2 += cumsumMM((params.EP - 1) * params.expertPerRank + groupIdx); + } + blockEpilogue.Finalize(); + } + + CATLASS_DEVICE + void CombineV2(Params const ¶ms, BlockEpilogue3 & blockEpilogue) { + BlockScheduler blockScheduler; + int32_t syncLoopIdx = 0; + uint32_t startCoreIdx = 0; + uint32_t aicCoreNum = coreNum / 2; + uint32_t aicCoreIdx = get_block_idx(); + uint32_t aivSubCoreIdx = get_subblockid(); + uint32_t preSrcExpertSum = 0; + uint32_t n2 = params.problemShape.k(); + uint32_t k2 = params.problemShape.n() / 2; + icache_preload(8); + for (uint32_t groupIdx = 0; groupIdx < params.expertPerRank; ++groupIdx) { + uint32_t currentExpertM = cumsumMM((params.EP - 1) * params.expertPerRank + groupIdx); + if (preSrcExpertSum >= params.maxOutputSize) { + currentExpertM = 0; + } else if (preSrcExpertSum + currentExpertM > params.maxOutputSize) { + currentExpertM = params.maxOutputSize - preSrcExpertSum; + } + GemmCoord inGroupProblemShape{currentExpertM, n2, k2}; // M N K + blockScheduler.Update(inGroupProblemShape, MakeCoord(L1TileShape::M, L1TileShape::N)); + uint32_t coreLoops = blockScheduler.GetCoreLoops(); + uint32_t startLoopIdx = ((aicCoreIdx < startCoreIdx) ? (aicCoreIdx + aicCoreNum) : aicCoreIdx) - startCoreIdx; + + for (uint32_t loopIdx = startLoopIdx; loopIdx < coreLoops; loopIdx += aicCoreNum) { + GemmCoord blockCoord = blockScheduler.GetBlockCoord(loopIdx); + GemmCoord actualBlockShape = blockScheduler.GetActualBlockShape(blockCoord); + int32_t m0 = 16; + // Block count, the shape of each block is (m0, actualBlockShape.n()) + int32_t m_rows = (actualBlockShape.m() + m0 - 1) / m0; + int32_t aiv_m_rows = m_rows / 2; + if (aivSubCoreIdx == 1 && aiv_m_rows * 2 < m_rows) { + aiv_m_rows += 1; + } + uint32_t m_offset = blockCoord.m() * L1TileShape::M;//blockOffset + if(aivSubCoreIdx == 1) { + m_offset += (m_rows / 2) * m0; + } + + + for (;syncLoopIdx <= groupIdx; syncLoopIdx ++) { + int32_t flag_id = syncLoopIdx / CROSS_CORE_FLAG_MAX_SET_COUNT; + AscendC::CrossCoreWaitFlag<0x2>(flag_id); + } + + for (int32_t cur_row = 0; cur_row < aiv_m_rows; cur_row ++) { + GemmCoord realTileCoord{m_offset, blockCoord.n() * L1TileShape::N, 1}; + uint32_t actualm = m0; + if(aivSubCoreIdx == 1 && cur_row == aiv_m_rows - 1){ + actualm = actualBlockShape.m() - (m_rows / 2) * m0 - cur_row * m0; + } + GemmCoord realTileShape{actualm, actualBlockShape.n(), 1}; + blockEpilogue(gmC2, gmPerTokenScale2, realTileCoord, realTileShape, groupIdx, preSrcExpertSum, preSumBeforeRank); + m_offset += m0; + } + } + preSrcExpertSum += currentExpertM; + startCoreIdx = (startCoreIdx + coreLoops) % aicCoreNum; + } + blockEpilogue.Finalize(); + } + +private: + struct WorkspaceInfo { + GM_ADDR ptrA; + GM_ADDR ptrPerTokenScale; + GM_ADDR ptrcumsumMM; + GM_ADDR ptrC; + GM_ADDR ptrC2; + GM_ADDR ptrPermutedToken; + GM_ADDR ptrPerTokenScale2; + GM_ADDR expandedRowIdx; + GM_ADDR ptrTokenPerExpert; + GM_ADDR ptrSumBeforeRank; + __gm__ float* ptrSoftFlagBase; + + + CATLASS_DEVICE + WorkspaceInfo(){} + + CATLASS_DEVICE + WorkspaceInfo(const Params & params) { + uint32_t k2 = params.problemShape.n() / 2; + uint32_t n2 = params.problemShape.k(); + int64_t workspaceOffset = 0; + expandedRowIdx = params.ptrWorkspace; + + workspaceOffset += AlignUp(params.problemShape.m(), 256) * params.topK * sizeof(int32_t); + ptrcumsumMM = params.ptrWorkspace + workspaceOffset; + + workspaceOffset += (params.EP * params.EP * params.expertPerRank) * sizeof(int32_t); + + ptrPerTokenScale = params.ptrWorkspace + workspaceOffset; + + workspaceOffset += params.maxOutputSize * sizeof(ElementPerTokenScale); + ptrPerTokenScale2 = params.ptrWorkspace + workspaceOffset; + + workspaceOffset += params.maxOutputSize * sizeof(ElementPerTokenScale); + ptrTokenPerExpert = params.ptrWorkspace + workspaceOffset; + + workspaceOffset += (params.EP * params.EP * params.expertPerRank) * sizeof(int32_t); + ptrC = params.ptrWorkspace + workspaceOffset; + + ptrC2 = params.ptrWorkspace + workspaceOffset; + + workspaceOffset += params.maxOutputSize * n2 * sizeof(ElementC); + ptrA = params.ptrWorkspace + workspaceOffset; + ptrPermutedToken = params.ptrWorkspace + workspaceOffset; + + workspaceOffset += params.maxOutputSize * params.problemShape.k() * sizeof(ElementA); + + ptrSumBeforeRank = params.ptrWorkspace + workspaceOffset; + + workspaceOffset += params.EP * sizeof(int32_t) * FLAGSTRIDE; + ptrSoftFlagBase = reinterpret_cast<__gm__ float*>(params.ptrWorkspace + workspaceOffset); + } + }; + + struct PeermemInfo { + int64_t offsetA; + int64_t offsetPeerPerTokenScale; + int64_t offsetPeerTokenPerExpert; + int64_t offsetD; + + CATLASS_DEVICE + PeermemInfo(){} + + CATLASS_DEVICE + PeermemInfo(const Params & params, const HcclShmem & shmem) { + offsetA = 0; // Occupies one third of BUFFSIZE + offsetPeerPerTokenScale = offsetA + AlignUp(shmem.SegmentSize() / 3, 512); // Occupies 1 MB + offsetD = offsetPeerPerTokenScale + MB_SIZE; // Occupies the remaining space + offsetPeerTokenPerExpert = shmem.SegmentSize() - 2 * MB_SIZE; // Occupies the final 2 MB + } + }; + + Arch::Resource resource; + + uint32_t coreIdx; + uint32_t coreNum; + + WorkspaceInfo workspaceInfo; + PeermemInfo peermemInfo; + + AscendC::GlobalTensor gmA; + AscendC::GlobalTensor gmC; + + AscendC::GlobalTensor gmPermutedToken; + AscendC::GlobalTensor gmC2; + + AscendC::GlobalTensor gmPerTokenScale1; + AscendC::GlobalTensor gmPerTokenScale2; + + AscendC::GlobalTensor gmXActiveMask; + + AscendC::GlobalTensor tokenPerExpert; + AscendC::GlobalTensor cumsumMM; + AscendC::GlobalTensor preSumBeforeRank; + Layout3D tokenPerExpertLayout; + HcclShmem shmem; + int32_t paddedExpertNumAligned; + bool isCombineV1; +}; + +} // namespace Catlass::Gemm::Kernel + +#endif // DISPATCH_FFN_COMBINE_KERNEL_HPP \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine_tiling.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine_tiling.h new file mode 100644 index 000000000..944d453b3 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine_tiling.h @@ -0,0 +1,58 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dispatch_ffn_combine_tiling.h + * \brief + */ + +#include "moe_init_routing_quant_v2/moe_init_routing_v2_tiling.h" +#include "moe_init_routing_quant_v2/moe_init_routing_quant_v2_tiling.h" + +#ifndef ASCENDC_DISPATCH_FFN_COMBINE_TILING_H +#define ASCENDC_DISPATCH_FFN_COMBINE_TILING_H +struct DispatchFFNCombineInfo { + uint32_t M; + uint32_t K; + uint32_t N; + uint32_t expertPerRank; + uint32_t maxOutputSize; + uint32_t isTransposeB; + uint32_t isWeightNz; + uint32_t aivNum; + uint32_t totalUbSize; + uint32_t topK; + uint32_t worldSize; + uint32_t listLen; + float swigluLimit; +}; + +struct CoCTiling { + int32_t m0 = -1; + int32_t k0 = -1; + int32_t n0 = -1; + int32_t swizzleDirect = -1; + int32_t swizzleOffset = -1; + int32_t ubMoveNum = -1; + int32_t pValue = -1; + int32_t commNpuSplit = -1; + int32_t commDataSplit = -1; + int32_t lenPerLoop = -1; + uint64_t initRoutingQuantTilingKey; + optiling::MoeInitRoutingQuantV2TilingData moeInitRoutingQuantV2TilingData; +}; + +struct DispatchFFNCombineTilingData { + Mc2InitTiling mc2InitTiling; + Mc2CcTiling mc2CcTiling; + DispatchFFNCombineInfo dispatchFFNCombineInfo; + CoCTiling cocTiling; +}; +#endif \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_init_routing_quant_v2.cpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_init_routing_quant_v2.cpp new file mode 100644 index 000000000..9180b06da --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_init_routing_quant_v2.cpp @@ -0,0 +1,133 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_init_routing_quant_v2.cpp + * \brief + */ +#include "moe_v2_sort_one_core.h" +#include "moe_v2_sort_multi_core.h" +#include "moe_v2_mrgsort_out.h" +#include "moe_v2_mrgsort.h" +#include "moe_v2_expert_token_out.h" +#include "moe_v2_src_to_dst_op.h" +#include "moe_v2_src_to_dst_with_capacity.h" +#include "moe_v2_fullload_quant.h" +#include "moe_v2_fullload_dynamic_quant.h" +#include "moe_v2_gather_quant.h" +#include "moe_v2_gather_dynamic_quant.h" +#include "moe_v2_src_to_dst_and_gather.h" + +using namespace AscendC; +using namespace MoeInitRoutingQuantV2; +using namespace optiling; + +template +__aicore__ inline void moe_init_routing_quant_v2( + GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR offset, GM_ADDR expandedX, GM_ADDR expandedRowIdx, + GM_ADDR expertTokensCountOrCumsum, GM_ADDR expertTokensBeforeCapacity, GM_ADDR dynamicQuantScale, GM_ADDR workspace, + const MoeInitRoutingQuantV2TilingData* tilingData, uint64_t tilingKey) { + + if (g_coreType == AIC) { + return; + } + + if (workspace == nullptr) { + return; + } + + if (tilingKey == 20000) { // quant full load + TPipe sortPipe; + MoeV2FullLoadQuant op; + op.Init(x, expertIdx, scale, offset, expandedX, expandedRowIdx, expertTokensCountOrCumsum, workspace, tilingData, &sortPipe); + op.Process(); + sortPipe.Destroy(); + return; + } + + + else if (tilingKey == 21000) { // dynamic quant full load + TPipe sortPipe; + MoeV2FullLoadDynamicQuant op; + op.Init(x, expertIdx, expandedX, expandedRowIdx, expertTokensCountOrCumsum, scale, dynamicQuantScale, workspace, tilingData, + &sortPipe); + op.Process(); + sortPipe.Destroy(); + return; + } + + // sort + if (tilingKey == 10000 || tilingKey == 10100 || tilingKey == 11000 || tilingKey == 11100) { + TPipe sortPipe; + MoeV2SortOneCore op; + op.Init(expertIdx, expertTokensCountOrCumsum, expertTokensBeforeCapacity, workspace, + tilingData, &sortPipe); + op.Process(); + sortPipe.Destroy(); + } else if (tilingKey == 10010 || tilingKey == 10110 || tilingKey == 11010 || tilingKey== 11110) { + TPipe sortPipe; + MoeV2SortMultiCore op; + op.Init(expertIdx, expertTokensCountOrCumsum, expertTokensBeforeCapacity, workspace, + tilingData, &sortPipe); + op.Process(); + sortPipe.Destroy(); + } + + if (tilingKey == 10000 || tilingKey == 10010 || tilingKey ==11000 || tilingKey ==11010) { // No drop scenario + if (tilingData->expertTokensCountOrCumsumFlag != EXERPT_TOKENS_NONE) { + TPipe expertTokenOutPipe; + MoeV2ExpertTokenOut expertTokenOutOp; + expertTokenOutOp.Init(expertTokensCountOrCumsum, expertTokensBeforeCapacity, + expandedRowIdx, workspace, tilingData, &expertTokenOutPipe); + expertTokenOutOp.Process(); + expertTokenOutPipe.Destroy(); + } + TPipe srcToDstPipe; + MoeV2SrcToDstOp srcToDstOp; + srcToDstOp.Init(expandedRowIdx, workspace, tilingData, &srcToDstPipe); + srcToDstOp.Process(); + srcToDstPipe.Destroy(); + } else if (tilingKey ==10100 || tilingKey ==10110 || tilingKey ==11100 || tilingKey ==11110) { // Drop scenario + TPipe expertTokenOutPipe; + MoeV2ExpertTokenOut expertTokenOutOp; + expertTokenOutOp.Init(expertTokensCountOrCumsum, expertTokensBeforeCapacity, + expandedRowIdx, workspace, tilingData, &expertTokenOutPipe); + expertTokenOutOp.Process(); + expertTokenOutPipe.Destroy(); + + if (tilingKey == 10100 || tilingKey == 10110) { + TPipe srcToDstPipe; + MoeV2SrcToDstWithCapacity srcToDstWithCapacityOp; + srcToDstWithCapacityOp.Init(expandedRowIdx, expandedX, workspace, tilingData, &srcToDstPipe); + srcToDstWithCapacityOp.Process(); + srcToDstPipe.Destroy(); + } else { + TPipe srcToDstGatherPipe; + MoeV2SrcToDstAndGather srcToDstAndGatherOp; + srcToDstAndGatherOp.Init(x, scale, expandedRowIdx, expandedX, dynamicQuantScale, workspace, tilingData, &srcToDstGatherPipe); + srcToDstAndGatherOp.Process(); + srcToDstGatherPipe.Destroy(); + } + } + + if (tilingKey == 10000 || tilingKey == 10010 || tilingKey == 10100 || tilingKey == 10110) { + TPipe gatherPipe; + MoeV2GatherQuant gatherQuantOp; + gatherQuantOp.Init(x, scale, offset, expandedRowIdx, expandedX, workspace, tilingData, &gatherPipe); + gatherQuantOp.Process(); + gatherPipe.Destroy(); + } else if (tilingKey == 11000 || tilingKey == 11010) { + TPipe gatherPipe; + MoeV2GatherDynamicQuant gatherDynamicQuantOp; + gatherDynamicQuantOp.Init(x, scale, expandedRowIdx, expandedX, dynamicQuantScale, workspace, tilingData, &gatherPipe); + gatherDynamicQuantOp.Process(); + gatherPipe.Destroy(); + } +} diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_init_routing_quant_v2_tiling.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_init_routing_quant_v2_tiling.h new file mode 100644 index 000000000..08ced8acb --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_init_routing_quant_v2_tiling.h @@ -0,0 +1,434 @@ +#pragma once +#ifndef ASCENDC_DISPATCH_FFN_COMBINE_INIT_ROUTING_QUANT_TILING_H +#define ASCENDC_DISPATCH_FFN_COMBINE_INIT_ROUTING_QUANT_TILING_H + +#include "moe_init_routing_v2_tiling.h" + +namespace optiling { + +const static int64_t ATTR_QUANT_MODE = 6; +const static int64_t TILING_KEY_BASE = 10000; +const static int64_t TILING_KEY_PERF_BASE = 20000; +const static int64_t TILING_KEY_QUANT_BASE = 1000; +const static int64_t TILING_KEY_DROP_MODE_BASE = 100; +const static int64_t TILING_KEY_SORT_BASE = 10; +const static int64_t FOUR_BLOCK_BYTE = 128; +const static int64_t MAX_COLS_ONE_LOOP_QUANT = 8192; +const static int64_t INDEX_SCALE = 2; +const static int64_t INDEX_OFFSET = 3; +const static int64_t SMOOTH_NONE = 0; +const static int64_t SMOOTH_1H = 1; +const static int64_t SMOOTH_EH = 2; +const static int64_t MAX_COLS_DYNAMIC_QUANT = 6144; +const static int64_t DYNAMIC_QUANT_SRC_TO_DST_BUFFER = 15; +const static int64_t DYNAMIC_QUANT_COLS_BUFFER = 21; +const static int64_t DYNAMIC_QUANT_FULLLOAD_COLS_BUFFER = 13; +const static int64_t DYNAMIC_QUANT_SCALE_SIZE_64 = 64; +const static int64_t DYNAMIC_QUANT_SCALE_SIZE_128 = 128; +const static int64_t OUTOUT_DYNAMIC_QUANT_SCALE = 4; +const static int64_t FULLLOAD_H_LIMIT = 7168; + + +inline static int64_t AlignOneBlockByte(int64_t x) { + return (x + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; +} + +inline static int64_t AlignOneBlockByteCeil(int64_t x) { + return x / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; +} + +struct MoeInitRoutingQuantV2TilingData { + int64_t coreNum; + int64_t n; + int64_t cols; + int64_t k; + int64_t expertCapacity; + int64_t expertNum; + int64_t dropPadMode; + int64_t expertTokensCountOrCumsumFlag; + int64_t expertTokensBeforeCapacityFlag; + int64_t smoothType; + InnerMoeV2VBSComputeTilingData vbsComputeParamsOp; + InnerMoeV2VMSMiddleComputeTilingData vmsMiddleComputeParamsOp; + InnerMoeV2SortOutComputeTilingData sortOutComputeParamsOp; + InnerMoeV2GatherOutComputeTilingData srcToDstComputeParamsOp; + InnerMoeV2GatherOutComputeTilingData srcToDstCapacityComputeParamsOp; + InnerMoeV2GatherOutComputeTilingData gatherOutComputeParamsOp; +}; + + + +class MoeInitRoutingQuantV2TilingBase : public InnerMoeInitRoutingV2TilingBase { +public: +protected: + + bool GetShapeAttrsInfo(int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, + int64_t expertNum, int64_t activeNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag, + bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0) override; + uint64_t GetTilingKey() const override; + bool GetWorkspaceSize() override; + bool PostTiling() override; +public: + //bool CheckOutShape() override; + bool IsFullLoadQuant(int64_t space); + bool IsFullLoadDynamicQuant(int64_t space); + bool IsFullLoad() override; + void SetGatherTilingData(InnerMoeV2GatherOutComputeTilingData* tilingData, int64_t perCoreRows, int64_t lastCoreRows, + int64_t cols); + void SetGatherTilingDataCols(InnerMoeV2GatherOutComputeTilingData* tilingData, int64_t baseMaxCols, int64_t cols); + void SetGatherTilingDataRows(InnerMoeV2GatherOutComputeTilingData* tilingData, int64_t perCoreRows, + int64_t lastCoreRows, int64_t basePerLoopMaxRows); + void Tiling4GatherQuant(); + void Tiling4GatherDynamicQuant(); + void Tiling4SrcToDstCapacityCompute() override; + void Tiling4GatherOutCompute() override; + void CopyGatherOutTiling(InnerMoeV2GatherOutComputeTilingData& dst, InnerMoeV2GatherOutComputeTilingData& src); + void CopyTilingData(); + + + int64_t quantMode; + MoeInitRoutingQuantV2TilingData quantTilingData; +}; + + +bool MoeInitRoutingQuantV2TilingBase::IsFullLoadQuant(int64_t space) { + int64_t perCoreXRows = moeInitRoutingTilingData.n / aivNum; + int64_t remainder = moeInitRoutingTilingData.n % aivNum; + // NUM_TWO is Max xRows need add 2 becauseof the left and right row may be another row. + perCoreXRows = remainder <= 1 ? perCoreXRows + 1 : perCoreXRows + NUM_TWO; + int64_t quantBaseSpace = AlignOneBlockByte(moeInitRoutingTilingData.cols); + int64_t quantSpace = + quantBaseSpace * (inuptXDtypeSize_ + sizeof(int8_t) + sizeof(float) + sizeof(int16_t)) * perCoreXRows; + int64_t remainUbAfterSort = aicoreParams_.ubSize - space - quantSpace; + return remainUbAfterSort > 0; +} + +bool MoeInitRoutingQuantV2TilingBase::IsFullLoadDynamicQuant(int64_t space) { + int64_t quantSpace = AlignOneBlockByte(moeInitRoutingTilingData.cols) * DYNAMIC_QUANT_FULLLOAD_COLS_BUFFER; + int64_t scaleOutSpace = 64; + int64_t remainUbAfterSort = aicoreParams_.ubSize - space - scaleOutSpace - quantSpace; + return remainUbAfterSort > 0; +} + +bool MoeInitRoutingQuantV2TilingBase::IsFullLoad() { + if (totalLength > sortLoopMaxElement || moeInitRoutingTilingData.cols > MAX_COLS_ONE_LOOP_QUANT || + this->dropPadMode == 1) { + return false; + } + int64_t sortSpace = AlignOneBlockByte(this->totalLength) * sizeof(int32_t) * ONE_CORE_SORT_BUFFER; + int64_t otherSpace = AlignOneBlockByte(this->totalLength) * sizeof(int32_t) * NUM_THREE; + int64_t expertSpace = AlignOneBlockByte(this->expertNum * sizeof(int32_t)); + if (quantMode == 0) { + return IsFullLoadQuant(sortSpace + otherSpace + expertSpace); + } else { + return IsFullLoadDynamicQuant(sortSpace + otherSpace + expertSpace); + } +} + +bool MoeInitRoutingQuantV2TilingBase::GetShapeAttrsInfo(int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, + int64_t expertNum, int64_t activeNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag, + bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0) { + + InnerMoeInitRoutingV2TilingBase::GetShapeAttrsInfo(m, cols, topK, expertCapacity, expertNum, activeNum, dropPadMode, + expertTokensCountOrCumsumFlag, expertTokensBeforeCapacityFlag, inuptXDtypeSize, quantMode, scaleDim0); + this -> quantMode = quantMode; + if (quantMode == 0) { + } else { + if (scaleDim0 > 0) { + quantTilingData.smoothType = ((scaleDim0 == 1) ? SMOOTH_1H : SMOOTH_EH); + } else { + quantTilingData.smoothType = SMOOTH_NONE; + } + } + return true; +} + + +uint64_t MoeInitRoutingQuantV2TilingBase::GetTilingKey() const { + if (isFullLoad) { + return TILING_KEY_PERF_BASE + quantMode * TILING_KEY_QUANT_BASE; + } + return TILING_KEY_BASE + quantMode * TILING_KEY_QUANT_BASE + dropPadMode * TILING_KEY_DROP_MODE_BASE + + (totalLength > sortLoopMaxElement) * TILING_KEY_SORT_BASE; +} + + +bool MoeInitRoutingQuantV2TilingBase::PostTiling() { + CopyTilingData(); + return true; +} +void MoeInitRoutingQuantV2TilingBase::CopyGatherOutTiling(InnerMoeV2GatherOutComputeTilingData& dst, + InnerMoeV2GatherOutComputeTilingData& src) { + dst.needCoreNum = (src.needCoreNum); + dst.activateRows = (src.activateRows); + dst.perCoreRows = (src.perCoreRows); + dst.perCorePerLoopRows = (src.perCorePerLoopRows); + dst.perCoreLastLoopRows = (src.perCoreLastLoopRows); + dst.lastCoreRows = (src.lastCoreRows); + dst.lastCorePerLoopRows = (src.lastCorePerLoopRows); + dst.lastCoreLastLoopRows = (src.lastCoreLastLoopRows); + dst.perCoreLoops = (src.perCoreLoops); + dst.lastCoreLoops = (src.lastCoreLoops); + dst.perLoopCols = (src.perLoopCols); + dst.lastLoopCols = (src.lastLoopCols); + dst.colLoops = (src.colLoops); +} + +void MoeInitRoutingQuantV2TilingBase::CopyTilingData() { + quantTilingData.coreNum = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.coreNum); + quantTilingData.n = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.n); + quantTilingData.cols = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.cols); + quantTilingData.k = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.k); + quantTilingData.expertCapacity = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.expertCapacity); + quantTilingData.expertNum = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.expertNum); + quantTilingData.dropPadMode = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.dropPadMode); + quantTilingData.expertTokensCountOrCumsumFlag = ( + InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.expertTokensCountOrCumsumFlag); + quantTilingData.expertTokensBeforeCapacityFlag = ( + InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.expertTokensBeforeCapacityFlag); + + auto vbsTilingData = &InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.vbsComputeParamsOp; + quantTilingData.vbsComputeParamsOp.needCoreNum = (vbsTilingData->needCoreNum); + quantTilingData.vbsComputeParamsOp.perCoreElements = (vbsTilingData->perCoreElements); + quantTilingData.vbsComputeParamsOp.perCoreLoops = (vbsTilingData->perCoreLoops); + quantTilingData.vbsComputeParamsOp.perCorePerLoopElements = (vbsTilingData->perCorePerLoopElements); + quantTilingData.vbsComputeParamsOp.perCoreLastLoopElements = (vbsTilingData->perCoreLastLoopElements); + quantTilingData.vbsComputeParamsOp.lastCoreElements = (vbsTilingData->lastCoreElements); + quantTilingData.vbsComputeParamsOp.lastCoreLoops = (vbsTilingData->lastCoreLoops); + quantTilingData.vbsComputeParamsOp.lastCorePerLoopElements = (vbsTilingData->lastCorePerLoopElements); + quantTilingData.vbsComputeParamsOp.lastCoreLastLoopElements = (vbsTilingData->lastCoreLastLoopElements); + quantTilingData.vbsComputeParamsOp.oneLoopMaxElements = (vbsTilingData->oneLoopMaxElements); + + quantTilingData.vmsMiddleComputeParamsOp.needCoreNum = ( + InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.vmsMiddleComputeParamsOp.needCoreNum); + quantTilingData.sortOutComputeParamsOp.oneLoopMaxElements = ( + InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.sortOutComputeParamsOp.oneLoopMaxElements); + + CopyGatherOutTiling(quantTilingData.srcToDstComputeParamsOp, + InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.srcToDstComputeParamsOp); + CopyGatherOutTiling(quantTilingData.srcToDstCapacityComputeParamsOp, + InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp); +} + + +bool MoeInitRoutingQuantV2TilingBase::GetWorkspaceSize() { + InnerMoeInitRoutingV2TilingBase::GetWorkspaceSize(); + bool useCols = + (dropPadMode == 0 && quantTilingData.gatherOutComputeParamsOp.colLoops > 1) || + (dropPadMode == 1 && + InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp.colLoops > 1); + if (quantMode == 1 && useCols) { + workspaceSize_ += aivNum * InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.cols * sizeof(float); + } + return true; +} + +void MoeInitRoutingQuantV2TilingBase::SetGatherTilingData(InnerMoeV2GatherOutComputeTilingData* tilingData, + int64_t perCoreRows, int64_t lastCoreRows, int64_t cols) { + tilingData->perCorePerLoopRows = perCoreRows; + tilingData->perCoreLastLoopRows = perCoreRows; + tilingData->lastCorePerLoopRows = lastCoreRows; + tilingData->lastCoreLastLoopRows = lastCoreRows; + tilingData->perCoreLoops = 1; + tilingData->lastCoreLoops = 1; + tilingData->perLoopCols = cols; + tilingData->lastLoopCols = cols; + tilingData->colLoops = 1; +} + +void MoeInitRoutingQuantV2TilingBase::SetGatherTilingDataCols(InnerMoeV2GatherOutComputeTilingData* tilingData, + int64_t baseMaxCols, int64_t cols) { + tilingData->perLoopCols = (std::min(baseMaxCols, cols)); + tilingData->lastLoopCols = (GetPerOrLastValue(cols, baseMaxCols)); + tilingData->colLoops = (baseMaxCols == 0 ? 0 : (cols + baseMaxCols - 1) / baseMaxCols); +} + +void MoeInitRoutingQuantV2TilingBase::SetGatherTilingDataRows(InnerMoeV2GatherOutComputeTilingData* tilingData, + int64_t perCoreRows, int64_t lastCoreRows, + int64_t basePerLoopMaxRows) { + tilingData->perCorePerLoopRows = (std::min(perCoreRows, basePerLoopMaxRows)); + tilingData->perCoreLastLoopRows = (GetPerOrLastValue(perCoreRows, basePerLoopMaxRows)); + tilingData->perCoreLoops = (basePerLoopMaxRows == 0 ? 0 + : (perCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows); + tilingData->lastCorePerLoopRows = (std::min(lastCoreRows, basePerLoopMaxRows)); + tilingData->lastCoreLastLoopRows = (GetPerOrLastValue(lastCoreRows, basePerLoopMaxRows)); + tilingData->lastCoreLoops = (basePerLoopMaxRows == 0 ? 0 + : (lastCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows); +} + +void MoeInitRoutingQuantV2TilingBase::Tiling4SrcToDstCapacityCompute() { + if (quantMode == 0 || dropPadMode == 0) { + InnerMoeInitRoutingV2TilingBase::Tiling4SrcToDstCapacityCompute(); + return; + } + + auto tilingData = &moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp; + int64_t perCoreRows = CeilDiv(totalLength, aivNum); + if (perCoreRows <= 0) { + tilingData->needCoreNum = 0; + return; + } + + tilingData->needCoreNum = CeilDiv(totalLength, perCoreRows); + int64_t cols = moeInitRoutingTilingData.cols; + tilingData->perCoreRows = perCoreRows; + int64_t lastCoreRows = totalLength - perCoreRows * (tilingData->needCoreNum - 1); + tilingData->lastCoreRows = lastCoreRows; + + int64_t rowSize = AlignOneBlockByte(perCoreRows * sizeof(int32_t)) * NUM_FOUR; + int64_t colSize = AlignOneBlockByte(cols * sizeof(int8_t)) * DYNAMIC_QUANT_SRC_TO_DST_BUFFER; + int64_t scaleSize = DYNAMIC_QUANT_SCALE_SIZE_64; + if (rowSize + colSize + scaleSize < static_cast(aicoreParams_.ubSize)) { + + SetGatherTilingData(tilingData, perCoreRows, lastCoreRows, cols); + } else { + + int64_t baseMaxCols = MAX_COLS_DYNAMIC_QUANT; + int64_t totalColSize = AlignOneBlockByte(baseMaxCols * sizeof(int8_t)) * DYNAMIC_QUANT_SRC_TO_DST_BUFFER; + int64_t ubSize = static_cast(aicoreParams_.ubSize); + int64_t basePerLoopMaxRows = + AlignOneBlockByteCeil((ubSize - totalColSize - scaleSize) / sizeof(int32_t)) / NUM_FOUR; + if (cols < MAX_COLS_DYNAMIC_QUANT) { + basePerLoopMaxRows = AlignOneBlockByteCeil((ubSize - colSize - scaleSize) / sizeof(int32_t)) / NUM_FOUR; + } else if (perCoreRows < basePerLoopMaxRows) { + baseMaxCols = AlignOneBlockByteCeil(ubSize - rowSize - scaleSize) / DYNAMIC_QUANT_SRC_TO_DST_BUFFER; + } + SetGatherTilingDataCols(tilingData, baseMaxCols, cols); + SetGatherTilingDataRows(tilingData, perCoreRows, lastCoreRows, basePerLoopMaxRows); + } +} + + +void MoeInitRoutingQuantV2TilingBase::Tiling4GatherQuant() { + auto tilingData = &quantTilingData.gatherOutComputeParamsOp; + tilingData->activateRows = totalLength; + if (dropPadMode == 0 && activateNum > 0) { + tilingData->activateRows = (std::min(activateNum, totalLength)); + } + int64_t perCoreRows = CeilDiv(totalLength, aivNum); + + if (perCoreRows <= 0) { + tilingData->needCoreNum = 0; + return; + } + + tilingData->needCoreNum = (CeilDiv(totalLength, perCoreRows)); + int64_t cols = moeInitRoutingTilingData.cols; + tilingData->perCoreRows = perCoreRows; + int64_t lastCoreRows = totalLength - perCoreRows * (tilingData->needCoreNum - 1); + tilingData->lastCoreRows = lastCoreRows; + int64_t sizeOfCol = sizeof(int8_t) * NUM_TWO + sizeof(float) + sizeof(int16_t) + inuptXDtypeSize_ * NUM_TWO; + int64_t rowSize = AlignOneBlockByte((perCoreRows * sizeof(int32_t) * NUM_TWO)); + int64_t colSize = AlignOneBlockByte(cols * sizeOfCol); + if (rowSize + colSize < static_cast(aicoreParams_.ubSize) / NUM_TWO) { + SetGatherTilingData(tilingData, perCoreRows, lastCoreRows, cols); + } else { + int64_t baseMaxCols = MAX_COLS_ONE_LOOP_QUANT; + int64_t baseMaxColsSize = AlignOneBlockByte(baseMaxCols * sizeOfCol); + int64_t ubSize = static_cast(aicoreParams_.ubSize); + int64_t basePerLoopMaxRows = AlignOneBlockByteCeil((ubSize - baseMaxColsSize) / NUM_TWO / sizeof(int32_t)); + if (cols < MAX_COLS_ONE_LOOP_QUANT) { + basePerLoopMaxRows = AlignOneBlockByteCeil((ubSize - colSize) / NUM_TWO / sizeof(int32_t)); + } else if (perCoreRows < basePerLoopMaxRows) { + baseMaxCols = AlignOneBlockByteCeil((ubSize - rowSize) / sizeOfCol); + } + SetGatherTilingDataCols(tilingData, baseMaxCols, cols); + SetGatherTilingDataRows(tilingData, perCoreRows, lastCoreRows, basePerLoopMaxRows); + } +} + + + +void SetGatherTilingDatawithloop(InnerMoeV2GatherOutComputeTilingData* tilingData, + int64_t perCorePerLoopRows, int64_t lastCorePerLoopRows, int64_t cols, + int64_t perCoreLastLoopRows = 1, int64_t lastCoreLastLoopRows = 1, + int64_t perCoreLoops = 1, int64_t lastCoreLoops = 1) { + tilingData-> perCorePerLoopRows = perCorePerLoopRows; + tilingData-> perCoreLastLoopRows = perCoreLastLoopRows; + tilingData-> lastCorePerLoopRows = lastCorePerLoopRows; + tilingData-> lastCoreLastLoopRows = lastCoreLastLoopRows; + tilingData-> perCoreLoops = perCoreLoops; + tilingData-> lastCoreLoops = lastCoreLoops; + tilingData-> perLoopCols = cols; + tilingData-> lastLoopCols = cols; + tilingData-> colLoops = 1; +} + +void MoeInitRoutingQuantV2TilingBase::Tiling4GatherDynamicQuant() { + + auto tilingData = &quantTilingData.gatherOutComputeParamsOp; + tilingData->activateRows = totalLength; + if (dropPadMode == 0 && activateNum > 0) { + tilingData->activateRows = (std::min(activateNum, totalLength)); + } + int64_t perCoreRows = CeilDiv(totalLength, aivNum); + + if (perCoreRows <= 0) { + tilingData->needCoreNum = 0; + return; + } + + tilingData->needCoreNum = (CeilDiv(totalLength, perCoreRows)); + + int64_t cols = InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.cols; + + tilingData->perCoreRows = perCoreRows; + int64_t lastCoreRows = totalLength - perCoreRows * (tilingData->needCoreNum - 1); + tilingData->lastCoreRows = lastCoreRows; + + + int64_t rowSize = AlignOneBlockByte(perCoreRows * sizeof(int32_t)) * NUM_FOUR; + int64_t colSize = AlignOneBlockByte(cols * sizeof(int8_t)) * DYNAMIC_QUANT_COLS_BUFFER; + int64_t scaleSize = DYNAMIC_QUANT_SCALE_SIZE_64; + int64_t onceRowSize = (static_cast(aicoreParams_.ubSize) - + colSize - scaleSize - + ONE_BLOCK_BYTE * NUM_FOUR * NUM_THREE) / + (sizeof(int32_t) * NUM_FOUR); + int64_t oneBlockNumInt = static_cast(ONE_BLOCK_BYTE) / static_cast(sizeof(int32_t)); + onceRowSize = onceRowSize / oneBlockNumInt * oneBlockNumInt; + bool ifOneLoop = ((static_cast(aicoreParams_.ubSize) > colSize + + scaleSize + ONE_BLOCK_BYTE * NUM_FOUR * NUM_FOUR) && + quantTilingData.smoothType == SMOOTH_NONE && + cols == FULLLOAD_H_LIMIT); + + int64_t perCoreOnceRowSize = ifOneLoop ? std::min(onceRowSize, perCoreRows) : perCoreRows; + int64_t lastCoreOnceRowSize = ifOneLoop ? std::min(onceRowSize, lastCoreRows) : lastCoreRows; + int64_t perCoreLoops = ifOneLoop ? CeilDiv(perCoreRows, perCoreOnceRowSize) : 1; + int64_t lastCoreLoops = ifOneLoop ? CeilDiv(lastCoreRows, lastCoreOnceRowSize) : 1; + int64_t perCoreLastLoopRows = ifOneLoop ? GetPerOrLastValue(perCoreRows, perCoreOnceRowSize) : perCoreRows; + int64_t lastCoreLastLoopRows = ifOneLoop ? GetPerOrLastValue(lastCoreRows, lastCoreOnceRowSize) : lastCoreRows; + + if (rowSize + colSize + scaleSize < static_cast(aicoreParams_.ubSize) || ifOneLoop) { + + SetGatherTilingDatawithloop(tilingData, perCoreOnceRowSize, lastCoreOnceRowSize, cols, + perCoreLastLoopRows, lastCoreLastLoopRows, + perCoreLoops, lastCoreLoops); + } else { + int64_t baseMaxCols = MAX_COLS_DYNAMIC_QUANT; + int64_t totalColSize = AlignOneBlockByte(baseMaxCols * sizeof(int8_t)) * DYNAMIC_QUANT_COLS_BUFFER; + int64_t ubSize = static_cast(aicoreParams_.ubSize); + int64_t basePerLoopMaxRows = + AlignOneBlockByteCeil((ubSize - totalColSize - scaleSize) / sizeof(int32_t)) / NUM_FOUR; + if (cols < MAX_COLS_DYNAMIC_QUANT) { + basePerLoopMaxRows = AlignOneBlockByteCeil((ubSize - colSize - scaleSize) / sizeof(int32_t)) / NUM_FOUR; + } else if (perCoreRows < basePerLoopMaxRows) { + baseMaxCols = AlignOneBlockByteCeil(ubSize - rowSize - scaleSize) / DYNAMIC_QUANT_COLS_BUFFER; + } + SetGatherTilingDataCols(tilingData, baseMaxCols, cols); + SetGatherTilingDataRows(tilingData, perCoreRows, lastCoreRows, basePerLoopMaxRows); + } +} + + +void MoeInitRoutingQuantV2TilingBase::Tiling4GatherOutCompute() { + if (quantMode == 0) { + Tiling4GatherQuant(); + } else { + Tiling4GatherDynamicQuant(); + } +} + + +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_init_routing_v2_tiling.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_init_routing_v2_tiling.h new file mode 100644 index 000000000..5f136a9a3 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_init_routing_v2_tiling.h @@ -0,0 +1,415 @@ +#pragma once + +#ifndef ASCENDC_DISPATCH_FFN_COMBINE_INIT_ROUTING_TILING_H +#define ASCENDC_DISPATCH_FFN_COMBINE_INIT_ROUTING_TILING_H + +#include "tiling_base.h" + + +namespace optiling { +const static int64_t TILING_KEY_DROPLESS_SORT_ONE_CORE = 10001; +const static int64_t TILING_KEY_DROPLESS_SORT_MULTI_CORE = 10002; +const static int64_t TILING_KEY_DROP_PAD_MODE_SORT_ONE_CORE = 10011; +const static int64_t TILING_KEY_DROP_PAD_MODE_SORT_MULTI_CORE = 10012; +const static int64_t TILING_KEY_HIGH_PERFORMANCE = 20000; +const static int64_t NUM_TWO = 2; +const static int64_t NUM_THREE = 3; +const static int64_t NUM_FOUR = 4; +const static int64_t MRG_LIST_NUM = 4; +const static int64_t SORT32_ALIGN_ELEMENT = 32; +const static int64_t ONE_BLOCK_BYTE = 32; +const static size_t DIM_ONE = 1; +const static size_t DIM_TWO = 2; +const static size_t DIM_THREE = 3; +const static int32_t SIZE_16 = 16; +const static int32_t LENGTH_1024 = 1024; +const static int64_t MAX_COLS_ONE_LOOP = 16376; +const static int64_t ASSIST_NUM = 256; +const static int64_t INDEX_INPUT_X = 0; +const static int64_t INDEX_INPUT_EXPERT_IDX = 1; +const static int64_t ATTR_ACTIVE_ROWS = 0; +const static int64_t ATTR_EXPERT_CAPACITY = 1; +const static int64_t ATTR_EXPERT_NUM = 2; +const static int64_t ATTR_DROP_PAD_MODE = 3; +const static int64_t ATTR_EXPERT_TOKENS_COUNT_OR_CUMSUM_FLAG = 4; +const static int64_t ATTR_EXPERT_TOKENS_BEFORE_CAPACITY_FLAG = 5; +const static int64_t OUTOUT_EXPANDED_X = 0; +const static int64_t OUTOUT_EXPANDED_ROW_IDX = 1; +const static int64_t OUTOUT_EXPERT_TOKENS_COUNT_OR_CUMSUM = 2; +const static int64_t OUTOUT_EXPERT_TOKENS_BEFORE_CAPACITY = 3; +const static int64_t KV_FACTOR = 2; +const static int64_t ONE_CORE_SORT_BUFFER = 6; +const static int64_t EXPERT_TOKENS_COUNT = 2; + + +inline static int64_t CeilLog4(int64_t x) { + return static_cast(std::ceil(std::log(x) / std::log(NUM_FOUR))); +} + +inline static int64_t GetPerOrLastValue(int64_t x, int64_t y) { + if (y == 0) { + return 0; + } + return x <= y ? x : x % y; +} + +template +constexpr T CeilDiv(const T dividend, const T divisor) +{ + return (dividend + divisor - 1) / divisor; +} + + +struct InnerMoeV2VBSComputeTilingData { + int64_t needCoreNum = 0; + int64_t perCoreElements = 0; + int64_t perCoreLoops = 0; + int64_t perCorePerLoopElements = 0; + int64_t perCoreLastLoopElements = 0; + int64_t lastCoreElements = 0; + int64_t lastCoreLoops = 0; + int64_t lastCorePerLoopElements = 0; + int64_t lastCoreLastLoopElements = 0; + int64_t oneLoopMaxElements = 0; +}; + +struct InnerMoeV2VMSMiddleComputeTilingData { + int64_t needCoreNum = 0; +}; + +struct InnerMoeV2SortOutComputeTilingData { + int64_t oneLoopMaxElements = 0; +}; + +struct InnerMoeV2GatherOutComputeTilingData { + int64_t needCoreNum = 0; + int64_t activateRows = 0; + int64_t perCoreRows = 0; + int64_t perCorePerLoopRows = 0; + int64_t perCoreLastLoopRows = 0; + int64_t lastCoreRows = 0; + int64_t lastCorePerLoopRows = 0; + int64_t lastCoreLastLoopRows = 0; + int64_t perCoreLoops = 0; + int64_t lastCoreLoops = 0; + int64_t perLoopCols = 0; + int64_t lastLoopCols = 0; + int64_t colLoops = 0; +}; + +struct InnerMoeInitRoutingV2TilingData { + int64_t coreNum; + int64_t n; + int64_t cols; + int64_t k; + int64_t expertCapacity; + int64_t expertNum; + int64_t dropPadMode; + int64_t expertTokensCountOrCumsumFlag; + int64_t expertTokensBeforeCapacityFlag; + InnerMoeV2VBSComputeTilingData vbsComputeParamsOp; + InnerMoeV2VMSMiddleComputeTilingData vmsMiddleComputeParamsOp; + InnerMoeV2SortOutComputeTilingData sortOutComputeParamsOp; + InnerMoeV2GatherOutComputeTilingData srcToDstComputeParamsOp; + InnerMoeV2GatherOutComputeTilingData srcToDstCapacityComputeParamsOp; + InnerMoeV2GatherOutComputeTilingData gatherOutComputeParamsOp; +}; + + +class InnerMoeInitRoutingV2TilingBase : public TilingBaseClass { + +protected: + bool GetPlatformInfo(int64_t aivCoreNum, int64_t ubSizePlatForm) override; + bool GetShapeAttrsInfo(int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, + int64_t expertNum, int64_t activeNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag, + bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0) override; + + bool DoOpTiling() override; + uint64_t GetTilingKey() const override; + bool GetWorkspaceSize() override; + + +protected: + bool CheckTokenCount(int64_t num, const char* tag); + + virtual void Tiling4GatherOutCompute() = 0; + void Tiling4SrcToDstCompute(); + virtual void Tiling4SrcToDstCapacityCompute(); + void Tiling4SortOutCompute(); + void Tiling4VMSMiddleCompute(); + void Tiling4VBSCompute(); + void ShowTilingData(); + void Tiling4VBSMultiCoreCompute(InnerMoeV2VBSComputeTilingData* tilingData); + void Tiling4VBSOneCoreCompute(InnerMoeV2VBSComputeTilingData* tilingData); + virtual bool IsFullLoad() = 0; + + + + + int64_t aivNum = 0; + int64_t sortLoopMaxElement = 0; + int64_t mrgSortListMaxElement = 2040; + int64_t totalLength = 0; + int64_t activateNum = 0; + int64_t expertCapacity = 0; + int64_t expertNum = 0; + int64_t dropPadMode = 0; + int64_t expertTokensCountOrCumsumFlag = 0; + bool expertTokensBeforeCapacityFlag = false; + int64_t inuptXDtypeSize_ = 0; + bool isFullLoad = false; + + InnerMoeInitRoutingV2TilingData moeInitRoutingTilingData; +}; + + +bool InnerMoeInitRoutingV2TilingBase::DoOpTiling() { + sortLoopMaxElement = + (aicoreParams_.ubSize) / (sizeof(int32_t) * NUM_TWO * NUM_FOUR) / SORT32_ALIGN_ELEMENT * SORT32_ALIGN_ELEMENT; + isFullLoad = IsFullLoad(); + Tiling4VBSCompute(); + Tiling4VMSMiddleCompute(); + Tiling4SortOutCompute(); + Tiling4SrcToDstCompute(); + Tiling4SrcToDstCapacityCompute(); + Tiling4GatherOutCompute(); + return true; +}; + +uint64_t InnerMoeInitRoutingV2TilingBase::GetTilingKey() const { + if (isFullLoad) { + return TILING_KEY_HIGH_PERFORMANCE; + } + if (dropPadMode == 0) { + if (totalLength <= sortLoopMaxElement) { // Sorting uses only one core + return TILING_KEY_DROPLESS_SORT_ONE_CORE; + } else { + return TILING_KEY_DROPLESS_SORT_MULTI_CORE; + } + } else { + if (totalLength <= sortLoopMaxElement) { + return TILING_KEY_DROP_PAD_MODE_SORT_ONE_CORE; + } else { + return TILING_KEY_DROP_PAD_MODE_SORT_MULTI_CORE; + } + } + return tilingKey_; +} + + + +bool InnerMoeInitRoutingV2TilingBase::GetShapeAttrsInfo(int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, + int64_t expertNum, int64_t activateNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag, + bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0) { + + this->activateNum = activateNum; + this->expertCapacity = expertCapacity; + this->expertNum = expertNum; + this->dropPadMode = dropPadMode; + this->expertTokensCountOrCumsumFlag = expertTokensCountOrCumsumFlag; + this->expertTokensBeforeCapacityFlag = expertTokensBeforeCapacityFlag; + if (dropPadMode == 1) { + // Do not output expertTokensCountOrCumsum in drop-pad mode + expertTokensCountOrCumsumFlag = 0; + } else { + // Do not output expertTokensBeforeCapacity in dropless mode + expertTokensBeforeCapacityFlag = false; + } + moeInitRoutingTilingData.cols = cols; + moeInitRoutingTilingData.n = m; + moeInitRoutingTilingData.k = topK; + moeInitRoutingTilingData.expertCapacity = expertCapacity; + moeInitRoutingTilingData.expertNum = expertNum; + moeInitRoutingTilingData.dropPadMode = dropPadMode; + moeInitRoutingTilingData.expertTokensCountOrCumsumFlag = expertTokensCountOrCumsumFlag; + moeInitRoutingTilingData.expertTokensBeforeCapacityFlag = expertTokensBeforeCapacityFlag; + totalLength = moeInitRoutingTilingData.n * moeInitRoutingTilingData.k; + inuptXDtypeSize_ = inuptXDtypeSize; + return true; +} + +bool InnerMoeInitRoutingV2TilingBase::GetPlatformInfo(int64_t aivCoreNum, int64_t ubSizePlatForm) { + aivNum = aivCoreNum; + aicoreParams_.blockDim = aivCoreNum; + aicoreParams_.ubSize = ubSizePlatForm; + moeInitRoutingTilingData.coreNum = aivCoreNum; + return true; +} + + +bool InnerMoeInitRoutingV2TilingBase::GetWorkspaceSize() { + // Calculate workspace size + size_t sortWorkspaceSize = totalLength * sizeof(float) * NUM_TWO * NUM_THREE; // Space needed for sorting + size_t scatterWorkspaceSize = totalLength * sizeof(int32_t) * NUM_TWO; + size_t expertTokenFlagSize = aivNum * 2 * sizeof(int32_t); + workspaceSize_ = sortWorkspaceSize + scatterWorkspaceSize + expertTokenFlagSize + SIZE_16 * LENGTH_1024 * LENGTH_1024; + return true; +} + +void InnerMoeInitRoutingV2TilingBase::Tiling4VBSOneCoreCompute(InnerMoeV2VBSComputeTilingData* tilingData) { + tilingData->needCoreNum = 1; + tilingData->perCoreElements = totalLength; + tilingData->perCoreLoops = 1; + tilingData->perCorePerLoopElements = tilingData->perCoreElements; + tilingData->perCoreLastLoopElements = tilingData->perCoreElements; + tilingData->lastCoreElements = tilingData->perCoreElements; + tilingData->lastCoreLoops = 1; + tilingData->lastCorePerLoopElements = tilingData->perCoreElements; + tilingData->lastCoreLastLoopElements = tilingData->perCoreElements; +} + +void InnerMoeInitRoutingV2TilingBase::Tiling4VBSMultiCoreCompute(InnerMoeV2VBSComputeTilingData* tilingData) { + //Tiling4VBSMultiCoreCompute + int64_t needCoreNum = CeilDiv(totalLength, sortLoopMaxElement); // Round up + needCoreNum = static_cast(std::pow(4, CeilLog4(needCoreNum))); + needCoreNum = std::min(needCoreNum, aivNum); // Cannot exceed physical core count + if (needCoreNum > 0) { + int64_t perCoreElements = totalLength / needCoreNum; // Elements handled per core + int64_t alineFloorPerCoreElements = perCoreElements - perCoreElements % SORT32_ALIGN_ELEMENT; + int64_t lastCoreElement = totalLength - (needCoreNum - 1) * alineFloorPerCoreElements; + int64_t alineCeilPerCoreElements = perCoreElements + SORT32_ALIGN_ELEMENT - perCoreElements % SORT32_ALIGN_ELEMENT; + if (lastCoreElement > alineCeilPerCoreElements) { + perCoreElements = alineCeilPerCoreElements; + needCoreNum = CeilDiv(totalLength, perCoreElements); + } else { + perCoreElements = alineFloorPerCoreElements; + } + tilingData->needCoreNum = needCoreNum; + do { + tilingData->perCoreElements = perCoreElements; + tilingData->perCoreLoops = CeilDiv(tilingData->perCoreElements, sortLoopMaxElement); // Loops handled per core + tilingData->perCorePerLoopElements = std::min(tilingData->perCoreElements, sortLoopMaxElement); + tilingData->perCoreLastLoopElements = tilingData->perCoreElements - (tilingData->perCoreLoops - 1) * tilingData->perCorePerLoopElements; + tilingData->lastCoreElements = totalLength - (tilingData->needCoreNum - 1) * tilingData->perCoreElements; + tilingData->lastCoreLoops = tilingData->perCoreLoops; + int64_t tmp = CeilDiv(tilingData->lastCoreElements, tilingData->lastCoreLoops); + int64_t lastCorePerLoopElements = + CeilDiv(CeilDiv(tilingData->lastCoreElements, tilingData->lastCoreLoops), SORT32_ALIGN_ELEMENT) * + SORT32_ALIGN_ELEMENT; + tilingData->lastCorePerLoopElements = lastCorePerLoopElements; + tilingData->lastCoreLastLoopElements = tilingData-> lastCoreElements - (tilingData->lastCoreLoops - 1) * tilingData->lastCorePerLoopElements; + perCoreElements -= SORT32_ALIGN_ELEMENT; + } while (tilingData->lastCoreLastLoopElements <= 0 && perCoreElements > 0); + } +} + + +void InnerMoeInitRoutingV2TilingBase::Tiling4VBSCompute() { + auto tilingData = &moeInitRoutingTilingData.vbsComputeParamsOp; + tilingData->oneLoopMaxElements = sortLoopMaxElement; + if (totalLength <= sortLoopMaxElement) { // Only one core is used + Tiling4VBSOneCoreCompute(tilingData); + return; + } + Tiling4VBSMultiCoreCompute(tilingData); +} + +void InnerMoeInitRoutingV2TilingBase::Tiling4VMSMiddleCompute() { + auto vbsComputeTilingData = &moeInitRoutingTilingData.vbsComputeParamsOp; + auto tilingData = &moeInitRoutingTilingData.vmsMiddleComputeParamsOp; + if (vbsComputeTilingData->needCoreNum <= MRG_LIST_NUM) { // No intermediate merge if queue count fits one VMS + tilingData->needCoreNum = 0; // Required core count + } else { + int64_t needCoreNum = CeilDiv(vbsComputeTilingData->needCoreNum, MRG_LIST_NUM); + tilingData->needCoreNum = needCoreNum; // Required core count + } +} + +void InnerMoeInitRoutingV2TilingBase::Tiling4SortOutCompute() { + auto tilingData = &moeInitRoutingTilingData.sortOutComputeParamsOp; + tilingData->oneLoopMaxElements = mrgSortListMaxElement; +} + + +void InnerMoeInitRoutingV2TilingBase::Tiling4SrcToDstCompute() { + auto tilingData = &moeInitRoutingTilingData.srcToDstComputeParamsOp; + + int64_t perLoopMaxRows = (aicoreParams_.ubSize - ASSIST_NUM * sizeof(float) - aivNum * SORT32_ALIGN_ELEMENT) / + (SORT32_ALIGN_ELEMENT * NUM_TWO) / NUM_TWO; + int64_t perCoreRows = CeilDiv(totalLength, aivNum); + if (perCoreRows <= 0) { + tilingData->needCoreNum = 0; + return; + } + + int64_t needCoreNum = CeilDiv(totalLength, perCoreRows); + tilingData->needCoreNum = needCoreNum; + int64_t lastCoreNum = totalLength - perCoreRows * (tilingData->needCoreNum - 1); + tilingData->perCoreRows = perCoreRows; + if (perLoopMaxRows >= tilingData->perCoreRows) { // One loop completes + tilingData->perCorePerLoopRows = tilingData->perCoreRows; + tilingData->perCoreLastLoopRows = tilingData->perCoreRows; + } else { + tilingData->perCorePerLoopRows = perLoopMaxRows; + tilingData->perCoreLastLoopRows = tilingData->perCoreRows - (CeilDiv(tilingData->perCoreRows, perLoopMaxRows) - 1) * perLoopMaxRows; + } + tilingData->lastCoreRows = lastCoreNum; + if (perLoopMaxRows >= tilingData->lastCoreRows) { + tilingData->lastCorePerLoopRows = tilingData->lastCoreRows; + tilingData->lastCoreLastLoopRows = tilingData->lastCoreRows; + } else { + tilingData->lastCorePerLoopRows = perLoopMaxRows; + tilingData->lastCoreLastLoopRows = tilingData->lastCoreRows - (CeilDiv(tilingData->lastCoreRows, perLoopMaxRows) - 1) * perLoopMaxRows; + } +} + + +void InnerMoeInitRoutingV2TilingBase::Tiling4SrcToDstCapacityCompute() { + auto tilingData = &moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp; + int64_t perCoreRows = CeilDiv(totalLength, aivNum); + + if (perCoreRows <= 0) { + tilingData->needCoreNum = 0; + return; + } + + int64_t needCoreNum = CeilDiv(totalLength, perCoreRows); + tilingData->needCoreNum = needCoreNum; + int64_t cols = moeInitRoutingTilingData.cols; + tilingData->perCoreRows = perCoreRows; + int64_t lastCoreRows = totalLength - perCoreRows * (needCoreNum - 1); + tilingData->lastCoreRows = lastCoreRows; + + + int64_t rowSize = + (perCoreRows * sizeof(int32_t) * 2 + ONE_BLOCK_BYTE + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; + int64_t colSize = (cols * inuptXDtypeSize_ + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; + + if (rowSize + colSize < static_cast(aicoreParams_.ubSize)) { + tilingData->perCorePerLoopRows = perCoreRows; + tilingData->perCoreLastLoopRows = perCoreRows; + tilingData->lastCorePerLoopRows = lastCoreRows; + tilingData->lastCoreLastLoopRows = lastCoreRows; + tilingData->perCoreLoops = 1; + tilingData->lastCoreLoops = 1; + tilingData->perLoopCols = cols; + tilingData->lastLoopCols = cols; + tilingData->colLoops = 1; + + } else { + int64_t baseMaxCols = MAX_COLS_ONE_LOOP; + int64_t baseMaxColsSize = (baseMaxCols * inuptXDtypeSize_ + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; + int64_t basePerLoopMaxRows = (static_cast(aicoreParams_.ubSize) - baseMaxColsSize - ONE_BLOCK_BYTE) / + static_cast(sizeof(int32_t)) / NUM_TWO / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; + if (cols < MAX_COLS_ONE_LOOP) { + basePerLoopMaxRows = (static_cast(aicoreParams_.ubSize) - colSize - ONE_BLOCK_BYTE) / + static_cast(sizeof(int32_t)) / NUM_TWO / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; + } else if (perCoreRows < basePerLoopMaxRows) { + baseMaxCols = + (static_cast(aicoreParams_.ubSize) - rowSize) / inuptXDtypeSize_ / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; + } + tilingData->perLoopCols = (std::min(baseMaxCols, cols)); + tilingData->lastLoopCols = (GetPerOrLastValue(cols, baseMaxCols)); + tilingData->colLoops = ((cols + baseMaxCols - 1) / baseMaxCols); + tilingData->perCorePerLoopRows = (std::min(perCoreRows, basePerLoopMaxRows)); + tilingData->perCoreLastLoopRows = (GetPerOrLastValue(perCoreRows, basePerLoopMaxRows)); + tilingData->perCoreLoops = ((perCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows); + tilingData->lastCorePerLoopRows = (std::min(lastCoreRows, basePerLoopMaxRows)); + tilingData->lastCoreLastLoopRows = (GetPerOrLastValue(lastCoreRows, basePerLoopMaxRows)); + tilingData->lastCoreLoops = ((lastCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows); + } +} + +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_common.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_common.h new file mode 100644 index 000000000..e362f50a3 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_common.h @@ -0,0 +1,96 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_v2_common.h + * \brief + */ +#ifndef INNER_MOE_V2_COMMON_H +#define INNER_MOE_V2_COMMON_H + +#include "kernel_operator.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +constexpr int64_t SPLIT_N = 0; +constexpr int64_t SPLIT_K = 1; +constexpr float MIN_FP32 = -3.4e38; +constexpr int64_t ONE_REPEAT_SORT_NUM = 32; +constexpr int64_t BLOCK_BYTES = 32; +constexpr int64_t INT32_ONE_BLOCK_NUM = 8; + +constexpr int64_t ASSIST_NUM = 256; +constexpr int64_t ASSIST_INDEX_NUM = 32; + +constexpr int64_t MERGE_LIST_TWO = 2; +constexpr int64_t MERGE_LIST_THREE = 3; +constexpr int64_t MERGE_LIST_FOUR = 4; + +constexpr int64_t MERGE_LIST_IDX_TWO = 2; +constexpr int64_t MERGE_LIST_IDX_THREE = 3; + +constexpr int64_t MAX_EXPERT_NUM = 5120; +constexpr int64_t DROPLESS_MODE = 0; +constexpr int64_t DROP_PAD_MODE = 1; +constexpr int64_t EXERPT_TOKENS_COUNT = 2; +constexpr int64_t EXERPT_TOKENS_CUMSUM = 1; +constexpr int64_t EXERPT_TOKENS_NONE = 0; +constexpr int64_t EXERPT_TOKENS_BEFORE_CAPACITY = 1; +constexpr int64_t ALIGN_512 = 512; +constexpr int64_t ALIGN_128 = 128; + +const __gm__ int32_t assist[256] = { + 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, + 4, 0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, + 8, 0, 0, 0, 0, 0, 0, 0, 9, 0, 0, 0, 0, 0, 0, 0, 10, 0, 0, 0, 0, 0, 0, 0, 11, 0, 0, 0, 0, 0, 0, 0, + 12, 0, 0, 0, 0, 0, 0, 0, 13, 0, 0, 0, 0, 0, 0, 0, 14, 0, 0, 0, 0, 0, 0, 0, 15, 0, 0, 0, 0, 0, 0, 0, + 16, 0, 0, 0, 0, 0, 0, 0, 17, 0, 0, 0, 0, 0, 0, 0, 18, 0, 0, 0, 0, 0, 0, 0, 19, 0, 0, 0, 0, 0, 0, 0, + 20, 0, 0, 0, 0, 0, 0, 0, 21, 0, 0, 0, 0, 0, 0, 0, 22, 0, 0, 0, 0, 0, 0, 0, 23, 0, 0, 0, 0, 0, 0, 0, + 24, 0, 0, 0, 0, 0, 0, 0, 25, 0, 0, 0, 0, 0, 0, 0, 26, 0, 0, 0, 0, 0, 0, 0, 27, 0, 0, 0, 0, 0, 0, 0, + 28, 0, 0, 0, 0, 0, 0, 0, 29, 0, 0, 0, 0, 0, 0, 0, 30, 0, 0, 0, 0, 0, 0, 0, 31, 0, 0, 0, 0, 0, 0, 0}; + +__aicore__ inline int64_t Ceil(int64_t a, int64_t b) { + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes) { + if (bytes == 0) { + return 0; + } + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes; +} + +__aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes) { + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES; +} + +template +__aicore__ inline T Min(T a, T b) { + return a > b ? b : a; +} + +template +__aicore__ inline T Max(T a, T b) { + return a < b ? b : a; +} + +template +__aicore__ inline void SetWaitFlag(HardEvent evt) { + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(evt)); + SetFlag(eventId); + WaitFlag(eventId); +} + +} // namespace MoeInitRoutingQuantV2 +#endif // INNER_MOE_V2_COMMON_H \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_expert_token_out.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_expert_token_out.h new file mode 100644 index 000000000..ab14c2400 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_expert_token_out.h @@ -0,0 +1,310 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_v2_expert_token_out.h + * \brief + */ +#ifndef INNER_MOE_V2_EXPERT_TOKEN_OUT_H +#define INNER_MOE_V2_EXPERT_TOKEN_OUT_H + +#include "moe_v2_common.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +constexpr int64_t EXPERT_ID_VALUE_NUM = 2; + +class MoeV2ExpertTokenOut { + public: + __aicore__ inline MoeV2ExpertTokenOut(){}; + template + __aicore__ inline void Init(GM_ADDR expertTokensCountOrCumsum, GM_ADDR expertTokensBeforeCapacity, + GM_ADDR expandedRowIdx, GM_ADDR workspace, const TilingData* tilingData, TPipe* tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void CopyIn(int64_t progress); + __aicore__ inline void Compute(int64_t progress); + __aicore__ inline void SyncAll(); + __aicore__ inline void InitLocal(); + __aicore__ inline void GetExpertTokenCount(int32_t curExpertId); + __aicore__ inline void CopyOutTokenGm(); + __aicore__ inline void CopyOutExpertTokensCumsum(bool isTail); + __aicore__ inline void CopyOutExpertTokensCount(bool isTail); + + private: + TPipe* pipe; + TQue copyInQueue; + TQue expertTokenIdxCopyInQueue; + TQue expertTokenIdxCopyOutQueue; + + GlobalTensor expertTokensCountOrCumsumGm; + GlobalTensor expertTokensBeforeCapacityGm; + GlobalTensor expandedExpertIdxGm; + GlobalTensor expertIdxValueGm; + GlobalTensor expandedRowIdxGm; + + LocalTensor expertTokenIdxOutLocal; + + const InnerMoeV2GatherOutComputeTilingData* srcToDstTilingData; + + int64_t coreNum; + int64_t blockIdx; + int64_t totalLength; + int64_t currentLoopRows; + int64_t coreRows; + int64_t perLoopRows; + int64_t lastLoopRows; + int64_t expertNum; + int64_t expertNumUbAlign; + int64_t dropPadMode = 0; + int64_t expertTokensCountOrCumsumFlag = 0; + int64_t expertTokensBeforeCapacityFlag = 0; + + int64_t tokenCount = 0; + int64_t expertIdx = 0; + int32_t lastExpertId = -1; + int32_t firstExpertId = -1; + + int32_t expertTokenValue = 0; +}; + +__aicore__ inline void MoeV2ExpertTokenOut::InitLocal() { + LocalTensor tokenIdxLocal = expertTokenIdxCopyOutQueue.AllocTensor(); + Duplicate(tokenIdxLocal, 0, this->expertNumUbAlign); + expertTokenIdxCopyOutQueue.EnQue(tokenIdxLocal); + + // expandedRowIdx initialized to -1, which is used in the src_to_dst_with_capacity step. + // use this step SyncAll to synchronize every core data + if (this->dropPadMode == 0) { + return; + } + LocalTensor outLocal = copyInQueue.AllocTensor(); + int64_t loops = (coreRows + perLoopRows - 1) / perLoopRows; + Duplicate(outLocal, -1, perLoopRows); + SetWaitFlag(HardEvent::V_MTE3); + for (int64_t loop = 0; loop < loops; loop++) { + int64_t copyLength = perLoopRows; + if (loop == loops - 1) { + copyLength = lastLoopRows; + } + DataCopyExtParams copyParams{static_cast(1), static_cast(copyLength * sizeof(int32_t)), 0, 0, + 0}; + DataCopyPad(expandedRowIdxGm[this->blockIdx * this->srcToDstTilingData->perCoreRows + loop * perLoopRows], outLocal, + copyParams); + } + SetWaitFlag(HardEvent::MTE3_MTE2); + copyInQueue.FreeTensor(outLocal); +} + +__aicore__ inline void MoeV2ExpertTokenOut::CopyIn(int64_t progress) { + LocalTensor inLocal = copyInQueue.AllocTensor(); + DataCopy(inLocal, expandedExpertIdxGm[progress * perLoopRows], Align(currentLoopRows, sizeof(int32_t))); + copyInQueue.EnQue(inLocal); +} + +__aicore__ inline void MoeV2ExpertTokenOut::GetExpertTokenCount(int32_t curExpertId) { + this->tokenCount++; + if (this->lastExpertId < curExpertId) { + this->expertTokenIdxOutLocal.SetValue(this->expertIdx, this->tokenCount - 1); + this->tokenCount = 1; + this->expertIdx += (curExpertId - this->lastExpertId); + while (curExpertId - this->firstExpertId + 1 > this->expertNumUbAlign) { + SetWaitFlag(HardEvent::S_MTE3); + CopyOutExpertTokensCumsum(false); + CopyOutExpertTokensCount(false); + SetWaitFlag(HardEvent::MTE3_V); + Duplicate(this->expertTokenIdxOutLocal, 0, this->expertNumUbAlign); + SetWaitFlag(HardEvent::V_S); + this->firstExpertId += this->expertNumUbAlign; + this->expertIdx = curExpertId - this->firstExpertId; + } + this->lastExpertId = curExpertId; + } +} + +__aicore__ inline void MoeV2ExpertTokenOut::Compute(int64_t progress) { + LocalTensor inLocal = copyInQueue.DeQue(); + SetWaitFlag(HardEvent::MTE2_S); + if (this->lastExpertId == -1) { + this->lastExpertId = inLocal.GetValue(0); + this->firstExpertId = this->lastExpertId; + } + for (int64_t i = 0; i < currentLoopRows; i++) { + int32_t expertId = inLocal.GetValue(i); + GetExpertTokenCount(expertId); + } + this->expertTokenIdxOutLocal.SetValue(this->expertIdx, this->tokenCount); + copyInQueue.FreeTensor(inLocal); +} + +__aicore__ inline void MoeV2ExpertTokenOut::CopyOutExpertTokensCumsum(bool isTail) { + if (this->dropPadMode != DROPLESS_MODE || expertTokensCountOrCumsumFlag != EXERPT_TOKENS_CUMSUM) { + return; + } +#ifdef __CCE_KT_TEST__ + // CPU twin debugging cannot use multi-core sync, so index may contain uninitialized dirty data; handle specially + if (this->firstExpertId > expertTokensCountOrCumsumGm.GetSize()) { + return; + } +#endif + int64_t copyLength = isTail ? this->lastExpertId - this->firstExpertId + 1 : this->expertNumUbAlign; + int64_t end = this->expertNum - this->firstExpertId; + for (int64_t i = 0; i < copyLength; i++) { + this->expertTokenValue += this->expertTokenIdxOutLocal.GetValue(i); + this->expertTokenIdxOutLocal.SetValue(i, this->expertTokenValue); + } + // if the remaining UB is sufficient, use the UB space to copy + // otherwise, copy the calculated data first, and then copy the last tokenValue to remaining expert position + if (isTail && end <= this->expertNumUbAlign) { + int64_t startAlign = Min(Align(copyLength, sizeof(int32_t)), end); + for (int64_t i = copyLength; i < startAlign; i++) { + this->expertTokenIdxOutLocal.SetValue(i, this->expertTokenValue); + } + if (startAlign < end) { + Duplicate(this->expertTokenIdxOutLocal[startAlign], this->expertTokenValue, end - startAlign); + } + copyLength = end; + SetWaitFlag(HardEvent::V_MTE3); + } + DataCopyExtParams copyParams{static_cast(1), static_cast(copyLength * sizeof(int32_t)), 0, 0, 0}; + SetAtomicAdd(); +#ifndef __CCE_KT_TEST__ + DataCopyPad(expertTokensCountOrCumsumGm[this->firstExpertId], this->expertTokenIdxOutLocal, copyParams); +#endif + SetAtomicNone(); + if (isTail && end > this->expertNumUbAlign) { + int64_t remainderLength = end - copyLength; + SetWaitFlag(HardEvent::MTE3_V); + Duplicate(this->expertTokenIdxOutLocal, this->expertTokenValue, this->expertNumUbAlign); + SetWaitFlag(HardEvent::V_MTE3); + int64_t loopTimes = remainderLength / this->expertNumUbAlign + 1; + for (int64_t i = 0; i < loopTimes; i++) { + copyLength = i == loopTimes - 1 ? remainderLength - this->expertNumUbAlign * i : this->expertNumUbAlign; + DataCopyExtParams params{static_cast(1), static_cast(copyLength * sizeof(int32_t)), 0, 0, 0}; + SetAtomicAdd(); + DataCopyPad(expertTokensCountOrCumsumGm[this->lastExpertId + 1 + this->expertNumUbAlign * i], + this->expertTokenIdxOutLocal, params); + SetAtomicNone(); + } + } +} + +__aicore__ inline void MoeV2ExpertTokenOut::CopyOutExpertTokensCount(bool isTail) { + int64_t copyLength = isTail ? this->lastExpertId - this->firstExpertId + 1 : this->expertNumUbAlign; + DataCopyExtParams copyParams{static_cast(1), static_cast(copyLength * sizeof(int32_t)), 0, 0, 0}; +#ifdef __CCE_KT_TEST__ + // CPU twin debugging skips output copies + return; +#endif + SetAtomicAdd(); + if (this->dropPadMode == DROP_PAD_MODE && expertTokensBeforeCapacityFlag > EXERPT_TOKENS_NONE) { + DataCopyPad(expertTokensBeforeCapacityGm[this->firstExpertId], this->expertTokenIdxOutLocal, copyParams); + } + if (this->dropPadMode == DROPLESS_MODE && expertTokensCountOrCumsumFlag == EXERPT_TOKENS_COUNT) { + DataCopyPad(expertTokensCountOrCumsumGm[this->firstExpertId], this->expertTokenIdxOutLocal, copyParams); + } + SetAtomicNone(); +} + +__aicore__ inline void MoeV2ExpertTokenOut::CopyOutTokenGm() { + if (this->dropPadMode == DROPLESS_MODE) { + SetWaitFlag(HardEvent::S_MTE3); + CopyOutExpertTokensCumsum(true); + CopyOutExpertTokensCount(true); + return; + } + this->expertTokenIdxOutLocal.SetValue(this->expertNumUbAlign, this->lastExpertId); + this->expertTokenIdxOutLocal.SetValue(this->expertNumUbAlign + 1, this->tokenCount); + DataCopyExtParams copyParams{static_cast(1), static_cast(EXPERT_ID_VALUE_NUM * sizeof(int32_t)), + 0, 0, 0}; + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(expertIdxValueGm[this->blockIdx * EXPERT_ID_VALUE_NUM], + this->expertTokenIdxOutLocal[this->expertNumUbAlign], copyParams); + CopyOutExpertTokensCount(true); +} + +__aicore__ inline void MoeV2ExpertTokenOut::SyncAll() { + if (coreNum == 1) { + return; + } +#ifndef __CCE_KT_TEST__ + AscendC::SyncAll(); +#endif +} + +template +__aicore__ inline void MoeV2ExpertTokenOut::Init(GM_ADDR expertTokensCountOrCumsum, GM_ADDR expertTokensBeforeCapacity, + GM_ADDR expandedRowIdx, GM_ADDR workspace, + const TilingData* tilingData, TPipe* tPipe) { + int64_t blockNum = GetBlockNum(); + this->pipe = tPipe; + //this->blockIdx = GetBlockIdx(); + this->blockIdx = get_block_idx() + get_subblockid() * get_block_num(); + this->coreNum = tilingData->coreNum; + this->totalLength = tilingData->n * tilingData->k; + this->srcToDstTilingData = &(tilingData->srcToDstComputeParamsOp); + this->expertNum = tilingData->expertNum; + this->dropPadMode = tilingData->dropPadMode; + this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag; + this->expertTokensBeforeCapacityFlag = tilingData->expertTokensBeforeCapacityFlag; + + if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) { + this->coreRows = this->srcToDstTilingData->lastCoreRows; + this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows; + this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows; + } else { + this->coreRows = this->srcToDstTilingData->perCoreRows; + this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows; + this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows; + } + + expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t*)expandedRowIdx, Align(this->totalLength, sizeof(int32_t))); + if (this->dropPadMode == DROPLESS_MODE && this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) { + expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensCountOrCumsum, this->expertNum); + } + if (this->dropPadMode == DROP_PAD_MODE && this->expertTokensBeforeCapacityFlag == EXERPT_TOKENS_BEFORE_CAPACITY) { + expertTokensBeforeCapacityGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensBeforeCapacity, this->expertNum); + } + + expandedExpertIdxGm.SetGlobalBuffer( + (__gm__ int32_t*)workspace + this->blockIdx * this->srcToDstTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + expertIdxValueGm.SetGlobalBuffer((__gm__ int32_t*)workspace + Align(this->totalLength, sizeof(int32_t)) * 2, + this->coreNum * 2); + + this->expertNumUbAlign = Min(Align(this->expertNum, sizeof(int32_t)), MAX_EXPERT_NUM); + pipe->InitBuffer(copyInQueue, 1, this->perLoopRows * BLOCK_BYTES); + pipe->InitBuffer(expertTokenIdxCopyInQueue, 1, this->expertNumUbAlign * sizeof(int32_t)); + pipe->InitBuffer(expertTokenIdxCopyOutQueue, 1, (this->expertNumUbAlign + EXPERT_ID_VALUE_NUM) * sizeof(int32_t)); +} + +__aicore__ inline void MoeV2ExpertTokenOut::Process() { + if (this->blockIdx < this->srcToDstTilingData->needCoreNum) { + int64_t loops = (coreRows + perLoopRows - 1) / perLoopRows; + currentLoopRows = perLoopRows; + InitLocal(); + this->expertTokenIdxOutLocal = expertTokenIdxCopyOutQueue.DeQue(); + for (int64_t loop = 0; loop < loops - 1; loop++) { + CopyIn(loop); + Compute(loop); + } + currentLoopRows = lastLoopRows; + CopyIn(loops - 1); + Compute(loops - 1); + CopyOutTokenGm(); + expertTokenIdxCopyOutQueue.FreeTensor(this->expertTokenIdxOutLocal); + } + this->SyncAll(); +} + +} // namespace MoeInitRoutingQuantV2 +#endif // INNER_MOE_V2_EXPERT_TOKEN_OUT_H diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_fullload_dynamic_quant.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_fullload_dynamic_quant.h new file mode 100644 index 000000000..60ab43f4c --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_fullload_dynamic_quant.h @@ -0,0 +1,398 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file moe_v2_fullload_dynamic_quant.h + * \brief + */ +#ifndef MOE_V2_FULL_LOAD_DYNAMIC_QUANT_H +#define MOE_V2_FULL_LOAD_DYNAMIC_QUANT_H + +#include "moe_v2_mrgsort.h" +#include "moe_v2_sort_base.h" +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +template +class MoeV2FullLoadDynamicQuant : public MoeV2SortBase { + public: + __aicore__ inline MoeV2FullLoadDynamicQuant(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX, GM_ADDR expandedRowIdx, + GM_ADDR expertTokensCountOrCumsum, GM_ADDR quantSmooth, GM_ADDR dynamicQuantScale, + GM_ADDR workspace, const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void CopyIn(); + __aicore__ inline void SortCompute(); + __aicore__ inline void CopyOutIdx(); + __aicore__ inline void CopyOutEmpty(); + __aicore__ inline void CopyOutXQuant1H(); + __aicore__ inline void ComputeExpertTokenCountOrCumsum(); + __aicore__ inline void Compute(LocalTensor& smoothLocal); + + private: + int64_t sortNum_; + const InnerMoeV2GatherOutComputeTilingData* gatherOutTilingData_; + int64_t blockIdx_; + int64_t needCoreNum_; + int64_t coreRows_; + int64_t perCoreRows_; + int64_t k_; + int64_t n_; + int64_t cols_; + int64_t cols_scale_; + int64_t activateRows_; + int64_t expertNum; + int64_t expertCapacity; + int64_t smoothType; + int64_t colsAlign; + + TQue xCopyInQueue_; + TQue expandedRowIdxCopyOutQueue_; + TQue expandedExpertIdxCopyOutQueue_; + TQue expandDstToSrcRowQueue_; + TQue expertTokensCopyOutQueue_; + TQue smoothInQueue; + TQue calcQueue; + TQue inputXOutQueue; + + GlobalTensor xGm_; + GlobalTensor expertIdxGm_; + GlobalTensor quantSmoothGm; + + GlobalTensor expandedXGm_; + GlobalTensor expandedRowIdxGm_; + GlobalTensor expandedExpertIdxGm_; + GlobalTensor expertTokensCountOrCumsumGm; + GlobalTensor expertTokensBeforeCapacityGm; + + int64_t expertTokensCountOrCumsumFlag = 0; + int64_t expertTokensBeforeCapacityFlag = 0; + int64_t dropPadMode = 0; + + LocalTensor expandDstToSrcRowLocal; + LocalTensor expandedExpertIdxLocal; +}; + +template +__aicore__ inline void MoeV2FullLoadDynamicQuant::CopyIn() { + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + DataCopyExtParams dataCopyParams{static_cast(1), static_cast(this->totalLength * sizeof(int32_t)), + 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal[0], expertIdxGm_, dataCopyParams, dataCopyPadParams); + ArithProgression(inLocal[this->sortNum_], 0, 1, this->totalLength); + sortDataCopyInQueue.EnQue(inLocal); +} + +template +__aicore__ inline void MoeV2FullLoadDynamicQuant::SortCompute() { + LocalTensor inLocal = sortDataCopyInQueue.DeQue(); + LocalTensor expertIdxLocal = inLocal[0]; + LocalTensor expertIdxLocalFp32 = expertIdxLocal.ReinterpretCast(); + Cast(expertIdxLocalFp32, expertIdxLocal, RoundMode::CAST_ROUND, this->totalLength); + AscendC::PipeBarrier(); + Muls(expertIdxLocalFp32, expertIdxLocalFp32, (float)-1, this->totalLength); + AscendC::PipeBarrier(); + int64_t duplicateNum = this->totalLength % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + int duplicateIndex = this->totalLength - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expertIdxLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + AscendC::PipeBarrier(); + } + LocalTensor concatLocal; + LocalTensor tempTensor = tempBuffer.Get(GetSortLen(this->sortNum_)); + Concat(concatLocal, expertIdxLocalFp32, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + LocalTensor rowIdxLocal = inLocal[this->sortNum_].template ReinterpretCast(); + LocalTensor sortedLocal = sortedBuffer.Get(GetSortLen(this->sortNum_)); + Sort(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + LocalTensor expandedExpertIdxLocal = expandedExpertIdxCopyOutQueue_.AllocTensor(); + expandDstToSrcRowLocal = expandDstToSrcRowQueue_.AllocTensor(); + LocalTensor expandDstToSrcRowLocalFp32 = expandDstToSrcRowLocal.ReinterpretCast(); + Extract(expandedExpertIdxLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum_ / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + Cast(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocal.ReinterpretCast(), RoundMode::CAST_ROUND, + this->totalLength); + AscendC::PipeBarrier(); + Muls(expandedExpertIdxLocal, expandedExpertIdxLocal, (float)-1, this->totalLength); + AscendC::PipeBarrier(); + LocalTensor expandedExpertIdxLocalInt32; + expandedExpertIdxLocalInt32 = expandedExpertIdxLocal.ReinterpretCast(); + Cast(expandedExpertIdxLocalInt32, expandedExpertIdxLocal, RoundMode::CAST_ROUND, this->totalLength); + AscendC::PipeBarrier(); + expandedExpertIdxCopyOutQueue_.EnQue(expandedExpertIdxLocalInt32); + + LocalTensor expandedRowIdx = expandedRowIdxCopyOutQueue_.AllocTensor(); + LocalTensor expandedRowIdxU32 = expandedRowIdx.ReinterpretCast(); + Muls(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocalFp32, (float)-1, this->totalLength); + AscendC::PipeBarrier(); + ArithProgression(inLocal[this->sortNum_], 0, 1, this->totalLength); + AscendC::PipeBarrier(); + if (duplicateNum > 0) { + int duplicateIndex = this->totalLength - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expandDstToSrcRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + AscendC::PipeBarrier(); + } + Concat(concatLocal, expandDstToSrcRowLocalFp32, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + Sort(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + Extract(tempTensor, expandedRowIdxU32, sortedLocal, this->sortNum_ / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + expandedRowIdxCopyOutQueue_.EnQue(expandedRowIdx); + sortDataCopyInQueue.FreeTensor(inLocal); +} + +template +__aicore__ inline void MoeV2FullLoadDynamicQuant::CopyOutIdx() { + LocalTensor expandedRowIdx = expandedRowIdxCopyOutQueue_.DeQue(); + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = this->totalLength * sizeof(int32_t); + DataCopyPad(expandedRowIdxGm_, expandedRowIdx, intriParams); + expandedRowIdxCopyOutQueue_.EnQue(expandedRowIdx); +} + +template +__aicore__ inline void MoeV2FullLoadDynamicQuant::ComputeExpertTokenCountOrCumsum() { + expandedExpertIdxLocal = expandedExpertIdxCopyOutQueue_.DeQue(); + LocalTensor expertTokensCount = expertTokensCopyOutQueue_.AllocTensor(); + + int64_t expertNumAlign = Align(this->expertNum, sizeof(int32_t)); + Duplicate(expertTokensCount, 0, expertNumAlign); + SetWaitFlag(HardEvent::V_S); + + int32_t lastExpertId = expandedExpertIdxLocal.GetValue(0); + int64_t tokenCount = 0; + int64_t lastExpertCount = 0; + for (int64_t i = 0; i < this->totalLength; i++) { + int32_t curExpertId = expandedExpertIdxLocal.GetValue(i); + tokenCount++; + while (lastExpertId < curExpertId) { + expertTokensCount.SetValue(lastExpertId, tokenCount - 1); + if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_COUNT) { + tokenCount = 1; + } + lastExpertId++; + } + } +#ifndef __CCE_KT_TEST__ + expertTokensCount.SetValue(lastExpertId, tokenCount); + if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_CUMSUM) { + lastExpertId++; + while (lastExpertId < this->expertNum) { + expertTokensCount.SetValue(lastExpertId, tokenCount); + lastExpertId++; + } + } + DataCopyExtParams copyParams{static_cast(1), static_cast(this->expertNum * sizeof(int32_t)), 0, 0, + 0}; + if (this->expertTokensCountOrCumsumFlag > 0) { + DataCopyPad(expertTokensCountOrCumsumGm, expertTokensCount, copyParams); + } + expertTokensCopyOutQueue_.FreeTensor(expertTokensCount); +#endif +} + +template +__aicore__ inline void MoeV2FullLoadDynamicQuant::CopyOutEmpty() { + expandedExpertIdxLocal = expandedExpertIdxCopyOutQueue_.DeQue(); +} + +template +__aicore__ inline void MoeV2FullLoadDynamicQuant::Compute(LocalTensor& smoothLocal) { + LocalTensor inLocal = xCopyInQueue_.DeQue(); + + LocalTensor tempLocal = calcQueue.AllocTensor(); + LocalTensor outLocal = inputXOutQueue.AllocTensor(); + LocalTensor dynamicQuantLocal = outLocal[this->cols_].template ReinterpretCast(); + + if constexpr (!IsSameType::value) { + Cast(inLocal, inLocal.ReinterpretCast()[colsAlign], RoundMode::CAST_NONE, this->cols_); + AscendC::PipeBarrier(); + } + + if (smoothType != 0) { + Mul(inLocal, inLocal, smoothLocal, this->cols_); + AscendC::PipeBarrier(); + } + + Abs(tempLocal, inLocal, this->cols_); + AscendC::PipeBarrier(); + + ReduceMax(dynamicQuantLocal, tempLocal, tempLocal, this->cols_); + AscendC::PipeBarrier(); + + float maxValue = dynamicQuantLocal.GetValue(0) / 127.0f; + + Duplicate(dynamicQuantLocal, maxValue, 8); + Duplicate(tempLocal, maxValue, this->cols_); + AscendC::PipeBarrier(); + + Div(tempLocal, inLocal, tempLocal, this->cols_); + AscendC::PipeBarrier(); + + Cast(tempLocal.ReinterpretCast(), tempLocal, RoundMode::CAST_TRUNC, this->cols_); + AscendC::PipeBarrier(); + + Cast(outLocal, tempLocal.ReinterpretCast(), RoundMode::CAST_ROUND, this->cols_); + + calcQueue.FreeTensor(tempLocal); + inputXOutQueue.EnQue(outLocal); +} + +template +__aicore__ inline void MoeV2FullLoadDynamicQuant::CopyOutXQuant1H() { + expandDstToSrcRowQueue_.FreeTensor(expandDstToSrcRowLocal); + expandedExpertIdxCopyOutQueue_.FreeTensor(expandedExpertIdxLocal); + + LocalTensor expandedRowIdx = expandedRowIdxCopyOutQueue_.DeQue(); + int64_t curRowsStart = this->blockIdx_ * this->perCoreRows_; + int64_t curRowsEnd = curRowsStart + this->coreRows_ - 1; + int64_t startXRow = curRowsStart / this->k_; + int64_t endXRow = curRowsEnd / this->k_; + + DataCopyExtParams dataXCopyParams{1, static_cast(this->cols_ * sizeof(T)), 0, 0, 0}; + DataCopyExtParams smoothCopyParams{1, static_cast(this->cols_ * sizeof(float)), 0, 0, 0}; + DataCopyExtParams intriParams{1, static_cast((this->cols_ + BLOCK_BYTES) * sizeof(int8_t)), 0, 0, 0}; + + LocalTensor smoothLocal; + if (smoothType == 1) { + smoothLocal = smoothInQueue.AllocTensor(); + DataCopyPad(smoothLocal, quantSmoothGm, smoothCopyParams, {false, 0, 0, 0}); + smoothInQueue.EnQue(smoothLocal); + smoothLocal = smoothInQueue.DeQue(); + } + for (int64_t row = startXRow; row <= endXRow; row++) { + LocalTensor xLocal = xCopyInQueue_.AllocTensor(); + if constexpr (IsSameType::value) { + DataCopyPad(xLocal, xGm_[row * this->cols_], dataXCopyParams, {false, 0, 0, 0}); + } else { + DataCopyPad(xLocal[colsAlign], xGm_[row * this->cols_], dataXCopyParams, {false, 0, 0, 0}); + } + + xCopyInQueue_.EnQue(xLocal); + Compute(smoothLocal); + + LocalTensor outLocal = inputXOutQueue.DeQue(); + while (curRowsStart <= curRowsEnd && curRowsStart / this->k_ == row) { + int32_t outIndex = expandedRowIdx.GetValue(curRowsStart); + curRowsStart++; + if (outIndex == -1 || (this->dropPadMode == DROPLESS_MODE && outIndex >= this->activateRows_)) { + continue; + } + DataCopyPad(expandedXGm_[outIndex * this->cols_scale_], outLocal, intriParams); + } + + xCopyInQueue_.FreeTensor(xLocal); + inputXOutQueue.FreeTensor(outLocal); + } + expandedRowIdxCopyOutQueue_.FreeTensor(expandedRowIdx); +} + +template +__aicore__ inline void MoeV2FullLoadDynamicQuant::Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX, + GM_ADDR expandedRowIdx, GM_ADDR expertTokensCountOrCumsum, + GM_ADDR quantSmooth, GM_ADDR dynamicQuantScale, + GM_ADDR workspace, + const MoeInitRoutingQuantV2TilingData* tilingData, + TPipe* tPipe) { + this->gatherOutTilingData_ = &(tilingData->gatherOutComputeParamsOp); + //this->blockIdx_ = GetBlockIdx(); + this->blockIdx_ = get_block_idx() + get_subblockid() * get_block_num(); + this->k_ = tilingData->k; + this->n_ = tilingData->n; + this->cols_ = tilingData->cols; + this->cols_scale_ = this->cols_ + UB_ALIGN; + this->needCoreNum_ = this->gatherOutTilingData_->needCoreNum; + this->perCoreRows_ = this->gatherOutTilingData_->perCoreRows; + this->activateRows_ = this->gatherOutTilingData_->activateRows; + if (this->blockIdx_ == this->gatherOutTilingData_->needCoreNum - 1) { + this->coreRows_ = this->gatherOutTilingData_->lastCoreRows; + } else { + this->coreRows_ = this->gatherOutTilingData_->perCoreRows; + } + this->expertNum = tilingData->expertNum; + this->dropPadMode = tilingData->dropPadMode; + this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag; + + this->tileLength = Align(tilingData->vbsComputeParamsOp.lastCorePerLoopElements, sizeof(int32_t)); + this->sortNum_ = Ceil(this->tileLength, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + this->totalLength = tilingData->n * tilingData->k; + this->smoothType = tilingData->smoothType; + this->colsAlign = Align(this->cols_, sizeof(T)); + this->pipe = tPipe; + + xGm_.SetGlobalBuffer((__gm__ T*)x); + expertIdxGm_.SetGlobalBuffer((__gm__ int32_t*)expertIdx, this->tileLength); + + expandedXGm_.SetGlobalBuffer((__gm__ int8_t*)expandedX); + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t*)expandedRowIdx, this->tileLength); + if (this->expertTokensCountOrCumsumFlag > 0) { + // dropless + expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensCountOrCumsum, + Align(this->expertNum, sizeof(int32_t))); + } + quantSmoothGm.SetGlobalBuffer((__gm__ float*)quantSmooth); + + int64_t kvFactor = 2; + int64_t buffSize = this->sortNum_ * sizeof(int32_t); + + int64_t curRowsStart = this->blockIdx_ * this->perCoreRows_; + int64_t startXRow = curRowsStart / this->k_; + int64_t endXRow = (curRowsStart + this->coreRows_ - 1) / this->k_; + + pipe->InitBuffer(expandedRowIdxCopyOutQueue_, bufferNum, buffSize); + pipe->InitBuffer(expandedExpertIdxCopyOutQueue_, bufferNum, buffSize); + pipe->InitBuffer(expertTokensCopyOutQueue_, bufferNum, AlignBytes(this->expertNum, sizeof(int32_t))); + pipe->InitBuffer(expandDstToSrcRowQueue_, bufferNum, buffSize); + pipe->InitBuffer(sortDataCopyInQueue, bufferNum, buffSize * kvFactor); + pipe->InitBuffer(tempBuffer, buffSize * kvFactor); + pipe->InitBuffer(sortedBuffer, buffSize * kvFactor); + + if constexpr (IsSameType::value) { + pipe->InitBuffer(xCopyInQueue_, 1, AlignBytes(this->cols_, sizeof(float))); + } else { + pipe->InitBuffer(xCopyInQueue_, 1, 2 * AlignBytes(this->cols_, sizeof(T))); + } + pipe->InitBuffer(smoothInQueue, 1, AlignBytes(this->cols_, sizeof(float))); + pipe->InitBuffer(calcQueue, 1, AlignBytes(this->cols_, sizeof(float))); + pipe->InitBuffer(inputXOutQueue, 1, AlignBytes(this->cols_scale_, sizeof(int8_t))); +} + +template +__aicore__ inline void MoeV2FullLoadDynamicQuant::Process() { + if (this->blockIdx_ < this->needCoreNum_) { + CopyIn(); + SortCompute(); + if (this->blockIdx_ == 0) { + CopyOutIdx(); + } + if (this->blockIdx_ == this->needCoreNum_ - 1 && this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) { + ComputeExpertTokenCountOrCumsum(); + } else { + CopyOutEmpty(); + } + CopyOutXQuant1H(); + } +} +} // namespace MoeInitRoutingQuantV2 +#endif // MOE_V2_DYNAMIC_QUANT_FULL_LOAD_H \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_fullload_quant.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_fullload_quant.h new file mode 100644 index 000000000..8884889ac --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_fullload_quant.h @@ -0,0 +1,155 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file moe_v2_fullload_quant.h + * \brief + */ +#ifndef MOE_V2_FULL_LOAD_QUANT_H +#define MOE_V2_FULL_LOAD_QUANT_H + +#include "moe_v2_fullload_quant_base.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +template +class MoeV2FullLoadQuant : public MoeV2FullLoadQuantBase { + public: + __aicore__ inline MoeV2FullLoadQuant(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR offset, GM_ADDR expandedX, + GM_ADDR expandedRowIdx, GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace, + const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void Compute(int64_t xLocalLength); + __aicore__ inline void CopyOutX(); + + private: + TQue floatQueue; + TQue halfQueue; + TQue inputXCopyOutQueue; + + GlobalTensor xGm; + GlobalTensor scaleGm; + GlobalTensor offsetGm; + + float scale; + float offset; +}; + +template +__aicore__ inline void MoeV2FullLoadQuant::Compute(int64_t xLocalLength) { + LocalTensor inLocal = xCopyInQueue.DeQue(); + LocalTensor outLocal = inputXCopyOutQueue.AllocTensor(); + LocalTensor floatLocal = floatQueue.AllocTensor(); + LocalTensor halfLocal = halfQueue.AllocTensor(); + + uint32_t elements = Align(this->cols, sizeof(int8_t)) * xLocalLength; + if constexpr (IsSameType::value) { + Cast(floatLocal, inLocal, RoundMode::CAST_NONE, elements); + AscendC::PipeBarrier(); + Cast(halfLocal, floatLocal, RoundMode::CAST_NONE, elements); + AscendC::PipeBarrier(); + Muls(halfLocal, halfLocal, static_cast(this->scale), elements); + AscendC::PipeBarrier(); + Adds(halfLocal, halfLocal, static_cast(this->offset), elements); + AscendC::PipeBarrier(); + LocalTensor intLocal = floatLocal.ReinterpretCast(); + Cast(intLocal, halfLocal, RoundMode::CAST_RINT, elements); + AscendC::PipeBarrier(); + SetDeqScale((half)1.000000e+00f); + AscendC::PipeBarrier(); + Cast(halfLocal, intLocal, RoundMode::CAST_RINT, elements); + AscendC::PipeBarrier(); + Cast(outLocal, halfLocal, RoundMode::CAST_RINT, elements); + } else if constexpr (IsSameType::value) { + Cast(halfLocal, inLocal, RoundMode::CAST_NONE, elements); + AscendC::PipeBarrier(); + Muls(halfLocal, halfLocal, static_cast(this->scale), elements); + AscendC::PipeBarrier(); + Adds(halfLocal, halfLocal, static_cast(this->offset), elements); + AscendC::PipeBarrier(); + Cast(outLocal, halfLocal, RoundMode::CAST_RINT, elements); + } else { + Muls(inLocal, inLocal, static_cast(this->scale), elements); + AscendC::PipeBarrier(); + Adds(inLocal, inLocal, static_cast(this->offset), elements); + AscendC::PipeBarrier(); + Cast(outLocal, inLocal, RoundMode::CAST_RINT, elements); + } + inputXCopyOutQueue.EnQue(outLocal); + xCopyInQueue.FreeTensor(inLocal); + floatQueue.FreeTensor(floatLocal); + halfQueue.FreeTensor(halfLocal); +} + +template +__aicore__ inline void MoeV2FullLoadQuant::CopyOutX() { + LocalTensor xLocal = xCopyInQueue.AllocTensor(); + LocalTensor expandedRowIdx = expandedRowIdxCopyOutQueue.DeQue(); + int64_t inFactor = Align(this->cols, sizeof(int8_t)); + int64_t curRowsStart = this->blockIdx * this->perCoreRows; + int64_t startXRow = curRowsStart / this->k; + int64_t endXRow = (curRowsStart + this->coreRows - 1) / this->k; + + uint32_t dstStride = (inFactor * sizeof(T) - AlignBytes(this->cols, sizeof(T))) / BLOCK_BYTES; + DataCopyExtParams dataXCopyParams{static_cast(endXRow - startXRow + 1), + static_cast(this->cols * sizeof(T)), 0, dstStride, 0}; + DataCopyPadExtParams dataXCopyPadParams{false, 0, 0, 0}; + DataCopyPad(xLocal, xGm[startXRow * this->cols], dataXCopyParams, dataXCopyPadParams); + xCopyInQueue.EnQue(xLocal); + Compute(endXRow - startXRow + 1); + LocalTensor outLocal = inputXCopyOutQueue.DeQue(); + int64_t k = 0; + DataCopyExtParams intriParams{1, static_cast(this->cols * sizeof(int8_t)), 0, 0, 0}; + for (int64_t i = startXRow; i <= endXRow; i++) { + for (; k < this->perCoreRows && curRowsStart / this->k == i; curRowsStart++, k++) { + int32_t outIndex = expandedRowIdx.GetValue(curRowsStart); + if (outIndex < this->activateRows) { + DataCopyPad(expandedXGm[outIndex * this->cols], outLocal[(i - startXRow) * inFactor], intriParams); + } + } + } + expandedRowIdxCopyOutQueue.FreeTensor(expandedRowIdx); + inputXCopyOutQueue.FreeTensor(outLocal); +} + +template +__aicore__ inline void MoeV2FullLoadQuant::Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR offset, + GM_ADDR expandedX, GM_ADDR expandedRowIdx, + GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace, + const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe) { + this->InitBase(x, expertIdx, expandedX, expandedRowIdx, expertTokensCountOrCumsum, workspace, tilingData, tPipe); + xGm.SetGlobalBuffer((__gm__ T*)x); + scaleGm.SetGlobalBuffer((__gm__ float*)scale, 1); + offsetGm.SetGlobalBuffer((__gm__ float*)offset, 1); + this->scale = scaleGm.GetValue(0); + this->offset = offsetGm.GetValue(0); + + int64_t curRowsStart = this->blockIdx * this->perCoreRows; + int64_t rowLength = (curRowsStart + this->coreRows - 1) / this->k - curRowsStart / this->k + 1; + int64_t xAlignedCount = Align(this->cols, sizeof(int8_t)); + pipe->InitBuffer(xCopyInQueue, bufferNum, xAlignedCount * sizeof(T) * rowLength); + pipe->InitBuffer(inputXCopyOutQueue, 1, xAlignedCount * sizeof(int8_t) * rowLength); + pipe->InitBuffer(floatQueue, 1, xAlignedCount * sizeof(float) * rowLength); + pipe->InitBuffer(halfQueue, 1, xAlignedCount * sizeof(half) * rowLength); +} + +template +__aicore__ inline void MoeV2FullLoadQuant::Process() { + if (this->blockIdx < this->needCoreNum) { + this->ProcessBase(); + CopyOutX(); + } +} +} // namespace MoeInitRoutingQuantV2 +#endif \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_fullload_quant_base.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_fullload_quant_base.h new file mode 100644 index 000000000..7cc61fda8 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_fullload_quant_base.h @@ -0,0 +1,279 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file moe_v2_fullload_quant_base.h + * \brief + */ +#ifndef MOE_V2_FULL_LOAD_QUANT_BASE_H +#define MOE_V2_FULL_LOAD_QUANT_BASE_H + +#include "kernel_operator.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +class MoeV2FullLoadQuantBase { + public: + __aicore__ inline MoeV2FullLoadQuantBase(){}; + + protected: + __aicore__ inline void InitBase(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX, GM_ADDR expandedRowIdx, + GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace, + const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe); + __aicore__ inline void ProcessBase(); + __aicore__ inline void CopyIn(); + __aicore__ inline void SortCompute(); + __aicore__ inline void CopyOutIdx(); + __aicore__ inline void CopyOutEmpty(); + __aicore__ inline void ComputeExpertTokenCountOrCumsum(); + + protected: + const InnerMoeV2GatherOutComputeTilingData* gatherOutTilingData; + + TPipe* pipe; + int64_t tileLength; + int64_t bufferNum = 1; + int64_t totalLength; + int64_t coreNum; + int64_t sortNum; + int64_t blockIdx; + int64_t needCoreNum; + int64_t coreRows; + int64_t perCoreRows; + int64_t k; + int64_t n; + int64_t cols; + int64_t activateRows; + int64_t expertNum; + int64_t expertCapacity; + + TQue sortDataCopyInQueue; + TBuf tempBuffer; + TBuf sortedBuffer; + TQue xCopyInQueue; + TQue expandedRowIdxCopyOutQueue; + TQue expandedExpertIdxCopyOutQueue; + TQue expandDstToSrcRowQueue; + TQue expertTokensCopyOutQueue; + + GlobalTensor expertIdxGm; + GlobalTensor expandedXGm; + GlobalTensor expandedRowIdxGm; + GlobalTensor expandedExpertIdxGm; + GlobalTensor expertTokensCountOrCumsumGm; + GlobalTensor expertTokensBeforeCapacityGm; + + int64_t expertTokensCountOrCumsumFlag = 0; + int64_t expertTokensBeforeCapacityFlag = 0; + int64_t dropPadMode = 0; + static constexpr int64_t DST_BLK_STRIDE = 1; + static constexpr int64_t DST_REP_STRIDE = 8; + static constexpr int64_t FOUR_BLOCK_BYTES = 128; +}; + +__aicore__ inline void MoeV2FullLoadQuantBase::CopyIn() { + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + DataCopyExtParams dataCopyParams{static_cast(1), static_cast(this->totalLength * sizeof(int32_t)), + 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal[0], expertIdxGm, dataCopyParams, dataCopyPadParams); + ArithProgression(inLocal[this->sortNum], 0, 1, this->totalLength); + sortDataCopyInQueue.EnQue(inLocal); +} + +__aicore__ inline void MoeV2FullLoadQuantBase::SortCompute() { + LocalTensor inLocal = sortDataCopyInQueue.DeQue(); + LocalTensor expertIdxLocal = inLocal[0]; + LocalTensor expertIdxLocalFp32 = expertIdxLocal.ReinterpretCast(); + Cast(expertIdxLocalFp32, expertIdxLocal, RoundMode::CAST_ROUND, this->totalLength); + AscendC::PipeBarrier(); + Muls(expertIdxLocalFp32, expertIdxLocalFp32, (float)-1, this->totalLength); + AscendC::PipeBarrier(); + int64_t duplicateNum = this->totalLength % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + int duplicateIndex = this->totalLength - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expertIdxLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + AscendC::PipeBarrier(); + } + LocalTensor concatLocal; + LocalTensor tempTensor = tempBuffer.Get(GetSortLen(this->sortNum)); + Concat(concatLocal, expertIdxLocalFp32, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + LocalTensor rowIdxLocal = inLocal[this->sortNum].template ReinterpretCast(); + LocalTensor sortedLocal = sortedBuffer.Get(GetSortLen(this->sortNum)); + Sort(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + LocalTensor expandedExpertIdxLocal = expandedExpertIdxCopyOutQueue.AllocTensor(); + LocalTensor expandDstToSrcRowLocal = expandDstToSrcRowQueue.AllocTensor(); + LocalTensor expandDstToSrcRowLocalFp32 = expandDstToSrcRowLocal.ReinterpretCast(); + Extract(expandedExpertIdxLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + Cast(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocal.ReinterpretCast(), RoundMode::CAST_ROUND, + this->totalLength); + AscendC::PipeBarrier(); + Muls(expandedExpertIdxLocal, expandedExpertIdxLocal, (float)-1, this->totalLength); + AscendC::PipeBarrier(); + LocalTensor expandedExpertIdxLocalInt32; + expandedExpertIdxLocalInt32 = expandedExpertIdxLocal.ReinterpretCast(); + Cast(expandedExpertIdxLocalInt32, expandedExpertIdxLocal, RoundMode::CAST_ROUND, this->totalLength); + AscendC::PipeBarrier(); + expandedExpertIdxCopyOutQueue.EnQue(expandedExpertIdxLocalInt32); + + LocalTensor expandedRowIdx = expandedRowIdxCopyOutQueue.AllocTensor(); + LocalTensor expandedRowIdxU32 = expandedRowIdx.ReinterpretCast(); + Muls(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocalFp32, (float)-1, this->totalLength); + AscendC::PipeBarrier(); + ArithProgression(inLocal[this->sortNum], 0, 1, this->totalLength); + AscendC::PipeBarrier(); + if (duplicateNum > 0) { + int duplicateIndex = this->totalLength - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expandDstToSrcRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + AscendC::PipeBarrier(); + } + Concat(concatLocal, expandDstToSrcRowLocalFp32, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + Sort(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + Extract(tempTensor, expandedRowIdxU32, sortedLocal, this->sortNum / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + expandedRowIdxCopyOutQueue.EnQue(expandedRowIdx); + sortDataCopyInQueue.FreeTensor(inLocal); + + expandDstToSrcRowQueue.FreeTensor(expandDstToSrcRowLocal); +} + +__aicore__ inline void MoeV2FullLoadQuantBase::CopyOutIdx() { + LocalTensor expandedRowIdx = expandedRowIdxCopyOutQueue.DeQue(); + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = this->totalLength * sizeof(int32_t); + DataCopyPad(expandedRowIdxGm, expandedRowIdx, intriParams); + expandedRowIdxCopyOutQueue.EnQue(expandedRowIdx); +} + +__aicore__ inline void MoeV2FullLoadQuantBase::ComputeExpertTokenCountOrCumsum() { + LocalTensor expandedExpertIdx = expandedExpertIdxCopyOutQueue.DeQue(); + LocalTensor expertTokensCount = expertTokensCopyOutQueue.AllocTensor(); + + int64_t expertNumAlign = Align(this->expertNum, sizeof(int32_t)); + Duplicate(expertTokensCount, 0, expertNumAlign); + SetWaitFlag(HardEvent::V_S); + + int32_t lastExpertId = expandedExpertIdx.GetValue(0); + int64_t tokenCount = 0; + int64_t lastExpertCount = 0; + for (int64_t i = 0; i < this->totalLength; i++) { + int32_t curExpertId = expandedExpertIdx.GetValue(i); + tokenCount++; + while (lastExpertId < curExpertId) { + expertTokensCount.SetValue(lastExpertId, tokenCount - 1); + if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_COUNT) { + tokenCount = 1; + } + lastExpertId++; + } + } + expertTokensCount.SetValue(lastExpertId, tokenCount); + if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_CUMSUM) { + lastExpertId++; + while (lastExpertId < this->expertNum) { + expertTokensCount.SetValue(lastExpertId, tokenCount); + lastExpertId++; + } + } + DataCopyExtParams copyParams{static_cast(1), static_cast(this->expertNum * sizeof(int32_t)), 0, 0, + 0}; + if (this->expertTokensCountOrCumsumFlag > 0) { + DataCopyPad(expertTokensCountOrCumsumGm, expertTokensCount, copyParams); + } + expertTokensCopyOutQueue.FreeTensor(expertTokensCount); + expandedExpertIdxCopyOutQueue.FreeTensor(expandedExpertIdx); +} + +__aicore__ inline void MoeV2FullLoadQuantBase::CopyOutEmpty() { + LocalTensor outLocal = expandedExpertIdxCopyOutQueue.DeQue(); + expandedExpertIdxCopyOutQueue.FreeTensor(outLocal); +} + +__aicore__ inline void MoeV2FullLoadQuantBase::InitBase(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX, + GM_ADDR expandedRowIdx, GM_ADDR expertTokensCountOrCumsum, + GM_ADDR workspace, + const MoeInitRoutingQuantV2TilingData* tilingData, + TPipe* tPipe) { + this->gatherOutTilingData = &(tilingData->gatherOutComputeParamsOp); + this->blockIdx = get_block_idx() + get_subblockid() * get_block_num(); + this->k = tilingData->k; + this->n = tilingData->n; + this->cols = tilingData->cols; + this->needCoreNum = this->gatherOutTilingData->needCoreNum; + this->perCoreRows = this->gatherOutTilingData->perCoreRows; + this->activateRows = this->gatherOutTilingData->activateRows; + if (this->blockIdx == this->gatherOutTilingData->needCoreNum - 1) { + this->coreRows = this->gatherOutTilingData->lastCoreRows; + } else { + this->coreRows = this->gatherOutTilingData->perCoreRows; + } + this->expertNum = tilingData->expertNum; + this->dropPadMode = tilingData->dropPadMode; + this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag; + + this->tileLength = Align(tilingData->vbsComputeParamsOp.lastCorePerLoopElements, sizeof(int32_t)); + this->sortNum = Ceil(this->tileLength, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + this->totalLength = tilingData->n * tilingData->k; + this->pipe = tPipe; + + expertIdxGm.SetGlobalBuffer((__gm__ int32_t*)expertIdx, this->tileLength); + + expandedXGm.SetGlobalBuffer((__gm__ int8_t*)expandedX); + expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t*)expandedRowIdx, this->tileLength); + if (this->expertTokensCountOrCumsumFlag > 0) { + // dropless + expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensCountOrCumsum, + Align(this->expertNum, sizeof(int32_t))); + } + + int64_t kvFactor = 2; + int64_t buffSize = this->sortNum * sizeof(int32_t); + + pipe->InitBuffer(expandedRowIdxCopyOutQueue, bufferNum, buffSize); + pipe->InitBuffer(expandedExpertIdxCopyOutQueue, bufferNum, buffSize); + pipe->InitBuffer(expertTokensCopyOutQueue, bufferNum, AlignBytes(this->expertNum, sizeof(int32_t))); + pipe->InitBuffer(expandDstToSrcRowQueue, bufferNum, buffSize); + pipe->InitBuffer(sortDataCopyInQueue, bufferNum, buffSize * kvFactor); + pipe->InitBuffer(tempBuffer, buffSize * kvFactor); + pipe->InitBuffer(sortedBuffer, buffSize * kvFactor); +} + +__aicore__ inline void MoeV2FullLoadQuantBase::ProcessBase() { + if (this->blockIdx < this->needCoreNum) { + CopyIn(); + SortCompute(); + if (this->blockIdx == 0) { + CopyOutIdx(); + } + if (this->blockIdx == this->needCoreNum - 1 && this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) { + ComputeExpertTokenCountOrCumsum(); + } else { + CopyOutEmpty(); + } + } +} + +} // namespace MoeInitRoutingQuantV2 +#endif // MOE_V2_FULL_LOAD_QUANT_BASE_H \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_gather_dynamic_quant.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_gather_dynamic_quant.h new file mode 100644 index 000000000..594f58407 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_gather_dynamic_quant.h @@ -0,0 +1,547 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_v2_gather_dynamic_quant.h + * \brief + */ +#ifndef MOE_V2_GATHER_DYNAMIC_QUANT_H +#define MOE_V2_GATHER_DYNAMIC_QUANT_H + +#include "moe_v2_common.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +template +class MoeV2GatherDynamicQuant { + public: + __aicore__ inline MoeV2GatherDynamicQuant(){}; + __aicore__ inline void Init(GM_ADDR inputX, GM_ADDR quantSmooth, GM_ADDR expandedRowIdx, GM_ADDR expandedX, + GM_ADDR dynamicQuantScale, GM_ADDR workspace, + const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void CopyInExpandedRowIdx(int64_t progress); + __aicore__ inline void CopyInExpandedExpertIdx(int64_t progress); + __aicore__ inline void CopyOutXQuant1H(int64_t progress); + __aicore__ inline void CopyOutXQuantEH(int64_t progress); + __aicore__ inline void Compute(LocalTensor& smoothLocal); + __aicore__ inline void CopyOutPartialXQuantEH(int64_t progress); + __aicore__ inline void CopyOutPartialXQuant1H(int64_t progress); + __aicore__ inline float ComputeMax(LocalTensor& inLocal, LocalTensor& tempLocal, + LocalTensor& dynamicQuantLocal, int32_t srcIdx, int32_t expertIdx, + int64_t j); + __aicore__ inline void ComputeScale(LocalTensor& inLocal, LocalTensor& tempLocal, float scaleTemp, + int64_t dstIndex, int64_t j); + + private: + TPipe* pipe; + TQue inputXInQueue; + TQue smoothInQueue; + TQue expandRowIdxInQueue; + TQue calcQueue; + TQue inputXOutQueue; + TQue scaleOutQueue; + + GlobalTensor inputXGm; + GlobalTensor expandedXGm; + GlobalTensor expandedRowIdxGm; + GlobalTensor quantSmoothGm; + GlobalTensor dynamicQuantScaleGm; + GlobalTensor quantSrcGm; + GlobalTensor expandedExpertIdxGm; + GlobalTensor sortedRowIdxGm; + + const InnerMoeV2GatherOutComputeTilingData* gatherOutTilingData; + + int64_t needCoreNum; + int64_t blockIdx; + int64_t cols; + int64_t cols_scale_; + int64_t n; + int64_t k; + int64_t totalLength; + int64_t activateRows; + int64_t currentLoopRows; + int64_t currentLoopRowsAlign; + int64_t coreRows; + int64_t perLoopRows; + int64_t lastLoopRows; + int64_t rowLoops; + int64_t colsTileLength; + int64_t perLoopCols; + int64_t perLoopColsAlign; + int64_t lastLoopCols; + int64_t colLoops; + int64_t dropPadMode; + int64_t smoothType; + + int64_t indicesOffset; + int64_t inputOffset; + int64_t outOffset; +}; + +template +__aicore__ inline void MoeV2GatherDynamicQuant::CopyInExpandedRowIdx(int64_t progress) { + this->indicesOffset = progress * this->perLoopRows; + LocalTensor indicesLocal = expandRowIdxInQueue.AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(this->currentLoopRows * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(indicesLocal, expandedRowIdxGm[indicesOffset], dataCopyParams, dataCopyPadParams); + expandRowIdxInQueue.EnQue(indicesLocal); +} + +template +__aicore__ inline void MoeV2GatherDynamicQuant::CopyInExpandedExpertIdx(int64_t progress) { + this->indicesOffset = progress * this->perLoopRows; + LocalTensor indicesLocal = expandRowIdxInQueue.AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(this->currentLoopRows * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(indicesLocal, sortedRowIdxGm[indicesOffset], dataCopyParams, dataCopyPadParams); + DataCopyPad(indicesLocal[currentLoopRowsAlign], expandedExpertIdxGm[indicesOffset], dataCopyParams, + dataCopyPadParams); + expandRowIdxInQueue.EnQue(indicesLocal); +} + +template +__aicore__ inline void MoeV2GatherDynamicQuant::Compute(LocalTensor& smoothLocal) { + LocalTensor inLocal = inputXInQueue.DeQue(); + + LocalTensor tempLocal = calcQueue.AllocTensor(); + LocalTensor outLocal = inputXOutQueue.AllocTensor(); + LocalTensor dynamicQuantLocal = outLocal[this->cols].template ReinterpretCast(); + + if constexpr (!IsSameType::value) { + Cast(inLocal, inLocal.ReinterpretCast()[perLoopColsAlign], RoundMode::CAST_NONE, this->cols); + AscendC::PipeBarrier(); + } + + if (smoothType != 0) { + Mul(inLocal, inLocal, smoothLocal, this->cols); + AscendC::PipeBarrier(); + } + + Abs(tempLocal, inLocal, this->cols); + AscendC::PipeBarrier(); + + ReduceMax(dynamicQuantLocal, tempLocal, tempLocal, this->cols); + AscendC::PipeBarrier(); + + float maxValue = dynamicQuantLocal.GetValue(0) / 127.0f; + + Duplicate(dynamicQuantLocal, maxValue, 8); + Duplicate(tempLocal, maxValue, this->cols); + AscendC::PipeBarrier(); + + Div(tempLocal, inLocal, tempLocal, this->cols); + AscendC::PipeBarrier(); + + Cast(tempLocal.ReinterpretCast(), tempLocal, RoundMode::CAST_TRUNC, this->cols); + AscendC::PipeBarrier(); + + Cast(outLocal, tempLocal.ReinterpretCast(), RoundMode::CAST_ROUND, this->cols); + + calcQueue.FreeTensor(tempLocal); + inputXOutQueue.EnQue(outLocal); +} + +template +__aicore__ inline void MoeV2GatherDynamicQuant::CopyOutXQuant1H(int64_t progress) { + LocalTensor indicesLocal = expandRowIdxInQueue.DeQue(); + + int64_t initialRow = this->gatherOutTilingData->perCoreRows * this->blockIdx + this->perLoopRows * progress; + int64_t curLoopRow = 0; + int64_t currentLoopStartRow = initialRow / this->k; + int64_t currentLoopLastRow = (initialRow + this->currentLoopRows - 1) / this->k; + DataCopyExtParams copyInParams{1, static_cast(this->cols * sizeof(T)), 0, 0, 0}; + DataCopyExtParams copyOutParams{1, static_cast((this->cols + BLOCK_BYTES) * sizeof(int8_t)), 0, 0, 0}; + DataCopyExtParams smoothParams{1, static_cast(this->cols * sizeof(float)), 0, 0, 0}; + + LocalTensor smoothLocal; + if (smoothType == 1) { + smoothLocal = smoothInQueue.AllocTensor(); + DataCopyPad(smoothLocal, quantSmoothGm, smoothParams, {false, 0, 0, 0}); + smoothInQueue.EnQue(smoothLocal); + smoothLocal = smoothInQueue.DeQue(); + } + + for (int64_t row = currentLoopStartRow; row <= currentLoopLastRow; row++) { + LocalTensor inLocal = inputXInQueue.AllocTensor(); + if constexpr (IsSameType::value) { + DataCopyPad(inLocal, inputXGm[row * this->cols], copyInParams, {false, 0, 0, 0}); + } else { + DataCopyPad(inLocal[perLoopColsAlign], inputXGm[row * this->cols], copyInParams, {false, 0, 0, 0}); + } + + inputXInQueue.EnQue(inLocal); + + // Compute quantization + Compute(smoothLocal); + + LocalTensor outLocal = inputXOutQueue.DeQue(); + + while (curLoopRow < this->currentLoopRows && initialRow / this->k == row) { + int32_t outIndex = indicesLocal.GetValue(curLoopRow); + curLoopRow++; + initialRow++; + if (outIndex == -1 || (this->dropPadMode == DROPLESS_MODE && outIndex >= this->activateRows)) { + continue; + } + // Scale is placed after the data position + DataCopyPad(expandedXGm[outIndex * cols_scale_], outLocal, copyOutParams); + } + inputXInQueue.FreeTensor(inLocal); + inputXOutQueue.FreeTensor(outLocal); + } + expandRowIdxInQueue.FreeTensor(indicesLocal); +} + +template +__aicore__ inline void MoeV2GatherDynamicQuant::CopyOutXQuantEH(int64_t progress) { + LocalTensor indicesLocal = expandRowIdxInQueue.DeQue(); + SetWaitFlag(HardEvent::MTE2_S); + + DataCopyExtParams copyInParams{1, static_cast(this->perLoopCols * sizeof(T)), 0, 0, 0}; + DataCopyExtParams smoothParams{1, static_cast(this->perLoopCols * sizeof(float)), 0, 0, 0}; + DataCopyExtParams copyOutParams{1, static_cast(this->perLoopCols * sizeof(int8_t)), 0, 0, 0}; + + int32_t lastExpertIdx = -1; + LocalTensor inLocal = inputXInQueue.AllocTensor(); + LocalTensor smoothLocal = smoothInQueue.AllocTensor(); + for (int64_t i = 0; i < this->currentLoopRows; i++) { + int64_t rowOffset = this->gatherOutTilingData->perCoreRows * this->blockIdx + this->perLoopRows * progress; + if (this->dropPadMode == DROPLESS_MODE && rowOffset + i >= this->activateRows) { + break; + } + int32_t srcIdx = indicesLocal.GetValue(i); + int32_t expertIdx = indicesLocal.GetValue(currentLoopRowsAlign + i); + + if constexpr (IsSameType::value) { + DataCopyPad(inLocal, inputXGm[srcIdx / this->k * this->cols], copyInParams, {false, 0, 0, 0}); + } else { + DataCopyPad(inLocal[perLoopColsAlign], inputXGm[srcIdx / this->k * this->cols], copyInParams, {false, 0, 0, 0}); + } + inputXInQueue.EnQue(inLocal); + + if (expertIdx != lastExpertIdx) { + DataCopyPad(smoothLocal, quantSmoothGm[expertIdx * this->cols], smoothParams, {false, 0, 0, 0}); + smoothInQueue.EnQue(smoothLocal); + smoothLocal = smoothInQueue.DeQue(); + lastExpertIdx = expertIdx; + } + + Compute(smoothLocal); + + LocalTensor quantScaleLocal = scaleOutQueue.DeQue(); + DataCopyPad(dynamicQuantScaleGm[(rowOffset + i)], quantScaleLocal, {1, 4, 0, 0, 0}); + + LocalTensor outLocal = inputXOutQueue.DeQue(); + DataCopyPad(expandedXGm[(rowOffset + i) * this->cols], outLocal, copyOutParams); + + inputXOutQueue.FreeTensor(outLocal); + scaleOutQueue.FreeTensor(quantScaleLocal); + } + + inputXInQueue.FreeTensor(inLocal); + smoothInQueue.FreeTensor(smoothLocal); + expandRowIdxInQueue.FreeTensor(indicesLocal); +} + +template +__aicore__ inline float MoeV2GatherDynamicQuant::ComputeMax(LocalTensor& inLocal, + LocalTensor& tempLocal, + LocalTensor& dynamicQuantLocal, int32_t srcIdx, + int32_t expertIdx, int64_t j) { + LocalTensor smoothLocal = smoothInQueue.AllocTensor(); + + DataCopyExtParams intriParamsT{1, static_cast(colsTileLength * sizeof(T)), 0, 0, 0}; + DataCopyExtParams intriParamsFp32{1, static_cast(colsTileLength * sizeof(float)), 0, 0, 0}; + + if constexpr (!IsSameType::value) { + DataCopyPad(inLocal.ReinterpretCast()[perLoopColsAlign], inputXGm[srcIdx * this->cols + j * this->perLoopCols], + intriParamsT, {false, 0, 0, 0}); + } else { + DataCopyPad(inLocal, inputXGm[srcIdx * this->cols + j * this->perLoopCols], intriParamsT, {false, 0, 0, 0}); + } + + inputXInQueue.EnQue(inLocal); + inLocal = inputXInQueue.DeQue(); + + if (smoothType != 0) { + DataCopyPad(smoothLocal, quantSmoothGm[expertIdx * this->cols + j * this->perLoopCols], intriParamsFp32, + {false, 0, 0, 0}); + smoothInQueue.EnQue(smoothLocal); + smoothLocal = smoothInQueue.DeQue(); + } + + if constexpr (!IsSameType::value) { + Cast(inLocal, inLocal.ReinterpretCast()[perLoopColsAlign], RoundMode::CAST_NONE, colsTileLength); + AscendC::PipeBarrier(); + } + + if (smoothType != 0) { + Mul(inLocal, inLocal, smoothLocal, colsTileLength); + AscendC::PipeBarrier(); + } + + Abs(tempLocal, inLocal, colsTileLength); + AscendC::PipeBarrier(); + + ReduceMax(dynamicQuantLocal[8], tempLocal, tempLocal, colsTileLength); + + DataCopyPad(quantSrcGm[j * this->perLoopCols], inLocal, intriParamsFp32); + smoothInQueue.FreeTensor(smoothLocal); + SetWaitFlag(HardEvent::MTE3_MTE2); + + return dynamicQuantLocal.GetValue(8); +} + +template +__aicore__ inline void MoeV2GatherDynamicQuant::ComputeScale(LocalTensor& inLocal, + LocalTensor& tempLocal, float scaleTemp, + int64_t dstIndex, int64_t j) { + DataCopyExtParams copyInParams{1, static_cast(colsTileLength * sizeof(float)), 0, 0, 0}; + DataCopyExtParams copyOutParams{1, static_cast(colsTileLength * sizeof(int8_t)), 0, 0, 0}; + + LocalTensor outLocal = inputXOutQueue.AllocTensor(); + + DataCopyPad(inLocal, quantSrcGm[j * this->perLoopCols], copyInParams, {false, 0, 0, 0}); + inputXInQueue.EnQue(inLocal); + inLocal = inputXInQueue.DeQue(); + + Duplicate(tempLocal, scaleTemp, colsTileLength); + AscendC::PipeBarrier(); + + Div(tempLocal, inLocal, tempLocal, colsTileLength); + AscendC::PipeBarrier(); + + Cast(tempLocal.ReinterpretCast(), tempLocal, RoundMode::CAST_TRUNC, colsTileLength); + AscendC::PipeBarrier(); + + Cast(outLocal, tempLocal.ReinterpretCast(), RoundMode::CAST_ROUND, colsTileLength); + + inputXOutQueue.EnQue(outLocal); + outLocal = inputXOutQueue.DeQue(); + DataCopyPad(expandedXGm[dstIndex * this->cols + j * this->perLoopCols], outLocal, copyOutParams); + + inputXOutQueue.FreeTensor(outLocal); + SetWaitFlag(HardEvent::MTE3_MTE2); +} + +template +__aicore__ inline void MoeV2GatherDynamicQuant::CopyOutPartialXQuantEH(int64_t progress) { + LocalTensor indicesLocal = expandRowIdxInQueue.DeQue(); + SetWaitFlag(HardEvent::MTE2_S); + + for (int64_t i = 0; i < this->currentLoopRows; i++) { + int64_t rowOffset = this->gatherOutTilingData->perCoreRows * this->blockIdx + this->perLoopRows * progress; + if (this->dropPadMode == DROPLESS_MODE && rowOffset + i >= this->activateRows) { + break; + } + int32_t srcIdx = indicesLocal.GetValue(i); + int32_t expertIdx = indicesLocal.GetValue(currentLoopRowsAlign + i); + + LocalTensor inLocal = inputXInQueue.AllocTensor(); + LocalTensor tempLocal = calcQueue.AllocTensor(); + LocalTensor quantScaleLocal = scaleOutQueue.AllocTensor(); + + uint32_t tmp = 0xFF7FFFFF; + float reduceMax = *((float*)&tmp); + for (int64_t j = 0; j < this->colLoops; j++) { + colsTileLength = this->perLoopCols; + if (j == this->colLoops - 1) { + colsTileLength = this->lastLoopCols; + } + float tileMax = ComputeMax(inLocal, tempLocal, quantScaleLocal, srcIdx / this->k, expertIdx, j); + reduceMax = (reduceMax > tileMax) ? reduceMax : tileMax; + } + + float scaleTemp = reduceMax / 127.0f; + Duplicate(quantScaleLocal, scaleTemp, 8); + scaleOutQueue.EnQue(quantScaleLocal); + quantScaleLocal = scaleOutQueue.DeQue(); + + DataCopyPad(dynamicQuantScaleGm[(rowOffset + i)], quantScaleLocal, {1, 4, 0, 0, 0}); + + for (int64_t j = 0; j < this->colLoops; j++) { + colsTileLength = this->perLoopCols; + if (j == this->colLoops - 1) { + colsTileLength = this->lastLoopCols; + } + + ComputeScale(inLocal, tempLocal, scaleTemp, rowOffset + i, j); + } + + inputXInQueue.FreeTensor(inLocal); + calcQueue.FreeTensor(tempLocal); + scaleOutQueue.FreeTensor(quantScaleLocal); + } + + expandRowIdxInQueue.FreeTensor(indicesLocal); +} + +template +__aicore__ inline void MoeV2GatherDynamicQuant::CopyOutPartialXQuant1H(int64_t progress) { + LocalTensor indicesLocal = expandRowIdxInQueue.DeQue(); + + int64_t initialRow = this->gatherOutTilingData->perCoreRows * this->blockIdx + this->perLoopRows * progress; + int64_t curLoopRow = 0; + + int64_t currentLoopStartRow = initialRow / this->k; + int64_t currentLoopLastRow = (initialRow + this->currentLoopRows - 1) / this->k; + + for (int64_t row = currentLoopStartRow; row <= currentLoopLastRow; row++) { + LocalTensor inLocal = inputXInQueue.AllocTensor(); + LocalTensor tempLocal = calcQueue.AllocTensor(); + LocalTensor quantScaleLocal = scaleOutQueue.AllocTensor(); + + uint32_t tmp = 0xFF7FFFFF; + float reduceMax = *((float*)&tmp); + for (int64_t j = 0; j < this->colLoops; j++) { + colsTileLength = this->perLoopCols; + if (j == this->colLoops - 1) { + colsTileLength = this->lastLoopCols; + } + + float tileMax = ComputeMax(inLocal, tempLocal, quantScaleLocal, row, 0, j); + reduceMax = (reduceMax > tileMax) ? reduceMax : tileMax; + } + + float scaleTemp = reduceMax / 127.0f; + Duplicate(quantScaleLocal, scaleTemp, 8); + scaleOutQueue.EnQue(quantScaleLocal); + quantScaleLocal = scaleOutQueue.DeQue(); + + while (curLoopRow < this->currentLoopRows && initialRow / this->k == row) { + int32_t outIndex = indicesLocal.GetValue(curLoopRow); + curLoopRow++; + initialRow++; + if (outIndex == -1 || (this->dropPadMode == DROPLESS_MODE && outIndex >= this->activateRows)) { + continue; + } + DataCopyPad(dynamicQuantScaleGm[outIndex], quantScaleLocal, {1, 4, 0, 0, 0}); + for (int64_t j = 0; j < this->colLoops; j++) { + colsTileLength = this->perLoopCols; + if (j == this->colLoops - 1) { + colsTileLength = this->lastLoopCols; + } + + ComputeScale(inLocal, tempLocal, scaleTemp, outIndex, j); + } + } + inputXInQueue.FreeTensor(inLocal); + calcQueue.FreeTensor(tempLocal); + scaleOutQueue.FreeTensor(quantScaleLocal); + } + + expandRowIdxInQueue.FreeTensor(indicesLocal); +} + +template +__aicore__ inline void MoeV2GatherDynamicQuant::Init(GM_ADDR inputX, GM_ADDR quantSmooth, GM_ADDR expandedRowIdx, + GM_ADDR expandedX, GM_ADDR dynamicQuantScale, GM_ADDR workspace, + const MoeInitRoutingQuantV2TilingData* tilingData, + TPipe* tPipe) { + this->pipe = tPipe; + this->blockIdx = get_block_idx() + get_subblockid() * get_block_num(); + this->gatherOutTilingData = &(tilingData->gatherOutComputeParamsOp); + + this->needCoreNum = this->gatherOutTilingData->needCoreNum; + this->activateRows = this->gatherOutTilingData->activateRows; + this->cols = tilingData->cols; + this->cols_scale_ = this->cols + UB_ALIGN; + this->n = tilingData->n; + this->k = tilingData->k; + this->totalLength = tilingData->n * tilingData->k; + this->dropPadMode = tilingData->dropPadMode; + this->smoothType = tilingData->smoothType; + + if (this->blockIdx == this->gatherOutTilingData->needCoreNum - 1) { + this->coreRows = this->gatherOutTilingData->lastCoreRows; + this->perLoopRows = this->gatherOutTilingData->lastCorePerLoopRows; + this->lastLoopRows = this->gatherOutTilingData->lastCoreLastLoopRows; + this->rowLoops = this->gatherOutTilingData->lastCoreLoops; + } else { + this->coreRows = this->gatherOutTilingData->perCoreRows; + this->perLoopRows = this->gatherOutTilingData->perCorePerLoopRows; + this->lastLoopRows = this->gatherOutTilingData->perCoreLastLoopRows; + this->rowLoops = this->gatherOutTilingData->perCoreLoops; + } + this->perLoopCols = this->gatherOutTilingData->perLoopCols; + this->lastLoopCols = this->gatherOutTilingData->lastLoopCols; + this->colLoops = this->gatherOutTilingData->colLoops; + this->perLoopColsAlign = Align(this->perLoopCols, sizeof(T)); + + inputXGm.SetGlobalBuffer((__gm__ T*)inputX); + expandedXGm.SetGlobalBuffer((__gm__ int8_t*)expandedX); + + expandedRowIdxGm.SetGlobalBuffer( + (__gm__ int32_t*)expandedRowIdx + this->blockIdx * this->gatherOutTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + + quantSmoothGm.SetGlobalBuffer((__gm__ float*)quantSmooth); + dynamicQuantScaleGm.SetGlobalBuffer((__gm__ float*)dynamicQuantScale); + + expandedExpertIdxGm.SetGlobalBuffer( + (__gm__ int32_t*)workspace + this->blockIdx * this->gatherOutTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + sortedRowIdxGm.SetGlobalBuffer((__gm__ int32_t*)workspace + Align(this->totalLength, sizeof(int32_t)) + + this->blockIdx * this->gatherOutTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + if (this->cols > 1) { + quantSrcGm.SetGlobalBuffer( + (__gm__ float*)workspace + Align(this->totalLength, sizeof(int32_t)) * 2 + this->blockIdx * this->cols, + this->cols * sizeof(float)); + } + + this->currentLoopRowsAlign = Align(this->perLoopRows, sizeof(int32_t)); + + int64_t perLoopColsAlignBytes = AlignBytes(this->perLoopCols, sizeof(T)); + perLoopColsAlignBytes = + Max(int64_t(perLoopColsAlignBytes * sizeof(float) / sizeof(T)), int64_t(BLOCK_BYTES + BLOCK_BYTES)); + + pipe->InitBuffer(expandRowIdxInQueue, BUFFER_NUM, 2 * AlignBytes(this->perLoopRows, sizeof(int32_t))); + pipe->InitBuffer(inputXInQueue, BUFFER_NUM, perLoopColsAlignBytes); + pipe->InitBuffer(smoothInQueue, BUFFER_NUM, AlignBytes(this->perLoopCols, sizeof(float))); + pipe->InitBuffer(calcQueue, 1, AlignBytes(this->perLoopCols, sizeof(float))); + pipe->InitBuffer(inputXOutQueue, 1, AlignBytes(this->perLoopCols, sizeof(int8_t))); +} + +template +__aicore__ inline void MoeV2GatherDynamicQuant::Process() { + if (this->blockIdx < this->needCoreNum) { + currentLoopRows = perLoopRows; + if (colLoops > 1) { // Cannot fit all data in one row, workspace is required + trap(); // Not supported + } else { // All data can fit in one row + if (smoothType == 2) { + for (int64_t loop = 0; loop < this->rowLoops - 1; loop++) { + CopyInExpandedExpertIdx(loop); + CopyOutXQuantEH(loop); + } + currentLoopRows = lastLoopRows; + CopyInExpandedExpertIdx(this->rowLoops - 1); + CopyOutXQuantEH(this->rowLoops - 1); + } else { + for (int64_t loop = 0; loop < this->rowLoops - 1; loop++) { + CopyInExpandedRowIdx(loop); + CopyOutXQuant1H(loop); + } + currentLoopRows = lastLoopRows; + CopyInExpandedRowIdx(this->rowLoops - 1); + CopyOutXQuant1H(this->rowLoops - 1); + } + } + } +} +} // namespace MoeInitRoutingQuantV2 +#endif // MOE_V2_GATHER_DYNAMIC_QUANT_H diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_gather_out.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_gather_out.h new file mode 100644 index 000000000..764272766 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_gather_out.h @@ -0,0 +1,181 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_v2_gather_out.h + * \brief + */ +#ifndef INNER_MOE_V2_GATHER_OUT_H +#define INNER_MOE_V2_GATHER_OUT_H + +#include "moe_v2_common.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +constexpr int64_t BUFFER_NUM = 2; + +template +class MoeV2GatherOut { + public: + __aicore__ inline MoeV2GatherOut(){}; + __aicore__ inline void Init(GM_ADDR inputX, GM_ADDR expandedRowIdx, GM_ADDR expandedX, GM_ADDR workspace, + const InnerMoeInitRoutingV2TilingData* tilingData, TPipe* tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void CopyInIndices(int64_t progress); + __aicore__ inline void CopyOut(int64_t progress); + + private: + TPipe* pipe; + TQueBind inputActivationsCopyInQueue; + TQue expandDstToSrcRowCopyInQueue; + + GlobalTensor inputXGm; + GlobalTensor expandedXGm; + GlobalTensor expandedRowIdxGm; + + const InnerMoeV2GatherOutComputeTilingData* gatherOutTilingData; + + int64_t needCoreNum; + int64_t blockIdx; + int64_t cols; + int64_t n; + int64_t k; + int64_t activateRows; + int64_t currentLoopRows; + int64_t coreRows; + int64_t perLoopRows; + int64_t lastLoopRows; + int64_t rowLoops; + int64_t colsTileLength; + int64_t perLoopCols; + int64_t lastLoopCols; + int64_t colLoops; + int64_t dropPadMode; + + int64_t indicesOffset; + int64_t inputOffset; + int64_t outOffset; +}; + +template +__aicore__ inline void MoeV2GatherOut::CopyInIndices(int64_t progress) { + this->indicesOffset = progress * this->perLoopRows; + LocalTensor indicesLocal = expandDstToSrcRowCopyInQueue.AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(this->currentLoopRows * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(indicesLocal, expandedRowIdxGm[indicesOffset], dataCopyParams, dataCopyPadParams); + + expandDstToSrcRowCopyInQueue.EnQue(indicesLocal); +} + +template +__aicore__ inline void MoeV2GatherOut::CopyOut(int64_t progress) { + LocalTensor indicesLocal = expandDstToSrcRowCopyInQueue.DeQue(); + SetWaitFlag(HardEvent::MTE2_S); + colsTileLength = this->perLoopCols; + for (int64_t colsLoop = 0; colsLoop < this->colLoops; colsLoop++) { + int64_t initialRow = this->gatherOutTilingData->perCoreRows * this->blockIdx + this->perLoopRows * progress; + int64_t curLoopRow = 0; + if (colsLoop == this->colLoops - 1) { + colsTileLength = this->lastLoopCols; + } + int64_t currentLoopStartRow = initialRow / this->k; + int64_t currentLoopLastRow = (initialRow + this->currentLoopRows - 1) / this->k; + for (int64_t row = currentLoopStartRow; row <= currentLoopLastRow; row++) { + LocalTensor inLocal = inputActivationsCopyInQueue.AllocTensor(); + // input row position + inputOffset = row * this->cols + colsLoop * this->perLoopCols; + DataCopyExtParams dataCopyParams{1, static_cast(this->colsTileLength * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal, inputXGm[inputOffset], dataCopyParams, dataCopyPadParams); + SetWaitFlag(HardEvent::MTE2_MTE3); + + DataCopyExtParams intriParams{1, static_cast(this->colsTileLength * sizeof(T)), 0, 0, 0}; + while (curLoopRow < this->currentLoopRows && initialRow / this->k == row) { + int32_t outIndex = indicesLocal.GetValue(curLoopRow); + curLoopRow++; + initialRow++; + if (outIndex == -1 || (this->dropPadMode == DROPLESS_MODE && outIndex >= this->activateRows)) { + continue; + } + outOffset = outIndex * cols + colsLoop * this->perLoopCols; +#ifdef __CCE_KT_TEST__ + // CPU twin debugging cannot use multi-core sync, so index may contain uninitialized dirty data; handle specially + if (outOffset > expandedXGm.GetSize()) { + continue; + } +#endif + DataCopyPad(expandedXGm[outOffset], inLocal, intriParams); + } + inputActivationsCopyInQueue.FreeTensor(inLocal); + } + } + expandDstToSrcRowCopyInQueue.FreeTensor(indicesLocal); +} + +template +__aicore__ inline void MoeV2GatherOut::Init(GM_ADDR inputX, GM_ADDR expandedRowIdx, GM_ADDR expandedX, + GM_ADDR workspace, const InnerMoeInitRoutingV2TilingData* tilingData, + TPipe* tPipe) { + this->pipe = tPipe; + this->blockIdx = get_block_idx() + get_subblockid() * get_block_num(); + this->gatherOutTilingData = &(tilingData->gatherOutComputeParamsOp); + + this->needCoreNum = this->gatherOutTilingData->needCoreNum; + this->activateRows = this->gatherOutTilingData->activateRows; + this->cols = tilingData->cols; + this->n = tilingData->n; + this->k = tilingData->k; + this->dropPadMode = tilingData->dropPadMode; + + if (this->blockIdx == this->gatherOutTilingData->needCoreNum - 1) { + this->coreRows = this->gatherOutTilingData->lastCoreRows; + this->perLoopRows = this->gatherOutTilingData->lastCorePerLoopRows; + this->lastLoopRows = this->gatherOutTilingData->lastCoreLastLoopRows; + this->rowLoops = this->gatherOutTilingData->lastCoreLoops; + } else { + this->coreRows = this->gatherOutTilingData->perCoreRows; + this->perLoopRows = this->gatherOutTilingData->perCorePerLoopRows; + this->lastLoopRows = this->gatherOutTilingData->perCoreLastLoopRows; + this->rowLoops = this->gatherOutTilingData->perCoreLoops; + } + this->perLoopCols = this->gatherOutTilingData->perLoopCols; + this->lastLoopCols = this->gatherOutTilingData->lastLoopCols; + this->colLoops = this->gatherOutTilingData->colLoops; + + inputXGm.SetGlobalBuffer((__gm__ T*)inputX, this->coreRows * this->cols); + expandedXGm.SetGlobalBuffer((__gm__ T*)expandedX, tilingData->n * tilingData->k * this->cols); + expandedRowIdxGm.SetGlobalBuffer( + (__gm__ int32_t*)expandedRowIdx + this->blockIdx * this->gatherOutTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + + pipe->InitBuffer(inputActivationsCopyInQueue, BUFFER_NUM, AlignBytes(this->perLoopCols, sizeof(T))); + pipe->InitBuffer(expandDstToSrcRowCopyInQueue, BUFFER_NUM, AlignBytes(this->perLoopRows, sizeof(int32_t))); +} + +template +__aicore__ inline void MoeV2GatherOut::Process() { + if (this->blockIdx < this->needCoreNum) { + currentLoopRows = perLoopRows; + for (int64_t loop = 0; loop < this->rowLoops; loop++) { + if (loop == this->rowLoops - 1) { + currentLoopRows = lastLoopRows; + } + CopyInIndices(loop); + CopyOut(loop); + } + } +} +} // namespace MoeInitRoutingQuantV2 +#endif // INNER_MOE_V2_GATHER_OUT_H diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_gather_quant.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_gather_quant.h new file mode 100644 index 000000000..f938fac94 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_gather_quant.h @@ -0,0 +1,235 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_v2_gather_quant.h + * \brief + */ +#ifndef MOE_V2_GATHER_QUANT_H +#define MOE_V2_GATHER_QUANT_H + +#include "moe_v2_common.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +constexpr int64_t BUFFER_NUM = 2; + +template +class MoeV2GatherQuant { + public: + __aicore__ inline MoeV2GatherQuant(){}; + __aicore__ inline void Init(GM_ADDR inputX, GM_ADDR scale, GM_ADDR offset, GM_ADDR expandedRowIdx, GM_ADDR expandedX, + GM_ADDR workspace, const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void CopyInIndices(int64_t progress); + __aicore__ inline void Compute(); + __aicore__ inline void CopyOut(int64_t progress); + + private: + TPipe* pipe; + TQue inputXCopyInQueue; + TQue expandRowIdxCopyInQueue; + TQue inputXCopyOutQueue; + TQue floatQueue; + TQue halfQueue; + + GlobalTensor inputXGm; + GlobalTensor expandedXGm; + GlobalTensor expandedRowIdxGm; + GlobalTensor scaleGm; + GlobalTensor offsetGm; + + const InnerMoeV2GatherOutComputeTilingData* gatherOutTilingData; + + int64_t needCoreNum; + int64_t blockIdx; + int64_t cols; + int64_t n; + int64_t k; + int64_t activateRows; + int64_t currentLoopRows; + int64_t coreRows; + int64_t perLoopRows; + int64_t lastLoopRows; + int64_t rowLoops; + int64_t colsTileLength; + int64_t perLoopCols; + int64_t lastLoopCols; + int64_t colLoops; + int64_t dropPadMode; + float scale; + float offset; + + int64_t indicesOffset; + int64_t inputOffset; + int64_t outOffset; +}; + +template +__aicore__ inline void MoeV2GatherQuant::CopyInIndices(int64_t progress) { + this->indicesOffset = progress * this->perLoopRows; + LocalTensor indicesLocal = expandRowIdxCopyInQueue.AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(this->currentLoopRows * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(indicesLocal, expandedRowIdxGm[indicesOffset], dataCopyParams, dataCopyPadParams); + expandRowIdxCopyInQueue.EnQue(indicesLocal); +} + +template +__aicore__ inline void MoeV2GatherQuant::Compute() { + LocalTensor inLocal = inputXCopyInQueue.DeQue(); + LocalTensor outLocal = inputXCopyOutQueue.AllocTensor(); + LocalTensor floatLocal = floatQueue.AllocTensor(); + LocalTensor halfLocal = halfQueue.AllocTensor(); + uint32_t elements = Align(this->colsTileLength, sizeof(T)); + if constexpr (IsSameType::value) { + Cast(floatLocal, inLocal, RoundMode::CAST_NONE, elements); + AscendC::PipeBarrier(); + Cast(halfLocal, floatLocal, RoundMode::CAST_NONE, elements); + AscendC::PipeBarrier(); + Muls(halfLocal, halfLocal, static_cast(this->scale), elements); + AscendC::PipeBarrier(); + Adds(halfLocal, halfLocal, static_cast(this->offset), elements); + AscendC::PipeBarrier(); + LocalTensor intLocal = floatLocal.ReinterpretCast(); + Cast(intLocal, halfLocal, RoundMode::CAST_RINT, elements); + AscendC::PipeBarrier(); + SetDeqScale((half)1.000000e+00f); + AscendC::PipeBarrier(); + Cast(halfLocal, intLocal, RoundMode::CAST_RINT, elements); + AscendC::PipeBarrier(); + Cast(outLocal, halfLocal, RoundMode::CAST_RINT, elements); + } else if constexpr (IsSameType::value) { + Cast(halfLocal, inLocal, RoundMode::CAST_NONE, elements); + AscendC::PipeBarrier(); + Muls(halfLocal, halfLocal, static_cast(this->scale), elements); + AscendC::PipeBarrier(); + Adds(halfLocal, halfLocal, static_cast(this->offset), elements); + AscendC::PipeBarrier(); + Cast(outLocal, halfLocal, RoundMode::CAST_RINT, elements); + } else { + Muls(inLocal, inLocal, static_cast(this->scale), elements); + AscendC::PipeBarrier(); + Adds(inLocal, inLocal, static_cast(this->offset), elements); + AscendC::PipeBarrier(); + Cast(outLocal, inLocal, RoundMode::CAST_RINT, elements); + } + inputXCopyOutQueue.EnQue(outLocal); + floatQueue.FreeTensor(floatLocal); + halfQueue.FreeTensor(halfLocal); +} + +template +__aicore__ inline void MoeV2GatherQuant::CopyOut(int64_t progress) { + LocalTensor indicesLocal = expandRowIdxCopyInQueue.DeQue(); + SetWaitFlag(HardEvent::MTE2_S); + colsTileLength = this->perLoopCols; + for (int64_t colsLoop = 0; colsLoop < this->colLoops; colsLoop++) { + int64_t initialRow = this->gatherOutTilingData->perCoreRows * this->blockIdx + this->perLoopRows * progress; + int64_t curLoopRow = 0; + if (colsLoop == this->colLoops - 1) { + colsTileLength = this->lastLoopCols; + } + int64_t currentLoopStartRow = initialRow / this->k; + int64_t currentLoopLastRow = (initialRow + this->currentLoopRows - 1) / this->k; + for (int64_t row = currentLoopStartRow; row <= currentLoopLastRow; row++) { + LocalTensor inLocal = inputXCopyInQueue.AllocTensor(); + // input row position + inputOffset = row * this->cols + colsLoop * this->perLoopCols; + DataCopyExtParams dataCopyParams{1, static_cast(this->colsTileLength * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal, inputXGm[inputOffset], dataCopyParams, dataCopyPadParams); + inputXCopyInQueue.EnQue(inLocal); + Compute(); + LocalTensor outLocal = inputXCopyOutQueue.DeQue(); + DataCopyExtParams intriParams{1, static_cast(this->colsTileLength * sizeof(int8_t)), 0, 0, 0}; + while (curLoopRow < this->currentLoopRows && initialRow / this->k == row) { + int32_t outIndex = indicesLocal.GetValue(curLoopRow); + curLoopRow++; + initialRow++; + if (outIndex == -1 || (this->dropPadMode == DROPLESS_MODE && outIndex >= this->activateRows)) { + continue; + } + outOffset = outIndex * cols + colsLoop * this->perLoopCols; + DataCopyPad(expandedXGm[outOffset], outLocal, intriParams); + } + inputXCopyInQueue.FreeTensor(inLocal); + inputXCopyOutQueue.FreeTensor(outLocal); + } + } + expandRowIdxCopyInQueue.FreeTensor(indicesLocal); +} + +template +__aicore__ inline void MoeV2GatherQuant::Init(GM_ADDR inputX, GM_ADDR scale, GM_ADDR offset, GM_ADDR expandedRowIdx, + GM_ADDR expandedX, GM_ADDR workspace, + const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe) { + this->pipe = tPipe; + this->blockIdx = get_block_idx() + get_subblockid() * get_block_num(); + this->gatherOutTilingData = &(tilingData->gatherOutComputeParamsOp); + + this->needCoreNum = this->gatherOutTilingData->needCoreNum; + this->activateRows = this->gatherOutTilingData->activateRows; + this->cols = tilingData->cols; + this->n = tilingData->n; + this->k = tilingData->k; + this->dropPadMode = tilingData->dropPadMode; + + if (this->blockIdx == this->gatherOutTilingData->needCoreNum - 1) { + this->coreRows = this->gatherOutTilingData->lastCoreRows; + this->perLoopRows = this->gatherOutTilingData->lastCorePerLoopRows; + this->lastLoopRows = this->gatherOutTilingData->lastCoreLastLoopRows; + this->rowLoops = this->gatherOutTilingData->lastCoreLoops; + } else { + this->coreRows = this->gatherOutTilingData->perCoreRows; + this->perLoopRows = this->gatherOutTilingData->perCorePerLoopRows; + this->lastLoopRows = this->gatherOutTilingData->perCoreLastLoopRows; + this->rowLoops = this->gatherOutTilingData->perCoreLoops; + } + this->perLoopCols = this->gatherOutTilingData->perLoopCols; + this->lastLoopCols = this->gatherOutTilingData->lastLoopCols; + this->colLoops = this->gatherOutTilingData->colLoops; + + inputXGm.SetGlobalBuffer((__gm__ T*)inputX); + expandedXGm.SetGlobalBuffer((__gm__ int8_t*)expandedX); + expandedRowIdxGm.SetGlobalBuffer( + (__gm__ int32_t*)expandedRowIdx + this->blockIdx * this->gatherOutTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + scaleGm.SetGlobalBuffer((__gm__ float*)scale, 1); + offsetGm.SetGlobalBuffer((__gm__ float*)offset, 1); + this->scale = scaleGm.GetValue(0); + this->offset = offsetGm.GetValue(0); + + pipe->InitBuffer(inputXCopyInQueue, BUFFER_NUM, AlignBytes(this->perLoopCols, sizeof(T))); + pipe->InitBuffer(inputXCopyOutQueue, BUFFER_NUM, AlignBytes(this->perLoopCols, sizeof(int8_t))); + pipe->InitBuffer(expandRowIdxCopyInQueue, BUFFER_NUM, AlignBytes(this->perLoopRows, sizeof(int32_t))); + pipe->InitBuffer(floatQueue, 1, AlignBytes(this->perLoopCols, sizeof(float))); + pipe->InitBuffer(halfQueue, 1, AlignBytes(this->perLoopCols, sizeof(half))); +} + +template +__aicore__ inline void MoeV2GatherQuant::Process() { + if (this->blockIdx < this->needCoreNum) { + currentLoopRows = perLoopRows; + for (int64_t loop = 0; loop < this->rowLoops; loop++) { + if (loop == this->rowLoops - 1) { + currentLoopRows = lastLoopRows; + } + CopyInIndices(loop); + CopyOut(loop); + } + } +} +} // namespace MoeInitRoutingQuantV2 +#endif // MOE_V2_GATHER_QUANT_H diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_init_routing_fullload.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_init_routing_fullload.h new file mode 100644 index 000000000..31709e06a --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_init_routing_fullload.h @@ -0,0 +1,312 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file moe_v2_init_routing_fullload.h + * \brief + */ +#ifndef INNER_MOE_V2_FULL_LOAD_H +#define INNER_MOE_V2_FULL_LOAD_H + +#include "moe_v2_mrgsort.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +template +class MoeV2FullLoad : public MoeV2SortBase { + public: + __aicore__ inline MoeV2FullLoad(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX, GM_ADDR expandedRowIdx, + GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace, + const InnerMoeInitRoutingV2TilingData* tilingData, TPipe* tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void CopyIn(); + __aicore__ inline void SortCompute(); + __aicore__ inline void CopyOutIdx(); + __aicore__ inline void CopyOutEmpty(); + __aicore__ inline void CopyOutX(); + __aicore__ inline void ComputeExpertTokenCountOrCumsum(); + + private: + int64_t sortNum_; + const InnerMoeV2GatherOutComputeTilingData* gatherOutTilingData_; + int64_t blockIdx_; + int64_t needCoreNum_; + int64_t coreRows_; + int64_t perCoreRows_; + int64_t k_; + int64_t n_; + int64_t cols_; + int64_t activateRows_; + int64_t expertNum; + int64_t expertCapacity; + + TQue xCopyInQueue_; + TQue expandedRowIdxCopyOutQueue_; + TQue expandedExpertIdxCopyOutQueue_; + TQue expandDstToSrcRowQueue_; + TQue expertTokensCopyOutQueue_; + + GlobalTensor xGm_; + GlobalTensor expertIdxGm_; + + GlobalTensor expandedXGm_; + GlobalTensor expandedRowIdxGm_; + GlobalTensor expandedExpertIdxGm_; + GlobalTensor expertTokensCountOrCumsumGm; + GlobalTensor expertTokensBeforeCapacityGm; + + int64_t expertTokensCountOrCumsumFlag = 0; + int64_t expertTokensBeforeCapacityFlag = 0; + int64_t dropPadMode = 0; +}; + +template +__aicore__ inline void MoeV2FullLoad::CopyIn() { + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + DataCopyExtParams dataCopyParams{static_cast(1), static_cast(this->totalLength * sizeof(int32_t)), + 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal[0], expertIdxGm_, dataCopyParams, dataCopyPadParams); + ArithProgression(inLocal[this->sortNum_], 0, 1, this->totalLength); + sortDataCopyInQueue.EnQue(inLocal); +} + +template +__aicore__ inline void MoeV2FullLoad::SortCompute() { + LocalTensor inLocal = sortDataCopyInQueue.DeQue(); + LocalTensor expertIdxLocal = inLocal[0]; + LocalTensor expertIdxLocalFp32 = expertIdxLocal.ReinterpretCast(); + Cast(expertIdxLocalFp32, expertIdxLocal, RoundMode::CAST_ROUND, this->totalLength); + AscendC::PipeBarrier(); + Muls(expertIdxLocalFp32, expertIdxLocalFp32, (float)-1, this->totalLength); + AscendC::PipeBarrier(); + int64_t duplicateNum = this->totalLength % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + int duplicateIndex = this->totalLength - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expertIdxLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + AscendC::PipeBarrier(); + } + LocalTensor concatLocal; + LocalTensor tempTensor = tempBuffer.Get(GetSortLen(this->sortNum_)); + Concat(concatLocal, expertIdxLocalFp32, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + LocalTensor rowIdxLocal = inLocal[this->sortNum_].template ReinterpretCast(); + LocalTensor sortedLocal = sortedBuffer.Get(GetSortLen(this->sortNum_)); + Sort(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + LocalTensor expandedExpertIdxLocal = expandedExpertIdxCopyOutQueue_.AllocTensor(); + LocalTensor expandDstToSrcRowLocal = expandDstToSrcRowQueue_.AllocTensor(); + LocalTensor expandDstToSrcRowLocalFp32 = expandDstToSrcRowLocal.ReinterpretCast(); + Extract(expandedExpertIdxLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum_ / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + Cast(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocal.ReinterpretCast(), RoundMode::CAST_ROUND, + this->totalLength); + AscendC::PipeBarrier(); + Muls(expandedExpertIdxLocal, expandedExpertIdxLocal, (float)-1, this->totalLength); + AscendC::PipeBarrier(); + LocalTensor expandedExpertIdxLocalInt32; + expandedExpertIdxLocalInt32 = expandedExpertIdxLocal.ReinterpretCast(); + Cast(expandedExpertIdxLocalInt32, expandedExpertIdxLocal, RoundMode::CAST_ROUND, this->totalLength); + AscendC::PipeBarrier(); + expandedExpertIdxCopyOutQueue_.EnQue(expandedExpertIdxLocalInt32); + + LocalTensor expandedRowIdx = expandedRowIdxCopyOutQueue_.AllocTensor(); + LocalTensor expandedRowIdxU32 = expandedRowIdx.ReinterpretCast(); + Muls(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocalFp32, (float)-1, this->totalLength); + AscendC::PipeBarrier(); + ArithProgression(inLocal[this->sortNum_], 0, 1, this->totalLength); + AscendC::PipeBarrier(); + if (duplicateNum > 0) { + int duplicateIndex = this->totalLength - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expandDstToSrcRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + AscendC::PipeBarrier(); + } + Concat(concatLocal, expandDstToSrcRowLocalFp32, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + Sort(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + Extract(tempTensor, expandedRowIdxU32, sortedLocal, this->sortNum_ / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + expandedRowIdxCopyOutQueue_.EnQue(expandedRowIdx); + sortDataCopyInQueue.FreeTensor(inLocal); + + expandDstToSrcRowQueue_.FreeTensor(expandDstToSrcRowLocal); +} + +template +__aicore__ inline void MoeV2FullLoad::CopyOutIdx() { + LocalTensor expandedRowIdx = expandedRowIdxCopyOutQueue_.DeQue(); + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = this->totalLength * sizeof(int32_t); + DataCopyPad(expandedRowIdxGm_, expandedRowIdx, intriParams); + expandedRowIdxCopyOutQueue_.EnQue(expandedRowIdx); +} + +template +__aicore__ inline void MoeV2FullLoad::ComputeExpertTokenCountOrCumsum() { + LocalTensor expandedExpertIdx = expandedExpertIdxCopyOutQueue_.DeQue(); + LocalTensor expertTokensCount = expertTokensCopyOutQueue_.AllocTensor(); + + int64_t expertNumAlign = Align(this->expertNum, sizeof(int32_t)); + Duplicate(expertTokensCount, 0, expertNumAlign); + SetWaitFlag(HardEvent::V_S); + + int32_t lastExpertId = expandedExpertIdx.GetValue(0); + int64_t tokenCount = 0; + int64_t lastExpertCount = 0; + for (int64_t i = 0; i < this->totalLength; i++) { + int32_t curExpertId = expandedExpertIdx.GetValue(i); + tokenCount++; + while (lastExpertId < curExpertId) { + expertTokensCount.SetValue(lastExpertId, tokenCount - 1); + if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_COUNT) { + tokenCount = 1; + } + lastExpertId++; + } + } + expertTokensCount.SetValue(lastExpertId, tokenCount); + if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_CUMSUM) { + lastExpertId++; + while (lastExpertId < this->expertNum) { + expertTokensCount.SetValue(lastExpertId, tokenCount); + lastExpertId++; + } + } + DataCopyExtParams copyParams{static_cast(1), static_cast(this->expertNum * sizeof(int32_t)), 0, 0, + 0}; + if (this->expertTokensCountOrCumsumFlag > 0) { + DataCopyPad(expertTokensCountOrCumsumGm, expertTokensCount, copyParams); + } + expertTokensCopyOutQueue_.FreeTensor(expertTokensCount); + expandedExpertIdxCopyOutQueue_.FreeTensor(expandedExpertIdx); +} + +template +__aicore__ inline void MoeV2FullLoad::CopyOutX() { + LocalTensor xLocal = xCopyInQueue_.AllocTensor(); + LocalTensor expandedRowIdx = expandedRowIdxCopyOutQueue_.DeQue(); + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = this->cols_ * sizeof(T); + int64_t inFactor = Align(this->cols_, sizeof(T)); + int64_t curRowsStart = this->blockIdx_ * this->perCoreRows_; + int64_t startXRow = curRowsStart / this->k_; + int64_t endXRow = (curRowsStart + this->coreRows_ - 1) / this->k_; + + DataCopyExtParams dataXCopyParams{static_cast(endXRow - startXRow + 1), + static_cast(this->cols_ * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataXCopyPadParams{false, 0, 0, 0}; + DataCopyPad(xLocal, xGm_[startXRow * this->cols_], dataXCopyParams, dataXCopyPadParams); + SetWaitFlag(HardEvent::MTE2_S); + + int64_t k = 0; + for (int64_t i = startXRow; i <= endXRow; i++) { + for (; k < this->perCoreRows_ && curRowsStart / this->k_ == i; curRowsStart++, k++) { + int32_t outIndex = expandedRowIdx.GetValue(curRowsStart); + if (outIndex < this->activateRows_) { + DataCopyPad(expandedXGm_[outIndex * this->cols_], xLocal[(i - startXRow) * inFactor], intriParams); + } + } + } + expandedRowIdxCopyOutQueue_.FreeTensor(expandedRowIdx); + xCopyInQueue_.FreeTensor(xLocal); +} + +template +__aicore__ inline void MoeV2FullLoad::CopyOutEmpty() { + LocalTensor outLocal = expandedExpertIdxCopyOutQueue_.DeQue(); + expandedExpertIdxCopyOutQueue_.FreeTensor(outLocal); +} + +template +__aicore__ inline void MoeV2FullLoad::Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX, GM_ADDR expandedRowIdx, + GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace, + const InnerMoeInitRoutingV2TilingData* tilingData, TPipe* tPipe) { + this->gatherOutTilingData_ = &(tilingData->gatherOutComputeParamsOp); + this->blockIdx_ = get_block_idx() + get_subblockid() * get_block_num(); + this->k_ = tilingData->k; + this->n_ = tilingData->n; + this->cols_ = tilingData->cols; + this->needCoreNum_ = this->gatherOutTilingData_->needCoreNum; + this->perCoreRows_ = this->gatherOutTilingData_->perCoreRows; + this->activateRows_ = this->gatherOutTilingData_->activateRows; + if (this->blockIdx_ == this->gatherOutTilingData_->needCoreNum - 1) { + this->coreRows_ = this->gatherOutTilingData_->lastCoreRows; + } else { + this->coreRows_ = this->gatherOutTilingData_->perCoreRows; + } + this->expertNum = tilingData->expertNum; + this->dropPadMode = tilingData->dropPadMode; + this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag; + + this->tileLength = Align(tilingData->vbsComputeParamsOp.lastCorePerLoopElements, sizeof(int32_t)); + this->sortNum_ = Ceil(this->tileLength, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + this->totalLength = tilingData->n * tilingData->k; + this->pipe = tPipe; + + xGm_.SetGlobalBuffer((__gm__ T*)x); + expertIdxGm_.SetGlobalBuffer((__gm__ int32_t*)expertIdx, this->tileLength); + + expandedXGm_.SetGlobalBuffer((__gm__ T*)expandedX); + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t*)expandedRowIdx, this->tileLength); + if (this->expertTokensCountOrCumsumFlag > 0) { + // dropless + expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensCountOrCumsum, + Align(this->expertNum, sizeof(int32_t))); + } + + int64_t kvFactor = 2; + int64_t buffSize = this->sortNum_ * sizeof(int32_t); + + int64_t curRowsStart = this->blockIdx_ * this->perCoreRows_; + int64_t startXRow = curRowsStart / this->k_; + int64_t endXRow = (curRowsStart + this->coreRows_ - 1) / this->k_; + pipe->InitBuffer(xCopyInQueue_, bufferNum, AlignBytes(this->cols_, sizeof(T)) * (endXRow - startXRow + 1)); + pipe->InitBuffer(expandedRowIdxCopyOutQueue_, bufferNum, buffSize); + pipe->InitBuffer(expandedExpertIdxCopyOutQueue_, bufferNum, buffSize); + pipe->InitBuffer(expertTokensCopyOutQueue_, bufferNum, AlignBytes(this->expertNum, sizeof(int32_t))); + pipe->InitBuffer(expandDstToSrcRowQueue_, bufferNum, buffSize); + pipe->InitBuffer(sortDataCopyInQueue, bufferNum, buffSize * kvFactor); + pipe->InitBuffer(tempBuffer, buffSize * kvFactor); + pipe->InitBuffer(sortedBuffer, buffSize * kvFactor); +} + +template +__aicore__ inline void MoeV2FullLoad::Process() { + if (this->blockIdx_ < this->needCoreNum_) { + CopyIn(); + SortCompute(); + if (this->blockIdx_ == 0) { + CopyOutIdx(); + } + if (this->blockIdx_ == this->needCoreNum_ - 1 && this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) { + ComputeExpertTokenCountOrCumsum(); + } else { + CopyOutEmpty(); + } + CopyOutX(); + } +} +} // namespace MoeInitRoutingQuantV2 +#endif // INNER_MOE_V2_FULL_LOAD_H \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_mrgsort.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_mrgsort.h new file mode 100644 index 000000000..f72b66373 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_mrgsort.h @@ -0,0 +1,189 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_v2_mrgsort.h + * \brief + */ +#ifndef INNER_MOE_V2_MRGSORT_H +#define INNER_MOE_V2_MRGSORT_H + +#include "moe_v2_common.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +struct MoeV2MrgsortParam { + int64_t perListElements; + int64_t lastListElements; + int64_t oneLoopMaxElements; +}; + +class MoeV2Mrgsort { + public: + __aicore__ inline MoeV2Mrgsort(){}; + __aicore__ inline void Init(MoeV2MrgsortParam* param); + __aicore__ inline void Process(); + __aicore__ inline void SetInput(GlobalTensor& gmInput, LocalTensor& ubInput); + __aicore__ inline void SetOutput(GlobalTensor& gmOutput, LocalTensor& ubOutput); + + private: + __aicore__ inline void CopyIn(); + __aicore__ inline void UpdateMrgParam(); + __aicore__ inline void MrgsortCompute(); + __aicore__ inline void UpdateSortInfo(); + __aicore__ inline void CopyOut(); + __aicore__ inline void ClearCache(); + + private: + MoeV2MrgsortParam* param = nullptr; + + GlobalTensor gmInputs[4]; + GlobalTensor gmOutput; + + LocalTensor ubInputs[4]; + LocalTensor ubOutput; + + int64_t listNum{0}; + int64_t remainListNum{0}; + int64_t outOffset{0}; + int64_t offsets[4]; + int64_t listRemainElements[4]; + int64_t lengths[4]; + int64_t allRemainElements{0}; + int64_t curLoopSortedNum{0}; + + // for MrgSort + uint16_t validBitTail{0}; + uint16_t elementCountListTail[4]; + uint32_t listSortedNums[4]; + LocalTensor tmpUbInputs[4]; +}; + +__aicore__ inline void MoeV2Mrgsort::ClearCache() { + this->listNum = 0; + this->allRemainElements = 0; + this->outOffset = 0; +} + +__aicore__ inline void MoeV2Mrgsort::SetInput(GlobalTensor& gmInput, LocalTensor& ubInput) { + this->gmInputs[listNum] = gmInput; + this->ubInputs[listNum] = ubInput; + this->listNum += 1; +} + +__aicore__ inline void MoeV2Mrgsort::SetOutput(GlobalTensor& gmOutput, LocalTensor& ubOutput) { + this->gmOutput = gmOutput; + this->ubOutput = ubOutput; +} + +__aicore__ inline void MoeV2Mrgsort::UpdateMrgParam() { + if (this->remainListNum == MERGE_LIST_TWO) { + elementCountListTail[MERGE_LIST_IDX_TWO] = 0; + elementCountListTail[MERGE_LIST_IDX_THREE] = 0; + validBitTail = 0b0011; + } else if (this->remainListNum == MERGE_LIST_THREE) { + elementCountListTail[MERGE_LIST_IDX_THREE] = 0; + validBitTail = 0b0111; + } else if (this->remainListNum == MERGE_LIST_FOUR) { + validBitTail = 0b1111; + } else { + validBitTail = 0b0001; + } +} + +__aicore__ inline void MoeV2Mrgsort::CopyIn() { + this->remainListNum = 0; + SetWaitFlag(HardEvent::MTE3_MTE2); + for (int64_t i = 0, j = 0; i < listNum; i++) { + lengths[i] = Min(param->oneLoopMaxElements, listRemainElements[i]); + if (lengths[i] > 0) { + DataCopy(this->ubInputs[i], this->gmInputs[i][offsets[i]], Align(GetSortLen(lengths[i]), sizeof(float))); + tmpUbInputs[j] = this->ubInputs[i]; + elementCountListTail[j] = lengths[i]; + this->remainListNum += 1; + j++; + } + } +} + +__aicore__ inline void MoeV2Mrgsort::MrgsortCompute() { + SetWaitFlag(HardEvent::MTE2_V); + if (this->remainListNum == MERGE_LIST_TWO) { + MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]); + MrgSort(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else if (this->remainListNum == MERGE_LIST_THREE) { + MrgSortSrcList sortListTail = + MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], tmpUbInputs[0]); + MrgSort(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else if (this->remainListNum == MERGE_LIST_FOUR) { + MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], + tmpUbInputs[MERGE_LIST_IDX_THREE]); + MrgSort(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else { + DataCopy(this->ubOutput, this->tmpUbInputs[0], Align(GetSortLen(elementCountListTail[0]), sizeof(float))); + listSortedNums[0] = elementCountListTail[0]; + } +} + +__aicore__ inline void MoeV2Mrgsort::UpdateSortInfo() { + curLoopSortedNum = 0; + for (int64_t i = 0, j = 0; i < listNum; i++) { + if (lengths[i] > 0) { + // update remain size + listRemainElements[i] -= listSortedNums[j]; + allRemainElements -= listSortedNums[j]; + // update offset + offsets[i] += GetSortOffset(listSortedNums[j]); + // update current loop sorted nums + curLoopSortedNum += listSortedNums[j]; + j += 1; + } + } +} + +__aicore__ inline void MoeV2Mrgsort::CopyOut() { + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = GetSortLen(curLoopSortedNum) * sizeof(float); + SetWaitFlag(HardEvent::V_MTE3); + DataCopyPad(this->gmOutput[outOffset], this->ubOutput, intriParams); + outOffset += GetSortLen(curLoopSortedNum); +} + +__aicore__ inline void MoeV2Mrgsort::Init(MoeV2MrgsortParam* param) { + this->param = param; + this->remainListNum = listNum; + + for (int64_t i = 0; i < listNum; i++) { + offsets[i] = GetSortOffset(param->perListElements * i); + if (i == listNum - 1) { + listRemainElements[i] = param->lastListElements; + } else { + listRemainElements[i] = param->perListElements; + } + allRemainElements += listRemainElements[i]; + } +} + +__aicore__ inline void MoeV2Mrgsort::Process() { + for (; allRemainElements > 0;) { + CopyIn(); + UpdateMrgParam(); + MrgsortCompute(); + UpdateSortInfo(); + CopyOut(); + } + + ClearCache(); +} +} // namespace MoeInitRoutingQuantV2 +#endif // INNER_MOE_V2_MRGSORT_H \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_mrgsort_out.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_mrgsort_out.h new file mode 100644 index 000000000..4c6ee76e7 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_mrgsort_out.h @@ -0,0 +1,213 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_v2_mrgsort_out.h + * \brief + */ +#ifndef INNER_MOE_V2_MRGSORT_OUT_H +#define INNER_MOE_V2_MRGSORT_OUT_H + +#include "moe_v2_mrgsort.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +class MoeV2MrgsortOut { + public: + __aicore__ inline MoeV2MrgsortOut(){}; + __aicore__ inline void Init(MoeV2MrgsortParam* param, TPipe* tPipe); + __aicore__ inline void Process(); + __aicore__ inline void SetInput(GlobalTensor& gmInput, LocalTensor& ubInput); + __aicore__ inline void SetOutput(GlobalTensor& gmOutput1, GlobalTensor& gmOutput2, + LocalTensor& ubOutput1, LocalTensor& ubOutput2); + __aicore__ inline void SetBuffer(LocalTensor& tempBuffer); + + private: + __aicore__ inline void CopyIn(); + __aicore__ inline void UpdateMrgParam(); + __aicore__ inline void MrgsortCompute(); + __aicore__ inline void UpdateSortInfo(); + __aicore__ inline void Extract(); + __aicore__ inline void CopyOut(); + __aicore__ inline void ClearCache(); + + private: + MoeV2MrgsortParam* param = nullptr; + + GlobalTensor gmInputs[4]; + GlobalTensor gmOutput1; + GlobalTensor gmOutput2; + + LocalTensor ubInputs[4]; + LocalTensor tempBuffer; + + // for extract + LocalTensor ubOutput1; + LocalTensor ubOutput2; + + // for copy out + LocalTensor ubOutputInt1; + LocalTensor ubOutputInt2; + + int64_t listNum{0}; + int64_t remainListNum{0}; + int64_t outOffset{0}; + int64_t offsets[4]; + int64_t listRemainElements[4]; + int64_t lengths[4]; + int64_t allRemainElements{0}; + int64_t curLoopSortedNum{0}; + + // for MrgSort + uint16_t validBitTail; + uint16_t elementCountListTail[4]; + uint32_t listSortedNums[4]; + LocalTensor tmpUbInputs[4]; +}; + +__aicore__ inline void MoeV2MrgsortOut::ClearCache() { + this->listNum = 0; + this->allRemainElements = 0; + this->outOffset = 0; +} + +__aicore__ inline void MoeV2MrgsortOut::SetInput(GlobalTensor& gmInput, LocalTensor& ubInput) { + this->gmInputs[listNum] = gmInput; + this->ubInputs[listNum] = ubInput; + this->listNum += 1; +} + +__aicore__ inline void MoeV2MrgsortOut::SetOutput(GlobalTensor& gmOutput1, GlobalTensor& gmOutput2, + LocalTensor& ubOutput1, LocalTensor& ubOutput2) { + this->gmOutput1 = gmOutput1; + this->ubOutput1 = ubOutput1; + this->ubOutputInt1 = ubOutput1.ReinterpretCast(); + + this->gmOutput2 = gmOutput2; + this->ubOutput2 = ubOutput2.ReinterpretCast(); + this->ubOutputInt2 = ubOutput2.ReinterpretCast(); +} + +__aicore__ inline void MoeV2MrgsortOut::SetBuffer(LocalTensor& tempBuffer) { + this->tempBuffer = tempBuffer; +} + +__aicore__ inline void MoeV2MrgsortOut::UpdateMrgParam() { + if (this->remainListNum == MERGE_LIST_TWO) { + elementCountListTail[MERGE_LIST_IDX_TWO] = 0; + elementCountListTail[MERGE_LIST_IDX_THREE] = 0; + validBitTail = 0b0011; + } else if (this->remainListNum == MERGE_LIST_THREE) { + elementCountListTail[MERGE_LIST_IDX_THREE] = 0; + validBitTail = 0b0111; + } else if (this->remainListNum == MERGE_LIST_FOUR) { + validBitTail = 0b1111; + } else { + validBitTail = 0b0001; + } +} + +__aicore__ inline void MoeV2MrgsortOut::CopyIn() { + this->remainListNum = 0; + SetWaitFlag(HardEvent::MTE3_MTE2); + for (int64_t i = 0, j = 0; i < listNum; i++) { + lengths[i] = Min(param->oneLoopMaxElements, listRemainElements[i]); + if (lengths[i] > 0) { + DataCopy(this->ubInputs[i], this->gmInputs[i][offsets[i]], Align(GetSortLen(lengths[i]), sizeof(float))); + tmpUbInputs[j] = this->ubInputs[i]; + elementCountListTail[j] = lengths[i]; + this->remainListNum += 1; + j++; + } + } +} + +__aicore__ inline void MoeV2MrgsortOut::MrgsortCompute() { + SetWaitFlag(HardEvent::MTE2_V); + if (this->remainListNum == MERGE_LIST_TWO) { + MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]); + MrgSort(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else if (this->remainListNum == MERGE_LIST_THREE) { + MrgSortSrcList sortListTail = + MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], tmpUbInputs[0]); + MrgSort(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else if (this->remainListNum == MERGE_LIST_FOUR) { + MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], + tmpUbInputs[MERGE_LIST_IDX_THREE]); + MrgSort(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else { + DataCopy(this->tempBuffer, this->tmpUbInputs[0], Align(GetSortLen(elementCountListTail[0]), sizeof(float))); + listSortedNums[0] = elementCountListTail[0]; + } +} + +__aicore__ inline void MoeV2MrgsortOut::UpdateSortInfo() { + curLoopSortedNum = 0; + for (int64_t i = 0, j = 0; i < listNum; i++) { + if (lengths[i] > 0) { + // update remain size + listRemainElements[i] -= listSortedNums[j]; + allRemainElements -= listSortedNums[j]; + // update offset + offsets[i] += GetSortOffset(listSortedNums[j]); + // update current loop sorted nums + curLoopSortedNum += listSortedNums[j]; + j += 1; + } + } +} + +__aicore__ inline void MoeV2MrgsortOut::Extract() { + AscendC::Extract(this->ubOutput1, this->ubOutput2, this->tempBuffer, Ceil(curLoopSortedNum, ONE_REPEAT_SORT_NUM)); + AscendC::PipeBarrier(); + Muls(this->ubOutput1, this->ubOutput1, (float)-1, Align(curLoopSortedNum, sizeof(float))); + AscendC::PipeBarrier(); + Cast(this->ubOutputInt1, this->ubOutput1, RoundMode::CAST_ROUND, Align(curLoopSortedNum, sizeof(float))); +} + +__aicore__ inline void MoeV2MrgsortOut::CopyOut() { + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = curLoopSortedNum * sizeof(int32_t); + SetWaitFlag(HardEvent::V_MTE3); + DataCopyPad(this->gmOutput1[outOffset], this->ubOutputInt1, intriParams); + DataCopyPad(this->gmOutput2[outOffset], this->ubOutputInt2, intriParams); + outOffset += curLoopSortedNum; +} + +__aicore__ inline void MoeV2MrgsortOut::Init(MoeV2MrgsortParam* param, TPipe* tPipe) { + this->param = param; + this->allRemainElements = 0; + for (int64_t i = 0; i < listNum; i++) { + offsets[i] = GetSortOffset(param->perListElements * i); + if (i == listNum - 1) { + listRemainElements[i] = param->lastListElements; + } else { + listRemainElements[i] = param->perListElements; + } + allRemainElements += listRemainElements[i]; + } +} + +__aicore__ inline void MoeV2MrgsortOut::Process() { + for (; allRemainElements > 0;) { + CopyIn(); + UpdateMrgParam(); + MrgsortCompute(); + UpdateSortInfo(); + Extract(); + CopyOut(); + } + ClearCache(); +} +} // namespace MoeInitRoutingQuantV2 +#endif // INNER_MOE_V2_MRGSORT_OUT_H \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_sort_base.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_sort_base.h new file mode 100644 index 000000000..203afb6a6 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_sort_base.h @@ -0,0 +1,70 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_v2_sort_base.h + * \brief + */ +#ifndef INNER_MOE_V2_SORT_BASE_H +#define INNER_MOE_V2_SORT_BASE_H + +#include "kernel_operator.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +class MoeV2SortBase { + public: + __aicore__ inline MoeV2SortBase(){}; + + protected: + __aicore__ inline void SyncAll(); + + protected: + TPipe* pipe; + TQue sortDataCopyInQueue; + TQue sortDataCopyOutQueue; + TBuf tempBuffer; + TBuf sortedBuffer; + + GlobalTensor expertIdxGm; + GlobalTensor sortedexpertIdxGm; + GlobalTensor expandDstToSrcRowGm; + GlobalTensor expertTokensCountOrCumsumGm; + GlobalTensor expertTokensBeforeCapacityGm; + + int64_t tileLength; + int64_t bufferNum = 1; + int64_t totalLength; + int64_t coreNum; + int64_t n; + int64_t k; + int64_t existRowIdx; + int64_t expertNum; + int64_t expertTokensCountOrCumsumFlag = 0; + int64_t expertTokensBeforeCapacityFlag = 0; + + static constexpr int64_t SYNC_GM_NUM = 2; + static constexpr int64_t WORK_GM_NUM = 2; + static constexpr int64_t DST_BLK_STRIDE = 1; + static constexpr int64_t DST_REP_STRIDE = 8; +}; + +__aicore__ inline void MoeV2SortBase::SyncAll() { + if (coreNum == 1) { + return; + } +#ifndef __CCE_KT_TEST__ + AscendC::SyncAll(); +#endif +} + +} // namespace MoeInitRoutingQuantV2 +#endif // INNER_MOE_V2_SORT_BASE_H \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_sort_multi_core.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_sort_multi_core.h new file mode 100644 index 000000000..8db224ef1 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_sort_multi_core.h @@ -0,0 +1,373 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_v2_sort_multi_core.h + * \brief + */ +#ifndef INNER_MOE_V2_VBS_ONE_CORE_H +#define INNER_MOE_V2_VBS_ONE_CORE_H + +#include "moe_v2_sort_base.h" +#include "moe_v2_mrgsort.h" +#include "moe_v2_mrgsort_out.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +class MoeV2SortMultiCore : public MoeV2SortBase { + public: + __aicore__ inline MoeV2SortMultiCore(){}; + template + __aicore__ inline void Init(GM_ADDR expertIdx, GM_ADDR expertTokensCountOrCumsum, GM_ADDR expertTokensBeforeCapacity, + GM_ADDR workspace, const TilingData* tilingData, TPipe* tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void VBSProcess(); + __aicore__ inline void UBSortProcess(int64_t progress, int64_t size, int64_t sortNum); + __aicore__ inline void OneCoreVMSProcess(int64_t listNum, int64_t perListElements, int64_t lastListElements); + __aicore__ inline void VMSProcess(); + __aicore__ inline void SortOutProcess(); + __aicore__ inline void VBSCopyIn(int64_t progress, int64_t size, int64_t sortNum); + __aicore__ inline void UBSortCompute(int64_t progress, int64_t size, int64_t sortNum); + __aicore__ inline void VBSCopyOut(int64_t progress, int64_t size, int64_t sortNum); + __aicore__ inline void InitMoeMrgSort(MoeV2Mrgsort* sorter, int64_t listNum, int64_t coreOffset, int64_t loopOffset); + __aicore__ inline void InitMoeMrgSortOut(MoeV2MrgsortOut* sorter, int64_t listNum, int64_t coreOffset); + __aicore__ inline void InitExpertTokensGlobalMemory(); + + private: + GlobalTensor workspaceGms[2]; + + const InnerMoeV2VBSComputeTilingData* vbsTilingData; + const InnerMoeV2VMSMiddleComputeTilingData* vmsTilingData; + const InnerMoeV2SortOutComputeTilingData* sortOutTilingData; + + // for MoeMrgsort + MoeV2Mrgsort mrgsorter; + MoeV2MrgsortParam mrgsortParam; + + int64_t coreNum; + int64_t blockIdx; + int64_t srcWsIndex = 0; + + int64_t listNum; + int64_t perListElements; + int64_t lastListElements; + + int64_t sortTotalLength; + int64_t sortCoreLoops; + int64_t sortCoreLoopElements; + int64_t sortCoreLastLoopElements; + + int64_t perCoreExpert; + int64_t needInitExpertCore; + int64_t currentCoreExpert; + + static constexpr int64_t MAX_MRGSORT_LIST = 4; +}; + +__aicore__ inline void MoeV2SortMultiCore::InitExpertTokensGlobalMemory() { + if (this->blockIdx < this->needInitExpertCore) { + if (this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) { + InitGlobalMemory(expertTokensCountOrCumsumGm, currentCoreExpert, 0); + } + if (this->expertTokensBeforeCapacityFlag == EXERPT_TOKENS_BEFORE_CAPACITY) { + InitGlobalMemory(expertTokensBeforeCapacityGm, currentCoreExpert, 0); + } + } +} + +__aicore__ inline void MoeV2SortMultiCore::VBSCopyIn(int64_t progress, int64_t size, int64_t sortNum) { + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + int64_t inOffset = progress * sortCoreLoopElements; + DataCopyExtParams dataCopyParams{static_cast(1), static_cast(size * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal[0], expertIdxGm[inOffset], dataCopyParams, dataCopyPadParams); + + LocalTensor rowIdxLocal = inLocal[sortNum]; + int64_t startValue = this->blockIdx * this->vbsTilingData->perCoreElements + inOffset; + SetWaitFlag(HardEvent::MTE3_S); + ArithProgression(rowIdxLocal, startValue, 1, size); + sortDataCopyInQueue.EnQue(inLocal); +} + +__aicore__ inline void MoeV2SortMultiCore::UBSortCompute(int64_t progress, int64_t size, int64_t sortNum) { + LocalTensor inLocal = sortDataCopyInQueue.DeQue(); + LocalTensor expertForSourceRowLocal = inLocal[0]; + LocalTensor expertForSourceRowLocalFp32; + + expertForSourceRowLocalFp32 = expertForSourceRowLocal.ReinterpretCast(); + Cast(expertForSourceRowLocalFp32, expertForSourceRowLocal, RoundMode::CAST_ROUND, sortNum); + AscendC::PipeBarrier(); + Muls(expertForSourceRowLocalFp32, expertForSourceRowLocalFp32, (float)-1, sortNum); + AscendC::PipeBarrier(); + + int64_t duplicateNum = size % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + int duplicateIndex = size - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expertForSourceRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + AscendC::PipeBarrier(); + } + + LocalTensor concatLocal = expertForSourceRowLocalFp32; + LocalTensor sortedLocal = sortedBuffer.Get(GetSortLen(sortNum)); + LocalTensor outLocal = sortDataCopyOutQueue.AllocTensor(); + LocalTensor sourceRowLocal; + sourceRowLocal = inLocal[sortNum].ReinterpretCast(); + Sort(outLocal, concatLocal, sourceRowLocal, sortedLocal, sortNum / ONE_REPEAT_SORT_NUM); + + sortDataCopyOutQueue.EnQue(outLocal); + sortDataCopyInQueue.FreeTensor(inLocal); +} + +__aicore__ inline void MoeV2SortMultiCore::VBSCopyOut(int64_t progress, int64_t size, int64_t sortNum) { + LocalTensor outLocal = sortDataCopyOutQueue.DeQue(); + DataCopy(workspaceGms[0][this->blockIdx * GetSortLen(this->vbsTilingData->perCoreElements) + + GetSortLen(progress * sortCoreLoopElements)], + outLocal, Align(GetSortLen(size), sizeof(float))); + sortDataCopyOutQueue.FreeTensor(outLocal); +} + +__aicore__ inline void MoeV2SortMultiCore::InitMoeMrgSort(MoeV2Mrgsort* sorter, int64_t listNum, int64_t coreOffset, + int64_t loopOffset) { + GlobalTensor srcWsGm = workspaceGms[srcWsIndex][blockIdx * coreOffset + loopOffset]; + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + LocalTensor outLocal = sortDataCopyOutQueue.AllocTensor(); + for (int64_t i = 0; i < listNum; i++) { + LocalTensor inLocalT = inLocal[GetSortLen(this->sortOutTilingData->oneLoopMaxElements) * i]; + sorter->SetInput(srcWsGm, inLocalT); + } + GlobalTensor dstWsGm = workspaceGms[1 - srcWsIndex][blockIdx * coreOffset + loopOffset]; + sorter->SetOutput(dstWsGm, outLocal); + sortDataCopyInQueue.FreeTensor(inLocal); + sortDataCopyOutQueue.FreeTensor(outLocal); +} + +__aicore__ inline void MoeV2SortMultiCore::InitMoeMrgSortOut(MoeV2MrgsortOut* sorter, int64_t listNum, + int64_t coreOffset) { + GlobalTensor srcWsGm = workspaceGms[srcWsIndex]; + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + LocalTensor outLocal = sortDataCopyOutQueue.AllocTensor(); + + for (int64_t i = 0; i < listNum; i++) { + LocalTensor inLocalT = inLocal[GetSortLen(this->sortOutTilingData->oneLoopMaxElements) * i]; + sorter->SetInput(srcWsGm, inLocalT); + } + + LocalTensor outLocalV = outLocal[this->sortOutTilingData->oneLoopMaxElements * MAX_MRGSORT_LIST]; + sorter->SetOutput(this->sortedexpertIdxGm, this->expandDstToSrcRowGm, outLocal, outLocalV); + + LocalTensor tempBuffer = + sortedBuffer.Get(GetSortLen(this->sortOutTilingData->oneLoopMaxElements) * MAX_MRGSORT_LIST); + sorter->SetBuffer(tempBuffer); + sortDataCopyInQueue.FreeTensor(inLocal); + sortDataCopyOutQueue.FreeTensor(outLocal); +} + +__aicore__ inline void MoeV2SortMultiCore::OneCoreVMSProcess(int64_t listNum, int64_t perListElements, + int64_t lastListElements) { + int64_t coreOffset = GetSortLen(this->vbsTilingData->perCoreElements); + mrgsortParam.oneLoopMaxElements = this->sortOutTilingData->oneLoopMaxElements; + + for (int64_t i = 0; listNum >= 1; i++) { + int64_t loops = (listNum + MAX_MRGSORT_LIST - 1) / MAX_MRGSORT_LIST; + int64_t remainListNum = listNum - (loops - 1) * MAX_MRGSORT_LIST; + + mrgsortParam.perListElements = perListElements; + mrgsortParam.lastListElements = perListElements; + + int64_t loopOffset = GetSortLen(mrgsortParam.perListElements * MAX_MRGSORT_LIST); + for (int64_t loop = 0; loop < loops - 1; loop++) { + InitMoeMrgSort(&mrgsorter, MAX_MRGSORT_LIST, coreOffset, loop * loopOffset); + mrgsorter.Init(&mrgsortParam); + mrgsorter.Process(); + } + + mrgsortParam.perListElements = perListElements; + mrgsortParam.lastListElements = lastListElements; + InitMoeMrgSort(&mrgsorter, remainListNum, coreOffset, (loops - 1) * loopOffset); + mrgsorter.Init(&mrgsortParam); + mrgsorter.Process(); + + listNum = loops; + lastListElements = perListElements * (remainListNum - 1) + lastListElements; + perListElements = perListElements * MAX_MRGSORT_LIST; + srcWsIndex = (srcWsIndex + 1) % WORK_GM_NUM; + + if (loops == 1) { + break; + } + } +} + +__aicore__ inline void MoeV2SortMultiCore::UBSortProcess(int64_t progress, int64_t size, int64_t sortNum) { + VBSCopyIn(progress, size, sortNum); + UBSortCompute(progress, size, sortNum); + VBSCopyOut(progress, size, sortNum); +} + +__aicore__ inline void MoeV2SortMultiCore::VBSProcess() { + if (this->blockIdx < this->vbsTilingData->needCoreNum) { + int64_t sortNum = Ceil(sortCoreLoopElements, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + for (int64_t loop = 0; loop < sortCoreLoops - 1; loop++) { + UBSortProcess(loop, sortCoreLoopElements, sortNum); + } + + sortNum = Ceil(sortCoreLastLoopElements, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + UBSortProcess(sortCoreLoops - 1, sortCoreLastLoopElements, sortNum); + if (sortCoreLoops > 1) { + OneCoreVMSProcess(sortCoreLoops, sortCoreLoopElements, sortCoreLastLoopElements); + } + } +#ifndef __CCE_KT_TEST__ + AscendC::SyncAll(); +#endif +} + +__aicore__ inline void MoeV2SortMultiCore::VMSProcess() { + int64_t currentStageNeedCoreNum = this->vmsTilingData->needCoreNum; + perListElements = this->vbsTilingData->perCoreElements; + lastListElements = this->vbsTilingData->lastCoreElements; + listNum = this->vbsTilingData->needCoreNum; + + for (; listNum > MAX_MRGSORT_LIST;) { + currentStageNeedCoreNum = Ceil(listNum, MAX_MRGSORT_LIST); + int64_t coreOffset = GetSortLen(perListElements * MAX_MRGSORT_LIST); + int64_t remainListNum = listNum - (currentStageNeedCoreNum - 1) * MAX_MRGSORT_LIST; + + if (this->blockIdx < currentStageNeedCoreNum - 1) { + mrgsortParam.perListElements = perListElements; + mrgsortParam.lastListElements = perListElements; + mrgsortParam.oneLoopMaxElements = this->sortOutTilingData->oneLoopMaxElements; + InitMoeMrgSort(&mrgsorter, MAX_MRGSORT_LIST, coreOffset, 0); + mrgsorter.Init(&mrgsortParam); + mrgsorter.Process(); + } else if (this->blockIdx == currentStageNeedCoreNum - 1) { + mrgsortParam.perListElements = perListElements; + mrgsortParam.lastListElements = lastListElements; + mrgsortParam.oneLoopMaxElements = this->sortOutTilingData->oneLoopMaxElements; + InitMoeMrgSort(&mrgsorter, remainListNum, coreOffset, 0); + mrgsorter.Init(&mrgsortParam); + mrgsorter.Process(); + } + listNum = currentStageNeedCoreNum; + currentStageNeedCoreNum = Ceil(listNum, MAX_MRGSORT_LIST); + srcWsIndex = (srcWsIndex + 1) % WORK_GM_NUM; + + lastListElements = perListElements * (remainListNum - 1) + lastListElements; + perListElements = perListElements * MAX_MRGSORT_LIST; +#ifndef __CCE_KT_TEST__ + AscendC::SyncAll(); +#endif + } +} + +__aicore__ inline void MoeV2SortMultiCore::SortOutProcess() { + if (this->blockIdx < 1) { + mrgsortParam.perListElements = perListElements; + mrgsortParam.lastListElements = lastListElements; + mrgsortParam.oneLoopMaxElements = this->sortOutTilingData->oneLoopMaxElements; + + MoeV2MrgsortOut sorter; + InitMoeMrgSortOut(&sorter, listNum, GetSortLen(perListElements)); + sorter.Init(&mrgsortParam, pipe); + sorter.Process(); + } +#ifndef __CCE_KT_TEST__ + AscendC::SyncAll(); +#endif +} + +template +__aicore__ inline void MoeV2SortMultiCore::Init(GM_ADDR expertIdx, GM_ADDR expertTokensCountOrCumsum, + GM_ADDR expertTokensBeforeCapacity, GM_ADDR workspace, + const TilingData* tilingData, TPipe* tPipe) { + this->totalLength = tilingData->n * tilingData->k; + this->coreNum = tilingData->coreNum; + this->vbsTilingData = &(tilingData->vbsComputeParamsOp); + this->vmsTilingData = &(tilingData->vmsMiddleComputeParamsOp); + this->sortOutTilingData = &(tilingData->sortOutComputeParamsOp); + + this->blockIdx = get_block_idx() + get_subblockid() * get_block_num(); + this->tileLength = this->vbsTilingData->perCorePerLoopElements; + this->sortTotalLength = this->vbsTilingData->perCoreElements; + if (this->blockIdx == tilingData->vbsComputeParamsOp.needCoreNum - 1) { + this->tileLength = this->vbsTilingData->lastCorePerLoopElements; + this->sortTotalLength = this->vbsTilingData->lastCoreElements; + } + this->n = tilingData->n; + this->k = tilingData->k; + this->expertNum = tilingData->expertNum; + this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag; + this->expertTokensBeforeCapacityFlag = tilingData->expertTokensBeforeCapacityFlag; + + // VBS param init + if (this->blockIdx == this->vbsTilingData->needCoreNum - 1) { + sortCoreLoops = this->vbsTilingData->lastCoreLoops; + sortCoreLoopElements = this->vbsTilingData->lastCorePerLoopElements; + sortCoreLastLoopElements = this->vbsTilingData->lastCoreLastLoopElements; + } else { + sortCoreLoops = this->vbsTilingData->perCoreLoops; + sortCoreLoopElements = this->vbsTilingData->perCorePerLoopElements; + sortCoreLastLoopElements = this->vbsTilingData->perCoreLastLoopElements; + } + + this->pipe = tPipe; + expertIdxGm.SetGlobalBuffer( + (__gm__ int32_t*)expertIdx + this->blockIdx * tilingData->vbsComputeParamsOp.perCoreElements, + this->sortTotalLength); + sortedexpertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(workspace), + Align(this->totalLength, sizeof(int32_t))); + expandDstToSrcRowGm.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t*>(workspace) + Align(this->totalLength, sizeof(int32_t)), + Align(this->totalLength, sizeof(int32_t))); + + this->perCoreExpert = Align((this->expertNum + this->coreNum - 1) / this->coreNum, sizeof(int32_t)); + this->needInitExpertCore = (this->expertNum + this->perCoreExpert - 1) / this->perCoreExpert; + this->currentCoreExpert = this->perCoreExpert; + if (this->blockIdx == needInitExpertCore - 1) { + this->currentCoreExpert = this->expertNum - (this->needInitExpertCore - 1) * this->perCoreExpert; + } + if (this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) { + expertTokensCountOrCumsumGm.SetGlobalBuffer( + (__gm__ int32_t*)expertTokensCountOrCumsum + this->blockIdx * this->perCoreExpert, this->currentCoreExpert); + } + if (this->expertTokensBeforeCapacityFlag == EXERPT_TOKENS_BEFORE_CAPACITY) { + expertTokensBeforeCapacityGm.SetGlobalBuffer( + (__gm__ int32_t*)expertTokensBeforeCapacity + this->blockIdx * this->perCoreExpert, this->currentCoreExpert); + } + // key and value + int64_t kvFactor = 2; + workspaceGms[0].SetGlobalBuffer((__gm__ float*)workspace + Align(this->totalLength, sizeof(int32_t)) * 2, + Align(this->totalLength, sizeof(int32_t)) * kvFactor); + workspaceGms[1].SetGlobalBuffer((__gm__ float*)workspace + Align(this->totalLength, sizeof(int32_t)) * (kvFactor + 2), + Align(this->totalLength, sizeof(int32_t)) * kvFactor); + + int64_t bufferSize = Ceil(Max(this->sortOutTilingData->oneLoopMaxElements * MAX_MRGSORT_LIST, sortCoreLoopElements), + ONE_REPEAT_SORT_NUM) * + ONE_REPEAT_SORT_NUM * sizeof(int32_t) * kvFactor; + pipe->InitBuffer(sortDataCopyInQueue, bufferNum, bufferSize); + pipe->InitBuffer(sortDataCopyOutQueue, bufferNum, bufferSize); + pipe->InitBuffer(sortedBuffer, bufferSize); +} + +__aicore__ inline void MoeV2SortMultiCore::Process() { + InitExpertTokensGlobalMemory(); + VBSProcess(); + VMSProcess(); + SortOutProcess(); +} +} // namespace MoeInitRoutingQuantV2 +#endif // INNER_MOE_V2_VBS_ONE_CORE_H \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_sort_one_core.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_sort_one_core.h new file mode 100644 index 000000000..c370ef0e2 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_sort_one_core.h @@ -0,0 +1,162 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_v2_sort_one_core.h + * \brief + */ +#ifndef INNER_MOE_V2_SORT_ONE_CORE_H +#define INNER_MOE_V2_SORT_ONE_CORE_H + +#include "moe_v2_mrgsort.h" +#include "moe_v2_sort_base.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +class MoeV2SortOneCore : public MoeV2SortBase { + public: + __aicore__ inline MoeV2SortOneCore(){}; + template + __aicore__ inline void Init(GM_ADDR expertIdx, GM_ADDR expertTokensCountOrCumsum, GM_ADDR expertTokensBeforeCapacity, + GM_ADDR workspace, const TilingData* tilingData, TPipe* tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void CopyIn(); + __aicore__ inline void SortCompute(); + __aicore__ inline void CopyOut(); + + private: + int64_t sortNum; + int64_t blockIdx; +}; + +__aicore__ inline void MoeV2SortOneCore::CopyIn() { + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + DataCopyExtParams dataCopyParams{static_cast(1), static_cast(this->totalLength * sizeof(int32_t)), + 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal[0], expertIdxGm, dataCopyParams, dataCopyPadParams); + + LocalTensor rowIdxLocal = inLocal[this->sortNum]; + ArithProgression(rowIdxLocal, 0, 1, this->sortNum); + sortDataCopyInQueue.EnQue(inLocal); +} + +__aicore__ inline void MoeV2SortOneCore::SortCompute() { + LocalTensor inLocal = sortDataCopyInQueue.DeQue(); + LocalTensor expertForSourceRowLocal = inLocal[0]; + LocalTensor expertForSourceRowLocalFp32 = expertForSourceRowLocal.ReinterpretCast(); + Cast(expertForSourceRowLocalFp32, expertForSourceRowLocal, RoundMode::CAST_ROUND, this->tileLength); + AscendC::PipeBarrier(); + Muls(expertForSourceRowLocalFp32, expertForSourceRowLocalFp32, (float)-1, this->tileLength); + AscendC::PipeBarrier(); + + int64_t duplicateNum = this->totalLength % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + int duplicateIndex = this->totalLength - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expertForSourceRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + AscendC::PipeBarrier(); + } + + LocalTensor concatLocal; + LocalTensor tempTensor = tempBuffer.Get(GetSortLen(this->sortNum)); + Concat(concatLocal, expertForSourceRowLocalFp32, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + + LocalTensor sortedLocal = sortedBuffer.Get(GetSortLen(this->sortNum)); + LocalTensor sourceRowLocal; + sourceRowLocal = inLocal[this->sortNum].ReinterpretCast(); + Sort(sortedLocal, concatLocal, sourceRowLocal, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + + LocalTensor outLocal = sortDataCopyOutQueue.AllocTensor(); + LocalTensor sortedExpertForSourceRowLocal = outLocal[0]; + LocalTensor expandDstToSrcRowLocal; + expandDstToSrcRowLocal = outLocal[this->sortNum].ReinterpretCast(); + Extract(sortedExpertForSourceRowLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum / ONE_REPEAT_SORT_NUM); + AscendC::PipeBarrier(); + Muls(sortedExpertForSourceRowLocal, sortedExpertForSourceRowLocal, (float)-1, this->tileLength); + AscendC::PipeBarrier(); + + LocalTensor expertForSourceRowLocalInt32; + expertForSourceRowLocalInt32 = sortedExpertForSourceRowLocal.ReinterpretCast(); + Cast(expertForSourceRowLocalInt32, sortedExpertForSourceRowLocal, RoundMode::CAST_ROUND, this->tileLength); + sortDataCopyOutQueue.EnQue(outLocal); + sortDataCopyInQueue.FreeTensor(inLocal); +} + +__aicore__ inline void MoeV2SortOneCore::CopyOut() { + LocalTensor outLocal = sortDataCopyOutQueue.DeQue(); + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = this->totalLength * sizeof(int32_t); + DataCopyPad(sortedexpertIdxGm, outLocal[0], intriParams); + DataCopyPad(expandDstToSrcRowGm, outLocal[this->sortNum], intriParams); + sortDataCopyOutQueue.FreeTensor(outLocal); +} + +template +__aicore__ inline void MoeV2SortOneCore::Init(GM_ADDR expertIdx, GM_ADDR expertTokensCountOrCumsum, + GM_ADDR expertTokensBeforeCapacity, GM_ADDR workspace, + const TilingData* tilingData, TPipe* tPipe) { + this->blockIdx = get_block_idx() + get_subblockid() * get_block_num(); + this->tileLength = Align(tilingData->vbsComputeParamsOp.lastCorePerLoopElements, sizeof(int32_t)); + this->sortNum = Ceil(this->tileLength, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + this->totalLength = tilingData->n * tilingData->k; + this->coreNum = tilingData->coreNum; + this->pipe = tPipe; + this->n = tilingData->n; + this->k = tilingData->k; + this->expertNum = tilingData->expertNum; + this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag; + this->expertTokensBeforeCapacityFlag = tilingData->expertTokensBeforeCapacityFlag; + + expertIdxGm.SetGlobalBuffer((__gm__ int32_t*)expertIdx, this->tileLength); + sortedexpertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(workspace), this->tileLength); + expandDstToSrcRowGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(workspace) + this->tileLength, + this->tileLength); + + if (this->blockIdx == this->coreNum - 1) { + if (this->expertTokensCountOrCumsumFlag > 0) { + expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensCountOrCumsum, + Align(this->expertNum, sizeof(int32_t))); + InitGlobalMemory(expertTokensCountOrCumsumGm, this->expertNum, 0); + } + if (this->expertTokensBeforeCapacityFlag == 1) { + expertTokensBeforeCapacityGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensBeforeCapacity, + Align(this->expertNum, sizeof(int32_t))); + InitGlobalMemory(expertTokensBeforeCapacityGm, this->expertNum, 0); + } + } + // key and value + int64_t kvFactor = 2; + int64_t buffSize = this->sortNum * sizeof(int32_t) * kvFactor; + pipe->InitBuffer(sortDataCopyInQueue, bufferNum, buffSize); + pipe->InitBuffer(sortDataCopyOutQueue, bufferNum, buffSize); + pipe->InitBuffer(tempBuffer, buffSize); + pipe->InitBuffer(sortedBuffer, buffSize); +} + +__aicore__ inline void MoeV2SortOneCore::Process() { + if (get_block_idx() + get_subblockid() * get_block_num() < 1) { + CopyIn(); + SortCompute(); + CopyOut(); + } + this->SyncAll(); +} +} // namespace MoeInitRoutingQuantV2 +#endif // INNER_MOE_V2_SORT_ONE_CORE_H \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_src_to_dst_and_gather.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_src_to_dst_and_gather.h new file mode 100644 index 000000000..a327eb4eb --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_src_to_dst_and_gather.h @@ -0,0 +1,560 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_v2_src_to_dst_and_gather.h + * \brief + */ +#ifndef MOE_V2_SRC_TO_DST_AND_GATHER_H +#define MOE_V2_SRC_TO_DST_AND_GATHER_H + +#include "moe_v2_common.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +template +class MoeV2SrcToDstAndGather { + public: + __aicore__ inline MoeV2SrcToDstAndGather(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR scale, GM_ADDR expandedRowIdx, GM_ADDR expandedX, + GM_ADDR dynamicQuantScale, GM_ADDR workspace, const TilingData* tilingData, TPipe* tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void CopyIn(int64_t progress); + __aicore__ inline void CopyOut(int64_t progress); + __aicore__ inline void CopyOutLoops(int64_t progress); + __aicore__ inline void Compute(int32_t srcIdx, int32_t dstIdx, int32_t expertIdx); + __aicore__ inline float ComputeMax(LocalTensor& inLocal, LocalTensor& tempLocal, + LocalTensor& dynamicQuantLocal, int32_t srcIdx, int32_t expertIdx, + int64_t j); + __aicore__ inline void ComputeScale(LocalTensor& inLocal, LocalTensor& tempLocal, float scaleTemp, + int64_t dstIndex, int64_t j); + __aicore__ inline void ComputeLoops(int32_t srcIdx, int32_t dstIdx, int32_t expertIdx); + + __aicore__ inline void CopyOutRemain(); + __aicore__ inline void SyncAll(); + __aicore__ inline void AssistInit(); + + private: + TPipe* pipe; + TQue copyInQueue; + TQue copyOutQueue; + TQue copyOutZeroQueue; + + TQue inputXInQueue; + TQue smoothInQueue; + TQue calcQueue; + TQue inputXOutQueue; + TQue scaleOutQueue; + TQue scaleOutZeroQueue; + + GlobalTensor expandDstToSrcRowGm; + GlobalTensor expandedRowIdxGm; + GlobalTensor expertIdxValueGm; + GlobalTensor expandedExpertIdxGm; + GlobalTensor expandedXGm; + + GlobalTensor inputXGm; + GlobalTensor quantSmoothGm; + GlobalTensor dynamicQuantScaleGm; + GlobalTensor quantSrcGm; + + LocalTensor outTmpLocal; + LocalTensor scaleOutTmpLocal; + LocalTensor smoothLocal; + + const InnerMoeV2GatherOutComputeTilingData* srcToDstTilingData; + + int64_t coreNum; + int64_t blockIdx; + int64_t totalLength; + int64_t currentLoopRows; + int64_t coreRows; + int64_t perLoopRows; + int64_t lastLoopRows; + int64_t rowLoops; + int64_t expertCapacity; + int64_t expertNum; + int64_t cols; + int64_t perLoopCols; + int64_t lastLoopCols; + int64_t colLoops; + int64_t perLoopColsAlign; + int64_t k; + int64_t colsTileLength; + int64_t smoothType; + + int64_t tokenCount = 0; + int32_t lastExpertId = -1; + int32_t lastCoreExpertId = 0; + int32_t lastCoreExpertIdNum = 0; +}; + +template +__aicore__ inline void MoeV2SrcToDstAndGather::AssistInit() { + LocalTensor outLocal = copyOutZeroQueue.AllocTensor(); + Duplicate(outLocal, static_cast(0), this->perLoopCols); + copyOutZeroQueue.EnQue(outLocal); + LocalTensor scaleOutLocal = scaleOutZeroQueue.AllocTensor(); + Duplicate(scaleOutLocal, 0.0f, 8); + scaleOutZeroQueue.EnQue(scaleOutLocal); + + if (this->blockIdx != 0) { + this->lastCoreExpertId = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2); + this->lastCoreExpertIdNum = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2 + 1); + for (int64_t i = this->blockIdx - 2; i >= 0; i--) { + int32_t lastExpertIdx = expertIdxValueGm.GetValue(i * 2); + if (lastExpertIdx < this->lastCoreExpertId) { + break; + } + int32_t lastExpertNum = expertIdxValueGm.GetValue(i * 2 + 1); + this->lastCoreExpertIdNum += lastExpertNum; + } + } +} + +template +__aicore__ inline void MoeV2SrcToDstAndGather::CopyIn(int64_t progress) { + LocalTensor inLocal = copyInQueue.AllocTensor(); + int64_t length = Align(currentLoopRows, sizeof(int32_t)); + DataCopy(inLocal, expandDstToSrcRowGm[progress * perLoopRows], length); + DataCopy(inLocal[length], expandedExpertIdxGm[progress * perLoopRows], length); + + copyInQueue.EnQue(inLocal); +} + +template +__aicore__ inline void MoeV2SrcToDstAndGather::Compute(int32_t srcIdx, int32_t dstIdx, + int32_t expertIdx) { + DataCopyExtParams copyInParams{1, static_cast(this->cols * sizeof(T)), 0, 0, 0}; + DataCopyExtParams smoothParams{1, static_cast(this->cols * sizeof(float)), 0, 0, 0}; + DataCopyExtParams copyOutParams{1, static_cast(this->cols * sizeof(int8_t)), 0, 0, 0}; + + LocalTensor inLocal = inputXInQueue.AllocTensor(); + + if constexpr (IsSameType::value) { + DataCopyPad(inLocal, inputXGm[srcIdx / this->k * this->cols], copyInParams, {false, 0, 0, 0}); + } else { + DataCopyPad(inLocal.template ReinterpretCast()[perLoopColsAlign], inputXGm[srcIdx / this->k * this->cols], + copyInParams, {false, 0, 0, 0}); + } + + if (smoothType == 2) { + DataCopyPad(smoothLocal, quantSmoothGm[expertIdx * this->cols], smoothParams, {false, 0, 0, 0}); + } + + inputXInQueue.EnQue(inLocal); + smoothInQueue.EnQue(smoothLocal); + smoothLocal = smoothInQueue.DeQue(); + + inLocal = inputXInQueue.DeQue(); + + LocalTensor tempLocal = calcQueue.AllocTensor(); + LocalTensor outLocal = inputXOutQueue.AllocTensor(); + LocalTensor dynamicQuantLocal = scaleOutQueue.AllocTensor(); + + if constexpr (!IsSameType::value) { + Cast(inLocal, inLocal.template ReinterpretCast()[perLoopColsAlign], RoundMode::CAST_NONE, this->cols); + AscendC::PipeBarrier(); + } + + if (smoothType != 0) { + Mul(inLocal, inLocal, smoothLocal, this->cols); + AscendC::PipeBarrier(); + } + + Abs(tempLocal, inLocal, this->cols); + AscendC::PipeBarrier(); + + ReduceMax(dynamicQuantLocal, tempLocal, tempLocal, this->cols); + AscendC::PipeBarrier(); + + float maxValue = dynamicQuantLocal.GetValue(0) / 127.0f; + + Duplicate(dynamicQuantLocal, maxValue, 8); + Duplicate(tempLocal, maxValue, this->cols); + AscendC::PipeBarrier(); + + Div(tempLocal, inLocal, tempLocal, this->cols); + AscendC::PipeBarrier(); + + Cast(tempLocal.ReinterpretCast(), tempLocal, RoundMode::CAST_TRUNC, this->cols); + AscendC::PipeBarrier(); + + Cast(outLocal, tempLocal.ReinterpretCast(), RoundMode::CAST_ROUND, this->cols); + + calcQueue.FreeTensor(tempLocal); + inputXOutQueue.EnQue(outLocal); + scaleOutQueue.EnQue(dynamicQuantLocal); + + LocalTensor quantScaleLocal = scaleOutQueue.DeQue(); + DataCopyPad(dynamicQuantScaleGm[dstIdx], quantScaleLocal, {1, 4, 0, 0, 0}); + + outLocal = inputXOutQueue.DeQue(); +#ifndef __CCE_KT_TEST__ + DataCopyPad(expandedXGm[dstIdx * this->cols], outLocal, copyOutParams); +#endif + inputXInQueue.FreeTensor(inLocal); + inputXOutQueue.FreeTensor(outLocal); + scaleOutQueue.FreeTensor(quantScaleLocal); +} + +template +__aicore__ inline void MoeV2SrcToDstAndGather::CopyOut(int64_t progress) { + LocalTensor inLocal = copyInQueue.DeQue(); + LocalTensor outLocal = copyOutQueue.AllocTensor(); + int64_t length = Align(currentLoopRows, sizeof(int32_t)); + DataCopyExtParams copyParams{static_cast(1), static_cast(sizeof(int32_t)), 0, 0, 0}; + DataCopyExtParams copyParams1{static_cast(1), static_cast(this->cols * sizeof(int8_t)), 0, 0, 0}; + + SetWaitFlag(HardEvent::MTE2_S); + if (this->lastExpertId == -1) { + this->lastExpertId = this->lastCoreExpertId; + this->tokenCount = this->lastCoreExpertIdNum; + } + for (int64_t idx = 0; idx < currentLoopRows; idx++) { + int32_t expertIdx = inLocal[length].GetValue(idx); + SetWaitFlag(HardEvent::S_MTE3); + int32_t index = 0; + while (this->lastExpertId < expertIdx) { + while (this->tokenCount < this->expertCapacity) { + index = this->lastExpertId * this->expertCapacity + this->tokenCount; + DataCopyPad(expandedXGm[index * this->cols], this->outTmpLocal, copyParams1); + DataCopyPad(dynamicQuantScaleGm[index], this->scaleOutTmpLocal, {1, 4, 0, 0, 0}); + this->tokenCount++; + } + this->tokenCount = 0; + this->lastExpertId++; + } + + if (this->tokenCount < this->expertCapacity) { + int32_t outOffset = inLocal.GetValue(idx); + index = expertIdx * this->expertCapacity + this->tokenCount; + outLocal.SetValue(0, index); + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(expandedRowIdxGm[outOffset], outLocal, copyParams); + Compute(outOffset, index, expertIdx); + SetWaitFlag(HardEvent::MTE3_S); + this->tokenCount++; + } + } + copyInQueue.FreeTensor(inLocal); + copyOutQueue.FreeTensor(outLocal); +} + +template +__aicore__ inline float MoeV2SrcToDstAndGather::ComputeMax(LocalTensor& inLocal, + LocalTensor& tempLocal, + LocalTensor& dynamicQuantLocal, + int32_t srcIdx, int32_t expertIdx, + int64_t j) { + LocalTensor smoothLocal = smoothInQueue.AllocTensor(); + + DataCopyExtParams intriParamsT{1, static_cast(colsTileLength * sizeof(T)), 0, 0, 0}; + DataCopyExtParams intriParamsFp32{1, static_cast(colsTileLength * sizeof(float)), 0, 0, 0}; + + if constexpr (!IsSameType::value) { + DataCopyPad(inLocal.ReinterpretCast()[perLoopColsAlign], inputXGm[srcIdx * this->cols + j * this->perLoopCols], + intriParamsT, {false, 0, 0, 0}); + } else { + DataCopyPad(inLocal, inputXGm[srcIdx * this->cols + j * this->perLoopCols], intriParamsT, {false, 0, 0, 0}); + } + + inputXInQueue.EnQue(inLocal); + inLocal = inputXInQueue.DeQue(); + + if constexpr (!IsSameType::value) { + Cast(inLocal, inLocal.ReinterpretCast()[perLoopColsAlign], RoundMode::CAST_NONE, colsTileLength); + AscendC::PipeBarrier(); + } + + if (smoothType != 0) { + DataCopyPad(smoothLocal, quantSmoothGm[expertIdx * this->cols + j * this->perLoopCols], intriParamsFp32, + {false, 0, 0, 0}); + smoothInQueue.EnQue(smoothLocal); + smoothLocal = smoothInQueue.DeQue(); + + Mul(inLocal, inLocal, smoothLocal, colsTileLength); + AscendC::PipeBarrier(); + } + + Abs(tempLocal, inLocal, colsTileLength); + AscendC::PipeBarrier(); + + ReduceMax(dynamicQuantLocal[8], tempLocal, tempLocal, colsTileLength); + + DataCopyPad(quantSrcGm[j * this->perLoopCols], inLocal, intriParamsFp32); + smoothInQueue.FreeTensor(smoothLocal); + SetWaitFlag(HardEvent::MTE3_MTE2); + + return dynamicQuantLocal.GetValue(8); +} + +template +__aicore__ inline void MoeV2SrcToDstAndGather::ComputeScale(LocalTensor& inLocal, + LocalTensor& tempLocal, + float scaleTemp, int64_t dstIndex, + int64_t j) { + DataCopyExtParams copyInParams{1, static_cast(colsTileLength * sizeof(float)), 0, 0, 0}; + DataCopyExtParams copyOutParams{1, static_cast(colsTileLength * sizeof(int8_t)), 0, 0, 0}; + + LocalTensor outLocal = inputXOutQueue.AllocTensor(); + + DataCopyPad(inLocal, quantSrcGm[j * this->perLoopCols], copyInParams, {false, 0, 0, 0}); + inputXInQueue.EnQue(inLocal); + inLocal = inputXInQueue.DeQue(); + + Duplicate(tempLocal, scaleTemp, colsTileLength); + AscendC::PipeBarrier(); + + Div(tempLocal, inLocal, tempLocal, colsTileLength); + AscendC::PipeBarrier(); + + Cast(tempLocal.ReinterpretCast(), tempLocal, RoundMode::CAST_TRUNC, colsTileLength); + AscendC::PipeBarrier(); + + Cast(outLocal, tempLocal.ReinterpretCast(), RoundMode::CAST_ROUND, colsTileLength); + + inputXOutQueue.EnQue(outLocal); + outLocal = inputXOutQueue.DeQue(); + DataCopyPad(expandedXGm[dstIndex * this->cols + j * this->perLoopCols], outLocal, copyOutParams); + + inputXOutQueue.FreeTensor(outLocal); + SetWaitFlag(HardEvent::MTE3_MTE2); +} + +template +__aicore__ inline void MoeV2SrcToDstAndGather::ComputeLoops(int32_t srcIdx, int32_t dstIdx, + int32_t expertIdx) { + LocalTensor inLocal = inputXInQueue.AllocTensor(); + LocalTensor tempLocal = calcQueue.AllocTensor(); + LocalTensor quantScaleLocal = scaleOutQueue.AllocTensor(); + + uint32_t tmp = 0xFF7FFFFF; + float reduceMax = *((float*)&tmp); + for (int64_t j = 0; j < this->colLoops; j++) { + colsTileLength = this->perLoopCols; + if (j == this->colLoops - 1) { + colsTileLength = this->lastLoopCols; + } + float tileMax = ComputeMax(inLocal, tempLocal, quantScaleLocal, srcIdx / this->k, expertIdx, j); + reduceMax = (reduceMax > tileMax) ? reduceMax : tileMax; + } + + float scaleTemp = reduceMax / 127.0f; + Duplicate(quantScaleLocal, scaleTemp, 8); + scaleOutQueue.EnQue(quantScaleLocal); + quantScaleLocal = scaleOutQueue.DeQue(); + + DataCopyPad(dynamicQuantScaleGm[dstIdx], quantScaleLocal, {1, 4, 0, 0, 0}); + + for (int64_t j = 0; j < this->colLoops; j++) { + colsTileLength = this->perLoopCols; + if (j == this->colLoops - 1) { + colsTileLength = this->lastLoopCols; + } + ComputeScale(inLocal, tempLocal, scaleTemp, dstIdx, j); + } + + inputXInQueue.FreeTensor(inLocal); + calcQueue.FreeTensor(tempLocal); + scaleOutQueue.FreeTensor(quantScaleLocal); +} + +template +__aicore__ inline void MoeV2SrcToDstAndGather::CopyOutLoops(int64_t progress) { + LocalTensor inLocal = copyInQueue.DeQue(); + LocalTensor outLocal = copyOutQueue.AllocTensor(); + int64_t length = Align(currentLoopRows, sizeof(int32_t)); + DataCopyExtParams copyParams{static_cast(1), static_cast(sizeof(int32_t)), 0, 0, 0}; + + SetWaitFlag(HardEvent::MTE2_S); + if (this->lastExpertId == -1) { + this->lastExpertId = this->lastCoreExpertId; + this->tokenCount = this->lastCoreExpertIdNum; + } + for (int64_t idx = 0; idx < currentLoopRows; idx++) { + int32_t expertIdx = inLocal[length].GetValue(idx); + SetWaitFlag(HardEvent::S_MTE3); + int32_t index = 0; + while (this->lastExpertId < expertIdx) { + while (this->tokenCount < this->expertCapacity) { + index = this->lastExpertId * this->expertCapacity + this->tokenCount; + int64_t col = this->perLoopCols; + DataCopyPad(dynamicQuantScaleGm[index], this->scaleOutTmpLocal, {1, 4, 0, 0, 0}); + for (int64_t i = 0; i < this->colLoops; i++) { + if (i == this->colLoops - 1) { + col = this->lastLoopCols; + } + DataCopyExtParams copyParams1{static_cast(1), static_cast(col * sizeof(int8_t)), 0, 0, 0}; + DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, copyParams1); + SetWaitFlag(HardEvent::MTE3_S); + } + this->tokenCount++; + } + this->tokenCount = 0; + this->lastExpertId++; + } + + if (this->tokenCount < this->expertCapacity) { + int32_t outOffset = inLocal.GetValue(idx); + index = expertIdx * this->expertCapacity + this->tokenCount; + outLocal.SetValue(0, index); + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(expandedRowIdxGm[outOffset], outLocal, copyParams); + if (smoothType == 2) { + ComputeLoops(outOffset, index, expertIdx); + } else { + ComputeLoops(outOffset, index, 0); + } + SetWaitFlag(HardEvent::MTE3_S); + this->tokenCount++; + } + } + copyInQueue.FreeTensor(inLocal); + copyOutQueue.FreeTensor(outLocal); +} + +template +__aicore__ inline void MoeV2SrcToDstAndGather::CopyOutRemain() { + if (this->blockIdx != this->srcToDstTilingData->needCoreNum - 1) { + copyOutZeroQueue.FreeTensor(this->outTmpLocal); + scaleOutZeroQueue.FreeTensor(this->scaleOutTmpLocal); + return; + } + while (this->lastExpertId < this->expertNum) { + while (this->tokenCount < this->expertCapacity) { + int32_t index = this->lastExpertId * this->expertCapacity + this->tokenCount; + int64_t col = this->perLoopCols; + DataCopyPad(dynamicQuantScaleGm[index], this->scaleOutTmpLocal, {1, 4, 0, 0, 0}); + for (int64_t i = 0; i < this->colLoops; i++) { + if (i == this->colLoops - 1) { + col = this->lastLoopCols; + } + DataCopyExtParams copyParams{static_cast(1), static_cast(col * sizeof(int8_t)), 0, 0, 0}; + DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, copyParams); + SetWaitFlag(HardEvent::MTE3_S); + } + this->tokenCount++; + } + this->tokenCount = 0; + this->lastExpertId++; + } + copyOutZeroQueue.FreeTensor(this->outTmpLocal); + scaleOutZeroQueue.FreeTensor(this->scaleOutTmpLocal); +} + +template +__aicore__ inline void MoeV2SrcToDstAndGather::Init(GM_ADDR x, GM_ADDR scale, GM_ADDR expandedRowIdx, + GM_ADDR expandedX, GM_ADDR dynamicQuantScale, + GM_ADDR workspace, const TilingData* tilingData, + TPipe* tPipe) { + int64_t blockNum = GetBlockNum(); + this->pipe = tPipe; + this->blockIdx = get_block_idx() + get_subblockid() * get_block_num(); + + this->coreNum = tilingData->coreNum; + this->totalLength = tilingData->n * tilingData->k; + this->srcToDstTilingData = &(tilingData->srcToDstCapacityComputeParamsOp); + this->expertNum = tilingData->expertNum; + this->expertCapacity = tilingData->expertCapacity; + this->cols = tilingData->cols; + this->k = tilingData->k; + this->smoothType = tilingData->smoothType; + + if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) { + this->coreRows = this->srcToDstTilingData->lastCoreRows; + this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows; + this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows; + this->rowLoops = this->srcToDstTilingData->lastCoreLoops; + } else { + this->coreRows = this->srcToDstTilingData->perCoreRows; + this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows; + this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows; + this->rowLoops = this->srcToDstTilingData->perCoreLoops; + } + this->perLoopCols = this->srcToDstTilingData->perLoopCols; + this->lastLoopCols = this->srcToDstTilingData->lastLoopCols; + this->colLoops = this->srcToDstTilingData->colLoops; + this->perLoopColsAlign = Align(this->perLoopCols, sizeof(T)); + + inputXGm.SetGlobalBuffer((__gm__ T*)x); + quantSmoothGm.SetGlobalBuffer((__gm__ float*)scale); + dynamicQuantScaleGm.SetGlobalBuffer((__gm__ float*)dynamicQuantScale); + + int64_t length = Align(this->totalLength, sizeof(int32_t)); + expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t*)expandedRowIdx, length); + expandedXGm.SetGlobalBuffer((__gm__ int8_t*)expandedX, this->expertNum * this->expertCapacity * this->cols); + + expandedExpertIdxGm.SetGlobalBuffer( + (__gm__ int32_t*)workspace + this->blockIdx * this->srcToDstTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + expandDstToSrcRowGm.SetGlobalBuffer( + (__gm__ int32_t*)workspace + length + this->blockIdx * this->srcToDstTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + expertIdxValueGm.SetGlobalBuffer((__gm__ int32_t*)workspace + length * 2, this->coreNum * 2); + if (this->colLoops > 1) { + quantSrcGm.SetGlobalBuffer((__gm__ float*)workspace + length * 2 + this->coreNum * 2 + this->blockIdx * this->cols, + this->cols * sizeof(float)); + } + + pipe->InitBuffer(copyInQueue, 1, AlignBytes(this->perLoopRows, sizeof(int32_t)) * 2); + pipe->InitBuffer(copyOutQueue, 1, AlignBytes(INT32_ONE_BLOCK_NUM, sizeof(int32_t))); + pipe->InitBuffer(copyOutZeroQueue, 1, AlignBytes(this->perLoopCols, sizeof(int16_t))); + + int64_t perLoopColsAlignBytes = AlignBytes(this->perLoopCols, sizeof(T)); + perLoopColsAlignBytes = + Max(int64_t(perLoopColsAlignBytes * sizeof(float) / sizeof(T)), int64_t(BLOCK_BYTES + BLOCK_BYTES)); + + pipe->InitBuffer(inputXInQueue, 1, perLoopColsAlignBytes); + pipe->InitBuffer(smoothInQueue, 1, AlignBytes(this->perLoopCols, sizeof(float))); + pipe->InitBuffer(calcQueue, 1, AlignBytes(this->perLoopCols, sizeof(float))); + pipe->InitBuffer(inputXOutQueue, 1, AlignBytes(this->perLoopCols, sizeof(int8_t))); + pipe->InitBuffer(scaleOutQueue, 1, BLOCK_BYTES + BLOCK_BYTES); + pipe->InitBuffer(scaleOutZeroQueue, 1, BLOCK_BYTES); +} + +template +__aicore__ inline void MoeV2SrcToDstAndGather::Process() { + if (this->blockIdx < this->srcToDstTilingData->needCoreNum) { + AssistInit(); + this->outTmpLocal = copyOutZeroQueue.DeQue(); + this->scaleOutTmpLocal = scaleOutZeroQueue.DeQue(); + currentLoopRows = perLoopRows; + if (colLoops > 1) { + for (int64_t loop = 0; loop < this->rowLoops; loop++) { + if (loop == this->rowLoops - 1) { + currentLoopRows = lastLoopRows; + } + CopyIn(loop); + CopyOutLoops(loop); + } + } else { + smoothLocal = smoothInQueue.AllocTensor(); + if (smoothType == 1) { + DataCopyExtParams smoothParams{1, static_cast(this->cols * sizeof(float)), 0, 0, 0}; + DataCopyPad(smoothLocal, quantSmoothGm, smoothParams, {false, 0, 0, 0}); + } + for (int64_t loop = 0; loop < this->rowLoops; loop++) { + if (loop == this->rowLoops - 1) { + currentLoopRows = lastLoopRows; + } + CopyIn(loop); + CopyOut(loop); + } + smoothInQueue.FreeTensor(smoothLocal); + } + CopyOutRemain(); + } +} +} // namespace MoeInitRoutingQuantV2 +#endif // MOE_V2_SRC_TO_DST_AND_GATHER_H \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_src_to_dst_op.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_src_to_dst_op.h new file mode 100644 index 000000000..657378b4d --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_src_to_dst_op.h @@ -0,0 +1,164 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_v2_src_to_dst_op.h + * \brief + */ +#ifndef INNER_MOE_V2_SRC_TO_DST_H +#define INNER_MOE_V2_SRC_TO_DST_H + +#include "moe_v2_common.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +class MoeV2SrcToDstOp { + public: + __aicore__ inline MoeV2SrcToDstOp(){}; + template + __aicore__ inline void Init(GM_ADDR expandSrcToDstRow, GM_ADDR workspace, const TilingData* tilingData, TPipe* tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void CopyIn(int64_t progress); + __aicore__ inline void Compute(int64_t progress); + __aicore__ inline void CopyOut(); + __aicore__ inline void SyncAll(); + __aicore__ inline void AssistInit(); + + private: + TPipe* pipe; + TQue copyInQueue; + TQue copyOutQueue; + TBuf assistBuffer; + + GlobalTensor expandDstToSrcRowGm; + GlobalTensor expandSrcToDstRowGm; + GlobalTensor assistGm; + + const InnerMoeV2GatherOutComputeTilingData* srcToDstTilingData; + + int64_t coreNum; + int64_t blockIdx; + int64_t totalLength; + int64_t currentLoopRows; + int64_t coreRows; + int64_t perLoopRows; + int64_t lastLoopRows; +}; + +__aicore__ inline void MoeV2SrcToDstOp::AssistInit() { +#if defined(ASCENDC_OOM) && ASCENDC_OOM == 1 + OOMCheckAddrRange(assistGm.GetPhyAddr(), ASSIST_NUM * sizeof(int32_t)); +#endif + LocalTensor assistTensor = assistBuffer.Get(ASSIST_NUM); + DataCopy(assistTensor, assistGm, ASSIST_NUM); + SetWaitFlag(HardEvent::MTE2_V); + Adds(assistTensor, assistTensor, (int32_t)(this->blockIdx * this->srcToDstTilingData->perCoreRows), ASSIST_NUM); +} + +__aicore__ inline void MoeV2SrcToDstOp::CopyIn(int64_t progress) { + LocalTensor inLocal = copyInQueue.AllocTensor(); + DataCopy(inLocal, expandDstToSrcRowGm[progress * perLoopRows], Align(currentLoopRows, sizeof(int32_t))); + copyInQueue.EnQue(inLocal); +} + +__aicore__ inline void MoeV2SrcToDstOp::Compute(int64_t progress) { + LocalTensor outLocal = copyOutQueue.AllocTensor(); + LocalTensor assistTensor = assistBuffer.Get(ASSIST_NUM); + + AscendC::PipeBarrier(); + int64_t loops = Ceil(currentLoopRows, ASSIST_INDEX_NUM); + for (int64_t i = 0; i < loops; i++) { + Adds(outLocal[i * ASSIST_NUM], assistTensor, + static_cast(this->perLoopRows * progress + i * ASSIST_INDEX_NUM), ASSIST_NUM); + } + AscendC::PipeBarrier(); + copyOutQueue.EnQue(outLocal); +} + +__aicore__ inline void MoeV2SrcToDstOp::CopyOut() { + LocalTensor inLocal = copyInQueue.DeQue(); + LocalTensor outLocal = copyOutQueue.DeQue(); + SetWaitFlag(HardEvent::MTE2_S); + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = sizeof(int32_t); + uint32_t outOffset; + for (int64_t idx = 0; idx < currentLoopRows; idx++) { + outOffset = inLocal.GetValue(idx); + DataCopyPad(expandSrcToDstRowGm[outOffset], outLocal[idx * INT32_ONE_BLOCK_NUM], intriParams); + } + + copyInQueue.FreeTensor(inLocal); + copyOutQueue.FreeTensor(outLocal); +} + +__aicore__ inline void MoeV2SrcToDstOp::SyncAll() { + if (coreNum == 1) { + return; + } +#ifndef __CCE_KT_TEST__ + AscendC::SyncAll(); +#endif +} + +template +__aicore__ inline void MoeV2SrcToDstOp::Init(GM_ADDR expandSrcToDstRow, GM_ADDR workspace, const TilingData* tilingData, + TPipe* tPipe) { + int64_t blockNum = GetBlockNum(); + this->pipe = tPipe; + this->blockIdx = get_block_idx() + get_subblockid() * get_block_num(); + + this->coreNum = tilingData->coreNum; + this->totalLength = tilingData->n * tilingData->k; + this->srcToDstTilingData = &(tilingData->srcToDstComputeParamsOp); + + if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) { + this->coreRows = this->srcToDstTilingData->lastCoreRows; + this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows; + this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows; + } else { + this->coreRows = this->srcToDstTilingData->perCoreRows; + this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows; + this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows; + } + + expandSrcToDstRowGm.SetGlobalBuffer((__gm__ int32_t*)expandSrcToDstRow, Align(this->totalLength, sizeof(int32_t))); + expandDstToSrcRowGm.SetGlobalBuffer((__gm__ int32_t*)workspace + Align(this->totalLength, sizeof(int32_t)) + + this->blockIdx * this->srcToDstTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + assistGm.SetGlobalBuffer((__gm__ int32_t*)assist, ASSIST_NUM); + + pipe->InitBuffer(copyInQueue, 1, this->perLoopRows * BLOCK_BYTES); + pipe->InitBuffer(copyOutQueue, 1, Ceil(this->perLoopRows, ASSIST_NUM) * ASSIST_NUM * BLOCK_BYTES); + pipe->InitBuffer(assistBuffer, ASSIST_NUM * sizeof(int32_t)); +} + +__aicore__ inline void MoeV2SrcToDstOp::Process() { + if (this->blockIdx < this->srcToDstTilingData->needCoreNum) { + int64_t loops = (coreRows + perLoopRows - 1) / perLoopRows; + currentLoopRows = perLoopRows; + AssistInit(); + for (int64_t loop = 0; loop < loops - 1; loop++) { + CopyIn(loop); + Compute(loop); + CopyOut(); + } + currentLoopRows = lastLoopRows; + CopyIn(loops - 1); + Compute(loops - 1); + CopyOut(); + } + this->SyncAll(); +} +} // namespace MoeInitRoutingQuantV2 +#endif // INNER_MOE_V2_SRC_TO_DST_H \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_src_to_dst_with_capacity.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_src_to_dst_with_capacity.h new file mode 100644 index 000000000..770163fc5 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_src_to_dst_with_capacity.h @@ -0,0 +1,269 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_v2_src_to_dst_with_capacity.h + * \brief + */ +#ifndef INNER_MOE_V2_SRC_TO_DST_WITH_CAPACITY_H +#define INNER_MOE_V2_SRC_TO_DST_WITH_CAPACITY_H + +#include "moe_v2_common.h" + +namespace MoeInitRoutingQuantV2 { +using namespace AscendC; +using namespace optiling; +template +class MoeV2SrcToDstWithCapacity { + public: + __aicore__ inline MoeV2SrcToDstWithCapacity(){}; + __aicore__ inline void Init(GM_ADDR expandedRowIdx, GM_ADDR expandedX, GM_ADDR workspace, + const TilingData* tilingData, TPipe* tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void CopyIn(int64_t progress); + __aicore__ inline void CopyOut(int64_t progress); + __aicore__ inline void CopyOutRemain(); + __aicore__ inline void SyncAll(); + __aicore__ inline void AssistInit(); + + private: + TPipe* pipe; + TQue copyInQueue; + TQue copyOutQueue; + TQue copyOutZeroQueue; + + GlobalTensor expandDstToSrcRowGm; + GlobalTensor expandedRowIdxGm; + GlobalTensor expertIdxValueGm; + GlobalTensor expandedExpertIdxGm; + GlobalTensor expandedXGm; + + LocalTensor outTmpLocal; + + const InnerMoeV2GatherOutComputeTilingData* srcToDstTilingData; + + int64_t coreNum; + int64_t blockIdx; + int64_t totalLength; + int64_t currentLoopRows; + int64_t coreRows; + int64_t perLoopRows; + int64_t lastLoopRows; + int64_t rowLoops; + int64_t expertCapacity; + int64_t expertNum; + int64_t cols; + int64_t perLoopCols; + int64_t lastLoopCols; + int64_t colLoops; + + int64_t tokenCount = 0; + int32_t lastExpertId = -1; + int32_t lastCoreExpertId = 0; + int32_t lastCoreExpertIdNum = 0; +}; + +template +__aicore__ inline void MoeV2SrcToDstWithCapacity::AssistInit() { + if constexpr (IsSameType::value) { + LocalTensor outLocal = copyOutZeroQueue.AllocTensor(); + Duplicate(outLocal, static_cast(0), this->perLoopCols); + copyOutZeroQueue.EnQue(outLocal); + } else { + LocalTensor outLocal = copyOutZeroQueue.AllocTensor(); + Duplicate(outLocal, static_cast(0), this->perLoopCols); + copyOutZeroQueue.EnQue(outLocal); + } + + if (this->blockIdx != 0) { + this->lastCoreExpertId = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2); + this->lastCoreExpertIdNum = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2 + 1); + for (int64_t i = this->blockIdx - 2; i >= 0; i--) { + int32_t lastExpertIdx = expertIdxValueGm.GetValue(i * 2); + if (lastExpertIdx < this->lastCoreExpertId) { + break; + } + int32_t lastExpertNum = expertIdxValueGm.GetValue(i * 2 + 1); + this->lastCoreExpertIdNum += lastExpertNum; + } + } +} + +template +__aicore__ inline void MoeV2SrcToDstWithCapacity::CopyIn(int64_t progress) { + LocalTensor inLocal = copyInQueue.AllocTensor(); + int64_t length = Align(currentLoopRows, sizeof(int32_t)); + DataCopy(inLocal, expandDstToSrcRowGm[progress * perLoopRows], length); + DataCopy(inLocal[length], expandedExpertIdxGm[progress * perLoopRows], length); + copyInQueue.EnQue(inLocal); +} + +template +__aicore__ inline void MoeV2SrcToDstWithCapacity::CopyOut(int64_t progress) { + LocalTensor inLocal = copyInQueue.DeQue(); + LocalTensor outLocal = copyOutQueue.AllocTensor(); + int64_t length = Align(currentLoopRows, sizeof(int32_t)); + DataCopyExtParams copyParams{static_cast(1), static_cast(sizeof(int32_t)), 0, 0, 0}; + + SetWaitFlag(HardEvent::MTE2_S); + if (this->lastExpertId == -1) { + this->lastExpertId = this->lastCoreExpertId; + this->tokenCount = this->lastCoreExpertIdNum; + } + for (int64_t idx = 0; idx < currentLoopRows; idx++) { + int32_t expertIdx = inLocal[length].GetValue(idx); + SetWaitFlag(HardEvent::S_MTE3); + int32_t index = 0; + while (this->lastExpertId < expertIdx) { + while (this->tokenCount < this->expertCapacity) { + index = this->lastExpertId * this->expertCapacity + this->tokenCount; + int64_t col = this->perLoopCols; + for (int64_t i = 0; i < this->colLoops; i++) { + if (i == this->colLoops - 1) { + col = this->lastLoopCols; + } +#ifdef __CCE_KT_TEST__ + // CPU twin debugging cannot use multi-core sync, so index may contain uninitialized dirty data; handle specially + if (index * this->cols + i * this->perLoopCols + col * sizeof(T) > expandedXGm.GetSize()) { + continue; + } +#endif + DataCopyExtParams copyParams1{static_cast(1), static_cast(col * sizeof(T)), 0, 0, 0}; + DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, copyParams1); + SetWaitFlag(HardEvent::MTE3_S); + } + this->tokenCount++; + } + this->tokenCount = 0; + this->lastExpertId++; + } + + if (this->tokenCount < this->expertCapacity) { + int32_t outOffset = inLocal.GetValue(idx); + index = expertIdx * this->expertCapacity + this->tokenCount; + outLocal.SetValue(0, index); + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(expandedRowIdxGm[outOffset], outLocal, copyParams); + SetWaitFlag(HardEvent::MTE3_S); + this->tokenCount++; + } + } + copyInQueue.FreeTensor(inLocal); + copyOutQueue.FreeTensor(outLocal); +} + +template +__aicore__ inline void MoeV2SrcToDstWithCapacity::CopyOutRemain() { + if (this->blockIdx != this->srcToDstTilingData->needCoreNum - 1) { + copyOutZeroQueue.FreeTensor(this->outTmpLocal); + return; + } + while (this->lastExpertId < this->expertNum) { + while (this->tokenCount < this->expertCapacity) { + int32_t index = this->lastExpertId * this->expertCapacity + this->tokenCount; + int64_t col = this->perLoopCols; + for (int64_t i = 0; i < this->colLoops; i++) { + if (i == this->colLoops - 1) { + col = this->lastLoopCols; + } + DataCopyExtParams copyParams{static_cast(1), static_cast(col * sizeof(T)), 0, 0, 0}; + DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, copyParams); + SetWaitFlag(HardEvent::MTE3_S); + } + this->tokenCount++; + } + this->tokenCount = 0; + this->lastExpertId++; + } + copyOutZeroQueue.FreeTensor(this->outTmpLocal); +} + +template +__aicore__ inline void MoeV2SrcToDstWithCapacity::SyncAll() { + if (coreNum == 1) { + return; + } +#ifndef __CCE_KT_TEST__ + AscendC::SyncAll(); +#endif +} + +template +__aicore__ inline void MoeV2SrcToDstWithCapacity::Init(GM_ADDR expandedRowIdx, GM_ADDR expandedX, + GM_ADDR workspace, const TilingData* tilingData, + TPipe* tPipe) { + int64_t blockNum = GetBlockNum(); + this->pipe = tPipe; + this->blockIdx = get_block_idx() + get_subblockid() * get_block_num(); + + this->coreNum = tilingData->coreNum; + this->totalLength = tilingData->n * tilingData->k; + this->srcToDstTilingData = &(tilingData->srcToDstCapacityComputeParamsOp); + this->expertNum = tilingData->expertNum; + this->expertCapacity = tilingData->expertCapacity; + this->cols = tilingData->cols; + + if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) { + this->coreRows = this->srcToDstTilingData->lastCoreRows; + this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows; + this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows; + this->rowLoops = this->srcToDstTilingData->lastCoreLoops; + } else { + this->coreRows = this->srcToDstTilingData->perCoreRows; + this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows; + this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows; + this->rowLoops = this->srcToDstTilingData->perCoreLoops; + } + this->perLoopCols = this->srcToDstTilingData->perLoopCols; + this->lastLoopCols = this->srcToDstTilingData->lastLoopCols; + this->colLoops = this->srcToDstTilingData->colLoops; + + int64_t length = Align(this->totalLength, sizeof(int32_t)); + expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t*)expandedRowIdx, length); + expandedXGm.SetGlobalBuffer((__gm__ T*)expandedX, this->expertNum * this->expertCapacity * this->cols); + + expandedExpertIdxGm.SetGlobalBuffer( + (__gm__ int32_t*)workspace + this->blockIdx * this->srcToDstTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + expandDstToSrcRowGm.SetGlobalBuffer( + (__gm__ int32_t*)workspace + length + this->blockIdx * this->srcToDstTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + expertIdxValueGm.SetGlobalBuffer((__gm__ int32_t*)workspace + length * 2, this->coreNum * 2); + + pipe->InitBuffer(copyInQueue, 1, AlignBytes(this->perLoopRows, sizeof(int32_t)) * 2); + pipe->InitBuffer(copyOutQueue, 1, AlignBytes(INT32_ONE_BLOCK_NUM, sizeof(int32_t))); + if constexpr (IsSameType::value) { + pipe->InitBuffer(copyOutZeroQueue, 1, AlignBytes(this->perLoopCols, sizeof(int16_t))); + } else { + pipe->InitBuffer(copyOutZeroQueue, 1, AlignBytes(this->perLoopCols, sizeof(T))); + } +} + +template +__aicore__ inline void MoeV2SrcToDstWithCapacity::Process() { + if (this->blockIdx < this->srcToDstTilingData->needCoreNum) { + AssistInit(); + this->outTmpLocal = copyOutZeroQueue.DeQue(); + currentLoopRows = perLoopRows; + for (int64_t loop = 0; loop < this->rowLoops; loop++) { + if (loop == this->rowLoops - 1) { + currentLoopRows = lastLoopRows; + } + CopyIn(loop); + CopyOut(loop); + } + CopyOutRemain(); + } + this->SyncAll(); +} +} // namespace MoeInitRoutingQuantV2 +#endif // INNER_MOE_V2_SRC_TO_DST_WITH_CAPACITY_H diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/tiling_base.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/tiling_base.h new file mode 100644 index 000000000..f60cedab6 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/tiling_base.h @@ -0,0 +1,69 @@ +#pragma once +#ifndef ASCENDC_DISPATCH_FFN_COMBINE_TILING_BASE_H +#define ASCENDC_DISPATCH_FFN_COMBINE_TILING_BASE_H +namespace optiling { +struct AiCoreParams { + uint64_t ubSize; + uint64_t blockDim; + uint64_t aicNum; + + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; +}; + +class TilingBaseClass { +public: + bool DoTiling( + int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, + int64_t expertNum, int64_t activeNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag, + bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0, + int64_t aivCoreNum, int64_t ubSizePlatForm) + { + bool ret = GetShapeAttrsInfo(m, cols, topK, expertCapacity, expertNum, activeNum, dropPadMode, expertTokensCountOrCumsumFlag, + expertTokensBeforeCapacityFlag, inuptXDtypeSize, quantMode, scaleDim0); + + if (!ret){ + return ret; + } + ret = GetPlatformInfo(aivCoreNum, ubSizePlatForm); + if (!ret){ + return ret; + } + ret = DoOpTiling(); + if (!ret){ + return ret; + } + ret = GetWorkspaceSize(); + if (!ret){ + return ret; + } + ret = PostTiling(); + if (!ret){ + return ret; + } + tilingKey_ = GetTilingKey(); + + return true; + } + +//protected: + virtual bool GetPlatformInfo(int64_t aivCoreNum, int64_t ubSizePlatForm) = 0; + virtual bool GetShapeAttrsInfo(int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, + int64_t expertNum, int64_t activeNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag, + bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0) = 0; + + virtual bool DoOpTiling() = 0; + virtual bool GetWorkspaceSize() = 0; + virtual bool PostTiling() = 0; + virtual uint64_t GetTilingKey() const = 0; +//protected: + uint32_t blockDim_{0}; + uint64_t workspaceSize_{0}; + uint64_t tilingKey_{0}; + AiCoreParams aicoreParams_{0, 0, 0, 0, 0, 0, 0}; +}; + +} +#endif \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/unpermute/moe_token_unpermute.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/unpermute/moe_token_unpermute.h new file mode 100644 index 000000000..d03d50fb7 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/unpermute/moe_token_unpermute.h @@ -0,0 +1,375 @@ +/** + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/*! + * \file moe_token_unpermute.h + * \brief + */ + +#ifndef MOE_TOKEN_UNPERMUTE +#define MOE_TOKEN_UNPERMUTE + +#include "kernel_operator.h" +#include "moe_token_unpermute_tiling.h" +using namespace AscendC; + + +template class KernelMoeTokenUnpermute { +public: + __aicore__ inline KernelMoeTokenUnpermute() + { + } + + __aicore__ inline void Init(GM_ADDR permuted_tokens, GM_ADDR sorted_indices, GM_ADDR probs, + GM_ADDR unpermuted_tokens, const MoeTokenUnpermuteTilingData *__restrict tiling_data); + __aicore__ inline void Process(); + +protected: + __aicore__ inline void CalMultiOutToken(const int64_t out_offset, const int64_t out_tokens_number); + __aicore__ inline void CalSingleOutToken(const int64_t start_token, const int64_t out_token_idx); + __aicore__ inline void CalPartOutToken(const int64_t start_token, const int64_t h_index, const int64_t h_length, + const int64_t out_token_index); + __aicore__ inline void CopyTokenIn(const T2 in_token_index, const int64_t h_index, const int64_t h_length); + __aicore__ inline void CalFirstToken(const float prob_value, const int64_t h_length); + __aicore__ inline void CalToken(const float prob_value, const int64_t h_length); + __aicore__ inline void CopyOut(const int64_t out_token_index, const int64_t h_index, const int64_t h_length); + + TPipe pipe; + TQue tokens_inque, indices_inque, probs_inque; + TBuf temp_buffer0, temp_buffer1, temp_buffer2; + TQue outque; + GlobalTensor tokensGM, outGM; + GlobalTensor indicesGM; + GlobalTensor probsGM; + LocalTensor indicesLocal; + LocalTensor token_tensor0, token_tensor1, probs_tensor; + DataCopyPadExtParams extParams1{false, 0, 0, 0}; + DataCopyPadExtParams extParams2{false, 0, 0, 0}; + DataCopyPadExtParams extParams3{false, 0, 0, 0}; + DataCopyExtParams copyParams{1, 0, 0, 0, 0}; + + constexpr static uint32_t BLOCK_SIZE = 32; + constexpr static uint32_t ALIGN_512 = 512; + + int64_t hidden_size; + int64_t top_k; + int64_t num_out_tokens; + int64_t hidden_splited_length; + int64_t hidden_splited_num; + int64_t hidden_splited_remain; + int64_t tokens_core_length; + int64_t tokens_core_remain; + int64_t tokens_splited_length; + int64_t tokens_splited_num; + int64_t tokens_splited_remain; + int32_t blockIdx; + int32_t blockNum; +}; + +template +__aicore__ inline void +KernelMoeTokenUnpermute::Init(GM_ADDR permuted_tokens, GM_ADDR sorted_indices, GM_ADDR probs, + GM_ADDR unpermuted_tokens, + const MoeTokenUnpermuteTilingData *__restrict tiling_data) +{ + this->blockIdx = get_block_idx() + get_subblockid() * get_block_num(); + this->blockNum = get_block_num() * get_subblockdim(); + if (blockIdx >= blockNum) { + return; + } + ASSERT(blockNum != 0 && "block dim can not be zero!"); + // row_input + this->hidden_size = tiling_data->hidden_size; + this->top_k = tiling_data->top_k; + this->num_out_tokens = tiling_data->num_out_tokens; + // hidden_tiling + this->hidden_splited_length = tiling_data->hidden_splited_length; + this->hidden_splited_num = tiling_data->hidden_splited_num; + this->hidden_splited_remain = tiling_data->hidden_splited_remain; + // token_tiling + this->tokens_core_length = tiling_data->tokens_core_length; + this->tokens_core_remain = tiling_data->tokens_core_remain; + this->tokens_splited_length = tiling_data->tokens_splited_length; + this->tokens_splited_num = tiling_data->tokens_splited_num; + this->tokens_splited_remain = tiling_data->tokens_splited_remain; + + // Handle the tail block for token_by_core + if (this->tokens_core_remain > 0 && blockIdx < this->tokens_core_remain) { + this->tokens_core_length += 1; + this->tokens_splited_remain += 1; + } + + int64_t hidden_splited_length_align512 = (this->hidden_splited_length + ALIGN_512 - 1) & ~(ALIGN_512 - 1); + + int64_t block_length = this->tokens_core_length * this->top_k; + int64_t block_splited_length = this->tokens_splited_length * this->top_k; + + int64_t block_offset; + if (this->tokens_core_remain > 0) { + if (blockIdx < this->tokens_core_remain) { + block_offset = block_length * blockIdx; + } else { + block_offset = (block_length + this->top_k) * this->tokens_core_remain + + block_length * (blockIdx - this->tokens_core_remain); + } + } else { + block_offset = block_length * blockIdx; + } + + this->tokensGM.SetGlobalBuffer((__gm__ T1 *)permuted_tokens); + this->indicesGM.SetGlobalBuffer((__gm__ T2 *)sorted_indices + block_offset, block_length); + + + int64_t out_block_offset; + if (this->tokens_core_remain > 0) { + if (blockIdx < this->tokens_core_remain) { + out_block_offset = this->tokens_core_length * blockIdx * hidden_size; + } else { + out_block_offset = (this->tokens_core_length + 1) * this->tokens_core_remain + + this->tokens_core_length * (blockIdx - this->tokens_core_remain); + out_block_offset *= this->hidden_size; + } + } else { + out_block_offset = this->tokens_core_length * blockIdx * hidden_size; + } + + this->outGM.SetGlobalBuffer((__gm__ T1 *)unpermuted_tokens + out_block_offset, + this->tokens_core_length * this->hidden_size); + + this->pipe.InitBuffer(tokens_inque, tiling_data->buffer_num, hidden_splited_length_align512 * sizeof(T1)); + this->pipe.InitBuffer(indices_inque, 1, block_splited_length * (sizeof(T2))); + this->pipe.InitBuffer(outque, 1, hidden_splited_length_align512 * sizeof(T1)); + + if constexpr (!IsSameType::value) { + this->pipe.InitBuffer(temp_buffer0, hidden_splited_length_align512 * sizeof(float) + 256); + this->pipe.InitBuffer(temp_buffer1, hidden_splited_length_align512 * sizeof(float)); + this->token_tensor0 = this->temp_buffer0.template Get(); + this->token_tensor1 = this->temp_buffer1.template Get(); + } + + if constexpr (PROBS) { + this->probsGM.SetGlobalBuffer((__gm__ T3 *)probs + block_offset, block_length); + this->pipe.InitBuffer(probs_inque, 1, block_splited_length * (sizeof(T3))); + if constexpr (!IsSameType::value) { + this->pipe.InitBuffer(temp_buffer2, block_splited_length * sizeof(float)); + this->probs_tensor = this->temp_buffer2.template Get(); + } + } +}; + +template +__aicore__ inline void KernelMoeTokenUnpermute::Process() +{ + + if (blockIdx >= blockNum) { + return; + } + for (int64_t i = 0; i < this->tokens_splited_num; ++i) { + CalMultiOutToken(i * this->tokens_splited_length, this->tokens_splited_length); + } + // Handle the tail block when tokens_num is not evenly divisible by core count + if (this->tokens_splited_remain > 0) { + CalMultiOutToken(this->tokens_splited_num * this->tokens_splited_length, this->tokens_splited_remain); + } +} + +template +__aicore__ inline void KernelMoeTokenUnpermute::CalMultiOutToken(const int64_t out_offset, + const int64_t out_tokens_number) +{ + this->indicesLocal = this->indices_inque.template AllocTensor(); + int64_t in_offset = out_offset * this->top_k; + this->copyParams.blockLen = out_tokens_number * this->top_k * sizeof(T2); + DataCopyPad(this->indicesLocal, this->indicesGM[in_offset], this->copyParams, this->extParams2); + this->indices_inque.template EnQue(this->indicesLocal); + + if constexpr (PROBS) { + LocalTensor temp_probs_tensor = this->probs_inque.template AllocTensor(); + this->copyParams.blockLen = out_tokens_number * this->top_k * sizeof(T3); + DataCopyPad(temp_probs_tensor, this->probsGM[in_offset], this->copyParams, this->extParams3); + this->probs_inque.template EnQue(temp_probs_tensor); + temp_probs_tensor = this->probs_inque.template DeQue(); + if constexpr (!IsSameType::value) { + Cast(this->probs_tensor, temp_probs_tensor, RoundMode::CAST_NONE, out_tokens_number * this->top_k); + this->probs_inque.FreeTensor(temp_probs_tensor); + PipeBarrier(); + } else { + this->probs_tensor = temp_probs_tensor; + } + } + this->indicesLocal = this->indices_inque.template DeQue(); + + + for (int64_t out_token_idx = 0; out_token_idx < out_tokens_number; ++out_token_idx) { + CalSingleOutToken(out_token_idx * this->top_k, out_offset + out_token_idx); + } + // Free Tensor + this->indices_inque.FreeTensor(this->indicesLocal); + if constexpr (PROBS && IsSameType::value) { + this->probs_inque.FreeTensor(this->probs_tensor); + } +} + +template +__aicore__ inline void KernelMoeTokenUnpermute::CalSingleOutToken(const int64_t start_token, + const int64_t out_token_idx) +{ + for (int64_t h_index = 0; h_index < this->hidden_splited_num; ++h_index) { + CalPartOutToken(start_token, h_index, this->hidden_splited_length, out_token_idx); + } + // Handle the tail block when a full hidden_size does not fit in one pass + if (this->hidden_splited_remain > 0) { + CalPartOutToken(start_token, this->hidden_splited_num, this->hidden_splited_remain, out_token_idx); + } +} + +template +__aicore__ inline void +KernelMoeTokenUnpermute::CalPartOutToken(const int64_t start_token, const int64_t h_index, + const int64_t h_length, const int64_t out_token_index) +{ + if constexpr (IsSameType::value) { + this->token_tensor0 = this->outque.template AllocTensor(); + } + int64_t end_token = start_token + this->top_k; + T2 cal_token_idx = this->indicesLocal.GetValue(start_token); + + // Handle the first token + if (cal_token_idx < this->num_out_tokens) { + float probsValue = 0; + if constexpr (PROBS) { + probsValue = this->probs_tensor.GetValue(start_token); + } + + CopyTokenIn(cal_token_idx, h_index, h_length); + PipeBarrier(); + CalFirstToken(probsValue, h_length); + } else { + PipeBarrier(); + Duplicate(this->token_tensor0, static_cast(0), h_length); + } + + // Handle the remaining tokens + for (int64_t token_index = start_token + 1; token_index < end_token; ++token_index) { + cal_token_idx = this->indicesLocal.GetValue(token_index); + if (cal_token_idx < this->num_out_tokens) { + float probsValue = 0; + if constexpr (PROBS) { + probsValue = this->probs_tensor.GetValue(token_index); + } + + CopyTokenIn(cal_token_idx, h_index, h_length); + PipeBarrier(); + CalToken(probsValue, h_length); + } + } + + // Write out the computed result + CopyOut(out_token_index, h_index, h_length); +} + +template +__aicore__ inline void KernelMoeTokenUnpermute::CopyTokenIn(const T2 in_token_index, + const int64_t h_index, + const int64_t h_length) +{ + LocalTensor tokensLocal = this->tokens_inque.template AllocTensor(); + int64_t offset = in_token_index * this->hidden_size + h_index * this->hidden_splited_length; + + if (likely((h_length * sizeof(T1)) % BLOCK_SIZE == 0)) { + DataCopy(tokensLocal, this->tokensGM[offset], h_length); + } else { + this->copyParams.blockLen = h_length * sizeof(T1); + DataCopyPad(tokensLocal, this->tokensGM[offset], this->copyParams, this->extParams1); + } + + this->tokens_inque.template EnQue(tokensLocal); +} + +template +__aicore__ inline void KernelMoeTokenUnpermute::CalFirstToken(const float prob_value, + const int64_t h_length) +{ + LocalTensor tokensLocal = this->tokens_inque.template DeQue(); + + if constexpr (!IsSameType::value) { + Cast(this->token_tensor0, tokensLocal, RoundMode::CAST_NONE, h_length); + } else { + uint64_t byteAlign32 = (h_length * sizeof(float) + BLOCK_SIZE - 1) & ~(BLOCK_SIZE - 1); + DataCopy(this->token_tensor0, tokensLocal, byteAlign32 / sizeof(float)); + } + + this->tokens_inque.FreeTensor(tokensLocal); + + if constexpr (PROBS) { + PipeBarrier(); + Muls(this->token_tensor0, this->token_tensor0, prob_value, h_length); + } +} + +template +__aicore__ inline void KernelMoeTokenUnpermute::CalToken(const float prob_value, + const int64_t h_length) +{ + LocalTensor tokensLocal = this->tokens_inque.template DeQue(); + + if constexpr (!IsSameType::value) { + Cast(this->token_tensor1, tokensLocal, RoundMode::CAST_NONE, h_length); + this->tokens_inque.FreeTensor(tokensLocal); + if constexpr (PROBS) { + PipeBarrier(); + Muls(this->token_tensor1, this->token_tensor1, prob_value, h_length); + } + PipeBarrier(); + Add(this->token_tensor0, this->token_tensor0, this->token_tensor1, h_length); + } else { + if constexpr (PROBS) { + Muls(tokensLocal, tokensLocal, prob_value, h_length); + PipeBarrier(); + } + Add(this->token_tensor0, this->token_tensor0, tokensLocal, h_length); + this->tokens_inque.FreeTensor(tokensLocal); + } +} + +template +__aicore__ inline void KernelMoeTokenUnpermute::CopyOut(const int64_t out_token_index, + const int64_t h_index, + const int64_t h_length) +{ + LocalTensor temp_out_tensors; + if constexpr (!IsSameType::value) { + temp_out_tensors = this->outque.template AllocTensor(); + PipeBarrier(); + Cast(temp_out_tensors, this->token_tensor0, RoundMode::CAST_RINT, h_length); + } else { + temp_out_tensors = this->token_tensor0; + } + + this->outque.template EnQue(temp_out_tensors); + temp_out_tensors = this->outque.template DeQue(); + + int64_t offset = out_token_index * this->hidden_size + h_index * this->hidden_splited_length; + if (likely((h_length * sizeof(T1)) % BLOCK_SIZE == 0)) { + DataCopy(this->outGM[offset], temp_out_tensors, h_length); + } else { + this->copyParams.blockLen = h_length * sizeof(T1); + DataCopyPad(this->outGM[offset], temp_out_tensors, this->copyParams); + } + + this->outque.FreeTensor(temp_out_tensors); +} +#endif // MOE_TOKEN_UNPERMUTE \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/unpermute/moe_token_unpermute_tiling.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/unpermute/moe_token_unpermute_tiling.h new file mode 100644 index 000000000..df47f6db8 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/unpermute/moe_token_unpermute_tiling.h @@ -0,0 +1,38 @@ +#ifndef MOE_TOKEN_UNPERMUTE_TILING +#define MOE_TOKEN_UNPERMUTE_TILING + +struct MoeTokenUnpermuteTilingData { + int64_t hidden_size; + int64_t top_k; + int64_t num_out_tokens; + int64_t hidden_splited_length; + int64_t hidden_splited_num; + int64_t hidden_splited_remain; + int64_t tokens_core_length; + int64_t tokens_core_remain; + int64_t tokens_splited_length; + int64_t tokens_splited_num; + int64_t tokens_splited_remain; + int64_t buffer_num; +}; + +__forceinline__ [host, aicore] void +MoeTokenUnpermuteTiling(int32_t m, int32_t n, int32_t topK, MoeTokenUnpermuteTilingData &tilingData, uint32_t coreNum) +{ + #define I64(x) static_cast(x) + tilingData.hidden_size = I64(n); + tilingData.top_k = I64(topK); + tilingData.num_out_tokens = I64(m); + tilingData.hidden_splited_length = tilingData.hidden_size; + tilingData.hidden_splited_num = 1; + tilingData.hidden_splited_remain = 0; + uint32_t outTokens = m / topK; + tilingData.tokens_core_length = I64(outTokens / coreNum); + tilingData.tokens_core_remain = I64(outTokens % coreNum); + tilingData.tokens_splited_length = I64(min(tilingData.tokens_core_length, 600)); + tilingData.tokens_splited_num = I64(tilingData.tokens_core_length / tilingData.tokens_splited_length); + tilingData.tokens_splited_remain = I64(tilingData.tokens_core_length % tilingData.tokens_splited_length); + tilingData.buffer_num = 4; +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/block_epilogue_pertoken_row.hpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/block_epilogue_pertoken_row.hpp new file mode 100644 index 000000000..41bee5057 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/block_epilogue_pertoken_row.hpp @@ -0,0 +1,216 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_ROW_HPP +#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_ROW_HPP + +#include "catlass/catlass.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/epilogue/dispatch_policy.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/matrix_coord.hpp" +#include "catlass/layout/layout.hpp" +#include "catlass/detail/callback.hpp" +#include "catlass/epilogue/block/block_epilogue.hpp" + +namespace Catlass::Epilogue::Block { + +// float scale, dequant per expert +template < + uint32_t UB_STAGES_, + class CType_, + class LayoutPerTokenScale_, + class DType_, + class TileCopy_ +> +class BlockEpilogue < + EpilogueAtlasA2PerTokenDequant, + CType_, + Gemm::GemmType, + DType_, + TileCopy_ +> { +public: + using DispatchPolicy = EpilogueAtlasA2PerTokenDequant; + using ArchTag = typename DispatchPolicy::ArchTag; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + + // Data infos + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using ElementPerTokenScale = float; + using LayoutPerTokenScale = LayoutPerTokenScale_; + using ElementD = typename DType_::Element; + using LayoutD = typename DType_::Layout; + + // Check data infos + static_assert( + std::is_same_v && (std::is_same_v || std::is_same_v), + "The element type template parameters of BlockEpilogue are wrong" + ); + static_assert( + std::is_same_v && + std::is_same_v && std::is_same_v, + "The layout template parameters of BlockEpilogue are wrong" + ); + + + // Tile copy + using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; + using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; + + struct Params { + __gm__ int32_t *ptrTokenPerExpert{nullptr}; + int32_t EP; + int32_t expertPerRank; + int32_t n2; + + CATLASS_DEVICE + Params() {}; + + CATLASS_DEVICE + Params(int32_t EP_, int32_t expertPerRank_, __gm__ int32_t *ptrTokenPerExpert_, int32_t n2_) : ptrTokenPerExpert(ptrTokenPerExpert_), EP(EP_), expertPerRank(expertPerRank_), n2(n2_) {} + }; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}) : params(params) + { + size_t ubOffset = 0; + int32_t eventVMTE2 = 0; + int32_t eventMTE2V = 0; + int32_t eventMTE3V = 0; + int32_t eventVMTE3 = 0; + int32_t blockN = params.n2; + for (uint32_t i = 0; i < UB_STAGES; ++i) { + ubCList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += blockN * sizeof(ElementC); + ubDList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += blockN * sizeof(ElementD); + + eventUbCVMTE2List[i] = eventVMTE2++; + eventUbCMTE2VList[i] = eventMTE2V++; + eventUbDMTE3VList[i] = eventMTE3V++; + eventUbDVMTE3List[i] = eventVMTE3++; + + + ubCFp32List[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += blockN * sizeof(float); + } + } + CATLASS_DEVICE + void SetFlag() + { + for (uint32_t i = 0; i < UB_STAGES; ++i) { + AscendC::SetFlag(eventUbCVMTE2List[i]); + AscendC::SetFlag(eventUbDMTE3VList[i]); + } + } + + CATLASS_DEVICE + void Finalize() + { + for (uint32_t i = 0; i < UB_STAGES; ++i) { + AscendC::WaitFlag(eventUbCVMTE2List[i]); + AscendC::WaitFlag(eventUbDMTE3VList[i]); + } + } + CATLASS_DEVICE + ~BlockEpilogue() + { + + } + + CATLASS_DEVICE + void UpdateParams(Params const ¶ms_) + { + params = params_; + } + + CATLASS_DEVICE + void operator() ( + AscendC::GlobalTensor const &gmC, + MatrixCoord const &shapeC, + AscendC::GlobalTensor const &gmPerTokenScale, + AscendC::GlobalTensor const &gmD + ) + { + uint32_t blockM = shapeC.row(); + uint32_t blockN = shapeC.column(); + + uint32_t tileLoops = blockM; + + for (uint32_t loopIdx = 0; loopIdx < tileLoops; loopIdx ++) { + auto gmTileC = gmC[loopIdx * blockN]; + auto &ubC = ubCList[ubListId]; + auto &ubCFp32 = ubCFp32List[ubListId]; + auto &ubMul = ubMulList[ubListId]; + auto &ubD = ubDList[ubListId]; + auto gmTileD = gmD[loopIdx * blockN]; + LayoutC layoutUbC{1, blockN}; + + // Move C from GM workspace to UB + AscendC::WaitFlag(eventUbCVMTE2List[ubListId]); + copyGmToUbC(ubC, gmTileC, layoutUbC, layoutUbC); + AscendC::SetFlag(eventUbCMTE2VList[ubListId]); + + // Cast C to FP32 in UB + AscendC::WaitFlag(eventUbCMTE2VList[ubListId]); + AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_NONE, blockN); + AscendC::SetFlag(eventUbCVMTE2List[ubListId]); + + // Get per-token scale from row loopIdx of gmPerTokenScale + ElementPerTokenScale perTokenScale = gmPerTokenScale(loopIdx); + + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + // Multiply FP32 C by the per-token scale + AscendC::PipeBarrier(); + AscendC::Muls(ubCFp32, ubCFp32, perTokenScale, blockN); + AscendC::PipeBarrier(); + + // Cast the muls result back to fp16/bf16 + LayoutD layoutUbD{1, blockN}; + AscendC::WaitFlag(eventUbDMTE3VList[ubListId]); + + AscendC::Cast(ubD, ubCFp32, AscendC::RoundMode::CAST_RINT, blockN); + AscendC::SetFlag(eventUbDVMTE3List[ubListId]); + + AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); + copyUbToGmD(gmTileD, ubD, layoutUbD, layoutUbD); + AscendC::SetFlag(eventUbDMTE3VList[ubListId]); + + ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; + } + } + +private: + Params params; + + AscendC::LocalTensor ubCList[UB_STAGES]; + AscendC::LocalTensor ubDList[UB_STAGES]; + + int32_t eventUbCVMTE2List[UB_STAGES]; + int32_t eventUbCMTE2VList[UB_STAGES]; + int32_t eventUbDMTE3VList[UB_STAGES]; + int32_t eventUbDVMTE3List[UB_STAGES]; + + uint32_t ubListId{0}; + + AscendC::LocalTensor ubCFp32List[UB_STAGES]; + AscendC::LocalTensor ubMulList[UB_STAGES]; + + + CopyGmToUbC copyGmToUbC; + CopyUbToGmD copyUbToGmD; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_ROW_HPP \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/block_epilogue_pertoken_swiglu.hpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/block_epilogue_pertoken_swiglu.hpp new file mode 100644 index 000000000..4e47ea33e --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/block_epilogue_pertoken_swiglu.hpp @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_SWIGLU_HPP +#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_SWIGLU_HPP + +#include "catlass/catlass.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/epilogue/dispatch_policy.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/matrix_coord.hpp" +#include "catlass/layout/layout.hpp" +#include "catlass/detail/callback.hpp" + +namespace Catlass::Epilogue::Block { + +// float scale, dequant per expert +template < + uint32_t UB_STAGES_, + class CType_, + class LayoutPerTokenScale_, + class DType_, + class TileElemWiseMuls_, + class TileCopy_ +> +class BlockEpilogue < + EpilogueAtlasA2PerTokenDequantSwigluQuant, + CType_, + Gemm::GemmType, + DType_, + TileElemWiseMuls_, + TileCopy_ +> { +public: + using DispatchPolicy = EpilogueAtlasA2PerTokenDequantSwigluQuant; + using ArchTag = typename DispatchPolicy::ArchTag; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + + // Data infos + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using ElementPerTokenScale = float; + using LayoutPerTokenScale = LayoutPerTokenScale_; + using ElementD = typename DType_::Element; + using LayoutD = typename DType_::Layout; + + // Check data infos + static_assert( + std::is_same_v && (std::is_same_v || std::is_same_v), + "The element type template parameters of BlockEpilogue are wrong" + ); + static_assert( + std::is_same_v && + std::is_same_v && std::is_same_v, + "The layout template parameters of BlockEpilogue are wrong" + ); + + // Tile copy + using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; + using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; + using CopyUbToGmDequantScale = Epilogue::Tile::CopyUb2Gm>; + + struct Params { + __gm__ ElementPerTokenScale *ptrPerTokenScale{nullptr}; + LayoutPerTokenScale layoutPerTokenScale{}; + __gm__ ElementD *ptrD{nullptr}; + LayoutD layoutD{}; + + CATLASS_DEVICE + Params() {}; + + CATLASS_DEVICE + Params(__gm__ ElementPerTokenScale *ptrPerTokenScale_, LayoutPerTokenScale const &layoutPerTokenScale_, + __gm__ ElementD *ptrD_, LayoutD const &layoutD_ + ) : ptrPerTokenScale(ptrPerTokenScale_), layoutPerTokenScale(layoutPerTokenScale_), + ptrD(ptrD_), layoutD(layoutD_) {} + }; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource const &resource, int32_t n, Params const ¶ms = Params{}) : params(params) + { + int32_t eventVMTE2 = 0; + int32_t eventMTE2V = 0; + int32_t eventMTE3V = 0; + int32_t eventVMTE3 = 0; + uint32_t blockN = n; + uint32_t ChunkTileLen = blockN / 2; + uint32_t HalfChunkTileLen = ChunkTileLen / 2; + + for (uint32_t i = 0; i < UB_STAGES; ++i) { + ubCList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += blockN * sizeof(ElementC); + ubDList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += blockN * sizeof(ElementD); + ubCFp32List[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += blockN * sizeof(float); + ubCFp32ChunkNList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += ChunkTileLen * sizeof(float); + ubCFp32ChunkNAbsList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += ChunkTileLen * sizeof(float); + ubCFp32ChunkNMaxList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += HalfChunkTileLen * sizeof(float); + ubQuantS32List[i] = ubCFp32ChunkNAbsList[i].template ReinterpretCast(); + ubQuantF16List[i] = ubCFp32ChunkNAbsList[i].template ReinterpretCast(); + + eventUbCVMTE2List[i] = eventVMTE2++; + eventUbCMTE2VList[i] = eventMTE2V++; + eventUbDMTE3VList[i] = eventMTE3V++; + eventUbDVMTE3List[i] = eventVMTE3++; + + AscendC::SetFlag(eventUbCVMTE2List[i]); + AscendC::SetFlag(eventUbDMTE3VList[i]); + } + ubPerTokenScaleOutput = resource.ubBuf.template GetBufferByByte(ubOffset); + } + CATLASS_DEVICE + void Finalize() + { + for (uint32_t i = 0; i < UB_STAGES; ++i) { + AscendC::WaitFlag(eventUbCVMTE2List[i]); + AscendC::WaitFlag(eventUbDMTE3VList[i]); + } + } + CATLASS_DEVICE + ~BlockEpilogue() + { + } + + CATLASS_DEVICE + void UpdateParams(Params const ¶ms_) + { + params = params_; + } + // Each tile is 1x7168, and each block covers all tokens for one expert = [group[i], 7168] + CATLASS_DEVICE + void operator() ( + AscendC::GlobalTensor const &gmC, + MatrixCoord const &shapeC, + AscendC::GlobalTensor const &gmPerTokenScale1, + AscendC::GlobalTensor const &gmD, + AscendC::GlobalTensor const &gmPerTokenScale2, + Arch::Resource const &resource, + + uint32_t epilogueCoreNum = 40, + float swigluLimit = 0.0f, + uint32_t blockK = 1, + Callback &&callback = Callback{} + ) + { + callback(); + uint32_t blockM = shapeC.row(); + uint32_t blockN = shapeC.column(); + + uint32_t tileLoops = blockM; + uint32_t subblockIdx = get_block_idx() + get_subblockid() * get_block_num(); + + uint32_t subblockNum = get_block_num() * 2; + uint32_t moveDataCoreNum = subblockNum - epilogueCoreNum; + + if (subblockIdx < moveDataCoreNum) { + return; + } + uint32_t epilogueCoreIdx = subblockIdx - moveDataCoreNum; + + uint32_t perCoreData = blockM / epilogueCoreNum; + uint32_t remainderData = blockM % epilogueCoreNum; + + uint32_t scaleBlock = (perCoreData * sizeof(ElementPerTokenScale) + 31) / 32 * 32; + sharedTmpBuffer = resource.ubBuf.template GetBufferByByte(ubOffset + scaleBlock); + + uint32_t tasksForIdx = epilogueCoreIdx < remainderData ? perCoreData + 1 : perCoreData; + uint32_t loopStartIdx = epilogueCoreIdx * perCoreData + (epilogueCoreIdx < remainderData? epilogueCoreIdx : remainderData); + + uint32_t alignedPerCoreData = RoundUp(perCoreData + 1); + + uint32_t ChunkTileLen = blockN / 2; + uint32_t HalfChunkTileLen = ChunkTileLen / 2; + + + for (uint32_t loopIdx = loopStartIdx; loopIdx < loopStartIdx + tasksForIdx; ++loopIdx) { + + auto gmTileC = gmC[loopIdx * blockK]; + + auto &ubC = ubCList[ubListId]; + auto &ubD = ubDList[ubListId]; + + auto &ubCFp32 = ubCFp32List[ubListId]; + auto &ubCFp32ChunkN = ubCFp32ChunkNList[ubListId]; + auto &ubAbs = ubCFp32ChunkNAbsList[ubListId]; + // auto &ubMax = ubCFp32ChunkNMaxList[ubListId]; + auto &ubReduceMax = ubCFp32ChunkNMaxList[ubListId]; + auto &ubOutputTmp = ubAbs; + auto &sharedUbTmpBuffer = ubReduceMax; + auto &ubQuantS32 = ubQuantS32List[ubListId]; + auto &ubQuantF16 = ubQuantF16List[ubListId]; + + auto gmTileD = gmD[loopIdx * ChunkTileLen]; + LayoutC layoutUbC{1, blockN}; + + // Move C from GM workspace to UB + AscendC::WaitFlag(eventUbCVMTE2List[ubListId]); + copyGmToUbC(ubC, gmTileC, layoutUbC, layoutUbC); + AscendC::SetFlag(eventUbCMTE2VList[ubListId]); + + // Cast C to FP32 in UB + AscendC::WaitFlag(eventUbCMTE2VList[ubListId]); + AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_NONE, blockN); + AscendC::SetFlag(eventUbCVMTE2List[ubListId]); + + // Get per-token scale from row loopIdx of gmPerTokenScale + ElementPerTokenScale perTokenScale = gmPerTokenScale1(loopIdx); + + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + // Multiply FP32 C by the per-token scale + AscendC::PipeBarrier(); + AscendC::Muls(ubCFp32, ubCFp32, perTokenScale, blockN); + AscendC::PipeBarrier(); + + // swiglu limit clamp + if (swigluLimit > 0.0f) { + AscendC::ClampMax(ubCFp32, ubCFp32, sharedTmpBuffer, swigluLimit, blockN); + AscendC::PipeBarrier(); + AscendC::ClampMin(ubCFp32[ChunkTileLen], ubCFp32[ChunkTileLen], sharedTmpBuffer, -1.0f * swigluLimit, ChunkTileLen); + //AscendC::ClampMin(ubCFp32, ubCFp32, sharedTmpBuffer, -1.0f * swigluLimit, ChunkTileLen); + AscendC::PipeBarrier(); + } + + // Swiglu computation process + AscendC::Muls(ubCFp32ChunkN, ubCFp32, -1.0f, ChunkTileLen); + AscendC::PipeBarrier(); + AscendC::Exp(ubCFp32ChunkN, ubCFp32ChunkN, ChunkTileLen); + AscendC::PipeBarrier(); + AscendC::Adds(ubCFp32ChunkN, ubCFp32ChunkN, 1.0f, ChunkTileLen); + AscendC::PipeBarrier(); + // TODO: confirm whether the division impacts subsequent data + AscendC::Div(ubCFp32ChunkN, ubCFp32, ubCFp32ChunkN, ChunkTileLen); + AscendC::PipeBarrier(); + AscendC::Mul(ubCFp32ChunkN, ubCFp32ChunkN, ubCFp32[ChunkTileLen], ChunkTileLen); + + // Quantization process; difference between the two approaches + AscendC::PipeBarrier(); + AscendC::Abs(ubAbs, ubCFp32ChunkN, ChunkTileLen); + AscendC::PipeBarrier(); + + AscendC::ReduceMax(ubReduceMax, ubAbs, sharedUbTmpBuffer, ChunkTileLen, false); + AscendC::PipeBarrier(); + + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + + // TODO: compare the efficiency of the two calculation methods + ElementPerTokenScale GMubDequantScale = ubReduceMax.GetValue(0); + AscendC::SetFlag(0); + + auto ubPerTokenScaleOutputOffset = loopIdx - loopStartIdx; + ubPerTokenScaleOutput.SetValue(ubPerTokenScaleOutputOffset, GMubDequantScale / 127.f); + + AscendC::WaitFlag(0); + AscendC::Muls(ubOutputTmp, ubCFp32ChunkN, 127.f / GMubDequantScale, ChunkTileLen); + AscendC::PipeBarrier(); + + AscendC::Cast(ubQuantS32, ubOutputTmp, AscendC::RoundMode::CAST_RINT, ChunkTileLen); + AscendC::PipeBarrier(); + AscendC::SetDeqScale(static_cast(1.0)); + AscendC::Cast(ubQuantF16, ubQuantS32, AscendC::RoundMode::CAST_RINT, ChunkTileLen); + AscendC::PipeBarrier(); + + AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); + AscendC::Cast(ubD, ubQuantF16, AscendC::RoundMode::CAST_RINT, ChunkTileLen); + // AscendC::Muls(ubD, ubCFp32ChunkN, 127.f / GMubDequantScale, ChunkTileLen); + AscendC::SetFlag(eventUbDMTE3VList[ubListId]); + + LayoutD layoutUbD{1, ChunkTileLen}; + AscendC::WaitFlag(eventUbDVMTE3List[ubListId]); + copyUbToGmD(gmTileD, ubD, layoutUbD, layoutUbD); + AscendC::SetFlag(eventUbDMTE3VList[ubListId]); + ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0; + } + + if(tasksForIdx > 0){ + LayoutPerTokenScale layoutGmPerTokenScale2{tasksForIdx}; + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + copyUbToGmDequantScale(gmPerTokenScale2[loopStartIdx], ubPerTokenScaleOutput[0], layoutGmPerTokenScale2, layoutGmPerTokenScale2); + } + + + } + +private: + Params params; + + AscendC::LocalTensor ubCList[UB_STAGES]; + AscendC::LocalTensor ubDList[UB_STAGES]; + + int32_t eventUbCVMTE2List[UB_STAGES]; + int32_t eventUbCMTE2VList[UB_STAGES]; + int32_t eventUbDMTE3VList[UB_STAGES]; + int32_t eventUbDVMTE3List[UB_STAGES]; + + uint32_t ubListId{0}; + size_t ubOffset = 0; + + AscendC::LocalTensor ubCFp32List[UB_STAGES]; + AscendC::LocalTensor ubCFp32ChunkNList[UB_STAGES]; + AscendC::LocalTensor ubCFp32ChunkNAbsList[UB_STAGES]; + AscendC::LocalTensor ubCFp32ChunkNMaxList[UB_STAGES]; + AscendC::LocalTensor ubQuantS32List[UB_STAGES]; + AscendC::LocalTensor ubQuantF16List[UB_STAGES]; + AscendC::LocalTensor ubPerTokenScaleOutput; + AscendC::LocalTensor sharedTmpBuffer; + + CopyGmToUbC copyGmToUbC; + CopyUbToGmD copyUbToGmD; + CopyUbToGmDequantScale copyUbToGmDequantScale; +}; + +} // namespace Catlass::Epilogue::Block + +#endif // CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_SWIGLU_HPP \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/block_epilogue_pertoken_v2.hpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/block_epilogue_pertoken_v2.hpp new file mode 100644 index 000000000..547c2365d --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/block_epilogue_pertoken_v2.hpp @@ -0,0 +1,255 @@ +#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_V2_ONLY_HPP +#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_V2_ONLY_HPP + +#include "catlass/catlass.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/epilogue/dispatch_policy.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/matrix_coord.hpp" +#include "catlass/layout/layout.hpp" +#include "catlass/detail/callback.hpp" + +#include "hccl_shmem.hpp" +#include "layout3d.hpp" + +namespace Catlass::Epilogue::Block { +template < + uint32_t UB_STAGES_, + class CType_, + class LayoutPerTokenScale_, + class DType_, + class TileCopy_ +> +class BlockEpilogue < + EpilogueAtlasA2PerTokenDequantV2, + CType_, + Gemm::GemmType, + DType_, + TileCopy_ +> { +public: + using DispatchPolicy = EpilogueAtlasA2PerTokenDequantV2; + using ArchTag = typename DispatchPolicy::ArchTag; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + + // Data infos + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using ElementPerTokenScale = float; + using LayoutPerTokenScale = LayoutPerTokenScale_; + using ElementD = typename DType_::Element; + using LayoutD = typename DType_::Layout; + + using CopyScaleGmToUb = Epilogue::Tile::CopyGm2Ub>; + using CopyGmToUbC = typename TileCopy_::CopyGmToUbC; + using CopyUbToGmD = typename TileCopy_::CopyUbToGmD; + + struct Params { + __gm__ int32_t *ptrTokenPerExpert{nullptr}; + int32_t EP; + int32_t expertPerRank; + int32_t n2; + LayoutC layoutC; + int32_t n0; + int32_t rank; + HcclShmem shmem; + int32_t offsetD; + Layout3D tokenPerExpertLayout; + CATLASS_DEVICE + Params() {}; + CATLASS_DEVICE + Params(int32_t EP_, int32_t expertPerRank_, int32_t rank_, __gm__ int32_t *ptrTokenPerExpert_, + LayoutC layoutC_, int32_t n2_, int32_t n0_, HcclShmem& shmem_, int32_t offsetD_, Layout3D tokenPerExpertLayout_) : + ptrTokenPerExpert(ptrTokenPerExpert_), EP(EP_), + expertPerRank(expertPerRank_),rank(rank_), layoutC(layoutC_), n2(n2_), n0(n0_), + shmem(shmem_), offsetD(offsetD_), tokenPerExpertLayout(tokenPerExpertLayout_) + {} + }; + + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource const &resource, Params const ¶ms = Params{}) : params(params) + { + //ub:192KB + n0 = params.n0; + size_t ubOffset = 0; + for(int32_t i = 0; i < 2; i++) { + ubCList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += max_len * sizeof(ElementC); + ubDList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += max_len * sizeof(ElementD); + ubFp32List[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += max_len * sizeof(float); + scaleUbList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += (max_len / n0) * sizeof(float); + source_scale_offset[i] = -1; + } + tokenPerExpert.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(params.ptrTokenPerExpert)); + tokenPerExpertLayout = params.tokenPerExpertLayout; + is_ping = true; + } + CATLASS_DEVICE + void SetFlag() + { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + } + + CATLASS_DEVICE + void Finalize() + { + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + + } + CATLASS_DEVICE + ~BlockEpilogue() + { + + } + CATLASS_DEVICE + void operator() ( + AscendC::GlobalTensor const &gmC, + AscendC::GlobalTensor const &gmPerTokenScale, + GemmCoord& blockCoord, + GemmCoord& actualBlockShape, + int32_t groupIdx, + int32_t preSrcExpertSum, + AscendC::GlobalTensor preSumBeforeRank + ){ + is_ping = !is_ping; + auto event_id = is_ping ? EVENT_ID0 : EVENT_ID1; + auto event_id_2 = is_ping ? EVENT_ID2 : EVENT_ID3; + + auto &ubC = ubCList[is_ping]; + auto &ubD = ubDList[is_ping]; + int32_t gmCOffset = preSrcExpertSum * params.n2 + blockCoord.m() * params.n2 + blockCoord.n(); + auto gmTileC = gmC[gmCOffset]; + auto &ubCFp32 = ubFp32List[is_ping]; + auto &scaleUb = scaleUbList[is_ping]; + + LayoutC layoutGM{actualBlockShape.m(), actualBlockShape.n(), params.n2}; + LayoutC layoutUB{actualBlockShape.m(), actualBlockShape.n(), n0}; + + + AscendC::WaitFlag(event_id); + copyGmToUbC(ubC, gmTileC, layoutUB, layoutGM); + AscendC::SetFlag(event_id); + + AscendC::WaitFlag(event_id); + AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_NONE, -1, repeat, {1, 1, 8, 4}); + AscendC::SetFlag(event_id); + + + AscendC::WaitFlag(event_id_2); + AscendC::WaitFlag(event_id_2); + + int32_t gmScaleOffset = preSrcExpertSum + blockCoord.m(); + layout::VectorLayout scaleLauout{actualBlockShape.m()}; + if (source_scale_offset[event_id] != gmScaleOffset) { + source_scale_offset[event_id] = gmScaleOffset; + copyScaleGmToUb(scaleUb, gmPerTokenScale[gmScaleOffset], scaleLauout, scaleLauout); + } + + AscendC::SetFlag(event_id_2); + AscendC::SetFlag(event_id_2); + + + + + AscendC::WaitFlag(event_id_2); + AscendC::WaitFlag(event_id_2); // Note that the value must be MTE2_S instead of MTE2_V. + // Otherwise, 0 will be read, causing garbled characters. + AscendC::PipeBarrier(); + for (int32_t row = 0; row < actualBlockShape.m(); ++row) { + float scale = scaleUb(row); + Muls(ubCFp32[n0* row], ubCFp32[n0 * row] , scale, -1, (actualBlockShape.n() + 127) / 128 * 2, {1, 1, 8, 8}); + } + AscendC::PipeBarrier(); + AscendC::WaitFlag(event_id); + AscendC::Cast(ubD, ubCFp32, AscendC::RoundMode::CAST_RINT, -1, repeat, {1, 1, 4, 8}); + AscendC::SetFlag(event_id_2); + AscendC::SetFlag(event_id_2); + AscendC::SetFlag(event_id); + + int32_t lenTile = actualBlockShape.m(); + int32_t stTile = blockCoord.m(); + int32_t edTile = stTile + lenTile; + int32_t preSumRankInExpert = 0; + int32_t tileOffset = 0; + + AscendC::WaitFlag(event_id); + for (int32_t dstEpIdx = 0; dstEpIdx < params.EP; dstEpIdx ++) { + int32_t lenRankInExpert = tokenPerExpert(tokenPerExpertLayout(dstEpIdx, params.rank, groupIdx)); + int32_t dstExpertOffset = preSumBeforeRank(dstEpIdx * params.expertPerRank + groupIdx); + int32_t stRankInExpert = preSumRankInExpert; + int32_t edRankInExpert = stRankInExpert + lenRankInExpert; + preSumRankInExpert += lenRankInExpert; + if (stRankInExpert >= edTile) { + break; + } + else if (edRankInExpert <= stTile) { + continue; + } + int32_t stData = max(stRankInExpert, stTile); + int32_t edData = min(edRankInExpert, edTile); + uint32_t lenData = edData - stData; + if (lenData <= 0){ + continue; + } + + uint32_t dstOffsetInExpert = 0; + if (stTile > stRankInExpert) { + dstOffsetInExpert = stTile - stRankInExpert; + } + AscendC::GlobalTensor gmRemotePeer; + __gm__ void* dstPeermemPtr = params.shmem(params.offsetD, dstEpIdx); + gmRemotePeer.SetGlobalBuffer(reinterpret_cast<__gm__ ElementD*>(dstPeermemPtr)); + MatrixCoord dstOffset{dstOffsetInExpert + dstExpertOffset, blockCoord.n()}; + int64_t gmDstOffset = params.layoutC.GetOffset(dstOffset); + auto gmTileD = gmRemotePeer[gmDstOffset]; + LayoutC layoutGM2{lenData, actualBlockShape.n(), params.n2}; + LayoutC layoutUB2{lenData, actualBlockShape.n(), n0}; + copyUbToGmD(gmTileD, ubD[tileOffset * n0], layoutGM2, layoutUB2); + tileOffset += lenData; + } + AscendC::SetFlag(event_id); + + } +private: + + Params params; + AscendC::LocalTensor ubCList[UB_STAGES]; + AscendC::LocalTensor ubDList[UB_STAGES]; + AscendC::LocalTensor ubFp32List[UB_STAGES]; + AscendC::LocalTensor scaleUbList[UB_STAGES]; + int32_t source_scale_offset[UB_STAGES]; + + int32_t max_len = 8 * 32 / 4 * 128; + int32_t n0; + bool is_ping = false; + + + int32_t repeat = 128; + + CopyGmToUbC copyGmToUbC; + CopyUbToGmD copyUbToGmD; + + CopyScaleGmToUb copyScaleGmToUb; + AscendC::GlobalTensor tokenPerExpert; + Layout3D tokenPerExpertLayout; +}; +} +#endif \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/block_mmad_preload_async_fixpipe_quant.hpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/block_mmad_preload_async_fixpipe_quant.hpp new file mode 100644 index 000000000..88a8ec53f --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/block_mmad_preload_async_fixpipe_quant.hpp @@ -0,0 +1,554 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_FIXPIPE_QUANT_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_FIXPIPE_QUANT_HPP + +#include "catlass/catlass.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/coord.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/helper.hpp" +#include "dispatch_policy_custom.hpp" + + +namespace Catlass::Gemm::Block { + +template +__aicore__ inline void SyncFlagFunc(int32_t eventID) +{ + AscendC::SetFlag(eventID); + AscendC::WaitFlag(eventID); +} + +template < + uint32_t PRELOAD_STAGES_, + uint32_t L1_STAGES_, + uint32_t L0A_STAGES_, + uint32_t L0B_STAGES_, + uint32_t L0C_STAGES_, + bool ENABLE_UNIT_FLAG_, + bool ENABLE_SHUFFLE_K_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_ +> +struct BlockMmad < + MmadAtlasA2PreloadAsyncFixpipe< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_ + >, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_ +> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2PreloadAsyncFixpipe< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_ + >; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyGmToL1S = Gemm::Tile::CopyGmToL1>; + using CopyL1ToFP = typename Gemm::Tile::QuantTileCopy::CopyL1ToFP; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using CopyL0CToGm = typename std::conditional< + std::is_same_v, + Gemm::Tile::CopyL0CToGm, + typename TileCopy_::CopyL0CToGm + >::type; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES; + static constexpr uint32_t L1_STAGES = DispatchPolicy::L1_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K; + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L1S_TILE_SIZE = L1TileShape::N * sizeof(int64_t); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1TileShape::M * L1TileShape::N * sizeof(ElementAccumulator); + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + // Check L1TileShape + static_assert( + (std::is_same_v + ? (L1A_TILE_SIZE + L1B_TILE_SIZE + L1S_TILE_SIZE) * L1_STAGES <= ArchTag::L1_SIZE + : (L1A_TILE_SIZE + L1B_TILE_SIZE) * L1_STAGES <= ArchTag::L1_SIZE), + "L1TileShape exceeding the L1 space for the given data type" + ); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + + static constexpr auto L1A_LAYOUT = LayoutAInL1::template MakeLayout( + L1TileShape::M, L1TileShape::K); + static constexpr auto L1B_LAYOUT = LayoutBInL1::template MakeLayout( + L1TileShape::K, L1TileShape::N); + + CATLASS_DEVICE + BlockMmad(Arch::Resource &resource, __gm__ int32_t* flagPtr = nullptr, int32_t expertPerRank = 0, + uint32_t l1BufAddrStart = 0, uint32_t FpAddrStart = 0) + { + syncGroupIdx = 0; + ptrSoftFlagBase_ = flagPtr; + expertPerRank_ = expertPerRank; + InitL1(resource, l1BufAddrStart); + InitFpBuf(resource, FpAddrStart); + InitL0A(resource); + InitL0B(resource); + InitL0C(resource); + } + + CATLASS_DEVICE + ~BlockMmad() + { + SynchronizeBlock(); + for (uint32_t i = 0; i < L1_STAGES; ++i) { + AscendC::WaitFlag(l1AEventList[i]); + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + AscendC::WaitFlag(l0BEventList[i]); + } + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + AscendC::WaitFlag(l0CEventList[i]); + } + if constexpr (std::is_same_v) { + AscendC::WaitFlag(0); + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmBlockA, LayoutA const &layoutA, + AscendC::GlobalTensor const &gmBlockB, LayoutB const &layoutB, + AscendC::GlobalTensor const &gmBlockC, LayoutC const &layoutC, + AscendC::GlobalTensor const &gmBlockS, layout::VectorLayout const &layoutScale, + GemmCoord const &actualShape, int32_t syncLoopIdx = -1, int32_t flag = 0 + ) + { + uint32_t kTileCount = CeilDiv(actualShape.k()); + + uint32_t mRound = RoundUp(actualShape.m()); + uint32_t nRound = RoundUp(actualShape.n()); + + uint32_t startTileIdx = 0; + if constexpr (ENABLE_SHUFFLE_K) { + startTileIdx = AscendC::GetBlockIdx() % kTileCount; + } + + for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; ++kLoopIdx) { + uint32_t kTileIdx = (startTileIdx + kLoopIdx < kTileCount) ? + (startTileIdx + kLoopIdx) : (startTileIdx + kLoopIdx - kTileCount); + + uint32_t kActual = (kTileIdx < kTileCount - 1) ? + L1TileShape::K : (actualShape.k() - kTileIdx * L1TileShape::K); + + // Emission load instruction from GM to L1 + MatrixCoord gmTileAOffset{0, kTileIdx * L1TileShape::K}; + MatrixCoord gmTileBOffset{kTileIdx * L1TileShape::K, 0}; + auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)]; + auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)]; + // Load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1ListId]); + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual)); + copyGmToL1A(l1ATensorList[l1ListId], gmTileA, L1A_LAYOUT, layoutTileA); + AscendC::SetFlag(l1AEventList[l1ListId]); + // Load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1ListId]); + auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n())); + copyGmToL1B(l1BTensorList[l1ListId], gmTileB, L1B_LAYOUT, layoutTileB); + AscendC::SetFlag(l1BEventList[l1ListId]); + + // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile + if (preloadCount == PRELOAD_STAGES) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + } + + // Store the current load status + uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ? + (l1TileMmadParamsId + preloadCount) : (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES); + auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId]; + l1TileMmadParams.l1ListId = l1ListId; + l1TileMmadParams.mRound = mRound; + l1TileMmadParams.nRound = nRound; + l1TileMmadParams.kActual = kActual; + l1TileMmadParams.isKLoopFirst = (kLoopIdx == 0); + l1TileMmadParams.isKLoopLast = (kLoopIdx == kTileCount - 1); + l1TileMmadParams.flag = flag; + if (kLoopIdx == kTileCount - 1) { + l1TileMmadParams.gmBlockC = gmBlockC; + l1TileMmadParams.gmBlockS = gmBlockS; + l1TileMmadParams.layoutCInGm = layoutC.GetTileLayout(actualShape.GetCoordMN()); + l1TileMmadParams.layoutScale = layoutScale; + l1TileMmadParams.syncLoopIdx = syncLoopIdx; + } + + if (preloadCount < PRELOAD_STAGES) { + ++preloadCount; + } else { + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + } + l1ListId = (l1ListId + 1 < L1_STAGES) ? (l1ListId + 1) : 0; + } + } + + CATLASS_DEVICE + void SynchronizeBlock() + { + while (preloadCount > 0) { + L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]); + l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0; + --preloadCount; + } + } + + CATLASS_DEVICE + void Finalize(int32_t target, int32_t flag = 0) + { + if (ptrSoftFlagBase_ != nullptr) { + if (target < 0) { + return; + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::GlobalTensor flagGlobal; + flagGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(ptrSoftFlagBase_) + (expertPerRank_ + AscendC::GetBlockIdx()) * FLAGSTRIDE); + AscendC::DataCopy(flagGlobal, l1FTensor[target * 16], FLAGSTRIDE); + } + else { + for(;syncGroupIdx <= target; syncGroupIdx++) { + int32_t flagId = syncGroupIdx / 15 + flag; + AscendC::CrossCoreSetFlag<0x2, PIPE_FIX>(flagId); + } + } + } +private: + struct L1TileMmadParams { + uint32_t l1ListId; + uint32_t mRound; + uint32_t nRound; + uint32_t kActual; + bool isKLoopFirst; + bool isKLoopLast; + AscendC::GlobalTensor gmBlockC; + AscendC::GlobalTensor gmBlockS; + LayoutC layoutCInGm; + layout::VectorLayout layoutScale; + int32_t syncLoopIdx; + int32_t flag; + CATLASS_DEVICE + L1TileMmadParams() = default; + }; + + CATLASS_DEVICE + void InitL1(Arch::Resource &resource, uint32_t l1BufAddrStart) + { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1_STAGES; + + for (uint32_t i = 0; i < L1_STAGES; ++i) { + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + l1AEventList[i] = i; + l1BEventList[i] = i + L1_STAGES; + AscendC::SetFlag(l1AEventList[i]); + AscendC::SetFlag(l1BEventList[i]); + } + uint32_t l1SOffset = l1BOffset + L1B_TILE_SIZE * L1_STAGES; + if constexpr (std::is_same_v) { + l1STensor = resource.l1Buf.template GetBufferByByte(l1SOffset); + AscendC::SetFlag(0); + } + if (ptrSoftFlagBase_ != nullptr) { + // Initialize the flag matrix (structure as below): + // 1 0 0 0 0 0 0 0 + // 2 0 0 0 0 0 0 0 + // ... + // 16 0 0 0 0 0 0 0 + // Then move it to L1 + uint32_t l1FOffset = l1SOffset + L1S_TILE_SIZE; + l1FTensor = resource.l1Buf.template GetBufferByByte(l1FOffset); + AscendC::GlobalTensor flagBase; + flagBase.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(ptrSoftFlagBase_)); + AscendC::DataCopy(l1FTensor, flagBase, expertPerRank_ * FLAGSTRIDE); + } + } + + CATLASS_DEVICE + void InitFpBuf(Arch::Resource &resource, uint32_t FpAddrStart) + { + uint32_t FpOffset = FpAddrStart; + fixpipeBuf = resource.fpBuf.template GetBufferByByte(FpOffset); + } + + CATLASS_DEVICE + void InitL0A(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0A_STAGES; ++i) { + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + l0AEventList[i] = i; + AscendC::SetFlag(l0AEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0B(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0B_STAGES; ++i) { + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + l0BEventList[i] = i + L0A_STAGES; + AscendC::SetFlag(l0BEventList[i]); + } + } + + CATLASS_DEVICE + void InitL0C(Arch::Resource &resource) + { + for (uint32_t i = 0; i < L0C_STAGES; ++i) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } + + CATLASS_DEVICE + void L1TileMmad(L1TileMmadParams const ¶ms) + { + uint32_t mPartLoop = CeilDiv(params.mRound); + uint32_t nPartLoop = CeilDiv(params.nRound); + uint32_t kPartLoop = CeilDiv(params.kActual); + auto &l1ATensor = l1ATensorList[params.l1ListId]; + auto &l1BTensor = l1BTensorList[params.l1ListId]; + + auto &l0CTensor = l0CTensorList[l0CListId]; + LayoutCInL0 layoutCInL0 = LayoutCInL0::MakeLayoutInL0C(MakeCoord(params.mRound, params.nRound)); + + if constexpr (!ENABLE_UNIT_FLAG) { + if (params.isKLoopFirst) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + } + + for (uint32_t mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) { + uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? + L0TileShape::M : (params.mRound - mPartIdx * L0TileShape::M); + + for (uint32_t kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) { + uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? + L0TileShape::K : (params.kActual - kPartIdx * L0TileShape::K); + + auto &l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = LayoutAInL0::template MakeLayout(mPartActual, kPartActual); + auto l1AOffset = MakeCoord(mPartIdx, kPartIdx) * L0TileShape::ToCoordMK(); + auto l1ATile = l1ATensor[L1A_LAYOUT.GetOffset(l1AOffset)]; + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mPartIdx == 0) && (kPartIdx == 0)) { + AscendC::WaitFlag(l1AEventList[params.l1ListId]); + } + copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, L1A_LAYOUT); + if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) { + AscendC::SetFlag(l1AEventList[params.l1ListId]); + } + + for (uint32_t nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) { + uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? + L0TileShape::N : (params.nRound - nPartIdx * L0TileShape::N); + + auto &l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = LayoutBInL0::template MakeLayout(kPartActual, nPartActual); + auto l1BOffset = MakeCoord(kPartIdx, nPartIdx) * L0TileShape::ToCoordKN(); + auto l1BTile = l1BTensor[L1B_LAYOUT.GetOffset(l1BOffset)]; + + AscendC::WaitFlag(l0BEventList[l0BListId]); + if ((kPartIdx == 0) && (nPartIdx == 0)) { + AscendC::WaitFlag(l1BEventList[params.l1ListId]); + } + copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, L1B_LAYOUT); + if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + AscendC::SetFlag(l1BEventList[params.l1ListId]); + } + + AscendC::SetFlag(EVENT_ID0); + + auto l0COffset = MakeCoord(mPartIdx, nPartIdx) * L0TileShape::ToCoordMN(); + auto l0CTile = l0CTensor[layoutCInL0.GetOffset(l0COffset)]; + + AscendC::WaitFlag(EVENT_ID0); + // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0 + bool initC = (params.isKLoopFirst && (kPartIdx == 0)); + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if (params.isKLoopLast && + (mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag); + + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + + if (params.isKLoopLast) { + auto layoutCInGm = params.layoutCInGm; + if constexpr (std::is_same_v) { + auto layoutScale = params.layoutScale; + auto layoutTileS = layoutScale.GetTileLayout(MakeCoord(layoutCInGm.shape(1))); + AscendC::WaitFlag(0); + copyGmToL1S(l1STensor, params.gmBlockS, layoutTileS, layoutTileS); + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + + copyL1ToFP(fixpipeBuf, l1STensor, layoutTileS, layoutTileS); + AscendC::SetFlag(0); + AscendC::PipeBarrier(); + } + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + if constexpr (std::is_same_v) { + copyL0CToGm(params.gmBlockC, l0CTensor, l1STensor, layoutCInGm, layoutCInL0); + } else { + copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0); + } + AscendC::SetFlag(l0CEventList[l0CListId]); + } else { + if constexpr (std::is_same_v) { + copyL0CToGm(params.gmBlockC, l0CTensor, l1STensor, layoutCInGm, layoutCInL0, 0b11); + } else { + copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0, 0b11); + } + } + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + #ifdef __TILE_SYNC__ + if (params.flag > 0) { + int32_t flagId = params.flag + params.syncLoopIdx / 8; + AscendC::CrossCoreSetFlag<0x2, PIPE_FIX>(flagId); + } + #else + Finalize(params.syncLoopIdx, params.flag); + #endif + } + } + + AscendC::LocalTensor fixpipeBuf; + + AscendC::LocalTensor l1ATensorList[L1_STAGES]; + AscendC::LocalTensor l1BTensorList[L1_STAGES]; + AscendC::LocalTensor l1STensor; + AscendC::LocalTensor l1FTensor; + int32_t syncGroupIdx; + int32_t l1AEventList[L1_STAGES]; + int32_t l1BEventList[L1_STAGES]; + uint32_t l1ListId{0}; + + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + uint32_t l0AListId{0}; + + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + uint32_t l0BListId{0}; + + AscendC::LocalTensor l0CTensorList[L0C_STAGES_]; + int32_t l0CEventList[L0C_STAGES_]; + uint32_t l0CListId{0}; + + L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES]; + uint32_t l1TileMmadParamsId{0}; + uint32_t preloadCount{0}; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyGmToL1S copyGmToL1S; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; + CopyL1ToFP copyL1ToFP; + + __gm__ int32_t* ptrSoftFlagBase_ = nullptr; + int32_t expertPerRank_; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_FIXPIPE_QUANT_HPP \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/const_args.hpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/const_args.hpp new file mode 100644 index 000000000..3249138e1 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/const_args.hpp @@ -0,0 +1,11 @@ + +#ifndef CONST_ARGS_HPP +#define CONST_ARGS_HPP +constexpr static uint64_t MB_SIZE = 1024 * 1024UL; +constexpr static int32_t NUMS_PER_FLAG = 16; +constexpr static int32_t CACHE_LINE = 512; +constexpr static int32_t RESET_VAL = 0xffff; +constexpr static int32_t FLAGSTRIDE = 16; +constexpr static int32_t UB_ALIGN = 32; +constexpr uint16_t CROSS_CORE_FLAG_MAX_SET_COUNT = 15; +#endif \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/copy_gm_to_l1_custom.hpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/copy_gm_to_l1_custom.hpp new file mode 100644 index 000000000..e2954f432 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/copy_gm_to_l1_custom.hpp @@ -0,0 +1,40 @@ +#ifndef COPY_GM_TO_L1_CUSTOM_HPP +#define COPY_GM_TO_L1_CUSTOM_HPP + +namespace Catlass::Gemm::Tile { + /// Partial specialization for nZ in and nZ out. + template < + class ArchTag, + class Element + > + struct CopyGmToL1> { + using LayoutDst = layout::VectorLayout; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); // int64, 32/8=4 + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t blockCount = 1; + uint32_t blockLen = CeilDiv(layoutSrc.shape(0)); + + AscendC::DataCopyParams repeatParams; + + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } + }; +} +#endif // COPY_GM_TO_L1_CUSTOM_HPP \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/copy_l0c_to_gm_custom.hpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/copy_l0c_to_gm_custom.hpp new file mode 100644 index 000000000..ba4779840 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/copy_l0c_to_gm_custom.hpp @@ -0,0 +1,47 @@ +#ifndef COPY_L0C_TO_GM_CUSTOM_HPP +#define COPY_L0C_TO_GM_CUSTOM_HPP + +namespace Catlass::Gemm::Tile { + template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_ + > + struct CopyL0CToGm, + ScaleGranularity::PER_CHANNEL, + ReluEnable_> + { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, AscendC::LocalTensor cbufWorkspace, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, cbufWorkspace, intriParams); + } + }; +} +#endif // COPY_L0C_TO_GM_CUSTOM_HPP \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/dispatch_policy_custom.hpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/dispatch_policy_custom.hpp new file mode 100644 index 000000000..7e30114e4 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/dispatch_policy_custom.hpp @@ -0,0 +1,47 @@ +#ifndef DISPATH_POLICY_CUSTOM_HPP +#define DISPATH_POLICY_CUSTOM_HPP + +namespace Catlass::Gemm { + template + struct MmadAtlasA2PreloadFixpipeQuant : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; + }; + + template + struct MmadAtlasA2PreloadAsyncFixpipe : + public MmadAtlasA2PreloadAsync< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_ + > { + }; +} + +namespace Catlass::Epilogue { + + template + struct EpilogueAtlasA2UnQuant { + using ArchTag = Arch::AtlasA2; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + }; + + template + struct EpilogueAtlasA2PerTokenDequantSwigluQuant { + using ArchTag = Arch::AtlasA2; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + }; + + template + struct EpilogueAtlasA2PerTokenDequantV2 { + using ArchTag = Arch::AtlasA2; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + }; +} +#endif // DISPATH_POLICY_CUSTOM_HPP \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/get_tensor_addr.hpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/get_tensor_addr.hpp new file mode 100644 index 000000000..67b32c25b --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/get_tensor_addr.hpp @@ -0,0 +1,16 @@ +#ifndef GET_TENSOR_ADDR_HPP +#define GET_TENSOR_ADDR_HPP +#include "kernel_operator.h" + +#define FORCE_INLINE_AICORE inline __attribute__((always_inline)) __aicore__ + +template +FORCE_INLINE_AICORE __gm__ T* GetTensorAddr(uint32_t index, GM_ADDR tensorPtr) { + __gm__ uint64_t* dataAddr = reinterpret_cast<__gm__ uint64_t*>(tensorPtr); + uint64_t tensorPtrOffset = *dataAddr; // The offset of the data address from the first address. + // Moving 3 bits to the right means dividing by sizeof(uint64 t). + __gm__ uint64_t* retPtr = dataAddr + (tensorPtrOffset >> 3); + return reinterpret_cast<__gm__ T*>(*(retPtr + index)); +} + +#endif // GET_TENSOR_ADDR_HPP \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/hccl_shmem.hpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/hccl_shmem.hpp new file mode 100644 index 000000000..1adc5e3a6 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/hccl_shmem.hpp @@ -0,0 +1,327 @@ +#ifndef SYNC_UTIL_HPP +#define SYNC_UTIL_HPP + + +#include "kernel_operator.h" +#include "const_args.hpp" + +#ifdef HCCL_COMM +#include "moe_distribute_base.h" +using namespace AscendC::HcclContextDef; + +#else +#include "shmem_api.h" +#endif + +#define FORCE_INLINE_AICORE inline __attribute__((always_inline)) __aicore__ +constexpr int32_t MAX_RANK_SIZE = 32; +constexpr int32_t SHMEM_MEM = 700 * MB_SIZE; + +constexpr uint16_t SEND_SYNC_EVENT_ID = 9; +constexpr uint16_t RECV_SYNC_EVENT_ID = 10; + +constexpr uint32_t SELF_STATE_OFFSET = 256 * 1024; +constexpr uint32_t STATE_OFFSET = 512; + +FORCE_INLINE_AICORE void AicSyncAll() { + AscendC::CrossCoreSetFlag<0x0, PIPE_FIX>(8); + AscendC::CrossCoreWaitFlag<0x0>(8); +} + +template +FORCE_INLINE_AICORE void gm_store(__gm__ T *addr, T val) { + *((__gm__ T *)addr) = val; +} + +template +FORCE_INLINE_AICORE T gm_load(__gm__ T *cache) { + return *((__gm__ T *)cache); +} + +template +FORCE_INLINE_AICORE void gm_dcci(__gm__ T * addr) { + using namespace AscendC; + GlobalTensor global; + global.SetGlobalBuffer(reinterpret_cast(addr)); + + // Important: add hint to avoid dcci being optimized by compiler + __asm__ __volatile__(""); + DataCacheCleanAndInvalid(global); + __asm__ __volatile__(""); +} + +FORCE_INLINE_AICORE int32_t gm_signal_wait_until_eq_for_barrier(__gm__ int32_t *sig_addr, int32_t cmp_val) { + do { + gm_dcci((__gm__ uint8_t *)sig_addr); + if (*sig_addr == cmp_val) { + return *sig_addr; + } + if (*sig_addr == cmp_val + 1) { + return *sig_addr; + } + } while (true); + return -1; +} + +FORCE_INLINE_AICORE void gm_signal_wait_until_ne(__gm__ int32_t *sig_addr, int32_t cmp_val) { + do { + AscendC::LocalTensor ub; + ub.address_.logicPos = static_cast(AscendC::TPosition::VECIN); + ub.address_.bufferAddr = 0; + AscendC::GlobalTensor sig; + sig.SetGlobalBuffer(sig_addr); + AscendC::DataCopy(ub, sig, 8); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + if (ub(0) != cmp_val) { + return; + } + } while (true); + return; +} + + +class HcclShmem { +public: + #ifdef HCCL_COMM // HCCL needs to initialize the HCCL context + __gm__ HcclOpResParamCustom *WinContext_{nullptr}; + Hccl hccl_; + AscendC::LocalTensor ub; + FORCE_INLINE_AICORE + HcclShmem(){ + auto contextGM0 = AscendC::GetHcclContext(); + WinContext_ = (__gm__ HcclOpResParamCustom *)contextGM0; + + m_rank = WinContext_->localUsrRankId; + m_rankSize = WinContext_->rankSize; + m_segmentSize = WinContext_->winSize; + } + #else + FORCE_INLINE_AICORE + HcclShmem(){ + m_segmentSize = SHMEM_MEM; + } + FORCE_INLINE_AICORE + void initShmem(GM_ADDR symmetricPtr_, size_t rank, size_t rankSize) { + symmetricPtr = symmetricPtr_; + m_rank = rank; + m_rankSize = rankSize; + } + #endif + + FORCE_INLINE_AICORE + GM_ADDR operator() () const { // No parameters: return pointer to local peermem + #ifdef HCCL_COMM + return (GM_ADDR)(WinContext_->localWindowsIn); + #else + return reinterpret_cast(shmem_ptr(symmetricPtr, m_rank)); + #endif + } + + FORCE_INLINE_AICORE + GM_ADDR operator() (int32_t index) const { // With index parameter: return pointer to the base address of remote peermem + #ifdef HCCL_COMM + return (GM_ADDR)((index == m_rank) ? WinContext_->localWindowsIn : + ((HcclRankRelationResV2Custom *)(WinContext_->remoteRes[index].nextDevicePtr))->windowsIn); + #else + return reinterpret_cast(shmem_ptr(symmetricPtr, index)); + #endif + } + + FORCE_INLINE_AICORE + GM_ADDR operator () (int64_t offset, int32_t rankId) const { + #ifdef HCCL_COMM + if (offset < 0 || offset >= m_segmentSize) { + return nullptr; + } + if (rankId < 0 || rankId >= m_rankSize) { + return nullptr; + } + return (GM_ADDR)((rankId == m_rank) ? WinContext_->localWindowsIn : + ((HcclRankRelationResV2Custom *)(WinContext_->remoteRes[rankId].nextDevicePtr))->windowsIn) + offset; + #else + return reinterpret_cast(shmem_ptr((symmetricPtr + offset), rankId)); + #endif + } + + + + FORCE_INLINE_AICORE + size_t SegmentSize() const { + return m_segmentSize; + } + + FORCE_INLINE_AICORE + int32_t RankSize() const { + return m_rankSize; + } + + + FORCE_INLINE_AICORE + ~HcclShmem() { + } + + + FORCE_INLINE_AICORE + void CrossRankSync() { + uint64_t flag_offset = (m_segmentSize - MB_SIZE) / sizeof(int32_t); + __gm__ int32_t* sync_counter = (__gm__ int32_t*)(*this)() + flag_offset; + __gm__ int32_t* sync_base = (__gm__ int32_t*)(*this)() + flag_offset + 2048; + int count = gm_load(sync_base) + 1; + int vec_id = AscendC::GetBlockIdx(); + int vec_size = AscendC::GetBlockNum() * AscendC::GetTaskRation(); + for(int i = vec_id; i < m_rankSize; i += vec_size) { + __gm__ int32_t* sync_remote = (__gm__ int32_t*)((*this)(i)) + flag_offset + m_rank * 16; + gm_store(sync_remote, count); + gm_dcci((__gm__ uint8_t*)sync_remote); + auto sync_check = sync_counter + i * 16; + gm_signal_wait_until_eq_for_barrier(sync_check, count); + } + + AscendC::SyncAll(); + gm_store(sync_base, count); + } + + + FORCE_INLINE_AICORE + void InitStatusTargetSum() + { + using namespace AscendC; + uint64_t flag_offset = (m_segmentSize - MB_SIZE) + SELF_STATE_OFFSET; + //uint64_t self_state_offset = (m_segmentSize - 2 * MB_SIZE); + // ep state + //uint32_t coreIdx = get_block_idx();; + uint32_t coreIdx = GetBlockIdx(); + GlobalTensor selfStatusTensor; + selfStatusTensor.SetGlobalBuffer((__gm__ int32_t *)((*this)() + flag_offset)); + __asm__ __volatile__(""); + DataCacheCleanAndInvalid(selfStatusTensor[coreIdx * UB_ALIGN]); + __asm__ __volatile__(""); + int32_t state = selfStatusTensor(coreIdx * UB_ALIGN); + if (state == 0) { + sumTarget_ = static_cast(1.0); + selfStatusTensor(coreIdx * UB_ALIGN) = 0x3F800000; // 1.0f + epStateValue_ = 0x3F800000; // 1.0f + } else { + sumTarget_ = static_cast(0.0); + selfStatusTensor(coreIdx * UB_ALIGN) = 0; + epStateValue_ = 0; + } + __asm__ __volatile__(""); + DataCacheCleanAndInvalid(selfStatusTensor[coreIdx * UB_ALIGN]); + __asm__ __volatile__(""); + } + + FORCE_INLINE_AICORE + void CrossRankSyncV2Set(AscendC::LocalTensor ctrBuffer) { + //subblockid = 0 + uint32_t stateOffset_ = STATE_OFFSET; + // uint32_t epStateOffsetOnWin_ = m_rank * stateOffset_; + + uint64_t flag_offset = (m_segmentSize - MB_SIZE) + m_rank * stateOffset_; + //uint64_t flag_offset = (m_segmentSize - MB_SIZE); + int vec_size = get_block_num(); + int vec_id = get_block_idx(); + + AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(RECV_SYNC_EVENT_ID); + AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(SEND_SYNC_EVENT_ID); + AscendC::CrossCoreWaitFlag(SEND_SYNC_EVENT_ID); + AscendC::PipeBarrier(); + + ctrBuffer.SetValue(0, epStateValue_); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + for (uint32_t dstEpIdx = vec_id; dstEpIdx < m_rankSize; dstEpIdx += vec_size) { + AscendC::GlobalTensor gmDstStates; + gmDstStates.SetGlobalBuffer((__gm__ int32_t*)((*this)(flag_offset, dstEpIdx))); + DataCopy(gmDstStates, ctrBuffer, 8); + } + AscendC::CrossCoreWaitFlag(RECV_SYNC_EVENT_ID); + } + + FORCE_INLINE_AICORE + void CrossRankSyncV2Wait(AscendC::LocalTensor statusTensor, AscendC::LocalTensor gatherMaskOutTensor, + AscendC::LocalTensor gatherTmpTensor, AscendC::LocalTensor statusSumOutTensor) { + + uint64_t flag_offset = (m_segmentSize - MB_SIZE); + int vec_size = get_block_num(); + int vec_id = get_block_idx(); + uint32_t stateOffset_ = STATE_OFFSET; + + uint32_t sendRankNum_ = m_rankSize / vec_size; + uint32_t remainderRankNum = m_rankSize % vec_size; + uint32_t startRankId_ = sendRankNum_ * vec_id; + if (vec_id < remainderRankNum) { + sendRankNum_++; + startRankId_ += vec_id; + } else { + startRankId_ += remainderRankNum; + } + uint32_t endRankId_ = startRankId_ + sendRankNum_; + AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(SEND_SYNC_EVENT_ID); + + AscendC::GlobalTensor epStatusSpaceGlobalTensor_; + epStatusSpaceGlobalTensor_.SetGlobalBuffer((__gm__ float *)((*this)() + flag_offset)); + + if (startRankId_ < m_rankSize) { + AscendC::PipeBarrier(); + gatherTmpTensor.SetValue(0, 1); + uint32_t mask = 1; // gatherMask + sum + uint64_t rsvdCnt = 0; + // DataCopyParams intriParams{static_cast(sendRankNum_), 1, + // static_cast((moeSendNum_ > 512) ? 7 : 15), 0}; + AscendC::DataCopyParams intriParams{static_cast(sendRankNum_), 1, + static_cast(15), 0}; + + float sumOfFlag = static_cast(-1.0); + float minTarget = (sumTarget_ * sendRankNum_) - (float)0.5; + float maxTarget = (sumTarget_ * sendRankNum_) + (float)0.5; + AscendC::SumParams sumParams{1, sendRankNum_, sendRankNum_}; + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + while ((sumOfFlag < minTarget) || (sumOfFlag > maxTarget)) { + AscendC::DataCopy(statusTensor, epStatusSpaceGlobalTensor_[startRankId_ * stateOffset_ / sizeof(float)], + intriParams); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + GatherMask(gatherMaskOutTensor, statusTensor, gatherTmpTensor, true, mask, + {1, (uint16_t)sendRankNum_, 1, 0}, rsvdCnt); + + AscendC::PipeBarrier(); + AscendC::Sum(statusSumOutTensor, gatherMaskOutTensor, sumParams); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + sumOfFlag = statusSumOutTensor.GetValue(0); + } + } + + AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(RECV_SYNC_EVENT_ID); + AscendC::CrossCoreWaitFlag(RECV_SYNC_EVENT_ID); + + //unpermute + AscendC::CrossCoreWaitFlag(SEND_SYNC_EVENT_ID); + } + + + FORCE_INLINE_AICORE + __gm__ int32_t* SyncBaseAddr() { + uint64_t flag_offset = (m_segmentSize - MB_SIZE) / sizeof(int32_t); + return (__gm__ int32_t*)(*this)() + flag_offset + 2048; + } + +private: + GM_ADDR symmetricPtr; + int32_t m_rank; + int32_t m_rankSize; + size_t m_segmentSize; + float sumTarget_{0.0}; + int32_t epStateValue_; +}; + + + + +#endif diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/layout3d.hpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/layout3d.hpp new file mode 100644 index 000000000..7cc3a9c1f --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/layout3d.hpp @@ -0,0 +1,20 @@ +#ifndef LAYOUT_3D_HPP +#define LAYOUT_3D_HPP +#include "kernel_operator.h" +#include "catlass/catlass.hpp" +class Layout3D { + int64_t strides[2]; + public: + CATLASS_DEVICE + Layout3D() {} + CATLASS_DEVICE + Layout3D(int64_t stride0, int64_t stride1) { + strides[0] = stride0; + strides[1] = stride1; + } + CATLASS_DEVICE + int64_t operator() (int64_t dim0, int64_t dim1, int64_t dim2) { + return dim0 * strides[0] + dim1 * strides[1] + dim2; + } +}; +#endif // LAYOUT_3D_HPP diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/moe_distribute_base.h b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/moe_distribute_base.h new file mode 100755 index 000000000..785eb3160 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/moe_distribute_base.h @@ -0,0 +1,365 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_distribute_base.h + * \brief + */ + +#ifndef MOE_DISTRIBUTE_BASE_H +#define MOE_DISTRIBUTE_BASE_H + +#include "kernel_operator.h" + +constexpr uint32_t LOCAL_NOTIFY_MAX_NUM = 64; +constexpr uint32_t LOCAL_STREAM_MAX_NUM = 19U; +constexpr uint32_t AICPU_OP_NOTIFY_MAX_NUM = 2; +constexpr uint32_t AICPU_MAX_RANK_NUM = 128 * 1024; +constexpr uint32_t TIME_CYCLE = 50; // 系统cycle数转换成时间的基准单位，固定为50 + +struct HcclSignalInfo { + uint64_t resId; // 在代表event时为eventid，notify时为notifyid + uint64_t addr; + uint32_t devId; + uint32_t tsId; + uint32_t rankId; + uint32_t flag; +}; + +struct ListCommon { + uint64_t nextHost; + uint64_t preHost; + uint64_t nextDevice; + uint64_t preDevice; +}; + +struct HcclStreamInfo { + int32_t streamIds; + uint32_t sqIds; + uint32_t cqIds; // 记录物理cqId + uint32_t logicCqids; // 记录逻辑cqId +}; + +struct LocalResInfoV2 { + uint32_t streamNum; + uint32_t signalNum; + HcclSignalInfo localSignals[LOCAL_NOTIFY_MAX_NUM]; + HcclStreamInfo streamInfo[LOCAL_STREAM_MAX_NUM]; + HcclStreamInfo mainStreamInfo; + HcclSignalInfo aicpuOpNotify[AICPU_OP_NOTIFY_MAX_NUM]; // 集合通信AICPU展开资源 + ListCommon nextTagRes; // HccltagLocalResV2 +}; + +enum class rtFloatOverflowMode_t { + RT_OVERFLOW_MODE_SATURATION = 0, + RT_OVERFLOW_MODE_INFNAN, + RT_OVERFLOW_MODE_UNDEF, +}; + +struct AlgoTopoInfo { + uint32_t userRank; // 通信域 RankID + uint32_t userRankSize; // 通信域的Rank数量 + int32_t deviceLogicId; + bool isSingleMeshAggregation; + uint32_t deviceNumPerAggregation; // 每个Module中的Device数量 + uint32_t superPodNum; // 集群中总的超节点数 + uint32_t devicePhyId; + uint32_t topoType; // TopoType + uint32_t deviceType; + uint32_t serverNum; + uint32_t meshAggregationRankSize; + uint32_t multiModuleDiffDeviceNumMode; + uint32_t multiSuperPodDiffServerNumMode; + uint32_t realUserRank; + bool isDiffDeviceModule; + bool isDiffDeviceType; + uint32_t gcdDeviceNumPerAggregation; + uint32_t moduleNum; + uint32_t isUsedRdmaRankPairNum; + uint64_t isUsedRdmaRankPair; + uint32_t pairLinkCounterNum; + uint64_t pairLinkCounter; + uint32_t nicNum; + uint64_t nicList; // niclist数组指针 + uint64_t complanRankLength; // complanRank占用的字节数 + uint64_t complanRank; // 指针 + uint64_t bridgeRankNum; // bridgeRank占用的个数 + uint64_t bridgeRank; // 指针 + uint64_t serverAndsuperPodRankLength; // serverAndsuperPodRank占用的字节数 + uint64_t serverAndsuperPodRank; // 指针 +}; + +struct HcclOpConfig { + uint8_t deterministic; //确定性计算开关 + uint8_t retryEnable; // 是否重执行 + uint8_t highPerfEnable; + uint8_t padding[5]; // 大小需要64By对齐，未来添加参数时减小padding + uint8_t linkTimeOut[8]; // 发送超时时长 + uint64_t notifyWaitTime; // 超时时长，同HCCL_EXEC_TIMEOUT + uint32_t retryHoldTime; + uint32_t retryIntervalTime; + bool interHccsDisable = false; //使能rdma开关 + rtFloatOverflowMode_t floatOverflowMode = rtFloatOverflowMode_t::RT_OVERFLOW_MODE_UNDEF; + uint32_t multiQpThreshold = 512; // 多QP每个QP分担数据量最小阈值 +}; + +struct HcclMC2WorkSpace { + uint64_t workSpace; + uint64_t workSpaceSize; +}; + +struct RemoteResPtr { + uint64_t nextHostPtr; + uint64_t nextDevicePtr; +}; + +struct HDCommunicateParams { + uint64_t hostAddr { 0 }; + uint64_t deviceAddr { 0 }; + uint64_t readCacheAddr { 0 }; + uint32_t devMemSize{ 0 }; + uint32_t buffLen{ 0 }; + uint32_t flag{ 0 }; +}; + +struct HcclRankRelationResV2Custom { + uint32_t remoteUsrRankId; + uint32_t remoteWorldRank; + uint64_t windowsIn; + uint64_t windowsOut; + uint64_t windowsExp; + ListCommon nextTagRes; +}; + +struct HcclOpResParamCustom { + // 本地资源 + HcclMC2WorkSpace mc2WorkSpace; + uint32_t localUsrRankId; // usrrankid + uint32_t rankSize; // 通信域内total rank个数 + uint64_t winSize; // 每个win大小，静态图时，可能是0，如果通信域内也有动态图，则可能为非0 + uint64_t localWindowsIn; // 全F为无效值 + uint64_t localWindowsOut; // 全F为无效值 + char hcomId[128]; + // aicore识别remote window + uint64_t winExpSize; + uint64_t localWindowsExp; + uint32_t rWinStart; // 为HcclRankRelationRes起始位置 + uint32_t rWinOffset; // 为HcclRemoteRes的大小 + uint64_t version; + LocalResInfoV2 localRes; + AlgoTopoInfo topoInfo; + + // 外部配置参数 + HcclOpConfig config; + uint64_t hostStateInfo; + uint64_t aicpuStateInfo; + uint64_t lockAddr; + uint32_t rsv[16]; + uint32_t notifysize; // RDMA场景使用，910B/910_93为4B，其余芯片为8B + uint32_t remoteResNum; // 有效的remoteResNum + RemoteResPtr remoteRes[AICPU_MAX_RANK_NUM]; //数组指针，指向HcclRankRelationResV2，下标为remoteUserRankId + + // communicate retry + HDCommunicateParams kfcControlTransferH2DParams; + HDCommunicateParams kfcStatusTransferD2HParams; + uint64_t tinyMem; // for all2all + uint64_t tinyMemSize; + // 零拷贝场景使用 + uint64_t zeroCopyHeadPtr; + uint64_t zeroCopyTailPtr; + uint64_t zeroCopyRingBuffer; + uint64_t zeroCopyIpcPtrs[16]; // 保存集合通信时每个对端的输入输出内存地址 + uint32_t zeroCopyDevicePhyId[16]; // 保存每个rank对应的物理卡Id + + bool utraceStatusFlag; +}; + +// Transport 内存类型 +enum class HcclAiRMAMemType : uint32_t { + LOCAL_INPUT = 0, + REMOTE_INPUT, + + LOCAL_OUTPUT, + REMOTE_OUTPUT, + + // 可透传更多的内存，可在MAX_NUM之前追加，例如： + // LOCAL_EXP, + // REMOTE_EXP, + MAX_NUM +}; + +// Transport 内存信息 +struct HcclAiRMAMemInfo { + uint32_t memMaxNum{0}; // 最大内存数量，等于 HcclAiRMAMemType::MAX_NUM + uint32_t sizeOfMemDetails{0}; // sizeof(MemDetails)，用于内存校验和偏移计算 + uint64_t memDetailPtr{0}; // MemDetails数组首地址, 个数: HcclAiRMAMemType::MAX_NUM + // 可往后追加字段 +}; + +// 全部 Transport QP/Mem 信息 +struct HcclAiRMAInfo { + uint32_t curRankId{0}; // 当前rankId + uint32_t rankNum{0}; // rank数量 + uint32_t qpNum{0}; // 单个Transport的QP数量 + + uint32_t sizeOfAiRMAWQ{0}; // sizeof(HcclAiRMAWQ) + uint32_t sizeOfAiRMACQ{0}; // sizeof(HcclAiRMACQ) + uint32_t sizeOfAiRMAMem{0}; // sizeof(HcclAiRMAMemInfo) + + // HcclAiRMAWQ二维数组首地址 + // QP个数: rankNum * qpNum + // 计算偏移获取SQ指针：sqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMAWQ + // 0 <= qpIndex < qpNum + uint64_t sqPtr{0}; + + // HcclAiRMACQ二维数组首地址 + // QP个数: rankNum * qpNum + // 计算偏移获取SCQ指针：scqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMACQ + // 0 <= qpIndex < qpNum + uint64_t scqPtr{0}; + + // HcclAiRMAWQ二维数组首地址 + // QP个数: rankNum * qpNum + // 计算偏移获取RQ指针：rqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMAWQ + // 0 <= qpIndex < qpNum + uint64_t rqPtr{0}; + + // HcclAiRMACQ二维数组首地址 + // QP个数: rankNum * qpNum + // 计算偏移获取RCQ指针: rcqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMACQ + // 0 <= qpIndex < qpNum + uint64_t rcqPtr{0}; + + // HcclAivMemInfo一维数组 + // 内存信息个数: rankNum + // 计算偏移获取内存信息指针: memPtr + rankId * sizeOfAiRMAMem + // srcRankId 获取自身内存信息，dstRankId 获取 Transport 内存信息 + uint64_t memPtr{0}; + // 可往后追加字段 +}; +struct CombinedCapability { + uint64_t dataplaneModeBitmap; +}; + +struct HcclA2CombineOpParam { + uint64_t workSpace; // Address for communication between client and server, + // hccl requests and clears + uint64_t workSpaceSize; // Space for communication between client and server + uint32_t rankId; // id of this rank + uint32_t rankNum; // num of ranks in this comm group + uint64_t winSize; // size of each windows memory + uint64_t windowsIn[AscendC::HCCL_MAX_RANK_NUM]; // windows address for input, windowsIn[rankId] corresponds + // to the local card address, + // and others are cross-card mapping addresses. + uint64_t windowsOut[AscendC::HCCL_MAX_RANK_NUM]; // windows address for output, windowsOut[rankId] corresponds + // to the local card address, + // and others are cross-card mapping addresses. + uint8_t res[8328]; + uint8_t multiFlag; + __gm__ AscendC::IbVerbsData *data; + uint64_t dataSize; + // 追加字段 + uint64_t sizeOfAiRMAInfo; // sizeof(HcclAiRMAInfo) + uint64_t aiRMAInfo; // HcclAiRMAInfo* 单个结构体指针 + + CombinedCapability* capability; // address of the communication capability information structure on the Device + uint64_t capabilitySize; // size of the communication capability information structure +}; +enum class DataplaneMode : uint32_t { + HOST = 0, + AICPU = 1, + AIV = 2, +}; + +enum class DBMode : int32_t { + INVALID_DB = -1, + HW_DB = 0, + SW_DB +}; + +struct HcclAiRMAWQ { + uint32_t wqn{0}; + uint64_t bufAddr{0}; + uint32_t wqeSize{0}; + uint32_t depth{0}; + uint64_t headAddr{0}; + uint64_t tailAddr{0}; + DBMode dbMode{DBMode::INVALID_DB}; // 0-hw/1-sw + uint64_t dbAddr{0}; + uint32_t sl{0}; +}; + +struct HcclAiRMACQ { + uint32_t cqn{0}; + uint64_t bufAddr{0}; + uint32_t cqeSize{0}; + uint32_t depth{0}; + uint64_t headAddr{0}; + uint64_t tailAddr{0}; + DBMode dbMode{DBMode::INVALID_DB}; // 0-hw/1-sw + uint64_t dbAddr{0}; +}; + +struct hns_roce_rc_sq_wqe { + uint32_t byte_4; + uint32_t msg_len; + uint32_t immtdata; + uint32_t byte_16; + uint32_t byte_20; + uint32_t rkey; + uint64_t remoteVA; +}; + + +struct hns_roce_lite_wqe_data_seg { + uint32_t len; + uint32_t lkey; + uint64_t localVA; +}; + +__aicore__ inline void cacheWriteThrough(__gm__ uint8_t* sourceAddr, uint64_t length) { + __gm__ uint8_t* start = + (__gm__ uint8_t*)((uint64_t)sourceAddr / AscendC::CACHE_LINE_SIZE * AscendC::CACHE_LINE_SIZE); + __gm__ uint8_t* end = + (__gm__ uint8_t*)(((uint64_t)sourceAddr + length) / AscendC::CACHE_LINE_SIZE * AscendC::CACHE_LINE_SIZE); + AscendC::GlobalTensor global; + global.SetGlobalBuffer(start); + for (uint32_t i = 0; i <= end - start; i += AscendC::CACHE_LINE_SIZE) { + AscendC::DataCacheCleanAndInvalid(global[i]); + } +} +__aicore__ inline DataplaneMode GetDataplaneMode(GM_ADDR contextGM0) { + __gm__ HcclA2CombineOpParam *winContext_ = (__gm__ HcclA2CombineOpParam *)contextGM0; + CombinedCapability* capability = winContext_->capability; + uint64_t capabilitySize = winContext_->capabilitySize; + DataplaneMode dataplaneMode = DataplaneMode::AICPU; + if (capability == 0) { + return dataplaneMode; + } + uint64_t dataplaneModeBitmap = capability->dataplaneModeBitmap; + if ((dataplaneModeBitmap & 0x04) == 0x04) { + dataplaneMode = DataplaneMode::AIV; + } + return dataplaneMode; +} + +__aicore__ inline int64_t GetCurrentTimestampUs() +{ + return AscendC::GetSystemCycle() / TIME_CYCLE; +} + +__aicore__ inline void RecordRankCommDuration(AscendC::LocalTensor performanceInfoU32Tensor, uint32_t rankId, int64_t startTime) +{ + int64_t endTime = GetCurrentTimestampUs(); + int32_t duration = static_cast(endTime - startTime); // int32_t可以表示2^31(us)，约35min在实际场景下满足需要 + performanceInfoU32Tensor.SetValue(rankId * sizeof(int64_t) / sizeof(int32_t), duration); // 使用int32_t是因为atomicAdd不支持int64_t类型，这里只赋值到int64_t的低32位。 +} +#endif // MOE_DISTRIBUTE_BASE_H \ No newline at end of file diff --git a/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/select_helper.hpp b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/select_helper.hpp new file mode 100644 index 000000000..574ab3351 --- /dev/null +++ b/csrc/ascend/mc2/dispatch_ffn_combine/op_kernel/utils/select_helper.hpp @@ -0,0 +1,25 @@ +#ifndef SELECT_HELPER_HPP +#define SELECT_HELPER_HPP + +#include "catlass/layout/layout.hpp" +using namespace AscendC; +using namespace Catlass; + +template +struct LayoutBInitializer { + CATLASS_DEVICE + static Layout create(uint32_t k, uint32_t n) { + return Layout{k, n}; + } +}; + +template +struct LayoutBInitializer> +> { + CATLASS_DEVICE + static Layout create(uint32_t k, uint32_t n) { + return Layout::template MakeLayout(k, n); + } +}; +#endif // SELECT_HELPER_HPP \ No newline at end of file diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/CMakeLists.txt b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/CMakeLists.txt new file mode 100644 index 000000000..298ed5677 --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/matmul_allreduce_add_rmsnorm_torch_adpt.h b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/matmul_allreduce_add_rmsnorm_torch_adpt.h new file mode 100644 index 000000000..102b02790 --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/matmul_allreduce_add_rmsnorm_torch_adpt.h @@ -0,0 +1,51 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef MATMUL_ALLREDUCE_ADD_RMSNORM_TORCH_ADPT_H +#define MATMUL_ALLREDUCE_ADD_RMSNORM_TORCH_ADPT_H +namespace vllm_fl { + +std::tuple matmul_allreduce_add_rmsnorm( + const at::Tensor &x1, + const at::Tensor &x2, + const at::Tensor &residual, + const at::Tensor &gamma, + c10::string_view group_tp, + int64_t tp_rank_size, + int64_t tp_rank_id, + double epsilon, + bool is_trans_b, + bool is_gather_add_out) + { + at::Tensor output = at::empty_like(residual); + at::Tensor add_out = at::empty_like(residual); + + std::string group_tp_str(group_tp); + + char *group_tp_ptr = group_tp_str.data(); + + float epsilon_f = static_cast(epsilon); + EXEC_NPU_CMD(aclnnMatmulAllreduceAddRmsnorm, + // input + x1, x2, residual, gamma, + // attr + group_tp_ptr, tp_rank_size, tp_rank_id, epsilon_f, is_trans_b, is_gather_add_out, + // output + output, add_out); + + return {output, add_out}; + } +} +#endif \ No newline at end of file diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/CMakeLists.txt b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/CMakeLists.txt new file mode 100644 index 000000000..7bf18ab75 --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/CMakeLists.txt @@ -0,0 +1,22 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnInner PRIVATE + matmul_allreduce_add_rmsnorm_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME MatmulAllreduceAddRmsnormTensorList + OPTIONS + --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE matmul_allreduce_add_rmsnorm ACLNNTYPE aclnn_inner) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/matmul_allreduce_add_rmsnorm_def.cpp b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/matmul_allreduce_add_rmsnorm_def.cpp new file mode 100644 index 000000000..9e44dee2c --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/matmul_allreduce_add_rmsnorm_def.cpp @@ -0,0 +1,68 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "register/op_def_registry.h" + +namespace ops{ +class MatmulAllreduceAddRmsnorm : public OpDef { +public: + explicit MatmulAllreduceAddRmsnorm(const char* name) : OpDef(name) + { + this->Input("x1") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("x2") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("residual") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("gamma") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("add_out") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + + this->Attr("group_tp").String(); + this->Attr("tp_rank_size").Int(); + this->Attr("tp_rank_id").Int(); + this->Attr("epsilon").AttrType(OPTIONAL).Float(1e-6); + this->Attr("is_trans_b").AttrType(OPTIONAL).Bool(false); + this->Attr("is_gather_add_out").AttrType(OPTIONAL).Bool(false); + + this->MC2().HcclGroup({"group_tp"}); + this->AICore().AddConfig("ascend910b"); + } +}; + +OP_ADD(MatmulAllreduceAddRmsnorm); +} \ No newline at end of file diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/matmul_allreduce_add_rmsnorm_proto.cpp b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/matmul_allreduce_add_rmsnorm_proto.cpp new file mode 100644 index 000000000..027ff88fc --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/matmul_allreduce_add_rmsnorm_proto.cpp @@ -0,0 +1,68 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include "graph/utils/type_utils.h" +#include "register/op_def_registry.h" + +namespace ge { + constexpr uint32_t RESIDUAL_INDEX = 3; + constexpr uint32_t OUTPUT_Y_INDEX = 0; + constexpr uint32_t OUTPUT_ADD_OUT_INDEX = 1; + constexpr int SHAPE_INDEX0 = 0; + constexpr int SHAPE_INDEX1 = 1; + constexpr int SHAPE_INDEX2 = 2; + constexpr int DIM_NUM_2 = 2; + constexpr int DIM_NUM_3 = 3; + +static void CloneShape(const gert::Shape* src, gert::Shape* dst) +{ + int ndim = src->GetDimNum(); + dst->SetDimNum(ndim); + for (int i = 0; i < ndim; ++i) { + dst->SetDim(i, src->GetDim(i)); + } +} + +static ge::graphStatus InferShape(gert::InferShapeContext* context) +{ + const gert::Shape* residualShape = context->GetInputShape(RESIDUAL_INDEX); + int residualDimNum = residualShape->GetDimNum(); + + if (residualDimNum != DIM_NUM_2 && residualDimNum != DIM_NUM_3) { + return GRAPH_FAILED; + } + + gert::Shape* x1OutShape = context->GetOutputShape(OUTPUT_Y_INDEX); + gert::Shape* addOutShape = context->GetOutputShape(OUTPUT_ADD_OUT_INDEX); + CloneShape(residualShape, x1OutShape); + CloneShape(residualShape, addOutShape); + + return GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataType(gert::InferDataTypeContext *context) +{ + const auto residualDataType = context->GetInputDataType(RESIDUAL_INDEX); + context->SetOutputDataType(OUTPUT_Y_INDEX, residualDataType); + context->SetOutputDataType(OUTPUT_ADD_OUT_INDEX, residualDataType); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP(MatmulAllreduceAddRmsnorm) + .InferShape(InferShape) + .InferDataType(InferDataType); +} \ No newline at end of file diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/matmul_allreduce_add_rmsnorm_tiling.cpp b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/matmul_allreduce_add_rmsnorm_tiling.cpp new file mode 100644 index 000000000..6a92e402a --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/matmul_allreduce_add_rmsnorm_tiling.cpp @@ -0,0 +1,619 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include + +#include "log/ops_log.h" +#include "error/ops_error.h" + +#include "graph/utils/type_utils.h" +#include "register/op_def_registry.h" +#include "tiling/tiling_api.h" +#include "../op_kernel/matmul_allreduce_add_rmsnorm_tiling.h" +#include "matmul_allreduce_add_rmsnorm_workspace.h" +#include "tiling/platform/platform_ascendc.h" +#include "tiling/hccl/hccl_tiling.h" + +typedef enum { + ATTR_TP_INDEX = 0, + ATTR_RANK_SIZE_INDEX, + ATTR_RANK_ID_INDEX, + ATTR_EPSILON_INDEX, + ATTR_IS_TRANS_B_INDEX, + ATTR_IS_GATHER_ADD_OUT_INDEX +} ATTR_TYPE; + +int32_t CeilDev(int32_t num, int32_t div) +{ + if (div == 0) { + return 0; + } + return (num + div - 1) / div; +} +static constexpr uint32_t OP_TYPE_ALL_TO_ALL = 8; +static constexpr uint32_t BATCH_SIZE_ONE = 1; +static constexpr uint32_t DEFAULT_ROW = 128; +static constexpr uint32_t DEFAULT_COL = 256; +static constexpr uint32_t DEFAULT_SWIZZLE_COUNT = 4; +static constexpr int32_t VALID_UB_MOVE_NUM = 20480; +static constexpr int32_t COMMDATASPLIT_ONE = 1; +static constexpr int32_t COMM_DATA_DIRECT = 0; +static constexpr uint32_t ALLREDUCE_EIGHT_RANK_FP16_M0_DEFAULT = 128; +static constexpr int32_t ALLREDUCE_EIGHT_RANK_FP16_DATASPLIT_DEFAULT = 16; +static constexpr int32_t ALLREDUCE_EIGHT_RANK_FP16_UBMOVENUM_DEFAULT = 100; +static constexpr int32_t HALF_KBYTE = 512; +static constexpr int32_t ALLREDUCE_EIGHT_RANK_FP16_PVALUE_DEFAULT = 14; +static constexpr int32_t SWIZZLE_DIRECT_ONE = 1; +static constexpr int32_t COMMNPUSPLIT_ONE = 1; +static constexpr int32_t COMMDATASPLIT_SIXTEEN = 16; +constexpr int32_t SECOND_TO_MS = 1000; +constexpr int64_t MATMUL_BASE_100US = static_cast(1024) * 8192 * 1024; +constexpr int64_t ALLREDUCE_BASE_100US = 4096 * 1024; +constexpr double ONE_K = 1024.0; +constexpr double B1_FLOP_PER_MS = (364 * 0.8) * 1e9; +constexpr double DOUBLE = 2.0; +constexpr double HALF_PROB = 0.5; +constexpr int32_t CONDITION_M_ST = 0; +constexpr int32_t CONDITION_M_END = 1; +constexpr int32_t CONDITION_K_ST = 2; +constexpr int32_t CONDITION_K_END = 3; +constexpr int32_t CONDITION_N_ST = 4; +constexpr int32_t CONDITION_N_END = 5; +constexpr int32_t RANKSIZE_FOUR = 4; +constexpr int32_t RANKSIZE_EIGHT = 8; +constexpr int32_t DIV_TWO = 2; +constexpr int32_t LENPERLOOP_DEFAULT = 5120; +constexpr int32_t MIN_UB_MOVE_NUM = 5120; +constexpr int32_t MAX_UB_NUM = 97280; // 190 * 1024 / 2 +constexpr int32_t MAX_P_VALUE = 15; + +constexpr int32_t DIM_NUM_TWO = 2; +constexpr int32_t DIM_NUM_THREE = 3; +constexpr int32_t DIM_INDEX_ZERO = 0; +constexpr int32_t DIM_INDEX_ONE = 1; +constexpr int32_t DIM_INDEX_TWO = 2; + +static constexpr uint32_t SYSTEM_NEED_WORKSPACE = 16 * 1024 * 1024; + +static constexpr uint32_t USE_CORE_NUM = 20; + +static std::vector ALLREDUCE_UBMOVENUM_COEF = {{-1.72352427e+01, + 2.56887672e-03, + -8.21819480e+00, + 8.70965589e+01, + -3.63853858e-01, + 1.27789264e+01, + 1.29782183e+02, + 1.90250023e-02, + -3.48175441e+00, + 6.18921914e+03, + 3.77072171e+03, + -5.86895290e+01, + -8.70740991e-01, + -1.40262280e-04, + -2.81910331e-08, + 3.22795486e-05, + -4.84522320e-03, + 2.94839177e-01, + 2.97260958e-03, + 9.08844709e+01, + -5.80426209e-10, + 38.183465184603484}}; + +static std::map>> ALLREDUCE_EIGHT_RANK_FP16_M0_MAP = { + {128, + {{-1, 31220, -1, 2147483647, -1, 768}, + {31220, 36980, 1280, 2147483647, -1, 768}, + {36980, 2147483647, -1, 2147483647, -1, 768}, + {-1, 2147483647, -1, 2147483647, 768, 2147483647}}}, + {256, {{31220, 36980, -1, 1280, -1, 768}}}}; + +static std::map>> ALLREDUCE_EIGHT_RANK_FP16_UBMOVENUM_MAP = { + {100, + {{-1, 3072, -1, 2147483647, -1, 768}, + {3072, 19680, -1, 3072, -1, 768}, + {-1, 3072, -1, 2147483647, 768, 1536}, + {3072, 19680, -1, 3072, 768, 1536}, + {-1, 2147483647, 1792, 2976, 1536, 13312}}}, + {30, + {{3072, 19680, 3072, 2147483647, -1, 768}, + {19680, 2147483647, -1, 3072, -1, 1536}, + {-1, 2147483647, -1, 1792, 1536, 13312}, + {-1, 768, 2976, 2147483647, 5376, 13312}, + {-1, 768, -1, 2147483647, 13312, 2147483647}, + {26880, 2147483647, -1, 3072, 13312, 2147483647}}}, + {20, + {{3072, 19680, 3072, 2147483647, 768, 1536}, + {19680, 2147483647, 3072, 2147483647, -1, 1536}, + {-1, 2147483647, 2976, 2147483647, 1536, 5376}, + {768, 2147483647, 2976, 2147483647, 5376, 13312}, + {768, 26880, -1, 2147483647, 13312, 2147483647}, + {26880, 2147483647, 3072, 2147483647, 13312, 2147483647}}}}; + +static std::vector ALLREDUCE_PVALUE_COEF = {{-4.23166350e+00, + 6.71137487e-04, + -1.33434156e+00, + 1.12915884e+01, + -7.85892737e-02, + 2.59059897e+00, + 3.22129881e+01, + -5.15776887e-02, + 9.15542742e-01, + 1.56322201e+03, + 3.61977421e+01, + -5.49544589e-01, + -2.66903417e-01, + -3.68521920e-05, + -6.40666333e-09, + 6.77406054e-06, + -9.92992099e-04, + 5.60658043e-02, + 2.69372863e-04, + 2.17222337e+01, + -1.17749660e-10, + 6.100544547671263}}; + +double GetMTETime(double mknGB, int32_t m0, int32_t n0, double aBindWidth = 3.0, double bBindWidth = 3.0) +{ + // 预估Matmul计算的MTE2搬运时间 + return DOUBLE * mknGB * (SECOND_TO_MS / ONE_K) * (1.0 / (n0 * aBindWidth) + 1.0 / (m0 * bBindWidth)); +} + +int32_t AllReduceUbMoveNum(int m, int k, int n) +{ + double commPredict = 1.0 * (m / ONE_K) * (n / ONE_K) * (SECOND_TO_MS / ONE_K) / 40; + double cubePredict = DOUBLE * m * k / B1_FLOP_PER_MS * n; + double mknGB = (m / ONE_K) * (k / ONE_K) * (n / ONE_K); + double mteTimePredict1 = GetMTETime(mknGB, DEFAULT_ROW, DEFAULT_COL); + double mteTimePredict2 = GetMTETime(mknGB, DEFAULT_COL, DEFAULT_ROW); + double mteTimePredict = std::min(mteTimePredict1, mteTimePredict2); + double matmulPredict = std::max(cubePredict, mteTimePredict); + double c0 = matmulPredict / commPredict; + double c1 = 1.0 * m * n / k; + double c2 = sqrt(c1); + double c3 = sqrt(1.0 * m * n) / k; + double c4 = c3 * c3; + double c5 = matmulPredict; + double c6 = commPredict; + double c7 = 1.0 * n / m; + double c8 = 1.0 * m * n / sqrt(k); + double c9 = 1.0 * m * n * sqrt(k); + double c10 = sqrt(1.0 * m * n) * k; + double c11 = sqrt(1.0 * m * n * k); + double c12 = sqrt(1.0 * m * n); + double c13 = 1.0 * k * k / sqrt(1.0 * m * n); + double c14 = 1.0 * k * k * sqrt(1.0 * m * n); + double ubMoveNumDouble = 0; + std::vector feats_update = {c0, + c1, + c2, + c3, + c4, + c5, + c6, + c7, + 1.0 / c0, + 1.0 / c1, + 1.0 / c2, + 1.0 / c3, + 1.0 / c4, + c8, + c9, + c10, + c11, + c12, + c13, + 1.0 / c13, + c14, + 1}; + for (uint32_t i = 0; i < feats_update.size(); i++) { + ubMoveNumDouble += feats_update[i] * ALLREDUCE_UBMOVENUM_COEF[i]; + } + + return std::min(std::max(static_cast(ubMoveNumDouble) * HALF_KBYTE, MIN_UB_MOVE_NUM), MAX_UB_NUM); +} + +int32_t AllReducePValue(int m, int k, int n) +{ + double commPredict = 1.0 * (m / ONE_K) * (n / ONE_K) * (SECOND_TO_MS / ONE_K) / 40; + double cubePredict = DOUBLE * m * k / B1_FLOP_PER_MS * n; + double mknGB = (m / ONE_K) * (k / ONE_K) * (n / ONE_K); + double mteTimePredict1 = GetMTETime(mknGB, DEFAULT_ROW, DEFAULT_COL); + double mteTimePredict2 = GetMTETime(mknGB, DEFAULT_COL, DEFAULT_ROW); + double mteTimePredict = std::min(mteTimePredict1, mteTimePredict2); + double matmulPredict = std::max(cubePredict, mteTimePredict); + double c0 = matmulPredict / commPredict; + double c1 = 1.0 * m * n / k; + double c2 = sqrt(c1); + double c3 = sqrt(1.0 * m * n) / k; + double c4 = c3 * c3; + double c5 = matmulPredict; + double c6 = commPredict; + double c7 = 1.0 * n / m; + double c8 = 1.0 * m * n / sqrt(k); + double c9 = 1.0 * m * n * sqrt(k); + double c10 = sqrt(1.0 * m * n) * k; + double c11 = sqrt(1.0 * m * n * k); + double c12 = sqrt(1.0 * m * n); + double c13 = 1.0 * k * k / sqrt(1.0 * m * n); + double c14 = 1.0 * k * k * sqrt(1.0 * m * n); + double pValueDouble = 0; + std::vector feats_update = {c0, + c1, + c2, + c3, + c4, + c5, + c6, + c7, + 1.0 / c0, + 1.0 / c1, + 1.0 / c2, + 1.0 / c3, + 1.0 / c4, + c8, + c9, + c10, + c11, + c12, + c13, + 1.0 / c13, + c14, + 1}; + for (uint32_t i = 0; i < feats_update.size(); i++) { + pValueDouble += feats_update[i] * ALLREDUCE_PVALUE_COEF[i]; + } + + return std::min(std::max(static_cast(pValueDouble), 1), MAX_P_VALUE); +} + +int32_t GetValueFromMKNConditionMap( + int32_t m, int32_t k, int32_t n, int32_t defaultValue, std::map>> conditionMap) +{ + int32_t value = defaultValue; + for (auto &item : conditionMap) { + for (auto &condition : item.second) { + bool inRange = m > condition[CONDITION_M_ST] && m <= condition[CONDITION_M_END] && + k > condition[CONDITION_K_ST] && k <= condition[CONDITION_K_END] && + n > condition[CONDITION_N_ST] && n <= condition[CONDITION_N_END]; + if (inRange) { + return item.first; + } + } + } + return value; +} + +void AllReduceEightRankFP16GetDefaultTiling( + gert::TilingContext *context, PPTilingData &ppTilingData, CommTilingData &commTilingData) +{ + int32_t m = ppTilingData.opShape.m; + int32_t k = ppTilingData.opShape.k; + int32_t n = ppTilingData.opShape.n; + + ppTilingData.m0 = + GetValueFromMKNConditionMap(m, k, n, ALLREDUCE_EIGHT_RANK_FP16_M0_DEFAULT, ALLREDUCE_EIGHT_RANK_FP16_M0_MAP); + + ppTilingData.k0 = DEFAULT_COL; + ppTilingData.n0 = ppTilingData.m0 == DEFAULT_ROW ? DEFAULT_COL : DEFAULT_ROW; + + ppTilingData.mLoop = CeilDev(m, ppTilingData.m0); + ppTilingData.nLoop = CeilDev(n, ppTilingData.n0); + ppTilingData.kLoop = CeilDev(k, ppTilingData.k0); + + ppTilingData.coreLoop = ppTilingData.opShape.batchSize * ppTilingData.mLoop * ppTilingData.nLoop; + ppTilingData.swizzlDirect = SWIZZLE_DIRECT_ONE; + ppTilingData.swizzlCount = DEFAULT_SWIZZLE_COUNT; + ppTilingData.tilingKey = 0; + ppTilingData.splitK = 0; + + uint32_t blockDim = 1U; + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + ppTilingData.blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, 0, aivNum); + + commTilingData.ubMoveNum = + GetValueFromMKNConditionMap( + m, k, n, ALLREDUCE_EIGHT_RANK_FP16_UBMOVENUM_DEFAULT, ALLREDUCE_EIGHT_RANK_FP16_UBMOVENUM_MAP) * + HALF_KBYTE; + commTilingData.pValue = ALLREDUCE_EIGHT_RANK_FP16_PVALUE_DEFAULT; + + commTilingData.commDirect = COMM_DATA_DIRECT; + commTilingData.commNpuSplit = COMMNPUSPLIT_ONE; + commTilingData.commDataSplit = COMMDATASPLIT_SIXTEEN; + commTilingData.is91093 = 0; + commTilingData.withSerialMode = 0; + commTilingData.tag = 0; + commTilingData.write2OtherRank = 0; +} + +void GetDefaultTiling(gert::TilingContext *context, PPTilingData &ppTilingData, CommTilingData &commTilingData) +{ + int32_t m = ppTilingData.opShape.m; + int32_t k = ppTilingData.opShape.k; + int32_t n = ppTilingData.opShape.n; + + ppTilingData.m0 = DEFAULT_ROW; + ppTilingData.n0 = DEFAULT_COL; + ppTilingData.k0 = DEFAULT_COL; + + ppTilingData.mLoop = CeilDev(m, ppTilingData.m0); + ppTilingData.nLoop = CeilDev(n, ppTilingData.n0); + ppTilingData.kLoop = CeilDev(k, ppTilingData.k0); + ppTilingData.coreLoop = ppTilingData.opShape.batchSize * ppTilingData.mLoop * ppTilingData.nLoop; + + ppTilingData.swizzlDirect = m > n ? 0 : 1; + ppTilingData.swizzlCount = DEFAULT_SWIZZLE_COUNT; + ppTilingData.tilingKey = 0; + ppTilingData.splitK = 0; + + uint32_t blockDim = 1U; + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); + ppTilingData.blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, 0, aivNum); + + commTilingData.ubMoveNum = AllReduceUbMoveNum(m, k, n); + commTilingData.pValue = AllReducePValue(m, k, n); + commTilingData.commNpuSplit = commTilingData.rankSize; + commTilingData.commDataSplit = COMMDATASPLIT_ONE; + commTilingData.commDirect = COMM_DATA_DIRECT; + commTilingData.lenPerLoop = ppTilingData.m0 * ppTilingData.n0 * commTilingData.pValue * ppTilingData.blockDim; + commTilingData.lenPerLoop = commTilingData.lenPerLoop / commTilingData.rankSize; + commTilingData.is91093 = 0; + commTilingData.withSerialMode = 0; + commTilingData.tag = 0; + commTilingData.write2OtherRank = 0; +} + +static inline void GetRmsnormTilingData(RmsNormTilingData &rmsnormtiling, std::vector &shapeVec, + std::vector &oriShapeVec, uint32_t calcBytes = 0, uint32_t loopCount = 1, float ep = 1e-5) +{ + ge::Shape srcShape(shapeVec); + ge::Shape oriSrcShape(oriShapeVec); + uint32_t minValue = 0; + uint32_t maxValue = 0; + AscendC::GetRmsNormMaxMinTmpSize(srcShape, sizeof(uint16_t), maxValue, minValue, false); + + if (calcBytes < minValue) { + rmsnormtiling.calcBytes = minValue; + } else if (calcBytes > maxValue) { + rmsnormtiling.calcBytes = maxValue; + } else { + rmsnormtiling.calcBytes = calcBytes; + } + + optiling::RmsNormTiling tilingdata; + AscendC::GetRmsNormTilingInfo(srcShape, oriSrcShape, rmsnormtiling.calcBytes, sizeof(uint16_t), tilingdata, false); + size_t tilingSize = tilingdata.GetDataSize(); + tilingdata.SaveToBuffer(&rmsnormtiling.tiling, tilingSize); + rmsnormtiling.epsilon = ep; + rmsnormtiling.loopCount = loopCount; +} + +static inline void GetQuantTilingData(QuantInfo &quantInfo) +{ + quantInfo.dequantGranularity = QuantGranularity::QUANT_GRANULARITY_UNDEFINED; + quantInfo.dequantGroupSize = -1; + quantInfo.quantGranularity = QuantGranularity::QUANT_GRANULARITY_UNDEFINED; + quantInfo.quantGroupSize = -1; +} + +static ge::graphStatus GetAttrAndSetTilingData( + gert::TilingContext *context, const char *nodeName, MatmulAllreduceAddRmsnormTilingData &tilingData) + +{ + CommTilingData &commTilingData = tilingData.matmulAllreduceAddRmsnormInfo.commTilingData; + PPTilingData &ppTilingData = tilingData.matmulAllreduceAddRmsnormInfo.ppTilingData; + RmsNormTilingData &rmsnormTilingData = tilingData.matmulAllreduceAddRmsnormInfo.rmsnormTilingData; + QuantInfo &quantInfo = tilingData.matmulAllreduceAddRmsnormInfo.quantInfo; + + auto attrs = context->GetAttrs(); + OPS_ERR_IF(attrs == nullptr, OPS_LOG_E(nodeName, "attrs is nullptr."), return ge::GRAPH_FAILED); + + auto RankSizePtr = attrs->GetAttrPointer(ATTR_RANK_SIZE_INDEX); + auto RankIdPtr = attrs->GetAttrPointer(ATTR_RANK_ID_INDEX); + + bool isTransB = *(attrs->GetAttrPointer(ATTR_IS_TRANS_B_INDEX)); + + ppTilingData.isTransA = false; + ppTilingData.isTransB = isTransB; + ppTilingData.isGatherAddOut = *(attrs->GetAttrPointer(ATTR_IS_GATHER_ADD_OUT_INDEX)); + + auto &opShape = ppTilingData.opShape; + auto &tensor0Shape = context->GetInputTensor(0)->GetOriginShape(); + uint32_t dimNum = tensor0Shape.GetDimNum(); + int64_t bs; + int64_t rankM; + int64_t rankK; + + if (dimNum == DIM_NUM_THREE) { + bs = tensor0Shape.GetDim(DIM_INDEX_ZERO); + rankM = tensor0Shape.GetDim(DIM_INDEX_ONE); + rankK = tensor0Shape.GetDim(DIM_INDEX_TWO); + } else if (dimNum == DIM_NUM_TWO) { + bs = BATCH_SIZE_ONE; + rankM = tensor0Shape.GetDim(DIM_INDEX_ZERO); + rankK = tensor0Shape.GetDim(DIM_INDEX_ONE); + } else { + const char *nodeName = context->GetNodeName(); + OPS_LOG_E(nodeName, "Tiling input dim error."); + return ge::GRAPH_FAILED; + } + + int64_t rankN = isTransB ? + context->GetInputTensor(1)->GetOriginShape().GetDim(DIM_INDEX_ZERO) : + context->GetInputTensor(1)->GetOriginShape().GetDim(DIM_INDEX_ONE); + + opShape.batchSize = BATCH_SIZE_ONE; + opShape.m = bs * rankM; + opShape.n = rankN; + opShape.k = rankK; + + commTilingData.rankSize = static_cast(*RankSizePtr); + commTilingData.rank = static_cast(*RankIdPtr); + if (commTilingData.rankSize == RANKSIZE_EIGHT) { + AllReduceEightRankFP16GetDefaultTiling(context, ppTilingData, commTilingData); + } else { + GetDefaultTiling(context, ppTilingData, commTilingData); + } + + uint32_t calcBytes = 0; + uint32_t sLength = 1; + std::vector shapeVec = {1, 1, rankN}; + std::vector oriShapeVec = shapeVec; + auto EpsilonPtr = attrs->GetAttrPointer(ATTR_EPSILON_INDEX); + float epsilon = static_cast(*EpsilonPtr); + GetRmsnormTilingData( + rmsnormTilingData, shapeVec, oriShapeVec, calcBytes, commTilingData.rankSize * sLength * rankN, epsilon); + GetQuantTilingData(quantInfo); + + return ge::GRAPH_SUCCESS; +} + +bool IsMatrixAligned(const int64_t &m, const int64_t &n, const bool &transpose, int nElemAlign) +{ + return (transpose ? m : n) % nElemAlign == 0; +} + +int64_t GetAlignedMatrixSize( + const int64_t &batchSize, const int64_t &m, const int64_t &n, const bool &transpose, int nElemAlign) +{ + int64_t nRow = transpose ? n : m; + int64_t nCol = transpose ? m : n; + int64_t nColAlign = (nCol + nElemAlign - 1) / nElemAlign * nElemAlign; + return batchSize * nRow * nColAlign; +} + +WorkspaceDetail GetWorkspaceDetail(CoCDataTypeDesc dataType, const MatMulInfo &mmInfo, const QuantInfo &quantInfo) +{ + WorkspaceDetail workspaceDetail; + + int32_t eleSize = COC_TYPE2ELE_SIZE.at(dataType); + int32_t nElemAlign = ALIGN_BYTES / eleSize; + + bool hasQuant = quantInfo.quantGranularity != QuantGranularity::QUANT_GRANULARITY_UNDEFINED; + if (hasQuant || (!IsMatrixAligned(mmInfo.m, mmInfo.k, mmInfo.transA, nElemAlign) && mmInfo.m != 1)) { + workspaceDetail.matrixActivationSize = + GetAlignedMatrixSize(mmInfo.batchSize, mmInfo.m, mmInfo.k, mmInfo.transA, nElemAlign) * eleSize; + } + + bool hasDequant = quantInfo.dequantGranularity != QuantGranularity::QUANT_GRANULARITY_UNDEFINED; + if ((hasDequant && !mmInfo.isInt8) || !IsMatrixAligned(mmInfo.k, mmInfo.n, mmInfo.transB, nElemAlign)) { + workspaceDetail.matrixWeightSize = + GetAlignedMatrixSize(mmInfo.batchSize, mmInfo.k, mmInfo.n, mmInfo.transB, nElemAlign) * eleSize; + } + + bool hasAccum = dataType == CoCDataTypeDesc::INT8INT8_INT32_BF16; + if (hasAccum) { + workspaceDetail.matrixIntermediateSize = + static_cast(mmInfo.batchSize) * mmInfo.m * mmInfo.n * sizeof(int32_t); + } + + if (mmInfo.isInt8) { + workspaceDetail.formatDequantParamSize = + mmInfo.k > mmInfo.n ? mmInfo.k * sizeof(float) : mmInfo.n * sizeof(float); + } + return workspaceDetail; +} + +void GetMmInfo(gert::TilingContext *context, MatmulAllreduceAddRmsnormTilingData *tiling, MatMulInfo *mmInfo) +{ + PPTilingData tempPPTilingData = tiling->matmulAllreduceAddRmsnormInfo.ppTilingData; + mmInfo->batchSize = tempPPTilingData.opShape.batchSize; + mmInfo->m = tempPPTilingData.opShape.m; + mmInfo->n = tempPPTilingData.opShape.n; + mmInfo->k = tempPPTilingData.opShape.k; + auto attrs = context->GetAttrs(); + mmInfo->transA = false; + mmInfo->transB = *(attrs->GetAttrPointer(ATTR_IS_TRANS_B_INDEX)); + mmInfo->withBias = false; + mmInfo->weightNz = false; + mmInfo->isInt8 = context->GetInputTensor(0)->GetDataType() == ge::DT_INT8; +} + +size_t GetUserWorkspaceSize(gert::TilingContext *context, MatmulAllreduceAddRmsnormTilingData *tiling) +{ + MatMulInfo mmInfo; + GetMmInfo(context, tiling, &mmInfo); + QuantInfo quantInfo = tiling->matmulAllreduceAddRmsnormInfo.quantInfo; + return GetWorkspaceDetail(FP16FP16_FP32_FP16, mmInfo, quantInfo).GetSize(); +} + +static ge::graphStatus SetWorkSpace(gert::TilingContext *context, const char *nodeName) +{ + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + size_t *workSpaces = context->GetWorkspaceSizes(1); + OPS_ERR_IF(workSpaces == nullptr, OPS_LOG_E(nodeName, "workSpaces is nullptr."), return ge::GRAPH_FAILED); + size_t systemWorkspaceSize = static_cast(ascendcPlatform.GetLibApiWorkSpaceSize()); + MatmulAllreduceAddRmsnormTilingData *tilingData = context->GetTilingData(); + size_t userWorkspaceSize = GetUserWorkspaceSize(context, tilingData); + workSpaces[0] = userWorkspaceSize + systemWorkspaceSize; + return ge::GRAPH_SUCCESS; +} + +static void SetHcommCfg( + const gert::TilingContext *context, MatmulAllreduceAddRmsnormTilingData *tiling, const std::string groupTp) +{ + const char *nodeName = context->GetNodeName(); + uint32_t opType = OP_TYPE_ALL_TO_ALL; + std::string algConfigAllToAllStr = "AlltoAll=level0:fullmesh"; + + AscendC::Mc2CcTilingConfig mc2CcTilingConfig(groupTp, opType, algConfigAllToAllStr); + mc2CcTilingConfig.GetTiling(tiling->mc2InitTiling); + mc2CcTilingConfig.GetTiling(tiling->mc2CcTiling); +} + +static ge::graphStatus MatmulAllreduceAddRmsnormTilingFuncImpl(gert::TilingContext *context) +{ + const char *nodeName = context->GetNodeName(); + MatmulAllreduceAddRmsnormTilingData *tilingData = context->GetTilingData(); + OPS_ERR_IF(tilingData == nullptr, OPS_LOG_E(nodeName, "tilingData is nullptr."), return ge::GRAPH_FAILED); + + OPS_ERR_IF(GetAttrAndSetTilingData(context, nodeName, *tilingData) != ge::GRAPH_SUCCESS, + OPS_LOG_E(nodeName, "Get attr and set tiling data failed."), + return ge::GRAPH_FAILED); + OPS_ERR_IF(SetWorkSpace(context, nodeName) != ge::GRAPH_SUCCESS, + OPS_LOG_E(nodeName, "Tiling set workspace failed."), + return ge::GRAPH_FAILED); + SetHcommCfg(context, tilingData, "hcomms"); + + fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo(); + OPS_LOG_E_IF_NULL(context, platformInfoPtr, return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + uint32_t aicNum_ = ascendcPlatform.GetCoreNumAic(); + context->SetBlockDim(aicNum_); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus MatmulAllreduceAddRmsnormTilingFunc(gert::TilingContext *context) +{ + ge::graphStatus ret = MatmulAllreduceAddRmsnormTilingFuncImpl(context); + return ret; +} + +struct MatmulAllreduceAddRmsnormCompileInfo {}; +ge::graphStatus TilingParseForMatmulAllreduceAddRmsnorm(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(MatmulAllreduceAddRmsnorm) + .Tiling(MatmulAllreduceAddRmsnormTilingFunc) + .TilingParse(TilingParseForMatmulAllreduceAddRmsnorm); diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/matmul_allreduce_add_rmsnorm_workspace.h b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/matmul_allreduce_add_rmsnorm_workspace.h new file mode 100644 index 000000000..7da928f82 --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/matmul_allreduce_add_rmsnorm_workspace.h @@ -0,0 +1,79 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MATMUL_ALLREDUCE_ADD_RMSNORM_WORKSPACE_H +#define MATMUL_ALLREDUCE_ADD_RMSNORM_WORKSPACE_H + +#include + +#pragma once +const constexpr uint32_t ALIGN_BYTES = 512; +const constexpr int32_t INT8_ELE_SIZE = 1; +const constexpr int32_t FP_BF_16_ELE_SIZE = 2; + +enum CoCDataTypeDesc : int { + COC_DATA_TYPE_UNDEFINED = -1, + FP16FP16_FP32_FP16 = 0, + BF16BF16_FP32_BF16 = 1, + INT8INT8_INT32_FP16 = 2, + INT8INT8_INT32_BF16 = 3, + FP16INT8_INT32_FP16 = 4, + BF16INT8_INT32_BF16 = 5, + FP16INT8_FP32_FP16 = 6, + BF16INT8_FP32_BF16 = 7, + FP16INT4_FP32_FP16 = 8, + BF16INT4_FP32_BF16 = 9, + COC_DATA_TYPE_DESC_MAX = 10, +}; + +const std::map COC_TYPE2ELE_SIZE = { + {FP16FP16_FP32_FP16, FP_BF_16_ELE_SIZE}, + {BF16BF16_FP32_BF16, FP_BF_16_ELE_SIZE}, + {INT8INT8_INT32_FP16, INT8_ELE_SIZE}, + {INT8INT8_INT32_BF16, INT8_ELE_SIZE}, + {FP16INT8_INT32_FP16, INT8_ELE_SIZE}, + {BF16INT8_INT32_BF16, INT8_ELE_SIZE}, + {FP16INT8_FP32_FP16, FP_BF_16_ELE_SIZE}, + {BF16INT8_FP32_BF16, FP_BF_16_ELE_SIZE}, + {FP16INT4_FP32_FP16, FP_BF_16_ELE_SIZE}, + {BF16INT4_FP32_BF16, FP_BF_16_ELE_SIZE} +}; + +struct MatMulInfo { + int64_t batchSize = 1; + int64_t m = -1; + int64_t n = -1; + int64_t k = -1; + bool transA = false; + bool transB = false; + bool withBias = false; + bool isInt8 = false; + bool weightNz = false; +}; + +struct WorkspaceDetail { + int64_t matrixActivationSize{0}; + int64_t matrixWeightSize{0}; + int64_t matrixIntermediateSize{0}; + int64_t formatDequantParamSize{0}; + + int64_t GetSize() const + { + return matrixActivationSize + matrixWeightSize + matrixIntermediateSize + formatDequantParamSize; + } +}; + +#endif \ No newline at end of file diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/op_api/aclnn_matmul_allreduce_add_rmsnorm.cpp b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/op_api/aclnn_matmul_allreduce_add_rmsnorm.cpp new file mode 100644 index 000000000..396da512f --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/op_api/aclnn_matmul_allreduce_add_rmsnorm.cpp @@ -0,0 +1,89 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include "graph/types.h" +#include "aclnn/opdev/platform.h" +#include "aclnn_matmul_allreduce_add_rmsnorm.h" + +enum NnopbaseHcclServerType { + NNOPBASE_HCCL_SERVER_TYPE_AICPU = 0, + NNOPBASE_HCCL_SERVER_TYPE_MTE, + NNOPBASE_HCCL_SERVER_TYPE_END +}; +extern "C" void __attribute__((weak)) NnopbaseSetHcclServerType(void *executor, NnopbaseHcclServerType sType); + +#ifdef __cplusplus +extern "C" { +#endif + +extern aclnnStatus aclnnInnerMatmulAllreduceAddRmsnormGetWorkspaceSize( + const aclTensor *x1, + const aclTensor *x2, + const aclTensor *residual, + const aclTensor *gamma, + char *groupTp, + int64_t tpRankSize, + int64_t tpRankId, + double epsilon, + bool isTransB, + bool isGatherAddOut, + const aclTensor *yOut, + const aclTensor *addOutOut, + uint64_t *workspaceSize, + aclOpExecutor **executor); + +extern aclnnStatus aclnnInnerMatmulAllreduceAddRmsnorm( + void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream); + +aclnnStatus aclnnMatmulAllreduceAddRmsnormGetWorkspaceSize( + const aclTensor *x1, + const aclTensor *x2, + const aclTensor *residual, + const aclTensor *gamma, + char *groupTp, + int64_t tpRankSize, + int64_t tpRankId, + double epsilon, + bool isTransB, + bool isGatherAddOut, + const aclTensor *y, + const aclTensor *addOut, + uint64_t *workspaceSize, + aclOpExecutor **executor) +{ + return aclnnInnerMatmulAllreduceAddRmsnormGetWorkspaceSize(x1, x2, residual, + gamma, groupTp, tpRankSize, tpRankId, epsilon, isTransB, isGatherAddOut, y, addOut, workspaceSize, executor); +} + +aclnnStatus aclnnMatmulAllreduceAddRmsnorm( + void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream) +{ + if (NnopbaseSetHcclServerType) { + NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE); + } + return aclnnInnerMatmulAllreduceAddRmsnorm(workspace, workspaceSize, executor, stream); +} + +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/op_api/aclnn_matmul_allreduce_add_rmsnorm.h b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/op_api/aclnn_matmul_allreduce_add_rmsnorm.h new file mode 100644 index 000000000..b2920a24f --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_host/op_api/aclnn_matmul_allreduce_add_rmsnorm.h @@ -0,0 +1,52 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef ACLNN_MATMUL_ALLREDUCE_ADD_RMSNORM +#define ACLNN_MATMUL_ALLREDUCE_ADD_RMSNORM + +#include "aclnn/acl_meta.h" + +#ifdef __cplusplus +extern "C" { +#endif + +__attribute__((visibility("default"))) aclnnStatus aclnnMatmulAllreduceAddRmsnormGetWorkspaceSize( + const aclTensor *x1, + const aclTensor *x2, + const aclTensor *residual, + const aclTensor *gamma, + char *groupTp, + int64_t tpRankSize, + int64_t tpRankId, + double epsilon, + bool isTransB, + bool isGatherAddOut, + const aclTensor *y, + const aclTensor *addOut, + uint64_t *workspaceSize, + aclOpExecutor **executor); + +__attribute__((visibility("default"))) aclnnStatus aclnnMatmulAllreduceAddRmsnorm( + void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm.cpp b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm.cpp new file mode 100644 index 000000000..ef907abf8 --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm.cpp @@ -0,0 +1,53 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "lib/matmul_intf.h" +#include +#include "matmul_allreduce_add_rmsnorm_aic_kernel.h" +#include "matmul_allreduce_add_rmsnorm_aiv_kernel.h" + +extern "C" __global__ __aicore__ void matmul_allreduce_add_rmsnorm( + GM_ADDR x1, GM_ADDR x2, GM_ADDR residual, + GM_ADDR gamma, GM_ADDR y, GM_ADDR add_out, GM_ADDR workspace, GM_ADDR tiling) +{ + REGISTER_TILING_DEFAULT(MatmulAllreduceAddRmsnormTilingData); + GET_TILING_DATA(tiling_data, tiling); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + + Hccl hccl_; + auto tilingData = (__gm__ MatmulAllreduceAddRmsnormTilingData*)tiling; + __gm__ void* mc2InitTiling = (__gm__ void*)(&(tilingData->mc2InitTiling)); + __gm__ void* mc2CcTiling = (__gm__ void*)(&(tilingData->mc2CcTiling)); + auto contextGM0 = AscendC::GetHcclContext(); + + hccl_.Init(contextGM0, mc2InitTiling); + hccl_.SetCcTiling(mc2CcTiling); + + if ASCEND_IS_AIC { + MatmulAllreduceAddRmsnormAicKernel op; + op.Init(x1, x2, residual, gamma, y, workspace, &tiling_data, hccl_); + op.Process(); + return; + } + + if ASCEND_IS_AIV { + MatmulAllreduceAddRmsnormAivKernel op; + + op.Init(x1, x2, residual, gamma, y, add_out, workspace, &tiling_data, hccl_); + op.Process(&tiling_data); + return; + } +} diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm_aic_kernel.h b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm_aic_kernel.h new file mode 100644 index 000000000..b64709c04 --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm_aic_kernel.h @@ -0,0 +1,360 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MATMUL_ALLREDUCE_ADD_RMSNORM_AIC_KERNEL_H +#define MATMUL_ALLREDUCE_ADD_RMSNORM_AIC_KERNEL_H + +#define ASCENDC_CUBE_ONLY +#define CATLASS_ARCH 2201 + +#include "catlass/catlass.hpp" +#include "catlass/arch/arch.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "catlass/gemm/block/block_swizzle.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/kernel/basic_matmul.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" + +#include "matmul_allreduce_add_rmsnorm_utils.h" +#include "matmul_allreduce_add_rmsnorm_tiling.h" + +constexpr int32_t SCALE_L1_SIZE_A = 256 * 8; +constexpr int32_t SCALE_L1_SIZE_B = 128 * 1024; +constexpr int32_t CUBE_MATRIX_SIZE_B16 = 256; // 16 * 16 +constexpr int32_t CUBE_MATRIX_SIZE_B8 = 16 * 32; // 16 * 32 +constexpr int32_t SCALE_L1_SIZE = 256 * 8; // 2 KB +constexpr int32_t BLOCK_SIZE_16 = 16; +constexpr int32_t BLOCK_SIZE_32 = 32; +constexpr int32_t DOUBLE_BUFFER_SIZE = 2; +constexpr uint32_t MM_L1_TILE_SHAPE_M = 128; +constexpr uint32_t MM_L1_TILE_SHAPE_N = 256; +constexpr uint32_t MM_L1_TILE_SHAPE_K = 256; +constexpr uint32_t MM_L0_TILE_SHAPE_M = MM_L1_TILE_SHAPE_M; +constexpr uint32_t MM_L0_TILE_SHAPE_N = MM_L1_TILE_SHAPE_N; +constexpr uint32_t MM_L0_TILE_SHAPE_K = 64; + +using namespace Catlass; + +template +struct GetAccumType { + using T = float; +}; + +__aicore__ inline bool IsQuant(const QuantGranularity &granularity) +{ + return (granularity > QuantGranularity::QUANT_GRANULARITY_UNDEFINED) && + (granularity < QuantGranularity::QUANT_GRANULARITY_MAX); +} + +template +class MatmulAllreduceAddRmsnormAicKernel { + using T_ACCUM = typename GetAccumType::T; +public: + int PIPE_DEPTH = 2; + Arch::Resource resource; + __aicore__ inline MatmulAllreduceAddRmsnormAicKernel() { } + + __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR residual, GM_ADDR gamma, GM_ADDR y, + GM_ADDR workspace, const MatmulAllreduceAddRmsnormTilingData* tilingData, + Hccl &hccl_) + { + this->hccl_ = hccl_; + this->gm_c = reinterpret_cast<__gm__ OutDtype *>(y); + + this->gm_dequant_scale = nullptr; + this->has_offset = false; + + auto ppTilingData = &tilingData->matmulAllreduceAddRmsnormInfo.ppTilingData; + auto commTilingData = &tilingData->matmulAllreduceAddRmsnormInfo.commTilingData; + auto quantInfo = &tilingData->matmulAllreduceAddRmsnormInfo.quantInfo; + + this->batch_size = ppTilingData->opShape.batchSize; + this->m = ppTilingData->opShape.m; + this->k = ppTilingData->opShape.k; + this->n = ppTilingData->opShape.n; + this->weight_nz = false; + + this->is_int8 = false; + this->cube_matrix_size = this->is_int8 ? CUBE_MATRIX_SIZE_B8 : CUBE_MATRIX_SIZE_B16; + + this->m_align = Block512B::AlignUp(m); + this->k_align = Block512B::AlignUp(k); + this->n_align = Block512B::AlignUp(n); + + this->m0 = ppTilingData->m0; + this->k0 = ppTilingData->k0; + this->n0 = ppTilingData->n0; + + int32_t tiling_key = ppTilingData->tilingKey; + this->trans_a = ppTilingData->isTransA; + this->trans_b = ppTilingData->isTransB; + + int32_t aligned_a; + int32_t aligned_b; + this->dequant_granularity = quantInfo->dequantGranularity; + AlignJudge(this->trans_a, this->trans_b, this->m, this->k, this->n, + this->m_align, this->k_align, this->n_align, aligned_a, aligned_b); + this->aligned_a = aligned_a; + this->aligned_b = aligned_b; + if (weight_nz) { + this->k_align16 = Block32B::AlignUp(k); + this->n_align16 = Block32B::AlignUp(n); + } + bool has_a_align = IsQuant(quantInfo->quantGranularity) || aligned_a; + bool has_b_align = IsQuant(this->dequant_granularity) && !this->is_int8 || aligned_b; + bool has_accum = IsQuant(this->dequant_granularity) && + this->is_int8 && std::is_same::value; + bool has_format_dequant_offset = + (this->dequant_granularity == QuantGranularity::PER_TENSOR) && this->is_int8 && this->has_offset; + auto workspace_info = GetWorkspaceInfo(workspace, this->batch_size, this->m, this->k, this->n, + this->m_align, this->k_align, this->n_align, this->trans_a, this->trans_b, + sizeof(MmadDtype), has_a_align, has_b_align, has_accum, has_format_dequant_offset); + this->gm_a_src = reinterpret_cast<__gm__ MmadDtype *>(x1); + this->gm_b_src = reinterpret_cast<__gm__ MmadDtype *>(x2); + this->gm_format_dequant_offset = reinterpret_cast<__gm__ int32_t *>(has_format_dequant_offset ? + workspace_info.gm_dequant_param : nullptr); + this->gm_workspace_src = workspace; + this->block_size = BLOCK_SIZE_32 / sizeof(MmadDtype); + + int32_t a_l1_size = this->m0 * this->k0 * sizeof(MmadDtype); + int32_t a_l1_size_round = AscendC::DivCeil(a_l1_size, 512) * 512; + int32_t b_l1_size = this->n0 * this->k0 * sizeof(MmadDtype); + int32_t b_l1_size_round = AscendC::DivCeil(b_l1_size, 512) * 512; + this->l1_base_a = reinterpret_cast<__cbuf__ MmadDtype *>((uintptr_t)(this->is_int8 ? SCALE_L1_SIZE : 0)); + this->l1_base_b = + reinterpret_cast<__cbuf__ MmadDtype *>(a_l1_size_round * (this->is_int8 ? DOUBLE_BUFFER_SIZE : 1) + + (uintptr_t) this->l1_base_a); + + this->core_num = get_block_num(); + this->core_idx = get_block_idx(); + + this->m_loop = ppTilingData->mLoop; + this->k_loop = ppTilingData->kLoop; + this->n_loop = ppTilingData->nLoop; + this->core_loop = ppTilingData->coreLoop; + this->swizzl_count = ppTilingData->swizzlCount; + this->swizzl_direct = ppTilingData->swizzlDirect; + this->is_91093 = commTilingData->is91093; + this->ping_flag = 1; + this->rank = hccl_.GetRankId(); + this->rank_size = hccl_.GetRankDim(); + this->withSerialMode = commTilingData->withSerialMode; + + this->gm_peer_mem = (__gm__ OutDtype *)hccl_.GetWindowsInAddr(this->rank); + } + + __aicore__ inline void MoveL0CToGM(__gm__ OutDtype *gm_dst, int64_t offset_c, + int32_t m_actual, int32_t n_actual, int32_t src_stride, int32_t dst_stride) { + if constexpr (std::is_same::value) { + copy_matrix_cc_to_gm( + gm_dst + offset_c, + l0c_buf, + 0, + n_actual, + m_actual, + dst_stride, + src_stride, + 0, + F322BF16, + 0, + false, + true + ); + } else { + copy_matrix_cc_to_gm( + gm_dst + offset_c, + l0c_buf, + 0, + n_actual, + m_actual, + dst_stride, + src_stride, + 0, + F322F16, + 0, + false, + true + ); + } + SetFlag(EVENT_ID0); + } + + __aicore__ inline void InitFlags() + { + WaitEvent(AIC_WAIT_AIV_FINISH_ALIGN_FLAG_ID); + } + + __aicore__ inline void Endflags() + { + } + + __aicore__ inline void Process() + { + // AIC matmul func, waits for AIV to complete [AllReduce & Add & RMSNorm]. + InitFlags(); + uint32_t m = this->m; + uint32_t k = this->k; + uint32_t n = this->n; + gmB.SetGlobalBuffer(gm_b_src, k * n); + + using LayoutA = layout::RowMajor; + using LayoutB = layout::ColumnMajor; + using LayoutC = layout::RowMajor; + LayoutB layoutB {(layout::ColumnMajor::Index)k, (layout::ColumnMajor::Index)n}; + + using L1TileShape = GemmShape; + using L0TileShape = GemmShape; + using AType = Gemm::GemmType; + using BType = Gemm::GemmType; + using CType = AType; + constexpr bool ENABLE_UNIT_FLAG = true; + using MmadDispatchPolicy = Gemm::MmadAtlasA2Pingpong; + using BlockMmad = Gemm::Block::BlockMmad; + GemmCoord blockShape = L1TileShape::ToCoord(); + + BlockMmad blockMmad(resource); + int mPerSplit = this->m0 * this->swizzl_count; + int mAvg = mPerSplit; + int splitM = AscendC::DivCeil(m, mPerSplit); + int flag_idx = 0; + icache_preload(8); // 8 corresponding to 16k + for (int splitIndex = 0; splitIndex < splitM; ++splitIndex) { + uint32_t mStart = splitIndex * mAvg; + uint32_t mActual = mAvg > (m - mStart) ? m - mStart:mAvg; + flag_idx = splitIndex % PIPE_DEPTH; + if (splitIndex >= PIPE_DEPTH) { + WaitEvent(flag_idx); + } + + __gm__ MmadDtype *gm_a_src_tmp = reinterpret_cast<__gm__ MmadDtype *>(gm_a_src) + mStart * k; + __gm__ MmadDtype *gm_c_src_tmp = reinterpret_cast<__gm__ MmadDtype *>(gm_peer_mem) + mStart * n; + gmA.SetGlobalBuffer(gm_a_src_tmp, mActual*k); + gmC.SetGlobalBuffer(gm_c_src_tmp, mActual*n); + + GemmCoord splitShape{mActual, n, k}; + using BlockScheduler = typename Gemm::Block::GemmIdentityBlockSwizzle<3, 1>; // SwizzleOffset=3 + BlockScheduler splitScheduler(splitShape, blockShape.GetCoordMN()); + uint32_t coreLoops = splitScheduler.GetCoreLoops(); + + LayoutA layoutA{mActual, k}; + LayoutC layoutC{mActual, n}; + + for (uint32_t loopIdx = core_idx; loopIdx < coreLoops; loopIdx += core_num) { + GemmCoord blockCoord = splitScheduler.GetBlockCoord(loopIdx); + GemmCoord actualBlockShape = splitScheduler.GetActualBlockShape(blockCoord); + GemmCoord offsetCoord = blockCoord * blockShape; + + MatrixCoord offsetA = offsetCoord.GetCoordMK(); + MatrixCoord offsetB = offsetCoord.GetCoordKN(); + MatrixCoord offsetC = offsetCoord.GetCoordMN(); + + int64_t gmOffsetA = layoutA.GetOffset(offsetA); + int64_t gmOffsetB = layoutB.GetOffset(offsetB); + int64_t gmOffsetC = layoutC.GetOffset(offsetC); + + blockMmad (gmA[gmOffsetA], layoutA, gmB[gmOffsetB], layoutB, gmC[gmOffsetC], layoutC, actualBlockShape); + } + + FFTSCrossCoreSync(FFTS_SYNC_AICORE_GROUP_MODE, flag_idx); + } + + Endflags(); + PipeBarrier(); + } + +private: + AscendC::GlobalTensor gmA; + AscendC::GlobalTensor gmB; + AscendC::GlobalTensor gmC; + __gm__ MmadDtype *gm_a_src{nullptr}; + __gm__ MmadDtype *gm_b_src{nullptr}; + + __gm__ OutDtype *gm_c{nullptr}; + __gm__ OutDtype *gm_peer_mem{nullptr}; + __gm__ int64_t *gm_dequant_scale{nullptr}; + __gm__ int32_t *gm_format_dequant_offset{nullptr}; + __gm__ int32_t *gm_accum{nullptr}; + __gm__ uint8_t *gm_workspace_src; + + __cbuf__ MmadDtype *l1_base_a = reinterpret_cast<__cbuf__ MmadDtype *>((uintptr_t) SCALE_L1_SIZE_A); + __cbuf__ MmadDtype *l1_base_b = reinterpret_cast<__cbuf__ MmadDtype *>((uintptr_t) SCALE_L1_SIZE_B); + + __ca__ MmadDtype *l0a_base = reinterpret_cast<__ca__ MmadDtype *>((uintptr_t) 0); + __cb__ MmadDtype *l0b_base = reinterpret_cast<__cb__ MmadDtype *>((uintptr_t) 0); + + __cc__ T_ACCUM *l0c_buf = reinterpret_cast<__cc__ T_ACCUM *>((uintptr_t) 0); + + __cbuf__ int64_t *scale_l1 = reinterpret_cast<__cbuf__ int64_t *>((uintptr_t) 0); + __fbuf__ int64_t *scale_FB = (__fbuf__ int64_t *)(0); + + __cbuf__ int32_t *bias_l1 = reinterpret_cast<__cbuf__ int32_t *>((uintptr_t)0); + uint16_t bias_bt = 0; + bool has_offset{false}; + + int32_t core_num; + + int32_t batch_size; + int32_t m; + int32_t k; + int32_t n; + int32_t m_align; + int32_t k_align; + int32_t n_align; + int32_t k_align16; + int32_t n_align16; + int32_t m0; + int32_t k0; + int32_t n0; + + int32_t m_loop; + int32_t n_loop; + int32_t k_loop; + int32_t core_loop; + int32_t core_idx; + int32_t ping_flag; + int32_t block_size; + int32_t cube_matrix_size; + + int32_t aligned_a; + int32_t aligned_b; + + int32_t swizzl_count; + int32_t swizzl_direct; + + int32_t rank; + int32_t rank_size; + + int32_t withSerialMode; + + int32_t ag_dim; + int32_t rs_dim; + bool inner_dim_is_Ag{false}; + int32_t ag_rank_idx; + int32_t rs_rank_idx; + bool weight_nz{false}; + + bool is_91093{false}; + QuantGranularity dequant_granularity; + + bool is_int8; + bool trans_a; + bool trans_b; + + Hccl hccl_; +}; + +#endif // MATMUL_ALLREDUCE_ADD_RMSNORM_AIC_KERNEL_H diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm_aiv_kernel.h b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm_aiv_kernel.h new file mode 100644 index 000000000..c76d39eb3 --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm_aiv_kernel.h @@ -0,0 +1,702 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MATMUL_ALLREDUCE_ADD_RMSNORM_AIV_KERNEL_H +#define MATMUL_ALLREDUCE_ADD_RMSNORM_AIV_KERNEL_H + +#include "kernel_operator.h" +#include "matmul_allreduce_add_rmsnorm_tiling.h" +#include "matmul_allreduce_add_rmsnorm_utils.h" + +using namespace AscendC; + +constexpr int32_t DIFUSION_ADD_LEN = 512; +constexpr int32_t TQUE_DEPTH = 1; +constexpr uint32_t TBUF_POOL_MAX_BUFID_SIZE = 8; +enum CrossRankSyncFlagEnum { + FLAG_ZERO_IDX, + FLAG_ONE_IDX, + FLAG_TWO_IDX, + FLAG_ADD_IDX, + FLAG_FOUR_IDX, + FLAG_GATHER_ADD_OUT_STEP1, + FLAG_GATHER_ADD_OUT_STEP2, + FLAG_NUM +}; +constexpr int32_t FLAG_VALUE = 1; +constexpr int32_t NUM_PER_REP_FP32 = 64; + +template +__aicore__ void CopyUbufToGmAlignB16(__gm__ T *dst, __ubuf__ T *src, uint16_t nBurst, uint32_t lenBurst, + uint16_t srcSTride, uint16_t dstStride) +{ + DataCopyExtParams dataCopyParams(nBurst, + lenBurst, + srcSTride, + dstStride, + 0); + LocalTensor ubTensor; + TBuffAddr ubAddr; + ubAddr.logicPos = static_cast(TPosition::VECIN); + ubAddr.bufferAddr = reinterpret_cast(src); + ubTensor.SetAddr(ubAddr); + GlobalTensor gmTensor; + gmTensor.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(dst)); + DataCopyPad(gmTensor, ubTensor, dataCopyParams); +} + +template +__aicore__ void CopyGmToUbufAlignB16(__ubuf__ T *dst, __gm__ T *src, uint16_t nBurst, uint32_t lenBurst, + uint16_t srcSTride, uint16_t dstStride) +{ + DataCopyExtParams dataCopyParams(nBurst, + lenBurst, + srcSTride, + dstStride, + 0); + LocalTensor ubTensor; + TBuffAddr ubAddr; + ubAddr.logicPos = static_cast(TPosition::VECIN); + ubAddr.bufferAddr = reinterpret_cast(dst); + ubTensor.SetAddr(ubAddr); + GlobalTensor gmTensor; + gmTensor.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(src)); + DataCopyPadExtParams padParams; + DataCopyPad(ubTensor, gmTensor, dataCopyParams, padParams); +} + +template +class MatmulAllreduceAddRmsnormAivKernel { + +public: + __aicore__ inline MatmulAllreduceAddRmsnormAivKernel() { } + __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR residual, GM_ADDR gamma, GM_ADDR y, GM_ADDR add_out, + GM_ADDR workspace, const MatmulAllreduceAddRmsnormTilingData *tilingData, + Hccl &hccl_) + { + this->hccl_ = hccl_; + is_deterministic = false; + auto ppTilingData = &tilingData->matmulAllreduceAddRmsnormInfo.ppTilingData; + auto commTilingData = &tilingData->matmulAllreduceAddRmsnormInfo.commTilingData; + auto quantInfo = &tilingData->matmulAllreduceAddRmsnormInfo.quantInfo; + + gm_out = reinterpret_cast<__gm__ MmadDtype *>(y); + gm_add_input = reinterpret_cast<__gm__ MmadDtype *>(residual); + gm_add_output = reinterpret_cast<__gm__ MmadDtype *>(add_out); + gm_gamma = reinterpret_cast<__gm__ MmadDtype *>(gamma); + + batch_size = ppTilingData->opShape.batchSize; + m = ppTilingData->opShape.m; + k = ppTilingData->opShape.k; + n = ppTilingData->opShape.n; + + m0 = ppTilingData->m0; + k0 = ppTilingData->k0; + n0 = ppTilingData->n0; + + m_loop = ppTilingData->mLoop; + k_loop = ppTilingData->kLoop; + n_loop = ppTilingData->nLoop; + + core_loop = ppTilingData->coreLoop; + swizzl_count = ppTilingData->swizzlCount; + tiling_key = ppTilingData->tilingKey; + rank = hccl_.GetRankId(); + rank_size = hccl_.GetRankDim(); + + max_ub_single_dma_size = commTilingData->ubMoveNum; + withSerialMode = false; + tag = commTilingData->tag; + comm_npu_split = commTilingData->commNpuSplit; + comm_data_split = commTilingData->commDataSplit; + comm_direct = commTilingData->commDirect; + is_91093 = false; + core_count = comm_npu_split * comm_data_split; + dequant_granularity = static_cast(quantInfo->dequantGranularity); + dequant_group_size = quantInfo->dequantGroupSize; + quant_granularity = static_cast(quantInfo->quantGranularity); + quant_group_size = quantInfo->quantGroupSize; + epsilon = tilingData->matmulAllreduceAddRmsnormInfo.rmsnormTilingData.epsilon; + is_gather_add_out = tilingData->matmulAllreduceAddRmsnormInfo.ppTilingData.isGatherAddOut; + + swizzl_direct = (tiling_key & SWIZZL_MASK) ? true : false; + trans_a = ppTilingData->isTransA; + trans_b = ppTilingData->isTransB; + is_int8 = false; + ag_dim = 0; + rs_dim = 0; + inner_dim_is_Ag = false; + weight_nz = false; + max_ub_ping_pong_size = max_ub_single_dma_size / 2; // 2 - double buffer + + core_idx = get_block_idx(); + core_num = get_block_num(); + aiv_idx = get_subblockid(); + other_rank = (core_idx < rank_size) ? core_idx : -1; + + // init ub usage + pipe.InitBuffer(ctrlBuf, AscendC::ONE_BLK_SIZE); + ub_ctrl_flag = reinterpret_cast<__ubuf__ int32_t *>(ctrlBuf.Get().GetPhyAddr()); + + pipe.InitBuffer(gammaBuf, n * sizeof(MmadDtype)); + + uint32_t step1_ub_usage = AscendC::AlignUp( + n * sizeof(MmadDtype) + + 2 * (rank_size * DIFUSION_ADD_LEN * sizeof(MmadDtype)) + + n * sizeof(MmadDtype) + + n * sizeof(MmadDtype) + + n * sizeof(float) + + n * sizeof(float) + + n * sizeof(float), + AscendC::ONE_BLK_SIZE); + + uint32_t step2_ub_usage = AscendC::AlignUp( + max_ub_ping_pong_size * sizeof(MmadDtype), + AscendC::ONE_BLK_SIZE) * 2; + uint32_t max_step_ub_usage = max(step1_ub_usage, step2_ub_usage); + + pipe.InitBufPool(step1BufPool, max_step_ub_usage); + pipe.InitBufPool(step2BufPool, max_step_ub_usage, step1BufPool); + + step1BufPool.InitBuffer(inQueueX, 1, n * sizeof(MmadDtype)); + step1BufPool.InitBuffer(inQueueY, 2, rank_size * DIFUSION_ADD_LEN * sizeof(MmadDtype)); + step1BufPool.InitBuffer(addOutQueue, 1, n * sizeof(MmadDtype)); + step1BufPool.InitBuffer(outQueue, 1, n * sizeof(MmadDtype)); + step1BufPool.InitBuffer(xFp32Buf, n * sizeof(float)); + step1BufPool.InitBuffer(sqxBuf, n * sizeof(float)); + step1BufPool.InitBuffer(reduceFp32Buf, n * sizeof(float)); + + step2BufPool.InitBuffer(allgatherBuf[0], max_ub_ping_pong_size * sizeof(MmadDtype)); + step2BufPool.InitBuffer(allgatherBuf[1], max_ub_ping_pong_size * sizeof(MmadDtype)); + + CopyInGamma(); + } + + __aicore__ inline void Process(const MatmulAllreduceAddRmsnormTilingData *tilingData) + { + // AIV AllReduce & Add & RMSNorm func, waits for AIC to complete [Matmul]. + FFTSCrossCoreSync(FFTS_SYNC_AICORE_GROUP_MODE, AIC_WAIT_AIV_FINISH_ALIGN_FLAG_ID); + PipeBarrier(); + + ResetIpcFlags(FLAG_NUM); + CrossRankSyncEx(FLAG_NUM); + constexpr int32_t allreduce_used_core = 16; + int32_t one_comm_count = swizzl_count; + int32_t loop_num_per_comm = one_comm_count * n_loop; + int32_t comm_count = DivCeil(core_loop, loop_num_per_comm); + int32_t pipe_depth = is_91093 ? BLOCK_COUNT_4 : MAX_BLOCK_COUNT; + + for (int cal_idx = 0; cal_idx < comm_count; ++cal_idx) { + uint64_t flag_idx = cal_idx % pipe_depth; + int32_t m_total = (cal_idx == comm_count - 1) ? + m - cal_idx * swizzl_count * m0 : swizzl_count * m0; + int32_t m_per_rank = DivCeil(m_total, rank_size); + int32_t loop_offset = cal_idx * swizzl_count * m0; + + WaitEvent(flag_idx); + SetAndWaitAivSync(flag_idx, is_91093 ? BLOCK_COUNT_4 : MAX_BLOCK_COUNT); + CrossRankSyncV1(FLAG_ZERO_IDX, cal_idx + 1); + SetAndWaitAivSync(flag_idx, is_91093 ? BLOCK_COUNT_4 : MAX_BLOCK_COUNT); + + if (aiv_idx == 0 && core_idx < allreduce_used_core) { + int32_t m_cur_rank = LimitRange(m_total - rank * m_per_rank, 0, m_per_rank); + int32_t m_per_core = DivCeil(m_cur_rank, allreduce_used_core); + int32_t m_cur_core = LimitRange(m_cur_rank - core_idx * m_per_core, 0, m_per_core); + int32_t core_offset_m = loop_offset + rank * m_per_rank + core_idx * m_per_core; + ParallelWithSplitStepOneAddNorm(core_offset_m * n, m_cur_core); + } + + PipeBarrier(); + + SetAndWaitAivSync(flag_idx, is_91093 ? BLOCK_COUNT_4 : MAX_BLOCK_COUNT); + CrossRankSyncV1(FLAG_ADD_IDX, cal_idx + 1); + SetAndWaitAivSync(flag_idx, is_91093 ? BLOCK_COUNT_4 : MAX_BLOCK_COUNT); + + { // ParallelWithSplitStepTwo + int32_t used_core_per_rank = allreduce_used_core / rank_size; + int32_t sub_core_idx = core_idx % used_core_per_rank; + int32_t gather_rank_id = core_idx / used_core_per_rank; + int32_t m_in_rank = LimitRange(m_total - gather_rank_id * m_per_rank, 0, m_per_rank); + int32_t m_per_core = DivCeil(m_in_rank, used_core_per_rank); + int32_t m_cur_core = LimitRange(m_in_rank - sub_core_idx * m_per_core, 0, m_per_core); + int32_t core_offset_m = loop_offset + gather_rank_id * m_per_rank + sub_core_idx * m_per_core; + auto gm_share_buff = (__gm__ MmadDtype *)hccl_.GetWindowsInAddr(gather_rank_id); + + bool filter_core_cond = aiv_idx == 0 && core_idx < allreduce_used_core && m_cur_core > 0; + if (filter_core_cond) { + ParallelAllGather(gm_out, gm_share_buff, core_offset_m * n, m_cur_core * n); + } + + SetAndWaitAivSync(flag_idx); + CrossRankSyncV2(FLAG_TWO_IDX, cal_idx + 1); + SetAndWaitAivSync(flag_idx); + + if (is_gather_add_out) { + if (filter_core_cond && gather_rank_id == rank) { + ParallelAllGather(gm_share_buff, gm_add_output, core_offset_m * n, m_cur_core * n); + } + + SetAndWaitAivSync(flag_idx); + CrossRankSyncV2(FLAG_GATHER_ADD_OUT_STEP1, cal_idx + 1); + SetAndWaitAivSync(flag_idx); + + if (filter_core_cond && gather_rank_id != rank) { + ParallelAllGather(gm_add_output, gm_share_buff, core_offset_m * n, m_cur_core * n); + } + + SetAndWaitAivSync(flag_idx); + CrossRankSyncV2(FLAG_GATHER_ADD_OUT_STEP2, cal_idx + 1); + SetAndWaitAivSync(flag_idx); + } + } + + if (cal_idx <= comm_count - pipe_depth) { + SetAicSync(flag_idx); + } + } + ResetIpcFlags(FLAG_NUM); + if (aiv_idx == 0 && core_idx < rank_size) { + __gm__ int32_t *state_buff = (__gm__ int32_t *)hccl_.GetWindowsOutAddr(other_rank); + CheckBuffFlag(ub_ctrl_flag, state_buff + FLAG_ZERO_IDX, 0); + } + } + +private: + __aicore__ void SetBuffFlag(__ubuf__ int32_t *ub_ctrl_flag, __gm__ int32_t *buff, int32_t flag) + { + *ub_ctrl_flag = flag; + SetFlag(EVENT_ID2); + WaitFlag(EVENT_ID2); + CopyUbufToGmAlignB16(buff, ub_ctrl_flag, 1, sizeof(int32_t), 0, 0); + } + + __aicore__ void SetBuffFlagByAdd(__ubuf__ int32_t *ub_ctrl_flag, __gm__ int32_t *buff, int32_t flag) + { + PipeBarrier(); + *ub_ctrl_flag = flag; + PipeBarrier(); + SetAtomicAdd(); + PipeBarrier(); + CopyUbufToGmAlignB16(buff, ub_ctrl_flag, 1, sizeof(int32_t), 0, 0); + PipeBarrier(); + SetAtomicNone(); + PipeBarrier(); + } + + __aicore__ void CheckBuffFlag(__ubuf__ int32_t *ub_ctrl_flag, __gm__ int32_t *buff, int32_t flag) + { + SetFlag(EVENT_ID1); + WaitFlag(EVENT_ID1); + while (true) { + CopyGmToUbufAlignB16(ub_ctrl_flag, buff, 1, sizeof(int32_t), 0, 0); + SetFlag(EVENT_ID3); + WaitFlag(EVENT_ID3); + if (*ub_ctrl_flag == flag) { + break; + } + } + } + + __aicore__ void SetAicSync(uint64_t flag_idx) + { + FFTSCrossCoreSync(FFTS_SYNC_AICORE_GROUP_MODE, flag_idx); + } + + __aicore__ void ResetIpcFlags(int32_t num_flags) + { + for (int32_t idx = 0; idx <= num_flags; ++idx) { + if (core_idx == 0 && aiv_idx == 0) { + __gm__ int32_t *state_buff = (__gm__ int32_t *)hccl_.GetWindowsOutAddr(rank); + SetBuffFlag(ub_ctrl_flag, state_buff + idx, 0); + } + } + } + + __aicore__ void CrossRankSyncV1(int32_t flag_idx, int32_t flag_data) + { + if (aiv_idx == 0 && core_idx == rank) { + __gm__ int32_t *state_buff = (__gm__ int32_t *)hccl_.GetWindowsOutAddr(rank); + SetBuffFlagByAdd(ub_ctrl_flag, state_buff + flag_idx, FLAG_VALUE); + } else if (aiv_idx == 0 && core_idx < rank_size) { + __gm__ int32_t *state_buff = (__gm__ int32_t *)hccl_.GetWindowsOutAddr(core_idx); + CheckBuffFlag(ub_ctrl_flag, state_buff + flag_idx, FLAG_VALUE * flag_data); + } + } + + __aicore__ void CrossRankSyncV2(int32_t flag_idx, int32_t flag_data) + { + if (aiv_idx == 0 && core_idx < rank_size) { + __gm__ int32_t *state_buff = (__gm__ int32_t *)hccl_.GetWindowsOutAddr(core_idx); + SetBuffFlagByAdd(ub_ctrl_flag, state_buff + flag_idx, FLAG_VALUE); + } + if (aiv_idx == 0 && core_idx == rank) { + __gm__ int32_t *state_buff = (__gm__ int32_t *)hccl_.GetWindowsOutAddr(rank); + CheckBuffFlag(ub_ctrl_flag, state_buff + flag_idx, FLAG_VALUE * rank_size * flag_data); + } + } + + __aicore__ void SetAndWaitAivSync(uint64_t flag_idx, int32_t pipe_depth = 2) + { + FFTSCrossCoreSync(0, flag_idx + pipe_depth); + WaitEvent(flag_idx + pipe_depth); + } + + __aicore__ inline uint32_t GetGmU32(GM_ADDR gm_addr) + { + copy_gm_to_ubuf_align_b32(ub_ctrl_flag, gm_addr, 0, 1, sizeof(uint32_t), 0, 0, 0, 0); + PipeSync(); + return *reinterpret_cast<__ubuf__ uint32_t *>(ub_ctrl_flag); + } + + __aicore__ inline void SetGmU32(GM_ADDR gm_addr, uint32_t data) + { + *reinterpret_cast<__ubuf__ uint32_t *>(ub_ctrl_flag) = data; + PipeSync(); + copy_ubuf_to_gm_align_b32(gm_addr, ub_ctrl_flag, 0, 1, sizeof(uint32_t), 0, 0, 0, 0); + } + + __aicore__ inline void CrossRankSyncEx(uint32_t flag_idx) + { + AscendC::SyncAll(); + __asm__ __volatile__(""); + if (aiv_idx == 0 && core_idx == 0) { + auto flag_addr = (GM_ADDR)hccl_.GetWindowsOutAddr(0) + flag_idx * AscendC::ONE_BLK_SIZE; + uint32_t old_flag_data = GetGmU32(flag_addr); + __asm__ __volatile__(""); + SetAtomicAdd(); + SetGmU32(flag_addr, 1); + PipeSync(); + SetAtomicNone(); + __asm__ __volatile__(""); + + uint32_t new_flag_data; + do { + new_flag_data = GetGmU32(flag_addr); + __asm__ __volatile__(""); + } while (new_flag_data - old_flag_data < rank_size); + __asm__ __volatile__(""); + SetAtomicAdd(); + SetGmU32(flag_addr, 1); + PipeSync(); + SetAtomicNone(); + } + __asm__ __volatile__(""); + AscendC::SyncAll(); + } + + template + __aicore__ inline T min(const T& a, const T& b) { + return (a < b) ? a : b; + } + + template + __aicore__ inline T max(const T& a, const T& b) { + return (a > b) ? a : b; + } + + template + __aicore__ inline T LimitRange(const T& val, const T& low, const T& high) { + return min(max(val, low), high); + } + + template + __aicore__ inline void PipeSync() + { + AscendC::TEventID event_id = static_cast(GetTPipePtr()->FetchEventID(EVENT)); + AscendC::SetFlag(event_id); + AscendC::WaitFlag(event_id); + } + + __aicore__ inline void CopyInGamma() + { + GlobalTensor gamma_global; + gamma_global.SetGlobalBuffer((__gm__ MmadDtype *)gm_gamma, n); + DataCopy(gammaBuf.Get(), gamma_global, n); + PipeSync(); + } + + __aicore__ void ParallelWithSplitStepOneAddNorm(uint32_t core_buf_offset, uint32_t m_cur_core) + { + if (m_cur_core <= 0) { + return; + } + + auto buff = (__gm__ MmadDtype *)hccl_.GetWindowsInAddr(rank); + + GlobalTensor x_global; + GlobalTensor y_global; + GlobalTensor out_global; + GlobalTensor add_out_global; + + x_global.SetGlobalBuffer(buff + core_buf_offset); + out_global.SetGlobalBuffer(buff + core_buf_offset); + add_out_global.SetGlobalBuffer(gm_add_output + core_buf_offset); + + uint32_t add_count = DivCeil(n, DIFUSION_ADD_LEN); + + LocalTensor x_local; + LocalTensor y_local; + + for (uint32_t i = 0; i < m_cur_core; i++) { + LocalTensor x_fp32 = xFp32Buf.Get(); + LocalTensor sqx = sqxBuf.Get(); + + x_local = inQueueX.AllocTensor(); + for (uint32_t j = 0; j < add_count; j++) { + uint32_t add_offset = j * DIFUSION_ADD_LEN; + uint32_t add_len = min(n - add_offset, DIFUSION_ADD_LEN); + + DataCopy(x_local[add_offset], x_global[i * n + add_offset], add_len); + inQueueX.EnQue(x_local); + + uint32_t iterate_end = (rank + 1) % rank_size; + y_local = inQueueY.AllocTensor(); + for (uint32_t k = 0; k < rank_size; ++k) { + uint32_t iterate_idx = iterate_end + k; + if (iterate_idx >= rank_size) { + iterate_idx -= rank_size; + } + + if (iterate_idx == rank) { + y_global.SetGlobalBuffer(gm_add_input + core_buf_offset); + } else { + auto other_buff = (__gm__ MmadDtype *)hccl_.GetWindowsInAddr(iterate_idx); + y_global.SetGlobalBuffer(other_buff + core_buf_offset); + } + DataCopy(y_local[k * add_len], y_global[i * n + add_offset], add_len); + } + inQueueY.EnQue(y_local); + x_local = inQueueX.DeQue(); + y_local = inQueueY.DeQue(); + + Cast(x_fp32[add_offset], x_local[add_offset], RoundMode::CAST_NONE, add_len); + PipeBarrier(); + for (uint32_t k = 0; k < rank_size; ++k) { + // use sqx as shared buf, required n >= add_len + Cast(sqx, y_local[k * add_len], RoundMode::CAST_NONE, add_len); + PipeBarrier(); + Add(x_fp32[add_offset], x_fp32[add_offset], sqx, add_len); + PipeBarrier(); + } + + inQueueY.FreeTensor(y_local); + } + inQueueX.FreeTensor(x_local); + + // copy add result out + LocalTensor add_out = addOutQueue.AllocTensor(); + Cast(add_out, x_fp32, RoundMode::CAST_RINT, n); + addOutQueue.EnQue(add_out); + add_out = addOutQueue.DeQue(); + DataCopy(add_out_global[i * n], add_out, n); + addOutQueue.FreeTensor(add_out); + + LocalTensor gamma_local = gammaBuf.Get(); + LocalTensor out_local = outQueue.AllocTensor(); + LocalTensor reduce_buf_local = reduceFp32Buf.Get(); + + // make sure precision is same in bf16 case + Cast(out_local, x_fp32, RoundMode::CAST_RINT, n); + PipeBarrier(); + + Cast(x_fp32, out_local, RoundMode::CAST_NONE, n); + PipeBarrier(); + + Mul(sqx, x_fp32, x_fp32, n); + PipeBarrier(); + + Muls(sqx, sqx, (float)1.0 / n, n); + PipeBarrier(); + + ReduceSum(sqx, sqx, reduce_buf_local, n); + PipeBarrier(); + + Adds(sqx, sqx, epsilon, 1); + PipeBarrier(); + + Sqrt(sqx, sqx, 1); + Duplicate(reduce_buf_local, (float)1.0, 1); + PipeBarrier(); + + Div(sqx, reduce_buf_local, sqx, 1); + PipeBarrier(); + + PipeSync(); + float rstd_value = sqx.GetValue(0); + PipeSync(); + PipeBarrier(); + + Muls(x_fp32, x_fp32, rstd_value, n); + PipeBarrier(); + + if constexpr (std::is_same::value) { + Cast(out_local, x_fp32, RoundMode::CAST_NONE, n); + PipeBarrier(); + Mul(out_local, gamma_local, out_local, n); + PipeBarrier(); + } else if constexpr (std::is_same::value) { + Cast(out_local, x_fp32, RoundMode::CAST_RINT, n); + PipeBarrier(); + Cast(x_fp32, out_local, RoundMode::CAST_NONE, n); + PipeBarrier(); + Cast(sqx, gamma_local, RoundMode::CAST_NONE, n); + PipeBarrier(); + + Mul(x_fp32, x_fp32, sqx, n); + PipeBarrier(); + Cast(out_local, x_fp32, RoundMode::CAST_RINT, n); + PipeBarrier(); + PipeSync(); + } + + outQueue.EnQue(out_local); + out_local = outQueue.DeQue(); + DataCopy(out_global[i * n], out_local, n); + outQueue.FreeTensor(out_local); + } + } + + __aicore__ void ParallelAllGather(__gm__ MmadDtype *gm_dst, __gm__ MmadDtype *gm_src, + uint32_t core_buf_offset, uint32_t data_len) + { + GlobalTensor src_global; + GlobalTensor dst_global; + src_global.SetGlobalBuffer(gm_src); + dst_global.SetGlobalBuffer(gm_dst); + + constexpr uint32_t PIPELINE_COPY_NUM = sizeof(allgatherBuf) / sizeof(allgatherBuf[0]); + TEventID ev_mte3_mte2[PIPELINE_COPY_NUM]; + TEventID ev_mte2_mte3[PIPELINE_COPY_NUM]; + LocalTensor local_tensors[PIPELINE_COPY_NUM]; + + for (uint32_t i = 0; i < PIPELINE_COPY_NUM; i++) { + ev_mte3_mte2[i] = GetTPipePtr()->AllocEventID(); + ev_mte2_mte3[i] = GetTPipePtr()->AllocEventID(); + SetFlag(ev_mte3_mte2[i]); + local_tensors[i] = allgatherBuf[i].Get(); + } + + uint32_t offset = core_buf_offset; + uint32_t copy_len = max_ub_ping_pong_size; // num of MmadDtype, not the byte length + uint32_t copy_count = DivCeil(data_len, copy_len); + uint32_t pipe_id = 0; + + for (uint32_t i = 0; i < copy_count; i++) { + uint32_t actual_copy_len = + (i == copy_count - 1) ? (data_len - i * copy_len) : copy_len; + + auto &local_tensor = local_tensors[pipe_id]; + + WaitFlag(ev_mte3_mte2[pipe_id]); + DataCopy(local_tensor, src_global[offset], actual_copy_len); + SetFlag(ev_mte2_mte3[pipe_id]); + WaitFlag(ev_mte2_mte3[pipe_id]); + DataCopy(dst_global[offset], local_tensor, actual_copy_len); + SetFlag(ev_mte3_mte2[pipe_id]); + + offset += actual_copy_len; + pipe_id = (pipe_id + 1) % PIPELINE_COPY_NUM; + } + + for (uint32_t i = 0; i < PIPELINE_COPY_NUM; i++) { + WaitFlag(ev_mte3_mte2[i]); + GetTPipePtr()->ReleaseEventID(ev_mte3_mte2[i]); + GetTPipePtr()->ReleaseEventID(ev_mte2_mte3[i]); + } + + PipeBarrier(); + } + + __gm__ MmadDtype *gm_out; + __gm__ MmadDtype *gm_add_input; + __gm__ MmadDtype *gm_add_output; + __gm__ MmadDtype *gm_gamma; + __ubuf__ int32_t *ub_ctrl_flag; + + int32_t batch_size; + int32_t m; + int32_t k; + int32_t n; + int32_t m0; + int32_t k0; + int32_t n0; + + int32_t m_loop; + int32_t n_loop; + int32_t k_loop; + int32_t core_loop; + int32_t core_idx; + + int32_t rank; + int32_t rank_size; + int32_t tiling_key; + int32_t swizzl_count; + bool swizzl_direct; + + bool trans_a; + bool trans_b; + bool is_int8; + bool is_91093; + bool is_gather_add_out; + + int32_t aiv_idx; + int32_t other_rank; + int32_t core_num; + int32_t max_ub_single_dma_size; + int32_t max_ub_ping_pong_size; + + int32_t gm_c_pingpong_size; + int32_t withSerialMode; + int32_t tag; + int32_t comm_npu_split; + int32_t comm_data_split; + int32_t comm_direct; + + int32_t core_count; + bool is_deterministic; + + QuantGranularity dequant_granularity; + int32_t dequant_group_size; + QuantGranularity quant_granularity; + int32_t quant_group_size; + + WorkspaceInfo workspace_info; + int32_t ag_dim; + int32_t rs_dim; + bool inner_dim_is_Ag; + bool weight_nz{false}; + + float epsilon; + + TPipe pipe; + AscendC::TBufPool step1BufPool; + AscendC::TBufPool step2BufPool; + + AscendC::TQue inQueueX, inQueueY; + AscendC::TQue outQueueZ; + AscendC::TQue addOutQueue; + AscendC::TQue outQueue; + + AscendC::TBuf ctrlBuf; + AscendC::TBuf gammaBuf; + AscendC::TBuf xFp32Buf; + AscendC::TBuf sqxBuf; + AscendC::TBuf reduceFp32Buf; + AscendC::TBuf allgatherBuf[2]; + + Hccl hccl_; +}; +#endif \ No newline at end of file diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm_tiling.h b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm_tiling.h new file mode 100644 index 000000000..88a4401e1 --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm_tiling.h @@ -0,0 +1,101 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MATMUL_ALLREDUCE_ADD_RMSNORM_TILING_H +#define MATMUL_ALLREDUCE_ADD_RMSNORM_TILING_H + +#include +#include "kernel_tiling/kernel_tiling.h" + +enum QuantGranularity : int { + QUANT_GRANULARITY_UNDEFINED = -1, + PER_TENSOR = 0, + PER_CHANNEL = 1, + PER_GROUP = 2, + QUANT_GRANULARITY_MAX = 3, +}; + +struct Opshape { + int32_t batchSize = 1; + int32_t m = -1; + int32_t k = -1; + int32_t n = -1; +}; + +struct PPTilingData { + Opshape opShape = {}; + int32_t m0 = 1; + int32_t k0 = 1; + int32_t n0 = 1; + int32_t mLoop = 1; + int32_t kLoop = 1; + int32_t nLoop = 1; + int32_t coreLoop = 1; + int32_t swizzlCount = 1; + int32_t swizzlDirect = 0; + uint32_t tilingKey = 0; + int32_t blockDim = 1; + int32_t splitK = 0; + bool weightNz = false; + bool isTransA = false; + bool isTransB = false; + bool isGatherAddOut = false; +}; + +struct CommTilingData { + int32_t rank = 1; + int32_t rankSize = 1; + int32_t pValue = 1; + int32_t ubMoveNum = 1; + int32_t write2OtherRank = 0; + int32_t withSerialMode = 0; + int32_t tag = 0; + int32_t commNpuSplit = 1; + int32_t commDataSplit = 1; + int32_t commDirect = 0; + int32_t lenPerLoop = 1; + int32_t is91093 = 0; + int32_t buffer_size = 0; +}; + +struct RmsNormTilingData { + RmsNormTiling tiling{}; + uint32_t loopCount; + uint32_t calcBytes; + float epsilon{}; +}; + +struct QuantInfo { + QuantGranularity dequantGranularity = QuantGranularity::QUANT_GRANULARITY_UNDEFINED; + int32_t dequantGroupSize = -1; + QuantGranularity quantGranularity = QuantGranularity::QUANT_GRANULARITY_UNDEFINED; + int32_t quantGroupSize = -1; +}; + +struct MatmulAllreduceAddRmsnormInfo { + PPTilingData ppTilingData{}; + CommTilingData commTilingData{}; + RmsNormTilingData rmsnormTilingData{}; + QuantInfo quantInfo{}; +}; + +struct MatmulAllreduceAddRmsnormTilingData { + Mc2InitTiling mc2InitTiling; + Mc2CcTiling mc2CcTiling; + MatmulAllreduceAddRmsnormInfo matmulAllreduceAddRmsnormInfo; +}; + +#endif // MATMUL_ALLREDUCE_ADD_RMSNORM_TILING_H \ No newline at end of file diff --git a/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm_utils.h b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm_utils.h new file mode 100644 index 000000000..9f2e88976 --- /dev/null +++ b/csrc/ascend/mc2/matmul_allreduce_add_rmsnorm/op_kernel/matmul_allreduce_add_rmsnorm_utils.h @@ -0,0 +1,414 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MATMUL_ALLREDUCE_ADD_RMSNORM_UTILS_H +#define MATMUL_ALLREDUCE_ADD_RMSNORM_UTILS_H + +#include +#include "kernel_operator.h" +using namespace AscendC; + +constexpr int64_t ND2NZ_STRIDE_LIMIT = 65536; +constexpr int32_t AIC_WAIT_AIV_FINISH_ALIGN_FLAG_ID = 12; +constexpr int32_t MAX_BLOCK_COUNT = 2; +constexpr int32_t BLOCK_COUNT_3 = 3; +constexpr int32_t BLOCK_COUNT_4 = 4; +constexpr int32_t TILE_BLOCK_MOD = 2; + +constexpr int32_t BLOCK_SIZE_32B = 32; +constexpr int32_t BLOCK_SIZE_256B = 256; +constexpr int32_t BLOCK_SIZE_512B = 512; + +constexpr int32_t FFTS_SYNC_INTERNEL_MODE = 0; +constexpr int32_t FFTS_SYNC_AICORE_GROUP_MODE = 2; + +constexpr int32_t SWIZZL_MASK = 0b100000; +constexpr int32_t TRANS_A_MASK = 0b010000; +constexpr int32_t TRANS_B_MASK = 0b001000; +constexpr int32_t INT8_MASK = 0b000100; +constexpr int32_t BIAS_MASK = 0b000010; + +template +struct BaseBlock { + static_assert((SIZE & (SIZE - 1)) == 0, "Invalid block size"); + static constexpr size_t size = SIZE / sizeof(T); + + static __aicore__ inline size_t Count(size_t len) + { + return (len + size - 1) / size; + } + + static __aicore__ inline bool IsAligned(size_t len) + { + return len % size == 0; + } + + static __aicore__ inline size_t AlignUp(size_t len) + { + return (len + size - 1) & ~(size - 1); + } + + static __aicore__ inline size_t AlignDown(size_t len) + { + return len & ~(size - 1); + } +}; + +template +using Block32B = BaseBlock; + +template +using Block256B = BaseBlock; + +template +using Block512B = BaseBlock; + +struct WorkspaceInfo { + __gm__ uint8_t *gm_a_align{ nullptr }; + __gm__ uint8_t *gm_b_align{ nullptr }; + __gm__ uint8_t *gm_accum{ nullptr }; + __gm__ uint8_t *gm_dequant_param{ nullptr }; +}; + +template +__aicore__ inline LocalTensor CreateLocalTensor(__ubuf__ T *addr) +{ + LocalTensor tensor; + TBuffAddr taddr; + taddr.bufferAddr = reinterpret_cast(addr); + tensor.SetAddr(taddr); + return tensor; +} + +template +__aicore__ inline LocalTensor CreateLocalTensor(uint32_t buffer_offset) +{ + LocalTensor tensor; + tensor.address_.bufferAddr = buffer_offset; + return tensor; +} + +template +__aicore__ inline LocalTensor CreateLocalTensor(uint32_t buffer_offset, uint8_t logic_pos) +{ + LocalTensor tensor; + tensor.address_.logicPos = logic_pos; + tensor.address_.bufferAddr = buffer_offset; + return tensor; +} + +template +struct IntrinsicCopyGmToL1Nd2Nz { + static __aicore__ inline void move( + __cbuf__ T *dst, __gm__ T *src, + uint8_t sid, uint16_t ndNum, uint16_t nValue, uint16_t dValue, + uint16_t srcNdMatrixStride, uint16_t srcDValue, uint16_t dstNzC0Stride, + uint16_t dstNzNStride, uint16_t dstNzMatrixStride) { + Nd2NzParams nd2nzParams( + ndNum, nValue, dValue, + srcNdMatrixStride, srcDValue, dstNzC0Stride, + dstNzNStride, dstNzMatrixStride + ); + uint32_t dst_buffer_offset = reinterpret_cast(dst); + uint8_t dst_logicpos = static_cast(TPosition::C1); + LocalTensor dstTensor; + dstTensor = CreateLocalTensor(dst_buffer_offset, dst_logicpos); + GlobalTensor srcTensor; + srcTensor.SetGlobalBuffer(src); + DataCopy(dstTensor, srcTensor, nd2nzParams); + } +}; + +template +struct CopyGmToL1Nd2zN { + static __aicore__ inline void move( + __cbuf__ T *dst, __gm__ T *src, + uint16_t nValue, uint16_t dValue, uint32_t srcDValue, uint16_t dstNzC0Stride) { + constexpr int BLOCK_LEN = 32 / sizeof(T); + if (srcDValue < ND2NZ_STRIDE_LIMIT) { + IntrinsicCopyGmToL1Nd2Nz::move( + dst, + src, + 0, + 1, + nValue, + dValue, + 0, + srcDValue, + dstNzC0Stride, + 1, + 0 + ); + } else { + for (int i = 0; i < nValue; i++) { + IntrinsicCopyGmToL1Nd2Nz::move( + dst + i * BLOCK_LEN, + src + i * srcDValue, + 0, + 1, + 1, + dValue, + 0, + 0, + dstNzC0Stride, + 0, + 0 + ); + } + } + } +}; + +__aicore__ inline void AlignJudge(bool trans_a, bool trans_b, int32_t m, int32_t k, int32_t n, int32_t m_align, + int32_t k_align, int32_t n_align, int32_t &aligned_a, int32_t &aligned_b) +{ + if (!trans_a) { + aligned_a = k != k_align; + } else { + aligned_a = (m != m_align && m != 1); + } + + if (!trans_b) { + aligned_b = (n != n_align); + } else { + aligned_b = (k != k_align); + } +} + +__aicore__ inline WorkspaceInfo GetWorkspaceInfo(__gm__ uint8_t *gm_workspace, int32_t batch_size, int32_t m, + int32_t k, int32_t n, int32_t m_align, int32_t k_align, int32_t n_align, bool trans_a, bool trans_b, + int32_t mmad_dsize, bool has_a_align, bool has_b_align, bool has_accum = false, bool has_dequant_param = false) +{ + WorkspaceInfo workspace_info; + uint64_t workspace_offset = 0; + + if (has_a_align) { + workspace_info.gm_a_align = gm_workspace + workspace_offset; + workspace_offset += static_cast(batch_size) * (trans_a ? k * m_align : m * k_align) * mmad_dsize; + } + + if (has_b_align) { + workspace_info.gm_b_align = gm_workspace + workspace_offset; + workspace_offset += static_cast(batch_size) * (trans_b ? n * k_align : k * n_align) * mmad_dsize; + } + + if (has_accum) { + workspace_info.gm_accum = gm_workspace + workspace_offset; + workspace_offset += static_cast(batch_size) * m * n * sizeof(int32_t); + } + + if (has_dequant_param) { + workspace_info.gm_dequant_param = gm_workspace + workspace_offset; + workspace_offset += n * sizeof(float32_t); + } + + return workspace_info; +} + + +template +__aicore__ inline void CopyCubfToBt(uint64_t dst, __cbuf__ T *src, uint16_t convControl, uint16_t nBurst, + uint16_t lenBurst, uint16_t sourceGap, uint16_t dstGap) +{ + DataCopyParams intriParams(nBurst, lenBurst, sourceGap, dstGap); + uint32_t src_buffer_offset = reinterpret_cast(src); + uint32_t dst_buffer_offset = reinterpret_cast(dst); + uint8_t src_logicpos = static_cast(TPosition::C1); + uint8_t dst_logicpos = static_cast(TPosition::C2); + LocalTensor srcTensor; + LocalTensor dstTensor; + srcTensor = CreateLocalTensor(src_buffer_offset, src_logicpos); + dstTensor = CreateLocalTensor(dst_buffer_offset, dst_logicpos); + DataCopy(dstTensor, srcTensor, intriParams); +} + +template +__aicore__ inline void CopyGmToCbuf(__cbuf__ T *dst, __gm__ T *src, uint8_t sid, uint16_t nBurst, + uint16_t lenBurst, uint16_t srcStride, uint16_t dstStride, pad_t padMode) +{ + DataCopyParams intriParams(nBurst, lenBurst, srcStride, dstStride); + GlobalTensor srcTensor; + srcTensor.SetGlobalBuffer(src); + uint32_t dst_buffer_offset = reinterpret_cast(dst); + uint8_t logicpos = static_cast(TPosition::C1); + LocalTensor dstTensor; + dstTensor = CreateLocalTensor(dst_buffer_offset, logicpos); + DataCopy(dstTensor, srcTensor, intriParams); +} + + +template +__aicore__ inline void SetFpc(__fbuf__ T *src) +{ + LocalTensor tensor; + uint32_t src_buffer_offset = reinterpret_cast(src); + tensor = CreateLocalTensor(src_buffer_offset); + SetFixPipeConfig(tensor); +} + + +template +__aicore__ inline void LoadCbufToCaTranspose(__ca__ T *dst, __cbuf__ T *src, uint16_t indexID, uint8_t repeat, + uint16_t srcStride, uint16_t dstStride, bool addrmode, + uint16_t dstFracStride) +{ + LoadData2dTransposeParams params( + indexID, + repeat, + srcStride, + dstStride, + dstFracStride, + addrmode + ); + uint32_t src_buffer_offset = reinterpret_cast(src); + uint32_t dst_buffer_offset = reinterpret_cast(dst); + uint8_t src_logicpos = static_cast(TPosition::C1); + uint8_t dst_logicpos = static_cast(TPosition::A2); + LocalTensor srcTensor; + LocalTensor dstTensor; + srcTensor = CreateLocalTensor(src_buffer_offset, src_logicpos); + dstTensor = CreateLocalTensor(dst_buffer_offset, dst_logicpos); + LoadDataWithTranspose(dstTensor, srcTensor, params); +} + +template +__aicore__ inline void LoadCbufToCbTranspose(__cb__ T *dst, __cbuf__ T *src, uint16_t indexID, uint8_t repeat, + uint16_t srcStride, uint16_t dstStride, bool addrmode, + uint16_t dstFracStride) +{ + LoadData2dTransposeParams params( + indexID, + repeat, + srcStride, + dstStride, + dstFracStride, + addrmode + ); + uint32_t src_buffer_offset = reinterpret_cast(src); + uint32_t dst_buffer_offset = reinterpret_cast(dst); + uint8_t src_logicpos = static_cast(TPosition::C1); + uint8_t dst_logicpos = static_cast(TPosition::B2); + LocalTensor srcTensor; + LocalTensor dstTensor; + srcTensor = CreateLocalTensor(src_buffer_offset, src_logicpos); + dstTensor = CreateLocalTensor(dst_buffer_offset, dst_logicpos); + LoadDataWithTranspose(dstTensor, srcTensor, params); +} + +template +__aicore__ inline void LoadCbufToCa(__ca__ T *dst, __cbuf__ T *src, uint16_t baseIdx, uint8_t repeat, + uint16_t srcStride, uint16_t dstStride, uint8_t sid, bool transpose, + uint8_t addr_cal_mode) +{ + LoadData2dParams params( + baseIdx, + repeat, + srcStride, + sid, + dstStride, + transpose, + addr_cal_mode + ); + uint32_t src_buffer_offset = reinterpret_cast(src); + uint32_t dst_buffer_offset = reinterpret_cast(dst); + uint8_t src_logicpos = static_cast(TPosition::C1); + uint8_t dst_logicpos = static_cast(TPosition::A2); + LocalTensor srcTensor; + LocalTensor dstTensor; + srcTensor = CreateLocalTensor(src_buffer_offset, src_logicpos); + dstTensor = CreateLocalTensor(dst_buffer_offset, dst_logicpos); + LoadData(dstTensor, srcTensor, params); +} + + +template +__aicore__ inline void LoadCbufToCb(__cb__ T *dst, __cbuf__ T *src, uint16_t baseIdx, uint8_t repeat, + uint16_t srcStride, uint16_t dstStride, uint8_t sid, bool transpose, + uint8_t addr_cal_mode) +{ + LoadData2dParams params( + baseIdx, + repeat, + srcStride, + sid, + dstStride, + transpose, + addr_cal_mode + ); + uint32_t src_buffer_offset = reinterpret_cast(src); + uint32_t dst_buffer_offset = reinterpret_cast(dst); + uint8_t src_logicpos = static_cast(TPosition::C1); + uint8_t dst_logicpos = static_cast(TPosition::B2); + LocalTensor srcTensor; + LocalTensor dstTensor; + srcTensor = CreateLocalTensor(src_buffer_offset, src_logicpos); + dstTensor = CreateLocalTensor(dst_buffer_offset, dst_logicpos); + LoadData(dstTensor, srcTensor, params); +} + + +__aicore__ inline void GetBlockIdx(int32_t loop_idx, int32_t m_loop, int32_t n_loop, int32_t swizzl_direction, + int32_t swizzl_count, int64_t &m_idx, int64_t &n_idx) +{ + uint32_t in_batch_idx = loop_idx % (m_loop * n_loop); + if (swizzl_direction == 0) { + uint32_t tile_block_loop = (m_loop + swizzl_count - 1) / swizzl_count; + uint32_t tile_block_idx = in_batch_idx / (swizzl_count * n_loop); + uint32_t in_tile_block_idx = in_batch_idx % (swizzl_count * n_loop); + + uint32_t n_row = swizzl_count; + if (tile_block_idx == tile_block_loop - 1) { + n_row = m_loop - swizzl_count * tile_block_idx; + } + m_idx = tile_block_idx * swizzl_count + in_tile_block_idx % n_row; + n_idx = in_tile_block_idx / n_row; + if (tile_block_idx % TILE_BLOCK_MOD != 0) { + n_idx = n_loop - n_idx - 1; + } + } else if (swizzl_direction == 1) { + uint32_t tile_block_loop = (n_loop + swizzl_count - 1) / swizzl_count; + uint32_t tile_block_idx = in_batch_idx / (swizzl_count * m_loop); + uint32_t in_tile_block_idx = in_batch_idx % (swizzl_count * m_loop); + + uint32_t n_col = swizzl_count; + if (tile_block_idx == tile_block_loop - 1) { + n_col = n_loop - swizzl_count * tile_block_idx; + } + m_idx = in_tile_block_idx / n_col; + n_idx = tile_block_idx * swizzl_count + in_tile_block_idx % n_col; + if (tile_block_idx % TILE_BLOCK_MOD != 0) { + m_idx = m_loop - m_idx - 1; + } + } +} + +template +__aicore__ inline void FFTSCrossCoreSync(uint64_t mode, uint64_t flag_id) +{ + uint64_t config = 1 | (mode << 4) | (flag_id << 8); + ffts_cross_core_sync(pipe, config); +} + + +template +__aicore__ GlobalTensor CreateGlobalTensor(__gm__ T *addr) +{ + GlobalTensor tensor; + tensor.SetGlobalBuffer(addr); + return tensor; +} + +#endif // MATMUL_ALLREDUCE_ADD_RMSNORM_H diff --git a/csrc/ascend/moe/CMakeLists.txt b/csrc/ascend/moe/CMakeLists.txt new file mode 100644 index 000000000..22a6bd86b --- /dev/null +++ b/csrc/ascend/moe/CMakeLists.txt @@ -0,0 +1,21 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB SUBDIRECTORIES LIST_DIRECTORIES true RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +foreach(SUBDIR ${SUBDIRECTORIES}) + if(DEFINED ASCEND950_SKIP_OPS AND ${SUBDIR} IN_LIST ASCEND950_SKIP_OPS) + message(STATUS "Skipping ${SUBDIR} because ASCEND950 is not available.") + continue() + endif() + + if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${SUBDIR}/CMakeLists.txt) + add_subdirectory(${SUBDIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/add_rms_norm_bias/CMakeLists.txt b/csrc/ascend/moe/add_rms_norm_bias/CMakeLists.txt new file mode 100644 index 000000000..298ed5677 --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/add_rms_norm_bias/add_rms_norm_bias_torch_adpt.h b/csrc/ascend/moe/add_rms_norm_bias/add_rms_norm_bias_torch_adpt.h new file mode 100644 index 000000000..748117394 --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/add_rms_norm_bias_torch_adpt.h @@ -0,0 +1,53 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef ADD_RMS_NORM_BIAS_TORCH_ADPT_H +#define ADD_RMS_NORM_BIAS_TORCH_ADPT_H + +namespace vllm_fl { + +std::tuple npu_add_rms_norm_bias( + const at::Tensor& x1, + const at::Tensor& x2, + const at::Tensor& gamma, + const c10::optional &beta, + double epsilon) +{ + int64_t dim_x = x1.dim(); + int64_t dim_gamma = gamma.dim(); + int64_t diff = dim_x - dim_gamma; + std::vector new_shape; + at::Tensor rstd; + + if (diff > 0) { + new_shape.reserve(dim_x); + auto x1_sizes = x1.sizes(); + for (int64_t i = 0; i < diff; ++i) { + new_shape.push_back(x1_sizes[i]); + } + for (int64_t i = 0; i < dim_gamma; ++i) { + new_shape.push_back(1); + } + } else { + new_shape.assign(dim_x, 1); + } + rstd = at::empty(new_shape, x1.options().dtype(at::kFloat)); + at::Tensor y = at::empty(x1.sizes(), x1.options()); + at::Tensor x = at::empty(x1.sizes(), x1.options()); + EXEC_NPU_CMD(aclnnAddRmsNormBias, x1, x2, gamma, beta, epsilon, y, rstd, x); + return std::tuple(y, rstd, x); +} +} +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/add_rms_norm_bias/op_host/CMakeLists.txt b/csrc/ascend/moe/add_rms_norm_bias/op_host/CMakeLists.txt new file mode 100644 index 000000000..50d68d58b --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/op_host/CMakeLists.txt @@ -0,0 +1,22 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + add_rms_norm_bias_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME AddRmsNormBias + OPTIONS + --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE add_rms_norm_bias ACLNNTYPE aclnn) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/moe/add_rms_norm_bias/op_host/add_rms_norm_bias_def.cpp b/csrc/ascend/moe/add_rms_norm_bias/op_host/add_rms_norm_bias_def.cpp new file mode 100644 index 000000000..95824de0f --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/op_host/add_rms_norm_bias_def.cpp @@ -0,0 +1,71 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_bias_def.cpp + * \brief + */ +#include "register/op_def_registry.h" + +namespace ops { +class AddRmsNormBias : public OpDef { +public: + explicit AddRmsNormBias(const char* name) : OpDef(name) + { + this->Input("x1") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("x2") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("gamma") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("beta") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("rstd") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Attr("epsilon").AttrType(OPTIONAL).Float(1e-6); + + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + } +}; +OP_ADD(AddRmsNormBias); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/moe/add_rms_norm_bias/op_host/add_rms_norm_bias_infershape.cpp b/csrc/ascend/moe/add_rms_norm_bias/op_host/add_rms_norm_bias_infershape.cpp new file mode 100644 index 000000000..60335ab9c --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/op_host/add_rms_norm_bias_infershape.cpp @@ -0,0 +1,84 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_bias_infershape.cpp + * \brief + */ +#include "tiling_base/error_log.h" +#include "util/shape_util.h" +#include "register/op_impl_registry.h" + +static constexpr int IDX_0 = 0; +static constexpr int IDX_1 = 1; +static constexpr int IDX_2 = 2; + +using namespace ge; +using namespace Ops::Base; + +namespace ops { + +static ge::graphStatus InferShape4AddRmsNormBias(gert::InferShapeContext* context) +{ + OP_LOGD(context, "Begin to do InferShape4AddRmsNormBias"); + + // get input shapes + const gert::Shape* x1Shape = context->GetInputShape(IDX_0); + OP_CHECK_NULL_WITH_CONTEXT(context, x1Shape); + const gert::Shape* gammaShape = context->GetInputShape(IDX_2); + OP_CHECK_NULL_WITH_CONTEXT(context, gammaShape); + // get output shapes + gert::Shape* yShape = context->GetOutputShape(IDX_0); + gert::Shape* rstdShape = context->GetOutputShape(IDX_1); + gert::Shape* xShape = context->GetOutputShape(IDX_2); + OP_CHECK_NULL_WITH_CONTEXT(context, yShape); + OP_CHECK_NULL_WITH_CONTEXT(context, rstdShape); + OP_CHECK_NULL_WITH_CONTEXT(context, xShape); + *yShape = *x1Shape; + *xShape = *x1Shape; + + size_t xDimNum = x1Shape->GetDimNum(); + size_t gammaDimNum = gammaShape->GetDimNum(); + + if (IsUnknownRank(*x1Shape) || IsUnknownRank(*gammaShape)) { + SetUnknownRank(*rstdShape); + OP_LOGD(context, "End to do InferShape4AddRmsNormBias with unknown rank."); + return GRAPH_SUCCESS; + } + + OP_CHECK_IF( + xDimNum < gammaDimNum, OP_LOGE(context, "x dim num should not be smaller than gamma dim num."), + return GRAPH_FAILED); + + rstdShape->SetDimNum(xDimNum); + for (size_t i = 0; i < xDimNum; i++) { + if (i < xDimNum - gammaDimNum) { + rstdShape->SetDim(i, x1Shape->GetDim(i)); + } else { + rstdShape->SetDim(i, 1); + } + } + + OP_LOGD(context, "End to do InferShape4AddRmsNormBias"); + return GRAPH_SUCCESS; +} + +static graphStatus InferDataType4AddRmsNormBias(gert::InferDataTypeContext* context) +{ + OP_LOGD(context, "Begin to do InferDataType4AddRmsNormBias"); + context->SetOutputDataType(IDX_0, context->GetInputDataType(IDX_0)); + context->SetOutputDataType(IDX_1, DT_FLOAT); + context->SetOutputDataType(IDX_2, context->GetInputDataType(IDX_0)); + OP_LOGD(context, "End to do InferDataType4AddRmsNormBias"); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(AddRmsNormBias).InferShape(InferShape4AddRmsNormBias).InferDataType(InferDataType4AddRmsNormBias); +} // namespace ops diff --git a/csrc/ascend/moe/add_rms_norm_bias/op_host/add_rms_norm_bias_tiling.cpp b/csrc/ascend/moe/add_rms_norm_bias/op_host/add_rms_norm_bias_tiling.cpp new file mode 100644 index 000000000..f88381ae1 --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/op_host/add_rms_norm_bias_tiling.cpp @@ -0,0 +1,443 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_bias_tiling.cpp + * \brief + */ +#include "add_rms_norm_bias_tiling.h" +#include "log/ops_log.h" + +namespace optiling { +constexpr uint32_t DTYPE_KEY_FP16 = 1; +constexpr uint32_t DTYPE_KEY_FP32 = 2; +constexpr uint32_t DTYPE_KEY_BF16 = 3; +constexpr uint32_t UB_USED = 1024; +constexpr uint32_t UB_FACTOR_B16 = 12288; +constexpr uint32_t UB_FACTOR_B32 = 10240; +constexpr uint32_t UB_FACTOR_B16_CUTD = 12096; +constexpr uint32_t UB_FACTOR_B32_CUTD = 9696; + +constexpr uint32_t UB_FACTOR_B32_WITH_BETA = 9216; +constexpr uint32_t UB_FACTOR_B16_WITH_BETA = 11264; +constexpr uint32_t UB_FACTOR_B32_CUTD_WITH_BETA = 8096; +constexpr uint32_t UB_FACTOR_B16_CUTD_WITH_BETA = 10752; +constexpr uint32_t SMALL_REDUCE_NUM_WITH_BETA = 1600; +constexpr uint32_t FP32_WEIGHT_WITH_BETA = 28; +constexpr uint32_t OTHER_WEIGHT_WITH_BETA = 20; +constexpr size_t NUM_WITH_BETA = 4; + +constexpr uint32_t BLOCK_ALIGN_NUM = 16; +constexpr uint32_t FLOAT_BLOCK_ALIGN_NUM = 8; +constexpr uint32_t SMALL_REDUCE_NUM = 2000; +constexpr uint32_t MODE_NORMAL = 0; +constexpr uint32_t MODE_SPLIT_D = 1; +constexpr uint32_t MODE_MERGE_N = 2; +constexpr uint32_t MODE_SINGLE_N = 3; +constexpr uint32_t MODE_MULTI_N = 4; +constexpr int32_t INPUT_X1_INDEX = 0; +constexpr int32_t INPUT_X2_INDEX = 1; +constexpr int32_t INPUT_GAMMA_INDEX = 2; +constexpr int32_t INPUT_BETA_INDEX = 3; +constexpr int32_t OUTPUT_Y_INDEX = 0; +constexpr int32_t OUTPUT_RSTD_INDEX = 1; +constexpr int32_t OUTPUT_X_INDEX = 2; +constexpr size_t MAX_DIM_NUM = 8; +constexpr size_t MIN_DIM_X = 1; +constexpr size_t MIN_DIM_GAMMA = 1; +constexpr size_t FP32_WEIGHT = 24; +constexpr size_t OTHER_WEIGHT = 18; +constexpr size_t DIV_FACTOR = 260; +constexpr size_t FLOAT_PER_REPEAT = 64; +constexpr size_t USE_SIZE = 256; +constexpr size_t NUM = 2; +constexpr int32_t TEN = 10; + +constexpr int32_t PERFORMANC_DIM_ZERO = 0; +constexpr int32_t PERFORMANC_DIM_ONE = 1; +constexpr int32_t PERFORMANC_DIM_TWO = 2; +constexpr int32_t PERFORMANC_DIM_THREE = 3; +constexpr int32_t PERFORMANC_DIM_ONE_MAX = 512; +constexpr int32_t PERFORMANC_DIM_TWO_MAX = 8; +constexpr int32_t PERFORMANC_DIM_THREE_MAX = 5120; + +platform_ascendc::SocVersion addRmsNormBiasSocVersion; + +uint8_t getPerformanceFlag(uint32_t num_col, gert::Shape x_shape, gert::Shape gamma_shape, uint32_t xDtypeKey) +{ + uint8_t isPerformance = 0; + if(addRmsNormBiasSocVersion != platform_ascendc::SocVersion::ASCEND910B) { + return isPerformance; + } + size_t xDimNum = x_shape.GetDimNum(); + size_t gammaDimNum = gamma_shape.GetDimNum(); + bool dimOK = ((xDimNum == PERFORMANC_DIM_TWO || xDimNum == PERFORMANC_DIM_THREE) && gammaDimNum == PERFORMANC_DIM_ONE); + bool sizeOk = num_col <= PERFORMANC_DIM_THREE_MAX && + ((xDimNum == PERFORMANC_DIM_TWO && x_shape.GetDim(PERFORMANC_DIM_ZERO) <= PERFORMANC_DIM_ONE_MAX) || + (xDimNum == PERFORMANC_DIM_THREE && x_shape.GetDim(PERFORMANC_DIM_ZERO) <= PERFORMANC_DIM_ONE_MAX && x_shape.GetDim(PERFORMANC_DIM_ONE) <= PERFORMANC_DIM_TWO_MAX)); + bool dtypeOk = (xDtypeKey == DTYPE_KEY_FP16 || xDtypeKey == DTYPE_KEY_BF16); + if(dimOK && sizeOk && dtypeOk) { + isPerformance = 1; + } + return isPerformance; +} + +static void SetByDtype(ge::DataType dataType, uint32_t& dtypeKey, uint32_t& dataPerBlock) +{ + switch (dataType) { + case ge::DT_FLOAT16: + dtypeKey = DTYPE_KEY_FP16; + dataPerBlock = BLOCK_ALIGN_NUM; + break; + case ge::DT_BF16: + dtypeKey = DTYPE_KEY_BF16; + dataPerBlock = BLOCK_ALIGN_NUM; + break; + default: + dtypeKey = DTYPE_KEY_FP32; + dataPerBlock = FLOAT_BLOCK_ALIGN_NUM; + break; + } +} + +static bool CheckInputOutputDim(const gert::TilingContext* context) +{ + const gert::StorageShape* x1_shape = context->GetInputShape(INPUT_X1_INDEX); + const gert::StorageShape* x2_shape = context->GetInputShape(INPUT_X2_INDEX); + const gert::StorageShape* gamma_shape = context->GetInputShape(INPUT_GAMMA_INDEX); + const gert::StorageShape* y_shape = context->GetOutputShape(OUTPUT_Y_INDEX); + const gert::StorageShape* rstd_shape = context->GetOutputShape(OUTPUT_RSTD_INDEX); + const gert::StorageShape* x_shape = context->GetOutputShape(OUTPUT_X_INDEX); + + OP_CHECK_NULL_WITH_CONTEXT(context, x1_shape); + OP_CHECK_NULL_WITH_CONTEXT(context, x2_shape); + OP_CHECK_NULL_WITH_CONTEXT(context, gamma_shape); + OP_CHECK_NULL_WITH_CONTEXT(context, y_shape); + OP_CHECK_NULL_WITH_CONTEXT(context, rstd_shape); + OP_CHECK_NULL_WITH_CONTEXT(context, x_shape); + + size_t x1DimNum = x1_shape->GetStorageShape().GetDimNum(); + size_t x2DimNum = x2_shape->GetStorageShape().GetDimNum(); + size_t gammaDimNum = gamma_shape->GetStorageShape().GetDimNum(); + size_t yDimNum = y_shape->GetStorageShape().GetDimNum(); + size_t rstdDimNum = rstd_shape->GetStorageShape().GetDimNum(); + size_t xDimNum = x_shape->GetStorageShape().GetDimNum(); + + OP_CHECK_IF( + x1DimNum > MAX_DIM_NUM || x1DimNum < MIN_DIM_X, + OP_LOGE(context, "Input x1's dim num should not greater than 8 or smaller than 1."), + return false); + OP_CHECK_IF( + gammaDimNum > MAX_DIM_NUM || gammaDimNum < MIN_DIM_GAMMA, + OP_LOGE(context, "Input gamma's dim num should not greater than 8 or smaller than 1."), + return false); + OP_CHECK_IF( + x1DimNum != yDimNum, OP_LOGE(context, "Input x's dim num must equal to output y's dim num."), + return false); + + OP_CHECK_IF( + x1DimNum != x2DimNum, + OP_LOGE(context, "Input x2/x1 shape invalid, dim num is not equal x1 dim."), return false); + OP_CHECK_IF( + (yDimNum != xDimNum) || (xDimNum != x1DimNum) || (rstdDimNum != x1DimNum), + OP_LOGE(context, "Output y/x/rstd shape invalid, dim num is not equal x1 dim."), return false); + OP_CHECK_IF( + x1DimNum < gammaDimNum, OP_LOGE(context, "X1 dim num should not be smaller than gamma dim num."), + return false); + return true; +} + +static bool CheckInputOutputShape(const gert::TilingContext* context) +{ + OP_CHECK_IF(!CheckInputOutputDim(context), OP_LOGE(context, "Input Dim invalid."), return false); + const gert::StorageShape* x1_shape = context->GetInputShape(INPUT_X1_INDEX); + const gert::StorageShape* x2_shape = context->GetInputShape(INPUT_X2_INDEX); + const gert::StorageShape* gamma_shape = context->GetInputShape(INPUT_GAMMA_INDEX); + const gert::StorageShape* y_shape = context->GetOutputShape(OUTPUT_Y_INDEX); + const gert::StorageShape* rstd_shape = context->GetOutputShape(OUTPUT_RSTD_INDEX); + const gert::StorageShape* x_shape = context->GetOutputShape(OUTPUT_X_INDEX); + + OP_CHECK_NULL_WITH_CONTEXT(context, x1_shape); + OP_CHECK_NULL_WITH_CONTEXT(context, x2_shape); + OP_CHECK_NULL_WITH_CONTEXT(context, gamma_shape); + OP_CHECK_NULL_WITH_CONTEXT(context, y_shape); + OP_CHECK_NULL_WITH_CONTEXT(context, rstd_shape); + OP_CHECK_NULL_WITH_CONTEXT(context, x_shape); + + size_t x1DimNum = x1_shape->GetStorageShape().GetDimNum(); + size_t gammaDimNum = gamma_shape->GetStorageShape().GetDimNum(); + + for (uint32_t i = 0; i < x1DimNum; i++) { + OP_CHECK_IF( + x1_shape->GetStorageShape().GetDim(i) == 0, OP_LOGE(context, "Input x1 shape can not be 0."), + return false); + OP_CHECK_IF( + x2_shape->GetStorageShape().GetDim(i) != x1_shape->GetStorageShape().GetDim(i), + OP_LOGE(context, "Input x2/x1 shape invalid, shape is not equal x1 shape."), return false); + OP_CHECK_IF( + (y_shape->GetStorageShape().GetDim(i) != x1_shape->GetStorageShape().GetDim(i)) || + (x_shape->GetStorageShape().GetDim(i) != x1_shape->GetStorageShape().GetDim(i)), + OP_LOGE(context, "Input y/x shape invalid, shape is not equal x1 shape."), return false); + } + for (uint32_t i = 0; i < x1DimNum - gammaDimNum; i++) { + OP_CHECK_IF( + rstd_shape->GetStorageShape().GetDim(i) != x2_shape->GetStorageShape().GetDim(i), + OP_LOGE(context, "Output rstd shape invalid, shape is not equal x1 first few dim."), + return false); + } + for (uint32_t i = 0; i < gammaDimNum; i++) { + OP_CHECK_IF( + gamma_shape->GetStorageShape().GetDim(i) != x1_shape->GetStorageShape().GetDim(x1DimNum - gammaDimNum + i), + OP_LOGE(context, "Input gamma shape invalid, gamma shape is not equal x1 last few dim."), + return false); + OP_CHECK_IF( + rstd_shape->GetStorageShape().GetDim(x1DimNum - 1 - i) != 1, + OP_LOGE(context, "Output rstd shape invalid, last few dim is not equal to 1."), + return false); + } + return true; +} + +static void GetCompileParameters( + gert::TilingContext* context, uint32_t& numCore, uint64_t& ubSize) +{ + auto ptrCompileInfo = reinterpret_cast(context->GetCompileInfo()); + if (ptrCompileInfo == nullptr) { + auto ascendc_platform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + addRmsNormBiasSocVersion = ascendc_platform.GetSocVersion(); + numCore = ascendc_platform.GetCoreNumAiv(); + ascendc_platform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); + } else { + numCore = ptrCompileInfo->totalCoreNum; + ubSize = ptrCompileInfo->totalUbSize; + addRmsNormBiasSocVersion = ptrCompileInfo->socVersion; + } + ubSize -= UB_USED; +} + +static void CalculateRowAndColParameters(gert::TilingContext* context, uint32_t& numRow, uint32_t& numCol) +{ + const gert::Shape x1_shape = context->GetInputShape(0)->GetStorageShape(); + const size_t gammaIndex = 2; + const gert::Shape gamma_shape = context->GetInputShape(gammaIndex)->GetStorageShape(); + numCol = gamma_shape.GetShapeSize(); + + const size_t x1DimNum = x1_shape.GetDimNum(); + const size_t gammaDimNum = gamma_shape.GetDimNum(); + numRow = 1U; + for (size_t i = 0; i < x1DimNum - gammaDimNum; ++i) { + numRow *= x1_shape.GetDim(i); + } +} + +static ge::graphStatus GetEpsilonParameter(gert::TilingContext* context, float& epsilon) +{ + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + epsilon = *attrs->GetFloat(0); + OP_CHECK_IF( + epsilon < 0, OP_LOGE(context, "Epsilon less than zero, please check."), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +static void CalculateBlockParameters( + uint32_t numRow, uint32_t numCore, uint32_t& blockFactor, uint32_t& latsBlockFactor, uint32_t& useCoreNum) +{ + blockFactor = 1U; + uint32_t tileNum = CeilDiv(numRow, numCore * blockFactor); + blockFactor *= tileNum; + useCoreNum = CeilDiv(numRow, blockFactor); + latsBlockFactor = numRow - blockFactor * (useCoreNum - 1); +} + +static ge::DataType SetDataTypeParameters(gert::TilingContext* context, uint32_t& dtype_key, uint32_t& data_per_block) +{ + auto data_type = context->GetInputDesc(0)->GetDataType(); + dtype_key = DTYPE_KEY_FP16; + SetByDtype(data_type, dtype_key, data_per_block); + return data_type; +} + +static void DetermineModeParameters( + AddRMSNormBiasTilingData* tiling, + uint32_t numCol, uint32_t& ubFactor, uint32_t& rowFactor, uint32_t blockFactor, + uint32_t latsBlockFactor, ge::DataType dataType, uint32_t dtypKey, uint64_t ubSize, + uint32_t dataPerBlock, uint32_t numColAlign, uint32_t& modeKey, uint32_t isPerformance) +{ + if (numCol > ubFactor) { + modeKey = MODE_SPLIT_D; + ubFactor = tiling->get_nullptr_beta() == 1 ? ((dataType == ge::DT_FLOAT) ? UB_FACTOR_B32_CUTD : UB_FACTOR_B16_CUTD) : ((dataType == ge::DT_FLOAT) ? UB_FACTOR_B32_CUTD_WITH_BETA : UB_FACTOR_B16_CUTD_WITH_BETA); + uint32_t colTileNum = CeilDiv(numCol, ubFactor); + ubFactor = CeilDiv(numCol, colTileNum * dataPerBlock) * dataPerBlock; + } else if (blockFactor == 1 && addRmsNormBiasSocVersion != platform_ascendc::SocVersion::ASCEND310P) { + modeKey = MODE_SINGLE_N; + } else if (((tiling->get_nullptr_beta() == 1 && numColAlign <= SMALL_REDUCE_NUM) || (tiling->get_nullptr_beta() == 0 && numColAlign <= SMALL_REDUCE_NUM_WITH_BETA)) && addRmsNormBiasSocVersion != platform_ascendc::SocVersion::ASCEND310P) { + modeKey = MODE_MERGE_N; + uint64_t numColAlignWeight = tiling->get_nullptr_beta() == 1 ? ((dtypKey == DTYPE_KEY_FP32) ? FP32_WEIGHT : OTHER_WEIGHT) : ((dtypKey == DTYPE_KEY_FP32) ? FP32_WEIGHT_WITH_BETA : OTHER_WEIGHT_WITH_BETA); + rowFactor = static_cast(ubSize) / + (numColAlign * static_cast(numColAlignWeight) + static_cast(DIV_FACTOR)); + ubFactor = rowFactor * numColAlign; + + uint32_t mulLoopFp32 = numColAlign / 64; + uint32_t mulTailFp32 = numColAlign - mulLoopFp32 * 64; + uint8_t dstRepStrideFp32 = numColAlign / 8; + + uint32_t mulLoopFp16 = numColAlign / 128; + uint32_t mulTailFp16 = numColAlign - mulLoopFp16 * 128; + uint8_t dstRepStrideFp16 = numColAlign / 16; + + tiling->set_is_performance(isPerformance); + tiling->set_mul_loop_fp32(mulLoopFp32); + tiling->set_mul_tail_fp32(mulTailFp32); + tiling->set_dst_rep_stride_fp32(dstRepStrideFp32); + tiling->set_mul_loop_fp16(mulLoopFp16); + tiling->set_mul_tail_fp16(mulTailFp16); + tiling->set_dst_rep_stride_fp16(dstRepStrideFp16); + } else if ((dataType == ge::DT_FLOAT16 || isPerformance == 1) && numCol == numColAlign) { + modeKey = MODE_MULTI_N; + rowFactor = (static_cast(ubSize) - static_cast(USE_SIZE) - + numColAlign * static_cast(tiling->get_nullptr_beta() == 1 ? NUM : NUM_WITH_BETA)) / + (numColAlign * BLOCK_ALIGN_NUM + static_cast(FLOAT_PER_REPEAT)); + ubFactor = rowFactor * numColAlign; + if (rowFactor == 0U) { + modeKey = MODE_NORMAL; + rowFactor = FLOAT_PER_REPEAT; + ubFactor = UB_FACTOR_B16; + } + } + uint32_t rowLoop = CeilDiv(blockFactor, rowFactor); + uint32_t lastBlockRowLoop = CeilDiv(latsBlockFactor, rowFactor); + uint32_t rowTail = blockFactor - (rowLoop - 1) * rowFactor; + uint32_t lastBlockRowTail = latsBlockFactor - (lastBlockRowLoop - 1) * rowFactor; + tiling->set_row_loop(rowLoop); + tiling->set_last_block_row_loop(lastBlockRowLoop); + tiling->set_row_tail(rowTail); + tiling->set_last_block_row_tail(lastBlockRowTail); +} + +static void SetTilingParameters( + AddRMSNormBiasTilingData* tiling, uint32_t num_row, uint32_t num_col, uint32_t numColAlign, + uint32_t block_factor, uint32_t latsBlockFactor, uint32_t row_factor, + uint32_t ub_factor, float epsilon) +{ + const float avg_factor = (num_col == 0) ? 0 : 1.0f / num_col; + tiling->set_num_row(num_row); + tiling->set_num_col(num_col); + tiling->set_num_col_align(numColAlign); + tiling->set_block_factor(block_factor); + tiling->set_last_block_factor(latsBlockFactor); + tiling->set_row_factor(row_factor); + tiling->set_ub_factor(ub_factor); + tiling->set_epsilon(epsilon); + tiling->set_avg_factor(avg_factor); +} + +static void SaveTilingData( + gert::TilingContext* context, AddRMSNormBiasTilingData* tiling, uint32_t dtype_key, uint32_t mode_key) +{ + const uint32_t tiling_key = dtype_key * 10 + mode_key; + context->SetTilingKey(tiling_key); + tiling->SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tiling->GetDataSize()); +} + +static void SetWorkspaceSize(gert::TilingContext* context) +{ + constexpr size_t sysWorkspaceSize = 16 * 1024 * 1024; + constexpr size_t usrSize = 256; + size_t* currentWorkspace = context->GetWorkspaceSizes(1); + currentWorkspace[0] = usrSize + sysWorkspaceSize; +} + +static void LogTilingResults( + gert::TilingContext* context, AddRMSNormBiasTilingData* tiling, uint32_t mode_key, uint32_t dtype_key, + uint32_t use_core_num, float epsilon) +{ + OPS_LOG_I(context, "Tiling Key: %u", dtype_key * TEN + mode_key); + OPS_LOG_I(context, "Block Dim: %u", use_core_num); + OPS_LOG_I(context, "usr Workspace: 256"); + OPS_LOG_I( + context, + "num_row: %d, num_col: %d, block_factor: %d, row_factor: %d, ub_factor: %d, epsilon: %f, avg_factor: %f", + tiling->get_num_row(), tiling->get_num_col(), tiling->get_block_factor(), tiling->get_row_factor(), + tiling->get_ub_factor(), epsilon, tiling->get_avg_factor()); +} + +static ge::graphStatus Tiling4AddRmsNormBias(gert::TilingContext* context) +{ + OP_LOGI("Tiling4AddRmsNormBias", "Enter Tiling4AddRmsNormBias"); + OPS_LOG_D(context, "Tiling4AddRmsNormBias1 running. \n"); + OP_CHECK_IF( + !CheckInputOutputShape(context), OP_LOGE(context, "Input shape invalid."), + return ge::GRAPH_FAILED); + + AddRMSNormBiasTilingData tiling; + + auto betaDesc = context->GetOptionalInputDesc(INPUT_BETA_INDEX); + tiling.set_nullptr_beta(betaDesc == nullptr ? 1 : 0); + + uint32_t num_core; + uint64_t ub_size; + GetCompileParameters(context, num_core, ub_size); + uint32_t num_row; + uint32_t num_col; + CalculateRowAndColParameters(context, num_row, num_col); + float epsilon = 0; + GetEpsilonParameter(context, epsilon); + if (epsilon < 0) { + return ge::GRAPH_FAILED; + } + uint32_t block_factor; + uint32_t latsBlockFactor; + uint32_t use_core_num; + CalculateBlockParameters(num_row, num_core, block_factor, latsBlockFactor, use_core_num); + context->SetBlockDim(use_core_num); + uint32_t dtype_key; + uint32_t data_per_block; + ge::DataType data_type = SetDataTypeParameters(context, dtype_key, data_per_block); + uint32_t mode_key = MODE_NORMAL; + uint32_t row_factor = 64; + uint32_t ub_factor = betaDesc == nullptr ? ((dtype_key == DTYPE_KEY_FP32) ? UB_FACTOR_B32 : UB_FACTOR_B16) : ((dtype_key == DTYPE_KEY_FP32) ? UB_FACTOR_B32_WITH_BETA : UB_FACTOR_B16_WITH_BETA); + uint32_t numColAlign = CeilDiv(num_col, data_per_block) * data_per_block; + const gert::Shape x1_shape = context->GetInputShape(0)->GetStorageShape(); + const gert::Shape gamma_shape = context->GetInputShape(2)->GetStorageShape(); + uint8_t isPerformance = getPerformanceFlag(num_col, x1_shape, gamma_shape, dtype_key); + DetermineModeParameters( + &tiling, + num_col, ub_factor, row_factor, block_factor, latsBlockFactor, + data_type, dtype_key, ub_size, data_per_block, + numColAlign, mode_key, isPerformance); + SetTilingParameters(&tiling, num_row, num_col, numColAlign, block_factor, latsBlockFactor, row_factor, ub_factor, epsilon); + SaveTilingData(context, &tiling, dtype_key, mode_key); + SetWorkspaceSize(context); + LogTilingResults(context, &tiling, mode_key, dtype_key, use_core_num, epsilon); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus TilingPrepare4AddRmsNormBias(gert::TilingParseContext* context) +{ + OPS_LOG_D(context, "TilingPrepare4AddRmsNormBias running. \n"); + OP_LOGI(context, "TilingPrepare4AddRmsNormBias running."); + auto compileInfo = context->GetCompiledInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo); + auto platformInfo = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + + compileInfo->socVersion = ascendcPlatform.GetSocVersion(); + compileInfo->totalCoreNum = ascendcPlatform.GetCoreNumAiv(); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, compileInfo->totalUbSize); + + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(AddRmsNormBias).Tiling(Tiling4AddRmsNormBias).TilingParse(TilingPrepare4AddRmsNormBias); + +} // namespace optiling diff --git a/csrc/ascend/moe/add_rms_norm_bias/op_host/add_rms_norm_bias_tiling.h b/csrc/ascend/moe/add_rms_norm_bias/op_host/add_rms_norm_bias_tiling.h new file mode 100644 index 000000000..3e370ce40 --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/op_host/add_rms_norm_bias_tiling.h @@ -0,0 +1,53 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OPS_BUILT_IN_OP_TILING_RUNTIME_ADD_RMS_NORM_H_ +#define OPS_BUILT_IN_OP_TILING_RUNTIME_ADD_RMS_NORM_H_ +#include "register/tilingdata_base.h" +#include "tiling_base/error_log.h" +#include "register/op_impl_registry.h" +#include "tiling/platform/platform_ascendc.h" +#include "platform/platform_infos_def.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(AddRMSNormBiasTilingData) +TILING_DATA_FIELD_DEF(uint32_t, num_row); +TILING_DATA_FIELD_DEF(uint32_t, num_col); +TILING_DATA_FIELD_DEF(uint32_t, block_factor); +TILING_DATA_FIELD_DEF(uint32_t, row_factor); +TILING_DATA_FIELD_DEF(uint32_t, ub_factor); +TILING_DATA_FIELD_DEF(float, epsilon); +TILING_DATA_FIELD_DEF(float, avg_factor); +TILING_DATA_FIELD_DEF(uint32_t, num_col_align); +TILING_DATA_FIELD_DEF(uint32_t, last_block_factor); +TILING_DATA_FIELD_DEF(uint32_t, row_loop); +TILING_DATA_FIELD_DEF(uint32_t, last_block_row_loop); +TILING_DATA_FIELD_DEF(uint32_t, row_tail); +TILING_DATA_FIELD_DEF(uint32_t, last_block_row_tail); +TILING_DATA_FIELD_DEF(uint32_t, mul_loop_fp32); +TILING_DATA_FIELD_DEF(uint32_t, mul_tail_fp32); +TILING_DATA_FIELD_DEF(uint32_t, dst_rep_stride_fp32); +TILING_DATA_FIELD_DEF(uint32_t, mul_loop_fp16); +TILING_DATA_FIELD_DEF(uint32_t, mul_tail_fp16); +TILING_DATA_FIELD_DEF(uint32_t, dst_rep_stride_fp16); +TILING_DATA_FIELD_DEF(uint32_t, is_performance); +TILING_DATA_FIELD_DEF(uint32_t, nullptr_beta); +END_TILING_DATA_DEF; + +struct AddRmsNormBiasCompileInfo { + uint32_t totalCoreNum = 0; + uint64_t totalUbSize = 0; + platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND950; +}; + +REGISTER_TILING_DATA_CLASS(AddRmsNormBias, AddRMSNormBiasTilingData) +} // namespace optiling + +#endif // OPS_BUILT_IN_OP_TILING_RUNTIME_ADD_RMS_NORM_BIAS_H_ diff --git a/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias.cpp b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias.cpp new file mode 100644 index 000000000..c09ed0a16 --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias.cpp @@ -0,0 +1,72 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_bias.cpp + * \brief + */ +#include "add_rms_norm_bias.h" +#include "add_rms_norm_bias_split_d.h" +#include "add_rms_norm_bias_merge_n.h" +#include "add_rms_norm_bias_multi_n.h" +#include "add_rms_norm_bias_single_n.h" + +using namespace AscendC; + +#define GENERAL_OP_IMPL(templateClass, ...) \ + do { \ + templateClass<__VA_ARGS__> op(&pipe); \ + op.Init(x1, x2, gamma, beta, y, rstd, x, &tilingData); \ + op.Process(); \ + } while (0) + +extern "C" __global__ __aicore__ void add_rms_norm_bias( + GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta, GM_ADDR y, GM_ADDR rstd, GM_ADDR x, GM_ADDR workspace, GM_ADDR tiling) +{ + TPipe pipe; + GET_TILING_DATA(tilingData, tiling); + if (TILING_KEY_IS(10)) { + GENERAL_OP_IMPL(KernelAddRmsNormBias, half); + } else if (TILING_KEY_IS(20)) { + GENERAL_OP_IMPL(KernelAddRmsNormBias, float); + } else if (TILING_KEY_IS(30)) { +#if !(defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + GENERAL_OP_IMPL(KernelAddRmsNormBias, bfloat16_t); +#endif + } else if (TILING_KEY_IS(11)) { + GENERAL_OP_IMPL(KernelAddRmsNormBiasSplitD, half); + } else if (TILING_KEY_IS(21)) { + GENERAL_OP_IMPL(KernelAddRmsNormBiasSplitD, float); + } else if (TILING_KEY_IS(31)) { +#if !(defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + GENERAL_OP_IMPL(KernelAddRmsNormBiasSplitD, bfloat16_t); +#endif + } else if (TILING_KEY_IS(12)) { + GENERAL_OP_IMPL(KernelAddRmsNormBiasMergeN, half); + } else if (TILING_KEY_IS(22)) { + GENERAL_OP_IMPL(KernelAddRmsNormBiasMergeN, float); + } else if (TILING_KEY_IS(32)) { +#if !(defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + GENERAL_OP_IMPL(KernelAddRmsNormBiasMergeN, bfloat16_t); +#endif + } else if (TILING_KEY_IS(13)) { + GENERAL_OP_IMPL(KernelAddRmsNormBiasSingleN, half); + } else if (TILING_KEY_IS(23)) { + GENERAL_OP_IMPL(KernelAddRmsNormBiasSingleN, float); + } else if (TILING_KEY_IS(33)) { +#if !(defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + GENERAL_OP_IMPL(KernelAddRmsNormBiasSingleN, bfloat16_t); +#endif + } else if (TILING_KEY_IS(14)) { + GENERAL_OP_IMPL(KernelAddRmsNormBiasMultiN, half); + } else if (TILING_KEY_IS(34)) { + GENERAL_OP_IMPL(KernelAddRmsNormBiasMultiN, bfloat16_t); + } +} \ No newline at end of file diff --git a/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias.h b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias.h new file mode 100644 index 000000000..b7699f6e1 --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias.h @@ -0,0 +1,368 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_bias.h + * \brief add rms norm bias file + */ +#ifndef ADD_RMS_NORM_H_ +#define ADD_RMS_NORM_H_ +#include "./rms_norm_base.h" + +using namespace AscendC; +using namespace RmsNorm; + +template +class KernelAddRmsNormBias { +public: + __aicore__ inline KernelAddRmsNormBias(TPipe* pipe) + { + Ppipe = pipe; + } + __aicore__ inline void Init( + GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta, GM_ADDR y, GM_ADDR rstd, GM_ADDR x, const AddRMSNormBiasTilingData* tiling) + { + ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!"); + this->numRow = tiling->num_row; + this->numCol = tiling->num_col; + this->blockFactor = tiling->block_factor; + this->rowFactor = tiling->row_factor; + this->ubFactor = tiling->ub_factor; + this->epsilon = tiling->epsilon; + this->avgFactor = (numCol != 0) ? (float)1.0 / numCol : 0; + this->nullptrBeta = tiling->nullptr_beta; + + blockIdx_ = GetBlockIdx(); + if (blockIdx_ < GetBlockNum() - 1) { + this->rowWork = blockFactor; + } else if (blockIdx_ == GetBlockNum() - 1) { + this->rowWork = numRow - (GetBlockNum() - 1) * blockFactor; + } + // get start index for current core, core parallel + x1Gm.SetGlobalBuffer((__gm__ T*)x1 + blockIdx_ * blockFactor * numCol, rowWork * numCol); + x2Gm.SetGlobalBuffer((__gm__ T*)x2 + blockIdx_ * blockFactor * numCol, rowWork * numCol); + gammaGm.SetGlobalBuffer((__gm__ T*)gamma, numCol); + if (!this->nullptrBeta) { + betaGm.SetGlobalBuffer((__gm__ T*)beta, numCol); + } + yGm.SetGlobalBuffer((__gm__ T*)y + blockIdx_ * blockFactor * numCol, rowWork * numCol); + rstdGm.SetGlobalBuffer((__gm__ float*)rstd + blockIdx_ * blockFactor, blockFactor); + xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * blockFactor * numCol, rowWork * numCol); + + // pipe alloc memory to queue, the unit is Bytes + Ppipe->InitBuffer(inQueueX, BUFFER_NUM, ubFactor * sizeof(T)); + Ppipe->InitBuffer(inQueueGamma, BUFFER_NUM, ubFactor * sizeof(T)); + if (!this->nullptrBeta) { + Ppipe->InitBuffer(inQueueBeta, BUFFER_NUM, ubFactor * sizeof(T)); + } + Ppipe->InitBuffer(outQueueY, BUFFER_NUM, ubFactor * sizeof(T)); + Ppipe->InitBuffer(outQueueRstd, BUFFER_NUM, rowFactor * sizeof(float)); + + if constexpr (is_same::value || is_same::value) { + Ppipe->InitBuffer(xFp32Buf, ubFactor * sizeof(float)); + } + Ppipe->InitBuffer(sqxBuf, ubFactor * sizeof(float)); + Ppipe->InitBuffer(reduceFp32Buf, NUM_PER_REP_FP32 * sizeof(float)); + } + + __aicore__ inline void Process() + { + CopyInGammaBeta(); + LocalTensor gammaLocal = inQueueGamma.DeQue(); + LocalTensor betaLocal; + if (!this->nullptrBeta) { + betaLocal = inQueueBeta.DeQue(); + } + uint32_t i_o_max = RmsNorm::CeilDiv(rowWork, rowFactor); + uint32_t row_tail = rowWork - (i_o_max - 1) * rowFactor; + + for (uint32_t i_o = 0; i_o < i_o_max - 1; i_o++) { + SubProcess(i_o, rowFactor, gammaLocal, betaLocal); + } + SubProcess(i_o_max - 1, row_tail, gammaLocal, betaLocal); + inQueueGamma.FreeTensor(gammaLocal); + if (!this->nullptrBeta) { + inQueueBeta.FreeTensor(betaLocal); + } + } + + __aicore__ inline void SubProcess(uint32_t i_o, uint32_t calc_row_num, LocalTensor& gammaLocal, LocalTensor& betaLocal) + { + LocalTensor rstdLocal = outQueueRstd.AllocTensor(); + for (uint32_t i_i = 0; i_i < calc_row_num; i_i++) { + uint32_t gm_bias = (i_o * rowFactor + i_i) * numCol; + CopyIn(gm_bias); + Compute(i_i, gammaLocal, betaLocal, rstdLocal); + CopyOutY(gm_bias); + } + outQueueRstd.EnQue(rstdLocal); + CopyOutRstd(i_o, calc_row_num); + } + +private: + __aicore__ inline void CopyIn(uint32_t gm_bias) + { + LocalTensor x1Local_in = inQueueX.AllocTensor(); + LocalTensor x2Local = sqxBuf.Get(); + LocalTensor xLocal = outQueueY.AllocTensor(); + + if constexpr (is_same::value || is_same::value) { + x2Local = x2Local[ubFactor]; + } + + DataCopyCustom(x1Local_in, x1Gm[gm_bias], numCol); + DataCopyCustom(x2Local, x2Gm[gm_bias], numCol); + inQueueX.EnQue(x1Local_in); + auto x1Local = inQueueX.DeQue(); + + if constexpr (is_same::value) { + LocalTensor x1_fp32 = xFp32Buf.Get(); + Add(xLocal, x1Local, x2Local, numCol); + PipeBarrier(); + Cast(x1_fp32, xLocal, RoundMode::CAST_NONE, numCol); + PipeBarrier(); + } else if constexpr (is_same::value) { + LocalTensor x1_fp32 = xFp32Buf.Get(); + LocalTensor x2_fp32 = sqxBuf.Get(); + Cast(x1_fp32, x1Local, RoundMode::CAST_NONE, numCol); + Cast(x2_fp32, x2Local, RoundMode::CAST_NONE, numCol); + PipeBarrier(); + Add(x1_fp32, x1_fp32, x2_fp32, numCol); + PipeBarrier(); + Cast(xLocal, x1_fp32, RoundMode::CAST_RINT, numCol); + PipeBarrier(); + } else { + Add(x1Local, x1Local, x2Local, numCol); + PipeBarrier(); + Adds(xLocal, x1Local, (float)0, numCol); + } + inQueueX.FreeTensor(x1Local); + + // CopyOut x1 + x2 + outQueueY.EnQue(xLocal); + auto x_out = outQueueY.DeQue(); + DataCopyCustom(xGm[gm_bias], x_out, numCol); + outQueueY.FreeTensor(x_out); + } + + __aicore__ inline void CopyInGammaBeta() + { + LocalTensor gammaLocal = inQueueGamma.AllocTensor(); + DataCopyCustom(gammaLocal, gammaGm, numCol); + inQueueGamma.EnQue(gammaLocal); + if (!this->nullptrBeta) { + LocalTensor betaLocal = inQueueBeta.AllocTensor(); + DataCopyCustom(betaLocal, betaGm, numCol); + inQueueBeta.EnQue(betaLocal); + } + } + + __aicore__ inline void Compute(uint32_t inner_progress, LocalTensor gammaLocal, LocalTensor betaLocal, LocalTensor rstdLocal) + { + LocalTensor xLocal = inQueueX.AllocTensor(); + LocalTensor sqx = sqxBuf.Get(); + LocalTensor reduce_buf_local = reduceFp32Buf.Get(); + Mul(sqx, xLocal, xLocal, numCol); + PipeBarrier(); + + Muls(sqx, sqx, avgFactor, numCol); + PipeBarrier(); + + ReduceSumCustom(sqx, sqx, reduce_buf_local, numCol); + PipeBarrier(); + Adds(sqx, sqx, epsilon, 1); + PipeBarrier(); + + Sqrt(sqx, sqx, 1); + Duplicate(reduce_buf_local, ONE, 1); + PipeBarrier(); + Div(sqx, reduce_buf_local, sqx, 1); + PipeBarrier(); + event_t event_v_s = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(event_v_s); + WaitFlag(event_v_s); + float rstdValue = sqx.GetValue(0); + event_t event_s_v = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(event_s_v); + WaitFlag(event_s_v); + rstdLocal.SetValue(inner_progress, rstdValue); + PipeBarrier(); + LocalTensor yLocal = outQueueY.AllocTensor(); + Muls(yLocal, xLocal, rstdValue, numCol); + inQueueX.FreeTensor(xLocal); + PipeBarrier(); + Mul(yLocal, gammaLocal, yLocal, numCol); + if (!this->nullptrBeta) { + PipeBarrier(); + Add(yLocal, betaLocal, yLocal, numCol); + } + PipeBarrier(); + outQueueY.EnQue(yLocal); + } + + __aicore__ inline void Compute( + uint32_t inner_progress, LocalTensor gammaLocal, LocalTensor betaLocal, LocalTensor rstdLocal) + { + LocalTensor x_fp32 = xFp32Buf.Get(); + LocalTensor sqx = sqxBuf.Get(); + LocalTensor reduce_buf_local = reduceFp32Buf.Get(); + + Mul(sqx, x_fp32, x_fp32, numCol); + PipeBarrier(); + + Muls(sqx, sqx, avgFactor, numCol); + PipeBarrier(); + ReduceSumCustom(sqx, sqx, reduce_buf_local, numCol); + PipeBarrier(); + + Adds(sqx, sqx, epsilon, 1); + PipeBarrier(); + + Sqrt(sqx, sqx, 1); + Duplicate(reduce_buf_local, ONE, 1); + PipeBarrier(); + Div(sqx, reduce_buf_local, sqx, 1); + PipeBarrier(); + event_t event_v_s = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(event_v_s); + WaitFlag(event_v_s); + float rstdValue = sqx.GetValue(0); + event_t event_s_v = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(event_s_v); + WaitFlag(event_s_v); + rstdLocal.SetValue(inner_progress, rstdValue); + PipeBarrier(); + Muls(x_fp32, x_fp32, rstdValue, numCol); + PipeBarrier(); + LocalTensor yLocal = outQueueY.AllocTensor(); + Cast(yLocal, x_fp32, RoundMode::CAST_RINT, numCol); + PipeBarrier(); + Cast(x_fp32, yLocal, RoundMode::CAST_NONE, numCol); + PipeBarrier(); + Cast(sqx, gammaLocal, RoundMode::CAST_NONE, numCol); // gamma_fp32 reuse sqx + PipeBarrier(); + Mul(x_fp32, x_fp32, sqx, numCol); + if (!this->nullptrBeta) { + PipeBarrier(); + Cast(sqx, betaLocal, RoundMode::CAST_NONE, numCol); + PipeBarrier(); + Add(x_fp32, x_fp32, sqx, numCol); + } + PipeBarrier(); + Cast(yLocal, x_fp32, RoundMode::CAST_RINT, numCol); + PipeBarrier(); + + event_t event_v_mte = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(event_v_mte); + WaitFlag(event_v_mte); + + outQueueY.EnQue(yLocal); + } + + __aicore__ inline void Compute(uint32_t inner_progress, LocalTensor gammaLocal, LocalTensor betaLocal, LocalTensor rstdLocal) + { + LocalTensor x_fp32 = xFp32Buf.Get(); + LocalTensor sqx = sqxBuf.Get(); + LocalTensor reduce_buf_local = reduceFp32Buf.Get(); + + Mul(sqx, x_fp32, x_fp32, numCol); + PipeBarrier(); + + Muls(sqx, sqx, avgFactor, numCol); + PipeBarrier(); + + ReduceSumCustom(sqx, sqx, reduce_buf_local, numCol); + PipeBarrier(); + + Adds(sqx, sqx, epsilon, 1); + PipeBarrier(); + + Sqrt(sqx, sqx, 1); + Duplicate(reduce_buf_local, ONE, 1); + PipeBarrier(); + Div(sqx, reduce_buf_local, sqx, 1); + PipeBarrier(); + event_t event_v_s = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(event_v_s); + WaitFlag(event_v_s); + float rstdValue = sqx.GetValue(0); + event_t event_s_v = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(event_s_v); + WaitFlag(event_s_v); + rstdLocal.SetValue(inner_progress, rstdValue); + PipeBarrier(); + Muls(x_fp32, x_fp32, rstdValue, numCol); + PipeBarrier(); + LocalTensor yLocal = outQueueY.AllocTensor(); + Cast(yLocal, x_fp32, RoundMode::CAST_NONE, numCol); + + event_t event_v_mte = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(event_v_mte); + WaitFlag(event_v_mte); + + PipeBarrier(); + Mul(yLocal, gammaLocal, yLocal, numCol); + if (!this->nullptrBeta) { + PipeBarrier(); + Add(yLocal, betaLocal, yLocal, numCol); + } + PipeBarrier(); + outQueueY.EnQue(yLocal); + } + + __aicore__ inline void CopyOutY(uint32_t progress) + { + LocalTensor yLocal = outQueueY.DeQue(); + DataCopyCustom(yGm[progress], yLocal, numCol); + outQueueY.FreeTensor(yLocal); + } + + __aicore__ inline void CopyOutRstd(uint32_t outer_progress, uint32_t num) + { + LocalTensor rstdLocal = outQueueRstd.DeQue(); +#if __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + DataCopyCustom(rstdGm[outer_progress * rowFactor], rstdLocal, num); +#endif + outQueueRstd.FreeTensor(rstdLocal); + } + +private: + TPipe* Ppipe = nullptr; + // create queues for input, in this case depth is equal to buffer num + TQue inQueueX; + TQue inQueueGamma; + TQue inQueueBeta; + // create queues for output, in this case depth is equal to buffer num + TQue outQueueY; + TQue outQueueRstd; + + TBuf xFp32Buf; + TBuf sqxBuf; + TBuf reduceFp32Buf; + GlobalTensor x1Gm; + GlobalTensor x2Gm; + GlobalTensor gammaGm; + GlobalTensor betaGm; + GlobalTensor yGm; + GlobalTensor rstdGm; + GlobalTensor xGm; + + uint32_t numRow; + uint32_t numCol; + uint32_t blockFactor; // number of calculations rows on each core + uint32_t rowFactor; + uint32_t ubFactor; + float epsilon; + float avgFactor; + int32_t blockIdx_; + uint32_t rowWork = 1; + uint32_t nullptrBeta = 0; +}; +#endif // ADD_RMS_NORM_BIAS_H_ \ No newline at end of file diff --git a/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias_merge_n.h b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias_merge_n.h new file mode 100644 index 000000000..135467823 --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias_merge_n.h @@ -0,0 +1,471 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_bias_merge_n.h + * \brief add rms norm bias merge n file + */ +#ifndef ADD_RMS_NORM_BIAS_MERGE_N_H_ +#define ADD_RMS_NORM_BIAS_MERGE_N_H_ +#include "./rms_norm_base.h" + +using namespace AscendC; +using namespace RmsNorm; + +template +class KernelAddRmsNormBiasMergeN { +public: + __aicore__ inline KernelAddRmsNormBiasMergeN(TPipe* pipe) + { + Ppipe = pipe; + } + __aicore__ inline void Init( + GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta, GM_ADDR y, GM_ADDR rstd, GM_ADDR x, const AddRMSNormBiasTilingData* tiling) + { + ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!"); + this->numRow = tiling->num_row; + this->numCol = tiling->num_col; + this->numColAlign = tiling->num_col_align; + this->blockFactor = tiling->block_factor; + this->rowFactor = tiling->row_factor; + this->ubFactor = tiling->ub_factor; + this->epsilon = tiling->epsilon; + this->avgFactor = tiling->avg_factor; + + blockIdx_ = GetBlockIdx(); + if (blockIdx_ < GetBlockNum() - 1) { + this->rowWork = blockFactor; + this->rowLoop = tiling->row_loop; + this->rowTail = tiling->row_tail; + } else if (blockIdx_ == GetBlockNum() - 1) { + this->rowWork = tiling->last_block_factor; + this->rowLoop = tiling->last_block_row_loop; + this->rowTail = tiling->last_block_row_tail; + } + this->mulLoopFp32 = tiling->mul_loop_fp32; + this->mulTailFp32 = tiling->mul_tail_fp32; + this->dstRepStrideFp32 = tiling->dst_rep_stride_fp32; + this->mulLoopFp16 = tiling->mul_loop_fp16; + this->mulTailFp16 = tiling->mul_tail_fp16; + this->dstRepStrideFp16 = tiling->dst_rep_stride_fp16; + this->isPerformance = tiling->is_performance; + this->nullptrBeta = tiling->nullptr_beta; + // get start index for current core, core parallel + x1Gm.SetGlobalBuffer((__gm__ T*)x1 + blockIdx_ * blockFactor * numCol, rowWork * numCol); + x2Gm.SetGlobalBuffer((__gm__ T*)x2 + blockIdx_ * blockFactor * numCol, rowWork * numCol); + gammaGm.SetGlobalBuffer((__gm__ T*)gamma, numCol); + if (!this->nullptrBeta) { + betaGm.SetGlobalBuffer((__gm__ T*)beta, numCol); + } + yGm.SetGlobalBuffer((__gm__ T*)y + blockIdx_ * blockFactor * numCol, rowWork * numCol); + rstdGm.SetGlobalBuffer((__gm__ float*)rstd + blockIdx_ * blockFactor, blockFactor); + xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * blockFactor * numCol, rowWork * numCol); + + // pipe alloc memory to queue, the unit is Bytes + Ppipe->InitBuffer(inQueueX, DOUBLE_BUFFER_NUM, ubFactor * sizeof(T)); + Ppipe->InitBuffer(inQueueGamma, BUFFER_NUM, ubFactor * sizeof(T)); + if (!this->nullptrBeta) { + Ppipe->InitBuffer(inQueueBeta, BUFFER_NUM, ubFactor * sizeof(T)); + } + Ppipe->InitBuffer(outQueueY, DOUBLE_BUFFER_NUM, ubFactor * sizeof(T)); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + Ppipe->InitBuffer(outQueueRstd, BUFFER_NUM, rowFactor * sizeof(float)); +#else + Ppipe->InitBuffer(rstdBuf, rowFactor * sizeof(float)); +#endif + if constexpr (is_same::value || is_same::value) { + Ppipe->InitBuffer(xFp32Buf, ubFactor * sizeof(float)); + } + Ppipe->InitBuffer(sqxBuf, ubFactor * sizeof(float)); + Ppipe->InitBuffer(tmpBuf, rowFactor * NUM_PER_REP_FP32 * sizeof(float)); + } + + __aicore__ inline void Process() + { + CopyInGammaBeta(); + LocalTensor gammaLocal = inQueueGamma.DeQue(); + LocalTensor betaLocal; + if (!this->nullptrBeta) { + betaLocal = inQueueBeta.DeQue(); + } + for (uint32_t i_o = 0; i_o < rowLoop - 1; i_o++) { + MainCompute(i_o, rowFactor, gammaLocal, betaLocal); + } + MainCompute(rowLoop - 1, rowTail, gammaLocal, betaLocal); + inQueueGamma.FreeTensor(gammaLocal); + if (!this->nullptrBeta) { + inQueueBeta.FreeTensor(betaLocal); + } + } + + __aicore__ inline void MainCompute(uint32_t i_o, uint32_t calc_row_num, LocalTensor& gammaLocal, LocalTensor& betaLocal) + { + uint32_t gm_bias = i_o * rowFactor * numCol; + uint32_t elementNum = calc_row_num * numColAlign; + CopyInX(gm_bias, calc_row_num); + LocalTensor xLocal = ComputeX(elementNum); + CopyOutX(gm_bias, calc_row_num); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + LocalTensor rstdLocal = outQueueRstd.AllocTensor(); + ComputeRstd(xLocal, rstdLocal, calc_row_num, elementNum); + outQueueRstd.EnQue(rstdLocal); + CopyOutRstd(i_o, calc_row_num); +#else + LocalTensor rstdLocal = rstdBuf.Get(); + ComputeRstd(xLocal, rstdLocal, calc_row_num, elementNum); +#endif + ComputeY(xLocal, gammaLocal, betaLocal, rstdLocal, calc_row_num, elementNum); + CopyOutY(gm_bias, calc_row_num); + } + +private: + __aicore__ inline void CopyInX(uint32_t gm_bias, uint32_t calc_row_num) + { + LocalTensor x1Local = inQueueX.AllocTensor(); + if (isNumColAlign) { + DataCopyCustom(x1Local, x1Gm[gm_bias], calc_row_num * numCol); + } else { + DataCopyCustom(x1Local, x1Gm[gm_bias], calc_row_num, numCol); + } + inQueueX.EnQue(x1Local); + LocalTensor x2Local = inQueueX.AllocTensor(); + if (isNumColAlign) { + DataCopyCustom(x2Local, x2Gm[gm_bias], calc_row_num * numCol); + } else { + DataCopyCustom(x2Local, x2Gm[gm_bias], calc_row_num, numCol); + } + inQueueX.EnQue(x2Local); + } + + __aicore__ inline LocalTensor ComputeX(uint32_t elementNum) + { + LocalTensor x1Local = inQueueX.DeQue(); + LocalTensor x2Local = inQueueX.DeQue(); + LocalTensor xLocal = outQueueY.AllocTensor(); + if constexpr (!is_same::value) { + Add(xLocal, x1Local, x2Local, elementNum); + } else { + LocalTensor x1Fp32 = xFp32Buf.Get(); + LocalTensor x2Fp32 = sqxBuf.Get(); + Cast(x1Fp32, x1Local, RoundMode::CAST_NONE, elementNum); + Cast(x2Fp32, x2Local, RoundMode::CAST_NONE, elementNum); + PipeBarrier(); + Add(x1Fp32, x1Fp32, x2Fp32, elementNum); + PipeBarrier(); + Cast(xLocal, x1Fp32, RoundMode::CAST_RINT, elementNum); + } + inQueueX.FreeTensor(x1Local); + inQueueX.FreeTensor(x2Local); + outQueueY.EnQue(xLocal); + PipeBarrier(); + return xLocal; + } + + __aicore__ inline void CopyOutX(uint32_t gm_bias, uint32_t calc_row_num) + { + // CopyOut x1 + x2 + auto xOut = outQueueY.DeQue(); + if (isNumColAlign) { + DataCopyCustom(xGm[gm_bias], xOut, calc_row_num * numCol); + } else { + DataCopyCustom(xGm[gm_bias], xOut, calc_row_num, numCol); + } + outQueueY.FreeTensor(xOut); + } + + __aicore__ inline void CopyInGammaBeta() + { + LocalTensor gammaLocal = inQueueGamma.AllocTensor(); + DataCopyCustom(gammaLocal, gammaGm, numCol); + inQueueGamma.EnQue(gammaLocal); + if (!this->nullptrBeta) { + LocalTensor betaLocal = inQueueBeta.AllocTensor(); + DataCopyCustom(betaLocal, betaGm, numCol); + inQueueBeta.EnQue(betaLocal); + } + } + + __aicore__ inline void ComputeRstd(LocalTensor xLocal, LocalTensor rstdLocal, uint32_t calc_row_num, uint32_t elementNum) + { + LocalTensor sqx = sqxBuf.Get(); + LocalTensor tmpLocal = tmpBuf.Get(); + if constexpr (!is_same::value) { + LocalTensor x_fp32 = xFp32Buf.Get(); + Cast(x_fp32, xLocal, RoundMode::CAST_NONE, elementNum); + PipeBarrier(); + Mul(sqx, x_fp32, x_fp32, elementNum); + } else { + Mul(sqx, xLocal, xLocal, elementNum); + } + PipeBarrier(); + + Muls(sqx, sqx, avgFactor, elementNum); + PipeBarrier(); + + ReduceSumMultiN(rstdLocal, sqx, tmpLocal, calc_row_num, numCol, numColAlign); + PipeBarrier(); + Adds(rstdLocal, rstdLocal, epsilon, calc_row_num); + PipeBarrier(); + + Sqrt(rstdLocal, rstdLocal, calc_row_num); + Duplicate(tmpLocal, ONE, calc_row_num); + PipeBarrier(); + + Div(rstdLocal, tmpLocal, rstdLocal, calc_row_num); + PipeBarrier(); + } + + __aicore__ inline void ComputeY( + LocalTensor xLocal, LocalTensor gammaLocal, LocalTensor betaLocal, LocalTensor rstdLocal, uint32_t calc_row_num, uint32_t elementNum) + { + LocalTensor tmpLocal = tmpBuf.Get(); + uint32_t splidRow = 240; + uint32_t rowRepeatLoop1 = calc_row_num / splidRow; + uint32_t rowRepeatTail1 = calc_row_num - rowRepeatLoop1 * splidRow; + for(uint32_t r_i = 0; r_i < rowRepeatLoop1; r_i ++) { + Brcb(tmpLocal[r_i * splidRow * MOV_8], rstdLocal[r_i * splidRow], splidRow, {1, 8}); + } + PipeBarrier(); + + if(rowRepeatTail1 > 0) { + Brcb(tmpLocal[rowRepeatLoop1 * splidRow * MOV_8], rstdLocal[rowRepeatLoop1 * splidRow], rowRepeatTail1, {1, 8}); + PipeBarrier(); + } + LocalTensor yLocal = outQueueY.AllocTensor(); + if constexpr (!is_same::value) { + LocalTensor x_fp32 = xFp32Buf.Get(); + repeatByRow(x_fp32, x_fp32, tmpLocal, calc_row_num, ONE_UINT); + if constexpr (is_same::value) { + Cast(yLocal, x_fp32, RoundMode::CAST_NONE, elementNum); + } else { + Cast(yLocal, x_fp32, RoundMode::CAST_RINT, elementNum); + } + } else { + repeatByRow(yLocal, xLocal, tmpLocal, calc_row_num, ONE_UINT); + } + PipeBarrier(); + if constexpr (is_same::value) { + repeatByRow(yLocal, yLocal, gammaLocal, calc_row_num, TWO_UINT); + if (!this->nullptrBeta) { + addRepeatByRow(yLocal, yLocal, betaLocal, calc_row_num, TWO_UINT); + } + } else if constexpr (is_same::value) { + LocalTensor sqx = sqxBuf.Get(); + LocalTensor x_fp32 = xFp32Buf.Get(); + Cast(x_fp32, yLocal, RoundMode::CAST_NONE, elementNum); + Cast(sqx, gammaLocal, RoundMode::CAST_NONE, elementNum); + PipeBarrier(); + repeatByRow(x_fp32, x_fp32, sqx, calc_row_num, THREE_UINT); + if (!this->nullptrBeta) { + Cast(sqx, betaLocal, RoundMode::CAST_NONE, elementNum); + PipeBarrier(); + addRepeatByRow(x_fp32, x_fp32, sqx, calc_row_num, THREE_UINT); + } + Cast(yLocal, x_fp32, RoundMode::CAST_RINT, elementNum); + } else { + repeatByRow(yLocal, yLocal, gammaLocal, calc_row_num, THREE_UINT); + if (!this->nullptrBeta) { + addRepeatByRow(yLocal, yLocal, betaLocal, calc_row_num, THREE_UINT); + } + } + PipeBarrier(); + outQueueY.EnQue(yLocal); + } + + __aicore__ inline void CopyOutY(uint32_t progress, uint32_t calc_row_num) + { + LocalTensor yLocal = outQueueY.DeQue(); + if (isNumColAlign) { + DataCopyCustom(yGm[progress], yLocal, calc_row_num * numCol); + } else { + DataCopyCustom(yGm[progress], yLocal, calc_row_num, numCol); + } + outQueueY.FreeTensor(yLocal); + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + __aicore__ inline void CopyOutRstd(uint32_t outer_progress, uint32_t num) + { + LocalTensor rstdLocal = outQueueRstd.DeQue(); + DataCopyCustom(rstdGm[outer_progress * rowFactor], rstdLocal, num); + outQueueRstd.FreeTensor(rstdLocal); + } +#endif + + template + __aicore__ inline void repeatByRow(const LocalTensor& dstLocal, const LocalTensor& src1Local, const LocalTensor& src2Local, uint32_t calc_row_num, uint32_t type) + { + // TWO_UINT=gammaFp16 ONE_UINT=rstd + uint32_t strideParams[6] = {mulLoopFp32, mulTailFp32, 64, 1, dstRepStrideFp32, 0}; + if (type == TWO_UINT) { + strideParams[0] = mulLoopFp16; + strideParams[1] = mulTailFp16; + strideParams[2] = 128; + strideParams[4] = dstRepStrideFp16; + } else if (type == ONE_UINT) { + strideParams[3] = 0; + strideParams[5] = 1; + } + uint32_t singlT = 255; + uint32_t rowRepeatLoop = calc_row_num / singlT; + uint32_t rowRepeatTail = calc_row_num - rowRepeatLoop * singlT; + uint32_t offset2 = 0; + for(uint32_t r_i = 0; r_i < rowRepeatLoop; r_i ++) { + offset2 = type == 1 ? (r_i * singlT * MOV_8) : 0; + mulRepeat(dstLocal[r_i * singlT * numColAlign], src1Local[r_i * singlT * numColAlign], src2Local[offset2], singlT, strideParams); + } + if(rowRepeatTail > 0) { + offset2 = type == 1 ? (rowRepeatLoop * singlT * MOV_8) : 0; + uint32_t offset1 = rowRepeatLoop * singlT * numColAlign; + mulRepeat(dstLocal[offset1], src1Local[offset1], src2Local[offset2], rowRepeatTail, strideParams); + } + } + + template + __aicore__ inline void mulRepeat(const LocalTensor& dstLocal, const LocalTensor& src1Local, const LocalTensor& src2Local, uint32_t calcRowNum, uint32_t strideParams[6]) + { + uint32_t mulLoop = strideParams[0]; + uint32_t mulTail = strideParams[1]; + uint32_t strideNum = strideParams[2]; + uint8_t src1BlkStride = static_cast(strideParams[3]); + uint8_t dstRepStride = static_cast(strideParams[4]); + uint8_t src1RepStride = static_cast(strideParams[5]); + if(src1BlkStride == 0) { + for (uint32_t m_i = 0; m_i < mulLoop; m_i++) { + Mul(dstLocal[m_i * strideNum], src1Local[m_i * strideNum], src2Local, strideNum, calcRowNum, {1, 1, src1BlkStride, dstRepStride, dstRepStride, src1RepStride}); + } + PipeBarrier(); + if(mulTail > 0) { + Mul(dstLocal[mulLoop * strideNum], src1Local[mulLoop * strideNum], src2Local, mulTail, calcRowNum, {1, 1, src1BlkStride, dstRepStride, dstRepStride, src1RepStride}); + } + PipeBarrier(); + } else { + for (uint32_t m_i = 0; m_i < mulLoop; m_i++) { + Mul(dstLocal[m_i * strideNum], src1Local[m_i * strideNum], src2Local[m_i * strideNum], strideNum, calcRowNum, {1, 1, src1BlkStride, dstRepStride, dstRepStride, src1RepStride}); + } + PipeBarrier(); + if(mulTail > 0) { + Mul(dstLocal[mulLoop * strideNum], src1Local[mulLoop * strideNum], src2Local[mulLoop * strideNum], mulTail, calcRowNum, {1, 1, src1BlkStride, dstRepStride, dstRepStride, src1RepStride}); + } + PipeBarrier(); + } + } + + template + __aicore__ inline void addRepeatByRow(const LocalTensor& dstLocal, const LocalTensor& src1Local, const LocalTensor& src2Local, uint32_t calc_row_num, uint32_t type) + { + // TWO_UINT=gammaFp16 ONE_UINT=rstd + uint32_t strideParams[6] = {mulLoopFp32, mulTailFp32, 64, 1, dstRepStrideFp32, 0}; + if (type == TWO_UINT) { + strideParams[0] = mulLoopFp16; + strideParams[1] = mulTailFp16; + strideParams[2] = 128; + strideParams[4] = dstRepStrideFp16; + } else if (type == ONE_UINT) { + strideParams[3] = 0; + strideParams[5] = 1; + } + uint32_t singlT = 255; + uint32_t rowRepeatLoop = calc_row_num / singlT; + uint32_t rowRepeatTail = calc_row_num - rowRepeatLoop * singlT; + uint32_t offset2 = 0; + for(uint32_t r_i = 0; r_i < rowRepeatLoop; r_i ++) { + offset2 = type == 1 ? (r_i * singlT * MOV_8) : 0; + addRepeat(dstLocal[r_i * singlT * numColAlign], src1Local[r_i * singlT * numColAlign], src2Local[offset2], singlT, strideParams); + } + if(rowRepeatTail > 0) { + offset2 = type == 1 ? (rowRepeatLoop * singlT * MOV_8) : 0; + uint32_t offset1 = rowRepeatLoop * singlT * numColAlign; + addRepeat(dstLocal[offset1], src1Local[offset1], src2Local[offset2], rowRepeatTail, strideParams); + } + } + + template + __aicore__ inline void addRepeat(const LocalTensor& dstLocal, const LocalTensor& src1Local, const LocalTensor& src2Local, uint32_t calcRowNum, uint32_t strideParams[6]) + { + uint32_t addLoop = strideParams[0]; + uint32_t addTail = strideParams[1]; + uint32_t strideNum = strideParams[2]; + uint8_t src1BlkStride = static_cast(strideParams[3]); + uint8_t dstRepStride = static_cast(strideParams[4]); + uint8_t src1RepStride = static_cast(strideParams[5]); + if(src1BlkStride == 0) { + for (uint32_t m_i = 0; m_i < addLoop; m_i++) { + Add(dstLocal[m_i * strideNum], src1Local[m_i * strideNum], src2Local, strideNum, calcRowNum, {1, 1, src1BlkStride, dstRepStride, dstRepStride, src1RepStride}); + } + PipeBarrier(); + if(addTail > 0) { + Add(dstLocal[addLoop * strideNum], src1Local[addLoop * strideNum], src2Local, addTail, calcRowNum, {1, 1, src1BlkStride, dstRepStride, dstRepStride, src1RepStride}); + } + PipeBarrier(); + } else { + for (uint32_t m_i = 0; m_i < addLoop; m_i++) { + Add(dstLocal[m_i * strideNum], src1Local[m_i * strideNum], src2Local[m_i * strideNum], strideNum, calcRowNum, {1, 1, src1BlkStride, dstRepStride, dstRepStride, src1RepStride}); + } + PipeBarrier(); + if(addTail > 0) { + Add(dstLocal[addLoop * strideNum], src1Local[addLoop * strideNum], src2Local[addLoop * strideNum], addTail, calcRowNum, {1, 1, src1BlkStride, dstRepStride, dstRepStride, src1RepStride}); + } + PipeBarrier(); + } + } + +private: + TPipe* Ppipe = nullptr; + // create queues for input, in this case depth is equal to buffer num + TQue inQueueGamma; + TQue inQueueBeta; + TQue inQueueX; + // create queues for output, in this case depth is equal to buffer num +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + TQue outQueueRstd; +#else + TBuf rstdBuf; +#endif + TQue outQueueY; + + TBuf xFp32Buf; + TBuf sqxBuf; + TBuf tmpBuf; + GlobalTensor x1Gm; + GlobalTensor x2Gm; + GlobalTensor gammaGm; + GlobalTensor betaGm; + GlobalTensor yGm; + GlobalTensor rstdGm; + GlobalTensor xGm; + + uint32_t numRow; + uint32_t numCol; + uint32_t numColAlign; + uint32_t blockFactor; // number of calculations rows on each core + uint32_t rowFactor; + uint32_t ubFactor; + float epsilon; + float avgFactor; + int32_t blockIdx_; + uint32_t rowWork = 1; +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + bool isNumColAlign = true; +#else + bool isNumColAlign = false; +#endif + uint8_t isPerformance = 0; + uint32_t rowLoop = 1; + uint32_t rowTail = 0; + uint32_t mulLoopFp32; + uint32_t mulTailFp32; + uint8_t dstRepStrideFp32; + uint32_t mulLoopFp16; + uint32_t mulTailFp16; + uint8_t dstRepStrideFp16; + uint32_t nullptrBeta = 0; +}; +#endif // _ADD_RMS_NORM_BIAS_MERGE_N_H_ \ No newline at end of file diff --git a/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias_multi_n.h b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias_multi_n.h new file mode 100644 index 000000000..379e2e147 --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias_multi_n.h @@ -0,0 +1,339 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_bias_multi_n.h + * \brief add rms norm bias multi n file + */ +#ifndef ADD_RMS_NORM_BIAS_MULTI_N_H_ +#define ADD_RMS_NORM_BIAS_MULTI_N_H_ +#include "./rms_norm_base.h" + +using namespace AscendC; +using namespace RmsNorm; + +template +class KernelAddRmsNormBiasMultiN { +public: + __aicore__ inline KernelAddRmsNormBiasMultiN(TPipe* pipe) + { + Ppipe = pipe; + } + __aicore__ inline void Init( + GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta, GM_ADDR y, GM_ADDR rstd, GM_ADDR x, const AddRMSNormBiasTilingData* tiling) + { + ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!"); + this->numRow = tiling->num_row; + this->numCol = tiling->num_col; + this->numColAlign = tiling->num_col_align; + this->blockFactor = tiling->block_factor; + this->rowFactor = tiling->row_factor; + this->ubFactor = tiling->ub_factor; + this->epsilon = tiling->epsilon; + this->avgFactor = tiling->avg_factor; + this->nullptrBeta = tiling->nullptr_beta; + + blockIdx_ = GetBlockIdx(); + if (blockIdx_ < GetBlockNum() - 1) { + this->rowWork = blockFactor; + this->rowLoop = tiling->row_loop; + this->rowTail = tiling->row_tail; + } else if (blockIdx_ == GetBlockNum() - 1) { + this->rowWork = tiling->last_block_factor; + this->rowLoop = tiling->last_block_row_loop; + this->rowTail = tiling->last_block_row_tail; + } + // get start index for current core, core parallel + x1Gm.SetGlobalBuffer((__gm__ T*)x1 + blockIdx_ * blockFactor * numCol, rowWork * numCol); + x2Gm.SetGlobalBuffer((__gm__ T*)x2 + blockIdx_ * blockFactor * numCol, rowWork * numCol); + gammaGm.SetGlobalBuffer((__gm__ T*)gamma, numCol); + if (!this->nullptrBeta) { + betaGm.SetGlobalBuffer((__gm__ T*)beta, numCol); + } + yGm.SetGlobalBuffer((__gm__ T*)y + blockIdx_ * blockFactor * numCol, rowWork * numCol); + rstdGm.SetGlobalBuffer((__gm__ float*)rstd + blockIdx_ * blockFactor, blockFactor); + xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * blockFactor * numCol, rowWork * numCol); + + // pipe alloc memory to queue, the unit is Bytes + Ppipe->InitBuffer(inQueueX, DOUBLE_BUFFER_NUM, ubFactor * sizeof(T)); + Ppipe->InitBuffer(inQueueGamma, BUFFER_NUM, numColAlign * sizeof(T)); + if (!this->nullptrBeta) { + Ppipe->InitBuffer(inQueueBeta, BUFFER_NUM, numColAlign * sizeof(T)); + } + Ppipe->InitBuffer(outQueueY, DOUBLE_BUFFER_NUM, ubFactor * sizeof(T)); +#if __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + Ppipe->InitBuffer(outQueueRstd, BUFFER_NUM, rowFactor * NUM_PER_BLK_FP32 * sizeof(float)); +#else + Ppipe->InitBuffer(rstdBuf, rowFactor * NUM_PER_BLK_FP32 * sizeof(float)); +#endif + if constexpr (is_same::value || is_same::value) { + Ppipe->InitBuffer(xFp32Buf, ubFactor * sizeof(float)); + } + Ppipe->InitBuffer(sqxBuf, ubFactor * sizeof(float)); + Ppipe->InitBuffer(reduceFp32Buf, NUM_PER_REP_FP32 * sizeof(float)); + Ppipe->InitBuffer(offsetBuf, rowFactor * NUM_PER_BLK_FP32 * sizeof(uint32_t)); + } + __aicore__ inline void Process() + { + CopyInGammaBeta(); + LocalTensor betaLocal; + if (!this->nullptrBeta) { + betaLocal = inQueueBeta.DeQue(); + } + LocalTensor gammaLocal = inQueueGamma.DeQue(); + LocalTensor offsetLocal = offsetBuf.Get(); + for (uint32_t i = 0; i < rowFactor; i++) { + Duplicate(offsetLocal[i * NUM_PER_BLK_FP32], i * ONE_BLK_SIZE, NUM_PER_BLK_FP32); + } + for (uint32_t i_o = 0; i_o < rowLoop - 1; i_o++) { + SubProcessHalf(i_o, rowFactor, gammaLocal, betaLocal); + } + SubProcessHalf(rowLoop - 1, rowTail, gammaLocal, betaLocal); + inQueueGamma.FreeTensor(gammaLocal); + if (!this->nullptrBeta) { + inQueueBeta.FreeTensor(betaLocal); + } + } + + __aicore__ inline void SubProcessHalf(uint32_t i_o, uint32_t calc_row_num, LocalTensor& gammaLocal, LocalTensor& betaLocal) + { + uint32_t gm_bias = i_o * rowFactor * numCol; + CopyInX(gm_bias, calc_row_num); + LocalTensor xLocal = ComputeX(calc_row_num); + CopyOutX(gm_bias, calc_row_num); +#if __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + LocalTensor rstdLocal = outQueueRstd.AllocTensor(); + ComputeRstd(xLocal, rstdLocal, calc_row_num); + outQueueRstd.EnQue(rstdLocal); + CopyOutRstd(i_o * rowFactor, calc_row_num); +#else + LocalTensor rstdLocal = rstdBuf.Get(); + ComputeRstd(xLocal, rstdLocal, calc_row_num); +#endif + ComputeY(xLocal, gammaLocal, betaLocal, rstdLocal, calc_row_num); + CopyOutY(gm_bias, calc_row_num); + } + +private: + __aicore__ inline void CopyInX(uint32_t gm_bias, uint32_t calc_row_num) + { + LocalTensor x1Local = inQueueX.AllocTensor(); + DataCopyCustom(x1Local, x1Gm[gm_bias], calc_row_num * numCol); + inQueueX.EnQue(x1Local); + LocalTensor x2Local = inQueueX.AllocTensor(); + DataCopyCustom(x2Local, x2Gm[gm_bias], calc_row_num * numCol); + inQueueX.EnQue(x2Local); + } + + __aicore__ inline LocalTensor ComputeX(uint32_t calc_row_num) + { + uint32_t calc_num = calc_row_num * numColAlign; + LocalTensor x1Local = inQueueX.DeQue(); + LocalTensor x2Local = inQueueX.DeQue(); + LocalTensor xLocal = outQueueY.AllocTensor(); + if constexpr (!is_same::value) { + Add(xLocal, x1Local, x2Local, calc_num); + } else { + LocalTensor x1Fp32 = xFp32Buf.Get(); + LocalTensor x2Fp32 = sqxBuf.Get(); + Cast(x1Fp32, x1Local, RoundMode::CAST_NONE, calc_num); + Cast(x2Fp32, x2Local, RoundMode::CAST_NONE, calc_num); + PipeBarrier(); + Add(x1Fp32, x1Fp32, x2Fp32, calc_num); + PipeBarrier(); + Cast(xLocal, x1Fp32, RoundMode::CAST_RINT, calc_num); + } + inQueueX.FreeTensor(x1Local); + inQueueX.FreeTensor(x2Local); + outQueueY.EnQue(xLocal); + PipeBarrier(); + return xLocal; + } + + __aicore__ inline void CopyOutX(uint32_t gm_bias, uint32_t calc_row_num) + { + // CopyOut x1 + x2 + auto x_out = outQueueY.DeQue(); + DataCopyCustom(xGm[gm_bias], x_out, calc_row_num * numCol); + outQueueY.FreeTensor(x_out); + } + + __aicore__ inline void CopyInGammaBeta() + { + LocalTensor gammaLocal = inQueueGamma.AllocTensor(); + DataCopyCustom(gammaLocal, gammaGm, numCol); + inQueueGamma.EnQue(gammaLocal); + if (!this->nullptrBeta) { + LocalTensor betaLocal = inQueueBeta.AllocTensor(); + DataCopyCustom(betaLocal, betaGm, numCol); + inQueueBeta.EnQue(betaLocal); + } + } + + __aicore__ inline void ComputeRstd(LocalTensor xLocal, LocalTensor rstdLocal, uint32_t calc_row_num) + { + LocalTensor x_fp32 = xFp32Buf.Get(); + LocalTensor sqx = sqxBuf.Get(); + LocalTensor reduce_buf_local = reduceFp32Buf.Get(); + Cast(x_fp32, xLocal, RoundMode::CAST_NONE, calc_row_num * numColAlign); + PipeBarrier(); + + Mul(sqx, x_fp32, x_fp32, calc_row_num * numColAlign); + PipeBarrier(); + + Muls(sqx, sqx, avgFactor, calc_row_num * numColAlign); + PipeBarrier(); + + for (uint32_t i_i = 0; i_i < calc_row_num; i_i++) { + ReduceSumCustom(rstdLocal[i_i * NUM_PER_BLK_FP32], sqx[i_i * numColAlign], reduce_buf_local, numCol); + } + Adds(rstdLocal, rstdLocal, epsilon, calc_row_num * NUM_PER_BLK_FP32); + PipeBarrier(); + + Sqrt(rstdLocal, rstdLocal, calc_row_num * NUM_PER_BLK_FP32); + Duplicate(reduce_buf_local, ONE, NUM_PER_BLK_FP32); + PipeBarrier(); + + int32_t repeatTimes = calc_row_num * NUM_PER_BLK_FP32 / NUM_PER_REP_FP32; + int32_t tailCount = calc_row_num * NUM_PER_BLK_FP32 % NUM_PER_REP_FP32; + int32_t bodyCount = repeatTimes * NUM_PER_REP_FP32; + + if (likely(repeatTimes > 0)) { + Div(rstdLocal, reduce_buf_local, rstdLocal, NUM_PER_REP_FP32, repeatTimes, {1, 0, 1, DEFAULT_REPEAT_STRIDE, 0, DEFAULT_REPEAT_STRIDE}); + } + if (unlikely(tailCount != 0)) { + Div(rstdLocal[bodyCount], reduce_buf_local, rstdLocal[bodyCount], tailCount, 1, {1, 0, 1, DEFAULT_REPEAT_STRIDE, 0, DEFAULT_REPEAT_STRIDE}); + } + PipeBarrier(); + } + + __aicore__ inline void ComputeY( + LocalTensor xLocal, LocalTensor gammaLocal, LocalTensor betaLocal, LocalTensor rstdLocal, uint32_t calc_row_num) + { + LocalTensor x_fp32 = xFp32Buf.Get(); + LocalTensor offsetLocal = offsetBuf.Get(); + Gather(rstdLocal, rstdLocal, offsetLocal, ZERO_UINT, calc_row_num * NUM_PER_BLK_FP32); + PipeBarrier(); + int32_t repeatTimes = numCol / NUM_PER_REP_FP32; + int32_t tailCount = numCol % NUM_PER_REP_FP32; + int32_t bodyCount = repeatTimes * NUM_PER_REP_FP32; + for (uint32_t i_i = 0; i_i < calc_row_num; i_i++) { + if (likely(repeatTimes > 0)) { + Mul(x_fp32[i_i * numColAlign], x_fp32[i_i * numColAlign], rstdLocal[i_i * NUM_PER_BLK_FP32], + NUM_PER_REP_FP32, repeatTimes, {1, 1, 0, DEFAULT_REPEAT_STRIDE, DEFAULT_REPEAT_STRIDE, 0}); + } + if (unlikely(tailCount != 0)) { + Mul(x_fp32[i_i * numColAlign + bodyCount], x_fp32[i_i * numColAlign + bodyCount], + rstdLocal[i_i * NUM_PER_BLK_FP32], tailCount, 1, + {1, 1, 0, DEFAULT_REPEAT_STRIDE, DEFAULT_REPEAT_STRIDE, 0}); + } + } + PipeBarrier(); + LocalTensor yLocal = outQueueY.AllocTensor(); + if constexpr (is_same::value) { + Cast(yLocal, x_fp32, RoundMode::CAST_NONE, calc_row_num * numColAlign); + PipeBarrier(); + + for (uint32_t i_i = 0; i_i < calc_row_num; i_i++) { + Mul(yLocal[i_i * numColAlign], gammaLocal, yLocal[i_i * numColAlign], numCol); + } + if (!this->nullptrBeta) { + PipeBarrier(); + for (uint32_t i_i = 0; i_i < calc_row_num; i_i++) { + Add(yLocal[i_i * numColAlign], betaLocal, yLocal[i_i * numColAlign], numCol); + } + } + } else { + Cast(yLocal, x_fp32, RoundMode::CAST_RINT, calc_row_num * numColAlign); + PipeBarrier(); + LocalTensor yfp32 = xFp32Buf.Get(); + Cast(yfp32, yLocal, RoundMode::CAST_NONE, calc_row_num * numColAlign); + PipeBarrier(); + LocalTensor gammaFp32 = sqxBuf.Get(); + Cast(gammaFp32, gammaLocal, RoundMode::CAST_NONE, numCol); + PipeBarrier(); + for (uint32_t i_i = 0; i_i < calc_row_num; i_i++) { + Mul(yfp32[i_i * numColAlign], gammaFp32, yfp32[i_i * numColAlign], numCol); + } + PipeBarrier(); + if (!this->nullptrBeta) { + Cast(gammaFp32, betaLocal, RoundMode::CAST_NONE, numCol); + PipeBarrier(); + for (uint32_t i_i = 0; i_i < calc_row_num; i_i++) { + Add(yfp32[i_i * numColAlign], gammaFp32, yfp32[i_i * numColAlign], numCol); + } + PipeBarrier(); + } + Cast(yLocal, yfp32, RoundMode::CAST_RINT, calc_row_num * numColAlign); + } + PipeBarrier(); + outQueueY.EnQue(yLocal); + } + + __aicore__ inline void CopyOutY(uint32_t progress, uint32_t calc_row_num) + { + LocalTensor yLocal = outQueueY.DeQue(); + DataCopyCustom(yGm[progress], yLocal, calc_row_num * numCol); + outQueueY.FreeTensor(yLocal); + } + +#if __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + __aicore__ inline void CopyOutRstd(uint32_t outer_progress, uint32_t num) + { + LocalTensor rstdLocal = outQueueRstd.DeQue(); + DataCopyParams copyParams; + copyParams.blockLen = sizeof(float); + copyParams.blockCount = num; + DataCopyPad(rstdGm[outer_progress], rstdLocal, copyParams); + outQueueRstd.FreeTensor(rstdLocal); + } +#endif + +private: + TPipe* Ppipe = nullptr; + // create queues for input, in this case depth is equal to buffer num + TQue inQueueGamma; + TQue inQueueBeta; + TQue inQueueX; + // create queues for output, in this case depth is equal to buffer num +#if __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + TQue outQueueRstd; +#else + TBuf rstdBuf; +#endif + TQue outQueueY; + + TBuf xFp32Buf; + TBuf sqxBuf; + TBuf reduceFp32Buf; + TBuf offsetBuf; + GlobalTensor x1Gm; + GlobalTensor x2Gm; + GlobalTensor gammaGm; + GlobalTensor betaGm; + GlobalTensor yGm; + GlobalTensor rstdGm; + GlobalTensor xGm; + + uint32_t numRow; + uint32_t numCol; + uint32_t blockFactor; // number of calculations rows on each core + uint32_t rowFactor; + uint32_t ubFactor; + float epsilon; + float avgFactor; + uint32_t numColAlign; + int32_t blockIdx_; + uint32_t rowWork = 1; + uint32_t rowLoop = 1; + uint32_t rowTail = 0; + uint32_t nullptrBeta = 0; +}; +#endif // ADD_RMS_NORM__BIAS_MULTI_N_H_ \ No newline at end of file diff --git a/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias_single_n.h b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias_single_n.h new file mode 100644 index 000000000..067fff050 --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias_single_n.h @@ -0,0 +1,376 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_bias_single_n.h + * \brief add rms norm bias single n file + */ +#ifndef ADD_RMS_NORM_BIAS_SINGLE_N_H_ +#define ADD_RMS_NORM_BIAS_SINGLE_N_H_ +#include "./rms_norm_base.h" + +using namespace AscendC; +using namespace RmsNorm; + +template +class KernelAddRmsNormBiasSingleN { + static constexpr int32_t MAXBUFFER = 195584; +public: + __aicore__ inline KernelAddRmsNormBiasSingleN(TPipe* pipe) + { + Ppipe = pipe; + } + __aicore__ inline void Init( + GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta, GM_ADDR y, GM_ADDR rstd, GM_ADDR x, const AddRMSNormBiasTilingData* tiling) + { + ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!"); + + this->numCol = tiling->num_col; + this->blockFactor = 1; // in this case, blockFactor = 1 + this->ubFactor = tiling->ub_factor; + this->epsilon = tiling->epsilon; + this->avgFactor = (numCol != 0) ? (float)1.0 / numCol : 0; + this->nullptrBeta = tiling->nullptr_beta; + + this->rowWork = 1; + blockIdx_ = GetBlockIdx(); + // get start index for current core, core parallel + x1Gm.SetGlobalBuffer((__gm__ T*)x1 + blockIdx_ * numCol, numCol); + x2Gm.SetGlobalBuffer((__gm__ T*)x2 + blockIdx_ * numCol, numCol); + gammaGm.SetGlobalBuffer((__gm__ T*)gamma, numCol); + if (!this->nullptrBeta) { + betaGm.SetGlobalBuffer((__gm__ T*)beta, numCol); + } + yGm.SetGlobalBuffer((__gm__ T*)y + blockIdx_ * numCol, numCol); + rstdGm.SetGlobalBuffer((__gm__ float*)rstd + blockIdx_, 1); + xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * numCol, numCol); + + Ppipe->InitBuffer(unitBuf, MAXBUFFER); // (192 - 1) * 1024 byte + } + + __aicore__ inline void Process() + { + if constexpr (is_same::value) { + ProcessFp16(); + } else if constexpr (is_same::value) { + ProcessFp32(); + } else { + ProcessBf16(); + } + } + +private: + __aicore__ inline void ProcessFp16() + { + LocalTensor ubLocal = unitBuf.Get(); + LocalTensor xLocal = ubLocal.template ReinterpretCast(); + LocalTensor x1Local = xLocal[0]; + LocalTensor x2Local = xLocal[ubFactor]; + LocalTensor xFp32Local = ubLocal[ubFactor]; + LocalTensor sqxLocal = ubLocal[ubFactor * 2]; + LocalTensor tmpLocal = ubLocal[ubFactor * 3]; + + DataCopyCustom(x1Local, x1Gm, numCol); + event_t eventMTE2V1 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventMTE2V1); + DataCopyCustom(x2Local, x2Gm, numCol); + event_t eventMTE2V2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventMTE2V2); + WaitFlag(eventMTE2V1); + WaitFlag(eventMTE2V2); + Add(x1Local, x1Local, x2Local, numCol); + PipeBarrier(); + + // copy gamma + event_t eventVMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventVMTE2); + WaitFlag(eventVMTE2); + + DataCopyCustom(x2Local, gammaGm, numCol); // gammaLocal use x2Local + SetFlag(eventMTE2V2); + + // copy x out + event_t eventVMTE3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventVMTE3); + WaitFlag(eventVMTE3); + DataCopyCustom(xGm, x1Local, numCol); + event_t eventMTE3V = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V)); + SetFlag(eventMTE3V); + + Cast(xFp32Local, x1Local, RoundMode::CAST_NONE, numCol); + PipeBarrier(); + Mul(sqxLocal, xFp32Local, xFp32Local, numCol); + PipeBarrier(); + Muls(sqxLocal, sqxLocal, avgFactor, numCol); + PipeBarrier(); + ReduceSumCustom(sqxLocal, sqxLocal, tmpLocal, numCol); + PipeBarrier(); + Adds(sqxLocal, sqxLocal, epsilon, 1); + PipeBarrier(); + Sqrt(sqxLocal, sqxLocal, 1); + Duplicate(tmpLocal, ONE, 1); + PipeBarrier(); + Div(sqxLocal, tmpLocal, sqxLocal, 1); + PipeBarrier(); + + // copyout rstd +#if (defined(__CCE_AICORE__) && __CCE_AICORE__ == 220) || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + SetFlag(eventVMTE3); + WaitFlag(eventVMTE3); + DataCopyCustom(rstdGm, sqxLocal, 1); +#endif + event_t eventVS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventVS); + WaitFlag(eventVS); + float rstdValue = sqxLocal.GetValue(0); + event_t eventSV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventSV); + WaitFlag(eventSV); + + Muls(xFp32Local, xFp32Local, rstdValue, numCol); + PipeBarrier(); + WaitFlag(eventMTE3V); + Cast(x1Local, xFp32Local, RoundMode::CAST_NONE, numCol); + PipeBarrier(); + WaitFlag(eventMTE2V2); + Mul(x1Local, x1Local, x2Local, numCol); + + if (!this->nullptrBeta) { + event_t eventVMTE2Beta = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventVMTE2Beta); + WaitFlag(eventVMTE2Beta); + DataCopyCustom(x2Local, betaGm, numCol); + event_t eventMTE2XBeta = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventMTE2XBeta); + WaitFlag(eventMTE2XBeta); + Add(x1Local, x1Local, x2Local, numCol); + } + SetFlag(eventVMTE3); + WaitFlag(eventVMTE3); + DataCopyCustom(yGm, x1Local, numCol); + } + + __aicore__ inline void ProcessFp32() + { + LocalTensor ubLocal = unitBuf.Get(); + LocalTensor x1Local = ubLocal[0]; + LocalTensor x2Local = ubLocal[ubFactor]; + LocalTensor sqxLocal = ubLocal[ubFactor * 2]; + LocalTensor tmpLocal = ubLocal[ubFactor * 3]; + + DataCopyCustom(x1Local, x1Gm, numCol); + event_t eventMTE2V1 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventMTE2V1); + DataCopyCustom(x2Local, x2Gm, numCol); + event_t eventMTE2V2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventMTE2V2); + WaitFlag(eventMTE2V1); + WaitFlag(eventMTE2V2); + Add(x1Local, x1Local, x2Local, numCol); + PipeBarrier(); + + // copy gamma + event_t eventVMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventVMTE2); + WaitFlag(eventVMTE2); + + DataCopyCustom(x2Local, gammaGm, numCol); // gammaLocal use x2Local + SetFlag(eventMTE2V2); + + // copy x out + event_t eventVMTE3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventVMTE3); + WaitFlag(eventVMTE3); + DataCopyCustom(xGm, x1Local, numCol); + event_t eventMTE3V = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V)); + SetFlag(eventMTE3V); + + Mul(sqxLocal, x1Local, x1Local, numCol); + PipeBarrier(); + Muls(sqxLocal, sqxLocal, avgFactor, numCol); + PipeBarrier(); + ReduceSumCustom(sqxLocal, sqxLocal, tmpLocal, numCol); + PipeBarrier(); + Adds(sqxLocal, sqxLocal, epsilon, 1); + PipeBarrier(); + Sqrt(sqxLocal, sqxLocal, 1); + Duplicate(tmpLocal, ONE, 1); + PipeBarrier(); + Div(sqxLocal, tmpLocal, sqxLocal, 1); + PipeBarrier(); + + // copyout rstd +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + SetFlag(eventVMTE3); + WaitFlag(eventVMTE3); + DataCopyCustom(rstdGm, sqxLocal, 1); +#endif + event_t eventVS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventVS); + WaitFlag(eventVS); + float rstdValue = sqxLocal.GetValue(0); + event_t eventSV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventSV); + WaitFlag(eventSV); + WaitFlag(eventMTE3V); + Muls(x1Local, x1Local, rstdValue, numCol); + PipeBarrier(); + WaitFlag(eventMTE2V2); + Mul(x1Local, x1Local, x2Local, numCol); + if (!this->nullptrBeta) { + event_t eventVMTE2Beta = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventVMTE2Beta); + WaitFlag(eventVMTE2Beta); + DataCopyCustom(x2Local, betaGm, numCol); + event_t eventMTE2XBeta = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventMTE2XBeta); + WaitFlag(eventMTE2XBeta); + Add(x1Local, x1Local, x2Local, numCol); + } + SetFlag(eventVMTE3); + WaitFlag(eventVMTE3); + DataCopyCustom(yGm, x1Local, numCol); + } + + __aicore__ inline void ProcessBf16() + { + LocalTensor ubLocal = unitBuf.Get(); + LocalTensor xLocal = ubLocal.template ReinterpretCast(); + LocalTensor x1Local = xLocal[0]; + LocalTensor x2Local = xLocal[ubFactor]; + LocalTensor xFp32Local = ubLocal[ubFactor]; + LocalTensor sqxLocal = ubLocal[ubFactor * 2]; + LocalTensor tmpLocal = ubLocal[ubFactor * 3]; + + DataCopyCustom(x1Local, x1Gm, numCol); + event_t eventMTE2V1_BF16_0 = static_cast(GetTPipePtr()->AllocEventID()); + SetFlag(eventMTE2V1_BF16_0); + DataCopyCustom(x2Local, x2Gm, numCol); + event_t eventMTE2V2_BF16_0 = static_cast(GetTPipePtr()->AllocEventID()); + SetFlag(eventMTE2V2_BF16_0); + WaitFlag(eventMTE2V1_BF16_0); + GetTPipePtr()->ReleaseEventID(eventMTE2V1_BF16_0); + Cast(xFp32Local, x1Local, RoundMode::CAST_NONE, numCol); + WaitFlag(eventMTE2V2_BF16_0); + GetTPipePtr()->ReleaseEventID(eventMTE2V2_BF16_0); + Cast(sqxLocal, x2Local, RoundMode::CAST_NONE, numCol); + PipeBarrier(); + Add(xFp32Local, xFp32Local, sqxLocal, numCol); + PipeBarrier(); + Cast(x1Local, xFp32Local, RoundMode::CAST_RINT, numCol); + PipeBarrier(); + // copy gamma + event_t eventVMTE2_BF16_0 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventVMTE2_BF16_0); + WaitFlag(eventVMTE2_BF16_0); + + DataCopyCustom(x2Local, gammaGm, numCol); // gammaLocal use x2Local + event_t eventMTE2V2_BF16_1 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventMTE2V2_BF16_1); + + // copy x out + event_t eventVMTE3_BF16_0 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventVMTE3_BF16_0); + WaitFlag(eventVMTE3_BF16_0); + DataCopyCustom(xGm, x1Local, numCol); + event_t eventMTE3V_BF16_0 = static_cast(GetTPipePtr()->AllocEventID()); + SetFlag(eventMTE3V_BF16_0); + + Cast(xFp32Local, x1Local, RoundMode::CAST_NONE, numCol); + PipeBarrier(); + Mul(sqxLocal, xFp32Local, xFp32Local, numCol); + PipeBarrier(); + Muls(sqxLocal, sqxLocal, avgFactor, numCol); + PipeBarrier(); + ReduceSumCustom(sqxLocal, sqxLocal, tmpLocal, numCol); + PipeBarrier(); + Adds(sqxLocal, sqxLocal, epsilon, 1); + PipeBarrier(); + Sqrt(sqxLocal, sqxLocal, 1); + Duplicate(tmpLocal, ONE, 1); + PipeBarrier(); + Div(sqxLocal, tmpLocal, sqxLocal, 1); + PipeBarrier(); + event_t eventVS_BF16_0 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventVS_BF16_0); + WaitFlag(eventVS_BF16_0); + float rstdValue = sqxLocal.GetValue(0); + event_t eventSV_BF16_0 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventSV_BF16_0); + WaitFlag(eventSV_BF16_0); + // copyout rstd +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + event_t eventVMTE3_BF16_1 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventVMTE3_BF16_1); + WaitFlag(eventVMTE3_BF16_1); + DataCopyCustom(rstdGm, sqxLocal, 1); + event_t eventMTE3V2_BF16_0 = static_cast(GetTPipePtr()->AllocEventID()); + SetFlag(eventMTE3V2_BF16_0); +#endif + + Muls(xFp32Local, xFp32Local, rstdValue, numCol); + PipeBarrier(); + WaitFlag(eventMTE3V_BF16_0); + GetTPipePtr()->ReleaseEventID(eventMTE3V_BF16_0); + Cast(x1Local, xFp32Local, RoundMode::CAST_RINT, numCol); + PipeBarrier(); + Cast(xFp32Local, x1Local, RoundMode::CAST_NONE, numCol); + PipeBarrier(); + WaitFlag(eventMTE2V2_BF16_1); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + WaitFlag(eventMTE3V2_BF16_0); + GetTPipePtr()->ReleaseEventID(eventMTE3V2_BF16_0); +#endif + Cast(sqxLocal, x2Local, RoundMode::CAST_NONE, numCol); + PipeBarrier(); + Mul(xFp32Local, xFp32Local, sqxLocal, numCol); + if (!this->nullptrBeta) { + event_t eventVMTE2Beta = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventVMTE2Beta); + WaitFlag(eventVMTE2Beta); + DataCopyCustom(x2Local, betaGm, numCol); + event_t eventMTE2XBeta = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventMTE2XBeta); + WaitFlag(eventMTE2XBeta); + Cast(sqxLocal, x2Local, RoundMode::CAST_NONE, numCol); + PipeBarrier(); + Add(xFp32Local, xFp32Local, sqxLocal, numCol); + } + PipeBarrier(); + Cast(x1Local, xFp32Local, RoundMode::CAST_RINT, numCol); + event_t eventVMTE3_BF16_2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventVMTE3_BF16_2); + WaitFlag(eventVMTE3_BF16_2); + DataCopyCustom(yGm, x1Local, numCol); + } + +private: + TPipe* Ppipe = nullptr; + + TBuf unitBuf; + GlobalTensor x1Gm; + GlobalTensor x2Gm; + GlobalTensor gammaGm; + GlobalTensor betaGm; + GlobalTensor yGm; + GlobalTensor rstdGm; + GlobalTensor xGm; + + uint32_t numRow; + uint32_t numCol; + uint32_t blockFactor; // number of calculations rows on each core + uint32_t ubFactor; + float epsilon; + float avgFactor; + int32_t blockIdx_; + uint32_t rowWork = 1; + uint32_t nullptrBeta = 0; +}; +#endif // _ADD_RMS_NORM_BIAS_SINGLE_N_H_ \ No newline at end of file diff --git a/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias_split_d.h b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias_split_d.h new file mode 100644 index 000000000..c1a7000aa --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/add_rms_norm_bias_split_d.h @@ -0,0 +1,395 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_rms_norm_bias_split_d.h + * \brief add rms norm bias split d file + */ +#ifndef ADD_RMS_NORM_BIAS_SPLIT_D_H_ +#define ADD_RMS_NORM_BIAS_SPLIT_D_H_ +#include "./rms_norm_base.h" + +using namespace AscendC; +using namespace RmsNorm; + +template +class KernelAddRmsNormBiasSplitD { +public: + __aicore__ inline KernelAddRmsNormBiasSplitD(TPipe* pipe) + { + Ppipe = pipe; + } + __aicore__ inline void Init( + GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta, GM_ADDR y, GM_ADDR rstd, GM_ADDR x, const AddRMSNormBiasTilingData* tiling) + { + ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!"); + this->numRow = tiling->num_row; + this->numCol = tiling->num_col; + this->blockFactor = tiling->block_factor; + this->rowFactor = tiling->row_factor; + this->ubFactor = tiling->ub_factor; + this->epsilon = tiling->epsilon; + this->avgFactor = (numCol != 0) ? (float)1.0 / numCol : 0; + this->nullptrBeta = tiling->nullptr_beta; + + blockIdx_ = GetBlockIdx(); + if (blockIdx_ < GetBlockNum() - 1) { + this->rowWork = blockFactor; + } else if (blockIdx_ == GetBlockNum() - 1) { + this->rowWork = numRow - (GetBlockNum() - 1) * blockFactor; + } else { + } + // get start index for current core, core parallel + x1Gm.SetGlobalBuffer((__gm__ T*)x1 + blockIdx_ * blockFactor * numCol, rowWork * numCol); + x2Gm.SetGlobalBuffer((__gm__ T*)x2 + blockIdx_ * blockFactor * numCol, rowWork * numCol); + gammaGm.SetGlobalBuffer((__gm__ T*)gamma, numCol); + if (!this->nullptrBeta) { + betaGm.SetGlobalBuffer((__gm__ T*)beta, numCol); + } + yGm.SetGlobalBuffer((__gm__ T*)y + blockIdx_ * blockFactor * numCol, rowWork * numCol); + rstdGm.SetGlobalBuffer((__gm__ float*)rstd + blockIdx_ * blockFactor, blockFactor); + xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * blockFactor * numCol, rowWork * numCol); + + // pipe alloc memory to queue, the unit is Bytes. + // We need 2 buffers here for both x1 and x2. + Ppipe->InitBuffer(inQueueX, BUFFER_NUM, 2 * ubFactor * sizeof(T)); + Ppipe->InitBuffer(inQueueGamma, BUFFER_NUM, ubFactor * sizeof(T)); + if (!this->nullptrBeta) { + Ppipe->InitBuffer(inQueueBeta, BUFFER_NUM, ubFactor * sizeof(T)); + } + Ppipe->InitBuffer(outQueueY, BUFFER_NUM, ubFactor * sizeof(T)); + Ppipe->InitBuffer(outQueueRstd, BUFFER_NUM, rowFactor * sizeof(float)); + + if constexpr (is_same::value || is_same::value) { + Ppipe->InitBuffer(xFp32Buf, ubFactor * sizeof(float)); + } + Ppipe->InitBuffer(sqxBuf, ubFactor * sizeof(float)); + Ppipe->InitBuffer(sumBuf, rowFactor * NUM_PER_BLK_FP32 * sizeof(float)); + Ppipe->InitBuffer(reduceFp32Buf, NUM_PER_REP_FP32 * sizeof(float)); + } + + __aicore__ inline void Process() + { + uint32_t i_o_max = RmsNorm::CeilDiv(rowWork, rowFactor); + uint32_t row_tail = rowWork - (i_o_max - 1) * rowFactor; + uint32_t j_max = RmsNorm::CeilDiv(numCol, ubFactor); + uint32_t col_tail = numCol - (j_max - 1) * ubFactor; + for (uint32_t i_o = 0; i_o < i_o_max - 1; i_o++) { + SubProcess(i_o, rowFactor, j_max, col_tail); + } + SubProcess(i_o_max - 1, row_tail, j_max, col_tail); + } + + __aicore__ inline void SubProcess(uint32_t i_o, uint32_t calc_row_num, uint32_t j_max, uint32_t col_tail) + { + LocalTensor sumLocal = sumBuf.Get(); + + LocalTensor rstdLocal = outQueueRstd.AllocTensor(); + Duplicate(rstdLocal, (float)0.0, calc_row_num); + PipeBarrier(); + for (uint32_t j = 0; j < j_max - 1; j++) { + ComputeFormer(i_o, calc_row_num, j, rstdLocal, sumLocal, ubFactor); + } + // do tail + ComputeFormer(i_o, calc_row_num, j_max - 1, rstdLocal, sumLocal, col_tail); + ComputeRstd(rstdLocal, calc_row_num); + + for (uint32_t j = 0; j < j_max - 1; j++) { + ComputeLatter(i_o, calc_row_num, j, rstdLocal, ubFactor); + } + ComputeLatter(i_o, calc_row_num, j_max - 1, rstdLocal, col_tail); + outQueueRstd.EnQue(rstdLocal); + CopyOutRstd(i_o, calc_row_num); + } + +private: + __aicore__ inline void CopyInAndAdd(uint32_t i_idx, uint32_t j_idx, uint32_t num) + { + LocalTensor x1x2_in = inQueueX.AllocTensor(); + LocalTensor x1_in = x1x2_in[0]; + LocalTensor x2_in = x1x2_in[ubFactor]; + DataCopyCustom(x1_in, x1Gm[i_idx * numCol + j_idx * ubFactor], num); + DataCopyCustom(x2_in, x2Gm[i_idx * numCol + j_idx * ubFactor], num); + inQueueX.EnQue(x1x2_in); + LocalTensor x1x2Local = inQueueX.DeQue(); + + auto x1Local = x1x2Local[0]; + auto x2Local = x1x2Local[ubFactor]; + + LocalTensor xLocal = outQueueY.AllocTensor(); + + if constexpr (is_same::value) { + LocalTensor x1_fp32 = xFp32Buf.Get(); + + Add(xLocal, x1Local, x2Local, num); + PipeBarrier(); + Cast(x1_fp32, xLocal, RoundMode::CAST_NONE, num); + PipeBarrier(); + // x1+x2 saved in x1_fp32 + } else if constexpr (is_same::value) { + LocalTensor x1_fp32 = xFp32Buf.Get(); + LocalTensor x2_fp32 = x1x2Local.template ReinterpretCast(); + + Cast(x1_fp32, x1Local, RoundMode::CAST_NONE, num); + PipeBarrier(); + Cast(x2_fp32, x2Local, RoundMode::CAST_NONE, num); + PipeBarrier(); + + Add(x1_fp32, x1_fp32, x2_fp32, num); + PipeBarrier(); + Cast(xLocal, x1_fp32, RoundMode::CAST_RINT, num); + PipeBarrier(); + // x1+x2 saved in x1_fp32 + } else { + Add(x1Local, x1Local, x2Local, num); + PipeBarrier(); + Adds(xLocal, x1Local, (float)0.0, num); + // x1+x2 saved in inQueueX + } + inQueueX.FreeTensor(x1x2Local); + + // copy out to workspace && x_out + outQueueY.EnQue(xLocal); + auto x_out = outQueueY.DeQue(); + DataCopyCustom(xGm[i_idx * numCol + j_idx * ubFactor], x_out, num); + outQueueY.FreeTensor(x_out); + } + + __aicore__ inline void ComputeFormer( + uint32_t i_o_idx, uint32_t calc_row_num, uint32_t j_idx, LocalTensor& rstdLocal, + LocalTensor& sumLocal, uint32_t num) + { + for (uint32_t i_i = 0; i_i < calc_row_num; i_i++) { + CopyInAndAdd(i_o_idx * rowFactor + i_i, j_idx, num); + ComputeSum(i_i, sumLocal, num); + } + BlockReduceSumFP32(sumLocal, sumLocal, calc_row_num * NUM_PER_BLK_FP32); + Add(rstdLocal, rstdLocal, sumLocal, calc_row_num); + PipeBarrier(); + } + + __aicore__ inline void ComputeSum(uint32_t i_i_idx, LocalTensor& sumLocal, uint32_t num) + { + LocalTensor sqx = sqxBuf.Get(); + LocalTensor reduce_buf_local = reduceFp32Buf.Get(); + if constexpr (is_same::value || is_same::value) { + LocalTensor x_fp32 = xFp32Buf.Get(); + PipeBarrier(); + Mul(sqx, x_fp32, x_fp32, num); + } else { + LocalTensor xLocal = inQueueX.AllocTensor(); + PipeBarrier(); + Mul(sqx, xLocal, xLocal, num); + inQueueX.FreeTensor(xLocal); + } + PipeBarrier(); + Muls(sqx, sqx, avgFactor, num); + PipeBarrier(); + // 8 means 8 fp32 pre block + ReduceSumFP32ToBlock(sumLocal[i_i_idx * 8], sqx, reduce_buf_local, num); + } + + __aicore__ inline void ComputeRstd(LocalTensor rstdLocal, uint32_t num) + { + LocalTensor reduce_buf_local = reduceFp32Buf.Get(); + Adds(rstdLocal, rstdLocal, epsilon, num); + PipeBarrier(); + Sqrt(rstdLocal, rstdLocal, num); + Duplicate(reduce_buf_local, ONE, num); + PipeBarrier(); + Div(rstdLocal, reduce_buf_local, rstdLocal, num); + PipeBarrier(); + } + + __aicore__ inline void ComputeLatter( + uint32_t i_o_idx, uint32_t calc_row_num, uint32_t j_idx, LocalTensor& rstdLocal, uint32_t num) + { + CopyInGammaBeta(j_idx, num); + LocalTensor gammaLocal = inQueueGamma.DeQue(); + LocalTensor betaLocal; + if (!this->nullptrBeta) { + betaLocal = inQueueBeta.DeQue(); + } + for (uint32_t i_i = 0; i_i < calc_row_num; i_i++) { + CopyInX(i_o_idx * rowFactor + i_i, j_idx, num); + ComputeY(i_i, gammaLocal, betaLocal, rstdLocal, num); + CopyOutY(i_o_idx * rowFactor + i_i, j_idx, num); + } + inQueueGamma.FreeTensor(gammaLocal); + if (!this->nullptrBeta) { + inQueueBeta.FreeTensor(betaLocal); + } + } + + __aicore__ inline void CopyInGammaBeta(uint32_t j_idx, uint32_t num) + { + LocalTensor gammaLocal = inQueueGamma.AllocTensor(); + DataCopyCustom(gammaLocal, gammaGm[j_idx * ubFactor], num); + inQueueGamma.EnQue(gammaLocal); + if (!this->nullptrBeta) { + LocalTensor betaLocal = inQueueBeta.AllocTensor(); + DataCopyCustom(betaLocal, betaGm[j_idx * ubFactor], num); + inQueueBeta.EnQue(betaLocal); + } + } + + __aicore__ inline void CopyInX(uint32_t i_idx, uint32_t j_idx, uint32_t num) + { + LocalTensor xLocal = inQueueX.AllocTensor(); + DataCopyCustom(xLocal, xGm[i_idx * numCol + j_idx * ubFactor], num); + inQueueX.EnQue(xLocal); + if constexpr (is_same::value || is_same::value) { + LocalTensor x_fp32 = xFp32Buf.Get(); + LocalTensor xLocal = inQueueX.DeQue(); + Cast(x_fp32, xLocal, RoundMode::CAST_NONE, num); + PipeBarrier(); + inQueueX.FreeTensor(xLocal); + } + } + + __aicore__ inline void ComputeY( + uint32_t i_i_idx, LocalTensor& gammaLocal, LocalTensor& betaLocal, LocalTensor& rstdLocal, uint32_t num) + { + LocalTensor x_fp32 = xFp32Buf.Get(); + LocalTensor sqx = sqxBuf.Get(); + event_t event_v_s = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(event_v_s); + WaitFlag(event_v_s); + float rstdValue = rstdLocal.GetValue(i_i_idx); + event_t event_s_v = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(event_s_v); + WaitFlag(event_s_v); + PipeBarrier(); + Muls(x_fp32, x_fp32, rstdValue, num); + PipeBarrier(); + LocalTensor yLocal = outQueueY.AllocTensor(); + Cast(yLocal, x_fp32, RoundMode::CAST_NONE, num); + PipeBarrier(); + Mul(yLocal, gammaLocal, yLocal, num); + PipeBarrier(); + if (!this->nullptrBeta) { + Add(yLocal, betaLocal, yLocal, num); + PipeBarrier(); + } + outQueueY.EnQue(yLocal); + } + + __aicore__ inline void ComputeY( + uint32_t i_i_idx, LocalTensor& gammaLocal, LocalTensor& betaLocal, LocalTensor& rstdLocal, uint32_t num) + { + LocalTensor xLocal = inQueueX.DeQue(); + LocalTensor sqx = sqxBuf.Get(); + event_t event_v_s = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(event_v_s); + WaitFlag(event_v_s); + float rstdValue = rstdLocal.GetValue(i_i_idx); + event_t event_s_v = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(event_s_v); + WaitFlag(event_s_v); + LocalTensor yLocal = outQueueY.AllocTensor(); + Muls(yLocal, xLocal, rstdValue, num); + inQueueX.FreeTensor(xLocal); + PipeBarrier(); + Mul(yLocal, gammaLocal, yLocal, num); + PipeBarrier(); + if (!this->nullptrBeta) { + Add(yLocal, betaLocal, yLocal, num); + PipeBarrier(); + } + outQueueY.EnQue(yLocal); + } + + __aicore__ inline void ComputeY( + uint32_t i_i_idx, LocalTensor& gammaLocal, LocalTensor& betaLocal, LocalTensor& rstdLocal, uint32_t num) + { + LocalTensor x_fp32 = xFp32Buf.Get(); + LocalTensor sqx = sqxBuf.Get(); + event_t event_v_s = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(event_v_s); + WaitFlag(event_v_s); + float rstdValue = rstdLocal.GetValue(i_i_idx); + event_t event_s_v = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(event_s_v); + WaitFlag(event_s_v); + PipeBarrier(); + Muls(x_fp32, x_fp32, rstdValue, num); + PipeBarrier(); + LocalTensor yLocal = outQueueY.AllocTensor(); + Cast(yLocal, x_fp32, RoundMode::CAST_RINT, num); + PipeBarrier(); + Cast(x_fp32, yLocal, RoundMode::CAST_NONE, num); + PipeBarrier(); + Cast(sqx, gammaLocal, RoundMode::CAST_NONE, num); + PipeBarrier(); + Mul(x_fp32, x_fp32, sqx, num); + PipeBarrier(); + if (!this->nullptrBeta) { + Cast(sqx, betaLocal, RoundMode::CAST_NONE, num); + PipeBarrier(); + Add(x_fp32, x_fp32, sqx, num); + PipeBarrier(); + } + Cast(yLocal, x_fp32, RoundMode::CAST_RINT, num); + PipeBarrier(); + outQueueY.EnQue(yLocal); + } + + __aicore__ inline void CopyOutY(uint32_t i_idx, uint32_t j_idx, uint32_t num) + { + LocalTensor yLocal = outQueueY.DeQue(); + DataCopyCustom(yGm[i_idx * numCol + j_idx * ubFactor], yLocal, num); + outQueueY.FreeTensor(yLocal); + } + + __aicore__ inline void CopyOutRstd(uint32_t i_o_idx, uint32_t num) + { + LocalTensor rstdLocal = outQueueRstd.DeQue(); +#if __CCE_AICORE__ == 220 || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + DataCopyCustom(rstdGm[i_o_idx * rowFactor], rstdLocal, num); +#endif + outQueueRstd.FreeTensor(rstdLocal); + } + +private: + TPipe* Ppipe = nullptr; + // create queues for input, in this case depth is equal to buffer num + TQue inQueueX; + TQue inQueueGamma; + TQue inQueueBeta; + // create queues for output, in this case depth is equal to buffer num + TQue outQueueY; + TQue outQueueRstd; + TBuf xFp32Buf; + TBuf sqxBuf; + TBuf sumBuf; + TBuf reduceFp32Buf; + + GlobalTensor x1Gm; + GlobalTensor x2Gm; + GlobalTensor gammaGm; + GlobalTensor betaGm; + GlobalTensor yGm; + GlobalTensor rstdGm; + GlobalTensor xGm; + + uint32_t numRow; + uint32_t numCol; + uint32_t blockFactor; // number of calculations rows on each core + uint32_t rowFactor; + uint32_t ubFactor; + float epsilon; + float avgFactor; + int32_t blockIdx_; + uint32_t rowWork = 1; + uint32_t nullptrBeta = 0; + + int tempbufNum; +}; +#endif // _ADD_RMS_NORM_BIAS_SPLIT_D_H_ \ No newline at end of file diff --git a/csrc/ascend/moe/add_rms_norm_bias/op_kernel/reduce_common.h b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/reduce_common.h new file mode 100644 index 000000000..8b4b268ba --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/reduce_common.h @@ -0,0 +1,179 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +/*! + * \file reduce_common.h + */ +#ifndef REDUCE_COMMON_H_RMS_NORM +#define REDUCE_COMMON_H_RMS_NORM +#include "kernel_operator.h" +using namespace AscendC; + +constexpr uint32_t MAX_REP_NUM = 255; +constexpr uint32_t ELEM_PER_REP_FP32 = 64; +constexpr uint32_t ELEM_PER_BLK_FP32 = 8; +constexpr float ZERO = 0; +constexpr int32_t HALf_INTERVAL = 2; +constexpr int32_t INDEX_TWO = 2; +constexpr int32_t INDEX_FOUR = 4; +constexpr int32_t INDEX_EIGHT = 8; +constexpr int32_t INDEX_SIXTEEN = 16; + +__aicore__ inline void ReduceSumForSmallReduceDimPreRepeat( + const LocalTensor& dstLocal, const LocalTensor& srcLocal, const LocalTensor& tmpLocal, + const uint32_t elemNum, const uint32_t numLastDim, const uint32_t tailCount, const uint32_t repeat, + const uint8_t repStride) +{ + uint32_t elemIndex = 0; + for (; elemIndex + ELEM_PER_REP_FP32 <= numLastDim; elemIndex += ELEM_PER_REP_FP32) { + Add(tmpLocal, srcLocal[elemIndex], tmpLocal, elemNum, repeat, + {1, 1, 1, ELEM_PER_BLK_FP32, repStride, ELEM_PER_BLK_FP32}); + PipeBarrier(); + } + if (unlikely(tailCount != 0)) { + Add(tmpLocal, srcLocal[elemIndex], tmpLocal, tailCount, repeat, + {1, 1, 1, ELEM_PER_BLK_FP32, repStride, ELEM_PER_BLK_FP32}); + } + PipeBarrier(); + AscendCUtils::SetMask(ELEM_PER_REP_FP32); // set mask = 64 +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + if ASCEND_IS_AIV { + WholeReduceSum(dstLocal, tmpLocal, elemNum, repeat, 1, 1, ELEM_PER_BLK_FP32); + } +#elif defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003 + WholeReduceSum(dstLocal, tmpLocal, elemNum, repeat, 1, 1, ELEM_PER_BLK_FP32); +#else + WholeReduceSum(dstLocal, tmpLocal, elemNum, repeat, 1, 1, ELEM_PER_BLK_FP32); +#endif +} + +/* + * reduce dim form (N, D) to (N, 1) + * this reduce sum is for small reduce dim. + */ +__aicore__ inline void ReduceSumForSmallReduceDim( + const LocalTensor& dstLocal, const LocalTensor& srcLocal, const LocalTensor& tmpLocal, + const uint32_t numLastDimAligned, const uint32_t numLastDim, const uint32_t tailCount, const uint32_t repeat, + const uint8_t repStride) +{ + uint32_t repeatTimes = repeat / MAX_REP_NUM; + if (repeatTimes == 0) { + ReduceSumForSmallReduceDimPreRepeat( + dstLocal, srcLocal, tmpLocal, ELEM_PER_REP_FP32, numLastDim, tailCount, repeat, repStride); + } else { + uint32_t repTailNum = repeat % MAX_REP_NUM; + uint32_t repIndex = 0; + uint32_t repElem; + for (; repIndex + MAX_REP_NUM <= repeat; repIndex += MAX_REP_NUM) { + ReduceSumForSmallReduceDimPreRepeat( + dstLocal[repIndex], srcLocal[repIndex * numLastDimAligned], tmpLocal[repIndex * ELEM_PER_REP_FP32], + ELEM_PER_REP_FP32, numLastDim, tailCount, MAX_REP_NUM, repStride); + } + if (repTailNum != 0) { + ReduceSumForSmallReduceDimPreRepeat( + dstLocal[repIndex], srcLocal[repIndex * numLastDimAligned], tmpLocal[repIndex * ELEM_PER_REP_FP32], + ELEM_PER_REP_FP32, numLastDim, tailCount, repTailNum, repStride); + } + } +} + +/* + * reduce dim form (N, D) to (N, 1) + * this reduce sum is for small reduce dim, require D < 255 * 8. + * size of tmpLocal: (N, 64) + */ +__aicore__ inline void ReduceSumMultiN( + const LocalTensor& dstLocal, const LocalTensor& srcLocal, const LocalTensor& tmpLocal, + const uint32_t numRow, const uint32_t numCol, const uint32_t numColAlign) +{ + const uint32_t tailCount = numCol % ELEM_PER_REP_FP32; + const uint32_t repeat = numRow; + const uint8_t repStride = numColAlign / ELEM_PER_BLK_FP32; + Duplicate(tmpLocal, ZERO, numRow * ELEM_PER_REP_FP32); + PipeBarrier(); + ReduceSumForSmallReduceDim(dstLocal, srcLocal, tmpLocal, numColAlign, numCol, tailCount, repeat, repStride); +} + +__aicore__ inline int32_t findPowerTwo(int32_t n) +{ + // find max power of 2 no more than n (32 bit) + n |= n >> 1; // Set the first digit of n's binary to 1 + n |= n >> INDEX_TWO; + n |= n >> INDEX_FOUR; + n |= n >> INDEX_EIGHT; + n |= n >> INDEX_SIXTEEN; + return (n + 1) >> 1; +} + +__aicore__ inline void ReduceSumHalfInterval( + const LocalTensor& dst_local, const LocalTensor& src_local, int32_t count) +{ + if (likely(count > ELEM_PER_REP_FP32)) { + int32_t bodyCount = findPowerTwo(count); + int32_t tailCount = count - bodyCount; + if (tailCount > 0) { + Add(src_local, src_local, src_local[bodyCount], tailCount); + PipeBarrier(); + } + while (bodyCount > ELEM_PER_REP_FP32) { + bodyCount = bodyCount / HALf_INTERVAL; + Add(src_local, src_local, src_local[bodyCount], bodyCount); + PipeBarrier(); + } + + AscendCUtils::SetMask(ELEM_PER_REP_FP32); + } else { + AscendCUtils::SetMask(count); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + if (g_coreType == AIV) { + WholeReduceSum(dst_local, src_local, ELEM_PER_REP_FP32, 1, 0, 1, 0); + } +#elif defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003 + WholeReduceSum(dst_local, src_local, ELEM_PER_REP_FP32, 1, 1, 1, ELEM_PER_BLK_FP32); +#else + WholeReduceSum(dst_local, src_local, ELEM_PER_REP_FP32, 1, 1, 1, DEFAULT_REPEAT_STRIDE); +#endif + PipeBarrier(); +} + +__aicore__ inline float ReduceSumHalfInterval(const LocalTensor& src_local, int32_t count) +{ + if (likely(count > ELEM_PER_REP_FP32)) { + int32_t bodyCount = findPowerTwo(count); + int32_t tailCount = count - bodyCount; + if (tailCount > 0) { + Add(src_local, src_local, src_local[bodyCount], tailCount); + PipeBarrier(); + } + while (bodyCount > ELEM_PER_REP_FP32) { + bodyCount = bodyCount / HALf_INTERVAL; + Add(src_local, src_local, src_local[bodyCount], bodyCount); + PipeBarrier(); + } + + AscendCUtils::SetMask(ELEM_PER_REP_FP32); + } else { + AscendCUtils::SetMask(count); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + if (g_coreType == AIV) { + WholeReduceSum(src_local, src_local, ELEM_PER_REP_FP32, 1, 0, 1, 0); + } +#elif defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003 + WholeReduceSum(src_local, src_local, ELEM_PER_REP_FP32, 1, 1, 1, ELEM_PER_BLK_FP32); +#else + WholeReduceSum(src_local, src_local, ELEM_PER_REP_FP32, 1, 1, 1, DEFAULT_REPEAT_STRIDE); +#endif + event_t event_v_s = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(event_v_s); + WaitFlag(event_v_s); + return src_local.GetValue(0); +} +#endif // _REDUCE_COMMON_H_ diff --git a/csrc/ascend/moe/add_rms_norm_bias/op_kernel/rms_norm_base.h b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/rms_norm_base.h new file mode 100644 index 000000000..d2b75ebba --- /dev/null +++ b/csrc/ascend/moe/add_rms_norm_bias/op_kernel/rms_norm_base.h @@ -0,0 +1,316 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef RMS_NORM_BASE_H_ +#define RMS_NORM_BASE_H_ +#include "kernel_operator.h" +#include "reduce_common.h" + +namespace RmsNorm { +using namespace AscendC; + + +/** + * Get the block size of unified buffer in bytes + */ +__aicore__ inline constexpr uint32_t GetUbBlockSize() +{ + return 32U; +} + +/** + * Get the size of vector registers in bytes + */ +__aicore__ inline constexpr uint32_t GetVRegSize() +{ +#if __CCE_AICORE__ == 310 + return AscendC::VECTOR_REG_WIDTH; +#else + return 256U; +#endif +} + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ != 220 && __CCE_AICORE__ != 310 +#define bfloat16_t int16_t +#endif +constexpr int32_t BUFFER_NUM = 1; // tensor num for each queue +constexpr int32_t DOUBLE_BUFFER_NUM = 2; +constexpr int32_t UNROLL_NUM = 2; +constexpr int32_t NUM_PER_REP_FP32 = 64; // ONE_REPEAT_BYTE_SIZE / sizeof(float); +constexpr int32_t NUM_PER_BLK_FP32 = 8; +constexpr int32_t FLOAT_BTYPE_SIZE = 4; +constexpr int32_t NUM_PER_BLK_FP16 = 16; +constexpr int32_t CONTINUE_STRIDE = 8; +constexpr int32_t BLOCK_SIZE = 32; +constexpr uint32_t ONCE_VECTOR_SIZE = 256; +constexpr float MINUS_HALF = -0.5f; +constexpr uint32_t ZERO_UINT = 0; +constexpr uint32_t ONE_UINT = 1; +constexpr uint32_t TWO_UINT = 2; +constexpr uint32_t THREE_UINT = 3; +constexpr float ONE = 1; +constexpr int32_t SECOND_LOOP = 2; +constexpr int32_t HALf_INTERVAL = 2; +constexpr int32_t MAX_REAPEAT = 255; +constexpr int32_t DIM_NUM = 2; +constexpr int32_t NDDMA_DIM = 5; + +constexpr uint32_t V_LENGTH = GetVRegSize() / sizeof(float); +constexpr uint64_t ALIGN_512_FACTOR = 512; +constexpr uint64_t ALIGN_32_FACTOR = 32; +constexpr int32_t CONST_FACTOR_2 = 2; +constexpr uint32_t SUM_COUNT = 2; +constexpr int32_t MOV_2 = 2; +constexpr int32_t MOV_4 = 4; +constexpr int32_t MOV_8 = 8; +constexpr int32_t MOV_16 = 16; + +template +__aicore__ inline T CeilDiv(T x, T y) +{ + return y == 0 ? x : (x + y - 1) / y; +} + +template +__aicore__ inline T Min(T left, T right) +{ + return (left < right ? left : right); +} + +template +struct integral_constant { + static constexpr Tp value = v; +}; +using true_type = integral_constant; +using false_type = integral_constant; +template +struct is_same : public false_type {}; +template +struct is_same : public true_type {}; + +template +class KernelRmsNormBase { +#define IS_X_FP32 (is_same::value) +#define IS_GAMMA_FP32 (is_same::value) +#define IS_MIX_DTYPE ((!IS_X_FP32) && IS_GAMMA_FP32) +}; + +__aicore__ inline int32_t findPowerTwo(int32_t n) +{ + // find max power of 2 no more than n (32 bit) + n |= n >> 1; // Set the first digit of n's binary to 1 + n |= n >> MOV_2; + n |= n >> MOV_4; + n |= n >> MOV_8; + n |= n >> MOV_16; + return (n + 1) >> 1; +} + +__aicore__ inline void ReduceSumHalfIntervalToRepeat( + const LocalTensor& dst_local, const LocalTensor& src_local, int32_t count, int32_t left) +{ + // count need smaller than 255 repeat + if (likely(count > NUM_PER_BLK_FP32)) { + int32_t bodyCount = count - left; + int32_t tailCount = left; + if (tailCount > 0) { + Add(src_local, src_local, src_local[bodyCount], tailCount); + PipeBarrier(); + } + while (bodyCount > SECOND_LOOP * NUM_PER_BLK_FP32) { + bodyCount = bodyCount / HALf_INTERVAL; + Add(src_local, src_local, src_local[bodyCount], bodyCount); + PipeBarrier(); + } + bodyCount = bodyCount / HALf_INTERVAL; + Add(dst_local, src_local, src_local[bodyCount], bodyCount); + PipeBarrier(); + } +} + +__aicore__ inline void ReduceSumFP32( + const LocalTensor& dst_local, const LocalTensor& src_local, const LocalTensor& work_local, + int32_t count) +{ + // count need smaller than 255 repeat + uint64_t mask = NUM_PER_REP_FP32; + int32_t repeatTimes = count / NUM_PER_REP_FP32; + int32_t tailCount = count % NUM_PER_REP_FP32; + int32_t bodyCount = repeatTimes * NUM_PER_REP_FP32; + BinaryRepeatParams repeatParams; + repeatParams.src0RepStride = ONE_REPEAT_BYTE_SIZE / ONE_BLK_SIZE; + repeatParams.src0BlkStride = 1; + repeatParams.src1RepStride = 0; + repeatParams.src1BlkStride = 1; + repeatParams.dstRepStride = 0; + repeatParams.dstBlkStride = 1; + Duplicate(work_local, ZERO, NUM_PER_REP_FP32); + PipeBarrier(); + if (likely(repeatTimes > 0)) { + Add(work_local, src_local, work_local, mask, repeatTimes, repeatParams); + PipeBarrier(); + } + if (unlikely(tailCount != 0)) { + Add(work_local, src_local[bodyCount], work_local, tailCount, 1, repeatParams); + PipeBarrier(); + } + AscendCUtils::SetMask(NUM_PER_REP_FP32); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + if (g_coreType == AIV) { + WholeReduceSum(dst_local, work_local, MASK_PLACEHOLDER, 1, 0, 1, 0); + } +#elif !(defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + WholeReduceSum(dst_local, work_local, MASK_PLACEHOLDER, 1, 1, 1, DEFAULT_REPEAT_STRIDE); +#endif + PipeBarrier(); +} + +__aicore__ inline void ReduceSumCustom( + const LocalTensor& dst_local, const LocalTensor& src_local, const LocalTensor& work_local, + int32_t count) +{ + ReduceSumFP32(dst_local, src_local, work_local, count); +} +__aicore__ inline void ReduceSumFP32ToBlock( + const LocalTensor& dst_local, const LocalTensor& src_local, const LocalTensor& work_local, + int32_t count) +{ + // count need smaller than 255 repeat + uint64_t mask = NUM_PER_REP_FP32; + int32_t repeatTimes = count / NUM_PER_REP_FP32; + int32_t tailCount = count % NUM_PER_REP_FP32; + int32_t bodyCount = repeatTimes * NUM_PER_REP_FP32; + BinaryRepeatParams repeatParams; + repeatParams.src0RepStride = ONCE_VECTOR_SIZE / BLOCK_SIZE; + repeatParams.src0BlkStride = 1; + repeatParams.src1RepStride = 0; + repeatParams.src1BlkStride = 1; + repeatParams.dstRepStride = 0; + repeatParams.dstBlkStride = 1; + Duplicate(work_local, ZERO, NUM_PER_REP_FP32); + PipeBarrier(); + if (likely(repeatTimes > 0)) { + Add(work_local, src_local, work_local, mask, repeatTimes, repeatParams); + PipeBarrier(); + } + if (unlikely(tailCount != 0)) { + Add(work_local, src_local[bodyCount], work_local, tailCount, 1, repeatParams); + PipeBarrier(); + } + BlockReduceSum(dst_local, work_local, 1, mask, 1, 1, DEFAULT_REPEAT_STRIDE); + PipeBarrier(); +} + +__aicore__ inline void BlockReduceSumFP32( + const LocalTensor& dst_local, const LocalTensor& src_local, int32_t count) +{ + // count need multiple of 8 + int32_t repeatTimes = count / NUM_PER_REP_FP32; + int32_t tailCount = count % NUM_PER_REP_FP32; + int32_t dstAddr = repeatTimes * 8; + int32_t srcAddr = repeatTimes * NUM_PER_REP_FP32; + if (likely(repeatTimes > 0)) { + BlockReduceSum(dst_local, src_local, repeatTimes, NUM_PER_REP_FP32, 1, 1, DEFAULT_REPEAT_STRIDE); + PipeBarrier(); + } + if (tailCount != 0) { + BlockReduceSum(dst_local[dstAddr], src_local[srcAddr], 1, tailCount, 1, 1, DEFAULT_REPEAT_STRIDE); + PipeBarrier(); + } +} + +template +__aicore__ inline void DataCopyCustom(const U& dstTensor, const R& srcTensor, const uint32_t count) +{ +#if (defined(__CCE_AICORE__) && __CCE_AICORE__ == 220) || (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3003) + DataCopyParams copyParams; + copyParams.blockLen = count * sizeof(T); + copyParams.blockCount = 1; + if constexpr (is_same>::value) { + DataCopyPadParams padParams; + DataCopyPad(dstTensor, srcTensor, copyParams, padParams); + } else { + DataCopyPad(dstTensor, srcTensor, copyParams); + } +#else + // only support count greater than 32byte + int32_t numPerBlock = ONE_BLK_SIZE / sizeof(T); + if (count % numPerBlock == 0) { + DataCopy(dstTensor, srcTensor, count); + } else { + if constexpr (is_same>::value) { + int32_t num = AlignUp(count, numPerBlock); + DataCopy(dstTensor, srcTensor, num); + } else { + if (count < numPerBlock) { + DataCopy(dstTensor, srcTensor, numPerBlock); + } else { + int32_t num = count / numPerBlock * numPerBlock; + DataCopy(dstTensor, srcTensor, num); + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + for (int32_t i = 0; i < numPerBlock; i++) { + T tensorValue = srcTensor.GetValue(count - numPerBlock + i); + srcTensor.SetValue(i, tensorValue); + } + SetFlag(EVENT_ID0); + WaitFlag(EVENT_ID0); + DataCopy(dstTensor[count - numPerBlock], srcTensor, numPerBlock); + } + } + } +#endif +} + +template +__aicore__ inline void DataCopyCustom( + const LocalTensor& dstTensor, const GlobalTensor& srcTensor, const uint32_t numRow, const uint32_t numCol) +{ +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + DataCopyParams copyParams; + copyParams.blockLen = numCol * sizeof(T); + copyParams.blockCount = numRow; + DataCopyPadParams padParams; + DataCopyPad(dstTensor, srcTensor, copyParams, padParams); +#endif +} + +template +__aicore__ inline void DataCopyCustom( + const GlobalTensor& dstTensor, const LocalTensor& srcTensor, const uint32_t numRow, const uint32_t numCol) +{ +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + DataCopyParams copyParams; + copyParams.blockLen = numCol * sizeof(T); + copyParams.blockCount = numRow; + DataCopyPad(dstTensor, srcTensor, copyParams); +#endif +} + +__aicore__ inline void RoundFloat2Int8(LocalTensor& dstTensor, LocalTensor& srcTensor, int32_t size) +{ + Cast(srcTensor.ReinterpretCast(), srcTensor, RoundMode::CAST_RINT, size); + PipeBarrier(); + SetDeqScale((half)1.000000e+00f); + PipeBarrier(); + Cast(srcTensor.ReinterpretCast(), srcTensor.ReinterpretCast(), RoundMode::CAST_NONE, size); + PipeBarrier(); + Cast(dstTensor, srcTensor.ReinterpretCast(), RoundMode::CAST_TRUNC, size); +} + +__aicore__ inline uint32_t ROUND_UP(uint32_t x, uint32_t block_number) +{ + if (block_number > 0) { + return (x + block_number - 1) / block_number * block_number; + } + return 0; +} +} // namespace RmsNorm +#endif // RMS_NORM_BASE_H_ diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/CMakeLists.txt b/csrc/ascend/moe/apply_top_k_top_p_custom/CMakeLists.txt new file mode 100644 index 000000000..298ed5677 --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/apply_top_k_top_p_custom_torch_adpt.h b/csrc/ascend/moe/apply_top_k_top_p_custom/apply_top_k_top_p_custom_torch_adpt.h new file mode 100644 index 000000000..c7b9a0fba --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/apply_top_k_top_p_custom_torch_adpt.h @@ -0,0 +1,40 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef APPLY_TOP_K_TOP_P_CUSTOM_TORCH_ADPT_H +#define APPLY_TOP_K_TOP_P_CUSTOM_TORCH_ADPT_H + +namespace vllm_fl { +at::Tensor npu_apply_top_k_top_p( + const at::Tensor& logits, + const c10::optional& p, + const c10::optional& k) +{ + TORCH_CHECK(p.has_value() || k.has_value(), + "apply_top_k_top_p: p and k cannot be None at the same time."); + + at::Tensor out = at::empty_like(logits); + + EXEC_NPU_CMD( + aclnnApplyTopKTopPCustom, + logits, + p, + k, + out); + + return out; +} +} +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/CMakeLists.txt b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/CMakeLists.txt new file mode 100644 index 000000000..2c62314c8 --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/CMakeLists.txt @@ -0,0 +1,22 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnExc PRIVATE + apply_top_k_top_p_custom_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME ApplyTopKTopPCustom + OPTIONS + --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE apply_top_k_top_p_custom ACLNNTYPE aclnn_exclude) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom.h b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom.h new file mode 100644 index 000000000..87fd50b52 --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom.h @@ -0,0 +1,24 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file apply_top_k_top_p_custom.h + * \brief + */ +#ifndef OP_API_INC_LEVEL0_OP_APPLY_TOP_K_TOP_P_CUSTOM_OP_H_ +#define OP_API_INC_LEVEL0_OP_APPLY_TOP_K_TOP_P_CUSTOM_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor* ApplyTopKTopPCustom(const aclTensor* sortedValue, const aclTensor* sortedIndices, + const aclTensor* p, const aclTensor* k, aclOpExecutor* executor); +} +#endif // OP_API_INC_LEVEL0_OP_APPLY_TOP_K_TOP_P_CUSTOM_OP_H_ \ No newline at end of file diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom_def.cpp b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom_def.cpp new file mode 100644 index 000000000..319953b36 --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom_def.cpp @@ -0,0 +1,100 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file apply_top_k_top_p_custom_def.cpp + * \brief + */ +#include "register/op_def_registry.h" + +namespace ops { +class ApplyTopKTopPCustom : public OpDef { +public: + explicit ApplyTopKTopPCustom(const char *name) : OpDef(name) + { + this->Input("sorted_value") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("sorted_indices") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("p") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("k") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("out") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(false) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true); + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + + OpAICoreConfig config_kirin = GetKirinCoreConfig(); + this->AICore().AddConfig("kirinx90", config_kirin); + } + +private: + OpAICoreConfig GetKirinCoreConfig() const + { + OpAICoreConfig config_kirin; + config_kirin.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true); + config_kirin.Input("sorted_value") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Input("sorted_indices") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Input("p") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Input("k") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Output("out") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + return config_kirin; + } +}; + +OP_ADD(ApplyTopKTopPCustom); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom_tiling.cpp b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom_tiling.cpp new file mode 100644 index 000000000..f16df2009 --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom_tiling.cpp @@ -0,0 +1,314 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file apply_top_k_top_p_custom_tiling.cpp + * \brief + */ + +#include +#include +#include "tiling_base/error_log.h" +#include "tiling/platform/platform_ascendc.h" +#include "platform/platform_infos_def.h" +#include "register/op_def_registry.h" +#include "register/op_impl_registry.h" +#include "tiling/tiling_api.h" +#include "apply_top_k_top_p_custom_tiling.h" + +namespace { + constexpr uint32_t SYS_RESERVED_UB = uint32_t(16 * 1024); + constexpr uint32_t SELECT_RESERVED_UB = uint32_t(8 * 1024); + constexpr uint32_t DIM_ONE = 1; + constexpr uint32_t DIM_TWO = 2; + constexpr int32_t SORTED_VALUE_INPUT_INDEX = 0; + constexpr int32_t SORTED_INDICES_INPUT_INDEX = 1; + constexpr int32_t P_INPUT_INDEX = 2; + constexpr int32_t K_INPUT_INDEX = 3; + constexpr uint32_t DIM_INDEX0 = 0; + constexpr uint32_t FLOAT_BYTES = 4; + static std::map DTYPE_MAP = {{ge::DT_BF16, 2}, {ge::DT_FLOAT16, 1}, {ge::DT_FLOAT, 0}}; + static std::map DATATYPE_LEN_MAP = { + {ge::DT_FLOAT16, 2}, {ge::DT_BF16, 2}, {ge::DT_FLOAT, 4}}; + const static uint32_t SYS_WORKSPACESIZE = uint32_t(16 * 1024 * 1024); + + constexpr uint32_t DATA_PER_BLOCK_B32 = 8; + constexpr uint32_t BYTES_B32 = 4; + constexpr uint32_t BLOCK_BYTES = 32; + constexpr uint32_t K_VALUE_MAX = 1024; + constexpr uint32_t ONLY_TOP_P_KEY = 2; + constexpr uint32_t ONLY_TOP_K_KEY = 1; + constexpr uint32_t BATCH_MODE = 1; +} // namespace + +namespace optiling { +class ApplyTopKTopPCustomTiling { +public: + explicit ApplyTopKTopPCustomTiling(gert::TilingContext* context) : tilingcontext(context){}; + ge::graphStatus Init(); + ge::graphStatus RunKernelTiling(); +private: + ApplyTopKTopPCustomTilingData tilingData; + gert::TilingContext* tilingcontext = nullptr; + ge::graphStatus CheckShape(); + void SetTilingKey(); + void GetUsedCore(); + void CalDataPerCore(); + void FillTilingData(); + void PrintTilingData(); + template + inline auto CeilAlign(T1 a, T1 b) const -> T1 + { + return b == 0 ? a : (a + b - 1) / b * b; + } + template + inline auto FloorAlign(T1 a, T1 b) const -> T1 + { + return b == 0 ? a : a / b * b; + } + + const char *opName_ = nullptr; + uint32_t coreNum_ = 0; + uint32_t calUbSize_ = 0; + uint32_t batchSize_ = 0; + uint32_t vocabSize_ = 0; + uint32_t tilingKey_ = 0; + uint32_t usedCoreNum_ = 0; + uint32_t batchPerCore_ = 1; + uint32_t tailBatch_ = 0; + uint32_t dataNumInit_ = 0; + uint32_t dataNumInitAligned_ = 0; + uint32_t ubFactorElement_ = 0; + uint32_t ubFactorElementAligned_ = 0; + uint32_t tailUbFactorElement_ = 0; + uint32_t tailUbFactorElementAligned_ = 0; + uint32_t iterateTimes_ = 0; + uint32_t onlyTopK_ = 0; + uint32_t onlyTopP_ = 0; + uint64_t platformUbSize_ = 0; +}; + +ge::graphStatus ApplyTopKTopPCustomTiling::CheckShape() { + auto sortedValueShapePtr = tilingcontext->GetInputShape(SORTED_VALUE_INPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(tilingcontext, sortedValueShapePtr); + auto sortedValueShape = sortedValueShapePtr->GetStorageShape(); + if (sortedValueShape.GetDimNum() != DIM_TWO) { + OP_LOGE(opName_, "the dimNum of sorted_value should be 2, but got %zu.", sortedValueShape.GetDimNum()); + return ge::GRAPH_FAILED; + } + auto sortedIndicesShapePtr = tilingcontext->GetInputShape(SORTED_INDICES_INPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(tilingcontext, sortedIndicesShapePtr); + auto sortedIndicesShape = sortedIndicesShapePtr->GetStorageShape(); + if (sortedIndicesShape.GetDimNum() != DIM_TWO) { + OP_LOGE(opName_, "the dimNum of sorted_indices should be 2, but got %zu.", sortedIndicesShape.GetDimNum()); + return ge::GRAPH_FAILED; + } + batchSize_ = sortedValueShape.GetDim(DIM_INDEX0); + vocabSize_ = sortedValueShape.GetDim(DIM_ONE); + if (sortedIndicesShape.GetDim(DIM_INDEX0) != batchSize_ || sortedIndicesShape.GetDim(DIM_ONE) != vocabSize_) { + OP_LOGE(opName_, "the shape of sorted_indices should be equal to sorted_value."); + return ge::GRAPH_FAILED; + } + + auto pShapePtr = tilingcontext->GetOptionalInputShape(P_INPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(tilingcontext, pShapePtr); + auto pShape = pShapePtr->GetStorageShape(); + auto pDimNum = pShape.GetDimNum(); + if (pDimNum != DIM_ONE && pDimNum != 0) { + OP_LOGE(opName_, "the dimNum of p should be 1 or 0, but got %zu.", pDimNum); + return ge::GRAPH_FAILED; + } + if (pDimNum != 0 && batchSize_ != pShape.GetDim(DIM_INDEX0)) { + OP_LOGE(opName_, "p.shape[0] should be equal to logits.shape[0]."); + return ge::GRAPH_FAILED; + } + + auto kShapePtr = tilingcontext->GetOptionalInputShape(K_INPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(tilingcontext, kShapePtr); + auto kShape = kShapePtr->GetStorageShape(); + auto kDimNum = kShape.GetDimNum(); + if (kDimNum != DIM_ONE && kDimNum != 0) { + OP_LOGE(opName_, "the dimNum of k should be 1 or 0, but got %zu.", kDimNum); + return ge::GRAPH_FAILED; + } + if (kDimNum != 0 && batchSize_ != kShape.GetDim(DIM_INDEX0)) { + OP_LOGE(opName_, "k.shape[0] should be equal to logits.shape[0]."); + return ge::GRAPH_FAILED; + } + if (kDimNum == 0 && pDimNum == 0) { + OP_LOGE(opName_, "the dimNum of q and k should be 0 at the same time."); + return ge::GRAPH_FAILED; + } + onlyTopK_ = (kDimNum != 0 && pDimNum == 0) ? ONLY_TOP_K_KEY : 0; + onlyTopP_ = (pDimNum != 0 && kDimNum == 0) ? ONLY_TOP_P_KEY : 0; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus ApplyTopKTopPCustomTiling::Init() { + opName_ = tilingcontext->GetNodeName(); + OP_LOGD(opName_, "TilingForApplyTopKTopPCustom init."); + auto platformInfo = platform_ascendc::PlatformAscendC(tilingcontext->GetPlatformInfo()); + coreNum_ = platformInfo.GetCoreNumAiv(); + platformInfo.GetCoreMemSize(platform_ascendc::CoreMemType::UB, platformUbSize_); + OP_LOGD(opName_, "platformUbSize: %lu.", platformUbSize_); + uint32_t avaliableUb = static_cast(platformUbSize_) - SYS_RESERVED_UB - SELECT_RESERVED_UB; + calUbSize_ = FloorAlign(avaliableUb, BLOCK_BYTES); + if (CheckShape() == ge::GRAPH_FAILED) { + OP_LOGE(opName_, "check shape failed."); + return ge::GRAPH_FAILED; + } + uint32_t tempValue = 1; + while (tempValue < vocabSize_) { + tempValue <<= 1; + iterateTimes_++; + } // ceil(log2(vocabSize_)) + return ge::GRAPH_SUCCESS; +} + +void ApplyTopKTopPCustomTiling::SetTilingKey() { + tilingKey_ += onlyTopK_; + tilingKey_ += onlyTopP_; + tilingcontext->SetTilingKey(tilingKey_); + if (tilingKey_ == ONLY_TOP_P_KEY){ + tilingcontext->SetScheduleMode(BATCH_MODE); + } +} + +void ApplyTopKTopPCustomTiling::GetUsedCore() +{ + if (coreNum_ > 0) { + batchPerCore_ = coreNum_ == uint32_t(0) ? batchSize_ : batchSize_ / coreNum_; + tailBatch_ = batchSize_ % coreNum_; + usedCoreNum_ = coreNum_; + } +} + +void ApplyTopKTopPCustomTiling::CalDataPerCore() +{ + uint32_t inputDataTypeByte = DATATYPE_LEN_MAP[tilingcontext->GetInputDesc(SORTED_VALUE_INPUT_INDEX)->GetDataType()]; + uint32_t dataPerBlock = BLOCK_BYTES / inputDataTypeByte; + dataNumInit_ = vocabSize_ < K_VALUE_MAX ? vocabSize_ : K_VALUE_MAX; + dataNumInitAligned_ = vocabSize_ < K_VALUE_MAX ? vocabSize_ : K_VALUE_MAX; + ubFactorElement_ = vocabSize_ < K_VALUE_MAX ? vocabSize_ : K_VALUE_MAX; + ubFactorElementAligned_ = CeilAlign(ubFactorElement_, dataPerBlock); + tailUbFactorElement_ = vocabSize_ % ubFactorElement_; + tailUbFactorElement_ = tailUbFactorElement_ == uint32_t(0) ? ubFactorElement_ : tailUbFactorElement_; + tailUbFactorElementAligned_ = CeilAlign(tailUbFactorElement_, dataPerBlock); + + uint32_t sortedValueBytes = ubFactorElementAligned_ * inputDataTypeByte + K_VALUE_MAX * inputDataTypeByte; + uint32_t sortedIndicesBytes = ubFactorElementAligned_ * BYTES_B32 + K_VALUE_MAX * BYTES_B32; + uint32_t pBytes = dataPerBlock * inputDataTypeByte; + uint32_t kBytes = DATA_PER_BLOCK_B32 * BYTES_B32; + uint32_t outTensorBytes = ubFactorElementAligned_ * inputDataTypeByte; + + calUbSize_ = calUbSize_ - sortedValueBytes - sortedIndicesBytes - pBytes - kBytes - outTensorBytes; + if (onlyTopP_ > 0) { + calUbSize_ = static_cast(platformUbSize_); + } +} + +void ApplyTopKTopPCustomTiling::FillTilingData() +{ + tilingData.set_batchSize(batchSize_); + tilingData.set_vocabSize(vocabSize_); + tilingData.set_batchPerCore(batchPerCore_); + tilingData.set_tailBatch(tailBatch_); + tilingData.set_blockNum(usedCoreNum_); + tilingData.set_dataNumInit(dataNumInit_); + tilingData.set_dataNumInitAligned(dataNumInitAligned_); + tilingData.set_ubFactorElement(ubFactorElement_); + tilingData.set_ubFactorElementAligned(ubFactorElementAligned_); + tilingData.set_tailUbFactorElement(tailUbFactorElement_); + tilingData.set_tailUbFactorElementAligned(tailUbFactorElementAligned_); + tilingData.set_calUbSize(calUbSize_); + tilingData.set_iterateTimes(iterateTimes_); +} + +void ApplyTopKTopPCustomTiling::PrintTilingData() +{ + OP_LOGD(opName_, "batchSize: %u.", tilingData.get_batchSize()); + OP_LOGD(opName_, "vocabSize: %u.", tilingData.get_vocabSize()); + OP_LOGD(opName_, "batchPerCore: %u.", tilingData.get_batchPerCore()); + OP_LOGD(opName_, "tailBatch: %u.", tilingData.get_tailBatch()); + OP_LOGD(opName_, "usedCoreNum: %u.", tilingData.get_blockNum()); + OP_LOGD(opName_, "dataNumInit_: %u.", tilingData.get_dataNumInit()); + OP_LOGD(opName_, "dataNumInitAligned_: %u.", tilingData.get_dataNumInitAligned()); + OP_LOGD(opName_, "ubFactorElement: %u.", tilingData.get_ubFactorElement()); + OP_LOGD(opName_, "ubFactorElementAligned: %u.", tilingData.get_ubFactorElementAligned()); + OP_LOGD(opName_, "tailUbFactorElement: %u.", tilingData.get_tailUbFactorElement()); + OP_LOGD(opName_, "tailUbFactorElementAligned: %u.", tilingData.get_tailUbFactorElementAligned()); + OP_LOGD(opName_, "calUbSize: %u.", tilingData.get_calUbSize()); + OP_LOGD(opName_, "iterateTimes: %u.", tilingData.get_iterateTimes()); +} + +ge::graphStatus ApplyTopKTopPCustomTiling::RunKernelTiling() +{ + OP_LOGD(opName_, "TilingForApplyTopKTopPCustom start."); + + SetTilingKey(); + GetUsedCore(); + CalDataPerCore(); + FillTilingData(); + PrintTilingData(); + + OP_LOGD(opName_, "tilingKey: %u.", tilingKey_); + uint32_t syncWorkspaceSize = SYS_WORKSPACESIZE; + size_t* currentWorkspace = tilingcontext->GetWorkspaceSizes(1); + currentWorkspace[0] = onlyTopP_ > 0 ? syncWorkspaceSize + batchSize_ * vocabSize_ * FLOAT_BYTES : syncWorkspaceSize; + + tilingData.SaveToBuffer(tilingcontext->GetRawTilingData()->GetData(), + tilingcontext->GetRawTilingData()->GetCapacity()); + tilingcontext->GetRawTilingData()->SetDataSize(tilingData.GetDataSize()); + tilingcontext->SetBlockDim(usedCoreNum_); + + OP_LOGD(opName_, "TilingForApplyTopKTopPCustom end."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus TilingForApplyTopKTopPCustom(gert::TilingContext* context) +{ + ApplyTopKTopPCustomTiling tilingObject(context); + auto ret = tilingObject.Init(); + if (ret != ge::GRAPH_SUCCESS) { + OP_LOGE(context->GetNodeName(), "tiling Init failed."); + return ge::GRAPH_FAILED; + } + ret = tilingObject.RunKernelTiling(); + OP_LOGD(context->GetNodeName(), "TilingForApplyTopKTopPCustom end."); + return ret; +} + +static ge::graphStatus TilingPrepareForApplyTopKTopPCustom(gert::TilingParseContext* context) +{ + OP_LOGD(context->GetNodeName(), "TilingPrepareForApplyTopKTopPCustom start"); + auto compileInfo = context->GetCompiledInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo); + auto platformInfo = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + compileInfo->totalCoreNum = ascendcPlatform.GetCoreNumAiv(); + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + compileInfo->ubSizePlatForm = static_cast(ubSizePlatForm); + OP_CHECK_IF(compileInfo->ubSizePlatForm <= 0, + OP_LOGE(context->GetNodeName(), "Failed to get ub size"), + return ge::GRAPH_FAILED); + OP_LOGD(context->GetNodeName(), "ub_size_platform is %lu", compileInfo->ubSizePlatForm); + uint64_t totalUbSize = 0; + platformInfo->GetLocalMemSize(fe::LocalMemType::UB, totalUbSize); + OP_LOGD(context->GetNodeName(), "total ub size is %lu", totalUbSize); + OP_LOGD(context->GetNodeName(), "TilingPrepareForApplyTopKTopPCustom end"); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(ApplyTopKTopPCustom) + .Tiling(TilingForApplyTopKTopPCustom) + .TilingParse(TilingPrepareForApplyTopKTopPCustom); +} // namespace optiling diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom_tiling.h b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom_tiling.h new file mode 100644 index 000000000..baf9e9ba1 --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom_tiling.h @@ -0,0 +1,51 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file apply_top_k_top_p_custom_tiling.h + * \brief + * ATTENTION: MAKE SURE 'BEGIN_TILING_DATA_DEF' STAY IN THE SAME LINE (28) USING BLANK LINES. + * + * + * + * + * + */ +#ifndef __APPLY_TOP_K_TOP_P_CUSTOM_TILINGDATA_H__ +#define __APPLY_TOP_K_TOP_P_CUSTOM_TILINGDATA_H__ + +#include "register/tilingdata_base.h" + +namespace optiling { + +BEGIN_TILING_DATA_DEF(ApplyTopKTopPCustomTilingData) + TILING_DATA_FIELD_DEF(uint32_t, batchSize); + TILING_DATA_FIELD_DEF(uint32_t, vocabSize); + TILING_DATA_FIELD_DEF(uint32_t, batchPerCore); + TILING_DATA_FIELD_DEF(uint32_t, tailBatch); + TILING_DATA_FIELD_DEF(uint32_t, blockNum); + TILING_DATA_FIELD_DEF(uint32_t, dataNumInit); + TILING_DATA_FIELD_DEF(uint32_t, dataNumInitAligned); + TILING_DATA_FIELD_DEF(uint32_t, ubFactorElement); + TILING_DATA_FIELD_DEF(uint32_t, ubFactorElementAligned); + TILING_DATA_FIELD_DEF(uint32_t, tailUbFactorElement); + TILING_DATA_FIELD_DEF(uint32_t, tailUbFactorElementAligned); + TILING_DATA_FIELD_DEF(uint32_t, calUbSize); + TILING_DATA_FIELD_DEF(uint32_t, iterateTimes); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(ApplyTopKTopPCustom, ApplyTopKTopPCustomTilingData) + +struct TilingForApplyTopKTopPCustomCompileInfo { + uint32_t totalCoreNum = 0; + uint64_t ubSizePlatForm = 0; +}; + +} // namespace optiling +#endif // __APPLY_TOP_K_TOP_P_CUSTOM_TILINGDATA_H__ \ No newline at end of file diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/op_api/aclnn_apply_top_k_top_p_custom.cpp b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/op_api/aclnn_apply_top_k_top_p_custom.cpp new file mode 100644 index 000000000..4bab96ae3 --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/op_api/aclnn_apply_top_k_top_p_custom.cpp @@ -0,0 +1,212 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file aclnn_apply_top_k_top_p_custom.cpp + * \brief + */ +#include "aclnn_apply_top_k_top_p_custom.h" +#include "../apply_top_k_top_p_custom.h" +#include "../sort.h" +#include "aclnn_kernels/contiguous.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "aclnn/aclnn_base.h" +#include "opdev/common_types.h" +#include "opdev/data_type_utils.h" +#include "opdev/make_op_executor.h" +#include "opdev/format_utils.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/tensor_view_utils.h" +#include "opdev/shape_utils.h" +#include "opdev/platform.h" + +using namespace op; +#ifdef __cplusplus +extern "C" { +#endif +namespace { +static const int64_t EXPECTED_DIM_ONE = 1; +static const int64_t EXPECTED_DIM_TWO = 2; +static constexpr size_t DIM_ONE = 1; + +// According to the API definition, all supported dtypes must be enumerated. +static const std::initializer_list DTYPE_SUPPORT_LIST = { + op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BF16}; + +static const std::initializer_list INT_DTYPE_SUPPORT_LIST = { + op::DataType::DT_INT32}; + +static bool CheckNotNull(const aclTensor* logits, const aclTensor* p, const aclTensor *k, const aclTensor* out) +{ + OP_CHECK_NULL(logits, return false); + if (p == nullptr && k == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The inputs, p and k, should not be nullptr at the same time."); + } + OP_CHECK_NULL(out, return false); + return true; +} + +static bool CheckDtypeValid(const aclTensor* logits, const aclTensor* p, const aclTensor *k, const aclTensor* out) +{ + // Check whether the data type is within the supported list. + OP_CHECK_DTYPE_NOT_SUPPORT(logits, DTYPE_SUPPORT_LIST, return false); + if (p != nullptr) { + OP_CHECK_DTYPE_NOT_SUPPORT(p, DTYPE_SUPPORT_LIST, return false); + } + if (k != nullptr) { + OP_CHECK_DTYPE_NOT_SUPPORT(k, INT_DTYPE_SUPPORT_LIST, return false); + } + OP_CHECK_DTYPE_NOT_SUPPORT(out, DTYPE_SUPPORT_LIST, return false); + + // Check whether the data types are identical. + if (p != nullptr) { + OP_CHECK_DTYPE_NOT_MATCH(p, logits->GetDataType(), return false); + } + OP_CHECK_DTYPE_NOT_MATCH(out, logits->GetDataType(), return false); + return true; +} + +static bool CheckShapeValid(const aclTensor* logits, const aclTensor* p, const aclTensor *k, const aclTensor* out) +{ + OP_CHECK_WRONG_DIMENSION(logits, EXPECTED_DIM_TWO, return false); + OP_CHECK_SHAPE_NOT_EQUAL(out, logits, return false); + if (p != nullptr) { + OP_CHECK_WRONG_DIMENSION(p, EXPECTED_DIM_ONE, return false); + } + if (k != nullptr) { + OP_CHECK_WRONG_DIMENSION(k, EXPECTED_DIM_ONE, return false); + } + if (p != nullptr && p->GetViewShape().GetDim(0) != logits->GetViewShape().GetDim(0)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "expected p.size(0) is equal to logits.size(0), but got %ld.", + p->GetViewShape().GetDim(0)); + return false; + } + if (k != nullptr && k->GetViewShape().GetDim(0) != logits->GetViewShape().GetDim(0)) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "expected k.size(0) is equal to logits.size(0), but got %ld.", + k->GetViewShape().GetDim(0)); + return false; + } + return true; +} + +static bool CheckFormatValid(const aclTensor* logits, const aclTensor* p, const aclTensor *k, const aclTensor* out) +{ + if (logits->GetStorageFormat() != Format::FORMAT_ND) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "logits format only support ND"); + return false; + } + if (p != nullptr && p->GetStorageFormat() != Format::FORMAT_ND) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "p format only support ND"); + return false; + } + if (k != nullptr && k->GetStorageFormat() != Format::FORMAT_ND) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "k format only support ND"); + return false; + } + if (out->GetStorageFormat() != Format::FORMAT_ND) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "out format only support ND"); + return false; + } + return true; +} + +static aclnnStatus CheckParams(const aclTensor* logits, const aclTensor* p, const aclTensor *k, const aclTensor* out) +{ + // Refresh after the DFX scheme for error codes, etc. is refined; error logs are printed inside the check interfaces. + // 1. Check whether any parameters are null pointers. + CHECK_RET(CheckNotNull(logits, p, k, out), ACLNN_ERR_PARAM_NULLPTR); + + // 2. Check whether the input data types are within the range supported by the API; validate according to the API definition. + CHECK_RET(CheckDtypeValid(logits, p, k, out), ACLNN_ERR_PARAM_INVALID); + + // 3. Check whether shapes satisfy the constraints. + CHECK_RET(CheckShapeValid(logits, p, k, out), ACLNN_ERR_PARAM_INVALID); + + // 4. Check whether formats satisfy the constraints. + CHECK_RET(CheckFormatValid(logits, p, k, out), ACLNN_ERR_PARAM_INVALID); + + return ACLNN_SUCCESS; +} +} // namespace + +aclnnStatus aclnnApplyTopKTopPCustomGetWorkspaceSize( + const aclTensor* logits, const aclTensor* p, const aclTensor* k, aclTensor* out, uint64_t* workspaceSize, + aclOpExecutor** executor) +{ + OP_CHECK_COMM_INPUT(workspaceSize, executor); + L2_DFX_PHASE_1(aclnnApplyTopKTopPCustom, DFX_IN(logits, p, k), DFX_OUT(out)); + // Fixed boilerplate: create OpExecutor. + auto uniqueExecutor = CREATE_EXECUTOR(); + CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + + // Fixed boilerplate: parameter validation. + auto ret = CheckParams(logits, p, k, out); + CHECK_RET(ret == ACLNN_SUCCESS, ret); + bool pIsEmpty = false; + bool kIsEmpty = false; + if (p != nullptr) { + pIsEmpty = p->IsEmpty(); + } + if (k != nullptr) { + kIsEmpty = k->IsEmpty(); + } + if (logits->IsEmpty() || pIsEmpty || kIsEmpty) { + // Supplement according to actual support status. + *workspaceSize = 0; + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; + } + // Fixed boilerplate: convert the input selfRef to a contiguous tensor. + auto logitsContiguous = l0op::Contiguous(logits, uniqueExecutor.get()); + CHECK_RET(logitsContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR); + const aclTensor* pContiguous = nullptr; + const aclTensor* kContiguous = nullptr; + if (p != nullptr) { + pContiguous = l0op::Contiguous(p, uniqueExecutor.get()); + CHECK_RET(pContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR); + } + if (k != nullptr) { + kContiguous = l0op::Contiguous(k, uniqueExecutor.get()); + CHECK_RET(kContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR); + } + bool isLastDimSizeOne = logits->GetViewShape()[DIM_ONE] == 1; + auto viewCopyResult = logitsContiguous; + if (isLastDimSizeOne) { + viewCopyResult = l0op::ViewCopy(logitsContiguous, out, uniqueExecutor.get()); + } else { + auto sortResult = l0op::Sort(logitsContiguous, -1, false, true, op::DataType::DT_INT32, uniqueExecutor.get()); + const aclTensor* sortedValue = std::get<0>(sortResult); + CHECK_RET(sortedValue != nullptr, ACLNN_ERR_INNER_NULLPTR); + const aclTensor* sortedIndices = std::get<1>(sortResult); + CHECK_RET(sortedIndices != nullptr, ACLNN_ERR_INNER_NULLPTR); + auto res = l0op::ApplyTopKTopPCustom(sortedValue, sortedIndices, pContiguous, kContiguous, uniqueExecutor.get()); + CHECK_RET(res != nullptr, ACLNN_ERR_INNER_NULLPTR); + // Fixed boilerplate: copy the computed result to the output 'out'; 'out' may be a non-contiguous tensor. + viewCopyResult = l0op::ViewCopy(res, out, uniqueExecutor.get()); + } + CHECK_RET(viewCopyResult != nullptr, ACLNN_ERR_INNER_NULLPTR); + // Fixed boilerplate: obtain the workspace size required during computation. + *workspaceSize = uniqueExecutor->GetWorkspaceSize(); + uniqueExecutor.ReleaseTo(executor); // Transfer ownership of the executor held by uniqueExecutor to executor. + return ACLNN_SUCCESS; +} + +aclnnStatus aclnnApplyTopKTopPCustom(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream) +{ + L2_DFX_PHASE_2(aclnnApplyTopKTopPCustom); + // Fixed boilerplate: invoke framework capability to complete the computation. + return CommonOpExecutorRun(workspace, workspaceSize, executor, stream); +} + +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/op_api/aclnn_apply_top_k_top_p_custom.h b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/op_api/aclnn_apply_top_k_top_p_custom.h new file mode 100644 index 000000000..95c7a1b97 --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/op_api/aclnn_apply_top_k_top_p_custom.h @@ -0,0 +1,54 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file aclnn_apply_top_k_top_p_custom.h + * \brief + */ +#ifndef OP_API_INC_APPLY_TOP_K_TOP_P_CUSTOM_H_ +#define OP_API_INC_APPLY_TOP_K_TOP_P_CUSTOM_H_ + +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @brief aclnnApplyTopKTopPCustom的第一段接口，根据具体的计算流程，计算workspace大小。 + * @domain aclnn_ops_infer + * @param [in] logits: npu device侧的aclTensor，数据类型支持FLOAT、FLOAT16、BFLOAT16，支持非连续的Tensor，数据格式支持ND。 + * @param [in] p: npu device侧的aclTensor，数据类型支持FLOAT、FLOAT16、BFLOAT16，支持非连续的Tensor，数据格式支持ND。 + * @param [in] k: npu device侧的aclTensor，数据类型支持INT32，支持非连续的Tensor，数据格式支持ND。 + * @param [in] out: npu device侧的aclTensor，数据类型支持FLOAT、FLOAT16、BFLOAT16，支持非连续的Tensor，数据格式支持ND。 + * @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。 + * @param [out] executor: 返回op执行器，包含算子计算流程。 + * @return aclnnStatus: 返回状态码。 + */ +aclnnStatus aclnnApplyTopKTopPCustomGetWorkspaceSize(const aclTensor* logits, const aclTensor* p, + const aclTensor* k, aclTensor* out, uint64_t* workspaceSize, + aclOpExecutor** executor); + +/** + * @brief aclnnApplyTopKTopPCustom的第二段接口，用于执行计算。 + * @param [in] workspace: 在npu device侧申请的workspace内存起址。 + * @param [in] workspaceSize: 在npu device侧申请的workspace大小，由第一段接口aaclnnApplyTopKTopPCustomGetWorkspaceSize获取。 + * @param [in] stream: acl stream流。 + * @param [in] executor: op执行器，包含了算子计算流程。 + * @return aclnnStatus: 返回状态码。 + */ +aclnnStatus aclnnApplyTopKTopPCustom(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, + aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif // OP_API_INC_APPLY_TOP_K_TOP_P_CUSTOM_H_ diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/op_api/apply_top_k_top_p_custom.cpp b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/op_api/apply_top_k_top_p_custom.cpp new file mode 100644 index 000000000..a88dbb753 --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/op_api/apply_top_k_top_p_custom.cpp @@ -0,0 +1,46 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file apply_top_k_top_p_custom.cpp + * \brief + */ +#include "../apply_top_k_top_p_custom.h" +#include "opdev/data_type_utils.h" +#include "opdev/format_utils.h" +#include "opdev/make_op_executor.h" +#include "opdev/op_def.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/shape_utils.h" +using namespace op; + +namespace l0op { + +OP_TYPE_REGISTER(ApplyTopKTopPCustom); + +const aclTensor* ApplyTopKTopPCustom( + const aclTensor* sortedValue, const aclTensor* sortedIndices, const aclTensor* p, const aclTensor* k, + aclOpExecutor* executor) +{ + L0_DFX(ApplyTopKTopPCustom, sortedValue, sortedIndices, p, k); + auto output = executor->AllocTensor(sortedValue->GetViewShape(), sortedValue->GetDataType()); + if (p == nullptr) { + p = executor->AllocTensor(sortedValue->GetDataType(), Format::FORMAT_ND, Format::FORMAT_ND); + } + if (k == nullptr) { + k = executor->AllocTensor(DataType::DT_INT32, Format::FORMAT_ND, Format::FORMAT_ND); + } + ADD_TO_LAUNCHER_LIST_AICORE(ApplyTopKTopPCustom, OP_INPUT(sortedValue, sortedIndices, p, k), OP_OUTPUT(output)); + + return output; +} +} // namespace l0op diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/sort.h b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/sort.h new file mode 100644 index 000000000..6ca937fc3 --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/op_host/sort.h @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file sort.h + * \brief + */ +#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_SORT_OP_H_ +#define PTA_NPU_OP_API_INC_LEVEL0_OP_SORT_OP_H_ + +#include "opdev/op_executor.h" +#include "opdev/fast_vector.h" + +namespace l0op { +const std::tuple Sort(const aclTensor* self, int64_t dim, bool descending, bool stable, + op::DataType indicesType, aclOpExecutor* executor); +} + +#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_SORT_OP_H_ diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/op_kernel/apply_top_k_top_p_custom.cpp b/csrc/ascend/moe/apply_top_k_top_p_custom/op_kernel/apply_top_k_top_p_custom.cpp new file mode 100644 index 000000000..a98f6cecc --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/op_kernel/apply_top_k_top_p_custom.cpp @@ -0,0 +1,42 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file apply_top_k_top_p_custom.cpp + * \brief + */ + +#include "apply_top_k_top_p_custom.h" +#include "apply_top_p_custom.h" +using namespace AscendC; +using namespace ApplyTopKTopPCustomOp; +using namespace ApplyTopPCustomOp; + +extern "C" __global__ __aicore__ void apply_top_k_top_p_custom(GM_ADDR sorted_value, GM_ADDR sorted_indices, + GM_ADDR p, GM_ADDR k, GM_ADDR out, GM_ADDR workSpace, GM_ADDR tiling) { + TPipe pipe; + GET_TILING_DATA(tilingData, tiling); + if (TILING_KEY_IS(0)) { + ApplyTopKTopPCustomOp::ApplyTopKTopPCustom op; + op.InitTilingData(tilingData, sorted_value, sorted_indices, p, k, out); + op.InitBuffer(&pipe); + op.Process(); + } else if (TILING_KEY_IS(1)) { + ApplyTopKTopPCustomOp::ApplyTopKTopPCustom op; + op.InitTilingData(tilingData, sorted_value, sorted_indices, p, k, out); + op.InitBuffer(&pipe); + op.ProcessTopK(); + } else if (TILING_KEY_IS(2)) { + ApplyTopPCustomOp::ApplyTopPCustom op; + op.InitTilingData(tilingData, sorted_value, sorted_indices, p, k, out, workSpace); + op.InitBuffer(&pipe); + op.ProcessTopP(); + } +} \ No newline at end of file diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/op_kernel/apply_top_k_top_p_custom.h b/csrc/ascend/moe/apply_top_k_top_p_custom/op_kernel/apply_top_k_top_p_custom.h new file mode 100644 index 000000000..e911399e1 --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/op_kernel/apply_top_k_top_p_custom.h @@ -0,0 +1,719 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file apply_top_k_top_p_custom.h + * \brief + */ +#ifndef APPLY_TOP_K_TOP_P_CUSTOM_H_KERNEL +#define APPLY_TOP_K_TOP_P_CUSTOM_H_KERNEL + +#include "kernel_operator.h" + +using namespace AscendC; +namespace ApplyTopKTopPCustomOp { +constexpr uint32_t BUFFER_NUM = 1; +constexpr uint16_t FLOAT16_NEG_INF = 0xFC00; // -inf 64512 +constexpr uint16_t BF16_NEG_INF = 0xFF80; // -inf 65408 +constexpr int32_t FLOAT32_NEG_INF = 0xFF800000; // -inf -2139095040 + +constexpr uint32_t BLOCK_BYTES = 32; +constexpr uint32_t DATA_PER_BLOCK_B32 = 8; +constexpr uint32_t DATA_PER_REPEAT_B32 = 64; +constexpr uint32_t K_MAX = 1024; +constexpr uint64_t MASK_64 = 64; +constexpr CumSumConfig CUMSUM_CONFIG{true, true, false}; + +template +class ApplyTopKTopPCustom { +public: + __aicore__ inline ApplyTopKTopPCustom(){}; + __aicore__ inline void InitTilingData( + const ApplyTopKTopPCustomTilingData &__restrict tilingData, GM_ADDR sorted_value, GM_ADDR sorted_indices, + GM_ADDR p, GM_ADDR k, GM_ADDR out); + __aicore__ inline void InitBuffer(TPipe *inputPipe); + __aicore__ inline void Process(); + __aicore__ inline void ProcessTopK(); +private: + __aicore__ inline void InitCopyIn(uint32_t loopBatch, int64_t currentGmIdx); + __aicore__ inline void InitProcess(uint32_t loopBatch); + __aicore__ inline void ProcessKLtKMax(uint32_t loopBatch); + __aicore__ inline void ScatterCumtomImpl(uint32_t loopBatch, uint32_t loopProbNum, uint32_t offset); + __aicore__ inline void ProcessRemain(uint32_t loopBatch); + __aicore__ inline void GetKthResult(uint32_t loopBatch, uint32_t offset, uint8_t repeatTimes); + __aicore__ inline void GetFirstKLoop(uint32_t loopBatch, int32_t &firstKLoop); + __aicore__ inline void ScatterFromFirstKLoop(uint32_t loopBatch, int32_t firstKLoop, float &cumsumData); + __aicore__ inline void ReduceSumWithAddsAndExpImpl(uint32_t offset, uint32_t loopDataNum); + __aicore__ inline void CumSumWithAddsAndExpImpl( + uint32_t offset, uint32_t loopDataNum, uint32_t cumsumInner, float cumsumData); + // topk func + __aicore__ inline void InitProcessTopK(uint32_t loopBatch); + __aicore__ inline void ProcessKLtKMaxTopK(uint32_t loopBatch); + __aicore__ inline void ProcessRemainTopK(uint32_t loopBatch); + __aicore__ inline void GetFirstKLoopTopK(uint32_t loopBatch, int32_t &firstKLoop); + __aicore__ inline void ScatterFromFirstKLoopTopK(uint32_t loopBatch, int32_t firstKLoop); + __aicore__ inline void ScatterCumtomImplTopK(uint32_t loopBatch, uint32_t loopProbNum, uint32_t offset); + __aicore__ inline void SToMTE3Sync() { + event_t eventIDSToMTE3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3)); + SetFlag(eventIDSToMTE3); + WaitFlag(eventIDSToMTE3); + } + __aicore__ inline void MTE3ToSSync() { + event_t eventIDMTE3ToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S)); + SetFlag(eventIDMTE3ToS); + WaitFlag(eventIDMTE3ToS); + } + __aicore__ inline void VToSSync() { + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + } + __aicore__ inline void MTE2ToVSync() { + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + } + __aicore__ inline void MTE2ToSSync() { + event_t eventIdMte2ToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S)); + SetFlag(eventIdMte2ToS); + WaitFlag(eventIdMte2ToS); + } + __aicore__ inline void MTE3ToVSync() { + event_t eventIdMte3ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V)); + SetFlag(eventIdMte3ToV); + WaitFlag(eventIdMte3ToV); + } + __aicore__ inline void SToMTE2Sync() + { + event_t eventIDSToMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_MTE2)); + SetFlag(eventIDSToMTE2); + WaitFlag(eventIDSToMTE2); + } + __aicore__ inline void VToMTE3Sync() { + event_t eventIDVToMTE3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventIDVToMTE3); + WaitFlag(eventIDVToMTE3); + } +private: + TPipe *pipe_; + // create queues for input, in this case depth is equal to buffer num + TQue sortedValueInQueue_; + TQue sortedIndicesInQueue_; + TQue pInQueue_; + TQue kInQueue_; + TQue outQueue_; + TBuf calBuf_; + + // tilingData + uint32_t batchSize_ = 0; + uint32_t vocabSize_ = 0; + uint32_t batchPerCore_ = 0; + uint32_t tailBatch_ = 0; + uint32_t blockNum_ = 0; + uint32_t dataNumInit_ = 0; + uint32_t dataNumInitAligned_ = 0; + uint32_t ubFactorElement_ = 0; + uint32_t ubFactorElementAligned_ = 0; + uint32_t tailUbFactorElement_ = 0; + uint32_t tailUbFactorElementAligned_ = 0; + uint32_t calUbSize_ = 0; + + uint32_t blockIdx_ = 0; + uint32_t loopBatch_ = 0; + uint32_t batchOffset_ = 0; + uint32_t bufOffsetLoop = 0; + uint32_t loopInner_ = 0; + uint32_t loopInnerOnlyP_ = 0; + int64_t baseGmIdx_ = 0; + + GlobalTensor mGmSortedValue_; + GlobalTensor mGmSortedIndices_; + GlobalTensor mGmP_; + GlobalTensor mGmK_; + GlobalTensor mGmOut_; + + LocalTensor kLocal; + LocalTensor pLocal; + LocalTensor outTensor; + LocalTensor sortedValueLocal; + LocalTensor sortedIndicesLocal; + + LocalTensor sortedValueLocalFp32; + LocalTensor negInfLocal; + + LocalTensor calLocalFp32; + LocalTensor kthValueLocal; + LocalTensor tmpLocal; + LocalTensor cumSumRes; + LocalTensor cumSumTmp; + LocalTensor reduceLocal; + LocalTensor softMaxRes; + LocalTensor scatterTensor; + LocalTensor sharedTmpBuffer; + + float kthValue = 0; + float pValue = 0; + float maxValue = 0; + float reduceSumValueInvert = 0; + float reduceSumValue = 0; + inputT kthTopKValue = 0; + BinaryRepeatParams repeatParams = {1, 0, 1, 8, 0, 8}; + DataCopyExtParams scatterCopyParams{1, (uint32_t)(sizeof(outputT)), 0, 0, 0}; +}; + +template +__aicore__ inline void ApplyTopKTopPCustom::InitTilingData( + const ApplyTopKTopPCustomTilingData &__restrict tilingData, GM_ADDR sorted_value, GM_ADDR sorted_indices, + GM_ADDR p, GM_ADDR k, GM_ADDR out) { + batchSize_ = tilingData.batchSize; + vocabSize_ = tilingData.vocabSize; + batchPerCore_ = tilingData.batchPerCore; + tailBatch_ = tilingData.tailBatch; + blockNum_ = tilingData.blockNum; + dataNumInit_ = tilingData.dataNumInit; + dataNumInitAligned_ = AscendC::AlignUp(dataNumInit_, DATA_PER_BLOCK_B32); + ubFactorElement_ = tilingData.ubFactorElement; + ubFactorElementAligned_ = tilingData.ubFactorElementAligned; + tailUbFactorElement_ = tilingData.tailUbFactorElement; + tailUbFactorElementAligned_ = tilingData.tailUbFactorElementAligned; + calUbSize_ = tilingData.calUbSize; + blockIdx_ = GetBlockIdx(); + + if (blockIdx_ < tailBatch_) + { + loopBatch_ = batchPerCore_ + 1; + batchOffset_ = blockIdx_ * loopBatch_; + } + else + { + loopBatch_ = batchPerCore_; + batchOffset_ = blockIdx_ * batchPerCore_ + tailBatch_; + } + loopInner_ = (vocabSize_ - dataNumInit_ + ubFactorElementAligned_ - 1) / ubFactorElementAligned_; + loopInnerOnlyP_ = (vocabSize_ + ubFactorElementAligned_ - 1) / ubFactorElementAligned_; + mGmSortedValue_.SetGlobalBuffer(reinterpret_cast<__gm__ inputT *>(sorted_value)); + mGmSortedIndices_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(sorted_indices)); + mGmP_.SetGlobalBuffer(reinterpret_cast<__gm__ inputT *>(p)); + mGmK_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(k)); + mGmOut_.SetGlobalBuffer(reinterpret_cast<__gm__ outputT *>(out)); +} + +// init used buffer +template +__aicore__ inline void ApplyTopKTopPCustom::InitBuffer(TPipe *inputPipe) { + pipe_ = inputPipe; + pipe_->InitBuffer(sortedValueInQueue_, BUFFER_NUM, sizeof(inputT) * (ubFactorElementAligned_ + K_MAX)); + pipe_->InitBuffer(sortedIndicesInQueue_, BUFFER_NUM, sizeof(int32_t) * (ubFactorElementAligned_ + K_MAX)); + pipe_->InitBuffer(pInQueue_, BUFFER_NUM, BLOCK_BYTES); + pipe_->InitBuffer(kInQueue_, BUFFER_NUM, BLOCK_BYTES); + pipe_->InitBuffer(outQueue_, BUFFER_NUM, sizeof(outputT) * ubFactorElementAligned_); + pipe_->InitBuffer(calBuf_, calUbSize_); + if constexpr (!IsSameType::value) { + sortedValueLocalFp32 = calBuf_.GetWithOffset(ubFactorElementAligned_ + K_MAX, bufOffsetLoop); + bufOffsetLoop = bufOffsetLoop + (ubFactorElementAligned_ + K_MAX) * sizeof(float); + } + kthValueLocal = calBuf_.GetWithOffset(DATA_PER_BLOCK_B32, bufOffsetLoop); + bufOffsetLoop = bufOffsetLoop + BLOCK_BYTES; + + negInfLocal = calBuf_.GetWithOffset(DATA_PER_BLOCK_B32, bufOffsetLoop); + bufOffsetLoop = bufOffsetLoop + BLOCK_BYTES; + + tmpLocal = calBuf_.GetWithOffset(ubFactorElementAligned_, bufOffsetLoop); + bufOffsetLoop = bufOffsetLoop + ubFactorElementAligned_ * sizeof(float); + cumSumRes = calBuf_.GetWithOffset(ubFactorElementAligned_, bufOffsetLoop); + bufOffsetLoop = bufOffsetLoop + ubFactorElementAligned_ * sizeof(float); + cumSumTmp = calBuf_.GetWithOffset(ubFactorElementAligned_, bufOffsetLoop); + bufOffsetLoop = bufOffsetLoop + ubFactorElementAligned_ * sizeof(float); + reduceLocal = calBuf_.GetWithOffset(ubFactorElementAligned_ * BLOCK_BYTES, bufOffsetLoop); + bufOffsetLoop = bufOffsetLoop + ubFactorElementAligned_ * BLOCK_BYTES * sizeof(float); + + softMaxRes = tmpLocal.template ReinterpretCast(); + scatterTensor = reduceLocal.template ReinterpretCast(); + sharedTmpBuffer = reduceLocal.template ReinterpretCast(); +} + +template +__aicore__ inline void ApplyTopKTopPCustom::Process() { + kLocal = kInQueue_.AllocTensor(); + pLocal = pInQueue_.AllocTensor(); + outTensor = outQueue_.AllocTensor(); + sortedValueLocal = sortedValueInQueue_.AllocTensor(); + sortedIndicesLocal = sortedIndicesInQueue_.AllocTensor(); + Duplicate(negInfLocal.template ReinterpretCast(), FLOAT32_NEG_INF, DATA_PER_BLOCK_B32); + if constexpr (IsSameType::value) { + calLocalFp32 = sortedValueLocal; + Duplicate(outTensor.template ReinterpretCast(), FLOAT32_NEG_INF, ubFactorElementAligned_); + } else if constexpr (IsSameType::value) { + calLocalFp32 = sortedValueLocalFp32; + Duplicate(outTensor.template ReinterpretCast(), FLOAT16_NEG_INF, ubFactorElementAligned_); + } else { + calLocalFp32 = sortedValueLocalFp32; + Duplicate(outTensor.template ReinterpretCast(), BF16_NEG_INF, ubFactorElementAligned_); + } + VToMTE3Sync(); + for (uint32_t loopBatch = 0; loopBatch < loopBatch_; loopBatch++) { + baseGmIdx_ = batchOffset_ * vocabSize_ + loopBatch * vocabSize_; + InitProcess(loopBatch); + if (calLocalFp32.GetValue(ubFactorElementAligned_) < kthValue) { + ProcessKLtKMax(loopBatch); + } else { + ProcessRemain(loopBatch); + } + } + kInQueue_.FreeTensor(kLocal); + pInQueue_.FreeTensor(pLocal); + sortedValueInQueue_.FreeTensor(sortedValueLocal); + sortedIndicesInQueue_.FreeTensor(sortedIndicesLocal); + outQueue_.FreeTensor(outTensor); +} + +template +__aicore__ inline void ApplyTopKTopPCustom::InitCopyIn(uint32_t loopBatch, + int64_t currentGmIdx) { + DataCopyPad(mGmOut_[currentGmIdx], outTensor, {1, (uint32_t)(dataNumInit_ * sizeof(outputT)), 0, 0, 0}); + DataCopyPad(sortedValueLocal[ubFactorElementAligned_], mGmSortedValue_[currentGmIdx], + {1, static_cast(dataNumInit_ * sizeof(inputT)), 0, 0, 0}, + {false, 0, 0, 0}); + DataCopyPad(pLocal, mGmP_[batchOffset_ + loopBatch], {1, static_cast(sizeof(inputT)), 0, 0, 0}, + {false, 0, 0, 0}); + if constexpr (!IsSameType::value) { + MTE2ToVSync(); + Cast(sortedValueLocalFp32[ubFactorElementAligned_], sortedValueLocal[ubFactorElementAligned_], + RoundMode::CAST_NONE, dataNumInit_); + Cast(tmpLocal, pLocal, RoundMode::CAST_NONE, DATA_PER_BLOCK_B32); + } + DataCopyPad(sortedIndicesLocal[ubFactorElementAligned_], mGmSortedIndices_[currentGmIdx], + {1, static_cast(dataNumInit_ * sizeof(int32_t)), 0, 0, 0}, + {false, 0, 0, 0}); + DataCopyPad(kLocal, mGmK_[batchOffset_ + loopBatch], {1, static_cast(sizeof(int32_t)), 0, 0, 0}, + {false, 0, 0, 0}); +} + +template +__aicore__ inline void ApplyTopKTopPCustom::GetKthResult(uint32_t loopBatch, + uint32_t offset, uint8_t repeatTimes){ + Compare(tmpLocal.template ReinterpretCast(), kthValueLocal, calLocalFp32[offset], + CMPMODE::GT, MASK_64, repeatTimes, repeatParams); + PipeBarrier(); + Select(calLocalFp32[offset], tmpLocal.template ReinterpretCast(), + negInfLocal, calLocalFp32[offset], SELMODE::VSEL_TENSOR_TENSOR_MODE, MASK_64, + repeatTimes, repeatParams); +} + +template +__aicore__ inline void ApplyTopKTopPCustom::ReduceSumWithAddsAndExpImpl(uint32_t offset, + uint32_t loopDataNum) { + Adds(softMaxRes, calLocalFp32[offset], maxValue, loopDataNum); + PipeBarrier(); + Exp(softMaxRes, softMaxRes, loopDataNum); + PipeBarrier(); + ReduceSum(reduceLocal, softMaxRes, reduceLocal, loopDataNum); +} + +template +__aicore__ inline void ApplyTopKTopPCustom::InitProcess(uint32_t loopBatch) { + int64_t initGmIdx = baseGmIdx_ + vocabSize_ - dataNumInit_; + InitCopyIn(loopBatch, initGmIdx); + MTE2ToSSync(); + + int32_t kValue = kLocal.GetValue(0); + if constexpr (IsSameType::value) { + pValue = float(1.0) - pLocal.GetValue(0); + } else { + pValue = float(1.0) - tmpLocal.GetValue(0); + } + maxValue = -calLocalFp32[ubFactorElementAligned_].GetValue(dataNumInit_ - 1); + if constexpr (IsSameType::value) { + kthValue = mGmSortedValue_[baseGmIdx_ + vocabSize_ - kValue].GetValue(0); + } else if constexpr (IsSameType::value) { + kthValue = static_cast(mGmSortedValue_[baseGmIdx_ + vocabSize_ - kValue].GetValue(0)); + } else { + kthValue = ToFloat(mGmSortedValue_[baseGmIdx_ + vocabSize_ - kValue].GetValue(0)); + } + + Duplicate(kthValueLocal, kthValue, 8); + PipeBarrier(); + + uint8_t repeatTimes = (dataNumInit_ + DATA_PER_REPEAT_B32 - 1) / DATA_PER_REPEAT_B32; + GetKthResult(loopBatch, ubFactorElementAligned_, repeatTimes); + PipeBarrier(); + DataCopyExtParams copyParams{1, (uint32_t)(ubFactorElementAligned_ * sizeof(outputT)), 0, 0, 0}; + + ReduceSumWithAddsAndExpImpl(ubFactorElementAligned_, dataNumInit_); + VToSSync(); + reduceSumValue = reduceLocal.GetValue(0); + reduceSumValueInvert = 1 / reduceSumValue; +} + +template +__aicore__ inline void ApplyTopKTopPCustom::ProcessKLtKMax(uint32_t loopBatch) { + DataCopyExtParams copyParams{1, (uint32_t)(ubFactorElementAligned_ * sizeof(outputT)), 0, 0, 0}; + for (int32_t loopInner = 0; loopInner < loopInner_; loopInner++) { + int64_t currentGmIdxInner = baseGmIdx_ + loopInner * ubFactorElementAligned_; + if (loopInner == loopInner_ - 1) { + DataCopyPad(mGmOut_[currentGmIdxInner], outTensor, + {1, (uint32_t)(tailUbFactorElement_ * sizeof(outputT)), 0, 0, 0}); + } else { + DataCopyPad(mGmOut_[currentGmIdxInner], outTensor, copyParams); + } + } + Muls(softMaxRes, softMaxRes, reduceSumValueInvert, dataNumInit_); + PipeBarrier(); + const CumSumInfo cumSumInfo{1, dataNumInitAligned_}; + CumSum(cumSumRes, cumSumTmp, softMaxRes, sharedTmpBuffer, cumSumInfo); + VToSSync(); + int32_t loopProb = dataNumInit_ - 1; + scatterTensor.SetValue(0, sortedValueLocal[ubFactorElementAligned_].GetValue(loopProb)); + SToMTE3Sync(); + int32_t gmIndex = sortedIndicesLocal[ubFactorElementAligned_].GetValue(loopProb); + PipeBarrier(); + DataCopyPad(mGmOut_[baseGmIdx_ + gmIndex], scatterTensor.template ReinterpretCast(), scatterCopyParams); + MTE3ToSSync(); + loopProb = loopProb - 1; + for (; loopProb >= 0; loopProb--) { + float cumsumData = cumSumRes.GetValue(loopProb); + if (cumsumData <= pValue) { + break; + } + scatterTensor.SetValue(0, sortedValueLocal[ubFactorElementAligned_].GetValue(loopProb)); + gmIndex = sortedIndicesLocal[ubFactorElementAligned_].GetValue(loopProb); + SToMTE3Sync(); + DataCopyPad(mGmOut_[baseGmIdx_ + gmIndex], + scatterTensor.template ReinterpretCast(), scatterCopyParams); + MTE3ToSSync(); + } +} + +template +__aicore__ inline void ApplyTopKTopPCustom::ScatterCumtomImpl(uint32_t loopBatch, + uint32_t loopProbNum, uint32_t offset) { + for (int32_t loopProb = 0; loopProb < static_cast(loopProbNum); loopProb++) { + float cumsumDataTmp = cumSumRes.GetValue(loopProb); + if (cumsumDataTmp <= pValue) { + continue; + } + scatterTensor.SetValue(0, sortedValueLocal[offset].GetValue(loopProb)); + int32_t gmIndex = sortedIndicesLocal[offset].GetValue(loopProb); + SToMTE3Sync(); + DataCopyPad(mGmOut_[baseGmIdx_ + gmIndex], scatterTensor.template ReinterpretCast(), + {1, (uint32_t)(1 * sizeof(outputT)), 0, 0, 0}); + MTE3ToSSync(); + } +} + +template +__aicore__ inline void ApplyTopKTopPCustom::GetFirstKLoop(uint32_t loopBatch, + int32_t &firstKLoop) { + uint8_t repeatTimes = (dataNumInit_ + DATA_PER_REPEAT_B32 - 1) / DATA_PER_REPEAT_B32; + uint32_t loopDataNum = ubFactorElementAligned_; + for (int32_t loopInner = 0; loopInner < loopInner_; loopInner++) { + int64_t currentGmIdx = baseGmIdx_ + loopInner * ubFactorElementAligned_; + if (loopInner == (loopInner_ - 1)) { + repeatTimes = ((tailUbFactorElement_) + DATA_PER_REPEAT_B32 - 1) / DATA_PER_REPEAT_B32; + loopDataNum = tailUbFactorElement_; + } + DataCopyPad(mGmOut_[currentGmIdx], outTensor, {1, (uint32_t)(loopDataNum * sizeof(outputT)), 0, 0, 0}); + DataCopyPad(sortedValueLocal.template ReinterpretCast(), mGmSortedValue_[currentGmIdx], + {1, static_cast(loopDataNum * sizeof(inputT)), 0, 0, 0}, + {false, 0, 0, 0}); + if constexpr (!IsSameType::value) { + MTE2ToVSync(); + Cast(sortedValueLocalFp32, sortedValueLocal, RoundMode::CAST_NONE, loopDataNum); + VToSSync(); + } else { + MTE2ToSSync(); + } + if (calLocalFp32.GetValue(loopDataNum - 1) < kthValue) { + firstKLoop += 1; + continue; + } + + GetKthResult(loopBatch, 0, repeatTimes); + PipeBarrier(); + + ReduceSumWithAddsAndExpImpl(0, loopDataNum); + VToSSync(); + reduceSumValue += reduceLocal.GetValue(0); + } +} + +template +__aicore__ inline void ApplyTopKTopPCustom::CumSumWithAddsAndExpImpl(uint32_t offset, + uint32_t loopDataNum, uint32_t cumsumInner, float cumsumData) { + Adds(softMaxRes, calLocalFp32[offset], maxValue, loopDataNum); + PipeBarrier(); + Exp(softMaxRes, softMaxRes, loopDataNum); + PipeBarrier(); + Muls(softMaxRes, softMaxRes, reduceSumValueInvert, loopDataNum); + PipeBarrier(); + const CumSumInfo cumSumInfo{1, cumsumInner}; + CumSum(cumSumRes, cumSumTmp, softMaxRes, sharedTmpBuffer, cumSumInfo); + PipeBarrier(); + Adds(cumSumRes, cumSumRes, cumsumData, loopDataNum); +} + +template +__aicore__ inline void ApplyTopKTopPCustom::ProcessRemain(uint32_t loopBatch) { + int32_t firstKLoop = 0; + GetFirstKLoop(loopBatch, firstKLoop); + reduceSumValueInvert = 1 / reduceSumValue; + float cumsumData = 0; + ScatterFromFirstKLoop(loopBatch, firstKLoop, cumsumData); + uint32_t loopProb = dataNumInit_ - 1; + scatterTensor.SetValue(0, sortedValueLocal[ubFactorElementAligned_].GetValue(loopProb)); + int32_t gmIndex = sortedIndicesLocal[ubFactorElementAligned_].GetValue(loopProb); + SToMTE3Sync(); + DataCopyPad(mGmOut_[baseGmIdx_ + gmIndex], + scatterTensor.template ReinterpretCast(), scatterCopyParams); + MTE3ToVSync(); + CumSumWithAddsAndExpImpl(ubFactorElementAligned_, dataNumInit_, dataNumInitAligned_, cumsumData); + VToSSync(); + ScatterCumtomImpl(loopBatch, dataNumInit_ - 1, ubFactorElementAligned_); +} + +template +__aicore__ inline void ApplyTopKTopPCustom::ScatterFromFirstKLoop(uint32_t loopBatch, + int32_t firstKLoop, float &cumsumData) { + uint32_t loopDataNum = ubFactorElementAligned_; + uint32_t cumsumInner = ubFactorElementAligned_; + uint8_t repeatTimes = ((ubFactorElementAligned_) + DATA_PER_REPEAT_B32 - 1) / DATA_PER_REPEAT_B32; + for (int32_t loopInner = firstKLoop; loopInner < loopInner_; loopInner++) { + int64_t currentGmIdx = baseGmIdx_ + loopInner * ubFactorElementAligned_; + if (loopInner == (loopInner_ - 1)) { + repeatTimes = (tailUbFactorElement_ + DATA_PER_REPEAT_B32 - 1) / DATA_PER_REPEAT_B32; + loopDataNum = tailUbFactorElement_; + cumsumInner = tailUbFactorElementAligned_; + } + DataCopyPad(sortedValueLocal.template ReinterpretCast(), mGmSortedValue_[currentGmIdx], + {1, static_cast(loopDataNum * sizeof(inputT)), 0, 0, 0}, + {false, 0, 0, 0}); + DataCopyPad(sortedIndicesLocal, mGmSortedIndices_[currentGmIdx], + {1, static_cast(loopDataNum * sizeof(int32_t)), 0, 0, 0}, + {false, 0, 0, 0}); + if constexpr (!IsSameType::value) { + MTE2ToVSync(); + Cast(sortedValueLocalFp32, sortedValueLocal, RoundMode::CAST_NONE, loopDataNum); + PipeBarrier(); + } else { + MTE2ToVSync(); + } + GetKthResult(loopBatch, 0, repeatTimes); + PipeBarrier(); + CumSumWithAddsAndExpImpl(0, loopDataNum, cumsumInner, cumsumData); + VToSSync(); + float cumsumDataTmp = cumSumRes.GetValue(loopDataNum - 1); + cumsumData = cumsumDataTmp; + if (cumsumDataTmp <= pValue) { + continue; + } + ScatterCumtomImpl(loopBatch, loopDataNum, 0); + } +} + +template +__aicore__ inline void ApplyTopKTopPCustom::ProcessTopK() { + kLocal = kInQueue_.AllocTensor(); + outTensor = outQueue_.AllocTensor(); + sortedValueLocal = sortedValueInQueue_.AllocTensor(); + sortedIndicesLocal = sortedIndicesInQueue_.AllocTensor(); + Duplicate(negInfLocal.template ReinterpretCast(), FLOAT32_NEG_INF, DATA_PER_BLOCK_B32); + if constexpr (IsSameType::value) { + calLocalFp32 = sortedValueLocal; + Duplicate(outTensor.template ReinterpretCast(), FLOAT32_NEG_INF, ubFactorElementAligned_); + } else if constexpr (IsSameType::value) { + calLocalFp32 = sortedValueLocalFp32; + Duplicate(outTensor.template ReinterpretCast(), FLOAT16_NEG_INF, ubFactorElementAligned_); + } else { + calLocalFp32 = sortedValueLocalFp32; + Duplicate(outTensor.template ReinterpretCast(), BF16_NEG_INF, ubFactorElementAligned_); + } + VToMTE3Sync(); + for (uint32_t loopBatch = 0; loopBatch < loopBatch_; loopBatch++) { + baseGmIdx_ = batchOffset_ * vocabSize_ + loopBatch * vocabSize_; + InitProcessTopK(loopBatch); + /* The difference lies in that for the max branch, some data is less than the kthvalue, + so part of the data can be filtered out in advance; + while for the remain branch, all data must undergo the topk calculation.*/ + if (calLocalFp32.GetValue(ubFactorElementAligned_) < kthValue) { + ProcessKLtKMaxTopK(loopBatch); + } else { + ProcessRemainTopK(loopBatch); + } + } + kInQueue_.FreeTensor(kLocal); + sortedValueInQueue_.FreeTensor(sortedValueLocal); + sortedIndicesInQueue_.FreeTensor(sortedIndicesLocal); + outQueue_.FreeTensor(outTensor); +} + +template +__aicore__ inline void ApplyTopKTopPCustom::ProcessRemainTopK(uint32_t loopBatch) { + int32_t firstKLoop = 0; + GetFirstKLoopTopK(loopBatch, firstKLoop); + // Start the scatter calculation from the first loop in the row where the value is ≥ kthValue. + ScatterFromFirstKLoopTopK(loopBatch, firstKLoop); + /* Perform scatter calculation on the maximum number of ubFactorElementAligned_, + which does not overlap with the previous ones.*/ + uint32_t loopProb = dataNumInit_ - 1; + scatterTensor.SetValue(0, sortedValueLocal[ubFactorElementAligned_].GetValue(loopProb)); + SToMTE3Sync(); + int32_t gmIndex = sortedIndicesLocal[ubFactorElementAligned_].GetValue(loopProb); + DataCopyPad(mGmOut_[baseGmIdx_ + gmIndex], + scatterTensor.template ReinterpretCast(), scatterCopyParams); + MTE3ToSSync(); + ScatterCumtomImplTopK(loopBatch, dataNumInit_ - 1, ubFactorElementAligned_); +} + +template +__aicore__ inline void ApplyTopKTopPCustom::GetFirstKLoopTopK(uint32_t loopBatch, + int32_t &firstKLoop) { + uint8_t repeatTimes = (dataNumInit_ + DATA_PER_REPEAT_B32 - 1) / DATA_PER_REPEAT_B32; + uint32_t loopDataNum = ubFactorElementAligned_; + for (int32_t loopInner = 0; loopInner < loopInner_; loopInner++) { + int64_t currentGmIdx = baseGmIdx_ + loopInner * ubFactorElementAligned_; + if (loopInner == (loopInner_ - 1)) { + repeatTimes = ((tailUbFactorElement_) + DATA_PER_REPEAT_B32 - 1) / DATA_PER_REPEAT_B32; + loopDataNum = tailUbFactorElement_; + } + DataCopyPad(mGmOut_[currentGmIdx], outTensor, {1, (uint32_t)(loopDataNum * sizeof(outputT)), 0, 0, 0}); + DataCopyPad(sortedValueLocal.template ReinterpretCast(), mGmSortedValue_[currentGmIdx], + {1, static_cast(loopDataNum * sizeof(inputT)), 0, 0, 0}, + {false, 0, 0, 0}); + MTE2ToSSync(); + float rightVlaue = 0; + // Make a judgment on the rightmost value of each loop to filter the data. + if constexpr (IsSameType::value) { + rightVlaue = ToFloat(sortedValueLocal.GetValue(loopDataNum - 1)); + } else { + rightVlaue = static_cast(sortedValueLocal.GetValue(loopDataNum - 1)); + } + SToMTE2Sync(); + if (rightVlaue < kthValue) { + firstKLoop += 1; + continue; + } + } +} + +template +__aicore__ inline void ApplyTopKTopPCustom::ScatterFromFirstKLoopTopK(uint32_t loopBatch, + int32_t firstKLoop) { + uint32_t loopDataNum = ubFactorElementAligned_; + uint32_t cumsumInner = ubFactorElementAligned_; + uint8_t repeatTimes = ((ubFactorElementAligned_) + DATA_PER_REPEAT_B32 - 1) / DATA_PER_REPEAT_B32; + for (int32_t loopInner = firstKLoop; loopInner < loopInner_; loopInner++) { + int64_t currentGmIdx = baseGmIdx_ + loopInner * ubFactorElementAligned_; + if (loopInner == (loopInner_ - 1)) { + repeatTimes = (tailUbFactorElement_ + DATA_PER_REPEAT_B32 - 1) / DATA_PER_REPEAT_B32; + loopDataNum = tailUbFactorElement_; + cumsumInner = tailUbFactorElementAligned_; + } + DataCopyPad(sortedValueLocal.template ReinterpretCast(), mGmSortedValue_[currentGmIdx], + {1, static_cast(loopDataNum * sizeof(inputT)), 0, 0, 0}, + {false, 0, 0, 0}); + if constexpr (!IsSameType::value) { + MTE2ToVSync(); + Cast(sortedValueLocalFp32, sortedValueLocal, RoundMode::CAST_NONE, loopDataNum); + VToSSync(); + } + DataCopyPad(sortedIndicesLocal, mGmSortedIndices_[currentGmIdx], + {1, static_cast(loopDataNum * sizeof(int32_t)), 0, 0, 0}, + {false, 0, 0, 0}); + MTE2ToSSync(); + ScatterCumtomImplTopK(loopBatch, loopDataNum, 0); + } +} + +template +__aicore__ inline void ApplyTopKTopPCustom::ScatterCumtomImplTopK(uint32_t loopBatch, + uint32_t loopProbNum, uint32_t offset) { + // Reverse traversal, returning early to improve performance. + for (int32_t loopProb = static_cast(loopProbNum) - 1; loopProb >= 0; loopProb--) { + float curValue = calLocalFp32[offset].GetValue(loopProb); + if (curValue < kthValue) { + break; + } + scatterTensor.SetValue(0, sortedValueLocal[offset].GetValue(loopProb)); + int32_t gmIndex = sortedIndicesLocal[offset].GetValue(loopProb); + SToMTE3Sync(); + DataCopyPad(mGmOut_[baseGmIdx_ + gmIndex], scatterTensor.template ReinterpretCast(), + {1, (uint32_t)(1 * sizeof(outputT)), 0, 0, 0}); + MTE3ToSSync(); + } +} + +template +__aicore__ inline void ApplyTopKTopPCustom::InitProcessTopK(uint32_t loopBatch) { + int64_t initGmIdx = baseGmIdx_ + vocabSize_ - dataNumInit_; + DataCopyPad(mGmOut_[initGmIdx], outTensor, {1, (uint32_t)(dataNumInit_ * sizeof(outputT)), 0, 0, 0}); + DataCopyPad(sortedValueLocal[ubFactorElementAligned_], mGmSortedValue_[initGmIdx], + {1, static_cast(dataNumInit_ * sizeof(inputT)), 0, 0, 0}, + {false, 0, 0, 0}); + if constexpr (!IsSameType::value) { + MTE2ToVSync(); + Cast(sortedValueLocalFp32[ubFactorElementAligned_], sortedValueLocal[ubFactorElementAligned_], + RoundMode::CAST_NONE, dataNumInit_); + } + DataCopyPad(sortedIndicesLocal[ubFactorElementAligned_], mGmSortedIndices_[initGmIdx], + {1, static_cast(dataNumInit_ * sizeof(int32_t)), 0, 0, 0}, + {false, 0, 0, 0}); + DataCopyPad(kLocal, mGmK_[batchOffset_ + loopBatch], {1, static_cast(sizeof(int32_t)), 0, 0, 0}, + {false, 0, 0, 0}); + MTE2ToSSync(); + int32_t kValue = mGmK_.GetValue(batchOffset_ + loopBatch); + maxValue = -calLocalFp32[ubFactorElementAligned_].GetValue(dataNumInit_ - 1); + if constexpr (IsSameType::value) { + kthValue = mGmSortedValue_[baseGmIdx_ + vocabSize_ - kValue].GetValue(0); + } else if constexpr (IsSameType::value) { + kthValue = static_cast(mGmSortedValue_[baseGmIdx_ + vocabSize_ - kValue].GetValue(0)); + } else { + kthValue = ToFloat(mGmSortedValue_[baseGmIdx_ + vocabSize_ - kValue].GetValue(0)); + } +} + +template +__aicore__ inline void ApplyTopKTopPCustom::ProcessKLtKMaxTopK(uint32_t loopBatch) { + DataCopyExtParams copyParams{1, (uint32_t)(ubFactorElementAligned_ * sizeof(outputT)), 0, 0, 0}; + // Move out -infinity to fill GM + for (int32_t loopInner = 0; loopInner < loopInner_; loopInner++) { + int64_t currentGmIdxInner = baseGmIdx_ + loopInner * ubFactorElementAligned_; + if (loopInner == loopInner_ - 1) { + DataCopyPad(mGmOut_[currentGmIdxInner], outTensor, + {1, (uint32_t)(tailUbFactorElement_ * sizeof(outputT)), 0, 0, 0}); + } else { + DataCopyPad(mGmOut_[currentGmIdxInner], outTensor, copyParams); + } + } + // Scatter calculation + int32_t loopProb = dataNumInit_ - 1; + scatterTensor.SetValue(0, sortedValueLocal[ubFactorElementAligned_].GetValue(loopProb)); + int32_t gmIndex = sortedIndicesLocal[ubFactorElementAligned_].GetValue(loopProb); + SToMTE3Sync(); + PipeBarrier(); + DataCopyPad(mGmOut_[baseGmIdx_ + gmIndex], scatterTensor.template ReinterpretCast(), scatterCopyParams); + loopProb = loopProb - 1; + + for (; loopProb >= 0; loopProb--) { + float curValue = calLocalFp32[ubFactorElementAligned_].GetValue(loopProb); + if (curValue < kthValue) { + break; + } + MTE3ToSSync(); + scatterTensor.SetValue(0, sortedValueLocal[ubFactorElementAligned_].GetValue(loopProb)); + gmIndex = sortedIndicesLocal[ubFactorElementAligned_].GetValue(loopProb); + SToMTE3Sync(); + DataCopyPad(mGmOut_[baseGmIdx_ + gmIndex], + scatterTensor.template ReinterpretCast(), scatterCopyParams); + } +} + +} // namespace + +#endif // APPLY_TOP_K_TOP_P_CUSTOM_H_KERNEL \ No newline at end of file diff --git a/csrc/ascend/moe/apply_top_k_top_p_custom/op_kernel/apply_top_p_custom.h b/csrc/ascend/moe/apply_top_k_top_p_custom/op_kernel/apply_top_p_custom.h new file mode 100644 index 000000000..f58619877 --- /dev/null +++ b/csrc/ascend/moe/apply_top_k_top_p_custom/op_kernel/apply_top_p_custom.h @@ -0,0 +1,468 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file apply_top_p_custom.h + * \brief + */ +#ifndef APPLY_TOP_P_CUSTOM_H_KERNEL +#define APPLY_TOP_P_CUSTOM_H_KERNEL + +#include "kernel_operator.h" + +using namespace AscendC; +namespace ApplyTopPCustomOp { +constexpr uint16_t FLOAT16_NEG_INF = 0xFC00; // -inf 64512 +constexpr uint16_t BF16_NEG_INF = 0xFF80; // -inf 65408 +constexpr int32_t FLOAT32_NEG_INF = 0xFF800000; // -inf -2139095040 + +constexpr uint32_t BLOCK_BYTES = 32; +constexpr uint32_t DATA_PER_BLOCK_B32 = 8; +constexpr uint32_t DATA_PER_REPEAT_B32 = 64; +constexpr uint32_t SCATTER_PART_LENGTH = 1024; +constexpr uint32_t RESERVED_UB = 1024; +constexpr uint32_t FLOAT_BYTES = 4; +constexpr uint32_t SOFTMAX_UB_NUM = 2; + +template +class ApplyTopPCustom { +public: + __aicore__ inline ApplyTopPCustom(){}; + __aicore__ inline void InitTilingData( + const ApplyTopKTopPCustomTilingData &__restrict tilingData, GM_ADDR sorted_value, GM_ADDR sorted_indices, GM_ADDR p, GM_ADDR k, GM_ADDR out, GM_ADDR workspace); + __aicore__ inline void InitBuffer(TPipe *inputPipe); + __aicore__ inline void ProcessTopP(); +private: + __aicore__ inline void ReduceSumWithAddsAndExpImpl(uint32_t loopDataNum); + // topp func + __aicore__ inline void ProcessPreSingleBatch(uint32_t loopBatch); + __aicore__ inline void GetSoftmaxSum(uint32_t loopBatch); + __aicore__ inline void CumsumKoggleStone(uint32_t loopBatch); + __aicore__ inline void GetPValue(uint32_t batchOffset); + __aicore__ inline void CumsumParamCompute(uint32_t iterateTime); + __aicore__ inline void GetMaxValue(int64_t baseGmIdx); + __aicore__ inline void GetSoftMaxRes(uint32_t loopBatch); + __aicore__ inline void ScatterSingleTask(uint32_t taskIndex); + + __aicore__ inline void SToMTE3Sync() { + event_t eventIDSToMTE3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3)); + SetFlag(eventIDSToMTE3); + WaitFlag(eventIDSToMTE3); + } + __aicore__ inline void VToMTE3Sync() { + event_t eventIDVToMTE3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventIDVToMTE3); + WaitFlag(eventIDVToMTE3); + } + __aicore__ inline void VToMTE2Sync() { + event_t eventIDVToMTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2)); + SetFlag(eventIDVToMTE2); + WaitFlag(eventIDVToMTE2); + } + __aicore__ inline void MTE3ToSSync() { + event_t eventIDMTE3ToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S)); + SetFlag(eventIDMTE3ToS); + WaitFlag(eventIDMTE3ToS); + } + __aicore__ inline void VToSSync() { + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + } + __aicore__ inline void SToVSync() { + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + } + __aicore__ inline void MTE2ToVSync() { + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + } + __aicore__ inline void MTE2ToSSync() { + event_t eventIdMte2ToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S)); + SetFlag(eventIdMte2ToS); + WaitFlag(eventIdMte2ToS); + } + + __aicore__ inline void MTE3ToMTE2Sync() { + event_t eventIdMte3ToMte2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventIdMte3ToMte2); + WaitFlag(eventIdMte3ToMte2); + } + + __aicore__ inline void MTE3ToVSync() { + event_t eventIdMte3ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V)); + SetFlag(eventIdMte3ToV); + WaitFlag(eventIdMte3ToV); + } + + __aicore__ inline uint32_t CeilDiv(uint32_t x, uint32_t y) { + return y == 0 ? x : (x + y - 1) / y; + } +private: + TPipe *pipe_; + TBuf calBuf_; + + // tilingData + uint32_t batchSize_ = 0; + uint32_t vocabSize_ = 0; + uint32_t batchPerCore_ = 0; + uint32_t tailBatch_ = 0; + uint32_t blockNum_ = 0; + uint32_t dataNumInitAligned_ = 0; + uint32_t calUbSize_ = 0; + uint32_t blockIdx_ = 0; + uint32_t loopBatch_ = 0; + uint32_t batchOffset_ = 0; + uint32_t bufOffsetLoop = 0; + int64_t baseGmIdx_ = 0; + + // topp scalar + uint32_t maxSoftmaxLength = 1; + uint32_t softmaxLength = 1; + uint32_t lineSfLoopTimes = 1; + uint32_t softmaxLengthTail = 1; + uint32_t scatterLength = 1; + + uint32_t singleCoreB = 0; + uint32_t singleCoreBTail = 0; + uint32_t vCnt = 0; + uint32_t bCnt = 0; + uint32_t singleCoreV = 1; + uint32_t singleCoreVTail = 1; + uint32_t iterateTimes = 1; + + GlobalTensor mGmSortedValue_; + GlobalTensor mGmSortedIndices_; + GlobalTensor mGmP_; + GlobalTensor mGmK_; + GlobalTensor mGmOut_; + GlobalTensor softMaxGm; + + LocalTensor totalUb; + + // softmax tensor + LocalTensor softMaxLocalFp32; + LocalTensor softMaxLocal; + LocalTensor softMaxResLocal; + LocalTensor reduceLocal; + LocalTensor outInfLocal; + + // cumsum tensor + LocalTensor cumSumInput1Local; + LocalTensor cumSumInput2Local; + + // scatter tensor + LocalTensor sortedValueLocal; + LocalTensor sortedIndicesLocal; + LocalTensor sortedValueLocalFp32; + LocalTensor scatterLocal; + LocalTensor cumsumLocal; + + float pValue = 0; + float maxValue = 0; + float reduceSumValueInvert = 0; + float reduceSumValue = 0; + BinaryRepeatParams repeatParams = {1, 0, 1, 8, 0, 8}; + DataCopyExtParams scatterCopyParams{1, (uint32_t)(sizeof(outputT)), 0, 0, 0}; +}; + +template +__aicore__ inline void ApplyTopPCustom::InitTilingData( + const ApplyTopKTopPCustomTilingData &__restrict tilingData, GM_ADDR sorted_value, GM_ADDR sorted_indices, + GM_ADDR p, GM_ADDR k, GM_ADDR out, GM_ADDR workspace) { + batchSize_ = tilingData.batchSize; + vocabSize_ = tilingData.vocabSize; + batchPerCore_ = tilingData.batchPerCore; + tailBatch_ = tilingData.tailBatch; + blockNum_ = tilingData.blockNum; + calUbSize_ = tilingData.calUbSize; + iterateTimes = tilingData.iterateTimes; + blockIdx_ = GetBlockIdx(); + if (blockIdx_ < tailBatch_) { + loopBatch_ = batchPerCore_ + 1; + batchOffset_ = blockIdx_ * loopBatch_; + } else { + loopBatch_ = batchPerCore_; + batchOffset_ = blockIdx_ * batchPerCore_ + tailBatch_; + } + maxSoftmaxLength = (calUbSize_ - RESERVED_UB) / SOFTMAX_UB_NUM / FLOAT_BYTES; + softmaxLength = maxSoftmaxLength < vocabSize_ ? maxSoftmaxLength : vocabSize_; + + lineSfLoopTimes = (vocabSize_ + softmaxLength - 1) / softmaxLength; + softmaxLengthTail = vocabSize_ - (lineSfLoopTimes - 1) * softmaxLength; + scatterLength = (calUbSize_ - RESERVED_UB - BLOCK_BYTES) / (SOFTMAX_UB_NUM * FLOAT_BYTES + sizeof(inputT)) / + SCATTER_PART_LENGTH * SCATTER_PART_LENGTH; + singleCoreB = CeilDiv(batchSize_, blockNum_); + vCnt = batchSize_ < blockNum_ ? blockNum_ / batchSize_ : 1; + bCnt = batchSize_; + singleCoreB = 1; + singleCoreBTail = 1; + singleCoreV = vocabSize_ / vCnt; + singleCoreVTail = vocabSize_ - vCnt * singleCoreV; + mGmSortedValue_.SetGlobalBuffer(reinterpret_cast<__gm__ inputT *>(sorted_value)); + mGmSortedIndices_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(sorted_indices)); + mGmP_.SetGlobalBuffer(reinterpret_cast<__gm__ inputT *>(p)); + mGmK_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(k)); + mGmOut_.SetGlobalBuffer(reinterpret_cast<__gm__ outputT *>(out)); + softMaxGm.SetGlobalBuffer((__gm__ float*)workspace, batchSize_ * vocabSize_); +} + +// init used buffer +template +__aicore__ inline void ApplyTopPCustom::InitBuffer(TPipe *inputPipe) { + pipe_ = inputPipe; + pipe_->InitBuffer(calBuf_, calUbSize_); + totalUb = calBuf_.Get(); + // softmax ub + uint32_t softmaxLengthAligned = CeilDiv(softmaxLength, BLOCK_BYTES / sizeof(inputT)) * BLOCK_BYTES / sizeof(inputT); + softMaxLocalFp32 = totalUb.ReinterpretCast(); + softMaxLocal = totalUb[softmaxLengthAligned * sizeof(inputT)].ReinterpretCast(); + softMaxResLocal = totalUb[softmaxLengthAligned * sizeof(float)].ReinterpretCast(); + reduceLocal = totalUb[softmaxLengthAligned * sizeof(float) * 2].ReinterpretCast(); // 32 bytes + outInfLocal = totalUb.ReinterpretCast(); // Take softmax ub + + // cumsum ub + cumSumInput1Local = totalUb.ReinterpretCast(); // Take softmax local + cumSumInput2Local = totalUb[softmaxLengthAligned * sizeof(float)].ReinterpretCast(); // Take softmax res ub + + // scatter ub + sortedValueLocal = totalUb[0].ReinterpretCast(); + sortedIndicesLocal = totalUb[scatterLength * sizeof(inputT)].ReinterpretCast(); + cumsumLocal = totalUb[scatterLength * (FLOAT_BYTES + sizeof(inputT))].ReinterpretCast(); + scatterLocal = totalUb[calUbSize_ - RESERVED_UB + BLOCK_BYTES].ReinterpretCast(); // 32 bytes +} + +template +__aicore__ inline void ApplyTopPCustom::GetMaxValue(int64_t baseGmIdx) { + int64_t initGmIdx = baseGmIdx + vocabSize_ - 1; + if constexpr (IsSameType::value) { + maxValue = -mGmSortedValue_[initGmIdx].GetValue(0); + } else if constexpr (IsSameType::value) { + maxValue = -static_cast(mGmSortedValue_[initGmIdx].GetValue(0)); + } else { + maxValue = -ToFloat(mGmSortedValue_[initGmIdx].GetValue(0)); + } +} + +template +__aicore__ inline void ApplyTopPCustom::GetPValue(uint32_t batchOffset) { + if constexpr (IsSameType::value) { + pValue = float(1.0) - mGmP_[batchOffset].GetValue(0); + } else if constexpr (IsSameType::value) { + pValue = float(1.0) - static_cast(mGmP_[batchOffset].GetValue(0)); + } else { + pValue = float(1.0) - ToFloat(mGmP_[batchOffset].GetValue(0)); + } +} + +template +__aicore__ inline void ApplyTopPCustom::ProcessPreSingleBatch(uint32_t loopBatch) { + reduceSumValue = 0; + GetSoftmaxSum(loopBatch); + GetSoftMaxRes(loopBatch); + CumsumKoggleStone(loopBatch); +} + +template +__aicore__ inline void ApplyTopPCustom::ProcessTopP() { + for (uint32_t loopBatch = 0; loopBatch < loopBatch_; loopBatch++) { + baseGmIdx_ = batchOffset_ * vocabSize_ + loopBatch * vocabSize_; + GetMaxValue(baseGmIdx_); // Get max value in softmax. + ProcessPreSingleBatch(loopBatch); // Softmax and cumsum. + } + SyncAll(); + for (uint32_t taskIndex = 0; taskIndex < bCnt * vCnt; taskIndex++) { + ScatterSingleTask(taskIndex); + } +} + +template +__aicore__ inline void ApplyTopPCustom::ScatterSingleTask(uint32_t taskIndex) { + if (GetBlockIdx() == taskIndex % blockNum_) { + uint32_t bCntIndex = taskIndex / vCnt; + uint32_t vCntIndex = taskIndex % vCnt; + uint32_t vCurSingleCore = vCntIndex < singleCoreVTail ? (singleCoreV + 1) : singleCoreV; + uint32_t copyTimes = CeilDiv(vCurSingleCore, scatterLength); + uint32_t copyLength = scatterLength; + uint32_t copyLengthTail = vCurSingleCore - (copyTimes - 1) * scatterLength; + GetPValue(bCntIndex); // Get maxPValue. + for (uint32_t cpIndex = 0; cpIndex < copyTimes; cpIndex++) { + uint32_t curCopyLength = cpIndex == (copyTimes - 1) ? copyLengthTail : copyLength; + int64_t gmOffset = vCntIndex < singleCoreVTail ? + bCntIndex * vocabSize_ + vCntIndex * (singleCoreV + 1) + cpIndex * copyLength : + bCntIndex * vocabSize_ + vCntIndex * singleCoreV + singleCoreVTail + cpIndex * copyLength; + DataCopyPad(cumsumLocal, softMaxGm[gmOffset], + {1, static_cast(curCopyLength * sizeof(float)), 0, 0, 0}, {false, 0, 0, 0}); + DataCopyPad(sortedIndicesLocal, mGmSortedIndices_[gmOffset], + {1, static_cast(curCopyLength * sizeof(float)), 0, 0, 0}, {false, 0, 0, 0}); + DataCopyPad(sortedValueLocal, mGmSortedValue_[gmOffset], + {1, static_cast(curCopyLength * sizeof(inputT)), 0, 0, 0}, {false, 0, 0, 0}); + MTE2ToSSync(); + + if (cumsumLocal.GetValue(curCopyLength - 1) <= pValue) { + continue; + } + uint32_t scatterLoop = CeilDiv(curCopyLength, SCATTER_PART_LENGTH); + uint32_t scatterNumsTail = curCopyLength - (scatterLoop - 1) * SCATTER_PART_LENGTH; + for (uint32_t scatterLoopIndex = 0; scatterLoopIndex < scatterLoop; scatterLoopIndex++) { + uint32_t curScatterNums = scatterLoopIndex == (scatterLoop - 1) ? scatterNumsTail : SCATTER_PART_LENGTH; + if (cumsumLocal.GetValue(scatterLoopIndex * SCATTER_PART_LENGTH + curScatterNums - 1) <= pValue) { + continue; + } + for (uint32_t scatterIndex = 0; scatterIndex < curScatterNums; scatterIndex++) { + int64_t scatterOffset = scatterLoopIndex * SCATTER_PART_LENGTH + scatterIndex; + if (cumsumLocal.GetValue(scatterOffset) <= pValue) { + continue; + } + scatterLocal.SetValue(0, sortedValueLocal.GetValue(scatterOffset)); + int32_t lineIndex = sortedIndicesLocal.GetValue(scatterOffset); + SToMTE3Sync(); + DataCopyPad(mGmOut_[bCntIndex * vocabSize_ + lineIndex], scatterLocal.template ReinterpretCast(), + {1, (uint32_t)(1 * sizeof(outputT)), 0, 0, 0}); + MTE3ToSSync(); + } + } + } + } +} +template +__aicore__ inline void ApplyTopPCustom::GetSoftMaxRes(uint32_t loopBatch) { + uint32_t loopDataNum = softmaxLength; + for (int32_t loopInner = 0; loopInner < lineSfLoopTimes; loopInner++) { + int64_t currentGmIdx = baseGmIdx_ + loopInner * softmaxLength; + if (loopInner == (lineSfLoopTimes - 1)) { + loopDataNum = softmaxLengthTail; + } + if constexpr (!IsSameType::value) { + DataCopyPad(softMaxLocal, mGmSortedValue_[currentGmIdx], + {1, static_cast(loopDataNum * sizeof(inputT)), 0, 0, 0}, + {false, 0, 0, 0}); + MTE2ToVSync(); + Cast(softMaxLocalFp32, softMaxLocal, RoundMode::CAST_NONE, loopDataNum); + PipeBarrier(); + } else { + DataCopyPad(softMaxLocalFp32, mGmSortedValue_[currentGmIdx], + {1, static_cast(loopDataNum * sizeof(float)), 0, 0, 0}, + {false, 0, 0, 0}); + MTE2ToVSync(); + } + Adds(softMaxResLocal, softMaxLocalFp32, maxValue, loopDataNum); + VToMTE2Sync(); + PipeBarrier(); + Exp(softMaxResLocal, softMaxResLocal, loopDataNum); + PipeBarrier(); + Muls(softMaxResLocal, softMaxResLocal, reduceSumValueInvert, loopDataNum); + VToMTE3Sync(); + DataCopyPad(softMaxGm[currentGmIdx], softMaxResLocal, + {1, static_cast(loopDataNum * sizeof(float)), 0, 0, 0}); + MTE3ToMTE2Sync(); + } +} + +template +__aicore__ inline void ApplyTopPCustom::CumsumKoggleStone(uint32_t loopBatch) { + uint32_t loopDataNum = softmaxLength; + for (uint32_t iterateTime = 0; iterateTime < iterateTimes; iterateTime++) { + int64_t iteratOffset = 1; + for (uint32_t powerIdx = 0; powerIdx < iterateTime; powerIdx++) { + iteratOffset = iteratOffset * 2; + } + uint32_t addLength = vocabSize_ - iteratOffset; + uint32_t innerLoopNum = addLength / softmaxLength; + uint32_t dataTail = addLength - innerLoopNum * softmaxLength; + loopDataNum = softmaxLength; + for (uint32_t innerLoopIdx = 0; innerLoopIdx < innerLoopNum; innerLoopIdx++) { + // Copy data from right + int64_t loopInnerOffset = dataTail + (innerLoopNum - 1 - innerLoopIdx) * softmaxLength; + DataCopyPad(cumSumInput1Local, softMaxGm[baseGmIdx_ + loopInnerOffset], + {1, static_cast(loopDataNum * sizeof(float)), 0, 0, 0}, {false, 0, 0, 0}); + DataCopyPad(cumSumInput2Local, softMaxGm[baseGmIdx_ + loopInnerOffset + iteratOffset], + {1, static_cast(loopDataNum * sizeof(float)), 0, 0, 0}, {false, 0, 0, 0}); + MTE2ToVSync(); + Add(cumSumInput1Local, cumSumInput1Local, cumSumInput2Local, loopDataNum); + VToMTE3Sync(); + DataCopyPad(softMaxGm[baseGmIdx_ + loopInnerOffset + iteratOffset], cumSumInput1Local, + {1, static_cast(loopDataNum * sizeof(float)), 0, 0, 0}); + MTE3ToMTE2Sync(); + } + if (dataTail > 0) { + loopDataNum = dataTail; + DataCopyPad(cumSumInput1Local, softMaxGm[baseGmIdx_], + {1, static_cast(loopDataNum * sizeof(float)), 0, 0, 0}, {false, 0, 0, 0}); + DataCopyPad(cumSumInput2Local, softMaxGm[baseGmIdx_ + iteratOffset], + {1, static_cast(loopDataNum * sizeof(float)), 0, 0, 0}, {false, 0, 0, 0}); + MTE2ToVSync(); + Add(cumSumInput1Local, cumSumInput1Local, cumSumInput2Local, loopDataNum); + VToMTE3Sync(); + DataCopyPad(softMaxGm[baseGmIdx_ + iteratOffset], cumSumInput1Local, + {1, static_cast(loopDataNum * sizeof(float)), 0, 0, 0}); + MTE3ToMTE2Sync(); + } + } + MTE3ToVSync(); +} + +template +__aicore__ inline void ApplyTopPCustom::ReduceSumWithAddsAndExpImpl( + uint32_t loopDataNum) { + Adds(softMaxResLocal, softMaxLocalFp32, maxValue, loopDataNum); + PipeBarrier(); + Exp(softMaxResLocal, softMaxResLocal, loopDataNum); + PipeBarrier(); + ReduceSum(reduceLocal, softMaxResLocal, reduceLocal, loopDataNum); +} + +template +__aicore__ inline void ApplyTopPCustom::GetSoftmaxSum(uint32_t loopBatch) { + uint32_t loopDataNum = softmaxLength; + for (int32_t loopInner = 0; loopInner < lineSfLoopTimes; loopInner++) { + int64_t currentGmIdx = baseGmIdx_ + loopInner * softmaxLength; + if (loopInner == (lineSfLoopTimes - 1)) { + loopDataNum = softmaxLengthTail; + } + if constexpr (IsSameType::value) { + Duplicate(outInfLocal.template ReinterpretCast(), FLOAT32_NEG_INF, loopDataNum); + } else if constexpr (IsSameType::value) { + Duplicate(outInfLocal.template ReinterpretCast(), FLOAT16_NEG_INF, loopDataNum); + } else { + Duplicate(outInfLocal.template ReinterpretCast(), BF16_NEG_INF, loopDataNum); + } + VToMTE3Sync(); + DataCopyPad(mGmOut_[currentGmIdx], outInfLocal, + {1, static_cast(loopDataNum * sizeof(inputT)), 0, 0, 0}); + MTE3ToMTE2Sync(); + if constexpr (!IsSameType::value) { + DataCopyPad(softMaxLocal, mGmSortedValue_[currentGmIdx], + {1, static_cast(loopDataNum * sizeof(inputT)), 0, 0, 0}, + {false, 0, 0, 0}); + MTE2ToVSync(); + Cast(softMaxLocalFp32, softMaxLocal, RoundMode::CAST_NONE, loopDataNum); + PipeBarrier(); + } else { + DataCopyPad(softMaxLocalFp32, mGmSortedValue_[currentGmIdx], + {1, static_cast(loopDataNum * sizeof(inputT)), 0, 0, 0}, + {false, 0, 0, 0}); + MTE2ToVSync(); + } + + ReduceSumWithAddsAndExpImpl(loopDataNum); + VToSSync(); + // Sum up to obtain the sum of exp reduce for the first x loops in the row. + reduceSumValue += reduceLocal.GetValue(0); + SToVSync(); + } + reduceSumValueInvert = 1 / reduceSumValue; + SToVSync(); +} +} // namespace + +#endif // APPLY_TOP_P_CUSTOM_H_KERNEL \ No newline at end of file diff --git a/csrc/ascend/moe/causal_conv1d/CMakeLists.txt b/csrc/ascend/moe/causal_conv1d/CMakeLists.txt new file mode 100644 index 000000000..298ed5677 --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/causal_conv1d/op_host/CMakeLists.txt b/csrc/ascend/moe/causal_conv1d/op_host/CMakeLists.txt new file mode 100644 index 000000000..1a8798b3c --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_host/CMakeLists.txt @@ -0,0 +1,22 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + causal_conv1d_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME CausalConv1d + OPTIONS + --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE causal_conv1d ACLNNTYPE aclnn) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_def.cpp b/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_def.cpp new file mode 100644 index 000000000..400c3a538 --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_def.cpp @@ -0,0 +1,89 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file causal_conv1d_def.cpp + * \brief + */ +#include "register/op_def_registry.h" + +namespace ops { + +class CausalConv1d : public OpDef { +public: + explicit CausalConv1d(const char* name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("weight") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("bias") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("convStates") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("queryStartLoc") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32, ge::DT_INT64}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("cacheIndices") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32, ge::DT_INT64}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("initialStateMode") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_BOOL, ge::DT_INT32, ge::DT_INT64}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("numAcceptedTokens") + .ParamType(OPTIONAL) + .DataTypeList({ge::DT_INT32, ge::DT_INT64}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .FormatList({ge::FORMAT_ND}) + .AutoContiguous(); + + this->Attr("activationMode").AttrType(OPTIONAL).Int(0); + this->Attr("padSlotId").AttrType(OPTIONAL).Int(-1); + this->Attr("runMode").AttrType(OPTIONAL).Int(0); + + OpAICoreConfig aicoreConfig; + aicoreConfig.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(false) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("coreType.value", "AiCore"); + this->AICore().AddConfig("ascend910b", aicoreConfig); + this->AICore().AddConfig("ascend910_93", aicoreConfig); + } +}; +OP_ADD(CausalConv1d); + +} // namespace ops diff --git a/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_infershape.cpp b/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_infershape.cpp new file mode 100644 index 000000000..2f83b9f87 --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_infershape.cpp @@ -0,0 +1,42 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file causal_conv1d_infershape.cpp + * \brief + */ +#include "register/op_impl_registry.h" +#include "tiling_base/error_log.h" + +using namespace ge; + +namespace ops { +static constexpr int64_t IDX_0 = 0; + +static ge::graphStatus InferShapeCausalConv1d(gert::InferShapeContext* context) +{ + OP_LOGD(context->GetNodeName(), "Begin to do InferShapeCausalConv1d"); + + // get input shapes + const gert::Shape* xShape = context->GetInputShape(IDX_0); + OP_CHECK_NULL_WITH_CONTEXT(context, xShape); + + // get output shapes + gert::Shape* yShape = context->GetOutputShape(IDX_0); + OP_CHECK_NULL_WITH_CONTEXT(context, yShape); + *yShape = *xShape; + + OP_LOGD(context->GetNodeName(), "End to do InferShapeCausalConv1d"); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(CausalConv1d).InferShape(InferShapeCausalConv1d); +} // namespace ops diff --git a/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_tiling.cpp b/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_tiling.cpp new file mode 100644 index 000000000..3acb544b3 --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_tiling.cpp @@ -0,0 +1,167 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file causal_conv1d_tiling.cpp + */ + + #include "tiling_base/tiling_templates_registry.h" + #include "causal_conv1d_tiling_utils.h" + #include "causal_conv1d_tiling_planner.h" + #include "causal_conv1d_tiling_validation.h" + + namespace optiling { + + using namespace Ops::Transformer::OpTiling; + using namespace causal_conv1d_host; + + static ge::graphStatus CausalConv1dTilingFunc(gert::TilingContext *context) + { + uint64_t ubSize = 0; + uint32_t coreNum = 0; + OP_CHECK_IF(GetPlatformInfo(context, ubSize, coreNum) != ge::GRAPH_SUCCESS, + OP_LOGE(context, "GetPlatformInfo error"), return ge::GRAPH_FAILED); + + CausalConv1dTilingData *tiling = context->GetTilingData(); + OP_CHECK_NULL_WITH_CONTEXT(context, tiling); + OP_CHECK_IF(memset_s(tiling, sizeof(CausalConv1dTilingData), 0, sizeof(CausalConv1dTilingData)) != EOK, + OP_LOGE(context, "set tiling data error"), return ge::GRAPH_FAILED); + + CausalConv1dAttrInfo attrInfo; + OP_CHECK_IF(GetAttrsInfo(context, attrInfo) != ge::GRAPH_SUCCESS, OP_LOGE(context, "GetAttrsInfo error"), + return ge::GRAPH_FAILED); + bool hasBias = false; + OP_CHECK_IF(GetShapeDtypeInfo(context, attrInfo, *tiling, hasBias) != ge::GRAPH_SUCCESS, + OP_LOGE(context, "GetShapeDtypeInfo error"), return ge::GRAPH_FAILED); + + const int64_t &dim = tiling->dim; + const int64_t &batch = tiling->batch; + OP_CHECK_IF(dim <= 0 || batch <= 0, OP_LOGE(context, "dim/batch must be positive"), return ge::GRAPH_FAILED); + + const uint32_t runModeKey = static_cast(attrInfo.runMode); + const bool &isFn = (runModeKey == CAUSAL_CONV1D_TPL_RUN_MODE_FN); + const bool &hasActivation = (attrInfo.activationMode != 0); + const char *plannerModeTag = "update"; + DimTileChoice baseDimChoice; + FnExecutionPlan fnExecutionPlan = FN_EXECUTION_PLAN_INVALID; + FnHostPlan fnHostPlan; + const int64_t *qslData = nullptr; + + if (isFn) { + fnHostPlan = ChooseFnHostPlan(context, *tiling, ubSize, coreNum); + plannerModeTag = GetFnTilingCaseName(fnHostPlan.caseKind); + baseDimChoice = fnHostPlan.baseDimChoice; + fnExecutionPlan = fnHostPlan.executionPlan; + } else { + baseDimChoice = ChooseCanonicalUpdateBaseDimChoice(context, tiling->batch, tiling->dim, coreNum); + } + + OP_CHECK_IF(baseDimChoice.baseDim <= 0 || baseDimChoice.baseDimCnt <= 0 || baseDimChoice.gridSize <= 0, + OP_LOGE(context, "invalid dim tile size selection"), return ge::GRAPH_FAILED); + + int64_t effectiveGridSize = baseDimChoice.gridSize; + + if (isFn) { + OP_CHECK_IF(fnHostPlan.caseKind == FN_TILING_CASE_INVALID || fnExecutionPlan == FN_EXECUTION_PLAN_INVALID || + !fnHostPlan.tokenBlockChoice.enabled || fnHostPlan.tokenBlockChoice.tokenBlockSize <= 0 || + fnHostPlan.tokenBlockChoice.tokenBlockCnt <= 0 || fnHostPlan.tokenBlockChoice.gridSize <= 0 || + fnHostPlan.tokenCoreMapping.tokenCoreBudget <= 0 || fnHostPlan.tokenCoreMapping.blockDim <= 0, + OP_LOGE(context, "runMode=0 must resolve a valid unified token tiling plan"), + return ge::GRAPH_FAILED); + + tiling->tokenBlockSize = fnHostPlan.tokenBlockChoice.tokenBlockSize; + tiling->tokenBlockCnt = fnHostPlan.tokenBlockChoice.tokenBlockCnt; + effectiveGridSize = fnHostPlan.tokenBlockChoice.gridSize; + if (tiling->inputMode == 0) { + fnHostPlan.tokenSeqRangePlan = + BuildFnTokenSeqRangePlan(qslData, tiling->batch, tiling->tokenBlockSize, tiling->tokenBlockCnt); + if (fnHostPlan.tokenSeqRangePlan.enabled) { + tiling->hasExplicitTokenSeqRanges = 1; + tiling->explicitTokenSeqRangeCount = fnHostPlan.tokenSeqRangePlan.rangeCount; + for (int64_t i = 0; i < fnHostPlan.tokenSeqRangePlan.rangeCount; ++i) { + tiling->tokenTileStartSeq[i] = fnHostPlan.tokenSeqRangePlan.tokenTileStartSeq[i]; + tiling->tokenTileEndSeq[i] = fnHostPlan.tokenSeqRangePlan.tokenTileEndSeq[i]; + } + } else if (qslData != nullptr && tiling->tokenBlockCnt > MAX_FN_TOKEN_SEQ_RANGE_COUNT) { + OP_LOGD(context, + "FnTokenSeqRanges disabled: tokenBlockCnt[%ld] exceeds fixed tiling capacity[%ld].", + tiling->tokenBlockCnt, MAX_FN_TOKEN_SEQ_RANGE_COUNT); + } + } + OP_LOGD(context, + "FnHostPlan(case=%s): inputMode[%ld], dim[%ld], cuSeqlen[%ld], baseDim[%ld], baseDimCnt[%ld], " + "tokenCoreBudget[%ld], tokenBlockSize[%ld], tokenBlockCnt[%ld], tokenBlocksPerCore[%ld], " + "tokenCoreTailCnt[%ld], explicitSeqRanges[%ld], baseGrid[%ld], phase1Grid[%ld], mappedBlockDim[%ld].", + plannerModeTag, tiling->inputMode, tiling->dim, tiling->cuSeqlen, baseDimChoice.baseDim, + baseDimChoice.baseDimCnt, fnHostPlan.tokenCoreMapping.tokenCoreBudget, + fnHostPlan.tokenBlockChoice.tokenBlockSize, fnHostPlan.tokenBlockChoice.tokenBlockCnt, + fnHostPlan.tokenCoreMapping.tokenBlocksPerCore, fnHostPlan.tokenCoreMapping.tokenCoreTailCnt, + tiling->hasExplicitTokenSeqRanges, + baseDimChoice.gridSize, fnHostPlan.tokenBlockChoice.gridSize, fnHostPlan.tokenCoreMapping.blockDim); + } + + uint32_t blockDim = + (effectiveGridSize < static_cast(coreNum)) ? static_cast(effectiveGridSize) : coreNum; + if (isFn) { + const int64_t mappedBlockDim = (effectiveGridSize < fnHostPlan.tokenCoreMapping.blockDim) ? effectiveGridSize : fnHostPlan.tokenCoreMapping.blockDim; + OP_CHECK_IF(mappedBlockDim <= 0, OP_LOGE(context, "invalid mapped blockDim for runMode=0"), + return ge::GRAPH_FAILED); + blockDim = static_cast(mappedBlockDim); + } + + OP_LOGD(context, + "Tiling result: mode[%s], batch[%ld], dim[%ld], baseDim[%ld], baseDimCnt[%ld], gridSize[%ld], " + "effectiveGrid[%ld], blockDim[%u], coreNum[%u], tokenTiling[%ld,%ld], hasActivation[%d], hasBias[%d], " + "fnPlan[%ld].", + plannerModeTag, batch, dim, baseDimChoice.baseDim, baseDimChoice.baseDimCnt, baseDimChoice.gridSize, + effectiveGridSize, blockDim, coreNum, tiling->tokenBlockSize, tiling->tokenBlockCnt, + static_cast(hasActivation), static_cast(hasBias), static_cast(fnExecutionPlan)); + + context->SetBlockDim(blockDim); + tiling->baseDim = baseDimChoice.baseDim; + tiling->baseDimCnt = baseDimChoice.baseDimCnt; + const uint32_t fnPlanKey = NormalizeFnPlanTilingKey(runModeKey, fnExecutionPlan); + const uint32_t widthKey = NormalizeWidthTilingKey(runModeKey, static_cast(tiling->width)); + if (isFn && tiling->hasInitialStateMode != 0) { + constexpr int64_t kDtypeSize = 2; + constexpr int64_t kSyncBytesPerBlock = 32; + const int64_t historyCount = (tiling->width - 1 > 0) ? tiling->width - 1 : 0; + const int64_t syncWorkspaceSize = static_cast(blockDim) * kSyncBytesPerBlock; + const int64_t snapshotWorkspaceSize = tiling->batch * historyCount * tiling->dim * kDtypeSize; + const int64_t workspaceSize = + ASCENDC_RESERVED_WORKSPACE_SIZE + syncWorkspaceSize + snapshotWorkspaceSize; + OP_CHECK_IF(SetWorkspaceSize(context, static_cast(workspaceSize)) != ge::GRAPH_SUCCESS, + OP_LOGE(context, "SetWorkspaceSize error"), return ge::GRAPH_FAILED); + OP_CHECK_IF(context->SetScheduleMode(1) != ge::GRAPH_SUCCESS, + OP_LOGE(context, "SetScheduleMode(1) error"), return ge::GRAPH_FAILED); + tiling->hasInitStateWorkspace = 1; + } else { + OP_CHECK_IF(SetWorkspaceSize(context, 0) != ge::GRAPH_SUCCESS, OP_LOGE(context, "SetWorkspaceSize error"), + return ge::GRAPH_FAILED); + tiling->hasInitStateWorkspace = 0; + } + + const uint64_t tilingKey = GET_TPL_TILING_KEY(runModeKey, widthKey, fnPlanKey); + context->SetTilingKey(tilingKey); + return ge::GRAPH_SUCCESS; + } + + static ge::graphStatus TilingParseForCausalConv1d(gert::TilingParseContext *context) + { + OP_LOGD(context, "Enter TilingParseForCausalConv1d."); + return ge::GRAPH_SUCCESS; + } + + IMPL_OP_OPTILING(CausalConv1d) + .Tiling(CausalConv1dTilingFunc) + .TilingParse(TilingParseForCausalConv1d); + + } diff --git a/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_tiling_planner.h b/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_tiling_planner.h new file mode 100644 index 000000000..3efefc35b --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_tiling_planner.h @@ -0,0 +1,312 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CAUSAL_CONV1D_TILING_PLANNER_H +#define CAUSAL_CONV1D_TILING_PLANNER_H + +#include "causal_conv1d_tiling_utils.h" +#include "../op_kernel/causal_conv1d_tiling_data.h" + +namespace optiling::causal_conv1d_host { + +using namespace Ops::Transformer::OpTiling; + +inline DimTileChoice ChooseCanonicalUpdateBaseDimChoice(gert::TilingContext *context, int64_t batch, int64_t dim, + uint32_t coreNum) +{ + const int64_t candidates[] = {4096, 2048, 1024, 512, 384, 192}; + + auto chooseOnce = [&](bool requireExactDiv) -> DimTileChoice { + DimTileChoice bestOver; + int64_t bestOverGap = std::numeric_limits::max(); + DimTileChoice bestUnder; + + for (int64_t baseDim : candidates) { + if (baseDim <= 0) { + continue; + } + if (requireExactDiv && (dim % baseDim != 0)) { + continue; + } + + const int64_t baseDimCnt = requireExactDiv ? (dim / baseDim) : CeilDivInt64(dim, baseDim); + const int64_t gridSize = batch * baseDimCnt; + if (gridSize <= 0) { + continue; + } + + OP_LOGD(context, + "DimTile(update) candidate[%s]: baseDim[%ld], baseDimCnt[%ld], gridSize[%ld], coreNum[%u].", + requireExactDiv ? "exact" : "tail", baseDim, baseDimCnt, gridSize, coreNum); + if (gridSize >= static_cast(coreNum)) { + const int64_t gap = gridSize - static_cast(coreNum); + if (gap < bestOverGap) { + // bestOver = {baseDim, baseDimCnt, gridSize}; + bestOver.baseDim = baseDim; + bestOver.baseDimCnt = baseDimCnt; + bestOver.gridSize = gridSize; + bestOverGap = gap; + } + } else if (gridSize > bestUnder.gridSize || + (gridSize == bestUnder.gridSize && baseDim < bestUnder.baseDim)) { + // bestUnder = {baseDim, baseDimCnt, gridSize}; + bestUnder.baseDim = baseDim; + bestUnder.baseDimCnt = baseDimCnt; + bestUnder.gridSize = gridSize; + } + } + + return (bestOver.baseDim != 0) ? bestOver : bestUnder; + }; + + DimTileChoice result = chooseOnce(true); + if (result.baseDim == 0) { + result = chooseOnce(false); + } + OP_LOGD(context, "DimTile(update) chosen: baseDim[%ld], baseDimCnt[%ld], gridSize[%ld].", result.baseDim, + result.baseDimCnt, result.gridSize); + return result; +} + +inline int64_t ResolveFnTokenCoreBudget(int64_t baseDimCnt, FnExecutionPlan fnExecutionPlan, uint32_t coreNum) +{ + if (baseDimCnt <= 0 || coreNum == 0 || fnExecutionPlan == FN_EXECUTION_PLAN_INVALID) { + return 0; + } + + int64_t tokenCoreBudget = static_cast(coreNum); + if (fnExecutionPlan == FN_EXECUTION_PLAN_CUTBSD) { + tokenCoreBudget = std::max(1, tokenCoreBudget / baseDimCnt); + } + return tokenCoreBudget; +} + +inline VarlenTokenTileChoice ChooseFnTokenBlockChoice(int64_t cuSeqlen, int64_t baseDimCnt, + FnExecutionPlan fnExecutionPlan, uint32_t coreNum); + +inline int64_t ComputeFnUbLimitedBaseDim(uint64_t ubSize) +{ + if (ubSize <= static_cast(FN_UB_RESERVED_BYTES)) { + return 0; + } + + const int64_t bytesPerElem = (RING_SLOT_CNT * BF16_FP16_ELEM_BYTES) + (FN_OUT_SLOT_CNT * BF16_FP16_ELEM_BYTES) + + (FN_CALC_FP32_SLOT_CNT * static_cast(sizeof(float))); + const int64_t budgetBytes = static_cast(ubSize) - FN_UB_RESERVED_BYTES; + const int64_t ubLimitedBaseDim = AlignDownInt64(budgetBytes / bytesPerElem, DIM_ALIGN_ELEMS); + return std::min(MAX_DIM_TILE_SIZE, ubLimitedBaseDim); +} + +inline DimTileChoice ChooseFnTokenFirstBaseDimChoice(int64_t dim) +{ + if (dim <= 0 || dim > MAX_DIM_TILE_SIZE) { + return {}; + } + DimTileChoice choice; + choice.baseDim = dim; + choice.baseDimCnt = 1; + choice.gridSize = 1; + return choice; +} + +inline DimTileChoice ChooseFnTokenDimCoSplitBaseDimChoice(gert::TilingContext *context, int64_t dim, uint64_t ubSize, + uint32_t coreNum) +{ + if (dim <= 0) { + return {}; + } + + const int64_t ubLimitedBaseDim = ComputeFnUbLimitedBaseDim(ubSize); + if (ubLimitedBaseDim <= 0) { + OP_LOGD(context, "FnDimCoSplit: UB budget is too small to form a valid baseDim."); + return {}; + } + + DimTileChoice result; + result.baseDim = ubLimitedBaseDim; + result.baseDimCnt = CeilDivInt64(dim, result.baseDim); + result.gridSize = result.baseDimCnt; + + if (coreNum == 0 || result.baseDimCnt <= 1 || result.baseDimCnt >= static_cast(coreNum) || + (coreNum % result.baseDimCnt == 0)) { + OP_LOGD(context, + "FnDimCoSplit: dim[%ld], ubLimitedBaseDim[%ld], baseDimCnt[%ld], coreNum[%u], adjusted[%d].", dim, + result.baseDim, result.baseDimCnt, coreNum, 0); + return result; + } + + int64_t adjustedBaseDimCnt = result.baseDimCnt; + while (adjustedBaseDimCnt < static_cast(coreNum) && (coreNum % adjustedBaseDimCnt != 0)) { + ++adjustedBaseDimCnt; + } + + if (adjustedBaseDimCnt >= static_cast(coreNum)) { + OP_LOGD(context, + "FnDimCoSplit: keep baseDimCnt[%ld] because no divisible adjustment exists under coreNum[%u].", + result.baseDimCnt, coreNum); + return result; + } + + const int64_t adjustedBaseDim = AlignUpInt64(CeilDivInt64(dim, adjustedBaseDimCnt), DIM_ALIGN_ELEMS); + if (adjustedBaseDim <= 0 || adjustedBaseDim > ubLimitedBaseDim || adjustedBaseDim > MAX_DIM_TILE_SIZE) { + OP_LOGD(context, + "FnDimCoSplit: rejected adjusted baseDim[%ld] with baseDimCnt[%ld], ubLimitedBaseDim[%ld].", + adjustedBaseDim, adjustedBaseDimCnt, ubLimitedBaseDim); + return result; + } + + result.baseDim = adjustedBaseDim; + result.baseDimCnt = CeilDivInt64(dim, result.baseDim); + result.gridSize = result.baseDimCnt; + OP_LOGD(context, + "FnDimCoSplit: dim[%ld], ubLimitedBaseDim[%ld], adjustedBaseDim[%ld], baseDimCnt[%ld], coreNum[%u].", + dim, ubLimitedBaseDim, result.baseDim, result.baseDimCnt, coreNum); + return result; +} + +inline TokenCoreMappingChoice BuildFnTokenCoreMappingChoice(int64_t tokenBlockCnt, int64_t baseDimCnt, + FnExecutionPlan fnExecutionPlan, uint32_t coreNum) +{ + TokenCoreMappingChoice mapping; + mapping.tokenCoreBudget = ResolveFnTokenCoreBudget(baseDimCnt, fnExecutionPlan, coreNum); + if (tokenBlockCnt <= 0 || mapping.tokenCoreBudget <= 0 || baseDimCnt <= 0) { + return mapping; + } + + mapping.tokenBlocksPerCore = CeilDivInt64(tokenBlockCnt, mapping.tokenCoreBudget); + mapping.tokenCoreTailCnt = + tokenBlockCnt - (std::max(0, mapping.tokenBlocksPerCore - 1) * mapping.tokenCoreBudget); + if (mapping.tokenCoreTailCnt <= 0) { + mapping.tokenCoreTailCnt = mapping.tokenCoreBudget; + } + mapping.blockDim = mapping.tokenCoreBudget * baseDimCnt; + return mapping; +} + +inline FnTokenSeqRangePlan BuildFnTokenSeqRangePlan(const int64_t *qslData, int64_t batch, int64_t tokenBlockSize, + int64_t tokenBlockCnt) +{ + FnTokenSeqRangePlan plan; + if (qslData == nullptr || batch <= 0 || tokenBlockSize <= 0 || tokenBlockCnt <= 0 || + tokenBlockCnt > MAX_FN_TOKEN_SEQ_RANGE_COUNT) { + return plan; + } + + plan.enabled = true; + plan.rangeCount = tokenBlockCnt; + int64_t seq = 0; + for (int64_t tokenTileId = 0; tokenTileId < tokenBlockCnt; ++tokenTileId) { + const int64_t tokenStart = tokenTileId * tokenBlockSize; + const int64_t tokenEnd = tokenStart + tokenBlockSize; + + while (seq < batch && qslData[seq + 1] <= tokenStart) { + ++seq; + } + + int64_t endSeq = seq; + while (endSeq < batch && qslData[endSeq] < tokenEnd) { + ++endSeq; + } + + plan.tokenTileStartSeq[tokenTileId] = seq; + plan.tokenTileEndSeq[tokenTileId] = endSeq; + } + return plan; +} + +inline VarlenTokenTileChoice ChooseUnifiedFnTokenBlockPlan(gert::TilingContext *context, + const CausalConv1dTilingData &tiling, + const DimTileChoice &baseDimChoice, + FnExecutionPlan fnExecutionPlan, + uint32_t coreNum) +{ + VarlenTokenTileChoice tokenBlockChoice; + if ((tiling.inputMode != 0 && tiling.inputMode != 1) || tiling.batch <= 0 || tiling.cuSeqlen <= 0 || + baseDimChoice.baseDimCnt <= 0 || coreNum == 0 || fnExecutionPlan == FN_EXECUTION_PLAN_INVALID) { + return tokenBlockChoice; + } + if (tiling.hasNumAcceptedTokens != 0) { + OP_LOGD(context, "Varlen token tiling disabled: speculative decode still uses the existing seq mapping."); + return tokenBlockChoice; + } + + tokenBlockChoice = ChooseFnTokenBlockChoice(tiling.cuSeqlen, baseDimChoice.baseDimCnt, fnExecutionPlan, coreNum); + + OP_LOGD(context, + "FnTokenTile(plan=%ld): cuSeqlen[%ld], baseDimCnt[%ld], tokenBlockSize[%ld], " + "tokenBlockCnt[%ld], gridSize[%ld].", + static_cast(fnExecutionPlan), tiling.cuSeqlen, baseDimChoice.baseDimCnt, + tokenBlockChoice.tokenBlockSize, tokenBlockChoice.tokenBlockCnt, tokenBlockChoice.gridSize); + return tokenBlockChoice; +} + +inline VarlenTokenTileChoice ChooseFnTokenBlockChoice(int64_t cuSeqlen, int64_t baseDimCnt, + FnExecutionPlan fnExecutionPlan, uint32_t coreNum) +{ + VarlenTokenTileChoice tokenBlockChoice; + const int64_t tokenCoreBudget = ResolveFnTokenCoreBudget(baseDimCnt, fnExecutionPlan, coreNum); + if (cuSeqlen <= 0 || tokenCoreBudget <= 0) { + return tokenBlockChoice; + } + + tokenBlockChoice.enabled = true; + const int64_t idealBlockSize = CeilDivInt64(cuSeqlen, tokenCoreBudget); + tokenBlockChoice.tokenBlockSize = (idealBlockSize > 0) ? idealBlockSize : 1; + tokenBlockChoice.tokenBlockCnt = CeilDivInt64(cuSeqlen, tokenBlockChoice.tokenBlockSize); + tokenBlockChoice.gridSize = tokenBlockChoice.tokenBlockCnt * baseDimCnt; + return tokenBlockChoice; +} + +inline FnHostPlan ChooseFnHostPlan(gert::TilingContext *context, const CausalConv1dTilingData &tiling, uint64_t ubSize, + uint32_t coreNum) +{ + FnHostPlan plan; + if ((tiling.inputMode != 0 && tiling.inputMode != 1) || tiling.batch <= 0 || tiling.cuSeqlen <= 0 || + tiling.dim <= 0 || coreNum == 0) { + return plan; + } + + if (tiling.dim <= MAX_DIM_TILE_SIZE) { + plan.caseKind = FN_TILING_CASE_TOKEN_FIRST; + plan.executionPlan = FN_EXECUTION_PLAN_CUTBS; + plan.baseDimChoice = ChooseFnTokenFirstBaseDimChoice(tiling.dim); + } else { + plan.caseKind = FN_TILING_CASE_TOKEN_DIM_CO_SPLIT; + plan.executionPlan = FN_EXECUTION_PLAN_CUTBSD; + plan.baseDimChoice = ChooseFnTokenDimCoSplitBaseDimChoice(context, tiling.dim, ubSize, coreNum); + } + + if (plan.baseDimChoice.baseDim <= 0 || plan.baseDimChoice.baseDimCnt <= 0) { + return {}; + } + + plan.baseDimChoice.gridSize = tiling.batch * plan.baseDimChoice.baseDimCnt; + plan.tokenBlockChoice = + ChooseUnifiedFnTokenBlockPlan(context, tiling, plan.baseDimChoice, plan.executionPlan, coreNum); + if (!plan.tokenBlockChoice.enabled || plan.tokenBlockChoice.tokenBlockSize <= 0 || + plan.tokenBlockChoice.tokenBlockCnt <= 0 || plan.tokenBlockChoice.gridSize <= 0) { + return {}; + } + + plan.tokenCoreMapping = BuildFnTokenCoreMappingChoice(plan.tokenBlockChoice.tokenBlockCnt, + plan.baseDimChoice.baseDimCnt, plan.executionPlan, coreNum); + if (plan.tokenCoreMapping.tokenCoreBudget <= 0 || plan.tokenCoreMapping.blockDim <= 0) { + return {}; + } + if (plan.tokenCoreMapping.blockDim > static_cast(coreNum)) { + plan.tokenCoreMapping.blockDim = static_cast(coreNum); + } + return plan; +} + +} // namespace optiling::causal_conv1d_host + +#endif // CAUSAL_CONV1D_TILING_PLANNER_H diff --git a/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_tiling_utils.h b/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_tiling_utils.h new file mode 100644 index 000000000..55adf540b --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_tiling_utils.h @@ -0,0 +1,165 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CAUSAL_CONV1D_TILING_UTILS_H +#define CAUSAL_CONV1D_TILING_UTILS_H + +#include "tiling_base/tiling_util.h" +#include "../op_kernel/causal_conv1d_tiling_key.h" + +namespace optiling::causal_conv1d_host { + +constexpr uint32_t X_INDEX = 0; +constexpr uint32_t WEIGHT_INDEX = 1; +constexpr uint32_t BIAS_INDEX = 2; +constexpr uint32_t CONV_STATES_INDEX = 3; +constexpr uint32_t QUERY_START_LOC_INDEX = 4; +constexpr uint32_t CACHE_INDICES_INDEX = 5; +constexpr uint32_t INITIAL_STATE_MODE_INDEX = 6; +constexpr uint32_t NUM_ACCEPTED_TOKENS_INDEX = 7; + +constexpr int32_t ATTR_ACTIVATION_MODE_INDEX = 0; +constexpr int32_t ATTR_PAD_SLOT_ID_INDEX = 1; +constexpr int32_t ATTR_RUN_MODE_INDEX = 2; +constexpr int64_t ASCENDC_RESERVED_WORKSPACE_SIZE = 16 * 1024 * 1024; + +struct CausalConv1dCompileInfo { + uint64_t ubSize = 0; + uint32_t coreNum = 0; +}; + +struct CausalConv1dAttrInfo { + int64_t activationMode = 0; + int64_t padSlotId = -1; + int64_t runMode = 0; +}; + +struct DimTileChoice { + int64_t baseDim = 0; + int64_t baseDimCnt = 0; + int64_t gridSize = 0; +}; + +struct VarlenTokenTileChoice { + bool enabled = false; + int64_t tokenBlockSize = 0; + int64_t tokenBlockCnt = 0; + int64_t gridSize = 0; +}; + +enum FnTilingCaseKind : int64_t { + FN_TILING_CASE_INVALID = 0, + FN_TILING_CASE_TOKEN_FIRST = 1, + FN_TILING_CASE_TOKEN_DIM_CO_SPLIT = 2, +}; + +struct TokenCoreMappingChoice { + int64_t tokenCoreBudget = 0; + int64_t tokenBlocksPerCore = 0; + int64_t tokenCoreTailCnt = 0; + int64_t blockDim = 0; +}; + +constexpr int64_t MAX_FN_TOKEN_SEQ_RANGE_COUNT = 128; + +struct FnTokenSeqRangePlan { + bool enabled = false; + int64_t rangeCount = 0; + int64_t tokenTileStartSeq[MAX_FN_TOKEN_SEQ_RANGE_COUNT] = {}; + int64_t tokenTileEndSeq[MAX_FN_TOKEN_SEQ_RANGE_COUNT] = {}; +}; + +struct FnHostPlan { + FnTilingCaseKind caseKind = FN_TILING_CASE_INVALID; + FnExecutionPlan executionPlan = FN_EXECUTION_PLAN_INVALID; + DimTileChoice baseDimChoice; + VarlenTokenTileChoice tokenBlockChoice; + TokenCoreMappingChoice tokenCoreMapping; + FnTokenSeqRangePlan tokenSeqRangePlan; +}; + +constexpr int64_t DIM_ALIGN_BYTES = 32; +constexpr int64_t BF16_FP16_ELEM_BYTES = 2; +constexpr int64_t DIM_ALIGN_ELEMS = DIM_ALIGN_BYTES / BF16_FP16_ELEM_BYTES; +constexpr int64_t MAX_DIM_TILE_SIZE = 4096; +constexpr int64_t FN_UB_RESERVED_BYTES = 512; +constexpr int64_t RING_SLOT_CNT = 5; +constexpr int64_t FN_OUT_SLOT_CNT = 2; +constexpr int64_t FN_CALC_FP32_SLOT_CNT = 8; + +inline uint32_t NormalizeFnPlanTilingKey(uint32_t runModeKey, FnExecutionPlan fnExecutionPlan) +{ + if (runModeKey != CAUSAL_CONV1D_TPL_RUN_MODE_FN) { + return CAUSAL_CONV1D_TPL_FN_PLAN_INVALID; + } + switch (fnExecutionPlan) { + case FN_EXECUTION_PLAN_CUTBS: + return CAUSAL_CONV1D_TPL_FN_PLAN_CUTBS; + case FN_EXECUTION_PLAN_CUTBSD: + return CAUSAL_CONV1D_TPL_FN_PLAN_CUTBSD; + default: + return CAUSAL_CONV1D_TPL_FN_PLAN_INVALID; + } +} + +inline uint32_t NormalizeWidthTilingKey(uint32_t runModeKey, int32_t width) +{ + if (runModeKey != CAUSAL_CONV1D_TPL_RUN_MODE_FN) { + return CAUSAL_CONV1D_TPL_WIDTH_RUNTIME; + } + switch (width) { + case 2: + return CAUSAL_CONV1D_TPL_WIDTH_2; + case 3: + return CAUSAL_CONV1D_TPL_WIDTH_3; + case 4: + return CAUSAL_CONV1D_TPL_WIDTH_4; + default: + return CAUSAL_CONV1D_TPL_WIDTH_RUNTIME; + } +} + +inline int64_t CeilDivInt64(int64_t x, int64_t y) +{ + return (x + y - 1) / y; +} + +inline int64_t AlignDownInt64(int64_t value, int64_t align) +{ + if (align <= 0 || value <= 0) { + return 0; + } + return (value / align) * align; +} + +inline int64_t AlignUpInt64(int64_t value, int64_t align) +{ + if (align <= 0 || value <= 0) { + return 0; + } + return CeilDivInt64(value, align) * align; +} + +inline const char *GetFnTilingCaseName(FnTilingCaseKind caseKind) +{ + switch (caseKind) { + case FN_TILING_CASE_TOKEN_FIRST: + return "token_first"; + case FN_TILING_CASE_TOKEN_DIM_CO_SPLIT: + return "token_dim_co_split"; + default: + return "invalid"; + } +} + +} // namespace optiling::causal_conv1d_host + +#endif // CAUSAL_CONV1D_TILING_UTILS_H diff --git a/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_tiling_validation.h b/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_tiling_validation.h new file mode 100644 index 000000000..cec825332 --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_host/causal_conv1d_tiling_validation.h @@ -0,0 +1,385 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + + #ifndef CAUSAL_CONV1D_TILING_VALIDATION_H + #define CAUSAL_CONV1D_TILING_VALIDATION_H + + #include "tiling_base/tiling_util.h" + #include "causal_conv1d_tiling_utils.h" + #include "../op_kernel/causal_conv1d_tiling_data.h" + + namespace optiling::causal_conv1d_host { + + using namespace Ops::Transformer::OpTiling; + + inline ge::graphStatus GetPlatformInfo(gert::TilingContext *context, uint64_t &ubSize, uint32_t &coreNum) + { + fe::PlatFormInfos *platformInfoPtr = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfoPtr); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + coreNum = ascendcPlatform.GetCoreNumAiv(); + OP_CHECK_IF(coreNum == 0, OP_LOGE(context, "coreNum is 0"), return ge::GRAPH_FAILED); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); + OP_CHECK_IF(ubSize == 0, OP_LOGE(context, "ubSize is 0"), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; + } + + inline ge::graphStatus SetWorkspaceSize(gert::TilingContext *context, size_t workspaceSize) + { + size_t *currentWorkspace = context->GetWorkspaceSizes(1); + OP_CHECK_NULL_WITH_CONTEXT(context, currentWorkspace); + currentWorkspace[0] = workspaceSize; + return ge::GRAPH_SUCCESS; + } + + inline ge::graphStatus GetAttrsInfo(gert::TilingContext *context, CausalConv1dAttrInfo &attrInfo) + { + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + + const int64_t *activationModePtr = attrs->GetAttrPointer(ATTR_ACTIVATION_MODE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, activationModePtr); + attrInfo.activationMode = *activationModePtr; + OP_CHECK_IF(attrInfo.activationMode != 0 && attrInfo.activationMode != 1, + OP_LOGE(context, "activationMode only supports 0/1"), + return ge::GRAPH_FAILED); + + const int64_t *padSlotIdPtr = attrs->GetAttrPointer(ATTR_PAD_SLOT_ID_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, padSlotIdPtr); + attrInfo.padSlotId = *padSlotIdPtr; + + const int64_t *runModePtr = attrs->GetAttrPointer(ATTR_RUN_MODE_INDEX); + attrInfo.runMode = (runModePtr == nullptr) ? 0 : *runModePtr; + OP_CHECK_IF(attrInfo.runMode != 0 && attrInfo.runMode != 1, OP_LOGE(context, "runMode only supports 0/1"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; + } + + inline ge::graphStatus ValidateAlignedDim(gert::TilingContext *context, int64_t dim) + { + OP_CHECK_IF(dim % DIM_ALIGN_ELEMS != 0, + OP_LOGE(context, + "dim must satisfy dim %% %ld == 0 for causal_conv1d; " + "x/weight/convStates last dimension and bias length must all use the same aligned dim, " + "got dim=%ld.", + DIM_ALIGN_ELEMS, dim), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; + } + + inline ge::graphStatus GetShapeDtypeInfo(gert::TilingContext *context, const CausalConv1dAttrInfo &attrInfo, + CausalConv1dTilingData &tiling, bool &hasBias) + { + const bool isDecodeMode = (attrInfo.runMode == 1); + tiling.activationMode = attrInfo.activationMode; + tiling.padSlotId = attrInfo.padSlotId; + + auto xShapePtr = context->GetInputShape(X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, xShapePtr); + auto xShape = EnsureNotScalar(xShapePtr->GetStorageShape()); + + int64_t dim = 0; + int64_t cuSeqlen = 0; + int64_t seqLen = 0; + int64_t batch = 0; + int64_t inputMode = 0; + + if (xShape.GetDimNum() == 2) { + if (isDecodeMode) { + inputMode = 2; + batch = xShape.GetDim(0); + dim = xShape.GetDim(1); + seqLen = 1; + cuSeqlen = batch; + OP_CHECK_IF(batch <= 0 || dim <= 0, OP_LOGE(context, "invalid x shape for 2D decode mode"), + return ge::GRAPH_FAILED); + } else { + inputMode = 0; + cuSeqlen = xShape.GetDim(0); + dim = xShape.GetDim(1); + seqLen = 0; + OP_CHECK_IF(dim <= 0 || cuSeqlen < 0, OP_LOGE(context, "invalid x shape for 2D varlen mode"), + return ge::GRAPH_FAILED); + } + } else if (xShape.GetDimNum() == 3) { + inputMode = 1; + batch = xShape.GetDim(0); + seqLen = xShape.GetDim(1); + dim = xShape.GetDim(2); + cuSeqlen = batch * seqLen; + OP_CHECK_IF(batch <= 0 || dim <= 0 || seqLen <= 0, OP_LOGE(context, "invalid x shape for 3D batch mode"), + return ge::GRAPH_FAILED); + } else { + OP_LOGE(context, "x must be 2D (cu_seqlen, dim) or 3D (batch, seqlen, dim)"); + return ge::GRAPH_FAILED; + } + OP_CHECK_IF(ValidateAlignedDim(context, dim) != ge::GRAPH_SUCCESS, + OP_LOGE(context, "dim alignment validation failed"), + return ge::GRAPH_FAILED); + + auto wShapePtr = context->GetInputShape(WEIGHT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, wShapePtr); + auto wShape = EnsureNotScalar(wShapePtr->GetStorageShape()); + OP_CHECK_IF(wShape.GetDimNum() != 2, OP_LOGE(context, "weight must be 2D: (width, dim)"), return ge::GRAPH_FAILED); + const int64_t width = wShape.GetDim(0); + const int64_t wDim = wShape.GetDim(1); + OP_CHECK_IF(wDim != dim, OP_LOGE(context, "weight.shape[1] must equal dim"), return ge::GRAPH_FAILED); + OP_CHECK_IF(width < 2 || width > 4, OP_LOGE(context, "Only support width in [2,4] now, actually is %ld.", width), + return ge::GRAPH_FAILED); + + auto sShapePtr = context->GetInputShape(CONV_STATES_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, sShapePtr); + auto sShape = EnsureNotScalar(sShapePtr->GetStorageShape()); + OP_CHECK_IF(sShape.GetDimNum() != 3, OP_LOGE(context, "convStates must be 3D: (num_cache_lines, state_len, dim)"), + return ge::GRAPH_FAILED); + const int64_t numCacheLines = sShape.GetDim(0); + const int64_t stateLen = sShape.GetDim(1); + const int64_t sDim = sShape.GetDim(2); + OP_CHECK_IF(numCacheLines <= 0, OP_LOGE(context, "convStates.shape[0] (num_cache_lines) must be > 0"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(sDim != dim, OP_LOGE(context, "convStates.shape[2] must equal dim"), return ge::GRAPH_FAILED); + OP_CHECK_IF(stateLen < (width - 1), OP_LOGE(context, "convStates.shape[1] must be >= width-1"), + return ge::GRAPH_FAILED); + + auto qslShapePtr = context->GetOptionalInputShape(QUERY_START_LOC_INDEX); + const gert::CompileTimeTensorDesc *qslDesc = context->GetOptionalInputDesc(QUERY_START_LOC_INDEX); + bool qslAbsent = true; + int64_t qslSize = 0; + if (qslShapePtr != nullptr) { + const auto qslStorageShape = qslShapePtr->GetStorageShape(); + const int64_t qslDimNum = qslStorageShape.GetDimNum(); + qslAbsent = (qslDimNum == 0) || (qslDimNum == 1 && qslStorageShape.GetDim(0) <= 0); + if (!qslAbsent) { + auto qslShape = EnsureNotScalar(qslStorageShape); + OP_CHECK_IF(qslShape.GetDimNum() != 1, OP_LOGE(context, "queryStartLoc must be 1D"), + return ge::GRAPH_FAILED); + qslSize = qslShape.GetDim(0); + OP_CHECK_IF(qslSize < 1, OP_LOGE(context, "queryStartLoc.size must be >= 1"), return ge::GRAPH_FAILED); + OP_CHECK_NULL_WITH_CONTEXT(context, qslDesc); + const ge::DataType qslDtype = qslDesc->GetDataType(); + OP_CHECK_IF(qslDtype != ge::DT_INT32 && qslDtype != ge::DT_INT64, + OP_LOGE(context, "queryStartLoc dtype must be int32 or int64"), return ge::GRAPH_FAILED); + } + } + + if (qslAbsent) { + OP_CHECK_IF(inputMode == 0, OP_LOGE(context, "queryStartLoc is required in 2D varlen mode (inputMode=0)"), + return ge::GRAPH_FAILED); + qslSize = batch + 1; + } + tiling.hasQueryStartLoc = qslAbsent ? 0 : 1; + + OP_CHECK_IF(cuSeqlen > static_cast(std::numeric_limits::max()), + OP_LOGE(context, "cuSeqlen is too large for int32 indexing, got %ld", cuSeqlen), + return ge::GRAPH_FAILED); + + if (!qslAbsent && isDecodeMode && inputMode == 2) { + const int64_t batchFromQsl = qslSize - 1; + if (batchFromQsl != batch) { + inputMode = 0; + cuSeqlen = xShape.GetDim(0); + batch = batchFromQsl; + seqLen = 0; + OP_CHECK_IF(dim <= 0 || cuSeqlen < 0 || batch < 0, + OP_LOGE(context, "invalid x/queryStartLoc shapes for 2D varlen decode mode"), + return ge::GRAPH_FAILED); + } + } + + if (inputMode == 0) { + batch = qslSize - 1; + } + if (!qslAbsent && (inputMode == 1 || inputMode == 2)) { + OP_CHECK_IF(qslSize != batch + 1, OP_LOGE(context, "queryStartLoc.size must equal batch + 1"), + return ge::GRAPH_FAILED); + } + if (isDecodeMode) { + const int64_t decodeSeqLen = (inputMode == 1) ? seqLen : 1; + OP_CHECK_IF(decodeSeqLen < 1, OP_LOGE(context, "decode mode requires seqlen >= 1, actual is %ld", decodeSeqLen), + return ge::GRAPH_FAILED); + } + + tiling.hasCacheIndices = 0; + tiling.cacheIndicesStride = 1; + bool ciAbsent = true; + auto ciShapePtr = context->GetOptionalInputShape(CACHE_INDICES_INDEX); + if (ciShapePtr != nullptr) { + const auto ciStorageShape = ciShapePtr->GetStorageShape(); + const int64_t ciDimNum = ciStorageShape.GetDimNum(); + ciAbsent = (ciDimNum == 0) || (ciDimNum == 1 && ciStorageShape.GetDim(0) <= 0); + if (!ciAbsent) { + auto ciShape = EnsureNotScalar(ciStorageShape); + // Spec decode passes cache indices as [batch, num_spec + 1]; + // kernels read the first column with cacheIndicesStride. + OP_CHECK_IF(ciShape.GetDimNum() != 1 && ciShape.GetDimNum() != 2, + OP_LOGE(context, "cacheIndices must be 1D or 2D"), return ge::GRAPH_FAILED); + OP_CHECK_IF(ciShape.GetDim(0) != batch, OP_LOGE(context, "cacheIndices first dim must equal batch"), + return ge::GRAPH_FAILED); + if (ciShape.GetDimNum() == 2) { + OP_CHECK_IF(ciShape.GetDim(1) <= 0, OP_LOGE(context, "cacheIndices second dim must be positive"), + return ge::GRAPH_FAILED); + tiling.cacheIndicesStride = ciShape.GetDim(1); + } + tiling.hasCacheIndices = 1; + } + } + if (ciAbsent) { + OP_CHECK_IF(numCacheLines < batch, + OP_LOGE(context, + "cacheIndices is absent, requires convStates.shape[0] (num_cache_lines) >= batch for " + "identity mapping, got num_cache_lines=%ld batch=%ld", + numCacheLines, batch), + return ge::GRAPH_FAILED); + } + + tiling.hasInitialStateMode = 0; + auto ismShapePtr = context->GetOptionalInputShape(INITIAL_STATE_MODE_INDEX); + if (ismShapePtr != nullptr) { + const auto ismStorageShape = ismShapePtr->GetStorageShape(); + const int64_t ismDimNum = ismStorageShape.GetDimNum(); + const bool ismAbsent = (ismDimNum == 0) || (ismDimNum == 1 && ismStorageShape.GetDim(0) <= 0); + if (!ismAbsent) { + OP_CHECK_IF(isDecodeMode, + OP_LOGE(context, "initialStateMode is only supported in runMode=0 (fn/prefill)"), + return ge::GRAPH_FAILED); + auto ismShape = EnsureNotScalar(ismStorageShape); + OP_CHECK_IF(ismShape.GetDimNum() != 1, OP_LOGE(context, "initialStateMode must be 1D"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(ismShape.GetDim(0) != batch, OP_LOGE(context, "initialStateMode.size must equal batch"), + return ge::GRAPH_FAILED); + + tiling.hasInitialStateMode = 1; + } + } + + tiling.hasNumAcceptedTokens = 0; + auto natShapePtr = context->GetOptionalInputShape(NUM_ACCEPTED_TOKENS_INDEX); + if (natShapePtr != nullptr) { + const auto natStorageShape = natShapePtr->GetStorageShape(); + const int64_t natDimNum = natStorageShape.GetDimNum(); + const bool natAbsent = (natDimNum == 0) || (natDimNum == 1 && natStorageShape.GetDim(0) <= 0); + if (!natAbsent) { + OP_CHECK_IF(!isDecodeMode, + OP_LOGE(context, "numAcceptedTokens is only supported in runMode=1 (decode/update)"), + return ge::GRAPH_FAILED); + auto natShape = EnsureNotScalar(natStorageShape); + OP_CHECK_IF(natShape.GetDimNum() != 1, OP_LOGE(context, "numAcceptedTokens must be 1D"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(natShape.GetDim(0) != batch, OP_LOGE(context, "numAcceptedTokens.size must equal batch"), + return ge::GRAPH_FAILED); + + if (inputMode == 1) { + const int64_t reqStateLen = (width - 1) + (seqLen - 1); + OP_CHECK_IF(stateLen < reqStateLen, + OP_LOGE(context, + "spec decode requires stateLen >= (width-1) + (seqlen-1), got stateLen=%ld req=%ld", + stateLen, reqStateLen), + return ge::GRAPH_FAILED); + } + + tiling.hasNumAcceptedTokens = 1; + } + } + + tiling.hasBias = 0; + hasBias = false; + auto biasShapePtr = context->GetOptionalInputShape(BIAS_INDEX); + if (biasShapePtr != nullptr) { + const auto biasStorageShape = biasShapePtr->GetStorageShape(); + const int64_t biasDimNum = biasStorageShape.GetDimNum(); + const bool biasAbsent = (biasDimNum == 0) || (biasDimNum == 1 && biasStorageShape.GetDim(0) <= 0); + if (!biasAbsent) { + auto biasShape = EnsureNotScalar(biasStorageShape); + OP_CHECK_IF(biasShape.GetDimNum() != 1, OP_LOGE(context, "bias must be 1D: (dim,)"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(biasShape.GetDim(0) != dim, OP_LOGE(context, "bias.size must equal dim"), + return ge::GRAPH_FAILED); + tiling.hasBias = 1; + hasBias = true; + } + } + + auto xDesc = context->GetInputDesc(X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, xDesc); + const ge::DataType xDtype = xDesc->GetDataType(); + OP_CHECK_IF(xDtype != ge::DT_BF16 && xDtype != ge::DT_FLOAT16, + OP_LOGE(context, "x dtype only supports bf16/fp16"), + return ge::GRAPH_FAILED); + + auto wDesc = context->GetInputDesc(WEIGHT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, wDesc); + OP_CHECK_IF(wDesc->GetDataType() != xDtype, OP_LOGE(context, "weight dtype must equal x dtype"), + return ge::GRAPH_FAILED); + + if (hasBias) { + auto biasDesc = context->GetOptionalInputDesc(BIAS_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, biasDesc); + OP_CHECK_IF(biasDesc->GetDataType() != xDtype, OP_LOGE(context, "bias dtype must equal x dtype"), + return ge::GRAPH_FAILED); + } + + auto sDesc = context->GetInputDesc(CONV_STATES_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, sDesc); + OP_CHECK_IF(sDesc->GetDataType() != xDtype, OP_LOGE(context, "convStates dtype must equal x dtype"), + return ge::GRAPH_FAILED); + + if (!qslAbsent) { + auto qslDesc2 = context->GetOptionalInputDesc(QUERY_START_LOC_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, qslDesc2); + const ge::DataType qslDtype = qslDesc2->GetDataType(); + OP_CHECK_IF(qslDtype != ge::DT_INT32 && qslDtype != ge::DT_INT64, + OP_LOGE(context, "queryStartLoc dtype must be int32 or int64"), return ge::GRAPH_FAILED); + tiling.queryStartLocUseInt64 = (qslDtype == ge::DT_INT64) ? 1 : 0; + } + tiling.cacheIndicesUseInt64 = 0; + if (tiling.hasCacheIndices == 1) { + auto ciDesc = context->GetOptionalInputDesc(CACHE_INDICES_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, ciDesc); + const ge::DataType ciDtype = ciDesc->GetDataType(); + OP_CHECK_IF(ciDtype != ge::DT_INT32 && ciDtype != ge::DT_INT64, + OP_LOGE(context, "cacheIndices dtype must be int32 or int64"), return ge::GRAPH_FAILED); + tiling.cacheIndicesUseInt64 = (ciDtype == ge::DT_INT64) ? 1 : 0; + } + if (tiling.hasInitialStateMode == 1) { + auto ismDesc = context->GetOptionalInputDesc(INITIAL_STATE_MODE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, ismDesc); + const ge::DataType ismDtype = ismDesc->GetDataType(); + OP_CHECK_IF(ismDtype != ge::DT_BOOL && ismDtype != ge::DT_INT32 && ismDtype != ge::DT_INT64, + OP_LOGE(context, "initialStateMode dtype must be bool, int32 or int64"), return ge::GRAPH_FAILED); + tiling.initialStateModeDtype = + (ismDtype == ge::DT_INT64) ? 2 : ((ismDtype == ge::DT_INT32) ? 1 : 0); + } + tiling.numAcceptedTokensUseInt64 = 0; + if (tiling.hasNumAcceptedTokens == 1) { + OP_CHECK_IF(width != 4, OP_LOGE(context, "numAcceptedTokens is only supported for width=4 currently"), + return ge::GRAPH_FAILED); + auto natDesc = context->GetOptionalInputDesc(NUM_ACCEPTED_TOKENS_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context, natDesc); + const ge::DataType natDtype = natDesc->GetDataType(); + OP_CHECK_IF(natDtype != ge::DT_INT32 && natDtype != ge::DT_INT64, + OP_LOGE(context, "numAcceptedTokens dtype must be int32 or int64"), return ge::GRAPH_FAILED); + tiling.numAcceptedTokensUseInt64 = (natDtype == ge::DT_INT64) ? 1 : 0; + } + + tiling.dim = dim; + tiling.cuSeqlen = cuSeqlen; + tiling.seqLen = seqLen; + tiling.inputMode = inputMode; + tiling.width = width; + tiling.stateLen = stateLen; + tiling.numCacheLines = numCacheLines; + tiling.batch = batch; + return ge::GRAPH_SUCCESS; + } + + } + + #endif diff --git a/csrc/ascend/moe/causal_conv1d/op_host/math_util.h b/csrc/ascend/moe/causal_conv1d/op_host/math_util.h new file mode 100644 index 000000000..14f90fc4c --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_host/math_util.h @@ -0,0 +1,61 @@ +/** +* Copyright (c) 2025 Huawei Technologies Co., Ltd. +* This program is free software, you can redistribute it and/or modify it under the terms and conditions of +* CANN Open Software License Agreement Version 2.0 (the "License"). +* Please refer to the License for details. You may not use this file except in compliance with the License. +* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +* See LICENSE in the root of the software repository for the full text of the License. +*/ + +/*! + * \file math_util.h + * \brief + */ + +#ifndef TILING_MATMUL_MATH_UTIL_H +#define TILING_MATMUL_MATH_UTIL_H + +#include +#include +#include +#include +namespace matmul_tiling { +class MathUtil { +public: + static bool IsEqual(float leftValue, float rightValue); + template + static auto CeilDivision(T num1, T num2) -> T + { + if (num2 == 0) { + return 0; + } + return static_cast((static_cast(num1) + static_cast(num2) - 1) / + static_cast(num2)); + } + template + static auto Align(T num1, T num2) -> T + { + return CeilDivision(num1, num2) * num2; + } + static int32_t AlignDown(int32_t num1, int32_t num2); + static bool CheckMulOverflow(int32_t a, int32_t b, int32_t &c); + static int32_t MapShape(int32_t shape, bool roundUpFlag = true); + static void AddFactor(std::vector &dimsFactors, int32_t dim); + static void GetFactorCnt(const int32_t shape, int32_t &factorCnt, const int32_t factorStart, + const int32_t factorEnd); + static void GetFactorLayerCnt(const int32_t shape, int32_t &factorCnt, const int32_t factorStart, + const int32_t factorEnd); + static bool CheckFactorNumSatisfy(const int32_t dim); + static int32_t FindBestSingleCore(const int32_t oriShape, const int32_t mappedShape, const int32_t coreNum, + bool isKDim); + static void GetFactors(std::vector &factorList, int32_t srcNum, int32_t minFactor, int32_t maxFactor); + static void GetFactors(std::vector &factorList, int32_t srcNum, int32_t maxFactor); + static void GetBlockFactors(std::vector &factorList, const int32_t oriShape, const int32_t mpShape, + const int32_t coreNum, const int32_t maxNum); + static int32_t GetNonFactorMap(std::vector &factorList, int32_t srcNum, int32_t maxFactor); + static std::vector> GetFactorPairs(int32_t num); + static std::pair DivideIntoMainAndTail(int32_t num, int32_t divisor); +}; +} // namespace matmul_tiling +#endif // _MATH_UTIL_H_ \ No newline at end of file diff --git a/csrc/ascend/moe/causal_conv1d/op_kernel/arch35/causal_conv1d_regbase.h b/csrc/ascend/moe/causal_conv1d/op_kernel/arch35/causal_conv1d_regbase.h new file mode 100644 index 000000000..cb12fb950 --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_kernel/arch35/causal_conv1d_regbase.h @@ -0,0 +1,176 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file causal_conv1d_regbase.h + * \brief + */ +#ifndef CAUSAL_CONV1D_REGBASE_H +#define CAUSAL_CONV1D_REGBASE_H + +namespace NsCausalConv1d { +using namespace AscendC; +using namespace AscendC::MicroAPI; + +constexpr uint16_t V_LENGTH = VECTOR_REG_WIDTH / sizeof(float); + +constexpr CastTrait castTraitB16ToB32 = { + RegLayout::ZERO, SatMode::UNKNOWN, MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + +template +__aicore__ inline void ComputeFnRollingOutputRegbase(LocalTensor ring, LocalTensor currF, + LocalTensor state0F, LocalTensor weightF, uint32_t dataCount) +{ + __ubuf__ T* ringAddr = (__ubuf__ T*)ring.GetPhyAddr(); + __ubuf__ float* currFAddr = (__ubuf__ float*)currF.GetPhyAddr(); + __ubuf__ float* state0FAddr = (__ubuf__ float*)state0F.GetPhyAddr(); + __ubuf__ float* weightFAddr = (__ubuf__ float*)weightF.GetPhyAddr(); + + uint16_t colLoopTimes = static_cast(Ceil(dataCount, V_LENGTH)); + __VEC_SCOPE__ + { + RegTensor ring; + RegTensor currF; + RegTensor state0F; + RegTensor weightF; + RegTensor tmp; + MaskReg pregLoop; + for (uint16_t j = 0; j < colLoopTimes; j++) { + pregLoop = UpdateMask(dataCount); + DataCopy(ring, ringAddr + j * V_LENGTH); + DataCopy(state0F, state0FAddr + j * V_LENGTH); + DataCopy(weightF, weightFAddr + j * V_LENGTH); + Cast(currF, ring, pregLoop); + Mul(currF, currF, weightF, pregLoop); + Add(state0F, state0F, currF, pregLoop); + if constexpr (hasActivation) { + Muls(tmp, state0F, -1.0f, pregLoop); + Exp(tmp, tmp, pregLoop); + Adds(tmp, tmp, 1.0f, pregLoop); + Div(currF, state0F, tmp, pregLoop); + DataCopy(currFAddr + j * V_LENGTH, currF, pregLoop); + } else { + DataCopy(state0FAddr + j * V_LENGTH, state0F, pregLoop); + } + } + } +} + +template +static __simd_vf__ inline void AdvanceFnLocalPartialsWidthTwo(__ubuf__ T* ringAddr, __ubuf__ float* weight0FAddr, + __ubuf__ float* state0FAddr, uint32_t dataCount, uint16_t colLoopTimes) +{ + RegTensor ring; + RegTensor currF; + RegTensor weight0F; + RegTensor state0F; + MaskReg pregLoop; + for (uint16_t j = 0; j < colLoopTimes; j++) { + pregLoop = UpdateMask(dataCount); + DataCopy(ring, ringAddr + j * V_LENGTH); + DataCopy(weight0F, weight0FAddr + j * V_LENGTH); + Cast(currF, ring, pregLoop); + Mul(state0F, currF, weight0F, pregLoop); + DataCopy(state0FAddr + j * V_LENGTH, state0F, pregLoop); + } +} + +template +static __simd_vf__ inline void AdvanceFnLocalPartialsWidthThree(__ubuf__ T* ringAddr, __ubuf__ float* weight0FAddr, + __ubuf__ float* weight1FAddr, __ubuf__ float* state0FAddr, __ubuf__ float* state1FAddr, uint32_t dataCount, + uint16_t colLoopTimes) +{ + RegTensor ring; + RegTensor currF; + RegTensor weight0F; + RegTensor weight1F; + RegTensor state0F; + RegTensor state1F; + MaskReg pregLoop; + for (uint16_t j = 0; j < colLoopTimes; j++) { + pregLoop = UpdateMask(dataCount); + DataCopy(ring, ringAddr + j * V_LENGTH); + DataCopy(state1F, state1FAddr + j * V_LENGTH); + Cast(currF, ring, pregLoop); + DataCopy(weight1F, weight1FAddr + j * V_LENGTH); + Mul(state0F, currF, weight1F, pregLoop); + DataCopy(weight0F, weight0FAddr + j * V_LENGTH); + Add(state0F, state0F, state1F, pregLoop); + Mul(state1F, currF, weight0F, pregLoop); + DataCopy(state0FAddr + j * V_LENGTH, state0F, pregLoop); + DataCopy(state1FAddr + j * V_LENGTH, state1F, pregLoop); + } +} + +template +static __simd_vf__ inline void AdvanceFnLocalPartialsWidthFour(__ubuf__ T* ringAddr, __ubuf__ float* weight0FAddr, + __ubuf__ float* weight1FAddr, __ubuf__ float* weight2FAddr, __ubuf__ float* state0FAddr, __ubuf__ float* state1FAddr, + __ubuf__ float* state2FAddr, uint32_t dataCount, uint16_t colLoopTimes) +{ + RegTensor ring; + RegTensor currF; + RegTensor weight0F; + RegTensor weight1F; + RegTensor weight2F; + RegTensor state0F; + RegTensor state1F; + RegTensor state2F; + MaskReg pregLoop; + for (uint16_t j = 0; j < colLoopTimes; j++) { + pregLoop = UpdateMask(dataCount); + DataCopy(ring, ringAddr + j * V_LENGTH); + DataCopy(state1F, state1FAddr + j * V_LENGTH); + DataCopy(state2F, state2FAddr + j * V_LENGTH); + Cast(currF, ring, pregLoop); + DataCopy(weight2F, weight2FAddr + j * V_LENGTH); + Mul(state0F, currF, weight2F, pregLoop); + DataCopy(weight1F, weight1FAddr + j * V_LENGTH); + Add(state0F, state0F, state1F, pregLoop); + Mul(state1F, currF, weight1F, pregLoop); + DataCopy(weight0F, weight0FAddr + j * V_LENGTH); + Add(state1F, state1F, state2F, pregLoop); + Mul(state2F, currF, weight0F, pregLoop); + DataCopy(state0FAddr + j * V_LENGTH, state0F, pregLoop); + DataCopy(state1FAddr + j * V_LENGTH, state1F, pregLoop); + DataCopy(state2FAddr + j * V_LENGTH, state2F, pregLoop); + } +} + +template +__aicore__ inline void AdvanceFnLocalPartialsRegbase(LocalTensor ring, LocalTensor weightF, + LocalTensor state0F, LocalTensor state1F, LocalTensor state2F, uint32_t dataCount, + uint32_t weightStep) +{ + uint16_t colLoopTimes = static_cast(Ceil(dataCount, V_LENGTH)); + + __ubuf__ T* ringAddr = (__ubuf__ T*)ring.GetPhyAddr(); + __ubuf__ float* weight0FAddr = (__ubuf__ float*)weightF.GetPhyAddr(); + __ubuf__ float* state0FAddr = (__ubuf__ float*)state0F.GetPhyAddr(); + if constexpr (kTemplateWidth == 2) { + AscendC::VF_CALL>(ringAddr, weight0FAddr, state0FAddr, dataCount, colLoopTimes); + } else if constexpr (kTemplateWidth == 3) { + __ubuf__ float* weight1FAddr = weight0FAddr + weightStep; + __ubuf__ float* state1FAddr = (__ubuf__ float*)state1F.GetPhyAddr(); + AscendC::VF_CALL>(ringAddr, weight0FAddr, weight1FAddr, state0FAddr, + state1FAddr, dataCount, colLoopTimes); + } else if constexpr (kTemplateWidth == 4) { + __ubuf__ float* weight1FAddr = weight0FAddr + weightStep; + __ubuf__ float* weight2FAddr = weight1FAddr + weightStep; + __ubuf__ float* state1FAddr = (__ubuf__ float*)state1F.GetPhyAddr(); + __ubuf__ float* state2FAddr = (__ubuf__ float*)state2F.GetPhyAddr(); + AscendC::VF_CALL>(ringAddr, weight0FAddr, weight1FAddr, weight2FAddr, + state0FAddr, state1FAddr, state2FAddr, dataCount, colLoopTimes); + } +} + +} // namespace NsCausalConv1d + +#endif // CAUSAL_CONV1D_REGBASE_H \ No newline at end of file diff --git a/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d.cpp b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d.cpp new file mode 100644 index 000000000..4e9902f5a --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d.cpp @@ -0,0 +1,57 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file causal_conv1d.cpp + * \brief + */ + +#include "causal_conv1d_fn.h" +#include "causal_conv1d_update.h" + +namespace { + +template +__aicore__ inline void RunCausalConv1d(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR convStates, + GM_ADDR queryStartLoc, GM_ADDR cacheIndices, GM_ADDR initialStateMode, + GM_ADDR numAcceptedTokens, GM_ADDR y, GM_ADDR workspace, + const CausalConv1dTilingData *tilingData) +{ + if constexpr (runModeKey == CAUSAL_CONV1D_TPL_RUN_MODE_FN) { + NsCausalConv1d::RunCausalConv1dFn( + x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, y, workspace, + tilingData); + } else { + NsCausalConv1d::RunCausalConv1dUpdate( + x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, y, workspace, + tilingData); + } +} + +} // namespace + +template +__global__ __aicore__ void causal_conv1d(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR convStates, + GM_ADDR queryStartLoc, GM_ADDR cacheIndices, GM_ADDR initialStateMode, + GM_ADDR numAcceptedTokens, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) +{ + REGISTER_TILING_DEFAULT(CausalConv1dTilingData); + GET_TILING_DATA(tilingData, tiling); + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0); + GM_ADDR userWorkspace = workspace; + if (workspace != nullptr) { + userWorkspace = AscendC::GetUserWorkspace(workspace); + } + + RunCausalConv1d( + x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, y, + userWorkspace, &tilingData); +} diff --git a/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d.h b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d.h new file mode 100644 index 000000000..c510a27fd --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d.h @@ -0,0 +1,1088 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file causal_conv1d.h + */ + +#ifndef CAUSAL_CONV1D_H +#define CAUSAL_CONV1D_H + +#include "kernel_operator.h" +#include "kernel_tiling/kernel_tiling.h" +#include "causal_conv1d_tiling_data.h" +#include "causal_conv1d_tiling_key.h" +#include "causal_conv1d_common.h" +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#include "arch35/causal_conv1d_regbase.h" +#endif + + namespace NsCausalConv1d { + + using namespace AscendC; + using namespace NsCausalConv1dCommon; + + #define CAUSAL_CONV1D_TEMPLATE_ARGS typename T, uint32_t runModeKey, uint32_t widthKey, uint32_t fnPlanKey + #define CAUSAL_CONV1D_CLASS CausalConv1d + + enum SeqTaskWindowMode : int32_t { + SEQ_TASK_WINDOW_MODE_VARLEN = 0, + SEQ_TASK_WINDOW_MODE_BATCH = 1, + SEQ_TASK_WINDOW_MODE_DECODE2D = 2, + }; + + inline constexpr int32_t INIT_STATE_SYNCALL_NEED_SIZE = 8; + inline constexpr int32_t INIT_STATE_SYNCALL_MAX_BLOCKS = 64; + inline constexpr int64_t INT32_MAX_VALUE = 2147483647LL; + + struct SeqTaskWindow { + bool valid = false; + int32_t start = 0; + int32_t len = 0; + }; + + __aicore__ inline int32_t GetSeqTaskWindowMode(int32_t inputMode) + { + if (inputMode == 0) { + return SEQ_TASK_WINDOW_MODE_VARLEN; + } + if (inputMode == 2) { + return SEQ_TASK_WINDOW_MODE_DECODE2D; + } + return SEQ_TASK_WINDOW_MODE_BATCH; + } + + __aicore__ inline SeqTaskWindow BuildSeqTaskWindowVarlen(int32_t startVal, int32_t endVal) + { + SeqTaskWindow window; + window.start = startVal; + window.len = endVal - startVal; + window.valid = (window.len > 0); + return window; + } + + __aicore__ inline int32_t RetreatRingSlot(int32_t slot, int32_t delta) + { + int32_t prev = slot - delta; + return (prev >= 0) ? prev : (prev + RING_SLOTS); + } + + __aicore__ inline SeqTaskWindow BuildSeqTaskWindowBatch(int32_t seq, int32_t seqLen) + { + SeqTaskWindow window; + window.start = seq * seqLen; + window.len = seqLen; + window.valid = (window.len > 0); + return window; + } + + __aicore__ inline SeqTaskWindow BuildSeqTaskWindowDecode2D(int32_t seq) + { + SeqTaskWindow window; + window.valid = true; + window.start = seq; + window.len = 1; + return window; + } + + __aicore__ inline constexpr int32_t DecodeWidthTplKey(uint32_t widthKey) + { + switch (widthKey) { + case CAUSAL_CONV1D_TPL_WIDTH_2: + return 2; + case CAUSAL_CONV1D_TPL_WIDTH_3: + return 3; + case CAUSAL_CONV1D_TPL_WIDTH_4: + return 4; + default: + return 0; + } + } + + template + class CausalConv1d { + public: + __aicore__ inline CausalConv1d() = default; + + protected: + static constexpr bool kIsUpdateMode = (runModeKey == CAUSAL_CONV1D_TPL_RUN_MODE_UPDATE); + static constexpr int32_t kTemplateWidth = DecodeWidthTplKey(widthKey); + static constexpr bool kHasCompileTimeWidth = + (runModeKey == CAUSAL_CONV1D_TPL_RUN_MODE_FN) && (kTemplateWidth >= 2) && (kTemplateWidth <= MAX_WIDTH); + static constexpr FnExecutionPlan kFnExecutionPlan = static_cast(fnPlanKey); + + __aicore__ inline void ResetRuntimeState(const CausalConv1dTilingData *tilingData); + __aicore__ inline void InitSharedBuffersAndEvents(); + __aicore__ inline void LoadWeightAndBias(int32_t channelStart, int32_t baseDim); + __aicore__ inline void InitRing(int32_t cacheIdx, bool hasInit, int32_t stateTokenOffset, int32_t start, + int32_t len, int32_t channelStart, int32_t baseDim, int32_t dim); + __aicore__ inline void InitRingSeqSplit(int32_t seq, int32_t cacheIdx, bool hasInit, int32_t seqStart, + int32_t tileStart, int32_t tileLen, int32_t channelStart, int32_t baseDim, + int32_t dim); + __aicore__ inline void PrefetchInitStatesToWorkspace(int32_t channelStart, int32_t baseDimSize); + __aicore__ inline void RestoreFnLocalPartials(int32_t baseDim); + __aicore__ inline void ComputeFnRollingOutput(int32_t slotCurr, int32_t baseDim); + __aicore__ inline void AdvanceFnLocalPartials(int32_t slotCurr, int32_t baseDim); + __aicore__ inline void RunSeqFnRolling(int32_t start, int32_t len, int32_t channelStart, int32_t baseDim, + int32_t dim); + __aicore__ inline void RunSeq(int32_t start, int32_t len, int32_t channelStart, int32_t baseDim, int32_t dim); + __aicore__ inline void WriteBackState(int32_t cacheIdx, int32_t len, int32_t channelStart, int32_t baseDim, + int32_t dim); + __aicore__ inline void WriteBackStateSpec(int32_t cacheIdx, bool hasInit, int32_t stateTokenOffset, int32_t start, + int32_t len, int32_t channelStart, int32_t baseDim, int32_t dim); + __aicore__ inline void DrainTaskMte3(); + __aicore__ inline void AllocEvents(); + __aicore__ inline void ReleaseEvents(); + __aicore__ inline int32_t FindVarlenSeqByToken(int32_t tokenIdx) const; + __aicore__ inline bool ResolveExplicitTokenTileSeqRange(int32_t tokenTileId, int32_t &startSeq, int32_t &endSeq) const; + __aicore__ inline bool ResolveSeqTaskWindow(int32_t seq, int32_t inputMode, int32_t seqLen, int32_t &start, + int32_t &len) const; + template + __aicore__ inline bool ResolveSeqTaskWindowByMode(int32_t seq, int32_t seqLen, int32_t &start, int32_t &len) const; + __aicore__ inline int32_t ReadQueryStartLocValue(int32_t index) const; + __aicore__ inline int64_t ReadCacheIndexValue(int32_t seq) const; + __aicore__ inline bool ReadInitialStateModeValue(int32_t seq) const; + __aicore__ inline int32_t ReadNumAcceptedTokensValue(int32_t seq) const; + __aicore__ inline bool ResolveSeqCacheIndex(int32_t seq, bool hasCacheIndices, int32_t &cacheIdx) const; + __aicore__ inline bool ResolveSeqHasInit(int32_t seq, bool hasInitialStateMode) const; + __aicore__ inline void MaybeWriteBackSeqSplitTailChunk(int32_t chunkStart, int32_t chunkLen, int32_t seqStart, + int32_t seqLen, int32_t cacheIdx, int32_t channelStart, + int32_t baseDim, int32_t dim); + __aicore__ inline void ProcessDefault(); + template + __aicore__ inline void ProcessDefaultByWindowMode(); + __aicore__ inline void ProcessVarlenTokenTiled(); + __aicore__ inline void ProcessFnChunk(int32_t seq, int32_t cacheIdx, bool hasInit, int32_t seqStart, + int32_t seqLen, int32_t chunkStart, int32_t chunkLen, int32_t channelStart, + int32_t baseDim, int32_t dim); + __aicore__ inline const CausalConv1dTilingData *GetTilingData() const; + __aicore__ inline bool HasActivation() const; + __aicore__ inline bool HasBias() const; + __aicore__ inline bool IsUpdateMode() const; + __aicore__ inline bool IsFnRollingFastPathEnabled() const; + __aicore__ inline bool HasExplicitFnTokenSeqRanges() const; + __aicore__ inline bool IsUpdateSpecDecodingEnabled() const; + + protected: + TPipe pipe; + TBuf inBuf; + TBuf outBuf; + TBuf calcBuf; + + TEventID weightBiasMte2ToVEvent_; + TEventID stateMte2ToVEvent_; + TEventID inputMte2ToVEvent_[RING_SLOTS]; + TEventID inputVToMte2Event_; + TEventID outMte3ToVEvent_[2]; + TEventID outVToMte3Event_[2]; + TEventID stateWritebackMte3ToVEvent_; + TEventID stateWritebackMte3ToMte2Event_; + TEventID stateShiftMte2ToMte3Event_; + TEventID stateShiftVToMte3Event_; + TEventID stateShiftMte3ToMte2Event_; + TEventID initSnapshotMte2ToMte3Event_; + TEventID initSnapshotMte3ToMte2Event_; + TEventID initSyncVToMte3Event_; + TEventID initSyncMte3ToVEvent_; + TEventID specWritebackMte2ToMte3Event_[2]; + TEventID specWritebackMte3ToMte2Event_[2]; + + GlobalTensor xGm; + GlobalTensor weightGm; + GlobalTensor biasGm; + GlobalTensor convStatesGm; + GlobalTensor queryStartLocGmInt32; + GlobalTensor queryStartLocGmInt64; + GlobalTensor cacheIndicesGmInt32; + GlobalTensor cacheIndicesGmInt64; + GlobalTensor initialStateModeGmBool; + GlobalTensor initialStateModeGmInt32; + GlobalTensor initialStateModeGmInt64; + GlobalTensor numAcceptedTokensGmInt32; + GlobalTensor numAcceptedTokensGmInt64; + GlobalTensor yGm; + GlobalTensor initStateSyncGm_; + GlobalTensor initStateWorkspaceGm_; + + const CausalConv1dTilingData *tilingData_{nullptr}; + }; + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::ResetRuntimeState(const CausalConv1dTilingData *tilingData) + { + tilingData_ = tilingData; + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::InitSharedBuffersAndEvents() + { + pipe.InitBuffer(inBuf, RING_SLOTS * MAX_BLOCK_DIM * sizeof(T)); + pipe.InitBuffer(outBuf, 2 * MAX_BLOCK_DIM * sizeof(T)); + pipe.InitBuffer(calcBuf, (MAX_WIDTH + 4) * MAX_BLOCK_DIM * sizeof(float)); + AllocEvents(); + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::AllocEvents() + { + weightBiasMte2ToVEvent_ = GetTPipePtr()->AllocEventID(); + stateMte2ToVEvent_ = GetTPipePtr()->AllocEventID(); + for (int32_t i = 0; i < RING_SLOTS; ++i) { + inputMte2ToVEvent_[i] = GetTPipePtr()->AllocEventID(); + } + inputVToMte2Event_ = GetTPipePtr()->AllocEventID(); + outMte3ToVEvent_[0] = GetTPipePtr()->AllocEventID(); + outMte3ToVEvent_[1] = GetTPipePtr()->AllocEventID(); + outVToMte3Event_[0] = GetTPipePtr()->AllocEventID(); + outVToMte3Event_[1] = GetTPipePtr()->AllocEventID(); + stateWritebackMte3ToVEvent_ = GetTPipePtr()->AllocEventID(); + stateWritebackMte3ToMte2Event_ = GetTPipePtr()->AllocEventID(); + stateShiftMte2ToMte3Event_ = GetTPipePtr()->AllocEventID(); + stateShiftVToMte3Event_ = GetTPipePtr()->AllocEventID(); + stateShiftMte3ToMte2Event_ = GetTPipePtr()->AllocEventID(); + initSnapshotMte2ToMte3Event_ = GetTPipePtr()->AllocEventID(); + initSnapshotMte3ToMte2Event_ = GetTPipePtr()->AllocEventID(); + initSyncVToMte3Event_ = GetTPipePtr()->AllocEventID(); + initSyncMte3ToVEvent_ = GetTPipePtr()->AllocEventID(); + specWritebackMte2ToMte3Event_[0] = GetTPipePtr()->AllocEventID(); + specWritebackMte2ToMte3Event_[1] = GetTPipePtr()->AllocEventID(); + specWritebackMte3ToMte2Event_[0] = GetTPipePtr()->AllocEventID(); + specWritebackMte3ToMte2Event_[1] = GetTPipePtr()->AllocEventID(); + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::ReleaseEvents() + { + GetTPipePtr()->ReleaseEventID(weightBiasMte2ToVEvent_); + GetTPipePtr()->ReleaseEventID(stateMte2ToVEvent_); + for (int32_t i = 0; i < RING_SLOTS; ++i) { + GetTPipePtr()->ReleaseEventID(inputMte2ToVEvent_[i]); + } + GetTPipePtr()->ReleaseEventID(inputVToMte2Event_); + GetTPipePtr()->ReleaseEventID(outMte3ToVEvent_[0]); + GetTPipePtr()->ReleaseEventID(outMte3ToVEvent_[1]); + GetTPipePtr()->ReleaseEventID(outVToMte3Event_[0]); + GetTPipePtr()->ReleaseEventID(outVToMte3Event_[1]); + GetTPipePtr()->ReleaseEventID(stateWritebackMte3ToVEvent_); + GetTPipePtr()->ReleaseEventID(stateWritebackMte3ToMte2Event_); + GetTPipePtr()->ReleaseEventID(stateShiftMte2ToMte3Event_); + GetTPipePtr()->ReleaseEventID(stateShiftVToMte3Event_); + GetTPipePtr()->ReleaseEventID(stateShiftMte3ToMte2Event_); + GetTPipePtr()->ReleaseEventID(initSnapshotMte2ToMte3Event_); + GetTPipePtr()->ReleaseEventID(initSnapshotMte3ToMte2Event_); + GetTPipePtr()->ReleaseEventID(initSyncVToMte3Event_); + GetTPipePtr()->ReleaseEventID(initSyncMte3ToVEvent_); + GetTPipePtr()->ReleaseEventID(specWritebackMte2ToMte3Event_[0]); + GetTPipePtr()->ReleaseEventID(specWritebackMte2ToMte3Event_[1]); + GetTPipePtr()->ReleaseEventID(specWritebackMte3ToMte2Event_[0]); + GetTPipePtr()->ReleaseEventID(specWritebackMte3ToMte2Event_[1]); + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::LoadWeightAndBias(int32_t channelStart, int32_t baseDim) + { + const int32_t dim = tilingData_->dim; + const int32_t width = static_cast(tilingData_->width); + const int32_t jStart = MAX_WIDTH - width; + const bool hasBias = HasBias(); + auto cl = CalcBufLayout::FromCalcBuf(calcBuf); + LocalTensor &weightF = cl.weightF; + LocalTensor &biasF = cl.biasF; + LocalTensor weightT; + LocalTensor biasT; + + if constexpr (!std::is_same::value) { + weightT = weightF.ReinterpretCast(); + biasT = biasF.ReinterpretCast(); + } + + for (int32_t j = 0; j < jStart; ++j) { + Duplicate(weightF[j * MAX_BLOCK_DIM], 0.0f, baseDim); + } + + for (int32_t j = 0; j < width; ++j) { + const int32_t jDst = jStart + j; + const int64_t weightOffset = static_cast(j) * dim + channelStart; + + if constexpr (std::is_same::value) { + DataCopy(weightF[jDst * MAX_BLOCK_DIM], weightGm[weightOffset], baseDim); + } else { + DataCopy(weightT[jDst * MAX_BLOCK_DIM * 2 + MAX_BLOCK_DIM], weightGm[weightOffset], baseDim); + } + } + + if (hasBias) { + if constexpr (std::is_same::value) { + DataCopy(biasF, biasGm[channelStart], baseDim); + } else { + DataCopy(biasT[MAX_BLOCK_DIM], biasGm[channelStart], baseDim); + } + } + + SetFlag(weightBiasMte2ToVEvent_); + WaitFlag(weightBiasMte2ToVEvent_); + + if constexpr (!std::is_same::value) { + for (int32_t j = 0; j < width; ++j) { + const int32_t jDst = jStart + j; + Cast(weightF[jDst * MAX_BLOCK_DIM], weightT[jDst * MAX_BLOCK_DIM * 2 + MAX_BLOCK_DIM], RoundMode::CAST_NONE, + baseDim); + } + if (hasBias) { + Cast(biasF, biasT[MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); + } + PipeBarrier(); + } + + if (!hasBias) { + Duplicate(biasF, 0.0f, baseDim); + } + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::InitRing(int32_t cacheIdx, bool hasInit, int32_t stateTokenOffset, + int32_t start, int32_t len, int32_t channelStart, + int32_t baseDim, int32_t dim) + { + const int32_t stateLen = tilingData_->stateLen; + const int32_t width = static_cast(tilingData_->width); + const int32_t ringStart = MAX_WIDTH - width; + LocalTensor ring = inBuf.Get(); + + for (int32_t i = 0; i < ringStart; ++i) { + Duplicate(ring[i * MAX_BLOCK_DIM], static_cast(0), baseDim); + } + if (ringStart > 0) { + PipeBarrier(); + } + + if (hasInit) { + for (int32_t i = 0; i < (width - 1); ++i) { + const int32_t pos = stateTokenOffset + i; + const int64_t stateOffset = + static_cast(cacheIdx) * stateLen * dim + static_cast(pos) * dim + channelStart; + DataCopy(ring[(ringStart + i) * MAX_BLOCK_DIM], convStatesGm[stateOffset], baseDim); + } + SetFlag(stateMte2ToVEvent_); + WaitFlag(stateMte2ToVEvent_); + } else { + for (int32_t i = 0; i < (width - 1); ++i) { + Duplicate(ring[(ringStart + i) * MAX_BLOCK_DIM], static_cast(0), baseDim); + } + PipeBarrier(); + } + + if (len > 0) { + const int32_t slot0 = SlotCurr(0); + const int64_t xOffset = static_cast(start) * dim + channelStart; + DataCopy(ring[slot0 * MAX_BLOCK_DIM], xGm[xOffset], baseDim); + SetFlag(inputMte2ToVEvent_[slot0]); + } + + if (len > 1) { + SetFlag(inputVToMte2Event_); + } + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::RunSeq(int32_t start, int32_t len, int32_t channelStart, + int32_t baseDim, int32_t dim) + { + if (IsFnRollingFastPathEnabled()) { + RunSeqFnRolling(start, len, channelStart, baseDim, dim); + return; + } + + const int32_t width = static_cast(tilingData_->width); + const int32_t jStart = MAX_WIDTH - width; + auto cl = CalcBufLayout::FromCalcBuf(calcBuf); + LocalTensor &weightF = cl.weightF; + LocalTensor &biasF = cl.biasF; + LocalTensor &accF = cl.accF; + LocalTensor &tmpF = cl.tmpF; + LocalTensor ring = inBuf.Get(); + LocalTensor outT = outBuf.Get(); + const bool hasBias = HasBias(); + const bool hasActivation = HasActivation(); + for (int32_t t = 0; t < len; ++t) { + const int32_t slotCurr = SlotCurr(t); + + WaitFlag(inputMte2ToVEvent_[slotCurr]); + + if (t + 1 < len) { + const int32_t slotNext = SlotPrefetch(t); + const int64_t xOffsetNext = static_cast(start + t + 1) * dim + channelStart; + WaitFlag(inputVToMte2Event_); + DataCopy(ring[slotNext * MAX_BLOCK_DIM], xGm[xOffsetNext], baseDim); + SetFlag(inputMte2ToVEvent_[slotNext]); + } + + bool accInitialized = false; + if (hasBias) { + Adds(accF, biasF, 0.0f, baseDim); + PipeBarrier(); + accInitialized = true; + } + + for (int32_t j = jStart; j < MAX_WIDTH; ++j) { + const int32_t tap = (MAX_WIDTH - 1) - j; + const int32_t slot = (tap == 0) ? slotCurr : SlotHist(t, tap); + Cast(tmpF, ring[slot * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); + PipeBarrier(); + if (!accInitialized) { + Mul(accF, tmpF, weightF[j * MAX_BLOCK_DIM], baseDim); + accInitialized = true; + } else { + MulAddDst(accF, tmpF, weightF[j * MAX_BLOCK_DIM], baseDim); + } + } + + PipeBarrier(); + + if (hasActivation) { + Silu(tmpF, accF, baseDim); + } + + const int32_t outSlot = t & 1; + LocalTensor outSlotT = outT[outSlot * MAX_BLOCK_DIM]; + if (t >= 2) { + WaitFlag(outMte3ToVEvent_[outSlot]); + } + + if constexpr (IsSameType::value) { + if (hasActivation) { + DataCopy(outSlotT, tmpF, baseDim); + } else { + DataCopy(outSlotT, accF, baseDim); + } + } else { + if (hasActivation) { + Cast(outSlotT, tmpF, RoundMode::CAST_RINT, baseDim); + } else { + Cast(outSlotT, accF, RoundMode::CAST_RINT, baseDim); + } + } + + SetFlag(outVToMte3Event_[outSlot]); + + const int64_t outOffset = static_cast(start + t) * dim + channelStart; + WaitFlag(outVToMte3Event_[outSlot]); + DataCopy(yGm[outOffset], outSlotT, baseDim); + if (t + 2 < len) { + SetFlag(outMte3ToVEvent_[outSlot]); + } + + if (t + 2 < len) { + SetFlag(inputVToMte2Event_); + } + } + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::RestoreFnLocalPartials(int32_t baseDim) + { + if constexpr (!kHasCompileTimeWidth) { + return; + } + + auto cl = CalcBufLayout::FromCalcBuf(calcBuf); + LocalTensor &weightF = cl.weightF; + LocalTensor &state2F = cl.biasF; + LocalTensor &state1F = cl.accF; + LocalTensor &state0F = cl.tmpF; + LocalTensor &currF = cl.currF; + LocalTensor ring = inBuf.Get(); + constexpr int32_t ringStart = MAX_WIDTH - kTemplateWidth; + constexpr int32_t w0Idx = MAX_WIDTH - kTemplateWidth; + + if constexpr (kTemplateWidth == 2) { + Duplicate(state2F, 0.0f, baseDim); + Duplicate(state1F, 0.0f, baseDim); + PipeBarrier(); + + Cast(currF, ring[ringStart * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); + PipeBarrier(); + Mul(state0F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + } else if constexpr (kTemplateWidth == 3) { + Duplicate(state2F, 0.0f, baseDim); + PipeBarrier(); + + Cast(currF, ring[ringStart * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); + PipeBarrier(); + Mul(state0F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + + Cast(currF, ring[(ringStart + 1) * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); + PipeBarrier(); + Mul(state1F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + MulAddDst(state0F, currF, weightF[(w0Idx + 1) * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + } else if constexpr (kTemplateWidth == 4) { + Cast(currF, ring[ringStart * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); + PipeBarrier(); + Mul(state0F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + + Cast(currF, ring[(ringStart + 1) * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); + PipeBarrier(); + Mul(state1F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + MulAddDst(state0F, currF, weightF[(w0Idx + 1) * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + + Cast(currF, ring[(ringStart + 2) * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); + PipeBarrier(); + Mul(state2F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + MulAddDst(state1F, currF, weightF[(w0Idx + 1) * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + MulAddDst(state0F, currF, weightF[(w0Idx + 2) * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + } + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::ComputeFnRollingOutput(int32_t slotCurr, int32_t baseDim) + { + if constexpr (!kHasCompileTimeWidth) { + return; + } + + auto cl = CalcBufLayout::FromCalcBuf(calcBuf); + LocalTensor &weightF = cl.weightF; + LocalTensor &state0F = cl.tmpF; + LocalTensor &currF = cl.currF; + LocalTensor ring = inBuf.Get(); + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + const bool hasActivation = HasActivation(); + if (hasActivation) { + ComputeFnRollingOutputRegbase(ring[slotCurr * MAX_BLOCK_DIM], currF, state0F, weightF[3 * MAX_BLOCK_DIM], baseDim); + } else { + ComputeFnRollingOutputRegbase(ring[slotCurr * MAX_BLOCK_DIM], currF, state0F, weightF[3 * MAX_BLOCK_DIM], baseDim); + } +#else + Cast(currF, ring[slotCurr * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); + PipeBarrier(); + MulAddDst(state0F, currF, weightF[3 * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + + const bool hasActivation = HasActivation(); + if (hasActivation) { + PipeBarrier(); + Silu(currF, state0F, baseDim); + } +#endif + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::AdvanceFnLocalPartials(int32_t slotCurr, int32_t baseDim) + { + if constexpr (!kHasCompileTimeWidth) { + return; + } + + auto cl = CalcBufLayout::FromCalcBuf(calcBuf); + LocalTensor &weightF = cl.weightF; + LocalTensor &state2F = cl.biasF; + LocalTensor &state1F = cl.accF; + LocalTensor &state0F = cl.tmpF; + LocalTensor &currF = cl.currF; + LocalTensor ring = inBuf.Get(); + constexpr int32_t w0Idx = MAX_WIDTH - kTemplateWidth; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + AdvanceFnLocalPartialsRegbase(ring[slotCurr * MAX_BLOCK_DIM], weightF[w0Idx * MAX_BLOCK_DIM], + state0F, state1F, state2F, baseDim, MAX_BLOCK_DIM); +#else + Cast(currF, ring[slotCurr * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); + PipeBarrier(); + + if constexpr (kTemplateWidth == 2) { + Mul(state0F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + } else if constexpr (kTemplateWidth == 3) { + Mul(state0F, currF, weightF[(w0Idx + 1) * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + Add(state0F, state0F, state1F, baseDim); + PipeBarrier(); + + Mul(state1F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + } else if constexpr (kTemplateWidth == 4) { + Mul(state0F, currF, weightF[(w0Idx + 2) * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + Add(state0F, state0F, state1F, baseDim); + PipeBarrier(); + + Mul(state1F, currF, weightF[(w0Idx + 1) * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + Add(state1F, state1F, state2F, baseDim); + PipeBarrier(); + + Mul(state2F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); + PipeBarrier(); + } +#endif + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::RunSeqFnRolling(int32_t start, int32_t len, int32_t channelStart, + int32_t baseDim, int32_t dim) + { + if constexpr (!kHasCompileTimeWidth) { + return; + } + + auto cl = CalcBufLayout::FromCalcBuf(calcBuf); + LocalTensor &state0F = cl.tmpF; + LocalTensor &currF = cl.currF; + LocalTensor ring = inBuf.Get(); + LocalTensor outT = outBuf.Get(); + const bool hasActivation = HasActivation(); + RestoreFnLocalPartials(baseDim); + + for (int32_t t = 0; t < len; ++t) { + const int32_t slotCurr = SlotCurr(t); + + WaitFlag(inputMte2ToVEvent_[slotCurr]); + + if (t + 1 < len) { + const int32_t slotNext = SlotPrefetch(t); + const int64_t xOffsetNext = static_cast(start + t + 1) * dim + channelStart; + WaitFlag(inputVToMte2Event_); + DataCopy(ring[slotNext * MAX_BLOCK_DIM], xGm[xOffsetNext], baseDim); + SetFlag(inputMte2ToVEvent_[slotNext]); + } + + ComputeFnRollingOutput(slotCurr, baseDim); + + const int32_t outSlot = t & 1; + LocalTensor outSlotT = outT[outSlot * MAX_BLOCK_DIM]; + if (t >= 2) { + WaitFlag(outMte3ToVEvent_[outSlot]); + } + + if constexpr (IsSameType::value) { + if (hasActivation) { + DataCopy(outSlotT, currF, baseDim); + } else { + DataCopy(outSlotT, state0F, baseDim); + } + } else { + if (hasActivation) { + Cast(outSlotT, currF, RoundMode::CAST_RINT, baseDim); + } else { + Cast(outSlotT, state0F, RoundMode::CAST_RINT, baseDim); + } + } + + AdvanceFnLocalPartials(slotCurr, baseDim); + + SetFlag(outVToMte3Event_[outSlot]); + + const int64_t outOffset = static_cast(start + t) * dim + channelStart; + WaitFlag(outVToMte3Event_[outSlot]); + DataCopy(yGm[outOffset], outSlotT, baseDim); + if (t + 2 < len) { + SetFlag(outMte3ToVEvent_[outSlot]); + } + + if (t + 2 < len) { + SetFlag(inputVToMte2Event_); + } + } + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::DrainTaskMte3() + { + SetFlag(stateWritebackMte3ToVEvent_); + WaitFlag(stateWritebackMte3ToVEvent_); + SetFlag(stateWritebackMte3ToMte2Event_); + WaitFlag(stateWritebackMte3ToMte2Event_); + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::WriteBackState(int32_t cacheIdx, int32_t len, int32_t channelStart, + int32_t baseDim, int32_t dim) + { + const int32_t stateLen = tilingData_->stateLen; + const int32_t width = static_cast(tilingData_->width); + if (len <= 0) { + return; + } + + const int32_t lastT = len - 1; + LocalTensor ring = inBuf.Get(); + const int32_t lastSlot = SlotCurr(lastT); + const int64_t stateBaseOffset = static_cast(cacheIdx) * stateLen * dim + channelStart; + + for (int32_t pos = 0; pos < (width - 1); ++pos) { + const int32_t tap = (width - 2) - pos; + const int32_t slot = RetreatRingSlot(lastSlot, tap); + const int64_t stateOffset = stateBaseOffset + static_cast(pos) * dim; + DataCopy(convStatesGm[stateOffset], ring[slot * MAX_BLOCK_DIM], baseDim); + } + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::WriteBackStateSpec(int32_t cacheIdx, bool hasInit, + int32_t stateTokenOffset, int32_t start, int32_t len, + int32_t channelStart, int32_t baseDim, int32_t dim) + { + const int32_t width = static_cast(tilingData_->width); + const int32_t stateLen = tilingData_->stateLen; + if (len <= 0) { + return; + } + + if (width != 4) { + WriteBackState(cacheIdx, len, channelStart, baseDim, dim); + return; + } + + constexpr int32_t keep = MAX_WIDTH - 2; + const int32_t reqStateLen = keep + len; + if (reqStateLen > stateLen) { + WriteBackState(cacheIdx, len, channelStart, baseDim, dim); + return; + } + + LocalTensor ring = inBuf.Get(); + LocalTensor buf0 = ring[0 * MAX_BLOCK_DIM]; + LocalTensor buf1 = ring[1 * MAX_BLOCK_DIM]; + + if (hasInit) { + const int32_t srcPos0 = stateTokenOffset + 1; + const int32_t srcPos1 = stateTokenOffset + 2; + const int64_t srcOffset0 = + static_cast(cacheIdx) * stateLen * dim + static_cast(srcPos0) * dim + channelStart; + const int64_t srcOffset1 = + static_cast(cacheIdx) * stateLen * dim + static_cast(srcPos1) * dim + channelStart; + DataCopy(buf0, convStatesGm[srcOffset0], baseDim); + DataCopy(buf1, convStatesGm[srcOffset1], baseDim); + SetFlag(stateShiftMte2ToMte3Event_); + WaitFlag(stateShiftMte2ToMte3Event_); + const int64_t dstOffset0 = + static_cast(cacheIdx) * stateLen * dim + static_cast(0) * dim + channelStart; + const int64_t dstOffset1 = + static_cast(cacheIdx) * stateLen * dim + static_cast(1) * dim + channelStart; + DataCopy(convStatesGm[dstOffset0], buf0, baseDim); + DataCopy(convStatesGm[dstOffset1], buf1, baseDim); + SetFlag(stateShiftMte3ToMte2Event_); + WaitFlag(stateShiftMte3ToMte2Event_); + } else { + Duplicate(buf0, static_cast(0), baseDim); + SetFlag(stateShiftVToMte3Event_); + WaitFlag(stateShiftVToMte3Event_); + const int64_t dstOffset0 = + static_cast(cacheIdx) * stateLen * dim + static_cast(0) * dim + channelStart; + const int64_t dstOffset1 = + static_cast(cacheIdx) * stateLen * dim + static_cast(1) * dim + channelStart; + DataCopy(convStatesGm[dstOffset0], buf0, baseDim); + DataCopy(convStatesGm[dstOffset1], buf0, baseDim); + SetFlag(stateShiftMte3ToMte2Event_); + WaitFlag(stateShiftMte3ToMte2Event_); + } + + const int64_t xOffset0 = static_cast(start) * dim + channelStart; + DataCopy(buf0, xGm[xOffset0], baseDim); + SetFlag(specWritebackMte2ToMte3Event_[0]); + + for (int32_t t = 0; t < len; ++t) { + const int32_t curr = t & 1; + const int32_t next = curr ^ 1; + LocalTensor currBuf = (curr == 0) ? buf0 : buf1; + LocalTensor nextBuf = (next == 0) ? buf0 : buf1; + + WaitFlag(specWritebackMte2ToMte3Event_[curr]); + + if (t + 1 < len) { + const int64_t xOffsetNext = static_cast(start + t + 1) * dim + channelStart; + if (t > 0) { + WaitFlag(specWritebackMte3ToMte2Event_[next]); + } + DataCopy(nextBuf, xGm[xOffsetNext], baseDim); + SetFlag(specWritebackMte2ToMte3Event_[next]); + } + + const int64_t dstOffset = + static_cast(cacheIdx) * stateLen * dim + static_cast(keep + t) * dim + channelStart; + DataCopy(convStatesGm[dstOffset], currBuf, baseDim); + SetFlag(specWritebackMte3ToMte2Event_[curr]); + } + + WaitFlag(specWritebackMte3ToMte2Event_[0]); + if (len > 1) { + WaitFlag(specWritebackMte3ToMte2Event_[1]); + } + } + + template + __aicore__ inline bool CAUSAL_CONV1D_CLASS::ResolveSeqTaskWindow(int32_t seq, int32_t inputMode, int32_t seqLen, + int32_t &start, int32_t &len) const + { + switch (GetSeqTaskWindowMode(inputMode)) { + case SEQ_TASK_WINDOW_MODE_VARLEN: + return ResolveSeqTaskWindowByMode(seq, seqLen, start, len); + case SEQ_TASK_WINDOW_MODE_DECODE2D: + return ResolveSeqTaskWindowByMode(seq, seqLen, start, len); + default: + return ResolveSeqTaskWindowByMode(seq, seqLen, start, len); + } + } + + template + template + __aicore__ inline bool CAUSAL_CONV1D_CLASS::ResolveSeqTaskWindowByMode(int32_t seq, int32_t seqLen, int32_t &start, + int32_t &len) const + { + SeqTaskWindow window; + if constexpr (kWindowMode == SEQ_TASK_WINDOW_MODE_VARLEN) { + const int32_t startVal = ReadQueryStartLocValue(seq); + const int32_t endVal = ReadQueryStartLocValue(seq + 1); + if (startVal < 0 || endVal < startVal || endVal > tilingData_->cuSeqlen) { + return false; + } + window = BuildSeqTaskWindowVarlen(startVal, endVal); + } else if constexpr (kWindowMode == SEQ_TASK_WINDOW_MODE_DECODE2D) { + window = BuildSeqTaskWindowDecode2D(seq); + } else { + window = BuildSeqTaskWindowBatch(seq, seqLen); + } + + if (!window.valid) { + return false; + } + start = window.start; + len = window.len; + return true; + } + + template + __aicore__ inline int32_t CAUSAL_CONV1D_CLASS::ReadQueryStartLocValue(int32_t index) const + { + if (tilingData_->queryStartLocUseInt64 != 0) { + const int64_t value = queryStartLocGmInt64.GetValue(index); + if (value < 0 || value > INT32_MAX_VALUE) { + return -1; + } + return static_cast(value); + } + return queryStartLocGmInt32.GetValue(index); + } + + template + __aicore__ inline int64_t CAUSAL_CONV1D_CLASS::ReadCacheIndexValue(int32_t seq) const + { + const int32_t offset = seq * static_cast(tilingData_->cacheIndicesStride); + if (tilingData_->cacheIndicesUseInt64 != 0) { + return cacheIndicesGmInt64.GetValue(offset); + } + return static_cast(cacheIndicesGmInt32.GetValue(offset)); + } + + template + __aicore__ inline bool CAUSAL_CONV1D_CLASS::ReadInitialStateModeValue(int32_t seq) const + { + if (tilingData_->initialStateModeDtype == 2) { + return initialStateModeGmInt64.GetValue(seq) != 0; + } + if (tilingData_->initialStateModeDtype == 1) { + return initialStateModeGmInt32.GetValue(seq) != 0; + } + return initialStateModeGmBool.GetValue(seq); + } + + template + __aicore__ inline int32_t CAUSAL_CONV1D_CLASS::ReadNumAcceptedTokensValue(int32_t seq) const + { + if (tilingData_->numAcceptedTokensUseInt64 != 0) { + const int64_t value = numAcceptedTokensGmInt64.GetValue(seq); + if (value <= 0) { + return 0; + } + if (value > INT32_MAX_VALUE) { + return static_cast(INT32_MAX_VALUE); + } + return static_cast(value); + } + return numAcceptedTokensGmInt32.GetValue(seq); + } + + template + __aicore__ inline bool CAUSAL_CONV1D_CLASS::ResolveSeqCacheIndex(int32_t seq, bool hasCacheIndices, + int32_t &cacheIdx) const + { + cacheIdx = seq; + if (!hasCacheIndices) { + return true; + } + + const int64_t cacheIdx64 = ReadCacheIndexValue(seq); + if (cacheIdx64 == tilingData_->padSlotId) { + return false; + } + if (cacheIdx64 < 0 || cacheIdx64 >= tilingData_->numCacheLines) { + return false; + } + cacheIdx = static_cast(cacheIdx64); + return true; + } + + template + __aicore__ inline bool CAUSAL_CONV1D_CLASS::ResolveSeqHasInit(int32_t seq, bool hasInitialStateMode) const + { + return hasInitialStateMode ? ReadInitialStateModeValue(seq) : false; + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::ProcessDefault() + { + switch (GetSeqTaskWindowMode(tilingData_->inputMode)) { + case SEQ_TASK_WINDOW_MODE_VARLEN: + ProcessDefaultByWindowMode(); + return; + case SEQ_TASK_WINDOW_MODE_DECODE2D: + ProcessDefaultByWindowMode(); + return; + default: + ProcessDefaultByWindowMode(); + return; + } + } + + template + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::ProcessDefaultByWindowMode() + { + const int32_t dim = tilingData_->dim; + const int32_t batch = tilingData_->batch; + const int32_t seqLen = tilingData_->seqLen; + const int32_t baseDim = static_cast(tilingData_->baseDim); + const int32_t baseDimCnt = static_cast(tilingData_->baseDimCnt); + const int32_t width = static_cast(tilingData_->width); + const bool hasCacheIndices = (tilingData_->hasCacheIndices != 0); + const bool hasInit = true; + const bool isSpecDecodingGlobal = IsUpdateSpecDecodingEnabled(); + + const uint32_t blockIdx = GetBlockIdx(); + const uint32_t blockNum = GetBlockNum(); + + if (baseDim <= 0 || baseDimCnt <= 0 || baseDim > MAX_BLOCK_DIM || width < 2 || width > MAX_WIDTH) { + ReleaseEvents(); + return; + } + + const int64_t gridSize = static_cast(batch) * baseDimCnt; + for (int64_t task = static_cast(blockIdx); task < gridSize; task += static_cast(blockNum)) { + const int32_t seq = static_cast(task / baseDimCnt); + const int32_t baseDimIdx = static_cast(task % baseDimCnt); + const int32_t channelStart = baseDimIdx * baseDim; + if (channelStart >= dim) { + continue; + } + const int32_t curBaseDim = (channelStart + baseDim <= dim) ? baseDim : (dim - channelStart); + + int32_t start = 0; + int32_t len = 0; + if (!ResolveSeqTaskWindowByMode(seq, seqLen, start, len)) { + continue; + } + + int32_t cacheIdx = 0; + if (!ResolveSeqCacheIndex(seq, hasCacheIndices, cacheIdx)) { + continue; + } + + int32_t stateTokenOffset = 0; + if (isSpecDecodingGlobal) { + int32_t accepted = ReadNumAcceptedTokensValue(seq); + stateTokenOffset = accepted - 1; + const int32_t maxOffset = static_cast(tilingData_->stateLen - (width - 1)); + if (stateTokenOffset < 0) { + stateTokenOffset = 0; + } else if (stateTokenOffset > maxOffset) { + stateTokenOffset = maxOffset; + } + } + + LoadWeightAndBias(channelStart, curBaseDim); + + InitRing(cacheIdx, hasInit, stateTokenOffset, start, len, channelStart, curBaseDim, dim); + RunSeq(start, len, channelStart, curBaseDim, dim); + + if (isSpecDecodingGlobal) { + DrainTaskMte3(); + WriteBackStateSpec(cacheIdx, hasInit, stateTokenOffset, start, len, channelStart, curBaseDim, dim); + } else { + WriteBackState(cacheIdx, len, channelStart, curBaseDim, dim); + } + + DrainTaskMte3(); + } + } + + template + __aicore__ inline const CausalConv1dTilingData *CAUSAL_CONV1D_CLASS::GetTilingData() const + { + return tilingData_; + } + + template + __aicore__ inline bool CAUSAL_CONV1D_CLASS::HasActivation() const + { + return (tilingData_ != nullptr) && (tilingData_->activationMode != 0); + } + + template + __aicore__ inline bool CAUSAL_CONV1D_CLASS::HasBias() const + { + return (tilingData_ != nullptr) && (tilingData_->hasBias != 0); + } + + template + __aicore__ inline bool CAUSAL_CONV1D_CLASS::IsUpdateMode() const + { + return kIsUpdateMode; + } + + template + __aicore__ inline bool CAUSAL_CONV1D_CLASS::IsFnRollingFastPathEnabled() const + { + return !kIsUpdateMode && (tilingData_ != nullptr) && (kFnExecutionPlan != FN_EXECUTION_PLAN_INVALID) && + (tilingData_->hasNumAcceptedTokens == 0) && !HasBias(); + } + + template + __aicore__ inline bool CAUSAL_CONV1D_CLASS::HasExplicitFnTokenSeqRanges() const + { + return !kIsUpdateMode && (tilingData_ != nullptr) && (tilingData_->inputMode == 0) && + (tilingData_->hasExplicitTokenSeqRanges != 0) && + (tilingData_->explicitTokenSeqRangeCount >= tilingData_->tokenBlockCnt); + } + + template + __aicore__ inline bool CAUSAL_CONV1D_CLASS::IsUpdateSpecDecodingEnabled() const + { + return kIsUpdateMode && (tilingData_->hasNumAcceptedTokens != 0) && (tilingData_->width == 4); + } + + #include "causal_conv1d_fn_tasks.h" + + #undef CAUSAL_CONV1D_CLASS + #undef CAUSAL_CONV1D_TEMPLATE_ARGS + +} // namespace NsCausalConv1d +#endif // CAUSAL_CONV1D_H diff --git a/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_common.h b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_common.h new file mode 100644 index 000000000..57302040f --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_common.h @@ -0,0 +1,66 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file causal_conv1d_common.h + */ + +#ifndef CAUSAL_CONV1D_COMMON_H +#define CAUSAL_CONV1D_COMMON_H + +#include "kernel_operator.h" + +namespace NsCausalConv1dCommon { + +constexpr int32_t MAX_WIDTH = 4; +constexpr int32_t MAX_BLOCK_DIM = 4096; +constexpr int32_t RING_SLOTS = 5; + +__aicore__ inline int32_t SlotCurr(int32_t t) +{ + return (t + 3) % RING_SLOTS; +} + +__aicore__ inline int32_t SlotHist(int32_t t, int32_t i) +{ + return (t + 3 - i) % RING_SLOTS; +} + +__aicore__ inline int32_t SlotPrefetch(int32_t t) +{ + return (t + 4) % RING_SLOTS; +} + +struct CalcBufLayout { + AscendC::LocalTensor weightF; + AscendC::LocalTensor biasF; + AscendC::LocalTensor accF; + AscendC::LocalTensor tmpF; + AscendC::LocalTensor currF; + + __aicore__ inline CalcBufLayout() = default; + + __aicore__ static inline CalcBufLayout FromCalcBuf(AscendC::TBuf &calcBuf) + { + CalcBufLayout layout; + AscendC::LocalTensor calc = calcBuf.template Get(); + layout.weightF = calc; + layout.biasF = calc[MAX_WIDTH * MAX_BLOCK_DIM]; + layout.accF = layout.biasF[MAX_BLOCK_DIM]; + layout.tmpF = layout.accF[MAX_BLOCK_DIM]; + layout.currF = layout.tmpF[MAX_BLOCK_DIM]; + return layout; + } +}; + +} // namespace NsCausalConv1dCommon + +#endif // CAUSAL_CONV1D_COMMON_H diff --git a/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_fn.h b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_fn.h new file mode 100644 index 000000000..dce8e8b13 --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_fn.h @@ -0,0 +1,93 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + + #ifndef CAUSAL_CONV1D_FN_H + #define CAUSAL_CONV1D_FN_H + + #include "causal_conv1d.h" + + namespace NsCausalConv1d { + + template + class CausalConv1dFn : public CausalConv1d { + public: + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR convStates, GM_ADDR queryStartLoc, + GM_ADDR cacheIndices, GM_ADDR initialStateMode, GM_ADDR numAcceptedTokens, GM_ADDR y, + GM_ADDR workspace, const CausalConv1dTilingData *tilingData) + { + (void)numAcceptedTokens; + this->ResetRuntimeState(tilingData); + this->xGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(x)); + this->weightGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(weight)); + this->biasGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(bias)); + this->convStatesGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(convStates)); + if (tilingData->hasQueryStartLoc != 0) { + if (tilingData->queryStartLocUseInt64 != 0) { + this->queryStartLocGmInt64.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(queryStartLoc)); + } else { + this->queryStartLocGmInt32.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(queryStartLoc)); + } + } + if (tilingData->hasCacheIndices != 0) { + if (tilingData->cacheIndicesUseInt64 != 0) { + this->cacheIndicesGmInt64.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(cacheIndices)); + } else { + this->cacheIndicesGmInt32.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(cacheIndices)); + } + } + if (tilingData->hasInitialStateMode != 0) { + if (tilingData->initialStateModeDtype == 2) { + this->initialStateModeGmInt64.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(initialStateMode)); + } else if (tilingData->initialStateModeDtype == 1) { + this->initialStateModeGmInt32.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(initialStateMode)); + } else { + this->initialStateModeGmBool.SetGlobalBuffer(reinterpret_cast<__gm__ bool *>(initialStateMode)); + } + } + this->yGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(y)); + if (tilingData->hasInitStateWorkspace != 0) { + const uint64_t syncElems = + static_cast(GetBlockNum()) * INIT_STATE_SYNCALL_NEED_SIZE; + const uint64_t syncBytes = syncElems * sizeof(int32_t); + const uint64_t workspaceElems = + static_cast(tilingData->batch) * + static_cast(tilingData->width - 1) * + static_cast(tilingData->dim); + this->initStateSyncGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace), syncElems); + auto *workspaceBytes = reinterpret_cast<__gm__ uint8_t *>(workspace); + this->initStateWorkspaceGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(workspaceBytes + syncBytes), + workspaceElems); + } + this->InitSharedBuffersAndEvents(); + } + + __aicore__ inline void Process() + { + this->ProcessVarlenTokenTiled(); + this->ReleaseEvents(); + } + }; + + template + __aicore__ inline void RunCausalConv1dFn(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR convStates, + GM_ADDR queryStartLoc, GM_ADDR cacheIndices, GM_ADDR initialStateMode, + GM_ADDR numAcceptedTokens, GM_ADDR y, GM_ADDR workspace, + const CausalConv1dTilingData *tilingData) + { + CausalConv1dFn op; + op.Init(x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, y, workspace, + tilingData); + op.Process(); + } + + } + + #endif diff --git a/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_fn_tasks.h b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_fn_tasks.h new file mode 100644 index 000000000..9cf2a2c40 --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_fn_tasks.h @@ -0,0 +1,306 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + + #ifndef CAUSAL_CONV1D_FN_TASKS_H + #define CAUSAL_CONV1D_FN_TASKS_H + + struct FnDirectBlockTask { + bool valid = false; + int32_t tokenTileId = 0; + int32_t baseDimIdx = 0; + int32_t tokenStart = 0; + int32_t tokenEnd = 0; + int32_t channelStart = 0; + int32_t baseDimSize = 0; + }; + + __aicore__ inline FnDirectBlockTask ResolveFnDirectBlockTask(int32_t blockIdx, int32_t tokenBlockCnt, int32_t tokenBlockSize, + int32_t cuSeqlen, int32_t baseDimCnt, int32_t baseDim, + int32_t dim) + { + FnDirectBlockTask task; + if (blockIdx < 0 || tokenBlockCnt <= 0 || tokenBlockSize <= 0 || cuSeqlen <= 0 || baseDimCnt <= 0 || baseDim <= 0 || + dim <= 0) { + return task; + } + + const int64_t phase1Grid = static_cast(tokenBlockCnt) * baseDimCnt; + if (phase1Grid <= 0 || static_cast(blockIdx) >= phase1Grid) { + return task; + } + + task.tokenTileId = blockIdx / baseDimCnt; + task.baseDimIdx = blockIdx % baseDimCnt; + task.channelStart = task.baseDimIdx * baseDim; + if (task.channelStart >= dim) { + return task; + } + + task.baseDimSize = (task.channelStart + baseDim <= dim) ? baseDim : (dim - task.channelStart); + task.tokenStart = task.tokenTileId * tokenBlockSize; + if (task.tokenStart >= cuSeqlen) { + return task; + } + + const int32_t tokenEndRaw = task.tokenStart + tokenBlockSize; + task.tokenEnd = (tokenEndRaw <= cuSeqlen) ? tokenEndRaw : cuSeqlen; + if (task.baseDimSize <= 0 || task.tokenEnd <= task.tokenStart) { + return {}; + } + + task.valid = true; + return task; + } + + __aicore__ inline bool IsFnInitStateSnapshotOwnerBlock(const FnDirectBlockTask &task) + { + return task.valid && task.tokenTileId == 0; + } + + template + __aicore__ inline int32_t CAUSAL_CONV1D_CLASS::FindVarlenSeqByToken(int32_t tokenIdx) const + { + int32_t left = 0; + int32_t right = static_cast(tilingData_->batch); + while (left < right) { + const int32_t mid = left + ((right - left) >> 1); + const int32_t endVal = ReadQueryStartLocValue(mid + 1); + if (tokenIdx < endVal) { + right = mid; + } else { + left = mid + 1; + } + } + return left; + } + + template + __aicore__ inline bool CAUSAL_CONV1D_CLASS::ResolveExplicitTokenTileSeqRange(int32_t tokenTileId, int32_t &startSeq, + int32_t &endSeq) const + { + if (!HasExplicitFnTokenSeqRanges() || tokenTileId < 0 || tokenTileId >= tilingData_->explicitTokenSeqRangeCount) { + return false; + } + startSeq = static_cast(tilingData_->tokenTileStartSeq[tokenTileId]); + endSeq = static_cast(tilingData_->tokenTileEndSeq[tokenTileId]); + return (startSeq >= 0) && (endSeq >= startSeq); + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::InitRingSeqSplit(int32_t seq, int32_t cacheIdx, bool hasInit, + int32_t seqStart, int32_t tileStart, int32_t tileLen, + int32_t channelStart, int32_t baseDim, int32_t dim) + { + const int32_t stateLen = tilingData_->stateLen; + const int32_t width = static_cast(tilingData_->width); + const int32_t historyCount = width - 1; + const int32_t ringStart = MAX_WIDTH - width; + const int32_t historyStartTok = tileStart - historyCount; + LocalTensor ring = inBuf.Get(); + bool hasGmHistoryCopy = false; + bool hasVectorInit = false; + const int64_t stateBaseOffset = static_cast(cacheIdx) * stateLen * dim + channelStart; + int64_t xHistoryOffset = static_cast(historyStartTok) * dim + channelStart; + + for (int32_t i = 0; i < ringStart; ++i) { + Duplicate(ring[i * MAX_BLOCK_DIM], static_cast(0), baseDim); + hasVectorInit = true; + } + + for (int32_t i = 0, srcTok = historyStartTok; i < historyCount; ++i, ++srcTok, xHistoryOffset += dim) { + LocalTensor histSlot = ring[(ringStart + i) * MAX_BLOCK_DIM]; + if (srcTok >= seqStart) { + DataCopy(histSlot, xGm[xHistoryOffset], baseDim); + hasGmHistoryCopy = true; + } else if (hasInit) { + const int32_t statePos = srcTok - seqStart + historyCount; + const int64_t stateOffset = stateBaseOffset + static_cast(statePos) * dim; + if (tilingData_->hasInitStateWorkspace != 0) { + const int64_t snapshotOffset = + (static_cast(seq) * historyCount + statePos) * dim + channelStart; + DataCopy(histSlot, initStateWorkspaceGm_[snapshotOffset], baseDim); + } else { + DataCopy(histSlot, convStatesGm[stateOffset], baseDim); + } + hasGmHistoryCopy = true; + } else { + Duplicate(histSlot, static_cast(0), baseDim); + hasVectorInit = true; + } + } + + if (hasGmHistoryCopy) { + SetFlag(stateMte2ToVEvent_); + WaitFlag(stateMte2ToVEvent_); + } + if (hasVectorInit) { + PipeBarrier(); + } + + if (tileLen > 0) { + const int32_t slot0 = SlotCurr(0); + const int64_t xOffset = static_cast(tileStart) * dim + channelStart; + DataCopy(ring[slot0 * MAX_BLOCK_DIM], xGm[xOffset], baseDim); + SetFlag(inputMte2ToVEvent_[slot0]); + } + + if (tileLen > 1) { + SetFlag(inputVToMte2Event_); + } + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::ProcessFnChunk(int32_t seq, int32_t cacheIdx, bool hasInit, + int32_t seqStart, int32_t seqLen, int32_t chunkStart, + int32_t chunkLen, int32_t channelStart, int32_t baseDim, + int32_t dim) + { + LoadWeightAndBias(channelStart, baseDim); + InitRingSeqSplit(seq, cacheIdx, hasInit, seqStart, chunkStart, chunkLen, channelStart, baseDim, dim); + + RunSeq(chunkStart, chunkLen, channelStart, baseDim, dim); + + MaybeWriteBackSeqSplitTailChunk(chunkStart, chunkLen, seqStart, seqLen, cacheIdx, channelStart, baseDim, dim); + DrainTaskMte3(); + } + + template + __aicore__ inline void + CAUSAL_CONV1D_CLASS::MaybeWriteBackSeqSplitTailChunk(int32_t chunkStart, int32_t chunkLen, int32_t seqStart, + int32_t seqLen, int32_t cacheIdx, int32_t channelStart, + int32_t baseDim, int32_t dim) + { + if (chunkStart + chunkLen != seqStart + seqLen) { + return; + } + + DrainTaskMte3(); + WriteBackState(cacheIdx, chunkLen, channelStart, baseDim, dim); + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::PrefetchInitStatesToWorkspace(int32_t channelStart, int32_t baseDimSize) + { + if (tilingData_->hasInitStateWorkspace == 0) { + return; + } + + const int32_t dim = tilingData_->dim; + const int32_t historyCount = static_cast(tilingData_->width - 1); + const int32_t batch = tilingData_->batch; + const bool hasCacheIndices = (tilingData_->hasCacheIndices != 0); + const bool hasInitialStateMode = (tilingData_->hasInitialStateMode != 0); + LocalTensor tmpBuf = inBuf.Get()[0 * MAX_BLOCK_DIM]; + + for (int32_t seq = 0; seq < batch; ++seq) { + if (!ResolveSeqHasInit(seq, hasInitialStateMode)) { + continue; + } + + int32_t cacheIdx = 0; + if (!ResolveSeqCacheIndex(seq, hasCacheIndices, cacheIdx)) { + continue; + } + + const int64_t stateBaseOffset = static_cast(cacheIdx) * tilingData_->stateLen * dim + channelStart; + const int64_t snapshotBaseOffset = static_cast(seq) * historyCount * dim + channelStart; + for (int32_t statePos = 0; statePos < historyCount; ++statePos) { + const int64_t stateOffset = stateBaseOffset + static_cast(statePos) * dim; + const int64_t snapshotOffset = snapshotBaseOffset + static_cast(statePos) * dim; + DataCopy(tmpBuf, convStatesGm[stateOffset], baseDimSize); + SetFlag(initSnapshotMte2ToMte3Event_); + WaitFlag(initSnapshotMte2ToMte3Event_); + DataCopy(initStateWorkspaceGm_[snapshotOffset], tmpBuf, baseDimSize); + SetFlag(initSnapshotMte3ToMte2Event_); + WaitFlag(initSnapshotMte3ToMte2Event_); + } + } + } + + template + __aicore__ inline void CAUSAL_CONV1D_CLASS::ProcessVarlenTokenTiled() + { + const int32_t dim = tilingData_->dim; + const int32_t batch = tilingData_->batch; + const int32_t seqLen = tilingData_->seqLen; + const int32_t cuSeqlen = tilingData_->cuSeqlen; + const int32_t baseDim = static_cast(tilingData_->baseDim); + const int32_t baseDimCnt = static_cast(tilingData_->baseDimCnt); + const int32_t tokenBlockSize = static_cast(tilingData_->tokenBlockSize); + const int32_t tokenBlockCnt = static_cast(tilingData_->tokenBlockCnt); + const bool hasCacheIndices = (tilingData_->hasCacheIndices != 0); + const bool hasInitialStateMode = (tilingData_->hasInitialStateMode != 0); + const bool isVarlenMode = (tilingData_->inputMode == 0); + + const int32_t blockIdx = static_cast(GetBlockIdx()); + const auto blockTask = ResolveFnDirectBlockTask(blockIdx, tokenBlockCnt, tokenBlockSize, cuSeqlen, baseDimCnt, + baseDim, dim); + if (tilingData_->hasInitStateWorkspace != 0) { + if (IsFnInitStateSnapshotOwnerBlock(blockTask)) { + PrefetchInitStatesToWorkspace(blockTask.channelStart, blockTask.baseDimSize); + } + SyncAll(); + } + if (!blockTask.valid) { + return; + } + + int32_t seq = 0; + int32_t seqUpperBound = batch; + if (isVarlenMode) { + if (!ResolveExplicitTokenTileSeqRange(blockTask.tokenTileId, seq, seqUpperBound)) { + seq = FindVarlenSeqByToken(blockTask.tokenStart); + } + } else { + seq = (seqLen > 0) ? (blockTask.tokenStart / seqLen) : 0; + } + + int32_t cursor = blockTask.tokenStart; + while (cursor < blockTask.tokenEnd && seq < seqUpperBound) { + int32_t seqStart = 0; + int32_t curSeqLen = 0; + if (!ResolveSeqTaskWindow(seq, tilingData_->inputMode, seqLen, seqStart, curSeqLen)) { + ++seq; + continue; + } + const int32_t curSeqEnd = seqStart + curSeqLen; + if (cursor < seqStart) { + cursor = seqStart; + } + if (cursor >= curSeqEnd) { + ++seq; + continue; + } + + const int32_t tileEnd = (blockTask.tokenEnd <= curSeqEnd) ? blockTask.tokenEnd : curSeqEnd; + const int32_t tileLen = tileEnd - cursor; + if (tileLen <= 0) { + ++seq; + continue; + } + + int32_t cacheIdx = 0; + if (!ResolveSeqCacheIndex(seq, hasCacheIndices, cacheIdx)) { + cursor = tileEnd; + ++seq; + continue; + } + + const bool hasInit = ResolveSeqHasInit(seq, hasInitialStateMode); + ProcessFnChunk(seq, cacheIdx, hasInit, seqStart, curSeqLen, cursor, tileLen, blockTask.channelStart, + blockTask.baseDimSize, dim); + + cursor = tileEnd; + ++seq; + } + } + + #endif diff --git a/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_tiling_data.h b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_tiling_data.h new file mode 100644 index 000000000..6e7ae889b --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_tiling_data.h @@ -0,0 +1,68 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file causal_conv1d_tiling_data.h + */ + +#ifndef CAUSAL_CONV1D_TILING_DATA_H_ +#define CAUSAL_CONV1D_TILING_DATA_H_ + +#include + +enum FnExecutionPlan : int64_t { + FN_EXECUTION_PLAN_INVALID = 0, + FN_EXECUTION_PLAN_CUTBS = 1, + FN_EXECUTION_PLAN_CUTBSD = 2, +}; + +inline constexpr int64_t ResolveFnExecutionPlan(int64_t baseDimCnt) +{ + return (baseDimCnt <= 0) ? FN_EXECUTION_PLAN_INVALID + : (baseDimCnt <= 1) ? FN_EXECUTION_PLAN_CUTBS + : FN_EXECUTION_PLAN_CUTBSD; +} + + +struct CausalConv1dTilingData { + int64_t dim; + int64_t cuSeqlen; + int64_t seqLen; + int64_t inputMode; + + int64_t width; + + int64_t stateLen; + int64_t numCacheLines; + int64_t batch; + int64_t activationMode; + int64_t padSlotId; + int64_t hasBias; + int64_t baseDim; + int64_t baseDimCnt; + int64_t hasNumAcceptedTokens; + int64_t hasQueryStartLoc; + int64_t hasCacheIndices; + int64_t hasInitialStateMode; + int64_t queryStartLocUseInt64; + int64_t cacheIndicesStride; + int64_t cacheIndicesUseInt64; + int64_t initialStateModeDtype; + int64_t numAcceptedTokensUseInt64; + int64_t tokenBlockSize; + int64_t tokenBlockCnt; + int64_t hasExplicitTokenSeqRanges; + int64_t explicitTokenSeqRangeCount; + int64_t tokenTileStartSeq[128]; + int64_t tokenTileEndSeq[128]; + int64_t hasInitStateWorkspace; +}; +#endif // CAUSAL_CONV1D_TILING_DATA_H_ diff --git a/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_tiling_key.h b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_tiling_key.h new file mode 100644 index 000000000..7a2120224 --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_tiling_key.h @@ -0,0 +1,66 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file causal_conv1d_tiling_key.h + * \brief causal_conv1d tiling key declare + */ + +#ifndef __CAUSAL_CONV1D_TILING_KEY_H__ +#define __CAUSAL_CONV1D_TILING_KEY_H__ + +#include "causal_conv1d_tiling_data.h" +#include "ascendc/host_api/tiling/template_argument.h" + +#define CAUSAL_CONV1D_TPL_RUN_MODE_FN 0 +#define CAUSAL_CONV1D_TPL_RUN_MODE_UPDATE 1 +#define CAUSAL_CONV1D_TPL_WIDTH_RUNTIME 0 +#define CAUSAL_CONV1D_TPL_WIDTH_2 1 +#define CAUSAL_CONV1D_TPL_WIDTH_3 2 +#define CAUSAL_CONV1D_TPL_WIDTH_4 3 +#define CAUSAL_CONV1D_TPL_FN_PLAN_INVALID 0 +#define CAUSAL_CONV1D_TPL_FN_PLAN_CUTBS 1 +#define CAUSAL_CONV1D_TPL_FN_PLAN_CUTBSD 2 +ASCENDC_TPL_ARGS_DECL(CausalConv1d, + ASCENDC_TPL_UINT_DECL(runModeKey, 1, ASCENDC_TPL_UI_LIST, CAUSAL_CONV1D_TPL_RUN_MODE_FN, + CAUSAL_CONV1D_TPL_RUN_MODE_UPDATE), + ASCENDC_TPL_UINT_DECL(widthKey, 2, ASCENDC_TPL_UI_LIST, CAUSAL_CONV1D_TPL_WIDTH_RUNTIME, + CAUSAL_CONV1D_TPL_WIDTH_2, CAUSAL_CONV1D_TPL_WIDTH_3, + CAUSAL_CONV1D_TPL_WIDTH_4), + ASCENDC_TPL_UINT_DECL(fnPlanKey, 2, ASCENDC_TPL_UI_LIST, CAUSAL_CONV1D_TPL_FN_PLAN_INVALID, + CAUSAL_CONV1D_TPL_FN_PLAN_CUTBS, CAUSAL_CONV1D_TPL_FN_PLAN_CUTBSD)); + +#define CAUSAL_CONV1D_TPL_SEL_ENTRY(RUN_MODE, WIDTH, FN_PLAN) \ + ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_UINT_SEL(runModeKey, ASCENDC_TPL_UI_LIST, RUN_MODE), \ + ASCENDC_TPL_UINT_SEL(widthKey, ASCENDC_TPL_UI_LIST, WIDTH), \ + ASCENDC_TPL_UINT_SEL(fnPlanKey, ASCENDC_TPL_UI_LIST, FN_PLAN), \ + ASCENDC_TPL_TILING_STRUCT_SEL(CausalConv1dTilingData)) + +// Keep entries in encoded tiling-key order: real-device sub-kernel dispatch is sensitive to declaration order. +ASCENDC_TPL_SEL( + CAUSAL_CONV1D_TPL_SEL_ENTRY(CAUSAL_CONV1D_TPL_RUN_MODE_UPDATE, CAUSAL_CONV1D_TPL_WIDTH_RUNTIME, + CAUSAL_CONV1D_TPL_FN_PLAN_INVALID), + CAUSAL_CONV1D_TPL_SEL_ENTRY(CAUSAL_CONV1D_TPL_RUN_MODE_FN, CAUSAL_CONV1D_TPL_WIDTH_2, + CAUSAL_CONV1D_TPL_FN_PLAN_CUTBS), + CAUSAL_CONV1D_TPL_SEL_ENTRY(CAUSAL_CONV1D_TPL_RUN_MODE_FN, CAUSAL_CONV1D_TPL_WIDTH_3, + CAUSAL_CONV1D_TPL_FN_PLAN_CUTBS), + CAUSAL_CONV1D_TPL_SEL_ENTRY(CAUSAL_CONV1D_TPL_RUN_MODE_FN, CAUSAL_CONV1D_TPL_WIDTH_4, + CAUSAL_CONV1D_TPL_FN_PLAN_CUTBS), + CAUSAL_CONV1D_TPL_SEL_ENTRY(CAUSAL_CONV1D_TPL_RUN_MODE_FN, CAUSAL_CONV1D_TPL_WIDTH_2, + CAUSAL_CONV1D_TPL_FN_PLAN_CUTBSD), + CAUSAL_CONV1D_TPL_SEL_ENTRY(CAUSAL_CONV1D_TPL_RUN_MODE_FN, CAUSAL_CONV1D_TPL_WIDTH_3, + CAUSAL_CONV1D_TPL_FN_PLAN_CUTBSD), + CAUSAL_CONV1D_TPL_SEL_ENTRY(CAUSAL_CONV1D_TPL_RUN_MODE_FN, CAUSAL_CONV1D_TPL_WIDTH_4, + CAUSAL_CONV1D_TPL_FN_PLAN_CUTBSD)); + +#undef CAUSAL_CONV1D_TPL_SEL_ENTRY + +#endif // __CAUSAL_CONV1D_TILING_KEY_H__ diff --git a/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_update.h b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_update.h new file mode 100644 index 000000000..f63f94fae --- /dev/null +++ b/csrc/ascend/moe/causal_conv1d/op_kernel/causal_conv1d_update.h @@ -0,0 +1,91 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING + * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CAUSAL_CONV1D_UPDATE_H +#define CAUSAL_CONV1D_UPDATE_H + +#include "causal_conv1d.h" + +namespace NsCausalConv1d { + +template +class CausalConv1dUpdate + : public CausalConv1d { +public: + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR convStates, GM_ADDR queryStartLoc, + GM_ADDR cacheIndices, GM_ADDR, GM_ADDR numAcceptedTokens, GM_ADDR y, GM_ADDR workspace, + const CausalConv1dTilingData *tilingData) + { + (void)workspace; + this->ResetRuntimeState(tilingData); + this->xGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(x)); + this->weightGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(weight)); + this->biasGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(bias)); + this->convStatesGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(convStates)); + if (tilingData->hasQueryStartLoc != 0) { + if (tilingData->queryStartLocUseInt64 != 0) { + this->queryStartLocGmInt64.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(queryStartLoc)); + } else { + this->queryStartLocGmInt32.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(queryStartLoc)); + } + } + if (tilingData->hasCacheIndices != 0) { + if (tilingData->cacheIndicesUseInt64 != 0) { + this->cacheIndicesGmInt64.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(cacheIndices)); + } else { + this->cacheIndicesGmInt32.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(cacheIndices)); + } + } + if (tilingData->hasNumAcceptedTokens != 0) { + if (tilingData->numAcceptedTokensUseInt64 != 0) { + this->numAcceptedTokensGmInt64.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(numAcceptedTokens)); + } else { + this->numAcceptedTokensGmInt32.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(numAcceptedTokens)); + } + } + this->yGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(y)); + this->InitSharedBuffersAndEvents(); + } + + __aicore__ inline void Process() + { + const CausalConv1dTilingData *tilingData = this->GetTilingData(); + const int32_t dim = tilingData->dim; + const int32_t baseDimCnt = static_cast(tilingData->baseDimCnt); + const int32_t width = static_cast(tilingData->width); + const int32_t baseDim = static_cast(tilingData->baseDim); + if (baseDim <= 0 || baseDimCnt <= 0 || baseDim > MAX_BLOCK_DIM || width < 2 || width > MAX_WIDTH || dim <= 0 || + tilingData->batch <= 0) { + this->ReleaseEvents(); + return; + } + + this->ProcessDefault(); + this->ReleaseEvents(); + } +}; + +template +__aicore__ inline void RunCausalConv1dUpdate(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR convStates, + GM_ADDR queryStartLoc, GM_ADDR cacheIndices, GM_ADDR initialStateMode, + GM_ADDR numAcceptedTokens, GM_ADDR y, GM_ADDR workspace, + const CausalConv1dTilingData *tilingData) +{ + CausalConv1dUpdate op; + op.Init(x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, y, workspace, + tilingData); + op.Process(); +} + +} // namespace NsCausalConv1d + +#endif // CAUSAL_CONV1D_UPDATE_H diff --git a/csrc/ascend/moe/chunk_fwd_o/CMakeLists.txt b/csrc/ascend/moe/chunk_fwd_o/CMakeLists.txt new file mode 100644 index 000000000..86b308249 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_fwd_o/op_host/CMakeLists.txt b/csrc/ascend/moe/chunk_fwd_o/op_host/CMakeLists.txt new file mode 100644 index 000000000..016cb0b95 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_host/CMakeLists.txt @@ -0,0 +1,32 @@ +set(CURRENT_CMAKE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) +set(CATLASS_INCLUDE_DIR "${CMAKE_SOURCE_DIR}/third_party/catlass/include") +get_filename_component(CATLASS_INCLUDE_DIR_ABS ${CATLASS_INCLUDE_DIR} ABSOLUTE) + +set(COMMON_KERNEL_UTILS_DIR "${CMAKE_SOURCE_DIR}/moe/common") +get_filename_component(COMMON_KERNEL_UTILS_DIR_ABS ${COMMON_KERNEL_UTILS_DIR} ABSOLUTE) + +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnExc PRIVATE + chunk_fwd_o_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME ChunkFwdO + OPTIONS + --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -I${CATLASS_INCLUDE_DIR_ABS} + -I${COMMON_KERNEL_UTILS_DIR_ABS} +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE chunk_fwd_o ACLNNTYPE aclnn_exclude) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ${CATLASS_INCLUDE_DIR_ABS} + ) +endif() diff --git a/csrc/ascend/moe/chunk_fwd_o/op_host/chunk_fwd_o_def.cpp b/csrc/ascend/moe/chunk_fwd_o/op_host/chunk_fwd_o_def.cpp new file mode 100644 index 000000000..13e9986f3 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_host/chunk_fwd_o_def.cpp @@ -0,0 +1,104 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +/*! + *\file chunk_fwd_o_def.cpp + *\brief + */ + +#include "register/op_def_registry.h" + +namespace ops { + +class ChunkFwdO : public OpDef { + +public: + explicit ChunkFwdO(const char *name) : OpDef(name) + { + // Define inputs + this->Input("q") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + this->Input("k") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + this->Input("v") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + this->Input("h") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + this->Input("g") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + this->Input("cu_seqlens") + .ParamType(OPTIONAL) + .ValueDepend(OPTIONAL) + .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + this->Input("chunk_offsets") + .ParamType(OPTIONAL) + .ValueDepend(OPTIONAL) + .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + this->Output("o") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + this->Attr("scale").AttrType(REQUIRED).Float(1.0); + this->Attr("chunk_size").AttrType(REQUIRED).Int(64); + + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("prebuildPattern.value", "Opaque") + .ExtendCfgInfo("coreType.value", "AiCore") + .ExtendCfgInfo("jitCompile.flag", "static_false,dynamic_false"); + + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + this->AICore().AddConfig("ascend310p", aicore_config); + } +}; + +OP_ADD(ChunkFwdO); + +} // namespace ops diff --git a/csrc/ascend/moe/chunk_fwd_o/op_host/chunk_fwd_o_tiling.cpp b/csrc/ascend/moe/chunk_fwd_o/op_host/chunk_fwd_o_tiling.cpp new file mode 100644 index 000000000..072189109 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_host/chunk_fwd_o_tiling.cpp @@ -0,0 +1,157 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +/*! + * \file chunk_fwd_o_tiling.cpp + * \brief + */ + +#include "chunk_fwd_o_tiling.h" +#include +#include "../tiling_base/data_copy_transpose_tiling.h" +#include "../tiling_base/tiling_templates_registry.h" + +namespace optiling { +static constexpr size_t INPUT_Q_IDX = 0; +static constexpr size_t INPUT_K_IDX = 1; +static constexpr size_t INPUT_V_IDX = 2; +static constexpr size_t INPUT_H_IDX = 3; +static constexpr size_t INPUT_G_IDX = 4; +static constexpr size_t INPUT_SEQLENS_IDX = 5; +static constexpr size_t INPUT_CHUNK_INDICES_IDX = 6; + +static constexpr size_t ATTR_SCALE_IDX = 0; +static constexpr size_t ATTR_CHUNK_SIZE_IDX = 1; + +static constexpr size_t DIM_BATCH = 0; +static constexpr size_t DIM_HEAD_NUM = 1; +static constexpr size_t DIM_SEQLEN = 2; +static constexpr size_t DIM_HEAD_DIM = 3; + + +static void ChunkFwdOTilingDataPrint(gert::TilingContext *context, ChunkFwdOTilingData &tiling) +{ + auto nodeName = context->GetNodeName(); + OP_LOGD(nodeName, ">>>>>>>>>>>>>>> Start to print ChunkFwdO tiling data <<<<<<<<<<<<<<<<"); + OP_LOGD(nodeName, "=== batch: %ld", tiling.get_shapeBatch()); + OP_LOGD(nodeName, "=== seqlen: %ld", tiling.get_seqlen()); + OP_LOGD(nodeName, "=== kNumHead: %ld", tiling.get_kNumHead()); + OP_LOGD(nodeName, "=== vNumHead: %ld", tiling.get_vNumHead()); + OP_LOGD(nodeName, "=== kHeadDim: %ld", tiling.get_kHeadDim()); + OP_LOGD(nodeName, "=== vHeadDim: %ld", tiling.get_vHeadDim()); + OP_LOGD(nodeName, "=== chunkSize: %ld", tiling.get_chunkSize()); + OP_LOGD(nodeName, "=== dataType: %ld", tiling.get_dataType()); + OP_LOGD(nodeName, "=== isVariedLen: %ld", tiling.get_isVariedLen()); + OP_LOGD(nodeName, "=== tokenBatch: %f", tiling.get_tokenBatch()); + OP_LOGD(nodeName, ">>>>>>>>>>>>>>> Print ChunkFwdO tiling data end <<<<<<<<<<<<<<<<"); +} + +ge::graphStatus Tiling4ChunkFwdO(gert::TilingContext *context) +{ + OP_LOGD(context->GetNodeName(), "Tiling4ChunkFwdO start."); + ChunkFwdOTilingData tiling; + + gert::Shape qStorageShape = context->GetOptionalInputShape(INPUT_Q_IDX)->GetStorageShape(); + gert::Shape vStorageShape = context->GetOptionalInputShape(INPUT_V_IDX)->GetStorageShape(); + + int64_t seqlen = qStorageShape.GetDim(DIM_SEQLEN); + int64_t kNumHead = qStorageShape.GetDim(DIM_HEAD_NUM); + int64_t vNumHead = vStorageShape.GetDim(DIM_HEAD_NUM); + int64_t kHeadDim = qStorageShape.GetDim(DIM_HEAD_DIM); + int64_t vHeadDim = vStorageShape.GetDim(DIM_HEAD_DIM); + int64_t isVariedLen, shapeBatch, tokenBatch; + + auto cuSeqlensTensor = context->GetOptionalInputTensor(INPUT_SEQLENS_IDX); + if (cuSeqlensTensor == nullptr) { + isVariedLen = false; + shapeBatch = qStorageShape.GetDim(DIM_BATCH); + tokenBatch = 1; + } else { + isVariedLen = true; + shapeBatch = 1; + tokenBatch = cuSeqlensTensor->GetStorageShape().GetDim(DIM_BATCH) - 1; + } + + auto attrPtr = context->GetAttrs(); + float scale = *(attrPtr->GetAttrPointer(ATTR_SCALE_IDX)); + int64_t chunkSize = *(attrPtr->GetAttrPointer(ATTR_CHUNK_SIZE_IDX)); + + auto dtype = context->GetInputTensor(0)->GetDataType(); + uint64_t dataType = dtype == ge::DT_BF16 ? 1 : 0; + + auto gDType = context->GetOptionalInputTensor(INPUT_G_IDX)->GetDataType(); + int64_t gDataType = 2; + if (gDType == ge::DT_BF16) { + gDataType = 1; + } else if (gDType == ge::DT_FLOAT16) { + gDataType = 0; + } + + const auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + uint32_t aicCoreNum = ascendcPlatform.GetCoreNumAic(); + context->SetBlockDim(aicCoreNum); + + constexpr size_t WORKSPACE_RSV_BYTE = 16 * 1024 * 1024; + constexpr size_t GM_ALIGN = 512; + int64_t pingpongStages = 2; + + size_t workspaceOffset = ascendcPlatform.GetLibApiWorkSpaceSize(); + workspaceOffset += WORKSPACE_RSV_BYTE; + + tiling.set_vWorkspaceOffset(workspaceOffset); + workspaceOffset += (aicCoreNum * chunkSize * vHeadDim * sizeof(float) * pingpongStages + GM_ALIGN) / GM_ALIGN * GM_ALIGN; + + tiling.set_hWorkspaceOffset(workspaceOffset); + workspaceOffset += (aicCoreNum * chunkSize * vHeadDim * sizeof(float) * pingpongStages + GM_ALIGN) / GM_ALIGN * GM_ALIGN; + + tiling.set_attnWorkspaceOffset(workspaceOffset); + workspaceOffset += (aicCoreNum * chunkSize * chunkSize * sizeof(float) * pingpongStages + GM_ALIGN) / GM_ALIGN * GM_ALIGN; + + tiling.set_aftermaskWorkspaceOffset(workspaceOffset); + workspaceOffset += (aicCoreNum * chunkSize * chunkSize * sizeof(float) * pingpongStages + GM_ALIGN) / GM_ALIGN * GM_ALIGN; + + tiling.set_maskWorkspaceOffset(workspaceOffset); + workspaceOffset += (chunkSize * chunkSize + GM_ALIGN) / GM_ALIGN * GM_ALIGN; + + workspaceOffset += WORKSPACE_RSV_BYTE; + size_t *currentWorkspace = context->GetWorkspaceSizes(1); + currentWorkspace[0] = (workspaceOffset - 0); + + tiling.set_shapeBatch(shapeBatch); + tiling.set_seqlen(seqlen); + tiling.set_kNumHead(kNumHead); + tiling.set_vNumHead(vNumHead); + tiling.set_kHeadDim(kHeadDim); + tiling.set_vHeadDim(vHeadDim); + tiling.set_scale(scale); + tiling.set_chunkSize(chunkSize); + tiling.set_isVariedLen(isVariedLen); + tiling.set_tokenBatch(tokenBatch); + tiling.set_dataType(dataType); + tiling.set_gDataType(gDataType); + + tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); + + ChunkFwdOTilingDataPrint(context, tiling); + OP_LOGD(context->GetNodeName(), "Tiling4ChunkFwdO end."); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingPrepareForChunkFwdO(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(ChunkFwdO) + .Tiling(Tiling4ChunkFwdO) + .TilingParse(TilingPrepareForChunkFwdO); + +} // namespace optiling diff --git a/csrc/ascend/moe/chunk_fwd_o/op_host/chunk_fwd_o_tiling.h b/csrc/ascend/moe/chunk_fwd_o/op_host/chunk_fwd_o_tiling.h new file mode 100644 index 000000000..e73706eca --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_host/chunk_fwd_o_tiling.h @@ -0,0 +1,46 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +/*! + * \file chunk_fwd_o_tiling.h + * \brief + */ + +#pragma once + +#include +#include +#include + +namespace optiling { + +BEGIN_TILING_DATA_DEF(ChunkFwdOTilingData) +TILING_DATA_FIELD_DEF(int64_t, shapeBatch); +TILING_DATA_FIELD_DEF(int64_t, seqlen); +TILING_DATA_FIELD_DEF(int64_t, kNumHead); +TILING_DATA_FIELD_DEF(int64_t, vNumHead); +TILING_DATA_FIELD_DEF(int64_t, kHeadDim); +TILING_DATA_FIELD_DEF(int64_t, vHeadDim); +TILING_DATA_FIELD_DEF(int64_t, chunkSize); +TILING_DATA_FIELD_DEF(int64_t, isVariedLen); +TILING_DATA_FIELD_DEF(int64_t, tokenBatch); +TILING_DATA_FIELD_DEF(int64_t, dataType); +TILING_DATA_FIELD_DEF(int64_t, gDataType); +TILING_DATA_FIELD_DEF(int64_t, vWorkspaceOffset); +TILING_DATA_FIELD_DEF(int64_t, hWorkspaceOffset); +TILING_DATA_FIELD_DEF(int64_t, attnWorkspaceOffset); +TILING_DATA_FIELD_DEF(int64_t, aftermaskWorkspaceOffset); +TILING_DATA_FIELD_DEF(int64_t, maskWorkspaceOffset); +TILING_DATA_FIELD_DEF(float, scale); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(ChunkFwdO, ChunkFwdOTilingData) + +struct ChunkFwdOCompileInfo {}; +} // namespace optiling diff --git a/csrc/ascend/moe/chunk_fwd_o/op_host/op_api/aclnn_chunk_fwd_o.cpp b/csrc/ascend/moe/chunk_fwd_o/op_host/op_api/aclnn_chunk_fwd_o.cpp new file mode 100644 index 000000000..1f61d1ce5 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_host/op_api/aclnn_chunk_fwd_o.cpp @@ -0,0 +1,162 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#include "aclnn_chunk_fwd_o.h" +#include "chunk_fwd_o.h" +#include +#include + +#include "aclnn_kernels/transdata.h" +#include "aclnn_kernels/contiguous.h" +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "opdev/common_types.h" +#include "opdev/data_type_utils.h" +#include "opdev/format_utils.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/platform.h" +#include "opdev/shape_utils.h" +#include "opdev/tensor_view_utils.h" +#include "opdev/make_op_executor.h" + + +using namespace op; + +#ifdef __cplusplus +extern "C" { +#endif + +struct ChunkFwdOParams { + const aclTensor *q = nullptr; + const aclTensor *k = nullptr; + const aclTensor *v = nullptr; + const aclTensor *h = nullptr; + const aclTensor *g = nullptr; + const aclIntArray *cuSeqlensOptional = nullptr; + const aclIntArray *chunkOffsetsOptional = nullptr; + double scale = 1.0; + int64_t chunkSize = 64; + const aclTensor *oOut = nullptr; +}; + +static aclnnStatus CheckNotNull(ChunkFwdOParams params) +{ + CHECK_COND(params.q != nullptr, ACLNN_ERR_PARAM_NULLPTR, "q must not be nullptr."); + CHECK_COND(params.k != nullptr, ACLNN_ERR_PARAM_NULLPTR, "k must not be nullptr."); + CHECK_COND(params.v != nullptr, ACLNN_ERR_PARAM_NULLPTR, "v must not be nullptr."); + CHECK_COND(params.h != nullptr, ACLNN_ERR_PARAM_NULLPTR, "h must not be nullptr."); + CHECK_COND(params.g != nullptr, ACLNN_ERR_PARAM_NULLPTR, "g must not be nullptr."); + + CHECK_COND(params.oOut != nullptr, ACLNN_ERR_PARAM_NULLPTR, "oOut must not be nullptr."); + return ACLNN_SUCCESS; +} + +static aclnnStatus CheckFormat(ChunkFwdOParams params) +{ + return ACLNN_SUCCESS; +} + +static aclnnStatus CheckShape(ChunkFwdOParams params) +{ + return ACLNN_SUCCESS; +} + +static aclnnStatus DataContiguous(const aclTensor *&tensor, aclOpExecutor *executor) +{ + tensor = l0op::Contiguous(tensor, executor); + CHECK_RET(tensor != nullptr, ACLNN_ERR_INNER_NULLPTR); + return ACLNN_SUCCESS; +} + +static aclnnStatus ParamsDataContiguous(ChunkFwdOParams ¶ms, aclOpExecutor *executorPtr) +{ + CHECK_COND(DataContiguous(params.q, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "Contiguous q failed."); + CHECK_COND(DataContiguous(params.k, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "Contiguous k failed."); + CHECK_COND(DataContiguous(params.v, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "Contiguous v failed."); + CHECK_COND(DataContiguous(params.h, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "Contiguous h failed."); + CHECK_COND(DataContiguous(params.g, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "Contiguous g failed."); + + return ACLNN_SUCCESS; +} + +static aclnnStatus CheckDtype(ChunkFwdOParams params) +{ + return ACLNN_SUCCESS; +} + +static aclnnStatus CheckParams(ChunkFwdOParams params) +{ + CHECK_RET(CheckNotNull(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + CHECK_RET(CheckFormat(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + CHECK_RET(CheckShape(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + CHECK_RET(CheckDtype(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + return ACLNN_SUCCESS; +} + +aclnnStatus aclnnChunkFwdOGetWorkspaceSize( + const aclTensor *q, + const aclTensor *k, + const aclTensor *v, + const aclTensor *h, + const aclTensor *g, + const aclIntArray *cuSeqlensOptional, + const aclIntArray *chunkOffsetsOptional, + double scale, + int64_t chunkSize, + const aclTensor *oOut, + uint64_t *workspaceSize, + aclOpExecutor **executor) +{ + ChunkFwdOParams params{q, k, v, h, g, cuSeqlensOptional, chunkOffsetsOptional, scale, chunkSize, oOut}; + // Standard syntax, Check parameters. + L2_DFX_PHASE_1(aclnnChunkFwdO, DFX_IN(q, k, v, h, g, cuSeqlensOptional, chunkOffsetsOptional), + DFX_OUT(oOut)); + // 固定写法，创建OpExecutor + auto uniqueExecutor = CREATE_EXECUTOR(); + CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + auto executorPtr = uniqueExecutor.get(); + // 固定写法，参数检查 + auto ret = CheckParams(params); + CHECK_RET(ret == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + CHECK_COND(ParamsDataContiguous(params, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "ParamsDataContiguous failed."); + auto result = l0op::ChunkFwdO(params.q, params.k, params.v, params.h, params.g, params.cuSeqlensOptional, params.chunkOffsetsOptional, params.scale, params.chunkSize, params.oOut, executorPtr); + CHECK_RET(result[0] != nullptr, ACLNN_ERR_PARAM_NULLPTR); + + // If the output tensor is non-contiguous, convert the calculated contiguous tensor to non-contiguous. + auto viewCopyResult = l0op::ViewCopy(result[0], params.oOut, executorPtr); + CHECK_RET(viewCopyResult != nullptr, ACLNN_ERR_INNER_NULLPTR); + + // Standard syntax, get the size of workspace needed during computation. + *workspaceSize = uniqueExecutor->GetWorkspaceSize(); + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; +} + + +aclnnStatus aclnnChunkFwdO(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream) +{ + L2_DFX_PHASE_2(aclnnChunkFwdO); + CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER, + "This is an error in ChunkFwdO launch aicore."); + return ACLNN_SUCCESS; +} + + +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/moe/chunk_fwd_o/op_host/op_api/aclnn_chunk_fwd_o.h b/csrc/ascend/moe/chunk_fwd_o/op_host/op_api/aclnn_chunk_fwd_o.h new file mode 100644 index 000000000..040f1bfee --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_host/op_api/aclnn_chunk_fwd_o.h @@ -0,0 +1,65 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_API_INC_ACLNN_CHUNK_FWD_O_H +#define OP_API_INC_ACLNN_CHUNK_FWD_O_H +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/* function: aclnnChunkFwdOGetWorkspaceSize + * parameters : + * q : required + * k : required + * v : required + * h : required + * g : required + * cuSeqlensOptional : optional + * chunkOffsetsOptional : optional + * scale : required + * chunkSize : required + * oOut : required + * workspaceSize : size of workspace(output). + * executor : executor context(output). + */ +__attribute__((visibility("default"))) +aclnnStatus aclnnChunkFwdOGetWorkspaceSize( + const aclTensor *q, + const aclTensor *k, + const aclTensor *v, + const aclTensor *h, + const aclTensor *g, + const aclIntArray *cuSeqlensOptional, + const aclIntArray *chunkOffsetsOptional, + double scale, + int64_t chunkSize, + const aclTensor *oOut, + uint64_t *workspaceSize, + aclOpExecutor **executor); + +/* function: aclnnChunkFwdO + * parameters : + * workspace : workspace memory addr(input). + * workspaceSize : size of workspace(input). + * executor : executor context(input). + * stream : acl stream. + */ +__attribute__((visibility("default"))) +aclnnStatus aclnnChunkFwdO( + void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream); +#ifdef __cplusplus +} +#endif + +#endif diff --git a/csrc/ascend/moe/chunk_fwd_o/op_host/op_api/chunk_fwd_o.cpp b/csrc/ascend/moe/chunk_fwd_o/op_host/op_api/chunk_fwd_o.cpp new file mode 100644 index 000000000..f29ddfce9 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_host/op_api/chunk_fwd_o.cpp @@ -0,0 +1,67 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "opdev/op_log.h" +#include "opdev/op_dfx.h" +#include "opdev/make_op_executor.h" +#include "chunk_fwd_o.h" + +using namespace op; + +namespace l0op { +OP_TYPE_REGISTER(ChunkFwdO); + +const std::array ChunkFwdO( + const aclTensor *q, + const aclTensor *k, + const aclTensor *v, + const aclTensor *h, + const aclTensor *g, + const aclIntArray *cuSeqlensOptional, + const aclIntArray *chunkOffsetsOptional, + double scale, + int64_t chunkSize, + const aclTensor *oOut, + aclOpExecutor *executor) +{ + L0_DFX(ChunkFwdO, q, k, v, h, g, cuSeqlensOptional, chunkOffsetsOptional, scale, chunkSize, oOut); + + const aclTensor *actualCuSeqlens = nullptr; + if (cuSeqlensOptional) { + actualCuSeqlens = executor->ConvertToTensor(cuSeqlensOptional, DataType::DT_INT64); + const_cast(actualCuSeqlens)->SetStorageFormat(Format::FORMAT_ND); + const_cast(actualCuSeqlens)->SetViewFormat(Format::FORMAT_ND); + const_cast(actualCuSeqlens)->SetOriginalFormat(Format::FORMAT_ND); + } else { + actualCuSeqlens = nullptr; + } + + const aclTensor *actualChunkOffsets = nullptr; + if (chunkOffsetsOptional) { + actualChunkOffsets = executor->ConvertToTensor(chunkOffsetsOptional, DataType::DT_INT64); + const_cast(actualChunkOffsets)->SetStorageFormat(Format::FORMAT_ND); + const_cast(actualChunkOffsets)->SetViewFormat(Format::FORMAT_ND); + const_cast(actualChunkOffsets)->SetOriginalFormat(Format::FORMAT_ND); + } else { + actualChunkOffsets = nullptr; + } + + auto ret = ADD_TO_LAUNCHER_LIST_AICORE(ChunkFwdO, + OP_INPUT(q, k, v, h, g, actualCuSeqlens, actualChunkOffsets), + OP_OUTPUT(oOut), + OP_ATTR(scale, chunkSize)); + if (ret != ACLNN_SUCCESS) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "ADD_TO_LAUNCHER_LIST_AICORE failed."); + return {nullptr}; + } + return {oOut}; +} + +} // namespace l0op diff --git a/csrc/ascend/moe/chunk_fwd_o/op_host/op_api/chunk_fwd_o.h b/csrc/ascend/moe/chunk_fwd_o/op_host/op_api/chunk_fwd_o.h new file mode 100644 index 000000000..728d0075f --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_host/op_api/chunk_fwd_o.h @@ -0,0 +1,30 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_API_INC_LEVEL0_OP_CHUNK_FWD_O_H +#define OP_API_INC_LEVEL0_OP_CHUNK_FWD_O_H + +#include "opdev/op_executor.h" + +namespace l0op { +const std::array ChunkFwdO( + const aclTensor *q, + const aclTensor *k, + const aclTensor *v, + const aclTensor *h, + const aclTensor *g, + const aclIntArray *cuSeqlensOptional, + const aclIntArray *chunkOffsetsOptional, + double scale, + int64_t chunkSize, + const aclTensor *oOut, + aclOpExecutor *executor); +} + +#endif diff --git a/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/compat_310p.h b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/compat_310p.h new file mode 100644 index 000000000..288c47764 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/compat_310p.h @@ -0,0 +1,39 @@ +#ifndef COMPAT_310P_H +#define COMPAT_310P_H + +#ifndef __CCE_KT_TEST__ +#include "kernel_operator.h" +#endif + +// Dummy bfloat16_t only needed on 310P (dav_m200) where the compiler +// doesn't provide a native bf16 type. On 910B/910C the compiler's +// __clang_cce_types.h already typedefs bfloat16_t from __bf16. +#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 200) && !defined(__bfloat16_t_defined) +#define __bfloat16_t_defined +#define __COMPAT_310P_ACTIVE__ +struct bfloat16_t { + uint16_t val; + bfloat16_t() = default; + bfloat16_t(float v) : val(0) { (void)v; } + operator float() const { return 0.f; } +}; +#endif + +// 310P has no fixpipe unit; post-matmul stores go through MTE3 +#ifndef PIPE_FIX +#define PIPE_FIX PIPE_MTE3 +#endif + +// 310P renames LoadDataWithSparse → LoadDataWithSparseCal +#ifdef __COMPAT_310P_ACTIVE__ +#define LoadDataWithSparse LoadDataWithSparseCal +#endif + +// 310P has no AscendC::ToFloat — dummy bfloat16_t already has operator float() +#ifdef __COMPAT_310P_ACTIVE__ +namespace AscendC { + inline float ToFloat(bfloat16_t v) { return (float)v; } +} +#endif + +#endif diff --git a/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/epilogue/block/block_epilogue_gdn_fwdo_output.hpp b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/epilogue/block/block_epilogue_gdn_fwdo_output.hpp new file mode 100644 index 000000000..86c75a874 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/epilogue/block/block_epilogue_gdn_fwdo_output.hpp @@ -0,0 +1,394 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDO_OUTPUT_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDO_OUTPUT_HPP + +#include "catlass/catlass.hpp" +#include "catlass/arch/resource.hpp" +#include "../gdn_fwd_o_epilogue_policies.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/matrix_coord.hpp" +#include "catlass/epilogue/tile/tile_copy.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class HOutputType_, + class GInputType_, + class AInputType_, + class HInputType_ +> +class BlockEpilogue < + EpilogueAtlasGDNFwdOOutput, + HOutputType_, + GInputType_, + AInputType_, + HInputType_ +> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasGDNFwdOOutput; + using ArchTag = typename DispatchPolicy::ArchTag; + + using HElementOutput = typename HOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using AElementInput = typename AInputType_::Element; + using HElementInput = typename HInputType_::Element; + + // using CopyGmToUbInput = Tile::CopyGm2Ub; + // using CopyUbToGmOutput = Tile::CopyUb2Gm; + + static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; + static constexpr uint32_t FLOAT_ELENUM_PER_BLK = 8; + static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; + static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; + static constexpr uint32_t UB_TILE_SIZE = 16384; // 64 * 128 * 2B + static constexpr uint32_t UB_LINE_SIZE = 512; // 128 * 2 * 2B + static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; // 128 * 2 + static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; // 128 + static constexpr uint32_t MULTIPLIER = 2; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + constexpr uint32_t BASE = 0; + constexpr uint32_t MASK_UB_TENSOR_SIZE = 32 * UB_LINE_SIZE; + constexpr uint32_t GBRCLEFTCAST_UB_TENSOR_SIZE = 40 * UB_LINE_SIZE; + constexpr uint32_t GBRCUP_UB_TENSOR_SIZE = 32 * UB_LINE_SIZE; + constexpr uint32_t FLOAT_UB_TENSOR_SIZE = 32 * UB_LINE_SIZE; + constexpr uint32_t HALF_UB_TENSOR_SIZE = 16 * UB_LINE_SIZE; + constexpr uint32_t G_HALF_UB_TENSOR_SIZE = 2 * UB_LINE_SIZE; + constexpr uint32_t G_FLOAT_UB_TENSOR_SIZE = 2 * UB_LINE_SIZE; + + constexpr uint32_t MASK_UB_TENSOR_OFFSET = BASE; + constexpr uint32_t GBRCLEFTCAST_UB_TENSOR_OFFSET = MASK_UB_TENSOR_OFFSET + MASK_UB_TENSOR_SIZE; + constexpr uint32_t GBRCUP_UB_TENSOR_OFFSET = GBRCLEFTCAST_UB_TENSOR_OFFSET + GBRCLEFTCAST_UB_TENSOR_SIZE; + constexpr uint32_t GCOMP_UB_TENSOR_OFFSET = GBRCUP_UB_TENSOR_OFFSET + GBRCUP_UB_TENSOR_SIZE; + constexpr uint32_t SHARE_UB_TENSOR_OFFSET = GCOMP_UB_TENSOR_OFFSET + G_FLOAT_UB_TENSOR_SIZE; + + maskUbTensor = resource.ubBuf.template GetBufferByByte(MASK_UB_TENSOR_OFFSET); + gbrcLeftcastUbTensor = resource.ubBuf.template GetBufferByByte(GBRCLEFTCAST_UB_TENSOR_OFFSET); + gbrcUpUbTensor = resource.ubBuf.template GetBufferByByte(GBRCUP_UB_TENSOR_OFFSET); + gcompUbTensor = resource.ubBuf.template GetBufferByByte(GCOMP_UB_TENSOR_OFFSET); + shareUbTensor = resource.ubBuf.template GetBufferByByte(SHARE_UB_TENSOR_OFFSET); + + constexpr uint32_t G_UB_TENSOR_OFFSET_PING = SHARE_UB_TENSOR_OFFSET + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t G_HALF_UB_TENSOR_OFFSET_PING = G_UB_TENSOR_OFFSET_PING + G_FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t A_UB_TENSOR_OFFSET_PING = G_HALF_UB_TENSOR_OFFSET_PING + G_HALF_UB_TENSOR_SIZE; + constexpr uint32_t H_UB_TENSOR_OFFSET_PING = A_UB_TENSOR_OFFSET_PING + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t OUT_UB_TENSOR_OFFSET_PING = H_UB_TENSOR_OFFSET_PING + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t OUT_HALF_UB_TENSOR_OFFSET_PING = OUT_UB_TENSOR_OFFSET_PING + FLOAT_UB_TENSOR_SIZE; + + gUbTensorPing = resource.ubBuf.template GetBufferByByte(G_UB_TENSOR_OFFSET_PING); + gUbFPTensorPing = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PING); + gUbBFTensorPing = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PING); + aUbTensorPing = resource.ubBuf.template GetBufferByByte(A_UB_TENSOR_OFFSET_PING); + hUbTensorPing = resource.ubBuf.template GetBufferByByte(H_UB_TENSOR_OFFSET_PING); + outUbTensorPing = resource.ubBuf.template GetBufferByByte(OUT_UB_TENSOR_OFFSET_PING); + outUbFPTensorPing = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PING); + outUbBFTensorPing = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PING); + + constexpr uint32_t G_UB_TENSOR_OFFSET_PONG = OUT_HALF_UB_TENSOR_OFFSET_PING + HALF_UB_TENSOR_SIZE; + constexpr uint32_t G_HALF_UB_TENSOR_OFFSET_PONG = G_UB_TENSOR_OFFSET_PONG + G_FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t A_UB_TENSOR_OFFSET_PONG = G_HALF_UB_TENSOR_OFFSET_PONG + G_HALF_UB_TENSOR_SIZE; + constexpr uint32_t H_UB_TENSOR_OFFSET_PONG = A_UB_TENSOR_OFFSET_PONG + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t OUT_UB_TENSOR_OFFSET_PONG = H_UB_TENSOR_OFFSET_PONG + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t OUT_HALF_UB_TENSOR_OFFSET_PONG = OUT_UB_TENSOR_OFFSET_PONG + FLOAT_UB_TENSOR_SIZE; + + gUbTensorPong = resource.ubBuf.template GetBufferByByte(G_UB_TENSOR_OFFSET_PONG); + gUbFPTensorPong = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PONG); + gUbBFTensorPong = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PONG); + aUbTensorPong = resource.ubBuf.template GetBufferByByte(A_UB_TENSOR_OFFSET_PONG); + hUbTensorPong = resource.ubBuf.template GetBufferByByte(H_UB_TENSOR_OFFSET_PONG); + outUbTensorPong = resource.ubBuf.template GetBufferByByte(OUT_UB_TENSOR_OFFSET_PONG); + outUbFPTensorPong = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PONG); + outUbBFTensorPong = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PONG); + } + CATLASS_DEVICE + ~BlockEpilogue() + {} + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor hOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor attnInput, + AscendC::GlobalTensor hInput, + float scale, + uint32_t chunkSize, + uint32_t kHeadDim, + uint32_t vHeadDim, + uint32_t &pingpongFlag + , uint32_t batchIdx, uint32_t headIdx, uint32_t chunkIdx + ) + { + uint32_t mActual = chunkSize; + uint32_t nActual = vHeadDim; + uint32_t alignedM = CeilDiv(nActual, 8) * 8; + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t blockIdx = AscendC::GetBlockIdx(); + uint32_t mActualPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t mActualThisSubBlock = (subBlockIdx == 0) ? mActualPerSubBlock : (mActual - mActualPerSubBlock); + uint32_t mOffset = subBlockIdx * mActualPerSubBlock; + uint32_t nOffset = 0; + int64_t offsetA = mOffset * nActual + nOffset; + + uint32_t gbrcStart, gbrcRealStart, gbrcRealEnd, gbrcRealProcess, gbrcEffStart, gbrcEffEnd, mulsRemain, mulsRemainIdx; + if(mActualThisSubBlock <= 32) + { + if(subBlockIdx == 0) + { + gbrcStart = 0; + gbrcRealStart = 0; + gbrcRealProcess = mActualThisSubBlock; + } + else + { + gbrcStart = mActualPerSubBlock; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActual - gbrcRealStart; + } + gbrcEffStart = gbrcStart - gbrcRealStart; + uint32_t dstShape_[2] = {gbrcRealProcess, nActual}; + uint32_t srcShape_[2] = {gbrcRealProcess, 1}; + + AscendC::ResetMask(); + AscendC::GlobalTensor attnInputThisSubBlock = attnInput[gbrcStart * nActual]; + AscendC::GlobalTensor hInputThisSubBlock = hInput[gbrcStart * nActual]; + AscendC::GlobalTensor gInputThisSubBlock = gInput; + AscendC::GlobalTensor hOutputThisSubBlock = hOutput[gbrcStart * nActual]; + + AscendC::DataCopyParams gfloatUbParams{1, (uint16_t)(mActual*sizeof(float)), 0, 0}; + AscendC::DataCopyParams ghalfUbParams{1, (uint16_t)(mActual*sizeof(half)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + + AscendC::LocalTensor aUbTensor = (pingpongFlag == 0) ? aUbTensorPing : aUbTensorPong; + AscendC::LocalTensor hUbTensor = (pingpongFlag == 0) ? hUbTensorPing : hUbTensorPong; + AscendC::LocalTensor outUbTensor = (pingpongFlag == 0) ? outUbTensorPing : outUbTensorPong; + AscendC::LocalTensor outUbFPTensor = (pingpongFlag == 0) ? outUbFPTensorPing : outUbFPTensorPong; + AscendC::LocalTensor outUbBFTensor = (pingpongFlag == 0) ? outUbBFTensorPing : outUbBFTensorPong; + AscendC::LocalTensor gUbTensor = (pingpongFlag == 0) ? gUbTensorPing : gUbTensorPong; + AscendC::LocalTensor gUbFPTensor = (pingpongFlag == 0) ? gUbFPTensorPing : gUbFPTensorPong; + AscendC::LocalTensor gUbBFTensor = (pingpongFlag == 0) ? gUbBFTensorPing : gUbBFTensorPong; + + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(std::is_same::value) { + AscendC::DataCopy(gUbTensor, gInputThisSubBlock, mActual); + } else { + AscendC::DataCopy(gUbFPTensor, gInputThisSubBlock, mActual); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(!std::is_same::value) { + AscendC::Cast(gUbTensor, gUbFPTensor, AscendC::RoundMode::CAST_NONE, mActual); + AscendC::PipeBarrier(); + } + AscendC::Adds(gcompUbTensor, gUbTensor, (float)0.0, mActual); + AscendC::PipeBarrier(); + + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(hUbTensor, hInputThisSubBlock, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::DataCopy(aUbTensor, attnInputThisSubBlock, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + + AscendC::Exp(gcompUbTensor, gcompUbTensor, mActual); + AscendC::PipeBarrier(); + AscendC::Broadcast(gbrcLeftcastUbTensor, gcompUbTensor[gbrcRealStart], dstShape_, srcShape_, shareUbTensor); + AscendC::PipeBarrier(); + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Mul(gbrcUpUbTensor, hUbTensor, gbrcLeftcastUbTensor[gbrcEffStart*nActual], mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + + AscendC::Add(gbrcUpUbTensor, aUbTensor, gbrcUpUbTensor, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::Muls(outUbTensor, gbrcUpUbTensor, (float)scale, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + if(std::is_same::value) + { + AscendC::Cast(outUbFPTensor, outUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(hOutputThisSubBlock, outUbFPTensor, mActualThisSubBlock * nActual); + } + else + { + AscendC::Cast(outUbBFTensor, outUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(hOutputThisSubBlock, outUbBFTensor, mActualThisSubBlock * nActual); + } + pingpongFlag = 1 - pingpongFlag; + } + else + { + AscendC::ResetMask(); + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + AscendC::DataCopyParams gfloatUbParams{1, (uint16_t)(mActual*sizeof(float)), 0, 0}; + AscendC::DataCopyParams ghalfUbParams{1, (uint16_t)(mActual*sizeof(half)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + + AscendC::LocalTensor gUbTensor = (pingpongFlag == 0) ? gUbTensorPing : gUbTensorPong; + AscendC::LocalTensor gUbFPTensor = (pingpongFlag == 0) ? gUbFPTensorPing : gUbFPTensorPong; + AscendC::LocalTensor gUbBFTensor = (pingpongFlag == 0) ? gUbBFTensorPing : gUbBFTensorPong; + + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(std::is_same::value) { + AscendC::DataCopy(gUbTensor, gInputThisSubBlock, mActual); + } else { + AscendC::DataCopy(gUbFPTensor, gInputThisSubBlock, mActual); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(!std::is_same::value) { + AscendC::Cast(gUbTensor, gUbFPTensor, AscendC::RoundMode::CAST_NONE, mActual); + AscendC::PipeBarrier(); + } + AscendC::Adds(gcompUbTensor, gUbTensor, (float)0.0, mActual); + AscendC::PipeBarrier(); + AscendC::Exp(gcompUbTensor, gcompUbTensor, mActual); + AscendC::PipeBarrier(); + uint32_t mActualPerStage = CeilDiv(mActualThisSubBlock, 2); + uint32_t mActualThisStage = 0; + for(uint32_t stage = 0; stage < 2; stage++) + { + if(stage == 0) mActualThisStage = mActualPerStage; + else mActualThisStage = mActualThisSubBlock - mActualPerStage; + + if(subBlockIdx == 0 && stage == 0) + { + gbrcStart = 0; + gbrcRealStart = 0; + gbrcRealProcess = mActualThisStage; + } + else if(subBlockIdx == 0 && stage == 1) + { + gbrcStart = mActualPerStage; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActualThisSubBlock - gbrcRealStart; + } + else if(subBlockIdx == 1 && stage == 0) + { + gbrcStart = mActualPerSubBlock; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActualPerSubBlock + mActualThisStage - gbrcRealStart; + } + else if(subBlockIdx == 1 && stage == 1) + { + gbrcStart = mActualPerSubBlock + mActualPerStage; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActual - gbrcRealStart; + } + gbrcEffStart = gbrcStart - gbrcRealStart; + uint32_t dstShape_[2] = {gbrcRealProcess, nActual}; + uint32_t srcShape_[2] = {gbrcRealProcess, 1}; + + AscendC::GlobalTensor hOutputThisSubBlock = hOutput[gbrcStart * nActual]; + AscendC::GlobalTensor attnInputThisSubBlock = attnInput[gbrcStart * nActual]; + AscendC::GlobalTensor hInputThisSubBlock = hInput[gbrcStart * nActual]; + + AscendC::LocalTensor aUbTensor = (pingpongFlag == 0) ? aUbTensorPing : aUbTensorPong; + AscendC::LocalTensor hUbTensor = (pingpongFlag == 0) ? hUbTensorPing : hUbTensorPong; + AscendC::LocalTensor outUbTensor = (pingpongFlag == 0) ? outUbTensorPing : outUbTensorPong; + AscendC::LocalTensor outUbFPTensor = (pingpongFlag == 0) ? outUbFPTensorPing : outUbFPTensorPong; + AscendC::LocalTensor outUbBFTensor = (pingpongFlag == 0) ? outUbBFTensorPing : outUbBFTensorPong; + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(hUbTensor, hInputThisSubBlock, mActualThisStage * nActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::DataCopy(aUbTensor, attnInputThisSubBlock, mActualThisStage * nActual); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + + AscendC::Broadcast(gbrcLeftcastUbTensor, gcompUbTensor[gbrcRealStart], dstShape_, srcShape_, shareUbTensor); + AscendC::PipeBarrier(); + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Mul(gbrcUpUbTensor, hUbTensor, gbrcLeftcastUbTensor[gbrcEffStart*nActual], mActualThisStage * nActual); + AscendC::PipeBarrier(); + + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::Add(gbrcUpUbTensor, aUbTensor, gbrcUpUbTensor, mActualThisStage * nActual); + AscendC::PipeBarrier(); + AscendC::Muls(outUbTensor, gbrcUpUbTensor, (float)scale, mActualThisStage * nActual); + AscendC::PipeBarrier(); + + if(std::is_same::value) + { + AscendC::Cast(outUbFPTensor, outUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisStage * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(hOutputThisSubBlock, outUbFPTensor, mActualThisStage * nActual); + } + else + { + AscendC::Cast(outUbBFTensor, outUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisStage * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(hOutputThisSubBlock, outUbBFTensor, mActualThisStage * nActual); + } + pingpongFlag = 1 - pingpongFlag; + } + } + } + +private: + AscendC::LocalTensor maskUbTensor; + AscendC::LocalTensor gbrcLeftcastUbTensor; + AscendC::LocalTensor gbrcUpUbTensor; + AscendC::LocalTensor gcompUbTensor; + AscendC::LocalTensor shareUbTensor; + + AscendC::LocalTensor gUbTensorPing; + AscendC::LocalTensor gUbFPTensorPing; + AscendC::LocalTensor gUbBFTensorPing; + AscendC::LocalTensor aUbTensorPing; + AscendC::LocalTensor hUbTensorPing; + AscendC::LocalTensor outUbTensorPing; + AscendC::LocalTensor outUbFPTensorPing; + AscendC::LocalTensor outUbBFTensorPing; + + AscendC::LocalTensor gUbTensorPong; + AscendC::LocalTensor gUbFPTensorPong; + AscendC::LocalTensor gUbBFTensorPong; + AscendC::LocalTensor aUbTensorPong; + AscendC::LocalTensor hUbTensorPong; + AscendC::LocalTensor outUbTensorPong; + AscendC::LocalTensor outUbFPTensorPong; + AscendC::LocalTensor outUbBFTensorPong; + + +}; +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/epilogue/block/block_epilogue_gdn_fwdo_qkmask.hpp b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/epilogue/block/block_epilogue_gdn_fwdo_qkmask.hpp new file mode 100644 index 000000000..02c037883 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/epilogue/block/block_epilogue_gdn_fwdo_qkmask.hpp @@ -0,0 +1,425 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDO_QKMASK_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDO_QKMASK_HPP + +#include "catlass/catlass.hpp" +#include "catlass/arch/resource.hpp" +#include "../gdn_fwd_o_epilogue_policies.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/matrix_coord.hpp" +#include "catlass/epilogue/tile/tile_copy.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class AOutputType_, + class GInputType_, + class AInputType_, + class MaskInputType_ +> +class BlockEpilogue < + EpilogueAtlasGDNFwdOQkmask, + AOutputType_, + GInputType_, + AInputType_, + MaskInputType_ +> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasGDNFwdOQkmask; + using ArchTag = typename DispatchPolicy::ArchTag; + + using AElementOutput = typename AOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using AElementInput = typename AInputType_::Element; + using MaskElementInput = typename MaskInputType_::Element; + + static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; + static constexpr uint32_t FLOAT_ELENUM_PER_BLK = 8; + static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; + static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; + static constexpr uint32_t UB_TILE_SIZE = 16384; // 64 * 128 * 2B + static constexpr uint32_t UB_LINE_SIZE = 512; // 128 * 2 * 2B + static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; // 128 * 2 + static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; // 128 + static constexpr uint32_t MULTIPLIER = 2; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + constexpr uint32_t BASE = 0; + constexpr uint32_t MASK_UB_TENSOR_SIZE = 32 * UB_LINE_SIZE; + constexpr uint32_t GBRCLEFTCAST_UB_TENSOR_SIZE = 40 * UB_LINE_SIZE; + constexpr uint32_t GBRCUP_UB_TENSOR_SIZE = 32 * UB_LINE_SIZE; + constexpr uint32_t FLOAT_UB_TENSOR_SIZE = 32 * UB_LINE_SIZE; + constexpr uint32_t HALF_UB_TENSOR_SIZE = 16 * UB_LINE_SIZE; + constexpr uint32_t G_HALF_UB_TENSOR_SIZE = 2 * UB_LINE_SIZE; + constexpr uint32_t G_FLOAT_UB_TENSOR_SIZE = 2 * UB_LINE_SIZE; + + constexpr uint32_t MASK_UB_TENSOR_OFFSET = BASE; + constexpr uint32_t GBRCLEFTCAST_UB_TENSOR_OFFSET = MASK_UB_TENSOR_OFFSET + MASK_UB_TENSOR_SIZE; + constexpr uint32_t GBRCUP_UB_TENSOR_OFFSET = GBRCLEFTCAST_UB_TENSOR_OFFSET + GBRCLEFTCAST_UB_TENSOR_SIZE; + constexpr uint32_t GCOMP_UB_TENSOR_OFFSET = GBRCUP_UB_TENSOR_OFFSET + GBRCUP_UB_TENSOR_SIZE; + constexpr uint32_t SHARE_UB_TENSOR_OFFSET = GCOMP_UB_TENSOR_OFFSET + G_FLOAT_UB_TENSOR_SIZE; + + maskUbTensor = resource.ubBuf.template GetBufferByByte(MASK_UB_TENSOR_OFFSET); + gbrcLeftcastUbTensor = resource.ubBuf.template GetBufferByByte(GBRCLEFTCAST_UB_TENSOR_OFFSET); + gbrcUpUbTensor = resource.ubBuf.template GetBufferByByte(GBRCUP_UB_TENSOR_OFFSET); + gcompUbTensor = resource.ubBuf.template GetBufferByByte(GCOMP_UB_TENSOR_OFFSET); + shareUbTensor = resource.ubBuf.template GetBufferByByte(SHARE_UB_TENSOR_OFFSET); + + constexpr uint32_t G_UB_TENSOR_OFFSET_PING = SHARE_UB_TENSOR_OFFSET + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t G_HALF_UB_TENSOR_OFFSET_PING = G_UB_TENSOR_OFFSET_PING + G_FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t A_UB_TENSOR_OFFSET_PING = G_HALF_UB_TENSOR_OFFSET_PING + G_HALF_UB_TENSOR_SIZE; + constexpr uint32_t OUT_UB_TENSOR_OFFSET_PING = A_UB_TENSOR_OFFSET_PING + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t OUT_HALF_UB_TENSOR_OFFSET_PING = OUT_UB_TENSOR_OFFSET_PING + FLOAT_UB_TENSOR_SIZE; + + gUbTensorPing = resource.ubBuf.template GetBufferByByte(G_UB_TENSOR_OFFSET_PING); + gUbFPTensorPing = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PING); + gUbBFTensorPing = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PING); + aUbTensorPing = resource.ubBuf.template GetBufferByByte(A_UB_TENSOR_OFFSET_PING); + outUbTensorPing = resource.ubBuf.template GetBufferByByte(OUT_UB_TENSOR_OFFSET_PING); + outUbFPTensorPing = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PING); + outUbBFTensorPing = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PING); + + constexpr uint32_t G_UB_TENSOR_OFFSET_PONG = 32 * UB_LINE_SIZE + OUT_HALF_UB_TENSOR_OFFSET_PING + HALF_UB_TENSOR_SIZE; + constexpr uint32_t G_HALF_UB_TENSOR_OFFSET_PONG = G_UB_TENSOR_OFFSET_PONG + G_FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t A_UB_TENSOR_OFFSET_PONG = G_HALF_UB_TENSOR_OFFSET_PONG + G_HALF_UB_TENSOR_SIZE; + constexpr uint32_t OUT_UB_TENSOR_OFFSET_PONG = A_UB_TENSOR_OFFSET_PONG + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t OUT_HALF_UB_TENSOR_OFFSET_PONG = OUT_UB_TENSOR_OFFSET_PONG + FLOAT_UB_TENSOR_SIZE; + + gUbTensorPong = resource.ubBuf.template GetBufferByByte(G_UB_TENSOR_OFFSET_PONG); + gUbFPTensorPong = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PONG); + gUbBFTensorPong = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PONG); + aUbTensorPong = resource.ubBuf.template GetBufferByByte(A_UB_TENSOR_OFFSET_PONG); + outUbTensorPong = resource.ubBuf.template GetBufferByByte(OUT_UB_TENSOR_OFFSET_PONG); + outUbFPTensorPong = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PONG); + outUbBFTensorPong = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PONG); + } + + CATLASS_DEVICE + ~BlockEpilogue() + {} + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor maskOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor attnInput, + AscendC::GlobalTensor boolInput, + uint32_t fullChunkSize, + uint32_t chunkSize, + uint32_t kHeadDim, + uint32_t vHeadDim, + uint32_t &pingpongFlag + , uint32_t batchIdx, uint32_t headIdx, uint32_t chunkIdx + ) + { + uint32_t mActual = chunkSize; + uint32_t nActual = chunkSize; + uint32_t alignedNActual = CeilDiv(nActual, 16) * 16; + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t blockIdx = AscendC::GetBlockIdx(); + uint32_t mActualPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t mActualThisSubBlock = (subBlockIdx == 0) ? mActualPerSubBlock : (mActual - mActualPerSubBlock); + uint32_t mOffset = subBlockIdx * mActualPerSubBlock; + uint32_t nOffset = 0; + int64_t offsetA = mOffset * nActual + nOffset; + uint16_t aInputDstStride; + if((nActual - 1) % 16 <= 7) aInputDstStride = 1; + else aInputDstStride = 0; + + uint32_t gbrcStart, gbrcRealStart, gbrcRealEnd, gbrcRealProcess, gbrcEffStart, gbrcEffEnd, mulsRemain, mulsRemainIdx; + if(mActualThisSubBlock <= 32) + { if(subBlockIdx == 0) + { + gbrcStart = 0; + gbrcRealStart = 0; + gbrcRealProcess = mActualThisSubBlock; + } + else + { + gbrcStart = mActualPerSubBlock; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActual - gbrcRealStart; + } + + gbrcEffStart = gbrcStart - gbrcRealStart; + gbrcEffEnd = gbrcEffStart + mActualThisSubBlock; + + uint32_t dstUpShape_[2] = {mActualThisSubBlock, alignedNActual}; + uint32_t srcUpShape_[2] = {1, alignedNActual}; + uint32_t dstLeftShape_[2] = {gbrcRealProcess, alignedNActual}; + uint32_t srcLeftShape_[2] = {gbrcRealProcess, 1}; + + AscendC::ResetMask(); + AscendC::GlobalTensor maskOutputThisSubBlock = maskOutput[gbrcStart * nActual]; + AscendC::GlobalTensor attnInputThisSubBlock = attnInput[gbrcStart * nActual]; + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + + AscendC::DataCopyParams aInputUbParams{(uint16_t)mActualThisSubBlock, (uint16_t)(nActual*sizeof(float)), 0, aInputDstStride}; + AscendC::DataCopyPadParams aInputUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams aOutputUbParams{(uint16_t)mActualThisSubBlock, (uint32_t)(nActual*sizeof(half)), 0, 0, 0}; + + AscendC::DataCopyParams gfloatUbParams{1, (uint16_t)(mActual*sizeof(float)), 0, 0}; + AscendC::DataCopyParams ghalfUbParams{1, (uint16_t)(mActual*sizeof(half)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + + AscendC::LocalTensor aUbTensor = (pingpongFlag == 0) ? aUbTensorPing : aUbTensorPong; + AscendC::LocalTensor outUbTensor = (pingpongFlag == 0) ? outUbTensorPing : outUbTensorPong; + AscendC::LocalTensor outUbFPTensor = (pingpongFlag == 0) ? outUbFPTensorPing : outUbFPTensorPong; + AscendC::LocalTensor outUbBFTensor = (pingpongFlag == 0) ? outUbBFTensorPing : outUbBFTensorPong; + AscendC::LocalTensor gUbTensor = (pingpongFlag == 0) ? gUbTensorPing : gUbTensorPong; + AscendC::LocalTensor gUbFPTensor = (pingpongFlag == 0) ? gUbFPTensorPing : gUbFPTensorPong; + AscendC::LocalTensor gUbBFTensor = (pingpongFlag == 0) ? gUbBFTensorPing : gUbBFTensorPong; + + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(std::is_same::value) { + AscendC::DataCopy(gUbTensor, gInputThisSubBlock, mActual); + } else { + AscendC::DataCopy(gUbFPTensor, gInputThisSubBlock, mActual); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(!std::is_same::value) { + AscendC::Cast(gUbTensor, gUbFPTensor, AscendC::RoundMode::CAST_NONE, mActual); + AscendC::PipeBarrier(); + } + AscendC::Adds(gcompUbTensor, gUbTensor, (float)0.0, mActual); + AscendC::PipeBarrier(); + + AscendC::Broadcast(gbrcUpUbTensor, gcompUbTensor, dstUpShape_, srcUpShape_, shareUbTensor); + AscendC::Broadcast(gbrcLeftcastUbTensor, gcompUbTensor[gbrcRealStart], dstLeftShape_, srcLeftShape_, shareUbTensor); + AscendC::PipeBarrier(); + AscendC::Sub(gbrcUpUbTensor, gbrcLeftcastUbTensor[gbrcEffStart*alignedNActual], gbrcUpUbTensor, mActualThisSubBlock * alignedNActual); + AscendC::PipeBarrier(); + AscendC::Mins(gbrcUpUbTensor, gbrcUpUbTensor, (float)0.0, mActualThisSubBlock * alignedNActual); + AscendC::PipeBarrier(); + AscendC::Exp(gbrcUpUbTensor, gbrcUpUbTensor, mActualThisSubBlock * alignedNActual); + AscendC::PipeBarrier(); + + gbrcRealEnd = CeilDiv(gbrcStart + mActualThisSubBlock, 8) * 8; + AscendC::Mul(gbrcUpUbTensor[gbrcRealStart], gbrcUpUbTensor[gbrcRealStart], maskUbTensor[gbrcEffStart * 64], gbrcRealEnd - gbrcRealStart, mActualThisSubBlock, + {1, 1, 1, static_cast(alignedNActual/8), static_cast(alignedNActual/8), static_cast(64/8)}); + AscendC::PipeBarrier(); + + mulsRemain = alignedNActual - gbrcRealEnd; + mulsRemainIdx = gbrcRealEnd; + while(mulsRemain > 64) + { + AscendC::Muls(gbrcUpUbTensor[mulsRemainIdx], gbrcUpUbTensor[mulsRemainIdx], (float)0.0, 64, mActualThisSubBlock, + {1, 1, static_cast(alignedNActual/8), static_cast(alignedNActual/8)}); + mulsRemain -= 64; + mulsRemainIdx += 64; + } + AscendC::Muls(gbrcUpUbTensor[mulsRemainIdx], gbrcUpUbTensor[mulsRemainIdx], (float)0.0, mulsRemain, mActualThisSubBlock, + {1, 1, static_cast(alignedNActual/8), static_cast(alignedNActual/8)}); + AscendC::PipeBarrier(); + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if(chunkSize==fullChunkSize) AscendC::DataCopy(aUbTensor, attnInputThisSubBlock, mActualThisSubBlock*nActual); + else AscendC::DataCopy(aUbTensor, attnInputThisSubBlock, mActualThisSubBlock*nActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Mul(outUbTensor, aUbTensor, gbrcUpUbTensor, mActualThisSubBlock * alignedNActual); + AscendC::PipeBarrier(); + + if(std::is_same::value) + { + AscendC::Cast(outUbFPTensor, outUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * alignedNActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(maskOutputThisSubBlock, outUbFPTensor, mActualThisSubBlock*nActual); + } + else + { + AscendC::Cast(outUbBFTensor, outUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * alignedNActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(maskOutputThisSubBlock, outUbBFTensor, mActualThisSubBlock*nActual); + } + pingpongFlag = 1 - pingpongFlag; + } + else // mActualThisSubBlock > 32 ; <=64 + { + AscendC::ResetMask(); + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + AscendC::DataCopyParams gfloatUbParams{1, (uint16_t)(mActual*sizeof(float)), 0, 0}; + AscendC::DataCopyParams ghalfUbParams{1, (uint16_t)(mActual*sizeof(half)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + + AscendC::LocalTensor gUbTensor = (pingpongFlag == 0) ? gUbTensorPing : gUbTensorPong; + AscendC::LocalTensor gUbFPTensor = (pingpongFlag == 0) ? gUbFPTensorPing : gUbFPTensorPong; + AscendC::LocalTensor gUbBFTensor = (pingpongFlag == 0) ? gUbBFTensorPing : gUbBFTensorPong; + + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(std::is_same::value) { + AscendC::DataCopy(gUbTensor, gInputThisSubBlock, mActual); + } else { + AscendC::DataCopy(gUbFPTensor, gInputThisSubBlock, mActual); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(!std::is_same::value) { + AscendC::Cast(gUbTensor, gUbFPTensor, AscendC::RoundMode::CAST_NONE, mActual); + AscendC::PipeBarrier(); + } + AscendC::Adds(gcompUbTensor, gUbTensor, (float)0.0, mActual); + AscendC::PipeBarrier(); + uint32_t mActualPerStage = CeilDiv(mActualThisSubBlock, 2); + uint32_t mActualThisStage = 0; + for(uint32_t stage = 0; stage < 2; ++stage) + { + if(stage==0) mActualThisStage = mActualPerStage; + else mActualThisStage = mActualThisSubBlock - mActualPerStage; + + if(subBlockIdx == 0 && stage == 0) + { + gbrcStart = 0; + gbrcRealStart = 0; + gbrcRealProcess = mActualThisStage; + } + else if(subBlockIdx == 0 && stage == 1) + { + gbrcStart = mActualPerStage; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActualThisSubBlock - gbrcRealStart; + } + else if(subBlockIdx == 1 && stage == 0) + { + gbrcStart = mActualPerSubBlock; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActualPerSubBlock + mActualThisStage - gbrcRealStart; + } + else if(subBlockIdx == 1 && stage == 1) + { + gbrcStart = mActualPerSubBlock + mActualPerStage; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActual - gbrcRealStart; + } + + gbrcEffStart = gbrcStart - gbrcRealStart; + + AscendC::GlobalTensor maskOutputThisSubBlock = maskOutput[gbrcStart * nActual]; + AscendC::GlobalTensor attnInputThisSubBlock = attnInput[gbrcStart * nActual]; + + AscendC::DataCopyParams aInputUbParams{(uint16_t)mActualThisStage, (uint16_t)(nActual*sizeof(float)), 0, aInputDstStride}; + AscendC::DataCopyPadParams aInputUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams aOutputUbParams{(uint16_t)mActualThisStage, (uint32_t)(nActual*sizeof(half)), 0, 0, 0}; + + AscendC::LocalTensor aUbTensor = (pingpongFlag == 0) ? aUbTensorPing : aUbTensorPong; + AscendC::LocalTensor outUbTensor = (pingpongFlag == 0) ? outUbTensorPing : outUbTensorPong; + AscendC::LocalTensor outUbFPTensor = (pingpongFlag == 0) ? outUbFPTensorPing : outUbFPTensorPong; + AscendC::LocalTensor outUbBFTensor = (pingpongFlag == 0) ? outUbBFTensorPing : outUbBFTensorPong; + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if(chunkSize==fullChunkSize) AscendC::DataCopy(aUbTensor, attnInputThisSubBlock, mActualThisStage*nActual); + else AscendC::DataCopy(aUbTensor, attnInputThisSubBlock, mActualThisStage*nActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + uint32_t dstUpShape_[2] = {mActualThisStage, alignedNActual}; + uint32_t srcUpShape_[2] = {1, alignedNActual}; + uint32_t dstLeftShape_[2] = {gbrcRealProcess, alignedNActual}; + uint32_t srcLeftShape_[2] = {gbrcRealProcess, 1}; + + // 310P: Broadcast + gating + causal mask via row loops (strided Mul/Muls banned) + AscendC::Broadcast(gbrcUpUbTensor, gcompUbTensor, dstUpShape_, srcUpShape_, shareUbTensor); + AscendC::Broadcast(gbrcLeftcastUbTensor, gcompUbTensor[gbrcRealStart], dstLeftShape_, srcLeftShape_, shareUbTensor); + AscendC::PipeBarrier(); + AscendC::Sub(gbrcUpUbTensor, gbrcLeftcastUbTensor[gbrcEffStart*alignedNActual], gbrcUpUbTensor, mActualThisStage * alignedNActual); + AscendC::PipeBarrier(); + AscendC::Mins(gbrcUpUbTensor, gbrcUpUbTensor, (float)0.0, mActualThisStage * alignedNActual); + AscendC::PipeBarrier(); + AscendC::Exp(gbrcUpUbTensor, gbrcUpUbTensor, mActualThisStage * alignedNActual); + AscendC::PipeBarrier(); + // Causal mask: zero upper triangle row by row + // Use Duplicate for count >= 8, skip for count < 8 + // (near-diagonal positions have negligible impact on the causal gate) + for (uint32_t row = 0; row < mActualThisStage; ++row) { + uint32_t globalRow = gbrcStart + row; + uint32_t validCols = globalRow + 1; + if (validCols > alignedNActual) validCols = alignedNActual; + uint32_t rowOff = row * alignedNActual; + uint32_t zeroLen = alignedNActual - validCols; + if (zeroLen >= 8) { + AscendC::Duplicate(gbrcUpUbTensor[rowOff + validCols], (float)0.0, zeroLen); + } else if (zeroLen > 0) { + for (uint32_t c = 0; c < zeroLen; ++c) { + gbrcUpUbTensor.SetValue(rowOff + validCols + c, (float)0.0); + } + } + } + AscendC::PipeBarrier(); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Mul(outUbTensor, aUbTensor, gbrcUpUbTensor, mActualThisStage * alignedNActual); + AscendC::PipeBarrier(); + if(std::is_same::value) + { + AscendC::Cast(outUbFPTensor, outUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisStage * alignedNActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(maskOutputThisSubBlock, outUbFPTensor, mActualThisStage*nActual); + } + else + { + AscendC::Cast(outUbBFTensor, outUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisStage * alignedNActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(maskOutputThisSubBlock, outUbBFTensor, mActualThisStage*nActual); + } + pingpongFlag = 1 - pingpongFlag; + } + } + + } + +private: + AscendC::LocalTensor maskUbTensor; + AscendC::LocalTensor gbrcLeftcastUbTensor; + AscendC::LocalTensor gbrcUpUbTensor; + AscendC::LocalTensor gcompUbTensor; + AscendC::LocalTensor shareUbTensor; + + AscendC::LocalTensor gUbTensorPing; + AscendC::LocalTensor gUbFPTensorPing; + AscendC::LocalTensor gUbBFTensorPing; + AscendC::LocalTensor aUbTensorPing; + AscendC::LocalTensor outUbTensorPing; + AscendC::LocalTensor outUbFPTensorPing; + AscendC::LocalTensor outUbBFTensorPing; + + AscendC::LocalTensor gUbTensorPong; + AscendC::LocalTensor gUbFPTensorPong; + AscendC::LocalTensor gUbBFTensorPong; + AscendC::LocalTensor aUbTensorPong; + AscendC::LocalTensor outUbTensorPong; + AscendC::LocalTensor outUbFPTensorPong; + AscendC::LocalTensor outUbBFTensorPong; + +}; +} + +#endif diff --git a/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/epilogue/gdn_fwd_o_epilogue_policies.hpp b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/epilogue/gdn_fwd_o_epilogue_policies.hpp new file mode 100644 index 000000000..71e3c5e33 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/epilogue/gdn_fwd_o_epilogue_policies.hpp @@ -0,0 +1,27 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_EPILOGUE_GDN_FWD_O_EPILOGUE_POLICIES_HPP +#define CATLASS_EPILOGUE_GDN_FWD_O_EPILOGUE_POLICIES_HPP + +#include "catlass/catlass.hpp" + +namespace Catlass::Epilogue { + +struct EpilogueAtlasGDNFwdOQkmask { + using ArchTag = Arch::AtlasA2; +}; + +struct EpilogueAtlasGDNFwdOOutput { + using ArchTag = Arch::AtlasA2; +}; + +} // namespace Catlass::Epilogue + +#endif // CATLASS_EPILOGUE_GDN_FWD_O_EPILOGUE_POLICIES_HPP diff --git a/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/gemm/block/block_scheduler_gdn_fwd_o.hpp b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/gemm/block/block_scheduler_gdn_fwd_o.hpp new file mode 100644 index 000000000..bfc2df9a0 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/gemm/block/block_scheduler_gdn_fwd_o.hpp @@ -0,0 +1,274 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_GEMM_SCHEDULER_GDN_FWD_O_HPP +#define CATLASS_GEMM_SCHEDULER_GDN_FWD_O_HPP + +// constexpr uint32_t PING_PONG_STAGES = 1; +constexpr uint32_t PING_PONG_STAGES = 2; +constexpr uint32_t BYTE_SIZE_16_BIT = 2; + +template +CATLASS_DEVICE T AlignUp(T a, T b) { + return (b == 0) ? 0 : (a + b - 1) / b * b; +} + +template +CATLASS_DEVICE T Min(T a, T b) { + return (a > b) ? b : a; +} + +template +CATLASS_DEVICE T Max(T a, T b) { + return (a > b) ? a : b; +} + +namespace Catlass::Gemm::Block { + + +struct GDNFwdOOffsets { + uint32_t qkOffset; + uint32_t ovOffset; + uint32_t hOffset; + uint32_t gOffset; + uint32_t attnWorkOffset; + uint32_t hvWorkOffset; + bool isFinalState; + uint32_t blockTokens; + uint32_t batchIdx; + uint32_t headIdx; + uint32_t chunkIdx; + +}; + +struct BlockSchedulerGdnFwdO { + uint32_t shapeBatch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + uint32_t isVariedLen; + uint32_t tokenBatch; + uint32_t numChunks{0}; + uint32_t vBlockSize{128}; + + uint32_t taskIdx; + uint32_t cubeCoreIdx; + uint32_t cubeCoreNum; + uint32_t vLoops; + uint32_t taskNum; + uint32_t headGroups; + + bool isRunning; + bool processNewTask {true}; + bool firstLoop {true}; + bool lastLoop {false}; + GDNFwdOOffsets offsets[PING_PONG_STAGES]; + int32_t currStage{PING_PONG_STAGES - 1}; + + uint32_t vIdx; + uint32_t batchIdx; + uint32_t baseHeadIdx; + uint32_t chunkIdx; + uint32_t headInnerIdx; + uint32_t vHeadIdx; + uint32_t kHeadIdx; + uint32_t shapeBatchIdx; + uint32_t tokenBatchIdx; + + uint32_t batchChunkIdx; + uint32_t batchChunkStartIdx; + uint32_t tokenOffset; + uint32_t batchChunks; + uint32_t batchTokens; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmChunkOffsets; + + Arch::CrossCoreFlag cube1Done{3}; + Arch::CrossCoreFlag vec1Done{4}; + Arch::CrossCoreFlag cube2Done{5}; + Arch::CrossCoreFlag cube3Done{6}; + Arch::CrossCoreFlag vec2Done{7}; + + CATLASS_DEVICE + BlockSchedulerGdnFwdO() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_offsets, GM_ADDR tiling, uint32_t coreIdx, uint32_t coreNum) { + __gm__ ChunkFwdOTilingData *__restrict gdnFwdOTilingData = reinterpret_cast<__gm__ ChunkFwdOTilingData *__restrict>(tiling); + shapeBatch = gdnFwdOTilingData->shapeBatch; + seqlen = gdnFwdOTilingData->seqlen; + kNumHead = gdnFwdOTilingData->kNumHead; + vNumHead = gdnFwdOTilingData->vNumHead; + kHeadDim = gdnFwdOTilingData->kHeadDim; + vHeadDim = gdnFwdOTilingData->vHeadDim; + chunkSize = gdnFwdOTilingData->chunkSize; + isVariedLen = gdnFwdOTilingData->isVariedLen; + tokenBatch = gdnFwdOTilingData->tokenBatch; + + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmChunkOffsets.SetGlobalBuffer((__gm__ int64_t *)chunk_offsets); + + if (isVariedLen) { + for (uint32_t b = 1; b <= tokenBatch; b++) { + numChunks += (gmSeqlen.GetValue(b) - gmSeqlen.GetValue(b - 1) + chunkSize - 1) / chunkSize; + } + } else { + numChunks = (seqlen + chunkSize - 1) / chunkSize; + } + + cubeCoreIdx = coreIdx; + cubeCoreNum = coreNum; + vLoops = vHeadDim / vBlockSize; + taskNum = vLoops * shapeBatch * numChunks * vNumHead; + headGroups = vNumHead / kNumHead; + taskIdx = cubeCoreIdx * PING_PONG_STAGES; + isRunning = taskIdx < taskNum; + + } + + CATLASS_DEVICE + void InitTask() { + if (processNewTask) { + if (unlikely(taskIdx >= taskNum)) { + isRunning = false; + } + vIdx = taskIdx / (shapeBatch * numChunks * vNumHead); + shapeBatchIdx = (taskIdx - vIdx * shapeBatch * numChunks * vNumHead) / (numChunks * vNumHead); + chunkIdx = (taskIdx - vIdx * shapeBatch * numChunks * vNumHead - shapeBatchIdx * numChunks * vNumHead) / vNumHead; + baseHeadIdx = taskIdx % vNumHead; + tokenBatchIdx = isVariedLen ? gmChunkOffsets.GetValue(2 * chunkIdx) : 0; + batchChunkIdx = isVariedLen ? gmChunkOffsets.GetValue(2 * chunkIdx + 1) : chunkIdx; + batchChunkStartIdx = chunkIdx - batchChunkIdx; + tokenOffset = isVariedLen ? gmSeqlen.GetValue(tokenBatchIdx) : 0; + batchTokens = isVariedLen ? (gmSeqlen.GetValue(tokenBatchIdx + 1) - tokenOffset) : seqlen; + headInnerIdx = 0; + } else { + headInnerIdx = (headInnerIdx + 1) % PING_PONG_STAGES; + } + + vHeadIdx = baseHeadIdx + headInnerIdx; + kHeadIdx = vHeadIdx / headGroups; + offsets[currStage].qkOffset = (shapeBatchIdx * kNumHead * seqlen + kHeadIdx * seqlen + tokenOffset + batchChunkIdx * chunkSize) * kHeadDim; + offsets[currStage].ovOffset = (shapeBatchIdx * vNumHead * seqlen + vHeadIdx * seqlen + tokenOffset + batchChunkIdx * chunkSize) * vHeadDim; + offsets[currStage].hOffset = (shapeBatchIdx * vNumHead * numChunks + vHeadIdx * numChunks + chunkIdx) * kHeadDim * vHeadDim; + offsets[currStage].gOffset = shapeBatchIdx * vNumHead * seqlen + vHeadIdx * seqlen + tokenOffset + batchChunkIdx * chunkSize; + offsets[currStage].attnWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + currStage) * chunkSize * chunkSize; + offsets[currStage].hvWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + currStage) * chunkSize * vHeadDim; + offsets[currStage].isFinalState = chunkIdx == (numChunks - 1) || (isVariedLen && gmChunkOffsets.GetValue(2 * chunkIdx + 3) == 0); + offsets[currStage].blockTokens = offsets[currStage].isFinalState ? (batchTokens - batchChunkIdx * chunkSize) : chunkSize; + offsets[currStage].batchIdx = batchIdx; + offsets[currStage].headIdx = vHeadIdx; + offsets[currStage].chunkIdx = chunkIdx; + + processNewTask = headInnerIdx == PING_PONG_STAGES - 1; + if (processNewTask) { + taskIdx += PING_PONG_STAGES * cubeCoreNum; + } + + currStage = (currStage + 1) % PING_PONG_STAGES; + } + + +}; + +struct BlockSchedulerGdnFwdOCube : public BlockSchedulerGdnFwdO { + CATLASS_DEVICE + BlockSchedulerGdnFwdOCube() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_offsets, GM_ADDR tiling) { + BlockSchedulerGdnFwdO::Init(cu_seqlens, chunk_offsets, tiling, AscendC::GetBlockIdx(), AscendC::GetBlockNum()); + } + + CATLASS_DEVICE + bool NeedProcessCube1() { + return true; + } + + CATLASS_DEVICE + GDNFwdOOffsets& GetCube1Offsets() { + return offsets[(currStage - 1) % PING_PONG_STAGES]; + } + + CATLASS_DEVICE + GemmCoord GetCube1Shape() { + GDNFwdOOffsets& cube1Offsets = GetCube1Offsets(); + return GemmCoord{cube1Offsets.blockTokens, cube1Offsets.blockTokens, kHeadDim}; + } + + CATLASS_DEVICE + bool NeedProcessCube23() { + if (unlikely(firstLoop)) { + firstLoop = false; + return false; + } + return true; + } + + CATLASS_DEVICE + GDNFwdOOffsets& GetCube23Offsets() { + return offsets[(currStage - 2) % PING_PONG_STAGES]; + } + + CATLASS_DEVICE + GemmCoord GetCube2Shape() { + GDNFwdOOffsets& cube2Offsets = GetCube23Offsets(); + return GemmCoord{kHeadDim, vHeadDim, cube2Offsets.blockTokens}; + } + + CATLASS_DEVICE + GemmCoord GetCube3Shape() { + GDNFwdOOffsets& cube2Offsets = GetCube23Offsets(); + return GemmCoord{kHeadDim, vHeadDim, cube2Offsets.blockTokens}; + } + +}; + +struct BlockSchedulerGdnFwdOVec : public BlockSchedulerGdnFwdO { + CATLASS_DEVICE + BlockSchedulerGdnFwdOVec() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_offsets, GM_ADDR tiling) { + BlockSchedulerGdnFwdO::Init(cu_seqlens, chunk_offsets, tiling, AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(), AscendC::GetBlockNum()); + } + + CATLASS_DEVICE + bool NeedProcessVec1() { + return isRunning; + } + + CATLASS_DEVICE + bool NeedProcessVec2() { + if (unlikely(firstLoop)) { + firstLoop = false; + return false; + } + return true; + } + + CATLASS_DEVICE + GDNFwdOOffsets& GetVec1Offsets() { + return offsets[(currStage - 1) % PING_PONG_STAGES]; + } + + CATLASS_DEVICE + GDNFwdOOffsets& GetVec2Offsets() { + return offsets[(currStage - 2) % PING_PONG_STAGES]; + } + +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_SCHEDULER_GDN_FWD_O_HPP \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/gemm/kernel/gdn_fwd_o_kernel.hpp b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/gemm/kernel/gdn_fwd_o_kernel.hpp new file mode 100644 index 000000000..2bef5cc16 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch20/gemm/kernel/gdn_fwd_o_kernel.hpp @@ -0,0 +1,551 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#define CATLASS_ARCH 2201 +#define CATLASS_UNIFIED_CORE 1 + +#include "catlass/arch/arch.hpp" +#include "catlass/arch/cross_core_sync.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/catlass.hpp" +#include "catlass/epilogue/block/block_epilogue.hpp" +#include "../../epilogue/block/block_epilogue_gdn_fwdo_qkmask.hpp" +#include "../../epilogue/block/block_epilogue_gdn_fwdo_output.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/gemm/block/block_swizzle.hpp" +#include "../block/block_scheduler_gdn_fwd_o.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" +#include "catlass/gemm_coord.hpp" +#include "tla/tensor.hpp" +#include "tla/layout.hpp" +#include "tla/tensor.hpp" + +using _0 = tla::Int<0>; +using _1 = tla::Int<1>; +using _2 = tla::Int<2>; +using _4 = tla::Int<4>; +using _8 = tla::Int<8>; +using _16 = tla::Int<16>; +using _32 = tla::Int<32>; +using _64 = tla::Int<64>; +using _128 = tla::Int<128>; +using _256 = tla::Int<256>; +using _512 = tla::Int<512>; +using _1024 = tla::Int<1024>; +using _2048 = tla::Int<2048>; +using _4096 = tla::Int<4096>; +using _8192 = tla::Int<8192>; +using _16384 = tla::Int<16384>; +using _32768 = tla::Int<32768>; +using _65536 = tla::Int<65536>; + + +#include "kernel_operator.h" +using namespace Catlass; +using namespace tla; + +namespace Catlass::Gemm::Kernel { + +template< + typename INPUT_TYPE, + typename G_TYPE, + typename WORKSPACE_TYPE +> +class GDNFwdOKernel { +public: + + using ArchTag = Arch::AtlasA2; + using GDNFwdOOffsets = Catlass::Gemm::Block::GDNFwdOOffsets; + + using CubeScheduler = typename Catlass::Gemm::Block::BlockSchedulerGdnFwdOCube; + using VecScheduler = typename Catlass::Gemm::Block::BlockSchedulerGdnFwdOVec; + + using DispatchPolicyTla = Gemm::MmadPingpongTlaMulti; + using L1TileShapeTla = Shape<_128, _128, _128>; + using L0TileShapeTla = L1TileShapeTla; + using QType = Gemm::GemmType; + using KType = Gemm::GemmType; + using AttenType = Gemm::GemmType; + using AttenMaskedType = Gemm::GemmType; + using HType = Gemm::GemmType; + using OinterType = Gemm::GemmType; + using VNEWType = Gemm::GemmType; + + using GType = Gemm::GemmType; + using OType = Gemm::GemmType; + using MaskType = Gemm::GemmType; + + // cube 1 + using TileCopyQK = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadQK = Gemm::Block::BlockMmadTla; + + // cube 2 + using TileCopyQH = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadQH = Gemm::Block::BlockMmadTla; + + // cube 3 + using TileCopyAttenVNEW = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadAttenVNEW = Gemm::Block::BlockMmadTla; + + // vec 1 + using DispatchPolicyGDNFwdOQkmask = Epilogue::EpilogueAtlasGDNFwdOQkmask; + using EpilogueGDNFwdOQkmask = Epilogue::Block::BlockEpilogue; + + // vec 2 + using DispatchPolicyGDNFwdOOutput = Epilogue::EpilogueAtlasGDNFwdOOutput; + using EpilogueGDNFwdOOutput = Epilogue::Block::BlockEpilogue; + + using ElementQ = typename BlockMmadQK::ElementA; + using LayoutQ = Catlass::layout::RowMajor; + + using ElementK = typename BlockMmadQK::ElementB; + using LayoutK = Catlass::layout::ColumnMajor; + + using ElementAtten = typename BlockMmadQK::ElementC; + using LayoutAtten = Catlass::layout::RowMajor; + + using ElementAttenMasked = typename BlockMmadQH::ElementA; + using LayoutAttenMasked = Catlass::layout::RowMajor; + + using ElementH = typename BlockMmadQH::ElementB; + using LayoutH = Catlass::layout::RowMajor; + + using ElementOinter = typename BlockMmadQH::ElementC; + using LayoutOinter = Catlass::layout::RowMajor; + + + using ElementVNEW = typename BlockMmadAttenVNEW::ElementB; + using LayoutVNEW = Catlass::layout::RowMajor; + + + using ElementG = G_TYPE; + using ElementMask = bool; + + using L1TileShape = typename BlockMmadQK::L1TileShape; + + uint32_t shapeBatch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + float scale; + uint32_t numChunks; + uint32_t isVariedLen; + uint32_t tokenBatch; + uint32_t vWorkspaceOffset; + uint32_t hWorkspaceOffset; + uint32_t attnWorkspaceOffset; + uint32_t aftermaskWorkspaceOffset; + uint32_t maskWorkspaceOffset; + + AscendC::GlobalTensor gmQ; + AscendC::GlobalTensor gmK; + AscendC::GlobalTensor gmV; + AscendC::GlobalTensor gmH; + AscendC::GlobalTensor gmG; + AscendC::GlobalTensor gmO; + AscendC::GlobalTensor gmVWorkspace; + AscendC::GlobalTensor gmHWorkspace; + AscendC::GlobalTensor gmAttnWorkspace; + AscendC::GlobalTensor gmAftermaskWorkspace; + AscendC::GlobalTensor gmMask; + + CubeScheduler cubeBlockScheduler; + VecScheduler vecBlockScheduler; + + Arch::Resource resource; + + __aicore__ inline GDNFwdOKernel() {} + + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR h, GM_ADDR g, + GM_ADDR cu_seqlens, GM_ADDR chunk_offsets, GM_ADDR o, GM_ADDR tiling, GM_ADDR user) { + + __gm__ ChunkFwdOTilingData *__restrict gdnFwdOTilingData = reinterpret_cast<__gm__ ChunkFwdOTilingData *__restrict>(tiling); + + shapeBatch = gdnFwdOTilingData->shapeBatch; + seqlen = gdnFwdOTilingData->seqlen; + kNumHead = gdnFwdOTilingData->kNumHead; + vNumHead = gdnFwdOTilingData->vNumHead; + kHeadDim = gdnFwdOTilingData->kHeadDim; + vHeadDim = gdnFwdOTilingData->vHeadDim; + scale = gdnFwdOTilingData->scale; + chunkSize = gdnFwdOTilingData->chunkSize; + isVariedLen = gdnFwdOTilingData->isVariedLen; + tokenBatch = gdnFwdOTilingData->tokenBatch; + vWorkspaceOffset = gdnFwdOTilingData->vWorkspaceOffset; + hWorkspaceOffset = gdnFwdOTilingData->hWorkspaceOffset; + attnWorkspaceOffset = gdnFwdOTilingData->attnWorkspaceOffset; + aftermaskWorkspaceOffset = gdnFwdOTilingData->aftermaskWorkspaceOffset; + maskWorkspaceOffset = gdnFwdOTilingData->maskWorkspaceOffset; + + gmQ.SetGlobalBuffer((__gm__ ElementQ *)q); + gmK.SetGlobalBuffer((__gm__ ElementK *)k); + gmV.SetGlobalBuffer((__gm__ ElementVNEW *)v); + gmH.SetGlobalBuffer((__gm__ ElementH *)h); + gmG.SetGlobalBuffer((__gm__ ElementG *)g); + gmO.SetGlobalBuffer((__gm__ ElementVNEW *)o); + gmVWorkspace.SetGlobalBuffer((__gm__ ElementOinter *)(user + vWorkspaceOffset)); + gmHWorkspace.SetGlobalBuffer((__gm__ ElementOinter *)(user + hWorkspaceOffset)); + gmAttnWorkspace.SetGlobalBuffer((__gm__ ElementAtten *)(user + attnWorkspaceOffset)); + gmAftermaskWorkspace.SetGlobalBuffer((__gm__ ElementAttenMasked *)(user + aftermaskWorkspaceOffset)); + gmMask.SetGlobalBuffer((__gm__ ElementMask *)(user + maskWorkspaceOffset)); + + cubeBlockScheduler.Init(cu_seqlens, chunk_offsets, tiling); + } + + __aicore__ inline void Process() { + ProcessUnifiedCore(); + } + + __aicore__ inline void InitCausalMask() { + AscendC::LocalTensor maskUbTensor = resource.ubBuf.template GetBufferByByte(0); + // 310P: Duplicate count must be >= 8 (vector width = 8 floats). + // Build lower-triangular mask: row i has 1.0 in cols [0..i], 0.0 elsewhere. + // Fill all 1.0 first, then zero the upper triangle with count >= 8. + AscendC::Duplicate(maskUbTensor, (float)1.0, 64 * 64); + AscendC::PipeBarrier(); + for (uint32_t i = 0; i < 64; ++i) { + uint32_t zeroStart = i + 1; + uint32_t zeroLen = 64 - zeroStart; + if (zeroLen >= 8) { + AscendC::Duplicate(maskUbTensor[i * 64 + zeroStart], (float)0.0, zeroLen); + } else { + for (uint32_t j = 0; j < zeroLen; ++j) { + maskUbTensor.SetValue(i * 64 + zeroStart + j, (float)0.0); + } + } + } + AscendC::PipeBarrier(); + } + + __aicore__ inline void ProcessUnifiedCore() { + uint32_t coreNum = AscendC::GetBlockNum(); + + BlockMmadQK blockMmadQK(resource); + BlockMmadQH blockMmadQH(resource); + BlockMmadAttenVNEW blockMmadAttenVNEW(resource); + + auto qLayout = tla::MakeLayout(shapeBatch * kNumHead * seqlen, kHeadDim); + auto kLayout = tla::MakeLayout(kHeadDim, shapeBatch * kNumHead * seqlen); + auto hLayout = tla::MakeLayout(shapeBatch * vNumHead * seqlen * kHeadDim, vHeadDim); + auto ointerLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, vHeadDim); + auto vnewLayout = tla::MakeLayout(shapeBatch * vNumHead * seqlen, vHeadDim); + + bool needRun = false; + uint32_t pingpongFlag = 0; + + while (cubeBlockScheduler.isRunning) { + cubeBlockScheduler.InitTask(); + + if (cubeBlockScheduler.isRunning) { + // CUBE1: attn = q @ k.T + GDNFwdOOffsets& cube1Offsets = cubeBlockScheduler.GetCube1Offsets(); + auto attenLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, cube1Offsets.blockTokens); + auto tensorQ = tla::MakeTensor(gmQ[cube1Offsets.qkOffset], qLayout, Catlass::Arch::PositionGM{}); + auto tensorK = tla::MakeTensor(gmK[cube1Offsets.qkOffset], kLayout, Catlass::Arch::PositionGM{}); + auto tensorAttn = tla::MakeTensor(gmAttnWorkspace[cube1Offsets.attnWorkOffset], attenLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape{cube1Offsets.blockTokens, cube1Offsets.blockTokens, kHeadDim}; + auto tensorBlockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + auto tensorBlockAttn = GetTile(tensorAttn, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.n())); + blockMmadQK.preSetFlags(); + blockMmadQK(tensorBlockQ, tensorBlockK, tensorBlockAttn, cube1Shape); + blockMmadQK.finalWaitFlags(); + + // Re-init causal mask after cube (cube overwrites UB[0]) + InitCausalMask(); + + // VEC1: qkmask epilogue + EpilogueGDNFwdOQkmask epilogueGDNFwdOQkmask(resource); + epilogueGDNFwdOQkmask( + gmAftermaskWorkspace[cube1Offsets.attnWorkOffset], + gmG[cube1Offsets.gOffset], gmAttnWorkspace[cube1Offsets.attnWorkOffset], gmMask, + chunkSize, cube1Offsets.blockTokens, kHeadDim, vHeadDim, pingpongFlag, + cube1Offsets.batchIdx, cube1Offsets.headIdx, cube1Offsets.chunkIdx + ); + } + + // GM fence: ensure Vec1 MTE3 writes are committed before Cube3 MTE2 reads + AscendC::PipeBarrier(); + + if (needRun) { + GDNFwdOOffsets& prevOffsets = cubeBlockScheduler.GetCube23Offsets(); + + // CUBE2: h_work = q @ h + auto tensorQ2 = tla::MakeTensor(gmQ[prevOffsets.qkOffset], qLayout, Catlass::Arch::PositionGM{}); + auto tensorH = tla::MakeTensor(gmH[prevOffsets.hOffset], hLayout, Catlass::Arch::PositionGM{}); + auto tensorHWork = tla::MakeTensor(gmHWorkspace[prevOffsets.hvWorkOffset], ointerLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube2Shape{prevOffsets.blockTokens, vHeadDim, kHeadDim}; + auto tensorBlockQ2 = GetTile(tensorQ2, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + auto tensorBlockHWork = GetTile(tensorHWork, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.m(), cube2Shape.n())); + blockMmadQH.preSetFlags(); + blockMmadQH(tensorBlockQ2, tensorBlockH, tensorBlockHWork, cube2Shape); + blockMmadQH.finalWaitFlags(); + + // CUBE3: v_work = attn_masked @ v + auto attenLayout3 = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, prevOffsets.blockTokens); + auto tensorAttnMask = tla::MakeTensor(gmAftermaskWorkspace[prevOffsets.attnWorkOffset], attenLayout3, Catlass::Arch::PositionGM{}); + auto tensorV = tla::MakeTensor(gmV[prevOffsets.ovOffset], vnewLayout, Catlass::Arch::PositionGM{}); + auto tensorVWork = tla::MakeTensor(gmVWorkspace[prevOffsets.hvWorkOffset], ointerLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube3Shape{prevOffsets.blockTokens, vHeadDim, prevOffsets.blockTokens}; + auto tensorBlockAttnMask = GetTile(tensorAttnMask, tla::MakeCoord(0, 0), tla::MakeShape(cube3Shape.m(), cube3Shape.k())); + auto tensorBlockV = GetTile(tensorV, tla::MakeCoord(0, 0), tla::MakeShape(cube3Shape.k(), cube3Shape.n())); + auto tensorBlockVWork = GetTile(tensorVWork, tla::MakeCoord(0, 0), tla::MakeShape(cube3Shape.m(), cube3Shape.n())); + blockMmadAttenVNEW.preSetFlags(); + blockMmadAttenVNEW(tensorBlockAttnMask, tensorBlockV, tensorBlockVWork, cube3Shape); + blockMmadAttenVNEW.finalWaitFlags(); + + // GM fence: ensure Cube2/3 L0C→UB→MTE3→GM writes are committed + AscendC::PipeBarrier(); + + // VEC2 inline for 310P: o = scale * (v_work + exp(g) * h_work) + // The epilogue class uses event-based MTE2 sync that breaks after cube matmul on 310P. + { + constexpr uint32_t STAGE_ROWS = 32; + uint32_t bt = prevOffsets.blockTokens; + uint32_t stageCnt = STAGE_ROWS * vHeadDim; + // UB layout: vwUb[0..stageCnt), hwUb[stageCnt..2*stageCnt), gUb[2*stageCnt..+64) + AscendC::LocalTensor vwUb = resource.ubBuf.template GetBufferByByte(0); + AscendC::LocalTensor hwUb = resource.ubBuf.template GetBufferByByte(stageCnt * sizeof(float)); + AscendC::LocalTensor gUb = resource.ubBuf.template GetBufferByByte(stageCnt * sizeof(float) * 2); + // outUb (half) after gUb, aligned to 512B + constexpr uint32_t G_RESERVE = 512; + AscendC::LocalTensor outUb = resource.ubBuf.template GetBufferByByte( + stageCnt * sizeof(float) * 2 + G_RESERVE); + + for (uint32_t row = 0; row < bt; row += STAGE_ROWS) { + uint32_t rows = (row + STAGE_ROWS <= bt) ? STAGE_ROWS : (bt - row); + uint32_t elems = rows * vHeadDim; + uint32_t gmOff = row * vHeadDim; + + // Load v_work, h_work, g from GM + AscendC::DataCopy(vwUb, gmVWorkspace[prevOffsets.hvWorkOffset + gmOff], elems); + AscendC::DataCopy(hwUb, gmHWorkspace[prevOffsets.hvWorkOffset + gmOff], elems); + // Load g (may be float or half) + if constexpr (std::is_same::value) { + AscendC::DataCopy(gUb, gmG[prevOffsets.gOffset + row], rows); + } else { + AscendC::LocalTensor gTyped = resource.ubBuf.template GetBufferByByte( + stageCnt * sizeof(float) * 2 + 256); + AscendC::DataCopy(gTyped, gmG[prevOffsets.gOffset + row], rows); + AscendC::PipeBarrier(); + AscendC::Cast(gUb, gTyped, AscendC::RoundMode::CAST_NONE, rows); + } + AscendC::PipeBarrier(); + + // exp(g) + AscendC::Exp(gUb, gUb, rows); + AscendC::PipeBarrier(); + + // Broadcast exp(g) into gBrc: each row r gets exp(g[r]) repeated Dv times + // gBrc lives after outUb in UB + AscendC::LocalTensor gBrc = resource.ubBuf.template GetBufferByByte( + stageCnt * sizeof(float) * 2 + G_RESERVE + stageCnt * sizeof(ElementVNEW)); + { + uint32_t dstShape[2] = {rows, vHeadDim}; + uint32_t srcShape[2] = {rows, 1}; + // Broadcast needs a shared temp buffer — use space after gBrc + AscendC::LocalTensor brcTmp = resource.ubBuf.template GetBufferByByte( + stageCnt * sizeof(float) * 2 + G_RESERVE + stageCnt * sizeof(ElementVNEW) + elems * sizeof(float)); + AscendC::Broadcast(gBrc, gUb, dstShape, srcShape, brcTmp); + } + AscendC::PipeBarrier(); + AscendC::Mul(hwUb, hwUb, gBrc, elems); + AscendC::PipeBarrier(); + + // v_work + exp(g)*h_work + AscendC::Add(vwUb, vwUb, hwUb, elems); + AscendC::PipeBarrier(); + // * scale + AscendC::Muls(vwUb, vwUb, (float)scale, elems); + AscendC::PipeBarrier(); + // Cast to output dtype + AscendC::Cast(outUb, vwUb, AscendC::RoundMode::CAST_NONE, elems); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopyParams cp{1, static_cast(elems * sizeof(ElementVNEW) / 32), 0, 0}; + AscendC::DataCopy(gmO[prevOffsets.ovOffset + gmOff], outUb, cp); + AscendC::PipeBarrier(); + } + } + } + + needRun = true; + } + } + + __aicore__ inline void ProcessSplitCore() { + if ASCEND_IS_AIC { + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = AscendC::GetBlockNum(); + + BlockMmadQK blockMmadQK(resource); + BlockMmadQH blockMmadQH(resource); + BlockMmadAttenVNEW blockMmadAttenVNEW(resource); + + auto qLayout = tla::MakeLayout(shapeBatch * kNumHead * seqlen, kHeadDim); + auto kLayout = tla::MakeLayout(kHeadDim, shapeBatch * kNumHead * seqlen); + auto hLayout = tla::MakeLayout(shapeBatch * vNumHead * seqlen * kHeadDim, vHeadDim); + auto ointerLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, vHeadDim); + auto vnewLayout = tla::MakeLayout(shapeBatch * vNumHead * seqlen, vHeadDim); + + bool needRun = false; + bool isFirstC3 = true; + + while (cubeBlockScheduler.isRunning) { + cubeBlockScheduler.InitTask(); + + if (cubeBlockScheduler.isRunning && coreIdx < coreNum) { + + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done); + + GDNFwdOOffsets& cube1Offsets = cubeBlockScheduler.GetCube1Offsets(); + int64_t cube1OffsetQ = cube1Offsets.qkOffset; + int64_t cube1OffsetK = cube1Offsets.qkOffset; + int64_t cube1OffsetAttn = cube1Offsets.attnWorkOffset; + auto attenLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, cube1Offsets.blockTokens); + auto tensorQ = tla::MakeTensor(gmQ[cube1OffsetQ], qLayout, Catlass::Arch::PositionGM{}); + auto tensorK = tla::MakeTensor(gmK[cube1OffsetK], kLayout, Catlass::Arch::PositionGM{}); + auto tensorAttn = tla::MakeTensor(gmAttnWorkspace[cube1OffsetAttn], attenLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape{cube1Offsets.blockTokens, cube1Offsets.blockTokens, kHeadDim}; + auto tensorBlockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + auto tensorBlockAttn = GetTile(tensorAttn, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.n())); + blockMmadQK.preSetFlags(); + blockMmadQK(tensorBlockQ, tensorBlockK, tensorBlockAttn, cube1Shape); + blockMmadQK.finalWaitFlags(); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube1Done); + + } + // AscendC::PipeBarrier(); + + if (needRun && coreIdx < coreNum) { + if(!cubeBlockScheduler.isRunning) Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done); + GDNFwdOOffsets& cube2Offsets = cubeBlockScheduler.GetCube23Offsets(); + int64_t cube2OffsetQ = cube2Offsets.qkOffset; + int64_t cube2OffsetH = cube2Offsets.hOffset; + int64_t cube2OffsetHWork = cube2Offsets.hvWorkOffset; + auto tensorQ = tla::MakeTensor(gmQ[cube2OffsetQ], qLayout, Catlass::Arch::PositionGM{}); + auto tensorH = tla::MakeTensor(gmH[cube2OffsetH], hLayout, Catlass::Arch::PositionGM{}); + auto tensorHWork = tla::MakeTensor(gmHWorkspace[cube2OffsetHWork], ointerLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube2Shape{cube2Offsets.blockTokens, vHeadDim, kHeadDim}; + auto tensorBlockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + auto tensorBlockHWork = GetTile(tensorHWork, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.m(), cube2Shape.n())); + blockMmadQH.preSetFlags(); + blockMmadQH(tensorBlockQ, tensorBlockH, tensorBlockHWork, cube2Shape); + blockMmadQH.finalWaitFlags(); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube2Done); + } + + if (needRun && coreIdx < coreNum) { + GDNFwdOOffsets& cube3Offsets = cubeBlockScheduler.GetCube23Offsets(); + if(isFirstC3) Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done); + int64_t cube3OffsetAttnMask = cube3Offsets.attnWorkOffset; + int64_t cube3OffsetV = cube3Offsets.ovOffset; + int64_t cube3OffsetVWork = cube3Offsets.hvWorkOffset; + auto attenLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, cube3Offsets.blockTokens); + auto tensorAttnMask = tla::MakeTensor(gmAftermaskWorkspace[cube3OffsetAttnMask], attenLayout, Catlass::Arch::PositionGM{}); + auto tensorV = tla::MakeTensor(gmV[cube3OffsetV], vnewLayout, Catlass::Arch::PositionGM{}); + auto tensorVWork = tla::MakeTensor(gmVWorkspace[cube3OffsetVWork], ointerLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube3Shape{cube3Offsets.blockTokens, vHeadDim, cube3Offsets.blockTokens}; + auto tensorBlockAttnMask = GetTile(tensorAttnMask, tla::MakeCoord(0, 0), tla::MakeShape(cube3Shape.m(), cube3Shape.k())); + auto tensorBlockV = GetTile(tensorV, tla::MakeCoord(0, 0), tla::MakeShape(cube3Shape.k(), cube3Shape.n())); + auto tensorBlockVWork = GetTile(tensorVWork, tla::MakeCoord(0, 0), tla::MakeShape(cube3Shape.m(), cube3Shape.n())); + blockMmadAttenVNEW.preSetFlags(); + blockMmadAttenVNEW(tensorBlockAttnMask, tensorBlockV, tensorBlockVWork, cube3Shape); + blockMmadAttenVNEW.finalWaitFlags(); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube3Done); + isFirstC3 = false; + } + needRun = true; + // AscendC::PipeBarrier(); + } + if (coreIdx < coreNum) { + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done); + + } + } + + if ASCEND_IS_AIV { + + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = AscendC::GetBlockNum(); + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + AscendC::LocalTensor maskUbTensor = resource.ubBuf.template GetBufferByByte(0); + AscendC::Duplicate(maskUbTensor, (float)0.0, 64*64); + AscendC::PipeBarrier(); + for(uint32_t i = 0; i < 64; ++ i) AscendC::Duplicate(maskUbTensor[i * 64], (float)1.0, i + 1); + AscendC::PipeBarrier(); + + bool needRun = false; + uint32_t pingpongFlag = 0; + + if (coreIdx < coreNum * subBlockNum) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec1Done); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec1Done); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done); + } + + while (vecBlockScheduler.isRunning) { + vecBlockScheduler.InitTask(); + + if (vecBlockScheduler.isRunning && coreIdx < coreNum * subBlockNum) { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube1Done); + GDNFwdOOffsets& vec1Offsets = vecBlockScheduler.GetVec1Offsets(); + int64_t vec1OffsetAttnMask = vec1Offsets.attnWorkOffset; + int64_t vec1OffsetG = vec1Offsets.gOffset; + int64_t vec1OffsetAttn = vec1Offsets.attnWorkOffset; + EpilogueGDNFwdOQkmask epilogueGDNFwdOQkmask(resource); + epilogueGDNFwdOQkmask( + gmAftermaskWorkspace[vec1OffsetAttnMask], + gmG[vec1OffsetG], gmAttnWorkspace[vec1OffsetAttn], gmMask, + chunkSize, vec1Offsets.blockTokens, kHeadDim, vHeadDim, pingpongFlag, vec1Offsets.batchIdx, vec1Offsets.headIdx, vec1Offsets.chunkIdx + ); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec1Done); + } + + // AscendC::PipeBarrier(); + + if (needRun && coreIdx < coreNum * subBlockNum) { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube2Done); + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube3Done); + GDNFwdOOffsets& vec2Offsets = vecBlockScheduler.GetVec2Offsets(); + int64_t vec2OffsetO = vec2Offsets.ovOffset; + int64_t vec2OffsetG = vec2Offsets.gOffset; + int64_t vec2OffsetVWork = vec2Offsets.hvWorkOffset; + int64_t vec2OffsetHWork = vec2Offsets.hvWorkOffset; + EpilogueGDNFwdOOutput epilogueGDNFwdOOutput(resource); + epilogueGDNFwdOOutput( + gmO[vec2OffsetO], + gmG[vec2OffsetG], gmVWorkspace[vec2OffsetVWork], gmHWorkspace[vec2OffsetHWork], + scale, vec2Offsets.blockTokens, kHeadDim, vHeadDim, pingpongFlag, vec2Offsets.batchIdx, vec2Offsets.headIdx, vec2Offsets.chunkIdx + ); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done); + } + + // AscendC::PipeBarrier(); + + needRun = true; + } + } + } + +}; + +} diff --git a/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/epilogue/block/block_epilogue_gdn_fwdo_output.hpp b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/epilogue/block/block_epilogue_gdn_fwdo_output.hpp new file mode 100644 index 000000000..e74a4c44a --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/epilogue/block/block_epilogue_gdn_fwdo_output.hpp @@ -0,0 +1,394 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDO_OUTPUT_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDO_OUTPUT_HPP + +#include "catlass/catlass.hpp" +#include "catlass/arch/resource.hpp" +#include "../gdn_fwd_o_epilogue_policies.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/matrix_coord.hpp" +#include "catlass/epilogue/tile/tile_copy.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class HOutputType_, + class GInputType_, + class AInputType_, + class HInputType_ +> +class BlockEpilogue < + EpilogueAtlasGDNFwdOOutput, + HOutputType_, + GInputType_, + AInputType_, + HInputType_ +> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasGDNFwdOOutput; + using ArchTag = typename DispatchPolicy::ArchTag; + + using HElementOutput = typename HOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using AElementInput = typename AInputType_::Element; + using HElementInput = typename HInputType_::Element; + + // using CopyGmToUbInput = Tile::CopyGm2Ub; + // using CopyUbToGmOutput = Tile::CopyUb2Gm; + + static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; + static constexpr uint32_t FLOAT_ELENUM_PER_BLK = 8; + static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; + static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; + static constexpr uint32_t UB_TILE_SIZE = 16384; // 64 * 128 * 2B + static constexpr uint32_t UB_LINE_SIZE = 512; // 128 * 2 * 2B + static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; // 128 * 2 + static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; // 128 + static constexpr uint32_t MULTIPLIER = 2; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + constexpr uint32_t BASE = 0; + constexpr uint32_t MASK_UB_TENSOR_SIZE = 32 * UB_LINE_SIZE; + constexpr uint32_t GBRCLEFTCAST_UB_TENSOR_SIZE = 40 * UB_LINE_SIZE; + constexpr uint32_t GBRCUP_UB_TENSOR_SIZE = 32 * UB_LINE_SIZE; + constexpr uint32_t FLOAT_UB_TENSOR_SIZE = 32 * UB_LINE_SIZE; + constexpr uint32_t HALF_UB_TENSOR_SIZE = 16 * UB_LINE_SIZE; + constexpr uint32_t G_HALF_UB_TENSOR_SIZE = 2 * UB_LINE_SIZE; + constexpr uint32_t G_FLOAT_UB_TENSOR_SIZE = 2 * UB_LINE_SIZE; + + constexpr uint32_t MASK_UB_TENSOR_OFFSET = BASE; + constexpr uint32_t GBRCLEFTCAST_UB_TENSOR_OFFSET = MASK_UB_TENSOR_OFFSET + MASK_UB_TENSOR_SIZE; + constexpr uint32_t GBRCUP_UB_TENSOR_OFFSET = GBRCLEFTCAST_UB_TENSOR_OFFSET + GBRCLEFTCAST_UB_TENSOR_SIZE; + constexpr uint32_t GCOMP_UB_TENSOR_OFFSET = GBRCUP_UB_TENSOR_OFFSET + GBRCUP_UB_TENSOR_SIZE; + constexpr uint32_t SHARE_UB_TENSOR_OFFSET = GCOMP_UB_TENSOR_OFFSET + G_FLOAT_UB_TENSOR_SIZE; + + maskUbTensor = resource.ubBuf.template GetBufferByByte(MASK_UB_TENSOR_OFFSET); + gbrcLeftcastUbTensor = resource.ubBuf.template GetBufferByByte(GBRCLEFTCAST_UB_TENSOR_OFFSET); + gbrcUpUbTensor = resource.ubBuf.template GetBufferByByte(GBRCUP_UB_TENSOR_OFFSET); + gcompUbTensor = resource.ubBuf.template GetBufferByByte(GCOMP_UB_TENSOR_OFFSET); + shareUbTensor = resource.ubBuf.template GetBufferByByte(SHARE_UB_TENSOR_OFFSET); + + constexpr uint32_t G_UB_TENSOR_OFFSET_PING = SHARE_UB_TENSOR_OFFSET + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t G_HALF_UB_TENSOR_OFFSET_PING = G_UB_TENSOR_OFFSET_PING + G_FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t A_UB_TENSOR_OFFSET_PING = G_HALF_UB_TENSOR_OFFSET_PING + G_HALF_UB_TENSOR_SIZE; + constexpr uint32_t H_UB_TENSOR_OFFSET_PING = A_UB_TENSOR_OFFSET_PING + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t OUT_UB_TENSOR_OFFSET_PING = H_UB_TENSOR_OFFSET_PING + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t OUT_HALF_UB_TENSOR_OFFSET_PING = OUT_UB_TENSOR_OFFSET_PING + FLOAT_UB_TENSOR_SIZE; + + gUbTensorPing = resource.ubBuf.template GetBufferByByte(G_UB_TENSOR_OFFSET_PING); + gUbFPTensorPing = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PING); + gUbBFTensorPing = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PING); + aUbTensorPing = resource.ubBuf.template GetBufferByByte(A_UB_TENSOR_OFFSET_PING); + hUbTensorPing = resource.ubBuf.template GetBufferByByte(H_UB_TENSOR_OFFSET_PING); + outUbTensorPing = resource.ubBuf.template GetBufferByByte(OUT_UB_TENSOR_OFFSET_PING); + outUbFPTensorPing = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PING); + outUbBFTensorPing = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PING); + + constexpr uint32_t G_UB_TENSOR_OFFSET_PONG = OUT_HALF_UB_TENSOR_OFFSET_PING + HALF_UB_TENSOR_SIZE; + constexpr uint32_t G_HALF_UB_TENSOR_OFFSET_PONG = G_UB_TENSOR_OFFSET_PONG + G_FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t A_UB_TENSOR_OFFSET_PONG = G_HALF_UB_TENSOR_OFFSET_PONG + G_HALF_UB_TENSOR_SIZE; + constexpr uint32_t H_UB_TENSOR_OFFSET_PONG = A_UB_TENSOR_OFFSET_PONG + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t OUT_UB_TENSOR_OFFSET_PONG = H_UB_TENSOR_OFFSET_PONG + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t OUT_HALF_UB_TENSOR_OFFSET_PONG = OUT_UB_TENSOR_OFFSET_PONG + FLOAT_UB_TENSOR_SIZE; + + gUbTensorPong = resource.ubBuf.template GetBufferByByte(G_UB_TENSOR_OFFSET_PONG); + gUbFPTensorPong = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PONG); + gUbBFTensorPong = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PONG); + aUbTensorPong = resource.ubBuf.template GetBufferByByte(A_UB_TENSOR_OFFSET_PONG); + hUbTensorPong = resource.ubBuf.template GetBufferByByte(H_UB_TENSOR_OFFSET_PONG); + outUbTensorPong = resource.ubBuf.template GetBufferByByte(OUT_UB_TENSOR_OFFSET_PONG); + outUbFPTensorPong = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PONG); + outUbBFTensorPong = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PONG); + } + CATLASS_DEVICE + ~BlockEpilogue() + {} + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor hOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor attnInput, + AscendC::GlobalTensor hInput, + float scale, + uint32_t chunkSize, + uint32_t kHeadDim, + uint32_t vHeadDim, + uint32_t &pingpongFlag + , uint32_t batchIdx, uint32_t headIdx, uint32_t chunkIdx + ) + { + uint32_t mActual = chunkSize; + uint32_t nActual = vHeadDim; + uint32_t alignedM = CeilDiv(nActual, 8) * 8; + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t blockIdx = AscendC::GetBlockIdx(); + uint32_t mActualPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t mActualThisSubBlock = (subBlockIdx == 0) ? mActualPerSubBlock : (mActual - mActualPerSubBlock); + uint32_t mOffset = subBlockIdx * mActualPerSubBlock; + uint32_t nOffset = 0; + int64_t offsetA = mOffset * nActual + nOffset; + + uint32_t gbrcStart, gbrcRealStart, gbrcRealEnd, gbrcRealProcess, gbrcEffStart, gbrcEffEnd, mulsRemain, mulsRemainIdx; + if(mActualThisSubBlock <= 32) + { + if(subBlockIdx == 0) + { + gbrcStart = 0; + gbrcRealStart = 0; + gbrcRealProcess = mActualThisSubBlock; + } + else + { + gbrcStart = mActualPerSubBlock; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActual - gbrcRealStart; + } + gbrcEffStart = gbrcStart - gbrcRealStart; + uint32_t dstShape_[2] = {gbrcRealProcess, nActual}; + uint32_t srcShape_[2] = {gbrcRealProcess, 1}; + + AscendC::ResetMask(); + AscendC::GlobalTensor attnInputThisSubBlock = attnInput[gbrcStart * nActual]; + AscendC::GlobalTensor hInputThisSubBlock = hInput[gbrcStart * nActual]; + AscendC::GlobalTensor gInputThisSubBlock = gInput; + AscendC::GlobalTensor hOutputThisSubBlock = hOutput[gbrcStart * nActual]; + + AscendC::DataCopyParams gfloatUbParams{1, (uint16_t)(mActual*sizeof(float)), 0, 0}; + AscendC::DataCopyParams ghalfUbParams{1, (uint16_t)(mActual*sizeof(half)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + + AscendC::LocalTensor aUbTensor = (pingpongFlag == 0) ? aUbTensorPing : aUbTensorPong; + AscendC::LocalTensor hUbTensor = (pingpongFlag == 0) ? hUbTensorPing : hUbTensorPong; + AscendC::LocalTensor outUbTensor = (pingpongFlag == 0) ? outUbTensorPing : outUbTensorPong; + AscendC::LocalTensor outUbFPTensor = (pingpongFlag == 0) ? outUbFPTensorPing : outUbFPTensorPong; + AscendC::LocalTensor outUbBFTensor = (pingpongFlag == 0) ? outUbBFTensorPing : outUbBFTensorPong; + AscendC::LocalTensor gUbTensor = (pingpongFlag == 0) ? gUbTensorPing : gUbTensorPong; + AscendC::LocalTensor gUbFPTensor = (pingpongFlag == 0) ? gUbFPTensorPing : gUbFPTensorPong; + AscendC::LocalTensor gUbBFTensor = (pingpongFlag == 0) ? gUbBFTensorPing : gUbBFTensorPong; + + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(std::is_same::value) { + AscendC::DataCopyPad(gUbTensor, gInputThisSubBlock, gfloatUbParams, gUbPadParams); + } else { + AscendC::DataCopyPad(gUbFPTensor, gInputThisSubBlock, ghalfUbParams, gUbPadParams); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(!std::is_same::value) { + AscendC::Cast(gUbTensor, gUbFPTensor, AscendC::RoundMode::CAST_NONE, mActual); + AscendC::PipeBarrier(); + } + AscendC::Copy(gcompUbTensor, gUbTensor, 64, 2, {1, 1, 8, 8}); + AscendC::PipeBarrier(); + + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(hUbTensor, hInputThisSubBlock, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::DataCopy(aUbTensor, attnInputThisSubBlock, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + + AscendC::Exp(gcompUbTensor, gcompUbTensor, mActual); + AscendC::PipeBarrier(); + AscendC::Broadcast(gbrcLeftcastUbTensor, gcompUbTensor[gbrcRealStart], dstShape_, srcShape_, shareUbTensor); + AscendC::PipeBarrier(); + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Mul(gbrcUpUbTensor, hUbTensor, gbrcLeftcastUbTensor[gbrcEffStart*nActual], mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + + AscendC::Add(gbrcUpUbTensor, aUbTensor, gbrcUpUbTensor, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::Muls(outUbTensor, gbrcUpUbTensor, (float)scale, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + if(std::is_same::value) + { + AscendC::Cast(outUbFPTensor, outUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(hOutputThisSubBlock, outUbFPTensor, mActualThisSubBlock * nActual); + } + else + { + AscendC::Cast(outUbBFTensor, outUbTensor, AscendC::RoundMode::CAST_RINT, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(hOutputThisSubBlock, outUbBFTensor, mActualThisSubBlock * nActual); + } + pingpongFlag = 1 - pingpongFlag; + } + else + { + AscendC::ResetMask(); + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + AscendC::DataCopyParams gfloatUbParams{1, (uint16_t)(mActual*sizeof(float)), 0, 0}; + AscendC::DataCopyParams ghalfUbParams{1, (uint16_t)(mActual*sizeof(half)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + + AscendC::LocalTensor gUbTensor = (pingpongFlag == 0) ? gUbTensorPing : gUbTensorPong; + AscendC::LocalTensor gUbFPTensor = (pingpongFlag == 0) ? gUbFPTensorPing : gUbFPTensorPong; + AscendC::LocalTensor gUbBFTensor = (pingpongFlag == 0) ? gUbBFTensorPing : gUbBFTensorPong; + + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(std::is_same::value) { + AscendC::DataCopyPad(gUbTensor, gInputThisSubBlock, gfloatUbParams, gUbPadParams); + } else { + AscendC::DataCopyPad(gUbFPTensor, gInputThisSubBlock, ghalfUbParams, gUbPadParams); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(!std::is_same::value) { + AscendC::Cast(gUbTensor, gUbFPTensor, AscendC::RoundMode::CAST_NONE, mActual); + AscendC::PipeBarrier(); + } + AscendC::Copy(gcompUbTensor, gUbTensor, 64, 2, {1, 1, 8, 8}); + AscendC::PipeBarrier(); + AscendC::Exp(gcompUbTensor, gcompUbTensor, mActual); + AscendC::PipeBarrier(); + uint32_t mActualPerStage = CeilDiv(mActualThisSubBlock, 2); + uint32_t mActualThisStage = 0; + for(uint32_t stage = 0; stage < 2; stage++) + { + if(stage == 0) mActualThisStage = mActualPerStage; + else mActualThisStage = mActualThisSubBlock - mActualPerStage; + + if(subBlockIdx == 0 && stage == 0) + { + gbrcStart = 0; + gbrcRealStart = 0; + gbrcRealProcess = mActualThisStage; + } + else if(subBlockIdx == 0 && stage == 1) + { + gbrcStart = mActualPerStage; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActualThisSubBlock - gbrcRealStart; + } + else if(subBlockIdx == 1 && stage == 0) + { + gbrcStart = mActualPerSubBlock; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActualPerSubBlock + mActualThisStage - gbrcRealStart; + } + else if(subBlockIdx == 1 && stage == 1) + { + gbrcStart = mActualPerSubBlock + mActualPerStage; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActual - gbrcRealStart; + } + gbrcEffStart = gbrcStart - gbrcRealStart; + uint32_t dstShape_[2] = {gbrcRealProcess, nActual}; + uint32_t srcShape_[2] = {gbrcRealProcess, 1}; + + AscendC::GlobalTensor hOutputThisSubBlock = hOutput[gbrcStart * nActual]; + AscendC::GlobalTensor attnInputThisSubBlock = attnInput[gbrcStart * nActual]; + AscendC::GlobalTensor hInputThisSubBlock = hInput[gbrcStart * nActual]; + + AscendC::LocalTensor aUbTensor = (pingpongFlag == 0) ? aUbTensorPing : aUbTensorPong; + AscendC::LocalTensor hUbTensor = (pingpongFlag == 0) ? hUbTensorPing : hUbTensorPong; + AscendC::LocalTensor outUbTensor = (pingpongFlag == 0) ? outUbTensorPing : outUbTensorPong; + AscendC::LocalTensor outUbFPTensor = (pingpongFlag == 0) ? outUbFPTensorPing : outUbFPTensorPong; + AscendC::LocalTensor outUbBFTensor = (pingpongFlag == 0) ? outUbBFTensorPing : outUbBFTensorPong; + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(hUbTensor, hInputThisSubBlock, mActualThisStage * nActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::DataCopy(aUbTensor, attnInputThisSubBlock, mActualThisStage * nActual); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + + AscendC::Broadcast(gbrcLeftcastUbTensor, gcompUbTensor[gbrcRealStart], dstShape_, srcShape_, shareUbTensor); + AscendC::PipeBarrier(); + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Mul(gbrcUpUbTensor, hUbTensor, gbrcLeftcastUbTensor[gbrcEffStart*nActual], mActualThisStage * nActual); + AscendC::PipeBarrier(); + + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::Add(gbrcUpUbTensor, aUbTensor, gbrcUpUbTensor, mActualThisStage * nActual); + AscendC::PipeBarrier(); + AscendC::Muls(outUbTensor, gbrcUpUbTensor, (float)scale, mActualThisStage * nActual); + AscendC::PipeBarrier(); + + if(std::is_same::value) + { + AscendC::Cast(outUbFPTensor, outUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisStage * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(hOutputThisSubBlock, outUbFPTensor, mActualThisStage * nActual); + } + else + { + AscendC::Cast(outUbBFTensor, outUbTensor, AscendC::RoundMode::CAST_RINT, mActualThisStage * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(hOutputThisSubBlock, outUbBFTensor, mActualThisStage * nActual); + } + pingpongFlag = 1 - pingpongFlag; + } + } + } + +private: + AscendC::LocalTensor maskUbTensor; + AscendC::LocalTensor gbrcLeftcastUbTensor; + AscendC::LocalTensor gbrcUpUbTensor; + AscendC::LocalTensor gcompUbTensor; + AscendC::LocalTensor shareUbTensor; + + AscendC::LocalTensor gUbTensorPing; + AscendC::LocalTensor gUbFPTensorPing; + AscendC::LocalTensor gUbBFTensorPing; + AscendC::LocalTensor aUbTensorPing; + AscendC::LocalTensor hUbTensorPing; + AscendC::LocalTensor outUbTensorPing; + AscendC::LocalTensor outUbFPTensorPing; + AscendC::LocalTensor outUbBFTensorPing; + + AscendC::LocalTensor gUbTensorPong; + AscendC::LocalTensor gUbFPTensorPong; + AscendC::LocalTensor gUbBFTensorPong; + AscendC::LocalTensor aUbTensorPong; + AscendC::LocalTensor hUbTensorPong; + AscendC::LocalTensor outUbTensorPong; + AscendC::LocalTensor outUbFPTensorPong; + AscendC::LocalTensor outUbBFTensorPong; + + +}; +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/epilogue/block/block_epilogue_gdn_fwdo_qkmask.hpp b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/epilogue/block/block_epilogue_gdn_fwdo_qkmask.hpp new file mode 100644 index 000000000..de1a6ede1 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/epilogue/block/block_epilogue_gdn_fwdo_qkmask.hpp @@ -0,0 +1,428 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDO_QKMASK_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDO_QKMASK_HPP + +#include "catlass/catlass.hpp" +#include "catlass/arch/resource.hpp" +#include "../gdn_fwd_o_epilogue_policies.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/matrix_coord.hpp" +#include "catlass/epilogue/tile/tile_copy.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class AOutputType_, + class GInputType_, + class AInputType_, + class MaskInputType_ +> +class BlockEpilogue < + EpilogueAtlasGDNFwdOQkmask, + AOutputType_, + GInputType_, + AInputType_, + MaskInputType_ +> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasGDNFwdOQkmask; + using ArchTag = typename DispatchPolicy::ArchTag; + + using AElementOutput = typename AOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using AElementInput = typename AInputType_::Element; + using MaskElementInput = typename MaskInputType_::Element; + + static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; + static constexpr uint32_t FLOAT_ELENUM_PER_BLK = 8; + static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; + static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; + static constexpr uint32_t UB_TILE_SIZE = 16384; // 64 * 128 * 2B + static constexpr uint32_t UB_LINE_SIZE = 512; // 128 * 2 * 2B + static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; // 128 * 2 + static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; // 128 + static constexpr uint32_t MULTIPLIER = 2; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + constexpr uint32_t BASE = 0; + constexpr uint32_t MASK_UB_TENSOR_SIZE = 32 * UB_LINE_SIZE; + constexpr uint32_t GBRCLEFTCAST_UB_TENSOR_SIZE = 40 * UB_LINE_SIZE; + constexpr uint32_t GBRCUP_UB_TENSOR_SIZE = 32 * UB_LINE_SIZE; + constexpr uint32_t FLOAT_UB_TENSOR_SIZE = 32 * UB_LINE_SIZE; + constexpr uint32_t HALF_UB_TENSOR_SIZE = 16 * UB_LINE_SIZE; + constexpr uint32_t G_HALF_UB_TENSOR_SIZE = 2 * UB_LINE_SIZE; + constexpr uint32_t G_FLOAT_UB_TENSOR_SIZE = 2 * UB_LINE_SIZE; + + constexpr uint32_t MASK_UB_TENSOR_OFFSET = BASE; + constexpr uint32_t GBRCLEFTCAST_UB_TENSOR_OFFSET = MASK_UB_TENSOR_OFFSET + MASK_UB_TENSOR_SIZE; + constexpr uint32_t GBRCUP_UB_TENSOR_OFFSET = GBRCLEFTCAST_UB_TENSOR_OFFSET + GBRCLEFTCAST_UB_TENSOR_SIZE; + constexpr uint32_t GCOMP_UB_TENSOR_OFFSET = GBRCUP_UB_TENSOR_OFFSET + GBRCUP_UB_TENSOR_SIZE; + constexpr uint32_t SHARE_UB_TENSOR_OFFSET = GCOMP_UB_TENSOR_OFFSET + G_FLOAT_UB_TENSOR_SIZE; + + maskUbTensor = resource.ubBuf.template GetBufferByByte(MASK_UB_TENSOR_OFFSET); + gbrcLeftcastUbTensor = resource.ubBuf.template GetBufferByByte(GBRCLEFTCAST_UB_TENSOR_OFFSET); + gbrcUpUbTensor = resource.ubBuf.template GetBufferByByte(GBRCUP_UB_TENSOR_OFFSET); + gcompUbTensor = resource.ubBuf.template GetBufferByByte(GCOMP_UB_TENSOR_OFFSET); + shareUbTensor = resource.ubBuf.template GetBufferByByte(SHARE_UB_TENSOR_OFFSET); + + constexpr uint32_t G_UB_TENSOR_OFFSET_PING = SHARE_UB_TENSOR_OFFSET + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t G_HALF_UB_TENSOR_OFFSET_PING = G_UB_TENSOR_OFFSET_PING + G_FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t A_UB_TENSOR_OFFSET_PING = G_HALF_UB_TENSOR_OFFSET_PING + G_HALF_UB_TENSOR_SIZE; + constexpr uint32_t OUT_UB_TENSOR_OFFSET_PING = A_UB_TENSOR_OFFSET_PING + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t OUT_HALF_UB_TENSOR_OFFSET_PING = OUT_UB_TENSOR_OFFSET_PING + FLOAT_UB_TENSOR_SIZE; + + gUbTensorPing = resource.ubBuf.template GetBufferByByte(G_UB_TENSOR_OFFSET_PING); + gUbFPTensorPing = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PING); + gUbBFTensorPing = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PING); + aUbTensorPing = resource.ubBuf.template GetBufferByByte(A_UB_TENSOR_OFFSET_PING); + outUbTensorPing = resource.ubBuf.template GetBufferByByte(OUT_UB_TENSOR_OFFSET_PING); + outUbFPTensorPing = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PING); + outUbBFTensorPing = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PING); + + constexpr uint32_t G_UB_TENSOR_OFFSET_PONG = 32 * UB_LINE_SIZE + OUT_HALF_UB_TENSOR_OFFSET_PING + HALF_UB_TENSOR_SIZE; + constexpr uint32_t G_HALF_UB_TENSOR_OFFSET_PONG = G_UB_TENSOR_OFFSET_PONG + G_FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t A_UB_TENSOR_OFFSET_PONG = G_HALF_UB_TENSOR_OFFSET_PONG + G_HALF_UB_TENSOR_SIZE; + constexpr uint32_t OUT_UB_TENSOR_OFFSET_PONG = A_UB_TENSOR_OFFSET_PONG + FLOAT_UB_TENSOR_SIZE; + constexpr uint32_t OUT_HALF_UB_TENSOR_OFFSET_PONG = OUT_UB_TENSOR_OFFSET_PONG + FLOAT_UB_TENSOR_SIZE; + + gUbTensorPong = resource.ubBuf.template GetBufferByByte(G_UB_TENSOR_OFFSET_PONG); + gUbFPTensorPong = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PONG); + gUbBFTensorPong = resource.ubBuf.template GetBufferByByte(G_HALF_UB_TENSOR_OFFSET_PONG); + aUbTensorPong = resource.ubBuf.template GetBufferByByte(A_UB_TENSOR_OFFSET_PONG); + outUbTensorPong = resource.ubBuf.template GetBufferByByte(OUT_UB_TENSOR_OFFSET_PONG); + outUbFPTensorPong = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PONG); + outUbBFTensorPong = resource.ubBuf.template GetBufferByByte(OUT_HALF_UB_TENSOR_OFFSET_PONG); + } + + CATLASS_DEVICE + ~BlockEpilogue() + {} + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor maskOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor attnInput, + AscendC::GlobalTensor boolInput, + uint32_t fullChunkSize, + uint32_t chunkSize, + uint32_t kHeadDim, + uint32_t vHeadDim, + uint32_t &pingpongFlag + , uint32_t batchIdx, uint32_t headIdx, uint32_t chunkIdx + ) + { + uint32_t mActual = chunkSize; + uint32_t nActual = chunkSize; + uint32_t alignedNActual = CeilDiv(nActual, 16) * 16; + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t blockIdx = AscendC::GetBlockIdx(); + uint32_t mActualPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t mActualThisSubBlock = (subBlockIdx == 0) ? mActualPerSubBlock : (mActual - mActualPerSubBlock); + uint32_t mOffset = subBlockIdx * mActualPerSubBlock; + uint32_t nOffset = 0; + int64_t offsetA = mOffset * nActual + nOffset; + uint16_t aInputDstStride; + if((nActual - 1) % 16 <= 7) aInputDstStride = 1; + else aInputDstStride = 0; + + uint32_t gbrcStart, gbrcRealStart, gbrcRealEnd, gbrcRealProcess, gbrcEffStart, gbrcEffEnd, mulsRemain, mulsRemainIdx; + if(mActualThisSubBlock <= 32) + { if(subBlockIdx == 0) + { + gbrcStart = 0; + gbrcRealStart = 0; + gbrcRealProcess = mActualThisSubBlock; + } + else + { + gbrcStart = mActualPerSubBlock; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActual - gbrcRealStart; + } + + gbrcEffStart = gbrcStart - gbrcRealStart; + gbrcEffEnd = gbrcEffStart + mActualThisSubBlock; + + uint32_t dstUpShape_[2] = {mActualThisSubBlock, alignedNActual}; + uint32_t srcUpShape_[2] = {1, alignedNActual}; + uint32_t dstLeftShape_[2] = {gbrcRealProcess, alignedNActual}; + uint32_t srcLeftShape_[2] = {gbrcRealProcess, 1}; + + AscendC::ResetMask(); + AscendC::GlobalTensor maskOutputThisSubBlock = maskOutput[gbrcStart * nActual]; + AscendC::GlobalTensor attnInputThisSubBlock = attnInput[gbrcStart * nActual]; + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + + AscendC::DataCopyParams aInputUbParams{(uint16_t)mActualThisSubBlock, (uint16_t)(nActual*sizeof(float)), 0, aInputDstStride}; + AscendC::DataCopyPadParams aInputUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams aOutputUbParams{(uint16_t)mActualThisSubBlock, (uint32_t)(nActual*sizeof(half)), 0, 0, 0}; + + AscendC::DataCopyParams gfloatUbParams{1, (uint16_t)(mActual*sizeof(float)), 0, 0}; + AscendC::DataCopyParams ghalfUbParams{1, (uint16_t)(mActual*sizeof(half)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + + AscendC::LocalTensor aUbTensor = (pingpongFlag == 0) ? aUbTensorPing : aUbTensorPong; + AscendC::LocalTensor outUbTensor = (pingpongFlag == 0) ? outUbTensorPing : outUbTensorPong; + AscendC::LocalTensor outUbFPTensor = (pingpongFlag == 0) ? outUbFPTensorPing : outUbFPTensorPong; + AscendC::LocalTensor outUbBFTensor = (pingpongFlag == 0) ? outUbBFTensorPing : outUbBFTensorPong; + AscendC::LocalTensor gUbTensor = (pingpongFlag == 0) ? gUbTensorPing : gUbTensorPong; + AscendC::LocalTensor gUbFPTensor = (pingpongFlag == 0) ? gUbFPTensorPing : gUbFPTensorPong; + AscendC::LocalTensor gUbBFTensor = (pingpongFlag == 0) ? gUbBFTensorPing : gUbBFTensorPong; + + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(std::is_same::value) { + AscendC::DataCopyPad(gUbTensor, gInputThisSubBlock, gfloatUbParams, gUbPadParams); + } else { + AscendC::DataCopyPad(gUbFPTensor, gInputThisSubBlock, ghalfUbParams, gUbPadParams); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(!std::is_same::value) { + AscendC::Cast(gUbTensor, gUbFPTensor, AscendC::RoundMode::CAST_NONE, mActual); + AscendC::PipeBarrier(); + } + AscendC::Copy(gcompUbTensor, gUbTensor, 64, 2, {1, 1, 8, 8}); + AscendC::PipeBarrier(); + + AscendC::Broadcast(gbrcUpUbTensor, gcompUbTensor, dstUpShape_, srcUpShape_, shareUbTensor); + AscendC::Broadcast(gbrcLeftcastUbTensor, gcompUbTensor[gbrcRealStart], dstLeftShape_, srcLeftShape_, shareUbTensor); + AscendC::PipeBarrier(); + AscendC::Sub(gbrcUpUbTensor, gbrcLeftcastUbTensor[gbrcEffStart*alignedNActual], gbrcUpUbTensor, mActualThisSubBlock * alignedNActual); + AscendC::PipeBarrier(); + AscendC::Mins(gbrcUpUbTensor, gbrcUpUbTensor, (float)0.0, mActualThisSubBlock * alignedNActual); + AscendC::PipeBarrier(); + AscendC::Exp(gbrcUpUbTensor, gbrcUpUbTensor, mActualThisSubBlock * alignedNActual); + AscendC::PipeBarrier(); + + gbrcRealEnd = CeilDiv(gbrcStart + mActualThisSubBlock, 8) * 8; + AscendC::Mul(gbrcUpUbTensor[gbrcRealStart], gbrcUpUbTensor[gbrcRealStart], maskUbTensor[gbrcEffStart * 64], gbrcRealEnd - gbrcRealStart, mActualThisSubBlock, + {1, 1, 1, static_cast(alignedNActual/8), static_cast(alignedNActual/8), static_cast(64/8)}); + AscendC::PipeBarrier(); + + mulsRemain = alignedNActual - gbrcRealEnd; + mulsRemainIdx = gbrcRealEnd; + while(mulsRemain > 64) + { + AscendC::Muls(gbrcUpUbTensor[mulsRemainIdx], gbrcUpUbTensor[mulsRemainIdx], (float)0.0, 64, mActualThisSubBlock, + {1, 1, static_cast(alignedNActual/8), static_cast(alignedNActual/8)}); + mulsRemain -= 64; + mulsRemainIdx += 64; + } + AscendC::Muls(gbrcUpUbTensor[mulsRemainIdx], gbrcUpUbTensor[mulsRemainIdx], (float)0.0, mulsRemain, mActualThisSubBlock, + {1, 1, static_cast(alignedNActual/8), static_cast(alignedNActual/8)}); + AscendC::PipeBarrier(); + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if(chunkSize==fullChunkSize) AscendC::DataCopy(aUbTensor, attnInputThisSubBlock, mActualThisSubBlock*nActual); + else AscendC::DataCopyPad(aUbTensor, attnInputThisSubBlock, aInputUbParams, aInputUbPadParams); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Mul(outUbTensor, aUbTensor, gbrcUpUbTensor, mActualThisSubBlock * alignedNActual); + AscendC::PipeBarrier(); + + if(std::is_same::value) + { + AscendC::Cast(outUbFPTensor, outUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * alignedNActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if(chunkSize==fullChunkSize) AscendC::DataCopy(maskOutputThisSubBlock, outUbFPTensor, mActualThisSubBlock*nActual); + else AscendC::DataCopyPad(maskOutputThisSubBlock, outUbFPTensor, aOutputUbParams); + } + else + { + AscendC::Cast(outUbBFTensor, outUbTensor, AscendC::RoundMode::CAST_RINT, mActualThisSubBlock * alignedNActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if(chunkSize==fullChunkSize) AscendC::DataCopy(maskOutputThisSubBlock, outUbBFTensor, mActualThisSubBlock*nActual); + else AscendC::DataCopyPad(maskOutputThisSubBlock, outUbBFTensor, aOutputUbParams); + } + pingpongFlag = 1 - pingpongFlag; + } + else // mActualThisSubBlock > 32 ; <=64 + { + AscendC::ResetMask(); + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + AscendC::DataCopyParams gfloatUbParams{1, (uint16_t)(mActual*sizeof(float)), 0, 0}; + AscendC::DataCopyParams ghalfUbParams{1, (uint16_t)(mActual*sizeof(half)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + + AscendC::LocalTensor gUbTensor = (pingpongFlag == 0) ? gUbTensorPing : gUbTensorPong; + AscendC::LocalTensor gUbFPTensor = (pingpongFlag == 0) ? gUbFPTensorPing : gUbFPTensorPong; + AscendC::LocalTensor gUbBFTensor = (pingpongFlag == 0) ? gUbBFTensorPing : gUbBFTensorPong; + + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(std::is_same::value) { + AscendC::DataCopyPad(gUbTensor, gInputThisSubBlock, gfloatUbParams, gUbPadParams); + } else { + AscendC::DataCopyPad(gUbFPTensor, gInputThisSubBlock, ghalfUbParams, gUbPadParams); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr(!std::is_same::value) { + AscendC::Cast(gUbTensor, gUbFPTensor, AscendC::RoundMode::CAST_NONE, mActual); + AscendC::PipeBarrier(); + } + AscendC::Copy(gcompUbTensor, gUbTensor, 64, 2, {1, 1, 8, 8}); + AscendC::PipeBarrier(); + + uint32_t mActualPerStage = CeilDiv(mActualThisSubBlock, 2); + uint32_t mActualThisStage = 0; + for(uint32_t stage = 0; stage < 2; ++stage) + { + if(stage==0) mActualThisStage = mActualPerStage; + else mActualThisStage = mActualThisSubBlock - mActualPerStage; + + if(subBlockIdx == 0 && stage == 0) + { + gbrcStart = 0; + gbrcRealStart = 0; + gbrcRealProcess = mActualThisStage; + } + else if(subBlockIdx == 0 && stage == 1) + { + gbrcStart = mActualPerStage; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActualThisSubBlock - gbrcRealStart; + } + else if(subBlockIdx == 1 && stage == 0) + { + gbrcStart = mActualPerSubBlock; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActualPerSubBlock + mActualThisStage - gbrcRealStart; + } + else if(subBlockIdx == 1 && stage == 1) + { + gbrcStart = mActualPerSubBlock + mActualPerStage; + gbrcRealStart = gbrcStart & ~7; + gbrcRealProcess = mActual - gbrcRealStart; + } + + gbrcEffStart = gbrcStart - gbrcRealStart; + + AscendC::GlobalTensor maskOutputThisSubBlock = maskOutput[gbrcStart * nActual]; + AscendC::GlobalTensor attnInputThisSubBlock = attnInput[gbrcStart * nActual]; + + AscendC::DataCopyParams aInputUbParams{(uint16_t)mActualThisStage, (uint16_t)(nActual*sizeof(float)), 0, aInputDstStride}; + AscendC::DataCopyPadParams aInputUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyExtParams aOutputUbParams{(uint16_t)mActualThisStage, (uint32_t)(nActual*sizeof(half)), 0, 0, 0}; + + AscendC::LocalTensor aUbTensor = (pingpongFlag == 0) ? aUbTensorPing : aUbTensorPong; + AscendC::LocalTensor outUbTensor = (pingpongFlag == 0) ? outUbTensorPing : outUbTensorPong; + AscendC::LocalTensor outUbFPTensor = (pingpongFlag == 0) ? outUbFPTensorPing : outUbFPTensorPong; + AscendC::LocalTensor outUbBFTensor = (pingpongFlag == 0) ? outUbBFTensorPing : outUbBFTensorPong; + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if(chunkSize==fullChunkSize) AscendC::DataCopy(aUbTensor, attnInputThisSubBlock, mActualThisStage*nActual); + else AscendC::DataCopyPad(aUbTensor, attnInputThisSubBlock, aInputUbParams, aInputUbPadParams); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + uint32_t dstUpShape_[2] = {mActualThisStage, alignedNActual}; + uint32_t srcUpShape_[2] = {1, alignedNActual}; + uint32_t dstLeftShape_[2] = {gbrcRealProcess, alignedNActual}; + uint32_t srcLeftShape_[2] = {gbrcRealProcess, 1}; + + AscendC::Broadcast(gbrcUpUbTensor, gcompUbTensor, dstUpShape_, srcUpShape_, shareUbTensor); + AscendC::Broadcast(gbrcLeftcastUbTensor, gcompUbTensor[gbrcRealStart], dstLeftShape_, srcLeftShape_, shareUbTensor); + AscendC::PipeBarrier(); + AscendC::Sub(gbrcUpUbTensor, gbrcLeftcastUbTensor[gbrcEffStart*alignedNActual], gbrcUpUbTensor, mActualThisStage * alignedNActual); + AscendC::PipeBarrier(); + AscendC::Mins(gbrcUpUbTensor, gbrcUpUbTensor, (float)0.0, mActualThisStage * alignedNActual); + AscendC::PipeBarrier(); + AscendC::Exp(gbrcUpUbTensor, gbrcUpUbTensor, mActualThisStage * alignedNActual); + AscendC::PipeBarrier(); + + gbrcRealEnd = CeilDiv(gbrcStart + mActualThisStage, 8) * 8; + AscendC::Mul(gbrcUpUbTensor[gbrcRealStart], gbrcUpUbTensor[gbrcRealStart], maskUbTensor[gbrcEffStart * 64], gbrcRealEnd - gbrcRealStart, mActualThisStage, + {1, 1, 1, static_cast(alignedNActual/8), static_cast(alignedNActual/8), static_cast(64/8)}); + AscendC::PipeBarrier(); + mulsRemain = alignedNActual - gbrcRealEnd; + mulsRemainIdx = gbrcRealEnd; + while(mulsRemain > 64) + { + AscendC::Muls(gbrcUpUbTensor[mulsRemainIdx], gbrcUpUbTensor[mulsRemainIdx], (float)0.0, 64, mActualThisStage, + {1, 1, static_cast(alignedNActual/8), static_cast(alignedNActual/8)}); + mulsRemain -= 64; + mulsRemainIdx += 64; + } + AscendC::Muls(gbrcUpUbTensor[mulsRemainIdx], gbrcUpUbTensor[mulsRemainIdx], (float)0.0, mulsRemain, mActualThisStage, + {1, 1, static_cast(alignedNActual/8), static_cast(alignedNActual/8)}); + AscendC::PipeBarrier(); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Mul(outUbTensor, aUbTensor, gbrcUpUbTensor, mActualThisStage * alignedNActual); + AscendC::PipeBarrier(); + if(std::is_same::value) + { + AscendC::Cast(outUbFPTensor, outUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisStage * alignedNActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if(chunkSize==fullChunkSize) AscendC::DataCopy(maskOutputThisSubBlock, outUbFPTensor, mActualThisStage*nActual); + else AscendC::DataCopyPad(maskOutputThisSubBlock, outUbFPTensor, aOutputUbParams); + } + else + { + AscendC::Cast(outUbBFTensor, outUbTensor, AscendC::RoundMode::CAST_RINT, mActualThisStage * alignedNActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if(chunkSize==fullChunkSize) AscendC::DataCopy(maskOutputThisSubBlock, outUbBFTensor, mActualThisStage*nActual); + else AscendC::DataCopyPad(maskOutputThisSubBlock, outUbBFTensor, aOutputUbParams); + } + pingpongFlag = 1 - pingpongFlag; + } + } + + } + +private: + AscendC::LocalTensor maskUbTensor; + AscendC::LocalTensor gbrcLeftcastUbTensor; + AscendC::LocalTensor gbrcUpUbTensor; + AscendC::LocalTensor gcompUbTensor; + AscendC::LocalTensor shareUbTensor; + + AscendC::LocalTensor gUbTensorPing; + AscendC::LocalTensor gUbFPTensorPing; + AscendC::LocalTensor gUbBFTensorPing; + AscendC::LocalTensor aUbTensorPing; + AscendC::LocalTensor outUbTensorPing; + AscendC::LocalTensor outUbFPTensorPing; + AscendC::LocalTensor outUbBFTensorPing; + + AscendC::LocalTensor gUbTensorPong; + AscendC::LocalTensor gUbFPTensorPong; + AscendC::LocalTensor gUbBFTensorPong; + AscendC::LocalTensor aUbTensorPong; + AscendC::LocalTensor outUbTensorPong; + AscendC::LocalTensor outUbFPTensorPong; + AscendC::LocalTensor outUbBFTensorPong; + +}; +} + +#endif diff --git a/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/epilogue/gdn_fwd_o_epilogue_policies.hpp b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/epilogue/gdn_fwd_o_epilogue_policies.hpp new file mode 100644 index 000000000..c800a3af3 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/epilogue/gdn_fwd_o_epilogue_policies.hpp @@ -0,0 +1,35 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_EPILOGUE_GDN_FWD_O_EPILOGUE_POLICIES_HPP +#define CATLASS_EPILOGUE_GDN_FWD_O_EPILOGUE_POLICIES_HPP + +#include "catlass/catlass.hpp" + +namespace Catlass::Epilogue { + +struct EpilogueAtlasGDNFwdOQkmask { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + using ArchTag = Arch::Ascend950; +#else + using ArchTag = Arch::AtlasA2; +#endif +}; + +struct EpilogueAtlasGDNFwdOOutput { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + using ArchTag = Arch::Ascend950; +#else + using ArchTag = Arch::AtlasA2; +#endif +}; + +} // namespace Catlass::Epilogue + +#endif // CATLASS_EPILOGUE_GDN_FWD_O_EPILOGUE_POLICIES_HPP diff --git a/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/gemm/block/block_scheduler_gdn_fwd_o.hpp b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/gemm/block/block_scheduler_gdn_fwd_o.hpp new file mode 100644 index 000000000..c8010d23f --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/gemm/block/block_scheduler_gdn_fwd_o.hpp @@ -0,0 +1,275 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_GEMM_SCHEDULER_GDN_FWD_O_HPP +#define CATLASS_GEMM_SCHEDULER_GDN_FWD_O_HPP + +// constexpr uint32_t PING_PONG_STAGES = 1; +constexpr uint32_t PING_PONG_STAGES = 2; +constexpr uint32_t BYTE_SIZE_16_BIT = 2; + +template +CATLASS_DEVICE T AlignUp(T a, T b) { + return (b == 0) ? 0 : (a + b - 1) / b * b; +} + +template +CATLASS_DEVICE T Min(T a, T b) { + return (a > b) ? b : a; +} + +template +CATLASS_DEVICE T Max(T a, T b) { + return (a > b) ? a : b; +} + +namespace Catlass::Gemm::Block { + + +struct GDNFwdOOffsets { + uint32_t qkOffset; + uint32_t ovOffset; + uint32_t hOffset; + uint32_t gOffset; + uint32_t attnWorkOffset; + uint32_t hvWorkOffset; + bool isFinalState; + uint32_t blockTokens; + // for debug + uint32_t batchIdx; + uint32_t headIdx; + uint32_t chunkIdx; + +}; + +struct BlockSchedulerGdnFwdO { + uint32_t shapeBatch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + uint32_t isVariedLen; + uint32_t tokenBatch; + uint32_t numChunks{0}; + uint32_t vBlockSize{128}; + + uint32_t taskIdx; + uint32_t cubeCoreIdx; + uint32_t cubeCoreNum; + uint32_t vLoops; + uint32_t taskNum; + uint32_t headGroups; + + bool isRunning; + bool processNewTask {true}; + bool firstLoop {true}; + bool lastLoop {false}; + GDNFwdOOffsets offsets[PING_PONG_STAGES]; + int32_t currStage{PING_PONG_STAGES - 1}; + + uint32_t vIdx; + uint32_t batchIdx; + uint32_t baseHeadIdx; + uint32_t chunkIdx; + uint32_t headInnerIdx; + uint32_t vHeadIdx; + uint32_t kHeadIdx; + uint32_t shapeBatchIdx; + uint32_t tokenBatchIdx; + + uint32_t batchChunkIdx; + uint32_t batchChunkStartIdx; + uint32_t tokenOffset; + uint32_t batchChunks; + uint32_t batchTokens; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmChunkOffsets; + + Arch::CrossCoreFlag cube1Done{3}; + Arch::CrossCoreFlag vec1Done{4}; + Arch::CrossCoreFlag cube2Done{5}; + Arch::CrossCoreFlag cube3Done{6}; + Arch::CrossCoreFlag vec2Done{7}; + + CATLASS_DEVICE + BlockSchedulerGdnFwdO() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_offsets, GM_ADDR tiling, uint32_t coreIdx, uint32_t coreNum) { + __gm__ ChunkFwdOTilingData *__restrict gdnFwdOTilingData = reinterpret_cast<__gm__ ChunkFwdOTilingData *__restrict>(tiling); + shapeBatch = gdnFwdOTilingData->shapeBatch; + seqlen = gdnFwdOTilingData->seqlen; + kNumHead = gdnFwdOTilingData->kNumHead; + vNumHead = gdnFwdOTilingData->vNumHead; + kHeadDim = gdnFwdOTilingData->kHeadDim; + vHeadDim = gdnFwdOTilingData->vHeadDim; + chunkSize = gdnFwdOTilingData->chunkSize; + isVariedLen = gdnFwdOTilingData->isVariedLen; + tokenBatch = gdnFwdOTilingData->tokenBatch; + + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmChunkOffsets.SetGlobalBuffer((__gm__ int64_t *)chunk_offsets); + + if (isVariedLen) { + for (uint32_t b = 1; b <= tokenBatch; b++) { + numChunks += (gmSeqlen.GetValue(b) - gmSeqlen.GetValue(b - 1) + chunkSize - 1) / chunkSize; + } + } else { + numChunks = (seqlen + chunkSize - 1) / chunkSize; + } + + cubeCoreIdx = coreIdx; + cubeCoreNum = coreNum; + vLoops = vHeadDim / vBlockSize; + taskNum = vLoops * shapeBatch * numChunks * vNumHead; + headGroups = vNumHead / kNumHead; + taskIdx = cubeCoreIdx * PING_PONG_STAGES; + isRunning = taskIdx < taskNum; + + } + + CATLASS_DEVICE + void InitTask() { + if (processNewTask) { + if (unlikely(taskIdx >= taskNum)) { + isRunning = false; + } + vIdx = taskIdx / (shapeBatch * numChunks * vNumHead); + shapeBatchIdx = (taskIdx - vIdx * shapeBatch * numChunks * vNumHead) / (numChunks * vNumHead); + chunkIdx = (taskIdx - vIdx * shapeBatch * numChunks * vNumHead - shapeBatchIdx * numChunks * vNumHead) / vNumHead; + baseHeadIdx = taskIdx % vNumHead; + tokenBatchIdx = isVariedLen ? gmChunkOffsets.GetValue(2 * chunkIdx) : 0; + batchChunkIdx = isVariedLen ? gmChunkOffsets.GetValue(2 * chunkIdx + 1) : chunkIdx; + batchChunkStartIdx = chunkIdx - batchChunkIdx; + tokenOffset = isVariedLen ? gmSeqlen.GetValue(tokenBatchIdx) : 0; + batchTokens = isVariedLen ? (gmSeqlen.GetValue(tokenBatchIdx + 1) - tokenOffset) : seqlen; + headInnerIdx = 0; + } else { + headInnerIdx = (headInnerIdx + 1) % PING_PONG_STAGES; + } + + vHeadIdx = baseHeadIdx + headInnerIdx; + kHeadIdx = vHeadIdx / headGroups; + offsets[currStage].qkOffset = (shapeBatchIdx * kNumHead * seqlen + kHeadIdx * seqlen + tokenOffset + batchChunkIdx * chunkSize) * kHeadDim; + offsets[currStage].ovOffset = (shapeBatchIdx * vNumHead * seqlen + vHeadIdx * seqlen + tokenOffset + batchChunkIdx * chunkSize) * vHeadDim; + offsets[currStage].hOffset = (shapeBatchIdx * vNumHead * numChunks + vHeadIdx * numChunks + chunkIdx) * kHeadDim * vHeadDim; + offsets[currStage].gOffset = shapeBatchIdx * vNumHead * seqlen + vHeadIdx * seqlen + tokenOffset + batchChunkIdx * chunkSize; + offsets[currStage].attnWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + currStage) * chunkSize * chunkSize; + offsets[currStage].hvWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + currStage) * chunkSize * vHeadDim; + offsets[currStage].isFinalState = chunkIdx == (numChunks - 1) || (isVariedLen && gmChunkOffsets.GetValue(2 * chunkIdx + 3) == 0); + offsets[currStage].blockTokens = offsets[currStage].isFinalState ? (batchTokens - batchChunkIdx * chunkSize) : chunkSize; + offsets[currStage].batchIdx = batchIdx; + offsets[currStage].headIdx = vHeadIdx; + offsets[currStage].chunkIdx = chunkIdx; + + processNewTask = headInnerIdx == PING_PONG_STAGES - 1; + if (processNewTask) { + taskIdx += PING_PONG_STAGES * cubeCoreNum; + } + + currStage = (currStage + 1) % PING_PONG_STAGES; + } + + +}; + +struct BlockSchedulerGdnFwdOCube : public BlockSchedulerGdnFwdO { + CATLASS_DEVICE + BlockSchedulerGdnFwdOCube() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_offsets, GM_ADDR tiling) { + BlockSchedulerGdnFwdO::Init(cu_seqlens, chunk_offsets, tiling, AscendC::GetBlockIdx(), AscendC::GetBlockNum()); + } + + CATLASS_DEVICE + bool NeedProcessCube1() { + return true; + } + + CATLASS_DEVICE + GDNFwdOOffsets& GetCube1Offsets() { + return offsets[(currStage - 1) % PING_PONG_STAGES]; + } + + CATLASS_DEVICE + GemmCoord GetCube1Shape() { + GDNFwdOOffsets& cube1Offsets = GetCube1Offsets(); + return GemmCoord{cube1Offsets.blockTokens, cube1Offsets.blockTokens, kHeadDim}; + } + + CATLASS_DEVICE + bool NeedProcessCube23() { + if (unlikely(firstLoop)) { + firstLoop = false; + return false; + } + return true; + } + + CATLASS_DEVICE + GDNFwdOOffsets& GetCube23Offsets() { + return offsets[(currStage - 2) % PING_PONG_STAGES]; + } + + CATLASS_DEVICE + GemmCoord GetCube2Shape() { + GDNFwdOOffsets& cube2Offsets = GetCube23Offsets(); + return GemmCoord{kHeadDim, vHeadDim, cube2Offsets.blockTokens}; + } + + CATLASS_DEVICE + GemmCoord GetCube3Shape() { + GDNFwdOOffsets& cube2Offsets = GetCube23Offsets(); + return GemmCoord{kHeadDim, vHeadDim, cube2Offsets.blockTokens}; + } + +}; + +struct BlockSchedulerGdnFwdOVec : public BlockSchedulerGdnFwdO { + CATLASS_DEVICE + BlockSchedulerGdnFwdOVec() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_offsets, GM_ADDR tiling) { + BlockSchedulerGdnFwdO::Init(cu_seqlens, chunk_offsets, tiling, AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(), AscendC::GetBlockNum()); + } + + CATLASS_DEVICE + bool NeedProcessVec1() { + return isRunning; + } + + CATLASS_DEVICE + bool NeedProcessVec2() { + if (unlikely(firstLoop)) { + firstLoop = false; + return false; + } + return true; + } + + CATLASS_DEVICE + GDNFwdOOffsets& GetVec1Offsets() { + return offsets[(currStage - 1) % PING_PONG_STAGES]; + } + + CATLASS_DEVICE + GDNFwdOOffsets& GetVec2Offsets() { + return offsets[(currStage - 2) % PING_PONG_STAGES]; + } + +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_SCHEDULER_GDN_FWD_O_HPP \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/gemm/kernel/gdn_fwd_o_kernel.hpp b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/gemm/kernel/gdn_fwd_o_kernel.hpp new file mode 100644 index 000000000..51237bfec --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_kernel/arch22/gemm/kernel/gdn_fwd_o_kernel.hpp @@ -0,0 +1,404 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#define CATLASS_ARCH 3510 + +#include "catlass/arch/arch.hpp" +#include "catlass/arch/cross_core_sync.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/catlass.hpp" +#include "catlass/debug.hpp" +#include "catlass/epilogue/block/block_epilogue.hpp" +#include "../../epilogue/block/block_epilogue_gdn_fwdo_qkmask.hpp" +#include "../../epilogue/block/block_epilogue_gdn_fwdo_output.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/gemm/block/block_swizzle.hpp" +#include "../block/block_scheduler_gdn_fwd_o.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" +#include "catlass/gemm_coord.hpp" +#include "tla/tensor.hpp" +#include "tla/layout.hpp" +#include "tla/tensor.hpp" + +using _0 = tla::Int<0>; +using _1 = tla::Int<1>; +using _2 = tla::Int<2>; +using _4 = tla::Int<4>; +using _8 = tla::Int<8>; +using _16 = tla::Int<16>; +using _32 = tla::Int<32>; +using _64 = tla::Int<64>; +using _128 = tla::Int<128>; +using _256 = tla::Int<256>; +using _512 = tla::Int<512>; +using _1024 = tla::Int<1024>; +using _2048 = tla::Int<2048>; +using _4096 = tla::Int<4096>; +using _8192 = tla::Int<8192>; +using _16384 = tla::Int<16384>; +using _32768 = tla::Int<32768>; +using _65536 = tla::Int<65536>; + +#else +#define CATLASS_ARCH 2201 + +#include "catlass/arch/arch.hpp" +#include "catlass/arch/cross_core_sync.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/catlass.hpp" +#include "catlass/debug.hpp" +#include "catlass/epilogue/block/block_epilogue.hpp" +#include "../../epilogue/block/block_epilogue_gdn_fwdo_qkmask.hpp" +#include "../../epilogue/block/block_epilogue_gdn_fwdo_output.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/gemm/block/block_swizzle.hpp" +#include "../block/block_scheduler_gdn_fwd_o.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" +#include "catlass/gemm_coord.hpp" +#include "tla/tensor.hpp" +#include "tla/layout.hpp" +#include "tla/tensor.hpp" +#endif + +#include "kernel_operator.h" +using namespace Catlass; +using namespace tla; + +// template <> +namespace Catlass::Gemm::Kernel { + +template< + typename INPUT_TYPE, + typename G_TYPE, + typename WORKSPACE_TYPE +> +class GDNFwdOKernel { +public: + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + using ArchTag = Arch::Ascend950; +#else + using ArchTag = Arch::AtlasA2; +#endif + using GDNFwdOOffsets = Catlass::Gemm::Block::GDNFwdOOffsets; + + using CubeScheduler = typename Catlass::Gemm::Block::BlockSchedulerGdnFwdOCube; + using VecScheduler = typename Catlass::Gemm::Block::BlockSchedulerGdnFwdOVec; + + using DispatchPolicyTla = Gemm::MmadPingpongTlaMulti; + using L1TileShapeTla = Shape<_128, _128, _128>; + using L0TileShapeTla = L1TileShapeTla; + using QType = Gemm::GemmType; + using KType = Gemm::GemmType; + using AttenType = Gemm::GemmType; + using AttenMaskedType = Gemm::GemmType; + using HType = Gemm::GemmType; + using OinterType = Gemm::GemmType; + using VNEWType = Gemm::GemmType; + + using GType = Gemm::GemmType; + using OType = Gemm::GemmType; + using MaskType = Gemm::GemmType; + + // cube 1 + using TileCopyQK = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadQK = Gemm::Block::BlockMmadTla; + + // cube 2 + using TileCopyQH = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadQH = Gemm::Block::BlockMmadTla; + + // cube 3 + using TileCopyAttenVNEW = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadAttenVNEW = Gemm::Block::BlockMmadTla; + + // vec 1 + using DispatchPolicyGDNFwdOQkmask = Epilogue::EpilogueAtlasGDNFwdOQkmask; + using EpilogueGDNFwdOQkmask = Epilogue::Block::BlockEpilogue; + + // vec 2 + using DispatchPolicyGDNFwdOOutput = Epilogue::EpilogueAtlasGDNFwdOOutput; + using EpilogueGDNFwdOOutput = Epilogue::Block::BlockEpilogue; + + using ElementQ = typename BlockMmadQK::ElementA; + using LayoutQ = Catlass::layout::RowMajor; + + using ElementK = typename BlockMmadQK::ElementB; + using LayoutK = Catlass::layout::ColumnMajor; + + using ElementAtten = typename BlockMmadQK::ElementC; + using LayoutAtten = Catlass::layout::RowMajor; + + using ElementAttenMasked = typename BlockMmadQH::ElementA; + using LayoutAttenMasked = Catlass::layout::RowMajor; + + using ElementH = typename BlockMmadQH::ElementB; + using LayoutH = Catlass::layout::RowMajor; + + using ElementOinter = typename BlockMmadQH::ElementC; + using LayoutOinter = Catlass::layout::RowMajor; + + + using ElementVNEW = typename BlockMmadAttenVNEW::ElementB; + using LayoutVNEW = Catlass::layout::RowMajor; + + + using ElementG = G_TYPE; + using ElementMask = bool; + + using L1TileShape = typename BlockMmadQK::L1TileShape; + + uint32_t shapeBatch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + float scale; + uint32_t numChunks; + uint32_t isVariedLen; + uint32_t tokenBatch; + uint32_t vWorkspaceOffset; + uint32_t hWorkspaceOffset; + uint32_t attnWorkspaceOffset; + uint32_t aftermaskWorkspaceOffset; + uint32_t maskWorkspaceOffset; + + AscendC::GlobalTensor gmQ; + AscendC::GlobalTensor gmK; + AscendC::GlobalTensor gmV; + AscendC::GlobalTensor gmH; + AscendC::GlobalTensor gmG; + AscendC::GlobalTensor gmO; + AscendC::GlobalTensor gmVWorkspace; + AscendC::GlobalTensor gmHWorkspace; + AscendC::GlobalTensor gmAttnWorkspace; + AscendC::GlobalTensor gmAftermaskWorkspace; + AscendC::GlobalTensor gmMask; + + CubeScheduler cubeBlockScheduler; + VecScheduler vecBlockScheduler; + + Arch::Resource resource; + + __aicore__ inline GDNFwdOKernel() {} + + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR h, GM_ADDR g, + GM_ADDR cu_seqlens, GM_ADDR chunk_offsets, GM_ADDR o, GM_ADDR tiling, GM_ADDR user) { + + __gm__ ChunkFwdOTilingData *__restrict gdnFwdOTilingData = reinterpret_cast<__gm__ ChunkFwdOTilingData *__restrict>(tiling); + + shapeBatch = gdnFwdOTilingData->shapeBatch; + seqlen = gdnFwdOTilingData->seqlen; + kNumHead = gdnFwdOTilingData->kNumHead; + vNumHead = gdnFwdOTilingData->vNumHead; + kHeadDim = gdnFwdOTilingData->kHeadDim; + vHeadDim = gdnFwdOTilingData->vHeadDim; + scale = gdnFwdOTilingData->scale; + chunkSize = gdnFwdOTilingData->chunkSize; + isVariedLen = gdnFwdOTilingData->isVariedLen; + tokenBatch = gdnFwdOTilingData->tokenBatch; + vWorkspaceOffset = gdnFwdOTilingData->vWorkspaceOffset; + hWorkspaceOffset = gdnFwdOTilingData->hWorkspaceOffset; + attnWorkspaceOffset = gdnFwdOTilingData->attnWorkspaceOffset; + aftermaskWorkspaceOffset = gdnFwdOTilingData->aftermaskWorkspaceOffset; + maskWorkspaceOffset = gdnFwdOTilingData->maskWorkspaceOffset; + + gmQ.SetGlobalBuffer((__gm__ ElementQ *)q); + gmK.SetGlobalBuffer((__gm__ ElementK *)k); + gmV.SetGlobalBuffer((__gm__ ElementVNEW *)v); + gmH.SetGlobalBuffer((__gm__ ElementH *)h); + gmG.SetGlobalBuffer((__gm__ ElementG *)g); + gmO.SetGlobalBuffer((__gm__ ElementVNEW *)o); + gmVWorkspace.SetGlobalBuffer((__gm__ ElementOinter *)(user + vWorkspaceOffset)); + gmHWorkspace.SetGlobalBuffer((__gm__ ElementOinter *)(user + hWorkspaceOffset)); + gmAttnWorkspace.SetGlobalBuffer((__gm__ ElementAtten *)(user + attnWorkspaceOffset)); + gmAftermaskWorkspace.SetGlobalBuffer((__gm__ ElementAttenMasked *)(user + aftermaskWorkspaceOffset)); + gmMask.SetGlobalBuffer((__gm__ ElementMask *)(user + maskWorkspaceOffset)); + + if ASCEND_IS_AIC { + cubeBlockScheduler.Init(cu_seqlens, chunk_offsets, tiling); + } + + if ASCEND_IS_AIV { + vecBlockScheduler.Init(cu_seqlens, chunk_offsets, tiling); + } + } + + __aicore__ inline void Process() { + if ASCEND_IS_AIC { + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = AscendC::GetBlockNum(); + + BlockMmadQK blockMmadQK(resource); + BlockMmadQH blockMmadQH(resource); + BlockMmadAttenVNEW blockMmadAttenVNEW(resource); + + auto qLayout = tla::MakeLayout(shapeBatch * kNumHead * seqlen, kHeadDim); + auto kLayout = tla::MakeLayout(kHeadDim, shapeBatch * kNumHead * seqlen); + auto hLayout = tla::MakeLayout(shapeBatch * vNumHead * seqlen * kHeadDim, vHeadDim); + auto ointerLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, vHeadDim); + auto vnewLayout = tla::MakeLayout(shapeBatch * vNumHead * seqlen, vHeadDim); + + bool needRun = false; + bool isFirstC3 = true; + + while (cubeBlockScheduler.isRunning) { + cubeBlockScheduler.InitTask(); + + if (cubeBlockScheduler.isRunning && coreIdx < coreNum) { + + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done); + + GDNFwdOOffsets& cube1Offsets = cubeBlockScheduler.GetCube1Offsets(); + int64_t cube1OffsetQ = cube1Offsets.qkOffset; + int64_t cube1OffsetK = cube1Offsets.qkOffset; + int64_t cube1OffsetAttn = cube1Offsets.attnWorkOffset; + auto attenLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, cube1Offsets.blockTokens); + auto tensorQ = tla::MakeTensor(gmQ[cube1OffsetQ], qLayout, Catlass::Arch::PositionGM{}); + auto tensorK = tla::MakeTensor(gmK[cube1OffsetK], kLayout, Catlass::Arch::PositionGM{}); + auto tensorAttn = tla::MakeTensor(gmAttnWorkspace[cube1OffsetAttn], attenLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape{cube1Offsets.blockTokens, cube1Offsets.blockTokens, kHeadDim}; + auto tensorBlockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + auto tensorBlockAttn = GetTile(tensorAttn, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.n())); + blockMmadQK.preSetFlags(); + blockMmadQK(tensorBlockQ, tensorBlockK, tensorBlockAttn, cube1Shape); + blockMmadQK.finalWaitFlags(); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube1Done); + + } + // AscendC::PipeBarrier(); + + if (needRun && coreIdx < coreNum) { + if(!cubeBlockScheduler.isRunning) Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done); + GDNFwdOOffsets& cube2Offsets = cubeBlockScheduler.GetCube23Offsets(); + int64_t cube2OffsetQ = cube2Offsets.qkOffset; + int64_t cube2OffsetH = cube2Offsets.hOffset; + int64_t cube2OffsetHWork = cube2Offsets.hvWorkOffset; + auto tensorQ = tla::MakeTensor(gmQ[cube2OffsetQ], qLayout, Catlass::Arch::PositionGM{}); + auto tensorH = tla::MakeTensor(gmH[cube2OffsetH], hLayout, Catlass::Arch::PositionGM{}); + auto tensorHWork = tla::MakeTensor(gmHWorkspace[cube2OffsetHWork], ointerLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube2Shape{cube2Offsets.blockTokens, vHeadDim, kHeadDim}; + auto tensorBlockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + auto tensorBlockHWork = GetTile(tensorHWork, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.m(), cube2Shape.n())); + blockMmadQH.preSetFlags(); + blockMmadQH(tensorBlockQ, tensorBlockH, tensorBlockHWork, cube2Shape); + blockMmadQH.finalWaitFlags(); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube2Done); + } + + if (needRun && coreIdx < coreNum) { + GDNFwdOOffsets& cube3Offsets = cubeBlockScheduler.GetCube23Offsets(); + if(isFirstC3) Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done); + int64_t cube3OffsetAttnMask = cube3Offsets.attnWorkOffset; + int64_t cube3OffsetV = cube3Offsets.ovOffset; + int64_t cube3OffsetVWork = cube3Offsets.hvWorkOffset; + auto attenLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, cube3Offsets.blockTokens); + auto tensorAttnMask = tla::MakeTensor(gmAftermaskWorkspace[cube3OffsetAttnMask], attenLayout, Catlass::Arch::PositionGM{}); + auto tensorV = tla::MakeTensor(gmV[cube3OffsetV], vnewLayout, Catlass::Arch::PositionGM{}); + auto tensorVWork = tla::MakeTensor(gmVWorkspace[cube3OffsetVWork], ointerLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube3Shape{cube3Offsets.blockTokens, vHeadDim, cube3Offsets.blockTokens}; + auto tensorBlockAttnMask = GetTile(tensorAttnMask, tla::MakeCoord(0, 0), tla::MakeShape(cube3Shape.m(), cube3Shape.k())); + auto tensorBlockV = GetTile(tensorV, tla::MakeCoord(0, 0), tla::MakeShape(cube3Shape.k(), cube3Shape.n())); + auto tensorBlockVWork = GetTile(tensorVWork, tla::MakeCoord(0, 0), tla::MakeShape(cube3Shape.m(), cube3Shape.n())); + blockMmadAttenVNEW.preSetFlags(); + blockMmadAttenVNEW(tensorBlockAttnMask, tensorBlockV, tensorBlockVWork, cube3Shape); + blockMmadAttenVNEW.finalWaitFlags(); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube3Done); + isFirstC3 = false; + } + needRun = true; + // AscendC::PipeBarrier(); + } + if (coreIdx < coreNum) { + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done); + + } + } + + if ASCEND_IS_AIV { + + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = AscendC::GetBlockNum(); + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + AscendC::LocalTensor maskUbTensor = resource.ubBuf.template GetBufferByByte(0); + AscendC::Duplicate(maskUbTensor, (float)0.0, 64*64); + AscendC::PipeBarrier(); + for(uint32_t i = 0; i < 64; ++ i) AscendC::Duplicate(maskUbTensor[i * 64], (float)1.0, i + 1); + AscendC::PipeBarrier(); + + bool needRun = false; + uint32_t pingpongFlag = 0; + + if (coreIdx < coreNum * subBlockNum) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec1Done); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec1Done); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done); + } + + while (vecBlockScheduler.isRunning) { + vecBlockScheduler.InitTask(); + + if (vecBlockScheduler.isRunning && coreIdx < coreNum * subBlockNum) { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube1Done); + GDNFwdOOffsets& vec1Offsets = vecBlockScheduler.GetVec1Offsets(); + int64_t vec1OffsetAttnMask = vec1Offsets.attnWorkOffset; + int64_t vec1OffsetG = vec1Offsets.gOffset; + int64_t vec1OffsetAttn = vec1Offsets.attnWorkOffset; + EpilogueGDNFwdOQkmask epilogueGDNFwdOQkmask(resource); + epilogueGDNFwdOQkmask( + gmAftermaskWorkspace[vec1OffsetAttnMask], + gmG[vec1OffsetG], gmAttnWorkspace[vec1OffsetAttn], gmMask, + chunkSize, vec1Offsets.blockTokens, kHeadDim, vHeadDim, pingpongFlag, vec1Offsets.batchIdx, vec1Offsets.headIdx, vec1Offsets.chunkIdx + ); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec1Done); + } + + // AscendC::PipeBarrier(); + + if (needRun && coreIdx < coreNum * subBlockNum) { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube2Done); + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube3Done); + GDNFwdOOffsets& vec2Offsets = vecBlockScheduler.GetVec2Offsets(); + int64_t vec2OffsetO = vec2Offsets.ovOffset; + int64_t vec2OffsetG = vec2Offsets.gOffset; + int64_t vec2OffsetVWork = vec2Offsets.hvWorkOffset; + int64_t vec2OffsetHWork = vec2Offsets.hvWorkOffset; + EpilogueGDNFwdOOutput epilogueGDNFwdOOutput(resource); + epilogueGDNFwdOOutput( + gmO[vec2OffsetO], + gmG[vec2OffsetG], gmVWorkspace[vec2OffsetVWork], gmHWorkspace[vec2OffsetHWork], + scale, vec2Offsets.blockTokens, kHeadDim, vHeadDim, pingpongFlag, vec2Offsets.batchIdx, vec2Offsets.headIdx, vec2Offsets.chunkIdx + ); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done); + } + + // AscendC::PipeBarrier(); + + needRun = true; + } + } + } + +}; + +} diff --git a/csrc/ascend/moe/chunk_fwd_o/op_kernel/chunk_fwd_o.cpp b/csrc/ascend/moe/chunk_fwd_o/op_kernel/chunk_fwd_o.cpp new file mode 100644 index 000000000..bc2130ffb --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/op_kernel/chunk_fwd_o.cpp @@ -0,0 +1,69 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +/*! + * \file chunk_fwd_o.cpp + * \brief + */ + +// #include "chunk_fwd_o.h" +#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 200) +#include "arch20/compat_310p.h" +#include "arch20/gemm/kernel/gdn_fwd_o_kernel.hpp" +#else +#include "arch22/gemm/kernel/gdn_fwd_o_kernel.hpp" +#endif +#include "lib/matmul_intf.h" + +using namespace Catlass; + +extern "C" __global__ __aicore__ void chunk_fwd_o(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR h, + GM_ADDR g, GM_ADDR cu_seqlens, GM_ADDR chunk_offsets, + GM_ADDR o, GM_ADDR workspace, GM_ADDR tiling) +{ +#ifdef CATLASS_UNIFIED_CORE + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC); +#else + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); +#endif + + GM_ADDR user = AscendC::GetUserWorkspace(workspace); + + __gm__ ChunkFwdOTilingData *__restrict gdnFwdOTilingData = reinterpret_cast<__gm__ ChunkFwdOTilingData *__restrict>(tiling); + using workspaceType = float; + // dtype: 0 - fp16, 1 - bf16, 2 - fp32 +#ifndef CATLASS_UNIFIED_CORE + if (gdnFwdOTilingData->dataType == 1) { + if (gdnFwdOTilingData->gDataType == 2) { + using GDNFwdOKernel = Catlass::Gemm::Kernel::GDNFwdOKernel; + GDNFwdOKernel gdnFwdO; + gdnFwdO.Init(q, k, v, h, g, cu_seqlens, chunk_offsets, o, tiling, user); + gdnFwdO.Process(); + } else { + using GDNFwdOKernel = Catlass::Gemm::Kernel::GDNFwdOKernel; + GDNFwdOKernel gdnFwdO; + gdnFwdO.Init(q, k, v, h, g, cu_seqlens, chunk_offsets, o, tiling, user); + gdnFwdO.Process(); + } + } else +#endif + { + if (gdnFwdOTilingData->gDataType == 2) { + using GDNFwdOKernel = Catlass::Gemm::Kernel::GDNFwdOKernel; + GDNFwdOKernel gdnFwdO; + gdnFwdO.Init(q, k, v, h, g, cu_seqlens, chunk_offsets, o, tiling, user); + gdnFwdO.Process(); + } else { + using GDNFwdOKernel = Catlass::Gemm::Kernel::GDNFwdOKernel; + GDNFwdOKernel gdnFwdO; + gdnFwdO.Init(q, k, v, h, g, cu_seqlens, chunk_offsets, o, tiling, user); + gdnFwdO.Process(); + } + } +} diff --git a/csrc/ascend/moe/chunk_fwd_o/tiling_base/data_copy_transpose_tiling.h b/csrc/ascend/moe/chunk_fwd_o/tiling_base/data_copy_transpose_tiling.h new file mode 100644 index 000000000..f7e7fce1c --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/tiling_base/data_copy_transpose_tiling.h @@ -0,0 +1,51 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file data_copy_transpose_tiling.h + * \brief + */ + +#pragma once + +#include +#include +#include "data_copy_transpose_tiling_def.h" + +namespace optiling { + +inline void GetDataCopyTransposeTiling(const ge::Shape &dstShape, const ge::Shape &srcShape, const uint32_t typeSize, + optiling::CopyTransposeTiling &tiling) +{ + constexpr int64_t B_INDEX = 0; + constexpr int64_t N_INDEX = 1; + constexpr int64_t S_INDEX = 2; + constexpr int64_t H_INDEX = 3; + std::vector dstShapeInfo = dstShape.GetDims(); + std::vector srcShapeInfo = srcShape.GetDims(); + + tiling.set_dstShapeB(dstShapeInfo[B_INDEX]); + tiling.set_dstShapeN(dstShapeInfo[N_INDEX]); + tiling.set_dstShapeS(dstShapeInfo[S_INDEX]); + tiling.set_dstShapeH(dstShapeInfo[H_INDEX]); + tiling.set_dstShapeHN(tiling.get_dstShapeH() / tiling.get_dstShapeN()); + + tiling.set_srcShapeB(srcShapeInfo[B_INDEX]); + tiling.set_srcShapeN(srcShapeInfo[N_INDEX]); + tiling.set_srcShapeS(srcShapeInfo[S_INDEX]); + tiling.set_srcShapeHN(srcShapeInfo[H_INDEX]); + tiling.set_originalShapeNLen(tiling.get_srcShapeHN() * typeSize); + tiling.set_shapeSHValue(tiling.get_dstShapeS() * tiling.get_dstShapeH()); + tiling.set_shapeNsValue(tiling.get_dstShapeN() * tiling.get_dstShapeS()); + tiling.set_shapeNsnValue(tiling.get_dstShapeN() * tiling.get_srcShapeS() * tiling.get_srcShapeN()); + tiling.set_shapeBHValue(tiling.get_dstShapeB() * tiling.get_dstShapeH()); +} + +} // namespace optiling diff --git a/csrc/ascend/moe/chunk_fwd_o/tiling_base/data_copy_transpose_tiling_def.h b/csrc/ascend/moe/chunk_fwd_o/tiling_base/data_copy_transpose_tiling_def.h new file mode 100644 index 000000000..18552c36b --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/tiling_base/data_copy_transpose_tiling_def.h @@ -0,0 +1,43 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file data_copy_transpose_tiling_def.h + * \brief + */ + +#pragma once + +#include +#include + +namespace optiling { + +BEGIN_TILING_DATA_DEF(CopyTransposeTiling) +TILING_DATA_FIELD_DEF(uint32_t, dstShapeB); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeN); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeS); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeHN); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeH); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeB); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeN); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeS); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeHN); +TILING_DATA_FIELD_DEF(uint32_t, originalShapeNLen); +TILING_DATA_FIELD_DEF(uint32_t, shapeSHValue); +TILING_DATA_FIELD_DEF(uint32_t, shapeNsValue); +TILING_DATA_FIELD_DEF(uint32_t, shapeNsnValue); +TILING_DATA_FIELD_DEF(uint32_t, invalidParamCopyTransposeTiling); +TILING_DATA_FIELD_DEF(uint32_t, shapeBHValue); +TILING_DATA_FIELD_DEF(uint32_t, paramsAlign); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(CopyTransposeTilingOp, CopyTransposeTiling) + +} // namespace optiling diff --git a/csrc/ascend/moe/chunk_fwd_o/tiling_base/error_log.h b/csrc/ascend/moe/chunk_fwd_o/tiling_base/error_log.h new file mode 100644 index 000000000..119189974 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/tiling_base/error_log.h @@ -0,0 +1,63 @@ +#ifndef OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_ +#define OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_ + +#include +#include +#include "toolchain/slog.h" + +#define OP_LOGI(opname, ...) +#define OP_LOGW(opname, ...) \ + do { \ + (void)(opname); \ + std::printf("[WARN] "); \ + std::printf(__VA_ARGS__); \ + std::printf("\n"); \ + } while (0) + +#define OP_LOGE_WITHOUT_REPORT(opname, ...) \ + do { \ + (void)(opname); \ + std::printf("[ERRORx] "); \ + std::printf(__VA_ARGS__); \ + std::printf("\n"); \ + } while (0) + +#define OP_LOGE(opname, ...) \ + do { \ + (void)(opname); \ + std::printf("[ERROR] "); \ + std::printf(__VA_ARGS__); \ + std::printf("\n"); \ + } while (0) + +#define OP_LOGD(opname, ...) + +namespace optiling { + +#define VECTOR_INNER_ERR_REPORT_TILIING(op_name, err_msg, ...) \ + do { \ + OP_LOGE_WITHOUT_REPORT(op_name, err_msg, ##__VA_ARGS__); \ + } while (0) + +// Modify OP_TILING_CHECK macro to ensure proper handling of expressions +#define OP_CHECK_IF(cond, log_func, expr) \ + do { \ + if (cond) { \ + log_func; \ + expr; \ + } \ + } while (0) + + + +#define OP_CHECK_NULL_WITH_CONTEXT(context, ptr) \ + do { \ + if ((ptr) == nullptr) { \ + OP_LOGE(context->GetNodeType(), "%s is null", #ptr); \ + return ge::GRAPH_FAILED; \ + } \ + } while (0) + +} // namespace optiling + +#endif // OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_ diff --git a/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_base.h b/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_base.h new file mode 100644 index 000000000..f0bbbdccc --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_base.h @@ -0,0 +1,256 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_base.h + * \brief + */ + +#pragma once + +#include +#include +#include +#include "tiling/platform/platform_ascendc.h" +#include "error_log.h" + +#ifdef ASCENDC_OP_TEST +#define ASCENDC_EXTERN_C extern "C" +#else +#define ASCENDC_EXTERN_C +#endif + +namespace Ops { +namespace Transformer { +namespace OpTiling { + +struct AiCoreParams { + uint64_t ubSize = 0; + uint64_t blockDim = 0; + uint64_t aicNum = 0; + uint64_t l1Size = 0; + uint64_t l0aSize = 0; + uint64_t l0bSize = 0; + uint64_t l0cSize = 0; +}; + +struct CompileInfoCommon { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; + int32_t socVersion; + uint32_t rsvd; +}; + +struct FlashAttentionScoreGradCompileInfo { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; + platform_ascendc::SocVersion socVersion; +}; + +struct FACompileInfoCommon { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; + int32_t socVersion; + uint32_t rsvd; +}; + +class TilingBaseClass { +public: + explicit TilingBaseClass(gert::TilingContext* context) : context_(context) + {} + + virtual ~TilingBaseClass() = default; + + // Tiling execution framework + // 1. GRAPH_SUCCESS: Success, and no need to continue executing subsequent Tiling class implementations + // 2. GRAPH_FAILED: Failure, abort the entire Tiling process + // 3. GRAPH_PARAM_INVALID: This class does not support, need to continue executing other Tiling class implementations + ge::graphStatus DoTiling() + { + auto ret = GetShapeAttrsInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetPlatformInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + if (!IsCapable()) { + return ge::GRAPH_PARAM_INVALID; + } + ret = DoOpTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = DoLibApiTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetWorkspaceSize(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = PostTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + context_->SetTilingKey(GetTilingKey()); + DumpTilingInfo(); + return ge::GRAPH_SUCCESS; + } + + // Update context + virtual void Reset(gert::TilingContext* context) + { + context_ = context; + } + +protected: + virtual bool IsCapable() = 0; + // 1. Get platform information such as CoreNum, UB/L1/L0C resource sizes + virtual ge::graphStatus GetPlatformInfo() = 0; + // 2. Get INPUT/OUTPUT/ATTR information + virtual ge::graphStatus GetShapeAttrsInfo() = 0; + // 3. Calculate data splitting TilingData + virtual ge::graphStatus DoOpTiling() = 0; + // 4. Calculate high-level API TilingData + virtual ge::graphStatus DoLibApiTiling() = 0; + // 5. Calculate TilingKey + [[nodiscard]] virtual uint64_t GetTilingKey() const = 0; + // 6. Calculate Workspace size + virtual ge::graphStatus GetWorkspaceSize() = 0; + // 7. Save Tiling data + virtual ge::graphStatus PostTiling() = 0; + // 8. Dump Tiling data + virtual void DumpTilingInfo() + { + int32_t enable = CheckLogLevel(static_cast(OP), DLOG_DEBUG); + if (enable != 1) { + return; + } + auto buf = (uint32_t*)context_->GetRawTilingData()->GetData(); + auto bufLen = context_->GetRawTilingData()->GetDataSize(); + std::ostringstream oss; + oss << "Start to dump tiling info. tilingkey:" << context_->GetTilingKey() << ", tiling data size:" << bufLen + << ", content:"; + for (size_t i = 0; i < bufLen / sizeof(uint32_t); i++) { + oss << *(buf + i) << ","; + if (oss.str().length() > 640) { // Split according to 640 to avoid truncation + OP_LOGD(context_, "%s", oss.str().c_str()); + oss.str(""); + } + } + OP_LOGD(context_, "%s", oss.str().c_str()); + } + + static uint32_t CalcTschBlockDim(uint32_t sliceNum, uint32_t aicCoreNum, uint32_t aivCoreNum) + { + uint32_t ration; + if (aicCoreNum == 0 || aivCoreNum == 0 || aicCoreNum > aivCoreNum) { + return sliceNum; + } + ration = aivCoreNum / aicCoreNum; + return (sliceNum + (ration - 1)) / ration; + } + + template + [[nodiscard]] std::string GetShapeDebugStr(const T& shape) const + { + std::ostringstream oss; + oss << "["; + if (shape.GetDimNum() > 0) { + for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) { + oss << shape.GetDim(i) << ", "; + } + oss << shape.GetDim(shape.GetDimNum() - 1); + } + oss << "]"; + return oss.str(); + } + + [[nodiscard]] std::string GetTensorDebugStr( + const gert::StorageShape* shape, const gert::CompileTimeTensorDesc* tensor) + { + if (shape == nullptr || tensor == nullptr) { + return "nil "; + } + std::ostringstream oss; + oss << "(dtype: " << ge::TypeUtils::DataTypeToSerialString(tensor->GetDataType()) << "),"; + oss << "(shape:" << GetShapeDebugStr(shape->GetStorageShape()) << "),"; + oss << "(ori_shape:" << GetShapeDebugStr(shape->GetOriginShape()) << "),"; + oss << "(format: " + << ge::TypeUtils::FormatToSerialString( + static_cast(ge::GetPrimaryFormat(tensor->GetStorageFormat()))) + << "),"; + oss << "(ori_format: " << ge::TypeUtils::FormatToSerialString(tensor->GetOriginFormat()) << ") "; + return oss.str(); + } + + [[nodiscard]] std::string GetTilingContextDebugStr() + { + std::ostringstream oss; + for (size_t i = 0; i < context_->GetComputeNodeInfo()->GetInputsNum(); ++i) { + oss << "input" << i << ": "; + oss << GetTensorDebugStr(context_->GetInputShape(i), context_->GetInputDesc(i)); + } + + for (size_t i = 0; i < context_->GetComputeNodeInfo()->GetOutputsNum(); ++i) { + oss << "output" << i << ": "; + oss << GetTensorDebugStr(context_->GetOutputShape(i), context_->GetOutputDesc(i)); + } + return oss.str(); + } + + [[nodiscard]] std::string GetTilingDataDebugStr() const + { + auto rawTilingData = context_->GetRawTilingData(); + auto rawTilingDataSize = rawTilingData->GetDataSize(); + auto data = reinterpret_cast(rawTilingData->GetData()); + size_t len = rawTilingDataSize / sizeof(int32_t); + std::ostringstream oss; + for (size_t i = 0; i < len; i++) { + oss << data[i] << ", "; + } + return oss.str(); + } + +protected: + gert::TilingContext* context_ = nullptr; + std::unique_ptr ascendcPlatform_{nullptr}; + uint32_t blockDim_{0}; + uint64_t workspaceSize_{0}; + uint64_t tilingKey_{0}; + AiCoreParams aicoreParams_; +}; + +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_key.h b/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_key.h new file mode 100644 index 000000000..607f965b0 --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_key.h @@ -0,0 +1,63 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_key.h + * \brief + */ + +#pragma once + +#include + +namespace Ops { +namespace Transformer { +namespace OpTiling { +constexpr uint64_t RecursiveSum() +{ + return 0; +} + +constexpr uint64_t kBase = 10; // Base-10 carry base +template constexpr uint64_t RecursiveSum(T templateId, Args... templateIds) +{ + return static_cast(templateId) + kBase * RecursiveSum(templateIds...); +} + +// TilingKey generation rules: +// FlashAttentionScore/FlashAttentionScoreGrad assembles tiling key using decimal digits, containing the following key parameters from low to high: Ub0, Ub1, +// Block, DataType, Format, Sparse. Specialized template Ub0, Ub1: +// Represents the axis for UB intra-core splitting, using AxisEnum. Since we allow at most two axes to be split, UB0 and UB1 exist. If there is no UB intra-core splitting, +// fill with AXIS_NONE. UB0 and UB1 each occupy one decimal digit; +// Block: Represents the axis used by UB for multi-core splitting, using AxisEnum, occupies one decimal digit; +// DataType: Represents the input/output data types supported by the current tiling key, using SupportedDtype enum, occupies one decimal digit +// Format: Represents the Format supported by the current tiling key, using InputLayout enum, occupies one decimal digit +// Sparse: Represents whether the current tiling key supports Sparse, using SparseCapability enum, occupies one decimal digit +// For other specialized scenarios, define your own bit fields and values +// usage: get tilingKey from inputted types +// uint64_t tilingKey = GET_FLASHATTENTION_TILINGKEY(AxisEnum::AXIS_S1, AxisEnum::AXIS_S2, AxisEnum::AXIS_N2, +// SupportedDtype::FLOAT32, InputLayout::BSH, SparseCapability::SUPPORT_ALL) + +constexpr uint64_t TILINGKEYOFFSET = uint64_t(10000000000000000000UL); // 10^19 +template constexpr uint64_t GET_TILINGKEY(Args... templateIds) +{ + return TILINGKEYOFFSET + RecursiveSum(templateIds...); +} + +// usage: get tilingKey from inputted types +// uint64_t tilingKey = TILINGKEY(S2, S1, N2, FLOAT32, BSND, ALL) + +#define TILINGKEY(ub2, ub1, block, dtype, layout, sparse) \ + (GET_TILINGKEY(AxisEnum::ub2, AxisEnum::ub1, AxisEnum::block, DtypeEnum::dtype, LayoutEnum::layout, \ + SparseEnum::sparse)) + +} // namespace Optiling +} // namespace Transformer +} // namespace Ops diff --git a/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_templates_registry.h b/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_templates_registry.h new file mode 100644 index 000000000..cbf4785ab --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_templates_registry.h @@ -0,0 +1,351 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_templates_registry.h + * \brief + */ + +#pragma once + +#include +#include +#include +#include "exe_graph/runtime/tiling_context.h" +#include "tiling_base.h" +#include "error_log.h" + +namespace Ops { +namespace Transformer { +namespace OpTiling { + +template +std::unique_ptr TILING_CLASS(gert::TilingContext* context) +{ + return std::unique_ptr(new (std::nothrow) T(context)); +} + +using TilingClassCase = std::unique_ptr (*)(gert::TilingContext*); + +class TilingCases { +public: + explicit TilingCases(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + void AddTiling(int32_t priority) + { + OP_CHECK_IF( + cases_.find(priority) != cases_.end(), OP_LOGE(op_type_, "There are duplicate registrations."), return); + cases_[priority] = TILING_CLASS; + OP_CHECK_IF( + cases_[priority] == nullptr, + OP_LOGE(op_type_, "Register op tiling func failed, please check the class name."), return); + } + + const std::map& GetTilingCases() + { + return cases_; + } + +private: + std::map cases_; + const std::string op_type_; +}; + +// --------------------------------Interfacce with soc version -------------------------------- +class TilingRegistryNew { +public: + TilingRegistryNew() = default; + +#ifdef ASCENDC_OP_TEST + static TilingRegistryNew& GetInstance(); +#else + static TilingRegistryNew& GetInstance() + { + static TilingRegistryNew registry_impl_; + return registry_impl_; + } +#endif + + std::shared_ptr RegisterOp(const std::string& op_type, int32_t soc_version) + { + auto soc_iter = registry_map_.find(soc_version); + if (soc_iter == registry_map_.end()) { + std::map> op_type_map; + op_type_map[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + registry_map_[soc_version] = op_type_map; + } else { + if (soc_iter->second.find(op_type) == soc_iter->second.end()) { + soc_iter->second[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + } + } + + OP_CHECK_IF( + registry_map_[soc_version][op_type] == nullptr, + OP_LOGE(op_type, "Register tiling func failed, please check the class name."), return nullptr); + return registry_map_[soc_version][op_type]; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context) + { + int32_t soc_version = (int32_t)platform_ascendc::SocVersion::RESERVED_VERSION; + const char* op_type = context->GetNodeType(); + fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + auto compileInfoPtr = static_cast(context->GetCompileInfo()); + OP_CHECK_IF( + compileInfoPtr == nullptr, OP_LOGE(op_type, "compileInfoPtr is null."), return ge::GRAPH_FAILED); + soc_version = compileInfoPtr->socVersion; + OP_LOGD(context, "soc version in compileInfo is %d", soc_version); + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + soc_version = static_cast(ascendcPlatform.GetSocVersion()); + OP_LOGD(context, "soc version is %d", soc_version); + if (soc_version == (int32_t)platform_ascendc::SocVersion::RESERVED_VERSION) { + OP_LOGE(op_type, "Do op tiling failed, cannot find soc version."); + return ge::GRAPH_FAILED; + } + } + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type, soc_version); + for (auto it = tilingTemplateRegistryMap.begin(); it != tilingTemplateRegistryMap.end(); ++it) { + auto tilingTemplate = it->second(context); + if (tilingTemplate != nullptr) { + ge::graphStatus status = tilingTemplate->DoTiling(); + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do general op tiling success priority=%d", it->first); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", it->first); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context, const std::vector& priorities) + { + int32_t soc_version; + const char* op_type = context->GetNodeType(); + auto platformInfoPtr = context->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + auto compileInfoPtr = reinterpret_cast(context->GetCompileInfo()); + OP_CHECK_IF( + compileInfoPtr == nullptr, OP_LOGE(op_type, "compileInfoPtr is null."), return ge::GRAPH_FAILED); + soc_version = compileInfoPtr->socVersion; + OP_LOGD(context, "soc version in compileInfo is %d", soc_version); + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + soc_version = static_cast(ascendcPlatform.GetSocVersion()); + OP_LOGD(context, "soc version is %d", soc_version); + } + + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type, soc_version); + for (auto priority_id : priorities) { + auto tilingCaseIter = tilingTemplateRegistryMap.find(priority_id); + if (tilingCaseIter != tilingTemplateRegistryMap.end()) { + auto templateFunc = tilingCaseIter->second(context); + if (templateFunc != nullptr) { + ge::graphStatus status = templateFunc->DoTiling(); + if (status == ge::GRAPH_SUCCESS) { + OP_LOGD(context, "Do general op tiling success priority=%d", priority_id); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", priority_id); + } + } + } + return ge::GRAPH_FAILED; + } + + const std::map& GetTilingTemplates(const std::string& op_type, int32_t soc_version) + { + auto soc_iter = registry_map_.find(soc_version); + OP_CHECK_IF( + soc_iter == registry_map_.end(), + OP_LOGE(op_type, "Get op tiling func failed, please check the soc version %d", soc_version), + return empty_tiling_case_); + auto op_iter = soc_iter->second.find(op_type); + OP_CHECK_IF( + op_iter == soc_iter->second.end(), OP_LOGE(op_type, "Get op tiling func failed, please check the op name."), + return empty_tiling_case_); + return op_iter->second->GetTilingCases(); + } + +private: + std::map>> registry_map_; // key is socversion + const std::map empty_tiling_case_{}; +}; + +class RegisterNew { +public: + explicit RegisterNew(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + RegisterNew& tiling(int32_t priority, int32_t soc_version) + { + auto tilingCases = TilingRegistryNew::GetInstance().RegisterOp(op_type_, soc_version); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), return *this); + tilingCases->AddTiling(priority); + return *this; + } + + template + RegisterNew& tiling(int32_t priority, const std::vector& soc_versions) + { + for (int32_t soc_version : soc_versions) { + auto tilingCases = TilingRegistryNew::GetInstance().RegisterOp(op_type_, soc_version); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), + return *this); + tilingCases->AddTiling(priority); + } + return *this; + } + +private: + const std::string op_type_; +}; + +// --------------------------------Interfacce without soc version -------------------------------- +class TilingRegistry { +public: + TilingRegistry() = default; + +#ifdef ASCENDC_OP_TEST + static TilingRegistry& GetInstance(); +#else + static TilingRegistry& GetInstance() + { + static TilingRegistry registry_impl_; + return registry_impl_; + } +#endif + + std::shared_ptr RegisterOp(const std::string& op_type) + { + if (registry_map_.find(op_type) == registry_map_.end()) { + registry_map_[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + } + OP_CHECK_IF( + registry_map_[op_type] == nullptr, + OP_LOGE(op_type, "Register tiling func failed, please check the class name."), return nullptr); + return registry_map_[op_type]; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context) + { + const char* op_type = context->GetNodeType(); + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type); + for (auto it = tilingTemplateRegistryMap.begin(); it != tilingTemplateRegistryMap.end(); ++it) { + auto tilingTemplate = it->second(context); + if (tilingTemplate != nullptr) { + ge::graphStatus status = tilingTemplate->DoTiling(); + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do general op tiling success priority=%d", it->first); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", it->first); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context, const std::vector& priorities) + { + const char* op_type = context->GetNodeType(); + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type); + for (auto priorityId : priorities) { + auto templateFunc = tilingTemplateRegistryMap[priorityId](context); + if (templateFunc != nullptr) { + ge::graphStatus status = templateFunc->DoTiling(); + if (status == ge::GRAPH_SUCCESS) { + OP_LOGD(context, "Do general op tiling success priority=%d", priorityId); + return status; + } + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do op tiling failed"); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", priorityId); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + const std::map& GetTilingTemplates(const std::string& op_type) + { + OP_CHECK_IF( + registry_map_.find(op_type) == registry_map_.end(), + OP_LOGE(op_type, "Get op tiling func failed, please check the op name."), return empty_tiling_case_); + return registry_map_[op_type]->GetTilingCases(); + } + +private: + std::map> registry_map_; + const std::map empty_tiling_case_; +}; + +class Register { +public: + explicit Register(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + Register& tiling(int32_t priority) + { + auto tilingCases = TilingRegistry::GetInstance().RegisterOp(op_type_); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), return *this); + tilingCases->AddTiling(priority); + return *this; + } + +private: + const std::string op_type_; +}; +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops + +// op_type: operator name, class_name: registered tiling class, soc_version: chip version number +// priority: priority of tiling class, smaller value means higher priority, i.e., this tiling class will be selected first +#define REGISTER_TILING_TEMPLATE_WITH_SOCVERSION(op_type, class_name, soc_versions, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::RegisterNew VAR_UNUSED##op_type##class_name##priority_register = \ + Ops::Transformer::OpTiling::RegisterNew(#op_type).tiling(priority, soc_versions) + +// op_type: operator name, class_name: registered tiling class +// priority: priority of tiling class, smaller value means higher priority, i.e., higher probability of being selected +#define REGISTER_TILING_TEMPLATE(op_type, class_name, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::Register VAR_UNUSED##op_type_##class_name##priority_register = \ + Ops::Transformer::OpTiling::Register(op_type).tiling(priority) + +// op_type: operator name, class_name: registered tiling class +// soc_version: SOC version, used to distinguish different SOCs +// priority: priority of tiling class, smaller value means higher priority, i.e., this tiling class will be selected first +#define REGISTER_TILING_TEMPLATE_NEW(op_type, class_name, soc_version, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::RegisterNew VAR_UNUSED##op_type##class_name##priority_register = \ + Ops::Transformer::OpTiling::RegisterNew(#op_type).tiling(priority, soc_version) + +// op_type: operator name, class_name: registered tiling class +// priority: priority of tiling class, smaller value means higher priority, i.e., higher probability of being selected +// Replaces REGISTER_TILING_TEMPLATE, if op_type is a string constant, remove the quotes +#define REGISTER_OPS_TILING_TEMPLATE(op_type, class_name, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::Register \ + __attribute__((unused)) tiling_##op_type##_##class_name##_##priority##_register = \ + Ops::Transformer::OpTiling::Register(#op_type).tiling(priority) diff --git a/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_type.h b/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_type.h new file mode 100644 index 000000000..7c781d19f --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_type.h @@ -0,0 +1,139 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_type.h + * \brief + */ + +#pragma once + +#include + +namespace optiling { + +enum class AxisEnum { + B = 0, + N2 = 1, + G = 2, + S1 = 3, + S2 = 4, + D = 5, + NONE = 9, +}; + +enum class DtypeEnum { + FLOAT16 = 0, + FLOAT32 = 1, + BFLOAT16 = 2, + FLOAT16_PRECISION = 3, +}; + +enum class PerformanceOrientedEnum { + BIG_BUFFER = 1, + BIG_DOUBLE_BUFFER = 2, +}; + +enum class MatmulConfig { + NULL_CONFIG = 0, + NORMAL_CONFIG = 1, + MDL_CONFIG = 2 +}; + +enum class PseConfig { + NO_PSE = 0, + EXIST_PSE = 1 +}; + +enum class AttenMaskConfig { + NO_ATTEN_MASK = 0, + EXIST_ATTEN_MASK = 1 +}; + +enum class DropOutConfig { + NO_DROP_OUT = 0, + EXIST_DROP_OUT = 1 +}; + +enum class CubeFormatEnum { + ND = 0, + NZ = 1 +}; +enum class LayoutEnum { + BSND = 0, + SBND = 1, + BNSD = 2, + TND = 3, + NTD_TND = 4 +}; + +enum class CubeInputSourceEnum { + GM = 0, + L1 = 1 +}; + +enum class OptionEnum { + DISABLE = 0, + ENABLE = 1 +}; + +enum class SparseEnum { + ALL = 0, + NONE = 1, + ANY = 2, + CAUSAL = 3, + BAND = 4, + PREFIX = 5, + BAND_COMPRESS = 6, + RIGHT_DOWN_CAUSAL = 7, + RIGHT_DOWN_CAUSAL_BAND = 8, + BAND_LEFT_UP_CAUSAL = 9 +}; + +constexpr uint64_t RecursiveSum() +{ + return 0; +} + +constexpr int64_t base10Multiplier = 10; + +template constexpr uint64_t RecursiveSum(T templateId, Args... templateIds) +{ + return static_cast(templateId) + base10Multiplier * RecursiveSum(templateIds...); +} + +// TilingKey generation rules: +// FlashAttentionScore/FlashAttentionScoreGrad assembles tiling key using decimal digits, containing the following key parameters from low to high: Ub0, Ub1, +// Block, DataType, Format, Sparse. Specialized template Ub0, Ub1: +// Represents the axis for UB intra-core splitting, using AxisEnum. Since we allow at most two axes to be split, UB0 and UB1 exist. If there is no UB intra-core splitting, +// fill with AXIS_NONE. UB0 and UB1 each occupy one decimal digit; +// Block: Represents the axis used by UB for multi-core splitting, using AxisEnum, occupies one decimal digit; +// DataType: Represents the input/output data types supported by the current tiling key, using SupportedDtype enum, occupies one decimal digit +// Format: Represents the Format supported by the current tiling key, using InputLayout enum, occupies one decimal digit +// Sparse: Represents whether the current tiling key supports Sparse, using SparseCapability enum, occupies one decimal digit +// For other specialized scenarios, define your own bit fields and values +// usage: get tilingKey from inputted types +// uint64_t tilingKey = GET_FLASHATTENTION_TILINGKEY(AxisEnum::AXIS_S1, AxisEnum::AXIS_S2, AxisEnum::AXIS_N2, +// SupportedDtype::FLOAT32, InputLayout::BSH, SparseCapability::SUPPORT_ALL) + +constexpr uint64_t TILINGKEYOFFSET = uint64_t(10000000000000000000UL); // 10^19 +template constexpr uint64_t GET_TILINGKEY(Args... templateIds) +{ + return TILINGKEYOFFSET + RecursiveSum(templateIds...); +} + +// usage: get tilingKey from inputted types +// uint64_t tilingKey = TILINGKEY(S2, S1, N2, FLOAT32, BSND, ALL) + +#define TILINGKEY(ub2, ub1, block, dtype, layout, sparse) \ + (GET_TILINGKEY(AxisEnum::ub2, AxisEnum::ub1, AxisEnum::block, DtypeEnum::dtype, LayoutEnum::layout, \ + SparseEnum::sparse)) + +} // namespace optiling diff --git a/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_util.h b/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_util.h new file mode 100644 index 000000000..fb6ffa2df --- /dev/null +++ b/csrc/ascend/moe/chunk_fwd_o/tiling_base/tiling_util.h @@ -0,0 +1,30 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_util.h + * \brief + */ + +#pragma once + +#include "register/op_impl_registry.h" + +namespace Ops { +namespace Transformer { +namespace OpTiling { +bool IsRegbaseSocVersion(const gert::TilingParseContext* context); + +bool IsRegbaseSocVersion(const gert::TilingContext* context); + +const gert::Shape& EnsureNotScalar(const gert::Shape& inShape); +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/CMakeLists.txt b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/CMakeLists.txt new file mode 100644 index 000000000..86b308249 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/CMakeLists.txt b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/CMakeLists.txt new file mode 100644 index 000000000..f58e9b7d4 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/CMakeLists.txt @@ -0,0 +1,32 @@ +set(CURRENT_CMAKE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) +set(CATLASS_INCLUDE_DIR "${CMAKE_SOURCE_DIR}/third_party/catlass/include") +get_filename_component(CATLASS_INCLUDE_DIR_ABS ${CATLASS_INCLUDE_DIR} ABSOLUTE) + +set(COMMON_KERNEL_UTILS_DIR "${CMAKE_SOURCE_DIR}/moe/common") +get_filename_component(COMMON_KERNEL_UTILS_DIR_ABS ${COMMON_KERNEL_UTILS_DIR} ABSOLUTE) + +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnExc PRIVATE + chunk_gated_delta_rule_fwd_h_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME ChunkGatedDeltaRuleFwdH + OPTIONS + --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -I${CATLASS_INCLUDE_DIR_ABS} + -I${COMMON_KERNEL_UTILS_DIR_ABS} +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE chunk_gated_delta_rule_fwd_h ACLNNTYPE aclnn_exclude) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ${CATLASS_INCLUDE_DIR_ABS} + ) +endif() diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/chunk_gated_delta_rule_fwd_h_def.cpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/chunk_gated_delta_rule_fwd_h_def.cpp new file mode 100644 index 000000000..4d16edc57 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/chunk_gated_delta_rule_fwd_h_def.cpp @@ -0,0 +1,116 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +/*! + * \file chunk_gated_delta_rule_fwd_h_def.cpp + * \brief + */ + +#include "register/op_def_registry.h" + +namespace ops { + +class ChunkGatedDeltaRuleFwdH : public OpDef { +public: + explicit ChunkGatedDeltaRuleFwdH(const char *name) : OpDef(name) + { + this->Input("k") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + this->Input("w") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + this->Input("u") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + this->Input("g") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + this->Input("inital_state") + .ParamType(OPTIONAL) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .IgnoreContiguous(); + + this->Input("cu_seqlens") + .ParamType(OPTIONAL) + .ValueDepend(OPTIONAL) + .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + this->Input("chunk_indices") + .ParamType(OPTIONAL) + .ValueDepend(OPTIONAL) + .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + this->Output("h") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + this->Output("v_new") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + this->Output("final_state") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + this->Attr("output_final_state").AttrType(REQUIRED).Bool(false); + this->Attr("chunk_size").AttrType(REQUIRED).Int(64); + this->Attr("inital_state_stride0").AttrType(REQUIRED).Int(0); + + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("prebuildPattern.value", "Opaque") + .ExtendCfgInfo("coreType.value", "AiCore") + .ExtendCfgInfo("jitCompile.flag", "static_false,dynamic_false"); + + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + this->AICore().AddConfig("ascend310p", aicore_config); + + } +}; + +OP_ADD(ChunkGatedDeltaRuleFwdH); + +} // namespace ops diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/chunk_gated_delta_rule_fwd_h_tiling.cpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/chunk_gated_delta_rule_fwd_h_tiling.cpp new file mode 100644 index 000000000..27172db94 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/chunk_gated_delta_rule_fwd_h_tiling.cpp @@ -0,0 +1,180 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +/*! + * \file chunk_gated_delta_rule_fwd_h_tiling.cpp + * \brief + */ + +#include "chunk_gated_delta_rule_fwd_h_tiling.h" +#include +#include "../tiling_base/data_copy_transpose_tiling.h" +#include "../tiling_base/tiling_templates_registry.h" + +namespace optiling { +static constexpr size_t INPUT_K_IDX = 0; +static constexpr size_t INPUT_W_IDX = 1; +static constexpr size_t INPUT_U_IDX = 2; +static constexpr size_t INPUT_G_IDX = 3; +static constexpr size_t INPUT_INITIAL_STATE_IDX = 4; +static constexpr size_t INPUT_SEQLENS_IDX = 5; +static constexpr size_t INPUT_CHUNK_INDICES_IDX = 6; + +static constexpr size_t ATTR_STORE_FINAL_STATE_IDX = 0; +static constexpr size_t ATTR_CHUNK_SIZE_IDX = 1; +static constexpr size_t ATTR_INITIAL_STATE_STRIDE_IDX = 2; + +static constexpr size_t DIM_BATCH = 0; +static constexpr size_t DIM_HEAD_NUM = 1; +static constexpr size_t DIM_SEQLEN = 2; +static constexpr size_t DIM_HEAD_DIM = 3; + + +static void ChunkGatedDeltaRuleFwdHTilingDataPrint(gert::TilingContext *context, ChunkGatedDeltaRuleFwdHTilingData &tiling) +{ + auto nodeName = context->GetNodeName(); + OP_LOGD(nodeName, ">>>>>>>>>>>>>>> Start to print ChunkGatedDeltaRuleFwdH tiling data <<<<<<<<<<<<<<<<"); + OP_LOGD(nodeName, "=== batch: %ld", tiling.get_batch()); + OP_LOGD(nodeName, "=== seqlen: %ld", tiling.get_seqlen()); + OP_LOGD(nodeName, "=== kNumHead: %ld", tiling.get_kNumHead()); + OP_LOGD(nodeName, "=== vNumHead: %ld", tiling.get_vNumHead()); + OP_LOGD(nodeName, "=== kHeadDim: %ld", tiling.get_kHeadDim()); + OP_LOGD(nodeName, "=== vHeadDim: %ld", tiling.get_vHeadDim()); + OP_LOGD(nodeName, "=== chunkSize: %ld", tiling.get_chunkSize()); + OP_LOGD(nodeName, "=== useInitialState: %ld", tiling.get_useInitialState()); + OP_LOGD(nodeName, "=== storeFinalState: %ld", tiling.get_storeFinalState()); + OP_LOGD(nodeName, "=== dataType: %ld", tiling.get_dataType()); + OP_LOGD(nodeName, "=== isVariedLen: %ld", tiling.get_isVariedLen()); + OP_LOGD(nodeName, "=== shapeBatch: %ld", tiling.get_shapeBatch()); + OP_LOGD(nodeName, "=== tokenBatch: %f", tiling.get_tokenBatch()); + OP_LOGD(nodeName, ">>>>>>>>>>>>>>> Print ChunkGatedDeltaRuleFwdH tiling data end <<<<<<<<<<<<<<<<"); +} + +ge::graphStatus Tiling4ChunkGatedDeltaRuleFwdH(gert::TilingContext *context) +{ + OP_LOGD(context->GetNodeName(), "Tiling4ChunkGatedDeltaRuleFwdH start."); + ChunkGatedDeltaRuleFwdHTilingData tiling; + + gert::Shape kStorageShape = context->GetOptionalInputShape(INPUT_K_IDX)->GetStorageShape(); + gert::Shape uStorageShape = context->GetOptionalInputShape(INPUT_U_IDX)->GetStorageShape(); + + int64_t seqlen = kStorageShape.GetDim(DIM_SEQLEN); + int64_t kNumHead = kStorageShape.GetDim(DIM_HEAD_NUM); + int64_t vNumHead = uStorageShape.GetDim(DIM_HEAD_NUM); + int64_t kHeadDim = kStorageShape.GetDim(DIM_HEAD_DIM); + int64_t vHeadDim = uStorageShape.GetDim(DIM_HEAD_DIM); + int64_t batch, isVariedLen, shapeBatch, tokenBatch; + + auto cuSeqlensTensor = context->GetOptionalInputTensor(INPUT_SEQLENS_IDX); + if (cuSeqlensTensor == nullptr) { + isVariedLen = false; + shapeBatch = kStorageShape.GetDim(DIM_BATCH); + tokenBatch = 1; + batch = shapeBatch; + } else { + isVariedLen = true; + shapeBatch = 1; + tokenBatch = cuSeqlensTensor->GetStorageShape().GetDim(DIM_BATCH) - 1; + batch = tokenBatch; + } + + auto initialStateTensor = context->GetOptionalInputTensor(INPUT_INITIAL_STATE_IDX); + bool useInitialState = initialStateTensor != nullptr; + int64_t stateDataType = 2; + if (useInitialState) { + auto stateDType = initialStateTensor->GetDataType(); + if (stateDType == ge::DT_BF16) { + stateDataType = 1; + } else if (stateDType == ge::DT_FLOAT16) { + stateDataType = 0; + } + } + + auto gDType = context->GetOptionalInputTensor(INPUT_G_IDX)->GetDataType(); + int64_t gDataType = 2; + if (gDType == ge::DT_BF16) { + gDataType = 1; + } else if (gDType == ge::DT_FLOAT16) { + gDataType = 0; + } + + auto attrPtr = context->GetAttrs(); + bool storeFinalState = *(attrPtr->GetAttrPointer(ATTR_STORE_FINAL_STATE_IDX)); + int64_t chunkSize = *(attrPtr->GetAttrPointer(ATTR_CHUNK_SIZE_IDX)); + int64_t initalStateStride0 = *(attrPtr->GetAttrPointer(ATTR_INITIAL_STATE_STRIDE_IDX)); + + auto dtype = context->GetInputTensor(0)->GetDataType(); + uint64_t dataType = dtype == ge::DT_BF16 ? 1 : 0; + + const auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + uint32_t aicCoreNum = ascendcPlatform.GetCoreNumAic(); + context->SetBlockDim(aicCoreNum); + + constexpr size_t WORKSPACE_RSV_BYTE = 16 * 1024 * 1024; + constexpr size_t GM_ALIGN = 512; + constexpr int64_t PING_PONG_STAGES = 2; + + size_t workspaceOffset = ascendcPlatform.GetLibApiWorkSpaceSize(); + workspaceOffset += WORKSPACE_RSV_BYTE; + + tiling.set_vWorkspaceOffset(workspaceOffset); + workspaceOffset += (aicCoreNum * chunkSize * vHeadDim * sizeof(float) * PING_PONG_STAGES + GM_ALIGN) / GM_ALIGN * GM_ALIGN; + + tiling.set_vUpdateWorkspaceOffset(workspaceOffset); + workspaceOffset += (aicCoreNum * chunkSize * vHeadDim * sizeof(float) * PING_PONG_STAGES + GM_ALIGN) / GM_ALIGN * GM_ALIGN; + + tiling.set_hWorkspaceOffset(workspaceOffset); + workspaceOffset += (aicCoreNum * kHeadDim * vHeadDim * sizeof(float) * PING_PONG_STAGES + GM_ALIGN) / GM_ALIGN * GM_ALIGN; + + tiling.set_numSeqWorkspaceOffset(workspaceOffset); + workspaceOffset += ((tokenBatch + 1) * sizeof(int64_t) + GM_ALIGN) / GM_ALIGN * GM_ALIGN; + + tiling.set_numChunksWorkspaceOffset(workspaceOffset); + workspaceOffset += ((tokenBatch + 1) * sizeof(int64_t) + GM_ALIGN) / GM_ALIGN * GM_ALIGN; + + workspaceOffset += WORKSPACE_RSV_BYTE; + size_t *currentWorkspace = context->GetWorkspaceSizes(1); + currentWorkspace[0] = (workspaceOffset - 0); + + tiling.set_batch(batch); + tiling.set_seqlen(seqlen); + tiling.set_kNumHead(kNumHead); + tiling.set_vNumHead(vNumHead); + tiling.set_kHeadDim(kHeadDim); + tiling.set_vHeadDim(vHeadDim); + tiling.set_chunkSize(chunkSize); + tiling.set_initalStateStride0(initalStateStride0); + tiling.set_useInitialState(useInitialState); + tiling.set_storeFinalState(storeFinalState); + tiling.set_dataType(dataType); + tiling.set_stateDataType(stateDataType); + tiling.set_gDataType(gDataType); + tiling.set_isVariedLen(isVariedLen); + tiling.set_shapeBatch(shapeBatch); + tiling.set_tokenBatch(tokenBatch); + + tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); + + ChunkGatedDeltaRuleFwdHTilingDataPrint(context, tiling); + OP_LOGD(context->GetNodeName(), "Tiling4ChunkGatedDeltaRuleFwdH end."); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingPrepareForChunkGatedDeltaRuleFwdH(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(ChunkGatedDeltaRuleFwdH) + .Tiling(Tiling4ChunkGatedDeltaRuleFwdH) + .TilingParse(TilingPrepareForChunkGatedDeltaRuleFwdH); + +} // namespace optiling diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/chunk_gated_delta_rule_fwd_h_tiling.h b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/chunk_gated_delta_rule_fwd_h_tiling.h new file mode 100644 index 000000000..44f3a01a9 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/chunk_gated_delta_rule_fwd_h_tiling.h @@ -0,0 +1,50 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +/*! + * \file chunk_gated_delta_rule_fwd_h_tiling.h + * \brief + */ + +#pragma once + +#include +#include +#include + +namespace optiling { + +BEGIN_TILING_DATA_DEF(ChunkGatedDeltaRuleFwdHTilingData) +TILING_DATA_FIELD_DEF(int64_t, batch); +TILING_DATA_FIELD_DEF(int64_t, seqlen); +TILING_DATA_FIELD_DEF(int64_t, kNumHead); +TILING_DATA_FIELD_DEF(int64_t, vNumHead); +TILING_DATA_FIELD_DEF(int64_t, kHeadDim); +TILING_DATA_FIELD_DEF(int64_t, vHeadDim); +TILING_DATA_FIELD_DEF(int64_t, chunkSize); +TILING_DATA_FIELD_DEF(int64_t, initalStateStride0); +TILING_DATA_FIELD_DEF(bool, useInitialState); +TILING_DATA_FIELD_DEF(bool, storeFinalState); +TILING_DATA_FIELD_DEF(int64_t, dataType); +TILING_DATA_FIELD_DEF(int64_t, gDataType); +TILING_DATA_FIELD_DEF(int64_t, stateDataType); +TILING_DATA_FIELD_DEF(int64_t, isVariedLen); +TILING_DATA_FIELD_DEF(int64_t, shapeBatch); +TILING_DATA_FIELD_DEF(int64_t, tokenBatch); +TILING_DATA_FIELD_DEF(int64_t, vWorkspaceOffset); +TILING_DATA_FIELD_DEF(int64_t, vUpdateWorkspaceOffset); +TILING_DATA_FIELD_DEF(int64_t, hWorkspaceOffset); +TILING_DATA_FIELD_DEF(int64_t, numSeqWorkspaceOffset); +TILING_DATA_FIELD_DEF(int64_t, numChunksWorkspaceOffset); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(ChunkGatedDeltaRuleFwdH, ChunkGatedDeltaRuleFwdHTilingData) + +struct ChunkGatedDeltaRuleFwdHCompileInfo {}; +} // namespace optiling diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/op_api/aclnn_chunk_gated_delta_rule_fwd_h.cpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/op_api/aclnn_chunk_gated_delta_rule_fwd_h.cpp new file mode 100644 index 000000000..55386a4c8 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/op_api/aclnn_chunk_gated_delta_rule_fwd_h.cpp @@ -0,0 +1,230 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#include "aclnn_chunk_gated_delta_rule_fwd_h.h" +#include "chunk_gated_delta_rule_fwd_h.h" +#include +#include +#include + +#include "aclnn_kernels/transdata.h" +#include "aclnn_kernels/contiguous.h" +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "opdev/common_types.h" +#include "opdev/data_type_utils.h" +#include "opdev/format_utils.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/platform.h" +#include "opdev/shape_utils.h" +#include "opdev/tensor_view_utils.h" +#include "opdev/make_op_executor.h" + + +using namespace op; + +#ifdef __cplusplus +extern "C" { +#endif + +struct ChunkGatedDeltaRuleFwdHParams { + const aclTensor *k = nullptr; + const aclTensor *w = nullptr; + const aclTensor *u = nullptr; + const aclTensor *gOptional = nullptr; + const aclTensor *gkOptional = nullptr; + const aclTensor *initalStateOptional = nullptr; + bool outputFinalState = false; + int64_t chunkSize = 64; + bool saveNewValue = true; + const aclIntArray *cuSeqlensOptional = nullptr; + const aclIntArray *chunkIndicesOptional = nullptr; + bool useExp2 = false; + bool transposeStateLayout = false; + const aclTensor *hOut = nullptr; + const aclTensor *vNewOut = nullptr; + const aclTensor *finalStateOut = nullptr; +}; + +static aclnnStatus CheckNotNull(ChunkGatedDeltaRuleFwdHParams params) +{ + CHECK_COND(params.k != nullptr, ACLNN_ERR_PARAM_NULLPTR, "k must not be nullptr."); + CHECK_COND(params.w != nullptr, ACLNN_ERR_PARAM_NULLPTR, "w must not be nullptr."); + CHECK_COND(params.u != nullptr, ACLNN_ERR_PARAM_NULLPTR, "u must not be nullptr."); + + CHECK_COND(params.hOut != nullptr, ACLNN_ERR_PARAM_NULLPTR, "hOut must not be nullptr."); + CHECK_COND(params.vNewOut != nullptr, ACLNN_ERR_PARAM_NULLPTR, "vNewOut must not be nullptr."); + return ACLNN_SUCCESS; +} + +static aclnnStatus CheckFormat(ChunkGatedDeltaRuleFwdHParams params) +{ + return ACLNN_SUCCESS; +} + +static aclnnStatus CheckShape(ChunkGatedDeltaRuleFwdHParams params) +{ + return ACLNN_SUCCESS; +} + +static aclnnStatus CheckDtype(ChunkGatedDeltaRuleFwdHParams params) +{ + return ACLNN_SUCCESS; +} + +static aclnnStatus DataContiguous(const aclTensor *&tensor, aclOpExecutor *executor) +{ + tensor = l0op::Contiguous(tensor, executor); + CHECK_RET(tensor != nullptr, ACLNN_ERR_INNER_NULLPTR); + return ACLNN_SUCCESS; +} + +static aclnnStatus ParamsDataContiguous(ChunkGatedDeltaRuleFwdHParams ¶ms, aclOpExecutor *executorPtr) +{ + CHECK_COND(DataContiguous(params.k, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "Contiguous k failed."); + CHECK_COND(DataContiguous(params.w, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "Contiguous w failed."); + CHECK_COND(DataContiguous(params.u, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "Contiguous u failed."); + CHECK_COND(DataContiguous(params.gOptional, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "Contiguous gOptional failed."); + if (params.initalStateOptional != nullptr) { + CHECK_COND(DataContiguous(params.initalStateOptional, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "Contiguous initalStateOptional failed."); + } + + return ACLNN_SUCCESS; +} + +static aclnnStatus CheckGOptionalNonNull(const ChunkGatedDeltaRuleFwdHParams ¶ms) +{ + CHECK_COND(params.gOptional != nullptr, ACLNN_ERR_PARAM_INVALID, + "g is an optional-parameter slot in the API but only a non-null aclTensor is supported; nullptr is not allowed until g=None is implemented."); + return ACLNN_SUCCESS; +} + +static aclnnStatus CheckReservedOptions(const ChunkGatedDeltaRuleFwdHParams ¶ms) +{ + CHECK_COND(params.gkOptional == nullptr, ACLNN_ERR_PARAM_INVALID, + "gk is reserved for ChunkGatedDeltaRuleFwdH and must be nullptr."); + CHECK_COND(params.saveNewValue, ACLNN_ERR_PARAM_INVALID, + "save_new_value is reserved and only true is supported."); + CHECK_COND(!params.useExp2, ACLNN_ERR_PARAM_INVALID, + "use_exp2 is reserved and only false is supported."); + CHECK_COND(!params.transposeStateLayout, ACLNN_ERR_PARAM_INVALID, + "transpose_state_layout is reserved and only false is supported."); + return ACLNN_SUCCESS; +} + +static aclnnStatus CheckParams(ChunkGatedDeltaRuleFwdHParams params) +{ + CHECK_RET(CheckNotNull(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + CHECK_RET(CheckGOptionalNonNull(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + CHECK_RET(CheckReservedOptions(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + CHECK_RET(CheckFormat(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + CHECK_RET(CheckShape(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + CHECK_RET(CheckDtype(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + return ACLNN_SUCCESS; +} + +aclnnStatus aclnnChunkGatedDeltaRuleFwdHGetWorkspaceSize( + const aclTensor *k, + const aclTensor *w, + const aclTensor *u, + const aclTensor *gOptional, + const aclTensor *gkOptional, + const aclTensor *initalStateOptional, + bool outputFinalState, + int64_t chunkSize, + bool saveNewValue, + const aclIntArray *cuSeqlensOptional, + const aclIntArray *chunkIndicesOptional, + bool useExp2, + bool transposeStateLayout, + const aclTensor *hOut, + const aclTensor *vNewOut, + const aclTensor *finalStateOut, + uint64_t *workspaceSize, + aclOpExecutor **executor) +{ + ChunkGatedDeltaRuleFwdHParams params{k, + w, + u, + gOptional, + gkOptional, + initalStateOptional, + outputFinalState, + chunkSize, + saveNewValue, + cuSeqlensOptional, + chunkIndicesOptional, + useExp2, + transposeStateLayout, + hOut, + vNewOut, + finalStateOut}; + // Standard syntax, Check parameters. + L2_DFX_PHASE_1(aclnnChunkGatedDeltaRuleFwdH, + DFX_IN(k, w, u, gOptional, gkOptional, initalStateOptional, cuSeqlensOptional, chunkIndicesOptional), + DFX_OUT(hOut, vNewOut, finalStateOut)); + auto uniqueExecutor = CREATE_EXECUTOR(); + CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + auto executorPtr = uniqueExecutor.get(); + auto ret = CheckParams(params); + CHECK_RET(ret == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID); + CHECK_COND(ParamsDataContiguous(params, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "ParamsDataContiguous failed."); + + // aclGetViewStrides obtains the strides and the number of strides corresponding to aclTensor + int64_t *initialStateStridesValuePtr = nullptr; + int64_t initialStateStridesValue = 0; + uint64_t initialStateStridesNum = 0; + + if (initalStateOptional != nullptr) { + ret = aclGetViewStrides(initalStateOptional, &initialStateStridesValuePtr, &initialStateStridesNum); + CHECK_RET(ret == ACLNN_SUCCESS, ret); + initialStateStridesValue = initialStateStridesValuePtr[initialStateStridesNum - 2]; + } + + auto result = l0op::ChunkGatedDeltaRuleFwdH(params.k, params.w, params.u, params.gOptional, params.initalStateOptional, params.cuSeqlensOptional, params.chunkIndicesOptional, params.outputFinalState, params.chunkSize, initialStateStridesValue, params.hOut, params.vNewOut, params.finalStateOut, executorPtr); + CHECK_RET(result[0] != nullptr, ACLNN_ERR_PARAM_NULLPTR); + + // If the output tensor is non-contiguous, convert the calculated contiguous tensor to non-contiguous. + auto viewCopyResult0 = l0op::ViewCopy(result[0], params.hOut, executorPtr); + CHECK_RET(viewCopyResult0 != nullptr, ACLNN_ERR_INNER_NULLPTR); + auto viewCopyResult1 = l0op::ViewCopy(result[1], params.vNewOut, executorPtr); + CHECK_RET(viewCopyResult1 != nullptr, ACLNN_ERR_INNER_NULLPTR); + if (outputFinalState && params.finalStateOut != nullptr) { + auto viewCopyResult2 = l0op::ViewCopy(result[2], params.finalStateOut, executorPtr); + CHECK_RET(viewCopyResult2 != nullptr, ACLNN_ERR_INNER_NULLPTR); + } + + // Standard syntax, get the size of workspace needed during computation. + *workspaceSize = uniqueExecutor->GetWorkspaceSize(); + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; +} + + +aclnnStatus aclnnChunkGatedDeltaRuleFwdH(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream) +{ + L2_DFX_PHASE_2(aclnnChunkGatedDeltaRuleFwdH); + CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER, + "This is an error in ChunkGatedDeltaRuleFwdH launch aicore."); + return ACLNN_SUCCESS; +} + + +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/op_api/aclnn_chunk_gated_delta_rule_fwd_h.h b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/op_api/aclnn_chunk_gated_delta_rule_fwd_h.h new file mode 100644 index 000000000..956c9f63f --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/op_api/aclnn_chunk_gated_delta_rule_fwd_h.h @@ -0,0 +1,77 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_API_INC_ACLNN_CHUNK_GATED_DELTA_RULE_FWD_H_H +#define OP_API_INC_ACLNN_CHUNK_GATED_DELTA_RULE_FWD_H_H +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/* function: aclnnChunkGatedDeltaRuleFwdHGetWorkspaceSize + * parameters (order aligned with chunk_gated_delta_rule_fwd_h Python API): + * k : required + * w : required + * u : required + * gOptional : optional, only non-null aclTensor is supported + * gkOptional : optional, reserved (must be nullptr) + * initalStateOptional : optional + * outputFinalState : required + * chunkSize : required + * saveNewValue : reserved (must be true) + * cuSeqlensOptional : optional + * chunkIndicesOptional : optional + * useExp2 : reserved (must be false) + * transposeStateLayout : reserved (must be false) + * hOut : required + * vNewOut : required + * finalStateOut : optional + * workspaceSize : size of workspace(output). + * executor : executor context(output). + */ +__attribute__((visibility("default"))) +aclnnStatus aclnnChunkGatedDeltaRuleFwdHGetWorkspaceSize( + const aclTensor *k, + const aclTensor *w, + const aclTensor *u, + const aclTensor *gOptional, + const aclTensor *gkOptional, + const aclTensor *initalStateOptional, + bool outputFinalState, + int64_t chunkSize, + bool saveNewValue, + const aclIntArray *cuSeqlensOptional, + const aclIntArray *chunkIndicesOptional, + bool useExp2, + bool transposeStateLayout, + const aclTensor *hOut, + const aclTensor *vNewOut, + const aclTensor *finalStateOut, + uint64_t *workspaceSize, + aclOpExecutor **executor); + +/* function: aclnnChunkGatedDeltaRuleFwdH + * parameters : + * workspace : workspace memory addr(input). + * workspaceSize : size of workspace(input). + * executor : executor context(input). + * stream : acl stream. + */ +__attribute__((visibility("default"))) +aclnnStatus aclnnChunkGatedDeltaRuleFwdH( + void *workspace, + uint64_t workspaceSize, + aclOpExecutor *executor, + aclrtStream stream); +#ifdef __cplusplus +} +#endif + +#endif diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/op_api/chunk_gated_delta_rule_fwd_h.cpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/op_api/chunk_gated_delta_rule_fwd_h.cpp new file mode 100644 index 000000000..eb05d74eb --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/op_api/chunk_gated_delta_rule_fwd_h.cpp @@ -0,0 +1,70 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "opdev/op_log.h" +#include "opdev/op_dfx.h" +#include "opdev/make_op_executor.h" +#include "chunk_gated_delta_rule_fwd_h.h" + +using namespace op; + +namespace l0op { +OP_TYPE_REGISTER(ChunkGatedDeltaRuleFwdH); + +const std::array ChunkGatedDeltaRuleFwdH( + const aclTensor *k, + const aclTensor *w, + const aclTensor *u, + const aclTensor *g, + const aclTensor *initalStateOptional, + const aclIntArray *cuSeqlensOptional, + const aclIntArray *chunkIndicesOptional, + bool outputFinalState, + int64_t chunkSize, + int64_t initialStateStridesValue, + const aclTensor *hOut, + const aclTensor *vNewOut, + const aclTensor *finalStateOut, + aclOpExecutor *executor) +{ + L0_DFX(ChunkGatedDeltaRuleFwdH, k, w, u, g, initalStateOptional, cuSeqlensOptional, chunkIndicesOptional, outputFinalState, chunkSize, initialStateStridesValue, hOut, vNewOut, finalStateOut); + + const aclTensor *actualCuSeqlens = nullptr; + if (cuSeqlensOptional) { + actualCuSeqlens = executor->ConvertToTensor(cuSeqlensOptional, DataType::DT_INT64); + const_cast(actualCuSeqlens)->SetStorageFormat(Format::FORMAT_ND); + const_cast(actualCuSeqlens)->SetViewFormat(Format::FORMAT_ND); + const_cast(actualCuSeqlens)->SetOriginalFormat(Format::FORMAT_ND); + } else { + actualCuSeqlens = nullptr; + } + + const aclTensor *actualChunkIndices = nullptr; + if (chunkIndicesOptional) { + actualChunkIndices = executor->ConvertToTensor(chunkIndicesOptional, DataType::DT_INT64); + const_cast(actualChunkIndices)->SetStorageFormat(Format::FORMAT_ND); + const_cast(actualChunkIndices)->SetViewFormat(Format::FORMAT_ND); + const_cast(actualChunkIndices)->SetOriginalFormat(Format::FORMAT_ND); + } else { + actualChunkIndices = nullptr; + } + + auto ret = ADD_TO_LAUNCHER_LIST_AICORE(ChunkGatedDeltaRuleFwdH, + OP_INPUT(k, w, u, g, initalStateOptional, actualCuSeqlens, actualChunkIndices), + OP_OUTPUT(hOut, vNewOut, finalStateOut), + OP_ATTR(outputFinalState, chunkSize, initialStateStridesValue)); + if (ret != ACLNN_SUCCESS) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "ADD_TO_LAUNCHER_LIST_AICORE failed."); + return {nullptr, nullptr, nullptr}; + } + return {hOut, vNewOut, finalStateOut}; +} + +} // namespace l0op diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/op_api/chunk_gated_delta_rule_fwd_h.h b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/op_api/chunk_gated_delta_rule_fwd_h.h new file mode 100644 index 000000000..98817016c --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_host/op_api/chunk_gated_delta_rule_fwd_h.h @@ -0,0 +1,33 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_API_INC_LEVEL0_OP_CHUNK_GATED_DELTA_RULE_FWD_H_H +#define OP_API_INC_LEVEL0_OP_CHUNK_GATED_DELTA_RULE_FWD_H_H + +#include "opdev/op_executor.h" + +namespace l0op { +const std::array ChunkGatedDeltaRuleFwdH( + const aclTensor *k, + const aclTensor *w, + const aclTensor *u, + const aclTensor *g, + const aclTensor *initalStateOptional, + const aclIntArray *cuSeqlensOptional, + const aclIntArray *chunkIndicesOptional, + bool outputFinalState, + int64_t chunkSize, + int64_t initialStateStridesValue, + const aclTensor *hOut, + const aclTensor *vNewOut, + const aclTensor *finalStateOut, + aclOpExecutor *executor); +} + +#endif diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/compat_310p.h b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/compat_310p.h new file mode 100644 index 000000000..288c47764 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/compat_310p.h @@ -0,0 +1,39 @@ +#ifndef COMPAT_310P_H +#define COMPAT_310P_H + +#ifndef __CCE_KT_TEST__ +#include "kernel_operator.h" +#endif + +// Dummy bfloat16_t only needed on 310P (dav_m200) where the compiler +// doesn't provide a native bf16 type. On 910B/910C the compiler's +// __clang_cce_types.h already typedefs bfloat16_t from __bf16. +#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 200) && !defined(__bfloat16_t_defined) +#define __bfloat16_t_defined +#define __COMPAT_310P_ACTIVE__ +struct bfloat16_t { + uint16_t val; + bfloat16_t() = default; + bfloat16_t(float v) : val(0) { (void)v; } + operator float() const { return 0.f; } +}; +#endif + +// 310P has no fixpipe unit; post-matmul stores go through MTE3 +#ifndef PIPE_FIX +#define PIPE_FIX PIPE_MTE3 +#endif + +// 310P renames LoadDataWithSparse → LoadDataWithSparseCal +#ifdef __COMPAT_310P_ACTIVE__ +#define LoadDataWithSparse LoadDataWithSparseCal +#endif + +// 310P has no AscendC::ToFloat — dummy bfloat16_t already has operator float() +#ifdef __COMPAT_310P_ACTIVE__ +namespace AscendC { + inline float ToFloat(bfloat16_t v) { return (float)v; } +} +#endif + +#endif diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/epilogue/block/block_epilogue_gdn_fwdh_update.hpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/epilogue/block/block_epilogue_gdn_fwdh_update.hpp new file mode 100644 index 000000000..514ee1b3b --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/epilogue/block/block_epilogue_gdn_fwdh_update.hpp @@ -0,0 +1,179 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDH_UPDATE_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDH_UPDATE_HPP +#include "catlass/catlass.hpp" +#include "catlass/arch/resource.hpp" +#include "../gdn_fwd_h_epilogue_policies.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/matrix_coord.hpp" +#include "catlass/epilogue/tile/tile_copy.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class HOutputType_, + class GInputType_, + class HInputType_, + class HUpdateInputType_, + class FinalStateType_ +> +class BlockEpilogue < + EpilogueAtlasGDNFwdHUpdate, + HOutputType_, + GInputType_, + HInputType_, + HUpdateInputType_, + FinalStateType_ +> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasGDNFwdHUpdate; + using ArchTag = typename DispatchPolicy::ArchTag; + + using HElementOutput = typename HOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using HElementInput = typename HInputType_::Element; + using HUpdateElementInput = typename HUpdateInputType_::Element; + using FinalStateElement = typename FinalStateType_::Element; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + + // Bumped layout to fit kHeadDim up to 256 with subBlockNum=2 (per-subblock M up to 128). + // Required: calc (fp32) up to 128*128*4=64KB; h (fp16) up to 128*128*2=32KB; + // hUpdate/hOutput/finalOutput at the same offset, max needed 64KB; glast small. + constexpr uint32_t CALC_BUF_OFFSET = 0; + constexpr uint32_t PING_BUF_0_OFFSET = 64 * 1024; + constexpr uint32_t PING_BUF_1_OFFSET = 96 * 1024; + constexpr uint32_t PING_BUF_2_OFFSET = 112 * 1024; + constexpr uint32_t PING_G_BUF_OFFSET = 160 * 1024; + + + calcUbTensor = resource.ubBuf.template GetBufferByByte(CALC_BUF_OFFSET); + + hUpdateUbTensor = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + hUbTensor = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + + hOutputUbTensor = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + finalOutputUbTensor = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + + glastUbTensor = resource.ubBuf.template GetBufferByByte(PING_G_BUF_OFFSET); + + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor hOutput, + AscendC::GlobalTensor finalState, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor hInput, + AscendC::GlobalTensor hUpdateInput, + uint32_t chunkSize, + uint32_t kHeadDim, + uint32_t vHeadDim, + Arch::CrossCoreFlag cube2Done, + bool isFinalState + ) + { + uint32_t mActual = kHeadDim; + uint32_t nActual = vHeadDim; + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t mActualPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t mActualThisSubBlock = (subBlockIdx == 0) ? mActualPerSubBlock : (mActual - mActualPerSubBlock); + uint32_t mOffset = subBlockIdx * mActualPerSubBlock; + uint32_t nOffset = 0; + int64_t offsetH = mOffset * nActual + nOffset; + + AscendC::ResetMask(); + + AscendC::GlobalTensor hOutputThisSubBlock = hOutput[offsetH]; + AscendC::GlobalTensor gInputThisSubBlock = gInput; + AscendC::GlobalTensor hInputThisSubBlock = hInput[offsetH]; + AscendC::GlobalTensor hUpdateInputThisSubBlock = hUpdateInput[offsetH]; + AscendC::GlobalTensor finalStateThisSubBlock = finalState[offsetH]; + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopy(hUbTensor, hInputThisSubBlock, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::Cast(calcUbTensor, hUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + + GElementInput gLastVal = gInputThisSubBlock.GetValue(chunkSize-1); + float gLastFloat = 0.0f; + if constexpr(std::is_same::value) { + gLastFloat = gLastVal; + } else if constexpr(std::is_same::value) { + gLastFloat = (float)gLastVal; + } else if constexpr(std::is_same::value) { + gLastFloat = AscendC::ToFloat(gLastVal); + } + glastUbTensor.SetValue(0, gLastFloat); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::Exp(glastUbTensor, glastUbTensor, 1); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + float muls = glastUbTensor.GetValue(0); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::Muls(calcUbTensor, calcUbTensor, muls, mActualThisSubBlock * nActual); + + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + AscendC::DataCopy(hUpdateUbTensor, hUpdateInputThisSubBlock, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + AscendC::Add(hUpdateUbTensor, calcUbTensor, hUpdateUbTensor, mActualThisSubBlock * nActual); + + if (isFinalState) { + if constexpr(!std::is_same::value) { + AscendC::PipeBarrier(); + AscendC::Cast(finalOutputUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::DataCopy(finalStateThisSubBlock, finalOutputUbTensor, mActualThisSubBlock * nActual); + } else { + AscendC::PipeBarrier(); + AscendC::DataCopy(finalStateThisSubBlock, hUpdateUbTensor, mActualThisSubBlock * nActual); + } + } else { + AscendC::PipeBarrier(); + AscendC::Cast(hOutputUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::DataCopy(hOutputThisSubBlock, hOutputUbTensor, mActualThisSubBlock * nActual); + } + } + +private: + AscendC::LocalTensor calcUbTensor; + + AscendC::LocalTensor hUbTensor; + AscendC::LocalTensor hUpdateUbTensor; + + AscendC::LocalTensor hOutputUbTensor; + AscendC::LocalTensor finalOutputUbTensor; + + AscendC::LocalTensor glastUbTensor; + +}; +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/epilogue/block/block_epilogue_gdn_fwdh_vnew.hpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/epilogue/block/block_epilogue_gdn_fwdh_vnew.hpp new file mode 100644 index 000000000..deae4f114 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/epilogue/block/block_epilogue_gdn_fwdh_vnew.hpp @@ -0,0 +1,258 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDH_VNEW_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDH_VNEW_HPP +#include "catlass/catlass.hpp" +#include "catlass/arch/resource.hpp" +#include "../gdn_fwd_h_epilogue_policies.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/matrix_coord.hpp" +#include "catlass/epilogue/tile/tile_copy.hpp" + + + +namespace Catlass::Epilogue::Block { + +template < + class VOutputType_, + class GInputType_, + class UInputType_, + class WSInputType_ +> +class BlockEpilogue < + EpilogueAtlasGDNFwdHVnew, + VOutputType_, + GInputType_, + UInputType_, + WSInputType_ +> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasGDNFwdHVnew; + using ArchTag = typename DispatchPolicy::ArchTag; + + using VElementOutput = typename VOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using UElementInput = typename UInputType_::Element; + using WSElementInput = typename WSInputType_::Element; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + + constexpr uint32_t CALC_BUF_OFFSET = 0; + constexpr uint32_t PING_BUF_0_OFFSET = 32 * 1024; + constexpr uint32_t PING_BUF_1_OFFSET = 64 * 1024; + constexpr uint32_t PONG_BUF_0_OFFSET = 96 * 1024; + constexpr uint32_t PONG_BUF_1_OFFSET = 128 * 1024; + constexpr uint32_t PING_G_BUF_OFFSET = 160 * 1024; + constexpr uint32_t PONG_G_BUF_OFFSET = 161 * 1024; + constexpr uint32_t PING_G_SUB_BUF_OFFSET = 162 * 1024; + constexpr uint32_t PONG_G_SUB_BUF_OFFSET = 163 * 1024; + constexpr uint32_t PING_G_INPUT_BUF_OFFSET = 164 * 1024; + constexpr uint32_t PONG_G_INPUT_BUF_OFFSET = 165 * 1024; + constexpr uint32_t SHARE_BUF_OFFSET = 166 * 1024; + + calcUbTensor = resource.ubBuf.template GetBufferByByte(CALC_BUF_OFFSET); + + uUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + uUbFloatTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + wsUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + gUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_BUF_OFFSET); + gLastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + gInputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_INPUT_BUF_OFFSET); + vNewOutputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + vNewDecayUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + + uUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_1_OFFSET); + uUbFloatTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_0_OFFSET); + wsUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_1_OFFSET); + gUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_BUF_OFFSET); + gLastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + gInputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_INPUT_BUF_OFFSET); + vNewOutputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_1_OFFSET); + vNewDecayUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_0_OFFSET); + + shareBuffer_ = resource.ubBuf.template GetBufferByByte(SHARE_BUF_OFFSET); + + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor vnewOutput, + AscendC::GlobalTensor vnewdecayOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor uInput, + AscendC::GlobalTensor wsInput, + uint32_t chunkSize, + uint32_t kHeadDim, + uint32_t vHeadDim, + Arch::CrossCoreFlag cube1Done + ) + { + uint32_t mActual = chunkSize; + uint32_t nkActual = kHeadDim; + uint32_t nvActual = vHeadDim; + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t mActualPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t mActualThisSubBlock = (subBlockIdx == 0) ? mActualPerSubBlock : (mActual - mActualPerSubBlock); + uint32_t mOffset = subBlockIdx * mActualPerSubBlock; + uint32_t nOffset = 0; + // 当前场景内部一定连续 + // k [B, H, T, D] + // g [B, H, T] + // 在外部offset的基础上进一步offset + // 当前asset kdim == vHeadDim + int64_t offsetK = mOffset * nvActual + nOffset; + int64_t offsetD = 0; // 因为要用最后一个数减去之前所有，所以全部读入 + + uint32_t gbrcStart, gbrcRealStart, gbrcReptime, gbrcEffStart, gbrcEffEnd; + if(subBlockIdx==0) + { + gbrcStart = 0; + gbrcRealStart = 0; + gbrcReptime = (mActualThisSubBlock + 8 - 1) / 8; + + } + else + { + gbrcStart = mActualPerSubBlock; + gbrcRealStart = gbrcStart & ~15; + gbrcReptime = (mActual - gbrcRealStart + 8 - 1) / 8; + } + gbrcEffStart = gbrcStart-gbrcRealStart; + gbrcEffEnd = gbrcEffStart + mActualThisSubBlock; + + AscendC::ResetMask(); + + AscendC::GlobalTensor vnewOutputThisSubBlock = vnewOutput[offsetK]; + AscendC::GlobalTensor vnewdecayOutputThisSubBlock = vnewdecayOutput[offsetK]; + AscendC::GlobalTensor gInputThisSubBlock = gInput; + AscendC::GlobalTensor uInputThisSubBlock = uInput[offsetK]; + AscendC::GlobalTensor wsInputThisSubBlock = wsInput[offsetK]; + + pingpongFlag = isFirst ? 0 : 4; + AscendC::LocalTensor uUbTensor = isFirst ? uUbTensor_ping : uUbTensor_pong; + AscendC::LocalTensor uUbFloatTensor = isFirst ? uUbFloatTensor_ping : uUbFloatTensor_pong; + AscendC::LocalTensor wsUbTensor = isFirst ? wsUbTensor_ping : wsUbTensor_pong; + AscendC::LocalTensor gUbTensor = isFirst ? gUbTensor_ping : gUbTensor_pong; + AscendC::LocalTensor gLastUbTensor = isFirst ? gLastUbTensor_ping : gLastUbTensor_pong; + AscendC::LocalTensor gInputUbTensor = isFirst ? gInputUbTensor_ping : gInputUbTensor_pong; + AscendC::LocalTensor vNewOutputUbTensor = isFirst ? vNewOutputUbTensor_ping : vNewOutputUbTensor_pong; + AscendC::LocalTensor vNewDecayUbTensor = isFirst ? vNewDecayUbTensor_ping : vNewDecayUbTensor_pong; + + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + if constexpr(std::is_same::value) { + AscendC::DataCopy(gUbTensor, gInputThisSubBlock, mActual); + } else { + AscendC::DataCopy(gInputUbTensor, gInputThisSubBlock, mActual); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + if constexpr(!std::is_same::value) { + AscendC::Cast(gUbTensor, gInputUbTensor, AscendC::RoundMode::CAST_NONE, mActual); + } + + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + float inputVal = gUbTensor.GetValue(mActual-1); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + + AscendC::PipeBarrier(); + AscendC::Duplicate(gLastUbTensor, inputVal, mActual); + AscendC::PipeBarrier(); + + AscendC::Sub(gUbTensor, gLastUbTensor, gUbTensor, mActual); + AscendC::PipeBarrier(); + + AscendC::Exp(gUbTensor, gUbTensor, mActual); + AscendC::PipeBarrier(); + + uint32_t dstShape_[2] = {gbrcReptime*8, nvActual}; + uint32_t srcShape_[2] = {gbrcReptime*8, 1}; + AscendC::Broadcast(calcUbTensor, gUbTensor[gbrcRealStart], dstShape_, srcShape_, shareBuffer_); + AscendC::PipeBarrier(); + + + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + + AscendC::DataCopy(uUbTensor, uInputThisSubBlock, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::Cast(uUbFloatTensor, uUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(wsUbTensor, wsInputThisSubBlock, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + + AscendC::Sub(uUbFloatTensor, uUbFloatTensor, wsUbTensor, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + AscendC::Cast(vNewOutputUbTensor, uUbFloatTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + AscendC::DataCopy(vnewOutputThisSubBlock, vNewOutputUbTensor, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + + AscendC::Mul(calcUbTensor[gbrcEffStart*nvActual], uUbFloatTensor, calcUbTensor[gbrcEffStart*nvActual], mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + AscendC::Cast(vNewDecayUbTensor, calcUbTensor[gbrcEffStart*nvActual], AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + AscendC::DataCopy(vnewdecayOutputThisSubBlock, vNewDecayUbTensor, mActualThisSubBlock * nvActual); + + if (isFirst) { + AscendC::PipeBarrier(); + } + + isFirst = false; + } + +private: + uint32_t pingpongFlag = 0; + bool isFirst = true; + + AscendC::LocalTensor calcUbTensor; + + AscendC::LocalTensor uUbTensor_ping; + AscendC::LocalTensor uUbFloatTensor_ping; + AscendC::LocalTensor wsUbTensor_ping; + AscendC::LocalTensor gUbTensor_ping; + AscendC::LocalTensor gLastUbTensor_ping; + AscendC::LocalTensor gInputUbTensor_ping; + AscendC::LocalTensor vNewOutputUbTensor_ping; + AscendC::LocalTensor vNewDecayUbTensor_ping; + + AscendC::LocalTensor uUbTensor_pong; + AscendC::LocalTensor uUbFloatTensor_pong; + AscendC::LocalTensor wsUbTensor_pong; + AscendC::LocalTensor gUbTensor_pong; + AscendC::LocalTensor gLastUbTensor_pong; + AscendC::LocalTensor gInputUbTensor_pong; + AscendC::LocalTensor vNewOutputUbTensor_pong; + AscendC::LocalTensor vNewDecayUbTensor_pong; + + AscendC::LocalTensor shareBuffer_; + +}; +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/epilogue/gdn_fwd_h_epilogue_policies.hpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/epilogue/gdn_fwd_h_epilogue_policies.hpp new file mode 100644 index 000000000..d5287d55d --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/epilogue/gdn_fwd_h_epilogue_policies.hpp @@ -0,0 +1,27 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_EPILOGUE_GDN_FWD_H_EPILOGUE_POLICIES_HPP +#define CATLASS_EPILOGUE_GDN_FWD_H_EPILOGUE_POLICIES_HPP + +#include "catlass/catlass.hpp" + +namespace Catlass::Epilogue { + +struct EpilogueAtlasGDNFwdHVnew { + using ArchTag = Arch::AtlasA2; +}; + +struct EpilogueAtlasGDNFwdHUpdate { + using ArchTag = Arch::AtlasA2; +}; + +} // namespace Catlass::Epilogue + +#endif // CATLASS_EPILOGUE_GDN_FWD_H_EPILOGUE_POLICIES_HPP diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/gemm/block/block_scheduler_gdn_fwd_h.hpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/gemm/block/block_scheduler_gdn_fwd_h.hpp new file mode 100644 index 000000000..e24d8283b --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/gemm/block/block_scheduler_gdn_fwd_h.hpp @@ -0,0 +1,285 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#include "catlass/gemm_coord.hpp" +using namespace Catlass; + +#ifndef CATLASS_GEMM_SCHEDULER_GDN_FWD_H_HPP +#define CATLASS_GEMM_SCHEDULER_GDN_FWD_H_HPP + +// constexpr uint32_t PING_PONG_STAGES = 1; +constexpr uint32_t PING_PONG_STAGES = 2; + +template +CATLASS_DEVICE T AlignUp(T a, T b) { + return (b == 0) ? 0 : (a + b - 1) / b * b; +} + +template +CATLASS_DEVICE T Min(T a, T b) { + return (a > b) ? b : a; +} + +template +CATLASS_DEVICE T Max(T a, T b) { + return (a > b) ? a : b; +} + +namespace Catlass::Gemm::Block { + +struct GDNFwdHOffsets { + uint32_t hSrcOffset; + uint32_t hDstOffset; + uint32_t uvOffset; + uint32_t wkOffset; + uint32_t wOffset; + uint32_t gOffset; + uint32_t hWorkOffset; + uint32_t vWorkOffset; + uint32_t initialStateOffset; + uint32_t finalStateOffset; + bool isInitialState; + bool isFinalState; + uint32_t blockTokens; + bool isDummyHead; + uint32_t batchIdx; + uint32_t headIdx; + uint32_t chunkIdx; + +}; + +struct BlockSchedulerGdnFwdH { + uint32_t batch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + uint32_t initalStateStride0; + uint32_t vBlockSize{128}; + uint32_t isVariedLen; + uint32_t shapeBatch; + uint32_t tokenBatch; + bool useInitialState; + bool storeFinalState; + uint32_t numSeqWorkspaceOffset; + uint32_t numChunksWorkspaceOffset; + + uint32_t taskIdx; + uint32_t taskLoops; + uint32_t cubeCoreIdx; + uint32_t cubeCoreNum; + uint32_t vLoops; + uint32_t taskNum; + uint32_t headGroups; + uint32_t totalChunks; + uint32_t totalTokens; + uint32_t headInnerLoop; + + uint32_t iterId {0}; + bool hasDummyHead; + bool isRunning; + bool processNewTask {true}; + bool firstLoop {true}; + bool lastLoop {false}; + GDNFwdHOffsets offsets[PING_PONG_STAGES]; + int32_t currStage{PING_PONG_STAGES - 1}; + + uint32_t vIdx; + uint32_t batchIdx; + uint32_t baseHeadIdx; + uint32_t chunkIdx; + uint32_t headInnerIdx; + uint32_t vHeadIdx; + uint32_t kHeadIdx; + uint32_t shapeBatchIdx; + uint32_t tokenBatchIdx; + + uint32_t chunkOffset; + uint32_t tokenOffset; + uint32_t batchChunks; + uint32_t batchTokens; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmNumSeq; + AscendC::GlobalTensor gmNumChunks; + + Arch::CrossCoreFlag cube1Done{0}; + Arch::CrossCoreFlag vec1Done{1}; + Arch::CrossCoreFlag cube2Done{2}; + Arch::CrossCoreFlag vec2Done{3}; + + CATLASS_DEVICE + BlockSchedulerGdnFwdH() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user, uint32_t coreIdx, uint32_t coreNum) { + __gm__ ChunkGatedDeltaRuleFwdHTilingData *__restrict gdnFwdHTilingData = reinterpret_cast<__gm__ ChunkGatedDeltaRuleFwdHTilingData *__restrict>(tiling); + + batch = gdnFwdHTilingData->batch; + seqlen = gdnFwdHTilingData->seqlen; + kNumHead = gdnFwdHTilingData->kNumHead; + vNumHead = gdnFwdHTilingData->vNumHead; + kHeadDim = gdnFwdHTilingData->kHeadDim; + vHeadDim = gdnFwdHTilingData->vHeadDim; + chunkSize = gdnFwdHTilingData->chunkSize; + initalStateStride0 = gdnFwdHTilingData->initalStateStride0; + isVariedLen = gdnFwdHTilingData->isVariedLen; + shapeBatch = gdnFwdHTilingData->shapeBatch; + tokenBatch = gdnFwdHTilingData->tokenBatch; + useInitialState = gdnFwdHTilingData->useInitialState; + storeFinalState = gdnFwdHTilingData->storeFinalState; + numSeqWorkspaceOffset = gdnFwdHTilingData->numSeqWorkspaceOffset; + numChunksWorkspaceOffset = gdnFwdHTilingData->numChunksWorkspaceOffset; + + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + if (isVariedLen) { + gmNumChunks.SetValue(0, 0); + gmNumSeq.SetValue(0, 0); + uint32_t actualBatch = 0; + int64_t prevSeq = 0, currSeq; + for (uint32_t b = 1; b <= tokenBatch; b++) { + currSeq = gmSeqlen.GetValue(b); + int64_t batchSeqLen = currSeq - prevSeq; + if (batchSeqLen > 0) { + actualBatch++; + gmNumSeq.SetValue(actualBatch, currSeq); + int64_t batchChunk = (batchSeqLen + chunkSize - 1) / chunkSize; + gmNumChunks.SetValue(actualBatch, gmNumChunks.GetValue(actualBatch - 1) + batchChunk); + } + prevSeq = currSeq; + } + tokenBatch = actualBatch; + batch = actualBatch; + totalChunks = gmNumChunks.GetValue(tokenBatch); + totalTokens = gmNumSeq.GetValue(tokenBatch); + } else { + totalChunks = (seqlen + chunkSize - 1) / chunkSize; + totalTokens = seqlen; + } + + cubeCoreIdx = coreIdx; + cubeCoreNum = coreNum; + vLoops = vHeadDim / vBlockSize; + taskNum = vLoops * batch * vNumHead; + headGroups = vNumHead / kNumHead; + hasDummyHead = (taskNum % (PING_PONG_STAGES * cubeCoreNum) <= cubeCoreNum) && (taskNum % (PING_PONG_STAGES * cubeCoreNum) > 0); + taskLoops = (taskNum + cubeCoreNum * PING_PONG_STAGES - 1) / (cubeCoreNum * PING_PONG_STAGES); + headInnerLoop = taskNum > cubeCoreNum ? PING_PONG_STAGES : 1; + taskIdx = cubeCoreIdx * headInnerLoop; + isRunning = taskIdx < taskNum; + + } + + CATLASS_DEVICE + void InitTask() { + iterId++; + currStage = (currStage + 1) % PING_PONG_STAGES; + if (processNewTask) { + if (taskIdx >= taskNum) { + lastLoop = true; + isRunning = false; + return; + } + vIdx = taskIdx / (batch * vNumHead); + batchIdx = (taskIdx - vIdx * batch * vNumHead) / vNumHead; + baseHeadIdx = taskIdx % vNumHead; + shapeBatchIdx = isVariedLen ? 0 : batchIdx; + tokenBatchIdx = isVariedLen ? batchIdx : 0; + chunkOffset = isVariedLen ? gmNumChunks.GetValue(tokenBatchIdx) : 0; + batchChunks = isVariedLen ? (gmNumChunks.GetValue(tokenBatchIdx + 1) - chunkOffset) : totalChunks; + tokenOffset = isVariedLen ? gmNumSeq.GetValue(tokenBatchIdx) : 0; + batchTokens = isVariedLen ? (gmNumSeq.GetValue(tokenBatchIdx + 1) - tokenOffset) : totalTokens; + chunkIdx = 0; + headInnerIdx = 0; + } else { + chunkIdx = headInnerIdx == PING_PONG_STAGES - 1 ? chunkIdx + 1 : chunkIdx; + headInnerIdx = (headInnerIdx + 1) % PING_PONG_STAGES; + } + + vHeadIdx = baseHeadIdx + headInnerIdx; + kHeadIdx = vHeadIdx / headGroups; + offsets[currStage].isInitialState = chunkIdx == 0; + offsets[currStage].isFinalState = chunkIdx == (batchChunks - 1); + offsets[currStage].initialStateOffset = (batchIdx * vNumHead + vHeadIdx) * kHeadDim * initalStateStride0; + offsets[currStage].finalStateOffset = (batchIdx * vNumHead + vHeadIdx) * kHeadDim * vHeadDim; + offsets[currStage].hSrcOffset = (shapeBatchIdx * vNumHead * totalChunks + vHeadIdx * totalChunks + chunkOffset + chunkIdx) * kHeadDim * vHeadDim; + offsets[currStage].hDstOffset = offsets[currStage].hSrcOffset + kHeadDim * vHeadDim; + offsets[currStage].uvOffset = (shapeBatchIdx * vNumHead * totalTokens + vHeadIdx * totalTokens + tokenOffset + chunkIdx * chunkSize) * vHeadDim; + offsets[currStage].wkOffset = (shapeBatchIdx * kNumHead * totalTokens + kHeadIdx * totalTokens + tokenOffset + chunkIdx * chunkSize) * kHeadDim; + offsets[currStage].wOffset = (shapeBatchIdx * vNumHead * totalTokens + vHeadIdx * totalTokens + tokenOffset + chunkIdx * chunkSize) * kHeadDim; + offsets[currStage].gOffset = shapeBatchIdx * vNumHead * totalTokens + vHeadIdx * totalTokens + tokenOffset + chunkIdx * chunkSize; + offsets[currStage].hWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + currStage) * kHeadDim * vHeadDim; + offsets[currStage].vWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + currStage) * chunkSize * vHeadDim; + offsets[currStage].blockTokens = offsets[currStage].isFinalState ? (batchTokens - chunkIdx * chunkSize) : chunkSize; + offsets[currStage].isDummyHead = headInnerLoop < PING_PONG_STAGES && headInnerIdx >= headInnerLoop; + offsets[currStage].batchIdx = batchIdx; + offsets[currStage].headIdx = vHeadIdx; + offsets[currStage].chunkIdx = chunkIdx; + + processNewTask = chunkIdx == batchChunks - 1 && headInnerIdx == PING_PONG_STAGES - 1; + if (processNewTask) { + uint32_t currLoopIdx = taskIdx / (PING_PONG_STAGES * cubeCoreNum); + headInnerLoop = ((currLoopIdx + 2 == taskLoops) && hasDummyHead) ? 1 : PING_PONG_STAGES; + taskIdx = (currLoopIdx + 1) * PING_PONG_STAGES * cubeCoreNum + headInnerLoop * cubeCoreIdx; + } + } + + CATLASS_DEVICE + GDNFwdHOffsets& GetStage1Offsets() { + return offsets[currStage]; + } + + CATLASS_DEVICE + bool NeedProcessStage1() { + GDNFwdHOffsets& stage1Offsets = GetStage1Offsets(); + return !(lastLoop || stage1Offsets.isDummyHead); + } + + CATLASS_DEVICE + GDNFwdHOffsets& GetStage2Offsets() { + return offsets[(currStage - 1) % PING_PONG_STAGES]; + } + + CATLASS_DEVICE + bool NeedProcessStage2() { + GDNFwdHOffsets& stage2Offsets = GetStage2Offsets(); + return !(iterId == 1 || (!storeFinalState && stage2Offsets.isFinalState) || stage2Offsets.isDummyHead); + } +}; + +struct BlockSchedulerGdnFwdHCube : public BlockSchedulerGdnFwdH { + CATLASS_DEVICE + BlockSchedulerGdnFwdHCube() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user) { + BlockSchedulerGdnFwdH::Init(cu_seqlens, chunk_indices, tiling, user, AscendC::GetBlockIdx(), AscendC::GetBlockNum()); + } + +}; + +struct BlockSchedulerGdnFwdHVec : public BlockSchedulerGdnFwdH { + CATLASS_DEVICE + BlockSchedulerGdnFwdHVec() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user) { + BlockSchedulerGdnFwdH::Init(cu_seqlens, chunk_indices, tiling, user, AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(), AscendC::GetBlockNum()); + } + +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_SCHEDULER_GDN_FWD_H_HPP \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/gemm/kernel/gdn_fwd_h_kernel.hpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/gemm/kernel/gdn_fwd_h_kernel.hpp new file mode 100644 index 000000000..4288e533d --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch20/gemm/kernel/gdn_fwd_h_kernel.hpp @@ -0,0 +1,511 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#define CATLASS_ARCH 2201 +#define CATLASS_UNIFIED_CORE 1 + +#include "catlass/arch/arch.hpp" +#include "catlass/arch/cross_core_sync.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/catlass.hpp" +#include "catlass/epilogue/block/block_epilogue.hpp" +#include "../../epilogue/block/block_epilogue_gdn_fwdh_update.hpp" +#include "../../epilogue/block/block_epilogue_gdn_fwdh_vnew.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/gemm/block/block_swizzle.hpp" +#include "../block/block_scheduler_gdn_fwd_h.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" +#include "catlass/gemm_coord.hpp" +#include "tla/tensor.hpp" +#include "tla/layout.hpp" +#include "tla/tensor.hpp" + +using _0 = tla::Int<0>; +using _1 = tla::Int<1>; +using _2 = tla::Int<2>; +using _4 = tla::Int<4>; +using _8 = tla::Int<8>; +using _16 = tla::Int<16>; +using _32 = tla::Int<32>; +using _64 = tla::Int<64>; +using _128 = tla::Int<128>; +using _256 = tla::Int<256>; +using _512 = tla::Int<512>; +using _1024 = tla::Int<1024>; +using _2048 = tla::Int<2048>; +using _4096 = tla::Int<4096>; +using _8192 = tla::Int<8192>; +using _16384 = tla::Int<16384>; +using _32768 = tla::Int<32768>; +using _65536 = tla::Int<65536>; + + + + +#include "kernel_operator.h" +using namespace Catlass; +using namespace tla; + +namespace Catlass::Gemm::Kernel { + +template< + typename INPUT_TYPE, + typename G_TYPE, + typename STATE_TYPE, + typename WORKSPACE_TYPE +> +class GDNFwdHKernel { +public: + + using ArchTag = Arch::AtlasA2; + using CubeScheduler = typename Catlass::Gemm::Block::BlockSchedulerGdnFwdHCube; + using VecScheduler = typename Catlass::Gemm::Block::BlockSchedulerGdnFwdHVec; + + using DispatchPolicyTla = Gemm::MmadPingpongTlaMulti; + using L1TileShapeTla = Shape<_128, _128, _128>; + using L0TileShapeTla = L1TileShapeTla; + + using WType = Gemm::GemmType; + using HType = Gemm::GemmType; + using VworkType = Gemm::GemmType; + using KType = Gemm::GemmType; + using HworkType = Gemm::GemmType; + using VType = Gemm::GemmType; + using GType = Gemm::GemmType; + using UType = Gemm::GemmType; + using FinalStateType = Gemm::GemmType; + + // cube 1 + using TileCopyWH = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadWH = Gemm::Block::BlockMmadTla; + + // cube 2 + using TileCopyKV = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadKV = Gemm::Block::BlockMmadTla; + + // vec 1 + using DispatchPolicyGDNFwdHVnew = Epilogue::EpilogueAtlasGDNFwdHVnew; + using EpilogueGDNFwdHVnew = Epilogue::Block::BlockEpilogue; + + // vec 2 + using DispatchPolicyGDNFwdHUpdate = Epilogue::EpilogueAtlasGDNFwdHUpdate; + using EpilogueGDNFwdHUpdate = Epilogue::Block::BlockEpilogue; + + using GDNFwdHOffsets = Catlass::Gemm::Block::GDNFwdHOffsets; + + using ElementK = INPUT_TYPE; + using ElementW = INPUT_TYPE; + using ElementU = INPUT_TYPE; + using ElementG = G_TYPE; + using ElementH = INPUT_TYPE; + using ElementV = INPUT_TYPE; + using ElementVWork = WORKSPACE_TYPE; + using ElementHWork = WORKSPACE_TYPE; + using ElementInitialState = STATE_TYPE; + using ElementFinalState = STATE_TYPE; + + using LayoutW = Catlass::layout::RowMajor; + using LayoutH = Catlass::layout::RowMajor; + using LayoutV = Catlass::layout::RowMajor; + using LayoutK = Catlass::layout::ColumnMajor; + + + uint32_t batch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + uint32_t initalStateStride0; + bool useInitialState; + bool storeFinalState; + uint32_t isVariedLen; + uint32_t shapeBatch; + uint32_t tokenBatch; + uint32_t vWorkspaceOffset; + uint32_t vUpdateWorkspaceOffset; + uint32_t hWorkspaceOffset; + uint32_t numSeqWorkspaceOffset; + uint32_t numChunksWorkspaceOffset; + + AscendC::GlobalTensor gmK; + AscendC::GlobalTensor gmW; + AscendC::GlobalTensor gmU; + AscendC::GlobalTensor gmG; + AscendC::GlobalTensor gmInitialState; + AscendC::GlobalTensor gmH; + AscendC::GlobalTensor gmV; + AscendC::GlobalTensor gmFinalState; + AscendC::GlobalTensor gmVWorkspace; + AscendC::GlobalTensor gmVUpdateWorkspace; + AscendC::GlobalTensor gmHWorkspace; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmNumSeq; + AscendC::GlobalTensor gmNumChunks; + + CubeScheduler cubeBlockScheduler; + VecScheduler vecBlockScheduler; + + Arch::Resource resource; + + + __aicore__ inline GDNFwdHKernel() {} + + __aicore__ inline void Init(GM_ADDR k, GM_ADDR w, GM_ADDR u, GM_ADDR g, GM_ADDR inital_state, GM_ADDR cu_seqlens, GM_ADDR chunk_indices, + GM_ADDR h, GM_ADDR v_new, GM_ADDR final_state, GM_ADDR tiling, GM_ADDR user) { + + __gm__ ChunkGatedDeltaRuleFwdHTilingData *__restrict gdnFwdHTilingData = reinterpret_cast<__gm__ ChunkGatedDeltaRuleFwdHTilingData *__restrict>(tiling); + + batch = gdnFwdHTilingData->batch; + seqlen = gdnFwdHTilingData->seqlen; + kNumHead = gdnFwdHTilingData->kNumHead; + vNumHead = gdnFwdHTilingData->vNumHead; + kHeadDim = gdnFwdHTilingData->kHeadDim; + vHeadDim = gdnFwdHTilingData->vHeadDim; + chunkSize = gdnFwdHTilingData->chunkSize; + initalStateStride0 = gdnFwdHTilingData->initalStateStride0; + useInitialState = gdnFwdHTilingData->useInitialState; + storeFinalState = gdnFwdHTilingData->storeFinalState; + isVariedLen = gdnFwdHTilingData->isVariedLen; + shapeBatch = gdnFwdHTilingData->shapeBatch; + tokenBatch = gdnFwdHTilingData->tokenBatch; + vWorkspaceOffset = gdnFwdHTilingData->vWorkspaceOffset; + vUpdateWorkspaceOffset = gdnFwdHTilingData->vUpdateWorkspaceOffset; + hWorkspaceOffset = gdnFwdHTilingData->hWorkspaceOffset; + numSeqWorkspaceOffset = gdnFwdHTilingData->numSeqWorkspaceOffset; + numChunksWorkspaceOffset = gdnFwdHTilingData->numChunksWorkspaceOffset; + + gmK.SetGlobalBuffer((__gm__ ElementK *)k); + gmW.SetGlobalBuffer((__gm__ ElementW *)w); + gmU.SetGlobalBuffer((__gm__ ElementU *)u); + gmG.SetGlobalBuffer((__gm__ ElementG *)g); + gmInitialState.SetGlobalBuffer((__gm__ ElementInitialState *)inital_state); + gmH.SetGlobalBuffer((__gm__ ElementH *)h); + gmV.SetGlobalBuffer((__gm__ ElementV *)v_new); + gmFinalState.SetGlobalBuffer((__gm__ ElementFinalState *)final_state); + gmVWorkspace.SetGlobalBuffer((__gm__ ElementVWork *)(user + vWorkspaceOffset)); + gmVUpdateWorkspace.SetGlobalBuffer((__gm__ ElementV *)(user + vUpdateWorkspaceOffset)); + gmHWorkspace.SetGlobalBuffer((__gm__ ElementHWork *)(user + hWorkspaceOffset)); + + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + cubeBlockScheduler.Init(cu_seqlens, chunk_indices, tiling, user); + } + + __aicore__ inline void Process() { + ProcessUnifiedCore(); + } + + __aicore__ inline void ProcessUnifiedCore() { + uint32_t coreNum = AscendC::GetBlockNum(); + + BlockMmadWH blockMmadWH(resource); + BlockMmadKV blockMmadKV(resource); + EpilogueGDNFwdHVnew epilogueGDNFwdHVnew(resource); + + auto wLayout = tla::MakeLayout(shapeBatch * kNumHead * cubeBlockScheduler.totalTokens, kHeadDim); + auto hLayout = tla::MakeLayout(shapeBatch * vNumHead * cubeBlockScheduler.totalChunks * kHeadDim, vHeadDim); + auto vLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, vHeadDim); + auto kLayout = tla::MakeLayout(kHeadDim, shapeBatch * kNumHead * cubeBlockScheduler.totalTokens); + auto vworkLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, vHeadDim); + auto hworkLayout = tla::MakeLayout(coreNum * kHeadDim * PING_PONG_STAGES, vHeadDim); + + if (useInitialState) { + AscendC::LocalTensor stateUbTensorPing = resource.ubBuf.template GetBufferByByte(0); + AscendC::LocalTensor stateUbTensorPong = resource.ubBuf.template GetBufferByByte(96 * 1024); + AscendC::LocalTensor hUbTensorPing = resource.ubBuf.template GetBufferByByte(64 * 1024); + AscendC::LocalTensor hUbTensorPong = resource.ubBuf.template GetBufferByByte(160 * 1024); + uint32_t totalChunks = isVariedLen ? cubeBlockScheduler.totalChunks : ((seqlen + chunkSize - 1) / chunkSize); + uint32_t stateBlockSize = kHeadDim * vHeadDim; + uint32_t pingpongFlag = 1; + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + for (uint32_t shapeBatchIdx = 0; shapeBatchIdx < shapeBatch; shapeBatchIdx++) { + for (uint32_t vHeadIdx = 0; vHeadIdx < vNumHead; vHeadIdx++) { + for (uint32_t tokenBatchIdx = 0; tokenBatchIdx < cubeBlockScheduler.tokenBatch; tokenBatchIdx++) { + uint32_t batchIdx = isVariedLen ? tokenBatchIdx : shapeBatchIdx; + uint32_t chunkOffset = isVariedLen ? gmNumChunks.GetValue(tokenBatchIdx) : 0; + uint32_t initialStateOffset = (batchIdx * vNumHead + vHeadIdx) * stateBlockSize; + uint32_t hOffset = (shapeBatchIdx * vNumHead * totalChunks + vHeadIdx * totalChunks + chunkOffset) * stateBlockSize; + AscendC::LocalTensor stateUbTensor = pingpongFlag ? stateUbTensorPing : stateUbTensorPong; + AscendC::LocalTensor hUbTensor = pingpongFlag ? hUbTensorPing : hUbTensorPong; + auto event_id = pingpongFlag ? EVENT_ID1 : EVENT_ID0; + AscendC::WaitFlag(event_id); + if constexpr(!std::is_same::value) { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateOffset], stateBlockSize); + AscendC::SetFlag(event_id); + AscendC::WaitFlag(event_id); + AscendC::Cast(hUbTensor, stateUbTensor, AscendC::RoundMode::CAST_NONE, stateBlockSize); + AscendC::SetFlag(event_id); + AscendC::WaitFlag(event_id); + AscendC::DataCopy(gmH[hOffset], hUbTensor, stateBlockSize); + } else { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateOffset], stateBlockSize); + AscendC::SetFlag(event_id); + AscendC::WaitFlag(event_id); + AscendC::DataCopy(gmH[hOffset], stateUbTensor, stateBlockSize); + } + AscendC::SetFlag(event_id); + pingpongFlag = 1 - pingpongFlag; + } + } + } + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + } + + while (cubeBlockScheduler.isRunning) { + cubeBlockScheduler.InitTask(); + GDNFwdHOffsets& stage1Offsets = cubeBlockScheduler.GetStage1Offsets(); + + // CUBE1: v_work = w @ h[i] + if (cubeBlockScheduler.NeedProcessStage1()) { + auto tensorW = tla::MakeTensor(gmW[stage1Offsets.wOffset], wLayout, Catlass::Arch::PositionGM{}); + auto tensorH = tla::MakeTensor(gmH[stage1Offsets.hSrcOffset], hLayout, Catlass::Arch::PositionGM{}); + auto tensorV = tla::MakeTensor(gmVWorkspace[stage1Offsets.vWorkOffset], vLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape{stage1Offsets.blockTokens, vHeadDim, kHeadDim}; + auto tensorBlockW = GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + auto tensorBlockV = GetTile(tensorV, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.n())); + blockMmadWH.preSetFlags(); + blockMmadWH(tensorBlockW, tensorBlockH, tensorBlockV, cube1Shape); + blockMmadWH.finalWaitFlags(); + } + + // VEC1: v_new epilogue + if (cubeBlockScheduler.NeedProcessStage1()) { + epilogueGDNFwdHVnew( + gmV[stage1Offsets.uvOffset], gmVUpdateWorkspace[stage1Offsets.vWorkOffset], + gmG[stage1Offsets.gOffset], gmU[stage1Offsets.uvOffset], gmVWorkspace[stage1Offsets.vWorkOffset], + stage1Offsets.blockTokens, kHeadDim, vHeadDim, cubeBlockScheduler.cube1Done + ); + } + + if (cubeBlockScheduler.iterId > 1) { + GDNFwdHOffsets& stage2Offsets = cubeBlockScheduler.GetStage2Offsets(); + + // CUBE2: h_work = k.T @ v_update + // BlockMmadTla has no outer M loop; m must be split when kHeadDim > L1_TILE_M. + if (cubeBlockScheduler.NeedProcessStage2()) { + auto tensorK = tla::MakeTensor(gmK[stage2Offsets.wkOffset], kLayout, Catlass::Arch::PositionGM{}); + auto tensorVwork = tla::MakeTensor(gmVUpdateWorkspace[stage2Offsets.vWorkOffset], vworkLayout, Catlass::Arch::PositionGM{}); + auto tensorHwork = tla::MakeTensor(gmHWorkspace[stage2Offsets.hWorkOffset], hworkLayout, Catlass::Arch::PositionGM{}); + constexpr uint32_t L1_TILE_M_C2 = tla::get<0>(L1TileShapeTla{}); + uint32_t mLoopC2 = (kHeadDim + L1_TILE_M_C2 - 1) / L1_TILE_M_C2; + for (uint32_t mIdx = 0; mIdx < mLoopC2; ++mIdx) { + uint32_t mOff = mIdx * L1_TILE_M_C2; + uint32_t mTail = kHeadDim - mOff; + uint32_t mActual = (mTail < L1_TILE_M_C2) ? mTail : L1_TILE_M_C2; + GemmCoord cube2Shape{mActual, vHeadDim, stage2Offsets.blockTokens}; + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(mOff, 0), tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockVwork = GetTile(tensorVwork, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + auto tensorBlockHwork = GetTile(tensorHwork, tla::MakeCoord(mOff, 0), tla::MakeShape(cube2Shape.m(), cube2Shape.n())); + blockMmadKV.preSetFlags(); + blockMmadKV(tensorBlockK, tensorBlockVwork, tensorBlockHwork, cube2Shape); + blockMmadKV.finalWaitFlags(); + } + } + + // VEC2: h update epilogue + if (cubeBlockScheduler.NeedProcessStage2()) { + EpilogueGDNFwdHUpdate epilogueGDNFwdHUpdate(resource); + epilogueGDNFwdHUpdate( + gmH[stage2Offsets.hDstOffset], gmFinalState[stage2Offsets.finalStateOffset], + gmG[stage2Offsets.gOffset], gmH[stage2Offsets.hSrcOffset], + gmHWorkspace[stage2Offsets.hWorkOffset], + stage2Offsets.blockTokens, kHeadDim, vHeadDim, cubeBlockScheduler.cube2Done, + (stage2Offsets.isFinalState && storeFinalState) + ); + } + } + } + } + + __aicore__ inline void ProcessSplitCore() { + + if ASCEND_IS_AIC { + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = AscendC::GetBlockNum(); + + BlockMmadWH blockMmadWH(resource); + BlockMmadKV blockMmadKV(resource); + + auto wLayout = tla::MakeLayout(shapeBatch * kNumHead * cubeBlockScheduler.totalTokens, kHeadDim); + auto hLayout = tla::MakeLayout(shapeBatch * vNumHead * cubeBlockScheduler.totalChunks * kHeadDim, vHeadDim); + auto vLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, vHeadDim); + + auto kLayout = tla::MakeLayout(kHeadDim, shapeBatch * kNumHead * cubeBlockScheduler.totalTokens); + auto vworkLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, vHeadDim); + auto hworkLayout = tla::MakeLayout(coreNum * kHeadDim * PING_PONG_STAGES, vHeadDim); + + while (cubeBlockScheduler.isRunning) { + cubeBlockScheduler.InitTask(); + // step 1: v_work = w @ h[i] + GDNFwdHOffsets& cube1Offsets = cubeBlockScheduler.GetStage1Offsets(); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done); + if (cubeBlockScheduler.NeedProcessStage1()) { + int64_t cube1OffsetW = cube1Offsets.wOffset; + int64_t cube1OffsetH = cube1Offsets.hSrcOffset; + int64_t cube1OffsetVwork = cube1Offsets.vWorkOffset; + auto tensorW = tla::MakeTensor(gmW[cube1OffsetW], wLayout, Catlass::Arch::PositionGM{}); + auto tensorH = tla::MakeTensor(gmH[cube1OffsetH], hLayout, Catlass::Arch::PositionGM{}); + auto tensorV = tla::MakeTensor(gmVWorkspace[cube1OffsetVwork], vLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape {cube1Offsets.blockTokens, vHeadDim, kHeadDim}; + auto tensorBlockW = GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + auto tensorBlockV = GetTile(tensorV, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.n())); + blockMmadWH.preSetFlags(); + blockMmadWH(tensorBlockW, tensorBlockH, tensorBlockV, cube1Shape); + blockMmadWH.finalWaitFlags(); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube1Done); + + if (cubeBlockScheduler.iterId > 1) { + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done); + GDNFwdHOffsets& cube2Offsets = cubeBlockScheduler.GetStage2Offsets(); + if (cubeBlockScheduler.NeedProcessStage2()) { + // step 3: h[i+1] = k.T @ v_work + // BlockMmadTla has no outer M loop; m must be split when kHeadDim > L1_TILE_M. + int64_t cube2OffsetK = cube2Offsets.wkOffset; + int64_t cube2OffsetVwork = cube2Offsets.vWorkOffset; + int64_t cube2OffsetH = cube2Offsets.hWorkOffset; + auto tensorK = tla::MakeTensor(gmK[cube2OffsetK], kLayout, Catlass::Arch::PositionGM{}); + auto tensorVwork = tla::MakeTensor(gmVUpdateWorkspace[cube2OffsetVwork], vworkLayout, Catlass::Arch::PositionGM{}); + auto tensorHwork = tla::MakeTensor(gmHWorkspace[cube2OffsetH], hworkLayout, Catlass::Arch::PositionGM{}); + constexpr uint32_t L1_TILE_M_C2 = tla::get<0>(L1TileShapeTla{}); + uint32_t mLoopC2 = (kHeadDim + L1_TILE_M_C2 - 1) / L1_TILE_M_C2; + for (uint32_t mIdx = 0; mIdx < mLoopC2; ++mIdx) { + uint32_t mOff = mIdx * L1_TILE_M_C2; + uint32_t mTail = kHeadDim - mOff; + uint32_t mActual = (mTail < L1_TILE_M_C2) ? mTail : L1_TILE_M_C2; + GemmCoord cube2Shape{mActual, vHeadDim, cube2Offsets.blockTokens}; + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(mOff, 0), tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockVwork = GetTile(tensorVwork, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + auto tensorBlockHwork = GetTile(tensorHwork, tla::MakeCoord(mOff, 0), tla::MakeShape(cube2Shape.m(), cube2Shape.n())); + blockMmadKV.preSetFlags(); + blockMmadKV(tensorBlockK, tensorBlockVwork, tensorBlockHwork, cube2Shape); + blockMmadKV.finalWaitFlags(); + } + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube2Done); + } + } + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done); + + } + + if ASCEND_IS_AIV { + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = AscendC::GetBlockNum(); + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + EpilogueGDNFwdHVnew epilogueGDNFwdHVnew(resource); + + if (useInitialState) { + AscendC::LocalTensor stateUbTensorPing = resource.ubBuf.template GetBufferByByte(0); + AscendC::LocalTensor stateUbTensorPong = resource.ubBuf.template GetBufferByByte(96 * 1024); + AscendC::LocalTensor hUbTensorPing = resource.ubBuf.template GetBufferByByte(64 * 1024); + AscendC::LocalTensor hUbTensorPong = resource.ubBuf.template GetBufferByByte(160 * 1024); + uint32_t totalChunks = isVariedLen ? vecBlockScheduler.totalChunks : ((seqlen + chunkSize - 1) / chunkSize); + uint32_t stateBlockSize = kHeadDim * vHeadDim; + uint32_t pingpongFlag = 1; + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::DataCopyParams repeatParams = {static_cast(kHeadDim), static_cast(vHeadDim * sizeof(ElementInitialState) / 32), + static_cast((initalStateStride0 - vHeadDim)* sizeof(ElementInitialState) / 32), static_cast(0)}; + for (uint32_t shapeBatchIdx = 0; shapeBatchIdx < shapeBatch; shapeBatchIdx++) { + for (uint32_t vHeadIdx = 0; vHeadIdx < vNumHead; vHeadIdx++) { + for (uint32_t tokenBatchIdx = 0; tokenBatchIdx < vecBlockScheduler.tokenBatch; tokenBatchIdx++) { + uint32_t batchIdx = isVariedLen ? tokenBatchIdx : shapeBatchIdx; + uint32_t chunkOffset = isVariedLen ? gmNumChunks.GetValue(tokenBatchIdx) : 0; + uint32_t initialStateSrcOffset = (batchIdx * vNumHead + vHeadIdx) * kHeadDim * initalStateStride0; + uint32_t hOffset = (shapeBatchIdx * vNumHead * totalChunks + vHeadIdx * totalChunks + chunkOffset) * stateBlockSize; + AscendC::LocalTensor stateUbTensor = pingpongFlag ? stateUbTensorPing : stateUbTensorPong; + AscendC::LocalTensor hUbTensor = pingpongFlag ? hUbTensorPing : hUbTensorPong; + auto event_id = pingpongFlag ? EVENT_ID1 : EVENT_ID0; + AscendC::WaitFlag(event_id); + if constexpr(!std::is_same::value) { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateSrcOffset], repeatParams); + AscendC::SetFlag(event_id); + AscendC::WaitFlag(event_id); + AscendC::Cast(hUbTensor, stateUbTensor, AscendC::RoundMode::CAST_RINT, stateBlockSize); + AscendC::SetFlag(event_id); + AscendC::WaitFlag(event_id); + AscendC::DataCopy(gmH[hOffset], hUbTensor, stateBlockSize); + } else { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateSrcOffset], repeatParams); + AscendC::SetFlag(event_id); + AscendC::WaitFlag(event_id); + AscendC::DataCopy(gmH[hOffset], stateUbTensor, stateBlockSize); + } + AscendC::SetFlag(event_id); + pingpongFlag = 1 - pingpongFlag; + } + + } + } + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + } + + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done); + while (vecBlockScheduler.isRunning) { + vecBlockScheduler.InitTask(); + // step 2: + GDNFwdHOffsets& vec1Offsets = vecBlockScheduler.GetStage1Offsets(); + // gmV = gmU - gmVWorkspace + // g_buf = gmG[-1] - gmG + // g_buf = exp(g_buf) + // gmVWorkspace = g_buf * gmV + if (vecBlockScheduler.NeedProcessStage1()) { + epilogueGDNFwdHVnew( + gmV[vec1Offsets.uvOffset], gmVUpdateWorkspace[vec1Offsets.vWorkOffset], + gmG[vec1Offsets.gOffset], gmU[vec1Offsets.uvOffset], gmVWorkspace[vec1Offsets.vWorkOffset], + vec1Offsets.blockTokens, kHeadDim, vHeadDim, vecBlockScheduler.cube1Done + ); + } else { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube1Done); + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec1Done); + + if (vecBlockScheduler.iterId > 1) { + GDNFwdHOffsets& vec2Offsets = vecBlockScheduler.GetStage2Offsets(); + if (vecBlockScheduler.NeedProcessStage2()) { + // step 4: h[i+1] += h_work if i < num_chunks - 1 else None + EpilogueGDNFwdHUpdate epilogueGDNFwdHUpdate(resource); + epilogueGDNFwdHUpdate( + gmH[vec2Offsets.hDstOffset], gmFinalState[vec2Offsets.finalStateOffset], + gmG[vec2Offsets.gOffset], + gmH[vec2Offsets.hSrcOffset], + gmHWorkspace[vec2Offsets.hWorkOffset], + vec2Offsets.blockTokens, kHeadDim, vHeadDim, vecBlockScheduler.cube2Done, + (vec2Offsets.isFinalState && storeFinalState) + ); + } else { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube2Done); + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done); + } + } + + } + } + +}; + +} \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/epilogue/block/block_epilogue_gdn_fwdh_update.hpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/epilogue/block/block_epilogue_gdn_fwdh_update.hpp new file mode 100644 index 000000000..0fb901d7d --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/epilogue/block/block_epilogue_gdn_fwdh_update.hpp @@ -0,0 +1,180 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDH_UPDATE_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDH_UPDATE_HPP +#include "catlass/catlass.hpp" +#include "catlass/arch/resource.hpp" +#include "../gdn_fwd_h_epilogue_policies.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/matrix_coord.hpp" +#include "catlass/epilogue/tile/tile_copy.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class HOutputType_, + class GInputType_, + class HInputType_, + class HUpdateInputType_, + class FinalStateType_ +> +class BlockEpilogue < + EpilogueAtlasGDNFwdHUpdate, + HOutputType_, + GInputType_, + HInputType_, + HUpdateInputType_, + FinalStateType_ +> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasGDNFwdHUpdate; + using ArchTag = typename DispatchPolicy::ArchTag; + + using HElementOutput = typename HOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using HElementInput = typename HInputType_::Element; + using HUpdateElementInput = typename HUpdateInputType_::Element; + using FinalStateElement = typename FinalStateType_::Element; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + + constexpr uint32_t CALC_BUF_OFFSET = 0; + constexpr uint32_t PING_BUF_0_OFFSET = 32 * 1024; + constexpr uint32_t PING_BUF_1_OFFSET = 64 * 1024; + constexpr uint32_t PING_BUF_2_OFFSET = 80 * 1024; + constexpr uint32_t PING_G_BUF_OFFSET = 160 * 1024; + + + calcUbTensor = resource.ubBuf.template GetBufferByByte(CALC_BUF_OFFSET); + + hUpdateUbTensor = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + hUbTensor = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + + hOutputUbTensor = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + finalOutputUbTensor = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + + glastUbTensor = resource.ubBuf.template GetBufferByByte(PING_G_BUF_OFFSET); + + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor hOutput, + AscendC::GlobalTensor finalState, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor hInput, + AscendC::GlobalTensor hUpdateInput, + uint32_t chunkSize, + uint32_t kHeadDim, + uint32_t vHeadDim, + Arch::CrossCoreFlag cube2Done, + bool isFinalState + ) + { + uint32_t mActual = kHeadDim; + uint32_t nActual = vHeadDim; + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t mActualPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t mActualThisSubBlock = (subBlockIdx == 0) ? mActualPerSubBlock : (mActual - mActualPerSubBlock); + uint32_t mOffset = subBlockIdx * mActualPerSubBlock; + uint32_t nOffset = 0; + int64_t offsetH = mOffset * nActual + nOffset; + + AscendC::ResetMask(); + + AscendC::GlobalTensor hOutputThisSubBlock = hOutput[offsetH]; + AscendC::GlobalTensor gInputThisSubBlock = gInput; + AscendC::GlobalTensor hInputThisSubBlock = hInput[offsetH]; + AscendC::GlobalTensor hUpdateInputThisSubBlock = hUpdateInput[offsetH]; + AscendC::GlobalTensor finalStateThisSubBlock = finalState[offsetH]; + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopy(hUbTensor, hInputThisSubBlock, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::Cast(calcUbTensor, hUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + + GElementInput gLastVal = gInputThisSubBlock.GetValue(chunkSize-1); + float gLastFloat = 0.0f; + if constexpr(std::is_same::value) { + gLastFloat = gLastVal; + } else if constexpr(std::is_same::value) { + gLastFloat = (float)gLastVal; + } else if constexpr(std::is_same::value) { + gLastFloat = AscendC::ToFloat(gLastVal); + } + glastUbTensor.SetValue(0, gLastFloat); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::Exp(glastUbTensor, glastUbTensor, 1); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + float muls = glastUbTensor.GetValue(0); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::Muls(calcUbTensor, calcUbTensor, muls, mActualThisSubBlock * nActual); + + Arch::CrossCoreWaitFlag(cube2Done); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + AscendC::DataCopy(hUpdateUbTensor, hUpdateInputThisSubBlock, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + AscendC::Add(hUpdateUbTensor, calcUbTensor, hUpdateUbTensor, mActualThisSubBlock * nActual); + + if (isFinalState) { + if constexpr(!std::is_same::value) { + AscendC::PipeBarrier(); + AscendC::Cast(finalOutputUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopy(finalStateThisSubBlock, finalOutputUbTensor, mActualThisSubBlock * nActual); + } else { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopy(finalStateThisSubBlock, hUpdateUbTensor, mActualThisSubBlock * nActual); + } + } else { + AscendC::PipeBarrier(); + AscendC::Cast(hOutputUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopy(hOutputThisSubBlock, hOutputUbTensor, mActualThisSubBlock * nActual); + } + } + +private: + AscendC::LocalTensor calcUbTensor; + + AscendC::LocalTensor hUbTensor; + AscendC::LocalTensor hUpdateUbTensor; + + AscendC::LocalTensor hOutputUbTensor; + AscendC::LocalTensor finalOutputUbTensor; + + AscendC::LocalTensor glastUbTensor; + +}; +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/epilogue/block/block_epilogue_gdn_fwdh_vnew.hpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/epilogue/block/block_epilogue_gdn_fwdh_vnew.hpp new file mode 100644 index 000000000..9f9349225 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/epilogue/block/block_epilogue_gdn_fwdh_vnew.hpp @@ -0,0 +1,268 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDH_VNEW_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_GDN_FWDH_VNEW_HPP +#include "catlass/catlass.hpp" +#include "catlass/arch/resource.hpp" +#include "../gdn_fwd_h_epilogue_policies.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/matrix_coord.hpp" +#include "catlass/epilogue/tile/tile_copy.hpp" + + + +namespace Catlass::Epilogue::Block { + +template < + class VOutputType_, + class GInputType_, + class UInputType_, + class WSInputType_ +> +class BlockEpilogue < + EpilogueAtlasGDNFwdHVnew, + VOutputType_, + GInputType_, + UInputType_, + WSInputType_ +> { +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasGDNFwdHVnew; + using ArchTag = typename DispatchPolicy::ArchTag; + + using VElementOutput = typename VOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using UElementInput = typename UInputType_::Element; + using WSElementInput = typename WSInputType_::Element; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + + constexpr uint32_t CALC_BUF_OFFSET = 0; + constexpr uint32_t PING_BUF_0_OFFSET = 32 * 1024; + constexpr uint32_t PING_BUF_1_OFFSET = 64 * 1024; + constexpr uint32_t PONG_BUF_0_OFFSET = 96 * 1024; + constexpr uint32_t PONG_BUF_1_OFFSET = 128 * 1024; + constexpr uint32_t PING_G_BUF_OFFSET = 160 * 1024; + constexpr uint32_t PONG_G_BUF_OFFSET = 161 * 1024; + constexpr uint32_t PING_G_SUB_BUF_OFFSET = 162 * 1024; + constexpr uint32_t PONG_G_SUB_BUF_OFFSET = 163 * 1024; + constexpr uint32_t PING_G_INPUT_BUF_OFFSET = 164 * 1024; + constexpr uint32_t PONG_G_INPUT_BUF_OFFSET = 165 * 1024; + constexpr uint32_t SHARE_BUF_OFFSET = 166 * 1024; + + calcUbTensor = resource.ubBuf.template GetBufferByByte(CALC_BUF_OFFSET); + + uUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + uUbFloatTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + wsUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + gUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_BUF_OFFSET); + gLastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + gInputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_INPUT_BUF_OFFSET); + vNewOutputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + vNewDecayUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + + uUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_1_OFFSET); + uUbFloatTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_0_OFFSET); + wsUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_1_OFFSET); + gUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_BUF_OFFSET); + gLastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + gInputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_INPUT_BUF_OFFSET); + vNewOutputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_1_OFFSET); + vNewDecayUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_0_OFFSET); + + shareBuffer_ = resource.ubBuf.template GetBufferByByte(SHARE_BUF_OFFSET); + + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor vnewOutput, + AscendC::GlobalTensor vnewdecayOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor uInput, + AscendC::GlobalTensor wsInput, + uint32_t chunkSize, + uint32_t kHeadDim, + uint32_t vHeadDim, + Arch::CrossCoreFlag cube1Done + // const LayoutOutput &layoutOutput, + // const LayoutInput &LayoutInput + ) + { + uint32_t mActual = chunkSize; + uint32_t nkActual = kHeadDim; + uint32_t nvActual = vHeadDim; + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t mActualPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t mActualThisSubBlock = (subBlockIdx == 0) ? mActualPerSubBlock : (mActual - mActualPerSubBlock); + uint32_t mOffset = subBlockIdx * mActualPerSubBlock; + uint32_t nOffset = 0; + // 当前场景内部一定连续 + // k [B, H, T, D] + // g [B, H, T] + // 在外部offset的基础上进一步offset + // 当前asset kdim == vHeadDim + int64_t offsetK = mOffset * nvActual + nOffset; + int64_t offsetD = 0; // 因为要用最后一个数减去之前所有，所以全部读入 + + uint32_t gbrcStart, gbrcRealStart, gbrcReptime, gbrcEffStart, gbrcEffEnd; + if(subBlockIdx==0) + { + gbrcStart = 0; + gbrcRealStart = 0; + gbrcReptime = (mActualThisSubBlock + 8 - 1) / 8; + + } + else + { + gbrcStart = mActualPerSubBlock; + gbrcRealStart = gbrcStart & ~15; + gbrcReptime = (mActual - gbrcRealStart + 8 - 1) / 8; + } + gbrcEffStart = gbrcStart-gbrcRealStart; + gbrcEffEnd = gbrcEffStart + mActualThisSubBlock; + + AscendC::ResetMask(); + + AscendC::GlobalTensor vnewOutputThisSubBlock = vnewOutput[offsetK]; + AscendC::GlobalTensor vnewdecayOutputThisSubBlock = vnewdecayOutput[offsetK]; + AscendC::GlobalTensor gInputThisSubBlock = gInput; + AscendC::GlobalTensor uInputThisSubBlock = uInput[offsetK]; + AscendC::GlobalTensor wsInputThisSubBlock = wsInput[offsetK]; + + pingpongFlag = isFirst ? 0 : 4; + AscendC::LocalTensor uUbTensor = isFirst ? uUbTensor_ping : uUbTensor_pong; + AscendC::LocalTensor uUbFloatTensor = isFirst ? uUbFloatTensor_ping : uUbFloatTensor_pong; + AscendC::LocalTensor wsUbTensor = isFirst ? wsUbTensor_ping : wsUbTensor_pong; + AscendC::LocalTensor gUbTensor = isFirst ? gUbTensor_ping : gUbTensor_pong; + AscendC::LocalTensor gLastUbTensor = isFirst ? gLastUbTensor_ping : gLastUbTensor_pong; + AscendC::LocalTensor gInputUbTensor = isFirst ? gInputUbTensor_ping : gInputUbTensor_pong; + AscendC::LocalTensor vNewOutputUbTensor = isFirst ? vNewOutputUbTensor_ping : vNewOutputUbTensor_pong; + AscendC::LocalTensor vNewDecayUbTensor = isFirst ? vNewDecayUbTensor_ping : vNewDecayUbTensor_pong; + + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + if constexpr(std::is_same::value) { + AscendC::DataCopyParams gUbParams{1, (uint16_t)(mActual * sizeof(float)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyPad(gUbTensor, gInputThisSubBlock, gUbParams, gUbPadParams); + } else { + AscendC::DataCopyParams gUbParams{1, (uint16_t)(mActual * sizeof(half)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyPad(gInputUbTensor, gInputThisSubBlock, gUbParams, gUbPadParams); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + if constexpr(!std::is_same::value) { + AscendC::Cast(gUbTensor, gInputUbTensor, AscendC::RoundMode::CAST_NONE, mActual); + } + + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + float inputVal = gUbTensor.GetValue(mActual-1); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + + AscendC::PipeBarrier(); + AscendC::Duplicate(gLastUbTensor, inputVal, mActual); + AscendC::PipeBarrier(); + + AscendC::Sub(gUbTensor, gLastUbTensor, gUbTensor, mActual); + AscendC::PipeBarrier(); + + AscendC::Exp(gUbTensor, gUbTensor, mActual); + AscendC::PipeBarrier(); + + uint32_t dstShape_[2] = {gbrcReptime*8, nvActual}; + uint32_t srcShape_[2] = {gbrcReptime*8, 1}; + AscendC::Broadcast(calcUbTensor, gUbTensor[gbrcRealStart], dstShape_, srcShape_, shareBuffer_); + AscendC::PipeBarrier(); + + Arch::CrossCoreWaitFlag(cube1Done); + + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + + AscendC::DataCopy(uUbTensor, uInputThisSubBlock, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::Cast(uUbFloatTensor, uUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(wsUbTensor, wsInputThisSubBlock, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + + AscendC::Sub(uUbFloatTensor, uUbFloatTensor, wsUbTensor, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + AscendC::Cast(vNewOutputUbTensor, uUbFloatTensor, AscendC::RoundMode::CAST_RINT, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(vnewOutputThisSubBlock, vNewOutputUbTensor, mActualThisSubBlock * nvActual); + + AscendC::PipeBarrier(); + AscendC::Mul(calcUbTensor[gbrcEffStart*nvActual], uUbFloatTensor, calcUbTensor[gbrcEffStart*nvActual], mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + AscendC::Cast(vNewDecayUbTensor, calcUbTensor[gbrcEffStart*nvActual], AscendC::RoundMode::CAST_RINT, mActualThisSubBlock * nvActual); + + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(vnewdecayOutputThisSubBlock, vNewDecayUbTensor, mActualThisSubBlock * nvActual); + + if (isFirst) { + AscendC::PipeBarrier(); + } + + isFirst = false; + } + +private: + uint32_t pingpongFlag = 0; + bool isFirst = true; + + AscendC::LocalTensor calcUbTensor; + + AscendC::LocalTensor uUbTensor_ping; + AscendC::LocalTensor uUbFloatTensor_ping; + AscendC::LocalTensor wsUbTensor_ping; + AscendC::LocalTensor gUbTensor_ping; + AscendC::LocalTensor gLastUbTensor_ping; + AscendC::LocalTensor gInputUbTensor_ping; + AscendC::LocalTensor vNewOutputUbTensor_ping; + AscendC::LocalTensor vNewDecayUbTensor_ping; + + AscendC::LocalTensor uUbTensor_pong; + AscendC::LocalTensor uUbFloatTensor_pong; + AscendC::LocalTensor wsUbTensor_pong; + AscendC::LocalTensor gUbTensor_pong; + AscendC::LocalTensor gLastUbTensor_pong; + AscendC::LocalTensor gInputUbTensor_pong; + AscendC::LocalTensor vNewOutputUbTensor_pong; + AscendC::LocalTensor vNewDecayUbTensor_pong; + + AscendC::LocalTensor shareBuffer_; + +}; +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/epilogue/gdn_fwd_h_epilogue_policies.hpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/epilogue/gdn_fwd_h_epilogue_policies.hpp new file mode 100644 index 000000000..eb5cc74c0 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/epilogue/gdn_fwd_h_epilogue_policies.hpp @@ -0,0 +1,35 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_EPILOGUE_GDN_FWD_H_EPILOGUE_POLICIES_HPP +#define CATLASS_EPILOGUE_GDN_FWD_H_EPILOGUE_POLICIES_HPP + +#include "catlass/catlass.hpp" + +namespace Catlass::Epilogue { + +struct EpilogueAtlasGDNFwdHVnew { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + using ArchTag = Arch::Ascend950; +#else + using ArchTag = Arch::AtlasA2; +#endif +}; + +struct EpilogueAtlasGDNFwdHUpdate { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + using ArchTag = Arch::Ascend950; +#else + using ArchTag = Arch::AtlasA2; +#endif +}; + +} // namespace Catlass::Epilogue + +#endif // CATLASS_EPILOGUE_GDN_FWD_H_EPILOGUE_POLICIES_HPP diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/gemm/block/block_scheduler_gdn_fwd_h.hpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/gemm/block/block_scheduler_gdn_fwd_h.hpp new file mode 100644 index 000000000..a83b08ef2 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/gemm/block/block_scheduler_gdn_fwd_h.hpp @@ -0,0 +1,286 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#include "catlass/gemm_coord.hpp" +using namespace Catlass; + +#ifndef CATLASS_GEMM_SCHEDULER_GDN_FWD_H_HPP +#define CATLASS_GEMM_SCHEDULER_GDN_FWD_H_HPP + +// constexpr uint32_t PING_PONG_STAGES = 1; +constexpr uint32_t PING_PONG_STAGES = 2; + +template +CATLASS_DEVICE T AlignUp(T a, T b) { + return (b == 0) ? 0 : (a + b - 1) / b * b; +} + +template +CATLASS_DEVICE T Min(T a, T b) { + return (a > b) ? b : a; +} + +template +CATLASS_DEVICE T Max(T a, T b) { + return (a > b) ? a : b; +} + +namespace Catlass::Gemm::Block { + +struct GDNFwdHOffsets { + uint32_t hSrcOffset; + uint32_t hDstOffset; + uint32_t uvOffset; + uint32_t wkOffset; + uint32_t wOffset; + uint32_t gOffset; + uint32_t hWorkOffset; + uint32_t vWorkOffset; + uint32_t initialStateOffset; + uint32_t finalStateOffset; + bool isInitialState; + bool isFinalState; + uint32_t blockTokens; + bool isDummyHead; + // for debug + uint32_t batchIdx; + uint32_t headIdx; + uint32_t chunkIdx; + +}; + +struct BlockSchedulerGdnFwdH { + uint32_t batch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + uint32_t initalStateStride0; + uint32_t vBlockSize{128}; + uint32_t isVariedLen; + uint32_t shapeBatch; + uint32_t tokenBatch; + bool useInitialState; + bool storeFinalState; + uint32_t numSeqWorkspaceOffset; + uint32_t numChunksWorkspaceOffset; + + uint32_t taskIdx; + uint32_t taskLoops; + uint32_t cubeCoreIdx; + uint32_t cubeCoreNum; + uint32_t vLoops; + uint32_t taskNum; + uint32_t headGroups; + uint32_t totalChunks; + uint32_t totalTokens; + uint32_t headInnerLoop; + + uint32_t iterId {0}; + bool hasDummyHead; + bool isRunning; + bool processNewTask {true}; + bool firstLoop {true}; + bool lastLoop {false}; + GDNFwdHOffsets offsets[PING_PONG_STAGES]; + int32_t currStage{PING_PONG_STAGES - 1}; + + uint32_t vIdx; + uint32_t batchIdx; + uint32_t baseHeadIdx; + uint32_t chunkIdx; + uint32_t headInnerIdx; + uint32_t vHeadIdx; + uint32_t kHeadIdx; + uint32_t shapeBatchIdx; + uint32_t tokenBatchIdx; + + uint32_t chunkOffset; + uint32_t tokenOffset; + uint32_t batchChunks; + uint32_t batchTokens; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmNumSeq; + AscendC::GlobalTensor gmNumChunks; + + Arch::CrossCoreFlag cube1Done{0}; + Arch::CrossCoreFlag vec1Done{1}; + Arch::CrossCoreFlag cube2Done{2}; + Arch::CrossCoreFlag vec2Done{3}; + + CATLASS_DEVICE + BlockSchedulerGdnFwdH() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user, uint32_t coreIdx, uint32_t coreNum) { + __gm__ ChunkGatedDeltaRuleFwdHTilingData *__restrict gdnFwdHTilingData = reinterpret_cast<__gm__ ChunkGatedDeltaRuleFwdHTilingData *__restrict>(tiling); + + batch = gdnFwdHTilingData->batch; + seqlen = gdnFwdHTilingData->seqlen; + kNumHead = gdnFwdHTilingData->kNumHead; + vNumHead = gdnFwdHTilingData->vNumHead; + kHeadDim = gdnFwdHTilingData->kHeadDim; + vHeadDim = gdnFwdHTilingData->vHeadDim; + chunkSize = gdnFwdHTilingData->chunkSize; + initalStateStride0 = gdnFwdHTilingData->initalStateStride0; + isVariedLen = gdnFwdHTilingData->isVariedLen; + shapeBatch = gdnFwdHTilingData->shapeBatch; + tokenBatch = gdnFwdHTilingData->tokenBatch; + useInitialState = gdnFwdHTilingData->useInitialState; + storeFinalState = gdnFwdHTilingData->storeFinalState; + numSeqWorkspaceOffset = gdnFwdHTilingData->numSeqWorkspaceOffset; + numChunksWorkspaceOffset = gdnFwdHTilingData->numChunksWorkspaceOffset; + + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + if (isVariedLen) { + gmNumChunks.SetValue(0, 0); + gmNumSeq.SetValue(0, 0); + uint32_t actualBatch = 0; + int64_t prevSeq = 0, currSeq; + for (uint32_t b = 1; b <= tokenBatch; b++) { + currSeq = gmSeqlen.GetValue(b); + int64_t batchSeqLen = currSeq - prevSeq; + if (batchSeqLen > 0) { + actualBatch++; + gmNumSeq.SetValue(actualBatch, currSeq); + int64_t batchChunk = (batchSeqLen + chunkSize - 1) / chunkSize; + gmNumChunks.SetValue(actualBatch, gmNumChunks.GetValue(actualBatch - 1) + batchChunk); + } + prevSeq = currSeq; + } + tokenBatch = actualBatch; + batch = actualBatch; + totalChunks = gmNumChunks.GetValue(tokenBatch); + totalTokens = gmNumSeq.GetValue(tokenBatch); + } else { + totalChunks = (seqlen + chunkSize - 1) / chunkSize; + totalTokens = seqlen; + } + + cubeCoreIdx = coreIdx; + cubeCoreNum = coreNum; + vLoops = vHeadDim / vBlockSize; + taskNum = vLoops * batch * vNumHead; + headGroups = vNumHead / kNumHead; + hasDummyHead = (taskNum % (PING_PONG_STAGES * cubeCoreNum) <= cubeCoreNum) && (taskNum % (PING_PONG_STAGES * cubeCoreNum) > 0); + taskLoops = (taskNum + cubeCoreNum * PING_PONG_STAGES - 1) / (cubeCoreNum * PING_PONG_STAGES); + headInnerLoop = taskNum > cubeCoreNum ? PING_PONG_STAGES : 1; + taskIdx = cubeCoreIdx * headInnerLoop; + isRunning = taskIdx < taskNum; + + } + + CATLASS_DEVICE + void InitTask() { + iterId++; + currStage = (currStage + 1) % PING_PONG_STAGES; + if (processNewTask) { + if (taskIdx >= taskNum) { + lastLoop = true; + isRunning = false; + return; + } + vIdx = taskIdx / (batch * vNumHead); + batchIdx = (taskIdx - vIdx * batch * vNumHead) / vNumHead; + baseHeadIdx = taskIdx % vNumHead; + shapeBatchIdx = isVariedLen ? 0 : batchIdx; + tokenBatchIdx = isVariedLen ? batchIdx : 0; + chunkOffset = isVariedLen ? gmNumChunks.GetValue(tokenBatchIdx) : 0; + batchChunks = isVariedLen ? (gmNumChunks.GetValue(tokenBatchIdx + 1) - chunkOffset) : totalChunks; + tokenOffset = isVariedLen ? gmNumSeq.GetValue(tokenBatchIdx) : 0; + batchTokens = isVariedLen ? (gmNumSeq.GetValue(tokenBatchIdx + 1) - tokenOffset) : totalTokens; + chunkIdx = 0; + headInnerIdx = 0; + } else { + chunkIdx = headInnerIdx == PING_PONG_STAGES - 1 ? chunkIdx + 1 : chunkIdx; + headInnerIdx = (headInnerIdx + 1) % PING_PONG_STAGES; + } + + vHeadIdx = baseHeadIdx + headInnerIdx; + kHeadIdx = vHeadIdx / headGroups; + offsets[currStage].isInitialState = chunkIdx == 0; + offsets[currStage].isFinalState = chunkIdx == (batchChunks - 1); + offsets[currStage].initialStateOffset = (batchIdx * vNumHead + vHeadIdx) * kHeadDim * initalStateStride0; + offsets[currStage].finalStateOffset = (batchIdx * vNumHead + vHeadIdx) * kHeadDim * vHeadDim; + offsets[currStage].hSrcOffset = (shapeBatchIdx * vNumHead * totalChunks + vHeadIdx * totalChunks + chunkOffset + chunkIdx) * kHeadDim * vHeadDim; + offsets[currStage].hDstOffset = offsets[currStage].hSrcOffset + kHeadDim * vHeadDim; + offsets[currStage].uvOffset = (shapeBatchIdx * vNumHead * totalTokens + vHeadIdx * totalTokens + tokenOffset + chunkIdx * chunkSize) * vHeadDim; + offsets[currStage].wkOffset = (shapeBatchIdx * kNumHead * totalTokens + kHeadIdx * totalTokens + tokenOffset + chunkIdx * chunkSize) * kHeadDim; + offsets[currStage].wOffset = (shapeBatchIdx * vNumHead * totalTokens + vHeadIdx * totalTokens + tokenOffset + chunkIdx * chunkSize) * kHeadDim; + offsets[currStage].gOffset = shapeBatchIdx * vNumHead * totalTokens + vHeadIdx * totalTokens + tokenOffset + chunkIdx * chunkSize; + offsets[currStage].hWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + currStage) * kHeadDim * vHeadDim; + offsets[currStage].vWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + currStage) * chunkSize * vHeadDim; + offsets[currStage].blockTokens = offsets[currStage].isFinalState ? (batchTokens - chunkIdx * chunkSize) : chunkSize; + offsets[currStage].isDummyHead = headInnerLoop < PING_PONG_STAGES && headInnerIdx >= headInnerLoop; + offsets[currStage].batchIdx = batchIdx; + offsets[currStage].headIdx = vHeadIdx; + offsets[currStage].chunkIdx = chunkIdx; + + processNewTask = chunkIdx == batchChunks - 1 && headInnerIdx == PING_PONG_STAGES - 1; + if (processNewTask) { + uint32_t currLoopIdx = taskIdx / (PING_PONG_STAGES * cubeCoreNum); + headInnerLoop = ((currLoopIdx + 2 == taskLoops) && hasDummyHead) ? 1 : PING_PONG_STAGES; + taskIdx = (currLoopIdx + 1) * PING_PONG_STAGES * cubeCoreNum + headInnerLoop * cubeCoreIdx; + } + } + + CATLASS_DEVICE + GDNFwdHOffsets& GetStage1Offsets() { + return offsets[currStage]; + } + + CATLASS_DEVICE + bool NeedProcessStage1() { + GDNFwdHOffsets& stage1Offsets = GetStage1Offsets(); + return !(lastLoop || stage1Offsets.isDummyHead); + } + + CATLASS_DEVICE + GDNFwdHOffsets& GetStage2Offsets() { + return offsets[(currStage - 1) % PING_PONG_STAGES]; + } + + CATLASS_DEVICE + bool NeedProcessStage2() { + GDNFwdHOffsets& stage2Offsets = GetStage2Offsets(); + return !(iterId == 1 || (!storeFinalState && stage2Offsets.isFinalState) || stage2Offsets.isDummyHead); + } +}; + +struct BlockSchedulerGdnFwdHCube : public BlockSchedulerGdnFwdH { + CATLASS_DEVICE + BlockSchedulerGdnFwdHCube() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user) { + BlockSchedulerGdnFwdH::Init(cu_seqlens, chunk_indices, tiling, user, AscendC::GetBlockIdx(), AscendC::GetBlockNum()); + } + +}; + +struct BlockSchedulerGdnFwdHVec : public BlockSchedulerGdnFwdH { + CATLASS_DEVICE + BlockSchedulerGdnFwdHVec() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user) { + BlockSchedulerGdnFwdH::Init(cu_seqlens, chunk_indices, tiling, user, AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(), AscendC::GetBlockNum()); + } + +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_SCHEDULER_GDN_FWD_H_HPP \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/gemm/kernel/gdn_fwd_h_kernel.hpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/gemm/kernel/gdn_fwd_h_kernel.hpp new file mode 100644 index 000000000..4a97a9888 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/arch22/gemm/kernel/gdn_fwd_h_kernel.hpp @@ -0,0 +1,408 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#define CATLASS_ARCH 3510 + +#include "catlass/arch/arch.hpp" +#include "catlass/arch/cross_core_sync.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/catlass.hpp" +#include "catlass/debug.hpp" +#include "catlass/epilogue/block/block_epilogue.hpp" +#include "../../epilogue/block/block_epilogue_gdn_fwdh_update.hpp" +#include "../../epilogue/block/block_epilogue_gdn_fwdh_vnew.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/gemm/block/block_swizzle.hpp" +#include "../block/block_scheduler_gdn_fwd_h.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" +#include "catlass/gemm_coord.hpp" +#include "tla/tensor.hpp" +#include "tla/layout.hpp" +#include "tla/tensor.hpp" + +using _0 = tla::Int<0>; +using _1 = tla::Int<1>; +using _2 = tla::Int<2>; +using _4 = tla::Int<4>; +using _8 = tla::Int<8>; +using _16 = tla::Int<16>; +using _32 = tla::Int<32>; +using _64 = tla::Int<64>; +using _128 = tla::Int<128>; +using _256 = tla::Int<256>; +using _512 = tla::Int<512>; +using _1024 = tla::Int<1024>; +using _2048 = tla::Int<2048>; +using _4096 = tla::Int<4096>; +using _8192 = tla::Int<8192>; +using _16384 = tla::Int<16384>; +using _32768 = tla::Int<32768>; +using _65536 = tla::Int<65536>; + +#else +#define CATLASS_ARCH 2201 + +#include "catlass/arch/arch.hpp" +#include "catlass/arch/cross_core_sync.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/catlass.hpp" +#include "catlass/debug.hpp" +#include "catlass/epilogue/block/block_epilogue.hpp" +#include "../../epilogue/block/block_epilogue_gdn_fwdh_update.hpp" +#include "../../epilogue/block/block_epilogue_gdn_fwdh_vnew.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/gemm/block/block_swizzle.hpp" +#include "../block/block_scheduler_gdn_fwd_h.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" +#include "catlass/gemm_coord.hpp" +#include "tla/tensor.hpp" +#include "tla/layout.hpp" +#include "tla/tensor.hpp" +#endif + + + +#include "kernel_operator.h" +using namespace Catlass; +using namespace tla; + +namespace Catlass::Gemm::Kernel { + +template< + typename INPUT_TYPE, + typename G_TYPE, + typename STATE_TYPE, + typename WORKSPACE_TYPE +> +class GDNFwdHKernel { +public: + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + using ArchTag = Arch::Ascend950; +#else + using ArchTag = Arch::AtlasA2; +#endif + using CubeScheduler = typename Catlass::Gemm::Block::BlockSchedulerGdnFwdHCube; + using VecScheduler = typename Catlass::Gemm::Block::BlockSchedulerGdnFwdHVec; + + using DispatchPolicyTla = Gemm::MmadPingpongTlaMulti; + using L1TileShapeTla = Shape<_128, _128, _128>; + using L0TileShapeTla = L1TileShapeTla; + + using WType = Gemm::GemmType; + using HType = Gemm::GemmType; + using VworkType = Gemm::GemmType; + using KType = Gemm::GemmType; + using HworkType = Gemm::GemmType; + using VType = Gemm::GemmType; + using GType = Gemm::GemmType; + using UType = Gemm::GemmType; + using FinalStateType = Gemm::GemmType; + + // cube 1 + using TileCopyWH = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadWH = Gemm::Block::BlockMmadTla; + + // cube 2 + using TileCopyKV = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadKV = Gemm::Block::BlockMmadTla; + + // vec 1 + using DispatchPolicyGDNFwdHVnew = Epilogue::EpilogueAtlasGDNFwdHVnew; + using EpilogueGDNFwdHVnew = Epilogue::Block::BlockEpilogue; + + // vec 2 + using DispatchPolicyGDNFwdHUpdate = Epilogue::EpilogueAtlasGDNFwdHUpdate; + using EpilogueGDNFwdHUpdate = Epilogue::Block::BlockEpilogue; + + using GDNFwdHOffsets = Catlass::Gemm::Block::GDNFwdHOffsets; + + using ElementK = INPUT_TYPE; + using ElementW = INPUT_TYPE; + using ElementU = INPUT_TYPE; + using ElementG = G_TYPE; + using ElementH = INPUT_TYPE; + using ElementV = INPUT_TYPE; + using ElementVWork = WORKSPACE_TYPE; + using ElementHWork = WORKSPACE_TYPE; + using ElementInitialState = STATE_TYPE; + using ElementFinalState = STATE_TYPE; + + using LayoutW = Catlass::layout::RowMajor; + using LayoutH = Catlass::layout::RowMajor; + using LayoutV = Catlass::layout::RowMajor; + using LayoutK = Catlass::layout::ColumnMajor; + + + uint32_t batch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + uint32_t initalStateStride0; + bool useInitialState; + bool storeFinalState; + uint32_t isVariedLen; + uint32_t shapeBatch; + uint32_t tokenBatch; + uint32_t vWorkspaceOffset; + uint32_t vUpdateWorkspaceOffset; + uint32_t hWorkspaceOffset; + uint32_t numSeqWorkspaceOffset; + uint32_t numChunksWorkspaceOffset; + + AscendC::GlobalTensor gmK; + AscendC::GlobalTensor gmW; + AscendC::GlobalTensor gmU; + AscendC::GlobalTensor gmG; + AscendC::GlobalTensor gmInitialState; + AscendC::GlobalTensor gmH; + AscendC::GlobalTensor gmV; + AscendC::GlobalTensor gmFinalState; + AscendC::GlobalTensor gmVWorkspace; + AscendC::GlobalTensor gmVUpdateWorkspace; + AscendC::GlobalTensor gmHWorkspace; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmNumSeq; + AscendC::GlobalTensor gmNumChunks; + + CubeScheduler cubeBlockScheduler; + VecScheduler vecBlockScheduler; + + Arch::Resource resource; + + + __aicore__ inline GDNFwdHKernel() {} + + __aicore__ inline void Init(GM_ADDR k, GM_ADDR w, GM_ADDR u, GM_ADDR g, GM_ADDR inital_state, GM_ADDR cu_seqlens, GM_ADDR chunk_indices, + GM_ADDR h, GM_ADDR v_new, GM_ADDR final_state, GM_ADDR tiling, GM_ADDR user) { + + __gm__ ChunkGatedDeltaRuleFwdHTilingData *__restrict gdnFwdHTilingData = reinterpret_cast<__gm__ ChunkGatedDeltaRuleFwdHTilingData *__restrict>(tiling); + + batch = gdnFwdHTilingData->batch; + seqlen = gdnFwdHTilingData->seqlen; + kNumHead = gdnFwdHTilingData->kNumHead; + vNumHead = gdnFwdHTilingData->vNumHead; + kHeadDim = gdnFwdHTilingData->kHeadDim; + vHeadDim = gdnFwdHTilingData->vHeadDim; + chunkSize = gdnFwdHTilingData->chunkSize; + initalStateStride0 = gdnFwdHTilingData->initalStateStride0; + useInitialState = gdnFwdHTilingData->useInitialState; + storeFinalState = gdnFwdHTilingData->storeFinalState; + isVariedLen = gdnFwdHTilingData->isVariedLen; + shapeBatch = gdnFwdHTilingData->shapeBatch; + tokenBatch = gdnFwdHTilingData->tokenBatch; + vWorkspaceOffset = gdnFwdHTilingData->vWorkspaceOffset; + vUpdateWorkspaceOffset = gdnFwdHTilingData->vUpdateWorkspaceOffset; + hWorkspaceOffset = gdnFwdHTilingData->hWorkspaceOffset; + numSeqWorkspaceOffset = gdnFwdHTilingData->numSeqWorkspaceOffset; + numChunksWorkspaceOffset = gdnFwdHTilingData->numChunksWorkspaceOffset; + + gmK.SetGlobalBuffer((__gm__ ElementK *)k); + gmW.SetGlobalBuffer((__gm__ ElementW *)w); + gmU.SetGlobalBuffer((__gm__ ElementU *)u); + gmG.SetGlobalBuffer((__gm__ ElementG *)g); + gmInitialState.SetGlobalBuffer((__gm__ ElementInitialState *)inital_state); + gmH.SetGlobalBuffer((__gm__ ElementH *)h); + gmV.SetGlobalBuffer((__gm__ ElementV *)v_new); + gmFinalState.SetGlobalBuffer((__gm__ ElementFinalState *)final_state); + gmVWorkspace.SetGlobalBuffer((__gm__ ElementVWork *)(user + vWorkspaceOffset)); + gmVUpdateWorkspace.SetGlobalBuffer((__gm__ ElementV *)(user + vUpdateWorkspaceOffset)); + gmHWorkspace.SetGlobalBuffer((__gm__ ElementHWork *)(user + hWorkspaceOffset)); + + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + if ASCEND_IS_AIC { + cubeBlockScheduler.Init(cu_seqlens, chunk_indices, tiling, user); + } + + if ASCEND_IS_AIV { + vecBlockScheduler.Init(cu_seqlens, chunk_indices, tiling, user); + } + } + + __aicore__ inline void Process() { + + if ASCEND_IS_AIC { + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = AscendC::GetBlockNum(); + + BlockMmadWH blockMmadWH(resource); + BlockMmadKV blockMmadKV(resource); + + auto wLayout = tla::MakeLayout(shapeBatch * kNumHead * cubeBlockScheduler.totalTokens, kHeadDim); + auto hLayout = tla::MakeLayout(shapeBatch * vNumHead * cubeBlockScheduler.totalChunks * kHeadDim, vHeadDim); + auto vLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, vHeadDim); + + auto kLayout = tla::MakeLayout(kHeadDim, shapeBatch * kNumHead * cubeBlockScheduler.totalTokens); + auto vworkLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, vHeadDim); + auto hworkLayout = tla::MakeLayout(coreNum * kHeadDim * PING_PONG_STAGES, vHeadDim); + + while (cubeBlockScheduler.isRunning) { + cubeBlockScheduler.InitTask(); + // step 1: v_work = w @ h[i] + GDNFwdHOffsets& cube1Offsets = cubeBlockScheduler.GetStage1Offsets(); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done); + if (cubeBlockScheduler.NeedProcessStage1()) { + int64_t cube1OffsetW = cube1Offsets.wOffset; + int64_t cube1OffsetH = cube1Offsets.hSrcOffset; + int64_t cube1OffsetVwork = cube1Offsets.vWorkOffset; + auto tensorW = tla::MakeTensor(gmW[cube1OffsetW], wLayout, Catlass::Arch::PositionGM{}); + auto tensorH = tla::MakeTensor(gmH[cube1OffsetH], hLayout, Catlass::Arch::PositionGM{}); + auto tensorV = tla::MakeTensor(gmVWorkspace[cube1OffsetVwork], vLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape {cube1Offsets.blockTokens, vHeadDim, kHeadDim}; + auto tensorBlockW = GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + auto tensorBlockV = GetTile(tensorV, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.n())); + blockMmadWH.preSetFlags(); + blockMmadWH(tensorBlockW, tensorBlockH, tensorBlockV, cube1Shape); + blockMmadWH.finalWaitFlags(); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube1Done); + + if (cubeBlockScheduler.iterId > 1) { + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done); + GDNFwdHOffsets& cube2Offsets = cubeBlockScheduler.GetStage2Offsets(); + if (cubeBlockScheduler.NeedProcessStage2()) { + // step 3: h[i+1] = k.T @ v_work + int64_t cube2OffsetK = cube2Offsets.wkOffset; + int64_t cube2OffsetVwork = cube2Offsets.vWorkOffset; + int64_t cube2OffsetH = cube2Offsets.hWorkOffset; + auto tensorK = tla::MakeTensor(gmK[cube2OffsetK], kLayout, Catlass::Arch::PositionGM{}); + auto tensorVwork = tla::MakeTensor(gmVUpdateWorkspace[cube2OffsetVwork], vworkLayout, Catlass::Arch::PositionGM{}); + auto tensorHwork = tla::MakeTensor(gmHWorkspace[cube2OffsetH], hworkLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube2Shape{kHeadDim, vHeadDim, cube2Offsets.blockTokens}; + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockVwork = GetTile(tensorVwork, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + auto tensorBlockHwork = GetTile(tensorHwork, tla::MakeCoord(0, 0), tla::MakeShape(cube2Shape.m(), cube2Shape.n())); + blockMmadKV.preSetFlags(); + blockMmadKV(tensorBlockK, tensorBlockVwork, tensorBlockHwork, cube2Shape); + blockMmadKV.finalWaitFlags(); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube2Done); + } + } + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done); + + } + + if ASCEND_IS_AIV { + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = AscendC::GetBlockNum(); + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + EpilogueGDNFwdHVnew epilogueGDNFwdHVnew(resource); + + if (useInitialState) { + AscendC::LocalTensor stateUbTensorPing = resource.ubBuf.template GetBufferByByte(0); + AscendC::LocalTensor stateUbTensorPong = resource.ubBuf.template GetBufferByByte(96 * 1024); + AscendC::LocalTensor hUbTensorPing = resource.ubBuf.template GetBufferByByte(64 * 1024); + AscendC::LocalTensor hUbTensorPong = resource.ubBuf.template GetBufferByByte(160 * 1024); + uint32_t totalChunks = isVariedLen ? vecBlockScheduler.totalChunks : ((seqlen + chunkSize - 1) / chunkSize); + uint32_t stateBlockSize = kHeadDim * vHeadDim; + uint32_t pingpongFlag = 1; + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::DataCopyParams repeatParams = {static_cast(kHeadDim), static_cast(vHeadDim * sizeof(ElementInitialState) / 32), + static_cast((initalStateStride0 - vHeadDim)* sizeof(ElementInitialState) / 32), static_cast(0)}; + for (uint32_t shapeBatchIdx = 0; shapeBatchIdx < shapeBatch; shapeBatchIdx++) { + for (uint32_t vHeadIdx = 0; vHeadIdx < vNumHead; vHeadIdx++) { + for (uint32_t tokenBatchIdx = 0; tokenBatchIdx < vecBlockScheduler.tokenBatch; tokenBatchIdx++) { + uint32_t batchIdx = isVariedLen ? tokenBatchIdx : shapeBatchIdx; + uint32_t chunkOffset = isVariedLen ? gmNumChunks.GetValue(tokenBatchIdx) : 0; + uint32_t initialStateSrcOffset = (batchIdx * vNumHead + vHeadIdx) * kHeadDim * initalStateStride0; + uint32_t hOffset = (shapeBatchIdx * vNumHead * totalChunks + vHeadIdx * totalChunks + chunkOffset) * stateBlockSize; + AscendC::LocalTensor stateUbTensor = pingpongFlag ? stateUbTensorPing : stateUbTensorPong; + AscendC::LocalTensor hUbTensor = pingpongFlag ? hUbTensorPing : hUbTensorPong; + auto event_id = pingpongFlag ? EVENT_ID1 : EVENT_ID0; + AscendC::WaitFlag(event_id); + if constexpr(!std::is_same::value) { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateSrcOffset], repeatParams); + AscendC::SetFlag(event_id); + AscendC::WaitFlag(event_id); + AscendC::Cast(hUbTensor, stateUbTensor, AscendC::RoundMode::CAST_RINT, stateBlockSize); + AscendC::SetFlag(event_id); + AscendC::WaitFlag(event_id); + AscendC::DataCopy(gmH[hOffset], hUbTensor, stateBlockSize); + } else { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateSrcOffset], repeatParams); + AscendC::SetFlag(event_id); + AscendC::WaitFlag(event_id); + AscendC::DataCopy(gmH[hOffset], stateUbTensor, stateBlockSize); + } + AscendC::SetFlag(event_id); + pingpongFlag = 1 - pingpongFlag; + } + + } + } + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + } + + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done); + while (vecBlockScheduler.isRunning) { + vecBlockScheduler.InitTask(); + // step 2: + GDNFwdHOffsets& vec1Offsets = vecBlockScheduler.GetStage1Offsets(); + // gmV = gmU - gmVWorkspace + // g_buf = gmG[-1] - gmG + // g_buf = exp(g_buf) + // gmVWorkspace = g_buf * gmV + if (vecBlockScheduler.NeedProcessStage1()) { + epilogueGDNFwdHVnew( + gmV[vec1Offsets.uvOffset], gmVUpdateWorkspace[vec1Offsets.vWorkOffset], + gmG[vec1Offsets.gOffset], gmU[vec1Offsets.uvOffset], gmVWorkspace[vec1Offsets.vWorkOffset], + vec1Offsets.blockTokens, kHeadDim, vHeadDim, vecBlockScheduler.cube1Done + ); + } else { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube1Done); + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec1Done); + + if (vecBlockScheduler.iterId > 1) { + GDNFwdHOffsets& vec2Offsets = vecBlockScheduler.GetStage2Offsets(); + if (vecBlockScheduler.NeedProcessStage2()) { + // step 4: h[i+1] += h_work if i < num_chunks - 1 else None + EpilogueGDNFwdHUpdate epilogueGDNFwdHUpdate(resource); + epilogueGDNFwdHUpdate( + gmH[vec2Offsets.hDstOffset], gmFinalState[vec2Offsets.finalStateOffset], + gmG[vec2Offsets.gOffset], + gmH[vec2Offsets.hSrcOffset], + gmHWorkspace[vec2Offsets.hWorkOffset], + vec2Offsets.blockTokens, kHeadDim, vHeadDim, vecBlockScheduler.cube2Done, + (vec2Offsets.isFinalState && storeFinalState) + ); + } else { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube2Done); + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done); + } + } + + } + } + +}; + +} \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/chunk_gated_delta_rule_fwd_h.cpp b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/chunk_gated_delta_rule_fwd_h.cpp new file mode 100644 index 000000000..9167f2bcd --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/op_kernel/chunk_gated_delta_rule_fwd_h.cpp @@ -0,0 +1,95 @@ +/** + * Copyright (c) 2026 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +/*! + * \file chunk_gated_delta_rule_fwd_h.cpp + * \brief + */ + +// #include "chunk_gated_delta_rule_fwd_h.h" +#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 200) +#include "arch20/compat_310p.h" +#include "arch20/gemm/kernel/gdn_fwd_h_kernel.hpp" +#else +#include "arch22/gemm/kernel/gdn_fwd_h_kernel.hpp" +#endif +#include "lib/matmul_intf.h" + +using namespace Catlass; + +extern "C" __global__ __aicore__ void chunk_gated_delta_rule_fwd_h(GM_ADDR k, GM_ADDR w, GM_ADDR u, GM_ADDR g, + GM_ADDR inital_state, GM_ADDR cu_seqlens, GM_ADDR chunk_indices, + GM_ADDR h, GM_ADDR v_new, GM_ADDR final_state, + GM_ADDR workspace, GM_ADDR tiling) +{ + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + + GM_ADDR user = AscendC::GetUserWorkspace(workspace); + + __gm__ ChunkGatedDeltaRuleFwdHTilingData *__restrict gdnFwdHTilingData = reinterpret_cast<__gm__ ChunkGatedDeltaRuleFwdHTilingData *__restrict>(tiling); + + using workspaceType = float; + // dtype: 0 - fp16, 1 - bf16, 2 - fp32 +#ifndef CATLASS_UNIFIED_CORE + if (gdnFwdHTilingData->dataType == 1) { + if (gdnFwdHTilingData->stateDataType == 2) { + if (gdnFwdHTilingData->gDataType == 2) { + using GDNFwdHKernel = Catlass::Gemm::Kernel::GDNFwdHKernel; + GDNFwdHKernel gdnFwdH; + gdnFwdH.Init(k, w, u, g, inital_state, cu_seqlens, chunk_indices, h, v_new, final_state, tiling, user); + gdnFwdH.Process(); + } else { + using GDNFwdHKernel = Catlass::Gemm::Kernel::GDNFwdHKernel; + GDNFwdHKernel gdnFwdH; + gdnFwdH.Init(k, w, u, g, inital_state, cu_seqlens, chunk_indices, h, v_new, final_state, tiling, user); + gdnFwdH.Process(); + } + } else { + if (gdnFwdHTilingData->gDataType == 2) { + using GDNFwdHKernel = Catlass::Gemm::Kernel::GDNFwdHKernel; + GDNFwdHKernel gdnFwdH; + gdnFwdH.Init(k, w, u, g, inital_state, cu_seqlens, chunk_indices, h, v_new, final_state, tiling, user); + gdnFwdH.Process(); + } else { + using GDNFwdHKernel = Catlass::Gemm::Kernel::GDNFwdHKernel; + GDNFwdHKernel gdnFwdH; + gdnFwdH.Init(k, w, u, g, inital_state, cu_seqlens, chunk_indices, h, v_new, final_state, tiling, user); + gdnFwdH.Process(); + } + } + } else +#endif + { + if (gdnFwdHTilingData->stateDataType == 2) { + if (gdnFwdHTilingData->gDataType == 2) { + using GDNFwdHKernel = Catlass::Gemm::Kernel::GDNFwdHKernel; + GDNFwdHKernel gdnFwdH; + gdnFwdH.Init(k, w, u, g, inital_state, cu_seqlens, chunk_indices, h, v_new, final_state, tiling, user); + gdnFwdH.Process(); + } else { + using GDNFwdHKernel = Catlass::Gemm::Kernel::GDNFwdHKernel; + GDNFwdHKernel gdnFwdH; + gdnFwdH.Init(k, w, u, g, inital_state, cu_seqlens, chunk_indices, h, v_new, final_state, tiling, user); + gdnFwdH.Process(); + } + } else { + if (gdnFwdHTilingData->gDataType == 2) { + using GDNFwdHKernel = Catlass::Gemm::Kernel::GDNFwdHKernel; + GDNFwdHKernel gdnFwdH; + gdnFwdH.Init(k, w, u, g, inital_state, cu_seqlens, chunk_indices, h, v_new, final_state, tiling, user); + gdnFwdH.Process(); + } else { + using GDNFwdHKernel = Catlass::Gemm::Kernel::GDNFwdHKernel; + GDNFwdHKernel gdnFwdH; + gdnFwdH.Init(k, w, u, g, inital_state, cu_seqlens, chunk_indices, h, v_new, final_state, tiling, user); + gdnFwdH.Process(); + } + } + } +} diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/data_copy_transpose_tiling.h b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/data_copy_transpose_tiling.h new file mode 100644 index 000000000..f7e7fce1c --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/data_copy_transpose_tiling.h @@ -0,0 +1,51 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file data_copy_transpose_tiling.h + * \brief + */ + +#pragma once + +#include +#include +#include "data_copy_transpose_tiling_def.h" + +namespace optiling { + +inline void GetDataCopyTransposeTiling(const ge::Shape &dstShape, const ge::Shape &srcShape, const uint32_t typeSize, + optiling::CopyTransposeTiling &tiling) +{ + constexpr int64_t B_INDEX = 0; + constexpr int64_t N_INDEX = 1; + constexpr int64_t S_INDEX = 2; + constexpr int64_t H_INDEX = 3; + std::vector dstShapeInfo = dstShape.GetDims(); + std::vector srcShapeInfo = srcShape.GetDims(); + + tiling.set_dstShapeB(dstShapeInfo[B_INDEX]); + tiling.set_dstShapeN(dstShapeInfo[N_INDEX]); + tiling.set_dstShapeS(dstShapeInfo[S_INDEX]); + tiling.set_dstShapeH(dstShapeInfo[H_INDEX]); + tiling.set_dstShapeHN(tiling.get_dstShapeH() / tiling.get_dstShapeN()); + + tiling.set_srcShapeB(srcShapeInfo[B_INDEX]); + tiling.set_srcShapeN(srcShapeInfo[N_INDEX]); + tiling.set_srcShapeS(srcShapeInfo[S_INDEX]); + tiling.set_srcShapeHN(srcShapeInfo[H_INDEX]); + tiling.set_originalShapeNLen(tiling.get_srcShapeHN() * typeSize); + tiling.set_shapeSHValue(tiling.get_dstShapeS() * tiling.get_dstShapeH()); + tiling.set_shapeNsValue(tiling.get_dstShapeN() * tiling.get_dstShapeS()); + tiling.set_shapeNsnValue(tiling.get_dstShapeN() * tiling.get_srcShapeS() * tiling.get_srcShapeN()); + tiling.set_shapeBHValue(tiling.get_dstShapeB() * tiling.get_dstShapeH()); +} + +} // namespace optiling diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/data_copy_transpose_tiling_def.h b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/data_copy_transpose_tiling_def.h new file mode 100644 index 000000000..18552c36b --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/data_copy_transpose_tiling_def.h @@ -0,0 +1,43 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file data_copy_transpose_tiling_def.h + * \brief + */ + +#pragma once + +#include +#include + +namespace optiling { + +BEGIN_TILING_DATA_DEF(CopyTransposeTiling) +TILING_DATA_FIELD_DEF(uint32_t, dstShapeB); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeN); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeS); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeHN); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeH); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeB); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeN); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeS); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeHN); +TILING_DATA_FIELD_DEF(uint32_t, originalShapeNLen); +TILING_DATA_FIELD_DEF(uint32_t, shapeSHValue); +TILING_DATA_FIELD_DEF(uint32_t, shapeNsValue); +TILING_DATA_FIELD_DEF(uint32_t, shapeNsnValue); +TILING_DATA_FIELD_DEF(uint32_t, invalidParamCopyTransposeTiling); +TILING_DATA_FIELD_DEF(uint32_t, shapeBHValue); +TILING_DATA_FIELD_DEF(uint32_t, paramsAlign); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(CopyTransposeTilingOp, CopyTransposeTiling) + +} // namespace optiling diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/error_log.h b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/error_log.h new file mode 100644 index 000000000..119189974 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/error_log.h @@ -0,0 +1,63 @@ +#ifndef OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_ +#define OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_ + +#include +#include +#include "toolchain/slog.h" + +#define OP_LOGI(opname, ...) +#define OP_LOGW(opname, ...) \ + do { \ + (void)(opname); \ + std::printf("[WARN] "); \ + std::printf(__VA_ARGS__); \ + std::printf("\n"); \ + } while (0) + +#define OP_LOGE_WITHOUT_REPORT(opname, ...) \ + do { \ + (void)(opname); \ + std::printf("[ERRORx] "); \ + std::printf(__VA_ARGS__); \ + std::printf("\n"); \ + } while (0) + +#define OP_LOGE(opname, ...) \ + do { \ + (void)(opname); \ + std::printf("[ERROR] "); \ + std::printf(__VA_ARGS__); \ + std::printf("\n"); \ + } while (0) + +#define OP_LOGD(opname, ...) + +namespace optiling { + +#define VECTOR_INNER_ERR_REPORT_TILIING(op_name, err_msg, ...) \ + do { \ + OP_LOGE_WITHOUT_REPORT(op_name, err_msg, ##__VA_ARGS__); \ + } while (0) + +// Modify OP_TILING_CHECK macro to ensure proper handling of expressions +#define OP_CHECK_IF(cond, log_func, expr) \ + do { \ + if (cond) { \ + log_func; \ + expr; \ + } \ + } while (0) + + + +#define OP_CHECK_NULL_WITH_CONTEXT(context, ptr) \ + do { \ + if ((ptr) == nullptr) { \ + OP_LOGE(context->GetNodeType(), "%s is null", #ptr); \ + return ge::GRAPH_FAILED; \ + } \ + } while (0) + +} // namespace optiling + +#endif // OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_ diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_base.h b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_base.h new file mode 100644 index 000000000..f0bbbdccc --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_base.h @@ -0,0 +1,256 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_base.h + * \brief + */ + +#pragma once + +#include +#include +#include +#include "tiling/platform/platform_ascendc.h" +#include "error_log.h" + +#ifdef ASCENDC_OP_TEST +#define ASCENDC_EXTERN_C extern "C" +#else +#define ASCENDC_EXTERN_C +#endif + +namespace Ops { +namespace Transformer { +namespace OpTiling { + +struct AiCoreParams { + uint64_t ubSize = 0; + uint64_t blockDim = 0; + uint64_t aicNum = 0; + uint64_t l1Size = 0; + uint64_t l0aSize = 0; + uint64_t l0bSize = 0; + uint64_t l0cSize = 0; +}; + +struct CompileInfoCommon { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; + int32_t socVersion; + uint32_t rsvd; +}; + +struct FlashAttentionScoreGradCompileInfo { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; + platform_ascendc::SocVersion socVersion; +}; + +struct FACompileInfoCommon { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; + int32_t socVersion; + uint32_t rsvd; +}; + +class TilingBaseClass { +public: + explicit TilingBaseClass(gert::TilingContext* context) : context_(context) + {} + + virtual ~TilingBaseClass() = default; + + // Tiling execution framework + // 1. GRAPH_SUCCESS: Success, and no need to continue executing subsequent Tiling class implementations + // 2. GRAPH_FAILED: Failure, abort the entire Tiling process + // 3. GRAPH_PARAM_INVALID: This class does not support, need to continue executing other Tiling class implementations + ge::graphStatus DoTiling() + { + auto ret = GetShapeAttrsInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetPlatformInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + if (!IsCapable()) { + return ge::GRAPH_PARAM_INVALID; + } + ret = DoOpTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = DoLibApiTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetWorkspaceSize(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = PostTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + context_->SetTilingKey(GetTilingKey()); + DumpTilingInfo(); + return ge::GRAPH_SUCCESS; + } + + // Update context + virtual void Reset(gert::TilingContext* context) + { + context_ = context; + } + +protected: + virtual bool IsCapable() = 0; + // 1. Get platform information such as CoreNum, UB/L1/L0C resource sizes + virtual ge::graphStatus GetPlatformInfo() = 0; + // 2. Get INPUT/OUTPUT/ATTR information + virtual ge::graphStatus GetShapeAttrsInfo() = 0; + // 3. Calculate data splitting TilingData + virtual ge::graphStatus DoOpTiling() = 0; + // 4. Calculate high-level API TilingData + virtual ge::graphStatus DoLibApiTiling() = 0; + // 5. Calculate TilingKey + [[nodiscard]] virtual uint64_t GetTilingKey() const = 0; + // 6. Calculate Workspace size + virtual ge::graphStatus GetWorkspaceSize() = 0; + // 7. Save Tiling data + virtual ge::graphStatus PostTiling() = 0; + // 8. Dump Tiling data + virtual void DumpTilingInfo() + { + int32_t enable = CheckLogLevel(static_cast(OP), DLOG_DEBUG); + if (enable != 1) { + return; + } + auto buf = (uint32_t*)context_->GetRawTilingData()->GetData(); + auto bufLen = context_->GetRawTilingData()->GetDataSize(); + std::ostringstream oss; + oss << "Start to dump tiling info. tilingkey:" << context_->GetTilingKey() << ", tiling data size:" << bufLen + << ", content:"; + for (size_t i = 0; i < bufLen / sizeof(uint32_t); i++) { + oss << *(buf + i) << ","; + if (oss.str().length() > 640) { // Split according to 640 to avoid truncation + OP_LOGD(context_, "%s", oss.str().c_str()); + oss.str(""); + } + } + OP_LOGD(context_, "%s", oss.str().c_str()); + } + + static uint32_t CalcTschBlockDim(uint32_t sliceNum, uint32_t aicCoreNum, uint32_t aivCoreNum) + { + uint32_t ration; + if (aicCoreNum == 0 || aivCoreNum == 0 || aicCoreNum > aivCoreNum) { + return sliceNum; + } + ration = aivCoreNum / aicCoreNum; + return (sliceNum + (ration - 1)) / ration; + } + + template + [[nodiscard]] std::string GetShapeDebugStr(const T& shape) const + { + std::ostringstream oss; + oss << "["; + if (shape.GetDimNum() > 0) { + for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) { + oss << shape.GetDim(i) << ", "; + } + oss << shape.GetDim(shape.GetDimNum() - 1); + } + oss << "]"; + return oss.str(); + } + + [[nodiscard]] std::string GetTensorDebugStr( + const gert::StorageShape* shape, const gert::CompileTimeTensorDesc* tensor) + { + if (shape == nullptr || tensor == nullptr) { + return "nil "; + } + std::ostringstream oss; + oss << "(dtype: " << ge::TypeUtils::DataTypeToSerialString(tensor->GetDataType()) << "),"; + oss << "(shape:" << GetShapeDebugStr(shape->GetStorageShape()) << "),"; + oss << "(ori_shape:" << GetShapeDebugStr(shape->GetOriginShape()) << "),"; + oss << "(format: " + << ge::TypeUtils::FormatToSerialString( + static_cast(ge::GetPrimaryFormat(tensor->GetStorageFormat()))) + << "),"; + oss << "(ori_format: " << ge::TypeUtils::FormatToSerialString(tensor->GetOriginFormat()) << ") "; + return oss.str(); + } + + [[nodiscard]] std::string GetTilingContextDebugStr() + { + std::ostringstream oss; + for (size_t i = 0; i < context_->GetComputeNodeInfo()->GetInputsNum(); ++i) { + oss << "input" << i << ": "; + oss << GetTensorDebugStr(context_->GetInputShape(i), context_->GetInputDesc(i)); + } + + for (size_t i = 0; i < context_->GetComputeNodeInfo()->GetOutputsNum(); ++i) { + oss << "output" << i << ": "; + oss << GetTensorDebugStr(context_->GetOutputShape(i), context_->GetOutputDesc(i)); + } + return oss.str(); + } + + [[nodiscard]] std::string GetTilingDataDebugStr() const + { + auto rawTilingData = context_->GetRawTilingData(); + auto rawTilingDataSize = rawTilingData->GetDataSize(); + auto data = reinterpret_cast(rawTilingData->GetData()); + size_t len = rawTilingDataSize / sizeof(int32_t); + std::ostringstream oss; + for (size_t i = 0; i < len; i++) { + oss << data[i] << ", "; + } + return oss.str(); + } + +protected: + gert::TilingContext* context_ = nullptr; + std::unique_ptr ascendcPlatform_{nullptr}; + uint32_t blockDim_{0}; + uint64_t workspaceSize_{0}; + uint64_t tilingKey_{0}; + AiCoreParams aicoreParams_; +}; + +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops \ No newline at end of file diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_key.h b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_key.h new file mode 100644 index 000000000..607f965b0 --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_key.h @@ -0,0 +1,63 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_key.h + * \brief + */ + +#pragma once + +#include + +namespace Ops { +namespace Transformer { +namespace OpTiling { +constexpr uint64_t RecursiveSum() +{ + return 0; +} + +constexpr uint64_t kBase = 10; // Base-10 carry base +template constexpr uint64_t RecursiveSum(T templateId, Args... templateIds) +{ + return static_cast(templateId) + kBase * RecursiveSum(templateIds...); +} + +// TilingKey generation rules: +// FlashAttentionScore/FlashAttentionScoreGrad assembles tiling key using decimal digits, containing the following key parameters from low to high: Ub0, Ub1, +// Block, DataType, Format, Sparse. Specialized template Ub0, Ub1: +// Represents the axis for UB intra-core splitting, using AxisEnum. Since we allow at most two axes to be split, UB0 and UB1 exist. If there is no UB intra-core splitting, +// fill with AXIS_NONE. UB0 and UB1 each occupy one decimal digit; +// Block: Represents the axis used by UB for multi-core splitting, using AxisEnum, occupies one decimal digit; +// DataType: Represents the input/output data types supported by the current tiling key, using SupportedDtype enum, occupies one decimal digit +// Format: Represents the Format supported by the current tiling key, using InputLayout enum, occupies one decimal digit +// Sparse: Represents whether the current tiling key supports Sparse, using SparseCapability enum, occupies one decimal digit +// For other specialized scenarios, define your own bit fields and values +// usage: get tilingKey from inputted types +// uint64_t tilingKey = GET_FLASHATTENTION_TILINGKEY(AxisEnum::AXIS_S1, AxisEnum::AXIS_S2, AxisEnum::AXIS_N2, +// SupportedDtype::FLOAT32, InputLayout::BSH, SparseCapability::SUPPORT_ALL) + +constexpr uint64_t TILINGKEYOFFSET = uint64_t(10000000000000000000UL); // 10^19 +template constexpr uint64_t GET_TILINGKEY(Args... templateIds) +{ + return TILINGKEYOFFSET + RecursiveSum(templateIds...); +} + +// usage: get tilingKey from inputted types +// uint64_t tilingKey = TILINGKEY(S2, S1, N2, FLOAT32, BSND, ALL) + +#define TILINGKEY(ub2, ub1, block, dtype, layout, sparse) \ + (GET_TILINGKEY(AxisEnum::ub2, AxisEnum::ub1, AxisEnum::block, DtypeEnum::dtype, LayoutEnum::layout, \ + SparseEnum::sparse)) + +} // namespace Optiling +} // namespace Transformer +} // namespace Ops diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_templates_registry.h b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_templates_registry.h new file mode 100644 index 000000000..cbf4785ab --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_templates_registry.h @@ -0,0 +1,351 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_templates_registry.h + * \brief + */ + +#pragma once + +#include +#include +#include +#include "exe_graph/runtime/tiling_context.h" +#include "tiling_base.h" +#include "error_log.h" + +namespace Ops { +namespace Transformer { +namespace OpTiling { + +template +std::unique_ptr TILING_CLASS(gert::TilingContext* context) +{ + return std::unique_ptr(new (std::nothrow) T(context)); +} + +using TilingClassCase = std::unique_ptr (*)(gert::TilingContext*); + +class TilingCases { +public: + explicit TilingCases(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + void AddTiling(int32_t priority) + { + OP_CHECK_IF( + cases_.find(priority) != cases_.end(), OP_LOGE(op_type_, "There are duplicate registrations."), return); + cases_[priority] = TILING_CLASS; + OP_CHECK_IF( + cases_[priority] == nullptr, + OP_LOGE(op_type_, "Register op tiling func failed, please check the class name."), return); + } + + const std::map& GetTilingCases() + { + return cases_; + } + +private: + std::map cases_; + const std::string op_type_; +}; + +// --------------------------------Interfacce with soc version -------------------------------- +class TilingRegistryNew { +public: + TilingRegistryNew() = default; + +#ifdef ASCENDC_OP_TEST + static TilingRegistryNew& GetInstance(); +#else + static TilingRegistryNew& GetInstance() + { + static TilingRegistryNew registry_impl_; + return registry_impl_; + } +#endif + + std::shared_ptr RegisterOp(const std::string& op_type, int32_t soc_version) + { + auto soc_iter = registry_map_.find(soc_version); + if (soc_iter == registry_map_.end()) { + std::map> op_type_map; + op_type_map[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + registry_map_[soc_version] = op_type_map; + } else { + if (soc_iter->second.find(op_type) == soc_iter->second.end()) { + soc_iter->second[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + } + } + + OP_CHECK_IF( + registry_map_[soc_version][op_type] == nullptr, + OP_LOGE(op_type, "Register tiling func failed, please check the class name."), return nullptr); + return registry_map_[soc_version][op_type]; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context) + { + int32_t soc_version = (int32_t)platform_ascendc::SocVersion::RESERVED_VERSION; + const char* op_type = context->GetNodeType(); + fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + auto compileInfoPtr = static_cast(context->GetCompileInfo()); + OP_CHECK_IF( + compileInfoPtr == nullptr, OP_LOGE(op_type, "compileInfoPtr is null."), return ge::GRAPH_FAILED); + soc_version = compileInfoPtr->socVersion; + OP_LOGD(context, "soc version in compileInfo is %d", soc_version); + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + soc_version = static_cast(ascendcPlatform.GetSocVersion()); + OP_LOGD(context, "soc version is %d", soc_version); + if (soc_version == (int32_t)platform_ascendc::SocVersion::RESERVED_VERSION) { + OP_LOGE(op_type, "Do op tiling failed, cannot find soc version."); + return ge::GRAPH_FAILED; + } + } + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type, soc_version); + for (auto it = tilingTemplateRegistryMap.begin(); it != tilingTemplateRegistryMap.end(); ++it) { + auto tilingTemplate = it->second(context); + if (tilingTemplate != nullptr) { + ge::graphStatus status = tilingTemplate->DoTiling(); + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do general op tiling success priority=%d", it->first); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", it->first); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context, const std::vector& priorities) + { + int32_t soc_version; + const char* op_type = context->GetNodeType(); + auto platformInfoPtr = context->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + auto compileInfoPtr = reinterpret_cast(context->GetCompileInfo()); + OP_CHECK_IF( + compileInfoPtr == nullptr, OP_LOGE(op_type, "compileInfoPtr is null."), return ge::GRAPH_FAILED); + soc_version = compileInfoPtr->socVersion; + OP_LOGD(context, "soc version in compileInfo is %d", soc_version); + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + soc_version = static_cast(ascendcPlatform.GetSocVersion()); + OP_LOGD(context, "soc version is %d", soc_version); + } + + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type, soc_version); + for (auto priority_id : priorities) { + auto tilingCaseIter = tilingTemplateRegistryMap.find(priority_id); + if (tilingCaseIter != tilingTemplateRegistryMap.end()) { + auto templateFunc = tilingCaseIter->second(context); + if (templateFunc != nullptr) { + ge::graphStatus status = templateFunc->DoTiling(); + if (status == ge::GRAPH_SUCCESS) { + OP_LOGD(context, "Do general op tiling success priority=%d", priority_id); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", priority_id); + } + } + } + return ge::GRAPH_FAILED; + } + + const std::map& GetTilingTemplates(const std::string& op_type, int32_t soc_version) + { + auto soc_iter = registry_map_.find(soc_version); + OP_CHECK_IF( + soc_iter == registry_map_.end(), + OP_LOGE(op_type, "Get op tiling func failed, please check the soc version %d", soc_version), + return empty_tiling_case_); + auto op_iter = soc_iter->second.find(op_type); + OP_CHECK_IF( + op_iter == soc_iter->second.end(), OP_LOGE(op_type, "Get op tiling func failed, please check the op name."), + return empty_tiling_case_); + return op_iter->second->GetTilingCases(); + } + +private: + std::map>> registry_map_; // key is socversion + const std::map empty_tiling_case_{}; +}; + +class RegisterNew { +public: + explicit RegisterNew(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + RegisterNew& tiling(int32_t priority, int32_t soc_version) + { + auto tilingCases = TilingRegistryNew::GetInstance().RegisterOp(op_type_, soc_version); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), return *this); + tilingCases->AddTiling(priority); + return *this; + } + + template + RegisterNew& tiling(int32_t priority, const std::vector& soc_versions) + { + for (int32_t soc_version : soc_versions) { + auto tilingCases = TilingRegistryNew::GetInstance().RegisterOp(op_type_, soc_version); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), + return *this); + tilingCases->AddTiling(priority); + } + return *this; + } + +private: + const std::string op_type_; +}; + +// --------------------------------Interfacce without soc version -------------------------------- +class TilingRegistry { +public: + TilingRegistry() = default; + +#ifdef ASCENDC_OP_TEST + static TilingRegistry& GetInstance(); +#else + static TilingRegistry& GetInstance() + { + static TilingRegistry registry_impl_; + return registry_impl_; + } +#endif + + std::shared_ptr RegisterOp(const std::string& op_type) + { + if (registry_map_.find(op_type) == registry_map_.end()) { + registry_map_[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + } + OP_CHECK_IF( + registry_map_[op_type] == nullptr, + OP_LOGE(op_type, "Register tiling func failed, please check the class name."), return nullptr); + return registry_map_[op_type]; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context) + { + const char* op_type = context->GetNodeType(); + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type); + for (auto it = tilingTemplateRegistryMap.begin(); it != tilingTemplateRegistryMap.end(); ++it) { + auto tilingTemplate = it->second(context); + if (tilingTemplate != nullptr) { + ge::graphStatus status = tilingTemplate->DoTiling(); + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do general op tiling success priority=%d", it->first); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", it->first); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context, const std::vector& priorities) + { + const char* op_type = context->GetNodeType(); + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type); + for (auto priorityId : priorities) { + auto templateFunc = tilingTemplateRegistryMap[priorityId](context); + if (templateFunc != nullptr) { + ge::graphStatus status = templateFunc->DoTiling(); + if (status == ge::GRAPH_SUCCESS) { + OP_LOGD(context, "Do general op tiling success priority=%d", priorityId); + return status; + } + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do op tiling failed"); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", priorityId); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + const std::map& GetTilingTemplates(const std::string& op_type) + { + OP_CHECK_IF( + registry_map_.find(op_type) == registry_map_.end(), + OP_LOGE(op_type, "Get op tiling func failed, please check the op name."), return empty_tiling_case_); + return registry_map_[op_type]->GetTilingCases(); + } + +private: + std::map> registry_map_; + const std::map empty_tiling_case_; +}; + +class Register { +public: + explicit Register(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + Register& tiling(int32_t priority) + { + auto tilingCases = TilingRegistry::GetInstance().RegisterOp(op_type_); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), return *this); + tilingCases->AddTiling(priority); + return *this; + } + +private: + const std::string op_type_; +}; +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops + +// op_type: operator name, class_name: registered tiling class, soc_version: chip version number +// priority: priority of tiling class, smaller value means higher priority, i.e., this tiling class will be selected first +#define REGISTER_TILING_TEMPLATE_WITH_SOCVERSION(op_type, class_name, soc_versions, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::RegisterNew VAR_UNUSED##op_type##class_name##priority_register = \ + Ops::Transformer::OpTiling::RegisterNew(#op_type).tiling(priority, soc_versions) + +// op_type: operator name, class_name: registered tiling class +// priority: priority of tiling class, smaller value means higher priority, i.e., higher probability of being selected +#define REGISTER_TILING_TEMPLATE(op_type, class_name, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::Register VAR_UNUSED##op_type_##class_name##priority_register = \ + Ops::Transformer::OpTiling::Register(op_type).tiling(priority) + +// op_type: operator name, class_name: registered tiling class +// soc_version: SOC version, used to distinguish different SOCs +// priority: priority of tiling class, smaller value means higher priority, i.e., this tiling class will be selected first +#define REGISTER_TILING_TEMPLATE_NEW(op_type, class_name, soc_version, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::RegisterNew VAR_UNUSED##op_type##class_name##priority_register = \ + Ops::Transformer::OpTiling::RegisterNew(#op_type).tiling(priority, soc_version) + +// op_type: operator name, class_name: registered tiling class +// priority: priority of tiling class, smaller value means higher priority, i.e., higher probability of being selected +// Replaces REGISTER_TILING_TEMPLATE, if op_type is a string constant, remove the quotes +#define REGISTER_OPS_TILING_TEMPLATE(op_type, class_name, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::Transformer::OpTiling::Register \ + __attribute__((unused)) tiling_##op_type##_##class_name##_##priority##_register = \ + Ops::Transformer::OpTiling::Register(#op_type).tiling(priority) diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_type.h b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_type.h new file mode 100644 index 000000000..7c781d19f --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_type.h @@ -0,0 +1,139 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_type.h + * \brief + */ + +#pragma once + +#include + +namespace optiling { + +enum class AxisEnum { + B = 0, + N2 = 1, + G = 2, + S1 = 3, + S2 = 4, + D = 5, + NONE = 9, +}; + +enum class DtypeEnum { + FLOAT16 = 0, + FLOAT32 = 1, + BFLOAT16 = 2, + FLOAT16_PRECISION = 3, +}; + +enum class PerformanceOrientedEnum { + BIG_BUFFER = 1, + BIG_DOUBLE_BUFFER = 2, +}; + +enum class MatmulConfig { + NULL_CONFIG = 0, + NORMAL_CONFIG = 1, + MDL_CONFIG = 2 +}; + +enum class PseConfig { + NO_PSE = 0, + EXIST_PSE = 1 +}; + +enum class AttenMaskConfig { + NO_ATTEN_MASK = 0, + EXIST_ATTEN_MASK = 1 +}; + +enum class DropOutConfig { + NO_DROP_OUT = 0, + EXIST_DROP_OUT = 1 +}; + +enum class CubeFormatEnum { + ND = 0, + NZ = 1 +}; +enum class LayoutEnum { + BSND = 0, + SBND = 1, + BNSD = 2, + TND = 3, + NTD_TND = 4 +}; + +enum class CubeInputSourceEnum { + GM = 0, + L1 = 1 +}; + +enum class OptionEnum { + DISABLE = 0, + ENABLE = 1 +}; + +enum class SparseEnum { + ALL = 0, + NONE = 1, + ANY = 2, + CAUSAL = 3, + BAND = 4, + PREFIX = 5, + BAND_COMPRESS = 6, + RIGHT_DOWN_CAUSAL = 7, + RIGHT_DOWN_CAUSAL_BAND = 8, + BAND_LEFT_UP_CAUSAL = 9 +}; + +constexpr uint64_t RecursiveSum() +{ + return 0; +} + +constexpr int64_t base10Multiplier = 10; + +template constexpr uint64_t RecursiveSum(T templateId, Args... templateIds) +{ + return static_cast(templateId) + base10Multiplier * RecursiveSum(templateIds...); +} + +// TilingKey generation rules: +// FlashAttentionScore/FlashAttentionScoreGrad assembles tiling key using decimal digits, containing the following key parameters from low to high: Ub0, Ub1, +// Block, DataType, Format, Sparse. Specialized template Ub0, Ub1: +// Represents the axis for UB intra-core splitting, using AxisEnum. Since we allow at most two axes to be split, UB0 and UB1 exist. If there is no UB intra-core splitting, +// fill with AXIS_NONE. UB0 and UB1 each occupy one decimal digit; +// Block: Represents the axis used by UB for multi-core splitting, using AxisEnum, occupies one decimal digit; +// DataType: Represents the input/output data types supported by the current tiling key, using SupportedDtype enum, occupies one decimal digit +// Format: Represents the Format supported by the current tiling key, using InputLayout enum, occupies one decimal digit +// Sparse: Represents whether the current tiling key supports Sparse, using SparseCapability enum, occupies one decimal digit +// For other specialized scenarios, define your own bit fields and values +// usage: get tilingKey from inputted types +// uint64_t tilingKey = GET_FLASHATTENTION_TILINGKEY(AxisEnum::AXIS_S1, AxisEnum::AXIS_S2, AxisEnum::AXIS_N2, +// SupportedDtype::FLOAT32, InputLayout::BSH, SparseCapability::SUPPORT_ALL) + +constexpr uint64_t TILINGKEYOFFSET = uint64_t(10000000000000000000UL); // 10^19 +template constexpr uint64_t GET_TILINGKEY(Args... templateIds) +{ + return TILINGKEYOFFSET + RecursiveSum(templateIds...); +} + +// usage: get tilingKey from inputted types +// uint64_t tilingKey = TILINGKEY(S2, S1, N2, FLOAT32, BSND, ALL) + +#define TILINGKEY(ub2, ub1, block, dtype, layout, sparse) \ + (GET_TILINGKEY(AxisEnum::ub2, AxisEnum::ub1, AxisEnum::block, DtypeEnum::dtype, LayoutEnum::layout, \ + SparseEnum::sparse)) + +} // namespace optiling diff --git a/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_util.h b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_util.h new file mode 100644 index 000000000..fb6ffa2df --- /dev/null +++ b/csrc/ascend/moe/chunk_gated_delta_rule_fwd_h/tiling_base/tiling_util.h @@ -0,0 +1,30 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_util.h + * \brief + */ + +#pragma once + +#include "register/op_impl_registry.h" + +namespace Ops { +namespace Transformer { +namespace OpTiling { +bool IsRegbaseSocVersion(const gert::TilingParseContext* context); + +bool IsRegbaseSocVersion(const gert::TilingContext* context); + +const gert::Shape& EnsureNotScalar(const gert::Shape& inShape); +} // namespace OpTiling +} // namespace Transformer +} // namespace Ops \ No newline at end of file diff --git a/csrc/ascend/moe/common/kernel_utils/block/block_mmad_pingpong_tla_multi.hpp b/csrc/ascend/moe/common/kernel_utils/block/block_mmad_pingpong_tla_multi.hpp new file mode 100644 index 000000000..79e5666b1 --- /dev/null +++ b/csrc/ascend/moe/common/kernel_utils/block/block_mmad_pingpong_tla_multi.hpp @@ -0,0 +1,751 @@ +/** + * Copyright (c) 2025 Tianjin University, Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * the BSD 3-Clause License (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_MULTI_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_MULTI_HPP + +#include "catlass/catlass.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/coord.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/helper.hpp" +#include "catlass/gemm/tile/tile_copy.hpp" +#include "catlass/gemm/tile/tile_mmad.hpp" +#include "tla/layout.hpp" +#include "tla/tensor.hpp" + +namespace Catlass::Gemm { + +template +struct MmadPingpongTlaMulti : public MmadBase { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool USE_HF32_MODE = USE_HF32_MODE_; + static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_; +}; + +} // namespace Catlass::Gemm + +namespace Catlass::Gemm::Block { + +template < + class ArchTag_, + bool ENABLE_UNIT_FLAG_, + bool USE_HF32_MODE_, + uint32_t L0C_STAGES_, + bool ENABLE_L1_RESIDENT_, + uint32_t L1A_STAGES_, + uint32_t L1B_STAGES_, + uint32_t L0A_STAGES_, + uint32_t L0B_STAGES_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_ +> +struct BlockMmadTla < + MmadPingpongTlaMulti, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_ +> { +public: + // Type Aliases + using DispatchPolicy = MmadPingpongTlaMulti; + using ArchTag = typename DispatchPolicy::ArchTag; + using TileCopy = TileCopy_; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + using ElementBias = ElementBias_; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL1ToBT = typename TileCopy::CopyL1ToBT; + + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool USE_HF32_MODE = DispatchPolicy::USE_HF32_MODE; + static constexpr bool ENABLE_L1_RESIDENT = DispatchPolicy::ENABLE_L1_RESIDENT; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + + // Check HF32_MODE + static_assert( + !USE_HF32_MODE || (USE_HF32_MODE && std::is_same_v && std::is_same_v), + "HF32 MODE only supports in float!" + ); + + // Check L0C_STAGES + static_assert(!(ENABLE_UNIT_FLAG && L0C_STAGES != 1), "L0C_STAGES must be 1 when UnitFlag is true!"); + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t _32B = 32*8; // in bits + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); +#if (defined (CATLASS_ARCH) && CATLASS_ARCH == 2201) + static_assert(L1_TILE_M * SizeOfBits::value % _32B == 0, "L1TileShape::M must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_N * SizeOfBits::value % _32B == 0, "L1TileShape::N must be 32B aligned."); + static_assert(L0_TILE_K * SizeOfBits::value % _32B == 0, "L0TileShape::K must be 32B aligned."); +#endif + + static_assert((!HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 8) || (HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 7), + "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); + + static_assert((!HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 8) || (HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 7), + "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); + + static constexpr auto L1A_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1B_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + + // When enabling L1 resident mode, restore the pointer and coordinates that record the last state + // to the initial state. if two blockmmad instances need to be consecutively invoked at the kernel layer, + // RestoreStatus() must be inserted between them. + CATLASS_DEVICE + void RestoreStatus() + { + for (int i = 0; i < L1A_STAGES; ++i) { + lastAddrA[i] = nullptr; + lastCoordA[i] = MatrixCoord{0U, 0U}; + } + for (int i = 0; i < L1B_STAGES; ++i) { + lastAddrB[i] = nullptr; + lastCoordB[i] = MatrixCoord{0U, 0U}; + } + } + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { +#ifdef CATLASS_UNIFIED_CORE + resourcePtr = &resource; + { +#else + if ASCEND_IS_AIC { +#endif + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; + // Init buffers + for (uint32_t i = 0; i < L1A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + // Assign event ID for each stages + l1AEventList[i] = i; + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + // Assign event ID for each stages + l1BEventList[i] = i + L1A_STAGES; + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + // Assign event ID for each stages + l0AEventList[i] = i; + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + // Assign event ID for each stages + l0BEventList[i] = i + L0A_STAGES; + } + if constexpr(!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + } + } else { + l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte(0); + } + if constexpr (HAS_BIAS) { + uint32_t l1BiasOffset = l1BOffset + L1B_TILE_SIZE * L1B_STAGES; + l1BiasTensor = resource.l1Buf.template GetBufferByByte(l1BiasOffset); + l0BiasTensor = resource.btBuf.template GetBufferByByte(0); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() {} + + CATLASS_DEVICE + void preSetFlags() { +#ifdef CATLASS_UNIFIED_CORE + { +#else + if ASCEND_IS_AIC { +#endif + // use HF32 when USE_HF32_MODE is true + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(true); + } else { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(true); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::SetFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::SetFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::SetFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::SetFlag(l0BEventList[i]); + } + if constexpr(!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::SetFlag(l0CEventList[i]); + } + } + if constexpr (HAS_BIAS) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } + + if constexpr (ENABLE_L1_RESIDENT) { + RestoreStatus(); + } + } + } + + CATLASS_DEVICE + void finalWaitFlags() { +#ifdef CATLASS_UNIFIED_CORE + { +#else + if ASCEND_IS_AIC { +#endif + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(false); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr(!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + if constexpr (HAS_BIAS) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + } + } + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, GemmCoord const &actualShape, + TensorBias const &tensorBias = {}) + { + // Check L1TileShape + if constexpr (HAS_BIAS) { + static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); + static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, + "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES + L1BIAS_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } + + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; +#ifdef CATLASS_UNIFIED_CORE + // 310P: no Fixpipe, no DataCopyCO12Dst. L0C exits via DataCopy L0C→UB then UB→GM. +#elif (defined (CATLASS_ARCH) && CATLASS_ARCH == 2201) + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyL0CToGm copyL0CToDst; +#elif (defined (CATLASS_ARCH) && CATLASS_ARCH == 3510) + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; +#endif + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Arch::PositionBias{}); + + uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Arch::PositionL1{}); + auto tensorTileA = GetTileA(tensorA, 0, 0, mBlockActual, kL1Actual); + if constexpr (ENABLE_L1_RESIDENT) { + // If the currently loaded GM pointer and block coordinates are the same as the last loaded ones, + // skip this loading. + if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() + || tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() + || tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListId] = MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr() + ); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(kL1Actual, nBlockActual)); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() + || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() + || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListId] = MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr() + ); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (HAS_BIAS && !std::is_same_v) { + using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; + CopyGmToL1Bias copyGmToL1Bias; + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + copyGmToL1Bias(tensorL1Bias, tensorBias); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + uint32_t mL0Loop = CeilDiv(mL1Actual); + uint32_t nL0Loop = CeilDiv(nL1Actual); + + // main loop + uint32_t kL1Loop = CeilDiv(kBlockActual); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + uint32_t kL1ActualNext{0}; + // preload next tile from GM to L1 + if (kL1Idx < kL1Loop - 1) { + uint32_t kL1IdxNext = kL1Idx + 1; + kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1AListIdNext]; + auto l1BTensor = l1BTensorList[l1BListIdNext]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileA = GetTileA(tensorA, 0, kL1IdxNext * L1_TILE_K, mBlockActual, kL1ActualNext); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), + tla::MakeShape(kL1ActualNext, nBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrA[l1AListIdNext] != tensorTileA.data().GetPhyAddr() + || tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].row() + || tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListIdNext] = + MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr() + ); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListIdNext] != tensorTileB.data().GetPhyAddr() + || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].row() + || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListIdNext] = + MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr() + ); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + auto l1BTensor = l1BTensorList[l1BListId]; + tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kL0Loop = CeilDiv(kL1Actual); + + for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); + + for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K, mL0Actual, kL0Actual); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mL0Idx == 0) && (kL0Idx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); + for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { + uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorL1B, + tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), + tla::MakeShape(kL0Actual, nL0Actual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + if (nL0Idx == 0) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + } + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + auto tensorTileL1Bias = GetTile(tensorL1Bias, + tla::MakeCoord(nL0Idx * L0_TILE_N), + tla::MakeShape(nL0Actual)); + // Load bias to l0 biasTable + copyL1ToBT(tensorL0Bias, tensorTileL1Bias); + if (nL0Idx == nL0Loop - 1) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + } + } + + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, + tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), + tla::MakeShape(mL0Actual, nL0Actual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kL1Idx == kL1Loop - 1) && (mL0Idx == mL0Loop - 1) && + (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tensorL0Bias, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + } + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + } + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + l1AListId = l1AListIdNext; + l1BListId = l1BListIdNext; + kL1Actual = kL1ActualNext; + } + + // copy block out +#ifdef CATLASS_UNIFIED_CORE + { + // 310P unified core: L0C→UB via DataCopy, then UB→GM. + // No Fixpipe or DataCopyCO12Dst on dav_m200. + uint32_t mAligned = (mBlockActual + 15) / 16 * 16; + uint32_t nAligned = (nBlockActual + 15) / 16 * 16; + uint32_t tileElems = mAligned * nAligned; + uint32_t tileBytes = tileElems * sizeof(ElementAccumulator); + + // UB temp for L0C→UB transfer. Offset 0 is safe: on unified core, + // mmad and epilogue run sequentially so UB is not shared concurrently. + // The epilogue allocates its own UB regions at higher offsets (≥32KB). + AscendC::LocalTensor co2Temp = + resourcePtr->ubBuf.template GetBufferByByte(0); + + AscendC::PipeBarrier(); + + // L0C → UB: BLOCK_MODE_MATRIX copies raw NZ fractals to UB + // For float: blockLen unit = 1024B (one 16×16 fractal) + AscendC::DataCopyParams l0c2ubParams; + l0c2ubParams.blockCount = static_cast(nAligned / 16); + l0c2ubParams.blockLen = static_cast(mAligned / 16); + l0c2ubParams.srcStride = 0; + l0c2ubParams.dstStride = 0; + AscendC::DataCopyEnhancedParams enhParams; + enhParams.blockMode = AscendC::BlockMode::BLOCK_MODE_MATRIX; + AscendC::DataCopy(co2Temp, l0CTensorList[l0CListId], l0c2ubParams, enhParams); + AscendC::PipeBarrier(); + + // UB → GM: fractal-by-fractal with strided DataCopy (NZ→ND deformat) + // NZ in UB: [N/16 Z-cols][M/16 fractals][16 rows][16 cols] + // ND in GM: [M rows][N cols] + auto dstOffset = tensorC.layout()(tensorC.coord()); + uint32_t gmStride = tla::get<0>(tensorC.stride()); + uint32_t mFracs = mAligned / 16; + uint32_t nFracs = nAligned / 16; + for (uint32_t nf = 0; nf < nFracs; nf++) { + for (uint32_t mf = 0; mf < mFracs; mf++) { + uint32_t ubOff = (nf * mFracs + mf) * 256; + uint32_t gmRow = mf * 16; + uint32_t gmCol = nf * 16; + uint32_t gmOff = dstOffset + gmRow * gmStride + gmCol; + AscendC::DataCopyParams fracParams; + fracParams.blockCount = 16; + fracParams.blockLen = static_cast(16 * sizeof(ElementAccumulator) / 32); + fracParams.srcStride = 0; + fracParams.dstStride = static_cast((gmStride - 16) * sizeof(ElementAccumulator) / 32); + AscendC::DataCopy(tensorC.data()[gmOff], co2Temp[ubOff], fracParams); + } + } + AscendC::PipeBarrier(); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } +#else + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToDst(tensorC, tensorL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + copyL0CToDst(tensorC, tensorL0C, 0b11); + } +#endif + } + +protected: + template + CATLASS_DEVICE auto GetTileA(TensorA &tensorA, uint32_t mIndex, uint32_t kIndex, uint32_t mSize, uint32_t kSize) + { + if constexpr(tla::detail::isVector::value) { + return GetTile(tensorA, tla::MakeCoord(kIndex), tla::MakeShape(kSize)); + } else { + return GetTile(tensorA, tla::MakeCoord(mIndex, kIndex), tla::MakeShape(mSize, kSize)); + } + } + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + AscendC::LocalTensor l0CTensorList[L0C_STAGES]; + AscendC::LocalTensor l1BiasTensor; + AscendC::LocalTensor l0BiasTensor; +#ifdef CATLASS_UNIFIED_CORE + Arch::Resource* resourcePtr{nullptr}; +#endif + + // Multi-stage event id list + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + int32_t l0CEventList[L0C_STAGES]; + + __gm__ typename AscendC::GlobalTensor::PrimType* lastAddrA[L1A_STAGES]; + __gm__ typename AscendC::GlobalTensor::PrimType* lastAddrB[L1B_STAGES]; + MatrixCoord lastCoordA[L1A_STAGES]; + MatrixCoord lastCoordB[L1B_STAGES]; + + // The id of current stage + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + uint32_t l0CListId{0}; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL1ToBT copyL1ToBT; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_MULTI_HPP diff --git a/csrc/ascend/moe/copy_and_expand_eagle_inputs/CMakeLists.txt b/csrc/ascend/moe/copy_and_expand_eagle_inputs/CMakeLists.txt new file mode 100644 index 000000000..298ed5677 --- /dev/null +++ b/csrc/ascend/moe/copy_and_expand_eagle_inputs/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/CMakeLists.txt b/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/CMakeLists.txt new file mode 100644 index 000000000..f0c631544 --- /dev/null +++ b/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/CMakeLists.txt @@ -0,0 +1,22 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + copy_and_expand_eagle_inputs_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME CopyAndExpandEagleInputs + OPTIONS + --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE copy_and_expand_eagle_inputs ACLNNTYPE aclnn) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_def.cpp b/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_def.cpp new file mode 100644 index 000000000..bb10b672c --- /dev/null +++ b/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_def.cpp @@ -0,0 +1,88 @@ +/** + * @file copy_and_expand_eagle_inputs_def.cpp + * @brief CopyAndExpandEagleInputs OpDef registration + */ + +#include "register/op_def_registry.h" + +namespace ops { + +class CopyAndExpandEagleInputs : public OpDef { +public: + explicit CopyAndExpandEagleInputs(const char* name) : OpDef(name) + { + // -------------------- Inputs -------------------- + this->Input("target_token_ids") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("target_positions") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("next_token_ids") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("query_start_loc") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("query_end_loc") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + + // -------------------- Outputs -------------------- + this->Output("out_input_ids") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("out_positions") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("out_is_rejected_token_mask") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("out_is_masked_token_mask") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("out_new_token_indices") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("out_hidden_state_mapping") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + + // -------------------- Attributes -------------------- + this->Attr("padding_token_id").Int(); + this->Attr("parallel_drafting_token_id").Int(); + this->Attr("num_padding_slots_per_request").Int(); + this->Attr("shift_input_ids").Bool(); + this->Attr("total_input_tokens").Int(); + + // -------------------- Platform -------------------- + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + } +}; + +OP_ADD(CopyAndExpandEagleInputs); + +} // namespace ops diff --git a/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_infershape.cpp b/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_infershape.cpp new file mode 100644 index 000000000..f2447ea19 --- /dev/null +++ b/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_infershape.cpp @@ -0,0 +1,111 @@ +/** + * @file copy_and_expand_eagle_inputs_infershape.cpp + * @brief InferShape and InferDataType for CopyAndExpandEagleInputs + */ + +#include "register/op_def_registry.h" +#include "log/ops_log.h" + +#define unlikely(x) __builtin_expect((x), 0) +#define OP_CHECK_NULL_WITH_CONTEXT(context, ptr) \ + do { \ + if (unlikely((ptr) == nullptr)) { \ + const char* name = (unlikely(((context) == nullptr) || (context)->GetNodeName() == nullptr)) ? \ + "nil" : \ + (context)->GetNodeName(); \ + OPS_LOG_E(name, "%s is nullptr!", #ptr); \ + return ge::GRAPH_FAILED; \ + } \ + } while (0) + +static constexpr int IDX_TARGET_TOKEN_IDS = 0; +static constexpr int IDX_TARGET_POSITIONS = 1; +static constexpr int IDX_NEXT_TOKEN_IDS = 2; +static constexpr int IDX_QUERY_START_LOC = 3; +static constexpr int IDX_QUERY_END_LOC = 4; + +static constexpr int OUT_INPUT_IDS = 0; +static constexpr int OUT_POSITIONS = 1; +static constexpr int OUT_REJECTED_MASK = 2; +static constexpr int OUT_MASKED_MASK = 3; +static constexpr int OUT_NEW_TOKEN_INDICES = 4; +static constexpr int OUT_HIDDEN_STATE_MAPPING = 5; +static constexpr int OUTPUT_NUM = 6; + +static constexpr int ATTR_NUM_PADDING_SLOTS = 2; +static constexpr int ATTR_SHIFT_INPUT_IDS = 3; + +using namespace ge; + +namespace ops { + +static ge::graphStatus InferShape4CopyAndExpandEagleInputs(gert::InferShapeContext* context) +{ + // Get input shapes + const gert::Shape* targetTokenIdsShape = context->GetInputShape(IDX_TARGET_TOKEN_IDS); + OP_CHECK_NULL_WITH_CONTEXT(context, targetTokenIdsShape); + const gert::Shape* queryStartLocShape = context->GetInputShape(IDX_QUERY_START_LOC); + OP_CHECK_NULL_WITH_CONTEXT(context, queryStartLocShape); + + // Derive dimensions from input shapes + int64_t totalInputTokens = targetTokenIdsShape->GetDim(0); + int64_t numReqs = queryStartLocShape->GetDim(0) - 1; + + // Get num_padding_slots_per_request and shift_input_ids from attributes. + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + int64_t numPaddingSlotsPerReq = *(attrs->GetAttrPointer(ATTR_NUM_PADDING_SLOTS)); + bool shiftInputIds = *(attrs->GetAttrPointer(ATTR_SHIFT_INPUT_IDS)); + + // total_input_tokens already includes rejected tokens. The only delta between + // shift=false and shift=true is whether each request keeps its first accepted + // token in the draft sequence. + int64_t totalDraftTokens = totalInputTokens + + (numPaddingSlotsPerReq - (shiftInputIds ? 1 : 0)) * numReqs; + + // Get and validate all output shapes + gert::Shape* outShapes[OUTPUT_NUM]; + for (int i = 0; i < OUTPUT_NUM; ++i) { + outShapes[i] = context->GetOutputShape(i); + OP_CHECK_NULL_WITH_CONTEXT(context, outShapes[i]); + outShapes[i]->SetDimNum(1); + } + + // out_input_ids, out_positions, out_rejected_mask, out_masked_mask: [total_draft_tokens] + outShapes[OUT_INPUT_IDS]->SetDim(0, totalDraftTokens); + outShapes[OUT_POSITIONS]->SetDim(0, totalDraftTokens); + outShapes[OUT_REJECTED_MASK]->SetDim(0, totalDraftTokens); + outShapes[OUT_MASKED_MASK]->SetDim(0, totalDraftTokens); + + // out_new_token_indices: [num_reqs * num_padding_slots_per_request] + outShapes[OUT_NEW_TOKEN_INDICES]->SetDim(0, numReqs * numPaddingSlotsPerReq); + + // out_hidden_state_mapping: [total_input_tokens] + outShapes[OUT_HIDDEN_STATE_MAPPING]->SetDim(0, totalInputTokens); + + return GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataType4CopyAndExpandEagleInputs(gert::InferDataTypeContext* context) +{ + // out_input_ids: INT32 + context->SetOutputDataType(OUT_INPUT_IDS, DT_INT32); + // out_positions: INT32 + context->SetOutputDataType(OUT_POSITIONS, DT_INT32); + // out_is_rejected_token_mask: INT8 + context->SetOutputDataType(OUT_REJECTED_MASK, DT_INT8); + // out_is_masked_token_mask: INT8 + context->SetOutputDataType(OUT_MASKED_MASK, DT_INT8); + // out_new_token_indices: INT32 + context->SetOutputDataType(OUT_NEW_TOKEN_INDICES, DT_INT32); + // out_hidden_state_mapping: INT32 + context->SetOutputDataType(OUT_HIDDEN_STATE_MAPPING, DT_INT32); + + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(CopyAndExpandEagleInputs) + .InferShape(InferShape4CopyAndExpandEagleInputs) + .InferDataType(InferDataType4CopyAndExpandEagleInputs); + +} // namespace ops diff --git a/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_tiling.cpp b/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_tiling.cpp new file mode 100644 index 000000000..f080d1757 --- /dev/null +++ b/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_tiling.cpp @@ -0,0 +1,123 @@ +/** + * @file copy_and_expand_eagle_inputs_tiling.cpp + * @brief CopyAndExpandEagleInputs TilingFunc implementation + */ + +#include "copy_and_expand_eagle_inputs_tiling.h" +#include "tiling_base/error_log.h" +#include "register/op_def_registry.h" +#include "log/ops_log.h" + +#include + +namespace optiling { + +static void GetCompileParameters( + gert::TilingContext* context, uint32_t& coreNum) +{ + auto ptrCompileInfo = reinterpret_cast(context->GetCompileInfo()); + if (ptrCompileInfo == nullptr) { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + coreNum = ascendcPlatform.GetCoreNum(); + } else { + coreNum = ptrCompileInfo->totalCoreNum; + } +} + +static ge::graphStatus TilingFunc(gert::TilingContext* context) +{ + OPS_LOG_I(context, "Enter TilingFunc for CopyAndExpandEagleInputs"); + OPS_LOG_D(context, "TilingFunc running."); + + // ========== 1. Get hardware core count ========== + uint32_t coreNum; + GetCompileParameters(context, coreNum); + + // ========== 2. Derive num_reqs from query_start_loc shape ========== + // query_start_loc is the 4th input (index 3), shape [num_reqs + 1] + auto queryStartLocShape = context->GetInputShape(3); + uint32_t numReqs = 0; + if (queryStartLocShape != nullptr && + queryStartLocShape->GetStorageShape().GetDimNum() > 0) { + int64_t dim0 = queryStartLocShape->GetStorageShape().GetDim(0); + numReqs = (dim0 > 1) ? static_cast(dim0 - 1) : 0; + } + + // ========== 3. Get operator attributes ========== + auto attrs = context->GetAttrs(); + + int32_t paddingTokenId = *(attrs->GetAttrPointer(0)); + int32_t parallelDraftingTokenId = *(attrs->GetAttrPointer(1)); + int32_t numPaddingSlotsPerReq = *(attrs->GetAttrPointer(2)); + bool shiftInputIds = *(attrs->GetAttrPointer(3)); + int32_t totalInputTokens = *(attrs->GetAttrPointer(4)); + + // ========== 4. Compute core distribution ========== + // Adjacent requests do not guarantee block-aligned output boundaries. + // The runtime write path emits block-granular stores, so request-parallel + // execution can cause neighboring segments to overlap at block tails. + // Run this op on a single core to preserve correctness. + uint32_t usedCoreNum = 1; + uint32_t reqsPerCore = numReqs / usedCoreNum; + uint32_t remainderReqs = numReqs % usedCoreNum; + + // ========== 5. Set tiling_key ========== + context->SetTilingKey(1); + + // ========== 6. Get output shape ========== + uint32_t totalDraftTokens = 0; + auto outShape = context->GetOutputShape(0); + if (outShape != nullptr && + outShape->GetStorageShape().GetDimNum() > 0) { + totalDraftTokens = static_cast(outShape->GetStorageShape().GetDim(0)); + } + + // ========== 7. Fill TilingData ========== + CopyAndExpandEagleInputsTilingData tiling; + tiling.set_usedCoreNum(usedCoreNum); + tiling.set_numReqs(numReqs); + tiling.set_reqsPerCore(reqsPerCore); + tiling.set_remainderReqs(remainderReqs); + tiling.set_paddingTokenId(paddingTokenId); + tiling.set_parallelDraftingTokenId(parallelDraftingTokenId); + tiling.set_numPaddingSlotsPerReq(static_cast(numPaddingSlotsPerReq)); + tiling.set_totalInputTokens(static_cast(totalInputTokens)); + tiling.set_shiftInputIds(shiftInputIds ? 1u : 0u); + tiling.set_totalDraftTokens(totalDraftTokens); + + tiling.SaveToBuffer( + context->GetRawTilingData()->GetData(), + context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); + + // ========== 8. Set block_dim ========== + context->SetBlockDim(usedCoreNum); + + OPS_LOG_I(context, "Block Dim: %u", usedCoreNum); + OPS_LOG_I(context, + "numReqs: %u, reqsPerCore: %u, remainderReqs: %u, totalInputTokens: %d, totalDraftTokens: %u", + numReqs, reqsPerCore, remainderReqs, totalInputTokens, totalDraftTokens); + + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus TilingPrepare4CopyAndExpandEagleInputs(gert::TilingParseContext* context) +{ + OPS_LOG_D(context, "TilingPrepare4CopyAndExpandEagleInputs running."); + OPS_LOG_I(context, "TilingPrepare4CopyAndExpandEagleInputs running."); + auto compileInfo = context->GetCompiledInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo); + auto platformInfo = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + + compileInfo->totalCoreNum = ascendcPlatform.GetCoreNum(); + + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(CopyAndExpandEagleInputs) + .Tiling(TilingFunc) + .TilingParse(TilingPrepare4CopyAndExpandEagleInputs); + +} // namespace optiling diff --git a/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_tiling.h b/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_tiling.h new file mode 100644 index 000000000..2975e9c42 --- /dev/null +++ b/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_host/copy_and_expand_eagle_inputs_tiling.h @@ -0,0 +1,37 @@ +#ifndef COPY_AND_EXPAND_EAGLE_INPUTS_TILING_H +#define COPY_AND_EXPAND_EAGLE_INPUTS_TILING_H + +#include "register/tilingdata_base.h" +#include "tiling_base/error_log.h" +#include "register/op_impl_registry.h" +#include "tiling/platform/platform_ascendc.h" + +namespace optiling { + +BEGIN_TILING_DATA_DEF(CopyAndExpandEagleInputsTilingData) + // ---- 分核参数 ---- + TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum); // 实际使用的核数 + TILING_DATA_FIELD_DEF(uint32_t, numReqs); // 总请求数 + TILING_DATA_FIELD_DEF(uint32_t, reqsPerCore); // 每核基础请求数 + TILING_DATA_FIELD_DEF(uint32_t, remainderReqs); // 余数（前 remainder 个核多处理 1 个请求） + + // ---- 算子属性 ---- + TILING_DATA_FIELD_DEF(int32_t, paddingTokenId); // 填充 token ID + TILING_DATA_FIELD_DEF(int32_t, parallelDraftingTokenId); // 并行推测解码 token ID + TILING_DATA_FIELD_DEF(uint32_t, numPaddingSlotsPerReq); // 每个请求的 padding 槽位数 + TILING_DATA_FIELD_DEF(uint32_t, totalInputTokens); // 输入 token 总数（用于 clamp） + TILING_DATA_FIELD_DEF(uint32_t, shiftInputIds); // 0 = false, 1 = true + + // ---- 输出尺寸 ---- + TILING_DATA_FIELD_DEF(uint32_t, totalDraftTokens); // 输出 token 总数 +END_TILING_DATA_DEF; + +struct CopyAndExpandEagleInputsCompileInfo { + uint32_t totalCoreNum = 0; +}; + +REGISTER_TILING_DATA_CLASS(CopyAndExpandEagleInputs, CopyAndExpandEagleInputsTilingData) + +} // namespace optiling + +#endif // COPY_AND_EXPAND_EAGLE_INPUTS_TILING_H diff --git a/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_kernel/copy_and_expand_eagle_inputs.cpp b/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_kernel/copy_and_expand_eagle_inputs.cpp new file mode 100644 index 000000000..e761c2359 --- /dev/null +++ b/csrc/ascend/moe/copy_and_expand_eagle_inputs/op_kernel/copy_and_expand_eagle_inputs.cpp @@ -0,0 +1,388 @@ +#include "kernel_operator.h" + +using namespace AscendC; + +// ONE_BLK_SIZE comes from AscendC namespace (32 bytes per block) + +class CopyAndExpandEagleInputsKernel { +public: + __aicore__ inline CopyAndExpandEagleInputsKernel() {} + + __aicore__ inline void Init(GM_ADDR targetTokenIds, GM_ADDR targetPositions, + GM_ADDR nextTokenIds, GM_ADDR queryStartLoc, + GM_ADDR queryEndLoc, + GM_ADDR outInputIds, GM_ADDR outPositions, + GM_ADDR outIsRejectedTokenMask, GM_ADDR outIsMaskedTokenMask, + GM_ADDR outNewTokenIndices, GM_ADDR outHiddenStateMapping, + const CopyAndExpandEagleInputsTilingData* tilingData) + { + usedCoreNum = tilingData->usedCoreNum; + numReqs = tilingData->numReqs; + reqsPerCore = tilingData->reqsPerCore; + remainderReqs = tilingData->remainderReqs; + paddingTokenId = tilingData->paddingTokenId; + parallelDraftingTokenId = tilingData->parallelDraftingTokenId; + numPaddingSlotsPerReq = tilingData->numPaddingSlotsPerReq; + totalInputTokens = tilingData->totalInputTokens; + totalDraftTokens = tilingData->totalDraftTokens; + + uint32_t coreId = GetBlockIdx(); + if (coreId < remainderReqs) { + myStartReq = coreId * (reqsPerCore + 1); + myNumReqs = reqsPerCore + 1; + } else { + myStartReq = remainderReqs * (reqsPerCore + 1) + (coreId - remainderReqs) * reqsPerCore; + myNumReqs = reqsPerCore; + } + + gmTargetTokenIds.SetGlobalBuffer((__gm__ int32_t*)targetTokenIds, totalInputTokens); + gmTargetPositions.SetGlobalBuffer((__gm__ int32_t*)targetPositions, totalInputTokens); + gmNextTokenIds.SetGlobalBuffer((__gm__ int32_t*)nextTokenIds, numReqs); + gmQueryStartLoc.SetGlobalBuffer((__gm__ int32_t*)queryStartLoc, numReqs + 1); + gmQueryEndLoc.SetGlobalBuffer((__gm__ int32_t*)queryEndLoc, numReqs); + gmOutInputIds.SetGlobalBuffer((__gm__ int32_t*)outInputIds, totalDraftTokens); + gmOutPositions.SetGlobalBuffer((__gm__ int32_t*)outPositions, totalDraftTokens); + gmOutIsRejectedTokenMask.SetGlobalBuffer((__gm__ int8_t*)outIsRejectedTokenMask, totalDraftTokens); + gmOutIsMaskedTokenMask.SetGlobalBuffer((__gm__ int8_t*)outIsMaskedTokenMask, totalDraftTokens); + gmOutNewTokenIndices.SetGlobalBuffer((__gm__ int32_t*)outNewTokenIndices, numPaddingSlotsPerReq * numReqs); + gmOutHiddenStateMapping.SetGlobalBuffer((__gm__ int32_t*)outHiddenStateMapping, totalInputTokens); + + uint32_t metaAligned = AlignUp((myNumReqs + 1) * sizeof(int32_t), ONE_BLK_SIZE); + pipe.InitBuffer(qsBuf, metaAligned); + pipe.InitBuffer(qeBuf, AlignUp(myNumReqs * sizeof(int32_t), ONE_BLK_SIZE)); + pipe.InitBuffer(ntBuf, AlignUp(myNumReqs * sizeof(int32_t), ONE_BLK_SIZE)); + + constexpr uint32_t MAX_PER_REQ = 4096; + pipe.InitBuffer(inputBuf, AlignUp(MAX_PER_REQ * sizeof(int32_t), ONE_BLK_SIZE)); + pipe.InitBuffer(outIdsBuf, AlignUp(MAX_PER_REQ * sizeof(int32_t), ONE_BLK_SIZE)); + pipe.InitBuffer(outPosBuf, AlignUp(MAX_PER_REQ * sizeof(int32_t), ONE_BLK_SIZE)); + pipe.InitBuffer(outRejBuf, AlignUp(MAX_PER_REQ * sizeof(int8_t), ONE_BLK_SIZE)); + pipe.InitBuffer(outMskBuf, AlignUp(MAX_PER_REQ * sizeof(int8_t), ONE_BLK_SIZE)); + pipe.InitBuffer(ntiBuf, AlignUp(64 * sizeof(int32_t), ONE_BLK_SIZE)); + pipe.InitBuffer(hsmBuf, AlignUp(MAX_PER_REQ * sizeof(int32_t), ONE_BLK_SIZE)); + + if (myNumReqs > 0) { + LocalTensor lqs = qsBuf.Get(); + DataCopyIn(lqs, gmQueryStartLoc, (int32_t)myStartReq, (int32_t)(myNumReqs + 1)); + + LocalTensor lqe = qeBuf.Get(); + DataCopyIn(lqe, gmQueryEndLoc, (int32_t)myStartReq, (int32_t)myNumReqs); + + LocalTensor lnt = ntBuf.Get(); + DataCopyIn(lnt, gmNextTokenIds, (int32_t)myStartReq, (int32_t)myNumReqs); + } + } + + __aicore__ inline void ProcessShiftFalse() + { + for (uint32_t rLocal = 0; rLocal < myNumReqs; rLocal++) { + ProcessOneRequestShiftFalse(myStartReq + rLocal, rLocal); + } + } + + __aicore__ inline void ProcessShiftTrue() + { + for (uint32_t rLocal = 0; rLocal < myNumReqs; rLocal++) { + ProcessOneRequestShiftTrue(myStartReq + rLocal, rLocal); + } + + if (myStartReq + myNumReqs == numReqs && totalDraftTokens > 0) { + LocalTensor lqs = qsBuf.Get(); + LocalTensor lqe = qeBuf.Get(); + uint32_t lastLocal = myNumReqs - 1; + int32_t lastQs = lqs.GetValue(lastLocal); + int32_t lastNqs = lqs.GetValue(lastLocal + 1); + int32_t lastQe = lqe.GetValue(lastLocal); + int32_t lastNumValid = lastQe - lastQs; + if (lastNumValid < 0) lastNumValid = 0; + int32_t lastNumRejected = lastNqs - lastQe - 1; + if (lastNumRejected < 0) lastNumRejected = 0; + int32_t lastOutputLen = lastNumValid + (int32_t)numPaddingSlotsPerReq + lastNumRejected; + int32_t lastOutputStart = lastQs + (int32_t)(myStartReq + lastLocal) * ((int32_t)numPaddingSlotsPerReq - 1); + int32_t writtenEnd = lastOutputStart + lastOutputLen; + int32_t tailCount = (int32_t)totalDraftTokens - writtenEnd; + if (tailCount > 0) { + LocalTensor lZero32 = outIdsBuf.Get(); + LocalTensor lZero8 = outRejBuf.Get(); + for (int32_t j = 0; j < tailCount; j++) { + lZero32.SetValue(j, (int32_t)0); + lZero8.SetValue(j, (int8_t)0); + } + DataCopyOut_int32(gmOutInputIds, lZero32, writtenEnd, tailCount); + DataCopyOut_int32(gmOutPositions, lZero32, writtenEnd, tailCount); + DataCopyOut_int8(gmOutIsRejectedTokenMask, lZero8, writtenEnd, tailCount); + DataCopyOut_int8(gmOutIsMaskedTokenMask, lZero8, writtenEnd, tailCount); + } + } + } + +private: + static __aicore__ inline uint32_t AlignUp(uint32_t x, uint32_t a) + { + return (x + a - 1) / a * a; + } + + __aicore__ inline void DataCopyIn(LocalTensor& dst, + GlobalTensor& src, + int32_t gmOffset, int32_t count) + { + if (count <= 0) return; + constexpr int32_t ELEMS_PER_BLK = ONE_BLK_SIZE / (int32_t)sizeof(int32_t); // 8 + int32_t aligned = (count + ELEMS_PER_BLK - 1) / ELEMS_PER_BLK * ELEMS_PER_BLK; + DataCopy(dst, src[gmOffset], aligned); + pipe_barrier(PIPE_ALL); + } + + __aicore__ inline void DataCopyOut_int32(GlobalTensor& dst, + LocalTensor& src, + int32_t gmOffset, int32_t count) + { + if (count <= 0) return; + uint32_t totalBytes = static_cast(count) * static_cast(sizeof(int32_t)); + pipe_barrier(PIPE_ALL); + DataCopyPad(dst[gmOffset], src, DataCopyExtParams(1, totalBytes, 0, 0, 0)); + pipe_barrier(PIPE_ALL); + } + + __aicore__ inline void DataCopyOut_int8(GlobalTensor& dst, + LocalTensor& src, + int32_t gmOffset, int32_t count) + { + if (count <= 0) return; + uint32_t totalBytes = static_cast(count) * static_cast(sizeof(int8_t)); + pipe_barrier(PIPE_ALL); + DataCopyPad(dst[gmOffset], src, DataCopyExtParams(1, totalBytes, 0, 0, 0)); + pipe_barrier(PIPE_ALL); + } + + __aicore__ inline int32_t ReadQS(uint32_t rLocal) { + return qsBuf.Get().GetValue(rLocal); + } + __aicore__ inline int32_t ReadNextQS(uint32_t rLocal) { + return qsBuf.Get().GetValue(rLocal + 1); + } + __aicore__ inline int32_t ReadQE(uint32_t rLocal) { + return qeBuf.Get().GetValue(rLocal); + } + __aicore__ inline int32_t ReadNT(uint32_t rLocal) { + return ntBuf.Get().GetValue(rLocal); + } + + // ============================================================ + // shift_input_ids = false + // ============================================================ + __aicore__ inline void ProcessOneRequestShiftFalse(uint32_t r, uint32_t rLocal) + { + int32_t queryStart = ReadQS(rLocal); + int32_t nextQueryStart = ReadNextQS(rLocal); + int32_t queryEnd = ReadQE(rLocal); + + int32_t numRejected = nextQueryStart - queryEnd - 1; + if (numRejected < 0) numRejected = 0; + int32_t numValid = queryEnd - queryStart + 1; + if (numValid < 0) numValid = 0; + + int32_t outputStart = queryStart + (int32_t)r * (int32_t)numPaddingSlotsPerReq; + int32_t outputLen = numValid + (int32_t)numPaddingSlotsPerReq + numRejected; + + int32_t numInputTokensForReq = nextQueryStart - queryStart; + LocalTensor localInput = inputBuf.Get(); + if (numInputTokensForReq > 0) { + DataCopyIn(localInput, gmTargetTokenIds, queryStart, numInputTokensForReq); + } + + LocalTensor localTmpPos = hsmBuf.Get(); + DataCopyIn(localTmpPos, gmTargetPositions, queryStart, 1); + int32_t startPos = localTmpPos.GetValue(0); + + int32_t nextTokenId = ReadNT(rLocal); + + LocalTensor lIds = outIdsBuf.Get(); + LocalTensor lPos = outPosBuf.Get(); + LocalTensor lRej = outRejBuf.Get(); + LocalTensor lMsk = outMskBuf.Get(); + + for (int32_t j = 0; j < numValid; j++) { + int32_t inIdx = j; + if (inIdx >= numInputTokensForReq) inIdx = numInputTokensForReq - 1; + lIds.SetValue(j, localInput.GetValue(inIdx)); + lPos.SetValue(j, startPos + j); + lRej.SetValue(j, (int8_t)0); + lMsk.SetValue(j, (int8_t)0); + } + // Bonus + lIds.SetValue(numValid, nextTokenId); + lPos.SetValue(numValid, startPos + numValid); + lRej.SetValue(numValid, (int8_t)0); + lMsk.SetValue(numValid, (int8_t)0); + // Parallel Draft + for (int32_t k = 1; k < (int32_t)numPaddingSlotsPerReq; k++) { + int32_t j = numValid + k; + lIds.SetValue(j, parallelDraftingTokenId); + lPos.SetValue(j, startPos + j); + lRej.SetValue(j, (int8_t)0); + lMsk.SetValue(j, (int8_t)1); + } + // Rejected + for (int32_t k = 0; k < numRejected; k++) { + int32_t j = numValid + (int32_t)numPaddingSlotsPerReq + k; + lIds.SetValue(j, paddingTokenId); + lPos.SetValue(j, (int32_t)0); + lRej.SetValue(j, (int8_t)1); + lMsk.SetValue(j, (int8_t)0); + } + + // UB → GM + DataCopyOut_int32(gmOutInputIds, lIds, outputStart, outputLen); + DataCopyOut_int32(gmOutPositions, lPos, outputStart, outputLen); + DataCopyOut_int8(gmOutIsRejectedTokenMask, lRej, outputStart, outputLen); + DataCopyOut_int8(gmOutIsMaskedTokenMask, lMsk, outputStart, outputLen); + + // NTI + LocalTensor lNti = ntiBuf.Get(); + lNti.SetValue(0, outputStart + numValid); + for (int32_t k = 1; k < (int32_t)numPaddingSlotsPerReq; k++) { + lNti.SetValue(k, outputStart + numValid + k); + } + int32_t ntiOff = (int32_t)r * (int32_t)numPaddingSlotsPerReq; + DataCopyOut_int32(gmOutNewTokenIndices, lNti, ntiOff, (int32_t)numPaddingSlotsPerReq); + } + + // ============================================================ + // shift_input_ids = true + // ============================================================ + __aicore__ inline void ProcessOneRequestShiftTrue(uint32_t r, uint32_t rLocal) + { + int32_t queryStart = ReadQS(rLocal); + int32_t nextQueryStart = ReadNextQS(rLocal); + int32_t queryEnd = ReadQE(rLocal); + + int32_t numRejected = nextQueryStart - queryEnd - 1; + if (numRejected < 0) numRejected = 0; + int32_t numValid = queryEnd - queryStart; + if (numValid < 0) numValid = 0; + + int32_t outputStart = queryStart + (int32_t)r * ((int32_t)numPaddingSlotsPerReq - 1); + int32_t outputLen = numValid + (int32_t)numPaddingSlotsPerReq + numRejected; + + int32_t numInputTokensForReq = nextQueryStart - queryStart; + LocalTensor localInput = inputBuf.Get(); + int32_t readStart = queryStart + 1; + int32_t readCount = numValid; + if (readStart + readCount > (int32_t)totalInputTokens) { + readCount = (int32_t)totalInputTokens - readStart; + if (readCount < 0) readCount = 0; + } + if (readCount > 0) { + DataCopyIn(localInput, gmTargetTokenIds, readStart, readCount); + } + + LocalTensor localTmpPos = hsmBuf.Get(); + DataCopyIn(localTmpPos, gmTargetPositions, queryStart, 1); + int32_t startPos = localTmpPos.GetValue(0); + + int32_t nextTokenId = ReadNT(rLocal); + + LocalTensor lIds = outIdsBuf.Get(); + LocalTensor lPos = outPosBuf.Get(); + LocalTensor lRej = outRejBuf.Get(); + LocalTensor lMsk = outMskBuf.Get(); + + for (int32_t j = 0; j < numValid; j++) { + int32_t inIdx = j; + if (inIdx >= readCount && readCount > 0) inIdx = readCount - 1; + lIds.SetValue(j, readCount > 0 ? localInput.GetValue(inIdx) : (int32_t)0); + lPos.SetValue(j, startPos + j); + lRej.SetValue(j, (int8_t)0); + lMsk.SetValue(j, (int8_t)0); + } + lIds.SetValue(numValid, nextTokenId); + lPos.SetValue(numValid, startPos + numValid); + lRej.SetValue(numValid, (int8_t)0); + lMsk.SetValue(numValid, (int8_t)0); + for (int32_t k = 1; k < (int32_t)numPaddingSlotsPerReq; k++) { + int32_t j = numValid + k; + lIds.SetValue(j, parallelDraftingTokenId); + lPos.SetValue(j, startPos + j); + lRej.SetValue(j, (int8_t)0); + lMsk.SetValue(j, (int8_t)1); + } + for (int32_t k = 0; k < numRejected; k++) { + int32_t j = numValid + (int32_t)numPaddingSlotsPerReq + k; + lIds.SetValue(j, paddingTokenId); + lPos.SetValue(j, (int32_t)0); + lRej.SetValue(j, (int8_t)1); + lMsk.SetValue(j, (int8_t)0); + } + + int32_t outputLenToWrite = outputLen; + if (nextQueryStart == queryEnd && (r + 1) < numReqs) { + outputLenToWrite = outputLen - 1; + } + DataCopyOut_int32(gmOutInputIds, lIds, outputStart, outputLenToWrite); + DataCopyOut_int32(gmOutPositions, lPos, outputStart, outputLenToWrite); + DataCopyOut_int8(gmOutIsRejectedTokenMask, lRej, outputStart, outputLenToWrite); + DataCopyOut_int8(gmOutIsMaskedTokenMask, lMsk, outputStart, outputLenToWrite); + + LocalTensor lNti = ntiBuf.Get(); + lNti.SetValue(0, outputStart + numValid); + for (int32_t k = 1; k < (int32_t)numPaddingSlotsPerReq; k++) { + lNti.SetValue(k, outputStart + numValid + k); + } + int32_t ntiOff = (int32_t)r * (int32_t)numPaddingSlotsPerReq; + DataCopyOut_int32(gmOutNewTokenIndices, lNti, ntiOff, (int32_t)numPaddingSlotsPerReq); + + // hidden_state_mapping + LocalTensor lHsm = hsmBuf.Get(); + for (int32_t j = 0; j < numInputTokensForReq; j++) { + lHsm.SetValue(j, outputStart + j); + } + DataCopyOut_int32(gmOutHiddenStateMapping, lHsm, queryStart, numInputTokensForReq); + } + +private: + GlobalTensor gmTargetTokenIds, gmTargetPositions, gmNextTokenIds; + GlobalTensor gmQueryStartLoc, gmQueryEndLoc; + GlobalTensor gmOutInputIds, gmOutPositions; + GlobalTensor gmOutIsRejectedTokenMask, gmOutIsMaskedTokenMask; + GlobalTensor gmOutNewTokenIndices, gmOutHiddenStateMapping; + + uint32_t usedCoreNum, numReqs, reqsPerCore, remainderReqs; + int32_t paddingTokenId, parallelDraftingTokenId; + uint32_t numPaddingSlotsPerReq, totalInputTokens, totalDraftTokens; + uint32_t myStartReq, myNumReqs; + + TPipe pipe; + TBuf qsBuf, qeBuf, ntBuf; + TBuf inputBuf, outIdsBuf, outPosBuf; + TBuf outRejBuf, outMskBuf, ntiBuf, hsmBuf; +}; + +extern "C" __global__ __aicore__ void copy_and_expand_eagle_inputs( + GM_ADDR targetTokenIds, GM_ADDR targetPositions, + GM_ADDR nextTokenIds, GM_ADDR queryStartLoc, + GM_ADDR queryEndLoc, + GM_ADDR outInputIds, GM_ADDR outPositions, + GM_ADDR outIsRejectedTokenMask, GM_ADDR outIsMaskedTokenMask, + GM_ADDR outNewTokenIndices, GM_ADDR outHiddenStateMapping, + GM_ADDR workspace, GM_ADDR tiling) +{ + GET_TILING_DATA(tilingData, tiling); + + if (GetBlockIdx() >= tilingData.usedCoreNum) { + return; + } + + if (TILING_KEY_IS(1)) { + CopyAndExpandEagleInputsKernel op; + op.Init(targetTokenIds, targetPositions, nextTokenIds, queryStartLoc, queryEndLoc, + outInputIds, outPositions, outIsRejectedTokenMask, outIsMaskedTokenMask, + outNewTokenIndices, outHiddenStateMapping, &tilingData); + + if (tilingData.shiftInputIds == 0) { + op.ProcessShiftFalse(); + } else { + op.ProcessShiftTrue(); + } + } +} diff --git a/csrc/ascend/moe/dequant_swiglu_quant/CMakeLists.txt b/csrc/ascend/moe/dequant_swiglu_quant/CMakeLists.txt new file mode 100644 index 000000000..be2a3b531 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/CMakeLists.txt @@ -0,0 +1,20 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ---------------------------------------------------------------------------- + + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_host/CMakeLists.txt b/csrc/ascend/moe/dequant_swiglu_quant/op_host/CMakeLists.txt new file mode 100644 index 000000000..9d37c0c09 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_host/CMakeLists.txt @@ -0,0 +1,28 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ---------------------------------------------------------------------------- +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + dequant_swiglu_quant_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME DequantSwigluQuant + OPTIONS --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror +) + + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE dequant_swiglu_quant ACLNNTYPE aclnn) +endif() \ No newline at end of file diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_def.cpp b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_def.cpp new file mode 100644 index 000000000..a131643e8 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_def.cpp @@ -0,0 +1,658 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_def.cpp + * \brief + */ + +#include +#include "register/op_def_registry.h" + +namespace ops +{ +constexpr uint32_t DEQUANT_SWIGLU_QUANT_VERSION_TWO = 2; +constexpr uint32_t DEQUANT_SWIGLU_QUANT_DEFAULT_VALUE = 2; +constexpr float CLAMP_LIMIT_DEFAULT_VALUE = 0.0; +constexpr float GLU_ALPHA_DEFAULT_VALUE =1.00;// 1.702; +class DequantSwigluQuant : public OpDef +{ + public: + explicit DequantSwigluQuant(const char* name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("weight_scale") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("activation_scale") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("bias") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_INT32}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("quant_scale") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("quant_offset") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("group_index") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + this->Attr("activate_left").AttrType(OPTIONAL).Bool(false); + this->Attr("quant_mode").AttrType(OPTIONAL).String("static"); + this->Attr("dst_type").AttrType(OPTIONAL).Version(DEQUANT_SWIGLU_QUANT_VERSION_TWO).Int(DEQUANT_SWIGLU_QUANT_DEFAULT_VALUE); // default value + this->Attr("round_mode").AttrType(OPTIONAL).Version(DEQUANT_SWIGLU_QUANT_VERSION_TWO).String("rint"); // default value + this->Attr("activate_dim").AttrType(OPTIONAL).Version(DEQUANT_SWIGLU_QUANT_VERSION_TWO).Int(-1); // default value + this->Attr("swiglu_mode").AttrType(OPTIONAL).Version(DEQUANT_SWIGLU_QUANT_VERSION_TWO).Int(0); // default value + this->Attr("clamp_limit").AttrType(OPTIONAL).Version(DEQUANT_SWIGLU_QUANT_VERSION_TWO).Float(CLAMP_LIMIT_DEFAULT_VALUE); // default value + this->Attr("glu_alpha").AttrType(OPTIONAL).Version(DEQUANT_SWIGLU_QUANT_VERSION_TWO).Float(GLU_ALPHA_DEFAULT_VALUE); // default value 1.702; + this->Attr("glu_bias").AttrType(OPTIONAL).Version(DEQUANT_SWIGLU_QUANT_VERSION_TWO).Float(0.0); // default value 1.0 + OpAICoreConfig aicoreConfig; + aicoreConfig.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(false) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("coreType.value", "AiCore"); + this->AICore().AddConfig("ascend910b", aicoreConfig); + this->AICore().AddConfig("ascend910_93", aicoreConfig); + // this->AICore().AddConfig("ascend910b",config_kirin); + // this->AICore().AddConfig("ascend910_93",config_kirin); + + // OpAICoreConfig config_950; + // config_950.Input("x") + // .ParamType(REQUIRED) + // .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, + // ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_FLOAT16, ge::DT_BF16, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, + // ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_FLOAT16, ge::DT_BF16}) + // .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}) + // .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}); + // config_950.Input("weight_scale") + // .ParamType(OPTIONAL) + // .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT}) + // .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}) + // .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}); + // config_950.Input("activation_scale") + // .ParamType(OPTIONAL) + // .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT}) + // .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}) + // .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}); + // config_950.Input("bias") + // .ParamType(OPTIONAL) + // .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, + // ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_INT32, + // ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, + // ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_INT32, + // ge::DT_FLOAT, ge::DT_FLOAT}) + // .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}) + // .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}); + // config_950.Input("quant_scale") + // .ParamType(OPTIONAL) + // .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT}) + // .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}) + // .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}); + // config_950.Input("quant_offset") + // .ParamType(OPTIONAL) + // .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT}) + // .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}) + // .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}); + // config_950.Input("group_index") + // .ParamType(OPTIONAL) + // .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + // ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + // ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + // ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + // ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + // ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + // ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + // ge::DT_INT64, ge::DT_INT64, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + // ge::DT_INT32, ge::DT_INT32}) + // .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}) + // .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}); + // config_950.Output("y") + // .ParamType(REQUIRED) + // .DataType({ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, + // ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, + // ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, + // ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, + // ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, + // ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, + // ge::DT_HIFLOAT8, ge::DT_HIFLOAT8, ge::DT_HIFLOAT8, ge::DT_HIFLOAT8, + // ge::DT_HIFLOAT8, ge::DT_HIFLOAT8, + // ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, + // ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, + // ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, + // ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, + // ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, + // ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, + // ge::DT_HIFLOAT8, ge::DT_HIFLOAT8, ge::DT_HIFLOAT8, ge::DT_HIFLOAT8, + // ge::DT_HIFLOAT8, ge::DT_HIFLOAT8}) + // .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}) + // .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}); + // config_950.Output("scale") + // .ParamType(REQUIRED) + // .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + // ge::DT_FLOAT, ge::DT_FLOAT}) + // .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}) + // .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + // ge::FORMAT_ND, ge::FORMAT_ND}); + // config_950.DynamicCompileStaticFlag(true) + // .DynamicFormatFlag(true) + // .DynamicRankSupportFlag(true) + // .DynamicShapeSupportFlag(true) + // .NeedCheckSupportFlag(false) + // .PrecisionReduceFlag(true) + // .ExtendCfgInfo("opFile.value", "dequant_swiglu_quant_apt"); + + // OpAICoreConfig config_kirin = GetKirinCoreConfig(); + // this->AICore().AddConfig("kirinx90", config_kirin); + // this->AICore().AddConfig("kirin9030", config_kirin); + } +/* + private: + OpAICoreConfig GetKirinCoreConfig() const + { + OpAICoreConfig config_kirin; + config_kirin.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true); + config_kirin.Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Input("weight_scale") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Input("activation_scale") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Input("bias") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Input("quant_scale") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Input("quant_offset") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Input("group_index") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + config_kirin.Output("scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + return config_kirin; + }*/ +}; + +OP_ADD(DequantSwigluQuant); +} // namespace ops diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_infershape.cpp b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_infershape.cpp new file mode 100644 index 000000000..da1c08074 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_infershape.cpp @@ -0,0 +1,97 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_infershape.cpp + * \brief + */ +#include "register/op_impl_registry.h" +#include "graph/utils/type_utils.h" +#include "util/shape_util.h" +#include "log/log.h" +#include "util/math_util.h" + +using namespace ge; +namespace ops { +constexpr size_t INPUT_IDX_X = 0; +constexpr size_t OUTPUT_IDX_Y = 0; +constexpr size_t OUTPUT_IDX_SCALE = 1; +constexpr int64_t CONST_UNKNOW_SHAPE = -1; +constexpr int64_t NUM_TWO = 2; +constexpr int64_t INDEX_ATTR_DST_TYPE = 2; +constexpr int64_t INDEX_ATTR_ACTIVATE_DIM = 4; +static const std::initializer_list Y_SUPPORT_DTYPE_SET = {ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, + ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2, + ge::DT_INT8, ge::DT_HIFLOAT8}; + +graphStatus InferShape4DequantSwigluQuant(gert::InferShapeContext* context) { + OP_LOGD(context, "Begin to do InferShape4DequantSwigluQuant."); + + const gert::Shape* xShape = context->GetInputShape(INPUT_IDX_X); + OP_CHECK_NULL_WITH_CONTEXT(context, xShape); + gert::Shape* yShape = context->GetOutputShape(OUTPUT_IDX_Y); + OP_CHECK_NULL_WITH_CONTEXT(context, yShape); + gert::Shape* scaleShape = context->GetOutputShape(OUTPUT_IDX_SCALE); + OP_CHECK_NULL_WITH_CONTEXT(context, scaleShape); + + *yShape = *xShape; + OP_CHECK_IF(Ops::Base::IsUnknownRank(*xShape), + OP_LOGD(context, "End to do InferShape4DequantSwigluQuant, inputx is [-2]."), + return GRAPH_SUCCESS); + + auto attrsPtr = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrsPtr); + const int64_t *activateDim = attrsPtr->GetAttrPointer(INDEX_ATTR_ACTIVATE_DIM); + const int64_t activateDimNum = (activateDim == nullptr) ? -1 : *activateDim; + + // 将切分轴转换为正数 + int64_t xShapeRank = static_cast(xShape->GetDimNum()); + int64_t selectDim = (activateDimNum >= 0) ? activateDimNum : (activateDimNum + xShapeRank); + OP_CHECK_IF(selectDim >= xShapeRank, + OP_LOGE(context, "activateDim must < xShapeRank, but is %ld, xShapeRank is %ld", selectDim, xShapeRank), + return ge::GRAPH_FAILED); + int64_t activateShape = xShape->GetDim(selectDim); + int64_t outActivateShape = activateShape == CONST_UNKNOW_SHAPE ? CONST_UNKNOW_SHAPE : activateShape / NUM_TWO; + OP_CHECK_IF((activateShape != CONST_UNKNOW_SHAPE) && (activateShape % NUM_TWO != 0), + OP_LOGE(context, "The active axis must be an even number， but is %ld", activateShape), + return ge::GRAPH_FAILED); + // 设置Y的shape + yShape->SetDim(selectDim, outActivateShape); + // 设置Scale的shape + *scaleShape = *yShape; + scaleShape->SetDimNum(xShapeRank - 1); + OP_LOGD(context, "End to do InferShape4DequantSwigluQuant"); + return ge::GRAPH_SUCCESS; +} + +graphStatus InferDtype4DequantSwigluQuant(gert::InferDataTypeContext* context) { + OP_LOGD(context, "InferDtype4DequantSwigluQuant enter"); + + auto attrsPtr = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrsPtr); + const int64_t *dstDtype = attrsPtr->GetAttrPointer(INDEX_ATTR_DST_TYPE); + const int64_t dstDtypeNum = (dstDtype == nullptr) ? NUM_TWO : *dstDtype; + + ge::DataType outDtype = static_cast(dstDtypeNum); + OP_CHECK_IF(std::find(Y_SUPPORT_DTYPE_SET.begin(), Y_SUPPORT_DTYPE_SET.end(), outDtype) == Y_SUPPORT_DTYPE_SET.end(), + OP_LOGE(context, "dst_type is illegal, only supports 2(INT8) 40(FLOAT4_E2M1), 41(FLOAT4_E1M2), 35(FLOAT8E5M2), 36(FLOAT8E4M3), 34(HiFloat8)"), + return ge::GRAPH_FAILED); + + context->SetOutputDataType(OUTPUT_IDX_Y, outDtype); + context->SetOutputDataType(OUTPUT_IDX_SCALE, DT_FLOAT); + OP_LOGD(context, "InferDtype4DequantSwigluQuant end"); + + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(DequantSwigluQuant) + .InferShape(InferShape4DequantSwigluQuant) + .InferDataType(InferDtype4DequantSwigluQuant); +} // namespace ops diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_proto.h b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_proto.h new file mode 100644 index 000000000..0d7314cda --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_proto.h @@ -0,0 +1,96 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_proto.h + * \brief + */ +#ifndef OPS_QUANT_DEQUANT_SWIGLU_QUANT_PROTO_H_ +#define OPS_QUANT_DEQUANT_SWIGLU_QUANT_PROTO_H_ +#include "graph/operator_reg.h" + +namespace ge { + +/** +* @brief Combine Dequant + Swiglu + Quant. + +* @par Inputs: +* Seven inputs, including: +* @li x: A tensor. Shape is (X..., H), dim must > 2, and H must be even. Type is int32, float16, bfloat16. +* @li weight_scale: Dequantization scale of weight. An optional tensor. Type is float32. Shape is (1..., H). +* @li activation_scale: Dequantization scale of activation. An optional tensor. Type is float32. Shape is (X..., 1). +* @li bias: Bias for matmul. An optional tensor. Type is float16/bfloat16/int32/float32. Shape is (X..., H). +* @li quant_scale: Quantized scale. An optional tensor. Type is float16/bfloat16/float32. Shape is (1..., H). +* @li quant_offset: Quantized offset. An optional tensor. Type is float16/bfloat16/float32. Shape is (1..., H). +* @li group_index: Mean group index. An optional tensor. Type is int32/int64. Shape is (1,). \n + +* @par Outputs: +* @li y: A tensor. Type is int8/fp8_e5m2/fp8_e4m3fn/fp4x2_e2m1/fp4x2_e1m2/hifloat8. +* @li scale: A tensor. Type is float32. + +* @par Attributes: +* @li activate_left: Type is bool. +* The swi activate_left algorithm to use: +* 'false'(activate right) or 'true'(activate left), default is 'false'(activate right). +* @li quant_mode: Type is string. The quant mode to use: 'static' or 'dynamic', default is 'static'. +* @li dst_type: Type is Int32. Declare the output y dtype. Support 2:int8, 35:fp8_e5m2, 36:fp8_e4m3fn, 40:fp4x2_e2m1, 41:fp4x2_e1m2, Defaults to 2, only used for Ascend 950 AI Processors. +* @li round_mode: Type is String. The round mode to use: 'rint', 'round, 'floor', 'ceil', 'trunc', default is 'rint', only used for Ascend 950 AI Processors. +* @li activate_dim: Type is Int32. Describing the split dimension in Glu algorithm: value in [-xDim, xDim-1], default is -1, only used for Ascend 950 AI Processors. +* @li swiglu_mode: Type is int. Optional parameter, default is 0. The SWIGLU computation mode to use: +* '0' (default) for standard SWIGLU, '1' for a variant using odd-even blocking, which requires support for clamp_limit, activation coefficient, and bias. This attribute is not supported in Ascend 950 AI Processors. +* @li clamp_limit: Type is float. Optional parameter, default is 0.0. The threshold limit for SWIGLU input. Use 0.0 to disable clamp. This attribute is not supported in Ascend 950 AI Processors. +* @li glu_alpha: Type is float. Optional parameter, default is 1.702. The activation coefficient for the GLU activation function. This attribute is not supported in Ascend 950 AI Processors. +* @li glu_bias: Type is float. Optional parameter, default is 1.0. The bias applied during SWIGLU linear computation. This attribute is not supported in Ascend 950 AI Processors. + +* @attention Constraints: +* @li When the type of x is int32, weight_scale must be input. +* @li When the type of x is float16, bfloat16, weight_scale, activation_scale and bias must be None. +* @li When dst_type is int8, fp8_e5m2, fp8_e4m3fn, round_mode only supports 'rint'. +* @li When dst_type is fp4x2_e2m1 or fp4x2_e1m2, round_mode supports 'rint', 'round, 'floor', 'ceil' and 'trunc'. +* @li When dst_type is hifloat8, round_mode supports 'round'. +* @li The shape of activate_dim corresponding to x must be divisible by 2. +* @li When activate_dim is not the last dimension of x, group_index must be None. +* @li The input quant_offset is not supported in Ascend 950 AI Processors only. +* @li The type of output y is fp8_e5m2, fp8_e4m3fn, fp4x2_e2m1 and fp4x2_e1m2 only supported in Ascend 950 AI Processors. +* @li The attribute quant_mode is only supported 'dynamic' in Ascend 950 AI Processors. +* @li The attribute dst_type is only supported in Ascend 950 AI Processors. +* @li The attribute round_mode is only supported in Ascend 950 AI Processors. +* @li The attribute activate_dim is only supported in Ascend 950 AI Processors. +* @li The attribute swiglu_mode is not supported in Ascend 950 AI Processors. +* @li The attribute clamp_limit is not supported in Ascend 950 AI Processors. +* @li The attribute glu_alpha is not supported in Ascend 950 AI Processors. +* @li The attribute glu_bias is not supported in Ascend 950 AI Processors. + +* @par Restrictions: +* Warning: THIS FUNCTION IS EXPERIMENTAL. Please do not use. +*/ +REG_OP(DequantSwigluQuant) + .INPUT(x, TensorType({DT_FLOAT16, DT_BF16, DT_INT32})) + .OPTIONAL_INPUT(weight_scale, TensorType({DT_FLOAT})) + .OPTIONAL_INPUT(activation_scale, TensorType({DT_FLOAT})) + .OPTIONAL_INPUT(bias, TensorType({DT_FLOAT16, DT_BF16, DT_INT32, DT_FLOAT})) + .OPTIONAL_INPUT(quant_scale, TensorType({DT_BF16, DT_FLOAT16, DT_FLOAT})) + .OPTIONAL_INPUT(quant_offset, TensorType({DT_BF16, DT_FLOAT16, DT_FLOAT})) + .OPTIONAL_INPUT(group_index, TensorType({DT_INT32, DT_INT64})) + .OUTPUT(y, TensorType({DT_INT8, DT_FP8_E4M3FN, DT_FP8_E5M2, DT_FP4X2_E2M1, DT_FP4X2_E1M2, DT_HIFLOAT8})) + .OUTPUT(scale, TensorType({DT_FLOAT})) + .ATTR(activate_left, Bool, false) + .ATTR(quant_mode, String, "static") + .ATTR(dst_type, Int, DT_INT8) + .ATTR(round_mode, String, "rint") + .ATTR(activate_dim, Int, -1) + .ATTR(swiglu_mode, Int, 0) + .ATTR(clamp_limit, Float, 0.0) + .ATTR(glu_alpha, Float, 1.702) + .ATTR(glu_bias, Float, 1.0) + .OP_END_FACTORY_REG(DequantSwigluQuant) +} // namespace ge + +#endif // OPS_QUANT_DEQUANT_SWIGLU_QUANT_PROTO_H_ diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling.cpp b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling.cpp new file mode 100644 index 000000000..1da6069c8 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling.cpp @@ -0,0 +1,778 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_tiling.cpp + * \brief + */ + +#include "dequant_swiglu_quant_tiling.h" +#include "../tiling_base/tiling_util.h" +#include "swi_glu_tiling.h" +#include "../tiling_base/tiling_templates_registry.h" + +// using namespace AscendC; +using namespace ge; +namespace optiling +{ +constexpr int64_t ATTR_ACTIVATE_LEFT_INDEX = 0; +constexpr int64_t ATTR_QUANT_MODE_INDEX = 1; +constexpr int64_t X_INDEX = 0; +constexpr int64_t WEIGHT_SCALE_INDEX = 1; +constexpr int64_t ACTIVATION_SCALE_INDEX = 2; +constexpr int64_t BIAS_INDEX = 3; +constexpr int64_t QUANT_SCALE_INDEX = 4; +constexpr int64_t QUANT_OFFSET_INDEX = 5; +constexpr int64_t INPUT_GROUP_INDEX = 6; +constexpr int64_t Y_INDEX = 0; +// attr index for SwiGLU used by GPT-OSS +constexpr int64_t SWIGLU_MODE_INDEX = 5; +constexpr int64_t CLAMP_LIMIT_INDEX = 6; +constexpr int64_t GLU_ALPHA_INDEX = 7; +constexpr int64_t GLU_BIAS_INDEX = 8; + +constexpr int64_t BLOCK_SIZE = 32; +constexpr int64_t BLOCK_ELEM = BLOCK_SIZE / static_cast(sizeof(float)); +constexpr uint64_t WORKSPACE_SIZE = 32; +// define tiling key offset +constexpr uint64_t TILING_KEY_HAS_GROUP = 100000000; +constexpr uint64_t TILING_KEY_NO_GROUP = 200000000; +// define cut by group +constexpr uint64_t TILING_KEY_CUT_GROUP = 10000000; +constexpr int64_t CUT_GROUP_LARGE_THAN_64 = 64; +constexpr int64_t CUT_GROUP_LARGE_THAN_32 = 32; +constexpr int64_t EACH_GROUP_TOKEN_LESS_THAN = 16; + +// quant_scale tiling offset +constexpr uint64_t TILING_KEY_QS_DTYPE = 100; +// bias tiling offset +constexpr uint64_t TILING_KEY_BIAS_DTYPE = 1000; + +constexpr int64_t UB_RESERVE = 1024; +constexpr int64_t SWI_FACTOR = 2; +constexpr int64_t QUANT_MODE_DYNAMIC = 1; +constexpr int64_t PERFORMANCE_H_2048 = 2048; +constexpr int64_t PERFORMANCE_H_4096 = 4096; +constexpr int64_t PERFORMANCE_CORE_NUM = 36; +constexpr int64_t PERFORMANCE_UB_FACTOR = static_cast(4096) * 4; + +constexpr int QUANT_SCALE_DTYPE_BF16 = 2; +constexpr int QUANT_SCALE_DTYPE_FP32 = 0; +constexpr int QUANT_SCALE_DTYPE_FP16 = 1; + +constexpr int BIAS_DTYPE_BF16 = 0; +constexpr int BIAS_DTYPE_FP16 = 1; +constexpr int BIAS_DTYPE_FP32 = 2; +constexpr int BIAS_DTYPE_INT32 = 3; + +constexpr int DIM_SIZE_2 = 2; + +constexpr float CLAMP_LIMIT_DEFAULT= 0.0; +constexpr float GLU_ALPHA_DEFAULT = 1.702; +constexpr float GLU_BIAS_DEFAULT = 1.0; + +static const std::set SUPPORT_DTYPE = {ge::DT_INT32, ge::DT_BF16}; +static const std::map SUPPORT_QUANT_MODE = {{"dynamic", 1},{"static", 0}}; + +bool DequantSwigluQuantDskTiling::CheckOptionalShapeExisting(const gert::StorageShape* storageShape){ + if(storageShape == nullptr){ + return false; + } + int64_t shapeSize = storageShape->GetOriginShape().GetShapeSize(); + if(shapeSize <= 0){ + return false; + } + return true; +} + +ge::graphStatus DequantSwigluQuantDskTiling::GetPlatformInfo() { + auto platformInfo = context_->GetPlatformInfo(); + if (platformInfo == nullptr) { + auto compileInfoPtr = context_->GetCompileInfo(); + OP_CHECK_IF(compileInfoPtr == nullptr, OP_LOGE(context_, "compile info is null"), + return ge::GRAPH_FAILED); + coreNum_ = compileInfoPtr->coreNum; + ubSize_ = compileInfoPtr->ubSize; + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + coreNum_ = ascendcPlatform.GetCoreNumAiv(); + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + ubSize_ = ubSizePlatForm; + socVersion = ascendcPlatform.GetSocVersion(); + } + + maxPreCore_ = static_cast(coreNum_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::CheckXAndGroupIndexDtype() { + auto xPtr = context_->GetInputDesc(X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, xPtr); + auto xDtype = xPtr->GetDataType(); + OP_CHECK_IF((SUPPORT_DTYPE.find(xDtype) == SUPPORT_DTYPE.end()), + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "x", + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str(), "int32 or bfloat16"), + return ge::GRAPH_FAILED); + tilingData_.set_groupIndexDtype(-1); + if (hasGroupIndex_) { + auto groupIndexPtr = context_->GetOptionalInputDesc(INPUT_GROUP_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, groupIndexPtr); + auto groupIndexDtype = groupIndexPtr->GetDataType(); + bool dtypeInValid = groupIndexDtype != ge::DT_INT64; + OP_CHECK_IF( + dtypeInValid, + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "group_index", + ge::TypeUtils::DataTypeToSerialString(groupIndexDtype).c_str(), "int64"), + return ge::GRAPH_FAILED); + tilingData_.set_groupIndexDtype(1); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::CheckBias() { + auto biasShapePtr = context_->GetOptionalInputShape(BIAS_INDEX); + if (biasShapePtr != nullptr) { + hasBias_ = true; + OP_CHECK_IF(CheckScaleShapeWithDim(BIAS_INDEX, inDimy_, "bias") != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "bias shape check failed."), + return ge::GRAPH_FAILED); + } + else { + hasBias_ = false; + } + + auto biasPtr = context_->GetOptionalInputDesc(BIAS_INDEX); + if (biasPtr != nullptr && hasBias_ == true) { + auto biasDtype = biasPtr->GetDataType(); + bool dtypeInValid = (biasDtype != ge::DT_INT32 && biasDtype != ge::DT_FLOAT && biasDtype != ge::DT_FLOAT16 && biasDtype != ge::DT_BF16); + OP_CHECK_IF( + dtypeInValid, + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "bias", + ge::TypeUtils::DataTypeToSerialString(biasDtype).c_str(), "bf16, fp16, float or int32"), + return ge::GRAPH_FAILED); + if (biasDtype == ge::DT_BF16) { + tilingData_.set_biasDtype(BIAS_DTYPE_BF16); + } else if (biasDtype == ge::DT_FLOAT16) { + tilingData_.set_biasDtype(BIAS_DTYPE_FP16); + } else if (biasDtype == ge::DT_FLOAT) { + tilingData_.set_biasDtype(BIAS_DTYPE_FP32); + } else if (biasDtype == ge::DT_INT32) { + tilingData_.set_biasDtype(BIAS_DTYPE_INT32); + } + } + else { + tilingData_.set_biasDtype(0); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::CheckWeightScale() { + auto weightScalePtr = context_->GetOptionalInputDesc(WEIGHT_SCALE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, weightScalePtr); + auto weightScaleDtype = weightScalePtr->GetDataType(); + bool dtypeInValid = weightScaleDtype != ge::DT_FLOAT; + OP_CHECK_IF(dtypeInValid, + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "weight_scale", + ge::TypeUtils::DataTypeToSerialString(weightScaleDtype).c_str(), "float32"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckScaleShapeWithDim(WEIGHT_SCALE_INDEX, inDimy_, "weight_scale") != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "weight scale shape check failed."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::CheckActivationScale() { + auto activationScaleShapePtr = context_->GetOptionalInputShape(ACTIVATION_SCALE_INDEX); + if(CheckOptionalShapeExisting(activationScaleShapePtr)) { + auto activationScalePtr = context_->GetOptionalInputDesc(ACTIVATION_SCALE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, activationScalePtr); + auto activationScaleDtype = activationScalePtr->GetDataType(); + bool dtypeInValid = activationScaleDtype != ge::DT_FLOAT; + + OP_CHECK_IF(dtypeInValid, + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "activation_scale", + ge::TypeUtils::DataTypeToSerialString(activationScaleDtype).c_str(), "float32"), + return ge::GRAPH_FAILED); + OP_CHECK_NULL_WITH_CONTEXT(context_, activationScaleShapePtr); + auto activationScaleShape = activationScaleShapePtr->GetStorageShape(); + int64_t activationScaleNum = activationScaleShape.GetShapeSize(); + + OP_CHECK_IF( + activationScaleNum != inDimx_, + OP_LOGE( + context_->GetNodeName(), + "activation_scale num(%ld) must be equal to the tokens num(%ld), please check.", + activationScaleNum, inDimx_), + return ge::GRAPH_FAILED); + tilingData_.set_activationScaleIsEmpty(0); + } + else { + tilingData_.set_activationScaleIsEmpty(1); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::CheckForDequant() { + // check weight scale, activation scale and bias + auto xPtr = context_->GetInputDesc(X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, xPtr); + auto xDtype = xPtr->GetDataType(); + if (xDtype == ge::DT_INT32) { + OP_CHECK_IF(CheckWeightScale() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "weight scale check failed."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckActivationScale() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "activation scale check failed."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckBias() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "bias check failed."), + return ge::GRAPH_FAILED); + } + + if (xDtype == ge::DT_BF16 && hasGroupIndex_) { + auto shapeGroupIndex = context_->GetOptionalInputShape(INPUT_GROUP_INDEX); + const gert::Shape& inputShapeGroupIndex = shapeGroupIndex->GetStorageShape(); + OP_CHECK_IF(inputShapeGroupIndex.GetDimNum() != 1, + OP_LOGE(context_->GetNodeName(), + "groupIndex only support 1D Tensor now, please check."), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::CheckForDynamicQuant() { + auto offsetPtr = context_->GetOptionalInputShape(QUANT_OFFSET_INDEX); + OP_CHECK_IF(offsetPtr != nullptr, + OP_LOGE(context_->GetNodeName(), + "quantOffSet only support None in dynamic quantization of group mode now, please check."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(CheckScaleShapeWithDim(QUANT_SCALE_INDEX, outDimy_, "quant_scale") != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "quant scale shape check failed."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::CheckForStaticQuant() { + // check quantOffset dtype + auto quantOffsetDescPtr = context_->GetOptionalInputDesc(QUANT_OFFSET_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, quantOffsetDescPtr); + auto quantScaleDescPtr = context_->GetOptionalInputDesc(QUANT_SCALE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, quantScaleDescPtr); + auto quantOffsetDtype = quantOffsetDescPtr->GetDataType(); + auto quantScaleDtype = quantScaleDescPtr->GetDataType(); + OP_CHECK_IF(quantOffsetDtype != quantScaleDtype, + OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON( + context_->GetNodeName(), "quant_offset and quant_scale", + (ge::TypeUtils::DataTypeToSerialString(quantOffsetDtype) + " and " + + ge::TypeUtils::DataTypeToSerialString(quantScaleDtype)).c_str(), + "quantOffset dtype must be same as quantScale dtype"), + return ge::GRAPH_FAILED); + + int64_t quantScaleColLen = 0; + int64_t quantOffsetColLen = 0; + OP_CHECK_IF(CheckStaticQuantShape(QUANT_SCALE_INDEX, quantScaleColLen, "quant_scale") != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "quant scale shape check failed."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckStaticQuantShape(QUANT_OFFSET_INDEX, quantOffsetColLen, "quant_offset") != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "quant offset shape check failed."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(quantScaleColLen != quantOffsetColLen, + OP_LOGE(context_->GetNodeName(), "quant offset shape is different from quant scale."), + return ge::GRAPH_FAILED); + if(quantScaleColLen == 1){ + tilingData_.set_quantIsOne(1); + } + else { + tilingData_.set_quantIsOne(0); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::CheckForQuant() { + // check and set quant scale dtype + OP_CHECK_IF(CheckQuantScaleDtype() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "Check QuantScale Dtype failed."), + return ge::GRAPH_FAILED); + + // check quant offset and quant scale shape in dynamic scenario + if(quantMode_ == QUANT_MODE_DYNAMIC){ + OP_CHECK_IF(CheckForDynamicQuant() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "Check For Dynamic Quant failed."), + return ge::GRAPH_FAILED); + } + // // check quant offset and quant scale shape in static scenario + else { + OP_CHECK_IF(CheckForStaticQuant() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "Check For Static Quant failed."), + return ge::GRAPH_FAILED); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::CheckQuantScaleDtype() { + bool dtypeInValid = false; + + auto quantScaleShapePtr = context_->GetOptionalInputShape(QUANT_SCALE_INDEX); + if (quantScaleShapePtr == nullptr) { + tilingData_.set_quantScaleDtype(0); + tilingData_.set_needSmoothScale(0); + } else { + auto quantScalePtr = context_->GetOptionalInputDesc(QUANT_SCALE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, quantScalePtr); + tilingData_.set_needSmoothScale(1); + auto quantScaleDtype = quantScalePtr->GetDataType(); + dtypeInValid = + quantScaleDtype != ge::DT_FLOAT && quantScaleDtype != ge::DT_FLOAT16 && quantScaleDtype != ge::DT_BF16; + OP_CHECK_IF( + dtypeInValid, + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "quant_scale", + ge::TypeUtils::DataTypeToSerialString(quantScaleDtype).c_str(), "float32, float16 or bfloat16"), + return ge::GRAPH_FAILED); + tilingData_.set_quantScaleDtype(QUANT_SCALE_DTYPE_BF16); + if (quantScaleDtype == ge::DT_FLOAT) { + tilingData_.set_quantScaleDtype(QUANT_SCALE_DTYPE_FP32); + } else if (quantScaleDtype == ge::DT_FLOAT16) { + tilingData_.set_quantScaleDtype(QUANT_SCALE_DTYPE_FP16); + } + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::GetAttr() { + auto* attrs = context_->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context_, attrs); + + auto* attrActivateLeft = attrs->GetAttrPointer(ATTR_ACTIVATE_LEFT_INDEX); + actRight_ = (attrActivateLeft == nullptr || *attrActivateLeft == false) ? 1 : 0; + std::string quantMode = attrs->GetAttrPointer(ATTR_QUANT_MODE_INDEX); + auto it = SUPPORT_QUANT_MODE.find(quantMode); + OP_CHECK_IF(it == SUPPORT_QUANT_MODE.end(), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + context_->GetNodeName(), "quant_mode", + quantMode.c_str(), "quant_mode only support dynamic(1) and static(0) currently"), + return ge::GRAPH_FAILED); + quantMode_ = it->second; + + auto* swigluMode = attrs->GetAttrPointer(SWIGLU_MODE_INDEX); + auto* clampLimit = attrs->GetAttrPointer(CLAMP_LIMIT_INDEX); + auto* gluAlpha = attrs->GetAttrPointer(GLU_ALPHA_INDEX); + auto* gluBias = attrs->GetAttrPointer(GLU_BIAS_INDEX); + + swigluMode_ = swigluMode == nullptr ? 0 : *swigluMode; + clampLimit_ = clampLimit == nullptr ? CLAMP_LIMIT_DEFAULT : *clampLimit; + gluAlpha_ = gluAlpha == nullptr ? GLU_ALPHA_DEFAULT : *gluAlpha; + gluBias_ = gluBias == nullptr ? GLU_BIAS_DEFAULT : *gluBias; + + OP_CHECK_IF(swigluMode_ != 0 && swigluMode_ != 1, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + context_->GetNodeName(), "swigluMode", + std::to_string(swigluMode_).c_str(), "swigluMode only support 0 or 1"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(!(clampLimit_ >= 0.0), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + context_->GetNodeName(), "clamp_limit", + std::to_string(clampLimit_).c_str(), "clamp_limit should be non-negative"), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::CheckScaleShapeWithDim(const int64_t scaleInputIdx, + const int64_t expectDim, + const char* paramName) { + auto scalePtr = context_->GetOptionalInputShape(scaleInputIdx); + if (scalePtr == nullptr) { + return ge::GRAPH_SUCCESS; + } + auto scaleShape = scalePtr->GetStorageShape(); + OP_CHECK_IF(scaleShape.GetDimNum() < 1, + OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), paramName, + std::to_string(scaleShape.GetDimNum()).c_str(), "greater than or equal to 1"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(scaleShape.GetDim(scaleShape.GetDimNum() - 1) != expectDim, + OP_LOGE_FOR_INVALID_SHAPE(context_->GetNodeName(), paramName, + Ops::Base::ToString(scaleShape).c_str(), + std::to_string(expectDim).c_str()), + return ge::GRAPH_FAILED); + if (groupNum_ > 1) { + // check with group index + OP_CHECK_IF( + scaleShape.GetDimNum() != DIM_SIZE_2, + OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), paramName, + std::to_string(scaleShape.GetDimNum()).c_str(), "2D"), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + scaleShape.GetDim(0) != groupNum_, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON( + context_->GetNodeName(), paramName, Ops::Base::ToString(scaleShape).c_str(), + ("the first dimension of " + std::string(paramName) + " (" + std::to_string(scaleShape.GetDim(0)) + + ") must be equal to the first dimension of group_index (" + std::to_string(groupNum_) + ")") + .c_str()), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF( + scaleShape.GetDimNum() > DIM_SIZE_2, + OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), paramName, + std::to_string(scaleShape.GetDimNum()).c_str(), "less than or equal to 2"), + return ge::GRAPH_FAILED); + int64_t groupNumFromScale = scaleShape.GetDimNum() <= 1 ? 1 : scaleShape.GetDim(0); + OP_CHECK_IF( + groupNumFromScale != 1, + OP_LOGE_FOR_INVALID_SHAPE(context_->GetNodeName(), paramName, + Ops::Base::ToString(scaleShape).c_str(), + ("[1," + std::to_string(expectDim) + "] or [" + std::to_string(expectDim) + "]").c_str()), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::CheckStaticQuantShape(const int64_t quantInputIdx, int64_t& colLen, const char* paramName) { + // check quant scale and quant offset shape + auto quantPtr = context_->GetOptionalInputShape(quantInputIdx); + if(quantPtr == nullptr){ + return ge::GRAPH_SUCCESS; + } + auto quantShape = quantPtr->GetStorageShape(); + OP_CHECK_IF(quantShape.GetDimNum() < 1, + OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), paramName, + std::to_string(quantShape.GetDimNum()).c_str(), "greater than or equal to 1"), + return ge::GRAPH_FAILED); + colLen = quantShape.GetDim(quantShape.GetDimNum() - 1); + if(quantShape.GetDimNum() == 1){ + OP_CHECK_IF(colLen != groupNum_, + OP_LOGE_FOR_INVALID_SHAPE(context_->GetNodeName(), paramName, + Ops::Base::ToString(quantShape).c_str(), + ("[" + std::to_string(groupNum_) + ", ] or [" + + std::to_string(groupNum_) + ", " + std::to_string(outDimy_) + "]").c_str()), + return ge::GRAPH_FAILED); + colLen = 1; + } + else { + OP_CHECK_IF(colLen != outDimy_ || quantShape.GetDim(0) != groupNum_, + OP_LOGE_FOR_INVALID_SHAPE(context_->GetNodeName(), paramName, + Ops::Base::ToString(quantShape).c_str(), + ("[" + std::to_string(groupNum_) + ", ] or [" + + std::to_string(groupNum_) + ", " + std::to_string(outDimy_) + "]").c_str()), + return ge::GRAPH_FAILED); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::GetShapeAttrsInfo() { + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::CheckIllegalParam() { + // if hasbias, speGroupType_ must be false + if (hasBias_) { + OP_CHECK_IF(speGroupType_ == true, + OP_LOGE(context_->GetNodeName(), "speGroupType_ only support false when using bias"), + return ge::GRAPH_FAILED); + } + + // if swigluMode is 1, speGroupType_ must be false + if (swigluMode_) { + OP_CHECK_IF(speGroupType_ == true, + OP_LOGE(context_->GetNodeName(), "speGroupType_ only support false when swiglu mode is 1"), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::GetShapeAttrsInfoInner() { + if (!IsPerformanceAndGroupIndexBrach()) { + return ge::GRAPH_SUCCESS; + } + // get 2H from x, get H from y, check if 2H can be divided by 64 + auto shapeX = context_->GetInputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context_, shapeX); + const gert::Shape& inputShapeX = shapeX->GetStorageShape(); + int64_t inputShapeXTotalNum = inputShapeX.GetShapeSize(); + int64_t inputShapeXRank = inputShapeX.GetDimNum(); + inDimy_ = inputShapeX.GetDim(inputShapeXRank - 1); + inDimx_ = inputShapeXTotalNum / inDimy_; + auto shapeY = context_->GetOutputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context_, shapeY); + const gert::Shape& outputShapeY = shapeY->GetStorageShape(); + outDimy_ = outputShapeY.GetDim(inputShapeXRank - 1); + OP_CHECK_IF(inDimy_ % (BLOCK_SIZE * SWI_FACTOR) != 0, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), "x", + std::to_string(inDimy_).c_str(), + "lastdimSize of x must be divisible by 64"), + return ge::GRAPH_FAILED); + + // set the relevant param of group, hasGroupIndex_, groupNum_ and speGroupType_ + auto shapeGroupIndex = context_->GetOptionalInputShape(INPUT_GROUP_INDEX); + hasGroupIndex_ = shapeGroupIndex != nullptr; + groupNum_ = 1; + speGroupType_ = false; + if (hasGroupIndex_) { + const gert::Shape& inputShapeGroupIndex = shapeGroupIndex->GetStorageShape(); + groupNum_ = inputShapeGroupIndex.GetDimNum() == 0 ? 1 : inputShapeGroupIndex.GetDim(0); + speGroupType_ = inputShapeGroupIndex.GetDimNum() == DIM_SIZE_2; + } + + OP_CHECK_IF(CheckXAndGroupIndexDtype() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "dtype check failed."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(GetAttr() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "get attr failed."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(CheckForDequant() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "check for dequant failed."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(CheckForQuant() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "check for quant failed."), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(CheckIllegalParam() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "check illegal param failed."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +bool DequantSwigluQuantDskTiling::IsPerformanceAndGroupIndexBrach() { + auto shapeGroupIndex = context_->GetOptionalInputShape(INPUT_GROUP_INDEX); + if (shapeGroupIndex != nullptr) { + return true; + } + + auto xPtr = context_->GetInputDesc(X_INDEX); + auto attrs = context_->GetAttrs(); + if (xPtr == nullptr || attrs == nullptr) { + return false; + } + auto* swigluMode = attrs->GetAttrPointer(SWIGLU_MODE_INDEX); + return xPtr->GetDataType() == ge::DT_INT32 && swigluMode != nullptr && *swigluMode == 1; +} + +bool DequantSwigluQuantDskTiling::IsCapable() { + return IsPerformanceAndGroupIndexBrach(); +} + +void DequantSwigluQuantDskTiling::CountTilingKey() { + auto xPtr = context_->GetInputDesc(X_INDEX); + auto xDtype = xPtr->GetDataType(); + tilingKey_ = hasGroupIndex_ ? TILING_KEY_HAS_GROUP : TILING_KEY_NO_GROUP; + // add quant scale offset to tilingKey_ + tilingKey_ += TILING_KEY_QS_DTYPE * tilingData_.get_quantScaleDtype(); + // add bias offset to tilingKey_ + tilingKey_ += TILING_KEY_BIAS_DTYPE * tilingData_.get_biasDtype(); + // tiling based on groupnum, pre cut num by coreNum_ and total tokens + bool cond1 = speGroupType_ && + (groupNum_ >= CUT_GROUP_LARGE_THAN_64) && + (inDimx_ / groupNum_ <= EACH_GROUP_TOKEN_LESS_THAN); + bool cond2 = !speGroupType_ && + (groupNum_ >= CUT_GROUP_LARGE_THAN_32) && + (inDimx_ / groupNum_ <= EACH_GROUP_TOKEN_LESS_THAN) && + !tilingData_.get_biasDtype() && !tilingData_.get_quantScaleDtype() && + (xDtype == ge::DT_INT32); + if (cond1 || cond2) { + tilingKey_ += TILING_KEY_CUT_GROUP; + maxPreCore_ = std::min(static_cast(coreNum_), static_cast(inDimx_)); + } +} + +ge::graphStatus DequantSwigluQuantDskTiling::CountMaxDim(int64_t& ubFactorDimx) { + /* + x used mem: [UbFactorDimx, outDimy_ * 2] dtype: float + activation_scale used mem: [UbFactorDimx, 8] dtype: float + weight_scale used mem: [1, outDimy_ * 2] dtype: float + quant_scale used mem: [1, outDimy_] dtype: float + y used mem: [UbFactorDimx, outDimy_] dtype: int8_t + scale used mem: [UbFactorDimx,] dtype: float + tmp used mem: [UbFactorDimx, outDimy_ * 2] dtype: float + x, activation_scale enable db + ub reserve 1024B + + optional buffer： + bias used mem: [1, outDimy_ * 2] dtype: float + + clamp tmp buffer: [UbFactorDimx, outDimy_] dtype: uint8 + + gather offset buffer: [UbFactorDimx, outDimy_] dtype: uint32 + + */ + int64_t db = 2; + int64_t maxOutDimy = 0; + int64_t biasBufferY = hasBias_ == false ? 0 : static_cast(SWI_FACTOR * sizeof(float)); + int64_t biasBufferX = hasBias_ == false ? 0 : outDimy_ * SWI_FACTOR * static_cast(sizeof(float)); + + int64_t SweiGLUBufferY = swigluMode_ == 0 ? 0 : static_cast(sizeof(int8_t) + sizeof(int32_t)); + int64_t SweiGLUBufferX = swigluMode_ == 0 ? 0 : outDimy_ * static_cast(sizeof(int8_t)) + outDimy_ * static_cast(sizeof(int32_t)); + + int64_t quantOffsetSpace = quantMode_ == QUANT_MODE_DYNAMIC ? 0 : static_cast(sizeof(float)); + + // UbFactorDimx is 1,compute maxOutDimy + int64_t numerator = static_cast(ubSize_) - UB_RESERVE - BLOCK_SIZE - db * BLOCK_SIZE - static_cast(sizeof(float)); + int64_t denominator = + 5 * static_cast(sizeof(float)) + db * SWI_FACTOR * static_cast(sizeof(float)) + static_cast(sizeof(int8_t)) + biasBufferY + SweiGLUBufferY + quantOffsetSpace; + maxOutDimy = static_cast(numerator / denominator); + maxOutDimy = maxOutDimy / BLOCK_SIZE * BLOCK_SIZE; + int64_t maxInDimy = static_cast(maxOutDimy * SWI_FACTOR); + OP_LOGI(context_->GetNodeName(), "Get maxInDimy[%ld]", maxInDimy); + OP_CHECK_IF(inDimy_ > maxInDimy, + OP_LOGE_FOR_INVALID_SHAPESIZE(context_->GetNodeName(), "x", + std::to_string(inDimy_).c_str(), + ("less than or equal to " + std::to_string(maxInDimy)).c_str()), + return ge::GRAPH_FAILED); + + // compute ubFactorDimx + quantOffsetSpace = quantMode_ == QUANT_MODE_DYNAMIC ? 0 : outDimy_ * sizeof(float); + numerator = static_cast(ubSize_) - UB_RESERVE - outDimy_ * static_cast(sizeof(float)) - BLOCK_SIZE - SWI_FACTOR * outDimy_ * static_cast(sizeof(float)) - biasBufferX - quantOffsetSpace; + + denominator = db * (outDimy_ * SWI_FACTOR + BLOCK_ELEM) * static_cast(sizeof(float)) + outDimy_ * static_cast(sizeof(int8_t)) + static_cast(sizeof(float)) + + outDimy_ * SWI_FACTOR * static_cast(sizeof(float)) + SweiGLUBufferX; + ubFactorDimx = static_cast(numerator / denominator); + ubFactorDimx = std::min(ubFactorDimx, inDimx_); + OP_LOGI(context_->GetNodeName(), "Get ubFactorDimx[%ld]", ubFactorDimx); + + // special ub cut for 2048 4096 + if (swigluMode_ == 0 && hasBias_ == false) { + ubFactorDimx = + (inDimy_ == PERFORMANCE_H_2048 || inDimy_ == PERFORMANCE_H_4096) ? PERFORMANCE_UB_FACTOR / inDimy_ : ubFactorDimx; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::DoOpTiling() { + if (GetShapeAttrsInfoInner() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + auto inputShapeX = context_->GetInputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context_, inputShapeX); + + int64_t ubFactorDimx = 0; + OP_CHECK_IF(CountMaxDim(ubFactorDimx) != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "Count MaxDim failed."), + return ge::GRAPH_FAILED); + + maxPreCore_ = (inDimx_ + ubFactorDimx - 1) / ubFactorDimx; + maxPreCore_ = std::min(maxPreCore_, static_cast(PERFORMANCE_CORE_NUM)); + maxPreCore_ = std::min(maxPreCore_, static_cast(coreNum_)); + + CountTilingKey(); + + tilingData_.set_inDimx(inDimx_); + tilingData_.set_inDimy(inDimy_); + tilingData_.set_outDimy(outDimy_); + tilingData_.set_UbFactorDimx(ubFactorDimx); + tilingData_.set_UbFactorDimy(outDimy_); + tilingData_.set_usedCoreNum(maxPreCore_); + tilingData_.set_maxCoreNum(maxPreCore_); + tilingData_.set_inGroupNum(groupNum_); + tilingData_.set_quantMode(quantMode_); + tilingData_.set_actRight(actRight_); + tilingData_.set_speGroupType(static_cast(speGroupType_)); + tilingData_.set_hasBias(hasBias_); + + tilingData_.set_swigluMode(swigluMode_); + tilingData_.set_clampLimit(clampLimit_); + tilingData_.set_gluAlpha(gluAlpha_); + tilingData_.set_gluBias(gluBias_); + return ge::GRAPH_SUCCESS; +} + +void DequantSwigluQuantDskTiling::DumpTilingInfo() { + std::ostringstream info; + info << "inDimx_: " << tilingData_.get_inDimx(); + info << ", inDimy_: " << tilingData_.get_inDimy(); + info << ", outDimy: " << tilingData_.get_outDimy(); + info << ", UbFactorDimx: " << tilingData_.get_UbFactorDimx(); + info << ", UbFactorDimy: " << tilingData_.get_UbFactorDimy(); + info << ", usedCoreNum: " << tilingData_.get_usedCoreNum(); + info << ", maxCoreNum: " << tilingData_.get_maxCoreNum(); + info << ", inGroupNum: " << tilingData_.get_inGroupNum(); + info << ", quantMode: " << tilingData_.get_quantMode(); + info << ", actRight: " << tilingData_.get_actRight(); + info << ", tilingKey: " << tilingKey_; + info << ", hasBias: " << hasBias_; + info << ", swigluMode: " << tilingData_.get_swigluMode(); + info << ", clampLimit: " << tilingData_.get_clampLimit(); + info << ", gluAlpha: " << tilingData_.get_gluAlpha(); + info << ", gluBias: " << tilingData_.get_gluBias(); + + OP_LOGI(context_->GetNodeName(), "%s", info.str().c_str()); +} + +ge::graphStatus DequantSwigluQuantDskTiling::DoLibApiTiling() { + return ge::GRAPH_SUCCESS; +} + +uint64_t DequantSwigluQuantDskTiling::GetTilingKey() const { + return tilingKey_; +} + +ge::graphStatus DequantSwigluQuantDskTiling::GetWorkspaceSize() { + workspaceSize_ = WORKSPACE_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantDskTiling::PostTiling() { + context_->SetTilingKey(GetTilingKey()); + context_->SetBlockDim(maxPreCore_); + size_t* workspaces = context_->GetWorkspaceSizes(1); + workspaces[0] = workspaceSize_; + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +REGISTER_TILING_TEMPLATE("DequantSwigluQuant", DequantSwigluQuantDskTiling, 0); + +ge::graphStatus TilingForDequantSwigluQuant(gert::TilingContext* context) { + return TilingRegistry::GetInstance().DoTilingImpl(context); +} + +ge::graphStatus TilingPrepareForDequantSwigluQuant(gert::TilingParseContext* context) { + OP_LOGD(context, "TilingPrepare4DequantSwigluQuant enter."); + auto compileInfo = context->GetCompiledInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo); + auto platformInfo = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + compileInfo->coreNum = ascendcPlatform.GetCoreNumAiv(); + OP_CHECK_IF((compileInfo->coreNum <= 0), + OP_LOGE(context->GetNodeName(), "Get core num failed, core num: %u", + static_cast(compileInfo->coreNum)), + return ge::GRAPH_FAILED); + + uint64_t ubSize; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); + compileInfo->ubSize = ubSize; + OP_CHECK_IF((compileInfo->ubSize <= 0), + OP_LOGE(context->GetNodeName(), "Get ub size failed, ub size: %u", + static_cast(compileInfo->ubSize)), + return ge::GRAPH_FAILED); + + OP_LOGD(context, "TilingPrepare4DequantSwigluQuant exit."); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(DequantSwigluQuant) + .Tiling(TilingForDequantSwigluQuant) + .TilingParse(TilingPrepareForDequantSwigluQuant); + +} // namespace optiling diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling.h b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling.h new file mode 100644 index 000000000..36f82141a --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling.h @@ -0,0 +1,478 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_tiling.h + * \brief + */ + +#ifndef DEQUANT_SWIGLU_QUANT_TILING_H +#define DEQUANT_SWIGLU_QUANT_TILING_H + + +#include +#include +#include "register/op_impl_registry.h" +#include "util/math_util.h" +#include "log/log.h" +#include "tiling/platform/platform_ascendc.h" +#include "platform/platform_infos_def.h" +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "dequant_swiglu_quant_proto.h" +#include "../tiling_base/tiling_base.h" +#include "../tiling_base/tiling_templates_registry.h" + +namespace optiling +{ +BEGIN_TILING_DATA_DEF(DequantSwigluQuantBaseTilingData) +TILING_DATA_FIELD_DEF(int64_t, inDimx); +TILING_DATA_FIELD_DEF(int64_t, inDimy); +TILING_DATA_FIELD_DEF(int64_t, outDimy); +TILING_DATA_FIELD_DEF(int64_t, UbFactorDimx); +TILING_DATA_FIELD_DEF(int64_t, UbFactorDimy); // cut for output dim +TILING_DATA_FIELD_DEF(int64_t, usedCoreNum); +TILING_DATA_FIELD_DEF(int64_t, maxCoreNum); +TILING_DATA_FIELD_DEF(int64_t, inGroupNum); +TILING_DATA_FIELD_DEF(int64_t, hasBias); +TILING_DATA_FIELD_DEF(int64_t, quantMode); +TILING_DATA_FIELD_DEF(int64_t, actRight); +TILING_DATA_FIELD_DEF(int64_t, quantScaleDtype); +TILING_DATA_FIELD_DEF(int64_t, groupIndexDtype); +TILING_DATA_FIELD_DEF(int64_t, needSmoothScale); +TILING_DATA_FIELD_DEF(int64_t, biasDtype); +TILING_DATA_FIELD_DEF(int64_t, speGroupType); +TILING_DATA_FIELD_DEF(int64_t, activationScaleIsEmpty); +TILING_DATA_FIELD_DEF(int64_t, quantIsOne); +// data field for SwiGLU used by GPT-OSS +TILING_DATA_FIELD_DEF(int64_t, swigluMode); +TILING_DATA_FIELD_DEF(float, clampLimit); +TILING_DATA_FIELD_DEF(float, gluAlpha); +TILING_DATA_FIELD_DEF(float, gluBias); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100000000, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100001000, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100002000, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100003000, DequantSwigluQuantBaseTilingData) + +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100000100, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100001100, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100002100, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100003100, DequantSwigluQuantBaseTilingData) + +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100000200, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100001200, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100002200, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100003200, DequantSwigluQuantBaseTilingData) + +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_200000000, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_200000100, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_200000200, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_110000000, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_110000100, DequantSwigluQuantBaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_110000200, DequantSwigluQuantBaseTilingData) + +BEGIN_TILING_DATA_DEF(DequantSwigluQuantV35BaseTilingData) +TILING_DATA_FIELD_DEF(int64_t, inDimx); +TILING_DATA_FIELD_DEF(int64_t, inDimy); +TILING_DATA_FIELD_DEF(int64_t, outDimy); +TILING_DATA_FIELD_DEF(int64_t, UbFactorDimx); +TILING_DATA_FIELD_DEF(int64_t, UbFactorDimy); // cut for output dim +TILING_DATA_FIELD_DEF(int64_t, usedCoreNum); +TILING_DATA_FIELD_DEF(int64_t, maxCoreNum); +TILING_DATA_FIELD_DEF(int64_t, inGroupNum); +TILING_DATA_FIELD_DEF(int64_t, quantMode); +TILING_DATA_FIELD_DEF(int64_t, actRight); // swish的激活与门控左右排布情况下生效，1表示右半部为激活 +TILING_DATA_FIELD_DEF(int64_t, dstType); +TILING_DATA_FIELD_DEF(int64_t, roundMode); +TILING_DATA_FIELD_DEF(int64_t, activateDim); +TILING_DATA_FIELD_DEF(int64_t, loopTimesPerRow); // 非全载模板下处理一行需要的UB循环次数 +TILING_DATA_FIELD_DEF(int64_t, tailPerRow); // 非全载模板UB循环最后一次的元素个数 +TILING_DATA_FIELD_DEF(int64_t, swiGluMode); // 0表示swish的激活与门控左右排布，1表示奇偶排布 +TILING_DATA_FIELD_DEF(int64_t, biasMode); // bias类型，0：不存在；1：int32；2：int64 +TILING_DATA_FIELD_DEF(int64_t, groupIndexMode); // group_index类型，0：不存在；1：int32；2：int64 +TILING_DATA_FIELD_DEF(int64_t, quantIsOne); // kernel侧计算时quant尾轴是否为单个元素 +TILING_DATA_FIELD_DEF(int64_t, speGroupType); //groupidx是否2维 +TILING_DATA_FIELD_DEF(int64_t, isSpecialCoreCut); // 是否多专家少token场景 +TILING_DATA_FIELD_DEF(float, clampLimit); +TILING_DATA_FIELD_DEF(float, gluAlpha); +TILING_DATA_FIELD_DEF(float, gluBias); +END_TILING_DATA_DEF; + +// static quant full +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_10000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_10001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_10010, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_10011, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_10100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_10101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_10110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_10111, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_11111, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_12111, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_13111, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_14111, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_11110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_12110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_13110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_14110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_11101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_12101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_13101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_14101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_11100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_12100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_13100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_14100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_11011, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_12011, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_13011, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_14011, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_11010, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_12010, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_13010, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_14010, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_11001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_12001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_13001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_14001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_11000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_12000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_13000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_14000, DequantSwigluQuantV35BaseTilingData) +// static quant not full +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1000000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1000010, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1000100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1000110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1001000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1001010, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1001100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1001110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1010000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1010010, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1010100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1010110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1011000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1011010, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1011100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1011110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1000001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1000011, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1000101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1000111, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1001001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1001011, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1001101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1001111, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1010001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1010011, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1010101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1010111, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1011001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1011011, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1011101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1011111, DequantSwigluQuantV35BaseTilingData) +// ## dynamic +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1100000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1100001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1100100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1100101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1101000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1101001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1101100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1101101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1110000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1110001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1110100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1110101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1111000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1111001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1111100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1111101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1120000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1120001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1120100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1120101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1121000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1121001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1121100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1121101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1130000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1130001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1130100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1130101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1131000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1131001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1131100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1131101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1140000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1140001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1140100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1140101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1141000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1141001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1141100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1141101, DequantSwigluQuantV35BaseTilingData) + +BEGIN_TILING_DATA_DEF(DequantSwigluQuantV35NlastTilingData) +TILING_DATA_FIELD_DEF(int64_t, inDim0); +TILING_DATA_FIELD_DEF(int64_t, inDim1); +TILING_DATA_FIELD_DEF(int64_t, inDim2); +TILING_DATA_FIELD_DEF(int64_t, outDim1); +TILING_DATA_FIELD_DEF(int64_t, blockNum0); +TILING_DATA_FIELD_DEF(int64_t, blockNum1); +TILING_DATA_FIELD_DEF(int64_t, blockFormer0); +TILING_DATA_FIELD_DEF(int64_t, blockFormer1); +TILING_DATA_FIELD_DEF(int64_t, ubFormer0); +TILING_DATA_FIELD_DEF(int64_t, ubFormer1); +TILING_DATA_FIELD_DEF(int64_t, ubLoopOfFormerBlock0); +TILING_DATA_FIELD_DEF(int64_t, ubLoopOfFormerBlock1); +TILING_DATA_FIELD_DEF(int64_t, ubLoopOfTailBlock0); +TILING_DATA_FIELD_DEF(int64_t, ubLoopOfTailBlock1); +TILING_DATA_FIELD_DEF(int64_t, ubTailOfFormerBlock0); +TILING_DATA_FIELD_DEF(int64_t, ubTailOfFormerBlock1); +TILING_DATA_FIELD_DEF(int64_t, ubTailOfTailBlock0); +TILING_DATA_FIELD_DEF(int64_t, ubTailOfTailBlock1); +TILING_DATA_FIELD_DEF(int64_t, actRight); +TILING_DATA_FIELD_DEF(int64_t, roundMode); +END_TILING_DATA_DEF; + +struct DequantSwigluQuantCompileInfo { + uint64_t coreNum = 0; + uint64_t ubSize = 0; +}; + +class DequantSwigluQuantDskTiling : public TilingBaseClass +{ + public: + explicit DequantSwigluQuantDskTiling(gert::TilingContext* tilingContext) : TilingBaseClass(tilingContext) + { + } + ~DequantSwigluQuantDskTiling() override + { + } + uint64_t coreNum_ = 0; + uint64_t ubSize_ = 0; + int64_t groupNum_ = 0; + int64_t actRight_ = 0; + int64_t quantMode_ = 0; + uint64_t workspaceSize_ = 0; + int64_t maxPreCore_ = 0; + bool hasWeightScale_ = false; + bool hasActivationScale_ = false; + bool hasBias_ = false; + bool hasQuantScale_ = false; + bool hasQuantOffset_ = false; + bool hasGroupIndex_ = false; + bool speGroupType_ = false; + + // variable for SwiGLU used by GPT-OSS + int64_t swigluMode_ = 0; + float clampLimit_ = 0.0; + float gluAlpha_ = 0.0; + float gluBias_ = 0.0; + + protected: + bool IsCapable() override; + ge::graphStatus GetPlatformInfo() override; + ge::graphStatus GetShapeAttrsInfo() override; + ge::graphStatus DoOpTiling() override; + ge::graphStatus DoLibApiTiling() override; + uint64_t GetTilingKey() const override; + ge::graphStatus GetWorkspaceSize() override; + ge::graphStatus PostTiling() override; + void DumpTilingInfo() override; + ge::graphStatus GetAttr(); + ge::graphStatus CheckBias(); + ge::graphStatus CheckWeightScale(); + ge::graphStatus CheckActivationScale(); + ge::graphStatus CheckXAndGroupIndexDtype(); + ge::graphStatus CheckForDequant(); + ge::graphStatus CheckForQuant(); + ge::graphStatus CheckForDynamicQuant(); + ge::graphStatus CheckForStaticQuant(); + ge::graphStatus CheckQuantScaleDtype(); + ge::graphStatus CheckStaticQuantShape(const int64_t quantInputIdx, int64_t& colLen, const char* paramName); + ge::graphStatus CheckIllegalParam(); + void CountTilingKey(); + ge::graphStatus CountMaxDim(int64_t& ubFactorDimx); + ge::graphStatus CheckScaleShapeWithDim(const int64_t scaleInputIdx, const int64_t expectDim, const char* paramName); + bool IsPerformanceAndGroupIndexBrach(); + ge::graphStatus GetShapeAttrsInfoInner(); + static bool CheckOptionalShapeExisting(const gert::StorageShape* storageShape); + + private: + uint64_t tilingKey_ = 0; + DequantSwigluQuantBaseTilingData tilingData_; + int64_t inDimx_ = 0; + int64_t inDimy_ = 0; + int64_t outDimy_ = 0; + platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND910B; +}; + +template +inline auto AlignUp(T num, T rnd) -> decltype(num) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd) * (rnd))); +} +// align num to multiples of rnd, round down +template +inline auto AlignDown(T num, T rnd) -> decltype(num) +{ + return ((((rnd) == 0) || ((num) < (rnd))) ? 0 : ((num) / (rnd) * (rnd))); +} + +template +inline auto DivCeil(T num, T div) -> decltype(num) +{ + return (((div) == 0) ? 0 : (((num) + (div)-1) / (div))); +} + +inline bool GetLengthByType(int32_t dtype, uint32_t& dsize) +{ + switch (dtype) { + case ge::DT_FLOAT16: + case ge::DT_INT16: + case ge::DT_UINT16: + case ge::DT_BF16: + dsize = sizeof(int16_t); + return true; + case ge::DT_FLOAT: + case ge::DT_INT32: + case ge::DT_UINT32: + dsize = sizeof(int32_t); + return true; + case ge::DT_DOUBLE: + case ge::DT_INT64: + case ge::DT_UINT64: + dsize = sizeof(int64_t); + return true; + default: + return false; + } +} + +class DequantSwigluQuantV35DskTiling : public TilingBaseClass { + public: + explicit DequantSwigluQuantV35DskTiling(gert::TilingContext* tilingContext) : TilingBaseClass(tilingContext) { + } + ~DequantSwigluQuantV35DskTiling() override { + } + uint64_t coreNum_ = 0; + uint64_t ubSize_ = 0; + int64_t actRight_ = 0; + int64_t quantMode_ = 0; + uint64_t workspaceSize_ = 0; + int64_t maxPreCore_ = 0; + int64_t groupNum_ = 0; + int64_t biasMode_ = 0; + int64_t groupIndexMode_ = 0; + int64_t swigluMode_ = 0; + int64_t speGroupType_ = 0; + int64_t isSpecialCoreCut_ = 0; + float clampLimit_ = 0; + float gluAlpha_ = 1.702; + float gluBias_ = 1.0; + bool hasWeightScale_ = false; + bool hasActivationScale_ = false; + bool hasBias_ = false; + bool hasQuantScale_ = false; + bool hasQuantOffset_ = false; + bool quantIsOne_ = true; + bool hasGroupIndex_ = false; + gert::Shape xShape_ = gert::Shape(); + size_t xDimNum_ = 0; + gert::Shape groupIndexShape_ = gert::Shape(); + int64_t dstType_ = 2; + int64_t roundMode_ = 0; + int64_t activateDim_ = -1UL; + + protected: + bool IsCapable() override; + ge::graphStatus GetPlatformInfo() override; + ge::graphStatus GetShapeAttrsInfo() override; + ge::graphStatus DoOpTiling() override; + ge::graphStatus DoLibApiTiling() override; + uint64_t GetTilingKey() const override; + ge::graphStatus GetWorkspaceSize() override; + ge::graphStatus PostTiling() override; + ge::graphStatus GetAttr(); + ge::graphStatus GetInputX(); + ge::graphStatus GetAttrActivateDim(); + ge::graphStatus CheckInputWeightScale(); + ge::graphStatus CheckInputActScale(); + ge::graphStatus CheckInputBias(); + ge::graphStatus CheckInputQuantScale(); + ge::graphStatus CheckInputQuantOffset(); + ge::graphStatus CheckForStaticQuant(); + ge::graphStatus GetInputGroupIndex(); + ge::graphStatus CheckOutputY(); + ge::graphStatus CheckOutputScale(); + ge::graphStatus DoOpTilingNotFull(); + void CalcTilingKeyForNotFull(); + + private: + uint64_t tilingKey_ = 0; + DequantSwigluQuantV35BaseTilingData tilingData_; + int64_t inDimx_ = 0; + int64_t inDimy_ = 0; + int64_t outDimy_ = 0; + platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND910B; + }; + +class DequantSwigluQuantV35NlastTiling : public TilingBaseClass { + public: + explicit DequantSwigluQuantV35NlastTiling(gert::TilingContext* tilingContext) : TilingBaseClass(tilingContext) { + } + ~DequantSwigluQuantV35NlastTiling() override { + } + uint64_t coreNum_ = 0; + uint64_t ubSize_ = 0; + protected: + bool IsCapable() override; + ge::graphStatus GetPlatformInfo() override; + ge::graphStatus GetShapeAttrsInfo() override; + ge::graphStatus DoOpTiling() override; + ge::graphStatus DoLibApiTiling() override; + uint64_t GetTilingKey() const override; + ge::graphStatus GetWorkspaceSize() override; + ge::graphStatus PostTiling() override; + void FusedShape(); + void DoBlockSplit(); + bool DoUbSplit(); + + private: + uint64_t tilingKey_ = 0; + uint64_t workspaceSize_ = 0; + int32_t actDimIndex_ = 0; + int64_t actRight_ = 0; + int64_t roundMode_ = 0; + gert::Shape xShape_ = gert::Shape(); + int64_t inDim0_ = 1; + int64_t inDim1_ = 1; + int64_t inDim2_ = 1; + int64_t outDim1_ = 1; + int64_t blockFormer0_ = 0; + int64_t blockNum0_ = 0; + int64_t blockFormer1_ = 0; + int64_t blockNum1_ = 0; + int64_t blockNum_ = 0; + int64_t ubFormer0_ = 0; + int64_t ubFormer1_ = 0; + int64_t biasDtypeValue_ = 0; + + DequantSwigluQuantV35NlastTilingData tilingData_; + platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND910B; +}; + +} // namespace optiling +#endif // DEQUANT_SWIGLU_QUANT_TILING_H \ No newline at end of file diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_arch35.cpp b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_arch35.cpp new file mode 100644 index 000000000..09459867c --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_arch35.cpp @@ -0,0 +1,1378 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_tiling_arch35.cpp + * \brief + */ + +#include +#include "dequant_swiglu_quant_tiling.h" +#include "../tiling_base/tiling_templates_registry.h" +#include "dequant_swiglu_quant_tiling.h" +#include "../tiling_base/tiling_base.h" +#include "../tiling_base/tiling_util.h" + +using namespace AscendC; +using namespace ge; +namespace optiling { + +constexpr int64_t ATTR_ACTIVATE_LEFT_INDEX = 0; +constexpr int64_t ATTR_QUANT_MODE_INDEX = 1; +constexpr int64_t ATTR_DST_TYPE_INDEX = 2; +constexpr int64_t ATTR_ROUND_MODE_INDEX = 3; +constexpr int64_t ATTR_ACTIVATE_DIM_INDEX = 4; +constexpr int64_t ATTR_SWIGLU_MODE_INDEX = 5; +constexpr int64_t ATTR_CLAMP_LIMIT_INDEX = 6; +constexpr int64_t ATTR_GLU_ALPHA_INDEX = 7; +constexpr int64_t ATTR_GLU_BIAS_INDEX = 8; +constexpr int64_t X_INDEX = 0; +constexpr int64_t WEIGHT_SCALE_INDEX = 1; +constexpr int64_t ACTIVATION_SCALE_INDEX = 2; +constexpr int64_t BIAS_INDEX = 3; +constexpr int64_t QUANT_SCALE_INDEX = 4; +constexpr int64_t QUANT_OFFSET_INDEX = 5; +constexpr int64_t INPUT_GROUP_INDEX = 6; +constexpr int64_t Y_INDEX = 0; +constexpr int64_t SCALE_INDEX = 1; +constexpr int64_t BLOCK_SIZE = 32; +constexpr int64_t BLOCK_ELEM_B32 = BLOCK_SIZE / static_cast(sizeof(float)); +constexpr int64_t BLOCK_ELEM_B16 = BLOCK_SIZE / static_cast(sizeof(int16_t)); +constexpr int64_t BLOCK_ELEM_B8 = BLOCK_SIZE / static_cast(sizeof(int8_t)); +constexpr size_t SYS_WORK_SPACE_SIZE = static_cast(16 * 1024 * 1024); +constexpr uint64_t WORKSPACE_SIZE = 32; +constexpr int64_t UB_REVERSE = 1024; +constexpr int64_t SWI_FACTOR = 2; +constexpr int64_t Y_LAST_DIM_FULL_LOAD_MAX_VALUE = 5120; // 能够命中UB全载模板的输出尾轴最大值 +constexpr int64_t QUANT_MODE_DYNAMIC = 1; +constexpr int64_t QUANT_MODE_INDEX = 1; +constexpr int64_t ACTIVATE_DIM_FACTOR = 100000; +constexpr int64_t INPUT_X_FACTOR = 10000; +constexpr int64_t BIAS_FACTOR = 1000; +constexpr int64_t ACT_SCALE_FACTOR = 100; +constexpr int64_t QUANT_SCALE_FACTOR = 10; +constexpr int64_t GROUP_INDEX_FACTOR = 1; +constexpr int64_t DIM_TWO = 2; +constexpr int64_t PLACEHOLDER = 1000000; +constexpr int64_t QUANT_MODE_FACTOR = 100000; +constexpr int64_t BIAS_FACTOR_FOR_NOT_FULL = 10000; +constexpr int64_t ACTIVATE_FACTOR_FOR_NOT_FULL = 1000; +constexpr int64_t QUANT_SCALE_FACTOR_FOR_NOT_FULL = 100; +constexpr int64_t QUANT_OFFSET_FACTOR_FOR_NOT_FULL = 10; +constexpr int64_t SPECIAL_GROUP_NUM_64 = 64; // groupIndex存在时走特殊分核的group条件 +constexpr int64_t SPECIAL_GROUP_NUM_32 = 32; // groupIndex不存在时走特殊分核的group条件 +constexpr int64_t SPECIAL_GROUP_NUM_16 = 16; // 走特殊分核场景的分组条件 +constexpr float CLAMP_LIMIT_DEFAULT= 0.0; +constexpr float GLU_ALPHA_DEFAULT = 1.702; +constexpr float GLU_BIAS_DEFAULT = 1.0; +static const gert::Shape g_vec_1_shape = {1}; + +inline const gert::Shape &EnsureNotScalar(const gert::Shape &in_shape) { + if (in_shape.IsScalar()) { + return g_vec_1_shape; + } + return in_shape; +} + +static const std::set SUPPORT_DTYPE = {ge::DT_INT32, ge::DT_BF16, ge::DT_FLOAT16}; +static const std::map SUPPORT_QUANT_MODE = {{"dynamic", 1}, {"static", 0}}; +// 定义bias支持的所有类型 +static const std::set BIAS_SUPPORT_DTYPE = {ge::DT_INT32, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT}; +static const std::map SUPPORT_BIAS_MODE = {{ge::DT_INT32, 1}, {ge::DT_BF16, 2}, {ge::DT_FLOAT16, 3}, {ge::DT_FLOAT, 4}}; +// 定义quant_scale支持的所有类型 +static const std::set QUANT_SCALE_SUPPORT_DTYPE = {ge::DT_FLOAT}; +// 定义quant_offset支持的所有类型 +static const std::set QUANT_OFFSET_SUPPORT_DTYPE = {ge::DT_FLOAT}; +// 定义输出y支持的所有类型:int8, hifloat8, float8的两种类型, float4的两种类型 +static const std::set OUTPUT_SUPPORT_DTYPE = {ge::DT_INT8, ge::DT_HIFLOAT8, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2}; +// 定义roundMode映射表。 +static const std::map SUPPORT_ROUND_MODE = {{"rint", 0}, {"round", 1}, {"floor", 2}, {"ceil", 3}, {"trunc", 4}}; + +ge::graphStatus DequantSwigluQuantV35DskTiling::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + if (platformInfo == nullptr) { + auto compileInfoPtr = context_->GetCompileInfo(); + OP_CHECK_IF(compileInfoPtr == nullptr, OP_LOGE(context_, "compile info is null"), + return ge::GRAPH_FAILED); + coreNum_ = compileInfoPtr->coreNum; + ubSize_ = compileInfoPtr->ubSize; + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + coreNum_ = ascendcPlatform.GetCoreNumAiv(); + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + ubSize_ = ubSizePlatForm; + socVersion = ascendcPlatform.GetSocVersion(); + } + + maxPreCore_ = static_cast(coreNum_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::GetInputX() +{ + auto xDesc = context_->GetInputDesc(X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, xDesc); + ge::DataType xDType = xDesc->GetDataType(); + // 校验x的数据类型是否合法 + OP_CHECK_IF((SUPPORT_DTYPE.find(xDType) == SUPPORT_DTYPE.end()), + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "x", + ge::TypeUtils::DataTypeToSerialString(xDType).c_str(), "int32, float16 or bf16"), + return ge::GRAPH_FAILED); + + auto xStorageShape = context_->GetInputShape(X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, xStorageShape); + xShape_ = EnsureNotScalar(xStorageShape->GetStorageShape()); + xDimNum_ = xShape_.GetDimNum(); + OP_CHECK_IF(xDimNum_ < 2, + OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), "x", + std::to_string(xDimNum_).c_str(), "greater than or equal to 2"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(xDimNum_ > 8, + OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), "x", + std::to_string(xDimNum_).c_str(), "less than or equal to 8"), + return ge::GRAPH_FAILED); + for (size_t i = 0; i < xDimNum_; i++) { + OP_CHECK_IF(xShape_.GetDim(i) <= 0, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), "x", + Ops::Base::ToString(xShape_).c_str(), + ("the dim[" + std::to_string(i) + "] of x must be positive").c_str()), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::GetInputGroupIndex() +{ + auto groupIndexDesc = context_->GetOptionalInputDesc(INPUT_GROUP_INDEX); + if (groupIndexDesc != nullptr) { + ge::DataType groupIndexDType = groupIndexDesc->GetDataType(); + OP_CHECK_IF(groupIndexDType != ge::DT_INT64 && groupIndexDType != ge::DT_INT32, + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "group_index", + ge::TypeUtils::DataTypeToSerialString(groupIndexDType).c_str(), "int32 or int64"), + return ge::GRAPH_FAILED); + + auto groupIndexStorageShape = context_->GetOptionalInputShape(INPUT_GROUP_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, groupIndexStorageShape); + groupIndexShape_ = EnsureNotScalar(groupIndexStorageShape->GetStorageShape()); + auto groupIndexDimNum = groupIndexShape_.GetDimNum(); + OP_CHECK_IF((groupIndexDimNum != 1) && (groupIndexDimNum != 2), + OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), "group_index", + std::to_string(groupIndexDimNum).c_str(), "1 or 2"), + return ge::GRAPH_FAILED); + + groupNum_ = groupIndexShape_.GetDim(0); + OP_CHECK_IF(groupNum_ < 1, + OP_LOGE_FOR_INVALID_SHAPESIZE(context_->GetNodeName(), "group_index", + std::to_string(groupNum_).c_str(), "group_index[0] must be greater than or equal to 1"), + return ge::GRAPH_FAILED); + + hasGroupIndex_ = true; + groupIndexMode_ = groupIndexDType == ge::DT_INT32 ? 1 : 2; // groupIndex：int32类型时mode设为1；int64类型时设为2；不存在则保持默认值0 + if (groupIndexDimNum == 2) { + speGroupType_ = 1; + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::GetAttrActivateDim() +{ + auto* attrs = context_->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context_, attrs); + // 校验activate_dim + auto* attrActivateDim = attrs->GetAttrPointer(ATTR_ACTIVATE_DIM_INDEX); + // 类型校验 + activateDim_ = (attrActivateDim != nullptr) ? *attrActivateDim : -1; + // 指定切分轴维度转换为正数 + activateDim_ = activateDim_ < 0 ? activateDim_ + static_cast(xDimNum_) : activateDim_; + + // 判断切分轴维度合法性 + OP_CHECK_IF(activateDim_ < 0 || activateDim_ >= static_cast(xDimNum_), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + context_->GetNodeName(), "activate_dim", + std::to_string(activateDim_).c_str(), + ("activate_dim must be in [-" + std::to_string(xDimNum_) + ", " + + std::to_string(xDimNum_ - 1) + "]").c_str()), + return ge::GRAPH_FAILED); + // 校验切分轴对应的shape是不是偶数 + OP_CHECK_IF(xShape_.GetDim(activateDim_) % 2 != 0, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), "x", + Ops::Base::ToString(xShape_).c_str(), + ("the split dim(" + std::to_string(activateDim_)+ "dimension) must be even").c_str()), + return ge::GRAPH_FAILED); + + //如果activateDim不是尾轴，则不允许输入group + if (activateDim_ != static_cast(xDimNum_ - 1)) { + OP_CHECK_IF(hasGroupIndex_ == true, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "group_index and activate_dim", + ("group_index is not None, and activate_dim is " + std::to_string(activateDim_)).c_str(), "group_index must be None when activate_dim is not the last dim of x"), + return ge::GRAPH_FAILED); + } + + // activate_dim对应在x的轴需要是偶数 + OP_CHECK_IF((xShape_.GetDim(activateDim_) % 2) != 0, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), "x", + Ops::Base::ToString(xShape_).c_str(), + "the x dimension of activateDim must be even"), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::CheckOutputY() +{ + auto yDesc = context_->GetOutputDesc(Y_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, yDesc); + ge::DataType yDType = yDesc->GetDataType(); + OP_CHECK_IF(OUTPUT_SUPPORT_DTYPE.find(yDType) == OUTPUT_SUPPORT_DTYPE.end(), + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "y", + ge::TypeUtils::DataTypeToSerialString(yDType).c_str(), + "int8, hifloat8, float8e4m3, float8e5m2, float4e2m1 or floate1m2"), + return ge::GRAPH_FAILED); + auto yStorageShape = context_->GetOutputShape(Y_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, yStorageShape); + auto& yShape = EnsureNotScalar(yStorageShape->GetStorageShape()); + const size_t yDimNum = yShape.GetDimNum(); + // 输出y是fp4类型时，y的尾轴对应的shape需要是偶数 + if (yDType == ge::DT_FLOAT4_E2M1 || yDType == ge::DT_FLOAT4_E1M2) { + OP_CHECK_IF((yShape.GetDim(xDimNum_ - 1) % 2) != 0, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), "y", + Ops::Base::ToString(yShape).c_str(), + "The last dim of y must be even when the type of y is FP4X2_E2M1 or FP4X2_E1M2"), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(yDimNum != xDimNum_, + OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), "y", + std::to_string(yDimNum).c_str(), + ("equal to x dimension " + std::to_string(xDimNum_)).c_str()), + return ge::GRAPH_FAILED); + for (size_t i = 0; i < yDimNum; i++) { + if (static_cast(i) != activateDim_){ + OP_CHECK_IF(yShape.GetDim(i) != xShape_.GetDim(i), + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(context_->GetNodeName(), "x and y", + (Ops::Base::ToString(xShape_) + "and" + Ops::Base::ToString(yShape)).c_str(), + ("dim[" + std::to_string(i) + "] of y must be equal to dim[" + std::to_string(i) + "] of x").c_str()), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF(yShape.GetDim(i) != xShape_.GetDim(i) / SWI_FACTOR, + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(context_->GetNodeName(), "x and y", + Ops::Base::ToString(yShape).c_str(), + ("dim[" + std::to_string(i) + "] of y must be equal to half of dim[" + std::to_string(i) + "] of x").c_str()), + return ge::GRAPH_FAILED); + } + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::CheckInputWeightScale() +{ + auto wScaleDesc = context_->GetOptionalInputDesc(WEIGHT_SCALE_INDEX); + auto xDesc = context_->GetInputDesc(X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, xDesc); + ge::DataType xDType = xDesc->GetDataType(); + // 如果输入x是bf16 or float16，则weight_scale需要为空，非法性校验 + if (wScaleDesc != nullptr) { + OP_CHECK_IF(xDType == ge::DT_FLOAT16 || xDType == ge::DT_BF16, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "weight_scale", + "not None", "weight_scale must be None when x's datatype is in [bfloat16, float16]"), + return ge::GRAPH_FAILED); + } + + // 如果输入x是int32，则weight_scale必须有值，合法性校验 + OP_CHECK_IF((xDType == ge::DT_INT32) && (wScaleDesc == nullptr), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "weight_scale", + "None", "weight_scale must be not None when x's datatype is int32"), + return ge::GRAPH_FAILED); + + // weight_scale不为空，进行判断 + if (wScaleDesc != nullptr) { + OP_CHECK_NULL_WITH_CONTEXT(context_, wScaleDesc); + ge::DataType wScaleDType = wScaleDesc->GetDataType(); + OP_CHECK_IF(wScaleDType != ge::DT_FLOAT, + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "weight_scale", + ge::TypeUtils::DataTypeToSerialString(wScaleDType).c_str(), "float32"), + return ge::GRAPH_FAILED); + + auto wScaleStorageShape = context_->GetOptionalInputShape(WEIGHT_SCALE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, wScaleStorageShape); + auto& wScaleShape = EnsureNotScalar(wScaleStorageShape->GetStorageShape()); + const size_t wScaleDimNum = wScaleShape.GetDimNum(); + OP_CHECK_IF(wScaleDimNum > 2, + OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), "weight_scale", + std::to_string(wScaleDimNum).c_str(), "less than or equal to 2"), + return ge::GRAPH_FAILED); + + if (wScaleDimNum == static_cast(1)) { + OP_CHECK_IF(hasGroupIndex_ == true, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "group_index", + "not None", "group_index should be none when weight_scale dimension is 1"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(wScaleShape.GetDim(0) != xShape_.GetDim(xDimNum_ - 1), + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), "weight_scale", + Ops::Base::ToString(wScaleShape).c_str(), + ("The first dim of weight_scale must be the same as the last dim of x: " + std::to_string(xShape_.GetDim(xDimNum_ - 1))).c_str()), + return ge::GRAPH_FAILED); + } + if (wScaleDimNum > static_cast(1)) { + if (hasGroupIndex_) { + OP_CHECK_IF(!(wScaleShape.GetDim(0) == groupIndexShape_.GetDim(0) && wScaleShape[wScaleDimNum - 1] == xShape_.GetDim(xDimNum_ - 1)), + OP_LOGE(context_->GetNodeName(), + "weight_scale shape[0] must be equal to group_index shape[0], and shape[-1] must be equal to x shape[-1] " + "when group_index exists, please check."), + return ge::GRAPH_FAILED); + } else { + OP_CHECK_IF(!(wScaleShape.GetDim(0) == 1 && wScaleShape.GetDim(wScaleDimNum - 1) == xShape_.GetDim(xDimNum_ - 1)) && + !(wScaleShape.GetDim(0) == xShape_.GetDim(xDimNum_ - 1) && wScaleShape.GetDim(wScaleDimNum - 1) == 1), + OP_LOGE(context_->GetNodeName(), "weight_scale shape must be in {[1, %ld], [%ld, 1]} when weight_scale dimension == 2 and group_index not exists," + "please check.", xShape_.GetDim(xDimNum_ - 1), xShape_.GetDim(xDimNum_ - 1)), + return ge::GRAPH_FAILED); + } + } + hasWeightScale_ = true; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::CheckInputActScale() +{ + auto aScaleDesc = context_->GetOptionalInputDesc(ACTIVATION_SCALE_INDEX); + auto xDesc = context_->GetInputDesc(X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, xDesc); + ge::DataType xDType = xDesc->GetDataType(); + // 当x：bfloat16 or float16时，activate_scale需要为空 + if (aScaleDesc != nullptr) { + OP_CHECK_IF(xDType == ge::DT_FLOAT16 || xDType == ge::DT_BF16, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "activation_scale", + "not None", "activate_scale must be None when x's datatype is in [bfloat16, float16]"), + return ge::GRAPH_FAILED); + + ge::DataType aScaleDType = aScaleDesc->GetDataType(); + OP_CHECK_IF(aScaleDType != ge::DT_FLOAT, + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "activation_scale", + ge::TypeUtils::DataTypeToSerialString(aScaleDType).c_str(), "float32"), + return ge::GRAPH_FAILED); + + auto aScaleStorageShape = context_->GetOptionalInputShape(ACTIVATION_SCALE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, aScaleStorageShape); + auto& aScaleShape = EnsureNotScalar(aScaleStorageShape->GetStorageShape()); + const size_t aScaleDimNum = aScaleShape.GetDimNum(); + + OP_CHECK_IF(aScaleDimNum <= 0, + OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), "activation_scale", + std::to_string(aScaleDimNum).c_str(), "greater than 0"), + return ge::GRAPH_FAILED); + + // shape校验 + // activation_scale的shape size与x除尾轴外的shape size一致 + int64_t aScaleSize = aScaleStorageShape->GetStorageShape().GetShapeSize(); + int64_t xSizeWithoutLastDim = xShape_.GetShapeSize() / xShape_.GetDim(xDimNum_ - 1); + OP_CHECK_IF(aScaleSize != xSizeWithoutLastDim, + OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON(context_->GetNodeName(), "activation_scale", + std::to_string(aScaleSize).c_str(), + ("The shape size of activation_scale should be the same as x's shape size without last dim " + std::to_string(xSizeWithoutLastDim)).c_str()), + return ge::GRAPH_FAILED); + hasActivationScale_ = true; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::CheckInputBias() +{ + auto biasDesc = context_->GetOptionalInputDesc(BIAS_INDEX); + auto xDesc = context_->GetInputDesc(X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, xDesc); + ge::DataType xDType = xDesc->GetDataType(); + // 首先判断bias不为空时，其数据类型是不是满足计算要求 + if (biasDesc != nullptr) { + ge::DataType biasDtype = biasDesc->GetDataType(); + auto it = SUPPORT_BIAS_MODE.find(biasDtype); + OP_CHECK_IF(it == SUPPORT_BIAS_MODE.end(), + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "bias", + ge::TypeUtils::DataTypeToSerialString(biasDtype).c_str(), "float16, float, bf16 or int32"), + return ge::GRAPH_FAILED); + biasMode_ = it->second; + // 当前bias支持四种数据类型，但是有些bias类型仅支持x的特定类型 + // x：bf16, float16，bias不支持输入 + if (xDType == ge::DT_BF16 or xDType == ge::DT_FLOAT16) { + OP_CHECK_IF(BIAS_SUPPORT_DTYPE.find(biasDtype) != BIAS_SUPPORT_DTYPE.end(), + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "bias", + ge::TypeUtils::DataTypeToSerialString(biasDtype).c_str(), + "bias not support when the type of x is bf16 or float16"), + return ge::GRAPH_FAILED); + } + + // 然后判断bias的维度是不是满足要求 + auto biasStorageShape = context_->GetOptionalInputShape(BIAS_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, biasStorageShape); + auto& biasShape = EnsureNotScalar(biasStorageShape->GetStorageShape()); + const size_t biasDimNum = biasShape.GetDimNum(); + OP_CHECK_IF(biasDimNum > static_cast(2) || biasDimNum == static_cast(0), + OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), "bias", + std::to_string(biasDimNum).c_str(), "1D or 2D"), + return ge::GRAPH_FAILED); + // 当biasDimNum=1时 + if (biasDimNum == static_cast(1)) { + OP_CHECK_IF(hasGroupIndex_ == true, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), "group_index", + "not None", "group_index should be none when bias dimension is 1"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(biasShape.GetDim(0) != xShape_.GetDim(xDimNum_ - 1), + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), "bias", + Ops::Base::ToString(biasShape).c_str(), + ("The last dimension of bias should be the same as the last dimension of x " + + std::to_string(xShape_.GetDim(xDimNum_ - 1))).c_str()), + return ge::GRAPH_FAILED); + } + + // 当biasDimNum=2时 + if (biasDimNum == static_cast(2)) { + OP_CHECK_IF(biasShape.GetDim(1) != xShape_.GetDim(xDimNum_ - 1), + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), "bias", + Ops::Base::ToString(biasShape).c_str(), + ("The last dimension of bias should be the same as the last dimension of x " + + std::to_string(xShape_.GetDim(xDimNum_ - 1))).c_str()), + return ge::GRAPH_FAILED); + if (hasGroupIndex_) { + if (biasShape.GetDim(0) != groupNum_) { + std::string reason = + "when the dimension of bias is 2 and group_index exists, the first dimension of bias (" + + std::to_string(biasShape.GetDim(0)) + + ") should be equal to the first dimension of group_index (" + std::to_string(groupNum_) + + "), please check"; + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON( + context_->GetNodeName(), "bias", Ops::Base::ToString(biasShape).c_str(), reason.c_str()); + return ge::GRAPH_FAILED; + } + } else { + OP_CHECK_IF(biasShape.GetDim(0) != 1, + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), "bias", + std::to_string(biasShape.GetDim(0)).c_str(), + "The first dimension of bias should be 1 when the dimension of bias is 2 and group_index does not exist"), + return ge::GRAPH_FAILED); + } + } + hasBias_ = true; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::CheckInputQuantScale() +{ + auto qScaleDesc = context_->GetOptionalInputDesc(QUANT_SCALE_INDEX); + if (qScaleDesc != nullptr) { + ge::DataType qScaleDType = qScaleDesc->GetDataType(); + OP_CHECK_IF(QUANT_SCALE_SUPPORT_DTYPE.find(qScaleDType) == QUANT_SCALE_SUPPORT_DTYPE.end(), + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "quant_scale", + ge::TypeUtils::DataTypeToSerialString(qScaleDType).c_str(), "float32"), + return ge::GRAPH_FAILED); + + auto qScaleStorageShape = context_->GetOptionalInputShape(QUANT_SCALE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, qScaleStorageShape); + auto& qScaleShape = EnsureNotScalar(qScaleStorageShape->GetStorageShape()); + const size_t qScaleDimNum = qScaleShape.GetDimNum(); + OP_CHECK_IF(qScaleDimNum > 2, + OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), "quant_scale", + std::to_string(qScaleDimNum).c_str(), "less than or equal to 2"), + return ge::GRAPH_FAILED); + + // 获取y的shape + auto yStorageShape = context_->GetOutputShape(Y_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, yStorageShape); + auto& yShape = EnsureNotScalar(yStorageShape->GetStorageShape()); + if (hasGroupIndex_) { + if (quantMode_ == 0) { + OP_CHECK_IF(qScaleShape.GetDim(0) != (groupIndexShape_.GetDim(0)), + OP_LOGE(context_->GetNodeName(), + "quant_scale shape[0] must be equal to group_index shape[0] when static_quant and group_index exists, please check."), + return ge::GRAPH_FAILED); + if (qScaleDimNum == DIM_TWO) { + OP_CHECK_IF(!((qScaleShape.GetDim(qScaleDimNum - 1) == yShape.GetDim(xDimNum_ - 1)) || (qScaleShape.GetDim(qScaleDimNum - 1) == 1)), + OP_LOGE(context_->GetNodeName(), + "quant_scale shape[-1] must be equal to or can be broadcast to y shape[-1] when static_quant and group_index exists, please check."), + return ge::GRAPH_FAILED); + } + } else if (quantMode_ == 1) { + OP_CHECK_IF(qScaleShape.GetDim(0) != groupIndexShape_.GetDim(0) || qScaleShape.GetDim(qScaleDimNum - 1) != yShape.GetDim(xDimNum_ - 1), + OP_LOGE(context_->GetNodeName(), + "quant_scale shape must be [ group_index_shape[0], y_shape[-1] ] when dynamic quant and group_index exists, please check."), + return ge::GRAPH_FAILED); + } + quantIsOne_ = (qScaleDimNum == DIM_TWO && qScaleShape.GetDim(qScaleDimNum - 1) == yShape.GetDim(xDimNum_ - 1)) ? 0 : 1; + } else { + if (qScaleDimNum == DIM_TWO) { + OP_CHECK_IF(qScaleShape.GetDim(0) != 1, + OP_LOGE(context_->GetNodeName(), + "if dim of quant_scale is 2, shape[0] must be [1] when group_index not exists, please check."), + return ge::GRAPH_FAILED); + } + if (quantMode_ == 0) { + OP_CHECK_IF(qScaleShape.GetDim(0) != 1 && qScaleShape.GetDim(0) != yShape.GetDim(xDimNum_ - 1), + OP_LOGE(context_->GetNodeName(), + "quant_scale shape[0] must be or can be broadcast to y shape[-1] when static_quant and group_index not exists, please check."), + return ge::GRAPH_FAILED); + } else if (quantMode_ == 1) { + OP_CHECK_IF(qScaleShape.GetDim(qScaleDimNum - 1) != yShape.GetDim(xDimNum_ - 1), + OP_LOGE(context_->GetNodeName(), + "quant_scale shape[-1] must be equal to y shape[-1] when dynamic_quant and group_index not exists, please check."), + return ge::GRAPH_FAILED); + } + quantIsOne_ = qScaleShape.GetDim(qScaleDimNum - 1) == yShape.GetDim(xDimNum_ - 1) ? 0 : 1; + } + hasQuantScale_ = true; + } + if (quantMode_ == 0) { + OP_CHECK_IF(!hasQuantScale_, + OP_LOGE(context_->GetNodeName(), "quant_scale must exist when static_quant, please check."), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::CheckInputQuantOffset() +{ + auto qOffsetDesc = context_->GetOptionalInputDesc(QUANT_OFFSET_INDEX); + /* + 校验： + 1、是静态量化 + 2、activate_dim=-1 + 3、有group_index时shape维度≤2 + 3.1、1维时：shape为[G] + 3.2、2维时：shape为[G,1]/[G,H] + 4、无group_index时shape为[1]/[H]/[1,H] + */ + if (qOffsetDesc != nullptr) { + ge::DataType qOffsetDType = qOffsetDesc->GetDataType(); + OP_CHECK_IF(QUANT_OFFSET_SUPPORT_DTYPE.find(qOffsetDType) == QUANT_OFFSET_SUPPORT_DTYPE.end(), + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "quant_offset", + ge::TypeUtils::DataTypeToSerialString(qOffsetDType).c_str(), "float32"), + return ge::GRAPH_FAILED); + OP_CHECK_IF(quantMode_ != 0, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "quant_offset", + "not None", "quant_offset only be supported when static quant, but current quant mode is dynamic quant, quant_offset should be None."), + return ge::GRAPH_FAILED); + auto qOffsetStorageShape = context_->GetOptionalInputShape(QUANT_OFFSET_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, qOffsetStorageShape); + auto& qOffsetShape = EnsureNotScalar(qOffsetStorageShape->GetStorageShape()); + const size_t qOffsetDimNum = qOffsetShape.GetDimNum(); + OP_CHECK_IF(qOffsetDimNum > 2, + OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), "quant_offset", + std::to_string(qOffsetDimNum).c_str(), "less than or equal to 2"), + return ge::GRAPH_FAILED); + auto yStorageShape = context_->GetOutputShape(Y_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, yStorageShape); + auto& yShape = EnsureNotScalar(yStorageShape->GetStorageShape()); + if (hasGroupIndex_) { + OP_CHECK_IF(qOffsetShape.GetDim(0) != (groupIndexShape_.GetDim(0)), + OP_LOGE(context_->GetNodeName(), + "quant_offset shape[0] must be equal to group_index shape[0] when group_index exists, please check."), + return ge::GRAPH_FAILED); + if (qOffsetDimNum == DIM_TWO) { + OP_CHECK_IF(qOffsetShape.GetDim(qOffsetDimNum - 1) != 1 && qOffsetShape.GetDim(qOffsetDimNum - 1) != yShape.GetDim(xDimNum_ - 1), + OP_LOGE(context_->GetNodeName(), + "quant_offset shape[-1] must be equal to or can be broadcast to y shape[-1] when group_index exists, please check."), + return ge::GRAPH_FAILED); + } + } else { + if (qOffsetDimNum == DIM_TWO) { + OP_CHECK_IF(qOffsetShape.GetDim(0) != 1, + OP_LOGE(context_->GetNodeName(), + "if dim of quant_offset is 2, shape[0] must be [1] when group_index not exists, please check."), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(qOffsetShape.GetDim(0) != 1 && qOffsetShape.GetDim(0) != yShape.GetDim(xDimNum_ - 1), + OP_LOGE(context_->GetNodeName(), + "quant_offset shape[0] must be or can be broadcast to y shape[-1] when static_quant and group_index not exists, please check."), + return ge::GRAPH_FAILED); + } + hasQuantOffset_ = true; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::CheckForStaticQuant() // 静态量化quant_scale和quant_offset的shape size需要一致 +{ + if (!hasQuantScale_ || !hasQuantOffset_) { + return ge::GRAPH_SUCCESS; + } + auto qScaleStorageShape = context_->GetOptionalInputShape(QUANT_SCALE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, qScaleStorageShape); + int64_t qScaleSize = qScaleStorageShape->GetStorageShape().GetShapeSize(); + auto qOffsetStorageShape = context_->GetOptionalInputShape(QUANT_OFFSET_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, qOffsetStorageShape); + int64_t qOffsetSize = qOffsetStorageShape->GetStorageShape().GetShapeSize(); + OP_CHECK_IF(qScaleSize != qOffsetSize, + OP_LOGE(context_->GetNodeName(), "quant_scale size should be equal to quant_offset size, please check."), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::CheckOutputScale() +{ + auto scaleDesc = context_->GetOutputDesc(SCALE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, scaleDesc); + ge::DataType scaleDType = scaleDesc->GetDataType(); + OP_CHECK_IF(scaleDType != ge::DT_FLOAT, + OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "scale", + ge::TypeUtils::DataTypeToSerialString(scaleDType).c_str(), "float32"), + return ge::GRAPH_FAILED); + auto scaleStorageShape = context_->GetOutputShape(SCALE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, scaleStorageShape); + auto& scaleShape = EnsureNotScalar(scaleStorageShape->GetStorageShape()); + const size_t scaleDimNum = scaleShape.GetDimNum(); + + auto yStorageShape = context_->GetOutputShape(Y_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, yStorageShape); + auto& yShape = EnsureNotScalar(yStorageShape->GetStorageShape()); + const size_t yDimNum = yShape.GetDimNum(); + + OP_CHECK_IF(scaleDimNum != (yDimNum - 1), + OP_LOGE(context_->GetNodeName(), + "scale dimension should be only 1 less than y dimension, please check."), + return ge::GRAPH_FAILED); + + for (size_t i = 0; i < scaleDimNum; i++) { + OP_CHECK_IF(scaleShape[i] != yShape[i], + OP_LOGE(context_->GetNodeName(), + "scale shape[%zu] must be equal to y shape[%zu], please check.", i, i), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::GetAttr() +{ + auto* attrs = context_->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context_, attrs); + + auto* attrActivateLeft = attrs->GetAttrPointer(ATTR_ACTIVATE_LEFT_INDEX); + actRight_ = (attrActivateLeft == nullptr || *attrActivateLeft == false) ? 1 : 0; + const char* attrQuantMode = attrs->GetAttrPointer(ATTR_QUANT_MODE_INDEX); + std::string quantMode = attrQuantMode == nullptr ? "static" : attrQuantMode; + auto it = SUPPORT_QUANT_MODE.find(quantMode); + OP_CHECK_IF(it == SUPPORT_QUANT_MODE.end(), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + context_->GetNodeName(), "quant_mode", + quantMode.c_str(), "quant_mode only support [dynamic] or [static]"), + return ge::GRAPH_FAILED); + quantMode_ = it->second; + // 校验dst_type + auto* attrDstType = attrs->GetAttrPointer(ATTR_DST_TYPE_INDEX); + // 类型校验,防止空指针 + dstType_ = (attrDstType != nullptr) ? *attrDstType : 2; // 默认是2，也即对应输出类型为int8 + OP_CHECK_IF(dstType_ != 2 && dstType_ != 34 && dstType_ != 35 && dstType_ != 36 && dstType_ != 40 && dstType_ != 41, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + context_->GetNodeName(), "dst_type", + std::to_string(dstType_).c_str(), "dst_type only support [2, 34, 35, 36, 40, 41]"), + return ge::GRAPH_FAILED); + // 校验round_mode + const char* attrRoundMode = attrs->GetAttrPointer(ATTR_ROUND_MODE_INDEX); + std::string roundMode = attrRoundMode == nullptr ? "rint" : attrRoundMode; + auto roundModeIt = SUPPORT_ROUND_MODE.find(roundMode); + OP_CHECK_IF(roundModeIt == SUPPORT_ROUND_MODE.end(), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + context_->GetNodeName(), "round_mode", + roundMode.c_str(), "round_mode only support [rint, round, floor, ceil, trunc]"), + return ge::GRAPH_FAILED); + roundMode_ = roundModeIt->second; + // y:[int8, float8]，仅支持rint，y:[float4]，五种类型都支持 + auto yDesc = context_->GetOutputDesc(Y_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, yDesc); + ge::DataType yDType = yDesc->GetDataType(); + // 校验y属于int8和float8时，roundMode是不是rint + if (yDType != ge::DT_HIFLOAT8) { + OP_CHECK_IF((yDType == ge::DT_INT8 || yDType == ge::DT_FLOAT8_E5M2 || yDType == ge::DT_FLOAT8_E4M3FN) && roundMode_ != 0, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + context_->GetNodeName(), "round_mode", + roundMode.c_str(), "round_mode only support [rint] when the type of y in [int8, float8]"), + return ge::GRAPH_FAILED); + } else { + // 校验y属于hifloat8时，roundMode是不是round + OP_CHECK_IF(roundMode_ != 1, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + context_->GetNodeName(), "round_mode", + roundMode.c_str(), "round_mode only support [round] when the type of y in [hifloat8]"), + return ge::GRAPH_FAILED); + } + auto* attrSwigluMode = attrs->GetAttrPointer(ATTR_SWIGLU_MODE_INDEX); + swigluMode_ = (attrSwigluMode == nullptr) ? 0 : *attrSwigluMode; + OP_CHECK_IF(swigluMode_ != 0 && swigluMode_ != 1, + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + context_->GetNodeName(), "swigluMode", + std::to_string(swigluMode_).c_str(), "swigluMode only support [0, 1]"), + return ge::GRAPH_FAILED); + auto* attrClampLimit = attrs->GetAttrPointer(ATTR_CLAMP_LIMIT_INDEX); + clampLimit_ = (attrClampLimit == nullptr) ? CLAMP_LIMIT_DEFAULT : *attrClampLimit; + OP_CHECK_IF(!(clampLimit_ >= 0.0), + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + context_->GetNodeName(), "clamp_limit", + std::to_string(clampLimit_).c_str(), "clamp_limit should be non-negative"), + return ge::GRAPH_FAILED); + auto* attrGluAlpha = attrs->GetAttrPointer(ATTR_GLU_ALPHA_INDEX); + gluAlpha_ = (attrGluAlpha == nullptr) ? GLU_ALPHA_DEFAULT : *attrGluAlpha; + auto* attrGluBias = attrs->GetAttrPointer(ATTR_GLU_BIAS_INDEX); + gluBias_ = (attrGluBias == nullptr) ? GLU_BIAS_DEFAULT : *attrGluBias; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::GetShapeAttrsInfo() { + OP_CHECK_IF(context_ == nullptr, OP_LOGE("DequantSwigluQuant", "context is null."), return ge::GRAPH_FAILED); + OP_CHECK_IF(GetInputX() != ge::GRAPH_SUCCESS, OP_LOGE(context_->GetNodeName(), "get input x failed."), + return ge::GRAPH_FAILED); + OP_CHECK_IF(GetInputGroupIndex() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "get input group_index failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(GetAttrActivateDim() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "get attr activate_dim failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(GetAttr() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "get attr failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckOutputY() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "check output y failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckInputWeightScale() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "check input weight_scale failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckInputActScale() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "check input activation_scale failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckInputBias() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "check input bias failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckInputQuantScale() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "check input quant_scale failed."), return ge::GRAPH_FAILED); + OP_CHECK_IF(CheckInputQuantOffset() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "check input quant_offset failed."), return ge::GRAPH_FAILED); + if (quantMode_ == 0) { + OP_CHECK_IF(CheckForStaticQuant() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "check input quant_scale and quant_offset size failed."), + return ge::GRAPH_FAILED); + } + OP_CHECK_IF(CheckOutputScale() != ge::GRAPH_SUCCESS, + OP_LOGE(context_->GetNodeName(), "check output scale failed."), return ge::GRAPH_FAILED); + + int64_t xTotalNum = xShape_.GetShapeSize(); + inDimy_ = xShape_.GetDim(xDimNum_ - 1); + inDimx_ = xTotalNum / inDimy_; + if ((speGroupType_ == 1) && (groupNum_ >= SPECIAL_GROUP_NUM_64) && + (inDimx_ / groupNum_ <= SPECIAL_GROUP_NUM_16) && + (activateDim_ == static_cast(xDimNum_-1))) { + isSpecialCoreCut_ = 1; + } + if ((speGroupType_ == 0) && (groupNum_ >= SPECIAL_GROUP_NUM_32) && + (inDimx_ / groupNum_ <= SPECIAL_GROUP_NUM_16) && + (activateDim_ == static_cast(xDimNum_-1)) && (!hasBias_) && (!hasQuantScale_)) { + isSpecialCoreCut_ = 1; + } + auto shapeY = context_->GetOutputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context_, shapeY); + const gert::Shape& outputShapeY = shapeY->GetStorageShape(); + outDimy_ = outputShapeY.GetDim(xDimNum_ - 1); // 输出y的-1轴对应的shape + return ge::GRAPH_SUCCESS; +} + +bool DequantSwigluQuantV35DskTiling::IsCapable() { + if (static_cast(activateDim_) != xDimNum_ - static_cast(1)) { + OP_LOGI(context_->GetNodeName(), "transform tiling template 2!"); + return false; + } + return true; +} + +void DequantSwigluQuantV35DskTiling::CalcTilingKeyForNotFull() { + // tilingkey含义：占位\quantMode_\bias\activate_scale\quant_scale\quant_offset\group + if (quantMode_ == 0) { + tilingKey_ = PLACEHOLDER + quantMode_ * QUANT_MODE_FACTOR + hasBias_ * BIAS_FACTOR_FOR_NOT_FULL + + hasActivationScale_ * ACTIVATE_FACTOR_FOR_NOT_FULL + hasQuantScale_ * QUANT_SCALE_FACTOR_FOR_NOT_FULL + + hasQuantOffset_ * QUANT_OFFSET_FACTOR_FOR_NOT_FULL + hasGroupIndex_; + } else { + auto biasDesc = context_->GetOptionalInputDesc(BIAS_INDEX); + int8_t biasDtypeValue = 0; + int8_t value_int32 = 1; + int8_t value_float = 2; + int8_t vlaue_float16 = 3; + int8_t value_bf16 = 4; + if (biasDesc != nullptr) { + ge::DataType biasDtype = biasDesc->GetDataType(); + if (biasDesc != nullptr) { + if (biasDtype == ge::DT_INT32) { + biasDtypeValue = value_int32; + } + if (biasDtype == ge::DT_BF16) { + biasDtypeValue = value_bf16; + } + if (biasDtype == ge::DT_FLOAT16) { + biasDtypeValue = vlaue_float16; + } + if (biasDtype == ge::DT_FLOAT) { + biasDtypeValue = value_float; + } + } + } + // tilingkey含义：占位\quantMode_\bias\activate_scale\quant_scale\quant\offset + tilingKey_ = PLACEHOLDER + quantMode_ * QUANT_MODE_FACTOR + biasDtypeValue * BIAS_FACTOR_FOR_NOT_FULL + + hasActivationScale_ * ACTIVATE_FACTOR_FOR_NOT_FULL + hasQuantScale_ * QUANT_SCALE_FACTOR_FOR_NOT_FULL + + hasQuantOffset_ * QUANT_OFFSET_FACTOR_FOR_NOT_FULL + hasGroupIndex_; + } +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::DoOpTilingNotFull() { + /* + 参数 行尾轴大小 类型占用(bit) 可选/必选 + x 2H 32/16 必选 + weight_scale 2H 32 可选(x为int32——必须，x为16——必须不存在) + activation_scale 1 32 可选(x为int32——可选，x为16——必须不存在) + bias 2H 32/16 可选(x为int32——可选，x为16——必须不存在) + quant_scale 1/H 32 可选(静态量化必选) + quant_offset 1/H 32 可选 + group_index 1 64/32 可选(控制参数不占用UB) + y H 8/4 必选 + scale 1 32 必选(静态量化不计算) + tmp_buffer + */ + int64_t ubFactorDimy = 1; + int64_t ubAvailable = ubSize_ - UB_REVERSE - 1 * BLOCK_SIZE; // 预留1个block给activation_scale（无论activation_scale输入是否存在） + auto yDesc = context_->GetOutputDesc(Y_INDEX); + ge::DataType yDtype = yDesc->GetDataType(); + if (yDtype == ge::DT_FLOAT4_E2M1 or yDtype == ge::DT_FLOAT4_E1M2) { + int64_t ubSplitNum = 1; // 以4bit类型的输出y为一份 + auto xDesc = context_->GetInputDesc(X_INDEX); + ge::DataType xDtype = xDesc->GetDataType(); + ubSplitNum = xDtype == ge::DT_INT32 ? ubSplitNum + 16 : ubSplitNum + 8; // x尾轴2H，32bit类型占16倍空间，否则占8倍 + ubSplitNum = hasWeightScale_ ? ubSplitNum + 16 : ubSplitNum; // weight_scale尾轴2H，输入存在32bit类型占16倍空间，否则不额外占空间 + ubSplitNum = ubSplitNum + 8; // 静态量化场景quant_scale为必选输入，尾轴可能为1或H，按照H预留8倍空间 + ubSplitNum = hasQuantOffset_ ? ubSplitNum + 8 : ubSplitNum; // quant_offset尾轴可能为1或H，输入存在32bit类型占8倍空间，否则不占空间 + if (hasBias_) { + auto biasDesc = context_->GetOptionalInputDesc(BIAS_INDEX); + ge::DataType biasDtype = biasDesc->GetDataType(); + ubSplitNum = (biasDtype == ge::DT_INT32 or biasDtype == ge::DT_FLOAT) ? ubSplitNum + 16 : ubSplitNum + 8; // bias尾轴2H，32bit类型占16倍空间，否则占8倍 + } + ubSplitNum += 1; // y为4bit类型时，ub分配内存实际按照1B/num分配，多加一倍空间 + int64_t doubleBuffer = 2; + ubSplitNum *= doubleBuffer; // x、weight_scale、bias、y doublebuffer + ubSplitNum += 16; // tmp buffer尾轴2H，固定为32bit，占16倍空间 + + int64_t ySize = (ubAvailable / ubSplitNum) / BLOCK_SIZE * BLOCK_SIZE; // 一份4bit类型y占用空间，此处32B对齐保证y的空间32B对齐，偶数 + + ubFactorDimy = ySize / 0.5; // 非全载模板，单核内一次循环处理4bit类型y的元素个数（单个元素0.5Byte） + } else { + int64_t ubSplitNum = 1; // 以uint8类型的输出y为一份 + auto xDesc = context_->GetInputDesc(X_INDEX); + ge::DataType xDtype = xDesc->GetDataType(); + ubSplitNum = xDtype == ge::DT_INT32 ? ubSplitNum + 8 : ubSplitNum + 4; // x尾轴2H，32bit类型占8倍空间，否则占4倍 + ubSplitNum = hasWeightScale_ ? ubSplitNum + 8 : ubSplitNum; // weight_scale尾轴2H，输入存在32bit类型占8倍空间，否则不额外占空间 + ubSplitNum = ubSplitNum + 4; // 静态量化场景quant_scale为必选输入，尾轴可能为1或H，按照H预留4倍空间 + ubSplitNum = hasQuantOffset_ ? ubSplitNum + 4 : ubSplitNum; // quant_offset尾轴可能为1或H，输入存在32bit类型占4倍空间，否则不占空间 + if (hasBias_) { + auto biasDesc = context_->GetOptionalInputDesc(BIAS_INDEX); + ge::DataType biasDtype = biasDesc->GetDataType(); + ubSplitNum = (biasDtype == ge::DT_INT32 or biasDtype == ge::DT_FLOAT) ? ubSplitNum + 8 : ubSplitNum + 4; // bias尾轴2H，32bit类型占8倍空间，否则占4倍 + } + int64_t doubleBuffer = 2; + ubSplitNum *= doubleBuffer; // x、weight_scale、bias、y doublebuffer + ubSplitNum += 8; // tmp buffer尾轴2H，固定为32bit，占8倍空间 + + int64_t ySize = (ubAvailable / ubSplitNum) / BLOCK_SIZE * BLOCK_SIZE; // 一份8bit类型y占用空间，此处32B对齐保证y的空间32B对齐，偶数 + + ubFactorDimy = ySize; // 非全载模板，单核内一次循环处理8bit类型y的元素个数（单个元素1Byte） + } + // 当实际输出尾轴长度小于搬运长度时的适配 + if (ubFactorDimy > outDimy_) { + //获取一个block内y的元素个数，当y为4bit类型时(0.5Byte)，一个block内元素个数等于BLOCK_SIZE * 2, 否则y为8bit类型，一个block内元素个数等于BLOCK_SIZE + int64_t numPerBlock = (yDtype == ge::DT_FLOAT4_E2M1 or yDtype == ge::DT_FLOAT4_E1M2) ? BLOCK_SIZE * 2 : BLOCK_SIZE; + ubFactorDimy = Ops::Base::CeilDiv(outDimy_, numPerBlock) * numPerBlock; + } + + int64_t loopTimesPerRow = (outDimy_ + ubFactorDimy - 1) / ubFactorDimy; // 非全载模板，单核处理一行需要的循环次数 + int64_t tailPerRow = outDimy_ - (loopTimesPerRow - 1) * ubFactorDimy; // 非全载模板，单核处理一行的尾块大小 + + maxPreCore_ = std::min(static_cast(coreNum_), inDimx_); // 非全载模板，按照行数分核，单核内将一行分段处理 + + CalcTilingKeyForNotFull(); + tilingData_.set_inDimx(inDimx_); + tilingData_.set_inDimy(inDimy_); + tilingData_.set_outDimy(outDimy_); + tilingData_.set_UbFactorDimx(1); // 非全载模板固定1 + tilingData_.set_UbFactorDimy(ubFactorDimy); // 非全载模板ub内一次循环处理的尾轴元素个数 + tilingData_.set_usedCoreNum(maxPreCore_); + tilingData_.set_maxCoreNum(coreNum_); + tilingData_.set_inGroupNum(groupNum_); + tilingData_.set_quantMode(quantMode_); // quantMode,0:静态量化,1:动态量化 + tilingData_.set_speGroupType(speGroupType_); + tilingData_.set_isSpecialCoreCut(isSpecialCoreCut_); + tilingData_.set_actRight(actRight_); + tilingData_.set_dstType(dstType_); + tilingData_.set_roundMode(roundMode_); + tilingData_.set_activateDim(activateDim_); + tilingData_.set_loopTimesPerRow(loopTimesPerRow); // 非全载模板处理一行需要的UB循环次数 + tilingData_.set_tailPerRow(tailPerRow); // 非全载模板UB循环最后一次的元素个数 + tilingData_.set_swiGluMode(swigluMode_); // swiGluMode + tilingData_.set_biasMode(biasMode_); // bias类型，0：不存在；1：int32；2：bf16；3：fp16；4：fp32 + tilingData_.set_groupIndexMode(groupIndexMode_); // group_index类型，0：不存在；1：int32；2：int64 + tilingData_.set_quantIsOne(quantIsOne_); + tilingData_.set_clampLimit(clampLimit_); + tilingData_.set_gluAlpha(gluAlpha_); + tilingData_.set_gluBias(gluBias_); + + OP_LOGI(context_->GetNodeName(), "inDimx is %ld, inDimy is %ld, outDimy is %ld, UbFactorDimx is %ld, UbFactorDimy is %ld, usedCoreNum is %ld, maxCoreNum is %ld, \ + inGroupNum is %ld, quantMode is %ld, actRight is %ld, dstType is %ld, roundMode is %ld, activateDim is %ld, loopTimesPerRow is %ld, \ + tailPerRow is %ld, swiGluMode is %ld, biasMode is %ld, groupIndexMode is %ld, quantIsOne is %ld, clampLimit is %f, gluAlpha is %f, gluBias is %f", \ + tilingData_.get_inDimx(), tilingData_.get_inDimy(), tilingData_.get_outDimy(), tilingData_.get_UbFactorDimx(), tilingData_.get_UbFactorDimy(), \ + tilingData_.get_usedCoreNum(), tilingData_.get_maxCoreNum(), tilingData_.get_inGroupNum(), tilingData_.get_quantMode(), tilingData_.get_actRight(), tilingData_.get_dstType(), \ + tilingData_.get_roundMode(), tilingData_.get_activateDim(), tilingData_.get_loopTimesPerRow(), tilingData_.get_tailPerRow(), tilingData_.get_swiGluMode(), \ + tilingData_.get_biasMode(), tilingData_.get_groupIndexMode(), tilingData_.get_quantIsOne(), tilingData_.get_clampLimit(), tilingData_.get_gluAlpha(), tilingData_.get_gluBias()); + OP_LOGI(context_->GetNodeName(), "tilingKey_ is %ld, speGroupType is %ld, isSpecialCoreCut is %ld", tilingKey_, speGroupType_, isSpecialCoreCut_); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::DoOpTiling() { + if (outDimy_ > Y_LAST_DIM_FULL_LOAD_MAX_VALUE) { + return DequantSwigluQuantV35DskTiling::DoOpTilingNotFull(); + } + auto xDesc = context_->GetInputDesc(X_INDEX); + ge::DataType xDtype = xDesc->GetDataType(); + size_t xBits = xDtype == ge::DT_INT32 ? sizeof(int32_t) : sizeof(int16_t); + int64_t xUbAlign32B = ((outDimy_ + BLOCK_ELEM_B32 - 1) / BLOCK_ELEM_B32) * BLOCK_ELEM_B32; + int64_t xUbAlign = xDtype == ge::DT_INT32 ? xUbAlign32B : + ((outDimy_ + BLOCK_ELEM_B16 - 1) / BLOCK_ELEM_B16) * BLOCK_ELEM_B16; + int64_t aScaleAlign32B_ = BLOCK_ELEM_B32; + int64_t yAlign8B = ((outDimy_ + BLOCK_ELEM_B8 - 1) / BLOCK_ELEM_B8) * BLOCK_ELEM_B8; + int64_t doubleBuffer = 2; + int64_t ubAvailable = ubSize_ - UB_REVERSE - (xUbAlign32B * SWI_FACTOR + xUbAlign32B) * sizeof(float); + int64_t denominator = doubleBuffer * (xUbAlign * SWI_FACTOR * xBits + aScaleAlign32B_ * sizeof(float)) + + doubleBuffer * yAlign8B* sizeof(int8_t) + aScaleAlign32B_ * sizeof(float) + + xUbAlign32B * sizeof(float); + + // swiglu_mode=1时，增加x weight_scale的尾轴128向上对齐ub + if (swigluMode_ == 1) { + int64_t tailSupply = (inDimy_ + 128 - 1) / 128 * 128 - inDimy_; + denominator += tailSupply * xBits + tailSupply * sizeof(int32_t); + } + + // 判断bias, bias=nullptr：biasDtypeValue = 0, 否则，biasDtypeValue = 1 + auto biasDesc = context_->GetOptionalInputDesc(BIAS_INDEX); + int64_t biasDtypeValue = 0; + int64_t value_int32 = 2; + int64_t value_bf16 = 4; + int64_t vlaue_float16 = 3; + int64_t value_float = 2; + // 判断bias是否合法存在，如果存在的话，则需要考虑bias在ub里面占用的内存 + if (biasDesc != nullptr) { + ge::DataType biasDtype = biasDesc->GetDataType(); + int64_t biasMemory = xUbAlign32B * sizeof(int16_t); // base为4B类型 + // swigluV2场景下，增加bias的尾块128B对齐，防止越界 + int64_t biasTailSupply = (inDimy_ + 128 - 1) / 128 * 128 - inDimy_; + if (biasDesc != nullptr) { + if (biasDtype == ge::DT_INT32) { + denominator += biasMemory * value_int32; + biasDtypeValue = 1; + if (swigluMode_ == 1) { + denominator += biasTailSupply * sizeof(int32_t); + } + } + if (biasDtype == ge::DT_BF16) { + denominator += biasMemory; + biasDtypeValue = value_bf16; + if (swigluMode_ == 1) { + denominator += biasTailSupply * sizeof(int16_t); + } + } + if (biasDtype == ge::DT_FLOAT16) { + denominator += biasMemory; + biasDtypeValue = vlaue_float16; + if (swigluMode_ == 1) { + denominator += biasTailSupply * sizeof(int16_t); + } + } + if (biasDtype == ge::DT_FLOAT) { + denominator += biasMemory * SWI_FACTOR; + biasDtypeValue = value_float; + if (swigluMode_ == 1) { + denominator += biasTailSupply * sizeof(int32_t); + } + } + } + } + + // 静态量化下，增加quant_offset的UB + if (hasQuantOffset_) { + denominator += xUbAlign32B * sizeof(float); + } + + // ubFactorDimX: ub最多可以处理多少行数据 + int64_t ubFactorDimx = ubAvailable / denominator; + ubFactorDimx = std::min(ubFactorDimx, inDimx_); + OP_CHECK_IF(ubFactorDimx < 1, + OP_LOGE(context_->GetNodeName(), "x last dim:%ld is too large to full load", inDimy_), + return ge::GRAPH_FAILED); + maxPreCore_ = std::min(maxPreCore_, (inDimx_ + ubFactorDimx - 1) / ubFactorDimx); + OP_LOGI(context_->GetNodeName(), "start maxPreCore_ is %ld ", maxPreCore_); + if (isSpecialCoreCut_ == static_cast(1)) { + maxPreCore_ = std::min(static_cast(coreNum_), static_cast(inDimx_)); + OP_LOGI(context_->GetNodeName(), "after maxPreCore_ is %ld ", maxPreCore_); + } + + auto quantScaleDesc = context_->GetOptionalInputDesc(QUANT_SCALE_INDEX); + auto actScaleDesc = context_->GetOptionalInputDesc(ACTIVATION_SCALE_INDEX); + auto groupIndexDesc = context_->GetOptionalInputDesc(INPUT_GROUP_INDEX); + + // 输入x的tiling key计算位 + int64_t hasXInt = quantMode_ == 1 ? 0: 1; + int64_t hasAScale = actScaleDesc != nullptr; + // hasQScale=0:无quant_scale，hasQScale=1:float + int64_t hasQScale = quantScaleDesc != nullptr; + int64_t hasGIndex = groupIndexDesc != nullptr; + // activateDim=-1时，hasActivateDim=0，否则为1; 当activateDim=xDim-1时，hashActivateDim=1 + int64_t hasActivateDim = static_cast(activateDim_) != xDimNum_ - static_cast(1); + + // 增加十万分位的tiling_key，hasActivateDim;增加万分位的tilling_key hasXInt:判断输入x是不是int32;千分位的tiling_key biasDtypeValue + tilingKey_ = hasActivateDim * ACTIVATE_DIM_FACTOR + hasXInt * INPUT_X_FACTOR + biasDtypeValue * BIAS_FACTOR + hasAScale * ACT_SCALE_FACTOR + hasQScale * QUANT_SCALE_FACTOR + hasGIndex * GROUP_INDEX_FACTOR; + tilingData_.set_inDimx(inDimx_); + tilingData_.set_inDimy(inDimy_); + tilingData_.set_outDimy(outDimy_); + tilingData_.set_UbFactorDimx(ubFactorDimx); + tilingData_.set_UbFactorDimy(outDimy_); + tilingData_.set_usedCoreNum(maxPreCore_); + tilingData_.set_maxCoreNum(maxPreCore_); + tilingData_.set_inGroupNum(groupNum_); + tilingData_.set_quantMode(quantMode_); + tilingData_.set_actRight(actRight_); + tilingData_.set_dstType(dstType_); + tilingData_.set_roundMode(roundMode_); + tilingData_.set_activateDim(activateDim_); + tilingData_.set_swiGluMode(swigluMode_); // swiGluMode + tilingData_.set_biasMode(biasMode_); // bias类型，0：不存在；1：int32；2：bf16；3：fp16；4：fp32 + tilingData_.set_groupIndexMode(groupIndexMode_); // group_index类型，0：不存在；1：int32；2：int64 + tilingData_.set_quantIsOne(quantIsOne_); + tilingData_.set_clampLimit(clampLimit_); + tilingData_.set_gluAlpha(gluAlpha_); + tilingData_.set_gluBias(gluBias_); + tilingData_.set_speGroupType(speGroupType_); + tilingData_.set_isSpecialCoreCut(isSpecialCoreCut_); + OP_LOGI(context_->GetNodeName(), "inDimx is %ld, inDimy is %ld, outDimy is %ld, UbFactorDimx is %ld, UbFactorDimy is %ld, usedCoreNum is %ld, maxCoreNum is %ld, \ + inGroupNum is %ld, quantMode is %ld, actRight is %ld, dstType is %ld, roundMode is %ld, activateDim is %ld, swiGluMode is %ld, \ + biasMode is %ld, groupIndexMode is %ld, biasMode is %ld, groupIndexMode is %ld, quantIsOne is %ld, clampLimit is %f, gluAlpha is %f, gluBias is %f", \ + tilingData_.get_inDimx(), tilingData_.get_inDimy(), tilingData_.get_outDimy(), tilingData_.get_UbFactorDimx(), tilingData_.get_UbFactorDimy(), \ + tilingData_.get_usedCoreNum(), tilingData_.get_maxCoreNum(), tilingData_.get_inGroupNum(), tilingData_.get_quantMode(), tilingData_.get_actRight(), tilingData_.get_dstType(), \ + tilingData_.get_roundMode(), tilingData_.get_activateDim(), tilingData_.get_swiGluMode(), tilingData_.get_biasMode(), tilingData_.get_groupIndexMode(), \ + tilingData_.get_biasMode(), tilingData_.get_groupIndexMode(), tilingData_.get_quantIsOne(), tilingData_.get_clampLimit(), tilingData_.get_gluAlpha(), tilingData_.get_gluBias()); + OP_LOGI(context_->GetNodeName(), "tilingKey_ is %ld, speGroupType is %ld, isSpecialCoreCut is %ld", tilingKey_, speGroupType_, isSpecialCoreCut_); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::DoLibApiTiling() { + return ge::GRAPH_SUCCESS; +} + +uint64_t DequantSwigluQuantV35DskTiling::GetTilingKey() const { + return tilingKey_; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::GetWorkspaceSize() { + // 如果是动态非全载，需要使用workspace存Dequant Swiglu的计算结果 + size_t usrSize = 0; + if (outDimy_ > Y_LAST_DIM_FULL_LOAD_MAX_VALUE && quantMode_ == 1) { + usrSize = maxPreCore_ * outDimy_ * sizeof(float); + } + OP_LOGI(context_->GetNodeName(), "usrSize is %u", usrSize); + workspaceSize_ = SYS_WORK_SPACE_SIZE + usrSize; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35DskTiling::PostTiling() { + context_->SetTilingKey(GetTilingKey()); + context_->SetBlockDim(maxPreCore_); + size_t* workspaces = context_->GetWorkspaceSizes(1); + OP_CHECK_NULL_WITH_CONTEXT(context_, workspaces); + workspaces[0] = workspaceSize_; + OP_LOGI(context_->GetNodeName(), "workspace is %lu, SetBlockDim is %ld", workspaceSize_, maxPreCore_); + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35NlastTiling::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + if (platformInfo == nullptr) { + auto compileInfoPtr = static_cast(context_->GetCompileInfo()); + OP_CHECK_IF(compileInfoPtr == nullptr, OP_LOGE(context_, "compile info is null"), + return ge::GRAPH_FAILED); + coreNum_ = compileInfoPtr->coreNum; + ubSize_ = compileInfoPtr->ubSize; + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + coreNum_ = ascendcPlatform.GetCoreNumAiv(); + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + ubSize_ = ubSizePlatForm; + socVersion = ascendcPlatform.GetSocVersion(); + } + + return ge::GRAPH_SUCCESS; +} + +void DequantSwigluQuantV35NlastTiling::FusedShape() +{ + inDim0_ = 1; + inDim1_ = 1; + inDim2_ = 1; + for (size_t i = 0; i < xShape_.GetDimNum(); i++) { + if (i < static_cast(actDimIndex_)) { + inDim0_ *= xShape_.GetDim(i); + } else if (i == xShape_.GetDimNum() - 1) { + inDim2_ *= xShape_.GetDim(i); + } else { + inDim1_ *= xShape_.GetDim(i); + } + } + return; +} + +ge::graphStatus DequantSwigluQuantV35NlastTiling::GetShapeAttrsInfo() +{ + auto xStorageShape = context_->GetInputShape(X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, xStorageShape); + xShape_ = EnsureNotScalar(xStorageShape->GetStorageShape()); + + auto* attrs = context_->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context_, attrs); + auto* attrActivateLeft = attrs->GetAttrPointer(ATTR_ACTIVATE_LEFT_INDEX); + actRight_ = (attrActivateLeft == nullptr || *attrActivateLeft == false) ? 1 : 0; + + const char* attrRoundMode = attrs->GetAttrPointer(ATTR_ROUND_MODE_INDEX); + std::string roundMode = attrRoundMode == nullptr ? "rint" : attrRoundMode; + // has checked + roundMode_ = SUPPORT_ROUND_MODE.find(roundMode)->second; + + auto* attrActivateDim = attrs->GetAttrPointer(ATTR_ACTIVATE_DIM_INDEX); + actDimIndex_ = (attrActivateDim != nullptr) ? *attrActivateDim : -1; + actDimIndex_ = actDimIndex_ < 0 ? actDimIndex_ + static_cast(xShape_.GetDimNum()) : actDimIndex_; + + FusedShape(); + + outDim1_ = inDim1_ / SWI_FACTOR; + + return ge::GRAPH_SUCCESS; +} + +bool DequantSwigluQuantV35NlastTiling::IsCapable() +{ + if (static_cast(actDimIndex_) == xShape_.GetDimNum() - 1) { + return false; + } + + return true; +} + +void DequantSwigluQuantV35NlastTiling::DoBlockSplit() +{ + int64_t maxCoreNum = static_cast(coreNum_); + blockFormer0_ = (inDim0_ + maxCoreNum - 1) / maxCoreNum; + blockNum0_ = (inDim0_ + blockFormer0_ - 1) / blockFormer0_; + int64_t needBlockNum1 = maxCoreNum / blockNum0_; + blockFormer1_ = (outDim1_ + needBlockNum1 - 1) / needBlockNum1; + blockNum1_ = (outDim1_ + blockFormer1_ - 1) / blockFormer1_; +} + +bool DequantSwigluQuantV35NlastTiling::DoUbSplit() +{ + int64_t xUbAlign32B = ((inDim2_ + BLOCK_ELEM_B32 - 1) / BLOCK_ELEM_B32) * BLOCK_ELEM_B32; + int64_t ubAvailable = ubSize_ - UB_REVERSE - (xUbAlign32B + xUbAlign32B) * sizeof(float); + int64_t doubleBuffer = 2; + int64_t aScaleAlign32B = BLOCK_ELEM_B32; + int64_t yAlign8B = ((inDim2_ + BLOCK_ELEM_B8 - 1) / BLOCK_ELEM_B8) * BLOCK_ELEM_B8; + auto xDesc = context_->GetInputDesc(X_INDEX); + size_t xBits = xDesc->GetDataType() == ge::DT_INT32 ? sizeof(int32_t) : sizeof(int16_t); + int64_t xUbAlign = xDesc->GetDataType() == ge::DT_INT32 ? xUbAlign32B : + ((inDim2_ + BLOCK_ELEM_B16 - 1) / BLOCK_ELEM_B16) * BLOCK_ELEM_B16; + int64_t denominator = doubleBuffer * (xUbAlign * SWI_FACTOR * xBits + + aScaleAlign32B * SWI_FACTOR * sizeof(float)) + + doubleBuffer * yAlign8B * sizeof(int8_t) + aScaleAlign32B * sizeof(float) + + xUbAlign32B * sizeof(float); + auto biasDesc = context_->GetOptionalInputDesc(BIAS_INDEX); + biasDtypeValue_ = 0; + int64_t value_int32 = 2; + int64_t value_bf16 = 4; + int64_t vlaue_float16 = 3; + int64_t value_float = 2; + if (biasDesc != nullptr) { + ge::DataType biasDtype = biasDesc->GetDataType(); + int64_t biasMemory = xUbAlign32B * sizeof(int16_t); + if (biasDtype == ge::DT_INT32) { + denominator += biasMemory * value_int32; + biasDtypeValue_ = 1; + } + if (biasDtype == ge::DT_BF16) { + denominator += biasMemory; + biasDtypeValue_ = value_bf16; + } + if (biasDtype == ge::DT_FLOAT16) { + denominator += biasMemory; + biasDtypeValue_ = vlaue_float16; + } + if (biasDtype == ge::DT_FLOAT) { + denominator += biasMemory * SWI_FACTOR; + biasDtypeValue_ = value_float; + } + } + ubFormer1_ = ubAvailable / denominator; + if (ubFormer1_ < 1) { + return false; + } + ubFormer0_ = 1; + if (ubFormer1_ > blockFormer1_) { + ubFormer0_ = ubFormer1_ / blockFormer1_; + ubFormer1_ = blockFormer1_; + } + + ubFormer0_ = std::min(ubFormer0_, blockFormer0_); + + return true; +} + +ge::graphStatus DequantSwigluQuantV35NlastTiling::DoOpTiling() +{ + DoBlockSplit(); + if (!DoUbSplit()) { + OP_LOGE(context_->GetNodeName(), "UB size cannot load last dim of input x, return failed."); + return ge::GRAPH_FAILED; + } + int64_t ubLoopOfFormerBlock0 = (blockFormer0_ + ubFormer0_ - 1) / ubFormer0_; + int64_t blockTail0 = inDim0_ - blockFormer0_ * (blockNum0_ - 1); + int64_t ubLoopOfTailBlock0 = (blockTail0 + ubFormer0_ - 1) / ubFormer0_; + int64_t ubTailOfFormerBlock0 = blockFormer0_ - (ubLoopOfFormerBlock0 - 1) * ubFormer0_; + int64_t ubTailOfTailBlock0 = blockTail0 - (ubLoopOfTailBlock0 - 1) * ubFormer0_; + + int64_t ubLoopOfFormerBlock1 = (blockFormer1_ + ubFormer1_ - 1) / ubFormer1_; + int64_t blockTail1 = outDim1_ - blockFormer1_ * (blockNum1_ - 1); + int64_t ubLoopOfTailBlock1 = (blockTail1 + ubFormer1_ - 1) / ubFormer1_; + int64_t ubTailOfFormerBlock1 = blockFormer1_ - (ubLoopOfFormerBlock1 - 1) * ubFormer1_; + int64_t ubTailOfTailBlock1 = blockTail1 - (ubLoopOfTailBlock1 - 1) * ubFormer1_; + + blockNum_ = blockNum0_ * blockNum1_; + tilingData_.set_inDim0(inDim0_); + tilingData_.set_inDim1(inDim1_); + tilingData_.set_inDim2(inDim2_); + tilingData_.set_outDim1(inDim1_ / SWI_FACTOR); + tilingData_.set_blockNum0(blockNum0_); + tilingData_.set_blockNum1(blockNum1_); + tilingData_.set_blockFormer0(blockFormer0_); + tilingData_.set_blockFormer1(blockFormer1_); + tilingData_.set_ubFormer0(ubFormer0_); + tilingData_.set_ubFormer1(ubFormer1_); + tilingData_.set_ubLoopOfFormerBlock0(ubLoopOfFormerBlock0); + tilingData_.set_ubLoopOfFormerBlock1(ubLoopOfFormerBlock1); + tilingData_.set_ubLoopOfTailBlock0(ubLoopOfTailBlock0); + tilingData_.set_ubLoopOfTailBlock1(ubLoopOfTailBlock1); + tilingData_.set_ubTailOfFormerBlock0(ubTailOfFormerBlock0); + tilingData_.set_ubTailOfFormerBlock1(ubTailOfFormerBlock1); + tilingData_.set_ubTailOfTailBlock0(ubTailOfTailBlock0); + tilingData_.set_ubTailOfTailBlock1(ubTailOfTailBlock1); + tilingData_.set_actRight(actRight_); + tilingData_.set_roundMode(roundMode_); + + auto quantScaleDesc = context_->GetOptionalInputDesc(QUANT_SCALE_INDEX); + auto actScaleDesc = context_->GetOptionalInputDesc(ACTIVATION_SCALE_INDEX); + + int64_t hasAScale = actScaleDesc != nullptr; + int64_t hasQScale = quantScaleDesc != nullptr; + + tilingKey_ = 1 * ACTIVATE_DIM_FACTOR + biasDtypeValue_ * BIAS_FACTOR + hasAScale * ACT_SCALE_FACTOR + + hasQScale * QUANT_SCALE_FACTOR + 0 * GROUP_INDEX_FACTOR; + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35NlastTiling::DoLibApiTiling() { + return ge::GRAPH_SUCCESS; +} + +uint64_t DequantSwigluQuantV35NlastTiling::GetTilingKey() const { + return tilingKey_; +} + +ge::graphStatus DequantSwigluQuantV35NlastTiling::GetWorkspaceSize() { + workspaceSize_ = WORKSPACE_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantV35NlastTiling::PostTiling() { + context_->SetTilingKey(GetTilingKey()); + context_->SetBlockDim(blockNum_); + size_t* workspaces = context_->GetWorkspaceSizes(1); + OP_CHECK_NULL_WITH_CONTEXT(context_, workspaces); + workspaces[0] = workspaceSize_; + OP_LOGI(context_->GetNodeName(), "SetBlockDim is %ld", blockNum_); + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +REGISTER_TILING_TEMPLATE("DequantSwigluQuant", DequantSwigluQuantV35DskTiling, 1000); +REGISTER_TILING_TEMPLATE("DequantSwigluQuant", DequantSwigluQuantV35NlastTiling, 2000); + +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_0, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_10, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_11, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_111, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1111, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_2111, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_3111, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_4111, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_2110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_3110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_4110, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_2101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_3101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_4101, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_2100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_3100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_4100, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1011, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_2011, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_3011, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_4011, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1010, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_2010, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_3010, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_4010, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_2001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_3001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_4001, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_1000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_2000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_3000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_4000, DequantSwigluQuantV35BaseTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100110, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100100, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100010, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_100000, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_101110, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_102110, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_103110, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_104110, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_101100, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_102100, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_103100, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_104100, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_101010, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_102010, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_103010, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_104010, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_101000, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_102000, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_103000, DequantSwigluQuantV35NlastTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant_104000, DequantSwigluQuantV35NlastTilingData) +} // namespace optiling diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_base.cpp b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_base.cpp new file mode 100644 index 000000000..612ff81c2 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_base.cpp @@ -0,0 +1,731 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_tiling_base.cpp + * \brief + */ + +#include +#include +#include "tiling/tiling_api.h" +#include "swi_glu_tiling.h" +#include "../tiling_base/tiling_util.h" +#include "dequant_swiglu_quant_tiling.h" +#include "../tiling_base/tiling_templates_registry.h" + +#define CHECK_FAIL(cont, cond, ...) \ + do { \ + if (cond) { \ + OP_LOGE(cont->GetNodeName(), ##__VA_ARGS__); \ + return ge::GRAPH_FAILED; \ + } \ + } while (0) + +namespace optiling { +constexpr uint32_t UB_RESERVED_BUFF = 0; // reserve 0k +constexpr uint32_t PACK_UINT_IN_CACHE_512B = 512; // pack unit in cache 512B +constexpr uint32_t ALIGN_UINT_IN_CACHE_32B = 32; // align unit in cache 32B +constexpr uint32_t ALIGN_UINT_IN_CACHE_64B = 64; // align unit in cache 64B +constexpr uint32_t ALIGN_TYPE_INT32 = 8; // int32 对齐32字节 +constexpr uint32_t DEFAULT_BUFFER_NUM = 2; +constexpr uint32_t MAX_BLOCK_COUNT = 4095; // datacopy指令包含的连续传输数据块的最大个数 +constexpr uint32_t MAX_BLOCK_LEN = 2097120; // 65535 * 32 datacopy指令每个连续传输数据块的最长长度为65535，单位为32bytes +constexpr uint32_t MAX_UINT32 = 4294967295; +constexpr uint32_t MAX_CORE_NUMBER = 64; +constexpr uint16_t DISCONTINE_COPY_MAX_BLOCKCNT = 4095; // 非连续拷贝，blockCount最大值,AscendC接口限制 +constexpr uint16_t DISCONTINE_COPY_MAX_BLOCKLEN = 65535; // 非连续拷贝，blockLen最大值,AscendC接口限制 +constexpr uint16_t DISCONTINE_COPY_MAX_STRIDE = 65535; // 非连续拷贝，srcStride/dstStride最大值,AscendC接口限制 + +static const uint32_t DYNAMIC_BF16_TBUF_NUM_HALF = 11; +static const uint32_t DYNAMIC_BF16_INT16_TBUF_NUM_HALF = 6; +static const uint32_t STATIC_BF16_TBUF_NUM_HALF = 12; +static const uint32_t STATIC_BF16_INT16_TBUF_NUM_HALF = 7; +static const uint32_t DYNAMIC_INT16_TBUF_NUM_HALF = 2; + +static const size_t INDEX_IN_WEIGHT_SCALE = 1; +static const size_t INDEX_IN_ACTIVATE_SCALE = 2; +static const size_t INDEX_IN_BIAS = 3; +static const size_t INDEX_IN_QUANT_SCALE = 4; +static const size_t INDEX_IN_QUANT_OFFSET = 5; +static const size_t NUMBER_OF_INPUT_SIZE = 10; +static const size_t USER_WORKSPACE = 16777216; // 16 * 1024 * 1024 +constexpr uint32_t PERFORMANCE_COL_LEN = 1536; +constexpr uint32_t PERFORMANCE_ROW_LEN = 128; +constexpr uint32_t MIN_CORE = 12; +const int64_t DYNAMIC_INT_X_FLOAT32_BIAS_QUANT_D_PERFORMANCE = 30013; + +// Tiling优选参数 +struct GluSingleTilingOptParam { + // Maximum amount of data that can be transferred by an operator UB at a time. Unit:element + uint32_t maxTileLen = 0; + uint32_t optBaseRowLen = 0; // 最优的BaseRowLen + uint32_t optBaseColLen = 0; // 最优的BaseColLen + uint64_t optTotalTileNum = 0; // 最优的分割后的数据块数量 + uint64_t optBaseSize = 0; // 最优的分割后的base shape数据块的大小， optBaseRowLen*optBaseColLen, Unit:element + uint64_t optBaseTileNum = 0; // 最优的分割后的base shape数据块数量，不包含尾块 + + uint32_t totalUsedCoreNum = 0; // 最终实际使用的核数 + uint64_t tileNumPerCore = 0; // 每个核需要处理的TileNum，如果不均匀，按照多的计算 +}; + +class DequantSwigluQuantTiling : public TilingBaseClass { +public: + explicit DequantSwigluQuantTiling(gert::TilingContext* cont) : TilingBaseClass(cont) + { + Reset(); + } + ~DequantSwigluQuantTiling() override = default; + + void Reset(gert::TilingContext* cont) override + { + TilingBaseClass::Reset(cont); + Reset(); + } + +protected: + bool IsCapable() override + { + auto shapeGroupIndex = context_->GetOptionalInputShape(6); + if (shapeGroupIndex == nullptr) { + return true; + } + return false; + } + + // 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小 + ge::graphStatus GetPlatformInfo() override; + // 2、获取INPUT/OUTPUT/ATTR信息 + ge::graphStatus GetShapeAttrsInfo() override; + // 3、计算数据切分TilingData + ge::graphStatus DoOpTiling() override; + // 4、计算高阶API的TilingData + ge::graphStatus DoLibApiTiling() override; + // 5、计算TilingKey + uint64_t GetTilingKey() const override; + // 6、计算Workspace 大小 + ge::graphStatus GetWorkspaceSize() override; + // 7、保存Tiling数据 + ge::graphStatus PostTiling() override; + void Reset(); + +private: + void ShowTilingData(); + + ge::graphStatus checkInputShape(gert::TilingContext* context, ge::DataType xDataType); + + ge::graphStatus checkWeightBiasActivate(gert::TilingContext* context); + + ge::graphStatus SetTotalShape(gert::TilingContext* cont, const gert::Shape& inShape); + + bool SetAttr(const gert::RuntimeAttrs* attrs); + + bool CalcTiling(const uint32_t totalCores, const uint64_t ubSize, const platform_ascendc::SocVersion socVersion_); + + bool CalcOptTiling(const uint64_t ubSize, const int32_t dtype, GluSingleTilingOptParam& optTiling); + + bool CalcUbMaxTileLen(uint64_t ubSize, int32_t dtype, GluSingleTilingOptParam& optTiling); + + bool GetBufferNumAndDataLenPerUB(uint64_t ubSize, int32_t dtype, uint64_t& dataLenPerUB); + + bool CalcOptBaseShape(GluSingleTilingOptParam& optTiling, int32_t dtype); + + uint32_t getBaseColLenUpBound(GluSingleTilingOptParam& optTiling); + + void SaveOptBaseShape(uint32_t baseRowLen_, uint32_t baseColLen_, GluSingleTilingOptParam& optTiling); + + int64_t getTilingKeyDynamic( + const int32_t inputDtype, const ge::DataType biasType, const int64_t scaleSize) const; + + bool isPerformanceBranch(); + + int64_t getTilingKeyStatic( + const int32_t inputDtype, const ge::DataType biasType, const int64_t scaleSize) const; + + ge::graphStatus GetShapeAttrsInfoInner(); + + uint32_t inputDTypeLen = 2; + uint32_t activateLeft = 0; // false <-> 0: activate right + int32_t quantMode = 0; + uint32_t maxTileLen = 0; + uint32_t optBaseRowLen = 0; // 最优的BaseRowLen + uint32_t optBaseColLen = 0; // 最优的BaseColLen + uint64_t optTotalTileNum = 0; // 最优的分割后的数据块数量 + uint64_t optBaseSize = 0; // 最优的分割后的base shape数据块的大小， optBaseRowLen*optBaseColLen, Unit:element + uint64_t optBaseTileNum = 0; // 最优的分割后的base shape数据块数量，不包含尾块 + uint32_t ubMinBlockLen = 0; + uint32_t cacheLineLen = 0; + uint32_t alignPackLen = 0; + uint32_t totalAvailableCore = 0; + uint32_t totalUsedCoreNum_ = 0; + uint32_t totalUsedCoreNum = 0; + uint32_t totalCore = 0; + ge::DataType xInputDataType; + + bool isPerfBranch = false; + + ge::DataType biasDataType = ge::DT_FLOAT; + uint64_t quantScaleShapeSize = 0; + platform_ascendc::SocVersion curShortSocName_; + + const char* opName = ""; + SwiGluTilingData tilingData; + platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND910B; +}; + +void DequantSwigluQuantTiling::Reset() +{ + opName = nullptr; + return; +} + +ge::graphStatus DequantSwigluQuantTiling::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + OP_CHECK_IF(platformInfo == nullptr, OP_LOGE(opName, "fail to get platform info"), return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + curShortSocName_ = ascendcPlatform.GetSocVersion(); + totalCore = ascendcPlatform.GetCoreNumAiv(); + aicoreParams_.numBlocks = totalCore; + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + aicoreParams_.ubSize = ubSizePlatForm; + socVersion = ascendcPlatform.GetSocVersion(); + return ge::GRAPH_SUCCESS; +} + +inline ge::graphStatus DequantSwigluQuantTiling::SetTotalShape(gert::TilingContext* cont, const gert::Shape& inShape) +{ + int64_t shapeBefore = 1; + int64_t shapeAfter = 1; + int64_t dimNum = inShape.GetDimNum(); + CHECK_FAIL(cont, dimNum <= 1, "The shape dim of x can not be less than 2"); + + int64_t splitDim = dimNum - 1; // inDim default -1 + for (int64_t i = 0; i < splitDim; i++) { + shapeBefore *= inShape.GetDim(i); + } + shapeAfter = inShape.GetDim(splitDim); + // 如果shape不是2的倍数,返回 + + CHECK_FAIL(cont, shapeAfter % 2 != 0, "The shape dim of x dim must be even number"); + + tilingData.set_rowLen(shapeBefore); + // colLen为原shape除以2 + tilingData.set_colLen(shapeAfter / 2); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantTiling::checkWeightBiasActivate(gert::TilingContext* context) +{ + auto biasShapeShapePtr = context->GetOptionalInputShape(3); + if (biasShapeShapePtr != nullptr) { + auto biasInputDesc = context->GetOptionalInputDesc(3); + OP_CHECK_NULL_WITH_CONTEXT(context, biasInputDesc); + biasDataType = biasInputDesc->GetDataType(); + + bool checkBiasRes = biasDataType != ge::DT_INT32 && biasDataType != ge::DT_FLOAT && + biasDataType != ge::DT_FLOAT16 && biasDataType != ge::DT_BF16; + OP_CHECK_IF(checkBiasRes, + OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "bias", + ge::TypeUtils::DataTypeToSerialString(biasDataType).c_str(), "int32, float, fp16 or bf16"), + return ge::GRAPH_FAILED); + + uint64_t biasShapeSize = biasShapeShapePtr->GetStorageShape().GetShapeSize(); + OP_CHECK_IF(biasShapeSize != tilingData.get_colLen() * 2, + OP_LOGE_FOR_INVALID_SHAPESIZE(context->GetNodeName(), "bias", + std::to_string(biasShapeSize).c_str(), + (std::to_string(tilingData.get_colLen() * 2)).c_str()), + return ge::GRAPH_FAILED); + } + tilingData.set_biasIsEmpty(biasShapeShapePtr == nullptr); + // int32时 weight_scale为必选项 + auto weightScaleShapePtr = context->GetOptionalInputShape(1); + OP_CHECK_NULL_WITH_CONTEXT(context, weightScaleShapePtr); + + auto weightScaleInputDesc = context->GetOptionalInputDesc(1); + OP_CHECK_NULL_WITH_CONTEXT(context, weightScaleInputDesc); + ge::DataType weightScaleDataType = weightScaleInputDesc->GetDataType(); + OP_CHECK_IF(weightScaleDataType != ge::DT_FLOAT, + OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "weight_scale", + ge::TypeUtils::DataTypeToSerialString(weightScaleDataType).c_str(), "float32"), + return ge::GRAPH_FAILED); + + uint64_t weightScaleShapeSize = weightScaleShapePtr->GetStorageShape().GetShapeSize(); + OP_CHECK_IF(weightScaleShapeSize != tilingData.get_colLen() * 2, + OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON(context->GetNodeName(), "weight_scale", + std::to_string(weightScaleShapeSize).c_str(), + ("The shapesize of the weight scale is not equal to the last dimension of the xshape " + + std::to_string(tilingData.get_colLen() * 2)).c_str()), + return ge::GRAPH_FAILED); + + // int32时 activate_scale为可选项 + auto activateScaleShapePtr = context->GetOptionalInputShape(2); + if (activateScaleShapePtr != nullptr) { + auto activateScaleInputDesc = context->GetOptionalInputDesc(2); + OP_CHECK_NULL_WITH_CONTEXT(context, activateScaleInputDesc); + ge::DataType activateScaleDataType = activateScaleInputDesc->GetDataType(); + OP_CHECK_IF(activateScaleDataType != ge::DT_FLOAT, + OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "activation_scale", + ge::TypeUtils::DataTypeToSerialString(activateScaleDataType).c_str(), "float32"), + return ge::GRAPH_FAILED); + + uint64_t activateScaleShapeSize = activateScaleShapePtr->GetStorageShape().GetShapeSize(); + OP_CHECK_IF(activateScaleShapeSize != tilingData.get_rowLen(), + OP_LOGE_FOR_INVALID_SHAPESIZE(context->GetNodeName(), "activation_scale", + std::to_string(activateScaleShapeSize).c_str(), + ("equal to " + std::to_string(tilingData.get_rowLen())).c_str()), + return ge::GRAPH_FAILED); + } + tilingData.set_activateScaleIsEmpty(activateScaleShapePtr == nullptr); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantTiling::checkInputShape(gert::TilingContext* context, ge::DataType xDataType) +{ + if (xDataType == ge::DT_INT32) { + if (checkWeightBiasActivate(context) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + } + // quant_scale + auto quantScaleShapePtr = context->GetOptionalInputShape(4); // 3: bias idx + if (quantScaleShapePtr == nullptr) { + tilingData.set_quantScaleIsEmpty(1); + return ge::GRAPH_SUCCESS; + } + auto quantScaleInputDesc = context->GetOptionalInputDesc(4); + OP_CHECK_NULL_WITH_CONTEXT(context, quantScaleInputDesc); + ge::DataType quantScaleDataType = quantScaleInputDesc->GetDataType(); + OP_CHECK_IF(quantScaleDataType != ge::DT_FLOAT, + OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "quant_scale", + ge::TypeUtils::DataTypeToSerialString(quantScaleDataType).c_str(), "float32"), + return ge::GRAPH_FAILED); + quantScaleShapeSize = quantScaleShapePtr->GetStorageShape().GetShapeSize(); + bool checkQuantScaleSize = (quantScaleShapeSize != tilingData.get_colLen()) && (quantScaleShapeSize != 1); + OP_CHECK_IF(checkQuantScaleSize, + OP_LOGE_FOR_INVALID_SHAPESIZE(context->GetNodeName(), "quant_scale", + std::to_string(quantScaleShapeSize).c_str(), + (std::to_string(tilingData.get_colLen()) + " or 1").c_str()), + return ge::GRAPH_FAILED); + if (quantMode == 0) { + auto quantOffsetShapePtr = context->GetOptionalInputShape(5); + auto quantOffsetInputDesc = context->GetOptionalInputDesc(5); + OP_CHECK_NULL_WITH_CONTEXT(context, quantOffsetInputDesc); + ge::DataType quantOffsetDataType = quantOffsetInputDesc->GetDataType(); + OP_CHECK_IF(quantOffsetDataType != ge::DT_FLOAT, + OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "quant_offset", + ge::TypeUtils::DataTypeToSerialString(quantOffsetDataType).c_str(), "float32"), + return ge::GRAPH_FAILED); + uint64_t quantOffsetShapeSize = quantOffsetShapePtr->GetStorageShape().GetShapeSize(); + bool checkQuantOffsetSize = (quantOffsetShapeSize != tilingData.get_colLen()) && (quantOffsetShapeSize != 1); + OP_CHECK_IF(checkQuantOffsetSize, + OP_LOGE_FOR_INVALID_SHAPESIZE(context->GetNodeName(), "quant_offset", + std::to_string(quantOffsetShapeSize).c_str(), + (std::to_string(tilingData.get_colLen()) + " or 1").c_str()), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +bool DequantSwigluQuantTiling::SetAttr(const gert::RuntimeAttrs* attrs) +{ + auto isActivateLeftAttr = *(attrs->GetBool(0)); + auto str = attrs->GetStr(1); + std::string quantModeAttr{str}; + std::transform(quantModeAttr.begin(), quantModeAttr.end(), quantModeAttr.begin(), ::tolower); + + if ((quantModeAttr != "static") && (quantModeAttr != "dynamic")) { + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON( + context_->GetNodeName(), "quant_mode", + quantModeAttr.c_str(), + "quant_mode should be static or dynamic with case insensitive"); + return false; + } + activateLeft = (isActivateLeftAttr ? 1 : 0); + quantMode = ((quantModeAttr == "static") ? 0 : 1); + tilingData.set_activateLeft(activateLeft); + return true; +} + +bool DequantSwigluQuantTiling::GetBufferNumAndDataLenPerUB(uint64_t ubSize, int32_t dtype, uint64_t& dataLenPerUB) +{ + uint32_t singleDataSize = 1; + if (quantMode == 1) { + if (dtype == ge::DT_FLOAT16 || dtype == ge::DT_BF16) { + singleDataSize = DYNAMIC_BF16_INT16_TBUF_NUM_HALF * static_cast(sizeof(float)) + + static_cast(sizeof(int8_t)); + } else if (dtype == ge::DT_INT32) { + if ((biasDataType == ge::DT_INT32 || biasDataType == ge::DT_FLOAT)) { + singleDataSize = DYNAMIC_BF16_TBUF_NUM_HALF * static_cast(sizeof(float)) + + static_cast(sizeof(int8_t)); + } else { + singleDataSize = DYNAMIC_BF16_TBUF_NUM_HALF * static_cast(sizeof(float)) + + DYNAMIC_INT16_TBUF_NUM_HALF * static_cast(sizeof(int16_t)) + + static_cast(sizeof(int8_t)); + } + } + } + if (quantMode == 0) { + if (dtype == ge::DT_INT32) { + if ((biasDataType == ge::DT_INT32 || biasDataType == ge::DT_FLOAT)) { + singleDataSize = STATIC_BF16_TBUF_NUM_HALF * static_cast(sizeof(float)) + + static_cast(sizeof(int8_t)); /* 11 -> float 块数量 */ + } else { + singleDataSize = STATIC_BF16_TBUF_NUM_HALF * static_cast(sizeof(float)) + + DYNAMIC_INT16_TBUF_NUM_HALF * static_cast(sizeof(int16_t)) + + static_cast(sizeof(int8_t)); /* 11 -> float 块数量 */ + } + } else if (dtype == ge::DT_FLOAT16 || dtype == ge::DT_BF16) { + singleDataSize = STATIC_BF16_INT16_TBUF_NUM_HALF * static_cast(sizeof(float)) + + static_cast(sizeof(int8_t)); + } + } + dataLenPerUB = ubSize / singleDataSize; + return true; +} + +bool DequantSwigluQuantTiling::CalcUbMaxTileLen(uint64_t ubSize, int32_t dtype, GluSingleTilingOptParam& optTiling) +{ + // get buffernum and maxTileLen + uint64_t maxTileLenPerUB = 1; + if (!GetBufferNumAndDataLenPerUB(ubSize, dtype, maxTileLenPerUB)) { + OP_LOGE("DequantSwigluQuant", "CalcTiling Get maxTileLenPerUB %lu failed", maxTileLenPerUB); + return false; + } + optTiling.maxTileLen = AlignDown(maxTileLenPerUB, ALIGN_UINT_IN_CACHE_32B); // 32个元素对齐 + OP_LOGI("DequantSwigluQuant", "CalcTiling ubSize:%lu, maxTileLenPerUB:%u", ubSize, optTiling.maxTileLen); + return true; +} + +uint32_t DequantSwigluQuantTiling::getBaseColLenUpBound(GluSingleTilingOptParam& optTiling) +{ + uint32_t upBound = std::min(tilingData.get_colLen(), static_cast(optTiling.maxTileLen)); + if (tilingData.get_is32BAligned() == 1) { + upBound = std::min(upBound, static_cast(DISCONTINE_COPY_MAX_BLOCKLEN)); + } else { + upBound = std::min(upBound, static_cast(DISCONTINE_COPY_MAX_BLOCKLEN / sizeof(xInputDataType))); + } + + if (upBound < tilingData.get_colLen() && upBound > cacheLineLen) { + // 该种场景，每一个colLen至少被切割成2块，需要保证baseColLen为512B整数倍才高效 + return AlignDown(upBound, cacheLineLen); + } else { + return upBound; + } +} + +void DequantSwigluQuantTiling::SaveOptBaseShape( + uint32_t baseRowLen_, uint32_t baseColLen_, GluSingleTilingOptParam& optTiling) +{ + uint64_t totalTileNum = + std::min(static_cast(tilingData.get_rowLen()), static_cast(totalAvailableCore)); + uint64_t baseSize = static_cast(baseRowLen_ * baseColLen_); + if (static_cast(baseRowLen_) == 0 || static_cast(baseColLen_) == 0) { + OP_LOGI("SaveOptBaseShape", "baseRowLen_:%u or baseColLen:%u is zero.", baseRowLen_, baseColLen_); + return; + } + uint64_t baseTileNum = (baseRowLen_ == 0 ? 0 : (tilingData.get_rowLen() / baseRowLen_)) * + (baseColLen_ == 0 ? 0 : (tilingData.get_colLen() / baseColLen_)); + totalUsedCoreNum_ = std::min(totalTileNum, static_cast(totalAvailableCore)); + if(tilingData.get_colLen() < PERFORMANCE_COL_LEN + && tilingData.get_rowLen() < PERFORMANCE_ROW_LEN) { + totalUsedCoreNum_ = std::min(totalUsedCoreNum_, static_cast(MIN_CORE)); + } + optTiling.optBaseRowLen = baseRowLen_; + optTiling.optBaseColLen = baseColLen_; + optTiling.optTotalTileNum = totalTileNum; + optTiling.optBaseSize = baseSize; + optTiling.optBaseTileNum = baseTileNum; + optTiling.totalUsedCoreNum = totalUsedCoreNum_; + optTiling.tileNumPerCore = DivCeil(totalTileNum, totalUsedCoreNum_); +} + +bool DequantSwigluQuantTiling::CalcOptBaseShape(GluSingleTilingOptParam& optTiling, int32_t dtype) +{ + uint32_t baseColLen_ = getBaseColLenUpBound(optTiling); + uint32_t baseRowlen_ = 1; + if ((quantMode == 1) && (dtype == ge::DT_FLOAT16 || dtype == ge::DT_BF16)) { + baseRowlen_ = std::min( + optTiling.maxTileLen / AlignUp(baseColLen_, ALIGN_UINT_IN_CACHE_32B), + static_cast(tilingData.get_rowLen())); + baseRowlen_ = std::min(DivCeil(tilingData.get_rowLen(), totalAvailableCore), baseRowlen_); + } + SaveOptBaseShape(baseRowlen_, baseColLen_, optTiling); + return true; +} + +bool DequantSwigluQuantTiling::CalcOptTiling( + const uint64_t ubSize, const int32_t dtype, GluSingleTilingOptParam& optTiling) +{ + // 计算maxTilingLen + if (!CalcUbMaxTileLen(ubSize, dtype, optTiling)) { + return false; + } + // 计算最优的base块形状 + if (!CalcOptBaseShape(optTiling, dtype)) { + return false; + } + return true; +} + +bool DequantSwigluQuantTiling::CalcTiling( + const uint32_t totalCores, const uint64_t ubSize, const platform_ascendc::SocVersion socVersion_) +{ + totalAvailableCore = totalCores; + if (!GetLengthByType(xInputDataType, inputDTypeLen)) { + OP_LOGI("DequantSwigluQuant", "CalcTiling Unsupported input data type %d", xInputDataType); + return false; + } + ubMinBlockLen = ALIGN_UINT_IN_CACHE_32B / inputDTypeLen; // min block size + cacheLineLen = PACK_UINT_IN_CACHE_512B / inputDTypeLen; // bandwidth max efficiency + alignPackLen = cacheLineLen; // 默认512对齐，策略可调整 + OP_LOGI( + "DequantSwigluQuant", "CalcTiling GetLengthByType:%u ubMinBlockLen:%u cacheLineLen:%u alignPackLen:%u", + inputDTypeLen, ubMinBlockLen, cacheLineLen, alignPackLen); + // Is 32-byte aligned for split colLen? + tilingData.set_is32BAligned(tilingData.get_colLen() % ubMinBlockLen == 0); + // 310p not support Non-64B + const uint32_t blockSizeOf64B = ALIGN_UINT_IN_CACHE_64B / inputDTypeLen; + if (((socVersion_ == platform_ascendc::SocVersion::ASCEND310P)) && + (tilingData.get_colLen() % blockSizeOf64B != 0)) { + OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), "x", + std::to_string(tilingData.get_colLen()).c_str(), + "colLen (the last dimension of x) must be 64B aligned on ASCEND310P"); + return false; + } + GluSingleTilingOptParam optTilingDb; + if (!CalcOptTiling(ubSize, xInputDataType, optTilingDb)) { + return false; + } + const GluSingleTilingOptParam* const optTiling = &optTilingDb; + // 记录最优的结果 + tilingData.set_baseRowLen(optTiling->optBaseRowLen); + tilingData.set_baseColLen(optTiling->optBaseColLen); + totalUsedCoreNum = optTiling->totalUsedCoreNum; + tilingData.set_usedCoreNum(totalUsedCoreNum); + OP_LOGI( + "DequantSwigluQuant", "CalcTilingRES baseRowLen:%u baseColLen:%u", optTiling->optBaseRowLen, + optTiling->optBaseColLen); + return true; +} + +ge::graphStatus DequantSwigluQuantTiling::GetShapeAttrsInfo() +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantTiling::GetShapeAttrsInfoInner() +{ + opName = context_->GetNodeName(); + // 获取输入shape + auto xShapePtr = context_->GetInputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context_, xShapePtr); + const gert::Shape xShape = xShapePtr->GetStorageShape(); + auto inputDesc = context_->GetInputDesc(0); + OP_CHECK_NULL_WITH_CONTEXT(context_, inputDesc); + xInputDataType = inputDesc->GetDataType(); + if (SetTotalShape(context_, xShape) == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + // 获取输入属性 + const gert::RuntimeAttrs* attrs = context_->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context_, attrs); + + if (!SetAttr(attrs)) { + return ge::GRAPH_FAILED; + } + + if (checkInputShape(context_, xInputDataType) == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + auto yShapePtr = context_->GetOutputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context_, yShapePtr); + const gert::Shape yShape = yShapePtr->GetStorageShape(); + + int32_t dimNum = xShape.GetDimNum(); + if(xShape.GetDimNum() != yShape.GetDimNum()){ + std::string incorrectDims = std::to_string(xShape.GetDimNum()) + " and " + std::to_string(yShape.GetDimNum()); + OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(opName, "x and y", + incorrectDims.c_str(), + "The shape of y must be equal to the shape of x"); + } + + if(xShape.GetDim(dimNum - 1) != yShape.GetDim(dimNum - 1) * 2){ + std::string incorrectDims = std::to_string(xShape.GetDimNum()) + " and " + std::to_string(yShape.GetDimNum()); + OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(opName, "x and y", + incorrectDims.c_str(), + "The last dimension of x must be twice the last dimension of y."); + } + + auto scaleShapePtr = context_->GetOutputShape(1); + OP_CHECK_NULL_WITH_CONTEXT(context_, scaleShapePtr); + const gert::Shape scaleShape = scaleShapePtr->GetStorageShape(); + + if (static_cast(scaleShape.GetShapeSize()) != tilingData.get_rowLen()) { + std::string incorrectSize = std::to_string(static_cast(scaleShape.GetShapeSize())); + std::string reason = + "scale's shapesize must be equal to row length" + std::to_string(tilingData.get_rowLen()) + + "(row length is total number of elements of x across all dimensions except the last one.)"; + OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON(opName, "scale", incorrectSize.c_str(), reason.c_str()); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantTiling::DoOpTiling() +{ + if (GetShapeAttrsInfoInner() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + if (!CalcTiling(totalCore, aicoreParams_.ubSize, curShortSocName_)) { + return ge::GRAPH_FAILED; + } + isPerfBranch = isPerformanceBranch(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantTiling::DoLibApiTiling() +{ + return ge::GRAPH_SUCCESS; +} + +int64_t DequantSwigluQuantTiling::getTilingKeyStatic( + const int32_t inputDtype, const ge::DataType biasType, const int64_t scaleSize) const +{ + if (inputDtype != ge::DT_INT32) { + if (scaleSize == 1) { + if (inputDtype == ge::DT_FLOAT16) { + return STATIC_FLOAT16_X; + } else { + return STATIC_BFLOAT16_X; + } + } else { + if (inputDtype == ge::DT_FLOAT16) { + return STATIC_FLOAT16_XD; + } else { + return STATIC_BFLOAT16_XD; + } + } + } + if (scaleSize == 1) { + if (biasType == ge::DT_INT32) { + return STATIC_INT_X_INT_BIAS_QUANT_ONE; + } else if (biasType == ge::DT_FLOAT) { + return STATIC_INT_X_FLOAT32_BIAS_QUANT_ONE; + } else if (biasType == ge::DT_FLOAT16) { + return STATIC_INT_X_FLOAT16_BIAS_QUANT_ONE; + } else { + return STATIC_INT_X_BFLOAT16_BIAS_QUANT_ONE; + } + } else { + if (biasType == ge::DT_INT32) { + return STATIC_INT_X_INT_BIAS_QUANT_D; + } else if (biasType == ge::DT_FLOAT) { + return STATIC_INT_X_FLOAT32_BIAS_QUANT_D; + } else if (biasType == ge::DT_FLOAT16) { + return STATIC_INT_X_FLOAT16_BIAS_QUANT_D; + } else { + return STATIC_INT_X_BFLOAT16_BIAS_QUANT_D; + } + } +} + +int64_t DequantSwigluQuantTiling::getTilingKeyDynamic( + const int32_t inputDtype, const ge::DataType biasType, const int64_t scaleSize) const +{ + if (inputDtype != ge::DT_INT32) { + if (inputDtype == ge::DT_FLOAT16) { + if (scaleSize == 1) { + return DYNAMIC_FLOAT16_X; + } else { + return DYNAMIC_FLOAT16_XD; + } + } else { + if (scaleSize == 1) { + return DYNAMIC_BFLOAT16_X; + } else { + return DYNAMIC_BFLOAT16_XD; + } + } + } + if (scaleSize == 1) { + if (biasType == ge::DT_INT32) { + return DYNAMIC_INT_X_INT_BIAS_QUANT_ONE; + } else if (biasType == ge::DT_FLOAT) { + return DYNAMIC_INT_X_FLOAT32_BIAS_QUANT_ONE; + } else if (biasType == ge::DT_FLOAT16) { + return DYNAMIC_INT_X_FLOAT16_BIAS_QUANT_ONE; + } else { + return DYNAMIC_INT_X_BFLOAT16_BIAS_QUANT_ONE; + } + } else { + if (biasType == ge::DT_INT32) { + return DYNAMIC_INT_X_INT_BIAS_QUANT_D; + } else if (biasType == ge::DT_FLOAT) { + if(isPerfBranch) { + return DYNAMIC_INT_X_FLOAT32_BIAS_QUANT_D_PERFORMANCE; + } + return DYNAMIC_INT_X_FLOAT32_BIAS_QUANT_D; + } else if (biasType == ge::DT_FLOAT16) { + return DYNAMIC_INT_X_FLOAT16_BIAS_QUANT_D; + } else { + return DYNAMIC_INT_X_BFLOAT16_BIAS_QUANT_D; + } + } +} + +bool DequantSwigluQuantTiling::isPerformanceBranch() { + if(tilingData.get_is32BAligned() == 1 + && tilingData.get_colLen() <= PERFORMANCE_COL_LEN + && tilingData.get_baseRowLen() == 1 + && tilingData.get_baseColLen() == tilingData.get_colLen() + && tilingData.get_biasIsEmpty() == 1 + && tilingData.get_activateScaleIsEmpty() == 0) { + return true; + } + return false; +} + +uint64_t DequantSwigluQuantTiling::GetTilingKey() const +{ + if (quantMode == 0) { // static + return getTilingKeyStatic(xInputDataType, biasDataType, quantScaleShapeSize); + } else { // dynamic + return getTilingKeyDynamic(xInputDataType, biasDataType, quantScaleShapeSize); + } +} + +ge::graphStatus DequantSwigluQuantTiling::GetWorkspaceSize() +{ + // 计算workspace大小，无需workspace临时空间，不存在多核同步，预留固定大小即可 + workspaceSize_ = USER_WORKSPACE; + if (quantMode == 1 && (tilingData.get_colLen() > tilingData.get_baseColLen())) { + workspaceSize_ += (totalUsedCoreNum * tilingData.get_colLen() * sizeof(float)); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus DequantSwigluQuantTiling::PostTiling() +{ + context_->SetBlockDim(totalCore); + size_t* currentWorkspace = context_->GetWorkspaceSizes(1); + currentWorkspace[0] = workspaceSize_; + OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetRawTilingData()); + + tilingData.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData.GetDataSize()); + context_->SetBlockDim(totalUsedCoreNum); + return ge::GRAPH_SUCCESS; +} + +REGISTER_TILING_TEMPLATE("DequantSwigluQuant", DequantSwigluQuantTiling, 1); +} // namespace optiling diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_host/swi_glu_grad_regbase_tiling.h b/csrc/ascend/moe/dequant_swiglu_quant/op_host/swi_glu_grad_regbase_tiling.h new file mode 100644 index 000000000..bdedb042a --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_host/swi_glu_grad_regbase_tiling.h @@ -0,0 +1,27 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file swi_glu_grad_regbase_tiling.h + * \brief + */ + +struct GluBaseTilingData { + int64_t rowTotal; + int64_t colTotal; + int64_t rowBase; + int64_t colBase; + int64_t rowTail; + int64_t colTail; + int64_t ubSize; + int64_t rowTileNum; + int64_t colTileNum; + int64_t usedCoreNum; +}; \ No newline at end of file diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_host/swi_glu_grad_tiling_regbase.h b/csrc/ascend/moe/dequant_swiglu_quant/op_host/swi_glu_grad_tiling_regbase.h new file mode 100644 index 000000000..a2194c060 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_host/swi_glu_grad_tiling_regbase.h @@ -0,0 +1,75 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file swi_glu_grad_tiling_regbase.h + * \brief + */ +#pragma once + +#include +#include "tiling/tiling_api.h" +#include "../tiling_base/tiling_base.h" +#include "swi_glu_tiling.h" + +namespace optiling { + +class GluBaseTiling4RegBase : public TilingBaseClass { +public: + explicit GluBaseTiling4RegBase(gert::TilingContext *context) : TilingBaseClass(context), opName_(context->GetNodeName()) {} + +protected: + constexpr static int64_t UB_RESERVED_BUFF {0}; + constexpr static int64_t BASE_BLOCK_SIZE {8192}; + constexpr static int64_t MOVE_ALIGN_LIMIT_BYTE {1024}; + constexpr static int64_t BASE_BLOCK_COPY_ALIGN {512}; + + bool IsCapable() override; + // 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小 + ge::graphStatus GetPlatformInfo() override; + // 2、获取INPUT/OUTPUT/ATTR信息 + ge::graphStatus GetShapeAttrsInfo() override; + // 3、计算数据切分TilingData + ge::graphStatus DoOpTiling() override; + // 4、计算高阶API的TilingData + ge::graphStatus DoLibApiTiling() override; + // 5、计算TilingKey + uint64_t GetTilingKey() const override; + // 6、计算Workspace 大小 + ge::graphStatus GetWorkspaceSize() override; + // 7、保存Tiling数据 + ge::graphStatus PostTiling() override; + + void DumpTilingInfo() override; + +private: + const std::string opName_; + uint64_t ubSize_ {0}; + GluBaseTilingData tilingData_; + int64_t rowTotalNum_ {0}; + int64_t colTotalNum_ {0}; + int64_t rowNormalNum_ {0}; + int64_t colNormalNum_ {0}; + int64_t rowTailNum_ {0}; + int64_t colTailNum_ {0}; + uint64_t usedCoreNum_ {0}; + uint32_t rowTileNum_ {0}; + uint32_t colTileNum_ {0}; + ge::DataType dataType_ {DT_FLOAT}; + uint64_t dataSize_ {0}; + + bool CalcShapeTo2D(const gert::Shape& inShape, const int64_t splitDim); + bool CheckShapeValid(const gert::Shape& gradYShape, const gert::Shape& xShape, const int64_t dim); + void AutoTiling(); + std::set FindUniqueCut(); + uint64_t ComputeTiling(const std::vector& args) const; + void SetTilingData(); +}; +} // namespace optiling \ No newline at end of file diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_host/swi_glu_tiling.h b/csrc/ascend/moe/dequant_swiglu_quant/op_host/swi_glu_tiling.h new file mode 100644 index 000000000..1bdedfe25 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_host/swi_glu_tiling.h @@ -0,0 +1,78 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file swi_glu_tiling.h + * \brief + */ +#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_SWIGLU_H_ +#define AIR_CXX_RUNTIME_V2_OP_IMPL_SWIGLU_H_ + +#include +#include "register/op_impl_registry.h" +#include "util/math_util.h" +#include "log/log.h" +#include "tiling/platform/platform_ascendc.h" +#include "platform/platform_infos_def.h" +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "../tiling_base/tiling_templates_registry.h" +#include "swi_glu_grad_regbase_tiling.h" + +namespace optiling { +const int64_t STATIC_FLOAT16_X = 10000; +const int64_t STATIC_BFLOAT16_X = 10001; +const int64_t STATIC_FLOAT16_XD = 10002; +const int64_t STATIC_BFLOAT16_XD = 10003; +const int64_t STATIC_INT_X_INT_BIAS_QUANT_ONE = 10004; +const int64_t STATIC_INT_X_INT_BIAS_QUANT_D = 10005; +const int64_t STATIC_INT_X_FLOAT16_BIAS_QUANT_ONE = 10006; +const int64_t STATIC_INT_X_FLOAT16_BIAS_QUANT_D = 10007; +const int64_t STATIC_INT_X_FLOAT32_BIAS_QUANT_ONE = 10008; +const int64_t STATIC_INT_X_FLOAT32_BIAS_QUANT_D = 10009; +const int64_t STATIC_INT_X_BFLOAT16_BIAS_QUANT_ONE = 10010; +const int64_t STATIC_INT_X_BFLOAT16_BIAS_QUANT_D = 10011; + +const int64_t DYNAMIC_FLOAT16_X = 30009; +const int64_t DYNAMIC_BFLOAT16_X = 30011; +const int64_t DYNAMIC_FLOAT16_XD = 30010; +const int64_t DYNAMIC_BFLOAT16_XD = 30012; +const int64_t DYNAMIC_INT_X_INT_BIAS_QUANT_ONE = 30001; +const int64_t DYNAMIC_INT_X_INT_BIAS_QUANT_D = 30005; +const int64_t DYNAMIC_INT_X_FLOAT16_BIAS_QUANT_ONE = 30003; +const int64_t DYNAMIC_INT_X_FLOAT16_BIAS_QUANT_D = 30007; +const int64_t DYNAMIC_INT_X_FLOAT32_BIAS_QUANT_ONE = 30002; +const int64_t DYNAMIC_INT_X_FLOAT32_BIAS_QUANT_D = 30006; +const int64_t DYNAMIC_INT_X_BFLOAT16_BIAS_QUANT_ONE = 30004; +const int64_t DYNAMIC_INT_X_BFLOAT16_BIAS_QUANT_D = 30008; + +BEGIN_TILING_DATA_DEF(SwiGluTilingData) + TILING_DATA_FIELD_DEF(uint32_t, is32BAligned); + TILING_DATA_FIELD_DEF(uint32_t, isDoubleBuffer); + TILING_DATA_FIELD_DEF(uint64_t, rowLen); + TILING_DATA_FIELD_DEF(uint64_t, colLen); + TILING_DATA_FIELD_DEF(uint32_t, baseRowLen); + TILING_DATA_FIELD_DEF(uint32_t, baseColLen); + TILING_DATA_FIELD_DEF(uint32_t, activateLeft); + TILING_DATA_FIELD_DEF(uint32_t, biasIsEmpty); + TILING_DATA_FIELD_DEF(uint32_t, quantScaleIsEmpty); + TILING_DATA_FIELD_DEF(uint32_t, activateScaleIsEmpty); + TILING_DATA_FIELD_DEF(uint64_t, swiColLen); + TILING_DATA_FIELD_DEF(uint64_t, perRowLen); + TILING_DATA_FIELD_DEF(uint64_t, modRowLen); + TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(SwiGlu, SwiGluTilingData) +REGISTER_TILING_DATA_CLASS(SwiGluGrad, SwiGluTilingData) +REGISTER_TILING_DATA_CLASS(DequantSwigluQuant, SwiGluTilingData) + +} // namespace optiling +#endif // AIR_CXX_RUNTIME_V2_OP_IMPL_SWIGLU_H_ diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.cpp b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.cpp new file mode 100644 index 000000000..fb645f455 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.cpp @@ -0,0 +1,433 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant.cpp + * \brief + */ +#include "kernel_operator.h" +#include "kernel_tiling/kernel_tiling.h" +#if (ORIG_DTYPE_X == DT_INT32) || (ORIG_DTYPE_X == DT_BF16) + #include "dequant_swiglu_quant.h" + #include "dequant_swiglu_quant_cut_group.h" +#endif + +#include "dequant_swiglu_quant_static_bf16.hpp" +#include "dequant_swiglu_quant_static_bias_int32.hpp" +#include "dequant_swiglu_quant_static_bias_float.hpp" +#include "dequant_swiglu_quant_dynamic_bf16.hpp" +#include "dequant_swiglu_quant_dynamic_bias_int32.hpp" +#include "dequant_swiglu_quant_dynamic_bias_float.hpp" +#include "dequant_swiglu_quant_dynamic_performance.hpp" + +using namespace AscendC; + +// DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP32_QS HAS_GROUP(100000000) + QS_OFFSET(100) * QS_FP32(0) +// DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP16_QS HAS_GROUP(100000000) + QS_OFFSET(100) * QS_FP16(1) +// DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_QS HAS_GROUP(100000000) + QS_OFFSET(100) * QS_BF16(2) +// DEQUANT_SWIGLU_QUANT_WITHOUT_GROUP_FP32_QS NO_GROUP(200000000) + QS_OFFSET(100) * QS_FP32(0) +// DEQUANT_SWIGLU_QUANT_WITHOUT_GROUP_FP16_QS NO_GROUP(200000000) + QS_OFFSET(100) * QS_FP16(1) +// DEQUANT_SWIGLU_QUANT_WITHOUT_GROUP_BF16_QS NO_GROUP(200000000) + QS_OFFSET(100) * QS_BF16(2) +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_BIAS_FP32_QS 100000000 +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_NO_BIAS_FP32_QS DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_BIAS_FP32_QS +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_BIAS_FP16_QS 100000100 +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_NO_BIAS_FP16_QS DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_BIAS_FP16_QS +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_BIAS_BF16_QS 100000200 +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_NO_BIAS_BF16_QS DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_BIAS_BF16_QS +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP16_BIAS_FP32_QS 100001000 +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP16_BIAS_FP16_QS 100001100 +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP16_BIAS_BF16_QS 100001200 +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP32_BIAS_FP32_QS 100002000 +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP32_BIAS_FP16_QS 100002100 +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP32_BIAS_BF16_QS 100002200 +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_INT32_BIAS_FP32_QS 100003000 +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_INT32_BIAS_FP16_QS 100003100 +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_INT32_BIAS_BF16_QS 100003200 +#define DEQUANT_SWIGLU_QUANT_WITHOUT_GROUP_FP32_QS 200000000 +#define DEQUANT_SWIGLU_QUANT_WITHOUT_GROUP_FP16_QS 200000100 +#define DEQUANT_SWIGLU_QUANT_WITHOUT_GROUP_BF16_QS 200000200 +// cut by groupnum +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP32_QS_GR 110000000 +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP16_QS_GR 110000100 +#define DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_QS_GR 110000200 + +extern "C" __global__ __aicore__ void dequant_swiglu_quant(GM_ADDR xGM, GM_ADDR weightSscaleGM, + GM_ADDR activationScaleGM, GM_ADDR biasGM, + GM_ADDR quantScaleGM, GM_ADDR quantOffsetGM, + GM_ADDR groupIndex, GM_ADDR yGM, GM_ADDR scaleGM, + GM_ADDR workspace, GM_ADDR tiling) +{ + if (workspace == nullptr) { + return; + } + + GM_ADDR userspace = GetUserWorkspace(workspace); + if (userspace == nullptr) { + return; + } + TPipe pipe; +#if (ORIG_DTYPE_X == DT_INT32) + if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_BIAS_FP32_QS)) { +#if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); +#endif + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_BIAS_FP16_QS)) { +#if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); +#endif + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_BIAS_BF16_QS)) { +#if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); +#endif + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP16_BIAS_FP32_QS)) { + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP16_BIAS_FP16_QS)) { + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP16_BIAS_BF16_QS)) { +#if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); +#endif + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP32_BIAS_FP32_QS)) { + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP32_BIAS_FP16_QS)) { + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP32_BIAS_BF16_QS)) { +#if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); +#endif + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_INT32_BIAS_FP32_QS)) { + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_INT32_BIAS_FP16_QS)) { + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_INT32_BIAS_BF16_QS)) { +#if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); +#endif + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITHOUT_GROUP_FP32_QS)) { + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + // DTYPE_GROUP_INDEX == float mean have no groupIndex + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, nullptr, quantScaleGM, nullptr, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITHOUT_GROUP_FP16_QS)) { + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + // DTYPE_GROUP_INDEX == float mean have no groupIndex + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, nullptr, quantScaleGM, nullptr, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITHOUT_GROUP_BF16_QS)) { +#if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + // DTYPE_GROUP_INDEX == float mean have no groupIndex + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, nullptr, quantScaleGM, nullptr, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); +#endif + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP32_QS_GR)) { + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantGroupOps::DequantSwigluQuantGroup op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, tilingData); + op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP16_QS_GR)) { + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantGroupOps::DequantSwigluQuantGroup op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, tilingData); + op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_QS_GR)) { +#if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantGroupOps::DequantSwigluQuantGroup op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, tilingData); + op.Process(); +#endif + } else if (TILING_KEY_IS(10004)) { + // ORIG_DTYPE_BIAS == DT_INT32 + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantStaticBiasInt32 op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, tilingData, + &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(10005)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantStaticBiasInt32 op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, tilingData, + &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(30001)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantDynamicBiasInt32 op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, + tilingData, &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(30005)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantDynamicBiasInt32 op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, + tilingData, &(pipe)); + op.Process(); + } + // ORIG_DTYPE_BIAS == DT_FLOAT16 + else if (TILING_KEY_IS(10006)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantStaticBiasFloat op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, tilingData, + &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(10007)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantStaticBiasFloat op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, tilingData, + &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(30003)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantDynamicBiasFloat op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, + tilingData, &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(30007)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantDynamicBiasFloat op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, + tilingData, &(pipe)); + op.Process(); + } + // ORIG_DTYPE_BIAS == DT_FLOAT + else if (TILING_KEY_IS(10008)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantStaticBiasFloat op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, tilingData, + &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(10009)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantStaticBiasFloat op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, tilingData, + &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(30002)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantDynamicBiasFloat op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, + tilingData, &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(30013)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantDynamicPerformance op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, + tilingData, &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(30006)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantDynamicBiasFloat op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, + tilingData, &(pipe)); + op.Process(); + } +#if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) + // ORIG_DTYPE_BIAS == DT_BF16 + else if (TILING_KEY_IS(10010)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantStaticBiasFloat op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, tilingData, + &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(10011)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantStaticBiasFloat op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, tilingData, + &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(30004)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantDynamicBiasFloat op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, + tilingData, &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(30008)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantDynamicBiasFloat op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, + tilingData, &(pipe)); + op.Process(); + } +#endif +#endif +#if (ORIG_DTYPE_X == DT_FLOAT16) + if (TILING_KEY_IS(10000)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantStaticBF16 op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, tilingData, + &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(10002)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantStaticBF16 op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, tilingData, + &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(30009)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantDynamicBF16 op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, + tilingData, &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(30010)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantDynamicBF16 op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, + tilingData, &(pipe)); + op.Process(); + } +#endif +#if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) && (ORIG_DTYPE_X == DT_BF16) + if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_NO_BIAS_FP32_QS)) { + // New tiling branch for BF16 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, nullptr, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_NO_BIAS_FP16_QS)) { + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, nullptr, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_WITH_GROUP_NO_BIAS_BF16_QS)) { + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantBaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantBaseTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuantOps::DequantSwigluQuantBase op(&pipe); + op.Init(xGM, weightSscaleGM, activationScaleGM, nullptr, quantScaleGM, quantOffsetGM, groupIndex, yGM, scaleGM, + tilingData); + op.Process(); + } else if (TILING_KEY_IS(10001)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantStaticBF16 op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, tilingData, + &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(10003)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantStaticBF16 op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, tilingData, + &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(30011)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantDynamicBF16 op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, + tilingData, &(pipe)); + op.Process(); + } else if (TILING_KEY_IS(30012)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantDynamicBF16 op; + op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, + tilingData, &(pipe)); + op.Process(); + } +#endif +} \ No newline at end of file diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.h b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.h new file mode 100644 index 000000000..1300daa74 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.h @@ -0,0 +1,812 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant.h + * \brief + */ + +#ifndef DEQUANT_SWIGLU_QUANT_H +#define DEQUANT_SWIGLU_QUANT_H + +#include "kernel_tiling/kernel_tiling.h" +#include "kernel_operator.h" + +#define TEMPLATE_DSQ_DECLARE template +#define TEMPLATE_DSQ_ARGS TBias, TQuantScale, TGroup, TXGm + +namespace DequantSwigluQuantOps { +using namespace AscendC; +constexpr static int64_t DB_BUFFER = 1; +constexpr static int64_t BLOCK_SIZE = 32; +constexpr static int64_t BLOCK_ELEM = BLOCK_SIZE / sizeof(float); +constexpr static int64_t MASK_NUM_T32 = 256 / sizeof(float); +constexpr static int64_t MASK_BLK_STRIDE = 8; +constexpr static int64_t SWI_FACTOR = 2; +constexpr static float DYNAMIC_QUANT_FACTOR = 1.0 / 127.0; + +__aicore__ inline void CopyLocalContiguousFloat( + const LocalTensor& dst, const LocalTensor& src, uint32_t count) +{ + constexpr uint32_t MAX_REPEAT_TIMES = 255; + constexpr uint32_t MAX_REPEAT_ELEMS = MASK_NUM_T32 * MAX_REPEAT_TIMES; + CopyRepeatParams copyParams{1, 1, MASK_BLK_STRIDE, MASK_BLK_STRIDE}; + uint32_t offset = 0; + + while (count >= MAX_REPEAT_ELEMS) { + Copy(dst[offset], src[offset], MASK_NUM_T32, MAX_REPEAT_TIMES, copyParams); + offset += MAX_REPEAT_ELEMS; + count -= MAX_REPEAT_ELEMS; + } + + if (count >= MASK_NUM_T32) { + uint8_t repeatTimes = static_cast(count / MASK_NUM_T32); + Copy(dst[offset], src[offset], MASK_NUM_T32, repeatTimes, copyParams); + offset += repeatTimes * MASK_NUM_T32; + count -= repeatTimes * MASK_NUM_T32; + } + + if (count > 0) { + Copy(dst[offset], src[offset], count, 1, copyParams); + } +} + +TEMPLATE_DSQ_DECLARE +class DequantSwigluQuantBase +{ +public: + static constexpr bool hasGroupIndex_ = !IsSameType::value; + __aicore__ inline DequantSwigluQuantBase(TPipe* pipe) + { + pipe_ = pipe; + }; + + __aicore__ inline void Init( + GM_ADDR x, GM_ADDR weightScale, GM_ADDR activationScale, GM_ADDR bias, GM_ADDR quantScale, GM_ADDR quantOffset, + GM_ADDR groupIndex, GM_ADDR y, GM_ADDR scale, const DequantSwigluQuantBaseTilingData* tilingData); + __aicore__ inline void Process(); + __aicore__ inline void ComputeReduceMax(const LocalTensor& tempRes, int32_t calCount); + __aicore__ inline void ProcessSingleGroup(int64_t groupIdx, int64_t realCount, int64_t globalOffset); + __aicore__ inline void ProcessSingleGroupPerCore(int64_t groupIdx, int64_t dimxCore, int64_t dimxCoreOffset); + __aicore__ inline void CreateOffsetLocalTensor(uint32_t tensorLen, int swigluMode); + __aicore__ inline void SwiGluGate( + int32_t proDimsx, const LocalTensor& xLocalF32); + __aicore__ inline void DynamicQuant( + const LocalTensor& tmpUbF32Act, const LocalTensor& tmpUbF32Gate, + const LocalTensor& inScaleLocal, uint32_t proDimsx); + __aicore__ inline void StaticQuant( + const LocalTensor& tmpUbF32Act, const LocalTensor& tmpUbF32Gate, + const LocalTensor& inScaleLocal, uint32_t proDimsx); + __aicore__ inline void CopyInWeightScale(int64_t groupIdx); + __aicore__ inline void CopyInQuantScale(int64_t groupIdx); + __aicore__ inline void CopyInBias(int64_t groupIdx); + __aicore__ inline void ParamDequeAndCast(); + __aicore__ inline void CopyInXAct(int32_t proDimsx, int64_t xDimxOffset); + __aicore__ inline void Compute(int32_t proDimsx); + __aicore__ inline void ComputeDequant(int32_t proDimsx); + __aicore__ inline void ComputeSwiGLU(int32_t proDimsx); + __aicore__ inline void ComputeQuant(int32_t proDimsx); + __aicore__ inline void CopyOut(int32_t proDimsx, int64_t xDimxOffset); + __aicore__ inline void ParamFree(); + __aicore__ inline void CastFloatToInt8( + const LocalTensor& tmpUbF32Act, const LocalTensor& tmpUbF32Gate, uint32_t proDimsx, LocalTensor& yOut); + template + __aicore__ inline void CopyReshape(LocalTensor& dstTensor, LocalTensor& oriTensor, uint32_t rowNum, uint32_t colNum, CopyRepeatParams param); + +protected: + /* global memory address */ + // input global mem + GlobalTensor xGm_; + GlobalTensor weightScaleGm_; + GlobalTensor activationScaleGm_; + GlobalTensor biasGm_; + GlobalTensor quantScaleGm_; + GlobalTensor quantOffsetGm_; + GlobalTensor groupIndexGm_; + + // output global mem + GlobalTensor yGm_; + GlobalTensor scaleGm_; + + /* ub memory tensor */ + LocalTensor weightScaleLocal_; + LocalTensor inScaleLocal_; // quant scale and quant offset + LocalTensor biasLocal_; + LocalTensor biasLocalF32_; + LocalTensor xOffsetLocalU32_; // offset for gather + + /* ascendc variable */ + TPipe* pipe_ = nullptr; + TQue xActQueue_; + TQue inScaleQueue_; + TQue weightScaleQueue_; + TQue biasQueue_; + TQue outQueue_; + + TBuf tmpBuf1_; + TBuf tmpBuf2_; // only use in swigluMode == 1 + + uint32_t blockIdx_ = GetBlockIdx(); + int64_t realDimx_ = 0; + int64_t groupOffset_ = 0; + float quantScale_ = 1.0f; + float quantOffset_ = 1.0f; + + uint32_t UbSingleOutSize_ = 0; + uint32_t TBufActSclInOfs_ = 0; + uint32_t TBufXLocalInOfs_ = 0; + + int32_t actOffset_; + int32_t gateOffset_; + + const DequantSwigluQuantBaseTilingData* tl_ = nullptr; +}; +// 公共函数实现 + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::Init( + GM_ADDR x, GM_ADDR weightScale, GM_ADDR activationScale, GM_ADDR bias, GM_ADDR quantScale, GM_ADDR quantOffset, + GM_ADDR groupIndex, GM_ADDR y, GM_ADDR scale, const DequantSwigluQuantBaseTilingData* tilingData) +{ + tl_ = tilingData; + xGm_.SetGlobalBuffer((__gm__ TXGm*)x); + weightScaleGm_.SetGlobalBuffer((__gm__ float*)weightScale); + activationScaleGm_.SetGlobalBuffer((__gm__ float*)activationScale); + biasGm_.SetGlobalBuffer((__gm__ TBias*)bias); + quantScaleGm_.SetGlobalBuffer((__gm__ TQuantScale*)quantScale); + if constexpr (hasGroupIndex_) { + groupIndexGm_.SetGlobalBuffer((__gm__ TGroup*)groupIndex); + } + // static quant + if (tl_->quantMode == 0) { + quantOffsetGm_.SetGlobalBuffer((__gm__ TQuantScale*)quantOffset); + } + yGm_.SetGlobalBuffer((__gm__ int8_t*)y); + scaleGm_.SetGlobalBuffer((__gm__ float*)scale); + + UbSingleOutSize_ = static_cast(tl_->UbFactorDimx * tl_->outDimy); + TBufActSclInOfs_ = static_cast(tl_->UbFactorDimx * tl_->inDimy); +#if (ORIG_DTYPE_X == DT_BF16) + TBufXLocalInOfs_ = TBufActSclInOfs_; +#endif + + // swiglu offset + actOffset_ = tl_->actRight * tl_->UbFactorDimy; + gateOffset_ = tl_->UbFactorDimy - actOffset_; + + // init buffer + pipe_->InitBuffer( + xActQueue_, DB_BUFFER, (UbSingleOutSize_ * SWI_FACTOR + tl_->UbFactorDimx * BLOCK_ELEM) * sizeof(int32_t)); + pipe_->InitBuffer(weightScaleQueue_, 1, tl_->inDimy * sizeof(float)); + + if (tl_->quantMode == 0) { + pipe_->InitBuffer(inScaleQueue_, 1, tl_->outDimy * SWI_FACTOR * sizeof(float)); + } else { + pipe_->InitBuffer(inScaleQueue_, 1, tl_->outDimy * sizeof(float)); + } + + if (tl_->hasBias == 1) { + pipe_->InitBuffer(biasQueue_, 1, tl_->inDimy * sizeof(float)); + } + pipe_->InitBuffer(outQueue_, 1, UbSingleOutSize_ * sizeof(int8_t) + tl_->UbFactorDimx * sizeof(float) + BLOCK_SIZE); + + pipe_->InitBuffer(tmpBuf1_, UbSingleOutSize_ * SWI_FACTOR * sizeof(float)); + if (tl_->swigluMode == 1) { + pipe_->InitBuffer( + tmpBuf2_, + UbSingleOutSize_ * sizeof(int32_t) + UbSingleOutSize_ * sizeof(uint8_t)); // for gather offset and clamp + } +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::Process() +{ + if constexpr (!hasGroupIndex_) { + realDimx_ = tl_->inDimx; + // do protect realDimx_ < 0, ignore this group + realDimx_ = (realDimx_ < 0) ? 0 : realDimx_; + ProcessSingleGroup(0, realDimx_, 0); + return; + } + + CreateOffsetLocalTensor(UbSingleOutSize_, tl_->swigluMode); + + groupOffset_ = 0; + for (int32_t groupIdx = 0; groupIdx < tl_->inGroupNum; ++groupIdx) { + int64_t realGroupIdx = + tl_->speGroupType == 0 ? static_cast(groupIdx) : static_cast(groupIndexGm_(groupIdx * 2)); + realDimx_ = tl_->speGroupType == 0 ? static_cast(groupIndexGm_(groupIdx)) : + static_cast(groupIndexGm_(groupIdx * 2 + 1)); + // do protect realDimx_ < 0, ignore this group + realDimx_ = (realDimx_ < 0) ? 0 : realDimx_; + if (realDimx_ > 0 && groupOffset_ < tl_->inDimx) { + ProcessSingleGroup(realGroupIdx, realDimx_, groupOffset_); + groupOffset_ += realDimx_; + } + // speGroupindex场景下出现异常值(realDimx_ < 0), 退出计算 + if (tl_->speGroupType == 1 && realDimx_ <= 0) { + break; + } + } +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::ProcessSingleGroup( + int64_t groupIdx, int64_t realCount, int64_t globalOffset) +{ + // do block tiling again + int32_t blockDimxFactor = (realCount + tl_->maxCoreNum - 1) / tl_->maxCoreNum; + int32_t realCoreDim = (realCount + blockDimxFactor - 1) / blockDimxFactor; + + if (blockIdx_ < realCoreDim) { + int32_t blockDimxTailFactor = realCount - blockDimxFactor * (realCoreDim - 1); + int32_t dimxCore = blockIdx_ == (realCoreDim - 1) ? blockDimxTailFactor : blockDimxFactor; + int64_t coreDimxOffset = blockDimxFactor * blockIdx_ + globalOffset; + ProcessSingleGroupPerCore(static_cast(groupIdx), static_cast(dimxCore), coreDimxOffset); + } +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::CopyInWeightScale(int64_t groupIdx) +{ + // copy weight scale [1, 2H] offset:0 + DataCopyPadParams padParams{false, 0, 0, 0}; + LocalTensor weightScaleLocal = weightScaleQueue_.AllocTensor(); + DataCopyParams dataCopyWeightScaleParams; + dataCopyWeightScaleParams.blockCount = 1; + dataCopyWeightScaleParams.blockLen = tl_->inDimy * sizeof(float); + dataCopyWeightScaleParams.srcStride = 0; + dataCopyWeightScaleParams.dstStride = 0; + if constexpr (std::is_same_v) { + DataCopyPad(weightScaleLocal, weightScaleGm_[groupIdx * tl_->inDimy], dataCopyWeightScaleParams, padParams); + } + weightScaleQueue_.EnQue(weightScaleLocal); +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::CopyInQuantScale(int64_t groupIdx) +{ + DataCopyPadParams padParams{false, 0, 0, 0}; + // copy static quant scale + LocalTensor inScaleLocal = inScaleQueue_.AllocTensor(); + if (tl_->quantIsOne) { + if constexpr (IsSameType::value) { + this->quantScale_ = 1 / ToFloat(this->quantScaleGm_.GetValue(groupIdx)); + this->quantOffset_ = ToFloat(this->quantOffsetGm_.GetValue(groupIdx)); + } else if constexpr (IsSameType::value) { + this->quantScale_ = 1 / static_cast(this->quantScaleGm_.GetValue(groupIdx)); + this->quantOffset_ = static_cast(this->quantOffsetGm_.GetValue(groupIdx)); + } else { + this->quantScale_ = 1 / this->quantScaleGm_.GetValue(groupIdx); + this->quantOffset_ = this->quantOffsetGm_.GetValue(groupIdx); + } + } + + // copy dynamic quant scale [1, H] offset:tl_->inDimy + if (tl_->needSmoothScale == 1 && !tl_->quantIsOne) { + DataCopyParams dataCopyQuantScaleParams; + dataCopyQuantScaleParams.blockCount = 1; + dataCopyQuantScaleParams.blockLen = tl_->outDimy * sizeof(TQuantScale); + dataCopyQuantScaleParams.srcStride = 0; + dataCopyQuantScaleParams.dstStride = 0; + if constexpr (std::is_same_v) { + DataCopyPad(inScaleLocal, quantScaleGm_[groupIdx * tl_->outDimy], dataCopyQuantScaleParams, padParams); + if (tl_->quantMode == 0) { + DataCopyPad( + inScaleLocal[tl_->outDimy], quantOffsetGm_[groupIdx * tl_->outDimy], dataCopyQuantScaleParams, + padParams); + } + + } else { + LocalTensor quantScaleLocalT16 = inScaleLocal.template ReinterpretCast(); + DataCopyPad( + quantScaleLocalT16[tl_->outDimy], quantScaleGm_[groupIdx * tl_->outDimy], dataCopyQuantScaleParams, + padParams); + if (tl_->quantMode == 0) { + DataCopyPad( + quantScaleLocalT16[tl_->outDimy + tl_->inDimy], quantOffsetGm_[groupIdx * tl_->outDimy], + dataCopyQuantScaleParams, padParams); + } + } + } + inScaleQueue_.EnQue(inScaleLocal); +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::CopyInBias(int64_t groupIdx) +{ + DataCopyPadParams padParams{false, 0, 0, 0}; + if constexpr (std::is_same_v) { + if (tl_->hasBias == 1) { + biasLocal_ = biasQueue_.AllocTensor(); + DataCopyParams dataCopyBiasParams; + dataCopyBiasParams.blockCount = 1; + dataCopyBiasParams.blockLen = tl_->inDimy * sizeof(TBias); + dataCopyBiasParams.srcStride = 0; + dataCopyBiasParams.dstStride = 0; + if constexpr (std::is_same_v || std::is_same_v) { + DataCopyPad(biasLocal_, biasGm_[groupIdx * tl_->inDimy], dataCopyBiasParams, padParams); + } else { + DataCopyPad(biasLocal_[tl_->inDimy], biasGm_[groupIdx * tl_->inDimy], dataCopyBiasParams, padParams); + } + biasQueue_.EnQue(biasLocal_); + } + } +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::CopyInXAct(int32_t proDimsx, int64_t xDimxOffset) +{ + // copyin x and Act scale + DataCopyPadParams padParams{false, 0, 0, 0}; + LocalTensor xActLocal = xActQueue_.AllocTensor(); + DataCopyParams dataCopyXParams; + dataCopyXParams.blockCount = proDimsx; + dataCopyXParams.blockLen = tl_->inDimy * sizeof(TXGm); + dataCopyXParams.srcStride = 0; + dataCopyXParams.dstStride = 0; + DataCopyPad(xActLocal[TBufXLocalInOfs_], xGm_[xDimxOffset * tl_->inDimy], dataCopyXParams, padParams); + + // copy act scale: [proDimsx,8] offset:tl_->UbFactorDimx * tl_->inDimy = TBufActSclInOfs_ + DataCopyParams dataCopyActScaleParams; + dataCopyActScaleParams.blockCount = proDimsx; + dataCopyActScaleParams.blockLen = sizeof(float); + dataCopyActScaleParams.srcStride = 0; + dataCopyActScaleParams.dstStride = 0; + LocalTensor xActLocalF32 = xActLocal.template ReinterpretCast(); + if (std::is_same_v && !tl_->activationScaleIsEmpty) { + DataCopyPad(xActLocalF32[TBufActSclInOfs_], activationScaleGm_[xDimxOffset], dataCopyActScaleParams, padParams); + } + xActQueue_.EnQue(xActLocal); +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::ComputeDequant(int32_t proDimsx) +{ + LocalTensor xActLocal = xActQueue_.DeQue(); + LocalTensor xActLocalF32 = xActLocal.template ReinterpretCast(); + LocalTensor xLocalF32 = xActLocalF32; + LocalTensor activationScaleLocal = xActLocalF32[TBufActSclInOfs_]; + LocalTensor tmpUbF32 = tmpBuf1_.AllocTensor(); // weight scale FP32 + LocalTensor tmpUbI32 = tmpUbF32.template ReinterpretCast(); + + if constexpr (std::is_same_v) { + if constexpr (std::is_same_v){ + // Copy bias: [1,2H] -> [proDimsx,2H] + // params: dstStride: 1, srcStride: 1, dstRepStride: tl_->UbFactorDimy * 2 / 8, srcRepStride: 0 + CopyReshape(tmpUbI32, biasLocal_, proDimsx, tl_->UbFactorDimy * SWI_FACTOR, + {1, 1, static_cast((tl_->UbFactorDimy * SWI_FACTOR) / BLOCK_ELEM), 0}); + PipeBarrier(); + Add(xActLocal, xActLocal, tmpUbI32, proDimsx * tl_->inDimy); + PipeBarrier(); + } + + // Copy weight scale: [1,2H] -> [proDimsx,2H] + // params: dstStride: 1, srcStride: 1, dstRepStride: tl_->UbFactorDimy * 2 / 8, srcRepStride: 0 + CopyReshape(tmpUbF32, weightScaleLocal_, proDimsx, tl_->UbFactorDimy * SWI_FACTOR, + {1, 1, static_cast((tl_->UbFactorDimy * SWI_FACTOR) / BLOCK_ELEM), 0}); + } + + // x 为 bf16时 + Cast(xLocalF32, xActLocal[TBufXLocalInOfs_], RoundMode::CAST_NONE, SWI_FACTOR * proDimsx * tl_->UbFactorDimy); + PipeBarrier(); + if constexpr (std::is_same_v) { + // Calc dequant: xLocalF32 = weightScaleLocal * xLocalF32 + Mul(xLocalF32, tmpUbF32, xLocalF32, tl_->UbFactorDimy * SWI_FACTOR * proDimsx); + PipeBarrier(); + if (!tl_->activationScaleIsEmpty) { + // Copy act scale: [proDimsx,8] -> [proDimsx,2H] + CopyReshape(tmpUbF32, activationScaleLocal, proDimsx, tl_->UbFactorDimy * SWI_FACTOR, + {1, 0, static_cast((tl_->UbFactorDimy * SWI_FACTOR) / BLOCK_ELEM), 1}); + PipeBarrier(); + // Calc dequant: xLocalF32 = activationScaleLocal * xLocalF32 + Mul(xLocalF32, tmpUbF32, xLocalF32, tl_->UbFactorDimy * SWI_FACTOR * proDimsx); + PipeBarrier(); + } + } + + if constexpr (std::is_same_v && !std::is_same_v) { + if (tl_->hasBias == 1) { + // Copy bias: [1,2H] -> [proDimsx,2H] + CopyReshape(tmpUbF32, biasLocalF32_, proDimsx, tl_->UbFactorDimy * SWI_FACTOR, + {1, 1, static_cast((tl_->UbFactorDimy * SWI_FACTOR) / BLOCK_ELEM), 0}); + PipeBarrier(); + Add(xLocalF32, xLocalF32, tmpUbF32, proDimsx * tl_->inDimy); + PipeBarrier(); + } + } + xActQueue_.EnQue(xLocalF32); +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::ComputeSwiGLU(int32_t proDimsx) +{ + LocalTensor xLocalF32 = xActQueue_.DeQue(); + if (tl_->swigluMode == 1) { + // do special swiglu + SwiGluGate(proDimsx, xLocalF32); + } else { + uint32_t calEleNum = tl_->UbFactorDimy * proDimsx; + LocalTensor tmpUbF32 = tmpBuf1_.AllocTensor(); + // do normal swi pre + LocalTensor tmpUbF32Act = tmpUbF32; + LocalTensor tmpUbF32Gate = tmpUbF32[calEleNum]; + // Copy dequant result: xLocalF32[actOffset] -> tmpUbF32Act, [proDimsx,H] + // Copy dequant result: xLocalF32[gateOffset] -> tmpUbF32Gate, [proDimsx,H] + SetMaskCount(); + SetVectorMask(tl_->UbFactorDimy); + Copy( + tmpUbF32Act, xLocalF32[actOffset_], AscendC::MASK_PLACEHOLDER, proDimsx, + {1, 1, static_cast(tl_->UbFactorDimy / BLOCK_ELEM), + static_cast(tl_->UbFactorDimy / BLOCK_ELEM * SWI_FACTOR)}); + Copy( + tmpUbF32Gate, xLocalF32[gateOffset_], AscendC::MASK_PLACEHOLDER, proDimsx, + {1, 1, static_cast(tl_->UbFactorDimy / BLOCK_ELEM), + static_cast(tl_->UbFactorDimy / BLOCK_ELEM * SWI_FACTOR)}); + SetMaskNorm(); + ResetMask(); + PipeBarrier(); + Muls(xLocalF32, tmpUbF32Act, static_cast(-1.0), calEleNum); + PipeBarrier(); + Exp(xLocalF32, xLocalF32, calEleNum); + PipeBarrier(); + Adds(xLocalF32, xLocalF32, static_cast(1.0), calEleNum); + PipeBarrier(); + Div(tmpUbF32Act, tmpUbF32Act, xLocalF32, calEleNum); + PipeBarrier(); + Mul(tmpUbF32Act, tmpUbF32Gate, tmpUbF32Act, calEleNum); + PipeBarrier(); + } + // x compute done, free + xActQueue_.FreeTensor(xLocalF32); +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::ComputeQuant(int32_t proDimsx) +{ + LocalTensor tmpUbF32 = tmpBuf1_.AllocTensor(); + LocalTensor tmpUbF32Act = tmpUbF32; + LocalTensor tmpUbF32Gate = tmpUbF32[tl_->UbFactorDimy * proDimsx]; + if (tl_->quantMode == 1) { + DynamicQuant(tmpUbF32Act, tmpUbF32Gate, inScaleLocal_, proDimsx); + } else { + StaticQuant(tmpUbF32Act, tmpUbF32Gate, inScaleLocal_, proDimsx); + } + tmpBuf1_.FreeTensor(tmpUbF32); +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::Compute(int32_t proDimsx) +{ + ComputeDequant(proDimsx); + ComputeSwiGLU(proDimsx); + ComputeQuant(proDimsx); +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::CopyOut(int32_t proDimsx, int64_t xDimxOffset) +{ + // copy out + LocalTensor outLocal = outQueue_.DeQue(); + LocalTensor scaleOut = outLocal[UbSingleOutSize_ * sizeof(int8_t) / sizeof(float)]; + LocalTensor yOut = outLocal.template ReinterpretCast(); + + if (tl_->quantMode == 1) { + DataCopyParams dataCopyOutScaleParams; + dataCopyOutScaleParams.blockCount = 1; + dataCopyOutScaleParams.blockLen = proDimsx * sizeof(float); + dataCopyOutScaleParams.srcStride = 0; + dataCopyOutScaleParams.dstStride = 0; + DataCopyPad(scaleGm_[xDimxOffset], scaleOut, dataCopyOutScaleParams); + } + DataCopyParams dataCopyOutyParams; + dataCopyOutyParams.blockCount = 1; + dataCopyOutyParams.blockLen = proDimsx * tl_->outDimy * sizeof(int8_t); + dataCopyOutyParams.srcStride = 0; + dataCopyOutyParams.dstStride = 0; + DataCopyPad(yGm_[xDimxOffset * tl_->outDimy], yOut, dataCopyOutyParams); + outQueue_.FreeTensor(outLocal); +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::ParamDequeAndCast() +{ + weightScaleLocal_ = weightScaleQueue_.DeQue(); + + // bias deque and cast bias to fp32 if needed + if constexpr (std::is_same_v) { + if (tl_->hasBias == 1) { + biasLocal_ = biasQueue_.DeQue(); + biasLocalF32_ = biasLocal_.template ReinterpretCast(); + if constexpr (std::is_same_v || std::is_same_v) { + Cast(biasLocalF32_, biasLocal_[tl_->inDimy], RoundMode::CAST_NONE, tl_->inDimy); + } + } + } + + // quant scale and quant offset deque, cast them to fp32 if needed + inScaleLocal_ = inScaleQueue_.DeQue(); + if (tl_->needSmoothScale == 1 && !tl_->quantIsOne) { + if (std::is_same_v || std::is_same_v) { + LocalTensor quantScaleLocalT16 = inScaleLocal_.template ReinterpretCast(); + Cast(inScaleLocal_, quantScaleLocalT16[tl_->outDimy], RoundMode::CAST_NONE, tl_->outDimy); + PipeBarrier(); + if (tl_->quantMode == 0) { + Cast( + inScaleLocal_[tl_->outDimy], quantScaleLocalT16[tl_->outDimy + tl_->inDimy], RoundMode::CAST_NONE, + tl_->outDimy); + PipeBarrier(); + } + } + } +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::ProcessSingleGroupPerCore( + int64_t groupIdx, int64_t dimxCore, int64_t coreDimxOffset) +{ + // do ub tiling again + int32_t ubDimxLoop = (dimxCore + tl_->UbFactorDimx - 1) / tl_->UbFactorDimx; + int32_t ubDimxTailFactor = dimxCore - tl_->UbFactorDimx * (ubDimxLoop - 1); + + // copyin 当前分组下使用的参数，weight scale, bias scale, quant scale+quant offset + CopyInWeightScale(groupIdx); + CopyInQuantScale(groupIdx); + CopyInBias(groupIdx); + ParamDequeAndCast(); + + /* + 1. copyin x, activation scale + 2. compute + 3. copyout y, scale + */ + for (uint32_t loopIdx = 0; loopIdx < ubDimxLoop; ++loopIdx) { + int64_t xDimxOffset = coreDimxOffset + loopIdx * tl_->UbFactorDimx; + int32_t proDimsx = loopIdx == (ubDimxLoop - 1) ? ubDimxTailFactor : tl_->UbFactorDimx; + CopyInXAct(proDimsx, xDimxOffset); + Compute(proDimsx); + CopyOut(proDimsx, xDimxOffset); + } + ParamFree(); +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::ParamFree() +{ + // 释放当前分组下使用的参数，weight scale, bias scale, quant scale，quant offset + inScaleQueue_.FreeTensor(inScaleLocal_); + weightScaleQueue_.FreeTensor(weightScaleLocal_); + if constexpr (std::is_same_v) { + if (tl_->hasBias == 1) { + biasQueue_.FreeTensor(biasLocal_); + } + } +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::ComputeReduceMax( + const LocalTensor& tempRes, int32_t calCount) +{ + uint32_t vectorCycles = calCount / MASK_NUM_T32; + uint32_t remainElements = calCount % MASK_NUM_T32; + + BinaryRepeatParams repeatParams; + repeatParams.dstBlkStride = 1; + repeatParams.src0BlkStride = 1; + repeatParams.src1BlkStride = 1; + repeatParams.dstRepStride = 0; + repeatParams.src0RepStride = MASK_BLK_STRIDE; + repeatParams.src1RepStride = 0; + + if (vectorCycles > 0 && remainElements > 0) { + Max(tempRes, tempRes, tempRes[vectorCycles * MASK_NUM_T32], remainElements, 1, repeatParams); + PipeBarrier(); + } + + if (vectorCycles > 1) { + Max(tempRes, tempRes[MASK_NUM_T32], tempRes, MASK_NUM_T32, vectorCycles - 1, repeatParams); + PipeBarrier(); + } +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::CreateOffsetLocalTensor( + uint32_t tensorLen, int swigluMode) +{ + // 不再需要创建偏移张量，因为直接使用前一半和后一半数据 +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::SwiGluGate( + int32_t proDimsx, const LocalTensor& xLocalF32) +{ + uint32_t calEleNum = tl_->UbFactorDimy * proDimsx; + LocalTensor tmpUbF32 = tmpBuf1_.AllocTensor(); + LocalTensor tmpUbF32Act = tmpUbF32; + LocalTensor tmpUbF32Gate = tmpUbF32[calEleNum]; + SetMaskCount(); + SetVectorMask(tl_->UbFactorDimy); + Copy( + tmpUbF32Act, xLocalF32[actOffset_], AscendC::MASK_PLACEHOLDER, proDimsx, + {1, 1, static_cast(tl_->UbFactorDimy / BLOCK_ELEM), + static_cast(tl_->UbFactorDimy / BLOCK_ELEM * SWI_FACTOR)}); + Copy( + tmpUbF32Gate, xLocalF32[gateOffset_], AscendC::MASK_PLACEHOLDER, proDimsx, + {1, 1, static_cast(tl_->UbFactorDimy / BLOCK_ELEM), + static_cast(tl_->UbFactorDimy / BLOCK_ELEM * SWI_FACTOR)}); + SetMaskNorm(); + ResetMask(); + PipeBarrier(); + if (tl_->clampLimit > 0.0f) { + // tmpUbF32Gate + Mins(tmpUbF32Gate, tmpUbF32Gate, tl_->clampLimit, calEleNum); + PipeBarrier(); + Maxs(tmpUbF32Gate, tmpUbF32Gate, -(tl_->clampLimit), calEleNum); + PipeBarrier(); + } + Adds(tmpUbF32Gate, tmpUbF32Gate, tl_->gluBias, calEleNum); + PipeBarrier(); + if (tl_->clampLimit > 0.0f) { + // tmpUbF32Act + Mins(tmpUbF32Act, tmpUbF32Act, tl_->clampLimit, calEleNum); + PipeBarrier(); + } + Muls(xLocalF32, tmpUbF32Act, -(tl_->gluAlpha), calEleNum); + PipeBarrier(); + Exp(xLocalF32, xLocalF32, calEleNum); + PipeBarrier(); + Adds(xLocalF32, xLocalF32, static_cast(1.0), calEleNum); + PipeBarrier(); + Div(tmpUbF32Act, tmpUbF32Act, xLocalF32, calEleNum); + PipeBarrier(); + Mul(tmpUbF32Act, tmpUbF32Gate, tmpUbF32Act, calEleNum); + PipeBarrier(); +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::DynamicQuant( + const LocalTensor& tmpUbF32Act, const LocalTensor& tmpUbF32Gate, + const LocalTensor& inScaleLocal, uint32_t proDimsx) +{ + if (tl_->needSmoothScale == 1) { + // Copy quant scale: [1,H] -> [proDimsx,H] + SetMaskCount(); + SetVectorMask(tl_->UbFactorDimy); + Copy( + tmpUbF32Gate, inScaleLocal, AscendC::MASK_PLACEHOLDER, proDimsx, + {1, 1, static_cast(tl_->UbFactorDimy / BLOCK_ELEM), 0}); + SetMaskNorm(); + ResetMask(); + PipeBarrier(); + // Calc quant: xLocalF32 = tmpUbF32Act * inScaleLocal + Mul(tmpUbF32Act, tmpUbF32Gate, tmpUbF32Act, tl_->UbFactorDimy * proDimsx); + PipeBarrier(); + } + + // Calc quant: tmpUbF32Gate = abs(tmpUbF32Act) + Abs(tmpUbF32Gate, tmpUbF32Act, tl_->UbFactorDimy * proDimsx); + + LocalTensor outLocal = outQueue_.AllocTensor(); + LocalTensor scaleOut = outLocal[UbSingleOutSize_ * sizeof(int8_t) / sizeof(float)]; + LocalTensor yOut = outLocal.template ReinterpretCast(); + PipeBarrier(); + // Calc quant: proDimsx * tl_->UbFactorDimy -> proDimsx * 64 + for (uint32_t i = 0; i < proDimsx; ++i) { + ComputeReduceMax(tmpUbF32Gate[i * tl_->UbFactorDimy], tl_->UbFactorDimy); + } + // Calc quant: proDimsx * 64 -> proDimsx + // repeatTimes:proDimsx, dstRepStride:1(dtype), srcBlkStride:1, srcRepStride:tl_->UbFactorDimy / 64 * 8 + WholeReduceMax( + tmpUbF32Gate, tmpUbF32Gate, MASK_NUM_T32, proDimsx, 1, 1, tl_->UbFactorDimy / MASK_NUM_T32 * MASK_BLK_STRIDE, + ReduceOrder::ORDER_ONLY_VALUE); + PipeBarrier(); + // Calc quant: scaleOut / 127.0 + Muls(scaleOut, tmpUbF32Gate, DYNAMIC_QUANT_FACTOR, proDimsx); + PipeBarrier(); + // Calc Broadcast: proDimsx -> proDimsx,8 + int64_t blockCount = (proDimsx + BLOCK_ELEM - 1) / BLOCK_ELEM; + Brcb(outLocal, scaleOut, blockCount, {1, MASK_BLK_STRIDE}); + PipeBarrier(); + // Copy scale: [proDimsx,8] -> [proDimsx,H] + SetMaskCount(); + SetVectorMask(tl_->UbFactorDimy); + Copy( + tmpUbF32Gate, outLocal, AscendC::MASK_PLACEHOLDER, proDimsx, + {1, 0, static_cast(tl_->UbFactorDimy / BLOCK_ELEM), 1}); + SetMaskNorm(); + ResetMask(); + PipeBarrier(); + // Calc y: tmpUbF32Act = tmpUbF32Act / scaleOut + Div(tmpUbF32Act, tmpUbF32Act, tmpUbF32Gate, tl_->UbFactorDimy * proDimsx); + PipeBarrier(); + + CastFloatToInt8(tmpUbF32Act, tmpUbF32Gate, proDimsx, yOut); + outQueue_.EnQue(outLocal); +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::StaticQuant( + const LocalTensor& tmpUbF32Act, const LocalTensor& tmpUbF32Gate, + const LocalTensor& inScaleLocal, uint32_t proDimsx) +{ + if (tl_->needSmoothScale == 1) { + if (tl_->quantIsOne) { + // Calc quant: y = tmpUbF32Act * quantScale + quantOffset + Muls(tmpUbF32Gate, tmpUbF32Act, this->quantScale_, tl_->UbFactorDimy * proDimsx); + PipeBarrier(); + Adds(tmpUbF32Act, tmpUbF32Gate, this->quantOffset_, tl_->UbFactorDimy * proDimsx); + PipeBarrier(); + } else { + // Copy quant scale: [1,H] -> [proDimsx,H] + SetMaskCount(); + SetVectorMask(tl_->UbFactorDimy); + Copy( + tmpUbF32Gate, inScaleLocal, AscendC::MASK_PLACEHOLDER, proDimsx, + {1, 1, static_cast(tl_->UbFactorDimy / BLOCK_ELEM), 0}); + SetMaskNorm(); + ResetMask(); + PipeBarrier(); + // Calc quant: y = tmpUbF32Act / quantScale + Div(tmpUbF32Act, tmpUbF32Act, tmpUbF32Gate, tl_->UbFactorDimy * proDimsx); + PipeBarrier(); + + // Copy quant offset: [1,H] -> [proDimsx,H] + SetMaskCount(); + SetVectorMask(tl_->UbFactorDimy); + Copy( + tmpUbF32Gate, inScaleLocal[tl_->outDimy], AscendC::MASK_PLACEHOLDER, proDimsx, + {1, 1, static_cast(tl_->UbFactorDimy / BLOCK_ELEM), 0}); + SetMaskNorm(); + ResetMask(); + PipeBarrier(); + // Calc quant: y = tmpUbF32Act + quantOffset + Add(tmpUbF32Act, tmpUbF32Act, tmpUbF32Gate, tl_->UbFactorDimy * proDimsx); + PipeBarrier(); + } + } + + // do cast float to int8 + LocalTensor outLocal = outQueue_.AllocTensor(); + LocalTensor yOut = outLocal.template ReinterpretCast(); + + CastFloatToInt8(tmpUbF32Act, tmpUbF32Gate, proDimsx, yOut); + outQueue_.EnQue(outLocal); +} + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantBase::CastFloatToInt8( + const LocalTensor& tmpUbF32Act, const LocalTensor& tmpUbF32Gate, uint32_t proDimsx, LocalTensor& yOut) +{ + LocalTensor tmpUbF32ActI32 = tmpUbF32Act.ReinterpretCast(); + Cast(tmpUbF32ActI32, tmpUbF32Act, RoundMode::CAST_RINT, tl_->UbFactorDimy * proDimsx); + SetDeqScale((half)1.000000e+00f); + + LocalTensor tmpUbF32Gate16 = tmpUbF32Gate.template ReinterpretCast(); + Cast(tmpUbF32Gate16, tmpUbF32ActI32, RoundMode::CAST_ROUND, tl_->UbFactorDimy * proDimsx); + PipeBarrier(); + + Cast(yOut, tmpUbF32Gate16, RoundMode::CAST_TRUNC, tl_->UbFactorDimy * proDimsx); + PipeBarrier(); +} + +TEMPLATE_DSQ_DECLARE +template +__aicore__ inline void DequantSwigluQuantBase::CopyReshape( + LocalTensor& dstTensor, LocalTensor& oriTensor, uint32_t rowNum, uint32_t colNum, CopyRepeatParams param) +{ + SetMaskCount(); + SetVectorMask(colNum); + Copy(dstTensor, oriTensor, AscendC::MASK_PLACEHOLDER, rowNum, param); + SetMaskNorm(); + ResetMask(); +} + +} // namespace DequantSwigluQuantOps +#endif // DEQUANT_SWIGLU_QUANT_H diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_apt.cpp b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_apt.cpp new file mode 100644 index 000000000..b0489b3df --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_apt.cpp @@ -0,0 +1,1411 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_apt.cpp + * \brief + */ +#if ASC_DEVKIT_MAJOR >= 9 + #include "basic_api/kernel_vec_intf.h" +#else + #include "kernel_operator.h" +#endif +#include "arch35/dequant_swiglu_quant.h" +#include "arch35/dequant_swiglu_quant_nlast.h" +#include "arch35/dequant_swiglu_quant_static.h" +#include "arch35/dequant_swiglu_quant_dynamic_not_full.h" +#include "arch35/dequant_swiglu_quant_static_not_full.h" + +#define ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE 111 +#define ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE 110 +#define ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE 101 +#define ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE 100 +#define ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE 11 +#define ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE 10 +#define ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE 1 +#define ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE 0 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE 1111 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE 2111 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE 3111 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE 4111 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE 1110 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE 2110 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE 3110 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE 4110 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE 1101 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE 2101 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE 3101 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE 4101 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE 1100 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE 2100 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE 3100 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE 4100 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE 1011 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE 2011 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE 3011 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE 4011 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE 1010 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE 2010 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE 3010 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE 4010 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE 1001 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE 2001 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE 3001 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE 4001 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE 1000 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE 2000 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE 3000 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE 4000 // bias = bfloat16 +#define ACTDIM_TRUE_X_FALSE_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE 100110 +#define ACTDIM_TRUE_X_FALSE_BIAS_FALSE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE 100100 +#define ACTDIM_TRUE_X_FALSE_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE 100010 +#define ACTDIM_TRUE_X_FALSE_BIAS_FALSE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE 100000 +#define ACTDIM_TRUE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE 101110 // bias = int32 +#define ACTDIM_TRUE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE 102110 // bias = float32 +#define ACTDIM_TRUE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE 103110 // bias = float16 +#define ACTDIM_TRUE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE 104110 // bias = bfloat16 +#define ACTDIM_TRUE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE 101100 // bias = int32 +#define ACTDIM_TRUE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE 102100 // bias = float32 +#define ACTDIM_TRUE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE 103100 // bias = float16 +#define ACTDIM_TRUE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE 104100 // bias = bfloat16 +#define ACTDIM_TRUE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE 101010 // bias = int32 +#define ACTDIM_TRUE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE 102010 // bias = float32 +#define ACTDIM_TRUE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE 103010 // bias = float16 +#define ACTDIM_TRUE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE 104010 // bias = bfloat16 +#define ACTDIM_TRUE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE 101000 // bias = int32 +#define ACTDIM_TRUE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE 102000 // bias = float32 +#define ACTDIM_TRUE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE 103000 // bias = float16 +#define ACTDIM_TRUE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE 104000 // bias = bfloat16 +// static quant not full load +#define QUANTMODE_STATIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1000100 +#define QUANTMODE_STATIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_TRUE_GROUP_FALSE 1000110 +#define QUANTMODE_STATIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1001100 +#define QUANTMODE_STATIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_TRUE_GROUP_FALSE 1001110 +#define QUANTMODE_STATIC_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1010100 +#define QUANTMODE_STATIC_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_OFFSET_TRUE_GROUP_FALSE 1010110 +#define QUANTMODE_STATIC_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1011100 +#define QUANTMODE_STATIC_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_OFFSET_TRUE_GROUP_FALSE 1011110 +#define QUANTMODE_STATIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1000101 +#define QUANTMODE_STATIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_TRUE_GROUP_TRUE 1000111 +#define QUANTMODE_STATIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1001101 +#define QUANTMODE_STATIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_TRUE_GROUP_TRUE 1001111 +#define QUANTMODE_STATIC_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1010101 +#define QUANTMODE_STATIC_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_OFFSET_TRUE_GROUP_TRUE 1010111 +#define QUANTMODE_STATIC_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1011101 +#define QUANTMODE_STATIC_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_OFFSET_TRUE_GROUP_TRUE 1011111 + +#define ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC 10111 +#define ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC 10110 +#define ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC 10101 +#define ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC 10100 +#define ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC 10011 +#define ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC 10010 +#define ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC 10001 +#define ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC 10000 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC 11111 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC 12111 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC 13111 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC 14111 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC 11110 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC 12110 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC 13110 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC 14110 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC 11101 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC 12101 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC 13101 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC 14101 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC 11100 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC 12100 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC 13100 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC 14100 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC 11011 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC 12011 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC 13011 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC 14011 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC 11010 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC 12010 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC 13010 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC 14010 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC 11001 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC 12001 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC 13001 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC 14001 // bias = bfloat16 +#define ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC 11000 // bias = int32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC 12000 // bias = float32 +#define ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC 13000 // bias = float16 +#define ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC 14000 // bias = bfloat16 +// ##Dynamic Not Full Load +#define QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE 1100000 +#define QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE 1100001 +#define QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1100100 +#define QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1100101 +#define QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE 1101000 +#define QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE 1101001 +#define QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1101100 +#define QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1101101 +#define QUANTMODE_DYNAMIC_BIAS_INT32_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE 1110000 +#define QUANTMODE_DYNAMIC_BIAS_INT32_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE 1110001 +#define QUANTMODE_DYNAMIC_BIAS_INT32_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1110100 +#define QUANTMODE_DYNAMIC_BIAS_INT32_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1110101 +#define QUANTMODE_DYNAMIC_BIAS_INT32_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE 1111000 +#define QUANTMODE_DYNAMIC_BIAS_INT32_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE 1111001 +#define QUANTMODE_DYNAMIC_BIAS_INT32_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1111100 +#define QUANTMODE_DYNAMIC_BIAS_INT32_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1111101 +#define QUANTMODE_DYNAMIC_BIAS_FP32_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE 1120000 +#define QUANTMODE_DYNAMIC_BIAS_FP32_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE 1120001 +#define QUANTMODE_DYNAMIC_BIAS_FP32_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1120100 +#define QUANTMODE_DYNAMIC_BIAS_FP32_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1120101 +#define QUANTMODE_DYNAMIC_BIAS_FP32_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE 1121000 +#define QUANTMODE_DYNAMIC_BIAS_FP32_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE 1121001 +#define QUANTMODE_DYNAMIC_BIAS_FP32_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1121100 +#define QUANTMODE_DYNAMIC_BIAS_FP32_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1121101 +#define QUANTMODE_DYNAMIC_BIAS_FP16_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE 1130000 +#define QUANTMODE_DYNAMIC_BIAS_FP16_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE 1130001 +#define QUANTMODE_DYNAMIC_BIAS_FP16_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1130100 +#define QUANTMODE_DYNAMIC_BIAS_FP16_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1130101 +#define QUANTMODE_DYNAMIC_BIAS_FP16_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE 1131000 +#define QUANTMODE_DYNAMIC_BIAS_FP16_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE 1131001 +#define QUANTMODE_DYNAMIC_BIAS_FP16_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1131100 +#define QUANTMODE_DYNAMIC_BIAS_FP16_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1131101 +#define QUANTMODE_DYNAMIC_BIAS_BF16_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE 1140000 +#define QUANTMODE_DYNAMIC_BIAS_BF16_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE 1140001 +#define QUANTMODE_DYNAMIC_BIAS_BF16_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1140100 +#define QUANTMODE_DYNAMIC_BIAS_BF16_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1140101 +#define QUANTMODE_DYNAMIC_BIAS_BF16_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE 1141000 +#define QUANTMODE_DYNAMIC_BIAS_BF16_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE 1141001 +#define QUANTMODE_DYNAMIC_BIAS_BF16_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE 1141100 +#define QUANTMODE_DYNAMIC_BIAS_BF16_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE 1141101 + + +using namespace AscendC; + +extern "C" __global__ __aicore__ void dequant_swiglu_quant(GM_ADDR x, GM_ADDR weightScale, GM_ADDR activationScale, + GM_ADDR bias, GM_ADDR quantScale, GM_ADDR quantOffset, + GM_ADDR groupIndex, GM_ADDR y, GM_ADDR scale, + GM_ADDR workspace, GM_ADDR tiling) { + if (g_coreType != AIV) { + return; + } + + GM_ADDR globalWorkGm = GetUserWorkspace(workspace); + TPipe pipe; + TILING_KEY_IS(ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_FALSE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_FALSE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE); + TILING_KEY_IS(ACTDIM_TRUE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE); + // static quant not full load + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_TRUE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_TRUE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_OFFSET_TRUE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_OFFSET_TRUE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_TRUE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_TRUE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_OFFSET_TRUE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_STATIC_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_OFFSET_TRUE_GROUP_TRUE); + + TILING_KEY_IS(ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC); + TILING_KEY_IS(ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC); + // ##Dynamic Not Full Load + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_INT32_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_INT32_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_INT32_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_INT32_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_INT32_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_INT32_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_INT32_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_INT32_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP32_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP32_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP32_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP32_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP32_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP32_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP32_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP32_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP16_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP16_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP16_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP16_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP16_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP16_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP16_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_FP16_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_BF16_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_BF16_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_BF16_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_BF16_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_BF16_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_BF16_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_BF16_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE); + TILING_KEY_IS(QUANTMODE_DYNAMIC_BIAS_BF16_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE); + #if TILING_KEY_VAR == ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale有值, quant_scale有值, group_index有值, 000111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale有值, quant_scale有值, group_index无值, 000110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale有值, quant_scale无值, group_index有值, 000101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale有值, quant_scale无值, group_index无值, 000100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale无值, quant_scale有值, group_index有值, 000011 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale无值, quant_scale有值, group_index无值, 000010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale无值, quant_scale无值, group_index有值, 000001 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale无值, quant_scale无值, group_index无值, 000000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale有值, quant_scale有值, group_index有值, 001111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale有值, quant_scale有值, group_index有值, 002111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale有值, quant_scale有值, group_index有值, 003111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale有值, quant_scale有值, group_index有值, 004111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale有值, quant_scale有值, group_index无值, 001110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale有值, quant_scale有值, group_index无值, 002110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale有值, quant_scale有值, group_index无值, 003110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale有值, quant_scale有值, group_index无值, 004110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale有值, quant_scale无值, group_index有值, 001101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale有值, quant_scale无值, group_index有值, 002101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale有值, quant_scale无值, group_index有值, 003101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale有值, quant_scale无值, group_index有值, 004101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale有值, quant_scale无值, group_index无值, 001100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale有值, quant_scale无值, group_index无值, 002100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale有值, quant_scale无值, group_index无值, 003100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale有值, quant_scale无值, group_index无值, 004100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale无值, quant_scale有值, group_index有值, 001011 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale无值, quant_scale有值, group_index有值, 002011 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale无值, quant_scale有值, group_index有值, 003011 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale无值, quant_scale有值, group_index有值, 004011 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale无值, quant_scale有值, group_index无值, 001010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale无值, quant_scale有值, group_index无值, 002010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale无值, quant_scale有值, group_index无值, 003010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale无值, quant_scale有值, group_index无值, 004010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale无值, quant_scale无值, group_index有值, 001001 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale无值, quant_scale无值, group_index有值, 002001 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale无值, quant_scale无值, group_index有值, 003001 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale无值, quant_scale无值, group_index有值, 004001 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale无值, quant_scale无值, group_index无值, 001000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale无值, quant_scale无值, group_index无值, 002000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale无值, quant_scale无值, group_index无值, 003000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale无值, quant_scale无值, group_index无值, 004000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBase op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=int32, activate_scale有值, quant_scale有值, group_index无值, 101110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=float32, activate_scale有值, quant_scale有值, group_index无值, 102110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=float16, activate_scale有值, quant_scale有值, group_index无值, 103110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale有值, quant_scale有值, group_index无值, 104110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=int32, activate_scale有值, quant_scale无值, group_index无值, 101100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=float32, activate_scale有值, quant_scale无值, group_index无值, 102100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=flaot16, activate_scale有值, quant_scale无值, group_index无值, 103100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale有值, quant_scale无值, group_index无值, 104100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=int32, activate_scale无值, quant_scale有值, group_index无值, 101010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=float32, activate_scale无值, quant_scale有值, group_index无值, 102010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=float16, activate_scale无值, quant_scale有值, group_index无值, 103010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale无值, quant_scale有值, group_index无值, 104010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=int32, activate_scale无值, quant_scale无值, group_index无值, 101000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=flaot32, activate_scale无值, quant_scale无值, group_index无值, 102000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=float16, activate_scale无值, quant_scale无值, group_index无值, 103000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale无值, quant_scale无值, group_index无值, 104000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_FALSE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias无值, activate_scale无值, quant_scale无值, group_index无值, 100000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias无值, activate_scale无值, quant_scale有值, group_index无值, 100010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_FALSE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias无值, activate_scale有值, quant_scale无值, group_index无值, 100100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_TRUE_X_FALSE_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE + // activate_dim!=-1, x=int32/fp16/bf6, bias无值, activate_scale有值, quant_scale有值, group_index无值, 100110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35NlastTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35NlastTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantNlast op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale有值, quant_scale有值, group_index有值, 000111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale有值, quant_scale有值, group_index无值, 000110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale有值, quant_scale无值, group_index有值, 000101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale有值, quant_scale无值, group_index无值, 000100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale无值, quant_scale有值, group_index有值, 000011 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale无值, quant_scale有值, group_index无值, 000010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale无值, quant_scale无值, group_index有值, 000001 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias无值, activate_scale无值, quant_scale无值, group_index无值, 000000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale有值, quant_scale有值, group_index有值, 001111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale有值, quant_scale有值, group_index有值, 002111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale有值, quant_scale有值, group_index有值, 003111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale有值, quant_scale有值, group_index有值, 004111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale有值, quant_scale有值, group_index无值, 001110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale有值, quant_scale有值, group_index无值, 002110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale有值, quant_scale有值, group_index无值, 003110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale有值, quant_scale有值, group_index无值, 004110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale有值, quant_scale无值, group_index有值, 001101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale有值, quant_scale无值, group_index有值, 002101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale有值, quant_scale无值, group_index有值, 003101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale有值, quant_scale无值, group_index有值, 004101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale有值, quant_scale无值, group_index无值, 001100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale有值, quant_scale无值, group_index无值, 002100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale有值, quant_scale无值, group_index无值, 003100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_TRUE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale有值, quant_scale无值, group_index无值, 004100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale无值, quant_scale有值, group_index有值, 001011 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale无值, quant_scale有值, group_index有值, 002011 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale无值, quant_scale有值, group_index有值, 003011 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale无值, quant_scale有值, group_index有值, 004011 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale无值, quant_scale有值, group_index无值, 001010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale无值, quant_scale有值, group_index无值, 002010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale无值, quant_scale有值, group_index无值, 003010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_TRUE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale无值, quant_scale有值, group_index无值, 004010 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale无值, quant_scale无值, group_index有值, 001001 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale无值, quant_scale无值, group_index有值, 002001 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale无值, quant_scale无值, group_index有值, 003001 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_TRUE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale无值, quant_scale无值, group_index有值, 004001 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_TRUE_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=int32, activate_scale无值, quant_scale无值, group_index无值, 001000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP32_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float32, activate_scale无值, quant_scale无值, group_index无值, 002000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_FP16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=float16, activate_scale无值, quant_scale无值, group_index无值, 003000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == ACTDIM_FALSE_X_FALSE_BIAS_BF16_ACT_FALSE_QUANT_FALSE_GROUPINDEX_FALSE_STATIC + // activate_dim=-1, x=int32/fp16/bf6, bias=bfloat16, activate_scale无值, quant_scale无值, group_index无值, 004000 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantBaseStatic op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_INT32_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_INT32_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_INT32_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_INT32_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_INT32_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_INT32_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_INT32_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_INT32_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP32_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP32_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP32_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP32_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP32_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP32_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP32_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP32_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); +#elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP16_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP16_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP16_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP16_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP16_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP16_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP16_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_FP16_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_BF16_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_BF16_ACT_FALSE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_BF16_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_BF16_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_BF16_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_BF16_ACT_TRUE_QUANT_FALSE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_BF16_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_DYNAMIC_BIAS_BF16_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantDynamicNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, globalWorkGm, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + // 静态量化, bias无值, activate_scale无值, quant_scale有值, quant_offset无值, 无group, 1000100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, nullptr, nullptr, quantScale, nullptr, nullptr, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_TRUE_GROUP_FALSE + // 静态量化, bias无值, activate_scale无值, quant_scale有值, quant_offset有值, 无group, 1000110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, nullptr, nullptr, quantScale, quantOffset, nullptr, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + // 静态量化, bias无值, activate_scale有值, quant_scale有值, quant_offset无值, 无group, 1001100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, nullptr, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_TRUE_GROUP_FALSE + // 静态量化, bias无值, activate_scale有值, quant_scale有值, quant_offset有值, 无group, 1001110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, quantOffset, nullptr, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + // 静态量化, bias有值, activate_scale无值, quant_scale有值, quant_offset无值, 无group, 1010100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, nullptr, bias, quantScale, nullptr, nullptr, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_OFFSET_TRUE_GROUP_FALSE + // 静态量化, bias有值, activate_scale无值, quant_scale有值, quant_offset有值, 无group, 1010110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, nullptr, bias, quantScale, quantOffset, nullptr, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_FALSE + // 静态量化, bias有值, activate_scale有值, quant_scale有值, quant_offset无值, 无group, 1011100 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, nullptr, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_OFFSET_TRUE_GROUP_FALSE + // 静态量化, bias有值, activate_scale有值, quant_scale有值, quant_offset有值, 无group, 1011110 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, nullptr, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + // 静态量化, bias无值, activate_scale无值, quant_scale有值, quant_offset无值, 有group, 1000101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, nullptr, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_FALSE_ACT_FALSE_QUANT_TRUE_OFFSET_TRUE_GROUP_TRUE + // 静态量化, bias无值, activate_scale无值, quant_scale有值, quant_offset有值, 有group, 1000111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, nullptr, nullptr, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + // 静态量化, bias无值, activate_scale有值, quant_scale有值, quant_offset无值, 有group, 1001101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_FALSE_ACT_TRUE_QUANT_TRUE_OFFSET_TRUE_GROUP_TRUE + // 静态量化, bias无值, activate_scale有值, quant_scale有值, quant_offset有值, 有group, 1001111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, activationScale, nullptr, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + // 静态量化, bias有值, activate_scale无值, quant_scale有值, quant_offset无值, 有group, 1010101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, nullptr, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_TRUE_ACT_FALSE_QUANT_TRUE_OFFSET_TRUE_GROUP_TRUE + // 静态量化, bias有值, activate_scale无值, quant_scale有值, quant_offset有值, 有group, 1010111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, nullptr, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_OFFSET_FALSE_GROUP_TRUE + // 静态量化, bias有值, activate_scale有值, quant_scale有值, quant_offset无值, 有group, 1011101 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, nullptr, groupIndex, y, scale, tilingData); + op.Process(); + #elif TILING_KEY_VAR == QUANTMODE_STATIC_BIAS_TRUE_ACT_TRUE_QUANT_TRUE_OFFSET_TRUE_GROUP_TRUE + // 静态量化, bias有值, activate_scale有值, quant_scale有值, quant_offset有值, 有group, 1011111 + GET_TILING_DATA_WITH_STRUCT(DequantSwigluQuantV35BaseTilingData, tilingDataIn, tiling); + const DequantSwigluQuantV35BaseTilingData* __restrict tilingData = &tilingDataIn; + DequantSwigluQuantV35Ops::DequantSwigluQuantStaticNotFull op(&pipe); + op.Init(x, weightScale, activationScale, bias, quantScale, quantOffset, groupIndex, y, scale, tilingData); + op.Process(); + #endif +} diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_cut_group.h b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_cut_group.h new file mode 100644 index 000000000..1571297ff --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_cut_group.h @@ -0,0 +1,62 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_cut_group.h + * \brief + */ + +#ifndef DEQUANT_SWIGLU_QUANT_CUT_GROUP_H +#define DEQUANT_SWIGLU_QUANT_CUT_GROUP_H + +#include "kernel_tiling/kernel_tiling.h" +#include "kernel_operator.h" +#include "dequant_swiglu_quant.h" + +namespace DequantSwigluQuantGroupOps { +using namespace AscendC; +constexpr static int64_t GROUPINDEX_STRIDE = 2; + +TEMPLATE_DSQ_DECLARE +class DequantSwigluQuantGroup : public DequantSwigluQuantOps::DequantSwigluQuantBase { + public: + __aicore__ inline DequantSwigluQuantGroup(TPipe* pipe) : DequantSwigluQuantOps::DequantSwigluQuantBase(pipe) + { + this->pipe_ = pipe; + }; + __aicore__ inline void Process(); +}; +// 公共函数实现 + +TEMPLATE_DSQ_DECLARE +__aicore__ inline void DequantSwigluQuantGroup::Process() { + this->CreateOffsetLocalTensor(this->UbSingleOutSize_, this->tl_->swigluMode); + this->groupOffset_ = 0; + int64_t cuGroupIdx = this->blockIdx_; + for (int32_t groupIdx = 0; groupIdx < this->tl_->inGroupNum; ++groupIdx) { + int64_t realGroupIdx = this->tl_->speGroupType == 0 ? static_cast(groupIdx) : + static_cast(this->groupIndexGm_(groupIdx*GROUPINDEX_STRIDE)); + this->realDimx_ = this->tl_->speGroupType == 0 ? static_cast(this->groupIndexGm_(groupIdx)) : + static_cast(this->groupIndexGm_(groupIdx*GROUPINDEX_STRIDE + 1)); + if (this->realDimx_ <= 0 && this->tl_->speGroupType) { + break; + } + if (groupIdx == cuGroupIdx) { + if (this->realDimx_ > 0) { + this->ProcessSingleGroupPerCore(realGroupIdx, this->realDimx_, this->groupOffset_); + } + cuGroupIdx += this->tl_->maxCoreNum; + } + this->groupOffset_ += this->realDimx_; + } +} + +} // namespace DequantSwigluQuantGroupOps +#endif // DEQUANT_SWIGLU_QUANT_CUT_GROUP_H diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_base.hpp b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_base.hpp new file mode 100644 index 000000000..7c28043d9 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_base.hpp @@ -0,0 +1,592 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_dynamic_base.hpp + * \brief + */ + +#ifndef CANN_DEQUANT_SWIGLU_QUANT_DYNAMIC_BASE_HPP +#define CANN_DEQUANT_SWIGLU_QUANT_DYNAMIC_BASE_HPP +#include "kernel_operator.h" + +#define TEMPLATE_DECLARE template +#define TEMPLATE_ARGS InType, CalcType, BiasType, OutType, bufferNum, quantIsOne +namespace DequantSwigluQuant { +constexpr uint32_t DOUBLE = 2; +using namespace AscendC; + +TEMPLATE_DECLARE +class DequantSwigluQuantDynamicBase { +public: + __aicore__ inline DequantSwigluQuantDynamicBase() {} + __aicore__ inline ~DequantSwigluQuantDynamicBase() {} + + __aicore__ inline void InitCommon(GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, + GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + GM_ADDR userspace, const SwiGluTilingData* tilingData, TPipe* pipe_) + { + pipe = pipe_; + curBlockIdx = GetBlockIdx(); + activateLeft = tilingData->activateLeft; + quantScaleIsEmpty = tilingData->quantScaleIsEmpty; + activateScaleIsEmpty = tilingData->activateScaleIsEmpty; + biasIsEmpty = tilingData->biasIsEmpty; + colNum = tilingData->colLen; + rowNum = tilingData->rowLen; + useCoreNum = tilingData->usedCoreNum; + // 每行为全载情况下每次最大拷贝行数 + baseRowLen = tilingData->baseRowLen; + baseColLen = tilingData->baseColLen; + if (rowNum < useCoreNum) { + useCoreNum = rowNum; + } + // 行全载和不全载情况下，分别判断是否对齐 + isMultiCols = baseRowLen == 1 && this->baseColLen < this->colNum; + if (isMultiCols) { + isOut32BAligned = baseColLen == Align(baseColLen, sizeof(InType)); + } else { + isOut32BAligned = (colNum % blockBytes == 0) || (baseRowLen == 1); + } + perRoundCnt = useCoreNum == 0 ? 0 : rowNum / useCoreNum; + uint32_t remainCnt = rowNum - useCoreNum * perRoundCnt; + numRound = perRoundCnt; + if (curBlockIdx < remainCnt) { + numRound = perRoundCnt + 1; + biasOffset = curBlockIdx * (perRoundCnt + 1); + } else { + biasOffset = (perRoundCnt + 1) * remainCnt + (curBlockIdx - remainCnt) * perRoundCnt; + } + xGm.SetGlobalBuffer((__gm__ InType*)x_gm + biasOffset * colNum * DOUBLE, colNum * numRound * DOUBLE); + scaleGm.SetGlobalBuffer((__gm__ float*)scale_gm, rowNum); + yGm.SetGlobalBuffer((__gm__ int8_t*)y_gm + biasOffset * colNum, numRound * colNum); + quantScaleGm.SetGlobalBuffer((__gm__ float*)quant_scale_gm, colNum); + if (this->quantScaleIsEmpty == 0) { + if constexpr (quantIsOne == 1) { + quant_scale = ((__gm__ float*)quant_scale_gm)[0]; + } + } + if (isMultiCols) { + swigluTmpGm.SetGlobalBuffer((__gm__ float*)userspace + curBlockIdx * colNum, colNum); + } + } + + __aicore__ inline void BaseProcess() + { + if (this->curBlockIdx >= this->useCoreNum) { + return; + } + this->maxTempLocal = this->outQueueS.template AllocTensor(); + uint32_t offset1 = 0; + uint32_t offset2 = this->colNum; + if (this->activateLeft == 0) { + offset1 = this->colNum; + offset2 = 0; + } + if (!this->isMultiCols) { + this->CanFullLocaOneRow(offset1, offset2); + } else { + this->CanNotFullLocaOneRow(offset1, offset2); + } + this->CopyOutScale(this->numRound); + } + + __aicore__ inline void InitUbBufferCommon(uint64_t tileLength, uint32_t realRowLen) + { + uint64_t alignTileLength = tileLength; + if (!isOut32BAligned) { + alignTileLength = Align(tileLength, sizeof(int8_t)); + } + pipe->InitBuffer(inputTempBufferBF16D, alignTileLength * sizeof(CalcType) * baseRowLen); + pipe->InitBuffer(outputTempBufferBF16D, alignTileLength * sizeof(CalcType) * baseRowLen); + pipe->InitBuffer(inQueueA, 1, alignTileLength * sizeof(InType) * baseRowLen); + pipe->InitBuffer(inQueueB, 1, alignTileLength * sizeof(InType) * baseRowLen); + pipe->InitBuffer(swiGluQueue, 1, alignTileLength * sizeof(float) * baseRowLen); + if (quantScaleIsEmpty == 0) { + if (quantIsOne == 0) { + pipe->InitBuffer(inQueueQuantScale, bufferNum, alignTileLength * sizeof(float)); + } + } + pipe->InitBuffer(outQueueF, 1, alignTileLength * sizeof(int8_t) * baseRowLen); + pipe->InitBuffer(outQueueS, 1, AlignBytes(realRowLen, sizeof(float))); + } + + __aicore__ inline float dynamicMultiColMax(uint64_t rowId, uint64_t tileLen, int64_t colLoop) + { + LocalTensor swiLocal = swiGluQueue.template DeQue(); + LocalTensor absTempLocal = inputTempBufferBF16D.Get(); + Abs(absTempLocal, swiLocal, tileLen); + PipeBarrier(); + ReduceMax(maxTempLocal[rowId], absTempLocal, absTempLocal, tileLen); + DataCopyExtParams dataCopyParams{1, static_cast(tileLen * sizeof(float)), 0, 0, 0}; + DataCopyPad(swigluTmpGm[colLoop * baseColLen], swiLocal, dataCopyParams); + swiGluQueue.FreeTensor(swiLocal); + return maxTempLocal.GetValue(rowId); + } + + __aicore__ inline void dynamicAllColOut(uint64_t rowId, uint64_t tileLen, int64_t colLoop, float scale) + { + LocalTensor swiLocal = swiGluQueue.template AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(tileLen * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(swiLocal, swigluTmpGm[colLoop * baseColLen], dataCopyParams, dataCopyPadParams); + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S)); + SetFlag(eventId); + WaitFlag(eventId); + Muls(swiLocal, swiLocal, scale, tileLen); + PipeBarrier(); + LocalTensor int16Local = outputTempBufferBF16D.Get(); + Cast(int16Local, swiLocal, RoundMode::CAST_RINT, tileLen); + PipeBarrier(); + swiGluQueue.FreeTensor(swiLocal); + // int16-> half + LocalTensor halfLocal = int16Local.ReinterpretCast(); + Cast(halfLocal, int16Local, RoundMode::CAST_NONE, tileLen); + PipeBarrier(); + // half -> int8_t + LocalTensor outLocal = outQueueF.template AllocTensor(); + Cast(outLocal, halfLocal, RoundMode::CAST_NONE, tileLen); + outQueueF.EnQue(outLocal); + outLocal = outQueueF.DeQue(); + event_t eventId2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S)); + SetFlag(eventId2); + WaitFlag(eventId2); + DataCopyExtParams intriParams{1, static_cast(tileLen), 0, 0, 0}; + DataCopyPad(yGm[rowId * colNum + colLoop * baseColLen], outLocal, intriParams); + outQueueF.FreeTensor(outLocal); + } + + __aicore__ inline void DynamicCompute(uint64_t rowId, uint64_t tileLen, uint64_t length) + { + LocalTensor swiLocal = swiGluQueue.template DeQue(); + LocalTensor absTempLocal = inputTempBufferBF16D.Get(); + Abs(absTempLocal, swiLocal, tileLen); + uint32_t offsetCalc = (length == 0 ? 0 : (tileLen / length)); + PipeBarrier(); + for (int64_t i = 0; i < length; i++) { + ReduceMax(maxTempLocal[rowId * baseRowLen + i], absTempLocal[i * offsetCalc], absTempLocal[i * offsetCalc], colNum); + PipeBarrier(); + + event_t eventIdV2S = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdV2S); + WaitFlag(eventIdV2S); + float value = maxTempLocal.GetValue(rowId * baseRowLen + i) / 127; + maxTempLocal.SetValue(rowId * baseRowLen + i, value); + float scale = 1 / value; + Muls(swiLocal[i * offsetCalc], swiLocal[i * offsetCalc], scale, colNum); + PipeBarrier(); + } + LocalTensor int16Local = outputTempBufferBF16D.Get(); + Cast(int16Local, swiLocal, RoundMode::CAST_RINT, tileLen); + PipeBarrier(); + + // int16-> half + LocalTensor halfLocal = int16Local.ReinterpretCast(); + Cast(halfLocal, int16Local, RoundMode::CAST_NONE, tileLen); + PipeBarrier(); + swiGluQueue.FreeTensor(swiLocal); + + // half -> int8_t + LocalTensor outLocal = outQueueF.template AllocTensor(); + Cast(outLocal, halfLocal, RoundMode::CAST_NONE, tileLen); + + event_t eventId1 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventId1); + WaitFlag(eventId1); + + if (isOut32BAligned) { + DataCopyExtParams intriParams{1, static_cast(tileLen), 0, 0, 0}; + DataCopyPad(yGm[rowId * colNum * baseRowLen], outLocal, intriParams); + } else { + DataCopyExtParams intriParams{static_cast(length), static_cast(colNum), 0, 0, 0}; + DataCopyPad(yGm[rowId * colNum * baseRowLen], outLocal, intriParams); + } + outQueueF.FreeTensor(outLocal); + } + + __aicore__ inline void BaseComputeWithQuantScale(LocalTensor &outTmpLocal, LocalTensor &bLocal, + uint64_t curTileLen, uint64_t blockCount) + { + LocalTensor swiLocal = swiGluQueue.template AllocTensor(); + Mul(swiLocal, outTmpLocal, bLocal, curTileLen); + if (quantScaleIsEmpty == 0) { + PipeBarrier(); + if constexpr (quantIsOne == 0) { + uint32_t calcOffset = (blockCount == 0 ? 0 : curTileLen / blockCount); + for (uint64_t idx = 0; idx < blockCount; idx++) { + Mul(swiLocal[idx * calcOffset], swiLocal[idx * calcOffset], quantScaleLocal, calcOffset); + } + } else { + Muls(swiLocal, swiLocal, quant_scale, curTileLen); + } + } + swiGluQueue.template EnQue(swiLocal); + } + + __aicore__ inline void BaseCompute(uint64_t curTileLen, uint64_t blockCount, uint64_t idx) + { + LocalTensor inALocal = this->inQueueA.template DeQue(); + LocalTensor outTmpLocal = this->outputTempBufferBF16D.template Get(); + LocalTensor inputTmpLocal = this->inputTempBufferBF16D.template Get(); + float value = this->getActivateScaleValue(idx); + if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + this->addBiasWithBiasInt(inALocal, this->biasLocalA, curTileLen); + } + } + Cast(inputTmpLocal, inALocal, RoundMode::CAST_NONE, curTileLen); + PipeBarrier(); + this->inQueueA.template FreeTensor(inALocal); + if constexpr (std::is_same_v) { + addWeightScaleAndActivateScale(inputTmpLocal, this->weightScaleLocalA, curTileLen, value); + if constexpr (std::is_same_v || std::is_same_v || std::is_same_v) { + if (this->biasIsEmpty == 0) { + addBiasWithBiasFloat(inputTmpLocal, this->biasLocalA, curTileLen); + } + } + } + Muls(outTmpLocal, inputTmpLocal, this->beta, curTileLen); + PipeBarrier(); + Exp(outTmpLocal, outTmpLocal, curTileLen); + PipeBarrier(); + Adds(outTmpLocal, outTmpLocal, CalcType(1.0), curTileLen); + PipeBarrier(); + Div(outTmpLocal, inputTmpLocal, outTmpLocal, curTileLen); + PipeBarrier(); + LocalTensor bLocal_ = this->inQueueB.template DeQue(); + if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + this->addBiasWithBiasInt(bLocal_, this->biasLocalB, curTileLen); + } + } + LocalTensor bLocal = this->inputTempBufferBF16D.template Get(); + Cast(bLocal, bLocal_, RoundMode::CAST_NONE, curTileLen); + PipeBarrier(); + if constexpr (std::is_same_v) { + addWeightScaleAndActivateScale(bLocal, this->weightScaleLocalB, curTileLen, value); + if constexpr (std::is_same_v || std::is_same_v || std::is_same_v) { + if (this->biasIsEmpty == 0) { + addBiasWithBiasFloat(bLocal, this->biasLocalB, curTileLen); + } + } + } + this->inQueueB.template FreeTensor(bLocal_); + BaseComputeWithQuantScale(outTmpLocal, bLocal, curTileLen, blockCount); + } + + __aicore__ inline void CanFullLocaOneRow(uint32_t offset1, uint32_t offset2) + { + if (this->quantScaleIsEmpty == 0) { + if constexpr (quantIsOne == 0) { + this->CopyInQuantScale(this->colNum, 0); + } + } + CopyInDequantBuffer(offset1, offset2, this->colNum); + this->alignSize = this->isOut32BAligned ? this->colNum : this->Align(this->colNum, sizeof(int8_t)); + + int64_t blockCount = this->baseRowLen; + int64_t loops = (this->numRound + blockCount - 1) / blockCount; + int64_t lastLoopBlockCount = this->numRound - (loops - 1) * blockCount; + int64_t lastLoopColSize = this->alignSize * lastLoopBlockCount; + int64_t perLoopColSize = this->alignSize * blockCount; + + uint32_t aligCalcNum = this->Align(this->colNum, sizeof(InType)); + uint32_t alig8CalcNum = this->Align(this->colNum, sizeof(int8_t)); + this->dstStride = this->isOut32BAligned ? 0 : (alig8CalcNum - this->colNum) * sizeof(InType) / this->blockBytes; + for (uint32_t i = 0; i < loops - 1; i++) { + uint32_t base = i * (this->colNum * DOUBLE) * this->baseRowLen; + this->CopyIn(this->colNum, offset1 + base, offset2 + base, blockCount); + this->BaseCompute(perLoopColSize, blockCount, i); + this->DynamicCompute(i, perLoopColSize, blockCount); + } + uint32_t base = (loops - 1) * (this->colNum * DOUBLE) * this->baseRowLen; + this->CopyIn(this->colNum, offset1 + base, offset2 + base, lastLoopBlockCount); + this->BaseCompute(lastLoopColSize, lastLoopBlockCount, (loops - 1)); + this->DynamicCompute((loops - 1), lastLoopColSize, lastLoopBlockCount); + + if (this->quantScaleIsEmpty == 0) { + if constexpr (quantIsOne == 0) { + this->inQueueQuantScale.FreeTensor(this->quantScaleLocal); + } + } + FreeDequantBuffer(); + } + + __aicore__ inline void CanNotFullLocaOneRow(uint32_t offset1, uint32_t offset2) + { + int64_t colLoops = (this->colNum + this->baseColLen - 1) / this->baseColLen; + int64_t lastColNum = this->colNum - (colLoops - 1) * this->baseColLen; + for (uint32_t i = 0; i < this->numRound; i++) { + uint32_t tmp = 0xFF7FFFFF; + float reduceMax = *((float*)&tmp); + for (uint32_t j = 0; j < colLoops; j++) { + int64_t curColNum = this->baseColLen; + if (j == colLoops - 1) { + curColNum = lastColNum; + } + if (this->quantScaleIsEmpty == 0) { + if constexpr (quantIsOne == 0) { + this->CopyInQuantScale(curColNum, j * this->baseColLen); + } + } + bool isOutAligned = curColNum == this->Align(curColNum, sizeof(InType)); + uint32_t alignColNum = isOutAligned ? curColNum : this->Align(curColNum, sizeof(OutType)); + uint32_t base = i * (this->colNum * DOUBLE) + j * this->baseColLen; + CopyInDequantBuffer(offset1 + j * this->baseColLen, offset2 + j * this->baseColLen, curColNum); + this->CopyIn(curColNum, offset1 + base, offset2 + base, 1); + this->BaseCompute(alignColNum, 1, i); + float maxValue = this->dynamicMultiColMax(i, curColNum, j); + if (maxValue > reduceMax) { + reduceMax = maxValue; + } + if (this->quantScaleIsEmpty == 0) { + if constexpr (quantIsOne == 0) { + this->inQueueQuantScale.FreeTensor(this->quantScaleLocal); + } + } + FreeDequantBuffer(); + } + float value = reduceMax / 127.0f; + this->maxTempLocal.SetValue(i, value); + float scale = 1 / value; + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventId); + WaitFlag(eventId); + for (uint32_t j = 0; j < colLoops; j++) { + int64_t curColNum = this->baseColLen; + if (j == colLoops - 1) { + curColNum = lastColNum; + } + bool isOutAligned = curColNum == this->Align(curColNum, sizeof(InType)); + uint32_t alignColNum = isOutAligned ? curColNum : this->Align(curColNum, sizeof(OutType)); + this->dynamicAllColOut(i, curColNum, j, scale); + } + } + } + + __aicore__ inline void CopyInDequantBuffer(uint32_t offset1, uint32_t offset2, uint32_t dataTileLen) + { + if constexpr (std::is_same_v) { + this->CopyInWeightAndBias(dataTileLen, offset1, offset2); + this->CopyInActivateScale(0, this->numRound); + if (this->biasIsEmpty == 0) { + this->biasLocalA = this->inBiasQueueA.template DeQue(); + this->biasLocalB = this->inBiasQueueB.template DeQue(); + } + this->weightScaleLocalA = this->weightScaleQueueA.template DeQue(); + this->weightScaleLocalB = this->weightScaleQueueB.template DeQue(); + if (this->activateScaleIsEmpty == 0) { + this->activateLocal = this->inQueueActivationScale.template DeQue(); + } + } + } + + __aicore__ inline void FreeDequantBuffer() + { + if constexpr (std::is_same_v) { + if (this->biasIsEmpty == 0) { + this->inBiasQueueA.FreeTensor(this->biasLocalA); + this->inBiasQueueB.FreeTensor(this->biasLocalB); + } + this->weightScaleQueueA.FreeTensor(this->weightScaleLocalA); + this->weightScaleQueueB.FreeTensor(this->weightScaleLocalB); + if (this->activateScaleIsEmpty == 0) { + this->inQueueActivationScale.FreeTensor(this->activateLocal); + } + } + } + + __aicore__ inline float getActivateScaleValue(uint64_t idx) + { + float value = 1; + if constexpr (std::is_same_v) { + if (activateScaleIsEmpty == 0) { + event_t eventIdM2S = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S)); + SetFlag(eventIdM2S); + WaitFlag(eventIdM2S); + value = activateLocal.GetValue(idx); + } + } + return value; + } + + __aicore__ inline void addWeightScaleAndActivateScale( + LocalTensor &dstLocal, LocalTensor &weightScaleLocal, uint64_t curTileLen, float value) + { + Mul(dstLocal, dstLocal, weightScaleLocal, curTileLen); + PipeBarrier(); + if (activateScaleIsEmpty == 0) { + Muls(dstLocal, dstLocal, value, curTileLen); + PipeBarrier(); + } + } + + __aicore__ inline void addBiasWithBiasInt(LocalTensor &dstLocal, LocalTensor &biasLocal, uint64_t curTileLen) + { + if (this->biasIsEmpty == 0) { + Add(dstLocal, dstLocal, biasLocal, curTileLen); + PipeBarrier(); + } + } + + __aicore__ inline void addBiasWithBiasFloat(LocalTensor &dstLocal, LocalTensor &biasLocal, uint64_t curTileLen) + { + if (this->biasIsEmpty == 0) { + if constexpr (std::is_same_v) { + Add(dstLocal, dstLocal, biasLocal, curTileLen); + } else { + Cast(biasFloatLocalB, biasLocal, RoundMode::CAST_NONE, curTileLen); + PipeBarrier(); + Add(dstLocal, dstLocal, biasFloatLocalB, curTileLen); + } + PipeBarrier(); + } + } + + __aicore__ inline void CopyOutScale(uint32_t realRowLen) + { + DataCopyExtParams intriParams{1, static_cast(sizeof(float) * realRowLen), 0, 0, 0}; + DataCopyPad(scaleGm[biasOffset], maxTempLocal, intriParams); + outQueueS.FreeTensor(maxTempLocal); + } + + __aicore__ inline void CopyInActivateScale(uint32_t offset3, uint32_t blockCount) + { + if (activateScaleIsEmpty == 0) { + DataCopyExtParams activateparams = {1, static_cast(blockCount * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + LocalTensor activateLocal1 = inQueueActivationScale.template AllocTensor(); + DataCopyPad(activateLocal1, activationScaleGm[offset3], activateparams, padParams); + inQueueActivationScale.EnQue(activateLocal1); + } + } + + __aicore__ inline void CopyInWeightAndBias(uint32_t dataTileLen, uint32_t offset1, uint32_t offset2) + { + DataCopyExtParams params = {1, static_cast(dataTileLen * sizeof(float)), 0, 0, 0}; + DataCopyExtParams paramsBias = {1, static_cast(dataTileLen * sizeof(BiasType)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPadExtParams padParams1{false, 0, 0, 0}; + + if (this->biasIsEmpty == 0) { + // copy bias A + LocalTensor biasLocalA1 = inBiasQueueA.template AllocTensor(); + DataCopyPad(biasLocalA1, biasGm[offset1], paramsBias, padParams1); + inBiasQueueA.EnQue(biasLocalA1); + // copy bias B + LocalTensor biasLocalB1 = inBiasQueueB.template AllocTensor(); + DataCopyPad(biasLocalB1, biasGm[offset2], paramsBias, padParams1); + inBiasQueueB.EnQue(biasLocalB1); + } + // copy ws A + LocalTensor wsLocalA1 = weightScaleQueueA.template AllocTensor(); + DataCopyPad(wsLocalA1, weightScaleGm[offset1], params, padParams); + weightScaleQueueA.EnQue(wsLocalA1); + // copy ws B + LocalTensor wsLocalB1 = weightScaleQueueB.template AllocTensor(); + DataCopyPad(wsLocalB1, weightScaleGm[offset2], params, padParams); + weightScaleQueueB.EnQue(wsLocalB1); + } + + __aicore__ inline void CopyInQuantScale(uint64_t dataTileLength, uint64_t offset) + { + DataCopyExtParams dataCopyParams{1, static_cast(dataTileLength * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + LocalTensor scaleLocal = inQueueQuantScale.template AllocTensor(); + DataCopyPad(scaleLocal, quantScaleGm[offset], dataCopyParams, dataCopyPadParams); + inQueueQuantScale.EnQue(scaleLocal); + quantScaleLocal = inQueueQuantScale.template DeQue(); + } + + __aicore__ inline void CopyIn(uint32_t dataTileLen, uint32_t offset1, uint32_t offset2, uint32_t blockCount) + { + uint32_t srcStride = dataTileLen * sizeof(InType); + DataCopyExtParams dataCopyParams{static_cast(blockCount), + static_cast(dataTileLen * sizeof(InType)), srcStride, dstStride, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + // Copy A + LocalTensor aLocal = inQueueA.template AllocTensor(); + DataCopyPad(aLocal, xGm[offset1], dataCopyParams, dataCopyPadParams); + inQueueA.EnQue(aLocal); + // Copy B + LocalTensor bLocal = inQueueB.template AllocTensor(); + DataCopyPad(bLocal, xGm[offset2], dataCopyParams, dataCopyPadParams); + inQueueB.EnQue(bLocal); + } + + __aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes) + { + if (bytes == 0) { + return 0; + } + return (elementNum * bytes + blockBytes - 1) / blockBytes * blockBytes / bytes; + } + + __aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes) + { + return (elementNum * bytes + blockBytes - 1) / blockBytes * blockBytes; + } +protected: + TPipe* pipe; + GlobalTensor xGm; + GlobalTensor quantScaleGm; + GlobalTensor weightScaleGm; + GlobalTensor activationScaleGm; + GlobalTensor biasGm; + GlobalTensor yGm; + GlobalTensor scaleGm; + GlobalTensor swigluTmpGm; + TBuf inputTempBufferBF16D; + TBuf outputTempBufferBF16D; + TBuf inputBiasTempBufferA; + TBuf inputBiasTempBufferB; + TQue inQueueA; + TQue inQueueB; + TQue inQueueQuantScale; + TQue swiGluQueue; + TQue outQueueF; + TQue outQueueS; + TQue inBiasQueueA; + TQue inBiasQueueB; + TQue weightScaleQueueA; + TQue weightScaleQueueB; + TQue inQueueActivationScale; + LocalTensor maxTempLocal; + LocalTensor quantScaleLocal; + LocalTensor weightScaleLocalA; + LocalTensor weightScaleLocalB; + LocalTensor biasLocalA; + LocalTensor biasLocalB; + LocalTensor activateLocal; + LocalTensor biasFloatLocalA; + LocalTensor biasFloatLocalB; + float beta = -1.0f; + float quant_scale = 1; + uint32_t quantScaleIsEmpty = 0; + uint32_t biasIsEmpty = 0; + uint32_t activateScaleIsEmpty = 0; + uint64_t perRoundCnt = 0; + uint64_t numRound = 0; + uint32_t colNum = 0; + uint32_t rowNum = 0; + uint32_t useCoreNum = 0; + uint32_t biasOffset = 0; + uint32_t curBlockIdx = 0; + uint32_t activateLeft = 0; + uint32_t baseRowLen = 0; + uint32_t baseColLen = 0; + uint32_t alignSize = 0; + bool isOut32BAligned = true; + uint32_t dstStride = 0; + bool isMultiCols = false; + int64_t blockBytes = 32; +}; +} +#endif // CANN_DEQUANT_SWIGLU_QUANT_DYNAMIC_BASE_HPP diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_bf16.hpp b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_bf16.hpp new file mode 100644 index 000000000..d9ff38b50 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_bf16.hpp @@ -0,0 +1,55 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_dynamic_bf16.hpp + * \brief + */ + +#ifndef CANN_DEQUANT_SWIGLU_QUANT_DYNAMIC_BF16_HPP +#define CANN_DEQUANT_SWIGLU_QUANT_DYNAMIC_BF16_HPP + +#include "kernel_operator.h" +#include "dequant_swiglu_quant_dynamic_base.hpp" + +namespace DequantSwigluQuant { +using namespace AscendC; + +TEMPLATE_DECLARE +class DequantSwigluQuantDynamicBF16 : public DequantSwigluQuantDynamicBase { +public: + __aicore__ inline DequantSwigluQuantDynamicBF16(){}; + __aicore__ inline ~DequantSwigluQuantDynamicBF16(){}; + + __aicore__ inline void Init(GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, + GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + GM_ADDR userspace, const SwiGluTilingData* tilingData, TPipe* pipe_); + __aicore__ inline void Process(); +}; + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicBF16::Init( + GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, GM_ADDR quant_scale_gm, + GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, GM_ADDR userspace, const SwiGluTilingData* tilingData, + TPipe* pipe_) { + this->InitCommon(x_gm, weight_scale_gm, activation_scale_gm, bias_gm, quant_scale_gm, quant_offset_gm, y_gm, scale_gm, userspace, tilingData, pipe_); + + if (this->numRound < this->baseRowLen) { + this->baseRowLen = this->numRound; + } + this->InitUbBufferCommon(this->baseColLen, this->numRound); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicBF16::Process() { + this->BaseProcess(); +} +} // namespace DequantSwigluQuant +#endif // CANN_DEQUANT_SWIGLU_QUANT_DYNAMIC_BF16_HPP diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_bias_float.hpp b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_bias_float.hpp new file mode 100644 index 000000000..57321efca --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_bias_float.hpp @@ -0,0 +1,94 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_dynamic_bias_float.hpp + * \brief + */ + +#ifndef CANN_DEQUANT_SWIGLU_QUANT_DYNAMIC_BIAS_FLOAT_HPP +#define CANN_DEQUANT_SWIGLU_QUANT_DYNAMIC_BIAS_FLOAT_HPP + +#include "kernel_operator.h" +#include "dequant_swiglu_quant_dynamic_base.hpp" + +namespace DequantSwigluQuant { +using namespace AscendC; + +TEMPLATE_DECLARE + +class DequantSwigluQuantDynamicBiasFloat : public DequantSwigluQuantDynamicBase { +public: + __aicore__ inline DequantSwigluQuantDynamicBiasFloat(){}; + __aicore__ inline ~DequantSwigluQuantDynamicBiasFloat(){}; + + __aicore__ inline void Init(GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, + GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + GM_ADDR userspace, const SwiGluTilingData* tilingData, TPipe* pipe_); + __aicore__ inline void Process(); + +private: + __aicore__ inline void InitUbBuffer(uint64_t tileLength, uint32_t realRowLen); +}; + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicBiasFloat::Init( + GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, GM_ADDR quant_scale_gm, + GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, GM_ADDR userspace, const SwiGluTilingData* tilingData, + TPipe* pipe_) +{ + this->InitCommon(x_gm, weight_scale_gm, activation_scale_gm, bias_gm, quant_scale_gm, quant_offset_gm, y_gm, scale_gm, userspace, tilingData, pipe_); + this->weightScaleGm.SetGlobalBuffer((__gm__ float*)weight_scale_gm, this->colNum); + if (this->biasIsEmpty == 0) { + this->biasGm.SetGlobalBuffer((__gm__ BiasType*)bias_gm, this->colNum); + } + + if (this->activateScaleIsEmpty == 0) { + this->activationScaleGm.SetGlobalBuffer((__gm__ float*) activation_scale_gm + this->biasOffset, + this->numRound); + } + + this->InitUbBufferCommon(this->baseColLen, this->numRound); + + InitUbBuffer(this->baseColLen, this->numRound); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicBiasFloat::Process() { + this->BaseProcess(); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicBiasFloat::InitUbBuffer(uint64_t tileLength, + uint32_t realRowLen) +{ + uint64_t alignTileLength = tileLength; + if (!this->isOut32BAligned) { + alignTileLength = this->Align(tileLength, sizeof(int8_t)); + } + if (this->biasIsEmpty == 0) { + this->pipe->InitBuffer(this->inBiasQueueA, 1, alignTileLength * sizeof(BiasType)); + this->pipe->InitBuffer(this->inBiasQueueB, 1, alignTileLength * sizeof(BiasType)); + if constexpr (std::is_same_v || std::is_same_v) { + this->pipe->InitBuffer(this->inputBiasTempBufferA, alignTileLength * sizeof(float)); + this->pipe->InitBuffer(this->inputBiasTempBufferB, alignTileLength * sizeof(float)); + this->biasFloatLocalA = this->inputBiasTempBufferA.template Get(); + this->biasFloatLocalB = this->inputBiasTempBufferB.template Get(); + } + } + + this->pipe->InitBuffer(this->weightScaleQueueA, 1, alignTileLength * sizeof(float)); + this->pipe->InitBuffer(this->weightScaleQueueB, 1, alignTileLength * sizeof(float)); + if (this->activateScaleIsEmpty == 0) { + this->pipe->InitBuffer(this->inQueueActivationScale, 1, this->baseRowLen * sizeof(float)); + } +} +} // namespace DequantSwigluQuant +#endif // CANN_DEQUANT_SWIGLU_QUANT_DYNAMIC_BIAS_FLOAT_HPP diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_bias_int32.hpp b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_bias_int32.hpp new file mode 100644 index 000000000..d267fa72b --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_bias_int32.hpp @@ -0,0 +1,83 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_dynamic_bias_int32.hpp + * \brief + */ + +#ifndef CANN_DEQUANT_SWIGLU_QUANT_DYNAMIC_BIAS_INT32_HPP +#define CANN_DEQUANT_SWIGLU_QUANT_DYNAMIC_BIAS_INT32_HPP + +#include "kernel_operator.h" +#include "dequant_swiglu_quant_dynamic_base.hpp" + +namespace DequantSwigluQuant { +using namespace AscendC; + +constexpr int64_t BLOCK_BYTES = 32; + +TEMPLATE_DECLARE +class DequantSwigluQuantDynamicBiasInt32 : public DequantSwigluQuantDynamicBase { +public: + __aicore__ inline DequantSwigluQuantDynamicBiasInt32(){}; + __aicore__ inline ~DequantSwigluQuantDynamicBiasInt32(){}; + + __aicore__ inline void Init(GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, + GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + GM_ADDR userspace, const SwiGluTilingData* tilingData, TPipe* pipe_); + __aicore__ inline void Process(); + +private: + __aicore__ inline void InitUbBuffer(uint64_t tileLength, uint32_t realRowLen); +}; + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicBiasInt32::Init( + GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, GM_ADDR quant_scale_gm, + GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, GM_ADDR userspace, const SwiGluTilingData* tilingData, + TPipe* pipe_) { + this->InitCommon(x_gm, weight_scale_gm, activation_scale_gm, bias_gm, quant_scale_gm, quant_offset_gm, y_gm, scale_gm, userspace, tilingData, pipe_); + if (this->activateScaleIsEmpty == 0) { + this->activationScaleGm.SetGlobalBuffer((__gm__ float*) activation_scale_gm + this->biasOffset, + this->numRound); + } + this->weightScaleGm.SetGlobalBuffer((__gm__ float*)weight_scale_gm, this->colNum); + if (this->biasIsEmpty == 0) { + this->biasGm.SetGlobalBuffer((__gm__ BiasType*)bias_gm, this->colNum); + } + this->InitUbBufferCommon(this->baseColLen, this->numRound); + InitUbBuffer(this->baseColLen, this->numRound); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicBiasInt32::Process() { + this->BaseProcess(); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicBiasInt32::InitUbBuffer(uint64_t tileLength, + uint32_t realRowLen) { + uint64_t alignTileLength = tileLength; + if (!this->isOut32BAligned) { + alignTileLength = this->Align(tileLength, sizeof(int8_t)); + } + if (this->biasIsEmpty == 0) { + this->pipe->InitBuffer(this->inBiasQueueA, 1, alignTileLength * sizeof(BiasType)); + this->pipe->InitBuffer(this->inBiasQueueB, 1, alignTileLength * sizeof(BiasType) ); + } + this->pipe->InitBuffer(this->weightScaleQueueA, 1, alignTileLength * sizeof(float)); + this->pipe->InitBuffer(this->weightScaleQueueB, 1, alignTileLength * sizeof(float)); + if (this->activateScaleIsEmpty == 0) { + this->pipe->InitBuffer(this->inQueueActivationScale, 1, this->baseRowLen * sizeof(float)); + } +} +} // namespace DequantSwigluQuant +#endif // CANN_DEQUANT_SWIGLU_QUANT_DYNAMIC_BIAS_INT32_HPP diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_performance.hpp b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_performance.hpp new file mode 100644 index 000000000..7a98172be --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_dynamic_performance.hpp @@ -0,0 +1,384 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_dynamic_performance.hpp + * \brief + */ + +#ifndef DEQUANT_SWIGLU_QUANT_DYNAMIC_PERFORMANCE_HPP +#define DEQUANT_SWIGLU_QUANT_DYNAMIC_PERFORMANCE_HPP + +#include "kernel_operator.h" +#include "dequant_swiglu_quant_dynamic_base.hpp" + +namespace DequantSwigluQuant { +using namespace AscendC; + +TEMPLATE_DECLARE +class DequantSwigluQuantDynamicPerformance : public DequantSwigluQuantDynamicBiasFloat { +public: + __aicore__ inline DequantSwigluQuantDynamicPerformance(){}; + __aicore__ inline ~DequantSwigluQuantDynamicPerformance(){}; + + __aicore__ inline void Init(GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, + GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + GM_ADDR userspace, const SwiGluTilingData* tilingData, TPipe* pipe_); + __aicore__ inline void Process(); + __aicore__ inline void BaseCompute1(uint64_t curTileLen, uint64_t blockCount, uint64_t idx, int32_t ppFlag); + __aicore__ inline void BaseCompute2(uint64_t curTileLen, uint64_t blockCount, uint64_t idx, int32_t ppFlag); + __aicore__ inline void CopyOutF(uint64_t rowId, uint64_t tileLen, uint64_t length, int32_t ppFlag); + __aicore__ inline void CanFullLocaOneRow(uint32_t offset1, uint32_t offset2); + __aicore__ inline void CopyIn(uint32_t dataTileLen, uint32_t offset1, uint32_t offset2, uint32_t blockCount, int32_t ppFlag); + __aicore__ inline void CopyInDequantBuffer(uint32_t offset1, uint32_t offset2, uint32_t dataTileLen); + __aicore__ inline void InitCommon(GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, + GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + GM_ADDR userspace, const SwiGluTilingData* tilingData, TPipe* pipe_); + __aicore__ inline void InitUbBufferCommon(uint64_t tileLength, uint32_t realRowLen); + __aicore__ inline void CopyOutScale(uint32_t realRowLen); + __aicore__ inline void CopyInQuantScale(uint64_t dataTileLength, uint64_t offset); + +public: + uint32_t offsetCalc; + + LocalTensor outTmpLocal; + LocalTensor inputTmpLocal; + LocalTensor absTempLocal; + LocalTensor int16Local; + LocalTensor bLocal; + LocalTensor swiLocal; + + TBuf calcSwiGluTmpBuf; + TBuf weightScaleBufA; + TBuf weightScaleBufB; + TBuf quantScaleBuf; + + TBuf inQueueAPingBuf; + TBuf inQueueAPongBuf; + TBuf inQueueBPingBuf; + TBuf inQueueBPongBuf; + TBuf outQueueFPingBuf; + TBuf outQueueFPongBuf; + TBuf outQueueSBuf; + + LocalTensor inALocalPing; + LocalTensor inALocalPong; + LocalTensor inBLocalPing; + LocalTensor inBLocalPong; + LocalTensor outFLocalPing; + LocalTensor outFLocalPong; + + int32_t pingPongFlag = 0; + event_t eventId = EVENT_ID0; +}; + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicPerformance::Init( + GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, GM_ADDR quant_scale_gm, + GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, GM_ADDR userspace, const SwiGluTilingData* tilingData, + TPipe* pipe_) +{ + this->InitCommon(x_gm, weight_scale_gm, activation_scale_gm, bias_gm, quant_scale_gm, quant_offset_gm, y_gm, scale_gm, userspace, tilingData, pipe_); + this->weightScaleGm.SetGlobalBuffer((__gm__ float*)weight_scale_gm, this->colNum); + if (this->biasIsEmpty == 0) { + this->biasGm.SetGlobalBuffer((__gm__ BiasType*)bias_gm, this->colNum); + } + + if (this->activateScaleIsEmpty == 0) { + this->activationScaleGm.SetGlobalBuffer((__gm__ float*) activation_scale_gm + this->biasOffset, + this->numRound); + } + + this->InitUbBufferCommon(this->baseColLen, this->numRound); + + uint64_t alignTileLength = this->baseColLen; + if (!this->isOut32BAligned) { + alignTileLength = this->Align(this->baseColLen, sizeof(int8_t)); + } + + this->pipe->InitBuffer(weightScaleBufA, alignTileLength * sizeof(float)); + this->pipe->InitBuffer(weightScaleBufB, alignTileLength * sizeof(float)); + this->pipe->InitBuffer(quantScaleBuf, alignTileLength * sizeof(float)); + + this->pipe->InitBuffer(calcSwiGluTmpBuf, alignTileLength * sizeof(float) * this->baseRowLen); + + this->pipe->InitBuffer(inQueueAPingBuf, alignTileLength * sizeof(InType) * this->baseRowLen); + this->pipe->InitBuffer(inQueueAPongBuf, alignTileLength * sizeof(InType) * this->baseRowLen); + this->pipe->InitBuffer(inQueueBPingBuf, alignTileLength * sizeof(InType) * this->baseRowLen); + this->pipe->InitBuffer(inQueueBPongBuf, alignTileLength * sizeof(InType) * this->baseRowLen); + this->pipe->InitBuffer(outQueueFPingBuf, alignTileLength * sizeof(int8_t) * this->baseRowLen); + this->pipe->InitBuffer(outQueueFPongBuf, alignTileLength * sizeof(int8_t) * this->baseRowLen); + this->pipe->InitBuffer(outQueueSBuf, this->AlignBytes(this->numRound, sizeof(float))); + + outTmpLocal = this->outputTempBufferBF16D.template Get(); + inputTmpLocal = this->inputTempBufferBF16D.template Get(); + absTempLocal = this->inputTempBufferBF16D.template Get(); + int16Local = this->outputTempBufferBF16D.template Get(); + bLocal = this->inputTempBufferBF16D.template Get(); + swiLocal = calcSwiGluTmpBuf.Get(); + + inALocalPing = inQueueAPingBuf.Get(); + inALocalPong = inQueueAPongBuf.Get(); + inBLocalPing = inQueueBPingBuf.Get(); + inBLocalPong = inQueueBPongBuf.Get(); + outFLocalPing = outQueueFPingBuf.Get(); + outFLocalPong = outQueueFPongBuf.Get(); + this->maxTempLocal = outQueueSBuf.Get(); + + this->weightScaleLocalA = weightScaleBufA.Get(); + this->weightScaleLocalB = weightScaleBufB.Get(); + this->quantScaleLocal = quantScaleBuf.Get(); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicPerformance::InitCommon(GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, + GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + GM_ADDR userspace, const SwiGluTilingData* tilingData, TPipe* pipe_) +{ + this->pipe = pipe_; + this->curBlockIdx = GetBlockIdx(); + this->activateLeft = tilingData->activateLeft; + this->quantScaleIsEmpty = tilingData->quantScaleIsEmpty; + this->activateScaleIsEmpty = tilingData->activateScaleIsEmpty; + this->biasIsEmpty = tilingData->biasIsEmpty; + this->colNum = tilingData->colLen; + this->rowNum = tilingData->rowLen; + this->useCoreNum = tilingData->usedCoreNum; + // 每行为全载情况下每次最大拷贝行数 + this->baseRowLen = tilingData->baseRowLen; + this->baseColLen = tilingData->baseColLen; + if (this->rowNum < this->useCoreNum) { + this->useCoreNum = this->rowNum; + } + // 行全载和不全载情况下，分别判断是否对齐 + this->isMultiCols = this->baseRowLen == 1 && this->baseColLen < this->colNum; + if (this->isMultiCols) { + this->isOut32BAligned = this->baseColLen == this->Align(this->baseColLen, sizeof(InType)); + } else { + this->isOut32BAligned = (this->colNum % this->blockBytes == 0) || (this->baseRowLen == 1); + } + this->perRoundCnt = this->useCoreNum == 0 ? 0 : this->rowNum / this->useCoreNum; + uint32_t remainCnt = this->rowNum - this->useCoreNum * this->perRoundCnt; + this->numRound = this->perRoundCnt; + if (this->curBlockIdx < remainCnt) { + this->numRound = this->perRoundCnt + 1; + this->biasOffset = this->curBlockIdx * (this->perRoundCnt + 1); + } else { + this->biasOffset = (this->perRoundCnt + 1) * remainCnt + (this->curBlockIdx - remainCnt) * this->perRoundCnt; + } + this->xGm.SetGlobalBuffer((__gm__ InType*)x_gm + this->biasOffset * this->colNum * DOUBLE, this->colNum * this->numRound * DOUBLE); + this->scaleGm.SetGlobalBuffer((__gm__ float*)scale_gm, this->rowNum); + this->yGm.SetGlobalBuffer((__gm__ int8_t*)y_gm + this->biasOffset * this->colNum, this->numRound * this->colNum); + this->quantScaleGm.SetGlobalBuffer((__gm__ float*)quant_scale_gm, this->colNum); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicPerformance::InitUbBufferCommon(uint64_t tileLength, uint32_t realRowLen) +{ + uint64_t alignTileLength = tileLength; + if (!this->isOut32BAligned) { + alignTileLength = this->Align(tileLength, sizeof(int8_t)); + } + this->pipe->InitBuffer(this->inputTempBufferBF16D, alignTileLength * sizeof(CalcType) * this->baseRowLen); + this->pipe->InitBuffer(this->outputTempBufferBF16D, alignTileLength * sizeof(CalcType) * this->baseRowLen); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicPerformance::Process() { + if (this->curBlockIdx >= this->useCoreNum) { + return; + } + uint32_t offset1 = 0; + uint32_t offset2 = this->colNum; + if (this->activateLeft == 0) { + offset1 = this->colNum; + offset2 = 0; + } + this->CanFullLocaOneRow(offset1, offset2); + this->CopyOutScale(this->numRound); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicPerformance::CopyOutScale(uint32_t realRowLen) +{ + DataCopyExtParams intriParams{1, static_cast(sizeof(float) * realRowLen), 0, 0, 0}; + DataCopyPad(this->scaleGm[this->biasOffset], this->maxTempLocal, intriParams); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicPerformance::CanFullLocaOneRow(uint32_t offset1, uint32_t offset2) +{ + if (this->quantScaleIsEmpty == 0) { + this->CopyInQuantScale(this->colNum, 0); + } + this->CopyInDequantBuffer(offset1, offset2, this->colNum); + + this->alignSize = this->isOut32BAligned ? this->colNum : this->Align(this->colNum, sizeof(int8_t)); + int64_t blockCount = this->baseRowLen; + int64_t loops = (this->numRound + blockCount - 1) / blockCount; + int64_t lastLoopBlockCount = this->numRound - (loops - 1) * blockCount; + int64_t perLoopColSize = this->alignSize * blockCount; + offsetCalc = (blockCount == 0 ? 0 : (perLoopColSize / blockCount)); + uint32_t alig8CalcNum = this->Align(this->colNum, sizeof(int8_t)); + this->dstStride = this->isOut32BAligned ? 0 : (alig8CalcNum - this->colNum) * sizeof(InType) / this->blockBytes; + uint32_t base = (this->colNum * DOUBLE) * this->baseRowLen; + + SetFlag(EVENT_ID0); + SetFlag(EVENT_ID1); + for (uint32_t i = 0; i < loops; i++) { + eventId = pingPongFlag ? EVENT_ID1 : EVENT_ID0; + event_t eventIdNext = pingPongFlag ? EVENT_ID0 : EVENT_ID1; + + if (i == 0) { + WaitFlag(eventId); + this->CopyIn(this->colNum, offset1 + i * base, offset2 + i * base, blockCount, pingPongFlag); + SetFlag(eventId); + } + + WaitFlag(eventId); + this->BaseCompute1(perLoopColSize, blockCount, i, pingPongFlag); + + if(i != loops -1) { + WaitFlag(eventIdNext); + this->CopyIn(this->colNum, offset1 + (i + 1) * base, offset2 + (i + 1) * base, blockCount, 1 - pingPongFlag); + SetFlag(eventIdNext); + } + + this->BaseCompute2(perLoopColSize, blockCount, i, pingPongFlag); + SetFlag(eventId); + + WaitFlag(eventId); + this->CopyOutF(i, perLoopColSize, blockCount, pingPongFlag); + SetFlag(eventId); + + pingPongFlag = 1 - pingPongFlag; + } + + WaitFlag(EVENT_ID0); + WaitFlag(EVENT_ID1); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicPerformance::CopyInQuantScale(uint64_t dataTileLength, uint64_t offset) +{ + DataCopyExtParams dataCopyParams{1, static_cast(dataTileLength * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(this->quantScaleLocal, this->quantScaleGm[offset], dataCopyParams, dataCopyPadParams); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicPerformance::CopyInDequantBuffer(uint32_t offset1, uint32_t offset2, uint32_t dataTileLen) +{ + DataCopyExtParams params = {1, static_cast(dataTileLen * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + + DataCopyPad(this->weightScaleLocalA, this->weightScaleGm[offset1], params, padParams); + DataCopyPad(this->weightScaleLocalB, this->weightScaleGm[offset2], params, padParams); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicPerformance::CopyIn(uint32_t dataTileLen, uint32_t offset1, uint32_t offset2, uint32_t blockCount, int32_t ppFlag) +{ + uint32_t srcStride = dataTileLen * sizeof(InType); + DataCopyExtParams dataCopyParams{static_cast(blockCount), + static_cast(dataTileLen * sizeof(InType)), srcStride, this->dstStride, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + LocalTensor aLocal = ppFlag ? inALocalPong : inALocalPing; + LocalTensor bLocal = ppFlag ? inBLocalPong : inBLocalPing; + DataCopyPad(aLocal, this->xGm[offset1], dataCopyParams, dataCopyPadParams); + DataCopyPad(bLocal, this->xGm[offset2], dataCopyParams, dataCopyPadParams); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicPerformance::BaseCompute1(uint64_t curTileLen, uint64_t blockCount, uint64_t rowId, int32_t ppFlag) +{ + LocalTensor inALocal = ppFlag ? inALocalPong : inALocalPing; + LocalTensor bLocal_ = ppFlag ? inBLocalPong : inBLocalPing; + + float value = this->activationScaleGm.GetValue(rowId); + Cast(inputTmpLocal, inALocal, RoundMode::CAST_NONE, curTileLen); + PipeBarrier(); + + Mul(inputTmpLocal, inputTmpLocal, this->weightScaleLocalA, curTileLen); + PipeBarrier(); + + Muls(inputTmpLocal, inputTmpLocal, value, curTileLen); + PipeBarrier(); + + Muls(outTmpLocal, inputTmpLocal, this->beta, curTileLen); + PipeBarrier(); + + Exp(outTmpLocal, outTmpLocal, curTileLen); + PipeBarrier(); + + Adds(outTmpLocal, outTmpLocal, CalcType(1.0), curTileLen); + PipeBarrier(); + + Div(outTmpLocal, inputTmpLocal, outTmpLocal, curTileLen); + PipeBarrier(); + + Cast(bLocal, bLocal_, RoundMode::CAST_NONE, curTileLen); + PipeBarrier(); + + Mul(bLocal, bLocal, this->weightScaleLocalB, curTileLen); + PipeBarrier(); + + Muls(bLocal, bLocal, value, curTileLen); + PipeBarrier(); + + Mul(swiLocal, outTmpLocal, bLocal, curTileLen); + PipeBarrier(); + + if (this->quantScaleIsEmpty == 0) { + Mul(swiLocal, swiLocal, this->quantScaleLocal, curTileLen); + PipeBarrier(); + } +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicPerformance::BaseCompute2(uint64_t curTileLen, uint64_t blockCount, uint64_t rowId, int32_t ppFlag) +{ + LocalTensor outLocal = ppFlag ? outFLocalPong : outFLocalPing; + + Abs(absTempLocal, swiLocal, curTileLen); + PipeBarrier(); + + ReduceMax(this->maxTempLocal[rowId * this->baseRowLen], absTempLocal, absTempLocal, this->colNum); + PipeBarrier(); + + float value = this->maxTempLocal.GetValue(rowId * this->baseRowLen) / 127; + this->maxTempLocal.SetValue(rowId * this->baseRowLen, value); + float scale = 1 / value; + Muls(swiLocal, swiLocal, scale, this->colNum); + PipeBarrier(); + + Cast(int16Local, swiLocal, RoundMode::CAST_RINT, curTileLen); + PipeBarrier(); + + // int16-> half + LocalTensor halfLocal = int16Local.ReinterpretCast(); + Cast(halfLocal, int16Local, RoundMode::CAST_NONE, curTileLen); + PipeBarrier(); + + // half -> int8_t + Cast(outLocal, halfLocal, RoundMode::CAST_NONE, curTileLen); + PipeBarrier(); +} + +TEMPLATE_DECLARE +__aicore__ inline void DequantSwigluQuantDynamicPerformance::CopyOutF(uint64_t rowId, uint64_t tileLen, uint64_t length, int32_t ppFlag) +{ + LocalTensor outLocal = ppFlag ? outFLocalPong : outFLocalPing; + DataCopyExtParams intriParams{1, static_cast(tileLen), 0, 0, 0}; + DataCopyPad(this->yGm[rowId * this->colNum * this->baseRowLen], outLocal, intriParams); +} + +} // namespace DequantSwigluQuant +#endif // DEQUANT_SWIGLU_QUANT_DYNAMIC_PERFORMANCE_HPP diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_static_base.hpp b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_static_base.hpp new file mode 100644 index 000000000..2eca32f9d --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_static_base.hpp @@ -0,0 +1,371 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_static_base.hpp + * \brief + */ + +#ifndef CANN_DEQUANT_SWIGLU_QUANT_STATIC_BASE_HPP +#define CANN_DEQUANT_SWIGLU_QUANT_STATIC_BASE_HPP +#include "kernel_operator.h" + +#define TEMPLATE_DECLARE_STATIC template +#define TEMPLATE_ARGS_STATIC InType, CalcType, BiasType, OutType, bufferNum, quantIsOne + +namespace DequantSwigluQuant { +constexpr uint32_t NUM2 = 2; +using namespace AscendC; + +TEMPLATE_DECLARE_STATIC +class DequantSwigluQuantStaticBase { +public: + __aicore__ inline DequantSwigluQuantStaticBase() {} + __aicore__ inline ~DequantSwigluQuantStaticBase() {} + + __aicore__ inline void InitCommon(GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, + GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + const SwiGluTilingData* tilingData, TPipe* pipe_) { + this->blockIdx = GetBlockIdx(); + this->activateLeft = tilingData->activateLeft; + this->quantScaleIsEmpty = tilingData->quantScaleIsEmpty; + this->activateScaleIsEmpty = tilingData->activateScaleIsEmpty; + this->biasIsEmpty = tilingData->biasIsEmpty; + this->colNum = tilingData->colLen; + this->rowNum = tilingData->rowLen; + this->usedCoreNum = tilingData->usedCoreNum; + + this->baseRowLen = tilingData->baseRowLen; + this->baseColLen = tilingData->baseColLen < this->colNum ? tilingData->baseColLen : this->colNum; + this->curColNum = this->baseColLen; + if (this->rowNum < this->usedCoreNum) { + this->usedCoreNum = this->rowNum; + } + int64_t perRoundCnt = this->usedCoreNum == 0 ? 0 : this->rowNum / this->usedCoreNum; + int64_t remainCnt = this->rowNum - this->usedCoreNum * perRoundCnt; + this->curCoreRowNum = perRoundCnt; + if (this->blockIdx < remainCnt) { + this->curCoreRowNum = perRoundCnt + 1; + this->inputCopyOffset = this->blockIdx * this->curCoreRowNum; + } else { + this->inputCopyOffset = remainCnt * (perRoundCnt + 1) + (this->blockIdx - remainCnt) * perRoundCnt; + } + + this->xGm.SetGlobalBuffer((__gm__ InType*)x_gm + this->inputCopyOffset * this->colNum * NUM2, this->curCoreRowNum * this->colNum * NUM2); + this->yGm.SetGlobalBuffer((__gm__ OutType*)y_gm + this->inputCopyOffset * this->colNum, this->curCoreRowNum * this->colNum); + if (quantScaleIsEmpty == 0) { + if constexpr(quantIsOne == 0) { + this->quantOffsetGm.SetGlobalBuffer((__gm__ float*) quant_offset_gm, this->colNum); + this->quantScaleGm.SetGlobalBuffer((__gm__ float*) quant_scale_gm, this->colNum); + } else { + this->quantScaleGm.SetGlobalBuffer((__gm__ float*) quant_scale_gm, 1); + this->quant_scale = 1 / this->quantScaleGm.GetValue(0); + this->quantOffsetGm.SetGlobalBuffer((__gm__ float*) quant_offset_gm, 1); + this->quant_offset = this->quantOffsetGm.GetValue(0); + } + } + } + + __aicore__ inline void InitUbBufferCommon() + { + int64_t alignColNum = curColNum == Align(curColNum, sizeof(InType)) ? curColNum : Align(curColNum, sizeof(OutType)); + pipe->InitBuffer(inputTempBufferInt32SD, alignColNum * sizeof(CalcType) * NUM2); + pipe->InitBuffer(swigluTempBuffer, alignColNum * sizeof(CalcType)); + pipe->InitBuffer(inQueue, bufferNum, alignColNum * sizeof(InType) * NUM2); + pipe->InitBuffer(outQueue, bufferNum, alignColNum * sizeof(OutType)); + if (quantScaleIsEmpty == 0) { + if constexpr(quantIsOne == 0) { + pipe->InitBuffer(inQueueQuant, bufferNum, alignColNum * sizeof(float) * NUM2); + } + } + } + + __aicore__ inline void CopyInWeightAndBias(int64_t offset) + { + DataCopyExtParams params = {1, static_cast(curColNum * sizeof(float)), 0, 0, 0}; + DataCopyExtParams paramsBias = {1, static_cast(curColNum * sizeof(BiasType)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPadExtParams padParams1{false, 0, 0, 0}; + + LocalTensor weightLocal = inQueueWeightScale.template AllocTensor(); + LocalTensor biasTensorLocal; + if (this->biasIsEmpty == 0) { + biasTensorLocal = inQueueBias.template AllocTensor(); + } + + if (activateLeft == 0) { + DataCopyPad(weightLocal, weightScaleGm[offset + colNum], params, padParams); + DataCopyPad(weightLocal[alignColNum], weightScaleGm[offset], params, padParams); + + if (this->biasIsEmpty == 0) { + DataCopyPad(biasTensorLocal, biasGm[offset + colNum], paramsBias, padParams1); + if constexpr (std::is_same_v || std::is_same_v) { + DataCopyPad(biasTensorLocal[alignColNum], biasGm[offset], paramsBias, padParams1); + } else { + DataCopyPad(biasTensorLocal[biasAlignColNum], biasGm[offset], paramsBias, padParams1); + } + } + } else { + DataCopyPad(weightLocal, weightScaleGm[offset], params, padParams); + DataCopyPad(weightLocal[alignColNum], weightScaleGm[offset + colNum], params, padParams); + if (this->biasIsEmpty == 0) { + DataCopyPad(biasTensorLocal, biasGm[offset], paramsBias, padParams1); + if constexpr (std::is_same_v || std::is_same_v) { + DataCopyPad(biasTensorLocal[alignColNum], biasGm[offset + colNum], paramsBias, padParams1); + } else { + DataCopyPad(biasTensorLocal[biasAlignColNum], biasGm[offset + colNum], paramsBias, padParams1); + } + } + } + if (activateScaleIsEmpty == 0) { + DataCopyExtParams activateparams = {1, static_cast(curCoreRowNum * sizeof(float)), 0, 0, 0}; + LocalTensor activateLocal = inQueueActivationScale.template AllocTensor(); + DataCopyPad(activateLocal, activationScaleGm, activateparams, padParams); + inQueueActivationScale.EnQue(activateLocal); + } + inQueueWeightScale.EnQue(weightLocal); + if (this->biasIsEmpty == 0) { + inQueueBias.EnQue(biasTensorLocal); + } + } + + __aicore__ inline void CopyInQuant(int64_t offset) + { + DataCopyExtParams params = {1, static_cast(curColNum * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + LocalTensor quantLocal = inQueueQuant.template AllocTensor(); + DataCopyPad(quantLocal, quantScaleGm[offset], params, padParams); + DataCopyPad(quantLocal[alignColNum], quantOffsetGm[offset], params, padParams); + inQueueQuant.EnQue(quantLocal); + } + + __aicore__ inline void CopyIn(int64_t offset1, int64_t offset2) + { + DataCopyExtParams params = {1, static_cast(curColNum * sizeof(InType)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + + LocalTensor aLocal = inQueue.template AllocTensor(); + if (activateLeft == 0) { + DataCopyPad(aLocal, xGm[offset2], params, padParams); + DataCopyPad(aLocal[alignColNum], xGm[offset1], params, padParams); + } else { + DataCopyPad(aLocal, xGm[offset1], params, padParams); + DataCopyPad(aLocal[alignColNum], xGm[offset2], params, padParams); + } + inQueue.EnQue(aLocal); + } + + __aicore__ inline void dequant(uint64_t tileLen, uint64_t i) + { + LocalTensor aLocal = this->inQueue.template DeQue(); + this->inputTmpELocal = this->inputTempBufferInt32SD.template Get(); + if constexpr (std::is_same_v) { + if (this->biasIsEmpty == 0) { + Add(aLocal, aLocal, this->biasLocal, tileLen); + PipeBarrier(); + } + } + + Cast(this->inputTmpELocal, aLocal, RoundMode::CAST_NONE, tileLen); + PipeBarrier(); + this->inQueue.template FreeTensor(aLocal); + + Mul(this->inputTmpELocal, this->inputTmpELocal, this->weightScaleLocal, tileLen); + PipeBarrier(); + + if (this->activateScaleIsEmpty == 0) { + float value = this->activateLocal.GetValue(i); + Muls(this->inputTmpELocal, this->inputTmpELocal, value, tileLen); + PipeBarrier(); + } + + if constexpr (std::is_same_v || std::is_same_v || std::is_same_v) { + if (this->biasIsEmpty == 0) { + if constexpr (std::is_same_v) { + Add(this->inputTmpELocal, this->inputTmpELocal, this->biasLocal, tileLen); + } else { + LocalTensor biasFloatLocal = this->inputBiasTempBuffer.template Get(); + Cast(biasFloatLocal, this->biasLocal, RoundMode::CAST_NONE, tileLen / NUM2); + PipeBarrier(); + Cast(biasFloatLocal[tileLen / NUM2], this->biasLocal[biasAlignColNum], RoundMode::CAST_NONE, tileLen / NUM2); + PipeBarrier(); + Add(this->inputTmpELocal, this->inputTmpELocal, biasFloatLocal, tileLen); + } + PipeBarrier(); + } + } + } + + __aicore__ inline void processComputeFree() + { + if (this->biasIsEmpty == 0) { + this->inQueueBias.FreeTensor(this->biasLocal); + } + this->inQueueWeightScale.FreeTensor(this->weightScaleLocal); + if (quantScaleIsEmpty == 0) { + if constexpr(quantIsOne == 0) { + this->inQueueQuant.FreeTensor(this->quantLocal); + } + } + if (this->activateScaleIsEmpty == 0) { + this->inQueueActivationScale.template FreeTensor(this->activateLocal); + } + } + + __aicore__ inline void processCompute() + { + int64_t lastColNum = this->baseColLen; + int64_t colLoops = 1; + if (this->baseColLen < this->colNum) { + colLoops = (this->colNum + this->baseColLen - 1) / this->baseColLen; + lastColNum = this->colNum - (colLoops - 1) * this->baseColLen; + } + for (int64_t colLoop = 0; colLoop < colLoops; colLoop++) { + if (colLoop == colLoops - 1) { + this->curColNum = lastColNum; + } + bool isAligned = this->curColNum == this->Align(this->curColNum, sizeof(InType)); + this->alignColNum = isAligned ? this->curColNum : this->Align(this->curColNum, sizeof(int8_t)); + if constexpr (std::is_same_v || std::is_same_v) { + bool biasIsAligned = this->curColNum == this->Align(this->curColNum, sizeof(BiasType)); + this->biasAlignColNum = biasIsAligned ? this->curColNum : this->Align(this->curColNum, sizeof(int8_t)); + } + this->CopyInWeightAndBias(colLoop * this->baseColLen); + if (this->biasIsEmpty == 0) { + this->biasLocal = this->inQueueBias.template DeQue(); + } + this->weightScaleLocal = this->inQueueWeightScale.template DeQue(); + if (this->activateScaleIsEmpty == 0) { + this->activateLocal = this->inQueueActivationScale.template DeQue(); + } + for (int64_t i = 0; i < this->curCoreRowNum; i++) { + this->CopyIn(i * this->colNum * NUM2 + colLoop * this->baseColLen, i * this->colNum * NUM2 + this->colNum + colLoop * this->baseColLen); + this->dequant(this->alignColNum * NUM2, i); + if (i == 0 && quantScaleIsEmpty == 0) { + if constexpr(quantIsOne == 0) { + this->CopyInQuant(colLoop * this->baseColLen); + this->quantLocal = this->inQueueQuant.template DeQue(); + } + } + this->swiglu(this->alignColNum, i); + this->CopyOut(colLoop, i); + } + processComputeFree(); + } + } + + __aicore__ inline void swiglu(uint64_t curTileLen, int64_t idx) + { + LocalTensor swigluLocal = swigluTempBuffer.Get(); + Muls(swigluLocal, inputTmpELocal, beta, curTileLen); + PipeBarrier(); + Exp(swigluLocal, swigluLocal, curTileLen); + PipeBarrier(); + Adds(swigluLocal, swigluLocal, CalcType(1.0), curTileLen); + PipeBarrier(); + Div(inputTmpELocal, inputTmpELocal, swigluLocal, curTileLen); + PipeBarrier(); + Mul(inputTmpELocal[curTileLen], inputTmpELocal, inputTmpELocal[curTileLen], curTileLen); + PipeBarrier(); + if (quantScaleIsEmpty == 0) { + if constexpr(quantIsOne == 0) { + Div(inputTmpELocal[curTileLen], inputTmpELocal[curTileLen], quantLocal, curTileLen); + PipeBarrier(); + Add(inputTmpELocal[curTileLen], inputTmpELocal[curTileLen], quantLocal[curTileLen], curTileLen); + PipeBarrier(); + } else { + Muls(inputTmpELocal[curTileLen], inputTmpELocal[curTileLen], quant_scale, curTileLen); + PipeBarrier(); + Adds(inputTmpELocal[curTileLen], inputTmpELocal[curTileLen], quant_offset, curTileLen); + PipeBarrier(); + } + } + // fp32->int16 + LocalTensor int16Local = swigluTempBuffer.Get(); + Cast(int16Local, inputTmpELocal[curTileLen], RoundMode::CAST_RINT, curTileLen); + PipeBarrier(); + // int16-> half + LocalTensor halfLocal = int16Local.ReinterpretCast(); + Cast(halfLocal, int16Local, RoundMode::CAST_NONE, curTileLen); + PipeBarrier(); + + LocalTensor outLocal = outQueue.template AllocTensor(); + // half -> int8 + Cast(outLocal, halfLocal, RoundMode::CAST_NONE, curTileLen); + outQueue.template EnQue(outLocal); + } + + __aicore__ inline void CopyOut(int64_t colLoop, int64_t idx) + { + LocalTensor outLocal = outQueue.template DeQue(); + DataCopyExtParams dataCopyParams{1, static_cast(curColNum * sizeof(OutType)), 0, 0, 0}; + DataCopyPad(yGm[idx * colNum + colLoop * baseColLen], outLocal, dataCopyParams); + outQueue.FreeTensor(outLocal); + } + +protected: + __aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes) + { + constexpr int64_t BLOCK_BYTES = 32; + if (bytes == 0) { + return 0; + } + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes; + } + +protected: + float beta = -1.0; + float quant_scale = 1; + float quant_offset = 1; + TPipe* pipe = nullptr; + int64_t biasIsEmpty = 0; + int64_t quantScaleIsEmpty = 0; + int64_t activateScaleIsEmpty = 0; + int64_t colNum = 0; + int64_t rowNum = 0; + int64_t curCoreRowNum = 0; + int64_t inputCopyOffset = 0; + int64_t alignColNum = 0; + int64_t biasAlignColNum = 0; + int64_t curColNum = 0; + int64_t activateLeft = 0; + int64_t blockIdx = 0; + int64_t usedCoreNum = 0; + int64_t baseRowLen = 0; + int64_t baseColLen = 0; + + GlobalTensor yGm; + GlobalTensor xGm; + GlobalTensor weightScaleGm; + GlobalTensor activationScaleGm; + GlobalTensor biasGm; + GlobalTensor quantScaleGm; + GlobalTensor quantOffsetGm; + + LocalTensor inputTmpELocal; + LocalTensor weightScaleLocal; + LocalTensor biasLocal; + LocalTensor quantLocal; + LocalTensor activateLocal; + + TQue inQueueWeightScale; + TQue inQueueActivationScale; + TQue inQueueBias; + TQue inQueueQuant; + TQue inQueue; + TQue outQueue; + + TBuf inputTempBufferInt32SD; + TBuf swigluTempBuffer; + TBuf inputBiasTempBuffer; +}; +} + +#endif // CANN_DEQUANT_SWIGLU_QUANT_STATIC_BASE_HPP diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_static_bf16.hpp b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_static_bf16.hpp new file mode 100644 index 000000000..585bfe1d5 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_static_bf16.hpp @@ -0,0 +1,98 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_static_bf16.hpp + * \brief + */ + +#ifndef CANN_DEQUANT_SWIGLU_QUANT_STATIC_BF16_HPP +#define CANN_DEQUANT_SWIGLU_QUANT_STATIC_BF16_HPP +#include "kernel_operator.h" +#include "dequant_swiglu_quant_static_base.hpp" +namespace DequantSwigluQuant { +using namespace AscendC; + +TEMPLATE_DECLARE_STATIC +class DequantSwigluQuantStaticBF16 : public DequantSwigluQuantStaticBase { +public: + __aicore__ inline DequantSwigluQuantStaticBF16() {} + __aicore__ inline ~DequantSwigluQuantStaticBF16() {} + __aicore__ inline void Init(GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, + GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + const SwiGluTilingData* tilingData, TPipe* pipe_); + __aicore__ inline void Process(); + +protected: + __aicore__ inline void convertFloat(uint64_t curTileLen, uint64_t i); +}; + +TEMPLATE_DECLARE_STATIC +__aicore__ inline void DequantSwigluQuantStaticBF16::Init(GM_ADDR x_gm, GM_ADDR weight_scale_gm, + GM_ADDR activation_scale_gm, GM_ADDR bias_gm, + GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + const SwiGluTilingData* tilingData, TPipe* pipe_) +{ + this->pipe = pipe_; + this->InitCommon(x_gm, weight_scale_gm, activation_scale_gm, bias_gm, quant_scale_gm, quant_offset_gm, y_gm, scale_gm, tilingData, pipe_); + this->InitUbBufferCommon(); +} + +TEMPLATE_DECLARE_STATIC +__aicore__ inline void DequantSwigluQuantStaticBF16::Process() +{ + if (this->blockIdx >= this->usedCoreNum) { + return; + } + int64_t colLoops = 1; + int64_t lastColNum = this->baseColLen; + if (this->baseColLen < this->colNum) { + colLoops = (this->colNum + this->baseColLen - 1) / this->baseColLen; + lastColNum = this->colNum - (colLoops - 1) * this->baseColLen; + } + for (int64_t colLoop = 0; colLoop < colLoops; colLoop++) { + if (colLoop == colLoops - 1) { + this->curColNum = lastColNum; + } + bool isOutAligned = this->curColNum == this->Align(this->curColNum, sizeof(InType)); + this->alignColNum = isOutAligned ? this->curColNum : this->Align(this->curColNum, sizeof(OutType)); + for (int64_t i = 0; i < this->curCoreRowNum; i++) { + this->CopyIn(i * this->colNum * NUM2 + colLoop * this->baseColLen, i * this->colNum * NUM2 + this->colNum + colLoop * this->baseColLen); + convertFloat(this->alignColNum * NUM2, i); + if (i == 0 && this->quantScaleIsEmpty == 0) { + if constexpr(quantIsOne == 0) { + this->CopyInQuant(colLoop * this->baseColLen); + this->quantLocal = this->inQueueQuant.template DeQue(); + } + } + this->swiglu(this->alignColNum, i); + this->CopyOut(colLoop, i); + } + if (this->quantScaleIsEmpty == 0) { + if constexpr(quantIsOne == 0) { + this->inQueueQuant.FreeTensor(this->quantLocal); + } + } + } +} + +TEMPLATE_DECLARE_STATIC +__aicore__ inline void DequantSwigluQuantStaticBF16::convertFloat(uint64_t tileLen, uint64_t i) +{ + LocalTensor aLocal = this->inQueue.template DeQue(); + this->inputTmpELocal = this->inputTempBufferInt32SD.template Get(); + + Cast(this->inputTmpELocal, aLocal, RoundMode::CAST_NONE, tileLen); + PipeBarrier(); + this->inQueue.template FreeTensor(aLocal); +} +} + +#endif // CANN_DEQUANT_SWIGLU_QUANT_STATIC_BF16_HPP diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_static_bias_float.hpp b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_static_bias_float.hpp new file mode 100644 index 000000000..6875ec284 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_static_bias_float.hpp @@ -0,0 +1,93 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_static_bias_float.hpp + * \brief + */ + +#ifndef CANN_DEQUANT_SWIGLU_QUANT_STATIC_BIAS_FLOAT_HPP +#define CANN_DEQUANT_SWIGLU_QUANT_STATIC_BIAS_FLOAT_HPP +#include "kernel_operator.h" +#include "dequant_swiglu_quant_static_base.hpp" +namespace DequantSwigluQuant { + +using namespace AscendC; + +TEMPLATE_DECLARE_STATIC +class DequantSwigluQuantStaticBiasFloat : public DequantSwigluQuantStaticBase { +public: + __aicore__ inline DequantSwigluQuantStaticBiasFloat() {} + __aicore__ inline ~DequantSwigluQuantStaticBiasFloat() {} + + __aicore__ inline void Init(GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, + GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + const SwiGluTilingData* tilingData, TPipe* pipe_); + __aicore__ inline void Process(); + +protected: + __aicore__ inline void InitUbBuffer(); + +private: +}; + +TEMPLATE_DECLARE_STATIC +__aicore__ inline void DequantSwigluQuantStaticBiasFloat::Init(GM_ADDR x_gm, GM_ADDR weight_scale_gm, + GM_ADDR activation_scale_gm, GM_ADDR bias_gm, GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, + GM_ADDR y_gm, GM_ADDR scale_gm, const SwiGluTilingData* tilingData, TPipe* pipe_) +{ + this->pipe = pipe_; + this->InitCommon(x_gm, weight_scale_gm, activation_scale_gm, bias_gm, quant_scale_gm, quant_offset_gm, y_gm, scale_gm, tilingData, pipe_); + if (this->biasIsEmpty == 0) { + this->biasGm.SetGlobalBuffer((__gm__ BiasType*)bias_gm, this->colNum); + } + if (this->activateScaleIsEmpty == 0) { + this->activationScaleGm.SetGlobalBuffer((__gm__ float*) activation_scale_gm + this->inputCopyOffset, + this->curCoreRowNum); + } + this->weightScaleGm.SetGlobalBuffer((__gm__ float*) weight_scale_gm, this->colNum); + + this->InitUbBufferCommon(); + + this->InitUbBuffer(); +} + +TEMPLATE_DECLARE_STATIC +__aicore__ inline void DequantSwigluQuantStaticBiasFloat::Process() +{ + if (this->blockIdx >= this->usedCoreNum) { + return; + } + this->processCompute(); +} +TEMPLATE_DECLARE_STATIC +__aicore__ inline void DequantSwigluQuantStaticBiasFloat::InitUbBuffer() +{ + // pipe alloc memory to queue, the unit is Bytes + int64_t alignNumCol = this->curColNum == this->Align(this->curColNum, sizeof(InType)) + ? this->curColNum + : this->Align(this->curColNum, sizeof(OutType)); + this->pipe->InitBuffer(this->inQueueWeightScale, bufferNum, alignNumCol * sizeof(float) * NUM2); + if (this->activateScaleIsEmpty == 0) { + this->pipe->InitBuffer(this->inQueueActivationScale, bufferNum, this->curCoreRowNum * sizeof(float)); + } + if (this->biasIsEmpty == 0) { + int64_t biasAlignNumCol = this->curColNum == this->Align(this->curColNum, sizeof(BiasType)) + ? this->curColNum + : this->Align(this->curColNum, sizeof(OutType)); + this->pipe->InitBuffer(this->inQueueBias, bufferNum, biasAlignNumCol * sizeof(BiasType) * NUM2); + if constexpr (std::is_same_v || std::is_same_v) { + this->pipe->InitBuffer(this->inputBiasTempBuffer, alignNumCol * sizeof(float) * NUM2); + } + } +} +} // using namespace DequantSwigluQuant + +#endif // CANN_DEQUANT_SWIGLU_QUANT_STATIC_BIAS_FLOAT_HPP diff --git a/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_static_bias_int32.hpp b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_static_bias_int32.hpp new file mode 100644 index 000000000..05463526e --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_static_bias_int32.hpp @@ -0,0 +1,84 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_static_bias_int32.hpp + * \brief + */ + +#ifndef CANN_DEQUANT_SWIGLU_QUANT_STATIC_BIAS_INT32_HPP +#define CANN_DEQUANT_SWIGLU_QUANT_STATIC_BIAS_INT32_HPP +#include "kernel_operator.h" +#include "dequant_swiglu_quant_static_base.hpp" +namespace DequantSwigluQuant { +using namespace AscendC; + +TEMPLATE_DECLARE_STATIC +class DequantSwigluQuantStaticBiasInt32 : public DequantSwigluQuantStaticBase { +public: + __aicore__ inline DequantSwigluQuantStaticBiasInt32() {} + __aicore__ inline ~DequantSwigluQuantStaticBiasInt32() {} + __aicore__ inline void Init(GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, GM_ADDR bias_gm, + GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + const SwiGluTilingData* tilingData, TPipe* pipe_); + __aicore__ inline void Process(); + +protected: + __aicore__ inline void InitUbBuffer(); +private: +}; + +TEMPLATE_DECLARE_STATIC +__aicore__ inline void DequantSwigluQuantStaticBiasInt32::Init(GM_ADDR x_gm, GM_ADDR weight_scale_gm, + GM_ADDR activation_scale_gm, GM_ADDR bias_gm, + GM_ADDR quant_scale_gm, GM_ADDR quant_offset_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + const SwiGluTilingData* tilingData, TPipe* pipe_) +{ + this->pipe = pipe_; + this->InitCommon(x_gm, weight_scale_gm, activation_scale_gm, bias_gm, quant_scale_gm, quant_offset_gm, y_gm, scale_gm, tilingData, pipe_); + + if (this->activateScaleIsEmpty == 0) { + this->activationScaleGm.SetGlobalBuffer((__gm__ float*) activation_scale_gm + this->inputCopyOffset, + this->curCoreRowNum); + } + this->weightScaleGm.SetGlobalBuffer((__gm__ float*) weight_scale_gm, this->colNum); + if (tilingData->biasIsEmpty == 0) { + this->biasGm.SetGlobalBuffer((__gm__ BiasType*)bias_gm, this->colNum); + } + this->InitUbBufferCommon(); + this->InitUbBuffer(); +} + +TEMPLATE_DECLARE_STATIC +__aicore__ inline void DequantSwigluQuantStaticBiasInt32::Process() +{ + if (this->blockIdx >= this->usedCoreNum) { + return; + } + this->processCompute(); +} +TEMPLATE_DECLARE_STATIC +__aicore__ inline void DequantSwigluQuantStaticBiasInt32::InitUbBuffer() +{ + int64_t alignColNumber = this->curColNum == this->Align(this->curColNum, sizeof(InType)) + ? this->curColNum + : this->Align(this->curColNum, sizeof(OutType)); + + this->pipe->InitBuffer(this->inQueueWeightScale, bufferNum, alignColNumber * sizeof(float) * NUM2); + if (this->activateScaleIsEmpty == 0) { + this->pipe->InitBuffer(this->inQueueActivationScale, bufferNum, this->curCoreRowNum * sizeof(float)); + } + if (this->biasIsEmpty == 0) { + this->pipe->InitBuffer(this->inQueueBias, bufferNum, alignColNumber * sizeof(float) * NUM2); + } +} + +} +#endif // CANN_DEQUANT_SWIGLU_QUANT_STATIC_BIAS_INT32_HPP diff --git a/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/error_log.h b/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/error_log.h new file mode 100644 index 000000000..1b09d2022 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/error_log.h @@ -0,0 +1,63 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#include "log/log.h" + +#ifndef OP_LOGE_FOR_INVALID_DTYPE +#define OP_LOGE_FOR_INVALID_DTYPE(opname, param, actual, expected) \ + OP_LOGE(opname, "Invalid dtype for %s, actual: %s, expected: %s", param, actual, expected) +#endif + +#ifndef OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON +#define OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid dtype for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPE +#define OP_LOGE_FOR_INVALID_SHAPE(opname, param, actual, expected) \ + OP_LOGE(opname, "Invalid shape for %s, actual: %s, expected: %s", param, actual, expected) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON +#define OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid shape for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON +#define OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid shapes for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPEDIM +#define OP_LOGE_FOR_INVALID_SHAPEDIM(opname, param, actual, expected) \ + OP_LOGE(opname, "Invalid shape dim for %s, actual: %s, expected: %s", param, actual, expected) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON +#define OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid shape dims for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPESIZE +#define OP_LOGE_FOR_INVALID_SHAPESIZE(opname, param, actual, expected) \ + OP_LOGE(opname, "Invalid shape size for %s, actual: %s, expected: %s", param, actual, expected) +#endif + +#ifndef OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON +#define OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid shape size for %s, actual: %s, reason: %s", param, actual, reason) +#endif + +#ifndef OP_LOGE_FOR_INVALID_VALUE_WITH_REASON +#define OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(opname, param, actual, reason) \ + OP_LOGE(opname, "Invalid value for %s, actual: %s, reason: %s", param, actual, reason) +#endif diff --git a/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/static_register_symbol.h b/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/static_register_symbol.h new file mode 100644 index 000000000..4228c142c --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/static_register_symbol.h @@ -0,0 +1,35 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file static_register_symbol.h + * \brief + */ + +#pragma once + +#include + +#define GLOBAL_REGISTER_SYMBOL_REAL(op_type, class_name, priority, counter, line) \ +[[maybe_unused]] std::string op_impl_register_template_##op_type##_##class_name##priority##counter##line = \ + std::string("op_impl_register_template_" #op_type) \ + + +#define GLOBAL_REGISTER_SYMBOL(op_type, class_name, priority, counter, line) \ +GLOBAL_REGISTER_SYMBOL_REAL(op_type, class_name, priority, counter, line) + + +#define GLOBAL_REGISTER_STR_SYMBOL_REAL(op_type, class_name, priority, counter, line) \ +[[maybe_unused]] std::string op_impl_register_template_##class_name##priority##counter##line = \ + std::string("op_impl_register_template_" op_type) \ + + +#define GLOBAL_REGISTER_STR_SYMBOL(op_type, class_name, priority, counter, line) \ +GLOBAL_REGISTER_STR_SYMBOL_REAL(op_type, class_name, priority, counter, line) diff --git a/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/tiling_base.h b/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/tiling_base.h new file mode 100644 index 000000000..d6bc7e4d0 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/tiling_base.h @@ -0,0 +1,238 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_base.h + * \brief + */ + +#pragma once + +#include +#include +#include +#include +#include "tiling/platform/platform_ascendc.h" +#include "platform/soc_spec.h" +#include "log/log.h" +#include "error_log.h" + +#ifdef ASCENDC_OP_TEST +#define ASCENDC_EXTERN_C extern "C" +#else +#define ASCENDC_EXTERN_C +#endif + +namespace Ops { +namespace NN { +namespace Optiling { + +struct AiCoreParams { + uint64_t ubSize = 0UL; + uint64_t blockDim = 0UL; + uint64_t numBlocks = 0UL; + uint64_t aicNum = 0UL; + uint64_t l1Size = 0UL; + uint64_t l0aSize = 0UL; + uint64_t l0bSize = 0UL; + uint64_t l0cSize = 0UL; +}; + +struct CompileInfoCommon { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; + int32_t socVersion; + uint32_t rsvd; +}; + +class TilingBaseClass +{ +public: + explicit TilingBaseClass(gert::TilingContext* context) : context_(context) + {} + + virtual ~TilingBaseClass() = default; + + // Tiling执行框架 + // 1、GRAPH_SUCCESS: 成功，并且不需要继续执行后续Tiling类的实现 + // 2、GRAPH_FAILED: 失败，中止整个Tiling流程 + // 3、GRAPH_PARAM_INVALID: 本类不支持，需要继续往下执行其他Tiling类的实现 + ge::graphStatus DoTiling() + { + auto ret = GetShapeAttrsInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetPlatformInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + if (!IsCapable()) { + return ge::GRAPH_PARAM_INVALID; + } + ret = DoOpTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = DoLibApiTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetWorkspaceSize(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = PostTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + context_->SetTilingKey(GetTilingKey()); + DumpTilingInfo(); + return ge::GRAPH_SUCCESS; + } + + // 更新 context + virtual void Reset(gert::TilingContext* context) + { + context_ = context; + } + +protected: + virtual bool IsCapable() = 0; + // 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小 + virtual ge::graphStatus GetPlatformInfo() = 0; + // 2、获取INPUT/OUTPUT/ATTR信息 + virtual ge::graphStatus GetShapeAttrsInfo() = 0; + // 3、计算数据切分TilingData + virtual ge::graphStatus DoOpTiling() = 0; + // 4、计算高阶API的TilingData + virtual ge::graphStatus DoLibApiTiling() = 0; + // 5、计算TilingKey + [[nodiscard]] virtual uint64_t GetTilingKey() const = 0; + // 6、计算Workspace 大小 + virtual ge::graphStatus GetWorkspaceSize() = 0; + // 7、保存Tiling数据 + virtual ge::graphStatus PostTiling() = 0; + // 8、Dump Tiling数据 + virtual void DumpTilingInfo() + { + int32_t enable = CheckLogLevel(static_cast(OP), DLOG_DEBUG); + if (enable != 1) { + return; + } + auto buf = (uint32_t*)context_->GetRawTilingData()->GetData(); + auto bufLen = context_->GetRawTilingData()->GetDataSize(); + std::ostringstream oss; + oss << "Start to dump tiling info. tilingkey:" << context_->GetTilingKey() << ", tiling data size:" << bufLen + << ", content:"; + for (size_t i = 0; i < bufLen / sizeof(uint32_t); i++) { + oss << *(buf + i) << ","; + if (oss.str().length() > 640) { // Split according to 640 to avoid truncation + OP_LOGD(context_, "%s", oss.str().c_str()); + oss.str(""); + } + } + OP_LOGD(context_, "%s", oss.str().c_str()); + } + + static uint32_t CalcTschBlockDim(uint32_t sliceNum, uint32_t aicCoreNum, uint32_t aivCoreNum) + { + uint32_t ration; + if (aicCoreNum == 0 || aivCoreNum == 0 || aicCoreNum > aivCoreNum) { + return sliceNum; + } + ration = aivCoreNum / aicCoreNum; + return (sliceNum + (ration - 1)) / ration; + } + + template + [[nodiscard]] std::string GetShapeDebugStr(const T& shape) const + { + std::ostringstream oss; + oss << "["; + if (shape.GetDimNum() > 0) { + for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) { + oss << shape.GetDim(i) << ", "; + } + oss << shape.GetDim(shape.GetDimNum() - 1); + } + oss << "]"; + return oss.str(); + } + + [[nodiscard]] std::string GetTensorDebugStr( + const gert::StorageShape* shape, const gert::CompileTimeTensorDesc* tensor) const + { + if (shape == nullptr || tensor == nullptr) { + return "nil "; + } + std::ostringstream oss; + oss << "(dtype: " << ge::TypeUtils::DataTypeToSerialString(tensor->GetDataType()) << "),"; + oss << "(shape:" << GetShapeDebugStr(shape->GetStorageShape()) << "),"; + oss << "(ori_shape:" << GetShapeDebugStr(shape->GetOriginShape()) << "),"; + oss << "(format: " + << ge::TypeUtils::FormatToSerialString( + static_cast(ge::GetPrimaryFormat(tensor->GetStorageFormat()))) + << "),"; + oss << "(ori_format: " << ge::TypeUtils::FormatToSerialString(tensor->GetOriginFormat()) << ") "; + return oss.str(); + } + + [[nodiscard]] std::string GetTilingContextDebugStr() + { + std::ostringstream oss; + for (size_t i = 0; i < context_->GetComputeNodeInfo()->GetInputsNum(); ++i) { + oss << "input" << i << ": "; + oss << GetTensorDebugStr(context_->GetInputShape(i), context_->GetInputDesc(i)); + } + + for (size_t i = 0; i < context_->GetComputeNodeInfo()->GetOutputsNum(); ++i) { + oss << "output" << i << ": "; + oss << GetTensorDebugStr(context_->GetOutputShape(i), context_->GetOutputDesc(i)); + } + return oss.str(); + } + + [[nodiscard]] std::string GetTilingDataDebugStr() const + { + auto rawTilingData = context_->GetRawTilingData(); + auto rawTilingDataSize = rawTilingData->GetDataSize(); + auto data = reinterpret_cast(rawTilingData->GetData()); + size_t len = rawTilingDataSize / sizeof(int32_t); + std::ostringstream oss; + for (size_t i = 0; i < len; i++) { + oss << data[i] << ", "; + } + return oss.str(); + } + +protected: + gert::TilingContext* context_ = nullptr; + std::unique_ptr ascendcPlatform_{nullptr}; + uint32_t blockDim_{0}; + uint64_t workspaceSize_{0}; + uint64_t tilingKey_{0}; + AiCoreParams aicoreParams_; +}; + +} // namespace Optiling +} // namespace NN +} // namespace Ops + +namespace optiling { +using Ops::NN::Optiling::TilingBaseClass; +} // namespace optiling diff --git a/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/tiling_key.h b/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/tiling_key.h new file mode 100644 index 000000000..0557a7ed0 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/tiling_key.h @@ -0,0 +1,63 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_key.h + * \brief + */ + +#pragma once + +#include + +namespace Ops { +namespace NN { +namespace Optiling { +constexpr uint64_t RecursiveSum() +{ + return 0; +} + +template constexpr uint64_t RecursiveSum(T templateId, Args... templateIds) +{ + const int carryCoefficient = 10; //进位系数 + return static_cast(templateId) + carryCoefficient * RecursiveSum(templateIds...); +} + +// TilingKey 的生成规则： +// FlashAttentionScore/FlashAttentionScoreGrad 十进制位组装tiling key，包含以下关键参数，从低位到高位依次是：Ub0, Ub1, +// Block, DataType, Format, Sparse, 特化模板 Ub0、Ub1: +// 表示Ub核内切分的轴，使用枚举AxisEnum表示，因为我们允许最多切分两根轴，所以存在UB0和UB1，如果没有UB核内切分， +// 那么填AXIS_NONE。UB0和UB1各占一个十进制位; +// Block: 表示UB用来分核的轴，使用枚举AxisEnum表示，占一个十进制位; +// DataType: 表示当前tiling key支持的输入输出的数据类型，使用枚举SupportedDtype来表示，占一个十进制位 +// Format: 表示当前tiling key支持的Format, 使用枚举InputLayout表示，占一个十进制位 +// Sparse: 表示当前tiling key是否支持Sparse，使用枚举SparseCapability表示，占一个十进制位 +// 其余特化场景，定义自己的位域和值 +// usage: get tilingKey from inputted types +// uint64_t tilingKey = GET_FLASHATTENTION_TILINGKEY(AxisEnum::AXIS_S1, AxisEnum::AXIS_S2, AxisEnum::AXIS_N2, +// SupportedDtype::FLOAT32, InputLayout::BSH, SparseCapability::SUPPORT_ALL) + +constexpr uint64_t TILINGKEYOFFSET = uint64_t(10000000000000000000UL); // 10^19 +template constexpr uint64_t GET_TILINGKEY(Args... templateIds) +{ + return TILINGKEYOFFSET + RecursiveSum(templateIds...); +} + +// usage: get tilingKey from inputted types +// uint64_t tilingKey = TILINGKEY(S2, S1, N2, FLOAT32, BSND, ALL) + +#define TILINGKEY(ub2, ub1, block, dtype, layout, sparse) \ + (GET_TILINGKEY(AxisEnum::ub2, AxisEnum::ub1, AxisEnum::block, DtypeEnum::dtype, LayoutEnum::layout, \ + SparseEnum::sparse)) + +} // namespace Optiling +} // namespace NN +} // namespace Ops diff --git a/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/tiling_templates_registry.h b/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/tiling_templates_registry.h new file mode 100644 index 000000000..aacd06f0d --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/tiling_templates_registry.h @@ -0,0 +1,486 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_templates_registry.h + * \brief + */ + +#pragma once + +#include +#include +#include +#include +#include "exe_graph/runtime/tiling_context.h" +#include "tiling_base.h" +#include "static_register_symbol.h" +#include "log/log.h" + +namespace Ops { +namespace NN { +namespace Optiling { + +template +std::unique_ptr TILING_CLASS(gert::TilingContext* context) +{ + return std::unique_ptr(new (std::nothrow) T(context)); +} + +using TilingClassCase = std::unique_ptr (*)(gert::TilingContext*); + +class TilingCases +{ +public: + explicit TilingCases(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + void AddTiling(int32_t priority) + { + OP_CHECK_IF( + cases_.find(priority) != cases_.end(), OP_LOGE(op_type_, "There are duplicate registrations."), return); + cases_[priority] = TILING_CLASS; + OP_CHECK_IF( + cases_[priority] == nullptr, + OP_LOGE(op_type_, "Register op tiling func failed, please check the class name."), return); + } + + const std::map& GetTilingCases() + { + return cases_; + } + +private: + std::map cases_; + const std::string op_type_; +}; + +// --------------------------------Interfacce with npu arch -------------------------------- +class TilingRegistryArch { +public: + TilingRegistryArch() = default; + +#ifdef ASCENDC_OP_TEST + static TilingRegistryArch& GetInstance(); +#else + static TilingRegistryArch& GetInstance() + { + static TilingRegistryArch registryImpl; + return registryImpl; + } +#endif + + std::shared_ptr RegisterOp(const std::string& opType, int32_t arch) + { + auto archIter = registryMap_.find(arch); + if (archIter == registryMap_.end()) { + std::map> opTypeMap; + opTypeMap[opType] = std::shared_ptr(new (std::nothrow) TilingCases(opType)); + registryMap_[arch] = opTypeMap; + } else { + if (archIter->second.find(opType) == archIter->second.end()) { + archIter->second[opType] = std::shared_ptr(new (std::nothrow) TilingCases(opType)); + } + } + + OP_CHECK_IF(registryMap_[arch][opType] == nullptr, + OP_LOGE(opType, "Register tiling func failed, please check the class name."), return nullptr); + return registryMap_[arch][opType]; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context) + { + int32_t arch = (int32_t)NpuArch::DAV_RESV; + const char* opType = context->GetNodeType(); + fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + OP_LOGE(opType, "Do op tiling failed, cannot get platformInfo."); + return ge::GRAPH_FAILED; + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + arch = static_cast(ascendcPlatform.GetCurNpuArch()); + OP_LOGD(context, "npu arch is %d", arch); + if (arch == (int32_t)NpuArch::DAV_RESV) { + OP_LOGE(opType, "Do op tiling failed, cannot find npu arch."); + return ge::GRAPH_FAILED; + } + } + auto tilingTemplateRegistryMap = GetTilingTemplates(opType, arch); + for (auto it = tilingTemplateRegistryMap.begin(); it != tilingTemplateRegistryMap.end(); ++it) { + auto tilingTemplate = it->second(context); + if (tilingTemplate != nullptr) { + ge::graphStatus status = tilingTemplate->DoTiling(); + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do general op tiling success priority=%d", it->first); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", it->first); + } + } + OP_LOGE(opType, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + const std::map& GetTilingTemplates(const std::string& opType, int32_t arch) + { + auto archIter = registryMap_.find(arch); + OP_CHECK_IF(archIter == registryMap_.end(), + OP_LOGE(opType, "Get op tiling func failed, please check the npu arch %d", arch), + return emptyTilingCase_); + auto opIter = archIter->second.find(opType); + OP_CHECK_IF( + opIter == archIter->second.end(), OP_LOGE(opType, "Get op tiling func failed, please check the op name."), + return emptyTilingCase_); + return opIter->second->GetTilingCases(); + } + +private: + std::map>> registryMap_; // key is npu-arch + const std::map emptyTilingCase_{}; +}; + +class RegisterArch { +public: + explicit RegisterArch(std::string opType) : opType_(std::move(opType)) + {} + + template + RegisterArch& tiling(int32_t priority, int32_t arch) + { + auto tilingCases = TilingRegistryArch::GetInstance().RegisterOp(opType_, arch); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(opType_, "Register op tiling failed, please check the op name."), + return *this); + tilingCases->AddTiling(priority); + return *this; + } + + template + RegisterArch& tiling(int32_t priority, const std::vector& archs) + { + for (int32_t arch : archs) { + auto tilingCases = TilingRegistryArch::GetInstance().RegisterOp(opType_, arch); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(opType_, "Register op tiling failed, please check the op name."), + return *this); + tilingCases->AddTiling(priority); + } + return *this; + } + +private: + const std::string opType_; +}; + +// --------------------------------Interfacce with soc version -------------------------------- +class TilingRegistryNew +{ +public: + TilingRegistryNew() = default; + +#ifdef ASCENDC_OP_TEST + static TilingRegistryNew& GetInstance(); +#else + static TilingRegistryNew& GetInstance() + { + static TilingRegistryNew registry_impl_; + return registry_impl_; + } +#endif + + std::shared_ptr RegisterOp(const std::string& op_type, int32_t soc_version) + { + auto soc_iter = registry_map_.find(soc_version); + if (soc_iter == registry_map_.end()) { + std::map> op_type_map; + op_type_map[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + registry_map_[soc_version] = op_type_map; + } else { + if (soc_iter->second.find(op_type) == soc_iter->second.end()) { + soc_iter->second[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + } + } + + OP_CHECK_IF( + registry_map_[soc_version][op_type] == nullptr, + OP_LOGE(op_type, "Register tiling func failed, please check the class name."), return nullptr); + return registry_map_[soc_version][op_type]; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context) + { + int32_t soc_version = (int32_t)platform_ascendc::SocVersion::RESERVED_VERSION; + const char* op_type = context->GetNodeType(); + fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + auto compileInfoPtr = context->GetCompileInfo(); + OP_CHECK_IF( + compileInfoPtr == nullptr, OP_LOGE(op_type, "compileInfoPtr is null."), return ge::GRAPH_FAILED); + soc_version = compileInfoPtr->socVersion; + OP_LOGD(context, "soc version in compileInfo is %d", soc_version); + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + soc_version = static_cast(ascendcPlatform.GetSocVersion()); + OP_LOGD(context, "soc version is %d", soc_version); + if (soc_version == (int32_t)platform_ascendc::SocVersion::RESERVED_VERSION) { + OP_LOGE(op_type, "Do op tiling failed, cannot find soc version."); + return ge::GRAPH_FAILED; + } + } + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type, soc_version); + for (auto it = tilingTemplateRegistryMap.begin(); it != tilingTemplateRegistryMap.end(); ++it) { + auto tilingTemplate = it->second(context); + if (tilingTemplate != nullptr) { + ge::graphStatus status = tilingTemplate->DoTiling(); + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do general op tiling success priority=%d", it->first); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", it->first); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context, const std::vector& priorities) + { + int32_t soc_version; + const char* op_type = context->GetNodeType(); + auto platformInfoPtr = context->GetPlatformInfo(); + if (platformInfoPtr == nullptr) { + auto compileInfoPtr = context->GetCompileInfo(); + OP_CHECK_IF( + compileInfoPtr == nullptr, OP_LOGE(op_type, "compileInfoPtr is null."), return ge::GRAPH_FAILED); + soc_version = compileInfoPtr->socVersion; + OP_LOGD(context, "soc version in compileInfo is %d", soc_version); + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + soc_version = static_cast(ascendcPlatform.GetSocVersion()); + OP_LOGD(context, "soc version is %d", soc_version); + } + + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type, soc_version); + for (auto priority_id : priorities) { + auto tilingCaseIter = tilingTemplateRegistryMap.find(priority_id); + if (tilingCaseIter != tilingTemplateRegistryMap.end()) { + auto templateFunc = tilingCaseIter->second(context); + if (templateFunc != nullptr) { + ge::graphStatus status = templateFunc->DoTiling(); + if (status == ge::GRAPH_SUCCESS) { + OP_LOGD(context, "Do general op tiling success priority=%d", priority_id); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", priority_id); + } + } + } + return ge::GRAPH_FAILED; + } + + const std::map& GetTilingTemplates(const std::string& op_type, int32_t soc_version) + { + auto soc_iter = registry_map_.find(soc_version); + OP_CHECK_IF( + soc_iter == registry_map_.end(), + OP_LOGE(op_type, "Get op tiling func failed, please check the soc version %d", soc_version), + return empty_tiling_case_); + auto op_iter = soc_iter->second.find(op_type); + OP_CHECK_IF( + op_iter == soc_iter->second.end(), OP_LOGE(op_type, "Get op tiling func failed, please check the op name."), + return empty_tiling_case_); + return op_iter->second->GetTilingCases(); + } + +private: + std::map>> registry_map_; // key is socversion + const std::map empty_tiling_case_{}; +}; + +class RegisterNew +{ +public: + explicit RegisterNew(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + RegisterNew& tiling(int32_t priority, int32_t soc_version) + { + auto tilingCases = TilingRegistryNew::GetInstance().RegisterOp(op_type_, soc_version); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), return *this); + tilingCases->AddTiling(priority); + return *this; + } + + template + RegisterNew& tiling(int32_t priority, const std::vector& soc_versions) + { + for (int32_t soc_version : soc_versions) { + auto tilingCases = TilingRegistryNew::GetInstance().RegisterOp(op_type_, soc_version); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), + return *this); + tilingCases->AddTiling(priority); + } + return *this; + } + +private: + const std::string op_type_; +}; + +// --------------------------------Interfacce without soc version -------------------------------- +class TilingRegistry +{ +public: + TilingRegistry() = default; + +#ifdef ASCENDC_OP_TEST + static TilingRegistry& GetInstance(); +#else + static TilingRegistry& GetInstance() + { + static TilingRegistry registry_impl_; + return registry_impl_; + } +#endif + + std::shared_ptr RegisterOp(const std::string& op_type) + { + if (registry_map_.find(op_type) == registry_map_.end()) { + registry_map_[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + } + OP_CHECK_IF( + registry_map_[op_type] == nullptr, + OP_LOGE(op_type, "Register tiling func failed, please check the class name."), return nullptr); + return registry_map_[op_type]; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context) + { + const char* op_type = context->GetNodeType(); + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type); + for (auto it = tilingTemplateRegistryMap.begin(); it != tilingTemplateRegistryMap.end(); ++it) { + auto tilingTemplate = it->second(context); + if (tilingTemplate != nullptr) { + ge::graphStatus status = tilingTemplate->DoTiling(); + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do general op tiling success priority=%d", it->first); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", it->first); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext* context, const std::vector& priorities) + { + const char* op_type = context->GetNodeType(); + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type); + for (auto priorityId : priorities) { + auto templateFunc = tilingTemplateRegistryMap[priorityId](context); + if (templateFunc != nullptr) { + ge::graphStatus status = templateFunc->DoTiling(); + if (status == ge::GRAPH_SUCCESS) { + OP_LOGD(context, "Do general op tiling success priority=%d", priorityId); + return status; + } + if (status != ge::GRAPH_PARAM_INVALID) { + OP_LOGD(context, "Do op tiling failed"); + return status; + } + OP_LOGD(context, "Ignore general op tiling priority=%d", priorityId); + } + } + OP_LOGE(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + const std::map& GetTilingTemplates(const std::string& op_type) + { + OP_CHECK_IF( + registry_map_.find(op_type) == registry_map_.end(), + OP_LOGE(op_type, "Get op tiling func failed, please check the op name."), return empty_tiling_case_); + return registry_map_[op_type]->GetTilingCases(); + } + +private: + std::map> registry_map_; + const std::map empty_tiling_case_; +}; + +class Register +{ +public: + explicit Register(std::string op_type) : op_type_(std::move(op_type)) + {} + + template + Register& tiling(int32_t priority) + { + auto tilingCases = TilingRegistry::GetInstance().RegisterOp(op_type_); + OP_CHECK_IF( + tilingCases == nullptr, OP_LOGE(op_type_, "Register op tiling failed, please the op name."), return *this); + tilingCases->AddTiling(priority); + return *this; + } + +private: + const std::string op_type_; +}; +// op_type: 算子名称， class_name: 注册的 tiling 类, arch：芯片架构号 +// priority: tiling 类的优先级, 越小表示优先级越高, 即会优先选择这个tiling类 +#define REGISTER_TILING_TEMPLATE_WITH_ARCH(op_type, class_name, archs, priority) \ + [[maybe_unused]] uint32_t op_impl_register_template_##op_type##_##class_name##priority; \ + static Ops::NN::Optiling::RegisterArch VAR_UNUSED##op_type##class_name##priority_register = \ + Ops::NN::Optiling::RegisterArch(#op_type).tiling(priority, archs) + +// op_type: 算子名称， class_name: 注册的 tiling 类, soc_version：芯片版本号 +// priority: tiling 类的优先级, 越小表示优先级越高, 即会优先选择这个tiling类 +#define REGISTER_TILING_TEMPLATE_WITH_SOCVERSION(op_type, class_name, soc_versions, priority) \ + GLOBAL_REGISTER_SYMBOL(op_type, class_name, priority, __COUNTER__, __LINE__); \ + static Ops::NN::Optiling::RegisterNew VAR_UNUSED##op_type##class_name##priority_register = \ + Ops::NN::Optiling::RegisterNew(#op_type).tiling(priority, soc_versions) + +// op_type: 算子名称， class_name: 注册的 tiling 类, +// priority: tiling 类的优先级, 越小表示优先级越高, 即被选中的概率越大 +#define REGISTER_TILING_TEMPLATE(op_type, class_name, priority) \ + GLOBAL_REGISTER_STR_SYMBOL(op_type, class_name, priority, __COUNTER__, __LINE__); \ + static Ops::NN::Optiling::Register VAR_UNUSED##op_type_##class_name##priority_register = \ + Ops::NN::Optiling::Register(op_type).tiling(priority) + +// op_type: 算子名称， class_name: 注册的 tiling 类, +// soc_version: soc版本，用于区分不同的soc +// priority: tiling 类的优先级, 越小表示优先级越高, 即会优先选择这个tiling类 +#define REGISTER_TILING_TEMPLATE_NEW(op_type, class_name, soc_version, priority) \ + GLOBAL_REGISTER_SYMBOL(op_type, class_name, priority, __COUNTER__, __LINE__); \ + static Ops::NN::Optiling::RegisterNew VAR_UNUSED##op_type##class_name##priority_register = \ + Ops::NN::Optiling::RegisterNew(#op_type).tiling(priority, soc_version) + +// op_type: 算子名称， class_name: 注册的 tiling 类, +// priority: tiling 类的优先级, 越小表示优先级越高, 即被选中的概率越大 +// 取代 REGISTER_TILING_TEMPLATE , 传入的op_type如果是字符串常量，需要去掉引号 +#define REGISTER_OPS_TILING_TEMPLATE(op_type, class_name, priority) \ + GLOBAL_REGISTER_SYMBOL(op_type, class_name, priority, __COUNTER__, __LINE__); \ + static Ops::NN::Optiling::Register \ + __attribute__((unused)) tiling_##op_type##_##class_name##_##priority##_register = \ + Ops::NN::Optiling::Register(#op_type).tiling(priority) +} // namespace Optiling +} // namespace NN +} // namespace Ops + +namespace optiling { +using Ops::NN::Optiling::TilingRegistry; +using Ops::NN::Optiling::TilingRegistryNew; +} // namespace optiling diff --git a/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/tiling_util.h b/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/tiling_util.h new file mode 100644 index 000000000..d72b6cf02 --- /dev/null +++ b/csrc/ascend/moe/dequant_swiglu_quant/tiling_base/tiling_util.h @@ -0,0 +1,61 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_util.h + * \brief + */ + +#pragma once + +#include "register/op_impl_registry.h" +#include "platform/platform_ascendc.h" +#include "platform/soc_spec.h" +#include "log/log.h" + +namespace Ops { +namespace NN { +namespace OpTiling { +static const gert::Shape g_vec_1_shape = {1}; + +static bool IsRegbaseNpuArch(NpuArch npuArch) +{ + const static std::set regbaseNpuArchs = { + NpuArch::DAV_3510, + NpuArch::DAV_5102}; + return regbaseNpuArchs.find(npuArch) != regbaseNpuArchs.end(); +} + +static inline bool IsRegbaseSocVersion(const gert::TilingParseContext* context) +{ + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + auto npuArch = ascendcPlatform.GetCurNpuArch(); + OP_LOGI(context, "Current NpuArch is %u", static_cast(npuArch)); + return IsRegbaseNpuArch(npuArch); +} + +static inline bool IsRegbaseSocVersion(const gert::TilingContext* context) +{ + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + auto npuArch = ascendcPlatform.GetCurNpuArch(); + OP_LOGI(context, "Current NpuArch is %u", static_cast(npuArch)); + return IsRegbaseNpuArch(npuArch); +} + +inline const gert::Shape& EnsureNotScalar(const gert::Shape& inShape) +{ + if (inShape.IsScalar()) { + return g_vec_1_shape; + } + return inShape; +} +} // namespace OpTiling +} // namespace NN +} // namespace Ops \ No newline at end of file diff --git a/csrc/ascend/moe/hamming_dist_top_k/CMakeLists.txt b/csrc/ascend/moe/hamming_dist_top_k/CMakeLists.txt new file mode 100644 index 000000000..549b5c14c --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_host/CMakeLists.txt b/csrc/ascend/moe/hamming_dist_top_k/op_host/CMakeLists.txt new file mode 100644 index 000000000..68f90c0a1 --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_host/CMakeLists.txt @@ -0,0 +1,29 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + hamming_dist_top_k_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME HammingDistTopK + OPTIONS + --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE hamming_dist_top_k ACLNNTYPE aclnn) + target_sources(${OPHOST_NAME}_tiling_obj PRIVATE + hamming_dist_top_k_tiling.cpp + hamming_dist_top_k.cpp + hamming_dist_top_k_split.cpp + ) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k.cpp b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k.cpp new file mode 100644 index 000000000..8063bf339 --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k.cpp @@ -0,0 +1,213 @@ + +#include "hamming_dist_top_k_tiling.h" +#include "hamming_dist_top_k.h" +#include "register/op_def_registry.h" +#include +namespace optiling { + +namespace { + +} + +bool HammingDistTopKTiling::IsCapable() +{ + return true; +} + +ge::graphStatus HammingDistTopKTiling::GetPlatformInfo() { return ge::GRAPH_SUCCESS; } + +ge::graphStatus HammingDistTopKTiling::GetShapeAttrsInfo() { + inputParams_.opName = context_->GetNodeName(); + opName_ = context_->GetNodeName(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HammingDistTopKTiling::DoOpTiling() { + auto keyBlockTablePtr = context_->GetOptionalInputShape(KEY_BLOCK_TABLE_INPUT_INDEX); + continFlag_ = keyBlockTablePtr != nullptr; + if (!this->SetPlatformInfoForTiling()) { + return ge::GRAPH_FAILED; + } + + uint32_t batch = GetShape(0).GetDim(0); + uint32_t head = GetShape(1).GetDim(1); + uint32_t qHead = GetShape(0).GetDim(1); + uint32_t headGroupNum = qHead / head; + seqLen_ = GetShape(1).GetDim(2); /* when continFlag==true, it is blockSize */ + if (continFlag_) { + uint32_t blockSize = GetShape(1).GetDim(2); /* when continFlag==true, it is blockSize */ + seqLen_ = GetInputAttrData(0); + seqLen_ = ops::CeilDiv(seqLen_, blockSize) * blockSize; + uint32_t blockCount = GetShape(KEY_BLOCK_TABLE_INPUT_INDEX).GetDim(1); + tilingData_.params.set_blockCount(blockCount); + } + uint32_t dimension = GetShape(0).GetDim(3) * COMPRESSED_RATE; + uint64_t nope_dimension = GetShape(1).GetDim(3) * COMPRESSED_RATE; + uint32_t reducedBatch = batch * head; + uint32_t usedCoreNum = std::min(reducedBatch, coreNum_); + uint32_t singleCoreBatch = ops::CeilDiv(reducedBatch, usedCoreNum); + tilingData_.params.set_batch(batch); + tilingData_.params.set_head(head); + tilingData_.params.set_maxSeqLen(seqLen_); + tilingData_.params.set_qHead(qHead); + tilingData_.params.set_headGroupNum(headGroupNum); + + tilingData_.params.set_maxK(maxK); + + tilingData_.params.set_dimension(dimension); + tilingData_.params.set_nope_dimension(nope_dimension); + tilingData_.params.set_reducedBatch(reducedBatch); + tilingData_.params.set_usedCoreNum(usedCoreNum); + tilingData_.params.set_tileN1(TILE_N1); + if (continFlag_) { + uint32_t blockSize = GetShape(1).GetDim(2); /* 2: the dim of blockSize */ + tilingData_.params.set_tileN1(blockSize); + } + tilingData_.params.set_tileN2(TILE_N2); + tilingData_.params.set_singleCoreBatch(singleCoreBatch); + + tilingData_.params.set_singleCoreSeqLen(seqLen_); + tilingData_.params.set_kNopeUnpackGmOffset(static_cast(reducedBatch) * seqLen_ * nope_dimension / 2); /* 2 : 1 / sizeof(int4b_t) */ + tilingData_.params.set_qUnpackGmOffset(static_cast(reducedBatch) * seqLen_ * dimension / 2); + tilingData_.params.set_mmGmOffset(static_cast(reducedBatch) * seqLen_ * dimension / 2 + /* 2 : 1 / sizeof(int4b_t) */ + static_cast(reducedBatch) * 1 * dimension / 2); + + this->SetMatmulTiling(); + + bool supportKeyRope = context_->GetOptionalInputShape(KEY_ROPE_INPUT_INDEX) != nullptr;; + tilingData_.params.set_supportKeyRope(supportKeyRope); + if (supportKeyRope) { + uint64_t rope_dimension = GetShape(KEY_ROPE_INPUT_INDEX).GetDim(3) * COMPRESSED_RATE; + tilingData_.params.set_rope_dimension(rope_dimension); + this->SetMatmulTilingRope(); + } + + this->SetTopKTiling(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HammingDistTopKTiling::DoLibApiTiling() { + return ge::GRAPH_SUCCESS; +} + +uint64_t HammingDistTopKTiling::GetTilingKey() { return 1; } + +ge::graphStatus HammingDistTopKTiling::GetWorkspaceSize() { + uint64_t *workspaces = context_->GetWorkspaceSizes(1); + uint64_t sysWorkspaceSize = WORKSIZE; + /* usrWorkspaceSize = workspace for Select + workspace for Topk */ + uint64_t usrWorkspaceSize = ops::CeilDiv(static_cast(tilingData_.params.get_reducedBatch()) * + tilingData_.params.get_maxSeqLen() * tilingData_.params.get_dimension() * + sizeof(int8_t), static_cast(2))*2 + /* 2: 1/2, size of int4 */ + ops::CeilDiv(static_cast(tilingData_.params.get_reducedBatch()) * + tilingData_.params.get_dimension() * sizeof(int8_t), static_cast(2)) + + static_cast(tilingData_.params.get_reducedBatch()) * + tilingData_.params.get_maxSeqLen() * sizeof(uint16_t); + workspaces[0] = sysWorkspaceSize + usrWorkspaceSize; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HammingDistTopKTiling::PostTiling() { + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + auto blockDim = tilingData_.params.get_usedCoreNum(); + context_->SetBlockDim(blockDim); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +void HammingDistTopKTiling::Reset() { + tilingData_.SetDataPtr(context_->GetRawTilingData()->GetData()); + inputParams_.mSize = 0UL; + inputParams_.kSize = 0UL; + inputParams_.nSize = 0UL; + inputParams_.queryDtype = ge::DT_INT4; + inputParams_.keyDtype = ge::DT_UINT8; + inputParams_.kDtype = ge::DT_INT32; + inputParams_.seqLenDtype = ge::DT_INT32; + inputParams_.indicesDtype = ge::DT_INT32; + inputParams_.libApiWorkSpaceSize = 0U; + inputParams_.opName = nullptr; + inputParams_.aFormat = ge::FORMAT_ND; + inputParams_.bFormat = ge::FORMAT_ND; + inputParams_.cFormat = ge::FORMAT_ND; +} + +void HammingDistTopKTiling::SetMatmulTiling() { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo()); + matmul_tiling::MultiCoreMatmulTiling tiling(ascendcPlatform); + tiling.SetDim(1); + tiling.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_INT4); + tiling.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_INT4); + tiling.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16); + if (seqLen_ >= SEQ_LEN_THRES) { + tiling.SetFixSplit(-1, 512, -1); /* 512: BaseN = 512 and BaseK = 128 can fully utilize L0B */ + } + uint64_t nope_dimension = GetShape(1).GetDim(3) * 8; + tiling.SetShape(1, seqLen_, nope_dimension); + tiling.SetSingleShape(1, seqLen_, nope_dimension); + tiling.SetOrgShape(1, seqLen_, nope_dimension); + tiling.SetBias(false); + tiling.GetTiling(tilingData_.matmulTiling); /* if ret = -1, get tiling failed */ +} + +void HammingDistTopKTiling::SetMatmulTilingRope() { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo()); + matmul_tiling::MultiCoreMatmulTiling tiling(ascendcPlatform); + tiling.SetDim(1); + tiling.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_INT4); + tiling.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_INT4); + tiling.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16); + if (seqLen_ >= SEQ_LEN_THRES) { + tiling.SetFixSplit(-1, 512, -1); /* 512: BaseN = 512 and BaseK = 128 can fully utilize L0B */ + } + uint64_t rope_dimension = GetShape(KEY_ROPE_INPUT_INDEX).GetDim(3) * 8; + tiling.SetShape(1, seqLen_, rope_dimension); + tiling.SetSingleShape(1, seqLen_, rope_dimension); + tiling.SetOrgShape(1, seqLen_, rope_dimension); + tiling.SetBias(false); + tiling.GetTiling(tilingData_.matmulTilingRope); /* if ret = -1, get tiling failed */ +} + +void HammingDistTopKTiling::SetTopKTiling() { + uint32_t inner = std::min(ops::CeilDiv(seqLen_, TOP_K_ALIGN_NUM) * TOP_K_ALIGN_NUM, tilingData_.params.get_tileN2()); + uint32_t outer = 1; + uint32_t k = std::min(seqLen_, maxK); + uint32_t maxSize = 0; + uint32_t minSize = 0; + uint32_t dTypeSize = 2; /* 2:size of float16 */ + const bool IS_REUSESOURCE = false; + const bool IS_INITINDEX = true; + const bool IS_LARGEST = true; + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo()); + tilingData_.params.set_outer(outer); + tilingData_.params.set_inner(inner); + tilingData_.params.set_topkN(inner); + AscendC::TopKTilingFunc(ascendcPlatform, inner, outer, k, dTypeSize, IS_INITINDEX, AscendC::TopKMode::TOPK_NORMAL, IS_LARGEST, tilingData_.topkTiling); + AscendC::GetTopKMaxMinTmpSize(ascendcPlatform, inner, outer, IS_REUSESOURCE, IS_INITINDEX, AscendC::TopKMode::TOPK_NORMAL, IS_LARGEST, dTypeSize, maxSize, minSize); +} + +const gert::Shape HammingDistTopKTiling::GetShape(const size_t index) { + return context_->GetInputShape(index)->GetStorageShape(); +} + +const gert::Shape HammingDistTopKTiling::GetOutShape(const size_t index) { + return context_->GetOutputShape(index)->GetStorageShape(); +} + +const uint32_t HammingDistTopKTiling::GetInputAttrData(const size_t index) { + if (auto attrPtr = context_->GetAttrs()) { + const int64_t* p = attrPtr->GetInt(index); + if (p != nullptr) { + return static_cast(*p); + } + } + return 0; +} + +bool HammingDistTopKTiling::SetPlatformInfoForTiling() { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo()); + coreNum_ = ascendcPlatform.GetCoreNumAic(); + return true; +} + +} \ No newline at end of file diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k.h b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k.h new file mode 100644 index 000000000..7f3fb5d1c --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k.h @@ -0,0 +1,88 @@ +#ifndef HAMMING_DIST_TOP_K_H +#define HAMMING_DIST_TOP_K_H + + +#include "hamming_dist_top_k_tiling.h" +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" + +namespace optiling { +class HammingDistTopKTiling { +public: + // from parent class + gert::TilingContext *context_ = nullptr; + std::unique_ptr ascendcPlatform_{nullptr}; + uint32_t blockDim_{0}; + uint64_t workspaceSize_{0}; + uint64_t tilingKey_{0}; + AiCoreParams aicoreParams_{0}; + + // from child class + HammingDistTopKMatmulInfo inputParams_; + uint32_t libApiWorkSpaceSize_ = 0; + uint32_t coreNum_ = 1; + const char *opName_ = ""; + int32_t dtypeByte_ = 2; /* 2: size of float16 */ + HammingDistTopKTilingData tilingData_; + bool compileInfoInit_ = false; + bool continFlag_ = false; + uint32_t seqLen_ = 1; + + HammingDistTopKTiling(gert::TilingContext *context) : context_(context) { + InitAttrParam(); + uint32_t dimNum = GetOutShape(0).GetDimNum(); + maxK = GetOutShape(0).GetDim(dimNum - 1); + } + + bool IsCapable(); + // 1. Obtain platform information such as CoreNum, UB/L1/L0C resource size + ge::graphStatus GetPlatformInfo(); + // 2. Obtain INPUT/OUTPUT/ATTR information + ge::graphStatus GetShapeAttrsInfo(); + // 3. Calculate data split TilingData + ge::graphStatus DoOpTiling(); + // 4. Calculate TilingData for high-level API + ge::graphStatus DoLibApiTiling(); + // 5. Calculate TilingKey + uint64_t GetTilingKey(); + // 6. Calculate Workspace size + ge::graphStatus GetWorkspaceSize(); + // 7. Save Tiling data + ge::graphStatus PostTiling(); + + void Reset(); + void SetMatmulTiling(); + void SetMatmulTilingRope(); + void SetTopKTiling(); + bool SetPlatformInfoForTiling(); + const gert::Shape GetShape(const size_t index); + // Get input data + const uint32_t GetInputAttrData(const size_t index); + // output shape + const gert::Shape GetOutShape(const size_t index); + + // Initialize sink and recent + const void InitAttrParam() { + uint32_t sink = GetInputAttrData(1); + uint32_t recent = GetInputAttrData(2); + uint32_t supportOffload = GetInputAttrData(3); + tilingData_.params.set_sink(sink); + tilingData_.params.set_recent(recent); + tilingData_.params.set_supportOffload(supportOffload); + } + + uint32_t maxK = 512; // default maxK + uint32_t TILE_N1 = 254; + uint32_t TILE_N2 = 3328; + uint32_t DIMENSION = 128; + uint32_t SEQ_LEN_THRES = 512; + uint64_t SUB_BLOCK_NUM_WITH_DB = 4; + uint64_t WORKSIZE = 16 * 1024 * 1024; + uint32_t TOP_K_ALIGN_NUM = 32; + uint32_t KEY_ROPE_INPUT_INDEX = 7; + uint32_t KEY_BLOCK_TABLE_INPUT_INDEX = 5; + uint32_t COMPRESSED_RATE = 8; +}; +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_def.cpp b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_def.cpp new file mode 100644 index 000000000..8f015091e --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_def.cpp @@ -0,0 +1,102 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hamming_dist_top_k_def.cpp + * \brief + */ +#include +#include "register/op_def_registry.h" + +namespace ops { +class HammingDistTopK : public OpDef { +public: + explicit HammingDistTopK(const char* name) : OpDef(name) + { + this->Input("query") + .ParamType(REQUIRED) + .DataType({ge::DT_UINT8}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("key_compressed") + .ParamType(REQUIRED) + .DataType({ge::DT_UINT8}) + .Format({ge::FORMAT_ND}); + this->Input("k") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("seq_len") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("chunk_size") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Attr("max_seq_len") + .AttrType(OPTIONAL) + .Int(0); + this->Attr("sink") + .AttrType(OPTIONAL) + .Int(0); + this->Attr("recent") + .AttrType(OPTIONAL) + .Int(0); + this->Attr("support_offload") + .AttrType(OPTIONAL) + .Int(0); + this->Input("key_block_table") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("indices_in") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("key_compressed_rope") + .ParamType(OPTIONAL) + .DataType({ge::DT_UINT8}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("mask") + .ParamType(OPTIONAL) + .DataType({ge::DT_BOOL}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("indices") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("aclnnSupport.value", "support_aclnn") + .ExtendCfgInfo("jitCompile.flag", "static_false,dynamic_false"); + + this->AICore().AddConfig("ascend910_93", aicore_config); + this->AICore().AddConfig("ascend910b", aicore_config); + + } +}; + +OP_ADD(HammingDistTopK); +} diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_proto.cpp b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_proto.cpp new file mode 100644 index 000000000..d8547409c --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_proto.cpp @@ -0,0 +1,40 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hamming_dist_top_k_proto.cpp + * \brief + */ +#include +#include +#include "error/ops_error.h" + +using namespace ge; + +namespace ops { +static ge::graphStatus InferShapeHammingDistTopK(gert::InferShapeContext *context) +{ + gert::Shape *outShape = context->GetOutputShape(0); + const gert::Shape *inputShape = context->GetInputShape(6); + *outShape = *inputShape; + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataTypeHammingDistTopK(gert::InferDataTypeContext *context) +{ + ge::DataType outputType = context->GetInputDataType(ge::DT_INT32); + context->SetOutputDataType(0, outputType); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(HammingDistTopK) + .InferShape(InferShapeHammingDistTopK) + .InferDataType(InferDataTypeHammingDistTopK); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_split.cpp b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_split.cpp new file mode 100644 index 000000000..2ecf99c3f --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_split.cpp @@ -0,0 +1,153 @@ + +#include "hamming_dist_top_k_tiling.h" +#include "hamming_dist_top_k.h" +#include "hamming_dist_top_k_split.h" +#include +#include + +namespace optiling { +namespace { + +} +bool HammingDistTopKSplitSTiling::IsCapable() { + SetPlatformInfoForTiling(); + bool isContinuousBatch = context_->GetOptionalInputShape(KEY_BLOCK_TABLE_INPUT_INDEX) != nullptr; + uint32_t batch = GetShape(0).GetDim(0); + uint32_t maxSeqLen = GetShape(1).GetDim(2); /* when continFlag==false, it is maxSeqLen */ + if (isContinuousBatch) { + uint32_t blockSize = GetShape(1).GetDim(2); /* when continFlag==true, it is blockSize */ + tilingData_.params.set_tileN1(blockSize); + uint32_t blockCount = GetShape(KEY_BLOCK_TABLE_INPUT_INDEX).GetDim(1); + tilingData_.params.set_blockCount(blockCount); + + maxSeqLen = GetInputAttrData(0); + maxSeqLen = ((maxSeqLen + blockSize - 1) / blockSize) * blockSize; + if (maxSeqLen == 0) { + maxSeqLen = blockCount * blockSize; + } + } else { + tilingData_.params.set_tileN1(TILE_N1); + } + tilingData_.params.set_maxSeqLen(maxSeqLen); + uint32_t head = GetShape(1).GetDim(1); + uint32_t usedCoreNum = coreNum_; + if (head > usedCoreNum) { + return false; + } + + if (maxSeqLen > SUPER_LONG_SEQLEN || (batch < MAX_BATCH && maxSeqLen > MIN_SPLIT_S_SEQLEN)) { + return true; + } + return false; +} + +ge::graphStatus HammingDistTopKSplitSTiling::GetWorkspaceSize() { + uint64_t *workspaces = context_->GetWorkspaceSizes(1); + uint64_t sysWorkspaceSize = WORKSIZE; + //usrWorkspaceSize = workspace for Select + workspace for Topk + uint64_t usrWorkspaceSize = ops::CeilDiv(static_cast(tilingData_.params.get_layerSize() * COMPRESSED_RATE * sizeof(int8_t)), static_cast(2)) + + ops::CeilDiv(static_cast(tilingData_.params.get_layerSizeRope() * COMPRESSED_RATE * sizeof(int8_t)), static_cast(2)) + + ops::CeilDiv(static_cast(tilingData_.params.get_matmulResultSize() * sizeof(float)), static_cast(2)) + + ops::CeilDiv(static_cast(tilingData_.params.get_topKValueSize() * sizeof(float)), static_cast(2)) + + static_cast(tilingData_.params.get_topKIdexSize() * sizeof(int32_t)) + + ops::CeilDiv(static_cast(tilingData_.params.get_batchN()) * + tilingData_.params.get_dimension() * sizeof(int8_t), static_cast(2)); + + workspaces[0] = sysWorkspaceSize + WORKSPACE_SCALE * usrWorkspaceSize; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HammingDistTopKSplitSTiling::DoOpTiling() { + uint64_t batch = GetShape(0).GetDim(0); + uint64_t qHead = GetShape(0).GetDim(1); + uint64_t head = GetShape(1).GetDim(1); + uint64_t dimension = GetShape(0).GetDim(3) * COMPRESSED_RATE; + uint64_t nope_dimension = GetShape(1).GetDim(3) * COMPRESSED_RATE; + uint64_t headGroupNum = qHead / head; + uint64_t maxSeqLen = tilingData_.params.get_maxSeqLen(); + uint64_t usedCoreNum = coreNum_; + uint64_t tileN2 = 4 * 1024; + tilingData_.params.set_batch(batch); + tilingData_.params.set_head(head); + tilingData_.params.set_qHead(qHead); + tilingData_.params.set_headGroupNum(headGroupNum); + tilingData_.params.set_batchN(batch * head); + tilingData_.params.set_dimension(dimension); + tilingData_.params.set_nope_dimension(nope_dimension); + tilingData_.params.set_layerSize(batch * head * maxSeqLen * nope_dimension / COMPRESSED_RATE); + tilingData_.params.set_matmulResultSize(batch * head * maxSeqLen); + tilingData_.params.set_topKValueSize(batch * head * ops::CeilDiv(maxSeqLen, tileN2) * maxK); + tilingData_.params.set_topKIdexSize(batch * head * ops::CeilDiv(maxSeqLen, tileN2) * maxK); + tilingData_.params.set_topKInnerSize(TOP_K_INNER_SIZE); + tilingData_.params.set_maxK(maxK); + tilingData_.params.set_usedCoreNum(usedCoreNum); + tilingData_.params.set_sBlockSize(S_BLOCK_SIZE); + tilingData_.params.set_tileN3(TILE_N3); + tilingData_.params.set_tileN2(tileN2); + bool supportKeyRope = context_->GetOptionalInputShape(KEY_ROPE_INPUT_INDEX) != nullptr;; + tilingData_.params.set_supportKeyRope(supportKeyRope); + SetMatmulTiling(); + if (supportKeyRope) { + uint64_t rope_dimension = GetShape(KEY_ROPE_INPUT_INDEX).GetDim(3) * COMPRESSED_RATE; + tilingData_.params.set_rope_dimension(rope_dimension); + tilingData_.params.set_layerSizeRope(batch * head * maxSeqLen * rope_dimension / COMPRESSED_RATE); + SetMatmulTilingRope(); + } + SetTopKTiling(); + return ge::GRAPH_SUCCESS; +} + +uint64_t HammingDistTopKSplitSTiling::GetTilingKey() { + return 10; +} + +void HammingDistTopKSplitSTiling::SetMatmulTiling() { + uint64_t nope_dimension = GetShape(1).GetDim(3) * COMPRESSED_RATE; + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo()); + matmul_tiling::MultiCoreMatmulTiling tiling(ascendcPlatform); + tiling.SetDim(1); + tiling.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_INT4); + tiling.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_INT4); + tiling.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16); + tiling.SetFixSplit(-1, L0B_BASE_SIZE, -1); + tiling.SetShape(1, VECTOR_CUBE_RATIO * tilingData_.params.get_tileN2(), nope_dimension); + tiling.SetSingleShape(1, VECTOR_CUBE_RATIO * tilingData_.params.get_tileN2(), nope_dimension); + tiling.SetOrgShape(1, VECTOR_CUBE_RATIO * tilingData_.params.get_tileN2(), nope_dimension); + tiling.SetBias(false); + tiling.GetTiling(tilingData_.matmulTiling); // if ret = -1, get tiling failed +} + +void HammingDistTopKSplitSTiling::SetMatmulTilingRope() { + uint64_t rope_dimension = GetShape(KEY_ROPE_INPUT_INDEX).GetDim(3) * COMPRESSED_RATE; + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo()); + matmul_tiling::MultiCoreMatmulTiling tiling(ascendcPlatform); + tiling.SetDim(1); + tiling.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_INT4); + tiling.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_INT4); + tiling.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16); + tiling.SetFixSplit(-1, L0B_BASE_SIZE, -1); + tiling.SetShape(1, VECTOR_CUBE_RATIO * tilingData_.params.get_tileN2(), rope_dimension); + tiling.SetSingleShape(1, VECTOR_CUBE_RATIO * tilingData_.params.get_tileN2(), rope_dimension); + tiling.SetOrgShape(1, VECTOR_CUBE_RATIO * tilingData_.params.get_tileN2(), rope_dimension); + tiling.SetBias(false); + tiling.GetTiling(tilingData_.matmulTilingRope); // if ret = -1, get tiling failed +} + +void HammingDistTopKSplitSTiling::SetTopKTiling() { + uint32_t inner = tilingData_.params.get_topKInnerSize(); + uint32_t outer = 1; + uint32_t dTypeSize = 2; // 2:size of float16 + const bool IS_REUSESOURCE = false; + const bool IS_INITINDEX = true; + const bool IS_LARGEST = true; + uint32_t maxSize = 0; + uint32_t minSize = 0; + auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo()); + tilingData_.params.set_outer(outer); + tilingData_.params.set_inner(inner); + tilingData_.params.set_topkN(inner); + AscendC::TopKTilingFunc(ascendcPlatform, inner, outer, maxK, dTypeSize, IS_INITINDEX, AscendC::TopKMode::TOPK_NORMAL, IS_LARGEST, tilingData_.topkTiling); + AscendC::GetTopKMaxMinTmpSize(ascendcPlatform, inner, outer, IS_REUSESOURCE, IS_INITINDEX, AscendC::TopKMode::TOPK_NORMAL, IS_LARGEST, dTypeSize, maxSize, minSize); +} + +} \ No newline at end of file diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_split.h b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_split.h new file mode 100644 index 000000000..84710e094 --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_split.h @@ -0,0 +1,45 @@ +#ifndef HAMMING_DIST_TOP_K_SPLIT_H +#define HAMMING_DIST_TOP_K_SPLIT_H + +#include "hamming_dist_top_k.h" +#include "hamming_dist_top_k_tiling.h" +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" + +namespace optiling { +class HammingDistTopKSplitSTiling : public HammingDistTopKTiling { +public: + HammingDistTopKSplitSTiling(gert::TilingContext *context) : HammingDistTopKTiling(context) {} + + bool IsCapable(); + + ge::graphStatus DoOpTiling(); + + uint64_t GetTilingKey(); + + void SetMatmulTiling(); + void SetMatmulTilingRope(); + + void SetTopKTiling(); + + ge::graphStatus GetWorkspaceSize(); + + float CORE_USE_RATIO = 0.8f; + uint64_t WORKSIZE = 16 * 1024 * 1024; + uint32_t COMPRESSED_RATE = 8; + uint32_t TILE_N1 = 128; + uint32_t TILE_N3 = 7 * 1024; + uint32_t TOP_K_INNER_SIZE = 4 * 1024; + uint32_t S_BLOCK_SIZE = 256; + uint32_t L0B_BASE_SIZE = 512; + uint32_t VECTOR_CUBE_RATIO = 2; + uint64_t WORKSPACE_SCALE = 2; + uint64_t MAX_BATCH = 16; + uint64_t SUPER_LONG_SEQLEN = 26 * 1024; + uint64_t MIN_SPLIT_S_SEQLEN = 8 * 1024; + + uint32_t KEY_ROPE_INPUT_INDEX = 7; + uint32_t KEY_BLOCK_TABLE_INPUT_INDEX = 5; +}; +} +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_tiling.cpp b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_tiling.cpp new file mode 100644 index 000000000..252fd3977 --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_tiling.cpp @@ -0,0 +1,43 @@ +#include "hamming_dist_top_k_tiling.h" +#include "hamming_dist_top_k.h" +#include "hamming_dist_top_k_split.h" +#include "register/op_def_registry.h" +#include "register/op_impl_registry.h" + + +namespace optiling { +static ge::graphStatus TilingPrepareForHammingDistTopK(gert::TilingParseContext *context) +{ + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus TilingFunc(gert::TilingContext* context) { + ge::graphStatus ret; + HammingDistTopKSplitSTiling hammingDistTopKSplitSTiling(context); + hammingDistTopKSplitSTiling.GetShapeAttrsInfo(); + hammingDistTopKSplitSTiling.GetPlatformInfo(); + auto can_split = hammingDistTopKSplitSTiling.IsCapable(); + if (can_split) { + hammingDistTopKSplitSTiling.DoOpTiling(); + hammingDistTopKSplitSTiling.DoLibApiTiling(); + hammingDistTopKSplitSTiling.GetWorkspaceSize(); + hammingDistTopKSplitSTiling.PostTiling(); + context->SetTilingKey(hammingDistTopKSplitSTiling.GetTilingKey()); + return ge::GRAPH_SUCCESS; + } + + HammingDistTopKTiling hammingDistTopKTiling(context); + hammingDistTopKTiling.GetShapeAttrsInfo(); + hammingDistTopKTiling.GetPlatformInfo(); + hammingDistTopKTiling.IsCapable(); + hammingDistTopKTiling.DoOpTiling(); + hammingDistTopKTiling.GetWorkspaceSize(); + hammingDistTopKTiling.PostTiling(); + context->SetTilingKey(hammingDistTopKTiling.GetTilingKey()); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(HammingDistTopK) + .Tiling(TilingFunc) + .TilingParse(TilingPrepareForHammingDistTopK); +} diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_tiling.h b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_tiling.h new file mode 100644 index 000000000..54317da7b --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_host/hamming_dist_top_k_tiling.h @@ -0,0 +1,92 @@ +#ifndef HAMMING_DIST_TOP_K_TILING_H +#define HAMMING_DIST_TOP_K_TILING_H + +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "op_host_util.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(HammingDistTopKTilingParams) + TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum); + TILING_DATA_FIELD_DEF(uint32_t, batch); + TILING_DATA_FIELD_DEF(uint32_t, batchN); + TILING_DATA_FIELD_DEF(uint32_t, head); + TILING_DATA_FIELD_DEF(uint32_t, dimension); + TILING_DATA_FIELD_DEF(uint32_t, nope_dimension); + TILING_DATA_FIELD_DEF(uint32_t, rope_dimension); + TILING_DATA_FIELD_DEF(uint32_t, reducedBatch); + TILING_DATA_FIELD_DEF(uint32_t, maxSeqLen); + TILING_DATA_FIELD_DEF(uint32_t, sink); + TILING_DATA_FIELD_DEF(uint32_t, recent); + TILING_DATA_FIELD_DEF(uint32_t, supportOffload); + TILING_DATA_FIELD_DEF(uint32_t, layerSize); + TILING_DATA_FIELD_DEF(uint32_t, layerSizeRope); + TILING_DATA_FIELD_DEF(uint32_t, matmulResultSize); + TILING_DATA_FIELD_DEF(uint32_t, topKValueSize); + TILING_DATA_FIELD_DEF(uint32_t, topKIdexSize); + TILING_DATA_FIELD_DEF(uint32_t, topKInnerSize); + TILING_DATA_FIELD_DEF(uint32_t, maxK); + TILING_DATA_FIELD_DEF(uint32_t, tileN1); + TILING_DATA_FIELD_DEF(uint32_t, sBlockSize); + TILING_DATA_FIELD_DEF(uint32_t, blockCount); + TILING_DATA_FIELD_DEF(uint32_t, tileN3); + TILING_DATA_FIELD_DEF(uint32_t, tileN2); + TILING_DATA_FIELD_DEF(uint32_t, singleCoreBatch); + TILING_DATA_FIELD_DEF(uint32_t, singleCoreSeqLen); + TILING_DATA_FIELD_DEF(uint32_t, outer); + TILING_DATA_FIELD_DEF(uint32_t, inner); + TILING_DATA_FIELD_DEF(uint32_t, topkN); + TILING_DATA_FIELD_DEF(uint64_t, kNopeUnpackGmOffset); + TILING_DATA_FIELD_DEF(uint64_t, mmGmOffset); + TILING_DATA_FIELD_DEF(uint32_t, qHead); + TILING_DATA_FIELD_DEF(uint64_t, qUnpackGmOffset); + TILING_DATA_FIELD_DEF(uint32_t, headGroupNum); + TILING_DATA_FIELD_DEF(uint32_t, supportKeyRope); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(HammingDistTopKTilingParamsOp, HammingDistTopKTilingParams) + +BEGIN_TILING_DATA_DEF(HammingDistTopKTilingData) + TILING_DATA_FIELD_DEF_STRUCT(HammingDistTopKTilingParams, params); + TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, matmulTiling); + TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, matmulTilingRope); + TILING_DATA_FIELD_DEF_STRUCT(TopkTiling, topkTiling); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(HammingDistTopK, HammingDistTopKTilingData) +REGISTER_TILING_DATA_CLASS(HammingDistTopKTilingDataOp, HammingDistTopKTilingData) + +struct HammingDistTopKMatmulInfo { + bool transA = false; + bool transB = false; + bool hasBias = false; + uint64_t mSize = 0UL; + uint64_t kSize = 0UL; + uint64_t nSize = 0UL; + ge::DataType queryDtype = ge::DT_INT4; + ge::DataType keyDtype = ge::DT_UINT8; + ge::DataType kDtype = ge::DT_INT32; + ge::DataType seqLenDtype = ge::DT_INT32; + ge::DataType indicesDtype = ge::DT_INT32; + int64_t outDtype = 0L; + uint32_t libApiWorkSpaceSize = 0U; + uint64_t bf16ExtreWorkSpaceSize = 0UL; + const char *opName = nullptr; + ge::Format aFormat = ge::FORMAT_ND; + ge::Format bFormat = ge::FORMAT_ND; + ge::Format cFormat = ge::FORMAT_ND; +}; + +struct AiCoreParams { + uint64_t ubSize; + uint64_t blockDim; + uint64_t aicNum; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; +}; +// using HammingDistTopKCompileInfo = gert::GemmCompileInfo; +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_host/op_host_util.h b/csrc/ascend/moe/hamming_dist_top_k/op_host/op_host_util.h new file mode 100644 index 000000000..a8fe8d6ab --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_host/op_host_util.h @@ -0,0 +1,192 @@ +/** + * Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/*! + * \file op_host_util.h + * \brief + */ + +#ifndef CANN_OPS_BUILT_IN_OP_UTIL_H_ +#define CANN_OPS_BUILT_IN_OP_UTIL_H_ + +#include +#include +#include + +namespace ops { + +/** + * if y is 0, return x + */ +template +typename std::enable_if ::value, T>::type CeilDiv(T x, T y) { + if (y != 0 && x != 0) { + const T quotient = x / y; + return (x % y != 0 && ((x ^ y) >= 0)) ? (quotient + 1) : quotient; + } + + return x; +} + +/** + * if y is 0, return x + */ +template +typename std::enable_if ::value, T>::type CeilDiv(T x, T y) { + if (y != 0 && x != 0) { + const T quotient = x / y; + return (x % y != 0) ? (quotient + 1) : quotient; + } + + return x; +} + + +/** + * if y is 0, return x + */ +template +typename std::enable_if ::value, T>::type FloorDiv(T x, T y) { + return y == 0 ? x : x / y; +} + +/** + * if align is 0, return 0 + */ +template +typename std::enable_if ::value, T>::type CeilAlign(T x, T align) { + return CeilDiv(x, align) * align; +} + +/** + * if align is 0, return 0 + */ +template +typename std::enable_if ::value, T>::type FloorAlign(T x, T align) { + return align == 0 ? 0 : x / align * align; +} + +} // namespace ops + + +namespace optiling { + +enum CubeTilingType { + CUBE_DYNAMIC_SHAPE_TILING, + CUBE_DEFAULT_TILING, + CUBE_BINARY_TILING, +}; + +constexpr uint64_t kInvalidTilingId = std::numeric_limits::max(); + +class CubeCompileInfo { + public: + bool correct_range_flag = false; + CubeTilingType tiling_type = CUBE_DYNAMIC_SHAPE_TILING; + uint64_t default_tiling_id = kInvalidTilingId; + std::vector default_range; + std::vector> repo_seeds; + std::vector> repo_range; + std::vector> cost_range; + std::vector> batch_range; // for dynamic batch + std::vector repo_tiling_ids; + std::vector cost_tiling_ids; + std::vector batch_tiling_ids; // for dynamic batch + std::map block_dim; + std::string soc_version = ""; + + uint32_t core_num = 0; + uint64_t ub_size = 0; + uint64_t l1_size = 0; + uint64_t l2_size = 0; + uint64_t l0a_size = 0; + uint64_t l0b_size = 0; + uint64_t l0c_size = 0; + uint64_t bt_size = 0; + int32_t cube_freq = 0; + bool load3d_constraints = true; + bool intrinsic_data_move_l12ub = true; + bool intrinsic_matmul_ub_to_ub = false; + bool intrinsic_conv_ub_to_ub = false; + bool intrinsic_data_move_l0c2ub = true; + bool intrinsic_fix_pipe_l0c2out = false; + bool intrinsic_fix_pipe_l0c2ub = false; + bool intrinsic_data_move_out2l1_nd2nz = false; + bool intrinsic_data_move_l12bt_bf16 = false; +}; + +struct BatchmatmulCompileParas { + bool binary_mode_flag = false; + bool bias_flag = false; + bool at_l1_flag = true; + bool split_k_flag = false; + bool pattern_flag = false; + bool zero_flag = false; + bool sparse_4to2_flag = false; + bool binary_constant_flag = false; + bool vector_pre_conv_mode = false; + float fused_double_operand_num = 0; + float aub_double_num = 0; + float bub_double_num = 0; + int64_t quant_scale = 0; + int64_t eltwise_src = 0; + int8_t enable_pad = 0; + bool enable_nz_fusion = false; + bool enable_rt_bank_cache = false; +}; + +struct Ub2UbBatchmatmulCompileParas { + int64_t block_m0 = 1; + int64_t block_n0 = 16; + int64_t block_a_k0 = 16; + int64_t block_b_k0 = 16; + bool is_batch_matmul = false; + bool bm_fusion_flag = false; + + std::string pre_conv = "None"; + std::string pre_activation = "None"; + std::string post_anti_quant = "None"; + std::string post_eltwise = "None"; + std::string post_activation = "None"; + std::string post_quant = "None"; + std::string post_transform = "None"; +}; + +enum DynamicMode { + DYNAMIC_MKN, + DYNAMIC_MKNB, + WEIGHT_QUANT_BMM +}; + +class HammingDistTopKCompileInfo : public CubeCompileInfo { + public: + HammingDistTopKCompileInfo() = default; + ~HammingDistTopKCompileInfo() = default; + + bool trans_a = false; + bool trans_b = false; + bool repo_seed_flag = false; + bool repo_costmodel_flag = false; + uint32_t workspace_num = 0; + uint32_t ub_size = 0; + BatchmatmulCompileParas params; + Ub2UbBatchmatmulCompileParas ub2ub_params; + DynamicMode dynamic_mode = DYNAMIC_MKN; +}; + +} + +#endif // CANN_OPS_BUILT_IN_OP_UTIL_H_ diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_kernel/hamming_dist_top_k.cpp b/csrc/ascend/moe/hamming_dist_top_k/op_kernel/hamming_dist_top_k.cpp new file mode 100644 index 000000000..56ab3a3af --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_kernel/hamming_dist_top_k.cpp @@ -0,0 +1,27 @@ +#include "hamming_dist_top_k_split_s.h" +#include "hamming_dist_top_k_parallel.h" + +using namespace AscendC; + +extern "C" __global__ __aicore__ void hamming_dist_top_k(GM_ADDR query, GM_ADDR keyCompressed, GM_ADDR k, + GM_ADDR seqLen, GM_ADDR chunkSize, GM_ADDR keyBlockTable, GM_ADDR indicesIn, GM_ADDR keyCompressedRope, GM_ADDR mask, GM_ADDR indices, GM_ADDR workspace, GM_ADDR tiling) { + TPipe tPipe; + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + GM_ADDR user1 = GetUserWorkspace(workspace); + if (user1 == nullptr) { + return; + } + + GET_TILING_DATA(tilingData, tiling); + if (TILING_KEY_IS(1)) { + HammingDistTopKParallelKernel op; + op.Init(query, keyCompressed, keyCompressedRope, k, seqLen, chunkSize, keyBlockTable, mask, indices, user1, tilingData, &tPipe); + op.Process(); + tPipe.Destroy(); + } else if (TILING_KEY_IS(10)) { + HammingDistTopKSplitSKernel op; + op.Init(query, keyCompressed, keyCompressedRope, k, seqLen, chunkSize, keyBlockTable, mask, indices, user1, tilingData, &tPipe); + op.Process(); + tPipe.Destroy(); + } +} \ No newline at end of file diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_kernel/hamming_dist_top_k_base.h b/csrc/ascend/moe/hamming_dist_top_k/op_kernel/hamming_dist_top_k_base.h new file mode 100644 index 000000000..3baaca87e --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_kernel/hamming_dist_top_k_base.h @@ -0,0 +1,363 @@ +/** + * Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/*! + * \file hamming_dist_top_k_base.h + * \brief + */ + +#ifndef HAMMING_DIST_TOP_K_BASE_H +#define HAMMING_DIST_TOP_K_BASE_H + +#include "kernel_operator.h" +#include "kernel_tiling/kernel_tiling.h" +#include "lib/matmul_intf.h" + +namespace AscendC { + +#define YF_LOG(format, ...) \ +if (false) { \ + printf("CoreIdx: %d on CoreType %d, " format, GetBlockIdx(), g_coreType, ##__VA_ARGS__); \ +} + +//constexpr uint32_t SKIP_HEAD_BLOCK_NUM = 1; +//constexpr uint32_t SKIP_TAIL_BLOCK_NUM = 2; +//constexpr uint32_t SKIP_HEAD_TOKEN_NUM = 128; +//constexpr uint32_t SKIP_TAIL_TOKEN_NUM = 256; + +constexpr uint32_t MAX_FP16_PROCESS_NUM = 128; +constexpr uint32_t MAX_INT32_PROCESS_NUM = 64; +constexpr float MIN_HALF_VALUE = -65535; +constexpr half MAX_HALF_VALUE = (half) 65504; + +// datablock bytes = 32bytes +constexpr uint32_t DATABLOCK_BYTES = 32; +// half 4 datablocks element size +constexpr uint32_t FOUR_DATABLOCKS_ELEMENT_SIZE = 64; +// half 8 datablocks element size +constexpr uint32_t EIGHT_DATABLOCKS_ELEMENT_SIZE = 128; + +constexpr MatmulConfig MM_CFG_NO_PRELOAD{false, false, true, 0, 0, 0, false, false, false, false, false, + 0, 0, 0, 0, 0, 0, 0, true}; + +struct TilingParam { + uint32_t usedCoreNum = 0; + uint32_t preCoreNum = 0; + uint32_t isBias = 0; + uint32_t M = 0; + uint32_t N = 0; + uint32_t baseM = 0; + uint32_t baseN = 0; + uint32_t singleCoreM = 0; + uint32_t singleCoreN = 0; + uint32_t singleCoreK = 0; + uint32_t ka = 0; + uint32_t kb = 0; + uint32_t rope_ka = 0; + uint32_t rope_kb = 0; + // tiling data for select + uint32_t layer = 0; + uint32_t batch = 0; + uint32_t head = 0; + uint32_t batchN = 0; + uint32_t selectUsedCoreNum = 0; + uint32_t layerSize = 0; + uint32_t layerSizeRope = 0; + uint32_t seqLen = 0; + uint32_t dimension = 0; + uint32_t nope_dimension = 0; + uint32_t rope_dimension = 0; + uint32_t reducedBatch = 0; + uint32_t tileN1 = 0; + uint32_t tileN2 = 0; + uint32_t singleCoreBatch = 0; + uint32_t singleCoreSeqLen = 0; + bool supportKeyRope; + // tiling data for matmul + uint32_t matmulResultSize = 0; + // tiling data for topk + uint32_t maxK = 0; + uint32_t maxSeqLen = 0; + uint32_t sink = 0; + uint32_t recent = 0; + uint32_t topKInnerSize = 0; + uint32_t topKValueSize = 0; + uint32_t topKIdexSize = 0; + uint32_t kNopeUnpackGmOffset = 0; + uint32_t mmGmOffset = 0; + uint32_t qHead = 0; + uint32_t headGroupNum = 0; + uint64_t qUnpackGmOffset = 0; + uint64_t blockCount = 0; + // support offload + bool supportOffload = false; +}; + +template +__aicore__ inline T Min(const T a, const T b) +{ + return a < b ? a : b; +} + +template +__aicore__ inline T Max(const T a, const T b) +{ + return a > b ? a : b; +} + +template +__aicore__ inline void SelectCustom(const LocalTensor &dstLocal, const LocalTensor &keyCompressed, const LocalTensor &src0Local, uint8_t repeatTimes) +{ + AscendC::BinaryRepeatParams repeatParams = {1, 1, 1, 8, 0, 8}; // {dstBlkStride, src0BlkStride, src1BlkStride, dstRepStride, src0RepStride, src1RepStride} src0 is reused, set repeat stride to 0. + uint64_t mask = MAX_FP16_PROCESS_NUM; + // DumpTensor(keyCompressed, 123, 256); + // DumpTensor(src0Local, 124, 256); + Select(dstLocal, keyCompressed, src0Local, static_cast(-1), AscendC::SELMODE::VSEL_TENSOR_SCALAR_MODE, mask, repeatTimes, repeatParams); + // DumpTensor(dstLocal, 126, 256); +} + +template +__aicore__ inline void TopKCustom(const LocalTensor &dstValueLocal, const LocalTensor &dstIndexLocal, + const LocalTensor &srcValueLocal, const LocalTensor &srcIndexLocal, const int32_t k, const HammingDistTopKTilingData &tiling, uint32_t n) +{ + LocalTensor finishLocal; + AscendC::TopKInfo topkInfo; + topkInfo.outter = tiling.params.outer; + topkInfo.n = n; + topkInfo.inner = matmul::CeilDiv(n, 32) * 32; /* 32: inner must be aligned to 32 */ + TopK(dstValueLocal, dstIndexLocal, srcValueLocal, srcIndexLocal, finishLocal, k, tiling.topkTiling, topkInfo, true); +} + +__aicore__ inline void ReduceMaxCustom(const GlobalTensor &inputGm, const LocalTensor &reduceInputLocal, + const LocalTensor &reduceOutputLocal, const uint16_t chunkNum, const uint8_t chunkSize) +{ + uint32_t dataBlockNum = (static_cast(chunkNum) * static_cast(chunkSize) + 15) / 16; // Each dataBlock contains 16 half elements + uint32_t blockLen = static_cast(16 * sizeof(half)); // Determine blockLen (more robust for copying by dataBlock unit), 32 bytes per dataBlock + // uint32_t blockLen = static_cast(chunkSize << 1); // Equivalent to blockLen=chunkSize*2=16*2=32, half type occupies 2 bytes, multiply by 2 to get the length in bytes + // DataCopyExtParams: Copy dataBlockCount dataBlocks to local, keep layout as continuous dataBlock list + DataCopyExtParams copyInParams{static_cast(dataBlockNum), blockLen, 0, 0, 0}; // {255, 32, 0, 0, 0}={number of blocks to copy, length per block, 0, 0, 0}, total data size 8160 + + if (chunkSize == 16 || chunkSize == 64 || chunkSize == 128) { // When chunkSize=16, 64, 128, copy one block directly at a time + copyInParams.blockCount = 1; + // copyInParams.blockLen = static_cast(chunkSize * chunkNum * sizeof(half)); // blockLen=16*255*2=8160 + copyInParams.blockLen = static_cast(dataBlockNum * blockLen); // Equivalent + } + + DataCopyPadExtParams copyInPadParams{false, 0, 0, 0}; // No padding + DataCopyPad(reduceInputLocal, inputGm, copyInParams, copyInPadParams); // DataCopyPad is an internal operator + // DumpTensor(reduceInputLocal, 158, 6 * chunkSize); + /* For positions where chunkNum tail is less than 8, fill with half minimum value, + so that the output of BlockReduceMax at corresponding positions will also be half minimum value, + which will not affect the subsequent TopK calculation */ + uint32_t dataBlockNumAligned = matmul::CeilDiv(dataBlockNum, 8) * 8; // dataBlockNumAligned=256, /* 8: BlockReduceMax processes 8 dataBlocks in parallel at one time */ + if (dataBlockNumAligned > dataBlockNum) { // 256>255 + Duplicate(reduceInputLocal[dataBlockNum * 16], static_cast(MIN_HALF_VALUE), (dataBlockNumAligned - dataBlockNum) * 16); /* 16: Each dataBlock is 32Bytes, containing 16 half values */ + } // Copy to 4080 + // printf("base.h dataBlockNumAligned: %d\n", dataBlockNumAligned); + + SetFlag(1); // Wait for copy completion + WaitFlag(1); + PipeBarrier(); + PipeBarrier(); + + if (chunkSize == 64) { + int32_t totalRepeat = dataBlockNumAligned / 8; + int32_t repeat = Min(MAX_REPEAT_TIMES, totalRepeat); + int32_t loopNum = matmul::CeilDiv(totalRepeat, repeat); // loopNum=1 + int32_t tailRepeat = totalRepeat - (loopNum - 1) * repeat; + uint64_t mask[2] = {0, 0}; /* 2: Set mask bit by bit, 2 64bit variables required */ + + mask[0] = UINT64_MAX; + uint32_t srcOffset = 0; + uint32_t dstOffset = 0; + + for (int32_t i = 0; i < loopNum - 1; i++) { + WholeReduceMax(reduceOutputLocal[dstOffset], reduceInputLocal[srcOffset], mask, repeat * 2, 1, 1, 4, ReduceOrder::ORDER_ONLY_VALUE); // (..., repeat, dstRepStride, srcBlkStride, srcRepStride) + srcOffset += repeat * 8 * 16; // Move repeat segments of 128-element: repeat * 128 + dstOffset += repeat * 2; // Advance dstOffset by output count: 1 value output per repeat + } + // First 4 datablocks of each repeat, 64 elements + // Iteration count tailRepeat * 2 + WholeReduceMax(reduceOutputLocal[dstOffset], reduceInputLocal[srcOffset], mask, tailRepeat * 2, 1, 1, 4, ReduceOrder::ORDER_ONLY_VALUE); + return; + } + + if (chunkSize == 128) { + int32_t totalRepeat = dataBlockNumAligned / 8; // (dataBlockNumAligned * 16) / 128 + int32_t repeat = Min(MAX_REPEAT_TIMES, totalRepeat); + int32_t loopNum = matmul::CeilDiv(totalRepeat, repeat); // loopNum=1 + int32_t tailRepeat = totalRepeat - (loopNum - 1) * repeat; + uint64_t mask[2]; /* 2: Set mask bit by bit, 2 64bit variables required */ + /* For chunkSize==128, we need to cover 128 consecutive half elements in each repeat, + so set all 128 bits to 1 (continuous participation in reduction) */ + mask[0] = UINT64_MAX; + mask[1] = UINT64_MAX; + + uint32_t srcOffset = 0; + uint32_t dstOffset = 0; + /* Explanation: + - One dataBlock contains 16 half elements (32 bytes) + - Process 8 dataBlocks in parallel at one time, so 1 repeat corresponds to 8 * 16 = 128 half elements + - WholeReduceMax(mask=128, repeat=k) outputs 1 maximum value for each repeat (128 half elements) + —— So dstOffset should only advance by 1 per repeat (not 8) + */ + for (int32_t i = 0; i < loopNum - 1; i++) { + WholeReduceMax(reduceOutputLocal[dstOffset], reduceInputLocal[srcOffset], mask, repeat, 1, 1, 8, ReduceOrder::ORDER_ONLY_VALUE); // (..., repeat, dstRepStride, srcBlkStride, srcRepStride) + srcOffset += repeat * 8 * 16; /* Move repeat segments of 128-element: repeat * 128 */ + dstOffset += repeat; // Advance dstOffset by output count: 1 value output per repeat + } + WholeReduceMax(reduceOutputLocal[dstOffset], reduceInputLocal[srcOffset], mask, tailRepeat, 1, 1, 8, ReduceOrder::ORDER_ONLY_VALUE); // (..., repeat, dstRepStride, srcBlkStride, srcRepStride) + // DumpTensor(reduceOutputLocal, 218, chunkNum); + } else { + int32_t totalRepeat = dataBlockNumAligned / 8; // totalRepeat=32, /* 8: BlockReduceMax processes 8 dataBlocks in parallel at one time */, repeat 32 times to complete + int32_t repeat = Min(MAX_REPEAT_TIMES, totalRepeat); // Internal parameter MAX_REPEAT_TIMES, unknown? Assume repeat=32 + int32_t loopNum = matmul::CeilDiv(totalRepeat, repeat); // loopNum=1 + int32_t tailRepeat = totalRepeat - (loopNum - 1) * repeat; // tailRepeat=32 + uint64_t mask[2]; /* 2: Set mask bit by bit, 2 64bit variables required */ + + if (chunkSize == 16) { /* chunkSize only supports 1, 8, 16 */ + mask[0] = UINT64_MAX; // 0xffffffffffffffff, bitwise mask, all 1, all participate in calculation + mask[1] = UINT64_MAX; // 0xffffffffffffffff + } else if (chunkSize == 8) { /* chunkSize only supports 1, 8, 16 */ + mask[0] = 0x00ff00ff00ff00ff; + mask[1] = 0x00ff00ff00ff00ff; + } + + uint32_t srcOffset = 0; + uint32_t dstOffset = 0; + for (int32_t i = 0; i < loopNum - 1; i++) { + BlockReduceMax(reduceOutputLocal[dstOffset], reduceInputLocal[srcOffset], repeat, mask, 1, 1, 8); // (..., mask, dstRepStride, srcBlkStride, srcRepStride) + srcOffset += repeat * 8 * 16; /* 8: BlockReduceMax processes 8 dataBlocks in parallel at one time, 16: Each dataBlock is 32Bytes, containing 16 half values */ + dstOffset += repeat * 8; /* 8: BlockReduceMax processes 8 dataBlocks in parallel at one time, outputs 8 points */ + } + BlockReduceMax(reduceOutputLocal[dstOffset], reduceInputLocal[srcOffset], tailRepeat, mask, 1, 1, 8); // (..., mask, dstRepStride, srcBlkStride, srcRepStride) + // repeat = 32, 8 elements one repeat, 256 elements total + // srcBlkStride = 1, no gap between blocks in one repeat + // dstRepStride = 1, srcRepStride = 8, no gap between repeats + } +} + +__aicore__ inline void SortInt32AscendingUB(LocalTensor& buf, uint32_t len) { + if ASCEND_IS_AIC { return; } + if (len <= 1) { return; } + __ubuf__ int32_t* data = reinterpret_cast<__ubuf__ int32_t*>(buf.GetPhyAddr()); + for (uint32_t i = 1; i < len; ++i) { + int32_t key = data[i]; + int32_t j = static_cast(i) - 1; + while (j >= 0 && data[j] > key) { + data[j + 1] = data[j]; + --j; + } + data[j + 1] = key; + } +} + +__aicore__ inline void WriteBlockTableFromTopK( + uint32_t curBatchIdx, + LocalTensor& topKIndexUb, // UB: Chunk index obtained by TopK (length ≥ curKScalar) + LocalTensor& blockIdUb, // UB: Temporary buffer allocated by the caller (length ≥ curKScalar) + uint32_t curKScalar, + uint64_t outGmOffset, // Offset for writing back to GM(indices) + LocalTensor& tableBlockTensor, + const GlobalTensor& indicesGm, + bool isContinuousBatch, + uint32_t blockCount // Number of blocks per batch (calculated by tileN1 or fixed BLOCK_SIZE) +) { + if ASCEND_IS_AIC { return; } + + //YF_LOG("ldeng WriteBlockTableFromTopK 245 curBatchIdx=%d, curKScalar=%d,blockCount=%d,\n", curBatchIdx,curKScalar,blockCount) + // DumpTensor(topKIndexUb, 246, topKIndexUb.GetSize()); + + // Sort chunk_id in ascending order in UB + SortInt32AscendingUB(topKIndexUb, curKScalar); + + // DumpTensor(topKIndexUb, 251, topKIndexUb.GetSize()); + + // Directly read and write in UB in scalar mode + __ubuf__ const int32_t* in_ptr = reinterpret_cast<__ubuf__ const int32_t*>(topKIndexUb.GetPhyAddr()); + __ubuf__ int32_t* out_ptr = reinterpret_cast<__ubuf__ int32_t*>(blockIdUb.GetPhyAddr()); + + for (uint32_t i = 0; i < curKScalar; ++i) { + const int32_t idx = in_ptr[i]; + out_ptr[i] = isContinuousBatch + ? tableBlockTensor.GetValue(static_cast(idx)) + : (idx + 1); // When no block_table exists,约定 block_id = idx + 1（1-based） + } + + // DumpTensor(blockIdUb, 264, blockIdUb.GetSize()); + // UB -> GM(indices) + DataCopyExtParams cpOut{1, static_cast(curKScalar * sizeof(int32_t)), 0, 0, 0}; + DataCopyPad(indicesGm[outGmOffset], blockIdUb, cpOut); + + // DumpTensor(blockIdUb, 269, 64); + +} + +// Used for setting tail top-k to MAX_HALF_VALUE +// tensorSize should be less than topKValueInTensor.GetSize() +// copyLen is the total number of elements that should be set to MAX_HALF_VALUE, starting from the actual tail +__aicore__ inline void FillMaxValueFromTail( + LocalTensor &topKValueInTensor, uint32_t tensorSize, uint32_t copyLen, uint32_t curChunkSize) +{ + if ASCEND_IS_AIC { + return; + } + + ASCENDC_ASSERT((copyLen <= tensorSize), { KERNEL_LOG(KERNEL_ERROR, "copyLen should be less tensorSize"); }); + // case1: tensorSize - copyLen % alignedElements = 0, address 32bytes aligned + // YF_LOG("tensorSize = %d, copyLen = %d\n", tensorSize, copyLen); + uint32_t alignedElements = DATABLOCK_BYTES / sizeof(half); + uint32_t offset = tensorSize - copyLen; + if (offset % alignedElements == 0) { + Duplicate(topKValueInTensor[offset], static_cast(MAX_HALF_VALUE), copyLen); + return; + } + + // case2: compute aligned address + uint32_t offsetAligned = offset / alignedElements * alignedElements; /* floor aligned for datacopy */ + ASCENDC_ASSERT((offsetAligned >= 0), { KERNEL_LOG(KERNEL_ERROR, "offsetAligned should be nonnegative"); }); + + // After 32Bytes alignment, number of elements to process + uint32_t alignedAddCopyElements = tensorSize - offsetAligned; + // Process 128 elements in one iteration, 8 datablocks, one datablock 32Bytes + // Use mask[] bitwise mode to control elements + uint64_t mask[2] = {0, 0}; + uint32_t needSkipElements = alignedAddCopyElements - copyLen; + // Directly process the remaining unprocessed elements, address is already 32bytes aligned + int32_t lastCopyLen = alignedAddCopyElements - EIGHT_DATABLOCKS_ELEMENT_SIZE; + // If one iteration can complete processing, use one iteration; otherwise, split into two Duplicate processes + if (lastCopyLen > 0) { + // Process 128 elements first + alignedAddCopyElements = EIGHT_DATABLOCKS_ELEMENT_SIZE; + } + // YF_LOG("tensorSize = %d, copyLen = %d offsetAligned = %d alignedAddCopyElements = %d needSkipElements = %d lastCopyLen = %d\n", tensorSize, copyLen, offsetAligned, alignedAddCopyElements, needSkipElements, lastCopyLen); + if (alignedAddCopyElements <= FOUR_DATABLOCKS_ELEMENT_SIZE) { + mask[0] = (UINT64_MAX << needSkipElements) & (UINT64_MAX >> (FOUR_DATABLOCKS_ELEMENT_SIZE - alignedAddCopyElements)); + } else if (alignedAddCopyElements <= EIGHT_DATABLOCKS_ELEMENT_SIZE) { + mask[0] = (UINT64_MAX << needSkipElements); + mask[1] = UINT64_MAX >> (EIGHT_DATABLOCKS_ELEMENT_SIZE - alignedAddCopyElements); + } + // YF_LOG("mask[0] = %x mask[1] = %x \n", mask[0], mask[1]); + + Duplicate(topKValueInTensor[offsetAligned], static_cast(MAX_HALF_VALUE), mask, 1, 1, 8); + + if (lastCopyLen > 0) { + Duplicate(topKValueInTensor[offsetAligned + EIGHT_DATABLOCKS_ELEMENT_SIZE], static_cast(MAX_HALF_VALUE), lastCopyLen); + } +} + +} // namespace AscendC +#endif // HAMMING_DIST_TOP_K_BASE_H diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_kernel/hamming_dist_top_k_parallel.h b/csrc/ascend/moe/hamming_dist_top_k/op_kernel/hamming_dist_top_k_parallel.h new file mode 100644 index 000000000..27456fcf8 --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_kernel/hamming_dist_top_k_parallel.h @@ -0,0 +1,917 @@ +/** + * Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/*! + * \file hamming_dist_top_k_parallel.h + * \brief + */ +#ifndef HAMMING_DIST_TOP_K_PARALLEL_H +#define HAMMING_DIST_TOP_K_PARALLEL_H + +#include "hamming_dist_top_k_base.h" + +namespace AscendC { +class HammingDistTopKParallelKernel { +public: + __aicore__ inline HammingDistTopKParallelKernel() {} + __aicore__ inline void Init(GM_ADDR query, GM_ADDR keyCompressed, GM_ADDR keyCompressedRope, GM_ADDR k, + GM_ADDR seqLen, GM_ADDR chunkSize, GM_ADDR keyBlockTable, GM_ADDR mask, + GM_ADDR indices, GM_ADDR workSpace, + const HammingDistTopKTilingData &tilingData, TPipe *pipe) + { + const TCubeTiling &tiling = tilingData.matmulTiling; + const TCubeTiling &tilingRope = tilingData.matmulTilingRope; + const TopkTiling &topkTiling = tilingData.topkTiling; + const HammingDistTopKTilingParams &tilingParam = tilingData.params; + pipe_ = pipe; + tilingData_ = tilingData; + InitTilingParams(tiling, tilingRope, topkTiling, tilingParam); + InitParams(); + InitGlobalBuffers(query, keyCompressed, keyCompressedRope, k, seqLen, chunkSize, keyBlockTable, mask, indices, workSpace); + + continFlag_ = keyBlockTableGm_.GetPhyAddr() != nullptr; + mm_.SetSubBlockIdx(0); + mm_.Init(&tiling, pipe_); + if (param_.supportKeyRope) { + mmRope_.SetSubBlockIdx(0); + mmRope_.Init(&tilingRope, pipe_); + } + } + + __aicore__ inline void Process() + { + uint32_t blockIndex = AscendC::GetBlockIdx(); + if ASCEND_IS_AIV { + blockIndex = blockIndex / 2; + } + if (blockIndex >= param_.usedCoreNum) { + return; + } + uint32_t batchNumPerLoop = 1; + uint32_t curCoreLoop = (curCoreBatch_ + batchNumPerLoop - 1) / batchNumPerLoop; + uint32_t tailBatchNumPerLoop = curCoreBatch_ - (curCoreLoop - 1) * batchNumPerLoop; + uint32_t computeBatch; + uint32_t nextLoopComputeBatch; + + if ASCEND_IS_AIV { + pipe_->Reset(); + } + + for (uint32_t loopIdx = 0; loopIdx < curCoreLoop; loopIdx++) { + ComputeUnpackMM(loopIdx * batchNumPerLoop, batchNumPerLoop, loopIdx % BATCH_PING_PONG_NUM); + } + if ASCEND_IS_AIV { + pipe_barrier(PIPE_ALL); + pipe_->Reset(); + InitLocalBuffersForTopK(); + } + + for (uint32_t loopIdx = 0; loopIdx < curCoreLoop; loopIdx++) { + if ASCEND_IS_AIV { + VectorWaitCube(SYNC_AIC_AIV_FLAG2 + loopIdx % BATCH_PING_PONG_NUM); + if (GetSubBlockIdx() == loopIdx % 2) { + ComputeTopK(loopIdx, loopIdx % BATCH_PING_PONG_NUM); + } + } + } + } + +protected: + __aicore__ inline int32_t GetCurSeqLen(const GlobalTensor &seqLenGm_, + const GlobalTensor &chunkSizeGm_, uint32_t realCurBatch) + { + int32_t curSeqLen = seqLenGm_.GetValue(realCurBatch); + int32_t curChunkSize = 1; + if (chunkSizeGm_.GetPhyAddr() != nullptr) { + curChunkSize = chunkSizeGm_.GetValue(realCurBatch); + } + if (curChunkSize == 1 || curChunkSize == 8 || curChunkSize == 16) { + if (curSeqLen <= 32) { + curSeqLen = 0; + } else { + curSeqLen = curSeqLen - Min(curSeqLen, static_cast(curSeqLen % curChunkSize + 16)); + } + } else if (curChunkSize == 64) { + curSeqLen = ((curSeqLen + 63) / 64) * 64; + } else if(curChunkSize == 128){ + curSeqLen = ((curSeqLen + 127)/128)*128; + } + else { + curSeqLen = 0; + } + return curSeqLen; + } + + __aicore__ inline void ComputeUnpackMM(uint32_t batchIdx, uint32_t batchNum, uint32_t pingPongFlag) + { + uint32_t maxLoopNum = (batchNum >> 1) << 1; + for (uint32_t j = 0; j < maxLoopNum; j++) { + uint32_t curReducedBatch = (curCoreBatchStartIdx_ + batchIdx + j); + uint32_t realCurBatch = curReducedBatch / param_.head; + if (supportMask_) { + bool batchMask = maskGm_.GetValue(realCurBatch); + //YF_LOG("realCurBatch = %d, batchMask = %d\n", realCurBatch, batchMask); + if (!batchMask) { + return; + } + } + int32_t curSeqLen = GetCurSeqLen(seqLenGm_, chunkSizeGm_, realCurBatch); + int32_t curK = kGm_.GetValue(realCurBatch); + if (curK == 0 || curSeqLen == 0) { + continue; + } + if ASCEND_IS_AIV { + if (GetSubBlockIdx() == (j % SUB_BLOCK_NUM)) { + UnpackOneBatch(curSeqLen, curReducedBatch, 0, true); + } + VectorNotifyCube(SYNC_AIV_AIC_FLAG + pingPongFlag); + } + if ASCEND_IS_AIC { + CubeWaitVector(SYNC_AIV_AIC_FLAG + pingPongFlag); + ComputeMM(batchIdx + j, curSeqLen); + } + } + if (batchNum > maxLoopNum) { + ComputeUnpackMMForLastBatch(batchIdx, pingPongFlag); + } + if ASCEND_IS_AIC { + CubeNotifyVector(SYNC_AIC_AIV_FLAG2 + pingPongFlag); + } + } + + __aicore__ inline void ComputeUnpackMMForLastBatch(uint32_t batchIdx, uint32_t pingPongFlag) + { + uint32_t curReducedBatch = (curCoreBatchStartIdx_ + batchIdx); + uint32_t realCurBatch = curReducedBatch / param_.head; + if (supportMask_) { + bool batchMask = maskGm_.GetValue(realCurBatch); + //YF_LOG("realCurBatch = %d, batchMask = %d\n", realCurBatch, batchMask); + if (!batchMask) { + return; + } + } + int32_t curSeqLen = GetCurSeqLen(seqLenGm_, chunkSizeGm_, realCurBatch); + int32_t curK = kGm_.GetValue(realCurBatch); + if (curK == 0 || curSeqLen == 0) { + return; + } + if ASCEND_IS_AIV { + uint32_t curBlockCount = matmul::CeilDiv(curSeqLen, param_.tileN1); + uint32_t subBlockSeqOffset = 0; + uint32_t subBlockSeqLen = 0; + if (GetSubBlockIdx() == 0) { + subBlockSeqLen = continFlag_ ? matmul::CeilDiv(curBlockCount, SUB_BLOCK_NUM) * param_.tileN1 : + matmul::CeilDiv(curSeqLen, SUB_BLOCK_NUM); + } else if (GetSubBlockIdx() == 1) { + subBlockSeqLen = continFlag_ ? (curBlockCount / SUB_BLOCK_NUM) * param_.tileN1 : curSeqLen / SUB_BLOCK_NUM; + if (subBlockSeqLen <= 0) { + VectorNotifyCube(SYNC_AIV_AIC_FLAG + pingPongFlag); + return; + } + subBlockSeqOffset += continFlag_ ? matmul::CeilDiv(curBlockCount, SUB_BLOCK_NUM) * param_.tileN1 : + matmul::CeilDiv(curSeqLen, SUB_BLOCK_NUM); + } + + UnpackOneBatch(subBlockSeqLen, curReducedBatch, subBlockSeqOffset, GetSubBlockIdx() == 0); + VectorNotifyCube(SYNC_AIV_AIC_FLAG + pingPongFlag); + } + if ASCEND_IS_AIC { + CubeWaitVector(SYNC_AIV_AIC_FLAG + pingPongFlag); + ComputeMM(batchIdx, curSeqLen); + } + } + + __aicore__ inline void UnpackOneBatch(uint32_t sequenceLen, uint32_t batchIdx, uint32_t subBlockSeqOffset, bool unpackQuery) + { + if ASCEND_IS_AIC { + return; + } + InitLocalBuffersForUnpackQuery(); + // unpack query + UnpackQuery(batchIdx); + pipe_->Reset(); + InitLocalBuffersForUnpackKey(); + // unpack key + UnpackKey(sequenceLen, batchIdx, subBlockSeqOffset, false); + if (param_.supportKeyRope) { + UnpackKey(sequenceLen, batchIdx, subBlockSeqOffset, true); + } + pipe_->Reset(); + } + + __aicore__ inline void UnpackQuery(uint32_t batchIdx) { + + LocalTensor constTensor = constBuf_.template Get(); + LocalTensor selectTensor = selectBuf_. template Get(); + LocalTensor qReduceSumTensor = qReduceSumBuf_. template Get(); + LocalTensor qReduceSumLastRowTensor = qReduceSumLastRowBuf_. template Get(); + + Duplicate(constTensor, 1, param_.dimension); + + uint32_t compressedDimension = param_.dimension / COMPRESS_RATE; + uint64_t queryGmOffset = batchIdx * param_.headGroupNum * compressedDimension; + LocalTensor queryCompressed = queryCompressedInQueue_.AllocTensor(); + DataCopyExtParams queryCopyInParams{1, param_.headGroupNum * static_cast(compressedDimension), 0, 0, 0}; + DataCopyPadExtParams queryCopyInPadParams{false, 0, 0, 0}; + DataCopyPad(queryCompressed, queryGm_[queryGmOffset], queryCopyInParams, queryCopyInPadParams); + //DumpTensor(queryCompressed, 210, queryCompressed.GetSize()); + queryCompressedInQueue_.EnQue(queryCompressed); + queryCompressed = queryCompressedInQueue_.DeQue(); + // Iterate 128 elements per iteration + uint32_t repeatedTimes = matmul::CeilDiv(param_.headGroupNum * param_.dimension, MAX_FP16_PROCESS_NUM); + SelectCustom(selectTensor, queryCompressed, constTensor, static_cast(repeatedTimes)); + PipeBarrier(); + if (param_.supportKeyRope) { + uint32_t pad_dim = param_.rope_dimension / 2; + uint32_t valid_dim = param_.nope_dimension + pad_dim; + // YF_LOG("param_.dimension=%d, valid_dim=%d, pad_dim=%d\n", param_.dimension, valid_dim, pad_dim); + for (uint32_t i = 0; i < param_.headGroupNum; i++) { + // DumpTensor(selectTensor[i * param_.dimension], 224, 672); + Duplicate(selectTensor[i * param_.dimension + valid_dim], 0, pad_dim); + // DumpTensor(selectTensor[i * param_.dimension], 226, 672); + } + } + queryCompressedInQueue_.FreeTensor(queryCompressed); + + uint64_t qMask = MAX_FP16_PROCESS_NUM; + uint32_t repeatTimes = matmul::CeilDiv(param_.dimension, MAX_FP16_PROCESS_NUM); + LocalTensor qHashTensor = selectTensor; + //DumpTensor(qHashTensor, 223, param_.dimension); + if (param_.headGroupNum > 1) { + static constexpr AscendC::CumSumConfig cumSumConfig{false, false, true}; + const AscendC::CumSumInfo cumSumInfo{param_.headGroupNum, param_.dimension}; + AscendC::CumSum(qReduceSumTensor, qReduceSumLastRowTensor, selectTensor, cumSumInfo); + PipeBarrier(); + + if (param_.headGroupNum > 8) { + uint32_t div = matmul::CeilDiv(param_.headGroupNum, 8); + half reciprocalDiv = static_cast((float)1.0 / div); + AscendC::Muls(qReduceSumLastRowTensor, qReduceSumTensor[(param_.headGroupNum - 1) * param_.dimension], + reciprocalDiv, qMask, repeatTimes, {1, 1, 8, 8}); + PipeBarrier(); + } + qHashTensor = qReduceSumLastRowTensor; + } + //DumpTensor(qHashTensor, 240, param_.dimension); + LocalTensor queryUnpacked = queryUnpackedOutQueue_.AllocTensor(); + Cast(queryUnpacked, qHashTensor, RoundMode::CAST_CEIL, qMask, repeatTimes, {1, 1, 2, 8}); + queryUnpackedOutQueue_.EnQue(queryUnpacked); + queryUnpacked = queryUnpackedOutQueue_.DeQue(); + uint64_t unpackQGmOffset = queryGmOffset * 8 / param_.headGroupNum; + DataCopyExtParams copyQOutParams{1, static_cast(param_.dimension / 2), 0, 0, 0}; /* 2: 1 / size of int4b_t */ + DataCopyPad(qUnpackGm_[unpackQGmOffset], queryUnpacked, copyQOutParams); + queryUnpackedOutQueue_.FreeTensor(queryUnpacked); + } + + __aicore__ inline void UnpackKey(uint32_t sequenceLen, + uint32_t batchIdx, + uint32_t subBlockSeqOffset, + bool isKeyRope) { + uint32_t realBatchIdx = batchIdx / param_.head; /* batchIdx without headNum */ + uint32_t headIdx = batchIdx % param_.head; + + uint32_t dimension = isKeyRope ? param_.rope_dimension : param_.nope_dimension; + GlobalTensor keyGm = isKeyRope ? keyRopeGm_ : keyGm_; + + uint32_t sequenceBlockNum = matmul::CeilDiv(sequenceLen, param_.tileN1); + uint32_t tailN1 = sequenceLen - (sequenceBlockNum - 1) * param_.tileN1; + uint32_t compressedDimension = dimension / COMPRESS_RATE; + + LocalTensor constTensor = constBuf_.template Get(); + LocalTensor selectTensor = selectBuf_.template Get(); + + Duplicate(constTensor, 1, dimension); + uint32_t selectRepeatedTimes = computeSelectRepeatedTimes(dimension); + for (uint32_t j = 0; j < sequenceBlockNum; j++) { + LocalTensor keyCompressed = keyCompressedInQueue_.AllocTensor(); + uint32_t copySeqLen = j == sequenceBlockNum - 1 ? tailN1 : param_.tileN1; + DataCopyExtParams copyInParams{1, static_cast(copySeqLen * compressedDimension), 0, 0, 0}; + DataCopyPadExtParams copyInPadParams{false, 0, 0, 0}; + uint64_t keyGmOffset = batchIdx * param_.maxSeqLen * compressedDimension + + j * param_.tileN1 * compressedDimension + subBlockSeqOffset * compressedDimension; + if (!continFlag_) { + DataCopyPad(keyCompressed, keyGm[keyGmOffset], copyInParams, copyInPadParams); + } else { + int32_t blockTableVal = keyBlockTableGm_.GetValue(realBatchIdx * param_.blockCount + j + + subBlockSeqOffset / param_.tileN1); + uint64_t keyGmOffsetConti = (headIdx + static_cast(blockTableVal) * param_.head) * + param_.tileN1 * compressedDimension; + DataCopyPad(keyCompressed, keyGm[keyGmOffsetConti], copyInParams, copyInPadParams); + } + keyCompressedInQueue_.EnQue(keyCompressed); + keyCompressed = keyCompressedInQueue_.DeQue(); + //DumpTensor(keyCompressed, 285, static_cast(copySeqLen * compressedDimension)); + //YF_LOG("sequenceBlockNum = %d j = %d copySeqLen = %d selectRepeatedTimes = %d subBlockSeqOffset = %d sequenceLen = %d\n", sequenceBlockNum, j, copySeqLen, selectRepeatedTimes, subBlockSeqOffset, sequenceLen); + + if (selectRepeatedTimes > MAX_SELECT_REPEATED_TIMES) { + UnpackKeyCompressedWithBigDim(selectRepeatedTimes, selectTensor, constTensor, keyCompressed, keyGmOffset); + } else { + UnpackKeyCompressedWithLittleDim(selectTensor, constTensor, keyCompressed, dimension, keyGmOffset, copySeqLen); + } + keyCompressedInQueue_.FreeTensor(keyCompressed); + } + } + + __aicore__ inline void UnpackKeyCompressedWithBigDim(uint32_t selectRepeatedTimes, + LocalTensor &selectTensor, + LocalTensor &constTensor, + LocalTensor &keyCompressed, + uint64_t keyGmOffset) { + uint32_t keySelectCycleCount = (selectRepeatedTimes + MAX_SELECT_REPEATED_TIMES - 1) / MAX_SELECT_REPEATED_TIMES; + uint32_t tailRepeateadTimes = selectRepeatedTimes % MAX_SELECT_REPEATED_TIMES; + LocalTensor keyUnpacked = keyUnpackedOutQueue_.AllocTensor(); + // key offset + uint64_t keyCompressedOffset = 0; + uint32_t keyOffset = MAX_FP16_PROCESS_NUM * MAX_SELECT_REPEATED_TIMES / COMPRESS_RATE; + for (uint32_t index = 0; index < keySelectCycleCount; index++) { + uint32_t selectAndCastOffset = index * MAX_FP16_PROCESS_NUM * MAX_SELECT_REPEATED_TIMES; + uint32_t maxRepeatedTimes = (tailRepeateadTimes > 0 && (index == keySelectCycleCount - 1)) + ? tailRepeateadTimes + : MAX_SELECT_REPEATED_TIMES; + SelectCustom(selectTensor, keyCompressed[keyCompressedOffset], constTensor, static_cast(maxRepeatedTimes)); + // DumpTensor(selectTensor, 366, 672); + Cast(keyUnpacked[selectAndCastOffset], selectTensor, RoundMode::CAST_CEIL, CAST_MASK, static_cast(maxRepeatedTimes), {1, 1, 2, 8}); + keyCompressedOffset = keyOffset * (index + 1); + // YF_LOG("keyGmOffset_ = %d tailRepeateadTimes = %d selectRepeatedTimes = %d \n", keyGmOffset_, tailRepeateadTimes, selectRepeatedTimes); + } + keyUnpackedOutQueue_.EnQue(keyUnpacked); + keyUnpacked = keyUnpackedOutQueue_.DeQue(); + DataCopyParams copyParams{1, static_cast(selectRepeatedTimes * MAX_FP16_PROCESS_NUM / 2 / BLOCK_CUBE), 0, 0}; // 2: 1/2, size of int4b_t + DataCopy(unpackGm_[keyGmOffset * COMPRESS_RATE], keyUnpacked, copyParams); // output to outQueue1 with DB_ON + keyUnpackedOutQueue_.FreeTensor(keyUnpacked); + } + + __aicore__ inline void UnpackKeyCompressedWithLittleDim(LocalTensor &selectTensor, + LocalTensor &constTensor, + const LocalTensor &keyCompressed, + uint32_t dimension, + uint64_t keyGmOffset, uint32_t repeateadTimes) { + SelectCustom(selectTensor, keyCompressed, constTensor, static_cast(repeateadTimes)); + PipeBarrier(); + LocalTensor keyUnpacked = keyUnpackedOutQueue_.AllocTensor(); + uint64_t mask = MAX_FP16_PROCESS_NUM; + Cast(keyUnpacked, selectTensor, RoundMode::CAST_CEIL, mask, static_cast(repeateadTimes), {1, 1, 2, 8}); + keyUnpackedOutQueue_.EnQue(keyUnpacked); + keyUnpacked = keyUnpackedOutQueue_.DeQue(); + uint64_t unpackGmOffset = keyGmOffset * 8; /* 8: Original Dimension / Compressed Dimension */ + DataCopyExtParams copyOutParams{1, static_cast(repeateadTimes * dimension / 2), 0, 0, 0}; /* 2: 1 / size of int4b_t */ + if (param_.supportKeyRope) { + DataCopyPad(kRopeUnpackGm_[unpackGmOffset], keyUnpacked, copyOutParams); + } else { + DataCopyPad(unpackGm_[unpackGmOffset], keyUnpacked, copyOutParams); + } + keyUnpackedOutQueue_.FreeTensor(keyUnpacked); + } + + __aicore__ inline void ComputeMM(uint32_t batchIdx, uint32_t seqLen) + { + if ASCEND_IS_AIV { + return; + } + mm_.SetOrgShape(param_.M, seqLen, param_.ka); + mm_.SetSingleShape(param_.M, seqLen, param_.ka); + float tmp = 1; + uint64_t quant_scalar = static_cast(*reinterpret_cast(&tmp)); + mm_.SetQuantScalar(quant_scalar); + + uint32_t realBatchIdx = curCoreBatchStartIdx_ + batchIdx; + mmOffsetA_ = realBatchIdx * (param_.ka + param_.rope_ka); + mmOffsetB_ = realBatchIdx * param_.maxSeqLen * param_.kb; + mmOffsetC_ = realBatchIdx * param_.maxSeqLen; + mm_.SetTensorA(qUnpackGm_[mmOffsetA_], AMatmulType::isTrans); + mm_.SetTensorB(unpackGm_[mmOffsetB_], BMatmulType::isTrans); + mm_.IterateAll(matmulGm_[mmOffsetC_]); // d + + if (param_.supportKeyRope) { + SetFlag(eventIDFIX_MTE2); + WaitFlag(eventIDFIX_MTE2); + // DumpTensor(matmulGm_[mmOffsetC_], 435, 64); + + mmRope_.SetOrgShape(param_.M, seqLen, param_.rope_ka); + mmRope_.SetSingleShape(param_.M, seqLen, param_.rope_ka); + mmRope_.SetQuantScalar(quant_scalar); + mmOffsetARope_ = mmOffsetA_ + param_.ka; + mmOffsetBRope_ = realBatchIdx * param_.maxSeqLen * param_.rope_kb; + mmRope_.SetTensorA(qUnpackGm_[mmOffsetARope_], AMatmulType::isTrans); + mmRope_.SetTensorB(kRopeUnpackGm_[mmOffsetBRope_], BMatmulType::isTrans); + mmRope_.IterateAll(matmulGm_[mmOffsetC_], 1); // c + // DumpTensor(matmulGm_[mmOffsetC_], 445, 64); + } + } + + __aicore__ inline void ComputeTopK(uint32_t batchIdx, uint32_t pingPongFlag) + { + if ASCEND_IS_AIC { + return; + } + + uint32_t realReducedBatchIdx = curCoreBatchStartIdx_ + batchIdx; + uint32_t realBatchIdx = realReducedBatchIdx / param_.head; /* batchIdx without headNum */ + + // Whether the current batch needs to be skipped + if (supportMask_) { + bool batchMask = maskGm_.GetValue(realBatchIdx); + //YF_LOG("realBatchIdx = %d, batchMask = %d\n", realBatchIdx, batchMask); + if (!batchMask) { + // TODO: Directly assign the block table to the output Indices + SetBlockTableForIndices(realBatchIdx, realReducedBatchIdx * param_.maxK); + //YF_LOG("realBatchIdx = %d SetBlockTableForIndices\n", realBatchIdx); + return; + } + } + uint32_t curSeqLen = seqLenGm_.GetValue(realBatchIdx); + uint32_t curK = kGm_.GetValue(realBatchIdx); + uint32_t curChunkSize = 1; + uint32_t curChunkNum = 0; + if (chunkSizeGm_.GetPhyAddr() != nullptr) { + curChunkSize = chunkSizeGm_.GetValue(realBatchIdx); + } + if (curChunkSize != 0) { + if(curChunkSize == 1 || curChunkSize == 8 || curChunkSize == 16) + { + if (curSeqLen <= 32) { + curSeqLen = 0; + } else { + curSeqLen = curSeqLen - Min(curSeqLen, static_cast(curSeqLen % curChunkSize + 16)); + } + curChunkNum = curSeqLen / curChunkSize; + } + else if(curChunkSize == 64) { + curChunkNum = ((curSeqLen + 63) / 64) * 64 / curChunkSize; + } + else if(curChunkSize == 128) + { + curChunkNum = ((curSeqLen + 127)/128)*128 / curChunkSize; + } + + } else { + curSeqLen = 0; + } + if (curK == 0 || curSeqLen == 0) { + return; + } + curK = Min(curK, curChunkNum); + curK = Min(curK, param_.maxK); + uint32_t topKBlockNum = 1; + /* param_.tileN2 > param_.maxK=> tileN2 >= curK */ + uint32_t tileN2 = Min(curChunkNum, param_.tileN2); //param.tileN2=3328, it is very large + + + uint32_t tailN2 = curChunkNum; + if (curK < tileN2) { + topKBlockNum = matmul::CeilDiv(curChunkNum - tileN2, tileN2 - curK) + 1; + uint32_t effectiveTailN2 = (curChunkNum - tileN2) % (tileN2 - curK) > 0 ? + (curChunkNum - tileN2) % (tileN2 - curK) : (tileN2 - curK); + tailN2 = topKBlockNum > 1 ? (effectiveTailN2 + curK) : curChunkNum; + } + + LocalTensor topKOutValueTensor; + LocalTensor topKOutIndexTensor; + uint64_t mmOffset = realReducedBatchIdx * param_.N; + uint32_t headChunkNum = 0; + uint32_t tailChunkNum = 0; + + uint32_t chunkPerBlock = param_.tileN1 / curChunkSize; + + + //uint32_t skipTailChunkNum = DivCeil(SKIP_TAIL_TOKEN_NUM, curChunkSize); + //uint32_t skipHeadChunkNum = DivCeil(SKIP_HEAD_TOKEN_NUM, curChunkSize); + uint32_t skipTailChunkNum = param_.recent; + uint32_t skipHeadChunkNum = param_.sink; + + uint32_t last_block_tail_num = 0; + uint32_t second_last_block_tail_num = 0; + if (tailN2 >= skipTailChunkNum) + { + last_block_tail_num = skipTailChunkNum; + } + else + { + last_block_tail_num = tailN2; + second_last_block_tail_num = skipTailChunkNum - tailN2; + } + + //YF_LOG("batchIdx=%d, curSeqLen=%d, curChunkSize=%d, curChunkNum=%d, tileN2=%d,tailN2=%d,topKBlockNum=%d,skipTailChunkNum=%d,skipHeadChunkNum=%d,last_block_tail_num=%d,second_last_block_tail_num=%d, param_.tileN1=%d, \n", batchIdx, curSeqLen, curChunkSize, curChunkNum, tileN2, tailN2, topKBlockNum, skipTailChunkNum, skipHeadChunkNum, last_block_tail_num, second_last_block_tail_num,param_.tileN1); + + for (uint32_t i = 0; i < topKBlockNum; i++) { + uint32_t copyLen = i == topKBlockNum - 1 ? tailN2 : tileN2; + uint64_t matmulGmOffset = i == 0 ? mmOffset : mmOffset + i * (tileN2 - curK); + GenerateTopKValueTensor(i, copyLen, tileN2, matmulGmOffset, curK, curChunkSize); + GenerateTopKIndexTensor(i, copyLen, tileN2, matmulGmOffset - mmOffset, curK); + LocalTensor topKInValueTensor = topKInValueQueue_.DeQue(); + + uint32_t chunkPerBlock = param_.tileN1 / curChunkSize; // blocksize / chunksize = chunknum per block + if (headChunkNum < skipHeadChunkNum) { + uint32_t curHeadChunkNum = min(skipHeadChunkNum - headChunkNum, copyLen); + headChunkNum += curHeadChunkNum; + //DumpTensor(topKInValueTensor, 400, topKInValueTensor.GetSize()); + Duplicate(topKInValueTensor, static_cast(MAX_HALF_VALUE), curHeadChunkNum); + //DumpTensor(topKInValueTensor, 403, topKInValueTensor.GetSize()); + } + + + if(i == topKBlockNum -2 && second_last_block_tail_num > 0) + { + //uint32_t offset = tileN2 - second_last_block_tail_num; + //Maxs(topKInValueTensor[offset], topKInValueTensor[offset], MAX_HALF_VALUE, second_last_block_tail_num); + FillMaxValueFromTail(topKInValueTensor, tileN2, second_last_block_tail_num, curChunkSize); + } + + if(i == topKBlockNum -1) + { + if(last_block_tail_num == tailN2) + { + //Maxs(topKInValueTensor, topKInValueTensor, MAX_HALF_VALUE, tailN2); + //DumpTensor(topKInValueTensor, 421, topKInValueTensor.GetSize()); + Duplicate(topKInValueTensor, static_cast(MAX_HALF_VALUE), tailN2); + //DumpTensor(topKInValueTensor, 424, topKInValueTensor.GetSize()); + } + else + { + //uint32_t offset = tailN2 - last_block_tail_num; + //Maxs(topKInValueTensor[offset], topKInValueTensor[offset], MAX_HALF_VALUE, last_block_tail_num); + //DumpTensor(topKInValueTensor, 431, topKInValueTensor.GetSize()); + FillMaxValueFromTail(topKInValueTensor, tailN2, last_block_tail_num, curChunkSize); + //DumpTensor(topKInValueTensor, 433, topKInValueTensor.GetSize()); + } + } + + LocalTensor topKInIndexTensor = topKInIndexQueue_.DeQue(); + topKOutValueTensor = topKOutValueQueue_.AllocTensor(); + topKOutIndexTensor = topKOutIndexQueue_.AllocTensor(); + TopKCustom(topKOutValueTensor, topKOutIndexTensor, topKInValueTensor, topKInIndexTensor, curK, tilingData_, copyLen); + topKInValueQueue_.FreeTensor(topKInValueTensor); + topKInIndexQueue_.FreeTensor(topKInIndexTensor); + topKOutValueQueue_.EnQue(topKOutValueTensor); + topKOutIndexQueue_.EnQue(topKOutIndexTensor); + } + + uint64_t topKOutGmOffset = static_cast(realReducedBatchIdx) * param_.maxK; + topKOutValueTensor = topKOutValueQueue_.DeQue(); + topKOutIndexTensor = topKOutIndexQueue_.DeQue(); + if (!param_.supportOffload && (curChunkSize == 64 || curChunkSize == 128)) { + // Map the TopK chunk indices to block_id and write back to GM(indices). + WriteBlockTableFromTopK(realBatchIdx, topKOutIndexTensor, curK, topKOutGmOffset); + } else { + DataCopyExtParams copyOutParams{1, static_cast(curK * sizeof(int32_t)), 0, 0, 0}; + DataCopyPad(indicesGm_[topKOutGmOffset], topKOutIndexTensor, copyOutParams); + } + topKOutValueQueue_.FreeTensor(topKOutValueTensor); + topKOutIndexQueue_.FreeTensor(topKOutIndexTensor); + } + + __aicore__ inline void SetBlockTableForIndices(uint32_t curBatchIdx, uint64_t outGmOffset) { + // Find the tableblock corresponding to the current batch + LocalTensor tableBlockTensor = tableBlockBuf_.template Get(); + DataCopyExtParams copyInParams{1, static_cast(param_.blockCount * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams copyInPadParams{false, 0, 0, 0}; + DataCopyPad(tableBlockTensor, keyBlockTableGm_[curBatchIdx * param_.blockCount], copyInParams, copyInPadParams); + //DumpTensor(tableBlockTensor, 500, 64); + SetFlag(0); + WaitFlag(0); + + uint32_t copyLen = param_.blockCount < param_.maxK ? param_.blockCount : param_.maxK; + DataCopyExtParams cpOut{1, static_cast(copyLen * sizeof(int32_t)), 0, 0, 0}; + DataCopyPad(indicesGm_[outGmOffset], tableBlockTensor, cpOut); + } + + __aicore__ inline void GenerateTopKValueTensor(uint32_t i, uint32_t copyLen, + uint32_t tileN2, uint64_t matmulGmOffset, uint32_t curK, uint32_t chunkSize) + { + LocalTensor topKInValueTensor = topKInValueQueue_.AllocTensor(); + uint32_t copyLenAligned = copyLen / BLOCK_CUBE * BLOCK_CUBE; /* floor aligned for datacopy */ + if (copyLenAligned < param_.tileN2) { + Duplicate(topKInValueTensor[copyLenAligned], static_cast(MIN_HALF_VALUE), param_.tileN2 - copyLenAligned); + SetFlag(0); + WaitFlag(0); + } + + if (chunkSize > 1) { + uint16_t chunkNum = static_cast(copyLen); + LocalTensor reduceInputTensor = topKInIndexQueue_.AllocTensor().ReinterpretCast(); + ReduceMaxCustom(matmulGm_[matmulGmOffset], reduceInputTensor, topKInValueTensor, chunkNum, static_cast(chunkSize)); + topKInValueQueue_.EnQue(topKInValueTensor); + topKInIndexQueue_.FreeTensor(reduceInputTensor); + } else { + uint32_t copyLenCeilAligned = matmul::CeilDiv(copyLen * sizeof(half), BLOCK_CUBE) + * BLOCK_CUBE / sizeof(half); + DataCopyExtParams copyInParams{1, static_cast(copyLen * sizeof(half)), 0, 0, 0}; + DataCopyPadExtParams copyInPadParams{true, 0, static_cast(copyLenCeilAligned - copyLen), + static_cast(MIN_HALF_VALUE)}; + DataCopyPad(topKInValueTensor, matmulGm_[matmulGmOffset], copyInParams, copyInPadParams); + topKInValueQueue_.EnQue(topKInValueTensor); + if (i > 0) { + LocalTensor topKOutValueTensor = topKOutValueQueue_.DeQue(); + topKInValueTensor = topKInValueQueue_.DeQue(); + uint64_t valueMask = curK > MAX_FP16_PROCESS_NUM ? MAX_FP16_PROCESS_NUM : curK; + uint8_t valueRepeatTimes = curK / MAX_FP16_PROCESS_NUM; + if (valueRepeatTimes > 0) { + Copy(topKInValueTensor, topKOutValueTensor, valueMask, valueRepeatTimes, {1, 1, 8, 8}); + } + if (curK % MAX_FP16_PROCESS_NUM != 0) { + Copy(topKInValueTensor[valueRepeatTimes * MAX_FP16_PROCESS_NUM], + topKOutValueTensor[valueRepeatTimes * MAX_FP16_PROCESS_NUM], + curK % MAX_FP16_PROCESS_NUM, 1, {1, 1, 8, 8}); + } + PipeBarrier(); + topKInValueQueue_.EnQue(topKInValueTensor); + topKOutValueQueue_.FreeTensor(topKOutValueTensor); + } + } + } + + __aicore__ inline void GenerateTopKIndexTensor(uint32_t i, uint32_t copyLen, + uint32_t tileN2, uint64_t startIndex, uint32_t curK) + { + LocalTensor topKInIndexTensor = topKInIndexQueue_.AllocTensor(); + ArithProgression(topKInIndexTensor, static_cast(startIndex), 1, static_cast(copyLen)); + topKInIndexQueue_.EnQue(topKInIndexTensor); + if (i > 0) { + LocalTensor topKOutIndexTensor = topKOutIndexQueue_.DeQue(); + topKInIndexTensor = topKInIndexQueue_.DeQue(); + uint64_t indexMask = curK > MAX_INT32_PROCESS_NUM ? MAX_INT32_PROCESS_NUM : curK; + uint8_t indexRepeatTimes = curK / MAX_INT32_PROCESS_NUM; + if (indexRepeatTimes > 0) { + Copy(topKInIndexTensor, topKOutIndexTensor, indexMask, indexRepeatTimes, {1, 1, 8, 8}); + } + if (curK % MAX_INT32_PROCESS_NUM != 0) { + Copy(topKInIndexTensor[indexRepeatTimes * MAX_INT32_PROCESS_NUM], + topKOutIndexTensor[indexRepeatTimes * MAX_INT32_PROCESS_NUM], + curK % MAX_INT32_PROCESS_NUM, 1, {1, 1, 8, 8}); + } + PipeBarrier(); + topKInIndexQueue_.EnQue(topKInIndexTensor); + topKOutIndexQueue_.FreeTensor(topKOutIndexTensor); + } + } + + template + __aicore__ inline void SyncAicOnly(uint16_t eventId) { + CrossCoreSetFlag(eventId); + CrossCoreWaitFlag(eventId); + } + + template + __aicore__ inline void SyncAivOnly(uint16_t eventId) { + CrossCoreSetFlag(eventId); + CrossCoreWaitFlag(eventId); + } + + template + __aicore__ inline void VectorNotifyCube(uint16_t aiv2AicEventId) { + CrossCoreSetFlag(aiv2AicEventId); + } + + __aicore__ inline void CubeWaitVector(uint16_t aiv2AicEventId) { + CrossCoreWaitFlag(aiv2AicEventId); + } + + template + __aicore__ inline void CubeNotifyVector(uint16_t aic2AivEventId) { + CrossCoreSetFlag(aic2AivEventId); + } + + __aicore__ inline void VectorWaitCube(uint16_t aic2AivEventId) { + CrossCoreWaitFlag(aic2AivEventId); + } + +protected: + uint64_t innerSplitLoopTimes_ = 0; + uint8_t innerSplitGMFlag_ = 0; + + static constexpr uint32_t BLOCK_CUBE = 32; + static constexpr uint64_t SYNC_MODE0 = 0; + static constexpr uint64_t SYNC_MODE2 = 2; + static constexpr uint64_t SYNC_AIC_ONLY_ALL_FLAG = 1; + static constexpr uint64_t SYNC_AIV_AIC_FLAG = 2; + static constexpr uint64_t SYNC_AIC_AIV_FLAG = 4; + static constexpr uint64_t SYNC_AIV_AIC_FLAG2 = 6; + static constexpr uint64_t SYNC_AIC_AIV_FLAG2 = 0; + static constexpr uint32_t DOUBLE_BUFFER_NUM = 2; + static constexpr uint32_t COMPRESSED_DIMENSION = 16; + static constexpr uint32_t BATCH_PING_PONG_NUM = 8; // The maximum depth of the inter-core synchronization flag is 15. + static constexpr uint32_t SUB_BLOCK_NUM = 2; + static constexpr uint32_t SUB_BLOCK_NUM_WITH_DB = 4; + static constexpr float MIN_HALF_VALUE = -65535; + int32_t eventIDFIX_MTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::FIX_MTE2)); + + static constexpr uint32_t MAX_SELECT_REPEATED_TIMES = 254; + + GlobalTensor queryGm_; + GlobalTensor keyGm_; + GlobalTensor keyRopeGm_; + GlobalTensor kGm_; + GlobalTensor seqLenGm_; + GlobalTensor chunkSizeGm_; + GlobalTensor keyBlockTableGm_; + GlobalTensor maskGm_; + GlobalTensor indicesGm_; + GlobalTensor unpackGm_; + GlobalTensor kRopeUnpackGm_; + GlobalTensor matmulGm_; + GlobalTensor qUnpackGm_; + + TPipe *pipe_; + TQue keyCompressedInQueue_; + TQue keyUnpackedOutQueue_; + TQue topKInValueQueue_; + TQue topKInIndexQueue_; + TQue topKOutValueQueue_; + TQue topKOutIndexQueue_; + TBuf constBuf_; + TBuf selectBuf_; + TBuf indexBuf_; + TQue queryCompressedInQueue_; + TQue queryUnpackedOutQueue_; + TBuf qReduceSumLastRowBuf_; + TBuf qReduceSumBuf_; + + TBuf tableBlockBuf_; + + TilingParam param_; + TopkTiling topkTiling_; + HammingDistTopKTilingData tilingData_; + + using AMatmulType = matmul::MatmulType; + using BMatmulType = matmul::MatmulType; + using BiasMatmulType = matmul::MatmulType; + // notice: the TPos of ctype must be ub given by mm api when iterate, + // but actually we can move data to gm then to ub. + using CMatmulType = matmul::MatmulType; + matmul::MatmulImpl mm_; + matmul::MatmulImpl mmRope_; + + uint64_t mmOffsetA_; + uint64_t mmOffsetB_; + uint64_t mmOffsetARope_; + uint64_t mmOffsetBRope_; + uint64_t mmOffsetC_; + uint32_t curCoreBatch_; + uint32_t curCoreBatchStartIdx_; + bool continFlag_; + bool supportMask_ = true; + + __aicore__ inline void InitParams() + { + mmOffsetA_ = 0; + mmOffsetB_ = 0; + mmOffsetC_ = 0; + + param_.preCoreNum = param_.reducedBatch % param_.usedCoreNum; + if (param_.preCoreNum == 0) { + param_.preCoreNum = param_.usedCoreNum; + } + uint32_t blockIndex = AscendC::GetBlockIdx(); + if ASCEND_IS_AIV { + blockIndex = blockIndex / 2; + } + if (blockIndex < param_.preCoreNum) { + curCoreBatch_ = param_.singleCoreBatch; + } else { + curCoreBatch_ = param_.singleCoreBatch - 1; + } + + curCoreBatchStartIdx_ = blockIndex * curCoreBatch_; + if (blockIndex >= param_.preCoreNum) { + curCoreBatchStartIdx_ += param_.preCoreNum; + } + } + + __aicore__ inline void InitTilingParams(const TCubeTiling &tiling, const TCubeTiling &tilingRope, const TopkTiling &topkTiling, + const HammingDistTopKTilingParams &tilingParam) + { + // tiling data for select + param_.usedCoreNum = tilingParam.usedCoreNum; + param_.batch = tilingParam.batch; + param_.head = tilingParam.head; + param_.dimension = tilingParam.dimension; + param_.nope_dimension = tilingParam.nope_dimension; + param_.rope_dimension = tilingParam.rope_dimension; + param_.reducedBatch = tilingParam.reducedBatch; + param_.tileN1 = tilingParam.tileN1; + param_.tileN2 = tilingParam.tileN2; + param_.singleCoreBatch = tilingParam.singleCoreBatch; + param_.qHead = tilingParam.qHead; + param_.headGroupNum = tilingParam.headGroupNum; + + param_.maxK = tilingParam.maxK; + + // support key rope + param_.supportKeyRope = tilingParam.supportKeyRope > 0; + + // tiling data for matmul + param_.M = tiling.M; + param_.N = tiling.N; + param_.ka = tiling.Ka; + param_.kb = tiling.Kb; + if (param_.supportKeyRope) { + param_.rope_ka = tilingRope.Ka; + param_.rope_kb = tilingRope.Kb; + } + + // tiling data for topk + param_.mmGmOffset = tilingParam.mmGmOffset; + param_.qUnpackGmOffset = tilingParam.qUnpackGmOffset; + param_.kNopeUnpackGmOffset = tilingParam.kNopeUnpackGmOffset; + topkTiling_ = topkTiling; + param_.maxSeqLen = tilingParam.maxSeqLen; + param_.sink = tilingParam.sink; + param_.recent = tilingParam.recent; + param_.blockCount = tilingParam.blockCount; + + // support offload + param_.supportOffload = tilingParam.supportOffload > 0; + } + + __aicore__ inline void InitGlobalBuffers(GM_ADDR query, GM_ADDR keyCompressed, GM_ADDR keyCompressedRope, GM_ADDR k, GM_ADDR seqLen, + GM_ADDR chunkSize, GM_ADDR keyBlockTable, GM_ADDR mask,GM_ADDR indices, GM_ADDR workSpace) + { + queryGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(query)); + keyGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(keyCompressed)); + keyRopeGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(keyCompressedRope)); + kGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(k)); + seqLenGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(seqLen)); + chunkSizeGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(chunkSize)); + keyBlockTableGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(keyBlockTable)); + maskGm_.SetGlobalBuffer(reinterpret_cast<__gm__ bool*>(mask)); + supportMask_ = maskGm_.GetPhyAddr() != nullptr; + indicesGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(indices)); + unpackGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int4b_t*>(workSpace)); + kRopeUnpackGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int4b_t*>(workSpace + param_.kNopeUnpackGmOffset)); + qUnpackGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int4b_t*>(workSpace + param_.qUnpackGmOffset)); + matmulGm_.SetGlobalBuffer(reinterpret_cast<__gm__ half*>(workSpace + param_.mmGmOffset)); + } + + __aicore__ inline void InitLocalBuffersForUnpackKey() + { + pipe_->InitBuffer(keyCompressedInQueue_, 1, param_.tileN1 * (param_.dimension / COMPRESS_RATE) * sizeof(int8_t)); /* 8: original dimension / compressed dimension */ + pipe_->InitBuffer(keyUnpackedOutQueue_, 1, param_.tileN1 * param_.dimension * sizeof(int8_t) / 2); /* 2: 1 / sizeof(int4b_t) */ + pipe_->InitBuffer(constBuf_, param_.dimension * sizeof(half)); + uint64_t selectRepeatedTimes = computeSelectRepeatedTimes(param_.dimension); + if (selectRepeatedTimes > MAX_SELECT_REPEATED_TIMES) { + pipe_->InitBuffer(selectBuf_, MAX_FP16_PROCESS_NUM * MAX_SELECT_REPEATED_TIMES * sizeof(half)); + } else { + pipe_->InitBuffer(selectBuf_, param_.tileN1 * param_.dimension * sizeof(half)); + } + } + + __aicore__ inline uint64_t computeSelectRepeatedTimes(uint32_t dimension) { + uint64_t selectElmentCount = param_.tileN1 * dimension; + uint64_t selectRepeatedTimes = (selectElmentCount + MAX_FP16_PROCESS_NUM - 1) / MAX_FP16_PROCESS_NUM; + return selectRepeatedTimes; + } + + __aicore__ inline void InitLocalBuffersForUnpackQuery() + { + pipe_->InitBuffer(constBuf_, param_.dimension * sizeof(half)); + pipe_->InitBuffer(selectBuf_, param_.headGroupNum * param_.dimension * sizeof(half)); + pipe_->InitBuffer(queryCompressedInQueue_, 1, param_.headGroupNum * (param_.dimension / 8) * sizeof(int8_t)); /* 8: original dimension / compressed dimension */ + pipe_->InitBuffer(queryUnpackedOutQueue_, 1, param_.headGroupNum * param_.dimension * sizeof(int8_t) / 2); /* 2: 1 / sizeof(int4b_t) */ + pipe_->InitBuffer(qReduceSumLastRowBuf_, param_.headGroupNum * param_.dimension * sizeof(half)); + pipe_->InitBuffer(qReduceSumBuf_, param_.headGroupNum * param_.dimension * sizeof(half)); + } + + __aicore__ inline void InitLocalBuffersForTopK() { + pipe_->InitBuffer(topKInValueQueue_, 1, param_.tileN2 * sizeof(half)); + pipe_->InitBuffer(topKInIndexQueue_, 1, param_.tileN2 * 8 * sizeof(int32_t)); /* 8: maximum chunkSize * sizeof(int32) / sizeof(half) */ + pipe_->InitBuffer(topKOutValueQueue_, 1, param_.maxK * sizeof(half)); + pipe_->InitBuffer(topKOutIndexQueue_, 1, param_.maxK * sizeof(int32_t)); + pipe_->InitBuffer(tableBlockBuf_, param_.blockCount * sizeof(int32_t)); + } + + __aicore__ inline void WriteBlockTableFromTopK( + uint32_t curBatchIdx, + LocalTensor& topKIndexUb, + uint32_t curKScalar, + uint64_t outGmOffset) + { + if ASCEND_IS_AIC { return; } + // Reuse the existing int32 queue to allocate a UB as a write-back intermediate buffer. + LocalTensor blockIdUb = topKInIndexQueue_.AllocTensor(); + + LocalTensor tableBlockTensor = tableBlockBuf_.template Get(); + DataCopyExtParams copyInParams{1, static_cast(param_.blockCount * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams copyInPadParams{false, 0, 0, 0}; + DataCopyPad(tableBlockTensor, keyBlockTableGm_[curBatchIdx * param_.blockCount], copyInParams, copyInPadParams); + + ::AscendC::WriteBlockTableFromTopK(curBatchIdx, topKIndexUb, blockIdUb, curKScalar, outGmOffset, + tableBlockTensor, indicesGm_, continFlag_, param_.blockCount); + + topKInIndexQueue_.FreeTensor(blockIdUb); + } +}; +} // namespace AscendC +#endif // HAMMING_DIST_TOP_K_PARALLEL_H \ No newline at end of file diff --git a/csrc/ascend/moe/hamming_dist_top_k/op_kernel/hamming_dist_top_k_split_s.h b/csrc/ascend/moe/hamming_dist_top_k/op_kernel/hamming_dist_top_k_split_s.h new file mode 100644 index 000000000..cfdb78cad --- /dev/null +++ b/csrc/ascend/moe/hamming_dist_top_k/op_kernel/hamming_dist_top_k_split_s.h @@ -0,0 +1,1317 @@ +/** + * Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/*! + * \file hamming_dist_top_k_split_s.h + * \brief + */ +#ifndef HAMMING_DIST_TOP_K_SPLIT_S_H +#define HAMMING_DIST_TOP_K_SPLIT_S_H + +#define KEEP_TAIL_TYPE_ONE_BLOCK 0 +#define KEEP_TAIL_TYPE_TWO_BLOCKS 1 +#define KEEP_TAIL_TYPE_TWO_TILES 2 + +#include +#include "hamming_dist_top_k_base.h" + +namespace AscendC { + +constexpr int32_t DOUBLE_BUFFER_NUM = 2; +constexpr uint32_t MAX_SELECT_AND_CAST_COUNT = 254; +constexpr uint32_t RESET_NUM = 0U; +constexpr uint32_t COMPRESS_RATE = 8; +constexpr uint32_t VECTOR_CUBE_RATIO = 2; +constexpr uint32_t INT4B_TYPE_SIZE_DIV_RATE = 2; +constexpr uint64_t CAST_MASK = 128; +constexpr uint32_t MAX_BATCH_SIZE = 150; +constexpr uint32_t MAX_CHUNK_SIZE = 16; +constexpr uint32_t MIN_CHUNK_SIZE = 1; +constexpr uint32_t CHUNK_TOPK_MIN_SEQ_LEN = 32; + +class HammingDistTopKSplitSKernel { +public: + __aicore__ inline HammingDistTopKSplitSKernel() {} + __aicore__ inline void Init(GM_ADDR query, GM_ADDR keyCompressed, GM_ADDR keyCompressedRope, GM_ADDR k, + GM_ADDR seqLen, GM_ADDR chunkSize, GM_ADDR keyBlockTable, GM_ADDR mask, + GM_ADDR indices, GM_ADDR workSpace, + const HammingDistTopKTilingData &tilingData, TPipe *que) + { + const TCubeTiling &tiling = tilingData.matmulTiling; + const TCubeTiling &tilingRope = tilingData.matmulTilingRope; + const TopkTiling &topkTiling = tilingData.topkTiling; + const HammingDistTopKTilingParams &tilingParam = tilingData.params; + pipe_ = que; + tilingData_ = tilingData; + InitTilingParams(tiling, tilingRope, topkTiling, tilingParam); + InitParams(); + InitGlobalBuffers(query, keyCompressed, keyCompressedRope, k, seqLen, chunkSize, keyBlockTable, mask, indices, workSpace); + mm_.SetSubBlockIdx(0); + mm_.Init(&tiling, pipe_); + if (param_.supportKeyRope) { + mmRope_.SetSubBlockIdx(0); + mmRope_.Init(&tilingRope, pipe_); + } + } + + __aicore__ inline void Process() + { + ComputeChunkTopKEffectLen(); + + ComputeBatchSeqLenTileInfo(); + + if ASCEND_IS_AIV { + InitLocalBuffersForUnpackQ(); + UnpackQueryCompressed(); + pipe_->Reset(); + InitLocalBuffersForUnpackK(); + UnpackKeyCompressed(false); + if (param_.supportKeyRope) { + UnpackKeyCompressed(true); + } + VectorNotifyCube(SYNC_AIV_ONLY_ALL_FLAG, SYNC_AIV_AIC_FLAG); + pipe_->Reset(); + InitLocalBuffersForTopK(); + if (!param_.supportOffload) { + InitLocalBuffersForTopKSort(); + } + InitLocalBuffersForTableBlock(); + } + + if ASCEND_IS_AIC { + CubeWaitVector(SYNC_AIV_AIC_FLAG); + } + + if ASCEND_IS_AIV { + blockIdx_ = blockIdx_ / VECTOR_CUBE_RATIO; + } + + ComputeMMTopKBatchNGroup(); + + if ASCEND_IS_AIV { + SyncAivOnly(SYNC_AIV_ONLY_ALL_FLAG); + MergeTopKInGroup(); + } + } + + __aicore__ inline void ProcessL2cache() + { + } + +protected: + __aicore__ inline void UnpackQueryCompressed() + { + if ASCEND_IS_AIC { + return; + } + LocalTensor constTensor = constBuf_.template Get(), selectTensor = selectBuf_. template Get(); + LocalTensor qReduceSumTensor = qReduceSumBuf_. template Get(); + LocalTensor qReduceSumLastRowTensor = qReduceSumLastRowBuf_. template Get(); + Duplicate(constTensor, 1, param_.dimension); + uint8_t repeatTimes = (param_.dimension + MAX_FP16_PROCESS_NUM - 1)/ MAX_FP16_PROCESS_NUM; + uint64_t qMask = MAX_FP16_PROCESS_NUM; + uint32_t compressedDimension = param_.dimension / COMPRESS_RATE; + uint32_t vCoreNum = 2 * uint32_t(param_.usedCoreNum); + uint32_t pad_dim = param_.rope_dimension / 2; + uint32_t valid_dim = param_.nope_dimension + pad_dim; + for (uint32_t curBatch = blockIdx_; curBatch < param_.batchN; curBatch += vCoreNum) { + uint64_t queryGmOffset = curBatch * param_.headGroupNum * compressedDimension; + LocalTensor queryCompressed = queryCompressedInQueue_.AllocTensor(); + DataCopyExtParams queryCopyInParams{1, param_.headGroupNum * compressedDimension, 0, 0, 0}; + DataCopyPadExtParams queryCopyInPadParams{false, 0, 0, 0}; + DataCopyPad(queryCompressed, queryGm_[queryGmOffset], queryCopyInParams, queryCopyInPadParams); + + queryCompressedInQueue_.EnQue(queryCompressed); + queryCompressed = queryCompressedInQueue_.DeQue(); + PipeBarrier(); + uint8_t selectRepeatTimes = param_.headGroupNum * param_.dimension / MAX_FP16_PROCESS_NUM; + SelectCustom(selectTensor, queryCompressed, constTensor, selectRepeatTimes); + PipeBarrier(); + if (param_.supportKeyRope) { + for (uint32_t i = 0; i < param_.headGroupNum; i++) { + // DumpTensor(selectTensor[i * param_.dimension], 139, 672); + Duplicate(selectTensor[i * param_.dimension + valid_dim], 0, pad_dim); + // DumpTensor(selectTensor[i * param_.dimension], 141, 672); + } + PipeBarrier(); + } + + queryCompressedInQueue_.FreeTensor(queryCompressed); + + LocalTensor qHashTensor = selectTensor; + if (param_.headGroupNum > 0) { + static constexpr AscendC::CumSumConfig cumSumConfig{false, false, true}; + const AscendC::CumSumInfo cumSumInfo{param_.headGroupNum, param_.dimension}; + AscendC::CumSum(qReduceSumTensor, qReduceSumLastRowTensor, selectTensor, cumSumInfo); + PipeBarrier(); + // DumpTensor(qReduceSumTensor, 152, 2048); + if (param_.headGroupNum > 8) { + uint32_t div = matmul::CeilDiv(param_.headGroupNum, 8); + half reciprocalDiv = static_cast((float)1.0 / div); + AscendC::Muls(qReduceSumLastRowTensor, qReduceSumTensor[(param_.headGroupNum - 1) * param_.dimension], + reciprocalDiv, qMask, repeatTimes, {1, 1, 8, 8}); + PipeBarrier(); + } + qHashTensor = qReduceSumLastRowTensor; + } + + LocalTensor queryUnpacked = queryUnpackedOutQueue_.AllocTensor(); + Cast(queryUnpacked, qHashTensor, RoundMode::CAST_CEIL, qMask, repeatTimes, {1, 1, 2, 8}); + PipeBarrier(); + queryUnpackedOutQueue_.EnQue(queryUnpacked); + queryUnpacked = queryUnpackedOutQueue_.DeQue(); + uint64_t unpackQGmOffset = queryGmOffset * COMPRESS_RATE / param_.headGroupNum; + DataCopyExtParams copyQOutParams{1, static_cast(param_.dimension / 2), 0, 0, 0}; /* 2: 1 / size of int4b_t */ + DataCopyPad(qUnpackGm_[unpackQGmOffset], queryUnpacked, copyQOutParams); + queryUnpackedOutQueue_.FreeTensor(queryUnpacked); + + SetFlag(0); + WaitFlag(0); + } + } + + __aicore__ inline void DataCopyInForKeyCompressed( + bool useCopyPad, uint32_t copySeqLen, uint32_t compressedDimension, LocalTensor &keyCompressedLocal, uint64_t keyGmOffset, GlobalTensor& keyCompressedGm) + { + if ASCEND_IS_AIC { + return; + } + if (useCopyPad) { + // copySeqLen <= TileN1, compressedDimension is fixed, block data is within uint16 range. + DataCopyExtParams copyInParams{1, static_cast(copySeqLen * compressedDimension), 0, 0, 0}; + DataCopyPadExtParams copyInPadParams{false, 0, 0, 0}; + DataCopyPad(keyCompressedLocal, keyCompressedGm[keyGmOffset], copyInParams, copyInPadParams); + } else { + DataCopy(keyCompressedLocal, keyCompressedGm[keyGmOffset], copySeqLen * compressedDimension); + } + } + + __aicore__ inline void DataCopyOutForKeyUnpacked(uint32_t copySeqLen, LocalTensor &keyUnpackedLocal, uint64_t keyGmOffset) + { + if ASCEND_IS_AIC { + return; + } + // copySeqLen <= TileN1，param_.dimension is fixed, block data is within uint16 range. + DataCopyParams copyParams{1, static_cast(copySeqLen * param_.dimension / 2 / BLOCK_CUBE), 0, 0}; // 2: 1/2, size of int4b_t + DataCopy(unpackGm_[keyGmOffset], keyUnpackedLocal, copyParams); // output to outQueue1 with DB_ON + } + + __aicore__ inline void UnpackKeyCompressed(bool isKeyRope) + { + if ASCEND_IS_AIC { + return; + } + uint32_t dimension = isKeyRope ? param_.rope_dimension : param_.nope_dimension;; + GlobalTensor unpackGm = isKeyRope ? kRopeUnpackGm_ : unpackGm_; + GlobalTensor keyCompressedGm = isKeyRope ? keyCompressedRopeGm_ : keyCompressedGm_; + LocalTensor constTensor = constBuf_.template Get(), selectTensor = selectBuf_.template Get(); + Duplicate(constTensor, 1, dimension); + uint32_t curBatchIdx = 0, curHeadIdx = 0, seqLens = 0, kScalar = 0, effectLen = 0, allSLoops = 0, sloops = 0, minSloops = 0, maxSloops = 0, tailSloops = 0, tailSeqLens = 0, preSloops = 0, blockIdx = 0, selectAndCastCount = 0, selectAndCastOffset = 0, compressedDimension = dimension / COMPRESS_RATE, vCoreNum = 2 * uint32_t(param_.usedCoreNum); + uint64_t keyUnpackInGmOffset = 0; + bool hasTailSeqLens = false; + for (uint32_t i = 0; i < param_.batchN; i++) { + curBatchIdx = i / param_.head; + if (supportMask_) { + bool batchMask = maskGm_.GetValue(curBatchIdx); + //YF_LOG("curBatchIdx = %d, batchMask = %d\n", curBatchIdx, batchMask); + if (!batchMask) { + continue; + } + } + curHeadIdx = i % param_.head; + seqLens = uint32_t(seqLenGm_.GetValue(curBatchIdx)); + kScalar = uint32_t(kGm_.GetValue(curBatchIdx)); + effectLen = effectLenArr_[curBatchIdx]; + if (seqLens == 0 || kScalar == 0 || effectLen == 0) { + continue; + } + allSLoops = DivCeil(seqLens, param_.tileN1), tailSeqLens = seqLens % param_.tileN1, hasTailSeqLens = tailSeqLens != 0, tailSloops = allSLoops % vCoreNum, minSloops = allSLoops / vCoreNum, maxSloops = tailSloops == 0 ? minSloops : minSloops + 1; + if (blockIdx_ < tailSloops) { + sloops = maxSloops, preSloops = blockIdx_ * sloops; + } else { + sloops = minSloops, preSloops = tailSloops == 0 ? blockIdx_ * minSloops : tailSloops * maxSloops + (blockIdx_ - tailSloops) * minSloops; + } + bool isTailSloopOwner = preSloops + sloops == allSLoops, isTailSeqLenNotAlignedForUint8 = isTailSloopOwner && (tailSeqLens % 2 != 0); + for (uint32_t j = 0; j < sloops; j++) { + bool isLastLoop = j == sloops - 1, useDataCopyPadForKeyCompressed = isLastLoop && isTailSeqLenNotAlignedForUint8; + uint32_t copySeqLen = (isTailSloopOwner && isLastLoop && hasTailSeqLens) ? tailSeqLens : param_.tileN1, selectAndCastLoops = copySeqLen / MAX_SELECT_AND_CAST_COUNT, selectAndCastTail = copySeqLen % MAX_SELECT_AND_CAST_COUNT; + if (isContinuousBatch_) { + blockIdx = keyBlockTableGm_.GetValue(curBatchIdx * param_.blockCount + preSloops + j), keyUnpackInGmOffset = (blockIdx * param_.head + curHeadIdx) * param_.tileN1 * compressedDimension; + } else { + keyUnpackInGmOffset = (i * param_.maxSeqLen + (preSloops + j) * param_.tileN1) * compressedDimension; + } + LocalTensor keyCompressedLocal = keyCompressedInBuf_.AllocTensor(); + LocalTensor keyUnpackedLocal = keyUnpackedOutBuf_.AllocTensor(); + DataCopyInForKeyCompressed(useDataCopyPadForKeyCompressed, copySeqLen, compressedDimension, keyCompressedLocal, keyUnpackInGmOffset, keyCompressedGm); + keyCompressedInBuf_.EnQue(keyCompressedLocal); + keyCompressedLocal = keyCompressedInBuf_.DeQue(); + uint64_t keyUnpackOutGmOffset = (i * param_.maxSeqLen + (preSloops + j) * param_.tileN1) * dimension; + // support key rope + if (dimension > 128 || isKeyRope) { + uint32_t tileN1RepeatTimes = (copySeqLen * dimension + MAX_FP16_PROCESS_NUM - 1) / MAX_FP16_PROCESS_NUM; + selectAndCastLoops = (tileN1RepeatTimes + MAX_SELECT_AND_CAST_COUNT - 1) / MAX_SELECT_AND_CAST_COUNT; + for (uint32_t k = 0; k < selectAndCastLoops; k++) { + uint32_t localStart = k * MAX_SELECT_AND_CAST_COUNT; + uint32_t remainRepeats = tileN1RepeatTimes - localStart; + uint8_t repeatTimes = remainRepeats > MAX_SELECT_AND_CAST_COUNT ? MAX_SELECT_AND_CAST_COUNT : remainRepeats; + uint32_t elemOffset = localStart * MAX_FP16_PROCESS_NUM; + uint32_t maskOffset = elemOffset / COMPRESS_RATE; + // YF_LOG("localStart=%d, remainRepeats=%d, repeatTimes=%d, elemOffset=%d, maskOffset=%d\n", localStart, remainRepeats, repeatTimes, elemOffset, maskOffset); + if (repeatTimes > 0) { + SelectCustom(selectTensor, keyCompressedLocal[maskOffset], constTensor, repeatTimes); + Cast(keyUnpackedLocal[elemOffset], selectTensor, RoundMode::CAST_CEIL, CAST_MASK, repeatTimes, {1, 1, 2, 8}); + } + } + keyUnpackedOutBuf_.EnQue(keyUnpackedLocal); + keyUnpackedLocal = keyUnpackedOutBuf_.DeQue(); + DataCopyParams copyParams{1, static_cast(tileN1RepeatTimes * MAX_FP16_PROCESS_NUM / 2 / BLOCK_CUBE), 0, 0}; // 2: 1/2, size of int4b_t + DataCopy(unpackGm[keyUnpackOutGmOffset], keyUnpackedLocal, copyParams); // output to outQueue1 with DB_ON + keyCompressedInBuf_.FreeTensor(keyCompressedLocal); + keyUnpackedOutBuf_.FreeTensor(keyUnpackedLocal); + } else { + bool isSelectAndCastHasTail = selectAndCastTail != 0; + selectAndCastLoops = isSelectAndCastHasTail ? selectAndCastLoops + 1 : selectAndCastLoops; + for (uint32_t k = 0; k < selectAndCastLoops; k++) { + selectAndCastCount = (isSelectAndCastHasTail && k == selectAndCastLoops - 1) ? selectAndCastTail : MAX_SELECT_AND_CAST_COUNT, selectAndCastOffset = k * MAX_SELECT_AND_CAST_COUNT * dimension; + SelectCustom(selectTensor[selectAndCastOffset], keyCompressedLocal[selectAndCastOffset / COMPRESS_RATE], constTensor, static_cast(selectAndCastCount)); + Cast(keyUnpackedLocal[selectAndCastOffset], selectTensor[selectAndCastOffset], RoundMode::CAST_CEIL, CAST_MASK, static_cast(selectAndCastCount), {1, 1, 2, 8}); + } + keyCompressedInBuf_.FreeTensor(keyCompressedLocal); + keyUnpackedOutBuf_.EnQue(keyUnpackedLocal); + keyUnpackedLocal = keyUnpackedOutBuf_.DeQue(); + DataCopyOutForKeyUnpacked(copySeqLen, keyUnpackedLocal, keyUnpackOutGmOffset); + keyUnpackedOutBuf_.FreeTensor(keyUnpackedLocal); + } + } + } + } + + __aicore__ inline void ComputeMM(uint32_t batchIdx, uint32_t curHeadIdx, uint32_t curCoreDealSize, uint32_t tileIdx) + { + if ASCEND_IS_AIV { + return; + } + + mm_.SetOrgShape(param_.M, curCoreDealSize, param_.ka); + mm_.SetSingleShape(param_.M, curCoreDealSize, param_.ka); + float tmp = 1; + uint64_t quant_scalar = static_cast(*reinterpret_cast(&tmp)); + mm_.SetQuantScalar(quant_scalar); + mmOffsetA_ = batchIdx * 1 * param_.head * (param_.ka + param_.rope_ka) + curHeadIdx * (param_.ka + param_.rope_ka); + mmOffsetB_ = (batchIdx * param_.maxSeqLen * param_.head + curHeadIdx * param_.maxSeqLen + tileIdx * tileSeqLenSize_) * param_.kb; + mmOffsetC_ = batchIdx * param_.maxSeqLen * param_.head + curHeadIdx * param_.maxSeqLen + tileIdx * tileSeqLenSize_; + mm_.SetTensorA(qUnpackGm_[mmOffsetA_], AMatmulType::isTrans); + mm_.SetTensorB(unpackGm_[mmOffsetB_], BMatmulType::isTrans); + mm_.IterateAll(matmulGm_[mmOffsetC_]); // d + + if (param_.supportKeyRope) { + SetFlag(eventIDFIX_MTE2); + WaitFlag(eventIDFIX_MTE2); + // DumpTensor(matmulGm_[mmOffsetC_], 377, curCoreDealSize); + + mmRope_.SetOrgShape(param_.M, curCoreDealSize, param_.rope_ka); + mmRope_.SetSingleShape(param_.M, curCoreDealSize, param_.rope_ka); + mmRope_.SetQuantScalar(quant_scalar); + mmOffsetARope_ = mmOffsetA_ + param_.ka; + mmOffsetBRope_ = (batchIdx * param_.maxSeqLen * param_.head + curHeadIdx * param_.maxSeqLen + tileIdx * tileSeqLenSize_) * param_.rope_kb; + mmRope_.SetTensorA(qUnpackGm_[mmOffsetARope_], AMatmulType::isTrans); + mmRope_.SetTensorB(kRopeUnpackGm_[mmOffsetBRope_], BMatmulType::isTrans); + mmRope_.IterateAll(matmulGm_[mmOffsetC_], 1); // c + // DumpTensor(matmulGm_[mmOffsetC_], 387, curCoreDealSize); + } + } + + __aicore__ inline void ComputeBatchSeqLenTileInfo() + { + uint32_t batchSeqLenTileNum = 0; + for (uint32_t batchIdx = 0; batchIdx < param_.batch; batchIdx++){ + uint32_t curSeqLen = uint32_t(seqLenGm_.GetValue(batchIdx)); + uint32_t curSeqLenTileNum = DivCeil(curSeqLen, tileSeqLenSize_); + batchSeqLenTileNum += curSeqLenTileNum; + batchSeqTileN[batchIdx] = batchSeqLenTileNum; + } + } + + __aicore__ inline void ComputeChunkTopKEffectLen() + { + uint32_t lastLen = 0, effectLen = 0, maxChunkSize = MIN_CHUNK_SIZE; + for (uint32_t bIdx = 0; bIdx < param_.batch; bIdx++) { + uint32_t seqLen = uint32_t(seqLenGm_.GetValue(bIdx)); + uint32_t chunkSize = isChunkTopK_ ? uint32_t(chunkSizeGm_.GetValue(bIdx)) : MIN_CHUNK_SIZE; + maxChunkSize = chunkSize > maxChunkSize ? chunkSize : maxChunkSize; + minChunkSize_ = chunkSize < minChunkSize_ ? chunkSize : minChunkSize_; + if(chunkSize == 1 || chunkSize == 8 || chunkSize == 16) + { + // When seqLen <= 32, set effectLen = 0 and skip hammingDistTopK. + if (seqLen <= CHUNK_TOPK_MIN_SEQ_LEN) { + lastLen = seqLen; + effectLen = 0; + } else { + lastLen = (chunkSize != 0) ? (seqLen % chunkSize + MAX_CHUNK_SIZE > seqLen ? seqLen : seqLen % chunkSize + MAX_CHUNK_SIZE) : 0; // min + effectLen = seqLen - lastLen > 0 ? seqLen - lastLen : 0; // max + } + effectLenArr_[bIdx] = effectLen; + } + else if (chunkSize == 64) + { + effectLenArr_[bIdx] = ((seqLen + 63)/64)*64; + } + else if (chunkSize == 128) + { + effectLenArr_[bIdx] = ((seqLen + 127)/128)*128; + } + else + { + effectLenArr_[bIdx] = 0; + } + + } + uint32_t maxSeqLenLastLen = param_.maxSeqLen % maxChunkSize + MAX_CHUNK_SIZE > param_.maxSeqLen ? + param_.maxSeqLen : param_.maxSeqLen % maxChunkSize + MAX_CHUNK_SIZE; + maxEffectLen_ = param_.maxSeqLen - maxSeqLenLastLen > 0 ? param_.maxSeqLen - maxSeqLenLastLen : 0; + } + + __aicore__ inline void ComputeMMTopKBatchNGroup() + { + /* + Example of batchN core grouping + + Example 1 + batch : 2 + headN = 8 = 1 * 8 + coreNum = 20 + SeqLen 86400 11296 ; totalSeqTileN 6 + 1 + coreNumInGroup=3 coreId:0 (0-2 16384 batch 0) 8 (3-4 16384 batch 0) 16 (5 4480 batch 0) (6 11296 batch 1) + coreNumInGroup=3 coreId:1 (0-2 16384 batch 0)9 (3-4 16384 batch 0) 17 (5 4480 batch 0) (6 11296 batch 1) + coreNumInGroup=3 coreId:2 (0-2 16384 batch 0)10 (3-4 16384 batch 0) 18 (5 4480 batch 0) (6 11296 batch 1) + coreNumInGroup=3 coreId:3 (0-2 16384 batch 0)11 (3-4 16384 batch 0) 19 (5 4480 batch 0) (6 11296 batch 1) + coreNumInGroup=2 coreId:4 (0-3 16384 batch 0) 12 (4-5 16384 4480 batch 0) (6 11296 batch 1) + coreNumInGroup=2 coreId:5 (0-3 16384 batch 0) 13 (4-5 16384 4480 batch 0) (6 11296 batch 1) + coreNumInGroup=2 coreId:6 (0-3 16384 batch 0) 14 (4-5 16384 4480 batch 0) (6 11296 batch 1) + coreNumInGroup=2 coreId:7 (0-3 16384 batch 0) 15 (4-5 16384 4480 batch 0) (6 11296 batch 1) + + Example 2 + batchN = 240 = 30 * 8 + coreNum = 20 + batchIdx = 0 + SeqLen 16384 + coreNumInGroup=3 coreId:0 (0-9 16384 batch 0 -9) 8 (10-19 16384 batch 10-19)16 (20-29 16384 batch 10-19) + coreNumInGroup=3 coreId:1 (0-9 16384 batch 0-9)9 (10-19 16384 batch 10-19) 17 (20-29 16384 batch 10-19) + coreNumInGroup=3 coreId:2 (0-9 16384 batch 0-9)10 (10-19 16384 batch 10-19) 18 (20-29 16384 batch 10-19) + coreNumInGroup=3 coreId:3 (0-9 16384 batch 0-9)11 (10-19 16384 batch 10-19) 19 (20-29 16384 batch 10-19) + coreNumInGroup=2 coreId:4 (0-14 16384 batch 0-14) 12 (15-29 16384 batch 15-29) + coreNumInGroup=2 coreId:5 (0-14 16384 batch 0-14) 13 (15-29 16384 batch 15-29) + coreNumInGroup=2 coreId:6 (0-14 16384 batch 0-14) 14 (15-29 16384 batch 15-29) + coreNumInGroup=2 coreId:7 (0-14 16384 batch 0-14) 15 (15-29 16384 batch 15-29) + */ + uint32_t cCoreNumInGroup = param_.usedCoreNum / param_.head; //param_.head is num_kv_heads + cCoreNumInGroup = (blockIdx_ % param_.head < param_.usedCoreNum % param_.head) ? cCoreNumInGroup + 1: cCoreNumInGroup; + uint32_t cCoreIdxInGroup = blockIdx_ / param_.head; + uint32_t curBatchIdx = 0; + uint32_t curHeadIdx = blockIdx_ % param_.head; + + + uint32_t totalSeqTileN = batchSeqTileN[param_.batch - 1]; + uint32_t minTileNumInCore = totalSeqTileN / cCoreNumInGroup; + uint32_t maxTileNumInCore = totalSeqTileN % cCoreNumInGroup == 0 ? minTileNumInCore : minTileNumInCore + 1; + uint32_t maxTileNumCoreIdxInGroup = totalSeqTileN % cCoreNumInGroup == 0 ? 0 : totalSeqTileN % cCoreNumInGroup - 1; + uint32_t tileNumInCore = cCoreIdxInGroup < totalSeqTileN % cCoreNumInGroup ? maxTileNumInCore : minTileNumInCore; + + uint32_t minTileNumStartTileIdx = (maxTileNumCoreIdxInGroup + 1) * maxTileNumInCore + (cCoreIdxInGroup - maxTileNumCoreIdxInGroup - 1) * minTileNumInCore; + uint32_t startTileIdx = cCoreIdxInGroup <= maxTileNumCoreIdxInGroup ? tileNumInCore * cCoreIdxInGroup : minTileNumStartTileIdx; + uint32_t endTileIdx = startTileIdx + tileNumInCore - 1; + + + // Outer loop of CV pipeline + for (uint32_t tileIdx = startTileIdx; tileIdx <= endTileIdx; tileIdx++) { + for (uint32_t i = 0; i < param_.batch; i++) { + if (tileIdx < batchSeqTileN[i]) { + curBatchIdx = i; + break; + } + } + if (supportMask_) { + bool batchMask = maskGm_.GetValue(curBatchIdx); + if (!batchMask) { + //YF_LOG("curBatchIdx = %d, batchMask = %d\n", curBatchIdx, batchMask); + continue; + } + } + uint32_t curSeqLen = uint32_t(seqLenGm_.GetValue(curBatchIdx)); + uint32_t curKScalar = uint32_t(kGm_.GetValue(curBatchIdx)); + uint32_t curEffectLen = effectLenArr_[curBatchIdx]; + if (curSeqLen == 0 || curKScalar == 0 || curEffectLen == 0) { + continue; + } + uint32_t tileSTailSize = curSeqLen % tileSeqLenSize_; + uint32_t realTileIdx = curBatchIdx == 0 ? tileIdx : tileIdx - batchSeqTileN[curBatchIdx - 1]; + uint32_t curCoreDealSize = ((realTileIdx + 1) * tileSeqLenSize_ <= curSeqLen) ? tileSeqLenSize_ : tileSTailSize; + + if ASCEND_IS_AIC { + ComputeMM(curBatchIdx, curHeadIdx, curCoreDealSize, realTileIdx); + CubeNotifyVector(SYNC_AIC_ONLY_ALL_FLAG, SYNC_AIC_AIV_FLAG2 + innerSplitGMPingPongFlag_); + } + if ASCEND_IS_AIV { + VectorWaitCube(SYNC_AIC_AIV_FLAG2 + innerSplitGMPingPongFlag_); + ComputeTopK(curCoreDealSize, curBatchIdx, curHeadIdx, realTileIdx, false); + } + innerSplitGMPingPongFlag_ &= 1; + innerSplitGMPingPongFlag_ ^= 1; + } + } + + __aicore__ inline void ComputeTopK(uint32_t curCoreDealSize, uint32_t curBatchIdx, uint32_t curHeadIdx, uint32_t tileIdx, bool isApproxiTopK) + { + if ASCEND_IS_AIC { + return; + } + + uint32_t curVectorDealSize; + uint32_t vectorTileSSize = tileSeqLenSize_ / VECTOR_CUBE_RATIO; // 8092/2=4096 + uint32_t curSeqLen = uint32_t(seqLenGm_.GetValue(curBatchIdx)); // 27761 + uint32_t curKScalar = uint32_t(kGm_.GetValue(curBatchIdx)); //125 + + if (subBlockIdx_ == 0) { + curVectorDealSize = curCoreDealSize < vectorTileSSize ? curCoreDealSize : vectorTileSSize; + } else if (subBlockIdx_ == 1 && curCoreDealSize > vectorTileSSize) { + curVectorDealSize = curCoreDealSize - vectorTileSSize; + } else { + return; // Return directly if the tail block from S split cannot be assigned to the second vector. + } + + + uint32_t curEffectLen = effectLenArr_[curBatchIdx]; + uint32_t beforeCurVectorDealSize = tileIdx * tileSeqLenSize_ + subBlockIdx_ * vectorTileSSize; + if (beforeCurVectorDealSize >= curEffectLen) { + return; + } + uint32_t curChunkSize = isChunkTopK_ ? uint32_t(chunkSizeGm_.GetValue(curBatchIdx)) : MIN_CHUNK_SIZE; + + uint32_t lastSubBlockSeqLenPadded = 0; // for chunk_size = 16 + uint32_t lastSubBlockSeqLenTruncated = 0; // for chunk_size = 128 + uint32_t curVectorTopKDealSize = 0; + + if(curChunkSize == 1 || curChunkSize == 8 || curChunkSize == 16) + { + lastSubBlockSeqLenTruncated = (curEffectLen - 1) % vectorTileSSize + 1; + //uint32_t curVectorTopKDealSize = beforeCurVectorDealSize + curVectorDealSize > curEffectLen ? curVectorTopKDealSize / + curVectorTopKDealSize = beforeCurVectorDealSize + curVectorDealSize > curEffectLen ? lastSubBlockSeqLenTruncated / curChunkSize : curVectorDealSize / curChunkSize; + } + else if (curChunkSize == 64 || curChunkSize == 128) + { + lastSubBlockSeqLenPadded = (curEffectLen - 1) % vectorTileSSize + 1; + curVectorTopKDealSize = DivCeil(curVectorDealSize, curChunkSize); + } + else + { + return; + } + + curKScalar = isApproxiTopK ? DivCeil(curVectorTopKDealSize * curKScalar, isChunkTopK_ ? curEffectLen / curChunkSize : curSeqLen) : curKScalar; + uint32_t startSeqIdx = curBatchIdx * param_.head * param_.maxSeqLen + curHeadIdx * param_.maxSeqLen + tileIdx * tileSeqLenSize_ + subBlockIdx_ * vectorTileSSize; + uint32_t startIdxInCurWholeSeq = tileIdx * tileSeqLenSize_ / curChunkSize + subBlockIdx_ * vectorTileSSize / curChunkSize; + uint32_t topKResultInterval = isApproxiTopK ? param_.maxK : DivCeil(maxEffectLen_, param_.tileN2) * param_.maxK; + uint32_t topkResultOffset = curBatchIdx * param_.head * topKResultInterval + curHeadIdx * topKResultInterval + tileIdx * VECTOR_CUBE_RATIO * curKScalar + subBlockIdx_ * curKScalar; + + + + uint32_t chunkPerBlock = param_.tileN1 / curChunkSize; //128/16=8 + uint32_t lastSubBlockChunkNum = 0; + if(curChunkSize == 1 || curChunkSize == 8 || curChunkSize == 16) + { + lastSubBlockChunkNum = lastSubBlockSeqLenTruncated / curChunkSize; + } + else + { + lastSubBlockChunkNum = lastSubBlockSeqLenPadded / curChunkSize; + } + + uint32_t curSeqLenTileNum = DivCeil(curSeqLen, tileSeqLenSize_); + + uint32_t lastTileSeqLen = (curEffectLen - 1) % tileSeqLenSize_ + 1; + //uint32_t skipTailChunkNum = DivCeil(SKIP_TAIL_TOKEN_NUM, curChunkSize); + //uint32_t skipHeadChunkNum = DivCeil(SKIP_HEAD_TOKEN_NUM, curChunkSize); + uint32_t skipTailChunkNum = param_.recent; + uint32_t skipHeadChunkNum = param_.sink; + + bool flagLastTileHasTwoSubBlocks = (lastTileSeqLen > vectorTileSSize); + int keepTailType = KEEP_TAIL_TYPE_ONE_BLOCK; + uint32_t lastSubBlockTailChunkNum = min(lastSubBlockChunkNum, skipTailChunkNum); //# of chunks that are kept in the last sub block (vector) + // YF_LOG("curEffectLen=%d, curSeqLen=%d, lastTileSeqLen=%d, tileSeqLenSize_=%d, lastSubBlockSeqLenPadded=%d, flagLastTileHasTwoSubBlocks=%d, lastSubBlockChunkNum=%d, lastSubBlockTailChunkNum=%d\n", curEffectLen,curSeqLen, lastTileSeqLen, tileSeqLenSize_, lastSubBlockSeqLenPadded, flagLastTileHasTwoSubBlocks, lastSubBlockChunkNum, lastSubBlockTailChunkNum); + uint32_t secondLastSubBlockTailChunkNum = 0; //# of chunks that are kept in the second last sub block (vector) + if(lastTileSeqLen < skipTailChunkNum*curChunkSize) //SKIP_TAIL_BLOCK_NUM assumes block_size=128 + { + keepTailType = KEEP_TAIL_TYPE_TWO_TILES; + secondLastSubBlockTailChunkNum = skipTailChunkNum - lastSubBlockTailChunkNum; + } + else if(flagLastTileHasTwoSubBlocks && (lastTileSeqLen % vectorTileSSize < skipTailChunkNum*curChunkSize) ) + { + keepTailType = KEEP_TAIL_TYPE_TWO_BLOCKS; + secondLastSubBlockTailChunkNum = skipTailChunkNum - lastSubBlockTailChunkNum; + } + else + { + keepTailType = KEEP_TAIL_TYPE_ONE_BLOCK; + secondLastSubBlockTailChunkNum = 0; + } + // YF_LOG("keepTailType=%d, lastSubBlockChunkNum=%d, secondLastSubBlockTailChunkNum=%d\n", keepTailType, lastSubBlockChunkNum, secondLastSubBlockTailChunkNum); + + + /* we do not use Case 1 such that we can easily add head and tail top-k chunks + // Case1 If curVectorTopKDealSize < curKScalar_, skip TopK. + if (curVectorTopKDealSize <= curKScalar) { // TODO + DataMoveWithoutTopK(curEffectLen, curVectorTopKDealSize, curChunkSize, startSeqIdx, startIdxInCurWholeSeq, topkResultOffset, isApproxiTopK); + return; + } + */ + + // Case2 Transfer curVectorTopKDealSize data in one go; TopK exhaustion mode. + LocalTensor topKValueInTensor = topKValueInQueue_.AllocTensor(); + + // YF_LOG("curBatchIdx = %d curHeadIdx = %d curVectorTopKDealSize = %d\n", curBatchIdx, curHeadIdx, curVectorTopKDealSize); + // 2.1 For chunkTopK, first compute the max value within the chunk, then perform TopK. + if (curChunkSize > MIN_CHUNK_SIZE) { + LocalTensor chunkReduceMaxValueInTensor = chunkReduceMaxValueInQueue_.AllocTensor(); + ReduceMaxCustom(matmulGm_[startSeqIdx], chunkReduceMaxValueInTensor, topKValueInTensor, curVectorTopKDealSize, curChunkSize); + chunkReduceMaxValueInQueue_.EnQue(chunkReduceMaxValueInTensor); + chunkReduceMaxValueInTensor = chunkReduceMaxValueInQueue_.DeQue(); + chunkReduceMaxValueInQueue_.FreeTensor(chunkReduceMaxValueInTensor); + } else { + // 2.2 For regular TopK, perform TopK directly. curVectorTopKDealSize is the processed length after core splitting, no conversion exception will be introduced. + DataCopyExtParams copyInParams{1, static_cast(curVectorTopKDealSize * sizeof(half)), 0, 0, 0}; + DataCopyPadExtParams copyInPadParams{false, 0, 0, 0}; + DataCopyPad(topKValueInTensor, matmulGm_[startSeqIdx], copyInParams, copyInPadParams); + } + + + + if( (tileIdx == 0) && (subBlockIdx_ == 0) ) + { + Duplicate(topKValueInTensor, static_cast(MAX_HALF_VALUE), min(skipHeadChunkNum, curVectorTopKDealSize)); + } + + switch(keepTailType) + { + case KEEP_TAIL_TYPE_ONE_BLOCK: + if(tileIdx == curSeqLenTileNum-1) + { + + // YF_LOG("skipTailChunkNum=%d, skipHeadChunkNum=%d, keepTailType=%d, lastSubBlockTailChunkNum=%d, secondLastSubBlockTailChunkNum=%d,curVectorTopKDealSize=%d, flagLastTileHasTwoSubBlocks=%d\n", skipTailChunkNum, skipHeadChunkNum, keepTailType, lastSubBlockTailChunkNum, secondLastSubBlockTailChunkNum, curVectorTopKDealSize, flagLastTileHasTwoSubBlocks); + + if(flagLastTileHasTwoSubBlocks) + { + if(subBlockIdx_ == 1) + { + //copy last + // DumpTensor(topKValueInTensor, 502, topKValueInTensor.GetSize()); + FillMaxValueFromTail(topKValueInTensor, curVectorTopKDealSize, lastSubBlockTailChunkNum,curChunkSize); + // DumpTensor(topKValueInTensor, 504, topKValueInTensor.GetSize()); + } + } + else + { + if(subBlockIdx_ == 0) + { + //copy last + // DumpTensor(topKValueInTensor, 513, topKValueInTensor.GetSize()); + FillMaxValueFromTail(topKValueInTensor, curVectorTopKDealSize, lastSubBlockTailChunkNum,curChunkSize); + // DumpTensor(topKValueInTensor, 515, topKValueInTensor.GetSize()); + } + } + } + break; + case KEEP_TAIL_TYPE_TWO_BLOCKS: + if(tileIdx == curSeqLenTileNum-1) + { + if(subBlockIdx_ == 1) + { + //copy last + FillMaxValueFromTail(topKValueInTensor, curVectorTopKDealSize, lastSubBlockTailChunkNum,curChunkSize); + } + else + { + //copy second last + FillMaxValueFromTail(topKValueInTensor, curVectorTopKDealSize, secondLastSubBlockTailChunkNum,curChunkSize); + } + } + break; + case KEEP_TAIL_TYPE_TWO_TILES: + if(tileIdx == curSeqLenTileNum-1) + { + if(subBlockIdx_ == 0) + { + //copy last + FillMaxValueFromTail(topKValueInTensor, curVectorTopKDealSize, lastSubBlockTailChunkNum,curChunkSize); + } + + } + else if(tileIdx == curSeqLenTileNum-2) + { + if(subBlockIdx_ == 1) + { + //copy second last + FillMaxValueFromTail(topKValueInTensor, curVectorTopKDealSize, secondLastSubBlockTailChunkNum,curChunkSize); + } + } + break; + } + + // DumpTensor(topKValueInTensor, 666, topKValueInTensor.GetSize()); + topKValueInQueue_.EnQue(topKValueInTensor); + topKValueInTensor = topKValueInQueue_.DeQue(); + // 2.3 chunkTopK, regular TopK exhaustion mode + TopKInExhaustionMode(curVectorTopKDealSize, curKScalar, topKValueInTensor, topkResultOffset, startIdxInCurWholeSeq, isApproxiTopK); + topKValueInQueue_.FreeTensor(topKValueInTensor); + } + + __aicore__ inline void DataMoveWithoutTopK(uint32_t curEffectLen, uint32_t curVectorTopKDealSize, uint32_t curChunkSize, uint32_t startSeqIdx, uint32_t startIdxInCurWholeSeq, uint32_t topkResultOffset, bool isApproxiTopK) + { + if ASCEND_IS_AIC { + return; + } + LocalTensor topKValueTensor = topKValueInnerOutQueue_.AllocTensor(); + LocalTensor topKIdexTensor = topKIndexInnerOutQueue_.AllocTensor(); + // For chunkTopK: perform chunk max first, generate index then write out. For regular TopK: load directly into UB, generate index then write out. + if (curChunkSize > MIN_CHUNK_SIZE) { + LocalTensor chunkReduceMaxValueInTensor = chunkReduceMaxValueInQueue_.AllocTensor(); + ReduceMaxCustom(matmulGm_[startSeqIdx], chunkReduceMaxValueInTensor, topKValueTensor, curVectorTopKDealSize, curChunkSize); + chunkReduceMaxValueInQueue_.EnQue(chunkReduceMaxValueInTensor); + chunkReduceMaxValueInTensor = chunkReduceMaxValueInQueue_.DeQue(); + chunkReduceMaxValueInQueue_.FreeTensor(chunkReduceMaxValueInTensor); + } else { + // curVectorTopKDealSize: length after core split, no conversion exception + DataCopyExtParams copyInParams{1, static_cast(curVectorTopKDealSize * sizeof(half)), 0, 0, 0}; + DataCopyPadExtParams copyInPadParams{false, 0, 0, 0}; + DataCopyPad(topKValueTensor, matmulGm_[startSeqIdx], copyInParams, copyInPadParams); + } + ArithProgression(topKIdexTensor, static_cast(startIdxInCurWholeSeq), 1, static_cast(curVectorTopKDealSize)); + topKValueInnerOutQueue_.EnQue(topKValueTensor); + topKIndexInnerOutQueue_.EnQue(topKIdexTensor); + topKValueTensor = topKValueInnerOutQueue_.DeQue(); + topKIdexTensor = topKIndexInnerOutQueue_.DeQue(); + DataCopyFromUBToGM(isApproxiTopK, curVectorTopKDealSize, topkResultOffset, topKValueTensor, topKIdexTensor); + topKValueInnerOutQueue_.FreeTensor(topKValueTensor); + topKIndexInnerOutQueue_.FreeTensor(topKIdexTensor); + } + + __aicore__ inline void DataCopyFromUBToGM(bool isApproxiTopK, uint32_t copyLen, uint32_t topkResultOffset, const LocalTensor &topKValueTensor, const LocalTensor &topKIdexTensor) + { + if ASCEND_IS_AIC { + return; + } + if (isApproxiTopK) { + DataCopyExtParams copyTopKIdxToOutParams{1, static_cast(copyLen * sizeof(uint32_t)), 0, 0, 0}; + DataCopyPad(indicesGm_[topkResultOffset], topKIdexTensor, copyTopKIdxToOutParams); + } else { + DataCopyExtParams copyTopKValueOutParams{1, static_cast(copyLen * sizeof(half)), 0, 0, 0}; + DataCopyPad(topkValueGm_[topkResultOffset], topKValueTensor, copyTopKValueOutParams); + DataCopyExtParams copyTopKIdxToWorkSpaceParams{1, static_cast(copyLen * sizeof(uint32_t)), 0, 0, 0}; + DataCopyPad(topkIdxGm_[topkResultOffset], topKIdexTensor, copyTopKIdxToWorkSpaceParams); + } + } + + __aicore__ inline void TopKInExhaustionMode(uint32_t curVectorTopKDealSize, uint32_t curKScalar, const LocalTensor &topKValueInTensor, uint32_t topkResultOffset, uint32_t startIdxInCurWholeSeq, bool isApproxiTopK) + { + /* Recalculate InnerSize, ensure both LocalTensor starts are 32-byte aligned for Copy */ + uint32_t effectiveInnerSize = curKScalar + (param_.topKInnerSize - curKScalar) / BLOCK_CUBE * BLOCK_CUBE; + uint32_t curDealSeqLen = 0, curExtendSeqLen = 0, topKLoop = 0, ubStartIdxInTileSeq = 0; + topKLoop = curVectorTopKDealSize < effectiveInnerSize ? 1 : DivCeil(curVectorTopKDealSize - effectiveInnerSize, effectiveInnerSize - curKScalar) + 1; + LocalTensor topKValueInnerOutTensor; + LocalTensor topKIdexInnerOutTensor; + for(uint32_t loopIdx = 0; loopIdx < topKLoop; loopIdx++) { + LocalTensor topKValueInnerInTensor = topKValueInnerInQueue_.AllocTensor(); + LocalTensor topKIdexInnerInTensor = topKIndexInnerInQueue_.AllocTensor(); + if (loopIdx == 0) { + curDealSeqLen = curVectorTopKDealSize < effectiveInnerSize ? curVectorTopKDealSize: effectiveInnerSize; + FillMinValue(curDealSeqLen, topKValueInnerInTensor); + DataCopyInUB(curDealSeqLen, topKValueInnerInTensor, topKValueInTensor[ubStartIdxInTileSeq], MAX_FP16_PROCESS_NUM); + ArithProgression(topKIdexInnerInTensor, static_cast(startIdxInCurWholeSeq), 1, static_cast(curDealSeqLen)); + ubStartIdxInTileSeq += effectiveInnerSize; + startIdxInCurWholeSeq += effectiveInnerSize; + } else { + topKValueInnerOutTensor = topKValueInnerOutQueue_.DeQue(); + topKIdexInnerOutTensor = topKIndexInnerOutQueue_.DeQue(); + uint32_t maxExtendSeqLen = effectiveInnerSize - curKScalar; + curExtendSeqLen = (curVectorTopKDealSize - ubStartIdxInTileSeq) >= maxExtendSeqLen ? maxExtendSeqLen : curVectorTopKDealSize - ubStartIdxInTileSeq; + curDealSeqLen = curKScalar + curExtendSeqLen; + FillMinValue(curDealSeqLen, topKValueInnerInTensor); + DataCopyInUB(curDealSeqLen, topKValueInnerInTensor, topKValueInTensor[ubStartIdxInTileSeq - curKScalar], MAX_FP16_PROCESS_NUM); + ArithProgression(topKIdexInnerInTensor, static_cast(startIdxInCurWholeSeq - curKScalar), 1, static_cast(curExtendSeqLen + curKScalar)); + DataCopyInUB(curKScalar, topKValueInnerInTensor, topKValueInnerOutTensor, MAX_FP16_PROCESS_NUM); + DataCopyInUB(curKScalar, topKIdexInnerInTensor, topKIdexInnerOutTensor, MAX_INT32_PROCESS_NUM); + ubStartIdxInTileSeq += maxExtendSeqLen; + startIdxInCurWholeSeq += maxExtendSeqLen; + topKValueInnerOutQueue_.FreeTensor(topKValueInnerOutTensor); + topKIndexInnerOutQueue_.FreeTensor(topKIdexInnerOutTensor); + } + topKValueInnerInQueue_.EnQue(topKValueInnerInTensor); + topKIndexInnerInQueue_.EnQue(topKIdexInnerInTensor); + topKValueInnerInTensor = topKValueInnerInQueue_.DeQue(); + topKIdexInnerInTensor = topKIndexInnerInQueue_.DeQue(); + topKValueInnerOutTensor = topKValueInnerOutQueue_.AllocTensor(); + topKIdexInnerOutTensor = topKIndexInnerOutQueue_.AllocTensor(); + TopKCustom(topKValueInnerOutTensor, topKIdexInnerOutTensor, topKValueInnerInTensor, topKIdexInnerInTensor, static_cast(curKScalar), tilingData_, curDealSeqLen); + topKValueInnerInQueue_.FreeTensor(topKValueInnerInTensor); + topKIndexInnerInQueue_.FreeTensor(topKIdexInnerInTensor); + topKValueInnerOutQueue_.EnQue(topKValueInnerOutTensor); + topKIndexInnerOutQueue_.EnQue(topKIdexInnerOutTensor); + } + topKValueInnerOutTensor = topKValueInnerOutQueue_.DeQue(); + topKIdexInnerOutTensor = topKIndexInnerOutQueue_.DeQue(); + DataCopyFromUBToGM(isApproxiTopK, curKScalar, topkResultOffset, topKValueInnerOutTensor, topKIdexInnerOutTensor); + topKValueInnerOutQueue_.FreeTensor(topKValueInnerOutTensor); + topKIndexInnerOutQueue_.FreeTensor(topKIdexInnerOutTensor); + } + + __aicore__ inline void FillMinValue(uint32_t copyLen, LocalTensor& topKValueInTensor) + { + if ASCEND_IS_AIC { + return; + } + uint32_t copyLenFloorAligned = copyLen / BLOCK_CUBE * BLOCK_CUBE; + if (copyLenFloorAligned < param_.topKInnerSize) { + Duplicate(topKValueInTensor[copyLenFloorAligned], static_cast(MIN_HALF_VALUE), param_.topKInnerSize - copyLenFloorAligned); + } + } + + template + __aicore__ inline void DataCopyInUB(uint32_t curDealSeqLen, const LocalTensor &topKToTensor, const LocalTensor &topKFromTensor, uint32_t maxProcessNum) + { + if ASCEND_IS_AIC { + return; + } + uint64_t mask; + uint8_t repeatTimes; + mask = curDealSeqLen > maxProcessNum ? maxProcessNum : curDealSeqLen; + repeatTimes = curDealSeqLen / maxProcessNum; + if (repeatTimes > 0) { + Copy(topKToTensor, topKFromTensor, mask, repeatTimes, {1, 1, 8, 8}); + } + if (curDealSeqLen % maxProcessNum != 0) { + mask = curDealSeqLen % maxProcessNum; + Copy(topKToTensor[repeatTimes * maxProcessNum], topKFromTensor[repeatTimes * maxProcessNum], mask, 1, {1, 1, 8, 8}); + } + } + + __aicore__ inline void MergeTopKInGroup() + { + if ASCEND_IS_AIC { + return; + } + LocalTensor topKValueOutTensor; + LocalTensor topKIndexOutTensor; + uint32_t vectorTotalNum = param_.usedCoreNum * VECTOR_CUBE_RATIO; + uint32_t iterTime = DivCeil(param_.batchN, vectorTotalNum), iterTail = param_.batchN % vectorTotalNum; + for(uint32_t iterIdx = 0; iterIdx < iterTime; iterIdx++) { + uint32_t curIterDealBatchN = (iterIdx + 1) * vectorTotalNum <= param_.batchN ? vectorTotalNum : iterTail; + if (blockIdx_ * VECTOR_CUBE_RATIO + subBlockIdx_ >= curIterDealBatchN) { + return; + } + uint32_t curVectorDealBatchNIdx = iterIdx * vectorTotalNum + blockIdx_ * VECTOR_CUBE_RATIO + subBlockIdx_; + uint32_t curBatchIdx = curVectorDealBatchNIdx / param_.head; + + if (supportMask_) { + bool batchMask = maskGm_.GetValue(curBatchIdx); + //YF_LOG("curBatchIdx = %d, batchMask = %d\n", curBatchIdx, batchMask); + if (!batchMask) { + SetBlockTableForIndices(curBatchIdx, curVectorDealBatchNIdx * param_.maxK); + //YF_LOG("curBatchIdx = %d SetBlockTableForIndices\n", curBatchIdx); + continue; + } + } + + uint32_t curKScalar = uint32_t(kGm_.GetValue(curBatchIdx)), curEffectLen = effectLenArr_[curBatchIdx], curSeqLen = uint32_t(seqLenGm_.GetValue(curBatchIdx)); + uint32_t curChunkSize = isChunkTopK_ ? uint32_t(chunkSizeGm_.GetValue(curBatchIdx)) : MIN_CHUNK_SIZE; + if (curSeqLen == 0 || curKScalar == 0 || curEffectLen == 0) { + continue; + } + uint32_t preTopkLenLast = ((curEffectLen / curChunkSize) % (param_.tileN2 / curChunkSize)) < curKScalar ? + (curEffectLen / curChunkSize) % (param_.tileN2 / curChunkSize) : curKScalar; + uint32_t preTopKLen = curEffectLen / param_.tileN2 * curKScalar + preTopkLenLast; + uint64_t mergeTopKOutGmOffset = curVectorDealBatchNIdx * param_.maxK; + uint64_t topKResultGmOffset = curVectorDealBatchNIdx * DivCeil(maxEffectLen_, param_.tileN2) * param_.maxK; + if (preTopKLen <= curKScalar) { + LocalTensor topKInIndexTensor = topKIndexInnerOutQueue_.AllocTensor(); + DataCopyExtParams copyIndexParams{1, static_cast(curKScalar * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams copyIndexPadParams{false, 0, 0, 0}; + DataCopyPad(topKInIndexTensor, topkIdxGm_[topKResultGmOffset], copyIndexParams, copyIndexPadParams); + PipeBarrier(); + if (!param_.supportOffload && (curChunkSize == 64 || curChunkSize == 128)) { + ReMappingBlockTableIndices(topKInIndexTensor, curBatchIdx, curKScalar, mergeTopKOutGmOffset); + } else { + topKIndexInnerOutQueue_.EnQue(topKInIndexTensor); + topKInIndexTensor = topKIndexInnerOutQueue_.DeQue(); + DataCopyExtParams copyTopKIdxParams{1, static_cast(curKScalar * sizeof(int32_t)), 0, 0, 0}; + DataCopyPad(indicesGm_[mergeTopKOutGmOffset], topKInIndexTensor, copyTopKIdxParams); + } + topKIndexInnerOutQueue_.FreeTensor(topKInIndexTensor); + continue; + } + uint32_t blockTile = preTopKLen > param_.topKInnerSize ? param_.topKInnerSize : preTopKLen; + uint32_t topKBlockNum = DivCeil(preTopKLen - blockTile, blockTile - curKScalar) + 1; + uint32_t blockTail = topKBlockNum > 1 ? (preTopKLen - ((topKBlockNum - 1) * (blockTile - curKScalar))): preTopKLen; + MergeTopKInExhaustionMode(topKBlockNum, topKResultGmOffset, blockTail, blockTile, curKScalar, topKValueOutTensor, topKIndexOutTensor); + topKValueOutTensor = topKValueInnerOutQueue_.DeQue(); + topKIndexOutTensor = topKIndexInnerOutQueue_.DeQue(); + if (!param_.supportOffload && (curChunkSize == 64 || curChunkSize == 128)) { + ReMappingBlockTableIndices(topKIndexOutTensor, curBatchIdx, curKScalar, mergeTopKOutGmOffset); + } else { + DataCopyExtParams copyOutParams{1, static_cast(curKScalar * sizeof(int32_t)), 0, 0, 0}; // {blockNum, blockLen, srcStride, dstStride, rsv(no concern)} + DataCopyPad(indicesGm_[mergeTopKOutGmOffset], topKIndexOutTensor, copyOutParams); + } + + topKValueInnerOutQueue_.FreeTensor(topKValueOutTensor); + topKIndexInnerOutQueue_.FreeTensor(topKIndexOutTensor); + } + } + + __aicore__ inline void ReMappingBlockTableIndices(LocalTensor& topKIndexOutTensor, uint32_t curBatchIdx, uint32_t curKScalar, uint64_t outGmOffset) { + if ASCEND_IS_AIC { return; } + + if (curKScalar < 32) { + useInnerSort = false; + } + + if (useInnerSort) { + CustomSort(topKIndexOutTensor, curKScalar); + SelectBlockTableFromTopK(curBatchIdx, curKScalar, outGmOffset); + return; + } + + WriteBlockTableFromTopK(curBatchIdx, topKIndexOutTensor, curKScalar, outGmOffset); + } + + __aicore__ inline void CustomSort(LocalTensor& topKIndexOutTensor, uint32_t len) { + if ASCEND_IS_AIC { return; } + if (len <= 1) { return; } + + LocalTensor valueLocal = topKIndexSortInQueue_.AllocTensor(); + Cast(valueLocal, topKIndexOutTensor, RoundMode::CAST_CEIL, static_cast(len));//, static_cast((len + 63) / 64), {1, 1, 8, 8}); + topKIndexSortInQueue_.EnQue(valueLocal); + LocalTensor indexLocal = topKIndexInnerInQueue_.AllocTensor(); + topKIndexInnerInQueue_.EnQue(indexLocal); + + uint32_t repeatTimes = (len + 31) / 32; + valueLocal = topKIndexSortInQueue_.DeQue(); + indexLocal = topKIndexInnerInQueue_.DeQue(); + LocalTensor sortedLocal = topKIndexSortCalcQueue_.AllocTensor(); + LocalTensor concatTmpLocal = topKIndexSortTmpQueue_.AllocTensor(); + LocalTensor concatLocal; + + Concat(concatLocal, valueLocal, concatTmpLocal, repeatTimes); + Sort(sortedLocal, concatLocal, indexLocal, concatTmpLocal, repeatTimes); + Extract(valueLocal, indexLocal, sortedLocal, repeatTimes); + + topKIndexSortTmpQueue_.FreeTensor(concatTmpLocal); + topKIndexInnerInQueue_.FreeTensor(indexLocal); + topKIndexSortCalcQueue_.FreeTensor(sortedLocal); + + topKIndexSortInQueue_.EnQue(valueLocal); + } + + __aicore__ inline void SelectBlockTableFromTopK(uint32_t curBatchIdx, uint32_t curKScalar, uint64_t outGmOffset) + { + if ASCEND_IS_AIC { + return; + } + + LocalTensor sortedTopKIndexTensor = topKIndexSortInQueue_.DeQue(); + LocalTensor blockIdUb = topKIndexInnerInQueue_.AllocTensor(); + + __ubuf__ const float *in_ptr = reinterpret_cast<__ubuf__ const float *>(sortedTopKIndexTensor.GetPhyAddr()); + __ubuf__ int32_t *out_ptr = reinterpret_cast<__ubuf__ int32_t *>(blockIdUb.GetPhyAddr()); + + LocalTensor tableBlockTensor = tableBlockBuf_.template Get(); + DataCopyExtParams copyInParams{1, static_cast(param_.blockCount * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams copyInPadParams{false, 0, 0, 0}; + DataCopyPad(tableBlockTensor, keyBlockTableGm_[curBatchIdx * param_.blockCount], copyInParams, copyInPadParams); + + for (uint32_t i = 0; i < curKScalar; ++i) { + const int32_t idx = static_cast(in_ptr[i]); + out_ptr[curKScalar - 1 - i] = isContinuousBatch_ + ? tableBlockTensor.GetValue(static_cast(idx)) + : (idx + 1); + } + + // DumpTensor(blockIdUb, 859, 32); + + DataCopyExtParams cpOut{1, static_cast(curKScalar * sizeof(int32_t)), 0, 0, 0}; + DataCopyPad(indicesGm_[outGmOffset], blockIdUb, cpOut); + + topKIndexInnerInQueue_.FreeTensor(blockIdUb); + topKIndexSortInQueue_.FreeTensor(sortedTopKIndexTensor); + } + + __aicore__ inline void SetBlockTableForIndices(uint32_t curBatchIdx, uint64_t outGmOffset) { + LocalTensor tableBlockTensor = tableBlockBuf_.template Get(); + DataCopyExtParams copyInParams{1, static_cast(param_.blockCount * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams copyInPadParams{false, 0, 0, 0}; + DataCopyPad(tableBlockTensor, keyBlockTableGm_[curBatchIdx * param_.blockCount], copyInParams, copyInPadParams); + //DumpTensor(tableBlockTensor, 921, 64); + SetFlag(0); + WaitFlag(0); + + uint32_t copyLen = param_.blockCount < param_.maxK ? param_.blockCount : param_.maxK; + DataCopyExtParams cpOut{1, static_cast(copyLen * sizeof(int32_t)), 0, 0, 0}; + DataCopyPad(indicesGm_[outGmOffset], tableBlockTensor, cpOut); + } + + __aicore__ inline void MergeTopKInExhaustionMode(uint32_t topKBlockNum, uint64_t topKResultGmOffset, uint32_t blockTail, uint32_t blockTile, uint32_t curKScalar, LocalTensor& topKValueOutTensor, LocalTensor& topKIndexOutTensor) + { + if ASCEND_IS_AIC { + return; + } + for (uint32_t loop = 0; loop < topKBlockNum; loop++) { + LocalTensor topKInValueTensor = topKValueInnerInQueue_.AllocTensor(); + LocalTensor topKInIndexTensor = topKIndexInnerInQueue_.AllocTensor(); + uint32_t copyLen = loop == topKBlockNum - 1 ? blockTail : blockTile; + uint64_t curLoopTopKResultGmOffset = loop == 0 ? topKResultGmOffset : topKResultGmOffset + loop * (blockTile - curKScalar); + GenerateTopKValueTensor(copyLen, blockTile, curLoopTopKResultGmOffset, curKScalar, topKInValueTensor); + GenerateTopKIndexTensor(copyLen, blockTile, curLoopTopKResultGmOffset, curKScalar, topKInIndexTensor); + if (loop > 0) { + topKValueOutTensor = topKValueInnerOutQueue_.DeQue(); + topKIndexOutTensor = topKIndexInnerOutQueue_.DeQue(); + topKInValueTensor = topKValueInnerInQueue_.DeQue(); + topKInIndexTensor = topKIndexInnerInQueue_.DeQue(); + uint64_t valueMask = curKScalar > MAX_FP16_PROCESS_NUM ? MAX_FP16_PROCESS_NUM : curKScalar; + uint8_t valueRepeatTimes = curKScalar / MAX_FP16_PROCESS_NUM; + uint64_t indexMask = curKScalar > MAX_INT32_PROCESS_NUM ? MAX_INT32_PROCESS_NUM : curKScalar; + uint8_t indexRepeatTimes = curKScalar / MAX_INT32_PROCESS_NUM; + if (valueRepeatTimes > 0) { + Copy(topKInValueTensor, topKValueOutTensor, valueMask, valueRepeatTimes, {1, 1, 8, 8}); + } + if (indexRepeatTimes > 0) { + Copy(topKInIndexTensor, topKIndexOutTensor, indexMask, indexRepeatTimes, {1, 1, 8, 8}); + } + if (curKScalar % MAX_FP16_PROCESS_NUM != 0) { + Copy(topKInValueTensor[valueRepeatTimes * MAX_FP16_PROCESS_NUM], topKValueOutTensor[valueRepeatTimes * MAX_FP16_PROCESS_NUM], curKScalar % MAX_FP16_PROCESS_NUM, 1, {1, 1, 8, 8}); + } + if (curKScalar % MAX_INT32_PROCESS_NUM != 0) { + Copy(topKInIndexTensor[indexRepeatTimes * MAX_INT32_PROCESS_NUM], topKIndexOutTensor[indexRepeatTimes * MAX_INT32_PROCESS_NUM], curKScalar % MAX_INT32_PROCESS_NUM, 1, {1, 1, 8, 8}); + } + PipeBarrier(); + topKValueInnerInQueue_.EnQue(topKInValueTensor); + topKIndexInnerInQueue_.EnQue(topKInIndexTensor); + topKValueInnerOutQueue_.FreeTensor(topKValueOutTensor); + topKIndexInnerOutQueue_.FreeTensor(topKIndexOutTensor); + } + topKInValueTensor = topKValueInnerInQueue_.DeQue(); + topKInIndexTensor = topKIndexInnerInQueue_.DeQue(); + topKValueOutTensor = topKValueInnerOutQueue_.AllocTensor(); + topKIndexOutTensor = topKIndexInnerOutQueue_.AllocTensor(); + TopKCustom(topKValueOutTensor, topKIndexOutTensor, topKInValueTensor, topKInIndexTensor, curKScalar, tilingData_, copyLen); + topKValueInnerInQueue_.FreeTensor(topKInValueTensor); + topKIndexInnerInQueue_.FreeTensor(topKInIndexTensor); + topKValueInnerOutQueue_.EnQue(topKValueOutTensor); + topKIndexInnerOutQueue_.EnQue(topKIndexOutTensor); + } + } + + __aicore__ inline void GenerateTopKValueTensor(uint32_t copyLen, uint32_t blockTile, uint64_t topkGmOffset, uint32_t curK, LocalTensor& topKInValueTensor) + { + if ASCEND_IS_AIC { + return; + } + uint32_t copyLenAligned = copyLen / BLOCK_CUBE * BLOCK_CUBE; /* floor aligned for datacopy */ + if (copyLenAligned < param_.topKInnerSize) { + Duplicate(topKInValueTensor[copyLenAligned], static_cast(MIN_HALF_VALUE), param_.topKInnerSize - copyLenAligned); + SetFlag(0); + WaitFlag(0); + } + uint32_t copyLenCeilAligned = DivCeil(copyLen * sizeof(half), BLOCK_CUBE) * BLOCK_CUBE / sizeof(half); + DataCopyExtParams copyInValueParams{1, static_cast(copyLen * sizeof(half)), 0, 0, 0}; + DataCopyPadExtParams copyInValuePadParams{true, 0, static_cast(copyLenCeilAligned - copyLen), static_cast(MIN_HALF_VALUE)}; + DataCopyPad(topKInValueTensor, topkValueGm_[topkGmOffset], copyInValueParams, copyInValuePadParams); + topKValueInnerInQueue_.EnQue(topKInValueTensor); + } + + __aicore__ inline void GenerateTopKIndexTensor(uint32_t copyLen, uint32_t blockTile, uint64_t topkGmOffset, uint32_t curK, LocalTensor& topKInIndexTensor) + { + if ASCEND_IS_AIC { + return; + } + DataCopyExtParams copyIndexParams{1, static_cast(copyLen * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams copyIndexPadParams{false, 0, 0, 0}; + DataCopyPad(topKInIndexTensor, topkIdxGm_[topkGmOffset], copyIndexParams, copyIndexPadParams); + topKIndexInnerInQueue_.EnQue(topKInIndexTensor); + } + + template + __aicore__ inline void SyncAivOnly(uint16_t eventId) { + CrossCoreSetFlag(eventId); + CrossCoreWaitFlag(eventId); + } + + template + __aicore__ inline void VectorNotifyCube(uint16_t aivOnlyEventId, uint16_t aiv2AicEventId) { + SyncAivOnly(aivOnlyEventId); + CrossCoreSetFlag(aiv2AicEventId); + } + + __aicore__ inline void CubeWaitVector(uint16_t aiv2AicEventId) { + CrossCoreWaitFlag(aiv2AicEventId); + } + + template + __aicore__ inline void CubeNotifyVector(uint16_t aicOnlyEventId, uint16_t aic2AivEventId) { + CrossCoreSetFlag(aic2AivEventId); + } + + __aicore__ inline void VectorWaitCube(uint16_t aic2AivEventId) { + CrossCoreWaitFlag(aic2AivEventId); + } + +protected: + uint64_t innerSplitLoopTimes_ = 0; + bool innerSplitGMPingPongFlag_ = 0; + uint32_t topKBlockNum = 1; + uint32_t batchSeqTileN[MAX_BATCH_SIZE] = {0}; + uint32_t effectLenArr_[MAX_BATCH_SIZE] = {0}; + bool isChunkTopK_ = true; + uint32_t maxEffectLen_ = 0; + uint32_t minChunkSize_ = MAX_CHUNK_SIZE; + uint32_t tileSeqLenSize_ = 0; + bool isContinuousBatch_ = true; + bool supportMask_ = true; + + static constexpr uint32_t BLOCK_CUBE = 32; + static constexpr uint64_t SYNC_MODE0 = 0; + static constexpr uint64_t SYNC_MODE2 = 2; + static constexpr uint64_t SYNC_AIV_ONLY_ALL_FLAG = 0; + static constexpr uint64_t SYNC_AIC_ONLY_ALL_FLAG = 1; + static constexpr uint64_t SYNC_AIV_AIC_FLAG = 2; + static constexpr uint64_t SYNC_AIC_AIV_FLAG = 4; + static constexpr uint64_t SYNC_AIV_AIC_FLAG2 = 6; + static constexpr uint64_t SYNC_AIC_AIV_FLAG2 = 8; + int32_t eventIDFIX_MTE2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::FIX_MTE2)); + + GlobalTensor queryGm_; + GlobalTensor keyCompressedGm_; + GlobalTensor keyCompressedRopeGm_; + GlobalTensor kGm_; + GlobalTensor seqLenGm_; + GlobalTensor chunkSizeGm_; + GlobalTensor keyBlockTableGm_; + GlobalTensor maskGm_; + GlobalTensor indicesGm_; + GlobalTensor unpackGm_; + GlobalTensor kRopeUnpackGm_; + GlobalTensor qUnpackGm_; + GlobalTensor matmulGm_; + GlobalTensor topkValueGm_; + GlobalTensor topkIdxGm_; + + TPipe *pipe_; + TQue keyCompressedInBuf_; + TQue keyUnpackedOutBuf_; + TQue chunkReduceMaxValueInQueue_; + TQue topKValueInQueue_; + TQue topKValueInnerInQueue_; + TQue topKIndexInnerInQueue_; + TQue topKValueInnerOutQueue_; + TQue topKIndexInnerOutQueue_; + TBuf constBuf_; + TBuf selectBuf_; + TQue queryCompressedInQueue_; + TQue queryUnpackedOutQueue_; + TBuf qReduceSumLastRowBuf_; + TBuf qReduceSumBuf_; + + TQue topKIndexSortInQueue_; + TQue topKIndexSortTmpQueue_; + TQue topKIndexSortCalcQueue_; + TBuf tableBlockBuf_; + + int32_t layerIdScalar_; + TilingParam param_; + TopkTiling topkTiling_; + HammingDistTopKTilingData tilingData_; + + using AMatmulType = matmul::MatmulType; + using BMatmulType = matmul::MatmulType; + using BiasMatmulType = matmul::MatmulType; + // notice: the TPos of ctype must be ub given by mm api when iterate, + // but actually we can move data to gm then to ub. + using CMatmulType = matmul::MatmulType; + matmul::MatmulImpl mm_; + matmul::MatmulImpl mmRope_; + + uint32_t blockIdx_; + uint32_t subBlockIdx_; + uint64_t mmOffsetA_; + uint64_t mmOffsetB_; + uint64_t mmOffsetARope_; + uint64_t mmOffsetBRope_; + uint64_t mmOffsetC_; + + uint64_t index_; + bool useInnerSort = true; + + __aicore__ inline void InitParams() + { + blockIdx_ = GetBlockIdx(); + subBlockIdx_ = GetSubBlockIdx(); + mmOffsetA_ = RESET_NUM; + mmOffsetB_ = RESET_NUM; + mmOffsetARope_ = RESET_NUM; + mmOffsetBRope_ = RESET_NUM; + mmOffsetC_ = RESET_NUM; + index_ = RESET_NUM; + } + + __aicore__ inline void InitTilingParams(const TCubeTiling &tiling, const TCubeTiling &tilingRope, const TopkTiling &topkTiling, const HammingDistTopKTilingParams &tilingParam) + { + // tiling data for select + param_.batch = tilingParam.batch; + param_.maxK = tilingParam.maxK; + param_.batchN = tilingParam.batchN; + param_.tileN1 = tilingParam.tileN1; + param_.tileN2 = tilingParam.tileN2; + param_.dimension = tilingParam.dimension; + param_.nope_dimension = tilingParam.nope_dimension; + param_.rope_dimension = tilingParam.rope_dimension; + param_.head = tilingParam.head; + param_.layerSize = tilingParam.layerSize; + param_.layerSizeRope = tilingParam.layerSizeRope; + param_.qHead = tilingParam.qHead; + param_.headGroupNum = tilingParam.headGroupNum; + + // support key rope + param_.supportKeyRope = tilingParam.supportKeyRope > 0; + + // tiling data for matmul + param_.usedCoreNum = tilingParam.usedCoreNum; + param_.M = tiling.M; + param_.N = tiling.N; + param_.ka = tiling.Ka; // dimension + param_.kb = tiling.Kb; // dimension + if (param_.supportKeyRope) { + param_.rope_ka = tilingRope.Ka; // dimension + param_.rope_kb = tilingRope.Kb; // dimension + } + // YF_LOG("param_.M=%d, param_.N=%d, param_.ka=%d, param_.kb=%d, param_.rope_ka=%d, param_.rope_kb=%d\n", param_.M, param_.N, param_.ka, param_.kb, param_.rope_ka, param_.rope_kb); + param_.matmulResultSize = tilingParam.matmulResultSize; // ceil(seq_len, chunk_size) * chunk_size = 8320 + + // tiling data for topk + param_.maxSeqLen = tilingParam.maxSeqLen; + param_.sink = tilingParam.sink; + param_.recent = tilingParam.recent; + param_.blockCount = tilingParam.blockCount; + param_.topKInnerSize = tilingParam.topKInnerSize; + param_.topKValueSize = tilingParam.topKValueSize; + param_.topKIdexSize = tilingParam.topKIdexSize; + topkTiling_ = topkTiling; + tileSeqLenSize_ = VECTOR_CUBE_RATIO * param_.tileN2; + + // support offload + param_.supportOffload = tilingParam.supportOffload > 0; + } + + __aicore__ inline void InitGlobalBuffers(GM_ADDR query, GM_ADDR keyCompressed, GM_ADDR keyCompressedRope, GM_ADDR k, GM_ADDR seqLen, + GM_ADDR chunkSize, GM_ADDR keyBlockTable, GM_ADDR mask, GM_ADDR indices, GM_ADDR workSpace) + { + queryGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(query)); + keyCompressedGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(keyCompressed)); + keyCompressedRopeGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(keyCompressedRope)); + kGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(k)); + seqLenGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(seqLen)); + chunkSizeGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(chunkSize)); + keyBlockTableGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(keyBlockTable)); + maskGm_.SetGlobalBuffer(reinterpret_cast<__gm__ bool*>(mask)); + supportMask_ = maskGm_.GetPhyAddr() != nullptr; + isChunkTopK_ = chunkSizeGm_.GetPhyAddr() != nullptr; + isContinuousBatch_ = keyBlockTableGm_.GetPhyAddr() != nullptr; + indicesGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(indices)); + + uint32_t unpackResultWorkspaceOffset = 0; + uint32_t kNopeUnpackResultWorkspaceOffset = unpackResultWorkspaceOffset + param_.layerSize * 8 /2; + uint32_t matmulResultWorkspaceOffset = kNopeUnpackResultWorkspaceOffset + param_.layerSizeRope * 8 /2; + uint32_t topkValueResultWorkspaceOffset = matmulResultWorkspaceOffset + param_.matmulResultSize * 2; + uint32_t topkIdxResultWorkspaceOffset = topkValueResultWorkspaceOffset + param_.topKValueSize * 2; + uint32_t queryUnpackResultWorkspaceOffset = topkIdxResultWorkspaceOffset + param_.topKIdexSize * 4; + + unpackGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int4b_t*>(workSpace + unpackResultWorkspaceOffset)); + kRopeUnpackGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int4b_t*>(workSpace + kNopeUnpackResultWorkspaceOffset)); + matmulGm_.SetGlobalBuffer(reinterpret_cast<__gm__ half*>(workSpace + matmulResultWorkspaceOffset)); + topkValueGm_.SetGlobalBuffer(reinterpret_cast<__gm__ half*>(workSpace + topkValueResultWorkspaceOffset)); + topkIdxGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(workSpace + topkIdxResultWorkspaceOffset)); + qUnpackGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int4b_t*>(workSpace + queryUnpackResultWorkspaceOffset)); + } + + __aicore__ inline void InitLocalBuffersForUnpackQ() + { + pipe_->InitBuffer(constBuf_, param_.dimension * sizeof(half)); + pipe_->InitBuffer(selectBuf_, param_.headGroupNum * param_.dimension * sizeof(half)); + pipe_->InitBuffer(queryCompressedInQueue_, DOUBLE_BUFFER_NUM, param_.headGroupNum * param_.dimension * sizeof(int8_t) / COMPRESS_RATE); + pipe_->InitBuffer(queryUnpackedOutQueue_, DOUBLE_BUFFER_NUM, param_.headGroupNum * param_.dimension * sizeof(int8_t) / INT4B_TYPE_SIZE_DIV_RATE); + pipe_->InitBuffer(qReduceSumLastRowBuf_, param_.dimension * sizeof(half)); + pipe_->InitBuffer(qReduceSumBuf_, param_.headGroupNum * param_.dimension * sizeof(half)); + } + + __aicore__ inline void InitLocalBuffersForUnpackK() + { + pipe_->InitBuffer(keyCompressedInBuf_, DOUBLE_BUFFER_NUM, param_.tileN1 * param_.dimension * sizeof(int8_t) / COMPRESS_RATE); + pipe_->InitBuffer(keyUnpackedOutBuf_, DOUBLE_BUFFER_NUM, param_.tileN1 * param_.dimension * sizeof(int8_t) / INT4B_TYPE_SIZE_DIV_RATE); + pipe_->InitBuffer(constBuf_, param_.dimension * sizeof(half)); + pipe_->InitBuffer(selectBuf_, MAX_FP16_PROCESS_NUM * MAX_SELECT_AND_CAST_COUNT * sizeof(half)); + } + + __aicore__ inline void InitLocalBuffersForTopK() + { + if (minChunkSize_ > MIN_CHUNK_SIZE) { + pipe_->InitBuffer(chunkReduceMaxValueInQueue_, DOUBLE_BUFFER_NUM, param_.tileN2 * DivCeil(MAX_CHUNK_SIZE, minChunkSize_) * sizeof(half)); + } + pipe_->InitBuffer(topKValueInQueue_, DOUBLE_BUFFER_NUM, DivCeil(param_.tileN2, minChunkSize_) * sizeof(half)); + pipe_->InitBuffer(topKValueInnerInQueue_, 1, param_.topKInnerSize * sizeof(half)); + pipe_->InitBuffer(topKIndexInnerInQueue_, 1, param_.topKInnerSize * sizeof(uint32_t)); + pipe_->InitBuffer(topKValueInnerOutQueue_, 1, param_.maxK * sizeof(half)); + pipe_->InitBuffer(topKIndexInnerOutQueue_, 1, param_.maxK * sizeof(uint32_t)); + } + + __aicore__ inline void InitLocalBuffersForTopKSort() + { + pipe_->InitBuffer(topKIndexSortInQueue_, 1, param_.maxK * sizeof(uint32_t)); + pipe_->InitBuffer(topKIndexSortTmpQueue_, 1, 3 * param_.maxK * sizeof(uint32_t)); + pipe_->InitBuffer(topKIndexSortCalcQueue_, 1, 3 * param_.maxK * sizeof(uint32_t)); + } + + __aicore__ inline void InitLocalBuffersForTableBlock() + { + pipe_->InitBuffer(tableBlockBuf_, param_.blockCount * sizeof(int32_t)); + } + + __aicore__ inline void WriteBlockTableFromTopK( + uint32_t curBatchIdx, + LocalTensor& topKIndexUb, + uint32_t curKScalar, + uint64_t outGmOffset) + { + if ASCEND_IS_AIC { return; } + + LocalTensor blockIdUb = topKIndexInnerInQueue_.AllocTensor(); + LocalTensor tableBlockTensor = tableBlockBuf_.template Get(); + DataCopyExtParams copyInParams{1, static_cast(param_.blockCount * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams copyInPadParams{false, 0, 0, 0}; + DataCopyPad(tableBlockTensor, keyBlockTableGm_[curBatchIdx * param_.blockCount], copyInParams, copyInPadParams); + + ::AscendC::WriteBlockTableFromTopK(curBatchIdx, topKIndexUb, blockIdUb, curKScalar, outGmOffset, + tableBlockTensor, indicesGm_, isContinuousBatch_, param_.blockCount); + + topKIndexInnerInQueue_.FreeTensor(blockIdUb); + } +}; +} // namespace AscendC +#endif // HAMMING_DIST_TOP_K_SPLIT_S_H diff --git a/csrc/ascend/moe/hc_post/CMakeLists.txt b/csrc/ascend/moe/hc_post/CMakeLists.txt new file mode 100644 index 000000000..86b308249 --- /dev/null +++ b/csrc/ascend/moe/hc_post/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/hc_post/op_host/CMakeLists.txt b/csrc/ascend/moe/hc_post/op_host/CMakeLists.txt new file mode 100644 index 000000000..016430d3b --- /dev/null +++ b/csrc/ascend/moe/hc_post/op_host/CMakeLists.txt @@ -0,0 +1,65 @@ +# This program is free software, you can redistribute it and/or modify it. +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This file is a part of the CANN Open Software. +# Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ====================================================================================================================== + +# add_ops_compile_options( +# OP_NAME HcPost +# OPTIONS --cce-auto-sync=off +# -Wno-deprecated-declarations +# -Werror +# -mllvm -cce-aicore-hoist-movemask=false +# --op_relocatable_kernel_binary=true +# ) + +# set(hc_post_depends transformer/attention/hc_post PARENT_SCOPE) + +# target_sources(op_host_aclnn PRIVATE +# op_host/hc_post_def.cpp +# ) + +# target_sources(optiling PRIVATE +# op_host/hc_post_tiling.cpp +# ) + +# if (NOT BUILD_OPEN_PROJECT) +# target_sources(opmaster_ct PRIVATE +# op_host/hc_post_tiling.cpp +# ) +# endif () + +# target_include_directories(optiling PRIVATE +# ${CMAKE_CURRENT_SOURCE_DIR}/op_host +# ) + +# target_sources(opsproto PRIVATE +# op_host/hc_post_proto.cpp +# ) + + + +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + hc_post_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME HcPost + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true +) + + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE hc_post ACLNNTYPE aclnn) +endif() \ No newline at end of file diff --git a/csrc/ascend/moe/hc_post/op_host/hc_post_def.cpp b/csrc/ascend/moe/hc_post/op_host/hc_post_def.cpp new file mode 100644 index 000000000..271a81212 --- /dev/null +++ b/csrc/ascend/moe/hc_post/op_host/hc_post_def.cpp @@ -0,0 +1,54 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_post_def.cpp + * \brief HcPost op post config + */ + +#include +#include "register/op_def_registry.h" + +namespace ops { +class HcPost : public OpDef { +public: + explicit HcPost(const char *name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("residual") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("post") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("comb") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + } +}; +OP_ADD(HcPost); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/moe/hc_post/op_host/hc_post_proto.cpp b/csrc/ascend/moe/hc_post/op_host/hc_post_proto.cpp new file mode 100644 index 000000000..875b5dd09 --- /dev/null +++ b/csrc/ascend/moe/hc_post/op_host/hc_post_proto.cpp @@ -0,0 +1,64 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_post_proto.cpp + * \brief + */ + +#include +#include +#include "error/ops_error.h" + +using namespace ge; +namespace ops { +const int32_t INPUT_IDX_X = 0; +const int32_t INPUT_IDX_RESIDUAL = 1; +const int32_t INPUT_IDX_POST = 2; +const int32_t INPUT_IDX_COMB = 3; +const int32_t INDEX_OUTPUT_Y = 0; + + +static ge::graphStatus InferShape4HcPost(gert::InferShapeContext* context) +{ + OPS_LOG_I(context->GetNodeName(), "Begin to do InferShape4HcPost."); + + const gert::Shape* xShape = context->GetInputShape(INPUT_IDX_X); + OPS_LOG_E_IF_NULL(context, xShape, return ge::GRAPH_FAILED); + + const gert::Shape* residualShape = context->GetInputShape(INPUT_IDX_RESIDUAL); + OPS_LOG_E_IF_NULL(context, residualShape, return ge::GRAPH_FAILED); + + const gert::Shape* postShape = context->GetInputShape(INPUT_IDX_POST); + OPS_LOG_E_IF_NULL(context, postShape, return ge::GRAPH_FAILED); + + const gert::Shape* combShape = context->GetInputShape(INPUT_IDX_COMB); + OPS_LOG_E_IF_NULL(context, combShape, return ge::GRAPH_FAILED); + + auto yShape = context->GetOutputShape(INDEX_OUTPUT_Y); + *yShape = *residualShape; + + OPS_LOG_I(context->GetNodeName(), "End to do InferShape4HcPost"); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDtype4HcPost(gert::InferDataTypeContext* context) +{ + OPS_LOG_I(context->GetNodeName(), "InferDtype4HcPost enter"); + const auto xDtype = context->GetInputDataType(INPUT_IDX_X); + context->SetOutputDataType(INDEX_OUTPUT_Y, xDtype); + OPS_LOG_I(context->GetNodeName(), "InferDtype4HcPost end"); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(HcPost) + .InferShape(InferShape4HcPost) + .InferDataType(InferDtype4HcPost); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/moe/hc_post/op_host/hc_post_tiling.cpp b/csrc/ascend/moe/hc_post/op_host/hc_post_tiling.cpp new file mode 100644 index 000000000..095c7605f --- /dev/null +++ b/csrc/ascend/moe/hc_post/op_host/hc_post_tiling.cpp @@ -0,0 +1,231 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_post_tiling.cpp + * \brief + */ + +#include "hc_post_tiling.h" +#include "hc_post_tiling_arch35.h" + +namespace optiling { +constexpr int64_t DEFAULT_DEAL_DPARAM = 2048; + +ge::graphStatus HcPostTiling::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, OPS_LOG_E(context_->GetNodeName(), "get platformInfo nullptr."), + return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + coreNum_ = ascendcPlatform.GetCoreNumAiv(); + OPS_ERR_IF( + coreNum_ <= 0, OPS_LOG_E(context_->GetNodeName(), "coreNum must be greater than 0."), + return ge::GRAPH_FAILED); + + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + ubSize_ = static_cast(ubSizePlatForm); + OPS_ERR_IF( + ubSize_ <= 0, OPS_LOG_E(context_->GetNodeName(), "ubSize must be greater than 0."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPostTiling::GetShapeInfo() +{ + OPS_ERR_IF( + context_ == nullptr, OPS_LOG_E("HcPostTiling", "context can not be nullptr."), + return ge::GRAPH_FAILED); + + if (GetInputShapeInfo() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + // dtype校验 + if (GetInputDtypeInfo() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPostTiling::GetInputShapeInfo() +{ + auto xInput = context_->GetInputShape(INPUT_IDX_X); + OPS_ERR_IF(xInput == nullptr, OPS_LOG_E(context_->GetNodeName(), "get xInput nullptr."), + return ge::GRAPH_FAILED); + gert::Shape xShape = xInput->GetStorageShape(); + size_t xDimsN = xShape.GetDimNum(); + OPS_ERR_IF((xDimsN != CONST2 && xDimsN != CONST3), + OPS_LOG_E(context_->GetNodeName(), "xInput dim:%lu should be 2 or 3.", xDimsN), + return ge::GRAPH_FAILED); + + auto residualInput = context_->GetInputShape(INPUT_IDX_RESIDUAL); + OPS_ERR_IF(residualInput == nullptr, OPS_LOG_E(context_->GetNodeName(), "get residualInput nullptr."), + return ge::GRAPH_FAILED); + gert::Shape residualShape = residualInput->GetStorageShape(); + size_t residualDimsN = residualShape.GetDimNum(); + OPS_ERR_IF((residualDimsN != xDimsN + 1), + OPS_LOG_E(context_->GetNodeName(), "residualInput dim:%lu should be %lu.", residualDimsN, xDimsN + 1), + return ge::GRAPH_FAILED); + + auto postInput = context_->GetInputShape(INPUT_IDX_POST); + OPS_ERR_IF(postInput == nullptr, OPS_LOG_E(context_->GetNodeName(), "get residualInput nullptr."), + return ge::GRAPH_FAILED); + gert::Shape postShape = postInput->GetStorageShape(); + size_t postDimsN = postShape.GetDimNum(); + OPS_ERR_IF((postDimsN != xDimsN), + OPS_LOG_E(context_->GetNodeName(), "postInput dim:%lu should be %lu.", postDimsN, xDimsN), + return ge::GRAPH_FAILED); + + auto combInput = context_->GetInputShape(INPUT_IDX_COMB); + OPS_ERR_IF(combInput == nullptr, OPS_LOG_E(context_->GetNodeName(), "get residualInput nullptr."), + return ge::GRAPH_FAILED); + gert::Shape combShape = combInput->GetStorageShape(); + size_t combDimsN = combShape.GetDimNum(); + OPS_ERR_IF((combDimsN != xDimsN + 1), + OPS_LOG_E(context_->GetNodeName(), "combInput dim:%lu should be %lu.", combDimsN, xDimsN + 1), + return ge::GRAPH_FAILED); + if (xDimsN == CONST2) { + bsParam_ = xShape.GetDim(DIM_INDEX_0); + dParam_ = xShape.GetDim(DIM_INDEX_1); + hcParam_ = residualShape.GetDim(DIM_INDEX_1); + } else { + bsParam_ = xShape.GetDim(DIM_INDEX_0) * xShape.GetDim(DIM_INDEX_1); + dParam_ = xShape.GetDim(DIM_INDEX_2); + hcParam_ = residualShape.GetDim(DIM_INDEX_2); + } + tilingData_.set_bsParam(bsParam_); + tilingData_.set_dParam(dParam_); + tilingData_.set_hcParam(hcParam_); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPostTiling::GetInputDtypeInfo() +{ + auto xDesc = context_->GetInputDesc(INPUT_IDX_X); + OPS_ERR_IF(xDesc == nullptr, OPS_LOG_E(context_->GetNodeName(), "get xDesc nullptr."), + return ge::GRAPH_FAILED); + auto xDtype = xDesc->GetDataType(); + OPS_ERR_IF( + (xDtype != ge::DT_FLOAT16 && xDtype != ge::DT_BF16 && xDtype != ge::DT_FLOAT), + OPS_LOG_E(context_->GetNodeName(), "xDtype is not supported."), + return ge::GRAPH_FAILED); + + auto residualDesc = context_->GetInputDesc(INPUT_IDX_RESIDUAL); + OPS_ERR_IF(residualDesc == nullptr, OPS_LOG_E(context_->GetNodeName(), "get residualDesc nullptr."), + return ge::GRAPH_FAILED); + ge::DataType residualDtype = residualDesc->GetDataType(); + OPS_ERR_IF( + (residualDtype != xDtype), + OPS_LOG_E(context_->GetNodeName(), "residualDtype is not equal to xDtype."), + return ge::GRAPH_FAILED); + + auto postDesc = context_->GetInputDesc(INPUT_IDX_POST); + OPS_ERR_IF(postDesc == nullptr, OPS_LOG_E(context_->GetNodeName(), "get postDesc nullptr."), + return ge::GRAPH_FAILED); + ge::DataType postDtype = postDesc->GetDataType(); + OPS_ERR_IF( + (postDtype != ge::DT_FLOAT16 && postDtype != ge::DT_BF16 && postDtype != ge::DT_FLOAT), + OPS_LOG_E(context_->GetNodeName(), "postDtype is not supported."), + return ge::GRAPH_FAILED); + + auto combDesc = context_->GetInputDesc(INPUT_IDX_COMB); + OPS_ERR_IF(combDesc == nullptr, OPS_LOG_E(context_->GetNodeName(), "get combDesc nullptr."), + return ge::GRAPH_FAILED); + ge::DataType combDtype = combDesc->GetDataType(); + OPS_ERR_IF( + (combDtype != postDtype), + OPS_LOG_E(context_->GetNodeName(), "combDtype is not equal to postDtype."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPostTiling::DoOpTiling() +{ + int64_t batchSize = bsParam_; + + int64_t useCoreNum = batchSize < coreNum_ ? batchSize : coreNum_; + int64_t batchOneCore = CeilDiv(batchSize, static_cast(useCoreNum)); + int64_t batchOneCoreTail = batchOneCore - 1; + int64_t frontCore = batchSize - batchOneCoreTail * useCoreNum; + tilingData_.set_usedCoreNum(useCoreNum); + tilingData_.set_batchOneCore(batchOneCore); + tilingData_.set_batchOneCoreTail(batchOneCoreTail); + tilingData_.set_frontCore(frontCore); + int64_t dSplitTime = dParam_ / DEFAULT_DEAL_DPARAM; + tilingData_.set_dSplitTime(dSplitTime); + tilingData_.set_dOnceDealing(DEFAULT_DEAL_DPARAM); + tilingData_.set_dLastDealing(dParam_ - dSplitTime * DEFAULT_DEAL_DPARAM); + context_->SetBlockDim(useCoreNum); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPostTiling::PostTiling() +{ + context_->SetTilingKey(0); + size_t* workspaces = context_->GetWorkspaceSizes(1); + workspaces[0] = WORKSPACE_SIZE; + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPostTiling::RunTiling() +{ + ge::graphStatus ret = GetShapeInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetPlatformInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = DoOpTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + return PostTiling(); +} + +ge::graphStatus Tiling4HcPost(gert::TilingContext* context) +{ + OPS_LOG_I(context->GetNodeName(), "TilingForHcPost running."); + OPS_ERR_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("TilingForHcPost", "Tiling context is null"), + return ge::GRAPH_FAILED); + auto platformInfo = context->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, OPS_REPORT_VECTOR_INNER_ERR("TilingForHcPost", "Tiling platformInfo is null"), + return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + auto socVersion = ascendcPlatform.GetSocVersion(); + if (socVersion == platform_ascendc::SocVersion::ASCEND950) { + OPS_LOG_I(context, "Using arch35 tiling for ASCEND950"); + HcPostTilingRegbase tiling(context); + return tiling.RunTilingRegbase(); + } + HcPostTiling tiling(context); + return tiling.RunTiling(); +} + +ge::graphStatus TilingPrepare4HcPost(gert::TilingParseContext* context) +{ + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(HcPost) + .Tiling(Tiling4HcPost) + .TilingParse(TilingPrepare4HcPost); + +} // namespace optiling diff --git a/csrc/ascend/moe/hc_post/op_host/hc_post_tiling.h b/csrc/ascend/moe/hc_post/op_host/hc_post_tiling.h new file mode 100644 index 000000000..c8710ed26 --- /dev/null +++ b/csrc/ascend/moe/hc_post/op_host/hc_post_tiling.h @@ -0,0 +1,112 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_post_tiling.h + * \brief + */ +#ifndef HC_POST_TILING_H_ +#define HC_POST_TILING_H_ + +#include "exe_graph/runtime/tiling_context.h" +#include "tiling/platform/platform_ascendc.h" +#include "register/op_def_registry.h" +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "error/ops_error.h" +#include "platform/platform_info.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(HcPostTilingData) +TILING_DATA_FIELD_DEF(int64_t, usedCoreNum); +TILING_DATA_FIELD_DEF(int64_t, bsParam); +TILING_DATA_FIELD_DEF(int64_t, hcParam); +TILING_DATA_FIELD_DEF(int64_t, dParam); +TILING_DATA_FIELD_DEF(int64_t, batchOneCore); +TILING_DATA_FIELD_DEF(int64_t, batchOneCoreTail); +TILING_DATA_FIELD_DEF(int64_t, frontCore); +TILING_DATA_FIELD_DEF(int64_t, dSplitTime); +TILING_DATA_FIELD_DEF(int64_t, dOnceDealing); +TILING_DATA_FIELD_DEF(int64_t, dLastDealing); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(HcPost, HcPostTilingData) + +struct HcPostCompileInfo { +}; + +class HcPostTilingRegbase { +public: + explicit HcPostTilingRegbase(gert::TilingContext* context) : context_(context) + {} + ~HcPostTilingRegbase() + {} + ge::graphStatus RunTilingRegbase(); + +protected: + ge::graphStatus DoOpTilingRegbase(); + ge::graphStatus GetPlatformInfoRegbase(); + ge::graphStatus GetShapeInfoRegbase(); + ge::graphStatus GetInputShapeInfoRegbase(); + ge::graphStatus PostTilingRegbase(); + +private: + ge::graphStatus GetInputDtypeInfoRegbase(); + +private: + HcPostTilingData tilingRegbaseData_; + + int64_t coreNum_ = 0; + int64_t ubSize_ = 0; + int64_t ubBlockSize_ = 0; + + int64_t bsParam_ = 0; + int64_t dParam_ = 0; + int64_t hcParam_ = 0; + int64_t batchSize_ = 0; + int64_t tilingKey_ = 0; + + gert::TilingContext *context_ = nullptr; +}; + +class HcPostTiling { +public: + explicit HcPostTiling(gert::TilingContext* context) : context_(context) + {} + ~HcPostTiling() + {} + ge::graphStatus RunTiling(); + +protected: + ge::graphStatus DoOpTiling(); + ge::graphStatus GetPlatformInfo(); + ge::graphStatus GetShapeInfo(); + ge::graphStatus GetInputShapeInfo(); + ge::graphStatus PostTiling(); + +private: + ge::graphStatus GetInputDtypeInfo(); + +private: + HcPostTilingData tilingData_; + + int64_t coreNum_ = 0; + int64_t ubSize_ = 0; + int64_t ubBlockSize_ = 0; + + int64_t bsParam_ = 0; + int64_t dParam_ = 0; + int64_t hcParam_ = 0; + + gert::TilingContext *context_ = nullptr; +}; + +} // namespace optiling + +#endif // HC_POST_TILING_H_ diff --git a/csrc/ascend/moe/hc_post/op_host/hc_post_tiling_arch35.h b/csrc/ascend/moe/hc_post/op_host/hc_post_tiling_arch35.h new file mode 100644 index 000000000..72a2b789c --- /dev/null +++ b/csrc/ascend/moe/hc_post/op_host/hc_post_tiling_arch35.h @@ -0,0 +1,234 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_post_tiling_arch35.cpp + * \brief + */ + +#include "hc_post_tiling.h" + +namespace optiling { + +constexpr int32_t INPUT_IDX_X = 0; +constexpr int32_t INPUT_IDX_RESIDUAL = 1; +constexpr int32_t INPUT_IDX_POST = 2; +constexpr int32_t INPUT_IDX_COMB = 3; +constexpr int32_t INDEX_OUTPUT_Y = 0; +constexpr int32_t DIM_INDEX_0 = 0; +constexpr int32_t DIM_INDEX_1 = 1; +constexpr int32_t DIM_INDEX_2 = 2; +constexpr int32_t DIM_INDEX_3 = 3; +constexpr size_t CONST1 = 1; +constexpr size_t CONST2 = 2; +constexpr size_t CONST3 = 3; +constexpr size_t CONST4 = 4; +constexpr size_t WORKSPACE_SIZE = static_cast(16 * 1024 * 1024); +constexpr int64_t ONCE_DEAL_DPARAM = 4096; + +template +static inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd))); +} + +template +static inline T CeilAlign(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)) * (rnd)); +} + +ge::graphStatus HcPostTilingRegbase::GetPlatformInfoRegbase() +{ + auto platformInfo = context_->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, OPS_LOG_E(context_->GetNodeName(), "get platformInfo nullptr."), + return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + coreNum_ = ascendcPlatform.GetCoreNumAiv(); + OPS_ERR_IF( + coreNum_ <= 0, OPS_LOG_E(context_->GetNodeName(), "coreNum must be greater than 0."), + return ge::GRAPH_FAILED); + + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + ubSize_ = static_cast(ubSizePlatForm); + OPS_ERR_IF( + ubSize_ <= 0, OPS_LOG_E(context_->GetNodeName(), "ubSize must be greater than 0."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPostTilingRegbase::GetShapeInfoRegbase() +{ + OPS_ERR_IF( + context_ == nullptr, OPS_LOG_E("HcPostTilingRegBase", "context can not be nullptr."), + return ge::GRAPH_FAILED); + + if (GetInputShapeInfoRegbase() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + // dtype校验 + if (GetInputDtypeInfoRegbase() != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPostTilingRegbase::GetInputShapeInfoRegbase() +{ + auto xInput = context_->GetInputShape(INPUT_IDX_X); + OPS_ERR_IF(xInput == nullptr, OPS_LOG_E(context_->GetNodeName(), "get xInput nullptr."), + return ge::GRAPH_FAILED); + gert::Shape xShape = xInput->GetStorageShape(); + size_t xDimsN = xShape.GetDimNum(); + OPS_ERR_IF((xDimsN != CONST2 && xDimsN != CONST3), + OPS_LOG_E(context_->GetNodeName(), "xInput dim:%lu should be 2 or 3.", xDimsN), + return ge::GRAPH_FAILED); + + auto residualInput = context_->GetInputShape(INPUT_IDX_RESIDUAL); + OPS_ERR_IF(residualInput == nullptr, OPS_LOG_E(context_->GetNodeName(), "get residualInput nullptr."), + return ge::GRAPH_FAILED); + gert::Shape residualShape = residualInput->GetStorageShape(); + size_t residualDimsN = residualShape.GetDimNum(); + OPS_ERR_IF((residualDimsN != xDimsN + 1), + OPS_LOG_E(context_->GetNodeName(), "residualInput dim:%lu should be %lu.", residualDimsN, xDimsN + 1), + return ge::GRAPH_FAILED); + + auto postInput = context_->GetInputShape(INPUT_IDX_POST); + OPS_ERR_IF(postInput == nullptr, OPS_LOG_E(context_->GetNodeName(), "get residualInput nullptr."), + return ge::GRAPH_FAILED); + gert::Shape postShape = postInput->GetStorageShape(); + size_t postDimsN = postShape.GetDimNum(); + OPS_ERR_IF((postDimsN != xDimsN), + OPS_LOG_E(context_->GetNodeName(), "postInput dim:%lu should be %lu.", postDimsN, xDimsN), + return ge::GRAPH_FAILED); + + auto combInput = context_->GetInputShape(INPUT_IDX_COMB); + OPS_ERR_IF(combInput == nullptr, OPS_LOG_E(context_->GetNodeName(), "get residualInput nullptr."), + return ge::GRAPH_FAILED); + gert::Shape combShape = combInput->GetStorageShape(); + size_t combDimsN = combShape.GetDimNum(); + OPS_ERR_IF((combDimsN != xDimsN + 1), + OPS_LOG_E(context_->GetNodeName(), "combInput dim:%lu should be %lu.", combDimsN, xDimsN + 1), + return ge::GRAPH_FAILED); + if (xDimsN == CONST2) { + bsParam_ = xShape.GetDim(DIM_INDEX_0); + dParam_ = xShape.GetDim(DIM_INDEX_1); + hcParam_ = residualShape.GetDim(DIM_INDEX_1); + } else { + bsParam_ = xShape.GetDim(DIM_INDEX_0) * xShape.GetDim(DIM_INDEX_1); + dParam_ = xShape.GetDim(DIM_INDEX_2); + hcParam_ = residualShape.GetDim(DIM_INDEX_2); + } + tilingRegbaseData_.set_bsParam(bsParam_); + tilingRegbaseData_.set_dParam(dParam_); + tilingRegbaseData_.set_hcParam(hcParam_); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPostTilingRegbase::GetInputDtypeInfoRegbase() +{ + auto xDesc = context_->GetInputDesc(INPUT_IDX_X); + OPS_ERR_IF(xDesc == nullptr, OPS_LOG_E(context_->GetNodeName(), "get xDesc nullptr."), + return ge::GRAPH_FAILED); + auto xDtype = xDesc->GetDataType(); + OPS_ERR_IF( + (xDtype != ge::DT_FLOAT16 && xDtype != ge::DT_BF16 && xDtype != ge::DT_FLOAT), + OPS_LOG_E(context_->GetNodeName(), "xDtype is not supported."), + return ge::GRAPH_FAILED); + + auto residualDesc = context_->GetInputDesc(INPUT_IDX_RESIDUAL); + OPS_ERR_IF(residualDesc == nullptr, OPS_LOG_E(context_->GetNodeName(), "get residualDesc nullptr."), + return ge::GRAPH_FAILED); + ge::DataType residualDtype = residualDesc->GetDataType(); + OPS_ERR_IF( + (residualDtype != xDtype), + OPS_LOG_E(context_->GetNodeName(), "residualDtype is not equal to xDtype."), + return ge::GRAPH_FAILED); + + auto postDesc = context_->GetInputDesc(INPUT_IDX_POST); + OPS_ERR_IF(postDesc == nullptr, OPS_LOG_E(context_->GetNodeName(), "get postDesc nullptr."), + return ge::GRAPH_FAILED); + ge::DataType postDtype = postDesc->GetDataType(); + OPS_ERR_IF( + (postDtype != ge::DT_FLOAT16 && postDtype != ge::DT_BF16 && postDtype != ge::DT_FLOAT), + OPS_LOG_E(context_->GetNodeName(), "postDtype is not supported."), + return ge::GRAPH_FAILED); + + auto combDesc = context_->GetInputDesc(INPUT_IDX_COMB); + OPS_ERR_IF(combDesc == nullptr, OPS_LOG_E(context_->GetNodeName(), "get combDesc nullptr."), + return ge::GRAPH_FAILED); + ge::DataType combDtype = combDesc->GetDataType(); + OPS_ERR_IF( + (combDtype != postDtype), + OPS_LOG_E(context_->GetNodeName(), "combDtype is not equal to postDtype."), + return ge::GRAPH_FAILED); + + if (xDtype == ge::DT_FLOAT) { + tilingKey_ = 0; + } else { + tilingKey_ = 1; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPostTilingRegbase::DoOpTilingRegbase() +{ + int64_t batchSize = bsParam_; + context_->SetTilingKey(tilingKey_); + int64_t useCoreNum = batchSize < coreNum_ ? batchSize : coreNum_; + int64_t batchOneCore = CeilDiv(batchSize, static_cast(useCoreNum)); + int64_t batchOneCoreTail = batchOneCore - 1; + int64_t frontCore = batchSize - batchOneCoreTail * useCoreNum; + tilingRegbaseData_.set_usedCoreNum(useCoreNum); + tilingRegbaseData_.set_batchOneCore(batchOneCore); + tilingRegbaseData_.set_batchOneCoreTail(batchOneCoreTail); + tilingRegbaseData_.set_frontCore(frontCore); + context_->SetBlockDim(useCoreNum); + + int64_t dSplitTime = dParam_ / ONCE_DEAL_DPARAM; + tilingRegbaseData_.set_dSplitTime(dSplitTime); + tilingRegbaseData_.set_dOnceDealing(ONCE_DEAL_DPARAM); + tilingRegbaseData_.set_dLastDealing(dParam_ - dSplitTime * ONCE_DEAL_DPARAM); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPostTilingRegbase::PostTilingRegbase() +{ + size_t* workspaces = context_->GetWorkspaceSizes(1); + workspaces[0] = WORKSPACE_SIZE; + tilingRegbaseData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingRegbaseData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPostTilingRegbase::RunTilingRegbase() +{ + ge::graphStatus ret = GetShapeInfoRegbase(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetPlatformInfoRegbase(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = DoOpTilingRegbase(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + return PostTilingRegbase(); +} +} +// namespace optiling \ No newline at end of file diff --git a/csrc/ascend/moe/hc_post/op_kernel/hc_post.cpp b/csrc/ascend/moe/hc_post/op_kernel/hc_post.cpp new file mode 100644 index 000000000..445ed926f --- /dev/null +++ b/csrc/ascend/moe/hc_post/op_kernel/hc_post.cpp @@ -0,0 +1,58 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_post_apt.cpp + * \brief + */ + +#include "kernel_operator.h" +#if defined(__DAV_C310__) + #include "hc_post_float32.h" + #include "hc_post_bfloat16.h" +#endif +#include "hc_post_d_split.h" + +using namespace AscendC; +using namespace HcPost; +#if defined(__DAV_C310__) + using namespace HcPostRegBase; +#endif + +#define HC_POST_FLOAT 0 +#define HC_POST_BFLOAT16 1 + +extern "C" __global__ __aicore__ void hc_post(GM_ADDR x, GM_ADDR residual, GM_ADDR post, + GM_ADDR comb, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) +{ + TPipe pipe; + #if defined(__DAV_C310__) + GET_TILING_DATA_WITH_STRUCT(HcPostTilingData, tilingData, tiling); + const HcPostTilingData *__restrict hcPostTilingData = &tilingData; + if (TILING_KEY_IS(HC_POST_FLOAT)) { + HcPostRegBaseFloat32 op; + op.Init(x, residual, post, comb, y, workspace, hcPostTilingData, &pipe); + op.Process(); + return; + } else if (TILING_KEY_IS(HC_POST_BFLOAT16)) { + HcPostRegBaseBfloat16 op; + op.Init(x, residual, post, comb, y, workspace, hcPostTilingData, &pipe); + op.Process(); + return; + } + #else + GET_TILING_DATA_WITH_STRUCT(HcPostTilingData, tilingData, tiling); + const HcPostTilingData *__restrict hcPostTilingData = &tilingData; + HcPostKernelDSplit op; + op.Init(x, residual, post, comb, y, workspace, hcPostTilingData, &pipe); + op.Process(); + return; + #endif +} \ No newline at end of file diff --git a/csrc/ascend/moe/hc_post/op_kernel/hc_post_bfloat16.h b/csrc/ascend/moe/hc_post/op_kernel/hc_post_bfloat16.h new file mode 100644 index 000000000..e23522066 --- /dev/null +++ b/csrc/ascend/moe/hc_post/op_kernel/hc_post_bfloat16.h @@ -0,0 +1,351 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_post_bfloat16.h + * \brief + */ +#ifndef HC_POST_BFLOAT16_H +#define HC_POST_BFLOAT16_H + +#include "kernel_operator.h" + +namespace HcPostRegBase { +using namespace AscendC; + +template +class HcPostRegBaseBfloat16 { +public: + __aicore__ inline HcPostRegBaseBfloat16() {}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR residual, GM_ADDR post, GM_ADDR comb, GM_ADDR y, GM_ADDR workspace, + const HcPostTilingData *tilingData, TPipe *pipe); + __aicore__ inline void Process(); + __aicore__ inline void DataCopyInX(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset); + __aicore__ inline void DataCopyInPost(int64_t batchIndex); + __aicore__ inline void DataCopyInResidual(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset); + __aicore__ inline void DataCopyInComb(int64_t batchIndex); + __aicore__ inline void DataCopyOut(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset); + __aicore__ inline void DoProcess(int64_t batchSize); + __aicore__ inline void DoCompute(LocalTensor sumTempBuf, LocalTensor postUb, LocalTensor combUb, int64_t batchIndex, int64_t dOffset, int64_t dDealing); + __aicore__ inline void DoMulAndAdd(LocalTensor xUb, LocalTensor postUb, LocalTensor residualUb, LocalTensor combUb, LocalTensor sumTempBuf, int64_t dOnceDealing); + +private: + TPipe* pipe_; + const HcPostTilingData* tiling_; + constexpr static AscendC::MicroAPI::CastTrait castB16ToB32 = { AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::UNKNOWN, AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::UNKNOWN }; + + int32_t blkIdx_ = -1; + int64_t batch_ = 0; + int64_t hcParam_ = 0; + int64_t dParam_ = 0; + int64_t batchOneCoreTail_ = 0; + int64_t batchOneCore_ = 0; + int64_t isFrontCore_ = 0; + int64_t dParamAlign_ = 0; + int64_t dOnceDealing_ = 0; + int64_t dLastDealing_ = 0; + int64_t dSplitTime_ = 0; + int64_t dParamOnceAlign_ = 0; + static constexpr int32_t ONE_BLOCK_SIZE = 32; + int32_t perBlock32 = ONE_BLOCK_SIZE / sizeof(float); + + GlobalTensor xGm_; + GlobalTensor residualGm_; + GlobalTensor postGm_; + GlobalTensor combGm_; + GlobalTensor yGm_; + + TQue xQue_; + TQue residualQue_; + TQue postQue_; + TQue combQue_; + TQue sumQue_; + TBuf sumTempBuf_; +}; + +template +__aicore__ inline void HcPostRegBaseBfloat16::Init(GM_ADDR x, GM_ADDR residual, GM_ADDR post, GM_ADDR comb, GM_ADDR y, + GM_ADDR workspace, const HcPostTilingData *tilingData, TPipe *pipe) +{ + blkIdx_ = GetBlockIdx(); + if (blkIdx_ >= tilingData->usedCoreNum) { + return; + } + tiling_ = tilingData; + pipe_ = pipe; + hcParam_ = tilingData->hcParam; + dParam_ = tilingData->dParam; + batchOneCoreTail_ = tilingData->batchOneCoreTail; + batchOneCore_ = tilingData->batchOneCore; + isFrontCore_ = blkIdx_ < tilingData->frontCore; + int64_t frontCore = tilingData->frontCore; + dOnceDealing_ = tilingData->dOnceDealing; + dLastDealing_ = tilingData->dLastDealing; + dSplitTime_ = tilingData->dSplitTime; + dParamAlign_ = (dParam_ + perBlock32 - 1) / perBlock32 * perBlock32; + dParamOnceAlign_ = (dOnceDealing_ + perBlock32 - 1) / perBlock32 * perBlock32; + + int64_t xOffset = blkIdx_ * batchOneCore_ * dParam_; + int64_t residualOffset = blkIdx_ * batchOneCore_ * hcParam_ * dParam_; + int64_t postOffset = blkIdx_ * batchOneCore_ * hcParam_; + int64_t combOffset = blkIdx_ * batchOneCore_ * hcParam_ * hcParam_; + int64_t yOffset = blkIdx_ * batchOneCore_ * hcParam_ * dParam_; + if (!isFrontCore_) { + xOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * dParam_; + residualOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_ * dParam_; + postOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_; + combOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_ * hcParam_; + yOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_ * dParam_; + } + xGm_.SetGlobalBuffer((__gm__ T1 *)x + xOffset); + residualGm_.SetGlobalBuffer((__gm__ T1 *)residual + residualOffset); + postGm_.SetGlobalBuffer((__gm__ T2 *)post + postOffset); + combGm_.SetGlobalBuffer((__gm__ T2 *)comb + combOffset); + yGm_.SetGlobalBuffer((__gm__ T1 *)y + yOffset); + + pipe_->InitBuffer(xQue_, 2, dParamOnceAlign_ * sizeof(T1)); + pipe_->InitBuffer(residualQue_, 2, hcParam_ * dParamOnceAlign_ * sizeof(T1)); + pipe_->InitBuffer(postQue_, 2, hcParam_ * sizeof(T2)); + pipe_->InitBuffer(combQue_, 2, hcParam_ * hcParam_ * sizeof(T2)); + pipe_->InitBuffer(sumQue_, 2, hcParam_* dParamOnceAlign_ * sizeof(T1)); + pipe_->InitBuffer(sumTempBuf_, hcParam_ * dParamOnceAlign_ * sizeof(float)); +} + +template +__aicore__ inline void HcPostRegBaseBfloat16::Process() +{ + if (blkIdx_ >= tiling_->usedCoreNum) { + return; + } + if (isFrontCore_) { + DoProcess(tiling_->batchOneCore); + } else { + DoProcess(tiling_->batchOneCoreTail); + } +} + +template +__aicore__ inline void HcPostRegBaseBfloat16::DataCopyInX(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset) +{ + LocalTensor xUb = xQue_.AllocTensor(); + DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = dOnceDealing * sizeof(T1); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(xUb, xGm_[batchIndex * dParam_ + dOffset], copyParams, dataCopyPadParams); + xQue_.EnQue(xUb); +} + +template +__aicore__ inline void HcPostRegBaseBfloat16::DataCopyInPost(int64_t batchIndex) +{ + LocalTensor postUb = postQue_.AllocTensor(); + DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = hcParam_ * sizeof(T2); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(postUb, postGm_[batchIndex * hcParam_], copyParams, dataCopyPadParams); + postQue_.EnQue(postUb); +} + +template +__aicore__ inline void HcPostRegBaseBfloat16::DataCopyInResidual(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset) +{ + LocalTensor residualUb = residualQue_.AllocTensor(); + DataCopyExtParams copyParams; + copyParams.blockCount = hcParam_; + copyParams.blockLen = dOnceDealing * sizeof(T1); + copyParams.srcStride = (dParamAlign_ - dOnceDealing) * sizeof(T1); + copyParams.dstStride = 0; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(residualUb, residualGm_[batchIndex * hcParam_ * dParam_ + dOffset], copyParams, dataCopyPadParams); + residualQue_.EnQue(residualUb); +} + +template +__aicore__ inline void HcPostRegBaseBfloat16::DataCopyInComb(int64_t batchIndex) +{ + LocalTensor combUb = combQue_.AllocTensor(); + DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = hcParam_ * hcParam_ * sizeof(T2); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(combUb, combGm_[batchIndex * hcParam_ * hcParam_], copyParams, dataCopyPadParams); + combQue_.EnQue(combUb); +} + +template +__aicore__ inline void HcPostRegBaseBfloat16::DataCopyOut(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset) +{ + LocalTensor outBuf = sumQue_.DeQue(); + DataCopyExtParams copyParams; + copyParams.blockCount = hcParam_; + copyParams.blockLen = dOnceDealing * sizeof(T1); + copyParams.srcStride = 0; + copyParams.dstStride = (dParamAlign_ - dOnceDealing) * sizeof(T1); + AscendC::DataCopyPad(yGm_[batchIndex * hcParam_ * dParam_ + dOffset], outBuf, copyParams); + sumQue_.FreeTensor(outBuf); +} + +template +__aicore__ inline void HcPostRegBaseBfloat16::DoMulAndAdd(LocalTensor xUb, LocalTensor postUb, LocalTensor residualUb, LocalTensor combUb, LocalTensor sumTempBuf, int64_t dOnceDealing) +{ + uint16_t aTimes = hcParam_; + uint32_t xDealNumAlign = (dOnceDealing + perBlock32 - 1) / perBlock32 * perBlock32; + uint32_t vfLen = 256 / sizeof(float); + uint16_t repeatTimes = dOnceDealing / vfLen; + uint16_t hcTimes = hcParam_; + uint32_t tailNum = dOnceDealing % vfLen; + uint16_t tailLoopTimes = tailNum == 0 ? 0 : 1; + + auto residualAddr = (__ubuf__ T1*)residualUb.GetPhyAddr(); + auto combAddr = (__ubuf__ T2*)combUb.GetPhyAddr(); + auto sumAddr = (__ubuf__ float*)sumTempBuf.GetPhyAddr(); + auto xAddr = (__ubuf__ T1*)xUb.GetPhyAddr(); + auto postAddr = (__ubuf__ T2*)postUb.GetPhyAddr(); + __VEC_SCOPE__ + { + uint32_t xDealNum = static_cast(hcParam_ * dOnceDealing); + AscendC::MicroAPI::RegTensor xReg; + AscendC::MicroAPI::RegTensor postReg; + AscendC::MicroAPI::RegTensor xRegFloat; + AscendC::MicroAPI::RegTensor postRegFloat; + AscendC::MicroAPI::RegTensor residualReg0; + AscendC::MicroAPI::RegTensor residualReg1; + AscendC::MicroAPI::RegTensor residualReg2; + AscendC::MicroAPI::RegTensor residualReg3; + AscendC::MicroAPI::RegTensor combReg0; + AscendC::MicroAPI::RegTensor combReg1; + AscendC::MicroAPI::RegTensor combReg2; + AscendC::MicroAPI::RegTensor combReg3; + AscendC::MicroAPI::RegTensor residualRegFloat0; + AscendC::MicroAPI::RegTensor residualRegFloat1; + AscendC::MicroAPI::RegTensor residualRegFloat2; + AscendC::MicroAPI::RegTensor residualRegFloat3; + AscendC::MicroAPI::RegTensor combRegFloat0; + AscendC::MicroAPI::RegTensor combRegFloat1; + AscendC::MicroAPI::RegTensor combRegFloat2; + AscendC::MicroAPI::RegTensor combRegFloat3; + AscendC::MicroAPI::RegTensor sumRegFloat; + AscendC::MicroAPI::RegTensor sumTempReg0; + AscendC::MicroAPI::RegTensor sumTempReg1; + AscendC::MicroAPI::RegTensor sumTempReg2; + AscendC::MicroAPI::RegTensor sumTempReg3; + AscendC::MicroAPI::MaskReg pMask; + AscendC::MicroAPI::MaskReg pregMain = AscendC::MicroAPI::CreateMask(); + for (uint16_t hcIndex = 0; hcIndex < hcTimes; hcIndex++) { + pMask = AscendC::MicroAPI::UpdateMask(xDealNum); + if constexpr (sizeof(T2) == 2) { + AscendC::MicroAPI::DataCopy(combReg0, combAddr+hcIndex); + AscendC::MicroAPI::DataCopy(combReg1, combAddr+hcParam_+hcIndex); + AscendC::MicroAPI::DataCopy(combReg2, combAddr+2*hcParam_+hcIndex); + AscendC::MicroAPI::DataCopy(combReg3, combAddr+3*hcParam_+hcIndex); + AscendC::MicroAPI::DataCopy(postReg, postAddr + hcIndex); + AscendC::MicroAPI::Cast(combRegFloat0, combReg0, pregMain); + AscendC::MicroAPI::Cast(combRegFloat1, combReg1, pregMain); + AscendC::MicroAPI::Cast(combRegFloat2, combReg2, pregMain); + AscendC::MicroAPI::Cast(combRegFloat3, combReg3, pregMain); + AscendC::MicroAPI::Cast(postRegFloat, postReg, pregMain); + } else { + AscendC::MicroAPI::DataCopy(combRegFloat0, combAddr+hcIndex); + AscendC::MicroAPI::DataCopy(combRegFloat1, combAddr+hcParam_+hcIndex); + AscendC::MicroAPI::DataCopy(combRegFloat2, combAddr+2*hcParam_+hcIndex); + AscendC::MicroAPI::DataCopy(combRegFloat3, combAddr+3*hcParam_+hcIndex); + AscendC::MicroAPI::DataCopy(postRegFloat, postAddr + hcIndex); + } + + for (uint16_t j = 0; j < repeatTimes; j++) { + AscendC::MicroAPI::DataCopy(xReg, xAddr+j*vfLen); + AscendC::MicroAPI::DataCopy(residualReg0, residualAddr+j*vfLen); + AscendC::MicroAPI::DataCopy(residualReg1, residualAddr+xDealNumAlign+j*vfLen); + AscendC::MicroAPI::DataCopy(residualReg2, residualAddr+2*xDealNumAlign+j*vfLen); + AscendC::MicroAPI::DataCopy(residualReg3, residualAddr+3*xDealNumAlign+j*vfLen); + AscendC::MicroAPI::Cast(residualRegFloat0, residualReg0, pMask); + AscendC::MicroAPI::Cast(residualRegFloat1, residualReg1, pMask); + AscendC::MicroAPI::Cast(residualRegFloat2, residualReg2, pMask); + AscendC::MicroAPI::Cast(residualRegFloat3, residualReg3, pMask); + AscendC::MicroAPI::Cast(xRegFloat, xReg, pMask); + AscendC::MicroAPI::Mul(sumTempReg0, residualRegFloat0, combRegFloat0, pMask); + AscendC::MicroAPI::MulAddDst(sumTempReg0, residualRegFloat3, combRegFloat3, pMask); + AscendC::MicroAPI::MulAddDst(sumTempReg0, residualRegFloat1, combRegFloat1, pMask); + AscendC::MicroAPI::MulAddDst(sumTempReg0, residualRegFloat2, combRegFloat2, pMask); + AscendC::MicroAPI::MulAddDst(sumTempReg0, xRegFloat, postRegFloat, pMask); + AscendC::MicroAPI::DataCopy(sumAddr+hcIndex*xDealNumAlign+j*vfLen, sumTempReg0, pMask); + } + for (uint16_t k = 0; k < tailLoopTimes; k++) { + AscendC::MicroAPI::DataCopy(xReg, xAddr+repeatTimes*vfLen); + AscendC::MicroAPI::DataCopy(residualReg0, residualAddr+repeatTimes*vfLen); + AscendC::MicroAPI::DataCopy(residualReg1, residualAddr+xDealNumAlign+repeatTimes*vfLen); + AscendC::MicroAPI::DataCopy(residualReg2, residualAddr+2*xDealNumAlign+repeatTimes*vfLen); + AscendC::MicroAPI::DataCopy(residualReg3, residualAddr+3*xDealNumAlign+repeatTimes*vfLen); + AscendC::MicroAPI::Cast(residualRegFloat0, residualReg0, pMask); + AscendC::MicroAPI::Cast(residualRegFloat1, residualReg1, pMask); + AscendC::MicroAPI::Cast(residualRegFloat2, residualReg2, pMask); + AscendC::MicroAPI::Cast(residualRegFloat3, residualReg3, pMask); + AscendC::MicroAPI::Cast(xRegFloat, xReg, pMask); + AscendC::MicroAPI::Mul(sumTempReg0, residualRegFloat0, combRegFloat0, pMask); + AscendC::MicroAPI::MulAddDst(sumTempReg0, residualRegFloat3, combRegFloat3, pMask); + AscendC::MicroAPI::MulAddDst(sumTempReg0, residualRegFloat1, combRegFloat1, pMask); + AscendC::MicroAPI::MulAddDst(sumTempReg0, residualRegFloat2, combRegFloat2, pMask); + AscendC::MicroAPI::MulAddDst(sumTempReg0, xRegFloat, postRegFloat, pMask); + AscendC::MicroAPI::DataCopy(sumAddr+hcIndex*xDealNumAlign+repeatTimes*vfLen, sumTempReg0, pMask); + } + } + } +} + +template +__aicore__ inline void HcPostRegBaseBfloat16::DoCompute(LocalTensor sumTempBuf, LocalTensor postUb, LocalTensor combUb, int64_t batchIndex, int64_t dOffset, int64_t dDealing) +{ + DataCopyInX(batchIndex, dDealing, dOffset); + LocalTensor xUb = xQue_.DeQue(); + DataCopyInResidual(batchIndex, dDealing, dOffset); + LocalTensor residualUb = residualQue_.DeQue(); + DoMulAndAdd(xUb, postUb, residualUb, combUb, sumTempBuf, dDealing); + LocalTensor sumUb = sumQue_.AllocTensor(); + AscendC::Cast(sumUb, sumTempBuf, AscendC::RoundMode::CAST_RINT, hcParam_ * dOnceDealing_); + sumQue_.EnQue(sumUb); + DataCopyOut(batchIndex, dDealing, dOffset); + residualQue_.FreeTensor(residualUb); + xQue_.FreeTensor(xUb); +} + +template +__aicore__ inline void HcPostRegBaseBfloat16::DoProcess(int64_t batchSize) +{ + LocalTensor sumTempBuf = sumTempBuf_.Get(); + for (int64_t batchIndex = 0; batchIndex < batchSize; batchIndex++) { + DataCopyInPost(batchIndex); + LocalTensor postUb = postQue_.DeQue(); + DataCopyInComb(batchIndex); + LocalTensor combUb = combQue_.DeQue(); + int64_t dOffset = 0; + for (int64_t dIndex = 0; dIndex < dSplitTime_; dIndex++) { + dOffset = dIndex*dOnceDealing_; + DoCompute(sumTempBuf, postUb, combUb, batchIndex, dOffset, dOnceDealing_); + } + if (dLastDealing_ != 0) { + dOffset = dSplitTime_ * dOnceDealing_; + DoCompute(sumTempBuf, postUb, combUb, batchIndex, dOffset, dLastDealing_); + } + combQue_.FreeTensor(combUb); + postQue_.FreeTensor(postUb); + } +} + +} +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/hc_post/op_kernel/hc_post_d_split.h b/csrc/ascend/moe/hc_post/op_kernel/hc_post_d_split.h new file mode 100644 index 000000000..421a0e60a --- /dev/null +++ b/csrc/ascend/moe/hc_post/op_kernel/hc_post_d_split.h @@ -0,0 +1,394 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_post_d_split.h + * \brief + */ +#ifndef HC_POST_D_SPLIT_H +#define HC_POST_D_SPLIT_H + +#include "kernel_operator.h" + +namespace HcPost { +using namespace AscendC; + +constexpr int64_t BLOCK_SIZE = 32; +constexpr int64_t DEFAULT_BLOCK_STRIDE = 1; +constexpr int64_t DEFAULT_REPEAT_STRIDE = 8; +constexpr int64_t ONE_REPEAT_BLOCK_NUMS = 8; +constexpr int64_t REPEAT_SIZE = 256; +constexpr int64_t MAX_REPEAT_STRIDE = 255; +constexpr int64_t REPEAT_NUM = 64; + +__aicore__ inline int32_t CeilDiv(int32_t a, int32_t b) +{ + if (b == 0) { + return a; + } + return (a + b - 1) / b; +} + +__aicore__ inline int32_t CeilAlign(int32_t a, int32_t b) +{ + return CeilDiv(a, b) * b; +} + +template +__aicore__ inline int32_t RoundUp(int32_t num) +{ + int32_t elemNum = BLOCK_SIZE / sizeof(T); + return CeilAlign(num, elemNum); +} + +template +class HcPostKernelDSplit { +public: + __aicore__ inline HcPostKernelDSplit() {}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR residual, GM_ADDR post, GM_ADDR comb, GM_ADDR y, GM_ADDR workspace, + const HcPostTilingData *tilingData, TPipe *pipe); + __aicore__ inline void Process(); + __aicore__ inline void DataCopyInX(int64_t batchIndex, int64_t dLoopTimes, int64_t dealNum); + __aicore__ inline void DataCopyInPost(int64_t batchIndex); + __aicore__ inline void DataCopyInResidual(int64_t batchIndex, int64_t dLoopTimes, int64_t dealNum); + __aicore__ inline void DataCopyInComb(int64_t batchIndex); + __aicore__ inline void DataCopyOut(int64_t batchIndex, int64_t dLoopTimes, int64_t dealNum); + __aicore__ inline void DoCompute(LocalTensor sumTempBuf, LocalTensor postBrcb, LocalTensor comBrcb, LocalTensor combUb, LocalTensor comCastBuf, int64_t batchIndex, int64_t dLoop, int64_t dealNum); + __aicore__ inline void DoProcess(int64_t batchSize); + +private: + TPipe* pipe_; + const HcPostTilingData* tiling_; + + int32_t blkIdx_ = -1; + int64_t batch_ = 0; + int64_t hcParam_ = 0; + int64_t dParam_ = 0; + int64_t dOnceDealing_ = 0; + int64_t dLastDealing_ = 0; + int64_t batchOneCoreTail_ = 0; + int64_t batchOneCore_ = 0; + int64_t dSplitTime_ = 0; + int64_t isFrontCore_ = 0; + int64_t hcParamAlign_ = 0; + static constexpr int32_t ONE_BLOCK_SIZE = 32; + int32_t perBlock32 = ONE_BLOCK_SIZE / sizeof(float); + + GlobalTensor xGm_; + GlobalTensor residualGm_; + GlobalTensor postGm_; + GlobalTensor combGm_; + GlobalTensor yGm_; + + TQue inputQue_; + TQue postQue_; + TQue combQue_; + TQue outQue_; + TBuf inputCastBuf_; + TBuf postCastBuf_; + TBuf combCastBuf_; + TBuf outCastBuf_; + TBuf tempSumBuf_; + TBuf postBrcbBuf_; + TBuf combBrcbBuf_; +}; + +template +__aicore__ inline void HcPostKernelDSplit::Init(GM_ADDR x, GM_ADDR residual, GM_ADDR post, GM_ADDR comb, GM_ADDR y, + GM_ADDR workspace, const HcPostTilingData *tilingData, TPipe *pipe) +{ + blkIdx_ = GetBlockIdx(); + if (blkIdx_ >= tilingData->usedCoreNum) { + return; + } + tiling_ = tilingData; + pipe_ = pipe; + hcParam_ = tilingData->hcParam; + dParam_ = tilingData->dParam; + batchOneCoreTail_ = tilingData->batchOneCoreTail; + batchOneCore_ = tilingData->batchOneCore; + dOnceDealing_ = tilingData->dOnceDealing; + dLastDealing_ = tilingData->dLastDealing; + dSplitTime_ = tilingData->dSplitTime; + isFrontCore_ = blkIdx_ < tilingData->frontCore; + int64_t frontCore = tilingData->frontCore; + hcParamAlign_ = RoundUp(hcParam_); + + int64_t xOffset = blkIdx_ * batchOneCore_ * dParam_; + int64_t residualOffset = blkIdx_ * batchOneCore_ * hcParam_ * dParam_; + int64_t postOffset = blkIdx_ * batchOneCore_ * hcParam_; + int64_t combOffset = blkIdx_ * batchOneCore_ * hcParam_ * hcParam_; + int64_t yOffset = blkIdx_ * batchOneCore_ * hcParam_ * dParam_; + if (!isFrontCore_) { + xOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * dParam_; + residualOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_ * dParam_; + postOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_; + combOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_ * hcParam_; + yOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_ * dParam_; + } + xGm_.SetGlobalBuffer((__gm__ T1 *)x + xOffset); + residualGm_.SetGlobalBuffer((__gm__ T1 *)residual + residualOffset); + postGm_.SetGlobalBuffer((__gm__ T2 *)post + postOffset); + combGm_.SetGlobalBuffer((__gm__ T2 *)comb + combOffset); + yGm_.SetGlobalBuffer((__gm__ T1 *)y + yOffset); + + pipe_->InitBuffer(postQue_, 2, hcParam_ * sizeof(T2)); + pipe_->InitBuffer(combQue_, 2, hcParam_ * hcParamAlign_ * sizeof(T2)); + pipe_->InitBuffer(outQue_, 2, hcParam_ * dOnceDealing_ * sizeof(T1)); + pipe_->InitBuffer(inputQue_, 2, hcParam_ * dOnceDealing_ * sizeof(T1)); + + if constexpr (sizeof(T1) == 2) { + pipe_->InitBuffer(outCastBuf_, hcParam_ * dOnceDealing_ * sizeof(float)); + pipe_->InitBuffer(inputCastBuf_, hcParam_ * dOnceDealing_ * sizeof(float)); + } + if constexpr (sizeof(T2) == 2) { + pipe_->InitBuffer(postCastBuf_, hcParam_ * sizeof(float)); + pipe_->InitBuffer(combCastBuf_, hcParam_ * RoundUp(hcParam_) * sizeof(float)); + } + pipe_->InitBuffer(postBrcbBuf_, 64 * sizeof(float)); + pipe_->InitBuffer(combBrcbBuf_, 64 * sizeof(float)); + pipe_->InitBuffer(tempSumBuf_, hcParam_ * dOnceDealing_ * sizeof(float)); +} + +template +__aicore__ inline void HcPostKernelDSplit::Process() +{ + if (blkIdx_ >= tiling_->usedCoreNum) { + return; + } + if (isFrontCore_) { + DoProcess(tiling_->batchOneCore); + } else { + DoProcess(tiling_->batchOneCoreTail); + } +} + +template +__aicore__ inline void HcPostKernelDSplit::DataCopyInX(int64_t batchIndex, int64_t dLoopTimes, int64_t dealNum) +{ + LocalTensor xUb = inputQue_.AllocTensor(); + DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = dealNum * sizeof(T1); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(xUb, xGm_[batchIndex * dParam_ + dLoopTimes * dOnceDealing_], copyParams, dataCopyPadParams); + inputQue_.EnQue(xUb); +} + +template +__aicore__ inline void HcPostKernelDSplit::DataCopyInPost(int64_t batchIndex) +{ + LocalTensor postUb = postQue_.AllocTensor(); + DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = hcParam_ * sizeof(T2); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(postUb, postGm_[batchIndex * hcParam_], copyParams, dataCopyPadParams); + postQue_.EnQue(postUb); +} + +template +__aicore__ inline void HcPostKernelDSplit::DataCopyInResidual(int64_t batchIndex, int64_t dLoopTimes, int64_t dealNum) +{ + LocalTensor residualUb = inputQue_.AllocTensor(); + DataCopyExtParams copyParams; + copyParams.blockCount = hcParam_; + copyParams.blockLen = dealNum * sizeof(T1); + copyParams.srcStride = (dParam_ - dealNum) * sizeof(T1); + copyParams.dstStride = 0; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(residualUb, residualGm_[batchIndex * hcParam_ * dParam_ + dLoopTimes * dOnceDealing_], copyParams, dataCopyPadParams); + inputQue_.EnQue(residualUb); +} + +template +__aicore__ inline void HcPostKernelDSplit::DataCopyInComb(int64_t batchIndex) +{ + uint8_t padNum = BLOCK_SIZE / sizeof(T2) - hcParam_; + LocalTensor combUb = combQue_.AllocTensor(); + DataCopyExtParams copyParams; + copyParams.blockCount = hcParam_; + copyParams.blockLen = hcParam_ * sizeof(T2); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + DataCopyPadExtParams dataCopyPadParams{true, 0, padNum, 0}; + DataCopyPad(combUb, combGm_[batchIndex * hcParam_ * hcParam_], copyParams, dataCopyPadParams); + combQue_.EnQue(combUb); +} + +template +__aicore__ inline void HcPostKernelDSplit::DataCopyOut(int64_t batchIndex, int64_t dLoopTimes, int64_t dealNum) +{ + LocalTensor outBuf = outQue_.DeQue(); + DataCopyExtParams copyParams; + copyParams.blockCount = hcParam_; + copyParams.blockLen = dealNum * sizeof(T1); + copyParams.srcStride = 0; + copyParams.dstStride = (dParam_ - dealNum) * sizeof(T1); + AscendC::DataCopyPad(yGm_[batchIndex * hcParam_ * dParam_ + dLoopTimes * dOnceDealing_], outBuf, copyParams); + outQue_.FreeTensor(outBuf); +} + +template +__aicore__ inline void DoBrcb(LocalTensor srcLocal, LocalTensor dstLocal, TBuf castLocalBuf, int64_t dealNum) +{ + uint32_t repeatTimes = CeilDiv(dealNum, REPEAT_NUM); + if constexpr (sizeof(T) == 2) { + LocalTensor castBuf = castLocalBuf.Get(); + Cast(castBuf, srcLocal, RoundMode::CAST_NONE, dealNum); + PipeBarrier(); + Brcb(dstLocal, castBuf, repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE}); + } else { + Brcb(dstLocal, srcLocal, repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE}); + } + PipeBarrier(); +} + +template +__aicore__ inline void DoMal(LocalTensor src0Local, LocalTensor src1Local, LocalTensor dstLocal, int64_t curRowNum, int64_t curColNum) +{ + int64_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + int64_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + int64_t curColNumAlign = RoundUp(curColNum); + int64_t numRepeatPerLine = curColNum / elemInOneRepeat; + BinaryRepeatParams instrParams; + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Mul(dstLocal[i*curColNumAlign], src0Local[0], src1Local[i*elemInOneBlock], elemInOneRepeat, numRepeatPerLine, instrParams); + } + PipeBarrier(); +} + +template +__aicore__ inline void DoAdd(LocalTensor src0Local, LocalTensor src1Local, LocalTensor dstLocal, int64_t curRowNum, int64_t curColNum) +{ + int64_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + int64_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + int64_t curColNumAlign = RoundUp(curColNum); + int64_t numRepeatPerLine = curColNum / elemInOneRepeat; + BinaryRepeatParams instrParams; + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = DEFAULT_REPEAT_STRIDE; + for (uint32_t i = 0; i < curRowNum; i++) { + Add(dstLocal[i*curColNumAlign], src0Local[i*curColNumAlign], src1Local[i*curColNumAlign], elemInOneRepeat, numRepeatPerLine, instrParams); + } + PipeBarrier(); +} + +template +__aicore__ inline void HcPostKernelDSplit::DoCompute(LocalTensor sumTempBuf, LocalTensor postBrcb, LocalTensor comBrcb, LocalTensor combUb, LocalTensor combCastBuf, int64_t batchIndex, int64_t dLoop, int64_t dealNum) +{ + LocalTensor outBuf; + if constexpr (sizeof(T1) == 2) { + outBuf = outCastBuf_.Get(); + } else { + outBuf = outQue_.AllocTensor(); + } + DataCopyInX(batchIndex, dLoop, dealNum); + LocalTensor xUb = inputQue_.DeQue(); + LocalTensor inputCastBuf; + if constexpr (sizeof(T1) == 2) { + inputCastBuf = inputCastBuf_.Get(); + Cast(inputCastBuf, xUb, RoundMode::CAST_NONE, dealNum); + PipeBarrier(); + DoMal(inputCastBuf, postBrcb, outBuf, hcParam_, dealNum); + } else { + DoMal(xUb, postBrcb, outBuf, hcParam_, dealNum); + } + inputQue_.FreeTensor(xUb); + + DataCopyInResidual(batchIndex, dLoop, dealNum); + LocalTensor residualUb = inputQue_.DeQue(); + if constexpr (sizeof(T1) == 2) { + for (int32_t i = 0; i < hcParam_; i++) { + Cast(inputCastBuf[i*RoundUp(dealNum)], residualUb[i*RoundUp(dealNum)], RoundMode::CAST_NONE, RoundUp(dealNum)); + } + PipeBarrier(); + } + + for (int64_t hcIndex = 0; hcIndex < hcParam_; hcIndex++) { + uint32_t repeatTimes = CeilDiv(hcParam_, REPEAT_NUM); + if constexpr (sizeof(T2) == 2) { + Brcb(comBrcb, combCastBuf[hcIndex*RoundUp(hcParam_)], repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE}); + } else { + Brcb(comBrcb, combUb[hcIndex*RoundUp(hcParam_)], repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE}); + } + PipeBarrier(); + if constexpr (sizeof(T1) == 2) { + DoMal(inputCastBuf[hcIndex*RoundUp(dealNum)], comBrcb, sumTempBuf, hcParam_, dealNum); + } else { + DoMal(residualUb[hcIndex*RoundUp(dealNum)], comBrcb, sumTempBuf, hcParam_, dealNum); + } + DoAdd(outBuf, sumTempBuf, outBuf, hcParam_, dealNum); + } + inputQue_.FreeTensor(residualUb); + if constexpr (sizeof(T1) == 2) { + LocalTensor outSumBuf = outQue_.AllocTensor(); + uint32_t outAlign = RoundUp(dealNum); + uint32_t inputAlign = RoundUp(dealNum); + for (int32_t i = 0; i < hcParam_; i++) { + Cast(outSumBuf[i*outAlign], outBuf[i*inputAlign], RoundMode::CAST_RINT, dealNum); + } + PipeBarrier(); + outQue_.EnQue(outSumBuf); + DataCopyOut(batchIndex, dLoop, dealNum); + } else { + outQue_.EnQue(outBuf); + DataCopyOut(batchIndex, dLoop, dealNum); + } +} + +template +__aicore__ inline void HcPostKernelDSplit::DoProcess(int64_t batchSize) +{ + LocalTensor sumTempBuf = tempSumBuf_.Get(); + for (int64_t batchIndex = 0; batchIndex < batchSize; batchIndex++) { + DataCopyInPost(batchIndex); + LocalTensor postUb = postQue_.DeQue(); + LocalTensor postBrcb = postBrcbBuf_.Get(); + DoBrcb(postUb, postBrcb, postCastBuf_, hcParam_); + LocalTensor comBrcb = combBrcbBuf_.Get(); + DataCopyInComb(batchIndex); + LocalTensor combUb = combQue_.DeQue(); + LocalTensor combCastBuf; + if constexpr (sizeof(T2) == 2) { + combCastBuf = combCastBuf_.Get(); + for (int32_t i = 0; i < hcParam_; i++) { + Cast(combCastBuf[i*RoundUp(hcParam_)], combUb[i*hcParamAlign_], RoundMode::CAST_NONE, hcParam_); + } + } + for (int64_t dLoop = 0; dLoop < dSplitTime_; dLoop++) { + DoCompute(sumTempBuf, postBrcb, comBrcb, combUb, combCastBuf, batchIndex, dLoop, dOnceDealing_); + PipeBarrier(); + } + if (dLastDealing_ != 0) { + DoCompute(sumTempBuf, postBrcb, comBrcb, combUb, combCastBuf, batchIndex, dSplitTime_, dLastDealing_); + PipeBarrier(); + } + combQue_.FreeTensor(combUb); + postQue_.FreeTensor(postUb); + } +} + +} +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/hc_post/op_kernel/hc_post_float32.h b/csrc/ascend/moe/hc_post/op_kernel/hc_post_float32.h new file mode 100644 index 000000000..e2515bbb0 --- /dev/null +++ b/csrc/ascend/moe/hc_post/op_kernel/hc_post_float32.h @@ -0,0 +1,322 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_post_float32.h + * \brief + */ +#ifndef HC_POST_FLOAT32_H +#define HC_POST_FLOAT32_H + +#include "kernel_operator.h" + +namespace HcPostRegBase { +using namespace AscendC; + +template +class HcPostRegBaseFloat32 { +public: + __aicore__ inline HcPostRegBaseFloat32() {}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR residual, GM_ADDR post, GM_ADDR comb, GM_ADDR y, GM_ADDR workspace, + const HcPostTilingData *tilingData, TPipe *pipe); + __aicore__ inline void Process(); + __aicore__ inline void DataCopyInX(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset); + __aicore__ inline void DataCopyInPost(int64_t batchIndex); + __aicore__ inline void DataCopyInResidual(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset); + __aicore__ inline void DataCopyInComb(int64_t batchIndex); + __aicore__ inline void DataCopyOut(int64_t batchIndex, int64_t hcIndex, int64_t dOnceDealing, int64_t dOffset); + __aicore__ inline void DoProcess(int64_t batchSize); + __aicore__ inline void DoCompute(LocalTensor sumTempBuf, LocalTensor postUb, LocalTensor combUb, int64_t batchIndex, int64_t dOffset, int64_t dDealing); + __aicore__ inline void DoMulAndAdd(LocalTensor xUb, LocalTensor postUb, LocalTensor residualUb, LocalTensor combUb, LocalTensor sumTempBuf, int64_t hcIndex, int64_t dDealing); + +private: + TPipe* pipe_; + const HcPostTilingData* tiling_; + constexpr static AscendC::MicroAPI::CastTrait castB16ToB32 = { AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::UNKNOWN, AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::UNKNOWN }; + + int32_t blkIdx_ = -1; + int64_t batch_ = 0; + int64_t hcParam_ = 0; + int64_t dParam_ = 0; + int64_t batchOneCoreTail_ = 0; + int64_t batchOneCore_ = 0; + int64_t isFrontCore_ = 0; + int64_t dParamAlign_ = 0; + int64_t dParamOnceAlign_ = 0; + int64_t dOnceDealing_ = 0; + int64_t dLastDealing_ = 0; + int64_t dSplitTime_ = 0; + static constexpr int32_t ONE_BLOCK_SIZE = 32; + int32_t perBlock32 = ONE_BLOCK_SIZE / sizeof(float); + + GlobalTensor xGm_; + GlobalTensor residualGm_; + GlobalTensor postGm_; + GlobalTensor combGm_; + GlobalTensor yGm_; + + TQue xQue_; + TQue residualQue_; + TQue postQue_; + TQue combQue_; + TQue sumQue_; + TBuf sumTempBuf_; +}; + +template +__aicore__ inline void HcPostRegBaseFloat32::Init(GM_ADDR x, GM_ADDR residual, GM_ADDR post, GM_ADDR comb, GM_ADDR y, + GM_ADDR workspace, const HcPostTilingData *tilingData, TPipe *pipe) +{ + blkIdx_ = GetBlockIdx(); + if (blkIdx_ >= tilingData->usedCoreNum) { + return; + } + tiling_ = tilingData; + pipe_ = pipe; + hcParam_ = tilingData->hcParam; + dParam_ = tilingData->dParam; + batchOneCoreTail_ = tilingData->batchOneCoreTail; + batchOneCore_ = tilingData->batchOneCore; + isFrontCore_ = blkIdx_ < tilingData->frontCore; + int64_t frontCore = tilingData->frontCore; + dOnceDealing_ = tilingData->dOnceDealing; + dLastDealing_ = tilingData->dLastDealing; + dSplitTime_ = tilingData->dSplitTime; + dParamAlign_ = (dParam_ + perBlock32 - 1) / perBlock32 * perBlock32; + dParamOnceAlign_ = (dOnceDealing_ + perBlock32 - 1) / perBlock32 * perBlock32; + + int64_t xOffset = blkIdx_ * batchOneCore_ * dParam_; + int64_t residualOffset = blkIdx_ * batchOneCore_ * hcParam_ * dParam_; + int64_t postOffset = blkIdx_ * batchOneCore_ * hcParam_; + int64_t combOffset = blkIdx_ * batchOneCore_ * hcParam_ * hcParam_; + int64_t yOffset = blkIdx_ * batchOneCore_ * hcParam_ * dParam_; + if (!isFrontCore_) { + xOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * dParam_; + residualOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_ * dParam_; + postOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_; + combOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_ * hcParam_; + yOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_ * dParam_; + } + xGm_.SetGlobalBuffer((__gm__ float *)x + xOffset); + residualGm_.SetGlobalBuffer((__gm__ float *)residual + residualOffset); + postGm_.SetGlobalBuffer((__gm__ T *)post + postOffset); + combGm_.SetGlobalBuffer((__gm__ T *)comb + combOffset); + yGm_.SetGlobalBuffer((__gm__ float *)y + yOffset); + + pipe_->InitBuffer(xQue_, 2, dParamOnceAlign_ * sizeof(float)); + pipe_->InitBuffer(residualQue_, 2, hcParam_ * dParamOnceAlign_ * sizeof(float)); + pipe_->InitBuffer(postQue_, 2, hcParam_ * sizeof(T)); + pipe_->InitBuffer(combQue_, 2, hcParam_ * hcParam_ * sizeof(T)); + pipe_->InitBuffer(sumQue_, 2, dParamOnceAlign_ * sizeof(float)); + pipe_->InitBuffer(sumTempBuf_, dParamOnceAlign_ * sizeof(float)); +} + +template +__aicore__ inline void HcPostRegBaseFloat32::Process() +{ + if (blkIdx_ >= tiling_->usedCoreNum) { + return; + } + if (isFrontCore_) { + DoProcess(tiling_->batchOneCore); + } else { + DoProcess(tiling_->batchOneCoreTail); + } +} + +template +__aicore__ inline void HcPostRegBaseFloat32::DataCopyInX(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset) +{ + LocalTensor xUb = xQue_.AllocTensor(); + DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = dOnceDealing * sizeof(float); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(xUb, xGm_[batchIndex * dParam_ + dOffset], copyParams, dataCopyPadParams); + xQue_.EnQue(xUb); +} + +template +__aicore__ inline void HcPostRegBaseFloat32::DataCopyInPost(int64_t batchIndex) +{ + LocalTensor postUb = postQue_.AllocTensor(); + DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = hcParam_ * sizeof(T); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(postUb, postGm_[batchIndex * hcParam_], copyParams, dataCopyPadParams); + postQue_.EnQue(postUb); +} + +template +__aicore__ inline void HcPostRegBaseFloat32::DataCopyInResidual(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset) +{ + LocalTensor residualUb = residualQue_.AllocTensor(); + DataCopyExtParams copyParams; + copyParams.blockCount = hcParam_; + copyParams.blockLen = dOnceDealing * sizeof(float); + copyParams.srcStride = (dParamAlign_ - dOnceDealing) * sizeof(float); + copyParams.dstStride = 0; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(residualUb, residualGm_[batchIndex * hcParam_ * dParam_ + dOffset], copyParams, dataCopyPadParams); + residualQue_.EnQue(residualUb); +} + +template +__aicore__ inline void HcPostRegBaseFloat32::DataCopyInComb(int64_t batchIndex) +{ + LocalTensor combUb = combQue_.AllocTensor(); + DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = hcParam_ * hcParam_ * sizeof(T); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(combUb, combGm_[batchIndex * hcParam_ * hcParam_], copyParams, dataCopyPadParams); + combQue_.EnQue(combUb); +} + +template +__aicore__ inline void HcPostRegBaseFloat32::DataCopyOut(int64_t batchIndex, int64_t hcIndex, int64_t dOnceDealing, int64_t dOffset) +{ + LocalTensor outBuf = sumQue_.DeQue(); + DataCopyExtParams copyParams; + copyParams.blockCount = 1; + copyParams.blockLen = dOnceDealing * sizeof(float); + copyParams.srcStride = 0; + copyParams.dstStride = 0; + AscendC::DataCopyPad(yGm_[batchIndex * hcParam_ * dParam_ + hcIndex * dParam_ + dOffset], outBuf, copyParams); + sumQue_.FreeTensor(outBuf); +} + +template +__aicore__ inline void HcPostRegBaseFloat32::DoMulAndAdd(LocalTensor xUb, LocalTensor postUb, LocalTensor residualUb, LocalTensor combUb, LocalTensor sumTempBuf, int64_t hcIndex, int64_t dOnceDealing) +{ + uint16_t aTimes = hcParam_; + uint32_t xDealNumAlign = (dOnceDealing + perBlock32 - 1) / perBlock32 * perBlock32; + uint32_t vfLen = 256 / sizeof(float); + uint16_t repeatTimes = (dOnceDealing + vfLen - 1) / vfLen; + + auto residualAddr = (__ubuf__ float*)residualUb.GetPhyAddr(); + auto combAddr = (__ubuf__ T*)combUb.GetPhyAddr(); + auto sumAddr = (__ubuf__ float*)sumTempBuf.GetPhyAddr(); + auto xAddr = (__ubuf__ float*)xUb.GetPhyAddr(); + auto postAddr = (__ubuf__ T*)postUb.GetPhyAddr(); + __VEC_SCOPE__ + { + uint32_t xDealNum = static_cast(dOnceDealing); + AscendC::MicroAPI::RegTensor combReg0; + AscendC::MicroAPI::RegTensor combReg1; + AscendC::MicroAPI::RegTensor combReg2; + AscendC::MicroAPI::RegTensor combReg3; + AscendC::MicroAPI::RegTensor residualRegFloat0; + AscendC::MicroAPI::RegTensor residualRegFloat1; + AscendC::MicroAPI::RegTensor residualRegFloat2; + AscendC::MicroAPI::RegTensor residualRegFloat3; + AscendC::MicroAPI::RegTensor combRegFloat0; + AscendC::MicroAPI::RegTensor combRegFloat1; + AscendC::MicroAPI::RegTensor combRegFloat2; + AscendC::MicroAPI::RegTensor combRegFloat3; + AscendC::MicroAPI::RegTensor sumRegFloat; + AscendC::MicroAPI::RegTensor sumTempReg0; + AscendC::MicroAPI::RegTensor sumTempReg1; + AscendC::MicroAPI::RegTensor sumTempReg2; + AscendC::MicroAPI::RegTensor sumTempReg3; + + AscendC::MicroAPI::RegTensor xReg; + AscendC::MicroAPI::RegTensor postReg; + AscendC::MicroAPI::RegTensor xRegFloat; + AscendC::MicroAPI::RegTensor postRegFloat; + AscendC::MicroAPI::MaskReg pMask; + AscendC::MicroAPI::MaskReg pregMain = AscendC::MicroAPI::CreateMask(); + if constexpr (sizeof(T) == 2) { + AscendC::MicroAPI::DataCopy(combReg0, combAddr+hcIndex); + AscendC::MicroAPI::DataCopy(combReg1, combAddr+hcParam_+hcIndex); + AscendC::MicroAPI::DataCopy(combReg2, combAddr+2*hcParam_+hcIndex); + AscendC::MicroAPI::DataCopy(combReg3, combAddr+3*hcParam_+hcIndex); + AscendC::MicroAPI::Cast(combRegFloat0, combReg0, pregMain); + AscendC::MicroAPI::Cast(combRegFloat1, combReg1, pregMain); + AscendC::MicroAPI::Cast(combRegFloat2, combReg2, pregMain); + AscendC::MicroAPI::Cast(combRegFloat3, combReg3, pregMain); + AscendC::MicroAPI::DataCopy(postReg, postAddr + hcIndex); + AscendC::MicroAPI::Cast(postRegFloat, postReg, pregMain); + } else { + AscendC::MicroAPI::DataCopy(combRegFloat0, combAddr+hcIndex); + AscendC::MicroAPI::DataCopy(combRegFloat1, combAddr+hcParam_+hcIndex); + AscendC::MicroAPI::DataCopy(combRegFloat2, combAddr+2*hcParam_+hcIndex); + AscendC::MicroAPI::DataCopy(combRegFloat3, combAddr+3*hcParam_+hcIndex); + AscendC::MicroAPI::DataCopy(postRegFloat, postAddr + hcIndex); + } + + for (uint16_t j = 0; j < repeatTimes; j++) { + pMask = AscendC::MicroAPI::UpdateMask(xDealNum); + AscendC::MicroAPI::DataCopy(xRegFloat, xAddr+j*vfLen); + AscendC::MicroAPI::DataCopy(residualRegFloat0, residualAddr+j*vfLen); + AscendC::MicroAPI::DataCopy(residualRegFloat1, residualAddr+xDealNumAlign+j*vfLen); + AscendC::MicroAPI::DataCopy(residualRegFloat2, residualAddr+2*xDealNumAlign+j*vfLen); + AscendC::MicroAPI::DataCopy(residualRegFloat3, residualAddr+3*xDealNumAlign+j*vfLen); + AscendC::MicroAPI::Mul(sumRegFloat, residualRegFloat0, combRegFloat0, pMask); + AscendC::MicroAPI::MulAddDst(sumRegFloat, residualRegFloat3, combRegFloat3, pMask); + AscendC::MicroAPI::MulAddDst(sumRegFloat, residualRegFloat1, combRegFloat1, pMask); + AscendC::MicroAPI::MulAddDst(sumRegFloat, residualRegFloat2, combRegFloat2, pMask); + AscendC::MicroAPI::MulAddDst(sumRegFloat, xRegFloat, postRegFloat, pMask); + AscendC::MicroAPI::DataCopy(sumAddr+j*vfLen, sumRegFloat, pMask); + } + } +} + +template +__aicore__ inline void HcPostRegBaseFloat32::DoCompute(LocalTensor sumTempBuf, LocalTensor postUb, LocalTensor combUb, int64_t batchIndex, int64_t dOffset, int64_t dDealing) +{ + DataCopyInX(batchIndex, dDealing, dOffset); + LocalTensor xUb = xQue_.DeQue(); + DataCopyInResidual(batchIndex, dDealing, dOffset); + LocalTensor residualUb = residualQue_.DeQue(); + for (int64_t hc1Index = 0; hc1Index < hcParam_; hc1Index++) { + DoMulAndAdd(xUb, postUb, residualUb, combUb, sumTempBuf, hc1Index, dDealing); + LocalTensor sumUb = sumQue_.AllocTensor(); + AscendC::Copy(sumUb, sumTempBuf, dDealing); + sumQue_.EnQue(sumUb); + DataCopyOut(batchIndex, hc1Index, dDealing, dOffset); + } + residualQue_.FreeTensor(residualUb); + xQue_.FreeTensor(xUb); +} + +template +__aicore__ inline void HcPostRegBaseFloat32::DoProcess(int64_t batchSize) +{ + LocalTensor sumTempBuf = sumTempBuf_.Get(); + for (int64_t batchIndex = 0; batchIndex < batchSize; batchIndex++) { + DataCopyInPost(batchIndex); + LocalTensor postUb = postQue_.DeQue(); + DataCopyInComb(batchIndex); + LocalTensor combUb = combQue_.DeQue(); + int64_t dOffset = 0; + for (int64_t dIndex = 0; dIndex < dSplitTime_; dIndex++) { + dOffset = dIndex*dOnceDealing_; + DoCompute(sumTempBuf, postUb, combUb, batchIndex, dOffset, dOnceDealing_); + } + if (dLastDealing_ != 0) { + dOffset = dSplitTime_*dOnceDealing_; + DoCompute(sumTempBuf, postUb, combUb, batchIndex, dOffset, dLastDealing_); + } + combQue_.FreeTensor(combUb); + postQue_.FreeTensor(postUb); + } +} + +} +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre/CMakeLists.txt b/csrc/ascend/moe/hc_pre/CMakeLists.txt new file mode 100644 index 000000000..86b308249 --- /dev/null +++ b/csrc/ascend/moe/hc_pre/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre/op_host/CMakeLists.txt b/csrc/ascend/moe/hc_pre/op_host/CMakeLists.txt new file mode 100644 index 000000000..4ac70b1ca --- /dev/null +++ b/csrc/ascend/moe/hc_pre/op_host/CMakeLists.txt @@ -0,0 +1,62 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ---------------------------------------------------------------------------- + +# add_ops_compile_options( +# OP_NAME HcPre +# OPTIONS --cce-auto-sync=off +# -Wno-deprecated-declarations +# -Werror +# -mllvm -cce-aicore-hoist-movemask=false +# --op_relocatable_kernel_binary=true +# ) + +# set(hc_pre_depends transformer/attention/hc_pre PARENT_SCOPE) + +# target_sources(op_host_aclnn PRIVATE +# op_host/hc_pre_def.cpp +# ) + +# target_sources(optiling PRIVATE +# op_host/hc_pre_tiling.cpp +# ) + +# if (NOT BUILD_OPEN_PROJECT) +# target_sources(opmaster_ct PRIVATE +# op_host/hc_pre_tiling.cpp +# ) +# endif () + +# target_include_directories(optiling PRIVATE +# ${CMAKE_CURRENT_SOURCE_DIR}/op_host +# ) + +# target_sources(opsproto PRIVATE +# op_host/hc_pre_proto.cpp +# ) + +if (BUILD_OPEN_PROJECT) + message(STATUS "=== Debug: BUILD is BUILD_OPEN_PROJECT") + target_sources(op_host_aclnn PRIVATE + hc_pre_def.cpp + ) + add_ops_compile_options( + OP_NAME HcPre + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true + ) +endif() + +if(NOT BUILD_OPS_RTY_KERNEL) + message(STATUS "=== Debug: BUILD is NOT BUILD_OPS_RTY_KERNEL") + add_modules_sources(OPTYPE hc_pre ACLNNTYPE aclnn) +endif() diff --git a/csrc/ascend/moe/hc_pre/op_host/hc_pre_def.cpp b/csrc/ascend/moe/hc_pre/op_host/hc_pre_def.cpp new file mode 100644 index 000000000..a6ad99dd2 --- /dev/null +++ b/csrc/ascend/moe/hc_pre/op_host/hc_pre_def.cpp @@ -0,0 +1,76 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_def.cpp + * \brief HcPre op host config + */ + +#include +#include "register/op_def_registry.h" + +namespace ops { +class HcPre : public OpDef { +public: + explicit HcPre(const char *name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("hc_fn") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("hc_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("hc_base") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("post") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("comb_frag") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + + this->Attr("hc_mult").AttrType(OPTIONAL).Int(4); + this->Attr("hc_sinkhorn_iters").AttrType(OPTIONAL).Int(20); + this->Attr("hc_eps").AttrType(OPTIONAL).Float(1e-6f); // default value + this->Attr("norm_eps").AttrType(OPTIONAL).Float(1e-6f); + + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + + OpAICoreConfig regbaseCfg; + regbaseCfg.DynamicCompileStaticFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .ExtendCfgInfo("opFile.value", "hc_pre"); + } +}; +OP_ADD(HcPre); +} // namespace ops diff --git a/csrc/ascend/moe/hc_pre/op_host/hc_pre_proto.cpp b/csrc/ascend/moe/hc_pre/op_host/hc_pre_proto.cpp new file mode 100644 index 000000000..e69de29bb diff --git a/csrc/ascend/moe/hc_pre/op_host/hc_pre_tiling.cpp b/csrc/ascend/moe/hc_pre/op_host/hc_pre_tiling.cpp new file mode 100644 index 000000000..e58ed6bc2 --- /dev/null +++ b/csrc/ascend/moe/hc_pre/op_host/hc_pre_tiling.cpp @@ -0,0 +1,393 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_tiling.cpp + * \brief + */ + +#include +#include "hc_pre_tiling.h" +#include "hc_pre_tiling_arch35.h" + +using namespace ge; +namespace optiling { +namespace { +constexpr uint64_t WORKSPACE_SIZE = 32; +int64_t CeilDiv(int64_t x, int64_t y) +{ + if (y != 0) { + return (x + y - 1) / y; + } + return x; +} +int64_t DownAlign(int64_t x, int64_t y) { + if (y == 0) { + return x; + } + return (x / y) * y; +} +int64_t RoundUp(int64_t x, int64_t y) { + return CeilDiv(x, y) * y; +} + +constexpr int64_t BLOCK_SIZE = 32; +constexpr int64_t REPEAT_SIZE = 256; +constexpr int64_t UB_RESEVED_SIZE = 8192; +constexpr int64_t DOUBLE_BUFFER = 2; +constexpr uint64_t M_L1_MAX_SIZE = 256; +constexpr uint64_t K_MULIT_CORE_SPLIT_BASE_SIZE = 256; +constexpr uint64_t A_L1_SIZE = 128 * 256; +constexpr uint64_t K_L1_MAX_SIZE = 1024; +constexpr int64_t K_L1_ALIGN_SIZE = 128; +constexpr int64_t HC_MULT_ATTR_IDX = 0; +constexpr int64_t ITER_TIMES_ATTR_IDX = 1; +constexpr int64_t HC_EPS_ATTR_IDX = 2; +constexpr int64_t NORM_EPS_ATTR_IDX = 3; +constexpr int64_t DEFAULT_ITER_TIMES = 20; +} + +ge::graphStatus HcPreTiling::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + if (platformInfo == nullptr) { + auto compileInfoPtr = context_->GetCompileInfo(); + OPS_ERR_IF(compileInfoPtr == nullptr, OPS_LOG_E(context_, "compile info is null"), + return ge::GRAPH_FAILED); + aivCoreNum_ = compileInfoPtr->coreNum; + ubSize_ = compileInfoPtr->ubSize; + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + aivCoreNum_ = ascendcPlatform.GetCoreNumAiv(); + aicCoreNum_ = ascendcPlatform.GetCoreNumAic(); + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + ubSize_ = ubSizePlatForm; + socVersion_ = ascendcPlatform.GetSocVersion(); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreTiling::GetAttr() +{ + auto* attrs = context_->GetAttrs(); + OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED); + + auto hcMultAttr = attrs->GetAttrPointer(HC_MULT_ATTR_IDX); + hcMult_ = hcMultAttr == nullptr ? 4 : *hcMultAttr; + + auto iterTimesAttr = attrs->GetAttrPointer(ITER_TIMES_ATTR_IDX); + iterTimes_ = iterTimesAttr == nullptr ? DEFAULT_ITER_TIMES : *iterTimesAttr; + + auto hcEpsAttr = attrs->GetAttrPointer(HC_EPS_ATTR_IDX); + hcEps_ = hcEpsAttr == nullptr ? 1e-6 : *hcEpsAttr; + + auto normEpsAttr = attrs->GetAttrPointer(NORM_EPS_ATTR_IDX); + normEps_ = normEpsAttr == nullptr ? 1e-6 : *normEpsAttr; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreTiling::GetShapeAttrsInfoInner() +{ + // (b, s, hc_mult, d) or (bs, hc_mult, d) + auto xShape = context_->GetInputShape(0); + OPS_LOG_E_IF_NULL(context_, xShape, return ge::GRAPH_FAILED); + size_t xDimNum = xShape->GetStorageShape().GetDimNum(); + if (xDimNum == 3) { + bs_ = xShape->GetStorageShape().GetDim(0); + hcMult_ = xShape->GetStorageShape().GetDim(1); + d_ = xShape->GetStorageShape().GetDim(2); + } else if (xDimNum == 4) { + int64_t b = xShape->GetStorageShape().GetDim(0); + int64_t s = xShape->GetStorageShape().GetDim(1); + bs_ = b * s; + hcMult_ = xShape->GetStorageShape().GetDim(2); + d_ = xShape->GetStorageShape().GetDim(3); + } + + auto shapeHcFn = context_->GetInputShape(1); + hcMix_ = shapeHcFn->GetStorageShape().GetDim(0); + OPS_ERR_IF(shapeHcFn->GetStorageShape().GetDim(1) != d_ * hcMult_, + OPS_LOG_E(context_->GetNodeName(), + "HcFn dim 1 should be equal with d_ * hcMult_ %ld, but is %ld", + d_ * hcMult_, shapeHcFn->GetStorageShape().GetDim(1)), + return ge::GRAPH_FAILED); + + auto shapeHcScale = context_->GetInputShape(2); + int64_t scaleFirstDim = shapeHcScale->GetStorageShape().GetDim(0); + OPS_ERR_IF(scaleFirstDim != 3, + OPS_LOG_E(context_->GetNodeName(), + "hc_scale size should be equal with 3, but is %ld", scaleFirstDim), + return ge::GRAPH_FAILED); + + auto shapeHcBase = context_->GetInputShape(3); + int64_t baseFirstDim = shapeHcBase->GetStorageShape().GetDim(0); + OPS_ERR_IF(baseFirstDim != hcMix_, + OPS_LOG_E(context_->GetNodeName(), + "hc_base size should be equal with mixhc, but is %ld", baseFirstDim), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(GetAttr() != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "get attr failed."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + + +ge::graphStatus HcPreTiling::CalcMKSplitCoreMembasePart2Tiling() +{ + rowOfFormerBlock_ = CeilDiv(bs_, static_cast(aivCoreNum_)); + usedAivCoreNums_ = std::min(CeilDiv(bs_, rowOfFormerBlock_), static_cast(aivCoreNum_)); + rowOfTailBlock_ = bs_ - (usedAivCoreNums_ - 1) * rowOfFormerBlock_; + + int64_t minRowPerCore = 1; + int64_t rowOnceLoop = std::min(rowOfFormerBlock_, minRowPerCore); + int64_t kBlockNum = tilingData_.get_cubeBlockDimK(); + + hcMultAlign_ = RoundUp(hcMult_, BLOCK_SIZE / sizeof(float)); + int64_t mix0OriginSize = kBlockNum * rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t mix1OriginSize = kBlockNum * rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t mix2OriginSize = kBlockNum * rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t mix0Size = rowOnceLoop * hcMultAlign_ * sizeof(float); + int64_t mix1Size = rowOnceLoop * hcMultAlign_ * sizeof(float); + int64_t mix2Size = rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float); + int64_t squareSumSize = kBlockNum * RoundUp(rowOnceLoop * 16, 16) * sizeof(float) * DOUBLE_BUFFER; + int64_t rsqrtSize = RoundUp(rowOnceLoop, BLOCK_SIZE / sizeof(float)) * sizeof(float) * DOUBLE_BUFFER; + int64_t xSize = rowOnceLoop * hcMult_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + int64_t ySize = rowOnceLoop * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; + int64_t postSize = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t combFragSize = rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t base0Size = hcMultAlign_ * sizeof(float); + int64_t base1Size = hcMultAlign_ * sizeof(float); + int64_t base2Size = hcMult_ * hcMultAlign_ * sizeof(float); + int64_t xCastSize = rowOnceLoop * hcMult_ * RoundUp(d_, 8) * sizeof(float); + int64_t yCastSize = rowOnceLoop * RoundUp(d_, 8) * sizeof(float); + int64_t rowBrcb0Size = RoundUp(rowOnceLoop, 8) * BLOCK_SIZE; + int64_t hcBrcb1Size = RoundUp(rowOnceLoop * hcMultAlign_, 8) * BLOCK_SIZE; + int64_t reduceBufSize = rowOnceLoop * hcMultAlign_ * sizeof(float); + int64_t maskPatternSize = BLOCK_SIZE * 16; + + int64_t totalSize = mix0OriginSize + mix1OriginSize + mix2OriginSize + + mix0Size + mix1Size + mix2Size + squareSumSize + rsqrtSize + + xSize + ySize + postSize + combFragSize + base0Size + base1Size + + base2Size + xCastSize + yCastSize + + rowBrcb0Size + hcBrcb1Size + reduceBufSize + maskPatternSize; + rowFactor_ = rowOnceLoop; + if (totalSize <= ubSize_) { + // row和d均可以在ub内全载 + dLoop_ = 1; + dFactor_ = d_; + tailDFactor_ = dFactor_; + } else { + int64_t usedUbSize = mix0OriginSize + mix1OriginSize + mix2OriginSize + + mix0Size + mix1Size + mix2Size + squareSumSize + rsqrtSize + + postSize + combFragSize + base0Size + base1Size + base2Size + + rowBrcb0Size + hcBrcb1Size + reduceBufSize + maskPatternSize; + int64_t ubRemain = ubSize_ - usedUbSize; + dFactor_ = d_; + int64_t base = 2; + while (1) { + dFactor_ = CeilDiv(d_, base); + xSize = rowOnceLoop * hcMult_ * RoundUp(dFactor_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + ySize = rowOnceLoop * RoundUp(dFactor_, 16) * 2 * DOUBLE_BUFFER; + xCastSize = rowOnceLoop * hcMult_ * RoundUp(dFactor_, 8) * sizeof(float); + yCastSize = rowOnceLoop * RoundUp(dFactor_, 8) * sizeof(float); + int64_t targetSize = xSize + ySize + xCastSize + yCastSize; + if (targetSize <= ubRemain) { + break; + } + base++; + } + if (dFactor_ > 32) { + dFactor_ = DownAlign(dFactor_, 32); + } + dLoop_ = CeilDiv(d_, dFactor_); + tailDFactor_ = d_ % dFactor_ == 0 ? dFactor_ : d_ % dFactor_; + } + + // d全载,尝试搬入更多的bs + if (dFactor_ == d_) { + while (rowFactor_ <= rowOfFormerBlock_) { + mix0OriginSize = kBlockNum * rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + mix1OriginSize = kBlockNum * rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + mix2OriginSize = kBlockNum * rowFactor_ * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + mix0Size = rowFactor_ * hcMultAlign_ * sizeof(float); + mix1Size = rowFactor_ * hcMultAlign_ * sizeof(float); + mix2Size = rowFactor_ * hcMult_ * hcMultAlign_ * sizeof(float); + squareSumSize = kBlockNum * RoundUp(rowFactor_ * 16, 16) * sizeof(float) * DOUBLE_BUFFER; + rsqrtSize = RoundUp(rowFactor_, BLOCK_SIZE / sizeof(float)) * sizeof(float) * DOUBLE_BUFFER; + xSize = rowFactor_ * hcMult_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + ySize = rowFactor_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; + postSize = rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + combFragSize = rowFactor_ * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + xCastSize = rowFactor_ * hcMult_ * RoundUp(d_, 8) * sizeof(float); + yCastSize = rowFactor_ * RoundUp(d_, 8) * sizeof(float); + rowBrcb0Size = RoundUp(rowFactor_, 8) * BLOCK_SIZE; + hcBrcb1Size = RoundUp(rowFactor_ * hcMultAlign_, 8) * BLOCK_SIZE; + reduceBufSize = rowFactor_ * hcMultAlign_ * sizeof(float); + maskPatternSize = BLOCK_SIZE; + totalSize = mix0OriginSize + mix1OriginSize + mix2OriginSize + + mix0Size + mix1Size + mix2Size + squareSumSize + rsqrtSize + + xSize + ySize + postSize + combFragSize + base0Size + base1Size + + base2Size + xCastSize + yCastSize + + rowBrcb0Size + hcBrcb1Size + reduceBufSize + maskPatternSize; + if (totalSize > ubSize_) { + rowFactor_ = rowFactor_ - 1; + break; + } + rowFactor_ = rowFactor_ + 1; + } + rowFactor_ = rowFactor_ > rowOfFormerBlock_ ? rowFactor_ - 1 : rowFactor_; + } + rowLoopOfFormerBlock_ = CeilDiv(rowOfFormerBlock_, rowFactor_); + rowLoopOfTailBlock_ = CeilDiv(rowOfTailBlock_, rowFactor_); + tailRowFactorOfFormerBlock_ = rowOfFormerBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfFormerBlock_ % rowFactor_; + tailRowFactorOfTailBlock_ = rowOfTailBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfTailBlock_ % rowFactor_; + + tilingData_.set_bs(bs_); + tilingData_.set_hcMix(hcMix_); + tilingData_.set_hcMult(hcMult_); + tilingData_.set_d(d_); + tilingData_.set_hcMultAlign(hcMultAlign_); + tilingData_.set_rowOfFormerBlock(rowOfFormerBlock_); + tilingData_.set_rowOfTailBlock(rowOfTailBlock_); + tilingData_.set_rowLoopOfFormerBlock(rowLoopOfFormerBlock_); + tilingData_.set_rowLoopOfTailBlock(rowLoopOfTailBlock_); + tilingData_.set_stage2RowFactor(rowFactor_); + tilingData_.set_secondUsedCoreNum(usedAivCoreNums_); + tilingData_.set_tailRowFactorOfFormerBlock(tailRowFactorOfFormerBlock_); + tilingData_.set_tailRowFactorOfTailBlock(tailRowFactorOfTailBlock_); + tilingData_.set_dLoop(dLoop_); + tilingData_.set_dFactor(dFactor_); + tilingData_.set_tailDFactor(tailDFactor_); + tilingData_.set_iterTimes(iterTimes_); + tilingData_.set_hcEps(hcEps_); + tilingData_.set_normEps(normEps_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreTiling::CalcOpTiling() { + uint64_t kSize = hcMult_ * d_; + tilingData_.set_k(kSize); + // 计算bs_轴切核 + uint64_t mDimNum = std::min(aicCoreNum_, static_cast(CeilDiv(bs_, M_L1_MAX_SIZE))); + uint64_t singleCoreM = RoundUp(CeilDiv(bs_, mDimNum), AscendC::BLOCK_CUBE); + uint64_t kDimNum = aicCoreNum_ / mDimNum; + uint64_t splitKSize = RoundUp(CeilDiv(kSize, kDimNum), K_MULIT_CORE_SPLIT_BASE_SIZE); + + tilingData_.set_cubeBlockDimM(mDimNum); + tilingData_.set_cubeBlockDimK(CeilDiv(kSize, splitKSize)); + tilingData_.set_multCoreSplitMSize(singleCoreM); + tilingData_.set_mL1Size(std::min(M_L1_MAX_SIZE, singleCoreM)); + tilingData_.set_multCoreSplitKSize(splitKSize); + tilingData_.set_kL1Size(std::min(A_L1_SIZE / tilingData_.get_mL1Size(), + static_cast(K_L1_MAX_SIZE)) / K_L1_ALIGN_SIZE * K_L1_ALIGN_SIZE); + + tilingData_.set_cvLoopKSize(1024); + + // vector stage1 tiling + tilingData_.set_cubeCoreNum(static_cast(aicCoreNum_)); + // x type bfloat16, y type float32 and double + // exit node 1 b16 input Queue and 1 b32 output Queue + int64_t lineByteSize = (sizeof(int16_t) + sizeof(int32_t)) * DOUBLE_BUFFER * tilingData_.get_cvLoopKSize(); + int64_t stage1MFactorValue = ubSize_ / lineByteSize; + tilingData_.set_stage1MFactor(stage1MFactorValue); + return CalcMKSplitCoreMembasePart2Tiling(); +} + + +ge::graphStatus HcPreTiling::DoOpTiling() +{ + if (GetPlatformInfo() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + if (GetShapeAttrsInfoInner() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + if (CalcOpTiling() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + if (GetWorkspaceSize() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + if (PostTiling() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreTiling::GetWorkspaceSize() +{ + uint64_t xCastFp32BufSize = tilingData_.get_mL1Size() * RoundUp(tilingData_.get_cvLoopKSize(), 128); + uint64_t workspaceSize1 = tilingData_.get_cubeCoreNum() * DOUBLE_BUFFER * xCastFp32BufSize * sizeof(float); + + uint64_t mmLastAxisSize = RoundUp(tilingData_.get_hcMix(), 128); + uint64_t workspaceSize2 = RoundUp(tilingData_.get_cubeBlockDimK() * + tilingData_.get_bs() * mmLastAxisSize * sizeof(float), 512); + + uint64_t squareSumSize = RoundUp(tilingData_.get_cubeBlockDimK() * + RoundUp(tilingData_.get_bs(), 16) * 16 * sizeof(float), 512); + + uint64_t requiredSize = workspaceSize1 + workspaceSize2 + squareSumSize + 16 * 1024 * 1024; // 16MB 预留缓冲 + + uint64_t defaultSize = 16 * 1024 * 1024 + 192 * 1024 * 1024; // 208MB + workspaceSize_ = requiredSize > defaultSize ? requiredSize : defaultSize; + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreTiling::PostTiling() +{ + context_->SetTilingKey(0); + context_->SetBlockDim(aicCoreNum_); + size_t* workspaces = context_->GetWorkspaceSizes(1); + workspaces[0] = workspaceSize_; + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingPrepareForHcPre(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingForHcPre(gert::TilingContext *context) +{ + OPS_ERR_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("HcPre", "Tiling context is null"), + return ge::GRAPH_FAILED); + + auto platformInfo = context->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, + OPS_REPORT_VECTOR_INNER_ERR("TilingForMoeGatingTopKHash", + "Tiling platformInfo is null"), return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + auto socVersion = ascendcPlatform.GetSocVersion(); + if (socVersion == platform_ascendc::SocVersion::ASCEND950) { + OPS_LOG_I(context, "Using arch35 tiling for ASCEND950"); + HcPreTilingRegbase::HcPreTilingRegbase hcPreTiling(context); + return hcPreTiling.DoOpTiling(); + } + HcPreTiling hcPreTiling(context); + return hcPreTiling.DoOpTiling(); +} + +IMPL_OP_OPTILING(HcPre) + .Tiling(TilingForHcPre) + .TilingParse(TilingPrepareForHcPre); + +} // namespace optiling diff --git a/csrc/ascend/moe/hc_pre/op_host/hc_pre_tiling.h b/csrc/ascend/moe/hc_pre/op_host/hc_pre_tiling.h new file mode 100644 index 000000000..83781f5b5 --- /dev/null +++ b/csrc/ascend/moe/hc_pre/op_host/hc_pre_tiling.h @@ -0,0 +1,161 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_tiling.h + * \brief + */ + +#ifndef HC_PRE_SINKHORN_TILING_H +#define HC_PRE_SINKHORN_TILING_H + + +#include +#include +#include "register/op_impl_registry.h" +#include "platform/platform_infos_def.h" +#include "exe_graph/runtime/tiling_context.h" +#include "tiling/platform/platform_ascendc.h" +#include "register/op_def_registry.h" +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "error/ops_error.h" +#include "platform/platform_info.h" + +namespace optiling { +// ----------公共定义---------- +struct TilingRequiredParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; +}; + +struct TilingOptionalParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::Tensor *tensor; +}; +// ----------算子TilingData定义---------- +BEGIN_TILING_DATA_DEF(HcPreTilingData) +TILING_DATA_FIELD_DEF(int64_t, bs); +TILING_DATA_FIELD_DEF(int64_t, hcMix); +TILING_DATA_FIELD_DEF(int64_t, hcMult); +TILING_DATA_FIELD_DEF(int64_t, d); +TILING_DATA_FIELD_DEF(int64_t, hcMultAlign); +TILING_DATA_FIELD_DEF(int64_t, rowOfFormerBlock); +TILING_DATA_FIELD_DEF(int64_t, rowOfTailBlock); +TILING_DATA_FIELD_DEF(int64_t, rowLoopOfFormerBlock); +TILING_DATA_FIELD_DEF(int64_t, rowLoopOfTailBlock); +TILING_DATA_FIELD_DEF(int64_t, rowFactor); +TILING_DATA_FIELD_DEF(int64_t, tailRowFactorOfFormerBlock); +TILING_DATA_FIELD_DEF(int64_t, tailRowFactorOfTailBlock); +TILING_DATA_FIELD_DEF(int64_t, dLoop); +TILING_DATA_FIELD_DEF(int64_t, dFactor); +TILING_DATA_FIELD_DEF(int64_t, tailDFactor); +TILING_DATA_FIELD_DEF(int64_t, iterTimes); +TILING_DATA_FIELD_DEF(float, hcEps); +TILING_DATA_FIELD_DEF(float, normEps); +TILING_DATA_FIELD_DEF(int64_t, kBlockFactor); +TILING_DATA_FIELD_DEF(int64_t, kFactor); +TILING_DATA_FIELD_DEF(int64_t, tailKFactor); +TILING_DATA_FIELD_DEF(int64_t, kLoop); +TILING_DATA_FIELD_DEF(int64_t, stage2RowFactor); + +TILING_DATA_FIELD_DEF(int64_t, k); +TILING_DATA_FIELD_DEF(int64_t, kLoopOfFormerBlock); +TILING_DATA_FIELD_DEF(int64_t, kLoopOfTailBlock); +TILING_DATA_FIELD_DEF(int64_t, stage1KFactor); +TILING_DATA_FIELD_DEF(int64_t, kFactorOfFormerBlock); +TILING_DATA_FIELD_DEF(int64_t, kL1Size); +TILING_DATA_FIELD_DEF(int64_t, mL1Size); +TILING_DATA_FIELD_DEF(int64_t, cubeBlockDimM); +TILING_DATA_FIELD_DEF(int64_t, cubeBlockDimK); +TILING_DATA_FIELD_DEF(int64_t, kUbSize); // ub一次处理k的size +TILING_DATA_FIELD_DEF(int64_t, cvLoopKSize); // cv同步使用的k轴长度 +TILING_DATA_FIELD_DEF(int64_t, multCoreSplitKSize); // 多核切K后的单核k处理长度 +TILING_DATA_FIELD_DEF(int64_t, multCoreSplitMSize); // 多核切K后的单核k处理长度 +TILING_DATA_FIELD_DEF(int64_t, tailKSizeOfFormerBlock); +TILING_DATA_FIELD_DEF(int64_t, tailKSizeOfTailBlock); + +TILING_DATA_FIELD_DEF(int64_t, mLoopOfFormerBlock); +TILING_DATA_FIELD_DEF(int64_t, mLoopOfTailBlock); +TILING_DATA_FIELD_DEF(int64_t, formerMSize); +TILING_DATA_FIELD_DEF(int64_t, tailMSizeOfFormerBlock); +TILING_DATA_FIELD_DEF(int64_t, tailMSizeOfTailBlock); + +TILING_DATA_FIELD_DEF(int64_t, firstUsedCoreNum); +TILING_DATA_FIELD_DEF(int64_t, secondUsedCoreNum); + +TILING_DATA_FIELD_DEF(int64_t, rowInnerFactor); + +TILING_DATA_FIELD_DEF(int64_t, cubeCoreNum); +TILING_DATA_FIELD_DEF(int64_t, stage1MFactor); + +TILING_DATA_FIELD_DEF(int64_t, bufferPool0Size); +TILING_DATA_FIELD_DEF(int64_t, bufferPool1Size); +TILING_DATA_FIELD_DEF(int64_t, mUbSize); + +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(HcPre, HcPreTilingData) + +// ----------算子CompileInfo定义---------- +struct HcPreCompileInfo { + uint64_t coreNum = 0; + uint64_t ubSize = 0; +}; + +// ----------算子Tiling入参信息解析及check类---------- +class HcPreTiling { +public: + explicit HcPreTiling(gert::TilingContext* tilingContext) : context_(tilingContext) + { + } + ~HcPreTiling() = default; + + ge::graphStatus GetPlatformInfo(); + ge::graphStatus DoOpTiling(); + ge::graphStatus GetWorkspaceSize(); + ge::graphStatus PostTiling(); + ge::graphStatus GetAttr(); + ge::graphStatus GetShapeAttrsInfoInner(); + ge::graphStatus CalcOpTiling(); + ge::graphStatus CalcMKSplitCoreMembasePart2Tiling(); + +private: + gert::TilingContext *context_ = nullptr; + uint64_t tilingKey_ = 0; + HcPreTilingData tilingData_; + uint64_t aivCoreNum_ = 0; + uint64_t aicCoreNum_ = 0; + uint64_t workspaceSize_ = 0; + uint64_t usedAivCoreNums_ = 0; + uint64_t ubSize_ = 0; + int64_t bs_ = 0; + int64_t hcMix_ = 0; + int64_t hcMult_ = 0; + int64_t d_ = 0; + int64_t hcMultAlign_ = 0; + int64_t rowOfFormerBlock_ = 0; + int64_t rowOfTailBlock_ = 0; + int64_t rowLoopOfFormerBlock_ = 0; + int64_t rowLoopOfTailBlock_ = 0; + int64_t rowFactor_ = 0; + int64_t tailRowFactorOfFormerBlock_ = 0; + int64_t tailRowFactorOfTailBlock_= 0; + int64_t dLoop_ = 0; + int64_t dFactor_ = 0; + int64_t tailDFactor_ = 0; + int64_t iterTimes_ = 0; + double hcEps_ = 0.0; + double normEps_ = 0.0; + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B; +}; + +} // namespace optiling +#endif // HC_PRE_SINKHORN_TILING_H \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre/op_host/hc_pre_tiling_arch35.h b/csrc/ascend/moe/hc_pre/op_host/hc_pre_tiling_arch35.h new file mode 100644 index 000000000..4b57cb7b0 --- /dev/null +++ b/csrc/ascend/moe/hc_pre/op_host/hc_pre_tiling_arch35.h @@ -0,0 +1,477 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_tiling_arch35.cpp + * \brief + */ + +#include +#include "hc_pre_tiling.h" + +using namespace ge; +namespace optiling { +namespace HcPreTilingRegbase { +namespace { +constexpr uint64_t WORKSPACE_SIZE = 32; +int64_t CeilDiv(int64_t x, int64_t y) +{ + if (y != 0) { + return (x + y - 1) / y; + } + return x; +} +int64_t DownAlign(int64_t x, int64_t y) { + if (y == 0) { + return x; + } + return (x / y) * y; +} +int64_t RoundUp(int64_t x, int64_t y) { + return CeilDiv(x, y) * y; +} + +constexpr int64_t BLOCK_SIZE = 32; +constexpr int64_t REPEAT_SIZE = 256; +constexpr int64_t DOUBLE_BUFFER = 2; +constexpr uint64_t M_L1_MAX_SIZE = 256; +constexpr uint64_t K_MULIT_CORE_SPLIT_BASE_SIZE = 256; +constexpr uint64_t A_L1_SIZE = 128 * 256; +constexpr uint64_t K_L1_MAX_SIZE = 1024; +} + +class HcPreTilingRegbase { +public: + explicit HcPreTilingRegbase(gert::TilingContext* tilingContext) : context_(tilingContext) + { + } + ~HcPreTilingRegbase() = default; + + ge::graphStatus GetPlatformInfo(); + ge::graphStatus DoOpTiling(); + ge::graphStatus GetWorkspaceSize(); + ge::graphStatus PostTiling(); + ge::graphStatus GetAttr(); + ge::graphStatus GetShapeAttrsInfoInner(); + ge::graphStatus CalcOpTiling(); + ge::graphStatus CalcRegbaseOpTiling(); + ge::graphStatus CalcMKSplitCorePart2Tiling(); +private: + gert::TilingContext *context_ = nullptr; + uint64_t tilingKey_ = 0; + HcPreTilingData tilingData_; + uint64_t aivCoreNum_ = 0; + uint64_t aicCoreNum_ = 0; + uint64_t workspaceSize_ = 0; + uint64_t usedCoreNums_ = 0; + uint64_t usedAivCoreNums_ = 0; + uint64_t ubSize_ = 0; + int64_t bs_ = 0; + int64_t hcMix_ = 0; + int64_t hcMult_ = 0; + int64_t d_ = 0; + int64_t hcMultAlign_ = 0; + int64_t rowOfFormerBlock_ = 0; + int64_t rowOfTailBlock_ = 0; + int64_t rowLoopOfFormerBlock_ = 0; + int64_t rowLoopOfTailBlock_ = 0; + int64_t rowFactor_ = 0; + int64_t tailRowFactorOfFormerBlock_ = 0; + int64_t tailRowFactorOfTailBlock_= 0; + int64_t dLoop_ = 0; + int64_t dFactor_ = 0; + int64_t tailDFactor_ = 0; + int64_t iterTimes_ = 0; + double hcEps_ = 0.0; + double normEps_ = 0.0; + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B; +}; + +ge::graphStatus HcPreTilingRegbase::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + if (platformInfo == nullptr) { + auto compileInfoPtr = context_->GetCompileInfo(); + OPS_ERR_IF(compileInfoPtr == nullptr, OPS_LOG_E(context_, "compile info is null"), + return ge::GRAPH_FAILED); + aivCoreNum_ = compileInfoPtr->coreNum; + ubSize_ = compileInfoPtr->ubSize; + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + aivCoreNum_ = ascendcPlatform.GetCoreNumAiv(); + aicCoreNum_ = ascendcPlatform.GetCoreNumAic(); + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + ubSize_ = ubSizePlatForm; + socVersion_ = ascendcPlatform.GetSocVersion(); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreTilingRegbase::GetAttr() +{ + auto* attrs = context_->GetAttrs(); + OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED); + + auto hcMultAttr = attrs->GetAttrPointer(0); + hcMult_ = hcMultAttr == nullptr ? 4 : *hcMultAttr; + + auto iterTimesAttr = attrs->GetAttrPointer(1); + iterTimes_ = iterTimesAttr == nullptr ? 20 : *iterTimesAttr; + + auto epsAttr = attrs->GetAttrPointer(2); + hcEps_ = epsAttr == nullptr ? 1e-6 : *epsAttr; + + auto normEpsAttr = attrs->GetAttrPointer(3); + normEps_ = normEpsAttr == nullptr ? 1e-6 : *normEpsAttr; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreTilingRegbase::GetShapeAttrsInfoInner() +{ + // (b, s, hc_mult, d) or (bs, hc_mult, d) + auto xShape = context_->GetInputShape(0); + OPS_LOG_E_IF_NULL(context_, xShape, return ge::GRAPH_FAILED); + size_t xDimNum = xShape->GetStorageShape().GetDimNum(); + if (xDimNum == 3) { + bs_ = xShape->GetStorageShape().GetDim(0); + hcMult_ = xShape->GetStorageShape().GetDim(1); + d_ = xShape->GetStorageShape().GetDim(2); + } else if (xDimNum == 4) { + int64_t b = xShape->GetStorageShape().GetDim(0); + int64_t s = xShape->GetStorageShape().GetDim(1); + bs_ = b * s; + hcMult_ = xShape->GetStorageShape().GetDim(2); + d_ = xShape->GetStorageShape().GetDim(3); + } + + auto shapeHcFn = context_->GetInputShape(1); + hcMix_ = shapeHcFn->GetStorageShape().GetDim(0); + OPS_ERR_IF(shapeHcFn->GetStorageShape().GetDim(1) != d_ * hcMult_, + OPS_LOG_E(context_->GetNodeName(), + "HcFn dim 1 should be equal with d_ * hcMult_ %ld, but is %ld", d_ * hcMult_, shapeHcFn->GetStorageShape().GetDim(1)), + return ge::GRAPH_FAILED); + + auto shapeHcScale = context_->GetInputShape(2); + int64_t scaleFirstDim = shapeHcScale->GetStorageShape().GetDim(0); + OPS_ERR_IF(scaleFirstDim != 3, + OPS_LOG_E(context_->GetNodeName(), + "hc_scale size should be equal with 3, but is %ld", scaleFirstDim), + return ge::GRAPH_FAILED); + + auto shapeHcBase = context_->GetInputShape(3); + int64_t baseFirstDim = shapeHcBase->GetStorageShape().GetDim(0); + OPS_ERR_IF(baseFirstDim != hcMix_, + OPS_LOG_E(context_->GetNodeName(), + "hc_base size should be equal with mixhc, but is %ld", baseFirstDim), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(GetAttr() != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "get attr failed."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + + +ge::graphStatus HcPreTilingRegbase::CalcRegbaseOpTiling() +{ + rowOfFormerBlock_ = CeilDiv(bs_, static_cast(aivCoreNum_)); + usedCoreNums_ = std::min(CeilDiv(bs_, rowOfFormerBlock_), static_cast(aivCoreNum_)); + rowOfTailBlock_ = bs_ - (usedCoreNums_ - 1) * rowOfFormerBlock_; + + int64_t minRowPerCore = 1; + int64_t rowOnceLoop = std::min(rowOfFormerBlock_, minRowPerCore); + + hcMultAlign_ = RoundUp(hcMult_, BLOCK_SIZE / sizeof(float)); + int64_t mixSize = rowOnceLoop * RoundUp(hcMix_, BLOCK_SIZE / sizeof(float)) * sizeof(float); + int64_t xSize = rowOnceLoop * hcMult_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + int64_t ySize = rowOnceLoop * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; + int64_t postSize = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + // comb is packed in UB as hcMult * hcMult, so size it by the actual footprint. + int64_t combFragSize = rowOnceLoop * hcMult_ * hcMult_ * sizeof(float) * DOUBLE_BUFFER; + int64_t base0Size = hcMultAlign_ * sizeof(float); + int64_t base1Size = hcMultAlign_ * sizeof(float); + int64_t base2Size = hcMult_ * hcMultAlign_ * sizeof(float); + + uint64_t kUbSize = tilingData_.get_kL1Size() / 2; // 先按2倍系数计算，m最大256，需保证kub小于256 + uint64_t mUbSize = CeilDiv(tilingData_.get_mL1Size(), 2); + + int64_t mmXBufSize = mUbSize * RoundUp(hcMix_, BLOCK_SIZE / sizeof(float)) * sizeof(float); + int64_t rmsNormBufSize = RoundUp(mUbSize, BLOCK_SIZE / sizeof(float)) * sizeof(float); + int64_t bufferPool0Size = ubSize_; + int64_t bufferPool1Size = DownAlign(bufferPool0Size - mmXBufSize - rmsNormBufSize - base0Size - base1Size - base2Size, BLOCK_SIZE); + + int64_t totalSize = mixSize + xSize + ySize + postSize + combFragSize; + rowFactor_ = rowOnceLoop; + if (totalSize <= bufferPool1Size) { + // row和d均可以在ub内全载 + dLoop_ = 1; + dFactor_ = d_; + tailDFactor_ = dFactor_; + } else { + int64_t usedUbSize = mixSize + postSize + combFragSize; + int64_t ubRemain = bufferPool1Size - usedUbSize; + dFactor_ = d_; + int64_t base = 2; + while (1) { + dFactor_ = CeilDiv(d_, base); + xSize = rowOnceLoop * hcMult_ * RoundUp(dFactor_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + ySize = rowOnceLoop * RoundUp(dFactor_, 16) * 2 * DOUBLE_BUFFER; + int64_t targetSize = xSize + ySize; + if (targetSize <= ubRemain) { + break; + } + base++; + } + if (dFactor_ > 32) { + dFactor_ = DownAlign(dFactor_, 32); + } + dLoop_ = CeilDiv(d_, dFactor_); + tailDFactor_ = d_ % dFactor_ == 0 ? dFactor_ : d_ % dFactor_; + } + + // d全载,尝试搬入更多的bs + if (dFactor_ == d_) { + while (rowFactor_ <= mUbSize) { + mixSize = rowFactor_ * RoundUp(hcMix_, BLOCK_SIZE / sizeof(float)) * sizeof(float); + xSize = rowFactor_ * hcMult_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + ySize = rowFactor_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; + postSize = rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + combFragSize = rowFactor_ * hcMult_ * hcMult_ * sizeof(float) * DOUBLE_BUFFER; + totalSize = mixSize + xSize + ySize + postSize + combFragSize; + if (totalSize > bufferPool1Size) { + rowFactor_ = rowFactor_ - 1; + break; + } + rowFactor_ = rowFactor_ + 1; + } + rowFactor_ = rowFactor_ > mUbSize ? rowFactor_ - 1 : rowFactor_; + } + + rowLoopOfFormerBlock_ = CeilDiv(rowOfFormerBlock_, rowFactor_); + rowLoopOfTailBlock_ = CeilDiv(rowOfTailBlock_, rowFactor_); + tailRowFactorOfFormerBlock_ = rowOfFormerBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfFormerBlock_ % rowFactor_; + tailRowFactorOfTailBlock_ = rowOfTailBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfTailBlock_ % rowFactor_; + + tilingData_.set_bs(bs_); + tilingData_.set_hcMix(hcMix_); + tilingData_.set_hcMult(hcMult_); + tilingData_.set_d(d_); + tilingData_.set_hcMultAlign(hcMultAlign_); + tilingData_.set_rowOfFormerBlock(rowOfFormerBlock_); + tilingData_.set_rowOfTailBlock(rowOfTailBlock_); + tilingData_.set_rowLoopOfFormerBlock(rowLoopOfFormerBlock_); + tilingData_.set_rowLoopOfTailBlock(rowLoopOfTailBlock_); + tilingData_.set_rowFactor(rowFactor_); + tilingData_.set_tailRowFactorOfFormerBlock(tailRowFactorOfFormerBlock_); + tilingData_.set_tailRowFactorOfTailBlock(tailRowFactorOfTailBlock_); + tilingData_.set_dLoop(dLoop_); + tilingData_.set_dFactor(dFactor_); + tilingData_.set_tailDFactor(tailDFactor_); + tilingData_.set_iterTimes(iterTimes_); + tilingData_.set_hcEps(hcEps_); + tilingData_.set_normEps(normEps_); + + tilingData_.set_bufferPool0Size(bufferPool0Size); + tilingData_.set_bufferPool1Size(bufferPool1Size); + + tilingData_.set_kUbSize(kUbSize); + tilingData_.set_mUbSize(mUbSize); + + tilingData_.set_kBlockFactor(tilingData_.get_cubeBlockDimK()); + + tilingData_.set_rowInnerFactor(rowFactor_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreTilingRegbase::CalcMKSplitCorePart2Tiling() +{ + uint64_t kUbSize = tilingData_.get_kL1Size() / 2; // 先按2倍系数计算，m最大256，需保证kub小于256 + uint64_t mUbSize = CeilDiv(tilingData_.get_mL1Size(), 2); + + rowOfFormerBlock_ = CeilDiv(bs_, static_cast(aivCoreNum_)); + usedAivCoreNums_ = std::min(CeilDiv(bs_, rowOfFormerBlock_), static_cast(aivCoreNum_)); + rowOfTailBlock_ = bs_ - (usedAivCoreNums_ - 1) * rowOfFormerBlock_; + + int64_t minRowPerCore = 1; + int64_t rowOnceLoop = std::min(rowOfFormerBlock_, minRowPerCore); + int64_t kBlockNum = tilingData_.get_cubeBlockDimK(); + + hcMultAlign_ = RoundUp(hcMult_, BLOCK_SIZE / sizeof(float)); + uint64_t hcMixAlign = RoundUp(hcMix_, BLOCK_SIZE / sizeof(float)); + int64_t mmSize = kBlockNum * rowOnceLoop * hcMixAlign * sizeof(float) * DOUBLE_BUFFER; + int64_t mixSize = rowOnceLoop * hcMixAlign * sizeof(float); + int64_t rmsSize = kBlockNum * RoundUp(rowOnceLoop, BLOCK_SIZE / sizeof(float)) * sizeof(float) * DOUBLE_BUFFER; + int64_t xSize = rowOnceLoop * hcMult_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + int64_t ySize = rowOnceLoop * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; + int64_t postSize = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + // comb is packed in UB as hcMult * hcMult, so size it by the actual footprint. + int64_t combFragSize = rowOnceLoop * hcMult_ * hcMult_ * sizeof(float) * DOUBLE_BUFFER; + int64_t base0Size = hcMultAlign_ * sizeof(float); + int64_t base1Size = hcMultAlign_ * sizeof(float); + int64_t base2Size = hcMult_ * hcMultAlign_ * sizeof(float); + + int64_t totalSize = mmSize + mixSize + rmsSize + xSize + ySize + postSize + combFragSize + base0Size + base1Size + base2Size; + rowFactor_ = rowOnceLoop; + if (totalSize <= ubSize_) { + // row和d均可以在ub内全载 + dLoop_ = 1; + dFactor_ = d_; + tailDFactor_ = dFactor_; + } else { + int64_t usedUbSize = mmSize + mixSize + rmsSize + postSize + combFragSize + base0Size + base1Size + base2Size; + int64_t ubRemain = ubSize_ - usedUbSize; + dFactor_ = d_; + int64_t base = 2; + while (1) { + dFactor_ = CeilDiv(d_, base); + xSize = rowOnceLoop * hcMult_ * RoundUp(dFactor_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + ySize = rowOnceLoop * RoundUp(dFactor_, 16) * 2 * DOUBLE_BUFFER; + int64_t targetSize = xSize + ySize; + if (targetSize <= ubRemain) { + break; + } + base++; + } + if (dFactor_ > 32) { + dFactor_ = DownAlign(dFactor_, 32); + } + dLoop_ = CeilDiv(d_, dFactor_); + tailDFactor_ = d_ % dFactor_ == 0 ? dFactor_ : d_ % dFactor_; + } + + // d全载,尝试搬入更多的bs + if (dFactor_ == d_) { + while (rowFactor_ <= rowOfFormerBlock_) { + mmSize = kBlockNum * rowFactor_ * hcMixAlign * sizeof(float) * DOUBLE_BUFFER; + mixSize = rowFactor_ * hcMixAlign * sizeof(float); + rmsSize = kBlockNum * RoundUp(rowFactor_, BLOCK_SIZE / sizeof(float)) * sizeof(float) * DOUBLE_BUFFER; + xSize = rowFactor_ * hcMult_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + ySize = rowFactor_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; + postSize = rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + combFragSize = rowFactor_ * hcMult_ * hcMult_ * sizeof(float) * DOUBLE_BUFFER; + base0Size = hcMultAlign_ * sizeof(float); + base1Size = hcMultAlign_ * sizeof(float); + base2Size = hcMult_ * hcMultAlign_ * sizeof(float); + + totalSize = mmSize + mixSize + rmsSize + xSize + ySize + postSize + combFragSize + base0Size + base1Size + base2Size; + if (totalSize > ubSize_) { + rowFactor_ = rowFactor_ - 1; + break; + } + rowFactor_ = rowFactor_ + 1; + } + rowFactor_ = rowFactor_ > rowOfFormerBlock_ ? rowFactor_ - 1 : rowFactor_; + } + rowLoopOfFormerBlock_ = CeilDiv(rowOfFormerBlock_, rowFactor_); + rowLoopOfTailBlock_ = CeilDiv(rowOfTailBlock_, rowFactor_); + tailRowFactorOfFormerBlock_ = rowOfFormerBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfFormerBlock_ % rowFactor_; + tailRowFactorOfTailBlock_ = rowOfTailBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfTailBlock_ % rowFactor_; + + tilingData_.set_bs(bs_); + tilingData_.set_hcMix(hcMix_); + tilingData_.set_hcMult(hcMult_); + tilingData_.set_d(d_); + tilingData_.set_hcMultAlign(hcMultAlign_); + tilingData_.set_rowOfFormerBlock(rowOfFormerBlock_); + tilingData_.set_rowOfTailBlock(rowOfTailBlock_); + tilingData_.set_rowLoopOfFormerBlock(rowLoopOfFormerBlock_); + tilingData_.set_rowLoopOfTailBlock(rowLoopOfTailBlock_); + tilingData_.set_stage2RowFactor(rowFactor_); + tilingData_.set_secondUsedCoreNum(usedAivCoreNums_); + tilingData_.set_tailRowFactorOfFormerBlock(tailRowFactorOfFormerBlock_); + tilingData_.set_tailRowFactorOfTailBlock(tailRowFactorOfTailBlock_); + tilingData_.set_dLoop(dLoop_); + tilingData_.set_dFactor(dFactor_); + tilingData_.set_tailDFactor(tailDFactor_); + tilingData_.set_iterTimes(iterTimes_); + tilingData_.set_hcEps(hcEps_); + tilingData_.set_normEps(normEps_); + tilingData_.set_kUbSize(kUbSize); + tilingData_.set_mUbSize(mUbSize); + tilingData_.set_kBlockFactor(tilingData_.get_cubeBlockDimK()); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreTilingRegbase::CalcOpTiling() { + uint64_t kSize = hcMult_ * d_; + tilingData_.set_k(kSize); + // 计算bs_轴切核 + uint64_t mDimNum = std::min(aicCoreNum_, static_cast(CeilDiv(bs_, M_L1_MAX_SIZE))); + uint64_t singleCoreM = RoundUp(CeilDiv(bs_, mDimNum), AscendC::BLOCK_CUBE); + uint64_t kDimNum = aicCoreNum_ / mDimNum; + + uint64_t splitKSize = RoundUp(CeilDiv(kSize, kDimNum), K_MULIT_CORE_SPLIT_BASE_SIZE); + uint64_t actualKBlockNum = CeilDiv(kSize, splitKSize); + + tilingData_.set_cubeBlockDimM(mDimNum); + tilingData_.set_cubeBlockDimK(actualKBlockNum); + tilingData_.set_multCoreSplitMSize(singleCoreM); // todo: 这个 tiling 根本没有使用，是否需要删掉 + tilingData_.set_mL1Size(std::min(M_L1_MAX_SIZE, singleCoreM)); + tilingData_.set_multCoreSplitKSize(splitKSize); + tilingData_.set_kL1Size(std::min(A_L1_SIZE / tilingData_.get_mL1Size(), static_cast(K_L1_MAX_SIZE)) / 128 * 128); + + tilingData_.set_cvLoopKSize(1024); + if (kDimNum != 1) { + tilingKey_ = 1000; + return CalcMKSplitCorePart2Tiling(); + } + tilingKey_ = 1001; + return CalcRegbaseOpTiling(); +} + + +ge::graphStatus HcPreTilingRegbase::DoOpTiling() +{ + if (GetPlatformInfo() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + if (GetShapeAttrsInfoInner() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + if (CalcOpTiling() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + if (GetWorkspaceSize() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + if (PostTiling() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreTilingRegbase::GetWorkspaceSize() +{ + if (tilingKey_ == 1000) { + // K分核模板需要预留Workspace大小 + workspaceSize_ = tilingData_.get_kBlockFactor() * tilingData_.get_bs() * tilingData_.get_hcMix() * 4 + tilingData_.get_kBlockFactor() * tilingData_.get_bs() * 4 + 16 * 1024 * 1024; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreTilingRegbase::PostTiling() +{ + context_->SetTilingKey(tilingKey_); + context_->SetBlockDim(aicCoreNum_); + size_t* workspaces = context_->GetWorkspaceSizes(1); + workspaces[0] = workspaceSize_; + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} +} +} // namespace optiling diff --git a/csrc/ascend/moe/hc_pre/op_kernel/hc_pre.cpp b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre.cpp new file mode 100644 index 000000000..7ccc1986c --- /dev/null +++ b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre.cpp @@ -0,0 +1,89 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre.cpp + * \brief + */ + +#include "kernel_operator.h" +#include "kernel_operator_intf.h" + +#if defined(__DAV_C310__) + #include "hc_pre_m_k_split_core_arch35.h" + #include "hc_pre_m_split_core_arch35.h" + #include "hc_pre_base_arch35.h" + using namespace HcPreNs; +#else + #include "lib/matmul_intf.h" + #include "hc_pre_m_k_split_core.h" + #include "hc_pre_base.h" + using namespace HcPre; +#endif + +using namespace AscendC; + +extern "C" __global__ __aicore__ void hc_pre(GM_ADDR x, GM_ADDR hc_fn, GM_ADDR hc_scale, GM_ADDR hc_base, + GM_ADDR y, GM_ADDR post, GM_ADDR comb_frag, GM_ADDR workspace, + GM_ADDR tiling) +{ + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + if (workspace == nullptr) { + return; + } + + GM_ADDR userWs = GetUserWorkspace(workspace); + if (userWs == nullptr) { + return; + } + TPipe pipe; + + // 950PR 950DT + #if defined(__DAV_C310__) + if (TILING_KEY_IS(1000)) { + GET_TILING_DATA_WITH_STRUCT(HcPreTilingData, tiling_data_in, tiling); + const HcPreTilingData *__restrict tilingData = &tiling_data_in; + HcPreNs::HcPreMKSplitCorePart1 op; + op.Init(x, hc_fn, userWs, tilingData, &pipe); + op.Process(); + pipe.Destroy(); + + TPipe pipeStage2; + HcPreNs::HcPreMKSplitCorePart2 op2; + op2.Init(x, hc_scale, hc_base, y, post, comb_frag, userWs, tilingData, &pipeStage2); + op2.Process(); + pipeStage2.Destroy(); + } else if (TILING_KEY_IS(1001)) { + GET_TILING_DATA_WITH_STRUCT(HcPreTilingData, tiling_data_in, tiling); + const HcPreTilingData *__restrict tilingData = &tiling_data_in; + HcPreNs::HcPreMSplitCorePart1 op; + op.Init(x, hc_fn, hc_scale, hc_base, y, post, comb_frag, tilingData, &pipe); + op.Process(); + } + #else + // A3 + if (TILING_KEY_IS(0)) { + GET_TILING_DATA_WITH_STRUCT(HcPreTilingData, tiling_data_in, tiling); + const HcPreTilingData *__restrict tilingData = &tiling_data_in; + HcPre::HcPreMembaseKSplitCorePart1 op; + op.Init(x, hc_fn, userWs, tilingData, &pipe); + op.Process(); + + pipe.Destroy(); + + TPipe pipeStage2; + HcPre::HcPreMembaseKSplitCorePart2 op2; + op2.Init(x, hc_scale, hc_base, y, post, comb_frag, userWs, tilingData, &pipeStage2); + op2.Process(); + + pipeStage2.Destroy(); + } + #endif +} \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_base.h b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_base.h new file mode 100644 index 000000000..65f15c6af --- /dev/null +++ b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_base.h @@ -0,0 +1,688 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_base.h + * \brief + */ + +#ifndef HC_PRE_VECTOR_BASE_H +#define HC_PRE_VECTOR_BASE_H + +#include "kernel_operator.h" + +namespace HcPre { +using namespace AscendC; +constexpr int32_t BLOCK_SIZE = 32; +constexpr int32_t DEFAULT_BLOCK_STRIDE = 1; +constexpr int32_t DEFAULT_REPEAT_STRIDE = 8; +constexpr int32_t ONE_REPEAT_BLOCK_NUMS = 8; +constexpr int32_t REPEAT_SIZE = 256; +constexpr int32_t MAX_REPEAT_STRIDE = 255; +constexpr int32_t WORKSPACE_ALIGN_SIZE = 512; +constexpr uint32_t ONE = 1; +constexpr uint32_t SHIFT_COEFF = 17; +constexpr int32_t DOUBLE_BUFFER = 2; +constexpr int32_t CV_RATIO = 2; +constexpr uint64_t N_SIZE = 24; +constexpr uint64_t NUM_TWO = 2; +constexpr uint64_t SQUARE_SUM_SIZE = 16; +constexpr uint64_t MASK_PATTERN_BASE_SIZE = 16; +constexpr uint64_t MASK_PATTERN_REPEAT_SIZE = 8; +constexpr uint64_t MASK_PATTERN_DIM_SIZE = 16; +constexpr uint64_t MM_CACHE_LINE_BYTES = 512; + +__aicore__ inline int32_t CeilDiv(int32_t a, int32_t b) +{ + if (b == 0) { + return a; + } + return (a + b - 1) / b; +} + +__aicore__ inline int32_t CeilAlign(int32_t a, int32_t b) +{ + return CeilDiv(a, b) * b; +} + +template +__aicore__ inline int32_t RoundUp(int32_t num) +{ + int32_t elemNum = BLOCK_SIZE / sizeof(T); + return CeilAlign(num, elemNum); +} + +__aicore__ inline void SetGatherMaskPattern(const LocalTensor& maskPattern) +{ + uint32_t base = ONE; + for (uint32_t i = 0; i < 16; i++) { + Duplicate(maskPattern[i * 8], base, 8); + base = (base << 1); + } + PipeBarrier(); +} + +__aicore__ inline void GatherMaskByDiagonal(const LocalTensor& output, const LocalTensor& input, + const LocalTensor maskPattern, uint16_t dim0) +{ + uint32_t totalCount = MASK_PATTERN_DIM_SIZE * MASK_PATTERN_DIM_SIZE; + uint32_t remainCount = (dim0 % MASK_PATTERN_DIM_SIZE) * MASK_PATTERN_DIM_SIZE; + uint32_t loopCount = dim0 / MASK_PATTERN_DIM_SIZE; + uint64_t rsvdCnt = 0; + for (uint32_t loopIdx = 0; loopIdx < loopCount; loopIdx++) { + GatherMask(output[loopIdx * MASK_PATTERN_DIM_SIZE], + input[loopIdx * totalCount], maskPattern, true, totalCount, + {1, 1, 8, 8}, rsvdCnt); + } + GatherMask(output[loopCount * MASK_PATTERN_DIM_SIZE], + input[loopCount * totalCount], maskPattern, true, remainCount, + {1, 1, 8, 8}, rsvdCnt); + PipeBarrier(); +} + + +template +__aicore__ inline void MulABLastDimBrcInline(const LocalTensor &output, const LocalTensor &input0, + const LocalTensor &input1, const LocalTensor &tmpBuffer, + const int32_t curRowNum, const int32_t curColNum) +{ + if constexpr (needBrc) { + uint32_t repeatTimes = CeilDiv(curRowNum, ONE_REPEAT_BLOCK_NUMS); + Brcb(tmpBuffer, input1, repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE}); + } + PipeBarrier(); + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + uint32_t curColNumAlign = RoundUp(curColNum); + if (curColNum <= elemInOneBlock) { + Mul(output, input0, tmpBuffer, curRowNum * curColNumAlign); + } else { + int32_t numRepeatPerLine = curColNum / elemInOneRepeat; + int32_t numRemainPerLine = curColNum % elemInOneRepeat; + int32_t dstRepStridePerLine = CeilDiv(curColNum, elemInOneBlock); + BinaryRepeatParams instrParams; + + if (numRepeatPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE || curRowNum < numRepeatPerLine) { + // 在Col方向开Repeat, 并且Repeat小于255 + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Mul(output[i * curColNumAlign], input0[i * curColNumAlign], tmpBuffer[i * elemInOneBlock], + elemInOneRepeat, numRepeatPerLine, instrParams); + } + } else { + // 在Row方向开Repeat + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 1; + for (uint32_t i = 0; i < numRepeatPerLine; i++) { + Mul(output[i * elemInOneRepeat], input0[i * elemInOneRepeat], tmpBuffer, elemInOneRepeat, curRowNum, + instrParams); + } + } + } + + if (numRemainPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE) { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = 0; + instrParams.src0RepStride = 0; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Mul(output[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], + input0[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], tmpBuffer[i * elemInOneBlock], + numRemainPerLine, 1, instrParams); + } + } else { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 0; + Mul(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], tmpBuffer, + numRemainPerLine, curRowNum, instrParams); + } + } + } + PipeBarrier(); +} + +template +__aicore__ inline void SubABLastDimBrcInline(const LocalTensor &output, const LocalTensor &input0, + const LocalTensor &input1, const LocalTensor &tmpBuffer, + const int32_t curRowNum, const int32_t curColNum) +{ + if constexpr (needBrc) { + uint32_t repeatTimes = CeilDiv(curRowNum, ONE_REPEAT_BLOCK_NUMS); + Brcb(tmpBuffer, input1, repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE}); + } + PipeBarrier(); + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + uint32_t curColNumAlign = RoundUp(curColNum); + if (curColNum <= elemInOneBlock) { + Sub(output, input0, tmpBuffer, curRowNum * curColNumAlign); + } else { + int32_t numRepeatPerLine = curColNum / elemInOneRepeat; + int32_t numRemainPerLine = curColNum % elemInOneRepeat; + int32_t dstRepStridePerLine = CeilDiv(curColNum, elemInOneBlock); + BinaryRepeatParams instrParams; + + if (numRepeatPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE || curRowNum < numRepeatPerLine) { + // 在Col方向开Repeat, 并且Repeat小于255 + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Sub(output[i * curColNumAlign], input0[i * curColNumAlign], tmpBuffer[i * elemInOneBlock], + elemInOneRepeat, numRepeatPerLine, instrParams); + } + } else { + // 在Row方向开Repeat + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 1; + for (uint32_t i = 0; i < numRepeatPerLine; i++) { + Sub(output[i * elemInOneRepeat], input0[i * elemInOneRepeat], tmpBuffer, elemInOneRepeat, curRowNum, + instrParams); + } + } + } + + if (numRemainPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE) { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = 0; + instrParams.src0RepStride = 0; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Sub(output[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], + input0[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], tmpBuffer[i * elemInOneBlock], + numRemainPerLine, 1, instrParams); + } + } else { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 0; + Sub(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], tmpBuffer, + numRemainPerLine, curRowNum, instrParams); + } + } + } + PipeBarrier(); +} + + +template +__aicore__ inline void DivABLastDimBrcInline(const LocalTensor &output, const LocalTensor &input0, + const LocalTensor &input1, const LocalTensor &tmpBuffer, + const int32_t curRowNum, const int32_t curColNum) +{ + if constexpr (needBrc) { + uint32_t repeatTimes = CeilDiv(curRowNum, ONE_REPEAT_BLOCK_NUMS); + Brcb(tmpBuffer, input1, repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE}); + } + PipeBarrier(); + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + uint32_t curColNumAlign = RoundUp(curColNum); + if (curColNum <= elemInOneBlock) { + Div(output, input0, tmpBuffer, curRowNum * curColNumAlign); + } else { + int32_t numRepeatPerLine = curColNum / elemInOneRepeat; + int32_t numRemainPerLine = curColNum % elemInOneRepeat; + int32_t dstRepStridePerLine = CeilDiv(curColNum, elemInOneBlock); + BinaryRepeatParams instrParams; + + if (numRepeatPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE || curRowNum < numRepeatPerLine) { + // 在Col方向开Repeat, 并且Repeat小于255 + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Div(output[i * curColNumAlign], input0[i * curColNumAlign], tmpBuffer[i * elemInOneBlock], + elemInOneRepeat, numRepeatPerLine, instrParams); + } + } else { + // 在Row方向开Repeat + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 1; + for (uint32_t i = 0; i < numRepeatPerLine; i++) { + Div(output[i * elemInOneRepeat], input0[i * elemInOneRepeat], tmpBuffer, elemInOneRepeat, curRowNum, + instrParams); + } + } + } + + if (numRemainPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE) { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = 0; + instrParams.src0RepStride = 0; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Div(output[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], + input0[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], tmpBuffer[i * elemInOneBlock], + numRemainPerLine, 1, instrParams); + } + } else { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 0; + Div(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], tmpBuffer, + numRemainPerLine, curRowNum, instrParams); + } + } + } + PipeBarrier(); +} + + +template +__aicore__ inline void AddBAFirstDimBrcInline(const LocalTensor &output, const LocalTensor &input0, + const LocalTensor &input1, const int32_t curRowNum, + const int32_t curColNum) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + uint32_t curColNumAlign = RoundUp(curColNum); + int32_t numRepeatPerLine = curColNum / elemInOneRepeat; + int32_t numRemainPerLine = curColNum % elemInOneRepeat; + int32_t dstRepStridePerLine = CeilDiv(curColNum, elemInOneBlock); + BinaryRepeatParams instrParams; + if (numRepeatPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE || curRowNum < numRepeatPerLine) { + // 在Col方向开Repeat, 并且Repeat小于255 + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = DEFAULT_REPEAT_STRIDE; + for (uint32_t i = 0; i < curRowNum; i++) { + Add(output[i * curColNumAlign], input0[i * curColNumAlign], input1, elemInOneRepeat, numRepeatPerLine, + instrParams); + } + } else { + // 在Row方向开Repeat + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < numRepeatPerLine; i++) { + Add(output[i * elemInOneRepeat], input0[i * elemInOneRepeat], input1[i * elemInOneRepeat], + elemInOneRepeat, curRowNum, instrParams); + } + } + } + + if (numRemainPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE) { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = 0; + instrParams.src0RepStride = 0; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Add(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], input1, + numRemainPerLine, 1, instrParams); + } + } else { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 0; + Add(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], input1, + numRemainPerLine, curRowNum, instrParams); + } + } + PipeBarrier(); +} + +template +__aicore__ inline void CalcDenominator(const LocalTensor &output, const LocalTensor &input, + const uint32_t calCount) +{ + Muls(output, input, static_cast(-1.0), calCount); + PipeBarrier(); + Exp(output, output, calCount); + PipeBarrier(); + Adds(output, output, static_cast(1.0), calCount); + PipeBarrier(); +} + +// 暂时不处理repeat超限场景 +template +__aicore__ inline void SigmoidPerf(const LocalTensor &output, const LocalTensor &input, + const LocalTensor &tmpBuffer, const int64_t calCount) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + Duplicate(tmpBuffer, static_cast(1.0), elemInOneBlock); + CalcDenominator(output, input, calCount); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + int32_t numRepeatPerLine = calCount / elemInOneRepeat; + int32_t numRemainPerLine = calCount % elemInOneRepeat; + BinaryRepeatParams instrParams; + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 0; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = 0; + instrParams.src1RepStride = DEFAULT_REPEAT_STRIDE; + Div(output, tmpBuffer, output, elemInOneRepeat, numRepeatPerLine, instrParams); + if (numRemainPerLine != 0) { + Div(output[numRepeatPerLine * elemInOneRepeat], tmpBuffer, output[numRepeatPerLine * elemInOneRepeat], + numRemainPerLine, 1, instrParams); + } + PipeBarrier(); +} + +__aicore__ inline void ProcessPre(const LocalTensor &preLocal, const LocalTensor &mixLocal, + const LocalTensor &hcBaseLocal, const LocalTensor &rsqrtLocal, + const LocalTensor &tmpBuffer0, const LocalTensor &tmpBuffer1, + float scale, float eps, const int32_t curRowNum, const int32_t curColNum) +{ + int32_t curColNumAlign = RoundUp(curColNum); + MulABLastDimBrcInline(mixLocal, mixLocal, rsqrtLocal, tmpBuffer0, curRowNum, curColNum); + Muls(mixLocal, mixLocal, scale, curRowNum * curColNumAlign); + PipeBarrier(); + AddBAFirstDimBrcInline(mixLocal, mixLocal, hcBaseLocal, curRowNum, curColNum); + SigmoidPerf(preLocal, mixLocal, tmpBuffer1, curRowNum * curColNumAlign); + Adds(preLocal, preLocal, eps, curRowNum * curColNumAlign); + PipeBarrier(); +} + +__aicore__ inline void ReduceSumARAPerf(const LocalTensor &output, const LocalTensor &input, + const uint32_t dim0, const uint32_t dim1, const uint32_t dim2) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(float); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(float); + uint32_t dim2Align = RoundUp(dim2); + + // 拷贝第一个R到output上 + DataCopyParams copyParams; + copyParams.blockCount = dim0; + copyParams.blockLen = dim2Align / elemInOneBlock; + copyParams.srcStride = (dim1 - 1) * (dim2Align / elemInOneBlock); + copyParams.dstStride = 0; + DataCopy(output, input, copyParams); + PipeBarrier(); + uint32_t dim2RepeatTimes = dim2 / elemInOneRepeat; + uint32_t dim2Reminder = dim2 % elemInOneRepeat; + // 沿着dim2方向开repeat + BinaryRepeatParams instrParams; + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = DEFAULT_REPEAT_STRIDE; + for (uint32_t i = 0; i < dim0; i++) { + for (uint32_t j = 1; j < dim1; j++) { + Add(output[i * dim2Align], output[i * dim2Align], input[i * dim1 * dim2Align + j * dim2Align], + elemInOneRepeat, dim2RepeatTimes, instrParams); + if (dim2Reminder != 0) { + Add(output[i * dim2Align + dim2RepeatTimes * elemInOneRepeat], + output[i * dim2Align + dim2RepeatTimes * elemInOneRepeat], + input[i * dim1 * dim2Align + j * dim2Align + +dim2RepeatTimes * elemInOneRepeat], dim2Reminder, 1, + instrParams); + } + PipeBarrier(); + } + } + PipeBarrier(); +} + +template +__aicore__ inline void CastTwoDim(const LocalTensor &output, const LocalTensor &input, const uint32_t dim0, + const uint32_t dim1) +{ + uint32_t dim1AlignT0 = RoundUp(dim1); + uint32_t dim1AlignT1 = RoundUp(dim1); + if constexpr (IsSameType::value && IsSameType::value) { + for (uint32_t i = 0; i < dim0; i++) { + Cast(output[i * dim1AlignT0], input[i * dim1AlignT1], AscendC::RoundMode::CAST_NONE, dim1); + } + } else { + for (uint32_t i = 0; i < dim0; i++) { + Cast(output[i * dim1AlignT0], input[i * dim1AlignT1], AscendC::RoundMode::CAST_RINT, dim1); + } + } + PipeBarrier(); +} + + +template +__aicore__ void inline ProcessY(const LocalTensor &yLocal, const LocalTensor &xLocal, + const LocalTensor &mix01Local, const LocalTensor &hcBrcbLocal1, + const LocalTensor &xCastLocal, const LocalTensor &yCastLocal, + const uint32_t dim0, const uint32_t dim1, const uint32_t dim2) +{ + CastTwoDim(xCastLocal, xLocal, dim0 * dim1, dim2); + MulABLastDimBrcInline(xCastLocal, xCastLocal, mix01Local, hcBrcbLocal1, dim0 * dim1, dim2); + ReduceSumARAPerf(yCastLocal, xCastLocal, dim0, dim1, dim2); + CastTwoDim(yLocal, yCastLocal, dim0, dim2); +} + + +__aicore__ inline void ProcessPost(const LocalTensor &postLocal, const LocalTensor &mixLocal, + const LocalTensor &hcBaseLocal, const LocalTensor &rsqrtLocal, + const LocalTensor &tmpBuffer0, const LocalTensor &tmpBuffer1, + float scale, const int32_t curRowNum, const int32_t curColNum) +{ + int32_t curColNumAlign = RoundUp(curColNum); + MulABLastDimBrcInline(mixLocal, mixLocal, rsqrtLocal, tmpBuffer0, curRowNum, curColNum); + Muls(mixLocal, mixLocal, scale, curRowNum * curColNumAlign); + PipeBarrier(); + AddBAFirstDimBrcInline(mixLocal, mixLocal, hcBaseLocal, curRowNum, curColNum); + SigmoidPerf(postLocal, mixLocal, tmpBuffer1, curRowNum * curColNumAlign); + Muls(postLocal, postLocal, static_cast(2.0f), curRowNum * curColNumAlign); + PipeBarrier(); +} + +__aicore__ inline void LastDimReduceMaxPerf(const LocalTensor &output, const LocalTensor &input, + const uint32_t curRowNum, const uint32_t curColNum) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(float); + WholeReduceMax(output, input, curColNum, curRowNum, 1, 1, CeilDiv(curColNum, elemInOneBlock), + ReduceOrder::ORDER_ONLY_VALUE); + PipeBarrier(); +} + +__aicore__ inline void LastDimReduceSumPerf(const LocalTensor &output, const LocalTensor &input, + const uint32_t curRowNum, const uint32_t curColNum) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(float); + WholeReduceSum(output, input, curColNum, curRowNum, 1, 1, CeilDiv(curColNum, elemInOneBlock)); + PipeBarrier(); +} + + +// 暂时只支持R轴小于64,既curColNum不能超过64 +__aicore__ inline void SoftmaxFP32Perf(const LocalTensor &output, const LocalTensor &input, + const LocalTensor &tmpReduceBuffer, + const LocalTensor tmpBrcbBuffer, const int32_t curRowNum, + const int32_t curColNum, float eps) +{ + LastDimReduceMaxPerf(tmpReduceBuffer, input, curRowNum, curColNum); + SubABLastDimBrcInline(output, input, tmpReduceBuffer, tmpBrcbBuffer, curRowNum, curColNum); + uint32_t curColNumAlign = RoundUp(curColNum); + Exp(output, output, curRowNum * curColNumAlign); + PipeBarrier(); + LastDimReduceSumPerf(tmpReduceBuffer, output, curRowNum, curColNum); + DivABLastDimBrcInline(output, output, tmpReduceBuffer, tmpBrcbBuffer, curRowNum, curColNum); + Adds(output, output, eps, curRowNum * curColNumAlign); + PipeBarrier(); +} + +// (bs, hc_mult, hc_mult) = (bs, hc_mult, hc_mult) + (bs, 1, hc_mult) + +template +__aicore__ inline void DivABABrcInline(const LocalTensor &output, const LocalTensor &input0, + const LocalTensor &input1, const uint32_t dim0, const uint32_t dim1, + const uint32_t dim2) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + uint32_t dim2Align = RoundUp(dim2); + uint32_t dim2RepeatTimes = dim2 / elemInOneRepeat; + uint32_t dim2Reminder = dim2 % elemInOneRepeat; + uint32_t dim2RepeatStride = CeilDiv(dim2, elemInOneBlock); + // 在dim1方向开repeat + BinaryRepeatParams instrParams; + if (dim1 >= dim2RepeatTimes) { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = dim2RepeatStride; + instrParams.src0RepStride = dim2RepeatStride; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < dim0; i++) { + for (uint32_t j = 0; j < dim2RepeatTimes; j++) { + Div(output[i * dim1 * dim2Align + j * elemInOneRepeat], + input0[i * dim1 * dim2Align + j * elemInOneRepeat], input1[i * dim2Align + j * elemInOneRepeat], + elemInOneRepeat, dim1, instrParams); + } + if (dim2Reminder != 0) { + Div(output[i * dim1 * dim2Align + dim2RepeatTimes * elemInOneRepeat], + input0[i * dim1 * dim2Align + dim2RepeatTimes * elemInOneRepeat], + input1[i * dim2Align + dim2RepeatTimes * elemInOneRepeat], dim2Reminder, dim1, instrParams); + } + } + } else { + // 在dim2方向开repeat + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = DEFAULT_REPEAT_STRIDE; + for (uint32_t i = 0; i < dim0; i++) { + for (uint32_t j = 0; j < dim1; j++) { + Div(output[i * dim1 * dim2Align + j * dim2Align], input0[i * dim1 * dim2Align + j * dim2Align], + input1[i * dim2Align], dim2); + } + } + } + PipeBarrier(); +} + +template +__aicore__ inline void CopyIn(const GlobalTensor &inputGm, const LocalTensor &inputTensor, const uint16_t nBurst, + const uint32_t copyLen, uint32_t srcStride = 0, uint32_t dstStride = 0) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + DataCopyPadExtParams dataCopyPadExtParams; + dataCopyPadExtParams.isPad = false; + dataCopyPadExtParams.leftPadding = 0; + dataCopyPadExtParams.rightPadding = 0; + dataCopyPadExtParams.paddingValue = 0; + + DataCopyExtParams dataCoptExtParams; + dataCoptExtParams.blockCount = nBurst; + dataCoptExtParams.blockLen = copyLen * sizeof(T); + dataCoptExtParams.srcStride = srcStride * sizeof(T); + dataCoptExtParams.dstStride = dstStride / elemInOneBlock; + DataCopyPad(inputTensor, inputGm, dataCoptExtParams, dataCopyPadExtParams); +} + +// (bs, hc_mult, hc_mult) --> (bs, hc_mult, hc_mult_align) +template +__aicore__ inline void CopyInWithOuterFor(const GlobalTensor &inputGm, const LocalTensor &inputTensor, + const uint16_t outerLoop, const uint16_t nBurst, const uint32_t copyLen, + const uint32_t gmLastDim) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + uint32_t ubLastDimAlign = RoundUp(copyLen); + + for (uint16_t i = 0; i < outerLoop; i++) { + CopyIn(inputGm[i * nBurst * gmLastDim], inputTensor[i * nBurst * ubLastDimAlign], nBurst, copyLen); + } +} + +template +__aicore__ inline void CopyInWithOuterFor(const GlobalTensor &inputGm, +const LocalTensor &inputTensor, const uint16_t outerLoop, +const uint16_t nBurst, const uint32_t copyLen, const uint32_t gmFirstDim, +const uint32_t gmLastDim) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + uint32_t ubLastDimAlign = RoundUp(copyLen); + if (outerLoop <= nBurst) { + for (uint32_t i = 0; i < outerLoop; i++) { + CopyIn(inputGm[i * gmFirstDim * gmLastDim], + inputTensor[i * nBurst * ubLastDimAlign], nBurst, copyLen, + gmLastDim - copyLen); + } + } else { + uint32_t srcStride = (gmLastDim - copyLen) + (gmFirstDim - 1) * gmLastDim; + uint32_t dstStride = (nBurst - 1) * ubLastDimAlign; + for (uint32_t i = 0; i < nBurst; i++) { + CopyIn(inputGm[i * gmLastDim], inputTensor[i * ubLastDimAlign], outerLoop, copyLen, srcStride, dstStride); + } + } +} + +template +__aicore__ inline void CopyOut(const LocalTensor &outputTensor, const GlobalTensor &outputGm, + const uint16_t nBurst, const uint32_t copyLen, uint32_t dstStride = 0) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = nBurst; + dataCopyParams.blockLen = copyLen * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = dstStride * sizeof(T); + DataCopyPad(outputGm, outputTensor, dataCopyParams); +} + +} // namespace HcPre + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_base_arch35.h b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_base_arch35.h new file mode 100644 index 000000000..4ef2fe85a --- /dev/null +++ b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_base_arch35.h @@ -0,0 +1,1147 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef HC_PRE_SINKHORN_RGEBASE_BASE_H +#define HC_PRE_SINKHORN_RGEBASE_BASE_H + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" + +namespace HcPreNs { +using namespace AscendC; +using namespace AscendC::MicroAPI; +using AscendC::MicroAPI::MaskReg; +using AscendC::MicroAPI::RegTensor; +using AscendC::MicroAPI::UnalignReg; +constexpr int32_t BLOCK_SIZE = 32; +constexpr int32_t VL_FP32 = 64; +constexpr int32_t C0_SIZE = 8; +constexpr int32_t FOUR_UNFOLD = 4; +constexpr int32_t DOUBLE_BUFFER = 2; +constexpr MatmulConfig MM_CFG = GetMDLConfig(); + +__aicore__ inline uint64_t Align(uint64_t a, uint64_t b) +{ + if (b == 0) { + return a; + } + return (a + b - 1) / b * b; +} + +__aicore__ inline uint64_t CeilDiv(uint64_t a, uint64_t b) +{ + if (b == 0) { + return a; + } + return (a + b - 1) / b; +} + +__aicore__ inline uint64_t CeilAlign(uint64_t a, uint64_t b) +{ + return CeilDiv(a, b) * b; +} + +template +__aicore__ inline int32_t RoundUp(int32_t num) +{ + int32_t elemNum = BLOCK_SIZE / sizeof(T); + return CeilAlign(num, elemNum); +} + +constexpr AscendC::MicroAPI::CastTrait castTraitB162B32Even = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::UNKNOWN, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB322B16Even = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_RINT, +}; + +template +__aicore__ inline void SetWaitFlag(HardEvent evt) +{ + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(evt)); + SetFlag(eventId); + WaitFlag(eventId); +} + +template +__aicore__ inline void LoadInputData(RegTensor &dst, __local_mem__ T *src, MaskReg pregLoop, uint32_t srcOffset) +{ + if constexpr (IsSameType::value) { + DataCopy(dst, src + srcOffset); + } else if constexpr (IsSameType::value || IsSameType::value) { + RegTensor tmp; + DataCopy(tmp, src + srcOffset); + Cast(dst, tmp, pregLoop); + } +} + +template +__aicore__ inline void StoreOutputData(__local_mem__ T *dst, RegTensor &src, MaskReg pregLoop, + uint32_t dstOffset) +{ + if constexpr (IsSameType::value) { + DataCopy(dst + dstOffset, src, pregLoop); + } else if constexpr (IsSameType::value || IsSameType::value) { + RegTensor tmp; + Cast(tmp, src, pregLoop); + DataCopy(dst + dstOffset, tmp, pregLoop); + } +} + +template +__aicore__ inline void LoadInputDataWithBrc(RegTensor &dst, __local_mem__ T *src, MaskReg pregLoop, + uint32_t srcOffset) +{ + if constexpr (IsSameType::value) { + DataCopy(dst, src + srcOffset); + } else if constexpr (IsSameType::value || IsSameType::value) { + RegTensor tmp; + DataCopy(tmp, src + srcOffset); + Cast(dst, tmp, pregLoop); + } +} + +template +__aicore__ inline void LoadInputDataUnalign( + RegTensor& dst, __local_mem__ T*& src, UnalignReg& uSrc, MaskReg pregLoop, uint32_t postUpdateStride) +{ + if constexpr (IsSameType::value) { + DataCopyUnAlign(dst, uSrc, src, postUpdateStride); + } else if constexpr (IsSameType::value || IsSameType::value) { + RegTensor tmp; + RegTensor tmpUnPack; + DataCopyUnAlign(tmp, uSrc, src, postUpdateStride); + UnPack((RegTensor&)tmpUnPack, (RegTensor&)tmp); + Cast(dst, tmpUnPack, pregLoop); + } +} + + +__aicore__ inline void VFSigmoid(RegTensor &y, RegTensor &x, RegTensor &one, MaskReg pregLoop) +{ + Muls(x, x, static_cast(-1), pregLoop); + Exp(x, x, pregLoop); + Adds(x, x, static_cast(1), pregLoop); + Div(y, one, x, pregLoop); +} + +__aicore__ inline void VFTransND2NZ(const LocalTensor &yLocal, const LocalTensor &xLocal, + const uint16_t curRowNum, const uint16_t curColNum) +{ + __local_mem__ float *yLocalAddr = (__local_mem__ float *)yLocal.GetPhyAddr(); + __local_mem__ float *xLocalAddr = (__local_mem__ float *)xLocal.GetPhyAddr(); + // NZ分享要求M,N方向按照16对齐 + uint16_t curRowNumAlign = CeilAlign(curRowNum, C0_SIZE); + uint16_t c1Size = BLOCK_SIZE / sizeof(float); + uint16_t curColNumAlign = RoundUp(curColNum); + uint16_t curRowMainCount = curRowNum / C0_SIZE; + uint16_t curRowReminder = curRowNum % C0_SIZE; + uint16_t tailBaseOffset = curRowReminder * c1Size; + uint32_t dataBlockStride = curColNumAlign / c1Size + 1; + uint16_t loopCount = curColNumAlign / c1Size; + if (curRowReminder == 0) { + __VEC_SCOPE__ + { + RegTensor x; + RegTensor y; + MaskReg pregMain = CreateMask(); + for (uint16_t i = 0; i < curRowMainCount; i++) { + for (uint16_t j = 0; j < loopCount; j++) { + DataCopy( + x, xLocalAddr + i * C0_SIZE * (curColNumAlign + BLOCK_SIZE / sizeof(float)) + j * c1Size, dataBlockStride, pregMain); + DataCopy(yLocalAddr + i * C0_SIZE * c1Size + j * curRowNumAlign * c1Size, x, pregMain); + } + } + } + } else { + __VEC_SCOPE__ + { + RegTensor x; + RegTensor y; + MaskReg pregMain = CreateMask(); + uint32_t sreg = curRowReminder * C0_SIZE; + MaskReg pregLoop = UpdateMask(sreg); + for (uint16_t i = 0; i < curRowMainCount; i++) { + for (uint16_t j = 0; j < loopCount; j++) { + DataCopy( + x, xLocalAddr + i * C0_SIZE * (curColNumAlign + BLOCK_SIZE / sizeof(float)) + j * c1Size, dataBlockStride, pregMain); + DataCopy(yLocalAddr + i * C0_SIZE * c1Size + j * curRowNumAlign * c1Size, x, pregMain); + } + } + xLocalAddr = xLocalAddr + curRowMainCount * C0_SIZE * (curColNumAlign + BLOCK_SIZE / sizeof(float)); + yLocalAddr = yLocalAddr + curRowMainCount * C0_SIZE * c1Size; + for (uint16_t i = 0; i < loopCount; i++) { + DataCopy(x, xLocalAddr + i * c1Size, + dataBlockStride, pregLoop); + DataCopy(yLocalAddr + i * curRowNumAlign * c1Size, x, pregLoop); + } + } + } +} + +template +__aicore__ inline void VFProcessCast(const LocalTensor &yLocal, const LocalTensor &xLocal, + const uint16_t curRowNum, const uint16_t curColNum) +{ + __local_mem__ float *yLocalAddr = (__local_mem__ float *)yLocal.GetPhyAddr(); + __local_mem__ T *xLocalAddr = (__local_mem__ T *)xLocal.GetPhyAddr(); + uint16_t loopCount = CeilDiv(curColNum, VL_FP32); + uint16_t curColNumAlign = RoundUp(curColNum); + uint16_t dstCurColNumAlign = RoundUp(curColNum) + BLOCK_SIZE / sizeof(float); + if (loopCount > 1) { + __VEC_SCOPE__ + { + RegTensor x; + RegTensor y; + MaskReg pregLoop; + uint32_t sreg; + for (uint16_t i = 0; i < curRowNum; i++) { + sreg = curColNum; + for (uint16_t j = 0; j < loopCount; j++) { + pregLoop = UpdateMask(sreg); + LoadInputData(x, xLocalAddr, pregLoop, i * curColNumAlign + j * VL_FP32); + StoreOutputData(yLocalAddr, x, pregLoop, i * dstCurColNumAlign + j * VL_FP32); + } + } + } + } else { + __VEC_SCOPE__ + { + RegTensor x; + RegTensor y; + MaskReg pregLoop = CreateMask(); + for (uint16_t i = 0; i < curRowNum; i++) { + LoadInputData(x, xLocalAddr, pregLoop, i * curColNumAlign); + StoreOutputData(yLocalAddr, x, pregLoop, i * dstCurColNumAlign); + } + } + } + +} + +template +__aicore__ inline void VFProcessCastAndInvRmsPart1(const LocalTensor &rmsNormLocal, const LocalTensor &xCastLocal, + const LocalTensor &xLocal, float coeff, + const uint16_t curRowNum, const uint16_t curColNum) +{ + __local_mem__ float *rmsNormLocalAddr = (__local_mem__ float *)rmsNormLocal.GetPhyAddr(); + __local_mem__ float *xCastLocalAddr = (__local_mem__ float *)xCastLocal.GetPhyAddr(); + __local_mem__ T *xLocalAddr = (__local_mem__ T *)xLocal.GetPhyAddr(); + uint16_t loopCount = CeilDiv(curColNum, VL_FP32); + uint16_t curColNumAlign = RoundUp(curColNum); + uint16_t dstCurColNumAlign = RoundUp(curColNum) + BLOCK_SIZE / sizeof(float); + if (loopCount > 1) { + __VEC_SCOPE__ + { + RegTensor x; + RegTensor x1; + RegTensor sum; + RegTensor one; + RegTensor y; + MaskReg pregLoop; + MaskReg pregMain = CreateMask(); + MaskReg pregMerge = CreateMask(); + uint32_t sreg; + for (uint16_t i = 0; i < curRowNum; i++) { + Duplicate(sum, 0.0f); + if constexpr (WithUbReduce) { + LoadInputDataWithBrc(y, rmsNormLocalAddr, pregMerge, i); + } + sreg = curColNum; + for (uint16_t j = 0; j < loopCount; j++) { + pregLoop = UpdateMask(sreg); + LoadInputData(x, xLocalAddr, pregLoop, i * curColNumAlign + j * VL_FP32); + Mul(x1, x, x, pregLoop); + Add(sum, sum, x1, pregMain); + StoreOutputData(xCastLocalAddr, x, pregLoop, i * dstCurColNumAlign + j * VL_FP32); + } + Muls(sum, sum, coeff, pregMain); + ReduceSum(sum, sum, pregMain); + if constexpr (WithUbReduce) { + Add(y, y, sum, pregMerge); + DataCopy(rmsNormLocalAddr + i, y, pregMerge); + } else { + DataCopy(rmsNormLocalAddr + i, sum, pregMerge); + } + } + + } + } else { + __VEC_SCOPE__ + { + RegTensor x; + RegTensor x1; + RegTensor sum; + RegTensor one; + RegTensor y; + MaskReg pregMain = CreateMask(); + MaskReg pregMerge = CreateMask(); + uint32_t sreg = curColNum; + MaskReg pregLoop = UpdateMask(sreg); + for (uint16_t i = 0; i < curRowNum; i++) { + Duplicate(sum, 0.0f); + if constexpr (WithUbReduce) { + LoadInputDataWithBrc(y, rmsNormLocalAddr, pregMerge, i); + } + LoadInputData(x, xLocalAddr, pregLoop, i * curColNumAlign); + StoreOutputData(xCastLocalAddr, x, pregLoop, i * dstCurColNumAlign); + Mul(x1, x, x, pregLoop); + Add(sum, sum, x1, pregLoop); + Muls(sum, sum, coeff, pregLoop); + ReduceSum(sum, sum, pregLoop); + if constexpr (WithUbReduce) { + Add(y, y, sum, pregMerge); + DataCopy(rmsNormLocalAddr + i, y, pregMerge); + } else { + DataCopy(rmsNormLocalAddr + i, sum, pregMerge); + } + } + } + } +} + + +template +__aicore__ inline void VFProcessInvRmsPart1(const LocalTensor &yLocal, const LocalTensor &xLocal, + float coeff, uint16_t curRowNum, uint32_t curColNum) +{ + __local_mem__ float *yLocalAddr = (__local_mem__ float *)yLocal.GetPhyAddr(); + __local_mem__ float *xLocalAddr = (__local_mem__ float *)xLocal.GetPhyAddr(); + + uint16_t loopCount = CeilDiv(curColNum, VL_FP32); + uint16_t curColNumAlign = RoundUp(curColNum) + BLOCK_SIZE / sizeof(float); + if (loopCount > 1) { + __VEC_SCOPE__ + { + RegTensor x; + RegTensor sum; + RegTensor one; + RegTensor y; + MaskReg pregLoop; + MaskReg pregMain = CreateMask(); + MaskReg pregMerge = CreateMask(); + uint32_t sreg; + for (uint16_t i = 0; i < curRowNum; i++) { + Duplicate(sum, 0.0f); + if constexpr (WithUbReduce) { + LoadInputDataWithBrc(y, yLocalAddr, pregMerge, i); + } + sreg = curColNum; + for (uint16_t j = 0; j < loopCount; j++) { + pregLoop = UpdateMask(sreg); + LoadInputData(x, xLocalAddr, pregLoop, i * curColNumAlign + j * VL_FP32); + Mul(x, x, x, pregLoop); + Add(sum, sum, x, pregMain); + } + Muls(sum, sum, coeff, pregMain); + ReduceSum(sum, sum, pregMain); + if constexpr (WithUbReduce) { + Add(y, y, sum, pregMerge); + DataCopy(yLocalAddr + i, y, pregMerge); + } else { + DataCopy(yLocalAddr + i, sum, pregMerge); + } + } + } + } else { + __VEC_SCOPE__ + { + RegTensor x; + RegTensor sum; + RegTensor one; + RegTensor y; + MaskReg pregMain = CreateMask(); + MaskReg pregMerge = CreateMask(); + uint32_t sreg = curColNum; + MaskReg pregLoop = UpdateMask(sreg); + for (uint16_t i = 0; i < curRowNum; i++) { + Duplicate(sum, 0.0f); + if constexpr (WithUbReduce) { + LoadInputDataWithBrc(y, yLocalAddr, pregMerge, i); + } + LoadInputData(x, xLocalAddr, pregLoop, i * curColNumAlign); + Mul(x, x, x, pregLoop); + Add(sum, sum, x, pregLoop); + Muls(sum, sum, coeff, pregLoop); + ReduceSum(sum, sum, pregLoop); + if constexpr (WithUbReduce) { + Add(y, y, sum, pregMerge); + DataCopy(yLocalAddr + i, y, pregMerge); + } else { + DataCopy(yLocalAddr + i, sum, pregMerge); + } + } + } + } +} + +// (bs, k) --> (bs, 1) +// k 为Matmul K轴切分时的分核数，必然小于64 +__aicore__ inline void VFProcessInvRmsPart2(const LocalTensor &yLocal, const LocalTensor &xLocal, + float eps, uint16_t curRowNum, uint32_t curColNum) +{ + __local_mem__ float *yLocalAddr = (__local_mem__ float *)yLocal.GetPhyAddr(); + __local_mem__ float *xLocalAddr = (__local_mem__ float *)xLocal.GetPhyAddr(); + uint16_t curColNumAlign = RoundUp(curColNum); + __VEC_SCOPE__ + { + RegTensor x; + RegTensor sum; + RegTensor one; + RegTensor y; + uint32_t sreg = curColNum; + MaskReg pregLoop = UpdateMask(sreg); + MaskReg pregMerge = CreateMask(); + Duplicate(one, static_cast(1.0), pregMerge); + for (uint16_t i = 0; i < curRowNum; i++) { + LoadInputData(x, xLocalAddr, pregLoop, i * curColNumAlign); + ReduceSum(sum, x, pregLoop); + Adds(sum, sum, eps, pregMerge); + Sqrt(sum, sum, pregMerge); + Div(y, one, sum, pregMerge); + DataCopy(yLocalAddr + i, y, pregMerge); + } + } +} + + +// (k, bs, hc_mix) * (k, bs, 1) = (bs, hc_mix) +// for循环组织形式如下: +/* + for (i, 0, bs) + + for (j, 0, k) + for (h, 0, hc_mix) +*/ +// hcMix小于64，因此直接去掉内层for循环 +__aicore__ inline void VFProcessInvRmsPart3WithGroupReduce(const LocalTensor &yLocal, + const LocalTensor &mmLocal, + const LocalTensor &xLocal, float eps, uint16_t groupK, + uint16_t bs, uint32_t hcMix) +{ + __local_mem__ float *yLocalAddr = (__local_mem__ float *)yLocal.GetPhyAddr(); + __local_mem__ float *mmLocalAddr = (__local_mem__ float *)mmLocal.GetPhyAddr(); + __local_mem__ float *xLocalAddr = (__local_mem__ float *)xLocal.GetPhyAddr(); + uint32_t hcMixAlign = RoundUp(hcMix); + uint32_t bsAlign = RoundUp(bs); + uint16_t fourLoopNum = groupK / FOUR_UNFOLD; + uint16_t tailLoopNum = groupK % FOUR_UNFOLD; + if (groupK < 4) { + __VEC_SCOPE__ + { + RegTensor x; + RegTensor sum1; + RegTensor sum2; + RegTensor one; + RegTensor rsqrt; + RegTensor y; + RegTensor mm; + uint32_t sreg = hcMix; + MaskReg pregLoop = UpdateMask(sreg); + MaskReg pregMerge = CreateMask(); + Duplicate(one, static_cast(1.0), pregMerge); + for (uint16_t i = 0; i < bs; i++) { + Duplicate(sum1, static_cast(0.0f), pregMerge); + Duplicate(sum2, static_cast(0.0f), pregLoop); + for (uint16_t j = 0; j < groupK; j++) { + LoadInputDataWithBrc(x, xLocalAddr, pregMerge, i + j * bsAlign); + Add(sum1, sum1, x, pregMerge); + LoadInputData(mm, mmLocalAddr, pregLoop, i * hcMixAlign + j * bs * hcMixAlign); + Add(sum2, sum2, mm, pregLoop); + } + Adds(sum1, sum1, eps, pregMerge); + Sqrt(sum1, sum1, pregMerge); + Div(rsqrt, one, sum1, pregMerge); + Duplicate(rsqrt, rsqrt, pregLoop); + Mul(y, sum2, rsqrt, pregLoop); + StoreOutputData(yLocalAddr, y, pregLoop, i * hcMixAlign); + } + } + } else { + __VEC_SCOPE__ + { + RegTensor x1; + RegTensor x2; + RegTensor x3; + RegTensor x4; + RegTensor mm1; + RegTensor mm2; + RegTensor mm3; + RegTensor mm4; + RegTensor sumX1; + RegTensor sumX2; + RegTensor sumX3; + RegTensor sumX4; + RegTensor sumM1; + RegTensor sumM2; + RegTensor sumM3; + RegTensor sumM4; + RegTensor one; + RegTensor rsqrt; + RegTensor y; + uint32_t sreg = hcMix; + MaskReg pregLoop = UpdateMask(sreg); + MaskReg pregMerge = CreateMask(); + Duplicate(one, static_cast(1.0), pregMerge); + for (uint16_t i = 0; i < bs; i++) { + Duplicate(sumX1, static_cast(0.0f), pregMerge); + Duplicate(sumX2, static_cast(0.0f), pregMerge); + Duplicate(sumX3, static_cast(0.0f), pregMerge); + Duplicate(sumX4, static_cast(0.0f), pregMerge); + Duplicate(sumM1, static_cast(0.0f), pregLoop); + Duplicate(sumM2, static_cast(0.0f), pregLoop); + Duplicate(sumM3, static_cast(0.0f), pregLoop); + Duplicate(sumM4, static_cast(0.0f), pregLoop); + for (uint16_t j = 0; j < fourLoopNum; j++) { + LoadInputDataWithBrc(x1, xLocalAddr, pregMerge, i + 4 * j * bsAlign); + Add(sumX1, sumX1, x1, pregMerge); + LoadInputData(mm1, mmLocalAddr, pregLoop, i * hcMixAlign + 4 * j * bs * hcMixAlign); + Add(sumM1, sumM1, mm1, pregLoop); + + LoadInputDataWithBrc(x2, xLocalAddr, pregMerge, i + (4 * j + 1) * bsAlign); + Add(sumX2, sumX2, x2, pregMerge); + LoadInputData(mm2, mmLocalAddr, pregLoop, i * hcMixAlign + (4 * j + 1) * bs * hcMixAlign); + Add(sumM2, sumM2, mm2, pregLoop); + + LoadInputDataWithBrc(x3, xLocalAddr, pregMerge, i + (4 * j + 2) * bsAlign); + Add(sumX3, sumX3, x3, pregMerge); + LoadInputData(mm3, mmLocalAddr, pregLoop, i * hcMixAlign + (4 * j + 2) * bs * hcMixAlign); + Add(sumM3, sumM3, mm3, pregLoop); + + LoadInputDataWithBrc(x4, xLocalAddr, pregMerge, i + (4 * j + 3) * bsAlign); + Add(sumX4, sumX4, x4, pregMerge); + LoadInputData(mm4, mmLocalAddr, pregLoop, i * hcMixAlign + (4 * j + 3) * bs * hcMixAlign); + Add(sumM4, sumM4, mm4, pregLoop); + } + for (uint16_t j = 0; j < tailLoopNum; j++) { + LoadInputDataWithBrc(x1, xLocalAddr, pregMerge, i + (fourLoopNum * FOUR_UNFOLD + j) * bsAlign); + Add(sumX1, sumX1, x1, pregMerge); + LoadInputData(mm1, mmLocalAddr, pregLoop, i * hcMixAlign + (fourLoopNum * FOUR_UNFOLD + j) * bs * hcMixAlign); + Add(sumM1, sumM1, mm1, pregLoop); + } + Add(sumX1, sumX1, sumX4, pregMerge); + Add(sumX2, sumX2, sumX3, pregMerge); + Add(sumX1, sumX1, sumX2, pregMerge); + Add(sumM1, sumM1, sumM4, pregLoop); + Add(sumM2, sumM2, sumM3, pregLoop); + Add(sumM1, sumM1, sumM2, pregLoop); + + Adds(sumX1, sumX1, eps, pregMerge); + Sqrt(sumX1, sumX1, pregMerge); + Div(rsqrt, one, sumX1, pregMerge); + Duplicate(rsqrt, rsqrt, pregLoop); + Mul(y, sumM1, rsqrt, pregLoop); + StoreOutputData(yLocalAddr, y, pregLoop, i * hcMixAlign); + } + } + } +} + + +__aicore__ inline void VFProcessInvRmsPart3(const LocalTensor &yLocal, + const LocalTensor &mmLocal, + const LocalTensor &xLocal, float eps, + uint16_t bs, uint32_t hcMix) +{ + __local_mem__ float *yLocalAddr = (__local_mem__ float *)yLocal.GetPhyAddr(); + __local_mem__ float *mmLocalAddr = (__local_mem__ float *)mmLocal.GetPhyAddr(); + __local_mem__ float *xLocalAddr = (__local_mem__ float *)xLocal.GetPhyAddr(); + uint32_t hcMixAlign = RoundUp(hcMix); + uint32_t bsAlign = RoundUp(bs); + __VEC_SCOPE__ + { + RegTensor x; + RegTensor sum1; + RegTensor sum2; + RegTensor one; + RegTensor rsqrt; + RegTensor y; + RegTensor mm; + uint32_t sreg = hcMix; + MaskReg pregLoop = UpdateMask(sreg); + Duplicate(one, static_cast(1.0), pregLoop); + for (uint16_t i = 0; i < bs; i++) { + LoadInputDataWithBrc(x, xLocalAddr, pregLoop, i); + LoadInputData(mm, mmLocalAddr, pregLoop, i * hcMixAlign); + Adds(x, x, eps, pregLoop); + Sqrt(x, x, pregLoop); + Div(rsqrt, one, x, pregLoop); + Mul(y, mm, rsqrt, pregLoop); + StoreOutputData(yLocalAddr, y, pregLoop, i * hcMixAlign); + } + } +} + + +// Matmul的结果会直接FixPipe到UB上，不会在搬运时完成Split动作，因此需要在UB内完成Split动作 +__aicore__ inline void VFProcessPre(const LocalTensor &preLocal, const LocalTensor &mixLocal, + const LocalTensor &hcBaseLocal, + float scale, float eps, uint16_t curRowNum, uint16_t curColNum, uint16_t hcMix) +{ + __local_mem__ float *preLocalAddr = (__local_mem__ float *)preLocal.GetPhyAddr(); + __local_mem__ float *mixLocalAddr = (__local_mem__ float *)mixLocal.GetPhyAddr(); + __local_mem__ float *hcBaseLocalAddr = (__local_mem__ float *)hcBaseLocal.GetPhyAddr(); + uint16_t loopCount = CeilDiv(curColNum, VL_FP32); + uint32_t curColNumAlign = RoundUp(curColNum); + uint32_t hcMixAlign = RoundUp(hcMix); + if (loopCount > 1) { + __VEC_SCOPE__ + { + RegTensor mix; + RegTensor base; + RegTensor one; + MaskReg pregLoop = CreateMask(); + uint32_t sreg = curColNum; + Duplicate(one, static_cast(1), pregLoop); + for (uint16_t i = 0; i < loopCount; i++) { + pregLoop = UpdateMask(sreg); + LoadInputData(base, hcBaseLocalAddr, pregLoop, i * VL_FP32); + for (uint16_t j = 0; j < curRowNum; j++) { + LoadInputData(mix, mixLocalAddr, pregLoop, i * VL_FP32 + j * hcMixAlign); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + VFSigmoid(mix, mix, one, pregLoop); + Adds(mix, mix, eps, pregLoop); + StoreOutputData(preLocalAddr, mix, pregLoop, i * VL_FP32 + j * hcMixAlign); + } + } + } + } else { + __VEC_SCOPE__ + { + RegTensor mix; + RegTensor base; + RegTensor one; + uint32_t sreg = curColNum; + MaskReg pregLoop = UpdateMask(sreg); + Duplicate(one, static_cast(1), pregLoop); + LoadInputData(base, hcBaseLocalAddr, pregLoop, 0); + for (uint16_t i = 0; i < curRowNum; i++) { + LoadInputData(mix, mixLocalAddr, pregLoop, i * hcMixAlign); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + VFSigmoid(mix, mix, one, pregLoop); + Adds(mix, mix, eps, pregLoop); + StoreOutputData(preLocalAddr, mix, pregLoop, i * hcMixAlign); + } + } + } +} + +__aicore__ inline void VFProcessPost(const LocalTensor &postLocal, const LocalTensor &mixLocal, + const LocalTensor &hcBaseLocal, + float scale, float eps, uint16_t curRowNum, uint16_t curColNum, uint16_t hcMix) +{ + __local_mem__ float *postLocalAddr = (__local_mem__ float *)postLocal.GetPhyAddr(); + __local_mem__ float *mixOriginLocalAddr = (__local_mem__ float *)mixLocal.GetPhyAddr(); + __local_mem__ float *hcBaseLocalAddr = (__local_mem__ float *)hcBaseLocal.GetPhyAddr(); + __local_mem__ float *mixLocalAddr = mixOriginLocalAddr; + uint16_t loopCount = CeilDiv(curColNum, VL_FP32); + uint32_t curColNumAlign = RoundUp(curColNum); + uint32_t hcMixAlign = RoundUp(hcMix); + if (loopCount > 1) { + __VEC_SCOPE__ + { + RegTensor mix; + RegTensor base; + RegTensor one; + UnalignReg uMix; + MaskReg pregLoop = CreateMask(); + uint32_t sreg = curColNum; + Duplicate(one, static_cast(1), pregLoop); + DataCopyUnAlignPre(uMix, mixLocalAddr); + for (uint16_t i = 0; i < loopCount; i++) { + mixLocalAddr = mixOriginLocalAddr + i * VL_FP32; + pregLoop = UpdateMask(sreg); + LoadInputData(base, hcBaseLocalAddr, pregLoop, i * VL_FP32); + for (uint16_t j = 0; j < curRowNum; j++) { + LoadInputDataUnalign(mix, mixLocalAddr, uMix, pregLoop, hcMixAlign); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + VFSigmoid(mix, mix, one, pregLoop); + Muls(mix, mix, static_cast(2.0), pregLoop); + StoreOutputData(postLocalAddr, mix, pregLoop, i * VL_FP32 + j * curColNumAlign); + } + } + } + } else { + __VEC_SCOPE__ + { + RegTensor mix; + RegTensor base; + RegTensor one; + UnalignReg uMix; + uint32_t sreg = curColNum; + MaskReg pregLoop = UpdateMask(sreg); + Duplicate(one, static_cast(1), pregLoop); + DataCopyUnAlignPre(uMix, mixLocalAddr); + LoadInputData(base, hcBaseLocalAddr, pregLoop, 0); + for (uint16_t i = 0; i < curRowNum; i++) { + LoadInputDataUnalign(mix, mixLocalAddr, uMix, pregLoop, hcMixAlign); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + VFSigmoid(mix, mix, one, pregLoop); + Muls(mix, mix, static_cast(2.0), pregLoop); + StoreOutputData(postLocalAddr, mix, pregLoop, i * curColNumAlign); + } + } + } +} + +// dim2是R轴，R轴小于64, 不需要回写UB +__aicore__ inline void VFProcessCombFragRLessVL(const LocalTensor &combFragLocal, + const LocalTensor &mixLocal, + const LocalTensor &hcBaseLocal, + float scale, float eps, + uint16_t iters, uint16_t dim0, uint16_t dim1, uint16_t dim2, uint16_t hcMix) +{ + __local_mem__ float *combFragLocalAddr = (__local_mem__ float *)combFragLocal.GetPhyAddr(); + __local_mem__ float *mixLocalOriginAddr = (__local_mem__ float *)mixLocal.GetPhyAddr(); + __local_mem__ float *hcBaseLocalAddr = (__local_mem__ float *)hcBaseLocal.GetPhyAddr(); + __local_mem__ float *mixLocalAddr = mixLocalOriginAddr; + uint32_t dim2Align = RoundUp(dim2); + uint32_t hcMixAlign = RoundUp(hcMix); + __VEC_SCOPE__ + { + RegTensor base; + RegTensor mix; + RegTensor rsqrt; + RegTensor max; + RegTensor sum; + RegTensor sum1; + UnalignReg uMix; + uint32_t sreg = dim2; + MaskReg pregLoop = UpdateMask(sreg); + DataCopyUnAlignPre(uMix, mixLocalAddr); + for (uint16_t i = 0; i < dim0; i++) { + Duplicate(sum1, static_cast(0), pregLoop); + for (uint16_t j = 0; j < dim1; j++) { + mixLocalAddr = mixLocalOriginAddr + i * hcMixAlign + j * dim2; + LoadInputData(base, hcBaseLocalAddr, pregLoop, j * dim2Align); + LoadInputDataUnalign(mix, mixLocalAddr, uMix, pregLoop, VL_FP32); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + ReduceMax(max, mix, pregLoop); + Duplicate(max, max, pregLoop); + Sub(mix, mix, max, pregLoop); + Exp(mix, mix, pregLoop); + ReduceSum(sum, mix, pregLoop); + Duplicate(sum, sum, pregLoop); + Div(mix, mix, sum, pregLoop); + Adds(mix, mix, eps, pregLoop); + Add(sum1, sum1, mix, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, i * dim1 * dim2Align + j * dim2Align); + } + LocalMemBar(); + Adds(sum1, sum1, eps, pregLoop); + for (uint16_t j = 0; j < dim1; j++) { + LoadInputData(mix, combFragLocalAddr, pregLoop, i * dim1 * dim2Align + j * dim2Align); + Div(mix, mix, sum1, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, i * dim1 * dim2Align + j * dim2Align); + } + } + for (uint16_t i = 0; i < iters; i++) { + LocalMemBar(); + for (uint16_t j = 0; j < dim0; j++) { + Duplicate(sum1, static_cast(0), pregLoop); + for (uint16_t k = 0; k < dim1; k++) { + LoadInputData(mix, combFragLocalAddr, pregLoop, j * dim1 * dim2Align + k * dim2Align); + ReduceSum(sum, mix, pregLoop); + Duplicate(sum, sum, pregLoop); + Adds(sum, sum, eps, pregLoop); + Div(mix, mix, sum, pregLoop); + Add(sum1, sum1, mix, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, j * dim1 * dim2Align + k * dim2Align); + } + LocalMemBar(); + Adds(sum1, sum1, eps, pregLoop); + for (uint16_t k = 0; k < dim1; k++) { + LoadInputData(mix, combFragLocalAddr, pregLoop, j * dim1 * dim2Align + k * dim2Align); + Div(mix, mix, sum1, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, j * dim1 * dim2Align + k * dim2Align); + } + } + } + } +} + +__aicore__ inline void VFProcessIteration(RegTensor &sum0, RegTensor &sum1, RegTensor &mix, + float eps, MaskReg pregLoop) +{ + ReduceSum(sum1, mix, pregLoop); + Duplicate(sum1, sum1, pregLoop); + Adds(sum1, sum1, eps, pregLoop); + Div(mix, mix, sum1, pregLoop); + Add(sum0, sum0, mix, pregLoop); +} + +__aicore__ inline void VFProcessCombFragRLessVLUseFourUnfold(const LocalTensor &combFragLocal, + const LocalTensor &mixLocal, + const LocalTensor &hcBaseLocal, + float scale, + float eps, uint16_t iters, uint16_t dim0, uint16_t dim1, + uint16_t dim2, uint16_t hcMix) +{ + __local_mem__ float *combFragLocalAddr = (__local_mem__ float *)combFragLocal.GetPhyAddr(); + __local_mem__ float *mixLocalOriginAddr = (__local_mem__ float *)mixLocal.GetPhyAddr(); + __local_mem__ float *hcBaseLocalAddr = (__local_mem__ float *)hcBaseLocal.GetPhyAddr(); + __local_mem__ float *mixLocalAddr = mixLocalOriginAddr; + uint32_t dim2Align = RoundUp(dim2); + uint32_t hcMixAlign = RoundUp(hcMix); + __VEC_SCOPE__ + { + RegTensor base; + RegTensor mix; + RegTensor mix1; + RegTensor mix2; + RegTensor mix3; + RegTensor mix4; + RegTensor max; + RegTensor sum; + RegTensor sum1; + RegTensor sum2; + RegTensor sum3; + RegTensor sum4; + UnalignReg uMix; + uint32_t sreg = dim2; + MaskReg pregLoop = UpdateMask(sreg); + DataCopyUnAlignPre(uMix, mixLocalAddr); + for (uint16_t i = 0; i < dim0; i++) { + Duplicate(sum1, static_cast(0), pregLoop); + for (uint16_t j = 0; j < dim1; j++) { + mixLocalAddr = mixLocalOriginAddr + i * hcMixAlign + j * dim2; + LoadInputData(base, hcBaseLocalAddr, pregLoop, j * dim2Align); + LoadInputDataUnalign(mix, mixLocalAddr, uMix, pregLoop, VL_FP32); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + ReduceMax(max, mix, pregLoop); + Duplicate(max, max, pregLoop); + Sub(mix, mix, max, pregLoop); + Exp(mix, mix, pregLoop); + ReduceSum(sum, mix, pregLoop); + Duplicate(sum, sum, pregLoop); + Div(mix, mix, sum, pregLoop); + Adds(mix, mix, eps, pregLoop); + Add(sum1, sum1, mix, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, i * dim1 * dim2Align + j * dim2Align); + } + LocalMemBar(); + Adds(sum1, sum1, eps, pregLoop); + for (uint16_t j = 0; j < dim1; j++) { + LoadInputData(mix, combFragLocalAddr, pregLoop, i * dim1 * dim2Align + j * dim2Align); + Div(mix, mix, sum1, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, i * dim1 * dim2Align + j * dim2Align); + } + } + LocalMemBar(); + for (uint16_t i = 0; i < dim0; i++) { + LoadInputData(mix1, combFragLocalAddr, pregLoop, i * dim1 * dim2Align); + LoadInputData(mix2, combFragLocalAddr, pregLoop, i * dim1 * dim2Align + 1 * dim2Align); + LoadInputData(mix3, combFragLocalAddr, pregLoop, i * dim1 * dim2Align + 2 * dim2Align); + LoadInputData(mix4, combFragLocalAddr, pregLoop, i * dim1 * dim2Align + 3 * dim2Align); + for (uint16_t j = 0; j < iters; j++) { + Duplicate(sum, static_cast(0), pregLoop); + VFProcessIteration(sum, sum1, mix1, eps, pregLoop); + VFProcessIteration(sum, sum2, mix2, eps, pregLoop); + VFProcessIteration(sum, sum3, mix3, eps, pregLoop); + VFProcessIteration(sum, sum4, mix4, eps, pregLoop); + Adds(sum, sum, eps, pregLoop); + Div(mix1, mix1, sum, pregLoop); + Div(mix2, mix2, sum, pregLoop); + Div(mix3, mix3, sum, pregLoop); + Div(mix4, mix4, sum, pregLoop); + } + StoreOutputData(combFragLocalAddr, mix1, pregLoop, i * dim1 * dim2Align); + StoreOutputData(combFragLocalAddr, mix2, pregLoop, i * dim1 * dim2Align + 1 * dim2Align); + StoreOutputData(combFragLocalAddr, mix3, pregLoop, i * dim1 * dim2Align + 2 * dim2Align); + StoreOutputData(combFragLocalAddr, mix4, pregLoop, i * dim1 * dim2Align + 3 * dim2Align); + } + } +} + +__aicore__ inline void RowGroupMaxBcast(RegTensor &out, RegTensor &in, RegTensor &t0, + RegTensor &t1, RegTensor &t2, RegTensor &t3, MaskReg preg) +{ + DeInterleave(t0, t1, in, in); // t0 has even lanes, t1 has odd lanes, paired inside each group. + Max(t0, t0, t1, preg); // Pairwise max for c0/c1 and c2/c3 inside each group. + DeInterleave(t2, t3, t0, t0); + Max(t2, t2, t3, preg); // lane0..3 holds the max for each row group. + Interleave(t0, t1, t2, t2); // [m0,m0,m1,m1,m2,m2,m3,m3] + Interleave(out, t1, t0, t0); // Broadcast each row max to 4 lanes. +} + +// Reduce the row sum over C in contiguous groups of 4 and broadcast it: out[4r+c] = sum_c M[r][c]. +__aicore__ inline void RowGroupSumBcast(RegTensor &out, RegTensor &in, RegTensor &t0, + RegTensor &t1, RegTensor &t2, RegTensor &t3, MaskReg preg) +{ + DeInterleave(t0, t1, in, in); + Add(t0, t0, t1, preg); + DeInterleave(t2, t3, t0, t0); + Add(t2, t2, t3, preg); + Interleave(t0, t1, t2, t2); + Interleave(out, t1, t0, t0); +} + +// In-place softmax over C: subtract row max, exp, divide by row sum, then add eps. +__aicore__ inline void SoftmaxRowBcastInplace(RegTensor &m, RegTensor &red, + RegTensor &t0, RegTensor &t1, + RegTensor &t2, RegTensor &t3, + float eps, MaskReg preg) +{ + RowGroupMaxBcast(red, m, t0, t1, t2, t3, preg); + Sub(m, m, red, preg); + Exp(m, m, preg); + RowGroupSumBcast(red, m, t0, t1, t2, t3, preg); + Div(m, m, red, preg); + Adds(m, m, eps, preg); +} + +// In-place divide by row sum plus eps, and accumulate the result into csum for Sinkhorn iterations. +__aicore__ inline void RowNormAccum(RegTensor &m, RegTensor &red, RegTensor &csum, + RegTensor &t0, RegTensor &t1, + RegTensor &t2, RegTensor &t3, + float eps, MaskReg preg) +{ + RowGroupSumBcast(red, m, t0, t1, t2, t3, preg); + Adds(red, red, eps, preg); + Div(m, m, red, preg); + Add(csum, csum, m, preg); +} + +__aicore__ inline void VFProcessCombFragPacked(const LocalTensor &combFragLocal, + const LocalTensor &mixLocal, + const LocalTensor &hcBaseLocal, + float scale, float eps, uint16_t iters, + uint16_t dim0, uint16_t hcMult, uint16_t hcMix) +{ + __local_mem__ float *combOutAddr = (__local_mem__ float *)combFragLocal.GetPhyAddr(); + __local_mem__ float *mixAddr = (__local_mem__ float *)mixLocal.GetPhyAddr(); + __local_mem__ float *hcBaseAddr = (__local_mem__ float *)hcBaseLocal.GetPhyAddr(); + + const uint16_t R = hcMult; + const uint16_t pack = VL_FP32 / R; + const uint16_t combLen = R * R; + const uint16_t nChunks = (dim0 + pack - 1) / pack; + __local_mem__ float *mBase = mixAddr; + __local_mem__ float *oBase = combOutAddr; + + __VEC_SCOPE__ + { + RegTensor mix0, mix1, mix2, mix3; + RegTensor base0, base1, base2, base3; + RegTensor red, csum; + RegTensor t0, t1, t2, t3; + RegTensor vL, vbs, vc, gIdx, sIdx, cShift, cMask; + + uint32_t uFull = VL_FP32; + MaskReg pIdx = UpdateMask(uFull); + + Arange(vL, 0); + Duplicate(cShift, (int32_t)2, pIdx); + Duplicate(cMask, (int32_t)(R - 1), pIdx); + ShiftRight(vbs, vL, cShift, pIdx); + And(vc, vL, cMask, pIdx); + Muls(gIdx, vbs, (int32_t)hcMix, pIdx); + Add(gIdx, gIdx, vc, pIdx); + Muls(sIdx, vbs, (int32_t)combLen, pIdx); + Add(sIdx, sIdx, vc, pIdx); + DataCopyGather(base0, hcBaseAddr + 0 * R, (RegTensor &)vc, pIdx); + DataCopyGather(base1, hcBaseAddr + 1 * R, (RegTensor &)vc, pIdx); + DataCopyGather(base2, hcBaseAddr + 2 * R, (RegTensor &)vc, pIdx); + DataCopyGather(base3, hcBaseAddr + 3 * R, (RegTensor &)vc, pIdx); + + uint32_t act = (uint32_t)dim0 * R; + for (uint16_t cIdx = 0; cIdx < nChunks; cIdx++) { + uint16_t i = cIdx * pack; + MaskReg preg = UpdateMask(act); + mBase = mixAddr + (uint32_t)i * hcMix; + oBase = combOutAddr + (uint32_t)i * combLen; + + DataCopyGather(mix0, mBase + 0 * R, (RegTensor &)gIdx, preg); + DataCopyGather(mix1, mBase + 1 * R, (RegTensor &)gIdx, preg); + DataCopyGather(mix2, mBase + 2 * R, (RegTensor &)gIdx, preg); + DataCopyGather(mix3, mBase + 3 * R, (RegTensor &)gIdx, preg); + Muls(mix0, mix0, scale, preg); Add(mix0, mix0, base0, preg); + Muls(mix1, mix1, scale, preg); Add(mix1, mix1, base1, preg); + Muls(mix2, mix2, scale, preg); Add(mix2, mix2, base2, preg); + Muls(mix3, mix3, scale, preg); Add(mix3, mix3, base3, preg); + + SoftmaxRowBcastInplace(mix0, red, t0, t1, t2, t3, eps, preg); + SoftmaxRowBcastInplace(mix1, red, t0, t1, t2, t3, eps, preg); + SoftmaxRowBcastInplace(mix2, red, t0, t1, t2, t3, eps, preg); + SoftmaxRowBcastInplace(mix3, red, t0, t1, t2, t3, eps, preg); + + Add(csum, mix0, mix1, preg); Add(csum, csum, mix2, preg); Add(csum, csum, mix3, preg); + Adds(csum, csum, eps, preg); + Div(mix0, mix0, csum, preg); Div(mix1, mix1, csum, preg); + Div(mix2, mix2, csum, preg); Div(mix3, mix3, csum, preg); + + for (uint16_t j = 0; j < iters; j++) { + Duplicate(csum, static_cast(0), preg); + RowNormAccum(mix0, red, csum, t0, t1, t2, t3, eps, preg); + RowNormAccum(mix1, red, csum, t0, t1, t2, t3, eps, preg); + RowNormAccum(mix2, red, csum, t0, t1, t2, t3, eps, preg); + RowNormAccum(mix3, red, csum, t0, t1, t2, t3, eps, preg); + Adds(csum, csum, eps, preg); + Div(mix0, mix0, csum, preg); Div(mix1, mix1, csum, preg); + Div(mix2, mix2, csum, preg); Div(mix3, mix3, csum, preg); + } + DataCopyScatter(oBase + 0 * R, mix0, (RegTensor &)sIdx, preg); + DataCopyScatter(oBase + 1 * R, mix1, (RegTensor &)sIdx, preg); + DataCopyScatter(oBase + 2 * R, mix2, (RegTensor &)sIdx, preg); + DataCopyScatter(oBase + 3 * R, mix3, (RegTensor &)sIdx, preg); + } + } +} + + +template +__aicore__ inline void VFProcessY(const LocalTensor &yLocal, const LocalTensor &mixLocal, + const LocalTensor &xLocal, uint16_t bs, uint16_t hcMult, uint16_t d, uint16_t hcMix) +{ + __local_mem__ T *yLocalAddr = (__local_mem__ T *)yLocal.GetPhyAddr(); + __local_mem__ float *mixLocalAddr = (__local_mem__ float *)mixLocal.GetPhyAddr(); + __local_mem__ T *xLocalAddr = (__local_mem__ T *)xLocal.GetPhyAddr(); + uint32_t dAlign = RoundUp(d); + uint16_t loopCount = CeilDiv(d, VL_FP32); + uint32_t hcMixAlign = RoundUp(hcMix); + if (loopCount > 1) { + __VEC_SCOPE__ + { + RegTensor x; + RegTensor mix; + RegTensor sum; + MaskReg pregLoop; + for (uint16_t i = 0; i < bs; i++) { + uint32_t sreg = d; + for (uint16_t j = 0; j < loopCount; j++) { + pregLoop = UpdateMask(sreg); + Duplicate(sum, static_cast(0), pregLoop); + for (uint16_t k = 0; k < hcMult; k++) { + LoadInputDataWithBrc(mix, mixLocalAddr, pregLoop, i * hcMixAlign + k); + LoadInputData(x, xLocalAddr, pregLoop, i * hcMult * dAlign + j * VL_FP32 + k * dAlign); + Mul(x, mix, x, pregLoop); + Add(sum, sum, x, pregLoop); + } + StoreOutputData(yLocalAddr, sum, pregLoop, i * dAlign + j * VL_FP32); + } + } + } + } else { + __VEC_SCOPE__ + { + RegTensor x; + RegTensor mix; + RegTensor sum; + uint32_t sreg = d; + MaskReg pregLoop = UpdateMask(sreg); + for (uint16_t i = 0; i < bs; i++) { + Duplicate(sum, static_cast(0), pregLoop); + for (uint16_t j = 0; j < hcMult; j++) { + LoadInputDataWithBrc(mix, mixLocalAddr, pregLoop, i * hcMixAlign + j); + LoadInputData(x, xLocalAddr, pregLoop, i * hcMult * dAlign + j * dAlign); + Mul(x, mix, x, pregLoop); + Add(sum, sum, x, pregLoop); + } + StoreOutputData(yLocalAddr, sum, pregLoop, i * dAlign); + } + } + } +} + +template +__aicore__ inline void CopyIn(const GlobalTensor &inputGm, const LocalTensor &inputTensor, const uint16_t nBurst, + const uint32_t copyLen, uint32_t srcStride = 0) +{ + DataCopyPadExtParams dataCopyPadExtParams; + dataCopyPadExtParams.isPad = false; + dataCopyPadExtParams.leftPadding = 0; + dataCopyPadExtParams.rightPadding = 0; + dataCopyPadExtParams.paddingValue = 0; + + DataCopyExtParams dataCoptExtParams; + dataCoptExtParams.blockCount = nBurst; + dataCoptExtParams.blockLen = copyLen * sizeof(T); + dataCoptExtParams.srcStride = srcStride * sizeof(T); + dataCoptExtParams.dstStride = 0; + DataCopyPad(inputTensor, inputGm, dataCoptExtParams, dataCopyPadExtParams); +} + +template +__aicore__ inline void CopyToL1(const LocalTensor &srcTensor, const LocalTensor &dstTensor, const DataCopyParams dataCopyXParams) +{ + DataCopy(dstTensor, srcTensor, dataCopyXParams); +} + +template +__aicore__ inline void CopyInWithLoopMode(const GlobalTensor &inputGm, const LocalTensor &inputTensor, + const uint16_t outerLoop, const uint16_t nBurst, const uint32_t copyLen, + const uint32_t gmLastDim, uint32_t srcStride = 0) +{ + uint16_t copyLenAlign = RoundUp(copyLen); + LoopModeParams loopParams; + loopParams.loop2Size = 1; + loopParams.loop1Size = outerLoop; + loopParams.loop2SrcStride = 0; + loopParams.loop1SrcStride = gmLastDim * sizeof(T); + loopParams.loop2DstStride = 0; + loopParams.loop1DstStride = nBurst * copyLenAlign * sizeof(T); + + DataCopyPadExtParams dataCopyPadExtParams; + dataCopyPadExtParams.isPad = false; + dataCopyPadExtParams.leftPadding = 0; + dataCopyPadExtParams.rightPadding = 0; + dataCopyPadExtParams.paddingValue = 0; + + DataCopyExtParams dataCoptExtParams; + dataCoptExtParams.blockCount = nBurst; + dataCoptExtParams.blockLen = copyLen * sizeof(T); + dataCoptExtParams.srcStride = srcStride * sizeof(T); + dataCoptExtParams.dstStride = 0; + SetLoopModePara(loopParams, DataCopyMVType::OUT_TO_UB); + DataCopyPad(inputTensor, inputGm, dataCoptExtParams, dataCopyPadExtParams); + ResetLoopModePara(DataCopyMVType::OUT_TO_UB); +} + +template +__aicore__ inline void CopyOut(const LocalTensor &outputTensor, const GlobalTensor &outputGm, + const uint16_t nBurst, const uint32_t copyLen, uint32_t dstStride = 0) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = nBurst; + dataCopyParams.blockLen = copyLen * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = dstStride * sizeof(T); + DataCopyPad(outputGm, outputTensor, dataCopyParams); +} + +template +__aicore__ inline void CopyOut(const LocalTensor &outputTensor, const LocalTensor &outputGm, + const uint16_t nBurst, const uint32_t copyLen) +{ + DataCopyParams dataCopyParams; + dataCopyParams.blockCount = nBurst; + dataCopyParams.blockLen = copyLen * sizeof(T) / BLOCK_SIZE; + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = 0; + DataCopy(outputGm, outputTensor, dataCopyParams); +} + + +} // namespace HcPreSinkhorn + +#endif diff --git a/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_cube_compute.h b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_cube_compute.h new file mode 100644 index 000000000..4502d3d93 --- /dev/null +++ b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_cube_compute.h @@ -0,0 +1,481 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_cube_compute.h + * \brief + */ +#ifndef HC_PRE_CUBE_COMPUTE_H +#define HC_PRE_CUBE_COMPUTE_H + +#include "kernel_operator.h" +#include "kernel_operator_intf.h" +#include "hc_pre_base.h" + +using AscendC::BLOCK_CUBE; +using AscendC::GlobalTensor; +using AscendC::HardEvent; +using AscendC::LocalTensor; +using AscendC::Nd2NzParams; +using AscendC::SetFlag; +using AscendC::TBuf; +using AscendC::TPipe; +using AscendC::TPosition; +using AscendC::WaitFlag; +using namespace AscendC; + +namespace HcPre { +struct MmParams { + uint64_t curML1; + uint64_t curKL1; + uint64_t curNL1; + uint64_t singleCoreK; + uint64_t kGmBaseOffset; + uint64_t nGmSize; + uint64_t kGmSize; + uint64_t nOutSize; + uint64_t xWsKSize; + bool isLastK; + bool isFirstK; +}; +#define HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM template + +#define HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS HcCubeCompute + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +class HcCubeCompute { +public: + __aicore__ inline HcCubeCompute(){}; + + __aicore__ inline void Init(const GlobalTensor& xGm, const GlobalTensor& fnGm, TPipe *tpipe); + __aicore__ inline void ComputeDecode(const AscendC::GlobalTensor &xGm, const AscendC::GlobalTensor &workspaceGlobalA2, + const AscendC::GlobalTensor &workspaceGlobalAB, const MmParams &mmParams); + __aicore__ inline void CopyInB1( + uint64_t mGmOffset, uint64_t kGmOffset, uint64_t kL1Size, const MmParams &mmParams); + __aicore__ inline void SetBL1Mte1ToMte2Flag(); + __aicore__ inline void WaitBL1Mte1ToMte2Flag(); + __aicore__ inline void End(); + +private: + __aicore__ inline void CopyInA1( + uint64_t kL1Size, + const GlobalTensor &aGlobal, const LocalTensor &al1Local, const MmParams &mmParams); + __aicore__ inline void CopyOut(const AscendC::GlobalTensor &workspaceGlobal, + const AscendC::LocalTensor &c1Local, uint64_t baseM, uint64_t baseN, bool enableNz2Nd, uint64_t N); + __aicore__ inline void Fixp(const AscendC::GlobalTensor &workspaceGlobalA2, + const AscendC::GlobalTensor &workspaceGlobalAB, const MmParams &mmParams); + __aicore__ inline void LoadAToL0A( + uint64_t kL1Offset, uint64_t kL0Size, uint64_t l1LoopIdx, const MmParams &mmParams); + __aicore__ inline void LoadAToL0B( + uint64_t kL1Offset, uint64_t kL0Size, uint64_t l1LoopIdx, const MmParams &mmParams); + __aicore__ inline void MmadA2(uint64_t kGmOffset, uint64_t kL0Size, bool isLastK, const MmParams &mmParams); + __aicore__ inline void LoadBToL0B( + uint64_t kL1Offset, uint64_t kL0Size, uint64_t l1LoopIdx, const MmParams &mmParams); + __aicore__ inline void MmadAB( + uint64_t kGmOffset, uint64_t kL0Size, bool isLastK, const MmParams &mmParams); + + TPipe *pipe_; + const HcPreTilingData *tiling_; + + int32_t blkIdx_ = -1; + int64_t batch_ = 0; + int64_t hcParam_ = 0; + int64_t dParam_ = 0; + + static constexpr int32_t ONE_BLOCK_SIZE = 32; + int32_t perBlock32 = ONE_BLOCK_SIZE / sizeof(float); + + GlobalTensor fnGm_; + GlobalTensor yGm_; + + static constexpr uint64_t MM1_MTE2_MTE1_EVENT = 2; + static constexpr uint64_t X_MTE1_MTE2_EVENT = 2; + static constexpr uint64_t B_MTE1_MTE2_EVENT = 4; + static constexpr uint64_t M_MTE1_EVENT_L0A = 3; + static constexpr uint64_t M_MTE1_EVENT_L0B = 5; + static constexpr uint64_t MTE1_M_EVENT = 2; + + static constexpr uint64_t L1_BUF_NUM = 2; + static constexpr uint64_t L0A_BUF_NUM = 2; + static constexpr uint64_t L0B_BUF_NUM = 2; + static constexpr uint64_t L0AB_BUF_NUM = 2; + static constexpr uint64_t L0C_BUF_NUM = 2; + + static constexpr uint64_t L1_BUF_OFFSET = 128 * 256; + static constexpr uint64_t L0AB_BUF_OFFSET = 32 * 256; + static constexpr uint64_t L0C_BUF_OFFSET = 64 * 256; + static constexpr uint64_t L0C_A2_BUF_OFFSET = 256 * 16; + + static constexpr uint16_t UNIT_FLAG_ENABLE = 2; + static constexpr uint16_t UNIT_FLAG_ENABLE_AUTO_CLOSE = 3; + constexpr static uint32_t FINAL_ACCUMULATION = 3; + constexpr static uint32_t NON_FINAL_ACCUMULATION = 2; + + static constexpr uint64_t K_L0_SIZE = 32UL; + static constexpr uint64_t FLOAT_C0_SIZE = 8UL; + + uint64_t l1aLoopIdx_ = 0; + uint64_t l1bLoopIdx_ = 0; + uint64_t l0aLoopIdx_ = 0; + uint64_t l0bLoopIdx_ = 0; + uint64_t l0cLoopIdx_ = 0; + + LocalTensor l1a_; + LocalTensor l1b_; + + LocalTensor l0a_; + LocalTensor l0b_; + LocalTensor l0c_; + + uint64_t k_ = 0; + uint64_t n_ = 0; +}; + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::Init(const GlobalTensor& xGm, const GlobalTensor& fnGm, TPipe *tpipe) +{ + fnGm_ = fnGm; + + TBuf l1aBuffer; + tpipe->InitBuffer(l1aBuffer, 256 * 1024); + l1a_ = l1aBuffer.Get(); + + TBuf l1bBuffer; + tpipe->InitBuffer(l1bBuffer, 256 * 1024); + l1b_ = l1bBuffer.Get(); + + TBuf l0aBuffer; + tpipe->InitBuffer(l0aBuffer, 64 * 1024); + l0a_ = l0aBuffer.Get(); + + TBuf l0bBuffer; + tpipe->InitBuffer(l0bBuffer, 64 * 1024); + l0b_ = l0bBuffer.Get(); + + // loc + TBuf l0cBuffer; + tpipe->InitBuffer(l0cBuffer, 64 * 1024); + l0c_ = l0cBuffer.Get(); + + for (int i = 0; i < L0A_BUF_NUM; i++) { + SetFlag(M_MTE1_EVENT_L0A + i); + SetFlag(M_MTE1_EVENT_L0B + i); + } + for (int i = 0; i < L1_BUF_NUM; i++) { + SetFlag(X_MTE1_MTE2_EVENT + i); + SetFlag(B_MTE1_MTE2_EVENT + i); + } +} + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::CopyInA1( + uint64_t kL1Size, + const GlobalTensor &aGlobal, const LocalTensor &al1Local, const MmParams &mmParams) +{ + Nd2NzParams nd2nzParams; + nd2nzParams.ndNum = 1; + nd2nzParams.nValue = mmParams.curML1; + nd2nzParams.dValue = kL1Size; + nd2nzParams.srcNdMatrixStride = 1; + nd2nzParams.srcDValue = mmParams.xWsKSize; // vec处理的singleK + nd2nzParams.dstNzC0Stride = (mmParams.curML1 + BLOCK_CUBE - 1) / BLOCK_CUBE * BLOCK_CUBE; + nd2nzParams.dstNzNStride = 1; + nd2nzParams.dstNzMatrixStride = 1; + DataCopy(al1Local, aGlobal, nd2nzParams); +} + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::CopyInB1( + uint64_t mGmOffset, uint64_t kGmOffset, uint64_t kL1Size, const MmParams &mmParams) +{ + AscendC::Nd2NzParams nd2nzParams; + nd2nzParams.ndNum = 1; + nd2nzParams.nValue = mmParams.curNL1; + nd2nzParams.dValue = kL1Size; + nd2nzParams.srcNdMatrixStride = 1; + nd2nzParams.srcDValue = mmParams.kGmSize; // 原始k + nd2nzParams.dstNzC0Stride = (mmParams.curNL1 + BLOCK_CUBE - 1) / BLOCK_CUBE * BLOCK_CUBE; + nd2nzParams.dstNzNStride = 1; + nd2nzParams.dstNzMatrixStride = 1; + DataCopy(l1b_[(l1bLoopIdx_ % L1_BUF_NUM) * L1_BUF_OFFSET], fnGm_[kGmOffset], nd2nzParams); +} + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::CopyOut(const AscendC::GlobalTensor &workspaceGlobal, + const AscendC::LocalTensor &c1Local, uint64_t baseM, uint64_t baseN, bool enableNz2Nd, uint64_t N) +{ + AscendC::DataCopyCO12DstParams intriParams; + intriParams.nSize = baseN; + intriParams.mSize = baseM; + // set mode to float32, then cast in ub + intriParams.quantPre = QuantMode_t::NoQuant; + intriParams.nz2ndEn = enableNz2Nd; + if (enableNz2Nd) { + intriParams.dstStride = N; // NZ -> ND + intriParams.srcStride = CeilAlign(baseM, AscendC::BLOCK_CUBE); + AscendC::SetFixpipeNz2ndFlag(1, 1, 1); + } else { + intriParams.dstStride = CeilAlign(intriParams.nSize, AscendC::BLOCK_CUBE); // NZ -> NZ + intriParams.srcStride = CeilAlign(baseM, AscendC::BLOCK_CUBE); + } + intriParams.unitFlag = UNIT_FLAG_ENABLE_AUTO_CLOSE; // 3 + + AscendC::DataCopy(workspaceGlobal, c1Local, intriParams); +} + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::Fixp(const AscendC::GlobalTensor &workspaceGlobalA2, + const AscendC::GlobalTensor &workspaceGlobalAB, const MmParams &mmParams) +{ + // Copy MmadA2 + CopyOut(workspaceGlobalA2, + l0c_[(l0cLoopIdx_ % L0C_BUF_NUM) * L0C_BUF_OFFSET], + mmParams.curML1, + BLOCK_CUBE, + false, + BLOCK_CUBE); // nz m,16 + // Copy MmadAB + CopyOut(workspaceGlobalAB, + l0c_[(l0cLoopIdx_ % L0C_BUF_NUM) * L0C_BUF_OFFSET + L0C_A2_BUF_OFFSET], + mmParams.curML1, + mmParams.curNL1, + true, + mmParams.nOutSize); // ND M/N are 512B aligned. +} + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::SetBL1Mte1ToMte2Flag() +{ + SetFlag(B_MTE1_MTE2_EVENT + l1bLoopIdx_ % L1_BUF_NUM); + l1bLoopIdx_++; +} + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::WaitBL1Mte1ToMte2Flag() +{ + WaitFlag(B_MTE1_MTE2_EVENT + l1bLoopIdx_ % L1_BUF_NUM); +} + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::ComputeDecode( + const AscendC::GlobalTensor &xGm, + const AscendC::GlobalTensor &workspaceGlobalA2, const AscendC::GlobalTensor &workspaceGlobalAB, + const MmParams &mmParams) +{ + WaitFlag(X_MTE1_MTE2_EVENT + l1aLoopIdx_ % L1_BUF_NUM); + uint64_t kGmOffset = 0; + uint64_t curKL1Size = (kGmOffset + mmParams.curKL1 >= mmParams.singleCoreK) ? (mmParams.singleCoreK - kGmOffset) : mmParams.curKL1; + CopyInA1(curKL1Size, xGm[kGmOffset], l1a_[(l1aLoopIdx_ % DOUBLE_BUFFER) * L1_BUF_OFFSET], mmParams); + + SetFlag(MM1_MTE2_MTE1_EVENT + l1aLoopIdx_ % L1_BUF_NUM); + kGmOffset += mmParams.curKL1; + while (kGmOffset < mmParams.singleCoreK) { + WaitFlag(X_MTE1_MTE2_EVENT + (l1aLoopIdx_ + 1) % L1_BUF_NUM); + uint64_t nextKL1Size = (kGmOffset + mmParams.curKL1 >= mmParams.singleCoreK) ? (mmParams.singleCoreK - kGmOffset) : mmParams.curKL1; + CopyInA1(nextKL1Size, xGm[kGmOffset], l1a_[((l1aLoopIdx_ + 1) % DOUBLE_BUFFER) * L1_BUF_OFFSET], mmParams); + + SetFlag(MM1_MTE2_MTE1_EVENT + (l1aLoopIdx_ + 1) % L1_BUF_NUM); + WaitFlag(MM1_MTE2_MTE1_EVENT + l1aLoopIdx_ % L1_BUF_NUM); + + for (uint64_t kL1Offset = 0; kL1Offset < curKL1Size; kL1Offset += K_L0_SIZE) { // K_L0_SIZE 32 + WaitFlag(M_MTE1_EVENT_L0A + l0aLoopIdx_ % L0A_BUF_NUM); + LoadAToL0A(kL1Offset, K_L0_SIZE, l1aLoopIdx_, mmParams); + + WaitFlag(M_MTE1_EVENT_L0B + l0bLoopIdx_ % L0B_BUF_NUM); + LoadAToL0B(kL1Offset, K_L0_SIZE, l1aLoopIdx_, mmParams); + MmadA2(kGmOffset + kL1Offset - mmParams.curKL1, K_L0_SIZE, + false, + mmParams); + SetFlag(M_MTE1_EVENT_L0B + l0bLoopIdx_ % L0B_BUF_NUM); + l0bLoopIdx_++; + + WaitFlag(M_MTE1_EVENT_L0B + l0bLoopIdx_ % L0B_BUF_NUM); + LoadBToL0B((kGmOffset + kL1Offset - mmParams.curKL1) % mmParams.xWsKSize, K_L0_SIZE, l1bLoopIdx_, mmParams); + MmadAB(kGmOffset + kL1Offset - mmParams.curKL1, K_L0_SIZE, + false, + mmParams); + SetFlag(M_MTE1_EVENT_L0B + l0bLoopIdx_ % L0B_BUF_NUM); + l0bLoopIdx_++; + + SetFlag(M_MTE1_EVENT_L0A + l0aLoopIdx_ % L0A_BUF_NUM); + l0aLoopIdx_++; + } + SetFlag(X_MTE1_MTE2_EVENT + l1aLoopIdx_ % L1_BUF_NUM); + l1aLoopIdx_++; + kGmOffset += mmParams.curKL1; + curKL1Size = nextKL1Size; + } + + WaitFlag(MM1_MTE2_MTE1_EVENT + l1aLoopIdx_ % L1_BUF_NUM); + for (uint64_t kL1Offset = 0; kL1Offset < curKL1Size; kL1Offset += K_L0_SIZE) { + WaitFlag(M_MTE1_EVENT_L0A + l0aLoopIdx_ % L0A_BUF_NUM); + LoadAToL0A(kL1Offset, K_L0_SIZE, l1aLoopIdx_, mmParams); // to l0a + + WaitFlag(M_MTE1_EVENT_L0B + l0bLoopIdx_ % L0B_BUF_NUM); + LoadAToL0B(kL1Offset, K_L0_SIZE, l1aLoopIdx_, mmParams); // to l0b + MmadA2(kGmOffset + kL1Offset - mmParams.curKL1, K_L0_SIZE, + mmParams.isLastK && kL1Offset + K_L0_SIZE >= curKL1Size, mmParams); + SetFlag(M_MTE1_EVENT_L0B + l0bLoopIdx_ % L0B_BUF_NUM); + l0bLoopIdx_++; + + WaitFlag(M_MTE1_EVENT_L0B + l0bLoopIdx_ % L0B_BUF_NUM); + LoadBToL0B((kGmOffset + kL1Offset - mmParams.curKL1) % mmParams.xWsKSize, K_L0_SIZE, l1bLoopIdx_, mmParams); // to l0b + MmadAB(kGmOffset + kL1Offset - mmParams.curKL1, K_L0_SIZE, + mmParams.isLastK && kL1Offset + K_L0_SIZE >= curKL1Size, + mmParams); + SetFlag(M_MTE1_EVENT_L0B + l0bLoopIdx_ % L0B_BUF_NUM); + l0bLoopIdx_++; + + SetFlag(M_MTE1_EVENT_L0A + l0aLoopIdx_ % L0A_BUF_NUM); + l0aLoopIdx_++; + } + SetFlag(X_MTE1_MTE2_EVENT + l1aLoopIdx_ % L1_BUF_NUM); + l1aLoopIdx_++; + if (mmParams.isLastK) { + Fixp(workspaceGlobalA2, workspaceGlobalAB, mmParams); // l0cLoopIdx_++; + } +} + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::End() +{ + for (int i = 0; i < L0A_BUF_NUM; i++) { + WaitFlag(M_MTE1_EVENT_L0A + i); + WaitFlag(M_MTE1_EVENT_L0B + i); + } + for (int i = 0; i < L1_BUF_NUM; i++) { + WaitFlag(X_MTE1_MTE2_EVENT + i); + WaitFlag(B_MTE1_MTE2_EVENT + i); + } +} + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::LoadAToL0A( + uint64_t kL1Offset, uint64_t kL0Size, uint64_t l1LoopIdx, const MmParams &mmParams) +{ + static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; + LoadData3DParamsV2 loadData3DParams; + // SetFmatrixParams + loadData3DParams.l1H = CeilDiv(mmParams.curML1, BLOCK_CUBE); // Hin=M1=8 + loadData3DParams.l1W = BLOCK_CUBE; // Win=M0 + loadData3DParams.channelSize = kL0Size; // Cin=K + + loadData3DParams.padList[0] = 0; + loadData3DParams.padList[1] = 0; + loadData3DParams.padList[2] = 0; + loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果 + + // SetLoadToA0Params + loadData3DParams.mExtension = CeilAlign(mmParams.curML1, BLOCK_CUBE); // M height维度目的 + loadData3DParams.kExtension = kL0Size; // K width维度目的 + loadData3DParams.mStartPt = 0; + loadData3DParams.kStartPt = 0; + loadData3DParams.strideW = 1; + loadData3DParams.strideH = 1; + loadData3DParams.filterW = 1; + loadData3DParams.filterSizeW = (1 >> 8) & 255; + loadData3DParams.filterH = 1; + loadData3DParams.filterSizeH = (1 >> 8) & 255; + loadData3DParams.dilationFilterW = 1; + loadData3DParams.dilationFilterH = 1; + loadData3DParams.enTranspose = 0; + loadData3DParams.fMatrixCtrl = 0; + + LoadData(l0a_[(l0aLoopIdx_ % L0AB_BUF_NUM) * L0AB_BUF_OFFSET], + l1a_[(l1LoopIdx % L1_BUF_NUM) * L1_BUF_OFFSET + + CeilAlign(mmParams.curML1, static_cast(BLOCK_CUBE)) * kL1Offset], + loadData3DParams); +} + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::LoadAToL0B( + uint64_t kL1Offset, uint64_t kL0Size, uint64_t l1LoopIdx, const MmParams &mmParams) +{ + // mk nz -> m,k zz + for (uint64_t mL0Offset = 0; mL0Offset < mmParams.curML1; mL0Offset += BLOCK_CUBE) { + LoadData2DParams l1ToL0bParams; + l1ToL0bParams.startIndex = 0; + l1ToL0bParams.repeatTimes = CeilDiv(kL0Size, (uint64_t)BLOCK_CUBE >> 1); + l1ToL0bParams.srcStride = CeilDiv(mmParams.curML1, (uint64_t)BLOCK_CUBE); + l1ToL0bParams.dstGap = 0; + LoadData(l0b_[(l0bLoopIdx_ % L0AB_BUF_NUM) * L0AB_BUF_OFFSET + + mL0Offset * CeilAlign(kL0Size, (uint64_t)BLOCK_CUBE >> 1)], + l1a_[(l1LoopIdx % L1_BUF_NUM) * L1_BUF_OFFSET + + CeilAlign(mmParams.curML1, static_cast(BLOCK_CUBE)) * kL1Offset + + mL0Offset * (uint64_t)(BLOCK_CUBE >> 1)], + l1ToL0bParams); + } +} + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::MmadA2( + uint64_t kGmOffset, uint64_t kL0Size, bool isLastK, const MmParams &mmParams) +{ + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + for (uint64_t mL0Offset = 0; mL0Offset < mmParams.curML1; mL0Offset += BLOCK_CUBE) { + MmadParams mmadParams; + mmadParams.m = BLOCK_CUBE; + mmadParams.n = BLOCK_CUBE; + mmadParams.k = kL0Size; + mmadParams.cmatrixInitVal = mmParams.isFirstK && kGmOffset == 0; + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = isLastK ? UNIT_FLAG_ENABLE_AUTO_CLOSE : UNIT_FLAG_ENABLE; + // mk zz @ mk zz + Mmad(l0c_[(l0cLoopIdx_ % L0C_BUF_NUM) * L0C_BUF_OFFSET + mL0Offset * BLOCK_CUBE], + l0a_[(l0aLoopIdx_ % L0A_BUF_NUM) * L0AB_BUF_OFFSET + mL0Offset * CeilAlign(kL0Size, 8)], + l0b_[(l0bLoopIdx_ % L0AB_BUF_NUM) * L0AB_BUF_OFFSET + mL0Offset * CeilAlign(kL0Size, 8)], + mmadParams); + } +} + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::LoadBToL0B( + uint64_t kL1Offset, uint64_t kL0Size, uint64_t l1LoopIdx, const MmParams &mmParams) +{ + LoadData2DParams l1ToL0bParams; + l1ToL0bParams.startIndex = 0; + l1ToL0bParams.repeatTimes = + CeilDiv(mmParams.curNL1, (uint64_t)BLOCK_CUBE) * CeilDiv(kL0Size, (uint64_t)BLOCK_CUBE >> 1); + l1ToL0bParams.srcStride = 1; + l1ToL0bParams.dstGap = 0; + // n,k nz -> n,k nz + LoadData(l0b_[(l0bLoopIdx_ % L0AB_BUF_NUM) * L0AB_BUF_OFFSET], + l1b_[(l1LoopIdx % L1_BUF_NUM) * L1_BUF_OFFSET + + CeilAlign(mmParams.curNL1, static_cast(BLOCK_CUBE)) * kL1Offset], + l1ToL0bParams); +} + +HC_PRE_CUBE_COMPUTE_TEMPLATE_PARAM +__aicore__ inline void HC_PRE_CUBE_COMPUTE_TEMPLATE_CLASS::MmadAB( + uint64_t kGmOffset, uint64_t kL0Size, bool isLastK, const MmParams &mmParams) +{ + SetFlag(MTE1_M_EVENT); + WaitFlag(MTE1_M_EVENT); + AscendC::SetHF32Mode(1); + AscendC::SetHF32TransMode(1); + MmadParams mmadParams; + mmadParams.m = CeilAlign(mmParams.curML1, BLOCK_CUBE); + mmadParams.n = mmParams.curNL1; + mmadParams.k = kL0Size; // kl0Size + mmadParams.cmatrixInitVal = mmParams.isFirstK && kGmOffset == 0; + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = isLastK ? UNIT_FLAG_ENABLE_AUTO_CLOSE : UNIT_FLAG_ENABLE; + // mk zz @ nk nz + Mmad(l0c_[(l0cLoopIdx_ % L0C_BUF_NUM) * L0C_BUF_OFFSET + L0C_A2_BUF_OFFSET], + l0a_[(l0aLoopIdx_ % L0A_BUF_NUM) * L0AB_BUF_OFFSET], + l0b_[(l0bLoopIdx_ % L0AB_BUF_NUM) * L0AB_BUF_OFFSET], + mmadParams); + AscendC::SetHF32Mode(0); +} + +} // namespace HcPre + +#endif // HC_PRE_CUBE_COMPUTE_H diff --git a/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_cube_compute_arch35.h b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_cube_compute_arch35.h new file mode 100644 index 000000000..ae0e51ff3 --- /dev/null +++ b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_cube_compute_arch35.h @@ -0,0 +1,241 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_cube_compute_arch35.h + * \brief + */ +#ifndef HC_PRE_CUBE_COMPUTE_ARCH35_H +#define HC_PRE_CUBE_COMPUTE_ARCH35_H + +#include "kernel_operator.h" +#include "hc_pre_base_arch35.h" + +namespace HcPreNs +{ + using namespace AscendC; + constexpr static uint32_t FINAL_ACCUMULATION = 3; + constexpr static uint32_t NON_FINAL_ACCUMULATION = 2; + // constexpr static int32_t C0_SIZE = AscendC::AuxGetC0Size(); + constexpr static bool splitM_ = true; + constexpr static uint64_t SPLIT_M_ALIGN = 2; + static constexpr uint64_t L1_ALLOC_SIZE = 512 * 1024; + static constexpr uint64_t L1_BUF_NUM = 2; + static constexpr uint64_t L1_BUF_OFFSET = 128 * 256; + constexpr static int SYNC_MODE4 = 4; + + class HcPreCubeCompute + { + public: + uint64_t m_{0}; + uint64_t n_{0}; + uint64_t k_{0}; + uint64_t baseM_{0}; + uint64_t baseN_{0}; + uint64_t baseK_{0}; + uint64_t kL1_{0}; + + public: + AscendC::LocalTensor aL0Ping_; + AscendC::LocalTensor aL0Pong_; + AscendC::LocalTensor bL0Ping_; + AscendC::LocalTensor bL0Pong_; + AscendC::LocalTensor cL0Ping_; + AscendC::LocalTensor cL0Pong_; + uint8_t bL1BufferID_{0}; + uint8_t l0PingPongID_{0}; + uint8_t crossPingPongID_{0}; + uint8_t cl0PingPongID_{0}; + + __aicore__ inline HcPreCubeCompute() + { + } + + __aicore__ inline uint8_t GetBL1BufferId() + { + return bL1BufferID_; + } + + __aicore__ inline void Init() + { + // L0 空间的分配 + uint32_t aL0OneBuffer = 256 * 32; + uint32_t bL0OneBuffer = 256 * 32; + uint32_t cL0OneBuffer = 256 * 128; + + aL0Ping_ = AscendC::LocalTensor(AscendC::TPosition::A2, 0, aL0OneBuffer); + aL0Pong_ = AscendC::LocalTensor(AscendC::TPosition::A2, aL0OneBuffer * sizeof(float), aL0OneBuffer); + bL0Ping_ = AscendC::LocalTensor(AscendC::TPosition::B2, 0, bL0OneBuffer); + bL0Pong_ = AscendC::LocalTensor(AscendC::TPosition::B2, bL0OneBuffer * sizeof(float), bL0OneBuffer); + cL0Ping_ = AscendC::LocalTensor(AscendC::TPosition::CO1, 0, cL0OneBuffer); + cL0Pong_ = AscendC::LocalTensor(AscendC::TPosition::CO1, cL0OneBuffer * sizeof(float), cL0OneBuffer); + // 同步 + // B 的 gm2L1 的 pingpong id + AscendC::SetFlag(0); + AscendC::SetFlag(1); + // l12l0a & l12l0b 的 pingpong id (共用) + AscendC::SetFlag(3); + AscendC::SetFlag(4); + + // hf32 compute + AscendC::SetHF32Mode(1); + AscendC::SetHF32TransMode(1); + } + + __aicore__ inline void CopyInB1Nd2Nz(uint64_t k, uint64_t currentK, uint64_t baseN, const AscendC::GlobalTensor &bGlobal, + const AscendC::LocalTensor &bl1Local) + { + AscendC::WaitFlag(bL1BufferID_); + k_ = k; + kL1_ = currentK; + baseN_ = baseN; + AscendC::Nd2NzParams nd2nzParam; + nd2nzParam.ndNum = 1; + nd2nzParam.srcNdMatrixStride = 1; + nd2nzParam.dstNzNStride = 1; + nd2nzParam.dstNzMatrixStride = 1; + nd2nzParam.nValue = baseN_; + nd2nzParam.dValue = kL1_; + nd2nzParam.srcDValue = k_; + nd2nzParam.dstNzC0Stride = (baseN_ + AscendC::BLOCK_CUBE - 1) / AscendC::BLOCK_CUBE * AscendC::BLOCK_CUBE; + AscendC::DataCopy(bl1Local, bGlobal, nd2nzParam); + AscendC::SetFlag(bL1BufferID_); + AscendC::WaitFlag(bL1BufferID_); + } + + // note: baseM * baseK must <= 256 * 32; baseN * baseK must <= 256 * 32; baseM * baseN must <= 256 * 128 + __aicore__ inline void Process(uint64_t m, uint64_t n, uint64_t baseM, uint64_t baseK, + bool isFirstKL1, bool isLastKL1, const AscendC::LocalTensor &al1Local, + const AscendC::LocalTensor &bl1Local) + { + m_ = m; + n_ = n; + baseM_ = baseM; + baseK_ = baseK; + uint64_t kL1Offset = 0; + for (uint64_t kb = 0; kb < kL1_; kb += baseK_) + { + bool isLastKL0 = (kb + baseK_) >= kL1_; + AscendC::WaitFlag(l0PingPongID_ + 3); + CopyInA2(kb, kL1Offset, al1Local); + CopyInB2(kb, kL1Offset, bl1Local); + AscendC::SetFlag(l0PingPongID_); + AscendC::WaitFlag(l0PingPongID_); + MmadBase(kb, isFirstKL1, isLastKL1, isLastKL0); + AscendC::SetFlag(l0PingPongID_ + 3); + l0PingPongID_ = l0PingPongID_ ^ 1; + kL1Offset += baseK_; + } + AscendC::SetFlag(bL1BufferID_); + bL1BufferID_ ^= 1; + } + + __aicore__ inline void CopyInA2(uint64_t kOffset, uint64_t kAL1Offset, const AscendC::LocalTensor &al1Local_) + { + uint64_t mAL1 = Align(baseM_, AscendC::BLOCK_CUBE); + uint64_t offsetAL1 = Align(kAL1Offset, C0_SIZE) * mAL1; + AscendC::LoadData2DParamsV2 loadData2dParams; + + uint64_t currM = baseM_; + uint64_t currK = AscendC::Std::min(baseK_, kL1_ - kOffset); + loadData2dParams.mStartPosition = 0; + loadData2dParams.kStartPosition = 0; + loadData2dParams.mStep = CeilDiv(currM, AscendC::BLOCK_CUBE); + loadData2dParams.kStep = CeilDiv(currK, C0_SIZE); + loadData2dParams.srcStride = CeilDiv(currM, AscendC::BLOCK_CUBE); + loadData2dParams.dstStride = loadData2dParams.mStep; + loadData2dParams.ifTranspose = false; + AscendC::LoadData(l0PingPongID_ == 0 ? aL0Ping_ : aL0Pong_, al1Local_[offsetAL1], loadData2dParams); + } + + __aicore__ inline void CopyInB2(uint64_t kOffset, uint64_t kBL1Offset, const AscendC::LocalTensor &bl1Local_) + { + uint64_t nBL1 = Align(baseN_, AscendC::BLOCK_CUBE); + uint64_t offsetBL1 = Align(kBL1Offset, C0_SIZE) * nBL1; + AscendC::LoadData2DParamsV2 loadData2dParams; + + uint64_t currN = baseN_; + uint64_t currK = AscendC::Std::min(baseK_, kL1_ - kOffset); + loadData2dParams.mStartPosition = 0; + loadData2dParams.kStartPosition = 0; + loadData2dParams.mStep = CeilDiv(currN, AscendC::BLOCK_CUBE); + loadData2dParams.kStep = CeilDiv(currK, C0_SIZE); + loadData2dParams.srcStride = CeilDiv(currN, AscendC::BLOCK_CUBE); + loadData2dParams.dstStride = loadData2dParams.mStep; + loadData2dParams.ifTranspose = false; + AscendC::LoadData(l0PingPongID_ == 0 ? bL0Ping_ : bL0Pong_, bl1Local_[offsetBL1], loadData2dParams); + } + + __aicore__ inline void MmadBase(uint64_t kOffset, bool isFirstKL1, bool isLastKL1, bool isLastKL0) + { + uint32_t mmadK = AscendC::Std::min(baseK_, kL1_ - kOffset); + AscendC::MmadParams mmadParams; + mmadParams.m = baseM_; + mmadParams.n = baseN_; + + mmadParams.k = mmadK; + mmadParams.disableGemv = true; + mmadParams.cmatrixInitVal = (isFirstKL1 && kOffset == 0); // kOffset == 0: isFirstKL0 + mmadParams.cmatrixSource = false; + mmadParams.unitFlag = (isLastKL1 && isLastKL0) ? FINAL_ACCUMULATION : NON_FINAL_ACCUMULATION; + AscendC::Mmad(cl0PingPongID_ == 0 ? cL0Ping_ : cL0Pong_, l0PingPongID_ == 0 ? aL0Ping_ : aL0Pong_, + l0PingPongID_ == 0 ? bL0Ping_ : bL0Pong_, mmadParams); + } + + // fixpipe CopyOut实现c01拷贝到UB + __aicore__ inline void CopyOut(const AscendC::LocalTensor& dstLocal) + { + AscendC::FixpipeParamsC310 fixpipeParams; // ROW_MAJOR默认使能NZ2ND + uint64_t c0 = AscendC::AuxGetC0Size(); + fixpipeParams.nSize = Align(baseN_, c0); + fixpipeParams.mSize = splitM_ ? Align(baseM_, SPLIT_M_ALIGN) : baseM_; // 切m需要m是2对齐 + fixpipeParams.dstStride = fixpipeParams.nSize; + fixpipeParams.srcStride = Align(baseM_, AscendC::BLOCK_CUBE); // 单位CO_SIZE (16*sizeof(C_T)) + + fixpipeParams.quantPre = QuantMode_t::NoQuant; + // fixpipeParams.quantPre = 0; + // set cvRatio=1:2 默认splitM + fixpipeParams.dualDstCtl = splitM_ ? static_cast(AscendC::McgShfMode::DUAL_DST_SPLIT_M) : 0; + fixpipeParams.unitFlag = FINAL_ACCUMULATION; // 3 unitflag + fixpipeParams.params.ndNum = 1; // ndNum + fixpipeParams.params.srcNdStride = 1; // srcNdStride + fixpipeParams.params.dstNdStride = 1; // dstNdStride + AscendC::Fixpipe(dstLocal, cl0PingPongID_ == 0 ? cL0Ping_ : cL0Pong_, fixpipeParams); + cl0PingPongID_ ^= 1; + } + + // fixpipe CopyOut实现c01拷贝到GM + __aicore__ inline void CopyOut(const AscendC::GlobalTensor &cGlobal) + { + AscendC::DataCopyCO12DstParams intriParams; + intriParams.nSize = baseN_; + intriParams.mSize = baseM_; + intriParams.dstStride = n_; + intriParams.srcStride = Align(baseM_, AscendC::BLOCK_CUBE); + // set mode according to dtype + intriParams.quantPre = QuantMode_t::NoQuant; + intriParams.nz2ndEn = true; + intriParams.unitFlag = FINAL_ACCUMULATION; // 3 unitflag + AscendC::SetFixpipeNz2ndFlag(1, 1, 1); + AscendC::DataCopy(cGlobal, cl0PingPongID_ == 0 ? cL0Ping_ : cL0Pong_, intriParams); + cl0PingPongID_ ^= 1; + } + + __aicore__ inline void End() + { + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + AscendC::WaitFlag(3); + AscendC::WaitFlag(4); + AscendC::SetHF32Mode(0); + } + }; +} +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_m_k_split_core.h b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_m_k_split_core.h new file mode 100644 index 000000000..412b06f4e --- /dev/null +++ b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_m_k_split_core.h @@ -0,0 +1,564 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_m_split_core.h + * \brief + */ + +#ifndef HC_PRE_M_K_SPLIT_A3_CORE_H +#define HC_PRE_M_K_SPLIT_A3_CORE_H + +#include "kernel_operator.h" +#include "hc_pre_base.h" +#include "hc_pre_cube_compute.h" + +namespace HcPre { +using namespace AscendC; +template +class HcPreMembaseKSplitCorePart1 { +public: + __aicore__ inline HcPreMembaseKSplitCorePart1() + { + } + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR hcFn, GM_ADDR workspace, + const HcPreTilingData *tilingDataPtr, TPipe *pipePtr) + { + pipe = pipePtr; + tilingData = tilingDataPtr; + + xGm.SetGlobalBuffer((__gm__ T *)x); + hcFnGm.SetGlobalBuffer((__gm__ float *)hcFn); + workspaceGm.SetGlobalBuffer((__gm__ float *)workspace); + + uint64_t curVectorBlockIdx = GetBlockIdx(); + uint64_t curCubeBlockIdx = curVectorBlockIdx; + if ASCEND_IS_AIV { + curCubeBlockIdx = curCubeBlockIdx / CV_RATIO; + } + + xCastFp32BufSize_ = tilingData->mL1Size * + CeilAlign(tilingData->cvLoopKSize, MM_CACHE_LINE_BYTES / sizeof(float)); + int64_t SingleCubeCoreXCastSize = xCastFp32BufSize_ * DOUBLE_BUFFER; + xCastFp32WsGm.SetGlobalBuffer((__gm__ float *)workspace + curCubeBlockIdx * SingleCubeCoreXCastSize); + + uint64_t wsOffset = tilingData->cubeCoreNum * SingleCubeCoreXCastSize; + mmOutFp32WsGm.SetGlobalBuffer((__gm__ float *)workspace + wsOffset); + mmOuterInnerSize_ = CeilAlign(N_SIZE, MM_CACHE_LINE_BYTES / sizeof(float)); + mmOutFp32BufSize_ = tilingData->bs * mmOuterInnerSize_; + wsOffset += CeilAlign(tilingData->cubeBlockDimK * mmOutFp32BufSize_ * sizeof(float), + static_cast(WORKSPACE_ALIGN_SIZE)) / sizeof(float); + squareSumFp32WsGm.SetGlobalBuffer((__gm__ float *)workspace + wsOffset); + squareSumFp32BufSize_ = CeilAlign(tilingData->bs, BLOCK_CUBE) * BLOCK_CUBE; + + if ASCEND_IS_AIC { + cubeCompute_.Init(xCastFp32WsGm, hcFnGm, pipePtr); + return; + } + // InQue + int64_t xQueNum = tilingData->stage1MFactor * RoundUp(tilingData->cvLoopKSize); + pipe->InitBuffer(xQue, NUM_TWO, xQueNum * sizeof(T)); + + // OutQue + pipe->InitBuffer(mmInQue, NUM_TWO, xQueNum * sizeof(float)); + } + + __aicore__ inline void Process() + { + if ASCEND_IS_AIC{ + // 初始设置 + CrossCoreSetFlag(SYNC_AIC_TO_AIV_FLAG); + CrossCoreSetFlag(SYNC_AIC_TO_AIV_FLAG); + } + + uint64_t curBlockIdx = GetBlockIdx(); + uint64_t curVectorBlockIdx = curBlockIdx; + if ASCEND_IS_AIV { + curBlockIdx = curBlockIdx / NUM_TWO; + } + uint64_t mBlkDimIdx = curBlockIdx / tilingData->cubeBlockDimK; + uint64_t kBlkDimIdx = curBlockIdx % tilingData->cubeBlockDimK; + uint64_t kGmStartOffset = tilingData->multCoreSplitKSize * kBlkDimIdx; + uint64_t kGmEndOffset = kGmStartOffset + tilingData->multCoreSplitKSize; + if (kGmEndOffset > tilingData->k) { + kGmEndOffset = tilingData->k; + } + + uint64_t mGmBaseOffset = mBlkDimIdx * tilingData->multCoreSplitMSize; + uint64_t mGmEndOffset = mGmBaseOffset + tilingData->multCoreSplitMSize; + if (mGmEndOffset > tilingData->bs) { + mGmEndOffset = tilingData->bs; + } + for (uint64_t mGmOffset = mGmBaseOffset; mGmOffset < mGmEndOffset; + mGmOffset += tilingData->mL1Size) { + uint64_t realMSize = mGmOffset + tilingData->mL1Size > mGmEndOffset ? + mGmEndOffset - mGmOffset : tilingData->mL1Size; + for (uint64_t kGmBaseOffset = kGmStartOffset; kGmBaseOffset < kGmEndOffset; + kGmBaseOffset += tilingData->cvLoopKSize) { + uint64_t realKGmSize = kGmBaseOffset + tilingData->cvLoopKSize > kGmEndOffset ? + kGmEndOffset - kGmBaseOffset : tilingData->cvLoopKSize; + if ASCEND_IS_AIC{ + MmParams mmParams; + mmParams.curML1 = realMSize; + mmParams.curKL1 = tilingData->kL1Size; + mmParams.curNL1 = N_SIZE; + mmParams.singleCoreK = realKGmSize; + mmParams.xWsKSize = tilingData->cvLoopKSize; + mmParams.nOutSize = mmOuterInnerSize_; + mmParams.kGmBaseOffset = kGmBaseOffset; + mmParams.nGmSize = N_SIZE; + mmParams.kGmSize = tilingData->k; + mmParams.isLastK = kGmBaseOffset + tilingData->cvLoopKSize >= kGmEndOffset; + mmParams.isFirstK = kGmBaseOffset == kGmStartOffset; + + cubeCompute_.WaitBL1Mte1ToMte2Flag(); + cubeCompute_.CopyInB1(mGmOffset, kGmBaseOffset, realKGmSize, mmParams); + + CrossCoreWaitFlag(SYNC_AIV_TO_AIC_FLAG); + cubeCompute_.ComputeDecode(xCastFp32WsGm[cvLoopIdx_ % DOUBLE_BUFFER * xCastFp32BufSize_], + squareSumFp32WsGm[kBlkDimIdx * squareSumFp32BufSize_ + mGmOffset * BLOCK_CUBE], + mmOutFp32WsGm[kBlkDimIdx * mmOutFp32BufSize_ + mGmOffset * mmOuterInnerSize_], + mmParams); + cubeCompute_.SetBL1Mte1ToMte2Flag(); + CrossCoreSetFlag(SYNC_AIC_TO_AIV_FLAG); + } else { + CrossCoreWaitFlag(SYNC_AIC_TO_AIV_FLAG); + // vec.compute(); + int64_t mVectorOffset = mGmOffset; + int64_t mVectorLength = (realMSize + 1) / NUM_TWO; + if ((curVectorBlockIdx % NUM_TWO) == 1) { + mVectorOffset = mGmOffset + (realMSize + 1) / NUM_TWO; + mVectorLength = realMSize - mVectorLength; + } + int64_t curUbLoops = (mVectorLength + tilingData->stage1MFactor - 1) / tilingData->stage1MFactor; + int64_t curUbMfactorTail = mVectorLength - ((curUbLoops - 1) * tilingData->stage1MFactor); + for (int64_t i = 0; i < curUbLoops; ++i) { + int64_t curUbMFactor = (i != (curUbLoops - 1)) ? tilingData->stage1MFactor : curUbMfactorTail; + xLocal = xQue.template AllocTensor(); + int64_t curGlobalxOffset = (mVectorOffset + i * tilingData->stage1MFactor) * + tilingData->k + kGmBaseOffset; + CopyIn(xGm[curGlobalxOffset], xLocal, curUbMFactor, realKGmSize, tilingData->k - realKGmSize); + xQue.template EnQue(xLocal); + xLocal = xQue.template DeQue(); + xCastLocal = mmInQue.AllocTensor(); + CastTwoDim(xCastLocal, xLocal, curUbMFactor, realKGmSize); + xQue.template FreeTensor(xLocal); + mmInQue.template EnQue(xCastLocal); + xCastLocal = mmInQue.template DeQue(); + int64_t cutMmInOffset = cvLoopIdx_ % DOUBLE_BUFFER * xCastFp32BufSize_ + + (i * tilingData->stage1MFactor + mVectorOffset - mGmOffset) * + tilingData->cvLoopKSize; + CopyOut(xCastLocal, xCastFp32WsGm[cutMmInOffset], curUbMFactor, realKGmSize, + tilingData->cvLoopKSize - realKGmSize); + mmInQue.FreeTensor(xCastLocal); + } + CrossCoreSetFlag(SYNC_AIV_TO_AIC_FLAG); + } + cvLoopIdx_++; + } + } + if ASCEND_IS_AIC { + cubeCompute_.End(); + } else { + CrossCoreWaitFlag(SYNC_AIC_TO_AIV_FLAG); + CrossCoreWaitFlag(SYNC_AIC_TO_AIV_FLAG); + } + SyncAll(); // cv全部同步 + } + +private: + TPipe *pipe; + const HcPreTilingData *tilingData; + GlobalTensor workspaceGm; + GlobalTensor xCastFp32WsGm; + GlobalTensor mmOutFp32WsGm; + GlobalTensor squareSumFp32WsGm; + GlobalTensor hcFnGm; + GlobalTensor xGm; + + TQue xQue; + + TQue mmInQue; + + LocalTensor xLocal; + LocalTensor xCastLocal; + LocalTensor yCastLocal; + LocalTensor mmOutLocal; + + HcCubeCompute cubeCompute_; + static constexpr uint64_t SYNC_AIV_TO_AIC_FLAG = 8; + static constexpr uint64_t SYNC_AIC_TO_AIV_FLAG = 9; + static constexpr uint64_t SYNC_MODE2 = NUM_TWO; + + uint64_t cvLoopIdx_ = 0; + uint64_t xCastFp32BufSize_; + uint64_t mmOuterInnerSize_; + uint64_t mmOutFp32BufSize_; + uint64_t squareSumFp32BufSize_; +}; + +template +class HcPreMembaseKSplitCorePart2 { +public: + __aicore__ inline HcPreMembaseKSplitCorePart2() + { + } + + __aicore__ inline void InitGlobalBuffers(GM_ADDR x, GM_ADDR hcScale, + GM_ADDR hcBase, GM_ADDR y, GM_ADDR post, GM_ADDR combFrag, + GM_ADDR workspace) + { + xGm.SetGlobalBuffer((__gm__ T *)x); + hcScaleGm.SetGlobalBuffer((__gm__ float *)hcScale); + hcBaseGm.SetGlobalBuffer((__gm__ float *)hcBase); + yGm.SetGlobalBuffer((__gm__ T *)y); + postGm.SetGlobalBuffer((__gm__ float *)post); + combFragGm.SetGlobalBuffer((__gm__ float *)combFrag); + workspaceGm.SetGlobalBuffer((__gm__ float *)workspace); + } + + __aicore__ inline void InitQueBuffers(int64_t stage1UsedCoreNum, + int64_t xQueNum2) + { + int64_t mixesQue01Size = stage1UsedCoreNum * tilingData->stage2RowFactor * + tilingData->hcMultAlign * NUM_TWO * sizeof(float); + pipe->InitBuffer(mixesQue01, NUM_TWO, mixesQue01Size); + pipe->InitBuffer(mixesQue2, NUM_TWO, + stage1UsedCoreNum * tilingData->stage2RowFactor * + tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(squareSumQue, NUM_TWO, stage1UsedCoreNum * + tilingData->stage2RowFactor * SQUARE_SUM_SIZE * sizeof(float)); + pipe->InitBuffer(xQue, NUM_TWO, xQueNum2 * sizeof(T)); + pipe->InitBuffer(squareSumQue, NUM_TWO, stage1UsedCoreNum * + tilingData->stage2RowFactor * SQUARE_SUM_SIZE * sizeof(float)); + pipe->InitBuffer(yQue, NUM_TWO, + tilingData->stage2RowFactor * RoundUp(tilingData->dFactor) * sizeof(T)); + pipe->InitBuffer(postQue, NUM_TWO, + tilingData->stage2RowFactor * tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(combFragQue, NUM_TWO, + tilingData->stage2RowFactor * tilingData->hcMult * + tilingData->hcMultAlign * sizeof(float)); + } + + __aicore__ inline void InitTBufBuffers(int64_t xQueNum2) + { + pipe->InitBuffer(hcBaseBuf0, tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(hcBaseBuf1, tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(hcBaseBuf2, + tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(rowBrcbBuf0, + RoundUp(tilingData->stage2RowFactor) * BLOCK_SIZE); + pipe->InitBuffer(hcBrcbBuf1, + RoundUp(tilingData->stage2RowFactor * + tilingData->hcMultAlign) * BLOCK_SIZE); + pipe->InitBuffer(reduceBuf, + tilingData->stage2RowFactor * tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(mixes01ReduceBuf, tilingData->stage2RowFactor * + tilingData->hcMultAlign * NUM_TWO * sizeof(float)); + pipe->InitBuffer(mixes02ReduceBuf, tilingData->stage2RowFactor * + tilingData->hcMultAlign * tilingData->hcMult * sizeof(float)); + pipe->InitBuffer(squareReduceBuf, + tilingData->stage2RowFactor * tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(xCastBuf, xQueNum2 * sizeof(float)); + pipe->InitBuffer(yCastBuf, tilingData->stage2RowFactor * + RoundUp(tilingData->dFactor) * sizeof(float)); + pipe->InitBuffer(rsqrtBuf, + RoundUp(tilingData->stage2RowFactor) * sizeof(float)); + pipe->InitBuffer(maskPatternBuf, + RoundUp(MASK_PATTERN_BASE_SIZE * MASK_PATTERN_REPEAT_SIZE) * + sizeof(uint32_t)); + } + + __aicore__ inline void GetLocalTensors() + { + hcBase0Local = hcBaseBuf0.Get(); + hcBase1Local = hcBaseBuf1.Get(); + hcBase2Local = hcBaseBuf2.Get(); + rowBrcbLocal0 = rowBrcbBuf0.Get(); + hcBrcbLocal1 = hcBrcbBuf1.Get(); + reduceLocal = reduceBuf.Get(); + squareReduceLocal = squareReduceBuf.Get(); + mixes01ReduceLocal = mixes01ReduceBuf.Get(); + mixes02ReduceLocal = mixes02ReduceBuf.Get(); + xCastLocal = xCastBuf.Get(); + yCastLocal = yCastBuf.Get(); + rsqrtLocal = rsqrtBuf.Get(); + maskPatternLocal = maskPatternBuf.Get(); + SetGatherMaskPattern(maskPatternLocal); + } + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR hcScale, GM_ADDR hcBase, + GM_ADDR y, GM_ADDR post, GM_ADDR combFrag, GM_ADDR workspace, + const HcPreTilingData *tilingDataPtr, TPipe *pipePtr) + { + pipe = pipePtr; + tilingData = tilingDataPtr; + InitGlobalBuffers(x, hcScale, hcBase, y, post, combFrag, workspace); + int64_t stage1UsedCoreNum = tilingData->cubeBlockDimK; + int64_t xQueNum2 = tilingData->stage2RowFactor * tilingData->hcMult * + RoundUp(tilingData->dFactor); + InitQueBuffers(stage1UsedCoreNum, xQueNum2); + InitTBufBuffers(xQueNum2); + GetLocalTensors(); + } + + __aicore__ inline void Process() + { + if ASCEND_IS_AIV { + int64_t stage1UsedCoreNum = tilingData->cubeBlockDimK;// todo check 此处不应该写死32 + int64_t stage2BlockIdx = GetBlockIdx(); + int64_t stage2UsedCoreNum = tilingData->secondUsedCoreNum; + if (stage2BlockIdx >= stage2UsedCoreNum) { + return; + } + int64_t mmLastAxisSize = CeilAlign(tilingData->hcMix, MM_CACHE_LINE_BYTES / sizeof(float)); + int64_t xCastFp32BufSize = tilingData->mL1Size * + CeilAlign(tilingData->cvLoopKSize, MM_CACHE_LINE_BYTES / sizeof(float)); + int64_t workspaceSize1 = tilingData->cubeCoreNum * DOUBLE_BUFFER * xCastFp32BufSize; + int64_t workspaceSize2 = CeilAlign(stage1UsedCoreNum * tilingData->bs * + mmLastAxisSize * sizeof(float), WORKSPACE_ALIGN_SIZE) / sizeof(float); + CopyIn(hcBaseGm, hcBase0Local, 1, tilingData->hcMult); + CopyIn(hcBaseGm[tilingData->hcMult], hcBase1Local, 1, tilingData->hcMult); + CopyIn(hcBaseGm[tilingData->hcMult * NUM_TWO], hcBase2Local, tilingData->hcMult, tilingData->hcMult); + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId); + WaitFlag(eventId); + + int64_t rowOuterLoop = + (stage2BlockIdx == stage2UsedCoreNum - 1) ? + tilingData->rowLoopOfTailBlock : tilingData->rowLoopOfFormerBlock; + int64_t tailRowFactor = (stage2BlockIdx == stage2UsedCoreNum - 1) ? tilingData->tailRowFactorOfTailBlock : + tilingData->tailRowFactorOfFormerBlock; + int64_t xGmBlockBaseOffsetPart2 = stage2BlockIdx * + tilingData->rowOfFormerBlock * tilingData->hcMult * tilingData->d; + + for (int64_t rowOuterIdx = 0; rowOuterIdx < rowOuterLoop; rowOuterIdx++) { + int64_t xGmBsBaseOffsetPart2 = rowOuterIdx * tilingData->stage2RowFactor * + tilingData->hcMult * tilingData->d; + int64_t curRowFactor = (rowOuterIdx == rowOuterLoop - 1) ? tailRowFactor : tilingData->stage2RowFactor; + squareSumOutLocal = squareSumQue.AllocTensor(); + //todo + CopyIn(workspaceGm[workspaceSize1 + workspaceSize2 + + stage2BlockIdx * tilingData->rowOfFormerBlock * SQUARE_SUM_SIZE + + rowOuterIdx * tilingData->stage2RowFactor * SQUARE_SUM_SIZE], + squareSumOutLocal, stage1UsedCoreNum, curRowFactor * SQUARE_SUM_SIZE, + CeilAlign(tilingData->bs, SQUARE_SUM_SIZE) * SQUARE_SUM_SIZE - + curRowFactor * SQUARE_SUM_SIZE); + squareSumQue.EnQue(squareSumOutLocal); + squareSumOutLocal = squareSumQue.DeQue(); + ReduceSumARAPerf(squareReduceLocal, squareSumOutLocal, 1, stage1UsedCoreNum, + curRowFactor * SQUARE_SUM_SIZE); +int64_t curBsIdxForAll = (stage2BlockIdx * tilingData->rowLoopOfFormerBlock + + rowOuterIdx) * tilingData->stage2RowFactor; + GatherMaskByDiagonal(rsqrtLocal, squareReduceLocal, + maskPatternLocal[(curBsIdxForAll % SQUARE_SUM_SIZE) * 8], curRowFactor); + float coeff = 1.0f / static_cast(tilingData->k); + Muls(rsqrtLocal, rsqrtLocal, coeff, curRowFactor); + PipeBarrier(); + Adds(rsqrtLocal, rsqrtLocal, tilingData->normEps, curRowFactor); + PipeBarrier(); + Sqrt(rsqrtLocal, rsqrtLocal, curRowFactor); + Duplicate(rowBrcbLocal0, static_cast(1.0f), curRowFactor); + PipeBarrier(); + Div(rsqrtLocal, rowBrcbLocal0, rsqrtLocal, curRowFactor); + + mixes01Local = mixesQue01.AllocTensor(); + + uint64_t mixBaseOffset = workspaceSize1 + + stage2BlockIdx * tilingData->rowOfFormerBlock * + CeilAlign(tilingData->hcMix, WORKSPACE_ALIGN_SIZE / sizeof(float)) + + rowOuterIdx * tilingData->stage2RowFactor * + CeilAlign(tilingData->hcMix, WORKSPACE_ALIGN_SIZE / sizeof(float)); + CopyInWithOuterFor(workspaceGm[mixBaseOffset], mixes01Local, stage1UsedCoreNum, + curRowFactor, tilingData->hcMult, tilingData->bs, + CeilAlign(tilingData->hcMix, WORKSPACE_ALIGN_SIZE / sizeof(float))); + CopyInWithOuterFor(workspaceGm[mixBaseOffset + tilingData->hcMult], + mixes01Local[stage1UsedCoreNum * tilingData->stage2RowFactor * + tilingData->hcMultAlign], stage1UsedCoreNum, curRowFactor, tilingData->hcMult, + tilingData->bs, + CeilAlign(tilingData->hcMix, WORKSPACE_ALIGN_SIZE / sizeof(float))); + + mixesQue01.EnQue(mixes01Local); + mixes01Local = mixesQue01.DeQue(); + ReduceSumARAPerf(mixes01ReduceLocal, mixes01Local, NUM_TWO, stage1UsedCoreNum, + curRowFactor * tilingData->hcMultAlign); + ProcessPre(mixes01ReduceLocal, mixes01ReduceLocal, hcBase0Local, rsqrtLocal, + rowBrcbLocal0, hcBrcbLocal1, hcScaleGm.GetValue(0), tilingData->hcEps, + curRowFactor, tilingData->hcMult); + for (int64_t dLoopIdx = 0; dLoopIdx < tilingData->dLoop; dLoopIdx++) { + int64_t curDFactor = + (dLoopIdx == tilingData->dLoop - 1) ? tilingData->tailDFactor : tilingData->dFactor; + xLocal = xQue.template AllocTensor(); + CopyIn(xGm[xGmBlockBaseOffsetPart2 + xGmBsBaseOffsetPart2 + + dLoopIdx * tilingData->dFactor], xLocal, + tilingData->stage2RowFactor * tilingData->hcMult, curDFactor, + tilingData->d - curDFactor); + xQue.template EnQue(xLocal); + xLocal = xQue.template DeQue(); + yLocal = yQue.template AllocTensor(); + ProcessY(yLocal, xLocal, mixes01ReduceLocal, hcBrcbLocal1, xCastLocal, yCastLocal, curRowFactor, + tilingData->hcMult, curDFactor); + xQue.template FreeTensor(xLocal); + yQue.template EnQue(yLocal); + yLocal = yQue.template DeQue(); + CopyOut(yLocal, + yGm[stage2BlockIdx * tilingData->rowOfFormerBlock * tilingData->d + + rowOuterIdx * tilingData->stage2RowFactor * tilingData->d + + dLoopIdx * tilingData->dFactor], + curRowFactor, curDFactor, tilingData->d - curDFactor); + yQue.template FreeTensor(yLocal); + } + // post + postLocal = postQue.AllocTensor(); + ProcessPost(postLocal, + mixes01ReduceLocal[tilingData->stage2RowFactor * tilingData->hcMultAlign], + hcBase1Local, rsqrtLocal, rowBrcbLocal0, hcBrcbLocal1, hcScaleGm.GetValue(1), + curRowFactor, tilingData->hcMult); + mixesQue01.template FreeTensor(mixes01Local); + postQue.EnQue(postLocal); + postLocal = postQue.DeQue(); + CopyOut(postLocal, + postGm[stage2BlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMult + + rowOuterIdx * tilingData->stage2RowFactor * tilingData->hcMult], + curRowFactor, tilingData->hcMult); + postQue.FreeTensor(postLocal); + + // combFrag + mixes2Local = mixesQue2.AllocTensor(); + for (int64_t i = 0; i < stage1UsedCoreNum; ++i) { + for (int64_t j = 0; j < curRowFactor; ++j) { + CopyIn(workspaceGm[workspaceSize1 + i * tilingData->bs * mmLastAxisSize + + j * mmLastAxisSize + + stage2BlockIdx * tilingData->rowOfFormerBlock * mmLastAxisSize + + rowOuterIdx * tilingData->stage2RowFactor * mmLastAxisSize + + tilingData->hcMult * NUM_TWO], + mixes2Local[(i * curRowFactor + j) * tilingData->hcMult * + tilingData->hcMultAlign], tilingData->hcMult, tilingData->hcMult); + } + } + mixesQue2.EnQue(mixes2Local); + mixes2Local = mixesQue2.DeQue(); + ReduceSumARAPerf(mixes02ReduceLocal, mixes2Local, 1, stage1UsedCoreNum, + curRowFactor * tilingData->hcMult * tilingData->hcMultAlign); + combFragLocal = combFragQue.AllocTensor(); + + MulABLastDimBrcInline(mixes02ReduceLocal, mixes02ReduceLocal, + rsqrtLocal, rowBrcbLocal0, curRowFactor, + tilingData->hcMult * tilingData->hcMultAlign); + Muls(mixes02ReduceLocal, mixes02ReduceLocal, hcScaleGm.GetValue(NUM_TWO), + curRowFactor * tilingData->hcMult * tilingData->hcMultAlign); + PipeBarrier(); + AddBAFirstDimBrcInline(mixes02ReduceLocal, mixes02ReduceLocal, hcBase2Local, curRowFactor, + tilingData->hcMult * tilingData->hcMultAlign); + SoftmaxFP32Perf(mixes02ReduceLocal, mixes02ReduceLocal, reduceLocal, hcBrcbLocal1, + curRowFactor * tilingData->hcMult, tilingData->hcMult, tilingData->hcEps); + ReduceSumARAPerf(reduceLocal, mixes02ReduceLocal, curRowFactor, tilingData->hcMult, tilingData->hcMult); + Adds(reduceLocal, reduceLocal, tilingData->hcEps, curRowFactor * tilingData->hcMult); + PipeBarrier(); + DivABABrcInline(combFragLocal, mixes02ReduceLocal, reduceLocal, curRowFactor, tilingData->hcMult, + tilingData->hcMult); + for (int64_t iter = 0; iter < tilingData->iterTimes - 1; iter++) { + LastDimReduceSumPerf(reduceLocal, combFragLocal, + curRowFactor * tilingData->hcMult, tilingData->hcMult); + Adds(reduceLocal, reduceLocal, tilingData->hcEps, + curRowFactor * tilingData->hcMult); + PipeBarrier(); + DivABLastDimBrcInline(combFragLocal, combFragLocal, + reduceLocal, hcBrcbLocal1, curRowFactor * tilingData->hcMult, + tilingData->hcMult); + ReduceSumARAPerf(reduceLocal, combFragLocal, curRowFactor, + tilingData->hcMult, tilingData->hcMult); + Adds(reduceLocal, reduceLocal, tilingData->hcEps, + curRowFactor * tilingData->hcMult); + PipeBarrier(); + DivABABrcInline(combFragLocal, combFragLocal, reduceLocal, + curRowFactor, tilingData->hcMult, tilingData->hcMult); + } + mixesQue2.FreeTensor(mixes2Local); + squareSumQue.template FreeTensor(squareSumOutLocal); + + combFragQue.EnQue(combFragLocal); + combFragLocal = combFragQue.DeQue(); + CopyOut(combFragLocal, + combFragGm[stage2BlockIdx * tilingData->rowOfFormerBlock * + tilingData->hcMult * tilingData->hcMult + + rowOuterIdx * tilingData->stage2RowFactor * tilingData->hcMult * + tilingData->hcMult], curRowFactor * tilingData->hcMult, + tilingData->hcMult); + combFragQue.FreeTensor(combFragLocal); + } + } + } + +private: + TPipe *pipe; + const HcPreTilingData *tilingData; + GlobalTensor mixesGm; + GlobalTensor rsqrtGm; + GlobalTensor hcScaleGm; + GlobalTensor hcBaseGm; + GlobalTensor workspaceGm; + GlobalTensor xGm; + GlobalTensor yGm; + GlobalTensor postGm; + GlobalTensor combFragGm; + + TQue mixesQue01; + TQue mixesQue2; + TQue xQue; + TQue yQue; + TQue postQue; + TQue combFragQue; + + TQue squareSumQue; + + TBuf hcBaseBuf0; + TBuf hcBaseBuf1; + TBuf hcBaseBuf2; + + TBuf rowBrcbBuf0; + TBuf hcBrcbBuf1; + TBuf reduceBuf; + + TBuf rsqrtBuf; + TBuf squareReduceBuf; + TBuf mixes01ReduceBuf; + TBuf mixes02ReduceBuf; + + TBuf xCastBuf; + TBuf yCastBuf; + TBuf maskPatternBuf; + + LocalTensor mixes01Local; + LocalTensor mixes2Local; + LocalTensor rsqrtLocal; + LocalTensor xLocal; + LocalTensor yLocal; + LocalTensor postLocal; + LocalTensor combFragLocal; + LocalTensor hcBase0Local; + LocalTensor hcBase1Local; + LocalTensor hcBase2Local; + LocalTensor rowBrcbLocal0; + LocalTensor hcBrcbLocal1; + LocalTensor reduceLocal; + LocalTensor squareReduceLocal; + LocalTensor mixes01ReduceLocal; + LocalTensor mixes02ReduceLocal; + LocalTensor xCastLocal; + LocalTensor yCastLocal; + LocalTensor squareSumOutLocal; + LocalTensor maskPatternLocal; +}; + +} // namespace HcPreSinkhorn + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_m_k_split_core_arch35.h b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_m_k_split_core_arch35.h new file mode 100644 index 000000000..84564f017 --- /dev/null +++ b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_m_k_split_core_arch35.h @@ -0,0 +1,428 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_m_k_split_core.h + * \brief + */ + +#ifndef HC_PRE_M_K_SPLIT_CORE_ARACH35_H +#define HC_PRE_M_K_SPLIT_CORE_ARACH35_H + +#include "kernel_operator.h" +#include "hc_pre_base_arch35.h" +#include "hc_pre_cube_compute_arch35.h" + +namespace HcPreNs { +using namespace AscendC; + +template +class HcPreMKSplitCorePart1 { +public: + __aicore__ inline HcPreMKSplitCorePart1() + {} + + __aicore__ inline void Init( + GM_ADDR x, GM_ADDR hcFn, GM_ADDR workspace, const HcPreTilingData* tilingDataPtr, TPipe* pipePtr) + { + pipe = pipePtr; + tilingData = tilingDataPtr; + xGm.SetGlobalBuffer((__gm__ T*)x); + hcFnGm.SetGlobalBuffer((__gm__ float*)hcFn); + mmGm.SetGlobalBuffer((__gm__ float*)workspace); + rmsGm.SetGlobalBuffer((__gm__ float*)workspace + tilingData->kBlockFactor * tilingData->bs * tilingData->hcMix); + + TBuf l1Buffer; + pipe->InitBuffer(l1Buffer, L1_ALLOC_SIZE); + xL1_ = l1Buffer.Get(); + wL1_ = l1Buffer.Get()[L1_BUF_NUM * L1_BUF_OFFSET]; + + // InQue + pipe->InitBuffer(xQue, 2, tilingData->mUbSize * RoundUp(tilingData->kUbSize) * sizeof(T)); + + // OutQue + pipe->InitBuffer(rmsQue, 2, RoundUp(tilingData->mUbSize) * sizeof(float)); + + // Calc Buf + pipe->InitBuffer(castBuf, tilingData->mUbSize * (RoundUp(tilingData->kUbSize) * sizeof(float) + BLOCK_SIZE)); + pipe->InitBuffer(nd2NzBuf, CeilAlign(tilingData->mUbSize, C0_SIZE) * RoundUp(tilingData->kUbSize) * sizeof(float) * DOUBLE_BUFFER); + + if ASCEND_IS_AIC { + mmService_.Init(); + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG); + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG + FLAG_ID_MAX); + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG); + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG + FLAG_ID_MAX); + } + xCastLocal = castBuf.Get(); + xNd2NzLocal = nd2NzBuf.Get(); + } + + __aicore__ inline void Process() + { + int64_t curBlockIdx = GetBlockIdx(); + int64_t totalBlockNum = GetBlockNum(); + + uint64_t mBlkDimIdx = curBlockIdx / tilingData->cubeBlockDimK; + uint64_t kBlkDimIdx = curBlockIdx % tilingData->cubeBlockDimK; + + // todo 移到tiling计算 + uint64_t mCnt = CeilDiv(tilingData->bs, tilingData->mL1Size); + uint64_t singleCoreMaxRound = CeilDiv(mCnt, tilingData->cubeBlockDimM); + uint64_t mainCoreCount = mCnt % tilingData->cubeBlockDimM; + uint64_t singleCoreRound = (mainCoreCount == 0 || curBlockIdx < mainCoreCount) ? singleCoreMaxRound : singleCoreMaxRound - 1; + uint64_t mGmOffset = 0; + uint64_t nd2NzBufSize = CeilAlign(tilingData->mUbSize, C0_SIZE) * RoundUp(tilingData->kUbSize); + if ASCEND_IS_AIC { + mGmOffset = (curBlockIdx / tilingData->cubeBlockDimK) * singleCoreMaxRound * tilingData->mL1Size; + } else { + mGmOffset = ((curBlockIdx / 2) / tilingData->cubeBlockDimK) * singleCoreMaxRound * tilingData->mL1Size; + } + int64_t xGmBaseOffset = 0; + int64_t rmsGmBaseOffset = 0; + if ASCEND_IS_AIV { + int64_t aivCurBlockIdx = GetBlockIdx(); + xGmBaseOffset = ((aivCurBlockIdx / 2) / tilingData->cubeBlockDimK) * singleCoreMaxRound * tilingData->mL1Size * tilingData->hcMult * tilingData->d + + ((aivCurBlockIdx / 2) % tilingData->cubeBlockDimK) * tilingData->multCoreSplitKSize; + rmsGmBaseOffset = ((aivCurBlockIdx / 2) / tilingData->cubeBlockDimK) * singleCoreMaxRound * tilingData->mL1Size; + } + + // todo 移到tiling计算 + int64_t xSplitOffset = 0; + int64_t rmsSplitOffset = 0; + // m轴切分 按照0 0 1 1..分核 + int64_t bufferIdx = 0; + int64_t curAicBlockIdx = 0; + if ASCEND_IS_AIC { + curAicBlockIdx = curBlockIdx; + } else { + curAicBlockIdx = curBlockIdx / 2; + } + + if (curAicBlockIdx < tilingData->cubeBlockDimK * tilingData->cubeBlockDimM) { + if ASCEND_IS_AIV { + SetFlag(static_cast(0)); + SetFlag(static_cast(1)); + } + for (uint64_t roundIdx = 0; roundIdx < singleCoreRound; mGmOffset += tilingData->mL1Size, ++roundIdx) + { + uint64_t mL1RealSize = AscendC::Std::min(tilingData->bs - mGmOffset, (uint64_t)tilingData->mL1Size); + uint64_t kGmStartOffset = 0; + uint64_t kGmEndOffset = AscendC::Std::min(tilingData->multCoreSplitKSize, + tilingData->k - (curAicBlockIdx % tilingData->cubeBlockDimK) * tilingData->multCoreSplitKSize); + if ASCEND_IS_AIV { + if (GetBlockIdx() % 2 != 0) { + xSplitOffset = (mL1RealSize / 2) * tilingData->hcMult * tilingData->d; + rmsSplitOffset = mL1RealSize / 2; + } + rmsNormLocal = rmsQue.template AllocTensor(); + } + + int64_t curRowFactor = 0; + for (int64_t kGmOffset = kGmStartOffset; kGmOffset < kGmEndOffset; kGmOffset += tilingData->kL1Size) { + uint64_t kL1RealSize = AscendC::Std::min(kGmEndOffset - kGmOffset, (uint64_t)tilingData->kL1Size); + if ASCEND_IS_AIC { + bool isFirstKL1 = kGmOffset == kGmStartOffset; + bool isLastKL1 = (kGmOffset + tilingData->kL1Size) >= kGmEndOffset; + mmService_.CopyInB1Nd2Nz(tilingData->hcMult * tilingData->d, kL1RealSize, + tilingData->hcMix, hcFnGm[kGmOffset + kBlkDimIdx * tilingData->multCoreSplitKSize], + wL1_[mmService_.GetBL1BufferId() * L1_BUF_OFFSET]); + CrossCoreWaitFlag(SYNC_AIV_AIC_FLAG + FLAG_ID_MAX); + CrossCoreWaitFlag(SYNC_AIV_AIC_FLAG); + uint64_t mL1AlignSize = Align(mL1RealSize, AscendC::BLOCK_CUBE); + uint64_t nL1AlignSize = Align((uint64_t)tilingData->hcMix, AscendC::BLOCK_CUBE); + + mmService_.Process(tilingData->bs, tilingData->hcMix, mL1RealSize, (256 / AscendC::Std::max(mL1AlignSize, nL1AlignSize)) * 32, + isFirstKL1, isLastKL1, xL1_[aL1BufferID_ * L1_BUF_OFFSET], wL1_[mmService_.GetBL1BufferId() * L1_BUF_OFFSET]); + if (isLastKL1) { + mmService_.CopyOut(mmGm[mBlkDimIdx * tilingData->mL1Size * singleCoreMaxRound * tilingData->hcMix + kBlkDimIdx * tilingData->bs * tilingData->hcMix + roundIdx * tilingData->mL1Size]); + } + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG); // 写出ub搬出，cv流水同步比较复杂，暂不讨论 + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG + FLAG_ID_MAX); + } else { + CrossCoreWaitFlag(SYNC_AIC_AIV_FLAG); + int64_t rowFactor = mL1RealSize / 2; + int64_t tailRowFactor = mL1RealSize - rowFactor; + curRowFactor = rowFactor; + int64_t mL1SizeAlign = CeilAlign(mL1RealSize, AscendC::BLOCK_CUBE); + if (curBlockIdx % 2 == 1) { + curRowFactor = tailRowFactor; + } + uint64_t cvLoopKSize = CeilDiv(kL1RealSize, tilingData->kUbSize); + uint64_t kReminderSize = kL1RealSize - (cvLoopKSize - 1) * tilingData->kUbSize; + float coeff = 1 / static_cast(tilingData->hcMult * tilingData->d); + for (int64_t cvLoopIdx = 0; cvLoopIdx < cvLoopKSize; cvLoopIdx++) { + uint64_t kRealSize = (cvLoopIdx == cvLoopKSize - 1) ? kReminderSize : tilingData->kUbSize; + xLocal = xQue.template AllocTensor(); + CopyIn(xGm[xGmBaseOffset + xSplitOffset + roundIdx * tilingData->mL1Size * tilingData->hcMult * tilingData->d + kGmOffset + cvLoopIdx * tilingData->kUbSize], + xLocal, curRowFactor, kRealSize, tilingData->hcMult * tilingData->d - kRealSize); + xQue.template EnQue(xLocal); + xLocal = xQue.template DeQue(); + if (kGmOffset == kGmStartOffset && cvLoopIdx == 0) { + VFProcessCastAndInvRmsPart1(rmsNormLocal, xCastLocal, xLocal, coeff, curRowFactor, kRealSize); + } else { + VFProcessCastAndInvRmsPart1(rmsNormLocal, xCastLocal, xLocal, coeff, curRowFactor, kRealSize); + } + xQue.template FreeTensor(xLocal); + + WaitFlag(static_cast(bufferIdx & 1)); + VFTransND2NZ(xNd2NzLocal[nd2NzBufSize * (bufferIdx & 1)], xCastLocal, curRowFactor, kRealSize); + SetFlag(static_cast(bufferIdx & 1)); + WaitFlag(static_cast(bufferIdx & 1)); + + if (curBlockIdx % 2 == 0) { + DataCopyParams dataCopyXParams; + dataCopyXParams.blockCount = CeilDiv(kRealSize, C0_SIZE); + dataCopyXParams.blockLen = curRowFactor * C0_SIZE * sizeof(float) / BLOCK_SIZE; + dataCopyXParams.srcStride = CeilAlign(curRowFactor, C0_SIZE) - curRowFactor; + dataCopyXParams.dstStride = CeilAlign(mL1RealSize, 16) - curRowFactor; + CopyToL1(xNd2NzLocal[nd2NzBufSize * (bufferIdx & 1)], xL1_[(aL1BufferID_ * L1_BUF_OFFSET) + cvLoopIdx * tilingData->kUbSize * mL1SizeAlign], dataCopyXParams); + } else { + DataCopyParams dataCopyXParams; + dataCopyXParams.blockCount = CeilDiv(kRealSize, C0_SIZE); + dataCopyXParams.blockLen = curRowFactor * C0_SIZE * sizeof(float) / BLOCK_SIZE; + dataCopyXParams.srcStride = CeilAlign(curRowFactor, C0_SIZE) - curRowFactor; + dataCopyXParams.dstStride = CeilAlign(mL1RealSize, 16) - curRowFactor; + CopyToL1(xNd2NzLocal[nd2NzBufSize * (bufferIdx & 1)], xL1_[(aL1BufferID_ * L1_BUF_OFFSET) + rowFactor * (BLOCK_SIZE / sizeof(float)) + cvLoopIdx * tilingData->kUbSize * mL1SizeAlign], dataCopyXParams); + } + SetFlag(static_cast(bufferIdx & 1)); + bufferIdx++; + } + CrossCoreSetFlag(SYNC_AIV_AIC_FLAG); + } + aL1BufferID_ ^= 1; + } + if ASCEND_IS_AIV { + int64_t kBaseOffset = (GetBlockIdx() / 2) % tilingData->kBlockFactor * tilingData->bs; + rmsQue.template EnQue(rmsNormLocal); + rmsNormLocal = rmsQue.template DeQue(); + CopyOut(rmsNormLocal, rmsGm[kBaseOffset + rmsGmBaseOffset + rmsSplitOffset + roundIdx * tilingData->mL1Size], 1, curRowFactor); + rmsQue.template FreeTensor(rmsNormLocal); + } + } + if ASCEND_IS_AIV { + WaitFlag(static_cast(0)); + WaitFlag(static_cast(1)); + } + } + SyncAll(); + } + +private: + TPipe* pipe; + const HcPreTilingData* tilingData; + // (M, K) * (N, K) + + GlobalTensor xGm; + GlobalTensor hcFnGm; + GlobalTensor mmGm; + GlobalTensor rmsGm; + + TQue xQue; + TQue rmsQue; + TBuf castBuf; + TBuf nd2NzBuf; + + LocalTensor xLocal; + LocalTensor mmXLocal; + LocalTensor rmsNormLocal; + LocalTensor xCastLocal; + LocalTensor xNd2NzLocal; + + HcPreCubeCompute mmService_; + LocalTensor xL1_; + LocalTensor wL1_; + static constexpr uint64_t SYNC_AIV_AIC_FLAG = 8; + static constexpr uint64_t SYNC_AIC_AIV_FLAG = 9; + static constexpr uint64_t SYNC_AIV_AIC_PRE_POST_FLAG = 10; + static constexpr uint64_t SYNC_AIC_AIV_PRE_POST_FLAG = 11; + static constexpr uint64_t FLAG_ID_MAX = 16; + uint64_t cvLoopIdx_ = 0; + uint8_t aL1BufferID_ = 0; +}; + + +template +class HcPreMKSplitCorePart2 { +public: + __aicore__ inline HcPreMKSplitCorePart2() + {} + + __aicore__ inline void Init( + GM_ADDR x, GM_ADDR hcScale, GM_ADDR hcBase, GM_ADDR y, GM_ADDR post, + GM_ADDR combFrag, GM_ADDR workspace, const HcPreTilingData* tilingDataPtr, TPipe* pipePtr) + { + pipe = pipePtr; + tilingData = tilingDataPtr; + + xGm.SetGlobalBuffer((__gm__ T*)x); + hcScaleGm.SetGlobalBuffer((__gm__ float*)hcScale); + hcBaseGm.SetGlobalBuffer((__gm__ float*)hcBase); + yGm.SetGlobalBuffer((__gm__ T*)y); + postGm.SetGlobalBuffer((__gm__ float*)post); + combFragGm.SetGlobalBuffer((__gm__ float*)combFrag); + mmGm.SetGlobalBuffer((__gm__ float*)workspace); + rmsGm.SetGlobalBuffer((__gm__ float*)workspace + tilingData->kBlockFactor * tilingData->bs * tilingData->hcMix); + + + // InQue + pipe->InitBuffer( + xQue, 2, tilingData->stage2RowFactor * tilingData->hcMult * RoundUp(tilingData->dFactor) * sizeof(T)); + int64_t rmsAndmmQueSize = tilingData->kBlockFactor * RoundUp(tilingData->stage2RowFactor) * sizeof(float) + + tilingData->kBlockFactor * tilingData->stage2RowFactor * RoundUp(tilingData->hcMix) * sizeof(float); + pipe->InitBuffer(rmsAndmmQue, 2, rmsAndmmQueSize); + + // OutQue + pipe->InitBuffer( + yQue, 2, tilingData->stage2RowFactor * RoundUp(tilingData->dFactor) * sizeof(T)); + pipe->InitBuffer(postQue, 2, tilingData->stage2RowFactor * tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(combFragQue, DOUBLE_BUFFER, + tilingData->stage2RowFactor * tilingData->hcMult * tilingData->hcMult * sizeof(float)); + + // TBuf + pipe->InitBuffer(hcBaseBuf0, tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(hcBaseBuf1, tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(hcBaseBuf2, tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(mixesBuf, tilingData->stage2RowFactor * RoundUp(tilingData->hcMix) * sizeof(float)); + + hcBase0Local = hcBaseBuf0.Get(); + hcBase1Local = hcBaseBuf1.Get(); + hcBase2Local = hcBaseBuf2.Get(); + mixesLocal = mixesBuf.Get(); + } + + __aicore__ inline void Process() + { + if ASCEND_IS_AIV { + int64_t stage1UsedCoreNum = tilingData->cubeBlockDimK; + int64_t curBlockIdx = GetBlockIdx(); + int64_t stage2UsedCoreNum = tilingData->secondUsedCoreNum; + if (curBlockIdx >= stage2UsedCoreNum) { + return; + } + int64_t rowOuterLoop = + (curBlockIdx == stage2UsedCoreNum - 1) ? tilingData->rowLoopOfTailBlock : tilingData->rowLoopOfFormerBlock; + int64_t tailRowFactor = (curBlockIdx == stage2UsedCoreNum - 1) ? tilingData->tailRowFactorOfTailBlock : + tilingData->tailRowFactorOfFormerBlock; + + CopyIn(hcBaseGm, hcBase0Local, 1, tilingData->hcMult); + CopyIn(hcBaseGm[tilingData->hcMult], hcBase1Local, 1, tilingData->hcMult); + CopyIn(hcBaseGm[tilingData->hcMult * 2], hcBase2Local, 1, tilingData->hcMult * tilingData->hcMult); + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId); + WaitFlag(eventId); + int64_t mmGmBaseOffset = curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMix; + int64_t rmsGmBaseOffset = curBlockIdx * tilingData->rowOfFormerBlock; + int64_t xGmBaseOffset = curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMult * tilingData->d; + int64_t mmLocalSize = stage1UsedCoreNum * tilingData->stage2RowFactor * RoundUp(tilingData->hcMix); + for (int64_t rowOuterIdx = 0; rowOuterIdx < rowOuterLoop; rowOuterIdx++) { + int64_t curRowFactor = (rowOuterIdx == rowOuterLoop - 1) ? tailRowFactor : tilingData->stage2RowFactor; + rmsAndmmLocal = rmsAndmmQue.AllocTensor(); + CopyInWithLoopMode( + mmGm[mmGmBaseOffset + rowOuterIdx * tilingData->stage2RowFactor * tilingData->hcMix], rmsAndmmLocal, tilingData->kBlockFactor, curRowFactor, tilingData->hcMix, tilingData->bs * tilingData->hcMix); + CopyIn( + rmsGm[rmsGmBaseOffset + rowOuterIdx * tilingData->stage2RowFactor], + rmsAndmmLocal[mmLocalSize], tilingData->kBlockFactor, curRowFactor, tilingData->bs - curRowFactor); + + rmsAndmmQue.EnQue(rmsAndmmLocal); + rmsAndmmLocal = rmsAndmmQue.DeQue(); + + VFProcessInvRmsPart3WithGroupReduce(mixesLocal, rmsAndmmLocal, rmsAndmmLocal[mmLocalSize], tilingData->normEps, tilingData->kBlockFactor, curRowFactor, tilingData->hcMix); + + VFProcessPre( + mixesLocal, mixesLocal, hcBase0Local, hcScaleGm.GetValue(0), tilingData->hcEps, + curRowFactor, tilingData->hcMult, tilingData->hcMix); + for (int64_t dLoopIdx = 0; dLoopIdx < tilingData->dLoop; dLoopIdx++) { + int64_t curDFactor = + (dLoopIdx == tilingData->dLoop - 1) ? tilingData->tailDFactor : tilingData->dFactor; + xLocal = xQue.template AllocTensor(); + CopyIn( + xGm[xGmBaseOffset + rowOuterIdx * tilingData->stage2RowFactor * tilingData->hcMult * tilingData->d + + dLoopIdx * tilingData->dFactor], + xLocal, curRowFactor * tilingData->hcMult, curDFactor, tilingData->d - curDFactor); + xQue.template EnQue(xLocal); + xLocal = xQue.template DeQue(); + + yLocal = yQue.template AllocTensor(); + VFProcessY(yLocal, mixesLocal, xLocal, curRowFactor, tilingData->hcMult, curDFactor, tilingData->hcMix); + xQue.template FreeTensor(xLocal); + yQue.template EnQue(yLocal); + yLocal = yQue.template DeQue(); + CopyOut(yLocal, yGm[curBlockIdx * tilingData->rowOfFormerBlock * tilingData->d + rowOuterIdx * tilingData->stage2RowFactor * tilingData->d + dLoopIdx * tilingData->dFactor], curRowFactor, curDFactor, tilingData->d - curDFactor); + yQue.template FreeTensor(yLocal); + } + + // post + postLocal = postQue.AllocTensor(); + VFProcessPost( + postLocal, mixesLocal[tilingData->hcMult], hcBase1Local, + hcScaleGm.GetValue(1), tilingData->hcEps, curRowFactor, tilingData->hcMult, tilingData->hcMix); + + postQue.EnQue(postLocal); + postLocal = postQue.DeQue(); + CopyOut(postLocal, postGm[curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMult + rowOuterIdx * tilingData->stage2RowFactor * tilingData->hcMult], curRowFactor, tilingData->hcMult); + postQue.FreeTensor(postLocal); + + // combFrag + combFragLocal = combFragQue.AllocTensor(); + VFProcessCombFragPacked( + combFragLocal, mixesLocal[tilingData->hcMult * 2], hcBase2Local, hcScaleGm.GetValue(2), tilingData->hcEps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMix); + rmsAndmmQue.FreeTensor(rmsAndmmLocal); + + combFragQue.EnQue(combFragLocal); + combFragLocal = combFragQue.DeQue(); + int64_t combLen = tilingData->hcMult * tilingData->hcMult; + int64_t combOutOffset = (curBlockIdx * tilingData->rowOfFormerBlock + + rowOuterIdx * tilingData->stage2RowFactor) * combLen; + CopyOut(combFragLocal, combFragGm[combOutOffset], curRowFactor, combLen); + combFragQue.FreeTensor(combFragLocal); + } + } + } + +private: + TPipe* pipe; + const HcPreTilingData* tilingData; + GlobalTensor hcScaleGm; + GlobalTensor hcBaseGm; + GlobalTensor xGm; + GlobalTensor yGm; + GlobalTensor postGm; + GlobalTensor combFragGm; + + GlobalTensor mmGm; + GlobalTensor rmsGm; + + TQue rmsAndmmQue; + TQue xQue; + + TQue yQue; + TQue postQue; + TQue combFragQue; + + TBuf mixesBuf; + TBuf hcBaseBuf0; + TBuf hcBaseBuf1; + TBuf hcBaseBuf2; + + LocalTensor mixesLocal; + LocalTensor rmsAndmmLocal; + LocalTensor xLocal; + LocalTensor yLocal; + LocalTensor postLocal; + LocalTensor combFragLocal; + LocalTensor hcBase0Local; + LocalTensor hcBase1Local; + LocalTensor hcBase2Local; +}; +} // namespace HCPreSinkhorn + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_m_split_core_arch35.h b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_m_split_core_arch35.h new file mode 100644 index 000000000..6af966c47 --- /dev/null +++ b/csrc/ascend/moe/hc_pre/op_kernel/hc_pre_m_split_core_arch35.h @@ -0,0 +1,411 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_m_k_split_core.h + * \brief + */ + +#ifndef HC_PRE_M_SPLIT_CORE_H +#define HC_PRE_M_SPLIT_CORE_H + +#include "kernel_operator.h" +#include "hc_pre_base_arch35.h" +#include "hc_pre_cube_compute_arch35.h" + +namespace HcPreNs { +using namespace AscendC; + +template +class HcPreMSplitCorePart1 { +public: + __aicore__ inline HcPreMSplitCorePart1() + {} + + __aicore__ inline void Init( + GM_ADDR x, GM_ADDR hcFn, GM_ADDR hcScale, GM_ADDR hcBase, + GM_ADDR y, GM_ADDR post, GM_ADDR combFrag, const HcPreTilingData* tilingDataPtr, TPipe* pipePtr) + { + pipe = pipePtr; + tilingData = tilingDataPtr; + xGm.SetGlobalBuffer((__gm__ T*)x); + hcFnGm.SetGlobalBuffer((__gm__ float*)hcFn); + yGm.SetGlobalBuffer((__gm__ T*)y); + + hcScaleGm.SetGlobalBuffer((__gm__ float*)hcScale); + hcBaseGm.SetGlobalBuffer((__gm__ float*)hcBase); + postGm.SetGlobalBuffer((__gm__ float*)post); + combFragGm.SetGlobalBuffer((__gm__ float*)combFrag); + + TBuf l1Buffer; + pipe->InitBuffer(l1Buffer, L1_ALLOC_SIZE); + xL1_ = l1Buffer.Get(); + wL1_ = l1Buffer.Get()[L1_BUF_NUM * L1_BUF_OFFSET]; + + pipe->InitBufPool(tbufPool0, tilingData->bufferPool0Size); + tbufPool0.InitBuffer(mmXBuf, CeilDiv(tilingData->mL1Size, 2) * RoundUp(tilingData->hcMix) * sizeof(float)); + mmXLocal = mmXBuf.Get(); + + if ASCEND_IS_AIC { + mmService_.Init(); + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG); + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG + FLAG_ID_MAX); + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG); + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG + FLAG_ID_MAX); + } else { + tbufPool0.InitBuffer(rmsNormBuf, RoundUp(CeilDiv(tilingData->mL1Size, 2)) * sizeof(float)); + tbufPool0.InitBufPool(tbufPool1, tilingData->bufferPool1Size); + + tbufPool0.InitBuffer(hcBaseBuf0, tilingData->hcMultAlign * sizeof(float)); + tbufPool0.InitBuffer(hcBaseBuf1, tilingData->hcMultAlign * sizeof(float)); + tbufPool0.InitBuffer(hcBaseBuf2, tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); + + hcBase0Local = hcBaseBuf0.Get(); + hcBase1Local = hcBaseBuf1.Get(); + hcBase2Local = hcBaseBuf2.Get(); + } + } + + __aicore__ inline void Process() + { + int64_t curBlockIdx = GetBlockIdx(); + int64_t logicalBlockIdx = curBlockIdx; + if ASCEND_IS_AIV { + logicalBlockIdx = curBlockIdx / 2; + } + if (logicalBlockIdx >= tilingData->cubeBlockDimM) { + return; + } + + if ASCEND_IS_AIV { + CopyIn(hcBaseGm, hcBase0Local, 1, tilingData->hcMult); + CopyIn(hcBaseGm[tilingData->hcMult], hcBase1Local, 1, tilingData->hcMult); + CopyIn(hcBaseGm[tilingData->hcMult * 2], hcBase2Local, 1, tilingData->hcMult * tilingData->hcMult); + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId); + WaitFlag(eventId); + } + + int64_t totalBlockNum = GetBlockNum(); + + uint64_t mBlkDimIdx = curBlockIdx % tilingData->cubeBlockDimM; + uint64_t kBlkDimIdx = curBlockIdx % tilingData->cubeBlockDimK; + + // todo 移到tiling计算 + uint64_t mCnt = CeilDiv(tilingData->bs, tilingData->mL1Size); + uint64_t singleCoreMaxRound = CeilDiv(mCnt, tilingData->cubeBlockDimM); + uint64_t mainCoreCount = mCnt % tilingData->cubeBlockDimM; + uint64_t singleCoreRound = (mainCoreCount == 0 || logicalBlockIdx < mainCoreCount) ? singleCoreMaxRound : singleCoreMaxRound - 1; + uint64_t mGmOffset = 0; + if ASCEND_IS_AIC { + if (mainCoreCount == 0 || curBlockIdx <= mainCoreCount) { + mGmOffset = curBlockIdx * singleCoreMaxRound * tilingData->mL1Size; + } else { + mGmOffset = (mainCoreCount * singleCoreMaxRound + (curBlockIdx - mainCoreCount) * (singleCoreMaxRound - 1)) * tilingData->mL1Size; + } + } else { + if (mainCoreCount == 0 || (curBlockIdx / 2) <= mainCoreCount) { + mGmOffset = curBlockIdx / 2 * singleCoreMaxRound * tilingData->mL1Size; + } else { + mGmOffset = (mainCoreCount * singleCoreMaxRound + (curBlockIdx / 2 - mainCoreCount) * (singleCoreMaxRound - 1)) * tilingData->mL1Size; + } + } + int64_t xGmBaseOffset = 0; + int64_t yGmBaseOffset = 0; + int64_t postGmBaseOffset = 0; + int64_t combFragGmBaseOffset = 0; + if ASCEND_IS_AIV { + xGmBaseOffset = mGmOffset * tilingData->hcMult * tilingData->d; + yGmBaseOffset = mGmOffset * tilingData->d; + postGmBaseOffset = mGmOffset * tilingData->hcMult; + combFragGmBaseOffset = mGmOffset * tilingData->hcMult * tilingData->hcMult; + SetFlag(static_cast(0)); + } + + uint64_t cvLoopKSize = tilingData->kL1Size / tilingData->kUbSize; + int64_t xSplitOffset = 0; + int64_t ySplitOffset = 0; + int64_t postSplitOffset = 0; + int64_t combFragSplitOffset = 0; + int64_t xOutSplitOffset = 0; + // m轴切分 按照0 0 1 1..分核 + for (uint64_t roundIdx = 0; roundIdx < singleCoreRound; mGmOffset += tilingData->mL1Size, ++roundIdx) + { + uint64_t mL1RealSize = AscendC::Std::min(tilingData->bs - mGmOffset, (uint64_t)tilingData->mL1Size); + uint64_t kGmStartOffset = 0; + uint64_t kGmEndOffset = tilingData->multCoreSplitKSize; + uint64_t nd2NzBufSize = CeilAlign(tilingData->mUbSize, C0_SIZE) * RoundUp(tilingData->kUbSize); + if ASCEND_IS_AIV { + tbufPool1.Reset(); + tbufPool1.InitBuffer(xQue, 2, tilingData->mUbSize * RoundUp(tilingData->kUbSize) * sizeof(T)); + tbufPool1.InitBuffer(castBuf, tilingData->mUbSize * (RoundUp(tilingData->kUbSize) * sizeof(float) + BLOCK_SIZE)); + tbufPool1.InitBuffer(nd2NzBuf, nd2NzBufSize * sizeof(float) * DOUBLE_BUFFER); + + xCastLocal = castBuf.Get(); + xNd2NzLocal = nd2NzBuf.Get(); + rmsNormLocal = rmsNormBuf.Get(); + WaitFlag(static_cast(0)); + if (GetBlockIdx() % 2 != 0) { + xSplitOffset = CeilDiv(mL1RealSize, 2) * tilingData->hcMult * tilingData->d; + xOutSplitOffset = CeilDiv(mL1RealSize, 2) * tilingData->hcMult * tilingData->d; + ySplitOffset = CeilDiv(mL1RealSize, 2) * tilingData->d; + postSplitOffset = CeilDiv(mL1RealSize, 2) * tilingData->hcMult; + combFragSplitOffset = CeilDiv(mL1RealSize, 2) * tilingData->hcMult * tilingData->hcMult; + } + } + // k轴切分（kCoreDim=1） + int64_t bufferIdx = 0; + if ASCEND_IS_AIV { + SetFlag(static_cast(0)); + SetFlag(static_cast(1)); + } + for (int64_t kGmOffset = kGmStartOffset; kGmOffset < kGmEndOffset; kGmOffset += tilingData->kL1Size) { + if ASCEND_IS_AIC { + bool isFirstKL1 = kGmOffset == kGmStartOffset; + bool isLastKL1 = (kGmOffset + tilingData->kL1Size) >= kGmEndOffset; + uint64_t kL1RealSize = AscendC::Std::min(kGmEndOffset - kGmOffset, (uint64_t)tilingData->kL1Size); + mmService_.CopyInB1Nd2Nz(tilingData->multCoreSplitKSize, kL1RealSize, + tilingData->hcMix, hcFnGm[kGmOffset], + wL1_[mmService_.GetBL1BufferId() * L1_BUF_OFFSET]); + CrossCoreWaitFlag(SYNC_AIV_AIC_FLAG + FLAG_ID_MAX); + CrossCoreWaitFlag(SYNC_AIV_AIC_FLAG); + uint64_t mL1AlignSize = Align(mL1RealSize, AscendC::BLOCK_CUBE); + uint64_t nL1AlignSize = Align((uint64_t)tilingData->hcMix, AscendC::BLOCK_CUBE); + mmService_.Process(tilingData->bs, tilingData->hcMix, mL1RealSize, (256 / AscendC::Std::max(mL1AlignSize, nL1AlignSize)) * 32, + isFirstKL1, isLastKL1, xL1_[aL1BufferID_ * L1_BUF_OFFSET], wL1_[mmService_.GetBL1BufferId() * L1_BUF_OFFSET]); + if (isLastKL1) { + mmService_.CopyOut(mmXLocal); + CrossCoreSetFlag(SYNC_AIC_AIV_PRE_POST_FLAG); + CrossCoreSetFlag(SYNC_AIC_AIV_PRE_POST_FLAG + FLAG_ID_MAX); + } + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG); // 写出ub搬出，cv流水同步比较复杂，暂不讨论 + CrossCoreSetFlag(SYNC_AIC_AIV_FLAG + FLAG_ID_MAX); + } else { + CrossCoreWaitFlag(SYNC_AIC_AIV_FLAG); + // Even cores take the first half with CeilDiv, and odd cores take the second half. + // This must match the Sinkhorn/output row split; otherwise odd mL1RealSize shifts the + // second half by one row. + int64_t rowFactor = CeilDiv(mL1RealSize, 2); + int64_t tailRowFactor = mL1RealSize - rowFactor; + int64_t curRowFactor = rowFactor; + int64_t mL1SizeAlign = CeilAlign(mL1RealSize, AscendC::BLOCK_CUBE); + if (curBlockIdx % 2 == 1) { + curRowFactor = tailRowFactor; + } + float coeff = 1 / static_cast(tilingData->hcMult * tilingData->d); + for (int64_t cvLoopIdx = 0; cvLoopIdx < cvLoopKSize; cvLoopIdx++) { + uint64_t kRealSize = kGmOffset + tilingData->kUbSize >= kGmEndOffset ? kGmEndOffset - kGmOffset : tilingData->kUbSize; + + xLocal = xQue.template AllocTensor(); + CopyIn(xGm[xGmBaseOffset + xSplitOffset + roundIdx * tilingData->mL1Size * tilingData->hcMult * tilingData->d + kGmOffset + cvLoopIdx * tilingData->kUbSize], + xLocal, curRowFactor, tilingData->kUbSize, tilingData->hcMult * tilingData->d - tilingData->kUbSize); + xQue.template EnQue(xLocal); + xLocal = xQue.template DeQue(); + if (kGmOffset == kGmStartOffset && cvLoopIdx == 0) { + VFProcessCastAndInvRmsPart1(rmsNormLocal, xCastLocal, xLocal, coeff, curRowFactor, tilingData->kUbSize); + } else { + VFProcessCastAndInvRmsPart1(rmsNormLocal, xCastLocal, xLocal, coeff, curRowFactor, tilingData->kUbSize); + } + xQue.template FreeTensor(xLocal); + + WaitFlag(static_cast(bufferIdx & 1)); + VFTransND2NZ(xNd2NzLocal[nd2NzBufSize * (bufferIdx & 1)], xCastLocal, curRowFactor, tilingData->kUbSize); + SetFlag(static_cast(bufferIdx & 1)); + WaitFlag(static_cast(bufferIdx & 1)); + + if (curBlockIdx % 2 == 0) { + DataCopyParams dataCopyXParams; + dataCopyXParams.blockCount = CeilDiv(tilingData->kUbSize, C0_SIZE); + dataCopyXParams.blockLen = curRowFactor * C0_SIZE * sizeof(float) / BLOCK_SIZE; + dataCopyXParams.srcStride = CeilAlign(curRowFactor, C0_SIZE) - curRowFactor; + dataCopyXParams.dstStride = CeilAlign(mL1RealSize, 16) - curRowFactor; + CopyToL1(xNd2NzLocal[nd2NzBufSize * (bufferIdx & 1)], xL1_[(aL1BufferID_ * L1_BUF_OFFSET) + cvLoopIdx * tilingData->kUbSize * mL1SizeAlign], dataCopyXParams); + } else { + DataCopyParams dataCopyXParams; + dataCopyXParams.blockCount = CeilDiv(tilingData->kUbSize, C0_SIZE); + dataCopyXParams.blockLen = curRowFactor * C0_SIZE * sizeof(float) / BLOCK_SIZE; + dataCopyXParams.srcStride = CeilAlign(curRowFactor, C0_SIZE) - curRowFactor; + dataCopyXParams.dstStride = CeilAlign(mL1RealSize, 16) - curRowFactor; + CopyToL1(xNd2NzLocal[nd2NzBufSize * (bufferIdx & 1)], xL1_[(aL1BufferID_ * L1_BUF_OFFSET) + rowFactor * (BLOCK_SIZE / sizeof(float)) + cvLoopIdx * tilingData->kUbSize * mL1SizeAlign], dataCopyXParams); + } + SetFlag(static_cast(bufferIdx & 1)); + bufferIdx++; + } + CrossCoreSetFlag(SYNC_AIV_AIC_FLAG); + } + aL1BufferID_ ^= 1; + } + + if ASCEND_IS_AIV { + WaitFlag(static_cast(0)); + WaitFlag(static_cast(1)); + CrossCoreWaitFlag(SYNC_AIC_AIV_PRE_POST_FLAG); + // mm计算结果存入mmXLocal，mmXLocal每轮循环需要累加; + tbufPool1.Reset(); + tbufPool1.InitBuffer(xQue, 2, tilingData->rowInnerFactor * tilingData->hcMult * RoundUp(tilingData->dFactor) * sizeof(T)); + tbufPool1.InitBuffer( + yQue, 2, tilingData->rowInnerFactor * RoundUp(tilingData->dFactor) * sizeof(T)); + tbufPool1.InitBuffer(postQue, 2, tilingData->rowInnerFactor * tilingData->hcMultAlign * sizeof(float)); + tbufPool1.InitBuffer(combFragQue, DOUBLE_BUFFER, + tilingData->rowInnerFactor * tilingData->hcMult * tilingData->hcMult * sizeof(float)); + + // TBuf + tbufPool1.InitBuffer(mixesBuf, tilingData->rowInnerFactor * RoundUp(tilingData->hcMix) * sizeof(float)); + + mixesLocal = mixesBuf.Get(); + + SetWaitFlag(HardEvent::V_MTE2); + + // m内层循环 + int64_t currentRow = mL1RealSize / 2; + if (mL1RealSize % 2 == 1 && curBlockIdx % 2 == 0) { + // m不整除时偶数核多处理一行 + currentRow += 1; + } + for (int64_t innerRowIdx = 0; innerRowIdx < currentRow; innerRowIdx += tilingData->rowInnerFactor) { + int64_t currentInnerRowFactor = innerRowIdx + tilingData->rowInnerFactor >= currentRow ? currentRow - innerRowIdx : + tilingData->rowInnerFactor; + VFProcessInvRmsPart3(mixesLocal, mmXLocal[innerRowIdx * tilingData->hcMix], rmsNormLocal[innerRowIdx], + tilingData->normEps, currentInnerRowFactor, tilingData->hcMix); + + VFProcessPre( + mixesLocal, mixesLocal, hcBase0Local, hcScaleGm.GetValue(0), tilingData->hcEps, + currentInnerRowFactor, tilingData->hcMult, tilingData->hcMix); + for (int64_t dLoopIdx = 0; dLoopIdx < tilingData->dLoop; dLoopIdx++) + { + int64_t curDFactor = + (dLoopIdx == tilingData->dLoop - 1) ? tilingData->tailDFactor : tilingData->dFactor; + xLocal = xQue.template AllocTensor(); + CopyIn( + xGm[xGmBaseOffset + xOutSplitOffset + roundIdx * tilingData->mL1Size * tilingData->hcMult * tilingData->d + + innerRowIdx * tilingData->hcMult * tilingData->d + dLoopIdx * tilingData->dFactor], + xLocal, currentInnerRowFactor * tilingData->hcMult, curDFactor, tilingData->d - curDFactor); + xQue.template EnQue(xLocal); + xLocal = xQue.template DeQue(); + + yLocal = yQue.template AllocTensor(); + VFProcessY(yLocal, mixesLocal, xLocal, currentInnerRowFactor, tilingData->hcMult, curDFactor, tilingData->hcMix); + xQue.template FreeTensor(xLocal); + yQue.template EnQue(yLocal); + yLocal = yQue.template DeQue(); + CopyOut(yLocal, yGm[yGmBaseOffset + ySplitOffset + roundIdx * tilingData->mL1Size * tilingData->d + innerRowIdx * tilingData->d + dLoopIdx * tilingData->dFactor], + currentInnerRowFactor, curDFactor, tilingData->d - curDFactor); + yQue.template FreeTensor(yLocal); + } + + // post + postLocal = postQue.AllocTensor(); + VFProcessPost( + postLocal, mixesLocal[tilingData->hcMult], hcBase1Local, + hcScaleGm.GetValue(1), tilingData->hcEps, currentInnerRowFactor, tilingData->hcMult, tilingData->hcMix); + + postQue.EnQue(postLocal); + postLocal = postQue.DeQue(); + CopyOut(postLocal, postGm[postGmBaseOffset + postSplitOffset + roundIdx * tilingData->mL1Size * tilingData->hcMult + innerRowIdx * tilingData->hcMult], currentInnerRowFactor, tilingData->hcMult); + postQue.FreeTensor(postLocal); + + // combFrag + combFragLocal = combFragQue.AllocTensor(); + VFProcessCombFragPacked( + combFragLocal, mixesLocal[tilingData->hcMult * 2], hcBase2Local, hcScaleGm.GetValue(2), tilingData->hcEps, + tilingData->iterTimes - 1, currentInnerRowFactor, tilingData->hcMult, tilingData->hcMix); + + combFragQue.EnQue(combFragLocal); + combFragLocal = combFragQue.DeQue(); + CopyOut(combFragLocal, combFragGm[combFragGmBaseOffset + combFragSplitOffset + roundIdx * tilingData->mL1Size * tilingData->hcMult * tilingData->hcMult + innerRowIdx * tilingData->hcMult * tilingData->hcMult], + currentInnerRowFactor, tilingData->hcMult * tilingData->hcMult); + combFragQue.FreeTensor(combFragLocal); + } + SetFlag(static_cast(0)); + } + } + if ASCEND_IS_AIV { + WaitFlag(static_cast(0)); + } else { + mmService_.End(); + } + } + +private: + TPipe *pipe; + const HcPreTilingData *tilingData; + // (M, K) * (N, K) + + GlobalTensor xGm; + GlobalTensor hcFnGm; + GlobalTensor workspaceGm; + GlobalTensor yGm; + GlobalTensor invRmsGm; + GlobalTensor hcScaleGm; + GlobalTensor hcBaseGm; + GlobalTensor postGm; + GlobalTensor combFragGm; + + TQue xQue; + TQue yQue; + TQue postQue; + TQue combFragQue; + + TBuf castBuf; + TBuf nd2NzBuf; + + TQue squareSumQue; + + TBuf hcBaseBuf0; + TBuf hcBaseBuf1; + TBuf hcBaseBuf2; + + TBuf rowBrcbBuf0; + TBuf hcBrcbBuf1; + TBuf reduceBuf; + + TBuf rsqrtBuf; + TBuf squareReduceBuf; + TBuf mixes01ReduceBuf; + + TBuf xCastBuf; + TBuf yCastBuf; + + TBuf mixesBuf; + TBuf rmsNormBuf; + TBuf mmXBuf; + + LocalTensor xLocal; + LocalTensor yLocal; + LocalTensor mmXLocal; + LocalTensor rmsNormLocal; + LocalTensor xCastLocal; + LocalTensor xNd2NzLocal; + + LocalTensor mixesLocal; + LocalTensor rmsAndmmLocal; + LocalTensor postLocal; + LocalTensor combFragLocal; + LocalTensor hcBase0Local; + LocalTensor hcBase1Local; + LocalTensor hcBase2Local; + + HcPreCubeCompute mmService_; + LocalTensor xL1_; + LocalTensor wL1_; + static constexpr uint64_t SYNC_AIV_AIC_FLAG = 8; + static constexpr uint64_t SYNC_AIC_AIV_FLAG = 9; + static constexpr uint64_t SYNC_AIC_AIV_PRE_POST_FLAG = 10; + static constexpr uint64_t FLAG_ID_MAX = 16; + uint64_t cvLoopIdx_ = 0; + uint8_t aL1BufferID_{0}; + + TBufPool tbufPool0; + TBufPool tbufPool1; +}; + +} // namespace HCPreSinkhorn + +#endif diff --git a/csrc/ascend/moe/hc_pre_inv_rms/CMakeLists.txt b/csrc/ascend/moe/hc_pre_inv_rms/CMakeLists.txt new file mode 100644 index 000000000..86b308249 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_inv_rms/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_inv_rms/op_host/CMakeLists.txt b/csrc/ascend/moe/hc_pre_inv_rms/op_host/CMakeLists.txt new file mode 100644 index 000000000..06c683007 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_inv_rms/op_host/CMakeLists.txt @@ -0,0 +1,63 @@ +# This program is free software, you can redistribute it and/or modify it. +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This file is a part of the CANN Open Software. +# Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ====================================================================================================================== + +# add_ops_compile_options( +# OP_NAME HcPreInvRms +# OPTIONS --cce-auto-sync=off +# -Wno-deprecated-declarations +# -Werror +# -mllvm -cce-aicore-hoist-movemask=false +# --op_relocatable_kernel_binary=true +# ) + +# set(hc_pre_inv_rms_depends nn/norm/hc_pre_inv_rms PARENT_SCOPE) + +# target_sources(op_host_aclnn PRIVATE +# op_host/hc_pre_inv_rms_def.cpp +# ) + +# target_sources(optiling PRIVATE +# op_host/hc_pre_inv_rms_tiling.cpp +# ) + +# if (NOT BUILD_OPEN_PROJECT) +# target_sources(opmaster_ct PRIVATE +# op_host/hc_pre_inv_rms_tiling.cpp +# ) +# endif () + +# target_include_directories(optiling PRIVATE +# ${CMAKE_CURRENT_SOURCE_DIR}/op_host +# ) + +# target_sources(opsproto PRIVATE +# op_host/hc_pre_inv_rms_proto.cpp +# ) + +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + hc_pre_inv_rms_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME HcPreInvRms + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true +) + + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE hc_pre_inv_rms ACLNNTYPE aclnn) +endif() \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_def.cpp b/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_def.cpp new file mode 100644 index 000000000..0af0bc721 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_def.cpp @@ -0,0 +1,50 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_inv_rms.cpp + * \brief HcPreInvRms op host config + */ + +#include +#include "register/op_def_registry.h" + +namespace ops { +class HcPreInvRms : public OpDef { +public: + explicit HcPreInvRms(const char *name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + + this->Attr("epsilon") + .AttrType(OPTIONAL) + .Float(1e-6f); + + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + + OpAICoreConfig regbaseCfg; + regbaseCfg.DynamicCompileStaticFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .ExtendCfgInfo("opFile.value", "hc_pre_inv_rms"); + } +}; +OP_ADD(HcPreInvRms); +} // namespace ops diff --git a/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_proto.cpp b/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_proto.cpp new file mode 100644 index 000000000..94daae592 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_proto.cpp @@ -0,0 +1,67 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_inv_rms_proto.cpp + * \brief + */ + +#include +#include +#include "error/ops_error.h" + +using namespace ge; +namespace ops { +const int32_t INPUT_IDX_X = 0; +const int32_t INDEX_OUTPUT_Y = 0; +const static int64_t DIM_0 = 0; +const static int64_t DIM_1 = 1; +const static int64_t DIM_2 = 2; +const static int64_t DIM_3 = 3; +const static int64_t BS_INPUT_DIM_NUM = 4; +const static int64_t TND_INPUT_DIM_NUM = 3; + +static ge::graphStatus InferShape4HcPreInvRms(gert::InferShapeContext* context) +{ + OPS_LOG_I(context->GetNodeName(), "Begin to do InferShape4HcPreInvRms."); + + const gert::Shape* xShape = context->GetInputShape(INPUT_IDX_X); + OPS_LOG_E_IF_NULL(context, xShape, return ge::GRAPH_FAILED); + auto xDimNum = xShape->GetDimNum(); + + auto yShape = context->GetOutputShape(INDEX_OUTPUT_Y); + // The first one or two dimensions of y match those of x, and the last dimension of y is 1. + // x: (b, s, hc, d) --> y: (b, s, 1) or x: (b * s, hc, d) --> y: (b * s, 1) + yShape->SetDimNum(xDimNum); + if (xDimNum == BS_INPUT_DIM_NUM) { + yShape->SetDim(DIM_0, xShape->GetDim(DIM_0)); + yShape->SetDim(DIM_1, xShape->GetDim(DIM_1)); + yShape->SetDim(DIM_2, 1); + } else if (xDimNum == TND_INPUT_DIM_NUM) { + yShape->SetDim(DIM_0, xShape->GetDim(DIM_0)); + yShape->SetDim(DIM_1, 1); + } + + OPS_LOG_I(context->GetNodeName(), "End to do InferShape4HcPreInvRms"); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDtype4HcPreInvRms(gert::InferDataTypeContext* context) +{ + OPS_LOG_I(context->GetNodeName(), "InferDtype4HcHost enter"); + context->SetOutputDataType(INDEX_OUTPUT_Y, ge::DT_FLOAT); + OPS_LOG_I(context->GetNodeName(), "InferDtype4HcPreInvRms end"); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(HcPreInvRms) + .InferShape(InferShape4HcPreInvRms) + .InferDataType(InferDtype4HcPreInvRms); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling.cpp b/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling.cpp new file mode 100644 index 000000000..9f23dd7d8 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling.cpp @@ -0,0 +1,359 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file hc_pre_inv_rms_tiling.cpp + * \brief + */ + #include +#include "hc_pre_inv_rms_tiling.h" +#include "hc_pre_inv_rms_tiling_arch35.h" +#include "hc_pre_inv_rms_tiling_large_d.h" + +namespace optiling { +const static int64_t DEFAULT_WORKSPACE_SIZE = 16777216; // 预留16M空间 +const static int64_t X_INPUT_INDEX = 0; +const static int64_t Y_OUTPUT_INDEX = 0; +const static int64_t EPS_ATTR_INDEX = 0; +const static size_t X_INPUT_BS_FUSED_DIMS = 3; +const static size_t X_INPUT_DIMS = 4; +const static int64_t UB_BLOCK_SIZE = 32; +const static uint64_t TILING_KEY_FULL_LOAD = 1000; +const static int64_t DIM_0 = 0; +const static int64_t DIM_1 = 1; +const static int64_t DIM_2 = 2; +const static int64_t DIM_3 = 3; +const static int64_t B16_TYPE_BYTE_SIZE = 2; +const static int64_t B32_TYPE_BYTE_SIZE = 4; + +template +static inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd))); +} + + +template +static inline T CeilAlign(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)) * (rnd)); +} + +class HcPreInvRmsTilingBase { +public: + explicit HcPreInvRmsTilingBase(gert::TilingContext *context) : context_(context) + { + Reset(); + } + ~HcPreInvRmsTilingBase() = default; + + bool IsCapable() + { + return true; + } + // 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小 + ge::graphStatus GetPlatformInfo(); + // 2、获取INPUT/OUTPUT/ATTR信息 + ge::graphStatus GetShapeAttrsInfo(); + // 3、计算数据切分TilingData + ge::graphStatus DoOpTiling(); + // 4、计算高阶API的TilingData + ge::graphStatus DoLibApiTiling(); + // 5、计算TilingKey + uint64_t GetTilingKey() const; + // 6、计算Workspace 大小 + ge::graphStatus GetWorkspaceSize(); + // 7、保存Tiling数据 + ge::graphStatus PostTiling(); + void Reset(); + +private: + ge::graphStatus CheckInputShape(); + ge::graphStatus CheckAttr(); + ge::graphStatus CheckOutShape(); + void SplitA(); + void CalUbFactorA(); + + const gert::Shape *xShape_ = nullptr; + const gert::Shape *yShape_ = nullptr; + + float eps_ = 1e-6f; + int64_t A_ = 0; + int64_t R_ = 0; + + int64_t inputDtypeSize_; + int64_t outputDtypeSize_; + const char *opName_ = ""; + HcPreInvRmsFullLoadTilingData invRmsTilingData_; + gert::TilingContext *context_ = nullptr; + uint64_t workspaceSize_ = 0; + + uint64_t coreNum_ = 0; + int64_t ubSize_ = 0; + int64_t ubBlockSize_ = 0; +}; + +ge::graphStatus HcPreInvRmsTilingBase::CheckInputShape() +{ + size_t xDimNum = xShape_->GetDimNum(); + OPS_ERR_IF(xDimNum != X_INPUT_DIMS && xDimNum != X_INPUT_BS_FUSED_DIMS, + OPS_LOG_E(context_, "The dim number of x is: %zu, but it should be %zu or %zu(bs fused)." + , xDimNum, X_INPUT_DIMS, X_INPUT_BS_FUSED_DIMS), + return ge::GRAPH_FAILED); + + if (xDimNum == X_INPUT_DIMS) { + A_ = xShape_->GetDim(DIM_0) * xShape_->GetDim(DIM_1); + R_ = xShape_->GetDim(DIM_2) * xShape_->GetDim(DIM_3); + } else if (xDimNum == X_INPUT_BS_FUSED_DIMS) { + A_ = xShape_->GetDim(DIM_0); + R_ = xShape_->GetDim(DIM_1) * xShape_->GetDim(DIM_2); + } + + invRmsTilingData_.set_A(A_); + invRmsTilingData_.set_R(R_); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingBase::CheckAttr() +{ + OPS_ERR_IF(eps_ <= 0, OPS_LOG_E(context_, "epsilon is: %ld, but it should not be less than 0.", eps_), return ge::GRAPH_FAILED); + invRmsTilingData_.set_epsilon(eps_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingBase::GetShapeAttrsInfo() +{ + opName_ = context_->GetNodeName(); + // 获取输入shape信息 + auto xShapePtr = context_->GetInputShape(X_INPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, xShapePtr, return ge::GRAPH_FAILED); + xShape_ = &xShapePtr->GetStorageShape(); + + // 获取输出shape + auto yShapePtr = context_->GetOutputShape(Y_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, yShapePtr, return ge::GRAPH_FAILED); + yShape_ = &yShapePtr->GetStorageShape(); + + // 获取输入dtype + auto xDesc = context_->GetInputDesc(X_INPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, xDesc, return ge::GRAPH_FAILED); + auto xDtype = xDesc->GetDataType(); + OPS_ERR_IF( + (xDtype != ge::DataType::DT_FLOAT && xDtype != ge::DataType::DT_FLOAT16 && xDtype != ge::DataType::DT_BF16), + OPS_LOG_E(context_, "x dtype %s error, only supports float32, float16 and bfloat16. please check.", + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + + // 获取输出dtype + auto yDesc = context_->GetOutputDesc(Y_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, yDesc, return ge::GRAPH_FAILED); + auto yDtype = yDesc->GetDataType(); + OPS_ERR_IF((yDtype != ge::DataType::DT_FLOAT), + OPS_LOG_E(context_, "y out dtype %s error, only support float32, please check", + ge::TypeUtils::DataTypeToSerialString(yDtype).c_str()), + return ge::GRAPH_FAILED); + + // 获取属性 + auto attrs = context_->GetAttrs(); + OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED); + + const float *epsPtr = attrs->GetAttrPointer(EPS_ATTR_INDEX); + if (epsPtr != nullptr) { + eps_ = *epsPtr; + } + OPS_LOG_I(context_, "Attr eps is: %f ", eps_); + + inputDtypeSize_ = static_cast(ge::GetSizeByDataType(context_->GetInputDesc(X_INPUT_INDEX)->GetDataType())); + outputDtypeSize_ = static_cast(ge::GetSizeByDataType(context_->GetOutputDesc(Y_OUTPUT_INDEX)->GetDataType())); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingBase::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, OPS_LOG_E(context_, "fail to get platform info"), return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + coreNum_ = ascendcPlatform.GetCoreNumAiv(); + OPS_ERR_IF( + coreNum_ <= 0, OPS_LOG_E(context_->GetNodeName(), "coreNum must be greater than 0."), + return ge::GRAPH_FAILED); + // 获取UB大小 + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + ubSize_ = static_cast(ubSizePlatForm); + OPS_ERR_IF( + ubSize_ <= 0, OPS_LOG_E(context_->GetNodeName(), "ubSize must be greater than 0."), + return ge::GRAPH_FAILED); + + ubBlockSize_ = UB_BLOCK_SIZE; // 32: ub block size + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingBase::CheckOutShape() +{ + OPS_ERR_IF((yShape_->GetDim(0) != xShape_->GetDim(0)), + OPS_LOG_E(context_, "y out dim[0] %ld not equal x dim[0] %ld, please check.", yShape_->GetDim(0), + xShape_->GetDim(0)), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void HcPreInvRmsTilingBase::SplitA() +{ + int64_t blockFactorA = CeilDiv(A_, static_cast(coreNum_)); + int64_t blockNumA = CeilDiv(A_, blockFactorA); + int64_t blockTailFactorA = A_ % blockFactorA == 0 ? blockFactorA : A_ % blockFactorA; + invRmsTilingData_.set_blockNumA(blockNumA); + invRmsTilingData_.set_blockFactorA(blockFactorA); + invRmsTilingData_.set_blockTailFactorA(blockTailFactorA); + int64_t ubFactorA = invRmsTilingData_.get_ubFactorA(); + if (ubFactorA > blockFactorA) { + invRmsTilingData_.set_ubFactorA(blockFactorA); + } +} + +void HcPreInvRmsTilingBase::CalUbFactorA() +{ + int64_t rAlignSize = CeilAlign(R_ * inputDtypeSize_, UB_BLOCK_SIZE); + int64_t ubFactorA = 1; + if (inputDtypeSize_ == B16_TYPE_BYTE_SIZE) { + ubFactorA = ubSize_ / (4 * rAlignSize + 2 * outputDtypeSize_ + R_ / 16); + } else if (inputDtypeSize_ == B32_TYPE_BYTE_SIZE) { + ubFactorA = ubSize_ / (2 * rAlignSize + 2 * outputDtypeSize_ + R_ / 16); + } + invRmsTilingData_.set_ubFactorA(ubFactorA); +} + +ge::graphStatus HcPreInvRmsTilingBase::DoOpTiling() +{ + auto ret = GetPlatformInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = GetShapeAttrsInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckInputShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckOutShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckAttr(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + CalUbFactorA(); + SplitA(); + + ret = PostTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingBase::DoLibApiTiling() +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingBase::GetWorkspaceSize() +{ + // 计算workspace大小 + workspaceSize_ = DEFAULT_WORKSPACE_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingBase::PostTiling() +{ + context_->SetTilingKey(GetTilingKey()); + context_->SetBlockDim(invRmsTilingData_.get_blockNumA()); + size_t *currentWorkspace = context_->GetWorkspaceSizes(1); + currentWorkspace[0] = workspaceSize_; + invRmsTilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), + context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(invRmsTilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +uint64_t HcPreInvRmsTilingBase::GetTilingKey() const +{ + return TILING_KEY_FULL_LOAD; +} + +void HcPreInvRmsTilingBase::Reset() +{ + opName_ = nullptr; + return; +} + +ge::graphStatus TilingForHcPreInvRms(gert::TilingContext *context) +{ + OPS_LOG_I(context, "TilingForHcPreInvRms start"); + OPS_ERR_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("TilingForHcPreInvRms", "Tiling context is null"), + return ge::GRAPH_FAILED); + + auto platformInfo = context->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, OPS_REPORT_VECTOR_INNER_ERR("TilingForHcPreInvRms", "Tiling platformInfo is null"), + return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + auto socVersion = ascendcPlatform.GetSocVersion(); + if (socVersion == platform_ascendc::SocVersion::ASCEND950) { + OPS_LOG_I(context, "Using arch35 tiling for ASCEND950"); + HcPreInvRmsRegbase::HcPreInvRmsTilingRegbase hcPreInvRmsTilingRegbase(context); + return hcPreInvRmsTilingRegbase.DoOpTiling(); + } + + auto xShapePtr = context->GetInputShape(0); + if (xShapePtr == nullptr) { + HcPreInvRmsTilingBase invRmsTilingBase(context); + return invRmsTilingBase.DoOpTiling(); + } + auto &xShape = xShapePtr->GetStorageShape(); + size_t xDimNum = xShape.GetDimNum(); + int64_t R = 0; + if (xDimNum == X_INPUT_DIMS) { + R = xShape.GetDim(DIM_2) * xShape.GetDim(DIM_3); + } else if (xDimNum == X_INPUT_BS_FUSED_DIMS) { + R = xShape.GetDim(DIM_1) * xShape.GetDim(DIM_2); + } + + if (R == 28672) { + OPS_LOG_I(context, "Using large_d tiling for R=28672"); + HcPreInvRmsLargeD::HcPreInvRmsTilingLargeD invRmsTilingLargeD(context); + return invRmsTilingLargeD.DoOpTiling(); + } + + HcPreInvRmsTilingBase invRmsTilingBase(context); + return invRmsTilingBase.DoOpTiling(); +} + +static ge::graphStatus TilingPrepareForHcPreInvRms(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(HcPreInvRms) + .Tiling(TilingForHcPreInvRms) + .TilingParse(TilingPrepareForHcPreInvRms); +} // namespace optiling \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling.h b/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling.h new file mode 100644 index 000000000..a3e556144 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling.h @@ -0,0 +1,52 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_inv_rms_tiling.h + * \brief + */ +#ifndef HC_PRE_INV_RMS_TILING_H_ +#define HC_PRE_INV_RMS_TILING_H_ + +#include +#include +#include +#include +#include +#include "register/op_impl_registry.h" +#include "platform/platform_infos_def.h" +#include "exe_graph/runtime/tiling_context.h" +#include "tiling/platform/platform_ascendc.h" +#include "register/op_def_registry.h" +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "error/ops_error.h" +#include "platform/platform_info.h" + +namespace optiling { + + +BEGIN_TILING_DATA_DEF(HcPreInvRmsFullLoadTilingData) +TILING_DATA_FIELD_DEF(int64_t, A); // A轴大小 +TILING_DATA_FIELD_DEF(int64_t, R); // R轴大小 +TILING_DATA_FIELD_DEF(int64_t, blockNumA); // 使用核数 +TILING_DATA_FIELD_DEF(int64_t, blockFactorA); // 每个核处理的A个数 +TILING_DATA_FIELD_DEF(int64_t, blockTailFactorA); // 尾核处理的A个数 +TILING_DATA_FIELD_DEF(int64_t, ubFactorA); // 每次UB循环处理的A个数 +TILING_DATA_FIELD_DEF(float, epsilon); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(HcPreInvRms, HcPreInvRmsFullLoadTilingData) + +struct HcPreInvRmsCompileInfo {}; + +} // namespace optiling + +#endif // HC_PRE_INV_RMS_TILING_H_ diff --git a/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling_arch35.h b/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling_arch35.h new file mode 100644 index 000000000..2694597b3 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling_arch35.h @@ -0,0 +1,310 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file hc_pre_inv_rms_tiling_arch35.cpp + * \brief + */ +#include +#include "hc_pre_inv_rms_tiling.h" + +namespace optiling { +namespace HcPreInvRmsRegbase{ + +const static int64_t DEFAULT_WORKSPACE_SIZE = 16777216; // 预留16M空间 +const static int64_t X_INPUT_INDEX = 0; +const static int64_t Y_OUTPUT_INDEX = 0; +const static int64_t EPS_ATTR_INDEX = 0; +const static size_t X_INPUT_BS_FUSED_DIMS = 3; +const static size_t X_INPUT_DIMS = 4; +const static int64_t UB_BLOCK_SIZE = 32; +const static uint64_t FULL_LOAD_REGBASE_TILING_KEY = 2000; +const static int64_t DIM_0 = 0; +const static int64_t DIM_1 = 1; +const static int64_t DIM_2 = 2; +const static int64_t DIM_3 = 3; + +template +static inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd))); +} + + +template +static inline T CeilAlign(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)) * (rnd)); +} + +class HcPreInvRmsTilingRegbase { +public: + explicit HcPreInvRmsTilingRegbase(gert::TilingContext *context) : context_(context) + { + Reset(); + } + ~HcPreInvRmsTilingRegbase() = default; + + bool IsCapable() + { + if (socVersion_ != platform_ascendc::SocVersion::ASCEND950) { + return false; + } + return true; + } + // 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小 + ge::graphStatus GetPlatformInfo(); + // 2、获取INPUT/OUTPUT/ATTR信息 + ge::graphStatus GetShapeAttrsInfo(); + // 3、计算数据切分TilingData + ge::graphStatus DoOpTiling(); + // 4、计算高阶API的TilingData + ge::graphStatus DoLibApiTiling(); + // 5、计算TilingKey + uint64_t GetTilingKey() const; + // 6、计算Workspace 大小 + ge::graphStatus GetWorkspaceSize(); + // 7、保存Tiling数据 + ge::graphStatus PostTiling(); + void Reset(); + +private: + ge::graphStatus CheckInputShape(); + ge::graphStatus CheckAttr(); + ge::graphStatus CheckOutShape(); + void SplitA(); + void CalUbFactorA(); + + const gert::Shape *xShape_ = nullptr; + const gert::Shape *yShape_ = nullptr; + + float eps_ = 1e-6f; + int64_t A_ = 0; + int64_t R_ = 0; + + int64_t inputDtypeSize_; + int64_t outputDtypeSize_; + const char *opName_ = ""; + HcPreInvRmsFullLoadTilingData invRmsTilingData_; + gert::TilingContext *context_ = nullptr; + uint64_t workspaceSize_ = 0; + + uint64_t coreNum_ = 0; + int64_t ubSize_ = 0; + int64_t ubBlockSize_ = 0; + platform_ascendc::SocVersion socVersion_; +}; + +ge::graphStatus HcPreInvRmsTilingRegbase::CheckInputShape() +{ + size_t xDimNum = xShape_->GetDimNum(); + OPS_ERR_IF(xDimNum != X_INPUT_DIMS && xDimNum != X_INPUT_BS_FUSED_DIMS, + OPS_LOG_E(context_, "The dim number of x is: %zu, but it should be %zu or %zu(bs fused)." + , xDimNum, X_INPUT_DIMS, X_INPUT_BS_FUSED_DIMS), + return ge::GRAPH_FAILED); + + if (xDimNum == X_INPUT_DIMS) { + A_ = xShape_->GetDim(DIM_0) * xShape_->GetDim(DIM_1); + R_ = xShape_->GetDim(DIM_2) * xShape_->GetDim(DIM_3); + } else if (xDimNum == X_INPUT_BS_FUSED_DIMS) { + A_ = xShape_->GetDim(DIM_0); + R_ = xShape_->GetDim(DIM_1) * xShape_->GetDim(DIM_2); + } + + invRmsTilingData_.set_A(A_); + invRmsTilingData_.set_R(R_); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingRegbase::CheckAttr() +{ + OPS_ERR_IF(eps_ <= 0, OPS_LOG_E(context_, "epsilon is: %ld, but it should not be less than 0.", eps_), return ge::GRAPH_FAILED); + invRmsTilingData_.set_epsilon(eps_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingRegbase::GetShapeAttrsInfo() +{ + opName_ = context_->GetNodeName(); + // 获取输入shape信息 + auto xShapePtr = context_->GetInputShape(X_INPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, xShapePtr, return ge::GRAPH_FAILED); + xShape_ = &xShapePtr->GetStorageShape(); + + // 获取输出shape + auto yShapePtr = context_->GetOutputShape(Y_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, yShapePtr, return ge::GRAPH_FAILED); + yShape_ = &yShapePtr->GetStorageShape(); + + // 获取输入dtype + auto xDesc = context_->GetInputDesc(X_INPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, xDesc, return ge::GRAPH_FAILED); + auto xDtype = xDesc->GetDataType(); + OPS_ERR_IF( + (xDtype != ge::DataType::DT_FLOAT && xDtype != ge::DataType::DT_FLOAT16 && xDtype != ge::DataType::DT_BF16), + OPS_LOG_E(context_, "x dtype %s error, only supports float32, float16 and bfloat16. please check.", + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + + // 获取输出dtype + auto yDesc = context_->GetOutputDesc(Y_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, yDesc, return ge::GRAPH_FAILED); + auto yDtype = yDesc->GetDataType(); + OPS_ERR_IF((yDtype != ge::DataType::DT_FLOAT), + OPS_LOG_E(context_, "y out dtype %s error, only support float32, please check", + ge::TypeUtils::DataTypeToSerialString(yDtype).c_str()), + return ge::GRAPH_FAILED); + + // 获取属性 + auto attrs = context_->GetAttrs(); + OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED); + + const float *epsPtr = attrs->GetAttrPointer(EPS_ATTR_INDEX); + if (epsPtr != nullptr) { + eps_ = *epsPtr; + } + OPS_LOG_I(context_, "Attr eps is: %f ", eps_); + + inputDtypeSize_ = static_cast(ge::GetSizeByDataType(context_->GetInputDesc(X_INPUT_INDEX)->GetDataType())); + outputDtypeSize_ = static_cast(ge::GetSizeByDataType(context_->GetOutputDesc(Y_OUTPUT_INDEX)->GetDataType())); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingRegbase::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, OPS_LOG_E(context_, "fail to get platform info"), return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + coreNum_ = ascendcPlatform.GetCoreNumAiv(); + socVersion_ = ascendcPlatform.GetSocVersion(); + OPS_ERR_IF( + coreNum_ <= 0, OPS_LOG_E(context_->GetNodeName(), "coreNum must be greater than 0."), + return ge::GRAPH_FAILED); + // 获取UB大小 + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + ubSize_ = static_cast(ubSizePlatForm); + OPS_ERR_IF( + ubSize_ <= 0, OPS_LOG_E(context_->GetNodeName(), "ubSize must be greater than 0."), + return ge::GRAPH_FAILED); + + ubBlockSize_ = UB_BLOCK_SIZE; // 32: ub block size + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingRegbase::CheckOutShape() +{ + OPS_ERR_IF((yShape_->GetDim(0) != xShape_->GetDim(0)), + OPS_LOG_E(context_, "y out dim[0] %ld not equal x dim[0] %ld, please check.", yShape_->GetDim(0), + xShape_->GetDim(0)), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void HcPreInvRmsTilingRegbase::SplitA() +{ + int64_t blockFactorA = CeilDiv(A_, static_cast(coreNum_)); + int64_t blockNumA = CeilDiv(A_, blockFactorA); + int64_t blockTailFactorA = A_ % blockFactorA == 0 ? blockFactorA : A_ % blockFactorA; + invRmsTilingData_.set_blockNumA(blockNumA); + invRmsTilingData_.set_blockFactorA(blockFactorA); + invRmsTilingData_.set_blockTailFactorA(blockTailFactorA); + int64_t ubFactorA = invRmsTilingData_.get_ubFactorA(); + if (ubFactorA > blockFactorA) { + invRmsTilingData_.set_ubFactorA(blockFactorA); + } +} + +void HcPreInvRmsTilingRegbase::CalUbFactorA() +{ + // ubFactorA * R_AlignSize * 2 + ubFactorA * outputDtypeSize_ + 32 = ubSize_ + int64_t rAlignSize = CeilAlign(R_ * inputDtypeSize_, UB_BLOCK_SIZE); + int64_t ubFactorA = (ubSize_ - UB_BLOCK_SIZE) / (2 * rAlignSize + outputDtypeSize_); + invRmsTilingData_.set_ubFactorA(ubFactorA); +} + +ge::graphStatus HcPreInvRmsTilingRegbase::DoOpTiling() +{ + auto ret = GetPlatformInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = GetShapeAttrsInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckInputShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckOutShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckAttr(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + CalUbFactorA(); + SplitA(); + + ret = PostTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingRegbase::DoLibApiTiling() +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingRegbase::GetWorkspaceSize() +{ + // 计算workspace大小 + workspaceSize_ = DEFAULT_WORKSPACE_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingRegbase::PostTiling() +{ + context_->SetTilingKey(GetTilingKey()); + context_->SetBlockDim(invRmsTilingData_.get_blockNumA()); + size_t *currentWorkspace = context_->GetWorkspaceSizes(1); + currentWorkspace[0] = workspaceSize_; + invRmsTilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), + context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(invRmsTilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +uint64_t HcPreInvRmsTilingRegbase::GetTilingKey() const +{ + // 全载场景 + return FULL_LOAD_REGBASE_TILING_KEY; +} + +void HcPreInvRmsTilingRegbase::Reset() +{ + opName_ = nullptr; + return; +} + +} // namespace HcPreInvRmsRegbase +} // namespace optiling \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling_large_d.h b/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling_large_d.h new file mode 100644 index 000000000..202cc3813 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_inv_rms/op_host/hc_pre_inv_rms_tiling_large_d.h @@ -0,0 +1,291 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file hc_pre_inv_rms_tiling_large_d.cpp + * \brief tiling for large d (R=28672, d=7168) + */ +#include +#include "hc_pre_inv_rms_tiling.h" + +namespace optiling { +namespace HcPreInvRmsLargeD{ + +const static int64_t DEFAULT_WORKSPACE_SIZE = 16777216; +const static int64_t X_INPUT_INDEX = 0; +const static int64_t Y_OUTPUT_INDEX = 0; +const static int64_t EPS_ATTR_INDEX = 0; +const static size_t X_INPUT_BS_FUSED_DIMS = 3; +const static size_t X_INPUT_DIMS = 4; +const static int64_t UB_BLOCK_SIZE = 32; +const static uint64_t FULL_LOAD_LARGE_D_TILING_KEY = 1001; +const static int64_t R_LARGE_D = 28672; +const static int64_t DIM_0 = 0; +const static int64_t DIM_1 = 1; +const static int64_t DIM_2 = 2; +const static int64_t DIM_3 = 3; + +template +static inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd))); +} + +template +static inline T CeilAlign(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)) * (rnd)); +} + +class HcPreInvRmsTilingLargeD { +public: + explicit HcPreInvRmsTilingLargeD(gert::TilingContext *context) : context_(context) + { + Reset(); + } + ~HcPreInvRmsTilingLargeD() = default; + + bool IsCapable() + { + return true; + } + ge::graphStatus GetPlatformInfo(); + ge::graphStatus GetShapeAttrsInfo(); + ge::graphStatus DoOpTiling(); + ge::graphStatus DoLibApiTiling(); + uint64_t GetTilingKey() const; + ge::graphStatus GetWorkspaceSize(); + ge::graphStatus PostTiling(); + void Reset(); + +private: + ge::graphStatus CheckInputShape(); + ge::graphStatus CheckAttr(); + ge::graphStatus CheckOutShape(); + void SplitA(); + void CalUbFactorA(); + + const gert::Shape *xShape_ = nullptr; + const gert::Shape *yShape_ = nullptr; + + float eps_ = 1e-6f; + int64_t A_ = 0; + int64_t R_ = 0; + + int64_t inputDtypeSize_; + int64_t outputDtypeSize_; + const char *opName_ = ""; + HcPreInvRmsFullLoadTilingData invRmsTilingData_; + gert::TilingContext *context_ = nullptr; + uint64_t workspaceSize_ = 0; + + uint64_t coreNum_ = 0; + int64_t ubSize_ = 0; + int64_t ubBlockSize_ = 0; +}; + +ge::graphStatus HcPreInvRmsTilingLargeD::CheckInputShape() +{ + size_t xDimNum = xShape_->GetDimNum(); + OPS_ERR_IF(xDimNum != X_INPUT_DIMS && xDimNum != X_INPUT_BS_FUSED_DIMS, + OPS_LOG_E(context_, "The dim number of x is: %zu, but it should be %zu or %zu(bs fused)." + , xDimNum, X_INPUT_DIMS, X_INPUT_BS_FUSED_DIMS), + return ge::GRAPH_FAILED); + + if (xDimNum == X_INPUT_DIMS) { + A_ = xShape_->GetDim(DIM_0) * xShape_->GetDim(DIM_1); + R_ = xShape_->GetDim(DIM_2) * xShape_->GetDim(DIM_3); + } else if (xDimNum == X_INPUT_BS_FUSED_DIMS) { + A_ = xShape_->GetDim(DIM_0); + R_ = xShape_->GetDim(DIM_1) * xShape_->GetDim(DIM_2); + } + + OPS_ERR_IF(R_ != R_LARGE_D, + OPS_LOG_E(context_, "R is: %ld, but large_d tiling only supports R=%ld.", R_, R_LARGE_D), + return ge::GRAPH_FAILED); + + invRmsTilingData_.set_A(A_); + invRmsTilingData_.set_R(R_); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingLargeD::CheckAttr() +{ + OPS_ERR_IF(eps_ <= 0, OPS_LOG_E(context_, "epsilon is: %ld, but it should not be less than 0.", eps_), return ge::GRAPH_FAILED); + invRmsTilingData_.set_epsilon(eps_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingLargeD::GetShapeAttrsInfo() +{ + opName_ = context_->GetNodeName(); + auto xShapePtr = context_->GetInputShape(X_INPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, xShapePtr, return ge::GRAPH_FAILED); + xShape_ = &xShapePtr->GetStorageShape(); + + auto yShapePtr = context_->GetOutputShape(Y_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, yShapePtr, return ge::GRAPH_FAILED); + yShape_ = &yShapePtr->GetStorageShape(); + + auto xDesc = context_->GetInputDesc(X_INPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, xDesc, return ge::GRAPH_FAILED); + auto xDtype = xDesc->GetDataType(); + OPS_ERR_IF( + (xDtype != ge::DataType::DT_FLOAT && xDtype != ge::DataType::DT_FLOAT16 && xDtype != ge::DataType::DT_BF16), + OPS_LOG_E(context_, "x dtype %s error, only supports float32, float16 and bfloat16. please check.", + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + + auto yDesc = context_->GetOutputDesc(Y_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, yDesc, return ge::GRAPH_FAILED); + auto yDtype = yDesc->GetDataType(); + OPS_ERR_IF((yDtype != ge::DataType::DT_FLOAT), + OPS_LOG_E(context_, "y out dtype %s error, only support float32, please check", + ge::TypeUtils::DataTypeToSerialString(yDtype).c_str()), + return ge::GRAPH_FAILED); + + auto attrs = context_->GetAttrs(); + OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED); + + const float *epsPtr = attrs->GetAttrPointer(EPS_ATTR_INDEX); + if (epsPtr != nullptr) { + eps_ = *epsPtr; + } + OPS_LOG_I(context_, "Attr eps is: %f ", eps_); + + inputDtypeSize_ = static_cast(ge::GetSizeByDataType(context_->GetInputDesc(X_INPUT_INDEX)->GetDataType())); + outputDtypeSize_ = static_cast(ge::GetSizeByDataType(context_->GetOutputDesc(Y_OUTPUT_INDEX)->GetDataType())); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingLargeD::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, OPS_LOG_E(context_, "fail to get platform info"), return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + coreNum_ = ascendcPlatform.GetCoreNumAiv(); + OPS_ERR_IF( + coreNum_ <= 0, OPS_LOG_E(context_->GetNodeName(), "coreNum must be greater than 0."), + return ge::GRAPH_FAILED); + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + ubSize_ = static_cast(ubSizePlatForm); + OPS_ERR_IF( + ubSize_ <= 0, OPS_LOG_E(context_->GetNodeName(), "ubSize must be greater than 0."), + return ge::GRAPH_FAILED); + + ubBlockSize_ = UB_BLOCK_SIZE; + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingLargeD::CheckOutShape() +{ + OPS_ERR_IF((yShape_->GetDim(0) != xShape_->GetDim(0)), + OPS_LOG_E(context_, "y out dim[0] %ld not equal x dim[0] %ld, please check.", yShape_->GetDim(0), + xShape_->GetDim(0)), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +void HcPreInvRmsTilingLargeD::SplitA() +{ + int64_t blockFactorA = CeilDiv(A_, static_cast(coreNum_)); + int64_t blockNumA = CeilDiv(A_, blockFactorA); + int64_t blockTailFactorA = A_ % blockFactorA == 0 ? blockFactorA : A_ % blockFactorA; + invRmsTilingData_.set_blockNumA(blockNumA); + invRmsTilingData_.set_blockFactorA(blockFactorA); + invRmsTilingData_.set_blockTailFactorA(blockTailFactorA); + int64_t ubFactorA = invRmsTilingData_.get_ubFactorA(); + if (ubFactorA > blockFactorA) { + invRmsTilingData_.set_ubFactorA(blockFactorA); + } +} + +void HcPreInvRmsTilingLargeD::CalUbFactorA() +{ + invRmsTilingData_.set_ubFactorA(1); +} + +ge::graphStatus HcPreInvRmsTilingLargeD::DoOpTiling() +{ + auto ret = GetPlatformInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = GetShapeAttrsInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckInputShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckOutShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckAttr(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + CalUbFactorA(); + SplitA(); + + ret = PostTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingLargeD::DoLibApiTiling() +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingLargeD::GetWorkspaceSize() +{ + workspaceSize_ = DEFAULT_WORKSPACE_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreInvRmsTilingLargeD::PostTiling() +{ + context_->SetTilingKey(GetTilingKey()); + context_->SetBlockDim(invRmsTilingData_.get_blockNumA()); + size_t *currentWorkspace = context_->GetWorkspaceSizes(1); + currentWorkspace[0] = workspaceSize_; + invRmsTilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), + context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(invRmsTilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +uint64_t HcPreInvRmsTilingLargeD::GetTilingKey() const +{ + return FULL_LOAD_LARGE_D_TILING_KEY; +} + +void HcPreInvRmsTilingLargeD::Reset() +{ + opName_ = nullptr; + return; +} + +} // namespace HcPreInvRmsLargeD +} // namespace optiling \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms.cpp b/csrc/ascend/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms.cpp new file mode 100644 index 000000000..8b242d1f7 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms.cpp @@ -0,0 +1,50 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_inv_rms_apt.cpp + * \brief hc_pre_inv_rms kernel file + */ +#include "hc_pre_inv_rms_full_load.h" +#include "hc_pre_inv_rms_full_load_large_d.h" +#if defined(__DAV_C310__) + #include "hc_pre_inv_rms_full_load_regbase.h" + using namespace HcPreInvRmsRegbase; +#endif +#include "kernel_operator.h" +using namespace AscendC; +using namespace HcPreInvRms; +using namespace HcPreInvRmsLargeD; + +#define FULL_LOAD_TILING_KEY 1000 +#define FULL_LOAD_LARGE_D_TILING_KEY 1001 +#define REGBASE_FULL_LOAD_TILING_KEY 2000 + +extern "C" __global__ __aicore__ void hc_pre_inv_rms(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) +{ + TPipe pipe; + GET_TILING_DATA(tilingData, tiling); + if (TILING_KEY_IS(FULL_LOAD_TILING_KEY)) { + HcPreInvRmsFullLoad op; + op.Init(x, y, workspace, &tilingData, &pipe); + op.Process(); + } else if (TILING_KEY_IS(FULL_LOAD_LARGE_D_TILING_KEY)) { + HcPreInvRmsFullLoadLargeD op; + op.Init(x, y, workspace, &tilingData, &pipe); + op.Process(); + } + #if defined(__DAV_C310__) + else if (TILING_KEY_IS(REGBASE_FULL_LOAD_TILING_KEY)) { + HcPreInvRmsFullLoadRegbase op; + op.Init(x, y, workspace, &tilingData, &pipe); + op.Process(); + } + #endif +} diff --git a/csrc/ascend/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load.h b/csrc/ascend/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load.h new file mode 100644 index 000000000..e741b9149 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load.h @@ -0,0 +1,248 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_inv_rms.h + * \brief inv rms file + */ +#ifndef ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_H_ +#define ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_H_ +#include "kernel_operator.h" + +namespace HcPreInvRms { +using namespace AscendC; +constexpr int32_t BUFFER_NUM = 2; +constexpr int32_t FLOAT_BTYPE_SIZE = 4; +constexpr uint32_t PER_REPEAT_LEN_B32 = 64; +constexpr uint32_t UB_BLOCK_SIZE = 32; +constexpr int32_t B16_TYPE_BYTE_SIZE = 2; +constexpr int32_t B32_TYPE_BYTE_SIZE = 4; +constexpr int32_t ONE_COUNT = 1; +constexpr int32_t FOUR_FOLD = 4; +constexpr int32_t DST_REP_STRIDE = 1; +constexpr int32_t SRC_BLK_STRIDE = 1; +constexpr int32_t SRC_REP_STRIDE = 8; + +template +class HcPreInvRmsFullLoad { +public: + __aicore__ inline HcPreInvRmsFullLoad() {}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe); + __aicore__ inline void Process(); + __aicore__ inline void CopyIn(uint64_t idx, uint64_t curUbFactorA); + __aicore__ inline void Compute(uint64_t curUbFactorA); + __aicore__ inline void ComputeB16(uint64_t curUbFactorA); + __aicore__ inline void ComputeB32(uint64_t curUbFactorA); + __aicore__ inline void CopyOut(uint64_t idx, uint64_t curUbFactorA); + +private: + TPipe* pipe_; + + TQue inQueueX; + TQue outQueueY; + TBuf castBuf; + TBuf reduceBuf; + + GlobalTensor xGm; + GlobalTensor yGm; + + int64_t A; // 输入数据 A 轴大小 + int64_t R; // 输入数据 R 轴大小 + int64_t blockNumA; // 使用的核数 + int64_t blockFactorA; // 每个核处理的A个数 + int64_t blockTailFactorA; // 尾核处理的A个数 + int64_t ubFactorA; // 每次ub循环处理的A个数 + int32_t blockIdx_; + float epsilon; // 算子参数 + uint32_t curBlockFactorA; // 当前核处理的A个数 + uint32_t rAlign; + uint32_t rAlignB32; + uint32_t reduceBufNum; +}; + +template +__aicore__ inline void HcPreInvRmsFullLoad::Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe) +{ + A = tiling->A; + R = tiling->R; + blockNumA = tiling->blockNumA; + blockFactorA = tiling->blockFactorA; + blockTailFactorA = tiling->blockTailFactorA; + ubFactorA = tiling->ubFactorA; + epsilon = tiling->epsilon; + + rAlign = ((R * sizeof(T) + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / sizeof(T)); + rAlignB32 = ((R * FLOAT_BTYPE_SIZE + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / FLOAT_BTYPE_SIZE); + + pipe_ = pipe; + + blockIdx_ = GetBlockIdx(); + + if (blockIdx_ < blockNumA - 1) { + this->curBlockFactorA = this->blockFactorA; + } else if (blockIdx_ == blockNumA - 1) { + this->curBlockFactorA = this->blockTailFactorA; + } else { + return; + } + xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * blockFactorA * R, curBlockFactorA * R); + yGm.SetGlobalBuffer((__gm__ float*)y + blockIdx_ * blockFactorA, curBlockFactorA); + // pipe alloc memory to queue, the unit is Bytes + pipe_->InitBuffer(inQueueX, BUFFER_NUM, ubFactorA * rAlign * sizeof(T)); + pipe_->InitBuffer(outQueueY, BUFFER_NUM, ubFactorA * FLOAT_BTYPE_SIZE); + + reduceBufNum = (rAlignB32 + PER_REPEAT_LEN_B32 - 1) / PER_REPEAT_LEN_B32; + pipe_->InitBuffer(reduceBuf, ubFactorA * reduceBufNum * FLOAT_BTYPE_SIZE); + if constexpr (sizeof(T) == B16_TYPE_BYTE_SIZE) { + pipe_->InitBuffer(castBuf, ubFactorA * rAlignB32 * FLOAT_BTYPE_SIZE); + } +} + +template +__aicore__ inline void HcPreInvRmsFullLoad::Process() +{ + if (blockIdx_ >= blockNumA) { + return; + } + uint64_t aUbLoopCount = (curBlockFactorA + ubFactorA - 1) / ubFactorA; // Ub循环次数 + uint64_t tailUbFactorA = curBlockFactorA - (aUbLoopCount - 1) * ubFactorA; // 最后一次Ub循环的A轴大小 + uint64_t curUbFactorA = ubFactorA; + for (uint64_t idx = 0; idx < aUbLoopCount; idx++) { + if (idx == aUbLoopCount - 1) { + curUbFactorA = tailUbFactorA; + } + CopyIn(idx, curUbFactorA); + Compute(curUbFactorA); + CopyOut(idx, curUbFactorA); + } +} + +template +__aicore__ inline void HcPreInvRmsFullLoad::CopyIn(uint64_t idx, uint64_t curUbFactorA) +{ + LocalTensor xLocal = inQueueX.AllocTensor(); + + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + int64_t xGmStartAddr = idx * R * ubFactorA; + DataCopyExtParams dataCopyParams{ + static_cast(curUbFactorA), static_cast(R * sizeof(T)), 0, 0, 0}; + DataCopyPad(xLocal, xGm[xGmStartAddr], dataCopyParams, dataCopyPadParams); + + inQueueX.EnQue(xLocal); +} + +template +__aicore__ inline void HcPreInvRmsFullLoad::Compute(uint64_t curUbFactorA) +{ + if constexpr (sizeof(T) == B16_TYPE_BYTE_SIZE) { + ComputeB16(curUbFactorA); + } else if constexpr (sizeof(T) == B32_TYPE_BYTE_SIZE) { + ComputeB32(curUbFactorA); + } +} + +template +__aicore__ inline void HcPreInvRmsFullLoad::ComputeB16(uint64_t curUbFactorA) +{ + LocalTensor xLocal = inQueueX.DeQue(); + LocalTensor yLocal = outQueueY.AllocTensor(); + + LocalTensor castLocal = castBuf.Get(); + LocalTensor reduceLocal = reduceBuf.Get(); + + int32_t perFoldElems = (rAlignB32 + FOUR_FOLD - 1) / FOUR_FOLD; // 4096 + int32_t perFoldRepTime = (perFoldElems + PER_REPEAT_LEN_B32 - 1) / PER_REPEAT_LEN_B32; // 64 + + AscendC::Cast(castLocal, xLocal, AscendC::RoundMode::CAST_NONE, R); + PipeBarrier(); + AscendC::Mul(castLocal, castLocal, castLocal, R); + + for (int idx = 0; idx < curUbFactorA; idx++) { + + for (int j = 0; j < FOUR_FOLD; j++) { + PipeBarrier(); + WholeReduceSum(reduceLocal[idx * reduceBufNum + j * perFoldRepTime], castLocal[idx * rAlignB32 + j * perFoldElems], PER_REPEAT_LEN_B32, perFoldRepTime, + DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE); + } + + PipeBarrier(); + WholeReduceSum(reduceLocal, reduceLocal, PER_REPEAT_LEN_B32, FOUR_FOLD, DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE); + PipeBarrier(); + WholeReduceSum(yLocal[idx], reduceLocal, FOUR_FOLD, 1, DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE); + } + + float meanCof = 1.0f / R; + PipeBarrier(); + AscendC::Muls(yLocal, yLocal, meanCof, curUbFactorA); + PipeBarrier(); + AscendC::Adds(yLocal, yLocal, epsilon, curUbFactorA); + PipeBarrier(); + AscendC::Duplicate(reduceLocal, 1.0f, curUbFactorA); + PipeBarrier(); + AscendC::Sqrt(yLocal, yLocal, curUbFactorA); + PipeBarrier(); + AscendC::Div(yLocal, reduceLocal, yLocal, curUbFactorA); + + outQueueY.EnQue(yLocal); + inQueueX.FreeTensor(xLocal); +} + +template +__aicore__ inline void HcPreInvRmsFullLoad::ComputeB32(uint64_t curUbFactorA) +{ + LocalTensor xLocal = inQueueX.DeQue(); + LocalTensor yLocal = outQueueY.AllocTensor(); + + int32_t perFoldElems = (rAlignB32 + FOUR_FOLD - 1) / FOUR_FOLD; // 4096 + int32_t perFoldRepTime = (perFoldElems + PER_REPEAT_LEN_B32 - 1) / PER_REPEAT_LEN_B32; // 64 + + LocalTensor reduceLocal = reduceBuf.Get(); + PipeBarrier(); + AscendC::Mul(xLocal, xLocal, xLocal, R); + + for (int idx = 0; idx < curUbFactorA; idx++) { + for (int j = 0; j < FOUR_FOLD; j++) { + PipeBarrier(); + WholeReduceSum(reduceLocal[idx * reduceBufNum + j * perFoldRepTime], xLocal[idx * rAlignB32 + j * perFoldElems], PER_REPEAT_LEN_B32, perFoldRepTime, + DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE); + } + PipeBarrier(); + WholeReduceSum(reduceLocal, reduceLocal, PER_REPEAT_LEN_B32, FOUR_FOLD, DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE); + PipeBarrier(); + WholeReduceSum(yLocal[idx], reduceLocal, FOUR_FOLD, 1, DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE); + } + + float meanCof = 1.0f / R; + PipeBarrier(); + AscendC::Muls(yLocal, yLocal, meanCof, curUbFactorA); + PipeBarrier(); + AscendC::Adds(yLocal, yLocal, epsilon, curUbFactorA); + PipeBarrier(); + AscendC::Duplicate(reduceLocal, 1.0f, curUbFactorA); + PipeBarrier(); + AscendC::Sqrt(yLocal, yLocal, curUbFactorA); + PipeBarrier(); + AscendC::Div(yLocal, reduceLocal, yLocal, curUbFactorA); + + outQueueY.EnQue(yLocal); + inQueueX.FreeTensor(xLocal); +} + +template +__aicore__ inline void HcPreInvRmsFullLoad::CopyOut(uint64_t idx, uint64_t curUbFactorA) +{ + LocalTensor yLocal = outQueueY.DeQue(); + AscendC::DataCopyExtParams copyParams{1, static_cast(curUbFactorA * sizeof(float)), 0, 0, 0}; + DataCopyPad(yGm[idx * ubFactorA], yLocal, copyParams); + outQueueY.FreeTensor(yLocal); +} + +} // namespace HcPreInvRms +#endif // ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_H_ diff --git a/csrc/ascend/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load_large_d.h b/csrc/ascend/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load_large_d.h new file mode 100644 index 000000000..dbb732c28 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load_large_d.h @@ -0,0 +1,262 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_inv_rms_full_load_large_d.h + * \brief inv rms for large d (R=28672, d=7168) + */ +#ifndef ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_LARGE_D_H_ +#define ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_LARGE_D_H_ +#include "kernel_operator.h" + +namespace HcPreInvRmsLargeD { +using namespace AscendC; +constexpr int32_t BUFFER_NUM = 1; +constexpr int32_t FLOAT_BTYPE_SIZE = 4; +constexpr uint32_t PER_REPEAT_LEN_B32 = 64; +constexpr uint32_t UB_BLOCK_SIZE = 32; +constexpr int32_t B16_TYPE_BYTE_SIZE = 2; +constexpr int32_t B32_TYPE_BYTE_SIZE = 4; +constexpr int32_t HALf_INTERVAL = 2; +constexpr int32_t INDEX_TWO = 2; +constexpr int32_t INDEX_FOUR = 4; +constexpr int32_t INDEX_EIGHT = 8; +constexpr int32_t INDEX_SIXTEEN = 16; + +__aicore__ inline int32_t FindPowerTwo(int32_t n) +{ + n |= n >> 1; + n |= n >> INDEX_TWO; + n |= n >> INDEX_FOUR; + n |= n >> INDEX_EIGHT; + n |= n >> INDEX_SIXTEEN; + return (n + 1) >> 1; +} + +__aicore__ inline void ReduceSumHalfInterval( + const LocalTensor& dst_local, const LocalTensor& src_local, int32_t count) +{ + if (likely(count > PER_REPEAT_LEN_B32)) { + int32_t bodyCount = FindPowerTwo(count); + int32_t tailCount = count - bodyCount; + if (tailCount > 0) { + Add(src_local, src_local, src_local[bodyCount], tailCount); + PipeBarrier(); + } + while (bodyCount > PER_REPEAT_LEN_B32) { + bodyCount = bodyCount / HALf_INTERVAL; + Add(src_local, src_local, src_local[bodyCount], bodyCount); + PipeBarrier(); + } + AscendCUtils::SetMask(PER_REPEAT_LEN_B32); + } else { + AscendCUtils::SetMask(count); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + if (g_coreType == AIV) { + WholeReduceSum(dst_local, src_local, MASK_PLACEHOLDER, 1, 0, 1, 0); + } +#else + WholeReduceSum(dst_local, src_local, MASK_PLACEHOLDER, 1, 1, 1, DEFAULT_REPEAT_STRIDE); +#endif + PipeBarrier(); +} + +template +class HcPreInvRmsFullLoadLargeD { +public: + __aicore__ inline HcPreInvRmsFullLoadLargeD() {}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe); + __aicore__ inline void Process(); + __aicore__ inline void CopyIn(uint64_t idx, uint64_t curUbFactorA); + __aicore__ inline void Compute(uint64_t curUbFactorA); + __aicore__ inline void ComputeB16(uint64_t curUbFactorA); + __aicore__ inline void ComputeB32(uint64_t curUbFactorA); + __aicore__ inline void CopyOut(uint64_t idx, uint64_t curUbFactorA); + +private: + TPipe* pipe_; + + TQue inQueueX; + TQue outQueueY; + TBuf castBuf; + TBuf reduceBuf; + + GlobalTensor xGm; + GlobalTensor yGm; + + int64_t A; + int64_t R; + int64_t blockNumA; + int64_t blockFactorA; + int64_t blockTailFactorA; + int64_t ubFactorA; + int32_t blockIdx_; + float epsilon; + uint32_t curBlockFactorA; + uint32_t rAlign; + uint32_t rAlignB32; + uint32_t reduceBufNum; +}; + +template +__aicore__ inline void HcPreInvRmsFullLoadLargeD::Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe) +{ + A = tiling->A; + R = tiling->R; + blockNumA = tiling->blockNumA; + blockFactorA = tiling->blockFactorA; + blockTailFactorA = tiling->blockTailFactorA; + ubFactorA = tiling->ubFactorA; + epsilon = tiling->epsilon; + + rAlign = ((R * sizeof(T) + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / sizeof(T)); + rAlignB32 = ((R * FLOAT_BTYPE_SIZE + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / FLOAT_BTYPE_SIZE); + + pipe_ = pipe; + + blockIdx_ = GetBlockIdx(); + + if (blockIdx_ < blockNumA - 1) { + this->curBlockFactorA = this->blockFactorA; + } else if (blockIdx_ == blockNumA - 1) { + this->curBlockFactorA = this->blockTailFactorA; + } else { + return; + } + xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * blockFactorA * R, curBlockFactorA * R); + yGm.SetGlobalBuffer((__gm__ float*)y + blockIdx_ * blockFactorA, curBlockFactorA); + pipe_->InitBuffer(inQueueX, BUFFER_NUM, ubFactorA * rAlign * sizeof(T)); + pipe_->InitBuffer(outQueueY, BUFFER_NUM, ubFactorA * FLOAT_BTYPE_SIZE); + + reduceBufNum = (rAlignB32 + PER_REPEAT_LEN_B32 - 1) / PER_REPEAT_LEN_B32; + pipe_->InitBuffer(reduceBuf, ubFactorA * reduceBufNum * FLOAT_BTYPE_SIZE); + if constexpr (sizeof(T) == B16_TYPE_BYTE_SIZE) { + pipe_->InitBuffer(castBuf, ubFactorA * rAlignB32 * FLOAT_BTYPE_SIZE); + } +} + +template +__aicore__ inline void HcPreInvRmsFullLoadLargeD::Process() +{ + if (blockIdx_ >= blockNumA) { + return; + } + uint64_t aUbLoopCount = (curBlockFactorA + ubFactorA - 1) / ubFactorA; + uint64_t tailUbFactorA = curBlockFactorA - (aUbLoopCount - 1) * ubFactorA; + uint64_t curUbFactorA = ubFactorA; + for (uint64_t idx = 0; idx < aUbLoopCount; idx++) { + if (idx == aUbLoopCount - 1) { + curUbFactorA = tailUbFactorA; + } + CopyIn(idx, curUbFactorA); + Compute(curUbFactorA); + CopyOut(idx, curUbFactorA); + } +} + +template +__aicore__ inline void HcPreInvRmsFullLoadLargeD::CopyIn(uint64_t idx, uint64_t curUbFactorA) +{ + LocalTensor xLocal = inQueueX.AllocTensor(); + + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + int64_t xGmStartAddr = idx * R * ubFactorA; + DataCopyExtParams dataCopyParams{ + static_cast(curUbFactorA), static_cast(R * sizeof(T)), 0, 0, 0}; + DataCopyPad(xLocal, xGm[xGmStartAddr], dataCopyParams, dataCopyPadParams); + + inQueueX.EnQue(xLocal); +} + +template +__aicore__ inline void HcPreInvRmsFullLoadLargeD::Compute(uint64_t curUbFactorA) +{ + if constexpr (sizeof(T) == B16_TYPE_BYTE_SIZE) { + ComputeB16(curUbFactorA); + } else if constexpr (sizeof(T) == B32_TYPE_BYTE_SIZE) { + ComputeB32(curUbFactorA); + } +} + +template +__aicore__ inline void HcPreInvRmsFullLoadLargeD::ComputeB16(uint64_t curUbFactorA) +{ + LocalTensor xLocal = inQueueX.DeQue(); + LocalTensor yLocal = outQueueY.AllocTensor(); + + LocalTensor castLocal = castBuf.Get(); + LocalTensor reduceLocal = reduceBuf.Get(); + + AscendC::Cast(castLocal, xLocal, AscendC::RoundMode::CAST_NONE, R); + PipeBarrier(); + AscendC::Mul(castLocal, castLocal, castLocal, R); + + for (int idx = 0; idx < curUbFactorA; idx++) { + ReduceSumHalfInterval(yLocal[idx], castLocal[idx * rAlignB32], R); + } + + float meanCof = 1.0f / R; + PipeBarrier(); + AscendC::Muls(yLocal, yLocal, meanCof, curUbFactorA); + PipeBarrier(); + AscendC::Adds(yLocal, yLocal, epsilon, curUbFactorA); + PipeBarrier(); + AscendC::Duplicate(reduceLocal, 1.0f, curUbFactorA); + PipeBarrier(); + AscendC::Sqrt(yLocal, yLocal, curUbFactorA); + PipeBarrier(); + AscendC::Div(yLocal, reduceLocal, yLocal, curUbFactorA); + + outQueueY.EnQue(yLocal); + inQueueX.FreeTensor(xLocal); +} + +template +__aicore__ inline void HcPreInvRmsFullLoadLargeD::ComputeB32(uint64_t curUbFactorA) +{ + LocalTensor xLocal = inQueueX.DeQue(); + LocalTensor yLocal = outQueueY.AllocTensor(); + + LocalTensor reduceLocal = reduceBuf.Get(); + PipeBarrier(); + AscendC::Mul(xLocal, xLocal, xLocal, R); + + for (int idx = 0; idx < curUbFactorA; idx++) { + ReduceSumHalfInterval(yLocal[idx], xLocal[idx * rAlignB32], R); + } + + float meanCof = 1.0f / R; + PipeBarrier(); + AscendC::Muls(yLocal, yLocal, meanCof, curUbFactorA); + PipeBarrier(); + AscendC::Adds(yLocal, yLocal, epsilon, curUbFactorA); + PipeBarrier(); + AscendC::Duplicate(reduceLocal, 1.0f, curUbFactorA); + PipeBarrier(); + AscendC::Sqrt(yLocal, yLocal, curUbFactorA); + PipeBarrier(); + AscendC::Div(yLocal, reduceLocal, yLocal, curUbFactorA); + + outQueueY.EnQue(yLocal); + inQueueX.FreeTensor(xLocal); +} + +template +__aicore__ inline void HcPreInvRmsFullLoadLargeD::CopyOut(uint64_t idx, uint64_t curUbFactorA) +{ + LocalTensor yLocal = outQueueY.DeQue(); + AscendC::DataCopyExtParams copyParams{1, static_cast(curUbFactorA * sizeof(float)), 0, 0, 0}; + DataCopyPad(yGm[idx * ubFactorA], yLocal, copyParams); + outQueueY.FreeTensor(yLocal); +} + +} // namespace HcPreInvRmsLargeD +#endif // ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_LARGE_D_H_ \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load_regbase.h b/csrc/ascend/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load_regbase.h new file mode 100644 index 000000000..cf6796115 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load_regbase.h @@ -0,0 +1,303 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_inv_rms.h + * \brief inv rms file + */ +#ifndef ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_REGBASE_H_ +#define ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_REGBASE_H_ +#include "kernel_operator.h" + +namespace HcPreInvRmsRegbase { +using namespace AscendC; +constexpr int32_t BUFFER_NUM = 2; +constexpr int32_t FLOAT_BTYPE_SIZE = 4; +constexpr uint32_t VF_LEN_B32 = 64; +constexpr uint32_t UB_BLOCK_SIZE = 32; +constexpr uint32_t FOLD_FOUR = 4; + +constexpr AscendC::MicroAPI::CastTrait castTraitB162B32Even = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::UNKNOWN, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN, +}; + +template +__aicore__ inline void LoadInputData(AscendC::MicroAPI::RegTensor& dst, __local_mem__ T* src, AscendC::MicroAPI::MaskReg pregLoop, uint32_t srcOffset) +{ + if constexpr (IsSameType::value) { + DataCopy(dst, src + srcOffset); + } else if constexpr (IsSameType::value || IsSameType::value) { + AscendC::MicroAPI::RegTensor tmp; + DataCopy(tmp, src + srcOffset); + Cast(dst, tmp, pregLoop); + } +} + +template +class HcPreInvRmsFullLoadRegbase { +public: + __aicore__ inline HcPreInvRmsFullLoadRegbase() {}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe); + __aicore__ inline void Process(); + __aicore__ inline void CopyIn(uint64_t idx, uint64_t curUbFactorA); + __aicore__ inline void Compute(uint64_t idx, uint64_t curUbFactorA); + __aicore__ inline void ComputeFullLoadVF(LocalTensor& yLocal, LocalTensor& xLocal, uint32_t rAlign, uint32_t rNum, uint64_t curUbFactorA); + __aicore__ inline void ComputeFullLoadVfPerf(LocalTensor& yLocal, LocalTensor& xLocal, uint32_t rAlign, uint32_t rNum, uint64_t curUbFactorA); + __aicore__ inline void CopyOut(uint64_t idx, uint64_t curUbFactorA); + +private: + TPipe* pipe_; + + TQue inQueueX; + TQue outQueueY; + + GlobalTensor xGm; + GlobalTensor yGm; + + int64_t A; // 输入数据 A 轴大小 + int64_t R; // 输入数据 R 轴大小 + int64_t blockNumA; // 使用的核数 + int64_t blockFactorA; // 每个核处理的A个数 + int64_t blockTailFactorA; // 尾核处理的A个数 + int64_t ubFactorA; // 每次ub循环处理的A个数 + int32_t blockIdx_; + float epsilon; // 算子参数 + uint32_t curBlockFactorA; // 当前核处理的A个数 + uint32_t rAlign; +}; + +template +__aicore__ inline void HcPreInvRmsFullLoadRegbase::Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe) +{ + A = tiling->A; + R = tiling->R; + blockNumA = tiling->blockNumA; + blockFactorA = tiling->blockFactorA; + blockTailFactorA = tiling->blockTailFactorA; + ubFactorA = tiling->ubFactorA; + epsilon = tiling->epsilon; + + rAlign = ((R * sizeof(T) + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / sizeof(T)); + + pipe_ = pipe; + + blockIdx_ = GetBlockIdx(); + + if (blockIdx_ < blockNumA - 1) { + this->curBlockFactorA = this->blockFactorA; + } else if (blockIdx_ == blockNumA - 1) { + this->curBlockFactorA = this->blockTailFactorA; + } else { + return; + } + xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * blockFactorA * R, curBlockFactorA * R); + yGm.SetGlobalBuffer((__gm__ float*)y + blockIdx_ * blockFactorA, curBlockFactorA); + // pipe alloc memory to queue, the unit is Bytes + pipe_->InitBuffer(inQueueX, BUFFER_NUM, ubFactorA * rAlign * sizeof(T)); + pipe_->InitBuffer(outQueueY, BUFFER_NUM, ubFactorA * FLOAT_BTYPE_SIZE); +} + +template +__aicore__ inline void HcPreInvRmsFullLoadRegbase::Process() +{ + if (blockIdx_ >= blockNumA) { + return; + } + uint64_t aUbLoopCount = (curBlockFactorA + ubFactorA - 1) / ubFactorA; // Ub循环次数 + uint64_t tailUbFactorA = curBlockFactorA - (aUbLoopCount - 1) * ubFactorA; // 最后一次Ub循环的A轴大小 + uint64_t curUbFactorA = ubFactorA; + for (uint64_t idx = 0; idx < aUbLoopCount; idx++) { + if (idx == aUbLoopCount - 1) { + curUbFactorA = tailUbFactorA; + } + CopyIn(idx, curUbFactorA); + Compute(idx, curUbFactorA); + CopyOut(idx, curUbFactorA); + } +} + +template +__aicore__ inline void HcPreInvRmsFullLoadRegbase::CopyIn(uint64_t idx, uint64_t curUbFactorA) +{ + LocalTensor xLocal = inQueueX.AllocTensor(); + + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + int64_t xGmStartAddr = idx * R * ubFactorA; + DataCopyExtParams dataCopyParams{ + static_cast(curUbFactorA), static_cast(R * sizeof(T)), 0, 0, 0}; + DataCopyPad(xLocal, xGm[xGmStartAddr], dataCopyParams, dataCopyPadParams); + + inQueueX.EnQue(xLocal); +} + +template +__aicore__ inline void HcPreInvRmsFullLoadRegbase::Compute(uint64_t idx, uint64_t curUbFactorA) +{ + LocalTensor xLocal = inQueueX.DeQue(); + LocalTensor yLocal = outQueueY.AllocTensor(); + + if (R % 256 == 0) { + ComputeFullLoadVfPerf(yLocal, xLocal, rAlign, R, curUbFactorA); + } else { + ComputeFullLoadVF(yLocal, xLocal, rAlign, R, curUbFactorA); + } + + outQueueY.EnQue(yLocal); + inQueueX.FreeTensor(xLocal); +} + +template +__aicore__ inline void HcPreInvRmsFullLoadRegbase::ComputeFullLoadVF(LocalTensor& yLocal, LocalTensor& xLocal, uint32_t rAlign, uint32_t rNum, uint64_t curUbFactorA) +{ + __ubuf__ T* xAddr = (__ubuf__ T*)xLocal.GetPhyAddr(); + __ubuf__ float* yAddr = (__ubuf__ float*)yLocal.GetPhyAddr(); + + uint32_t vfLen = VF_LEN_B32; + uint16_t iLoopNum = curUbFactorA; + uint16_t needLoopNum = (rAlign + vfLen - 1) / vfLen; // 需要VF循环次数 + uint16_t fourLoopNum = (needLoopNum + FOLD_FOUR - 1) / FOLD_FOUR; // 需要四循环次数 + uint16_t formerFourLoopNum = fourLoopNum - 1; + uint16_t tailFourLoop = needLoopNum - formerFourLoopNum * FOLD_FOUR; // 最后一次四循环需处理VF数 + uint16_t formerFourLoopElems = formerFourLoopNum * FOLD_FOUR * vfLen; // 除尾块外，处理元素个数 + uint32_t tailFourLoopElems = rNum - formerFourLoopElems; // 最后一次四循环需处理实际元素个数 + + __VEC_SCOPE__ + { + AscendC::MicroAPI::RegTensor vregX1; + AscendC::MicroAPI::RegTensor vregX2; + AscendC::MicroAPI::RegTensor vregX3; + AscendC::MicroAPI::RegTensor vregX4; + AscendC::MicroAPI::RegTensor vregX; + AscendC::MicroAPI::RegTensor vregSum; + AscendC::MicroAPI::RegTensor vregR; + AscendC::MicroAPI::RegTensor vregOne; + AscendC::MicroAPI::MaskReg preg; + AscendC::MicroAPI::MaskReg pregAll = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg pregOne = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::Duplicate(vregOne, 1.0f); + + for (uint16_t i = 0; i < iLoopNum; i++) { + AscendC::MicroAPI::Duplicate(vregSum, 0.0f); // 用于累加的vreg + + for (uint16_t j = 0; j < formerFourLoopNum; j++) { + uint32_t srcOffset1 = i * rAlign + 4 * j * vfLen; + LoadInputData(vregX1, xAddr, pregAll, srcOffset1); + AscendC::MicroAPI::Mul(vregX1, vregX1, vregX1, pregAll); + AscendC::MicroAPI::Add(vregSum, vregSum, vregX1, pregAll); + + uint32_t srcOffset2 = i * rAlign + (4 * j + 1) * vfLen; + LoadInputData(vregX2, xAddr, pregAll, srcOffset2); + AscendC::MicroAPI::Mul(vregX2, vregX2, vregX2, pregAll); + AscendC::MicroAPI::Add(vregSum, vregSum, vregX2, pregAll); + + uint32_t srcOffset3 = i * rAlign + (4 * j + 2) * vfLen; + LoadInputData(vregX3, xAddr, pregAll, srcOffset3); + AscendC::MicroAPI::Mul(vregX3, vregX3, vregX3, pregAll); + AscendC::MicroAPI::Add(vregSum, vregSum, vregX3, pregAll); + + uint32_t srcOffset4 = i * rAlign + (4 * j + 3) * vfLen; + LoadInputData(vregX4, xAddr, pregAll, srcOffset4); + AscendC::MicroAPI::Mul(vregX4, vregX4, vregX4, pregAll); + AscendC::MicroAPI::Add(vregSum, vregSum, vregX4, pregAll); + } + + tailFourLoopElems = rNum - formerFourLoopElems; + for (uint16_t j = 0; j < tailFourLoop; j++) { + preg = AscendC::MicroAPI::UpdateMask(tailFourLoopElems); + uint32_t srcOffset = i * rAlign + formerFourLoopElems + j * vfLen; + LoadInputData(vregX, xAddr, preg, srcOffset); + AscendC::MicroAPI::Mul(vregX, vregX, vregX, preg); + AscendC::MicroAPI::Add(vregSum, vregSum, vregX, pregAll); + } + + Reduce(vregSum, vregSum, pregAll); + + AscendC::MicroAPI::Duplicate(vregR, (float)rNum); + AscendC::MicroAPI::Div(vregSum, vregSum, vregR, pregOne); + AscendC::MicroAPI::Adds(vregSum, vregSum, epsilon, pregOne); + AscendC::MicroAPI::Sqrt(vregSum, vregSum, pregOne); + AscendC::MicroAPI::Div(vregSum, vregOne, vregSum, pregOne); + AscendC::MicroAPI::DataCopy(yAddr + i, vregSum, pregOne); + } + } +} + +template +__aicore__ inline void HcPreInvRmsFullLoadRegbase::ComputeFullLoadVfPerf(LocalTensor& yLocal, LocalTensor& xLocal, uint32_t rAlign, uint32_t rNum, uint64_t curUbFactorA) +{ + __ubuf__ T* xAddr = (__ubuf__ T*)xLocal.GetPhyAddr(); + __ubuf__ float* yAddr = (__ubuf__ float*)yLocal.GetPhyAddr(); + + uint32_t vfLen = VF_LEN_B32; + uint16_t iLoopNum = curUbFactorA; + uint16_t jLoopNum = (rAlign / vfLen) / FOLD_FOUR; + + __VEC_SCOPE__ + { + AscendC::MicroAPI::RegTensor vregX1; + AscendC::MicroAPI::RegTensor vregX2; + AscendC::MicroAPI::RegTensor vregX3; + AscendC::MicroAPI::RegTensor vregX4; + AscendC::MicroAPI::RegTensor vregSum; + AscendC::MicroAPI::RegTensor vregR; + AscendC::MicroAPI::RegTensor vregOne; + AscendC::MicroAPI::MaskReg pregAll = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg pregOne = AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::Duplicate(vregOne, 1.0f); + + for (uint16_t i = 0; i < iLoopNum; i++) { + AscendC::MicroAPI::Duplicate(vregSum, 0.0f); // 用于累加的vreg + for (uint16_t j = 0; j < jLoopNum; j++) { + uint32_t srcOffset1 = i * rAlign + 4 * j * vfLen; + LoadInputData(vregX1, xAddr, pregAll, srcOffset1); + AscendC::MicroAPI::Mul(vregX1, vregX1, vregX1, pregAll); + AscendC::MicroAPI::Add(vregSum, vregSum, vregX1, pregAll); + + uint32_t srcOffset2 = i * rAlign + (4 * j + 1) * vfLen; + LoadInputData(vregX2, xAddr, pregAll, srcOffset2); + AscendC::MicroAPI::Mul(vregX2, vregX2, vregX2, pregAll); + AscendC::MicroAPI::Add(vregSum, vregSum, vregX2, pregAll); + + uint32_t srcOffset3 = i * rAlign + (4 * j + 2) * vfLen; + LoadInputData(vregX3, xAddr, pregAll, srcOffset3); + AscendC::MicroAPI::Mul(vregX3, vregX3, vregX3, pregAll); + AscendC::MicroAPI::Add(vregSum, vregSum, vregX3, pregAll); + + uint32_t srcOffset4 = i * rAlign + (4 * j + 3) * vfLen; + LoadInputData(vregX4, xAddr, pregAll, srcOffset4); + AscendC::MicroAPI::Mul(vregX4, vregX4, vregX4, pregAll); + AscendC::MicroAPI::Add(vregSum, vregSum, vregX4, pregAll); + } + + Reduce(vregSum, vregSum, pregAll); + AscendC::MicroAPI::Duplicate(vregR, (float)rNum); + AscendC::MicroAPI::Div(vregSum, vregSum, vregR, pregOne); + AscendC::MicroAPI::Adds(vregSum, vregSum, epsilon, pregOne); + AscendC::MicroAPI::Sqrt(vregSum, vregSum, pregOne); + AscendC::MicroAPI::Div(vregSum, vregOne, vregSum, pregOne); + AscendC::MicroAPI::DataCopy(yAddr + i, vregSum, pregOne); + } + } +} + +template +__aicore__ inline void HcPreInvRmsFullLoadRegbase::CopyOut(uint64_t idx, uint64_t curUbFactorA) +{ + LocalTensor yLocal = outQueueY.DeQue(); + AscendC::DataCopyExtParams copyParams{1, static_cast(curUbFactorA * sizeof(float)), 0, 0, 0}; + DataCopyPad(yGm[idx * ubFactorA], yLocal, copyParams); + outQueueY.FreeTensor(yLocal); +} + +} // namespace HcPreInvRmsRegbase +#endif // ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_REGBASE_H_ \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_sinkhorn/CMakeLists.txt b/csrc/ascend/moe/hc_pre_sinkhorn/CMakeLists.txt new file mode 100644 index 000000000..86b308249 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_sinkhorn/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_sinkhorn/op_host/CMakeLists.txt b/csrc/ascend/moe/hc_pre_sinkhorn/op_host/CMakeLists.txt new file mode 100644 index 000000000..54d974764 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_sinkhorn/op_host/CMakeLists.txt @@ -0,0 +1,63 @@ +# This program is free software, you can redistribute it and/or modify it. +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This file is a part of the CANN Open Software. +# Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ====================================================================================================================== + +# add_ops_compile_options( +# OP_NAME HcPreSinkhorn +# OPTIONS --cce-auto-sync=off +# -Wno-deprecated-declarations +# -Werror +# -mllvm -cce-aicore-hoist-movemask=false +# --op_relocatable_kernel_binary=true +# ) + +# set(hc_pre_sinkhorn_depends transformer/attention/hc_pre_sinkhorn PARENT_SCOPE) + +# target_sources(op_host_aclnn PRIVATE +# op_host/hc_pre_sinkhorn_def.cpp +# ) + +# target_sources(optiling PRIVATE +# op_host/hc_pre_sinkhorn_tiling.cpp +# ) + +# if (NOT BUILD_OPEN_PROJECT) +# target_sources(opmaster_ct PRIVATE +# op_host/hc_pre_sinkhorn_tiling.cpp +# ) +# endif () + +# target_include_directories(optiling PRIVATE +# ${CMAKE_CURRENT_SOURCE_DIR}/op_host +# ) + +# target_sources(opsproto PRIVATE +# op_host/hc_pre_sinkhorn_proto.cpp +# ) + +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + hc_pre_sinkhorn_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME HcPreSinkhorn + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true +) + + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE hc_pre_sinkhorn ACLNNTYPE aclnn) +endif() \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_sinkhorn/op_host/hc_pre_sinkhorn_def.cpp b/csrc/ascend/moe/hc_pre_sinkhorn/op_host/hc_pre_sinkhorn_def.cpp new file mode 100644 index 000000000..1ffb7570b --- /dev/null +++ b/csrc/ascend/moe/hc_pre_sinkhorn/op_host/hc_pre_sinkhorn_def.cpp @@ -0,0 +1,81 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_sinkhorn_def.cpp + * \brief + */ + +#include +#include "register/op_def_registry.h" + +namespace ops { +class HcPreSinkhorn : public OpDef { +public: + explicit HcPreSinkhorn(const char* name) : OpDef(name) + { + this->Input("mixes") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("rsqrt") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("hc_scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("hc_base") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("post") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("comb_frag") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + + this->Attr("hc_mult").AttrType(OPTIONAL).Int(4); + this->Attr("hc_sinkhorn_iters").AttrType(OPTIONAL).Int(20); + this->Attr("hc_eps").AttrType(OPTIONAL).Float(1e-6f); // default value + + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + + OpAICoreConfig regbaseCfg; + regbaseCfg.DynamicCompileStaticFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .ExtendCfgInfo("opFile.value", "hc_pre_sinkhorn"); + } +}; + +OP_ADD(HcPreSinkhorn); +} // namespace ops diff --git a/csrc/ascend/moe/hc_pre_sinkhorn/op_host/hc_pre_sinkhorn_proto.cpp b/csrc/ascend/moe/hc_pre_sinkhorn/op_host/hc_pre_sinkhorn_proto.cpp new file mode 100644 index 000000000..06e4dfd9c --- /dev/null +++ b/csrc/ascend/moe/hc_pre_sinkhorn/op_host/hc_pre_sinkhorn_proto.cpp @@ -0,0 +1,91 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_sinkhorn_proto.cpp + * \brief + */ +#include +#include + +#include "error/ops_error.h" + +using namespace ge; +namespace ops { +graphStatus InferShape4HcPreSinkhorn(gert::InferShapeContext* context) +{ + OPS_LOG_D(context->GetNodeName(), "Begin to do InferShape4HcPreSinkhorn."); + + const gert::Shape* mixShape = context->GetInputShape(0); + OPS_LOG_E_IF_NULL(context, mixShape, ge::GRAPH_FAILED); + auto mixDimNum = mixShape->GetDimNum(); + const gert::Shape* xShape = context->GetInputShape(4); + OPS_LOG_E_IF_NULL(context, xShape, ge::GRAPH_FAILED); + + gert::Shape* yShape = context->GetOutputShape(0); + OPS_LOG_E_IF_NULL(context, yShape, ge::GRAPH_FAILED); + gert::Shape* postShape = context->GetOutputShape(1); + OPS_LOG_E_IF_NULL(context, postShape, ge::GRAPH_FAILED); + gert::Shape* combFragShape = context->GetOutputShape(2); + OPS_LOG_E_IF_NULL(context, combFragShape, ge::GRAPH_FAILED); + + auto attrs = context->GetAttrs(); + auto *hcMult = attrs->GetAttrPointer(0); + + yShape->SetDimNum(mixDimNum); + postShape->SetDimNum(mixDimNum); + combFragShape->SetDimNum(mixDimNum + 1); + if (mixDimNum == 2) { + yShape->SetDim(0, mixShape->GetDim(0)); + yShape->SetDim(1, xShape->GetDim(2)); + + postShape->SetDim(0, mixShape->GetDim(0)); + postShape->SetDim(1, *hcMult); + + combFragShape->SetDim(0, mixShape->GetDim(0)); + combFragShape->SetDim(1, *hcMult); + combFragShape->SetDim(2, *hcMult); + + } else { + yShape->SetDim(0, mixShape->GetDim(0)); + yShape->SetDim(1, mixShape->GetDim(1)); + yShape->SetDim(2, xShape->GetDim(3)); + + postShape->SetDim(0, mixShape->GetDim(0)); + postShape->SetDim(1, mixShape->GetDim(1)); + postShape->SetDim(2, *hcMult); + + combFragShape->SetDim(0, mixShape->GetDim(0)); + combFragShape->SetDim(1, mixShape->GetDim(1)); + combFragShape->SetDim(2, *hcMult); + combFragShape->SetDim(3, *hcMult); + } + + OPS_LOG_D(context->GetNodeName(), "End to do InferShape4HcPreSinkhorn"); + return ge::GRAPH_SUCCESS; +} + +graphStatus InferDtype4HcPreSinkhorn(gert::InferDataTypeContext* context) +{ + OPS_LOG_D(context->GetNodeName(), "InferDtype4HcPreSinkhorn enter"); + + const auto xDataType = context->GetInputDataType(4); + context->SetOutputDataType(0, xDataType); + context->SetOutputDataType(1, DT_FLOAT); + context->SetOutputDataType(2, DT_FLOAT); + OPS_LOG_D(context->GetNodeName(), "InferDtype4HcPreSinkhorn end"); + + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(HcPreSinkhorn) + .InferShape(InferShape4HcPreSinkhorn) + .InferDataType(InferDtype4HcPreSinkhorn); +} // namespace ops diff --git a/csrc/ascend/moe/hc_pre_sinkhorn/op_host/hc_pre_sinkhorn_tiling.cpp b/csrc/ascend/moe/hc_pre_sinkhorn/op_host/hc_pre_sinkhorn_tiling.cpp new file mode 100644 index 000000000..dc8c39a1d --- /dev/null +++ b/csrc/ascend/moe/hc_pre_sinkhorn/op_host/hc_pre_sinkhorn_tiling.cpp @@ -0,0 +1,402 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_sinkhorn_tiling.cpp + * \brief + */ + +#include +#include "hc_pre_sinkhorn_tiling.h" + +using namespace ge; +namespace optiling { +namespace { +constexpr uint64_t WORKSPACE_SIZE = 32; +int64_t CeilDiv(int64_t x, int64_t y) +{ + if (y != 0) { + return (x + y - 1) / y; + } + return x; +} +int64_t DownAlign(int64_t x, int64_t y) { + if (y == 0) { + return x; + } + return (x / y) * y; +} +int64_t RoundUp(int64_t x, int64_t y) { + return CeilDiv(x, y) * y; +} + +constexpr int64_t BLOCK_SIZE = 32; +constexpr int64_t REPEAT_SIZE = 256; +constexpr int64_t DOUBLE_BUFFER = 2; +} + +ge::graphStatus HcPreSinkhornTiling::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + if (platformInfo == nullptr) { + auto compileInfoPtr = context_->GetCompileInfo(); + OPS_ERR_IF(compileInfoPtr == nullptr, OPS_LOG_E(context_, "compile info is null"), + return ge::GRAPH_FAILED); + coreNum_ = compileInfoPtr->coreNum; + ubSize_ = compileInfoPtr->ubSize; + } else { + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + coreNum_ = ascendcPlatform.GetCoreNumAiv(); + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + ubSize_ = ubSizePlatForm; + socVersion_ = ascendcPlatform.GetSocVersion(); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreSinkhornTiling::GetAttr() +{ + auto* attrs = context_->GetAttrs(); + OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED); + + auto hcMultAttr = attrs->GetAttrPointer(0); + hcMult_ = hcMultAttr == nullptr ? 4 : *hcMultAttr; + + auto iterTimesAttr = attrs->GetAttrPointer(1); + iterTimes_ = iterTimesAttr == nullptr ? 20 : *iterTimesAttr; + + auto epsAttr = attrs->GetAttrPointer(2); + eps_ = epsAttr == nullptr ? 1e-5 : *epsAttr; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreSinkhornTiling::GetShapeAttrsInfoInner() +{ + // (b, s, hc_mix) or (bs, hc_mix) + auto shapeMixes = context_->GetInputShape(0); + OPS_LOG_E_IF_NULL(context_, shapeMixes, return ge::GRAPH_FAILED); + size_t mixerDimNum = shapeMixes->GetStorageShape().GetDimNum(); + if (mixerDimNum == 2) { + bs_ = shapeMixes->GetStorageShape().GetDim(0); + hcMix_ = shapeMixes->GetStorageShape().GetDim(1); + } else if (mixerDimNum == 3) { + int64_t b = shapeMixes->GetStorageShape().GetDim(0); + int64_t s = shapeMixes->GetStorageShape().GetDim(1); + bs_ = b * s; + hcMix_ = shapeMixes->GetStorageShape().GetDim(2); + } + + auto shapeHcScale = context_->GetInputShape(2); + int64_t scaleFirstDim = shapeHcScale->GetStorageShape().GetDim(0); + OPS_ERR_IF(scaleFirstDim != 3, + OPS_LOG_E(context_->GetNodeName(), + "hc_scale size should be equal with 3, but is %ld", scaleFirstDim), + return ge::GRAPH_FAILED); + + auto shapeHcBase = context_->GetInputShape(3); + int64_t baseFirstDim = shapeHcBase->GetStorageShape().GetDim(0); + OPS_ERR_IF(baseFirstDim != hcMix_, + OPS_LOG_E(context_->GetNodeName(), + "hc_base size should be equal with mixhc, but is %ld", baseFirstDim), + return ge::GRAPH_FAILED); + + auto shapeX = context_->GetInputShape(4); + d_ = (mixerDimNum == 2 ? shapeX->GetStorageShape().GetDim(2) : shapeX->GetStorageShape().GetDim(3)); + + OPS_ERR_IF(GetAttr() != ge::GRAPH_SUCCESS, + OPS_LOG_E(context_->GetNodeName(), "get attr failed."), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + + +ge::graphStatus HcPreSinkhornTiling::CalcRegbaseOpTiling() +{ + rowOfFormerBlock_ = CeilDiv(bs_, static_cast(coreNum_)); + usedCoreNums_ = std::min(CeilDiv(bs_, rowOfFormerBlock_), static_cast(coreNum_)); + rowOfTailBlock_ = bs_ - (usedCoreNums_ - 1) * rowOfFormerBlock_; + + int64_t minRowPerCore = 1; + int64_t rowOnceLoop = std::min(rowOfFormerBlock_, minRowPerCore); + + hcMultAlign_ = RoundUp(hcMult_, BLOCK_SIZE / sizeof(float)); + int64_t mix0Size = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t mix1Size = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t mix2Size = rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t rsqrtSize = RoundUp(rowOnceLoop, BLOCK_SIZE / sizeof(float)) * sizeof(float) * DOUBLE_BUFFER; + int64_t xSize = rowOnceLoop * hcMult_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + int64_t ySize = rowOnceLoop * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; + int64_t postSize = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t combFragSize = rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t base0Size = hcMultAlign_ * sizeof(float); + int64_t base1Size = hcMultAlign_ * sizeof(float); + int64_t base2Size = hcMult_ * hcMultAlign_ * sizeof(float); + + int64_t totalSize = mix0Size + mix1Size + mix2Size + rsqrtSize + xSize + ySize + postSize + combFragSize + + base0Size + base1Size + base2Size; + rowFactor_ = rowOnceLoop; + if (totalSize <= ubSize_) { + // row和d均可以在ub内全载 + dLoop_ = 1; + dFactor_ = d_; + tailDFactor_ = dFactor_; + } else { + int64_t usedUbSize = mix0Size + mix1Size + mix2Size + rsqrtSize + postSize + combFragSize + + base0Size + base1Size + base2Size; + int64_t ubRemain = ubSize_ - usedUbSize; + dFactor_ = d_; + int64_t base = 2; + while (1) { + dFactor_ = CeilDiv(d_, base); + xSize = rowOnceLoop * hcMult_ * RoundUp(dFactor_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + ySize = rowOnceLoop * RoundUp(dFactor_, 16) * 2 * DOUBLE_BUFFER; + int64_t targetSize = xSize + ySize; + if (targetSize <= ubRemain) { + break; + } + base++; + } + if (dFactor_ > 32) { + dFactor_ = DownAlign(dFactor_, 32); + } + dLoop_ = CeilDiv(d_, dFactor_); + tailDFactor_ = d_ % dFactor_ == 0 ? dFactor_ : d_ % dFactor_; + } + + // d全载,尝试搬入更多的bs + if (dFactor_ == d_) { + while (rowFactor_ <= rowOfFormerBlock_) { + mix0Size = rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + mix1Size = rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + mix2Size = rowFactor_ * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + rsqrtSize = RoundUp(rowFactor_, BLOCK_SIZE / sizeof(float)) * sizeof(float) * DOUBLE_BUFFER; + xSize = rowFactor_ * hcMult_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + ySize = rowFactor_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; + postSize = rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + combFragSize = rowFactor_ * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + totalSize = mix0Size + mix1Size + mix2Size + rsqrtSize + xSize + ySize + postSize + combFragSize + + base0Size + base1Size + base2Size; + if (totalSize > ubSize_) { + rowFactor_ = rowFactor_ - 1; + break; + } + rowFactor_ = rowFactor_ + 1; + } + rowFactor_ = rowFactor_ > rowOfFormerBlock_ ? rowFactor_ - 1 : rowFactor_; + } + rowLoopOfFormerBlock_ = CeilDiv(rowOfFormerBlock_, rowFactor_); + rowLoopOfTailBlock_ = CeilDiv(rowOfTailBlock_, rowFactor_); + tailRowFactorOfFormerBlock_ = rowOfFormerBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfFormerBlock_ % rowFactor_; + tailRowFactorOfTailBlock_ = rowOfTailBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfTailBlock_ % rowFactor_; + + tilingData_.set_bs(bs_); + tilingData_.set_hcMix(hcMix_); + tilingData_.set_hcMult(hcMult_); + tilingData_.set_d(d_); + tilingData_.set_hcMultAlign(hcMultAlign_); + tilingData_.set_rowOfFormerBlock(rowOfFormerBlock_); + tilingData_.set_rowOfTailBlock(rowOfTailBlock_); + tilingData_.set_rowLoopOfFormerBlock(rowLoopOfFormerBlock_); + tilingData_.set_rowLoopOfTailBlock(rowLoopOfTailBlock_); + tilingData_.set_rowFactor(rowFactor_); + tilingData_.set_tailRowFactorOfFormerBlock(tailRowFactorOfFormerBlock_); + tilingData_.set_tailRowFactorOfTailBlock(tailRowFactorOfTailBlock_); + tilingData_.set_dLoop(dLoop_); + tilingData_.set_dFactor(dFactor_); + tilingData_.set_tailDFactor(tailDFactor_); + tilingData_.set_iterTimes(iterTimes_); + tilingData_.set_eps(eps_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreSinkhornTiling::CalcMembaseOpTiling() +{ + rowOfFormerBlock_ = CeilDiv(bs_, static_cast(coreNum_)); + usedCoreNums_ = std::min(CeilDiv(bs_, rowOfFormerBlock_), static_cast(coreNum_)); + rowOfTailBlock_ = bs_ - (usedCoreNums_ - 1) * rowOfFormerBlock_; + + int64_t minRowPerCore = 1; + int64_t rowOnceLoop = std::min(rowOfFormerBlock_, minRowPerCore); + + hcMultAlign_ = RoundUp(hcMult_, BLOCK_SIZE / sizeof(float)); + int64_t mix0Size = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t mix1Size = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t mix2Size = rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t rsqrtSize = RoundUp(rowOnceLoop, BLOCK_SIZE / sizeof(float)) * sizeof(float) * DOUBLE_BUFFER; + int64_t xSize = rowOnceLoop * hcMult_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + int64_t ySize = rowOnceLoop * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; + int64_t postSize = rowOnceLoop * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t combFragSize = rowOnceLoop * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + int64_t base0Size = hcMultAlign_ * sizeof(float); + int64_t base1Size = hcMultAlign_ * sizeof(float); + int64_t base2Size = hcMult_ * hcMultAlign_ * sizeof(float); + int64_t xCastSize = rowOnceLoop * hcMult_ * RoundUp(d_, 8) * sizeof(float); + int64_t yCastSize = rowOnceLoop * RoundUp(d_, 8) * sizeof(float); + int64_t rowBrcb0Size = RoundUp(rowOnceLoop, 8) * BLOCK_SIZE; + int64_t hcBrcb1Size = RoundUp(rowOnceLoop * hcMultAlign_, 8) * BLOCK_SIZE; + int64_t reduceBufSize = rowOnceLoop * hcMultAlign_ * sizeof(float); + + int64_t totalSize = mix0Size + mix1Size + mix2Size + rsqrtSize + xSize + ySize + postSize + combFragSize + + base0Size + base1Size + base2Size + xCastSize + yCastSize + rowBrcb0Size + hcBrcb1Size + reduceBufSize; + rowFactor_ = rowOnceLoop; + if (totalSize <= ubSize_) { + // row和d均可以在ub内全载 + dLoop_ = 1; + dFactor_ = d_; + tailDFactor_ = dFactor_; + } else { + int64_t usedUbSize = mix0Size + mix1Size + mix2Size + rsqrtSize + postSize + combFragSize + + base0Size + base1Size + base2Size + rowBrcb0Size + hcBrcb1Size + reduceBufSize; + int64_t ubRemain = ubSize_ - usedUbSize; + dFactor_ = d_; + int64_t base = 2; + while (1) { + dFactor_ = CeilDiv(d_, base); + xSize = rowOnceLoop * hcMult_ * RoundUp(dFactor_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + ySize = rowOnceLoop * RoundUp(dFactor_, 16) * 2 * DOUBLE_BUFFER; + xCastSize = rowOnceLoop * hcMult_ * RoundUp(dFactor_, 8) * sizeof(float); + yCastSize = rowOnceLoop * RoundUp(dFactor_, 8) * sizeof(float); + int64_t targetSize = xSize + ySize + xCastSize + yCastSize; + if (targetSize <= ubRemain) { + break; + } + base++; + } + if (dFactor_ > 32) { + dFactor_ = DownAlign(dFactor_, 32); + } + dLoop_ = CeilDiv(d_, dFactor_); + tailDFactor_ = d_ % dFactor_ == 0 ? dFactor_ : d_ % dFactor_; + } + + // d全载,尝试搬入更多的bs + if (dFactor_ == d_) { + while (rowFactor_ <= rowOfFormerBlock_) { + mix0Size = rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + mix1Size = rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + mix2Size = rowFactor_ * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + rsqrtSize = RoundUp(rowFactor_, BLOCK_SIZE / sizeof(float)) * sizeof(float) * DOUBLE_BUFFER; + xSize = rowFactor_ * hcMult_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; // x是bfloat16_t 类型 + ySize = rowFactor_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER; + postSize = rowFactor_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + combFragSize = rowFactor_ * hcMult_ * hcMultAlign_ * sizeof(float) * DOUBLE_BUFFER; + xCastSize = rowFactor_ * hcMult_ * RoundUp(d_, 8) * sizeof(float); + yCastSize = rowFactor_ * RoundUp(d_, 8) * sizeof(float); + rowBrcb0Size = RoundUp(rowFactor_, 8) * BLOCK_SIZE; + hcBrcb1Size = RoundUp(rowFactor_ * hcMultAlign_, 8) * BLOCK_SIZE; + reduceBufSize = rowFactor_ * hcMultAlign_ * sizeof(float); + totalSize = mix0Size + mix1Size + mix2Size + rsqrtSize + xSize + ySize + postSize + combFragSize + + base0Size + base1Size + base2Size + xCastSize + yCastSize + rowBrcb0Size + hcBrcb1Size + reduceBufSize;; + if (totalSize > ubSize_) { + rowFactor_ = rowFactor_ - 1; + break; + } + rowFactor_ = rowFactor_ + 1; + } + rowFactor_ = rowFactor_ > rowOfFormerBlock_ ? rowFactor_ - 1 : rowFactor_; + } + rowLoopOfFormerBlock_ = CeilDiv(rowOfFormerBlock_, rowFactor_); + rowLoopOfTailBlock_ = CeilDiv(rowOfTailBlock_, rowFactor_); + tailRowFactorOfFormerBlock_ = rowOfFormerBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfFormerBlock_ % rowFactor_; + tailRowFactorOfTailBlock_ = rowOfTailBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfTailBlock_ % rowFactor_; + + tilingData_.set_bs(bs_); + tilingData_.set_hcMix(hcMix_); + tilingData_.set_hcMult(hcMult_); + tilingData_.set_d(d_); + tilingData_.set_hcMultAlign(hcMultAlign_); + tilingData_.set_rowOfFormerBlock(rowOfFormerBlock_); + tilingData_.set_rowOfTailBlock(rowOfTailBlock_); + tilingData_.set_rowLoopOfFormerBlock(rowLoopOfFormerBlock_); + tilingData_.set_rowLoopOfTailBlock(rowLoopOfTailBlock_); + tilingData_.set_rowFactor(rowFactor_); + tilingData_.set_tailRowFactorOfFormerBlock(tailRowFactorOfFormerBlock_); + tilingData_.set_tailRowFactorOfTailBlock(tailRowFactorOfTailBlock_); + tilingData_.set_dLoop(dLoop_); + tilingData_.set_dFactor(dFactor_); + tilingData_.set_tailDFactor(tailDFactor_); + tilingData_.set_iterTimes(iterTimes_); + tilingData_.set_eps(eps_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreSinkhornTiling::CalcOpTiling() { + if (socVersion_ == platform_ascendc::SocVersion::ASCEND950) { + return CalcRegbaseOpTiling(); + } + return CalcMembaseOpTiling(); +} + + +ge::graphStatus HcPreSinkhornTiling::DoOpTiling() +{ + if (GetPlatformInfo() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + if (GetShapeAttrsInfoInner() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + if (CalcOpTiling() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + if (GetWorkspaceSize() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + if (PostTiling() == ge::GRAPH_FAILED) { + return ge::GRAPH_FAILED; + } + + context_->SetTilingKey(0); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreSinkhornTiling::GetWorkspaceSize() +{ + workspaceSize_ = WORKSPACE_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus HcPreSinkhornTiling::PostTiling() +{ + context_->SetTilingKey(0); + context_->SetBlockDim(usedCoreNums_); + size_t* workspaces = context_->GetWorkspaceSizes(1); + workspaces[0] = workspaceSize_; + tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingPrepareForHcPreSinkhorn(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingForHcPreSinkhorn(gert::TilingContext *context) +{ + OPS_ERR_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("HcPreSinkhorn", "Tiling context is null"), + return ge::GRAPH_FAILED); + HcPreSinkhornTiling hcPreTiling(context); + return hcPreTiling.DoOpTiling(); +} + +IMPL_OP_OPTILING(HcPreSinkhorn) + .Tiling(TilingForHcPreSinkhorn) + .TilingParse(TilingPrepareForHcPreSinkhorn); + +} // namespace optiling diff --git a/csrc/ascend/moe/hc_pre_sinkhorn/op_host/hc_pre_sinkhorn_tiling.h b/csrc/ascend/moe/hc_pre_sinkhorn/op_host/hc_pre_sinkhorn_tiling.h new file mode 100644 index 000000000..7e3fcbdab --- /dev/null +++ b/csrc/ascend/moe/hc_pre_sinkhorn/op_host/hc_pre_sinkhorn_tiling.h @@ -0,0 +1,118 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_sinkhorn_tiling.h + * \brief + */ + +#ifndef HC_PRE_SINKHORN_TILING_H +#define HC_PRE_SINKHORN_TILING_H + + +#include +#include +#include "register/op_impl_registry.h" +#include "platform/platform_infos_def.h" +#include "exe_graph/runtime/tiling_context.h" +#include "tiling/platform/platform_ascendc.h" +#include "register/op_def_registry.h" +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "error/ops_error.h" +#include "platform/platform_info.h" + +namespace optiling { +// ----------公共定义---------- +struct TilingRequiredParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::StorageShape *shape; +}; + +struct TilingOptionalParaInfo { + const gert::CompileTimeTensorDesc *desc; + const gert::Tensor *tensor; +}; +// ----------算子TilingData定义---------- +BEGIN_TILING_DATA_DEF(HcPreSinkhornTilingData) +TILING_DATA_FIELD_DEF(int64_t, bs); +TILING_DATA_FIELD_DEF(int64_t, hcMix); +TILING_DATA_FIELD_DEF(int64_t, hcMult); +TILING_DATA_FIELD_DEF(int64_t, d); +TILING_DATA_FIELD_DEF(int64_t, hcMultAlign); +TILING_DATA_FIELD_DEF(int64_t, rowOfFormerBlock); +TILING_DATA_FIELD_DEF(int64_t, rowOfTailBlock); +TILING_DATA_FIELD_DEF(int64_t, rowLoopOfFormerBlock); +TILING_DATA_FIELD_DEF(int64_t, rowLoopOfTailBlock); +TILING_DATA_FIELD_DEF(int64_t, rowFactor); +TILING_DATA_FIELD_DEF(int64_t, tailRowFactorOfFormerBlock); +TILING_DATA_FIELD_DEF(int64_t, tailRowFactorOfTailBlock); +TILING_DATA_FIELD_DEF(int64_t, dLoop); +TILING_DATA_FIELD_DEF(int64_t, dFactor); +TILING_DATA_FIELD_DEF(int64_t, tailDFactor); +TILING_DATA_FIELD_DEF(int64_t, iterTimes); +TILING_DATA_FIELD_DEF(float, eps); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(HcPreSinkhorn, HcPreSinkhornTilingData) + +// ----------算子CompileInfo定义---------- +struct HcPreSinkhornCompileInfo { + uint64_t coreNum = 0; + uint64_t ubSize = 0; +}; + +// ----------算子Tiling入参信息解析及check类---------- +class HcPreSinkhornTiling { +public: + explicit HcPreSinkhornTiling(gert::TilingContext* tilingContext) : context_(tilingContext) + { + } + ~HcPreSinkhornTiling() = default; + + ge::graphStatus GetPlatformInfo(); + ge::graphStatus DoOpTiling(); + ge::graphStatus GetWorkspaceSize(); + ge::graphStatus PostTiling(); + ge::graphStatus GetAttr(); + ge::graphStatus GetShapeAttrsInfoInner(); + ge::graphStatus CalcOpTiling(); + ge::graphStatus CalcMembaseOpTiling(); + ge::graphStatus CalcRegbaseOpTiling(); +private: + gert::TilingContext *context_ = nullptr; + uint64_t tilingKey_ = 0; + HcPreSinkhornTilingData tilingData_; + uint64_t coreNum_ = 0; + uint64_t workspaceSize_ = 0; + uint64_t usedCoreNums_ = 0; + uint64_t ubSize_ = 0; + int64_t bs_ = 0; + int64_t hcMix_ = 0; + int64_t hcMult_ = 0; + int64_t d_ = 0; + int64_t hcMultAlign_ = 0; + int64_t rowOfFormerBlock_ = 0; + int64_t rowOfTailBlock_ = 0; + int64_t rowLoopOfFormerBlock_ = 0; + int64_t rowLoopOfTailBlock_ = 0; + int64_t rowFactor_ = 0; + int64_t tailRowFactorOfFormerBlock_ = 0; + int64_t tailRowFactorOfTailBlock_= 0; + int64_t dLoop_ = 0; + int64_t dFactor_ = 0; + int64_t tailDFactor_ = 0; + int64_t iterTimes_ = 0; + double eps_ = 0.0; + platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B; +}; + +} // namespace optiling +#endif // HC_PRE_SINKHORN_TILING_H diff --git a/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn.cpp b/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn.cpp new file mode 100644 index 000000000..b4660807e --- /dev/null +++ b/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn.cpp @@ -0,0 +1,45 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_sinkhorn.cpp + * \brief + */ + +#if defined(__DAV_C310__) + #include "hc_pre_sinkhorn_regbase_perf.h" + #include "hc_pre_sinkhorn_regbase_base.h" +#else + #include "hc_pre_sinkhorn_perf.h" + #include "hc_pre_sinkhorn_base.h" +#endif + +using namespace HcPreSinkhorn; + +extern "C" __global__ __aicore__ void hc_pre_sinkhorn(GM_ADDR mixes, GM_ADDR rsqrt, GM_ADDR hcScale, GM_ADDR hcBase, + GM_ADDR x, GM_ADDR y, GM_ADDR post, GM_ADDR combFrag, GM_ADDR workspace, + GM_ADDR tiling) +{ + if (workspace == nullptr) { + return; + } + + GM_ADDR userWs = GetUserWorkspace(workspace); + if (userWs == nullptr) { + return; + } + GET_TILING_DATA(tilingData, tiling); + TPipe pipe; + if (TILING_KEY_IS(0)) { + HcPreSinkhorn::HcPreSinkhornPerf op; + op.Init(mixes, rsqrt, hcScale, hcBase, x, y, post, combFrag, userWs, &tilingData, &pipe); + op.Process(); + } +} \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_base.h b/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_base.h new file mode 100644 index 000000000..c04ff9700 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_base.h @@ -0,0 +1,623 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_perf.h + * \brief + */ + +#ifndef HC_PRE_SINKHORN_BASE_H +#define HC_PRE_SINKHORN_BASE_H + +#include "kernel_operator.h" + +namespace HcPreSinkhorn { +using namespace AscendC; +constexpr int32_t BLOCK_SIZE = 32; +constexpr int32_t DEFAULT_BLOCK_STRIDE = 1; +constexpr int32_t DEFAULT_REPEAT_STRIDE = 8; +constexpr int32_t ONE_REPEAT_BLOCK_NUMS = 8; +constexpr int32_t REPEAT_SIZE = 256; +constexpr int32_t MAX_REPEAT_STRIDE = 255; + +__aicore__ inline int32_t CeilDiv(int32_t a, int32_t b) +{ + if (b == 0) { + return a; + } + return (a + b - 1) / b; +} + +__aicore__ inline int32_t CeilAlign(int32_t a, int32_t b) +{ + return CeilDiv(a, b) * b; +} + +template +__aicore__ inline int32_t RoundUp(int32_t num) +{ + int32_t elemNum = BLOCK_SIZE / sizeof(T); + return CeilAlign(num, elemNum); +} + +template +__aicore__ inline void MulABLastDimBrcInline(const LocalTensor &output, const LocalTensor &input0, + const LocalTensor &input1, const LocalTensor &tmpBuffer, + const int32_t curRowNum, const int32_t curColNum) +{ + if constexpr (needBrc) { + uint32_t repeatTimes = CeilDiv(curRowNum, ONE_REPEAT_BLOCK_NUMS); + Brcb(tmpBuffer, input1, repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE}); + } + PipeBarrier(); + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + uint32_t curColNumAlign = RoundUp(curColNum); + if (curColNum <= elemInOneBlock) { + Mul(output, input0, tmpBuffer, curRowNum * curColNumAlign); + } else { + int32_t numRepeatPerLine = curColNum / elemInOneRepeat; + int32_t numRemainPerLine = curColNum % elemInOneRepeat; + int32_t dstRepStridePerLine = CeilDiv(curColNum, elemInOneBlock); + BinaryRepeatParams instrParams; + + if (numRepeatPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE || curRowNum < numRepeatPerLine) { + // 在Col方向开Repeat, 并且Repeat小于255 + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Mul(output[i * curColNumAlign], input0[i * curColNumAlign], tmpBuffer[i * elemInOneBlock], + elemInOneRepeat, numRepeatPerLine, instrParams); + } + } else { + // 在Row方向开Repeat + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 1; + for (uint32_t i = 0; i < numRepeatPerLine; i++) { + Mul(output[i * elemInOneRepeat], input0[i * elemInOneRepeat], tmpBuffer, elemInOneRepeat, curRowNum, + instrParams); + } + } + } + + if (numRemainPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE) { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = 0; + instrParams.src0RepStride = 0; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Mul(output[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], + input0[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], tmpBuffer[i * elemInOneBlock], + numRemainPerLine, 1, instrParams); + } + } else { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 0; + Mul(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], tmpBuffer, + numRemainPerLine, curRowNum, instrParams); + } + } + } + PipeBarrier(); +} + +template +__aicore__ inline void SubABLastDimBrcInline(const LocalTensor &output, const LocalTensor &input0, + const LocalTensor &input1, const LocalTensor &tmpBuffer, + const int32_t curRowNum, const int32_t curColNum) +{ + if constexpr (needBrc) { + uint32_t repeatTimes = CeilDiv(curRowNum, ONE_REPEAT_BLOCK_NUMS); + Brcb(tmpBuffer, input1, repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE}); + } + PipeBarrier(); + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + uint32_t curColNumAlign = RoundUp(curColNum); + if (curColNum <= elemInOneBlock) { + Sub(output, input0, tmpBuffer, curRowNum * curColNumAlign); + } else { + int32_t numRepeatPerLine = curColNum / elemInOneRepeat; + int32_t numRemainPerLine = curColNum % elemInOneRepeat; + int32_t dstRepStridePerLine = CeilDiv(curColNum, elemInOneBlock); + BinaryRepeatParams instrParams; + + if (numRepeatPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE || curRowNum < numRepeatPerLine) { + // 在Col方向开Repeat, 并且Repeat小于255 + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Sub(output[i * curColNumAlign], input0[i * curColNumAlign], tmpBuffer[i * elemInOneBlock], + elemInOneRepeat, numRepeatPerLine, instrParams); + } + } else { + // 在Row方向开Repeat + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 1; + for (uint32_t i = 0; i < numRepeatPerLine; i++) { + Sub(output[i * elemInOneRepeat], input0[i * elemInOneRepeat], tmpBuffer, elemInOneRepeat, curRowNum, + instrParams); + } + } + } + + if (numRemainPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE) { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = 0; + instrParams.src0RepStride = 0; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Sub(output[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], + input0[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], tmpBuffer[i * elemInOneBlock], + numRemainPerLine, 1, instrParams); + } + } else { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 0; + Sub(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], tmpBuffer, + numRemainPerLine, curRowNum, instrParams); + } + } + } + PipeBarrier(); +} + + +template +__aicore__ inline void DivABLastDimBrcInline(const LocalTensor &output, const LocalTensor &input0, + const LocalTensor &input1, const LocalTensor &tmpBuffer, + const int32_t curRowNum, const int32_t curColNum) +{ + if constexpr (needBrc) { + uint32_t repeatTimes = CeilDiv(curRowNum, ONE_REPEAT_BLOCK_NUMS); + Brcb(tmpBuffer, input1, repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE}); + } + PipeBarrier(); + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + uint32_t curColNumAlign = RoundUp(curColNum); + if (curColNum <= elemInOneBlock) { + Div(output, input0, tmpBuffer, curRowNum * curColNumAlign); + } else { + int32_t numRepeatPerLine = curColNum / elemInOneRepeat; + int32_t numRemainPerLine = curColNum % elemInOneRepeat; + int32_t dstRepStridePerLine = CeilDiv(curColNum, elemInOneBlock); + BinaryRepeatParams instrParams; + + if (numRepeatPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE || curRowNum < numRepeatPerLine) { + // 在Col方向开Repeat, 并且Repeat小于255 + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Div(output[i * curColNumAlign], input0[i * curColNumAlign], tmpBuffer[i * elemInOneBlock], + elemInOneRepeat, numRepeatPerLine, instrParams); + } + } else { + // 在Row方向开Repeat + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 1; + for (uint32_t i = 0; i < numRepeatPerLine; i++) { + Div(output[i * elemInOneRepeat], input0[i * elemInOneRepeat], tmpBuffer, elemInOneRepeat, curRowNum, + instrParams); + } + } + } + + if (numRemainPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE) { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = 0; + instrParams.src0RepStride = 0; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Div(output[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], + input0[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], tmpBuffer[i * elemInOneBlock], + numRemainPerLine, 1, instrParams); + } + } else { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 0; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 0; + Div(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], tmpBuffer, + numRemainPerLine, curRowNum, instrParams); + } + } + } + PipeBarrier(); +} + + +template +__aicore__ inline void AddBAFirstDimBrcInline(const LocalTensor &output, const LocalTensor &input0, + const LocalTensor &input1, const int32_t curRowNum, + const int32_t curColNum) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + uint32_t curColNumAlign = RoundUp(curColNum); + int32_t numRepeatPerLine = curColNum / elemInOneRepeat; + int32_t numRemainPerLine = curColNum % elemInOneRepeat; + int32_t dstRepStridePerLine = CeilDiv(curColNum, elemInOneBlock); + BinaryRepeatParams instrParams; + if (numRepeatPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE || curRowNum < numRepeatPerLine) { + // 在Col方向开Repeat, 并且Repeat小于255 + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = DEFAULT_REPEAT_STRIDE; + for (uint32_t i = 0; i < curRowNum; i++) { + Add(output[i * curColNumAlign], input0[i * curColNumAlign], input1, elemInOneRepeat, numRepeatPerLine, + instrParams); + } + } else { + // 在Row方向开Repeat + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < numRepeatPerLine; i++) { + Add(output[i * elemInOneRepeat], input0[i * elemInOneRepeat], input1[i * elemInOneRepeat], + elemInOneRepeat, curRowNum, instrParams); + } + } + } + + if (numRemainPerLine > 0) { + if (dstRepStridePerLine > MAX_REPEAT_STRIDE) { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = 0; + instrParams.src0RepStride = 0; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < curRowNum; i++) { + Add(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], input1, + numRemainPerLine, 1, instrParams); + } + } else { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = dstRepStridePerLine; + instrParams.src0RepStride = dstRepStridePerLine; + instrParams.src1RepStride = 0; + Add(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], input1, + numRemainPerLine, curRowNum, instrParams); + } + } + PipeBarrier(); +} + +template +__aicore__ inline void CalcDenominator(const LocalTensor &output, const LocalTensor &input, + const uint32_t calCount) +{ + Muls(output, input, static_cast(-1.0), calCount); + PipeBarrier(); + Exp(output, output, calCount); + PipeBarrier(); + Adds(output, output, static_cast(1.0), calCount); + PipeBarrier(); +} + +// 暂时不处理repeat超限场景 +template +__aicore__ inline void SigmoidPerf(const LocalTensor &output, const LocalTensor &input, + const LocalTensor &tmpBuffer, const int64_t calCount) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + Duplicate(tmpBuffer, static_cast(1.0), elemInOneBlock); + CalcDenominator(output, input, calCount); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + int32_t numRepeatPerLine = calCount / elemInOneRepeat; + int32_t numRemainPerLine = calCount % elemInOneRepeat; + BinaryRepeatParams instrParams; + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 0; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = 0; + instrParams.src1RepStride = DEFAULT_REPEAT_STRIDE; + Div(output, tmpBuffer, output, elemInOneRepeat, numRepeatPerLine, instrParams); + if (numRemainPerLine != 0) { + Div(output[numRepeatPerLine * elemInOneRepeat], tmpBuffer, output[numRepeatPerLine * elemInOneRepeat], + numRemainPerLine, 1, instrParams); + } + PipeBarrier(); +} + +__aicore__ inline void ProcessPre(const LocalTensor &preLocal, const LocalTensor &mixLocal, + const LocalTensor &hcBaseLocal, const LocalTensor &rsqrtLocal, + const LocalTensor &tmpBuffer0, const LocalTensor &tmpBuffer1, + float scale, float eps, const int32_t curRowNum, const int32_t curColNum) +{ + int32_t curColNumAlign = RoundUp(curColNum); + MulABLastDimBrcInline(mixLocal, mixLocal, rsqrtLocal, tmpBuffer0, curRowNum, curColNum); + Muls(mixLocal, mixLocal, scale, curRowNum * curColNumAlign); + PipeBarrier(); + AddBAFirstDimBrcInline(mixLocal, mixLocal, hcBaseLocal, curRowNum, curColNum); + SigmoidPerf(preLocal, mixLocal, tmpBuffer1, curRowNum * curColNumAlign); + Adds(preLocal, preLocal, eps, curRowNum * curColNumAlign); + PipeBarrier(); +} + +__aicore__ inline void ReduceSumARAPerf(const LocalTensor &output, const LocalTensor &input, + const uint32_t dim0, const uint32_t dim1, const uint32_t dim2) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(float); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(float); + uint32_t dim2Align = RoundUp(dim2); + + // 拷贝第一个R到output上 + DataCopyParams copyParams; + copyParams.blockCount = dim0; + copyParams.blockLen = dim2Align / elemInOneBlock; + copyParams.srcStride = (dim1 - 1) * (dim2Align / elemInOneBlock); + copyParams.dstStride = 0; + DataCopy(output, input, copyParams); + PipeBarrier(); + uint32_t dim2RepeatTimes = dim2 / elemInOneRepeat; + uint32_t dim2Reminder = dim2 % elemInOneRepeat; + // 沿着dim2方向开repeat + BinaryRepeatParams instrParams; + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = DEFAULT_REPEAT_STRIDE; + for (uint32_t i = 0; i < dim0; i++) { + for (uint32_t j = 1; j < dim1; j++) { + Add(output[i * dim2Align], output[i * dim2Align], input[i * dim1 * dim2Align + j * dim2Align], + elemInOneRepeat, dim2RepeatTimes, instrParams); + if (dim2Reminder != 0) { + Add(output[i * dim2Align + dim2RepeatTimes * elemInOneRepeat], + output[i * dim2Align + dim2RepeatTimes * elemInOneRepeat], + input[i * dim1 * dim2Align + j * dim2Align + +dim2RepeatTimes * elemInOneRepeat], dim2Reminder, 1, + instrParams); + } + PipeBarrier(); + } + } + PipeBarrier(); +} + +template +__aicore__ inline void CastTwoDim(const LocalTensor &output, const LocalTensor &input, const uint32_t dim0, + const uint32_t dim1) +{ + uint32_t dim1AlignT0 = RoundUp(dim1); + uint32_t dim1AlignT1 = RoundUp(dim1); + if constexpr (IsSameType::value && IsSameType::value) { + for (uint32_t i = 0; i < dim0; i++) { + Cast(output[i * dim1AlignT0], input[i * dim1AlignT1], AscendC::RoundMode::CAST_NONE, dim1); + } + } else { + for (uint32_t i = 0; i < dim0; i++) { + Cast(output[i * dim1AlignT0], input[i * dim1AlignT1], AscendC::RoundMode::CAST_RINT, dim1); + } + } + PipeBarrier(); +} + + +template +__aicore__ void inline ProcessY(const LocalTensor &yLocal, const LocalTensor &xLocal, + const LocalTensor &mix01Local, const LocalTensor &hcBrcbLocal1, + const LocalTensor &xCastLocal, const LocalTensor &yCastLocal, + const uint32_t dim0, const uint32_t dim1, const uint32_t dim2) +{ + CastTwoDim(xCastLocal, xLocal, dim0 * dim1, dim2); + MulABLastDimBrcInline(xCastLocal, xCastLocal, mix01Local, hcBrcbLocal1, dim0 * dim1, dim2); + ReduceSumARAPerf(yCastLocal, xCastLocal, dim0, dim1, dim2); + CastTwoDim(yLocal, yCastLocal, dim0, dim2); +} + + +__aicore__ inline void ProcessPost(const LocalTensor &postLocal, const LocalTensor &mixLocal, + const LocalTensor &hcBaseLocal, const LocalTensor &rsqrtLocal, + const LocalTensor &tmpBuffer0, const LocalTensor &tmpBuffer1, + float scale, const int32_t curRowNum, const int32_t curColNum) +{ + int32_t curColNumAlign = RoundUp(curColNum); + MulABLastDimBrcInline(mixLocal, mixLocal, rsqrtLocal, tmpBuffer0, curRowNum, curColNum); + Muls(mixLocal, mixLocal, scale, curRowNum * curColNumAlign); + PipeBarrier(); + AddBAFirstDimBrcInline(mixLocal, mixLocal, hcBaseLocal, curRowNum, curColNum); + SigmoidPerf(postLocal, mixLocal, tmpBuffer1, curRowNum * curColNumAlign); + Muls(postLocal, postLocal, static_cast(2.0f), curRowNum * curColNumAlign); + PipeBarrier(); +} + +__aicore__ inline void LastDimReduceMaxPerf(const LocalTensor &output, const LocalTensor &input, + const uint32_t curRowNum, const uint32_t curColNum) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(float); + WholeReduceMax(output, input, curColNum, curRowNum, 1, 1, CeilDiv(curColNum, elemInOneBlock), + ReduceOrder::ORDER_ONLY_VALUE); + PipeBarrier(); +} + +__aicore__ inline void LastDimReduceSumPerf(const LocalTensor &output, const LocalTensor &input, + const uint32_t curRowNum, const uint32_t curColNum) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(float); + WholeReduceSum(output, input, curColNum, curRowNum, 1, 1, CeilDiv(curColNum, elemInOneBlock)); + PipeBarrier(); +} + + +// 暂时只支持R轴小于64,既curColNum不能超过64 +__aicore__ inline void SoftmaxFP32Perf(const LocalTensor &output, const LocalTensor &input, + const LocalTensor &tmpReduceBuffer, + const LocalTensor tmpBrcbBuffer, const int32_t curRowNum, + const int32_t curColNum, float eps) +{ + LastDimReduceMaxPerf(tmpReduceBuffer, input, curRowNum, curColNum); + SubABLastDimBrcInline(output, input, tmpReduceBuffer, tmpBrcbBuffer, curRowNum, curColNum); + uint32_t curColNumAlign = RoundUp(curColNum); + Exp(output, output, curRowNum * curColNumAlign); + PipeBarrier(); + LastDimReduceSumPerf(tmpReduceBuffer, output, curRowNum, curColNum); + DivABLastDimBrcInline(output, output, tmpReduceBuffer, tmpBrcbBuffer, curRowNum, curColNum); + Adds(output, output, eps, curRowNum * curColNumAlign); + PipeBarrier(); +} + +// (bs, hc_mult, hc_mult) = (bs, hc_mult, hc_mult) + (bs, 1, hc_mult) + +template +__aicore__ inline void DivABABrcInline(const LocalTensor &output, const LocalTensor &input0, + const LocalTensor &input1, const uint32_t dim0, const uint32_t dim1, + const uint32_t dim2) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T); + uint32_t dim2Align = RoundUp(dim2); + uint32_t dim2RepeatTimes = dim2 / elemInOneRepeat; + uint32_t dim2Reminder = dim2 % elemInOneRepeat; + uint32_t dim2RepeatStride = CeilDiv(dim2, elemInOneBlock); + // 在dim1方向开repeat + BinaryRepeatParams instrParams; + if (dim1 >= dim2RepeatTimes) { + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = dim2RepeatStride; + instrParams.src0RepStride = dim2RepeatStride; + instrParams.src1RepStride = 0; + for (uint32_t i = 0; i < dim0; i++) { + for (uint32_t j = 0; j < dim2RepeatTimes; j++) { + Div(output[i * dim1 * dim2Align + j * elemInOneRepeat], + input0[i * dim1 * dim2Align + j * elemInOneRepeat], input1[i * dim2Align + j * elemInOneRepeat], + elemInOneRepeat, dim1, instrParams); + } + if (dim2Reminder != 0) { + Div(output[i * dim1 * dim2Align + dim2RepeatTimes * elemInOneRepeat], + input0[i * dim1 * dim2Align + dim2RepeatTimes * elemInOneRepeat], + input1[i * dim2Align + dim2RepeatTimes * elemInOneRepeat], dim2Reminder, dim1, instrParams); + } + } + } else { + // 在dim2方向开repeat + instrParams.dstBlkStride = 1; + instrParams.src0BlkStride = 1; + instrParams.src1BlkStride = 1; + instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE; + instrParams.src1RepStride = DEFAULT_REPEAT_STRIDE; + for (uint32_t i = 0; i < dim0; i++) { + for (uint32_t j = 0; j < dim1; j++) { + Div(output[i * dim1 * dim2Align + j * dim2Align], input0[i * dim1 * dim2Align + j * dim2Align], + input1[i * dim2Align], dim2); + } + } + } + PipeBarrier(); +} + +template +__aicore__ inline void CopyIn(const GlobalTensor &inputGm, const LocalTensor &inputTensor, const uint16_t nBurst, + const uint32_t copyLen, uint32_t srcStride = 0) +{ + DataCopyPadExtParams dataCopyPadExtParams; + dataCopyPadExtParams.isPad = false; + dataCopyPadExtParams.leftPadding = 0; + dataCopyPadExtParams.rightPadding = 0; + dataCopyPadExtParams.paddingValue = 0; + + DataCopyExtParams dataCoptExtParams; + dataCoptExtParams.blockCount = nBurst; + dataCoptExtParams.blockLen = copyLen * sizeof(T); + dataCoptExtParams.srcStride = srcStride * sizeof(T); + dataCoptExtParams.dstStride = 0; + DataCopyPad(inputTensor, inputGm, dataCoptExtParams, dataCopyPadExtParams); +} + +// (bs, hc_mult, hc_mult) --> (bs, hc_mult, hc_mult_align) +template +__aicore__ inline void CopyInWithOuterFor(const GlobalTensor &inputGm, const LocalTensor &inputTensor, + const uint16_t outerLoop, const uint16_t nBurst, const uint32_t copyLen, + const uint32_t gmLastDim) +{ + uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T); + uint32_t ubLastDimAlign = RoundUp(copyLen); + + for (uint16_t i = 0; i < outerLoop; i++) { + CopyIn(inputGm[i * nBurst * gmLastDim], inputTensor[i * nBurst * ubLastDimAlign], nBurst, copyLen); + } +} + +template +__aicore__ inline void CopyOut(const LocalTensor &outputTensor, const GlobalTensor &outputGm, + const uint16_t nBurst, const uint32_t copyLen, uint32_t dstStride = 0) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = nBurst; + dataCopyParams.blockLen = copyLen * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = dstStride * sizeof(T); + DataCopyPad(outputGm, outputTensor, dataCopyParams); +} + +} // namespace HcPreSinkhorn + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_perf.h b/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_perf.h new file mode 100644 index 000000000..b8f8e59ce --- /dev/null +++ b/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_perf.h @@ -0,0 +1,256 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_sinkhorn_perf.h + * \brief + */ + +#ifndef HC_PRE_SINKHORN_PERF_H +#define HC_PRE_SINKHORN_PERF_H + +#include "kernel_operator.h" +#include "hc_pre_sinkhorn_base.h" + +namespace HcPreSinkhorn { +using namespace AscendC; +template +class HcPreSinkhornPerf { +public: + __aicore__ inline HcPreSinkhornPerf() + { + } + + __aicore__ inline void Init(GM_ADDR mixes, GM_ADDR rsqrt, GM_ADDR hcScale, GM_ADDR hcBase, GM_ADDR x, GM_ADDR y, + GM_ADDR post, GM_ADDR combFrag, GM_ADDR workspace, + const HcPreSinkhornTilingData *tilingDataPtr, TPipe *pipePtr) + { + pipe = pipePtr; + tilingData = tilingDataPtr; + + mixesGm.SetGlobalBuffer((__gm__ float *)mixes); + rsqrtGm.SetGlobalBuffer((__gm__ float *)rsqrt); + hcScaleGm.SetGlobalBuffer((__gm__ float *)hcScale); + hcBaseGm.SetGlobalBuffer((__gm__ float *)hcBase); + xGm.SetGlobalBuffer((__gm__ T *)x); + yGm.SetGlobalBuffer((__gm__ T *)y); + postGm.SetGlobalBuffer((__gm__ float *)post); + combFragGm.SetGlobalBuffer((__gm__ float *)combFrag); + + // InQue + int64_t mixesQue01Size = tilingData->rowFactor * tilingData->hcMultAlign * 2 * sizeof(float); + pipe->InitBuffer(mixesQue01, 2, mixesQue01Size); + pipe->InitBuffer(mixesQue2, 2, + tilingData->rowFactor * tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(rsqrtQue, 2, RoundUp(tilingData->rowFactor) * sizeof(float)); + pipe->InitBuffer(xQue, 2, + tilingData->rowFactor * tilingData->hcMult * RoundUp(tilingData->dFactor) * sizeof(T)); + + // OutQue + pipe->InitBuffer(yQue, 2, + tilingData->rowFactor * RoundUp(tilingData->dFactor) * sizeof(T)); + pipe->InitBuffer(postQue, 2, tilingData->rowFactor * tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(combFragQue, 2, + tilingData->rowFactor * tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); + + // TBuf + pipe->InitBuffer(hcBaseBuf0, tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(hcBaseBuf1, tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(hcBaseBuf2, tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(rowBrcbBuf0, RoundUp(tilingData->rowFactor) * BLOCK_SIZE); + pipe->InitBuffer(hcBrcbBuf1, RoundUp(tilingData->rowFactor * tilingData->hcMultAlign) * BLOCK_SIZE); + pipe->InitBuffer(reduceBuf, tilingData->rowFactor * tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(xCastBuf, tilingData->rowFactor * tilingData->hcMult * RoundUp(tilingData->dFactor) * + sizeof(float)); + pipe->InitBuffer(yCastBuf, tilingData->rowFactor * RoundUp(tilingData->dFactor) * + sizeof(float)); + + hcBase0Local = hcBaseBuf0.Get(); + hcBase1Local = hcBaseBuf1.Get(); + hcBase2Local = hcBaseBuf2.Get(); + rowBrcbLocal0 = rowBrcbBuf0.Get(); + hcBrcbLocal1 = hcBrcbBuf1.Get(); + reduceLocal = reduceBuf.Get(); + xCastLocal = xCastBuf.Get(); + yCastLocal = yCastBuf.Get(); + } + + __aicore__ inline void Process() + { + int64_t curBlockIdx = GetBlockIdx(); + int64_t totalBlockNum = GetBlockNum(); + + int64_t rowOuterLoop = + (curBlockIdx == totalBlockNum - 1) ? tilingData->rowLoopOfTailBlock : tilingData->rowLoopOfFormerBlock; + int64_t tailRowFactor = (curBlockIdx == totalBlockNum - 1) ? tilingData->tailRowFactorOfTailBlock : + tilingData->tailRowFactorOfFormerBlock; + + CopyIn(hcBaseGm, hcBase0Local, 1, tilingData->hcMult); + CopyIn(hcBaseGm[tilingData->hcMult], hcBase1Local, 1, tilingData->hcMult); + CopyIn(hcBaseGm[tilingData->hcMult * 2], hcBase2Local, tilingData->hcMult, tilingData->hcMult); + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId); + WaitFlag(eventId); + int64_t mixGmBaseOffset = curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMix; + int64_t xGmBaseOffset = curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMult * tilingData->d; + for (int64_t rowOuterIdx = 0; rowOuterIdx < rowOuterLoop; rowOuterIdx++) { + int64_t curRowFactor = (rowOuterIdx == rowOuterLoop - 1) ? tailRowFactor : tilingData->rowFactor; + mixes01Local = mixesQue01.AllocTensor(); + CopyIn(mixesGm[mixGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMix], mixes01Local, + curRowFactor, tilingData->hcMult, tilingData->hcMix - tilingData->hcMult); + CopyIn( + mixesGm[mixGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMix + tilingData->hcMult], + mixes01Local[tilingData->rowFactor * tilingData->hcMultAlign], curRowFactor, tilingData->hcMult, + tilingData->hcMix - tilingData->hcMult); + mixesQue01.EnQue(mixes01Local); + rsqrtLocal = rsqrtQue.AllocTensor(); + CopyIn(rsqrtGm[curBlockIdx * tilingData->rowOfFormerBlock + rowOuterIdx * tilingData->rowFactor], + rsqrtLocal, 1, curRowFactor); + rsqrtQue.EnQue(rsqrtLocal); + + mixes01Local = mixesQue01.DeQue(); + rsqrtLocal = rsqrtQue.DeQue(); + ProcessPre(mixes01Local, mixes01Local, hcBase0Local, rsqrtLocal, rowBrcbLocal0, hcBrcbLocal1, + hcScaleGm.GetValue(0), tilingData->eps, curRowFactor, tilingData->hcMult); + for (int64_t dLoopIdx = 0; dLoopIdx < tilingData->dLoop; dLoopIdx++) { + int64_t curDFactor = + (dLoopIdx == tilingData->dLoop - 1) ? tilingData->tailDFactor : tilingData->dFactor; + xLocal = xQue.template AllocTensor(); + CopyIn(xGm[xGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMult * tilingData->d + + dLoopIdx * tilingData->dFactor], + xLocal, tilingData->rowFactor * tilingData->hcMult, curDFactor, tilingData->d - curDFactor); + xQue.template EnQue(xLocal); + xLocal = xQue.template DeQue(); + yLocal = yQue.template AllocTensor(); + ProcessY(yLocal, xLocal, mixes01Local, hcBrcbLocal1, xCastLocal, yCastLocal, curRowFactor, + tilingData->hcMult, curDFactor); + xQue.template FreeTensor(xLocal); + yQue.template EnQue(yLocal); + yLocal = yQue.template DeQue(); + CopyOut(yLocal, + yGm[curBlockIdx * tilingData->rowOfFormerBlock * tilingData->d + + rowOuterIdx * tilingData->rowFactor * tilingData->d + dLoopIdx * tilingData->dFactor], + curRowFactor, curDFactor, tilingData->d - curDFactor); + yQue.template FreeTensor(yLocal); + } + // post + postLocal = postQue.AllocTensor(); + ProcessPost(postLocal, mixes01Local[tilingData->rowFactor * tilingData->hcMultAlign], hcBase1Local, + rsqrtLocal, rowBrcbLocal0, hcBrcbLocal1, hcScaleGm.GetValue(1), curRowFactor, + tilingData->hcMult); + mixesQue01.template FreeTensor(mixes01Local); + postQue.EnQue(postLocal); + postLocal = postQue.DeQue(); + CopyOut(postLocal, + postGm[curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMult + + rowOuterIdx * tilingData->rowFactor * tilingData->hcMult], + curRowFactor, tilingData->hcMult); + postQue.FreeTensor(postLocal); + + // combFrag + mixes2Local = mixesQue2.AllocTensor(); + CopyInWithOuterFor(mixesGm[mixGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMix + + tilingData->hcMult * 2], + mixes2Local, curRowFactor, tilingData->hcMult, tilingData->hcMult, tilingData->hcMix); + mixesQue2.EnQue(mixes2Local); + mixes2Local = mixesQue2.DeQue(); + + combFragLocal = combFragQue.AllocTensor(); + + MulABLastDimBrcInline(mixes2Local, mixes2Local, rsqrtLocal, rowBrcbLocal0, curRowFactor, + tilingData->hcMult * tilingData->hcMultAlign); + Muls(mixes2Local, mixes2Local, hcScaleGm.GetValue(2), + curRowFactor * tilingData->hcMult * tilingData->hcMultAlign); + PipeBarrier(); + AddBAFirstDimBrcInline(mixes2Local, mixes2Local, hcBase2Local, curRowFactor, + tilingData->hcMult * tilingData->hcMultAlign); + SoftmaxFP32Perf(mixes2Local, mixes2Local, reduceLocal, hcBrcbLocal1, curRowFactor * tilingData->hcMult, + tilingData->hcMult, tilingData->eps); + ReduceSumARAPerf(reduceLocal, mixes2Local, curRowFactor, tilingData->hcMult, tilingData->hcMult); + Adds(reduceLocal, reduceLocal, tilingData->eps, curRowFactor * tilingData->hcMult); + PipeBarrier(); + DivABABrcInline(combFragLocal, mixes2Local, reduceLocal, curRowFactor, tilingData->hcMult, + tilingData->hcMult); + for (int64_t iter = 0; iter < tilingData->iterTimes - 1; iter++) { + LastDimReduceSumPerf(reduceLocal, combFragLocal, curRowFactor * tilingData->hcMult, tilingData->hcMult); + Adds(reduceLocal, reduceLocal, tilingData->eps, curRowFactor * tilingData->hcMult); + PipeBarrier(); + DivABLastDimBrcInline(combFragLocal, combFragLocal, reduceLocal, hcBrcbLocal1, + curRowFactor * tilingData->hcMult, tilingData->hcMult); + ReduceSumARAPerf(reduceLocal, combFragLocal, curRowFactor, tilingData->hcMult, tilingData->hcMult); + Adds(reduceLocal, reduceLocal, tilingData->eps, curRowFactor * tilingData->hcMult); + PipeBarrier(); + DivABABrcInline(combFragLocal, combFragLocal, reduceLocal, curRowFactor, tilingData->hcMult, + tilingData->hcMult); + } + mixesQue2.FreeTensor(mixes2Local); + rsqrtQue.FreeTensor(rsqrtLocal); + + combFragQue.EnQue(combFragLocal); + combFragLocal = combFragQue.DeQue(); + CopyOut(combFragLocal, + combFragGm[curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMult * tilingData->hcMult + + rowOuterIdx * tilingData->rowFactor * tilingData->hcMult * tilingData->hcMult], + curRowFactor * tilingData->hcMult, tilingData->hcMult); + combFragQue.FreeTensor(combFragLocal); + } + } + +private: + TPipe *pipe; + const HcPreSinkhornTilingData *tilingData; + GlobalTensor mixesGm; + GlobalTensor rsqrtGm; + GlobalTensor hcScaleGm; + GlobalTensor hcBaseGm; + GlobalTensor xGm; + GlobalTensor yGm; + GlobalTensor postGm; + GlobalTensor combFragGm; + + TQue mixesQue01; + TQue mixesQue2; + TQue rsqrtQue; + TQue xQue; + TQue yQue; + TQue postQue; + TQue combFragQue; + + TBuf hcBaseBuf0; + TBuf hcBaseBuf1; + TBuf hcBaseBuf2; + + TBuf rowBrcbBuf0; + TBuf hcBrcbBuf1; + TBuf reduceBuf; + + TBuf xCastBuf; + TBuf yCastBuf; + + LocalTensor mixes01Local; + LocalTensor mixes2Local; + LocalTensor rsqrtLocal; + LocalTensor xLocal; + LocalTensor yLocal; + LocalTensor postLocal; + LocalTensor combFragLocal; + LocalTensor hcBase0Local; + LocalTensor hcBase1Local; + LocalTensor hcBase2Local; + LocalTensor rowBrcbLocal0; + LocalTensor hcBrcbLocal1; + LocalTensor reduceLocal; + LocalTensor xCastLocal; + LocalTensor yCastLocal; +}; + +} // namespace HcPreSinkhorn + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_base.h b/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_base.h new file mode 100644 index 000000000..f84d3d4bf --- /dev/null +++ b/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_base.h @@ -0,0 +1,511 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_sinkhorn_regbase_base.h + * \brief + */ + +#ifndef HC_PRE_SINKHORN_RGEBASE_BASE_H +#define HC_PRE_SINKHORN_RGEBASE_BASE_H + +#include "kernel_operator.h" + +namespace HcPreSinkhorn { +using namespace AscendC; +using namespace AscendC::MicroAPI; +using AscendC::MicroAPI::MaskReg; +using AscendC::MicroAPI::RegTensor; +using AscendC::MicroAPI::UnalignReg; +constexpr int32_t BLOCK_SIZE = 32; +constexpr int32_t VL_FP32 = 64; + +__aicore__ inline int32_t CeilDiv(int32_t a, int b) +{ + if (b == 0) { + return a; + } + return (a + b - 1) / b; +} + +__aicore__ inline int32_t CeilAlign(int32_t a, int b) +{ + return CeilDiv(a, b) * b; +} + +template +__aicore__ inline int32_t RoundUp(int32_t num) +{ + int32_t elemNum = BLOCK_SIZE / sizeof(T); + return CeilAlign(num, elemNum); +} + +constexpr AscendC::MicroAPI::CastTrait castTraitB162B32Even = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::UNKNOWN, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB322B16Even = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_RINT, +}; + +template +__aicore__ inline void LoadInputData(RegTensor& dst, __local_mem__ T* src, MaskReg pregLoop, uint32_t srcOffset) +{ + if constexpr (IsSameType::value) { + DataCopy(dst, src + srcOffset); + } else if constexpr (IsSameType::value || IsSameType::value) { + RegTensor tmp; + DataCopy(tmp, src + srcOffset); + Cast(dst, tmp, pregLoop); + } +} + +template +__aicore__ inline void StoreOutputData( + __local_mem__ T* dst, RegTensor& src, MaskReg pregLoop, uint32_t dstOffset) +{ + if constexpr (IsSameType::value) { + DataCopy(dst + dstOffset, src, pregLoop); + } else if constexpr (IsSameType::value || IsSameType::value) { + RegTensor tmp; + Cast(tmp, src, pregLoop); + DataCopy(dst + dstOffset, tmp, pregLoop); + } +} + +template +__aicore__ inline void LoadInputDataWithBrc( + RegTensor& dst, __local_mem__ T* src, MaskReg pregLoop, uint32_t srcOffset) +{ + if constexpr (IsSameType::value) { + DataCopy(dst, src + srcOffset); + } else if constexpr (IsSameType::value || IsSameType::value) { + RegTensor tmp; + DataCopy(tmp, src + srcOffset); + Cast(dst, tmp, pregLoop); + } +} + +__aicore__ inline void VFSigmoid( + RegTensor& y, RegTensor& x, RegTensor& one, MaskReg pregLoop) +{ + Muls(x, x, static_cast(-1), pregLoop); + Exp(x, x, pregLoop); + Adds(x, x, static_cast(1), pregLoop); + Div(y, one, x, pregLoop); +} + +__aicore__ inline void VFProcessPre( + const LocalTensor& preLocal, const LocalTensor& mixLocal, const LocalTensor& hcBaseLocal, + const LocalTensor& rsqrtLocal, float scale, float eps, uint16_t curRowNum, uint16_t curColNum) +{ + __local_mem__ float* preLocalAddr = (__local_mem__ float*)preLocal.GetPhyAddr(); + __local_mem__ float* mixLocalAddr = (__local_mem__ float*)mixLocal.GetPhyAddr(); + __local_mem__ float* hcBaseLocalAddr = (__local_mem__ float*)hcBaseLocal.GetPhyAddr(); + __local_mem__ float* rsqrtLocalAddr = (__local_mem__ float*)rsqrtLocal.GetPhyAddr(); + uint16_t loopCount = CeilDiv(curColNum, VL_FP32); + uint32_t curColNumAlign = RoundUp(curColNum); + if (loopCount > 1) { + __VEC_SCOPE__ + { + RegTensor mix; + RegTensor base; + RegTensor rsqrt; + RegTensor one; + MaskReg pregLoop = CreateMask(); + uint32_t sreg = curColNum; + Duplicate(one, static_cast(1), pregLoop); + for (uint16_t i = 0; i < loopCount; i++) { + pregLoop = UpdateMask(sreg); + LoadInputData(base, hcBaseLocalAddr, pregLoop, i * VL_FP32); + for (uint16_t j = 0; j < curRowNum; j++) { + LoadInputDataWithBrc(rsqrt, rsqrtLocalAddr, pregLoop, j); + LoadInputData(mix, mixLocalAddr, pregLoop, i * VL_FP32 + j * curColNumAlign); + Mul(mix, mix, rsqrt, pregLoop); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + VFSigmoid(mix, mix, one, pregLoop); + Adds(mix, mix, eps, pregLoop); + StoreOutputData(preLocalAddr, mix, pregLoop, i * VL_FP32 + j * curColNumAlign); + } + } + } + } else { + __VEC_SCOPE__ + { + RegTensor mix; + RegTensor base; + RegTensor rsqrt; + RegTensor one; + uint32_t sreg = curColNum; + MaskReg pregLoop = UpdateMask(sreg); + Duplicate(one, static_cast(1), pregLoop); + LoadInputData(base, hcBaseLocalAddr, pregLoop, 0); + for (uint16_t i = 0; i < curRowNum; i++) { + LoadInputData(mix, mixLocalAddr, pregLoop, i * curColNumAlign); + LoadInputDataWithBrc(rsqrt, rsqrtLocalAddr, pregLoop, i); + Mul(mix, mix, rsqrt, pregLoop); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + VFSigmoid(mix, mix, one, pregLoop); + Adds(mix, mix, eps, pregLoop); + StoreOutputData(preLocalAddr, mix, pregLoop, i * curColNumAlign); + } + } + } +} + +__aicore__ inline void VFProcessPost( + const LocalTensor& postLocal, const LocalTensor& mixLocal, const LocalTensor& hcBaseLocal, + const LocalTensor& rsqrtLocal, float scale, float eps, uint16_t curRowNum, uint16_t curColNum) +{ + __local_mem__ float* postLocalAddr = (__local_mem__ float*)postLocal.GetPhyAddr(); + __local_mem__ float* mixLocalAddr = (__local_mem__ float*)mixLocal.GetPhyAddr(); + __local_mem__ float* hcBaseLocalAddr = (__local_mem__ float*)hcBaseLocal.GetPhyAddr(); + __local_mem__ float* rsqrtLocalAddr = (__local_mem__ float*)rsqrtLocal.GetPhyAddr(); + uint16_t loopCount = CeilDiv(curColNum, VL_FP32); + uint32_t curColNumAlign = RoundUp(curColNum); + if (loopCount > 1) { + __VEC_SCOPE__ + { + RegTensor mix; + RegTensor base; + RegTensor rsqrt; + RegTensor one; + MaskReg pregLoop = CreateMask(); + uint32_t sreg = curColNum; + Duplicate(one, static_cast(1), pregLoop); + for (uint16_t i = 0; i < loopCount; i++) { + pregLoop = UpdateMask(sreg); + LoadInputData(base, hcBaseLocalAddr, pregLoop, i * VL_FP32); + for (uint16_t j = 0; j < curRowNum; j++) { + LoadInputData(mix, mixLocalAddr, pregLoop, i * VL_FP32 + j * curColNumAlign); + LoadInputDataWithBrc(rsqrt, rsqrtLocalAddr, pregLoop, i); + Mul(mix, mix, rsqrt, pregLoop); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + VFSigmoid(mix, mix, one, pregLoop); + Muls(mix, mix, static_cast(2.0), pregLoop); + StoreOutputData(postLocalAddr, mix, pregLoop, i * VL_FP32 + j * curColNumAlign); + } + } + } + } else { + __VEC_SCOPE__ + { + RegTensor mix; + RegTensor base; + RegTensor rsqrt; + RegTensor one; + uint32_t sreg = curColNum; + MaskReg pregLoop = UpdateMask(sreg); + Duplicate(one, static_cast(1), pregLoop); + LoadInputData(base, hcBaseLocalAddr, pregLoop, 0); + for (uint16_t i = 0; i < curRowNum; i++) { + LoadInputData(mix, mixLocalAddr, pregLoop, i * curColNumAlign); + LoadInputDataWithBrc(rsqrt, rsqrtLocalAddr, pregLoop, i); + Mul(mix, mix, rsqrt, pregLoop); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + VFSigmoid(mix, mix, one, pregLoop); + Muls(mix, mix, static_cast(2.0), pregLoop); + StoreOutputData(postLocalAddr, mix, pregLoop, i * curColNumAlign); + } + } + } +} + +// dim2是R轴，R轴小于64, 不需要回写UB +__aicore__ inline void VFProcessCombFragRLessVL( + const LocalTensor& combFragLocal, const LocalTensor& mixLocal, const LocalTensor& hcBaseLocal, + const LocalTensor& rsqrtLocal, float scale, float eps, uint16_t iters, uint16_t dim0, uint16_t dim1, + uint16_t dim2) +{ + __local_mem__ float* combFragLocalAddr = (__local_mem__ float*)combFragLocal.GetPhyAddr(); + __local_mem__ float* mixLocalAddr = (__local_mem__ float*)mixLocal.GetPhyAddr(); + __local_mem__ float* hcBaseLocalAddr = (__local_mem__ float*)hcBaseLocal.GetPhyAddr(); + __local_mem__ float* rsqrtLocalAddr = (__local_mem__ float*)rsqrtLocal.GetPhyAddr(); + uint32_t dim2Align = RoundUp(dim2); + __VEC_SCOPE__ + { + RegTensor base; + RegTensor mix; + RegTensor rsqrt; + RegTensor max; + RegTensor sum; + RegTensor sum1; + uint32_t sreg = dim2; + MaskReg pregLoop = UpdateMask(sreg); + for (uint16_t i = 0; i < dim0; i++) { + Duplicate(sum1, static_cast(0), pregLoop); + LoadInputDataWithBrc(rsqrt, rsqrtLocalAddr, pregLoop, i); + for (uint16_t j = 0; j < dim1; j++) { + LoadInputData(base, hcBaseLocalAddr, pregLoop, j * dim2Align); + LoadInputData(mix, mixLocalAddr, pregLoop, i * dim1 * dim2Align + j * dim2Align); + Mul(mix, mix, rsqrt, pregLoop); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + ReduceMax(max, mix, pregLoop); + Duplicate(max, max, pregLoop); + Sub(mix, mix, max, pregLoop); + Exp(mix, mix, pregLoop); + ReduceSum(sum, mix, pregLoop); + Duplicate(sum, sum, pregLoop); + Div(mix, mix, sum, pregLoop); + Adds(mix, mix, eps, pregLoop); + Add(sum1, sum1, mix, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, i * dim1 * dim2Align + j * dim2Align); + } + LocalMemBar(); + Adds(sum1, sum1, eps, pregLoop); + for (uint16_t j = 0; j < dim1; j++) { + LoadInputData(mix, combFragLocalAddr, pregLoop, i * dim1 * dim2Align + j * dim2Align); + Div(mix, mix, sum1, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, i * dim1 * dim2Align + j * dim2Align); + } + } + for (uint16_t i = 0; i < iters; i++) { + LocalMemBar(); + for (uint16_t j = 0; j < dim0; j++) { + Duplicate(sum1, static_cast(0), pregLoop); + for (uint16_t k = 0; k < dim1; k++) { + LoadInputData(mix, combFragLocalAddr, pregLoop, j * dim1 * dim2Align + k * dim2Align); + ReduceSum(sum, mix, pregLoop); + Duplicate(sum, sum, pregLoop); + Adds(sum, sum, eps, pregLoop); + Div(mix, mix, sum, pregLoop); + Add(sum1, sum1, mix, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, j * dim1 * dim2Align + k * dim2Align); + } + LocalMemBar(); + Adds(sum1, sum1, eps, pregLoop); + for (uint16_t k = 0; k < dim1; k++) { + LoadInputData(mix, combFragLocalAddr, pregLoop, j * dim1 * dim2Align + k * dim2Align); + Div(mix, mix, sum1, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, j * dim1 * dim2Align + k * dim2Align); + } + } + } + } +} + +__aicore__ inline void VFProcessIteration(RegTensor& sum0, RegTensor& sum1, RegTensor& mix, float eps, MaskReg pregLoop) +{ + ReduceSum(sum1, mix, pregLoop); + Duplicate(sum1, sum1, pregLoop); + Adds(sum1, sum1, eps, pregLoop); + Div(mix, mix, sum1, pregLoop); + Add(sum0, sum0, mix, pregLoop); +} + +__aicore__ inline void VFProcessCombFragRLessVLUseFourUnfold( + const LocalTensor& combFragLocal, const LocalTensor& mixLocal, const LocalTensor& hcBaseLocal, + const LocalTensor& rsqrtLocal, float scale, float eps, uint16_t iters, uint16_t dim0, uint16_t dim1, + uint16_t dim2) +{ + __local_mem__ float* combFragLocalAddr = (__local_mem__ float*)combFragLocal.GetPhyAddr(); + __local_mem__ float* mixLocalAddr = (__local_mem__ float*)mixLocal.GetPhyAddr(); + __local_mem__ float* hcBaseLocalAddr = (__local_mem__ float*)hcBaseLocal.GetPhyAddr(); + __local_mem__ float* rsqrtLocalAddr = (__local_mem__ float*)rsqrtLocal.GetPhyAddr(); + uint32_t dim2Align = RoundUp(dim2); + __VEC_SCOPE__ + { + RegTensor base; + RegTensor mix; + RegTensor mix1; + RegTensor mix2; + RegTensor mix3; + RegTensor mix4; + RegTensor rsqrt; + RegTensor max; + RegTensor sum; + RegTensor sum1; + RegTensor sum2; + RegTensor sum3; + RegTensor sum4; + uint32_t sreg = dim2; + MaskReg pregLoop = UpdateMask(sreg); + for (uint16_t i = 0; i < dim0; i++) { + Duplicate(sum1, static_cast(0), pregLoop); + LoadInputDataWithBrc(rsqrt, rsqrtLocalAddr, pregLoop, i); + for (uint16_t j = 0; j < dim1; j++) { + LoadInputData(base, hcBaseLocalAddr, pregLoop, j * dim2Align); + LoadInputData(mix, mixLocalAddr, pregLoop, i * dim1 * dim2Align + j * dim2Align); + Mul(mix, mix, rsqrt, pregLoop); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + ReduceMax(max, mix, pregLoop); + Duplicate(max, max, pregLoop); + Sub(mix, mix, max, pregLoop); + Exp(mix, mix, pregLoop); + ReduceSum(sum, mix, pregLoop); + Duplicate(sum, sum, pregLoop); + Div(mix, mix, sum, pregLoop); + Adds(mix, mix, eps, pregLoop); + Add(sum1, sum1, mix, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, i * dim1 * dim2Align + j * dim2Align); + } + LocalMemBar(); + Adds(sum1, sum1, eps, pregLoop); + for (uint16_t j = 0; j < dim1; j++) { + LoadInputData(mix, combFragLocalAddr, pregLoop, i * dim1 * dim2Align + j * dim2Align); + Div(mix, mix, sum1, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, i * dim1 * dim2Align + j * dim2Align); + } + } + LocalMemBar(); + for (uint16_t i = 0; i < dim0; i++) { + LoadInputData(mix1, combFragLocalAddr, pregLoop, i * dim1 * dim2Align); + LoadInputData(mix2, combFragLocalAddr, pregLoop, i * dim1 * dim2Align + 1 * dim2Align); + LoadInputData(mix3, combFragLocalAddr, pregLoop, i * dim1 * dim2Align + 2 * dim2Align); + LoadInputData(mix4, combFragLocalAddr, pregLoop, i * dim1 * dim2Align + 3 * dim2Align); + for (uint16_t j = 0; j < iters; j++) { + Duplicate(sum, static_cast(0), pregLoop); + VFProcessIteration(sum, sum1, mix1, eps, pregLoop); + VFProcessIteration(sum, sum2, mix2, eps, pregLoop); + VFProcessIteration(sum, sum3, mix3, eps, pregLoop); + VFProcessIteration(sum, sum4, mix4, eps, pregLoop); + Adds(sum, sum, eps, pregLoop); + Div(mix1, mix1, sum, pregLoop); + Div(mix2, mix2, sum, pregLoop); + Div(mix3, mix3, sum, pregLoop); + Div(mix4, mix4, sum, pregLoop); + } + StoreOutputData(combFragLocalAddr, mix1, pregLoop, i * dim1 * dim2Align); + StoreOutputData(combFragLocalAddr, mix2, pregLoop, i * dim1 * dim2Align + 1 * dim2Align); + StoreOutputData(combFragLocalAddr, mix3, pregLoop, i * dim1 * dim2Align + 2 * dim2Align); + StoreOutputData(combFragLocalAddr, mix4, pregLoop, i * dim1 * dim2Align + 3 * dim2Align); + } + } +} + +template +__aicore__ inline void VFProcessY( + const LocalTensor& yLocal, const LocalTensor& mixLocal, const LocalTensor& xLocal, uint16_t bs, + uint16_t hcMult, uint16_t d) +{ + __local_mem__ T* yLocalAddr = (__local_mem__ T*)yLocal.GetPhyAddr(); + __local_mem__ float* mixLocalAddr = (__local_mem__ float*)mixLocal.GetPhyAddr(); + __local_mem__ T* xLocalAddr = (__local_mem__ T*)xLocal.GetPhyAddr(); + uint32_t dAlign = RoundUp(d); + uint16_t loopCount = CeilDiv(d, VL_FP32); + uint32_t hcMultAlign = RoundUp(hcMult); + if (loopCount > 1) { + __VEC_SCOPE__ + { + RegTensor x; + RegTensor mix; + RegTensor sum; + MaskReg pregLoop; + for (uint16_t i = 0; i < bs; i++) { + uint32_t sreg = d; + for (uint16_t j = 0; j < loopCount; j++) { + pregLoop = UpdateMask(sreg); + Duplicate(sum, static_cast(0), pregLoop); + for (uint16_t k = 0; k < hcMult; k++) { + LoadInputDataWithBrc(mix, mixLocalAddr, pregLoop, i * hcMultAlign + k); + LoadInputData(x, xLocalAddr, pregLoop, i * hcMult * dAlign + j * VL_FP32 + k * dAlign); + Mul(x, mix, x, pregLoop); + Add(sum, sum, x, pregLoop); + } + StoreOutputData(yLocalAddr, sum, pregLoop, i * dAlign + j * VL_FP32); + } + } + } + } else { + __VEC_SCOPE__ + { + RegTensor x; + RegTensor mix; + RegTensor sum; + uint32_t sreg = d; + MaskReg pregLoop = UpdateMask(sreg); + for (uint16_t i = 0; i < bs; i++) { + Duplicate(sum, static_cast(0), pregLoop); + for (uint16_t j = 0; j < hcMult; j++) { + LoadInputDataWithBrc(mix, mixLocalAddr, pregLoop, i * hcMultAlign + j); + LoadInputData(x, xLocalAddr, pregLoop, i * hcMult * dAlign + j * dAlign); + Mul(x, mix, x, pregLoop); + Add(sum, sum, x, pregLoop); + } + StoreOutputData(yLocalAddr, sum, pregLoop, i * dAlign); + } + } + } +} + +template +__aicore__ inline void CopyIn( + const GlobalTensor& inputGm, const LocalTensor& inputTensor, const uint16_t nBurst, const uint32_t copyLen, uint32_t srcStride = 0) +{ + DataCopyPadExtParams dataCopyPadExtParams; + dataCopyPadExtParams.isPad = false; + dataCopyPadExtParams.leftPadding = 0; + dataCopyPadExtParams.rightPadding = 0; + dataCopyPadExtParams.paddingValue = 0; + + DataCopyExtParams dataCoptExtParams; + dataCoptExtParams.blockCount = nBurst; + dataCoptExtParams.blockLen = copyLen * sizeof(T); + dataCoptExtParams.srcStride = srcStride * sizeof(T); + dataCoptExtParams.dstStride = 0; + DataCopyPad(inputTensor, inputGm, dataCoptExtParams, dataCopyPadExtParams); +} + +template +__aicore__ inline void CopyInWithLoopMode( + const GlobalTensor& inputGm, const LocalTensor& inputTensor, const uint16_t outerLoop, const uint16_t nBurst, const uint32_t copyLen, const uint32_t gmLastDim, uint32_t srcStride = 0) +{ + uint16_t copyLenAlign = RoundUp(copyLen); + LoopModeParams loopParams; + loopParams.loop2Size = 1; + loopParams.loop1Size = outerLoop; + loopParams.loop2SrcStride = 0; + loopParams.loop1SrcStride = gmLastDim * sizeof(T); + loopParams.loop2DstStride = 0; + loopParams.loop1DstStride = nBurst * copyLenAlign * sizeof(T); + + DataCopyPadExtParams dataCopyPadExtParams; + dataCopyPadExtParams.isPad = false; + dataCopyPadExtParams.leftPadding = 0; + dataCopyPadExtParams.rightPadding = 0; + dataCopyPadExtParams.paddingValue = 0; + + DataCopyExtParams dataCoptExtParams; + dataCoptExtParams.blockCount = nBurst; + dataCoptExtParams.blockLen = copyLen * sizeof(T); + dataCoptExtParams.srcStride = srcStride * sizeof(T); + dataCoptExtParams.dstStride = 0; + SetLoopModePara(loopParams, DataCopyMVType::OUT_TO_UB); + DataCopyPad(inputTensor, inputGm, dataCoptExtParams, dataCopyPadExtParams); + ResetLoopModePara(DataCopyMVType::OUT_TO_UB); +} + +template +__aicore__ inline void CopyOut( + const LocalTensor& outputTensor, const GlobalTensor& outputGm, const uint16_t nBurst, const uint32_t copyLen, uint32_t dstStride = 0) +{ + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = nBurst; + dataCopyParams.blockLen = copyLen * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = dstStride * sizeof(T); + DataCopyPad(outputGm, outputTensor, dataCopyParams); +} + +} // namespace HCPreSinkhorn + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_perf.h b/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_perf.h new file mode 100644 index 000000000..56e9c52f4 --- /dev/null +++ b/csrc/ascend/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_perf.h @@ -0,0 +1,204 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file hc_pre_sinkhorn_regbase_perf.h + * \brief + */ + +#ifndef HC_PRE_SINKHORN_REGBASE_PERF_H +#define HC_PRE_SINKHORN_REGBASE_PERF_H + +#include "kernel_operator.h" +#include "hc_pre_sinkhorn_regbase_base.h" + +namespace HcPreSinkhorn { +using namespace AscendC; +template +class HcPreSinkhornPerf { +public: + __aicore__ inline HcPreSinkhornPerf() + {} + + __aicore__ inline void Init( + GM_ADDR mixes, GM_ADDR rsqrt, GM_ADDR hcScale, GM_ADDR hcBase, GM_ADDR x, GM_ADDR y, GM_ADDR post, + GM_ADDR combFrag, GM_ADDR workspace, const HcPreSinkhornTilingData* tilingDataPtr, TPipe* pipePtr) + { + pipe = pipePtr; + tilingData = tilingDataPtr; + + mixesGm.SetGlobalBuffer((__gm__ float*)mixes); + rsqrtGm.SetGlobalBuffer((__gm__ float*)rsqrt); + hcScaleGm.SetGlobalBuffer((__gm__ float*)hcScale); + hcBaseGm.SetGlobalBuffer((__gm__ float*)hcBase); + xGm.SetGlobalBuffer((__gm__ T*)x); + yGm.SetGlobalBuffer((__gm__ T*)y); + postGm.SetGlobalBuffer((__gm__ float*)post); + combFragGm.SetGlobalBuffer((__gm__ float*)combFrag); + + // InQue + int64_t mixesQue01Size = tilingData->rowFactor * tilingData->hcMultAlign * 2 * sizeof(float); + pipe->InitBuffer(mixesQue01, 2, mixesQue01Size); + pipe->InitBuffer( + mixesQue2, 2, tilingData->rowFactor * tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(rsqrtQue, 2, RoundUp(tilingData->rowFactor) * sizeof(float)); + pipe->InitBuffer( + xQue, 2, tilingData->rowFactor * tilingData->hcMult * RoundUp(tilingData->dFactor) * sizeof(T)); + + // OutQue + pipe->InitBuffer( + yQue, 2, tilingData->rowFactor * RoundUp(tilingData->dFactor) * sizeof(T)); + pipe->InitBuffer(postQue, 2, tilingData->rowFactor * tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer( + combFragQue, 2, tilingData->rowFactor * tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); + + // TBuf + pipe->InitBuffer(hcBaseBuf0, tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(hcBaseBuf1, tilingData->hcMultAlign * sizeof(float)); + pipe->InitBuffer(hcBaseBuf2, tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); + + hcBase0Local = hcBaseBuf0.Get(); + hcBase1Local = hcBaseBuf1.Get(); + hcBase2Local = hcBaseBuf2.Get(); + } + + __aicore__ inline void Process() + { + int64_t curBlockIdx = GetBlockIdx(); + int64_t totalBlockNum = GetBlockNum(); + + int64_t rowOuterLoop = + (curBlockIdx == totalBlockNum - 1) ? tilingData->rowLoopOfTailBlock : tilingData->rowLoopOfFormerBlock; + int64_t tailRowFactor = (curBlockIdx == totalBlockNum - 1) ? tilingData->tailRowFactorOfTailBlock : + tilingData->tailRowFactorOfFormerBlock; + + CopyIn(hcBaseGm, hcBase0Local, 1, tilingData->hcMult); + CopyIn(hcBaseGm[tilingData->hcMult], hcBase1Local, 1, tilingData->hcMult); + CopyIn(hcBaseGm[tilingData->hcMult * 2], hcBase2Local, tilingData->hcMult, tilingData->hcMult); + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventId); + WaitFlag(eventId); + int64_t mixGmBaseOffset = curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMix; + int64_t xGmBaseOffset = curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMult * tilingData->d; + for (int64_t rowOuterIdx = 0; rowOuterIdx < rowOuterLoop; rowOuterIdx++) { + int64_t curRowFactor = (rowOuterIdx == rowOuterLoop - 1) ? tailRowFactor : tilingData->rowFactor; + mixes01Local = mixesQue01.AllocTensor(); + CopyIn( + mixesGm[mixGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMix], mixes01Local, + curRowFactor, tilingData->hcMult, tilingData->hcMix - tilingData->hcMult); + CopyIn( + mixesGm[mixGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMix + tilingData->hcMult], + mixes01Local[tilingData->rowFactor * tilingData->hcMultAlign], curRowFactor, tilingData->hcMult, + tilingData->hcMix - tilingData->hcMult); + mixesQue01.EnQue(mixes01Local); + + rsqrtLocal = rsqrtQue.AllocTensor(); + CopyIn( + rsqrtGm[curBlockIdx * tilingData->rowOfFormerBlock + rowOuterIdx * tilingData->rowFactor], rsqrtLocal, + 1, curRowFactor); + rsqrtQue.EnQue(rsqrtLocal); + + mixes01Local = mixesQue01.DeQue(); + rsqrtLocal = rsqrtQue.DeQue(); + VFProcessPre( + mixes01Local, mixes01Local, hcBase0Local, rsqrtLocal, hcScaleGm.GetValue(0), tilingData->eps, + curRowFactor, tilingData->hcMult); + for (int64_t dLoopIdx = 0; dLoopIdx < tilingData->dLoop; dLoopIdx++) { + int64_t curDFactor = + (dLoopIdx == tilingData->dLoop - 1) ? tilingData->tailDFactor : tilingData->dFactor; + xLocal = xQue.template AllocTensor(); + CopyIn( + xGm[xGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMult * tilingData->d + + dLoopIdx * tilingData->dFactor], + xLocal, tilingData->rowFactor * tilingData->hcMult, curDFactor, tilingData->d - curDFactor); + xQue.template EnQue(xLocal); + xLocal = xQue.template DeQue(); + yLocal = yQue.template AllocTensor(); + VFProcessY(yLocal, mixes01Local, xLocal, curRowFactor, tilingData->hcMult, curDFactor); + xQue.template FreeTensor(xLocal); + yQue.template EnQue(yLocal); + yLocal = yQue.template DeQue(); + CopyOut(yLocal, yGm[curBlockIdx * tilingData->rowOfFormerBlock * tilingData->d + rowOuterIdx * tilingData->rowFactor * tilingData->d + dLoopIdx * tilingData->dFactor], curRowFactor, curDFactor, tilingData->d - curDFactor); + yQue.template FreeTensor(yLocal); + } + + // post + postLocal = postQue.AllocTensor(); + VFProcessPost( + postLocal, mixes01Local[tilingData->rowFactor * tilingData->hcMultAlign], hcBase1Local, rsqrtLocal, + hcScaleGm.GetValue(1), tilingData->eps, curRowFactor, tilingData->hcMult); + mixesQue01.template FreeTensor(mixes01Local); + postQue.EnQue(postLocal); + postLocal = postQue.DeQue(); + CopyOut(postLocal, postGm[curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMult + rowOuterIdx * tilingData->rowFactor * tilingData->hcMult], curRowFactor, tilingData->hcMult); + postQue.FreeTensor(postLocal); + + // combFrag + mixes2Local = mixesQue2.AllocTensor(); + CopyInWithLoopMode( + mixesGm + [mixGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMix + tilingData->hcMult * 2], + mixes2Local, curRowFactor, tilingData->hcMult, tilingData->hcMult, tilingData->hcMix); + mixesQue2.EnQue(mixes2Local); + mixes2Local = mixesQue2.DeQue(); + + combFragLocal = combFragQue.AllocTensor(); + VFProcessCombFragRLessVLUseFourUnfold( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + mixesQue2.FreeTensor(mixes2Local); + rsqrtQue.FreeTensor(rsqrtLocal); + + combFragQue.EnQue(combFragLocal); + combFragLocal = combFragQue.DeQue(); + CopyOut(combFragLocal, combFragGm[curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMult * tilingData->hcMult + rowOuterIdx * tilingData->rowFactor * tilingData->hcMult * tilingData->hcMult], curRowFactor * tilingData->hcMult, tilingData->hcMult); + combFragQue.FreeTensor(combFragLocal); + } + } + +private: + TPipe* pipe; + const HcPreSinkhornTilingData* tilingData; + GlobalTensor mixesGm; + GlobalTensor rsqrtGm; + GlobalTensor hcScaleGm; + GlobalTensor hcBaseGm; + GlobalTensor xGm; + GlobalTensor yGm; + GlobalTensor postGm; + GlobalTensor combFragGm; + + TQue mixesQue01; + TQue mixesQue2; + TQue rsqrtQue; + TQue xQue; + TQue yQue; + TQue postQue; + TQue combFragQue; + + TBuf hcBaseBuf0; + TBuf hcBaseBuf1; + TBuf hcBaseBuf2; + + LocalTensor mixes01Local; + LocalTensor mixes2Local; + LocalTensor rsqrtLocal; + LocalTensor xLocal; + LocalTensor yLocal; + LocalTensor postLocal; + LocalTensor combFragLocal; + LocalTensor hcBase0Local; + LocalTensor hcBase1Local; + LocalTensor hcBase2Local; +}; + +} // namespace HCPreSinkhorn + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k/CMakeLists.txt b/csrc/ascend/moe/moe_gating_top_k/CMakeLists.txt new file mode 100644 index 000000000..298ed5677 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k/moe_gating_top_k_torch_adpt.h b/csrc/ascend/moe/moe_gating_top_k/moe_gating_top_k_torch_adpt.h new file mode 100644 index 000000000..bf7fe0221 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/moe_gating_top_k_torch_adpt.h @@ -0,0 +1,74 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef MOE_GATING_TOP_K_TORCH_ADPT_H +#define MOE_GATING_TOP_K_TORCH_ADPT_H +namespace vllm_fl { +std::tuple moe_gating_top_k( + const at::Tensor& x, + int64_t k, + int64_t k_group, + int64_t group_count, + int64_t group_select_mode, + int64_t renorm, + int64_t norm_type, + bool out_flag, + double routed_scaling_factor, + double eps, + const c10::optional& bias_opt + ) +{ + TORCH_CHECK(x.dim() == 2, "The x should be 2D"); + TORCH_CHECK( + x.scalar_type() == at::kHalf || x.scalar_type() == at::kFloat || x.scalar_type() == at::kBFloat16, + "float16、float32 or bfloat16 tensor expected but got a tensor with dtype: ", + x.scalar_type()); + + auto x_size = x.sizes(); + auto rows = x_size[0]; + auto expert_num = x_size[1]; + const at::Tensor &bias = c10::value_or_else(bias_opt, [] { return at::Tensor(); }); + if (bias.defined()) { + TORCH_CHECK(x.scalar_type() == bias.scalar_type(), "The dtype of x and bias should be same"); + TORCH_CHECK(bias.dim() == 1, "The bias should be 1D"); + auto bias_size = bias.sizes(); + TORCH_CHECK(bias_size[0] == expert_num, "The bias first dim should be same as x second dim"); + } + at::Tensor y = at::empty({rows, k}, x.options()); + at::Tensor expert_idx = at::empty({rows, k}, x.options().dtype(at::kInt)); + at::Tensor out = at::empty({rows, expert_num}, x.options().dtype(at::kFloat)); + + EXEC_NPU_CMD(aclnnMoeGatingTopK, + x, + bias, + k, + k_group, + group_count, + group_select_mode, + renorm, + norm_type, + out_flag, + routed_scaling_factor, + eps, + y, + expert_idx, + out + ); + + return std::tuple(y,expert_idx,out); +} + +} +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k/op_host/CMakeLists.txt b/csrc/ascend/moe/moe_gating_top_k/op_host/CMakeLists.txt new file mode 100644 index 000000000..e483f9be9 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_host/CMakeLists.txt @@ -0,0 +1,22 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + moe_gating_top_k_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME MoeGatingTopK + OPTIONS + --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE moe_gating_top_k ACLNNTYPE aclnn) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/moe/moe_gating_top_k/op_host/math_util.h b/csrc/ascend/moe/moe_gating_top_k/op_host/math_util.h new file mode 100644 index 000000000..edc1c8ea1 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_host/math_util.h @@ -0,0 +1,61 @@ +/** +* Copyright (c) 2025 Huawei Technologies Co., Ltd. +* This program is free software, you can redistribute it and/or modify it under the terms and conditions of +* CANN Open Software License Agreement Version 2.0 (the "License"). +* Please refer to the License for details. You may not use this file except in compliance with the License. +* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +* See LICENSE in the root of the software repository for the full text of the License. +*/ + +/*! + * \file math_util.h + * \brief + */ + +#ifndef TILING_MATMUL_MATH_UTIL_H +#define TILING_MATMUL_MATH_UTIL_H + +#include +#include +#include +#include +namespace matmul_tiling { +class MathUtil { +public: + static bool IsEqual(float leftValue, float rightValue); + template + static auto CeilDivision(T num1, T num2) -> T + { + if (num2 == 0) { + return 0; + } + return static_cast((static_cast(num1) + static_cast(num2) - 1) / + static_cast(num2)); + } + template + static auto Align(T num1, T num2) -> T + { + return CeilDivision(num1, num2) * num2; + } + static int32_t AlignDown(int32_t num1, int32_t num2); + static bool CheckMulOverflow(int32_t a, int32_t b, int32_t &c); + static int32_t MapShape(int32_t shape, bool roundUpFlag = true); + static void AddFactor(std::vector &dimsFactors, int32_t dim); + static void GetFactorCnt(const int32_t shape, int32_t &factorCnt, const int32_t factorStart, + const int32_t factorEnd); + static void GetFactorLayerCnt(const int32_t shape, int32_t &factorCnt, const int32_t factorStart, + const int32_t factorEnd); + static bool CheckFactorNumSatisfy(const int32_t dim); + static int32_t FindBestSingleCore(const int32_t oriShape, const int32_t mappedShape, const int32_t coreNum, + bool isKDim); + static void GetFactors(std::vector &factorList, int32_t srcNum, int32_t minFactor, int32_t maxFactor); + static void GetFactors(std::vector &factorList, int32_t srcNum, int32_t maxFactor); + static void GetBlockFactors(std::vector &factorList, const int32_t oriShape, const int32_t mpShape, + const int32_t coreNum, const int32_t maxNum); + static int32_t GetNonFactorMap(std::vector &factorList, int32_t srcNum, int32_t maxFactor); + static std::vector> GetFactorPairs(int32_t num); + static std::pair DivideIntoMainAndTail(int32_t num, int32_t divisor); +}; +} // namespace matmul_tiling +#endif // _MATH_UTIL_H_ diff --git a/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_def.cpp b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_def.cpp new file mode 100644 index 000000000..6eacd7656 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_def.cpp @@ -0,0 +1,70 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_def.cpp + * \brief + */ +#include "register/op_def_registry.h" + +namespace ops { +class MoeGatingTopK : public OpDef { +public: + explicit MoeGatingTopK(const char *name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("bias") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("expert_idx") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("out") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Attr("k").Int(); + this->Attr("k_group").AttrType(OPTIONAL).Int(1); + this->Attr("group_count").AttrType(OPTIONAL).Int(1); + this->Attr("group_select_mode").AttrType(OPTIONAL).Int(0); + this->Attr("renorm").AttrType(OPTIONAL).Int(0); + this->Attr("norm_type").AttrType(OPTIONAL).Int(0); + this->Attr("out_flag").AttrType(OPTIONAL).Bool(false); + this->Attr("routed_scaling_factor").AttrType(OPTIONAL).Float(1.0); + this->Attr("eps").AttrType(OPTIONAL).Float(1e-20f); + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + + OpAICoreConfig regbaseCfg; + regbaseCfg.DynamicCompileStaticFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .ExtendCfgInfo("opFile.value", "moe_gating_top_k_apt"); + } +}; + +OP_ADD(MoeGatingTopK); +} // namespace ops diff --git a/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_infershape.cpp b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_infershape.cpp new file mode 100644 index 000000000..0ed40da5d --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_infershape.cpp @@ -0,0 +1,147 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file moe_gating_top_k_infershape.cpp + * \brief + */ + +#include "exe_graph/runtime/infer_shape_context.h" +#include "register/op_impl_registry.h" +#include "tiling_base/error_log.h" + +#include + +#include +#define TO_STRING(x) std::string(#x) + +using namespace ge; +namespace ops { +static constexpr size_t DIM_ONE = 1; +static constexpr size_t DIM_TWO = 2; +static constexpr int64_t NEG_ONE = -1; +static constexpr int64_t X_INDEX = 0; +static constexpr int64_t BIAS_INDEX = 1; +static constexpr int64_t Y_INDEX = 0; +static constexpr int64_t EXPERT_IDX_INDEX = 1; +static constexpr int64_t OUT_INDEX = 2; + +static ge::graphStatus CheckInputShape(gert::InferShapeContext *context, const gert::Shape *xShape) +{ + int64_t XRows = xShape->GetDimNum() == 1U ? NEG_ONE : xShape->GetDim(0); + int64_t expertNum = xShape->GetDimNum() == 1U ? NEG_ONE : xShape->GetDim(1); + if (XRows < NEG_ONE || expertNum < NEG_ONE) { + OP_LOGE(context, "Invalid x shape, shape is %s.", TO_STRING(*xShape).c_str()); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus CheckInputDimsAndAttr(gert::InferShapeContext *context, const gert::Shape *xShape, + const int64_t k) +{ + if (xShape->GetDimNum() == 1U) { + if (xShape->GetDim(0) != ge::UNKNOWN_DIM_NUM) { + OP_LOGE(context, "The dynamic dim of x should be -2, current shape is %s.", + TO_STRING(*xShape).c_str()); + return ge::GRAPH_FAILED; + } + } else if (xShape->GetDimNum() != DIM_TWO) { + OP_LOGE(context, "The dim of x should be 2 or dynamic, current shape is %s.", + TO_STRING(*xShape).c_str()); + return ge::GRAPH_FAILED; + } + + if (k < 0) { + OP_LOGE(context, "k must be a non-negative number."); + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +static void ShowInputShapeInfo(gert::InferShapeContext *context, const gert::Shape *xShape, const int64_t k) +{ + OP_LOGD(context, "x shape is: %s.", TO_STRING(*xShape).c_str()); + OP_LOGD(context, "k is: %ld.", k); +} + +static void ShowOutputShapeInfo(gert::InferShapeContext *context, const gert::Shape *yShape, + const gert::Shape *expertIdxShape, const gert::Shape *outShape) +{ + OP_LOGD(context, "y shape is: %s after infershape.", TO_STRING(*yShape).c_str()); + OP_LOGD(context, "expert_idx shape is: %s after infershape.", TO_STRING(*expertIdxShape).c_str()); + OP_LOGD(context, "out shape is: %s after infershape.", TO_STRING(*outShape).c_str()); +} + +static ge::graphStatus InferShape4MoeGatingTopK(gert::InferShapeContext *context) +{ + OP_LOGD(context, "Begin to do MoeGatingTopKInfershape."); + + // 获取输入shape + const gert::Shape *xShape = context->GetInputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context, xShape); + gert::Shape *yShape = context->GetOutputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context, yShape); + gert::Shape *expertIdxShape = context->GetOutputShape(1); + OP_CHECK_NULL_WITH_CONTEXT(context, expertIdxShape); + gert::Shape *outShape = context->GetOutputShape(2); + OP_CHECK_NULL_WITH_CONTEXT(context, outShape); + + // 获取attr + auto attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + const int64_t *kPtr = attrs->GetAttrPointer(0); + OP_CHECK_NULL_WITH_CONTEXT(context, kPtr); + const int64_t k = *kPtr; + ShowInputShapeInfo(context, xShape, k); + + // 参数校验 + if (CheckInputDimsAndAttr(context, xShape, k) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + if (CheckInputShape(context, xShape) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + int64_t rows = xShape->GetDimNum() == 1U ? NEG_ONE : xShape->GetDim(0); + int64_t expertNum = xShape->GetDimNum() == 1U ? NEG_ONE : xShape->GetDim(1); + + yShape->SetDimNum(DIM_TWO); + yShape->SetDim(0U, rows); + yShape->SetDim(1U, k); + + expertIdxShape->SetDimNum(DIM_TWO); + expertIdxShape->SetDim(0U, rows); + expertIdxShape->SetDim(1U, k); + + outShape->SetDimNum(DIM_TWO); + outShape->SetDim(0U, rows); + outShape->SetDim(1U, expertNum); + + ShowOutputShapeInfo(context, yShape, expertIdxShape, outShape); + OP_LOGD(context, "End to do MoeGatingTopKInfershape."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataType4MoeGatingTopK(gert::InferDataTypeContext *context) +{ + OP_LOGD(context, "Begin to do MoeGatingTopKInferDataType."); + auto xDtype = context->GetInputDataType(0); + context->SetOutputDataType(Y_INDEX, xDtype); + context->SetOutputDataType(EXPERT_IDX_INDEX, ge::DT_INT32); + context->SetOutputDataType(OUT_INDEX, ge::DT_FLOAT); + OP_LOGD(context, "End to do MoeGatingTopKInferDataType."); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(MoeGatingTopK).InferShape(InferShape4MoeGatingTopK).InferDataType(InferDataType4MoeGatingTopK); +} // namespace ops diff --git a/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_proto.cpp b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_proto.cpp new file mode 100644 index 000000000..f10adf717 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_proto.cpp @@ -0,0 +1,15 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_proto.h + * \brief + */ +#include "moe_gating_top_k_proto.h" \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_proto.h b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_proto.h new file mode 100644 index 000000000..ed50d3449 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_proto.h @@ -0,0 +1,66 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_proto.h + * \brief + */ +#ifndef OPS_OP_PROTO_INC_MOEGATINGTOPK_H_ +#define OPS_OP_PROTO_INC_MOEGATINGTOPK_H_ + +#include "graph/operator_reg.h" + +namespace ge { + +/** + * @brief Compute renorm(sigmoid) and topk for moe input. + * + * @par Inputs: + * @li x: A 2D tensor which moe gating topk is applied, The shape is: (B*S, E), format supports ND, and data type must be float16, float or bfloat16. E(Expert num) can not be greater than 2048. E(Expert num) should be divisible by group_count. + * @li bias: A 1D tensor which is "bias" in moe gating topk. The shape is: (E), format supports ND, and data type must be the same as that of x. + * + * @par Outputs: + * @li y: A 2D tensor which is the topk value result of moe gating topk, format supports ND, and data type must be the same as that of x. + The size of the non-1 axis must be the same as that of the corresponding axis of x. + The size of the -1 axis must be the same as that of k. + * @li expert_idx: A 2D tensor which is the topk index result of moe gating topk, format supports ND, and data type must be int. The shape must be the same as that of y. + * @li out: A 2D tensor which is the renorm result of moe gating topk, format supports ND, and data type must be float. The shape must be the same as that of x. + * + * @par Attributes: + * @li k: A required attribute of type int. The value must greater than 0 and less than or equal to expert_num / group_count * k_group, indicating the topk value. + * @li k_group: An optional attribute of type int. It can not be less than 1, and can not be greater than group_count, indicating the topk group value. The default value is 1. + * @li group_count: An optional attribute of type int. It can not be less than 1, indicating the group count. The group_count * align_32(expert_num / group_count) can not be greater than 2048. The default value is 1. + * @li group_select_mode: An optional attribute of type int. 0 indicating that sort group by max values, 1 indicating that sort group by sum of top-2 values. The default value is 0. + * @li renorm: An optional attribute of type int. It can only be 0 now, indicating that norm firstly and then topk. The default value is 0. + * @li norm_type: An optional attribute of type int. 0 indicating that the softmax function is used, 1 indicating that the sigmoid function is used. The default value is 0. + * @li out_flag: An optional attribute of type bool. true indicating that has renorm output, false indicating that does not have renorm output. The default value is false. + * @li routed_scaling_factor: An optional attribute of type float, indicating the routed_scaling_factor coefficient in use. The default value is 1.0. + * @li eps: An optional attribute of type float, indicating the eps coefficient in use. The default value is 1e-20. + */ +REG_OP(MoeGatingTopK) + .INPUT(x, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16})) + .OPTIONAL_INPUT(bias, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16})) + .OUTPUT(y, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16})) + .OUTPUT(expert_idx, TensorType({DT_INT32})) + .OUTPUT(out, TensorType({DT_FLOAT})) + .REQUIRED_ATTR(k, Int) + .ATTR(k_group, Int, 1) + .ATTR(group_count, Int, 1) + .ATTR(group_select_mode, Int, 0) + .ATTR(renorm, Int, 0) + .ATTR(norm_type, Int, 0) + .ATTR(out_flag, Bool, false) + .ATTR(routed_scaling_factor, Float, 1.0) + .ATTR(eps, Float, 1e-20f) + .OP_END_FACTORY_REG(MoeGatingTopK) + +} // namespace ge + +#endif // OPS_OP_PROTO_INC_MOEGATINGTOPK_H_ \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_tiling.cpp b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_tiling.cpp new file mode 100644 index 000000000..e8b695f83 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_tiling.cpp @@ -0,0 +1,573 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file moe_gating_top_k_tiling.cpp + * \brief + */ +#include +#include "register/op_def_registry.h" +#include "exe_graph/runtime/infer_shape_context.h" +#include "register/op_impl_registry.h" +#include "tiling_base/tiling_base.h" +#include "tiling_base/tiling_templates_registry.h" +#include "platform/platform_info.h" + + + +#include "tiling_base/error_log.h" +#include "moe_gating_top_k_tiling.h" + + +#ifndef CEIL_ALIGN +#define CEIL_ALIGN(val, align) ((((val) + (align) - 1) / (align)) * (align)) +#endif + +#ifndef CEIL_DIV +#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b)) +#endif +namespace optiling { +const static int64_t GROUP_SELECT_MODE_MAX = 0; +const static int64_t GROUP_SELECT_MODE_SUM = 1; +const static int64_t RENORM_NO = 0; +const static int64_t RENORM_L1 = 1; +const static int64_t NORM_TYPE_SOFTMAX = 0; +const static int64_t NORM_TYPE_SIGMOID = 1; +const static int64_t OUT_FLAG_FALSE = 0; +const static int64_t OUT_FLAG_TRUE = 1; +const static size_t X_INPUT_DIMS = 2; +const static size_t BIAS_INPUT_DIMS = 1; +const static size_t Y_OUTPUT_DIMS = 2; +const static size_t EXPERT_IDX_OUTPUY_DIMS = 2; +const static size_t OUT_OUTPUT_DIMS = 2; +const static int64_t MAX_EXPERT_COUNT = 2048; + +const static int64_t X_INPUT_INDEX = 0; +const static int64_t BIAS_INPUT_INDEX = 1; +const static int64_t Y_OUTPUT_INDEX = 0; +const static int64_t EXPERT_IDX_OUTPUT_INDEX = 1; +const static int64_t OUT_OUTPUT_INDEX = 2; +const static int64_t K_ATTR_INDEX = 0; +const static int64_t K_GROUP_ATTR_INDEX = 1; +const static int64_t GROUP_COUNT_ATTR_INDEX = 2; +const static int64_t GROUP_SELECT_MODE_ATTR_INDEX = 3; +const static int64_t RENORM_ATTR_INDEX = 4; +const static int64_t NORM_TYPE_ATTR_INDEX = 5; +const static int64_t OUT_FLAG_ATTR_INDEX = 6; +const static int64_t ROUTED_SCALING_FACTOR_ATTR_INDEX = 7; +const static int64_t EPS_ATTR_INDEX = 8; +const static int64_t DEFAULT_WORKSPACE_SIZE = 16777216; +const static uint32_t DATATYPESIZE_FLOAT = 4; +const static bool IS_LARGEST = true; +const static bool IS_INITINDEX = false; +const static bool IS_REUSESOURCE = false; +const static uint64_t WITH_GROUP_CONDITION = 1; +const static uint64_t WITHOUT_GROUP_CONDITION = 2; +const static uint64_t MAX_IN_GROUP_CONDITION = 3; +constexpr int32_t ROW_COUNT_PER_TASK = 1; + +const static uint64_t TILING_KEY_EXPERTNUM_GROUPNUM_ALIGN_HIGH_PERF = 0; +const static uint64_t TILING_KEY_WITHOUT_GROUP = 1; +const static uint64_t TILING_KEY_GENERALIZED = 2; + +inline static int64_t CeilLog4(int64_t x) +{ + return static_cast(std::ceil(std::log(x) / std::log(4))); // 4 for four +} + +class MoeGatingTopKTilingBase : public Ops::Transformer::OpTiling::TilingBaseClass { +public: + explicit MoeGatingTopKTilingBase(gert::TilingContext *context) : Ops::Transformer::OpTiling::TilingBaseClass(context) + { + Reset(); + } + ~MoeGatingTopKTilingBase() override = default; + + void Reset(gert::TilingContext *context) override + { + TilingBaseClass::Reset(context); + Reset(); + } + +protected: + bool IsCapable() override + { + return true; + } + + ge::graphStatus GetPlatformInfo() override; + + ge::graphStatus GetShapeAttrsInfo() override; + + ge::graphStatus DoOpTiling() override; + + ge::graphStatus DoLibApiTiling() override; + + uint64_t GetTilingKey() const override; + + ge::graphStatus GetWorkspaceSize() override; + + ge::graphStatus PostTiling() override; + void Reset(); + +private: + ge::graphStatus CheckInputShape(); + ge::graphStatus CheckAttr(); + ge::graphStatus CheckOutShape(); + void SplitRows(); + void CalTmpBufUbSize(); + + const gert::Shape *xShape_ = nullptr; + const gert::Shape *biasShape_ = nullptr; + const gert::Shape *yShape_ = nullptr; + const gert::Shape *expertIdxShape_ = nullptr; + const gert::Shape *outShape_ = nullptr; + + int64_t rows_ = 0; + int64_t expertCount_ = 0; + int64_t addBias_ = 0; + + int64_t k_ = 0; + int64_t kGroup_ = 0; + int64_t groupCount_ = 0; + int64_t perGroupExpertCount_ = 0; + int64_t groupSelectMode_ = GROUP_SELECT_MODE_MAX; + int64_t renorm_ = RENORM_NO; + int64_t normType_ = NORM_TYPE_SOFTMAX; + int64_t outFlag_ = OUT_FLAG_FALSE; + float routedScalingFactor_ = 1.0; + float eps_ = 1e-20f; + + int64_t inputDtypeSize_; + const char *opName_ = ""; + MoeGatingTopKTilingData moeGatingTopKTilingData_; +}; + +ge::graphStatus MoeGatingTopKTilingBase::CheckInputShape() +{ + size_t xDimNum = xShape_->GetDimNum(); + + OP_CHECK_IF(xDimNum != X_INPUT_DIMS, + + OP_LOGE(context_, "The dim number of x is: %zu, but should be %zu.", xDimNum, X_INPUT_DIMS), + return ge::GRAPH_FAILED); + + + rows_ = xShape_->GetDim(0); + expertCount_ = xShape_->GetDim(1); + + moeGatingTopKTilingData_.set_rowCount(rows_); + moeGatingTopKTilingData_.set_expertCount(expertCount_); + if (biasShape_ != nullptr) { + addBias_ = 1; + size_t biasDimNum = biasShape_->GetDimNum(); + OP_CHECK_IF(biasDimNum != BIAS_INPUT_DIMS, + OP_LOGE(context_, "The dim number of bias is: %zu, but should be %zu.", biasDimNum, BIAS_INPUT_DIMS), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + biasShape_->GetDim(0) != expertCount_, + OP_LOGE(context_, "The first dim of bias is: %ld, but should be %ld.", biasShape_->GetDim(0), expertCount_), + return ge::GRAPH_FAILED); + + } + moeGatingTopKTilingData_.set_addBias(addBias_); + + OP_CHECK_IF(k_ > expertCount_, + OP_LOGE(context_, "k is: %ld, expert num is: %ld, k cannot be greater than expert num.", k_, expertCount_), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingBase::CheckAttr() +{ + OP_CHECK_IF( + expertCount_ > MAX_EXPERT_COUNT, + OP_LOGE(context_, "expert count is: %ld, but should not greater than %ld.", expertCount_, MAX_EXPERT_COUNT), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(k_ <= 0, OP_LOGE(context_, "k is: %ld, but should be greater than 0.", k_), return ge::GRAPH_FAILED); + + OP_CHECK_IF(kGroup_ <= 0, OP_LOGE(context_, "k_group is: %ld, but should be greater than 0.", kGroup_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(kGroup_ > groupCount_, + OP_LOGE(context_, "k_group is: %ld, but should not greater than %ld.", kGroup_, groupCount_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(groupCount_ <= 0, OP_LOGE(context_, "group_count is: %ld, but should be greater than 0.", groupCount_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(normType_ != NORM_TYPE_SOFTMAX && normType_ != NORM_TYPE_SIGMOID, + OP_LOGE(context_, "norm type is: %ld, but currently only support %ld and %ld.", normType_, + NORM_TYPE_SOFTMAX, NORM_TYPE_SIGMOID), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(groupSelectMode_ != GROUP_SELECT_MODE_SUM && groupSelectMode_ != GROUP_SELECT_MODE_MAX, + OP_LOGE(context_, "group select mode is: %ld, but currently only support %ld and %ld.", groupSelectMode_, + GROUP_SELECT_MODE_SUM, GROUP_SELECT_MODE_MAX), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(renorm_ != RENORM_NO && renorm_ != RENORM_L1, + OP_LOGE(context_, "renorm is: %ld, but currently only support %ld.", renorm_, RENORM_NO), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(expertCount_ % groupCount_ != 0, + OP_LOGE(context_, "Expert count : %ld is not divisible by k_group: %ld", expertCount_, groupCount_), + return ge::GRAPH_FAILED); + + perGroupExpertCount_ = expertCount_ / groupCount_; + + OP_LOGI(context_, "perGroupExpertCount_: %ld", perGroupExpertCount_); + + OP_CHECK_IF(perGroupExpertCount_ < 1, + OP_LOGE(context_, "group expert count is: %ld, but should be greater than 1.", perGroupExpertCount_), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + groupSelectMode_ == GROUP_SELECT_MODE_SUM && perGroupExpertCount_ < 2, + OP_LOGE(context_, + "group expert count is: %ld, if group select mode is: %ld, group expert count should be greater than 1.", + perGroupExpertCount_, groupSelectMode_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(k_ > kGroup_ * perGroupExpertCount_, + OP_LOGE(context_, "k is: %ld, but should be smaller than %ld.", k_, kGroup_ * perGroupExpertCount_), + return ge::GRAPH_FAILED); + int64_t groupExpertCountAlign = CEIL_ALIGN(perGroupExpertCount_, 32L); + OP_LOGI(context_, "333groupExpertCountAlign: %ld", groupExpertCountAlign); + if (groupCount_ != 1 && groupCount_ != expertCount_ && kGroup_ != groupCount_) { + + OP_CHECK_IF(groupCount_ * groupExpertCountAlign > MAX_EXPERT_COUNT, + OP_LOGE(context_, "group count * group expert count align is: %ld, but should not greater than %ld.", + groupCount_ * groupExpertCountAlign, MAX_EXPERT_COUNT), + return ge::GRAPH_FAILED); + } + + moeGatingTopKTilingData_.set_perGroupExpertCount(perGroupExpertCount_); + moeGatingTopKTilingData_.set_perGroupExpertCountAlign(groupExpertCountAlign); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingBase::GetShapeAttrsInfo() +{ + opName_ = context_->GetNodeName(); + OP_LOGI(context_, "111GetShapeAttrsInfo: opName = %s", opName_); + auto xShapePtr = context_->GetInputShape(X_INPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, xShapePtr); + xShape_ = &xShapePtr->GetStorageShape(); + OP_LOGI(context_, "112xShape: %s", xShape_->ToString().c_str()); + + auto biasShapePtr = context_->GetOptionalInputShape(BIAS_INPUT_INDEX); + biasShape_ = biasShapePtr == nullptr ? nullptr : &biasShapePtr->GetStorageShape(); + if (biasShape_ != nullptr) { + OP_LOGI(context_, "113biasShape: %s", biasShape_->ToString().c_str()); + } + + auto yShapePtr = context_->GetOutputShape(Y_OUTPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, yShapePtr); + yShape_ = &yShapePtr->GetStorageShape(); + OP_LOGI(context_, "115yShape: %s", yShape_->ToString().c_str()); + auto expertIdxPtr = context_->GetOutputShape(EXPERT_IDX_OUTPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, expertIdxPtr); + expertIdxShape_ = &expertIdxPtr->GetStorageShape(); + OP_LOGI(context_, "116expertIdxShape: %s", expertIdxShape_->ToString().c_str()); + auto outPtr = context_->GetOutputShape(OUT_OUTPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, outPtr); + outShape_ = &outPtr->GetStorageShape(); + if (outShape_ != nullptr) { + OP_LOGI(context_, "117outShape: %s", outShape_->ToString().c_str()); + } + + auto x = context_->GetInputDesc(X_INPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, x); + auto xDtype = x->GetDataType(); + OP_CHECK_IF( + (xDtype != ge::DataType::DT_FLOAT && xDtype != ge::DataType::DT_FLOAT16 && xDtype != ge::DataType::DT_BF16), + OP_LOGE(context_, "x dtype %s error, only supports float32, half, bf16. please check.", + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + + if (biasShapePtr != nullptr) { + auto biasDtype = context_->GetOptionalInputDesc(BIAS_INPUT_INDEX)->GetDataType(); + OP_LOGI(context_, "118bias dtype: %s", ge::TypeUtils::DataTypeToSerialString(biasDtype).c_str()); + OP_CHECK_IF((biasDtype != xDtype), + OP_LOGE(context_, "bias dtype %s not equal x dtype %s, please check.", + ge::TypeUtils::DataTypeToSerialString(biasDtype).c_str(), + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + } + + auto yDesc = context_->GetOutputDesc(Y_OUTPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, yDesc); + auto yDtype = yDesc->GetDataType(); + OP_LOGI(context_, "119y dtype: %s", ge::TypeUtils::DataTypeToSerialString(yDtype).c_str()); + OP_CHECK_IF((yDtype != xDtype), + OP_LOGE(context_, "y out dtype %s must be the same with x dtype %s.", + ge::TypeUtils::DataTypeToSerialString(yDtype).c_str(), + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + + auto expertIdDesc = context_->GetOutputDesc(EXPERT_IDX_OUTPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, expertIdDesc); + auto expertIdDtype = expertIdDesc->GetDataType(); + OP_LOGI(context_, "120expertId dtype: %s", ge::TypeUtils::DataTypeToSerialString(expertIdDtype).c_str()); + OP_CHECK_IF((expertIdDtype != ge::DataType::DT_INT32), + OP_LOGE(context_, "expertId out dtype %s error, only supports int32. please check.", + ge::TypeUtils::DataTypeToSerialString(expertIdDtype).c_str()), + return ge::GRAPH_FAILED); + + auto normOutDesc = context_->GetOutputDesc(OUT_OUTPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, normOutDesc); + auto normOutDtype = normOutDesc->GetDataType(); + OP_CHECK_IF((normOutDtype != ge::DataType::DT_FLOAT), + OP_LOGE(context_, "norm out dtype %s error, only supports float. please check.", + ge::TypeUtils::DataTypeToSerialString(normOutDtype).c_str()), + return ge::GRAPH_FAILED); + + + auto attrs = context_->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context_, attrs); + + const int64_t *kPtr = attrs->GetAttrPointer(K_ATTR_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, kPtr); + k_ = *kPtr; + OP_LOGI(context_, "Attr k is: %ld", k_); + moeGatingTopKTilingData_.set_k(k_); + + OP_LOGI(context_, "Attr k is: %ld ", k_); + + const int64_t *kGroupPtr = attrs->GetAttrPointer(K_GROUP_ATTR_INDEX); + if (kGroupPtr != nullptr) { + kGroup_ = *kGroupPtr; + OP_LOGI(context_, "Attr k_group is: %ld", kGroup_); + moeGatingTopKTilingData_.set_kGroup(kGroup_); + } + OP_LOGI(context_, "Attr k_group is: %ld ", kGroup_); + + const int64_t *groupCountPtr = attrs->GetAttrPointer(GROUP_COUNT_ATTR_INDEX); + if (groupCountPtr != nullptr) { + groupCount_ = *groupCountPtr; + OP_LOGI(context_, "Attr group_count is: %ld", groupCount_); + moeGatingTopKTilingData_.set_groupCount(groupCount_); + } + OP_LOGI(context_, "Attr group_count is: %ld ", groupCount_); + + const int64_t *groupSelectModePtr = attrs->GetAttrPointer(GROUP_SELECT_MODE_ATTR_INDEX); + if (groupSelectModePtr != nullptr) { + groupSelectMode_ = *groupSelectModePtr; + OP_LOGI(context_, "Attr group_select_mode is: %ld", groupSelectMode_); + moeGatingTopKTilingData_.set_groupSelectMode(groupSelectMode_); + } + OP_LOGI(context_, "Attr group_select_mode is: %ld ", groupSelectMode_); + + const int64_t *renormPtr = attrs->GetAttrPointer(RENORM_ATTR_INDEX); + if (renormPtr != nullptr) { + renorm_ = *renormPtr; + OP_LOGI(context_, "Attr renorm is: %ld", renorm_); + moeGatingTopKTilingData_.set_renorm(renorm_); + } + OP_LOGI(context_, "Attr renorm is: %ld ", renorm_); + + const int64_t *normTypePtr = attrs->GetAttrPointer(NORM_TYPE_ATTR_INDEX); + if (normTypePtr != nullptr) { + normType_ = *normTypePtr; + OP_LOGI(context_, "Attr norm_type is: %ld", normType_); + moeGatingTopKTilingData_.set_normType(normType_); + } + OP_LOGI(context_, "Attr norm_type is: %ld ", normType_); + + const bool *outFlagPtr = attrs->GetAttrPointer(OUT_FLAG_ATTR_INDEX); + if (outFlagPtr != nullptr) { + outFlag_ = (*outFlagPtr) ? 1 : 0; + OP_LOGI(context_, "Attr out_flag is: %ld", outFlag_); + moeGatingTopKTilingData_.set_outFlag(outFlag_); + } + OP_LOGI(context_, "Attr out_flag is: %ld ", outFlag_); + + const float *routedScalingFactorPtr = attrs->GetAttrPointer(ROUTED_SCALING_FACTOR_ATTR_INDEX); + if (routedScalingFactorPtr != nullptr) { + routedScalingFactor_ = *routedScalingFactorPtr; + OP_LOGI(context_, "Attr routed_scaling_factor is: %f", routedScalingFactor_); + moeGatingTopKTilingData_.set_routedScalingFactor(routedScalingFactor_); + } + OP_LOGI(context_, "Attr routed_scaling_factor is: %f ", routedScalingFactor_); + + const float *epsPtr = attrs->GetAttrPointer(EPS_ATTR_INDEX); + if (epsPtr != nullptr) { + eps_ = *epsPtr; + OP_LOGI(context_, "Attr eps is: %f", eps_); + moeGatingTopKTilingData_.set_eps(eps_); + } + OP_LOGI(context_, "Attr eps is: %f ", eps_); + + inputDtypeSize_ = static_cast(ge::GetSizeByDataType(context_->GetInputDesc(0)->GetDataType())); + OP_LOGI(context_, "inputDtypeSize_: %ld", inputDtypeSize_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingBase::GetPlatformInfo() + +{ + auto platformInfo = context_->GetPlatformInfo(); + OP_CHECK_IF(platformInfo == nullptr, OP_LOGE(context_, "fail to get platform info"), return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + aicoreParams_.blockDim = ascendcPlatform.GetCoreNumAiv(); + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + aicoreParams_.ubSize = ubSizePlatForm; + OP_LOGI(context_, "GetPlatformInfo: blockDim = %ld, ubSize = %lu", aicoreParams_.blockDim, aicoreParams_.ubSize); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingBase::CheckOutShape() +{ + OP_LOGI(context_, "555CheckOutShape: yShape_: %s, xShape_: %s", yShape_->ToString().c_str(), xShape_->ToString().c_str()); + OP_CHECK_IF((yShape_->GetDimNum() != xShape_->GetDimNum()), + OP_LOGE(context_, "y out shape num %zu and x shape num %zu not equal, please check.", yShape_->GetDimNum(), + xShape_->GetDimNum()), + return ge::GRAPH_FAILED); + OP_CHECK_IF((expertIdxShape_->GetDimNum() != xShape_->GetDimNum()), + OP_LOGE(context_, "expertId out shape num %zu and x shape num %zu not equal, please check.", + expertIdxShape_->GetDimNum(), xShape_->GetDimNum()), + return ge::GRAPH_FAILED); + if (outShape_ != nullptr) { + OP_CHECK_IF((outShape_->GetDimNum() != xShape_->GetDimNum()), + OP_LOGE(context_, "norm out shape num %zu and x shape num %zu not equal, please check.", + outShape_->GetDimNum(), xShape_->GetDimNum()), + return ge::GRAPH_FAILED); + } + + OP_CHECK_IF((yShape_->GetDim(0) != xShape_->GetDim(0)), + OP_LOGE(context_, "y out dim[0] %ld not equal x dim[0] %ld, please check.", yShape_->GetDim(0), + xShape_->GetDim(0)), + return ge::GRAPH_FAILED); + OP_CHECK_IF((expertIdxShape_->GetDim(0) != xShape_->GetDim(0)), + OP_LOGE(context_, "expertId out dim[0] %ld not equal x dim[0] %ld, please check.", + expertIdxShape_->GetDim(0), xShape_->GetDim(0)), + return ge::GRAPH_FAILED); + if (outFlag_ && outShape_ != nullptr) { + OP_CHECK_IF((outShape_->GetDim(0) != xShape_->GetDim(0)), + OP_LOGE(context_, "norm out dim[0] %ld and x dim[0] %ld not equal, please check.", + outShape_->GetDim(0), outShape_->GetDim(0)), + return ge::GRAPH_FAILED); + } + + OP_CHECK_IF((yShape_->GetDim(1) != k_), + OP_LOGE(context_, "y dim[1] %ld not equal k %ld, please check.", yShape_->GetDim(1), k_), + return ge::GRAPH_FAILED); + OP_CHECK_IF((expertIdxShape_->GetDim(1) != k_), + OP_LOGE(context_, "expertId dim[1] %ld not equal k %ld, please check.", expertIdxShape_->GetDim(1), k_), + return ge::GRAPH_FAILED); + if (outFlag_ && outShape_ != nullptr) { + OP_CHECK_IF((outShape_->GetDim(1) != xShape_->GetDim(1)), + OP_LOGE(context_, "normOut dim[1] %ld and x dim[1] %ld not equal, please check.", outShape_->GetDim(1), + xShape_->GetDim(1)), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +void MoeGatingTopKTilingBase::SplitRows() +{ + int64_t perCoreRows = CEIL_DIV(rows_, static_cast(aicoreParams_.blockDim)); + int64_t needCoreNum = CEIL_DIV(rows_, perCoreRows); + // perCoreRows cannot be 0 + int64_t lastCoreRows = rows_ % perCoreRows == 0 ? perCoreRows : rows_ % perCoreRows; + moeGatingTopKTilingData_.set_needCoreNum(needCoreNum); + moeGatingTopKTilingData_.set_perCoreRowCount(perCoreRows); + moeGatingTopKTilingData_.set_lastCoreRowCount(lastCoreRows); + int64_t vmsCount = CeilLog4(CEIL_DIV(kGroup_, 4L)); + OP_LOGI(context_, "vms count is: %ld", vmsCount); + moeGatingTopKTilingData_.set_vmsCount(vmsCount); +} + +void MoeGatingTopKTilingBase::CalTmpBufUbSize() + +{ + + std::vector shape_vec = {expertCount_}; + ge::Shape shape(shape_vec); + uint32_t maxValue = 0; + uint32_t minValue = 0; + AscendC::GetSigmoidMaxMinTmpSize(shape, sizeof(float), false, maxValue, minValue); + + int64_t indexTmpBuf = (expertCount_ + 31) / 32 * 32 * static_cast(sizeof(float)); + moeGatingTopKTilingData_.set_calTmpBufUbSize(std::max(indexTmpBuf, static_cast(minValue))); +} + +ge::graphStatus MoeGatingTopKTilingBase::DoOpTiling() +{ + + OP_LOGI(context_, "DoOpTiling: start"); + auto ret = CheckInputShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckOutShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckAttr(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + CalTmpBufUbSize(); + SplitRows(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingBase::DoLibApiTiling() +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingBase::GetWorkspaceSize() +{ + + workspaceSize_ = DEFAULT_WORKSPACE_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingBase::PostTiling() +{ + context_->SetBlockDim(moeGatingTopKTilingData_.get_needCoreNum()); + size_t *currentWorkspace = context_->GetWorkspaceSizes(1); + currentWorkspace[0] = workspaceSize_; + moeGatingTopKTilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), + context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(moeGatingTopKTilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +uint64_t MoeGatingTopKTilingBase::GetTilingKey() const +{ + + if (expertCount_ == 256 && groupCount_ == 8 && kGroup_ == 4 && k_ <= 32 && addBias_ && + groupSelectMode_ == GROUP_SELECT_MODE_SUM && renorm_ == RENORM_NO && normType_ == NORM_TYPE_SIGMOID && + !outFlag_) { + + return TILING_KEY_EXPERTNUM_GROUPNUM_ALIGN_HIGH_PERF; + } else if (groupCount_ == 1 || groupCount_ == expertCount_ || kGroup_ == groupCount_) { + return TILING_KEY_WITHOUT_GROUP; + } else { + return TILING_KEY_GENERALIZED; + } +} + +void MoeGatingTopKTilingBase::Reset() +{ + opName_ = nullptr; + return; +} + +REGISTER_OPS_TILING_TEMPLATE(MoeGatingTopK, MoeGatingTopKTilingBase, 2000); +} // namespace optiling diff --git a/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_tiling.h b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_tiling.h new file mode 100644 index 000000000..0152a6dfb --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_tiling.h @@ -0,0 +1,86 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_tiling.h + * \brief + */ + +#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_MOE_GATING_TOP_K_H +#define AIR_CXX_RUNTIME_V2_OP_IMPL_MOE_GATING_TOP_K_H + +#include +#include +#include +#include + + +#include "tiling_base/tiling_base.h" +#include "tiling_base/tiling_templates_registry.h" +#include "register/op_def_registry.h" +#include "register/op_impl_registry.h" +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "tiling_base/error_log.h" + +#include "register/op_impl_registry.h" +#include "platform/platform_infos_def.h" +#include "math_util.h" +//#include "util/extern_math_util.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(MoeGatingTopKTilingData) +TILING_DATA_FIELD_DEF(int64_t, needCoreNum); +TILING_DATA_FIELD_DEF(int64_t, rowCount); +TILING_DATA_FIELD_DEF(int64_t, perCoreRowCount); +TILING_DATA_FIELD_DEF(int64_t, lastCoreRowCount); +TILING_DATA_FIELD_DEF(int64_t, expertCount); +TILING_DATA_FIELD_DEF(int64_t, addBias); +TILING_DATA_FIELD_DEF(int64_t, k); +TILING_DATA_FIELD_DEF(int64_t, kGroup); +TILING_DATA_FIELD_DEF(int64_t, groupCount); +TILING_DATA_FIELD_DEF(int64_t, perGroupExpertCount); +TILING_DATA_FIELD_DEF(int64_t, perGroupExpertCountAlign); +TILING_DATA_FIELD_DEF(int64_t, groupSelectMode); +TILING_DATA_FIELD_DEF(int64_t, renorm); +TILING_DATA_FIELD_DEF(int64_t, normType); +TILING_DATA_FIELD_DEF(int64_t, outFlag); +TILING_DATA_FIELD_DEF(int64_t, vmsCount); +TILING_DATA_FIELD_DEF(float, routedScalingFactor); +TILING_DATA_FIELD_DEF(float, eps); +TILING_DATA_FIELD_DEF(int64_t, calTmpBufUbSize); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(MoeGatingTopK, MoeGatingTopKTilingData) + +BEGIN_TILING_DATA_DEF(MoeGatingTopKRegbaseTilingData) +TILING_DATA_FIELD_DEF(int64_t, needCoreNum); +TILING_DATA_FIELD_DEF(int64_t, rowCount); +TILING_DATA_FIELD_DEF(int64_t, perCoreRowCount); +TILING_DATA_FIELD_DEF(int64_t, lastCoreRowCount); +TILING_DATA_FIELD_DEF(int64_t, expertCount); +TILING_DATA_FIELD_DEF(int64_t, addBias); +TILING_DATA_FIELD_DEF(int64_t, k); +TILING_DATA_FIELD_DEF(int64_t, kGroup); +TILING_DATA_FIELD_DEF(int64_t, groupCount); +TILING_DATA_FIELD_DEF(int64_t, perGroupExpertCount); +TILING_DATA_FIELD_DEF(int64_t, perGroupExpertCountAlign); +TILING_DATA_FIELD_DEF(int64_t, groupSelectMode); +TILING_DATA_FIELD_DEF(int64_t, renorm); +TILING_DATA_FIELD_DEF(int64_t, normType); +TILING_DATA_FIELD_DEF(int64_t, outFlag); +TILING_DATA_FIELD_DEF(int64_t, vmsCount); +TILING_DATA_FIELD_DEF(float, routedScalingFactor); +TILING_DATA_FIELD_DEF(float, eps); +TILING_DATA_FIELD_DEF_STRUCT(SoftMaxTiling, softmaxTilingData); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(MoeGatingTopK_10000, MoeGatingTopKRegbaseTilingData) +struct MoeGatingTopKCompileInfo {}; +} // namespace optiling +#endif // AIR_CXX_RUNTIME_V2_OP_IMPL_MOE_GATING_TOP_K_H diff --git a/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_tiling_arch35.cpp b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_tiling_arch35.cpp new file mode 100644 index 000000000..9beb336d4 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_tiling_arch35.cpp @@ -0,0 +1,521 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file moe_gating_top_k_tiling_arch35.cpp + * \brief + */ + +#include "tiling_base/error_log.h" +#include "moe_gating_top_k_tiling.h" +#include "register/op_def_registry.h" +#include "platform/platform_info.h" +#include "tiling_base/tiling_base.h" +#include "tiling_base/tiling_templates_registry.h" + +#ifndef CEIL_ALIGN +#define CEIL_ALIGN(val, align) ((((val) + (align) - 1) / (align)) * (align)) +#endif +#ifndef CEIL_DIV +#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b)) +#endif +namespace optiling { +const static uint64_t MOE_GATING_TOP_K_ASCEND_950_TILING_KEY = 10000; + +const static int64_t GROUP_SELECT_MODE_MAX = 0; +const static int64_t GROUP_SELECT_MODE_SUM = 1; +const static int64_t RENORM_NO = 0; +const static int64_t RENORM_L1 = 1; +const static int64_t NORM_TYPE_SOFTMAX = 0; +const static int64_t NORM_TYPE_SIGMOID = 1; +const static int64_t OUT_FLAG_FALSE = 0; +const static int64_t OUT_FLAG_TRUE = 1; +const static size_t X_INPUT_DIMS = 2; +const static size_t BIAS_INPUT_DIMS = 1; +const static size_t Y_OUTPUT_DIMS = 2; +const static size_t EXPERT_IDX_OUTPUY_DIMS = 2; +const static size_t OUT_OUTPUT_DIMS = 2; +const static int64_t MAX_EXPERT_COUNT = 2048; + +const static int64_t X_INPUT_INDEX = 0; +const static int64_t BIAS_INPUT_INDEX = 1; +const static int64_t Y_OUTPUT_INDEX = 0; +const static int64_t EXPERT_IDX_OUTPUT_INDEX = 1; +const static int64_t OUT_OUTPUT_INDEX = 2; +const static int64_t K_ATTR_INDEX = 0; +const static int64_t K_GROUP_ATTR_INDEX = 1; +const static int64_t GROUP_COUNT_ATTR_INDEX = 2; +const static int64_t GROUP_SELECT_MODE_ATTR_INDEX = 3; +const static int64_t RENORM_ATTR_INDEX = 4; +const static int64_t MRGSORT_SIZE = 4; +const static int64_t NORM_TYPE_ATTR_INDEX = 5; +const static int64_t OUT_FLAG_ATTR_INDEX = 6; +const static int64_t ROUTED_SCALING_FACTOR_ATTR_INDEX = 7; +const static int64_t EPS_ATTR_INDEX = 8; +const static int64_t DEFAULT_WORKSPACE_SIZE = static_cast(16 * 1024 * 1024); // 预留16M空间 + + +class MoeGatingTopKTilingRegbase : public Ops::Transformer::OpTiling::TilingBaseClass { +public: + explicit MoeGatingTopKTilingRegbase(gert::TilingContext *context) : Ops::Transformer::OpTiling::TilingBaseClass(context) + { + Reset(); + } + ~MoeGatingTopKTilingRegbase() override = default; + + void Reset(gert::TilingContext *context) override + { + TilingBaseClass::Reset(context); + Reset(); + } + +protected: + bool IsCapable() override + { + if (socVersion != platform_ascendc::SocVersion::ASCEND950) { + return false; + } + return true; + } + + ge::graphStatus GetPlatformInfo() override; + + ge::graphStatus GetShapeAttrsInfo() override; + + ge::graphStatus DoOpTiling() override; + + ge::graphStatus DoLibApiTiling() override; + + uint64_t GetTilingKey() const override; + + ge::graphStatus GetWorkspaceSize() override; + + ge::graphStatus PostTiling() override; + void Reset(); + +private: + ge::graphStatus CheckInputShape(); + ge::graphStatus CheckAttr(); + ge::graphStatus CheckOutShape(); + void CalTmpBufUbSize(); + void SplitRows(); + void Tiling4GatherOutComputeSplitK(); + + const gert::Shape *xShape_ = nullptr; + const gert::Shape *biasShape_ = nullptr; + const gert::Shape *yShape_ = nullptr; + const gert::Shape *expertIdxShape_ = nullptr; + const gert::Shape *outShape_ = nullptr; + + int64_t rows_; + int64_t expertCount_; + int64_t addBias_ = 0; + + int64_t k_; + int64_t kGroup_ = 1; + int64_t groupCount_ = 1; + int64_t groupSelectMode_ = GROUP_SELECT_MODE_MAX; + int64_t renorm_ = RENORM_NO; + int64_t normType_ = NORM_TYPE_SOFTMAX; + int64_t outFlag_ = OUT_FLAG_FALSE; + float routedScalingFactor_ = 1.0; + float eps_ = 1e-20f; + + int64_t inputDtypeSize_; + const char *opName_ = ""; + MoeGatingTopKRegbaseTilingData moeGatingTopKTilingData_; + platform_ascendc::SocVersion socVersion; +}; + +ge::graphStatus MoeGatingTopKTilingRegbase::CheckInputShape() +{ + size_t xDimNum = xShape_->GetDimNum(); + OP_CHECK_IF(xDimNum != X_INPUT_DIMS, + OP_LOGE(context_, "The dim number of x is: %zu, but should be %zu.", xDimNum, X_INPUT_DIMS), + return ge::GRAPH_FAILED); + + + rows_ = xShape_->GetDim(0); + expertCount_ = xShape_->GetDim(1); + moeGatingTopKTilingData_.set_rowCount(rows_); + moeGatingTopKTilingData_.set_expertCount(expertCount_); + OP_CHECK_IF( + expertCount_ > MAX_EXPERT_COUNT, + OP_LOGE(context_, "expert count is: %ld, but should not greater than %ld.", expertCount_, MAX_EXPERT_COUNT), + return ge::GRAPH_FAILED); + + if (biasShape_ != nullptr) { + addBias_ = 1; + size_t biasDimNum = biasShape_->GetDimNum(); + OP_CHECK_IF(biasDimNum != BIAS_INPUT_DIMS, + OP_LOGE(context_, "The number of bias dim is: %zu, but should be %zu.", biasDimNum, BIAS_INPUT_DIMS), + return ge::GRAPH_FAILED); + OP_CHECK_IF(biasShape_->GetDim(0) != expertCount_, + OP_LOGE(context_, "The first dim of bias is: %ld, but should be expert num: %ld.", + biasShape_->GetDim(0), expertCount_), + return ge::GRAPH_FAILED); + } + moeGatingTopKTilingData_.set_addBias(addBias_); + + OP_CHECK_IF(k_ > expertCount_, + OP_LOGE(context_, "k is: %ld, expert num is: %ld, k cannot be greater than expert num.", k_, expertCount_), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingRegbase::CheckAttr() +{ + OP_CHECK_IF(k_ <= 0, OP_LOGE(context_, "k is: %ld, but should be greater than 0.", k_), return ge::GRAPH_FAILED); + OP_CHECK_IF(kGroup_ <= 0, OP_LOGE(context_, "k_group is: %ld, but should be greater than 0.", kGroup_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(groupCount_ <= 0, OP_LOGE(context_, "group_count is: %ld, but should be greater than 0.", groupCount_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(expertCount_ % groupCount_ != 0, + OP_LOGE(context_, "expert num : %ld is not divisible by group_count: %ld", expertCount_, groupCount_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(kGroup_ > groupCount_, + OP_LOGE(context_, "k_group is: %ld, but should not greater than group_count: %ld", kGroup_, groupCount_), + return ge::GRAPH_FAILED); + OP_CHECK_IF(groupCount_ == expertCount_ && kGroup_ < k_, + OP_LOGE(context_, "k_group * group expert count is: %ld, but it must be greater than or equal to k: %ld.", + kGroup_, k_), + return ge::GRAPH_FAILED); + + if (kGroup_ == groupCount_ || groupCount_ == expertCount_) { + kGroup_ = 1; + groupCount_ = 1; + } + moeGatingTopKTilingData_.set_kGroup(kGroup_); + moeGatingTopKTilingData_.set_groupCount(groupCount_); + int64_t groupExpertCount = expertCount_ / groupCount_; + int64_t groupExpertCountAlign = CEIL_ALIGN(groupExpertCount, 32L); + moeGatingTopKTilingData_.set_perGroupExpertCount(expertCount_ / groupCount_); + moeGatingTopKTilingData_.set_perGroupExpertCountAlign(groupExpertCountAlign); + + OP_CHECK_IF(groupCount_ * groupExpertCountAlign > MAX_EXPERT_COUNT, + OP_LOGE(context_, "group count * group expert count align is: %ld, but should not greater than %ld.", + groupCount_ * groupExpertCountAlign, MAX_EXPERT_COUNT), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(kGroup_ * groupExpertCount < k_, + OP_LOGE(context_, "k_group * group expert count is: %ld, but it must be greater than or equal to k: %ld.", + kGroup_ * groupExpertCount, k_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(groupExpertCount < 1, + OP_LOGE(context_, "per group expert count is: %ld, but should be greater than 0.", groupExpertCount), + return ge::GRAPH_FAILED); + OP_CHECK_IF( + groupSelectMode_ != GROUP_SELECT_MODE_SUM && groupSelectMode_ != GROUP_SELECT_MODE_MAX, + OP_LOGE(context_, "group select mode is: %ld, but currently only support %ld and %ld.", groupSelectMode_, + GROUP_SELECT_MODE_SUM, GROUP_SELECT_MODE_MAX), + return ge::GRAPH_FAILED); + OP_CHECK_IF(groupSelectMode_ == GROUP_SELECT_MODE_SUM && groupExpertCount < 2, + OP_LOGE(context_, + "group expert count is: %ld, if group select mode is: %ld, group expert count should be greater than 1.", + groupExpertCount, groupSelectMode_), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(renorm_ != RENORM_NO, + OP_LOGE(context_, "renorm is: %ld, but currently only support %ld.", renorm_, RENORM_NO), + return ge::GRAPH_FAILED); + + OP_CHECK_IF(normType_ != NORM_TYPE_SOFTMAX && normType_ != NORM_TYPE_SIGMOID, + OP_LOGE(context_, "norm type is: %ld, but currently only support %ld and %ld.", normType_, + NORM_TYPE_SOFTMAX, NORM_TYPE_SIGMOID), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingRegbase::GetShapeAttrsInfo() +{ + opName_ = context_->GetNodeName(); + + auto xShapePtr = context_->GetInputShape(X_INPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, xShapePtr); + xShape_ = &xShapePtr->GetStorageShape(); + auto biasShapePtr = context_->GetOptionalInputShape(BIAS_INPUT_INDEX); + biasShape_ = biasShapePtr == nullptr ? nullptr : &biasShapePtr->GetStorageShape(); + + + auto yShapePtr = context_->GetOutputShape(Y_OUTPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, yShapePtr); + yShape_ = &yShapePtr->GetStorageShape(); + auto expertIdxPtr = context_->GetOutputShape(EXPERT_IDX_OUTPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, expertIdxPtr); + expertIdxShape_ = &expertIdxPtr->GetStorageShape(); + auto outPtr = context_->GetOutputShape(OUT_OUTPUT_INDEX); + if (outPtr != nullptr) { + outShape_ = &outPtr->GetStorageShape(); + } + + auto x = context_->GetInputDesc(X_INPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, x); + auto xDtype = x->GetDataType(); + OP_CHECK_IF( + (xDtype != ge::DataType::DT_FLOAT && xDtype != ge::DataType::DT_FLOAT16 && xDtype != ge::DataType::DT_BF16), + OP_LOGE(context_, "x dtype %s error, only supports float32, half, bf16. please check.", + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + + if (biasShapePtr != nullptr) { + auto biasDtype = context_->GetOptionalInputDesc(BIAS_INPUT_INDEX)->GetDataType(); + OP_CHECK_IF((biasDtype != xDtype), + OP_LOGE(context_, "bias dtype %s not equal x dtype %s, please check.", + ge::TypeUtils::DataTypeToSerialString(biasDtype).c_str(), + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + } + + auto yDesc = context_->GetOutputDesc(Y_OUTPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, yDesc); + auto yDtype = yDesc->GetDataType(); + OP_CHECK_IF((yDtype != xDtype), + OP_LOGE(context_, "y out dtype %s must be the same with x dtype %s.", + ge::TypeUtils::DataTypeToSerialString(yDtype).c_str(), + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + + auto expertIdDesc = context_->GetOutputDesc(EXPERT_IDX_OUTPUT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, expertIdDesc); + auto expertIdDtype = expertIdDesc->GetDataType(); + OP_CHECK_IF((expertIdDtype != ge::DataType::DT_INT32), + OP_LOGE(context_, "expertId out dtype %s error, only supports int32. please check.", + ge::TypeUtils::DataTypeToSerialString(expertIdDtype).c_str()), + return ge::GRAPH_FAILED); + + + auto attrs = context_->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context_, attrs); + + const int64_t *kPtr = attrs->GetAttrPointer(K_ATTR_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, kPtr); + k_ = *kPtr; + moeGatingTopKTilingData_.set_k(k_); + OP_LOGI(context_, "Attr k is: %ld ", k_); + + const int64_t *kGroupPtr = attrs->GetAttrPointer(K_GROUP_ATTR_INDEX); + if (kGroupPtr != nullptr) { + kGroup_ = *kGroupPtr; + } + OP_LOGI(context_, "Attr k_group is: %ld ", kGroup_); + + const int64_t *groupCountPtr = attrs->GetAttrPointer(GROUP_COUNT_ATTR_INDEX); + if (groupCountPtr != nullptr) { + groupCount_ = *groupCountPtr; + } + OP_LOGI(context_, "Attr group_count is: %ld ", groupCount_); + + const int64_t *groupSelectModePtr = attrs->GetAttrPointer(GROUP_SELECT_MODE_ATTR_INDEX); + if (groupSelectModePtr != nullptr) { + groupSelectMode_ = *groupSelectModePtr; + } + moeGatingTopKTilingData_.set_groupSelectMode(groupSelectMode_); + OP_LOGI(context_, "Attr group_select_mode is: %ld ", groupSelectMode_); + + const int64_t *renormPtr = attrs->GetAttrPointer(RENORM_ATTR_INDEX); + if (renormPtr != nullptr) { + renorm_ = *renormPtr; + } + moeGatingTopKTilingData_.set_renorm(renorm_); + OP_LOGI(context_, "Attr renorm is: %ld ", renorm_); + + const int64_t *normTypePtr = attrs->GetAttrPointer(NORM_TYPE_ATTR_INDEX); + if (normTypePtr != nullptr) { + normType_ = *normTypePtr; + } + moeGatingTopKTilingData_.set_normType(normType_); + OP_LOGI(context_, "Attr norm_type is: %ld ", normType_); + + const bool *outFlagPtr = attrs->GetAttrPointer(OUT_FLAG_ATTR_INDEX); + if (outFlagPtr != nullptr) { + outFlag_ = (*outFlagPtr) ? 1 : 0; + } + moeGatingTopKTilingData_.set_outFlag(outFlag_); + OP_LOGI(context_, "Attr out_flag is: %ld ", outFlag_); + + const float *routedScalingFactorPtr = attrs->GetAttrPointer(ROUTED_SCALING_FACTOR_ATTR_INDEX); + if (routedScalingFactorPtr != nullptr) { + routedScalingFactor_ = *routedScalingFactorPtr; + } + moeGatingTopKTilingData_.set_routedScalingFactor(routedScalingFactor_); + OP_LOGI(context_, "Attr routed_scaling_factor is: %f ", routedScalingFactor_); + + const float *epsPtr = attrs->GetAttrPointer(EPS_ATTR_INDEX); + if (epsPtr != nullptr) { + eps_ = *epsPtr; + } + moeGatingTopKTilingData_.set_eps(eps_); + OP_LOGI(context_, "Attr eps is: %f ", eps_); + + auto outDesc = context_->GetOutputDesc(OUT_OUTPUT_INDEX); + if (outFlag_ && outDesc != nullptr) { + auto outDtype = outDesc->GetDataType(); + OP_CHECK_IF((outDtype != ge::DataType::DT_FLOAT), + OP_LOGE(context_, "norm out dtype %s error, only supports float32. please check.", + ge::TypeUtils::DataTypeToSerialString(outDtype).c_str()), + return ge::GRAPH_FAILED); + } + + inputDtypeSize_ = static_cast(ge::GetSizeByDataType(context_->GetInputDesc(0)->GetDataType())); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingRegbase::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + OP_CHECK_IF(platformInfo == nullptr, OP_LOGE(context_, "fail to get platform info"), return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + aicoreParams_.blockDim = ascendcPlatform.GetCoreNumAiv(); + socVersion = ascendcPlatform.GetSocVersion(); + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + aicoreParams_.ubSize = ubSizePlatForm; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingRegbase::CheckOutShape() +{ + OP_CHECK_IF((yShape_->GetDimNum() != xShape_->GetDimNum()), + OP_LOGE(context_, "y out shape num %zu and x shape num %zu not equal, please check.", yShape_->GetDimNum(), + xShape_->GetDimNum()), + return ge::GRAPH_FAILED); + OP_CHECK_IF((expertIdxShape_->GetDimNum() != xShape_->GetDimNum()), + OP_LOGE(context_, "expertId out shape num %zu and x shape num %zu not equal, please check.", + expertIdxShape_->GetDimNum(), xShape_->GetDimNum()), + return ge::GRAPH_FAILED); + if (outShape_ != nullptr) { + OP_CHECK_IF((outShape_->GetDimNum() != xShape_->GetDimNum()), + OP_LOGE(context_, "norm out shape num %zu and x shape num %zu not equal, please check.", + outShape_->GetDimNum(), xShape_->GetDimNum()), + return ge::GRAPH_FAILED); + } + + OP_CHECK_IF((yShape_->GetDim(0) != xShape_->GetDim(0)), + OP_LOGE(context_, "y out dim[0] %ld not equal x dim[0] %ld, please check.", yShape_->GetDim(0), + xShape_->GetDim(0)), + return ge::GRAPH_FAILED); + OP_CHECK_IF((expertIdxShape_->GetDim(0) != xShape_->GetDim(0)), + OP_LOGE(context_, "expertId out dim[0] %ld not equal x dim[0] %ld, please check.", + expertIdxShape_->GetDim(0), xShape_->GetDim(0)), + return ge::GRAPH_FAILED); + if (outFlag_ && outShape_ != nullptr) { + OP_CHECK_IF((outShape_->GetDim(0) != xShape_->GetDim(0)), + OP_LOGE(context_, "norm out dim[0] %ld and x dim[0] %ld not equal, please check.", + outShape_->GetDim(0), outShape_->GetDim(0)), + return ge::GRAPH_FAILED); + } + + OP_CHECK_IF((yShape_->GetDim(1) != k_), + OP_LOGE(context_, "y dim[1] %ld not equal k %ld, please check.", yShape_->GetDim(1), k_), + return ge::GRAPH_FAILED); + OP_CHECK_IF((expertIdxShape_->GetDim(1) != k_), + OP_LOGE(context_, "expertId dim[1] %ld not equal k %ld, please check.", expertIdxShape_->GetDim(1), k_), + return ge::GRAPH_FAILED); + if (outFlag_ && outShape_ != nullptr) { + OP_CHECK_IF((outShape_->GetDim(1) != xShape_->GetDim(1)), + OP_LOGE(context_, "normOut dim[1] %ld and x dim[1] %ld not equal, please check.", outShape_->GetDim(1), + xShape_->GetDim(1)), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +void MoeGatingTopKTilingRegbase::CalTmpBufUbSize() { + std::vector shape_vec = {groupCount_ * moeGatingTopKTilingData_.get_perGroupExpertCountAlign()}; + ge::Shape softmaxShape(shape_vec); + + uint32_t softmaxTmpSize = AscendC::GetSoftMaxMaxTmpSize(softmaxShape, sizeof(float), true); + AscendC::SoftMaxTilingFunc(softmaxShape, sizeof(float), softmaxTmpSize, moeGatingTopKTilingData_.softmaxTilingData); +} + +void MoeGatingTopKTilingRegbase::SplitRows() +{ + int64_t perCoreRows = CEIL_DIV(rows_, static_cast(aicoreParams_.blockDim)); + int64_t needCoreNum = CEIL_DIV(rows_, perCoreRows); + if (perCoreRows == 0) { + OP_LOGE(context_, "perCoreRows can't be 0."); + return; + } + int64_t lastCoreRows = rows_ % perCoreRows == 0 ? perCoreRows : rows_ % perCoreRows; + moeGatingTopKTilingData_.set_needCoreNum(needCoreNum); + moeGatingTopKTilingData_.set_perCoreRowCount(perCoreRows); + moeGatingTopKTilingData_.set_lastCoreRowCount(lastCoreRows); + + int64_t vmsCount = 0; + if (kGroup_ > MRGSORT_SIZE) { + int64_t index = MRGSORT_SIZE; + while (index < kGroup_) { + index = index * MRGSORT_SIZE; + vmsCount++; + } + } + moeGatingTopKTilingData_.set_vmsCount(vmsCount); +} + +ge::graphStatus MoeGatingTopKTilingRegbase::DoOpTiling() +{ + auto ret = CheckInputShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckAttr(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckOutShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + CalTmpBufUbSize(); + SplitRows(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingRegbase::DoLibApiTiling() +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingRegbase::GetWorkspaceSize() +{ + + workspaceSize_ = DEFAULT_WORKSPACE_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKTilingRegbase::PostTiling() +{ + context_->SetBlockDim(moeGatingTopKTilingData_.get_needCoreNum()); + size_t *currentWorkspace = context_->GetWorkspaceSizes(1); + currentWorkspace[0] = workspaceSize_; + moeGatingTopKTilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), + context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(moeGatingTopKTilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +uint64_t MoeGatingTopKTilingRegbase::GetTilingKey() const +{ + return MOE_GATING_TOP_K_ASCEND_950_TILING_KEY; +} + +void MoeGatingTopKTilingRegbase::Reset() +{ + opName_ = nullptr; + return; +} + +REGISTER_OPS_TILING_TEMPLATE(MoeGatingTopK, MoeGatingTopKTilingRegbase, 1000); +} // namespace optiling diff --git a/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_tiling_base.cpp b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_tiling_base.cpp new file mode 100644 index 000000000..ae9f43542 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_host/moe_gating_top_k_tiling_base.cpp @@ -0,0 +1,38 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file moe_gating_top_k_tiling_base.cpp + * \brief + */ +#include "moe_gating_top_k_tiling.h" +#include "register/op_def_registry.h" +#include "tiling_base/tiling_base.h" +#include "tiling_base/tiling_templates_registry.h" +#include "tiling_base/error_log.h" +#include "kernel_tiling/kernel_tiling.h" + +namespace optiling { +static ge::graphStatus TilingForMoeGatingTopK(gert::TilingContext *context) +{ + return Ops::Transformer::OpTiling::TilingRegistry::GetInstance().DoTilingImpl(context); +} + +static ge::graphStatus TilingPrepareForMoeGatingTopK(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(MoeGatingTopK) + .Tiling(TilingForMoeGatingTopK) + .TilingParse(TilingPrepareForMoeGatingTopK); + +} // namespace optiling diff --git a/csrc/ascend/moe/moe_gating_top_k/op_kernel/common.h b/csrc/ascend/moe/moe_gating_top_k/op_kernel/common.h new file mode 100644 index 000000000..0847ee2b0 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_kernel/common.h @@ -0,0 +1,89 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file common.h + * \brief + */ +#ifndef MOE_GATING_TOP_K_COMMON_H +#define MOE_GATING_TOP_K_COMMON_H + +#include "kernel_operator.h" + +namespace MoeGatingTopK { +using namespace AscendC; +const float MIN_FP32 = *(float *)(&F32_NEG_INF); +constexpr int32_t FLOAT32_NEG_INF = 0xFF800000; // -inf -2139095040 +constexpr int64_t ONE_REPEAT_SORT_NUM = 32; +constexpr int64_t BLOCK_BYTES = 32; +constexpr int64_t REPEAT_BYTES = 256; +constexpr int64_t REPEAT_BLOCKS = 8; + +constexpr int32_t CONSTANT_TWO = 2; +constexpr int32_t CONSTANT_THREE = 3; +constexpr int32_t CONSTANT_FOUR = 4; +constexpr int32_t CONSTANT_EIGHT = 8; + +constexpr int64_t MERGE_LIST_TWO = 2; +constexpr int64_t MERGE_LIST_THREE = 3; +constexpr int64_t MERGE_LIST_FOUR = 4; + +constexpr int64_t MERGE_LIST_IDX_TWO = 2; +constexpr int64_t MERGE_LIST_IDX_THREE = 3; + +constexpr int64_t NORM_TYPE_SOFTMAX = 0; +constexpr int64_t NORM_TYPE_SIGMOID = 1; + +__aicore__ inline int64_t Ceil(int64_t a, int64_t b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes) +{ + if (bytes == 0) { + return 0; + } + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes; +} + +__aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes) +{ + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES; +} + +template +__aicore__ inline T Min(T a, T b) +{ + return a > b ? b : a; +} + +template +__aicore__ inline T Max(T a, T b) +{ + return a < b ? b : a; +} + +template +__aicore__ inline T1 CeilDiv(T1 x, T2 y) +{ + if (y != 0 && x != 0) { + const T1 quotient = x / y; + return (x % y != 0 && ((x ^ y) >= 0)) ? (quotient + 1) : quotient; + } + + return x; +} + +} // namespace MoeGatingTopK +#endif // MOE_GATING_TOP_K_COMMON_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k/op_kernel/error_log.h b/csrc/ascend/moe/moe_gating_top_k/op_kernel/error_log.h new file mode 100644 index 000000000..f48985a40 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_kernel/error_log.h @@ -0,0 +1,55 @@ +#ifndef OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_ +#define OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_ + +#include +#include "toolchain/slog.h" + +#define OP_LOGI(opname, ...) +#define OP_LOGW(opname, ...) \ + do { \ + printf("[WARN][%s] ", (opname), ##__VA_ARGS__); \ + printf("\n"); \ + } while (0) + +#define OP_LOGE_WITHOUT_REPORT(opname, ...) \ + do { \ + printf("[ERRORx][%s] ", (opname), ##__VA_ARGS__); \ + printf("\n"); \ + } while (0) + +#define OP_LOGE(opname, ...) \ + do { \ + printf("[ERROR][%s] ", (opname), ##__VA_ARGS__); \ + printf("\n"); \ + } while (0) + +#define OP_LOGD(opname, ...) + +namespace optiling { + +#define VECTOR_INNER_ERR_REPORT_TILIING(op_name, err_msg, ...) \ + do { \ + OP_LOGE_WITHOUT_REPORT(op_name, err_msg, ##__VA_ARGS__); \ + } while (0) + +#define OP_CHECK_IF(cond, log_func, expr) \ + do { \ + if (cond) { \ + log_func; \ + expr; \ + } \ + } while (0) + + + +#define OP_CHECK_NULL_WITH_CONTEXT(context, ptr) \ + do { \ + if ((ptr) == nullptr) { \ + OP_LOGE(context->GetNodeType(), "%s is null", #ptr); \ + return ge::GRAPH_FAILED; \ + } \ + } while (0) + +} // namespace optiling + +#endif // OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_ diff --git a/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k.cpp b/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k.cpp new file mode 100644 index 000000000..4390fee4a --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k.cpp @@ -0,0 +1,63 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k.cpp + * \brief + */ + +#include "moe_gating_top_k_e_k_fullload.h" +#include "moe_gating_top_k_without_group.h" +#include "moe_gating_top_k_generalized.h" +#include "error_log.h" + +#define TILING_KEY_PER_GROUP_COUNT_32 0 +#define TILING_KEY_WITHOUT_GROUP 1 +#define TILING_KEY_GENERALIZED 2 + +using namespace AscendC; +using namespace MoeGatingTopK; +extern "C" __global__ __aicore__ void moe_gating_top_k(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, + GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling) +{ + + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); + if (g_coreType == AIC) { + return; + } + + + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKTilingData, tilingData, tiling); + if (workspace == nullptr) { + return; + } + + GM_ADDR userWS = GetUserWorkspace(workspace); + if (userWS == nullptr) { + return; + } + + const MoeGatingTopKTilingData *__restrict t = &tilingData; + TPipe tPipe; + if (TILING_KEY_IS(TILING_KEY_PER_GROUP_COUNT_32)) { + MoeGatingTopKEKFullload op; + op.Init(x, bias, y, expertIdx, out, userWS, t, &tPipe); + op.Process(); + } else if (TILING_KEY_IS(TILING_KEY_WITHOUT_GROUP)) { + MoeGatingTopKWithoutGroup op; + op.Init(x, bias, y, expertIdx, out, userWS, t, &tPipe); + op.Process(); + } else if (TILING_KEY_IS(TILING_KEY_GENERALIZED)) { + MoeGatingTopKGenerlized op; + op.Init(x, bias, y, expertIdx, out, userWS, t, &tPipe); + op.Process(); + } + +} diff --git a/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k_apt.cpp b/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k_apt.cpp new file mode 100644 index 000000000..5abaf2baa --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k_apt.cpp @@ -0,0 +1,46 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_apt.cpp + * \brief + */ + +#include "arch35/moe_gating_top_k_regbase.h" +using namespace AscendC; +using namespace MoeGatingTopK; + +#define TILING_KEY_REGBASE 10000 + +extern "C" __global__ __aicore__ void moe_gating_top_k(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, + GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling) +{ + if (g_coreType == AIC) { + return; + } + + if (workspace == nullptr) { + return; + } + + GM_ADDR userWS = GetUserWorkspace(workspace); + if (userWS == nullptr) { + return; + } + + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKRegbaseTilingData, tiling_data_in, tiling); + const MoeGatingTopKRegbaseTilingData *__restrict tilingData = &tiling_data_in; + TPipe tPipe; + if (TILING_KEY_IS(TILING_KEY_REGBASE)) { + MoeGatingTopKRegbase op; + op.Init(x, bias, y, expertIdx, out, userWS, tilingData, &tPipe); + op.Process(); + } +} diff --git a/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k_e_k_fullload.h b/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k_e_k_fullload.h new file mode 100644 index 000000000..6891cc0e7 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k_e_k_fullload.h @@ -0,0 +1,404 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_e_k_fullload.h + * \brief + */ +#ifndef MOE_GATING_TOP_K_E_K_FULLLOAD_H +#define MOE_GATING_TOP_K_E_K_FULLLOAD_H +#include "kernel_operator.h" +#include "common.h" +namespace MoeGatingTopK { +using namespace AscendC; + +template +class MoeGatingTopKEKFullload { +public: + __aicore__ inline MoeGatingTopKEKFullload(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyInBias(); + __aicore__ inline void CopyInX(int64_t progress); + __aicore__ inline void ComputeX(); + __aicore__ inline void SortInGroup(); + __aicore__ inline void SelectTopKGroupIndex(); + __aicore__ inline void SelectTopKExpertIdx(); + __aicore__ inline void SelectTopKExpertScore(); + __aicore__ inline void CopyOut(int64_t progress); + +private: + TPipe *pipe_; + TQue xInQueue_; + TBuf biasInQueue_; + TQue yOutQueue_; + TQue expertIdxOutQueue_; + TQue outOutQueue_; + + TQue xBiasQueue_; + TQue xSigmoidQueue_; + TQue sigmoidTmpQueue_; + TQue sortedInGroupQueue_; + TQue sortedGroupQueue_; + TBuf calcTmpBuffer_; + + GlobalTensor xGm_; + GlobalTensor biasGm_; + GlobalTensor yGm_; + GlobalTensor expertIdxGm_; + GlobalTensor outGm_; + + int64_t blockIdx_; + int64_t perCoreRowCount_; + int64_t curCoreRowCount_; + int64_t expertCount_; + bool addBias_; + int64_t k_; + int64_t kGroup_; + int64_t groupCount_; + int64_t groupSelectMode_; + int64_t renorm_; + int64_t normType_; + int64_t outFlag_; + float routedScalingFactor_; + float eps_; + + int64_t expertCountAlign_; + int64_t kAlign_; + int64_t perGroupExpertCount_; + + const MoeGatingTopKTilingData *tilingData_; +}; + +template +__aicore__ inline void MoeGatingTopKEKFullload::CopyInBias() +{ + LocalTensor biasTensor = biasInQueue_.Get(); + DataCopyExtParams dataCopyParams{1, static_cast(expertCount_ * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if constexpr (IsSameType::value) { + DataCopyPad(biasTensor, biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + } else { + DataCopyPad(biasTensor[expertCountAlign_].ReinterpretCast(), biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(biasTensor, biasTensor[expertCountAlign_].ReinterpretCast(), RoundMode::CAST_NONE, expertCount_); + } +} + +template +__aicore__ inline void MoeGatingTopKEKFullload::CopyInX(int64_t row) +{ + LocalTensor xInLocalTensor = xInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(expertCount_ * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if constexpr (IsSameType::value) { + DataCopyPad(xInLocalTensor, xGm_[row * expertCount_], dataCopyParams, dataCopyPadParams); + } else { + DataCopyPad(xInLocalTensor[expertCountAlign_].ReinterpretCast(), xGm_[row * expertCount_], dataCopyParams, + dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(xInLocalTensor, xInLocalTensor[expertCountAlign_].ReinterpretCast(), RoundMode::CAST_NONE, + expertCount_); + } + + xInQueue_.EnQue(xInLocalTensor); +} + +template +__aicore__ inline void MoeGatingTopKEKFullload::ComputeX() +{ + LocalTensor xSigmoidTensor = xSigmoidQueue_.AllocTensor(); + LocalTensor xInLocalTensor = xInQueue_.DeQue(); + LocalTensor xBiasTensor = xBiasQueue_.AllocTensor(); + LocalTensor biasTensor = biasInQueue_.Get(); + LocalTensor sharedTmpBuffer = sigmoidTmpQueue_.AllocTensor(); // 临时空间可以复用 + Sigmoid(xSigmoidTensor, xInLocalTensor, sharedTmpBuffer, expertCount_); + PipeBarrier(); + if (addBias_) { + Add(xBiasTensor, xSigmoidTensor, biasTensor, expertCount_); + } else { + Adds(xBiasTensor, xSigmoidTensor, static_cast(0), expertCount_); + } + + xSigmoidQueue_.EnQue(xSigmoidTensor); + xBiasQueue_.EnQue(xBiasTensor); + xInQueue_.FreeTensor(xInLocalTensor); + sigmoidTmpQueue_.FreeTensor(sharedTmpBuffer); +} + +template +__aicore__ inline void MoeGatingTopKEKFullload::SortInGroup() +{ + LocalTensor xBiasTensor = xBiasQueue_.DeQue(); + LocalTensor sortedInGroupTensor = sortedInGroupQueue_.AllocTensor(); // 组内排序的结果, 后续归并需要 + LocalTensor indexTensor = calcTmpBuffer_.Get(); // 用于存储排序时的索引 + ArithProgression(indexTensor.ReinterpretCast(), 0, 1, expertCount_); // 生成组索引0 1 2 ...... + PipeBarrier(); + Sort32(sortedInGroupTensor, xBiasTensor, indexTensor, expertCount_ / ONE_REPEAT_SORT_NUM); // 组内排序 + sortedInGroupQueue_.EnQue(sortedInGroupTensor); + xBiasQueue_.FreeTensor(xBiasTensor); +} + +template +__aicore__ inline void MoeGatingTopKEKFullload::SelectTopKGroupIndex() +{ + LocalTensor sortedInGroupTensor = sortedInGroupQueue_.DeQue(); + LocalTensor indexTensor = calcTmpBuffer_.Get(); + LocalTensor top2ValueInGroupTensor = sigmoidTmpQueue_.AllocTensor(); // 这个临时空间可以复用 + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + + indexTensor.SetValue(0, static_cast(5)); // b0101 + indexTensor.SetValue(1, static_cast(0)); + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + + uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数 + GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = 8; + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = 8; + gatherMaskParams.src1RepeatStride = 0; + GatherMask(top2ValueInGroupTensor, sortedInGroupTensor, indexTensor, true, static_cast(64), + gatherMaskParams, rsvdCnt); + PipeBarrier(); + LocalTensor groupTop2SumTensor = top2ValueInGroupTensor; + PairReduceSum(groupTop2SumTensor, top2ValueInGroupTensor, 1, groupCount_ * 2, 1, 1, + 1); // 计算每个组内最大的两个数之和 + PipeBarrier(); + + LocalTensor groupIndexTensor = indexTensor; + ArithProgression(groupIndexTensor.ReinterpretCast(), 0, 1, groupCount_); // 生成组索引 + PipeBarrier(); + // 用最小值补到32个数 + int64_t duplicateNum = ONE_REPEAT_SORT_NUM - groupCount_; + if (duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX << groupCount_; + uint64_t mask[2] = {mask0, 0}; + Duplicate(groupTop2SumTensor, MIN_FP32, mask, 1, 1, 8); + PipeBarrier(); + } + // 排序，将kgroup选出来 + LocalTensor sortedGroupTensor = sortedGroupQueue_.AllocTensor(); + Sort32(sortedGroupTensor, groupTop2SumTensor, groupIndexTensor, 1); + + PipeBarrier(); + LocalTensor sortedGroupIndexTensor = indexTensor.ReinterpretCast(); + // 提取组序号 + uint8_t src1Pattern = 2; // 内置固定模式 + GatherMask(sortedGroupIndexTensor, sortedGroupTensor.template ReinterpretCast(), src1Pattern, false, + static_cast(0), {1, 1, 0, 0}, rsvdCnt); + + // 需要将组排序(这里是降序，所以下mrgsor的时候反着取，3、2、1、0) + Cast(sortedGroupTensor, sortedGroupIndexTensor, RoundMode::CAST_ROUND, kGroup_); + PipeBarrier(); + duplicateNum = ONE_REPEAT_SORT_NUM - kGroup_; + if (duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX << kGroup_; + uint64_t mask[2] = {mask0, 0}; + Duplicate(sortedGroupTensor, MIN_FP32, mask, 1, 1, 8); + PipeBarrier(); + } + Sort32(top2ValueInGroupTensor, sortedGroupTensor, sortedGroupIndexTensor.template ReinterpretCast(), 1); + PipeBarrier(); + src1Pattern = 1; + GatherMask(sortedGroupTensor, top2ValueInGroupTensor, src1Pattern, false, static_cast(0), {1, 1, 0, 0}, + rsvdCnt); + PipeBarrier(); + Cast(sortedGroupIndexTensor, sortedGroupTensor, RoundMode::CAST_ROUND, kGroup_); + + sortedGroupQueue_.FreeTensor(sortedGroupTensor); + sortedInGroupQueue_.EnQue(sortedInGroupTensor); + sigmoidTmpQueue_.FreeTensor(top2ValueInGroupTensor); +} + +template +__aicore__ inline void MoeGatingTopKEKFullload::SelectTopKExpertIdx() +{ + LocalTensor expertIdxTensor = expertIdxOutQueue_.AllocTensor(); + LocalTensor topKGroupIndexTensor = calcTmpBuffer_.Get(); + LocalTensor sortedInGroupTensor = sortedInGroupQueue_.DeQue(); + LocalTensor sortedExpertTensor = xInQueue_.AllocTensor(); + AscendC::MrgSort4Info params; + params.elementLengths[0] = k_; + params.elementLengths[1] = k_; + params.elementLengths[2] = k_; + params.elementLengths[3] = k_; + params.ifExhaustedSuspension = true; + params.validBit = 0b1111; + params.repeatTimes = 1; + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + int64_t listOffset1 = topKGroupIndexTensor.GetValue(3) * perGroupExpertCount_ * 2; + int64_t listOffset2 = topKGroupIndexTensor.GetValue(2) * perGroupExpertCount_ * 2; + int64_t listOffset3 = topKGroupIndexTensor.GetValue(1) * perGroupExpertCount_ * 2; + int64_t listOffset4 = topKGroupIndexTensor.GetValue(0) * perGroupExpertCount_ * 2; + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + AscendC::MrgSortSrcList srcList; + srcList.src1 = sortedInGroupTensor[listOffset1]; + srcList.src2 = sortedInGroupTensor[listOffset2]; + srcList.src3 = sortedInGroupTensor[listOffset3]; + srcList.src4 = sortedInGroupTensor[listOffset4]; + MrgSort(sortedExpertTensor, srcList, params); + PipeBarrier(); + uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数 + uint8_t src1Pattern = 2; // 内置固定模式 + GatherMask(expertIdxTensor, sortedExpertTensor.template ReinterpretCast(), src1Pattern, false, + static_cast(0), {1, 1, 0, 0}, rsvdCnt); + xInQueue_.FreeTensor(sortedExpertTensor); + expertIdxOutQueue_.EnQue(expertIdxTensor); + sortedInGroupQueue_.FreeTensor(sortedInGroupTensor); +} + +template +__aicore__ inline void MoeGatingTopKEKFullload::SelectTopKExpertScore() +{ + LocalTensor expertIdxTensor = expertIdxOutQueue_.DeQue(); + LocalTensor expertByteIdxTensor = calcTmpBuffer_.Get(); + LocalTensor xSigmoidTensor = xSigmoidQueue_.DeQue(); + LocalTensor yTensor = yOutQueue_.AllocTensor(); + LocalTensor yOutTensor; + if constexpr (!IsSameType::value) { + yOutTensor = yTensor.template ReinterpretCast()[kAlign_]; + } else { + yOutTensor = yTensor; + } + Muls(expertByteIdxTensor, expertIdxTensor, static_cast(sizeof(float)), k_); + PipeBarrier(); + Gather(yOutTensor, xSigmoidTensor, expertByteIdxTensor.template ReinterpretCast(), + static_cast(0), k_); + + LocalTensor calTensor = calcTmpBuffer_.Get(); + PipeBarrier(); + ReduceSum(calTensor, yOutTensor, xSigmoidTensor, k_); + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float sumValue = calTensor.GetValue(0) + eps_; + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Duplicate(calTensor, sumValue, k_); + PipeBarrier(); + Div(yOutTensor, yOutTensor, calTensor, k_); + PipeBarrier(); + Muls(yOutTensor, yOutTensor, routedScalingFactor_, k_); + + if constexpr (!IsSameType::value) { + PipeBarrier(); + Cast(yTensor, yOutTensor, RoundMode::CAST_RINT, k_); + } + + xSigmoidQueue_.EnQue(xSigmoidTensor); + expertIdxOutQueue_.EnQue(expertIdxTensor); + yOutQueue_.EnQue(yTensor); +} + +template +__aicore__ inline void MoeGatingTopKEKFullload::CopyOut(int64_t row) +{ + LocalTensor yOutTensor = yOutQueue_.DeQue(); + LocalTensor expertIdxTensor = expertIdxOutQueue_.DeQue(); + LocalTensor xSigmoidTensor = xSigmoidQueue_.DeQue(); + DataCopyExtParams dataCopyParams{1, static_cast(k_ * sizeof(T)), 0, 0, 0}; + DataCopyPad(yGm_[row * k_], yOutTensor, dataCopyParams); + dataCopyParams.blockLen = k_ * sizeof(int32_t); + DataCopyPad(expertIdxGm_[row * k_], expertIdxTensor, dataCopyParams); + xSigmoidQueue_.FreeTensor(xSigmoidTensor); + expertIdxOutQueue_.FreeTensor(expertIdxTensor); + yOutQueue_.FreeTensor(yOutTensor); +} + +template +__aicore__ inline void MoeGatingTopKEKFullload::Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, + GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKTilingData *tilingData, TPipe *tPipe) +{ + tilingData_ = tilingData; + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + perCoreRowCount_ = tilingData_->perCoreRowCount; + if (blockIdx_ == GetBlockNum() - 1) { + curCoreRowCount_ = tilingData_->lastCoreRowCount; + } else { + curCoreRowCount_ = tilingData_->perCoreRowCount; + } + expertCount_ = tilingData_->expertCount; + addBias_ = tilingData_->addBias == 1; + k_ = tilingData_->k; + kGroup_ = tilingData_->kGroup; + groupCount_ = tilingData_->groupCount; + perGroupExpertCount_ = tilingData_->perGroupExpertCount; + routedScalingFactor_ = tilingData_->routedScalingFactor; + eps_ = tilingData_->eps; + + expertCountAlign_ = Align(expertCount_, sizeof(float)); + kAlign_ = Align(expertCount_, sizeof(float)); + + // init input gm buf + xGm_.SetGlobalBuffer((__gm__ T *)x + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_); + biasGm_.SetGlobalBuffer((__gm__ T *)bias, expertCount_); + + // init output gm buf + yGm_.SetGlobalBuffer((__gm__ T *)y + perCoreRowCount_ * k_ * blockIdx_, k_); + expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx + perCoreRowCount_ * k_ * blockIdx_, k_); + outGm_.SetGlobalBuffer((__gm__ T *)out + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_); + + // init que + pipe_->InitBuffer(xInQueue_, 2, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(biasInQueue_, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + + pipe_->InitBuffer(xSigmoidQueue_, 1, AlignBytes(expertCount_, sizeof(float))); + pipe_->InitBuffer(xBiasQueue_, 2, AlignBytes(expertCount_, sizeof(float))); + + pipe_->InitBuffer(yOutQueue_, 2, kAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(expertIdxOutQueue_, 2, AlignBytes(k_, sizeof(int32_t))); + pipe_->InitBuffer(outOutQueue_, 2, AlignBytes(expertCount_, sizeof(float))); + + pipe_->InitBuffer(sigmoidTmpQueue_, 2, AlignBytes(expertCount_, sizeof(float))); + pipe_->InitBuffer(sortedInGroupQueue_, 2, AlignBytes(expertCount_, sizeof(float)) * 2); + pipe_->InitBuffer(sortedGroupQueue_, 2, + (groupCount_ + ONE_REPEAT_SORT_NUM - 1) / ONE_REPEAT_SORT_NUM * ONE_REPEAT_SORT_NUM * + sizeof(float) * 2); + + pipe_->InitBuffer(calcTmpBuffer_, tilingData_->calTmpBufUbSize); +} + +template +__aicore__ inline void MoeGatingTopKEKFullload::Process() +{ + CopyInBias(); + for (int64_t row = 0; row < curCoreRowCount_; row++) { + CopyInX(row); + ComputeX(); + SortInGroup(); + SelectTopKGroupIndex(); + SelectTopKExpertIdx(); + SelectTopKExpertScore(); + CopyOut(row); + } +} +} // namespace MoeGatingTopK +#endif // MOE_GATING_TOP_K_E_K_FULLLOAD_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k_generalized.h b/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k_generalized.h new file mode 100644 index 000000000..b1cc14e6f --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k_generalized.h @@ -0,0 +1,669 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_generalized.h + * \brief + */ +#ifndef MOE_GATING_TOP_K_E_K_GENERALIZED_H +#define MOE_GATING_TOP_K_E_K_GENERALIZED_H +#include "kernel_operator.h" +#include "common.h" +#include "kernel_utils.h" +namespace MoeGatingTopK { +using namespace AscendC; + +template +class MoeGatingTopKGenerlized { +public: + __aicore__ inline MoeGatingTopKGenerlized(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyInBiasAndInitExpertId(); + __aicore__ inline void CopyInX(int64_t progress); + __aicore__ inline void ComputeX(); + __aicore__ inline void CopuOutXNorm(int64_t row); + __aicore__ inline void SortInGroup(); + __aicore__ inline void SelectTopKGroupIndex(); + __aicore__ inline void SelectTopKExpertIdx(); + __aicore__ inline void SelectTopKExpertScore(); + __aicore__ inline void CumputeActualTopKExpertId(); + __aicore__ inline void CopyOut(int64_t row); + +private: + TPipe *pipe_; + TQue xInQueue_; + TQue yOutQueue_; + TQue expertIdxOutQueue_; + TQue outOutQueue_; + + TBuf biasBuf_; // Store input bias + TBuf expertIdBuf_; // Expert ID + TBuf xNormWithBiasBuf_; // Store value after adding bias + TBuf xNormBuf_; // Store value after computing sigmoid or softmax + TBuf sortedInGroupBuf_; // Store sorted results within groups + TBuf topKExpertIdBuf_; + TBuf sortedGroupIndexBuf_; + TBuf calcTmpBuf_; + + GlobalTensor xGm_; + GlobalTensor biasGm_; + GlobalTensor yGm_; + GlobalTensor expertIdxGm_; + GlobalTensor outGm_; + + int64_t blockIdx_ = 0; + int64_t perCoreRowCount_ = 0; + int64_t curCoreRowCount_ = 0; + int64_t expertCount_ = 0; + bool addBias_ = false; + int64_t k_ = 0; + int64_t kGroup_ = 0; + int64_t groupCount_ = 0; + int64_t groupCountAlign_ = 0; + int64_t perGroupExpertCount_ = 0; + int64_t perGroupExpertCountAlign_ = 0; + int64_t groupSelectMode_ = 0; + int64_t renorm_ = 0; + int64_t normType_ = 0; + int64_t outFlag_ = 0; + + int64_t expertCountAlign_ = 0; + int64_t kAlign_ = 0; + bool isAlign_ = false; + + const MoeGatingTopKTilingData *tilingData_; +}; + +template +__aicore__ inline void MoeGatingTopKGenerlized::CopyInBiasAndInitExpertId() +{ + LocalTensor biasTensor = biasBuf_.Get(); + LocalTensor expertIdTensor = expertIdBuf_.Get(); + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = groupCount_; + dataCopyParams.blockLen = perGroupExpertCount_ * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(T) / BLOCK_BYTES; + + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if (addBias_) { + if constexpr (IsSameType::value) { + DataCopyPad(biasTensor, biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + } else { + DataCopyPad(biasTensor[expertCountAlign_].ReinterpretCast(), biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(biasTensor, biasTensor[expertCountAlign_].ReinterpretCast(), RoundMode::CAST_NONE, + expertCountAlign_); + PipeBarrier(); + } + + if (!isAlign_) { + int64_t duplicateNum = perGroupExpertCount_ % ONE_REPEAT_SORT_NUM; + int duplicateIndex = perGroupExpertCount_ - duplicateNum; + if (duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(biasTensor.ReinterpretCast()[duplicateIndex], FLOAT32_NEG_INF, mask, groupCount_, 1, + perGroupExpertCountAlign_ * sizeof(float) / BLOCK_BYTES); + } + } + } + ArithProgression(expertIdTensor, static_cast(0), static_cast(1), expertCountAlign_); +} + +template +__aicore__ inline void MoeGatingTopKGenerlized::CopyInX(int64_t row) +{ + LocalTensor xInLocalTensor = xInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = groupCount_; + dataCopyParams.blockLen = perGroupExpertCount_ * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(T) / BLOCK_BYTES; + + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if constexpr (IsSameType::value) { + DataCopyPad(xInLocalTensor, xGm_[row * expertCount_], dataCopyParams, dataCopyPadParams); + } else { + DataCopyPad(xInLocalTensor[expertCountAlign_].ReinterpretCast(), xGm_[row * expertCount_], dataCopyParams, + dataCopyPadParams); + } + xInQueue_.EnQue(xInLocalTensor); +} + +template +__aicore__ inline void MoeGatingTopKGenerlized::ComputeX() +{ + LocalTensor xNormTensor = xNormBuf_.Get(); + LocalTensor xInLocalTensor = xInQueue_.DeQue(); + LocalTensor xNormWithBiasTensor = xNormWithBiasBuf_.Get(); + LocalTensor biasTensor = biasBuf_.Get(); + + if constexpr (!IsSameType::value) { + Cast(xInLocalTensor, xInLocalTensor[expertCountAlign_].ReinterpretCast(), RoundMode::CAST_NONE, + expertCountAlign_); + PipeBarrier(); + } + + int64_t duplicateNum = perGroupExpertCount_ % ONE_REPEAT_SORT_NUM; + int duplicateIndex = perGroupExpertCount_ - duplicateNum; + if (!isAlign_ && duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(xInLocalTensor.ReinterpretCast()[duplicateIndex], FLOAT32_NEG_INF, mask, groupCount_, 1, + (perGroupExpertCountAlign_ * sizeof(float)) / BLOCK_BYTES); + PipeBarrier(); + } + if (normType_ == 1) { // sigmoid + LocalTensor calcNormTmpTensor = calcTmpBuf_.Get(); + Sigmoid(xNormTensor, xInLocalTensor, calcNormTmpTensor, expertCountAlign_); + PipeBarrier(); + } + else if (normType_ == 0) { // softmax + LocalTensor reduceValueTensor = calcTmpBuf_.Get(); + LocalTensor calcTmp = calcTmpBuf_.Get()[BLOCK_BYTES]; + ReduceMax(reduceValueTensor, xInLocalTensor, calcTmp, expertCountAlign_); + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float maxValue = reduceValueTensor.GetValue(0); + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Adds(xNormTensor, xInLocalTensor, -maxValue, expertCountAlign_); + PipeBarrier(); + Exp(xNormTensor, xNormTensor, expertCountAlign_); + PipeBarrier(); + ReduceSum(reduceValueTensor, xNormTensor, calcTmp, expertCountAlign_); + eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float sumValue = reduceValueTensor.GetValue(0); + eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Muls(xNormTensor, xNormTensor, 1.0f / sumValue, expertCountAlign_); + PipeBarrier(); + } + if (addBias_) { + Add(xNormWithBiasTensor, xNormTensor, biasTensor, expertCountAlign_); + } else { + DataCopy(xNormWithBiasTensor, xNormTensor, expertCountAlign_); + } + + if (!isAlign_ && duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + PipeBarrier(); + Duplicate(xNormWithBiasTensor.ReinterpretCast()[duplicateIndex], + FLOAT32_NEG_INF, // MIN_FP32, + mask, groupCount_, 1, perGroupExpertCountAlign_ * sizeof(float) / BLOCK_BYTES); + } + xInQueue_.FreeTensor(xInLocalTensor); +} + +template +__aicore__ inline void MoeGatingTopKGenerlized::CopuOutXNorm(int64_t row) +{ + LocalTensor outOutTensor = outOutQueue_.AllocTensor(); + LocalTensor xNormTensor = xNormBuf_.Get(); + DataCopy(outOutTensor, xNormTensor, expertCountAlign_); + outOutQueue_.EnQue(outOutTensor); + outOutTensor = outOutQueue_.DeQue(); + DataCopyExtParams dataCopyParams{ + static_cast(groupCount_), static_cast(perGroupExpertCount_ * sizeof(float)), + static_cast((perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(float) / BLOCK_BYTES), 0, 0}; + DataCopyPad(outGm_[row * expertCount_], outOutTensor, dataCopyParams); + outOutQueue_.FreeTensor(outOutTensor); +} + +template +__aicore__ inline void MoeGatingTopKGenerlized::SortInGroup() +{ + LocalTensor xNormWithBiasTensor = xNormWithBiasBuf_.Get(); + LocalTensor expertIdTensor = expertIdBuf_.Get(); + LocalTensor sortedInGroupTensor = sortedInGroupBuf_.Get(); + LocalTensor tmpLocal = calcTmpBuf_.Get(); + if (perGroupExpertCountAlign_ == ONE_REPEAT_SORT_NUM) { + PipeBarrier(); + Sort32(sortedInGroupTensor, xNormWithBiasTensor, expertIdTensor, groupCount_); + } else { + for (int64_t group = 0; group < groupCount_; group++) { + PipeBarrier(); + Sort(sortedInGroupTensor[group * perGroupExpertCountAlign_ * CONSTANT_TWO], + xNormWithBiasTensor[group * perGroupExpertCountAlign_], + expertIdTensor[group * perGroupExpertCountAlign_], tmpLocal, + perGroupExpertCountAlign_ / ONE_REPEAT_SORT_NUM); + } + } +} + +template +__aicore__ inline void MoeGatingTopKGenerlized::SelectTopKGroupIndex() +{ + LocalTensor sortedInGroupTensor = sortedInGroupBuf_.Get(); + LocalTensor valueSelectedFromGroupTensor = calcTmpBuf_.GetWithOffset(groupCountAlign_ * 2, 0); + LocalTensor maskTensor = + calcTmpBuf_.GetWithOffset(groupCountAlign_, groupCountAlign_ * 2 * sizeof(float)); + LocalTensor topValueInGroupTensor = + calcTmpBuf_.GetWithOffset(groupCountAlign_, groupCountAlign_ * 3 * sizeof(float)); + LocalTensor groupIndex = + calcTmpBuf_.GetWithOffset(groupCountAlign_, groupCountAlign_ * 4 * sizeof(float)); + LocalTensor sortedTopValue = + calcTmpBuf_.GetWithOffset(groupCountAlign_ * 2, groupCountAlign_ * 5 * sizeof(float)); + LocalTensor sortTmp = + calcTmpBuf_.GetWithOffset(groupCountAlign_ * 2, groupCountAlign_ * 7 * sizeof(float)); + + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + + uint64_t rsvdCnt = 0; // Used to store the number of elements retained after filtering + PipeBarrier(); + if (groupSelectMode_ == 1) { // top2 sum + // Extract the first two elements of each group + maskTensor.SetValue(0, static_cast(5)); // b0101 + maskTensor.SetValue(1, static_cast(0)); + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + + GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = groupCount_; + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = + Ceil(perGroupExpertCountAlign_ * (sizeof(float) + sizeof(uint32_t)), BLOCK_BYTES); + gatherMaskParams.src1RepeatStride = 0; + GatherMask(valueSelectedFromGroupTensor, sortedInGroupTensor, maskTensor, true, + static_cast(ONE_REPEAT_SORT_NUM * CONSTANT_TWO), gatherMaskParams, rsvdCnt); + PipeBarrier(); + + // Calculate the sum of the first two numbers in each group + PairReduceSum(topValueInGroupTensor, valueSelectedFromGroupTensor, + Ceil(groupCount_ * sizeof(float) * 2, REPEAT_BYTES), REPEAT_BYTES / sizeof(float), 1, 1, + CONSTANT_EIGHT); // Calculate the sum of the two largest numbers in each group + } else { + maskTensor.SetValue(0, static_cast(1)); // b0101 + maskTensor.SetValue(1, static_cast(0)); + + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + uint64_t rsvdCnt = 0; // Used to store the number of elements retained after filtering + GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = groupCount_; + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = Ceil(perGroupExpertCountAlign_ * (sizeof(float) + sizeof(uint32_t)), 32); + gatherMaskParams.src1RepeatStride = 0; + GatherMask(topValueInGroupTensor, sortedInGroupTensor, maskTensor, true, + static_cast(ONE_REPEAT_SORT_NUM * CONSTANT_TWO), gatherMaskParams, rsvdCnt); + } + + PipeBarrier(); + // Generate group indices + ArithProgression(groupIndex.ReinterpretCast(), static_cast(0), static_cast(1), + groupCount_); // Generate group indices + PipeBarrier(); + + int64_t duplicateNum = groupCount_ % ONE_REPEAT_SORT_NUM; + int duplicateIndex = groupCount_ - duplicateNum; + if (duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(topValueInGroupTensor.ReinterpretCast()[duplicateIndex], FLOAT32_NEG_INF, mask, 1, 1, + REPEAT_BLOCKS); + PipeBarrier(); + } + PipeBarrier(); + + // Sort + Sort(sortedTopValue, topValueInGroupTensor, groupIndex, sortTmp, Ceil(groupCount_, 32)); + PipeBarrier(); + + // Extract group indices + uint8_t src1Pattern = 2; // Built-in fixed pattern + GatherMask(groupIndex, sortedTopValue.template ReinterpretCast(), src1Pattern, false, + static_cast(0), + {1, static_cast(Ceil(kGroup_ * sizeof(float) * CONSTANT_TWO, 256)), REPEAT_BLOCKS, 0}, rsvdCnt); + PipeBarrier(); + duplicateNum = kGroup_ % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + duplicateIndex = kGroup_ - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + PipeBarrier(); + Duplicate(groupIndex.ReinterpretCast()[duplicateIndex], FLOAT32_NEG_INF, mask, 1, 1, REPEAT_BLOCKS); + } + + // Sort the selected group indices in descending order + LocalTensor sortedGroupIndex = sortedGroupIndexBuf_.Get(); + PipeBarrier(); + Sort(sortedGroupIndex, groupIndex.ReinterpretCast(), groupIndex, sortTmp, Ceil(kGroup_, 32)); +} + +template +__aicore__ inline void MoeGatingTopKGenerlized::SelectTopKExpertIdx() +{ + LocalTensor sortedInGroupTensor = sortedInGroupBuf_.Get(); + LocalTensor sortedGroupIndex = sortedGroupIndexBuf_.Get(); + LocalTensor topKExpertId = topKExpertIdBuf_.Get(); + LocalTensor mrgSort0Tensor = calcTmpBuf_.Get(); + + uint32_t offset[CONSTANT_FOUR] = {0, 0, 0, 0}; + uint16_t lenArr[CONSTANT_FOUR] = { + static_cast(perGroupExpertCount_), static_cast(perGroupExpertCount_), + static_cast(perGroupExpertCount_), static_cast(perGroupExpertCount_)}; + MrgSort4Info params{lenArr, false, 0b1111, 1}; + MrgSortSrcList srcList; + + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + + for (int32_t i = kGroup_ - 1; i >= 0; i -= CONSTANT_FOUR) { + int64_t mrgLen = Min(i + 1, CONSTANT_FOUR); + if (mrgLen > 1) { + if (mrgLen == MERGE_LIST_FOUR) { + offset[0] = sortedGroupIndex.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + offset[1] = sortedGroupIndex.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2; + offset[2] = sortedGroupIndex.GetValue((i - 2) * 2) * perGroupExpertCountAlign_ * 2; + offset[3] = sortedGroupIndex.GetValue((i - 3) * 2) * perGroupExpertCountAlign_ * 2; + } else if (mrgLen == MERGE_LIST_THREE) { + offset[0] = sortedGroupIndex.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + offset[1] = sortedGroupIndex.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2; + offset[2] = sortedGroupIndex.GetValue((i - 2) * 2) * perGroupExpertCountAlign_ * 2; + offset[3] = 0; + params.elementLengths[3] = 0; + params.validBit = 0b111; + } else { + offset[0] = sortedGroupIndex.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + offset[1] = sortedGroupIndex.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2; + offset[2] = 0; + offset[3] = 0; + params.elementLengths[2] = 0; + params.elementLengths[3] = 0; + params.validBit = 0b11; + } + + srcList.src1 = sortedInGroupTensor[offset[0]]; + srcList.src2 = sortedInGroupTensor[offset[1]]; + srcList.src3 = sortedInGroupTensor[offset[2]]; + srcList.src4 = sortedInGroupTensor[offset[3]]; + + PipeBarrier(); + MrgSort(mrgSort0Tensor[(kGroup_ - 1 - i) * perGroupExpertCountAlign_ * 2], srcList, params); + } else { + offset[0] = sortedGroupIndex.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + PipeBarrier(); + DataCopy(mrgSort0Tensor[(kGroup_ - 1 - i) * perGroupExpertCountAlign_ * 2], sortedInGroupTensor[offset[0]], + perGroupExpertCountAlign_ * 2); + } + } + int32_t baseLoop = 4; + LocalTensor srcTensor = mrgSort0Tensor; + LocalTensor dstTensor = mrgSort0Tensor; + for (int i = 0; i < tilingData_->vmsCount; i++) { + if (i % 2 == 0) { + srcTensor = mrgSort0Tensor; + dstTensor = sortedInGroupTensor; + } else { + srcTensor = sortedInGroupTensor; + dstTensor = mrgSort0Tensor; + } + + int32_t nextBaseRow = baseLoop * MERGE_LIST_FOUR; + int32_t quotient = kGroup_ / nextBaseRow; + int32_t remainder = kGroup_ - quotient * nextBaseRow; + if (quotient > 0) { + MrgSort4Info params; + MrgSortSrcList srcList; + params.ifExhaustedSuspension = false; + params.elementLengths[0] = perGroupExpertCount_ * baseLoop; + params.elementLengths[1] = perGroupExpertCount_ * baseLoop; + params.elementLengths[2] = perGroupExpertCount_ * baseLoop; + params.elementLengths[3] = perGroupExpertCount_ * baseLoop; + params.validBit = 0b1111; + params.repeatTimes = 1; + for (int j = 0; j < quotient; j++) { + srcList.src1 = srcTensor[perGroupExpertCountAlign_ * baseLoop * 8 * j]; + srcList.src2 = srcTensor[perGroupExpertCountAlign_ * baseLoop * (8 * j + 2)]; + srcList.src3 = srcTensor[perGroupExpertCountAlign_ * baseLoop * (8 * j + 4)]; + srcList.src4 = srcTensor[perGroupExpertCountAlign_ * baseLoop * (8 * j + 6)]; + PipeBarrier(); + MrgSort(dstTensor[perGroupExpertCountAlign_ * baseLoop * 8 * j], srcList, params); + } + } + + if (remainder > 0) { + int32_t baseOffset = quotient * nextBaseRow * perGroupExpertCountAlign_ * 2; + int32_t mrgLen = CeilDiv(remainder, baseLoop); + int32_t tailRow = remainder - (mrgLen - 1) * baseLoop; + if (mrgLen > 1) { + MrgSort4Info params; + MrgSortSrcList srcList; + params.repeatTimes = 1; + params.ifExhaustedSuspension = false; + params.elementLengths[0] = perGroupExpertCount_ * baseLoop; + params.elementLengths[1] = perGroupExpertCount_ * baseLoop; + params.elementLengths[2] = perGroupExpertCount_ * baseLoop; + params.elementLengths[3] = perGroupExpertCount_ * baseLoop; + srcList.src1 = srcTensor[baseOffset]; + srcList.src2 = srcTensor[baseOffset + perGroupExpertCountAlign_ * baseLoop * 2]; + if (mrgLen == MERGE_LIST_FOUR) { + srcList.src3 = srcTensor[baseOffset + perGroupExpertCountAlign_ * baseLoop * 2 * 2]; + srcList.src4 = srcTensor[baseOffset + perGroupExpertCountAlign_ * baseLoop * 2 * 3]; + params.elementLengths[3] = perGroupExpertCount_ * tailRow; + params.validBit = 0b1111; + } else if (mrgLen == MERGE_LIST_THREE) { + srcList.src3 = srcTensor[baseOffset + perGroupExpertCountAlign_ * baseLoop * 2 * 2]; + params.elementLengths[2] = perGroupExpertCount_ * tailRow; + params.elementLengths[3] = 0; + params.validBit = 0b111; + } else { + params.elementLengths[1] = perGroupExpertCount_ * tailRow; + params.elementLengths[2] = 0; + params.elementLengths[3] = 0; + params.validBit = 0b11; + } + PipeBarrier(); + MrgSort(dstTensor[baseOffset], srcList, params); + } else { + PipeBarrier(); + DataCopy(dstTensor[baseOffset], srcTensor[baseOffset], tailRow * perGroupExpertCountAlign_ * 2); + } + } + baseLoop = nextBaseRow; + } + + GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = Ceil(k_ * sizeof(float) * 2, REPEAT_BYTES); + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = REPEAT_BLOCKS; + gatherMaskParams.src1RepeatStride = 0; + + uint64_t rsvdCnt = 0; // Used to store the number of elements retained after filtering + uint8_t src1Pattern = 2; // Built-in fixed pattern + PipeBarrier(); + GatherMask(topKExpertId, dstTensor.template ReinterpretCast(), src1Pattern, false, + static_cast(0), gatherMaskParams, rsvdCnt); +} + +template +__aicore__ inline void MoeGatingTopKGenerlized::SelectTopKExpertScore() +{ + LocalTensor xNormTensor = xNormBuf_.Get(); + LocalTensor yOutTensor = yOutQueue_.AllocTensor(); + LocalTensor topKExpertId = topKExpertIdBuf_.Get(); + LocalTensor topKExpertIdWithByte = calcTmpBuf_.Get(); + PipeBarrier(); + Muls(topKExpertIdWithByte, topKExpertId, static_cast(sizeof(float)), k_); + PipeBarrier(); + Gather(yOutTensor, xNormTensor, topKExpertIdWithByte.template ReinterpretCast(), static_cast(0), + k_); + bool needRenorm = (normType_ == 1 ) || // Case 1: sigmoid + renorm + (normType_ == 0 && renorm_ == 1); // Case 3: softmax + renorm + if (needRenorm) { + LocalTensor maxValueTensor = calcTmpBuf_.Get(); + LocalTensor tmpTensor = calcTmpBuf_.Get()[32]; + PipeBarrier(); + ReduceSum(maxValueTensor, yOutTensor, tmpTensor, k_); + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float sumValue = maxValueTensor.GetValue(0) + tilingData_->eps; + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Duplicate(tmpTensor, sumValue, k_); + PipeBarrier(); + Div(yOutTensor, yOutTensor, tmpTensor, k_); + } + PipeBarrier(); + Muls(yOutTensor, yOutTensor, tilingData_->routedScalingFactor, k_); + + if constexpr (!IsSameType::value) { + PipeBarrier(); + Cast(yOutTensor.ReinterpretCast(), yOutTensor, RoundMode::CAST_RINT, k_); + } + + yOutQueue_.EnQue(yOutTensor); +} + +template +__aicore__ inline void MoeGatingTopKGenerlized::CumputeActualTopKExpertId() +{ + LocalTensor expertIdxOut = expertIdxOutQueue_.AllocTensor(); + LocalTensor topKExpertId = topKExpertIdBuf_.Get(); + LocalTensor topKExpertIdFp32 = calcTmpBuf_.Get(); + + PipeBarrier(); + Cast(topKExpertIdFp32, topKExpertId, RoundMode::CAST_ROUND, k_); + PipeBarrier(); + Muls(topKExpertIdFp32, topKExpertIdFp32, 1.0f / (float)perGroupExpertCountAlign_, k_); + PipeBarrier(); + Cast(expertIdxOut, topKExpertIdFp32, RoundMode::CAST_TRUNC, k_); + PipeBarrier(); + Muls(expertIdxOut, expertIdxOut, static_cast(perGroupExpertCountAlign_ - perGroupExpertCount_), k_); + PipeBarrier(); + Sub(expertIdxOut, topKExpertId, expertIdxOut, k_); + expertIdxOutQueue_.EnQue(expertIdxOut); +} + +template +__aicore__ inline void MoeGatingTopKGenerlized::CopyOut(int64_t row) +{ + LocalTensor yOutTensor = yOutQueue_.DeQue(); + LocalTensor expertIdxOut = expertIdxOutQueue_.DeQue(); + DataCopyExtParams dataCopyParams{1, static_cast(k_ * sizeof(T)), 0, 0, 0}; + DataCopyPad(yGm_[row * k_], yOutTensor, dataCopyParams); + dataCopyParams.blockLen = k_ * sizeof(int32_t); + DataCopyPad(expertIdxGm_[row * k_], expertIdxOut, dataCopyParams); + yOutQueue_.FreeTensor(yOutTensor); + expertIdxOutQueue_.FreeTensor(expertIdxOut); +} + +template +__aicore__ inline void MoeGatingTopKGenerlized::Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, + GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKTilingData *tilingData, TPipe *tPipe) +{ + tilingData_ = tilingData; + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + perCoreRowCount_ = tilingData_->perCoreRowCount; + if (blockIdx_ == GetBlockNum() - 1) { + curCoreRowCount_ = tilingData_->lastCoreRowCount; + } else { + curCoreRowCount_ = tilingData_->perCoreRowCount; + } + expertCount_ = tilingData_->expertCount; + addBias_ = tilingData_->addBias == 1; + k_ = tilingData_->k; + kGroup_ = tilingData_->kGroup; + groupCount_ = tilingData_->groupCount; + groupCountAlign_ = Ceil(groupCount_, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + perGroupExpertCount_ = tilingData_->perGroupExpertCount; + perGroupExpertCountAlign_ = tilingData_->perGroupExpertCountAlign; + renorm_ = tilingData_->renorm; + normType_ = tilingData_->normType; + groupSelectMode_ = tilingData_->groupSelectMode; + + expertCountAlign_ = Align(perGroupExpertCountAlign_ * groupCount_, sizeof(float)); + kAlign_ = Align(k_, sizeof(float)); + + isAlign_ = perGroupExpertCount_ == perGroupExpertCountAlign_; + + // init input gm buf + xGm_.SetGlobalBuffer((__gm__ T *)x + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_); + biasGm_.SetGlobalBuffer((__gm__ T *)bias, expertCount_); + + // init output gm buf + yGm_.SetGlobalBuffer((__gm__ T *)y + perCoreRowCount_ * k_ * blockIdx_, k_); + expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx + perCoreRowCount_ * k_ * blockIdx_, k_); + outGm_.SetGlobalBuffer((__gm__ float *)out + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_); + + // init que + pipe_->InitBuffer(xInQueue_, 1, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(yOutQueue_, 1, kAlign_ * sizeof(float)); + pipe_->InitBuffer(expertIdxOutQueue_, 1, kAlign_ * sizeof(int32_t)); + pipe_->InitBuffer(outOutQueue_, 1, expertCountAlign_ * sizeof(float)); + + pipe_->InitBuffer(biasBuf_, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(expertIdBuf_, expertCountAlign_ * sizeof(int32_t)); + + pipe_->InitBuffer(xNormBuf_, expertCountAlign_ * sizeof(float)); + + pipe_->InitBuffer(xNormWithBiasBuf_, expertCountAlign_ * sizeof(float)); + pipe_->InitBuffer(sortedInGroupBuf_, expertCountAlign_ * (sizeof(float) + sizeof(uint32_t))); + + pipe_->InitBuffer(sortedGroupIndexBuf_, groupCountAlign_ * sizeof(float) * CONSTANT_TWO); + pipe_->InitBuffer(topKExpertIdBuf_, kAlign_ * sizeof(int32_t)); + pipe_->InitBuffer(calcTmpBuf_, expertCountAlign_ * sizeof(float) * 10); +} + +template +__aicore__ inline void MoeGatingTopKGenerlized::Process() +{ + CopyInBiasAndInitExpertId(); + for (int64_t row = 0; row < curCoreRowCount_; row++) { + CopyInX(row); + ComputeX(); + if (tilingData_->outFlag) { + CopuOutXNorm(row); + } + SortInGroup(); + SelectTopKGroupIndex(); + SelectTopKExpertIdx(); + SelectTopKExpertScore(); + CumputeActualTopKExpertId(); + CopyOut(row); + } +} +} // namespace MoeGatingTopK +#endif // MOE_GATING_TOP_K_E_K_GENERALIZED_H diff --git a/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k_without_group.h b/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k_without_group.h new file mode 100644 index 000000000..a28ea3da1 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k/op_kernel/moe_gating_top_k_without_group.h @@ -0,0 +1,338 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_without_group.h + * \brief + */ +#ifndef MOE_GATING_TOP_K_E_K_WITHOUT_GROUP_H +#define MOE_GATING_TOP_K_E_K_WITHOUT_GROUP_H +#include "kernel_operator.h" +#include "common.h" +#include "kernel_utils.h" +namespace MoeGatingTopK { +using namespace AscendC; + +template +class MoeGatingTopKWithoutGroup { +public: + __aicore__ inline MoeGatingTopKWithoutGroup(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyInBiasAndInitExpertId(); + __aicore__ inline void CopyInX(int64_t progress); + __aicore__ inline void ComputeX(); + __aicore__ inline void CopuOutXNorm(int64_t row); + __aicore__ inline void SelectTopKExpertIdx(); + __aicore__ inline void SelectTopKExpertScore(); + __aicore__ inline void CopyOut(int64_t row); + +private: + TPipe *pipe_; + TQue xInQueue_; + TQue yOutQueue_; + TQue expertIdxOutQueue_; + TQue outOutQueue_; + + TBuf biasBuf_; // Store input bias + TBuf expertIdBuf_; // Expert ID + TBuf xNormWithBiasBuf_; // Store value after adding bias + TBuf xNormBuf_; // Store value after computing sigmoid or softmax + TBuf topKExpertIdBuf_; + TBuf calcTmpBuf_; + + GlobalTensor xGm_; + GlobalTensor biasGm_; + GlobalTensor yGm_; + GlobalTensor expertIdxGm_; + GlobalTensor outGm_; + + int64_t blockIdx_ = 0; + int64_t perCoreRowCount_ = 0; + int64_t curCoreRowCount_ = 0; + int64_t expertCount_ = 0; + bool addBias_ = false; + bool outFlag_ = false; + int64_t k_ = 0; + int64_t renorm_ = 0; + int64_t normType_ = 0; + int64_t expertCountAlign_ = 0; + const MoeGatingTopKTilingData *tilingData_; +}; + +template +__aicore__ inline void MoeGatingTopKWithoutGroup::CopyInBiasAndInitExpertId() +{ + LocalTensor biasTensor = biasBuf_.Get(); + LocalTensor expertIdTensor = expertIdBuf_.Get(); + DataCopyExtParams dataCopyParams{1, static_cast(expertCount_ * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if (addBias_) { + if constexpr (IsSameType::value) { + DataCopyPad(biasTensor, biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + } else { + DataCopyPad(biasTensor[expertCountAlign_].ReinterpretCast(), biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(biasTensor, biasTensor[expertCountAlign_].ReinterpretCast(), RoundMode::CAST_NONE, + expertCountAlign_); + PipeBarrier(); + } + } + ArithProgression(expertIdTensor, static_cast(0), static_cast(1), expertCount_); +} + +template +__aicore__ inline void MoeGatingTopKWithoutGroup::CopyInX(int64_t row) +{ + LocalTensor xInLocalTensor = xInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(expertCount_ * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if constexpr (IsSameType::value) { + DataCopyPad(xInLocalTensor, xGm_[row * expertCount_], dataCopyParams, dataCopyPadParams); + } else { + DataCopyPad(xInLocalTensor[expertCountAlign_].ReinterpretCast(), xGm_[row * expertCount_], dataCopyParams, + dataCopyPadParams); + } + xInQueue_.EnQue(xInLocalTensor); +} + +template +__aicore__ inline void MoeGatingTopKWithoutGroup::ComputeX() +{ + LocalTensor xNormTensor = xNormBuf_.Get(); + LocalTensor xInLocalTensor = xInQueue_.DeQue(); + LocalTensor xNormWithBiasTensor = xNormWithBiasBuf_.Get(); + LocalTensor biasTensor = biasBuf_.Get(); + + if constexpr (!IsSameType::value) { + Cast(xInLocalTensor, xInLocalTensor[expertCountAlign_].ReinterpretCast(), RoundMode::CAST_NONE, + expertCount_); + PipeBarrier(); + } + + if (normType_ == 1) { // sigmoid + LocalTensor calcNormTmpTensor = calcTmpBuf_.Get(); + Sigmoid(xNormTensor, xInLocalTensor, calcNormTmpTensor, expertCount_); + PipeBarrier(); + } else if (normType_ == 0) { // sigmoid + LocalTensor reduceValueTensor = calcTmpBuf_.Get(); + LocalTensor calcTmp = calcTmpBuf_.Get()[8]; + ReduceMax(reduceValueTensor, xInLocalTensor, calcTmp, expertCount_); + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float maxValue = reduceValueTensor.GetValue(0); + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Adds(xNormTensor, xInLocalTensor, -maxValue, expertCount_); + PipeBarrier(); + Exp(xNormTensor, xNormTensor, expertCount_); + PipeBarrier(); + ReduceSum(reduceValueTensor, xNormTensor, calcTmp, expertCount_); + eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float sumValue = reduceValueTensor.GetValue(0); + eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Muls(xNormTensor, xNormTensor, 1.0f / sumValue, expertCount_); + PipeBarrier(); + } + if (addBias_) { + Add(xNormWithBiasTensor, xNormTensor, biasTensor, expertCount_); + } else { + DataCopy(xNormWithBiasTensor, xNormTensor, expertCountAlign_); + } + + int64_t duplicateNum = expertCount_ % ONE_REPEAT_SORT_NUM; + int duplicateIndex = expertCount_ - duplicateNum; + if (duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(xNormWithBiasTensor.ReinterpretCast()[duplicateIndex], FLOAT32_NEG_INF, mask, 1, 1, 1); + PipeBarrier(); + } + xInQueue_.FreeTensor(xInLocalTensor); +} + +template +__aicore__ inline void MoeGatingTopKWithoutGroup::CopuOutXNorm(int64_t row) +{ + LocalTensor outOutTensor = outOutQueue_.AllocTensor(); + LocalTensor xNormTensor = xNormBuf_.Get(); + DataCopy(outOutTensor, xNormTensor, expertCountAlign_); + outOutQueue_.EnQue(outOutTensor); + outOutTensor = outOutQueue_.DeQue(); + DataCopyExtParams dataCopyParams{1, static_cast(expertCount_ * sizeof(float)), 0, 0, 0}; + DataCopyPad(outGm_[row * expertCount_], outOutTensor, dataCopyParams); + outOutQueue_.FreeTensor(outOutTensor); +} + +template +__aicore__ inline void MoeGatingTopKWithoutGroup::SelectTopKExpertIdx() +{ + LocalTensor expertIdxOut = expertIdxOutQueue_.AllocTensor(); + LocalTensor xNormWithBiasTensor = xNormWithBiasBuf_.Get(); + LocalTensor expertIdTensor = expertIdBuf_.Get(); + LocalTensor topKExpertId = topKExpertIdBuf_.Get(); + LocalTensor sortedScore = calcTmpBuf_.Get(); + LocalTensor sortTmp = calcTmpBuf_.Get()[expertCountAlign_ * CONSTANT_TWO]; + PipeBarrier(); + Sort(sortedScore, xNormWithBiasTensor, expertIdTensor, sortTmp, + expertCountAlign_ / ONE_REPEAT_SORT_NUM); + + GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = Ceil(k_ * sizeof(float) * CONSTANT_TWO, REPEAT_BYTES); + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = REPEAT_BLOCKS; + gatherMaskParams.src1RepeatStride = 0; + + uint64_t rsvdCnt = 0; // Used to store the number of elements retained after filtering + uint8_t src1Pattern = 2; // Built-in fixed pattern + PipeBarrier(); + GatherMask(topKExpertId, sortedScore.template ReinterpretCast(), src1Pattern, false, + static_cast(0), gatherMaskParams, rsvdCnt); + + DataCopy(expertIdxOut, topKExpertId, expertCountAlign_); + expertIdxOutQueue_.EnQue(expertIdxOut); +} + +template +__aicore__ inline void MoeGatingTopKWithoutGroup::SelectTopKExpertScore() +{ + LocalTensor xNormTensor = xNormBuf_.Get(); + LocalTensor yOutTensor = yOutQueue_.AllocTensor(); + LocalTensor topKExpertId = topKExpertIdBuf_.Get(); + LocalTensor topKExpertIdWithByte = calcTmpBuf_.Get(); + PipeBarrier(); + Muls(topKExpertIdWithByte, topKExpertId, static_cast(sizeof(float)), k_); + PipeBarrier(); + Gather(yOutTensor, xNormTensor, topKExpertIdWithByte.template ReinterpretCast(), static_cast(0), + k_); + + bool needRenorm = (normType_ == 1 ) || // Case 1: sigmoid + renorm + (normType_ == 0 && renorm_ == 1); // Case 3: softmax + renorm + if (needRenorm == 1) { + LocalTensor maxValueTensor = calcTmpBuf_.Get(); + LocalTensor tmpTensor = calcTmpBuf_.Get()[BLOCK_BYTES]; + PipeBarrier(); + ReduceSum(maxValueTensor, yOutTensor, tmpTensor, k_); + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float sumValue = maxValueTensor.GetValue(0) + tilingData_->eps; + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Duplicate(tmpTensor, sumValue, k_); + PipeBarrier(); + Div(yOutTensor, yOutTensor, tmpTensor, k_); + } + PipeBarrier(); + Muls(yOutTensor, yOutTensor, tilingData_->routedScalingFactor, k_); + + if constexpr (!IsSameType::value) { + PipeBarrier(); + Cast(yOutTensor.ReinterpretCast(), yOutTensor, RoundMode::CAST_RINT, k_); + } + + yOutQueue_.EnQue(yOutTensor); +} + +template +__aicore__ inline void MoeGatingTopKWithoutGroup::CopyOut(int64_t row) +{ + LocalTensor yOutTensor = yOutQueue_.DeQue(); + LocalTensor expertIdxOut = expertIdxOutQueue_.DeQue(); + DataCopyExtParams dataCopyParams{1, static_cast(k_ * sizeof(T)), 0, 0, 0}; + DataCopyPad(yGm_[row * k_], yOutTensor, dataCopyParams); + dataCopyParams.blockLen = k_ * sizeof(int32_t); + DataCopyPad(expertIdxGm_[row * k_], expertIdxOut, dataCopyParams); + yOutQueue_.FreeTensor(yOutTensor); + expertIdxOutQueue_.FreeTensor(expertIdxOut); +} + +template +__aicore__ inline void MoeGatingTopKWithoutGroup::Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, + GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKTilingData *tilingData, TPipe *tPipe) +{ + tilingData_ = tilingData; + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + perCoreRowCount_ = tilingData_->perCoreRowCount; + if (blockIdx_ == GetBlockNum() - 1) { + curCoreRowCount_ = tilingData_->lastCoreRowCount; + } else { + curCoreRowCount_ = tilingData_->perCoreRowCount; + } + expertCount_ = tilingData_->expertCount; + addBias_ = tilingData_->addBias == 1; + outFlag_ = tilingData_->outFlag == 1; + k_ = tilingData_->k; + renorm_ = tilingData_->renorm; + normType_ = tilingData_->normType; + expertCountAlign_ = Ceil(expertCount_, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + + // init input gm buf + xGm_.SetGlobalBuffer((__gm__ T *)x + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_); + biasGm_.SetGlobalBuffer((__gm__ T *)bias, expertCount_); + + // init output gm buf + yGm_.SetGlobalBuffer((__gm__ T *)y + perCoreRowCount_ * k_ * blockIdx_, k_); + expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx + perCoreRowCount_ * k_ * blockIdx_, k_); + outGm_.SetGlobalBuffer((__gm__ float *)out + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_); + + // init que + pipe_->InitBuffer(xInQueue_, 1, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(yOutQueue_, 1, Align(k_, sizeof(float)) * sizeof(float)); + pipe_->InitBuffer(expertIdxOutQueue_, 1, Align(k_, sizeof(float)) * sizeof(int32_t)); + pipe_->InitBuffer(outOutQueue_, 1, expertCountAlign_ * sizeof(float)); + + // init calc buf + pipe_->InitBuffer(biasBuf_, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(expertIdBuf_, expertCountAlign_ * sizeof(int32_t)); + pipe_->InitBuffer(xNormBuf_, expertCountAlign_ * sizeof(float)); + pipe_->InitBuffer(xNormWithBiasBuf_, expertCountAlign_ * sizeof(float)); + pipe_->InitBuffer(topKExpertIdBuf_, Align(k_, sizeof(float)) * sizeof(int32_t)); + + // init tmp buf + pipe_->InitBuffer(calcTmpBuf_, expertCountAlign_ * sizeof(float) * CONSTANT_EIGHT); +} + +template +__aicore__ inline void MoeGatingTopKWithoutGroup::Process() +{ + CopyInBiasAndInitExpertId(); + for (int64_t row = 0; row < curCoreRowCount_; row++) { + CopyInX(row); + ComputeX(); + if (outFlag_) { + CopuOutXNorm(row); + } + SelectTopKExpertIdx(); + SelectTopKExpertScore(); + CopyOut(row); + } +} +} // namespace MoeGatingTopK +#endif // MOE_GATING_TOP_K_E_K_WITHOUT_GROUP_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/CMakeLists.txt b/csrc/ascend/moe/moe_gating_top_k_hash/CMakeLists.txt new file mode 100644 index 000000000..86b308249 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_host/CMakeLists.txt b/csrc/ascend/moe/moe_gating_top_k_hash/op_host/CMakeLists.txt new file mode 100644 index 000000000..51419ec5d --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_host/CMakeLists.txt @@ -0,0 +1,62 @@ +# This program is free software, you can redistribute it and/or modify it. +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This file is a part of the CANN Open Software. +# Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ====================================================================================================================== + +# add_ops_compile_options( +# OP_NAME MoeGatingTopKHash +# OPTIONS --cce-auto-sync=off +# -Wno-deprecated-declarations +# -Werror +# -mllvm -cce-aicore-hoist-movemask=false +# --op_relocatable_kernel_binary=true +# ) + +# set(moe_gating_top_k_hash_depends transformer/moe/moe_gating_top_k_hash PARENT_SCOPE) + +# target_sources(op_host_aclnn PRIVATE +# op_host/moe_gating_top_k_hash_def.cpp +# ) + +# target_sources(optiling PRIVATE +# op_host/moe_gating_top_k_hash_tiling.cpp +# ) + +# if (NOT BUILD_OPEN_PROJECT) +# target_sources(opmaster_ct PRIVATE +# op_host/moe_gating_top_k_hash_tiling.cpp +# ) +# endif () + +# target_include_directories(optiling PRIVATE +# ${CMAKE_CURRENT_SOURCE_DIR}/op_host +# ) + +# target_sources(opsproto PRIVATE +# op_host/moe_gating_top_k_hash_proto.cpp +# ) + + +if (BUILD_OPEN_PROJECT) + message(STATUS "=== Debug: BUILD is BUILD_OPEN_PROJECT") + target_sources(op_host_aclnn PRIVATE + moe_gating_top_k_hash_def.cpp + ) + add_ops_compile_options( + OP_NAME MoeGatingTopKHash + OPTIONS --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true + ) +endif() + +if(NOT BUILD_OPS_RTY_KERNEL) + message(STATUS "=== Debug: BUILD is NOT BUILD_OPS_RTY_KERNEL") + add_modules_sources(OPTYPE moe_gating_top_k_hash ACLNNTYPE aclnn) +endif() diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_def.cpp b/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_def.cpp new file mode 100644 index 000000000..beeb35cb3 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_def.cpp @@ -0,0 +1,145 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_hash_def.cpp + * \brief + */ +#include "register/op_def_registry.h" + +namespace ops { +class MoeGatingTopKHash : public OpDef { +public: + explicit MoeGatingTopKHash(const char *name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, + ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, + ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, + ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("bias") + .ParamType(OPTIONAL) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, + ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, + ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, + ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("input_ids") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("tid2eid") + .ParamType(OPTIONAL) + .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, + ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, + ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, + ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, + ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("expert_idx") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("out") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, + ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Attr("k").Int(); + this->Attr("k_group").AttrType(OPTIONAL).Int(1); + this->Attr("group_count").AttrType(OPTIONAL).Int(1); + this->Attr("group_select_mode").AttrType(OPTIONAL).Int(0); + this->Attr("renorm").AttrType(OPTIONAL).Int(0); + this->Attr("norm_type").AttrType(OPTIONAL).Int(0); + this->Attr("out_flag").AttrType(OPTIONAL).Bool(false); + this->Attr("routed_scaling_factor").AttrType(OPTIONAL).Float(1.0); + this->Attr("eps").AttrType(OPTIONAL).Float(1e-20f); + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + + OpAICoreConfig regbaseCfg; + regbaseCfg.DynamicCompileStaticFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .ExtendCfgInfo("opFile.value", "moe_gating_top_k_hash"); + } +}; + +OP_ADD(MoeGatingTopKHash); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_proto.cpp b/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_proto.cpp new file mode 100644 index 000000000..f1b1d26ec --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_proto.cpp @@ -0,0 +1,68 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_hash_proto.h + * \brief + */ +#ifndef OPS_OP_PROTO_INC_MOEGATINGTOPK_H_ +#define OPS_OP_PROTO_INC_MOEGATINGTOPK_H_ + +#include "graph/operator_reg.h" + +namespace ge { + +/** + * @brief Compute renorm(sigmoid) and topk for moe input. + * + * @par Inputs: + * @li x: A 2D tensor which moe gating topk is applied, The shape is: (B*S, E), format supports ND, and data type must be float16, float or bfloat16. E(Expert num) can not be greater than 2048. E(Expert num) should be divisible by group_count. + * @li bias: A 1D tensor which is "bias" in moe gating topk. The shape is: (E), format supports ND, and data type must be the same as that of x. + * + * @par Outputs: + * @li y: A 2D tensor which is the topk value result of moe gating topk, format supports ND, and data type must be the same as that of x. + The size of the non-1 axis must be the same as that of the corresponding axis of x. + The size of the -1 axis must be the same as that of k. + * @li expert_idx: A 2D tensor which is the topk index result of moe gating topk, format supports ND, and data type must be int. The shape must be the same as that of y. + * @li out: A 2D tensor which is the renorm result of moe gating topk, format supports ND, and data type must be float. The shape must be the same as that of x. + * + * @par Attributes: + * @li k: A required attribute of type int. The value must greater than 0 and less than or equal to expert_num / group_count * k_group, indicating the topk value. + * @li k_group: An optional attribute of type int. It can not be less than 1, and can not be greater than group_count, indicating the topk group value. The default value is 1. + * @li group_count: An optional attribute of type int. It can not be less than 1, indicating the group count. The group_count * align_32(expert_num / group_count) can not be greater than 2048. The default value is 1. + * @li group_select_mode: An optional attribute of type int. 0 indicating that sort group by max values, 1 indicating that sort group by sum of top-2 values. The default value is 0. + * @li renorm: An optional attribute of type int. It can only be 0 now, indicating that norm firstly and then topk. The default value is 0. + * @li norm_type: An optional attribute of type int. 0 indicating that the softmax function is used, 1 indicating that the sigmoid function is used. The default value is 0. + * @li out_flag: An optional attribute of type bool. true indicating that has renorm output, false indicating that does not have renorm output. The default value is false. + * @li routed_scaling_factor: An optional attribute of type float, indicating the routed_scaling_factor coefficient in use. The default value is 1.0. + * @li eps: An optional attribute of type float, indicating the eps coefficient in use. The default value is 1e-20. + */ +REG_OP(MoeGatingTopKHash) + .INPUT(x, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16})) + .OPTIONAL_INPUT(bias, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16})) + .OPTIONAL_INPUT(input_ids, TensorType({DT_INT64, DT_INT32})) + .OPTIONAL_INPUT(tid2eid, TensorType({DT_INT64, DT_INT32})) + .OUTPUT(y, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16})) + .OUTPUT(expert_idx, TensorType({DT_INT32})) + .OUTPUT(out, TensorType({DT_FLOAT})) + .REQUIRED_ATTR(k, Int) + .ATTR(k_group, Int, 1) + .ATTR(group_count, Int, 1) + .ATTR(group_select_mode, Int, 0) + .ATTR(renorm, Int, 0) + .ATTR(norm_type, Int, 0) + .ATTR(out_flag, Bool, false) + .ATTR(routed_scaling_factor, Float, 1.0) + .ATTR(eps, Float, 1e-20f) + .OP_END_FACTORY_REG(MoeGatingTopKHash) + +} // namespace ge + +#endif // OPS_OP_PROTO_INC_MOEGATINGTOPK_H_ \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_tiling.cpp b/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_tiling.cpp new file mode 100644 index 000000000..b2f4ae5be --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_tiling.cpp @@ -0,0 +1,640 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file moe_gating_top_k_hash_tiling.cpp + * \brief + */ +#include "moe_gating_top_k_hash_tiling.h" +#include "moe_gating_top_k_hash_tiling_arch35.h" +#include + +namespace optiling { +const static int64_t GROUP_SELECT_MODE_MAX = 0; +const static int64_t GROUP_SELECT_MODE_SUM = 1; +const static int64_t RENORM_NO = 0; +const static int64_t RENORM_L1 = 1; +const static int64_t NORM_TYPE_SOFTMAX = 0; +const static int64_t NORM_TYPE_SIGMOID = 1; +const static int64_t NORM_TYPE_SOFTPLUS = 2; +const static int64_t OUT_FLAG_FALSE = 0; +const static int64_t OUT_FLAG_TRUE = 1; +const static size_t X_INPUT_DIMS = 2; +const static size_t BIAS_INPUT_DIMS = 1; +const static size_t Y_OUTPUT_DIMS = 2; +const static size_t EXPERT_IDX_OUTPUY_DIMS = 2; +const static size_t OUT_OUTPUT_DIMS = 2; +const static int64_t MAX_EXPERT_COUNT = 2048; + +const static int64_t X_INPUT_INDEX = 0; +const static int64_t BIAS_INPUT_INDEX = 1; +const static int64_t INPUT_IDS_INPUT_INDEX = 2; +const static int64_t TID_TO_EID_INPUT_INDEX = 3; +const static int64_t Y_OUTPUT_INDEX = 0; +const static int64_t EXPERT_IDX_OUTPUT_INDEX = 1; +const static int64_t OUT_OUTPUT_INDEX = 2; +const static int64_t K_ATTR_INDEX = 0; +const static int64_t K_GROUP_ATTR_INDEX = 1; +const static int64_t GROUP_COUNT_ATTR_INDEX = 2; +const static int64_t GROUP_SELECT_MODE_ATTR_INDEX = 3; +const static int64_t RENORM_ATTR_INDEX = 4; +const static int64_t NORM_TYPE_ATTR_INDEX = 5; +const static int64_t OUT_FLAG_ATTR_INDEX = 6; +const static int64_t ROUTED_SCALING_FACTOR_ATTR_INDEX = 7; +const static int64_t EPS_ATTR_INDEX = 8; +const static int64_t DEFAULT_WORKSPACE_SIZE = 16777216; // 预留16M空间 +const static uint32_t DATATYPESIZE_FLOAT = 4; +const static bool IS_LARGEST = true; +const static bool IS_INITINDEX = false; +const static bool IS_REUSESOURCE = false; +const static uint64_t WITH_GROUP_CONDITION = 1; +const static uint64_t WITHOUT_GROUP_CONDITION = 2; +const static uint64_t MAX_IN_GROUP_CONDITION = 3; +constexpr int32_t ROW_COUNT_PER_TASK = 1; + +const static uint64_t TILING_KEY_EXPERTNUM_GROUPNUM_ALIGN_HIGH_PERF = 0; +const static uint64_t TILING_KEY_WITHOUT_GROUP = 1; +const static uint64_t TILING_KEY_GENERALIZED = 2; +const static uint64_t TILING_KEY_WITHOUT_GROUP_0 = 3; +const static uint64_t TILING_KEY_WITHOUT_GROUP_1 = 4; +const static uint64_t TILING_KEY_WITHOUT_GROUP_2 = 5; +const static uint64_t TILING_KEY_WITHOUT_GROUP_3 = 6; + +template +static inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd))); +} + + +template +static inline T CeilAlign(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)) * (rnd)); +} + +inline static int64_t CeilLog4(int64_t x) +{ + return static_cast(std::ceil(std::log(x) / std::log(4))); // 4 for four +} + +class MoeGatingTopKHashTilingBase { +public: + explicit MoeGatingTopKHashTilingBase(gert::TilingContext *context) : context_(context) + { + Reset(); + } + ~MoeGatingTopKHashTilingBase() = default; + + bool IsCapable() + { + return true; + } + // 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小 + ge::graphStatus GetPlatformInfo() ; + // 2、获取INPUT/OUTPUT/ATTR信息 + ge::graphStatus GetShapeAttrsInfo() ; + // 3、计算数据切分TilingData + ge::graphStatus DoOpTiling() ; + // 4、计算高阶API的TilingData + ge::graphStatus DoLibApiTiling() ; + // 5、计算TilingKey + uint64_t GetTilingKey() const ; + // 6、计算Workspace 大小 + ge::graphStatus GetWorkspaceSize() ; + // 7、保存Tiling数据 + ge::graphStatus PostTiling() ; + void Reset(); + +private: + ge::graphStatus CheckInputShape(); + ge::graphStatus CheckAttr(); + ge::graphStatus CheckOutShape(); + void SplitRows(); + void CalTmpBufUbSize(); + + const gert::Shape *xShape_ = nullptr; + const gert::Shape *biasShape_ = nullptr; + const gert::Shape *inputIdsShape_ = nullptr; + const gert::Shape *tid2eidShape_ = nullptr; + const gert::Shape *yShape_ = nullptr; + const gert::Shape *expertIdxShape_ = nullptr; + const gert::Shape *outShape_ = nullptr; + ge::DataType inputIdsDtype; + ge::DataType tid2eidDtype; + + uint64_t coreNum_ = 0; + int64_t rows_ = 0; + int64_t expertCount_ = 0; + int64_t addBias_ = 0; + + int64_t k_ = 0; + int64_t kGroup_ = 0; + int64_t groupCount_ = 0; + int64_t perGroupExpertCount_ = 0; + int64_t groupSelectMode_ = GROUP_SELECT_MODE_MAX; + int64_t renorm_ = RENORM_NO; + int64_t normType_ = NORM_TYPE_SOFTMAX; + int64_t outFlag_ = OUT_FLAG_FALSE; + int64_t hashFlag_ = 0; + float routedScalingFactor_ = 1.0; + float eps_ = 1e-20f; + + int64_t inputDtypeSize_; + const char *opName_ = ""; + MoeGatingTopKHashTilingData moeGatingTopKTilingData_; + gert::TilingContext *context_ = nullptr; + uint64_t workspaceSize_ = 0; +}; + +ge::graphStatus MoeGatingTopKHashTilingBase::CheckInputShape() +{ + size_t xDimNum = xShape_->GetDimNum(); + OPS_ERR_IF(xDimNum != X_INPUT_DIMS, + OPS_LOG_E(context_, "The dim number of x is: %zu, but should be %zu.", xDimNum, X_INPUT_DIMS), + return ge::GRAPH_FAILED); + + // 通过输入获取rows 和 expertCount + rows_ = xShape_->GetDim(0); + expertCount_ = xShape_->GetDim(1); + moeGatingTopKTilingData_.set_rowCount(rows_); + moeGatingTopKTilingData_.set_expertCount(expertCount_); + if (biasShape_ != nullptr) { + addBias_ = 1; + size_t biasDimNum = biasShape_->GetDimNum(); + OPS_ERR_IF(biasDimNum != BIAS_INPUT_DIMS, + OPS_LOG_E(context_, "The dim number of bias is: %zu, but should be %zu.", biasDimNum, BIAS_INPUT_DIMS), + return ge::GRAPH_FAILED); + OPS_ERR_IF( + biasShape_->GetDim(0) != expertCount_, + OPS_LOG_E(context_, "The first dim of bias is: %ld, but should be %ld.", biasShape_->GetDim(0), expertCount_), + return ge::GRAPH_FAILED); + } + moeGatingTopKTilingData_.set_addBias(addBias_); + + if (inputIdsShape_ != nullptr) { + OPS_ERR_IF( + tid2eidShape_ == nullptr, + OPS_LOG_E(context_, "The tid2eid should not be empty when inputIds has value."), + return ge::GRAPH_FAILED); + } + if (tid2eidShape_ != nullptr) { + OPS_ERR_IF( + inputIdsShape_ == nullptr, + OPS_LOG_E(context_, "The inputIds should not be empty when tid2eid has value."), + return ge::GRAPH_FAILED); + } + if (inputIdsShape_ != nullptr && tid2eidShape_ != nullptr) { + hashFlag_ = 1; + OPS_LOG_I(context_, "hashFlag_ is 1."); + } + moeGatingTopKTilingData_.set_hashFlag(hashFlag_); + + OPS_ERR_IF(k_ > expertCount_, + OPS_LOG_E(context_, "k is: %ld, expert num is: %ld, k cannot be greater than expert num.", k_, expertCount_), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingBase::CheckAttr() +{ + OPS_ERR_IF( + expertCount_ > MAX_EXPERT_COUNT, + OPS_LOG_E(context_, "expert count is: %ld, but should not greater than %ld.", expertCount_, MAX_EXPERT_COUNT), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(k_ <= 0, OPS_LOG_E(context_, "k is: %ld, but should be greater than 0.", k_), return ge::GRAPH_FAILED); + + OPS_ERR_IF(kGroup_ <= 0, OPS_LOG_E(context_, "k_group is: %ld, but should be greater than 0.", kGroup_), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(kGroup_ > groupCount_, + OPS_LOG_E(context_, "k_group is: %ld, but should not greater than %ld.", kGroup_, groupCount_), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(groupCount_ <= 0, OPS_LOG_E(context_, "group_count is: %ld, but should be greater than 0.", groupCount_), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(normType_ != NORM_TYPE_SOFTMAX && normType_ != NORM_TYPE_SIGMOID && normType_ != NORM_TYPE_SOFTPLUS, + OPS_LOG_E(context_, "norm type is: %ld, but currently only support %ld, %ld and %ld.", normType_, + NORM_TYPE_SOFTMAX, NORM_TYPE_SIGMOID, NORM_TYPE_SOFTPLUS), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(normType_ == NORM_TYPE_SOFTPLUS && groupCount_ != 1, + OPS_LOG_E(context_, "norm type softplus only supported when groupCount equals 1, but got %ld.", groupCount_), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(groupSelectMode_ != GROUP_SELECT_MODE_SUM && groupSelectMode_ != GROUP_SELECT_MODE_MAX, + OPS_LOG_E(context_, "group select mode is: %ld, but currently only support %ld and %ld.", groupSelectMode_, + GROUP_SELECT_MODE_SUM, GROUP_SELECT_MODE_MAX), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(renorm_ != RENORM_NO, + OPS_LOG_E(context_, "renorm is: %ld, but currently only support %ld.", renorm_, RENORM_NO), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(expertCount_ % groupCount_ != 0, + OPS_LOG_E(context_, "Expert count : %ld is not divisible by k_group: %ld", expertCount_, groupCount_), + return ge::GRAPH_FAILED); + perGroupExpertCount_ = expertCount_ / groupCount_; + + OPS_ERR_IF(perGroupExpertCount_ < 1, + OPS_LOG_E(context_, "group expert count is: %ld, but should be greater than 1.", perGroupExpertCount_), + return ge::GRAPH_FAILED); + OPS_ERR_IF( + groupSelectMode_ == GROUP_SELECT_MODE_SUM && perGroupExpertCount_ < 2, + OPS_LOG_E(context_, + "group expert count is: %ld, if group select mode is: %ld, group expert count should be greater than 1.", + perGroupExpertCount_, groupSelectMode_), + return ge::GRAPH_FAILED); + OPS_ERR_IF(k_ > kGroup_ * perGroupExpertCount_, + OPS_LOG_E(context_, "k is: %ld, but should be smaller than %ld.", k_, kGroup_ * perGroupExpertCount_), + return ge::GRAPH_FAILED); + int64_t groupExpertCountAlign = CeilAlign(perGroupExpertCount_, 32L); + if (groupCount_ != 1 && groupCount_ != expertCount_ && kGroup_ != groupCount_) { + // 分组场景下才需要校验对齐后的数量 + OPS_ERR_IF(groupCount_ * groupExpertCountAlign > MAX_EXPERT_COUNT, + OPS_LOG_E(context_, "group count * group expert count align is: %ld, but should not greater than %ld.", + groupCount_ * groupExpertCountAlign, MAX_EXPERT_COUNT), + return ge::GRAPH_FAILED); + } + + moeGatingTopKTilingData_.set_perGroupExpertCount(perGroupExpertCount_); + moeGatingTopKTilingData_.set_perGroupExpertCountAlign(groupExpertCountAlign); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingBase::GetShapeAttrsInfo() +{ + opName_ = context_->GetNodeName(); + // 获取输入shape信息 + auto xShapePtr = context_->GetInputShape(X_INPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, xShapePtr, return ge::GRAPH_FAILED); + xShape_ = &xShapePtr->GetStorageShape(); + + auto biasShapePtr = context_->GetOptionalInputShape(BIAS_INPUT_INDEX); + biasShape_ = biasShapePtr == nullptr ? nullptr : &biasShapePtr->GetStorageShape(); + + auto inputIdsShapePtr = context_->GetOptionalInputShape(INPUT_IDS_INPUT_INDEX); + inputIdsShape_ = inputIdsShapePtr == nullptr ? nullptr : &inputIdsShapePtr->GetStorageShape(); + + auto tid2eidShapePtr = context_->GetOptionalInputShape(TID_TO_EID_INPUT_INDEX); + tid2eidShape_ = tid2eidShapePtr == nullptr ? nullptr : &tid2eidShapePtr->GetStorageShape(); + + // 获取输出shape + auto yShapePtr = context_->GetOutputShape(Y_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, yShapePtr, return ge::GRAPH_FAILED); + yShape_ = &yShapePtr->GetStorageShape(); + auto expertIdxPtr = context_->GetOutputShape(EXPERT_IDX_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, expertIdxPtr, return ge::GRAPH_FAILED); + expertIdxShape_ = &expertIdxPtr->GetStorageShape(); + auto outPtr = context_->GetOutputShape(OUT_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, outPtr, return ge::GRAPH_FAILED); + outShape_ = &outPtr->GetStorageShape(); + + auto x = context_->GetInputDesc(X_INPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, x, return ge::GRAPH_FAILED); + auto xDtype = x->GetDataType(); + OPS_ERR_IF( + (xDtype != ge::DataType::DT_FLOAT && xDtype != ge::DataType::DT_FLOAT16 && xDtype != ge::DataType::DT_BF16), + OPS_LOG_E(context_, "x dtype %s error, only supports float32, half, bf16. please check.", + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + + if (biasShapePtr != nullptr) { + auto biasDtype = context_->GetOptionalInputDesc(BIAS_INPUT_INDEX)->GetDataType(); + OPS_ERR_IF((biasDtype != xDtype), + OPS_LOG_E(context_, "bias dtype %s not equal x dtype %s, please check.", + ge::TypeUtils::DataTypeToSerialString(biasDtype).c_str(), + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + } + if (inputIdsShapePtr != nullptr) { + inputIdsDtype = context_->GetOptionalInputDesc(INPUT_IDS_INPUT_INDEX)->GetDataType(); + OPS_ERR_IF((inputIdsDtype != ge::DataType::DT_INT32 && inputIdsDtype != ge::DataType::DT_INT64), + OPS_LOG_E(context_, "inputIds dtype %s error, only supports int32 and int64. please check.", + ge::TypeUtils::DataTypeToSerialString(inputIdsDtype).c_str()), + return ge::GRAPH_FAILED); + } + if (tid2eidShapePtr != nullptr) { + tid2eidDtype = context_->GetOptionalInputDesc(TID_TO_EID_INPUT_INDEX)->GetDataType(); + OPS_ERR_IF((tid2eidDtype != ge::DataType::DT_INT32 && tid2eidDtype != ge::DataType::DT_INT64), + OPS_LOG_E(context_, "tid2eid dtype %s error, only supports int32 and int64. please check.", + ge::TypeUtils::DataTypeToSerialString(tid2eidDtype).c_str()), + return ge::GRAPH_FAILED); + } + + auto yDesc = context_->GetOutputDesc(Y_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, yDesc, return ge::GRAPH_FAILED); + auto yDtype = yDesc->GetDataType(); + OPS_ERR_IF((yDtype != xDtype), + OPS_LOG_E(context_, "y out dtype %s must be the same with x dtype %s.", + ge::TypeUtils::DataTypeToSerialString(yDtype).c_str(), + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + + auto expertIdDesc = context_->GetOutputDesc(EXPERT_IDX_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, expertIdDesc, return ge::GRAPH_FAILED); + auto expertIdDtype = expertIdDesc->GetDataType(); + OPS_ERR_IF((expertIdDtype != ge::DataType::DT_INT32), + OPS_LOG_E(context_, "expertId out dtype %s error, only supports int32. please check.", + ge::TypeUtils::DataTypeToSerialString(expertIdDtype).c_str()), + return ge::GRAPH_FAILED); + + auto normOutDesc = context_->GetOutputDesc(OUT_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, normOutDesc, return ge::GRAPH_FAILED); + auto normOutDtype = normOutDesc->GetDataType(); + OPS_ERR_IF((normOutDtype != ge::DataType::DT_FLOAT), + OPS_LOG_E(context_, "norm out dtype %s error, only supports float. please check.", + ge::TypeUtils::DataTypeToSerialString(normOutDtype).c_str()), + return ge::GRAPH_FAILED); + + // 获取属性 + auto attrs = context_->GetAttrs(); + OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED); + + const int64_t *kPtr = attrs->GetAttrPointer(K_ATTR_INDEX); + OPS_LOG_E_IF_NULL(context_, kPtr, return ge::GRAPH_FAILED); + k_ = *kPtr; + moeGatingTopKTilingData_.set_k(k_); + OPS_LOG_I(context_, "Attr k is: %ld ", k_); + + const int64_t *kGroupPtr = attrs->GetAttrPointer(K_GROUP_ATTR_INDEX); + if (kGroupPtr != nullptr) { + kGroup_ = *kGroupPtr; + moeGatingTopKTilingData_.set_kGroup(kGroup_); + } + OPS_LOG_I(context_, "Attr k_group is: %ld ", kGroup_); + + const int64_t *groupCountPtr = attrs->GetAttrPointer(GROUP_COUNT_ATTR_INDEX); + if (groupCountPtr != nullptr) { + groupCount_ = *groupCountPtr; + moeGatingTopKTilingData_.set_groupCount(groupCount_); + } + OPS_LOG_I(context_, "Attr group_count is: %ld ", groupCount_); + + const int64_t *groupSelectModePtr = attrs->GetAttrPointer(GROUP_SELECT_MODE_ATTR_INDEX); + if (groupSelectModePtr != nullptr) { + groupSelectMode_ = *groupSelectModePtr; + moeGatingTopKTilingData_.set_groupSelectMode(groupSelectMode_); + } + OPS_LOG_I(context_, "Attr group_select_mode is: %ld ", groupSelectMode_); + + const int64_t *renormPtr = attrs->GetAttrPointer(RENORM_ATTR_INDEX); + if (renormPtr != nullptr) { + renorm_ = *renormPtr; + moeGatingTopKTilingData_.set_renorm(renorm_); + } + OPS_LOG_I(context_, "Attr renorm is: %ld ", renorm_); + + const int64_t *normTypePtr = attrs->GetAttrPointer(NORM_TYPE_ATTR_INDEX); + if (normTypePtr != nullptr) { + normType_ = *normTypePtr; + moeGatingTopKTilingData_.set_normType(normType_); + } + OPS_LOG_I(context_, "Attr norm_type is: %ld ", normType_); + + const bool *outFlagPtr = attrs->GetAttrPointer(OUT_FLAG_ATTR_INDEX); + if (outFlagPtr != nullptr) { + outFlag_ = (*outFlagPtr) ? 1 : 0; + moeGatingTopKTilingData_.set_outFlag(outFlag_); + } + OPS_LOG_I(context_, "Attr out_flag is: %ld ", outFlag_); + + const float *routedScalingFactorPtr = attrs->GetAttrPointer(ROUTED_SCALING_FACTOR_ATTR_INDEX); + if (routedScalingFactorPtr != nullptr) { + routedScalingFactor_ = *routedScalingFactorPtr; + moeGatingTopKTilingData_.set_routedScalingFactor(routedScalingFactor_); + } + OPS_LOG_I(context_, "Attr routed_scaling_factor is: %f ", routedScalingFactor_); + + const float *epsPtr = attrs->GetAttrPointer(EPS_ATTR_INDEX); + if (epsPtr != nullptr) { + eps_ = *epsPtr; + moeGatingTopKTilingData_.set_eps(eps_); + } + OPS_LOG_I(context_, "Attr eps is: %f ", eps_); + + inputDtypeSize_ = static_cast(ge::GetSizeByDataType(context_->GetInputDesc(0)->GetDataType())); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingBase::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, OPS_LOG_E(context_, "fail to get platform info"), return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + coreNum_ = ascendcPlatform.GetCoreNumAiv(); + uint64_t ubSizePlatForm; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingBase::CheckOutShape() +{ + OPS_ERR_IF((yShape_->GetDimNum() != xShape_->GetDimNum()), + OPS_LOG_E(context_, "y out shape num %zu and x shape num %zu not equal, please check.", yShape_->GetDimNum(), + xShape_->GetDimNum()), + return ge::GRAPH_FAILED); + OPS_ERR_IF((expertIdxShape_->GetDimNum() != xShape_->GetDimNum()), + OPS_LOG_E(context_, "expertId out shape num %zu and x shape num %zu not equal, please check.", + expertIdxShape_->GetDimNum(), xShape_->GetDimNum()), + return ge::GRAPH_FAILED); + if (outShape_ != nullptr) { + OPS_ERR_IF((outShape_->GetDimNum() != xShape_->GetDimNum()), + OPS_LOG_E(context_, "norm out shape num %zu and x shape num %zu not equal, please check.", + outShape_->GetDimNum(), xShape_->GetDimNum()), + return ge::GRAPH_FAILED); + } + + OPS_ERR_IF((yShape_->GetDim(0) != xShape_->GetDim(0)), + OPS_LOG_E(context_, "y out dim[0] %ld not equal x dim[0] %ld, please check.", yShape_->GetDim(0), + xShape_->GetDim(0)), + return ge::GRAPH_FAILED); + OPS_ERR_IF((expertIdxShape_->GetDim(0) != xShape_->GetDim(0)), + OPS_LOG_E(context_, "expertId out dim[0] %ld not equal x dim[0] %ld, please check.", + expertIdxShape_->GetDim(0), xShape_->GetDim(0)), + return ge::GRAPH_FAILED); + if (outFlag_ && outShape_ != nullptr) { + OPS_ERR_IF((outShape_->GetDim(0) != xShape_->GetDim(0)), + OPS_LOG_E(context_, "norm out dim[0] %ld and x dim[0] %ld not equal, please check.", + outShape_->GetDim(0), outShape_->GetDim(0)), + return ge::GRAPH_FAILED); + } + + OPS_ERR_IF((yShape_->GetDim(1) != k_), + OPS_LOG_E(context_, "y dim[1] %ld not equal k %ld, please check.", yShape_->GetDim(1), k_), + return ge::GRAPH_FAILED); + OPS_ERR_IF((expertIdxShape_->GetDim(1) != k_), + OPS_LOG_E(context_, "expertId dim[1] %ld not equal k %ld, please check.", expertIdxShape_->GetDim(1), k_), + return ge::GRAPH_FAILED); + if (outFlag_ && outShape_ != nullptr) { + OPS_ERR_IF((outShape_->GetDim(1) != xShape_->GetDim(1)), + OPS_LOG_E(context_, "normOut dim[1] %ld and x dim[1] %ld not equal, please check.", outShape_->GetDim(1), + xShape_->GetDim(1)), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +void MoeGatingTopKHashTilingBase::SplitRows() +{ + int64_t perCoreRows = CeilDiv(rows_, static_cast(coreNum_)); + int64_t needCoreNum = CeilDiv(rows_, perCoreRows); + // perCoreRows cannot be 0 + int64_t lastCoreRows = rows_ % perCoreRows == 0 ? perCoreRows : rows_ % perCoreRows; + moeGatingTopKTilingData_.set_needCoreNum(needCoreNum); + moeGatingTopKTilingData_.set_perCoreRowCount(perCoreRows); + moeGatingTopKTilingData_.set_lastCoreRowCount(lastCoreRows); + int64_t vmsCount = CeilLog4(CeilDiv(kGroup_, 4L)); + OPS_LOG_I(context_, "vms count is: %ld", vmsCount); + moeGatingTopKTilingData_.set_vmsCount(vmsCount); // 需要归并的轮数 +} + +void MoeGatingTopKHashTilingBase::CalTmpBufUbSize() +{ + std::vector shape_vec = {expertCount_}; + ge::Shape shape(shape_vec); + uint32_t maxValue = 0; + uint32_t minValue = 0; + AscendC::GetSigmoidMaxMinTmpSize(shape, sizeof(float), false, maxValue, minValue); + + int64_t indexTmpBuf = (expertCount_ + 31) / 32 * 32 * static_cast(sizeof(float)); + moeGatingTopKTilingData_.set_calTmpBufUbSize(std::max(indexTmpBuf, static_cast(minValue))); +} + +ge::graphStatus MoeGatingTopKHashTilingBase::DoOpTiling() +{ + auto ret = GetPlatformInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = GetShapeAttrsInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckInputShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckOutShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckAttr(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + CalTmpBufUbSize(); + SplitRows(); + + ret = PostTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingBase::DoLibApiTiling() +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingBase::GetWorkspaceSize() +{ + // 计算workspace大小 + workspaceSize_ = DEFAULT_WORKSPACE_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingBase::PostTiling() +{ + context_->SetTilingKey(GetTilingKey()); + context_->SetBlockDim(moeGatingTopKTilingData_.get_needCoreNum()); + size_t *currentWorkspace = context_->GetWorkspaceSizes(1); + currentWorkspace[0] = workspaceSize_; + moeGatingTopKTilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), + context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(moeGatingTopKTilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +uint64_t MoeGatingTopKHashTilingBase::GetTilingKey() const +{ + // DeepSeekV3排序对齐高性能场景 + if (expertCount_ == 256 && groupCount_ == 8 && kGroup_ == 4 && k_ <= 32 && addBias_ && + groupSelectMode_ == GROUP_SELECT_MODE_SUM && renorm_ == RENORM_NO && normType_ == NORM_TYPE_SIGMOID && + !outFlag_) { + // DeepSeekV3排序对齐高性能场景 + return TILING_KEY_EXPERTNUM_GROUPNUM_ALIGN_HIGH_PERF; + } else if (groupCount_ == 1 || groupCount_ == expertCount_ || kGroup_ == groupCount_) { + /** + * 不分组场景： + * 1. 分组数为 1 + * 2. 分组数等于专家数（每个组只有一个专家） + * 3. 选择所有组 + */ + if (inputIdsShape_ == nullptr) { + return TILING_KEY_WITHOUT_GROUP; + } else if (inputIdsDtype == ge::DataType::DT_INT32 && tid2eidDtype == ge::DataType::DT_INT64) { + return TILING_KEY_WITHOUT_GROUP_0; + } else if (inputIdsDtype == ge::DataType::DT_INT32 && tid2eidDtype == ge::DataType::DT_INT32) { + return TILING_KEY_WITHOUT_GROUP_1; + } else if (inputIdsDtype == ge::DataType::DT_INT64 && tid2eidDtype == ge::DataType::DT_INT64) { + return TILING_KEY_WITHOUT_GROUP_2; + } else if (inputIdsDtype == ge::DataType::DT_INT64 && tid2eidDtype == ge::DataType::DT_INT32) { + return TILING_KEY_WITHOUT_GROUP_3; + } + } else { + return TILING_KEY_GENERALIZED; + } +} + +void MoeGatingTopKHashTilingBase::Reset() +{ + opName_ = nullptr; + return; +} + +ge::graphStatus TilingForMoeGatingTopKHash(gert::TilingContext *context) +{ + OPS_LOG_I(context, "TilingForMoeGatingTopKHash start"); + OPS_ERR_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("TilingForMoeGatingTopKHash", "Tiling context is null"), + return ge::GRAPH_FAILED); + + auto platformInfo = context->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, OPS_REPORT_VECTOR_INNER_ERR("TilingForMoeGatingTopKHash", "Tiling platformInfo is null"), + return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + auto socVersion = ascendcPlatform.GetSocVersion(); + if (socVersion == platform_ascendc::SocVersion::ASCEND950) { + OPS_LOG_I(context, "Using arch35 tiling for ASCEND950"); + MoeGatingTopKHashRegBase::MoeGatingTopKHashTilingRegbase moeGatingTopKTilingRegbase(context); + return moeGatingTopKTilingRegbase.DoOpTiling(); + } + + MoeGatingTopKHashTilingBase moeGatingTopKTilingBase(context); + return moeGatingTopKTilingBase.DoOpTiling(); +} + +static ge::graphStatus TilingPrepareForMoeGatingTopKHash(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(MoeGatingTopKHash) + .Tiling(TilingForMoeGatingTopKHash) + .TilingParse(TilingPrepareForMoeGatingTopKHash); +} // namespace optiling \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_tiling.h b/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_tiling.h new file mode 100644 index 000000000..8c47b37cf --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_tiling.h @@ -0,0 +1,84 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_hash_tiling.h + * \brief + */ + +#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_MOE_GATING_TOP_K_H +#define AIR_CXX_RUNTIME_V2_OP_IMPL_MOE_GATING_TOP_K_H + +#include +#include +#include +#include +#include +#include "register/op_impl_registry.h" +#include "platform/platform_infos_def.h" +#include "exe_graph/runtime/tiling_context.h" +#include "tiling/platform/platform_ascendc.h" +#include "register/op_def_registry.h" +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" +#include "error/ops_error.h" +#include "platform/platform_info.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(MoeGatingTopKHashTilingData) +TILING_DATA_FIELD_DEF(int64_t, needCoreNum); +TILING_DATA_FIELD_DEF(int64_t, rowCount); +TILING_DATA_FIELD_DEF(int64_t, perCoreRowCount); +TILING_DATA_FIELD_DEF(int64_t, lastCoreRowCount); +TILING_DATA_FIELD_DEF(int64_t, expertCount); +TILING_DATA_FIELD_DEF(int64_t, addBias); +TILING_DATA_FIELD_DEF(int64_t, k); +TILING_DATA_FIELD_DEF(int64_t, kGroup); +TILING_DATA_FIELD_DEF(int64_t, groupCount); +TILING_DATA_FIELD_DEF(int64_t, perGroupExpertCount); +TILING_DATA_FIELD_DEF(int64_t, perGroupExpertCountAlign); +TILING_DATA_FIELD_DEF(int64_t, groupSelectMode); +TILING_DATA_FIELD_DEF(int64_t, renorm); +TILING_DATA_FIELD_DEF(int64_t, normType); +TILING_DATA_FIELD_DEF(int64_t, outFlag); +TILING_DATA_FIELD_DEF(int64_t, hashFlag); +TILING_DATA_FIELD_DEF(int64_t, vmsCount); +TILING_DATA_FIELD_DEF(float, routedScalingFactor); +TILING_DATA_FIELD_DEF(float, eps); +TILING_DATA_FIELD_DEF(int64_t, calTmpBufUbSize); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(MoeGatingTopKHash, MoeGatingTopKHashTilingData) + +BEGIN_TILING_DATA_DEF(MoeGatingTopKHashRegbaseTilingData) +TILING_DATA_FIELD_DEF(int64_t, needCoreNum); +TILING_DATA_FIELD_DEF(int64_t, rowCount); +TILING_DATA_FIELD_DEF(int64_t, perCoreRowCount); +TILING_DATA_FIELD_DEF(int64_t, lastCoreRowCount); +TILING_DATA_FIELD_DEF(int64_t, expertCount); +TILING_DATA_FIELD_DEF(int64_t, addBias); +TILING_DATA_FIELD_DEF(int64_t, k); +TILING_DATA_FIELD_DEF(int64_t, kGroup); +TILING_DATA_FIELD_DEF(int64_t, groupCount); +TILING_DATA_FIELD_DEF(int64_t, perGroupExpertCount); +TILING_DATA_FIELD_DEF(int64_t, perGroupExpertCountAlign); +TILING_DATA_FIELD_DEF(int64_t, groupSelectMode); +TILING_DATA_FIELD_DEF(int64_t, renorm); +TILING_DATA_FIELD_DEF(int64_t, normType); +TILING_DATA_FIELD_DEF(int64_t, outFlag); +TILING_DATA_FIELD_DEF(int64_t, hashFlag); +TILING_DATA_FIELD_DEF(int64_t, vmsCount); +TILING_DATA_FIELD_DEF(float, routedScalingFactor); +TILING_DATA_FIELD_DEF(float, eps); +TILING_DATA_FIELD_DEF_STRUCT(SoftMaxTiling, softmaxTilingData); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(MoeGatingTopKHash_10000, MoeGatingTopKHashRegbaseTilingData) +struct MoeGatingTopKHashCompileInfo {}; +} // namespace optiling +#endif // AIR_CXX_RUNTIME_V2_OP_IMPL_MOE_GATING_TOP_K_H diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_tiling_arch35.h b/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_tiling_arch35.h new file mode 100644 index 000000000..3c13c0856 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_host/moe_gating_top_k_hash_tiling_arch35.h @@ -0,0 +1,594 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file moe_gating_top_k_hash_tiling_arch35.cpp + * \brief + */ +#include "moe_gating_top_k_hash_tiling.h" +#include + +namespace optiling { +namespace MoeGatingTopKHashRegBase { + +const static uint64_t MOE_GATING_TOP_K_REGBASE_TILING_KEY = 10000; +const static uint64_t MOE_GATING_TOP_K_REGBASE_TILING_KEY_1 = 10001; +const static uint64_t MOE_GATING_TOP_K_REGBASE_TILING_KEY_2 = 10002; +const static uint64_t MOE_GATING_TOP_K_REGBASE_TILING_KEY_3 = 10003; +const static uint64_t MOE_GATING_TOP_K_REGBASE_TILING_KEY_4 = 10004; + +const static int64_t GROUP_SELECT_MODE_MAX = 0; +const static int64_t GROUP_SELECT_MODE_SUM = 1; +const static int64_t RENORM_NO = 0; +const static int64_t RENORM_L1 = 1; +const static int64_t NORM_TYPE_SOFTMAX = 0; +const static int64_t NORM_TYPE_SIGMOID = 1; +const static int64_t NORM_TYPE_SOFTPLUS = 2; +const static int64_t OUT_FLAG_FALSE = 0; +const static int64_t OUT_FLAG_TRUE = 1; +const static size_t X_INPUT_DIMS = 2; +const static size_t BIAS_INPUT_DIMS = 1; +const static size_t Y_OUTPUT_DIMS = 2; +const static size_t EXPERT_IDX_OUTPUY_DIMS = 2; +const static size_t OUT_OUTPUT_DIMS = 2; +const static int64_t MAX_EXPERT_COUNT = 2048; + +const static int64_t X_INPUT_INDEX = 0; +const static int64_t BIAS_INPUT_INDEX = 1; +const static int64_t INPUT_IDS_INPUT_INDEX = 2; +const static int64_t TID_TO_EID_INPUT_INDEX = 3; +const static int64_t Y_OUTPUT_INDEX = 0; +const static int64_t EXPERT_IDX_OUTPUT_INDEX = 1; +const static int64_t OUT_OUTPUT_INDEX = 2; +const static int64_t K_ATTR_INDEX = 0; +const static int64_t K_GROUP_ATTR_INDEX = 1; +const static int64_t GROUP_COUNT_ATTR_INDEX = 2; +const static int64_t GROUP_SELECT_MODE_ATTR_INDEX = 3; +const static int64_t RENORM_ATTR_INDEX = 4; +const static int64_t MRGSORT_SIZE = 4; +const static int64_t NORM_TYPE_ATTR_INDEX = 5; +const static int64_t OUT_FLAG_ATTR_INDEX = 6; +const static int64_t ROUTED_SCALING_FACTOR_ATTR_INDEX = 7; +const static int64_t EPS_ATTR_INDEX = 8; +const static int64_t DEFAULT_WORKSPACE_SIZE = static_cast(16 * 1024 * 1024); // 预留16M空间 + +template +static inline T CeilDiv(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd))); +} + +template +static inline T CeilAlign(T num, T rnd) +{ + return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd)) * (rnd)); +} + +class MoeGatingTopKHashTilingRegbase { +public: + explicit MoeGatingTopKHashTilingRegbase(gert::TilingContext *context) : context_(context) + { + Reset(); + } + ~MoeGatingTopKHashTilingRegbase() = default; + + bool IsCapable() + { + if (socVersion_ != platform_ascendc::SocVersion::ASCEND950) { + return false; + } + return true; + } + // 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小 + ge::graphStatus GetPlatformInfo(); + // 2、获取INPUT/OUTPUT/ATTR信息 + ge::graphStatus GetShapeAttrsInfo(); + // 3、计算数据切分TilingData + ge::graphStatus DoOpTiling(); + // 4、计算高阶API的TilingData + ge::graphStatus DoLibApiTiling(); + // 5、计算TilingKey + uint64_t GetTilingKey() const; + // 6、计算Workspace 大小 + ge::graphStatus GetWorkspaceSize(); + // 7、保存Tiling数据 + ge::graphStatus PostTiling(); + void Reset(); + +private: + ge::graphStatus CheckInputShape(); + ge::graphStatus CheckAttr(); + ge::graphStatus CheckOutShape(); + void CalTmpBufUbSize(); + void SplitRows(); + void Tiling4GatherOutComputeSplitK(); + + const gert::Shape *xShape_ = nullptr; + const gert::Shape *biasShape_ = nullptr; + const gert::Shape *yShape_ = nullptr; + const gert::Shape *expertIdxShape_ = nullptr; + const gert::Shape *outShape_ = nullptr; + const gert::Shape *inputIdsShape_ = nullptr; + const gert::Shape *tid2eidShape_ = nullptr; + + ge::DataType inputIdsDtype; + ge::DataType tid2eidDtype; + + uint64_t coreNum_ = 0; + int64_t rows_ = 0; + int64_t expertCount_ = 0; + int64_t addBias_ = 0; + + int64_t k_ = 0; + int64_t kGroup_ = 1; + int64_t groupCount_ = 1; + int64_t groupSelectMode_ = GROUP_SELECT_MODE_MAX; + int64_t renorm_ = RENORM_NO; + int64_t normType_ = NORM_TYPE_SOFTMAX; + int64_t outFlag_ = OUT_FLAG_FALSE; + int64_t hashFlag_ = 0; + float routedScalingFactor_ = 1.0; + float eps_ = 1e-20f; + + int64_t inputDtypeSize_; + const char *opName_ = ""; + MoeGatingTopKHashRegbaseTilingData moeGatingTopKTilingData_; + gert::TilingContext *context_ = nullptr; + uint64_t workspaceSize_ = 0; + platform_ascendc::SocVersion socVersion_; +}; + +ge::graphStatus MoeGatingTopKHashTilingRegbase::CheckInputShape() +{ + size_t xDimNum = xShape_->GetDimNum(); + OPS_ERR_IF(xDimNum != X_INPUT_DIMS, + OPS_LOG_E(context_, "The dim number of x is: %zu, but should be %zu.", xDimNum, X_INPUT_DIMS), + return ge::GRAPH_FAILED); + + // 通过输入获取rows 和 expertCount + rows_ = xShape_->GetDim(0); + expertCount_ = xShape_->GetDim(1); + moeGatingTopKTilingData_.set_rowCount(rows_); + moeGatingTopKTilingData_.set_expertCount(expertCount_); + OPS_ERR_IF( + expertCount_ > MAX_EXPERT_COUNT, + OPS_LOG_E(context_, "expert count is: %ld, but should not greater than %ld.", expertCount_, MAX_EXPERT_COUNT), + return ge::GRAPH_FAILED); + + if (biasShape_ != nullptr) { + addBias_ = 1; + size_t biasDimNum = biasShape_->GetDimNum(); + OPS_ERR_IF(biasDimNum != BIAS_INPUT_DIMS, + OPS_LOG_E(context_, "The number of bias dim is: %zu, but should be %zu.", biasDimNum, BIAS_INPUT_DIMS), + return ge::GRAPH_FAILED); + OPS_ERR_IF(biasShape_->GetDim(0) != expertCount_, + OPS_LOG_E(context_, "The first dim of bias is: %ld, but should be expert num: %ld.", + biasShape_->GetDim(0), expertCount_), + return ge::GRAPH_FAILED); + } + moeGatingTopKTilingData_.set_addBias(addBias_); + + if (inputIdsShape_ != nullptr) { + OPS_ERR_IF( + tid2eidShape_ == nullptr, + OPS_LOG_E(context_, "The tid2eid should not be empty when inputIds has value."), + return ge::GRAPH_FAILED); + } + if (tid2eidShape_ != nullptr) { + OPS_ERR_IF( + inputIdsShape_ == nullptr, + OPS_LOG_E(context_, "The inputIds should not be empty when tid2eid has value."), + return ge::GRAPH_FAILED); + } + + if (inputIdsShape_ != nullptr && tid2eidShape_ != nullptr) { + hashFlag_ = 1; + OPS_LOG_I(context_, "hashFlag_ is 1."); + } + moeGatingTopKTilingData_.set_hashFlag(hashFlag_); + OPS_ERR_IF(k_ > expertCount_, + OPS_LOG_E(context_, "k is: %ld, expert num is: %ld, k cannot be greater than expert num.", k_, expertCount_), + return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingRegbase::CheckAttr() +{ + OPS_ERR_IF(k_ <= 0, OPS_LOG_E(context_, "k is: %ld, but should be greater than 0.", k_), return ge::GRAPH_FAILED); + OPS_ERR_IF(kGroup_ <= 0, OPS_LOG_E(context_, "k_group is: %ld, but should be greater than 0.", kGroup_), + return ge::GRAPH_FAILED); + OPS_ERR_IF(groupCount_ <= 0, OPS_LOG_E(context_, "group_count is: %ld, but should be greater than 0.", groupCount_), + return ge::GRAPH_FAILED); + OPS_ERR_IF(expertCount_ % groupCount_ != 0, + OPS_LOG_E(context_, "expert num : %ld is not divisible by group_count: %ld", expertCount_, groupCount_), + return ge::GRAPH_FAILED); + OPS_ERR_IF(kGroup_ > groupCount_, + OPS_LOG_E(context_, "k_group is: %ld, but should not greater than group_count: %ld", kGroup_, groupCount_), + return ge::GRAPH_FAILED); + OPS_ERR_IF(groupCount_ == expertCount_ && kGroup_ < k_, + OPS_LOG_E(context_, "k_group * group expert count is: %ld, but it must be greater than or equal to k: %ld.", + kGroup_, k_), + return ge::GRAPH_FAILED); + + if (kGroup_ == groupCount_ || groupCount_ == expertCount_) { + kGroup_ = 1; + groupCount_ = 1; + } + moeGatingTopKTilingData_.set_kGroup(kGroup_); + moeGatingTopKTilingData_.set_groupCount(groupCount_); + int64_t groupExpertCount = expertCount_ / groupCount_; + int64_t groupExpertCountAlign = CeilAlign(groupExpertCount, 32L); + moeGatingTopKTilingData_.set_perGroupExpertCount(expertCount_ / groupCount_); + moeGatingTopKTilingData_.set_perGroupExpertCountAlign(groupExpertCountAlign); + + OPS_ERR_IF(groupCount_ * groupExpertCountAlign > MAX_EXPERT_COUNT, + OPS_LOG_E(context_, "group count * group expert count align is: %ld, but should not greater than %ld.", + groupCount_ * groupExpertCountAlign, MAX_EXPERT_COUNT), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(kGroup_ * groupExpertCount < k_, + OPS_LOG_E(context_, "k_group * group expert count is: %ld, but it must be greater than or equal to k: %ld.", + kGroup_ * groupExpertCount, k_), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(groupExpertCount < 1, + OPS_LOG_E(context_, "per group expert count is: %ld, but should be greater than 0.", groupExpertCount), + return ge::GRAPH_FAILED); + OPS_ERR_IF( + groupSelectMode_ != GROUP_SELECT_MODE_SUM && groupSelectMode_ != GROUP_SELECT_MODE_MAX, + OPS_LOG_E(context_, "group select mode is: %ld, but currently only support %ld and %ld.", groupSelectMode_, + GROUP_SELECT_MODE_SUM, GROUP_SELECT_MODE_MAX), + return ge::GRAPH_FAILED); + OPS_ERR_IF(groupSelectMode_ == GROUP_SELECT_MODE_SUM && groupExpertCount < 2, + OPS_LOG_E(context_, + "group expert count is: %ld, if group select mode is: %ld, group expert count should be greater than 1.", + groupExpertCount, groupSelectMode_), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(renorm_ != RENORM_NO, + OPS_LOG_E(context_, "renorm is: %ld, but currently only support %ld.", renorm_, RENORM_NO), + return ge::GRAPH_FAILED); + + OPS_ERR_IF(normType_ != NORM_TYPE_SOFTMAX && normType_ != NORM_TYPE_SIGMOID && normType_ != NORM_TYPE_SOFTPLUS, + OPS_LOG_E(context_, "norm type is: %ld, but currently only support %ld, %ld and %ld.", normType_, + NORM_TYPE_SOFTMAX, NORM_TYPE_SIGMOID, NORM_TYPE_SOFTPLUS), + return ge::GRAPH_FAILED); + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingRegbase::GetShapeAttrsInfo() +{ + opName_ = context_->GetNodeName(); + // 获取输入shape信息 + auto xShapePtr = context_->GetInputShape(X_INPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, xShapePtr, return ge::GRAPH_FAILED); + xShape_ = &xShapePtr->GetStorageShape(); + auto biasShapePtr = context_->GetOptionalInputShape(BIAS_INPUT_INDEX); + biasShape_ = biasShapePtr == nullptr ? nullptr : &biasShapePtr->GetStorageShape(); + auto inputIdsShapePtr = context_->GetOptionalInputShape(INPUT_IDS_INPUT_INDEX); + inputIdsShape_ = inputIdsShapePtr == nullptr ? nullptr : &inputIdsShapePtr->GetStorageShape(); + auto tid2eidShapePtr = context_->GetOptionalInputShape(TID_TO_EID_INPUT_INDEX); + tid2eidShape_ = tid2eidShapePtr == nullptr ? nullptr : &tid2eidShapePtr->GetStorageShape(); + + // 获取输出shape + auto yShapePtr = context_->GetOutputShape(Y_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, yShapePtr, return ge::GRAPH_FAILED); + yShape_ = &yShapePtr->GetStorageShape(); + auto expertIdxPtr = context_->GetOutputShape(EXPERT_IDX_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, expertIdxPtr, return ge::GRAPH_FAILED); + expertIdxShape_ = &expertIdxPtr->GetStorageShape(); + auto outPtr = context_->GetOutputShape(OUT_OUTPUT_INDEX); + if (outPtr != nullptr) { + outShape_ = &outPtr->GetStorageShape(); + } + + auto x = context_->GetInputDesc(X_INPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, x, return ge::GRAPH_FAILED); + auto xDtype = x->GetDataType(); + OPS_ERR_IF( + (xDtype != ge::DataType::DT_FLOAT && xDtype != ge::DataType::DT_FLOAT16 && xDtype != ge::DataType::DT_BF16), + OPS_LOG_E(context_, "x dtype %s error, only supports float32, half, bf16. please check.", + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + + if (biasShapePtr != nullptr) { + auto biasDtype = context_->GetOptionalInputDesc(BIAS_INPUT_INDEX)->GetDataType(); + OPS_ERR_IF((biasDtype != xDtype), + OPS_LOG_E(context_, "bias dtype %s not equal x dtype %s, please check.", + ge::TypeUtils::DataTypeToSerialString(biasDtype).c_str(), + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + } + if (inputIdsShapePtr != nullptr) { + inputIdsDtype = context_->GetOptionalInputDesc(INPUT_IDS_INPUT_INDEX)->GetDataType(); + OPS_ERR_IF((inputIdsDtype != ge::DataType::DT_INT32 && inputIdsDtype != ge::DataType::DT_INT64), + OPS_LOG_E(context_, "inputIds dtype %s error, only supports int32 and int64. please check.", + ge::TypeUtils::DataTypeToSerialString(inputIdsDtype).c_str()), + return ge::GRAPH_FAILED); + } + if (tid2eidShapePtr != nullptr) { + tid2eidDtype = context_->GetOptionalInputDesc(TID_TO_EID_INPUT_INDEX)->GetDataType(); + OPS_ERR_IF((tid2eidDtype != ge::DataType::DT_INT32 && tid2eidDtype != ge::DataType::DT_INT64), + OPS_LOG_E(context_, "tid2eid dtype %s error, only supports int32 and int64. please check.", + ge::TypeUtils::DataTypeToSerialString(tid2eidDtype).c_str()), + return ge::GRAPH_FAILED); + } + + auto yDesc = context_->GetOutputDesc(Y_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, yDesc, return ge::GRAPH_FAILED); + auto yDtype = yDesc->GetDataType(); + OPS_ERR_IF((yDtype != xDtype), + OPS_LOG_E(context_, "y out dtype %s must be the same with x dtype %s.", + ge::TypeUtils::DataTypeToSerialString(yDtype).c_str(), + ge::TypeUtils::DataTypeToSerialString(xDtype).c_str()), + return ge::GRAPH_FAILED); + + auto expertIdDesc = context_->GetOutputDesc(EXPERT_IDX_OUTPUT_INDEX); + OPS_LOG_E_IF_NULL(context_, expertIdDesc, return ge::GRAPH_FAILED); + auto expertIdDtype = expertIdDesc->GetDataType(); + OPS_ERR_IF((expertIdDtype != ge::DataType::DT_INT32), + OPS_LOG_E(context_, "expertId out dtype %s error, only supports int32. please check.", + ge::TypeUtils::DataTypeToSerialString(expertIdDtype).c_str()), + return ge::GRAPH_FAILED); + + // 获取属性 + auto attrs = context_->GetAttrs(); + OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED); + + const int64_t *kPtr = attrs->GetAttrPointer(K_ATTR_INDEX); + OPS_LOG_E_IF_NULL(context_, kPtr, return ge::GRAPH_FAILED); + k_ = *kPtr; + moeGatingTopKTilingData_.set_k(k_); + OPS_LOG_I(context_, "Attr k is: %ld ", k_); + + const int64_t *kGroupPtr = attrs->GetAttrPointer(K_GROUP_ATTR_INDEX); + if (kGroupPtr != nullptr) { + kGroup_ = *kGroupPtr; + } + OPS_LOG_I(context_, "Attr k_group is: %ld ", kGroup_); + + const int64_t *groupCountPtr = attrs->GetAttrPointer(GROUP_COUNT_ATTR_INDEX); + if (groupCountPtr != nullptr) { + groupCount_ = *groupCountPtr; + } + OPS_LOG_I(context_, "Attr group_count is: %ld ", groupCount_); + + const int64_t *groupSelectModePtr = attrs->GetAttrPointer(GROUP_SELECT_MODE_ATTR_INDEX); + if (groupSelectModePtr != nullptr) { + groupSelectMode_ = *groupSelectModePtr; + } + moeGatingTopKTilingData_.set_groupSelectMode(groupSelectMode_); + OPS_LOG_I(context_, "Attr group_select_mode is: %ld ", groupSelectMode_); + + const int64_t *renormPtr = attrs->GetAttrPointer(RENORM_ATTR_INDEX); + if (renormPtr != nullptr) { + renorm_ = *renormPtr; + } + moeGatingTopKTilingData_.set_renorm(renorm_); + OPS_LOG_I(context_, "Attr renorm is: %ld ", renorm_); + + const int64_t *normTypePtr = attrs->GetAttrPointer(NORM_TYPE_ATTR_INDEX); + if (normTypePtr != nullptr) { + normType_ = *normTypePtr; + } + moeGatingTopKTilingData_.set_normType(normType_); + OPS_LOG_I(context_, "Attr norm_type is: %ld ", normType_); + + const bool *outFlagPtr = attrs->GetAttrPointer(OUT_FLAG_ATTR_INDEX); + if (outFlagPtr != nullptr) { + outFlag_ = (*outFlagPtr) ? 1 : 0; + } + moeGatingTopKTilingData_.set_outFlag(outFlag_); + OPS_LOG_I(context_, "Attr out_flag is: %ld ", outFlag_); + + const float *routedScalingFactorPtr = attrs->GetAttrPointer(ROUTED_SCALING_FACTOR_ATTR_INDEX); + if (routedScalingFactorPtr != nullptr) { + routedScalingFactor_ = *routedScalingFactorPtr; + } + moeGatingTopKTilingData_.set_routedScalingFactor(routedScalingFactor_); + OPS_LOG_I(context_, "Attr routed_scaling_factor is: %f ", routedScalingFactor_); + + const float *epsPtr = attrs->GetAttrPointer(EPS_ATTR_INDEX); + if (epsPtr != nullptr) { + eps_ = *epsPtr; + } + moeGatingTopKTilingData_.set_eps(eps_); + OPS_LOG_I(context_, "Attr eps is: %f ", eps_); + + auto outDesc = context_->GetOutputDesc(OUT_OUTPUT_INDEX); + if (outFlag_ && outDesc != nullptr) { + auto outDtype = outDesc->GetDataType(); + OPS_ERR_IF((outDtype != ge::DataType::DT_FLOAT), + OPS_LOG_E(context_, "norm out dtype %s error, only supports float32. please check.", + ge::TypeUtils::DataTypeToSerialString(outDtype).c_str()), + return ge::GRAPH_FAILED); + } + + inputDtypeSize_ = static_cast(ge::GetSizeByDataType(context_->GetInputDesc(0)->GetDataType())); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingRegbase::GetPlatformInfo() +{ + auto platformInfo = context_->GetPlatformInfo(); + OPS_ERR_IF(platformInfo == nullptr, OPS_LOG_E(context_, "fail to get platform info"), return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + coreNum_ = ascendcPlatform.GetCoreNumAiv(); + socVersion_ = ascendcPlatform.GetSocVersion(); + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingRegbase::CheckOutShape() +{ + OPS_ERR_IF((yShape_->GetDimNum() != xShape_->GetDimNum()), + OPS_LOG_E(context_, "y out shape num %zu and x shape num %zu not equal, please check.", yShape_->GetDimNum(), + xShape_->GetDimNum()), + return ge::GRAPH_FAILED); + OPS_ERR_IF((expertIdxShape_->GetDimNum() != xShape_->GetDimNum()), + OPS_LOG_E(context_, "expertId out shape num %zu and x shape num %zu not equal, please check.", + expertIdxShape_->GetDimNum(), xShape_->GetDimNum()), + return ge::GRAPH_FAILED); + if (outShape_ != nullptr) { + OPS_ERR_IF((outShape_->GetDimNum() != xShape_->GetDimNum()), + OPS_LOG_E(context_, "norm out shape num %zu and x shape num %zu not equal, please check.", + outShape_->GetDimNum(), xShape_->GetDimNum()), + return ge::GRAPH_FAILED); + } + + OPS_ERR_IF((yShape_->GetDim(0) != xShape_->GetDim(0)), + OPS_LOG_E(context_, "y out dim[0] %ld not equal x dim[0] %ld, please check.", yShape_->GetDim(0), + xShape_->GetDim(0)), + return ge::GRAPH_FAILED); + OPS_ERR_IF((expertIdxShape_->GetDim(0) != xShape_->GetDim(0)), + OPS_LOG_E(context_, "expertId out dim[0] %ld not equal x dim[0] %ld, please check.", + expertIdxShape_->GetDim(0), xShape_->GetDim(0)), + return ge::GRAPH_FAILED); + if (outFlag_ && outShape_ != nullptr) { + OPS_ERR_IF((outShape_->GetDim(0) != xShape_->GetDim(0)), + OPS_LOG_E(context_, "norm out dim[0] %ld and x dim[0] %ld not equal, please check.", + outShape_->GetDim(0), outShape_->GetDim(0)), + return ge::GRAPH_FAILED); + } + + OPS_ERR_IF((yShape_->GetDim(1) != k_), + OPS_LOG_E(context_, "y dim[1] %ld not equal k %ld, please check.", yShape_->GetDim(1), k_), + return ge::GRAPH_FAILED); + OPS_ERR_IF((expertIdxShape_->GetDim(1) != k_), + OPS_LOG_E(context_, "expertId dim[1] %ld not equal k %ld, please check.", expertIdxShape_->GetDim(1), k_), + return ge::GRAPH_FAILED); + if (outFlag_ && outShape_ != nullptr) { + OPS_ERR_IF((outShape_->GetDim(1) != xShape_->GetDim(1)), + OPS_LOG_E(context_, "normOut dim[1] %ld and x dim[1] %ld not equal, please check.", outShape_->GetDim(1), + xShape_->GetDim(1)), + return ge::GRAPH_FAILED); + } + return ge::GRAPH_SUCCESS; +} + +void MoeGatingTopKHashTilingRegbase::CalTmpBufUbSize() +{ + std::vector shape_vec = {groupCount_ * moeGatingTopKTilingData_.get_perGroupExpertCountAlign()}; + ge::Shape softmaxShape(shape_vec); + + uint32_t softmaxTmpSize = AscendC::GetSoftMaxMaxTmpSize(softmaxShape, sizeof(float), true); + AscendC::SoftMaxTilingFunc(softmaxShape, sizeof(float), softmaxTmpSize, moeGatingTopKTilingData_.softmaxTilingData); +} + +void MoeGatingTopKHashTilingRegbase::SplitRows() +{ + int64_t perCoreRows = CeilDiv(rows_, static_cast(coreNum_)); + int64_t needCoreNum = CeilDiv(rows_, perCoreRows); + if (perCoreRows == 0) { + OPS_LOG_E(context_, "perCoreRows can't be 0."); + return; + } + int64_t lastCoreRows = rows_ % perCoreRows == 0 ? perCoreRows : rows_ % perCoreRows; + moeGatingTopKTilingData_.set_needCoreNum(needCoreNum); + moeGatingTopKTilingData_.set_perCoreRowCount(perCoreRows); + moeGatingTopKTilingData_.set_lastCoreRowCount(lastCoreRows); + + int64_t vmsCount = 0; + if (kGroup_ > MRGSORT_SIZE) { + int64_t index = MRGSORT_SIZE; + while (index < kGroup_) { + index = index * MRGSORT_SIZE; + vmsCount++; + } + } + moeGatingTopKTilingData_.set_vmsCount(vmsCount); +} + +ge::graphStatus MoeGatingTopKHashTilingRegbase::DoOpTiling() +{ + auto ret = GetPlatformInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = GetShapeAttrsInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckInputShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckAttr(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckOutShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + CalTmpBufUbSize(); + SplitRows(); + + ret = PostTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingRegbase::DoLibApiTiling() +{ + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingRegbase::GetWorkspaceSize() +{ + // 计算workspace大小 + workspaceSize_ = DEFAULT_WORKSPACE_SIZE; + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeGatingTopKHashTilingRegbase::PostTiling() +{ + context_->SetTilingKey(GetTilingKey()); + context_->SetBlockDim(moeGatingTopKTilingData_.get_needCoreNum()); + size_t *currentWorkspace = context_->GetWorkspaceSizes(1); + currentWorkspace[0] = workspaceSize_; + moeGatingTopKTilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), + context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(moeGatingTopKTilingData_.GetDataSize()); + return ge::GRAPH_SUCCESS; +} + +uint64_t MoeGatingTopKHashTilingRegbase::GetTilingKey() const +{ + if (inputIdsShape_ == nullptr) { + return MOE_GATING_TOP_K_REGBASE_TILING_KEY; + } else if (inputIdsDtype == ge::DataType::DT_INT32 && tid2eidDtype == ge::DataType::DT_INT64) { + return MOE_GATING_TOP_K_REGBASE_TILING_KEY_1; + } else if (inputIdsDtype == ge::DataType::DT_INT32 && tid2eidDtype == ge::DataType::DT_INT32) { + return MOE_GATING_TOP_K_REGBASE_TILING_KEY_2; + } else if (inputIdsDtype == ge::DataType::DT_INT64 && tid2eidDtype == ge::DataType::DT_INT64) { + return MOE_GATING_TOP_K_REGBASE_TILING_KEY_3; + } else if (inputIdsDtype == ge::DataType::DT_INT64 && tid2eidDtype == ge::DataType::DT_INT32) { + return MOE_GATING_TOP_K_REGBASE_TILING_KEY_4; + } +} + +void MoeGatingTopKHashTilingRegbase::Reset() +{ + opName_ = nullptr; + return; +} + +} // namespace MoeGatingTopKHashRegBase +} // namespace optiling diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/arch35/common.h b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/arch35/common.h new file mode 100644 index 000000000..9438c4afa --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/arch35/common.h @@ -0,0 +1,214 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file common.h + * \brief + */ +#ifndef MOE_GATING_TOP_K_COMMON_H_V35 +#define MOE_GATING_TOP_K_COMMON_H_V35 + +#include "kernel_operator.h" + +namespace MoeGatingTopKHashRegbaseNS { +using namespace AscendC; +constexpr int32_t MIN_FP32 = 0xFF800000; +constexpr int64_t ONE_REPEAT_SORT_NUM = 32; +constexpr int64_t BLOCK_BYTES = 32; + +constexpr int64_t MERGE_LIST_TWO = 2; +constexpr int64_t MERGE_LIST_THREE = 3; +constexpr int64_t MERGE_LIST_FOUR = 4; + +constexpr int64_t MERGE_LIST_IDX_TWO = 2; +constexpr int64_t MERGE_LIST_IDX_THREE = 3; + +__aicore__ inline int64_t Ceil(int64_t a, int64_t b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes) +{ + if (bytes == 0) { + return 0; + } + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes; +} + +__aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes) +{ + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES; +} + +template +__aicore__ inline T Min(T a, T b) +{ + return a > b ? b : a; +} + +template +__aicore__ inline T Max(T a, T b) +{ + return a < b ? b : a; +} + +template +__aicore__ inline T1 CeilAlign(T1 a, T2 b) { + if (b == 0) { + return 0; + } + return (a + b - 1) / b * b; +} + +template +__aicore__ inline T1 CeilDiv(T1 x, T2 y) +{ + if (y != 0 && x != 0) { + const T1 quotient = x / y; + return (x % y != 0 && ((x ^ y) >= 0)) ? (quotient + 1) : quotient; + } + + return x; +} + +template +__aicore__ inline void SetWaitFlag(HardEvent evt) +{ + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(evt)); + SetFlag(eventId); + WaitFlag(eventId); +} + +constexpr AscendC::MicroAPI::CastTrait castTraitB162B32 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::UNKNOWN, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB322B16 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_RINT, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB322Int32 = { + AscendC::MicroAPI::RegLayout::UNKNOWN, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_TRUNC, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB322Int16 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_TRUNC, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB162Int8 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_TRUNC, +}; + +// load 对齐的 bfloat16,float16,bfloat32类型的 input(ub中)数据到 float32类型的dst(寄存器)中 +template +__aicore__ inline void LoadOneTensorForDtypeT(__local_mem__ T *input, MicroAPI::RegTensor &dst, + MicroAPI::MaskReg &preg, uint32_t offset) +{ + if constexpr (IsSameType::value) { + MicroAPI::RegTensor xFp16; + DataCopy(xFp16, ((__local_mem__ half *)(input) + (offset))); + Cast(dst, xFp16, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor xBf16; + DataCopy(xBf16, + ((__local_mem__ bfloat16_t *)(input) + (offset))); + Cast(dst, xBf16, preg); + } else { + DataCopy(dst, ((__local_mem__ float *)(input) + (offset))); + } +} + +// load 2个对齐的Tensor 到寄存器中 +template +__aicore__ inline void LoadTwoTensorForDtypeT(__local_mem__ T *src1, __local_mem__ T *src2, + MicroAPI::RegTensor &dst1, MicroAPI::RegTensor &dst2, + MicroAPI::MaskReg &dst1Preg, MicroAPI::MaskReg &dst2Preg, + uint32_t src1Offset, uint32_t src2Offset) +{ + if constexpr (IsSameType::value) { + MicroAPI::RegTensor xFp16Q; + MicroAPI::RegTensor xFp16R; + DataCopy(xFp16Q, ((__local_mem__ half *)(src1) + (src1Offset))); + DataCopy(xFp16R, ((__local_mem__ half *)(src2) + (src2Offset))); + Cast(dst1, xFp16Q, dst1Preg); + Cast(dst2, xFp16R, dst2Preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor xFp16Q; + MicroAPI::RegTensor xFp16R; + DataCopy(xFp16Q, ((__local_mem__ bfloat16_t *)(src1) + (src1Offset))); + DataCopy(xFp16R, ((__local_mem__ bfloat16_t *)(src2) + (src2Offset))); + Cast(dst1, xFp16Q, dst1Preg); + Cast(dst2, xFp16R, dst2Preg); + } else { + DataCopy(dst1, ((__local_mem__ float *)(src1) + (src1Offset))); + DataCopy(dst2, ((__local_mem__ float *)(src2) + (src2Offset))); + } +} + +// store 对齐的float32类型的src(寄存器)数据到output(ub)中，output数据类型支持bfloat16,float16,bfloat32,int32_t,int16_t,int8_t,uint8_t +template +__aicore__ inline void StoreOneTensorForDtypeT(__local_mem__ T *output, MicroAPI::RegTensor &src, + MicroAPI::MaskReg &preg, uint32_t offset) +{ + if constexpr (IsSameType::value) { + MicroAPI::RegTensor yFp16; + Cast(yFp16, src, preg); + DataCopy(((__local_mem__ half *)output + offset), yFp16, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor xBf16; + Cast(xBf16, src, preg); + DataCopy(((__local_mem__ bfloat16_t *)output + offset), + xBf16, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor zInt32; + MicroAPI::Cast(zInt32, src, preg); + DataCopy(((__local_mem__ int32_t *)output + offset), zInt32, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor zInt16; + MicroAPI::Cast(zInt16, src, preg); + DataCopy(((__local_mem__ int16_t *)output + offset), zInt16, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor yFp16; + MicroAPI::RegTensor zInt8; + Cast(yFp16, src, preg); + Cast(zInt8, yFp16, preg); + DataCopy(((__local_mem__ int8_t *)output + offset), zInt8, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor yFp16; + MicroAPI::RegTensor zUint8; + Cast(yFp16, src, preg); + Cast(zUint8, yFp16, preg); + DataCopy(((__local_mem__ uint8_t *)output + offset), zUint8, preg); + } else { + DataCopy(((__local_mem__ float *)output + offset), src, preg); + } +} + +} // namespace MoeGatingTopKHash +#endif // MOE_GATING_TOP_K_COMMON_H_V35 \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/arch35/moe_gating_top_k_hash_regbase.h b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/arch35/moe_gating_top_k_hash_regbase.h new file mode 100644 index 000000000..ae40058c7 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/arch35/moe_gating_top_k_hash_regbase.h @@ -0,0 +1,1577 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_hash_regbase.h + * \brief + */ +#ifndef MOE_GATING_TOP_K_REGBASE_H +#define MOE_GATING_TOP_K_REGBASE_H + +#include +#include "common.h" +#include "kernel_operator.h" +#include "kernel_utils.h" + +namespace MoeGatingTopKHashRegbaseNS +{ + using namespace AscendC; + using MicroAPI::RegTensor; + + constexpr int32_t CONSTANT_TWO = 2; + constexpr int32_t CONSTANT_THREE = 3; + constexpr int32_t CONSTANT_FOUR = 4; + constexpr int32_t CONSTANT_EIGHT = 8; + constexpr uint32_t VL_FLOAT_SIZE = VECTOR_REG_WIDTH / sizeof(float); + constexpr MicroAPI::DivSpecificMode mode = {MicroAPI::MaskMergeMode::ZEROING, true}; + + template + class MoeGatingTopKHashRegbase + { + public: + __aicore__ inline MoeGatingTopKHashRegbase(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR bias, GM_ADDR inputIds, GM_ADDR tid2eid, GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKHashRegbaseTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void CopyInBias(); + __aicore__ inline void CopyInX(int64_t progress); + __aicore__ inline void ComputeX(); + __aicore__ inline void ComputeSoftmax(); + __aicore__ inline void ComputeSigmoid(); + __aicore__ inline void ComputeSoftplus(); + __aicore__ inline void CopyOutXNorm(int64_t progress); + __aicore__ inline void SortInGroup(); + __aicore__ inline void SelectTopKGroupIndex(); + __aicore__ inline void FinalSortByKGroup(); + __aicore__ inline void FinalSortAfterKGroup(); + __aicore__ inline void SelectTopKExpertScore(); + __aicore__ inline void TopKCompute(); + __aicore__ inline void HashCompute(int64_t row); + __aicore__ inline void CopyOut(int64_t progress); + + __aicore__ inline void TopKGatherWithSmallKAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor); + __aicore__ inline void TopKGatherWithLargeKAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor); + __aicore__ inline void TopKGatherWithSmallKNotAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor, int32_t expertIdxPad, + int32_t perGroupExpertCountAlign); + __aicore__ inline void TopKGatherWithLargeKNotAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor, int32_t expertIdxPad, + int32_t perGroupExpertCountAlign); + __aicore__ inline void HashGatherWithSmallKAlignE(LocalTensor xSigmoidTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor); + __aicore__ inline void HashGatherWithSmallKNotAlignE(LocalTensor xSigmoidTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor, int32_t expertIdxPad, + int32_t perGroupExpertCountAlign); + + private: + TPipe *pipe_; + TQue xInQueue_; + TQue yOutQueue_; + TQue expertIdxOutQueue_; + TQue outOutQueue_; + + TBuf biasBuf_; + TBuf xBiasBuf_; + TBuf xSigmoidBuf_; + TBuf groupBuf_; + TBuf sortedInGroupBuf_; + TBuf sortedGroupBuf_; + TBuf indexBuffer_; + TBuf finalSortBuffer_; + + GlobalTensor xGm_; + GlobalTensor biasGm_; + GlobalTensor yGm_; + GlobalTensor expertIdxGm_; + GlobalTensor outGm_; + GlobalTensor inputIdsGm_; + GlobalTensor tid2eidGm_; + + LocalTensor indexTensor; + LocalTensor sortedInGroupTensor; + LocalTensor sortedGroupTensor; + LocalTensor mrgSortTensor; + + int64_t blockIdx_; + int64_t curCoreRowCount_; + int64_t expertCount_; + int64_t k_; + int64_t kGroup_; + int64_t groupCount_; + float routedScalingFactor_; + float eps_; + bool hasBias_ = false; + bool hashFlag_ = false; + + int64_t perGroupExpertCount_; + int64_t perGroupExpertCountAlign_; + int64_t expertCountAlign_; + int64_t groupSelectMode_; + const MoeGatingTopKHashRegbaseTilingData *tilingData_; + }; + + template + __aicore__ inline void MoeGatingTopKHashRegbase::CopyInBias() + { + if (!hasBias_) + { + return; + } + LocalTensor biasTensor = biasBuf_.Get(); + + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = groupCount_; + dataCopyParams.blockLen = perGroupExpertCount_ * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(T) / BLOCK_BYTES; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + + DataCopyPad(biasTensor, biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::CopyInX(int64_t row) + { + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = groupCount_; + dataCopyParams.blockLen = perGroupExpertCount_ * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(T) / BLOCK_BYTES; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if (tilingData_->normType == 0) + { + LocalTensor xInLocalTensor = xInQueue_.AllocTensor(); + if constexpr (IsSameType::value) + { + DataCopyPad(xInLocalTensor, xGm_[row * expertCount_], dataCopyParams, dataCopyPadParams); + } + else + { + DataCopyPad(xInLocalTensor[expertCountAlign_].ReinterpretCast(), xGm_[row * expertCount_], dataCopyParams, + dataCopyPadParams); + } + xInQueue_.EnQue(xInLocalTensor); + } + else + { + LocalTensor xInLocalTensor = xInQueue_.AllocTensor(); + DataCopyPad(xInLocalTensor, xGm_[row * expertCount_], dataCopyParams, dataCopyPadParams); + xInQueue_.EnQue(xInLocalTensor); + } + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::ComputeSoftmax() + { + LocalTensor xSoftmaxTensor = xSigmoidBuf_.Get(); + LocalTensor xBiasTensor = xBiasBuf_.Get(); + indexTensor = indexBuffer_.Get(); + + uint32_t size = perGroupExpertCountAlign_ * groupCount_; + uint32_t perGroupExpertCount0 = perGroupExpertCount_; + uint32_t perGroupExpertCountAlign0 = perGroupExpertCountAlign_; + uint16_t groupCount0 = groupCount_; + LocalTensor xInLocalTensor = xInQueue_.DeQue(); + if constexpr (!IsSameType::value) + { + Cast(xInLocalTensor, xInLocalTensor[expertCountAlign_].template ReinterpretCast(), RoundMode::CAST_NONE, + expertCountAlign_); + } + + int64_t duplicateNum = perGroupExpertCount_ % ONE_REPEAT_SORT_NUM; + int duplicateIndex = perGroupExpertCount_ - duplicateNum; + if (perGroupExpertCount_ != perGroupExpertCountAlign_ && duplicateNum > 0) + { + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(xInLocalTensor.ReinterpretCast()[duplicateIndex], MIN_FP32, mask, groupCount_, 1, + (perGroupExpertCountAlign_ * sizeof(float)) / BLOCK_BYTES); + } + + SoftMaxShapeInfo softmaxShapeInfo; + softmaxShapeInfo.srcM = 1; + softmaxShapeInfo.srcK = groupCount_ * perGroupExpertCountAlign_; + softmaxShapeInfo.oriSrcM = 1; + softmaxShapeInfo.oriSrcK = groupCount_ * perGroupExpertCountAlign_; + SoftMax(xSoftmaxTensor, xInLocalTensor, tilingData_->softmaxTilingData, softmaxShapeInfo); + + __local_mem__ float *softmaxOutAddr = (__local_mem__ float *)xSoftmaxTensor.GetPhyAddr(); + __local_mem__ int32_t *indexOutAddr = (__local_mem__ int32_t *)indexTensor.GetPhyAddr(); + __local_mem__ float *addBiasOutAddr = (__local_mem__ float *)xBiasTensor.GetPhyAddr(); + if (hasBias_) + { + LocalTensor biasTensor = biasBuf_.Get(); + __VEC_SCOPE__ + { + RegTensor vregBiasFp32; + RegTensor vregIndex; + RegTensor vregSoftmaxResult; + RegTensor vregBiasResult; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + __local_mem__ T *biasAddr = (__local_mem__ T *)biasTensor.GetPhyAddr(); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + LoadOneTensorForDtypeT(softmaxOutAddr, vregSoftmaxResult, preg0, i * VL_FLOAT_SIZE); + LoadOneTensorForDtypeT(biasAddr, vregBiasFp32, preg0, i * VL_FLOAT_SIZE); + // add bias + MicroAPI::Add(vregBiasResult, vregSoftmaxResult, vregBiasFp32, preg0); + // 使用Arange生成排序索引, 起始值为i乘每个循环的veclen + MicroAPI::Arange(vregIndex, static_cast(i * VL_FLOAT_SIZE)); + MicroAPI::DataCopy(addBiasOutAddr + i * VL_FLOAT_SIZE, vregBiasResult, preg0); + MicroAPI::DataCopy(indexOutAddr + i * VL_FLOAT_SIZE, vregIndex, preg0); + } + MicroAPI::LocalMemBar(); + + // pad min fp32 + RegTensor vregPad; + MicroAPI::UnalignReg u0; + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + auto padUbAddr = addBiasOutAddr + i * perGroupExpertCountAlign0 + perGroupExpertCount0; + MicroAPI::DataCopyUnAlign(padUbAddr, vregPad, u0, perGroupExpertCountAlign0 - perGroupExpertCount0); + MicroAPI::DataCopyUnAlignPost(padUbAddr, u0, 0); + } + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vregIndex; + RegTensor vregSoftmaxResult; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + LoadOneTensorForDtypeT(softmaxOutAddr, vregSoftmaxResult, preg0, i * VL_FLOAT_SIZE); + // 使用Arange生成排序索引, 起始值为i乘每个循环的veclen + MicroAPI::Arange(vregIndex, static_cast(i * VL_FLOAT_SIZE)); + MicroAPI::DataCopy(addBiasOutAddr + i * VL_FLOAT_SIZE, vregSoftmaxResult, preg0); + MicroAPI::DataCopy(indexOutAddr + i * VL_FLOAT_SIZE, vregIndex, preg0); + } + MicroAPI::LocalMemBar(); + + // pad min fp32 + RegTensor vregPad; + MicroAPI::UnalignReg u0; + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + auto padUbAddr = addBiasOutAddr + i * perGroupExpertCountAlign0 + perGroupExpertCount0; + MicroAPI::DataCopyUnAlign(padUbAddr, vregPad, u0, perGroupExpertCountAlign0 - perGroupExpertCount0); + MicroAPI::DataCopyUnAlignPost(padUbAddr, u0, 0); + } + } + } + xInQueue_.FreeTensor(xInLocalTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::ComputeSigmoid() + { + LocalTensor xSigmoidTensor = xSigmoidBuf_.Get(); + LocalTensor xBiasTensor = xBiasBuf_.Get(); + indexTensor = indexBuffer_.Get(); + + uint32_t size = perGroupExpertCountAlign_ * groupCount_; + uint32_t perGroupExpertCount0 = perGroupExpertCount_; + uint32_t perGroupExpertCountAlign0 = perGroupExpertCountAlign_; + uint16_t groupCount0 = groupCount_; + LocalTensor xInLocalTensor = xInQueue_.DeQue(); + if (hasBias_) + { + LocalTensor biasTensor = biasBuf_.Get(); + __VEC_SCOPE__ + { + RegTensor vregBiasFp32; + RegTensor vregIndex; + RegTensor vregSigmoidResult; + RegTensor vregBiasResult; + RegTensor vregOne; + RegTensor vregInFp32; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + + __local_mem__ T *inputAddr = (__local_mem__ T *)xInLocalTensor.GetPhyAddr(); + __local_mem__ T *biasAddr = (__local_mem__ T *)biasTensor.GetPhyAddr(); + __local_mem__ float *sigmoidOutAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ int32_t *indexOutAddr = (__local_mem__ int32_t *)indexTensor.GetPhyAddr(); + __local_mem__ float *addBiasOutAddr = (__local_mem__ float *)xBiasTensor.GetPhyAddr(); + + // sigmoid + MicroAPI::Duplicate(vregOne, static_cast(1), preg0); + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + LoadTwoTensorForDtypeT(inputAddr, biasAddr, vregInFp32, vregBiasFp32, preg0, preg0, + i * VL_FLOAT_SIZE, i * VL_FLOAT_SIZE); + MicroAPI::Muls(vreg1, vregInFp32, static_cast(-1), preg0); + MicroAPI::Exp(vreg2, vreg1, preg0); + MicroAPI::Adds(vreg3, vreg2, static_cast(1), preg0); + MicroAPI::Div(vregSigmoidResult, vregOne, vreg3, preg0); + // add bias + MicroAPI::Add(vregBiasResult, vregSigmoidResult, vregBiasFp32, preg0); + // 使用Arange生成排序索引, 起始值为i乘每个循环的veclen + MicroAPI::Arange(vregIndex, static_cast(i * VL_FLOAT_SIZE)); + MicroAPI::DataCopy(sigmoidOutAddr + i * VL_FLOAT_SIZE, vregSigmoidResult, preg0); + MicroAPI::DataCopy(addBiasOutAddr + i * VL_FLOAT_SIZE, vregBiasResult, preg0); + MicroAPI::DataCopy(indexOutAddr + i * VL_FLOAT_SIZE, vregIndex, preg0); + } + MicroAPI::LocalMemBar(); + + // pad min fp32 + RegTensor vregPad; + MicroAPI::UnalignReg u0; + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + auto padUbAddr = addBiasOutAddr + i * perGroupExpertCountAlign0 + perGroupExpertCount0; + MicroAPI::DataCopyUnAlign(padUbAddr, vregPad, u0, perGroupExpertCountAlign0 - perGroupExpertCount0); + MicroAPI::DataCopyUnAlignPost(padUbAddr, u0, 0); + } + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vregBiasFp32; + RegTensor vregIndex; + RegTensor vregSigmoidResult; + RegTensor vregOne; + RegTensor vregInFp32; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + + __local_mem__ T *inputAddr = (__local_mem__ T *)xInLocalTensor.GetPhyAddr(); + __local_mem__ float *sigmoidOutAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ int32_t *indexOutAddr = (__local_mem__ int32_t *)indexTensor.GetPhyAddr(); + __local_mem__ float *addBiasOutAddr = (__local_mem__ float *)xBiasTensor.GetPhyAddr(); + + // sigmoid + MicroAPI::Duplicate(vregOne, static_cast(1), preg0); + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + LoadOneTensorForDtypeT(inputAddr, vregInFp32, preg0, i * VL_FLOAT_SIZE); + MicroAPI::Muls(vreg1, vregInFp32, static_cast(-1), preg0); + MicroAPI::Exp(vreg2, vreg1, preg0); + MicroAPI::Adds(vreg3, vreg2, static_cast(1), preg0); + MicroAPI::Div(vregSigmoidResult, vregOne, vreg3, preg0); + // 使用Arange生成排序索引, 起始值为i乘每个循环的veclen + MicroAPI::Arange(vregIndex, static_cast(i * VL_FLOAT_SIZE)); + MicroAPI::DataCopy(sigmoidOutAddr + i * VL_FLOAT_SIZE, vregSigmoidResult, preg0); + MicroAPI::DataCopy(addBiasOutAddr + i * VL_FLOAT_SIZE, vregSigmoidResult, preg0); + MicroAPI::DataCopy(indexOutAddr + i * VL_FLOAT_SIZE, vregIndex, preg0); + } + MicroAPI::LocalMemBar(); + + // pad inf + RegTensor vregPad; + MicroAPI::UnalignReg u0; + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + auto padUbAddr = addBiasOutAddr + i * perGroupExpertCountAlign0 + perGroupExpertCount0; + MicroAPI::DataCopyUnAlign(padUbAddr, vregPad, u0, perGroupExpertCountAlign0 - perGroupExpertCount0); + MicroAPI::DataCopyUnAlignPost(padUbAddr, u0, 0); + } + } + } + xInQueue_.FreeTensor(xInLocalTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::ComputeSoftplus() + { + LocalTensor xSoftplusTensor = xSigmoidBuf_.Get(); + LocalTensor xBiasTensor = xBiasBuf_.Get(); + indexTensor = indexBuffer_.Get(); + + uint32_t size = perGroupExpertCountAlign_ * groupCount_; + uint32_t perGroupExpertCount0 = perGroupExpertCount_; + uint32_t perGroupExpertCountAlign0 = perGroupExpertCountAlign_; + uint16_t groupCount0 = groupCount_; + LocalTensor xInLocalTensor = xInQueue_.DeQue(); + if (hasBias_) + { + LocalTensor biasTensor = biasBuf_.Get(); + __VEC_SCOPE__ + { + RegTensor vregBiasFp32; + RegTensor vregIndex; + RegTensor vregSoftplusResult; + RegTensor vregBiasResult; + RegTensor vregOne; + RegTensor vregInFp32; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + + __local_mem__ T *inputAddr = (__local_mem__ T *)xInLocalTensor.GetPhyAddr(); + __local_mem__ T *biasAddr = (__local_mem__ T *)biasTensor.GetPhyAddr(); + __local_mem__ float *softplusOutAddr = (__local_mem__ float *)xSoftplusTensor.GetPhyAddr(); + __local_mem__ int32_t *indexOutAddr = (__local_mem__ int32_t *)indexTensor.GetPhyAddr(); + __local_mem__ float *addBiasOutAddr = (__local_mem__ float *)xBiasTensor.GetPhyAddr(); + + // softplus + MicroAPI::Duplicate(vregOne, static_cast(1), preg0); + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + LoadTwoTensorForDtypeT(inputAddr, biasAddr, vregInFp32, vregBiasFp32, preg0, preg0, + i * VL_FLOAT_SIZE, i * VL_FLOAT_SIZE); + MicroAPI::Exp(vreg1, vregInFp32, preg0); + MicroAPI::Adds(vreg2, vreg1, static_cast(1), preg0); + MicroAPI::Ln(vreg3, vreg2, preg0); + MicroAPI::Sqrt(vregSoftplusResult, vreg3, preg0); + // add bias + MicroAPI::Add(vregBiasResult, vregSoftplusResult, vregBiasFp32, preg0); + // 使用Arange生成排序索引, 起始值为i乘每个循环的veclen + MicroAPI::Arange(vregIndex, static_cast(i * VL_FLOAT_SIZE)); + MicroAPI::DataCopy(softplusOutAddr + i * VL_FLOAT_SIZE, vregSoftplusResult, preg0); + MicroAPI::DataCopy(addBiasOutAddr + i * VL_FLOAT_SIZE, vregBiasResult, preg0); + MicroAPI::DataCopy(indexOutAddr + i * VL_FLOAT_SIZE, vregIndex, preg0); + } + MicroAPI::LocalMemBar(); + + // pad min fp32 + RegTensor vregPad; + MicroAPI::UnalignReg u0; + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + auto padUbAddr = addBiasOutAddr + i * perGroupExpertCountAlign0 + perGroupExpertCount0; + MicroAPI::DataCopyUnAlign(padUbAddr, vregPad, u0, perGroupExpertCountAlign0 - perGroupExpertCount0); + MicroAPI::DataCopyUnAlignPost(padUbAddr, u0, 0); + } + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vregBiasFp32; + RegTensor vregIndex; + RegTensor vregSoftplusResult; + RegTensor vregOne; + RegTensor vregInFp32; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + + __local_mem__ T *inputAddr = (__local_mem__ T *)xInLocalTensor.GetPhyAddr(); + __local_mem__ float *softplusOutAddr = (__local_mem__ float *)xSoftplusTensor.GetPhyAddr(); + __local_mem__ int32_t *indexOutAddr = (__local_mem__ int32_t *)indexTensor.GetPhyAddr(); + __local_mem__ float *addBiasOutAddr = (__local_mem__ float *)xBiasTensor.GetPhyAddr(); + + // sigmoid + MicroAPI::Duplicate(vregOne, static_cast(1), preg0); + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + LoadOneTensorForDtypeT(inputAddr, vregInFp32, preg0, i * VL_FLOAT_SIZE); + MicroAPI::Exp(vreg1, vregInFp32, preg0); + MicroAPI::Adds(vreg2, vreg1, static_cast(1), preg0); + MicroAPI::Ln(vreg3, vreg2, preg0); + MicroAPI::Sqrt(vregSoftplusResult, vreg3, preg0); + // 使用Arange生成排序索引, 起始值为i乘每个循环的veclen + MicroAPI::Arange(vregIndex, static_cast(i * VL_FLOAT_SIZE)); + MicroAPI::DataCopy(softplusOutAddr + i * VL_FLOAT_SIZE, vregSoftplusResult, preg0); + MicroAPI::DataCopy(addBiasOutAddr + i * VL_FLOAT_SIZE, vregSoftplusResult, preg0); + MicroAPI::DataCopy(indexOutAddr + i * VL_FLOAT_SIZE, vregIndex, preg0); + } + MicroAPI::LocalMemBar(); + + // pad inf + RegTensor vregPad; + MicroAPI::UnalignReg u0; + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + auto padUbAddr = addBiasOutAddr + i * perGroupExpertCountAlign0 + perGroupExpertCount0; + MicroAPI::DataCopyUnAlign(padUbAddr, vregPad, u0, perGroupExpertCountAlign0 - perGroupExpertCount0); + MicroAPI::DataCopyUnAlignPost(padUbAddr, u0, 0); + } + } + } + xInQueue_.FreeTensor(xInLocalTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::ComputeX() + { + if (tilingData_->normType == 0) + { + ComputeSoftmax(); + return; + } + if (tilingData_->normType == 1) + { + ComputeSigmoid(); + return; + } + ComputeSoftplus(); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::CopyOutXNorm(int64_t progress) + { + if (tilingData_->outFlag == 0) + { + return; + } + LocalTensor outOutTensor = outOutQueue_.AllocTensor(); + LocalTensor xSigmoidTensor = xSigmoidBuf_.Get(); + DataCopy(outOutTensor, xSigmoidTensor, expertCountAlign_); + outOutQueue_.EnQue(outOutTensor); + outOutTensor = outOutQueue_.DeQue(); + DataCopyExtParams dataCopyParams{ + static_cast(groupCount_), static_cast(perGroupExpertCount_ * sizeof(float)), + static_cast((perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(float) / BLOCK_BYTES), 0, 0}; + DataCopyPad(outGm_[progress * expertCount_], outOutTensor, dataCopyParams); + outOutQueue_.FreeTensor(outOutTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::SortInGroup() + { + LocalTensor xBiasTensor = xBiasBuf_.Get(); + LocalTensor sortedInGroupTensor = sortedInGroupBuf_.Get(); // 组内排序的结果, 后续归并需要 + LocalTensor tmpLocal = finalSortBuffer_.Get(); + + if (perGroupExpertCountAlign_ == ONE_REPEAT_SORT_NUM) + { + Sort32(sortedInGroupTensor, xBiasTensor, indexTensor, groupCount_); + } + else + { + for (uint16_t i = 0; i < groupCount_; i++) + { + Sort(sortedInGroupTensor[i * perGroupExpertCountAlign_ * CONSTANT_TWO], + xBiasTensor[i * perGroupExpertCountAlign_], indexTensor[i * perGroupExpertCountAlign_], + tmpLocal, perGroupExpertCountAlign_ / ONE_REPEAT_SORT_NUM); + } + } + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::TopKCompute() + { + LocalTensor xBiasTensor = xBiasBuf_.Get(); + LocalTensor sortedInGroupTensor = sortedInGroupBuf_.Get(); // 组内排序的结果, 后续归并需要 + LocalTensor tmpLocal = finalSortBuffer_.Get(); + + Sort(sortedInGroupTensor, xBiasTensor, indexTensor, tmpLocal, + perGroupExpertCountAlign_ * groupCount_ / ONE_REPEAT_SORT_NUM); + + LocalTensor xSigmoidTensor = xSigmoidBuf_.Get(); + LocalTensor yTensor = yOutQueue_.AllocTensor(); + LocalTensor expertIdxTensor = expertIdxOutQueue_.AllocTensor(); + + LocalTensor sortedInGroupTensorCast = sortedInGroupTensor.template ReinterpretCast(); + + int32_t expertIdxPad = perGroupExpertCountAlign_ - perGroupExpertCount_; + if (k_ <= VL_FLOAT_SIZE) + { + if (expertIdxPad != 0) + { + TopKGatherWithSmallKNotAlignE(xSigmoidTensor, sortedInGroupTensorCast, expertIdxTensor, yTensor, k_, eps_, + routedScalingFactor_, expertIdxPad, perGroupExpertCountAlign_); + } + else + { + TopKGatherWithSmallKAlignE(xSigmoidTensor, sortedInGroupTensorCast, expertIdxTensor, yTensor, k_, eps_, + routedScalingFactor_); + } + } + else + { + if (expertIdxPad != 0) + { + TopKGatherWithLargeKNotAlignE(xSigmoidTensor, sortedInGroupTensorCast, expertIdxTensor, yTensor, k_, eps_, + routedScalingFactor_, expertIdxPad, perGroupExpertCountAlign_); + } + else + { + TopKGatherWithLargeKAlignE(xSigmoidTensor, sortedInGroupTensorCast, expertIdxTensor, yTensor, k_, eps_, + routedScalingFactor_); + } + } + + yOutQueue_.EnQue(yTensor); + expertIdxOutQueue_.EnQue(expertIdxTensor); + } + + template + __aicore__ inline void + MoeGatingTopKHashRegbase::HashGatherWithSmallKAlignE(LocalTensor xSigmoidTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor) + { + __local_mem__ float *inputAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ uint32_t *expertIdxAddr = (__local_mem__ uint32_t *)expertIdxTensor.GetPhyAddr(); + __local_mem__ T *outputAddr = (__local_mem__ T *)yTensor.GetPhyAddr(); + + if (tilingData_->normType == 0) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg1, expertIdxAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Muls(vreg2, vreg2, routedScalingFactor, preg0); + StoreOneTensorForDtypeT(outputAddr, vreg2, preg0, 0); + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vreg4; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg1, expertIdxAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::ReduceSum(vreg3, vreg2, preg0); + MicroAPI::Adds(vreg3, vreg3, eps, preg0); + MicroAPI::Duplicate(vreg4, vreg3, preg0); + MicroAPI::Div(vreg4, vreg2, vreg4, preg0); + MicroAPI::Muls(vreg4, vreg4, routedScalingFactor, preg0); + StoreOneTensorForDtypeT(outputAddr, vreg4, preg0, 0); + } + } + } + + template + __aicore__ inline void + MoeGatingTopKHashRegbase::HashGatherWithSmallKNotAlignE(LocalTensor xSigmoidTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor, int32_t expertIdxPad, + int32_t perGroupExpertCountAlign) + { + __local_mem__ float *inputAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ uint32_t *expertIdxAddr = (__local_mem__ uint32_t *)expertIdxTensor.GetPhyAddr(); + __local_mem__ T *outputAddr = (__local_mem__ T *)yTensor.GetPhyAddr(); + + if (tilingData_->normType == 0) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vregAlign; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + + MicroAPI::Duplicate(vregAlign, perGroupExpertCountAlign, preg0); + + MicroAPI::DataCopy(vreg1, expertIdxAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Muls(vreg2, vreg2, routedScalingFactor, preg0); + + // compute expertIdx: id = id - floor_div(id, perGroupExpertCountAlign) * pad + MicroAPI::Div(vregAlign, vreg1, vregAlign, preg0); + MicroAPI::Muls(vregAlign, vregAlign, expertIdxPad, preg0); + MicroAPI::Sub(vreg1, vreg1, vregAlign, preg0); + + StoreOneTensorForDtypeT(outputAddr, vreg2, preg0, 0); + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vreg4; + RegTensor vregAlign; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + + MicroAPI::Duplicate(vregAlign, perGroupExpertCountAlign, preg0); + + MicroAPI::DataCopy(vreg1, expertIdxAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::ReduceSum(vreg3, vreg2, preg0); + MicroAPI::Adds(vreg3, vreg3, eps, preg0); + MicroAPI::Duplicate(vreg4, vreg3, preg0); + MicroAPI::Div(vreg4, vreg2, vreg4, preg0); + MicroAPI::Muls(vreg4, vreg4, routedScalingFactor, preg0); + + // compute expertIdx: id = id - floor_div(id, perGroupExpertCountAlign) * pad + MicroAPI::Div(vregAlign, vreg1, vregAlign, preg0); + MicroAPI::Muls(vregAlign, vregAlign, expertIdxPad, preg0); + MicroAPI::Sub(vreg1, vreg1, vregAlign, preg0); + + StoreOneTensorForDtypeT(outputAddr, vreg4, preg0, 0); + } + } + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::HashCompute(int64_t row) + { + LocalTensor expertIdxTensor = expertIdxOutQueue_.AllocTensor(); + LocalTensor hashExpertId = sortedInGroupBuf_.Get(); + LocalTensor hashExpertIdInt32 = hashExpertId.template ReinterpretCast(); + + U1 key = inputIdsGm_.GetValue(row); + SetWaitFlag(HardEvent::S_MTE2); + DataCopyExtParams dataCopyParams{1, static_cast(k_ * sizeof(U2)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + DataCopyPad(hashExpertId, tid2eidGm_[key * k_], dataCopyParams, dataCopyPadParams); + SetWaitFlag(HardEvent::MTE2_V); + if constexpr (IsSameType::value) { + DataCopy(expertIdxTensor, hashExpertId, Align(k_, sizeof(int32_t))); + } else { + Cast(hashExpertIdInt32, hashExpertId, RoundMode::CAST_NONE, Align(k_, sizeof(U2))); + PipeBarrier(); + DataCopy(expertIdxTensor, hashExpertIdInt32, Align(k_, sizeof(int32_t))); + } + PipeBarrier(); + + LocalTensor xSigmoidTensor = xSigmoidBuf_.Get(); + LocalTensor yTensor = yOutQueue_.AllocTensor(); + LocalTensor sortedInGroupTensorCast = sortedInGroupTensor.template ReinterpretCast(); + int32_t expertIdxPad = perGroupExpertCountAlign_ - perGroupExpertCount_; + // 当前只支持k小于64场景 + if (k_ <= VL_FLOAT_SIZE) + { + if (expertIdxPad != 0) + { + HashGatherWithSmallKNotAlignE(xSigmoidTensor, expertIdxTensor, yTensor, k_, eps_, + routedScalingFactor_, expertIdxPad, perGroupExpertCountAlign_); + } + else + { + HashGatherWithSmallKAlignE(xSigmoidTensor, expertIdxTensor, yTensor, k_, eps_, + routedScalingFactor_); + } + } + yOutQueue_.EnQue(yTensor); + expertIdxOutQueue_.EnQue(expertIdxTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::SelectTopKGroupIndex() + { + sortedInGroupTensor = sortedInGroupBuf_.Get(); + LocalTensor top2InGroupTensor = groupBuf_.Get(); + LocalTensor tmpLocal = xBiasBuf_.Get(); + // 排序，将kgroup选出来 + sortedGroupTensor = sortedGroupBuf_.Get(); + + uint16_t groupCount0 = groupCount_; + uint32_t perGroupExpertCountAlign0 = perGroupExpertCountAlign_; + int32_t groupCountNumAlign = (groupCount_ + 31) / 32 * 32; + uint32_t padNegInfNum = groupCountNumAlign - groupCount_; + if (groupSelectMode_ == 1) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vregPad; + + MicroAPI::UnalignReg u0; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + + __local_mem__ float *inputAddr = (__local_mem__ float *)sortedInGroupTensor.GetPhyAddr(); + __local_mem__ float *outputAddr = (__local_mem__ float *)top2InGroupTensor.GetPhyAddr(); + + // pair reduce sum + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + MicroAPI::DataCopy( + vreg0, vreg1, inputAddr + i * perGroupExpertCountAlign0 * 2); + MicroAPI::PairReduceSum(vreg2, vreg0, preg0); + MicroAPI::DataCopyUnAlign(outputAddr, vreg2, u0, 1); + } + MicroAPI::DataCopyUnAlign(outputAddr, vregPad, u0, + padNegInfNum); + MicroAPI::DataCopyUnAlignPost(outputAddr, u0, 0); + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vregPad; + + MicroAPI::UnalignReg u0; + + __local_mem__ float *inputAddr = (__local_mem__ float *)sortedInGroupTensor.GetPhyAddr(); + __local_mem__ float *outputAddr = (__local_mem__ float *)top2InGroupTensor.GetPhyAddr(); + + // max + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + MicroAPI::DataCopy( + vreg0, vreg1, inputAddr + i * perGroupExpertCountAlign0 * 2); + MicroAPI::DataCopyUnAlign(outputAddr, vreg0, u0, 1); + } + MicroAPI::DataCopyUnAlign(outputAddr, vregPad, u0, + padNegInfNum); + MicroAPI::DataCopyUnAlignPost(outputAddr, u0, 0); + } + } + + Sort(sortedGroupTensor, top2InGroupTensor, indexTensor, tmpLocal, + groupCountNumAlign / ONE_REPEAT_SORT_NUM); + + uint32_t size = groupCountNumAlign; + int32_t kGroup0 = kGroup_; + int32_t kGroupNumAlign = (kGroup0 + 31) / 32 * 32; + uint32_t padkGroupNum = kGroupNumAlign - kGroup0; + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vregPad; + MicroAPI::UnalignReg u0; + + __local_mem__ int32_t *inputAddr = (__local_mem__ int32_t *)sortedGroupTensor.GetPhyAddr(); + __local_mem__ int32_t *outputAddr = (__local_mem__ int32_t *)top2InGroupTensor.GetPhyAddr(); + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + MicroAPI::DataCopy(vreg0, vreg1, + inputAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::DataCopy(outputAddr + i * VL_FLOAT_SIZE, vreg1, preg0); + } + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + outputAddr = outputAddr + kGroup0; + MicroAPI::DataCopyUnAlign(outputAddr, (RegTensor &)vregPad, u0, padkGroupNum); + MicroAPI::DataCopyUnAlignPost(outputAddr, u0, 0); + } + + Sort(sortedGroupTensor, top2InGroupTensor, indexTensor, tmpLocal, + kGroupNumAlign / ONE_REPEAT_SORT_NUM); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::FinalSortByKGroup() + { + mrgSortTensor = finalSortBuffer_.Get(); + LocalTensor tmpLocal = sortedGroupTensor.template ReinterpretCast(); + uint32_t offset[MRG_SORT_ELEMENT_LEN] = {0, 0, 0, 0}; + + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + + uint16_t lenArr[CONSTANT_FOUR] = { + static_cast(perGroupExpertCount_), static_cast(perGroupExpertCount_), + static_cast(perGroupExpertCount_), static_cast(perGroupExpertCount_)}; + MrgSort4Info params{lenArr, false, 0b1111, 1}; + MrgSortSrcList srcList; + + for (int32_t i = kGroup_ - 1; i >= 0; i -= CONSTANT_FOUR) + { +#if defined(__CCE_KT_TEST__) + int32_t mrgLen = std::min(i + 1, CONSTANT_FOUR); +#else + int32_t mrgLen = min(i + 1, CONSTANT_FOUR); +#endif + if (mrgLen > 1) + { + if (mrgLen == CONSTANT_FOUR) + { + offset[0] = tmpLocal.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + offset[1] = tmpLocal.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2; + offset[CONSTANT_TWO] = tmpLocal.GetValue((i - 2) * 2) * perGroupExpertCountAlign_ * 2; + offset[CONSTANT_THREE] = tmpLocal.GetValue((i - 3) * 2) * perGroupExpertCountAlign_ * 2; + } + else if (mrgLen == CONSTANT_THREE) + { + offset[0] = tmpLocal.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + offset[1] = tmpLocal.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2; + offset[CONSTANT_TWO] = tmpLocal.GetValue((i - 2) * 2) * perGroupExpertCountAlign_ * 2; + offset[CONSTANT_THREE] = 0; + params.elementLengths[CONSTANT_THREE] = 0; + params.validBit = 0b111; + } + else + { + offset[0] = tmpLocal.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + offset[1] = tmpLocal.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2; + offset[CONSTANT_TWO] = 0; + offset[CONSTANT_THREE] = 0; + params.elementLengths[CONSTANT_TWO] = 0; + params.elementLengths[CONSTANT_THREE] = 0; + params.validBit = 0b11; + } + srcList.src1 = sortedInGroupTensor[offset[0]]; + srcList.src2 = sortedInGroupTensor[offset[1]]; + srcList.src3 = sortedInGroupTensor[offset[CONSTANT_TWO]]; + srcList.src4 = sortedInGroupTensor[offset[CONSTANT_THREE]]; + MrgSort(mrgSortTensor[(kGroup_ - 1 - i) * perGroupExpertCountAlign_ * 2], srcList, params); + } + else + { + offset[0] = tmpLocal.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + DataCopy(mrgSortTensor[(kGroup_ - 1 - i) * perGroupExpertCountAlign_ * 2], sortedInGroupTensor[offset[0]], + perGroupExpertCountAlign_ * 2); + } + } + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::FinalSortAfterKGroup() + { + LocalTensor srcTensor; + LocalTensor dstTensor; + int32_t sortedBaseRow = CONSTANT_FOUR; + for (int32_t i = 0; i < tilingData_->vmsCount; i++) + { + if (i % CONSTANT_TWO == 0) + { + srcTensor = mrgSortTensor; + dstTensor = sortedInGroupTensor; + } + else + { + srcTensor = sortedInGroupTensor; + dstTensor = mrgSortTensor; + } + int32_t nextBaseRow = sortedBaseRow * CONSTANT_FOUR; + int32_t quotient = kGroup_ / nextBaseRow; + int32_t remainder = kGroup_ - quotient * nextBaseRow; + + if (quotient > 0) + { + MrgSort4Info params; + MrgSortSrcList srcList; + params.ifExhaustedSuspension = false; + params.elementLengths[0] = perGroupExpertCount_ * sortedBaseRow; + params.elementLengths[1] = perGroupExpertCount_ * sortedBaseRow; + params.elementLengths[CONSTANT_TWO] = perGroupExpertCount_ * sortedBaseRow; + params.elementLengths[CONSTANT_THREE] = perGroupExpertCount_ * sortedBaseRow; + params.validBit = 0b1111; + params.repeatTimes = 1; + for (int j = 0; j < quotient; j++) + { + srcList.src1 = srcTensor[perGroupExpertCountAlign_ * sortedBaseRow * 8 * j]; + srcList.src2 = srcTensor[perGroupExpertCountAlign_ * sortedBaseRow * (8 * j + 2)]; + srcList.src3 = srcTensor[perGroupExpertCountAlign_ * sortedBaseRow * (8 * j + 4)]; + srcList.src4 = srcTensor[perGroupExpertCountAlign_ * sortedBaseRow * (8 * j + 6)]; + MrgSort(dstTensor[perGroupExpertCountAlign_ * sortedBaseRow * 8 * j], srcList, params); + } + } + if (remainder > 0) + { + int32_t baseOffset = quotient * nextBaseRow * perGroupExpertCountAlign_ * 2; + int32_t mrgLen = CeilDiv(remainder, sortedBaseRow); + int32_t tailRow = remainder - (mrgLen - 1) * sortedBaseRow; + if (mrgLen > 1) + { + MrgSort4Info params; + MrgSortSrcList srcList; + params.repeatTimes = 1; + params.ifExhaustedSuspension = false; + params.elementLengths[0] = perGroupExpertCount_ * sortedBaseRow; + params.elementLengths[1] = perGroupExpertCount_ * sortedBaseRow; + params.elementLengths[CONSTANT_TWO] = perGroupExpertCount_ * sortedBaseRow; + params.elementLengths[CONSTANT_THREE] = perGroupExpertCount_ * sortedBaseRow; + srcList.src1 = srcTensor[baseOffset]; + srcList.src2 = srcTensor[baseOffset + perGroupExpertCountAlign_ * sortedBaseRow * 2]; + if (mrgLen == CONSTANT_FOUR) + { + srcList.src3 = srcTensor[baseOffset + perGroupExpertCountAlign_ * sortedBaseRow * 2 * 2]; + srcList.src4 = srcTensor[baseOffset + perGroupExpertCountAlign_ * sortedBaseRow * 2 * 3]; + params.elementLengths[CONSTANT_THREE] = perGroupExpertCount_ * tailRow; + params.validBit = 0b1111; + } + else if (mrgLen == CONSTANT_THREE) + { + srcList.src3 = srcTensor[baseOffset + perGroupExpertCountAlign_ * sortedBaseRow * 2 * 2]; + params.elementLengths[CONSTANT_TWO] = perGroupExpertCount_ * tailRow; + params.elementLengths[CONSTANT_THREE] = 0; + params.validBit = 0b111; + } + else + { + params.elementLengths[1] = perGroupExpertCount_ * tailRow; + params.elementLengths[CONSTANT_TWO] = 0; + params.elementLengths[CONSTANT_THREE] = 0; + params.validBit = 0b11; + } + MrgSort(dstTensor[baseOffset], srcList, params); + } + else + { + DataCopy(dstTensor[baseOffset], srcTensor[baseOffset], tailRow * perGroupExpertCountAlign_ * 2); + } + } + + sortedBaseRow = nextBaseRow; + } + } + + template + __aicore__ inline void + MoeGatingTopKHashRegbase::TopKGatherWithSmallKAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor) + { + __local_mem__ float *inputAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ uint32_t *mrgSortAddr = (__local_mem__ uint32_t *)mrgSortTensor.GetPhyAddr(); + __local_mem__ uint32_t *expertIdxAddr = (__local_mem__ uint32_t *)expertIdxTensor.GetPhyAddr(); + __local_mem__ T *outputAddr = (__local_mem__ T *)yTensor.GetPhyAddr(); + + if (tilingData_->normType == 0) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg0, vreg1, mrgSortAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Muls(vreg2, vreg2, routedScalingFactor, preg0); + StoreOneTensorForDtypeT(outputAddr, vreg2, preg0, 0); + MicroAPI::DataCopy(expertIdxAddr, vreg1, preg0); + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vreg4; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg0, vreg1, mrgSortAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::ReduceSum(vreg3, vreg2, preg0); + MicroAPI::Adds(vreg3, vreg3, eps, preg0); + MicroAPI::Duplicate(vreg4, vreg3, preg0); + MicroAPI::Div(vreg4, vreg2, vreg4, preg0); + MicroAPI::Muls(vreg4, vreg4, routedScalingFactor, preg0); + StoreOneTensorForDtypeT(outputAddr, vreg4, preg0, 0); + MicroAPI::DataCopy(expertIdxAddr, vreg1, preg0); + } + } + } + + template + __aicore__ inline void + MoeGatingTopKHashRegbase::TopKGatherWithLargeKAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor) + { + uint32_t k1 = k_; + __local_mem__ float *inputAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ uint32_t *mrgSortAddr = (__local_mem__ uint32_t *)mrgSortTensor.GetPhyAddr(); + __local_mem__ uint32_t *expertIdxAddr = (__local_mem__ uint32_t *)expertIdxTensor.GetPhyAddr(); + __local_mem__ T *outputAddr = (__local_mem__ T *)yTensor.GetPhyAddr(); + + if (tilingData_->normType == 0) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vregSum; + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg preg1 = MicroAPI::CreateMask(); + MicroAPI::Duplicate(vregSum, static_cast(0), preg0); + uint16_t vfLoopNum = static_cast(CeilDiv(k, VL_FLOAT_SIZE)); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::Duplicate(vreg2, static_cast(0), preg1); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Add(vregSum, vregSum, vreg2, preg1); + } + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg1 = MicroAPI::UpdateMask(k1); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg1); + MicroAPI::Muls(vreg2, vreg2, routedScalingFactor, preg1); + StoreOneTensorForDtypeT(outputAddr, vreg2, preg1, i * VL_FLOAT_SIZE); + MicroAPI::DataCopy(expertIdxAddr + i * VL_FLOAT_SIZE, vreg1, preg1); + } + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vreg4; + RegTensor vreg5; + RegTensor vregSum; + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg preg1 = MicroAPI::CreateMask(); + + MicroAPI::Duplicate(vregSum, static_cast(0), preg0); + uint16_t vfLoopNum = static_cast(CeilDiv(k, VL_FLOAT_SIZE)); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::Duplicate(vreg2, static_cast(0), preg1); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Add(vregSum, vregSum, vreg2, preg1); + } + MicroAPI::ReduceSum(vregSum, vregSum, preg1); + MicroAPI::Adds(vregSum, vregSum, eps, preg1); + MicroAPI::Duplicate(vreg4, vregSum, preg1); + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg1 = MicroAPI::UpdateMask(k1); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg1); + MicroAPI::Div(vreg5, vreg2, vreg4, preg1); + MicroAPI::Muls(vreg5, vreg5, routedScalingFactor, preg1); + StoreOneTensorForDtypeT(outputAddr, vreg5, preg1, i * VL_FLOAT_SIZE); + MicroAPI::DataCopy(expertIdxAddr + i * VL_FLOAT_SIZE, vreg1, preg1); + } + } + } + } + + template + __aicore__ inline void + MoeGatingTopKHashRegbase::TopKGatherWithSmallKNotAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor, int32_t expertIdxPad, + int32_t perGroupExpertCountAlign) + { + __local_mem__ float *inputAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ uint32_t *mrgSortAddr = (__local_mem__ uint32_t *)mrgSortTensor.GetPhyAddr(); + __local_mem__ uint32_t *expertIdxAddr = (__local_mem__ uint32_t *)expertIdxTensor.GetPhyAddr(); + __local_mem__ T *outputAddr = (__local_mem__ T *)yTensor.GetPhyAddr(); + + if (tilingData_->normType == 0) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vregAlign; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + + MicroAPI::Duplicate(vregAlign, perGroupExpertCountAlign, preg0); + + MicroAPI::DataCopy(vreg0, vreg1, mrgSortAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Muls(vreg2, vreg2, routedScalingFactor, preg0); + + // compute expertIdx: id = id - floor_div(id, perGroupExpertCountAlign) * pad + MicroAPI::Div(vregAlign, vreg1, vregAlign, preg0); + MicroAPI::Muls(vregAlign, vregAlign, expertIdxPad, preg0); + MicroAPI::Sub(vreg1, vreg1, vregAlign, preg0); + + StoreOneTensorForDtypeT(outputAddr, vreg2, preg0, 0); + MicroAPI::DataCopy(expertIdxAddr, vreg1, preg0); + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vreg4; + RegTensor vregAlign; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + + MicroAPI::Duplicate(vregAlign, perGroupExpertCountAlign, preg0); + + MicroAPI::DataCopy(vreg0, vreg1, mrgSortAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::ReduceSum(vreg3, vreg2, preg0); + MicroAPI::Adds(vreg3, vreg3, eps, preg0); + MicroAPI::Duplicate(vreg4, vreg3, preg0); + MicroAPI::Div(vreg4, vreg2, vreg4, preg0); + MicroAPI::Muls(vreg4, vreg4, routedScalingFactor, preg0); + + // compute expertIdx: id = id - floor_div(id, perGroupExpertCountAlign) * pad + MicroAPI::Div(vregAlign, vreg1, vregAlign, preg0); + MicroAPI::Muls(vregAlign, vregAlign, expertIdxPad, preg0); + MicroAPI::Sub(vreg1, vreg1, vregAlign, preg0); + + StoreOneTensorForDtypeT(outputAddr, vreg4, preg0, 0); + MicroAPI::DataCopy(expertIdxAddr, vreg1, preg0); + } + } + } + + template + __aicore__ inline void + MoeGatingTopKHashRegbase::TopKGatherWithLargeKNotAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor, int32_t expertIdxPad, + int32_t perGroupExpertCountAlign) + { + uint32_t k1 = k_; + __local_mem__ float *inputAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ uint32_t *mrgSortAddr = (__local_mem__ uint32_t *)mrgSortTensor.GetPhyAddr(); + __local_mem__ uint32_t *expertIdxAddr = (__local_mem__ uint32_t *)expertIdxTensor.GetPhyAddr(); + __local_mem__ T *outputAddr = (__local_mem__ T *)yTensor.GetPhyAddr(); + + if (tilingData_->normType == 0) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vregSum; + RegTensor vregAlign; + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg preg1 = MicroAPI::CreateMask(); + + MicroAPI::Duplicate(vregSum, static_cast(0), preg0); + MicroAPI::Duplicate(vregAlign, perGroupExpertCountAlign, preg0); + + uint16_t vfLoopNum = static_cast(CeilDiv(k, VL_FLOAT_SIZE)); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::Duplicate(vreg2, static_cast(0), preg1); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Add(vregSum, vregSum, vreg2, preg1); + } + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg1 = MicroAPI::UpdateMask(k1); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg1); + MicroAPI::Muls(vreg2, vreg2, routedScalingFactor, preg1); + + // compute expertIdx: id = id - floor_div(id, perGroupExpertCountAlign) * pad + MicroAPI::Div(vreg3, vreg1, vregAlign, preg1); + MicroAPI::Muls(vreg3, vreg3, expertIdxPad, preg1); + MicroAPI::Sub(vreg1, vreg1, vreg3, preg1); + StoreOneTensorForDtypeT(outputAddr, vreg2, preg1, i * VL_FLOAT_SIZE); + MicroAPI::DataCopy(expertIdxAddr + i * VL_FLOAT_SIZE, vreg1, preg1); + } + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vreg4; + RegTensor vreg5; + RegTensor vreg6; + RegTensor vregSum; + RegTensor vregAlign; + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg preg1 = MicroAPI::CreateMask(); + + MicroAPI::Duplicate(vregSum, static_cast(0), preg0); + MicroAPI::Duplicate(vregAlign, perGroupExpertCountAlign, preg0); + + uint16_t vfLoopNum = static_cast(CeilDiv(k, VL_FLOAT_SIZE)); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::Duplicate(vreg2, static_cast(0), preg1); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Add(vregSum, vregSum, vreg2, preg1); + } + MicroAPI::ReduceSum(vregSum, vregSum, preg1); + MicroAPI::Adds(vregSum, vregSum, eps, preg1); + MicroAPI::Duplicate(vreg4, vregSum, preg1); + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg1 = MicroAPI::UpdateMask(k1); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg1); + MicroAPI::Div(vreg6, vreg2, vreg4, preg1); + MicroAPI::Muls(vreg6, vreg6, routedScalingFactor, preg1); + + // compute expertIdx: id = id - floor_div(id, perGroupExpertCountAlign) * pad + MicroAPI::Div(vreg5, vreg1, vregAlign, preg1); + MicroAPI::Muls(vreg5, vreg5, expertIdxPad, preg1); + MicroAPI::Sub(vreg1, vreg1, vreg5, preg1); + StoreOneTensorForDtypeT(outputAddr, vreg6, preg1, i * VL_FLOAT_SIZE); + MicroAPI::DataCopy(expertIdxAddr + i * VL_FLOAT_SIZE, vreg1, preg1); + } + } + } + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::SelectTopKExpertScore() + { + LocalTensor expertIdxTensor = expertIdxOutQueue_.AllocTensor(); + LocalTensor mrgSortTensor = finalSortBuffer_.Get(); + + if (tilingData_->vmsCount % CONSTANT_TWO == 1) + { + mrgSortTensor = sortedInGroupTensor.ReinterpretCast(); + } + + LocalTensor xSigmoidTensor = xSigmoidBuf_.Get(); + LocalTensor yTensor = yOutQueue_.AllocTensor(); + + int32_t expertIdxPad = perGroupExpertCountAlign_ - perGroupExpertCount_; + if (k_ <= VL_FLOAT_SIZE) + { + if (expertIdxPad != 0) + { + TopKGatherWithSmallKNotAlignE(xSigmoidTensor, mrgSortTensor, expertIdxTensor, yTensor, k_, eps_, routedScalingFactor_, + expertIdxPad, perGroupExpertCountAlign_); + } + else + { + TopKGatherWithSmallKAlignE(xSigmoidTensor, mrgSortTensor, expertIdxTensor, yTensor, k_, eps_, routedScalingFactor_); + } + } + else + { + if (expertIdxPad != 0) + { + TopKGatherWithLargeKNotAlignE(xSigmoidTensor, mrgSortTensor, expertIdxTensor, yTensor, k_, eps_, routedScalingFactor_, + expertIdxPad, perGroupExpertCountAlign_); + } + else + { + TopKGatherWithLargeKAlignE(xSigmoidTensor, mrgSortTensor, expertIdxTensor, yTensor, k_, eps_, routedScalingFactor_); + } + } + + yOutQueue_.EnQue(yTensor); + expertIdxOutQueue_.EnQue(expertIdxTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::CopyOut(int64_t row) + { + LocalTensor yOutTensor = yOutQueue_.DeQue(); + LocalTensor expertIdxTensor = expertIdxOutQueue_.DeQue(); + DataCopyExtParams dataCopyParams{1, static_cast(k_ * sizeof(T)), 0, 0, 0}; + DataCopyPad(yGm_[row * k_], yOutTensor, dataCopyParams); + + dataCopyParams.blockLen = k_ * sizeof(int32_t); + DataCopyPad(expertIdxGm_[row * k_], expertIdxTensor, dataCopyParams); + + expertIdxOutQueue_.FreeTensor(expertIdxTensor); + yOutQueue_.FreeTensor(yOutTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::Init(GM_ADDR x, GM_ADDR bias, GM_ADDR inputIds, GM_ADDR tid2eid, GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, + GM_ADDR workspace, + const MoeGatingTopKHashRegbaseTilingData *tilingData, TPipe *tPipe) + { + tilingData_ = tilingData; + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + if (blockIdx_ == GetBlockNum() - 1) + { + curCoreRowCount_ = tilingData_->lastCoreRowCount; + } + else + { + curCoreRowCount_ = tilingData_->perCoreRowCount; + } + expertCount_ = tilingData_->expertCount; + k_ = tilingData_->k; + kGroup_ = tilingData_->kGroup; + groupCount_ = tilingData_->groupCount; + perGroupExpertCount_ = tilingData_->perGroupExpertCount; + perGroupExpertCountAlign_ = tilingData_->perGroupExpertCountAlign; + expertCountAlign_ = Align(perGroupExpertCountAlign_ * groupCount_, sizeof(float)); + routedScalingFactor_ = tilingData_->routedScalingFactor; + groupSelectMode_ = tilingData_->groupSelectMode; + eps_ = tilingData_->eps; + hashFlag_ = tilingData_->hashFlag == 1; + + // init input gm buf + xGm_.SetGlobalBuffer((__gm__ T *)x + tilingData_->perCoreRowCount * expertCount_ * blockIdx_, expertCount_); + if (bias != nullptr) + { + hasBias_ = true; + biasGm_.SetGlobalBuffer((__gm__ T *)bias, expertCount_); + } + yGm_.SetGlobalBuffer((__gm__ T *)y + tilingData_->perCoreRowCount * k_ * blockIdx_, k_); + expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx + tilingData_->perCoreRowCount * k_ * blockIdx_, k_); + outGm_.SetGlobalBuffer((__gm__ float *)out + tilingData_->perCoreRowCount * expertCount_ * blockIdx_, expertCount_); + + inputIdsGm_.SetGlobalBuffer((__gm__ U1 *)inputIds); + tid2eidGm_.SetGlobalBuffer((__gm__ U2 *)tid2eid); + + // init queue + int32_t expertGroupAlign = groupCount_ * perGroupExpertCountAlign_; + int32_t groupAlign = static_cast(CeilAlign(groupCount_, ONE_REPEAT_SORT_NUM)); + pipe_->InitBuffer(xInQueue_, CONSTANT_TWO, expertGroupAlign * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(yOutQueue_, CONSTANT_TWO, AlignBytes(k_, sizeof(T))); + pipe_->InitBuffer(expertIdxOutQueue_, CONSTANT_TWO, AlignBytes(k_, sizeof(int32_t))); + pipe_->InitBuffer(outOutQueue_, CONSTANT_TWO, expertGroupAlign * sizeof(float)); + + pipe_->InitBuffer(biasBuf_, expertGroupAlign * sizeof(T)); + pipe_->InitBuffer(xSigmoidBuf_, expertGroupAlign * sizeof(float)); + pipe_->InitBuffer(xBiasBuf_, expertGroupAlign * sizeof(float)); + pipe_->InitBuffer(indexBuffer_, expertGroupAlign * sizeof(int32_t)); + pipe_->InitBuffer(sortedInGroupBuf_, expertGroupAlign * sizeof(U2) * 2); + pipe_->InitBuffer(finalSortBuffer_, expertGroupAlign * sizeof(float) * 2); + pipe_->InitBuffer(groupBuf_, groupAlign * sizeof(float)); + pipe_->InitBuffer(sortedGroupBuf_, groupAlign * sizeof(float) * 2); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::Process() + { + CopyInBias(); + if (kGroup_ == groupCount_ || groupCount_ == expertCount_) + { + CopyInX(0); + for (int64_t row = 1; row < curCoreRowCount_; row++) + { + ComputeX(); + CopyOutXNorm(row - 1); + CopyInX(row); + if (hashFlag_) + { + HashCompute(row - 1 + tilingData_->perCoreRowCount * blockIdx_); + } + else + { + TopKCompute(); + } + CopyOut(row - 1); + } + ComputeX(); + CopyOutXNorm(curCoreRowCount_ - 1); + if (hashFlag_) + { + HashCompute(curCoreRowCount_ - 1 + tilingData_->perCoreRowCount * blockIdx_); + } + else + { + TopKCompute(); + } + CopyOut(curCoreRowCount_ - 1); + return; + } + + CopyInX(0); + for (int64_t row = 1; row < curCoreRowCount_; row++) + { + ComputeX(); + CopyOutXNorm(row - 1); + SortInGroup(); + SelectTopKGroupIndex(); + CopyInX(row); + FinalSortByKGroup(); + FinalSortAfterKGroup(); + SelectTopKExpertScore(); + CopyOut(row - 1); + } + ComputeX(); + CopyOutXNorm(curCoreRowCount_ - 1); + SortInGroup(); + SelectTopKGroupIndex(); + FinalSortByKGroup(); + FinalSortAfterKGroup(); + SelectTopKExpertScore(); + CopyOut(curCoreRowCount_ - 1); + } +} // namespace MoeGatingTopKHash +#endif // MOE_GATING_TOP_K_REGBASE_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/common.h b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/common.h new file mode 100644 index 000000000..89cc2cd1b --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/common.h @@ -0,0 +1,90 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file common.h + * \brief + */ +#ifndef MOE_GATING_TOP_K_COMMON_H +#define MOE_GATING_TOP_K_COMMON_H + +#include "kernel_operator.h" + +namespace MoeGatingTopKHash { +using namespace AscendC; +const float MIN_FP32 = *(float *)(&F32_NEG_INF); +constexpr int32_t FLOAT32_NEG_INF = 0xFF800000; // -inf -2139095040 +constexpr int64_t ONE_REPEAT_SORT_NUM = 32; +constexpr int64_t BLOCK_BYTES = 32; +constexpr int64_t REPEAT_BYTES = 256; +constexpr int64_t REPEAT_BLOCKS = 8; + +constexpr int32_t CONSTANT_TWO = 2; +constexpr int32_t CONSTANT_THREE = 3; +constexpr int32_t CONSTANT_FOUR = 4; +constexpr int32_t CONSTANT_EIGHT = 8; + +constexpr int64_t MERGE_LIST_TWO = 2; +constexpr int64_t MERGE_LIST_THREE = 3; +constexpr int64_t MERGE_LIST_FOUR = 4; + +constexpr int64_t MERGE_LIST_IDX_TWO = 2; +constexpr int64_t MERGE_LIST_IDX_THREE = 3; + +constexpr int64_t NORM_TYPE_SOFTMAX = 0; +constexpr int64_t NORM_TYPE_SIGMOID = 1; +constexpr int64_t NORM_TYPE_SOFTPLUS = 2; + +__aicore__ inline int64_t Ceil(int64_t a, int64_t b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes) +{ + if (bytes == 0) { + return 0; + } + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes; +} + +__aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes) +{ + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES; +} + +template +__aicore__ inline T Min(T a, T b) +{ + return a > b ? b : a; +} + +template +__aicore__ inline T Max(T a, T b) +{ + return a < b ? b : a; +} + +template +__aicore__ inline T1 CeilDiv(T1 x, T2 y) +{ + if (y != 0 && x != 0) { + const T1 quotient = x / y; + return (x % y != 0 && ((x ^ y) >= 0)) ? (quotient + 1) : quotient; + } + + return x; +} + +} // namespace MoeGatingTopKHash +#endif // MOE_GATING_TOP_K_COMMON_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/common_regbase.h b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/common_regbase.h new file mode 100644 index 000000000..9438c4afa --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/common_regbase.h @@ -0,0 +1,214 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file common.h + * \brief + */ +#ifndef MOE_GATING_TOP_K_COMMON_H_V35 +#define MOE_GATING_TOP_K_COMMON_H_V35 + +#include "kernel_operator.h" + +namespace MoeGatingTopKHashRegbaseNS { +using namespace AscendC; +constexpr int32_t MIN_FP32 = 0xFF800000; +constexpr int64_t ONE_REPEAT_SORT_NUM = 32; +constexpr int64_t BLOCK_BYTES = 32; + +constexpr int64_t MERGE_LIST_TWO = 2; +constexpr int64_t MERGE_LIST_THREE = 3; +constexpr int64_t MERGE_LIST_FOUR = 4; + +constexpr int64_t MERGE_LIST_IDX_TWO = 2; +constexpr int64_t MERGE_LIST_IDX_THREE = 3; + +__aicore__ inline int64_t Ceil(int64_t a, int64_t b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes) +{ + if (bytes == 0) { + return 0; + } + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes; +} + +__aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes) +{ + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES; +} + +template +__aicore__ inline T Min(T a, T b) +{ + return a > b ? b : a; +} + +template +__aicore__ inline T Max(T a, T b) +{ + return a < b ? b : a; +} + +template +__aicore__ inline T1 CeilAlign(T1 a, T2 b) { + if (b == 0) { + return 0; + } + return (a + b - 1) / b * b; +} + +template +__aicore__ inline T1 CeilDiv(T1 x, T2 y) +{ + if (y != 0 && x != 0) { + const T1 quotient = x / y; + return (x % y != 0 && ((x ^ y) >= 0)) ? (quotient + 1) : quotient; + } + + return x; +} + +template +__aicore__ inline void SetWaitFlag(HardEvent evt) +{ + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(evt)); + SetFlag(eventId); + WaitFlag(eventId); +} + +constexpr AscendC::MicroAPI::CastTrait castTraitB162B32 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::UNKNOWN, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB322B16 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_RINT, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB322Int32 = { + AscendC::MicroAPI::RegLayout::UNKNOWN, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_TRUNC, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB322Int16 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_TRUNC, +}; + +constexpr AscendC::MicroAPI::CastTrait castTraitB162Int8 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_TRUNC, +}; + +// load 对齐的 bfloat16,float16,bfloat32类型的 input(ub中)数据到 float32类型的dst(寄存器)中 +template +__aicore__ inline void LoadOneTensorForDtypeT(__local_mem__ T *input, MicroAPI::RegTensor &dst, + MicroAPI::MaskReg &preg, uint32_t offset) +{ + if constexpr (IsSameType::value) { + MicroAPI::RegTensor xFp16; + DataCopy(xFp16, ((__local_mem__ half *)(input) + (offset))); + Cast(dst, xFp16, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor xBf16; + DataCopy(xBf16, + ((__local_mem__ bfloat16_t *)(input) + (offset))); + Cast(dst, xBf16, preg); + } else { + DataCopy(dst, ((__local_mem__ float *)(input) + (offset))); + } +} + +// load 2个对齐的Tensor 到寄存器中 +template +__aicore__ inline void LoadTwoTensorForDtypeT(__local_mem__ T *src1, __local_mem__ T *src2, + MicroAPI::RegTensor &dst1, MicroAPI::RegTensor &dst2, + MicroAPI::MaskReg &dst1Preg, MicroAPI::MaskReg &dst2Preg, + uint32_t src1Offset, uint32_t src2Offset) +{ + if constexpr (IsSameType::value) { + MicroAPI::RegTensor xFp16Q; + MicroAPI::RegTensor xFp16R; + DataCopy(xFp16Q, ((__local_mem__ half *)(src1) + (src1Offset))); + DataCopy(xFp16R, ((__local_mem__ half *)(src2) + (src2Offset))); + Cast(dst1, xFp16Q, dst1Preg); + Cast(dst2, xFp16R, dst2Preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor xFp16Q; + MicroAPI::RegTensor xFp16R; + DataCopy(xFp16Q, ((__local_mem__ bfloat16_t *)(src1) + (src1Offset))); + DataCopy(xFp16R, ((__local_mem__ bfloat16_t *)(src2) + (src2Offset))); + Cast(dst1, xFp16Q, dst1Preg); + Cast(dst2, xFp16R, dst2Preg); + } else { + DataCopy(dst1, ((__local_mem__ float *)(src1) + (src1Offset))); + DataCopy(dst2, ((__local_mem__ float *)(src2) + (src2Offset))); + } +} + +// store 对齐的float32类型的src(寄存器)数据到output(ub)中，output数据类型支持bfloat16,float16,bfloat32,int32_t,int16_t,int8_t,uint8_t +template +__aicore__ inline void StoreOneTensorForDtypeT(__local_mem__ T *output, MicroAPI::RegTensor &src, + MicroAPI::MaskReg &preg, uint32_t offset) +{ + if constexpr (IsSameType::value) { + MicroAPI::RegTensor yFp16; + Cast(yFp16, src, preg); + DataCopy(((__local_mem__ half *)output + offset), yFp16, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor xBf16; + Cast(xBf16, src, preg); + DataCopy(((__local_mem__ bfloat16_t *)output + offset), + xBf16, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor zInt32; + MicroAPI::Cast(zInt32, src, preg); + DataCopy(((__local_mem__ int32_t *)output + offset), zInt32, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor zInt16; + MicroAPI::Cast(zInt16, src, preg); + DataCopy(((__local_mem__ int16_t *)output + offset), zInt16, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor yFp16; + MicroAPI::RegTensor zInt8; + Cast(yFp16, src, preg); + Cast(zInt8, yFp16, preg); + DataCopy(((__local_mem__ int8_t *)output + offset), zInt8, preg); + } else if constexpr (IsSameType::value) { + MicroAPI::RegTensor yFp16; + MicroAPI::RegTensor zUint8; + Cast(yFp16, src, preg); + Cast(zUint8, yFp16, preg); + DataCopy(((__local_mem__ uint8_t *)output + offset), zUint8, preg); + } else { + DataCopy(((__local_mem__ float *)output + offset), src, preg); + } +} + +} // namespace MoeGatingTopKHash +#endif // MOE_GATING_TOP_K_COMMON_H_V35 \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash.cpp b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash.cpp new file mode 100644 index 000000000..807713b1c --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash.cpp @@ -0,0 +1,132 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_hash.cpp + * \brief + */ + +#include "moe_gating_top_k_hash_e_k_fullload.h" +#include "moe_gating_top_k_hash_without_group.h" +#include "moe_gating_top_k_hash_generalized.h" +#if defined(__DAV_C310__) + #include "moe_gating_top_k_hash_regbase.h" + using namespace MoeGatingTopKHashRegbaseNS; +#endif +#define TILING_KEY_PER_GROUP_COUNT_32 0 +#define TILING_KEY_WITHOUT_GROUP 1 +#define TILING_KEY_GENERALIZED 2 +#define TILING_KEY_WITHOUT_GROUP_1 3 +#define TILING_KEY_WITHOUT_GROUP_2 4 +#define TILING_KEY_WITHOUT_GROUP_3 5 +#define TILING_KEY_WITHOUT_GROUP_4 6 +#define TILING_KEY_REGBASE 10000 +#define TILING_KEY_REGBASE_1 10001 +#define TILING_KEY_REGBASE_2 10002 +#define TILING_KEY_REGBASE_3 10003 +#define TILING_KEY_REGBASE_4 10004 + +using namespace AscendC; +using namespace MoeGatingTopKHash; +extern "C" __global__ __aicore__ void moe_gating_top_k_hash(GM_ADDR x, GM_ADDR bias, GM_ADDR inputIds, GM_ADDR tid2eid, GM_ADDR y, GM_ADDR expertIdx, + GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling) +{ + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); + if (g_coreType == AIC) { + return; + } + + if (workspace == nullptr) { + return; + } + + GM_ADDR userWS = GetUserWorkspace(workspace); + if (userWS == nullptr) { + return; + } + + TPipe tPipe; + if (TILING_KEY_IS(TILING_KEY_PER_GROUP_COUNT_32)) { + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling); + const MoeGatingTopKHashTilingData *__restrict t = &tilingData; + MoeGatingTopKHashEKFullload op; + op.Init(x, bias, y, expertIdx, out, userWS, t, &tPipe); + op.Process(); + } else if (TILING_KEY_IS(TILING_KEY_WITHOUT_GROUP)) { + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling); + const MoeGatingTopKHashTilingData *__restrict t = &tilingData; + MoeGatingTopKHashWithoutGroup op; + op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, t, &tPipe); + op.Process(); + } else if (TILING_KEY_IS(TILING_KEY_WITHOUT_GROUP_1)) { + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling); + const MoeGatingTopKHashTilingData *__restrict t = &tilingData; + MoeGatingTopKHashWithoutGroup op; + op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, t, &tPipe); + op.Process(); + } else if (TILING_KEY_IS(TILING_KEY_WITHOUT_GROUP_2)) { + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling); + const MoeGatingTopKHashTilingData *__restrict t = &tilingData; + MoeGatingTopKHashWithoutGroup op; + op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, t, &tPipe); + op.Process(); + } else if (TILING_KEY_IS(TILING_KEY_WITHOUT_GROUP_3)) { + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling); + const MoeGatingTopKHashTilingData *__restrict t = &tilingData; + MoeGatingTopKHashWithoutGroup op; + op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, t, &tPipe); + op.Process(); + } else if (TILING_KEY_IS(TILING_KEY_WITHOUT_GROUP_4)) { + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling); + const MoeGatingTopKHashTilingData *__restrict t = &tilingData; + MoeGatingTopKHashWithoutGroup op; + op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, t, &tPipe); + op.Process(); + } else if (TILING_KEY_IS(TILING_KEY_GENERALIZED)) { + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashTilingData, tilingData, tiling); + const MoeGatingTopKHashTilingData *__restrict t = &tilingData; + MoeGatingTopKHashGenerlized op; + op.Init(x, bias, y, expertIdx, out, userWS, t, &tPipe); + op.Process(); + } + #if defined(__DAV_C310__) + else if (TILING_KEY_IS(TILING_KEY_REGBASE)) { + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashRegbaseTilingData, tiling_data_in, tiling); + const MoeGatingTopKHashRegbaseTilingData *__restrict tilingData = &tiling_data_in; + MoeGatingTopKHashRegbase op; + op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe); + op.Process(); + } else if (TILING_KEY_IS(TILING_KEY_REGBASE_1)) { + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashRegbaseTilingData, tiling_data_in, tiling); + const MoeGatingTopKHashRegbaseTilingData *__restrict tilingData = &tiling_data_in; + MoeGatingTopKHashRegbase op; + op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe); + op.Process(); + } else if (TILING_KEY_IS(TILING_KEY_REGBASE_2)) { + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashRegbaseTilingData, tiling_data_in, tiling); + const MoeGatingTopKHashRegbaseTilingData *__restrict tilingData = &tiling_data_in; + MoeGatingTopKHashRegbase op; + op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe); + op.Process(); + } else if (TILING_KEY_IS(TILING_KEY_REGBASE_3)) { + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashRegbaseTilingData, tiling_data_in, tiling); + const MoeGatingTopKHashRegbaseTilingData *__restrict tilingData = &tiling_data_in; + MoeGatingTopKHashRegbase op; + op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe); + op.Process(); + } else if (TILING_KEY_IS(TILING_KEY_REGBASE_4)) { + GET_TILING_DATA_WITH_STRUCT(MoeGatingTopKHashRegbaseTilingData, tiling_data_in, tiling); + const MoeGatingTopKHashRegbaseTilingData *__restrict tilingData = &tiling_data_in; + MoeGatingTopKHashRegbase op; + op.Init(x, bias, inputIds, tid2eid, y, expertIdx, out, userWS, tilingData, &tPipe); + op.Process(); + } + #endif +} diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash_e_k_fullload.h b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash_e_k_fullload.h new file mode 100644 index 000000000..f756dc8fe --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash_e_k_fullload.h @@ -0,0 +1,404 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_hash_e_k_fullload.h + * \brief + */ +#ifndef MOE_GATING_TOP_K_E_K_FULLLOAD_H +#define MOE_GATING_TOP_K_E_K_FULLLOAD_H +#include "kernel_operator.h" +#include "common.h" +namespace MoeGatingTopKHash { +using namespace AscendC; + +template +class MoeGatingTopKHashEKFullload { +public: + __aicore__ inline MoeGatingTopKHashEKFullload(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKHashTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyInBias(); + __aicore__ inline void CopyInX(int64_t progress); + __aicore__ inline void ComputeX(); + __aicore__ inline void SortInGroup(); + __aicore__ inline void SelectTopKGroupIndex(); + __aicore__ inline void SelectTopKExpertIdx(); + __aicore__ inline void SelectTopKExpertScore(); + __aicore__ inline void CopyOut(int64_t progress); + +private: + TPipe *pipe_; + TQue xInQueue_; + TBuf biasInQueue_; + TQue yOutQueue_; + TQue expertIdxOutQueue_; + TQue outOutQueue_; + + TQue xBiasQueue_; + TQue xSigmoidQueue_; + TQue sigmoidTmpQueue_; + TQue sortedInGroupQueue_; + TQue sortedGroupQueue_; + TBuf calcTmpBuffer_; + + GlobalTensor xGm_; + GlobalTensor biasGm_; + GlobalTensor yGm_; + GlobalTensor expertIdxGm_; + GlobalTensor outGm_; + + int64_t blockIdx_; + int64_t perCoreRowCount_; + int64_t curCoreRowCount_; + int64_t expertCount_; + bool addBias_; + int64_t k_; + int64_t kGroup_; + int64_t groupCount_; + int64_t groupSelectMode_; + int64_t renorm_; + int64_t normType_; + int64_t outFlag_; + float routedScalingFactor_; + float eps_; + + int64_t expertCountAlign_; + int64_t kAlign_; + int64_t perGroupExpertCount_; + + const MoeGatingTopKHashTilingData *tilingData_; +}; + +template +__aicore__ inline void MoeGatingTopKHashEKFullload::CopyInBias() +{ + LocalTensor biasTensor = biasInQueue_.Get(); + DataCopyExtParams dataCopyParams{1, static_cast(expertCount_ * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if constexpr (IsSameType::value) { + DataCopyPad(biasTensor, biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + } else { + DataCopyPad(biasTensor[expertCountAlign_].ReinterpretCast(), biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(biasTensor, biasTensor[expertCountAlign_].ReinterpretCast(), RoundMode::CAST_NONE, expertCount_); + } +} + +template +__aicore__ inline void MoeGatingTopKHashEKFullload::CopyInX(int64_t row) +{ + LocalTensor xInLocalTensor = xInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(expertCount_ * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if constexpr (IsSameType::value) { + DataCopyPad(xInLocalTensor, xGm_[row * expertCount_], dataCopyParams, dataCopyPadParams); + } else { + DataCopyPad(xInLocalTensor[expertCountAlign_].ReinterpretCast(), xGm_[row * expertCount_], dataCopyParams, + dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(xInLocalTensor, xInLocalTensor[expertCountAlign_].ReinterpretCast(), RoundMode::CAST_NONE, + expertCount_); + } + + xInQueue_.EnQue(xInLocalTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashEKFullload::ComputeX() +{ + LocalTensor xSigmoidTensor = xSigmoidQueue_.AllocTensor(); + LocalTensor xInLocalTensor = xInQueue_.DeQue(); + LocalTensor xBiasTensor = xBiasQueue_.AllocTensor(); + LocalTensor biasTensor = biasInQueue_.Get(); + LocalTensor sharedTmpBuffer = sigmoidTmpQueue_.AllocTensor(); // 临时空间可以复用 + Sigmoid(xSigmoidTensor, xInLocalTensor, sharedTmpBuffer, expertCount_); + PipeBarrier(); + if (addBias_) { + Add(xBiasTensor, xSigmoidTensor, biasTensor, expertCount_); + } else { + Adds(xBiasTensor, xSigmoidTensor, static_cast(0), expertCount_); + } + + xSigmoidQueue_.EnQue(xSigmoidTensor); + xBiasQueue_.EnQue(xBiasTensor); + xInQueue_.FreeTensor(xInLocalTensor); + sigmoidTmpQueue_.FreeTensor(sharedTmpBuffer); +} + +template +__aicore__ inline void MoeGatingTopKHashEKFullload::SortInGroup() +{ + LocalTensor xBiasTensor = xBiasQueue_.DeQue(); + LocalTensor sortedInGroupTensor = sortedInGroupQueue_.AllocTensor(); // 组内排序的结果, 后续归并需要 + LocalTensor indexTensor = calcTmpBuffer_.Get(); // 用于存储排序时的索引 + ArithProgression(indexTensor.ReinterpretCast(), 0, 1, expertCount_); // 生成组索引0 1 2 ...... + PipeBarrier(); + Sort32(sortedInGroupTensor, xBiasTensor, indexTensor, expertCount_ / ONE_REPEAT_SORT_NUM); // 组内排序 + sortedInGroupQueue_.EnQue(sortedInGroupTensor); + xBiasQueue_.FreeTensor(xBiasTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashEKFullload::SelectTopKGroupIndex() +{ + LocalTensor sortedInGroupTensor = sortedInGroupQueue_.DeQue(); + LocalTensor indexTensor = calcTmpBuffer_.Get(); + LocalTensor top2ValueInGroupTensor = sigmoidTmpQueue_.AllocTensor(); // 这个临时空间可以复用 + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + + indexTensor.SetValue(0, static_cast(5)); // b0101 + indexTensor.SetValue(1, static_cast(0)); + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + + uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数 + GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = 8; + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = 8; + gatherMaskParams.src1RepeatStride = 0; + GatherMask(top2ValueInGroupTensor, sortedInGroupTensor, indexTensor, true, static_cast(64), + gatherMaskParams, rsvdCnt); + PipeBarrier(); + LocalTensor groupTop2SumTensor = top2ValueInGroupTensor; + PairReduceSum(groupTop2SumTensor, top2ValueInGroupTensor, 1, groupCount_ * 2, 1, 1, + 1); // 计算每个组内最大的两个数之和 + PipeBarrier(); + + LocalTensor groupIndexTensor = indexTensor; + ArithProgression(groupIndexTensor.ReinterpretCast(), 0, 1, groupCount_); // 生成组索引 + PipeBarrier(); + // 用最小值补到32个数 + int64_t duplicateNum = ONE_REPEAT_SORT_NUM - groupCount_; + if (duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX << groupCount_; + uint64_t mask[2] = {mask0, 0}; + Duplicate(groupTop2SumTensor, MIN_FP32, mask, 1, 1, 8); + PipeBarrier(); + } + // 排序，将kgroup选出来 + LocalTensor sortedGroupTensor = sortedGroupQueue_.AllocTensor(); + Sort32(sortedGroupTensor, groupTop2SumTensor, groupIndexTensor, 1); + + PipeBarrier(); + LocalTensor sortedGroupIndexTensor = indexTensor.ReinterpretCast(); + // 提取组序号 + uint8_t src1Pattern = 2; // 内置固定模式 + GatherMask(sortedGroupIndexTensor, sortedGroupTensor.template ReinterpretCast(), src1Pattern, false, + static_cast(0), {1, 1, 0, 0}, rsvdCnt); + + // 需要将组排序(这里是降序，所以下mrgsor的时候反着取，3、2、1、0) + Cast(sortedGroupTensor, sortedGroupIndexTensor, RoundMode::CAST_ROUND, kGroup_); + PipeBarrier(); + duplicateNum = ONE_REPEAT_SORT_NUM - kGroup_; + if (duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX << kGroup_; + uint64_t mask[2] = {mask0, 0}; + Duplicate(sortedGroupTensor, MIN_FP32, mask, 1, 1, 8); + PipeBarrier(); + } + Sort32(top2ValueInGroupTensor, sortedGroupTensor, sortedGroupIndexTensor.template ReinterpretCast(), 1); + PipeBarrier(); + src1Pattern = 1; + GatherMask(sortedGroupTensor, top2ValueInGroupTensor, src1Pattern, false, static_cast(0), {1, 1, 0, 0}, + rsvdCnt); + PipeBarrier(); + Cast(sortedGroupIndexTensor, sortedGroupTensor, RoundMode::CAST_ROUND, kGroup_); + + sortedGroupQueue_.FreeTensor(sortedGroupTensor); + sortedInGroupQueue_.EnQue(sortedInGroupTensor); + sigmoidTmpQueue_.FreeTensor(top2ValueInGroupTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashEKFullload::SelectTopKExpertIdx() +{ + LocalTensor expertIdxTensor = expertIdxOutQueue_.AllocTensor(); + LocalTensor topKGroupIndexTensor = calcTmpBuffer_.Get(); + LocalTensor sortedInGroupTensor = sortedInGroupQueue_.DeQue(); + LocalTensor sortedExpertTensor = xInQueue_.AllocTensor(); + AscendC::MrgSort4Info params; + params.elementLengths[0] = k_; + params.elementLengths[1] = k_; + params.elementLengths[2] = k_; + params.elementLengths[3] = k_; + params.ifExhaustedSuspension = true; + params.validBit = 0b1111; + params.repeatTimes = 1; + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + int64_t listOffset1 = topKGroupIndexTensor.GetValue(3) * perGroupExpertCount_ * 2; + int64_t listOffset2 = topKGroupIndexTensor.GetValue(2) * perGroupExpertCount_ * 2; + int64_t listOffset3 = topKGroupIndexTensor.GetValue(1) * perGroupExpertCount_ * 2; + int64_t listOffset4 = topKGroupIndexTensor.GetValue(0) * perGroupExpertCount_ * 2; + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + AscendC::MrgSortSrcList srcList; + srcList.src1 = sortedInGroupTensor[listOffset1]; + srcList.src2 = sortedInGroupTensor[listOffset2]; + srcList.src3 = sortedInGroupTensor[listOffset3]; + srcList.src4 = sortedInGroupTensor[listOffset4]; + MrgSort(sortedExpertTensor, srcList, params); + PipeBarrier(); + uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数 + uint8_t src1Pattern = 2; // 内置固定模式 + GatherMask(expertIdxTensor, sortedExpertTensor.template ReinterpretCast(), src1Pattern, false, + static_cast(0), {1, 1, 0, 0}, rsvdCnt); + xInQueue_.FreeTensor(sortedExpertTensor); + expertIdxOutQueue_.EnQue(expertIdxTensor); + sortedInGroupQueue_.FreeTensor(sortedInGroupTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashEKFullload::SelectTopKExpertScore() +{ + LocalTensor expertIdxTensor = expertIdxOutQueue_.DeQue(); + LocalTensor expertByteIdxTensor = calcTmpBuffer_.Get(); + LocalTensor xSigmoidTensor = xSigmoidQueue_.DeQue(); + LocalTensor yTensor = yOutQueue_.AllocTensor(); + LocalTensor yOutTensor; + if constexpr (!IsSameType::value) { + yOutTensor = yTensor.template ReinterpretCast()[kAlign_]; + } else { + yOutTensor = yTensor; + } + Muls(expertByteIdxTensor, expertIdxTensor, static_cast(sizeof(float)), k_); + PipeBarrier(); + Gather(yOutTensor, xSigmoidTensor, expertByteIdxTensor.template ReinterpretCast(), + static_cast(0), k_); + + LocalTensor calTensor = calcTmpBuffer_.Get(); + PipeBarrier(); + ReduceSum(calTensor, yOutTensor, xSigmoidTensor, k_); + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float sumValue = calTensor.GetValue(0) + eps_; + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Duplicate(calTensor, sumValue, k_); + PipeBarrier(); + Div(yOutTensor, yOutTensor, calTensor, k_); + PipeBarrier(); + Muls(yOutTensor, yOutTensor, routedScalingFactor_, k_); + + if constexpr (!IsSameType::value) { + PipeBarrier(); + Cast(yTensor, yOutTensor, RoundMode::CAST_RINT, k_); + } + + xSigmoidQueue_.EnQue(xSigmoidTensor); + expertIdxOutQueue_.EnQue(expertIdxTensor); + yOutQueue_.EnQue(yTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashEKFullload::CopyOut(int64_t row) +{ + LocalTensor yOutTensor = yOutQueue_.DeQue(); + LocalTensor expertIdxTensor = expertIdxOutQueue_.DeQue(); + LocalTensor xSigmoidTensor = xSigmoidQueue_.DeQue(); + DataCopyExtParams dataCopyParams{1, static_cast(k_ * sizeof(T)), 0, 0, 0}; + DataCopyPad(yGm_[row * k_], yOutTensor, dataCopyParams); + dataCopyParams.blockLen = k_ * sizeof(int32_t); + DataCopyPad(expertIdxGm_[row * k_], expertIdxTensor, dataCopyParams); + xSigmoidQueue_.FreeTensor(xSigmoidTensor); + expertIdxOutQueue_.FreeTensor(expertIdxTensor); + yOutQueue_.FreeTensor(yOutTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashEKFullload::Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, + GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKHashTilingData *tilingData, TPipe *tPipe) +{ + tilingData_ = tilingData; + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + perCoreRowCount_ = tilingData_->perCoreRowCount; + if (blockIdx_ == GetBlockNum() - 1) { + curCoreRowCount_ = tilingData_->lastCoreRowCount; + } else { + curCoreRowCount_ = tilingData_->perCoreRowCount; + } + expertCount_ = tilingData_->expertCount; + addBias_ = tilingData_->addBias == 1; + k_ = tilingData_->k; + kGroup_ = tilingData_->kGroup; + groupCount_ = tilingData_->groupCount; + perGroupExpertCount_ = tilingData_->perGroupExpertCount; + routedScalingFactor_ = tilingData_->routedScalingFactor; + eps_ = tilingData_->eps; + + expertCountAlign_ = Align(expertCount_, sizeof(float)); + kAlign_ = Align(expertCount_, sizeof(float)); + + // init input gm buf + xGm_.SetGlobalBuffer((__gm__ T *)x + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_); + biasGm_.SetGlobalBuffer((__gm__ T *)bias, expertCount_); + + // init output gm buf + yGm_.SetGlobalBuffer((__gm__ T *)y + perCoreRowCount_ * k_ * blockIdx_, k_); + expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx + perCoreRowCount_ * k_ * blockIdx_, k_); + outGm_.SetGlobalBuffer((__gm__ T *)out + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_); + + // init que + pipe_->InitBuffer(xInQueue_, 2, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(biasInQueue_, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + + pipe_->InitBuffer(xSigmoidQueue_, 1, AlignBytes(expertCount_, sizeof(float))); + pipe_->InitBuffer(xBiasQueue_, 2, AlignBytes(expertCount_, sizeof(float))); + + pipe_->InitBuffer(yOutQueue_, 2, kAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(expertIdxOutQueue_, 2, AlignBytes(k_, sizeof(int32_t))); + pipe_->InitBuffer(outOutQueue_, 2, AlignBytes(expertCount_, sizeof(float))); + + pipe_->InitBuffer(sigmoidTmpQueue_, 2, AlignBytes(expertCount_, sizeof(float))); + pipe_->InitBuffer(sortedInGroupQueue_, 2, AlignBytes(expertCount_, sizeof(float)) * 2); + pipe_->InitBuffer(sortedGroupQueue_, 2, + (groupCount_ + ONE_REPEAT_SORT_NUM - 1) / ONE_REPEAT_SORT_NUM * ONE_REPEAT_SORT_NUM * + sizeof(float) * 2); + + pipe_->InitBuffer(calcTmpBuffer_, tilingData_->calTmpBufUbSize); +} + +template +__aicore__ inline void MoeGatingTopKHashEKFullload::Process() +{ + CopyInBias(); + for (int64_t row = 0; row < curCoreRowCount_; row++) { + CopyInX(row); + ComputeX(); + SortInGroup(); + SelectTopKGroupIndex(); + SelectTopKExpertIdx(); + SelectTopKExpertScore(); + CopyOut(row); + } +} +} // namespace MoeGatingTopKHash +#endif // MOE_GATING_TOP_K_E_K_FULLLOAD_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash_generalized.h b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash_generalized.h new file mode 100644 index 000000000..f71738770 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash_generalized.h @@ -0,0 +1,664 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_hash_generalized.h + * \brief + */ +#ifndef MOE_GATING_TOP_K_E_K_GENERALIZED_H +#define MOE_GATING_TOP_K_E_K_GENERALIZED_H +#include "kernel_operator.h" +#include "common.h" +#include "kernel_utils.h" +namespace MoeGatingTopKHash { +using namespace AscendC; + +template +class MoeGatingTopKHashGenerlized { +public: + __aicore__ inline MoeGatingTopKHashGenerlized(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKHashTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyInBiasAndInitExpertId(); + __aicore__ inline void CopyInX(int64_t progress); + __aicore__ inline void ComputeX(); + __aicore__ inline void CopuOutXNorm(int64_t row); + __aicore__ inline void SortInGroup(); + __aicore__ inline void SelectTopKGroupIndex(); + __aicore__ inline void SelectTopKExpertIdx(); + __aicore__ inline void SelectTopKExpertScore(); + __aicore__ inline void CumputeActualTopKExpertId(); + __aicore__ inline void CopyOut(int64_t row); + +private: + TPipe *pipe_; + TQue xInQueue_; + TQue yOutQueue_; + TQue expertIdxOutQueue_; + TQue outOutQueue_; + + TBuf biasBuf_; // 存放输入bias + TBuf expertIdBuf_; // 专家编号 + TBuf xNormWithBiasBuf_; // 存放加了bias之后的值 + TBuf xNormBuf_; // 存放计算sigmoid或softmax的值 + TBuf sortedInGroupBuf_; // 存放组内排序后的结果 + TBuf topKExpertIdBuf_; + TBuf sortedGroupIndexBuf_; + TBuf calcTmpBuf_; + + GlobalTensor xGm_; + GlobalTensor biasGm_; + GlobalTensor yGm_; + GlobalTensor expertIdxGm_; + GlobalTensor outGm_; + + int64_t blockIdx_ = 0; + int64_t perCoreRowCount_ = 0; + int64_t curCoreRowCount_ = 0; + int64_t expertCount_ = 0; + bool addBias_ = false; + int64_t k_ = 0; + int64_t kGroup_ = 0; + int64_t groupCount_ = 0; + int64_t groupCountAlign_ = 0; + int64_t perGroupExpertCount_ = 0; + int64_t perGroupExpertCountAlign_ = 0; + int64_t groupSelectMode_ = 0; + int64_t renorm_ = 0; + int64_t normType_ = 0; + int64_t outFlag_ = 0; + + int64_t expertCountAlign_ = 0; + int64_t kAlign_ = 0; + bool isAlign_ = false; + + const MoeGatingTopKHashTilingData *tilingData_; +}; + +template +__aicore__ inline void MoeGatingTopKHashGenerlized::CopyInBiasAndInitExpertId() +{ + LocalTensor biasTensor = biasBuf_.Get(); + LocalTensor expertIdTensor = expertIdBuf_.Get(); + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = groupCount_; + dataCopyParams.blockLen = perGroupExpertCount_ * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(T) / BLOCK_BYTES; + + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if (addBias_) { + if constexpr (IsSameType::value) { + DataCopyPad(biasTensor, biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + } else { + DataCopyPad(biasTensor[expertCountAlign_].ReinterpretCast(), biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(biasTensor, biasTensor[expertCountAlign_].ReinterpretCast(), RoundMode::CAST_NONE, + expertCountAlign_); + PipeBarrier(); + } + + if (!isAlign_) { + int64_t duplicateNum = perGroupExpertCount_ % ONE_REPEAT_SORT_NUM; + int duplicateIndex = perGroupExpertCount_ - duplicateNum; + if (duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(biasTensor.ReinterpretCast()[duplicateIndex], FLOAT32_NEG_INF, mask, groupCount_, 1, + perGroupExpertCountAlign_ * sizeof(float) / BLOCK_BYTES); + } + } + } + ArithProgression(expertIdTensor, static_cast(0), static_cast(1), expertCountAlign_); +} + +template +__aicore__ inline void MoeGatingTopKHashGenerlized::CopyInX(int64_t row) +{ + LocalTensor xInLocalTensor = xInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = groupCount_; + dataCopyParams.blockLen = perGroupExpertCount_ * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(T) / BLOCK_BYTES; + + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if constexpr (IsSameType::value) { + DataCopyPad(xInLocalTensor, xGm_[row * expertCount_], dataCopyParams, dataCopyPadParams); + } else { + DataCopyPad(xInLocalTensor[expertCountAlign_].ReinterpretCast(), xGm_[row * expertCount_], dataCopyParams, + dataCopyPadParams); + } + xInQueue_.EnQue(xInLocalTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashGenerlized::ComputeX() +{ + LocalTensor xNormTensor = xNormBuf_.Get(); + LocalTensor xInLocalTensor = xInQueue_.DeQue(); + LocalTensor xNormWithBiasTensor = xNormWithBiasBuf_.Get(); + LocalTensor biasTensor = biasBuf_.Get(); + + if constexpr (!IsSameType::value) { + Cast(xInLocalTensor, xInLocalTensor[expertCountAlign_].ReinterpretCast(), RoundMode::CAST_NONE, + expertCountAlign_); + PipeBarrier(); + } + + int64_t duplicateNum = perGroupExpertCount_ % ONE_REPEAT_SORT_NUM; + int duplicateIndex = perGroupExpertCount_ - duplicateNum; + if (!isAlign_ && duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(xInLocalTensor.ReinterpretCast()[duplicateIndex], FLOAT32_NEG_INF, mask, groupCount_, 1, + (perGroupExpertCountAlign_ * sizeof(float)) / BLOCK_BYTES); + PipeBarrier(); + } + if (tilingData_->normType == 1) { // sigmoid + LocalTensor calcNormTmpTensor = calcTmpBuf_.Get(); + Sigmoid(xNormTensor, xInLocalTensor, calcNormTmpTensor, expertCountAlign_); + PipeBarrier(); + } else { // softmax + LocalTensor reduceValueTensor = calcTmpBuf_.Get(); + LocalTensor calcTmp = calcTmpBuf_.Get()[BLOCK_BYTES]; + ReduceMax(reduceValueTensor, xInLocalTensor, calcTmp, expertCountAlign_); + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float maxValue = reduceValueTensor.GetValue(0); + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Adds(xNormTensor, xInLocalTensor, -maxValue, expertCountAlign_); + PipeBarrier(); + Exp(xNormTensor, xNormTensor, expertCountAlign_); + PipeBarrier(); + ReduceSum(reduceValueTensor, xNormTensor, calcTmp, expertCountAlign_); + eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float sumValue = reduceValueTensor.GetValue(0); + eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Muls(xNormTensor, xNormTensor, 1.0f / sumValue, expertCountAlign_); + PipeBarrier(); + } + if (addBias_) { + Add(xNormWithBiasTensor, xNormTensor, biasTensor, expertCountAlign_); + } else { + DataCopy(xNormWithBiasTensor, xNormTensor, expertCountAlign_); + } + + if (!isAlign_ && duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + PipeBarrier(); + Duplicate(xNormWithBiasTensor.ReinterpretCast()[duplicateIndex], + FLOAT32_NEG_INF, // MIN_FP32, + mask, groupCount_, 1, perGroupExpertCountAlign_ * sizeof(float) / BLOCK_BYTES); + } + xInQueue_.FreeTensor(xInLocalTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashGenerlized::CopuOutXNorm(int64_t row) +{ + LocalTensor outOutTensor = outOutQueue_.AllocTensor(); + LocalTensor xNormTensor = xNormBuf_.Get(); + DataCopy(outOutTensor, xNormTensor, expertCountAlign_); + outOutQueue_.EnQue(outOutTensor); + outOutTensor = outOutQueue_.DeQue(); + DataCopyExtParams dataCopyParams{ + static_cast(groupCount_), static_cast(perGroupExpertCount_ * sizeof(float)), + static_cast((perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(float) / BLOCK_BYTES), 0, 0}; + DataCopyPad(outGm_[row * expertCount_], outOutTensor, dataCopyParams); + outOutQueue_.FreeTensor(outOutTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashGenerlized::SortInGroup() +{ + LocalTensor xNormWithBiasTensor = xNormWithBiasBuf_.Get(); + LocalTensor expertIdTensor = expertIdBuf_.Get(); + LocalTensor sortedInGroupTensor = sortedInGroupBuf_.Get(); + LocalTensor tmpLocal = calcTmpBuf_.Get(); + if (perGroupExpertCountAlign_ == ONE_REPEAT_SORT_NUM) { + PipeBarrier(); + Sort32(sortedInGroupTensor, xNormWithBiasTensor, expertIdTensor, groupCount_); + } else { + for (int64_t group = 0; group < groupCount_; group++) { + PipeBarrier(); + Sort(sortedInGroupTensor[group * perGroupExpertCountAlign_ * CONSTANT_TWO], + xNormWithBiasTensor[group * perGroupExpertCountAlign_], + expertIdTensor[group * perGroupExpertCountAlign_], tmpLocal, + perGroupExpertCountAlign_ / ONE_REPEAT_SORT_NUM); + } + } +} + +template +__aicore__ inline void MoeGatingTopKHashGenerlized::SelectTopKGroupIndex() +{ + LocalTensor sortedInGroupTensor = sortedInGroupBuf_.Get(); + LocalTensor valueSelectedFromGroupTensor = calcTmpBuf_.GetWithOffset(groupCountAlign_ * 2, 0); + LocalTensor maskTensor = + calcTmpBuf_.GetWithOffset(groupCountAlign_, groupCountAlign_ * 2 * sizeof(float)); + LocalTensor topValueInGroupTensor = + calcTmpBuf_.GetWithOffset(groupCountAlign_, groupCountAlign_ * 3 * sizeof(float)); + LocalTensor groupIndex = + calcTmpBuf_.GetWithOffset(groupCountAlign_, groupCountAlign_ * 4 * sizeof(float)); + LocalTensor sortedTopValue = + calcTmpBuf_.GetWithOffset(groupCountAlign_ * 2, groupCountAlign_ * 5 * sizeof(float)); + LocalTensor sortTmp = + calcTmpBuf_.GetWithOffset(groupCountAlign_ * 2, groupCountAlign_ * 7 * sizeof(float)); + + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + + uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数 + PipeBarrier(); + if (tilingData_->groupSelectMode == 1) { // top2 sum + // 提取每组组前两个元素 + maskTensor.SetValue(0, static_cast(5)); // b0101 + maskTensor.SetValue(1, static_cast(0)); + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + + GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = groupCount_; + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = + Ceil(perGroupExpertCountAlign_ * (sizeof(float) + sizeof(uint32_t)), BLOCK_BYTES); + gatherMaskParams.src1RepeatStride = 0; + GatherMask(valueSelectedFromGroupTensor, sortedInGroupTensor, maskTensor, true, + static_cast(ONE_REPEAT_SORT_NUM * CONSTANT_TWO), gatherMaskParams, rsvdCnt); + PipeBarrier(); + + // 计算每个组前两个数的和 + PairReduceSum(topValueInGroupTensor, valueSelectedFromGroupTensor, + Ceil(groupCount_ * sizeof(float) * 2, REPEAT_BYTES), REPEAT_BYTES / sizeof(float), 1, 1, + CONSTANT_EIGHT); // 计算每个组内最大的两个数之和 + } else { + maskTensor.SetValue(0, static_cast(1)); // b0101 + maskTensor.SetValue(1, static_cast(0)); + + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数 + GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = groupCount_; + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = Ceil(perGroupExpertCountAlign_ * (sizeof(float) + sizeof(uint32_t)), 32); + gatherMaskParams.src1RepeatStride = 0; + GatherMask(topValueInGroupTensor, sortedInGroupTensor, maskTensor, true, + static_cast(ONE_REPEAT_SORT_NUM * CONSTANT_TWO), gatherMaskParams, rsvdCnt); + } + + PipeBarrier(); + // 生成组索引 + ArithProgression(groupIndex.ReinterpretCast(), static_cast(0), static_cast(1), + groupCount_); // 生成组索引 + PipeBarrier(); + + int64_t duplicateNum = groupCount_ % ONE_REPEAT_SORT_NUM; + int duplicateIndex = groupCount_ - duplicateNum; + if (duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(topValueInGroupTensor.ReinterpretCast()[duplicateIndex], FLOAT32_NEG_INF, mask, 1, 1, + REPEAT_BLOCKS); + PipeBarrier(); + } + PipeBarrier(); + + // 排序 + Sort(sortedTopValue, topValueInGroupTensor, groupIndex, sortTmp, Ceil(groupCount_, 32)); + PipeBarrier(); + + // 提取组序号 + uint8_t src1Pattern = 2; // 内置固定模式 + GatherMask(groupIndex, sortedTopValue.template ReinterpretCast(), src1Pattern, false, + static_cast(0), + {1, static_cast(Ceil(kGroup_ * sizeof(float) * CONSTANT_TWO, 256)), REPEAT_BLOCKS, 0}, rsvdCnt); + PipeBarrier(); + duplicateNum = kGroup_ % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + duplicateIndex = kGroup_ - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + PipeBarrier(); + Duplicate(groupIndex.ReinterpretCast()[duplicateIndex], FLOAT32_NEG_INF, mask, 1, 1, REPEAT_BLOCKS); + } + + // 将筛选出来的组序号降序排列 + LocalTensor sortedGroupIndex = sortedGroupIndexBuf_.Get(); + PipeBarrier(); + Sort(sortedGroupIndex, groupIndex.ReinterpretCast(), groupIndex, sortTmp, Ceil(kGroup_, 32)); +} + +template +__aicore__ inline void MoeGatingTopKHashGenerlized::SelectTopKExpertIdx() +{ + LocalTensor sortedInGroupTensor = sortedInGroupBuf_.Get(); + LocalTensor sortedGroupIndex = sortedGroupIndexBuf_.Get(); + LocalTensor topKExpertId = topKExpertIdBuf_.Get(); + LocalTensor mrgSort0Tensor = calcTmpBuf_.Get(); + + uint32_t offset[CONSTANT_FOUR] = {0, 0, 0, 0}; + uint16_t lenArr[CONSTANT_FOUR] = { + static_cast(perGroupExpertCount_), static_cast(perGroupExpertCount_), + static_cast(perGroupExpertCount_), static_cast(perGroupExpertCount_)}; + MrgSort4Info params{lenArr, false, 0b1111, 1}; + MrgSortSrcList srcList; + + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + + for (int32_t i = kGroup_ - 1; i >= 0; i -= CONSTANT_FOUR) { + int64_t mrgLen = Min(i + 1, CONSTANT_FOUR); + if (mrgLen > 1) { + if (mrgLen == MERGE_LIST_FOUR) { + offset[0] = sortedGroupIndex.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + offset[1] = sortedGroupIndex.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2; + offset[2] = sortedGroupIndex.GetValue((i - 2) * 2) * perGroupExpertCountAlign_ * 2; + offset[3] = sortedGroupIndex.GetValue((i - 3) * 2) * perGroupExpertCountAlign_ * 2; + } else if (mrgLen == MERGE_LIST_THREE) { + offset[0] = sortedGroupIndex.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + offset[1] = sortedGroupIndex.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2; + offset[2] = sortedGroupIndex.GetValue((i - 2) * 2) * perGroupExpertCountAlign_ * 2; + offset[3] = 0; + params.elementLengths[3] = 0; + params.validBit = 0b111; + } else { + offset[0] = sortedGroupIndex.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + offset[1] = sortedGroupIndex.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2; + offset[2] = 0; + offset[3] = 0; + params.elementLengths[2] = 0; + params.elementLengths[3] = 0; + params.validBit = 0b11; + } + + srcList.src1 = sortedInGroupTensor[offset[0]]; + srcList.src2 = sortedInGroupTensor[offset[1]]; + srcList.src3 = sortedInGroupTensor[offset[2]]; + srcList.src4 = sortedInGroupTensor[offset[3]]; + + PipeBarrier(); + MrgSort(mrgSort0Tensor[(kGroup_ - 1 - i) * perGroupExpertCountAlign_ * 2], srcList, params); + } else { + offset[0] = sortedGroupIndex.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + PipeBarrier(); + DataCopy(mrgSort0Tensor[(kGroup_ - 1 - i) * perGroupExpertCountAlign_ * 2], sortedInGroupTensor[offset[0]], + perGroupExpertCountAlign_ * 2); + } + } + int32_t baseLoop = 4; + LocalTensor srcTensor = mrgSort0Tensor; + LocalTensor dstTensor = mrgSort0Tensor; + for (int i = 0; i < tilingData_->vmsCount; i++) { + if (i % 2 == 0) { + srcTensor = mrgSort0Tensor; + dstTensor = sortedInGroupTensor; + } else { + srcTensor = sortedInGroupTensor; + dstTensor = mrgSort0Tensor; + } + + int32_t nextBaseRow = baseLoop * MERGE_LIST_FOUR; + int32_t quotient = kGroup_ / nextBaseRow; + int32_t remainder = kGroup_ - quotient * nextBaseRow; + if (quotient > 0) { + MrgSort4Info params; + MrgSortSrcList srcList; + params.ifExhaustedSuspension = false; + params.elementLengths[0] = perGroupExpertCount_ * baseLoop; + params.elementLengths[1] = perGroupExpertCount_ * baseLoop; + params.elementLengths[2] = perGroupExpertCount_ * baseLoop; + params.elementLengths[3] = perGroupExpertCount_ * baseLoop; + params.validBit = 0b1111; + params.repeatTimes = 1; + for (int j = 0; j < quotient; j++) { + srcList.src1 = srcTensor[perGroupExpertCountAlign_ * baseLoop * 8 * j]; + srcList.src2 = srcTensor[perGroupExpertCountAlign_ * baseLoop * (8 * j + 2)]; + srcList.src3 = srcTensor[perGroupExpertCountAlign_ * baseLoop * (8 * j + 4)]; + srcList.src4 = srcTensor[perGroupExpertCountAlign_ * baseLoop * (8 * j + 6)]; + PipeBarrier(); + MrgSort(dstTensor[perGroupExpertCountAlign_ * baseLoop * 8 * j], srcList, params); + } + } + + if (remainder > 0) { + int32_t baseOffset = quotient * nextBaseRow * perGroupExpertCountAlign_ * 2; + int32_t mrgLen = CeilDiv(remainder, baseLoop); + int32_t tailRow = remainder - (mrgLen - 1) * baseLoop; + if (mrgLen > 1) { + MrgSort4Info params; + MrgSortSrcList srcList; + params.repeatTimes = 1; + params.ifExhaustedSuspension = false; + params.elementLengths[0] = perGroupExpertCount_ * baseLoop; + params.elementLengths[1] = perGroupExpertCount_ * baseLoop; + params.elementLengths[2] = perGroupExpertCount_ * baseLoop; + params.elementLengths[3] = perGroupExpertCount_ * baseLoop; + srcList.src1 = srcTensor[baseOffset]; + srcList.src2 = srcTensor[baseOffset + perGroupExpertCountAlign_ * baseLoop * 2]; + if (mrgLen == MERGE_LIST_FOUR) { + srcList.src3 = srcTensor[baseOffset + perGroupExpertCountAlign_ * baseLoop * 2 * 2]; + srcList.src4 = srcTensor[baseOffset + perGroupExpertCountAlign_ * baseLoop * 2 * 3]; + params.elementLengths[3] = perGroupExpertCount_ * tailRow; + params.validBit = 0b1111; + } else if (mrgLen == MERGE_LIST_THREE) { + srcList.src3 = srcTensor[baseOffset + perGroupExpertCountAlign_ * baseLoop * 2 * 2]; + params.elementLengths[2] = perGroupExpertCount_ * tailRow; + params.elementLengths[3] = 0; + params.validBit = 0b111; + } else { + params.elementLengths[1] = perGroupExpertCount_ * tailRow; + params.elementLengths[2] = 0; + params.elementLengths[3] = 0; + params.validBit = 0b11; + } + PipeBarrier(); + MrgSort(dstTensor[baseOffset], srcList, params); + } else { + PipeBarrier(); + DataCopy(dstTensor[baseOffset], srcTensor[baseOffset], tailRow * perGroupExpertCountAlign_ * 2); + } + } + baseLoop = nextBaseRow; + } + + GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = Ceil(k_ * sizeof(float) * 2, REPEAT_BYTES); + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = REPEAT_BLOCKS; + gatherMaskParams.src1RepeatStride = 0; + + uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数 + uint8_t src1Pattern = 2; // 内置固定模式 + PipeBarrier(); + GatherMask(topKExpertId, dstTensor.template ReinterpretCast(), src1Pattern, false, + static_cast(0), gatherMaskParams, rsvdCnt); +} + +template +__aicore__ inline void MoeGatingTopKHashGenerlized::SelectTopKExpertScore() +{ + LocalTensor xNormTensor = xNormBuf_.Get(); + LocalTensor yOutTensor = yOutQueue_.AllocTensor(); + LocalTensor topKExpertId = topKExpertIdBuf_.Get(); + LocalTensor topKExpertIdWithByte = calcTmpBuf_.Get(); + PipeBarrier(); + Muls(topKExpertIdWithByte, topKExpertId, static_cast(sizeof(float)), k_); + PipeBarrier(); + Gather(yOutTensor, xNormTensor, topKExpertIdWithByte.template ReinterpretCast(), static_cast(0), + k_); + + if (tilingData_->normType == 1) { + LocalTensor maxValueTensor = calcTmpBuf_.Get(); + LocalTensor tmpTensor = calcTmpBuf_.Get()[32]; + PipeBarrier(); + ReduceSum(maxValueTensor, yOutTensor, tmpTensor, k_); + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float sumValue = maxValueTensor.GetValue(0) + tilingData_->eps; + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Duplicate(tmpTensor, sumValue, k_); + PipeBarrier(); + Div(yOutTensor, yOutTensor, tmpTensor, k_); + } + PipeBarrier(); + Muls(yOutTensor, yOutTensor, tilingData_->routedScalingFactor, k_); + + if constexpr (!IsSameType::value) { + PipeBarrier(); + Cast(yOutTensor.ReinterpretCast(), yOutTensor, RoundMode::CAST_RINT, k_); + } + + yOutQueue_.EnQue(yOutTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashGenerlized::CumputeActualTopKExpertId() +{ + LocalTensor expertIdxOut = expertIdxOutQueue_.AllocTensor(); + LocalTensor topKExpertId = topKExpertIdBuf_.Get(); + LocalTensor topKExpertIdFp32 = calcTmpBuf_.Get(); + + PipeBarrier(); + Cast(topKExpertIdFp32, topKExpertId, RoundMode::CAST_ROUND, k_); + PipeBarrier(); + Muls(topKExpertIdFp32, topKExpertIdFp32, 1.0f / (float)perGroupExpertCountAlign_, k_); + PipeBarrier(); + Cast(expertIdxOut, topKExpertIdFp32, RoundMode::CAST_TRUNC, k_); + PipeBarrier(); + Muls(expertIdxOut, expertIdxOut, static_cast(perGroupExpertCountAlign_ - perGroupExpertCount_), k_); + PipeBarrier(); + Sub(expertIdxOut, topKExpertId, expertIdxOut, k_); + expertIdxOutQueue_.EnQue(expertIdxOut); +} + +template +__aicore__ inline void MoeGatingTopKHashGenerlized::CopyOut(int64_t row) +{ + LocalTensor yOutTensor = yOutQueue_.DeQue(); + LocalTensor expertIdxOut = expertIdxOutQueue_.DeQue(); + DataCopyExtParams dataCopyParams{1, static_cast(k_ * sizeof(T)), 0, 0, 0}; + DataCopyPad(yGm_[row * k_], yOutTensor, dataCopyParams); + dataCopyParams.blockLen = k_ * sizeof(int32_t); + DataCopyPad(expertIdxGm_[row * k_], expertIdxOut, dataCopyParams); + yOutQueue_.FreeTensor(yOutTensor); + expertIdxOutQueue_.FreeTensor(expertIdxOut); +} + +template +__aicore__ inline void MoeGatingTopKHashGenerlized::Init(GM_ADDR x, GM_ADDR bias, GM_ADDR y, GM_ADDR expertIdx, + GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKHashTilingData *tilingData, TPipe *tPipe) +{ + tilingData_ = tilingData; + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + perCoreRowCount_ = tilingData_->perCoreRowCount; + if (blockIdx_ == GetBlockNum() - 1) { + curCoreRowCount_ = tilingData_->lastCoreRowCount; + } else { + curCoreRowCount_ = tilingData_->perCoreRowCount; + } + expertCount_ = tilingData_->expertCount; + addBias_ = tilingData_->addBias == 1; + k_ = tilingData_->k; + kGroup_ = tilingData_->kGroup; + groupCount_ = tilingData_->groupCount; + groupCountAlign_ = Ceil(groupCount_, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + perGroupExpertCount_ = tilingData_->perGroupExpertCount; + perGroupExpertCountAlign_ = tilingData_->perGroupExpertCountAlign; + + expertCountAlign_ = Align(perGroupExpertCountAlign_ * groupCount_, sizeof(float)); + kAlign_ = Align(k_, sizeof(float)); + + isAlign_ = perGroupExpertCount_ == perGroupExpertCountAlign_; + + // init input gm buf + xGm_.SetGlobalBuffer((__gm__ T *)x + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_); + biasGm_.SetGlobalBuffer((__gm__ T *)bias, expertCount_); + + // init output gm buf + yGm_.SetGlobalBuffer((__gm__ T *)y + perCoreRowCount_ * k_ * blockIdx_, k_); + expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx + perCoreRowCount_ * k_ * blockIdx_, k_); + outGm_.SetGlobalBuffer((__gm__ float *)out + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_); + + // init que + pipe_->InitBuffer(xInQueue_, 1, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(yOutQueue_, 1, kAlign_ * sizeof(float)); + pipe_->InitBuffer(expertIdxOutQueue_, 1, kAlign_ * sizeof(int32_t)); + pipe_->InitBuffer(outOutQueue_, 1, expertCountAlign_ * sizeof(float)); + + pipe_->InitBuffer(biasBuf_, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(expertIdBuf_, expertCountAlign_ * sizeof(int32_t)); + + pipe_->InitBuffer(xNormBuf_, expertCountAlign_ * sizeof(float)); + + pipe_->InitBuffer(xNormWithBiasBuf_, expertCountAlign_ * sizeof(float)); + pipe_->InitBuffer(sortedInGroupBuf_, expertCountAlign_ * (sizeof(float) + sizeof(uint32_t))); + + pipe_->InitBuffer(sortedGroupIndexBuf_, groupCountAlign_ * sizeof(float) * CONSTANT_TWO); + pipe_->InitBuffer(topKExpertIdBuf_, kAlign_ * sizeof(int32_t)); + pipe_->InitBuffer(calcTmpBuf_, expertCountAlign_ * sizeof(float) * 10); +} + +template +__aicore__ inline void MoeGatingTopKHashGenerlized::Process() +{ + CopyInBiasAndInitExpertId(); + for (int64_t row = 0; row < curCoreRowCount_; row++) { + CopyInX(row); + ComputeX(); + if (tilingData_->outFlag) { + CopuOutXNorm(row); + } + SortInGroup(); + SelectTopKGroupIndex(); + SelectTopKExpertIdx(); + SelectTopKExpertScore(); + CumputeActualTopKExpertId(); + CopyOut(row); + } +} +} // namespace MoeGatingTopKHash +#endif // MOE_GATING_TOP_K_E_K_GENERALIZED_H diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash_regbase.h b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash_regbase.h new file mode 100644 index 000000000..ff94cad2e --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash_regbase.h @@ -0,0 +1,1577 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_hash_regbase.h + * \brief + */ +#ifndef MOE_GATING_TOP_K_REGBASE_H +#define MOE_GATING_TOP_K_REGBASE_H + +#include +#include "common_regbase.h" +#include "kernel_operator.h" +#include "kernel_utils.h" + +namespace MoeGatingTopKHashRegbaseNS +{ + using namespace AscendC; + using MicroAPI::RegTensor; + + constexpr int32_t CONSTANT_TWO = 2; + constexpr int32_t CONSTANT_THREE = 3; + constexpr int32_t CONSTANT_FOUR = 4; + constexpr int32_t CONSTANT_EIGHT = 8; + constexpr uint32_t VL_FLOAT_SIZE = VECTOR_REG_WIDTH / sizeof(float); + constexpr MicroAPI::DivSpecificMode mode = {MicroAPI::MaskMergeMode::ZEROING, true}; + + template + class MoeGatingTopKHashRegbase + { + public: + __aicore__ inline MoeGatingTopKHashRegbase(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR bias, GM_ADDR inputIds, GM_ADDR tid2eid, GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKHashRegbaseTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + + private: + __aicore__ inline void CopyInBias(); + __aicore__ inline void CopyInX(int64_t progress); + __aicore__ inline void ComputeX(); + __aicore__ inline void ComputeSoftmax(); + __aicore__ inline void ComputeSigmoid(); + __aicore__ inline void ComputeSoftplus(); + __aicore__ inline void CopyOutXNorm(int64_t progress); + __aicore__ inline void SortInGroup(); + __aicore__ inline void SelectTopKGroupIndex(); + __aicore__ inline void FinalSortByKGroup(); + __aicore__ inline void FinalSortAfterKGroup(); + __aicore__ inline void SelectTopKExpertScore(); + __aicore__ inline void TopKCompute(); + __aicore__ inline void HashCompute(int64_t row); + __aicore__ inline void CopyOut(int64_t progress); + + __aicore__ inline void TopKGatherWithSmallKAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor); + __aicore__ inline void TopKGatherWithLargeKAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor); + __aicore__ inline void TopKGatherWithSmallKNotAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor, int32_t expertIdxPad, + int32_t perGroupExpertCountAlign); + __aicore__ inline void TopKGatherWithLargeKNotAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor, int32_t expertIdxPad, + int32_t perGroupExpertCountAlign); + __aicore__ inline void HashGatherWithSmallKAlignE(LocalTensor xSigmoidTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor); + __aicore__ inline void HashGatherWithSmallKNotAlignE(LocalTensor xSigmoidTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor, int32_t expertIdxPad, + int32_t perGroupExpertCountAlign); + + private: + TPipe *pipe_; + TQue xInQueue_; + TQue yOutQueue_; + TQue expertIdxOutQueue_; + TQue outOutQueue_; + + TBuf biasBuf_; + TBuf xBiasBuf_; + TBuf xSigmoidBuf_; + TBuf groupBuf_; + TBuf sortedInGroupBuf_; + TBuf sortedGroupBuf_; + TBuf indexBuffer_; + TBuf finalSortBuffer_; + + GlobalTensor xGm_; + GlobalTensor biasGm_; + GlobalTensor yGm_; + GlobalTensor expertIdxGm_; + GlobalTensor outGm_; + GlobalTensor inputIdsGm_; + GlobalTensor tid2eidGm_; + + LocalTensor indexTensor; + LocalTensor sortedInGroupTensor; + LocalTensor sortedGroupTensor; + LocalTensor mrgSortTensor; + + int64_t blockIdx_; + int64_t curCoreRowCount_; + int64_t expertCount_; + int64_t k_; + int64_t kGroup_; + int64_t groupCount_; + float routedScalingFactor_; + float eps_; + bool hasBias_ = false; + bool hashFlag_ = false; + + int64_t perGroupExpertCount_; + int64_t perGroupExpertCountAlign_; + int64_t expertCountAlign_; + int64_t groupSelectMode_; + const MoeGatingTopKHashRegbaseTilingData *tilingData_; + }; + + template + __aicore__ inline void MoeGatingTopKHashRegbase::CopyInBias() + { + if (!hasBias_) + { + return; + } + LocalTensor biasTensor = biasBuf_.Get(); + + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = groupCount_; + dataCopyParams.blockLen = perGroupExpertCount_ * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(T) / BLOCK_BYTES; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + + DataCopyPad(biasTensor, biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::CopyInX(int64_t row) + { + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = groupCount_; + dataCopyParams.blockLen = perGroupExpertCount_ * sizeof(T); + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(T) / BLOCK_BYTES; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if (tilingData_->normType == 0) + { + LocalTensor xInLocalTensor = xInQueue_.AllocTensor(); + if constexpr (IsSameType::value) + { + DataCopyPad(xInLocalTensor, xGm_[row * expertCount_], dataCopyParams, dataCopyPadParams); + } + else + { + DataCopyPad(xInLocalTensor[expertCountAlign_].ReinterpretCast(), xGm_[row * expertCount_], dataCopyParams, + dataCopyPadParams); + } + xInQueue_.EnQue(xInLocalTensor); + } + else + { + LocalTensor xInLocalTensor = xInQueue_.AllocTensor(); + DataCopyPad(xInLocalTensor, xGm_[row * expertCount_], dataCopyParams, dataCopyPadParams); + xInQueue_.EnQue(xInLocalTensor); + } + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::ComputeSoftmax() + { + LocalTensor xSoftmaxTensor = xSigmoidBuf_.Get(); + LocalTensor xBiasTensor = xBiasBuf_.Get(); + indexTensor = indexBuffer_.Get(); + + uint32_t size = perGroupExpertCountAlign_ * groupCount_; + uint32_t perGroupExpertCount0 = perGroupExpertCount_; + uint32_t perGroupExpertCountAlign0 = perGroupExpertCountAlign_; + uint16_t groupCount0 = groupCount_; + LocalTensor xInLocalTensor = xInQueue_.DeQue(); + if constexpr (!IsSameType::value) + { + Cast(xInLocalTensor, xInLocalTensor[expertCountAlign_].template ReinterpretCast(), RoundMode::CAST_NONE, + expertCountAlign_); + } + + int64_t duplicateNum = perGroupExpertCount_ % ONE_REPEAT_SORT_NUM; + int duplicateIndex = perGroupExpertCount_ - duplicateNum; + if (perGroupExpertCount_ != perGroupExpertCountAlign_ && duplicateNum > 0) + { + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(xInLocalTensor.ReinterpretCast()[duplicateIndex], MIN_FP32, mask, groupCount_, 1, + (perGroupExpertCountAlign_ * sizeof(float)) / BLOCK_BYTES); + } + + SoftMaxShapeInfo softmaxShapeInfo; + softmaxShapeInfo.srcM = 1; + softmaxShapeInfo.srcK = groupCount_ * perGroupExpertCountAlign_; + softmaxShapeInfo.oriSrcM = 1; + softmaxShapeInfo.oriSrcK = groupCount_ * perGroupExpertCountAlign_; + SoftMax(xSoftmaxTensor, xInLocalTensor, tilingData_->softmaxTilingData, softmaxShapeInfo); + + __local_mem__ float *softmaxOutAddr = (__local_mem__ float *)xSoftmaxTensor.GetPhyAddr(); + __local_mem__ int32_t *indexOutAddr = (__local_mem__ int32_t *)indexTensor.GetPhyAddr(); + __local_mem__ float *addBiasOutAddr = (__local_mem__ float *)xBiasTensor.GetPhyAddr(); + if (hasBias_) + { + LocalTensor biasTensor = biasBuf_.Get(); + __VEC_SCOPE__ + { + RegTensor vregBiasFp32; + RegTensor vregIndex; + RegTensor vregSoftmaxResult; + RegTensor vregBiasResult; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + __local_mem__ T *biasAddr = (__local_mem__ T *)biasTensor.GetPhyAddr(); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + LoadOneTensorForDtypeT(softmaxOutAddr, vregSoftmaxResult, preg0, i * VL_FLOAT_SIZE); + LoadOneTensorForDtypeT(biasAddr, vregBiasFp32, preg0, i * VL_FLOAT_SIZE); + // add bias + MicroAPI::Add(vregBiasResult, vregSoftmaxResult, vregBiasFp32, preg0); + // 使用Arange生成排序索引, 起始值为i乘每个循环的veclen + MicroAPI::Arange(vregIndex, static_cast(i * VL_FLOAT_SIZE)); + MicroAPI::DataCopy(addBiasOutAddr + i * VL_FLOAT_SIZE, vregBiasResult, preg0); + MicroAPI::DataCopy(indexOutAddr + i * VL_FLOAT_SIZE, vregIndex, preg0); + } + MicroAPI::LocalMemBar(); + + // pad min fp32 + RegTensor vregPad; + MicroAPI::UnalignReg u0; + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + auto padUbAddr = addBiasOutAddr + i * perGroupExpertCountAlign0 + perGroupExpertCount0; + MicroAPI::DataCopyUnAlign(padUbAddr, vregPad, u0, perGroupExpertCountAlign0 - perGroupExpertCount0); + MicroAPI::DataCopyUnAlignPost(padUbAddr, u0, 0); + } + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vregIndex; + RegTensor vregSoftmaxResult; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + LoadOneTensorForDtypeT(softmaxOutAddr, vregSoftmaxResult, preg0, i * VL_FLOAT_SIZE); + // 使用Arange生成排序索引, 起始值为i乘每个循环的veclen + MicroAPI::Arange(vregIndex, static_cast(i * VL_FLOAT_SIZE)); + MicroAPI::DataCopy(addBiasOutAddr + i * VL_FLOAT_SIZE, vregSoftmaxResult, preg0); + MicroAPI::DataCopy(indexOutAddr + i * VL_FLOAT_SIZE, vregIndex, preg0); + } + MicroAPI::LocalMemBar(); + + // pad min fp32 + RegTensor vregPad; + MicroAPI::UnalignReg u0; + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + auto padUbAddr = addBiasOutAddr + i * perGroupExpertCountAlign0 + perGroupExpertCount0; + MicroAPI::DataCopyUnAlign(padUbAddr, vregPad, u0, perGroupExpertCountAlign0 - perGroupExpertCount0); + MicroAPI::DataCopyUnAlignPost(padUbAddr, u0, 0); + } + } + } + xInQueue_.FreeTensor(xInLocalTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::ComputeSigmoid() + { + LocalTensor xSigmoidTensor = xSigmoidBuf_.Get(); + LocalTensor xBiasTensor = xBiasBuf_.Get(); + indexTensor = indexBuffer_.Get(); + + uint32_t size = perGroupExpertCountAlign_ * groupCount_; + uint32_t perGroupExpertCount0 = perGroupExpertCount_; + uint32_t perGroupExpertCountAlign0 = perGroupExpertCountAlign_; + uint16_t groupCount0 = groupCount_; + LocalTensor xInLocalTensor = xInQueue_.DeQue(); + if (hasBias_) + { + LocalTensor biasTensor = biasBuf_.Get(); + __VEC_SCOPE__ + { + RegTensor vregBiasFp32; + RegTensor vregIndex; + RegTensor vregSigmoidResult; + RegTensor vregBiasResult; + RegTensor vregOne; + RegTensor vregInFp32; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + + __local_mem__ T *inputAddr = (__local_mem__ T *)xInLocalTensor.GetPhyAddr(); + __local_mem__ T *biasAddr = (__local_mem__ T *)biasTensor.GetPhyAddr(); + __local_mem__ float *sigmoidOutAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ int32_t *indexOutAddr = (__local_mem__ int32_t *)indexTensor.GetPhyAddr(); + __local_mem__ float *addBiasOutAddr = (__local_mem__ float *)xBiasTensor.GetPhyAddr(); + + // sigmoid + MicroAPI::Duplicate(vregOne, static_cast(1), preg0); + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + LoadTwoTensorForDtypeT(inputAddr, biasAddr, vregInFp32, vregBiasFp32, preg0, preg0, + i * VL_FLOAT_SIZE, i * VL_FLOAT_SIZE); + MicroAPI::Muls(vreg1, vregInFp32, static_cast(-1), preg0); + MicroAPI::Exp(vreg2, vreg1, preg0); + MicroAPI::Adds(vreg3, vreg2, static_cast(1), preg0); + MicroAPI::Div(vregSigmoidResult, vregOne, vreg3, preg0); + // add bias + MicroAPI::Add(vregBiasResult, vregSigmoidResult, vregBiasFp32, preg0); + // 使用Arange生成排序索引, 起始值为i乘每个循环的veclen + MicroAPI::Arange(vregIndex, static_cast(i * VL_FLOAT_SIZE)); + MicroAPI::DataCopy(sigmoidOutAddr + i * VL_FLOAT_SIZE, vregSigmoidResult, preg0); + MicroAPI::DataCopy(addBiasOutAddr + i * VL_FLOAT_SIZE, vregBiasResult, preg0); + MicroAPI::DataCopy(indexOutAddr + i * VL_FLOAT_SIZE, vregIndex, preg0); + } + MicroAPI::LocalMemBar(); + + // pad min fp32 + RegTensor vregPad; + MicroAPI::UnalignReg u0; + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + auto padUbAddr = addBiasOutAddr + i * perGroupExpertCountAlign0 + perGroupExpertCount0; + MicroAPI::DataCopyUnAlign(padUbAddr, vregPad, u0, perGroupExpertCountAlign0 - perGroupExpertCount0); + MicroAPI::DataCopyUnAlignPost(padUbAddr, u0, 0); + } + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vregBiasFp32; + RegTensor vregIndex; + RegTensor vregSigmoidResult; + RegTensor vregOne; + RegTensor vregInFp32; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + + __local_mem__ T *inputAddr = (__local_mem__ T *)xInLocalTensor.GetPhyAddr(); + __local_mem__ float *sigmoidOutAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ int32_t *indexOutAddr = (__local_mem__ int32_t *)indexTensor.GetPhyAddr(); + __local_mem__ float *addBiasOutAddr = (__local_mem__ float *)xBiasTensor.GetPhyAddr(); + + // sigmoid + MicroAPI::Duplicate(vregOne, static_cast(1), preg0); + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + LoadOneTensorForDtypeT(inputAddr, vregInFp32, preg0, i * VL_FLOAT_SIZE); + MicroAPI::Muls(vreg1, vregInFp32, static_cast(-1), preg0); + MicroAPI::Exp(vreg2, vreg1, preg0); + MicroAPI::Adds(vreg3, vreg2, static_cast(1), preg0); + MicroAPI::Div(vregSigmoidResult, vregOne, vreg3, preg0); + // 使用Arange生成排序索引, 起始值为i乘每个循环的veclen + MicroAPI::Arange(vregIndex, static_cast(i * VL_FLOAT_SIZE)); + MicroAPI::DataCopy(sigmoidOutAddr + i * VL_FLOAT_SIZE, vregSigmoidResult, preg0); + MicroAPI::DataCopy(addBiasOutAddr + i * VL_FLOAT_SIZE, vregSigmoidResult, preg0); + MicroAPI::DataCopy(indexOutAddr + i * VL_FLOAT_SIZE, vregIndex, preg0); + } + MicroAPI::LocalMemBar(); + + // pad inf + RegTensor vregPad; + MicroAPI::UnalignReg u0; + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + auto padUbAddr = addBiasOutAddr + i * perGroupExpertCountAlign0 + perGroupExpertCount0; + MicroAPI::DataCopyUnAlign(padUbAddr, vregPad, u0, perGroupExpertCountAlign0 - perGroupExpertCount0); + MicroAPI::DataCopyUnAlignPost(padUbAddr, u0, 0); + } + } + } + xInQueue_.FreeTensor(xInLocalTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::ComputeSoftplus() + { + LocalTensor xSoftplusTensor = xSigmoidBuf_.Get(); + LocalTensor xBiasTensor = xBiasBuf_.Get(); + indexTensor = indexBuffer_.Get(); + + uint32_t size = perGroupExpertCountAlign_ * groupCount_; + uint32_t perGroupExpertCount0 = perGroupExpertCount_; + uint32_t perGroupExpertCountAlign0 = perGroupExpertCountAlign_; + uint16_t groupCount0 = groupCount_; + LocalTensor xInLocalTensor = xInQueue_.DeQue(); + if (hasBias_) + { + LocalTensor biasTensor = biasBuf_.Get(); + __VEC_SCOPE__ + { + RegTensor vregBiasFp32; + RegTensor vregIndex; + RegTensor vregSoftplusResult; + RegTensor vregBiasResult; + RegTensor vregOne; + RegTensor vregInFp32; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + + __local_mem__ T *inputAddr = (__local_mem__ T *)xInLocalTensor.GetPhyAddr(); + __local_mem__ T *biasAddr = (__local_mem__ T *)biasTensor.GetPhyAddr(); + __local_mem__ float *softplusOutAddr = (__local_mem__ float *)xSoftplusTensor.GetPhyAddr(); + __local_mem__ int32_t *indexOutAddr = (__local_mem__ int32_t *)indexTensor.GetPhyAddr(); + __local_mem__ float *addBiasOutAddr = (__local_mem__ float *)xBiasTensor.GetPhyAddr(); + + // softplus + MicroAPI::Duplicate(vregOne, static_cast(1), preg0); + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + LoadTwoTensorForDtypeT(inputAddr, biasAddr, vregInFp32, vregBiasFp32, preg0, preg0, + i * VL_FLOAT_SIZE, i * VL_FLOAT_SIZE); + MicroAPI::Exp(vreg1, vregInFp32, preg0); + MicroAPI::Adds(vreg2, vreg1, static_cast(1), preg0); + MicroAPI::Ln(vreg3, vreg2, preg0); + MicroAPI::Sqrt(vregSoftplusResult, vreg3, preg0); + // add bias + MicroAPI::Add(vregBiasResult, vregSoftplusResult, vregBiasFp32, preg0); + // 使用Arange生成排序索引, 起始值为i乘每个循环的veclen + MicroAPI::Arange(vregIndex, static_cast(i * VL_FLOAT_SIZE)); + MicroAPI::DataCopy(softplusOutAddr + i * VL_FLOAT_SIZE, vregSoftplusResult, preg0); + MicroAPI::DataCopy(addBiasOutAddr + i * VL_FLOAT_SIZE, vregBiasResult, preg0); + MicroAPI::DataCopy(indexOutAddr + i * VL_FLOAT_SIZE, vregIndex, preg0); + } + MicroAPI::LocalMemBar(); + + // pad min fp32 + RegTensor vregPad; + MicroAPI::UnalignReg u0; + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + auto padUbAddr = addBiasOutAddr + i * perGroupExpertCountAlign0 + perGroupExpertCount0; + MicroAPI::DataCopyUnAlign(padUbAddr, vregPad, u0, perGroupExpertCountAlign0 - perGroupExpertCount0); + MicroAPI::DataCopyUnAlignPost(padUbAddr, u0, 0); + } + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vregBiasFp32; + RegTensor vregIndex; + RegTensor vregSoftplusResult; + RegTensor vregOne; + RegTensor vregInFp32; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + + __local_mem__ T *inputAddr = (__local_mem__ T *)xInLocalTensor.GetPhyAddr(); + __local_mem__ float *softplusOutAddr = (__local_mem__ float *)xSoftplusTensor.GetPhyAddr(); + __local_mem__ int32_t *indexOutAddr = (__local_mem__ int32_t *)indexTensor.GetPhyAddr(); + __local_mem__ float *addBiasOutAddr = (__local_mem__ float *)xBiasTensor.GetPhyAddr(); + + // sigmoid + MicroAPI::Duplicate(vregOne, static_cast(1), preg0); + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + LoadOneTensorForDtypeT(inputAddr, vregInFp32, preg0, i * VL_FLOAT_SIZE); + MicroAPI::Exp(vreg1, vregInFp32, preg0); + MicroAPI::Adds(vreg2, vreg1, static_cast(1), preg0); + MicroAPI::Ln(vreg3, vreg2, preg0); + MicroAPI::Sqrt(vregSoftplusResult, vreg3, preg0); + // 使用Arange生成排序索引, 起始值为i乘每个循环的veclen + MicroAPI::Arange(vregIndex, static_cast(i * VL_FLOAT_SIZE)); + MicroAPI::DataCopy(softplusOutAddr + i * VL_FLOAT_SIZE, vregSoftplusResult, preg0); + MicroAPI::DataCopy(addBiasOutAddr + i * VL_FLOAT_SIZE, vregSoftplusResult, preg0); + MicroAPI::DataCopy(indexOutAddr + i * VL_FLOAT_SIZE, vregIndex, preg0); + } + MicroAPI::LocalMemBar(); + + // pad inf + RegTensor vregPad; + MicroAPI::UnalignReg u0; + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + auto padUbAddr = addBiasOutAddr + i * perGroupExpertCountAlign0 + perGroupExpertCount0; + MicroAPI::DataCopyUnAlign(padUbAddr, vregPad, u0, perGroupExpertCountAlign0 - perGroupExpertCount0); + MicroAPI::DataCopyUnAlignPost(padUbAddr, u0, 0); + } + } + } + xInQueue_.FreeTensor(xInLocalTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::ComputeX() + { + if (tilingData_->normType == 0) + { + ComputeSoftmax(); + return; + } + if (tilingData_->normType == 1) + { + ComputeSigmoid(); + return; + } + ComputeSoftplus(); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::CopyOutXNorm(int64_t progress) + { + if (tilingData_->outFlag == 0) + { + return; + } + LocalTensor outOutTensor = outOutQueue_.AllocTensor(); + LocalTensor xSigmoidTensor = xSigmoidBuf_.Get(); + DataCopy(outOutTensor, xSigmoidTensor, expertCountAlign_); + outOutQueue_.EnQue(outOutTensor); + outOutTensor = outOutQueue_.DeQue(); + DataCopyExtParams dataCopyParams{ + static_cast(groupCount_), static_cast(perGroupExpertCount_ * sizeof(float)), + static_cast((perGroupExpertCountAlign_ - perGroupExpertCount_) * sizeof(float) / BLOCK_BYTES), 0, 0}; + DataCopyPad(outGm_[progress * expertCount_], outOutTensor, dataCopyParams); + outOutQueue_.FreeTensor(outOutTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::SortInGroup() + { + LocalTensor xBiasTensor = xBiasBuf_.Get(); + LocalTensor sortedInGroupTensor = sortedInGroupBuf_.Get(); // 组内排序的结果, 后续归并需要 + LocalTensor tmpLocal = finalSortBuffer_.Get(); + + if (perGroupExpertCountAlign_ == ONE_REPEAT_SORT_NUM) + { + Sort32(sortedInGroupTensor, xBiasTensor, indexTensor, groupCount_); + } + else + { + for (uint16_t i = 0; i < groupCount_; i++) + { + Sort(sortedInGroupTensor[i * perGroupExpertCountAlign_ * CONSTANT_TWO], + xBiasTensor[i * perGroupExpertCountAlign_], indexTensor[i * perGroupExpertCountAlign_], + tmpLocal, perGroupExpertCountAlign_ / ONE_REPEAT_SORT_NUM); + } + } + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::TopKCompute() + { + LocalTensor xBiasTensor = xBiasBuf_.Get(); + LocalTensor sortedInGroupTensor = sortedInGroupBuf_.Get(); // 组内排序的结果, 后续归并需要 + LocalTensor tmpLocal = finalSortBuffer_.Get(); + + Sort(sortedInGroupTensor, xBiasTensor, indexTensor, tmpLocal, + perGroupExpertCountAlign_ * groupCount_ / ONE_REPEAT_SORT_NUM); + + LocalTensor xSigmoidTensor = xSigmoidBuf_.Get(); + LocalTensor yTensor = yOutQueue_.AllocTensor(); + LocalTensor expertIdxTensor = expertIdxOutQueue_.AllocTensor(); + + LocalTensor sortedInGroupTensorCast = sortedInGroupTensor.template ReinterpretCast(); + + int32_t expertIdxPad = perGroupExpertCountAlign_ - perGroupExpertCount_; + if (k_ <= VL_FLOAT_SIZE) + { + if (expertIdxPad != 0) + { + TopKGatherWithSmallKNotAlignE(xSigmoidTensor, sortedInGroupTensorCast, expertIdxTensor, yTensor, k_, eps_, + routedScalingFactor_, expertIdxPad, perGroupExpertCountAlign_); + } + else + { + TopKGatherWithSmallKAlignE(xSigmoidTensor, sortedInGroupTensorCast, expertIdxTensor, yTensor, k_, eps_, + routedScalingFactor_); + } + } + else + { + if (expertIdxPad != 0) + { + TopKGatherWithLargeKNotAlignE(xSigmoidTensor, sortedInGroupTensorCast, expertIdxTensor, yTensor, k_, eps_, + routedScalingFactor_, expertIdxPad, perGroupExpertCountAlign_); + } + else + { + TopKGatherWithLargeKAlignE(xSigmoidTensor, sortedInGroupTensorCast, expertIdxTensor, yTensor, k_, eps_, + routedScalingFactor_); + } + } + + yOutQueue_.EnQue(yTensor); + expertIdxOutQueue_.EnQue(expertIdxTensor); + } + + template + __aicore__ inline void + MoeGatingTopKHashRegbase::HashGatherWithSmallKAlignE(LocalTensor xSigmoidTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor) + { + __local_mem__ float *inputAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ uint32_t *expertIdxAddr = (__local_mem__ uint32_t *)expertIdxTensor.GetPhyAddr(); + __local_mem__ T *outputAddr = (__local_mem__ T *)yTensor.GetPhyAddr(); + + if (tilingData_->normType == 0) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg1, expertIdxAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Muls(vreg2, vreg2, routedScalingFactor, preg0); + StoreOneTensorForDtypeT(outputAddr, vreg2, preg0, 0); + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vreg4; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg1, expertIdxAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::ReduceSum(vreg3, vreg2, preg0); + MicroAPI::Adds(vreg3, vreg3, eps, preg0); + MicroAPI::Duplicate(vreg4, vreg3, preg0); + MicroAPI::Div(vreg4, vreg2, vreg4, preg0); + MicroAPI::Muls(vreg4, vreg4, routedScalingFactor, preg0); + StoreOneTensorForDtypeT(outputAddr, vreg4, preg0, 0); + } + } + } + + template + __aicore__ inline void + MoeGatingTopKHashRegbase::HashGatherWithSmallKNotAlignE(LocalTensor xSigmoidTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor, int32_t expertIdxPad, + int32_t perGroupExpertCountAlign) + { + __local_mem__ float *inputAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ uint32_t *expertIdxAddr = (__local_mem__ uint32_t *)expertIdxTensor.GetPhyAddr(); + __local_mem__ T *outputAddr = (__local_mem__ T *)yTensor.GetPhyAddr(); + + if (tilingData_->normType == 0) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vregAlign; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + + MicroAPI::Duplicate(vregAlign, perGroupExpertCountAlign, preg0); + + MicroAPI::DataCopy(vreg1, expertIdxAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Muls(vreg2, vreg2, routedScalingFactor, preg0); + + // compute expertIdx: id = id - floor_div(id, perGroupExpertCountAlign) * pad + MicroAPI::Div(vregAlign, vreg1, vregAlign, preg0); + MicroAPI::Muls(vregAlign, vregAlign, expertIdxPad, preg0); + MicroAPI::Sub(vreg1, vreg1, vregAlign, preg0); + + StoreOneTensorForDtypeT(outputAddr, vreg2, preg0, 0); + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vreg4; + RegTensor vregAlign; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + + MicroAPI::Duplicate(vregAlign, perGroupExpertCountAlign, preg0); + + MicroAPI::DataCopy(vreg1, expertIdxAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::ReduceSum(vreg3, vreg2, preg0); + MicroAPI::Adds(vreg3, vreg3, eps, preg0); + MicroAPI::Duplicate(vreg4, vreg3, preg0); + MicroAPI::Div(vreg4, vreg2, vreg4, preg0); + MicroAPI::Muls(vreg4, vreg4, routedScalingFactor, preg0); + + // compute expertIdx: id = id - floor_div(id, perGroupExpertCountAlign) * pad + MicroAPI::Div(vregAlign, vreg1, vregAlign, preg0); + MicroAPI::Muls(vregAlign, vregAlign, expertIdxPad, preg0); + MicroAPI::Sub(vreg1, vreg1, vregAlign, preg0); + + StoreOneTensorForDtypeT(outputAddr, vreg4, preg0, 0); + } + } + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::HashCompute(int64_t row) + { + LocalTensor expertIdxTensor = expertIdxOutQueue_.AllocTensor(); + LocalTensor hashExpertId = sortedInGroupBuf_.Get(); + LocalTensor hashExpertIdInt32 = hashExpertId.template ReinterpretCast(); + + U1 key = inputIdsGm_.GetValue(row); + SetWaitFlag(HardEvent::MTE3_MTE2); + DataCopyExtParams dataCopyParams{1, static_cast(k_ * sizeof(U2)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + DataCopyPad(hashExpertId, tid2eidGm_[key * k_], dataCopyParams, dataCopyPadParams); + SetWaitFlag(HardEvent::MTE2_V); + if constexpr (IsSameType::value) { + DataCopy(expertIdxTensor, hashExpertId, Align(k_, sizeof(int32_t))); + } else { + Cast(hashExpertIdInt32, hashExpertId, RoundMode::CAST_NONE, Align(k_, sizeof(U2))); + PipeBarrier(); + DataCopy(expertIdxTensor, hashExpertIdInt32, Align(k_, sizeof(int32_t))); + } + PipeBarrier(); + + LocalTensor xSigmoidTensor = xSigmoidBuf_.Get(); + LocalTensor yTensor = yOutQueue_.AllocTensor(); + LocalTensor sortedInGroupTensorCast = sortedInGroupTensor.template ReinterpretCast(); + int32_t expertIdxPad = perGroupExpertCountAlign_ - perGroupExpertCount_; + // 当前只支持k小于64场景 + if (k_ <= VL_FLOAT_SIZE) + { + if (expertIdxPad != 0) + { + HashGatherWithSmallKNotAlignE(xSigmoidTensor, expertIdxTensor, yTensor, k_, eps_, + routedScalingFactor_, expertIdxPad, perGroupExpertCountAlign_); + } + else + { + HashGatherWithSmallKAlignE(xSigmoidTensor, expertIdxTensor, yTensor, k_, eps_, + routedScalingFactor_); + } + } + yOutQueue_.EnQue(yTensor); + expertIdxOutQueue_.EnQue(expertIdxTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::SelectTopKGroupIndex() + { + sortedInGroupTensor = sortedInGroupBuf_.Get(); + LocalTensor top2InGroupTensor = groupBuf_.Get(); + LocalTensor tmpLocal = xBiasBuf_.Get(); + // 排序，将kgroup选出来 + sortedGroupTensor = sortedGroupBuf_.Get(); + + uint16_t groupCount0 = groupCount_; + uint32_t perGroupExpertCountAlign0 = perGroupExpertCountAlign_; + int32_t groupCountNumAlign = (groupCount_ + 31) / 32 * 32; + uint32_t padNegInfNum = groupCountNumAlign - groupCount_; + if (groupSelectMode_ == 1) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vregPad; + + MicroAPI::UnalignReg u0; + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + + __local_mem__ float *inputAddr = (__local_mem__ float *)sortedInGroupTensor.GetPhyAddr(); + __local_mem__ float *outputAddr = (__local_mem__ float *)top2InGroupTensor.GetPhyAddr(); + + // pair reduce sum + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + MicroAPI::DataCopy( + vreg0, vreg1, inputAddr + i * perGroupExpertCountAlign0 * 2); + MicroAPI::PairReduceSum(vreg2, vreg0, preg0); + MicroAPI::DataCopyUnAlign(outputAddr, vreg2, u0, 1); + } + MicroAPI::DataCopyUnAlign(outputAddr, vregPad, u0, + padNegInfNum); + MicroAPI::DataCopyUnAlignPost(outputAddr, u0, 0); + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vregPad; + + MicroAPI::UnalignReg u0; + + __local_mem__ float *inputAddr = (__local_mem__ float *)sortedInGroupTensor.GetPhyAddr(); + __local_mem__ float *outputAddr = (__local_mem__ float *)top2InGroupTensor.GetPhyAddr(); + + // max + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + for (uint16_t i = 0; i < groupCount0; i++) + { + MicroAPI::DataCopy( + vreg0, vreg1, inputAddr + i * perGroupExpertCountAlign0 * 2); + MicroAPI::DataCopyUnAlign(outputAddr, vreg0, u0, 1); + } + MicroAPI::DataCopyUnAlign(outputAddr, vregPad, u0, + padNegInfNum); + MicroAPI::DataCopyUnAlignPost(outputAddr, u0, 0); + } + } + + Sort(sortedGroupTensor, top2InGroupTensor, indexTensor, tmpLocal, + groupCountNumAlign / ONE_REPEAT_SORT_NUM); + + uint32_t size = groupCountNumAlign; + int32_t kGroup0 = kGroup_; + int32_t kGroupNumAlign = (kGroup0 + 31) / 32 * 32; + uint32_t padkGroupNum = kGroupNumAlign - kGroup0; + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vregPad; + MicroAPI::UnalignReg u0; + + __local_mem__ int32_t *inputAddr = (__local_mem__ int32_t *)sortedGroupTensor.GetPhyAddr(); + __local_mem__ int32_t *outputAddr = (__local_mem__ int32_t *)top2InGroupTensor.GetPhyAddr(); + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + uint16_t vfLoopNum = static_cast(CeilDiv(size, VL_FLOAT_SIZE)); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(size); + MicroAPI::DataCopy(vreg0, vreg1, + inputAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::DataCopy(outputAddr + i * VL_FLOAT_SIZE, vreg1, preg0); + } + MicroAPI::Duplicate(vregPad, *((float *)&MIN_FP32)); + outputAddr = outputAddr + kGroup0; + MicroAPI::DataCopyUnAlign(outputAddr, (RegTensor &)vregPad, u0, padkGroupNum); + MicroAPI::DataCopyUnAlignPost(outputAddr, u0, 0); + } + + Sort(sortedGroupTensor, top2InGroupTensor, indexTensor, tmpLocal, + kGroupNumAlign / ONE_REPEAT_SORT_NUM); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::FinalSortByKGroup() + { + mrgSortTensor = finalSortBuffer_.Get(); + LocalTensor tmpLocal = sortedGroupTensor.template ReinterpretCast(); + uint32_t offset[MRG_SORT_ELEMENT_LEN] = {0, 0, 0, 0}; + + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + + uint16_t lenArr[CONSTANT_FOUR] = { + static_cast(perGroupExpertCount_), static_cast(perGroupExpertCount_), + static_cast(perGroupExpertCount_), static_cast(perGroupExpertCount_)}; + MrgSort4Info params{lenArr, false, 0b1111, 1}; + MrgSortSrcList srcList; + + for (int32_t i = kGroup_ - 1; i >= 0; i -= CONSTANT_FOUR) + { +#if defined(__CCE_KT_TEST__) + int32_t mrgLen = std::min(i + 1, CONSTANT_FOUR); +#else + int32_t mrgLen = min(i + 1, CONSTANT_FOUR); +#endif + if (mrgLen > 1) + { + if (mrgLen == CONSTANT_FOUR) + { + offset[0] = tmpLocal.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + offset[1] = tmpLocal.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2; + offset[CONSTANT_TWO] = tmpLocal.GetValue((i - 2) * 2) * perGroupExpertCountAlign_ * 2; + offset[CONSTANT_THREE] = tmpLocal.GetValue((i - 3) * 2) * perGroupExpertCountAlign_ * 2; + } + else if (mrgLen == CONSTANT_THREE) + { + offset[0] = tmpLocal.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + offset[1] = tmpLocal.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2; + offset[CONSTANT_TWO] = tmpLocal.GetValue((i - 2) * 2) * perGroupExpertCountAlign_ * 2; + offset[CONSTANT_THREE] = 0; + params.elementLengths[CONSTANT_THREE] = 0; + params.validBit = 0b111; + } + else + { + offset[0] = tmpLocal.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + offset[1] = tmpLocal.GetValue((i - 1) * 2) * perGroupExpertCountAlign_ * 2; + offset[CONSTANT_TWO] = 0; + offset[CONSTANT_THREE] = 0; + params.elementLengths[CONSTANT_TWO] = 0; + params.elementLengths[CONSTANT_THREE] = 0; + params.validBit = 0b11; + } + srcList.src1 = sortedInGroupTensor[offset[0]]; + srcList.src2 = sortedInGroupTensor[offset[1]]; + srcList.src3 = sortedInGroupTensor[offset[CONSTANT_TWO]]; + srcList.src4 = sortedInGroupTensor[offset[CONSTANT_THREE]]; + MrgSort(mrgSortTensor[(kGroup_ - 1 - i) * perGroupExpertCountAlign_ * 2], srcList, params); + } + else + { + offset[0] = tmpLocal.GetValue(i * 2) * perGroupExpertCountAlign_ * 2; + DataCopy(mrgSortTensor[(kGroup_ - 1 - i) * perGroupExpertCountAlign_ * 2], sortedInGroupTensor[offset[0]], + perGroupExpertCountAlign_ * 2); + } + } + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::FinalSortAfterKGroup() + { + LocalTensor srcTensor; + LocalTensor dstTensor; + int32_t sortedBaseRow = CONSTANT_FOUR; + for (int32_t i = 0; i < tilingData_->vmsCount; i++) + { + if (i % CONSTANT_TWO == 0) + { + srcTensor = mrgSortTensor; + dstTensor = sortedInGroupTensor; + } + else + { + srcTensor = sortedInGroupTensor; + dstTensor = mrgSortTensor; + } + int32_t nextBaseRow = sortedBaseRow * CONSTANT_FOUR; + int32_t quotient = kGroup_ / nextBaseRow; + int32_t remainder = kGroup_ - quotient * nextBaseRow; + + if (quotient > 0) + { + MrgSort4Info params; + MrgSortSrcList srcList; + params.ifExhaustedSuspension = false; + params.elementLengths[0] = perGroupExpertCount_ * sortedBaseRow; + params.elementLengths[1] = perGroupExpertCount_ * sortedBaseRow; + params.elementLengths[CONSTANT_TWO] = perGroupExpertCount_ * sortedBaseRow; + params.elementLengths[CONSTANT_THREE] = perGroupExpertCount_ * sortedBaseRow; + params.validBit = 0b1111; + params.repeatTimes = 1; + for (int j = 0; j < quotient; j++) + { + srcList.src1 = srcTensor[perGroupExpertCountAlign_ * sortedBaseRow * 8 * j]; + srcList.src2 = srcTensor[perGroupExpertCountAlign_ * sortedBaseRow * (8 * j + 2)]; + srcList.src3 = srcTensor[perGroupExpertCountAlign_ * sortedBaseRow * (8 * j + 4)]; + srcList.src4 = srcTensor[perGroupExpertCountAlign_ * sortedBaseRow * (8 * j + 6)]; + MrgSort(dstTensor[perGroupExpertCountAlign_ * sortedBaseRow * 8 * j], srcList, params); + } + } + if (remainder > 0) + { + int32_t baseOffset = quotient * nextBaseRow * perGroupExpertCountAlign_ * 2; + int32_t mrgLen = CeilDiv(remainder, sortedBaseRow); + int32_t tailRow = remainder - (mrgLen - 1) * sortedBaseRow; + if (mrgLen > 1) + { + MrgSort4Info params; + MrgSortSrcList srcList; + params.repeatTimes = 1; + params.ifExhaustedSuspension = false; + params.elementLengths[0] = perGroupExpertCount_ * sortedBaseRow; + params.elementLengths[1] = perGroupExpertCount_ * sortedBaseRow; + params.elementLengths[CONSTANT_TWO] = perGroupExpertCount_ * sortedBaseRow; + params.elementLengths[CONSTANT_THREE] = perGroupExpertCount_ * sortedBaseRow; + srcList.src1 = srcTensor[baseOffset]; + srcList.src2 = srcTensor[baseOffset + perGroupExpertCountAlign_ * sortedBaseRow * 2]; + if (mrgLen == CONSTANT_FOUR) + { + srcList.src3 = srcTensor[baseOffset + perGroupExpertCountAlign_ * sortedBaseRow * 2 * 2]; + srcList.src4 = srcTensor[baseOffset + perGroupExpertCountAlign_ * sortedBaseRow * 2 * 3]; + params.elementLengths[CONSTANT_THREE] = perGroupExpertCount_ * tailRow; + params.validBit = 0b1111; + } + else if (mrgLen == CONSTANT_THREE) + { + srcList.src3 = srcTensor[baseOffset + perGroupExpertCountAlign_ * sortedBaseRow * 2 * 2]; + params.elementLengths[CONSTANT_TWO] = perGroupExpertCount_ * tailRow; + params.elementLengths[CONSTANT_THREE] = 0; + params.validBit = 0b111; + } + else + { + params.elementLengths[1] = perGroupExpertCount_ * tailRow; + params.elementLengths[CONSTANT_TWO] = 0; + params.elementLengths[CONSTANT_THREE] = 0; + params.validBit = 0b11; + } + MrgSort(dstTensor[baseOffset], srcList, params); + } + else + { + DataCopy(dstTensor[baseOffset], srcTensor[baseOffset], tailRow * perGroupExpertCountAlign_ * 2); + } + } + + sortedBaseRow = nextBaseRow; + } + } + + template + __aicore__ inline void + MoeGatingTopKHashRegbase::TopKGatherWithSmallKAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor) + { + __local_mem__ float *inputAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ uint32_t *mrgSortAddr = (__local_mem__ uint32_t *)mrgSortTensor.GetPhyAddr(); + __local_mem__ uint32_t *expertIdxAddr = (__local_mem__ uint32_t *)expertIdxTensor.GetPhyAddr(); + __local_mem__ T *outputAddr = (__local_mem__ T *)yTensor.GetPhyAddr(); + + if (tilingData_->normType == 0) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg0, vreg1, mrgSortAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Muls(vreg2, vreg2, routedScalingFactor, preg0); + StoreOneTensorForDtypeT(outputAddr, vreg2, preg0, 0); + MicroAPI::DataCopy(expertIdxAddr, vreg1, preg0); + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vreg4; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg0, vreg1, mrgSortAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::ReduceSum(vreg3, vreg2, preg0); + MicroAPI::Adds(vreg3, vreg3, eps, preg0); + MicroAPI::Duplicate(vreg4, vreg3, preg0); + MicroAPI::Div(vreg4, vreg2, vreg4, preg0); + MicroAPI::Muls(vreg4, vreg4, routedScalingFactor, preg0); + StoreOneTensorForDtypeT(outputAddr, vreg4, preg0, 0); + MicroAPI::DataCopy(expertIdxAddr, vreg1, preg0); + } + } + } + + template + __aicore__ inline void + MoeGatingTopKHashRegbase::TopKGatherWithLargeKAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor) + { + uint32_t k1 = k_; + __local_mem__ float *inputAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ uint32_t *mrgSortAddr = (__local_mem__ uint32_t *)mrgSortTensor.GetPhyAddr(); + __local_mem__ uint32_t *expertIdxAddr = (__local_mem__ uint32_t *)expertIdxTensor.GetPhyAddr(); + __local_mem__ T *outputAddr = (__local_mem__ T *)yTensor.GetPhyAddr(); + + if (tilingData_->normType == 0) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vregSum; + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg preg1 = MicroAPI::CreateMask(); + MicroAPI::Duplicate(vregSum, static_cast(0), preg0); + uint16_t vfLoopNum = static_cast(CeilDiv(k, VL_FLOAT_SIZE)); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::Duplicate(vreg2, static_cast(0), preg1); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Add(vregSum, vregSum, vreg2, preg1); + } + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg1 = MicroAPI::UpdateMask(k1); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg1); + MicroAPI::Muls(vreg2, vreg2, routedScalingFactor, preg1); + StoreOneTensorForDtypeT(outputAddr, vreg2, preg1, i * VL_FLOAT_SIZE); + MicroAPI::DataCopy(expertIdxAddr + i * VL_FLOAT_SIZE, vreg1, preg1); + } + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vreg4; + RegTensor vreg5; + RegTensor vregSum; + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg preg1 = MicroAPI::CreateMask(); + + MicroAPI::Duplicate(vregSum, static_cast(0), preg0); + uint16_t vfLoopNum = static_cast(CeilDiv(k, VL_FLOAT_SIZE)); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::Duplicate(vreg2, static_cast(0), preg1); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Add(vregSum, vregSum, vreg2, preg1); + } + MicroAPI::ReduceSum(vregSum, vregSum, preg1); + MicroAPI::Adds(vregSum, vregSum, eps, preg1); + MicroAPI::Duplicate(vreg4, vregSum, preg1); + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg1 = MicroAPI::UpdateMask(k1); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg1); + MicroAPI::Div(vreg5, vreg2, vreg4, preg1); + MicroAPI::Muls(vreg5, vreg5, routedScalingFactor, preg1); + StoreOneTensorForDtypeT(outputAddr, vreg5, preg1, i * VL_FLOAT_SIZE); + MicroAPI::DataCopy(expertIdxAddr + i * VL_FLOAT_SIZE, vreg1, preg1); + } + } + } + } + + template + __aicore__ inline void + MoeGatingTopKHashRegbase::TopKGatherWithSmallKNotAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor, int32_t expertIdxPad, + int32_t perGroupExpertCountAlign) + { + __local_mem__ float *inputAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ uint32_t *mrgSortAddr = (__local_mem__ uint32_t *)mrgSortTensor.GetPhyAddr(); + __local_mem__ uint32_t *expertIdxAddr = (__local_mem__ uint32_t *)expertIdxTensor.GetPhyAddr(); + __local_mem__ T *outputAddr = (__local_mem__ T *)yTensor.GetPhyAddr(); + + if (tilingData_->normType == 0) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vregAlign; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + + MicroAPI::Duplicate(vregAlign, perGroupExpertCountAlign, preg0); + + MicroAPI::DataCopy(vreg0, vreg1, mrgSortAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Muls(vreg2, vreg2, routedScalingFactor, preg0); + + // compute expertIdx: id = id - floor_div(id, perGroupExpertCountAlign) * pad + MicroAPI::Div(vregAlign, vreg1, vregAlign, preg0); + MicroAPI::Muls(vregAlign, vregAlign, expertIdxPad, preg0); + MicroAPI::Sub(vreg1, vreg1, vregAlign, preg0); + + StoreOneTensorForDtypeT(outputAddr, vreg2, preg0, 0); + MicroAPI::DataCopy(expertIdxAddr, vreg1, preg0); + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vreg4; + RegTensor vregAlign; + + MicroAPI::MaskReg preg0 = MicroAPI::UpdateMask(k); + + MicroAPI::Duplicate(vregAlign, perGroupExpertCountAlign, preg0); + + MicroAPI::DataCopy(vreg0, vreg1, mrgSortAddr); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::ReduceSum(vreg3, vreg2, preg0); + MicroAPI::Adds(vreg3, vreg3, eps, preg0); + MicroAPI::Duplicate(vreg4, vreg3, preg0); + MicroAPI::Div(vreg4, vreg2, vreg4, preg0); + MicroAPI::Muls(vreg4, vreg4, routedScalingFactor, preg0); + + // compute expertIdx: id = id - floor_div(id, perGroupExpertCountAlign) * pad + MicroAPI::Div(vregAlign, vreg1, vregAlign, preg0); + MicroAPI::Muls(vregAlign, vregAlign, expertIdxPad, preg0); + MicroAPI::Sub(vreg1, vreg1, vregAlign, preg0); + + StoreOneTensorForDtypeT(outputAddr, vreg4, preg0, 0); + MicroAPI::DataCopy(expertIdxAddr, vreg1, preg0); + } + } + } + + template + __aicore__ inline void + MoeGatingTopKHashRegbase::TopKGatherWithLargeKNotAlignE(LocalTensor xSigmoidTensor, LocalTensor mrgSortTensor, + LocalTensor expertIdxTensor, LocalTensor yTensor, uint32_t k, + float eps, float routedScalingFactor, int32_t expertIdxPad, + int32_t perGroupExpertCountAlign) + { + uint32_t k1 = k_; + __local_mem__ float *inputAddr = (__local_mem__ float *)xSigmoidTensor.GetPhyAddr(); + __local_mem__ uint32_t *mrgSortAddr = (__local_mem__ uint32_t *)mrgSortTensor.GetPhyAddr(); + __local_mem__ uint32_t *expertIdxAddr = (__local_mem__ uint32_t *)expertIdxTensor.GetPhyAddr(); + __local_mem__ T *outputAddr = (__local_mem__ T *)yTensor.GetPhyAddr(); + + if (tilingData_->normType == 0) + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vregSum; + RegTensor vregAlign; + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg preg1 = MicroAPI::CreateMask(); + + MicroAPI::Duplicate(vregSum, static_cast(0), preg0); + MicroAPI::Duplicate(vregAlign, perGroupExpertCountAlign, preg0); + + uint16_t vfLoopNum = static_cast(CeilDiv(k, VL_FLOAT_SIZE)); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::Duplicate(vreg2, static_cast(0), preg1); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Add(vregSum, vregSum, vreg2, preg1); + } + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg1 = MicroAPI::UpdateMask(k1); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg1); + MicroAPI::Muls(vreg2, vreg2, routedScalingFactor, preg1); + + // compute expertIdx: id = id - floor_div(id, perGroupExpertCountAlign) * pad + MicroAPI::Div(vreg3, vreg1, vregAlign, preg1); + MicroAPI::Muls(vreg3, vreg3, expertIdxPad, preg1); + MicroAPI::Sub(vreg1, vreg1, vreg3, preg1); + StoreOneTensorForDtypeT(outputAddr, vreg2, preg1, i * VL_FLOAT_SIZE); + MicroAPI::DataCopy(expertIdxAddr + i * VL_FLOAT_SIZE, vreg1, preg1); + } + } + } + else + { + __VEC_SCOPE__ + { + RegTensor vreg0; + RegTensor vreg1; + RegTensor vreg2; + RegTensor vreg3; + RegTensor vreg4; + RegTensor vreg5; + RegTensor vreg6; + RegTensor vregSum; + RegTensor vregAlign; + + MicroAPI::MaskReg preg0 = MicroAPI::CreateMask(); + MicroAPI::MaskReg preg1 = MicroAPI::CreateMask(); + + MicroAPI::Duplicate(vregSum, static_cast(0), preg0); + MicroAPI::Duplicate(vregAlign, perGroupExpertCountAlign, preg0); + + uint16_t vfLoopNum = static_cast(CeilDiv(k, VL_FLOAT_SIZE)); + + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg0 = MicroAPI::UpdateMask(k); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::Duplicate(vreg2, static_cast(0), preg1); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg0); + MicroAPI::Add(vregSum, vregSum, vreg2, preg1); + } + MicroAPI::ReduceSum(vregSum, vregSum, preg1); + MicroAPI::Adds(vregSum, vregSum, eps, preg1); + MicroAPI::Duplicate(vreg4, vregSum, preg1); + for (uint16_t i = 0; i < vfLoopNum; i++) + { + preg1 = MicroAPI::UpdateMask(k1); + MicroAPI::DataCopy(vreg0, vreg1, + mrgSortAddr + i * 2 * VL_FLOAT_SIZE); + MicroAPI::DataCopyGather(vreg2, inputAddr, vreg1, preg1); + MicroAPI::Div(vreg6, vreg2, vreg4, preg1); + MicroAPI::Muls(vreg6, vreg6, routedScalingFactor, preg1); + + // compute expertIdx: id = id - floor_div(id, perGroupExpertCountAlign) * pad + MicroAPI::Div(vreg5, vreg1, vregAlign, preg1); + MicroAPI::Muls(vreg5, vreg5, expertIdxPad, preg1); + MicroAPI::Sub(vreg1, vreg1, vreg5, preg1); + StoreOneTensorForDtypeT(outputAddr, vreg6, preg1, i * VL_FLOAT_SIZE); + MicroAPI::DataCopy(expertIdxAddr + i * VL_FLOAT_SIZE, vreg1, preg1); + } + } + } + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::SelectTopKExpertScore() + { + LocalTensor expertIdxTensor = expertIdxOutQueue_.AllocTensor(); + LocalTensor mrgSortTensor = finalSortBuffer_.Get(); + + if (tilingData_->vmsCount % CONSTANT_TWO == 1) + { + mrgSortTensor = sortedInGroupTensor.ReinterpretCast(); + } + + LocalTensor xSigmoidTensor = xSigmoidBuf_.Get(); + LocalTensor yTensor = yOutQueue_.AllocTensor(); + + int32_t expertIdxPad = perGroupExpertCountAlign_ - perGroupExpertCount_; + if (k_ <= VL_FLOAT_SIZE) + { + if (expertIdxPad != 0) + { + TopKGatherWithSmallKNotAlignE(xSigmoidTensor, mrgSortTensor, expertIdxTensor, yTensor, k_, eps_, routedScalingFactor_, + expertIdxPad, perGroupExpertCountAlign_); + } + else + { + TopKGatherWithSmallKAlignE(xSigmoidTensor, mrgSortTensor, expertIdxTensor, yTensor, k_, eps_, routedScalingFactor_); + } + } + else + { + if (expertIdxPad != 0) + { + TopKGatherWithLargeKNotAlignE(xSigmoidTensor, mrgSortTensor, expertIdxTensor, yTensor, k_, eps_, routedScalingFactor_, + expertIdxPad, perGroupExpertCountAlign_); + } + else + { + TopKGatherWithLargeKAlignE(xSigmoidTensor, mrgSortTensor, expertIdxTensor, yTensor, k_, eps_, routedScalingFactor_); + } + } + + yOutQueue_.EnQue(yTensor); + expertIdxOutQueue_.EnQue(expertIdxTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::CopyOut(int64_t row) + { + LocalTensor yOutTensor = yOutQueue_.DeQue(); + LocalTensor expertIdxTensor = expertIdxOutQueue_.DeQue(); + DataCopyExtParams dataCopyParams{1, static_cast(k_ * sizeof(T)), 0, 0, 0}; + DataCopyPad(yGm_[row * k_], yOutTensor, dataCopyParams); + + dataCopyParams.blockLen = k_ * sizeof(int32_t); + DataCopyPad(expertIdxGm_[row * k_], expertIdxTensor, dataCopyParams); + + expertIdxOutQueue_.FreeTensor(expertIdxTensor); + yOutQueue_.FreeTensor(yOutTensor); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::Init(GM_ADDR x, GM_ADDR bias, GM_ADDR inputIds, GM_ADDR tid2eid, GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, + GM_ADDR workspace, + const MoeGatingTopKHashRegbaseTilingData *tilingData, TPipe *tPipe) + { + tilingData_ = tilingData; + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + if (blockIdx_ == GetBlockNum() - 1) + { + curCoreRowCount_ = tilingData_->lastCoreRowCount; + } + else + { + curCoreRowCount_ = tilingData_->perCoreRowCount; + } + expertCount_ = tilingData_->expertCount; + k_ = tilingData_->k; + kGroup_ = tilingData_->kGroup; + groupCount_ = tilingData_->groupCount; + perGroupExpertCount_ = tilingData_->perGroupExpertCount; + perGroupExpertCountAlign_ = tilingData_->perGroupExpertCountAlign; + expertCountAlign_ = Align(perGroupExpertCountAlign_ * groupCount_, sizeof(float)); + routedScalingFactor_ = tilingData_->routedScalingFactor; + groupSelectMode_ = tilingData_->groupSelectMode; + eps_ = tilingData_->eps; + hashFlag_ = tilingData_->hashFlag == 1; + + // init input gm buf + xGm_.SetGlobalBuffer((__gm__ T *)x + tilingData_->perCoreRowCount * expertCount_ * blockIdx_, expertCount_); + if (bias != nullptr) + { + hasBias_ = true; + biasGm_.SetGlobalBuffer((__gm__ T *)bias, expertCount_); + } + yGm_.SetGlobalBuffer((__gm__ T *)y + tilingData_->perCoreRowCount * k_ * blockIdx_, k_); + expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx + tilingData_->perCoreRowCount * k_ * blockIdx_, k_); + outGm_.SetGlobalBuffer((__gm__ float *)out + tilingData_->perCoreRowCount * expertCount_ * blockIdx_, expertCount_); + + inputIdsGm_.SetGlobalBuffer((__gm__ U1 *)inputIds); + tid2eidGm_.SetGlobalBuffer((__gm__ U2 *)tid2eid); + + // init queue + int32_t expertGroupAlign = groupCount_ * perGroupExpertCountAlign_; + int32_t groupAlign = static_cast(CeilAlign(groupCount_, ONE_REPEAT_SORT_NUM)); + pipe_->InitBuffer(xInQueue_, CONSTANT_TWO, expertGroupAlign * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(yOutQueue_, CONSTANT_TWO, AlignBytes(k_, sizeof(T))); + pipe_->InitBuffer(expertIdxOutQueue_, CONSTANT_TWO, AlignBytes(k_, sizeof(int32_t))); + pipe_->InitBuffer(outOutQueue_, CONSTANT_TWO, expertGroupAlign * sizeof(float)); + + pipe_->InitBuffer(biasBuf_, expertGroupAlign * sizeof(T)); + pipe_->InitBuffer(xSigmoidBuf_, expertGroupAlign * sizeof(float)); + pipe_->InitBuffer(xBiasBuf_, expertGroupAlign * sizeof(float)); + pipe_->InitBuffer(indexBuffer_, expertGroupAlign * sizeof(int32_t)); + pipe_->InitBuffer(sortedInGroupBuf_, expertGroupAlign * sizeof(U2) * 2); + pipe_->InitBuffer(finalSortBuffer_, expertGroupAlign * sizeof(float) * 2); + pipe_->InitBuffer(groupBuf_, groupAlign * sizeof(float)); + pipe_->InitBuffer(sortedGroupBuf_, groupAlign * sizeof(float) * 2); + } + + template + __aicore__ inline void MoeGatingTopKHashRegbase::Process() + { + CopyInBias(); + if (kGroup_ == groupCount_ || groupCount_ == expertCount_) + { + CopyInX(0); + for (int64_t row = 1; row < curCoreRowCount_; row++) + { + ComputeX(); + CopyOutXNorm(row - 1); + CopyInX(row); + if (hashFlag_) + { + HashCompute(row - 1 + tilingData_->perCoreRowCount * blockIdx_); + } + else + { + TopKCompute(); + } + CopyOut(row - 1); + } + ComputeX(); + CopyOutXNorm(curCoreRowCount_ - 1); + if (hashFlag_) + { + HashCompute(curCoreRowCount_ - 1 + tilingData_->perCoreRowCount * blockIdx_); + } + else + { + TopKCompute(); + } + CopyOut(curCoreRowCount_ - 1); + return; + } + + CopyInX(0); + for (int64_t row = 1; row < curCoreRowCount_; row++) + { + ComputeX(); + CopyOutXNorm(row - 1); + SortInGroup(); + SelectTopKGroupIndex(); + CopyInX(row); + FinalSortByKGroup(); + FinalSortAfterKGroup(); + SelectTopKExpertScore(); + CopyOut(row - 1); + } + ComputeX(); + CopyOutXNorm(curCoreRowCount_ - 1); + SortInGroup(); + SelectTopKGroupIndex(); + FinalSortByKGroup(); + FinalSortAfterKGroup(); + SelectTopKExpertScore(); + CopyOut(curCoreRowCount_ - 1); + } +} // namespace MoeGatingTopKHash +#endif // MOE_GATING_TOP_K_REGBASE_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash_without_group.h b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash_without_group.h new file mode 100644 index 000000000..03b5e1303 --- /dev/null +++ b/csrc/ascend/moe/moe_gating_top_k_hash/op_kernel/moe_gating_top_k_hash_without_group.h @@ -0,0 +1,385 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_gating_top_k_hash_without_group.h + * \brief + */ +#ifndef MOE_GATING_TOP_K_E_K_WITHOUT_GROUP_H +#define MOE_GATING_TOP_K_E_K_WITHOUT_GROUP_H +#include "kernel_operator.h" +#include "common.h" +#include "kernel_utils.h" +namespace MoeGatingTopKHash { +using namespace AscendC; + +template +class MoeGatingTopKHashWithoutGroup { +public: + __aicore__ inline MoeGatingTopKHashWithoutGroup(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR bias, GM_ADDR inputIds, GM_ADDR tid2eid, GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKHashTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyInBiasAndInitExpertId(); + __aicore__ inline void CopyInX(int64_t progress); + __aicore__ inline void ComputeX(); + __aicore__ inline void CopuOutXNorm(int64_t row); + __aicore__ inline void SelectTopKExpertIdx(); + __aicore__ inline void SelectExpertIdxByHash(int64_t row); + __aicore__ inline void SelectTopKExpertScore(); + __aicore__ inline void CopyOut(int64_t row); + +private: + TPipe *pipe_; + TQue xInQueue_; + TQue yOutQueue_; + TQue expertIdxOutQueue_; + TQue outOutQueue_; + + TBuf biasBuf_; // 存放输入bias + TBuf expertIdBuf_; // 专家编号 + TBuf xNormWithBiasBuf_; // 存放加了bias之后的值 + TBuf xNormBuf_; // 存放计算sigmoid或softmax的值 + TBuf topKExpertIdBuf_; + TBuf calcTmpBuf_; + + GlobalTensor xGm_; + GlobalTensor biasGm_; + GlobalTensor inputIdsGm_; + GlobalTensor tid2eidGm_; + GlobalTensor yGm_; + GlobalTensor expertIdxGm_; + GlobalTensor outGm_; + + int64_t blockIdx_ = 0; + int64_t perCoreRowCount_ = 0; + int64_t curCoreRowCount_ = 0; + int64_t expertCount_ = 0; + bool addBias_ = false; + bool outFlag_ = false; + bool hashFlag_ = false; + int64_t k_ = 0; + + int64_t expertCountAlign_ = 0; + const MoeGatingTopKHashTilingData *tilingData_; + + template + __aicore__ inline void SetWaitFlag(HardEvent evt) + { + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(evt)); + SetFlag(eventId); + WaitFlag(eventId); + } +}; + +template +__aicore__ inline void MoeGatingTopKHashWithoutGroup::CopyInBiasAndInitExpertId() +{ + LocalTensor biasTensor = biasBuf_.Get(); + LocalTensor expertIdTensor = expertIdBuf_.Get(); + DataCopyExtParams dataCopyParams{1, static_cast(expertCount_ * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if (addBias_) { + if constexpr (IsSameType::value) { + DataCopyPad(biasTensor, biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + } else { + DataCopyPad(biasTensor[expertCountAlign_].ReinterpretCast(), biasGm_, dataCopyParams, dataCopyPadParams); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(biasTensor, biasTensor[expertCountAlign_].ReinterpretCast(), RoundMode::CAST_NONE, + expertCountAlign_); + PipeBarrier(); + } + } + ArithProgression(expertIdTensor, static_cast(0), static_cast(1), expertCount_); +} + +template +__aicore__ inline void MoeGatingTopKHashWithoutGroup::CopyInX(int64_t row) +{ + LocalTensor xInLocalTensor = xInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(expertCount_ * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + if constexpr (IsSameType::value) { + DataCopyPad(xInLocalTensor, xGm_[row * expertCount_], dataCopyParams, dataCopyPadParams); + } else { + DataCopyPad(xInLocalTensor[expertCountAlign_].ReinterpretCast(), xGm_[row * expertCount_], dataCopyParams, + dataCopyPadParams); + } + xInQueue_.EnQue(xInLocalTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashWithoutGroup::ComputeX() +{ + LocalTensor xNormTensor = xNormBuf_.Get(); + LocalTensor xInLocalTensor = xInQueue_.DeQue(); + LocalTensor xNormWithBiasTensor = xNormWithBiasBuf_.Get(); + LocalTensor biasTensor = biasBuf_.Get(); + + if constexpr (!IsSameType::value) { + Cast(xInLocalTensor, xInLocalTensor[expertCountAlign_].ReinterpretCast(), RoundMode::CAST_NONE, + expertCount_); + PipeBarrier(); + } + + if (tilingData_->normType == NORM_TYPE_SIGMOID) { // sigmoid + LocalTensor calcNormTmpTensor = calcTmpBuf_.Get(); + Sigmoid(xNormTensor, xInLocalTensor, calcNormTmpTensor, expertCount_); + PipeBarrier(); + } else if (tilingData_->normType == NORM_TYPE_SOFTMAX) { // softmax + LocalTensor reduceValueTensor = calcTmpBuf_.Get(); + LocalTensor calcTmp = calcTmpBuf_.Get()[8]; + ReduceMax(reduceValueTensor, xInLocalTensor, calcTmp, expertCount_); + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float maxValue = reduceValueTensor.GetValue(0); + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Adds(xNormTensor, xInLocalTensor, -maxValue, expertCount_); + PipeBarrier(); + Exp(xNormTensor, xNormTensor, expertCount_); + PipeBarrier(); + ReduceSum(reduceValueTensor, xNormTensor, calcTmp, expertCount_); + eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float sumValue = reduceValueTensor.GetValue(0); + eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Muls(xNormTensor, xNormTensor, 1.0f / sumValue, expertCount_); + PipeBarrier(); + } else { + LocalTensor calcNormTmpTensor = calcTmpBuf_.Get(); + Exp(calcNormTmpTensor, xInLocalTensor, expertCount_); + PipeBarrier(); + Adds(calcNormTmpTensor, calcNormTmpTensor, float(1.0), expertCount_); + PipeBarrier(); + Ln(calcNormTmpTensor, calcNormTmpTensor, expertCount_); + PipeBarrier(); + Sqrt(xNormTensor, calcNormTmpTensor, expertCount_); + PipeBarrier(); + } + if (addBias_) { + Add(xNormWithBiasTensor, xNormTensor, biasTensor, expertCount_); + } else { + DataCopy(xNormWithBiasTensor, xNormTensor, expertCountAlign_); + } + + int64_t duplicateNum = expertCount_ % ONE_REPEAT_SORT_NUM; + int duplicateIndex = expertCount_ - duplicateNum; + if (duplicateNum > 0) { + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(xNormWithBiasTensor.ReinterpretCast()[duplicateIndex], FLOAT32_NEG_INF, mask, 1, 1, 1); + PipeBarrier(); + } + xInQueue_.FreeTensor(xInLocalTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashWithoutGroup::CopuOutXNorm(int64_t row) +{ + LocalTensor outOutTensor = outOutQueue_.AllocTensor(); + LocalTensor xNormTensor = xNormBuf_.Get(); + DataCopy(outOutTensor, xNormTensor, expertCountAlign_); + outOutQueue_.EnQue(outOutTensor); + outOutTensor = outOutQueue_.DeQue(); + DataCopyExtParams dataCopyParams{1, static_cast(expertCount_ * sizeof(float)), 0, 0, 0}; + DataCopyPad(outGm_[row * expertCount_], outOutTensor, dataCopyParams); + outOutQueue_.FreeTensor(outOutTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashWithoutGroup::SelectTopKExpertIdx() +{ + LocalTensor expertIdxOut = expertIdxOutQueue_.AllocTensor(); + LocalTensor xNormWithBiasTensor = xNormWithBiasBuf_.Get(); + LocalTensor expertIdTensor = expertIdBuf_.Get(); + LocalTensor topKExpertId = topKExpertIdBuf_.Get(); + LocalTensor sortedScore = calcTmpBuf_.Get(); + LocalTensor sortTmp = calcTmpBuf_.Get()[expertCountAlign_ * CONSTANT_TWO]; + PipeBarrier(); + Sort(sortedScore, xNormWithBiasTensor, expertIdTensor, sortTmp, + expertCountAlign_ / ONE_REPEAT_SORT_NUM); + + GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = Ceil(k_ * sizeof(float) * CONSTANT_TWO, REPEAT_BYTES); + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = REPEAT_BLOCKS; + gatherMaskParams.src1RepeatStride = 0; + + uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数 + uint8_t src1Pattern = 2; // 内置固定模式 + PipeBarrier(); + GatherMask(topKExpertId, sortedScore.template ReinterpretCast(), src1Pattern, false, + static_cast(0), gatherMaskParams, rsvdCnt); + + DataCopy(expertIdxOut, topKExpertId, expertCountAlign_); + expertIdxOutQueue_.EnQue(expertIdxOut); +} + +template +__aicore__ inline void MoeGatingTopKHashWithoutGroup::SelectTopKExpertScore() +{ + LocalTensor xNormTensor = xNormBuf_.Get(); + LocalTensor yOutTensor = yOutQueue_.AllocTensor(); + LocalTensor topKExpertId = topKExpertIdBuf_.Get(); + LocalTensor topKExpertIdWithByte = calcTmpBuf_.Get(); + PipeBarrier(); + Muls(topKExpertIdWithByte, topKExpertId, static_cast(sizeof(float)), k_); + PipeBarrier(); + Gather(yOutTensor, xNormTensor, topKExpertIdWithByte.template ReinterpretCast(), static_cast(0), + k_); + + if (tilingData_->normType == NORM_TYPE_SIGMOID || tilingData_->normType == NORM_TYPE_SOFTPLUS) { + LocalTensor maxValueTensor = calcTmpBuf_.Get(); + LocalTensor tmpTensor = calcTmpBuf_.Get()[BLOCK_BYTES]; + PipeBarrier(); + ReduceSum(maxValueTensor, yOutTensor, tmpTensor, k_); + event_t eventIdVToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + SetFlag(eventIdVToS); + WaitFlag(eventIdVToS); + float sumValue = maxValueTensor.GetValue(0) + tilingData_->eps; + event_t eventIdSToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + SetFlag(eventIdSToV); + WaitFlag(eventIdSToV); + Duplicate(tmpTensor, sumValue, k_); + PipeBarrier(); + Div(yOutTensor, yOutTensor, tmpTensor, k_); + } + PipeBarrier(); + Muls(yOutTensor, yOutTensor, tilingData_->routedScalingFactor, k_); + + if constexpr (!IsSameType::value) { + PipeBarrier(); + Cast(yOutTensor.ReinterpretCast(), yOutTensor, RoundMode::CAST_RINT, k_); + } + + yOutQueue_.EnQue(yOutTensor); +} + +template +__aicore__ inline void MoeGatingTopKHashWithoutGroup::CopyOut(int64_t row) +{ + LocalTensor yOutTensor = yOutQueue_.DeQue(); + LocalTensor expertIdxOut = expertIdxOutQueue_.DeQue(); + DataCopyExtParams dataCopyParams{1, static_cast(k_ * sizeof(T)), 0, 0, 0}; + DataCopyPad(yGm_[row * k_], yOutTensor, dataCopyParams); + dataCopyParams.blockLen = k_ * sizeof(int32_t); + DataCopyPad(expertIdxGm_[row * k_], expertIdxOut, dataCopyParams); + yOutQueue_.FreeTensor(yOutTensor); + expertIdxOutQueue_.FreeTensor(expertIdxOut); +} + +template +__aicore__ inline void MoeGatingTopKHashWithoutGroup::SelectExpertIdxByHash(int64_t row) +{ + LocalTensor expertIdxOut = expertIdxOutQueue_.AllocTensor(); + LocalTensor hashExpertId = topKExpertIdBuf_.Get(); + LocalTensor hashExpertIdInt32 = hashExpertId.template ReinterpretCast(); + U1 key = inputIdsGm_.GetValue(row); + SetWaitFlag(HardEvent::S_MTE2); + DataCopyExtParams dataCopyParams{1, static_cast(k_ * sizeof(U2)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, static_cast(0)}; + DataCopyPad(hashExpertId, tid2eidGm_[key * k_], dataCopyParams, dataCopyPadParams); + SetWaitFlag(HardEvent::MTE2_V); + if constexpr (IsSameType::value) { + DataCopy(expertIdxOut, hashExpertId, Align(k_, sizeof(int32_t))); + } else { + Cast(hashExpertIdInt32, hashExpertId, RoundMode::CAST_NONE, Align(k_, sizeof(U2))); + PipeBarrier(); + DataCopy(expertIdxOut, hashExpertIdInt32, Align(k_, sizeof(int32_t))); + } + expertIdxOutQueue_.EnQue(expertIdxOut); +} + +template +__aicore__ inline void MoeGatingTopKHashWithoutGroup::Init(GM_ADDR x, GM_ADDR bias, GM_ADDR inputIds, GM_ADDR tid2eid, + GM_ADDR y, GM_ADDR expertIdx, GM_ADDR out, GM_ADDR workspace, + const MoeGatingTopKHashTilingData *tilingData, TPipe *tPipe) +{ + tilingData_ = tilingData; + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + perCoreRowCount_ = tilingData_->perCoreRowCount; + if (blockIdx_ == GetBlockNum() - 1) { + curCoreRowCount_ = tilingData_->lastCoreRowCount; + } else { + curCoreRowCount_ = tilingData_->perCoreRowCount; + } + expertCount_ = tilingData_->expertCount; + addBias_ = tilingData_->addBias == 1; + outFlag_ = tilingData_->outFlag == 1; + hashFlag_ = tilingData_->hashFlag == 1; + k_ = tilingData_->k; + + expertCountAlign_ = Ceil(expertCount_, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + + // init input gm buf + xGm_.SetGlobalBuffer((__gm__ T *)x + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_); + biasGm_.SetGlobalBuffer((__gm__ T *)bias, expertCount_); + inputIdsGm_.SetGlobalBuffer((__gm__ U1 *)inputIds); + tid2eidGm_.SetGlobalBuffer((__gm__ U2 *)tid2eid); + + // init output gm buf + yGm_.SetGlobalBuffer((__gm__ T *)y + perCoreRowCount_ * k_ * blockIdx_, k_); + expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx + perCoreRowCount_ * k_ * blockIdx_, k_); + outGm_.SetGlobalBuffer((__gm__ float *)out + perCoreRowCount_ * expertCount_ * blockIdx_, expertCount_); + + // init que + pipe_->InitBuffer(xInQueue_, 1, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(yOutQueue_, 1, Align(k_, sizeof(float)) * sizeof(float)); + pipe_->InitBuffer(expertIdxOutQueue_, 1, Align(k_, sizeof(float)) * sizeof(int32_t)); + pipe_->InitBuffer(outOutQueue_, 1, expertCountAlign_ * sizeof(float)); + + // init calc buf + pipe_->InitBuffer(biasBuf_, expertCountAlign_ * sizeof(float) * (sizeof(float) / sizeof(T))); + pipe_->InitBuffer(expertIdBuf_, expertCountAlign_ * sizeof(int32_t)); + pipe_->InitBuffer(xNormBuf_, expertCountAlign_ * sizeof(float)); + pipe_->InitBuffer(xNormWithBiasBuf_, expertCountAlign_ * sizeof(float)); + pipe_->InitBuffer(topKExpertIdBuf_, Align(k_, sizeof(U2)) * sizeof(U2)); + + // init tmp buf + pipe_->InitBuffer(calcTmpBuf_, expertCountAlign_ * sizeof(float) * CONSTANT_EIGHT); +} + +template +__aicore__ inline void MoeGatingTopKHashWithoutGroup::Process() +{ + CopyInBiasAndInitExpertId(); + for (int64_t row = 0; row < curCoreRowCount_; row++) { + CopyInX(row); + ComputeX(); + if (outFlag_) { + CopuOutXNorm(row); + } + if (hashFlag_) { + SelectExpertIdxByHash(row + perCoreRowCount_ * blockIdx_); + } else { + SelectTopKExpertIdx(); + } + SelectTopKExpertScore(); + CopyOut(row); + } +} +} // namespace MoeGatingTopKHash +#endif // MOE_GATING_TOP_K_E_K_WITHOUT_GROUP_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_grouped_matmul/CMakeLists.txt b/csrc/ascend/moe/moe_grouped_matmul/CMakeLists.txt new file mode 100644 index 000000000..298ed5677 --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/moe_grouped_matmul/op_host/CMakeLists.txt b/csrc/ascend/moe/moe_grouped_matmul/op_host/CMakeLists.txt new file mode 100644 index 000000000..7d9dc8334 --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/op_host/CMakeLists.txt @@ -0,0 +1,22 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnInner PRIVATE + moe_grouped_matmul_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME MoeGroupedMatmulCustom + OPTIONS + --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE moe_grouped_matmul ACLNNTYPE aclnn_inner) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/moe/moe_grouped_matmul/op_host/moe_grouped_matmul_def.cpp b/csrc/ascend/moe/moe_grouped_matmul/op_host/moe_grouped_matmul_def.cpp new file mode 100644 index 000000000..2ed0f124a --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/op_host/moe_grouped_matmul_def.cpp @@ -0,0 +1,53 @@ +/** +* This program is free software, you can redistribute it and/or modify. +* Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_grouped_matmul_def.cpp + * \brief + */ + +#include "register/op_def_registry.h" +#include "moe_grouped_matmul_infershape.cpp" + +namespace ops { +class MoeGroupedMatmul : public OpDef { +public: + explicit MoeGroupedMatmul(const char *name) : OpDef(name) { + this->Input("x") + .ParamType(DYNAMIC) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("weight") + .ParamType(DYNAMIC) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ}); + this->Input("group_list") + .ParamType(REQUIRED) + .DataTypeList({ge::DT_INT64, ge::DT_INT32}) + .FormatList({ge::FORMAT_ND}); + this->Output("y") + .ParamType(DYNAMIC) + .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Attr("transpose_weight") + .AttrType(OPTIONAL) + .Bool(false); + this->SetInferShape(ge::InferShape); + this->SetInferDataType(ge::InferDataType); + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + } + + +}; + +OP_ADD(MoeGroupedMatmul); + +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/moe/moe_grouped_matmul/op_host/moe_grouped_matmul_infershape.cpp b/csrc/ascend/moe/moe_grouped_matmul/op_host/moe_grouped_matmul_infershape.cpp new file mode 100644 index 000000000..3633a59fa --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/op_host/moe_grouped_matmul_infershape.cpp @@ -0,0 +1,35 @@ +#include "register/op_impl_registry.h" + + +#include + +namespace ge { +constexpr uint32_t X_INDEX = 0; +constexpr uint32_t WEIGHT_INDEX = 1; +constexpr uint32_t GROUPLIST_INDEX = 2; + +static ge::graphStatus InferShape(gert::InferShapeContext *context) { + const gert::Shape* x_shape = context->GetDynamicInputShape(X_INDEX, 0); + const gert::Shape* weight_shape = context->GetDynamicInputShape(WEIGHT_INDEX, 0); + bool transpose_weight = static_cast(*(context->GetAttrs()->GetAttrPointer(0))); + gert::Shape* y_shape = context->GetOutputShape(0); + *y_shape = *x_shape; + auto weight_desc = context->GetDynamicInputDesc(WEIGHT_INDEX, 0); + auto weight_format = static_cast(ge::GetPrimaryFormat(weight_desc->GetStorageFormat())); + bool weight_nz = weight_format == ge::FORMAT_FRACTAL_NZ; + int64_t dim_n; + if (weight_nz) { + dim_n = transpose_weight ? (weight_shape->GetDim(1) * weight_shape->GetDim(3)) : + (weight_shape->GetDim(2) * weight_shape->GetDim(4)); + } else { + dim_n = transpose_weight ? weight_shape->GetDim(1) : weight_shape->GetDim(2); + } + y_shape->SetDim(1, dim_n); +} + +static ge::graphStatus InferDataType(gert::InferDataTypeContext *context) { + const auto input_dtype = context->GetDynamicInputDataType(X_INDEX, 0); + context->SetOutputDataType(0, input_dtype); + return ge::GRAPH_SUCCESS; +} +} // namespace ge \ No newline at end of file diff --git a/csrc/ascend/moe/moe_grouped_matmul/op_host/moe_grouped_matmul_tiling.cpp b/csrc/ascend/moe/moe_grouped_matmul/op_host/moe_grouped_matmul_tiling.cpp new file mode 100644 index 000000000..d7bdd81c3 --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/op_host/moe_grouped_matmul_tiling.cpp @@ -0,0 +1,336 @@ +/** +* This program is free software, you can redistribute it and/or modify. +* Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "moe_grouped_matmul_tiling.h" +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" + +#define OP_LOGD(nodeName, fmt, ...) printf(fmt, ##__VA_ARGS__); printf("\n") +#define OP_LOGE(nodeName, fmt, ...) printf(fmt, ##__VA_ARGS__); printf("\n") + +constexpr uint32_t X_INDEX = 0; +constexpr uint32_t WEIGHT_INDEX = 1; +constexpr uint32_t GROUPLIST_INDEX = 2; +namespace optiling { +constexpr uint64_t BEST_L1_PARTA = 256UL * 1024UL; +constexpr uint64_t BEST_L1_PARTB = 128UL * 1024UL; +constexpr uint64_t L1_PARTA_SIZE = 256UL * 1024UL; +constexpr int32_t BEST_BASEN = 256; +constexpr uint64_t DOUBLE_BUFFER_L0A_L0B = 2; +constexpr uint64_t DOUBLE_BUFFER_STEPKA_STEPKB = 2; +constexpr uint32_t FP32_DATATYPE_SIZE = 4; +constexpr int32_t MAX_BASEM = 256; + +static inline uint32_t SixteenAlign(uint32_t a, bool up = false) { + if (up) { + a += 15U; // 15: 16 bytes up-align + } + return a & ~15U; // ~15: 16 bytes down-align +} + +static inline int64_t SixteenAlign(int64_t a, bool up = false) { + if (up) { + a += 15; // 15: 16 bytes up-align + } + return a & ~15; // ~15: 16 bytes down-align +} + +class TilingMoeGroupedMatmulFunc { + public: + explicit TilingMoeGroupedMatmulFunc(gert::TilingContext* tiling_context) + : tiling_context_(tiling_context) {} + + ge::graphStatus Init(); + ge::graphStatus RunKernelTiling(); + + private: + MoeGroupedMatmulTilingData tiling_data_; + gert::TilingContext* tiling_context_ = nullptr; + + void SetTilingKey(); + void FillTilingData(); + ge::graphStatus CalMMTiling(); + ge::graphStatus GMMSetMMTiling(); + ge::graphStatus CalcStepKaKb(uint32_t& mm_step_ka, uint32_t& mm_step_kb); + ge::graphStatus DynamicTIlingSingleN(); + + void InitPlatformInfo(matmul_tiling::PlatformInfo& platformInfo); + void GMMGetPlatformInfo(); + int64_t m_ = 0L; + int64_t n_ = 0L; + int64_t k_ = 0L; + bool transpose_weight = false; + bool weight_nz = false; + uint32_t single_m_ = 0; + uint32_t single_n_ = 0; + int32_t baseM_ = 0; + int32_t baseN_ = 0; + int32_t baseK_ = 0; + int32_t nz_factor_ = 1; + uint32_t group_num_ = 0; + uint32_t core_num_ = 0; + uint32_t mmDataTypeSize_ = 0; + size_t sync_workspace_size_ = 0; + ge::DataType x_dtype; + uint64_t l1_size, l0a_size, l0b_size, l0c_size, ub_size; + uint32_t aic_num, aiv_num; + // SocVersion soc_version; +}; + +ge::graphStatus TilingMoeGroupedMatmulFunc::CalcStepKaKb(uint32_t& mm_step_ka, uint32_t& mm_step_kb) { + uint64_t available_l1_size = l1_size; + if (available_l1_size < L1_PARTA_SIZE) { + OP_LOGE(tiling_context_->GetNodeName(), "available_l1_size is less than 256k."); + return ge::GRAPH_FAILED; + } + // according to double buffer, recompute the params used for data movement from GM to L1 + uint64_t l1_a_size = baseM_ > baseN_ ? L1_PARTA_SIZE : available_l1_size - L1_PARTA_SIZE; + uint64_t l1_b_size = available_l1_size - l1_a_size; + // 2: double buffer + mm_step_ka = (l1_a_size / 2UL) / (static_cast(baseM_) * baseK_ * mmDataTypeSize_); + // 2: double buffer + mm_step_kb = (l1_b_size / 2UL) / (static_cast(baseN_) * baseK_ * mmDataTypeSize_); + if (mm_step_ka == 0 || mm_step_kb == 0) { + OP_LOGE(tiling_context_->GetNodeName(), "stepka or stepkb cannot be 0."); + return ge::GRAPH_FAILED; + } + + if (mm_step_ka > mm_step_kb) { + mm_step_ka = mm_step_ka / mm_step_kb * mm_step_kb; + } else if (mm_step_ka < mm_step_kb) { + mm_step_kb = mm_step_kb / mm_step_ka * mm_step_ka; + } + return ge::GRAPH_SUCCESS; +} + +void TilingMoeGroupedMatmulFunc::GMMGetPlatformInfo() { + auto platform_info = platform_ascendc::PlatformAscendC(tiling_context_->GetPlatformInfo()); + platform_info.GetCoreMemSize(platform_ascendc::CoreMemType::L1, l1_size); + platform_info.GetCoreMemSize(platform_ascendc::CoreMemType::L0_A, l0a_size); + platform_info.GetCoreMemSize(platform_ascendc::CoreMemType::L0_B, l0b_size); + platform_info.GetCoreMemSize(platform_ascendc::CoreMemType::L0_C, l0c_size); + platform_info.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size); + aic_num = platform_info.GetCoreNumAic(); + aiv_num = platform_info.GetCoreNumAiv(); +} + +void TilingMoeGroupedMatmulFunc::InitPlatformInfo(matmul_tiling::PlatformInfo& platformInfo) { + auto platform_info = platform_ascendc::PlatformAscendC(tiling_context_->GetPlatformInfo()); + platformInfo.socVersion = platform_info.GetSocVersion(); + platformInfo.l1Size = l1_size; + platformInfo.l0CSize = l0c_size; + platformInfo.ubSize = ub_size; + platformInfo.l0ASize = l0a_size; + platformInfo.l0BSize = l0b_size; +} + +ge::graphStatus TilingMoeGroupedMatmulFunc::GMMSetMMTiling() { + matmul_tiling::DataType matmul_dtype = static_cast(x_dtype); + matmul_tiling::PlatformInfo platformInfo; + InitPlatformInfo(platformInfo); + // matmul_tiling::MatmulApiTiling mm(platformInfo); + matmul_tiling::MultiCoreMatmulTiling mm(platformInfo); + mm.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_dtype, false); + if (weight_nz) { + mm.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::NZ, matmul_dtype, transpose_weight); + } else { + mm.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_dtype, transpose_weight); + } + mm.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_dtype); + mm.SetOrgShape(m_, n_, k_); + mm.SetShape(m_, baseN_, k_); + // mm.SetShape(single_m_, single_n_, k_); + mm.SetFixSplit(baseM_, baseN_, baseK_); + mm.SetBufferSpace(l1_size, l0c_size, ub_size); + + if (mm.GetTiling(tiling_data_.mm_tiling) == -1) { + OP_LOGE(tiling_context_->GetNodeName(), "matmul getTiling failed."); + return ge::GRAPH_FAILED; + } + uint32_t mm_step_ka = 1; + uint32_t mm_step_kb = 1; + + auto ret = CalcStepKaKb(mm_step_ka, mm_step_kb); + if (ret != ge::GRAPH_SUCCESS) { + OP_LOGE(tiling_context_->GetNodeName(), "matmul calc stepka or stepkb failed."); + return ge::GRAPH_FAILED; + } + + constexpr uint32_t step_m = 1; // 1: step_m set fixed value 1 + constexpr uint32_t step_n = 1; // 1: step_n set fixed value 1 + uint32_t mm_depth_a1 = mm_step_ka * DOUBLE_BUFFER_STEPKA_STEPKB * step_m; + uint32_t mm_depth_b1 = mm_step_kb * DOUBLE_BUFFER_STEPKA_STEPKB * step_n; + tiling_data_.mm_tiling.set_shareMode(0); + tiling_data_.mm_tiling.set_dbL0C(1); // disable double buffer for LOC + tiling_data_.mm_tiling.set_baseM(baseM_); // set precomputed baseM + tiling_data_.mm_tiling.set_baseN(baseN_); // set precomputed baseN + tiling_data_.mm_tiling.set_baseK(baseK_); // set precomputed baseK + tiling_data_.mm_tiling.set_stepKa(mm_step_ka); // set precomputed mmStepKa + tiling_data_.mm_tiling.set_depthA1(mm_depth_a1); // set precomputed mmDepthA1 + tiling_data_.mm_tiling.set_stepKb(mm_step_kb); // set precomputed mmStepKb + tiling_data_.mm_tiling.set_depthB1(mm_depth_b1); // set precomputed mmDepthB1 + tiling_data_.mm_tiling.set_stepM(step_m); // set precomputed stepM + tiling_data_.mm_tiling.set_stepN(step_n); // set precomputed stepN + OP_LOGD(context->GetNodeName(), "GMM_tiling: baseM is %d, baseK is %d, baseN is %d, transpose_weight is %d, weight_nz is %d", baseM_, baseK_, baseN_, transpose_weight, weight_nz); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingMoeGroupedMatmulFunc::CalMMTiling() { + baseN_ = BEST_BASEN; + if (x_dtype == ge::DT_BF16 || x_dtype == ge::DT_FLOAT16) { + mmDataTypeSize_ = 2; + } else { + OP_LOGE(tiling_context_->GetNodeName(), "only support bf16 or fp16."); + return ge::GRAPH_FAILED; + } + + baseK_ = static_cast((l0b_size / DOUBLE_BUFFER_L0A_L0B) / (static_cast(baseN_) * mmDataTypeSize_)); + baseK_ = static_cast(SixteenAlign(static_cast(baseK_))); + uint32_t max_base_m = static_cast(l0c_size / + (static_cast(baseN_) * FP32_DATATYPE_SIZE)); + baseM_ = std::min((l0a_size / DOUBLE_BUFFER_L0A_L0B) / + (static_cast(baseK_) * mmDataTypeSize_), max_base_m); + baseM_ = baseM_ > m_ ? SixteenAlign(m_, true) : SixteenAlign(static_cast(baseM_)); + + if (baseM_ > MAX_BASEM) { + baseM_ = MAX_BASEM; + } + + if (baseM_ == 0 || baseK_ == 0) { + OP_LOGE(tiling_context_->GetNodeName(), "baseM_ or baseN_ cannot be 0."); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus TilingMoeGroupedMatmulFunc::Init() { + GMMGetPlatformInfo(); + // only support singlex、singleweight、singley + bool is_single_x = (tiling_context_->GetDynamicInputTensor(X_INDEX, 1) == nullptr); + bool is_single_weight = (tiling_context_->GetDynamicInputTensor(WEIGHT_INDEX, 1) == nullptr); + bool is_single_y = (tiling_context_->GetOutputShape(1) == nullptr); + transpose_weight = static_cast(*(tiling_context_->GetAttrs()->GetAttrPointer(0))); + if (!(is_single_x && is_single_weight && is_single_y)) { + OP_LOGE(tiling_context_->GetNodeName(), "only support singlex and singleweight and singley."); + return ge::GRAPH_FAILED; + } + + auto x_shape = tiling_context_->GetDynamicInputShape(X_INDEX, 0)->GetOriginShape(); + auto weight_shape = tiling_context_->GetDynamicInputShape(WEIGHT_INDEX, 0)->GetOriginShape(); + auto group_list_shape = tiling_context_->GetInputShape(GROUPLIST_INDEX)->GetOriginShape(); + auto y_shape = tiling_context_->GetOutputShape(0)->GetOriginShape(); + auto weight_desc = tiling_context_->GetDynamicInputDesc(WEIGHT_INDEX, 0); + auto weight_format = static_cast(ge::GetPrimaryFormat(weight_desc->GetStorageFormat())); + x_dtype = tiling_context_->GetDynamicInputDesc(X_INDEX, 0)->GetDataType(); + + // printf("weight_format %d\n", weight_format); + weight_nz = weight_format == ge::FORMAT_FRACTAL_NZ; + + // check input shape + if (x_shape.GetDimNum() != 2 || y_shape.GetDimNum() != 2) { + OP_LOGE(tiling_context_->GetNodeName(), "the dimNum of input and output should be 2, but got %zu, %zu.", static_cast(x_shape.GetDimNum()), static_cast(y_shape.GetDimNum())); + return ge::GRAPH_FAILED; + } + uint32_t weight_dim1, weight_dim2; + n_ = transpose_weight ? weight_shape.GetDim(1) : weight_shape.GetDim(2); + + if (group_list_shape.GetDimNum() != 2) { + OP_LOGE(tiling_context_->GetNodeName(), "only support key-value mode of groupList, the dimNum of groupList should be 2, but got %zu.", static_cast(group_list_shape.GetDimNum())); + return ge::GRAPH_FAILED; + } + + m_ = x_shape.GetDim(x_shape.GetDimNum() - 2); + k_ = x_shape.GetDim(x_shape.GetDimNum() - 1); + group_num_ = weight_shape.GetDim(0); + + if (weight_shape.GetDim(0) != group_num_) { + OP_LOGE(tiling_context_->GetNodeName(), "the dim0 of input weight should be equal to input groupList, but got %zu, %zu.", static_cast(weight_shape.GetDim(0)), static_cast(group_list_shape.GetDim(0))); + } + single_m_ = 128; + single_n_ = 256; + + core_num_ = aic_num; + auto n_task_num = (n_ + single_n_ - 1) / single_n_; + auto task_num = m_ * n_task_num; + if (task_num < core_num_) { + core_num_ = task_num; + } + + auto platform_info = platform_ascendc::PlatformAscendC(tiling_context_->GetPlatformInfo()); + sync_workspace_size_ = static_cast(platform_info.GetLibApiWorkSpaceSize()); + return ge::GRAPH_SUCCESS; +} + +void TilingMoeGroupedMatmulFunc::SetTilingKey() { + uint64_t tiling_key = 10UL; + if (transpose_weight) { + tiling_key = tiling_key + 1UL; + } + tiling_context_->SetTilingKey(tiling_key); +} + +void TilingMoeGroupedMatmulFunc::FillTilingData() { + tiling_data_.set_m(static_cast(m_)); + tiling_data_.set_n(static_cast(n_)); + tiling_data_.set_k(static_cast(k_)); + tiling_data_.set_single_m(single_m_); + tiling_data_.set_single_n(single_n_); + tiling_data_.set_group_num(group_num_); + tiling_data_.set_core_num(core_num_); +} + +ge::graphStatus TilingMoeGroupedMatmulFunc::RunKernelTiling() { + auto ret = CalMMTiling(); + if (ret != ge::GRAPH_SUCCESS) { + OP_LOGE(context->GetNodeName(), "cal mmtiling failed."); + return ge::GRAPH_FAILED; + } + ret = GMMSetMMTiling(); + if (ret != ge::GRAPH_SUCCESS) { + OP_LOGE(context->GetNodeName(), "gmm set mmtiling failed."); + return ge::GRAPH_FAILED; + } + SetTilingKey(); + FillTilingData(); + size_t userWorkspaceSize = 0; + size_t *currentWorkspace = tiling_context_->GetWorkspaceSizes(1); + currentWorkspace[0] = userWorkspaceSize + sync_workspace_size_; + tiling_data_.SaveToBuffer(tiling_context_->GetRawTilingData()->GetData(), + tiling_context_->GetRawTilingData()->GetCapacity()); + tiling_context_->GetRawTilingData()->SetDataSize(tiling_data_.GetDataSize()); + tiling_context_->SetBlockDim(core_num_); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus TilingForMoeGroupedMatmulFunc(gert::TilingContext *context){ + TilingMoeGroupedMatmulFunc tilingObject(context); + auto ret = tilingObject.Init(); + if(ret != ge::GRAPH_SUCCESS){ + OP_LOGE(context->GetNodeName(), "tiling Init failed."); + return ge::GRAPH_FAILED; + } + ret = tilingObject.RunKernelTiling(); + return ret; +} + +struct MatmulAllreduceAddRmsnormCompileInfo1 {}; +ge::graphStatus TilingParseForMatmulAllreduceAddRmsnorm1(gert::TilingParseContext *context) +{ + // (void)context; + return ge::GRAPH_SUCCESS; +} + + +IMPL_OP_OPTILING(MoeGroupedMatmul) + .Tiling(TilingForMoeGroupedMatmulFunc) + .TilingParse(TilingParseForMatmulAllreduceAddRmsnorm1); + +} // namespace optiling + diff --git a/csrc/ascend/moe/moe_grouped_matmul/op_host/moe_grouped_matmul_tiling.h b/csrc/ascend/moe/moe_grouped_matmul/op_host/moe_grouped_matmul_tiling.h new file mode 100644 index 000000000..a4c994317 --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/op_host/moe_grouped_matmul_tiling.h @@ -0,0 +1,26 @@ +/** +* This program is free software, you can redistribute it and/or modify. +* Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(MoeGroupedMatmulTilingData) +TILING_DATA_FIELD_DEF(uint32_t, group_num); +TILING_DATA_FIELD_DEF(uint32_t, core_num); +TILING_DATA_FIELD_DEF(uint32_t, m); +TILING_DATA_FIELD_DEF(uint32_t, n); +TILING_DATA_FIELD_DEF(uint32_t, k); +TILING_DATA_FIELD_DEF(uint32_t, single_m); +TILING_DATA_FIELD_DEF(uint32_t, single_n); +TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, mm_tiling); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(MoeGroupedMatmul, MoeGroupedMatmulTilingData) +} // namespace optiling diff --git a/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/aclnn_moe_grouped_matmul.cpp b/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/aclnn_moe_grouped_matmul.cpp new file mode 100644 index 000000000..0878230a8 --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/aclnn_moe_grouped_matmul.cpp @@ -0,0 +1,248 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#include "aclnn_moe_grouped_matmul.h" +#include "aclnn_moe_grouped_matmul_weight_nz.h" + +#include +#include + +#include "aclnn_kernels/transdata.h" +#include "moe_grouped_matmul_l0.h" +#include "aclnn_kernels/contiguous.h" +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "opdev/common_types.h" +#include "opdev/data_type_utils.h" +#include "opdev/format_utils.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/platform.h" +#include "opdev/shape_utils.h" +#include "opdev/tensor_view_utils.h" +#include "opdev/make_op_executor.h" + +using namespace op; + +#ifdef __cplusplus +extern "C" { +#endif + +namespace { + static constexpr size_t ALIGN_NZ_4BIT_N = 64UL; + static constexpr size_t ALIGN_NZ_4BIT_K = 64UL; + static constexpr size_t ALIGN_NZ_INT8_N = 32UL; + static constexpr size_t ALIGN_NZ_K = 16UL; + static constexpr size_t DIMS_THREE_FOR_GMM = 3UL; + static constexpr size_t LAST_FIRST_DIM_INDEX = 1; + static constexpr size_t LAST_SECOND_DIM_INDEX = 2; + static constexpr size_t LAST_THIRD_DIM_INDEX = 3; + + enum class GMMWeightVersion : uint32_t { + WeightNd = 1U, + WeightNz = 2U + }; + + struct MoeGroupedMatmulParams { + const aclTensorList *x = nullptr; + const aclTensorList *weight = nullptr; + const aclTensor *groupTensor = nullptr; + bool transposeWeight = false; + bool isSingleWeight = false; + GMMWeightVersion weightVersion = GMMWeightVersion::WeightNd; + const aclTensorList *y = nullptr; + DataType xDtype = DataType::DT_BF16; + }; +} + + +namespace { + +static aclnnStatus CheckNotNull(const aclTensorList *x, const aclTensorList *weight, const aclTensorList *y) { + CHECK_COND(x != nullptr, ACLNN_ERR_PARAM_NULLPTR, "x must not be nullptr."); + CHECK_COND(x->Size() != 0, ACLNN_ERR_PARAM_INVALID, "x must not be empty tensorlist."); + CHECK_COND(weight != nullptr, ACLNN_ERR_PARAM_NULLPTR, "weight must not be nullptr."); + CHECK_COND(weight->Size() != 0, ACLNN_ERR_PARAM_INVALID, "weight must not be empty tensorlist."); + CHECK_COND(y != nullptr, ACLNN_ERR_PARAM_NULLPTR, "y must not be nullptr."); + CHECK_COND(y->Size() != 0, ACLNN_ERR_PARAM_INVALID, "y must not be empty tensorlist."); + return ACLNN_SUCCESS; +} + +static aclnnStatus TransWeightToNzCheckAlign(MoeGroupedMatmulParams &gmmParams, const aclTensor *weight) +{ + size_t viewDimNum = weight->GetViewShape().GetDimNum(); + uint64_t k = gmmParams.transposeWeight ? weight->GetViewShape().GetDim(viewDimNum - 1) : + weight->GetViewShape().GetDim(viewDimNum - LAST_SECOND_DIM_INDEX); + uint64_t n = gmmParams.transposeWeight ? weight->GetViewShape().GetDim(viewDimNum - LAST_SECOND_DIM_INDEX) : + weight->GetViewShape().GetDim(viewDimNum - 1); + bool k_align = false; + bool n_align = false; + if (weight->GetDataType() == DataType::DT_BF16 || weight->GetDataType() == DataType::DT_FLOAT16) { + k_align = k % ALIGN_NZ_K == 0; + n_align = n % ALIGN_NZ_K == 0; + } + CHECK_COND(k_align == true && n_align == true, ACLNN_ERR_PARAM_INVALID, + "When weight format is FRACTAL_NZ, weight'shape(k[%lu], n[%lu]) should be divisible by the " + "following shape: BF16/FP16[16, 16]). If the weight is transposed," + "the k/n need to be reversed.", + k, n); + return ACLNN_SUCCESS; +} + +static aclnnStatus TransWeightToNz(MoeGroupedMatmulParams &gmmParams, aclOpExecutor *executor) { + const aclTensorList *&weights = gmmParams.weight; + const aclTensorList *&x = gmmParams.x; + CHECK_COND((*x)[0] != nullptr, ACLNN_ERR_PARAM_INVALID, "The first tensor of x is nullptr!"); + size_t wLength = weights->Size(); + for (size_t i(0); i < wLength; ++i) { + const aclTensor* weight = (*weights)[i]; + if (weight->GetStorageFormat() != op::Format::FORMAT_FRACTAL_NZ && + weight->GetStorageFormat() != op::Format::FORMAT_FRACTAL_NZ_C0_16 && + weight->GetStorageFormat() != op::Format::FORMAT_FRACTAL_NZ_C0_32) { + break; + } + TransWeightToNzCheckAlign(gmmParams, weight); + continue; + } + return ACLNN_SUCCESS; +} + +static const aclTensor *SetTensorToNZFormat(const aclTensor *input, op::Shape &shape, aclOpExecutor *executor) { + auto formatTensor = executor->CreateView(input, shape, input->GetViewOffset()); + formatTensor->SetStorageFormat(op::Format::FORMAT_FRACTAL_NZ); + formatTensor->SetOriginalFormat(input->GetViewFormat()); + formatTensor->SetViewShape(input->GetViewShape()); + return formatTensor; +} + +static aclnnStatus DataContiguous(const aclTensorList *&tensors, aclOpExecutor *executor) { + std::vector tensorsVec; + const aclTensor *contiguousTensor = nullptr; + for (size_t i = 0; i < tensors->Size(); ++i) { + const aclTensor *tensor = (*tensors)[i]; + contiguousTensor = l0op::Contiguous(tensor, executor); + CHECK_RET(contiguousTensor != nullptr, ACLNN_ERR_INNER_NULLPTR); + tensorsVec.push_back(contiguousTensor); + } + tensors = executor->AllocTensorList(tensorsVec.data(), tensorsVec.size()); + return ACLNN_SUCCESS; +} + +static aclnnStatus ParamsDataContiguous(MoeGroupedMatmulParams ¶ms, aclOpExecutor *executorPtr) { + CHECK_COND(DataContiguous(params.x, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "Contiguous x failed."); // make x contiguous + DataType xDtype = (*params.x)[0]->GetDataType(); + DataType weightDtype = (*params.weight)[0]->GetDataType(); + CHECK_COND(DataContiguous(params.weight, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "Contiguous weight failed."); // make w contiguous + params.groupTensor = l0op::Contiguous(params.groupTensor, executorPtr); + CHECK_COND(params.groupTensor != nullptr, ACLNN_ERR_PARAM_INVALID, + "Contiguous groupTensor failed."); + return ACLNN_SUCCESS; +} + +static aclnnStatus GetGMMResultByL0Api(MoeGroupedMatmulParams ¶ms, uint64_t *workspaceSize, aclOpExecutor **executor) { + auto uniqueExecutor = CREATE_EXECUTOR(); // fixed written style, create OpExecutor + aclOpExecutor *executorPtr = uniqueExecutor.get(); + CHECK_RET(executorPtr != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + // op::Shape wqbmmNzShape = (*params.weight)[0]->GetStorageShape(); + + CHECK_COND(ParamsDataContiguous(params, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "ParamsDataContiguous failed."); + if (params.weightVersion == GMMWeightVersion::WeightNz) { + std::vector tensorsVec; + for (size_t i = 0; i < params.weight->Size(); ++i) { + const aclTensor *tensor = (*params.weight)[i]; + op::Shape weightNzShape = tensor->GetViewShape(); + tensor = SetTensorToNZFormat(tensor, weightNzShape, executorPtr); + tensorsVec.push_back(tensor); + } + params.weight = executorPtr->AllocTensorList(tensorsVec.data(), tensorsVec.size()); + } + CHECK_COND(TransWeightToNz(params, executorPtr) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, + "TransWeightToNz failed."); + // Invoke l0 operator MoeGroupedMatmul for calculation. + auto result = l0op::MoeGroupedMatmul(params.x, params.weight, + params.groupTensor, params.transposeWeight, + (*params.y)[0]->GetViewShape(), params.y->Size(), + (*params.y)[0]->GetDataType(), executorPtr); + CHECK_RET(result != nullptr, ACLNN_ERR_INNER_NULLPTR); + for (size_t i(0); i < params.y->Size(); ++i) { + auto viewCopyResult = l0op::ViewCopy((*result)[i], (*params.y)[i], executorPtr); + CHECK_RET(viewCopyResult != nullptr, ACLNN_ERR_INNER_NULLPTR); + } + // Standard syntax, get the size of workspace needed during computation. + *workspaceSize = uniqueExecutor->GetWorkspaceSize(); + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; +} + + +static aclnnStatus aclnnMoeGroupedMatmulGetWorkspaceSizeCommon(const aclTensorList *x, const aclTensorList *weight, + const aclTensor *groupList, bool transposeWeight, GMMWeightVersion weightVersion, + const aclTensorList *y, uint64_t *workspaceSize, aclOpExecutor **executor) { + DataType xDtype = DataType::DT_UNDEFINED; + for (size_t i = 0; i < x->Size(); ++i) { + if ((*x)[i] != nullptr) { + xDtype = (*x)[i]->GetDataType(); + break; + } + } + MoeGroupedMatmulParams moeGmmParams{x, weight, groupList, transposeWeight, true, weightVersion, y, xDtype}; + aclnnStatus ret = GetGMMResultByL0Api(moeGmmParams, workspaceSize, executor); + return ret; +} +} + +aclnnStatus aclnnMoeGroupedMatmulWeightNzGetWorkspaceSize(const aclTensorList *x, const aclTensorList *weight, + const aclTensor *groupList, bool transposeWeight, aclTensorList *out, + uint64_t *workspaceSize, aclOpExecutor **executor) { + CHECK_COND(CheckNotNull(x, weight, out) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR, + "one of required inputs is nullptr."); + // Standard syntax, Check parameters. + L2_DFX_PHASE_1(aclnnMoeGroupedMatmulWeightNz, + DFX_IN(x, weight, groupList), + DFX_OUT(out)); + return aclnnMoeGroupedMatmulGetWorkspaceSizeCommon(x, weight, groupList, transposeWeight, GMMWeightVersion::WeightNz, out, workspaceSize, executor); +} + +aclnnStatus aclnnMoeGroupedMatmulGetWorkspaceSize(const aclTensorList *x, const aclTensorList *weight, + const aclTensor *groupList, bool transposeWeight, aclTensorList *out, + uint64_t *workspaceSize, aclOpExecutor **executor) { + CHECK_COND(CheckNotNull(x, weight, out) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR, + "one of required inputs is nullptr."); + // Standard syntax, Check parameters. + L2_DFX_PHASE_1(aclnnMoeGroupedMatmul, + DFX_IN(x, weight, groupList), + DFX_OUT(out)); + CHECK_COND(weight->Size() != 0, ACLNN_ERR_PARAM_INVALID, "weight should not be null tensorlist "); + return aclnnMoeGroupedMatmulGetWorkspaceSizeCommon(x, weight, groupList, transposeWeight, GMMWeightVersion::WeightNd, out, workspaceSize, executor); +} + +aclnnStatus aclnnMoeGroupedMatmul(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, + aclrtStream stream) { + L2_DFX_PHASE_2(aclnnMoeGroupedMatmul); + CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER, + "This is an error in GMM launch aicore"); + return ACLNN_SUCCESS; +} + +aclnnStatus aclnnMoeGroupedMatmulWeightNz(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, + aclrtStream stream) { + L2_DFX_PHASE_2(aclnnMoeGroupedMatmulWeightNz); + CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER, + "This is an error in GMM launch aicore"); + return ACLNN_SUCCESS; +} + +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/aclnn_moe_grouped_matmul.h b/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/aclnn_moe_grouped_matmul.h new file mode 100644 index 000000000..cad1ecccc --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/aclnn_moe_grouped_matmul.h @@ -0,0 +1,30 @@ +/** +* This program is free software, you can redistribute it and/or modify. +* Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_API_INC_MOE_GROUPED_MATMUL_H +#define OP_API_INC_MOE_GROUPED_MATMUL_H +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +__attribute__((visibility("default"))) aclnnStatus aclnnMoeGroupedMatmulGetWorkspaceSize( + const aclTensorList *x, const aclTensorList *weight, const aclTensor *groupList, + bool transposeWeight, aclTensorList *out, uint64_t *workspaceSize, aclOpExecutor **executor); + +__attribute__((visibility("default"))) aclnnStatus aclnnMoeGroupedMatmul(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, + aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/aclnn_moe_grouped_matmul_weight_nz.h b/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/aclnn_moe_grouped_matmul_weight_nz.h new file mode 100644 index 000000000..e1a915a78 --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/aclnn_moe_grouped_matmul_weight_nz.h @@ -0,0 +1,30 @@ +/** +* This program is free software, you can redistribute it and/or modify. +* Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_API_INC_MOE_GROUPED_MATMUL_WEIGHT_NZ_H +#define OP_API_INC_MOE_GROUPED_MATMUL_WEIGHT_NZ_H +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +__attribute__((visibility("default"))) aclnnStatus aclnnMoeGroupedMatmulWeightNzGetWorkspaceSize( + const aclTensorList *x, const aclTensorList *weight, const aclTensor *groupList, + bool transposeWeight, aclTensorList *out, uint64_t *workspaceSize, aclOpExecutor **executor); + +__attribute__((visibility("default"))) aclnnStatus aclnnMoeGroupedMatmulWeightNz(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, + aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/moe_grouped_matmul_l0.cpp b/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/moe_grouped_matmul_l0.cpp new file mode 100644 index 000000000..397a9ae73 --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/moe_grouped_matmul_l0.cpp @@ -0,0 +1,50 @@ +#include "moe_grouped_matmul_l0.h" +#include "opdev/op_log.h" +#include "opdev/op_dfx.h" +#include "opdev/shape_utils.h" +#include "opdev/make_op_executor.h" + +using namespace op; + +namespace l0op { +OP_TYPE_REGISTER(MoeGroupedMatmul); + +const aclTensorList *MoeGroupedMatmul(const aclTensorList *x, + const aclTensorList *weight, + const aclTensor *groupList, + bool transposeWeight, + op::Shape yShape, + size_t outLength, + op::DataType yDtype, + aclOpExecutor *executor) { + L0_DFX(MoeGroupedMatmul, x, weight, groupList, transposeWeight, outLength); + + std::vector tensorsVec; + const aclTensor *x0 = x->Size() > 0 ? (*x)[0] : nullptr; + if (x0 == nullptr) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "(*x)[0] is nullptr."); + return nullptr; + } + + for (size_t i(0); i < outLength; ++i) { + tensorsVec.emplace_back(executor->AllocTensor(yShape, yDtype)); + } + auto out = executor->AllocTensorList(tensorsVec.data(), outLength); + + auto x0_dim_num = x0->GetStorageShape().GetDimNum(); + auto x0_dim0 = x0->GetStorageShape().GetDim(0); + printf("x0_dim_num %d x0_dim0 %d\n", x0_dim_num, x0_dim0); + + auto ret = ADD_TO_LAUNCHER_LIST_AICORE(MoeGroupedMatmul, + OP_INPUT(x, weight, groupList), + OP_OUTPUT(out), + OP_ATTR(transposeWeight)); + if (ret != ACLNN_SUCCESS) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "ADD_TO_LAUNCHER_LIST_AICORE failed."); + return nullptr; + } + + return out; +} + +} // namespace l0op \ No newline at end of file diff --git a/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/moe_grouped_matmul_l0.h b/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/moe_grouped_matmul_l0.h new file mode 100644 index 000000000..6dff4dc80 --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/op_host/op_api/moe_grouped_matmul_l0.h @@ -0,0 +1,27 @@ +/** +* This program is free software, you can redistribute it and/or modify. +* Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_API_INC_LEVEL0_OP_MOE_GROUPED_MATMUL_OP_H +#define OP_API_INC_LEVEL0_OP_MOE_GROUPED_MATMUL_OP_H + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensorList *MoeGroupedMatmul(const aclTensorList *x, + const aclTensorList *weight, + const aclTensor *groupList, + bool transposeWeight, + op::Shape yShape, + size_t outLength, + op::DataType yDtype, + aclOpExecutor *executor); +} + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/moe_grouped_matmul/op_kernel/moe_grouped_matmul.cpp b/csrc/ascend/moe/moe_grouped_matmul/op_kernel/moe_grouped_matmul.cpp new file mode 100644 index 000000000..579dbe1ab --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/op_kernel/moe_grouped_matmul.cpp @@ -0,0 +1,42 @@ +/** +* This program is free software, you can redistribute it and/or modify. +* Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "moe_grouped_matmul.h" +#include "kernel_operator.h" + +#if defined(FORMAT_WEIGHT) && FORMAT_WEIGHT == FORMAT_FRACTAL_NZ +constexpr CubeFormat formatWeight = CubeFormat::NZ; +#else +constexpr CubeFormat formatWeight = CubeFormat::ND; +#endif + +//using namespace matmul; +#define GMM_CUBE_IMP(transWeight) \ + do { \ + if ASCEND_IS_AIV { \ + return; \ + } \ + GET_TILING_DATA_WITH_STRUCT(MoeGroupedMatmulTilingData, tiling_data, tiling); \ + AscendC::TPipe pipe; \ + KernelMoeGMMNoQuant op(&pipe); \ + op.Init(x, weight, group_list, y, &tiling_data); \ + op.Process(); \ + } while (0) + +extern "C" __global__ __aicore__ void moe_grouped_matmul(GM_ADDR x, GM_ADDR weight, GM_ADDR group_list, GM_ADDR y, + GM_ADDR workSpace, GM_ADDR tiling) { + REGISTER_TILING_DEFAULT(MoeGroupedMatmulTilingData); + + if (TILING_KEY_IS(10UL)) { + GMM_CUBE_IMP(false); + } else if (TILING_KEY_IS(11UL)) { + GMM_CUBE_IMP(true); + } +} diff --git a/csrc/ascend/moe/moe_grouped_matmul/op_kernel/moe_grouped_matmul.h b/csrc/ascend/moe/moe_grouped_matmul/op_kernel/moe_grouped_matmul.h new file mode 100644 index 000000000..ddd86bf4e --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/op_kernel/moe_grouped_matmul.h @@ -0,0 +1,186 @@ +/** +* This program is free software, you can redistribute it and/or modify. +* Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "kernel_operator.h" +#include "kernel_operator_list_tensor_intf.h" +#include "lib/matmul_intf.h" +#include "moe_grouped_matmul_tiling_data.h" +using namespace AscendC; + + +constexpr MatmulConfig matmulCFGUnitFlag{false, false, true, 0, 0, 0, false, false, false, false, false, 0, 0, 0, + 0, 0, 0, 0, true}; +struct GMMConfig { + uint32_t m = 0; + uint32_t k = 0; + uint32_t n = 0; + uint32_t baseM = 0; + uint32_t baseN = 0; + uint32_t mIdx = 0; + uint32_t nIdx = 0; + uint32_t blockDimM = 0; + uint32_t blockDimN = 0; + uint32_t singleM = 0; + uint32_t singleN = 0; + uint64_t wBaseOffset = 0; + uint64_t nAxisBaseOffset = 0; + uint64_t mAxisBaseOffset = 0; + uint64_t xBaseOffset = 0; + uint64_t yBaseOffset = 0; + uint64_t wOutOffset = 0; +}; + + +template +class KernelMoeGMMNoQuant { + +protected: + using xType = MatmulType; + using weightType = MatmulType; + using yType = MatmulType; + using biasType = MatmulType; + using mmT = matmul::MatmulImpl; + mmT mm; + + MoeGroupedMatmulTilingData tiling_; + AscendC::TPipe *pipe_ = nullptr; + + GlobalTensor x_gm_; + GlobalTensor weight_gm_; + GlobalTensor y_gm_; + GlobalTensor group_list_gm_; + ListTensorDesc x_list_; + ListTensorDesc weight_list_; + ListTensorDesc y_list_; + + uint32_t core_idx; + uint32_t used_core_num; + constexpr static bool transposeW = transWeight; + constexpr static uint32_t UB_BLOCK_UNIT_SIZE = 32; + +public: + __aicore__ inline KernelMoeGMMNoQuant(AscendC::TPipe *pipe) {pipe_ = pipe;} + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR group_list, GM_ADDR y, const MoeGroupedMatmulTilingData *tiling) { + core_idx = GetBlockIdx(); + tiling_ = *tiling; + used_core_num = GetBlockNum(); + group_list_gm_.SetGlobalBuffer((__gm__ T2*)group_list); + x_list_.Init((__gm__ void*)x); + weight_list_.Init((__gm__ void*)weight); + y_list_.Init((__gm__ void*)y); + GM_ADDR x_first_addr = (__gm__ uint8_t*)x_list_.GetDataPtr<__gm__ uint8_t>(0); + GM_ADDR weight_first_addr = (__gm__ uint8_t*)weight_list_.GetDataPtr<__gm__ uint8_t>(0); + GM_ADDR y_first_addr = (__gm__ uint8_t*)y_list_.GetDataPtr<__gm__ uint8_t>(0); + x_gm_.SetGlobalBuffer((__gm__ T*)x_first_addr); + weight_gm_.SetGlobalBuffer((__gm__ T*)weight_first_addr); + y_gm_.SetGlobalBuffer((__gm__ T*)y_first_addr); + + mm.Init(&tiling_.mm_tiling, pipe_); + } + + __aicore__ inline void Process() { + + uint32_t group_list_inner_shape = 2u; + uint32_t group_list_shape_size = tiling_.group_num * group_list_inner_shape; + GMMConfig mn_config; + + for (uint32_t loop = 0, count = 0; loop < group_list_shape_size; loop += group_list_inner_shape) { + int32_t split_value = static_cast(group_list_gm_.GetValue(loop + 1)); + if (split_value <= 0) { + break; + } + uint32_t group_idx = static_cast(group_list_gm_.GetValue(loop)); + mn_config.mAxisBaseOffset += mn_config.m; + mn_config.xBaseOffset += mn_config.m * mn_config.k; + mn_config.yBaseOffset += mn_config.m * mn_config.n; + this->SetMNConfig(split_value, mn_config); + mn_config.nAxisBaseOffset = group_idx * mn_config.n; + if constexpr (formatWeight == CubeFormat::NZ) { + mn_config.wBaseOffset = AlignUp(mn_config.k, 16) * AlignUp(mn_config.nAxisBaseOffset, 16); + } else { + mn_config.wBaseOffset = mn_config.k * mn_config.nAxisBaseOffset; + } + mn_config.blockDimM = Ceil(mn_config.m, mn_config.singleM); + mn_config.blockDimN = Ceil(mn_config.n, mn_config.singleN); + uint32_t cur_count = count + mn_config.blockDimM * mn_config.blockDimN; + uint32_t cur_block = this->core_idx >= count ? this->core_idx : this->core_idx + used_core_num; + while (cur_block < cur_count) { + mn_config.mIdx = (cur_block - count) / mn_config.blockDimN; + mn_config.nIdx = (cur_block - count) % mn_config.blockDimN; + this->MMCompute(group_idx, mn_config, this->core_idx); + cur_block += used_core_num; + } + count = cur_count % used_core_num; + } + } + +protected: + __aicore__ inline uint32_t AlignUp(uint32_t a, uint32_t base) { + return (a + base - 1) / base * base; + } + + __aicore__ inline uint32_t Ceil(uint32_t a, uint32_t base) { + if (base == 0) { + return a; + } + return (a + base - 1) / base; + } + + __aicore__ inline void SetMNConfig(const int32_t split_value, GMMConfig & mn_config) { + mn_config.m = split_value; + mn_config.k = tiling_.k; + mn_config.n = tiling_.n; + mn_config.baseM = tiling_.single_m; + mn_config.baseN = tiling_.single_n; + mn_config.singleM = mn_config.baseM; + mn_config.singleN = mn_config.baseN; + } + + __aicore__ inline void MMCompute(uint32_t group_idx, GMMConfig& mn_config, uint32_t core_idx) { + uint32_t tail_n = mn_config.nIdx * mn_config.singleN; + uint32_t cur_single_n = mn_config.nIdx < mn_config.blockDimN - 1 ? mn_config.singleN : mn_config.n - tail_n; + uint32_t cur_single_m = mn_config.mIdx < mn_config.blockDimM - 1 ? mn_config.singleM + : mn_config.m - mn_config.mIdx * mn_config.singleM; + uint64_t x_offset = mn_config.mIdx * mn_config.singleM * mn_config.k; + uint64_t out_offset = mn_config.mIdx * mn_config.singleM * mn_config.n + tail_n; + GlobalTensor weight_gm_local = GetGlobalBufferW(group_idx, tail_n, mn_config); + + mm.SetOrgShape(mn_config.m, mn_config.n, mn_config.k); + mm.SetSingleShape(cur_single_m, cur_single_n, mn_config.k); + mm.SetTensorA(x_gm_[mn_config.xBaseOffset + x_offset], false); + mm.SetTensorB(weight_gm_local, transposeW); + mm.template IterateAll(y_gm_[mn_config.yBaseOffset + out_offset], 0); + } + + __aicore__ inline GlobalTensor GetGlobalBufferW(uint32_t group_idx, uint32_t tail_n, GMMConfig& mn_config) { + uint64_t w_offset = SetWOffset(tail_n, mn_config.k); + GlobalTensor weight_gm_local; + weight_gm_local = weight_gm_[mn_config.wBaseOffset + w_offset]; + if (mn_config.blockDimM == 1) { + weight_gm_local.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE); + } + return weight_gm_local; + } + + __aicore__ inline uint64_t SetWOffset(uint32_t tail_n, uint32_t k) { + uint64_t w_offset = 0; + if constexpr (formatWeight == CubeFormat::NZ && transposeW) { + w_offset = tail_n * (UB_BLOCK_UNIT_SIZE / sizeof(T)); // 32: quant is 32, float16 is 16 + } else if constexpr (formatWeight == CubeFormat::NZ) { + w_offset = tail_n * AlignUp(k, 16); // 16: nz format last two dim size + } else if constexpr (transposeW) { + w_offset = tail_n * k; + } else { + w_offset = tail_n; + } + return w_offset; + } +}; diff --git a/csrc/ascend/moe/moe_grouped_matmul/op_kernel/moe_grouped_matmul_tiling_data.h b/csrc/ascend/moe/moe_grouped_matmul/op_kernel/moe_grouped_matmul_tiling_data.h new file mode 100644 index 000000000..61747b090 --- /dev/null +++ b/csrc/ascend/moe/moe_grouped_matmul/op_kernel/moe_grouped_matmul_tiling_data.h @@ -0,0 +1,21 @@ +#ifndef MOE_GROUPED_MATMUL_TILING_DATA_H +#define MOE_GROUPED_MATMUL_TILING_DATA_H + +#include + +#include "kernel_tiling/kernel_tiling.h" + +#pragma pack(push, 8) +struct alignas(8) MoeGroupedMatmulTilingData { + uint32_t group_num; + uint32_t core_num; + uint32_t m; + uint32_t n; + uint32_t k; + uint32_t single_m; + uint32_t single_n; + TCubeTiling mm_tiling; +}; +#pragma pack(pop) + +#endif diff --git a/csrc/ascend/moe/moe_init_routing_custom/CMakeLists.txt b/csrc/ascend/moe/moe_init_routing_custom/CMakeLists.txt new file mode 100644 index 000000000..298ed5677 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/moe_init_routing_custom_torch_adpt.h b/csrc/ascend/moe/moe_init_routing_custom/moe_init_routing_custom_torch_adpt.h new file mode 100644 index 000000000..eb8d066b7 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/moe_init_routing_custom_torch_adpt.h @@ -0,0 +1,119 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef MOE_INIT_ROUTING_CUSTOM_TORCH_ADPT_H +#define MOE_INIT_ROUTING_CUSTOM_TORCH_ADPT_H +namespace vllm_fl { +std::tuple npu_moe_init_routing_custom( + const at::Tensor &x, const at::Tensor &expert_idx, + const c10::optional &scale, const c10::optional &offset, int64_t active_num, + int64_t expert_capacity, int64_t expert_num, int64_t drop_pad_mode, int64_t expert_tokens_num_type, + bool expert_tokens_num_flag, int64_t quant_mode, at::IntArrayRef active_expert_range, int64_t row_idx_type) +{ + constexpr int64_t DIM_X = 2; + constexpr int64_t DIM_EXPERT_IDX = 2; + constexpr int64_t LENGTH_ACTIVE_EXPERT_RANGE = 2; + constexpr int64_t EXPERT_TOKENS_COUNT = 1; + constexpr int64_t EXPERT_TOKENS_KEY_VALUE = 2; + constexpr int64_t QUANT_MODE_UNQUANT = -1; + constexpr int64_t QUANT_MODE_DYNAMIC_QUANT = 1; + constexpr int64_t CUMSUM = 0; + constexpr int64_t COUNT = 1; + constexpr int64_t KEY_VALUE = 2; + + if (active_expert_range.empty()) { + active_expert_range = at::IntArrayRef({0, expert_num}); + } + + int64_t x_dim = x.dim(); + TORCH_CHECK(x_dim == DIM_X, "The x should be ", DIM_X, + "-Dimension, current is ", x_dim, "-Dimension."); + + int64_t expert_idx_dim = expert_idx.dim(); + TORCH_CHECK(expert_idx_dim == DIM_EXPERT_IDX, "The expert_idx should be ", DIM_EXPERT_IDX, + "-Dimension, current is ", expert_idx_dim, "-Dimension."); + + int64_t active_expert_range_length = active_expert_range.size(); + TORCH_CHECK(active_expert_range_length == LENGTH_ACTIVE_EXPERT_RANGE, "The active_expert_range should be ", LENGTH_ACTIVE_EXPERT_RANGE, + "-Dimension, current is ", expert_idx_dim, "-Dimension."); + + int expert_length = active_expert_range[1] - active_expert_range[0]; + auto x_size = x.sizes(); + auto expert_idx_size = expert_idx.sizes(); + + int bs = x_size[0]; + int h = x_size[1]; + int k = expert_idx_size[1]; + int64_t expanded_scale_len = 0; + at::Tensor expanded_x; + + if (drop_pad_mode == 1) { // Drop/Pad + if (quant_mode == QUANT_MODE_UNQUANT) { + expanded_x = at::empty({expert_num, expert_capacity, h}, x.options()); + } else { + expanded_x = at::empty({expert_num, expert_capacity, h}, x.options().dtype(at::kChar)); + } + expanded_scale_len = expert_num * expert_capacity; + } else { // Dropless / Active + if (active_num > 0) { // Active + int64_t num_out_tokens = std::min((int64_t)bs * k, active_num); + if (quant_mode == QUANT_MODE_UNQUANT) { + expanded_x = at::empty({num_out_tokens, h}, x.options()); + } else { + expanded_x = at::empty({num_out_tokens, h}, x.options().dtype(at::kChar)); + } + expanded_scale_len = num_out_tokens; + } else { // Dropless + if (quant_mode == QUANT_MODE_UNQUANT) { + expanded_x = at::empty({bs * k, h}, x.options()); + } else { + expanded_x = at::empty({bs * k, h}, x.options().dtype(at::kChar)); + } + expanded_scale_len = bs * k; + } + } + + at::Tensor expanded_row_idx = at::empty({bs * k}, expert_idx.options()); + at::Tensor expert_tokens_count_or_cumsum; + if (expert_tokens_num_type >= CUMSUM && expert_tokens_num_type <= COUNT) { + // expert_tokens_count_or_cumsum in [end-start, ] + expert_tokens_count_or_cumsum = at::empty({expert_length}, x.options().dtype(at::kLong)); + } else if (expert_tokens_num_type == KEY_VALUE) { + // key_value in [2, end-start] + expert_tokens_count_or_cumsum = at::empty({expert_num, 2}, x.options().dtype(at::kLong)); + } + at::Tensor expanded_scale = at::empty({expanded_scale_len}, x.options().dtype(at::kFloat)); + EXEC_NPU_CMD(aclnnMoeInitRoutingCustom, + x, + expert_idx, + scale, + offset, + active_num, + expert_capacity, + expert_num, + drop_pad_mode, + expert_tokens_num_type, + expert_tokens_num_flag, + quant_mode, + active_expert_range, + row_idx_type, + expanded_x, + expanded_row_idx, + expert_tokens_count_or_cumsum, + expanded_scale); + return std::tie(expanded_x, expanded_row_idx, expert_tokens_count_or_cumsum, expanded_scale); +} +} +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_host/CMakeLists.txt b/csrc/ascend/moe/moe_init_routing_custom/op_host/CMakeLists.txt new file mode 100644 index 000000000..56f3ad2cb --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_host/CMakeLists.txt @@ -0,0 +1,22 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + moe_init_routing_custom_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME MoeInitRoutingCustom + OPTIONS + --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE moe_init_routing_custom ACLNNTYPE aclnn) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_host/aclnn_moe_init_routing_custom.cpp b/csrc/ascend/moe/moe_init_routing_custom/op_host/aclnn_moe_init_routing_custom.cpp new file mode 100644 index 000000000..6564a58e9 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_host/aclnn_moe_init_routing_custom.cpp @@ -0,0 +1,143 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include +#include +#include +#include "opdev/make_op_executor.h" +#include "aclnn_kernels/contiguous.h" +#include "opdev/tensor_view_utils.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "opdev/op_log.h" +#include "aclnn_kernels/cast.h" +#include "opdev/common_types.h" +#include "moe_init_routing_custom.h" +#include "aclnn_moe_init_routing_custom.h" + +using namespace op; + +#ifdef __cplusplus +extern "C" { +#endif + +namespace { + static const int64_t MOE_DIM_2 = 2; + static const int64_t MOE_DIM_1 = 1; +} + +static const std::initializer_list DTYPE_SUPPORT_LIST_X= {DataType::DT_FLOAT16, DataType::DT_BF16, DataType::DT_FLOAT, DataType::DT_INT8}; +static const std::initializer_list DTYPE_SUPPORT_LIST_EXPERT_IDX = {DataType::DT_INT32}; +static const std::initializer_list DTYPE_SUPPORT_LIST_SCALE = {DataType::DT_FLOAT}; +static const std::initializer_list DTYPE_SUPPORT_LIST_OFFSET= {DataType::DT_FLOAT}; +static const std::initializer_list DTYPE_SUPPORT_LIST_EXPANDED_X_OUT = {DataType::DT_FLOAT16, DataType::DT_BF16, DataType::DT_FLOAT, DataType::DT_INT8}; +static const std::initializer_list DTYPE_SUPPORT_LIST_EXPANDED_ROW_IDX_OUT = {DataType::DT_INT32}; +static const std::initializer_list DTYPE_SUPPORT_LIST_EXPERT_TOKENS_COUNT_OR_CUMSUMOUT = {DataType::DT_INT64}; +static const std::initializer_list DTYPE_SUPPORT_LIST_EXPANDED_SCALE_OUT = {DataType::DT_FLOAT}; + +static inline bool CheckNotNull(const aclTensor *x, + const aclTensor *expertIdx, + const aclTensor *expandedXOut, + const aclTensor *expandedRowIdxOut, + const aclTensor *expertTokensCountOrCumsumOut, + const aclTensor *expandedScaleOut) { + OP_CHECK_NULL(x, return false); + OP_CHECK_NULL(expertIdx, return false); + OP_CHECK_NULL(expandedXOut, return false); + OP_CHECK_NULL(expandedRowIdxOut, return false); + OP_CHECK_NULL(expertTokensCountOrCumsumOut, return false); + OP_CHECK_NULL(expandedScaleOut, return false); + + return true; +} + +aclnnStatus aclnnMoeInitRoutingCustomGetWorkspaceSize(const aclTensor *x, + const aclTensor *expertIdx, + const aclTensor *scaleOptional, + const aclTensor *offsetOptional, + int64_t activeNum, + int64_t expertCapacity, + int64_t expertNum, + int64_t dropPadMode, + int64_t expertTokensNumType, + bool expertTokensNumFlag, + int64_t quantMode, + const aclIntArray *activeExpertRangeOptional, + int64_t rowIdxType, + const aclTensor *expandedXOut, + const aclTensor *expandedRowIdxOut, + const aclTensor *expertTokensCountOrCumsumOut, + const aclTensor *expandedScaleOut, + uint64_t *workspaceSize, + aclOpExecutor **executor) +{ + L2_DFX_PHASE_1(aclnnMoeInitRoutingCustom, + DFX_IN(x, expertIdx, scaleOptional, offsetOptional, + activeNum, expertCapacity, expertNum, dropPadMode, + expertTokensNumType, expertTokensNumFlag, quantMode, activeExpertRangeOptional, rowIdxType), + DFX_OUT(expandedXOut, expandedRowIdxOut, expertTokensCountOrCumsumOut, expandedScaleOut)); + auto ret = CheckNotNull(x, expertIdx, expandedXOut, expandedRowIdxOut, + expertTokensCountOrCumsumOut, expandedScaleOut); + + CHECK_RET(ret, ACLNN_ERR_PARAM_NULLPTR); + + auto uniqueExecutor = CREATE_EXECUTOR(); + CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + + auto xContiguous = l0op::Contiguous(x, uniqueExecutor.get()); + CHECK_RET(xContiguous != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + auto expertIdxContiguous = l0op::Contiguous(expertIdx, uniqueExecutor.get()); + CHECK_RET(expertIdxContiguous != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + + const aclTensor* scaleContiguous = nullptr; + const aclTensor* offsetContiguous = nullptr; + if (scaleOptional != nullptr) { + scaleContiguous = l0op::Contiguous(scaleOptional, uniqueExecutor.get()); + CHECK_RET(scaleContiguous != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + } + + if (offsetOptional != nullptr) { + offsetContiguous = l0op::Contiguous(offsetOptional, uniqueExecutor.get()); + CHECK_RET(offsetContiguous != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + } + + auto routingResult = std::tuple(nullptr, nullptr, nullptr, nullptr); + routingResult = l0op::MoeInitRoutingCustom(xContiguous, expertIdxContiguous, scaleContiguous, offsetContiguous, + activeNum, expertCapacity, expertNum, dropPadMode, expertTokensNumType, expertTokensNumFlag, + quantMode, activeExpertRangeOptional, rowIdxType, expandedXOut, expandedRowIdxOut, + expertTokensCountOrCumsumOut, expandedScaleOut, uniqueExecutor.get()); + auto [expandedXOut_, expandedRowIdxOut_, expertTokensCountOrCumsumOut_, expandedScaleOut_] = routingResult; + bool hasNullptr = (expandedXOut_ == nullptr) || (expandedRowIdxOut_ == nullptr) || (expertTokensCountOrCumsumOut_ == nullptr) || (expandedScaleOut_ == nullptr); + CHECK_RET(hasNullptr != true, ACLNN_ERR_INNER_NULLPTR); + + auto viewCopyExpandedXOutResult = l0op::ViewCopy(expandedXOut_, expandedXOut, uniqueExecutor.get()); + CHECK_RET(viewCopyExpandedXOutResult != nullptr, ACLNN_ERR_INNER_NULLPTR); + auto viewCopyExpandedRowIdxOutResult = l0op::ViewCopy(expandedRowIdxOut_, expandedRowIdxOut, uniqueExecutor.get()); + CHECK_RET(viewCopyExpandedRowIdxOutResult != nullptr, ACLNN_ERR_INNER_NULLPTR); + + auto viewCopyExpertTokensCountOrCumsumOutResult = l0op::ViewCopy(expertTokensCountOrCumsumOut_, expertTokensCountOrCumsumOut, uniqueExecutor.get()); + CHECK_RET(viewCopyExpertTokensCountOrCumsumOutResult != nullptr, ACLNN_ERR_INNER_NULLPTR); + + auto viewCopyExpandedScaleOutResult = l0op::ViewCopy(expandedScaleOut_, expandedScaleOut, uniqueExecutor.get()); + CHECK_RET(viewCopyExpandedScaleOutResult != nullptr, ACLNN_ERR_INNER_NULLPTR); + + *workspaceSize = uniqueExecutor->GetWorkspaceSize(); + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; +} +aclnnStatus aclnnMoeInitRoutingCustom(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, + aclrtStream stream) +{ + L2_DFX_PHASE_2(aclnnMoeInitRoutingCustom); + return CommonOpExecutorRun(workspace, workspaceSize, executor, stream); +} + +#ifdef __cplusplus +} +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_host/aclnn_moe_init_routing_custom.h b/csrc/ascend/moe/moe_init_routing_custom/op_host/aclnn_moe_init_routing_custom.h new file mode 100644 index 000000000..5c7106b5d --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_host/aclnn_moe_init_routing_custom.h @@ -0,0 +1,47 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_API_INC_MOE_INIT_ROUTING_CUSTOM_H_ +#define OP_API_INC_MOE_INIT_ROUTING_CUSTOM_H_ + +#include "aclnn/aclnn_base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +__attribute__((visibility("default"))) aclnnStatus aclnnMoeInitRoutingCustomGetWorkspaceSize(const aclTensor *x, + const aclTensor *expertIdx, + const aclTensor *scaleOptional, + const aclTensor *offsetOptional, + int64_t activeNum, + int64_t expertCapacity, + int64_t expertNum, + int64_t dropPadMode, + int64_t expertTokensNumType, + bool expertTokensNumFlag, + int64_t quantMode, + const aclIntArray *activeExpertRangeOptional, + int64_t rowIdxType, + const aclTensor *expandedXOut, + const aclTensor *expandedRowIdxOut, + const aclTensor *expertTokensCountOrCumsumOut, + const aclTensor *expandedScaleOut, + uint64_t *workspaceSize, + aclOpExecutor **executor); + +__attribute__((visibility("default"))) aclnnStatus aclnnMoeInitRoutingCustom(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, + aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom.cpp b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom.cpp new file mode 100644 index 000000000..df36f9d41 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom.cpp @@ -0,0 +1,50 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include +#include "moe_init_routing_custom.h" +#include "opdev/make_op_executor.h" +#include "opdev/op_def.h" +#include "opdev/op_dfx.h" +#include "opdev/op_executor.h" +#include "opdev/op_log.h" +#include "opdev/shape_utils.h" +#include "aclnn_kernels/common/op_error_check.h" + +using namespace op; + +namespace l0op { +OP_TYPE_REGISTER(MoeInitRoutingCustom); + +std::tuple MoeInitRoutingCustom(const aclTensor *x, const aclTensor *expertIdx, const aclTensor *scale, + const aclTensor *offset, int64_t activeNum, int64_t expertCapacity, + int64_t expertNum, int64_t dropPadMode, int64_t expertTokensNumType, + bool expertTokensNumFlag, int64_t quantMode, const aclIntArray *activeExpertRange, + int64_t rowIdxType, const aclTensor *expandedX, const aclTensor *expandedRowIdx, + const aclTensor *expertTokensCountOrCumsum, const aclTensor *expandedScale, aclOpExecutor *executor) +{ + L0_DFX(MoeInitRoutingCustom, x, expertIdx, scale, offset, activeNum, expertCapacity, expertNum, dropPadMode, expertTokensNumType, expertTokensNumFlag, + quantMode, activeExpertRange, rowIdxType, expandedX, expandedRowIdx, expertTokensCountOrCumsum, expandedScale); + + auto expandedXOut = executor->AllocTensor(expandedX->GetViewShape(), expandedX->GetDataType(), Format::FORMAT_ND); + auto expandedRowIdxOut = executor->AllocTensor(expandedRowIdx->GetViewShape(), expandedRowIdx->GetDataType(), Format::FORMAT_ND); + auto expertTokensCountOrCumsumOut = executor->AllocTensor(expertTokensCountOrCumsum->GetViewShape(), expertTokensCountOrCumsum->GetDataType(), Format::FORMAT_ND); + auto expandedScaleOut = executor->AllocTensor(expandedScale->GetViewShape(), expandedScale->GetDataType(), Format::FORMAT_ND); + if (expandedXOut == nullptr || expandedRowIdxOut == nullptr || expertTokensCountOrCumsumOut == nullptr || expandedScaleOut == nullptr) { + OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "alloc expandedXOut or expandedRowIdxOut or expertTokensCountOrCumsumOut or expandedScaleOut tensor failed."); + return std::tuple(nullptr, nullptr, nullptr, nullptr); + } + + ADD_TO_LAUNCHER_LIST_AICORE( + MoeInitRoutingCustom, OP_INPUT(x, expertIdx, scale, offset), OP_OUTPUT(expandedXOut, expandedRowIdxOut, expertTokensCountOrCumsumOut, expandedScaleOut), OP_ATTR(activeNum, expertCapacity, expertNum, dropPadMode, expertTokensNumType, expertTokensNumFlag, quantMode, activeExpertRange, rowIdxType)); + return std::tuple(expandedXOut, expandedRowIdxOut, expertTokensCountOrCumsumOut, expandedScaleOut); //OP_OUTPUT +} + +} // namespace l0op \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom.h b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom.h new file mode 100644 index 000000000..65da3ff33 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom.h @@ -0,0 +1,25 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef OP_API_INC_LEVEL0_MOE_INIT_ROUTING_CUSTOM_H +#define OP_API_INC_LEVEL0_MOE_INIT_ROUTING_CUSTOM_H + +#include +#include "opdev/op_executor.h" + +namespace l0op { +std::tuple MoeInitRoutingCustom(const aclTensor *x, const aclTensor *expertIdx, const aclTensor *scale, + const aclTensor *offset, int64_t activeNum, int64_t expertCapacity, + int64_t expertNum, int64_t dropPadMode, int64_t expertTokensNumType, + bool expertTokensNumFlag, int64_t quantMode, const aclIntArray *activeExpertRange, + int64_t rowIdxType, const aclTensor *expandedX, const aclTensor *expandedRowIdx, + const aclTensor *expertTokensCountOrCumsum, const aclTensor *expandedScale, aclOpExecutor *executor); +} // namespace l0op +#endif // OP_API_INC_LEVEL0_MOE_INIT_ROUTING_CUSTOM_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_def.cpp b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_def.cpp new file mode 100644 index 000000000..c1d980b9d --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_def.cpp @@ -0,0 +1,105 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_init_routing_v3_def.cpp + * \brief + */ +#include "register/op_def_registry.h" + +namespace ops { +class MoeInitRoutingCustom : public OpDef { +public: + explicit MoeInitRoutingCustom(const char *name) : OpDef(name) + { + this->Input("x") + .ParamType(REQUIRED) + .DataType( + {ge::DT_INT8, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("expert_idx") + .ParamType(REQUIRED) + .DataType( + {ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("scale") + .ParamType(OPTIONAL) + .DataType( + {ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Input("offset") + .ParamType(OPTIONAL) + .DataType( + {ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("expanded_x") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("expanded_row_idx") + .ParamType(REQUIRED) + .DataType( + {ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("expert_tokens_count_or_cumsum") + .ParamType(REQUIRED) + .DataType( + {ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("expanded_scale") + .ParamType(REQUIRED) + .DataType( + {ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND}); + this->Attr("active_num").AttrType(OPTIONAL).Int(-1); + this->Attr("expert_capacity").AttrType(OPTIONAL).Int(-1); + this->Attr("expert_num").AttrType(OPTIONAL).Int(-1); + this->Attr("drop_pad_mode").AttrType(OPTIONAL).Int(0); + this->Attr("expert_tokens_num_type").AttrType(OPTIONAL).Int(0); + this->Attr("expert_tokens_num_flag").AttrType(OPTIONAL).Bool(false); + this->Attr("quant_mode").AttrType(OPTIONAL).Int(-1); + this->Attr("active_expert_range").AttrType(OPTIONAL).ListInt({}); + this->Attr("row_idx_type").AttrType(OPTIONAL).Int(0); + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + + } +}; + +OP_ADD(MoeInitRoutingCustom); +} // namespace ops diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_infershape.cpp b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_infershape.cpp new file mode 100644 index 000000000..77e3d2833 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_infershape.cpp @@ -0,0 +1,797 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/* ! + * \file moe_init_routing_custom_infershape.cpp + * \brief + */ + +#include +#include +#include +#include "register/op_def_registry.h" +#include "log/ops_log.h" +#include "platform/platform_info.h" + +#define unlikely(x) __builtin_expect((x), 0) +#define OP_CHECK_NULL_WITH_CONTEXT(context, ptr) \ + do { \ + if (unlikely((ptr) == nullptr)) { \ + const char* name = (unlikely(((context) == nullptr) || (context)->GetNodeName() == nullptr)) ? \ + "nil" : \ + (context)->GetNodeName(); \ + OPS_LOG_E(name, "%s is nullptr!", #ptr); \ + return ge::GRAPH_FAILED; \ + } \ + } while (0) + +using namespace ge; +namespace ops { +static constexpr size_t DIM_ONE = 1U; +static constexpr size_t DIM_TWO = 2U; +static constexpr size_t DIM_THREE = 3U; +static constexpr int64_t NEG_ONE = static_cast(-1); +static constexpr int64_t NEG_TWO = static_cast(-2); +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_INPUT_X = 0; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_INPUT_EXPERT_IDX = 1; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_INPUT_SCALE = 2; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_INPUT_OFFSET = 3; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_ATTR_ACTIVE_NUM = 0; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_ATTR_EXPERT_CAPACITY = 1; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_ATTR_EXPERT_NUM = 2; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_ATTR_DROP_PAD_MODE = 3; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_ATTR_EXPERT_TOKEN_NUM_TYPE = 4; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_ATTR_EXPERT_TOKEN_NUM_FLAG = 5; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_ATTR_QUANT_MODE = 6; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_ATTR_ACTIVE_EXPERT_RANGE = 7; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_ATTR_ROW_IDX_TYPE = 8; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPANDED_X = 0; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPANDED_ROW_IDX = 1; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPERT_TOKEN_CUMSUM_OR_COUNT = 2; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPANDED_SCALE = 3; +static constexpr int64_t MOE_INIT_ROUTING_CUSTOM_EXPERT_END_BOUND = 10240; +static constexpr int64_t KEY_VALUE_MODE_DIM0_NUM = 2; +enum DropPadMode : int8_t { + NO_DROP_PAD = 0, + DROP_PAD = 1, +}; +enum QuantMode : int8_t { + NON_QUANT = -1, + STATIC_QUANT = 0, + DYNAMIC_QUANT = 1 +}; +enum ExpertTokenNumType : int8_t { + CUMSUM = 0, + COUNT = 1, + KEY_VALUE = 2 +}; + +static bool isSameDim(int64_t dim1, int64_t dim2) +{ + if (dim1 <= NEG_ONE || dim2 <= NEG_ONE) { + return true; + } + return dim1 == dim2; +} + +static ge::graphStatus GetAndCheckAttrActiveExpertRange(const gert::RuntimeAttrs *attrs, + gert::InferShapeContext *context, int64_t &expertStart, + int64_t &expertEnd, int64_t &experNum) +{ + OPS_LOG_D(context->GetNodeName(), "Begin to do GetAndCheckAttrActiveExpertRange."); + // Check if active_expert_range size is 2 and if expert_start < expert_end + auto activeExpertRangePtr = attrs->GetListInt(MOE_INIT_ROUTING_CUSTOM_ATTR_ACTIVE_EXPERT_RANGE); + if (nullptr == activeExpertRangePtr) { + OPS_LOG_E(context->GetNodeName(), "The active_expert_range should be list int. But it is none."); + return ge::GRAPH_FAILED; + } + int64_t activeExpertRangeSize = activeExpertRangePtr->GetSize(); + if (activeExpertRangePtr->GetSize() == DIM_TWO) { + expertStart = activeExpertRangePtr->GetData()[0]; + expertEnd = activeExpertRangePtr->GetData()[1]; + if (expertStart >= expertEnd || expertStart < 0 || expertEnd > MOE_INIT_ROUTING_CUSTOM_EXPERT_END_BOUND) { + OPS_LOG_E(context->GetNodeName(), + "The active_expert_range should be in [0, %ld), but the active_expert_range is [%ld, %ld).", + MOE_INIT_ROUTING_CUSTOM_EXPERT_END_BOUND, expertStart, expertEnd); + return ge::GRAPH_FAILED; + } + } else if (activeExpertRangePtr->GetSize() == 0) { + expertStart = 0; + expertEnd = experNum; + } else { + OPS_LOG_E(context->GetNodeName(), "The active_expert_range size should be 2, but its size is %ld.", activeExpertRangeSize); + return ge::GRAPH_FAILED; + } + + OPS_LOG_D(context->GetNodeName(), "End to do GetAndCheckAttrActiveExpertRange."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus GetAndCheckAttrActiveNum(const gert::RuntimeAttrs *attrs, gert::InferShapeContext *context, + int64_t &activeNum, int64_t &dropPadMode) +{ + OPS_LOG_D(context->GetNodeName(), "Begin to do GetAndCheckAttrActiveNum."); + const int64_t *activeNumPtr = attrs->GetAttrPointer(MOE_INIT_ROUTING_CUSTOM_ATTR_ACTIVE_NUM); + if (nullptr == activeNumPtr) { + OPS_LOG_E(context->GetNodeName(), "The active_num should not be none."); + return ge::GRAPH_FAILED; + } + activeNum = *activeNumPtr; + if (dropPadMode == DropPadMode::NO_DROP_PAD && activeNum < -1) { + OPS_LOG_E(context->GetNodeName(), "The active_num should be greater than or equal to 0. But it is %ld.", activeNum); + return ge::GRAPH_FAILED; + } + + OPS_LOG_D(context->GetNodeName(), "End to do GetAndCheckAttrActiveNum."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus GetAndCheckAttrExpertCapacity(const gert::RuntimeAttrs *attrs, gert::InferShapeContext *context, + const gert::Shape *xShape, int64_t &expertCapacity, + int64_t &dropPadMode) +{ + OPS_LOG_D(context->GetNodeName(), "Begin to do GetAndCheckAttrExpertCapacity."); + const int64_t *expertCapacityPtr = attrs->GetAttrPointer(MOE_INIT_ROUTING_CUSTOM_ATTR_EXPERT_CAPACITY); + if (nullptr == expertCapacityPtr) { + OPS_LOG_E(context->GetNodeName(), "The expert_capacity should not be none."); + return ge::GRAPH_FAILED; + } + expertCapacity = *expertCapacityPtr; + if (dropPadMode == DropPadMode::DROP_PAD && xShape->GetDim(0) > 0 && expertCapacity > xShape->GetDim(0)) { + OPS_LOG_E(context->GetNodeName(), "The expert_capacity should be between 0 and n. But it is %ld.", expertCapacity); + return ge::GRAPH_FAILED; + } + + OPS_LOG_D(context->GetNodeName(), "End to do GetAndCheckAttrExpertCapacity."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus GetAndCheckAttrExpertNum(const gert::RuntimeAttrs *attrs, gert::InferShapeContext *context, + int64_t &experNum) +{ + OPS_LOG_D(context->GetNodeName(), "Begin to do GetAndCheckexperNum."); + const int64_t *experNumPtr = attrs->GetAttrPointer(MOE_INIT_ROUTING_CUSTOM_ATTR_EXPERT_NUM); + if (nullptr == experNumPtr) { + OPS_LOG_E(context->GetNodeName(), "The expert_num should not be none."); + return ge::GRAPH_FAILED; + } + experNum = *experNumPtr; + if (experNum <= 0 || experNum > MOE_INIT_ROUTING_CUSTOM_EXPERT_END_BOUND) { + OPS_LOG_E(context->GetNodeName(), "The expert_num should be greater than 0. But it is %ld.", experNum); + return ge::GRAPH_FAILED; + } + + OPS_LOG_D(context->GetNodeName(), "End to do GetAndCheckAttrExpertNum."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus GetAndCheckAttrDropPadMode(const gert::RuntimeAttrs *attrs, gert::InferShapeContext *context, + int64_t &dropPadMode) +{ + OPS_LOG_D(context->GetNodeName(), "Begin to do GetAndCheckAttrDropPadMode."); + const int64_t *dropPadModePtr = attrs->GetAttrPointer(MOE_INIT_ROUTING_CUSTOM_ATTR_DROP_PAD_MODE); + if (nullptr == dropPadModePtr) { + OPS_LOG_E(context->GetNodeName(), "The RuntimeAttrs for drop_pad_mode is none."); + return ge::GRAPH_FAILED; + } + + dropPadMode = *dropPadModePtr; + if (dropPadMode < DropPadMode::NO_DROP_PAD || dropPadMode > DropPadMode::DROP_PAD) { + OPS_LOG_E(context->GetNodeName(), "The drop_pad_mode should be %d or %d. But it is %ld.", DropPadMode::NO_DROP_PAD, + DropPadMode::DROP_PAD, dropPadMode); + return ge::GRAPH_FAILED; + } + + OPS_LOG_D(context->GetNodeName(), "End to do GetAndCheckAttrDropPadMode."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus GetAndCheckAttrExpertTokenNumType(const gert::RuntimeAttrs *attrs, gert::InferShapeContext* context, + int64_t &experTokenNumType) +{ + OPS_LOG_D(context->GetNodeName(), "Begin to do GetAndCheckexperTokenNumType."); + const int64_t *experTokenNumTypePtr = + attrs->GetAttrPointer(MOE_INIT_ROUTING_CUSTOM_ATTR_EXPERT_TOKEN_NUM_TYPE); + if (nullptr == experTokenNumTypePtr) { + OPS_LOG_E(context->GetNodeName(), "The expert_token_num_type should not be none."); + return ge::GRAPH_FAILED; + } + experTokenNumType = *experTokenNumTypePtr; + if (experTokenNumType < ExpertTokenNumType::CUMSUM || experTokenNumType > ExpertTokenNumType::KEY_VALUE) { + OPS_LOG_E(context->GetNodeName(), "The expert_token_num_type should be %d, %d or %d. But it is %ld.", + ExpertTokenNumType::CUMSUM, ExpertTokenNumType::COUNT, ExpertTokenNumType::KEY_VALUE, + experTokenNumType); + return ge::GRAPH_FAILED; + } + + OPS_LOG_D(context->GetNodeName(), "End to do GetAndCheckAttrExpertTokenNumType."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus GetAndCheckAttrExpertTokenNumFlag(const gert::RuntimeAttrs *attrs, + gert::InferShapeContext *context, bool &experTokenNumFlag) +{ + OPS_LOG_D(context->GetNodeName(), "Begin to do GetAndCheckexperTokenNumType."); + const bool *experTokenNumFlagPtr = attrs->GetAttrPointer(MOE_INIT_ROUTING_CUSTOM_ATTR_EXPERT_TOKEN_NUM_FLAG); + if (nullptr == experTokenNumFlagPtr) { + OPS_LOG_E(context->GetNodeName(), "The expert_token_num_flag should not be none."); + return ge::GRAPH_FAILED; + } + experTokenNumFlag = *experTokenNumFlagPtr; + OPS_LOG_D(context->GetNodeName(), "End to do GetAndCheckAttrExpertTokenNumType."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus GetAndCheckAttrQuantMode(const gert::RuntimeAttrs *attrs, gert::InferShapeContext *context, + int64_t &quantMode) +{ + OPS_LOG_D(context->GetNodeName(), "Begin to do GetAndCheckQuantMode."); + if (nullptr == attrs) { + OPS_LOG_E(context->GetNodeName(), "The RuntimeAttrs for quant_mode is none."); + return ge::GRAPH_FAILED; + } + const int64_t *quantModePtr = attrs->GetAttrPointer(MOE_INIT_ROUTING_CUSTOM_ATTR_QUANT_MODE); + if (nullptr == quantModePtr) { + OPS_LOG_E(context->GetNodeName(), "The quant_mode should be %d, %d or %d. But it is none.", QuantMode::NON_QUANT, + QuantMode::STATIC_QUANT, QuantMode::DYNAMIC_QUANT); + return ge::GRAPH_FAILED; + } + quantMode = *quantModePtr; + if (quantMode < QuantMode::NON_QUANT || quantMode > QuantMode::DYNAMIC_QUANT) { + OPS_LOG_E(context->GetNodeName(), "The quant_mode should be %d, %d or %d. But it is %ld.", QuantMode::NON_QUANT, + QuantMode::STATIC_QUANT, QuantMode::DYNAMIC_QUANT, quantMode); + return ge::GRAPH_FAILED; + } + OPS_LOG_D(context->GetNodeName(), "End to do GetAndCheckQuantMode."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus GetAndCheckAttrRowIdxType(const gert::RuntimeAttrs *attrs, gert::InferShapeContext *context, + int64_t &rowIdxType, int64_t &dropPadMode) +{ + OPS_LOG_D(context->GetNodeName(), "Begin to do GetAndCheckAttrRowIdxType."); + if (nullptr == attrs) { + OPS_LOG_E(context->GetNodeName(), "The RuntimeAttrs for row_Idx_type is none."); + return ge::GRAPH_FAILED; + } + const int64_t *dropPadModePtr = attrs->GetAttrPointer(MOE_INIT_ROUTING_CUSTOM_ATTR_DROP_PAD_MODE); + dropPadMode = *dropPadModePtr; + + const int64_t *rowIdxTypePtr = attrs->GetAttrPointer(MOE_INIT_ROUTING_CUSTOM_ATTR_ROW_IDX_TYPE); + if (nullptr == rowIdxTypePtr) { + OPS_LOG_E(context->GetNodeName(), "The row_Idx_type should be 0 or 1. But it is none."); + return ge::GRAPH_FAILED; + } + rowIdxType = *rowIdxTypePtr; + if (dropPadMode == DropPadMode::DROP_PAD && rowIdxType != 0) { + OPS_LOG_E(context->GetNodeName(), "The row_Idx_type should be 0 when dropPadMode is equal to 1 But it is %ld.", rowIdxType); + return ge::GRAPH_FAILED; + } + + if (rowIdxType < 0 || rowIdxType > 1) { + OPS_LOG_E(context->GetNodeName(), "The row_Idx_type should be 0 or 1 But it is %ld.", rowIdxType); + return ge::GRAPH_FAILED; + } + + OPS_LOG_D(context->GetNodeName(), "End to do GetAndCheckAttrRowIdxType."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus CheckInputScaleShape(gert::InferShapeContext *context, const gert::Shape *xShape, + const gert::Shape *scaleShape, const int64_t expertStart, + const int64_t expertEnd, const int64_t quantMode) +{ + // When quant_mode is STATIC_QUANT, scale cannot be none. + OP_CHECK((nullptr == scaleShape && QuantMode::STATIC_QUANT == quantMode), + OPS_LOG_E(context->GetNodeName(), "The scale cannot be none when quant_mode is %ld.", quantMode), + return ge::GRAPH_FAILED); + + // When quant_mode is NON_QUANT or DYNAMIC_QUANT, scale can be none. + OP_CHECK((nullptr == scaleShape && (QuantMode::NON_QUANT == quantMode || QuantMode::DYNAMIC_QUANT == quantMode)), + OPS_LOG_I(context->GetNodeName(), "When quant_mode is NON_QUANT or DYNAMIC_QUANT, scale can be none."), + return ge::GRAPH_SUCCESS); + + if (QuantMode::NON_QUANT == quantMode) { + if (scaleShape->GetDimNum() == DIM_ONE) { + OP_CHECK(scaleShape->GetDim(0) < 0 && scaleShape->GetDim(0) != NEG_ONE && scaleShape->GetDim(0) != NEG_TWO, + OPS_LOG_E(context->GetNodeName(), + "When quant_mode is %ld and use scale in dynamic graph, The shape of scale should be (-1) or (-2), current shape is (%s).", + quantMode, ops::Shape2String(*scaleShape).c_str()), + return ge::GRAPH_FAILED); + OP_CHECK(scaleShape->GetDim(0) > 0 && !isSameDim(scaleShape->GetDim(0), xShape->GetDim(0)), + OPS_LOG_E(context->GetNodeName(), + "When quant_mode is %ld and use scale in static graph, The shape of scale should be (%ld,), current shape is (%s).", + quantMode, xShape->GetDim(0), ops::Shape2String(*scaleShape).c_str()), + return ge::GRAPH_FAILED); + } else { + OPS_LOG_E(context->GetNodeName(), "When quant_mode is %ld, The dimNum of scale should be 1, current shape is (%ld).", quantMode, + scaleShape->GetDimNum()); + return ge::GRAPH_FAILED; + } + } else if (QuantMode::STATIC_QUANT == quantMode) { + if (scaleShape->GetDimNum() == DIM_ONE) { + OP_CHECK( + scaleShape->GetDim(0) != NEG_ONE && scaleShape->GetDim(0) != NEG_TWO && + !isSameDim(scaleShape->GetDim(0), DIM_ONE), + OPS_LOG_E( + context->GetNodeName(), + "When quant_mode is %ld, the shape of scale should be (-1) or (-2) or (1,), current shape is (%s).", + quantMode, ops::Shape2String(*scaleShape).c_str()), + return ge::GRAPH_FAILED); + } else { + OPS_LOG_E(context->GetNodeName(), "When quant_mode is %ld, the dimNum of scale should be (1,), current shape is (%ld).", + quantMode, scaleShape->GetDimNum()); + return ge::GRAPH_FAILED; + } + } else if (QuantMode::DYNAMIC_QUANT == quantMode) { + int64_t activeExpertRange = expertEnd - expertStart; + if (scaleShape->GetDimNum() == DIM_ONE) { + OP_CHECK(scaleShape->GetDim(0) != NEG_TWO, + OPS_LOG_E(context->GetNodeName(), + "When quant_mode is %ld and scale dim is 1 in dynamic graph, the first dim of scale should be -2, but " + "its shape is (%ld).", + quantMode, scaleShape->GetDim(0)), + return ge::GRAPH_FAILED); + } else if (scaleShape->GetDimNum() == DIM_TWO) { + if (scaleShape->GetDim(0) > 0) { + OP_CHECK( + !isSameDim(scaleShape->GetDim(0), activeExpertRange) && !isSameDim(scaleShape->GetDim(0), DIM_ONE), + OPS_LOG_E( + context->GetNodeName(), + "When quant_mode is %ld in static graph, the first dim of scale should be 1 or %ld, but its shape is (%ld).", + quantMode, activeExpertRange, scaleShape->GetDim(0)), + return ge::GRAPH_FAILED); + OP_CHECK( + !isSameDim(scaleShape->GetDim(1), xShape->GetDim(1)), + OPS_LOG_E( + context->GetNodeName(), + "When quant_mode is %ld in static graph, the second dim of scale should or %ld, but its shape is (%ld).", + quantMode, xShape->GetDim(1), scaleShape->GetDim(0)), + return ge::GRAPH_FAILED); + } else { + OP_CHECK( + scaleShape->GetDim(0) != NEG_ONE || (scaleShape->GetDim(1) != NEG_ONE && scaleShape->GetDim(1) != xShape->GetDim(1)), + OPS_LOG_E(context->GetNodeName(), + "When quant_mode is %ld and scale dim is 2 in dynamic graph, the shape of scale should be (-1, -1) or (-1, %d), but its shape is (%s).", + quantMode, xShape->GetDim(1), ops::Shape2String(*scaleShape).c_str()), + return ge::GRAPH_FAILED); + } + } else { + OPS_LOG_E( + context->GetNodeName(), + "When quant_mode is %ld, the dimNum of scale should be 1(dynamic graph) or 2, but its shape is (%ld).", + scaleShape->GetDimNum()); + return ge::GRAPH_FAILED; + } + } + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus CheckInputOffsetShape(gert::InferShapeContext *context, + const gert::Shape *offsetShape, const int64_t expertStart, + const int64_t expertEnd, const int64_t quantMode) +{ + // The shape of offset can be none. + if (quantMode != QuantMode::STATIC_QUANT) { + return ge::GRAPH_SUCCESS; + } else if (nullptr == offsetShape) { + return ge::GRAPH_FAILED; + } + + if (offsetShape->GetDimNum() != DIM_ONE) { + OPS_LOG_E(context->GetNodeName(), "The dimNum of offset should be 1, current shape is (%ld).", offsetShape->GetDimNum()); + return ge::GRAPH_FAILED; + } + if (offsetShape->GetDim(0) != NEG_ONE && offsetShape->GetDim(0) != NEG_TWO && !isSameDim(offsetShape->GetDim(0), DIM_ONE)) { + OPS_LOG_E(context->GetNodeName(), + "The shape of offset should be (1,) in static graph or (-2), (-1,) in dynamic graph, current shape is (%s).", + ops::Shape2String(*offsetShape).c_str()); + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus CheckInputShape(gert::InferShapeContext *context, const gert::Shape *xShape, + const gert::Shape *expertIdxShape, const gert::Shape *scaleShape, + const gert::Shape *offsetShape, const int64_t expertStart, + const int64_t expertEnd, const int64_t quantMode) +{ + // Check the shape of input_x + if (xShape->GetDimNum() == DIM_ONE) { + if (xShape->GetDim(0) != ge::UNKNOWN_DIM_NUM) { + OPS_LOG_E(context->GetNodeName(), "The dynamic dim of x should be -2, current shape is %s.", + ops::Shape2String(*xShape).c_str()); + return ge::GRAPH_FAILED; + } + } else if (xShape->GetDimNum() != DIM_TWO) { + OPS_LOG_E(context->GetNodeName(), "The dim of x should be 2 or dynamic, current shape is %s.", + ops::Shape2String(*xShape).c_str()); + return ge::GRAPH_FAILED; + } + + int64_t x_n = xShape->GetDimNum() == DIM_ONE ? NEG_ONE : xShape->GetDim(0); + int64_t cols = xShape->GetDimNum() == DIM_ONE ? NEG_ONE : xShape->GetDim(1); + if (x_n < NEG_ONE || cols < NEG_ONE) { + OPS_LOG_E(context->GetNodeName(), "Invalid x shape, shape is %s.", ops::Shape2String(*xShape).c_str()); + return ge::GRAPH_FAILED; + } + + // Check the shape of expert_idx + if (expertIdxShape->GetDimNum() == DIM_ONE) { + if (expertIdxShape->GetDim(0) != ge::UNKNOWN_DIM_NUM) { + OPS_LOG_E(context->GetNodeName(), "The dynamic dim of expert_idx should be -2, current shape is %s.", + ops::Shape2String(*expertIdxShape).c_str()); + return ge::GRAPH_FAILED; + } + } else if (expertIdxShape->GetDimNum() != DIM_TWO) { + OPS_LOG_E(context->GetNodeName(), "The dim of expert_idx should be 2 or dynamic, current shape is %s.", + ops::Shape2String(*expertIdxShape).c_str()); + return ge::GRAPH_FAILED; + } + + int64_t expert_idx_n = expertIdxShape->GetDimNum() == DIM_ONE ? NEG_ONE : expertIdxShape->GetDim(0); + int64_t expert_idx_k = expertIdxShape->GetDimNum() == DIM_ONE ? NEG_ONE : expertIdxShape->GetDim(1); + if (expert_idx_n < NEG_ONE || expert_idx_k < NEG_ONE) { + OPS_LOG_E(context->GetNodeName(), "Invalid expert_idx shape, shape is %s.", + ops::Shape2String(*expertIdxShape).c_str()); + return ge::GRAPH_FAILED; + } + + if (!isSameDim(x_n, expert_idx_n)) { + OPS_LOG_E(context->GetNodeName(), "The first dim of x and expert_idx should be same."); + return ge::GRAPH_FAILED; + } + // Check the shape of scale + if (CheckInputScaleShape(context, xShape, scaleShape, expertStart, expertEnd, quantMode) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + // Check the shape of offset + if (CheckInputOffsetShape(context, offsetShape, expertStart, expertEnd, quantMode) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +static void ShowInputShapeAndAttrInfo(gert::InferShapeContext *context, const gert::Shape *xShape, + const gert::Shape *expertIdxShape, const gert::Shape *scaleShape, + const gert::Shape *offsetShape, const int64_t expertStart, + const int64_t expertEnd, const int64_t quantMode, const int64_t rowIdxType) +{ + // input_x and expert_idx are all required. + OPS_LOG_D(context->GetNodeName(), "x shape is: %s.", ops::Shape2String(*xShape).c_str()); + OPS_LOG_D(context->GetNodeName(), "expert_idx shape is: %s.", ops::Shape2String(*expertIdxShape).c_str()); + + // scale is optional and can be none. + if (nullptr == scaleShape) { + OPS_LOG_D(context->GetNodeName(), "scale_shape is: none."); + } else { + OPS_LOG_D(context->GetNodeName(), "scale_shape is: %s.", ops::Shape2String(*scaleShape).c_str()); + } + + // offset is optional and can be none. + OPS_LOG_D(context->GetNodeName(), "Begin print offset_shape."); + if (nullptr == offsetShape) { + OPS_LOG_D(context->GetNodeName(), "offset_shape is: none."); + } else { + OPS_LOG_D(context->GetNodeName(), "offset_shape is: %s.", ops::Shape2String(*offsetShape).c_str()); + } + OPS_LOG_D(context->GetNodeName(), "End print offset_shape."); + + // Attrs are all required. + OPS_LOG_D(context->GetNodeName(), "active_expert_range is: [%ld, %ld).", expertStart, expertEnd); + OPS_LOG_D(context->GetNodeName(), "quant_mode is: %ld.", quantMode); + OPS_LOG_D(context->GetNodeName(), "row_Idx_type is: %ld.", rowIdxType); +} + +static void ShowOutputShapeInfo(gert::InferShapeContext *context, const gert::Shape *expandedXShape, + const gert::Shape *expandedRowIdxShape, + const gert::Shape *expertTokenCumsumOrCountShape, const gert::Shape *expandedScaleShape) +{ + OPS_LOG_D(context->GetNodeName(), "expanded_x shape is: %s after infershape.", + ops::Shape2String(*expandedXShape).c_str()); + OPS_LOG_D(context->GetNodeName(), "expanded_row_idx shape is: %s after infershape.", + ops::Shape2String(*expandedRowIdxShape).c_str()); + OPS_LOG_D(context->GetNodeName(), "expert_token_cumsum_or_count shape is: %s after infershape.", + ops::Shape2String(*expertTokenCumsumOrCountShape).c_str()); + OPS_LOG_D(context->GetNodeName(), "expanded_scale shape is: %s after infershape.", + ops::Shape2String(*expandedScaleShape).c_str()); +} + +static ge::graphStatus InferShape4MoeInitRoutingCustom(gert::InferShapeContext *context) +{ + OPS_LOG_D(context->GetNodeName(), "Begin to do MoeInitRoutingCustomInfershape."); + // 1. Get and check input shape + // 1.1 Get and check input_x + const gert::Shape *xShape = context->GetInputShape(MOE_INIT_ROUTING_CUSTOM_INPUT_X); + OP_CHECK_NULL_WITH_CONTEXT(context, xShape); + + // 1.2 Get and check expert_idx + const gert::Shape *expertIdxShape = context->GetInputShape(MOE_INIT_ROUTING_CUSTOM_INPUT_EXPERT_IDX); + OP_CHECK_NULL_WITH_CONTEXT(context, expertIdxShape); + + // 1.3 Get scale shape without checking null, because scale is optional and can be none. + const gert::Shape *scaleShape = context->GetOptionalInputShape(MOE_INIT_ROUTING_CUSTOM_INPUT_SCALE); + + // 1.4 Get offset shape without checking null, because offset is optional and can be none. + const gert::Shape *offsetShape = context->GetOptionalInputShape(MOE_INIT_ROUTING_CUSTOM_INPUT_OFFSET); + // 2. Get and check attrs + const gert::RuntimeAttrs *attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + + // 2.1 Get and check expert_num attr + int64_t experNum = static_cast(-1); + if (GetAndCheckAttrExpertNum(attrs, context, experNum) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + // 2.2 Get and check active_expert_range attr + int64_t expertStart = static_cast(-1); + int64_t expertEnd = static_cast(-1); + if (GetAndCheckAttrActiveExpertRange(attrs, context, expertStart, expertEnd, experNum) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + if (nullptr == attrs) { + OPS_LOG_E(context->GetNodeName(), "The attrs is none."); + return ge::GRAPH_FAILED; + } + + // 2.3 Get and check drop_pad_mode attr + int64_t dropPadMode = static_cast(-1); + if (GetAndCheckAttrDropPadMode(attrs, context, dropPadMode) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + // 2.4 Get and check active_num attr + int64_t activeNum = static_cast(-1); + if (GetAndCheckAttrActiveNum(attrs, context, activeNum, dropPadMode) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + // 2.5 Get and check expert_capacity attr + int64_t expertCapacity = static_cast(-1); + if (GetAndCheckAttrExpertCapacity(attrs, context, xShape, expertCapacity, dropPadMode) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + // 2.6 Get and check expert_token_num_type attr + int64_t expertTokenNumType = static_cast(-1); + if (GetAndCheckAttrExpertTokenNumType(attrs, context, expertTokenNumType) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + // 2.7 Get and check expert_token_num_type attr + bool expertTokenNumFlag = false; + if (GetAndCheckAttrExpertTokenNumFlag(attrs, context, expertTokenNumFlag) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + // 2.8 Get and check quant_mode attr + int64_t quantMode = static_cast(-1); + if (GetAndCheckAttrQuantMode(attrs, context, quantMode) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + // 2.9 Get and check row_Idx_type attr + int64_t rowIdxType = static_cast(-1); + if (GetAndCheckAttrRowIdxType(attrs, context, rowIdxType, dropPadMode) != ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + // Check input shape + if (CheckInputShape(context, xShape, expertIdxShape, scaleShape, offsetShape, expertStart, expertEnd, quantMode) != + ge::GRAPH_SUCCESS) { + return ge::GRAPH_FAILED; + } + + // 3. Infer output shape + // 3.1 Prepare output shape + gert::Shape *expandedXShape = context->GetOutputShape(MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPANDED_X); + OP_CHECK_NULL_WITH_CONTEXT(context, expandedXShape); + gert::Shape *expandedRowIdxShape = context->GetOutputShape(MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPANDED_ROW_IDX); + OP_CHECK_NULL_WITH_CONTEXT(context, expandedRowIdxShape); + gert::Shape *expertTokenCumsumOrCountShape = + context->GetOutputShape(MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPERT_TOKEN_CUMSUM_OR_COUNT); + OP_CHECK_NULL_WITH_CONTEXT(context, expertTokenCumsumOrCountShape); + gert::Shape *expandedScaleShape = context->GetOutputShape(MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPANDED_SCALE); + OP_CHECK_NULL_WITH_CONTEXT(context, expandedScaleShape); + + int64_t x_n = xShape->GetDimNum() == DIM_ONE ? NEG_ONE : xShape->GetDim(0); + int64_t cols = xShape->GetDimNum() == DIM_ONE ? NEG_ONE : xShape->GetDim(1); + + int64_t expert_idx_n = expertIdxShape->GetDimNum() == DIM_ONE ? NEG_ONE : expertIdxShape->GetDim(0); + int64_t k = expertIdxShape->GetDimNum() == DIM_ONE ? NEG_ONE : expertIdxShape->GetDim(1); + int64_t n = x_n > expert_idx_n ? x_n : expert_idx_n; + if (activeNum == 0 || activeNum == -1) { + activeNum = n * k; + } else { + activeNum = std::min(activeNum, n * k); + } + + int64_t xOutDimNum = activeNum < n * k ? activeNum : n * k; + int64_t outNum = (n == NEG_ONE || k == NEG_ONE) ? NEG_ONE : n * k; + int64_t xOutNum = (n == NEG_ONE || k == NEG_ONE) ? NEG_ONE : xOutDimNum; + // 3.2 Set output expanded_x shape + if (dropPadMode == DropPadMode::NO_DROP_PAD) { + expandedXShape->SetDimNum(DIM_TWO); + expandedXShape->SetDim(0U, xOutNum); + expandedXShape->SetDim(DIM_ONE, cols); + } else { + expandedXShape->SetDimNum(DIM_THREE); + expandedXShape->SetDim(0U, experNum); + expandedXShape->SetDim(DIM_ONE, expertCapacity); + expandedXShape->SetDim(DIM_TWO, cols); + } + + // 3.3 Set output expanded_row_idx shape + expandedRowIdxShape->SetDimNum(DIM_ONE); + expandedRowIdxShape->SetDim(0U, outNum); + + // 3.4 Set output expert_token_cumsum_or_count shape + if (expertTokenNumFlag) { + if (expertTokenNumType == ExpertTokenNumType::KEY_VALUE) { + expertTokenCumsumOrCountShape->SetDimNum(DIM_TWO); + expertTokenCumsumOrCountShape->SetDim(0U, experNum); + expertTokenCumsumOrCountShape->SetDim(DIM_ONE, KEY_VALUE_MODE_DIM0_NUM); + } else { + expertTokenCumsumOrCountShape->SetDimNum(DIM_ONE); + expertTokenCumsumOrCountShape->SetDim(0U, expertEnd - expertStart); + } + } + + // 3.5 Set output expanded_scale shape + // When scale_shape=(b*s) and non-quant, or it is dynamic quant mode, the shape of expanded_scale should be (b*s*k) + if (QuantMode::NON_QUANT == quantMode || QuantMode::DYNAMIC_QUANT == quantMode) { + expandedScaleShape->SetDimNum(DIM_ONE); + if (dropPadMode == DropPadMode::NO_DROP_PAD) { + expandedScaleShape->SetDim(0U, xOutNum); + } else { + expandedScaleShape->SetDim(0U, experNum * expertCapacity); + } + } + + ShowOutputShapeInfo(context, expandedXShape, expandedRowIdxShape, expertTokenCumsumOrCountShape, + expandedScaleShape); + OPS_LOG_D(context->GetNodeName(), "End to do MoeInitRoutingCustomInfershape."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataType4MoeInitRoutingCustom(gert::InferDataTypeContext *context) +{ + OPS_LOG_D(context->GetNodeName(), "Begin to do MoeInitRoutingCustomInferDataType."); + + // Get and check quant_mode attr + const gert::RuntimeAttrs *attrs = context->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); + int64_t quantMode = static_cast(-1); + const int64_t *quantModePtr = attrs->GetAttrPointer(MOE_INIT_ROUTING_CUSTOM_ATTR_QUANT_MODE); + if (nullptr == quantModePtr) { + OPS_LOG_E(context->GetNodeName(), "The quant_mode should be %d, %d or %d. But it is none.", QuantMode::NON_QUANT, + QuantMode::STATIC_QUANT, QuantMode::DYNAMIC_QUANT); + return ge::GRAPH_FAILED; + } + quantMode = *quantModePtr; + // Infer output dtype according quant_mode + auto xDtype = context->GetInputDataType(MOE_INIT_ROUTING_CUSTOM_INPUT_X); + if (QuantMode::NON_QUANT == quantMode) { + context->SetOutputDataType(MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPANDED_X, xDtype); + } else if (QuantMode::STATIC_QUANT == quantMode || QuantMode::DYNAMIC_QUANT == quantMode) { + if (ge::DT_INT8 == xDtype) { + OPS_LOG_E(context->GetNodeName(), "When quant_mode=%ld, xDtype cannot be int_8.", quantMode); + return ge::GRAPH_FAILED; + } + context->SetOutputDataType(MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPANDED_X, ge::DT_INT8); + } + context->SetOutputDataType(MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPANDED_ROW_IDX, ge::DT_INT32); + context->SetOutputDataType(MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPERT_TOKEN_CUMSUM_OR_COUNT, ge::DT_INT64); + context->SetOutputDataType(MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPANDED_SCALE, ge::DT_FLOAT); + OPS_LOG_D(context->GetNodeName(), "End to do MoeInitRoutingCustomInferDataType."); + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferShapeRange4MoeInitRoutingCustom(gert::InferShapeRangeContext *context) +{ + OPS_LOG_D(context->GetNodeName(), "Begin to do MoeInitRoutingCustomInferRange."); + + // Get and check the pointers of all the outputs' shape range object + auto expanded_x = context->GetOutputShapeRange(MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPANDED_X); + OP_CHECK_NULL_WITH_CONTEXT(context, expanded_x); + auto expanded_row_idx = context->GetOutputShapeRange(MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPANDED_ROW_IDX); + OP_CHECK_NULL_WITH_CONTEXT(context, expanded_row_idx); + auto count = context->GetOutputShapeRange(MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPERT_TOKEN_CUMSUM_OR_COUNT); + OP_CHECK_NULL_WITH_CONTEXT(context, count); + auto expanded_scale = context->GetOutputShapeRange(MOE_INIT_ROUTING_CUSTOM_OUTPUT_EXPANDED_SCALE); + OP_CHECK_NULL_WITH_CONTEXT(context, expanded_scale); + + // Print the shape ranges of the outputs before InferShapeRange + OPS_LOG_D(context->GetNodeName(), "Before InferShapeRange, expanded_x->GetMin() = %s", + ops::Shape2String(*(expanded_x->GetMin())).c_str()); + OPS_LOG_D(context->GetNodeName(), "Before InferShapeRange, expanded_x->GetMax() = %s", + ops::Shape2String(*(expanded_x->GetMax())).c_str()); + + OPS_LOG_D(context->GetNodeName(), "Before InferShapeRange, expanded_row_idx->GetMin() = %s", + ops::Shape2String(*(expanded_row_idx->GetMin())).c_str()); + OPS_LOG_D(context->GetNodeName(), "Before InferShapeRange, expanded_row_idx->GetMax() = %s", + ops::Shape2String(*(expanded_row_idx->GetMax())).c_str()); + + OPS_LOG_D(context->GetNodeName(), "Before InferShapeRange, count->GetMin() = %s", + ops::Shape2String(*(count->GetMin())).c_str()); + OPS_LOG_D(context->GetNodeName(), "Before InferShapeRange, count->GetMax() = %s", + ops::Shape2String(*(count->GetMax())).c_str()); + + OPS_LOG_D(context->GetNodeName(), "Before InferShapeRange, expanded_scale->GetMin() = %s", + ops::Shape2String(*(expanded_scale->GetMin())).c_str()); + OPS_LOG_D(context->GetNodeName(), "Before InferShapeRange, expanded_scale->GetMax() = %s", + ops::Shape2String(*(expanded_scale->GetMax())).c_str()); + + // Set the dim num and dim of the outputs' shape range object + if (expanded_x->GetMin() != nullptr && expanded_x->GetMax() != nullptr) { + expanded_x->GetMin()->SetDimNum(DIM_TWO); + expanded_x->GetMax()->SetDimNum(DIM_TWO); + for (size_t i = 0; i < DIM_TWO; i++) { + expanded_x->GetMin()->SetDim(i, 0); + expanded_x->GetMax()->SetDim(i, -1); + } + } + + if (expanded_row_idx->GetMin() != nullptr && expanded_row_idx->GetMax() != nullptr) { + expanded_row_idx->GetMin()->SetDimNum(DIM_ONE); + expanded_row_idx->GetMax()->SetDimNum(DIM_ONE); + expanded_row_idx->GetMin()->SetDim(0, 0); + expanded_row_idx->GetMax()->SetDim(0, -1); + } + + if (count->GetMin() != nullptr && count->GetMax() != nullptr) { + count->GetMin()->SetDimNum(DIM_ONE); + count->GetMax()->SetDimNum(DIM_ONE); + count->GetMin()->SetDim(0, 0); + count->GetMax()->SetDim(0, -1); + } + + if (expanded_scale->GetMin() != nullptr && expanded_scale->GetMax() != nullptr) { + expanded_scale->GetMin()->SetDimNum(DIM_ONE); + expanded_scale->GetMax()->SetDimNum(DIM_ONE); + expanded_scale->GetMin()->SetDim(0, 0); + expanded_scale->GetMax()->SetDim(0, -1); + } + + // Print the shape ranges of the outputs after InferShapeRange + OPS_LOG_D(context->GetNodeName(), "After InferShapeRange, expanded_x->GetMin() = %s", + ops::Shape2String(*(expanded_x->GetMin())).c_str()); + OPS_LOG_D(context->GetNodeName(), "After InferShapeRange, expanded_x->GetMax() = %s", + ops::Shape2String(*(expanded_x->GetMax())).c_str()); + + OPS_LOG_D(context->GetNodeName(), "After InferShapeRange, expanded_row_idx->GetMin() = %s", + ops::Shape2String(*(expanded_row_idx->GetMin())).c_str()); + OPS_LOG_D(context->GetNodeName(), "After InferShapeRange, expanded_row_idx->GetMax() = %s", + ops::Shape2String(*(expanded_row_idx->GetMax())).c_str()); + + OPS_LOG_D(context->GetNodeName(), "After InferShapeRange, count->GetMin() = %s", + ops::Shape2String(*(count->GetMin())).c_str()); + OPS_LOG_D(context->GetNodeName(), "After InferShapeRange, count->GetMax() = %s", + ops::Shape2String(*(count->GetMax())).c_str()); + + OPS_LOG_D(context->GetNodeName(), "After InferShapeRange, expanded_scale->GetMin() = %s", + ops::Shape2String(*(expanded_scale->GetMin())).c_str()); + OPS_LOG_D(context->GetNodeName(), "After InferShapeRange, expanded_scale->GetMax() = %s", + ops::Shape2String(*(expanded_scale->GetMax())).c_str()); + + OPS_LOG_D(context->GetNodeName(), "End to do MoeInitRoutingCustomInferRange."); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(MoeInitRoutingCustom) + .InferShape(InferShape4MoeInitRoutingCustom) + .InferDataType(InferDataType4MoeInitRoutingCustom) + .InferShapeRange(InferShapeRange4MoeInitRoutingCustom); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_tiling.cpp b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_tiling.cpp new file mode 100644 index 000000000..e05f1ec07 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_tiling.cpp @@ -0,0 +1,1278 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_init_routing_custom_tiling.cpp + * \brief + */ +#include "moe_init_routing_custom_tiling.h" +#include "register/op_def_registry.h" +#include "tiling/tiling_templates_registry.h" + +namespace optiling { +const static int64_t NUM_TWO = 2; +const static int64_t NUM_THREE = 3; +const static int64_t NUM_FOUR = 4; +const static int64_t NUM_FIVE = 5; +const static int64_t MRG_LIST_NUM = 4; +const static int64_t SORT32_ALIGN_ELEMENT = 32; +const static int64_t ONE_BLOCK_BYTE = 32; +const static size_t DIM_ONE = 1; +const static size_t DIM_TWO = 2; +const static int32_t SIZE_16 = 16; +const static int32_t SIZE_31 = 31; +const static int32_t LENGTH_1024 = 1024; +const static int64_t MAX_COLS_ONE_LOOP = 16376; +const static int64_t ASSIST_NUM = 256; +const static int64_t SPLIT_K_THRESHOLD = 512; +const static int64_t KV_FACTOR = 2; +const static int64_t ONE_CORE_SORT_BUFFER = 6; +const static int64_t EXPERT_IDX_MAX = 10240; +const static int64_t KV_MODE_EXPERT_IDX_MAX = EXPERT_IDX_MAX / KV_FACTOR; +const static int64_t ACTIVE_NUM_MIN_VALUE = static_cast(-1); +const static int64_t EXPERT_CAPACITY_MIN_VALUE = static_cast(0); + +const static int64_t INPUT_X_INDEX = 0; +const static int64_t INPUT_EXPERT_IDX_INDEX = 1; +const static int64_t INPUT_SCALE_INDEX = 2; +const static int64_t INPUT_OFFSET_INDEX = 3; +const static int64_t OUTPUT_EXPANDED_X_INDEX = 0; +const static int64_t OUTPUT_EXPANDED_ROW_IDX_INDEX = 1; +const static int64_t OUTPUT_EXPERT_TOKENS_COUNT_INDEX = 2; +const static int64_t OUTPUT_EXPANDED_SCALE_INDEX = 3; +const static int64_t ATTR_ACTIVE_NUM_INDEX = 0; +const static int64_t ATTR_EXPERT_CAPACITY_INDEX = 1; +const static int64_t ATTR_EXPERT_NUM_INDEX = 2; +const static int64_t ATTR_DROP_PAD_MODE_INDEX = 3; +const static int64_t ATTR_EXPERT_TOKEN_NUM_TYPE_INDEX = 4; +const static int64_t ATTR_EXPERT_TOKEN_NUM_FLAG_INDEX = 5; +const static int64_t ATTR_QUANT_MODE_INDEX = 6; +const static int64_t ATTR_EXPERT_RANGE_INDEX = 7; +const static int64_t ATTR_ROW_IDX_TYPE_INDEX = 8; +const static int64_t ATTR_EXPERT_RANGE_DIM = 2; +const static int64_t GATHER = 0; +const static int64_t SCATTER = 1; +const static int64_t UN_QUANT = -1L; +const static int64_t STATIC_QUANT = 0; +const static int64_t DYNAMIC_QUANT = 1; +const static int64_t CUMSUM = 0; +const static int64_t COUNT = 1; +const static int64_t KEY_VALUE = 2; +const static int64_t DROP_LESS = 0; +const static int64_t DROP_PAD = 1; +const static int64_t DYNAMIC_QUANT_COLS_BUFFER = 21; +const static int64_t DYNAMIC_QUANT_FULLLOAD_COLS_BUFFER = 13; +const static int64_t STATIC_QUANT_FULLLOAD_COLS_BUFFER = 11; + +const static int64_t DYNAMIC_QUANT_SRC_TO_DST_BUFFER = 15; +const static int64_t DYNAMIC_QUANT_SCALE_SIZE_64 = 64; +const static int64_t MAX_COLS_DYNAMIC_QUANT = 6144; +const static int64_t SIZE_INT32 = 4; +const static int64_t SIZE_INT16 = 2; +const static int64_t SIZE_INT8 = 1; +const static int64_t SIZE_FP32 = 4; + +const static uint64_t TILINGKEY_BASE = 1000000; +const static uint64_t SORT_CORE_TILINGKEY_BASE = 100000; +const static uint64_t QUANT_MODE_TILINGKEY_BASE = 10000; +const static uint64_t ROWIDX_TYPE_TILINGKEY_BASE = 1000; +const static uint64_t DROP_MODE_TILINGKEY_BASE = 100; + +// Tiling Key for performance puncturing +const static uint64_t PERFORMANCE_TILINGKEY_X_1_7168_EXPERT_IDX_1_8_SCALE_256_7168 = 2000000; +const static uint64_t UNQUANTIZED_FULLLOAD_TILINGKEY = 2100000; +const static uint64_t STATIC_QUANT_FULLLOAD_TILINGKEY = 2200000; +const static uint64_t DYNAMIC_QUANT_FULLLOAD_TILINGKEY = 2300000; +const static uint64_t DYNAMIC_QUANT_EPFULLLOAD_TILINGKEY = 10000; +const static uint64_t DYNAMIC_QUANT_SMOOTHTYPE_FULLLOAD_TILINGKEY = 1000; + +const static int64_t PERFORMANCE_MODE_TOP_K = 8; +const static int64_t PERFORMANCE_MODE_BS_MIN = 384; +const static int64_t PERFORMANCE_MODE_BS_MAX = 8192; +const static int64_t PERFORMANCE_MODE_RANGE_MAX = 32; +const static int64_t PERFORMANCE_MODE_MAX_BATCH_SIZE_TOP_K = PERFORMANCE_MODE_BS_MAX * PERFORMANCE_MODE_TOP_K; +const static int64_t PERFORMANCE_MODE_MAX_ONE_CORE_GATHER = 21845; + +const static int64_t gatherFirstN = 100; +const static int64_t gatherFirstScale = 8; +const static int64_t scale1H = 1; +const static int64_t scaleEH = 2; +const static int64_t ONE_REPEAT_SORT_NUM = 32; + +enum class PerformanceMode : int32_t { + COMMON = 0, + ONE_CORE_GATHER_SORT = 1, + MULTI_CORE_GATHER_SORT = 2, +}; + +static constexpr int64_t KEY_VALUE_MODE_DIM0_NUM = 2; + +#define unlikely(x) __builtin_expect((x), 0) + +#define CHECK_FAIL(context, cond, ...) \ + do { \ + if (cond) { \ + OPS_LOG_E(context->GetNodeName(), ##__VA_ARGS__); \ + return ge::GRAPH_FAILED; \ + } \ + } while (0) + +#define OP_CHECK_NULL_WITH_CONTEXT(context, ptr) \ + do { \ + if (unlikely((ptr) == nullptr)) { \ + const char* name = (unlikely(((context) == nullptr) || (context)->GetNodeName() == nullptr)) ? \ + "nil" : \ + (context)->GetNodeName(); \ + OPS_LOG_E(name, "%s is nullptr!", #ptr); \ + return ge::GRAPH_FAILED; \ + } \ + } while (0) + +template +static T1 CeilDiv(T1 a, T2 b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +template +typename std::enable_if ::value, T>::type CeilAlign(T x, T align) { + return CeilDiv(x, align) * align; +} + +inline static int64_t CeilLog4(int64_t x) +{ + return static_cast(std::ceil(std::log(x) / std::log(NUM_FOUR))); +} + +inline static int64_t Align(int64_t elementNum, int64_t bytes) +{ + if (bytes == 0) { + return 0; + } + return (elementNum * bytes + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE / bytes; +} + +inline static int64_t AlignBytes(int64_t elementNum, int64_t bytes) +{ + return (elementNum * bytes + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; +} + +inline static int64_t GetPerOrLastValue(int64_t x, int64_t y) +{ + if (y == 0) { + return 0; + } + return x <= y ? x : x % y; +} + +inline static int64_t AlignOneBlockByteCeil(int64_t x) +{ + return x / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; +} + +class MoeInitRountingCustomTilingBase : public TilingBaseClass { +public: + explicit MoeInitRountingCustomTilingBase(gert::TilingContext *context) : TilingBaseClass(context) + { + Reset(); + } + ~MoeInitRountingCustomTilingBase() override = default; + + void Reset(gert::TilingContext *context) override + { + TilingBaseClass::Reset(context); + Reset(); + } + +protected: + bool IsCapable() override + { + return true; + } + ge::graphStatus GetPlatformInfo() override; + ge::graphStatus GetShapeAttrsInfo() override; + ge::graphStatus DoOpTiling() override; + ge::graphStatus DoLibApiTiling() override; + uint64_t GetTilingKey() const override; + ge::graphStatus GetWorkspaceSize() override; + ge::graphStatus PostTiling() override; + void Reset(); + +private: + ge::graphStatus CheckAttr(); + ge::graphStatus CheckOutShape(); + ge::graphStatus CheckInputShape(); + ge::graphStatus CheckDtype(); + void Tiling4GatherOutCompute(); + void Tiling4SortOutCompute(); + void Tiling4VMSMiddleCompute(); + void Tiling4VBSCompute(); + void Tiling4ExpertTokensCountCompute(); + void ShowTilingData(); + void Tinlig4VBSMultiCoreCompute(MoeCustomVBSComputeTilingData *tilingData); + void Tinlig4VBSOneCoreCompute(MoeCustomVBSComputeTilingData *tilingData); + bool IsPerformanceMode_X_1_7168_EXPERT_IDX_1_8_SCALE_256_7168() const; + bool IsFullLoad(); + int64_t IsGatherFirstFullLoad(); + void SetGatherTilingData(MoeCustomSrcToDstCapacityComputeTilingData *tilingData, int64_t perCoreRows, + int64_t lastCoreRows, int64_t cols); + void SetGatherTilingDataCols(MoeCustomSrcToDstCapacityComputeTilingData *tilingData, int64_t baseMaxCols, int64_t cols); + void SetGatherTilingDataRows(MoeCustomSrcToDstCapacityComputeTilingData *tilingData, int64_t perCoreRows, + int64_t lastCoreRows, int64_t basePerLoopMaxRows); + void Tiling4SrcToDstDropPadCompute(); + void Tiling4SrcToDstDropPadDynamicCompute(); + void Tiling4SrcToDstCompute(); + PerformanceMode GetPerformanceMode() const; + + int64_t aivNum; + int64_t sortLoopMaxElement = 0; + int64_t mrgSortListMaxElement = 1504; + int64_t totalLength_ = 0; + int64_t n_ = 0; + int64_t k_ = 0; + int64_t cols_ = 0; + int64_t inuptXDtypeSize_; + + int64_t expertStart_ = 0; + int64_t expertEnd_ = 0; + int64_t isInputScale_ = 0; + int64_t isInputOffset_ = 0; + + int64_t sortMode_ = 0; + int64_t rowIdxTytpe_ = 0; + int64_t activeNum_ = -1L; + int64_t expertCapacity_ = -1L; + int64_t expertNum_ = -1L; + int64_t dropPadMode_ = -1L; + int64_t expertTokensNumType_ = -1L; + bool expertTokensNumFlag_ = false; + int64_t quantMode_ = 0; + int64_t rowIdxType_ = -1L; + + bool isFullload_ = false; + int64_t gatherFirstFullload_ = 0; + int64_t ep_ = 0; + int64_t smoothType_ = 0; + + const gert::StorageShape *xShapePtr_ = nullptr; + const gert::StorageShape *expertIdxShapePtr_ = nullptr; + const gert::StorageShape *scaleShapePtr_ = nullptr; + const gert::StorageShape *offsetShapePtr_ = nullptr; + + const int64_t *activeNumPtr_ = nullptr; + const int64_t *expertCapacityPtr_ = nullptr; + const int64_t *expertNumPtr_ = nullptr; + const int64_t *dropPadModePtr_ = nullptr; + const int64_t *expertTokensNumTypePtr_ = nullptr; + const bool *expertTokensNumFlagPtr_ = nullptr; + const int64_t *quantModePtr_ = nullptr; + const gert::ContinuousVector *activeExpertRangeListPtr_; + const int64_t *rowIdxTypePtr_ = nullptr; + + const gert::StorageShape *expandedXShapePtr_ = nullptr; + const gert::StorageShape *expandedRowIdxShapePtr_ = nullptr; + const gert::StorageShape *expertTokensCountOrCumsumShapePtr_ = nullptr; + const gert::StorageShape *expandedScaleShapePtr_ = nullptr; + + const gert::Shape performXShape = gert::Shape({1, 7168}); + const gert::Shape performExpertIdxShape = gert::Shape({1, 8}); + const gert::Shape performScaleShape = gert::Shape({256, 7168}); + + const char *opName = ""; + MoeInitRoutingCustomTilingData moeInitRoutingCustomTilingData; +}; + +void MoeInitRountingCustomTilingBase::Reset() +{ + opName = nullptr; + return; +} + +ge::graphStatus MoeInitRountingCustomTilingBase::GetPlatformInfo() +{ + auto compileInfoPtr = reinterpret_cast(context_->GetCompileInfo()); + if (compileInfoPtr == nullptr) { + auto platformInfo = context_->GetPlatformInfo(); + CHECK_FAIL(context_, platformInfo == nullptr, "fail to get platform info"); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + aivNum = ascendcPlatform.GetCoreNumAiv(); + uint64_t ubSize = 0; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); + aicoreParams_.blockDim = aivNum; + aicoreParams_.ubSize = static_cast(ubSize); + } else { + aivNum = compileInfoPtr->aivNum; + aicoreParams_.blockDim = aivNum; + aicoreParams_.ubSize = compileInfoPtr->ubSize; + } + CHECK_FAIL(context_, aivNum <= 0 || aicoreParams_.ubSize <= 0, "fail to init platform info"); + moeInitRoutingCustomTilingData.set_coreNum(aivNum); + OPS_LOG_I(context_->GetNodeName(), "---PlatformInfo--- aivNum is: %ld, ubSizePlatForm is: %ld ", aivNum, aicoreParams_.ubSize); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeInitRountingCustomTilingBase::CheckAttr() +{ + quantMode_ = *quantModePtr_; + moeInitRoutingCustomTilingData.set_quantMode(quantMode_); + OPS_LOG_I(context_->GetNodeName(), "quant_mode is: %ld ", quantMode_); + + dropPadMode_ = *dropPadModePtr_; + moeInitRoutingCustomTilingData.set_dropPadMode(dropPadMode_); + CHECK_FAIL(context_, (dropPadMode_ != DROP_LESS) && (dropPadMode_ != DROP_PAD), + "drop_pad_mode should be %ld or %ld", DROP_LESS, DROP_PAD); + + rowIdxTytpe_ = *rowIdxTypePtr_; + moeInitRoutingCustomTilingData.set_rowIdxType(rowIdxTytpe_); + OPS_LOG_I(context_->GetNodeName(), "row_idx_type is: %ld ", rowIdxTytpe_); + + activeNum_ = *activeNumPtr_; + if (dropPadMode_ == DROP_LESS) { + CHECK_FAIL(context_, activeNum_ < ACTIVE_NUM_MIN_VALUE, + "active_num should be greater than or equal to 0"); + } + + expertNum_ = *expertNumPtr_; + moeInitRoutingCustomTilingData.set_expertNum(expertNum_); + if (expertNum_ <= 0) { + OPS_LOG_E(context_->GetNodeName(), "expert_num should be greater than 0"); + return ge::GRAPH_FAILED; + } + if (activeExpertRangeListPtr_->GetSize() != ATTR_EXPERT_RANGE_DIM && activeExpertRangeListPtr_->GetSize() != 0) { + OPS_LOG_E(context_, "The dim number of expert_range should be %ld or 0(no input)", ATTR_EXPERT_RANGE_DIM); + return ge::GRAPH_FAILED; + } + if (activeExpertRangeListPtr_->GetSize() == 0) { + expertStart_ = 0; + expertEnd_ = expertNum_; + } else { + const int64_t *expertRangeList = reinterpret_cast(activeExpertRangeListPtr_->GetData()); + expertStart_ = expertRangeList[0]; + expertEnd_ = expertRangeList[1]; + } + moeInitRoutingCustomTilingData.set_expertStart(expertStart_); + moeInitRoutingCustomTilingData.set_expertEnd(expertEnd_); + moeInitRoutingCustomTilingData.set_actualExpertNum(expertEnd_ - expertStart_); + OPS_LOG_I(context_, "expert_start is: %ld, expert_end is: %ld, actualExpertNum is: %ld ", expertStart_, expertEnd_, + expertEnd_ - expertStart_); + + n_ = xShapePtr_->GetStorageShape().GetDim(0); + expertCapacity_ = *expertCapacityPtr_; + moeInitRoutingCustomTilingData.set_expertCapacity(expertCapacity_); + if (dropPadMode_ == DROP_PAD) { + CHECK_FAIL(context_, expertCapacity_ <= EXPERT_CAPACITY_MIN_VALUE || expertCapacity_ > n_, + "expert_Capacity should be greater than 0 and less than %ld", n_); + CHECK_FAIL(context_, rowIdxTytpe_ == SCATTER, "rowIdxTytpe should be 0 when droppadmode is 1"); + CHECK_FAIL(context_, expertStart_ != 0 || expertEnd_ != expertNum_, + "expert_range should be [0, %ld] when droppadmode is 1", expertNum_); + } + + expertTokensNumType_ = *expertTokensNumTypePtr_; + moeInitRoutingCustomTilingData.set_expertTokensNumType(expertTokensNumType_); + CHECK_FAIL(context_, (expertTokensNumType_ != COUNT) && (expertTokensNumType_ != KEY_VALUE) && (expertTokensNumType_ != CUMSUM), + "expert_tokens_num_type currently not support %ld", expertTokensNumType_); + + expertTokensNumFlag_ = *expertTokensNumFlagPtr_; + if (dropPadMode_ == DROP_PAD && expertTokensNumFlag_) { + CHECK_FAIL(context_, expertTokensNumType_ != COUNT, "In DROP_PAD mode and expert_tokens_num_flag is true, expert_tokens_num_type only supports COUNT, but got %ld", expertTokensNumType_);} + if (expertTokensNumFlag_) { + moeInitRoutingCustomTilingData.set_expertTokensNumFlag(1); + } else { + moeInitRoutingCustomTilingData.set_expertTokensNumFlag(0); + } + + CHECK_FAIL(context_, expertStart_ < 0, "expert_start should be greater than or equal to 0"); + CHECK_FAIL(context_, expertStart_ >= expertEnd_, "expert_start should be less than expert_end"); + CHECK_FAIL(context_, expertEnd_ > expertNum_, "expert_end should be less than or equal to %ld", expertNum_); + if (expertTokensNumType_ == KEY_VALUE) { + CHECK_FAIL(context_, expertEnd_ > KV_MODE_EXPERT_IDX_MAX, "expert_end should be less than or equal to %ld in KEY_VALUE mode", + KV_MODE_EXPERT_IDX_MAX); + } else { + CHECK_FAIL(context_, expertEnd_ > EXPERT_IDX_MAX, "expert_end should be less than or equal to %ld", EXPERT_IDX_MAX); + } + CHECK_FAIL(context_, quantMode_ != UN_QUANT && quantMode_ != DYNAMIC_QUANT && quantMode_ != STATIC_QUANT, "quant_mode currently support %ld, %ld or %ld", UN_QUANT, DYNAMIC_QUANT, STATIC_QUANT); + CHECK_FAIL(context_, rowIdxTytpe_ != SCATTER && rowIdxTytpe_ != GATHER, "row_idx_type currently support %ld or %ld", SCATTER, GATHER); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeInitRountingCustomTilingBase::CheckInputShape() +{ + const gert::Shape xShape = xShapePtr_->GetStorageShape(); + OPS_LOG_I(context_->GetNodeName(), "input x shape: %s ", ops::Shape2String(xShape).c_str()); + const gert::Shape expertIdxShape = expertIdxShapePtr_->GetStorageShape(); + OPS_LOG_I(context_->GetNodeName(), "input expert_idx shape: %s.", ops::Shape2String(expertIdxShape).c_str()); + + // 参数校验 + CHECK_FAIL(context_, xShape.GetDimNum() != DIM_TWO, "The dim number of x should be %lu.", DIM_TWO); + CHECK_FAIL(context_, expertIdxShape.GetDimNum() != DIM_TWO, "The dim number of expert_idx should be %lu.", DIM_TWO); + CHECK_FAIL(context_, xShape.GetDim(0) != expertIdxShape.GetDim(0), context_->GetNodeName(), "Input rows should be same."); + + n_ = expertIdxShape.GetDim(0); + k_ = expertIdxShape.GetDim(1); + cols_ = xShape.GetDim(1); + moeInitRoutingCustomTilingData.set_n(n_); + moeInitRoutingCustomTilingData.set_k(k_); + moeInitRoutingCustomTilingData.set_cols(cols_); + totalLength_ = n_ * k_; + if (activeNum_ == 0 || activeNum_ == ACTIVE_NUM_MIN_VALUE) { + activeNum_ = totalLength_; + } else { + activeNum_ = std::min(activeNum_, totalLength_); + } + moeInitRoutingCustomTilingData.set_activeNum(activeNum_); + + inuptXDtypeSize_ = + static_cast(ge::GetSizeByDataType(context_->GetInputDesc(INPUT_X_INDEX)->GetDataType())); + OPS_LOG_I(context_->GetNodeName(), "Input x dtype size is: %ld. ", inuptXDtypeSize_); + + if (quantMode_ == UN_QUANT && scaleShapePtr_ != nullptr) { + auto scaleShape = scaleShapePtr_->GetStorageShape(); + OPS_LOG_I(context_->GetNodeName(), "input scale shape: %s", ops::Shape2String(scaleShape).c_str()); + auto scaleDimNum = static_cast(scaleShape.GetDimNum()); + CHECK_FAIL(context_, + scaleDimNum != 1, + context_->GetNodeName(), "The dim number of scale should be 1, current is %ld", scaleDimNum); + auto scaleDim0 = static_cast(scaleShape.GetDim(0)); + CHECK_FAIL(context_, + scaleDim0 != n_, + "The first dim of scale should be n_, current is %ld", scaleDim0); + } + + if (quantMode_ == STATIC_QUANT) { + CHECK_FAIL(context_, scaleShapePtr_ == nullptr, "scale is null"); + CHECK_FAIL(context_, offsetShapePtr_ == nullptr, "offset is null"); + auto scaleShape = scaleShapePtr_->GetStorageShape(); + OPS_LOG_I(context_->GetNodeName(), "input scale shape: %s", ops::Shape2String(scaleShape).c_str()); + auto scaleDimNum = static_cast(scaleShape.GetDimNum()); + CHECK_FAIL(context_, + scaleDimNum != 1, + "The dim number of scale should be 1, current is %ld", scaleDimNum); + auto scaleDim0 = static_cast(scaleShape.GetDim(0)); + CHECK_FAIL(context_, + scaleDim0 != 1, + "The first dim of scale should be 1, current is %ld", scaleDim0); + auto offsetShape = offsetShapePtr_->GetStorageShape(); + OPS_LOG_I(context_->GetNodeName(), "input offset shape: %s", ops::Shape2String(offsetShape).c_str()); + auto offsetDimNum = static_cast(offsetShape.GetDimNum()); + CHECK_FAIL(context_, + offsetDimNum != 1, + "The dim number of offset should be 1, current is %ld", offsetDimNum); + auto offsetDim0 = static_cast(offsetShape.GetDim(0)); + CHECK_FAIL(context_, + offsetDim0 != 1, + "The first dim of offset should be 1, current is %ld", offsetDim0); + } + + if (quantMode_ == DYNAMIC_QUANT && scaleShapePtr_ != nullptr) { + auto scaleShape = scaleShapePtr_->GetStorageShape(); + OPS_LOG_I(context_->GetNodeName(), "input scale shape: %s", ops::Shape2String(scaleShape).c_str()); + auto scaleDimNum = static_cast(scaleShape.GetDimNum()); + CHECK_FAIL(context_, + scaleDimNum != NUM_TWO, + "The dim number of scale should be 2, current is %ld", scaleDimNum); + auto scaleDim0 = static_cast(scaleShape.GetDim(0)); + CHECK_FAIL(context_, + scaleDim0 != (expertEnd_ - expertStart_) && scaleDim0 != 1, + "The first dim of scale should be %ld or 1, current is %ld", (expertEnd_ - expertStart_), scaleDim0); + auto scaleDim1 = static_cast(scaleShape.GetDim(1)); + CHECK_FAIL(context_, + scaleDim1 != cols_, + "The second dim of scale should be %ld, current is %ld", cols_, scaleDim0); + if (scaleDim0 == 1) { + smoothType_ = scale1H; + } else { + smoothType_ = scaleEH; + } + moeInitRoutingCustomTilingData.set_smoothType(smoothType_); + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeInitRountingCustomTilingBase::CheckOutShape() +{ + const gert::Shape expandedXShape = context_->GetOutputShape(0)->GetStorageShape(); + OPS_LOG_I(context_->GetNodeName(), "expanded_x shape: %s.", ops::Shape2String(expandedXShape).c_str()); + const gert::Shape expandedRowIdxShape = context_->GetOutputShape(1)->GetStorageShape(); + OPS_LOG_I(context_->GetNodeName(), "expanded_row_idx shape: %s.", ops::Shape2String(expandedRowIdxShape).c_str()); + const gert::Shape expertTokensCountOrCumsumShape = context_->GetOutputShape(NUM_TWO)->GetStorageShape(); + OPS_LOG_I(context_->GetNodeName(), "expert_tokens_count_or_cumsum shape: %s.", ops::Shape2String(expertTokensCountOrCumsumShape).c_str()); + + size_t expandedXDimNum = expandedXShape.GetDimNum(); + if (dropPadMode_ > 0) { + CHECK_FAIL(context_, expandedXDimNum != NUM_THREE, "The dim number of expandedX should be 3."); + CHECK_FAIL(context_, expandedXShape.GetDim(0) != expertNum_, "The first dim of expandedX should be %ld.", expertNum_); + CHECK_FAIL(context_, expandedXShape.GetDim(1) != expertCapacity_, "The second dim of expandedX should be %ld.", + expertCapacity_); + CHECK_FAIL(context_, + expandedXShape.GetDim(NUM_TWO) != cols_, + "The third dim of expandedX should be %ld.", cols_); + } else { + CHECK_FAIL(context_,expandedXDimNum != DIM_TWO, "The dim number of expandedX should be 2."); + int64_t firstDim = totalLength_; + firstDim = activeNum_ == 0 ? firstDim : std::min(firstDim, activeNum_); + CHECK_FAIL(context_, expandedXShape.GetDim(0) != firstDim, "The first dim of expandedX should be %ld.", firstDim); + CHECK_FAIL(context_, expandedXShape.GetDim(1) != cols_, + "The second dim of expandedX should be %ld.", cols_); + } + + CHECK_FAIL(context_, expandedRowIdxShape.GetDimNum() != DIM_ONE, + "The dim number of expanded_row_idx should be 1."); + CHECK_FAIL(context_, + expandedRowIdxShape.GetDim(0) != totalLength_, + "The first dim of expanded_row_idx and expanded_expert_idx should be %ld.", totalLength_); + + if(expertTokensNumFlag_){ + if (expertTokensNumType_ == KEY_VALUE) { + CHECK_FAIL(context_, + expertTokensCountOrCumsumShape.GetDimNum() != DIM_TWO, + "The dim number of expert_tokens_count_or_cumsum should be 2 when in KEY_VALUE mode."); + CHECK_FAIL(context_, expertTokensCountOrCumsumShape.GetDim(0) != expertNum_, + "The first dim of expert_tokens_count_or_cumsum should be %ld.", expertNum_); + CHECK_FAIL(context_, expertTokensCountOrCumsumShape.GetDim(1) != KEY_VALUE_MODE_DIM0_NUM, + "The second dim of expert_tokens_count_or_cumsum should be %ld.", + KEY_VALUE_MODE_DIM0_NUM); + } else { + CHECK_FAIL(context_, expertTokensCountOrCumsumShape.GetDimNum() != DIM_ONE, + "The dim number of expert_tokens_count_or_cumsum should be 1 when not in KEY_VALUE mode."); + CHECK_FAIL(context_, expertTokensCountOrCumsumShape.GetDim(0) != (expertEnd_ - expertStart_), + "The first dim of expert_tokens_count_or_cumsum should be %ld.", (expertEnd_ - expertStart_)); + } + } + + if (quantMode_ != STATIC_QUANT && scaleShapePtr_ != nullptr) { + const gert::Shape expandedScaleShape = context_->GetOutputShape(3)->GetStorageShape(); + OPS_LOG_I(context_->GetNodeName(), "expanded_scale shape: %s.", ops::Shape2String(expandedScaleShape).c_str()); + size_t expandedScaleDimNum = expandedScaleShape.GetDimNum(); + CHECK_FAIL(context_, expandedScaleDimNum != DIM_ONE, "The dim number of expanded_scale should be 1."); + if (dropPadMode_ > 0) { + CHECK_FAIL(context_, expandedScaleShape.GetDim(0) != expertNum_ * expertCapacity_, + "The first dim of expanded_scale should be %ld.", expertNum_ * expertCapacity_); + } else { + int64_t firstDim = totalLength_; + firstDim = activeNum_ == 0 ? firstDim : std::min(firstDim, activeNum_); + CHECK_FAIL(context_, expandedScaleShape.GetDim(0) != firstDim, + "The first dim of expanded_scale should be %ld.", firstDim); + } + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeInitRountingCustomTilingBase::GetShapeAttrsInfo() +{ + OPS_LOG_I(context_->GetNodeName(), "TilingContext: %s.", context_->GetNodeName()); + + xShapePtr_ = context_->GetInputShape(INPUT_X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, xShapePtr_); + + expertIdxShapePtr_ = context_->GetInputShape(INPUT_EXPERT_IDX_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, expertIdxShapePtr_); + + scaleShapePtr_ = context_->GetOptionalInputShape(INPUT_SCALE_INDEX); + if (scaleShapePtr_ == nullptr) { + OPS_LOG_I(context_->GetNodeName(), "optional input scale is null"); + } else { + isInputScale_ = 1; + } + moeInitRoutingCustomTilingData.set_isInputScale(isInputScale_); + + offsetShapePtr_ = context_->GetOptionalInputShape(INPUT_OFFSET_INDEX); + if (offsetShapePtr_ == nullptr) { + OPS_LOG_I(context_->GetNodeName(), "optional input offset is null"); + } else { + isInputOffset_ = 1; + } + moeInitRoutingCustomTilingData.set_isInputOffset(isInputOffset_); + + expandedXShapePtr_ = context_->GetOutputShape(OUTPUT_EXPANDED_X_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, expandedXShapePtr_); + expandedRowIdxShapePtr_ = context_->GetOutputShape(OUTPUT_EXPANDED_ROW_IDX_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, expandedRowIdxShapePtr_); + expertTokensCountOrCumsumShapePtr_ = context_->GetOutputShape(OUTPUT_EXPERT_TOKENS_COUNT_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, expertTokensCountOrCumsumShapePtr_); + expandedScaleShapePtr_ = context_->GetOutputShape(OUTPUT_EXPANDED_SCALE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, expandedScaleShapePtr_); + + auto attrs = context_->GetAttrs(); + OP_CHECK_NULL_WITH_CONTEXT(context_, attrs); + activeNumPtr_ = attrs->GetAttrPointer(ATTR_ACTIVE_NUM_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, activeNumPtr_); + expertCapacityPtr_ = attrs->GetAttrPointer(ATTR_EXPERT_CAPACITY_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, expertCapacityPtr_); + expertNumPtr_ = attrs->GetAttrPointer(ATTR_EXPERT_NUM_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, expertNumPtr_); + dropPadModePtr_ = attrs->GetAttrPointer(ATTR_DROP_PAD_MODE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, dropPadModePtr_); + expertTokensNumTypePtr_ = attrs->GetAttrPointer(ATTR_EXPERT_TOKEN_NUM_TYPE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, expertTokensNumTypePtr_); + expertTokensNumFlagPtr_ = attrs->GetAttrPointer(ATTR_EXPERT_TOKEN_NUM_FLAG_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, expertTokensNumFlagPtr_); + quantModePtr_ = attrs->GetAttrPointer(ATTR_QUANT_MODE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, quantModePtr_); + activeExpertRangeListPtr_ = attrs->GetAttrPointer(ATTR_EXPERT_RANGE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, activeExpertRangeListPtr_); + rowIdxTypePtr_ = attrs->GetAttrPointer(ATTR_ROW_IDX_TYPE_INDEX); + OP_CHECK_NULL_WITH_CONTEXT(context_, rowIdxTypePtr_); + return ge::GRAPH_SUCCESS; +} + +void MoeInitRountingCustomTilingBase::ShowTilingData() +{ + int64_t isFullloadInt = 1 ? isFullload_ == true : 0; + OPS_LOG_I(context_->GetNodeName(), "isFullload: %ld, gatherFirstFullload: %ld, ep: %ld", isFullloadInt, gatherFirstFullload_, ep_); +} + +int64_t MoeInitRountingCustomTilingBase::IsGatherFirstFullLoad() { + if (ep_ == 0) { + return 0; + } else if (n_ >= gatherFirstN && (expertEnd_-expertStart_) * gatherFirstScale <= expertNum_) { + return 1; + } + return 0; +} + +bool MoeInitRountingCustomTilingBase::IsFullLoad() { + int64_t perCoreTokens = 1; + if (expertStart_ == 0 && expertEnd_ == expertNum_) { + ep_ = 0; + if (quantMode_ != 1) { + perCoreTokens = n_ / aivNum; + int64_t remainder = n_ % aivNum; + // NUM_TWO is Max xRows need add 2 becauseof the left and right row may be another row. + perCoreTokens = remainder <= 1 ? perCoreTokens + 1 : perCoreTokens + NUM_TWO; + } + } else { + ep_ = 1; + perCoreTokens = 1; + } + moeInitRoutingCustomTilingData.set_ep(ep_); + + if (totalLength_ > sortLoopMaxElement || this->dropPadMode_ == 1) { + return false; + } + + gatherFirstFullload_ = IsGatherFirstFullLoad(); + moeInitRoutingCustomTilingData.set_gatherFirstFullload(gatherFirstFullload_); + int64_t tileLength = Align(this->totalLength_, int64_t(sizeof(int32_t))); + int64_t sortNum = CeilDiv(tileLength, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + + int64_t sortSpace = sortNum * sizeof(int32_t) * ONE_CORE_SORT_BUFFER; + int64_t rowIdxSpace = sortNum * sizeof(int32_t) * NUM_THREE; + int64_t expertSpace = CeilDiv(this->expertNum_ * int64_t(sizeof(int64_t)), ONE_BLOCK_BYTE) * ONE_BLOCK_BYTE * NUM_TWO; + int64_t gatherSpace = CeilDiv(cols_ * inuptXDtypeSize_, ONE_BLOCK_BYTE) * ONE_BLOCK_BYTE * perCoreTokens; + int64_t remainUb = aicoreParams_.ubSize - sortSpace - rowIdxSpace - expertSpace - LENGTH_1024; + + if (quantMode_ == -1) { + remainUb -= (gatherSpace + ONE_BLOCK_BYTE); + } else if (quantMode_ == 0) { + int64_t quantSpace = 0; + int64_t xAlignedCount = Align(this->cols_, int64_t(sizeof(int8_t))); + quantSpace = xAlignedCount * STATIC_QUANT_FULLLOAD_COLS_BUFFER * perCoreTokens; + remainUb -= (gatherSpace + quantSpace); + } else { + int64_t quantSpace = CeilDiv(cols_, ONE_BLOCK_BYTE) * ONE_BLOCK_BYTE * DYNAMIC_QUANT_FULLLOAD_COLS_BUFFER; + int64_t scaleOutSpace = ONE_BLOCK_BYTE * NUM_TWO; + remainUb -= (quantSpace + scaleOutSpace); + } + return remainUb > 0; +} + +bool MoeInitRountingCustomTilingBase::IsPerformanceMode_X_1_7168_EXPERT_IDX_1_8_SCALE_256_7168() const +{ + OPS_LOG_I(context_->GetNodeName(), "Begin IsPerformanceMode_X_1_7168_EXPERT_IDX_1_8_SCALE_256_7168() ..."); + bool result = false; + + // expert_range [0,256), quant_mode=DYNAMIC_QUANT + const gert::Shape performXShape_X_1_7168 = gert::Shape({1, 7168}); + const gert::Shape performExpertIdxShape_X_1_7168 = gert::Shape({1, 8}); + const gert::Shape performScaleShape_X_1_7168 = gert::Shape({256, 7168}); + + OP_CHECK_NULL_WITH_CONTEXT(context_, xShapePtr_); + OP_CHECK_NULL_WITH_CONTEXT(context_, expertIdxShapePtr_); + if (nullptr == scaleShapePtr_) { + result = false; + } else if (xShapePtr_->GetStorageShape() == performXShape_X_1_7168 && + expertIdxShapePtr_->GetStorageShape() == performExpertIdxShape_X_1_7168 && + scaleShapePtr_->GetStorageShape() == performScaleShape_X_1_7168 && offsetShapePtr_ == nullptr && + context_->GetInputDesc(INPUT_X_INDEX)->GetDataType() == ge::DT_BF16 && expertStart_ == 0 && + expertEnd_ == ASSIST_NUM && quantMode_ == DYNAMIC_QUANT && expertTokensNumType_ == KEY_VALUE) { + result = true; + } + OPS_LOG_I(context_->GetNodeName(), "End IsPerformanceMode_X_1_7168_EXPERT_IDX_1_8_SCALE_256_7168() ..."); + return result; +} + +PerformanceMode MoeInitRountingCustomTilingBase::GetPerformanceMode() const +{ + PerformanceMode result = PerformanceMode::COMMON; + if (expertNum_ != ASSIST_NUM || (expertEnd_ - expertStart_) > PERFORMANCE_MODE_RANGE_MAX || + n_ < PERFORMANCE_MODE_BS_MIN || n_ > PERFORMANCE_MODE_BS_MAX || k_ != PERFORMANCE_MODE_TOP_K) { + return result; + } + + // Judge performance mode according to totalLength_ + if (totalLength_ < PERFORMANCE_MODE_MAX_ONE_CORE_GATHER) { + OPS_LOG_I(context_->GetNodeName(), "totalLength_: %ld, PerformanceMode::ONE_CORE_GATHER_SORT", totalLength_); + result = PerformanceMode::ONE_CORE_GATHER_SORT; + } else if (totalLength_ <= PERFORMANCE_MODE_MAX_BATCH_SIZE_TOP_K) { + OPS_LOG_I(context_->GetNodeName(), "totalLength_: %ld, PerformanceMode::MULTI_CORE_GATHER_SORT", totalLength_); + result = PerformanceMode::MULTI_CORE_GATHER_SORT; + } + return result; +} + +ge::graphStatus MoeInitRountingCustomTilingBase::CheckDtype() +{ + auto inputXDtype_ = context_->GetInputDesc(INPUT_X_INDEX)->GetDataType(); + CHECK_FAIL(context_, inputXDtype_ != ge::DT_INT8 && inputXDtype_ != ge::DT_FLOAT16 && inputXDtype_ != ge::DT_BF16 && inputXDtype_ != ge::DT_FLOAT, + "The data type of input_X should be INT8, FLOAT16, BF16, FLOAT."); + CHECK_FAIL(context_, inputXDtype_ == ge::DT_INT8 && quantMode_ != UN_QUANT, + "When input_X is INT8, quantization is not supported."); + + auto expertIdxDtype_ = context_->GetInputDesc(INPUT_EXPERT_IDX_INDEX)->GetDataType(); + CHECK_FAIL(context_, expertIdxDtype_ != ge::DT_INT32, + "The data type of input_expertIdx should be INT32."); + + if (quantMode_ == STATIC_QUANT) { + auto scaleDtype_ = context_->GetOptionalInputDesc(INPUT_SCALE_INDEX)->GetDataType(); + CHECK_FAIL(context_, scaleDtype_ != ge::DT_FLOAT, + "The data type of input_scale should be FLOAT."); + + auto offsetDtype_ = context_->GetOptionalInputDesc(INPUT_OFFSET_INDEX)->GetDataType(); + CHECK_FAIL(context_, offsetDtype_ != ge::DT_FLOAT, + "The data type of input_offset should be FLOAT."); + } else { + if (scaleShapePtr_ != nullptr) { + auto scaleDtype_ = context_->GetOptionalInputDesc(INPUT_SCALE_INDEX)->GetDataType(); + CHECK_FAIL(context_, scaleDtype_ != ge::DT_FLOAT, + "The data type of input_scale should be FLOAT."); + } + } + + auto expandedXDtype_ = context_->GetOutputDesc(OUTPUT_EXPANDED_X_INDEX)->GetDataType(); + CHECK_FAIL(context_,expandedXDtype_ != ge::DT_INT8 && expandedXDtype_ != ge::DT_FLOAT16 && expandedXDtype_ != ge::DT_BF16 && expandedXDtype_ != ge::DT_FLOAT, + "The data type of output_expanded_X should be INT8, FLOAT16, BF16, FLOAT."); + + auto expandedRowIdxDtype_ = context_->GetOutputDesc(OUTPUT_EXPANDED_ROW_IDX_INDEX)->GetDataType(); + CHECK_FAIL(context_,expandedRowIdxDtype_ != ge::DT_INT32, + "The data type of output_expanded_row_idx should be INT32."); + + auto expertTokensCountOrCusumDtype_ = context_->GetOutputDesc(OUTPUT_EXPERT_TOKENS_COUNT_INDEX)->GetDataType(); + CHECK_FAIL(context_,expertTokensCountOrCusumDtype_ != ge::DT_INT64, + "The data type of output_expert_tokens_count_or_cumsum should be INT64."); + + if (quantMode_ == DYNAMIC_QUANT || (quantMode_ == UN_QUANT && scaleShapePtr_ != nullptr)) { + auto expandedScaleDtype_ = context_->GetOutputDesc(OUTPUT_EXPANDED_SCALE_INDEX)->GetDataType(); + CHECK_FAIL(context_,expandedScaleDtype_ != ge::DT_FLOAT, + "The data type of input_expanded_scale should be FLOAT."); + } + + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeInitRountingCustomTilingBase::DoOpTiling() +{ + auto ret = CheckAttr(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckInputShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckOutShape(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + ret = CheckDtype(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + + if (IsPerformanceMode_X_1_7168_EXPERT_IDX_1_8_SCALE_256_7168()) { + aivNum = totalLength_; + } + + sortLoopMaxElement = (aicoreParams_.ubSize - aivNum * ONE_BLOCK_BYTE) / (NUM_FOUR * NUM_TWO * NUM_FOUR) / + SORT32_ALIGN_ELEMENT * SORT32_ALIGN_ELEMENT; + + Tiling4VBSCompute(); + Tiling4VMSMiddleCompute(); + Tiling4SortOutCompute(); + Tiling4ExpertTokensCountCompute(); + Tiling4SrcToDstCompute(); + Tiling4SrcToDstDropPadCompute(); + Tiling4GatherOutCompute(); + isFullload_ = IsFullLoad(); + ShowTilingData(); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeInitRountingCustomTilingBase::DoLibApiTiling() +{ + return ge::GRAPH_SUCCESS; +} + +uint64_t MoeInitRountingCustomTilingBase::GetTilingKey() const +{ + if (isFullload_) { + if (quantMode_ == UN_QUANT) { + return UNQUANTIZED_FULLLOAD_TILINGKEY; + } else if (quantMode_ == STATIC_QUANT) { + return STATIC_QUANT_FULLLOAD_TILINGKEY; + } else { + return (DYNAMIC_QUANT_FULLLOAD_TILINGKEY + ep_ * DYNAMIC_QUANT_EPFULLLOAD_TILINGKEY + + smoothType_ * DYNAMIC_QUANT_SMOOTHTYPE_FULLLOAD_TILINGKEY); + } + } + else if (IsPerformanceMode_X_1_7168_EXPERT_IDX_1_8_SCALE_256_7168()) { + return PERFORMANCE_TILINGKEY_X_1_7168_EXPERT_IDX_1_8_SCALE_256_7168; + } else if (PerformanceMode::ONE_CORE_GATHER_SORT == GetPerformanceMode() && quantMode_ == UN_QUANT && + rowIdxTytpe_ == SCATTER && expertTokensNumType_ == COUNT) { + uint64_t sortMode = NUM_TWO; + return static_cast(TILINGKEY_BASE + sortMode * SORT_CORE_TILINGKEY_BASE + + static_cast(quantMode_ + 1) * QUANT_MODE_TILINGKEY_BASE + + static_cast(rowIdxTytpe_) * ROWIDX_TYPE_TILINGKEY_BASE + + static_cast(dropPadMode_) * DROP_MODE_TILINGKEY_BASE); + } else if (PerformanceMode::MULTI_CORE_GATHER_SORT == GetPerformanceMode() && quantMode_ == UN_QUANT && + rowIdxTytpe_ == SCATTER && expertTokensNumType_ == COUNT) { + uint64_t sortMode = 3; + return static_cast(TILINGKEY_BASE + sortMode * SORT_CORE_TILINGKEY_BASE + + static_cast(quantMode_ + 1) * QUANT_MODE_TILINGKEY_BASE + + static_cast(rowIdxTytpe_) * ROWIDX_TYPE_TILINGKEY_BASE + + static_cast(dropPadMode_) * DROP_MODE_TILINGKEY_BASE); + } + return static_cast(TILINGKEY_BASE + static_cast(sortMode_) * SORT_CORE_TILINGKEY_BASE + + static_cast(quantMode_ + 1) * QUANT_MODE_TILINGKEY_BASE + + static_cast(rowIdxTytpe_) * ROWIDX_TYPE_TILINGKEY_BASE + + static_cast(dropPadMode_) * DROP_MODE_TILINGKEY_BASE); +} + +ge::graphStatus MoeInitRountingCustomTilingBase::GetWorkspaceSize() +{ + size_t sortWorkspaceSize = + sizeof(float) * static_cast(totalLength_ * NUM_TWO * NUM_THREE); + size_t coreSyncWorkspaceSize = + moeInitRoutingCustomTilingData.get_coreNum() * SORT32_ALIGN_ELEMENT * NUM_TWO; + size_t scatterWorkspaceSize = sizeof(int32_t) * static_cast(totalLength_); + size_t expertIdxValueWorkspaceSize = sizeof(int32_t) * static_cast(aivNum) * 2U; + size_t expertTokensCountWorkspaceSize = sizeof(int32_t) * static_cast((expertEnd_ - expertStart_)); + int64_t expertTokenTotalCountWorkspace = AlignBytes(1, static_cast(sizeof(int32_t))); + int64_t quantTempWorkspaceSize = aivNum * cols_ * static_cast(sizeof(float)); + workspaceSize_ = sortWorkspaceSize + coreSyncWorkspaceSize + scatterWorkspaceSize + expertTokensCountWorkspaceSize + + expertTokenTotalCountWorkspace + SIZE_16 * LENGTH_1024 * LENGTH_1024; + if (quantMode_ == DYNAMIC_QUANT) { + workspaceSize_ += quantTempWorkspaceSize; + } + if (dropPadMode_ == DROP_PAD) { + workspaceSize_ += expertIdxValueWorkspaceSize; + } + OPS_LOG_I(context_->GetNodeName(), "Allocate workspaceSize is: %ld.", workspaceSize_); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus MoeInitRountingCustomTilingBase::PostTiling() +{ + context_->SetBlockDim(aivNum); + size_t *currentWorkspace = context_->GetWorkspaceSizes(1); + currentWorkspace[0] = workspaceSize_; + moeInitRoutingCustomTilingData.SaveToBuffer(context_->GetRawTilingData()->GetData(), + context_->GetRawTilingData()->GetCapacity()); + context_->GetRawTilingData()->SetDataSize(moeInitRoutingCustomTilingData.GetDataSize()); + return ge::GRAPH_SUCCESS; +} +void MoeInitRountingCustomTilingBase::Tinlig4VBSOneCoreCompute(MoeCustomVBSComputeTilingData *tilingData) +{ + tilingData->set_needCoreNum(1); + tilingData->set_perCoreElements(totalLength_); + tilingData->set_perCoreLoops(1); + tilingData->set_perCorePerLoopElements(tilingData->get_perCoreElements()); + tilingData->set_perCoreLastLoopElements(tilingData->get_perCoreElements()); + tilingData->set_lastCoreElements(tilingData->get_perCoreElements()); + tilingData->set_lastCoreLoops(1); + tilingData->set_lastCorePerLoopElements(tilingData->get_perCoreElements()); + tilingData->set_lastCoreLastLoopElements(tilingData->get_perCoreElements()); +} + +void MoeInitRountingCustomTilingBase::Tinlig4VBSMultiCoreCompute(MoeCustomVBSComputeTilingData *tilingData) +{ + int64_t needCoreNum = CeilDiv(totalLength_, sortLoopMaxElement); + needCoreNum = static_cast(std::pow(NUM_FOUR, CeilLog4(needCoreNum))); + needCoreNum = std::min(needCoreNum, aivNum); + + if (needCoreNum == 0) { + OPS_LOG_E(context_->GetNodeName(), "Variate needCoreNum cannot be 0."); + return; + } + int64_t perCoreElements = (needCoreNum == 0) ? 0 : (totalLength_ / needCoreNum); + int64_t alineFloorPerCoreElements = perCoreElements - perCoreElements % SORT32_ALIGN_ELEMENT; + int64_t lastCoreElement = totalLength_ - (needCoreNum - 1) * alineFloorPerCoreElements; + int64_t alineCeilPerCoreElements = perCoreElements + SORT32_ALIGN_ELEMENT - perCoreElements % SORT32_ALIGN_ELEMENT; + if (lastCoreElement > alineCeilPerCoreElements) { + perCoreElements = alineCeilPerCoreElements; + needCoreNum = CeilDiv(totalLength_, perCoreElements); + } else { + perCoreElements = alineFloorPerCoreElements; + } + + tilingData->set_needCoreNum(needCoreNum); + do { + tilingData->set_perCoreElements(perCoreElements); + tilingData->set_perCoreLoops( + CeilDiv(tilingData->get_perCoreElements(), sortLoopMaxElement)); + tilingData->set_perCorePerLoopElements(std::min(tilingData->get_perCoreElements(), sortLoopMaxElement)); + + tilingData->set_perCoreLastLoopElements(tilingData->get_perCoreElements() - + (tilingData->get_perCoreLoops() - 1) * + tilingData->get_perCorePerLoopElements()); + + tilingData->set_lastCoreElements(totalLength_ - + (tilingData->get_needCoreNum() - 1) * tilingData->get_perCoreElements()); + tilingData->set_lastCoreLoops(tilingData->get_perCoreLoops()); + int64_t lastCorePerLoopElements = + CeilDiv(CeilDiv(tilingData->get_lastCoreElements(), tilingData->get_lastCoreLoops()), + SORT32_ALIGN_ELEMENT) * + SORT32_ALIGN_ELEMENT; + tilingData->set_lastCorePerLoopElements(lastCorePerLoopElements); + tilingData->set_lastCoreLastLoopElements(tilingData->get_lastCoreElements() - + (tilingData->get_lastCoreLoops() - 1) * + tilingData->get_lastCorePerLoopElements()); + perCoreElements -= SORT32_ALIGN_ELEMENT; + } while (tilingData->get_lastCoreLastLoopElements() <= 0 && perCoreElements > 0); + if (tilingData->get_lastCoreLastLoopElements() <= 0) { + OPS_LOG_E(context_->GetNodeName(), "vbs tiling failed"); + return; + } +} + +void MoeInitRountingCustomTilingBase::Tiling4VBSCompute() +{ + if (totalLength_ <= sortLoopMaxElement) { + sortMode_ = 0; + } else { + sortMode_ = 1; + } + + auto tilingData = &moeInitRoutingCustomTilingData.vbsComputeParamsOp; + tilingData->set_oneLoopMaxElements(sortLoopMaxElement); + if (sortMode_ == 0UL) { + Tinlig4VBSOneCoreCompute(tilingData); + return; + } + Tinlig4VBSMultiCoreCompute(tilingData); +} + +void MoeInitRountingCustomTilingBase::Tiling4VMSMiddleCompute() +{ + auto vbsComputeTilingData = &moeInitRoutingCustomTilingData.vbsComputeParamsOp; + auto tilingData = &moeInitRoutingCustomTilingData.vmsMiddleComputeParamsOp; + if (vbsComputeTilingData->get_needCoreNum() <= MRG_LIST_NUM) { + tilingData->set_needCoreNum(0); + return; + } + int64_t needCoreNum = CeilDiv(vbsComputeTilingData->get_needCoreNum(), MRG_LIST_NUM); + tilingData->set_needCoreNum(needCoreNum); +} + +void MoeInitRountingCustomTilingBase::Tiling4SortOutCompute() +{ + auto tilingData = &moeInitRoutingCustomTilingData.sortOutComputeParamsOp; + tilingData->set_oneLoopMaxElements(mrgSortListMaxElement); +} + +void MoeInitRountingCustomTilingBase::Tiling4ExpertTokensCountCompute() +{ + auto tilingData = &moeInitRoutingCustomTilingData.expertTokensCountTilingDataOp; + int64_t totalElements = moeInitRoutingCustomTilingData.get_n() * moeInitRoutingCustomTilingData.get_k(); + int64_t perCoreElements = CeilDiv(totalElements, aivNum); + int64_t needCoreNum = CeilDiv(totalElements, perCoreElements); + int64_t lastCoreElements = totalElements - (needCoreNum - 1) * perCoreElements; + tilingData->set_needCoreNum(needCoreNum); + tilingData->set_perCoreElements(perCoreElements); + tilingData->set_lastCoreElements(lastCoreElements); + + int64_t expertNumElement = (moeInitRoutingCustomTilingData.get_expertTokensNumType() != KEY_VALUE) ? + moeInitRoutingCustomTilingData.get_actualExpertNum() : + (moeInitRoutingCustomTilingData.get_actualExpertNum() + 1) * DIM_TWO; + + int64_t maxElementsPerLoop = + (static_cast(aicoreParams_.ubSize) - + CeilAlign(expertNumElement, ONE_BLOCK_BYTE) * + (static_cast(sizeof(int32_t)) * NUM_TWO + static_cast(sizeof(int64_t))) - + ONE_BLOCK_BYTE) / static_cast(sizeof(int32_t)); + int64_t perCoreLoops = CeilDiv(perCoreElements, maxElementsPerLoop); + int64_t perCorePerLoopElements = CeilDiv(perCoreElements, perCoreLoops); + int64_t perCoreLastLoopElements = perCoreElements - (perCoreLoops - 1) * perCorePerLoopElements; + + tilingData->set_perCoreLoops(perCoreLoops); + tilingData->set_perCorePerLoopElements(perCorePerLoopElements); + tilingData->set_perCoreLastLoopElements(perCoreLastLoopElements); + + int64_t lastCoreLoops = CeilDiv(lastCoreElements, maxElementsPerLoop); + int64_t lastCorePerLoopElements = CeilDiv(lastCoreElements, lastCoreLoops); + int64_t lastCoreLastLoopElements = lastCoreElements - (lastCoreLoops - 1) * lastCorePerLoopElements; + + tilingData->set_lastCoreLoops(lastCoreLoops); + tilingData->set_lastCorePerLoopElements(lastCorePerLoopElements); + tilingData->set_lastCoreLastLoopElements(lastCoreLastLoopElements); + + OPS_LOG_I(context_->GetNodeName(), + "ExpertTokensCountCompute Tilingdata, needCoreNum is: %ld, perCoreElements is: %ld, lastCoreElements is: " + "%ld, maxElementsPerLoop is: %ld, perCoreLoops is: %ld, perCorePerLoopElements is: %ld, " + "perCoreLastLoopElements " + "is: %ld, lastCoreLoops is: %ld, lastCorePerLoopElements is: %ld, lastCoreLastLoopElements is: %ld.", + needCoreNum, perCoreElements, lastCoreElements, maxElementsPerLoop, perCoreLoops, perCorePerLoopElements, + perCoreLastLoopElements, lastCoreLoops, lastCorePerLoopElements, lastCoreLastLoopElements); +} + +void MoeInitRountingCustomTilingBase::Tiling4SrcToDstDropPadCompute() +{ + if (quantMode_ == DYNAMIC_QUANT && dropPadMode_ == DROP_PAD) { + MoeInitRountingCustomTilingBase::Tiling4SrcToDstDropPadDynamicCompute(); + return; + } + + auto tilingData = &moeInitRoutingCustomTilingData.srcToDstDropPadParamsOp; + + int64_t perCoreRows = CeilDiv(totalLength_, aivNum); + if (perCoreRows <= 0) { + tilingData->set_needCoreNum(0); + return; + } + int64_t needCoreNum = CeilDiv(totalLength_, perCoreRows); + tilingData->set_needCoreNum(needCoreNum); + int64_t cols = moeInitRoutingCustomTilingData.get_cols(); + tilingData->set_perCoreRows(perCoreRows); + int64_t lastCoreRows = totalLength_ - perCoreRows * (needCoreNum - 1); + tilingData->set_lastCoreRows(lastCoreRows); + bool needScaleCopy = (isInputScale_ != 0 && quantMode_ == -1); + int64_t inuptXDtypeSize = inuptXDtypeSize_ == SIZE_INT8 ? SIZE_INT16 : inuptXDtypeSize_; + + int64_t rowSize = + (perCoreRows * sizeof(int32_t) * NUM_TWO + ONE_BLOCK_BYTE + ONE_BLOCK_BYTE * needScaleCopy + ONE_BLOCK_BYTE - 1) / + ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; + int64_t colSize = (cols * inuptXDtypeSize + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; + + if (rowSize + colSize < static_cast(aicoreParams_.ubSize)) { + SetGatherTilingData(tilingData, perCoreRows, lastCoreRows, cols); + } else { + int64_t baseMaxCols = MAX_COLS_ONE_LOOP; + int64_t baseMaxColsSize = + (baseMaxCols * inuptXDtypeSize + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; + int64_t basePerLoopMaxRows = (static_cast(aicoreParams_.ubSize) - baseMaxColsSize - ONE_BLOCK_BYTE - + ONE_BLOCK_BYTE * needScaleCopy) /static_cast(sizeof(int32_t)) + / NUM_TWO / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; + if (cols < MAX_COLS_ONE_LOOP) { + basePerLoopMaxRows = (static_cast(aicoreParams_.ubSize) - colSize - ONE_BLOCK_BYTE - + ONE_BLOCK_BYTE * needScaleCopy) / static_cast(sizeof(int32_t)) + / NUM_TWO / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE; + } else if (perCoreRows < basePerLoopMaxRows) { + baseMaxCols = (static_cast(aicoreParams_.ubSize) - rowSize) / inuptXDtypeSize / ONE_BLOCK_BYTE * + ONE_BLOCK_BYTE; + } + tilingData->set_perLoopCols(std::min(baseMaxCols, cols)); + tilingData->set_lastLoopCols(GetPerOrLastValue(cols, baseMaxCols)); + tilingData->set_colLoops((cols + baseMaxCols - 1) / baseMaxCols); + + tilingData->set_perCorePerLoopRows(std::min(perCoreRows, basePerLoopMaxRows)); + tilingData->set_perCoreLastLoopRows(GetPerOrLastValue(perCoreRows, basePerLoopMaxRows)); + tilingData->set_perCoreLoops((perCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows); + + tilingData->set_lastCorePerLoopRows(std::min(lastCoreRows, basePerLoopMaxRows)); + tilingData->set_lastCoreLastLoopRows(GetPerOrLastValue(lastCoreRows, basePerLoopMaxRows)); + tilingData->set_lastCoreLoops((lastCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows); + } +} + +void MoeInitRountingCustomTilingBase::SetGatherTilingData( + MoeCustomSrcToDstCapacityComputeTilingData* tilingData, int64_t perCoreRows, int64_t lastCoreRows, int64_t cols) +{ + tilingData->set_perCorePerLoopRows(perCoreRows); + tilingData->set_perCoreLastLoopRows(perCoreRows); + tilingData->set_lastCorePerLoopRows(lastCoreRows); + tilingData->set_lastCoreLastLoopRows(lastCoreRows); + tilingData->set_perCoreLoops(1); + tilingData->set_lastCoreLoops(1); + tilingData->set_perLoopCols(cols); + tilingData->set_lastLoopCols(cols); + tilingData->set_colLoops(1); +} + +void MoeInitRountingCustomTilingBase::SetGatherTilingDataCols( + MoeCustomSrcToDstCapacityComputeTilingData* tilingData, int64_t baseMaxCols, int64_t cols) +{ + tilingData->set_perLoopCols(std::min(baseMaxCols, cols)); + tilingData->set_lastLoopCols(GetPerOrLastValue(cols, baseMaxCols)); + tilingData->set_colLoops(baseMaxCols == 0 ? 0 : (cols + baseMaxCols - 1) / baseMaxCols); +} + +void MoeInitRountingCustomTilingBase::SetGatherTilingDataRows( + MoeCustomSrcToDstCapacityComputeTilingData* tilingData, int64_t perCoreRows, int64_t lastCoreRows, + int64_t basePerLoopMaxRows) +{ + tilingData->set_perCorePerLoopRows(std::min(perCoreRows, basePerLoopMaxRows)); + tilingData->set_perCoreLastLoopRows(GetPerOrLastValue(perCoreRows, basePerLoopMaxRows)); + tilingData->set_perCoreLoops( + basePerLoopMaxRows == 0 ? 0 : (perCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows); + + tilingData->set_lastCorePerLoopRows(std::min(lastCoreRows, basePerLoopMaxRows)); + tilingData->set_lastCoreLastLoopRows(GetPerOrLastValue(lastCoreRows, basePerLoopMaxRows)); + tilingData->set_lastCoreLoops( + basePerLoopMaxRows == 0 ? 0 : (lastCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows); +} + +void MoeInitRountingCustomTilingBase::Tiling4SrcToDstDropPadDynamicCompute() +{ + auto tilingData = &moeInitRoutingCustomTilingData.srcToDstDropPadDynamicParamsOp; + + int64_t perCoreRows = CeilDiv(totalLength_, aivNum); + if (perCoreRows <= 0) { + tilingData->set_needCoreNum(0); + return; + } + tilingData->set_needCoreNum(CeilDiv(totalLength_, perCoreRows)); + int64_t cols = moeInitRoutingCustomTilingData.get_cols(); + tilingData->set_perCoreRows(perCoreRows); + int64_t lastCoreRows = totalLength_ - perCoreRows * (tilingData->get_needCoreNum() - 1); + tilingData->set_lastCoreRows(lastCoreRows); + + int64_t rowSize = AlignBytes(perCoreRows, static_cast(sizeof(int32_t))) * NUM_FOUR; + int64_t colSize = AlignBytes(cols, static_cast(sizeof(int8_t))) * DYNAMIC_QUANT_SRC_TO_DST_BUFFER; + int64_t scaleSize = DYNAMIC_QUANT_SCALE_SIZE_64; + if (rowSize + colSize + scaleSize < static_cast(aicoreParams_.ubSize)) { + SetGatherTilingData(tilingData, perCoreRows, lastCoreRows, cols); + } else { + int64_t baseMaxCols = MAX_COLS_DYNAMIC_QUANT; + int64_t totalColSize = AlignBytes(baseMaxCols, static_cast(sizeof(int8_t))) * DYNAMIC_QUANT_SRC_TO_DST_BUFFER; + int64_t ubSize = static_cast(aicoreParams_.ubSize); + int64_t basePerLoopMaxRows = AlignOneBlockByteCeil((ubSize - totalColSize - scaleSize) / SIZE_INT32) / NUM_FOUR; + if (cols < MAX_COLS_DYNAMIC_QUANT) { + basePerLoopMaxRows = AlignOneBlockByteCeil((ubSize - colSize - scaleSize) / SIZE_INT32) / NUM_FOUR; + } else if (perCoreRows < basePerLoopMaxRows) { + baseMaxCols = AlignOneBlockByteCeil(ubSize - rowSize - scaleSize) / DYNAMIC_QUANT_SRC_TO_DST_BUFFER; + } + SetGatherTilingDataCols(tilingData, baseMaxCols, cols); + SetGatherTilingDataRows(tilingData, perCoreRows, lastCoreRows, basePerLoopMaxRows); + } +} + +void MoeInitRountingCustomTilingBase::Tiling4SrcToDstCompute() +{ + auto tilingData = &moeInitRoutingCustomTilingData.srcToDstComputeParamsOp; + + int64_t useCore = aivNum; + int64_t remainUbSize = aicoreParams_.ubSize - ASSIST_NUM * sizeof(int32_t) - ONE_BLOCK_BYTE * (ASSIST_NUM + 1); + int64_t perLoopMaxElements = remainUbSize / (ONE_BLOCK_BYTE + SIZE_INT32); + int64_t perCoreElements = CeilDiv(totalLength_, useCore); + if (perCoreElements <= 0) { + tilingData->set_needCoreNum(0); + return; + } + int64_t needCoreNum = CeilDiv(totalLength_, perCoreElements); + tilingData->set_needCoreNum(needCoreNum); + int64_t lastCoreElements = totalLength_ - perCoreElements * (needCoreNum - 1); + + tilingData->set_perCoreElements(perCoreElements); + tilingData->set_lastCoreElements(lastCoreElements); + int64_t perCoreLoops = CeilDiv(perCoreElements, perLoopMaxElements); + int64_t perCorePerLoopElements = CeilDiv(perCoreElements, perCoreLoops); + int64_t perCoreLastLoopElements = perCoreElements - (perCoreLoops - 1) * perCorePerLoopElements; + + int64_t lastCoreLoops = CeilDiv(lastCoreElements, perLoopMaxElements); + int64_t lastCorePerLoopElements = CeilDiv(lastCoreElements, lastCoreLoops); + int64_t lastCoreLastLoopElements = lastCoreElements - (lastCoreLoops - 1) * lastCorePerLoopElements; + + tilingData->set_perCoreLoops(perCoreLoops); + tilingData->set_perCorePerLoopElements(perCorePerLoopElements); + tilingData->set_perCoreLastLoopElements(perCoreLastLoopElements); + tilingData->set_lastCoreLoops(lastCoreLoops); + tilingData->set_lastCorePerLoopElements(lastCorePerLoopElements); + tilingData->set_lastCoreLastLoopElements(lastCoreLastLoopElements); +} + +void MoeInitRountingCustomTilingBase::Tiling4GatherOutCompute() +{ + auto tilingData = &moeInitRoutingCustomTilingData.gatherOutComputeParamsOp; + int64_t perCoreIndicesElements = CeilDiv(totalLength_, aivNum); + if (perCoreIndicesElements <= 0) { + tilingData->set_needCoreNum(0); + return; + } + int64_t needCoreNum = CeilDiv(totalLength_, perCoreIndicesElements); + int64_t lastCoreIndicesElements = totalLength_ - (needCoreNum - 1) * perCoreIndicesElements; + + int64_t perLoopCols = moeInitRoutingCustomTilingData.get_cols(); + int64_t colMultiple = NUM_TWO * inuptXDtypeSize_; + int64_t rowMultiple = NUM_TWO; + if (quantMode_ == DYNAMIC_QUANT) { + colMultiple = DYNAMIC_QUANT_COLS_BUFFER; + rowMultiple = NUM_FOUR; + } + if (quantMode_ == STATIC_QUANT) { + colMultiple = SIZE_INT8 * NUM_TWO + SIZE_FP32 + SIZE_INT16 + inuptXDtypeSize_ * NUM_TWO; + rowMultiple = NUM_TWO; + } + int64_t perLoopMaxIndicesElements = + (static_cast(aicoreParams_.ubSize) - Align(perLoopCols, inuptXDtypeSize_) * colMultiple - + ONE_BLOCK_BYTE * NUM_TWO) / + rowMultiple / static_cast(sizeof(int32_t)); + while (perLoopMaxIndicesElements <= 0) { + perLoopCols = CeilDiv(perLoopCols, NUM_TWO); + perLoopMaxIndicesElements = (static_cast(aicoreParams_.ubSize) - + Align(perLoopCols, inuptXDtypeSize_) * colMultiple - ONE_BLOCK_BYTE * NUM_TWO) / + rowMultiple / static_cast(sizeof(int32_t)); + OPS_LOG_I(context_->GetNodeName(), "perLoopCols is: %ld, perLoopMaxIndicesElements is: %ld", perLoopCols, + perLoopMaxIndicesElements); + } + int64_t colsLoops = CeilDiv(moeInitRoutingCustomTilingData.get_cols(), perLoopCols); + int64_t lastLoopCols = moeInitRoutingCustomTilingData.get_cols() - (colsLoops - 1) * perLoopCols; + tilingData->set_needCoreNum(needCoreNum); + tilingData->set_perCoreIndicesElements(perCoreIndicesElements); + tilingData->set_lastCoreIndicesElements(lastCoreIndicesElements); + tilingData->set_colsLoops(colsLoops); + tilingData->set_perLoopCols(perLoopCols); + tilingData->set_lastLoopCols(lastLoopCols); + + int64_t perCorePerLoopIndicesElements = std::min(perLoopMaxIndicesElements, perCoreIndicesElements); + int64_t perCoreIndicesLoops = CeilDiv(perCoreIndicesElements, perCorePerLoopIndicesElements); + int64_t perCoreLastLoopIndicesElements = + perCoreIndicesElements - (perCoreIndicesLoops - 1) * perCorePerLoopIndicesElements; + tilingData->set_perCoreIndicesLoops(perCoreIndicesLoops); + tilingData->set_perCorePerLoopIndicesElements(perCorePerLoopIndicesElements); + tilingData->set_perCoreLastLoopIndicesElements(perCoreLastLoopIndicesElements); + + int64_t lastCorePerLoopIndicesElements = std::min(perLoopMaxIndicesElements, lastCoreIndicesElements); + int64_t lastCoreIndicesLoops = CeilDiv(lastCoreIndicesElements, lastCorePerLoopIndicesElements); + int64_t lastCoreLastLoopIndicesElements = + lastCoreIndicesElements - (lastCoreIndicesLoops - 1) * lastCorePerLoopIndicesElements; + tilingData->set_lastCoreIndicesLoops(lastCoreIndicesLoops); + tilingData->set_lastCorePerLoopIndicesElements(lastCorePerLoopIndicesElements); + tilingData->set_lastCoreLastLoopIndicesElements(lastCoreLastLoopIndicesElements); + + OPS_LOG_I( + context_->GetNodeName(), + "GatherOut Tilingdata, needCoreNum is: %ld, perCoreIndicesElements is: %ld, lastCoreIndicesElements is: %ld, " + "colsLoops is: %ld, perLoopCols is: %ld, lastLoopCols is: %ld, perCoreIndicesLoops is: %ld, " + "perCorePerLoopIndicesElements is: %ld, perCoreLastLoopIndicesElements is: %ld, lastCoreIndicesLoops is: " + "%ld, lastCorePerLoopIndicesElements is: " + "%ld, lastCoreLastLoopIndicesElements is: %ld.", + needCoreNum, perCoreIndicesElements, lastCoreIndicesElements, colsLoops, perLoopCols, lastLoopCols, + perCoreIndicesLoops, perCorePerLoopIndicesElements, perCoreLastLoopIndicesElements, lastCoreIndicesLoops, + lastCorePerLoopIndicesElements, lastCoreLastLoopIndicesElements); +} + +REGISTER_TILING_TEMPLATE("MoeInitRoutingCustom", MoeInitRountingCustomTilingBase, 10000); // Fallback template for non-Ascend950 SoC. +} // namespace optiling diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_tiling.h b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_tiling.h new file mode 100644 index 000000000..64a72d947 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_tiling.h @@ -0,0 +1,143 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_init_routing_custom_tiling.h + * \brief + */ +#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_MOE_INIT_ROUTING_CUSTOM_H +#define AIR_CXX_RUNTIME_V2_OP_IMPL_MOE_INIT_ROUTING_CUSTOM_H +#include "register/tilingdata_base.h" +#include "tiling/tiling_api.h" + + +namespace optiling { +BEGIN_TILING_DATA_DEF(MoeCustomVBSComputeTilingData) +TILING_DATA_FIELD_DEF(int64_t, needCoreNum); +TILING_DATA_FIELD_DEF(int64_t, perCoreElements); +TILING_DATA_FIELD_DEF(int64_t, perCoreLoops); +TILING_DATA_FIELD_DEF(int64_t, perCorePerLoopElements); +TILING_DATA_FIELD_DEF(int64_t, perCoreLastLoopElements); +TILING_DATA_FIELD_DEF(int64_t, lastCoreElements); +TILING_DATA_FIELD_DEF(int64_t, lastCoreLoops); +TILING_DATA_FIELD_DEF(int64_t, lastCorePerLoopElements); +TILING_DATA_FIELD_DEF(int64_t, lastCoreLastLoopElements); +TILING_DATA_FIELD_DEF(int64_t, oneLoopMaxElements); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(MoeCustomVBSComputeTilingDataOp, MoeCustomVBSComputeTilingData) + +BEGIN_TILING_DATA_DEF(MoeCustomVMSMiddleComputeTilingData) +TILING_DATA_FIELD_DEF(int64_t, needCoreNum); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(MoeCustomVMSMiddleComputeTilingDataOp, MoeCustomVMSMiddleComputeTilingData) + +BEGIN_TILING_DATA_DEF(MoeCustomSortOutComputeTilingData) +TILING_DATA_FIELD_DEF(int64_t, oneLoopMaxElements); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(MoeCustomSortOutComputeTilingDataOp, MoeCustomSortOutComputeTilingData) + +BEGIN_TILING_DATA_DEF(MoeCustomExpertTokensCountTilingData) +TILING_DATA_FIELD_DEF(int64_t, needCoreNum); +TILING_DATA_FIELD_DEF(int64_t, perCoreElements); +TILING_DATA_FIELD_DEF(int64_t, lastCoreElements); +TILING_DATA_FIELD_DEF(int64_t, perCoreLoops); +TILING_DATA_FIELD_DEF(int64_t, perCorePerLoopElements); +TILING_DATA_FIELD_DEF(int64_t, perCoreLastLoopElements); +TILING_DATA_FIELD_DEF(int64_t, lastCoreLoops); +TILING_DATA_FIELD_DEF(int64_t, lastCorePerLoopElements); +TILING_DATA_FIELD_DEF(int64_t, lastCoreLastLoopElements); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(MoeCustomExpertTokensCountTilingDataOp, MoeCustomExpertTokensCountTilingData) + +BEGIN_TILING_DATA_DEF(MoeCustomGatherOutComputeTilingData) +TILING_DATA_FIELD_DEF(int64_t, needCoreNum); +TILING_DATA_FIELD_DEF(int64_t, perCoreIndicesElements); +TILING_DATA_FIELD_DEF(int64_t, lastCoreIndicesElements); +TILING_DATA_FIELD_DEF(int64_t, perCoreIndicesLoops); +TILING_DATA_FIELD_DEF(int64_t, perCorePerLoopIndicesElements); +TILING_DATA_FIELD_DEF(int64_t, perCoreLastLoopIndicesElements); +TILING_DATA_FIELD_DEF(int64_t, lastCoreIndicesLoops); +TILING_DATA_FIELD_DEF(int64_t, lastCorePerLoopIndicesElements); +TILING_DATA_FIELD_DEF(int64_t, lastCoreLastLoopIndicesElements); +TILING_DATA_FIELD_DEF(int64_t, colsLoops); +TILING_DATA_FIELD_DEF(int64_t, perLoopCols); +TILING_DATA_FIELD_DEF(int64_t, lastLoopCols); +TILING_DATA_FIELD_DEF(int64_t, activeNum); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(MoeCustomGatherOutComputeTilingDataOp, MoeCustomGatherOutComputeTilingData) + +BEGIN_TILING_DATA_DEF(MoeCustomSrcToDstCapacityComputeTilingData) +TILING_DATA_FIELD_DEF(int64_t, needCoreNum); +TILING_DATA_FIELD_DEF(int64_t, perCoreRows); +TILING_DATA_FIELD_DEF(int64_t, perCorePerLoopRows); +TILING_DATA_FIELD_DEF(int64_t, perCoreLastLoopRows); +TILING_DATA_FIELD_DEF(int64_t, lastCoreRows); +TILING_DATA_FIELD_DEF(int64_t, lastCorePerLoopRows); +TILING_DATA_FIELD_DEF(int64_t, lastCoreLastLoopRows); +TILING_DATA_FIELD_DEF(int64_t, perCoreLoops); +TILING_DATA_FIELD_DEF(int64_t, lastCoreLoops); +TILING_DATA_FIELD_DEF(int64_t, perLoopCols); +TILING_DATA_FIELD_DEF(int64_t, lastLoopCols); +TILING_DATA_FIELD_DEF(int64_t, colLoops); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(MoeCustomSrcToDstCapacityComputeTilingDataOp, MoeCustomSrcToDstCapacityComputeTilingData) + +BEGIN_TILING_DATA_DEF(MoeCustomSrcToDstComputeTilingData) +TILING_DATA_FIELD_DEF(int64_t, needCoreNum); +TILING_DATA_FIELD_DEF(int64_t, perCoreElements); +TILING_DATA_FIELD_DEF(int64_t, perCorePerLoopElements); +TILING_DATA_FIELD_DEF(int64_t, perCoreLastLoopElements); +TILING_DATA_FIELD_DEF(int64_t, lastCoreElements); +TILING_DATA_FIELD_DEF(int64_t, lastCorePerLoopElements); +TILING_DATA_FIELD_DEF(int64_t, lastCoreLastLoopElements); +TILING_DATA_FIELD_DEF(int64_t, perCoreLoops); +TILING_DATA_FIELD_DEF(int64_t, lastCoreLoops) +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(MoeCustomSrcToDstComputeTilingDataOp, MoeCustomSrcToDstComputeTilingData) + +BEGIN_TILING_DATA_DEF(MoeInitRoutingCustomTilingData) +TILING_DATA_FIELD_DEF(int64_t, coreNum); +TILING_DATA_FIELD_DEF(int64_t, n); +TILING_DATA_FIELD_DEF(int64_t, cols); +TILING_DATA_FIELD_DEF(int64_t, k); +TILING_DATA_FIELD_DEF(int64_t, expertStart); +TILING_DATA_FIELD_DEF(int64_t, expertEnd); +TILING_DATA_FIELD_DEF(int64_t, actualExpertNum); +TILING_DATA_FIELD_DEF(int64_t, quantMode); +TILING_DATA_FIELD_DEF(int64_t, rowIdxType); +TILING_DATA_FIELD_DEF(int64_t, isInputScale); +TILING_DATA_FIELD_DEF(int64_t, isInputOffset); +TILING_DATA_FIELD_DEF(int64_t, expertNum); +TILING_DATA_FIELD_DEF(int64_t, expertTokensNumType); +TILING_DATA_FIELD_DEF(int64_t, expertTokensNumFlag); +TILING_DATA_FIELD_DEF(int64_t, gatherFirstFullload); +TILING_DATA_FIELD_DEF(int64_t, ep); +TILING_DATA_FIELD_DEF(int64_t, activeNum); +TILING_DATA_FIELD_DEF(int64_t, dropPadMode); +TILING_DATA_FIELD_DEF(int64_t, smoothType); +TILING_DATA_FIELD_DEF(int64_t, expertCountElements); +TILING_DATA_FIELD_DEF(int64_t, expertCapacity); +TILING_DATA_FIELD_DEF_STRUCT(MoeCustomVBSComputeTilingData, vbsComputeParamsOp); +TILING_DATA_FIELD_DEF_STRUCT(MoeCustomVMSMiddleComputeTilingData, vmsMiddleComputeParamsOp); +TILING_DATA_FIELD_DEF_STRUCT(MoeCustomSortOutComputeTilingData, sortOutComputeParamsOp); +TILING_DATA_FIELD_DEF_STRUCT(MoeCustomExpertTokensCountTilingData, expertTokensCountTilingDataOp); +TILING_DATA_FIELD_DEF_STRUCT(MoeCustomGatherOutComputeTilingData, gatherOutComputeParamsOp); +TILING_DATA_FIELD_DEF_STRUCT(MoeCustomSrcToDstCapacityComputeTilingData, srcToDstDropPadParamsOp); +TILING_DATA_FIELD_DEF_STRUCT(MoeCustomSrcToDstCapacityComputeTilingData, srcToDstDropPadDynamicParamsOp); +TILING_DATA_FIELD_DEF_STRUCT(MoeCustomSrcToDstComputeTilingData, srcToDstComputeParamsOp); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(MoeInitRoutingCustom, MoeInitRoutingCustomTilingData) +struct MoeInitRoutingCustomCompileInfo { + int32_t aivNum = 0; + uint64_t ubSize = 0; + platform_ascendc::SocVersion socVersion = platform_ascendc::SocVersion::ASCEND910B; + }; +} // namespace optiling +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_tiling_base.cpp b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_tiling_base.cpp new file mode 100644 index 000000000..65e3442ae --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_host/moe_init_routing_custom_tiling_base.cpp @@ -0,0 +1,68 @@ +/** + * This program is free software, you can redistribute it and/or modify. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_init_routing_custom_tiling_base.cpp + * \brief + */ +#include "moe_init_routing_custom_tiling.h" +#include "register/op_def_registry.h" +#include "tiling/tiling_templates_registry.h" + +#define unlikely(x) __builtin_expect((x), 0) + +#define OP_CHECK_NULL_WITH_CONTEXT(context, ptr) \ + do { \ + if (unlikely((ptr) == nullptr)) { \ + const char* name = (unlikely(((context) == nullptr) || (context)->GetNodeName() == nullptr)) ? \ + "nil" : \ + (context)->GetNodeName(); \ + OPS_LOG_E(name, "%s is nullptr!", #ptr); \ + return ge::GRAPH_FAILED; \ + } \ + } while (0) + +namespace optiling { +static ge::graphStatus TilingForMoeInitRoutingCustom(gert::TilingContext *context) +{ + return TilingRegistry::GetInstance().DoTilingImpl(context); +} + +static ge::graphStatus TilingPrepareForMoeInitRountingCustom(gert::TilingParseContext* context) +{ + OPS_LOG_D(context, "TilingPrepareForMoeInitRountingCustom enter."); + + auto compileInfo = context->GetCompiledInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo); + auto platformInfo = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + compileInfo->aivNum = ascendcPlatform.GetCoreNumAiv(); + if (compileInfo->aivNum <= 0) { + OPS_LOG_E(context, "TilingPrepareForMoeInitRountingCustom fail to get core num."); + return ge::GRAPH_FAILED; + } + + uint64_t ubSize; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); + compileInfo->ubSize = static_cast(ubSize); + compileInfo->socVersion = ascendcPlatform.GetSocVersion(); + if (compileInfo->ubSize <= 0) { + OPS_LOG_E(context, "TilingPrepareForMoeInitRountingCustom fail to get ub size."); + return ge::GRAPH_FAILED; + } + + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(MoeInitRoutingCustom) + .Tiling(TilingForMoeInitRoutingCustom) + .TilingParse(TilingPrepareForMoeInitRountingCustom); +} // namespace optiling \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_common.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_common.h new file mode 100644 index 000000000..5afeebaf9 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_common.h @@ -0,0 +1,110 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_common.h + * \brief + */ +#ifndef MOE_CUSTOM_COMMON_H +#define MOE_CUSTOM_COMMON_H + +#include "kernel_operator.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; +constexpr int64_t SPLIT_N = 0; +constexpr int64_t SPLIT_K = 1; +constexpr float MIN_FP32 = -3.4e38f; +constexpr int64_t FP32_ONE_REPEAT_NUM = 64; +constexpr int64_t ONE_REPEAT_SORT_NUM = 32; +constexpr int64_t ONE_REPEAT_COMPARE_NUM = 64; +constexpr int64_t BLOCK_BYTES = 32; +constexpr int64_t INT32_ONE_BLOCK_NUM = 8; +constexpr int64_t FP32_ONE_BLOCK_NUM = 8; +constexpr int64_t DROPLESS_MODE = 0; +constexpr int64_t DROP_PAD_MODE = 1; +constexpr int64_t ASSIST_NUM = 256; +constexpr int64_t ASSIST_INDEX_NUM = 32; +constexpr int64_t MRGSORT_LIST_MAX_ELEMENT = 2040; +constexpr float MAX_INT8 = 127.0f; +constexpr uint32_t INF = 0xFF7FFFFF; + +constexpr int64_t MERGE_LIST_TWO = 2; +constexpr int64_t MERGE_LIST_THREE = 3; +constexpr int64_t MERGE_LIST_FOUR = 4; + +constexpr int64_t MERGE_LIST_IDX_TWO = 2; +constexpr int64_t MERGE_LIST_IDX_THREE = 3; + +constexpr int64_t GATHER = 0; +constexpr int64_t SCATTER = 1; + +static constexpr int64_t NO_SCALE = 0; +static constexpr int64_t SCALE_1H = 1; +static constexpr int64_t SCALE_EH = 2; + +constexpr int64_t EXERPT_TOKENS_CUMSUM = 0; +constexpr int64_t EXERPT_TOKENS_COUNT = 1; +constexpr int64_t EXERPT_TOKENS_KEY_VALUE = 2; +constexpr int64_t EXERPT_TOKENS_NONE = 0; + +const __gm__ int32_t assist[256] = { + 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, + 4, 0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, + 8, 0, 0, 0, 0, 0, 0, 0, 9, 0, 0, 0, 0, 0, 0, 0, 10, 0, 0, 0, 0, 0, 0, 0, 11, 0, 0, 0, 0, 0, 0, 0, + 12, 0, 0, 0, 0, 0, 0, 0, 13, 0, 0, 0, 0, 0, 0, 0, 14, 0, 0, 0, 0, 0, 0, 0, 15, 0, 0, 0, 0, 0, 0, 0, + 16, 0, 0, 0, 0, 0, 0, 0, 17, 0, 0, 0, 0, 0, 0, 0, 18, 0, 0, 0, 0, 0, 0, 0, 19, 0, 0, 0, 0, 0, 0, 0, + 20, 0, 0, 0, 0, 0, 0, 0, 21, 0, 0, 0, 0, 0, 0, 0, 22, 0, 0, 0, 0, 0, 0, 0, 23, 0, 0, 0, 0, 0, 0, 0, + 24, 0, 0, 0, 0, 0, 0, 0, 25, 0, 0, 0, 0, 0, 0, 0, 26, 0, 0, 0, 0, 0, 0, 0, 27, 0, 0, 0, 0, 0, 0, 0, + 28, 0, 0, 0, 0, 0, 0, 0, 29, 0, 0, 0, 0, 0, 0, 0, 30, 0, 0, 0, 0, 0, 0, 0, 31, 0, 0, 0, 0, 0, 0, 0}; + +__aicore__ inline int64_t Ceil(int64_t a, int64_t b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes) +{ + if (bytes == 0) { + return 0; + } + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes; +} + +__aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes) +{ + return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES; +} + +template +__aicore__ inline T Min(T a, T b) +{ + return a > b ? b : a; +} + +template +__aicore__ inline T Max(T a, T b) +{ + return a < b ? b : a; +} + +template +__aicore__ inline void SetWaitFlag(HardEvent evt) +{ + event_t eventId = static_cast(GetTPipePtr()->FetchEventID(evt)); + SetFlag(eventId); + WaitFlag(eventId); +} + +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_COMMON_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_expert_tokens_count.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_expert_tokens_count.h new file mode 100644 index 000000000..c4fb6ba2f --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_expert_tokens_count.h @@ -0,0 +1,371 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_expert_tokens_count.h + * \brief + */ +#ifndef MOE_CUSTOM_EXPERT_TOKENS_COUNT_H +#define MOE_CUSTOM_EXPERT_TOKENS_COUNT_H + +#include "moe_custom_common.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +constexpr int64_t EXPERT_ID_VALUE_NUM = 2; +constexpr int64_t CUMSUM_MODE = 0; +constexpr int64_t COUNT_MODE = 1; +constexpr int64_t KEY_VALUE_MODE = 2; +constexpr int64_t KEY_VALUE_MODE_DIM_NUM = 2; +constexpr int64_t GATHER_SORT_CORE_NUM = 16; +constexpr int64_t DROP_LESS = 0; +constexpr int64_t DROP_PAD = 1; + +template +class ExpertTokensCount { +public: + __aicore__ inline ExpertTokensCount(){}; + template + __aicore__ inline void Init(GM_ADDR expandedRowIdx, GM_ADDR expertTokensCount, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyIn(int64_t loop, int64_t curLoopElements); + __aicore__ inline void Compute(int64_t curLoopElements); + __aicore__ inline void CopyOut(); + __aicore__ inline void CopyOutExpertTotalCount(); + + __aicore__ inline void expertCountCopyIn(); + __aicore__ inline void expertCountCompute(); + __aicore__ inline void expertCountCopyOut(); + +private: + GlobalTensor sortedexpertIdxGm_; + GlobalTensor expertCountTempGm_; + GlobalTensor expertTokensCountGm_; + GlobalTensor expertTotalCountGm_; + GlobalTensor expandedRowIdxGm_; + GlobalTensor expertIdxValueGm_; + TPipe *pipe_; + + TQue sortedExpertIdxInQueue_; + TQue expertCountOutToTempQueue_; + TQue expertCountTempInQueue_; + TQue expertIdxCountOutQueue_; + TQue expertTotalCountQueue_; + + const MoeCustomExpertTokensCountTilingData *expertTokensCountTilingData_; + int64_t coreNum_; + int64_t blockIdx_; + int64_t needCoreNum_; + int64_t perCoreElements_; + int64_t curCoreElements_ = 0; + int64_t expertStart_ = 0; + int64_t expertEnd_ = 0; + int64_t actualExpertNum_ = 0; + int64_t coreLoopsNum_ = 0; + int64_t perCorePerLoopElements_ = 0; + int64_t perCoreLastLoopElements_ = 0; + int64_t actualExpertTotalNum_ = 0; + int64_t expertNum_ = 0; + int64_t expertCountElements_ = 0; + bool expertTokensNumFlag_ = false; + int64_t dropPadMode_ = 0; + int32_t finalExpertId = -1; + int32_t expertTokenValue = 0; + int64_t ep_ = 0; + int64_t rowIdxType_ = 0; +}; + +template +template +__aicore__ inline void +ExpertTokensCount::Init(GM_ADDR expandedRowIdx, GM_ADDR expertTokensCount, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + coreNum_ = tilingData->coreNum; + pipe_ = tPipe; + expertTokensCountTilingData_ = &(tilingData->expertTokensCountTilingDataOp); + blockIdx_ = GetBlockIdx(); + needCoreNum_ = expertTokensCountTilingData_->needCoreNum; + perCoreElements_ = expertTokensCountTilingData_->perCoreElements; + expertStart_ = tilingData->expertStart; + expertEnd_ = tilingData->expertEnd; + actualExpertNum_ = tilingData->actualExpertNum; + expertNum_ = tilingData->expertNum; + expertTokensNumFlag_ = tilingData->expertTokensNumFlag; + dropPadMode_ = tilingData->dropPadMode; + ep_ = tilingData->ep; + rowIdxType_ = tilingData->rowIdxType; + + if (blockIdx_ == needCoreNum_ - 1) { + curCoreElements_ = expertTokensCountTilingData_->lastCoreElements; + coreLoopsNum_ = expertTokensCountTilingData_->lastCoreLoops; + perCorePerLoopElements_ = expertTokensCountTilingData_->lastCorePerLoopElements; + perCoreLastLoopElements_ = expertTokensCountTilingData_->lastCoreLastLoopElements; + } else { + curCoreElements_ = expertTokensCountTilingData_->perCoreElements; + coreLoopsNum_ = expertTokensCountTilingData_->perCoreLoops; + perCorePerLoopElements_ = expertTokensCountTilingData_->perCorePerLoopElements; + perCoreLastLoopElements_ = expertTokensCountTilingData_->perCoreLastLoopElements; + } + + if (CALC_ACTUAL_EXPERT_NUM) { + // key and value + int64_t kvFactor = 2; + GlobalTensor sortedNumGm; + sortedNumGm.SetGlobalBuffer((__gm__ int32_t *)workspace + + Align(tilingData->n * tilingData->k, sizeof(int32_t)) * kvFactor * kvFactor); + int32_t totalSortedNum = 0; + for (int32_t i = 0; i < 16; i++) { + totalSortedNum += sortedNumGm.GetValue(i); + } + perCoreElements_ = Ceil(totalSortedNum, GetBlockNum()); + needCoreNum_ = Ceil(totalSortedNum, perCoreElements_); + int64_t lastCoreElements = totalSortedNum - (needCoreNum_ - 1) * perCoreElements_; + if (blockIdx_ == needCoreNum_ - 1) { + curCoreElements_ = lastCoreElements; + } else { + curCoreElements_ = perCoreElements_; + } + coreLoopsNum_ = Ceil(curCoreElements_, expertTokensCountTilingData_->perCorePerLoopElements); + perCorePerLoopElements_ = Ceil(curCoreElements_, coreLoopsNum_); + perCoreLastLoopElements_ = curCoreElements_ - (coreLoopsNum_ - 1) * perCorePerLoopElements_; + } + + if constexpr (HISTOGRAMTYPE == KEY_VALUE_MODE) { + expertCountElements_ = ((actualExpertNum_ + 1) < expertNum_) ? (actualExpertNum_ + 1) * KEY_VALUE_MODE_DIM_NUM : + expertNum_ * KEY_VALUE_MODE_DIM_NUM; + } else { + expertCountElements_ = actualExpertNum_; + } + sortedexpertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + blockIdx_ * perCoreElements_, curCoreElements_); + expertTokensCountGm_.SetGlobalBuffer((__gm__ int64_t *)expertTokensCount, expertCountElements_); + expertCountTempGm_.SetGlobalBuffer( + (__gm__ int32_t *)workspace + Align(tilingData->n * tilingData->k, sizeof(int32_t)) * 2, actualExpertNum_); + expertTotalCountGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + + Align(tilingData->n * tilingData->k, sizeof(int32_t)) * 2 + + Align(actualExpertNum_, sizeof(int32_t)), + actualExpertNum_); + expertIdxValueGm_.SetGlobalBuffer( + (__gm__ int32_t *)workspace + Align(tilingData->n * tilingData->k, sizeof(int32_t)) * 2 + + Align((actualExpertNum_), sizeof(int32_t)) + Align((actualExpertNum_), sizeof(int32_t)), + coreNum_ * 2); + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx + blockIdx_ * perCoreElements_, + curCoreElements_); + + if ((tilingData->rowIdxType == GATHER) && (blockIdx_ < needCoreNum_)) { + InitGlobalMemory(expandedRowIdxGm_, curCoreElements_, -1); + SetWaitFlag(HardEvent::MTE3_MTE2); + } + int64_t sortedExpertIdxInLen = Max(perCorePerLoopElements_, perCoreLastLoopElements_); + + pipe_->InitBuffer(sortedExpertIdxInQueue_, 1, AlignBytes(sortedExpertIdxInLen, sizeof(int32_t))); + pipe_->InitBuffer(expertCountOutToTempQueue_, 1, AlignBytes(actualExpertNum_, sizeof(int32_t))); + pipe_->InitBuffer(expertCountTempInQueue_, 1, AlignBytes(actualExpertNum_, sizeof(int32_t))); + + pipe_->InitBuffer(expertIdxCountOutQueue_, 1, AlignBytes(expertCountElements_, sizeof(int64_t))); + pipe_->InitBuffer(expertTotalCountQueue_, 1, AlignBytes(1, sizeof(int32_t))); + + if (blockIdx_ == 0) { + InitGlobalMemory(expertTotalCountGm_, 1, 0); + SetWaitFlag(HardEvent::MTE3_MTE2); + } + SyncAll(); +} + +template +__aicore__ inline void ExpertTokensCount::Process() +{ + if (blockIdx_ < needCoreNum_) { + for (int64_t i = 0; i < coreLoopsNum_; i++) { + int64_t perLoopElements = (i == (coreLoopsNum_ - 1)) ? perCoreLastLoopElements_ : perCorePerLoopElements_; + CopyIn(i, perLoopElements); + Compute(perLoopElements); + CopyOut(); + } + if (ep_ == 1) { + CopyOutExpertTotalCount(); + } + } + if (ep_ == 1 || expertTokensNumFlag_ || dropPadMode_ == 1) { + SyncAll(); + } + /* copy expert tokens count result from worksapce to output GM. */ + if (blockIdx_ == 0 && expertTokensNumFlag_) { + expertCountCopyIn(); + expertCountCompute(); + expertCountCopyOut(); + } +} + +template +__aicore__ inline void ExpertTokensCount::CopyIn(int64_t loop, int64_t curLoopElements) +{ + LocalTensor sortedExpertIdxInLocal = sortedExpertIdxInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{static_cast(1), static_cast(curLoopElements * sizeof(int32_t)), + 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + int64_t sortedexpertIdxOffset = loop * perCorePerLoopElements_; + DataCopyPad(sortedExpertIdxInLocal, sortedexpertIdxGm_[sortedexpertIdxOffset], dataCopyParams, dataCopyPadParams); + sortedExpertIdxInQueue_.EnQue(sortedExpertIdxInLocal); +} + +template +__aicore__ inline void ExpertTokensCount::Compute(int64_t curLoopElements) +{ + LocalTensor sortedExpertIdxInLocal = sortedExpertIdxInQueue_.DeQue(); + LocalTensor expertCountOutLocal = expertCountOutToTempQueue_.AllocTensor(); + Duplicate(expertCountOutLocal.ReinterpretCast(), static_cast(0), + static_cast(actualExpertNum_)); + SetWaitFlag(HardEvent::V_S); + int64_t i = 0; + int32_t lastExpertId = sortedExpertIdxInLocal.GetValue(0); + int32_t lastIndex = 0; + int64_t loopTokenCount = 0; + int32_t lastlastExpertId = lastExpertId; + for (i = 1; i < curLoopElements; i++) { + if ((lastExpertId >= expertEnd_) || (lastExpertId < expertStart_)) { + break; + } + int32_t curExpertId = sortedExpertIdxInLocal.GetValue(i); + if (curExpertId != lastExpertId || curExpertId >= expertEnd_) { + if constexpr (HISTOGRAMTYPE == COUNT_MODE || HISTOGRAMTYPE == KEY_VALUE_MODE) { + expertCountOutLocal.SetValue(lastExpertId - expertStart_, i - lastIndex); + loopTokenCount += i - lastIndex; + } else { + for (int64_t j = lastlastExpertId; j < lastExpertId; j++) { + expertCountOutLocal.SetValue(j - expertStart_, loopTokenCount); + } + loopTokenCount += i - lastIndex; + expertCountOutLocal.SetValue(lastExpertId - expertStart_, loopTokenCount); + } + lastIndex = i; + lastlastExpertId = lastExpertId; + lastExpertId = curExpertId; + } + } + if ((i == curLoopElements) && ((lastExpertId >= expertStart_) && (lastExpertId < expertEnd_))) { + if constexpr (HISTOGRAMTYPE == COUNT_MODE || HISTOGRAMTYPE == KEY_VALUE_MODE) { + expertCountOutLocal.SetValue(lastExpertId - expertStart_, i - lastIndex); + loopTokenCount += i - lastIndex; + } else { + for (int64_t j = lastlastExpertId; j < lastExpertId; j++) { + expertCountOutLocal.SetValue(j - expertStart_, loopTokenCount); + } + loopTokenCount += i - lastIndex; + expertCountOutLocal.SetValue(lastExpertId - expertStart_, loopTokenCount); + for (int64_t j = lastExpertId; j < expertEnd_; j++) { + expertCountOutLocal.SetValue(j - expertStart_, loopTokenCount); + } + } + } else { + if constexpr (HISTOGRAMTYPE == EXERPT_TOKENS_CUMSUM) { + for (int64_t j = lastlastExpertId; j < expertEnd_; j++) { + expertCountOutLocal.SetValue(j - expertStart_, loopTokenCount); + } + } + } + actualExpertTotalNum_ += loopTokenCount; + finalExpertId = lastExpertId; + expertTokenValue = (i - lastIndex); + + expertCountOutToTempQueue_.EnQue(expertCountOutLocal); + sortedExpertIdxInQueue_.FreeTensor(sortedExpertIdxInLocal); +} + +template +__aicore__ inline void ExpertTokensCount::CopyOutExpertTotalCount() +{ + LocalTensor expertTotalCountLocal = expertTotalCountQueue_.AllocTensor(); + DataCopyExtParams copyTotalCountParams{static_cast(1), static_cast(sizeof(int32_t)), 0, 0, 0}; + expertTotalCountLocal.SetValue(0, static_cast(actualExpertTotalNum_)); + SetWaitFlag(HardEvent::S_MTE3); + SetAtomicAdd(); + DataCopyPad(expertTotalCountGm_, expertTotalCountLocal, copyTotalCountParams); + SetAtomicNone(); + expertTotalCountQueue_.FreeTensor(expertTotalCountLocal); +} + +template +__aicore__ inline void ExpertTokensCount::CopyOut() +{ + LocalTensor expertCountOutLocal = expertCountOutToTempQueue_.DeQue(); + DataCopyExtParams copyParams{static_cast(1), static_cast((actualExpertNum_) * sizeof(int32_t)), + 0, 0, 0}; + SetWaitFlag(HardEvent::S_MTE3); + SetAtomicAdd(); + DataCopyPad(expertCountTempGm_, expertCountOutLocal, copyParams); + SetAtomicNone(); + + if (dropPadMode_ == DROP_PAD) { + expertCountOutLocal.SetValue(0, finalExpertId); + expertCountOutLocal.SetValue(1, expertTokenValue); + DataCopyExtParams copyParams{static_cast(1), + static_cast(EXPERT_ID_VALUE_NUM * sizeof(int32_t)), 0, 0, 0}; + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(expertIdxValueGm_[blockIdx_ * EXPERT_ID_VALUE_NUM], expertCountOutLocal, copyParams); + } + expertCountOutToTempQueue_.FreeTensor(expertCountOutLocal); +} + +template +__aicore__ inline void ExpertTokensCount::expertCountCopyIn() +{ + LocalTensor expertCountTempInLocal = expertCountTempInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{static_cast(1), + static_cast((actualExpertNum_) * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(expertCountTempInLocal, expertCountTempGm_, dataCopyParams, dataCopyPadParams); + expertCountTempInQueue_.EnQue(expertCountTempInLocal); +} + +template +__aicore__ inline void ExpertTokensCount::expertCountCompute() +{ + LocalTensor expertCountTempInLocal = expertCountTempInQueue_.DeQue(); + LocalTensor expertCountOutLocal = expertIdxCountOutQueue_.AllocTensor(); + if constexpr (HISTOGRAMTYPE == KEY_VALUE_MODE) { + int64_t expertOffset = 0; + Duplicate(expertCountOutLocal.ReinterpretCast(), static_cast(0), + static_cast(expertCountElements_ * KEY_VALUE_MODE)); + SetWaitFlag(HardEvent::V_S); + for (int64_t i = 0; i < actualExpertNum_; i++) { + int64_t expertCount = static_cast(expertCountTempInLocal.GetValue(i)); + if (expertCount != 0) { + expertCountOutLocal.SetValue(expertOffset * KEY_VALUE_MODE_DIM_NUM, i + expertStart_); + expertCountOutLocal.SetValue(expertOffset * KEY_VALUE_MODE_DIM_NUM + 1, expertCount); + expertOffset++; + } + } + } else { + Cast(expertCountOutLocal, expertCountTempInLocal, RoundMode::CAST_NONE, actualExpertNum_); + } + + expertIdxCountOutQueue_.EnQue(expertCountOutLocal); + expertCountTempInQueue_.FreeTensor(expertCountTempInLocal); +} + +template +__aicore__ inline void ExpertTokensCount::expertCountCopyOut() +{ + LocalTensor expertCountOutLocal = expertIdxCountOutQueue_.DeQue(); + DataCopyExtParams copyParams{static_cast(1), + static_cast(expertCountElements_ * sizeof(int64_t)), 0, 0, 0}; + DataCopyPad(expertTokensCountGm_, expertCountOutLocal, copyParams); + copyParams.blockLen = sizeof(int32_t); + expertIdxCountOutQueue_.FreeTensor(expertCountOutLocal); +} + +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_EXPERT_TOKENS_COUNT_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load.h new file mode 100644 index 000000000..6b985ec08 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load.h @@ -0,0 +1,280 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_full_load.h + * \brief + */ +#ifndef MOE_CUSTOM_FULL_LOAD_H +#define MOE_CUSTOM_FULL_LOAD_H + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +class MoeCustomFullLoad { +public: + __aicore__ inline MoeCustomFullLoad(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR offset, GM_ADDR expandedX, + GM_ADDR expandedRowIdx, GM_ADDR expertTokensCountOrCumsum, GM_ADDR expandedScale, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyIn(); + __aicore__ inline void SortCompute(); + __aicore__ inline void ExpertCountCompute(); + __aicore__ inline void CopyOutDynamicQuant(); + +private: + int64_t sortNum; + + TPipe *pipe; + TQue sortDataCopyInQueue; + TQue sortDataCopyOutQueue; + TQue expertTokensCountOrCumsumOutQueue; + TQue smoothInQueue; + TQue inputXInQueue; + TQue inputXOutQueue; + TQue scaleOutQueue; + TQue rowIdxOutQueue; + + TBuf tempBuffer; + TBuf sortedBuffer; + TBuf quantTempBuffer; + + GlobalTensor inputXGm; + GlobalTensor smoothGm; + GlobalTensor expandedXGm; + GlobalTensor expandedScaleGm; + GlobalTensor expertIdxGm; + GlobalTensor expendedRowIdxGm; + GlobalTensor sortedExpertForSourceRowGm; + GlobalTensor expandDstToSrcRowGm; + GlobalTensor sortedexpertIdxGm; + GlobalTensor expertCountTempGm; + GlobalTensor expandedRowIdxGm; + GlobalTensor expertTokensCountOrCumsumGm; + + int64_t blockIdx = 0; + int64_t tileLength; + int64_t bufferNum = 1; + int64_t totalLength; + int64_t n; + int64_t k; + int64_t cols_; + int64_t expertNum_ = 256; + int64_t rowIdxType_; + int64_t kvFactor = 2; + static constexpr int64_t DST_BLK_STRIDE = 1; + static constexpr int64_t DST_REP_STRIDE = 8; +}; + +__aicore__ inline void MoeCustomFullLoad::CopyIn() +{ + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + DataCopyExtParams dataCopyParams{static_cast(1), + static_cast(this->totalLength * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal[0], expertIdxGm, dataCopyParams, dataCopyPadParams); + LocalTensor rowIdxLocal = inLocal[this->sortNum]; + ArithProgression(rowIdxLocal, 0, 1, this->sortNum); + sortDataCopyInQueue.EnQue(inLocal); +} + +__aicore__ inline void MoeCustomFullLoad::SortCompute() +{ + LocalTensor inLocal = sortDataCopyInQueue.DeQue(); + LocalTensor expertIdx = inLocal[0]; + LocalTensor expertIdxFp32 = expertIdx.ReinterpretCast(); + Cast(expertIdxFp32, expertIdx, RoundMode::CAST_ROUND, this->tileLength); + Muls(expertIdxFp32, expertIdxFp32, (float)-1, this->tileLength); + int64_t duplicateNum = this->totalLength % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + int duplicateIndex = this->totalLength - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expertIdxFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + } + + LocalTensor concatLocal; + LocalTensor tempTensor = tempBuffer.Get(GetSortLen(this->sortNum)); + Concat(concatLocal, expertIdxFp32, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM); + + LocalTensor sortedLocal = sortedBuffer.Get(GetSortLen(this->sortNum)); + LocalTensor sourceRowLocal; + sourceRowLocal = inLocal[this->sortNum].ReinterpretCast(); + Sort(sortedLocal, concatLocal, sourceRowLocal, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM); + + LocalTensor outLocal = sortDataCopyOutQueue.AllocTensor(); + LocalTensor sortedExpertForSourceRowLocal = outLocal[0]; + LocalTensor expandDstToSrcRowLocal; + expandDstToSrcRowLocal = outLocal[this->sortNum].ReinterpretCast(); + Extract(sortedExpertForSourceRowLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum / ONE_REPEAT_SORT_NUM); + Muls(sortedExpertForSourceRowLocal, sortedExpertForSourceRowLocal, (float)-1, this->tileLength); + + LocalTensor expertForSourceRowLocalInt32; + expertForSourceRowLocalInt32 = sortedExpertForSourceRowLocal.ReinterpretCast(); + Cast(expertForSourceRowLocalInt32, sortedExpertForSourceRowLocal, RoundMode::CAST_ROUND, this->tileLength); + sortDataCopyOutQueue.EnQue(outLocal); + sortDataCopyInQueue.FreeTensor(inLocal); +} + +__aicore__ inline void MoeCustomFullLoad::ExpertCountCompute() +{ + LocalTensor outLocal = sortDataCopyOutQueue.DeQue(); + LocalTensor sortedExpertId = outLocal; + LocalTensor expertTokensLocalTensor = expertTokensCountOrCumsumOutQueue.AllocTensor(); + + int64_t i = 0; + int32_t lastExpertId = sortedExpertId.GetValue(0); + int32_t lastIndex = 0; + int64_t index = 0; + for (i = 1; i < this->totalLength; i++) { + int32_t curExpertId = sortedExpertId.GetValue(i); + if (curExpertId != lastExpertId) { + expertTokensLocalTensor.SetValue(index * kvFactor, lastExpertId); + expertTokensLocalTensor.SetValue(index * kvFactor + 1, i - lastIndex); + index++; + lastIndex = i; + lastExpertId = curExpertId; + } + } + if (i == this->totalLength) { + expertTokensLocalTensor.SetValue(index * kvFactor, lastExpertId); + expertTokensLocalTensor.SetValue(index * kvFactor + 1, i - lastIndex); + index++; + } + // totalLength < 256 + expertTokensLocalTensor.SetValue(index * kvFactor, 0); + expertTokensLocalTensor.SetValue(index * kvFactor + 1, 0); + SetWaitFlag(HardEvent::S_MTE3); + + expertTokensCountOrCumsumOutQueue.EnQue(expertTokensLocalTensor); + sortDataCopyOutQueue.EnQue(outLocal); +} + +__aicore__ inline void MoeCustomFullLoad::CopyOutDynamicQuant() +{ + LocalTensor expertTokensLocalTensor = expertTokensCountOrCumsumOutQueue.DeQue(); + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = expertNum_ * sizeof(int64_t); + DataCopyPad(expertTokensCountOrCumsumGm, expertTokensLocalTensor, intriParams); + expertTokensCountOrCumsumOutQueue.FreeTensor(expertTokensLocalTensor); + LocalTensor outLocal = sortDataCopyOutQueue.DeQue(); + + int64_t expertIdx = outLocal.GetValue(blockIdx); + LocalTensor xInLocal = inputXInQueue.AllocTensor(); + LocalTensor xOutLocal = inputXOutQueue.AllocTensor(); + LocalTensor smoothLocal = smoothInQueue.AllocTensor(); + LocalTensor scaleLocal = scaleOutQueue.AllocTensor(); + LocalTensor tempLocal = quantTempBuffer.Get(); + DataCopyExtParams copyInParams{1, static_cast(cols_ * sizeof(bfloat16_t)), 0, 0, 0}; + DataCopyExtParams smoothParams{1, static_cast(cols_ * sizeof(float)), 0, 0, 0}; + DataCopyExtParams copyOutParams{1, static_cast(cols_ * sizeof(int8_t)), 0, 0, 0}; + DataCopyPad(xInLocal, inputXGm, copyInParams, {false, 0, 0, 0}); + DataCopyPad(smoothLocal, smoothGm[expertIdx * cols_], smoothParams, {false, 0, 0, 0}); + smoothInQueue.EnQue(smoothLocal); + smoothLocal = smoothInQueue.DeQue(); + Cast(tempLocal, xInLocal, RoundMode::CAST_NONE, cols_); + Mul(smoothLocal, tempLocal, smoothLocal, cols_); + // compute scale + Abs(tempLocal, smoothLocal, cols_); + ReduceMax(scaleLocal, tempLocal, tempLocal, cols_); + float scaleValue = scaleLocal.GetValue(0) / 127.0f; + Duplicate(scaleLocal, scaleValue, DST_REP_STRIDE); + Duplicate(tempLocal, scaleValue, cols_); + // compute quant + Div(tempLocal, smoothLocal, tempLocal, cols_); + Cast(tempLocal.ReinterpretCast(), tempLocal, RoundMode::CAST_ODD, cols_); // fp32->fp16 + Cast(xOutLocal, tempLocal.ReinterpretCast(), RoundMode::CAST_RINT, cols_); // fp16->int8 + inputXOutQueue.EnQue(xOutLocal); + xOutLocal = inputXOutQueue.DeQue(); + scaleOutQueue.EnQue(scaleLocal); + scaleLocal = scaleOutQueue.DeQue(); + DataCopyPad(expandedXGm[blockIdx * cols_], xOutLocal, copyOutParams); + DataCopyPad(expandedScaleGm[blockIdx], scaleLocal, {1, 4, 0, 0, 0}); + smoothInQueue.FreeTensor(smoothLocal); + inputXInQueue.FreeTensor(xInLocal); + inputXOutQueue.FreeTensor(xOutLocal); + scaleOutQueue.FreeTensor(scaleLocal); + + if (blockIdx == 0) { + intriParams.blockLen = this->totalLength * sizeof(int32_t); + if (rowIdxType_ == 1) { + DataCopyPad(expandedRowIdxGm, outLocal[this->sortNum], intriParams); + } else if (rowIdxType_ == 0) { + LocalTensor rowIdxLocalTensor = rowIdxOutQueue.AllocTensor(); + for (int i = 0; i < this->totalLength; i++) { + int32_t dstIdx = outLocal[this->sortNum].GetValue(i); + rowIdxLocalTensor.SetValue(dstIdx, i); + } + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(expandedRowIdxGm, rowIdxLocalTensor, intriParams); + rowIdxOutQueue.FreeTensor(rowIdxLocalTensor); + } + } + sortDataCopyOutQueue.FreeTensor(outLocal); +} + +__aicore__ inline void MoeCustomFullLoad::Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR offset, + GM_ADDR expandedX, GM_ADDR expandedRowIdx, GM_ADDR expertTokensCountOrCumsum, + GM_ADDR expandedScale, const MoeInitRoutingCustomTilingData *tilingData, + TPipe *tPipe) +{ + this->pipe = tPipe; + this->blockIdx = GetBlockIdx(); + this->n = tilingData->n; + this->k = tilingData->k; + this->tileLength = Align(tilingData->vbsComputeParamsOp.lastCorePerLoopElements, sizeof(int32_t)); + this->sortNum = Ceil(this->tileLength, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + this->totalLength = tilingData->n * tilingData->k; + cols_ = tilingData->cols; + rowIdxType_ = tilingData->rowIdxType; + + expertIdxGm.SetGlobalBuffer((__gm__ int32_t *)expertIdx, this->tileLength); + + expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx, this->tileLength); + expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int64_t *)expertTokensCountOrCumsum, this->tileLength); + + inputXGm.SetGlobalBuffer((__gm__ bfloat16_t *)x, this->n * cols_); + smoothGm.SetGlobalBuffer((__gm__ float *)scale, expertNum_ * cols_); + expandedXGm.SetGlobalBuffer((__gm__ int8_t *)expandedX, this->n * cols_ * this->k); + expandedScaleGm.SetGlobalBuffer((__gm__ float *)expandedScale, this->n * this->k); + + // key and value + int64_t buffSize = this->sortNum * sizeof(int32_t) * kvFactor; + pipe->InitBuffer(sortDataCopyInQueue, bufferNum, buffSize); + pipe->InitBuffer(sortDataCopyOutQueue, bufferNum, buffSize); + pipe->InitBuffer(tempBuffer, buffSize); + pipe->InitBuffer(sortedBuffer, buffSize); + pipe->InitBuffer(expertTokensCountOrCumsumOutQueue, bufferNum, Align(expertNum_ * kvFactor, sizeof(int32_t))); + + pipe->InitBuffer(smoothInQueue, bufferNum, AlignBytes(cols_, sizeof(float))); + pipe->InitBuffer(inputXInQueue, bufferNum, AlignBytes(cols_, sizeof(bfloat16_t))); + pipe->InitBuffer(inputXOutQueue, bufferNum, AlignBytes(cols_, sizeof(int8_t))); + pipe->InitBuffer(quantTempBuffer, AlignBytes(cols_, sizeof(float))); + pipe->InitBuffer(scaleOutQueue, bufferNum, AlignBytes(1, sizeof(float))); + pipe->InitBuffer(rowIdxOutQueue, bufferNum, AlignBytes(this->totalLength, sizeof(int32_t))); +} + +__aicore__ inline void MoeCustomFullLoad::Process() +{ + if (this->blockIdx < GetBlockNum()) { + CopyIn(); + SortCompute(); + ExpertCountCompute(); + CopyOutDynamicQuant(); + } +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_FULL_LOAD_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load_base.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load_base.h new file mode 100644 index 000000000..897c33226 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load_base.h @@ -0,0 +1,512 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_base_full_load.h + * \brief + */ +#ifndef MOE_CUSTOM_FULL_LOAD_BASE_H +#define MOE_CUSTOM_FULL_LOAD_BASE_H + +#include "moe_custom_common.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +template +class MoeCustomFullLoadBase { +public: + __aicore__ inline MoeCustomFullLoadBase(){}; + __aicore__ inline void Init(GM_ADDR expertIdx, GM_ADDR expandedRowIdx, GM_ADDR expertTokensCountOrCumsum, + GM_ADDR workspace, const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe); + +protected: + __aicore__ inline void CopyIn(); + __aicore__ inline void Compute(); + __aicore__ inline void TilingInKernel(); + __aicore__ inline void SortComputeWithRange(); + __aicore__ inline void SortCompute(); + __aicore__ inline void CopyOutIdx(); + __aicore__ inline void CopyOutDefaultGatherIdx(); + __aicore__ inline void CopyOutDefaultTokenCountOrCumsum(); + __aicore__ inline void ComputeExpertTokenCountOrCumsum(); + +protected: + int64_t sortNum_; + const MoeCustomGatherOutComputeTilingData *gatherOutTilingData_; + int64_t blockIdx_; + int64_t needCoreNum_; + int64_t coreIndicesElements_; + int64_t perCoreIndicesElements_; + int64_t k_; + int64_t n_; + int64_t cols_; + int64_t dropPadMode_; + int64_t activeNum_; + int64_t expertNum_; + int64_t expertStart_ = 0; + int64_t expertEnd_ = 0; + int64_t bufferNum_ = 1; + int64_t kvFactor_ = 2; + int64_t totalLength_; + int64_t tileLength_; + int64_t expertTokensNumType_ = 0; + int64_t expertTokensNumFlag_ = 0; + uint64_t actual_idx_num_ = 0; + int64_t ep_ = 0; + int64_t gatherFirstFullload_ = 0; + int64_t isInputScale_ = 0; + int64_t rowIdxType_ = 0; + int64_t actualExpertNum_ = 0; + int64_t expertCountElements_ = 0; + int64_t curIndexStart_; + int64_t startXRow_; + int64_t endXRow_; + int64_t quantMode_ = -1; + + static constexpr int64_t DST_BLK_STRIDE = 1; + static constexpr int64_t DST_REP_STRIDE = 8; + static constexpr int64_t MASK_STRIDE = 64; + + TQue expandedRowIdxCopyOutQueue_; + TQue expandedExpertIdxCopyOutQueue_; + TQue expandDstToSrcRowQueue_; + TQue expertTokensCopyOutQueue_; + TQue sortDataCopyInQueue_; + + TBuf tempBuffer_; + TBuf sortedBuffer_; + + GlobalTensor expertIdxGm_; + GlobalTensor expandedRowIdxGm_; + GlobalTensor expertTokensCountOrCumsumGm_; + + TPipe *pipe_; +}; + +template +__aicore__ inline void MoeCustomFullLoadBase::Init(GM_ADDR expertIdx, GM_ADDR expandedRowIdx, + GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + this->gatherOutTilingData_ = &(tilingData->gatherOutComputeParamsOp); + this->blockIdx_ = GetBlockIdx(); + this->n_ = tilingData->n; + this->k_ = tilingData->k; + this->cols_ = tilingData->cols; + this->expertStart_ = tilingData->expertStart; + this->expertEnd_ = tilingData->expertEnd; + this->needCoreNum_ = this->gatherOutTilingData_->needCoreNum; + + this->perCoreIndicesElements_ = this->gatherOutTilingData_->perCoreIndicesElements; + this->dropPadMode_ = tilingData->dropPadMode; + this->activeNum_ = tilingData->activeNum; + this->quantMode_ = tilingData->quantMode; + if (this->blockIdx_ == this->gatherOutTilingData_->needCoreNum - 1) { + this->coreIndicesElements_ = this->gatherOutTilingData_->lastCoreIndicesElements; + } else { + this->coreIndicesElements_ = this->gatherOutTilingData_->perCoreIndicesElements; + } + this->expertTokensNumType_ = tilingData->expertTokensNumType; + this->expertTokensNumFlag_ = tilingData->expertTokensNumFlag; + this->expertNum_ = tilingData->expertNum; + this->totalLength_ = tilingData->n * tilingData->k; + this->ep_ = tilingData->ep; + this->gatherFirstFullload_ = tilingData->gatherFirstFullload; + this->isInputScale_ = tilingData->isInputScale; + this->tileLength_ = Align(tilingData->vbsComputeParamsOp.lastCorePerLoopElements, sizeof(int32_t)); + this->sortNum_ = Ceil(this->tileLength_, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + this->actual_idx_num_ = this->totalLength_; + this->rowIdxType_ = tilingData->rowIdxType; + this->actualExpertNum_ = tilingData->actualExpertNum; + this->pipe_ = tPipe; + + expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx, this->tileLength_); + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx, this->tileLength_); + if (this->expertTokensNumFlag_ > 0) { + expertTokensCountOrCumsumGm_.SetGlobalBuffer((__gm__ int64_t *)expertTokensCountOrCumsum); + } + + if (expertTokensNumType_ == EXERPT_TOKENS_KEY_VALUE) { + expertCountElements_ = expertNum_ * EXERPT_TOKENS_KEY_VALUE; + } else { + expertCountElements_ = actualExpertNum_; + } + int64_t buffSize = this->sortNum_ * sizeof(int32_t); + + curIndexStart_ = this->blockIdx_ * this->perCoreIndicesElements_; + startXRow_ = curIndexStart_ / this->k_; + endXRow_ = (curIndexStart_ + this->coreIndicesElements_ - 1) / this->k_; + + pipe_->InitBuffer(expandedExpertIdxCopyOutQueue_, bufferNum_, buffSize); + pipe_->InitBuffer(expertTokensCopyOutQueue_, bufferNum_, AlignBytes(expertCountElements_, sizeof(int64_t))); + pipe_->InitBuffer(expandDstToSrcRowQueue_, bufferNum_, buffSize); + pipe_->InitBuffer(expandedRowIdxCopyOutQueue_, bufferNum_, buffSize); + pipe_->InitBuffer(sortDataCopyInQueue_, bufferNum_, buffSize * kvFactor_); + pipe_->InitBuffer(tempBuffer_, buffSize * kvFactor_); + pipe_->InitBuffer(sortedBuffer_, buffSize * kvFactor_); +} + +template +__aicore__ inline void MoeCustomFullLoadBase::CopyIn() +{ + LocalTensor inLocal = sortDataCopyInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{static_cast(1), static_cast(totalLength_ * sizeof(int32_t)), 0, + 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal[0], expertIdxGm_, dataCopyParams, dataCopyPadParams); + ArithProgression(inLocal[this->sortNum_], 0, 1, totalLength_); + sortDataCopyInQueue_.EnQue(inLocal); +} + +template +__aicore__ inline void MoeCustomFullLoadBase::Compute() +{ + if (ep_) { + SortComputeWithRange(); + } else { + SortCompute(); + } +} + +template +__aicore__ inline void MoeCustomFullLoadBase::SortComputeWithRange() +{ + LocalTensor inLocal = sortDataCopyInQueue_.DeQue(); + LocalTensor expertIdxLocal = inLocal[0]; + LocalTensor expertIdxLocalFp32 = expertIdxLocal.ReinterpretCast(); + LocalTensor rowIdxLocal = inLocal[this->sortNum_].template ReinterpretCast(); + Cast(expertIdxLocalFp32, expertIdxLocal, RoundMode::CAST_ROUND, totalLength_); + PipeBarrier(); + Muls(expertIdxLocalFp32, expertIdxLocalFp32, (float)-1, totalLength_); + PipeBarrier(); + if (gatherFirstFullload_) { + int64_t maskOffset = AlignBytes(Ceil(totalLength_, MASK_STRIDE) * MASK_STRIDE / DST_REP_STRIDE, sizeof(int8_t)); + LocalTensor compareScalarMaskLocalTensor0 = tempBuffer_.Get()[maskOffset]; + LocalTensor compareScalarMaskLocalTensor1 = tempBuffer_.Get()[maskOffset * kvFactor_]; + LocalTensor gatherMaskLocalTensor = tempBuffer_.Get(); + + // Find elements >= expertStart_, which means -elements <= -expertStart_ + AscendC::CompareScalar( + compareScalarMaskLocalTensor0, expertIdxLocalFp32, static_cast(-expertStart_), AscendC::CMPMODE::LE, + (totalLength_ + ONE_REPEAT_COMPARE_NUM - 1) / ONE_REPEAT_COMPARE_NUM * ONE_REPEAT_COMPARE_NUM); + PipeBarrier(); + + // Find elements < expertEnd_, which means -elements > -expertEnd_ + AscendC::CompareScalar( + compareScalarMaskLocalTensor1, expertIdxLocalFp32, static_cast(-expertEnd_), AscendC::CMPMODE::GT, + (totalLength_ + ONE_REPEAT_COMPARE_NUM - 1) / ONE_REPEAT_COMPARE_NUM * ONE_REPEAT_COMPARE_NUM); + PipeBarrier(); + + And(gatherMaskLocalTensor.ReinterpretCast(), + compareScalarMaskLocalTensor0.ReinterpretCast(), + compareScalarMaskLocalTensor1.ReinterpretCast(), + Ceil(totalLength_, MASK_STRIDE) * MASK_STRIDE / DST_REP_STRIDE / kvFactor_); + PipeBarrier(); + + uint64_t rsvdCnt = 0; + GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = 1; + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = DST_REP_STRIDE; + gatherMaskParams.src1RepeatStride = DST_REP_STRIDE; + GatherMask(expertIdxLocalFp32, expertIdxLocalFp32, gatherMaskLocalTensor.ReinterpretCast(), true, + static_cast(totalLength_), gatherMaskParams, rsvdCnt); + PipeBarrier(); + actual_idx_num_ = rsvdCnt; + sortNum_ = Ceil(actual_idx_num_, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + + GatherMask(rowIdxLocal, rowIdxLocal, gatherMaskLocalTensor.ReinterpretCast(), true, + static_cast(totalLength_), gatherMaskParams, actual_idx_num_); + PipeBarrier(); + TilingInKernel(); + } else { + LocalTensor maskLocalTensor = tempBuffer_.Get(); + AscendC::CompareScalar( + maskLocalTensor, expertIdxLocalFp32, static_cast(-expertStart_), AscendC::CMPMODE::GT, + (totalLength_ + ONE_REPEAT_COMPARE_NUM - 1) / ONE_REPEAT_COMPARE_NUM * ONE_REPEAT_COMPARE_NUM); + LocalTensor floatMinLocalTensor = sortedBuffer_.Get(); + Duplicate(floatMinLocalTensor, MIN_FP32, totalLength_); + PipeBarrier(); + Select(expertIdxLocalFp32, maskLocalTensor, floatMinLocalTensor, expertIdxLocalFp32, + SELMODE::VSEL_TENSOR_TENSOR_MODE, totalLength_); + PipeBarrier(); + } + // handle actual_idx_num_ == 0 + if (actual_idx_num_ < 1) { + sortDataCopyInQueue_.FreeTensor(inLocal); + return; + } + int64_t duplicateNum = actual_idx_num_ % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + int duplicateIndex = actual_idx_num_ - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> (FP32_ONE_REPEAT_NUM - ONE_REPEAT_SORT_NUM)); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expertIdxLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + PipeBarrier(); + } + + LocalTensor concatLocal = expertIdxLocalFp32; + LocalTensor tempTensor = tempBuffer_.Get(GetSortLen(this->sortNum_)); + Concat(concatLocal, expertIdxLocalFp32, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + PipeBarrier(); + + LocalTensor sortedLocal = sortedBuffer_.Get(GetSortLen(this->sortNum_)); + Sort(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + PipeBarrier(); + LocalTensor expandedExpertIdxLocal = expandedExpertIdxCopyOutQueue_.AllocTensor(); + LocalTensor expandDstToSrcRowLocal = expandDstToSrcRowQueue_.AllocTensor(); + Extract(expandedExpertIdxLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum_ / ONE_REPEAT_SORT_NUM); + PipeBarrier(); + + Muls(expandedExpertIdxLocal, expandedExpertIdxLocal, (float)-1, actual_idx_num_); + PipeBarrier(); + LocalTensor expandedExpertIdxLocalInt32; + expandedExpertIdxLocalInt32 = expandedExpertIdxLocal.ReinterpretCast(); + Cast(expandedExpertIdxLocalInt32, expandedExpertIdxLocal, RoundMode::CAST_ROUND, actual_idx_num_); + PipeBarrier(); + expandedExpertIdxCopyOutQueue_.EnQue(expandedExpertIdxLocalInt32); + expandDstToSrcRowQueue_.EnQue(expandDstToSrcRowLocal); + sortDataCopyInQueue_.FreeTensor(inLocal); +} + +template +__aicore__ inline void MoeCustomFullLoadBase::SortCompute() +{ + LocalTensor inLocal = sortDataCopyInQueue_.DeQue(); + LocalTensor expertIdxLocal = inLocal[0]; + LocalTensor expertIdxLocalFp32 = expertIdxLocal.ReinterpretCast(); + Cast(expertIdxLocalFp32, expertIdxLocal, RoundMode::CAST_ROUND, totalLength_); + PipeBarrier(); + Muls(expertIdxLocalFp32, expertIdxLocalFp32, (float)-1, totalLength_); + PipeBarrier(); + int64_t duplicateNum = totalLength_ % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + int duplicateIndex = totalLength_ - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> (FP32_ONE_REPEAT_NUM - ONE_REPEAT_SORT_NUM)); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expertIdxLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + PipeBarrier(); + } + LocalTensor concatLocal = expertIdxLocalFp32; + LocalTensor tempTensor = tempBuffer_.Get(GetSortLen(this->sortNum_)); + Concat(concatLocal, expertIdxLocalFp32, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + PipeBarrier(); + LocalTensor rowIdxLocal = inLocal[this->sortNum_].template ReinterpretCast(); + LocalTensor sortedLocal = sortedBuffer_.Get(GetSortLen(this->sortNum_)); + Sort(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + PipeBarrier(); + LocalTensor expandedExpertIdxLocal = expandedExpertIdxCopyOutQueue_.AllocTensor(); + LocalTensor expandDstToSrcRowLocal = expandDstToSrcRowQueue_.AllocTensor(); + LocalTensor expandDstToSrcRowLocalFp32 = expandDstToSrcRowLocal.ReinterpretCast(); + Extract(expandedExpertIdxLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum_ / ONE_REPEAT_SORT_NUM); + PipeBarrier(); + + LocalTensor expandedRowIdx = expandedRowIdxCopyOutQueue_.AllocTensor(); + Muls(expandedExpertIdxLocal, expandedExpertIdxLocal, (float)-1, totalLength_); + PipeBarrier(); + LocalTensor expandedExpertIdxLocalInt32; + expandedExpertIdxLocalInt32 = expandedExpertIdxLocal.ReinterpretCast(); + Cast(expandedExpertIdxLocalInt32, expandedExpertIdxLocal, RoundMode::CAST_ROUND, totalLength_); + PipeBarrier(); + + Cast(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocal.ReinterpretCast(), RoundMode::CAST_ROUND, + totalLength_); + PipeBarrier(); + Muls(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocalFp32, (float)-1, totalLength_); + PipeBarrier(); + ArithProgression(inLocal[this->sortNum_], 0, 1, totalLength_); + PipeBarrier(); + if (duplicateNum > 0) { + int duplicateIndex = totalLength_ - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> (FP32_ONE_REPEAT_NUM - ONE_REPEAT_SORT_NUM)); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expandDstToSrcRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + PipeBarrier(); + } + Concat(concatLocal, expandDstToSrcRowLocalFp32, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + PipeBarrier(); + Sort(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM); + PipeBarrier(); + Extract(tempTensor, expandedRowIdx, sortedLocal, this->sortNum_ / ONE_REPEAT_SORT_NUM); + PipeBarrier(); + + if (rowIdxType_ == SCATTER or quantMode_ == 1) { + Muls(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocalFp32, (float)-1, totalLength_); + PipeBarrier(); + Cast(expandDstToSrcRowLocal.ReinterpretCast(), expandDstToSrcRowLocalFp32, RoundMode::CAST_RINT, + totalLength_); + } + expandedExpertIdxCopyOutQueue_.EnQue(expandedExpertIdxLocalInt32); + expandedRowIdxCopyOutQueue_.EnQue(expandedRowIdx); + expandDstToSrcRowQueue_.EnQue(expandDstToSrcRowLocal); + sortDataCopyInQueue_.FreeTensor(inLocal); +} + +template +__aicore__ inline void MoeCustomFullLoadBase::CopyOutDefaultGatherIdx() +{ + LocalTensor expandedRowIdx = expandedRowIdxCopyOutQueue_.AllocTensor(); + Duplicate(expandedRowIdx, static_cast(-1), static_cast(totalLength_)); + SetWaitFlag(HardEvent::V_MTE3); + DataCopyExtParams copyParams{static_cast(1), static_cast(totalLength_ * sizeof(int32_t)), 0, 0, + 0}; + DataCopyPad(expandedRowIdxGm_, expandedRowIdx, copyParams); + expandedRowIdxCopyOutQueue_.FreeTensor(expandedRowIdx); +} + +template +__aicore__ inline void MoeCustomFullLoadBase::CopyOutDefaultTokenCountOrCumsum() +{ + LocalTensor expertTokensOut = expertTokensCopyOutQueue_.AllocTensor(); + Duplicate(expertTokensOut.ReinterpretCast(), static_cast(0), + static_cast(expertCountElements_ * EXERPT_TOKENS_KEY_VALUE)); + SetWaitFlag(HardEvent::V_MTE3); + DataCopyExtParams copyParams{static_cast(1), + static_cast(expertCountElements_ * sizeof(int64_t)), 0, 0, 0}; + DataCopyPad(expertTokensCountOrCumsumGm_, expertTokensOut, copyParams); + expertTokensCopyOutQueue_.FreeTensor(expertTokensOut); +} + +template +__aicore__ inline void MoeCustomFullLoadBase::CopyOutIdx() +{ + LocalTensor expandedExpertIdx = expandedExpertIdxCopyOutQueue_.DeQue(); + LocalTensor expandDstToSrcRowLocal = expandDstToSrcRowQueue_.DeQue(); + if (rowIdxType_ == SCATTER) { + DataCopyExtParams copyParams{static_cast(1), static_cast(actual_idx_num_ * sizeof(int32_t)), + 0, 0, 0}; + DataCopyPad(expandedRowIdxGm_, expandDstToSrcRowLocal, copyParams); + } else if (ep_) { + LocalTensor expandedRowIdx = expandedRowIdxCopyOutQueue_.AllocTensor(); + Duplicate(expandedRowIdx, static_cast(-1), static_cast(totalLength_)); + SetWaitFlag(HardEvent::V_S); + for (int64_t i = 0; i < actual_idx_num_; i++) { + int32_t curExpertId = expandedExpertIdx.GetValue(i); + if (curExpertId < expertStart_ || curExpertId >= expertEnd_) { + break; + } + int64_t outIndices = expandDstToSrcRowLocal.GetValue(i); + expandedRowIdx.SetValue(outIndices, i); + } + SetWaitFlag(HardEvent::S_MTE3); + DataCopyExtParams copyParams{static_cast(1), static_cast(totalLength_ * sizeof(int32_t)), 0, + 0, 0}; + DataCopyPad(expandedRowIdxGm_, expandedRowIdx, copyParams); + expandedRowIdxCopyOutQueue_.FreeTensor(expandedRowIdx); + } else { + LocalTensor expandedRowIdx = expandedRowIdxCopyOutQueue_.DeQue(); + DataCopyExtParams copyParams{static_cast(1), static_cast(totalLength_ * sizeof(int32_t)), 0, + 0, 0}; + DataCopyPad(expandedRowIdxGm_, expandedRowIdx, copyParams); + expandedRowIdxCopyOutQueue_.EnQue(expandedRowIdx); + } + expandedExpertIdxCopyOutQueue_.EnQue(expandedExpertIdx); + expandDstToSrcRowQueue_.EnQue(expandDstToSrcRowLocal); +} + +template +__aicore__ inline void MoeCustomFullLoadBase::ComputeExpertTokenCountOrCumsum() +{ + // compute + LocalTensor expandedExpertIdx = expandedExpertIdxCopyOutQueue_.DeQue(); + LocalTensor expertTokensOut = expertTokensCopyOutQueue_.AllocTensor(); + Duplicate(expertTokensOut.ReinterpretCast(), static_cast(0), + static_cast(expertCountElements_ * EXERPT_TOKENS_KEY_VALUE)); + SetWaitFlag(HardEvent::V_S); + int64_t i = 0; + int32_t lastExpertId = expandedExpertIdx.GetValue(0); + int32_t lastLastId = lastExpertId; + int64_t tokenCount = 0; + int64_t lastIndex = 0; + int64_t Offset = 0; + for (i = 1; i < actual_idx_num_; i++) { + if ((lastExpertId >= expertEnd_) || (lastExpertId < expertStart_)) { + break; + } + int32_t curExpertId = expandedExpertIdx.GetValue(i); + if (curExpertId != lastExpertId || curExpertId >= expertEnd_) { + int64_t expertOffset = lastExpertId - expertStart_; + if (expertTokensNumType_ == EXERPT_TOKENS_KEY_VALUE) { + expertTokensOut.SetValue(Offset * EXERPT_TOKENS_KEY_VALUE, lastExpertId); + expertTokensOut.SetValue(Offset * EXERPT_TOKENS_KEY_VALUE + 1, i - lastIndex); + Offset += 1; + } else if (expertTokensNumType_ == EXERPT_TOKENS_COUNT) { + expertTokensOut.SetValue(expertOffset, i - lastIndex); + } else { + for (int64_t j = lastLastId; j < lastExpertId; j++) { + expertTokensOut.SetValue(j - expertStart_, tokenCount); + } + tokenCount += i - lastIndex; + expertTokensOut.SetValue(expertOffset, tokenCount); + } + lastIndex = i; + lastLastId = lastExpertId; + lastExpertId = curExpertId; + } + } + if ((i == actual_idx_num_) && ((lastExpertId >= expertStart_) && (lastExpertId < expertEnd_))) { + int64_t expertOffset = lastExpertId - expertStart_; + if (expertTokensNumType_ == EXERPT_TOKENS_KEY_VALUE) { + expertTokensOut.SetValue(Offset * EXERPT_TOKENS_KEY_VALUE, lastExpertId); + expertTokensOut.SetValue(Offset * EXERPT_TOKENS_KEY_VALUE + 1, i - lastIndex); + } else if (expertTokensNumType_ == EXERPT_TOKENS_COUNT) { + expertTokensOut.SetValue(expertOffset, i - lastIndex); + } else { + for (int64_t j = lastLastId; j < lastExpertId; j++) { + expertTokensOut.SetValue(j - expertStart_, tokenCount); + } + tokenCount += i - lastIndex; + expertTokensOut.SetValue(expertOffset, tokenCount); + for (int64_t j = lastExpertId; j < expertEnd_; j++) { + expertTokensOut.SetValue(j - expertStart_, tokenCount); + } + } + } else { + if (expertTokensNumType_ == EXERPT_TOKENS_CUMSUM) { + for (int64_t j = lastLastId; j < expertEnd_; j++) { + expertTokensOut.SetValue(j - expertStart_, tokenCount); + } + } + } + expandedExpertIdxCopyOutQueue_.EnQue(expandedExpertIdx); + SetWaitFlag(HardEvent::S_MTE3); + DataCopyExtParams copyParams{static_cast(1), + static_cast(expertCountElements_ * sizeof(int64_t)), 0, 0, 0}; + DataCopyPad(expertTokensCountOrCumsumGm_, expertTokensOut, copyParams); + SetWaitFlag(HardEvent::MTE3_V); + expertTokensCopyOutQueue_.FreeTensor(expertTokensOut); +} + +template +__aicore__ inline void MoeCustomFullLoadBase::TilingInKernel() +{ + int64_t coreNum = needCoreNum_; + perCoreIndicesElements_ = Ceil(actual_idx_num_, coreNum); + needCoreNum_ = Ceil(actual_idx_num_, perCoreIndicesElements_); + int64_t lastCoreIndicesElements = actual_idx_num_ - (needCoreNum_ - 1) * perCoreIndicesElements_; + if (blockIdx_ == needCoreNum_ - 1) { + coreIndicesElements_ = lastCoreIndicesElements; + } else { + coreIndicesElements_ = perCoreIndicesElements_; + } + curIndexStart_ = this->blockIdx_ * this->perCoreIndicesElements_; + startXRow_ = curIndexStart_ / this->k_; + endXRow_ = (curIndexStart_ + this->coreIndicesElements_ - 1) / this->k_; +} + +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_FULL_LOAD_BASE_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load_dynamic_quant.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load_dynamic_quant.h new file mode 100644 index 000000000..3919445e9 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load_dynamic_quant.h @@ -0,0 +1,300 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_full_load_dynamic_quant.h + * \brief + */ +#ifndef MOE_CUSTOM_FULL_LOAD_DYNAMIC_QUANT_H +#define MOE_CUSTOM_FULL_LOAD_DYNAMIC_QUANT_H + +#include "moe_custom_full_load_base.h" +#include "moe_custom_common.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +template +class MoeCustomFullLoadDynamicQuant : public MoeCustomFullLoadBase { +public: + __aicore__ inline MoeCustomFullLoadDynamicQuant(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR expandedX, GM_ADDR expandedRowIdx, + GM_ADDR expertTokensCountOrCumsum, GM_ADDR expandedScale, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyOutXDynamicQuantFromGather(); + __aicore__ inline void CopyOutXDynamicQuantFromScatter(); + __aicore__ inline void FreeLocalTensor(); + __aicore__ inline void ComputeQuant(LocalTensor &smoothLocal); + +private: + TQue xCopyInQueue_; + TQue smoothInQueue_; + TBuf tmpBuff_; + TQue inputXOutQueue_; + TQue scaleOutQueue_; + + GlobalTensor xGm_; + GlobalTensor expandedXGm_; + GlobalTensor quantSmoothGm_; + GlobalTensor expandedScaleGm_; + + int64_t colsAlign_ = 0; +}; + +template +__aicore__ inline void MoeCustomFullLoadDynamicQuant::Init( + GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR expandedX, GM_ADDR expandedRowIdx, + GM_ADDR expertTokensCountOrCumsum, GM_ADDR expandedScale, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + MoeCustomFullLoadBase::Init(expertIdx, expandedRowIdx, expertTokensCountOrCumsum, workspace, tilingData, tPipe); + + xGm_.SetGlobalBuffer((__gm__ T *)x); + expandedXGm_.SetGlobalBuffer((__gm__ int8_t *)expandedX); + quantSmoothGm_.SetGlobalBuffer((__gm__ float *)scale); + expandedScaleGm_.SetGlobalBuffer((__gm__ float *)expandedScale); + this->colsAlign_ = Align(this->cols_, sizeof(T)); + if constexpr (IsSameType::value) { + this->pipe_->InitBuffer(xCopyInQueue_, 1, AlignBytes(this->cols_, sizeof(float))); + } else { + this->pipe_->InitBuffer(xCopyInQueue_, 1, 2 * AlignBytes(this->cols_, sizeof(T))); + } + this->pipe_->InitBuffer(inputXOutQueue_, 1, AlignBytes(this->cols_, sizeof(int8_t))); + this->pipe_->InitBuffer(smoothInQueue_, 1, AlignBytes(this->cols_, sizeof(float))); + this->pipe_->InitBuffer(tmpBuff_, AlignBytes(this->cols_, sizeof(float))); + this->pipe_->InitBuffer(scaleOutQueue_, 1, BLOCK_BYTES + BLOCK_BYTES); +} + +template +__aicore__ inline void MoeCustomFullLoadDynamicQuant::Process() +{ + if (this->blockIdx_ < this->needCoreNum_) { + this->CopyIn(); + this->Compute(); + + // valid expert equal zero + if (this->needCoreNum_ < 1) { + if (this->blockIdx_ == 0) { + if (this->rowIdxType_ == GATHER) { + this->CopyOutDefaultGatherIdx(); + } + if (this->expertTokensNumFlag_ == 1) { + this->CopyOutDefaultTokenCountOrCumsum(); + } + } + return; + } + + if (this->blockIdx_ == 0) { + this->CopyOutIdx(); + } + + if (this->blockIdx_ == this->needCoreNum_ - 1 && this->expertTokensNumFlag_ == 1) { + this->ComputeExpertTokenCountOrCumsum(); + } + + if (this->blockIdx_ < this->needCoreNum_) { + if constexpr (!COPYOUTTYPE && SMOOTHTYPE != SCALE_EH) { + CopyOutXDynamicQuantFromGather(); + } else { + CopyOutXDynamicQuantFromScatter(); + } + } + + FreeLocalTensor(); + } +} + +template +__aicore__ inline void +MoeCustomFullLoadDynamicQuant::ComputeQuant(LocalTensor &smoothLocal) +{ + LocalTensor tempLocal = tmpBuff_.Get(); + LocalTensor outLocal = inputXOutQueue_.AllocTensor(); + LocalTensor dynamicQuantLocal = scaleOutQueue_.AllocTensor(); + LocalTensor inLocal = xCopyInQueue_.DeQue(); + + if constexpr (!IsSameType::value && !IsSameType::value) { + Cast(inLocal, inLocal.ReinterpretCast()[colsAlign_], RoundMode::CAST_NONE, this->cols_); + PipeBarrier(); + } + + if constexpr (SMOOTHTYPE != NO_SCALE) { + Mul(inLocal, inLocal, smoothLocal, this->cols_); + PipeBarrier(); + } + + Abs(tempLocal, inLocal, this->cols_); + PipeBarrier(); + + ReduceMax(dynamicQuantLocal, tempLocal, tempLocal, this->cols_); + PipeBarrier(); + + float maxValue = dynamicQuantLocal.GetValue(0) / MAX_INT8; + + Duplicate(dynamicQuantLocal, maxValue, INT32_ONE_BLOCK_NUM); + PipeBarrier(); + Duplicate(tempLocal, maxValue, this->cols_); + PipeBarrier(); + + Div(tempLocal, inLocal, tempLocal, this->cols_); + PipeBarrier(); + + LocalTensor intLocal = tempLocal.ReinterpretCast(); + Cast(intLocal, tempLocal, RoundMode::CAST_RINT, this->cols_); + PipeBarrier(); + SetDeqScale((half)1.000000e+00f); + Cast(intLocal.ReinterpretCast(), intLocal, RoundMode::CAST_ROUND, this->cols_); + PipeBarrier(); + Cast(outLocal, intLocal.ReinterpretCast(), RoundMode::CAST_TRUNC, this->cols_); + + inputXOutQueue_.EnQue(outLocal); + scaleOutQueue_.EnQue(dynamicQuantLocal); +} + +template +__aicore__ inline void MoeCustomFullLoadDynamicQuant::CopyOutXDynamicQuantFromScatter() +{ + LocalTensor sortedRowIdx = this->expandDstToSrcRowQueue_.template DeQue(); + LocalTensor expandedExpertIdx = this->expandedExpertIdxCopyOutQueue_.template DeQue(); + + DataCopyExtParams dataXCopyParams{1, static_cast(this->cols_ * sizeof(T)), 0, 0, 0}; + DataCopyExtParams smoothCopyParams{1, static_cast(this->cols_ * sizeof(float)), 0, 0, 0}; + DataCopyExtParams intriParams{1, static_cast(this->cols_ * sizeof(int8_t)), 0, 0, 0}; + DataCopyExtParams quantScaleParams{1, static_cast(sizeof(int32_t)), 0, 0, 0}; + + LocalTensor smoothLocal = smoothInQueue_.AllocTensor(); + ; + + if constexpr (SMOOTHTYPE == SCALE_1H) { + DataCopyPad(smoothLocal, quantSmoothGm_, smoothCopyParams, {false, 0, 0, 0}); + smoothInQueue_.EnQue(smoothLocal); + smoothLocal = smoothInQueue_.DeQue(); + } + + int64_t dstIndexStart = this->curIndexStart_; + int64_t dstIndexEnd = dstIndexStart + this->coreIndicesElements_ - 1; + int32_t lastExpertIdx = -1; + + for (int64_t dstIndex = dstIndexStart; dstIndex <= dstIndexEnd; dstIndex++) { + if (this->dropPadMode_ == DROPLESS_MODE && dstIndex >= this->activeNum_) { + break; + } + int32_t srcIdx = sortedRowIdx.GetValue(dstIndex); + int32_t expertIdx = expandedExpertIdx.GetValue(dstIndex); + if (expertIdx < this->expertStart_ || expertIdx >= this->expertEnd_) { + break; + } + expertIdx = expertIdx - this->expertStart_; + LocalTensor xLocal = this->xCopyInQueue_.template AllocTensor(); + // copy in single x + if constexpr (IsSameType::value) { + DataCopyPad(xLocal, this->xGm_[srcIdx / this->k_ * this->cols_], dataXCopyParams, {false, 0, 0, 0}); + } else { + DataCopyPad(xLocal[colsAlign_], this->xGm_[srcIdx / this->k_ * this->cols_], dataXCopyParams, + {false, 0, 0, 0}); + } + xCopyInQueue_.EnQue(xLocal); + + // copyin dynamic scale + if constexpr (SMOOTHTYPE == SCALE_EH) { + if (expertIdx != lastExpertIdx) { + DataCopyPad(smoothLocal, quantSmoothGm_[expertIdx * this->cols_], smoothCopyParams, {false, 0, 0, 0}); + smoothInQueue_.EnQue(smoothLocal); + smoothLocal = smoothInQueue_.DeQue(); + lastExpertIdx = expertIdx; + } + } + + ComputeQuant(smoothLocal); + + LocalTensor quantScaleLocal = scaleOutQueue_.DeQue(); + DataCopyPad(expandedScaleGm_[dstIndex], quantScaleLocal, quantScaleParams); + + LocalTensor outLocal = inputXOutQueue_.DeQue(); + DataCopyPad(this->expandedXGm_[dstIndex * this->cols_], outLocal, intriParams); + + inputXOutQueue_.FreeTensor(outLocal); + scaleOutQueue_.FreeTensor(quantScaleLocal); + this->xCopyInQueue_.FreeTensor(xLocal); + } + smoothInQueue_.FreeTensor(smoothLocal); + this->expandDstToSrcRowQueue_.EnQue(sortedRowIdx); + this->expandedExpertIdxCopyOutQueue_.EnQue(expandedExpertIdx); +} + +template +__aicore__ inline void MoeCustomFullLoadDynamicQuant::CopyOutXDynamicQuantFromGather() +{ + DataCopyExtParams dataXCopyParams{1, static_cast(this->cols_ * sizeof(T)), 0, 0, 0}; + DataCopyExtParams smoothCopyParams{1, static_cast(this->cols_ * sizeof(float)), 0, 0, 0}; + DataCopyExtParams intriParams{1, static_cast(this->cols_ * sizeof(int8_t)), 0, 0, 0}; + DataCopyExtParams quantScaleParams{1, static_cast(sizeof(int32_t)), 0, 0, 0}; + + LocalTensor expandedRowIdx = this->expandedRowIdxCopyOutQueue_.template DeQue(); + LocalTensor smoothLocal = smoothInQueue_.AllocTensor(); + int64_t curIndex = this->blockIdx_ * this->perCoreIndicesElements_; + int64_t curIndexEnd = curIndex + this->coreIndicesElements_ - 1; + + if constexpr (SMOOTHTYPE == SCALE_1H) { + DataCopyPad(smoothLocal, quantSmoothGm_, smoothCopyParams, {false, 0, 0, 0}); + smoothInQueue_.EnQue(smoothLocal); + smoothLocal = smoothInQueue_.DeQue(); + } + + for (int64_t row = this->startXRow_; row <= this->endXRow_; row++) { + LocalTensor xLocal = xCopyInQueue_.AllocTensor(); + if constexpr (IsSameType::value) { + DataCopyPad(xLocal, this->xGm_[row * this->cols_], dataXCopyParams, {false, 0, 0, 0}); + } else { + DataCopyPad(xLocal[colsAlign_], this->xGm_[row * this->cols_], dataXCopyParams, {false, 0, 0, 0}); + } + xCopyInQueue_.EnQue(xLocal); + ComputeQuant(smoothLocal); + + LocalTensor quantScaleLocal = scaleOutQueue_.DeQue(); + LocalTensor outLocal = inputXOutQueue_.DeQue(); + while (curIndex <= curIndexEnd && curIndex / this->k_ == row) { + int32_t outIndex = expandedRowIdx.GetValue(curIndex); + curIndex++; + if (outIndex == -1 || this->dropPadMode_ == DROPLESS_MODE && outIndex >= this->activeNum_) { + continue; + } + DataCopyPad(expandedXGm_[outIndex * this->cols_], outLocal, intriParams); + DataCopyPad(expandedScaleGm_[outIndex], quantScaleLocal, quantScaleParams); + } + + xCopyInQueue_.FreeTensor(xLocal); + inputXOutQueue_.FreeTensor(outLocal); + scaleOutQueue_.FreeTensor(quantScaleLocal); + } + + smoothInQueue_.FreeTensor(smoothLocal); + this->expandedRowIdxCopyOutQueue_.EnQue(expandedRowIdx); +} + +template +__aicore__ inline void MoeCustomFullLoadDynamicQuant::FreeLocalTensor() +{ + if constexpr (!COPYOUTTYPE) { + LocalTensor expandedRowIdx = this->expandedRowIdxCopyOutQueue_.template DeQue(); + this->expandedRowIdxCopyOutQueue_.FreeTensor(expandedRowIdx); + } + LocalTensor sortedRowIdx = this->expandDstToSrcRowQueue_.template DeQue(); + LocalTensor expandedExpertIdx = this->expandedExpertIdxCopyOutQueue_.template DeQue(); + this->expandDstToSrcRowQueue_.FreeTensor(sortedRowIdx); + this->expandedExpertIdxCopyOutQueue_.FreeTensor(expandedExpertIdx); +} + +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_FULL_LOAD_DYNAMIC_QUANT_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load_static_quant.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load_static_quant.h new file mode 100644 index 000000000..6a8991b0e --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load_static_quant.h @@ -0,0 +1,229 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_static_quant_full_load.h + * \brief + */ +#ifndef MOE_CUSTOM_FULL_LOAD_STATIC_QUANT_H +#define MOE_CUSTOM_FULL_LOAD_STATIC_QUANT_H + +#include "moe_custom_full_load_base.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +template +class MoeCustomFullLoadStaticQuant : public MoeCustomFullLoadBase { +public: + __aicore__ inline MoeCustomFullLoadStaticQuant(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR offset, GM_ADDR expandedX, + GM_ADDR expandedRowIdx, GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyOutXStaticQuant(); + __aicore__ inline void FreeLocalTensor(); + __aicore__ inline void ComputeQuant(int64_t xLocalLength); + +private: + TQue xCopyInQueue_; + TQue floatQueue_; + TQue halfQueue_; + TQue inputXOutQueue_; + + GlobalTensor xGm_; + GlobalTensor expandedXGm_; + GlobalTensor scaleGm_; + GlobalTensor offsetGm_; + + float scale_; + float offset_; +}; + +template +__aicore__ inline void MoeCustomFullLoadStaticQuant::Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR offset, + GM_ADDR expandedX, GM_ADDR expandedRowIdx, + GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + MoeCustomFullLoadBase::Init(expertIdx, expandedRowIdx, expertTokensCountOrCumsum, workspace, tilingData, tPipe); + + xGm_.SetGlobalBuffer((__gm__ T *)x); + expandedXGm_.SetGlobalBuffer((__gm__ int8_t *)expandedX); + scaleGm_.SetGlobalBuffer((__gm__ float *)scale, 1); + offsetGm_.SetGlobalBuffer((__gm__ float *)offset, 1); + this->scale_ = scaleGm_.GetValue(0); + this->offset_ = offsetGm_.GetValue(0); + SetWaitFlag(HardEvent::S_V); + int64_t curIndexStart = this->blockIdx_ * this->perCoreIndicesElements_; + int64_t rowLength = 0; + if (this->ep_) { + rowLength = 1; + } else { + rowLength = (curIndexStart + this->coreIndicesElements_ - 1) / this->k_ - curIndexStart / this->k_ + 1; + } + int64_t xAlignedCount = Align(this->cols_, sizeof(int8_t)); + this->pipe_->InitBuffer(xCopyInQueue_, this->bufferNum_, xAlignedCount * sizeof(T) * rowLength); + this->pipe_->InitBuffer(inputXOutQueue_, 1, xAlignedCount * sizeof(int8_t) * rowLength); + this->pipe_->InitBuffer(floatQueue_, 1, xAlignedCount * sizeof(float) * rowLength); + this->pipe_->InitBuffer(halfQueue_, 1, xAlignedCount * sizeof(half) * rowLength); +} + +template +__aicore__ inline void MoeCustomFullLoadStaticQuant::Process() +{ + if (this->blockIdx_ < this->needCoreNum_) { + this->CopyIn(); + this->Compute(); + + // valid expert equal zero + if (this->needCoreNum_ < 1) { + if (this->blockIdx_ == 0) { + if (this->rowIdxType_ == GATHER) { + this->CopyOutDefaultGatherIdx(); + } + if (this->expertTokensNumFlag_ == 1) { + this->CopyOutDefaultTokenCountOrCumsum(); + } + } + return; + } + + if (this->blockIdx_ == 0) { + this->CopyOutIdx(); + } + if (this->blockIdx_ == this->needCoreNum_ - 1 && this->expertTokensNumFlag_ == 1) { + this->ComputeExpertTokenCountOrCumsum(); + } + if (this->blockIdx_ < this->needCoreNum_) { + CopyOutXStaticQuant(); + } + FreeLocalTensor(); + } +} + +template +__aicore__ inline void MoeCustomFullLoadStaticQuant::ComputeQuant(int64_t xLocalLength) +{ + LocalTensor floatLocal; + LocalTensor inLocal; + LocalTensor outLocal = inputXOutQueue_.AllocTensor(); + LocalTensor halfLocal = halfQueue_.AllocTensor(); + uint64_t elements = Align(this->cols_, sizeof(int8_t)) * xLocalLength; + if constexpr (IsSameType::value) { + floatLocal = this->xCopyInQueue_.template DeQue(); + } else { + inLocal = this->xCopyInQueue_.template DeQue(); + floatLocal = floatQueue_.AllocTensor(); + Cast(floatLocal, inLocal, RoundMode::CAST_NONE, elements); + PipeBarrier(); + } + Muls(floatLocal, floatLocal, this->scale_, elements); + PipeBarrier(); + Adds(floatLocal, floatLocal, this->offset_, elements); + PipeBarrier(); + LocalTensor intLocal = floatLocal.ReinterpretCast(); + Cast(intLocal, floatLocal, RoundMode::CAST_RINT, elements); + PipeBarrier(); + SetDeqScale((half)1.000000e+00f); + Cast(halfLocal, intLocal, RoundMode::CAST_ROUND, elements); + PipeBarrier(); + Cast(outLocal, halfLocal, RoundMode::CAST_TRUNC, elements); + inputXOutQueue_.EnQue(outLocal); + if constexpr (IsSameType::value) { + this->xCopyInQueue_.FreeTensor(floatLocal); + } else { + this->xCopyInQueue_.FreeTensor(inLocal); + floatQueue_.FreeTensor(floatLocal); + } + + halfQueue_.FreeTensor(halfLocal); +} + +template +__aicore__ inline void MoeCustomFullLoadStaticQuant::CopyOutXStaticQuant() +{ + int64_t curIndex = this->curIndexStart_; + int64_t curIndexEnd = curIndex + this->coreIndicesElements_ - 1; + + if (this->ep_) { + LocalTensor sortedRowIdx = this->expandDstToSrcRowQueue_.template DeQue(); + LocalTensor expandedExpertIdx = this->expandedExpertIdxCopyOutQueue_.template DeQue(); + + DataCopyExtParams dataXCopyParams{1, static_cast(this->cols_ * sizeof(T)), 0, 0, 0}; + DataCopyExtParams intriParams{1, static_cast(this->cols_ * sizeof(int8_t)), 0, 0, 0}; + + for (int64_t dstIndex = curIndex; dstIndex <= curIndexEnd; dstIndex++) { + if (this->dropPadMode_ == DROPLESS_MODE && dstIndex >= this->activeNum_) { + break; + } + int32_t srcIdx = sortedRowIdx.GetValue(dstIndex); + int32_t expertIdx = expandedExpertIdx.GetValue(dstIndex); + if (expertIdx < this->expertStart_ || expertIdx >= this->expertEnd_) { + break; + } + LocalTensor inLocal = this->xCopyInQueue_.template AllocTensor(); + // copyinx + DataCopyPad(inLocal, this->xGm_[srcIdx / this->k_ * this->cols_], dataXCopyParams, {false, 0, 0, 0}); + this->xCopyInQueue_.template EnQue(inLocal); + ComputeQuant(1); + + LocalTensor outLocal = inputXOutQueue_.DeQue(); + DataCopyPad(this->expandedXGm_[dstIndex * this->cols_], outLocal, intriParams); + inputXOutQueue_.FreeTensor(outLocal); + } + this->expandDstToSrcRowQueue_.EnQue(sortedRowIdx); + this->expandedExpertIdxCopyOutQueue_.EnQue(expandedExpertIdx); + } else { + LocalTensor xLocal = this->xCopyInQueue_.template AllocTensor(); + LocalTensor expandedRowIdx = this->expandedRowIdxCopyOutQueue_.template DeQue(); + int64_t inFactor = Align(this->cols_, sizeof(int8_t)); + uint32_t dstStride = (inFactor * sizeof(T) - AlignBytes(this->cols_, sizeof(T))) / BLOCK_BYTES; + DataCopyExtParams dataXCopyParams{static_cast(this->endXRow_ - this->startXRow_ + 1), + static_cast(this->cols_ * sizeof(T)), 0, dstStride, 0}; + DataCopyPad(xLocal, this->xGm_[this->startXRow_ * this->cols_], dataXCopyParams, {false, 0, 0, 0}); + this->xCopyInQueue_.EnQue(xLocal); + SetWaitFlag(HardEvent::MTE2_V); + ComputeQuant(this->endXRow_ - this->startXRow_ + 1); + + LocalTensor outLocal = inputXOutQueue_.DeQue(); + int64_t k = 0; + DataCopyExtParams intriParams{1, static_cast(this->cols_ * sizeof(int8_t)), 0, 0, 0}; + for (int64_t i = this->startXRow_; i <= this->endXRow_; i++) { + for (; k < this->coreIndicesElements_ && curIndex / this->k_ == i; curIndex++, k++) { + int32_t outIndex = expandedRowIdx.GetValue(curIndex); + if (outIndex < this->activeNum_) { + DataCopyPad(this->expandedXGm_[outIndex * this->cols_], outLocal[(i - this->startXRow_) * inFactor], + intriParams); + } + } + } + inputXOutQueue_.FreeTensor(outLocal); + this->expandedRowIdxCopyOutQueue_.EnQue(expandedRowIdx); + } +} + +template +__aicore__ inline void MoeCustomFullLoadStaticQuant::FreeLocalTensor() +{ + if (!this->ep_) { + LocalTensor expandedRowIdx = this->expandedRowIdxCopyOutQueue_.template DeQue(); + this->expandedRowIdxCopyOutQueue_.FreeTensor(expandedRowIdx); + } + LocalTensor expandedExpertIdx = this->expandedExpertIdxCopyOutQueue_.template DeQue(); + this->expandedExpertIdxCopyOutQueue_.FreeTensor(expandedExpertIdx); + LocalTensor sortedRowIdx = this->expandDstToSrcRowQueue_.template DeQue(); + this->expandDstToSrcRowQueue_.FreeTensor(sortedRowIdx); +} + +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_FULL_LOAD_STATIC_QUANT_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load_unquantized.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load_unquantized.h new file mode 100644 index 000000000..99cd28ac2 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_full_load_unquantized.h @@ -0,0 +1,224 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_full_load_unquantized.h + * \brief + */ +#ifndef MOE_CUSTOM_FULL_LOAD_UNQUANTIZED_H +#define MOE_CUSTOM_FULL_LOAD_UNQUANTIZED_H + +#include "moe_custom_full_load_base.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +template +class MoeCustomFullLoadUnquantized : public MoeCustomFullLoadBase { +public: + __aicore__ inline MoeCustomFullLoadUnquantized(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR expandedX, GM_ADDR expandedRowIdx, + GM_ADDR expertTokensCountOrCumsum, GM_ADDR expandedScale, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +protected: + __aicore__ inline void FreeLocalTensor(); + __aicore__ inline void GatherOutX(); + __aicore__ inline void CopyOutScale(); + +protected: + TQue xCopyInQueue_; + TQue scaleCopyInQueue_; + + GlobalTensor xGm_; + GlobalTensor scaleGm_; + GlobalTensor expandedXGm_; + GlobalTensor expandedRowIdxGm_; + GlobalTensor expandedScaleGm_; +}; + +template +__aicore__ inline void MoeCustomFullLoadUnquantized::Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR expandedX, + GM_ADDR expandedRowIdx, GM_ADDR expertTokensCountOrCumsum, + GM_ADDR expandedScale, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + MoeCustomFullLoadBase::Init(expertIdx, expandedRowIdx, expertTokensCountOrCumsum, workspace, tilingData, tPipe); + xGm_.SetGlobalBuffer((__gm__ T *)x); + if (this->isInputScale_) { + scaleGm_.SetGlobalBuffer((__gm__ float *)scale); + expandedScaleGm_.SetGlobalBuffer((__gm__ float *)expandedScale); + } + + expandedXGm_.SetGlobalBuffer((__gm__ T *)expandedX); + int64_t buffSize = this->sortNum_ * sizeof(int32_t); + int64_t row_length = + (this->curIndexStart_ + this->coreIndicesElements_ - 1) / this->k_ - this->curIndexStart_ / this->k_ + 1; + + if (this->ep_) { + this->pipe_->InitBuffer(xCopyInQueue_, this->bufferNum_, AlignBytes(this->cols_, sizeof(T))); + } else { + this->pipe_->InitBuffer(xCopyInQueue_, this->bufferNum_, AlignBytes(this->cols_, sizeof(T)) * row_length); + } + this->pipe_->InitBuffer(scaleCopyInQueue_, 1, AlignBytes(1, sizeof(float))); +} + +template +__aicore__ inline void MoeCustomFullLoadUnquantized::Process() +{ + if (this->blockIdx_ < this->needCoreNum_) { + this->CopyIn(); + this->Compute(); + + // valid expert equal zero + if (this->needCoreNum_ < 1) { + if (this->blockIdx_ == 0) { + if (this->rowIdxType_ == GATHER) { + this->CopyOutDefaultGatherIdx(); + } + if (this->expertTokensNumFlag_ == 1) { + this->CopyOutDefaultTokenCountOrCumsum(); + } + } + return; + } + + if (this->blockIdx_ == 0) { + this->CopyOutIdx(); + } + + if (this->blockIdx_ == this->needCoreNum_ - 1 && this->expertTokensNumFlag_ == 1) { + this->ComputeExpertTokenCountOrCumsum(); + } + + if (this->blockIdx_ < this->needCoreNum_) { + this->GatherOutX(); + if (this->isInputScale_) { + this->CopyOutScale(); + } + } + + this->FreeLocalTensor(); + } +} + +template +__aicore__ inline void MoeCustomFullLoadUnquantized::GatherOutX() +{ + if (this->ep_) { + LocalTensor expandedExpertIdx = this->expandedExpertIdxCopyOutQueue_.template DeQue(); + LocalTensor expandDstToSrcRowLocal = this->expandDstToSrcRowQueue_.template DeQue(); + int64_t startRowIdx = this->blockIdx_ * this->perCoreIndicesElements_; + int64_t endRowIdx = startRowIdx + this->coreIndicesElements_; + LocalTensor xLocal = xCopyInQueue_.AllocTensor(); + DataCopyExtParams copyParams{static_cast(1), static_cast(this->cols_ * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + for (int64_t i = startRowIdx; i < endRowIdx && i < this->activeNum_; i++) { + int32_t curExpertId = expandedExpertIdx.GetValue(i); + if (curExpertId < this->expertStart_ || curExpertId >= this->expertEnd_) { + break; + } + int64_t rowIdx = expandDstToSrcRowLocal.GetValue(i); + int64_t srcOffset = rowIdx / this->k_ * this->cols_; + int64_t dstOffset = i * this->cols_; + SetWaitFlag(HardEvent::MTE3_MTE2); + DataCopyPad(xLocal, xGm_[srcOffset], copyParams, padParams); + SetWaitFlag(HardEvent::MTE2_MTE3); + DataCopyPad(expandedXGm_[dstOffset], xLocal, copyParams); + } + xCopyInQueue_.FreeTensor(xLocal); + this->expandedExpertIdxCopyOutQueue_.template EnQue(expandedExpertIdx); + this->expandDstToSrcRowQueue_.template EnQue(expandDstToSrcRowLocal); + } else { + LocalTensor xLocal = xCopyInQueue_.AllocTensor(); + DataCopyExtParams dataXCopyParams{static_cast(this->endXRow_ - this->startXRow_ + 1), + static_cast(this->cols_ * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataXCopyPadParams{false, 0, 0, 0}; + DataCopyPad(xLocal, xGm_[this->startXRow_ * this->cols_], dataXCopyParams, dataXCopyPadParams); + SetWaitFlag(HardEvent::MTE2_MTE3); + int64_t inFactor = Align(this->cols_, sizeof(T)); + DataCopyExtParams copyParams{static_cast(1), static_cast(this->cols_ * sizeof(T)), 0, 0, 0}; + LocalTensor expandedRowIdx = this->expandedRowIdxCopyOutQueue_.template DeQue(); + int64_t curIndexStart = this->curIndexStart_; + int64_t k = 0; + for (int64_t i = this->startXRow_; i <= this->endXRow_; i++) { + for (; k < this->coreIndicesElements_ && curIndexStart / this->k_ == i; curIndexStart++, k++) { + int32_t outIndex = expandedRowIdx.GetValue(curIndexStart); + if (outIndex < this->activeNum_) { + DataCopyPad(expandedXGm_[outIndex * this->cols_], xLocal[(i - this->startXRow_) * inFactor], + copyParams); + } + } + } + xCopyInQueue_.FreeTensor(xLocal); + this->expandedRowIdxCopyOutQueue_.template EnQue(expandedRowIdx); + } +} + +template +__aicore__ inline void MoeCustomFullLoadUnquantized::FreeLocalTensor() +{ + LocalTensor expandedExpertIdx = this->expandedExpertIdxCopyOutQueue_.template DeQue(); + LocalTensor expandDstToSrcRowLocal = this->expandDstToSrcRowQueue_.template DeQue(); + this->expandedExpertIdxCopyOutQueue_.FreeTensor(expandedExpertIdx); + this->expandDstToSrcRowQueue_.FreeTensor(expandDstToSrcRowLocal); + if (!this->ep_) { + LocalTensor expandedRowIdx = this->expandedRowIdxCopyOutQueue_.template DeQue(); + this->expandedRowIdxCopyOutQueue_.FreeTensor(expandedRowIdx); + } +} + +template +__aicore__ inline void MoeCustomFullLoadUnquantized::CopyOutScale() +{ + LocalTensor scaleLocal = scaleCopyInQueue_.AllocTensor(); + DataCopyExtParams copyParams{static_cast(1), static_cast(sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + if (this->ep_) { + LocalTensor expandedExpertIdx = this->expandedExpertIdxCopyOutQueue_.template DeQue(); + LocalTensor expandDstToSrcRowLocal = this->expandDstToSrcRowQueue_.template DeQue(); + int64_t startRowIdx = this->blockIdx_ * this->perCoreIndicesElements_; + int64_t endRowIdx = startRowIdx + this->coreIndicesElements_; + for (int64_t i = startRowIdx; i < endRowIdx && i < this->activeNum_; i++) { + int32_t curExpertId = expandedExpertIdx.GetValue(i); + if (curExpertId < this->expertStart_ || curExpertId >= this->expertEnd_) { + break; + } + int64_t rowIdx = expandDstToSrcRowLocal.GetValue(i); + SetWaitFlag(HardEvent::MTE3_MTE2); + DataCopyPad(scaleLocal, scaleGm_[rowIdx / this->k_], copyParams, padParams); + SetWaitFlag(HardEvent::MTE2_MTE3); + DataCopyPad(expandedScaleGm_[i], scaleLocal, copyParams); + } + this->expandedExpertIdxCopyOutQueue_.template EnQue(expandedExpertIdx); + this->expandDstToSrcRowQueue_.template EnQue(expandDstToSrcRowLocal); + } else { + LocalTensor expandedRowIdx = this->expandedRowIdxCopyOutQueue_.template DeQue(); + int64_t curIndexStart = this->curIndexStart_; + int64_t k = 0; + for (int64_t i = this->startXRow_; i <= this->endXRow_; i++) { + SetWaitFlag(HardEvent::MTE3_MTE2); + DataCopyPad(scaleLocal, scaleGm_[i], copyParams, padParams); + SetWaitFlag(HardEvent::MTE2_MTE3); + for (; k < this->coreIndicesElements_ && curIndexStart / this->k_ == i; curIndexStart++, k++) { + int32_t outIndex = expandedRowIdx.GetValue(curIndexStart); + if (outIndex < this->activeNum_) { + DataCopyPad(expandedScaleGm_[outIndex], scaleLocal, copyParams); + } + } + } + this->expandedRowIdxCopyOutQueue_.template EnQue(expandedRowIdx); + } + scaleCopyInQueue_.FreeTensor(scaleLocal); +} + +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_FULL_LOAD_UNQUANTIZED_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_droppad_static_quant.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_droppad_static_quant.h new file mode 100644 index 000000000..e0c5f00b8 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_droppad_static_quant.h @@ -0,0 +1,238 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_gather_droppad_static_quant.h + * \brief + */ +#ifndef MOE_CUSTOM_GATHER_DROPPAD_STATIC_QUANT_H +#define MOE_CUSTOM_GATHER_DROPPAD_STATIC_QUANT_H + +#include "moe_custom_common.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +constexpr int64_t GATHER_OUT_DROPPAD_QUANT_BUFFER_NUM = 2; + +template +class MoeGatherDroppadQuant { +public: + __aicore__ inline MoeGatherDroppadQuant(){}; + __aicore__ inline void Init(GM_ADDR inputX, GM_ADDR scale, GM_ADDR offset, GM_ADDR expandedRowIdx, + GM_ADDR expandedX, GM_ADDR workspace, const MoeInitRoutingCustomTilingData *tilingData, + TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyExpertIn(int64_t progress); + __aicore__ inline void Compute(); + __aicore__ inline void CopyXIn(int64_t xSrcOffset, int64_t curLoopCols); + __aicore__ inline void CopyOut(int64_t progress); + +private: + TPipe *pipe_; + TQue inputXCopyInQueue_; + TQue expandRowIdxCopyInQueue_; + TQue inputXCopyOutQueue_; + TQue floatQueue_; + TQue halfQueue_; + + GlobalTensor inputXGm_; + GlobalTensor expandedXGm_; + GlobalTensor expandedRowIdxGm_; + GlobalTensor scaleGm_; + GlobalTensor offsetGm_; + + const MoeCustomGatherOutComputeTilingData *gatherOutTilingData_; + + int64_t needCoreNum_; + int64_t blockIdx_; + int64_t cols_; + int64_t n_; + int64_t k_; + int64_t currentLoopRows_; + int64_t coreRows_; + int64_t perLoopRows_; + int64_t lastLoopRows_; + int64_t rowLoops_; + int64_t colsTileLength_; + int64_t perLoopCols_; + int64_t lastLoopCols_; + int64_t colLoops_; + float scale_; + float offset_; + + int64_t indicesOffset_; + int64_t inputOffset_; + int64_t outOffset_; +}; + +template +__aicore__ inline void MoeGatherDroppadQuant::CopyExpertIn(int64_t progress) +{ + indicesOffset_ = progress * perLoopRows_; + LocalTensor indicesLocal = expandRowIdxCopyInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(currentLoopRows_ * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(indicesLocal, expandedRowIdxGm_[indicesOffset_], dataCopyParams, dataCopyPadParams); + expandRowIdxCopyInQueue_.EnQue(indicesLocal); +} + +template +__aicore__ inline void MoeGatherDroppadQuant::CopyXIn(int64_t xSrcOffset, int64_t curLoopCols) +{ + LocalTensor inLocal = inputXCopyInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{static_cast(1), static_cast(curLoopCols * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal, inputXGm_[xSrcOffset], dataCopyParams, dataCopyPadParams); + inputXCopyInQueue_.EnQue(inLocal); +} + +template +__aicore__ inline void MoeGatherDroppadQuant::Compute() +{ + LocalTensor floatLocal; + LocalTensor inLocal; + LocalTensor outLocal = inputXCopyOutQueue_.AllocTensor(); + LocalTensor halfLocal = halfQueue_.AllocTensor(); + uint32_t elements = Align(colsTileLength_, sizeof(T)); + if constexpr (IsSameType::value) { + floatLocal = inputXCopyInQueue_.DeQue(); + } else { + inLocal = inputXCopyInQueue_.DeQue(); + floatLocal = floatQueue_.AllocTensor(); + Cast(floatLocal, inLocal, RoundMode::CAST_NONE, elements); + PipeBarrier(); + } + Muls(floatLocal, floatLocal, scale_, elements); + PipeBarrier(); + Adds(floatLocal, floatLocal, offset_, elements); + PipeBarrier(); + LocalTensor intLocal = floatLocal.ReinterpretCast(); + Cast(intLocal, floatLocal, RoundMode::CAST_RINT, elements); + PipeBarrier(); + SetDeqScale((half)1.000000e+00f); + PipeBarrier(); + Cast(halfLocal, intLocal, RoundMode::CAST_ROUND, elements); + PipeBarrier(); + Cast(outLocal, halfLocal, RoundMode::CAST_TRUNC, elements); + inputXCopyOutQueue_.EnQue(outLocal); + if constexpr (IsSameType::value) { + inputXCopyInQueue_.FreeTensor(floatLocal); + } else { + inputXCopyInQueue_.FreeTensor(inLocal); + floatQueue_.FreeTensor(floatLocal); + } + halfQueue_.FreeTensor(halfLocal); +} + +template +__aicore__ inline void MoeGatherDroppadQuant::CopyOut(int64_t progress) +{ + LocalTensor indicesLocal = expandRowIdxCopyInQueue_.DeQue(); + SetWaitFlag(HardEvent::MTE2_S); + colsTileLength_ = perLoopCols_; + for (int64_t colsLoop = 0; colsLoop < colLoops_; colsLoop++) { + int64_t initialRow = gatherOutTilingData_->perCoreIndicesElements * blockIdx_ + perLoopRows_ * progress; + int64_t curLoopRow = 0; + if (colsLoop == colLoops_ - 1) { + colsTileLength_ = lastLoopCols_; + } + int64_t currentLoopStartRow = initialRow / k_; + int64_t currentLoopLastRow = (initialRow + currentLoopRows_ - 1) / k_; + for (int64_t row = currentLoopStartRow; row <= currentLoopLastRow; row++) { + inputOffset_ = row * cols_ + colsLoop * perLoopCols_; + // input row position + CopyXIn(inputOffset_, colsTileLength_); + Compute(); + LocalTensor outLocal = inputXCopyOutQueue_.DeQue(); + DataCopyExtParams intriParams{1, static_cast(colsTileLength_ * sizeof(int8_t)), 0, 0, 0}; + while (curLoopRow < currentLoopRows_ && initialRow / k_ == row) { + int32_t outIndex = indicesLocal.GetValue(curLoopRow); + curLoopRow++; + initialRow++; + if (outIndex == -1) { + continue; + } + outOffset_ = outIndex * cols_ + colsLoop * perLoopCols_; + DataCopyPad(expandedXGm_[outOffset_], outLocal, intriParams); + } + inputXCopyOutQueue_.FreeTensor(outLocal); + } + } + expandRowIdxCopyInQueue_.FreeTensor(indicesLocal); +} + +template +__aicore__ inline void MoeGatherDroppadQuant::Init(GM_ADDR inputX, GM_ADDR scale, GM_ADDR offset, + GM_ADDR expandedRowIdx, GM_ADDR expandedX, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + gatherOutTilingData_ = &(tilingData->gatherOutComputeParamsOp); + + needCoreNum_ = gatherOutTilingData_->needCoreNum; + cols_ = tilingData->cols; + n_ = tilingData->n; + k_ = tilingData->k; + + if (blockIdx_ == needCoreNum_ - 1) { + coreRows_ = gatherOutTilingData_->lastCoreIndicesElements; + perLoopRows_ = gatherOutTilingData_->lastCorePerLoopIndicesElements; + lastLoopRows_ = gatherOutTilingData_->lastCoreLastLoopIndicesElements; + rowLoops_ = gatherOutTilingData_->lastCoreIndicesLoops; + } else { + coreRows_ = gatherOutTilingData_->perCoreIndicesElements; + perLoopRows_ = gatherOutTilingData_->perCorePerLoopIndicesElements; + lastLoopRows_ = gatherOutTilingData_->perCoreLastLoopIndicesElements; + rowLoops_ = gatherOutTilingData_->perCoreIndicesLoops; + } + perLoopCols_ = gatherOutTilingData_->perLoopCols; + lastLoopCols_ = gatherOutTilingData_->lastLoopCols; + colLoops_ = gatherOutTilingData_->colsLoops; + + inputXGm_.SetGlobalBuffer((__gm__ T *)inputX); + expandedXGm_.SetGlobalBuffer((__gm__ int8_t *)expandedX); + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx + + blockIdx_ * gatherOutTilingData_->perCoreIndicesElements, + Align(coreRows_, sizeof(int32_t))); + scaleGm_.SetGlobalBuffer((__gm__ float *)scale, 1); + offsetGm_.SetGlobalBuffer((__gm__ float *)offset, 1); + scale_ = scaleGm_.GetValue(0); + offset_ = offsetGm_.GetValue(0); + + pipe_->InitBuffer(inputXCopyInQueue_, GATHER_OUT_DROPPAD_QUANT_BUFFER_NUM, AlignBytes(perLoopCols_, sizeof(T))); + pipe_->InitBuffer(inputXCopyOutQueue_, GATHER_OUT_DROPPAD_QUANT_BUFFER_NUM, + AlignBytes(perLoopCols_, sizeof(int8_t))); + pipe_->InitBuffer(expandRowIdxCopyInQueue_, GATHER_OUT_DROPPAD_QUANT_BUFFER_NUM, + AlignBytes(perLoopRows_, sizeof(int32_t))); + pipe_->InitBuffer(floatQueue_, 1, AlignBytes(perLoopCols_, sizeof(float))); + pipe_->InitBuffer(halfQueue_, 1, AlignBytes(perLoopCols_, sizeof(half))); +} + +template +__aicore__ inline void MoeGatherDroppadQuant::Process() +{ + if (blockIdx_ < needCoreNum_) { + currentLoopRows_ = perLoopRows_; + for (int64_t loop = 0; loop < rowLoops_; loop++) { + if (loop == rowLoops_ - 1) { + currentLoopRows_ = lastLoopRows_; + } + CopyExpertIn(loop); + CopyOut(loop); + } + } +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_GATHER_DROPPAD_STATIC_QUANT_H diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_dynamic_quant.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_dynamic_quant.h new file mode 100644 index 000000000..be6abc8ab --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_dynamic_quant.h @@ -0,0 +1,602 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_gather_dynamic_quant.h + * \brief + */ +#ifndef MOE_CUSTOM_GATHER_DYNAMIC_QUANT_H +#define MOE_CUSTOM_GATHER_DYNAMIC_QUANT_H + +#include "moe_custom_common.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; +constexpr int64_t GATHER_OUT_DYNAMIC_QUANT_BUFFER_NUM = 2; + +template +class MoeGatherOutDynamicQuant { +public: + __aicore__ inline MoeGatherOutDynamicQuant(){}; + __aicore__ inline void Init(GM_ADDR inputX, GM_ADDR quantSmooth, GM_ADDR expandedRowIdx, GM_ADDR expandedX, + GM_ADDR expandedScale, GM_ADDR sortedExpertIdx, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyOutXDynamicQuantFromGather(int64_t progress); + __aicore__ inline void CopyOutXDynamicQuantFromScatter(int64_t progress); + __aicore__ inline void CopyOutXPartialDynamicQuantFromGather(int64_t progress); + __aicore__ inline void CopyOutXPartialDynamicQuantFromScatter(int64_t progress); + __aicore__ inline void CopyInExpandedExpertIdx(int64_t progress); + __aicore__ inline void Compute(LocalTensor &smoothLocal); + __aicore__ inline float ComputeMax(LocalTensor &inLocal, LocalTensor &tempLocal, + LocalTensor &scaleLocal, int32_t srcIdx, int32_t expertIdx, int64_t j); + __aicore__ inline void ComputeScale(LocalTensor &inLocal, LocalTensor &tempLocal, float scaleTemp, + int64_t dstIndex, int64_t j); + +private: + TPipe *pipe_; + TQue inputXInQueue_; + TQue smoothInQueue_; + TQue expandRowIdxInQueue_; + TQue calcQueue_; + TQue inputXOutQueue_; + TQue scaleOutQueue_; + + GlobalTensor inputXGm_; + GlobalTensor expandedXGm_; + GlobalTensor expandedRowIdxGm_; + GlobalTensor quantSmoothGm_; + GlobalTensor expandedScaleGm_; + GlobalTensor quantTempGm_; + GlobalTensor expandedExpertIdxGm_; + GlobalTensor expertTotalCountGm_; + + const MoeCustomGatherOutComputeTilingData *gatherOutTilingData_; + + int64_t needCoreNum_; + int64_t blockIdx_; + int64_t cols_; + int64_t n_; + int64_t k_; + int64_t totalLength_; + int64_t perCoreRow_; + int64_t currentLoopRows_; + int64_t currentLoopRowsAlign_; + int64_t coreRows_; + int64_t perLoopRows_; + int64_t lastLoopRows_; + int64_t rowLoops_; + int64_t colsTileLength_; + int64_t perLoopCols_; + int64_t perLoopColsAlign_; + int64_t lastLoopCols_; + int64_t colLoops_; + int64_t isInputScale_; + int64_t expertStart_; + + int64_t indicesOffset_; + int64_t rowIdxType_ = 0; + int64_t dropPadMode_; + int64_t activeNum_; + int64_t ep_; + int64_t smoothType_; + int64_t coreNum_; + int64_t expertTotalCount_ = 0; +}; + +template +__aicore__ inline void MoeGatherOutDynamicQuant::CopyInExpandedExpertIdx(int64_t progress) +{ + indicesOffset_ = progress * perLoopRows_; + LocalTensor indicesLocal = expandRowIdxInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(currentLoopRows_ * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(indicesLocal, expandedRowIdxGm_[indicesOffset_], dataCopyParams, dataCopyPadParams); + DataCopyPad(indicesLocal[currentLoopRowsAlign_], expandedExpertIdxGm_[indicesOffset_], dataCopyParams, + dataCopyPadParams); + expandRowIdxInQueue_.EnQue(indicesLocal); +} + +template +__aicore__ inline void MoeGatherOutDynamicQuant::Compute(LocalTensor &smoothLocal) +{ + LocalTensor inLocal = inputXInQueue_.DeQue(); + + LocalTensor tempLocal = calcQueue_.AllocTensor(); + LocalTensor outLocal = inputXOutQueue_.AllocTensor(); + LocalTensor scaleLocal = scaleOutQueue_.AllocTensor(); + + if constexpr (!IsSameType::value) { + Cast(inLocal, inLocal.ReinterpretCast()[perLoopColsAlign_], RoundMode::CAST_NONE, cols_); + PipeBarrier(); + } + + if (isInputScale_) { + Mul(inLocal, inLocal, smoothLocal, cols_); + PipeBarrier(); + } + + Abs(tempLocal, inLocal, cols_); + PipeBarrier(); + + ReduceMax(scaleLocal, tempLocal, tempLocal, cols_); // get max value and index [0,1] + + float scaleValue = scaleLocal.GetValue(0) / MAX_INT8; + + Duplicate(scaleLocal, scaleValue, INT32_ONE_BLOCK_NUM); + PipeBarrier(); + Duplicate(tempLocal, scaleValue, cols_); + PipeBarrier(); + + Div(tempLocal, inLocal, tempLocal, cols_); + PipeBarrier(); + + LocalTensor intLocal = tempLocal.ReinterpretCast(); + Cast(intLocal, tempLocal, RoundMode::CAST_RINT, cols_); + PipeBarrier(); + SetDeqScale((half)1.000000e+00f); + Cast(intLocal.ReinterpretCast(), intLocal, RoundMode::CAST_ROUND, cols_); + PipeBarrier(); + Cast(outLocal, intLocal.ReinterpretCast(), RoundMode::CAST_TRUNC, cols_); + + calcQueue_.FreeTensor(tempLocal); + inputXOutQueue_.EnQue(outLocal); + scaleOutQueue_.EnQue(scaleLocal); +} + +template +__aicore__ inline void MoeGatherOutDynamicQuant::CopyOutXDynamicQuantFromScatter(int64_t progress) +{ + DataCopyExtParams copyInParams{1, static_cast(perLoopCols_ * sizeof(T)), 0, 0, 0}; + DataCopyExtParams smoothParams{1, static_cast(perLoopCols_ * sizeof(float)), 0, 0, 0}; + DataCopyExtParams copyOutParams{1, static_cast(perLoopCols_ * sizeof(int8_t)), 0, 0, 0}; + DataCopyExtParams quantScaleParams{1, static_cast(sizeof(int32_t)), 0, 0, 0}; + LocalTensor indicesLocal = expandRowIdxInQueue_.DeQue(); + LocalTensor smoothLocal = smoothInQueue_.AllocTensor(); + + // copyin [1,H] scale + if (smoothType_ == SCALE_1H) { + DataCopyPad(smoothLocal, quantSmoothGm_, smoothParams, {false, 0, 0, 0}); + smoothInQueue_.EnQue(smoothLocal); + smoothLocal = smoothInQueue_.DeQue(); + } + + int32_t lastExpertIdx = -1; + for (int64_t i = 0; i < currentLoopRows_; i++) { + int64_t rowOffset = perCoreRow_ * blockIdx_ + perLoopRows_ * progress; + if (dropPadMode_ == DROPLESS_MODE && (rowOffset + i) >= activeNum_) { + break; + } + LocalTensor inLocal = inputXInQueue_.AllocTensor(); + int32_t srcIdx = indicesLocal.GetValue(i); + + int32_t expertIdx = indicesLocal.GetValue(currentLoopRowsAlign_ + i) - expertStart_; + if constexpr (IsSameType::value) { + DataCopyPad(inLocal, inputXGm_[srcIdx / k_ * cols_], copyInParams, {false, 0, 0, 0}); + } else { + DataCopyPad(inLocal[perLoopColsAlign_], inputXGm_[srcIdx / k_ * cols_], copyInParams, {false, 0, 0, 0}); + } + inputXInQueue_.EnQue(inLocal); + + // copyin dynamic scale + if (smoothType_ == SCALE_EH && expertIdx != lastExpertIdx) { + DataCopyPad(smoothLocal, quantSmoothGm_[expertIdx * this->cols_], smoothParams, {false, 0, 0, 0}); + smoothInQueue_.EnQue(smoothLocal); + smoothLocal = smoothInQueue_.DeQue(); + lastExpertIdx = expertIdx; + } + Compute(smoothLocal); + inputXInQueue_.FreeTensor(inLocal); + LocalTensor scaleLocal = scaleOutQueue_.DeQue(); + DataCopyPad(expandedScaleGm_[(rowOffset + i)], scaleLocal, quantScaleParams); + LocalTensor outLocal = inputXOutQueue_.DeQue(); + DataCopyPad(expandedXGm_[(rowOffset + i) * cols_], outLocal, copyOutParams); + + inputXOutQueue_.FreeTensor(outLocal); + scaleOutQueue_.FreeTensor(scaleLocal); + } + + smoothInQueue_.FreeTensor(smoothLocal); + expandRowIdxInQueue_.FreeTensor(indicesLocal); +} + +template +__aicore__ inline void MoeGatherOutDynamicQuant::CopyOutXDynamicQuantFromGather(int64_t progress) +{ + DataCopyExtParams copyInParams{1, static_cast(perLoopCols_ * sizeof(T)), 0, 0, 0}; + DataCopyExtParams smoothParams{1, static_cast(perLoopCols_ * sizeof(float)), 0, 0, 0}; + DataCopyExtParams copyOutParams{1, static_cast(perLoopCols_ * sizeof(int8_t)), 0, 0, 0}; + DataCopyExtParams quantScaleParams{1, static_cast(sizeof(int32_t)), 0, 0, 0}; + + LocalTensor indicesLocal = expandRowIdxInQueue_.DeQue(); + LocalTensor smoothLocal = smoothInQueue_.AllocTensor(); + + int64_t rowOffset = blockIdx_ * perCoreRow_ + progress * perLoopRows_; + int64_t startXRow = rowOffset / k_; + int64_t endXRow = (rowOffset + currentLoopRows_ - 1) / k_; + int64_t curIndex = 0; + + if (smoothType_ == SCALE_1H) { + DataCopyPad(smoothLocal, quantSmoothGm_, smoothParams, {false, 0, 0, 0}); + smoothInQueue_.EnQue(smoothLocal); + smoothLocal = smoothInQueue_.DeQue(); + } + + for (int64_t row = startXRow; row <= endXRow; row++) { + LocalTensor inLocal = inputXInQueue_.AllocTensor(); + if constexpr (IsSameType::value) { + DataCopyPad(inLocal, inputXGm_[row * cols_], copyInParams, {false, 0, 0, 0}); + } else { + DataCopyPad(inLocal[perLoopColsAlign_], inputXGm_[row * cols_], copyInParams, {false, 0, 0, 0}); + } + inputXInQueue_.EnQue(inLocal); + Compute(smoothLocal); + LocalTensor scaleLocal = scaleOutQueue_.DeQue(); + LocalTensor outLocal = inputXOutQueue_.DeQue(); + + while (curIndex < currentLoopRows_ && (rowOffset + curIndex) / this->k_ == row) { + int32_t outIndex = indicesLocal.GetValue(curIndex); + curIndex++; + if (outIndex == -1 || dropPadMode_ == DROPLESS_MODE && outIndex >= this->activeNum_) { + continue; + } + DataCopyPad(expandedXGm_[outIndex * cols_], outLocal, copyOutParams); + DataCopyPad(expandedScaleGm_[outIndex], scaleLocal, quantScaleParams); + } + + inputXInQueue_.FreeTensor(inLocal); + inputXOutQueue_.FreeTensor(outLocal); + scaleOutQueue_.FreeTensor(scaleLocal); + } + + smoothInQueue_.FreeTensor(smoothLocal); + expandRowIdxInQueue_.FreeTensor(indicesLocal); +} + +template +__aicore__ inline float +MoeGatherOutDynamicQuant::ComputeMax(LocalTensor &inLocal, LocalTensor &tempLocal, + LocalTensor &scaleLocal, int32_t srcIdx, int32_t expertIdx, + int64_t j) +{ + LocalTensor smoothLocal = smoothInQueue_.AllocTensor(); + + DataCopyExtParams intriParamsT{1, static_cast(colsTileLength_ * sizeof(T)), 0, 0, 0}; + DataCopyExtParams intriParamsFp32{1, static_cast(colsTileLength_ * sizeof(float)), 0, 0, 0}; + + if constexpr (!IsSameType::value) { + DataCopyPad(inLocal.ReinterpretCast()[perLoopColsAlign_], inputXGm_[srcIdx * cols_ + j * perLoopCols_], + intriParamsT, {false, 0, 0, 0}); + } else { + DataCopyPad(inLocal, inputXGm_[srcIdx * cols_ + j * perLoopCols_], intriParamsT, {false, 0, 0, 0}); + } + + inputXInQueue_.EnQue(inLocal); + inLocal = inputXInQueue_.DeQue(); + + if (isInputScale_) { + DataCopyPad(smoothLocal, quantSmoothGm_[expertIdx * cols_ + j * perLoopCols_], intriParamsFp32, + {false, 0, 0, 0}); + smoothInQueue_.EnQue(smoothLocal); + smoothLocal = smoothInQueue_.DeQue(); + } + + if constexpr (!IsSameType::value) { + Cast(inLocal, inLocal.ReinterpretCast()[perLoopColsAlign_], RoundMode::CAST_NONE, colsTileLength_); + PipeBarrier(); + } + + if (isInputScale_) { + Mul(inLocal, inLocal, smoothLocal, colsTileLength_); + PipeBarrier(); + } + + Abs(tempLocal, inLocal, colsTileLength_); + PipeBarrier(); + + ReduceMax(scaleLocal[INT32_ONE_BLOCK_NUM], tempLocal, tempLocal, colsTileLength_); + + DataCopyPad(quantTempGm_[j * perLoopCols_], inLocal, intriParamsFp32); + smoothInQueue_.FreeTensor(smoothLocal); + SetWaitFlag(HardEvent::MTE3_MTE2); + return scaleLocal.GetValue(INT32_ONE_BLOCK_NUM); +} + +template +__aicore__ inline void +MoeGatherOutDynamicQuant::ComputeScale(LocalTensor &inLocal, LocalTensor &tempLocal, + float scaleTemp, int64_t dstIndex, int64_t j) +{ + DataCopyExtParams copyInParams{1, static_cast(colsTileLength_ * sizeof(float)), 0, 0, 0}; + DataCopyExtParams copyOutParams{1, static_cast(colsTileLength_ * sizeof(int8_t)), 0, 0, 0}; + + LocalTensor outLocal = inputXOutQueue_.AllocTensor(); + + DataCopyPad(inLocal, quantTempGm_[j * perLoopCols_], copyInParams, {false, 0, 0, 0}); + inputXInQueue_.EnQue(inLocal); + inLocal = inputXInQueue_.DeQue(); + + Duplicate(tempLocal, scaleTemp, colsTileLength_); + PipeBarrier(); + + Div(tempLocal, inLocal, tempLocal, colsTileLength_); + PipeBarrier(); + + Cast(tempLocal.ReinterpretCast(), tempLocal, RoundMode::CAST_TRUNC, colsTileLength_); + PipeBarrier(); + + Cast(outLocal, tempLocal.ReinterpretCast(), RoundMode::CAST_ROUND, colsTileLength_); + + inputXOutQueue_.EnQue(outLocal); + outLocal = inputXOutQueue_.DeQue(); + DataCopyPad(expandedXGm_[dstIndex * cols_ + j * perLoopCols_], outLocal, copyOutParams); + + inputXOutQueue_.FreeTensor(outLocal); + SetWaitFlag(HardEvent::MTE3_MTE2); +} + +template +__aicore__ inline void +MoeGatherOutDynamicQuant::CopyOutXPartialDynamicQuantFromScatter(int64_t progress) +{ + LocalTensor indicesLocal = expandRowIdxInQueue_.DeQue(); + for (int64_t i = 0; i < currentLoopRows_; i++) { + int64_t rowOffset = perCoreRow_ * blockIdx_ + perLoopRows_ * progress; + if (dropPadMode_ == DROPLESS_MODE && (rowOffset + i) >= activeNum_) { + break; + } + int32_t srcIdx = indicesLocal.GetValue(i); + int32_t expertIdx = indicesLocal.GetValue(currentLoopRowsAlign_ + i) - expertStart_; + LocalTensor inLocal = inputXInQueue_.AllocTensor(); + LocalTensor tempLocal = calcQueue_.AllocTensor(); + LocalTensor scaleLocal = scaleOutQueue_.AllocTensor(); + + float tileMax; + float reduceMax = *((float *)&INF); + for (int64_t j = 0; j < colLoops_; j++) { + colsTileLength_ = perLoopCols_; + if (j == colLoops_ - 1) { + colsTileLength_ = lastLoopCols_; + } + + if (smoothType_ == SCALE_1H) { + // 1H + tileMax = ComputeMax(inLocal, tempLocal, scaleLocal, srcIdx / k_, 0, j); + } else { + // EH + tileMax = ComputeMax(inLocal, tempLocal, scaleLocal, srcIdx / k_, expertIdx, j); + } + reduceMax = (reduceMax > tileMax) ? reduceMax : tileMax; + } + + float scaleTemp = reduceMax / MAX_INT8; + Duplicate(scaleLocal, scaleTemp, INT32_ONE_BLOCK_NUM); + scaleOutQueue_.EnQue(scaleLocal); + scaleLocal = scaleOutQueue_.DeQue(); + + DataCopyPad(expandedScaleGm_[(rowOffset + i)], scaleLocal, {1, 4, 0, 0, 0}); + + for (int64_t j = 0; j < colLoops_; j++) { + colsTileLength_ = perLoopCols_; + if (j == colLoops_ - 1) { + colsTileLength_ = lastLoopCols_; + } + ComputeScale(inLocal, tempLocal, scaleTemp, rowOffset + i, j); + } + inputXInQueue_.FreeTensor(inLocal); + calcQueue_.FreeTensor(tempLocal); + scaleOutQueue_.FreeTensor(scaleLocal); + } + expandRowIdxInQueue_.FreeTensor(indicesLocal); +} + +template +__aicore__ inline void MoeGatherOutDynamicQuant::CopyOutXPartialDynamicQuantFromGather(int64_t progress) +{ + LocalTensor indicesLocal = expandRowIdxInQueue_.DeQue(); + int64_t rowOffset = blockIdx_ * perCoreRow_ + progress * perLoopRows_; + int64_t startXRow = rowOffset / k_; + int64_t endXRow = (rowOffset + currentLoopRows_ - 1) / k_; + int64_t curIndex = 0; + + DataCopyExtParams quantScaleParams{1, static_cast(sizeof(int32_t)), 0, 0, 0}; + + for (int64_t row = startXRow; row <= endXRow; row++) { + LocalTensor inLocal = inputXInQueue_.AllocTensor(); + LocalTensor tempLocal = calcQueue_.AllocTensor(); + LocalTensor quantScaleLocal = scaleOutQueue_.AllocTensor(); + + float reduceMax = *((float *)&INF); + for (int64_t j = 0; j < colLoops_; j++) { + colsTileLength_ = perLoopCols_; + if (j == colLoops_ - 1) { + colsTileLength_ = lastLoopCols_; + } + + float tileMax = ComputeMax(inLocal, tempLocal, quantScaleLocal, row, 0, j); + reduceMax = (reduceMax > tileMax) ? reduceMax : tileMax; + } + + float scaleTemp = reduceMax / MAX_INT8; + Duplicate(quantScaleLocal, scaleTemp, INT32_ONE_BLOCK_NUM); + scaleOutQueue_.EnQue(quantScaleLocal); + quantScaleLocal = scaleOutQueue_.DeQue(); + + while (curIndex < currentLoopRows_ && (curIndex + rowOffset) / k_ == row) { + int32_t outIndex = indicesLocal.GetValue(curIndex); + curIndex++; + if (outIndex == -1 || (dropPadMode_ == DROPLESS_MODE && outIndex >= activeNum_)) { + continue; + } + DataCopyPad(expandedScaleGm_[outIndex], quantScaleLocal, quantScaleParams); + for (int64_t j = 0; j < colLoops_; j++) { + colsTileLength_ = perLoopCols_; + if (j == colLoops_ - 1) { + colsTileLength_ = lastLoopCols_; + } + ComputeScale(inLocal, tempLocal, scaleTemp, outIndex, j); + } + } + inputXInQueue_.FreeTensor(inLocal); + calcQueue_.FreeTensor(tempLocal); + scaleOutQueue_.FreeTensor(quantScaleLocal); + } + expandRowIdxInQueue_.FreeTensor(indicesLocal); +} + +template +__aicore__ inline void +MoeGatherOutDynamicQuant::Init(GM_ADDR inputX, GM_ADDR quantSmooth, GM_ADDR sortedExpertIdx, + GM_ADDR expandedRowIdx, GM_ADDR expandedX, GM_ADDR expandedScale, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + gatherOutTilingData_ = &(tilingData->gatherOutComputeParamsOp); + cols_ = tilingData->cols; + n_ = tilingData->n; + k_ = tilingData->k; + totalLength_ = n_ * k_; + isInputScale_ = tilingData->isInputScale; + expertStart_ = tilingData->expertStart; + rowIdxType_ = tilingData->rowIdxType; + dropPadMode_ = tilingData->dropPadMode; + activeNum_ = tilingData->activeNum; + ep_ = tilingData->ep; + smoothType_ = tilingData->smoothType; + coreNum_ = tilingData->coreNum; + + // core split + int64_t actualExpertNum_ = tilingData->actualExpertNum; + if (ep_) { + expertTotalCountGm_.SetGlobalBuffer((__gm__ int32_t *)sortedExpertIdx + Align(n_ * k_, sizeof(int32_t)) * 2 + + Align(actualExpertNum_, sizeof(int32_t)), + 1); + AscendC::DataCacheCleanAndInvalid(expertTotalCountGm_); + expertTotalCount_ = expertTotalCountGm_.GetValue(0); + } else { + expertTotalCount_ = totalLength_; + } + + perCoreRow_ = Ceil(expertTotalCount_, tilingData->coreNum); + needCoreNum_ = Ceil(expertTotalCount_, perCoreRow_); + int64_t lastCoreIndicesElements = expertTotalCount_ - (needCoreNum_ - 1) * perCoreRow_; + + // inner core split + int64_t originPerLoopElements; + if (blockIdx_ == needCoreNum_ - 1) { + coreRows_ = lastCoreIndicesElements; + originPerLoopElements = gatherOutTilingData_->lastCorePerLoopIndicesElements; + } else { + coreRows_ = perCoreRow_; + originPerLoopElements = gatherOutTilingData_->perCorePerLoopIndicesElements; + } + perLoopRows_ = Min(coreRows_, originPerLoopElements); + rowLoops_ = Ceil(coreRows_, perLoopRows_); + lastLoopRows_ = coreRows_ - (rowLoops_ - 1) * perLoopRows_; + + // cols split + perLoopCols_ = gatherOutTilingData_->perLoopCols; + lastLoopCols_ = gatherOutTilingData_->lastLoopCols; + colLoops_ = gatherOutTilingData_->colsLoops; + + perLoopColsAlign_ = Align(perLoopCols_, sizeof(T)); + + inputXGm_.SetGlobalBuffer((__gm__ T *)inputX); + expandedXGm_.SetGlobalBuffer((__gm__ int8_t *)expandedX); + + expandedExpertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)sortedExpertIdx + blockIdx_ * perCoreRow_, + Align(coreRows_, sizeof(int32_t))); + + if constexpr (COPYOUTTYPE == SCATTER) { + if (rowIdxType_ == SCATTER) { + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx + blockIdx_ * perCoreRow_, + Align(perCoreRow_, sizeof(int32_t))); + } else { + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)sortedExpertIdx + Align(n_ * k_, sizeof(int32_t)) + + blockIdx_ * perCoreRow_, + Align(perCoreRow_, sizeof(int32_t))); + } + } else { + if (rowIdxType_ == GATHER) { + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx + blockIdx_ * perCoreRow_, + Align(perCoreRow_, sizeof(int32_t))); + } else { + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)sortedExpertIdx + Align(n_ * k_, sizeof(int32_t)) + + blockIdx_ * perCoreRow_, + Align(perCoreRow_, sizeof(int32_t))); + } + } + + if (isInputScale_) { + quantSmoothGm_.SetGlobalBuffer((__gm__ float *)quantSmooth); + } + expandedScaleGm_.SetGlobalBuffer((__gm__ float *)expandedScale); + + if (colLoops_ > 1) { + quantTempGm_.SetGlobalBuffer((__gm__ float *)sortedExpertIdx + Align(totalLength_, sizeof(int32_t)) * 2 + + Align(actualExpertNum_, sizeof(int32_t)) * 2 + + Align(totalLength_, sizeof(int32_t)) + blockIdx_ * cols_, + cols_ * sizeof(float)); + } + + currentLoopRowsAlign_ = Align(perLoopRows_, sizeof(int32_t)); + + int64_t perLoopColsAlignBytes = AlignBytes(this->perLoopCols_, sizeof(T)); + perLoopColsAlignBytes = + Max(int64_t(perLoopColsAlignBytes * sizeof(float) / sizeof(T)), int64_t(BLOCK_BYTES + BLOCK_BYTES)); + pipe_->InitBuffer(expandRowIdxInQueue_, GATHER_OUT_DYNAMIC_QUANT_BUFFER_NUM, + 2 * AlignBytes(perLoopRows_, sizeof(int32_t))); + pipe_->InitBuffer(inputXInQueue_, GATHER_OUT_DYNAMIC_QUANT_BUFFER_NUM, perLoopColsAlignBytes); // percols * 2 * 4 + pipe_->InitBuffer(smoothInQueue_, GATHER_OUT_DYNAMIC_QUANT_BUFFER_NUM, + AlignBytes(perLoopCols_, sizeof(float))); // percols * 2 * 4 + pipe_->InitBuffer(calcQueue_, 1, AlignBytes(perLoopCols_, sizeof(float))); // percols * 1 * 4 + pipe_->InitBuffer(inputXOutQueue_, 1, AlignBytes(perLoopCols_, sizeof(int8_t))); // percols * 1 + pipe_->InitBuffer(scaleOutQueue_, 1, BLOCK_BYTES + BLOCK_BYTES); // 32 + 32 +} + +template +__aicore__ inline void MoeGatherOutDynamicQuant::Process() +{ + if (blockIdx_ < needCoreNum_) { + currentLoopRows_ = perLoopRows_; + if (colLoops_ > 1) { + for (int64_t loop = 0; loop < rowLoops_; loop++) { + if (loop == rowLoops_ - 1) { + currentLoopRows_ = lastLoopRows_; + } + CopyInExpandedExpertIdx(loop); + if constexpr (COPYOUTTYPE == GATHER) { + CopyOutXPartialDynamicQuantFromGather(loop); + } else { + CopyOutXPartialDynamicQuantFromScatter(loop); + } + } + } else { + for (int64_t loop = 0; loop < rowLoops_; loop++) { + if (loop == rowLoops_ - 1) { + currentLoopRows_ = lastLoopRows_; + } + CopyInExpandedExpertIdx(loop); + if constexpr (COPYOUTTYPE == GATHER) { + CopyOutXDynamicQuantFromGather(loop); + } else { + CopyOutXDynamicQuantFromScatter(loop); + } + } + } + } +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_GATHER_DYNAMIC_QUANT_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_out.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_out.h new file mode 100644 index 000000000..0ba44f76f --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_out.h @@ -0,0 +1,321 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_gather_out.h + * \brief + */ +#ifndef MOE_CUSTOM_GATHER_OUT_H +#define MOE_CUSTOM_GATHER_OUT_H + +#include "moe_custom_common.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +constexpr int64_t GATHER_OUT_BUFFER_NUM = 2; + +template +class MoeGatherOut { +public: + __aicore__ inline MoeGatherOut(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR scale, GM_ADDR workspace, GM_ADDR expandedRowIdx, GM_ADDR expandedX, + GM_ADDR expandedScale, const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + __aicore__ inline void CopyExpertIn(int64_t progress); + __aicore__ inline void CopyXIn(int64_t xSrcOffset, int64_t curLoopCols); + __aicore__ inline void CopyXOut(int64_t xDstOffset, int64_t curLoopCols); + __aicore__ inline void CopyScaleIn(int64_t scaleSrcOffset); + __aicore__ inline void CopyScaleOut(int64_t scaleDstOffset); + __aicore__ inline void GatherCopyOut(int64_t progress); + __aicore__ inline void ScatterCopyOut(int64_t progress); + +private: + TPipe *pipe_; + TQueBind xCopyInQueue_; + TQueBind scaleCopyInQueue_; + TQue expandedRowIdxCopyInQueue_; + + GlobalTensor xGm_; + GlobalTensor xGscaleGm_; + GlobalTensor sortedExpertIdxGm_; + GlobalTensor expandedXGm_; + GlobalTensor expandedRowIdxGm_; + GlobalTensor expandedScaleGm_; + GlobalTensor expertTotalCountGm_; + + int64_t blockIdx_; + int64_t cols_; + int64_t n_; + int64_t k_; + int64_t activeNum_; + int64_t dropPadMode_; + + int64_t colsLoops_; + int64_t perLoopCols_; + int64_t lastLoopCols_; + + int64_t indicesLoops_; + int64_t curLoopElements_; + + int64_t perCoreIndicesElements_; + int64_t lastCoreIndicesElements_; + int64_t perCorePerLoopIndicesElements_; + int64_t lastCorePerLoopIndicesElements_; + int64_t curCorePerLoopIndicesElements_; + int64_t curCoreLastLoopIndicesElements_; + int64_t needCoreNum_; + int64_t curCoreIndicesElements_; + + int64_t actualExpertNum_; + int64_t expertTotalCount_; + + int64_t rowIdxType_; + int64_t isInputScale_; + int64_t coreNum_; +}; + +template +__aicore__ inline void MoeGatherOut::Init(GM_ADDR x, GM_ADDR scale, GM_ADDR workspace, GM_ADDR expandedRowIdx, + GM_ADDR expandedX, GM_ADDR expandedScale, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + + cols_ = tilingData->cols; + n_ = tilingData->n; + k_ = tilingData->k; + coreNum_ = tilingData->coreNum; + dropPadMode_ = tilingData->dropPadMode; + activeNum_ = tilingData->activeNum; + + isInputScale_ = tilingData->isInputScale; + rowIdxType_ = tilingData->rowIdxType; + + colsLoops_ = tilingData->gatherOutComputeParamsOp.colsLoops; + perLoopCols_ = tilingData->gatherOutComputeParamsOp.perLoopCols; + lastLoopCols_ = tilingData->gatherOutComputeParamsOp.lastLoopCols; + + actualExpertNum_ = tilingData->actualExpertNum; + + if constexpr (EP) { + expertTotalCountGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(n_ * k_, sizeof(int32_t)) * 2 + + Align(actualExpertNum_, sizeof(int32_t)), + 1); + AscendC::DataCacheCleanAndInvalid(expertTotalCountGm_); + expertTotalCount_ = expertTotalCountGm_.GetValue(0); + } else { + expertTotalCount_ = n_ * k_; + } + + perCorePerLoopIndicesElements_ = tilingData->gatherOutComputeParamsOp.perCorePerLoopIndicesElements; + lastCorePerLoopIndicesElements_ = tilingData->gatherOutComputeParamsOp.lastCorePerLoopIndicesElements; + perCoreIndicesElements_ = Ceil(expertTotalCount_, tilingData->coreNum); + needCoreNum_ = Ceil(expertTotalCount_, perCoreIndicesElements_); + lastCoreIndicesElements_ = expertTotalCount_ - (needCoreNum_ - 1) * perCoreIndicesElements_; + + if (blockIdx_ == needCoreNum_ - 1) { + curCoreIndicesElements_ = lastCoreIndicesElements_; + curCorePerLoopIndicesElements_ = Min(lastCorePerLoopIndicesElements_, curCoreIndicesElements_); + } else { + curCoreIndicesElements_ = perCoreIndicesElements_; + curCorePerLoopIndicesElements_ = Min(perCorePerLoopIndicesElements_, curCoreIndicesElements_); + } + indicesLoops_ = Ceil(curCoreIndicesElements_, curCorePerLoopIndicesElements_); + curCoreLastLoopIndicesElements_ = curCoreIndicesElements_ - (indicesLoops_ - 1) * curCorePerLoopIndicesElements_; + + xGm_.SetGlobalBuffer((__gm__ T *)x, n_ * cols_); + xGscaleGm_.SetGlobalBuffer((__gm__ float *)scale, n_); + + expandedXGm_.SetGlobalBuffer((__gm__ T *)expandedX); + expandedScaleGm_.SetGlobalBuffer((__gm__ float *)expandedScale); + + pipe_->InitBuffer(expandedRowIdxCopyInQueue_, GATHER_OUT_BUFFER_NUM, + AlignBytes(curCorePerLoopIndicesElements_, sizeof(int32_t))); + pipe_->InitBuffer(xCopyInQueue_, GATHER_OUT_BUFFER_NUM, AlignBytes(perLoopCols_, sizeof(T))); + pipe_->InitBuffer(scaleCopyInQueue_, GATHER_OUT_BUFFER_NUM, AlignBytes(1, sizeof(float))); + + sortedExpertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + blockIdx_ * perCoreIndicesElements_, + Align(curCoreIndicesElements_, sizeof(int32_t))); + + if constexpr (EP) { + if (rowIdxType_ == SCATTER) { + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx + blockIdx_ * perCoreIndicesElements_, + Align(curCoreIndicesElements_, sizeof(int32_t))); + } else { + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(n_ * k_, sizeof(int32_t)) + + blockIdx_ * perCoreIndicesElements_, + Align(curCoreIndicesElements_, sizeof(int32_t))); + } + } else { + if (rowIdxType_ == GATHER) { + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx + blockIdx_ * perCoreIndicesElements_, + Align(curCoreIndicesElements_, sizeof(int32_t))); + } else { + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(n_ * k_, sizeof(int32_t)) + + blockIdx_ * perCoreIndicesElements_, + Align(curCoreIndicesElements_, sizeof(int32_t))); + } + } +} + +template +__aicore__ inline void MoeGatherOut::CopyExpertIn(int64_t progress) +{ + LocalTensor subRowIdxLocal = expandedRowIdxCopyInQueue_.AllocTensor(); + DataCopyExtParams copyParams{1, static_cast(curLoopElements_ * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(subRowIdxLocal, expandedRowIdxGm_[progress * curCorePerLoopIndicesElements_], copyParams, padParams); + expandedRowIdxCopyInQueue_.EnQue(subRowIdxLocal); +} + +template +__aicore__ inline void MoeGatherOut::CopyXIn(int64_t xSrcOffset, int64_t curLoopCols) +{ + LocalTensor xLocal = xCopyInQueue_.AllocTensor(); + DataCopyExtParams copyParams0{static_cast(1), static_cast(curLoopCols * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams padParams0{false, 0, 0, 0}; + DataCopyPad(xLocal, xGm_[xSrcOffset], copyParams0, padParams0); + xCopyInQueue_.EnQue(xLocal); +} + +template +__aicore__ inline void MoeGatherOut::CopyXOut(int64_t xDstOffset, int64_t curLoopCols) +{ + LocalTensor xLocal = xCopyInQueue_.DeQue(); + DataCopyExtParams copyParams2{1, static_cast(curLoopCols * sizeof(T)), 0, 0, 0}; + DataCopyPad(expandedXGm_[xDstOffset], xLocal, copyParams2); + xCopyInQueue_.FreeTensor(xLocal); +} + +template +__aicore__ inline void MoeGatherOut::CopyScaleIn(int64_t scaleSrcOffset) +{ + LocalTensor scaleLocal = scaleCopyInQueue_.AllocTensor(); + DataCopyExtParams copyParams1{static_cast(1), static_cast(1 * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams padParams1{false, 0, 0, 0}; + DataCopyPad(scaleLocal, xGscaleGm_[scaleSrcOffset], copyParams1, padParams1); + scaleCopyInQueue_.EnQue(scaleLocal); +} + +template +__aicore__ inline void MoeGatherOut::CopyScaleOut(int64_t scaleDstOffset) +{ + LocalTensor scaleLocal = scaleCopyInQueue_.DeQue(); + DataCopyExtParams copyParams3{1, static_cast(sizeof(float)), 0, 0, 0}; + DataCopyPad(expandedScaleGm_[scaleDstOffset], scaleLocal, copyParams3); + scaleCopyInQueue_.FreeTensor(scaleLocal); +} + +template +__aicore__ inline void MoeGatherOut::GatherCopyOut(int64_t progress) +{ + LocalTensor subRowIdxLocal = expandedRowIdxCopyInQueue_.DeQue(); + SetWaitFlag(HardEvent::MTE2_S); + int64_t curLoopCols = perLoopCols_; + for (int64_t colsLoop = 0; colsLoop < colsLoops_; colsLoop++) { + int64_t initialRow = blockIdx_ * perCoreIndicesElements_ + curCorePerLoopIndicesElements_ * progress; + int64_t curLoopRow = 0; + if (colsLoop == colsLoops_ - 1) { + curLoopCols = lastLoopCols_; + } + int64_t currentLoopStartRow = initialRow / k_; + int64_t currentLoopLastRow = (initialRow + this->curLoopElements_ - 1) / k_; + for (int64_t row = currentLoopStartRow; row <= currentLoopLastRow; row++) { + LocalTensor inLocal = xCopyInQueue_.AllocTensor(); + int64_t inputOffset = row * cols_ + colsLoop * perLoopCols_; + DataCopyExtParams xCopyParams{1, static_cast(curLoopCols * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal, xGm_[inputOffset], xCopyParams, dataCopyPadParams); + // copy in scale + LocalTensor scaleLocal = scaleCopyInQueue_.AllocTensor(); + DataCopyExtParams scaleCopyParams{1, static_cast(sizeof(float)), 0, 0, 0}; + if (isInputScale_ == 1 && colsLoop == 0) { + DataCopyPadExtParams scalePadParams{false, 0, 0, 0}; + DataCopyPad(scaleLocal, xGscaleGm_[row], scaleCopyParams, scalePadParams); + } + SetWaitFlag(HardEvent::MTE2_MTE3); + DataCopyExtParams intriParams{1, static_cast(curLoopCols * sizeof(T)), 0, 0, 0}; + while (curLoopRow < this->curLoopElements_ && initialRow / k_ == row) { + int32_t outIndex = subRowIdxLocal.GetValue(curLoopRow); + curLoopRow++; + initialRow++; + if (outIndex == -1 || (dropPadMode_ == DROPLESS_MODE && outIndex >= activeNum_)) { + continue; + } + int64_t outOffset = outIndex * this->cols_ + colsLoop * this->perLoopCols_; + DataCopyPad(expandedXGm_[outOffset], inLocal, intriParams); + // copy out scale + if (isInputScale_ == 1 && colsLoop == 0) { + DataCopyPad(expandedScaleGm_[outIndex], scaleLocal, scaleCopyParams); + } + } + scaleCopyInQueue_.FreeTensor(scaleLocal); + xCopyInQueue_.FreeTensor(inLocal); + } + } + expandedRowIdxCopyInQueue_.FreeTensor(subRowIdxLocal); +} + +template +__aicore__ inline void MoeGatherOut::ScatterCopyOut(int64_t progress) +{ + int64_t curExpertLoopOffset = progress * curCorePerLoopIndicesElements_; + LocalTensor subRowIdxLocal = expandedRowIdxCopyInQueue_.DeQue(); + for (int64_t indicesIndex = 0; indicesIndex < curLoopElements_; indicesIndex++) { + int64_t rowIdx = subRowIdxLocal.GetValue(indicesIndex); + int64_t rowOffset = curExpertLoopOffset + indicesIndex + blockIdx_ * perCoreIndicesElements_; + if (activeNum_ > 0 && dropPadMode_ == DROPLESS_MODE && rowOffset >= activeNum_) { + break; + } + SetWaitFlag(HardEvent::S_MTE2); + if (isInputScale_ == 1) { + int64_t scaleSrcOffset = rowIdx / k_; + CopyScaleIn(scaleSrcOffset); + CopyScaleOut(indicesIndex + curExpertLoopOffset + blockIdx_ * perCoreIndicesElements_); + } + int64_t curLoopCols = perLoopCols_; + for (int64_t colsLoop = 0; colsLoop < colsLoops_; colsLoop++) { + if (colsLoop == colsLoops_ - 1) { + curLoopCols = lastLoopCols_; + } + int64_t xSrcOffset = rowIdx / k_ * cols_; + int64_t xDstOffset = (blockIdx_ * perCoreIndicesElements_ + curExpertLoopOffset + indicesIndex) * cols_; + int64_t colsLoopOffset = colsLoop * perLoopCols_; + CopyXIn(xSrcOffset + colsLoopOffset, curLoopCols); + CopyXOut(xDstOffset + colsLoopOffset, curLoopCols); + } + } + expandedRowIdxCopyInQueue_.FreeTensor(subRowIdxLocal); +} + +template +__aicore__ inline void MoeGatherOut::Process() +{ + if (blockIdx_ < needCoreNum_) { + curLoopElements_ = curCorePerLoopIndicesElements_; + for (int64_t loop = 0; loop < indicesLoops_; loop++) { + if (loop == indicesLoops_ - 1) { + curLoopElements_ = curCoreLastLoopIndicesElements_; + } + CopyExpertIn(loop); + if constexpr (!EP) { + GatherCopyOut(loop); + } else { + ScatterCopyOut(loop); + } + } + } +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_GATHER_OUT_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_out_droppad.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_out_droppad.h new file mode 100644 index 000000000..d5229ba59 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_out_droppad.h @@ -0,0 +1,210 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_gather_out_droppad.h + * \brief + */ +#ifndef MOE_CUSTOM_GATHER_OUT_DROPPAD_H +#define MOE_CUSTOM_GATHER_OUT_DROPPAD_H + +#include "moe_custom_common.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +constexpr int64_t GATHER_OUT_DROPPAD_BUFFER_NUM = 2; + +template +class MoeGatherOutDroppad { +public: + __aicore__ inline MoeGatherOutDroppad(){}; + __aicore__ inline void Init(GM_ADDR inputX, GM_ADDR scale, GM_ADDR expandedRowIdx, GM_ADDR expandedX, + GM_ADDR expandedScale, GM_ADDR workspace, const MoeInitRoutingCustomTilingData *tilingData, + TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyInIndices(int64_t progress); + __aicore__ inline void CopyOut(int64_t progress); + __aicore__ inline void CopyScaleIn(int64_t scaleSrcOffset, LocalTensor scaleLocal); + __aicore__ inline void CopyScaleOut(int64_t scaleDstOffset, LocalTensor scaleLocal); + +private: + TPipe *pipe_; + TQueBind xCopyInQueue_; + TQueBind scaleCopyInQueue_; + TQue expandedRowIdxCopyInQueue_; + + GlobalTensor inputXGm_; + GlobalTensor xGscaleGm_; + GlobalTensor expandedXGm_; + GlobalTensor expandedRowIdxGm_; + GlobalTensor expandedScaleGm_; + + const MoeCustomGatherOutComputeTilingData *gatherOutTilingData_; + + int64_t needCoreNum_; + int64_t blockIdx_; + int64_t cols_; + int64_t n_; + int64_t k_; + int64_t currentLoopRows_; + int64_t coreRows_; + int64_t perLoopRows_; + int64_t lastLoopRows_; + int64_t rowLoops_; + int64_t colsTileLength_; + int64_t perLoopCols_; + int64_t lastLoopCols_; + int64_t colLoops_; + int64_t isInputScale_; + + int64_t indicesOffset_; + int64_t inputOffset_; + int64_t outOffset_; +}; + +template +__aicore__ inline void MoeGatherOutDroppad::CopyInIndices(int64_t progress) +{ + indicesOffset_ = progress * perLoopRows_; + LocalTensor indicesLocal = expandedRowIdxCopyInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(currentLoopRows_ * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(indicesLocal, expandedRowIdxGm_[indicesOffset_], dataCopyParams, dataCopyPadParams); + expandedRowIdxCopyInQueue_.EnQue(indicesLocal); +} + +template +__aicore__ inline void MoeGatherOutDroppad::CopyScaleIn(int64_t scaleSrcOffset, LocalTensor scaleLocal) +{ + DataCopyExtParams copyParams1{static_cast(1), static_cast(1 * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams padParams1{false, 0, 0, 0}; + DataCopyPad(scaleLocal, xGscaleGm_[scaleSrcOffset], copyParams1, padParams1); + scaleCopyInQueue_.EnQue(scaleLocal); +} + +template +__aicore__ inline void MoeGatherOutDroppad::CopyScaleOut(int64_t scaleDstOffset, LocalTensor scaleLocal) +{ + DataCopyExtParams copyParams3{1, static_cast(sizeof(float)), 0, 0, 0}; + DataCopyPad(expandedScaleGm_[scaleDstOffset], scaleLocal, copyParams3); +} + +template +__aicore__ inline void MoeGatherOutDroppad::CopyOut(int64_t progress) +{ + LocalTensor indicesLocal = expandedRowIdxCopyInQueue_.DeQue(); + SetWaitFlag(HardEvent::MTE2_S); + colsTileLength_ = perLoopCols_; + for (int64_t colsLoop = 0; colsLoop < colLoops_; colsLoop++) { + int64_t initialRow = gatherOutTilingData_->perCoreIndicesElements * blockIdx_ + perLoopRows_ * progress; + int64_t curLoopRow = 0; + if (colsLoop == colLoops_ - 1) { + colsTileLength_ = lastLoopCols_; + } + int64_t currentLoopStartRow = initialRow / k_; + int64_t currentLoopLastRow = (initialRow + currentLoopRows_ - 1) / k_; + for (int64_t row = currentLoopStartRow; row <= currentLoopLastRow; row++) { + LocalTensor scaleLocal = scaleCopyInQueue_.AllocTensor(); + if (isInputScale_ == 1) { + CopyScaleIn(row, scaleLocal); + LocalTensor scaleLocal = scaleCopyInQueue_.DeQue(); + } + inputOffset_ = row * cols_ + colsLoop * perLoopCols_; + // input row position + LocalTensor inLocal = xCopyInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(colsTileLength_ * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal, inputXGm_[inputOffset_], dataCopyParams, dataCopyPadParams); + SetWaitFlag(HardEvent::MTE2_MTE3); + DataCopyExtParams intriParams{1, static_cast(colsTileLength_ * sizeof(T)), 0, 0, 0}; + while (curLoopRow < currentLoopRows_ && initialRow / k_ == row) { + int32_t outIndex = indicesLocal.GetValue(curLoopRow); + curLoopRow++; + initialRow++; + if (outIndex == -1) { + continue; + } + outOffset_ = outIndex * cols_ + colsLoop * perLoopCols_; + DataCopyPad(expandedXGm_[outOffset_], inLocal, intriParams); + if (isInputScale_ == 1) { + CopyScaleOut(outIndex, scaleLocal); + } + } + xCopyInQueue_.FreeTensor(inLocal); + scaleCopyInQueue_.FreeTensor(scaleLocal); + } + } + expandedRowIdxCopyInQueue_.FreeTensor(indicesLocal); +} + +template +__aicore__ inline void MoeGatherOutDroppad::Init(GM_ADDR inputX, GM_ADDR scale, GM_ADDR expandedRowIdx, + GM_ADDR expandedX, GM_ADDR expandedScale, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + gatherOutTilingData_ = &(tilingData->gatherOutComputeParamsOp); + + needCoreNum_ = gatherOutTilingData_->needCoreNum; + cols_ = tilingData->cols; + n_ = tilingData->n; + k_ = tilingData->k; + isInputScale_ = tilingData->isInputScale; + + if (blockIdx_ == needCoreNum_ - 1) { + coreRows_ = gatherOutTilingData_->lastCoreIndicesElements; + perLoopRows_ = gatherOutTilingData_->lastCorePerLoopIndicesElements; + lastLoopRows_ = gatherOutTilingData_->lastCoreLastLoopIndicesElements; + rowLoops_ = gatherOutTilingData_->lastCoreIndicesLoops; + } else { + coreRows_ = gatherOutTilingData_->perCoreIndicesElements; + perLoopRows_ = gatherOutTilingData_->perCorePerLoopIndicesElements; + lastLoopRows_ = gatherOutTilingData_->perCoreLastLoopIndicesElements; + rowLoops_ = gatherOutTilingData_->perCoreIndicesLoops; + } + perLoopCols_ = gatherOutTilingData_->perLoopCols; + lastLoopCols_ = gatherOutTilingData_->lastLoopCols; + colLoops_ = gatherOutTilingData_->colsLoops; + + inputXGm_.SetGlobalBuffer((__gm__ T *)inputX, coreRows_ * cols_); + xGscaleGm_.SetGlobalBuffer((__gm__ float *)scale, n_); + expandedXGm_.SetGlobalBuffer((__gm__ T *)expandedX, n_ * k_ * cols_); + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx + + blockIdx_ * gatherOutTilingData_->perCoreIndicesElements, + Align(coreRows_, sizeof(int32_t))); + expandedScaleGm_.SetGlobalBuffer((__gm__ float *)expandedScale); + + pipe_->InitBuffer(xCopyInQueue_, GATHER_OUT_DROPPAD_BUFFER_NUM, AlignBytes(perLoopCols_, sizeof(T))); + pipe_->InitBuffer(expandedRowIdxCopyInQueue_, GATHER_OUT_DROPPAD_BUFFER_NUM, + AlignBytes(perLoopRows_, sizeof(int32_t))); + pipe_->InitBuffer(scaleCopyInQueue_, GATHER_OUT_DROPPAD_BUFFER_NUM, AlignBytes(1, sizeof(float))); +} + +template +__aicore__ inline void MoeGatherOutDroppad::Process() +{ + if (blockIdx_ < needCoreNum_) { + currentLoopRows_ = perLoopRows_; + for (int64_t loop = 0; loop < rowLoops_; loop++) { + if (loop == rowLoops_ - 1) { + currentLoopRows_ = lastLoopRows_; + } + CopyInIndices(loop); + CopyOut(loop); + } + } +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_GATHER_OUT_DROPPAD_H diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_sort_multi_core.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_sort_multi_core.h new file mode 100644 index 000000000..6ed3f2d72 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_sort_multi_core.h @@ -0,0 +1,242 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_gather_sort_multi_core.h + * \brief + */ +#ifndef MOE_CUSTOM_GATHER_SORT_MULTI_CORE_H +#define MOE_CUSTOM_GATHER_SORT_MULTI_CORE_H + +#include "moe_custom_common.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +constexpr int64_t SORT32_ALIGN_ELEMENT = 32; +constexpr int64_t PARALLEL_GATHERED_SORT_NEED_CORE_NUM = 16; +constexpr int64_t MULTI_GATHERED_MAX_NUM = 4096; // 8192 * 8 / 16 + +class MoeGatherSortMultiCore { +public: + __aicore__ inline MoeGatherSortMultiCore(){}; + __aicore__ inline void Init(GM_ADDR expertIdx, GM_ADDR expendedRowIdx, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyIn(); + __aicore__ inline void Compute(); + __aicore__ inline void CopyOut(); + +private: + TPipe *pipe_; + TBuf buffer_; + GlobalTensor workspaceGm_; + GlobalTensor expendedRowIdxGm_; + GlobalTensor expertIdxGm_; + GlobalTensor sortedExpertIdxGm_; + GlobalTensor sortedExpertIndexGm_; + GlobalTensor sortedNumGm_; + + TQue sortedNumCopyOutQueue_; + + int64_t expertIdxOffset_ = 0; + int64_t expertIndexOffset_ = 0; + int64_t compareScalarMask0Offset_ = 0; + int64_t compareScalarMask1Offset_ = 0; + int64_t gatherMaskOffset_ = 0; + + int64_t totalLength_; + int64_t expertStart_ = 0; + int64_t expertEnd_ = 0; + int64_t actual_expert_num_ = 0; + int64_t needCoreNum_ = 0; + int64_t perCoreElements_ = 0; + int64_t blockIdx_; + int64_t currentCoreElements_ = 0; + int64_t needSortNum_ = 0; + int64_t kvFactor = 2; + + static constexpr int64_t DST_BLK_STRIDE = 1; + static constexpr int64_t DST_REP_STRIDE = 8; + static constexpr int64_t MASK_STRIDE = 64; +}; + +__aicore__ inline void MoeGatherSortMultiCore::CopyIn() +{ + LocalTensor expertIdx = buffer_.Get()[expertIdxOffset_ / sizeof(int32_t)]; + + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyExtParams dataCopyParams{static_cast(1), + static_cast(currentCoreElements_ * sizeof(int32_t)), 0, 0, 0}; + + DataCopyPad(expertIdx, expertIdxGm_[blockIdx_ * perCoreElements_], dataCopyParams, dataCopyPadParams); + SetWaitFlag(HardEvent::MTE2_V); +} + +__aicore__ inline void MoeGatherSortMultiCore::Compute() +{ + LocalTensor expertIdx = buffer_.Get()[expertIdxOffset_ / sizeof(int32_t)]; + LocalTensor expertIdxFp32 = expertIdx.ReinterpretCast(); + LocalTensor gatheredExpertIdx = buffer_.Get(); + LocalTensor gatheredExpertIdxFp32 = gatheredExpertIdx.ReinterpretCast(); + + Cast(expertIdxFp32, expertIdx, RoundMode::CAST_ROUND, currentCoreElements_); + PipeBarrier(); + Muls(expertIdxFp32, expertIdxFp32, (float)-1, currentCoreElements_); + PipeBarrier(); + + LocalTensor compareScalarMaskLocalTensor0 = buffer_.Get()[compareScalarMask0Offset_]; + LocalTensor compareScalarMaskLocalTensor1 = buffer_.Get()[compareScalarMask1Offset_]; + LocalTensor gatherMaskLocalTensor = buffer_.Get()[gatherMaskOffset_]; + + // Find elements >= expertStart_, which means -elements <= -expertStart_ + AscendC::CompareScalar( + compareScalarMaskLocalTensor0, expertIdxFp32, static_cast(-expertStart_), AscendC::CMPMODE::LE, + (currentCoreElements_ + ONE_REPEAT_COMPARE_NUM - 1) / ONE_REPEAT_COMPARE_NUM * ONE_REPEAT_COMPARE_NUM); + PipeBarrier(); + + // Find elements < expertEnd_, which means -elements > -expertEnd_ + AscendC::CompareScalar( + compareScalarMaskLocalTensor1, expertIdxFp32, static_cast(-expertEnd_), AscendC::CMPMODE::GT, + (currentCoreElements_ + ONE_REPEAT_COMPARE_NUM - 1) / ONE_REPEAT_COMPARE_NUM * ONE_REPEAT_COMPARE_NUM); + PipeBarrier(); + + // Get experts between [expert_start, expert_end) + And(gatherMaskLocalTensor.ReinterpretCast(), compareScalarMaskLocalTensor0.ReinterpretCast(), + compareScalarMaskLocalTensor1.ReinterpretCast(), + Ceil(currentCoreElements_, MASK_STRIDE) * MASK_STRIDE / DST_REP_STRIDE / kvFactor); + PipeBarrier(); + + uint64_t sortedNum = 0; + GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = 1; + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = DST_REP_STRIDE; + gatherMaskParams.src1RepeatStride = DST_REP_STRIDE; + GatherMask(gatheredExpertIdxFp32, expertIdxFp32, gatherMaskLocalTensor.ReinterpretCast(), true, + static_cast(currentCoreElements_), gatherMaskParams, sortedNum); + PipeBarrier(); + actual_expert_num_ = sortedNum; + int64_t needSortNum = Ceil(static_cast(sortedNum), ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + needSortNum_ = needSortNum; + + // Handle actual_expert_num_ == 0 + if (actual_expert_num_ < 1) { + return; + } + + LocalTensor expertIndex = buffer_.Get()[expertIdxOffset_ / sizeof(int32_t)]; + LocalTensor gatheredExpertIndex = buffer_.Get()[needSortNum]; + ArithProgression(expertIndex, blockIdx_ * perCoreElements_, 1, currentCoreElements_); + GatherMask(gatheredExpertIndex, expertIndex, gatherMaskLocalTensor.ReinterpretCast(), true, + static_cast(currentCoreElements_), gatherMaskParams, sortedNum); + PipeBarrier(); + int64_t duplicateNum = sortedNum % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + int duplicateIndex = sortedNum - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(gatheredExpertIdxFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + } + PipeBarrier(); + + LocalTensor concatLocal; + LocalTensor sortTempTensor = buffer_.Get()[needSortNum * kvFactor]; + Concat(concatLocal, gatheredExpertIdxFp32, sortTempTensor, needSortNum / ONE_REPEAT_SORT_NUM); + LocalTensor sortedLocal = buffer_.Get()[needSortNum * kvFactor + needSortNum * kvFactor * kvFactor]; + Sort(sortedLocal, concatLocal, gatheredExpertIndex.ReinterpretCast(), sortTempTensor, + needSortNum / ONE_REPEAT_SORT_NUM); + SetWaitFlag(HardEvent::V_MTE3); +} + +__aicore__ inline void MoeGatherSortMultiCore::CopyOut() +{ + // Copy out sortedLocal for MergeSort + if (actual_expert_num_ > 0) { + LocalTensor sortedLocal = + buffer_.Get()[needSortNum_ * kvFactor + needSortNum_ * kvFactor * kvFactor]; + DataCopyExtParams extParams{static_cast(1), + static_cast(2 * actual_expert_num_ * sizeof(float)), 0, 0, 0}; + int64_t curCoreStartIndex = 2 * GetBlockIdx() * perCoreElements_; + DataCopyPad(sortedExpertIdxGm_[curCoreStartIndex], sortedLocal, extParams); + } + + // Copyout actual_expert_num_ + LocalTensor sortedNumOutLocal = sortedNumCopyOutQueue_.AllocTensor(); + sortedNumOutLocal.SetValue(0, actual_expert_num_); + SetWaitFlag(HardEvent::S_MTE3); + DataCopyExtParams copyParams3{static_cast(1), static_cast(sizeof(uint32_t)), 0, 0, 0}; + DataCopyPad(sortedNumGm_[GetBlockIdx()], sortedNumOutLocal, copyParams3); + + sortedNumCopyOutQueue_.FreeTensor(sortedNumOutLocal); +} + +__aicore__ inline void MoeGatherSortMultiCore::Init(GM_ADDR expertIdx, GM_ADDR expendedRowIdx, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + totalLength_ = tilingData->n * tilingData->k; + + expertStart_ = tilingData->expertStart; + expertEnd_ = tilingData->expertEnd; + + expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx); + + expendedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expendedRowIdx); + + workspaceGm_.SetGlobalBuffer((__gm__ int32_t *)workspace); + + sortedExpertIdxGm_.SetGlobalBuffer((__gm__ float *)workspace); + sortedExpertIndexGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(totalLength_, sizeof(int32_t))); + + // key and value + sortedNumGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + + Align(totalLength_, sizeof(int32_t)) * kvFactor * kvFactor); + + needCoreNum_ = PARALLEL_GATHERED_SORT_NEED_CORE_NUM; + perCoreElements_ = Ceil(totalLength_, needCoreNum_); + + int32_t lastCoreElements = totalLength_ - (needCoreNum_ - 1) * perCoreElements_; + if (blockIdx_ == (needCoreNum_ - 1)) { + currentCoreElements_ = lastCoreElements; + } else { + currentCoreElements_ = perCoreElements_; + } + + // expertIdxOffset_ + expertIdxOffset_ = AlignBytes(currentCoreElements_, sizeof(int32_t)); + expertIndexOffset_ = expertIdxOffset_; + + gatherMaskOffset_ = expertIdxOffset_ * kvFactor; + int64_t maskOffset = + AlignBytes(Ceil(currentCoreElements_, MASK_STRIDE) * MASK_STRIDE / DST_REP_STRIDE, sizeof(int8_t)); + compareScalarMask0Offset_ = gatherMaskOffset_ + maskOffset; + compareScalarMask1Offset_ = compareScalarMask0Offset_ + maskOffset; + int64_t bufferSize = MULTI_GATHERED_MAX_NUM * kvFactor * kvFactor * kvFactor * sizeof(int32_t); + pipe_->InitBuffer(sortedNumCopyOutQueue_, 1, AlignBytes(1, sizeof(int32_t))); + pipe_->InitBuffer(buffer_, bufferSize); // 73728 Bytes +} + +__aicore__ inline void MoeGatherSortMultiCore::Process() +{ + if (blockIdx_ < PARALLEL_GATHERED_SORT_NEED_CORE_NUM) { + CopyIn(); + Compute(); + CopyOut(); + } + SyncAll(); +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_GATHER_SORT_MULTI_CORE_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_static_quant.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_static_quant.h new file mode 100644 index 000000000..78bd9e1c9 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_gather_static_quant.h @@ -0,0 +1,329 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_gather_quant.h + * \brief + */ +#ifndef MOE_CUSTOM_GATHER_STATIC_QUANT_H +#define MOE_CUSTOM_GATHER_STATIC_QUANT_H + +#include "moe_custom_common.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +constexpr int64_t GATHER_OUT_QUANT_BUFFER_NUM = 2; + +template +class MoeGatherOutQuant { +public: + __aicore__ inline MoeGatherOutQuant(){}; + __aicore__ inline void Init(GM_ADDR inputX, GM_ADDR scale, GM_ADDR offset, GM_ADDR expandedRowIdx, + GM_ADDR expandedX, GM_ADDR workspace, const MoeInitRoutingCustomTilingData *tilingData, + TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyExpertIn(int64_t progress); + __aicore__ inline void Compute(int64_t curLoopCols); + __aicore__ inline void CopyXIn(int64_t xSrcOffset, int64_t curLoopCols); + __aicore__ inline void CopyXOut(int64_t xDstOffset, int64_t curLoopCols); + __aicore__ inline void ScatterCopyOut(int64_t progress); + __aicore__ inline void GatherCopyOut(int64_t progress); + +private: + TPipe *pipe_; + TQue inputXCopyInQueue_; + TQue expandRowIdxCopyInQueue_; + TQue inputXCopyOutQueue_; + TQue floatQueue_; + TQue halfQueue_; + + GlobalTensor inputXGm_; + GlobalTensor expandedXGm_; + GlobalTensor expandedRowIdxGm_; + GlobalTensor scaleGm_; + GlobalTensor offsetGm_; + GlobalTensor expertTotalCountGm_; + + const MoeCustomGatherOutComputeTilingData *gatherOutTilingData_; + + int64_t needCoreNum_; + int64_t blockIdx_; + int64_t cols_; + int64_t n_; + int64_t k_; + int64_t perCoreRow_; + int64_t currentLoopRows_; + int64_t coreRows_; + int64_t perLoopRows_; + int64_t lastLoopRows_; + int64_t rowLoops_; + int64_t colsTileLength_; + int64_t perLoopCols_; + int64_t lastLoopCols_; + int64_t colLoops_; + float scale_; + float offset_; + int64_t rowIdxType_; + int64_t dropPadMode_; + int64_t activeNum_; + int64_t indicesOffset_; + int64_t coreNum_; + int64_t inputOffset_; + int64_t outOffset_; + int64_t expertTotalCount_; +}; + +template +__aicore__ inline void MoeGatherOutQuant::Init(GM_ADDR inputX, GM_ADDR scale, GM_ADDR offset, + GM_ADDR expandedRowIdx, GM_ADDR expandedX, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + pipe_ = tPipe; + blockIdx_ = GetBlockIdx(); + + gatherOutTilingData_ = &(tilingData->gatherOutComputeParamsOp); + cols_ = tilingData->cols; + n_ = tilingData->n; + k_ = tilingData->k; + rowIdxType_ = tilingData->rowIdxType; + dropPadMode_ = tilingData->dropPadMode; + activeNum_ = tilingData->activeNum; + coreNum_ = tilingData->coreNum; + + // core split + int64_t actualExpertNum_ = tilingData->actualExpertNum; + + if constexpr (EP) { + expertTotalCountGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(n_ * k_, sizeof(int32_t)) * 2 + + Align(actualExpertNum_, sizeof(int32_t)), + 1); + AscendC::DataCacheCleanAndInvalid(expertTotalCountGm_); + expertTotalCount_ = expertTotalCountGm_.GetValue(0); + } else { + expertTotalCount_ = n_ * k_; + } + + perCoreRow_ = Ceil(expertTotalCount_, tilingData->coreNum); + needCoreNum_ = Ceil(expertTotalCount_, perCoreRow_); + int64_t lastCoreIndicesElements_ = expertTotalCount_ - (needCoreNum_ - 1) * perCoreRow_; + + // inner core split + int64_t originPerLoopElements; + if (blockIdx_ == needCoreNum_ - 1) { + coreRows_ = lastCoreIndicesElements_; + originPerLoopElements = gatherOutTilingData_->lastCorePerLoopIndicesElements; + } else { + coreRows_ = perCoreRow_; + originPerLoopElements = gatherOutTilingData_->perCorePerLoopIndicesElements; + } + perLoopRows_ = Min(coreRows_, originPerLoopElements); + rowLoops_ = Ceil(coreRows_, perLoopRows_); + lastLoopRows_ = coreRows_ - (rowLoops_ - 1) * perLoopRows_; + + // cols split + perLoopCols_ = gatherOutTilingData_->perLoopCols; + lastLoopCols_ = gatherOutTilingData_->lastLoopCols; + colLoops_ = gatherOutTilingData_->colsLoops; + + inputXGm_.SetGlobalBuffer((__gm__ T *)inputX); + expandedXGm_.SetGlobalBuffer((__gm__ int8_t *)expandedX); + + if constexpr (EP) { + if (rowIdxType_ == SCATTER) { + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx + blockIdx_ * perCoreRow_, + Align(coreRows_, sizeof(int32_t))); + } else { + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(n_ * k_, sizeof(int32_t)) + + blockIdx_ * perCoreRow_, + Align(coreRows_, sizeof(int32_t))); + } + } else { + if (rowIdxType_ == GATHER) { + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx + blockIdx_ * perCoreRow_, + Align(coreRows_, sizeof(int32_t))); + } else { + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(n_ * k_, sizeof(int32_t)) + + blockIdx_ * perCoreRow_, + Align(coreRows_, sizeof(int32_t))); + } + } + + + scaleGm_.SetGlobalBuffer((__gm__ float *)scale, 1); + offsetGm_.SetGlobalBuffer((__gm__ float *)offset, 1); + scale_ = scaleGm_.GetValue(0); + offset_ = offsetGm_.GetValue(0); + + pipe_->InitBuffer(inputXCopyInQueue_, GATHER_OUT_QUANT_BUFFER_NUM, AlignBytes(perLoopCols_, sizeof(T))); + pipe_->InitBuffer(inputXCopyOutQueue_, GATHER_OUT_QUANT_BUFFER_NUM, AlignBytes(perLoopCols_, sizeof(int8_t))); + pipe_->InitBuffer(expandRowIdxCopyInQueue_, GATHER_OUT_QUANT_BUFFER_NUM, AlignBytes(perLoopRows_, sizeof(int32_t))); + pipe_->InitBuffer(floatQueue_, 1, AlignBytes(perLoopCols_, sizeof(float))); + pipe_->InitBuffer(halfQueue_, 1, AlignBytes(perLoopCols_, sizeof(half))); +} + +template +__aicore__ inline void MoeGatherOutQuant::CopyExpertIn(int64_t progress) +{ + indicesOffset_ = progress * perLoopRows_; + LocalTensor indicesLocal = expandRowIdxCopyInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{1, static_cast(currentLoopRows_ * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(indicesLocal, expandedRowIdxGm_[indicesOffset_], dataCopyParams, dataCopyPadParams); + expandRowIdxCopyInQueue_.EnQue(indicesLocal); +} + +template +__aicore__ inline void MoeGatherOutQuant::CopyXIn(int64_t xSrcOffset, int64_t curLoopCols) +{ + LocalTensor inLocal = inputXCopyInQueue_.AllocTensor(); + DataCopyExtParams copyParams0{static_cast(1), static_cast(curLoopCols * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams padParams0{false, 0, 0, 0}; + DataCopyPad(inLocal, inputXGm_[xSrcOffset], copyParams0, padParams0); + inputXCopyInQueue_.EnQue(inLocal); +} + +template +__aicore__ inline void MoeGatherOutQuant::CopyXOut(int64_t xDstOffset, int64_t curLoopCols) +{ + LocalTensor outLocal = inputXCopyOutQueue_.DeQue(); + DataCopyExtParams copyParams2{1, static_cast(curLoopCols * sizeof(int8_t)), 0, 0, 0}; + DataCopyPad(expandedXGm_[xDstOffset], outLocal, copyParams2); + inputXCopyOutQueue_.FreeTensor(outLocal); +} + +template +__aicore__ inline void MoeGatherOutQuant::Compute(int64_t curLoopCols) +{ + LocalTensor floatLocal; + LocalTensor inLocal; + LocalTensor outLocal = inputXCopyOutQueue_.AllocTensor(); + LocalTensor halfLocal = halfQueue_.AllocTensor(); + uint32_t elements = Align(curLoopCols, sizeof(T)); + if constexpr (IsSameType::value) { + floatLocal = inputXCopyInQueue_.DeQue(); + } else { + inLocal = inputXCopyInQueue_.DeQue(); + floatLocal = floatQueue_.AllocTensor(); + Cast(floatLocal, inLocal, RoundMode::CAST_NONE, elements); + PipeBarrier(); + } + Muls(floatLocal, floatLocal, scale_, elements); + PipeBarrier(); + Adds(floatLocal, floatLocal, offset_, elements); + PipeBarrier(); + LocalTensor intLocal = floatLocal.ReinterpretCast(); + Cast(intLocal, floatLocal, RoundMode::CAST_RINT, elements); + PipeBarrier(); + SetDeqScale((half)1.000000e+00f); + PipeBarrier(); + Cast(halfLocal, intLocal, RoundMode::CAST_ROUND, elements); + PipeBarrier(); + Cast(outLocal, halfLocal, RoundMode::CAST_TRUNC, elements); + inputXCopyOutQueue_.EnQue(outLocal); + if constexpr (IsSameType::value) { + inputXCopyInQueue_.FreeTensor(floatLocal); + } else { + inputXCopyInQueue_.FreeTensor(inLocal); + floatQueue_.FreeTensor(floatLocal); + } + halfQueue_.FreeTensor(halfLocal); +} + +template +__aicore__ inline void MoeGatherOutQuant::ScatterCopyOut(int64_t progress) +{ + LocalTensor indicesLocal = expandRowIdxCopyInQueue_.DeQue(); + SetWaitFlag(HardEvent::MTE2_S); + for (int64_t indicesIndex = 0; indicesIndex < currentLoopRows_; indicesIndex++) { + int64_t rowOffset = perCoreRow_ * blockIdx_ + perLoopRows_ * progress; + int64_t rowIdx = indicesLocal.GetValue(indicesIndex); + int64_t xSrcOffset = rowIdx / k_ * cols_; + int64_t xDstOffset = (rowOffset + indicesIndex) * cols_; + int64_t curLoopCols = perLoopCols_; + if (activeNum_ > 0 && dropPadMode_ == DROPLESS_MODE && (rowOffset + indicesIndex) >= activeNum_) { + break; + } + SetWaitFlag(HardEvent::S_MTE2); + for (int64_t colsLoop = 0; colsLoop < colLoops_; colsLoop++) { + if (colsLoop == colLoops_ - 1) { + curLoopCols = lastLoopCols_; + } + int64_t colsLoopOffset = colsLoop * perLoopCols_; + CopyXIn(xSrcOffset + colsLoopOffset, curLoopCols); + Compute(curLoopCols); + CopyXOut(xDstOffset + colsLoopOffset, curLoopCols); + } + } + expandRowIdxCopyInQueue_.FreeTensor(indicesLocal); +} + +template +__aicore__ inline void MoeGatherOutQuant::GatherCopyOut(int64_t progress) +{ + LocalTensor indicesLocal = expandRowIdxCopyInQueue_.DeQue(); + SetWaitFlag(HardEvent::MTE2_S); + colsTileLength_ = perLoopCols_; + for (int64_t colsLoop = 0; colsLoop < colLoops_; colsLoop++) { + int64_t initialRow = perCoreRow_ * blockIdx_ + perLoopRows_ * progress; + int64_t curLoopRow = 0; + if (colsLoop == colLoops_ - 1) { + colsTileLength_ = lastLoopCols_; + } + int64_t currentLoopStartRow = initialRow / k_; + int64_t currentLoopLastRow = (initialRow + currentLoopRows_ - 1) / k_; + for (int64_t row = currentLoopStartRow; row <= currentLoopLastRow; row++) { + inputOffset_ = row * cols_ + colsLoop * perLoopCols_; + // input row position + CopyXIn(inputOffset_, colsTileLength_); + Compute(colsTileLength_); + LocalTensor outLocal = inputXCopyOutQueue_.DeQue(); + DataCopyExtParams intriParams{1, static_cast(colsTileLength_ * sizeof(int8_t)), 0, 0, 0}; + SetWaitFlag(HardEvent::MTE2_MTE3); + while (curLoopRow < currentLoopRows_ && initialRow / k_ == row) { + int32_t outIndex = indicesLocal.GetValue(curLoopRow); + curLoopRow++; + initialRow++; + if (outIndex == -1 || (dropPadMode_ == DROPLESS_MODE && outIndex >= activeNum_)) { + continue; + } + outOffset_ = outIndex * cols_ + colsLoop * perLoopCols_; + DataCopyPad(expandedXGm_[outOffset_], outLocal, intriParams); + } + inputXCopyOutQueue_.FreeTensor(outLocal); + } + } + expandRowIdxCopyInQueue_.FreeTensor(indicesLocal); +} + +template +__aicore__ inline void MoeGatherOutQuant::Process() +{ + if (blockIdx_ < needCoreNum_) { + currentLoopRows_ = perLoopRows_; + for (int64_t loop = 0; loop < rowLoops_; loop++) { + if (loop == rowLoops_ - 1) { + currentLoopRows_ = lastLoopRows_; + } + CopyExpertIn(loop); + if constexpr (EP) { + ScatterCopyOut(loop); + } else { + GatherCopyOut(loop); + } + } + } +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_GATHER_STATIC_QUANT_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_mrgsort.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_mrgsort.h new file mode 100644 index 000000000..4ae95fa2b --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_mrgsort.h @@ -0,0 +1,207 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_mrgsort.h + * \brief + */ +#ifndef MOE_CUSTOM_MRGSORT_H +#define MOE_CUSTOM_MRGSORT_H + +#include "moe_custom_common.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +struct MoeMrgsortParam { + int64_t perListElements; + int64_t lastListElements; + int64_t oneLoopMaxElements; +}; + +class MoeMrgsort { +public: + __aicore__ inline MoeMrgsort(){}; + __aicore__ inline void Init(MoeMrgsortParam *param); + __aicore__ inline void Process(); + __aicore__ inline void SetInput(GlobalTensor &gmInput, LocalTensor &ubInput); + __aicore__ inline void SetOutput(GlobalTensor &gmOutput, LocalTensor &ubOutput); + +private: + __aicore__ inline void CopyIn(); + __aicore__ inline void UpdateMrgParam(); + __aicore__ inline void MrgsortCompute(); + __aicore__ inline void UpdateSortInfo(); + __aicore__ inline void CopyOut(); + __aicore__ inline void ClearCache(); + +private: + MoeMrgsortParam *param = nullptr; + + GlobalTensor gmInputs[4]; + GlobalTensor gmOutput; + + LocalTensor ubInputs[4]; + LocalTensor ubOutput; + + int64_t listNum{0}; + int64_t remainListNum{0}; + int64_t outOffset{0}; + int64_t offsets[4]; + int64_t listRemainElements[4]; + int64_t lengths[4]; + int64_t allRemainElements{0}; + int64_t curLoopSortedNum{0}; + + // for MrgSort + uint16_t validBitTail{0}; + uint16_t elementCountListTail[4]; + uint32_t listSortedNums[4]; + LocalTensor tmpUbInputs[4]; +}; + +__aicore__ inline void MoeMrgsort::ClearCache() +{ + this->listNum = 0; + this->allRemainElements = 0; + this->outOffset = 0; +} + +__aicore__ inline void MoeMrgsort::SetInput(GlobalTensor &gmInput, LocalTensor &ubInput) +{ + this->gmInputs[listNum] = gmInput; + this->ubInputs[listNum] = ubInput; + this->listNum += 1; +} + +__aicore__ inline void MoeMrgsort::SetOutput(GlobalTensor &gmOutput, LocalTensor &ubOutput) +{ + this->gmOutput = gmOutput; + this->ubOutput = ubOutput; +} + +__aicore__ inline void MoeMrgsort::UpdateMrgParam() +{ + if (this->remainListNum == MERGE_LIST_TWO) { + elementCountListTail[MERGE_LIST_IDX_TWO] = 0; + elementCountListTail[MERGE_LIST_IDX_THREE] = 0; + validBitTail = 0b0011; + } else if (this->remainListNum == MERGE_LIST_THREE) { + elementCountListTail[MERGE_LIST_IDX_THREE] = 0; + validBitTail = 0b0111; + } else if (this->remainListNum == MERGE_LIST_FOUR) { + validBitTail = 0b1111; + } else { + validBitTail = 0b0001; + } +} + +__aicore__ inline void MoeMrgsort::CopyIn() +{ + this->remainListNum = 0; + event_t eventIdMte3ToMte2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventIdMte3ToMte2); + WaitFlag(eventIdMte3ToMte2); + for (int64_t i = 0, j = 0; i < listNum; i++) { + lengths[i] = Min(param->oneLoopMaxElements, listRemainElements[i]); + if (lengths[i] > 0) { + DataCopy(this->ubInputs[i], this->gmInputs[i][offsets[i]], + Align(GetSortLen(lengths[i]), sizeof(float))); + tmpUbInputs[j] = this->ubInputs[i]; + elementCountListTail[j] = lengths[i]; + this->remainListNum += 1; + j++; + } + } +} + +__aicore__ inline void MoeMrgsort::MrgsortCompute() +{ + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + if (this->remainListNum == MERGE_LIST_TWO) { + MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]); + MrgSort(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else if (this->remainListNum == MERGE_LIST_THREE) { + MrgSortSrcList sortListTail = + MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], tmpUbInputs[0]); + MrgSort(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else if (this->remainListNum == MERGE_LIST_FOUR) { + MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], + tmpUbInputs[MERGE_LIST_IDX_THREE]); + MrgSort(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else { + DataCopy(this->ubOutput, this->tmpUbInputs[0], + Align(GetSortLen(elementCountListTail[0]), sizeof(float))); + listSortedNums[0] = elementCountListTail[0]; + } +} + +__aicore__ inline void MoeMrgsort::UpdateSortInfo() +{ + curLoopSortedNum = 0; + for (int64_t i = 0, j = 0; i < listNum; i++) { + if (lengths[i] > 0) { + // update remain size + listRemainElements[i] -= listSortedNums[j]; + allRemainElements -= listSortedNums[j]; + // update offset + offsets[i] += GetSortOffset(listSortedNums[j]); + // update current loop sorted nums + curLoopSortedNum += listSortedNums[j]; + j += 1; + } + } +} + +__aicore__ inline void MoeMrgsort::CopyOut() +{ + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = GetSortLen(curLoopSortedNum) * sizeof(float); + event_t eventIdVToMte3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventIdVToMte3); + WaitFlag(eventIdVToMte3); + DataCopyPad(this->gmOutput[outOffset], this->ubOutput, intriParams); + outOffset += GetSortLen(curLoopSortedNum); +} + +__aicore__ inline void MoeMrgsort::Init(MoeMrgsortParam *param) +{ + this->param = param; + this->remainListNum = listNum; + + for (int64_t i = 0; i < listNum; i++) { + offsets[i] = GetSortOffset(param->perListElements * i); + if (i == listNum - 1) { + listRemainElements[i] = param->lastListElements; + } else { + listRemainElements[i] = param->perListElements; + } + allRemainElements += listRemainElements[i]; + } +} + +__aicore__ inline void MoeMrgsort::Process() +{ + for (; allRemainElements > 0;) { + CopyIn(); + UpdateMrgParam(); + MrgsortCompute(); + UpdateSortInfo(); + CopyOut(); + } + + ClearCache(); +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_MRGSORT_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_mrgsort_out.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_mrgsort_out.h new file mode 100644 index 000000000..84fb4b672 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_mrgsort_out.h @@ -0,0 +1,232 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_mrgsort_out.h + * \brief + */ +#ifndef MOE_CUSTOM_MRGSORT_OUT_H +#define MOE_CUSTOM_MRGSORT_OUT_H + +#include "moe_custom_mrgsort.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +class MoeMrgsortOut { +public: + __aicore__ inline MoeMrgsortOut(){}; + __aicore__ inline void Init(MoeMrgsortParam *param, TPipe *tPipe); + __aicore__ inline void Process(); + __aicore__ inline void SetInput(GlobalTensor &gmInput, LocalTensor &ubInput); + __aicore__ inline void SetOutput(GlobalTensor &gmOutput1, GlobalTensor &gmOutput2, + LocalTensor &ubOutput1, LocalTensor &ubOutput2); + __aicore__ inline void SetBuffer(LocalTensor &tempBuffer); + +private: + __aicore__ inline void CopyIn(); + __aicore__ inline void UpdateMrgParam(); + __aicore__ inline void MrgsortCompute(); + __aicore__ inline void UpdateSortInfo(); + __aicore__ inline void Extract(); + __aicore__ inline void CopyOut(); + __aicore__ inline void ClearCache(); + +private: + MoeMrgsortParam *param = nullptr; + + GlobalTensor gmInputs[4]; + GlobalTensor gmOutput1; + GlobalTensor gmOutput2; + + LocalTensor ubInputs[4]; + LocalTensor tempBuffer; + + // for extract + LocalTensor ubOutput1; + LocalTensor ubOutput2; + + // for copy out + LocalTensor ubOutputInt1; + LocalTensor ubOutputInt2; + + int64_t listNum{0}; + int64_t remainListNum{0}; + int64_t outOffset{0}; + int64_t offsets[4]; + int64_t listRemainElements[4]; + int64_t lengths[4]; + int64_t allRemainElements{0}; + int64_t curLoopSortedNum{0}; + + // for MrgSort + uint16_t validBitTail; + uint16_t elementCountListTail[4]; + uint32_t listSortedNums[4]; + LocalTensor tmpUbInputs[4]; +}; + +__aicore__ inline void MoeMrgsortOut::ClearCache() +{ + this->listNum = 0; + this->allRemainElements = 0; + this->outOffset = 0; +} + +__aicore__ inline void MoeMrgsortOut::SetInput(GlobalTensor &gmInput, LocalTensor &ubInput) +{ + this->gmInputs[listNum] = gmInput; + this->ubInputs[listNum] = ubInput; + this->listNum += 1; +} + +__aicore__ inline void MoeMrgsortOut::SetOutput(GlobalTensor &gmOutput1, GlobalTensor &gmOutput2, + LocalTensor &ubOutput1, LocalTensor &ubOutput2) +{ + this->gmOutput1 = gmOutput1; + this->ubOutput1 = ubOutput1; + this->ubOutputInt1 = ubOutput1.ReinterpretCast(); + + this->gmOutput2 = gmOutput2; + this->ubOutput2 = ubOutput2.ReinterpretCast(); + this->ubOutputInt2 = ubOutput2.ReinterpretCast(); +} + +__aicore__ inline void MoeMrgsortOut::SetBuffer(LocalTensor &tempBuffer) +{ + this->tempBuffer = tempBuffer; +} + +__aicore__ inline void MoeMrgsortOut::UpdateMrgParam() +{ + if (this->remainListNum == MERGE_LIST_TWO) { + elementCountListTail[MERGE_LIST_IDX_TWO] = 0; + elementCountListTail[MERGE_LIST_IDX_THREE] = 0; + validBitTail = 0b0011; + } else if (this->remainListNum == MERGE_LIST_THREE) { + elementCountListTail[MERGE_LIST_IDX_THREE] = 0; + validBitTail = 0b0111; + } else if (this->remainListNum == MERGE_LIST_FOUR) { + validBitTail = 0b1111; + } else { + validBitTail = 0b0001; + } +} + +__aicore__ inline void MoeMrgsortOut::CopyIn() +{ + this->remainListNum = 0; + event_t eventIdMte3ToMte2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventIdMte3ToMte2); + WaitFlag(eventIdMte3ToMte2); + for (int64_t i = 0, j = 0; i < listNum; i++) { + lengths[i] = Min(param->oneLoopMaxElements, listRemainElements[i]); + if (lengths[i] > 0) { + DataCopy(this->ubInputs[i], this->gmInputs[i][offsets[i]], + Align(GetSortLen(lengths[i]), sizeof(float))); + tmpUbInputs[j] = this->ubInputs[i]; + elementCountListTail[j] = lengths[i]; + this->remainListNum += 1; + j++; + } + } +} + +__aicore__ inline void MoeMrgsortOut::MrgsortCompute() +{ + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + if (this->remainListNum == MERGE_LIST_TWO) { + MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]); + MrgSort(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else if (this->remainListNum == MERGE_LIST_THREE) { + MrgSortSrcList sortListTail = + MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], tmpUbInputs[0]); + MrgSort(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else if (this->remainListNum == MERGE_LIST_FOUR) { + MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], + tmpUbInputs[MERGE_LIST_IDX_THREE]); + MrgSort(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else { + DataCopy(this->tempBuffer, this->tmpUbInputs[0], + Align(GetSortLen(elementCountListTail[0]), sizeof(float))); + listSortedNums[0] = elementCountListTail[0]; + } +} + +__aicore__ inline void MoeMrgsortOut::UpdateSortInfo() +{ + curLoopSortedNum = 0; + for (int64_t i = 0, j = 0; i < listNum; i++) { + if (lengths[i] > 0) { + // update remain size + listRemainElements[i] -= listSortedNums[j]; + allRemainElements -= listSortedNums[j]; + // update offset + offsets[i] += GetSortOffset(listSortedNums[j]); + // update current loop sorted nums + curLoopSortedNum += listSortedNums[j]; + j += 1; + } + } +} + +__aicore__ inline void MoeMrgsortOut::Extract() +{ + AscendC::Extract(this->ubOutput1, this->ubOutput2, this->tempBuffer, Ceil(curLoopSortedNum, ONE_REPEAT_SORT_NUM)); + Muls(this->ubOutput1, this->ubOutput1, (float)-1, Align(curLoopSortedNum, sizeof(float))); + Cast(this->ubOutputInt1, this->ubOutput1, RoundMode::CAST_ROUND, Align(curLoopSortedNum, sizeof(float))); +} + +__aicore__ inline void MoeMrgsortOut::CopyOut() +{ + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = curLoopSortedNum * sizeof(int32_t); + event_t eventIdVToMte3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventIdVToMte3); + WaitFlag(eventIdVToMte3); + DataCopyPad(this->gmOutput1[outOffset], this->ubOutputInt1, intriParams); + DataCopyPad(this->gmOutput2[outOffset], this->ubOutputInt2, intriParams); + + outOffset += curLoopSortedNum; +} + +__aicore__ inline void MoeMrgsortOut::Init(MoeMrgsortParam *param, TPipe *tPipe) +{ + this->param = param; + this->allRemainElements = 0; + for (int64_t i = 0; i < listNum; i++) { + offsets[i] = GetSortOffset(param->perListElements * i); + if (i == listNum - 1) { + listRemainElements[i] = param->lastListElements; + } else { + listRemainElements[i] = param->perListElements; + } + allRemainElements += listRemainElements[i]; + } +} + +__aicore__ inline void MoeMrgsortOut::Process() +{ + for (; allRemainElements > 0;) { + CopyIn(); + UpdateMrgParam(); + MrgsortCompute(); + UpdateSortInfo(); + Extract(); + CopyOut(); + } + ClearCache(); +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_MRGSORT_OUT_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_mrgsort_out_performance.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_mrgsort_out_performance.h new file mode 100644 index 000000000..650c90bcb --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_mrgsort_out_performance.h @@ -0,0 +1,239 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_mrgsort_out_performance.h + * \brief + */ +#ifndef MOE_CUSTOM_MRGSORT_OUT_PERFORMANCE_H +#define MOE_CUSTOM_MRGSORT_OUT_PERFORMANCE_H + +#include "moe_custom_mrgsort_performance.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +constexpr int64_t MAX_MRGSORT_LIST = 4; +constexpr int64_t MAX_MRGSORT_LIST_TOTAL = 16; + +class MoeMrgsortOutPerformance { +public: + __aicore__ inline MoeMrgsortOutPerformance(){}; + __aicore__ inline void Init(MoeMrgsortPerformanceParam *param, TPipe *tPipe); + __aicore__ inline void Process(); + __aicore__ inline void SetInput(GlobalTensor &gmInput, LocalTensor &ubInput, + GlobalTensor &gmActualSortNum); + __aicore__ inline void SetOutput(GlobalTensor &gmOutput1, GlobalTensor &gmOutput2, + LocalTensor &ubOutput1, LocalTensor &ubOutput2); + __aicore__ inline void SetBuffer(LocalTensor &tempBuffer); + +private: + __aicore__ inline void CopyIn(); + __aicore__ inline void UpdateMrgParam(); + __aicore__ inline void MrgsortCompute(); + __aicore__ inline void UpdateSortInfo(); + __aicore__ inline void Extract(); + __aicore__ inline void CopyOut(); + __aicore__ inline void ClearCache(); + +private: + MoeMrgsortPerformanceParam *param = nullptr; + + GlobalTensor gmInputs[4]; + GlobalTensor gmOutput1; + GlobalTensor gmOutput2; + GlobalTensor gmActualSortNum; + + LocalTensor ubInputs[4]; + LocalTensor tempBuffer; + + // for extract + LocalTensor ubOutput1; + LocalTensor ubOutput2; + + // for copy out + LocalTensor ubOutputInt1; + LocalTensor ubOutputInt2; + + int64_t listNum{0}; + int64_t remainListNum{0}; + int64_t outOffset{0}; + int64_t offsets[4] = {0}; + int64_t listRemainElements[4] = {0}; + int64_t lengths[4] = {0}; + int64_t allRemainElements{0}; + int64_t curLoopSortedNum{0}; + + // for MrgSort + uint16_t validBitTail; + uint16_t elementCountListTail[4] = {0}; + uint32_t listSortedNums[4] = {0}; + LocalTensor tmpUbInputs[4]; +}; + +__aicore__ inline void MoeMrgsortOutPerformance::ClearCache() +{ + this->listNum = 0; + this->allRemainElements = 0; + this->outOffset = 0; +} + +__aicore__ inline void MoeMrgsortOutPerformance::SetInput(GlobalTensor &gmInput, LocalTensor &ubInput, + GlobalTensor &gmActualSortNum) +{ + if (this->listNum == 0) { + this->gmActualSortNum = gmActualSortNum; + } + this->gmInputs[listNum] = gmInput; + this->ubInputs[listNum] = ubInput; + this->listNum += 1; +} + +__aicore__ inline void MoeMrgsortOutPerformance::SetOutput(GlobalTensor &gmOutput1, + GlobalTensor &gmOutput2, + LocalTensor &ubOutput1, LocalTensor &ubOutput2) +{ + this->gmOutput1 = gmOutput1; + this->ubOutput1 = ubOutput1; + this->ubOutputInt1 = ubOutput1.ReinterpretCast(); + + this->gmOutput2 = gmOutput2; + this->ubOutput2 = ubOutput2.ReinterpretCast(); + this->ubOutputInt2 = ubOutput2.ReinterpretCast(); +} + +__aicore__ inline void MoeMrgsortOutPerformance::SetBuffer(LocalTensor &tempBuffer) +{ + this->tempBuffer = tempBuffer; +} + +__aicore__ inline void MoeMrgsortOutPerformance::UpdateMrgParam() +{ + if (this->remainListNum == MERGE_LIST_TWO) { + elementCountListTail[MERGE_LIST_IDX_TWO] = 0; + elementCountListTail[MERGE_LIST_IDX_THREE] = 0; + validBitTail = 0b0011; + } else if (this->remainListNum == MERGE_LIST_THREE) { + elementCountListTail[MERGE_LIST_IDX_THREE] = 0; + validBitTail = 0b0111; + } else if (this->remainListNum == MERGE_LIST_FOUR) { + validBitTail = 0b1111; + } else { + validBitTail = 0b0001; + } +} + +__aicore__ inline void MoeMrgsortOutPerformance::CopyIn() +{ + this->remainListNum = 0; + event_t eventIdMte3ToMte2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventIdMte3ToMte2); + WaitFlag(eventIdMte3ToMte2); + for (int64_t i = 0, j = 0; i < listNum; i++) { + lengths[i] = Min(param->oneLoopMaxElements, listRemainElements[i]); + if (lengths[i] > 0) { + DataCopy(this->ubInputs[i], this->gmInputs[i][offsets[i]], + Align(GetSortLen(lengths[i]), sizeof(float))); + tmpUbInputs[j] = this->ubInputs[i]; + elementCountListTail[j] = lengths[i]; + this->remainListNum += 1; + j++; + } + } +} + +__aicore__ inline void MoeMrgsortOutPerformance::MrgsortCompute() +{ + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + if (this->remainListNum == MERGE_LIST_TWO) { + MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]); + MrgSort(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else if (this->remainListNum == MERGE_LIST_THREE) { + MrgSortSrcList sortListTail = + MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], tmpUbInputs[0]); + MrgSort(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else if (this->remainListNum == MERGE_LIST_FOUR) { + MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], + tmpUbInputs[MERGE_LIST_IDX_THREE]); + MrgSort(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else { + DataCopy(this->tempBuffer, this->tmpUbInputs[0], + Align(GetSortLen(elementCountListTail[0]), sizeof(float))); + listSortedNums[0] = elementCountListTail[0]; + } +} + +__aicore__ inline void MoeMrgsortOutPerformance::UpdateSortInfo() +{ + curLoopSortedNum = 0; + for (int64_t i = 0, j = 0; i < listNum; i++) { + if (lengths[i] > 0) { + // update remain size + listRemainElements[i] -= listSortedNums[j]; + allRemainElements -= listSortedNums[j]; + // update offset + offsets[i] += GetSortOffset(listSortedNums[j]); + // update current loop sorted nums + curLoopSortedNum += listSortedNums[j]; + j += 1; + } + } +} + +__aicore__ inline void MoeMrgsortOutPerformance::Extract() +{ + AscendC::Extract(this->ubOutput1, this->ubOutput2, this->tempBuffer, Ceil(curLoopSortedNum, ONE_REPEAT_SORT_NUM)); + Muls(this->ubOutput1, this->ubOutput1, (float)-1, Align(curLoopSortedNum, sizeof(float))); + Cast(this->ubOutputInt1, this->ubOutput1, RoundMode::CAST_ROUND, Align(curLoopSortedNum, sizeof(float))); +} + +__aicore__ inline void MoeMrgsortOutPerformance::CopyOut() +{ + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = curLoopSortedNum * sizeof(int32_t); + event_t eventIdVToMte3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventIdVToMte3); + WaitFlag(eventIdVToMte3); + DataCopyPad(this->gmOutput1[outOffset], this->ubOutputInt1, intriParams); + DataCopyPad(this->gmOutput2[outOffset], this->ubOutputInt2, intriParams); + + outOffset += curLoopSortedNum; +} + +__aicore__ inline void MoeMrgsortOutPerformance::Init(MoeMrgsortPerformanceParam *param, TPipe *tPipe) +{ + this->param = param; + for (int64_t i = 0; i < MAX_MRGSORT_LIST_TOTAL; i++) { + listRemainElements[i / MAX_MRGSORT_LIST] += static_cast(gmActualSortNum.GetValue(i)); + } + for (int64_t i = 0; i < listNum; i++) { + offsets[i] = GetSortOffset(param->perListElements * i * MAX_MRGSORT_LIST); + allRemainElements += listRemainElements[i]; + } +} + +__aicore__ inline void MoeMrgsortOutPerformance::Process() +{ + for (; allRemainElements > 0;) { + CopyIn(); + UpdateMrgParam(); + MrgsortCompute(); + UpdateSortInfo(); + Extract(); + CopyOut(); + } + ClearCache(); +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_MRGSORT_OUT_PERFORMANCE_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_mrgsort_performance.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_mrgsort_performance.h new file mode 100644 index 000000000..5f5cc97a0 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_mrgsort_performance.h @@ -0,0 +1,206 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_mrgsort_performance.h + * \brief + */ +#ifndef MOE_CUSTOM_MRGSORT_PERFORMANCE_H +#define MOE_CUSTOM_MRGSORT_PERFORMANCE_H + +#include "moe_custom_common.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +struct MoeMrgsortPerformanceParam { + int64_t perListElements; + int64_t oneLoopMaxElements; +}; + +class MoeMrgsortPerformance { +public: + __aicore__ inline MoeMrgsortPerformance(){}; + __aicore__ inline void Init(MoeMrgsortPerformanceParam *param); + __aicore__ inline void Process(); + __aicore__ inline void SetInput(GlobalTensor &gmInput, LocalTensor &ubInput, + GlobalTensor &gmActualSortNum); + __aicore__ inline void SetOutput(GlobalTensor &gmOutput, LocalTensor &ubOutput); + +private: + __aicore__ inline void CopyIn(); + __aicore__ inline void UpdateMrgParam(); + __aicore__ inline void MrgsortCompute(); + __aicore__ inline void UpdateSortInfo(); + __aicore__ inline void CopyOut(); + __aicore__ inline void ClearCache(); + +private: + MoeMrgsortPerformanceParam *param = nullptr; + + GlobalTensor gmInputs[4]; + GlobalTensor gmOutput; + GlobalTensor gmActualSortNum; + + LocalTensor ubInputs[4]; + LocalTensor ubOutput; + + int64_t listNum{0}; + int64_t remainListNum{0}; + int64_t outOffset{0}; + int64_t offsets[4]; + int64_t listRemainElements[4]; + int64_t lengths[4]; + int64_t allRemainElements{0}; + int64_t curLoopSortedNum{0}; + + // for MrgSort + uint16_t validBitTail{0}; + uint16_t elementCountListTail[4]; + uint32_t listSortedNums[4]; + LocalTensor tmpUbInputs[4]; +}; + +__aicore__ inline void MoeMrgsortPerformance::ClearCache() +{ + this->listNum = 0; + this->allRemainElements = 0; + this->outOffset = 0; +} + +__aicore__ inline void MoeMrgsortPerformance::SetInput(GlobalTensor &gmInput, LocalTensor &ubInput, + GlobalTensor &gmActualSortNum) +{ + if (this->listNum == 0) { + this->gmActualSortNum = gmActualSortNum; + } + this->gmInputs[listNum] = gmInput; + this->ubInputs[listNum] = ubInput; + this->listNum += 1; +} + +__aicore__ inline void MoeMrgsortPerformance::SetOutput(GlobalTensor &gmOutput, LocalTensor &ubOutput) +{ + this->gmOutput = gmOutput; + this->ubOutput = ubOutput; +} + +__aicore__ inline void MoeMrgsortPerformance::UpdateMrgParam() +{ + if (this->remainListNum == MERGE_LIST_TWO) { + elementCountListTail[MERGE_LIST_IDX_TWO] = 0; + elementCountListTail[MERGE_LIST_IDX_THREE] = 0; + validBitTail = 0b0011; + } else if (this->remainListNum == MERGE_LIST_THREE) { + elementCountListTail[MERGE_LIST_IDX_THREE] = 0; + validBitTail = 0b0111; + } else if (this->remainListNum == MERGE_LIST_FOUR) { + validBitTail = 0b1111; + } else { + validBitTail = 0b0001; + } +} + +__aicore__ inline void MoeMrgsortPerformance::CopyIn() +{ + this->remainListNum = 0; + event_t eventIdMte3ToMte2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + SetFlag(eventIdMte3ToMte2); + WaitFlag(eventIdMte3ToMte2); + for (int64_t i = 0, j = 0; i < listNum; i++) { + lengths[i] = Min(param->oneLoopMaxElements, listRemainElements[i]); + if (lengths[i] > 0) { + DataCopy(this->ubInputs[i], this->gmInputs[i][offsets[i]], + Align(GetSortLen(lengths[i]), sizeof(float))); + tmpUbInputs[j] = this->ubInputs[i]; + elementCountListTail[j] = lengths[i]; + this->remainListNum += 1; + j++; + } + } +} + +__aicore__ inline void MoeMrgsortPerformance::MrgsortCompute() +{ + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + if (this->remainListNum == MERGE_LIST_TWO) { + MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]); + MrgSort(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else if (this->remainListNum == MERGE_LIST_THREE) { + MrgSortSrcList sortListTail = + MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], tmpUbInputs[0]); + MrgSort(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else if (this->remainListNum == MERGE_LIST_FOUR) { + MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], + tmpUbInputs[MERGE_LIST_IDX_THREE]); + MrgSort(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1); + } else { + DataCopy(this->ubOutput, this->tmpUbInputs[0], + Align(GetSortLen(elementCountListTail[0]), sizeof(float))); + listSortedNums[0] = elementCountListTail[0]; + } +} + +__aicore__ inline void MoeMrgsortPerformance::UpdateSortInfo() +{ + curLoopSortedNum = 0; + for (int64_t i = 0, j = 0; i < listNum; i++) { + if (lengths[i] > 0) { + // update remain size + listRemainElements[i] -= listSortedNums[j]; + allRemainElements -= listSortedNums[j]; + // update offset + offsets[i] += GetSortOffset(listSortedNums[j]); + // update current loop sorted nums + curLoopSortedNum += listSortedNums[j]; + j += 1; + } + } +} + +__aicore__ inline void MoeMrgsortPerformance::CopyOut() +{ + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = GetSortLen(curLoopSortedNum) * sizeof(float); + event_t eventIdVToMte3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventIdVToMte3); + WaitFlag(eventIdVToMte3); + DataCopyPad(this->gmOutput[outOffset], this->ubOutput, intriParams); + outOffset += GetSortLen(curLoopSortedNum); +} + +__aicore__ inline void MoeMrgsortPerformance::Init(MoeMrgsortPerformanceParam *param) +{ + this->param = param; + for (int64_t i = 0; i < listNum; i++) { + offsets[i] = GetSortOffset(param->perListElements * i); + listRemainElements[i] = static_cast(gmActualSortNum.GetValue(i)); + allRemainElements += listRemainElements[i]; + } +} + +__aicore__ inline void MoeMrgsortPerformance::Process() +{ + for (; allRemainElements > 0;) { + CopyIn(); + UpdateMrgParam(); + MrgsortCompute(); + UpdateSortInfo(); + CopyOut(); + } + + ClearCache(); +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_MRGSORT_PERFORMANCE_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_row_idx_gather.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_row_idx_gather.h new file mode 100644 index 000000000..03f35a2e7 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_row_idx_gather.h @@ -0,0 +1,204 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_row_idx_gather.h + * \brief + */ +#ifndef MOE_CUSTOM_ROW_IDX_GATHER_H +#define MOE_CUSTOM_ROW_IDX_GATHER_H + +#include "moe_custom_common.h" +#include "kernel_operator.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +class RowIdxGather { +public: + __aicore__ inline RowIdxGather(){}; + __aicore__ inline void Init(GM_ADDR expandedRowIdx, GM_ADDR workspace, const MoeInitRoutingCustomTilingData *tilingData, + TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyIn(int64_t loop, int64_t elements); + __aicore__ inline void Compute(int64_t loop, int64_t elements); + __aicore__ inline void CopyOut(int64_t loop, int64_t elements, GlobalTensor &RowIdxDstGm_); + __aicore__ inline void AssistInit(); + +private: + GlobalTensor expandedRowIdxGm_; + GlobalTensor sortedExpertIndicesGm_; + GlobalTensor expertTokensCountGm_; + GlobalTensor expertTotalCountGm_; + GlobalTensor assistGm_; + GlobalTensor gatherIndicesGm_; + + TPipe *pipe_; + + TQue sortedExpertIndicesInQueue_; + TQue copyOutQueue_; + TBuf assistBuffer_; + + const MoeCustomSrcToDstComputeTilingData *srcToDstComputeTilingData_; + int64_t blockIdx_; + int64_t needCoreNum_; + int64_t perCoreElements_; + int64_t actualExpertNum_ = 0; + int64_t ep_ = 0; + int64_t rowIdxType_ = 0; + int64_t expertTotalCount_ = 0; + + int64_t loops_ = 0; + int64_t perLoopElements_ = 0; + int64_t lastLoopElements_ = 0; +}; + +__aicore__ inline void RowIdxGather::AssistInit() +{ + LocalTensor assistTensor = assistBuffer_.Get(ASSIST_NUM); + DataCopy(assistTensor, assistGm_, ASSIST_NUM); + SetWaitFlag(HardEvent::MTE2_V); + Adds(assistTensor, assistTensor, (int32_t)(blockIdx_ * perCoreElements_), ASSIST_NUM); +} + +__aicore__ inline void RowIdxGather::Init(GM_ADDR expandedRowIdx, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + pipe_ = tPipe; + srcToDstComputeTilingData_ = &(tilingData->srcToDstComputeParamsOp); + blockIdx_ = GetBlockIdx(); + actualExpertNum_ = tilingData->actualExpertNum; + ep_ = tilingData->ep; + rowIdxType_ = tilingData->rowIdxType; + + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx, actualExpertNum_); + + if (ep_) { + expertTotalCountGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + + Align(tilingData->n * tilingData->k, sizeof(int32_t)) * 2 + + Align(actualExpertNum_, sizeof(int32_t)), + actualExpertNum_); + AscendC::DataCacheCleanAndInvalid(expertTotalCountGm_); + expertTotalCount_ = expertTotalCountGm_.GetValue(0); + } else { + expertTotalCount_ = tilingData->n * tilingData->k; + } + assistGm_.SetGlobalBuffer((__gm__ int32_t *)assist, ASSIST_NUM); + perCoreElements_ = Ceil(expertTotalCount_, srcToDstComputeTilingData_->needCoreNum); + needCoreNum_ = Ceil(expertTotalCount_, perCoreElements_); + + int64_t lastCoreElements = expertTotalCount_ - (needCoreNum_ - 1) * perCoreElements_; + int64_t perCoreLoops = Ceil(perCoreElements_, srcToDstComputeTilingData_->perCorePerLoopElements); + int64_t perCorePerLoopElements = Ceil(perCoreElements_, perCoreLoops); + int64_t perCoreLastLoopElements = perCoreElements_ - (perCoreLoops - 1) * perCorePerLoopElements; + + int64_t lastCoreLoops = Ceil(lastCoreElements, srcToDstComputeTilingData_->perCorePerLoopElements); + int64_t lastCorePerLoopElements = Ceil(lastCoreElements, lastCoreLoops); + int64_t lastCoreLastLoopELements = lastCoreElements - (lastCoreLoops - 1) * lastCorePerLoopElements; + + loops_ = perCoreLoops; + if (blockIdx_ == needCoreNum_ - 1) { + loops_ = lastCoreLoops; + perLoopElements_ = lastCorePerLoopElements; + lastLoopElements_ = lastCoreLastLoopELements; + } else { + loops_ = perCoreLoops; + perLoopElements_ = perCorePerLoopElements; + lastLoopElements_ = perCoreLastLoopElements; + } + + if (rowIdxType_ == SCATTER) { + sortedExpertIndicesGm_.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx + blockIdx_ * perCoreElements_, + actualExpertNum_); + } else { + sortedExpertIndicesGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + + Align(tilingData->n * tilingData->k, sizeof(int32_t)) + + blockIdx_ * perCoreElements_, + actualExpertNum_); + } + + if ((ep_ == 0 && rowIdxType_ == SCATTER) && (blockIdx_ < needCoreNum_)) { + expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + + Align(tilingData->n * tilingData->k, sizeof(int32_t))); + } + pipe_->InitBuffer(sortedExpertIndicesInQueue_, 1, AlignBytes(perLoopElements_, sizeof(int32_t))); + pipe_->InitBuffer(copyOutQueue_, 1, Ceil(perLoopElements_, ASSIST_NUM) * ASSIST_NUM * BLOCK_BYTES); + pipe_->InitBuffer(assistBuffer_, ASSIST_NUM * sizeof(int32_t)); +} + +__aicore__ inline void RowIdxGather::Process() +{ + if (ep_ == 1 && rowIdxType_ == SCATTER) { + return; + } else { + if (blockIdx_ < needCoreNum_) { + AssistInit(); + for (int64_t loop = 0; loop < loops_; loop++) { + int64_t elements = perLoopElements_; + if (loop == loops_ - 1) { + elements = lastLoopElements_; + } + CopyIn(loop, elements); + Compute(loop, elements); + CopyOut(loop, elements, expandedRowIdxGm_); + } + } + } + AscendC::SyncAll(); +} + +__aicore__ inline void RowIdxGather::CopyIn(int64_t loop, int64_t elements) +{ + LocalTensor sortedExpertIndicesInLocal = sortedExpertIndicesInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams{static_cast(1), static_cast(elements * sizeof(int32_t)), 0, 0, + 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(sortedExpertIndicesInLocal, sortedExpertIndicesGm_[loop * perLoopElements_], dataCopyParams, + dataCopyPadParams); + sortedExpertIndicesInQueue_.EnQue(sortedExpertIndicesInLocal); +} + +__aicore__ inline void RowIdxGather::Compute(int64_t loop, int64_t elements) +{ + LocalTensor outLocal = copyOutQueue_.AllocTensor(); + LocalTensor assistTensor = assistBuffer_.Get(ASSIST_NUM); + PipeBarrier(); + int64_t loops = Ceil(elements, ASSIST_INDEX_NUM); + for (int64_t i = 0; i < loops; i++) { + Adds(outLocal[i * ASSIST_NUM], assistTensor, + static_cast(perLoopElements_ * loop + i * ASSIST_INDEX_NUM), ASSIST_NUM); + } + PipeBarrier(); + copyOutQueue_.EnQue(outLocal); +} + +__aicore__ inline void RowIdxGather::CopyOut(int64_t loop, int64_t elements, GlobalTensor &RowIdxDstGm_) +{ + LocalTensor inLocal = sortedExpertIndicesInQueue_.DeQue(); + LocalTensor outLocal = copyOutQueue_.DeQue(); + SetWaitFlag(HardEvent::MTE2_S); + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = sizeof(int32_t); + uint32_t outOffset; + for (int64_t idx = 0; idx < elements; idx++) { + outOffset = inLocal.GetValue(idx); + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(RowIdxDstGm_[outOffset], outLocal[idx * INT32_ONE_BLOCK_NUM], intriParams); + } + + sortedExpertIndicesInQueue_.FreeTensor(inLocal); + copyOutQueue_.FreeTensor(outLocal); +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_ROW_IDX_GATHER_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_row_idx_gather_droppad.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_row_idx_gather_droppad.h new file mode 100644 index 000000000..b8a94e536 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_row_idx_gather_droppad.h @@ -0,0 +1,306 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_row_idx_gather_droppad.h + * \brief + */ +#ifndef MOE_CUSTOM_ROW_IDX_GATHER_DROPPAD_H +#define MOE_CUSTOM_ROW_IDX_GATHER_DROPPAD_H + +#include "moe_custom_common.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +template +class MoeCustomSrcToDstWithCapacity { +public: + __aicore__ inline MoeCustomSrcToDstWithCapacity(){}; + __aicore__ inline void Init(GM_ADDR expandedRowIdx, GM_ADDR expandedX, GM_ADDR expandedScale, GM_ADDR workspace, + const TilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyIn(int64_t progress); + __aicore__ inline void CopyOut(int64_t progress); + __aicore__ inline void CopyOutRemain(); + __aicore__ inline void SyncAll(); + __aicore__ inline void AssistInit(); + +private: + TPipe *pipe; + TQue copyInQueue; + TQue copyOutQueue; + TQue copyOutZeroQueue; + TQue scaleOutZeroQueue; + + GlobalTensor expandDstToSrcRowGm; + GlobalTensor expandedRowIdxGm; + GlobalTensor expertIdxValueGm; + GlobalTensor expandedExpertIdxGm; + GlobalTensor expandedXGm; + GlobalTensor expandedScaleGm; + + LocalTensor outTmpLocal; + LocalTensor scaleLocal; + + const MoeCustomSrcToDstCapacityComputeTilingData *srcToDstTilingData; + int64_t coreNum; + int64_t blockIdx; + int64_t totalLength; + int64_t currentLoopRows; + int64_t coreRows; + int64_t perLoopRows; + int64_t lastLoopRows; + int64_t rowLoops; + int64_t expertCapacity; + int64_t expertNum; + int64_t cols; + int64_t perLoopCols; + int64_t lastLoopCols; + int64_t colLoops; + int64_t isInputScale_; + int64_t quantMode_; + + int64_t tokenCount = 0; + int32_t lastExpertId = -1; + int32_t lastCoreExpertId = 0; + int32_t lastCoreExpertIdNum = 0; + bool needScaleCopy = false; +}; + +template +__aicore__ inline void MoeCustomSrcToDstWithCapacity::AssistInit() +{ + if constexpr (IsSameType::value) { + LocalTensor outLocal = copyOutZeroQueue.AllocTensor(); + Duplicate(outLocal, static_cast(0), this->perLoopCols); + copyOutZeroQueue.EnQue(outLocal); + } else { + LocalTensor outLocal = copyOutZeroQueue.AllocTensor(); + Duplicate(outLocal, static_cast(0), this->perLoopCols); + copyOutZeroQueue.EnQue(outLocal); + } + if (this->needScaleCopy) { + LocalTensor scaleOutLocal = scaleOutZeroQueue.AllocTensor(); + Duplicate(scaleOutLocal, 0.0f, FP32_ONE_BLOCK_NUM); + scaleOutZeroQueue.EnQue(scaleOutLocal); + } + + if (this->blockIdx != 0) { + this->lastCoreExpertId = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2); + this->lastCoreExpertIdNum = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2 + 1); + for (int64_t i = this->blockIdx - 2; i >= 0; i--) { + int32_t lastExpertIdx = expertIdxValueGm.GetValue(i * 2); + if (lastExpertIdx < this->lastCoreExpertId) { + break; + } + int32_t lastExpertNum = expertIdxValueGm.GetValue(i * 2 + 1); + this->lastCoreExpertIdNum += lastExpertNum; + } + } +} + +template +__aicore__ inline void MoeCustomSrcToDstWithCapacity::CopyIn(int64_t progress) +{ + LocalTensor inLocal = copyInQueue.AllocTensor(); + int64_t length = Align(currentLoopRows, sizeof(int32_t)); + DataCopy(inLocal, expandDstToSrcRowGm[progress * perLoopRows], length); + DataCopy(inLocal[length], expandedExpertIdxGm[progress * perLoopRows], length); + copyInQueue.EnQue(inLocal); +} + +template +__aicore__ inline void MoeCustomSrcToDstWithCapacity::CopyOut(int64_t progress) +{ + LocalTensor inLocal = copyInQueue.DeQue(); + LocalTensor outLocal = copyOutQueue.AllocTensor(); + int64_t length = Align(currentLoopRows, sizeof(int32_t)); + DataCopyExtParams copyParams{static_cast(1), static_cast(sizeof(int32_t)), 0, 0, 0}; + DataCopyExtParams ScaleParams{1, static_cast(sizeof(float)), 0, 0, 0}; + + SetWaitFlag(HardEvent::MTE2_S); + if (this->lastExpertId == -1) { + this->lastExpertId = this->lastCoreExpertId; + this->tokenCount = this->lastCoreExpertIdNum; + } + for (int64_t idx = 0; idx < currentLoopRows; idx++) { + int32_t expertIdx = inLocal[length].GetValue(idx); + int32_t index = 0; + while (this->lastExpertId < expertIdx) { + while (this->tokenCount < this->expertCapacity) { + index = this->lastExpertId * this->expertCapacity + this->tokenCount; + if (this->needScaleCopy) { + DataCopyPad(expandedScaleGm[index], this->scaleLocal, ScaleParams); + } + int64_t col = this->perLoopCols; + for (int64_t i = 0; i < this->colLoops; i++) { + if (i == this->colLoops - 1) { + col = this->lastLoopCols; + } + DataCopyExtParams copyParams1{static_cast(1), static_cast(col * sizeof(T)), 0, + 0, 0}; + DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, + copyParams1); + } + this->tokenCount++; + } + this->tokenCount = 0; + this->lastExpertId++; + } + + if (this->tokenCount < this->expertCapacity) { + int32_t outOffset = inLocal.GetValue(idx); + index = expertIdx * this->expertCapacity + this->tokenCount; + outLocal.SetValue(0, index); + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(expandedRowIdxGm[outOffset], outLocal, copyParams); + this->tokenCount++; + } + } + copyInQueue.FreeTensor(inLocal); + copyOutQueue.FreeTensor(outLocal); +} + +template +__aicore__ inline void MoeCustomSrcToDstWithCapacity::CopyOutRemain() +{ + if (this->blockIdx != this->srcToDstTilingData->needCoreNum - 1) { + copyOutZeroQueue.FreeTensor(this->outTmpLocal); + if (this->needScaleCopy) { + scaleOutZeroQueue.FreeTensor(this->scaleLocal); + } + return; + } + DataCopyExtParams ScaleParams{1, static_cast(sizeof(float)), 0, 0, 0}; + while (this->lastExpertId < this->expertNum) { + while (this->tokenCount < this->expertCapacity) { + int32_t index = this->lastExpertId * this->expertCapacity + this->tokenCount; + if (this->needScaleCopy) { + DataCopyPad(expandedScaleGm[index], this->scaleLocal, ScaleParams); + } + int64_t col = this->perLoopCols; + for (int64_t i = 0; i < this->colLoops; i++) { + if (i == this->colLoops - 1) { + col = this->lastLoopCols; + } + DataCopyExtParams copyParams{static_cast(1), static_cast(col * sizeof(T)), 0, 0, 0}; + DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, copyParams); + SetWaitFlag(HardEvent::MTE3_S); + } + this->tokenCount++; + } + this->tokenCount = 0; + this->lastExpertId++; + } + copyOutZeroQueue.FreeTensor(this->outTmpLocal); + if (this->needScaleCopy) { + scaleOutZeroQueue.FreeTensor(this->scaleLocal); + } +} + +template +__aicore__ inline void MoeCustomSrcToDstWithCapacity::SyncAll() +{ + if (coreNum == 1) { + return; + } +#ifndef __CCE_KT_TEST__ + AscendC::SyncAll(); +#endif +} + +template +__aicore__ inline void MoeCustomSrcToDstWithCapacity::Init(GM_ADDR expandedRowIdx, GM_ADDR expandedX, + GM_ADDR expandedScale, GM_ADDR workspace, + const TilingData *tilingData, + TPipe *tPipe) +{ + int64_t blockNum = GetBlockNum(); + this->pipe = tPipe; + this->blockIdx = GetBlockIdx(); + + this->coreNum = tilingData->coreNum; + this->totalLength = tilingData->n * tilingData->k; + this->srcToDstTilingData = &(tilingData->srcToDstDropPadParamsOp); + this->expertNum = tilingData->expertNum; + this->expertCapacity = tilingData->expertCapacity; + this->cols = tilingData->cols; + this->isInputScale_ = tilingData->isInputScale; + this->quantMode_ = tilingData->quantMode; + + if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) { + this->coreRows = this->srcToDstTilingData->lastCoreRows; + this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows; + this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows; + this->rowLoops = this->srcToDstTilingData->lastCoreLoops; + } else { + this->coreRows = this->srcToDstTilingData->perCoreRows; + this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows; + this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows; + this->rowLoops = this->srcToDstTilingData->perCoreLoops; + } + this->perLoopCols = this->srcToDstTilingData->perLoopCols; + this->lastLoopCols = this->srcToDstTilingData->lastLoopCols; + this->colLoops = this->srcToDstTilingData->colLoops; + this->needScaleCopy = (this->isInputScale_ != 0 && this->quantMode_ == -1); + + expandedScaleGm.SetGlobalBuffer((__gm__ float *)expandedScale); + + int64_t length = Align(this->totalLength, sizeof(int32_t)); + expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx, length); + expandedXGm.SetGlobalBuffer((__gm__ T *)expandedX, this->expertNum * this->expertCapacity * this->cols); + + expandedExpertIdxGm.SetGlobalBuffer((__gm__ int32_t *)workspace + + this->blockIdx * this->srcToDstTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + expandDstToSrcRowGm.SetGlobalBuffer((__gm__ int32_t *)workspace + length + + this->blockIdx * this->srcToDstTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + expertIdxValueGm.SetGlobalBuffer( + (__gm__ int32_t *)workspace + length * 2 + Align(this->expertNum, sizeof(int32_t)) * 2, this->coreNum * 2); + + pipe->InitBuffer(copyInQueue, 1, AlignBytes(this->perLoopRows, sizeof(int32_t)) * 2); + pipe->InitBuffer(copyOutQueue, 1, AlignBytes(INT32_ONE_BLOCK_NUM, sizeof(int32_t))); + if constexpr (IsSameType::value) { + pipe->InitBuffer(copyOutZeroQueue, 1, AlignBytes(this->perLoopCols, sizeof(int16_t))); + } else { + pipe->InitBuffer(copyOutZeroQueue, 1, AlignBytes(this->perLoopCols, sizeof(T))); + } + if (this->needScaleCopy) { + pipe->InitBuffer(scaleOutZeroQueue, 1, BLOCK_BYTES); + } +} + +template +__aicore__ inline void MoeCustomSrcToDstWithCapacity::Process() +{ + if (this->blockIdx < this->srcToDstTilingData->needCoreNum) { + AssistInit(); + this->outTmpLocal = copyOutZeroQueue.DeQue(); + if (this->needScaleCopy) { + this->scaleLocal = scaleOutZeroQueue.DeQue(); + } + currentLoopRows = perLoopRows; + for (int64_t loop = 0; loop < this->rowLoops; loop++) { + if (loop == this->rowLoops - 1) { + currentLoopRows = lastLoopRows; + } + CopyIn(loop); + CopyOut(loop); + } + CopyOutRemain(); + } + this->SyncAll(); +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_ROW_IDX_GATHER_DROPPAD_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_row_idx_gather_droppad_dynamic.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_row_idx_gather_droppad_dynamic.h new file mode 100644 index 000000000..3a1800dc7 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_row_idx_gather_droppad_dynamic.h @@ -0,0 +1,582 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_row_idx_gather_droppad_dynamic.h + * \brief + */ +#ifndef MOE_CUSTOM_ROW_IDX_GATHER_DROPPAD_DYNAMIC_H +#define MOE_CUSTOM_ROW_IDX_GATHER_DROPPAD_DYNAMIC_H + +#include "moe_custom_common.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +template +class MoeCustomSrcToDstAndGather { +public: + __aicore__ inline MoeCustomSrcToDstAndGather(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR scale, GM_ADDR expandedRowIdx, GM_ADDR expandedX, + GM_ADDR dynamicQuantScale, GM_ADDR workspace, const TilingData *tilingData, + TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyIn(int64_t progress); + __aicore__ inline void CopyOut(int64_t progress); + __aicore__ inline void CopyOutLoops(int64_t progress); + __aicore__ inline void Compute(int32_t srcIdx, int32_t dstIdx, int32_t expertIdx); + __aicore__ inline float ComputeMax(LocalTensor &inLocal, LocalTensor &tempLocal, + LocalTensor &dynamicQuantLocal, int32_t srcIdx, int32_t expertIdx, + int64_t j); + __aicore__ inline void ComputeScale(LocalTensor &inLocal, LocalTensor &tempLocal, float scaleTemp, + int64_t dstIndex, int64_t j); + __aicore__ inline void ComputeLoops(int32_t srcIdx, int32_t dstIdx, int32_t expertIdx); + + __aicore__ inline void CopyOutRemain(); + __aicore__ inline void SyncAll(); + __aicore__ inline void AssistInit(); + +private: + TPipe *pipe; + TQue copyInQueue; + TQue copyOutQueue; + TQue copyOutZeroQueue; + + TQue inputXInQueue; + TQue smoothInQueue; + TQue calcQueue; + TQue inputXOutQueue; + TQue scaleOutQueue; + TQue scaleOutZeroQueue; + + GlobalTensor expandDstToSrcRowGm; + GlobalTensor expandedRowIdxGm; + GlobalTensor expertIdxValueGm; + GlobalTensor expandedExpertIdxGm; + GlobalTensor expandedXGm; + + GlobalTensor inputXGm; + GlobalTensor quantSmoothGm; + GlobalTensor dynamicQuantScaleGm; + GlobalTensor quantSrcGm; + + LocalTensor outTmpLocal; + LocalTensor scaleOutTmpLocal; + LocalTensor smoothLocal; + + const MoeCustomSrcToDstCapacityComputeTilingData *srcToDstTilingData; + + int64_t coreNum; + int64_t blockIdx; + int64_t totalLength; + int64_t currentLoopRows; + int64_t coreRows; + int64_t perLoopRows; + int64_t lastLoopRows; + int64_t rowLoops; + int64_t expertCapacity; + int64_t expertNum; + int64_t cols; + int64_t perLoopCols; + int64_t lastLoopCols; + int64_t colLoops; + int64_t perLoopColsAlign; + int64_t k; + int64_t colsTileLength; + int64_t smoothType; + + int64_t tokenCount = 0; + int32_t lastExpertId = -1; + int32_t lastCoreExpertId = 0; + int32_t lastCoreExpertIdNum = 0; +}; + +template +__aicore__ inline void MoeCustomSrcToDstAndGather::AssistInit() +{ + LocalTensor outLocal = copyOutZeroQueue.AllocTensor(); + Duplicate(outLocal, static_cast(0), this->perLoopCols); + copyOutZeroQueue.EnQue(outLocal); + LocalTensor scaleOutLocal = scaleOutZeroQueue.AllocTensor(); + Duplicate(scaleOutLocal, 0.0f, 8); + scaleOutZeroQueue.EnQue(scaleOutLocal); + + if (this->blockIdx != 0) { + this->lastCoreExpertId = expertIdxValueGm.GetValue((this->blockIdx - 1) * EXPERT_ID_VALUE_NUM); + this->lastCoreExpertIdNum = expertIdxValueGm.GetValue((this->blockIdx - 1) * EXPERT_ID_VALUE_NUM + 1); + for (int64_t i = this->blockIdx - 2; i >= 0; i--) { + int32_t lastExpertIdx = expertIdxValueGm.GetValue(i * EXPERT_ID_VALUE_NUM); + if (lastExpertIdx < this->lastCoreExpertId) { + break; + } + int32_t lastExpertNum = expertIdxValueGm.GetValue(i * EXPERT_ID_VALUE_NUM + 1); + this->lastCoreExpertIdNum += lastExpertNum; + } + } +} + +template +__aicore__ inline void MoeCustomSrcToDstAndGather::CopyIn(int64_t progress) +{ + LocalTensor inLocal = copyInQueue.AllocTensor(); + int64_t length = Align(currentLoopRows, sizeof(int32_t)); + DataCopy(inLocal, expandDstToSrcRowGm[progress * perLoopRows], length); + DataCopy(inLocal[length], expandedExpertIdxGm[progress * perLoopRows], length); + + copyInQueue.EnQue(inLocal); +} + +template +__aicore__ inline void MoeCustomSrcToDstAndGather::Compute(int32_t srcIdx, int32_t dstIdx, int32_t expertIdx) +{ + DataCopyExtParams copyInParams{1, static_cast(this->cols * sizeof(T)), 0, 0, 0}; + DataCopyExtParams smoothParams{1, static_cast(this->cols * sizeof(float)), 0, 0, 0}; + DataCopyExtParams copyOutParams{1, static_cast(this->cols * sizeof(int8_t)), 0, 0, 0}; + DataCopyExtParams quantScaleParams{1, static_cast(sizeof(int32_t)), 0, 0, 0}; + + LocalTensor inLocal = inputXInQueue.AllocTensor(); + + if constexpr (IsSameType::value) { + DataCopyPad(inLocal, inputXGm[srcIdx / this->k * this->cols], copyInParams, {false, 0, 0, 0}); + } else { + DataCopyPad(inLocal.template ReinterpretCast()[perLoopColsAlign], inputXGm[srcIdx / this->k * this->cols], + copyInParams, {false, 0, 0, 0}); + } + + if (smoothType == SCALE_EH) { + DataCopyPad(smoothLocal, quantSmoothGm[expertIdx * this->cols], smoothParams, {false, 0, 0, 0}); + } + + inputXInQueue.EnQue(inLocal); + smoothInQueue.EnQue(smoothLocal); + smoothLocal = smoothInQueue.DeQue(); + + inLocal = inputXInQueue.DeQue(); + + LocalTensor tempLocal = calcQueue.AllocTensor(); + LocalTensor outLocal = inputXOutQueue.AllocTensor(); + LocalTensor dynamicQuantLocal = scaleOutQueue.AllocTensor(); + + if constexpr (!IsSameType::value) { + Cast(inLocal, inLocal.template ReinterpretCast()[perLoopColsAlign], RoundMode::CAST_NONE, this->cols); + PipeBarrier(); + } + + if (smoothType != NO_SCALE) { + Mul(inLocal, inLocal, smoothLocal, this->cols); + PipeBarrier(); + } + + Abs(tempLocal, inLocal, this->cols); + PipeBarrier(); + + ReduceMax(dynamicQuantLocal, tempLocal, tempLocal, this->cols); + PipeBarrier(); + + float maxValue = dynamicQuantLocal.GetValue(0) / MAX_INT8; + + Duplicate(dynamicQuantLocal, maxValue, FP32_ONE_BLOCK_NUM); + Duplicate(tempLocal, maxValue, this->cols); + PipeBarrier(); + + Div(tempLocal, inLocal, tempLocal, this->cols); + PipeBarrier(); + + Cast(tempLocal.ReinterpretCast(), tempLocal, RoundMode::CAST_RINT, this->cols); + PipeBarrier(); + SetDeqScale((half)1.000000e+00f); + Cast(tempLocal.ReinterpretCast(), tempLocal.ReinterpretCast(), RoundMode::CAST_ROUND, this->cols); + PipeBarrier(); + Cast(outLocal, tempLocal.ReinterpretCast(), RoundMode::CAST_TRUNC, this->cols); + + calcQueue.FreeTensor(tempLocal); + inputXOutQueue.EnQue(outLocal); + scaleOutQueue.EnQue(dynamicQuantLocal); + + LocalTensor quantScaleLocal = scaleOutQueue.DeQue(); + DataCopyPad(dynamicQuantScaleGm[dstIdx], quantScaleLocal, quantScaleParams); + + outLocal = inputXOutQueue.DeQue(); + DataCopyPad(expandedXGm[dstIdx * this->cols], outLocal, copyOutParams); + + inputXInQueue.FreeTensor(inLocal); + inputXOutQueue.FreeTensor(outLocal); + scaleOutQueue.FreeTensor(quantScaleLocal); +} + +template +__aicore__ inline void MoeCustomSrcToDstAndGather::CopyOut(int64_t progress) +{ + LocalTensor inLocal = copyInQueue.DeQue(); + LocalTensor outLocal = copyOutQueue.AllocTensor(); + int64_t length = Align(currentLoopRows, sizeof(int32_t)); + DataCopyExtParams copyParams{static_cast(1), static_cast(sizeof(int32_t)), 0, 0, 0}; + DataCopyExtParams copyParams1{static_cast(1), static_cast(this->cols * sizeof(int8_t)), 0, 0, + 0}; + DataCopyExtParams quantScaleParams{1, static_cast(sizeof(int32_t)), 0, 0, 0}; + + SetWaitFlag(HardEvent::MTE2_S); + if (this->lastExpertId == -1) { + this->lastExpertId = this->lastCoreExpertId; + this->tokenCount = this->lastCoreExpertIdNum; + } + for (int64_t idx = 0; idx < currentLoopRows; idx++) { + int32_t expertIdx = inLocal[length].GetValue(idx); + int32_t index = 0; + while (this->lastExpertId < expertIdx) { + while (this->tokenCount < this->expertCapacity) { + index = this->lastExpertId * this->expertCapacity + this->tokenCount; + DataCopyPad(expandedXGm[index * this->cols], this->outTmpLocal, copyParams1); + DataCopyPad(dynamicQuantScaleGm[index], this->scaleOutTmpLocal, quantScaleParams); + this->tokenCount++; + } + this->tokenCount = 0; + this->lastExpertId++; + } + + if (this->tokenCount < this->expertCapacity) { + int32_t outOffset = inLocal.GetValue(idx); + index = expertIdx * this->expertCapacity + this->tokenCount; + outLocal.SetValue(0, index); + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(expandedRowIdxGm[outOffset], outLocal, copyParams); + Compute(outOffset, index, expertIdx); + this->tokenCount++; + } + } + copyInQueue.FreeTensor(inLocal); + copyOutQueue.FreeTensor(outLocal); +} + +template +__aicore__ inline float MoeCustomSrcToDstAndGather::ComputeMax(LocalTensor &inLocal, + LocalTensor &tempLocal, + LocalTensor &dynamicQuantLocal, + int32_t srcIdx, int32_t expertIdx, int64_t j) +{ + LocalTensor smoothLocal = smoothInQueue.AllocTensor(); + + DataCopyExtParams intriParamsT{1, static_cast(colsTileLength * sizeof(T)), 0, 0, 0}; + DataCopyExtParams intriParamsFp32{1, static_cast(colsTileLength * sizeof(float)), 0, 0, 0}; + + if constexpr (!IsSameType::value) { + DataCopyPad(inLocal.ReinterpretCast()[perLoopColsAlign], + inputXGm[srcIdx * this->cols + j * this->perLoopCols], intriParamsT, {false, 0, 0, 0}); + } else { + DataCopyPad(inLocal, inputXGm[srcIdx * this->cols + j * this->perLoopCols], intriParamsT, {false, 0, 0, 0}); + } + + inputXInQueue.EnQue(inLocal); + inLocal = inputXInQueue.DeQue(); + + if constexpr (!IsSameType::value) { + Cast(inLocal, inLocal.ReinterpretCast()[perLoopColsAlign], RoundMode::CAST_NONE, colsTileLength); + PipeBarrier(); + } + + if (smoothType != NO_SCALE) { + DataCopyPad(smoothLocal, quantSmoothGm[expertIdx * this->cols + j * this->perLoopCols], intriParamsFp32, + {false, 0, 0, 0}); + smoothInQueue.EnQue(smoothLocal); + smoothLocal = smoothInQueue.DeQue(); + + Mul(inLocal, inLocal, smoothLocal, colsTileLength); + PipeBarrier(); + } + + Abs(tempLocal, inLocal, colsTileLength); + PipeBarrier(); + + ReduceMax(dynamicQuantLocal[FP32_ONE_BLOCK_NUM], tempLocal, tempLocal, colsTileLength); + + DataCopyPad(quantSrcGm[j * this->perLoopCols], inLocal, intriParamsFp32); + smoothInQueue.FreeTensor(smoothLocal); + SetWaitFlag(HardEvent::MTE3_MTE2); + + return dynamicQuantLocal.GetValue(FP32_ONE_BLOCK_NUM); +} + +template +__aicore__ inline void MoeCustomSrcToDstAndGather::ComputeScale(LocalTensor &inLocal, + LocalTensor &tempLocal, + float scaleTemp, int64_t dstIndex, int64_t j) +{ + DataCopyExtParams copyInParams{1, static_cast(colsTileLength * sizeof(float)), 0, 0, 0}; + DataCopyExtParams copyOutParams{1, static_cast(colsTileLength * sizeof(int8_t)), 0, 0, 0}; + + LocalTensor outLocal = inputXOutQueue.AllocTensor(); + + DataCopyPad(inLocal, quantSrcGm[j * this->perLoopCols], copyInParams, {false, 0, 0, 0}); + inputXInQueue.EnQue(inLocal); + inLocal = inputXInQueue.DeQue(); + + Duplicate(tempLocal, scaleTemp, colsTileLength); + PipeBarrier(); + + Div(tempLocal, inLocal, tempLocal, colsTileLength); + PipeBarrier(); + + Cast(tempLocal.ReinterpretCast(), tempLocal, RoundMode::CAST_RINT, colsTileLength); + PipeBarrier(); + SetDeqScale((half)1.000000e+00f); + Cast(tempLocal.ReinterpretCast(), tempLocal.ReinterpretCast(), RoundMode::CAST_ROUND, + colsTileLength); + PipeBarrier(); + Cast(outLocal, tempLocal.ReinterpretCast(), RoundMode::CAST_TRUNC, colsTileLength); + + inputXOutQueue.EnQue(outLocal); + outLocal = inputXOutQueue.DeQue(); + DataCopyPad(expandedXGm[dstIndex * this->cols + j * this->perLoopCols], outLocal, copyOutParams); + + inputXOutQueue.FreeTensor(outLocal); + SetWaitFlag(HardEvent::MTE3_MTE2); +} + +template +__aicore__ inline void MoeCustomSrcToDstAndGather::ComputeLoops(int32_t srcIdx, int32_t dstIdx, + int32_t expertIdx) +{ + LocalTensor inLocal = inputXInQueue.AllocTensor(); + LocalTensor tempLocal = calcQueue.AllocTensor(); + LocalTensor quantScaleLocal = scaleOutQueue.AllocTensor(); + DataCopyExtParams quantScaleParams{1, static_cast(sizeof(int32_t)), 0, 0, 0}; + + uint32_t tmp = 0xFF7FFFFF; + float reduceMax = *((float *)&tmp); + for (int64_t j = 0; j < this->colLoops; j++) { + colsTileLength = this->perLoopCols; + if (j == this->colLoops - 1) { + colsTileLength = this->lastLoopCols; + } + float tileMax = ComputeMax(inLocal, tempLocal, quantScaleLocal, srcIdx / this->k, expertIdx, j); + reduceMax = (reduceMax > tileMax) ? reduceMax : tileMax; + } + + float scaleTemp = reduceMax / 127.0f; + Duplicate(quantScaleLocal, scaleTemp, 8); + scaleOutQueue.EnQue(quantScaleLocal); + quantScaleLocal = scaleOutQueue.DeQue(); + + DataCopyPad(dynamicQuantScaleGm[dstIdx], quantScaleLocal, quantScaleParams); + + for (int64_t j = 0; j < this->colLoops; j++) { + colsTileLength = this->perLoopCols; + if (j == this->colLoops - 1) { + colsTileLength = this->lastLoopCols; + } + ComputeScale(inLocal, tempLocal, scaleTemp, dstIdx, j); + } + + inputXInQueue.FreeTensor(inLocal); + calcQueue.FreeTensor(tempLocal); + scaleOutQueue.FreeTensor(quantScaleLocal); +} + +template +__aicore__ inline void MoeCustomSrcToDstAndGather::CopyOutLoops(int64_t progress) +{ + LocalTensor inLocal = copyInQueue.DeQue(); + LocalTensor outLocal = copyOutQueue.AllocTensor(); + int64_t length = Align(currentLoopRows, sizeof(int32_t)); + DataCopyExtParams copyParams{static_cast(1), static_cast(sizeof(int32_t)), 0, 0, 0}; + DataCopyExtParams quantScaleParams{1, static_cast(sizeof(int32_t)), 0, 0, 0}; + + SetWaitFlag(HardEvent::MTE2_S); + if (this->lastExpertId == -1) { + this->lastExpertId = this->lastCoreExpertId; + this->tokenCount = this->lastCoreExpertIdNum; + } + for (int64_t idx = 0; idx < currentLoopRows; idx++) { + int32_t expertIdx = inLocal[length].GetValue(idx); + SetWaitFlag(HardEvent::S_MTE3); + int32_t index = 0; + while (this->lastExpertId < expertIdx) { + while (this->tokenCount < this->expertCapacity) { + index = this->lastExpertId * this->expertCapacity + this->tokenCount; + int64_t col = this->perLoopCols; + DataCopyPad(dynamicQuantScaleGm[index], this->scaleOutTmpLocal, quantScaleParams); + for (int64_t i = 0; i < this->colLoops; i++) { + if (i == this->colLoops - 1) { + col = this->lastLoopCols; + } + DataCopyExtParams copyParams1{static_cast(1), static_cast(col * sizeof(int8_t)), + 0, 0, 0}; + DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, + copyParams1); + } + this->tokenCount++; + } + this->tokenCount = 0; + this->lastExpertId++; + } + + if (this->tokenCount < this->expertCapacity) { + int32_t outOffset = inLocal.GetValue(idx); + index = expertIdx * this->expertCapacity + this->tokenCount; + outLocal.SetValue(0, index); + SetWaitFlag(HardEvent::S_MTE3); + DataCopyPad(expandedRowIdxGm[outOffset], outLocal, copyParams); + if (smoothType == SCALE_EH) { + ComputeLoops(outOffset, index, expertIdx); + } else { + ComputeLoops(outOffset, index, 0); + } + SetWaitFlag(HardEvent::MTE3_S); + this->tokenCount++; + } + } + copyInQueue.FreeTensor(inLocal); + copyOutQueue.FreeTensor(outLocal); +} + +template +__aicore__ inline void MoeCustomSrcToDstAndGather::CopyOutRemain() +{ + DataCopyExtParams quantScaleParams{1, static_cast(sizeof(int32_t)), 0, 0, 0}; + if (this->blockIdx != this->srcToDstTilingData->needCoreNum - 1) { + copyOutZeroQueue.FreeTensor(this->outTmpLocal); + scaleOutZeroQueue.FreeTensor(this->scaleOutTmpLocal); + return; + } + while (this->lastExpertId < this->expertNum) { + while (this->tokenCount < this->expertCapacity) { + int32_t index = this->lastExpertId * this->expertCapacity + this->tokenCount; + int64_t col = this->perLoopCols; + DataCopyPad(dynamicQuantScaleGm[index], this->scaleOutTmpLocal, quantScaleParams); + for (int64_t i = 0; i < this->colLoops; i++) { + if (i == this->colLoops - 1) { + col = this->lastLoopCols; + } + DataCopyExtParams copyParams{static_cast(1), static_cast(col * sizeof(int8_t)), 0, + 0, 0}; + DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, copyParams); + SetWaitFlag(HardEvent::MTE3_S); + } + this->tokenCount++; + } + this->tokenCount = 0; + this->lastExpertId++; + } + copyOutZeroQueue.FreeTensor(this->outTmpLocal); + scaleOutZeroQueue.FreeTensor(this->scaleOutTmpLocal); +} + +template +__aicore__ inline void MoeCustomSrcToDstAndGather::Init(GM_ADDR x, GM_ADDR scale, GM_ADDR expandedRowIdx, + GM_ADDR expandedX, GM_ADDR dynamicQuantScale, + GM_ADDR workspace, const TilingData *tilingData, + TPipe *tPipe) +{ + int64_t blockNum = GetBlockNum(); + this->pipe = tPipe; + this->blockIdx = GetBlockIdx(); + + this->coreNum = tilingData->coreNum; + this->totalLength = tilingData->n * tilingData->k; + this->srcToDstTilingData = &(tilingData->srcToDstDropPadDynamicParamsOp); + this->expertNum = tilingData->expertNum; + this->expertCapacity = tilingData->expertCapacity; + this->cols = tilingData->cols; + this->k = tilingData->k; + this->smoothType = tilingData->smoothType; + + if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) { + this->coreRows = this->srcToDstTilingData->lastCoreRows; + this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows; + this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows; + this->rowLoops = this->srcToDstTilingData->lastCoreLoops; + } else { + this->coreRows = this->srcToDstTilingData->perCoreRows; + this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows; + this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows; + this->rowLoops = this->srcToDstTilingData->perCoreLoops; + } + this->perLoopCols = this->srcToDstTilingData->perLoopCols; + this->lastLoopCols = this->srcToDstTilingData->lastLoopCols; + this->colLoops = this->srcToDstTilingData->colLoops; + this->perLoopColsAlign = Align(this->perLoopCols, sizeof(T)); + + inputXGm.SetGlobalBuffer((__gm__ T *)x); + quantSmoothGm.SetGlobalBuffer((__gm__ float *)scale); + dynamicQuantScaleGm.SetGlobalBuffer((__gm__ float *)dynamicQuantScale); + + int64_t length = Align(this->totalLength, sizeof(int32_t)); + expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx, length); + expandedXGm.SetGlobalBuffer((__gm__ int8_t *)expandedX, this->expertNum * this->expertCapacity * this->cols); + + expandedExpertIdxGm.SetGlobalBuffer((__gm__ int32_t *)workspace + + this->blockIdx * this->srcToDstTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + expandDstToSrcRowGm.SetGlobalBuffer((__gm__ int32_t *)workspace + length + + this->blockIdx * this->srcToDstTilingData->perCoreRows, + Align(this->coreRows, sizeof(int32_t))); + expertIdxValueGm.SetGlobalBuffer( + (__gm__ int32_t *)workspace + length * 2 + Align(this->expertNum, sizeof(int32_t)) * 2, this->coreNum * 2); + if (this->colLoops > 1) { + quantSrcGm.SetGlobalBuffer((__gm__ float *)workspace + length * 2 + + Align(this->expertNum, sizeof(int32_t)) * 2 + this->coreNum * 2 + + this->blockIdx * this->cols, + this->cols * sizeof(float)); + } + + pipe->InitBuffer(copyInQueue, 1, AlignBytes(this->perLoopRows, sizeof(int32_t)) * 2); + pipe->InitBuffer(copyOutQueue, 1, AlignBytes(INT32_ONE_BLOCK_NUM, sizeof(int32_t))); + pipe->InitBuffer(copyOutZeroQueue, 1, AlignBytes(this->perLoopCols, sizeof(int16_t))); + + int64_t perLoopColsAlignBytes = AlignBytes(this->perLoopCols, sizeof(T)); + perLoopColsAlignBytes = + Max(int64_t(perLoopColsAlignBytes * sizeof(float) / sizeof(T)), int64_t(BLOCK_BYTES + BLOCK_BYTES)); + + pipe->InitBuffer(inputXInQueue, 1, perLoopColsAlignBytes); + pipe->InitBuffer(smoothInQueue, 1, AlignBytes(this->perLoopCols, sizeof(float))); + pipe->InitBuffer(calcQueue, 1, AlignBytes(this->perLoopCols, sizeof(float))); + pipe->InitBuffer(inputXOutQueue, 1, AlignBytes(this->perLoopCols, sizeof(int8_t))); + pipe->InitBuffer(scaleOutQueue, 1, BLOCK_BYTES + BLOCK_BYTES); + pipe->InitBuffer(scaleOutZeroQueue, 1, BLOCK_BYTES); +} + +template +__aicore__ inline void MoeCustomSrcToDstAndGather::Process() +{ + if (this->blockIdx < this->srcToDstTilingData->needCoreNum) { + AssistInit(); + this->outTmpLocal = copyOutZeroQueue.DeQue(); + this->scaleOutTmpLocal = scaleOutZeroQueue.DeQue(); + currentLoopRows = perLoopRows; + if (colLoops > 1) { + for (int64_t loop = 0; loop < this->rowLoops; loop++) { + if (loop == this->rowLoops - 1) { + currentLoopRows = lastLoopRows; + } + CopyIn(loop); + CopyOutLoops(loop); + } + } else { + smoothLocal = smoothInQueue.AllocTensor(); + if (smoothType == SCALE_1H) { + DataCopyExtParams smoothParams{1, static_cast(this->cols * sizeof(float)), 0, 0, 0}; + DataCopyPad(smoothLocal, quantSmoothGm, smoothParams, {false, 0, 0, 0}); + } + for (int64_t loop = 0; loop < this->rowLoops; loop++) { + if (loop == this->rowLoops - 1) { + currentLoopRows = lastLoopRows; + } + CopyIn(loop); + CopyOut(loop); + } + smoothInQueue.FreeTensor(smoothLocal); + } + CopyOutRemain(); + } +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_ROW_IDX_GATHER_DROPPAD_DYNAMIC_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_actual_expert.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_actual_expert.h new file mode 100644 index 000000000..b8c7355ae --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_actual_expert.h @@ -0,0 +1,430 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_sort_actual_expert.h + * \brief + */ +#ifndef MOE_CUSTOM_SORT_ACTUAL_EXPERT_H +#define MOE_CUSTOM_SORT_ACTUAL_EXPERT_H + +namespace MoeInitRoutingCustom { +using namespace AscendC; +constexpr int64_t MULTI_GATHERED_SORT_CORE_NUM = 16; +constexpr int64_t MULTI_GATHERED_SORT_THRSHOLD = 5632; +constexpr int64_t SINGLE_GATHERED_BUFFER_NUM = 2; +constexpr int64_t SINGLE_GATHERED_MAX_NUM = 21845; + +template +class MoeSortActualExpert { +public: + __aicore__ inline MoeSortActualExpert(){}; + __aicore__ inline void Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR expandedX, GM_ADDR expendedRowIdx, + GM_ADDR expertTokensCountOrCumsum, GM_ADDR expandedScale, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe); + __aicore__ inline bool Process(); + __aicore__ inline void multiCoreGatheredSort(); + __aicore__ inline void CopyOutExpandRowIdx(); + +private: + __aicore__ inline void CopyIn(); + __aicore__ inline void SortCompute(); + __aicore__ inline void TilingInKernel(); + __aicore__ inline void ExpertCountCompute(); + __aicore__ inline void CopyOut(); + __aicore__ inline void CopyOutExpertCount(); + +private: + TPipe *pipe; + TBuf buffer_; + TQueBind scaleCopyInQueue_; + TQue sortedNumCopyOutQueue_; + + GlobalTensor xGm_; + GlobalTensor scaleGm_; + GlobalTensor expandedXGm_; + GlobalTensor expertTokensCountOrCumsumGm_; + GlobalTensor expandedScaleGm_; + GlobalTensor expendedRowIdxGm_; + GlobalTensor expertIdxGm_; + GlobalTensor workspaceGm_; + GlobalTensor workspaceExpertIdxGm_; + GlobalTensor workspaceGatheredSortNumGm_; + GlobalTensor workspaceGatheredExpertIdxGm_; + GlobalTensor workspaceGatheredExpertIndexGm_; + + int64_t expertIdxOffset_ = 0; + int64_t expertIndexOffset_ = 0; + int64_t compareScalarMaskOffset_ = 0; + int64_t compareScalarMask0Offset_ = 0; + int64_t compareScalarMask1Offset_ = 0; + int64_t gatherMaskOffset_ = 0; + + int64_t totalLength_; + int64_t expertStart_ = 0; + int64_t expertEnd_ = 0; + int64_t actual_expert_num_ = 0; + int64_t cols_ = 0; + int64_t rowIdxType_ = 0; + int64_t isInputScale_ = 0; + int64_t k_ = 0; + + int64_t needSortNum_ = 0; + + int64_t needCoreNum_ = 0; + int64_t perCoreElements_ = 0; + int64_t lastCoreElements_ = 0; + int64_t curCoreElements_ = 0; + int64_t curCoreStartIndex_ = 0; + + bool needMultiSort = false; + + int64_t kvFactor = 2; + + static constexpr int64_t DST_BLK_STRIDE = 1; + static constexpr int64_t DST_REP_STRIDE = 8; + static constexpr int64_t MASK_STRIDE = 64; +}; + +template +__aicore__ inline void MoeSortActualExpert::CopyIn() +{ + LocalTensor expertIdx = buffer_.Get()[expertIdxOffset_ / sizeof(int32_t)]; + DataCopyExtParams dataCopyParams{static_cast(1), + static_cast(this->totalLength_ * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(expertIdx, expertIdxGm_, dataCopyParams, dataCopyPadParams); + SetWaitFlag(HardEvent::MTE2_V); +} + +template +__aicore__ inline void MoeSortActualExpert::SortCompute() +{ + LocalTensor expertIdx = buffer_.Get()[expertIdxOffset_ / sizeof(int32_t)]; + LocalTensor expertIdxFp32 = expertIdx.ReinterpretCast(); + LocalTensor gatheredExpertIdx = buffer_.Get(); + LocalTensor gatheredExpertIdxFp32 = gatheredExpertIdx.ReinterpretCast(); + + Cast(expertIdxFp32, expertIdx, RoundMode::CAST_ROUND, this->totalLength_); + PipeBarrier(); + Muls(expertIdxFp32, expertIdxFp32, (float)-1, this->totalLength_); + PipeBarrier(); + + LocalTensor compareScalarMaskLocalTensor0 = buffer_.Get()[compareScalarMask0Offset_]; + LocalTensor compareScalarMaskLocalTensor1 = buffer_.Get()[compareScalarMask1Offset_]; + LocalTensor gatherMaskLocalTensor = buffer_.Get()[gatherMaskOffset_]; + + AscendC::CompareScalar( + compareScalarMaskLocalTensor0, expertIdxFp32, static_cast(-expertStart_), AscendC::CMPMODE::LE, + (this->totalLength_ + ONE_REPEAT_COMPARE_NUM - 1) / ONE_REPEAT_COMPARE_NUM * ONE_REPEAT_COMPARE_NUM); + PipeBarrier(); + + AscendC::CompareScalar( + compareScalarMaskLocalTensor1, expertIdxFp32, static_cast(-expertEnd_), AscendC::CMPMODE::GT, + (this->totalLength_ + ONE_REPEAT_COMPARE_NUM - 1) / ONE_REPEAT_COMPARE_NUM * ONE_REPEAT_COMPARE_NUM); + PipeBarrier(); + And(gatherMaskLocalTensor.ReinterpretCast(), compareScalarMaskLocalTensor0.ReinterpretCast(), + compareScalarMaskLocalTensor1.ReinterpretCast(), + Ceil(this->totalLength_, MASK_STRIDE) * MASK_STRIDE / DST_REP_STRIDE / kvFactor); + PipeBarrier(); + + uint64_t rsvdCnt = 0; + GatherMaskParams gatherMaskParams; + gatherMaskParams.repeatTimes = 1; + gatherMaskParams.src0BlockStride = 1; + gatherMaskParams.src0RepeatStride = 8; + gatherMaskParams.src1RepeatStride = 8; + GatherMask(gatheredExpertIdxFp32, expertIdxFp32, gatherMaskLocalTensor.ReinterpretCast(), true, + static_cast(this->totalLength_), gatherMaskParams, rsvdCnt); + PipeBarrier(); + actual_expert_num_ = rsvdCnt; + // Handle actual_expert_num_ == 0 + if (actual_expert_num_ < 1) { + return; + } + int64_t needSortNum = Ceil(static_cast(rsvdCnt), ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + needSortNum_ = needSortNum; + + LocalTensor expertIndex = buffer_.Get()[expertIdxOffset_ / sizeof(int32_t)]; + LocalTensor gatheredExpertIndex = buffer_.Get()[needSortNum]; + ArithProgression(expertIndex, 0, 1, this->totalLength_); + GatherMask(gatheredExpertIndex, expertIndex, gatherMaskLocalTensor.ReinterpretCast(), true, + static_cast(this->totalLength_), gatherMaskParams, rsvdCnt); + PipeBarrier(); + if (rsvdCnt > MULTI_GATHERED_SORT_THRSHOLD) { + if (GetBlockIdx() == 0) { + SetWaitFlag(HardEvent::V_MTE3); + DataCopyExtParams copyParams{1, static_cast(rsvdCnt * sizeof(int32_t)), 0, 0, 0}; + DataCopyPad(workspaceGatheredExpertIdxGm_, gatheredExpertIdxFp32, copyParams); + DataCopyPad(workspaceGatheredExpertIndexGm_, gatheredExpertIndex, copyParams); + } + needMultiSort = true; + return; + } + int64_t duplicateNum = rsvdCnt % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + int duplicateIndex = rsvdCnt - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(gatheredExpertIdxFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + } + + PipeBarrier(); + LocalTensor concatLocal; + LocalTensor sortTempTensor = buffer_.Get()[needSortNum * kvFactor]; + Concat(concatLocal, gatheredExpertIdxFp32, sortTempTensor, needSortNum / ONE_REPEAT_SORT_NUM); + LocalTensor sortedLocal = buffer_.Get()[needSortNum * kvFactor + needSortNum * kvFactor * kvFactor]; + Sort(sortedLocal, concatLocal, gatheredExpertIndex.ReinterpretCast(), sortTempTensor, + needSortNum / ONE_REPEAT_SORT_NUM); + PipeBarrier(); + LocalTensor sortedExpertIdx = gatheredExpertIdxFp32; + LocalTensor sortedExpertIndex = gatheredExpertIndex.ReinterpretCast(); + + Extract(sortedExpertIdx, sortedExpertIndex.ReinterpretCast(), sortedLocal, + needSortNum / ONE_REPEAT_SORT_NUM); + PipeBarrier(); + + LocalTensor sortedExpertIdxInt32 = sortedExpertIdx.ReinterpretCast(); + + Muls(sortedExpertIdx, sortedExpertIdx, (float)-1, rsvdCnt); + Cast(sortedExpertIdxInt32, sortedExpertIdx, RoundMode::CAST_ROUND, rsvdCnt); +} + +template +__aicore__ inline void MoeSortActualExpert::TilingInKernel() +{ + int64_t coreNum = needMultiSort ? MULTI_GATHERED_SORT_CORE_NUM : GetBlockNum(); + perCoreElements_ = Ceil(actual_expert_num_, coreNum); + needCoreNum_ = Ceil(actual_expert_num_, perCoreElements_); + lastCoreElements_ = actual_expert_num_ - (needCoreNum_ - 1) * perCoreElements_; + if (GetBlockIdx() == needCoreNum_ - 1) { + curCoreElements_ = lastCoreElements_; + } else { + curCoreElements_ = perCoreElements_; + } + curCoreStartIndex_ = GetBlockIdx() * perCoreElements_; +} + +template +__aicore__ inline void MoeSortActualExpert::multiCoreGatheredSort() +{ + needSortNum_ = Ceil(static_cast(curCoreElements_), ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + perCoreElements_ = Ceil(this->totalLength_, MULTI_GATHERED_SORT_CORE_NUM); + + LocalTensor sortedNumOutLocal = sortedNumCopyOutQueue_.AllocTensor(); + LocalTensor gatheredExpertIdxFp32 = buffer_.Get(); + LocalTensor gatheredExpertIndex = buffer_.Get()[needSortNum_]; + DataCopyExtParams dataCopyParams{static_cast(1), static_cast(curCoreElements_ * sizeof(float)), + 0, 0, 0}; + DataCopyPadExtParams expertIdxPadParams{false, 0, 0, 0}; + DataCopyPad(gatheredExpertIdxFp32, workspaceGatheredExpertIdxGm_[curCoreStartIndex_], dataCopyParams, + expertIdxPadParams); + DataCopyPadExtParams expertIndexPadParams{false, 0, 0, 0}; + DataCopyPad(gatheredExpertIndex, workspaceGatheredExpertIndexGm_[curCoreStartIndex_], dataCopyParams, + expertIndexPadParams); + SetWaitFlag(HardEvent::MTE2_V); + + LocalTensor concatLocal; + LocalTensor sortTempTensor = buffer_.Get()[needSortNum_ * kvFactor]; + // Duplicate MIN_FP32 + int64_t duplicateNum = curCoreElements_ % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + int duplicateIndex = curCoreElements_ - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(gatheredExpertIdxFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + } + Concat(concatLocal, gatheredExpertIdxFp32, sortTempTensor, needSortNum_ / ONE_REPEAT_SORT_NUM); + LocalTensor sortedLocal = buffer_.Get()[needSortNum_ * kvFactor + needSortNum_ * kvFactor * kvFactor]; + Sort(sortedLocal, concatLocal, gatheredExpertIndex.ReinterpretCast(), sortTempTensor, + needSortNum_ / ONE_REPEAT_SORT_NUM); + + // Copy out sortedLocal for MergeSort + SetWaitFlag(HardEvent::V_MTE3); + int64_t curCoreSortedStartIndex = kvFactor * GetBlockIdx() * perCoreElements_; + dataCopyParams.blockLen = static_cast(kvFactor * curCoreElements_ * sizeof(float)); + DataCopyPad(workspaceExpertIdxGm_[curCoreSortedStartIndex], sortedLocal, dataCopyParams); + // Copyout sortedNum + sortedNumOutLocal.SetValue(0, curCoreElements_); + SetWaitFlag(HardEvent::S_MTE3); + dataCopyParams.blockLen = static_cast(sizeof(int32_t)); + DataCopyPad(workspaceGatheredSortNumGm_[GetBlockIdx()], sortedNumOutLocal, dataCopyParams); + sortedNumCopyOutQueue_.FreeTensor(sortedNumOutLocal); +} + +template +__aicore__ inline void MoeSortActualExpert::CopyOutExpandRowIdx() +{ + LocalTensor sortedExpertIndex = buffer_.Get()[needSortNum_]; + SetWaitFlag(HardEvent::V_MTE3); + if (GetBlockIdx() == 0) { + DataCopyExtParams copyParams{1, static_cast(actual_expert_num_ * sizeof(int32_t)), 0, 0, 0}; + DataCopyPad(expendedRowIdxGm_, sortedExpertIndex, copyParams); + } +} + +template +__aicore__ inline void MoeSortActualExpert::ExpertCountCompute() +{ + LocalTensor sortedExpertIdx = buffer_.Get()[curCoreStartIndex_]; + LocalTensor expertCountLocalTensor = buffer_.Get()[needSortNum_ * kvFactor]; + Duplicate(expertCountLocalTensor, 0, expertEnd_ - expertStart_); + + for (int64_t i = 0; i < curCoreElements_; i++) { + int64_t expertIdx = sortedExpertIdx.GetValue(i) - expertStart_; + int32_t curExpertCount = expertCountLocalTensor.GetValue(expertIdx); + expertCountLocalTensor.SetValue(expertIdx, curExpertCount + 1); + } + SetWaitFlag(HardEvent::S_MTE3); + DataCopyExtParams copyOutParams1{1, static_cast((expertEnd_ - expertStart_) * sizeof(int32_t)), 0, 0, 0}; + SetAtomicAdd(); + DataCopyPad(workspaceGm_, expertCountLocalTensor, copyOutParams1); + SetAtomicNone(); +} + +template +__aicore__ inline void MoeSortActualExpert::CopyOut() +{ + LocalTensor sortedExpertIndex = buffer_.Get()[needSortNum_ + curCoreStartIndex_]; + int64_t xLocalOffset = (needSortNum_ * kvFactor + ASSIST_NUM) * sizeof(int32_t) / sizeof(T); + LocalTensor xLocalTensor = buffer_.Get()[xLocalOffset]; + + for (int64_t i = 0; i < curCoreElements_; i++) { + int64_t srcRow = sortedExpertIndex.GetValue(i) / k_; + int64_t dstRow = i + curCoreStartIndex_; + SetWaitFlag(HardEvent::S_MTE2); + + LocalTensor scaleLocalTensor; + DataCopyExtParams dataCopyParams{static_cast(1), static_cast(cols_ * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(xLocalTensor, xGm_[srcRow * cols_], dataCopyParams, dataCopyPadParams); + if (isInputScale_ == 1) { + scaleLocalTensor = scaleCopyInQueue_.AllocTensor(); + DataCopyExtParams dataCopyParams2{static_cast(1), static_cast(sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams2{false, 0, 0, 0}; + DataCopyPad(scaleLocalTensor, scaleGm_[srcRow], dataCopyParams2, dataCopyPadParams2); + scaleCopyInQueue_.EnQue(scaleLocalTensor); + } + SetWaitFlag(HardEvent::MTE2_MTE3); + DataCopyExtParams copyOutParams1{1, static_cast(cols_ * sizeof(T)), 0, 0, 0}; + DataCopyPad(expandedXGm_[dstRow * cols_], xLocalTensor, copyOutParams1); + if (isInputScale_ == 1) { + scaleLocalTensor = scaleCopyInQueue_.DeQue(); + DataCopyExtParams copyOutParams2{1, static_cast(sizeof(float)), 0, 0, 0}; + DataCopyPad(expandedScaleGm_[dstRow], scaleLocalTensor, copyOutParams2); + scaleCopyInQueue_.FreeTensor(scaleLocalTensor); + } + } +} + +template +__aicore__ inline void MoeSortActualExpert::CopyOutExpertCount() +{ + LocalTensor expertCountLocalTensor = buffer_.Get()[needSortNum_ * kvFactor]; + LocalTensor expertCountLocalTensorInt64 = + buffer_.Get()[needSortNum_ * kvFactor + ASSIST_NUM].ReinterpretCast(); + DataCopyExtParams dataCopyParams{static_cast(1), + static_cast((expertEnd_ - expertStart_) * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(expertCountLocalTensor, workspaceGm_, dataCopyParams, dataCopyPadParams); + SetWaitFlag(HardEvent::MTE2_V); + Cast(expertCountLocalTensorInt64, expertCountLocalTensor, RoundMode::CAST_NONE, (expertEnd_ - expertStart_)); + SetWaitFlag(HardEvent::V_MTE3); + DataCopyExtParams copyOutParams1{1, static_cast((expertEnd_ - expertStart_) * sizeof(int64_t)), 0, 0, 0}; + DataCopyPad(expertTokensCountOrCumsumGm_, expertCountLocalTensorInt64, copyOutParams1); +} + +template +__aicore__ inline void MoeSortActualExpert::Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR expandedX, + GM_ADDR expendedRowIdx, GM_ADDR expertTokensCountOrCumsum, + GM_ADDR expandedScale, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + this->pipe = tPipe; + this->totalLength_ = tilingData->n * tilingData->k; + cols_ = tilingData->cols; + expertStart_ = tilingData->expertStart; + expertEnd_ = tilingData->expertEnd; + rowIdxType_ = tilingData->rowIdxType; + isInputScale_ = tilingData->isInputScale; + k_ = tilingData->k; + + expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx); + + expendedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expendedRowIdx); + + xGm_.SetGlobalBuffer((__gm__ T *)x); + scaleGm_.SetGlobalBuffer((__gm__ float *)scale); + expandedXGm_.SetGlobalBuffer((__gm__ T *)expandedX); + expertTokensCountOrCumsumGm_.SetGlobalBuffer((__gm__ int64_t *)expertTokensCountOrCumsum); + expandedScaleGm_.SetGlobalBuffer((__gm__ float *)expandedScale); + workspaceGm_.SetGlobalBuffer((__gm__ int32_t *)workspace, ASSIST_NUM); + if (GetBlockIdx() == 0) { + InitGlobalMemory(workspaceGm_, ASSIST_NUM, 0); + SetWaitFlag(HardEvent::MTE3_MTE2); + } + workspaceExpertIdxGm_.SetGlobalBuffer((__gm__ float *)workspace); + int64_t offset = kvFactor * Align(this->totalLength_, sizeof(int32_t)); + workspaceGatheredExpertIdxGm_.SetGlobalBuffer((__gm__ float *)workspace + offset); + offset += Align(this->totalLength_, sizeof(float)); + workspaceGatheredExpertIndexGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + offset); + offset += Align(this->totalLength_, sizeof(float)); + workspaceGatheredSortNumGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + offset); + + expertIdxOffset_ = AlignBytes(this->totalLength_, sizeof(int32_t)); + expertIndexOffset_ = expertIdxOffset_; + + gatherMaskOffset_ = expertIdxOffset_ * kvFactor; + int64_t maskOffset = + AlignBytes(Ceil(this->totalLength_, MASK_STRIDE) * MASK_STRIDE / DST_REP_STRIDE, sizeof(int8_t)); + compareScalarMask0Offset_ = gatherMaskOffset_ + maskOffset; + compareScalarMask1Offset_ = compareScalarMask0Offset_ + maskOffset; + int64_t maskOffsetMax = Ceil(SINGLE_GATHERED_MAX_NUM, MASK_STRIDE) * MASK_STRIDE / DST_REP_STRIDE; + int64_t bufferSize = + AlignBytes(SINGLE_GATHERED_MAX_NUM, sizeof(int32_t)) * kvFactor + maskOffsetMax + maskOffsetMax + maskOffsetMax; + pipe->InitBuffer(scaleCopyInQueue_, SINGLE_GATHERED_BUFFER_NUM, 32); + pipe->InitBuffer(sortedNumCopyOutQueue_, SINGLE_GATHERED_BUFFER_NUM, 32); + pipe->InitBuffer(buffer_, bufferSize); // 182992 Bytes +} + +template +__aicore__ inline bool MoeSortActualExpert::Process() +{ + CopyIn(); + SortCompute(); + TilingInKernel(); + if (needMultiSort) { + SyncAll(); + if (GetBlockIdx() < needCoreNum_) { + multiCoreGatheredSort(); + } + SyncAll(); + return false; + } + + if (GetBlockIdx() < needCoreNum_) { + CopyOutExpandRowIdx(); + } + if (GetBlockIdx() < needCoreNum_) { + ExpertCountCompute(); + CopyOut(); + } + SyncAll(); + if (GetBlockIdx() == GetBlockNum() - 1) { + CopyOutExpertCount(); + } + return true; +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_SORT_ACTUAL_EXPERT_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_base.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_base.h new file mode 100644 index 000000000..98db4b120 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_base.h @@ -0,0 +1,71 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_sort_base.h + * \brief + */ +#ifndef MOE_CUSTOM_SORT_BASE_H +#define MOE_CUSTOM_SORT_BASE_H + +#include "kernel_operator.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +class MoeSortBase { +public: + __aicore__ inline MoeSortBase(){}; + __aicore__ inline int64_t GetSyncRound(); + +protected: + __aicore__ inline void CleanWSCache(); + __aicore__ inline void SyncAll(); + +protected: + TPipe *pipe; + TQue sortDataCopyInQueue; + TQue sortDataCopyOutQueue; + TBuf tempBuffer; + TBuf sortedBuffer; + + GlobalTensor expertIdxGm; + GlobalTensor expendedRowIdxGm; + GlobalTensor sortedExpertForSourceRowGm; + GlobalTensor expandDstToSrcRowGm; + GlobalTensor sortedexpertIdxGm; + GlobalTensor expertCountTempGm; + + int64_t tileLength; + int64_t bufferNum = 1; + int64_t totalLength; + int64_t coreNum; + + int64_t expertStart_ = 0; + int64_t expertEnd_ = 0; + int64_t n; + int64_t k; + int64_t ep_ = 0; + int64_t oneLoopMaxElements_; + int64_t rowIdxType_ = 0; + + static constexpr int64_t SYNC_GM_NUM = 2; + static constexpr int64_t WORK_GM_NUM = 2; + static constexpr int64_t DST_BLK_STRIDE = 1; + static constexpr int64_t DST_REP_STRIDE = 8; +}; + +__aicore__ inline void MoeSortBase::SyncAll() +{ + AscendC::SyncAll(); +} + +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_SORT_BASE_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_multi_core.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_multi_core.h new file mode 100644 index 000000000..a3985f841 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_multi_core.h @@ -0,0 +1,377 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_sort_multi_core.h + * \brief + */ +#ifndef MOE_CUSTOM_VBS_ONE_CORE_H +#define MOE_CUSTOM_VBS_ONE_CORE_H + +#include "moe_custom_sort_base.h" +#include "moe_custom_mrgsort.h" +#include "moe_custom_mrgsort_out.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +class MoeSortMultiCore : public MoeSortBase { +public: + __aicore__ inline MoeSortMultiCore(){}; + __aicore__ inline void Init(GM_ADDR expertIdx, GM_ADDR expendedRowIdx, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void VBSProcess(); + __aicore__ inline void UBSortProcess(int64_t progress, int64_t size, int64_t sortNum); + __aicore__ inline void OneCoreVMSProcess(int64_t listNum, int64_t perListElements, int64_t lastListElements); + __aicore__ inline void VMSProcess(); + __aicore__ inline void SortOutProcess(); + __aicore__ inline void VBSCopyIn(int64_t progress, int64_t size, int64_t sortNum); + __aicore__ inline void UBSortCompute(int64_t progress, int64_t size, int64_t sortNum); + __aicore__ inline void VBSCopyOut(int64_t progress, int64_t size, int64_t sortNum); + __aicore__ inline void InitMoeMrgSort(MoeMrgsort *sorter, int64_t listNum, int64_t coreOffset, int64_t loopOffset); + __aicore__ inline void InitMoeMrgSortOut(MoeMrgsortOut *sorter, int64_t listNum, int64_t coreOffset); + +private: + GlobalTensor workspaceGms[2]; + // GlobalTensor expertTokensCountGm_; + + const MoeCustomVBSComputeTilingData *vbsTilingData; + const MoeCustomVMSMiddleComputeTilingData *vmsTilingData; + const MoeCustomSortOutComputeTilingData *sortOutTilingData; + + // for MoeMrgsort + MoeMrgsort mrgsorter; + MoeMrgsortParam mrgsortParam; + + int64_t coreNum; + int64_t blockIdx; + int64_t srcWsIndex = 0; + + int64_t listNum; + int64_t perListElements; + int64_t lastListElements; + + int64_t sortTotalLength; + int64_t sortCoreLoops; + int64_t sortCoreLoopElements; + int64_t sortCoreLastLoopElements; + + int64_t perCoreExpert; + int64_t needInitExpertCore; + int64_t currentCoreExpert; + + static constexpr int64_t MAX_MRGSORT_LIST = 4; +}; + +__aicore__ inline void MoeSortMultiCore::VBSCopyIn(int64_t progress, int64_t size, int64_t sortNum) +{ + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + int64_t inOffset = progress * sortCoreLoopElements; + DataCopyExtParams dataCopyParams{static_cast(1), static_cast(size * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal[0], expertIdxGm[inOffset], dataCopyParams, dataCopyPadParams); + + LocalTensor rowIdxLocal = inLocal[sortNum]; + int64_t startValue = this->blockIdx * this->vbsTilingData->perCoreElements + inOffset; + SetWaitFlag(HardEvent::MTE3_S); + ArithProgression(rowIdxLocal, startValue, 1, size); + sortDataCopyInQueue.EnQue(inLocal); +} + +__aicore__ inline void MoeSortMultiCore::UBSortCompute(int64_t progress, int64_t size, int64_t sortNum) +{ + LocalTensor inLocal = sortDataCopyInQueue.DeQue(); + LocalTensor expertForSourceRowLocal = inLocal[0]; + LocalTensor expertForSourceRowLocalFp32; + + expertForSourceRowLocalFp32 = expertForSourceRowLocal.ReinterpretCast(); + Cast(expertForSourceRowLocalFp32, expertForSourceRowLocal, RoundMode::CAST_ROUND, sortNum); + + Muls(expertForSourceRowLocalFp32, expertForSourceRowLocalFp32, (float)-1, sortNum); + + if (ep_) { + LocalTensor maskLocalTensor = sortedBuffer.Get(); + AscendC::CompareScalar( + maskLocalTensor, expertForSourceRowLocalFp32, static_cast(-expertStart_), AscendC::CMPMODE::GT, + (sortNum + ONE_REPEAT_COMPARE_NUM - 1) / ONE_REPEAT_COMPARE_NUM * ONE_REPEAT_COMPARE_NUM); + LocalTensor floatMinLocalTensor = tempBuffer.Get(); + Duplicate(floatMinLocalTensor, MIN_FP32, sortNum); + Select(expertForSourceRowLocalFp32, maskLocalTensor, floatMinLocalTensor, expertForSourceRowLocalFp32, + SELMODE::VSEL_TENSOR_TENSOR_MODE, sortNum); + } + + int64_t duplicateNum = size % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + int duplicateIndex = size - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expertForSourceRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + } + + LocalTensor concatLocal = expertForSourceRowLocalFp32; + LocalTensor sortedLocal = sortedBuffer.Get(GetSortLen(sortNum)); + LocalTensor outLocal = sortDataCopyOutQueue.AllocTensor(); + LocalTensor sourceRowLocal; + sourceRowLocal = inLocal[sortNum].ReinterpretCast(); + Sort(outLocal, concatLocal, sourceRowLocal, sortedLocal, sortNum / ONE_REPEAT_SORT_NUM); + + sortDataCopyOutQueue.EnQue(outLocal); + sortDataCopyInQueue.FreeTensor(inLocal); +} + +__aicore__ inline void MoeSortMultiCore::VBSCopyOut(int64_t progress, int64_t size, int64_t sortNum) +{ + LocalTensor outLocal = sortDataCopyOutQueue.DeQue(); + DataCopy(workspaceGms[0][this->blockIdx * GetSortLen(this->vbsTilingData->perCoreElements) + + GetSortLen(progress * sortCoreLoopElements)], + outLocal, Align(GetSortLen(size), sizeof(float))); + sortDataCopyOutQueue.FreeTensor(outLocal); +} + +__aicore__ inline void MoeSortMultiCore::InitMoeMrgSort(MoeMrgsort *sorter, int64_t listNum, int64_t coreOffset, + int64_t loopOffset) +{ + GlobalTensor srcWsGm = workspaceGms[srcWsIndex][blockIdx * coreOffset + loopOffset]; + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + LocalTensor outLocal = sortDataCopyOutQueue.AllocTensor(); + for (int64_t i = 0; i < listNum; i++) { + LocalTensor inLocalT = inLocal[GetSortLen(oneLoopMaxElements_) * i]; + sorter->SetInput(srcWsGm, inLocalT); + } + GlobalTensor dstWsGm = workspaceGms[1 - srcWsIndex][blockIdx * coreOffset + loopOffset]; + sorter->SetOutput(dstWsGm, outLocal); + sortDataCopyInQueue.FreeTensor(inLocal); + sortDataCopyOutQueue.FreeTensor(outLocal); +} + +__aicore__ inline void MoeSortMultiCore::InitMoeMrgSortOut(MoeMrgsortOut *sorter, int64_t listNum, int64_t coreOffset) +{ + GlobalTensor srcWsGm = workspaceGms[srcWsIndex]; + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + LocalTensor outLocal = sortDataCopyOutQueue.AllocTensor(); + + for (int64_t i = 0; i < listNum; i++) { + LocalTensor inLocalT = inLocal[GetSortLen(oneLoopMaxElements_) * i]; + sorter->SetInput(srcWsGm, inLocalT); + } + + LocalTensor outLocalV = outLocal[oneLoopMaxElements_ * MAX_MRGSORT_LIST]; + sorter->SetOutput(this->sortedexpertIdxGm, this->expendedRowIdxGm, outLocal, outLocalV); + + LocalTensor tempBuffer = sortedBuffer.Get(GetSortLen(oneLoopMaxElements_) * MAX_MRGSORT_LIST); + sorter->SetBuffer(tempBuffer); + sortDataCopyInQueue.FreeTensor(inLocal); + sortDataCopyOutQueue.FreeTensor(outLocal); +} + +__aicore__ inline void MoeSortMultiCore::OneCoreVMSProcess(int64_t listNum, int64_t perListElements, + int64_t lastListElements) +{ + int64_t coreOffset = GetSortLen(this->vbsTilingData->perCoreElements); + mrgsortParam.oneLoopMaxElements = oneLoopMaxElements_; + + for (int64_t i = 0; listNum >= 1; i++) { + int64_t loops = (listNum + MAX_MRGSORT_LIST - 1) / MAX_MRGSORT_LIST; + int64_t remainListNum = listNum - (loops - 1) * MAX_MRGSORT_LIST; + + mrgsortParam.perListElements = perListElements; + mrgsortParam.lastListElements = perListElements; + + int64_t loopOffset = GetSortLen(mrgsortParam.perListElements * MAX_MRGSORT_LIST); + for (int64_t loop = 0; loop < loops - 1; loop++) { + InitMoeMrgSort(&mrgsorter, MAX_MRGSORT_LIST, coreOffset, loop * loopOffset); + mrgsorter.Init(&mrgsortParam); + mrgsorter.Process(); + } + + mrgsortParam.perListElements = perListElements; + mrgsortParam.lastListElements = lastListElements; + InitMoeMrgSort(&mrgsorter, remainListNum, coreOffset, (loops - 1) * loopOffset); + mrgsorter.Init(&mrgsortParam); + mrgsorter.Process(); + + listNum = loops; + lastListElements = perListElements * (remainListNum - 1) + lastListElements; + perListElements = perListElements * MAX_MRGSORT_LIST; + srcWsIndex = (srcWsIndex + 1) % WORK_GM_NUM; + if (loops == 1) { + break; + } + } +} + +__aicore__ inline void MoeSortMultiCore::UBSortProcess(int64_t progress, int64_t size, int64_t sortNum) +{ + VBSCopyIn(progress, size, sortNum); + UBSortCompute(progress, size, sortNum); + VBSCopyOut(progress, size, sortNum); +} + +__aicore__ inline void MoeSortMultiCore::VBSProcess() +{ + if (this->blockIdx < this->vbsTilingData->needCoreNum) { + int64_t sortNum = Ceil(sortCoreLoopElements, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + for (int64_t loop = 0; loop < sortCoreLoops - 1; loop++) { + UBSortProcess(loop, sortCoreLoopElements, sortNum); + } + + sortNum = Ceil(sortCoreLastLoopElements, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + UBSortProcess(sortCoreLoops - 1, sortCoreLastLoopElements, sortNum); + + if (sortCoreLoops > 1) { + OneCoreVMSProcess(sortCoreLoops, sortCoreLoopElements, sortCoreLastLoopElements); + } + } + SyncAll(); +} + +__aicore__ inline void MoeSortMultiCore::VMSProcess() +{ + int64_t currentStageNeedCoreNum = this->vmsTilingData->needCoreNum; + perListElements = this->vbsTilingData->perCoreElements; + lastListElements = this->vbsTilingData->lastCoreElements; + listNum = this->vbsTilingData->needCoreNum; + + for (; listNum > MAX_MRGSORT_LIST;) { + currentStageNeedCoreNum = Ceil(listNum, MAX_MRGSORT_LIST); + int64_t coreOffset = GetSortLen(perListElements * MAX_MRGSORT_LIST); + int64_t remainListNum = listNum - (currentStageNeedCoreNum - 1) * MAX_MRGSORT_LIST; + + if (this->blockIdx < currentStageNeedCoreNum - 1) { + mrgsortParam.perListElements = perListElements; + mrgsortParam.lastListElements = perListElements; + mrgsortParam.oneLoopMaxElements = oneLoopMaxElements_; + InitMoeMrgSort(&mrgsorter, MAX_MRGSORT_LIST, coreOffset, 0); + mrgsorter.Init(&mrgsortParam); + mrgsorter.Process(); + } else if (this->blockIdx == currentStageNeedCoreNum - 1) { + mrgsortParam.perListElements = perListElements; + mrgsortParam.lastListElements = lastListElements; + mrgsortParam.oneLoopMaxElements = oneLoopMaxElements_; + InitMoeMrgSort(&mrgsorter, remainListNum, coreOffset, 0); + mrgsorter.Init(&mrgsortParam); + mrgsorter.Process(); + } + listNum = currentStageNeedCoreNum; + currentStageNeedCoreNum = Ceil(listNum, MAX_MRGSORT_LIST); + srcWsIndex = (srcWsIndex + 1) % WORK_GM_NUM; + + lastListElements = perListElements * (remainListNum - 1) + lastListElements; + perListElements = perListElements * MAX_MRGSORT_LIST; + + SyncAll(); + } +} + +__aicore__ inline void MoeSortMultiCore::SortOutProcess() +{ + if (this->blockIdx < 1) { + mrgsortParam.perListElements = perListElements; + mrgsortParam.lastListElements = lastListElements; + mrgsortParam.oneLoopMaxElements = oneLoopMaxElements_; + + MoeMrgsortOut sorter; + InitMoeMrgSortOut(&sorter, listNum, GetSortLen(perListElements)); + sorter.Init(&mrgsortParam, pipe); + sorter.Process(); + } + SyncAll(); +} + +__aicore__ inline void MoeSortMultiCore::Init(GM_ADDR expertIdx, GM_ADDR expendedRowIdx, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + this->totalLength = tilingData->n * tilingData->k; + this->coreNum = tilingData->coreNum; + this->vbsTilingData = &(tilingData->vbsComputeParamsOp); + this->vmsTilingData = &(tilingData->vmsMiddleComputeParamsOp); + this->sortOutTilingData = &(tilingData->sortOutComputeParamsOp); + + this->blockIdx = GetBlockIdx(); + this->tileLength = this->vbsTilingData->perCorePerLoopElements; + this->sortTotalLength = this->vbsTilingData->perCoreElements; + if (this->blockIdx == tilingData->vbsComputeParamsOp.needCoreNum - 1) { + this->tileLength = this->vbsTilingData->lastCorePerLoopElements; + this->sortTotalLength = this->vbsTilingData->lastCoreElements; + } + this->n = tilingData->n; + this->k = tilingData->k; + this->ep_ = tilingData->ep; + this->oneLoopMaxElements_ = ep_ ? this->sortOutTilingData->oneLoopMaxElements : MRGSORT_LIST_MAX_ELEMENT; + + expertStart_ = tilingData->expertStart; + expertEnd_ = tilingData->expertEnd; + rowIdxType_ = tilingData->rowIdxType; + + // VBS param init + if (this->blockIdx == this->vbsTilingData->needCoreNum - 1) { + sortCoreLoops = this->vbsTilingData->lastCoreLoops; + sortCoreLoopElements = this->vbsTilingData->lastCorePerLoopElements; + sortCoreLastLoopElements = this->vbsTilingData->lastCoreLastLoopElements; + } else { + sortCoreLoops = this->vbsTilingData->perCoreLoops; + sortCoreLoopElements = this->vbsTilingData->perCorePerLoopElements; + sortCoreLastLoopElements = this->vbsTilingData->perCoreLastLoopElements; + } + + this->pipe = tPipe; + expertIdxGm.SetGlobalBuffer((__gm__ int32_t *)expertIdx + + this->blockIdx * tilingData->vbsComputeParamsOp.perCoreElements, + this->sortTotalLength); + sortedexpertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace), + Align(this->totalLength, sizeof(int32_t))); + if (rowIdxType_ == SCATTER) { + expendedRowIdxGm.SetGlobalBuffer((__gm__ int32_t *)expendedRowIdx, Align(this->totalLength, sizeof(int32_t))); + } else { + expendedRowIdxGm.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(this->totalLength, sizeof(int32_t)), + Align(this->totalLength, sizeof(int32_t))); + } + + if (GetBlockIdx() == 0) { + expertCountTempGm.SetGlobalBuffer((__gm__ int32_t *)workspace + + Align(tilingData->n * tilingData->k, sizeof(int32_t)) * 2, + tilingData->actualExpertNum); + InitGlobalMemory(expertCountTempGm, tilingData->actualExpertNum, 0); + SetWaitFlag(HardEvent::MTE3_MTE2); + } + + // key and value + int64_t kvFactor = 2; + workspaceGms[0].SetGlobalBuffer((__gm__ float *)workspace + Align(this->totalLength, sizeof(int32_t)) * 2 + + tilingData->actualExpertNum, + Align(this->totalLength, sizeof(int32_t)) * kvFactor); + workspaceGms[1].SetGlobalBuffer((__gm__ float *)workspace + + Align(this->totalLength, sizeof(int32_t)) * (kvFactor + 2) + + tilingData->actualExpertNum, + Align(this->totalLength, sizeof(int32_t)) * kvFactor); + + int64_t bufferSize = Ceil(Max(oneLoopMaxElements_ * MAX_MRGSORT_LIST, sortCoreLoopElements), ONE_REPEAT_SORT_NUM) * + ONE_REPEAT_SORT_NUM * sizeof(int32_t) * kvFactor; + pipe->InitBuffer(sortDataCopyInQueue, bufferNum, bufferSize); + pipe->InitBuffer(sortDataCopyOutQueue, bufferNum, bufferSize); + pipe->InitBuffer(sortedBuffer, bufferSize); + if (ep_) { + pipe->InitBuffer(tempBuffer, bufferSize); + } +} + +__aicore__ inline void MoeSortMultiCore::Process() +{ + VBSProcess(); + VMSProcess(); + SortOutProcess(); +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_VBS_ONE_CORE_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_multi_core_performance.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_multi_core_performance.h new file mode 100644 index 000000000..1a678bd86 --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_multi_core_performance.h @@ -0,0 +1,171 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_sort_multi_core_performance.h + * \brief + */ +#ifndef MOE_CUSTOM_VBS_ONE_CORE_PERFORMANCE_H +#define MOE_CUSTOM_VBS_ONE_CORE_PERFORMANCE_H + +#include "moe_custom_sort_base.h" +#include "moe_custom_mrgsort_performance.h" +#include "moe_custom_mrgsort_out_performance.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +class MoeSortMultiCorePerformance : public MoeSortBase { +public: + __aicore__ inline MoeSortMultiCorePerformance(){}; + __aicore__ inline void Init(GM_ADDR expendedRowIdx, GM_ADDR workspace, const MoeInitRoutingCustomTilingData *tilingData, + TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void VMSProcess(); + __aicore__ inline void SortOutProcess(); + __aicore__ inline void InitMoeMrgSort(MoeMrgsortPerformance *sorter, int64_t coreOffset); + __aicore__ inline void InitMoeMrgSortOut(MoeMrgsortOutPerformance *sorter); + +private: + GlobalTensor workspaceGms[2]; + GlobalTensor workspaceGatheredSortNumGm_; + + const MoeCustomSortOutComputeTilingData *sortOutTilingData; + const MoeCustomVBSComputeTilingData *vbsTilingData; + + // for MoeMrgsortPerformance + MoeMrgsortPerformance mrgsorter; + MoeMrgsortPerformanceParam mrgsortParam; + + int64_t blockIdx; + + int64_t perListElements; + int64_t maxPerListElements; +}; + +__aicore__ inline void MoeSortMultiCorePerformance::InitMoeMrgSort(MoeMrgsortPerformance *sorter, int64_t coreOffset) +{ + GlobalTensor srcWsGm = workspaceGms[0][this->blockIdx * coreOffset]; // 0-3 + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + LocalTensor outLocal = sortDataCopyOutQueue.AllocTensor(); + GlobalTensor sortNumGm = workspaceGatheredSortNumGm_[this->blockIdx * MAX_MRGSORT_LIST]; + for (int64_t i = 0; i < MAX_MRGSORT_LIST; i++) { + LocalTensor inLocalT = inLocal[GetSortLen(maxPerListElements) * i]; + sorter->SetInput(srcWsGm, inLocalT, sortNumGm); + } + GlobalTensor dstWsGm = workspaceGms[1][this->blockIdx * coreOffset]; + sorter->SetOutput(dstWsGm, outLocal); + sortDataCopyInQueue.FreeTensor(inLocal); + sortDataCopyOutQueue.FreeTensor(outLocal); +} + +__aicore__ inline void MoeSortMultiCorePerformance::InitMoeMrgSortOut(MoeMrgsortOutPerformance *sorter) +{ + GlobalTensor srcWsGm = workspaceGms[1]; + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + LocalTensor outLocal = sortDataCopyOutQueue.AllocTensor(); + GlobalTensor sortNumGm = workspaceGatheredSortNumGm_; + for (int64_t i = 0; i < MAX_MRGSORT_LIST; i++) { + LocalTensor inLocalT = inLocal[GetSortLen(maxPerListElements) * i]; + sorter->SetInput(srcWsGm, inLocalT, sortNumGm); + } + + LocalTensor outLocalV = outLocal[maxPerListElements * MAX_MRGSORT_LIST]; + sorter->SetOutput(this->sortedexpertIdxGm, this->expendedRowIdxGm, outLocal, outLocalV); + + LocalTensor tempBuffer = sortedBuffer.Get(GetSortLen(maxPerListElements) * MAX_MRGSORT_LIST); + sorter->SetBuffer(tempBuffer); + sortDataCopyInQueue.FreeTensor(inLocal); + sortDataCopyOutQueue.FreeTensor(outLocal); +} + +__aicore__ inline void MoeSortMultiCorePerformance::VMSProcess() +{ + int64_t currentStageNeedCoreNum = MAX_MRGSORT_LIST; + int64_t coreOffset = GetSortLen(perListElements * MAX_MRGSORT_LIST); + if (this->blockIdx <= currentStageNeedCoreNum - 1) { + mrgsortParam.perListElements = perListElements; + mrgsortParam.oneLoopMaxElements = maxPerListElements; + InitMoeMrgSort(&mrgsorter, coreOffset); + mrgsorter.Init(&mrgsortParam); + mrgsorter.Process(); + } + SyncAll(); +} + +__aicore__ inline void MoeSortMultiCorePerformance::SortOutProcess() +{ + if (this->blockIdx < 1) { + mrgsortParam.perListElements = perListElements; + mrgsortParam.oneLoopMaxElements = maxPerListElements; + MoeMrgsortOutPerformance sorter; + InitMoeMrgSortOut(&sorter); + sorter.Init(&mrgsortParam, pipe); + sorter.Process(); + InitGlobalMemory(expertCountTempGm, expertEnd_ - expertStart_, 0); + SetWaitFlag(HardEvent::MTE3_MTE2); + } + SyncAll(); +} + +__aicore__ inline void MoeSortMultiCorePerformance::Init(GM_ADDR expendedRowIdx, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + this->totalLength = tilingData->n * tilingData->k; + this->blockIdx = GetBlockIdx(); + this->n = tilingData->n; + this->k = tilingData->k; + this->vbsTilingData = &(tilingData->vbsComputeParamsOp); + this->sortOutTilingData = &(tilingData->sortOutComputeParamsOp); + this->perListElements = Ceil(this->totalLength, MAX_MRGSORT_LIST_TOTAL); + this->maxPerListElements = this->sortOutTilingData->oneLoopMaxElements; + + expertStart_ = tilingData->expertStart; + expertEnd_ = tilingData->expertEnd; + rowIdxType_ = tilingData->rowIdxType; + + this->pipe = tPipe; + sortedexpertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace), + Align(this->totalLength, sizeof(int32_t))); + if (rowIdxType_ == SCATTER) { + expendedRowIdxGm.SetGlobalBuffer((__gm__ int32_t *)expendedRowIdx, Align(this->totalLength, sizeof(int32_t))); + } else { + expendedRowIdxGm.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(this->totalLength, sizeof(int32_t)), + Align(this->totalLength, sizeof(int32_t))); + } + + // key and value + int64_t kvFactor = 2; + workspaceGms[0].SetGlobalBuffer((__gm__ float *)workspace, Align(this->totalLength, sizeof(float)) * kvFactor); + workspaceGms[1].SetGlobalBuffer((__gm__ float *)workspace + Align(this->totalLength, sizeof(float)) * kvFactor, + Align(this->totalLength, sizeof(float)) * kvFactor); + workspaceGatheredSortNumGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + + Align(this->totalLength, sizeof(int32_t)) * kvFactor * kvFactor, + MAX_MRGSORT_LIST_TOTAL); + expertCountTempGm.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(this->totalLength, sizeof(int32_t)) * 2, + expertEnd_ - expertStart_); + + int64_t bufferSize = Ceil(maxPerListElements * MAX_MRGSORT_LIST, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM * + sizeof(float) * kvFactor; + pipe->InitBuffer(sortDataCopyInQueue, bufferNum, bufferSize); + pipe->InitBuffer(sortDataCopyOutQueue, bufferNum, bufferSize); + pipe->InitBuffer(sortedBuffer, bufferSize); + pipe->InitBuffer(tempBuffer, bufferSize); +} + +__aicore__ inline void MoeSortMultiCorePerformance::Process() +{ + VMSProcess(); + SortOutProcess(); +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_VBS_ONE_CORE_PERFORMANCE_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_one_core.h b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_one_core.h new file mode 100644 index 000000000..a83ee7b4c --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_custom_sort_one_core.h @@ -0,0 +1,167 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_custom_sort_one_core.h + * \brief + */ +#ifndef MOE_CUSTOM_SORT_ONE_CORE_H +#define MOE_CUSTOM_SORT_ONE_CORE_H + +#include "moe_custom_sort_base.h" + +namespace MoeInitRoutingCustom { +using namespace AscendC; + +class MoeSortOneCore : public MoeSortBase { +public: + __aicore__ inline MoeSortOneCore(){}; + __aicore__ inline void Init(GM_ADDR expertIdx, GM_ADDR expendedRowIdx, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyIn(); + __aicore__ inline void SortCompute(); + __aicore__ inline void ExpertCountCompute(); + __aicore__ inline void CopyOut(); + +private: + int64_t sortNum; +}; + +__aicore__ inline void MoeSortOneCore::CopyIn() +{ + LocalTensor inLocal = sortDataCopyInQueue.AllocTensor(); + DataCopyExtParams dataCopyParams{static_cast(1), + static_cast(this->totalLength * sizeof(int32_t)), 0, 0, 0}; + DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; + DataCopyPad(inLocal[0], expertIdxGm, dataCopyParams, dataCopyPadParams); + LocalTensor rowIdxLocal = inLocal[this->sortNum]; + ArithProgression(rowIdxLocal, 0, 1, this->sortNum); + sortDataCopyInQueue.EnQue(inLocal); +} + +__aicore__ inline void MoeSortOneCore::SortCompute() +{ + LocalTensor inLocal = sortDataCopyInQueue.DeQue(); + LocalTensor expertIdx = inLocal[0]; + LocalTensor expertIdxFp32 = expertIdx.ReinterpretCast(); + Cast(expertIdxFp32, expertIdx, RoundMode::CAST_ROUND, this->tileLength); + Muls(expertIdxFp32, expertIdxFp32, (float)-1, this->tileLength); + + if (ep_) { + LocalTensor maskLocalTensor = sortedBuffer.Get(); + AscendC::CompareScalar(maskLocalTensor, expertIdxFp32, static_cast(-expertStart_), AscendC::CMPMODE::GT, + (this->totalLength + ONE_REPEAT_COMPARE_NUM - 1) / ONE_REPEAT_COMPARE_NUM * + ONE_REPEAT_COMPARE_NUM); + LocalTensor floatMinLocalTensor = tempBuffer.Get(); + Duplicate(floatMinLocalTensor, MIN_FP32, this->tileLength); + Select(expertIdxFp32, maskLocalTensor, floatMinLocalTensor, expertIdxFp32, SELMODE::VSEL_TENSOR_TENSOR_MODE, + this->totalLength); + } + + int64_t duplicateNum = this->totalLength % ONE_REPEAT_SORT_NUM; + if (duplicateNum > 0) { + int duplicateIndex = this->totalLength - duplicateNum; + uint64_t mask0 = UINT64_MAX; + mask0 = mask0 << duplicateNum; + mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM); + uint64_t mask[2] = {mask0, 0}; + Duplicate(expertIdxFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE); + } + + LocalTensor concatLocal; + LocalTensor tempTensor = tempBuffer.Get(GetSortLen(this->sortNum)); + Concat(concatLocal, expertIdxFp32, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM); + + LocalTensor sortedLocal = sortedBuffer.Get(GetSortLen(this->sortNum)); + LocalTensor sourceRowLocal; + sourceRowLocal = inLocal[this->sortNum].ReinterpretCast(); + Sort(sortedLocal, concatLocal, sourceRowLocal, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM); + + LocalTensor outLocal = sortDataCopyOutQueue.AllocTensor(); + LocalTensor sortedExpertForSourceRowLocal = outLocal[0]; + LocalTensor expandDstToSrcRowLocal; + expandDstToSrcRowLocal = outLocal[this->sortNum].ReinterpretCast(); + Extract(sortedExpertForSourceRowLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum / ONE_REPEAT_SORT_NUM); + Muls(sortedExpertForSourceRowLocal, sortedExpertForSourceRowLocal, (float)-1, this->tileLength); + + LocalTensor expertForSourceRowLocalInt32; + expertForSourceRowLocalInt32 = sortedExpertForSourceRowLocal.ReinterpretCast(); + Cast(expertForSourceRowLocalInt32, sortedExpertForSourceRowLocal, RoundMode::CAST_ROUND, this->tileLength); + sortDataCopyOutQueue.EnQue(outLocal); + sortDataCopyInQueue.FreeTensor(inLocal); +} + +__aicore__ inline void MoeSortOneCore::CopyOut() +{ + LocalTensor outLocal = sortDataCopyOutQueue.DeQue(); + DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = this->totalLength * sizeof(int32_t); + DataCopyPad(sortedexpertIdxGm, outLocal[0], intriParams); + DataCopyPad(expendedRowIdxGm, outLocal[this->sortNum], intriParams); + sortDataCopyOutQueue.FreeTensor(outLocal); +} + +__aicore__ inline void MoeSortOneCore::Init(GM_ADDR expertIdx, GM_ADDR expendedRowIdx, GM_ADDR workspace, + const MoeInitRoutingCustomTilingData *tilingData, TPipe *tPipe) +{ + this->pipe = tPipe; + this->tileLength = Align(tilingData->vbsComputeParamsOp.lastCorePerLoopElements, sizeof(int32_t)); + this->sortNum = Ceil(this->tileLength, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM; + this->totalLength = tilingData->n * tilingData->k; + this->coreNum = tilingData->coreNum; + this->ep_ = tilingData->ep; + expertStart_ = tilingData->expertStart; + expertEnd_ = tilingData->expertEnd; + rowIdxType_ = tilingData->rowIdxType; + + expertIdxGm.SetGlobalBuffer((__gm__ int32_t *)expertIdx, this->tileLength); + sortedexpertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace), + Align(this->totalLength, sizeof(int32_t))); + if (rowIdxType_ == SCATTER) { + expendedRowIdxGm.SetGlobalBuffer((__gm__ int32_t *)expendedRowIdx, this->tileLength); + } else { + expendedRowIdxGm.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(this->tileLength, sizeof(int32_t)), + Align(this->tileLength, sizeof(int32_t))); + } + + if (GetBlockIdx() == 0) { + expertCountTempGm.SetGlobalBuffer((__gm__ int32_t *)workspace + + Align(tilingData->n * tilingData->k, sizeof(int32_t)) * 2, + tilingData->actualExpertNum); + InitGlobalMemory(expertCountTempGm, tilingData->actualExpertNum, 0); + SetWaitFlag(HardEvent::MTE3_MTE2); + } + + int64_t coreNum = GetBlockNum(); + + // key and value + int64_t kvFactor = 2; + int64_t buffSize = this->sortNum * sizeof(int32_t) * kvFactor; + pipe->InitBuffer(sortDataCopyInQueue, bufferNum, buffSize); + pipe->InitBuffer(sortDataCopyOutQueue, bufferNum, buffSize); + pipe->InitBuffer(tempBuffer, buffSize); + pipe->InitBuffer(sortedBuffer, buffSize); +} + +__aicore__ inline void MoeSortOneCore::Process() +{ + if (GetBlockIdx() < 1) { + CopyIn(); + SortCompute(); + CopyOut(); + } + this->SyncAll(); +} +} // namespace MoeInitRoutingCustom +#endif // MOE_CUSTOM_SORT_ONE_CORE_H \ No newline at end of file diff --git a/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_init_routing_custom.cpp b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_init_routing_custom.cpp new file mode 100644 index 000000000..b91983aec --- /dev/null +++ b/csrc/ascend/moe/moe_init_routing_custom/op_kernel/moe_init_routing_custom.cpp @@ -0,0 +1,412 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_init_routing_custom.cpp + * \brief + */ +#include "moe_custom_mrgsort_out.h" +#include "moe_custom_mrgsort.h" +#include "moe_custom_sort_one_core.h" +#include "moe_custom_sort_multi_core.h" +#include "moe_custom_gather_sort_multi_core.h" +#include "moe_custom_expert_tokens_count.h" +#include "moe_custom_row_idx_gather.h" +#include "moe_custom_gather_out.h" +#include "moe_custom_gather_dynamic_quant.h" +#include "moe_custom_gather_static_quant.h" +#include "moe_custom_full_load.h" +#include "moe_custom_full_load_dynamic_quant.h" +#include "moe_custom_full_load_static_quant.h" +#include "moe_custom_full_load_unquantized.h" +#include "moe_custom_sort_actual_expert.h" +#include "moe_custom_sort_multi_core_performance.h" +#include "moe_custom_row_idx_gather_droppad_dynamic.h" +#include "moe_custom_row_idx_gather_droppad.h" +#include "moe_custom_gather_out_droppad.h" +#include "moe_custom_gather_droppad_static_quant.h" + +#define MOE_INIT_ROUTING_CUSTOM_PERFORMANCE 2000000 +#define UNQUANTIZED_FULLLOAD 2100000 +#define STATIC_QUANT_FULLLOAD 2200000 +#define DYNAMIC_QUANT_GATHER_NO_SCALE_FULLLOAD 2300000 +#define DYNAMIC_QUANT_GATHER_1H_DIM_SCALE_FULLLOAD 2301000 +#define DYNAMIC_QUANT_GATHER_EH_SCALE_FULLLOAD 2302000 +#define DYNAMIC_QUANT_SCATTER_NO_SCALE_FULLLOAD 2310000 +#define DYNAMIC_QUANT_SCATTER_1H_SCALE_FULLLOAD 2311000 +#define DYNAMIC_QUANT_SCATTER_EH_SCALE_FULLLOAD 2312000 + +#define MOE_INIT_ROUTING_CUSTOM_SORTONECORE_GATHER_NODROP 1000000 +#define MOE_INIT_ROUTING_CUSTOM_SORTONECORE_SCATTER_NODROP 1001000 +#define MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_GATHER_NODROP 1100000 +#define MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_SCATTER_NODROP 1101000 + +#define MOE_INIT_ROUTING_CUSTOM_SORTONECORE_DYNAMICQUANT_GATHER_NODROP 1020000 +#define MOE_INIT_ROUTING_CUSTOM_SORTONECORE_DYNAMICQUANT_SCATTER_NODROP 1021000 +#define MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_DYNAMICQUANT_GATHER_NODROP 1120000 +#define MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_DYNAMICQUANT_SCATTER_NODROP 1121000 + +#define MOE_INIT_ROUTING_CUSTOM_SORTONECORE_QUANT_GATHER_NODROP 1010000 +#define MOE_INIT_ROUTING_CUSTOM_SORTONECORE_QUANT_SCATTER_NODROP 1011000 +#define MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_QUANT_GATHER_NODROP 1110000 +#define MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_QUANT_SCATTER_NODROP 1111000 + +#define MOE_INIT_ROUTING_CUSTOM_SORTONECORE_GATHER_DROP 1000100 +#define MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_GATHER_DROP 1100100 +#define MOE_INIT_ROUTING_CUSTOM_SORTONECORE_DYNAMICQUANT_GATHER_DROP 1020100 +#define MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_DYNAMICQUANT_GATHER_DROP 1120100 +#define MOE_INIT_ROUTING_CUSTOM_SORTONECORE_QUANT_GATHER_DROP 1010100 +#define MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_QUANT_GATHER_DROP 1110100 + +#define MOE_INIT_ROUTING_CUSTOM_GATHER_SORTONECORE_GATHER 1200000 +#define MOE_INIT_ROUTING_CUSTOM_GATHER_SORTONECORE_SCATTER 1201000 +#define MOE_INIT_ROUTING_CUSTOM_GATHER_SORTMULTICORE_GATHER 1300000 +#define MOE_INIT_ROUTING_CUSTOM_GATHER_SORTMULTICORE_SCATTER 1301000 + + +using namespace AscendC; +using namespace MoeInitRoutingCustom; +extern "C" __global__ __aicore__ void moe_init_routing_custom(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR offset, + GM_ADDR expandedX, GM_ADDR expandedRowIdx, + GM_ADDR expertTokensCountOrCumsum, GM_ADDR expandedScale, + GM_ADDR workspace, GM_ADDR tiling) +{ + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0); + if (g_coreType == AIC) { + return; + } + + GET_TILING_DATA(tilingData, tiling); + if (workspace == nullptr) { + return; + } + + GM_ADDR userWS = GetUserWorkspace(workspace); + if (userWS == nullptr) { + return; + } + + auto t = &tilingData; + + if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_PERFORMANCE)) { + TPipe fullLoadPipe; + MoeCustomFullLoad op; + op.Init(x, expertIdx, scale, offset, expandedX, expandedRowIdx, expertTokensCountOrCumsum, expandedScale, t, + &fullLoadPipe); + op.Process(); + fullLoadPipe.Destroy(); + return; + } + + if (TILING_KEY_IS(DYNAMIC_QUANT_GATHER_NO_SCALE_FULLLOAD)) { + if constexpr (!IsSameType::value) { + TPipe fullLoadPipe; + MoeCustomFullLoadDynamicQuant op; + op.Init(x, expertIdx, scale, expandedX, expandedRowIdx, expertTokensCountOrCumsum, expandedScale, userWS, t, + &fullLoadPipe); + op.Process(); + fullLoadPipe.Destroy(); + } + return; + } + + if (TILING_KEY_IS(DYNAMIC_QUANT_GATHER_1H_DIM_SCALE_FULLLOAD)) { + if constexpr (!IsSameType::value) { + TPipe fullLoadPipe; + MoeCustomFullLoadDynamicQuant op; + op.Init(x, expertIdx, scale, expandedX, expandedRowIdx, expertTokensCountOrCumsum, expandedScale, userWS, t, + &fullLoadPipe); + op.Process(); + fullLoadPipe.Destroy(); + } + return; + } + + if (TILING_KEY_IS(DYNAMIC_QUANT_GATHER_EH_SCALE_FULLLOAD)) { + if constexpr (!IsSameType::value) { + TPipe fullLoadPipe; + MoeCustomFullLoadDynamicQuant op; + op.Init(x, expertIdx, scale, expandedX, expandedRowIdx, expertTokensCountOrCumsum, expandedScale, userWS, t, + &fullLoadPipe); + op.Process(); + fullLoadPipe.Destroy(); + } + return; + } + + if (TILING_KEY_IS(DYNAMIC_QUANT_SCATTER_NO_SCALE_FULLLOAD)) { + if constexpr (!IsSameType::value) { + TPipe fullLoadPipe; + MoeCustomFullLoadDynamicQuant op; + op.Init(x, expertIdx, scale, expandedX, expandedRowIdx, expertTokensCountOrCumsum, expandedScale, userWS, t, + &fullLoadPipe); + op.Process(); + fullLoadPipe.Destroy(); + } + return; + } + + if (TILING_KEY_IS(DYNAMIC_QUANT_SCATTER_1H_SCALE_FULLLOAD)) { + if constexpr (!IsSameType::value) { + TPipe fullLoadPipe; + MoeCustomFullLoadDynamicQuant op; + op.Init(x, expertIdx, scale, expandedX, expandedRowIdx, expertTokensCountOrCumsum, expandedScale, userWS, t, + &fullLoadPipe); + op.Process(); + fullLoadPipe.Destroy(); + } + return; + } + + if (TILING_KEY_IS(DYNAMIC_QUANT_SCATTER_EH_SCALE_FULLLOAD)) { + if constexpr (!IsSameType::value) { + TPipe fullLoadPipe; + MoeCustomFullLoadDynamicQuant op; + op.Init(x, expertIdx, scale, expandedX, expandedRowIdx, expertTokensCountOrCumsum, expandedScale, userWS, t, + &fullLoadPipe); + op.Process(); + fullLoadPipe.Destroy(); + } + return; + } + + if (TILING_KEY_IS(UNQUANTIZED_FULLLOAD)) { + TPipe fullLoadPipe; + MoeCustomFullLoadUnquantized op; + op.Init(x, expertIdx, scale, expandedX, expandedRowIdx, expertTokensCountOrCumsum, expandedScale, userWS, t, + &fullLoadPipe); + op.Process(); + fullLoadPipe.Destroy(); + return; + } + + if (TILING_KEY_IS(STATIC_QUANT_FULLLOAD)) { + if constexpr (!IsSameType::value) { + TPipe fullLoadPipe; + MoeCustomFullLoadStaticQuant op; + op.Init(x, expertIdx, scale, offset, expandedX, expandedRowIdx, expertTokensCountOrCumsum, userWS, t, + &fullLoadPipe); + op.Process(); + fullLoadPipe.Destroy(); + } + return; + } + + if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTONECORE_GATHER) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTONECORE_SCATTER)) { + TPipe sortActualExpertPipe; + MoeSortActualExpert op; + bool isFinished = false; + op.Init(x, expertIdx, scale, expandedX, expandedRowIdx, expertTokensCountOrCumsum, expandedScale, userWS, t, + &sortActualExpertPipe); + isFinished = op.Process(); + sortActualExpertPipe.Destroy(); + if (isFinished) { + return; + } + } + + if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTMULTICORE_GATHER) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTMULTICORE_SCATTER)) { + TPipe gatherSortMultiCorePipe; + MoeGatherSortMultiCore op; + op.Init(expertIdx, expandedRowIdx, userWS, t, &gatherSortMultiCorePipe); + op.Process(); + gatherSortMultiCorePipe.Destroy(); + } + + if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTONECORE_GATHER) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTONECORE_SCATTER) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTMULTICORE_GATHER) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTMULTICORE_SCATTER)) { + TPipe mergeSortMultiCorePipe; + MoeSortMultiCorePerformance op; + op.Init(expandedRowIdx, userWS, t, &mergeSortMultiCorePipe); + op.Process(); + mergeSortMultiCorePipe.Destroy(); + } + + TPipe sortPipe; + if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_GATHER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_SCATTER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_QUANT_GATHER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_QUANT_SCATTER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_DYNAMICQUANT_GATHER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_DYNAMICQUANT_SCATTER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_QUANT_GATHER_DROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_DYNAMICQUANT_GATHER_DROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_GATHER_DROP)) { + MoeSortOneCore op; + op.Init(expertIdx, expandedRowIdx, userWS, t, &sortPipe); + op.Process(); + } else if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_DYNAMICQUANT_GATHER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_DYNAMICQUANT_SCATTER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_QUANT_SCATTER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_QUANT_GATHER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_SCATTER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_GATHER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_QUANT_GATHER_DROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_DYNAMICQUANT_GATHER_DROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_GATHER_DROP)) { + MoeSortMultiCore op; + op.Init(expertIdx, expandedRowIdx, userWS, t, &sortPipe); + op.Process(); + } + sortPipe.Destroy(); + + if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTONECORE_GATHER) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTONECORE_SCATTER) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTMULTICORE_GATHER) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTMULTICORE_SCATTER)) { + TPipe histogramPipe; + if (t->expertTokensNumType == CUMSUM_MODE) { + ExpertTokensCount countOp; + countOp.Init(expandedRowIdx, expertTokensCountOrCumsum, userWS, t, &histogramPipe); + countOp.Process(); + histogramPipe.Destroy(); + } else if (t->expertTokensNumType == COUNT_MODE) { + ExpertTokensCount countOp; + countOp.Init(expandedRowIdx, expertTokensCountOrCumsum, userWS, t, &histogramPipe); + countOp.Process(); + histogramPipe.Destroy(); + } else { + ExpertTokensCount countOp; + countOp.Init(expandedRowIdx, expertTokensCountOrCumsum, userWS, t, &histogramPipe); + countOp.Process(); + histogramPipe.Destroy(); + } + + } else { + if (t->dropPadMode == 1 || t->ep == 1 || t->expertTokensNumFlag != EXERPT_TOKENS_NONE) { + TPipe histogramPipe; + if (t->expertTokensNumType == CUMSUM_MODE) { + ExpertTokensCount countOp; + countOp.Init(expandedRowIdx, expertTokensCountOrCumsum, userWS, t, &histogramPipe); + countOp.Process(); + histogramPipe.Destroy(); + } else if (t->expertTokensNumType == COUNT_MODE) { + ExpertTokensCount countOp; + countOp.Init(expandedRowIdx, expertTokensCountOrCumsum, userWS, t, &histogramPipe); + countOp.Process(); + histogramPipe.Destroy(); + } else { + ExpertTokensCount countOp; + countOp.Init(expandedRowIdx, expertTokensCountOrCumsum, userWS, t, &histogramPipe); + countOp.Process(); + histogramPipe.Destroy(); + } + } + } + + if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_GATHER_DROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_GATHER_DROP)) { + TPipe rowIdxGatherDropPadPipe; + MoeCustomSrcToDstWithCapacity rowIdxGatherDropPadOp; + rowIdxGatherDropPadOp.Init(expandedRowIdx, expandedX, expandedScale, userWS, t, &rowIdxGatherDropPadPipe); + rowIdxGatherDropPadOp.Process(); + rowIdxGatherDropPadPipe.Destroy(); + } else if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_QUANT_GATHER_DROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_QUANT_GATHER_DROP)) { + TPipe rowIdxGatherDropPadPipe; + MoeCustomSrcToDstWithCapacity rowIdxGatherDropPadOp; + rowIdxGatherDropPadOp.Init(expandedRowIdx, expandedX, expandedScale, userWS, t, &rowIdxGatherDropPadPipe); + rowIdxGatherDropPadOp.Process(); + rowIdxGatherDropPadPipe.Destroy(); + } else if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_DYNAMICQUANT_GATHER_DROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_DYNAMICQUANT_GATHER_DROP)) { + if constexpr (!IsSameType::value) { + TPipe gatherPipe; + MoeCustomSrcToDstAndGather gatherDroppadDynamicQuantOp; + gatherDroppadDynamicQuantOp.Init(x, scale, expandedRowIdx, expandedX, expandedScale, userWS, t, + &gatherPipe); + gatherDroppadDynamicQuantOp.Process(); + gatherPipe.Destroy(); + } + } else { + TPipe rowIdxPipe; + RowIdxGather rowIdxGatherOp; + rowIdxGatherOp.Init(expandedRowIdx, userWS, t, &rowIdxPipe); + rowIdxGatherOp.Process(); + rowIdxPipe.Destroy(); + } + + if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_GATHER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_SCATTER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_GATHER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_SCATTER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTONECORE_SCATTER) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_GATHER_SORTMULTICORE_SCATTER)) { + TPipe gatherPipe; + if (t->ep == 1) { + MoeGatherOut gatherOp; + gatherOp.Init(x, scale, userWS, expandedRowIdx, expandedX, expandedScale, t, &gatherPipe); + gatherOp.Process(); + gatherPipe.Destroy(); + } else { + MoeGatherOut gatherOp; + gatherOp.Init(x, scale, userWS, expandedRowIdx, expandedX, expandedScale, t, &gatherPipe); + gatherOp.Process(); + gatherPipe.Destroy(); + } + + } else if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_DYNAMICQUANT_SCATTER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_DYNAMICQUANT_GATHER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_DYNAMICQUANT_SCATTER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_DYNAMICQUANT_GATHER_NODROP)) { + if constexpr (!IsSameType::value) { + TPipe gatherPipe; + if (t->ep == 0 and t->smoothType != SCALE_EH) { + MoeGatherOutDynamicQuant gatherDynamicQuantOp; + gatherDynamicQuantOp.Init(x, scale, userWS, expandedRowIdx, expandedX, expandedScale, t, &gatherPipe); + gatherDynamicQuantOp.Process(); + gatherPipe.Destroy(); + } else { + MoeGatherOutDynamicQuant gatherDynamicQuantOp; + gatherDynamicQuantOp.Init(x, scale, userWS, expandedRowIdx, expandedX, expandedScale, t, &gatherPipe); + gatherDynamicQuantOp.Process(); + gatherPipe.Destroy(); + } + } + } else if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_QUANT_SCATTER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_QUANT_GATHER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_QUANT_SCATTER_NODROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_QUANT_GATHER_NODROP)) { + if constexpr (!IsSameType::value) { + TPipe gatherPipe; + if (t->ep == 1) { + MoeGatherOutQuant gatherStaticQuantOp; + gatherStaticQuantOp.Init(x, scale, offset, expandedRowIdx, expandedX, userWS, t, &gatherPipe); + gatherStaticQuantOp.Process(); + gatherPipe.Destroy(); + } else { + MoeGatherOutQuant gatherStaticQuantOp; + gatherStaticQuantOp.Init(x, scale, offset, expandedRowIdx, expandedX, userWS, t, &gatherPipe); + gatherStaticQuantOp.Process(); + gatherPipe.Destroy(); + } + } + } else if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_GATHER_DROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_GATHER_DROP)) { + TPipe gatherPipe; + MoeGatherOutDroppad gatherDroppadOp; + gatherDroppadOp.Init(x, scale, expandedRowIdx, expandedX, expandedScale, userWS, t, &gatherPipe); + gatherDroppadOp.Process(); + gatherPipe.Destroy(); + } else if (TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTONECORE_QUANT_GATHER_DROP) || + TILING_KEY_IS(MOE_INIT_ROUTING_CUSTOM_SORTMULTICORE_QUANT_GATHER_DROP)) { + if constexpr (!IsSameType::value) { + TPipe gatherPipe; + MoeGatherDroppadQuant gatherDroppadStaticQuantOp; + gatherDroppadStaticQuantOp.Init(x, scale, offset, expandedRowIdx, expandedX, userWS, t, &gatherPipe); + gatherDroppadStaticQuantOp.Process(); + gatherPipe.Destroy(); + } + } +} \ No newline at end of file diff --git a/csrc/ascend/moe/scatter_nd_update_v2/CMakeLists.txt b/csrc/ascend/moe/scatter_nd_update_v2/CMakeLists.txt new file mode 100644 index 000000000..739a1b08e --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/CMakeLists.txt @@ -0,0 +1,19 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +#/ + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_host/CMakeLists.txt b/csrc/ascend/moe/scatter_nd_update_v2/op_host/CMakeLists.txt new file mode 100644 index 000000000..a71ea3b7d --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_host/CMakeLists.txt @@ -0,0 +1,33 @@ +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +#/ + +#optiling and opapi +#message(STATUS "=== Debug: start ops.index.scatter_nd_update_v2.CMakeLists.txt ") +#add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE scatter_nd_update_v2 ACLNNTYPE aclnn_exclude) + + +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnnExc PRIVATE + scatter_nd_update_v2_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME ScatterNdUpdateV2 + OPTIONS --cce-auto-sync=on + -Wno-deprecated-declarations + -Werror +) + + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE scatter_nd_update_v2 ACLNNTYPE aclnn_exclude) +endif() \ No newline at end of file diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_host/config/ascend910_93/scatter_nd_update_v2_binary.json b/csrc/ascend/moe/scatter_nd_update_v2/op_host/config/ascend910_93/scatter_nd_update_v2_binary.json new file mode 100644 index 000000000..649d4a13a --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_host/config/ascend910_93/scatter_nd_update_v2_binary.json @@ -0,0 +1,821 @@ +{ + "op_type": "ScatterNdUpdateV2", + "op_list": [ + { + "bin_filename": "ScatterNdUpdateV2_float16_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "float16", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "float16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "float16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_float16_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "float16", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "float16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "float16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_float32_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "float32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "float32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "float32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_float32_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "float32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "float32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "float32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_bfloat16_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "bfloat16", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "bfloat16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "bfloat16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_bfloat16_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "bfloat16", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "bfloat16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "bfloat16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int8_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int8", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int8", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int8", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int8_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int8", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int8", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int8", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int16_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int16", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int16_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int16", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int32_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int32_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int64_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int64_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_bool_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "bool", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "bool", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "bool", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_bool_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "bool", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "bool", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "bool", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + } + ] +} diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_host/config/ascend910_93/scatter_nd_update_v2_simplified_key.ini b/csrc/ascend/moe/scatter_nd_update_v2/op_host/config/ascend910_93/scatter_nd_update_v2_simplified_key.ini new file mode 100644 index 000000000..a7ebbffa8 --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_host/config/ascend910_93/scatter_nd_update_v2_simplified_key.ini @@ -0,0 +1,13 @@ +; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示： +; [某算子] +; default=xx +; ascendxx=xx +; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置； +; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理 +; 2)如果仅有default配置：各个版本按default配置 +; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理 +; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。 +; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式 +; 6)是否是AscendC算子，由 ops/built-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定 +[ScatterNdUpdateV2] +default=0 diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_host/config/ascend910b/scatter_nd_update_v2_binary.json b/csrc/ascend/moe/scatter_nd_update_v2/op_host/config/ascend910b/scatter_nd_update_v2_binary.json new file mode 100644 index 000000000..649d4a13a --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_host/config/ascend910b/scatter_nd_update_v2_binary.json @@ -0,0 +1,821 @@ +{ + "op_type": "ScatterNdUpdateV2", + "op_list": [ + { + "bin_filename": "ScatterNdUpdateV2_float16_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "float16", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "float16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "float16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_float16_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "float16", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "float16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "float16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_float32_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "float32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "float32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "float32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_float32_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "float32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "float32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "float32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_bfloat16_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "bfloat16", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "bfloat16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "bfloat16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_bfloat16_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "bfloat16", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "bfloat16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "bfloat16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int8_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int8", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int8", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int8", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int8_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int8", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int8", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int8", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int16_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int16", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int16_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int16", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int16", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int32_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int32_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int64_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_int64_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_bool_int32", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "bool", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int32", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "bool", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "bool", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + }, + { + "bin_filename": "ScatterNdUpdateV2_bool_int64", + "inputs": [ + { + "name": "var", + "index": 0, + "dtype": "bool", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "indices", + "index": 1, + "dtype": "int64", + "format": "ND", + "paramType": "required", + "shape": [-2] + }, + { + "name": "updates", + "index": 2, + "dtype": "bool", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "outputs": [ + { + "name": "var", + "index": 0, + "dtype": "bool", + "format": "ND", + "paramType": "required", + "shape": [-2] + } + ], + "attrs": [ + { + "name": "strides", + "dtype": "list_int", + "value": null + }, + { + "name": "use_locking", + "dtype": "bool", + "value": null + } + ] + } + ] +} diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_host/config/ascend910b/scatter_nd_update_v2_simplified_key.ini b/csrc/ascend/moe/scatter_nd_update_v2/op_host/config/ascend910b/scatter_nd_update_v2_simplified_key.ini new file mode 100644 index 000000000..a7ebbffa8 --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_host/config/ascend910b/scatter_nd_update_v2_simplified_key.ini @@ -0,0 +1,13 @@ +; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示： +; [某算子] +; default=xx +; ascendxx=xx +; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置； +; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理 +; 2)如果仅有default配置：各个版本按default配置 +; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理 +; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。 +; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式 +; 6)是否是AscendC算子，由 ops/built-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定 +[ScatterNdUpdateV2] +default=0 diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_host/op_api/aclnn_scatter_nd_update_v2.cpp b/csrc/ascend/moe/scatter_nd_update_v2/op_host/op_api/aclnn_scatter_nd_update_v2.cpp new file mode 100644 index 000000000..4486036e2 --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_host/op_api/aclnn_scatter_nd_update_v2.cpp @@ -0,0 +1,142 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file aclnn_scatter_nd_update_v2.cpp + * \brief + */ + +#include "aclnn_scatter_nd_update_v2.h" +#include "scatter_nd_update_v2.h" +#include "level0/broadcast_to.h" +#include "aclnn_kernels/contiguous.h" +#include "level0/squeeze.h" +#include "level0/unsqueeze.h" +#include "aclnn_kernels/common/op_error_check.h" +#include "opdev/make_op_executor.h" +#include "opdev/platform.h" +#include "opdev/op_dfx.h" +#include "opdev/tensor_view_utils.h" +#include "opdev/op_executor.h" +using namespace op; +#ifdef __cplusplus +extern "C" { +#endif + +// 根据API定义，需要列出所能支持的所有dtype +static const std::initializer_list ASCEND910_DTYPE_DTYPE_SUPPORT_LIST = { + op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BOOL}; + +static const std::initializer_list ASCEND910B_DTYPE_DTYPE_SUPPORT_LIST = { + op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BOOL, op::DataType::DT_INT16, + op::DataType::DT_BF16, op::DataType::DT_INT64, op::DataType::DT_INT8, op::DataType::DT_INT32}; + +static const std::initializer_list ASCEND950_DTYPE_DTYPE_SUPPORT_LIST = { + op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BOOL, + op::DataType::DT_BF16, op::DataType::DT_INT64, op::DataType::DT_INT8, + op::DataType::DT_FLOAT8_E5M2, op::DataType::DT_FLOAT8_E4M3FN, op::DataType::DT_FLOAT8_E8M0}; + +static const std::initializer_list INDEX_DTYPE_SUPPORT_LIST = {op::DataType::DT_INT64, + op::DataType::DT_INT32}; + +static bool CheckNotNull(aclTensor *varRef, const aclTensor *indices, const aclTensor *updates) { + OP_CHECK_NULL(varRef, return false); + OP_CHECK_NULL(indices, return false); + OP_CHECK_NULL(updates, return false); + return true; +} + +static const std::initializer_list& GetDtypeSupportList() { + // if (Ops::NN::AclnnUtil::IsRegbase()) { + // return ASCEND950_DTYPE_DTYPE_SUPPORT_LIST; + // } + if (GetCurrentPlatformInfo().GetSocVersion() >= SocVersion::ASCEND910B && + GetCurrentPlatformInfo().GetSocVersion() <= SocVersion::ASCEND910E) { + return ASCEND910B_DTYPE_DTYPE_SUPPORT_LIST; + } else { + return ASCEND910_DTYPE_DTYPE_SUPPORT_LIST; + } +} + +static bool CheckDtypeValid(aclTensor *varRef, const aclTensor *indices, const aclTensor *updates) { + // 检查self的数据类型是否在算子的支持列表内 + auto supportList = GetDtypeSupportList(); + OP_CHECK_DTYPE_NOT_SUPPORT(varRef, supportList, return false); + // 检查index的数据类型是否在算子的支持列表内 + OP_CHECK_DTYPE_NOT_SUPPORT(indices, INDEX_DTYPE_SUPPORT_LIST, return false); + // varRef和updates的数据类型要一致 + if (varRef->GetDataType() != updates->GetDataType()) { + OP_LOGE(ACLNN_ERR_PARAM_INVALID, "updates dtype %s should be in same with varRef dtype %s.", + op::ToString(updates->GetDataType()).GetString(), op::ToString(varRef->GetDataType()).GetString()); + return false; + } + + return true; +} + +static aclnnStatus CheckParams(aclTensor *varRef, const aclTensor *indices, const aclTensor *updates) { + // 1. 检查参数是否为空指针 + CHECK_RET(CheckNotNull(varRef, indices, updates), ACLNN_ERR_PARAM_NULLPTR); + + // 2. 检查输入的数据类型是否在API支持的数据类型范围之内 + CHECK_RET(CheckDtypeValid(varRef, indices, updates), ACLNN_ERR_PARAM_INVALID); + + return ACLNN_SUCCESS; +} + +aclnnStatus aclnnScatterNdUpdateV2GetWorkspaceSize(aclTensor *varRef, const aclTensor *indices, const aclTensor *updates, + const aclIntArray *strides, uint64_t* workspaceSize, aclOpExecutor** executor) { + L2_DFX_PHASE_1(aclnnScatterNdUpdateV2, DFX_IN(varRef, indices, updates), DFX_OUT(varRef)); + + // 固定写法，创建OpExecutor + auto uniqueExecutor = CREATE_EXECUTOR(); + CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); + + // 固定写法，参数检查 + auto ret = CheckParams(varRef, indices, updates); + CHECK_RET(ret == ACLNN_SUCCESS, ret); + + if (varRef->IsEmpty() || indices->IsEmpty() || updates->IsEmpty()) { + *workspaceSize = 0; + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; + } + + // 将输入varRef转换成连续的tensor + // auto varRefContiguous = l0op::Contiguous(varRef, uniqueExecutor.get()); + // CHECK_RET(varRefContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR); + varRef->SetStorageShape(varRef->GetViewShape()); + // 将输入indices转换成连续的tensor + auto indicesContiguous = l0op::Contiguous(indices, uniqueExecutor.get()); + CHECK_RET(indicesContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR); + + // 将输入updates转换成连续的tensor + auto updatesContiguous = l0op::Contiguous(updates, uniqueExecutor.get()); + CHECK_RET(updatesContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR); + + // 执行L0算子 (kernel直接原地更新varRef, 无需Contiguous和ViewCopy) + auto scatterUpdateRes = l0op::ScatterNdUpdateV2(varRef, indicesContiguous, updatesContiguous, strides, false, uniqueExecutor.get()); + CHECK_RET(scatterUpdateRes != nullptr, ACLNN_ERR_INNER_NULLPTR); + + // 获取计算过程中需要使用的workspace大小 + *workspaceSize = uniqueExecutor->GetWorkspaceSize(); + uniqueExecutor.ReleaseTo(executor); + return ACLNN_SUCCESS; +} + +aclnnStatus aclnnScatterNdUpdateV2(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream) { + L2_DFX_PHASE_2(aclnnScatterNdUpdateV2); + // 固定写法，调用框架能力，完成计算 + return CommonOpExecutorRun(workspace, workspaceSize, executor, stream); +} + +#ifdef __cplusplus +} +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_host/op_api/aclnn_scatter_nd_update_v2.h b/csrc/ascend/moe/scatter_nd_update_v2/op_host/op_api/aclnn_scatter_nd_update_v2.h new file mode 100644 index 000000000..64e4baf98 --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_host/op_api/aclnn_scatter_nd_update_v2.h @@ -0,0 +1,56 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_API_INC_SCATTER_ND_UPDATE_V2_H_ +#define OP_API_INC_SCATTER_ND_UPDATE_V2_H_ + +#include "aclnn/aclnn_base.h" +#include "aclnn_util.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @brief aclnnScatterNdUpdateV2的第一段接口，根据具体的计算流程，计算workspace大小。 + * @domain aclnnop_ops_infer + * @domain aclnnop_ops_train + * 算子功能： 将tensor updates中的值按指定的索引indices逐个更新tensor var中的值。 + * @param [in] varRef: npu device侧的aclTensor, 数据类型支持FLOAT16, FLOAT32, BOOL + * INT64，BFLOAT16，支持非连续的Tensor，数据格式支持ND。 + * @param [in] indices: npu device侧的aclTensor，数据类型支持INT32, INT64类型。支持非连续的Tensor，数据格式支持ND。 + * @param [in] updates: npu device侧的aclTensor，数据类型支持FLOAT16, FLOAT32, BOOL + * INT64，BFLOAT16，支持非连续的Tensor，数据格式支持ND, + * @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。 + * @param [out] executor: 返回op执行器，包含算子计算流程。 + * @return aclnnStatus: 返回状态码 + */ +ACLNN_API aclnnStatus aclnnScatterNdUpdateV2GetWorkspaceSize(aclTensor* varRef, const aclTensor* indices, + const aclTensor* updates, const aclIntArray* strides, uint64_t* workspaceSize, + aclOpExecutor** executor); + +/** + * @brief: aclnnScatterNdUpdateV2的第二段接口，用于执行计算 + * @domain aclnnop_ops_infer + * @domain aclnnop_ops_train + * 算子功能： 将tensor updates中的值按指定的索引indices逐个更新tensor var中的值。 + * @param [in] workspace: 在npu device侧申请的workspace内存起址。 + * @param [in] workspaceSize: 在npu device侧申请的workspace大小，由第一段接口aclnnScatterNdUpdateV2GetWorkspaceSize获取。 + * @param [in] stream: acl stream流。 + * @param [in] executor: op执行器，包含了算子计算流程。 + * @return aclnnStatus: 返回状态码。 + */ +ACLNN_API aclnnStatus aclnnScatterNdUpdateV2(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, + aclrtStream stream); + +#ifdef __cplusplus +} +#endif + +#endif // OP_API_INC_SCATTER_ND_UPDATE_V2_H_ \ No newline at end of file diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_host/op_api/scatter_nd_update_v2.cpp b/csrc/ascend/moe/scatter_nd_update_v2/op_host/op_api/scatter_nd_update_v2.cpp new file mode 100644 index 000000000..80ea91b08 --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_host/op_api/scatter_nd_update_v2.cpp @@ -0,0 +1,81 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +/*! + * \file scatter_nd_update_v2.cpp + * \brief + */ + +#include "scatter_nd_update_v2.h" +#include "opdev/make_op_executor.h" +#include "opdev/op_dfx.h" +#include "opdev/op_log.h" +#include "opdev/aicpu/aicpu_task.h" +#include "opdev/op_def.h" +#include "opdev/op_executor.h" +// #include "op_api/aclnn_util.h" +#include "aclnn_kernels/common/op_error_check.h" + +using namespace op; +namespace l0op { +OP_TYPE_REGISTER(ScatterNdUpdateV2); + +// AiCore支持的ScatterUpdate类型 +static const std::initializer_list AICORE_DTYPE_SUPPORT_LIST = { + op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BOOL}; + +static const std::initializer_list ASCEND910B_AICORE_DTYPE_SUPPORT_LIST = { + op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BOOL, op::DataType::DT_BF16, + op::DataType::DT_INT64, op::DataType::DT_INT8}; + +inline static bool IsAiCoreSupport(const aclTensor* self) { + // ScatterNdUpdateV2只需要判断self + if (GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910B || + GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910_93) { + return CheckType(self->GetDataType(), ASCEND910B_AICORE_DTYPE_SUPPORT_LIST); + } + return CheckType(self->GetDataType(), AICORE_DTYPE_SUPPORT_LIST); +} + +// AiCore的执行逻辑 +inline static const aclTensor* ScatterNdUpdateV2AiCore(const aclTensor* self, const aclTensor* indices, + const aclTensor* updates, const aclIntArray* strides, bool use_locking, + aclOpExecutor* executor) { + L0_DFX(ScatterNdUpdateV2AiCore, self, indices, updates, use_locking); + auto retAicore = + ADD_TO_LAUNCHER_LIST_AICORE(ScatterNdUpdateV2, + OP_INPUT(self, indices, updates), OP_OUTPUT(self), OP_ATTR(strides, use_locking)); + CHECK_RET(retAicore == ACLNN_SUCCESS, nullptr); + return self; +} + +// AiCPU的执行逻辑 +inline static const aclTensor* ScatterNdUpdateV2AiCPU(const aclTensor* self, const aclTensor* indices, + const aclTensor* updates, bool use_locking, + aclOpExecutor* executor) { + L0_DFX(ScatterNdUpdateV2AiCPU, self, indices, updates, use_locking); + + static internal::AicpuTaskSpace space("ScatterNdUpdateV2", ge::DEPEND_IN_SHAPE, true); + space.SetRef(0); + auto ret = ADD_TO_LAUNCHER_LIST_AICPU(ScatterNdUpdateV2, OP_ATTR_NAMES({"Tindices", "T", "use_locking"}), + OP_INPUT(self, indices, updates), OP_OUTPUT(self), + OP_ATTR(indices->GetDataType(), updates->GetDataType(), use_locking)); + CHECK_RET(ret == ACLNN_SUCCESS, nullptr); + return self; +} + +const aclTensor* ScatterNdUpdateV2(const aclTensor* self, const aclTensor* indices, const aclTensor* updates, + const aclIntArray* strides, bool use_locking, aclOpExecutor* executor) { + if (IsAiCoreSupport(self)) { + return ScatterNdUpdateV2AiCore(self, indices, updates, strides, use_locking, executor); + } else { + return ScatterNdUpdateV2AiCPU(self, indices, updates, use_locking, executor); + } +} +} // namespace l0op \ No newline at end of file diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_host/op_api/scatter_nd_update_v2.h b/csrc/ascend/moe/scatter_nd_update_v2/op_host/op_api/scatter_nd_update_v2.h new file mode 100644 index 000000000..f7543906d --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_host/op_api/scatter_nd_update_v2.h @@ -0,0 +1,24 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +/*! + * \file scatter_nd_update_v2.h + * \brief + */ +#ifndef OP_API_INC_LEVEL0_OP_SCATTER_ND_UPDATE_V2_OP_H_ +#define OP_API_INC_LEVEL0_OP_SCATTER_ND_UPDATE_V2_OP_H_ + +#include "opdev/op_executor.h" + +namespace l0op { +const aclTensor* ScatterNdUpdateV2(const aclTensor* self, const aclTensor* indices, const aclTensor* updates, + const aclIntArray* strides, bool use_locking, aclOpExecutor* executor); +} // namespace l0op + +#endif // OP_API_INC_LEVEL0_OP_SCATTER_ND_UPDATE_V2_OP_H_ diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_host/scatter_nd_update_infershape.cpp b/csrc/ascend/moe/scatter_nd_update_v2/op_host/scatter_nd_update_infershape.cpp new file mode 100644 index 000000000..51fec7c0e --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_host/scatter_nd_update_infershape.cpp @@ -0,0 +1,50 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file scatter_nd_update_v2_infershape.cpp + * \brief + */ +#include "register/op_impl_registry.h" +#include "log/log.h" +#include "util/shape_util.h" + +using namespace ge; +namespace ops { +static graphStatus InferDataType4ScatterNdUpdateV2(gert::InferDataTypeContext* context) +{ + OP_LOGD(context->GetNodeName(), "Begin to do ScatterNdUpdateV2InferDtype."); + auto var_dtype = context->GetInputDataType(0); + context->SetOutputDataType(0, var_dtype); + OP_LOGD(context->GetNodeName(), "End to do ScatterNdUpdateV2InferDtype."); + return GRAPH_SUCCESS; +} + +static ge::graphStatus InferShape4ScatterNdUpdateV2(gert::InferShapeContext* context) +{ + OP_LOGD(context->GetNodeName(), "Begin to do ScatterNdUpdateV2InferShape."); + const gert::Shape* var_in_shape = context->GetInputShape(0); + gert::Shape* var_out_shape = context->GetOutputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context, var_in_shape); + OP_CHECK_NULL_WITH_CONTEXT(context, var_out_shape); + if (Ops::Base::IsUnknownRank(*var_in_shape)) { + OP_LOGD(context->GetNodeName(), "input shape is UnknownRank, set output shape to (-2, )"); + Ops::Base::SetUnknownRank(*var_out_shape); + return ge::GRAPH_SUCCESS; + } + *var_out_shape = *var_in_shape; + OP_LOGD(context->GetNodeName(), "End to do ScatterNdUpdateV2InferShape."); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(ScatterNdUpdateV2) + .InferShape(InferShape4ScatterNdUpdateV2) + .InferDataType(InferDataType4ScatterNdUpdateV2); +} // namespace ops diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_host/scatter_nd_update_v2_def.cpp b/csrc/ascend/moe/scatter_nd_update_v2/op_host/scatter_nd_update_v2_def.cpp new file mode 100644 index 000000000..77f4fe116 --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_host/scatter_nd_update_v2_def.cpp @@ -0,0 +1,84 @@ +/** + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/*! + * \file scatter_nd_update_v2.cpp + * \brief ScatterNdUpdateV2 ophost + */ +#include "register/op_def_registry.h" + +namespace ops { +class ScatterNdUpdateV2 : public OpDef { + public: + explicit ScatterNdUpdateV2(const char* name) : OpDef(name) { + this->Input("var") + .ParamType(REQUIRED) + .DataType( + {ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BOOL, ge::DT_INT64, ge::DT_INT32, ge::DT_INT16, ge::DT_INT8, + ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BOOL, ge::DT_INT64, ge::DT_INT32, ge::DT_INT16, ge::DT_INT8}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("indices") + .ParamType(REQUIRED) + .DataType( + {ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, + ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("updates") + .ParamType(REQUIRED) + .DataType( + {ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BOOL, ge::DT_INT64, ge::DT_INT32, ge::DT_INT16, ge::DT_INT8, + ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BOOL, ge::DT_INT64, ge::DT_INT32, ge::DT_INT16, ge::DT_INT8}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Output("var") + .ParamType(REQUIRED) + .DataType( + {ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BOOL, ge::DT_INT64, ge::DT_INT32, ge::DT_INT16, ge::DT_INT8, + ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_BOOL, ge::DT_INT64, ge::DT_INT32, ge::DT_INT16, ge::DT_INT8}) + .Format( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat( + {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, + ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); + this->Attr("strides").AttrType(REQUIRED).ListInt(); + this->Attr("use_locking").AttrType(OPTIONAL).Bool(false); + OpAICoreConfig aicore_config; + aicore_config.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(true) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true); + this->AICore().AddConfig("ascend910b", aicore_config); + this->AICore().AddConfig("ascend910_93", aicore_config); + } +}; + +OP_ADD(ScatterNdUpdateV2); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_host/scatter_nd_update_v2_tiling.cpp b/csrc/ascend/moe/scatter_nd_update_v2/op_host/scatter_nd_update_v2_tiling.cpp new file mode 100644 index 000000000..1ffe1f993 --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_host/scatter_nd_update_v2_tiling.cpp @@ -0,0 +1,403 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file scatter_nd_update_v2_tiling.cpp + * \brief + */ + +#include "register/op_impl_registry.h" +#include "util/math_util.h" +#include "platform/platform_infos_def.h" +#include "log/log.h" +#include "tiling/platform/platform_ascendc.h" +#include "tiling_base/tiling_util.h" +#include "tiling_base/tiling_key.h" +#include "scatter_nd_update_v2_tiling.h" + +namespace optiling { +// using namespace Ops::NN::Optiling; +constexpr uint64_t MAX_DIM_NUM = 8; +constexpr uint64_t MAX_LENGTH_INT32 = (1LL << 31) - 1; +constexpr uint64_t MAX_FLOAT_EXPRESS_INT32 = (1LL << 24) - 1; +constexpr uint64_t SORT_USE_GM_NUM = 2; +constexpr uint64_t SORT_BLOCK_LENGTH = 4096; +constexpr uint64_t GATHER_USE_NUM = 2; +constexpr uint64_t ALIGNED_NUM = 8; +constexpr uint64_t ALIGNED_SIZE = 32; +constexpr uint64_t ATTR_STRIDE = 0; +class ScatterNdUpdateV2Tiling { +public: + explicit ScatterNdUpdateV2Tiling(gert::TilingContext* context) : tilingContext_(context){} + ge::graphStatus Init(); + ge::graphStatus SetKernelTiling(); + void TilingDataPrint() const; + +private: + inline bool IsSort(uint64_t totalLength, uint64_t indexRow); + inline bool IsLinearIndex(uint64_t totalLength); + inline size_t CalcWorkSpaceSize(uint64_t indexRow); + inline void SetTilingKeyMode(); + inline void GetDtypeSize(); + inline void Tiling4Scatter(uint64_t totalLength, uint64_t indexRow); + inline void Tiling4LinearIndex(uint64_t indexRow, uint64_t indexDim); + + ScatterNdUpdateV2TilingData tilingData_; + gert::TilingContext* tilingContext_ = nullptr; + + uint64_t coreNum_ = 0; + uint64_t tilingKey_ = 0; + uint64_t ubSize_ = 0; + uint64_t isLinearIndex_ = false; + uint64_t isSort_ = false; + uint64_t sortWorkspace_ = 0; + uint64_t dataTypeSize_ = 0; + uint64_t isInt64Indices_ = false; + uint64_t needLargeIndexKernel_ = false; + +private: + // LinearIndex + uint64_t indexDim_ = 0; + uint64_t blockLength_ = 0; + uint64_t blockNum_ = 0; + uint64_t blockRemainLength_ = 0; + uint64_t tailBlockNum_ = 0; + uint64_t frontBlockNum_ = 0; + uint64_t frontCoreNum_ = 0; + uint64_t tailCoreNum_ = 0; + uint64_t indicesMask_[MAX_DIM_NUM] = {0}; + + // Scatter + uint64_t scatterLength_ = 1; + uint64_t tailRow_ = 0; + uint64_t frontRow_ = 0; + uint64_t frontNum_ = 0; + uint64_t tailNum_ = 0; + uint64_t ubLengthForUpdates_ = 0; + uint64_t scatterAlignLength_ = 0; + uint64_t formDim_ = 0; + uint64_t copyRow_ = 0; + uint64_t scatterTileNum_ = 1; + uint64_t scatterTileLength_ = 0; + uint64_t scatterTileTail_ = 0; + uint64_t scatterTileAlignLength_ = 0; +}; + +inline void ScatterNdUpdateV2Tiling::SetTilingKeyMode() +{ + // tilingKey: indexType * 10 + sortFlag (indexType: 1=int32, 2=int64(cast), 3=int64(large)) + uint64_t indexType; + if (!isInt64Indices_) { + indexType = 1; + } else if (needLargeIndexKernel_) { + indexType = 3; + } else { + indexType = 2; + } + uint64_t sortFlag = (indexType == 3) ? 0 : (isSort_ ? 1 : 0); + tilingKey_ = indexType * 10 + sortFlag; + + tilingContext_->SetTilingKey(tilingKey_); + OP_LOGD(tilingContext_, "isLinearIndex=%lu, isSort=%lu, isInt64Indices=%lu, needLargeIndexKernel=%lu, tilingKey=%lu (indexType=%lu, sortFlag=%lu)", + isLinearIndex_, isSort_, isInt64Indices_, needLargeIndexKernel_, tilingKey_, indexType, sortFlag); +} + +inline bool ScatterNdUpdateV2Tiling::IsLinearIndex(uint64_t totalLength) +{ + return totalLength <= MAX_LENGTH_INT32; +} + +inline bool ScatterNdUpdateV2Tiling::IsSort(uint64_t totalLength, uint64_t indexRow) +{ + return totalLength <= MAX_FLOAT_EXPRESS_INT32; +} + +inline void ScatterNdUpdateV2Tiling::Tiling4LinearIndex(uint64_t indexRow, uint64_t indexDim) +{ + OP_LOGD(tilingContext_, "linearIndexTiling start"); + auto attrs = tilingContext_->GetAttrs(); + auto stridesPtr = attrs->GetListInt(ATTR_STRIDE); + for (uint64_t i = 0; i < indexDim; ++i) { + indicesMask_[i] = static_cast(stridesPtr->GetData()[i]); + } + uint64_t coeff = isInt64Indices_ ? (2 * indexDim + 3) : (indexDim + 3); + uint64_t maxBlockLength = ubSize_ / coeff / sizeof(int); + blockLength_ = (maxBlockLength / ALIGNED_SIZE) * ALIGNED_SIZE; + blockLength_ = std::min(blockLength_, (uint64_t)SORT_BLOCK_LENGTH); + blockNum_ = indexRow / blockLength_; + blockRemainLength_ = indexRow % blockLength_; + + if (blockNum_ == 0) { + tailBlockNum_ = 0; + frontBlockNum_ = 0; + frontCoreNum_ = 1; + tailCoreNum_ = 0; + } else { + tailBlockNum_ = blockNum_ / coreNum_; + frontBlockNum_ = tailBlockNum_ + 1; + frontCoreNum_ = blockNum_ % coreNum_; + tailCoreNum_ = tailBlockNum_ == 0 ? 0 : coreNum_ - frontCoreNum_; + } + OP_LOGD(tilingContext_, "linearIndexTiling finish"); +} + +inline void ScatterNdUpdateV2Tiling::Tiling4Scatter(uint64_t totalLength, uint64_t indexRow) +{ + OP_LOGD(tilingContext_, "scatterTiling start new"); + uint64_t scatterAlignNum = ALIGNED_SIZE / dataTypeSize_; + tailRow_ = totalLength / coreNum_; + frontRow_ = tailRow_ + 1; + frontNum_ = totalLength % coreNum_; + tailNum_ = tailRow_ == 0 ? 0 : coreNum_ - frontNum_; + ubLengthForUpdates_ = ((ubSize_ - SORT_BLOCK_LENGTH * SORT_USE_GM_NUM * sizeof(int)) / ALIGNED_SIZE * ALIGNED_SIZE) / dataTypeSize_; + scatterAlignLength_ = (scatterLength_ + scatterAlignNum - 1) & ~(scatterAlignNum - 1); + formDim_ = scatterAlignLength_ / ubLengthForUpdates_; + + scatterTileLength_ = std::min(scatterLength_, ubLengthForUpdates_); + if (scatterTileLength_ == 0) { + scatterTileLength_ = 1; + } + scatterTileNum_ = (scatterLength_ + scatterTileLength_ - 1) / scatterTileLength_; + scatterTileTail_ = scatterLength_ - (scatterTileNum_ - 1) * scatterTileLength_; + scatterTileAlignLength_ = (scatterTileLength_ + scatterAlignNum - 1) & ~(scatterAlignNum - 1); + + if (scatterTileNum_ > 1) { + copyRow_ = 1; + } else { + copyRow_ = formDim_ == 0 ? ubLengthForUpdates_ / scatterAlignLength_ : 1; + } + OP_LOGD(tilingContext_, "scatterTiling finish"); +} + +inline void ScatterNdUpdateV2Tiling::GetDtypeSize() +{ + uint64_t varDtype = tilingContext_->GetInputDesc(0)->GetDataType(); + switch (varDtype){ + case ge::DT_FLOAT: + dataTypeSize_ = 4; + break; + case ge::DT_BF16: + dataTypeSize_ = 2; + break; + case ge::DT_FLOAT16: + dataTypeSize_ = 2; + break; + case ge::DT_BOOL: + dataTypeSize_ = 1; + break; + case ge::DT_INT64: + dataTypeSize_ = 8; + break; + case ge::DT_INT32: + dataTypeSize_ = 4; + break; + case ge::DT_INT16: + dataTypeSize_ = 2; + break; + case ge::DT_INT8: + dataTypeSize_ = 1; + break; + default: + break; + } +} + + +ge::graphStatus ScatterNdUpdateV2Tiling::SetKernelTiling() +{ + tilingContext_->SetBlockDim(coreNum_); + tilingData_.linearIndexTiling.set_indexDim(indexDim_); + tilingData_.linearIndexTiling.set_ubSize(ubSize_); + tilingData_.linearIndexTiling.set_indicesMask(indicesMask_); + tilingData_.linearIndexTiling.set_coreNum(coreNum_); + tilingData_.linearIndexTiling.set_blockLength(blockLength_); + tilingData_.linearIndexTiling.set_blockNum(blockNum_); + tilingData_.linearIndexTiling.set_blockRemainLength(blockRemainLength_); + tilingData_.linearIndexTiling.set_tailBlockNum(tailBlockNum_); + tilingData_.linearIndexTiling.set_frontBlockNum(frontBlockNum_); + tilingData_.linearIndexTiling.set_frontCoreNum(frontCoreNum_); + tilingData_.linearIndexTiling.set_tailCoreNum(tailCoreNum_); + tilingData_.linearIndexTiling.set_sortWorkspace(sortWorkspace_); + tilingData_.linearIndexTiling.set_isInt64Indices(isInt64Indices_); + tilingData_.linearIndexTiling.set_needLargeIndexKernel(needLargeIndexKernel_); + tilingData_.scatterTiling.set_scatterLength(scatterLength_); + tilingData_.scatterTiling.set_tailRow(tailRow_); + tilingData_.scatterTiling.set_frontRow(frontRow_); + tilingData_.scatterTiling.set_frontNum(frontNum_); + tilingData_.scatterTiling.set_tailNum(tailNum_); + tilingData_.scatterTiling.set_ubLengthForUpdates(ubLengthForUpdates_); + tilingData_.scatterTiling.set_scatterAlignLength(scatterAlignLength_); + tilingData_.scatterTiling.set_formDim(formDim_); + tilingData_.scatterTiling.set_copyRow(copyRow_); + tilingData_.scatterTiling.set_scatterTileNum(scatterTileNum_); + tilingData_.scatterTiling.set_scatterTileLength(scatterTileLength_); + tilingData_.scatterTiling.set_scatterTileTail(scatterTileTail_); + tilingData_.scatterTiling.set_scatterTileAlignLength(scatterTileAlignLength_); + tilingData_.SaveToBuffer( + tilingContext_->GetRawTilingData()->GetData(), tilingContext_->GetRawTilingData()->GetCapacity()); + tilingContext_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize()); + TilingDataPrint(); + return ge::GRAPH_SUCCESS; +} + +inline size_t ScatterNdUpdateV2Tiling::CalcWorkSpaceSize(uint64_t indexRow) +{ + auto ascendcPlatform = platform_ascendc::PlatformAscendC(tilingContext_->GetPlatformInfo()); + size_t sysWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize(); + size_t indexRowAligned = (indexRow + ALIGNED_NUM - 1) & ~(ALIGNED_NUM - 1); + sortWorkspace_ = indexRowAligned; + size_t totalWorkspace = sysWorkspaceSize; + if (isLinearIndex_) { + totalWorkspace += sortWorkspace_ * SORT_USE_GM_NUM * sizeof(int); + } + if (isSort_) { + totalWorkspace += sortWorkspace_ * SORT_USE_GM_NUM * sizeof(int); + } + return totalWorkspace; +} + +void ScatterNdUpdateV2Tiling::TilingDataPrint() const +{ + OP_LOGD(tilingContext_, "coreNum: %lu", coreNum_); + OP_LOGD(tilingContext_, "tilingKey: %lu", tilingKey_); + OP_LOGD(tilingContext_, "isInt64Indices: %lu", isInt64Indices_); + OP_LOGD(tilingContext_, "needLargeIndexKernel: %lu", needLargeIndexKernel_); + OP_LOGD(tilingContext_, "tiling for LinearIndex--------"); + OP_LOGD(tilingContext_, "indexDim: %lu", indexDim_); + OP_LOGD(tilingContext_, "ubSize: %lu", ubSize_); + OP_LOGD(tilingContext_, "blockLength: %lu", blockLength_); + OP_LOGD(tilingContext_, "blockNum: %lu", blockNum_); + OP_LOGD(tilingContext_, "blockRemainLength: %lu", blockRemainLength_); + OP_LOGD(tilingContext_, "tailBlockNum: %lu", tailBlockNum_); + OP_LOGD(tilingContext_, "frontBlockNum: %lu", frontBlockNum_); + OP_LOGD(tilingContext_, "frontCoreNum: %lu", frontCoreNum_); + OP_LOGD(tilingContext_, "tailCoreNum: %lu", tailCoreNum_); + OP_LOGD(tilingContext_, "sortWorkspace: %lu", sortWorkspace_); + for (size_t i = 0; i < indexDim_; i++) { + OP_LOGD(tilingContext_, "indicesMask[%lu]: %lu", i, indicesMask_[i]); + } + OP_LOGD(tilingContext_, "tiling for Scatter------------"); + OP_LOGD(tilingContext_, "scatterLength: %lu", scatterLength_); + OP_LOGD(tilingContext_, "tailRow: %lu", tailRow_); + OP_LOGD(tilingContext_, "frontRow: %lu", frontRow_); + OP_LOGD(tilingContext_, "frontNum: %lu", frontNum_); + OP_LOGD(tilingContext_, "tailNum: %lu", tailNum_); + OP_LOGD(tilingContext_, "ubLengthForUpdates: %lu", ubLengthForUpdates_); + OP_LOGD(tilingContext_, "scatterAlignLength: %lu", scatterAlignLength_); + OP_LOGD(tilingContext_, "formDim: %lu", formDim_); + OP_LOGD(tilingContext_, "copyRow: %lu", copyRow_); + OP_LOGD(tilingContext_, "scatterTileNum: %lu", scatterTileNum_); + OP_LOGD(tilingContext_, "scatterTileLength: %lu", scatterTileLength_); + OP_LOGD(tilingContext_, "scatterTileTail: %lu", scatterTileTail_); + OP_LOGD(tilingContext_, "scatterTileAlignLength: %lu", scatterTileAlignLength_); +} + +ge::graphStatus ScatterNdUpdateV2Tiling::Init() +{ + OP_LOGD(tilingContext_, "Tiling initing"); + auto compileInfo = static_cast(tilingContext_->GetCompileInfo()); + auto varRefShape = tilingContext_->GetInputShape(0)->GetStorageShape(); + auto indicesShape = tilingContext_->GetInputShape(1)->GetStorageShape(); + auto updatesShape = tilingContext_->GetInputShape(2)->GetStorageShape(); + uint64_t varDimNum = varRefShape.GetDimNum(); + indexDim_ = indicesShape.GetDim(indicesShape.GetDimNum() - 1); + + auto indicesDtype = tilingContext_->GetInputDesc(1)->GetDataType(); + isInt64Indices_ = (indicesDtype == ge::DT_INT64); + OP_LOGD(tilingContext_, "indicesDtype=%d, isInt64Indices=%lu", indicesDtype, isInt64Indices_); + + uint64_t totalLength = 1; + for (uint64_t i = 0; i < indexDim_; ++i) { + totalLength *= varRefShape.GetDim(i); + } + + if (isInt64Indices_) { + needLargeIndexKernel_ = !IsLinearIndex(totalLength); + } + + if (varDimNum > indexDim_) { + for (uint64_t i = indexDim_; i < varDimNum; i++) { + scatterLength_ *= varRefShape.GetDim(i); + } + } + uint64_t indexRow = 1; + for (uint64_t i = 0; i < indicesShape.GetDimNum() - 1; i++) { + indexRow *= indicesShape.GetDim(i); + } + + if (needLargeIndexKernel_) { + isSort_ = false; + isLinearIndex_ = false; + } else { + isSort_ = false; + isLinearIndex_ = IsLinearIndex(totalLength); + } + coreNum_ = std::min(compileInfo->totalCoreNum, + std::min(static_cast(totalLength), static_cast(indexRow))); + coreNum_ = coreNum_ == 0 ? 1 : coreNum_; + ubSize_ = compileInfo->ubSizePlatForm; + GetDtypeSize(); + Tiling4LinearIndex(indexRow, indexDim_); + uint64_t maxPhysicalOffset = 0; + for (uint64_t i = 0; i < indexDim_; ++i) { + maxPhysicalOffset += (varRefShape.GetDim(i) - 1) * indicesMask_[i]; + } + uint64_t totalPhysicalRange = maxPhysicalOffset + scatterLength_; + if (!needLargeIndexKernel_) { + isSort_ = IsSort(totalPhysicalRange, indexRow); + } + SetTilingKeyMode(); + tilingContext_->SetScheduleMode(1); + Tiling4Scatter(totalPhysicalRange, indexRow); + size_t* currentWorkSpace = tilingContext_->GetWorkspaceSizes(1); + currentWorkSpace[0] = CalcWorkSpaceSize(indexRow); + OP_LOGD(tilingContext_, "Tiling inited"); + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus Tiling4ScatterNdUpdateV2(gert::TilingContext* context) +{ + if (context == nullptr) { + OP_LOGE("ScatterNdUpdateV2", "The context is nullptr."); + return ge::GRAPH_FAILED; + } + OP_LOGD(context, "Tiling for ScatterNdUpdateV2 start."); + ScatterNdUpdateV2Tiling tilingOp(context); + if (tilingOp.Init() != ge::GRAPH_SUCCESS) { + OP_LOGE(context, "Tiling init fail"); + return ge::GRAPH_FAILED; + } + OP_LOGD(context, "Tiling for ScatterNdUpdateV2 end."); + return tilingOp.SetKernelTiling(); +} + +ge::graphStatus TilingPrepare4ScatterNdUpdateV2(gert::TilingParseContext* context) +{ + OP_LOGD(context, "Tiling Prepare For ScatterNdUpdateV2 start."); + auto compileInfo = context->GetCompiledInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo); + auto platformInfo = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); + compileInfo->totalCoreNum = ascendcPlatform.GetCoreNumAiv(); + if (compileInfo->totalCoreNum == 0) { + OP_LOGE(context, "coreNum %lu", compileInfo->totalCoreNum); + return ge::GRAPH_FAILED; + } + uint64_t ubSizePlatForm; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm); + compileInfo->ubSizePlatForm = ubSizePlatForm; + OP_LOGD(context, "ubSizePlatForm is %lu.", compileInfo->ubSizePlatForm); + OP_LOGD(context, "Tiling Prepare For ScatterNdUpdateV2 end."); + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(ScatterNdUpdateV2).Tiling(Tiling4ScatterNdUpdateV2).TilingParse(TilingPrepare4ScatterNdUpdateV2); +} diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_host/scatter_nd_update_v2_tiling.h b/csrc/ascend/moe/scatter_nd_update_v2/op_host/scatter_nd_update_v2_tiling.h new file mode 100644 index 000000000..200d2bced --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_host/scatter_nd_update_v2_tiling.h @@ -0,0 +1,71 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file scatter_nd_update_v2_tiling.h + * \brief + */ + +#ifndef SCATTER_ND_UPDATE_V2_TILING_H +#define SCATTER_ND_UPDATE_V2_TILING_H +#include "register/tilingdata_base.h" +constexpr uint64_t MAX_DIM_NUM = 8; +namespace optiling { +BEGIN_TILING_DATA_DEF(ScatterNdUpdateV2ScatterTiling) +TILING_DATA_FIELD_DEF(uint64_t, scatterLength) +TILING_DATA_FIELD_DEF(uint64_t, tailRow) +TILING_DATA_FIELD_DEF(uint64_t, frontRow) +TILING_DATA_FIELD_DEF(uint64_t, frontNum) +TILING_DATA_FIELD_DEF(uint64_t, tailNum) +TILING_DATA_FIELD_DEF(uint64_t, ubLengthForUpdates) +TILING_DATA_FIELD_DEF(uint64_t, scatterAlignLength) +TILING_DATA_FIELD_DEF(uint64_t, formDim) +TILING_DATA_FIELD_DEF(uint64_t, copyRow) +TILING_DATA_FIELD_DEF(uint64_t, scatterTileNum) +TILING_DATA_FIELD_DEF(uint64_t, scatterTileLength) +TILING_DATA_FIELD_DEF(uint64_t, scatterTileTail) +TILING_DATA_FIELD_DEF(uint64_t, scatterTileAlignLength) +END_TILING_DATA_DEF + +REGISTER_TILING_DATA_CLASS(ScatterNdUpdateV2ScatterTilingOp, ScatterNdUpdateV2ScatterTiling) + +BEGIN_TILING_DATA_DEF(ScatterNdUpdateV2LinearIndexTiling) +TILING_DATA_FIELD_DEF(uint64_t, coreNum) +TILING_DATA_FIELD_DEF(uint64_t, ubSize) +TILING_DATA_FIELD_DEF(uint64_t, indexDim) +TILING_DATA_FIELD_DEF(uint64_t, blockLength) +TILING_DATA_FIELD_DEF(uint64_t, blockNum) +TILING_DATA_FIELD_DEF(uint64_t, blockRemainLength) +TILING_DATA_FIELD_DEF(uint64_t, tailBlockNum) +TILING_DATA_FIELD_DEF(uint64_t, frontBlockNum) +TILING_DATA_FIELD_DEF(uint64_t, frontCoreNum) +TILING_DATA_FIELD_DEF(uint64_t, tailCoreNum) +TILING_DATA_FIELD_DEF(uint64_t, sortWorkspace) +TILING_DATA_FIELD_DEF_ARR(uint64_t, MAX_DIM_NUM, indicesMask) +TILING_DATA_FIELD_DEF(uint64_t, isInt64Indices) +TILING_DATA_FIELD_DEF(uint64_t, needLargeIndexKernel) +END_TILING_DATA_DEF + +REGISTER_TILING_DATA_CLASS(ScatterNdUpdateV2LinearIndexTilingOp, ScatterNdUpdateV2LinearIndexTiling) + +BEGIN_TILING_DATA_DEF(ScatterNdUpdateV2TilingData) +TILING_DATA_FIELD_DEF_STRUCT(ScatterNdUpdateV2ScatterTiling, scatterTiling) +TILING_DATA_FIELD_DEF_STRUCT(ScatterNdUpdateV2LinearIndexTiling, linearIndexTiling) +END_TILING_DATA_DEF + +REGISTER_TILING_DATA_CLASS(ScatterNdUpdateV2, ScatterNdUpdateV2TilingData) +REGISTER_TILING_DATA_CLASS(ScatterNdUpdateV2TilingDataOp, ScatterNdUpdateV2TilingData) + +struct ScatterNdUpdateV2CompileInfo { + uint64_t totalCoreNum = 0; + uint64_t ubSizePlatForm = 0; +}; +} // namespace optiling +#endif // SCATTER_ND_UPDATE_V2_TILING_H diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_common.h b/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_common.h new file mode 100644 index 000000000..14a34c10c --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_common.h @@ -0,0 +1,70 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file scatter_nd_update_common.h + * \brief ScatterNdUpdateV2 公共定义和工具函数 + */ + +#ifndef SCATTER_ND_UPDATE_V2_COMMON_H +#define SCATTER_ND_UPDATE_V2_COMMON_H + +#include "kernel_operator.h" + +namespace ScatterNdUpdateV2 { +using namespace AscendC; + +// 公共常量定义 +constexpr uint64_t DOUBLE_BUFFER = 1; +constexpr uint64_t SORT_RES_NUM = 2; +constexpr uint64_t SORT_TMP_NUM = 3; +constexpr uint64_t ALIGNED_BLOCK_NUM = 32; +constexpr uint64_t ALIGN_NUM = 8; // 32 字节对齐 = 8 个 int32 +constexpr uint64_t ALIGNED_SIZE = 512; + +// 公共同步函数 +__aicore__ inline void PipeMte2ToS() +{ + event_t eventID = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S)); + SetFlag(eventID); + WaitFlag(eventID); +} + +__aicore__ inline void PipeMte3ToS() +{ + event_t eventID = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S)); + SetFlag(eventID); + WaitFlag(eventID); +} + +__aicore__ inline void PipeVToMte3() +{ + event_t eventID = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + SetFlag(eventID); + WaitFlag(eventID); +} + +// 计算 block 分布参数 +__aicore__ inline void CalcBlockDistribution( + uint64_t blockIdx, uint64_t frontNum, uint64_t frontRow, uint64_t tailRow, + uint64_t& computeRow, uint64_t& start) +{ + if (blockIdx >= frontNum) { + computeRow = tailRow; + start = frontNum * frontRow + (blockIdx - frontNum) * computeRow; + } else { + computeRow = frontRow; + start = blockIdx * computeRow; + } +} + +} // namespace ScatterNdUpdateV2 + +#endif // SCATTER_ND_UPDATE_V2_COMMON_H diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_large_index.h b/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_large_index.h new file mode 100644 index 000000000..efa0c967b --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_large_index.h @@ -0,0 +1,175 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file scatter_nd_update_large_index.h + * \brief LargeIndex Kernel (index > 2^31-1) + */ + +#include "kernel_operator.h" +#include "kernel_tiling/kernel_tiling.h" +#include "scatter_nd_update_common.h" + +namespace ScatterNdUpdateV2 { + +template +class LargeIndexKernel { +public: + __aicore__ inline LargeIndexKernel() = delete; + __aicore__ inline LargeIndexKernel( + GM_ADDR indices, GM_ADDR updates, GM_ADDR output, + const ScatterNdUpdateV2TilingData& tiling, TPipe& pipe) + { + InitParams(tiling); + InitBuffers(pipe); + SetGmAddr(indices, updates, output, tiling); + } + + __aicore__ inline void InitParams(const ScatterNdUpdateV2TilingData& tiling) + { + blockIdx_ = GetBlockIdx(); + + CalcBlockDistribution(blockIdx_, tiling.scatterTiling.frontNum, tiling.scatterTiling.frontRow, + tiling.scatterTiling.tailRow, computeRow_, start_); + end_ = start_ + computeRow_; + + startInt64_ = static_cast(start_); + endInt64_ = static_cast(end_); + + indexDim_ = tiling.linearIndexTiling.indexDim; + blockLength_ = tiling.linearIndexTiling.blockLength; + blockNum_ = tiling.linearIndexTiling.blockNum; + blockRemainLength_ = tiling.linearIndexTiling.blockRemainLength; + + scatterLength_ = tiling.scatterTiling.scatterLength; + ubLengthForUpdates_ = tiling.scatterTiling.ubLengthForUpdates; + scatterTileNum_ = tiling.scatterTiling.scatterTileNum; + scatterTileLength_ = tiling.scatterTiling.scatterTileLength; + scatterTileTail_ = tiling.scatterTiling.scatterTileTail; + + for (uint64_t i = 0; i < indexDim_; ++i) { + indicesMask_[i] = tiling.linearIndexTiling.indicesMask[i]; + } + } + + __aicore__ inline void InitBuffers(TPipe& pipe) + { + uint64_t indicesInt64Size = ((blockLength_ * indexDim_ * 2) + ALIGN_NUM - 1) & ~(ALIGN_NUM - 1); + uint64_t updateBufBytes = (ubLengthForUpdates_ * sizeof(T) + 31) & ~31ULL; + + pipe.InitBuffer(indicesBuf, indicesInt64Size * sizeof(int)); + pipe.InitBuffer(updateBuf, updateBufBytes); + + indicesInt64Local = indicesBuf.Get().ReinterpretCast(); + updateLocal = updateBuf.Get(); + } + + __aicore__ inline void SetGmAddr(GM_ADDR indices, GM_ADDR updates, GM_ADDR output, + const ScatterNdUpdateV2TilingData& tiling) + { + indicesGmInt64_.SetGlobalBuffer((__gm__ int64_t*)indices); + updatesGm_.SetGlobalBuffer((__gm__ T*)updates); + outputGm_.SetGlobalBuffer((__gm__ T*)output); + } + + __aicore__ inline void Process() + { + for (uint64_t blockIdx = 0; blockIdx < blockNum_; ++blockIdx) { + ProcessOneBlock(blockIdx, false); + } + if (blockRemainLength_ != 0) { + ProcessOneBlock(blockNum_, true); + } + } + + __aicore__ inline void ProcessOneBlock(uint64_t blockIdx, bool isTail) + { + uint64_t copyRow = isTail ? blockRemainLength_ : blockLength_; + CopyInInt64(blockIdx, isTail); + + for (uint64_t i = 0; i < copyRow; ++i) { + int64_t linearIndex = ComputeLinearIndex(i); + if (linearIndex >= startInt64_ && linearIndex < endInt64_) { + ScatterUpdate(i, linearIndex); + } + } + } + + __aicore__ inline void CopyInInt64(uint64_t blockIdx, bool isTail) + { + uint64_t indicesOffset = blockIdx * blockLength_ * indexDim_; + uint64_t copyRow = isTail ? blockRemainLength_ : blockLength_; + + DataCopyExtParams copyParams{1, static_cast(copyRow * indexDim_ * sizeof(int64_t)), 0, 0, 0}; + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(indicesInt64Local, indicesGmInt64_[indicesOffset], copyParams, padParams); + PipeMte2ToS(); + } + + __aicore__ inline int64_t ComputeLinearIndex(uint64_t rowIdx) + { + int64_t linearIndex = 0; + for (uint64_t dim = 0; dim < indexDim_; ++dim) { + int64_t idxValue = indicesInt64Local.GetValue(rowIdx * indexDim_ + dim); + int64_t stride = static_cast(indicesMask_[dim]); + linearIndex += idxValue * stride; + } + return linearIndex; + } + + __aicore__ inline void ScatterUpdate(uint64_t rowIdx, int64_t linearIndex) + { + for (uint64_t tileIdx = 0; tileIdx < scatterTileNum_; ++tileIdx) { + uint64_t tileLength = (tileIdx == scatterTileNum_ - 1) ? scatterTileTail_ : scatterTileLength_; + uint64_t gmOffset = rowIdx * scatterLength_ + tileIdx * scatterTileLength_; + DataCopyExtParams updateCopyParams{1, static_cast(tileLength * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(updateLocal, updatesGm_[gmOffset], updateCopyParams, padParams); + PipeMte2ToS(); + + uint64_t outOffset = static_cast(linearIndex) + tileIdx * scatterTileLength_; + DataCopyExtParams outParams{1, static_cast(tileLength * sizeof(T)), 0, 0, 0}; + DataCopyPad(outputGm_[outOffset], updateLocal, outParams); + PipeMte3ToS(); + } + } + +private: + GlobalTensor indicesGmInt64_; + GlobalTensor updatesGm_; + GlobalTensor outputGm_; + + TBuf indicesBuf; + TBuf updateBuf; + + LocalTensor indicesInt64Local; + LocalTensor updateLocal; + + uint64_t blockIdx_; + uint64_t computeRow_; + uint64_t start_; + uint64_t end_; + int64_t startInt64_; + int64_t endInt64_; + + uint64_t indexDim_; + uint64_t blockLength_; + uint64_t blockNum_; + uint64_t blockRemainLength_; + uint64_t indicesMask_[8]; + + uint64_t scatterLength_; + uint64_t ubLengthForUpdates_; + uint64_t scatterTileNum_; + uint64_t scatterTileLength_; + uint64_t scatterTileTail_; +}; + +} // namespace ScatterNdUpdateV2 diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_linear_index.h b/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_linear_index.h new file mode 100644 index 000000000..08224aa45 --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_linear_index.h @@ -0,0 +1,300 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file scatter_nd_update_linear_index.h + * \brief LinearIndex Kernel + */ + +#include "kernel_operator.h" +#include "kernel_tiling/kernel_tiling.h" +#include "scatter_nd_update_common.h" + +namespace ScatterNdUpdateV2 { + +template +class LinearIndexKernel { +public: + __aicore__ inline LinearIndexKernel() = delete; + __aicore__ inline LinearIndexKernel( + GM_ADDR indices, GM_ADDR workSpace, const ScatterNdUpdateV2TilingData& tiling, TPipe& pipe) + { + InitParams(tiling); + InitBuffers(pipe); + SetGmAddr(indices, workSpace, tiling); + } + + __aicore__ inline void InitParams(const ScatterNdUpdateV2TilingData& tiling) + { + blockIdx_ = GetBlockIdx(); + frontCoreNum_ = tiling.linearIndexTiling.frontCoreNum; + tailCoreNum_ = tiling.linearIndexTiling.tailCoreNum; + frontBlockNum_ = tiling.linearIndexTiling.frontBlockNum; + tailBlockNum_ = tiling.linearIndexTiling.tailBlockNum; + if (blockIdx_ >= frontCoreNum_) { + computeNum_ = tailBlockNum_; + } else { + computeNum_ = frontBlockNum_; + } + ubSize_ = tiling.linearIndexTiling.ubSize; + coreNum_ = tiling.linearIndexTiling.coreNum; + blockNum_ = tiling.linearIndexTiling.blockNum; + blockLength_ = tiling.linearIndexTiling.blockLength; + blockRemainLength_ = tiling.linearIndexTiling.blockRemainLength; + indexDim_ = tiling.linearIndexTiling.indexDim; + indicesMask_ = tiling.linearIndexTiling.indicesMask; + } + + template + __aicore__ inline void InitBuffersUnified() + { + uint64_t offset = 0; + indicesLocal = allUbLocal[offset]; + offset += blockLength_; + uint64_t indicesOffset = offset; + if constexpr (isInt64) { + indicesInt64Local = allUbLocal[offset].ReinterpretCast(); + indicesOriginLocal = allUbLocal[offset]; + offset += blockLength_ * indexDim_ * 2; + } else { + indicesOriginLocal = allUbLocal[offset]; + offset += blockLength_ * indexDim_; + } + addTmpLocal = allUbLocal[offset]; + offset += blockLength_; + rangeLocal = allUbLocal[offset]; + offset += blockLength_; + if constexpr (isSort) { + resLocal = allUbLocal[indicesOffset].ReinterpretCast(); + indicesOffset += blockLength_ * 2; + posIdxLocal = allUbLocal[indicesOffset]; + indicesOffset += blockLength_; + sortTmpLocal = allUbLocal[indicesOffset].ReinterpretCast(); + } + } + + __aicore__ inline void InitBuffers(TPipe& pipe) + { + pipe.InitBuffer(allUbBuf, ubSize_); + allUbLocal = allUbBuf.Get(); + if constexpr (isSort) { + if constexpr (std::is_same_v) { + InitBuffersUnified(); + } else { + InitBuffersUnified(); + } + } else { + if constexpr (std::is_same_v) { + InitBuffersUnified(); + } else { + InitBuffersUnified(); + } + } + } + + __aicore__ inline void SetGmAddr(GM_ADDR indices, GM_ADDR workSpace, const ScatterNdUpdateV2TilingData& tiling) + { + indiceAddrOffset_ = + blockIdx_ < tiling.linearIndexTiling.frontCoreNum ? + tiling.linearIndexTiling.frontBlockNum * blockLength_ * blockIdx_ : + tiling.linearIndexTiling.frontCoreNum * tiling.linearIndexTiling.frontBlockNum * blockLength_ + + (blockIdx_ - tiling.linearIndexTiling.frontCoreNum) * tiling.linearIndexTiling.tailBlockNum * + blockLength_; + + sortedIndicesGm_.SetGlobalBuffer((__gm__ int*)workSpace + indiceAddrOffset_); + if constexpr (isSort) { + posIndicesGm_.SetGlobalBuffer((__gm__ int*)workSpace + tiling.linearIndexTiling.sortWorkspace + indiceAddrOffset_); + } + if constexpr (std::is_same_v) { + indicesGmInt64_.SetGlobalBuffer((__gm__ int64_t*)indices + indiceAddrOffset_ * indexDim_); + } else { + indicesGm_.SetGlobalBuffer((__gm__ int*)indices + indiceAddrOffset_ * indexDim_); + } + } + + __aicore__ inline void Process() + { + if constexpr (isSort) { + for (uint64_t i = 0; i < computeNum_; i++) { + ProcessOneWithSort(i, false); + } + uint64_t lastActiveCore = (blockNum_ == 0) ? 0 : + (tailCoreNum_ == 0 ? frontCoreNum_ - 1 : frontCoreNum_ + tailCoreNum_ - 1); + if (blockIdx_ == lastActiveCore && blockRemainLength_ != 0) { + ProcessOneWithSort(computeNum_, true); + } + } else { + for (uint64_t i = 0; i < computeNum_; i++) { + ProcessOne(i, false); + } + uint64_t lastActiveCore = (blockNum_ == 0) ? 0 : + (tailCoreNum_ == 0 ? frontCoreNum_ - 1 : frontCoreNum_ + tailCoreNum_ - 1); + if (blockIdx_ == lastActiveCore && blockRemainLength_ != 0) { + ProcessOne(computeNum_, true); + } + } + } + + __aicore__ inline void ProcessOne(uint64_t idx, bool isTail) + { + CopyIn(idx, isTail); + if constexpr (std::is_same_v) { + CastToInt32(idx, isTail); + } + Compute4LinearIndex(idx, isTail); + CopyOut(idx, isTail); + } + + __aicore__ inline void ProcessOneWithSort(uint64_t idx, bool isTail) + { + CopyIn(idx, isTail); + if constexpr (std::is_same_v) { + CastToInt32(idx, isTail); + } + Compute4LinearIndex(idx, isTail); + ComputeForSort(idx, isTail); + CopyOut(idx, isTail); + } + + __aicore__ inline void CopyIn(uint64_t process, bool isTail) + { + uint64_t indicesOffset = process * blockLength_ * indexDim_; + uint64_t copyRow = isTail ? blockRemainLength_ : blockLength_; + + if constexpr (std::is_same_v) { + DataCopyExtParams copyParams{1, static_cast(copyRow * indexDim_ * sizeof(int64_t)), 0, 0, 0}; + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(indicesInt64Local, indicesGmInt64_[indicesOffset], copyParams, padParams); + } else { + DataCopyExtParams copyParams{1, static_cast(copyRow * indexDim_ * sizeof(int)), 0, 0, 0}; + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(indicesOriginLocal, indicesGm_[indicesOffset], copyParams, padParams); + } + PipeMte2ToS(); + } + + __aicore__ inline void CastToInt32(uint64_t process, bool isTail) + { + uint64_t computeRow = isTail ? blockRemainLength_ : blockLength_; + uint64_t totalElements = computeRow * indexDim_; + Cast(indicesOriginLocal, indicesInt64Local, RoundMode::CAST_NONE, totalElements); + PipeBarrier(); + } + + __aicore__ inline void Compute4LinearIndex(uint64_t process, bool isTail) + { + uint64_t computeRow = isTail ? blockRemainLength_ : blockLength_; + int32_t malValue = indexDim_ * sizeof(int); + Duplicate(indicesLocal, 0, computeRow); + CreateVecIndex(rangeLocal, (int)0, computeRow); + PipeBarrier(); + Muls(rangeLocal, rangeLocal, malValue, computeRow); + PipeBarrier(); + for (int i = 0; i < indexDim_; ++i) { + if (i != 0) { + Adds(rangeLocal, rangeLocal, (int)(sizeof(int)), computeRow); + PipeBarrier(); + } + LocalTensor rangeLocalCasted = rangeLocal.ReinterpretCast(); + Gather(addTmpLocal, indicesOriginLocal, rangeLocalCasted, (uint32_t)0, (uint32_t)computeRow); + PipeBarrier(); + Muls(addTmpLocal, addTmpLocal, (int)indicesMask_[i], computeRow); + PipeBarrier(); + Add(indicesLocal, indicesLocal, addTmpLocal, computeRow); + PipeBarrier(); + } + if constexpr (!isSort) { + PipeVToMte3(); + } + } + + __aicore__ inline void ComputeForSort(uint64_t process, bool isTail) + { + LocalTensor indicesLocalFp32 = indicesLocal.ReinterpretCast(); + uint64_t computeRow = isTail ? blockRemainLength_ : blockLength_; + uint64_t computeRowAligned = (computeRow + ALIGNED_BLOCK_NUM - 1) & ~(ALIGNED_BLOCK_NUM - 1); + uint64_t repeatTimes = computeRowAligned / ALIGNED_BLOCK_NUM; + + uint64_t repeatId = computeRow / ALIGNED_BLOCK_NUM; + uint64_t repeatRemain = computeRow % ALIGNED_BLOCK_NUM; + int addValue = indiceAddrOffset_ + process * blockLength_; + + Cast(indicesLocalFp32, indicesLocal, RoundMode::CAST_ROUND, computeRowAligned); + if (repeatRemain != 0) { + // 对齐处理：不足32的部分设为-1 + Duplicate(rangeLocal, -1, (uint32_t)ALIGNED_BLOCK_NUM); + PipeBarrier(); + Cast(rangeLocal, indicesLocalFp32[ALIGNED_BLOCK_NUM * repeatId], RoundMode::CAST_ROUND, (uint32_t)repeatRemain); + PipeBarrier(); + Cast(indicesLocalFp32[ALIGNED_BLOCK_NUM * repeatId], rangeLocal, RoundMode::CAST_ROUND, (uint32_t)ALIGNED_BLOCK_NUM); + PipeBarrier(); + } + Duplicate(posIdxLocal, -1, computeRowAligned); + PipeBarrier(); + CreateVecIndex(posIdxLocal, 0U, computeRow); + LocalTensor posIdxULocal = posIdxLocal.ReinterpretCast(); + PipeBarrier(); + Sort(resLocal, indicesLocalFp32, posIdxULocal, sortTmpLocal, repeatTimes); + PipeBarrier(); + Extract(indicesLocalFp32, posIdxULocal, resLocal, repeatTimes); + PipeBarrier(); + Cast(indicesLocal, indicesLocalFp32, RoundMode::CAST_ROUND, computeRowAligned); + PipeBarrier(); + Adds(posIdxLocal, posIdxLocal, addValue, computeRow); + PipeBarrier(); + PipeVToMte3(); + } + + __aicore__ inline void CopyOut(uint64_t process, bool isTail) + { + uint64_t outOffset = process * blockLength_; + uint64_t copyRow = isTail ? blockRemainLength_ : blockLength_; + DataCopyExtParams copyParams{1, static_cast(copyRow * sizeof(int)), 0, 0, 0}; + DataCopyPad(sortedIndicesGm_[outOffset], indicesLocal, copyParams); + if constexpr (isSort) { + DataCopyPad(posIndicesGm_[outOffset], posIdxLocal, copyParams); + } + PipeMte3ToS(); + } + +private: + GlobalTensor indicesGm_; + GlobalTensor indicesGmInt64_; + GlobalTensor sortedIndicesGm_; + GlobalTensor posIndicesGm_; + TBuf allUbBuf; + + LocalTensor allUbLocal; + LocalTensor indicesLocal; + LocalTensor indicesOriginLocal; + LocalTensor indicesInt64Local; + LocalTensor addTmpLocal; + LocalTensor rangeLocal; + LocalTensor posIdxLocal; + LocalTensor sortTmpLocal; + LocalTensor resLocal; + + uint64_t ubSize_; + uint64_t coreNum_; + uint64_t blockIdx_; + uint64_t indexDim_; + uint64_t computeNum_; + uint64_t blockNum_; + uint64_t blockLength_; + uint64_t blockRemainLength_; + uint64_t frontCoreNum_; + uint64_t tailCoreNum_; + uint64_t frontBlockNum_; + uint64_t tailBlockNum_; + const uint64_t* indicesMask_; + uint64_t indiceAddrOffset_; +}; + +} // namespace ScatterNdUpdateV2 diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_no_sort.h b/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_no_sort.h new file mode 100644 index 000000000..cc1e3657c --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_no_sort.h @@ -0,0 +1,152 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file scatter_nd_update_no_sort.h + * \brief Scatter Kernel (NoSort) + */ + +#include "kernel_operator.h" +#include "kernel_tiling/kernel_tiling.h" +#include "scatter_nd_update_common.h" + +namespace ScatterNdUpdateV2 { + +constexpr uint64_t ALIGNED_SIZE_INDEX = 8; + +template +class ScatterNdUpdateV2KernelNoSort { +public: + __aicore__ inline ScatterNdUpdateV2KernelNoSort() = delete; + __aicore__ inline ScatterNdUpdateV2KernelNoSort( + GM_ADDR updates, GM_ADDR output, GM_ADDR workSpace, const ScatterNdUpdateV2TilingData& tiling, TPipe& pipe) + { + InitParam(tiling); + InitBuffers(pipe); + SetGmAddr(updates, output, workSpace, tiling); + } + + __aicore__ inline void InitParam(const ScatterNdUpdateV2TilingData& tiling) + { + blockIdx_ = GetBlockIdx(); + CalcBlockDistribution(blockIdx_, tiling.scatterTiling.frontNum, tiling.scatterTiling.frontRow, + tiling.scatterTiling.tailRow, computeRow_, start_); + end_ = start_ + computeRow_; + totalIndexRow_ = tiling.linearIndexTiling.blockNum * tiling.linearIndexTiling.blockLength + + tiling.linearIndexTiling.blockRemainLength; + + scatterLength_ = tiling.scatterTiling.scatterLength; + scatterAlignLength_ = tiling.scatterTiling.scatterAlignLength; + ubLengthForUpdates_ = tiling.scatterTiling.ubLengthForUpdates; + scatterTileNum_ = tiling.scatterTiling.scatterTileNum; + scatterTileLength_ = tiling.scatterTiling.scatterTileLength; + scatterTileTail_ = tiling.scatterTiling.scatterTileTail; + scatterTileAlignLength_ = tiling.scatterTiling.scatterTileAlignLength; + + CalcIndexTileParams(); + } + + __aicore__ inline void CalcIndexTileParams() + { + uint64_t ubSizeBytes = ubLengthForUpdates_ * sizeof(T); + uint64_t updateSizeBytes = scatterTileLength_ * sizeof(T); + uint64_t remainBytes = ubSizeBytes - updateSizeBytes; + indexTileLength_ = (remainBytes / sizeof(int) / ALIGNED_SIZE_INDEX) * ALIGNED_SIZE_INDEX; + if (indexTileLength_ == 0) { + indexTileLength_ = ALIGNED_SIZE_INDEX; + } + indexTileNum_ = (totalIndexRow_ + indexTileLength_ - 1) / indexTileLength_; + indexTileTail_ = totalIndexRow_ - (indexTileNum_ - 1) * indexTileLength_; + } + + __aicore__ inline void InitBuffers(TPipe& pipe) + { + pipe.InitBuffer(indexQue_, DOUBLE_BUFFER, indexTileLength_ * sizeof(int)); + pipe.InitBuffer(updateQue_, DOUBLE_BUFFER, scatterTileLength_ * sizeof(T)); + } + + __aicore__ inline void SetGmAddr(GM_ADDR updates, GM_ADDR output, GM_ADDR workSpace, const ScatterNdUpdateV2TilingData& tiling) + { + linearIndicesGm_.SetGlobalBuffer((__gm__ int*)workSpace); + updatesGm_.SetGlobalBuffer((__gm__ T*)updates); + outputGm_.SetGlobalBuffer((__gm__ T*)output); + } + + __aicore__ inline void Process() + { + for (uint64_t tileIdx = 0; tileIdx < indexTileNum_; ++tileIdx) { + uint64_t curTileLen = (tileIdx == indexTileNum_ - 1) ? indexTileTail_ : indexTileLength_; + uint64_t gmOffset = tileIdx * indexTileLength_; + + LocalTensor indexLocal = indexQue_.AllocTensor(); + DataCopyExtParams indexCopyParams{1, static_cast(curTileLen * sizeof(int)), 0, 0, 0}; + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(indexLocal, linearIndicesGm_[gmOffset], indexCopyParams, padParams); + indexQue_.EnQue(indexLocal); + PipeMte2ToS(); + + LocalTensor indexData = indexQue_.DeQue(); + for (uint64_t i = 0; i < curTileLen; ++i) { + int64_t linearIndex = static_cast(indexData.GetValue(i)); + if (linearIndex >= (int64_t)start_ && linearIndex < (int64_t)end_) { + uint64_t idx = gmOffset + i; + ProcessOneIndex(idx, linearIndex); + } + } + indexQue_.FreeTensor(indexLocal); + } + } + + __aicore__ inline void ProcessOneIndex(uint64_t idx, int64_t linearIndex) + { + for (uint64_t tileIdx = 0; tileIdx < scatterTileNum_; ++tileIdx) { + uint64_t tileLength = (tileIdx == scatterTileNum_ - 1) ? scatterTileTail_ : scatterTileLength_; + + LocalTensor updateLocal = updateQue_.AllocTensor(); + uint64_t gmOffset = idx * scatterLength_ + tileIdx * scatterTileLength_; + DataCopyExtParams updateCopyParams{1, static_cast(tileLength * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(updateLocal, updatesGm_[gmOffset], updateCopyParams, padParams); + PipeMte2ToS(); + + uint64_t outOffset = linearIndex + tileIdx * scatterTileLength_; + DataCopyExtParams outParams{1, static_cast(tileLength * sizeof(T)), 0, 0, 0}; + DataCopyPad(outputGm_[outOffset], updateLocal, outParams); + PipeMte3ToS(); + + updateQue_.FreeTensor(updateLocal); + } + } + +private: + GlobalTensor linearIndicesGm_; + GlobalTensor updatesGm_; + GlobalTensor outputGm_; + TQue indexQue_; + TQue updateQue_; + + uint64_t blockIdx_; + uint64_t computeRow_; + uint64_t start_; + uint64_t end_; + uint64_t totalIndexRow_; + uint64_t scatterLength_; + uint64_t scatterAlignLength_; + uint64_t ubLengthForUpdates_; + uint64_t scatterTileNum_; + uint64_t scatterTileLength_; + uint64_t scatterTileTail_; + uint64_t scatterTileAlignLength_; + uint64_t indexTileLength_; + uint64_t indexTileNum_; + uint64_t indexTileTail_; +}; + +} // namespace ScatterNdUpdateV2 diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_v2.cpp b/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_v2.cpp new file mode 100644 index 000000000..57b25ab74 --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_v2.cpp @@ -0,0 +1,71 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + + /*! + * \file scatter_nd_update_v2.cpp + * \brief ScatterNdUpdateV2 算子入口 + */ +#include "scatter_nd_update_v2.h" +#include "scatter_nd_update_linear_index.h" +#include "scatter_nd_update_no_sort.h" +#include "scatter_nd_update_large_index.h" + +extern "C" __global__ __aicore__ void scatter_nd_update_v2(GM_ADDR varRef, GM_ADDR indices, + GM_ADDR updates, GM_ADDR output, GM_ADDR workSpace, GM_ADDR tiling) { + if (workSpace == nullptr) { + return; + } + GM_ADDR user = AscendC::GetUserWorkspace(workSpace); + if (user == nullptr) { + return; + } + GET_TILING_DATA(tilingData, tiling); + AscendC::TPipe tpipe; +#if (defined(DTYPE_VAR)) + // tilingKey: indexType * 10 + sortFlag + // indexType: 1=int32, 2=int64(cast), 3=int64(large); sortFlag: 0=非排序, 1=排序 + if (TILING_KEY_IS(11)) { + ScatterNdUpdateV2::LinearIndexKernel op1(indices, workSpace, tilingData, tpipe); + op1.Process(); + AscendC::SyncAll(); + tpipe.Destroy(); + AscendC::TPipe pipe; + ScatterNdUpdateV2::ScatterNdUpdateV2Kernel op2(updates, output, workSpace, tilingData, pipe); + op2.Process(); + } else if (TILING_KEY_IS(10)) { + ScatterNdUpdateV2::LinearIndexKernel op1(indices, workSpace, tilingData, tpipe); + op1.Process(); + AscendC::SyncAll(); + tpipe.Destroy(); + AscendC::TPipe pipe; + ScatterNdUpdateV2::ScatterNdUpdateV2KernelNoSort op2(updates, output, workSpace, tilingData, pipe); + op2.Process(); + } else if (TILING_KEY_IS(21)) { + ScatterNdUpdateV2::LinearIndexKernel op1(indices, workSpace, tilingData, tpipe); + op1.Process(); + AscendC::SyncAll(); + tpipe.Destroy(); + AscendC::TPipe pipe; + ScatterNdUpdateV2::ScatterNdUpdateV2Kernel op2(updates, output, workSpace, tilingData, pipe); + op2.Process(); + } else if (TILING_KEY_IS(20)) { + ScatterNdUpdateV2::LinearIndexKernel op1(indices, workSpace, tilingData, tpipe); + op1.Process(); + AscendC::SyncAll(); + tpipe.Destroy(); + AscendC::TPipe pipe; + ScatterNdUpdateV2::ScatterNdUpdateV2KernelNoSort op2(updates, output, workSpace, tilingData, pipe); + op2.Process(); + } else if (TILING_KEY_IS(30)) { + ScatterNdUpdateV2::LargeIndexKernel op(indices, updates, output, tilingData, tpipe); + op.Process(); + } +#endif +} diff --git a/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_v2.h b/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_v2.h new file mode 100644 index 000000000..4bf8171be --- /dev/null +++ b/csrc/ascend/moe/scatter_nd_update_v2/op_kernel/scatter_nd_update_v2.h @@ -0,0 +1,249 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file scatter_nd_update_v2.h + * \brief Scatter Kernel (Sort) + */ + +#include "kernel_operator.h" +#include "kernel_tiling/kernel_tiling.h" +#include "scatter_nd_update_common.h" + +namespace ScatterNdUpdateV2 { +template +class ScatterNdUpdateV2Kernel { +public: + __aicore__ inline ScatterNdUpdateV2Kernel() = delete; + __aicore__ inline ScatterNdUpdateV2Kernel( + GM_ADDR updates, GM_ADDR output, GM_ADDR workSpace, const ScatterNdUpdateV2TilingData& tiling, TPipe& pipe) + { + InitParams(tiling); + InitBuffers(pipe); + SetGmAddr(updates, output, workSpace, tiling); + } + + __aicore__ inline void InitParams(const ScatterNdUpdateV2TilingData& tiling) + { + blockIdx_ = GetBlockIdx(); + CalcBlockDistribution(blockIdx_, tiling.scatterTiling.frontNum, tiling.scatterTiling.frontRow, + tiling.scatterTiling.tailRow, computeRow_, start_); + end_ = start_ + computeRow_; + blockNum_ = tiling.linearIndexTiling.blockNum; + blockLength_ = tiling.linearIndexTiling.blockLength; + blockRemainLength_ = tiling.linearIndexTiling.blockRemainLength; + coreNum_ = tiling.linearIndexTiling.coreNum; + + scatterLength_ = tiling.scatterTiling.scatterLength; + scatterAlignLength_ = tiling.scatterTiling.scatterAlignLength; + ubLengthForUpdates_ = tiling.scatterTiling.ubLengthForUpdates; + formDim_ = tiling.scatterTiling.formDim; + copyRow_ = tiling.scatterTiling.copyRow; + scatterTileNum_ = tiling.scatterTiling.scatterTileNum; + scatterTileLength_ = tiling.scatterTiling.scatterTileLength; + scatterTileTail_ = tiling.scatterTiling.scatterTileTail; + scatterTileAlignLength_ = tiling.scatterTiling.scatterTileAlignLength; + } + + __aicore__ inline void InitBuffers(TPipe& pipe) + { + pipe.InitBuffer(indiceQue_, DOUBLE_BUFFER, blockLength_ * sizeof(int)); + pipe.InitBuffer(posIdxQue_, DOUBLE_BUFFER, blockLength_ * sizeof(int)); + pipe.InitBuffer(updateQue_, DOUBLE_BUFFER, ubLengthForUpdates_ * sizeof(T)); + } + + __aicore__ inline void SetGmAddr(GM_ADDR updates, GM_ADDR output, GM_ADDR workSpace, const ScatterNdUpdateV2TilingData& tiling) + { + sortedIndicesGm_.SetGlobalBuffer((__gm__ int*)workSpace); + posIndicesGm_.SetGlobalBuffer((__gm__ int*)workSpace + tiling.linearIndexTiling.sortWorkspace); + updatesGm_.SetGlobalBuffer((__gm__ T*)updates); + outputGm_.SetGlobalBuffer((__gm__ T*)output); + } + + __aicore__ inline void Process() + { + for (uint64_t i = 0; i < blockNum_; ++i) { + CopyIndicesIn(i, false); + Compute(i, false); + PipeMte3ToS(); + } + if (blockRemainLength_ != 0) { + CopyIndicesIn(blockNum_, true); + Compute(blockNum_, true); + PipeMte3ToS(); + } + } + + + __aicore__ inline void CopyIndicesIn(uint64_t process, bool isTail) + { + uint64_t copyNum = isTail ? blockRemainLength_ : blockLength_; + LocalTensor indiceLocal = indiceQue_.AllocTensor(); + LocalTensor posIdxLocal = posIdxQue_.AllocTensor(); + uint64_t indicesOffset = isTail ? (blockNum_ * blockLength_) : (process * blockLength_); + DataCopyParams indiceCopyParams{1, static_cast(copyNum * sizeof(int)), 0, 0}; + DataCopyPadParams padParams{true, 0, 0, 0}; + DataCopyPad(indiceLocal, sortedIndicesGm_[indicesOffset], indiceCopyParams, padParams); + DataCopyPad(posIdxLocal, posIndicesGm_[indicesOffset], indiceCopyParams, padParams); + PipeMte2ToS(); + PipeBarrier(); + UpdateSearchParam(indiceLocal, isTail); + indiceQue_.EnQue(indiceLocal); + posIdxQue_.EnQue(posIdxLocal); + } + + __aicore__ inline void CopyUpdateIn(LocalTensor &updateLocal, uint64_t gmIdx, uint64_t ubIdx, uint64_t tileIdx, uint64_t tileLength) + { + uint64_t gmOffset = gmIdx * scatterLength_ + tileIdx * scatterTileLength_; + uint64_t ubOffset = ubIdx * scatterTileAlignLength_; + DataCopyExtParams updateCopyParams{1, static_cast(tileLength * sizeof(T)), 0, 0, 0}; + DataCopyPadExtParams padParams{true, 0, 0, 0}; + DataCopyPad(updateLocal[ubOffset], updatesGm_[gmOffset], updateCopyParams, padParams); + PipeMte2ToS(); + } + + // 降序数组：二分查找边界 + __aicore__ inline int64_t findFirstLt(LocalTensor &indiceLocal, int64_t target, bool isTail) + { + int64_t left = 0; + int64_t right = (isTail ? blockRemainLength_ : blockLength_) - 1; + int64_t res = isTail ? blockRemainLength_ : blockLength_; + while (left <= right) { + int64_t mid = left + (right - left) / 2; + int64_t value = indiceLocal.GetValue(mid); + if (value < target) { + res = mid; + right = mid - 1; + } else { + left = mid + 1; + } + } + return res; + } + + __aicore__ inline int64_t findLastGe(LocalTensor &indiceLocal, int64_t target, bool isTail) + { + int64_t left = 0; + int64_t right = (isTail ? blockRemainLength_ : blockLength_) - 1; + int64_t res = -1; + while (left <= right) { + int64_t mid = left + (right - left) / 2; + int64_t value = indiceLocal.GetValue(mid); + if (value >= target) { + res = mid; + left = mid + 1; + } else { + right = mid - 1; + } + } + return res; + } + + __aicore__ inline void UpdateSearchParam(LocalTensor &indiceLocal, bool isTail) + { + int64_t searchNum = isTail ? blockRemainLength_ : blockLength_; + leftBound_ = findFirstLt(indiceLocal, end_, isTail); + rightBound_ = findLastGe(indiceLocal, start_, isTail); + isValidBound_ = (leftBound_ < searchNum && rightBound_ != -1 && leftBound_ <= rightBound_); + } + + __aicore__ inline void Compute(uint64_t process, bool isTail) + { + LocalTensor indiceLocal = indiceQue_.DeQue(); + LocalTensor posIdxLocal = posIdxQue_.DeQue(); + + if (!isValidBound_) { + indiceQue_.FreeTensor(indiceLocal); + posIdxQue_.FreeTensor(posIdxLocal); + return; + } + + for (uint64_t tileIdx = 0; tileIdx < scatterTileNum_; ++tileIdx) { + uint64_t tileLength = (tileIdx == scatterTileNum_ - 1) ? scatterTileTail_ : scatterTileLength_; + + uint64_t inUbNum = 0; + LocalTensor updateLocal; + + lastProcessedIdx_ = -1; + + for (int64_t i = rightBound_; i >= leftBound_; --i) { + if (inUbNum == 0) { + updateLocal = updateQue_.AllocTensor(); + } + int64_t posIdx = posIdxLocal.GetValue(i); + CopyUpdateIn(updateLocal, posIdx, inUbNum, tileIdx, tileLength); + inUbNum++; + + if (inUbNum == copyRow_) { + updateQue_.EnQue(updateLocal); + CopyOut(inUbNum, i, indiceLocal, posIdxLocal, tileIdx, tileLength); + inUbNum = 0; + } + if (i == leftBound_ && inUbNum != 0) { + updateQue_.EnQue(updateLocal); + CopyOut(inUbNum, i, indiceLocal, posIdxLocal, tileIdx, tileLength); + } + } + } + + indiceQue_.FreeTensor(indiceLocal); + posIdxQue_.FreeTensor(posIdxLocal); + } + + __aicore__ inline void CopyOut(uint64_t inUbNum, int64_t curIdx, LocalTensor &indiceLocal, + LocalTensor &posIdxLocal, uint64_t tileIdx, uint64_t tileLength) + { + LocalTensor updateLocal = updateQue_.DeQue(); + DataCopyExtParams outParams{1, static_cast(tileLength * sizeof(T)), 0, 0, 0}; + for (int64_t i = curIdx + inUbNum - 1; i >= curIdx; --i) { + int64_t curIdxValue = indiceLocal.GetValue(i); + if (curIdxValue == lastProcessedIdx_) continue; + lastProcessedIdx_ = curIdxValue; + uint64_t outOffset = curIdxValue + tileIdx * scatterTileLength_; + uint64_t updateOffset = (curIdx + inUbNum - 1 - i) * scatterTileAlignLength_; + DataCopyPad(outputGm_[outOffset], updateLocal[updateOffset], outParams); + } + PipeMte3ToS(); + updateQue_.FreeTensor(updateLocal); + } + +private: + + GlobalTensor sortedIndicesGm_; + GlobalTensor posIndicesGm_; + GlobalTensor updatesGm_; + GlobalTensor outputGm_; + TQue indiceQue_; + TQue posIdxQue_; + TQue updateQue_; + + uint64_t blockIdx_; + uint64_t computeRow_; + uint64_t start_; + uint64_t end_; + uint64_t blockNum_; + uint64_t blockLength_; + uint64_t blockRemainLength_; + uint64_t scatterLength_; + uint64_t scatterAlignLength_; + uint64_t ubLengthForUpdates_; + uint64_t formDim_; + uint64_t copyRow_; + uint64_t coreNum_; + uint64_t scatterTileNum_; + uint64_t scatterTileLength_; + uint64_t scatterTileTail_; + uint64_t scatterTileAlignLength_; + int64_t leftBound_; + int64_t rightBound_; + bool isValidBound_; + int64_t lastProcessedIdx_; +}; +} // namespace ScatterNdUpdateV2 diff --git a/csrc/ascend/moe/transpose_kv_cache_by_block/CMakeLists.txt b/csrc/ascend/moe/transpose_kv_cache_by_block/CMakeLists.txt new file mode 100644 index 000000000..298ed5677 --- /dev/null +++ b/csrc/ascend/moe/transpose_kv_cache_by_block/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() \ No newline at end of file diff --git a/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/CMakeLists.txt b/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/CMakeLists.txt new file mode 100644 index 000000000..734810a74 --- /dev/null +++ b/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/CMakeLists.txt @@ -0,0 +1,24 @@ +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + transpose_kv_cache_by_block_def.cpp + ) +endif() + +add_ops_compile_options( + OP_NAME TransposeKvCacheByBlock + OPTIONS + --cce-auto-sync=off + -Wno-deprecated-declarations + -Werror + -mllvm -cce-aicore-hoist-movemask=false + --op_relocatable_kernel_binary=true +) + +if (NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE transpose_kv_cache_by_block ACLNNTYPE aclnn) + target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ) +endif() diff --git a/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/transpose_kv_cache_by_block_def.cpp b/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/transpose_kv_cache_by_block_def.cpp new file mode 100644 index 000000000..c4c11a4e4 --- /dev/null +++ b/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/transpose_kv_cache_by_block_def.cpp @@ -0,0 +1,36 @@ +#include "register/op_def_registry.h" + +namespace ops { +class TransposeKvCacheByBlock : public OpDef { +public: + explicit TransposeKvCacheByBlock(const char* name) : OpDef(name) + { + this->Input("KCache") + .ParamType(DYNAMIC) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("VCache") + .ParamType(DYNAMIC) + .DataType({ge::DT_FLOAT16, ge::DT_BF16}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Input("blockIDs") + .ParamType(REQUIRED) + .DataType({ge::DT_INT64, ge::DT_INT64}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}); + this->Attr("blockSize").Int(); + this->Attr("headNum").Int(); + this->Attr("headDim").Int(); + this->Attr("splitNum").Int(); + this->Attr("layerNum").Int(); + + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + + } +}; + +OP_ADD(TransposeKvCacheByBlock); +} diff --git a/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/transpose_kv_cache_by_block_proto.cpp b/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/transpose_kv_cache_by_block_proto.cpp new file mode 100644 index 000000000..24f2d84d5 --- /dev/null +++ b/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/transpose_kv_cache_by_block_proto.cpp @@ -0,0 +1,36 @@ +/** + * This program is free software, you can redistribute it and/or modify it. + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file transpose_kv_cache_by_block_proto.cpp + * \brief + */ +#include +#include +#include "error/ops_error.h" + +using namespace ge; + +namespace ops { + +static ge::graphStatus InferShapeTransposeKvCacheByBlock(gert::InferShapeContext* context) +{ + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus InferDataTypeTransposeKvCacheByBlock(gert::InferDataTypeContext *context) +{ + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(TransposeKvCacheByBlock) + .InferShape(InferShapeTransposeKvCacheByBlock) + .InferDataType(InferDataTypeTransposeKvCacheByBlock); +} // namespace ops diff --git a/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/transpose_kv_cache_by_block_tiling.cpp b/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/transpose_kv_cache_by_block_tiling.cpp new file mode 100644 index 000000000..5e8ae289f --- /dev/null +++ b/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/transpose_kv_cache_by_block_tiling.cpp @@ -0,0 +1,182 @@ +#include "transpose_kv_cache_by_block_tiling.h" +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" +#include "log/ops_log.h" +#include + +namespace optiling { + +constexpr uint64_t DATA_SIZE = 2; +constexpr uint64_t BLOCK_SIZE = 32; +constexpr uint64_t DB_ON = 2; + +constexpr uint32_t FULL_LOAD = 0; +constexpr uint32_t SPLIT_BLOCK_SIZE_ALIGNED_AND_DB = 1; +constexpr uint32_t SPLIT_BLOCK_SIZE_UNALIGNED_AND_DB = 3; +constexpr uint32_t SPLIT_BLOCK_SIZE_ALIGNED_AND_NOT_DB = 2; +constexpr uint32_t SPLIT_BLOCK_SIZE_UNALIGNED_AND_NOT_DB = 4; + +void findFactorsOptimized(std::vector &factors, int64_t n) { + + for (int64_t i = 1; i * i <= n; i++) { + if (n % i == 0) { + factors.push_back(i); + + if (i != n / i) { + factors.push_back(n / i); + } + } + } + + sort(factors.begin(), factors.end()); +} + +ge::graphStatus CalTiling(gert::TilingContext* context, TransposeKvCacheByBlockTilingData &tiling) +{ + fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo(); + OPS_LOG_E_IF_NULL(context, platformInfoPtr, return ge::GRAPH_FAILED); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + int64_t useCoreNum = ascendcPlatform.GetCoreNumAiv(); + uint64_t ubSize; + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); + + auto attr = context->GetAttrs(); + OPS_LOG_E_IF_NULL(context, attr, return ge::GRAPH_FAILED); + const int64_t* blockSizePtr = attr->GetAttrPointer(0); + const int64_t* headNumPtr = attr->GetAttrPointer(1); + const int64_t* headDimPtr = attr->GetAttrPointer(2); + const int64_t* splitNumPtr = attr->GetAttrPointer(3); + const int64_t* layerNumPtr = attr->GetAttrPointer(4); + OPS_CHECK(blockSizePtr == nullptr || headNumPtr == nullptr || headDimPtr == nullptr || + splitNumPtr == nullptr || layerNumPtr == nullptr, + OPS_LOG_E(context->GetNodeName(), "Get attr failed."), + return ge::GRAPH_FAILED); + + auto blockIDsTensor = context->GetDynamicInputTensor(2, 0); + OPS_LOG_E_IF_NULL(context, blockIDsTensor, return ge::GRAPH_FAILED); + + gert::Shape blockIDsTensorShape = blockIDsTensor->GetStorageShape(); + int64_t calBlockNum = static_cast(blockIDsTensorShape.GetDim(0)); + + tiling.set_calBlockNum(static_cast(calBlockNum)); + + int64_t blockSize = *blockSizePtr; + int64_t headNum = *headNumPtr; + int64_t headDim = *headDimPtr; + int64_t splitNum = *splitNumPtr; + int64_t layerNum = *layerNumPtr; + uint32_t tilingKey = FULL_LOAD; + + if (headDim * DATA_SIZE % BLOCK_SIZE != 0) { + OPS_LOG_E(context, "headDim * DATA_SIZE must be a multiple of 32 bytes."); + return ge::GRAPH_FAILED; + } + + std::vector factors; + findFactorsOptimized(factors, useCoreNum); + + uint32_t factorIndex = 0; + bool findSplitNum = true; + int64_t blockSizeSplitNum = factors[factorIndex]; + uint64_t dataSizeloadOnce = blockSize * headNum * headDim * DATA_SIZE; + // if can full load, not split blockSize and db + if (dataSizeloadOnce > ubSize) { + tilingKey = SPLIT_BLOCK_SIZE_ALIGNED_AND_DB; + // split blockSize and db + while (dataSizeloadOnce > (ubSize / DB_ON)) { + factorIndex += 1; + if (factorIndex == factors.size()) { + tilingKey = FULL_LOAD; + findSplitNum = false; + break; + } + blockSizeSplitNum = factors[factorIndex]; + dataSizeloadOnce = ((blockSize + blockSizeSplitNum - 1) / blockSizeSplitNum) * headNum * headDim * DATA_SIZE; + } + if (tilingKey == SPLIT_BLOCK_SIZE_ALIGNED_AND_DB && (blockSize % blockSizeSplitNum != 0)) { + tilingKey = SPLIT_BLOCK_SIZE_UNALIGNED_AND_DB; + } + } + + if (!findSplitNum) { + tilingKey = SPLIT_BLOCK_SIZE_ALIGNED_AND_NOT_DB; + // split blockSize but not db + findSplitNum = true; + factorIndex = 0; + blockSizeSplitNum = factors[factorIndex]; + dataSizeloadOnce = blockSize * headNum * headDim * DATA_SIZE; + while (dataSizeloadOnce > ubSize) { + factorIndex += 1; + if (factorIndex == factors.size()) { + tilingKey = FULL_LOAD; + findSplitNum = false; + break; + } + blockSizeSplitNum = factors[factorIndex]; + dataSizeloadOnce = ((blockSize + blockSizeSplitNum - 1) / blockSizeSplitNum) * headNum * headDim * DATA_SIZE; + } + if (tilingKey == SPLIT_BLOCK_SIZE_ALIGNED_AND_NOT_DB && (blockSize % blockSizeSplitNum != 0)) { + tilingKey = SPLIT_BLOCK_SIZE_UNALIGNED_AND_NOT_DB; + } + } + + // headNum * headDim too large + if (!findSplitNum) { + OPS_LOG_E(context, "headNum * headDim * sizeof(half) > ubSize " + "or blockSize * headNum * headDim * sizeof(half) > ubSize * vectorCoreNum. " + "Currently, splitting headNum or headDim is not supported."); + return ge::GRAPH_FAILED; + } + tiling.set_blockSizePerTime(static_cast((blockSize + blockSizeSplitNum - 1) / blockSizeSplitNum)); + tiling.set_blockSizePerTimeTail(static_cast(blockSize % blockSizeSplitNum)); + tiling.set_blockSizeSplitNum(static_cast(blockSizeSplitNum)); + + tiling.set_blockSize(static_cast(blockSize)); + tiling.set_headNum(static_cast(headNum)); + tiling.set_headDim(static_cast(headDim)); + tiling.set_splitNum(static_cast(splitNum)); + tiling.set_layerNum(static_cast(layerNum)); + + int64_t totalRound = layerNum * calBlockNum; + + if ((totalRound * blockSizeSplitNum) < useCoreNum) { + useCoreNum = totalRound * blockSizeSplitNum; + } + int64_t blockPerCore = totalRound / (useCoreNum / blockSizeSplitNum); + int64_t tailCoreNum = totalRound % (useCoreNum / blockSizeSplitNum); + + tiling.set_useCoreNum(static_cast(useCoreNum)); + tiling.set_blockPerCore(static_cast(blockPerCore)); + tiling.set_tailCoreNum(static_cast(tailCoreNum)); + context->SetBlockDim(useCoreNum); + context->SetTilingKey(tilingKey); + + return ge::GRAPH_SUCCESS; +} + + +static ge::graphStatus TransposeKvCacheByBlockTilingFunc(gert::TilingContext* context) +{ + + TransposeKvCacheByBlockTilingData tiling; + auto status = CalTiling(context, tiling); + OP_CHECK(status != ge::GRAPH_SUCCESS, OPS_LOG_E(context->GetNodeName(), "Cal tiling failed."), + return ge::GRAPH_FAILED); + + tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); + + return ge::GRAPH_SUCCESS; +} + +struct TransposeKvCacheByBlockCompileInfo {}; +ge::graphStatus TilingParseForTransposeKvCacheByBlock(gert::TilingParseContext *context) +{ + (void)context; + return ge::GRAPH_SUCCESS; +} + +IMPL_OP_OPTILING(TransposeKvCacheByBlock) + .Tiling(TransposeKvCacheByBlockTilingFunc) + .TilingParse(TilingParseForTransposeKvCacheByBlock); +} \ No newline at end of file diff --git a/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/transpose_kv_cache_by_block_tiling.h b/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/transpose_kv_cache_by_block_tiling.h new file mode 100644 index 000000000..c3b0901b9 --- /dev/null +++ b/csrc/ascend/moe/transpose_kv_cache_by_block/op_host/transpose_kv_cache_by_block_tiling.h @@ -0,0 +1,23 @@ +#include "register/tilingdata_base.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(TransposeKvCacheByBlockTilingData) + // shape info + // TILING_DATA_FIELD_DEF(uint32_t, blockNum); + TILING_DATA_FIELD_DEF(uint32_t, blockSize); + TILING_DATA_FIELD_DEF(uint32_t, headNum); + TILING_DATA_FIELD_DEF(uint32_t, headDim); + TILING_DATA_FIELD_DEF(uint32_t, splitNum); + TILING_DATA_FIELD_DEF(uint32_t, layerNum); + // tiling info + TILING_DATA_FIELD_DEF(uint32_t, useCoreNum); + TILING_DATA_FIELD_DEF(uint32_t, blockPerCore); + TILING_DATA_FIELD_DEF(uint32_t, tailCoreNum); + TILING_DATA_FIELD_DEF(uint32_t, calBlockNum); + TILING_DATA_FIELD_DEF(uint32_t, blockSizePerTime); + TILING_DATA_FIELD_DEF(uint32_t, blockSizePerTimeTail); + TILING_DATA_FIELD_DEF(uint32_t, blockSizeSplitNum); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(TransposeKvCacheByBlock, TransposeKvCacheByBlockTilingData) +} diff --git a/csrc/ascend/moe/transpose_kv_cache_by_block/op_kernel/common.h b/csrc/ascend/moe/transpose_kv_cache_by_block/op_kernel/common.h new file mode 100644 index 000000000..e74089fa5 --- /dev/null +++ b/csrc/ascend/moe/transpose_kv_cache_by_block/op_kernel/common.h @@ -0,0 +1,16 @@ +#include "kernel_operator.h" +using namespace AscendC; + +#ifndef __OP_KERNEL_KV_CACHE_TRANSPOSE_H__ +#define __OP_KERNEL_KV_CACHE_TRANSPOSE_H__ + +template +__aicore__ inline __gm__ T* GetTensorAddr(uint16_t index, GM_ADDR tensorPtr) { + __gm__ uint64_t* dataAddr = reinterpret_cast<__gm__ uint64_t*>(tensorPtr); + // The offset of the data address from the first address. + uint64_t tensorPtrOffset = *dataAddr; + // Moving 3 bits to the right means dividing by sizeof(uint64 t). + __gm__ uint64_t* retPtr = dataAddr + (tensorPtrOffset >> 3); + return reinterpret_cast<__gm__ T*>(*(retPtr + index)); +} +#endif \ No newline at end of file diff --git a/csrc/ascend/moe/transpose_kv_cache_by_block/op_kernel/full_load.h b/csrc/ascend/moe/transpose_kv_cache_by_block/op_kernel/full_load.h new file mode 100644 index 000000000..86951f699 --- /dev/null +++ b/csrc/ascend/moe/transpose_kv_cache_by_block/op_kernel/full_load.h @@ -0,0 +1,141 @@ +#include "common.h" + +template +class TransposeKvCacheByBlockKernelFullLoad { + protected: + TQueBind vecInQueue_; + GlobalTensor kCacheGm_; + GlobalTensor vCacheGm_; + GlobalTensor blockIDsGm_; + + GM_ADDR kCachePtr_; + GM_ADDR vCachePtr_; + + // shape info + uint32_t blockNum_; + uint32_t blockSize_; + uint32_t headNum_; + uint32_t headDim_; + uint32_t splitNum_; + uint32_t layerNum_; + // tiling info + uint32_t useCoreNum_; + uint32_t blockPerCore_; + uint32_t tailCoreNum_; + uint32_t calBlockNum_; + + uint32_t srcFactor_; + uint32_t dstFactor_; + uint32_t copyOutLength_; + uint32_t dataBlockSize_; + + __aicore__ inline void CopyIn(GlobalTensor &cacheGm, uint32_t offsetBlock, DataCopyParams &repeatParams) { + LocalTensor cacheLocal = vecInQueue_.AllocTensor(); + for (uint32_t i = 0; i < splitNum_; ++i) { + DataCopy(cacheLocal[i * dstFactor_], cacheGm[i * srcFactor_ + offsetBlock], repeatParams); + } + vecInQueue_.EnQue(cacheLocal); + } + + __aicore__ inline void CopyOut(GlobalTensor &cacheGm, uint32_t offsetBlock) { + LocalTensor cacheLocal = vecInQueue_.DeQue(); + DataCopy(cacheGm[offsetBlock], cacheLocal, copyOutLength_); + vecInQueue_.FreeTensor(cacheLocal); + } + + __aicore__ inline void SetGlobalBuffers(uint32_t layerId) { + kCacheGm_.SetGlobalBuffer(GetTensorAddr(layerId, kCachePtr_)); + vCacheGm_.SetGlobalBuffer(GetTensorAddr(layerId, vCachePtr_)); + } + + __aicore__ inline void Caloffset(uint32_t &startBlock, uint32_t &endBlock, uint32_t &startLayer, uint32_t &endLayer) { + uint32_t blockIdx = GetBlockIdx(); + uint32_t curBlockStart; + uint32_t curBlocknum; + + if (blockIdx < tailCoreNum_) { + curBlockStart = blockIdx * (blockPerCore_ + 1); + curBlocknum = blockPerCore_ + 1; + } else { + curBlockStart = blockIdx * blockPerCore_ + tailCoreNum_; + curBlocknum = blockPerCore_; + } + uint32_t curBlockEnd = curBlockStart + curBlocknum; + startBlock = curBlockStart / layerNum_; + startLayer = curBlockStart % layerNum_; + endBlock = (curBlockEnd + layerNum_ - 1) / layerNum_; + endLayer = curBlockEnd % layerNum_; + if (endLayer == 0) { + endLayer = layerNum_; + } + } + + + public: + __aicore__ inline void Init(GM_ADDR KCache, GM_ADDR VCache, GM_ADDR blockIDs, + TransposeKvCacheByBlockTilingData* tilingData, TPipe* tPipe) { + kCachePtr_ = KCache; + vCachePtr_ = VCache; + blockIDsGm_.SetGlobalBuffer((__gm__ int64_t*)blockIDs); + + // shape info + blockSize_ = tilingData->blockSize; + headNum_ = tilingData->headNum; + headDim_ = tilingData->headDim; + splitNum_ = tilingData->splitNum; + layerNum_ = tilingData->layerNum; + // tiling info + useCoreNum_ = tilingData->useCoreNum; + blockPerCore_ = tilingData->blockPerCore; + tailCoreNum_ = tilingData->tailCoreNum; + calBlockNum_ = tilingData->calBlockNum; + + tPipe->InitBuffer(vecInQueue_, 1, TOTAL_UB_SIZE); + srcFactor_ = blockSize_ * headNum_ / splitNum_ * headDim_; + dstFactor_ = headNum_ / splitNum_ * headDim_; + copyOutLength_ = blockSize_ * headNum_ * headDim_; + dataBlockSize_ = static_cast(AscendC::GetDataBlockSizeInBytes()); + } + + __aicore__ inline void Process() { + DataCopyParams repeatParams; + repeatParams.blockCount = blockSize_; + repeatParams.blockLen = headNum_ / splitNum_ * headDim_ * sizeof(T) / dataBlockSize_; + repeatParams.srcStride = 0; + repeatParams.dstStride = (headNum_ * headDim_ - headNum_ / splitNum_ * headDim_) * sizeof(T) / dataBlockSize_; + + uint32_t startBlock; + uint32_t endBlock; + uint32_t startLayer; + uint32_t endLayer; + + Caloffset(startBlock, endBlock, startLayer, endLayer); + for (uint32_t i = startBlock; i < endBlock; ++i) { + int64_t blockId = blockIDsGm_.GetValue(i); + uint32_t offsetBlock = blockId * blockSize_ * headNum_ * headDim_; + uint32_t realStartLayer; + uint32_t realEndLayer; + if (i == startBlock) { + realStartLayer = startLayer; + } else { + realStartLayer = 0; + } + + if (i == (endBlock - 1)) { + realEndLayer = endLayer; + } else { + realEndLayer = layerNum_; + } + for (uint32_t layerId = realStartLayer; layerId < realEndLayer; ++layerId) { + SetGlobalBuffers(layerId); + + CopyIn(kCacheGm_, offsetBlock, repeatParams); + CopyOut(kCacheGm_, offsetBlock); + + CopyIn(vCacheGm_, offsetBlock, repeatParams); + CopyOut(vCacheGm_, offsetBlock); + } + + } + } +}; \ No newline at end of file diff --git a/csrc/ascend/moe/transpose_kv_cache_by_block/op_kernel/general.h b/csrc/ascend/moe/transpose_kv_cache_by_block/op_kernel/general.h new file mode 100644 index 000000000..1526a4db0 --- /dev/null +++ b/csrc/ascend/moe/transpose_kv_cache_by_block/op_kernel/general.h @@ -0,0 +1,190 @@ +#include "common.h" + +template +class TransposeKvCacheByBlockKernelGeneral { + protected: + TQueBind queBind_; + GlobalTensor kCacheGm_; + GlobalTensor vCacheGm_; + GlobalTensor blockIDsGm_; + + GM_ADDR kCachePtr_; + GM_ADDR vCachePtr_; + + // shape info + uint32_t blockNum_; + uint32_t blockSize_; + uint32_t headNum_; + uint32_t headDim_; + uint32_t splitNum_; + uint32_t layerNum_; + uint32_t headNumSplited_; + uint32_t blockSizeSplitNum_; + + // tiling info + uint32_t useCoreNum_; + uint32_t blockPerCore_; + uint32_t tailCoreNum_; + uint32_t calBlockNum_; + + uint32_t srcFactor_; + uint32_t dstFactor_; + uint32_t copyOutLength_; + + uint32_t blockSizePerTime_; + uint32_t blockSizePerTimeTail_; + + uint32_t blockIdx_; + uint32_t dataBlockSize_; + bool needSync_; + + __aicore__ inline void CopyIn(GlobalTensor &cacheGm, uint32_t offsetBlock, DataCopyParams &repeatParams) { + LocalTensor cacheLocal = queBind_.AllocTensor(); + for (uint32_t i = 0; i < splitNum_; ++i) { + DataCopy(cacheLocal[i * dstFactor_], cacheGm[i * srcFactor_ + offsetBlock], repeatParams); + } + queBind_.EnQue(cacheLocal); + } + + __aicore__ inline void CopyOut(GlobalTensor &cacheGm, uint32_t offsetBlock) { + LocalTensor cacheLocal = queBind_.DeQue(); + AscendC::CrossCoreSetFlag<0x0, PIPE_MTE2>(0x8); + AscendC::CrossCoreWaitFlag(0x8); + DataCopy(cacheGm[offsetBlock], cacheLocal, copyOutLength_); + queBind_.FreeTensor(cacheLocal); + } + + __aicore__ inline void SetGlobalBuffers(uint32_t layerId) { + kCacheGm_.SetGlobalBuffer(GetTensorAddr(layerId, kCachePtr_)); + vCacheGm_.SetGlobalBuffer(GetTensorAddr(layerId, vCachePtr_)); + } + + __aicore__ inline void Caloffset(uint32_t &startBlock, uint32_t &endBlock, uint32_t &startLayer, uint32_t &endLayer) { + + uint32_t curBlockStart; + uint32_t curBlocknum; + uint32_t groupBlockIdx = blockIdx_ / blockSizeSplitNum_; + if (groupBlockIdx < tailCoreNum_) { + needSync_ = false; + curBlockStart = groupBlockIdx * (blockPerCore_ + 1); + curBlocknum = blockPerCore_ + 1; + } else { + needSync_ = true; + curBlockStart = groupBlockIdx * blockPerCore_ + tailCoreNum_; + curBlocknum = blockPerCore_; + } + uint32_t curBlockEnd = curBlockStart + curBlocknum; + startBlock = curBlockStart / layerNum_; + startLayer = curBlockStart % layerNum_; + endBlock = (curBlockEnd + layerNum_ - 1) / layerNum_; + endLayer = curBlockEnd % layerNum_; + if (endLayer == 0) { + endLayer = layerNum_; + } + } + + + public: + __aicore__ inline void Init(GM_ADDR KCache, GM_ADDR VCache, GM_ADDR blockIDs, + TransposeKvCacheByBlockTilingData* tilingData, TPipe* tPipe) { + kCachePtr_ = KCache; + vCachePtr_ = VCache; + blockIDsGm_.SetGlobalBuffer((__gm__ int64_t*)blockIDs); + blockIdx_ = GetBlockIdx(); + // shape info + blockSize_ = tilingData->blockSize; + headNum_ = tilingData->headNum; + headDim_ = tilingData->headDim; + splitNum_ = tilingData->splitNum; + layerNum_ = tilingData->layerNum; + // tiling info + useCoreNum_ = tilingData->useCoreNum; + blockPerCore_ = tilingData->blockPerCore; + tailCoreNum_ = tilingData->tailCoreNum; + calBlockNum_ = tilingData->calBlockNum; + blockSizeSplitNum_ = tilingData->blockSizeSplitNum; + blockSizePerTime_ = tilingData->blockSizePerTime; + blockSizePerTimeTail_ = tilingData->blockSizePerTimeTail; + headNumSplited_ = headNum_ / splitNum_; + + tPipe->InitBuffer(queBind_, DB, TOTAL_UB_SIZE / DB); + srcFactor_ = blockSize_ * headNumSplited_ * headDim_; + dstFactor_ = headNumSplited_ * headDim_; + copyOutLength_ = blockSizePerTime_ * headNum_ * headDim_; + dataBlockSize_ = static_cast(AscendC::GetDataBlockSizeInBytes()); + } + + __aicore__ inline void Process() { + DataCopyParams repeatParams; + repeatParams.blockCount = blockSizePerTime_; + repeatParams.blockLen = headNumSplited_ * headDim_ * sizeof(T) / dataBlockSize_; + repeatParams.srcStride = 0; + repeatParams.dstStride = (headNum_ * headDim_ - headNumSplited_ * headDim_) * sizeof(T) / dataBlockSize_; + + uint32_t startBlock; + uint32_t endBlock; + uint32_t startLayer; + uint32_t endLayer; + + Caloffset(startBlock, endBlock, startLayer, endLayer); + + for (uint32_t i = startBlock; i < endBlock; ++i) { + int64_t blockId = blockIDsGm_.GetValue(i); + uint32_t offsetBlock = blockId * blockSize_ * headNum_ * headDim_; + uint32_t realStartLayer; + uint32_t realEndLayer; + if (i == startBlock) { + realStartLayer = startLayer; + } else { + realStartLayer = 0; + } + + if (i == (endBlock - 1)) { + realEndLayer = endLayer; + } else { + realEndLayer = layerNum_; + } + for (uint32_t layerId = realStartLayer; layerId < realEndLayer; ++layerId) { + SetGlobalBuffers(layerId); + uint32_t blockSizeIndex = blockIdx_ % blockSizeSplitNum_; + uint32_t srcOffset; + uint32_t dstOffset; + if constexpr (needHandleUnFactorSplit) { + // handle tail + if (blockSizeIndex >= blockSizePerTimeTail_) { + repeatParams.blockCount = (blockSizePerTime_ - 1); + copyOutLength_ = (blockSizePerTime_ - 1) * headNum_ * headDim_; + srcOffset = (blockSizeIndex * blockSizePerTime_ - (blockSizeIndex - blockSizePerTimeTail_)) * headNumSplited_ * headDim_; + dstOffset = (blockSizeIndex * blockSizePerTime_ - (blockSizeIndex - blockSizePerTimeTail_)) * headNum_ * headDim_; + } else { + repeatParams.blockCount = blockSizePerTime_; + copyOutLength_ = blockSizePerTime_ * headNum_ * headDim_; + srcOffset = blockSizeIndex * blockSizePerTime_ * headNumSplited_ * headDim_; + dstOffset = blockSizeIndex * blockSizePerTime_ * headNum_ * headDim_; + } + } else { + repeatParams.blockCount = blockSizePerTime_; + copyOutLength_ = blockSizePerTime_ * headNum_ * headDim_; + srcOffset = blockSizeIndex * blockSizePerTime_ * headNumSplited_ * headDim_; + dstOffset = blockSizeIndex * blockSizePerTime_ * headNum_ * headDim_; + } + + CopyIn(kCacheGm_, offsetBlock + srcOffset, repeatParams); + CopyOut(kCacheGm_, offsetBlock + dstOffset); + + CopyIn(vCacheGm_, offsetBlock + srcOffset, repeatParams); + CopyOut(vCacheGm_, offsetBlock + dstOffset); + } + + } + + if (needSync_) { + AscendC::CrossCoreSetFlag<0x0, PIPE_MTE2>(0x8); + AscendC::CrossCoreWaitFlag(0x8); + + AscendC::CrossCoreSetFlag<0x0, PIPE_MTE2>(0x8); + AscendC::CrossCoreWaitFlag(0x8); + } + + } +}; \ No newline at end of file diff --git a/csrc/ascend/moe/transpose_kv_cache_by_block/op_kernel/transpose_kv_cache_by_block.cpp b/csrc/ascend/moe/transpose_kv_cache_by_block/op_kernel/transpose_kv_cache_by_block.cpp new file mode 100644 index 000000000..39ced7449 --- /dev/null +++ b/csrc/ascend/moe/transpose_kv_cache_by_block/op_kernel/transpose_kv_cache_by_block.cpp @@ -0,0 +1,36 @@ +#include "kernel_operator.h" +#include "full_load.h" +#include "general.h" + + +extern "C" __global__ __aicore__ void transpose_kv_cache_by_block(GM_ADDR KCache, GM_ADDR VCache, GM_ADDR blockIDs, GM_ADDR workspace, GM_ADDR tiling) { + GET_TILING_DATA(tiling_data, tiling); + TPipe tPipe; + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0); + if (TILING_KEY_IS(0)) { + // full load not db + TransposeKvCacheByBlockKernelFullLoad kernel; + kernel.Init(KCache, VCache, blockIDs, &tiling_data, &tPipe); + kernel.Process(); + } else if (TILING_KEY_IS(1)) { + // db \ align split blockSize + TransposeKvCacheByBlockKernelGeneral kernel; + kernel.Init(KCache, VCache, blockIDs, &tiling_data, &tPipe); + kernel.Process(); + } else if (TILING_KEY_IS(2)) { + // not db \ align split blockSize + TransposeKvCacheByBlockKernelGeneral kernel; + kernel.Init(KCache, VCache, blockIDs, &tiling_data, &tPipe); + kernel.Process(); + } else if (TILING_KEY_IS(3)) { + // db \ unalign split blockSize + TransposeKvCacheByBlockKernelGeneral kernel; + kernel.Init(KCache, VCache, blockIDs, &tiling_data, &tPipe); + kernel.Process(); + } else if (TILING_KEY_IS(4)) { + // not db \ unalign split blockSize + TransposeKvCacheByBlockKernelGeneral kernel; + kernel.Init(KCache, VCache, blockIDs, &tiling_data, &tPipe); + kernel.Process(); + } +} \ No newline at end of file diff --git a/csrc/ascend/ops.h b/csrc/ascend/ops.h new file mode 100644 index 000000000..f65993df9 --- /dev/null +++ b/csrc/ascend/ops.h @@ -0,0 +1,164 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include +#include "kernels/types.h" +#include "torch_npu/csrc/aten/common/from_blob.h" + +namespace vllm_fl { + extern void get_masked_input_and_mask_impl( + void* stream, + void* input, + void* masked_input, + void* mask_out, + const int64_t org_vocab_start_index, + const int64_t org_vocab_end_index, + const int64_t num_org_vocab_padding, + const int64_t added_vocab_start_index, + const int64_t added_vocab_end_index, + const int64_t size, + const uint32_t loop_cnt, + const uint32_t aiv_num); + + torch::Tensor weak_ref_tensor(torch::Tensor& tensor) { + if (!tensor.is_privateuseone()) { + throw std::runtime_error("Tensor must be on NPU device"); + } + // Get the raw data pointer + void* data_ptr = tensor.data_ptr(); + // Get tensor sizes and strides + std::vector sizes = tensor.sizes().vec(); + std::vector strides = tensor.strides().vec(); + // Get tensor options (dtype, device) + auto options = tensor.options(); + // Create a new tensor from the raw data pointer + auto new_tensor = at_npu::native::from_blob(data_ptr, sizes, strides, options); + return new_tensor; + } + + extern void bgmv_shrink_impl( + AscendType type, + void *stream, + void *x, + void *weight, + void *indices, + uint32_t indicesSize, + void *y, + uint32_t batch_size, + uint32_t num_tokens_per_core, + uint32_t input_hidden_dim, + uint32_t lora_rank, + float scale); + + extern void bgmv_expand_impl( + AscendType type, + void *stream, + void *x, + void *weight, + void *indices, + uint32_t indicesSize, + void *y, + void *y_out, + uint32_t batch_size, + uint32_t num_tokens_per_core, + uint32_t lora_rank, + uint32_t output_hidden_dim, + uint32_t slice_offset, + uint32_t output_full_dim); + + extern void sgmv_shrink_impl( + AscendType type, + void *stream, + void *x, + void *weight, + void *loraIndices, + uint32_t loraIndicesSize, + void *seqLen, + uint32_t seqLenSize, + void *y, + uint32_t batch_size, + uint32_t num_tokens_per_core, + uint32_t input_hidden_dim, + uint32_t lora_rank, + float scale); + + extern void sgmv_expand_impl( + AscendType type, + void *stream, + void *x, + void *weight, + void *loraIndices, + uint32_t loraIndicesSize, + void *seqLen, + uint32_t seqLenSize, + void *y, + void *y_out, + uint32_t batch_size, + uint32_t num_tokens_per_core, + uint32_t lora_rank, + uint32_t output_hidden_dim, + uint32_t slice_offset, + uint32_t output_full_dim); + + extern void mla_preprocess_impl( + void* stream, + void* hidden_state, + void* quant_scale1, + void* quant_offset1, + void* wdqkv, + void* bias1, + void* gamma2, + void* beta2, + void* quant_scale2, + void* quant_offset2, + void* gamma3, + void* sin1, + void* cos1, + void* sin2, + void* cos2, + void* keycache, + void* slot_mapping, + void* wuq, + void* bias2, + void* wuk, + void* descale1, + void* descale2, + void* ctkv_scale, + void* qnope_scale, + void* q, + void* keycache_out, + void* q2, + void* keycache_out2, + void* inner_out, + void* workspace, + void* tiling, + const uint32_t block_dim + ); + + extern void batch_matmul_transpose_impl( + void* stream, + void* gm_a, + void* gm_b, + void* gm_c, + void* gm_tiling_data, + const uint32_t block_dim + ); +} diff --git a/csrc/ascend/patch_cann_extract_host_stub.py b/csrc/ascend/patch_cann_extract_host_stub.py new file mode 100644 index 000000000..14bd0e7d0 --- /dev/null +++ b/csrc/ascend/patch_cann_extract_host_stub.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 BAAI. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-FL project +# +# Small helper to work around a CANN 8.5.x bug in extract_host_stub.py: +# compile_commands.json may contain object paths with a leading "./" while +# the object paths passed to extract_host_stub.py are normalized. The lookup +# of enable_ascendc_time_stamp then raises KeyError. +# +# This script is invoked by setup.py immediately before building AscendC +# kernels and is asked to restore the original file afterwards. + +import argparse +import os +import shutil +import sys +from pathlib import Path + + +PATCH_MARKER = "# vllm-fl: normalize path before enable_ascendc_time_stamp lookup" + + +def _get_ascend_home() -> str: + ascend_home = os.environ.get("ASCEND_HOME_PATH") + if ascend_home and os.path.isdir(ascend_home): + return ascend_home + default = "/usr/local/Ascend/ascend-toolkit/latest" + if os.path.isdir(default): + return default + raise RuntimeError( + "ASCEND_HOME_PATH is not set and the default CANN path does not exist." + ) + + +def find_extract_host_stub() -> Path: + """Locate the CANN extract_host_stub.py that matches the current env.""" + ascend_home = _get_ascend_home() + # Architecture names used by CANN toolkit installations. + arch_dirs = ["tools", "compiler", "aarch64-linux", "x86_64-linux"] + for arch in arch_dirs: + candidate = Path(ascend_home) / arch / "tikcpp" / "ascendc_kernel_cmake" / "legacy_modules" / "util" / "extract_host_stub.py" + if candidate.is_file(): + return candidate.resolve() + # Some installations use ascendc_devkit layout. + candidate2 = Path(ascend_home) / arch / "ascendc_devkit" / "tikcpp" / "samples" / "cmake" / "util" / "extract_host_stub.py" + if candidate2.is_file(): + return candidate2.resolve() + raise RuntimeError( + f"Could not find extract_host_stub.py under ASCEND_HOME_PATH={ascend_home}. " + "Please check your CANN installation." + ) + + +def _original_lookup_block_get() -> str: + """CANN 8.5.0 style: uses dict.get(...).""" + return """ for func_group in func_sign_groups: + enbale_flag = enable_ascendc_time_stamp.get(func_group.filepath) + if enbale_flag is True:""" + + +def _original_lookup_block_subscript() -> str: + """CANN 8.5.1 style: uses direct dict subscription.""" + return """ for func_group in func_sign_groups: + enbale_flag = enable_ascendc_time_stamp[func_group.filepath] + if enbale_flag is True:""" + + +def _patched_lookup_block_get() -> str: + """Replacement for the .get(...) style block.""" + return """ for func_group in func_sign_groups: + {marker} + enbale_flag = enable_ascendc_time_stamp.get(func_group.filepath) + if enbale_flag is None: + enbale_flag = enable_ascendc_time_stamp.get(os.path.normpath(func_group.filepath)) + if enbale_flag is None: + enbale_flag = enable_ascendc_time_stamp.get(os.path.abspath(func_group.filepath)) + if enbale_flag is True:""".format(marker=PATCH_MARKER) + + +def _patched_lookup_block_subscript() -> str: + """Replacement for the direct subscription style block.""" + return """ for func_group in func_sign_groups: + {marker} + _filepath = func_group.filepath + if _filepath in enable_ascendc_time_stamp: + enbale_flag = enable_ascendc_time_stamp[_filepath] + elif os.path.normpath(_filepath) in enable_ascendc_time_stamp: + enbale_flag = enable_ascendc_time_stamp[os.path.normpath(_filepath)] + elif os.path.abspath(_filepath) in enable_ascendc_time_stamp: + enbale_flag = enable_ascendc_time_stamp[os.path.abspath(_filepath)] + else: + enbale_flag = False + if enbale_flag is True:""".format(marker=PATCH_MARKER) + + +def apply_patch(backup_path: Path | None = None) -> Path: + """Apply the workaround patch and return the path to the backup file.""" + target = find_extract_host_stub() + original_text = target.read_text(encoding="utf-8") + + if PATCH_MARKER in original_text: + print(f"extract_host_stub.py already patched ({target}); skipping apply.") + # No backup needed because we did not change anything. + return Path("__already_patched__") + + # CANN 9.0+ stores keys as absolute_obj, so the 8.5.x path-normalization + # workaround is no longer required. + if "enable_ascendc_time_stamp[absolute_obj]" in original_text: + print(f"extract_host_stub.py uses absolute_obj keys ({target}); skipping apply.") + return Path("__already_patched__") + + old_get = _original_lookup_block_get() + old_sub = _original_lookup_block_subscript() + if old_get in original_text: + old_block, new_block = old_get, _patched_lookup_block_get() + elif old_sub in original_text: + old_block, new_block = old_sub, _patched_lookup_block_subscript() + else: + raise RuntimeError( + f"Could not find the expected lookup block in {target}. " + "The CANN version may have changed; please review the patch script." + ) + + if backup_path is None: + backup_path = target.with_suffix(target.suffix + ".vllm_fl_bak") + else: + backup_path = Path(backup_path) + + shutil.copy2(target, backup_path) + + patched_text = original_text.replace(old_block, new_block, 1) + target.write_text(patched_text, encoding="utf-8") + print(f"Patched {target}") + print(f"Backup saved to {backup_path}") + return backup_path + + +def restore_patch(backup_path: Path | None = None) -> None: + """Restore the original file from backup and remove the backup.""" + target = find_extract_host_stub() + + if backup_path is None: + backup_path = target.with_suffix(target.suffix + ".vllm_fl_bak") + else: + backup_path = Path(backup_path) + + if not backup_path.is_file(): + if PATCH_MARKER in target.read_text(encoding="utf-8"): + print( + f"Warning: backup {backup_path} missing, cannot restore {target}", + file=sys.stderr, + ) + return + + shutil.copy2(backup_path, target) + backup_path.unlink() + print(f"Restored {target} and removed backup {backup_path}") + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Apply/restore a small workaround patch for CANN extract_host_stub.py" + ) + parser.add_argument( + "action", + choices=["apply", "restore"], + help="Apply the patch or restore the original file.", + ) + parser.add_argument( + "--backup", + type=str, + default=None, + help="Path to use for the backup file. Defaults to a file next to the target.", + ) + args = parser.parse_args() + + backup_path = Path(args.backup) if args.backup else None + + if args.action == "apply": + apply_patch(backup_path) + else: + restore_patch(backup_path) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/csrc/ascend/pto_chunk_gdn/CMakeLists.txt b/csrc/ascend/pto_chunk_gdn/CMakeLists.txt new file mode 100644 index 000000000..1b9c39afe --- /dev/null +++ b/csrc/ascend/pto_chunk_gdn/CMakeLists.txt @@ -0,0 +1,148 @@ +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# ---------------------------------------------------------------------------- +# pto_chunk_gdn — pre-compile the fused PTO GDN megakernel via Bisheng +# ---------------------------------------------------------------------------- +# +# This target pre-compiles mega_kernel.cpp for a set of common model +# configurations (H, Hg, D, C=128) used by Qwen3.5/3.6 models. +# The compiled .so files are installed beside the Python module and loaded +# directly at runtime, avoiding JIT compilation on the first request. +# +# Usage: +# cmake -S csrc -B csrc/build -DBUILD_PTO_CHUNK_GDN=ON +# cmake --build csrc/build --target pto_chunk_gdn_kernels -j8 +# +# Or via pip: +# BUILD_PTO_CHUNK_GDN=ON pip install -e . +# ---------------------------------------------------------------------------- + +option(BUILD_PTO_CHUNK_GDN "Pre-compile PTO GDN megakernel via Bisheng" OFF) + +if(NOT BUILD_PTO_CHUNK_GDN) + return() +endif() + +# ── Locate bisheng ───────────────────────────────────────────────────────────── +find_program(BISHENG_EXECUTABLE bisheng + HINTS ENV ASCEND_HOME_PATH + PATH_SUFFIXES bin + DOC "Bisheng C++ compiler for Ascend AI-Core kernels" +) +if(NOT BISHENG_EXECUTABLE) + message(WARNING "bisheng not found — PTO GDN megakernel will not be pre-compiled." + "Set BISHENG_EXECUTABLE or ensure bisheng is in PATH.") + return() +endif() +message(STATUS "Found bisheng: ${BISHENG_EXECUTABLE}") + +# ── Paths ────────────────────────────────────────────────────────────────────── +set(PTO_SRC_DIR "${CMAKE_CURRENT_SOURCE_DIR}") +set(PTO_INC_DIR "${PTO_SRC_DIR}/include") +set(PTO_ISA_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../third_party/pto-isa") +if(DEFINED ENV{ASCEND_DRIVER_PATH}) + set(DRIVER_INC "$ENV{ASCEND_DRIVER_PATH}") +else() + set(DRIVER_INC "/usr/local/Ascend/driver/kernel/inc") +endif() + +if(DEFINED ENV{ASCEND_HOME_PATH}) + set(ASCEND_HOME "$ENV{ASCEND_HOME_PATH}") +elseif(DEFINED ENV{ASCEND_TOOLKIT_HOME}) + set(ASCEND_HOME "$ENV{ASCEND_TOOLKIT_HOME}") +else() + message(WARNING "ASCEND_HOME_PATH not set — PTO kernel may not compile correctly.") + set(ASCEND_HOME "") +endif() + +if(NOT DEFINED Python_SITELIB) + execute_process( + COMMAND python3 -c "import sysconfig; print(sysconfig.get_paths()['purelib'])" + OUTPUT_VARIABLE Python_SITELIB OUTPUT_STRIP_TRAILING_WHITESPACE + ) +endif() +set(PTO_COMPILED_INSTALL_DIR + "${Python_SITELIB}/vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib" +) + +# ── Common compiler flags ────────────────────────────────────────────────────── +set(PTO_COMMON_FLAGS + -fPIC -shared -xcce -DMEMORY_BASE -O2 -std=gnu++17 + --cce-aicore-arch=dav-c220 + -mllvm -cce-aicore-stack-size=0x8000 + -mllvm -cce-aicore-function-stack-size=0x8000 + -mllvm -cce-aicore-record-overflow=true + -mllvm -cce-aicore-dcci-insert-for-scalar=false + -Wno-macro-redefined -Wno-ignored-attributes + -I${PTO_INC_DIR} + -I${PTO_ISA_DIR}/include + -I${ASCEND_HOME}/include + -I${ASCEND_HOME}/pkg_inc + -I${ASCEND_HOME}/pkg_inc/runtime + -I${ASCEND_HOME}/pkg_inc/profiling +) +if(EXISTS "${DRIVER_INC}") + list(APPEND PTO_COMMON_FLAGS -I${DRIVER_INC}) +endif() + +# ── Model configurations to pre-compile ─────────────────────────────────────── +# Format: "H,Hg,D,C" (H=value heads, Hg=Q/K heads, D=head dim, C=chunk size) +# Qwen3.5-0.8B / 1.5B / 3B: H=16, Hg=8, D=128 +# Qwen3.5-7B / 14B: H=16, Hg=16, D=128 +# Qwen3.6-27B: H=32, Hg=16, D=128 +# Qwen3.6-35B-A3B-MoE: H=32, Hg=16, D=128 +# Common extras for other models: +set(PTO_MODEL_CONFIGS + "16,8,128,128" + "16,16,128,128" + "32,16,128,128" + "48,16,128,128" + "64,16,128,128" + "64,64,128,128" +) + +set(_pto_outputs "") +file(MAKE_DIRECTORY "${PTO_COMPILED_INSTALL_DIR}") + +set(_mega_src "${PTO_SRC_DIR}/mega_kernel.cpp") + +foreach(cfg ${PTO_MODEL_CONFIGS}) + string(REPLACE "," ";" _cfg_list "${cfg}") + list(GET _cfg_list 0 H) + list(GET _cfg_list 1 Hg) + list(GET _cfg_list 2 D) + list(GET _cfg_list 3 C) + + set(_def_flags -DGDN_H=${H} -DGDN_HG=${Hg} -DGDN_D=${D} -DGDN_C=${C}) + set(_out "${PTO_COMPILED_INSTALL_DIR}/mega_kernel_H${H}_Hg${Hg}_D${D}_C${C}.so") + + add_custom_command( + OUTPUT "${_out}" + COMMAND "${BISHENG_EXECUTABLE}" + ${PTO_COMMON_FLAGS} ${_def_flags} + "${_mega_src}" -o "${_out}" + DEPENDS "${_mega_src}" + COMMENT "PTO: compiling mega_kernel H=${H} Hg=${Hg} D=${D} C=${C}" + VERBATIM + ) + list(APPEND _pto_outputs "${_out}") +endforeach() + +# ── Top-level target ─────────────────────────────────────────────────────────── +add_custom_target(pto_chunk_gdn_kernels ALL + DEPENDS ${_pto_outputs} + COMMENT "Building all PTO GDN megakernel Bisheng binaries" +) diff --git a/csrc/ascend/pto_chunk_gdn/chunk_cumsum.cpp b/csrc/ascend/pto_chunk_gdn/chunk_cumsum.cpp new file mode 100644 index 000000000..126434dbb --- /dev/null +++ b/csrc/ascend/pto_chunk_gdn/chunk_cumsum.cpp @@ -0,0 +1,426 @@ +// ============================================================================ +// chunk_cumsum_kernel.cpp — Prefix sum of gate values G along time dimension +// +// Mathematical operation (per chunk of C tokens, independently per head h): +// g_sum[t, h] = Σ_{i=0}^{t} g[i, h] for t = 0 .. valid-1 +// +// Input: g [total_tokens, H] float, BSND layout — raw gate values +// Output: g_sum [total_tokens, H] float — cumulative sums +// +// The prefix sum enables downstream kernels to compute exponential decay +// coefficients: exp(g_sum[i] - g_sum[j]) gives the cumulative gate +// from token j to token i within a chunk. +// +// Architecture: Vec-only kernel (no Cube/GEMM). Single Vec sub-block. +// Pipeline: MTE2(load) → Vec(compute) → MTE3(store), serialized per chunk. +// +// NPU memory hierarchy used: +// GM (Global Memory) → UB (Unified Buffer, on-chip SRAM, Vec-accessible) +// +// ─── PTO / NPU Primer for This Kernel ────────────────────────────────────── +// +// AI Core: The basic processing unit of an NPU, analogous to a Streaming +// Multiprocessor (SM) on a GPU. A single chip has many AI cores, and each +// core runs the same kernel code on different data (SPMD model). +// +// Memory hierarchy (outer → inner): +// GM (Global Memory) — Off-chip DRAM, like GPU HBM. Large (several GB) +// but high latency. All AI cores share GM. +// UB (Unified Buffer) — On-chip SRAM, ~256 KB per AI core. Like GPU +// shared memory. Very fast, but small. The Vec engine can only operate +// on data that lives in UB, so every tensor must be DMA'd in first. +// +// Hardware pipes (execute in parallel, like independent GPU warps): +// Vec — SIMD vector processor. Performs element-wise math (add, mul, etc.) +// on data already in UB. Think of it as a wide SIMD ALU. +// MTE2 — DMA engine for loads: copies data from GM → UB. +// MTE3 — DMA engine for stores: copies data from UB → GM. +// Cube — Matrix engine for GEMMs (not used in this kernel). +// +// Synchronization (set_flag / wait_flag): +// Because Vec, MTE2, and MTE3 run in parallel on separate hardware, you +// must explicitly synchronize them to ensure data is ready: +// set_flag(SRC_PIPE, DST_PIPE, event): SRC signals that it is done. +// wait_flag(SRC_PIPE, DST_PIPE, event): DST blocks until the signal. +// Example: After MTE2 loads data into UB, Vec must wait_flag before reading +// it. This is like a fine-grained torch.cuda.synchronize() between pipes. +// Events (EVENT_ID0 .. EVENT_ID7) are semaphore indices. +// +// ============================================================================ + +#include +#include "acl/acl.h" +#include +using namespace pto; + +// GDN_H, GDN_C: Compile-time constants injected by the build system. +// GDN_H = number of attention heads (e.g., 16) +// GDN_C = chunk size in tokens (e.g., 128) +// Using compile-time constants allows the compiler to optimize tile sizes, +// unroll loops, and compute UB addresses at compile time. +#ifndef GDN_H +#define GDN_H 16 +#endif + +#ifndef GDN_C +#define GDN_C 128 +#endif + +// ── PTO type aliases (device-only, guarded by __CCE_AICORE__) ─────────────── +// UB tile in row-major (ND) layout, used by Vec engine. +// T=dtype, R×C=static shape, RV×CV=valid region, P=pad value for TLOAD. +// +// Think of UbND as: torch.empty((R, C), dtype=T) allocated in on-chip SRAM (UB). +// - TileType::Vec = this tile lives in UB, operated on by the Vec (SIMD) engine +// - BLayout::RowMajor = row-major storage, like C arrays or numpy default +// - RV, CV = "valid" region within the R×C buffer (for handling partial/tail chunks) +// - PadValue = what to fill outside the valid region during TLOAD (Zero or Null) +// - 512 = alignment in bytes (hardware requirement for efficient DMA) +#ifdef __CCE_AICORE__ +template +using UbND = pto::Tile; +#endif + +template +AICORE void cumsum_kernel( + __gm__ float *g_ptr, __gm__ float *g_sum_ptr, + __gm__ int32_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, + uint64_t ffts_addr) +{ + // get_block_idx(): Returns this AI core's index (0..block_num-1). + // Like blockIdx.x in CUDA — identifies which core this code runs on. + // get_block_num(): Total number of AI cores launched (like gridDim.x in CUDA). + // get_subblockid(): Returns 0 or 1 — selects which Vec sub-block within the core. + // Each AI core has 2 Vec sub-blocks that can run in parallel. + auto cid = get_block_idx(); + auto block_num = get_block_num(); + auto vid = get_subblockid(); + // set_ffts_base_addr(ffts_addr): Configure the base address for FFTS + // (Fast Fine-grained Task Synchronization) — the cross-core signaling mechanism. + // Required before any cross-core sync (ffts_cross_core_sync / wait_flag_dev). + set_ffts_base_addr(ffts_addr); + +// #if defined(__DAV_C220_VEC__): This block only compiles for the Vec core pass. +// The bisheng compiler makes 3 passes over the same source file: +// Pass 1: __DAV_C220_VEC__ defined → compiles Vec (SIMD) code +// Pass 2: __DAV_C220_CUBE__ defined → compiles Cube (matrix) code +// Pass 3: neither defined → compiles host (CPU) launcher code +// Using these guards lets us put Vec, Cube, and host code in one file. +#if defined(__DAV_C220_VEC__) + if (vid != 0) return; + + // set_mask_norm(): Reset Vec mask to normal mode (all lanes active). + // set_vector_mask(-1, -1): Enable all SIMD lanes (128 lanes for fp32). + // The -1 sets all 64 bits to 1 in each of the two 64-bit mask registers. + // This is like setting torch's computation to operate on all elements. + set_mask_norm(); + set_vector_mask(-1, -1); + + // HeadTileCols: NumHeads rounded up to 8-element alignment (32B for float) + // HTC = NumHeads rounded up to nearest multiple of 8. + // Why? The Vec engine processes data in 32-byte granularity. + // For float (4 bytes), that's 8 elements per SIMD "word". + // Rounding up ensures every row is a whole number of SIMD words, + // avoiding partial-lane issues. The extra columns are zero-padded. + // Example: NumHeads=16 → HTC=16 (already aligned), NumHeads=13 → HTC=16. + constexpr int32_t HTC = ((NumHeads + 7) / 8) * 8; + constexpr int32_t BlockBytes = ChunkSize * HTC * + static_cast(sizeof(float)); + constexpr int32_t RowBytes = HTC * static_cast(sizeof(float)); + + // ── UB memory layout ────────────────────────────────────────────────── + // [0 .. BlockBytes) = g input (ChunkSize × HTC floats) + // [BlockBytes .. 2*BlockBytes) = g_sum output + // [2*BlockBytes .. 2*BlockBytes+RowBytes) = row accumulator (1 × HTC) + constexpr int32_t GUbAddr = 0; + constexpr int32_t SUbAddr = BlockBytes; + constexpr int32_t AccUbAddr = BlockBytes * 2; + + // GlobalTensor types for g/g_sum in [total_tokens, NumHeads] layout. + // 5D shape with last two dims dynamic; stride encodes row pitch. + // + // GlobalTensor is a "view" into GM (Global Memory), like torch.as_strided(). + // GlobalTensor(base_ptr, shape) + // Shape<1,1,1,DYNAMIC,DYNAMIC> = 5D shape where first 3 dims are 1 (unused), + // last 2 dims are set at runtime (valid rows × NumHeads). + // Stride<1,1,1,NumHeads,1> = stride between elements. The 4th stride = NumHeads + // means consecutive rows in GM are NumHeads elements apart (BSND layout: + // token[t] at offset t*NumHeads, head[h] at offset h within that token). + // This is equivalent to: + // g_gm = torch.as_strided(g_ptr, size=[valid, NumHeads], stride=[NumHeads, 1]) + using GmShape = Shape<1, 1, 1, DYNAMIC, DYNAMIC>; + using GmStride = Stride<1, 1, 1, NumHeads, 1>; + using GmFloat = GlobalTensor; + + // Pre-assign row accumulator at fixed UB address + // TASSIGN(tile, address): Binds a tile descriptor to a fixed byte address in UB. + // Think of it as: tile = ub_memory[address:address+sizeof(tile)] + // This does NOT allocate or move data — it just tells the hardware where the tile lives. + // We manually manage UB memory layout (like a memory pool) via compile-time addresses. + UbND acc_ub; + TASSIGN(acc_ub, AccUbAddr); + + int64_t num_seqs = batch_size; + + // ── Fixed-length sequence path (cu_seqlens == nullptr) ──────────────── + if (cu_seqlens == nullptr) { + int64_t chunks_per_seq = (seq_len + ChunkSize - 1) / ChunkSize; + int64_t total_chunks = num_seqs * chunks_per_seq; + + // Work distribution: Each AI core processes chunks in a round-robin pattern. + // Core `cid` handles chunks cid, cid+block_num, cid+2*block_num, ... + // This is the NPU equivalent of CUDA's grid-stride loop: + // for (int i = blockIdx.x; i < total; i += gridDim.x) + for (int64_t gi = static_cast(cid); gi < total_chunks; + gi += static_cast(block_num)) { + int64_t seq_idx = gi / chunks_per_seq; + int64_t local_chunk = gi % chunks_per_seq; + int64_t bos = seq_idx * seq_len; + int64_t chunk_start = bos + local_chunk * ChunkSize; + int64_t remaining = seq_len - local_chunk * ChunkSize; + int32_t valid = static_cast( + remaining < ChunkSize ? remaining : ChunkSize); + + // ── DMA: load g[chunk_start .. +valid] from GM → UB (MTE2 pipe) ── + // Constructs a GlobalTensor view over the g array, loads into UB, + // then zero-pads the tail region (rows beyond `valid`, cols beyond + // NumHeads up to the 8-aligned HTC) so downstream Vec ops see zeros. + { + GmShape gs; gs.shape[3] = valid; gs.shape[4] = NumHeads; + GmFloat g_gm(g_ptr + chunk_start * NumHeads, gs); + UbND + g_load(valid, NumHeads); + TASSIGN(g_load, GUbAddr); + // TLOAD(ub_tile, gm_tensor): DMA transfer from GM → UB. + // Equivalent to: ub_tile[:valid, :NumHeads] = gm_tensor[:valid, :NumHeads] + // This is an ASYNC operation on the MTE2 pipe — the CPU/Vec engine can do + // other work while DMA is in progress. You must call set_flag/wait_flag + // before reading the loaded data. + TLOAD(g_load, g_gm); + if (valid != ChunkSize || NumHeads != HTC) { + UbND g_pad; + TASSIGN(g_pad, GUbAddr); + // TFILLPAD_INPLACE(full_tile, partial_tile): Zero-fills the region outside + // the valid area of partial_tile. + // Equivalent to: + // full_tile[valid:ChunkSize, :] = 0 # zero rows beyond valid + // full_tile[:, NumHeads:HTC] = 0 # zero cols beyond NumHeads (alignment padding) + // This ensures downstream Vec operations see clean zeros in padded regions. + TFILLPAD_INPLACE(g_pad, g_load); + } + } + // ── Synchronization: MTE2 → Vec ──────────────────────────────────── + // set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0): Signal from MTE2 (DMA load + // engine) to Vec (SIMD engine) that the DMA transfer is complete. + // wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0): Vec waits here until MTE2 + // has set the flag. After this, UB data from TLOAD is safe to read. + // Think of it as: torch.cuda.synchronize() but fine-grained per pipe. + // EVENT_ID0 is a semaphore index (0-7 available). + // MTE2 → Vec sync: wait for DMA load to finish before Vec reads UB + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // ── Vec compute: prefix sum over rows (all H heads in parallel) ─── + // Row 0: acc[h] = g[0,h]; g_sum[0,h] = acc[h] + UbND g_row_0; + TASSIGN(g_row_0, GUbAddr); + // TMOV(dst, src): Element-wise copy, like dst = src.clone() in UB. + TMOV(acc_ub, g_row_0); + // pipe_barrier(PIPE_V): Ensures all pending Vec (SIMD) operations complete + // before the next Vec instruction begins. Needed because Vec ops are pipelined + // and may not finish in order. Think of it as a local __syncthreads() for the + // Vec engine only. Much lighter than set_flag/wait_flag (which sync across + // different hardware units). + pipe_barrier(PIPE_V); + + UbND s_row_0; + TASSIGN(s_row_0, SUbAddr); + TMOV(s_row_0, acc_ub); + pipe_barrier(PIPE_V); + + // Rows 1..valid-1: acc[h] += g[i,h]; g_sum[i,h] = acc[h] + for (int32_t i = 1; i < valid; ++i) { + UbND g_row_i; + TASSIGN(g_row_i, GUbAddr + i * RowBytes); + // TADD(dst, a, b): Element-wise add, like dst = a + b. All in UB. + // Operates on all HTC elements in parallel (SIMD). + TADD(acc_ub, acc_ub, g_row_i); + pipe_barrier(PIPE_V); + + UbND s_row_i; + TASSIGN(s_row_i, SUbAddr + i * RowBytes); + TMOV(s_row_i, acc_ub); + pipe_barrier(PIPE_V); + } + + // Zero-fill rows beyond valid (tail padding for downstream kernels) + // TEXPANDS(tile, scalar): Fill entire tile with a scalar value. + // Equivalent to: tile[:] = scalar (like torch.full_like(tile, scalar)) + TEXPANDS(acc_ub, 0.0f); + pipe_barrier(PIPE_V); + for (int32_t i = valid; i < ChunkSize; ++i) { + UbND s_row_i; + TASSIGN(s_row_i, SUbAddr + i * RowBytes); + TMOV(s_row_i, acc_ub); + pipe_barrier(PIPE_V); + } + + // ── DMA: store g_sum from UB → GM (MTE3 pipe) ──────────────────── + // ── Synchronization: Vec → MTE3 ─────────────────────────────────── + // Vec signals MTE3 that computation is done and UB data is ready to store. + // MTE3 (DMA store engine) waits for this before reading UB for TSTORE. + // Without this sync, MTE3 might read stale/partial data from UB. + // Vec → MTE3 sync: ensure Vec writes to UB are visible before DMA + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + { + GmShape ss; ss.shape[3] = valid; ss.shape[4] = NumHeads; + GmFloat gs_gm(g_sum_ptr + chunk_start * NumHeads, ss); + UbND + s_store(valid, NumHeads); + TASSIGN(s_store, SUbAddr); + // TSTORE(gm_tensor, ub_tile): DMA transfer from UB → GM. + // Equivalent to: gm_tensor[:valid, :NumHeads] = ub_tile[:valid, :NumHeads] + // Async on MTE3 pipe. Must sync (Vec→MTE3) before calling, and sync + // (MTE3→Vec) after if reusing the same UB region. + TSTORE(gs_gm, s_store); + } + // ── Synchronization: MTE3 → Vec ─────────────────────────────────── + // MTE3 signals Vec that the DMA store is complete and UB can be reused. + // Vec waits before starting the next iteration's TLOAD into the same UB region. + // Without this, the next TLOAD could overwrite data still being stored. + // MTE3 → Vec sync: wait for DMA store before reusing UB next iter + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + } + } + // ── Variable-length sequence path (cu_seqlens != nullptr) ───────────── + else { + int64_t gi = 0; + for (int64_t si = 0; si < num_seqs; ++si) { + int64_t bos = static_cast(cu_seqlens[si]); + int64_t eos = static_cast(cu_seqlens[si + 1]); + int64_t slen = eos - bos; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t c = 0; c < nc; ++c) { + if (gi % static_cast(block_num) == + static_cast(cid)) { + int64_t chunk_start = bos + c * ChunkSize; + int64_t remaining = slen - c * ChunkSize; + int32_t valid = static_cast( + remaining < ChunkSize ? remaining : ChunkSize); + + // Load g chunk from GM → UB, zero-padded + { + GmShape gs; gs.shape[3] = valid; gs.shape[4] = NumHeads; + GmFloat g_gm(g_ptr + chunk_start * NumHeads, gs); + UbND + g_load(valid, NumHeads); + TASSIGN(g_load, GUbAddr); + TLOAD(g_load, g_gm); + if (valid != ChunkSize || NumHeads != HTC) { + UbND + g_pad; + TASSIGN(g_pad, GUbAddr); + TFILLPAD_INPLACE(g_pad, g_load); + } + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Prefix sum: acc = g[0]; g_sum[0] = acc + UbND g_row_0; + TASSIGN(g_row_0, GUbAddr); + TMOV(acc_ub, g_row_0); + pipe_barrier(PIPE_V); + + UbND s_row_0; + TASSIGN(s_row_0, SUbAddr); + TMOV(s_row_0, acc_ub); + pipe_barrier(PIPE_V); + + // acc += g[i]; g_sum[i] = acc + for (int32_t i = 1; i < valid; ++i) { + UbND g_row_i; + TASSIGN(g_row_i, GUbAddr + i * RowBytes); + TADD(acc_ub, acc_ub, g_row_i); + pipe_barrier(PIPE_V); + + UbND s_row_i; + TASSIGN(s_row_i, SUbAddr + i * RowBytes); + TMOV(s_row_i, acc_ub); + pipe_barrier(PIPE_V); + } + + // Zero-fill padding rows + TEXPANDS(acc_ub, 0.0f); + pipe_barrier(PIPE_V); + for (int32_t i = valid; i < ChunkSize; ++i) { + UbND s_row_i; + TASSIGN(s_row_i, SUbAddr + i * RowBytes); + TMOV(s_row_i, acc_ub); + pipe_barrier(PIPE_V); + } + + // Store g_sum to GM + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + { + GmShape ss; ss.shape[3] = valid; ss.shape[4] = NumHeads; + GmFloat gs_gm(g_sum_ptr + chunk_start * NumHeads, ss); + UbND + s_store(valid, NumHeads); + TASSIGN(s_store, SUbAddr); + TSTORE(gs_gm, s_store); + } + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + } + gi++; + } + } + } +#endif +} + +// ── Device-side kernel entry point ───────────────────────────────── +// extern "C" __global__ AICORE: marks this as an NPU kernel function +// (like __global__ in CUDA). Each AI core runs one instance of this function. +// Parameters are passed as uint8_t* (raw bytes) and reinterpret_cast'd to +// typed pointers — this is the standard NPU kernel calling convention. +extern "C" __global__ AICORE void launch_cumsum( + __gm__ uint8_t *g_ptr, __gm__ uint8_t *g_sum_ptr, + __gm__ uint8_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, + uint64_t ffts_addr) +{ + cumsum_kernel( + reinterpret_cast<__gm__ float *>(g_ptr), + reinterpret_cast<__gm__ float *>(g_sum_ptr), + reinterpret_cast<__gm__ int32_t *>(cu_seqlens), + batch_size, seq_len, ffts_addr); +} + +// ── Host-side launcher (called from Python via ctypes) ──────────── +// call_kernel(): CPU function that launches the NPU kernel. +// block_dim = number of AI cores to use (like CUDA grid size) +// stream = NPU stream for async execution (like CUDA stream) +// rtGetC2cCtrlAddr: gets the FFTS control address for cross-core sync +// <<>>: NPU kernel launch syntax (like CUDA <<<>>>) +extern "C" void call_kernel( + uint32_t block_dim, void *stream, + uint8_t *g_ptr, uint8_t *g_sum_ptr, uint8_t *cu_seqlens, + int64_t batch_size, int64_t seq_len) +{ + uint32_t fftsLen{0}; + uint64_t fftsAddr{0}; + rtGetC2cCtrlAddr(&fftsAddr, &fftsLen); + launch_cumsum<<>>( + g_ptr, g_sum_ptr, cu_seqlens, batch_size, seq_len, fftsAddr); +} diff --git a/csrc/ascend/pto_chunk_gdn/chunk_h.cpp b/csrc/ascend/pto_chunk_gdn/chunk_h.cpp new file mode 100644 index 000000000..3039f1e0d --- /dev/null +++ b/csrc/ascend/pto_chunk_gdn/chunk_h.cpp @@ -0,0 +1,925 @@ +// ============================================================================ +// chunk_h_kernel.cpp — Recurrent hidden state update for GatedDeltaNet +// +// Mathematical recurrence per chunk c: +// S_{c+1} = exp(g_last) * S_c + K^T @ V +// +// where g_last = exp(g[valid-1]) is the chunk's final gate value, S is the +// D×D hidden state, K ∈ ℝ^{C×D}, V ∈ ℝ^{C×D}, and g ∈ ℝ^C is the per-token +// gate. +// +// ── Cube phase (two GEMMs per chunk, sequentially): ────────────────────── +// 1. WS = W @ S project current state through W (wy_fast output) +// W ∈ ℝ^{C×D}, S ∈ ℝ^{D×D} → WS ∈ ℝ^{C×D} +// 2. KV = K^T @ V outer product of keys and values (transpose_A!) +// K stored as D×C, V ∈ ℝ^{C×D} → KV ∈ ℝ^{D×D} +// +// ── Vec phase (two sub-blocks handle upper/lower C/2 rows): ───────────── +// For each chunk: +// 1. Load K, G (pre-transposed), U (from wy_fast) +// 2. Compute coeff[i] = exp(g[i] - g[valid-1]) — time-decay scaling +// Uses TROWEXPAND to broadcast coefficients across D columns +// 3. Scale K: K_scaled[i,:] = K[i,:] * coeff[i] +// 4. Load WS from Cube workspace, compute V_new = U - WS (residual) +// 5. Store V_new and K_scaled to workspace for Cube's next iteration +// 6. Update state: S = exp(g_last) * S + KV (from Cube workspace) +// 7. Store final state FS after last chunk +// +// Cross-core sync: Cube→Vec flags for WS/KV ready, Vec→Cube flags for +// K/S ready. +// +// Inputs: +// K [total_tokens, Hg, D] half — keys (BSND layout; GQA/MQA group heads) +// W [total_tokens, H, D] half — wy_fast output (BSND layout) +// U [total_tokens, H, D] half — values pre-residual (BSND layout) +// G [H, total_tokens] float — pre-transposed cumulative gates +// S [total_chunks, H, D, D] half — per-chunk state snapshots (output) +// V [total_tokens, H, D] half — residual-corrected values (output) +// FS [batch, H, D, D] half — final state per sequence (output) +// workspace [per-core scratch] — Cube↔Vec communication buffer +// +// NPU memory hierarchy: +// GM → L1 (Cube-accessible) → L0A/L0B/L0C (Cube GEMM registers) +// GM → UB (Vec-accessible, on-chip SRAM) +// Cross-core sync via FFTS (Fast Fine-grained Task Synchronization) +// +// ── PTO / NPU Primer ────────────────────────────────────────────────── +// This is the most complex kernel in the GDN suite. It implements the +// recurrent state update, requiring sequential chunk processing (chunks +// within a sequence CANNOT be parallelized — each depends on the previous). +// +// Key PTO APIs (numpy/torch equivalents): +// TLOAD(dst, gm) — dst = gm_data (DMA: GM→L1 or GM→UB) +// TSTORE(gm, src) — gm_data = src (DMA: UB/L0C→GM) +// TASSIGN(tile, addr) — tile = memory[addr] (bind tile to buffer address) +// TCVT(dst, src, mode) — dst = src.float()/.half() +// TMOV(dst, src) — dst = src.clone() +// TADD(d, a, b) — d = a + b +// TSUB(d, a, b) — d = a - b +// TMUL(d, a, b) — d = a * b +// TMULS(d, s, scalar) — d = s * scalar (scalar multiply) +// TADDS(d, s, scalar) — d = s + scalar (scalar add) +// TEXP(d, s) — d = torch.exp(s) +// TEXPANDS(tile, scalar) — tile[:] = scalar (fill with constant) +// TROWEXPAND(2d, col) — 2d[i,j] = col[i] (broadcast col across row dim) +// TFILLPAD(dst, src) — zero-fill L1 tile padding (for tail chunks) +// TEXTRACT(l0, l1, r, c) — L1 sub-tile → L0A/L0B +// TRESHAPE(zn, nz) — reinterpret layout NZ↔ZN (logical transpose, free) +// TMATMUL(C, A, B) — C = A @ B (Cube GEMM, fp16 inputs → fp32 accum) +// set_flag/wait_flag — pipe sync within same core +// ffts_cross_core_sync — cross-core signal Cube↔Vec +// wait_flag_dev(flag) — wait for cross-core signal +// GetValue(idx) — read a single scalar from a UB tile (slow, use sparingly) +// +// ── Workspace memory layout (shared between Cube and Vec via GM) ────── +// Each AI core has its own workspace region to avoid contention: +// WS_WS [C×D]: Cube writes WS = W @ S here → Vec reads it +// WS_K [D×C]: Vec writes K_scaled here → Cube reads it for KV = K^T @ V +// WS_S [D×D]: Vec writes current state S here → Cube reads it for GEMM 1 +// WS_KV [D×D]: Cube writes KV = K^T @ V here → Vec reads it to update S +// +// Data flow per chunk (think of it as a ping-pong between Cube and Vec): +// Vec: write S₀ to WS_S → signal Cube (flag 3) +// Cube: read S from WS_S, load W → compute WS = W@S → write WS_WS → signal Vec (flag 0) +// Vec: read WS, compute V_new = U - WS, compute K_scaled → write WS_K → signal Cube (flag 1) +// Cube: read K from WS_K, load V → compute KV = K^T@V → write WS_KV → signal Vec (flag 2) +// Vec: read KV, update S = exp(g_last)*S + KV → write S to WS_S → signal Cube (flag 3) +// ... repeat for next chunk ... +// ============================================================================ + +#include +#include +#include "acl/acl.h" +#include +using namespace pto; + +#ifdef __CCE_AICORE__ + +namespace { + +using GmShape2D = pto::Shape<1, 1, 1, pto::DYNAMIC, pto::DYNAMIC>; +using GmStride2D = pto::Stride<1, 1, 1, pto::DYNAMIC, 1>; + +template +using GmTensor2D = pto::GlobalTensor; + +template +using DynMatL1 = pto::Tile; + +template +using DynVecTile = pto::Tile; + +template +using DynAccTile = pto::TileAcc; + +template +using TileMatL1 = pto::Tile; + +template +using TileMatL1ZN = pto::Tile; + +template +using TileMatL0A = pto::Tile; + +template +using TileMatL0B = pto::Tile; + +template +using TileUbDataND = pto::Tile; + +template +using TileUbDataDN = pto::Tile; + +// PTO cheat sheet for the recurrent kernel: +// - `GlobalTensor` is a GM tensor view with explicit runtime shape/stride. +// - `Tile<..., Mat, ...>` lives in L1 and feeds Cube matmul instructions. +// - `Tile<..., Vec, ...>` lives in UB for elementwise vector work. +// - `TileAcc` is a Cube accumulator tile. +// - `TLOAD` / `TSTORE` are DMA copies between GM and on-chip memory. +// - `TROWEXPAND` broadcasts a column vector across the feature dimension. +// - `TFILLPAD(_INPLACE)` zero-pads tail rows so full-tile code can still run. + +template +AICORE PTO_INLINE void +gemm_v0(std::conditional_t, + TileMatL1> &A, + std::conditional_t, + TileMatL1> &B, + pto::TileAcc &C, bool clear) +{ + // Local K-sliced matmul helper: + // C = A @ B + // PTO exposes the L1/L0 staging explicitly, so this stays as a tiny file- + // local helper instead of a shared wrapper. + // + // PyTorch mental model: + // C = 0 + // for k0 in range(0, K, kL0Size): + // C += A[:, k0:k1] @ B[k0:k1, :] + constexpr uint32_t kL0Size = 128; + const uint32_t kL0split = (K + kL0Size - 1) / kL0Size; + + auto war_event_id = (event_t)(((int)EVENT_ID0 + 1) % 8); + set_flag(PIPE_MTE2, PIPE_MTE1, war_event_id); + wait_flag(PIPE_MTE2, PIPE_MTE1, war_event_id); + + for (uint32_t kL0Idx = 0; kL0Idx < kL0split; ++kL0Idx) { + const bool initflag = clear && (kL0Idx == 0); + const bool is_tail_block = (kL0Idx == kL0split - 1); + + if (is_tail_block) { + TileMatL0A l0a; + TileMatL0B l0b; + pto::TASSIGN(l0a, 0x0); + pto::TASSIGN(l0b, 0x0); + + set_flag(PIPE_M, PIPE_MTE1, war_event_id); + wait_flag(PIPE_M, PIPE_MTE1, war_event_id); + + if constexpr (!transpose_A) { + pto::TEXTRACT(l0a, A, 0, kL0Idx * K_tail); + } else { + TileMatL1ZN A_t; + pto::TRESHAPE(A_t, A); + pto::TEXTRACT(l0a, A_t, 0, kL0Idx * K_tail); + } + + if constexpr (!transpose_B) { + pto::TEXTRACT(l0b, B, kL0Idx * K_tail, 0); + } else { + TileMatL1ZN B_t; + pto::TRESHAPE(B_t, B); + pto::TEXTRACT(l0b, B_t, kL0Idx * K_tail, 0); + } + + set_flag(PIPE_MTE1, PIPE_M, war_event_id); + wait_flag(PIPE_MTE1, PIPE_M, war_event_id); + + if (initflag) { + pto::TMATMUL(C, l0a, l0b); + } else { + pto::TMATMUL_ACC(C, C, l0a, l0b); + } + } else { + TileMatL0A l0a; + TileMatL0B l0b; + pto::TASSIGN(l0a, 0x0); + pto::TASSIGN(l0b, 0x0); + + set_flag(PIPE_M, PIPE_MTE1, war_event_id); + wait_flag(PIPE_M, PIPE_MTE1, war_event_id); + + set_flag(PIPE_FIX, PIPE_M, war_event_id); + wait_flag(PIPE_FIX, PIPE_M, war_event_id); + + if constexpr (!transpose_A) { + pto::TEXTRACT(l0a, A, 0, kL0Idx * kL0Size); + } else { + TileMatL1ZN A_t; + pto::TRESHAPE(A_t, A); + pto::TEXTRACT(l0a, A_t, 0, kL0Idx * kL0Size); + } + + if constexpr (!transpose_B) { + pto::TEXTRACT(l0b, B, kL0Idx * kL0Size, 0); + } else { + TileMatL1ZN B_t; + pto::TRESHAPE(B_t, B); + pto::TEXTRACT(l0b, B_t, kL0Idx * kL0Size, 0); + } + + set_flag(PIPE_MTE1, PIPE_M, war_event_id); + wait_flag(PIPE_MTE1, PIPE_M, war_event_id); + + if (initflag) { + pto::TMATMUL(C, l0a, l0b); + } else { + pto::TMATMUL_ACC(C, C, l0a, l0b); + } + + set_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + wait_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + } + } + + set_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + wait_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + + set_flag(PIPE_M, PIPE_FIX, war_event_id); + wait_flag(PIPE_M, PIPE_FIX, war_event_id); +} + +} // namespace + +#endif + +template +AICORE void chunk_h_kernel( + __gm__ half *K_handle, __gm__ half *W_handle, __gm__ half *U_handle, + __gm__ float *G_handle, + __gm__ half *S_handle, __gm__ half *V_handle, __gm__ half *FS_handle, + __gm__ half *workspace_handle, + __gm__ int32_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, int64_t total_tokens, + uint64_t ffts_addr) +{ + // chunk_h advances the recurrent hidden state chunk by chunk: + // ws_i = W_i @ S_i + // v_i_new = U_i - ws_i + // k_i_tilde = exp(g_last - g_i) * K_i + // S_{i+1} = exp(g_last) * S_i + k_i_tilde^T @ v_i_new. + // + // Shapes for one (sequence, head, chunk): + // W_i, U_i, K_i, V_i_new : [valid, D] + // S_i, S_{i+1} : [D, D] + // + // PyTorch / NumPy sketch: + // ws = W_i @ S_i + // v_new = U_i - ws + // decay = exp(g_last - g_i)[:, None] + // k_tilde = decay * K_i + // kv = k_tilde.T @ v_new + // S = exp(g_last) * S + kv + // + // PTO split: + // Cube forms the two matmuls (`W_i @ S_i` and `K_i^T @ V_i_new`). + // Vec does the elementwise gating/decay and carries the running state. + auto cid = get_block_idx(); + auto block_num = get_block_num(); + set_ffts_base_addr(ffts_addr); + + constexpr int32_t D = HiddenSize; + constexpr int32_t C = ChunkSize; + constexpr int32_t H = NumHeads; + constexpr int32_t Hg = NumKeyHeads; + static_assert(Hg > 0 && H % Hg == 0, + "NumHeads must be divisible by NumKeyHeads"); + constexpr int32_t GROUP = H / Hg; + constexpr int32_t HalfC = C / 2; + constexpr int32_t BSND_QKV_STRIDE = H * D; + constexpr int32_t BSND_K_STRIDE = Hg * D; + constexpr int32_t DD = D * D; + + constexpr int32_t WS_WS = 0; + constexpr int32_t WS_K = DD; + constexpr int32_t WS_S = DD * 2; + constexpr int32_t WS_KV = DD * 3; + constexpr int32_t WS_PER_CORE = DD * 4; + + TileMatL1 s_l1; + TASSIGN(s_l1, 0); + TileMatL1 w_l1; + TASSIGN(w_l1, D * D * sizeof(half)); + TileAcc ws_l0; + TASSIGN(ws_l0, 0); + TileMatL1 k_l1; + TASSIGN(k_l1, (DD + C * D) * sizeof(half)); + TileMatL1 v_l1; + TASSIGN(v_l1, (DD + C * D + D * C) * sizeof(half)); + TileAcc kv_l0; + TASSIGN(kv_l0, C * D * sizeof(float)); + + constexpr int32_t G_BLOCK_UB = 0; + // Leading UB scratch: legacy kernels used ``C * NumHeads * sizeof(float)``, which overflows UB when + // ``NumHeads`` is 32/48/64. Keep the same slack as the historical ``GDN_H=16`` build (8192 bytes). + constexpr int32_t ZERO_UB = + ChunkSize * 16 * static_cast(sizeof(float)); + constexpr int32_t S_UB = ZERO_UB + 64 * sizeof(float); + constexpr int32_t K_UB_HALF = S_UB + HalfC * D * sizeof(float); + constexpr int32_t G_UB = K_UB_HALF + HalfC * D * sizeof(half); + constexpr int32_t U_UB_HALF = G_UB + C * sizeof(float); + constexpr int32_t K_UB = U_UB_HALF + HalfC * D * sizeof(half); + constexpr int32_t G_V_UB = K_UB + HalfC * D * sizeof(float); + constexpr int32_t COEFF_UB = G_V_UB + 64 * sizeof(float); + constexpr int32_t U_UB = COEFF_UB + 64 * sizeof(float); + constexpr int32_t WS_UB = U_UB + HalfC * D * sizeof(float); + constexpr int32_t KV_UB = U_UB_HALF; + constexpr int32_t S_UB_HALF = WS_UB + HalfC * D * sizeof(float); + + TileUbDataND zero_ub; + TASSIGN(zero_ub, ZERO_UB); + TileUbDataND s_ub; + TASSIGN(s_ub, S_UB); + TileUbDataND k_ub_half; + TASSIGN(k_ub_half, K_UB_HALF); + TileUbDataND g_ub; + TASSIGN(g_ub, G_UB); + TileUbDataND s_ub_half; + TASSIGN(s_ub_half, S_UB_HALF); + TileUbDataND u_ub_half; + TASSIGN(u_ub_half, U_UB_HALF); + TileUbDataND k_ub; + TASSIGN(k_ub, K_UB); + TileUbDataND g_v_ub; + TASSIGN(g_v_ub, G_V_UB); + TileUbDataND coeff_ub; + TASSIGN(coeff_ub, COEFF_UB); + TileUbDataND u_ub; + TASSIGN(u_ub, U_UB); + TileUbDataND ws_ub; + TASSIGN(ws_ub, WS_UB); + TileUbDataND kv_ub; + TASSIGN(kv_ub, KV_UB); + + auto vid = get_subblockid(); + + int64_t num_seqs = batch_size; + int64_t total_work = num_seqs * H; + +#if defined(__DAV_C220_CUBE__) + for (int64_t wi = 0; wi < (total_work + block_num - 1) / block_num; ++wi) { + int64_t pid = wi * block_num + cid; + if (pid >= total_work) break; + + int64_t head = pid % H; + int64_t seq_idx = pid / H; + + int64_t bos, slen; + int64_t chunk_offset = 0; + if (cu_seqlens != nullptr) { + bos = static_cast(cu_seqlens[seq_idx]); + int64_t eos = static_cast(cu_seqlens[seq_idx + 1]); + slen = eos - bos; + for (int64_t si = 0; si < seq_idx; ++si) { + int64_t sb = static_cast(cu_seqlens[si]); + int64_t se = static_cast(cu_seqlens[si + 1]); + chunk_offset += (se - sb + C - 1) / C; + } + } else { + bos = seq_idx * seq_len; + slen = seq_len; + chunk_offset = seq_idx * ((seq_len + C - 1) / C); + } + int64_t num_chunks = (slen + C - 1) / C; + int64_t ws_base = static_cast(cid) * WS_PER_CORE; + // One per-core scratch region stores: + // WS_WS : ws = W_i @ S_i + // WS_K : k_tilde + // WS_S : running state S_i + // WS_KV : k_tilde^T @ v_i_new + + for (int32_t ci = 0; ci < num_chunks; ++ci) { + wait_flag_dev(3); + + int64_t chunk_start = bos + static_cast(ci) * C; + int64_t valid = slen - static_cast(ci) * C; + if (valid > C) valid = C; + + { + GmShape2D s_shape(D, D); + GmStride2D s_stride(D); + GmTensor2D s_global(workspace_handle + ws_base + WS_S, s_shape, + s_stride); + DynMatL1 s_l1_load(D, D); + TASSIGN(s_l1_load, 0); + // Load the previous recurrent state S_i from per-core workspace. + TLOAD(s_l1_load, s_global); + } + + int64_t w_offset = ((chunk_start) * H + head) * D; + { + GmShape2D w_shape(static_cast(valid), D); + GmStride2D w_stride(BSND_QKV_STRIDE); + GmTensor2D w_global(W_handle + w_offset, w_shape, w_stride); + DynMatL1 w_l1_load(static_cast(valid), D); + TASSIGN(w_l1_load, D * D * static_cast(sizeof(half))); + TLOAD(w_l1_load, w_global); + if (valid != C) { + TFILLPAD(w_l1_load, w_l1_load); + } + } + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + // Apply the carried recurrent state to every token in this chunk. + gemm_v0( + w_l1, s_l1, ws_l0, (bool)1); + + { + GmShape2D ws_shape(C, D); + GmStride2D ws_stride(D); + GmTensor2D ws_global(workspace_handle + ws_base + WS_WS, + ws_shape, ws_stride); + DynAccTile ws_store(C, D); + TASSIGN(ws_store, 0); + // Save ws_i so the Vec phase can do `v_new = U_i - ws_i`. + TSTORE(ws_global, ws_store); + } + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (0 << 8)); + + wait_flag_dev(1); + + { + GmShape2D k_shape(D, C); + GmStride2D k_stride(C); + GmTensor2D k_global(workspace_handle + ws_base + WS_K, k_shape, + k_stride); + DynMatL1 k_l1_load(D, C); + TASSIGN(k_l1_load, (DD + C * D) * static_cast(sizeof(half))); + TLOAD(k_l1_load, k_global); + } + + int64_t v_offset = ((chunk_start) * H + head) * D; + { + GmShape2D v_shape(static_cast(valid), D); + GmStride2D v_stride(BSND_QKV_STRIDE); + GmTensor2D v_global(V_handle + v_offset, v_shape, v_stride); + DynMatL1 v_l1_load(static_cast(valid), D); + TASSIGN(v_l1_load, + (DD + C * D + D * C) * static_cast(sizeof(half))); + TLOAD(v_l1_load, v_global); + if (valid != C) { + TFILLPAD(v_l1_load, v_l1_load); + } + } + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + // This chunk contributes the additive update K_i^T V_i to the state recurrence. + gemm_v0( + k_l1, v_l1, kv_l0, (bool)1); + + { + GmShape2D kv_shape(D, D); + GmStride2D kv_stride(D); + GmTensor2D kv_global(workspace_handle + ws_base + WS_KV, + kv_shape, kv_stride); + DynAccTile kv_store(D, D); + TASSIGN(kv_store, C * D * static_cast(sizeof(float))); + // Save kv = k_tilde^T @ v_i_new so Vec can finish the state update. + TSTORE(kv_global, kv_store); + } + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (2 << 8)); + } + } +#endif +#if defined(__DAV_C220_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + + // Vec owns the running recurrent state S_i and updates it after every chunk. + for (int64_t wi = 0; wi < (total_work + block_num - 1) / block_num; ++wi) { + int64_t pid = wi * block_num + cid; + if (pid >= total_work) break; + + int64_t head = pid % H; + int64_t head_g = head / GROUP; + int64_t seq_idx = pid / H; + + int64_t bos, slen; + int64_t chunk_offset = 0; + if (cu_seqlens != nullptr) { + bos = static_cast(cu_seqlens[seq_idx]); + int64_t eos = static_cast(cu_seqlens[seq_idx + 1]); + slen = eos - bos; + for (int64_t si = 0; si < seq_idx; ++si) { + int64_t sb = static_cast(cu_seqlens[si]); + int64_t se = static_cast(cu_seqlens[si + 1]); + chunk_offset += (se - sb + C - 1) / C; + } + } else { + bos = seq_idx * seq_len; + slen = seq_len; + chunk_offset = seq_idx * ((seq_len + C - 1) / C); + } + int64_t num_chunks = (slen + C - 1) / C; + int64_t ws_base = static_cast(cid) * WS_PER_CORE; + + set_flag(PIPE_V, PIPE_S, EVENT_ID0); + wait_flag(PIPE_V, PIPE_S, EVENT_ID0); + TEXPANDS(zero_ub, 0.0f); + set_flag(PIPE_V, PIPE_S, EVENT_ID0); + wait_flag(PIPE_V, PIPE_S, EVENT_ID0); + // Start each sequence/head recurrence from S_0 = 0. + TEXPANDS(s_ub, 0.0f); + + TCVT(s_ub_half, s_ub, pto::RoundMode::CAST_NONE); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + // `workspace_handle` is a `half*`, so all offsets here are in half elements. + GmShape2D s_shape(HalfC, D); + GmStride2D s_stride(D); + GmTensor2D s_global( + workspace_handle + ws_base + WS_S + vid * HalfC * D, + s_shape, s_stride); + DynVecTile s_store(HalfC, D); + TASSIGN(s_store, S_UB_HALF); + TSTORE(s_global, s_store); + } + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + + int64_t chunk_start_0 = bos; + int64_t valid0 = slen; + if (valid0 > C) valid0 = C; + // Vec work is split by row stripe, not by individual token. For the first + // chunk we compute exactly how many live rows belong to this sub-block's + // HalfC stripe so short tails do not overrun the packed BSND input. + int32_t valid_rows_0 = + static_cast(valid0 - static_cast(vid) * HalfC); + if (valid_rows_0 < 0) valid_rows_0 = 0; + if (valid_rows_0 > HalfC) valid_rows_0 = HalfC; + + int64_t k_offset_0 = + (chunk_start_0 * Hg + head_g) * D + vid * HalfC * BSND_K_STRIDE; + if (valid_rows_0 > 0) { + GmShape2D k_shape(valid_rows_0, D); + GmStride2D k_stride(BSND_K_STRIDE); + GmTensor2D k_global(K_handle + k_offset_0, k_shape, k_stride); + DynVecTile k_load(valid_rows_0, D); + TASSIGN(k_load, K_UB_HALF); + TLOAD(k_load, k_global); + if (valid_rows_0 != HalfC) { + TFILLPAD_INPLACE(k_ub_half, k_load); + } + } else { + // Empty stripe (typically vid=1 on a very short tail chunk): synthesize + // a zero tile so later full-width vector math and workspace stores still + // observe proper padding semantics. + TEXPANDS(k_ub, 0.0f); + TCVT(k_ub_half, k_ub, pto::RoundMode::CAST_NONE); + } + + { + GmShape2D g_shape(1, static_cast(valid0)); + GmStride2D g_stride(1); + GmTensor2D g_global(G_handle + head * total_tokens + chunk_start_0, + g_shape, g_stride); + DynVecTile g_load( + 1, static_cast(valid0)); + TASSIGN(g_load, G_UB); + TLOAD(g_load, g_global); + if (valid0 != C) { + TFILLPAD_INPLACE(g_ub, g_load); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + for (int32_t ci = 0; ci < static_cast(num_chunks); ++ci) { + int64_t chunk_start = bos + static_cast(ci) * C; + int64_t valid = slen - static_cast(ci) * C; + if (valid > C) valid = C; + int32_t valid_rows = + static_cast(valid - static_cast(vid) * HalfC); + if (valid_rows < 0) valid_rows = 0; + if (valid_rows > HalfC) valid_rows = HalfC; + // Each Vec subblock owns one contiguous HalfC-row stripe of the chunk. + // For short tail chunks, `valid_rows` may be smaller or even zero. This + // is the key fix that keeps ragged tails and dense varlen boundary mixes + // from reading or writing beyond the live rows in this stripe. + + int64_t u_offset = (chunk_start * H + head) * D + vid * HalfC * BSND_QKV_STRIDE; + if (valid_rows > 0) { + GmShape2D u_shape(valid_rows, D); + GmStride2D u_stride(BSND_QKV_STRIDE); + GmTensor2D u_global(U_handle + u_offset, u_shape, u_stride); + DynVecTile u_load(valid_rows, D); + TASSIGN(u_load, U_UB_HALF); + TLOAD(u_load, u_global); + if (valid_rows != HalfC) { + TFILLPAD_INPLACE(u_ub_half, u_load); + } + } else { + // No live rows for this stripe in the current chunk; keep the tile + // explicitly zero-padded so the remainder of the recurrence logic can + // run in full-tile form without special-casing every later step. + TEXPANDS(u_ub, 0.0f); + TCVT(u_ub_half, u_ub, pto::RoundMode::CAST_NONE); + } + + TCVT(k_ub, k_ub_half, pto::RoundMode::CAST_NONE); + + TileUbDataND g_ub_temp; + TASSIGN(g_ub_temp, G_UB + vid * 64 * sizeof(float)); + TMOV(g_v_ub, g_ub_temp); + + set_flag(PIPE_V, PIPE_S, EVENT_ID0); + wait_flag(PIPE_V, PIPE_S, EVENT_ID0); + float g_last = g_ub.GetValue(static_cast(valid) - 1); + // Rebase the chunk gate around g_last so the intra-chunk decay stays numerically local. + // Torch-like: + // coeff = exp(g_last - g_rows_owned_by_this_subblock) + TADDS(coeff_ub, g_v_ub, -g_last); + pipe_barrier(PIPE_V); + TSUB(coeff_ub, zero_ub, coeff_ub); + pipe_barrier(PIPE_V); + TEXP(coeff_ub, coeff_ub); + + TEXP(g_ub, g_ub); + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(u_ub, u_ub_half, pto::RoundMode::CAST_NONE); + + TileUbDataDN coeff_col_ub; + TASSIGN(coeff_col_ub, COEFF_UB); + TileUbDataND coeff_2d_ub; + TASSIGN(coeff_2d_ub, WS_UB); + // Broadcast one decay scalar per token row across the D feature columns: + // coeff_2d[row, :] = coeff[row] + TROWEXPAND(coeff_2d_ub, coeff_col_ub); + pipe_barrier(PIPE_V); + // `k_ub` now holds k_tilde = exp(g_last - g_i) * K_i. + TMUL(k_ub, k_ub, coeff_2d_ub); + pipe_barrier(PIPE_V); + + wait_flag_dev(0); + { + GmShape2D ws_shape(HalfC, D); + GmStride2D ws_stride(D); + GmTensor2D ws_global( + workspace_handle + ws_base + WS_WS + vid * HalfC * D, + ws_shape, ws_stride); + DynVecTile ws_load(HalfC, D); + TASSIGN(ws_load, U_UB_HALF); + TLOAD(ws_load, ws_global); + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(ws_ub, u_ub_half, pto::RoundMode::CAST_NONE); + // v_i_new = U_i - W_i @ S_i. + // In PyTorch notation: + // u_ub = u_ub - ws_ub + TSUB(u_ub, u_ub, ws_ub); + TCVT(u_ub_half, u_ub, pto::RoundMode::CAST_NONE); + TCVT(k_ub_half, k_ub, pto::RoundMode::CAST_NONE); + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + int64_t v_offset = (chunk_start * H + head) * D + vid * HalfC * BSND_QKV_STRIDE; + if (valid_rows > 0) { + GmShape2D v_shape(valid_rows, D); + GmStride2D v_stride(BSND_QKV_STRIDE); + GmTensor2D v_global(V_handle + v_offset, v_shape, v_stride); + DynVecTile v_store(valid_rows, D); + TASSIGN(v_store, U_UB_HALF); + TSTORE(v_global, v_store); + } + + // Spill both V_i_new and k_i_tilde so the Cube stage can form + // k_i_tilde^T @ V_i_new for this chunk. + { + GmShape2D k_shape(HalfC, D); + GmStride2D k_stride(D); + GmTensor2D k_global( + workspace_handle + ws_base + WS_K + vid * HalfC * D, + k_shape, k_stride); + DynVecTile k_store(HalfC, D); + TASSIGN(k_store, K_UB_HALF); + TSTORE(k_global, k_store); + } + + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + float exp_g_last = g_ub.GetValue(static_cast(valid) - 1); + // Carry the recurrence across chunks: S_{i+1} = exp(g_last) * S_i + K_i^T V_i. + TMULS(s_ub, s_ub, exp_g_last); + + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + if (ci + 1 < static_cast(num_chunks)) { + int64_t next_start = bos + static_cast(ci + 1) * C; + int64_t next_valid = slen - static_cast(ci + 1) * C; + if (next_valid > C) next_valid = C; + int32_t next_valid_rows = static_cast( + next_valid - static_cast(vid) * HalfC); + if (next_valid_rows < 0) next_valid_rows = 0; + if (next_valid_rows > HalfC) next_valid_rows = HalfC; + + int64_t nk_off = + (next_start * Hg + head_g) * D + vid * HalfC * BSND_K_STRIDE; + if (next_valid_rows > 0) { + GmShape2D k_shape(next_valid_rows, D); + GmStride2D k_stride(BSND_K_STRIDE); + GmTensor2D k_global(K_handle + nk_off, k_shape, k_stride); + DynVecTile k_load( + next_valid_rows, D); + TASSIGN(k_load, K_UB_HALF); + TLOAD(k_load, k_global); + if (next_valid_rows != HalfC) { + TFILLPAD_INPLACE(k_ub_half, k_load); + } + } else { + // Same tail-safe zero materialization for the prefetch path: the next + // chunk may have no rows in this stripe even though the other stripe + // is still active. + TEXPANDS(k_ub, 0.0f); + TCVT(k_ub_half, k_ub, pto::RoundMode::CAST_NONE); + } + + { + GmShape2D g_shape(1, static_cast(next_valid)); + GmStride2D g_stride(1); + GmTensor2D g_global(G_handle + head * total_tokens + next_start, + g_shape, g_stride); + DynVecTile g_load( + 1, static_cast(next_valid)); + TASSIGN(g_load, G_UB); + TLOAD(g_load, g_global); + if (next_valid != C) { + TFILLPAD_INPLACE(g_ub, g_load); + } + } + } + + wait_flag_dev(2); + { + GmShape2D kv_shape(HalfC, D); + GmStride2D kv_stride(D); + GmTensor2D kv_global( + workspace_handle + ws_base + WS_KV + vid * HalfC * D, + kv_shape, kv_stride); + DynVecTile kv_load(HalfC, D); + TASSIGN(kv_load, S_UB_HALF); + TLOAD(kv_load, kv_global); + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(kv_ub, s_ub_half, pto::RoundMode::CAST_NONE); + pipe_barrier(PIPE_ALL); + // Finish S_{i+1} = exp(g_last) * S_i + k_i_tilde^T @ v_i_new. + // Torch-like: + // s_ub = s_ub + kv_ub + TADD(s_ub, s_ub, kv_ub); + TCVT(s_ub_half, s_ub, pto::RoundMode::CAST_NONE); + + if (ci + 1 < static_cast(num_chunks)) { + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + GmShape2D s_shape(HalfC, D); + GmStride2D s_stride(D); + GmTensor2D s_global( + workspace_handle + ws_base + WS_S + vid * HalfC * D, + s_shape, s_stride); + DynVecTile s_store(HalfC, D); + TASSIGN(s_store, S_UB_HALF); + TSTORE(s_global, s_store); + } + + // Expose the post-chunk state so the next chunk (and debug/verification + // outputs) can see S_{i+1}. Conceptually: + // S_handle[chunk_idx + 1, head] = S_{i+1} + int64_t s_out_offset = ((chunk_offset + ci + 1) * H + head) * DD; + { + GmShape2D s_out_shape(HalfC, D); + GmStride2D s_out_stride(D); + GmTensor2D s_out_global( + S_handle + s_out_offset + vid * HalfC * D, s_out_shape, + s_out_stride); + DynVecTile s_out_store(HalfC, D); + TASSIGN(s_out_store, S_UB_HALF); + TSTORE(s_out_global, s_out_store); + } + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + } + + if (ci + 1 < static_cast(num_chunks)) { + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + } + } + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + int64_t fs_offset = (seq_idx * H + head) * DD; + { + GmShape2D fs_shape(HalfC, D); + GmStride2D fs_stride(D); + GmTensor2D fs_global(FS_handle + fs_offset + vid * HalfC * D, + fs_shape, fs_stride); + DynVecTile fs_store(HalfC, D); + TASSIGN(fs_store, S_UB_HALF); + TSTORE(fs_global, fs_store); + } + // Drain the fs TSTORE before the next work item reuses S_UB_HALF; + // otherwise the MTE3 read of the UB races with the zeroing/TCVT at the + // start of the next item and corrupts the stored final state (only the + // last work item per core was previously correct). + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + } +#endif +} + +#ifndef GDN_HG +#define GDN_HG GDN_H +#endif + +extern "C" __global__ AICORE void launch_chunk_h( + __gm__ uint8_t *K, __gm__ uint8_t *W, __gm__ uint8_t *U, + __gm__ uint8_t *G, + __gm__ uint8_t *S, __gm__ uint8_t *V, __gm__ uint8_t *FS, + __gm__ uint8_t *workspace, + __gm__ uint8_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, int64_t total_tokens, + uint64_t ffts_addr) +{ + chunk_h_kernel( + reinterpret_cast<__gm__ half *>(K), + reinterpret_cast<__gm__ half *>(W), + reinterpret_cast<__gm__ half *>(U), + reinterpret_cast<__gm__ float *>(G), + reinterpret_cast<__gm__ half *>(S), + reinterpret_cast<__gm__ half *>(V), + reinterpret_cast<__gm__ half *>(FS), + reinterpret_cast<__gm__ half *>(workspace), + reinterpret_cast<__gm__ int32_t *>(cu_seqlens), + batch_size, seq_len, total_tokens, ffts_addr); +} + +extern "C" void call_kernel( + uint32_t block_dim, void *stream, + uint8_t *K, uint8_t *W, uint8_t *U, uint8_t *G, + uint8_t *S, uint8_t *V, uint8_t *FS, + uint8_t *workspace, + uint8_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, int64_t total_tokens) +{ + uint32_t fftsLen{0}; + uint64_t fftsAddr{0}; + rtGetC2cCtrlAddr(&fftsAddr, &fftsLen); + launch_chunk_h<<>>( + K, W, U, G, S, V, FS, workspace, cu_seqlens, + batch_size, seq_len, total_tokens, fftsAddr); +} diff --git a/csrc/ascend/pto_chunk_gdn/chunk_o.cpp b/csrc/ascend/pto_chunk_gdn/chunk_o.cpp new file mode 100644 index 000000000..a1b23f446 --- /dev/null +++ b/csrc/ascend/pto_chunk_gdn/chunk_o.cpp @@ -0,0 +1,1249 @@ +// ============================================================================ +// chunk_o_kernel.cpp — Output computation for GatedDeltaNet (chunk-wise) +// +// Mathematical operation (per chunk of C tokens, per head h): +// +// O = (QK_gated @ V) + exp(g) * (Q @ S) +// = intra_chunk_attention + inter_chunk_state_contribution +// +// where: +// Q, K, V ∈ ℝ^{C×D} — query/key/value projections for this chunk +// S ∈ ℝ^{D×D} — accumulated hidden state entering this chunk +// G ∈ ℝ^{C} — cumulative gate values (pre-transposed [H,T]) +// Msk ∈ ℝ^{C×C} — lower-triangular causal mask +// +// Cube phase (3 GEMMs per chunk): +// 1. QK = Q @ K^T — intra-chunk attention scores +// 2. QS = Q @ S — query applied to accumulated state +// 3. QKV = QK_gated @ V — gated attention applied to values +// +// Vec phase (two sub-blocks process upper/lower C/2 rows): +// a. Load G → compute gating coefficients: +// coeff[i,j] = exp(min(g[i] - g[j], 0)) * mask[i,j] +// b. Apply gating to QK: QK_gated = QK * coeff +// c. Scale QS by exp(g): QS_gated = QS * exp(g_row) +// d. Combine: O = QS_gated + QKV +// e. Store O to GM in BSND layout +// +// Cross-core sync protocol (Cube ↔ Vec via FFTS): +// flag 0: Cube→Vec — QK and QS results ready in workspace +// flag 1: Vec→Cube — QK_gated written back, Cube can proceed to GEMM 3 +// flag 2: Cube→Vec — QKV result ready in workspace +// flag 3: Vec→Cube — Vec done with this chunk, Cube can reuse workspace +// +// NPU memory hierarchy used: +// GM → L1 (Cube-accessible) → L0A/L0B (matrix engines) → L0C (accumulator) +// GM → UB (Vec-accessible, on-chip SRAM) +// +// ── PTO / NPU Primer ────────────────────────────────────────────────── +// This kernel combines matrix multiplication (Cube) with element-wise gating +// (Vec) in a tightly coordinated 3-GEMM + gating pipeline per chunk. +// +// Execution timeline for one chunk: +// Cube: GEMM1(Q@K^T) → GEMM2(Q@S) → store QK,QS → signal Vec ──────┐ +// Vec: (meanwhile) load G, compute gating coefficients │ +// Vec: ←── wait for Cube signal ──── apply gating to QK → QK_gated │ +// Vec: store QK_gated → signal Cube ────────────────────────────────┐│ +// Cube: ←── wait for Vec signal ──── GEMM3(QK_gated@V) → store QKV ─┘│ +// Vec: ←── wait for Cube signal ──── scale QS, combine O=QKV+QS_g │ +// Vec: store O → signal Cube "done" ─────────────────────────────────┘ +// +// numpy pseudocode for the entire chunk computation: +// QK = Q @ K.T # GEMM 1 +// QS = Q @ S # GEMM 2 +// coeff = exp(min(g_row - g_col, 0)) * mask # gating (dynamic PTO) +// (``static_baseline/run_chunk_o_static.py`` uses exp(g_row-g_col) without min.) +// QK_gated = QK * coeff # apply gating +// QKV = QK_gated @ V # GEMM 3 +// O = QKV + QS * np.exp(g_row).reshape(-1, 1) # final output +// +// Key PTO APIs (with numpy/torch equivalents): +// TLOAD(dst, gm) — dst = gm_data (DMA: GM→UB/L1, async) +// TSTORE(gm, src) — gm = src (DMA: UB/L0C→GM, async) +// TASSIGN(tile, addr) — bind tile descriptor to buffer address +// TCVT(dst, src, mode) — type cast: dst = src.float() or .half() +// TMOV(dst, src) — copy: dst = src.clone() +// TADD(d, a, b) — d = a + b +// TSUB(d, a, b) — d = a - b +// TMUL(d, a, b) — d = a * b +// TMINS(d, s, val) — d = torch.clamp(s, max=val) +// TEXP(d, s) — d = torch.exp(s) +// TROWEXPAND(2d, col) — 2d[i,j] = col[i] (broadcast column→rows) +// TCOLEXPAND(2d, row) — 2d[i,j] = row[j] (broadcast row→columns) +// TEXTRACT(l0, l1, r, c) — copy L1 sub-tile → L0A/L0B (Cube input regs) +// TRESHAPE(zn, nz) — reinterpret L1 fractal layout (transpose, free) +// TMATMUL(C, A, B) — C = A @ B (Cube engine, fp16→fp32 accum) +// set_flag / wait_flag — synchronize pipes within same AI core +// ffts_cross_core_sync — signal across Cube↔Vec cores +// wait_flag_dev(flag) — wait for cross-core signal +// ============================================================================ + +#include +#include "acl/acl.h" +#include +using namespace pto; + +// ── Compile-time configuration (overridable at build time via -D flags) ── +// GDN_H: number of attention heads (default 16) +// GDN_D: hidden dimension per head (default 128) +// GDN_C: chunk size in tokens (default 128) +#ifndef GDN_H +#define GDN_H 16 +#endif + +#ifndef GDN_HG +#define GDN_HG GDN_H +#endif + +#ifndef GDN_D +#define GDN_D 128 +#endif + +#ifndef GDN_C +#define GDN_C 128 +#endif + +// ── PTO type aliases (device-only, guarded for host pass safety) ──────────── +// The bisheng compiler performs 3 passes: vec core, cube core (__CCE_AICORE__ +// defined), and host (__CCE_AICORE__ NOT defined). Type aliases using PTO +// tile types must be guarded so the host pass never sees them. +#ifdef __CCE_AICORE__ + +// UbND = Unified Buffer tile, row-major (ND) layout, for Vec SIMD ops. +// Like torch.empty((R, C), dtype=T) in fast on-chip SRAM (~256KB). +// RV, CV = valid region (handles dynamic shapes, partial chunks). +// PadValue::Zero = fill with 0 outside valid region during TLOAD. +// T=dtype, R×C=static shape, RV×CV=valid region, P=pad fill for TLOAD. +template +using UbND = pto::Tile; + +// UbDN = UB tile in column-major (DN) layout. +// Needed as source for TROWEXPAND which requires column-format input. +// TROWEXPAND takes a column vector and broadcasts it across all columns +// of a destination ND tile: dst[i,j] = col[i] for all j. +template +using UbDN = pto::Tile; + +// L1Mat = L1 cache tile in NZ fractal format — standard Cube GEMM input. +// Data is loaded here from GM via TLOAD, then fed to L0A/L0B via TEXTRACT. +template +using L1Mat = pto::Tile; + +// L1MatZN = ZN fractal format — used for transposed GEMM operands. +// TRESHAPE(l1_zn, l1_nz) converts NZ→ZN = logical matrix transpose (free, no data movement). +template +using L1MatZN = pto::Tile; + +#endif // __CCE_AICORE__ + +template +AICORE void chunk_o_kernel( + __gm__ half *Q_handle, __gm__ half *K_handle, __gm__ half *V_handle, + __gm__ half *S_handle, __gm__ float *G_handle, + __gm__ float *Msk_handle, + __gm__ half *workspace_qk_handle, + __gm__ half *workspace_qs_qkv_handle, + __gm__ half *workspace_qk_gated_handle, + __gm__ half *O_handle, + __gm__ int32_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, + int64_t total_tokens, + uint64_t ffts_addr) +{ + // Half the chunk — each Vec sub-block handles C/2 rows independently. + constexpr int32_t HalfChunk = ChunkSize / 2; + // KTail / CTail: the number of valid elements in the last 128-element tile + // when D or C isn't a multiple of 128. Used internally by PTO for partial tiles. + constexpr uint32_t KTail = + (HiddenSize % 128 == 0) ? 128 : (HiddenSize % 128); + constexpr uint32_t CTail = + (ChunkSize % 128 == 0) ? 128 : (ChunkSize % 128); + + constexpr int32_t H = NumHeads; + constexpr int32_t Hg = NumKeyHeads; + static_assert(Hg > 0 && H % Hg == 0, + "NumHeads must be divisible by NumKeyHeads"); + constexpr int32_t GROUP = H / Hg; + constexpr int32_t BSND_V_STRIDE = H * HiddenSize; + constexpr int32_t BSND_QK_STRIDE = Hg * HiddenSize; + + // Workspace sizes (in elements) shared between Cube and Vec via GM + constexpr int32_t WsQKSize = ChunkSize * ChunkSize; + constexpr int32_t WsQSSize = ChunkSize * HiddenSize; + constexpr int32_t WsGatedSize = ChunkSize * ChunkSize; + + // ── UB memory map (byte addresses within Unified Buffer) ───────────── + constexpr int32_t GUbAddr = 0; + constexpr int32_t MskUbAddr = 512; + constexpr int32_t QKUbAddr = 33280; + constexpr int32_t GvUbAddr = 66048; + constexpr int32_t CoeffUbAddr = 66304; + constexpr int32_t QKHalfUbAddr = 99072; + constexpr int32_t QSHalfUbAddr = 115456; + constexpr int32_t QSUbAddr = 131840; + constexpr int32_t OHalfUbAddr = 164608; + constexpr int32_t OUbAddr = QKUbAddr; + + // Initialize the cross-core FFTS signaling base address for this AI core. + set_ffts_base_addr(ffts_addr); + // cid = which AI core am I? (0..block_num-1). Used to partition work items. + auto cid = get_block_idx(); + // block_num = total number of AI cores running this kernel in parallel. + auto block_num = get_block_num(); + // vid = Vec sub-block ID (0 or 1). Each Vec core has 2 sub-blocks that + // process the upper (vid=0) and lower (vid=1) halves of C/2 rows. + auto vid = get_subblockid(); + + int64_t num_seqs = batch_size; + + // ── L1 tiles for Cube GEMM operands ────────────────────────────────── + // L1 holds matrices in NZ (col-major fractal) format for the matrix engine. + // Each tile is assigned a fixed L1 byte address to avoid runtime allocation. + // + // ── L1 tile layout for Cube GEMMs ──────────────────────────────────── + // L1 cache (~1MB) is manually partitioned for the 3 GEMMs: + // q_l1 at 0: Q [C×D] — shared by GEMM 1 and GEMM 2 + // k_l1 at 32768: K [C×D] — used in GEMM 1 (transposed via TRESHAPE) + // s_l1 at 65536: S [D×D] — accumulated state, used in GEMM 2 + // qk_gated at 98304: QK_gated [C×C] — from Vec, used in GEMM 3 + // v_l1 at 131072: V [C×D] — values, used in GEMM 3 + L1Mat q_l1; + TASSIGN(q_l1, 0); + L1Mat k_l1; + TASSIGN(k_l1, 32768); + TileAcc qk_l0; + TASSIGN(qk_l0, 0); + L1Mat s_l1; + TASSIGN(s_l1, 65536); + TileAcc qs_l0; + TASSIGN(qs_l0, 65536); + L1Mat qk_gated_l1; + TASSIGN(qk_gated_l1, 98304); + L1Mat v_l1; + TASSIGN(v_l1, 131072); + TileAcc qkv_l0; + TASSIGN(qkv_l0, 0); + + // ── UB tiles for Vec element-wise operations ───────────────────────── + // UB (Unified Buffer) is on-chip SRAM accessible by the Vec engine. + // Tiles here are row-major (ND) for standard element-wise ops. + // + // ── UB tile layout for Vec element-wise ops ────────────────────────── + // Each Vec sub-block (vid=0 or vid=1) processes C/2 rows of the C×C or C×D + // matrices. The UB layout (byte addresses) is designed so all needed tiles + // fit simultaneously in the ~256KB UB without overlapping: + // g_ub: gate values [1, C] float @ 0 + // msk_ub: causal mask [C/2, C] float @ 512 (loaded once, reused) + // qk_ub: QK scores in float [C/2, C] @ 33280 (after cast from half) + // g_v_ub: this sub-block's gate slice [1, C/2] @ 66048 + // coeff_ub: gating coefficients [C/2, C] float @ 66304 + // qk_ub_half: QK in half [C/2, C] @ 99072 + // qs_ub_half: QS in half [C/2, D] @ 115456 + // qs_ub: QS in float [C/2, D] @ 131840 + // o_ub_half: output O in half [C/2, D] @ 164608 + // o_ub: output O in float [C/2, D] @ QKUbAddr (reuses qk_ub space) + UbND g_ub; + TASSIGN(g_ub, GUbAddr); + UbND msk_ub; + TASSIGN(msk_ub, MskUbAddr); + UbND qk_ub; + TASSIGN(qk_ub, QKUbAddr); + UbND g_v_ub; + TASSIGN(g_v_ub, GvUbAddr); + UbND coeff_ub; + TASSIGN(coeff_ub, CoeffUbAddr); + UbND qk_ub_half; + TASSIGN(qk_ub_half, QKHalfUbAddr); + UbND qs_ub_half; + TASSIGN(qs_ub_half, QSHalfUbAddr); + UbND qs_ub; + TASSIGN(qs_ub, QSUbAddr); + UbND o_ub_half; + TASSIGN(o_ub_half, OHalfUbAddr); + UbND o_ub; + TASSIGN(o_ub, OUbAddr); + + // Total work items = (batches * chunks_per_sequence * heads). + // Each AI core (cid) picks every block_num-th work item (round-robin). + int64_t total_work = 0; + if (cu_seqlens == nullptr) { + int64_t chunks_per_seq = (seq_len + ChunkSize - 1) / ChunkSize; + total_work = num_seqs * chunks_per_seq * NumHeads; + } + +// ===================================================================== +// CUBE CORE — Three GEMMs per chunk: QK, QS, QKV +// Each AI core processes a different (chunk, head) pair. The Cube engine +// performs the heavy matmuls, then writes results to GM workspace for +// the Vec engine to apply gating and produce the final output. +// ===================================================================== +#if defined(__DAV_C220_CUBE__) + if (cu_seqlens == nullptr) { + // ── Fixed-length sequence path ────────────────────────────────────── + int64_t chunks_per_seq = (seq_len + ChunkSize - 1) / ChunkSize; + int64_t global_chunk_base = 0; + bool first_cube_iter = true; + + for (int64_t work_idx = static_cast(cid); + work_idx < total_work; + work_idx += static_cast(block_num)) { + // Wait for Vec to finish with previous chunk's workspace (flag 3) + if (!first_cube_iter) wait_flag_dev(3); + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + + int32_t head_idx = static_cast(work_idx % NumHeads); + int32_t head_g = head_idx / GROUP; + int64_t chunk_head_idx = work_idx / NumHeads; + int64_t seq_idx = chunk_head_idx / chunks_per_seq; + int64_t ci = chunk_head_idx % chunks_per_seq; + + int64_t bos = seq_idx * seq_len; + int64_t slen = seq_len; + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast( + remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + int32_t row_offset = static_cast(vid) * HalfChunk; + int32_t local_rows = valid_rows - row_offset; + if (local_rows < 0) local_rows = 0; + if (local_rows > HalfChunk) local_rows = HalfChunk; + + int64_t qk_off = + (chunk_token_start * static_cast(Hg) + + static_cast(head_g)) * + static_cast(HiddenSize); + int64_t v_off = + (chunk_token_start * static_cast(H) + + static_cast(head_idx)) * + static_cast(HiddenSize); + + int64_t chunk_global_idx = seq_idx * chunks_per_seq + ci; + int64_t s_offset = + (chunk_global_idx * NumHeads + head_idx) * + static_cast(HiddenSize) * + static_cast(HiddenSize); + + // ── Load Q [valid_rows × D] from GM → L1 ──────────────────────── + // GlobalTensor describes the GM layout with BSND strides. + // TLOAD performs DMA (MTE2 pipe). TFILLPAD zero-pads tail rows so + // downstream GEMMs see a clean C×D matrix. + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 0); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = valid_rows; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm(Q_handle + qk_off, _gs); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + // ── Load K [valid_rows × D] from GM → L1 ──────────────────────── + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 32768); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = valid_rows; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm(K_handle + qk_off, _gs); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + + // ── GEMM 1: QK = Q @ K^T (intra-chunk attention scores) ──────── + // ── GEMM 1: QK = Q @ K^T ───────────────────────────────────────── + // numpy: QK = Q @ K.T → [C×D] @ [D×C] = [C×C] + // + // How transpose works on NPU: + // K is loaded into L1 in NZ (col-major fractal) format. + // TRESHAPE(l1_zn, k_l1) reinterprets it as ZN (row-major fractal) = K^T. + // This is a ZERO-COST operation — no data movement, just metadata change. + // TEXTRACT then loads the transposed view into L0B. + // + // Cube GEMM pipeline: + // TEXTRACT(l0a, q_l1, 0, 0) — Q → L0A (left operand) + // TEXTRACT(l0b, k_zn, 0, 0) — K^T → L0B (right operand) + // TMATMUL(qk_l0, l0a, l0b) — QK = L0A × L0B → L0C accumulator + // + // transpose_B: TRESHAPE converts k_l1 from NZ → ZN fractal layout, + // effectively transposing K before TEXTRACT loads it into L0B. + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); wait_flag(PIPE_M, PIPE_MTE1, _we); + TEXTRACT(_l0a, q_l1, 0, 0); + L1MatZN _bzn; TRESHAPE(_bzn, k_l1); TEXTRACT(_l0b, _bzn, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(qk_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); wait_flag(PIPE_M, PIPE_FIX, _we); + } + + // ── Load S [D × D] from GM → L1 (accumulated hidden state) ───── + { + L1Mat _l1(HiddenSize, HiddenSize); + TASSIGN(_l1, 65536); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = HiddenSize; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm(S_handle + s_offset, _gs); + TLOAD(_l1, _gm); + } + + // ── GEMM 2: QS = Q @ S (query applied to accumulated state) ──── + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); wait_flag(PIPE_M, PIPE_MTE1, _we); + TEXTRACT(_l0a, q_l1, 0, 0); + TEXTRACT(_l0b, s_l1, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(qs_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); wait_flag(PIPE_M, PIPE_FIX, _we); + } + + // ── Store QK [C × C] from L0C → GM workspace (fp32→fp16 cast) ─── + // TSTORE on TileAcc triggers MTE3 DMA with implicit type conversion. + { + TileAcc _l0(ChunkSize, ChunkSize); + TASSIGN(_l0, 0); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_handle + + static_cast(cid) * WsQKSize, _gs); + TSTORE(_gm, _l0); + } + + // ── Store QS [C × D] from L0C → GM workspace ──────────────────── + { + TileAcc _l0(ChunkSize, HiddenSize); + TASSIGN(_l0, 65536); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + + static_cast(cid) * WsQSSize, _gs); + TSTORE(_gm, _l0); + } + + // Signal Vec: QK and QS are ready (flag 0, Cube→Vec) + // ── Cross-core sync protocol ────────────────────────────────────── + // Cube and Vec are SEPARATE physical cores. They exchange data through GM + // and coordinate via FFTS flags. Think of it as two processes communicating + // through shared memory with semaphores. + // + // ffts_cross_core_sync(PIPE_FIX, config): + // config = 1 | (mode << 4) | (flag_id << 8) + // mode=2: broadcast signal to all cores in this block + // flag_id: identifies which signal (0, 1, 2, 3) + // + // Protocol for this kernel: + // flag 0: Cube→Vec "QK and QS are ready in workspace" + // flag 1: Vec→Cube "QK_gated is ready for GEMM 3" + // flag 2: Cube→Vec "QKV (GEMM 3 result) is ready" + // flag 3: Vec→Cube "I'm done with this chunk, you can reuse workspace" + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (0 << 8)); + + // Wait for Vec to write QK_gated back (flag 1, Vec→Cube) + wait_flag_dev(1); + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + + // ── Load QK_gated [C × C] from GM workspace → L1 ──────────────── + { + L1Mat _l1(ChunkSize, ChunkSize); + TASSIGN(_l1, 98304); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_gated_handle + + static_cast(cid) * WsGatedSize, _gs); + TLOAD(_l1, _gm); + } + // ── Load V [valid_rows × D] from GM → L1 ──────────────────────── + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 131072); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = valid_rows; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm(V_handle + v_off, _gs); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + + // ── GEMM 3: QKV = QK_gated @ V (gated attention → values) ────── + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); wait_flag(PIPE_M, PIPE_MTE1, _we); + TEXTRACT(_l0a, qk_gated_l1, 0, 0); + TEXTRACT(_l0b, v_l1, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(qkv_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); wait_flag(PIPE_M, PIPE_FIX, _we); + } + + // ── Store QKV [C × D] from L0C → GM workspace ─────────────────── + // ── Workspace buffer reuse ──────────────────────────────────────── + // workspace_qs_qkv_handle is shared between QS (GEMM 2 output) and QKV + // (GEMM 3 output). This is safe because: + // 1. Vec reads QS BEFORE Cube writes QKV to the same buffer + // 2. The cross-core flags ensure proper ordering: + // - flag 0: QS ready (Vec reads QS) + // - flag 1: QK_gated ready (Vec done reading QS, Cube can write QKV) + // - flag 2: QKV ready (Vec reads QKV from same buffer) + { + TileAcc _l0(ChunkSize, HiddenSize); + TASSIGN(_l0, 0); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + + static_cast(cid) * WsQSSize, _gs); + TSTORE(_gm, _l0); + } + + // Signal Vec: QKV is ready (flag 2, Cube→Vec) + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (2 << 8)); + first_cube_iter = false; + } + } else { + // ── Variable-length sequence path (cu_seqlens != nullptr) ────────── + int64_t gi = 0; + int64_t chunk_global_idx = 0; + bool first_cube_iter_v = true; + for (int64_t si = 0; si < num_seqs; ++si) { + int64_t bos = static_cast(cu_seqlens[si]); + int64_t eos = static_cast(cu_seqlens[si + 1]); + int64_t slen = eos - bos; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < nc; ++ci) { + for (int32_t h = 0; h < NumHeads; ++h) { + if (gi % static_cast(block_num) == + static_cast(cid)) { + if (!first_cube_iter_v) wait_flag_dev(3); + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast( + remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + int32_t head_idx = h; + int32_t head_g = head_idx / GROUP; + + int64_t qk_off = + (chunk_token_start * static_cast(Hg) + + static_cast(head_g)) * + static_cast(HiddenSize); + int64_t v_off = + (chunk_token_start * static_cast(H) + + static_cast(head_idx)) * + static_cast(HiddenSize); + int64_t s_offset = + (chunk_global_idx * NumHeads + head_idx) * + static_cast(HiddenSize) * + static_cast(HiddenSize); + + // Load Q + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 0); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = valid_rows; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm(Q_handle + qk_off, _gs); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + // Load K + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 32768); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = valid_rows; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm(K_handle + qk_off, _gs); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + + // GEMM 1: QK = Q @ K^T (transpose_B via TRESHAPE NZ→ZN) + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); wait_flag(PIPE_M, PIPE_MTE1, _we); + TEXTRACT(_l0a, q_l1, 0, 0); + L1MatZN _bzn; TRESHAPE(_bzn, k_l1); TEXTRACT(_l0b, _bzn, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(qk_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); wait_flag(PIPE_M, PIPE_FIX, _we); + } + + // Load S + { + L1Mat _l1(HiddenSize, HiddenSize); + TASSIGN(_l1, 65536); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = HiddenSize; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm(S_handle + s_offset, _gs); + TLOAD(_l1, _gm); + } + + // GEMM 2: QS = Q @ S + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); wait_flag(PIPE_M, PIPE_MTE1, _we); + TEXTRACT(_l0a, q_l1, 0, 0); + TEXTRACT(_l0b, s_l1, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(qs_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); wait_flag(PIPE_M, PIPE_FIX, _we); + } + + // Store QK → workspace + { + TileAcc _l0(ChunkSize, ChunkSize); + TASSIGN(_l0, 0); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_handle + + static_cast(cid) * WsQKSize, _gs); + TSTORE(_gm, _l0); + } + + // Store QS → workspace + { + TileAcc _l0(ChunkSize, HiddenSize); + TASSIGN(_l0, 65536); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + + static_cast(cid) * WsQSSize, _gs); + TSTORE(_gm, _l0); + } + + // Cube→Vec: QK & QS ready (flag 0) + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (0 << 8)); + + // Wait Vec→Cube: QK_gated ready (flag 1) + wait_flag_dev(1); + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + + // Load QK_gated + { + L1Mat _l1(ChunkSize, ChunkSize); + TASSIGN(_l1, 98304); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_gated_handle + + static_cast(cid) * WsGatedSize, _gs); + TLOAD(_l1, _gm); + } + // Load V + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 131072); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = valid_rows; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm(V_handle + v_off, _gs); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + + // GEMM 3: QKV = QK_gated @ V + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); wait_flag(PIPE_M, PIPE_MTE1, _we); + TEXTRACT(_l0a, qk_gated_l1, 0, 0); + TEXTRACT(_l0b, v_l1, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(qkv_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); wait_flag(PIPE_M, PIPE_FIX, _we); + } + + { + TileAcc _l0(ChunkSize, HiddenSize); + TASSIGN(_l0, 0); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + + static_cast(cid) * WsQSSize, _gs); + TSTORE(_gm, _l0); + } + + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (2 << 8)); + first_cube_iter_v = false; + } + gi++; + } + chunk_global_idx++; + } + } + } +#endif + +// ===================================================================== +// VEC CORE — Gating, element-wise ops, output assembly +// Two Vec sub-blocks (vid=0,1) process upper/lower C/2 rows in parallel. +// Each sub-block independently: +// 1. Computes gating coefficients from G and the causal mask +// 2. Applies gating to the Cube's QK result → QK_gated +// 3. Scales the Cube's QS result by exp(g) +// 4. Combines QKV + scaled QS → final output O +// ===================================================================== +#if defined(__DAV_C220_VEC__) + // Vec engine initialization: set_mask_norm selects "normal" masking mode, + // and set_vector_mask(-1, -1) enables ALL SIMD lanes (no masking). + set_mask_norm(); + set_vector_mask(-1, -1); + + // ── Load causal mask once (reused across all chunks) ───────────────── + // ── Causal mask (loaded once, reused) ───────────────────────────────── + // The causal mask is a C×C lower-triangular matrix of 0s and 1s: + // mask[i,j] = 1 if i >= j else 0 + // Each sub-block loads its C/2 rows. Applied via TMUL to zero out + // non-causal (future) attention scores. + // + // Each sub-block (vid=0,1) loads its C/2 rows of the C×C lower-tri mask. + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = HalfChunk; _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + Msk_handle + + static_cast(vid) * HalfChunk * ChunkSize, _gs); + UbND _ld(HalfChunk, ChunkSize); + TASSIGN(_ld, MskUbAddr); + TLOAD(_ld, _gm); + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + if (cu_seqlens == nullptr) { + // ── Fixed-length sequence path ────────────────────────────────────── + int64_t chunks_per_seq = (seq_len + ChunkSize - 1) / ChunkSize; + + for (int64_t work_idx = static_cast(cid); + work_idx < total_work; + work_idx += static_cast(block_num)) { + int32_t head_idx = static_cast(work_idx % NumHeads); + int64_t chunk_head_idx = work_idx / NumHeads; + int64_t seq_idx = chunk_head_idx / chunks_per_seq; + int64_t ci = chunk_head_idx % chunks_per_seq; + + int64_t bos = seq_idx * seq_len; + int64_t slen = seq_len; + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast( + remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + int32_t row_offset = static_cast(vid) * HalfChunk; + int32_t local_rows = valid_rows - row_offset; + if (local_rows < 0) local_rows = 0; + if (local_rows > HalfChunk) local_rows = HalfChunk; + + if (local_rows > 0) { + // ── Load G [1 × valid_rows] — gate values for this chunk ──────── + // G is pre-transposed to [H, total_tokens], contiguous per head. + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = 1; _gs.shape[4] = valid_rows; + GlobalTensor> _gm( + G_handle + static_cast(head_idx) * total_tokens + + chunk_token_start, _gs); + UbND _ld(1, valid_rows); + TASSIGN(_ld, GUbAddr); + TLOAD(_ld, _gm); + if (valid_rows != ChunkSize) { + UbND _pd; + TASSIGN(_pd, GUbAddr); + TFILLPAD_INPLACE(_pd, _ld); + } + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // ── Compute gating coefficients ────────────────────────────────── + // ── Gating coefficient computation (numpy pseudocode) ───────────── + // For this sub-block's rows (vid=0: rows 0..C/2-1, vid=1: rows C/2..C-1): + // + // g_row = g[my_start:my_start+C/2] # my gates (shape [C/2]) + // g_col = g[0:C] # full chunk gates (shape [C]) + // + // # Broadcast to 2D matrices: + // g_r_2d = g_row[:, None] * np.ones((1, C)) # TROWEXPAND: [C/2, C] + // g_c_2d = np.ones((C/2, 1)) * g_col[None, :] # TCOLEXPAND: [C/2, C] + // coeff = exp(min(g_r_2d - g_c_2d, 0)) * mask + // + // # Also compute exp(g_row) for QS scaling: + // exp_g_row = np.exp(g_row) # TEXP + UbND g_ub_temp_0; + TASSIGN(g_ub_temp_0, + GUbAddr + static_cast(vid) * HalfChunk * + static_cast(sizeof(float))); + TMOV(g_v_ub, g_ub_temp_0); + + // Broadcast g_row into [C/2 × C] and g_col into [C/2 × C] + UbND g_r_2d; + TASSIGN(g_r_2d, QSUbAddr); + UbDN g_v_col; + TASSIGN(g_v_col, GvUbAddr); + TROWEXPAND(g_r_2d, g_v_col); // g_r_2d[i,j] = g_row[i] + TCOLEXPAND(coeff_ub, g_ub); // coeff[i,j] = g_col[j] + TSUB(coeff_ub, g_r_2d, coeff_ub); // d = g_row - g_col + pipe_barrier(PIPE_V); + TMINS(coeff_ub, coeff_ub, 0.0f); + pipe_barrier(PIPE_V); + TEXP(coeff_ub, coeff_ub); + pipe_barrier(PIPE_V); + TMUL(coeff_ub, coeff_ub, msk_ub); + pipe_barrier(PIPE_V); + TEXP(g_v_ub, g_v_ub); // exp(g_row) for QS scaling + } + + // ── Wait for Cube→Vec flag 0: QK & QS ready ───────────────────── + wait_flag_dev(0); + if (local_rows == 0) { + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + wait_flag_dev(2); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + continue; + } + + // ── Load QK [C/2 × C] from workspace → UB ─────────────────────── + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_handle + + static_cast(cid) * WsQKSize + + static_cast(vid) * HalfChunk * ChunkSize, _gs); + UbND _ld(local_rows, ChunkSize); + TASSIGN(_ld, QKHalfUbAddr); + TLOAD(_ld, _gm); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(qk_ub_half, _ld); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(qk_ub, qk_ub_half, pto::RoundMode::CAST_NONE); + + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + + // ── Load QS [C/2 × D] from workspace → UB ─────────────────────── + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + + static_cast(cid) * WsQSSize + + static_cast(vid) * HalfChunk * HiddenSize, _gs); + UbND _ld(local_rows, HiddenSize); + TASSIGN(_ld, QSHalfUbAddr); + TLOAD(_ld, _gm); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(qs_ub_half, _ld); + } + } + + // ── Apply gating: QK_gated = QK * exp(d*mask)*mask + TMUL(qk_ub, qk_ub, coeff_ub); + TCVT(qk_ub_half, qk_ub, pto::RoundMode::CAST_NONE); + + // ── Store QK_gated [C/2 × C] → workspace for Cube's GEMM 3 ───── + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_gated_handle + + static_cast(cid) * WsGatedSize + + static_cast(vid) * HalfChunk * ChunkSize, _gs); + UbND _st(local_rows, ChunkSize); + TASSIGN(_st, QKHalfUbAddr); + TSTORE(_gm, _st); + } + // Vec→Cube: QK_gated ready (flag 1) + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + + // ── Scale QS by exp(g): QS_gated = QS * exp(g_row) ────────────── + // ── Scale QS by exp(g): inter-chunk state contribution ──────────── + // numpy: QS_scaled = QS * np.exp(g_row)[:, None] (broadcast across D columns) + // TROWEXPAND broadcasts the scalar exp(g[i]) for each row i across all D columns, + // then TMUL applies it element-wise. This gates how much the accumulated state + // contributes to each token's output. + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(qs_ub, qs_ub_half, pto::RoundMode::CAST_NONE); + UbND g_exp_2d; + TASSIGN(g_exp_2d, CoeffUbAddr); + UbDN g_v_col2; + TASSIGN(g_v_col2, GvUbAddr); + TROWEXPAND(g_exp_2d, g_v_col2); // broadcast exp(g_row) across columns + pipe_barrier(PIPE_V); + TMUL(qs_ub, qs_ub, g_exp_2d); // QS_gated = QS * exp(g_row) + + // ── Wait for Cube→Vec flag 2: QKV ready ───────────────────────── + wait_flag_dev(2); + + // ── Load QKV [C/2 × D] from workspace → UB ────────────────────── + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + + static_cast(cid) * WsQSSize + + static_cast(vid) * HalfChunk * HiddenSize, _gs); + UbND _ld(local_rows, HiddenSize); + TASSIGN(_ld, OHalfUbAddr); + TLOAD(_ld, _gm); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(o_ub_half, _ld); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // ── Combine: O = QS_gated + QKV ───────────────────────────────── + // ── Final output: O = QKV + QS_scaled ───────────────────────────── + // numpy: O = (QK_gated @ V) + (Q @ S) * exp(g)[:, None] + // = intra_chunk_attention + inter_chunk_state_contribution + // TCVT half→float for QKV, then TADD, then TCVT float→half for output. + TCVT(o_ub, o_ub_half, pto::RoundMode::CAST_NONE); + TADD(o_ub, qs_ub, o_ub); + TCVT(o_ub_half, o_ub, pto::RoundMode::CAST_NONE); + + // ── Store O [C/2 × D] → GM in BSND layout ─────────────────────── + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + int64_t o_offset = + (chunk_token_start * static_cast(H) + + static_cast(head_idx)) * + static_cast(HiddenSize) + + static_cast(vid) * HalfChunk * + static_cast(BSND_V_STRIDE); + + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + O_handle + o_offset, _gs); + UbND _st(local_rows, HiddenSize); + TASSIGN(_st, OHalfUbAddr); + TSTORE(_gm, _st); + } + + // Vec→Cube: done with this chunk (flag 3) + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + } + } else { + // ── Variable-length sequence path (cu_seqlens != nullptr) ────────── + int64_t gi = 0; + for (int64_t si = 0; si < num_seqs; ++si) { + int64_t bos = static_cast(cu_seqlens[si]); + int64_t eos = static_cast(cu_seqlens[si + 1]); + int64_t slen = eos - bos; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < nc; ++ci) { + for (int32_t h = 0; h < NumHeads; ++h) { + if (gi % static_cast(block_num) == + static_cast(cid)) { + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast( + remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + int32_t head_idx = h; + int32_t row_offset = static_cast(vid) * HalfChunk; + int32_t local_rows = valid_rows - row_offset; + if (local_rows < 0) local_rows = 0; + if (local_rows > HalfChunk) local_rows = HalfChunk; + + if (local_rows > 0) { + // Load G + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = 1; _gs.shape[4] = valid_rows; + GlobalTensor> _gm( + G_handle + static_cast(head_idx) * total_tokens + + chunk_token_start, _gs); + UbND _ld(1, valid_rows); + TASSIGN(_ld, GUbAddr); + TLOAD(_ld, _gm); + if (valid_rows != ChunkSize) { + UbND _pd; + TASSIGN(_pd, GUbAddr); + TFILLPAD_INPLACE(_pd, _ld); + } + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Compute gating coefficients (same math as fixed-length path — see detailed pseudocode above) + UbND g_ub_temp_v; + TASSIGN(g_ub_temp_v, + GUbAddr + + static_cast(vid) * HalfChunk * + static_cast(sizeof(float))); + TMOV(g_v_ub, g_ub_temp_v); + + UbND g_r_2d_v; + TASSIGN(g_r_2d_v, QSUbAddr); + UbDN g_v_col_v; + TASSIGN(g_v_col_v, GvUbAddr); + TROWEXPAND(g_r_2d_v, g_v_col_v); + TCOLEXPAND(coeff_ub, g_ub); + TSUB(coeff_ub, g_r_2d_v, coeff_ub); // d = g_row - g_col + pipe_barrier(PIPE_V); + TMINS(coeff_ub, coeff_ub, 0.0f); + pipe_barrier(PIPE_V); + TEXP(coeff_ub, coeff_ub); + pipe_barrier(PIPE_V); + TMUL(coeff_ub, coeff_ub, msk_ub); + pipe_barrier(PIPE_V); + TEXP(g_v_ub, g_v_ub); + } + + wait_flag_dev(0); + if (local_rows == 0) { + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + wait_flag_dev(2); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + } else { + // Load QK from workspace + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_handle + + static_cast(cid) * WsQKSize + + static_cast(vid) * HalfChunk * ChunkSize, _gs); + UbND _ld(local_rows, ChunkSize); + TASSIGN(_ld, QKHalfUbAddr); + TLOAD(_ld, _gm); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(qk_ub_half, _ld); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(qk_ub, qk_ub_half, pto::RoundMode::CAST_NONE); + + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + + // Load QS from workspace + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + + static_cast(cid) * WsQSSize + + static_cast(vid) * HalfChunk * HiddenSize, _gs); + UbND _ld(local_rows, HiddenSize); + TASSIGN(_ld, QSHalfUbAddr); + TLOAD(_ld, _gm); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(qs_ub_half, _ld); + } + } + + TMUL(qk_ub, qk_ub, coeff_ub); + TCVT(qk_ub_half, qk_ub, pto::RoundMode::CAST_NONE); // float→half for GM store + + // Store QK_gated → workspace + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_qk_gated_handle + + static_cast(cid) * WsGatedSize + + static_cast(vid) * HalfChunk * ChunkSize, _gs); + UbND _st(local_rows, ChunkSize); + TASSIGN(_st, QKHalfUbAddr); + TSTORE(_gm, _st); + } + // Vec→Cube: QK_gated ready (flag 1) + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + + // Scale QS by exp(g): QS_scaled = QS * exp(g_row)[:, None] + // (same inter-chunk state scaling as fixed-length path) + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(qs_ub, qs_ub_half, pto::RoundMode::CAST_NONE); // half→float for Vec math + + UbND g_exp_2d_v; + TASSIGN(g_exp_2d_v, CoeffUbAddr); + UbDN g_v_col2_v; + TASSIGN(g_v_col2_v, GvUbAddr); + TROWEXPAND(g_exp_2d_v, g_v_col2_v); + pipe_barrier(PIPE_V); + TMUL(qs_ub, qs_ub, g_exp_2d_v); + + wait_flag_dev(2); + + // Load QKV from workspace + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + workspace_qs_qkv_handle + + static_cast(cid) * WsQSSize + + static_cast(vid) * HalfChunk * HiddenSize, _gs); + UbND _ld(local_rows, HiddenSize); + TASSIGN(_ld, OHalfUbAddr); + TLOAD(_ld, _gm); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(o_ub_half, _ld); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // O = QS_gated + QKV (final output: intra-chunk attention + inter-chunk state) + TCVT(o_ub, o_ub_half, pto::RoundMode::CAST_NONE); // half→float + TADD(o_ub, qs_ub, o_ub); // O = QS_scaled + QKV + TCVT(o_ub_half, o_ub, pto::RoundMode::CAST_NONE); // float→half for GM store + + // Store O → GM + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + int64_t o_offset = + (chunk_token_start * static_cast(H) + + static_cast(head_idx)) * + static_cast(HiddenSize) + + static_cast(vid) * HalfChunk * + static_cast(BSND_V_STRIDE); + + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_rows; _gs.shape[4] = HiddenSize; + GlobalTensor> _gm( + O_handle + o_offset, _gs); + UbND _st(local_rows, HiddenSize); + TASSIGN(_st, OHalfUbAddr); + TSTORE(_gm, _st); + } + + // Vec→Cube: done with this chunk (flag 3) + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + } + } + gi++; + } + } + } + } +#endif +} + +// ── Device kernel entry point ───────────────────────────────────────── +// extern "C" __global__ AICORE: NPU kernel function. +// Runs on each AI core independently. Args are uint8_t* (type-erased) +// because the NPU launch ABI passes all pointers as raw bytes; we +// reinterpret_cast them to the correct types before calling the template. +extern "C" __global__ AICORE void launch_chunk_o( + __gm__ uint8_t *Q_handle, __gm__ uint8_t *K_handle, + __gm__ uint8_t *V_handle, __gm__ uint8_t *S_handle, + __gm__ uint8_t *G_handle, __gm__ uint8_t *Msk_handle, + __gm__ uint8_t *workspace_qk, __gm__ uint8_t *workspace_qs_qkv, + __gm__ uint8_t *workspace_qk_gated, + __gm__ uint8_t *O_handle, + __gm__ uint8_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, + int64_t total_tokens, + uint64_t ffts_addr) +{ + chunk_o_kernel( + reinterpret_cast<__gm__ half *>(Q_handle), + reinterpret_cast<__gm__ half *>(K_handle), + reinterpret_cast<__gm__ half *>(V_handle), + reinterpret_cast<__gm__ half *>(S_handle), + reinterpret_cast<__gm__ float *>(G_handle), + reinterpret_cast<__gm__ float *>(Msk_handle), + reinterpret_cast<__gm__ half *>(workspace_qk), + reinterpret_cast<__gm__ half *>(workspace_qs_qkv), + reinterpret_cast<__gm__ half *>(workspace_qk_gated), + reinterpret_cast<__gm__ half *>(O_handle), + reinterpret_cast<__gm__ int32_t *>(cu_seqlens), + batch_size, seq_len, total_tokens, ffts_addr); +} + +// ── Host launcher (called from Python ctypes) ───────────────────────── +// Launches kernel on block_dim AI cores via NPU stream. +// rtGetC2cCtrlAddr obtains the FFTS (cross-core sync) control address that +// the kernel needs for Cube↔Vec flag signaling. +extern "C" void call_kernel( + uint32_t block_dim, void *stream, + uint8_t *q, uint8_t *k, uint8_t *v, uint8_t *s, uint8_t *g_sum, + uint8_t *mask, + uint8_t *workspace_qk, uint8_t *workspace_qs_qkv, + uint8_t *workspace_qk_gated, + uint8_t *o, + uint8_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, + int64_t total_tokens) +{ + uint32_t fftsLen{0}; + uint64_t fftsAddr{0}; + rtGetC2cCtrlAddr(&fftsAddr, &fftsLen); + launch_chunk_o<<>>( + q, k, v, s, g_sum, mask, + workspace_qk, workspace_qs_qkv, workspace_qk_gated, + o, + cu_seqlens, + batch_size, seq_len, total_tokens, fftsAddr); +} diff --git a/csrc/ascend/pto_chunk_gdn/include/kernel_utils.h b/csrc/ascend/pto_chunk_gdn/include/kernel_utils.h new file mode 100644 index 000000000..d371e9b05 --- /dev/null +++ b/csrc/ascend/pto_chunk_gdn/include/kernel_utils.h @@ -0,0 +1,47 @@ +/** +Copyright (c) 2026 Huawei Technologies Co., Ltd. +All rights reserved. + +See LICENSE in the root of the software repository: +https://github.com/huawei-csl/pto-kernels/ +for the full License text. +*/ +#pragma once + +#define MEMORY_BASE +#include +#include + +namespace kernel_utils { +/** + * @brief Do a sync step (set-wait flag) between two pipes. + * + * @tparam SrcPipe The pipe that sets the flag. + * @tparam DstPipe The pipe that waits for the flag. + * @param [in] id The event id to sync for. + */ +template +AICORE inline void SetWaitFlag(uint32_t id) { + set_flag(SrcPipe, DstPipe, static_cast(id)); + wait_flag(SrcPipe, DstPipe, static_cast(id)); +} + +/** + * @brief Performs a division on two integral numbers and rounds the result up + * to the nearest integer. + * + * @tparam T1 Data type of dividend. + * @tparam T2 Data type of divisor. + * @param [in] value Dividend. + * @param [in] divisor Divisor. + * @return Result of division. + */ +template ::value && + std::is_integral::value, + int>::type = 0> +AICORE inline T1 CeilDiv(T1 value, T2 divisor) { + return (value + divisor - 1) / divisor; +} + +} // namespace kernel_utils diff --git a/csrc/ascend/pto_chunk_gdn/mega_kernel.cpp b/csrc/ascend/pto_chunk_gdn/mega_kernel.cpp new file mode 100644 index 000000000..64562f0e3 --- /dev/null +++ b/csrc/ascend/pto_chunk_gdn/mega_kernel.cpp @@ -0,0 +1,501 @@ +// mega_kernel.cpp — GDN Mega-Kernel (group-value / GQA): all PTO stages in one launch +// +// Same pipeline as pto_mega_kernel, but scaled_dot_kkt / wy_fast / chunk_h / chunk_o use +// templates (H, Hg) from dynamic_bsnd_groupvalue; cumsum still uses H (value heads) like +// dynamic_bsnd. +// +// Stages: +// 1. cumsum (Vec) +// 2. transpose (Vec) +// 3. kkt (Cube+Vec) — K has Hg heads; β,g,A use H value heads +// 4. solve_tril (Cube) +// 5. wy_fast (Vec+Cube) +// 6. chunk_h (Cube+Vec) +// 7. chunk_o (Cube+Vec) + +#ifndef GDN_H +#define GDN_H 16 +#endif +#ifndef GDN_HG +#define GDN_HG GDN_H +#endif +#ifndef GDN_D +#define GDN_D 128 +#endif +#ifndef GDN_C +#define GDN_C 128 +#endif +#ifndef MEMORY_BASE +#define MEMORY_BASE +#endif + +#include +#include "acl/acl.h" +#include +#include +using namespace pto; + +// =================================================================== +// Device-only helpers (shared with standard mega-kernel) +// =================================================================== +#ifdef __CCE_AICORE__ + +constexpr uint16_t SYNC_MODE_SHIFT_VALUE = 4; +constexpr uint16_t SYNC_FLAG_SHIFT_VALUE = 8; + +AICORE inline uint16_t GetffstMsg(uint16_t mode, uint16_t flagId) +{ + return (0x1 + ((mode & 0x3) << SYNC_MODE_SHIFT_VALUE) + + ((flagId & 0xf) << SYNC_FLAG_SHIFT_VALUE)); +} + +template +AICORE inline void SyncAllImpl() +{ + pipe_barrier(PIPE_ALL); + if constexpr (isAIVOnly) { + ffts_cross_core_sync(PIPE_MTE3, GetffstMsg(0x0, SYNC_AIV_ONLY_ALL)); + wait_flag_dev(SYNC_AIV_ONLY_ALL); + return; + } +#if defined(__DAV_C220_CUBE__) + wait_flag_dev(SYNC_AIV_FLAG); + ffts_cross_core_sync(PIPE_FIX, GetffstMsg(0x0, SYNC_AIC_FLAG)); + wait_flag_dev(SYNC_AIC_FLAG); + ffts_cross_core_sync(PIPE_MTE3, GetffstMsg(0x02, SYNC_AIC_AIV_FLAG)); +#elif defined(__DAV_C220_VEC__) + ffts_cross_core_sync(PIPE_MTE3, GetffstMsg(0x02, SYNC_AIV_FLAG)); + wait_flag_dev(SYNC_AIC_AIV_FLAG); +#endif +} + +template +AICORE void mega_transpose_TH_to_HT( + __gm__ T *src, __gm__ T *dst, int64_t T_len) +{ +#if defined(__DAV_C220_VEC__) + if (get_subblockid() != 0) return; + set_mask_norm(); + set_vector_mask(-1, -1); + + auto cid = get_block_idx(); + auto block_num = get_block_num(); + + constexpr int32_t BLOCK = 128; + constexpr int32_t H = static_cast(H_val); + constexpr int32_t ES = static_cast(sizeof(T)); + // Head count padded up so that HA * sizeof(T) is 32B-aligned; the + // padding columns are zero-filled before the transpose. + constexpr int32_t HA = ((H * ES + 31) / 32) * 32 / ES; + constexpr int32_t SRC_UB = 0; + constexpr int32_t DST_UB = SRC_UB + BLOCK * HA * ES; + constexpr int32_t TMP_UB = DST_UB + HA * BLOCK * ES; + + using UBSrcFull = Tile; + using UBSrcDyn = Tile; + using UBDst = Tile; + using UBDstDyn = Tile; + using UBTmp = Tile; + + using UBRow = Tile; + using UBRowDyn = Tile; + + using Gm2D = Shape<1, 1, 1, DYNAMIC, DYNAMIC>; + using Gm1D = Shape<1, 1, 1, 1, DYNAMIC>; + using GmSrcS = Stride<1, 1, 1, H, 1>; + using GmS1 = Stride<1, 1, 1, 1, 1>; + + UBSrcFull ub_src; TASSIGN(ub_src, SRC_UB); + UBDst ub_dst; TASSIGN(ub_dst, DST_UB); + UBTmp ub_tmp; TASSIGN(ub_tmp, TMP_UB); + + int64_t num_tok_blocks = (T_len + BLOCK - 1) / BLOCK; + + for (int64_t bi = static_cast(cid); bi < num_tok_blocks; + bi += static_cast(block_num)) { + int64_t t0 = bi * BLOCK; + int32_t valid = (t0 + BLOCK <= T_len) + ? BLOCK + : static_cast(T_len - t0); + + { + Gm2D gs; gs.shape[3] = valid; gs.shape[4] = H; + GlobalTensor gm(src + t0 * H, gs); + UBSrcDyn ld(valid, H); + TASSIGN(ld, SRC_UB); + TLOAD(ld, gm); + if (valid != BLOCK || HA != H) TFILLPAD_INPLACE(ub_src, ld); + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TTRANS(ub_dst, ub_src, ub_tmp); + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + for (int32_t h = 0; h < H; ++h) { + Gm1D gs; gs.shape[4] = valid; + GlobalTensor gm(dst + h * T_len + t0, gs); + UBRowDyn st(1, valid); + TASSIGN(st, DST_UB + h * BLOCK * ES); + TSTORE(gm, st); + } + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + } +#endif +} + +template +AICORE void mega_cast_fp32_to_fp16_bsnd( + __gm__ float *src, __gm__ half *dst, + uint32_t num_matrices, int64_t total_tokens) +{ +#if defined(__DAV_C220_VEC__) + if (get_subblockid() != 0) return; + set_mask_norm(); + set_vector_mask(-1, -1); + + auto cid = get_block_idx(); + auto block_num = get_block_num(); + + constexpr int32_t F32_UB = 0; + constexpr int32_t F16_UB = C * static_cast(sizeof(float)); + + using SrcUB = Tile; + using DynSrcUB = Tile; + using DstUB = Tile; + using DynDstUB = Tile; + using Gm1D = Shape<1, 1, 1, 1, DYNAMIC>; + using GmS1 = Stride<1, 1, 1, 1, 1>; + + SrcUB src_ub; TASSIGN(src_ub, F32_UB); + DstUB dst_ub; TASSIGN(dst_ub, F16_UB); + + for (uint32_t m = cid; m < num_matrices; m += block_num) { + uint32_t h = m % static_cast(H); + uint32_t chunk_idx = m / static_cast(H); + + for (int64_t t = 0; t < total_tokens; ++t) { + int64_t off = t * static_cast(H * C) + + static_cast(h * C); + + { + Gm1D gs; gs.shape[4] = C; + GlobalTensor gm(src + off, gs); + SrcUB ld; TASSIGN(ld, F32_UB); + TLOAD(ld, gm); + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TCVT(dst_ub, src_ub, RoundMode::CAST_NONE); + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + Gm1D gs; gs.shape[4] = C; + GlobalTensor gm(dst + off, gs); + DstUB st; TASSIGN(st, F16_UB); + TSTORE(gm, st); + } + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + } + } +#endif +} + +#endif // __CCE_AICORE__ + +// =================================================================== +// Include original kernel implementations in separate namespaces. +// =================================================================== + +#define call_kernel _mk_unused_gv_ck_cumsum +namespace mk_cumsum { +#include "chunk_cumsum.cpp" +} +#undef call_kernel + +#define call_kernel _mk_unused_gv_ck_kkt +namespace mk_kkt { +#include "scaled_dot_kkt.cpp" +} +#undef call_kernel + +namespace mk_solve { +#include "tri_inverse_impl.cpp" +} + +#define call_kernel _mk_unused_gv_ck_wy +namespace mk_wy { +#include "wy_fast.cpp" +} +#undef call_kernel + +#define call_kernel _mk_unused_gv_ck_h +namespace mk_h { +#include "chunk_h.cpp" +} +#undef call_kernel + +#define call_kernel _mk_unused_gv_ck_o +namespace mk_o { +#include "chunk_o.cpp" +} +#undef call_kernel + +AICORE void mega_solve_tril( + __gm__ half *out, __gm__ half *in, __gm__ half *minus_id, + uint32_t matrix_size, uint32_t num_matrices, + uint32_t num_bsnd_heads, + __gm__ int32_t *cu_seqlens, uint32_t is_lower) +{ + if (num_matrices <= get_block_num()) + mk_solve::runKernelTriInvRecUnroll( + out, in, minus_id, num_matrices, + num_bsnd_heads, cu_seqlens, is_lower); + else if (num_matrices <= 2u * get_block_num()) + mk_solve::runKernelTriInvRecUnroll( + out, in, minus_id, num_matrices, + num_bsnd_heads, cu_seqlens, is_lower); + else + mk_solve::runKernelTriInvRecUnroll( + out, in, minus_id, num_matrices, + num_bsnd_heads, cu_seqlens, is_lower); +} + +extern "C" __global__ AICORE void launch_mega_kernel( + __gm__ uint8_t *q_ptr, + __gm__ uint8_t *k_ptr, + __gm__ uint8_t *v_ptr, + __gm__ uint8_t *g_in_ptr, + __gm__ uint8_t *beta_ptr, + __gm__ uint8_t *msk_lower_ptr, + __gm__ uint8_t *msk_full_ptr, + __gm__ uint8_t *minus_id_ptr, + __gm__ uint8_t *cu_seqlens_ptr, + __gm__ uint8_t *o_ptr, + __gm__ uint8_t *g_sum_ptr, + __gm__ uint8_t *g_t_ptr, + __gm__ uint8_t *beta_t_ptr, + __gm__ uint8_t *A_ptr, + __gm__ uint8_t *A_inv_f32_ptr, + __gm__ uint8_t *A_inv_ptr, + __gm__ uint8_t *w_ptr, + __gm__ uint8_t *u_ptr, + __gm__ uint8_t *s_ptr, + __gm__ uint8_t *v_new_ptr, + __gm__ uint8_t *fs_ptr, + __gm__ uint8_t *kkt_ws_ptr, + __gm__ uint8_t *wy_ws_a1_ptr, + __gm__ uint8_t *wy_ws_a2_ptr, + __gm__ uint8_t *h_ws_ptr, + __gm__ uint8_t *o_ws_qk_ptr, + __gm__ uint8_t *o_ws_qs_ptr, + __gm__ uint8_t *o_ws_gated_ptr, + int64_t batch_size, + int64_t seq_len, + int64_t total_tokens, + uint32_t num_matrices, + uint64_t ffts_addr) +{ + set_ffts_base_addr(ffts_addr); + + constexpr int32_t H = GDN_H; + constexpr int32_t HG = GDN_HG; + constexpr int32_t D = GDN_D; + constexpr int32_t C = GDN_C; + + mk_cumsum::cumsum_kernel( + reinterpret_cast<__gm__ float *>(g_in_ptr), + reinterpret_cast<__gm__ float *>(g_sum_ptr), + reinterpret_cast<__gm__ int32_t *>(cu_seqlens_ptr), + batch_size, seq_len, ffts_addr); + +#ifdef MEGA_STOP_AFTER_CUMSUM + pipe_barrier(PIPE_ALL); + return; +#endif + + SyncAllImpl(); + +#ifdef MEGA_STOP_AFTER_SYNC1 + return; +#endif + + mega_transpose_TH_to_HT( + reinterpret_cast<__gm__ float *>(g_sum_ptr), + reinterpret_cast<__gm__ float *>(g_t_ptr), + total_tokens); + mega_transpose_TH_to_HT( + reinterpret_cast<__gm__ half *>(beta_ptr), + reinterpret_cast<__gm__ half *>(beta_t_ptr), + total_tokens); + +#ifdef MEGA_STOP_AFTER_TRANSPOSE + pipe_barrier(PIPE_ALL); + return; +#endif + + SyncAllImpl(); + + mk_kkt::kkt_kernel( + reinterpret_cast<__gm__ half *>(k_ptr), + reinterpret_cast<__gm__ half *>(beta_t_ptr), + reinterpret_cast<__gm__ float *>(g_t_ptr), + reinterpret_cast<__gm__ float *>(msk_lower_ptr), + reinterpret_cast<__gm__ half *>(kkt_ws_ptr), + reinterpret_cast<__gm__ half *>(A_ptr), + reinterpret_cast<__gm__ int32_t *>(cu_seqlens_ptr), + batch_size, seq_len, total_tokens, ffts_addr); + +#if defined(__DAV_C220_CUBE__) + pipe_barrier(PIPE_ALL); + wait_flag_dev(2); + wait_flag_dev(3); +#endif + +#ifdef MEGA_STOP_AFTER_KKT + pipe_barrier(PIPE_ALL); + return; +#endif + + SyncAllImpl(); + + mega_solve_tril( + reinterpret_cast<__gm__ half *>(A_inv_ptr), + reinterpret_cast<__gm__ half *>(A_ptr), + reinterpret_cast<__gm__ half *>(minus_id_ptr), + C, num_matrices, H, + reinterpret_cast<__gm__ int32_t *>(cu_seqlens_ptr), 1); + +#ifdef MEGA_STOP_AFTER_SOLVE + pipe_barrier(PIPE_ALL); + return; +#endif + + SyncAllImpl(); + +#ifdef MEGA_STOP_AFTER_CAST + pipe_barrier(PIPE_ALL); + return; +#endif + + SyncAllImpl(); + +#ifdef MEGA_STOP_AFTER_SYNC_BEFORE_WY + return; +#endif + + mk_wy::wy_fast_kernel( + reinterpret_cast<__gm__ half *>(k_ptr), + reinterpret_cast<__gm__ half *>(v_ptr), + reinterpret_cast<__gm__ half *>(beta_t_ptr), + reinterpret_cast<__gm__ float *>(g_t_ptr), + reinterpret_cast<__gm__ half *>(A_inv_ptr), + reinterpret_cast<__gm__ half *>(wy_ws_a1_ptr), + reinterpret_cast<__gm__ half *>(wy_ws_a2_ptr), + reinterpret_cast<__gm__ half *>(w_ptr), + reinterpret_cast<__gm__ half *>(u_ptr), + reinterpret_cast<__gm__ int32_t *>(cu_seqlens_ptr), + batch_size, seq_len, total_tokens, ffts_addr); + +#if defined(__DAV_C220_VEC__) + if (get_block_idx() < num_matrices) { + pipe_barrier(PIPE_ALL); + wait_flag_dev(3); + wait_flag_dev(4); + } +#endif + +#ifdef MEGA_STOP_AFTER_WY + pipe_barrier(PIPE_ALL); + return; +#endif + + SyncAllImpl(); + + mk_h::chunk_h_kernel( + reinterpret_cast<__gm__ half *>(k_ptr), + reinterpret_cast<__gm__ half *>(w_ptr), + reinterpret_cast<__gm__ half *>(u_ptr), + reinterpret_cast<__gm__ float *>(g_t_ptr), + reinterpret_cast<__gm__ half *>(s_ptr), + reinterpret_cast<__gm__ half *>(v_new_ptr), + reinterpret_cast<__gm__ half *>(fs_ptr), + reinterpret_cast<__gm__ half *>(h_ws_ptr), + reinterpret_cast<__gm__ int32_t *>(cu_seqlens_ptr), + batch_size, seq_len, total_tokens, ffts_addr); + +#ifdef MEGA_STOP_AFTER_H + pipe_barrier(PIPE_ALL); + return; +#endif + + SyncAllImpl(); + + mk_o::chunk_o_kernel( + reinterpret_cast<__gm__ half *>(q_ptr), + reinterpret_cast<__gm__ half *>(k_ptr), + reinterpret_cast<__gm__ half *>(v_new_ptr), + reinterpret_cast<__gm__ half *>(s_ptr), + reinterpret_cast<__gm__ float *>(g_t_ptr), + reinterpret_cast<__gm__ float *>(msk_full_ptr), + reinterpret_cast<__gm__ half *>(o_ws_qk_ptr), + reinterpret_cast<__gm__ half *>(o_ws_qs_ptr), + reinterpret_cast<__gm__ half *>(o_ws_gated_ptr), + reinterpret_cast<__gm__ half *>(o_ptr), + reinterpret_cast<__gm__ int32_t *>(cu_seqlens_ptr), + batch_size, seq_len, total_tokens, ffts_addr); + +#if defined(__DAV_C220_CUBE__) + if (get_block_idx() < num_matrices) { + pipe_barrier(PIPE_ALL); + wait_flag_dev(3); + } +#endif +} + +extern "C" void call_kernel( + uint32_t block_dim, void *stream, + uint8_t *q, uint8_t *k, uint8_t *v, + uint8_t *g_in, uint8_t *beta, + uint8_t *msk_lower, uint8_t *msk_full, + uint8_t *minus_id, uint8_t *cu_seqlens, + uint8_t *o, + uint8_t *g_sum, uint8_t *g_t, uint8_t *beta_t, + uint8_t *A, uint8_t *A_inv_f32, uint8_t *A_inv, + uint8_t *w, uint8_t *u, uint8_t *s, uint8_t *v_new, uint8_t *fs, + uint8_t *kkt_ws, uint8_t *wy_ws_a1, uint8_t *wy_ws_a2, + uint8_t *h_ws, + uint8_t *o_ws_qk, uint8_t *o_ws_qs, uint8_t *o_ws_gated, + int64_t batch_size, int64_t seq_len, int64_t total_tokens, + uint32_t num_matrices) +{ + uint32_t fftsLen{0}; + uint64_t fftsAddr{0}; + rtGetC2cCtrlAddr(&fftsAddr, &fftsLen); + launch_mega_kernel<<>>( + q, k, v, g_in, beta, msk_lower, msk_full, minus_id, cu_seqlens, + o, + g_sum, g_t, beta_t, A, A_inv_f32, A_inv, + w, u, s, v_new, fs, + kkt_ws, wy_ws_a1, wy_ws_a2, h_ws, + o_ws_qk, o_ws_qs, o_ws_gated, + batch_size, seq_len, total_tokens, num_matrices, + fftsAddr); +} diff --git a/csrc/ascend/pto_chunk_gdn/scaled_dot_kkt.cpp b/csrc/ascend/pto_chunk_gdn/scaled_dot_kkt.cpp new file mode 100644 index 000000000..8b0a4cd4a --- /dev/null +++ b/csrc/ascend/pto_chunk_gdn/scaled_dot_kkt.cpp @@ -0,0 +1,699 @@ +// ============================================================================ +// scaled_dot_kkt_kernel.cpp — Intra-chunk attention matrix for GatedDeltaNet +// +// Computes A = mask(KK^T · gating_coeff) per chunk, where: +// KK^T ∈ ℝ^{C×C} = K @ K^T (Cube engine, GEMM) +// coeff[i,j] = exp(clamp(g[i]+log(β[i]) - g[j], max=0)) (Vec engine) +// A[i,j] = KK^T[i,j] · coeff[i,j] · causal_mask[i,j] +// +// Inputs: +// K [total_tokens, Hg, D] half — key vectors (BSND along seq; stride Hg * D) +// Beta [H, total_tokens] half — gate bias per **value** head (pre-transposed) +// G [H, total_tokens] float — cumulative gate sum per **value** head +// Msk [C, C] float — lower-triangular causal mask +// +// Output: +// A [total_tokens, H, C] half — gated attention matrix in BSND +// +// Architecture: Cube + Vec cross-core kernel. +// Cube phase: K→L1, GEMM K@K^T→L0C, store to workspace (GM) +// Vec phase: load workspace KK^T, compute gating coefficients, apply mask +// +// Cross-core sync: Cube signals Vec via FFTS flag after each chunk's KK^T +// is written to workspace. Vec signals back when workspace buffer is free. +// Two workspace slots alternate (double-buffering via slot = ci & 1). +// +// Vec sub-blocks: Two sub-blocks (vid=0,1) process upper/lower halves of +// the C×C attention matrix in parallel (HalfChunk rows each). +// +// NPU memory hierarchy: +// GM → L1 (Cube-accessible) → L0A/L0B (GEMM operands) → L0C (accumulator) +// GM → UB (Vec-accessible SRAM) +// +// ── PTO / NPU Primer for This Kernel ────────────────────────────────── +// NPU Architecture (simplified): +// Each "AI Core" (like a GPU SM) has: +// - Cube engine: matrix multiply unit (like GPU Tensor Cores), works on L0A/L0B/L0C +// - Vec engine: SIMD vector unit (like GPU CUDA cores), works on UB (Unified Buffer) +// - MTE2: DMA engine for loading data: GM → L1 or GM → UB +// - MTE3: DMA engine for storing data: UB → GM or L0C → GM +// - MTE1: DMA engine for L1 → L0A/L0B transfers (internal to Cube pipeline) +// Memory hierarchy (fast→slow): L0 registers > L1 cache > UB (SRAM) > GM (HBM) +// Cube and Vec run on SEPARATE cores — they communicate via GM + cross-core flags. +// +// Key PTO APIs used in this kernel (with numpy/torch equivalents): +// TASSIGN(tile, addr) — Bind tile to UB/L1/L0 address (tile = memory[addr]) +// TLOAD(dst, gm_tensor) — DMA load: dst = gm_tensor (async, MTE2 pipe) +// TSTORE(gm, src) — DMA store: gm = src (async, MTE3 pipe) +// TFILLPAD(dst, src) — Zero-fill padding: dst[outside valid] = 0 +// TFILLPAD_INPLACE(d, s) — Same but in-place for UB tiles +// TEXTRACT(l0, l1, r, c) — Copy L1 sub-block → L0A or L0B (MTE1 pipe) +// TRESHAPE(dst, src) — Reinterpret L1 tile layout (NZ↔ZN for transpose) +// TMATMUL(C, A, B) — Matrix multiply: C = A @ B in Cube engine +// TCVT(dst, src, mode) — Type conversion: like dst = src.float() or src.half() +// TMOV(dst, src) — Copy: dst = src.clone() +// TADD(d, a, b) — Element-wise add: d = a + b +// TSUB(d, a, b) — Element-wise subtract: d = a - b +// TMUL(d, a, b) — Element-wise multiply: d = a * b +// TMINS(d, s, val) — Clamp max: d = torch.clamp(s, max=val) +// TEXP(d, s) — Element-wise exp: d = torch.exp(s) +// TLOG(d, s) — Element-wise log: d = torch.log(s) +// TROWEXPAND(2d, col) — Broadcast column → rows: 2d[i,j] = col[i] +// TCOLEXPAND(2d, row) — Broadcast row → cols: 2d[i,j] = row[j] +// set_flag(P1, P2, EVT) — Signal from pipe P1 to pipe P2 (like a semaphore post) +// wait_flag(P1, P2, EVT) — Wait for signal from P1 (like a semaphore wait) +// pipe_barrier(PIPE_V) — Local Vec barrier (ensure all Vec ops complete) +// pipe_barrier(PIPE_ALL) — Barrier for all local pipes +// ffts_cross_core_sync() — Cross-core signal (Cube↔Vec, different physical cores) +// wait_flag_dev(flag) — Wait for cross-core signal +// ============================================================================ + +#include // PTO (Performance Tile Operator): NPU kernel API +#include "acl/acl.h" // ACL (Ascend Computing Language): runtime API +#include // FFTS: cross-core synchronization primitives +using namespace pto; + +// ── Compile-time constants (set by the JIT compiler from Python) ────── +// These are typically passed as -DGDN_H=16 -DGDN_D=128 -DGDN_C=128 on the +// compiler command line. The #ifndef guards provide defaults for IDE tooling. +#ifndef GDN_H +#define GDN_H 16 // H = number of value heads (gates A β,g index here) +#endif + +#ifndef GDN_HG +#define GDN_HG GDN_H // Hg = shared key-query heads (GQA); default MHA +#endif + +#ifndef GDN_D +#define GDN_D 128 // D = hidden dimension per head +#endif + +#ifndef GDN_C +#define GDN_C 128 // C = chunk size (tokens processed per chunk) +#endif + +// ── PTO type aliases (device-only, guarded by __CCE_AICORE__) ─────────────── +// These are only compiled for the NPU device compiler (__CCE_AICORE__ is defined +// when compiling for AI Core hardware, similar to __CUDA_ARCH__ in CUDA). +#ifdef __CCE_AICORE__ +// UbND = UB tile in row-major (ND) layout for Vec engine. +// Think of it as: torch.empty((R, C), dtype=T) in on-chip SRAM. +// RV, CV = valid region (for dynamic shapes, like a[:valid_rows, :valid_cols]) +// The Vec engine (SIMD unit) reads/writes these tiles for element-wise ops. +template +using UbND = pto::Tile; + +// UbDN = UB tile in column-major (DN) layout — needed for TROWEXPAND source. +// TROWEXPAND requires its source vector in column-major (transposed) format. +// Same physical memory (UB SRAM), just different indexing convention. +template +using UbDN = pto::Tile; + +// L1Mat = L1 cache tile in NZ fractal format (col-major blocks, row-major within). +// This is the standard input format for the Cube matrix engine. +// Think of it as a matrix in L1 cache ready for GEMM. +// NZ = "Normal-Z": the default fractal layout that Cube expects for left/right operands. +template +using L1Mat = pto::Tile; + +// L1MatZN = L1 tile in ZN fractal format (row-major blocks, col-major within). +// Used when you need to transpose a matrix before GEMM: +// TRESHAPE(l1_zn, l1_nz) reinterprets NZ→ZN layout = logical transpose. +// This is FREE (no data movement) — it just changes how the Cube reads the bits. +template +using L1MatZN = pto::Tile; +#endif + +// ── Main kernel function (runs on each AI core) ────────────────────── +// Template parameters: NumHeads (H value), NumKeyHeads (Hg), HiddenSize, ChunkSize. +// GROUP = H/Hg; Cube loads K at head_g = head_idx / GROUP. +// +// __gm__: Marks pointers as Global Memory (HBM) — the NPU equivalent of +// CUDA's device memory. All input/output tensors live in GM. +template +AICORE void kkt_kernel( + __gm__ half *K_handle, __gm__ half *Beta_handle, + __gm__ float *G_handle, __gm__ float *Msk_handle, + __gm__ half *workspace_handle, __gm__ half *A_handle, + __gm__ int32_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, + int64_t total_tokens, + uint64_t ffts_addr) +{ + constexpr int32_t HalfChunk = ChunkSize / 2; + constexpr int32_t ChunkSquare = ChunkSize * ChunkSize; + static_assert(NumHeads % NumKeyHeads == 0, + "NumHeads must be divisible by NumKeyHeads (GQA grouping)"); + constexpr int32_t GROUP = NumHeads / NumKeyHeads; + constexpr int32_t BSND_QK_STRIDE = NumKeyHeads * HiddenSize; + // KTail: number of valid columns in the last 128-wide fractal block of K. + // If HiddenSize is a multiple of 128, the last block is fully used (128). + // Otherwise it's the remainder. Used internally by TLOAD for partial blocks. + constexpr uint32_t KTail = + (HiddenSize % 128 == 0) ? 128 : (HiddenSize % 128); + + // ── UB address map (manual memory planning) ───────────────────────── + // The UB is a flat SRAM; we manually assign byte offsets for each tile. + // This is like malloc'ing fixed regions — no dynamic allocator on NPU. + constexpr int32_t GUbAddr = 0; // g_ub: cumulative gates [1×C] + constexpr int32_t BetaHalfUbAddr = 512; // beta_ub_half: gate bias fp16 [1×C/2] + constexpr int32_t BetaUbAddr = 640; // beta_ub: gate bias fp32 [1×C/2] + constexpr int32_t GvUbAddr = 896; // g_v_ub: combined gate+bias [1×C/2] + constexpr int32_t AUbAddr = 1152; // a_ub: attention sub-block fp32 [C/2×C] + constexpr int32_t GRUbAddr = 33920; // g_r_ub: row gates [1×C/2] + constexpr int32_t GCUbAddr = 34176; // g_c_ub: column gates [1×C] + constexpr int32_t MskUbAddr = 34688; // msk_ub: causal mask [C/2×C] + constexpr int32_t GR2dUbAddr = 67456; // g_r_2d_ub: broadcast row gates [C/2×C] + constexpr int32_t GC2dUbAddr = 124800; // g_c_2d_ub: broadcast col gates [C/2×C] + constexpr int32_t CoeffUbAddr = 157568; // coeff_ub: gating coefficient [C/2×C] + // a_ub_half overlaps g_r_2d — safe because they're never live simultaneously + constexpr int32_t AUbHalfAddr = GR2dUbAddr; + + // set_ffts_base_addr: Tell the hardware where the cross-core flag table lives. + // This is a one-time setup so ffts_cross_core_sync / wait_flag_dev know + // which memory region to read/write for inter-core signaling. + set_ffts_base_addr(ffts_addr); + auto cid = get_block_idx(); // Which AI core am I? (like CUDA blockIdx.x) + auto block_num = get_block_num(); // Total AI cores launched (like CUDA gridDim.x) + // ── Vec sub-block parallelism ───────────────────────────────────────── + // Each AI core has 2 Vec sub-blocks (vid=0 and vid=1). + // They share the same UB memory but run independently in parallel. + // Here, vid=0 processes rows [0, C/2) and vid=1 processes rows [C/2, C). + // This halves the per-sub-block work and doubles Vec throughput. + auto vid = get_subblockid(); // 0 or 1: which Vec sub-block am I? + + // Work distribution: each (sequence, head) pair is one "work item". + // AI cores split work round-robin, just like CUDA blocks split a grid. + int64_t num_seqs = batch_size; + int64_t total_work = num_seqs * NumHeads; + + // ── Cube-side tile declarations ───────────────────────────────────── + // Cube-side tiles: K in L1 (NZ format), accumulator in L0C + L1Mat k_l1; + TASSIGN(k_l1, 0); + // TileAcc: L0C accumulator tile for GEMM results. + // The Cube engine always accumulates in float32 for precision, even when + // inputs are fp16. Think of it as: result = torch.matmul(a.half(), b.half()).float() + // When stored to GM via TSTORE with a half GlobalTensor, automatic fp32→fp16 cast occurs. + TileAcc a_l0; + TASSIGN(a_l0, 0); + + // ── Vec-side UB tile declarations ──────────────────────────────────── + // These tiles live in UB (Unified Buffer, the Vec engine's SRAM scratchpad). + // Each TASSIGN binds a tile handle to a fixed UB byte offset (our manual alloc). + // Vec-side UB tiles for gating computation + UbND g_ub; + TASSIGN(g_ub, GUbAddr); + UbND beta_ub_half; + TASSIGN(beta_ub_half, BetaHalfUbAddr); + UbND beta_ub; + TASSIGN(beta_ub, BetaUbAddr); + UbND g_v_ub; + TASSIGN(g_v_ub, GvUbAddr); + UbND a_ub; + TASSIGN(a_ub, AUbAddr); + UbND g_r_ub; + TASSIGN(g_r_ub, GRUbAddr); + UbND g_c_ub; + TASSIGN(g_c_ub, GCUbAddr); + UbND msk_ub; + TASSIGN(msk_ub, MskUbAddr); + UbND g_r_2d_ub; + TASSIGN(g_r_2d_ub, GR2dUbAddr); + UbND g_c_2d_ub; + TASSIGN(g_c_2d_ub, GC2dUbAddr); + UbND coeff_ub; + TASSIGN(coeff_ub, CoeffUbAddr); + UbND a_ub_half; + TASSIGN(a_ub_half, AUbHalfAddr); + + // ======================================================================== + // CUBE PHASE: Compute KK^T = K @ K^T for each chunk via GEMM + // + // ── How GEMM works on NPU (the "Cube pipeline") ────────────────────── + // The matrix multiply pipeline has 3 stages: + // Step 1: TLOAD loads data from GM → L1 (MTE2 pipe) + // Step 2: TEXTRACT copies sub-blocks from L1 → L0A/L0B (MTE1 pipe) + // L0A holds the left operand, L0B holds the right operand + // Step 3: TMATMUL multiplies L0A × L0B → L0C accumulator (M pipe) + // + // For K @ K^T: (numpy: KK_T = K @ K.T) + // Left operand: K [C×D] loaded into L1 in NZ format + // Right operand: K^T — same data, but we TRESHAPE to ZN format + // (TRESHAPE is FREE — it just reinterprets the fractal layout as transposed) + // Result: KK^T [C×C] in L0C (float32 accumulator, even though inputs are fp16) + // ======================================================================== + // __DAV_C220_CUBE__: This code only compiles for the Cube core. + // On NPU, Cube and Vec are separate compilation targets (like two different GPUs). +#if defined(__DAV_C220_CUBE__) + // Outer loop: iterate over all (sequence, head) work items assigned to this core + for (int64_t work_idx = 0; + work_idx < (total_work + block_num - 1) / block_num; ++work_idx) { + int64_t pid = work_idx * static_cast(block_num) + + static_cast(cid); + if (pid >= total_work) continue; + + // Map linear work index → (sequence, head) pair + int32_t head_idx = static_cast(pid % NumHeads); + int64_t seq_idx = pid / NumHeads; + + // Resolve sequence boundaries: cu_seqlens for variable-length, else fixed stride + int64_t bos, slen; + if (cu_seqlens != nullptr) { + // Variable-length sequences (packed tensor): cu_seqlens = [0, len0, len0+len1, ...] + bos = static_cast(cu_seqlens[seq_idx]); + slen = static_cast(cu_seqlens[seq_idx + 1]) - bos; + } else { + // Fixed-length sequences: each is seq_len tokens starting at seq_idx*seq_len + bos = seq_idx * seq_len; + slen = seq_len; + } + // Ceiling division: how many ChunkSize-sized chunks cover this sequence + int64_t num_chunks = (slen + ChunkSize - 1) / ChunkSize; + + // ── Double-buffering via workspace slots ────────────────────────── + // slot = ci & 1: alternates between 0 and 1 each chunk iteration. + // Cube writes KK^T to workspace[slot], then signals Vec. + // While Vec processes slot[0], Cube can write slot[1] (next chunk). + // This overlaps Cube computation with Vec computation for pipelining. + for (int64_t ci = 0; ci < num_chunks; ++ci) { + int32_t slot = static_cast(ci & 1); + // Wait for Vec to finish reading the previous KK^T from this slot + wait_flag_dev(2 + slot); + pipe_barrier(PIPE_ALL); + + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast( + remaining < ChunkSize ? remaining : ChunkSize); + + // BSND key layout [Seq, Hg, D]: token stride Hg * D (see BSND_QK_STRIDE). + // Value head head_idx maps to head_g = head_idx / GROUP for shared K rows. + int32_t head_g = head_idx / GROUP; + int64_t k_offset = + ((bos + chunk_start) * static_cast(NumKeyHeads) + + static_cast(head_g)) * + static_cast(HiddenSize); + + // ── Load K chunk from GM → L1 (MTE2 pipe) ────────────────────── + // DYNAMIC shape: valid_rows may be < ChunkSize for the last chunk. + // GlobalTensor describes the GM layout with strides (BSND interleaved). + // TLOAD triggers the MTE2 DMA engine to copy from GM (HBM) → L1 (on-chip cache). + // If the chunk is partial, TFILLPAD zero-fills the padding region + // so the GEMM doesn't produce garbage from uninitialized memory. + { + L1Mat _l1(valid_rows, HiddenSize); + TASSIGN(_l1, 0); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = valid_rows; _gs.shape[4] = HiddenSize; + GlobalTensor> + _gm(K_handle + k_offset, _gs); + TLOAD(_l1, _gm); + if (valid_rows != ChunkSize) TFILLPAD(_l1, _l1); + } + + // ── GEMM: KK^T = K @ K^T (L1→L0A/L0B→L0C) ──────────────────── + // K is [C×D] in L1 NZ; K^T obtained via ZN reshape of same tile. + // + // ── WAR (Write-After-Read) synchronization ──────────────────────── + // Before TEXTRACT (MTE1) writes new data to L0A/L0B, we must ensure: + // 1. MTE2 has finished loading L1 (MTE2→MTE1 sync) + // 2. Cube M pipe has finished reading previous L0A/L0B data (M→MTE1 sync) + // After TEXTRACT, before TMATMUL: + // 3. MTE1→M sync ensures L0A/L0B data is ready for the matrix engine + // After TMATMUL completes: + // 4. M→FIX sync ensures the L0C accumulator can be read + // This is like ensuring a producer-consumer chain is properly ordered. + // WAR sync: MTE2→MTE1, M→MTE1 before extract; MTE1→M before matmul. + { + TileLeft _l0a; + TileRight _l0b; + TASSIGN(_l0a, 0x0); + TASSIGN(_l0b, 0x0); + auto _we = EVENT_ID1; + set_flag(PIPE_MTE2, PIPE_MTE1, _we); + wait_flag(PIPE_MTE2, PIPE_MTE1, _we); + set_flag(PIPE_M, PIPE_MTE1, _we); + wait_flag(PIPE_M, PIPE_MTE1, _we); + // Left operand: K in NZ format, extract directly to L0A + TEXTRACT(_l0a, k_l1, 0, 0); + // Right operand: K^T via ZN reshape of same L1 tile, extract to L0B + L1MatZN _bzn; + TRESHAPE(_bzn, k_l1); + TEXTRACT(_l0b, _bzn, 0, 0); + set_flag(PIPE_MTE1, PIPE_M, _we); + wait_flag(PIPE_MTE1, PIPE_M, _we); + TMATMUL(a_l0, _l0a, _l0b); + set_flag(PIPE_MTE1, PIPE_MTE2, _we); + wait_flag(PIPE_MTE1, PIPE_MTE2, _we); + set_flag(PIPE_M, PIPE_FIX, _we); + wait_flag(PIPE_M, PIPE_FIX, _we); + } + + // ── Store KK^T from L0C → workspace GM (with fp32→fp16 cast) ─── + { + TileAcc _l0(ChunkSize, ChunkSize); + TASSIGN(_l0, 0); + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = ChunkSize; _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_handle + + (static_cast(cid) * 2 + slot) * ChunkSquare, + _gs); + TSTORE(_gm, _l0); + } + + // ── Cross-core synchronization (Cube → Vec) ────────────────────── + // ffts_cross_core_sync(pipe, config): Signal across physical cores. + // Unlike set_flag/wait_flag (which sync pipes within ONE core), this syncs + // between the Cube core and Vec core (they are separate hardware units). + // + // Config encoding: 1 | (mode << 4) | (flag_id << 8) + // mode=2: broadcast to all cores on same block + // flag_id: which flag to set (0,1,2,3...) + // + // The receiving side calls wait_flag_dev(flag_id) to wait for this signal. + // + // In this kernel: + // Cube sets flag 0/1 → Vec waits on wait_flag_dev(0/1) (KK^T ready) + // Vec sets flag 2/3 → Cube waits on wait_flag_dev(2/3) (workspace free) + // + // Signal Vec that this slot's KK^T is ready + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (slot << 8)); + } + } +#endif + + // ======================================================================== + // VEC PHASE: Apply gating and causal mask to KK^T + // coeff[i,j] = exp(min(g[i]+log(β[i]) - g[j], 0)) + // A[i,j] = KK^T[i,j] · coeff[i,j] · mask[i,j] + // Each sub-block (vid=0,1) handles HalfChunk rows of the C×C matrix. + // + // ── Gating computation (numpy pseudocode) ───────────────────────────── + // # For each sub-block's C/2 rows (vid selects upper or lower half): + // g_row = g_sum[row_offset:row_offset+C/2] # this sub-block's gates + // g_v = g_row + np.log(beta[row_offset:row_offset+C/2]) # combined gate+bias + // g_col = g_sum[0:C] # full chunk gates + // + // # Broadcast to 2D matrices for element-wise ops: + // g_r_2d = np.tile(g_v.reshape(-1, 1), (1, C)) # TROWEXPAND + // g_c_2d = np.tile(g_col.reshape(1, -1), (C/2, 1)) # TCOLEXPAND + // + // # Gating coefficient: exponential decay, clamped to ≤ 1 + // coeff = np.exp(np.minimum(g_r_2d - g_c_2d, 0)) # TSUB → TMINS → TEXP + // + // # Final: A = KK_T * coeff * causal_mask + // A = KK_T[my_rows] * coeff * mask[my_rows] # TMUL × 2 + // ======================================================================== + // __DAV_C220_VEC__: This code only compiles for the Vec core. +#if defined(__DAV_C220_VEC__) + // set_mask_norm / set_vector_mask: configure the SIMD mask for Vec ops. + // (-1, -1) means "all lanes active" — process every element. + // (Like CUDA's __activemask() returning all 1s for a full warp.) + set_mask_norm(); + set_vector_mask(-1, -1); + + // ── Load causal mask (lower triangular) once, reused across all chunks ── + // vid=0 loads the top half (rows 0..C/2-1), vid=1 loads the bottom half. + // The mask is [C×C] in GM; each sub-block loads its [C/2×C] portion. + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = HalfChunk; _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + Msk_handle + + static_cast(vid) * HalfChunk * ChunkSize, + _gs); + UbND _ld(HalfChunk, ChunkSize); + TASSIGN(_ld, MskUbAddr); + TLOAD(_ld, _gm); + } + // MTE2→V sync: ensure mask DMA is complete before Vec reads it + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Initial cross-core sync: release both workspace slots so Cube can start. + // Vec tells Cube "slots 0 and 1 are free" by setting flags 2 and 3. + // Without this, Cube would hang on wait_flag_dev(2/3) at the first iteration. + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (2 << 8)); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (3 << 8)); + + for (int64_t work_idx = 0; + work_idx < (total_work + block_num - 1) / block_num; ++work_idx) { + int64_t pid = work_idx * static_cast(block_num) + + static_cast(cid); + if (pid >= total_work) continue; + + int32_t head_idx = static_cast(pid % NumHeads); + int64_t seq_idx = pid / NumHeads; + + int64_t bos, slen; + if (cu_seqlens != nullptr) { + bos = static_cast(cu_seqlens[seq_idx]); + slen = static_cast(cu_seqlens[seq_idx + 1]) - bos; + } else { + bos = seq_idx * seq_len; + slen = seq_len; + } + int64_t num_chunks = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < num_chunks; ++ci) { + int32_t slot = static_cast(ci & 1); + + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast( + remaining < ChunkSize ? remaining : ChunkSize); + // row_offset: which half of the C×C matrix this sub-block handles + // vid=0 → rows [0, C/2), vid=1 → rows [C/2, C) + int32_t row_offset = static_cast(vid) * HalfChunk; + // local_valid: how many rows in this sub-block are real (not padding) + // Handles the case where the last chunk has fewer than C valid rows + int32_t local_valid = + valid_rows > row_offset + ? (valid_rows - row_offset < HalfChunk + ? valid_rows - row_offset + : HalfChunk) + : 0; + + if (local_valid > 0) { + // ── Load G (full chunk, 1×C) and Beta (sub-block rows, 1×HalfC) ── + // G is [H, total_tokens] float — contiguous per head + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = 1; _gs.shape[4] = valid_rows; + GlobalTensor> _gm( + G_handle + static_cast(head_idx) * total_tokens + + (bos + chunk_start), + _gs); + UbND _ld(1, valid_rows); + TASSIGN(_ld, GUbAddr); + TLOAD(_ld, _gm); + if (valid_rows != ChunkSize) { + UbND _pd; + TASSIGN(_pd, GUbAddr); + TFILLPAD_INPLACE(_pd, _ld); + } + } + + // Beta is [H, total_tokens] half — contiguous per head + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = 1; _gs.shape[4] = local_valid; + GlobalTensor> _gm( + Beta_handle + static_cast(head_idx) * total_tokens + + (bos + chunk_start + row_offset), + _gs); + UbND _ld(1, local_valid); + TASSIGN(_ld, BetaHalfUbAddr); + TLOAD(_ld, _gm); + if (local_valid != HalfChunk) { + UbND _pd; + TASSIGN(_pd, BetaHalfUbAddr); + TFILLPAD_INPLACE(_pd, _ld); + } + } + } + + // Wait for Cube to finish writing KK^T for this slot + wait_flag_dev(slot); + pipe_barrier(PIPE_ALL); + + if (local_valid > 0) { + // ── Compute gating coefficient ──────────────────────────────── + // Step 1: Convert beta from fp16→fp32 for precision + // Step 2: g_v[i] = g[row_offset+i] + log(β[i]) — combined row gate + // Step 3: Broadcast g_v (rows) and g (cols) to 2D matrices + // Step 4: coeff = exp(min(g_v_2d - g_2d, 0)) — clamped exponential gating + // g_v[i] = g[row_offset+i] + log(β[i]) — combined row gate + TCVT(beta_ub, beta_ub_half, pto::RoundMode::CAST_NONE); + // g_ub_temp points to the sub-block's portion of g within the full g_ub. + // row_offset * sizeof(float) is the byte offset into the g_ub tile. + UbND + g_ub_temp; + TASSIGN(g_ub_temp, + GUbAddr + row_offset * + static_cast(sizeof(float))); + TMOV(g_v_ub, g_ub_temp); // g_v = g[row_offset:row_offset+C/2] + pipe_barrier(PIPE_V); // Wait for TMOV to complete + + TLOG(beta_ub, beta_ub); // beta_ub = log(beta) in-place + pipe_barrier(PIPE_V); + TADD(g_v_ub, g_v_ub, beta_ub); // g_v = g_sub + log(beta) — the combined gate + pipe_barrier(PIPE_V); + TMOV(g_r_ub, g_v_ub); // Copy to g_r for row-broadcast + TMOV(g_c_ub, g_ub); // Copy full g to g_c for col-broadcast + pipe_barrier(PIPE_V); + + // Broadcast g_v to rows, g to columns → 2D gating matrix + // coeff[i,j] = exp(min(g_v[i] - g[j], 0)) + // + // g_r_ub_temp is a column-major (DN) alias of g_r_ub, required because + // TROWEXPAND expects its source in column-major layout. + UbDN g_r_ub_temp; + TASSIGN(g_r_ub_temp, GRUbAddr); + TROWEXPAND(g_r_2d_ub, g_r_ub_temp); // g_r_2d[i,j] = g_v[i] for all j + TCOLEXPAND(g_c_2d_ub, g_c_ub); // g_c_2d[i,j] = g[j] for all i + pipe_barrier(PIPE_V); + TSUB(coeff_ub, g_r_2d_ub, g_c_2d_ub); // coeff[i,j] = g_v[i] - g[j] + pipe_barrier(PIPE_V); + TMINS(coeff_ub, coeff_ub, 0.0f); // clamp to ≤ 0 (coeff will be ≤ 1 after exp) + pipe_barrier(PIPE_V); + TEXP(coeff_ub, coeff_ub); // coeff = exp(clamped_diff) ∈ (0, 1] + + // V→MTE2 sync: ensure gating computation is done before we start + // loading KK^T from workspace (we need coeff ready for the multiply later, + // and we want to overlap the DMA load with the preceding Vec work). + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + + // ── Load KK^T sub-block from workspace (fp16) ──────────────── + // workspace layout: [core_id * 2 + slot][C×C], we load our sub-block's + // [C/2×C] portion (offset by vid * HalfChunk * ChunkSize elements). + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = HalfChunk; _gs.shape[4] = ChunkSize; + GlobalTensor> _gm( + workspace_handle + + (static_cast(cid) * 2 + slot) * ChunkSquare + + static_cast(vid) * HalfChunk * ChunkSize, + _gs); + UbND _ld(HalfChunk, ChunkSize); + TASSIGN(_ld, AUbHalfAddr); + TLOAD(_ld, _gm); + } + + // MTE2→V sync: KK^T data is now in UB, safe for Vec to read + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // ── Apply gating and mask: A = KK^T · coeff · mask ─────────── + // 1. Convert KK^T from fp16 → fp32 (Cube stored it as fp16 to save GM bandwidth) + TCVT(a_ub, a_ub_half, pto::RoundMode::CAST_NONE); + // 2. Element-wise multiply by gating coefficient + TMUL(a_ub, a_ub, coeff_ub); + // 3. Element-wise multiply by causal mask (lower triangular, zeros above diagonal) + TMUL(a_ub, a_ub, msk_ub); + // 4. Convert result back to fp16 for output + TCVT(a_ub_half, a_ub, pto::RoundMode::CAST_NONE); + + // V→MTE3 sync: Vec computation done, safe for DMA store to begin + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + // ── Store A sub-block to output GM ──────────────────────────── + // Output A is in BSND layout: [total_tokens, NumHeads, ChunkSize] + // Each row of A corresponds to one token's attention weights for this head. + // Stride between consecutive tokens = NumHeads * ChunkSize (BSND interleaved). + int64_t a_gm_offset = + ((bos + chunk_start + row_offset) * NumHeads + + head_idx) * + static_cast(ChunkSize); + + { + Shape<1, 1, 1, DYNAMIC, DYNAMIC> _gs; + _gs.shape[3] = local_valid; _gs.shape[4] = ChunkSize; + GlobalTensor> _gm(A_handle + a_gm_offset, _gs); + UbND _st(local_valid, ChunkSize); + TASSIGN(_st, AUbHalfAddr); + TSTORE(_gm, _st); + } + } + + pipe_barrier(PIPE_ALL); + // Signal Cube that this workspace slot is free for reuse. + // Flag (2+slot): slot 0 → flag 2, slot 1 → flag 3. + // Cube is waiting on wait_flag_dev(2+slot) before writing the next chunk. + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | ((2 + slot) << 8)); + } + } +#endif +} + +// ── NPU kernel entry point ──────────────────────────────────────────── +// extern "C" __global__ AICORE: NPU kernel entry point (like CUDA __global__). +// Parameters passed as uint8_t* and reinterpret_cast'd — standard NPU convention. +// The NPU runtime passes raw byte pointers; we cast them to typed pointers here. +// GDN_H, GDN_D, GDN_C are compile-time constants set by #define at the top. +extern "C" __global__ AICORE void launch_scaled_dot_kkt( + __gm__ uint8_t *K_handle, __gm__ uint8_t *Beta_handle, + __gm__ uint8_t *G_handle, __gm__ uint8_t *Msk_handle, + __gm__ uint8_t *workspace_handle, __gm__ uint8_t *A_handle, + __gm__ uint8_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, + int64_t total_tokens, + uint64_t ffts_addr) +{ + kkt_kernel( + reinterpret_cast<__gm__ half *>(K_handle), + reinterpret_cast<__gm__ half *>(Beta_handle), + reinterpret_cast<__gm__ float *>(G_handle), + reinterpret_cast<__gm__ float *>(Msk_handle), + reinterpret_cast<__gm__ half *>(workspace_handle), + reinterpret_cast<__gm__ half *>(A_handle), + reinterpret_cast<__gm__ int32_t *>(cu_seqlens), + batch_size, seq_len, total_tokens, ffts_addr); +} + +// ── Host-side launcher ──────────────────────────────────────────────── +// call_kernel(): Host-side launcher invoked from Python via ctypes. +// block_dim = number of AI cores (like CUDA grid size) +// <<>>: NPU kernel launch syntax +// - block_dim: how many AI cores to use (each runs kkt_kernel independently) +// - nullptr: no shared memory (NPU doesn't have CUDA-style shared mem) +// - stream: async execution stream (like CUDA streams) +// +// rtGetC2cCtrlAddr: Get the hardware address of the cross-core (Cube↔Vec) flag +// table. This address is passed to the kernel so it can call ffts_cross_core_sync. +extern "C" void call_kernel( + uint32_t block_dim, void *stream, + uint8_t *K_handle, uint8_t *Beta_handle, + uint8_t *G_handle, uint8_t *Msk_handle, + uint8_t *workspace_handle, uint8_t *A_handle, + uint8_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, + int64_t total_tokens) +{ + uint32_t fftsLen{0}; + uint64_t fftsAddr{0}; + rtGetC2cCtrlAddr(&fftsAddr, &fftsLen); + launch_scaled_dot_kkt<<>>( + K_handle, Beta_handle, G_handle, Msk_handle, + workspace_handle, A_handle, cu_seqlens, + batch_size, seq_len, total_tokens, fftsAddr); +} diff --git a/csrc/ascend/pto_chunk_gdn/tri_inverse_impl.cpp b/csrc/ascend/pto_chunk_gdn/tri_inverse_impl.cpp new file mode 100644 index 000000000..54a79b1cb --- /dev/null +++ b/csrc/ascend/pto_chunk_gdn/tri_inverse_impl.cpp @@ -0,0 +1,824 @@ +/** +Copyright (c) 2026 Huawei Technologies Co., Ltd. +All rights reserved. + +See LICENSE in the root of the software repository: +https://github.com/huawei-csl/pto-kernels/ +for the full License text. +*/ + +#ifndef MEMORY_BASE +#define MEMORY_BASE +#endif +#include + +#include "kernel_utils.h" + +#define GM_ADDR __gm__ uint8_t* // To avoid #include "kernel_operator.h" +using namespace pto; +using namespace kernel_utils; + +#define BSND_OFFSET(tile_id, N, S, D) \ + (((tile_id) / (N)) * (S) * (N) * (D) + ((tile_id) % (N)) * (D)) + +/* + * For aligned BSND, tile_id enumerates chunk-major then head-major and maps to + * a fixed-stride address inside the dense BSND tensor. + */ +AICORE inline uint32_t GetBSNDFixedTileOffset(uint32_t tile_id, + uint32_t num_bsnd_heads, + uint32_t matrix_size) { + return BSND_OFFSET(tile_id, num_bsnd_heads, matrix_size, matrix_size); +} + +/** + * @brief Struct containing starting address and size of a single tile + */ +struct BSNDVarlenTileInfo { + uint32_t bsnd_offset; /**< Contains the starting index in the global tensor */ + uint32_t valid_size; /**< This is the size (num_rows/cols) of the tile */ +}; + +/* + * For cu_seqlens-based varlen BSND, tile_id still enumerates chunk-major then + * head-major. We recover the owning sequence by scanning cu_seqlens and + * counting chunks per sequence. + */ +AICORE inline BSNDVarlenTileInfo GetBSNDVarlenTileInfoFromCuSeqlens( + uint32_t tile_id, uint32_t num_bsnd_heads, uint32_t matrix_size, + __gm__ int32_t* cu_seqlens) { + const uint32_t head_idx = tile_id % num_bsnd_heads; + const uint32_t chunk_idx = tile_id / num_bsnd_heads; + + uint32_t seq_start = static_cast(cu_seqlens[0]); + uint32_t accumulated_chunks = 0; + for (uint32_t seq_idx = 0;; ++seq_idx) { + const uint32_t seq_end = static_cast(cu_seqlens[seq_idx + 1]); + const uint32_t seq_len = seq_end - seq_start; + const uint32_t seq_num_chunks = CeilDiv(seq_len, matrix_size); + if (chunk_idx < accumulated_chunks + seq_num_chunks) { + const uint32_t local_chunk_idx = chunk_idx - accumulated_chunks; + const uint32_t row_start = seq_start + local_chunk_idx * matrix_size; + const uint32_t valid_size = + min(static_cast(seq_end - row_start), matrix_size); + return {row_start * num_bsnd_heads * matrix_size + head_idx * matrix_size, + valid_size}; + } + accumulated_chunks += seq_num_chunks; + seq_start = seq_end; + } +} + +/* + * @brief: Takes as input two matrices of size MatrixSize * MatrixSize each. + * The src matrix lies in L1, while the dst matrix lies either in L0A or L0B. + * This kernel copies only the diagonal blocks (fractals) of size FractalSize * + * FractalSize from the src matrix to the dst matrix. + * + * @tparam InputT Input data type (fp16). + * @tparam FractalSize Size of each fractal matrix (diagonal block). + * @tparam MatrixSize Size of the entire input/output matrices. + * @tparam SrcL1TileT The actual tile type of the src matrix. + * @tparam DstL0TileT The actual tile type of the dst matrix. + * + * @param src Tile in L1 memory. + * @param dst Tile in L0A or L0B memory. + */ +template +AICORE inline void CopyDiagonalFractalsL1ToL0(SrcL1TileT src, DstL0TileT dst) { + constexpr uint32_t NumFractals = MatrixSize / FractalSize; + constexpr bool is_left = + std::is_same_v>; + constexpr TileType LeftOrRight = is_left ? TileType::Left : TileType::Right; + constexpr SLayout InnerLayout = + is_left ? SLayout::RowMajor : SLayout::ColMajor; + + Tile + fractals[NumFractals]; + const std::uintptr_t starting_address = + reinterpret_cast(dst.data()); + for (uint32_t i = 0; i < NumFractals; ++i) { + TASSIGN(fractals[i], starting_address + i * FractalSize * + (MatrixSize + FractalSize) * + sizeof(InputT)); + TEXTRACT(fractals[i], src, i * FractalSize, i * FractalSize); + } +} + +/* + * @brief: Takes as input two matrices of size MatrixSize * MatrixSize each, + * and an integer block_size. The src matrix lies in L1, while the dst matrix + * either in L0A or L0B. This method copies some of the diagonal blocks from the + * input to the output as follows: + * - If dst is in L0A (left): copy even diagonal blocks 0, 2, 4, ... + * - If dst is in L0B (right): copy odd blocks 1, 3, 5, ... + * Important note: the dst matrix should be initialized to all-zeros before + * calling this method + * + * @tparam InputT Input data type (fp16). + * @tparam FractalSize Size of each fractal matrix (diagonal block). + * @tparam MatrixSize Size of the entire input/output matrices. + * @tparam SrcL1TileT The actual tile type of the src matrix. + * @tparam DstL0TileT The actual tile type of the dst matrix. + * + * @param src Tile in L1 memory. + * @param dst Tile in L0A or L0B memory. + * @param block_size Size of diagonal blocks. Needs: block_size >= FractalSize. + */ +template +AICORE inline void CopyOddOrEvenBlocksL1ToL0(SrcL1TileT src, DstL0TileT dst, + uint32_t block_size, + bool swap_parity = false) { + constexpr bool is_left = + std::is_same_v>; + constexpr TileType LeftOrRight = is_left ? TileType::Left : TileType::Right; + constexpr SLayout InnerLayout = + is_left ? SLayout::RowMajor : SLayout::ColMajor; + + // Default: left→even(0), right→odd(1). swap_parity flips this. + const uint32_t starting_block_index = (is_left ? 0u : 1u) ^ (swap_parity ? 1u : 0u); + + const uint32_t num_blocks = MatrixSize / block_size; + const uint32_t num_fractals_per_block = block_size / FractalSize; + + // might need fewer fractals if block_size < FractalSize + Tile + fractals[MatrixSize / FractalSize]; + + const std::uintptr_t starting_address = + reinterpret_cast(dst.data()); + for (uint32_t i = 0; i < num_fractals_per_block; ++i) { + for (uint32_t j = 0; j < num_fractals_per_block; ++j) { + for (uint32_t b = starting_block_index; b < num_blocks; b += 2) { + const uint32_t offset = + b * (MatrixSize + FractalSize) * block_size /* block_offset */ + + i * MatrixSize * FractalSize /* col_fractal_offset */ + + j * FractalSize * FractalSize /* row_fractal_offset */; + TASSIGN(fractals[b], starting_address + offset * sizeof(InputT)); + TEXTRACT(fractals[b], src, b * block_size + i * FractalSize, + b * block_size + j * FractalSize); + } + } + } +} + +/* + * @brief: Prepares Identity and Zeros matrix. + * + * @tparam TileL1AB The type of the input tiles in L1. + * @tparam TileL0A The type of the input tiles in L0A. + * @tparam TileL0B The type of the input tiles in L0B. + * @tparam TileL0C The type of the input tiles in L0C. + * + * @param I_neg_l1_tile Tile containing the -I (negative identity) matrix. + * @param Zero_l1_tile Tile to store the all-zero matrix. + * @param I_l1_tile Tile to store the identity matrix. + * @param a_l0_tile Tile in L0A for matmuls. + * @param b_l0_tile Tile in L0B for matmuls. + * @param c_l0_tile Tile in L0C for matmuls. + */ +template +AICORE inline void PrepareAuxiliaryMatrices( + TileL1AB I_neg_l1_tile, TileL1AB Zero_l1_tile, TileL1AB I_l1_tile, + TileL0A a_l0_tile, TileL0B b_l0_tile, TileL0C c_l0_tile) { + TMOV(a_l0_tile, I_neg_l1_tile); // a_l0 initialized with I_neg + TMOV(b_l0_tile, I_neg_l1_tile); // b_l0 initialized with I_neg + set_flag(PIPE_MTE1, PIPE_M, static_cast(0)); + wait_flag(PIPE_MTE1, PIPE_M, static_cast(0)); + + TMATMUL(c_l0_tile, a_l0_tile, b_l0_tile); // c_l0 contains I + set_flag(PIPE_M, PIPE_FIX, static_cast(0)); + wait_flag(PIPE_M, PIPE_FIX, static_cast(0)); + + TMOV(I_l1_tile, c_l0_tile); // I_l1 now contains I + set_flag(PIPE_FIX, PIPE_MTE1, static_cast(0)); + wait_flag(PIPE_FIX, PIPE_MTE1, static_cast(0)); + + TMOV(b_l0_tile, I_l1_tile); // b_l0 contains I + set_flag(PIPE_MTE1, PIPE_M, static_cast(0)); + wait_flag(PIPE_MTE1, PIPE_M, static_cast(0)); + + TMATMUL_ACC(c_l0_tile, c_l0_tile, a_l0_tile, + b_l0_tile); // c_l0 contains zeros + set_flag(PIPE_M, PIPE_FIX, static_cast(0)); + wait_flag(PIPE_M, PIPE_FIX, static_cast(0)); + + TMOV(Zero_l1_tile, c_l0_tile); // Zeros_l1 now contains zeros + set_flag(PIPE_FIX, PIPE_MTE1, static_cast(0)); + wait_flag(PIPE_FIX, PIPE_MTE1, static_cast(0)); +} + +/* + * @brief: Inverts a single matrix / tile of the global tensor. + * The first part of the algorithm inverts the FractalSize * FractalSize + * diagonal blocks of the input matrix (inv_trick part). The second phase + * assembles the partial inverses using the cube unig (recursive part). + * + * @tparam InputT The type of the input elements. + * @tparam TileL1AB The type of the input tiles in L1. + * @tparam TileL0A The type of the input tiles in L0A. + * @tparam TileL0B The type of the input tiles in L0B. + * @tparam TileL0C The type of the input tiles in L0C. + * @tparam MatrixSize Size of the entire input/output matrices. + * @tparam FractalSize Size of matrix fractals. + * @tparam NumTilesPerCubeIter How many matrices to load and invert in a single + * cube iteration. + * + * @param X_l1_tile Tile in L1 used for intermediate computations. + * @param I_l1_tile Tile containing the identity matrix. + * @param I_neg_l1_tile Tile containing the negative identity matrix. + * @param M_neg_l1_tile Tile containing the negative input matrix. + * @param Zero_l1_tile Tile containing the all-zero matrix. + * @param Y_l1_tile Tile in L1 used for intermediate computations. + * @param a_l0_tile* Array of two tiles in L0A (for double-buffering). + * @param b_l0_tile* Array of two tiles in L0B (for double-buffering). + * @param c_l0_tile* Tile in L0C for matmuls. + * @param tile_id Index of the current tile (used for sync). + */ +template +AICORE inline void InvertSingleTile(TileL1AB X_l1_tile, TileL1AB I_l1_tile, + TileL1AB I_neg_l1_tile, + TileL1AB M_neg_l1_tile, + TileL1AB Zero_l1_tile, TileL1AB Y_l1_tile, + TileL0A* a_l0_tile, TileL0B* b_l0_tile, + TileL0C* c_l0_tile, + const uint32_t tile_id, + const bool swap_parity = false) { + const event_t event_0 = static_cast(tile_id); + const event_t event_1 = static_cast(tile_id + NumTilesPerCubeIter); + + TMOV(b_l0_tile[0], Y_l1_tile); // b_l0[0] contains M + TMOV(a_l0_tile[0], I_neg_l1_tile); // a_l0[0] contains I_neg + set_flag(PIPE_MTE1, PIPE_M, event_0); + TMOV(a_l0_tile[1], Zero_l1_tile); + TMOV(b_l0_tile[1], Zero_l1_tile); + set_flag(PIPE_MTE1, PIPE_M, event_1); + wait_flag(PIPE_MTE1, PIPE_M, event_1); + set_flag(PIPE_M, PIPE_MTE1, event_1); + wait_flag(PIPE_M, PIPE_MTE1, event_1); + CopyDiagonalFractalsL1ToL0( + Y_l1_tile, a_l0_tile[1]); // a_l0[1] = diag_fractals(M) + CopyDiagonalFractalsL1ToL0( + Y_l1_tile, b_l0_tile[1]); // b_l0[1] = diag_fractals(M) + set_flag(PIPE_MTE1, PIPE_M, event_1); + + /* First Matmul: event_0 */ + wait_flag(PIPE_MTE1, PIPE_M, event_0); + TMATMUL(c_l0_tile[0], a_l0_tile[0], b_l0_tile[0]); // c_l0[0] contains M_neg + set_flag(PIPE_M, PIPE_FIX, event_0); + set_flag(PIPE_M, PIPE_MTE1, event_0); + + wait_flag(PIPE_M, PIPE_FIX, event_0); + TMOV(M_neg_l1_tile, c_l0_tile[0]); // M_neg_l1 now contains M_neg + set_flag(PIPE_FIX, PIPE_M, event_0); + + /* Second Matmul: event_1 */ + wait_flag(PIPE_MTE1, PIPE_M, event_1); + set_flag(PIPE_MTE1, PIPE_M, event_1); + TMATMUL(c_l0_tile[1], a_l0_tile[1], + b_l0_tile[1]); // c_l0[1] contains diag_fractals(M)^2 + set_flag(PIPE_M, PIPE_FIX, event_1); + wait_flag(PIPE_M, PIPE_FIX, event_1); + TMOV(Y_l1_tile, + c_l0_tile[1]); // Y_l1 now contains diag_fractals(M)^2 + set_flag(PIPE_FIX, PIPE_M, event_1); + wait_flag(PIPE_FIX, PIPE_M, event_1); + + /* Third Matmul: event_0*/ + wait_flag(PIPE_M, PIPE_MTE1, event_0); + TMOV(b_l0_tile[0], I_neg_l1_tile); // b_l0[0] contains I_neg + TMOV(a_l0_tile[0], I_neg_l1_tile); // a_l0[0] contains I_neg + set_flag(PIPE_MTE1, PIPE_M, event_0); + + wait_flag(PIPE_MTE1, PIPE_M, event_0); + wait_flag(PIPE_FIX, PIPE_M, event_0); + wait_flag(PIPE_MTE1, PIPE_M, event_1); + TMATMUL(c_l0_tile[0], a_l0_tile[1], + b_l0_tile[0]); // c_l0[0] = diag_fractals(M_neg) + set_flag(PIPE_M, PIPE_FIX, event_0); + wait_flag(PIPE_M, PIPE_FIX, event_0); + set_flag(PIPE_FIX, PIPE_M, event_0); + wait_flag(PIPE_FIX, PIPE_M, event_0); + + TMATMUL_ACC(c_l0_tile[0], c_l0_tile[0], a_l0_tile[0], + b_l0_tile[0]); // c_l0[0] has I-diag_fractals(M) + set_flag(PIPE_M, PIPE_FIX, event_1); + wait_flag(PIPE_M, PIPE_FIX, event_1); + TMOV(X_l1_tile, c_l0_tile[0]); // X_l1 now contains I-diag_fractals(M) + + /* + * Inv Trick part: + * X = I - M + * Y = M + * block_size = 1 + * while block_size < FractalSize / 2: + * Y = Y @ Y + * X = X + X @ Y + * block_size *= 2 + */ + set_flag(PIPE_FIX, PIPE_M, event_0); // store c + set_flag(PIPE_M, PIPE_MTE1, event_0); // load matrices for matmuls + set_flag(PIPE_FIX, PIPE_MTE1, event_0); + set_flag(PIPE_FIX, PIPE_M, event_1); // only for update Y + set_flag(PIPE_M, PIPE_MTE1, event_1); // only for update Y + set_flag(PIPE_FIX, PIPE_MTE1, event_1); // only for update Y + for (uint32_t block_size = 1; block_size < FractalSize / 2; block_size *= 2) { + wait_flag(PIPE_M, PIPE_MTE1, event_0); + TMOV(b_l0_tile[0], I_l1_tile); + wait_flag(PIPE_FIX, PIPE_MTE1, event_0); + TMOV(a_l0_tile[0], X_l1_tile); + set_flag(PIPE_MTE1, PIPE_M, event_0); + + wait_flag(PIPE_FIX, PIPE_MTE1, event_1); + TMOV(b_l0_tile[1], Y_l1_tile); + set_flag(PIPE_MTE1, PIPE_M, event_1); + + wait_flag(PIPE_FIX, PIPE_M, event_0); // from previous iter + wait_flag(PIPE_MTE1, PIPE_M, event_0); // from loading a_l0[0], b_l0[0] + TMATMUL(c_l0_tile[0], a_l0_tile[0], b_l0_tile[0]); // c_l0[0] contains X + set_flag(PIPE_M, PIPE_FIX, event_0); + wait_flag(PIPE_M, PIPE_FIX, event_0); + set_flag(PIPE_FIX, PIPE_M, event_0); + wait_flag(PIPE_FIX, PIPE_M, event_0); + + if (block_size < FractalSize / 4) { // Update Y except in last iteration + wait_flag(PIPE_M, PIPE_MTE1, event_1); // from previous iter + TMOV(a_l0_tile[1], Y_l1_tile); + wait_flag(PIPE_MTE1, PIPE_M, event_1); + set_flag(PIPE_MTE1, PIPE_M, event_1); + + wait_flag(PIPE_MTE1, PIPE_M, event_1); + wait_flag(PIPE_FIX, PIPE_M, event_1); // from previous iter + TMATMUL(c_l0_tile[1], a_l0_tile[1], b_l0_tile[1]); + set_flag(PIPE_M, PIPE_MTE1, event_1); // for next iter + set_flag(PIPE_M, PIPE_FIX, event_1); + set_flag(PIPE_MTE1, PIPE_M, event_1); + + wait_flag(PIPE_M, PIPE_FIX, event_1); + TMOV(Y_l1_tile, c_l0_tile[1]); + set_flag(PIPE_FIX, PIPE_M, event_1); // for next iter + } + set_flag(PIPE_FIX, PIPE_MTE1, event_1); // for next iter + + wait_flag(PIPE_MTE1, PIPE_M, event_1); + TMATMUL_ACC(c_l0_tile[0], c_l0_tile[0], a_l0_tile[0], + b_l0_tile[1]); // c_l0[0] has X + X @ Y + set_flag(PIPE_M, PIPE_MTE1, event_0); + set_flag(PIPE_M, PIPE_FIX, event_0); + + wait_flag(PIPE_M, PIPE_FIX, event_0); + TMOV(X_l1_tile, c_l0_tile[0]); + set_flag(PIPE_FIX, PIPE_M, event_0); // for next iter + set_flag(PIPE_FIX, PIPE_MTE1, event_0); // for next iter + } + wait_flag(PIPE_FIX, PIPE_MTE1, event_1); // only for update Y + wait_flag(PIPE_M, PIPE_MTE1, event_1); // only for update Y + wait_flag(PIPE_FIX, PIPE_M, event_1); // only for update Y + wait_flag(PIPE_FIX, PIPE_MTE1, event_0); + wait_flag(PIPE_M, PIPE_MTE1, event_0); + wait_flag(PIPE_FIX, PIPE_M, event_0); + + /* + * Unrolled recursion part: + * Upper-tri (swap_parity=false): + * LX = even_blocks(X), RX = odd_blocks(X) + * Y = LX @ (-M) + I, X = Y @ RX + LX + * Lower-tri (swap_parity=true): + * RX = even→L0A(odd via swap), LX = odd→L0B(even via swap) + * Y = RX @ (-M) + I, X = Y @ LX + RX + */ + TMOV(b_l0_tile[1], M_neg_l1_tile); // b_l0[1] contains M_neg + TMOV(a_l0_tile[0], I_l1_tile); // a_l0[0] contains I + + if constexpr (MatrixSize > FractalSize) { + set_flag(PIPE_FIX, PIPE_M, event_1); + } + set_flag(PIPE_M, PIPE_MTE1, event_1); + set_flag(PIPE_M, PIPE_MTE1, event_0); + set_flag(PIPE_FIX, PIPE_MTE1, event_1); + set_flag(PIPE_FIX, PIPE_M, event_0); + for (uint32_t block_size = FractalSize; block_size < MatrixSize; + block_size *= 2) { + wait_flag(PIPE_M, PIPE_MTE1, event_0); // Wait for last iter a_l0[1] + TMOV(a_l0_tile[1], Zero_l1_tile); + + wait_flag(PIPE_M, PIPE_MTE1, event_1); + TMOV(b_l0_tile[0], I_l1_tile); + set_flag(PIPE_MTE1, PIPE_M, event_0); + + wait_flag(PIPE_FIX, PIPE_MTE1, event_1); // Wait to write last X + CopyOddOrEvenBlocksL1ToL0( + X_l1_tile, a_l0_tile[1], block_size, swap_parity); // a_l0[1]: even(LX) or odd(RX) + set_flag(PIPE_MTE1, PIPE_M, event_1); + + wait_flag(PIPE_MTE1, PIPE_M, event_0); + wait_flag(PIPE_FIX, PIPE_M, event_0); // Wait c_l0[0] from previous iter + TMATMUL(c_l0_tile[0], a_l0_tile[0], b_l0_tile[0]); // c_l0[0] has I + + wait_flag(PIPE_MTE1, PIPE_M, event_1); + wait_flag(PIPE_FIX, PIPE_M, event_1); // Wait c_l0[1] from previous iter + TMATMUL(c_l0_tile[1], a_l0_tile[1], b_l0_tile[0]); // c_l0[1] contains LX + set_flag(PIPE_M, PIPE_MTE1, event_1); // allow to load RX on b_l0[0] + + TMATMUL_ACC(c_l0_tile[0], c_l0_tile[0], a_l0_tile[1], + b_l0_tile[1]); // c_l0[0] <- LX * M_neg + I + set_flag(PIPE_M, PIPE_FIX, event_0); + set_flag(PIPE_M, PIPE_MTE1, event_0); + + wait_flag(PIPE_M, PIPE_FIX, event_0); + TMOV(Y_l1_tile, c_l0_tile[0]); // Y_l1 contains LX * M_neg + I + set_flag(PIPE_FIX, PIPE_MTE1, event_0); + set_flag(PIPE_FIX, PIPE_M, event_0); + + /* Load complementary blocks of X in L0B */ + wait_flag(PIPE_M, PIPE_MTE1, event_1); + TMOV(b_l0_tile[0], Zero_l1_tile); + CopyOddOrEvenBlocksL1ToL0( + X_l1_tile, b_l0_tile[0], block_size, swap_parity); // b_l0[0]: odd(RX) or even(LX) + + wait_flag(PIPE_M, PIPE_MTE1, event_0); // Wait for previous use of a_l0[1] + wait_flag(PIPE_FIX, PIPE_MTE1, event_0); // Wait for Y_l1 + TMOV(a_l0_tile[1], Y_l1_tile); // a_l0[1] contains LX * M_neg + I + set_flag(PIPE_MTE1, PIPE_M, event_0); + + wait_flag(PIPE_MTE1, PIPE_M, event_0); + TMATMUL_ACC(c_l0_tile[1], c_l0_tile[1], a_l0_tile[1], b_l0_tile[0]); + set_flag(PIPE_M, PIPE_MTE1, event_0); // next iter can read on a_l0[1] + set_flag(PIPE_M, PIPE_MTE1, event_1); // next iter can read on b_l0[0] + set_flag(PIPE_M, PIPE_FIX, event_0); + wait_flag(PIPE_M, PIPE_FIX, event_0); + + if (block_size < MatrixSize / 2) { // Update X_l1 except in last iteration + TMOV(X_l1_tile, c_l0_tile[1]); + set_flag(PIPE_FIX, PIPE_M, event_1); // release c_l0[1] for next iter + } + set_flag(PIPE_FIX, PIPE_MTE1, event_1); + } + wait_flag(PIPE_M, PIPE_MTE1, event_0); + wait_flag(PIPE_M, PIPE_MTE1, event_1); + wait_flag(PIPE_FIX, PIPE_M, event_0); + wait_flag(PIPE_FIX, PIPE_MTE1, event_1); // Write c_l0[1] to X_l1 +} + +/* + * @brief: Runs the main kernel (inverts all matrices in the tensor) + * + * @tparam InputT The type of the input elements. + * @tparam OutputT The type of the output elements. + * @tparam MatrixSize Size of the entire input/output matrices. + * @tparam NumTilesPerCubeIter How many matrices to load and invert in a single + * cube iteration. + * @tparam IsBSND If IsBSND is false, then the last two dimensions represent a + * 2D triangular matrix in row-major format, while the other dimensions are + * batch dimensions. If IsBSND is true, then the dimensions represent in order: + * B batch size, S sequence length (which is chunked in tiles of size D), N + * number of heads (equivalent to a second batch dimension for this kernel), and + * D chunk size. The inverse is over the dimensions S (chunked) and D, row-major + * within each tile. + * + * @param M_inv pointer to the global memory to store the final inverse. + * @param M Pointer to the global tensor matrix in global memory. + * @param I_neg Pointer to global memory that contains the negative identity. + * @param total_tiles The total number of matrices to invert. + * @param num_bsnd_heads The number of heads, only for BSND format. + */ +template +AICORE inline void TriInvRecUnrollKernel(__gm__ StoreT* M_inv, + __gm__ InputT* M, __gm__ InputT* I_neg, + uint32_t total_tiles, + uint32_t num_bsnd_heads = 0, + __gm__ int32_t* cu_seqlens = nullptr, + uint32_t is_lower = 0) { + /* Initializations */ + constexpr uint32_t TileLen = MatrixSize * MatrixSize; + constexpr uint32_t FractalSize = 16; // fractal size for half + constexpr uint32_t NumFractalsRowWise = MatrixSize / FractalSize; + constexpr uint32_t NumL0Buffers = 2; + + if (get_block_idx() * NumTilesPerCubeIter >= total_tiles) { + return; + } + + using GlobalTileShapeIn = + TileShape2D; + using GlobalTileStridesIn = typename std::conditional< + !IsBSND, BaseShape2D, + Stride<1, 1, 1, -1, 1>>::type; + using GlobalTileIn = + GlobalTensor; + using GlobalTileDynamicShape = Shape<1, 1, 1, DYNAMIC, DYNAMIC>; + using GlobalTileDynamicStride = Stride<1, 1, 1, DYNAMIC, 1>; + using GlobalTileDynamicIn = GlobalTensor; + using GlobalTileStridesINeg = + BaseShape2D; + using GlobalTileINeg = GlobalTensor; + + using GlobalTileShapeOut = + TileShape2D; + using GlobalTileStridesOut = typename std::conditional< + !IsBSND, BaseShape2D, + Stride<1, 1, 1, -1, 1>>::type; + using GlobalTileOut = GlobalTensor; + using GlobalTileDynamicOut = + GlobalTensor; + using TileL1AB = + Tile; + using TileL1ABDynamic = + Tile; + + // L0 Memory + using TileL0A = TileLeft; + using TileL0B = TileRight; + using TileL0C = TileAcc; + using TileL0CDynamic = + TileAcc; + + GlobalTileINeg I_neg_global_in(I_neg); + + TileL1AB X_l1_tile; + TileL1AB I_l1_tile; + TileL1AB I_neg_l1_tile; + TileL1AB M_neg_l1_tile; + TileL1AB Zero_l1_tile; + TileL1AB Y_l1_tile[NumTilesPerCubeIter]; + + TileL0A a_l0_tile[NumL0Buffers]; + TileL0B b_l0_tile[NumL0Buffers]; + TileL0C c_l0_tile[NumL0Buffers]; + + TASSIGN(I_l1_tile, 0x0); + TASSIGN(I_neg_l1_tile, 0x0 + TileLen * sizeof(InputT)); + TASSIGN(Zero_l1_tile, 0x0 + 2 * TileLen * sizeof(InputT)); + TASSIGN(M_neg_l1_tile, 0x0 + 3 * TileLen * sizeof(InputT)); + TASSIGN(X_l1_tile, 0x0 + 4 * TileLen * sizeof(InputT)); + for (uint32_t tile_id = 0; tile_id < NumTilesPerCubeIter; ++tile_id) { + TASSIGN(Y_l1_tile[tile_id], 0x0 + (5 + tile_id) * TileLen * sizeof(InputT)); + } + + for (uint32_t buffer_num = 0; buffer_num < NumL0Buffers; ++buffer_num) { + TASSIGN(a_l0_tile[buffer_num], 0x0 + buffer_num * TileLen * sizeof(InputT)); + TASSIGN(b_l0_tile[buffer_num], 0x0 + buffer_num * TileLen * sizeof(InputT)); + TASSIGN(c_l0_tile[buffer_num], + 0x0 + buffer_num * TileLen * sizeof(OutputT)); + } + TLOAD(I_neg_l1_tile, I_neg_global_in); + set_flag(PIPE_MTE2, PIPE_MTE1, static_cast(0)); + wait_flag(PIPE_MTE2, PIPE_MTE1, static_cast(0)); + + PrepareAuxiliaryMatrices( + I_neg_l1_tile, Zero_l1_tile, I_l1_tile, a_l0_tile[0], b_l0_tile[0], + c_l0_tile[0]); + + const uint32_t max_iters_per_aic = + CeilDiv(total_tiles, (uint32_t)(NumTilesPerCubeIter * get_block_num())); + + /* Main iteration - Compute all tiles */ + uint32_t bsnd_tile_offsets[NumTilesPerCubeIter] = {0}; + uint32_t bsnd_tile_valid_sizes[NumTilesPerCubeIter] = {0}; + uint32_t next_tile_id_that_waits_for_pipe_fix_pipe_m = 0; + set_flag(PIPE_FIX, PIPE_M, + static_cast(next_tile_id_that_waits_for_pipe_fix_pipe_m)); + for (uint32_t tile_id = 0; tile_id < NumTilesPerCubeIter; ++tile_id) { + set_flag(PIPE_M, PIPE_MTE2, static_cast(tile_id)); + } + for (uint32_t cube_iter = 0; cube_iter < max_iters_per_aic; ++cube_iter) { + const uint32_t global_index = + (cube_iter * get_block_num() + get_block_idx()) * NumTilesPerCubeIter; + if (global_index >= total_tiles) { + break; + } + for (uint32_t tile_id = 0; (tile_id < NumTilesPerCubeIter) && + (global_index + tile_id < total_tiles); + ++tile_id) { + if constexpr (IsBSND) { + const uint32_t global_tile_id = global_index + tile_id; + if (cu_seqlens != nullptr) { + const BSNDVarlenTileInfo tile_info = + GetBSNDVarlenTileInfoFromCuSeqlens(global_tile_id, num_bsnd_heads, + MatrixSize, cu_seqlens); + bsnd_tile_offsets[tile_id] = tile_info.bsnd_offset; + bsnd_tile_valid_sizes[tile_id] = tile_info.valid_size; + } else { + bsnd_tile_offsets[tile_id] = GetBSNDFixedTileOffset( + global_tile_id, num_bsnd_heads, MatrixSize); + bsnd_tile_valid_sizes[tile_id] = MatrixSize; + } + const uint32_t bsnd_offset = bsnd_tile_offsets[tile_id]; + const uint32_t valid_size = bsnd_tile_valid_sizes[tile_id]; + const int row_stride = static_cast(MatrixSize * num_bsnd_heads); + wait_flag(PIPE_M, PIPE_MTE2, static_cast(tile_id)); + if (valid_size < MatrixSize) { + TileL1ABDynamic Y_dyn_l1_tile(valid_size, valid_size); + TASSIGN(Y_dyn_l1_tile, + 0x0 + (5 + tile_id) * TileLen * sizeof(InputT)); + GlobalTileDynamicIn M_global_in_dyn( + M + bsnd_offset, + {1, 1, 1, static_cast(valid_size), + static_cast(valid_size)}, + {1, 1, 1, row_stride, 1}); + TLOAD(Y_dyn_l1_tile, M_global_in_dyn); + set_flag(PIPE_MTE2, PIPE_MTE1, static_cast(tile_id)); + wait_flag(PIPE_MTE2, PIPE_MTE1, static_cast(tile_id)); + TFILLPAD(Y_dyn_l1_tile, Y_dyn_l1_tile); + } else { + GlobalTileIn M_global_in(M + bsnd_offset, {}, {row_stride}); + TLOAD(Y_l1_tile[tile_id], M_global_in); + } + } else { + GlobalTileIn M_global_in(M + (global_index + tile_id) * TileLen); + wait_flag(PIPE_M, PIPE_MTE2, static_cast(tile_id)); + TLOAD(Y_l1_tile[tile_id], + M_global_in); // Copies NumTilesPerCubeIter tiles at once + } + set_flag(PIPE_MTE2, PIPE_MTE1, static_cast(tile_id)); + } + + constexpr uint32_t final_c_buffer_index = MatrixSize > FractalSize ? 1 : 0; + for (uint32_t tile_id = 0; (tile_id < NumTilesPerCubeIter) && + (global_index + tile_id < total_tiles); + ++tile_id) { + // Wait for previous cube iter to write result + wait_flag(PIPE_FIX, PIPE_M, static_cast(tile_id)); + // Wait for loading new matrices from GM + wait_flag(PIPE_MTE2, PIPE_MTE1, static_cast(tile_id)); + + InvertSingleTile( + X_l1_tile, I_l1_tile, I_neg_l1_tile, M_neg_l1_tile, Zero_l1_tile, + Y_l1_tile[tile_id], a_l0_tile, b_l0_tile, c_l0_tile, tile_id, + is_lower != 0); + + // Allow next cube_iter to proceed for this tile_id + set_flag(PIPE_M, PIPE_MTE2, static_cast(tile_id)); + + /* Store result */ + if constexpr (IsBSND) { + const uint32_t bsnd_offset = bsnd_tile_offsets[tile_id]; + const uint32_t valid_size = bsnd_tile_valid_sizes[tile_id]; + const int row_stride = static_cast(MatrixSize * num_bsnd_heads); + if (valid_size < MatrixSize) { + TileL0CDynamic c_l0_tail_tile(valid_size, valid_size); + TASSIGN(c_l0_tail_tile, + 0x0 + final_c_buffer_index * TileLen * sizeof(OutputT)); + GlobalTileDynamicOut M_inv_global_out_dyn( + M_inv + bsnd_offset, + {1, 1, 1, static_cast(valid_size), + static_cast(valid_size)}, + {1, 1, 1, row_stride, 1}); + TSTORE(M_inv_global_out_dyn, c_l0_tail_tile); + } else { + GlobalTileOut M_inv_global_out(M_inv + bsnd_offset, {}, {row_stride}); + TSTORE(M_inv_global_out, c_l0_tile[final_c_buffer_index]); + } + } else { + GlobalTileOut M_inv_global_out(M_inv + + (global_index + tile_id) * TileLen); + TSTORE(M_inv_global_out, c_l0_tile[final_c_buffer_index]); + } + next_tile_id_that_waits_for_pipe_fix_pipe_m = + (tile_id + 1) % NumTilesPerCubeIter; + set_flag( + PIPE_FIX, PIPE_M, + static_cast(next_tile_id_that_waits_for_pipe_fix_pipe_m)); + } + } + for (uint32_t tile_id = 0; tile_id < NumTilesPerCubeIter; ++tile_id) { + wait_flag(PIPE_M, PIPE_MTE2, static_cast(tile_id)); + } + wait_flag(PIPE_FIX, PIPE_M, + static_cast(next_tile_id_that_waits_for_pipe_fix_pipe_m)); +} + +/* + * @brief: Computes the inverses of the blocks of tensor M + */ +template +AICORE void runKernelTriInvRecUnroll(__gm__ StoreT* M_inv, __gm__ InputT* M, + __gm__ InputT* I_neg, uint32_t total_tiles, + uint32_t num_bsnd_heads = 0, + __gm__ int32_t* cu_seqlens = nullptr, + uint32_t is_lower = 0) { +#if (__CHECK_FEATURE_AT_PRECOMPILE) || \ + (__CCE_AICORE__ == 220 && defined(__DAV_C220_CUBE__)) // Cube compilation + + TriInvRecUnrollKernel(M_inv, M, I_neg, total_tiles, num_bsnd_heads, + cu_seqlens, is_lower); +#else +// Nothing to do on AIV +#endif +} + +template +AICORE void run_tri_inv_rec_unroll(__gm__ float* tensor_out, + __gm__ InputT* tensor_in, + __gm__ InputT* minus_identity_in, + uint32_t matrix_size, uint32_t num_matrices, + uint32_t num_bsnd_heads, + __gm__ int32_t* cu_seqlens = nullptr, + uint32_t is_lower = 0) { + static_assert(std::is_same_v, + "tri_inv_rec_unroll supports only fp16."); + switch (matrix_size) { + case 16: + runKernelTriInvRecUnroll( + tensor_out, tensor_in, minus_identity_in, num_matrices, + num_bsnd_heads, cu_seqlens, is_lower); + break; + case 32: + runKernelTriInvRecUnroll( + tensor_out, tensor_in, minus_identity_in, num_matrices, + num_bsnd_heads, cu_seqlens, is_lower); + break; + case 64: + runKernelTriInvRecUnroll( + tensor_out, tensor_in, minus_identity_in, num_matrices, + num_bsnd_heads, cu_seqlens, is_lower); + break; + case 128: + runKernelTriInvRecUnroll( + tensor_out, tensor_in, minus_identity_in, num_matrices, + num_bsnd_heads, cu_seqlens, is_lower); + break; + } +} + +/* + * @brief: Wrapper for the kernel, "half" type (fp16). + * + * @param tensor_out pointer to the global memory to store the final inverse. + * @param tensor_in Pointer to the global tensor matrix in global memory. + * @param minus_identity_in Pointer to global memory that contains the negative + * identity. + * @param matrix_size The size if each individual matrix / tile. Can take + * values: {16, 32, 64, 128}. + * @param num_matrices The total number of matrices / tiles in the global + * tensor. + * @param num_bsnd_heads The number of heads, which is only greater than zero + * if the matrix is in BSND format, that is, the tiles need to be loaded with + * strided accesses. If each tile is stored consecutively (and row-wise) in + * memory, then num_bsnd_heads=0. + */ +extern "C" __global__ AICORE void tri_inv_rec_unroll_fp16( + __gm__ void* tensor_out, __gm__ void* tensor_in, + __gm__ void* minus_identity_in, uint32_t matrix_size, uint32_t num_matrices, + uint32_t num_bsnd_heads, __gm__ void* cu_seqlens) { + const uint32_t is_lower = (num_bsnd_heads >> 16) & 1u; + const uint32_t actual_heads = num_bsnd_heads & 0xFFFFu; + if (actual_heads == 0) { + if (num_matrices <= get_block_num()) { + run_tri_inv_rec_unroll( + (__gm__ float*)tensor_out, (__gm__ half*)tensor_in, + (__gm__ half*)minus_identity_in, matrix_size, num_matrices, + actual_heads, (__gm__ int32_t*)cu_seqlens, is_lower); + } else if (num_matrices <= 2 * get_block_num()) { + run_tri_inv_rec_unroll( + (__gm__ float*)tensor_out, (__gm__ half*)tensor_in, + (__gm__ half*)minus_identity_in, matrix_size, num_matrices, + actual_heads, (__gm__ int32_t*)cu_seqlens, is_lower); + } else { + run_tri_inv_rec_unroll( + (__gm__ float*)tensor_out, (__gm__ half*)tensor_in, + (__gm__ half*)minus_identity_in, matrix_size, num_matrices, + actual_heads, (__gm__ int32_t*)cu_seqlens, is_lower); + } + } else { + if (num_matrices <= get_block_num()) { + run_tri_inv_rec_unroll( + (__gm__ float*)tensor_out, (__gm__ half*)tensor_in, + (__gm__ half*)minus_identity_in, matrix_size, num_matrices, + actual_heads, (__gm__ int32_t*)cu_seqlens, is_lower); + } else if (num_matrices <= 2 * get_block_num()) { + run_tri_inv_rec_unroll( + (__gm__ float*)tensor_out, (__gm__ half*)tensor_in, + (__gm__ half*)minus_identity_in, matrix_size, num_matrices, + actual_heads, (__gm__ int32_t*)cu_seqlens, is_lower); + } else { + run_tri_inv_rec_unroll( + (__gm__ float*)tensor_out, (__gm__ half*)tensor_in, + (__gm__ half*)minus_identity_in, matrix_size, num_matrices, + actual_heads, (__gm__ int32_t*)cu_seqlens, is_lower); + } + } +} diff --git a/csrc/ascend/pto_chunk_gdn/wy_fast.cpp b/csrc/ascend/pto_chunk_gdn/wy_fast.cpp new file mode 100644 index 000000000..418c0574b --- /dev/null +++ b/csrc/ascend/pto_chunk_gdn/wy_fast.cpp @@ -0,0 +1,1013 @@ +// ============================================================================ +// wy_fast_kernel.cpp — WY representation for GatedDeltaNet chunk recurrence +// +// Computes the WY update matrices U and W for each chunk of C tokens: +// U = A2 @ V where A2 = A * beta_2d (beta-scaled attention) +// W = A1 @ K where A1 = A * (exp(g)*beta)_2d (gate+beta-scaled attention) +// +// beta is the decay factor, g is the gate value, A is the triangular attention +// matrix (from the kkt kernel). The column-broadcast notation x_2d means +// expanding a 1xC vector into a C/2 x C matrix by replicating across rows. +// +// Architecture: Vec+Cube cooperative kernel using cross-core synchronization. +// +// Vec core (two sub-blocks for upper/lower C/2 rows): +// For each chunk: +// 1. Load beta [H,T] and A [B,S,H,C], compute A2 = A * beta_2d -> ws +// 2. Load G [H,T], compute A1 = A * (exp(g)*beta)_2d -> ws +// 3. Signal Cube via cross-core flags when workspaces are ready +// +// Cube core (waits for Vec signals): +// For each chunk: +// 1. Load K, V from BSND layout into L1 +// 2. Load A2 from workspace -> GEMM: U = A2 @ V +// 3. Load A1 from workspace -> GEMM: W = A1 @ K +// 4. Store U, W back to BSND layout +// +// NPU memory hierarchy used: +// GM -> UB (Vec), GM -> L1 -> L0A/L0B -> L0C -> GM (Cube) +// +// ── PTO / NPU Primer ────────────────────────────────────────────────── +// This kernel uses BOTH the Cube engine (matrix multiply) and Vec engine +// (SIMD element-wise ops), running on SEPARATE physical cores that +// communicate via Global Memory (GM) + cross-core flags (FFTS). +// +// Execution flow: +// Vec core: load A,beta,G → compute A2,A1 → store to GM workspace +// Cube core: wait for workspace → load A2/A1 + K/V → GEMM → store U,W +// +// Key PTO APIs (with numpy/torch equivalents): +// TLOAD(ub_tile, gm) — ub_tile = gm[...] (DMA: GM→UB, async MTE2) +// TSTORE(gm, ub_tile) — gm[...] = ub_tile (DMA: UB→GM, async MTE3) +// TCVT(dst, src, mode) — dst = src.float() or .half() (type conversion) +// TMOV(dst, src) — dst = src.clone() +// TMUL(d, a, b) — d = a * b (element-wise) +// TEXP(d, s) — d = torch.exp(s) +// TCOLEXPAND(2d, row) — 2d[i,j] = row[j] (broadcast row across all rows) +// TEXTRACT(l0, l1, r, c) — L1 sub-block → L0A/L0B (MTE1 for Cube GEMM) +// TMATMUL(C, A, B) — C = A @ B in Cube engine (fp16→fp32 accumulate) +// set_flag / wait_flag — sync between pipes on SAME core +// ffts_cross_core_sync — signal ACROSS Cube↔Vec cores +// wait_flag_dev(flag) — wait for cross-core signal +// ============================================================================ + +#include +#include "acl/acl.h" +#include +#include +using namespace pto; + +#ifndef GDN_H +#define GDN_H 16 +#endif + +#ifndef GDN_HG +#define GDN_HG GDN_H +#endif + +#ifndef GDN_D +#define GDN_D 128 +#endif + +#ifndef GDN_C +#define GDN_C 128 +#endif + +#ifdef __CCE_AICORE__ + +namespace { + +template +using TileMatL1 = pto::Tile; + +template +using TileMatL1ZN = pto::Tile; + +template +using TileMatL0A = pto::Tile; + +template +using TileMatL0B = pto::Tile; + +template +using TileUbDataND = + pto::Tile; + +template +using TileUbDataDN = + pto::Tile; + +using GmShape2D = pto::Shape<1, 1, 1, pto::DYNAMIC, pto::DYNAMIC>; +using GmStride2D = pto::Stride<1, 1, 1, pto::DYNAMIC, 1>; + +template +using GmTensor2D = pto::GlobalTensor; + +template +using DynMatL1 = pto::Tile; + +template +using DynVecTile = pto::Tile; + +template +using DynAccTile = pto::TileAcc; + +// PTO cheat sheet for readers coming from PyTorch / NumPy: +// - `GlobalTensor` is a GM tensor view with explicit shape/stride metadata. +// - `Tile<..., Mat, ...>` is an on-chip matrix tile used by Cube kernels. +// - `Tile<..., Vec, ...>` is an on-chip UB tile used by SIMD vector kernels. +// - `TileAcc` is the matmul accumulator tile. +// - `TLOAD` / `TSTORE` are DMA copies between GM and local memory. +// - `TCOLEXPAND` is broadcast like `x[None, :].expand(rows, -1)`. +// - `TMUL`, `TEXP`, `TCVT` are vector ops on UB tiles. + +template +AICORE PTO_INLINE void +gemm_v0(std::conditional_t, + TileMatL1> &A, + std::conditional_t, + TileMatL1> &B, + pto::TileAcc &C, bool clear) +{ + // Local K-sliced matmul helper: + // C = A @ B + // PTO exposes the L1 -> L0 -> Cube movement explicitly, so keeping this tiny + // helper local lets readers see the schedule without hiding it in a repo-wide + // wrapper layer. + // + // PyTorch mental model: + // C = 0 + // for k0 in range(0, K, kL0Size): + // C += A[:, k0:k1] @ B[k0:k1, :] + constexpr uint32_t kL0Size = 128; + const uint32_t kL0split = (K + kL0Size - 1) / kL0Size; + + auto war_event_id = (event_t)(((int)EVENT_ID0 + 1) % 8); + set_flag(PIPE_MTE2, PIPE_MTE1, war_event_id); + wait_flag(PIPE_MTE2, PIPE_MTE1, war_event_id); + + for (uint32_t kL0Idx = 0; kL0Idx < kL0split; ++kL0Idx) { + const bool initflag = clear && (kL0Idx == 0); + const bool is_tail_block = (kL0Idx == kL0split - 1); + + if (is_tail_block) { + TileMatL0A l0a; + TileMatL0B l0b; + pto::TASSIGN(l0a, 0x0); + pto::TASSIGN(l0b, 0x0); + + set_flag(PIPE_M, PIPE_MTE1, war_event_id); + wait_flag(PIPE_M, PIPE_MTE1, war_event_id); + + if constexpr (!transpose_A) { + pto::TEXTRACT(l0a, A, 0, kL0Idx * K_tail); + } else { + TileMatL1ZN A_t; + pto::TRESHAPE(A_t, A); + pto::TEXTRACT(l0a, A_t, 0, kL0Idx * K_tail); + } + + if constexpr (!transpose_B) { + pto::TEXTRACT(l0b, B, kL0Idx * K_tail, 0); + } else { + TileMatL1ZN B_t; + pto::TRESHAPE(B_t, B); + pto::TEXTRACT(l0b, B_t, kL0Idx * K_tail, 0); + } + + set_flag(PIPE_MTE1, PIPE_M, war_event_id); + wait_flag(PIPE_MTE1, PIPE_M, war_event_id); + + if (initflag) { + pto::TMATMUL(C, l0a, l0b); + } else { + pto::TMATMUL_ACC(C, C, l0a, l0b); + } + } else { + TileMatL0A l0a; + TileMatL0B l0b; + pto::TASSIGN(l0a, 0x0); + pto::TASSIGN(l0b, 0x0); + + set_flag(PIPE_M, PIPE_MTE1, war_event_id); + wait_flag(PIPE_M, PIPE_MTE1, war_event_id); + + set_flag(PIPE_FIX, PIPE_M, war_event_id); + wait_flag(PIPE_FIX, PIPE_M, war_event_id); + + if constexpr (!transpose_A) { + pto::TEXTRACT(l0a, A, 0, kL0Idx * kL0Size); + } else { + TileMatL1ZN A_t; + pto::TRESHAPE(A_t, A); + pto::TEXTRACT(l0a, A_t, 0, kL0Idx * kL0Size); + } + + if constexpr (!transpose_B) { + pto::TEXTRACT(l0b, B, kL0Idx * kL0Size, 0); + } else { + TileMatL1ZN B_t; + pto::TRESHAPE(B_t, B); + pto::TEXTRACT(l0b, B_t, kL0Idx * kL0Size, 0); + } + + set_flag(PIPE_MTE1, PIPE_M, war_event_id); + wait_flag(PIPE_MTE1, PIPE_M, war_event_id); + + if (initflag) { + pto::TMATMUL(C, l0a, l0b); + } else { + pto::TMATMUL_ACC(C, C, l0a, l0b); + } + + set_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + wait_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + } + } + + set_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + wait_flag(PIPE_MTE1, PIPE_MTE2, war_event_id); + + set_flag(PIPE_M, PIPE_FIX, war_event_id); + wait_flag(PIPE_M, PIPE_FIX, war_event_id); +} + +} // namespace + +#endif + +template +AICORE void wy_fast_kernel( + __gm__ half *K_handle, __gm__ half *V_handle, + __gm__ half *Beta_handle, __gm__ float *G_handle, + __gm__ half *A_handle, + __gm__ half *workspace_a1_handle, __gm__ half *workspace_a2_handle, + __gm__ half *W_handle, __gm__ half *U_handle, + __gm__ int32_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, int64_t total_tokens, + uint64_t ffts_addr) +{ + // WY recompute materializes two diagonal reweightings of the same A tile: + // A2[:, j] = A[:, j] * beta_j + // A1[:, j] = A[:, j] * exp(g_j) * beta_j + // and then forms the two branch outputs + // U = A2 @ V, W = A1 @ K. + // + // Shapes for one (sequence, head, chunk): + // A_chunk : [valid, valid] + // beta : [valid] + // g : [valid] + // K, V : [valid, D] + // + // PyTorch / NumPy sketch: + // A2 = A_chunk * beta[None, :] + // A1 = A_chunk * (exp(g) * beta)[None, :] + // U = A2 @ V_chunk + // W = A1 @ K_chunk + // + // PTO split: + // Vec builds the two reweighted A tiles in workspace. + // Cube later consumes those workspaces in two GEMMs. + constexpr int32_t HalfChunk = ChunkSize / 2; + constexpr uint32_t KTail = + (HiddenSize % 128 == 0) ? 128 : (HiddenSize % 128); + + constexpr int32_t H = NumHeads; + constexpr int32_t Hg = NumKeyHeads; + static_assert(Hg > 0 && H % Hg == 0, + "NumHeads must be divisible by NumKeyHeads"); + constexpr int32_t GROUP = H / Hg; + constexpr int32_t BSND_V_STRIDE = H * HiddenSize; + constexpr int32_t BSND_QK_STRIDE = Hg * HiddenSize; + + constexpr int32_t GHeadTileCols = ((NumHeads + 7) / 8) * 8; + constexpr int32_t BetaHeadTileCols = ((NumHeads + 15) / 16) * 16; + + constexpr int32_t BetaHalfUbAddr = 0; + constexpr int32_t A1HalfUbAddr = 256; + constexpr int32_t BetaUbAddr = 16640; + constexpr int32_t BetaRUbAddr = 17152; + constexpr int32_t Beta2dUbAddr = 17664; + constexpr int32_t TmpUbAddr = 50432; + constexpr int32_t A1UbAddr = 75008; + constexpr int32_t A2UbAddr = 107776; + constexpr int32_t A2HalfUbAddr = 140544; + constexpr int32_t GUbAddr = 156928; + constexpr int32_t GRUbAddr = 157440; + constexpr int32_t G2dUbAddr = 157952; + + constexpr int32_t GBlockUbAddr = TmpUbAddr; + constexpr int32_t BetaBlockUbAddr = TmpUbAddr; + + constexpr int32_t WsA1Size = ChunkSize * ChunkSize; + constexpr int32_t WsA2Size = ChunkSize * ChunkSize; + + set_ffts_base_addr(ffts_addr); + auto cid = get_block_idx(); + auto block_num = get_block_num(); + auto vid = get_subblockid(); + + int64_t num_seqs = batch_size; + + TileUbDataND beta_ub_half; + TASSIGN(beta_ub_half, BetaHalfUbAddr); + TileUbDataND a1_ub_half; + TASSIGN(a1_ub_half, A1HalfUbAddr); + TileUbDataND beta_ub; + TASSIGN(beta_ub, BetaUbAddr); + TileUbDataND beta_r_ub; + TASSIGN(beta_r_ub, BetaRUbAddr); + TileUbDataND beta_2d_ub; + TASSIGN(beta_2d_ub, Beta2dUbAddr); + TileUbDataND tmp_ub; + TASSIGN(tmp_ub, TmpUbAddr); + TileUbDataND a1_ub; + TASSIGN(a1_ub, A1UbAddr); + TileUbDataND a2_ub; + TASSIGN(a2_ub, A2UbAddr); + TileUbDataND a2_ub_half; + TASSIGN(a2_ub_half, A2HalfUbAddr); + TileUbDataND g_ub; + TASSIGN(g_ub, GUbAddr); + TileUbDataND g_r_ub; + TASSIGN(g_r_ub, GRUbAddr); + TileUbDataND g_2d_ub; + TASSIGN(g_2d_ub, G2dUbAddr); + + TileMatL1 k_l1; + TASSIGN(k_l1, 0); + TileMatL1 v_l1; + TASSIGN(v_l1, 32768); + TileMatL1 a2_l1; + TASSIGN(a2_l1, 65536); + TileAcc u_l0; + TASSIGN(u_l0, 0); + TileMatL1 a1_l1; + TASSIGN(a1_l1, 98304); + TileAcc w_l0; + TASSIGN(w_l0, 65536); + + int64_t total_work = 0; + if (cu_seqlens == nullptr) { + int64_t chunks_per_seq = (seq_len + ChunkSize - 1) / ChunkSize; + total_work = num_seqs * chunks_per_seq * NumHeads; + } + +#if defined(__DAV_C220_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + + // Vec prepares the two reweighted A workspaces (`A2` and `A1`) that the + // Cube phase consumes later. + if (cu_seqlens == nullptr) { + bool first_iter = true; + int64_t gi = 0; + for (int64_t seq_idx = 0; seq_idx < num_seqs; ++seq_idx) { + int64_t bos = seq_idx * seq_len; + int64_t slen = seq_len; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < nc; ++ci) { + for (int32_t head_idx = 0; head_idx < NumHeads; ++head_idx) { + if (gi % static_cast(block_num) == + static_cast(cid)) { + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast( + remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + // Each Vec sub-block owns one HalfChunk-row stripe of the chunk. + // For a tail chunk, the upper stripe (vid=0) may hold fewer than + // 64 rows, and the lower stripe (vid=1) may hold only a suffix or + // no rows at all. `local_rows` is the exact number of live rows in + // THIS sub-block's stripe. + int32_t local_rows = valid_rows - + static_cast(vid) * HalfChunk; + if (local_rows < 0) local_rows = 0; + if (local_rows > HalfChunk) local_rows = HalfChunk; + + // Beta is pre-transposed to [H, total_tokens] for contiguous loads. + { + GmShape2D beta_shape(1, valid_rows); + GmStride2D beta_stride(1); + GmTensor2D beta_global( + Beta_handle + static_cast(head_idx) * total_tokens + + chunk_token_start, + beta_shape, beta_stride); + DynVecTile beta_load( + 1, valid_rows); + TASSIGN(beta_load, BetaHalfUbAddr); + TLOAD(beta_load, beta_global); + if (valid_rows != ChunkSize) { + TFILLPAD_INPLACE(beta_ub_half, beta_load); + } + } + + // Load only the live rows for this sub-block, then zero-pad the + // remainder of the HalfChunk tile. The Cube phase always consumes + // a full [HalfChunk, ChunkSize] workspace tile, so stale rows here + // would leak garbage into ragged tails and cross-sequence boundaries. + if (local_rows > 0) { + int64_t a_gm_offset = + ((chunk_token_start + + static_cast(vid) * HalfChunk) * + NumHeads + head_idx) * + static_cast(ChunkSize); + GmShape2D a_shape(local_rows, ChunkSize); + GmStride2D a_stride(NumHeads * ChunkSize); + GmTensor2D a_global(A_handle + a_gm_offset, a_shape, + a_stride); + DynVecTile a_load( + local_rows, ChunkSize); + TASSIGN(a_load, A1HalfUbAddr); + TLOAD(a_load, a_global); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(a1_ub_half, a_load); + } + } else { + // Fully empty lower-half tail: materialize an all-zero tile so the + // workspace still looks like a correctly padded HalfChunk block. + TEXPANDS(a1_ub, 0.0f); + pipe_barrier(PIPE_V); + TCVT(a1_ub_half, a1_ub, pto::RoundMode::CAST_NONE); + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TCVT(beta_ub, beta_ub_half, pto::RoundMode::CAST_NONE); + pipe_barrier(PIPE_V); + TMOV(beta_r_ub, beta_ub); + pipe_barrier(PIPE_V); + // Replicate beta_j across rows so every column j of A gets the same beta. + // PyTorch-like: + // beta_2d = beta[None, :].expand(HalfChunk, ChunkSize) + TCOLEXPAND(beta_2d_ub, beta_r_ub); + + TCVT(a1_ub, a1_ub_half, pto::RoundMode::CAST_NONE); + // Form the beta-scaled tile that the later U = A2 * V matmul consumes. + // a2_ub = a1_ub * beta_2d_ub + TMUL(a2_ub, a1_ub, beta_2d_ub); + TCVT(a2_ub_half, a2_ub, pto::RoundMode::CAST_NONE); + + if (!first_iter) wait_flag_dev(3); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + GmShape2D a2_shape(HalfChunk, ChunkSize); + GmStride2D a2_stride(ChunkSize); + GmTensor2D workspace_a2_global( + workspace_a2_handle + + static_cast(cid) * WsA2Size + + static_cast(vid) * HalfChunk * ChunkSize, + a2_shape, a2_stride); + TSTORE(workspace_a2_global, a2_ub_half); + } + pipe_barrier(PIPE_ALL); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (2 << 8)); + + // G is pre-transposed to [H, total_tokens] for contiguous loads. + { + GmShape2D g_shape(1, valid_rows); + GmStride2D g_stride(1); + GmTensor2D g_global( + G_handle + static_cast(head_idx) * total_tokens + + chunk_token_start, + g_shape, g_stride); + DynVecTile g_load( + 1, valid_rows); + TASSIGN(g_load, GUbAddr); + TLOAD(g_load, g_global); + if (valid_rows != ChunkSize) { + TFILLPAD_INPLACE(g_ub, g_load); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Build the g-based column weights before forming the W = A1 * K branch. + // Torch-like: + // g_weight = exp(g) * beta + TEXP(g_ub, g_ub); + pipe_barrier(PIPE_V); + TMUL(g_ub, g_ub, beta_ub); + pipe_barrier(PIPE_V); + TMOV(g_r_ub, g_ub); + pipe_barrier(PIPE_V); + TCOLEXPAND(g_2d_ub, g_r_ub); + // A1 keeps the same A columns but multiplies each one by exp(g_j) * beta_j. + // a1_ub = a1_ub * g_weight[None, :] + TMUL(a1_ub, a1_ub, g_2d_ub); + TCVT(a1_ub_half, a1_ub, pto::RoundMode::CAST_NONE); + + if (!first_iter) wait_flag_dev(4); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + GmShape2D a1_shape(HalfChunk, ChunkSize); + GmStride2D a1_stride(ChunkSize); + GmTensor2D workspace_a1_global( + workspace_a1_handle + + static_cast(cid) * WsA1Size + + static_cast(vid) * HalfChunk * ChunkSize, + a1_shape, a1_stride); + TSTORE(workspace_a1_global, a1_ub_half); + } + pipe_barrier(PIPE_ALL); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + first_iter = false; + } + gi++; + } + } + } + } else { + // Same WY math as above; only the work enumeration changes for varlen input. + int64_t gi = 0; + bool first_iter_v = true; + for (int64_t si = 0; si < num_seqs; ++si) { + int64_t bos = static_cast(cu_seqlens[si]); + int64_t eos = static_cast(cu_seqlens[si + 1]); + int64_t slen = eos - bos; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < nc; ++ci) { + for (int32_t h = 0; h < NumHeads; ++h) { + if (gi % static_cast(block_num) == + static_cast(cid)) { + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast( + remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + // Same HalfChunk ownership rule as the fixed-length path above: + // each Vec sub-block handles one 64-row stripe, and ragged varlen + // tails may leave that stripe partially full or fully empty. + int32_t local_rows = valid_rows - + static_cast(vid) * HalfChunk; + if (local_rows < 0) local_rows = 0; + if (local_rows > HalfChunk) local_rows = HalfChunk; + int32_t head_idx = h; + + // Beta is pre-transposed to [H, total_tokens] for contiguous loads. + { + GmShape2D beta_shape(1, valid_rows); + GmStride2D beta_stride(1); + GmTensor2D beta_global( + Beta_handle + static_cast(head_idx) * total_tokens + + chunk_token_start, + beta_shape, beta_stride); + DynVecTile beta_load( + 1, valid_rows); + TASSIGN(beta_load, BetaHalfUbAddr); + TLOAD(beta_load, beta_global); + if (valid_rows != ChunkSize) { + TFILLPAD_INPLACE(beta_ub_half, beta_load); + } + } + + // Tail-safe A loading is especially important in varlen mode because + // the final chunk of one sequence may be immediately followed by the + // first chunk of the next sequence in packed storage. + if (local_rows > 0) { + int64_t a_gm_offset = + ((chunk_token_start + + static_cast(vid) * HalfChunk) * + NumHeads + head_idx) * + static_cast(ChunkSize); + GmShape2D a_shape(local_rows, ChunkSize); + GmStride2D a_stride(NumHeads * ChunkSize); + GmTensor2D a_global(A_handle + a_gm_offset, a_shape, + a_stride); + DynVecTile a_load( + local_rows, ChunkSize); + TASSIGN(a_load, A1HalfUbAddr); + TLOAD(a_load, a_global); + if (local_rows != HalfChunk) { + TFILLPAD_INPLACE(a1_ub_half, a_load); + } + } else { + // Empty stripe for this sub-block: write zeros so the downstream + // full-tile Cube GEMM sees valid padding rather than old workspace. + TEXPANDS(a1_ub, 0.0f); + pipe_barrier(PIPE_V); + TCVT(a1_ub_half, a1_ub, pto::RoundMode::CAST_NONE); + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TCVT(beta_ub, beta_ub_half, pto::RoundMode::CAST_NONE); + pipe_barrier(PIPE_V); + TMOV(beta_r_ub, beta_ub); + pipe_barrier(PIPE_V); + TCOLEXPAND(beta_2d_ub, beta_r_ub); + + TCVT(a1_ub, a1_ub_half, pto::RoundMode::CAST_NONE); + // Form the beta-scaled tile that the later U = A2 * V matmul consumes. + TMUL(a2_ub, a1_ub, beta_2d_ub); + TCVT(a2_ub_half, a2_ub, pto::RoundMode::CAST_NONE); + + if (!first_iter_v) wait_flag_dev(3); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + GmShape2D a2_shape(HalfChunk, ChunkSize); + GmStride2D a2_stride(ChunkSize); + GmTensor2D workspace_a2_global( + workspace_a2_handle + + static_cast(cid) * WsA2Size + + static_cast(vid) * HalfChunk * ChunkSize, + a2_shape, a2_stride); + TSTORE(workspace_a2_global, a2_ub_half); + } + pipe_barrier(PIPE_ALL); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (2 << 8)); + + // G is pre-transposed to [H, total_tokens] for contiguous loads. + { + GmShape2D g_shape(1, valid_rows); + GmStride2D g_stride(1); + GmTensor2D g_global( + G_handle + static_cast(head_idx) * total_tokens + + chunk_token_start, + g_shape, g_stride); + DynVecTile g_load( + 1, valid_rows); + TASSIGN(g_load, GUbAddr); + TLOAD(g_load, g_global); + if (valid_rows != ChunkSize) { + TFILLPAD_INPLACE(g_ub, g_load); + } + } + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Build the g-based column weights before forming the W = A1 * K branch. + TEXP(g_ub, g_ub); + pipe_barrier(PIPE_V); + TMUL(g_ub, g_ub, beta_ub); + pipe_barrier(PIPE_V); + TMOV(g_r_ub, g_ub); + pipe_barrier(PIPE_V); + TCOLEXPAND(g_2d_ub, g_r_ub); + TMUL(a1_ub, a1_ub, g_2d_ub); + TCVT(a1_ub_half, a1_ub, pto::RoundMode::CAST_NONE); + + if (!first_iter_v) wait_flag_dev(4); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + { + GmShape2D a1_shape(HalfChunk, ChunkSize); + GmStride2D a1_stride(ChunkSize); + GmTensor2D workspace_a1_global( + workspace_a1_handle + + static_cast(cid) * WsA1Size + + static_cast(vid) * HalfChunk * ChunkSize, + a1_shape, a1_stride); + TSTORE(workspace_a1_global, a1_ub_half); + } + pipe_barrier(PIPE_ALL); + ffts_cross_core_sync(PIPE_MTE3, 1 | (2 << 4) | (1 << 8)); + first_iter_v = false; + } + gi++; + } + } + } + } +#endif + +#if defined(__DAV_C220_CUBE__) + // Cube consumes the two Vec-generated workspaces and turns them into the + // branch outputs U and W. + if (cu_seqlens == nullptr) { + int64_t gi = 0; + for (int64_t seq_idx = 0; seq_idx < num_seqs; ++seq_idx) { + int64_t bos = seq_idx * seq_len; + int64_t slen = seq_len; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < nc; ++ci) { + for (int32_t head_idx = 0; head_idx < NumHeads; ++head_idx) { + if (gi % static_cast(block_num) == + static_cast(cid)) { + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast( + remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + + int32_t head_g = head_idx / GROUP; + int64_t k_off = + (chunk_token_start * static_cast(Hg) + + static_cast(head_g)) * + static_cast(HiddenSize); + int64_t v_off = + (chunk_token_start * static_cast(H) + + static_cast(head_idx)) * + static_cast(HiddenSize); + + { + GmShape2D k_shape(valid_rows, HiddenSize); + GmStride2D k_stride(BSND_QK_STRIDE); + GmTensor2D k_global(K_handle + k_off, k_shape, k_stride); + DynMatL1 k_l1_load(valid_rows, + HiddenSize); + TASSIGN(k_l1_load, 0); + TLOAD(k_l1_load, k_global); + if (valid_rows != ChunkSize) { + TFILLPAD(k_l1_load, k_l1_load); + } + } + { + GmShape2D v_shape(valid_rows, HiddenSize); + GmStride2D v_stride(BSND_V_STRIDE); + GmTensor2D v_global(V_handle + v_off, v_shape, v_stride); + DynMatL1 v_l1_load(valid_rows, + HiddenSize); + TASSIGN(v_l1_load, 32768); + TLOAD(v_l1_load, v_global); + if (valid_rows != ChunkSize) { + TFILLPAD(v_l1_load, v_l1_load); + } + } + + wait_flag_dev(2); + { + GmShape2D a2_shape(ChunkSize, ChunkSize); + GmStride2D a2_stride(ChunkSize); + GmTensor2D workspace_a2_global( + workspace_a2_handle + static_cast(cid) * WsA2Size, + a2_shape, a2_stride); + // Load the Vec-prepared A2 tile: + // A2 = A * beta[None, :] + TLOAD(a2_l1, workspace_a2_global); + } + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + // U = A2 * V keeps the beta-scaled path separate from the K-side update. + gemm_v0(a2_l1, v_l1, u_l0, true); + + { + GmShape2D u_shape(valid_rows, HiddenSize); + GmStride2D u_stride(BSND_V_STRIDE); + GmTensor2D u_global(U_handle + v_off, u_shape, u_stride); + DynAccTile u_store(valid_rows, + HiddenSize); + TASSIGN(u_store, 0); + // Store only the valid token rows even though the accumulator tile is + // physically ChunkSize x HiddenSize. + TSTORE(u_global, u_store); + } + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (3 << 8)); + + wait_flag_dev(1); + { + GmShape2D a1_shape(ChunkSize, ChunkSize); + GmStride2D a1_stride(ChunkSize); + GmTensor2D workspace_a1_global( + workspace_a1_handle + static_cast(cid) * WsA1Size, + a1_shape, a1_stride); + // Load the Vec-prepared A1 tile: + // A1 = A * (exp(g) * beta)[None, :] + TLOAD(a1_l1, workspace_a1_global); + } + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + // W = A1 * K uses the g-reweighted path for the complementary WY factor. + gemm_v0(a1_l1, k_l1, w_l0, true); + + { + GmShape2D w_shape(valid_rows, HiddenSize); + GmStride2D w_stride(BSND_V_STRIDE); + GmTensor2D w_global(W_handle + v_off, w_shape, w_stride); + DynAccTile w_store(valid_rows, + HiddenSize); + TASSIGN(w_store, 65536); + TSTORE(w_global, w_store); + } + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (4 << 8)); + } + gi++; + } + } + } + } else { + int64_t gi = 0; + for (int64_t si = 0; si < num_seqs; ++si) { + int64_t bos = static_cast(cu_seqlens[si]); + int64_t eos = static_cast(cu_seqlens[si + 1]); + int64_t slen = eos - bos; + int64_t nc = (slen + ChunkSize - 1) / ChunkSize; + + for (int64_t ci = 0; ci < nc; ++ci) { + for (int32_t h = 0; h < NumHeads; ++h) { + if (gi % static_cast(block_num) == + static_cast(cid)) { + int64_t chunk_start = ci * ChunkSize; + int64_t remaining = slen - chunk_start; + int32_t valid_rows = static_cast( + remaining < ChunkSize ? remaining : ChunkSize); + int64_t chunk_token_start = bos + chunk_start; + int32_t head_idx = h; + + int32_t head_g = head_idx / GROUP; + int64_t k_off = + (chunk_token_start * static_cast(Hg) + + static_cast(head_g)) * + static_cast(HiddenSize); + int64_t v_off = + (chunk_token_start * static_cast(H) + + static_cast(head_idx)) * + static_cast(HiddenSize); + + { + GmShape2D k_shape(valid_rows, HiddenSize); + GmStride2D k_stride(BSND_QK_STRIDE); + GmTensor2D k_global(K_handle + k_off, k_shape, + k_stride); + DynMatL1 k_l1_load(valid_rows, + HiddenSize); + TASSIGN(k_l1_load, 0); + TLOAD(k_l1_load, k_global); + if (valid_rows != ChunkSize) { + TFILLPAD(k_l1_load, k_l1_load); + } + } + { + GmShape2D v_shape(valid_rows, HiddenSize); + GmStride2D v_stride(BSND_V_STRIDE); + GmTensor2D v_global(V_handle + v_off, v_shape, + v_stride); + DynMatL1 v_l1_load(valid_rows, + HiddenSize); + TASSIGN(v_l1_load, 32768); + TLOAD(v_l1_load, v_global); + if (valid_rows != ChunkSize) { + TFILLPAD(v_l1_load, v_l1_load); + } + } + + wait_flag_dev(2); + { + GmShape2D a2_shape(ChunkSize, ChunkSize); + GmStride2D a2_stride(ChunkSize); + GmTensor2D workspace_a2_global( + workspace_a2_handle + static_cast(cid) * WsA2Size, + a2_shape, a2_stride); + TLOAD(a2_l1, workspace_a2_global); + } + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + // U = A2 * V keeps the beta-scaled path separate from the K-side update. + gemm_v0(a2_l1, v_l1, u_l0, true); + + { + GmShape2D u_shape(valid_rows, HiddenSize); + GmStride2D u_stride(BSND_V_STRIDE); + GmTensor2D u_global(U_handle + v_off, u_shape, + u_stride); + DynAccTile u_store(valid_rows, + HiddenSize); + TASSIGN(u_store, 0); + TSTORE(u_global, u_store); + } + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (3 << 8)); + + wait_flag_dev(1); + { + GmShape2D a1_shape(ChunkSize, ChunkSize); + GmStride2D a1_stride(ChunkSize); + GmTensor2D workspace_a1_global( + workspace_a1_handle + static_cast(cid) * WsA1Size, + a1_shape, a1_stride); + TLOAD(a1_l1, workspace_a1_global); + } + + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + // W = A1 * K uses the g-reweighted path for the complementary WY factor. + gemm_v0(a1_l1, k_l1, w_l0, true); + + { + GmShape2D w_shape(valid_rows, HiddenSize); + GmStride2D w_stride(BSND_V_STRIDE); + GmTensor2D w_global(W_handle + v_off, w_shape, + w_stride); + DynAccTile w_store(valid_rows, + HiddenSize); + TASSIGN(w_store, 65536); + TSTORE(w_global, w_store); + } + ffts_cross_core_sync(PIPE_FIX, 1 | (2 << 4) | (4 << 8)); + } + gi++; + } + } + } + } +#endif +} + +extern "C" __global__ AICORE void launch_wy_fast( + __gm__ uint8_t *K_handle, __gm__ uint8_t *V_handle, + __gm__ uint8_t *Beta_handle, __gm__ uint8_t *G_handle, + __gm__ uint8_t *A_handle, + __gm__ uint8_t *workspace_a1_handle, __gm__ uint8_t *workspace_a2_handle, + __gm__ uint8_t *W_handle, __gm__ uint8_t *U_handle, + __gm__ uint8_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, int64_t total_tokens, + uint64_t ffts_addr) +{ + wy_fast_kernel( + reinterpret_cast<__gm__ half *>(K_handle), + reinterpret_cast<__gm__ half *>(V_handle), + reinterpret_cast<__gm__ half *>(Beta_handle), + reinterpret_cast<__gm__ float *>(G_handle), + reinterpret_cast<__gm__ half *>(A_handle), + reinterpret_cast<__gm__ half *>(workspace_a1_handle), + reinterpret_cast<__gm__ half *>(workspace_a2_handle), + reinterpret_cast<__gm__ half *>(W_handle), + reinterpret_cast<__gm__ half *>(U_handle), + reinterpret_cast<__gm__ int32_t *>(cu_seqlens), + batch_size, seq_len, total_tokens, ffts_addr); +} + +extern "C" void call_kernel( + uint32_t block_dim, void *stream, + uint8_t *k, uint8_t *v, uint8_t *beta, uint8_t *g_sum, uint8_t *A, + uint8_t *workspace_a1, uint8_t *workspace_a2, + uint8_t *w, uint8_t *u, + uint8_t *cu_seqlens, + int64_t batch_size, int64_t seq_len, int64_t total_tokens) +{ + uint32_t fftsLen{0}; + uint64_t fftsAddr{0}; + rtGetC2cCtrlAddr(&fftsAddr, &fftsLen); + launch_wy_fast<<>>( + k, v, beta, g_sum, A, + workspace_a1, workspace_a2, + w, u, + cu_seqlens, + batch_size, seq_len, total_tokens, fftsAddr); +} diff --git a/csrc/ascend/scripts/opgen/opgen_standalone.py b/csrc/ascend/scripts/opgen/opgen_standalone.py new file mode 100644 index 000000000..8c89c71a8 --- /dev/null +++ b/csrc/ascend/scripts/opgen/opgen_standalone.py @@ -0,0 +1,163 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import argparse +import logging +import os +import shutil +import sys + +import regex as re + + +class OpGenerator: + """算子工程生成器""" + + def __init__(self, op_type, op_name, output_path): + self.op_type = op_type + self.op_name = op_name + self.output_path = output_path + self.template_name = "add_example" + + self.script_dir = os.path.dirname(os.path.abspath(__file__)) + self.template_dir = os.path.abspath(os.path.join(self.script_dir, "template", "add")) + self.dest_dir = os.path.abspath(os.path.join(self.output_path, self.op_type, self.op_name)) + + def run(self): + """执行生成流程""" + self._validate_inputs() + self._copy_template() + self._rename_files() + self._replace_content() + logging.info("成功为 %s/%s 创建算子工程！", self.op_type, self.op_name) + logging.info("工程路径: %s", self.dest_dir) + logging.info("Create the initial directory for %s under %s success", self.op_name, self.op_type) + + def _validate_inputs(self): + """校验输入参数的有效性和安全性""" + if not self.op_type or not self.op_name: + raise ValueError("算子类型和算子名称均不能为空。") + + if not re.match(r"^[a-zA-Z0-9_]+$", self.op_type): + raise ValueError(f"算子类型 '{self.op_type}' 包含无效字符。只允许字母、数字和下划线。") + + if not re.match(r"^[a-zA-Z0-9_]+$", self.op_name): + raise ValueError(f"算子名称 '{self.op_name}' 包含无效字符。只允许字母、数字和下划线。") + + if os.path.exists(self.dest_dir): + raise FileExistsError(f"目标目录 '{self.dest_dir}' 已存在。") + + def _copy_template(self): + """复制模板文件到目标目录""" + logging.info("使用模板在 '%s' 创建算子工程...", self.dest_dir) + if not os.path.exists(self.template_dir): + raise FileNotFoundError(f"找不到模板目录 '{self.template_dir}'。请确保 'template/add' 目录存在。") + + try: + shutil.copytree(self.template_dir, self.dest_dir) + if not os.path.isfile(os.path.join(os.path.dirname(self.dest_dir), "CMakeLists.txt")): + cmake_src = os.path.join(os.path.dirname(self.template_dir), "CMakeLists.txt") + cmake_dest = os.path.join(os.path.dirname(self.dest_dir), "CMakeLists.txt") + shutil.copy2(cmake_src, cmake_dest) + except OSError as e: + raise OSError(f"复制模板文件失败: {e}") from e + + def _rename_files(self): + """重命名文件和目录中的占位符""" + for root, dirs, files in os.walk(self.dest_dir, topdown=False): + for name in files + dirs: + if self.template_name not in name: + continue + + old_path = os.path.join(root, name) + new_name = name.replace(self.template_name, self.op_name) + new_path = os.path.join(root, new_name) + try: + os.rename(old_path, new_path) + except OSError as e: + raise OSError(f"重命名 '{old_path}' 到 '{new_path}' 失败: {e}") from e + + def _replace_content_in_file(self, file_path, replacements): + """Helper to replace content in a single file.""" + try: + with open(file_path, encoding="utf-8", errors="ignore") as f: + content = f.read() + except OSError as e: + logging.warning("读取文件 '%s' 失败: %s", file_path, e) + return + + original_content = content + for old, new in replacements.items(): + content = content.replace(old, new) + + if content == original_content: + return + + try: + with open(file_path, "w", encoding="utf-8") as f: + f.write(content) + except OSError as e: + logging.warning("写入文件 '%s' 失败: %s", file_path, e) + + def _replace_content(self): + """替换文件内容中的占位符""" + op_name_capitalized = "".join(word.capitalize() for word in self.op_name.split("_")) + template_name_capitalized = "".join(word.capitalize() for word in self.template_name.split("_")) + + replacements = { + self.template_name: self.op_name, + self.template_name.upper(): self.op_name.upper(), + template_name_capitalized: op_name_capitalized, + "add_example": self.op_name, + } + for root, _, files in os.walk(self.dest_dir): + for file in files: + if file.endswith((".pyc", ".pyo")): + continue + + file_path = os.path.join(root, file) + self._replace_content_in_file(file_path, replacements) + + +def execute(args): + """根据命令行参数执行算子生成""" + generator = OpGenerator(op_type=args.op_type, op_name=args.op_name, output_path=args.output_path) + generator.run() + + +def register_parser(subparsers): + """为 opgen 命令注册解析器。""" + parser_opgen = subparsers.add_parser("opgen", help="生成项目骨架") + parser_opgen.add_argument("--op_type", "-t", required=True, help="算子分类，例如 math") + parser_opgen.add_argument("--op_name", "-n", required=True, help="新算子的名称，例如 asinh") + parser_opgen.add_argument("--output_path", "-p", default=".", help="生成工程的根路径") + parser_opgen.set_defaults(func=execute) + + +def main(): + """主函数，用于独立执行""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s", stream=sys.stdout) + parser = argparse.ArgumentParser(description="生成项目骨架") + + parser.add_argument("--op_type", "-t", required=True, help="算子分类，例如 math") + parser.add_argument("--op_name", "-n", required=True, help="新算子的名称，例如 asinh") + parser.add_argument("--output_path", "-p", default=".", help="生成工程的根路径") + + args = parser.parse_args() + + try: + execute(args) + except Exception as e: + logging.error("发生非预期的错误，退出。错误信息: %s", e) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/csrc/ascend/scripts/opgen/template/CMakeLists.txt b/csrc/ascend/scripts/opgen/template/CMakeLists.txt new file mode 100644 index 000000000..fbf471e76 --- /dev/null +++ b/csrc/ascend/scripts/opgen/template/CMakeLists.txt @@ -0,0 +1,29 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +set(OPTEST_NAME optest_${PKG_NAME}) + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +foreach(SUB_DIR ${CURRENT_DIRS}) + if (DEFINED ASCEND_OP_NAME AND NOT "${ASCEND_OP_NAME}" STREQUAL "") + if (NOT "${ASCEND_OP_NAME}" STREQUAL "all" AND NOT "${ASCEND_OP_NAME}" STREQUAL "ALL") + if (NOT ${SUB_DIR} IN_LIST ASCEND_OP_NAME) + continue() + endif () + endif () + endif () + + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + else() + if (EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/op_host/CMakeLists.txt") + add_subdirectory(${SUB_DIR}/op_host) + endif() + endif() +endforeach() diff --git a/csrc/ascend/scripts/opgen/template/add/CMakeLists.txt b/csrc/ascend/scripts/opgen/template/add/CMakeLists.txt new file mode 100644 index 000000000..549b5c14c --- /dev/null +++ b/csrc/ascend/scripts/opgen/template/add/CMakeLists.txt @@ -0,0 +1,19 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +if(NOT ENABLE_TEST AND NOT BENCHMARK) + list(REMOVE_ITEM CURRENT_DIRS tests) +endif() +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() diff --git a/csrc/ascend/scripts/opgen/template/add/examples/test_aclnn_add_example.cpp b/csrc/ascend/scripts/opgen/template/add/examples/test_aclnn_add_example.cpp new file mode 100644 index 000000000..f5a485aa6 --- /dev/null +++ b/csrc/ascend/scripts/opgen/template/add/examples/test_aclnn_add_example.cpp @@ -0,0 +1,166 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +// 本文件参考example下示例自动生成 +// 您可自由修改此文件满足需求 + +#include +#include +#include "acl/acl.h" +#include "aclnnop/aclnn_add_example.h" + +#define CHECK_RET(cond, return_expr) \ + do { \ + if (!(cond)) { \ + return_expr; \ + } \ + } while (0) + +#define LOG_PRINT(message, ...) \ + do { \ + printf(message, ##__VA_ARGS__); \ + } while (0) + +int64_t GetShapeSize(const std::vector& shape) +{ + int64_t shapeSize = 1; + for (auto i : shape) { + shapeSize *= i; + } + return shapeSize; +} + +void PrintOutResult(std::vector& shape, void** deviceAddr) +{ + auto size = GetShapeSize(shape); + std::vector resultData(size, 0); + auto ret = aclrtMemcpy( + resultData.data(), resultData.size() * sizeof(resultData[0]), *deviceAddr, size * sizeof(resultData[0]), + ACL_MEMCPY_DEVICE_TO_HOST); + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return); + for (int64_t i = 0; i < size; i++) { + LOG_PRINT("mean result[%ld] is: %f\n", i, resultData[i]); + } +} + +int Init(int32_t deviceId, aclrtStream* stream) +{ + // 固定写法，初始化 + auto ret = aclInit(nullptr); + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret); + ret = aclrtSetDevice(deviceId); + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret); + ret = aclrtCreateStream(stream); + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret); + return 0; +} + +template +int CreateAclTensor( + const std::vector& hostData, const std::vector& shape, void** deviceAddr, aclDataType dataType, + aclTensor** tensor) +{ + auto size = GetShapeSize(shape) * sizeof(T); + // 2. 申请device侧内存 + auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret); + // 3. 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 + ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret); + + // 计算连续tensor的strides + std::vector strides(shape.size(), 1); + for (int64_t i = shape.size() - 2; i >= 0; i--) { + strides[i] = shape[i + 1] * strides[i + 1]; + } + + // 调用aclCreateTensor接口创建aclTensor + *tensor = aclCreateTensor( + shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), + *deviceAddr); + return 0; +} + +int main() +{ + // 1. 调用acl进行device/stream初始化 + int32_t deviceId = 0; + aclrtStream stream; + auto ret = Init(deviceId, &stream); + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret); + + // 2. 构造输入与输出，需要根据API的接口自定义构造 + aclTensor* selfX = nullptr; + void* selfXDeviceAddr = nullptr; + std::vector selfXShape = {32, 4, 4, 4}; + std::vector selfXHostData(2048, 1); // 2048：创建包含32*4*4*4=2048个元素的向量 + ret = CreateAclTensor(selfXHostData, selfXShape, &selfXDeviceAddr, aclDataType::ACL_FLOAT, &selfX); + CHECK_RET(ret == ACL_SUCCESS, return ret); + + aclTensor* selfY = nullptr; + void* selfYDeviceAddr = nullptr; + std::vector selfYShape = {32, 4, 4, 4}; + std::vector selfYHostData(2048, 1); + ret = CreateAclTensor(selfYHostData, selfYShape, &selfYDeviceAddr, aclDataType::ACL_FLOAT, &selfY); + CHECK_RET(ret == ACL_SUCCESS, return ret); + + aclTensor* out = nullptr; + void* outDeviceAddr = nullptr; + std::vector outShape = {32, 4, 4, 4}; + std::vector outHostData(2048, 1); + ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out); + CHECK_RET(ret == ACL_SUCCESS, return ret); + + // 3. 调用CANN算子库API，需要修改为具体的Api名称 + uint64_t workspaceSize = 0; + aclOpExecutor* executor; + + // 4. 调用aclnnAddExample第一段接口 + ret = aclnnAddExampleGetWorkspaceSize(selfX, selfY, out, &workspaceSize, &executor); + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAddExampleGetWorkspaceSize failed. ERROR: %d\n", ret); return ret); + + // 根据第一段接口计算出的workspaceSize申请device内存 + void* workspaceAddr = nullptr; + if (workspaceSize > static_cast(0)) { + ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret); + } + + // 5. 调用aclnnAddExample第二段接口 + ret = aclnnAddExample(workspaceAddr, workspaceSize, executor, stream); + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAddExample failed. ERROR: %d\n", ret); return ret); + + // 6. （固定写法）同步等待任务执行结束 + ret = aclrtSynchronizeStream(stream); + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret); + + // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改 + PrintOutResult(outShape, &outDeviceAddr); + + // 7. 释放aclTensor，需要根据具体API的接口定义修改 + aclDestroyTensor(selfX); + aclDestroyTensor(selfY); + aclDestroyTensor(out); + + // 8. 释放device资源 + aclrtFree(selfXDeviceAddr); + aclrtFree(selfYDeviceAddr); + aclrtFree(outDeviceAddr); + if (workspaceSize > static_cast(0)) { + aclrtFree(workspaceAddr); + } + aclrtDestroyStream(stream); + aclrtResetDevice(deviceId); + + // 9. acl去初始化 + aclFinalize(); + + return 0; +} \ No newline at end of file diff --git a/csrc/ascend/scripts/opgen/template/add/op_host/CMakeLists.txt b/csrc/ascend/scripts/opgen/template/add/op_host/CMakeLists.txt new file mode 100644 index 000000000..11aca0e28 --- /dev/null +++ b/csrc/ascend/scripts/opgen/template/add/op_host/CMakeLists.txt @@ -0,0 +1,18 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +add_op_to_compiled_list() + +if (BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + add_example_def.cpp + ) +endif() + +add_modules_sources(OPTYPE add_example ACLNNTYPE aclnn) diff --git a/csrc/ascend/scripts/opgen/template/add/op_host/add_example_def.cpp b/csrc/ascend/scripts/opgen/template/add/op_host/add_example_def.cpp new file mode 100644 index 000000000..5f655c959 --- /dev/null +++ b/csrc/ascend/scripts/opgen/template/add/op_host/add_example_def.cpp @@ -0,0 +1,55 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_example_def.cpp + * \brief + */ +#include "register/op_def_registry.h" + +namespace ops { +class AddExample : public OpDef { +public: + explicit AddExample(const char* name) : OpDef(name) + { + // 输入参数说明 + this->Input("x1") // 输入x1定义 + .ParamType(REQUIRED) // 必选输入 + .DataType({ge::DT_FLOAT, ge::DT_INT32}) // 支持数据类型 + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) // 支持format格式 + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}) // 未确定大小shape对应format格式 + .AutoContiguous(); // 内存自动连续化 + this->Input("x2") // 输入x2定义 + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + this->Output("y") // 输出y定义 + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT, ge::DT_INT32}) + .Format({ge::FORMAT_ND, ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND}) + .AutoContiguous(); + + OpAICoreConfig aicoreConfig; + aicoreConfig.DynamicCompileStaticFlag(true) + .DynamicFormatFlag(false) + .DynamicRankSupportFlag(true) + .DynamicShapeSupportFlag(true) + .NeedCheckSupportFlag(false) + .PrecisionReduceFlag(true) + .ExtendCfgInfo("opFile.value", "add_example"); // 这里制定的值会对应到kernel入口文件名.cpp + this->AICore().AddConfig("ascend910b", aicoreConfig); // 其他的soc版本补充部分配置项 + this->AICore().AddConfig("ascend910_93", aicoreConfig); + } +}; +OP_ADD(AddExample); // 添加算子信息库 +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/scripts/opgen/template/add/op_host/add_example_infershape.cpp b/csrc/ascend/scripts/opgen/template/add/op_host/add_example_infershape.cpp new file mode 100644 index 000000000..58fb96952 --- /dev/null +++ b/csrc/ascend/scripts/opgen/template/add/op_host/add_example_infershape.cpp @@ -0,0 +1,48 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_example_infer.cpp + * \brief + */ +#include "register/op_impl_registry.h" +#include "log/log.h" + +using namespace ge; + +namespace ops { +static constexpr int64_t IDX_0 = 0; + +static ge::graphStatus InferShapeAddExample(gert::InferShapeContext* context) +{ + OP_LOGD(context->GetNodeName(), "Begin to do InferShapeAddExample"); + + // get input shapes + const gert::Shape* xShape = context->GetInputShape(IDX_0); + OP_CHECK_NULL_WITH_CONTEXT(context, xShape); + + // get output shapes + gert::Shape* yShape = context->GetOutputShape(IDX_0); + OP_CHECK_NULL_WITH_CONTEXT(context, yShape); + + // 填充输出shape大小 + auto xShapeSize = xShape->GetDimNum(); + yShape->SetDimNum(xShapeSize); + for (size_t i = 0; i < xShapeSize; i++) { + int64_t dim = xShape->GetDim(i); + yShape->SetDim(i, dim); + } + + OP_LOGD(context->GetNodeName(), "End to do InferShapeAddExample"); + return GRAPH_SUCCESS; +} + +IMPL_OP_INFERSHAPE(AddExample).InferShape(InferShapeAddExample); +} // namespace ops \ No newline at end of file diff --git a/csrc/ascend/scripts/opgen/template/add/op_host/add_example_tiling.cpp b/csrc/ascend/scripts/opgen/template/add/op_host/add_example_tiling.cpp new file mode 100644 index 000000000..4571fa509 --- /dev/null +++ b/csrc/ascend/scripts/opgen/template/add/op_host/add_example_tiling.cpp @@ -0,0 +1,156 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_example_tiling.cpp + * \brief + */ + +#include "log/log.h" +#include "util/math_util.h" +#include "tiling_base/tiling_util.h" +#include "tiling_base/tiling_templates_registry.h" +#include "../op_kernel/add_example_tiling_data.h" +#include "../op_kernel/add_example_tiling_key.h" + +namespace optiling { + +using namespace Ops::Transformer::OpTiling; + +const uint32_t BLOCK_DIM = 8; +const int64_t TILE_NUM = 8; +const uint32_t WS_SYS_SIZE = 16U * 1024U * 1024U; +const int32_t DIMS_LIMIT = 4; +constexpr int32_t ATTRPOS0 = 0; +constexpr uint32_t INDEXZERO = 0; +constexpr uint32_t INDEXONE = 1; +constexpr uint32_t INDEXTWO = 2; +constexpr uint32_t INDEXTHREE = 3; + +struct AddExampleCompileInfo {}; + +// 获取平台信息如ubSize, coreNum +static ge::graphStatus GetPlatformInfo(gert::TilingContext* context, uint64_t& ubSize, int64_t& coreNum) +{ + // 获取ubsize coreNum + fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo(); + OP_CHECK_NULL_WITH_CONTEXT(context, platformInfoPtr); + auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); + coreNum = ascendcPlatform.GetCoreNumAiv(); + OP_CHECK_IF(coreNum == 0, OP_LOGE(context, "coreNum is 0"), return ge::GRAPH_FAILED); + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); + OP_CHECK_IF(ubSize == 0, OP_LOGE(context, "ubSize is 0"), return ge::GRAPH_FAILED); + return ge::GRAPH_SUCCESS; +} + +// 获取属性，shape信息 +ge::graphStatus GetShapeAttrsInfo(gert::TilingContext* context, int64_t& totalIdx, ge::DataType& dataType) +{ + // 获取输入shape信息 + auto inputX = context->GetInputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context, inputX); + // 如果输入shape 是标量 转换为{1}，否则保持原 shape 不变 + auto inputShapeX = EnsureNotScalar(inputX->GetStorageShape()); + auto inputY = context->GetInputShape(1); + OP_CHECK_NULL_WITH_CONTEXT(context, inputY); + auto inputShapeY = EnsureNotScalar(inputY->GetStorageShape()); + auto outZ = context->GetOutputShape(0); + OP_CHECK_NULL_WITH_CONTEXT(context, outZ); + auto outShapeZ = EnsureNotScalar(outZ->GetStorageShape()); + + // shape校验 + OP_CHECK_IF( + inputShapeX.GetDimNum() != DIMS_LIMIT || inputShapeY.GetDimNum() != DIMS_LIMIT || + outShapeZ.GetDimNum() != DIMS_LIMIT, + OP_LOGE( + context, "AddExample: inputx,inputy,outputz shape dim = %zu, %zu, %zu, should be equal 4", + inputShapeX.GetDimNum(), inputShapeY.GetDimNum(), outShapeZ.GetDimNum()), + return ge::GRAPH_FAILED); + + // 获取shape dim值 + auto nDim = inputShapeX.GetDim(INDEXZERO); + auto cDim = inputShapeX.GetDim(INDEXONE); + auto hDim = inputShapeX.GetDim(INDEXTWO); + auto wDim = inputShapeX.GetDim(INDEXTHREE); + totalIdx = nDim * cDim * hDim * wDim; + // dtype校验 + const std::set supportedDtype = {ge::DT_FLOAT, ge::DT_INT32}; + auto inputDesc = context->GetInputDesc(0); + OP_CHECK_NULL_WITH_CONTEXT(context, inputDesc); + dataType = inputDesc->GetDataType(); + if (supportedDtype.count(dataType) == 0) { + OP_LOGE(context, "invalid dtype"); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +ge::graphStatus GetWorkspaceSize(gert::TilingContext* context) +{ + size_t* currentWorkspace = context->GetWorkspaceSizes(1); + OP_CHECK_NULL_WITH_CONTEXT(context, currentWorkspace); + currentWorkspace[0] = WS_SYS_SIZE; + return ge::GRAPH_SUCCESS; +} + +// tiling 分发入口 +static ge::graphStatus AddExampleTilingFunc(gert::TilingContext* context) +{ + // 1、获取平台运行信息 + uint64_t ubSize; + int64_t coreNum; + OP_CHECK_IF( + GetPlatformInfo(context, ubSize, coreNum) != ge::GRAPH_SUCCESS, OP_LOGE(context, "GetPlatformInfo error"), + return ge::GRAPH_FAILED); + // 2、获取shape、属性信息 + int64_t totalIdx; + ge::DataType dataType; + + OP_CHECK_IF( + GetShapeAttrsInfo(context, totalIdx, dataType) != ge::GRAPH_SUCCESS, + OP_LOGE(context, "GetShapeAttrsInfo error"), return ge::GRAPH_FAILED); + // 3、获取WorkspaceSize信息 + OP_CHECK_IF( + GetWorkspaceSize(context) != ge::GRAPH_SUCCESS, OP_LOGE(context, "GetWorkspaceSize error"), + return ge::GRAPH_FAILED); + + // 4、设置tiling信息 + AddExampleTilingData* tiling = context->GetTilingData(); + OP_CHECK_NULL_WITH_CONTEXT(context, tiling); + OP_CHECK_IF( + memset_s(tiling, sizeof(AddExampleTilingData), 0, sizeof(AddExampleTilingData)) != EOK, + OP_LOGE(context, "set tiling data error"), return ge::GRAPH_FAILED); + tiling->totalLength = totalIdx; + tiling->tileNum = TILE_NUM; + + context->SetBlockDim(BLOCK_DIM); + uint64_t tilingKey = 0; + // 区分dtype走不同得tiling key分支. + if (dataType == ge::DT_FLOAT) { + tilingKey = GET_TPL_TILING_KEY(ELEMENTWISE_TPL_SCH_MODE_0); + context->SetTilingKey(tilingKey); + } else if (dataType == ge::DT_INT32) { + tilingKey = GET_TPL_TILING_KEY(ELEMENTWISE_TPL_SCH_MODE_1); + context->SetTilingKey(tilingKey); + } else { + OP_LOGE(context, "get dtype error"); + return ge::GRAPH_FAILED; + } + return ge::GRAPH_SUCCESS; +} + +static ge::graphStatus TilingParseForAddExample([[maybe_unused]] gert::TilingParseContext* context) +{ + return ge::GRAPH_SUCCESS; +} + +// tiling注册入口. +IMPL_OP_OPTILING(AddExample).Tiling(AddExampleTilingFunc).TilingParse(TilingParseForAddExample); +} // namespace optiling diff --git a/csrc/ascend/scripts/opgen/template/add/op_kernel/add_example.cpp b/csrc/ascend/scripts/opgen/template/add/op_kernel/add_example.cpp new file mode 100644 index 000000000..77a59f8b9 --- /dev/null +++ b/csrc/ascend/scripts/opgen/template/add/op_kernel/add_example.cpp @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_example.cpp + * \brief + */ + +#include "add_example.h" + +enum class AddExampleTilingKey : uint32_t +{ + TILING_KEY_EXAMPLE_FLOAT = 0, + TILING_KEY_EXAMPLE_INT32 = 1, +}; + +template +__global__ __aicore__ void add_example(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) +{ + REGISTER_TILING_DEFAULT(AddExampleTilingData); + GET_TILING_DATA_WITH_STRUCT(AddExampleTilingData, tilingData, tiling); + + // 场景1 + if constexpr (schMode == static_cast(AddExampleTilingKey::TILING_KEY_EXAMPLE_FLOAT)) { + NsAddExample::AddExample op; // 算子kernel实例获取 + op.Init(x, y, z, &tilingData); // 算子kernel实例初始化 + op.Process(); // 算子kernel实例执行 + } + + // 场景2 + if constexpr (schMode == static_cast(AddExampleTilingKey::TILING_KEY_EXAMPLE_INT32)) { + NsAddExample::AddExample op; // 算子kernel实例获取 + op.Init(x, y, z, &tilingData); // 算子kernel实例初始化 + op.Process(); // 算子kernel实例执行 + } +} diff --git a/csrc/ascend/scripts/opgen/template/add/op_kernel/add_example.h b/csrc/ascend/scripts/opgen/template/add/op_kernel/add_example.h new file mode 100644 index 000000000..58fd10427 --- /dev/null +++ b/csrc/ascend/scripts/opgen/template/add/op_kernel/add_example.h @@ -0,0 +1,117 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_example.h + * \brief + */ +#ifndef ADD_EXAMPLE_H +#define ADD_EXAMPLE_H + +#include "kernel_operator.h" +#include "kernel_tiling/kernel_tiling.h" +#include "add_example_tiling_data.h" +#include "add_example_tiling_key.h" + +namespace NsAddExample { + +using namespace AscendC; + +constexpr int32_t BUFFER_NUM = 2; + +template +class AddExample +{ +public: + __aicore__ inline AddExample(){}; + + __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, const AddExampleTilingData* tilingData); + __aicore__ inline void Process(); + +private: + __aicore__ inline void CopyIn(int32_t progress); + __aicore__ inline void CopyOut(int32_t progress); + __aicore__ inline void Compute(const int32_t dataLength); + +private: + TPipe pipe; + TQue inputQueueX; + TQue inputQueueY; + TQue outputQueueZ; + + GlobalTensor inputGMX; + GlobalTensor inputGMY; + GlobalTensor outputGMZ; + + int64_t blockLength_ = 0; + int64_t tileNum_ = 0; + uint32_t tileLength_ = 0; +}; + +template +__aicore__ inline void AddExample::Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, const AddExampleTilingData* tilingData) +{ + blockLength_ = tilingData->totalLength / AscendC::GetBlockNum(); + tileNum_ = tilingData->tileNum; + tileLength_ = blockLength_ / tileNum_ / BUFFER_NUM; + + inputGMX.SetGlobalBuffer((__gm__ T*)x + blockLength_ * AscendC::GetBlockIdx(), blockLength_); + inputGMY.SetGlobalBuffer((__gm__ T*)y + blockLength_ * AscendC::GetBlockIdx(), blockLength_); + outputGMZ.SetGlobalBuffer((__gm__ T*)z + blockLength_ * AscendC::GetBlockIdx(), blockLength_); + + pipe.InitBuffer(inputQueueX, BUFFER_NUM, tileLength_ * sizeof(T)); + pipe.InitBuffer(inputQueueY, BUFFER_NUM, tileLength_ * sizeof(T)); + pipe.InitBuffer(outputQueueZ, BUFFER_NUM, tileLength_ * sizeof(T)); +} + +template +__aicore__ inline void AddExample::CopyIn(int32_t progress) +{ + AscendC::LocalTensor xLocal = inputQueueX.AllocTensor(); + AscendC::LocalTensor yLocal = inputQueueY.AllocTensor(); + AscendC::DataCopy(xLocal, inputGMX[progress * tileLength_], tileLength_); + AscendC::DataCopy(yLocal, inputGMY[progress * tileLength_], tileLength_); + inputQueueX.EnQue(xLocal); + inputQueueY.EnQue(yLocal); +} + +template +__aicore__ inline void AddExample::CopyOut(int32_t progress) +{ + AscendC::LocalTensor zLocal = outputQueueZ.DeQue(); + AscendC::DataCopy(outputGMZ[progress * tileLength_], zLocal, tileLength_); + outputQueueZ.FreeTensor(zLocal); +} + +template +__aicore__ inline void AddExample::Compute(int32_t progress) +{ + AscendC::LocalTensor xLocal = inputQueueX.DeQue(); + AscendC::LocalTensor yLocal = inputQueueY.DeQue(); + AscendC::LocalTensor zLocal = outputQueueZ.AllocTensor(); + AscendC::Add(zLocal, xLocal, yLocal, tileLength_); + outputQueueZ.EnQue(zLocal); + inputQueueX.FreeTensor(xLocal); + inputQueueY.FreeTensor(yLocal); +} + +template +__aicore__ inline void AddExample::Process() +{ + int32_t loopCount = tileNum_ * BUFFER_NUM; + for (int32_t i = 0; i < loopCount; i++) { + CopyIn(i); + Compute(i); + CopyOut(i); + } +} + +} // namespace NsAddExample +#endif // ADD_EXAMPLE_H \ No newline at end of file diff --git a/csrc/ascend/scripts/opgen/template/add/op_kernel/add_example_tiling_data.h b/csrc/ascend/scripts/opgen/template/add/op_kernel/add_example_tiling_data.h new file mode 100644 index 000000000..31c439730 --- /dev/null +++ b/csrc/ascend/scripts/opgen/template/add/op_kernel/add_example_tiling_data.h @@ -0,0 +1,23 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_example_tiling_data.h + * \brief tiling data struct + */ + +#ifndef __ADD_EXAMPLE_TILLING_DATA_H__ +#define __ADD_EXAMPLE_TILLING_DATA_H__ + +struct AddExampleTilingData { + int64_t totalLength; + int64_t tileNum; +}; +#endif diff --git a/csrc/ascend/scripts/opgen/template/add/op_kernel/add_example_tiling_key.h b/csrc/ascend/scripts/opgen/template/add/op_kernel/add_example_tiling_key.h new file mode 100644 index 000000000..9e4ea3247 --- /dev/null +++ b/csrc/ascend/scripts/opgen/template/add/op_kernel/add_example_tiling_key.h @@ -0,0 +1,35 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file add_example_tiling_key.h + * \brief add_example tiling key declare + */ + +#ifndef __ADD_EXAMPLE_TILING_KEY_H__ +#define __ADD_EXAMPLE_TILING_KEY_H__ + +#include "ascendc/host_api/tiling/template_argument.h" + +/* Mode场景定义 */ +#define ELEMENTWISE_TPL_SCH_MODE_0 0 +#define ELEMENTWISE_TPL_SCH_MODE_1 1 +/* 继续定义其他Mode场景... */ + +/* 模板参数 */ +ASCENDC_TPL_ARGS_DECL(AddExample, + ASCENDC_TPL_UINT_DECL(schMode, 1, ASCENDC_TPL_UI_LIST, ELEMENTWISE_TPL_SCH_MODE_0, ELEMENTWISE_TPL_SCH_MODE_1) +); + +ASCENDC_TPL_SEL( + ASCENDC_TPL_ARGS_SEL( + ASCENDC_TPL_UINT_SEL(schMode, ASCENDC_TPL_UI_LIST, ELEMENTWISE_TPL_SCH_MODE_0, ELEMENTWISE_TPL_SCH_MODE_1))); + +#endif \ No newline at end of file diff --git a/csrc/ascend/scripts/opgen/template/add/tests/ut/.gitkeep b/csrc/ascend/scripts/opgen/template/add/tests/ut/.gitkeep new file mode 100644 index 000000000..e69de29bb diff --git a/csrc/ascend/scripts/package/common/__init__.py b/csrc/ascend/scripts/package/common/__init__.py new file mode 100644 index 000000000..59663dc4b --- /dev/null +++ b/csrc/ascend/scripts/package/common/__init__.py @@ -0,0 +1,11 @@ +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- diff --git a/csrc/ascend/scripts/package/common/cfg/path.cfg b/csrc/ascend/scripts/package/common/cfg/path.cfg new file mode 100644 index 000000000..4e87f6e8e --- /dev/null +++ b/csrc/ascend/scripts/package/common/cfg/path.cfg @@ -0,0 +1,2 @@ +CUSTOM_DATA_PATH=$HOME/Ascend/latest/data/ +CUSTOM_CONF_PATH=$HOME/Ascend/latest/conf/ diff --git a/csrc/ascend/scripts/package/common/py/__init__.py b/csrc/ascend/scripts/package/common/py/__init__.py new file mode 100644 index 000000000..59663dc4b --- /dev/null +++ b/csrc/ascend/scripts/package/common/py/__init__.py @@ -0,0 +1,11 @@ +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- diff --git a/csrc/ascend/scripts/package/common/py/filelist.py b/csrc/ascend/scripts/package/common/py/filelist.py new file mode 100644 index 000000000..db15bc5b9 --- /dev/null +++ b/csrc/ascend/scripts/package/common/py/filelist.py @@ -0,0 +1,515 @@ +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""filelist相关类。""" + +import itertools +import os +from collections import Counter +from collections.abc import Callable, Iterator +from enum import IntEnum +from functools import partial +from itertools import chain, repeat +from operator import and_, attrgetter, contains, itemgetter, lt, methodcaller, ne, not_ +from typing import NamedTuple + +from .utils.comm_log import CommLog +from .utils.funcbase import any_, constant, dispatch, identity, invoke, pipe, side_effect, star_apply +from .utils.pkg_utils import ( + TOP_DIR, + FilelistError, + GenerateFilelistError, + conditional_apply, + config_feature_to_string, + pairwise, + swap_args, +) + + +class FileItem(NamedTuple): + """文件条目""" + + module: str + operation: str + relative_path_in_pkg: str + relative_install_path: str + is_in_docker: str + permission: str + owner_group: str + install_type: str + softlink: list[str] + feature: set[str] + is_common_path: str + configurable: str + hash_value: str + block: str + pkg_inner_softlink: list[str] + chip: set[str] + is_dir: bool + + +def create_file_item(*args, **kwargs) -> FileItem: + """创建文件条目。""" + file_item = FileItem(*args, **kwargs) + + if not isinstance(file_item.feature, set): + raise TypeError("The feature parameter should be a set.") + if not isinstance(file_item.chip, set): + raise TypeError("The chip parameter should be a set.") + if not isinstance(file_item.softlink, list): + raise TypeError("The softlink parameter should be a list.") + if not isinstance(file_item.pkg_inner_softlink, list): + raise TypeError("The pkg_inner_softlink parameter should be a list.") + + return file_item + + +# 文件列表 +FileList = list[FileItem] + + +def soft_links_to_string(soft_links: list[str]) -> str: + """软链接转换为字符串。""" + if not soft_links: + return "NA" + return ";".join(soft_links) + + +def file_item_to_string(item: FileItem) -> str: + """文件条目转换为字符串。""" + return ",".join( + [ + item.module, + item.operation, + item.relative_path_in_pkg, + item.relative_install_path, + item.is_in_docker, + item.permission, + item.owner_group, + item.install_type, + soft_links_to_string(item.softlink), + config_feature_to_string(item.feature), + item.is_common_path, + item.configurable, + item.hash_value, + item.block, + soft_links_to_string(item.pkg_inner_softlink), + config_feature_to_string(item.chip), + ] + ) + + +def get_filelist_header_string() -> str: + """获取文件列表表头。""" + return ",".join( + [ + "module", + "operation", + "relative_path_in_pkg", + "relative_install_path", + "is_in_docker", + "permission", + "owner:group", + "install_type", + "softlink", + "feature", + "is_common_path", + "configurable", + "hash", + "block", + "pkg_inner_softlink", + "chip", + ] + ) + + +def get_soft_links_not_in_common_paths(filelist: FileList, target_env: str) -> Iterator[list[str]]: + for file_item_t in filelist: + if file_item_t.relative_install_path.startswith(target_env): + for softlink in file_item_t.softlink: + if not softlink.startswith(target_env): + yield softlink + + +def fill_is_common_path(filelist: FileList, target_env: str) -> Iterator[FileItem]: + """填充文件条目中是否为公共目录字段。""" + soft_links = set(get_soft_links_not_in_common_paths(filelist, target_env)) + for file_item in filelist: + if file_item.relative_install_path.startswith(target_env): + yield file_item._replace(is_common_path="Y") + else: + is_soft_links_prefix = map(methodcaller("startswith", f"{file_item.relative_install_path}/"), soft_links) + if any(is_soft_links_prefix): + yield file_item._replace(is_common_path="YY") + else: + yield file_item + + +def is_relative_install_path(path: str) -> bool: + """是否为相对路径。""" + return not path.startswith("/") + + +def is_specific_operations(file_item: FileItem, operations: list[str]) -> bool: + """是否为特定的操作类型。""" + return file_item.operation in operations + + +def is_specific_install_type(file_item: FileItem, install_types: set[str]) -> bool: + """是否为特定的安装类型。""" + item_install_types = set(file_item.install_type.split(";")) + if "all" in item_install_types: + return True + return bool(item_install_types & install_types) + + +def get_install_path_dirs(install_path: str) -> Iterator[str]: + """获取安装路径父目录。""" + install_path = os.path.dirname(install_path) + while install_path not in ("", "/"): + yield install_path + install_path = os.path.dirname(install_path) + + +def get_missing_dir_set(filelist: FileList) -> set[str]: + """获取缺失目录集合。 + + 文件列表可能出现某一级目录缺失情况。 + 如配置了file_info:aaa/bbb/ccc.txt，但只配置了dir_info:aaa， + 那么缺失dir_info:aaa/bbb + """ + parent_dirs: set[str] = invoke( + pipe( + dispatch( + pipe( + partial( + filter, + partial(is_specific_operations, operations={"copy", "copy_entity"}), + ), + partial(map, attrgetter("relative_install_path")), + partial(filter, is_relative_install_path), + set, + partial(map, get_install_path_dirs), + chain.from_iterable, + ), + pipe( + partial(map, attrgetter("softlink")), + chain.from_iterable, + partial( + filter, + pipe( + dispatch( + bool, + is_relative_install_path, + partial(ne, "NA"), + ), + all, + ), + ), + set, + partial(map, get_install_path_dirs), + chain.from_iterable, + ), + pipe( + partial(map, attrgetter("pkg_inner_softlink")), + chain.from_iterable, + partial( + filter, + pipe( + dispatch( + bool, + partial(ne, "NA"), + ), + all, + ), + ), + set, + partial(map, get_install_path_dirs), + chain.from_iterable, + ), + ), + chain.from_iterable, + set, + ), + filelist, + ) + mkdir_installs: set[str] = { + file_item.relative_install_path + for file_item in filter(partial(is_specific_operations, operations={"mkdir"}), filelist) + if is_relative_install_path(file_item.relative_install_path) + } + + mkdir_parent_dirs: set[str] = set(itertools.chain.from_iterable(map(get_install_path_dirs, mkdir_installs))) + + missing_dir_set = sorted((parent_dirs | mkdir_parent_dirs) - mkdir_installs) + return set(missing_dir_set) + + +def print_missing_dir_set(missing_dir_set: set[str], in_msg: str = None) -> set[str]: + """打印缺失目录集合。""" + if in_msg: + tail_msg = f" {in_msg}" + else: + tail_msg = "" + for path in sorted(missing_dir_set): + CommLog.cilog_error(f'missing dir info path "{path}"{tail_msg}') + return missing_dir_set + + +def print_unsafe_paths(unsafe_paths: tuple[str, ...]) -> tuple[str, ...]: + """打印非安全路径。""" + for path in unsafe_paths: + CommLog.cilog_error(f'unsafe path "{path}" in move scene.') + return unsafe_paths + + +# 获取filelist中所有的特性集合 +get_features_in_filelist = pipe( + partial(map, attrgetter("feature")), + chain.from_iterable, # 展开集合序列为元素序列 + set, # 去重 + partial(filter, partial(ne, "comm")), # 排除comm特性 + set, +) + +# 获取filelist中所有的芯片集合 +get_chips_in_filelist = pipe( + partial(map, attrgetter("chip")), + chain.from_iterable, # 展开集合序列为元素序列 + set, # 去重 +) + + +def check_features_in_filelist(features: set[str], filelist: FileList) -> set[str]: + """检查文件列表中特性配置目录规范。""" + return invoke( + pipe( + # 过滤指定features的file_item + partial(filter, pipe(attrgetter("feature"), partial(and_, features), bool)), + list, + get_missing_dir_set, + partial(print_missing_dir_set, in_msg=f"in features {features}"), + ), + filelist, + ) + + +def check_chip_in_filelist(chip: str, filelist: FileList) -> set[str]: + """检查文件列表中芯片配置目录规范。""" + return invoke( + pipe( + # 过滤指定chip的file_item + partial( + filter, + any_( + pipe(attrgetter("chip"), not_), # 没有配置chip + pipe(attrgetter("chip"), partial(swap_args(contains), chip), bool), # 配置了指定chip + ), + ), + list, + get_missing_dir_set, + partial(print_missing_dir_set, in_msg=f"in chip {chip}"), + ), + filelist, + ) + + +check_filelist_features = any_( + pipe( + dispatch( + pipe( + get_features_in_filelist, + # 对于每个feature，与comm组成一个set + partial(map, lambda x: {x, "comm"}), + # 此时为feature集合序列 + ), + repeat, # 重复filelist + ), + tuple, + star_apply(zip), + # 此时为元组序列，元组的第1个元素是过滤的feature集合，第2个元素是filelist + partial(itertools.starmap, check_features_in_filelist), + # 此时为集合序列；合并为一个集合 + chain.from_iterable, + set, + ), + pipe( + dispatch( + get_chips_in_filelist, + repeat, # 重复filelist + ), + tuple, + star_apply(zip), + # 此时为元组序列，元组的第1个元素是chip集合，第2个元素是filelist + partial(itertools.starmap, check_chip_in_filelist), + # 此时为集合序列；合并为一个集合 + chain.from_iterable, + set, + ), +) + + +# 检查move是否安全，是否存在同一个源路径被mv多次 +check_move_safe = pipe( + partial( + filter, + partial(is_specific_operations, operations={"copy", "copy_entity", "move"}), + ), + partial(map, attrgetter("relative_path_in_pkg")), + Counter, + methodcaller("items"), + partial(filter, pipe(itemgetter(1), partial(lt, 1))), + partial(map, itemgetter(0)), + tuple, + print_unsafe_paths, +) + + +def check_filelist(filelist: FileList, check_features: bool, check_move: bool): + """检查文件列表是否符合规范。""" + if check_features: + check_features_func = check_filelist_features + else: + check_features_func = constant(set()) + + if check_move: + check_move_func = check_move_safe + else: + check_move_func = constant(tuple()) + + # 此处使用any_，短路部分报错 + check_func = any_( + pipe( + get_missing_dir_set, + print_missing_dir_set, + ), + pipe( + partial(filter, partial(is_specific_install_type, install_types={"run"})), + list, + get_missing_dir_set, + partial(print_missing_dir_set, in_msg="in run install type"), + ), + check_features_func, + check_move_func, + ) + missing = check_func(filelist) + + if missing: + raise FilelistError() + + +def get_common_path(args: list[str]) -> str: + """公共路径前缀。""" + try: + return os.path.commonpath(args) + except ValueError: + return "" + + +class FileItemRelation(IntEnum): + """文件条目之间的关系。""" + + NOT_NESTED = 0 # 不是嵌套文件 + NESTED = 1 # 嵌套文件 + SAME = 2 # 相同文件 + + +def is_nested_file_item(item: FileItem, base_item: FileItem) -> FileItemRelation: + """是否为嵌套的文件。""" + if base_item is None: + return FileItemRelation.NOT_NESTED + + if item == base_item: + return FileItemRelation.SAME + + install_path = item.relative_install_path + base_install_path = base_item.relative_install_path + + common_install_path = get_common_path([install_path, base_install_path]) + if common_install_path != base_install_path: + return FileItemRelation.NOT_NESTED + + pkg_path = item.relative_path_in_pkg + base_pkg_path = base_item.relative_path_in_pkg + + install_rel_path = os.path.relpath(install_path, base_install_path) + pkg_rel_path = os.path.relpath(pkg_path, base_pkg_path) + if install_rel_path != pkg_rel_path: + # 确保打包与安装相对路径一致 + raise FilelistError(f"nested paths {item} and {base_item} are illegal.") + return FileItemRelation.NESTED + + +def found_nested_file_item(item: FileItem, base_item: FileItem): + """发现嵌套元素。""" + raise FilelistError(f"found nested paths {item} and {base_item}!") + + +def convert_nested_path_in_filelist(filelist: FileList): + """filelist中嵌套路径元素转为del。""" + pre_item = None + for item in filelist: + ret = is_nested_file_item(item, pre_item) + if ret == FileItemRelation.NESTED: + yield item._replace(operation="del") + elif any((ret == FileItemRelation.NOT_NESTED, (ret == FileItemRelation.SAME and not item.is_dir))): + yield item + pre_item = item + + +# 检查文件列表中的嵌套路径。入参: filelist +check_nested_path_in_filelist = pipe( + partial(filter, partial(is_specific_operations, operations={"copy", "copy_entity"})), + partial(sorted, key=attrgetter("relative_install_path")), + pairwise, + partial(map, conditional_apply(star_apply(is_nested_file_item), star_apply(found_nested_file_item))), + list, +) + + +# 变换文件列表中嵌套路径。入参: filelist +transform_nested_path_in_filelist = pipe( + dispatch( + partial(itertools.filterfalse, partial(is_specific_operations, operations={"copy"})), + pipe( + partial(filter, partial(is_specific_operations, operations={"copy"})), + partial(sorted, key=attrgetter("relative_install_path")), + convert_nested_path_in_filelist, + ), + ), + chain.from_iterable, + list, + side_effect(check_nested_path_in_filelist), +) + + +def generate_filelist(filelist: FileList, filename: str): + """生成文件列表文件。""" + content_list = list( + itertools.chain([get_filelist_header_string()], [file_item_to_string(item) for item in filelist]) + ) + content = "\n".join(content_list) + filepath = os.path.join(TOP_DIR, "build", filename) + try: + with open(filepath, "w", encoding="utf-8") as file: + file.write(content) + # filelist.csv文件末尾补充一个换行符 + file.write("\n") + except OSError as ex: + raise GenerateFilelistError(filename) from ex + + +def get_transform_nested_path_func(parallel: bool) -> Callable[[FileList], FileList]: + """获取转换嵌套路径函数。""" + if parallel: + return transform_nested_path_in_filelist + return identity diff --git a/csrc/ascend/scripts/package/common/py/merge_binary_info_config.py b/csrc/ascend/scripts/package/common/py/merge_binary_info_config.py new file mode 100644 index 000000000..150fa2970 --- /dev/null +++ b/csrc/ascend/scripts/package/common/py/merge_binary_info_config.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""合并算子binary_info_config.json。""" + +import argparse +import json +import os +import sys + + +def load_json_file(json_file: str): + """加载json文件。""" + with open(json_file, encoding="utf-8") as file: + json_content = json.load(file) + return json_content + + +def save_json_file(output_file: str, content): + """保存json文件。""" + output_dir = os.path.dirname(output_file) + if not os.path.exists(output_dir): + os.makedirs(output_dir, exist_ok=True) + + with open(output_file, "w", encoding="utf-8") as file: + json.dump(content, file, ensure_ascii=True, indent=2) + + +def update_config(base_content, update_content): + """更新配置。""" + new_content = base_content.copy() + new_content.update(update_content) + return dict(sorted(new_content.items())) + + +def parse_args(argv: list[str]): + """入参解析。""" + parser = argparse.ArgumentParser() + parser.add_argument("--base-file", required=True, help="the basic binary_info_config file") + parser.add_argument("--update-file", required=True, help="the update binary_info_config file") + parser.add_argument( + "--output-file", required=True, type=os.path.realpath, help="the output binary_info_config file" + ) + args = parser.parse_args(argv) + return args + + +def main(argv: list[str]) -> bool: + """主流程。""" + args = parse_args(argv) + base_content = load_json_file(args.base_file) + update_content = load_json_file(args.update_file) + result = update_config(base_content, update_content) + save_json_file(args.output_file, result) + return True + + +if __name__ == "__main__": + if not main(sys.argv[1:]): # pragma: no cover + sys.exit(1) # pragma: no cover diff --git a/csrc/ascend/scripts/package/common/py/packer.py b/csrc/ascend/scripts/package/common/py/packer.py new file mode 100644 index 000000000..9e433a431 --- /dev/null +++ b/csrc/ascend/scripts/package/common/py/packer.py @@ -0,0 +1,240 @@ +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import os +import shutil +import subprocess +from argparse import Namespace +from collections.abc import Callable +from itertools import chain +from subprocess import PIPE, STDOUT +from typing import NamedTuple + +from .utils.comm_log import CommLog +from .utils.pkg_utils import CompressError + + +class PackageName: + """包名。""" + + def __init__(self, package_attr, args: Namespace, version: str): + self.product_name = package_attr.get("product_name") + self.chip_name = args.chip_name or package_attr.get("chip_name") + self.suffix = args.suffix or package_attr.get("suffix") + self.func_name = get_func_name(args.func_name, package_attr) + self.chip_plat = package_attr.get("chip_plat") + self.deploy_type = package_attr.get("deploy_type") + self.version = version.lower() + self.not_in_name_list = args.not_in_name.split(",") + self.os_arch = args.os_arch + self.package_suffix = args.package_suffix + self.ext_name = args.ext_name + if args.pkg_name_style == "underline": + self.name_sep = "_" + else: + self.name_sep = "-" + + def get_attribute(self, name: str) -> str | None: + """获取属性。""" + if name in self.not_in_name_list: + return None + return getattr(self, name) + + def getvalue(self) -> str: + product_name = self.get_attribute("product_name") + chip_name = self.get_attribute("chip_name") + func_name = self.get_attribute("func_name") + version = self.get_attribute("version") + os_arch = self.get_attribute("os_arch") + chip_plat = self.get_attribute("chip_plat") + deploy_type = self.get_attribute("deploy_type") + ext_name = self.get_attribute("ext_name") + package_suffix = "debug" if self.package_suffix == "debug" else None + + region1 = "-".join(filter(None, [product_name, remove_ascend(chip_name), func_name])) + region2 = ".".join(filter(None, [version])) + region3 = "-".join(filter(None, [os_arch, chip_plat, deploy_type, package_suffix, ext_name])) + package_name = "_".join(filter(None, [region1, region2, region3])) + + return f"{package_name}.{self.suffix}" + + +class MakeselfPkgParams(NamedTuple): + """run包打包参数。""" + + package_name: str + comments: str + makeself_tool: str | None = None + makeself_header: str | None = None + help_info: str | None = None + source_target: str | None = None + + install_script: str | None = None + independent_pkg: bool | None = False + cleanup: str | None = None + + +def remove_ascend(text): + if text is None: + return None + text_lower = text.lower() + if text_lower == "ascend910_93": + return "A3" + if "ascend" in text_lower: + return text_lower.replace("ascend", "") + return text_lower + + +def get_func_name(func_name: str, package_attr) -> str: + """获取包func_name。""" + return func_name or package_attr.get("func_name") + + +def get_compress_tool() -> str: + tools = ["pigz", "gzip", "bzip2", "xz"] + for tool in tools: + path = shutil.which(tool) + if path: + return "--" + tool + CommLog.cilog_error( + "The system does not come with a compression tool pre-installed." + "Please ensure at least one of the following compression tools is available: %s", + tools, + ) + return "" + + +def get_compress_format() -> str: + tar_format = "gnu" + path = shutil.which("bsdtar") + if path: + tar_format = "ustar" + return tar_format + + +def compose_makeself_command(params: MakeselfPkgParams) -> str: + """组装makeself包打包命令。""" + + def get_cleanup_commands() -> list[str]: + if params.cleanup: + return ["--cleanup", params.cleanup] + return [] + + independent_pkg = params.independent_pkg + compress_tool = get_compress_tool() + tar_format = get_compress_format() + if independent_pkg: + commands = chain( + [ + "TMPDIR=$pwd", + params.makeself_tool, + "--header", + params.makeself_header, + "--help-header", + params.help_info, + compress_tool, + "--complevel", + "4", + "--nomd5", + "--sha256", + "--nooverwrite", + "--chown", + "--tar-format", + tar_format, + "--tar-extra", + "--numeric-owner", + "--tar-quietly", + ], + get_cleanup_commands(), + [params.source_target, params.package_name, params.comments, params.install_script], + ) + else: + commands = chain( + [ + compress_tool, + "--complevel", + "4", + "--nomd5", + "--sha256", + "--nooverwrite", + "--chown", + "--tar-format", + tar_format, + "--tar-extra", + "--numeric-owner", + "--tar-quietly", + ], + get_cleanup_commands(), + [params.package_name, params.comments], + ) + + command = " ".join(commands) + return command + + +def create_makeself_pkg_params_factory( + source_target: str, package_name: str, comments: str +) -> Callable[[str, dict, bool], MakeselfPkgParams]: + """创建Makeself打包参数工厂。""" + + def create_makeself_pkg_params(makeself_dir: str, package_attr: dict, independent_pkg=False) -> MakeselfPkgParams: + """创建Makeself打包参数。""" + cleanup = package_attr.get("cleanup") + + if independent_pkg: + install_script = str(package_attr.get("install_script")) + help_info = str(package_attr.get("help")) + makeself_tool = os.path.join(makeself_dir, "makeself.sh") + makeself_header = os.path.join(makeself_dir, "makeself-header.sh") + params = MakeselfPkgParams( + package_name=package_name, + comments=comments, + makeself_tool=makeself_tool, + makeself_header=makeself_header, + help_info=help_info, + source_target=source_target, + install_script=install_script, + independent_pkg=independent_pkg, + cleanup=cleanup, + ) + else: + params = MakeselfPkgParams( + package_name=package_name, + comments=comments, + cleanup=cleanup, + ) + return params + + return create_makeself_pkg_params + + +def create_run_package_command(params: MakeselfPkgParams) -> tuple[str | None, str | None]: + """ + 功能描述: 组装打run包命令 + 返回值: command + """ + return compose_makeself_command(params), None + + +def exec_pack_cmd(delivery_dir: str, pack_cmd: str, package_name: str) -> str: + """执行打包命令""" + if delivery_dir: + cmd = f"cd {delivery_dir} && {pack_cmd}" + else: + cmd = pack_cmd + CommLog.cilog_info("package cmd:%s", cmd) + result = subprocess.run(cmd, shell=True, check=False, stdout=PIPE, stderr=STDOUT) + output = result.stdout.decode() + if result.returncode != 0: + CommLog.cilog_error(__file__, "compress package(%s) failed! %s.", package_name, output) + raise CompressError(package_name) + return package_name diff --git a/csrc/ascend/scripts/package/common/py/pkg_parser.py b/csrc/ascend/scripts/package/common/py/pkg_parser.py new file mode 100644 index 000000000..476f4ba3b --- /dev/null +++ b/csrc/ascend/scripts/package/common/py/pkg_parser.py @@ -0,0 +1,1058 @@ +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import copy +import glob +import hashlib +import itertools +import os +import sys +import xml.etree.ElementTree as ET +from argparse import Namespace +from collections.abc import Callable, Iterable, Iterator +from functools import partial +from io import StringIO +from itertools import chain +from operator import attrgetter, itemgetter, methodcaller +from typing import Any, NamedTuple + +import regex as re + +from .filelist import FileItem, FileList, fill_is_common_path +from .utils import pkg_utils +from .utils.comm_log import CommLog +from .utils.funcbase import constant, dispatch, invoke, pipe, star_apply +from .utils.pkg_utils import ( + BLOCK_CONFIG_PATH, + FAIL, + BlockConfigError, + ContainAsteriskError, + EnvNotSupported, + IllegalVersionDir, + PackageError, + ParseOsArchError, + config_feature_to_set, + flatten, + merge_dict, + star_pipe, + yield_if, +) +from .version_info import VersionFormatNotMatch, VersionInfo, VersionXml, is_multi_version + +# 环境变量字典 +EnvDict = dict[str, str] + +# 文件信息 +FileInfo = dict[str, str] + +# 包属性 +PackageAttr = dict[str, str | bool] + +# 生成信息 +GenerateInfo = dict[str, str] + + +class ParseOption(NamedTuple): + """解析参数。""" + + os_arch: str | None + pkg_version: str | None + build_type: str | None + package_check: bool + ext_name: str = "" + + +def parse_os_arch(os_arch: str) -> tuple[str, str, str]: + """解析系统和架构。""" + match = re.match("^([a-z]+)(\\d+(\\.\\d+)*)?[.-]?(\\S*)", os_arch) + if match: + os_name = match.group(1) + os_ver = match.group(2) + if match.group(4): + arch = match.group(4) + else: + # 如果os_arch中没有配置ARCH，ARCH默认值为aarch64 + arch = "aarch64" + + return os_name, os_ver, arch + + raise ParseOsArchError() + + +def replace_env(env_dict: EnvDict, in_str: str): + """替换环境变量为实际值。""" + env_list = re.findall(".*?\\$\\((.*?)\\).*?", in_str) + for env in env_list: + if env == "FILE": + continue + if env in env_dict: + if env_dict[env] is not None: + in_str = in_str.replace(f"$({env})", env_dict[env]) + else: + in_str = in_str.replace(f"$({env})", "") + else: + raise EnvNotSupported(f"Error: {env} not supported.") + return in_str + + +class ParseEnv(NamedTuple): + """解析上下文环境。""" + + env_dict: EnvDict + parse_option: ParseOption + delivery_dir: str + top_dir: str + + +class BlockElement(NamedTuple): + """块配置。""" + + name: str + block_conf_path: str + dst_path: str + chips: set[str] + features: set[str] + attrs: dict[str, str] + + +# BlockElement直接透传给LoadedBlockElement的参数列表 +BLOCK_ELEMENT_PASS_THROUGH_ARGS = ["dst_path", "chips", "features", "attrs"] + + +class LoadedBlockElement(NamedTuple): + """加载后的块配置。""" + + root_ele: ET.Element + use_move: bool + dst_path: str + chips: set[str] + features: set[str] + attrs: dict[str, str] + + +class FileInfoParsedResult(NamedTuple): + """file_info元素解析结果。""" + + file_info: FileInfo + move_infos: list[FileInfo] + dir_infos: list[dict[str, str]] + expand_infos: list[dict[str, str]] + + +class BlockConfig(NamedTuple): + """块配置。""" + + dir_install_list: list[dict] + move_files: list[FileInfo] + expand_content_list: list[dict] + package_content_list: list[dict] + generate_infos: list[GenerateInfo] + + +class PackerConfig(NamedTuple): + """安装相关配置。""" + + fill_is_common_path: Callable[[FileList], Iterator[FileItem]] + + +class XmlConfig(NamedTuple): + """安装xml配置。""" + + default_config: dict[str, str] + package_attr: PackageAttr + version_info: VersionInfo + blocks: list[BlockConfig] + version: str + version_xml: VersionXml | None + packer_config: PackerConfig + + def _collect_list(self, list_name): + result = [] + for block in self.blocks: + result.extend(getattr(block, list_name)) + return result + + @property + def dir_install_list(self): + return self._collect_list("dir_install_list") + + @property + def move_content_list(self): + return self._collect_list("move_files") + + @property + def expand_content_list(self): + return self._collect_list("expand_content_list") + + @property + def package_content_list(self): + return self._collect_list("package_content_list") + + @property + def generate_infos(self) -> list[GenerateInfo]: + return self._collect_list("generate_infos") + + +# 默认包属性 +DEFAULT_PACKAGE_ATTR = { + "gen_version_info": True, +} + + +def parse_package_info(package_info_ele: ET.Element | None) -> dict: + """解析package_info元素。""" + + def get_package_info_attrs(ele: ET.Element) -> Iterator[tuple[str, str | bool]]: + # expand_asterisk: 展开配置中星号 + # parallel: 并行复制文件 + # parallel_limit: 限制并发数 + # package_check: 检查filelist.csv中配置目录是否完整 + # check_features: 检查filelist.csv中所有feature是否符合package_check + # gen_version_info: 是否生成version.info文件 + bool_attrs = ( + "expand_asterisk", + "parallel", + "parallel_limit", + "package_check", + "check_features", + "use_move", + "gen_version_info", + ) + bool_values = ("t", "true", "y", "yes") + if ele.tag in bool_attrs: + if ele.text.lower() in bool_values: + yield ele.tag, True + else: + yield ele.tag, False + else: + yield ele.tag, ele.text + + if not package_info_ele: + return {} + + attr = dict(chain.from_iterable(map(get_package_info_attrs, list(package_info_ele)))) + + return attr + + +def parse_package_attr_by_args(args: Namespace) -> dict: + """通过命令行参数解析""" + + def pairs(): + if hasattr(args, "chip_name") and args.chip_name: + yield "chip_name", args.chip_name + if hasattr(args, "suffix") and args.suffix: + yield "suffix", args.suffix + if hasattr(args, "func_name") and args.func_name: + yield "func_name", args.func_name + + return dict(pairs()) + + +def parse_package_attr(root_ele: ET.Element, args: Namespace) -> dict: + """通过根元素解析package_info元素。""" + package_info_ele = root_ele.find("package_info") + return merge_dict( + DEFAULT_PACKAGE_ATTR, + parse_package_info(package_info_ele), + parse_package_attr_by_args(args), + ) + + +def render_cann_version( + a_ver: int, b_ver: int, c_ver: int | None, d_ver: int | None, e_ver: int | None, f_ver: int | None +) -> str: + """渲染CANN版本号。""" + buffer = StringIO() + buffer.write("(") + buffer.write(f"({a_ver + 1} * 100000000) + ({b_ver + 1} * 1000000)") + if c_ver is not None: + buffer.write(f" + ({c_ver + 1} * 10000)") + if d_ver is not None: + buffer.write(f" + (({d_ver + 1} * 100) + 5000)") + if e_ver is not None: + buffer.write(f" + ({e_ver + 1} * 100)") + if f_ver is not None: + buffer.write(f" + {f_ver}") + buffer.write(")") + return buffer.getvalue() + + +def render_semver(package_name: str, version: str) -> Iterator[tuple[str, str]]: + """ + 将语义化版本号转换为可比较的整数表达式，严格遵循SemVer规范 + + 排序规则： + 1. 正式版本 > 所有对应预发布版本（如 8.0.5 > 8.0.5-rc.1） + 2. 预发布类型优先级：rc > beta > alpha > 其他类型（如 rc.1 > beta.100） + 3. 同类型预发布版本：序号越大优先级越高（如 alpha.10 > alpha.2） + 4. 多段序号比较：从左到右逐段比较（如 alpha.1.2 > alpha.1.1） + """ + expr_buffer = StringIO() + expr_buffer.write("(") + + # 移除构建元数据（+后面的内容不影响版本优先级） + version = version.split("+")[0] + + # 分离正式版本和预发布版本 + pre_release = None + if "-" in version: + release_part, pre_release = version.split("-", 1) + release_part = release_part.split(".") + else: + release_part = version.split(".") + if len(release_part) > 3: + pre_release = ".".join(release_part[3:]) + release_part = release_part[:3] + + # 解析正式版本号（主版本.次版本.修订号） + try: + major, minor, patch = map(int, release_part) + except (ValueError, TypeError) as ex: + raise IllegalVersionDir(f"无效的版本号格式: {version}") from ex + + yield f"{package_name}_VERSION_STR", f'"{version}"' + yield f"{package_name}_MAJOR", str(major) + yield f"{package_name}_MINOR", str(minor) + yield f"{package_name}_PATCH", str(patch) + # 计算基础版本值（主版本*10^7 + 次版本*10^5 + 修订号*10^3） + # 预留10^3空间用于预发布版本，确保不同正式版本区间不重叠 + expr_buffer.write(f"({major} * 10000000) + ({minor} * 100000) + ({patch} * 1000)") + + # 处理正式版本（无预发布部分） + if not pre_release: + expr_buffer.write(")") + yield f"{package_name}_PRERELEASE", '""' + yield f"{package_name}_VERSION_NUM", expr_buffer.getvalue() + return + + yield f"{package_name}_PRERELEASE", f'"{pre_release}"' + + # 预发布类型权重（值越小优先级越高） + type_weights = { + "rc": 100, # rc优先级最高 + "beta": 200, # beta次之 + "alpha": 300, # alpha最低 + } + + def calc_pre_release() -> tuple[int, int]: + """计算预发布版本。""" + if "." in pre_release: + pre_parts = pre_release.split(".") + pre_type = pre_parts[0] # 提取预发布类型（rc/beta/alpha等） + + # 提取序号部分（支持多段序号，非数字部分忽略） + pre_nums = [] + for part in pre_parts[1:]: + if part.isdigit(): + pre_nums.append(int(part)) + if not pre_nums: # 无序号时默认0 + pre_nums = [0] + + # 未知类型权重设为400（优先级低于alpha） + pre_type_weight = type_weights.get(pre_type, 400) + + # 计算序号值（支持多段和多位数） + num_str = "".join(map(str, pre_nums)) + # 转换为整数并返回 + num_value = int(num_str) + return pre_type_weight, num_value + + for pre_type in type_weights: + if pre_release.startswith(pre_type): + pre_type_weight = type_weights[pre_type] + num_value = int(pre_release[len(pre_type) :]) + return pre_type_weight, num_value + + return None, None + + try: + pre_type_weight, num_value = calc_pre_release() + except (ValueError, TypeError) as ex: + raise IllegalVersionDir(f"无效的预发布版本: {pre_release}") from ex + + if not pre_type_weight: + raise IllegalVersionDir(f"无效的预发布版本: {pre_release}") + + # 预发布版本最终值 = 基础值 - 类型权重 + 序号值 + # 确保：预发布值 < 基础值（正式版本） + expr_buffer.write(f" - {pre_type_weight} + {num_value}") + expr_buffer.write(")") + + yield f"{package_name}_VERSION_NUM", expr_buffer.getvalue() + + +def get_cann_version_info(name: str, version: str) -> Iterator[tuple[str, str]]: + """获取CANN版本号信息。""" + # 删除字符串中的_VERSION + package_name = name[:-8] + + if not version: + yield f"{package_name}_VERSION_STR", '"0"' + return + + yield from render_semver(package_name, version) + + +def get_default_env_items() -> Iterator[tuple[str, str]]: + """获取默认环境字典条目。""" + yield "VERSION_DIR", "" + yield "HOME", os.environ.get("HOME") + + +def get_env_items_by_version(version: str | None) -> Iterator[tuple[str, str]]: + """根据version获取环境字典条目。""" + if version: + yield "ASCEND_VER", version + + version_parts = version.split(".") + for idx in range(1, len(version_parts) + 1): + yield f"CUR_VER[{idx}]", ".".join(version_parts[:idx]) + yield "CUR_VER", version + yield "LOWER_CUR_VER", version.lower() + + +def get_env_items_by_version_dir(version_dir: str | None) -> Iterator[tuple[str, str]]: + """根据version_dir获取环境字典条目。""" + if version_dir: + yield "VERSION_DIR", version_dir + + +def get_os_arch_default_env_items() -> Iterator[tuple[str, str]]: + """获取系统相关默认环境字典条目。""" + yield "OS_NAME", "linux" + yield "OS_VER", "" + yield "ARM", "aarch64" + yield "TARGET_ENV", "$(TARGET_ENV)" + + +def get_env_items_by_os_arch(os_arch: str) -> Iterator[tuple[str, str]]: + """根据os_arch获取环境字典条目。""" + if os_arch: + os_name, os_ver, arch = parse_os_arch(os_arch) + yield "OS_NAME", os_name + yield "OS_VER", os_ver + yield "ARCH", arch + yield "OS_ARCH", os_arch + if arch in ("arm", "sw_64"): + yield "ARM", arch + else: + yield "ARM", "aarch64" + yield "TARGET_ENV", f"{arch}-linux" + else: + yield from get_os_arch_default_env_items() + + +def get_env_items_by_timestamp(timestamp: str | None) -> Iterator[tuple[str, str]]: + """根据timestamp获取环境字典条目。""" + if timestamp: + yield "TIMESTAMP", timestamp + yield "TIMESTAMP_NO", timestamp.replace("_", "") + else: + yield "TIMESTAMP", "0" + yield "TIMESTAMP_NO", "0" + + +def parse_env_dict( + os_arch: str, package_attr: PackageAttr, version: str | None, version_dir: str | None, timestamp: str | None +) -> EnvDict: + """解析环境变量字典。""" + env_dict = dict( + chain( + get_default_env_items(), + yield_if(("ARCH", package_attr.get("default_arch")), itemgetter(1)), + get_env_items_by_os_arch(os_arch), + get_env_items_by_version(version), + get_env_items_by_version_dir(version_dir), + yield_if(("VERSION_DIR", version_dir), constant(version_dir)), + get_env_items_by_timestamp(timestamp), + ) + ) + + return env_dict + + +def get_timestamp(args: Namespace) -> str | None: + """获取触发时间戳。""" + if "tag" not in args: + return None + + tag = args.tag + if tag: + timestamp_re = r"\d{8}_\d{9}" + timestamp_list = re.findall(timestamp_re, tag) + if not timestamp_list: + raise PackageError("The {} format is incorrect.".format(tag)) + timestamp = timestamp_list[-1] + else: + timestamp = None + return timestamp + + +def extract_element_attrib(ele: ET.Element) -> dict: + """提取元素属性。""" + return ele.attrib.copy() + + +def extract_generate_info_content(generate_info_ele: ET.Element, env_dict: EnvDict) -> dict: + """提取生成信息内容。""" + file_content = {sub_item.tag: replace_env(env_dict, sub_item.text) for sub_item in list(generate_info_ele)} + return {"content": file_content} + + +def parse_generate_infos_by_loaded_block( + loaded_block: LoadedBlockElement, default_config: dict[str, str], env_dict: EnvDict +) -> list[dict]: + """根据根元素解析生成信息列表。""" + return invoke( + pipe( + partial( + map, + pipe( + dispatch( + pipe( + extract_element_attrib, + partial(merge_dict, default_config), + partial(evaluate_info, loaded_block=loaded_block, env_dict=env_dict), + ), + partial(extract_generate_info_content, env_dict=env_dict), + ), + star_apply(merge_dict), + ), + ), + list, + ), + loaded_block.root_ele.findall("generate_info"), + ) + + +def join_pkg_inner_softlink(link_str_list: list[str]) -> str: + """合并pkg_inner_softlink""" + path = "/".join(link_str_list) + return os.path.normpath(path) + + +def check_contain_asterisk(value: str) -> bool: + """检查串是否包含星号。""" + return "*" in value + + +def check_value(value: str, package_check: bool, package_attr: PackageAttr): + """检查元素value属性。""" + if package_check and package_attr.get("suffix") == "run": + if check_contain_asterisk(value): + raise ContainAsteriskError(value) + + +def get_dst_prefix(file_info: FileInfo, env: ParseEnv) -> str: + """获取文件的前缀。""" + return os.path.join(env.delivery_dir, file_info["dst_path"]) + + +def get_dst_target(file_info: FileInfo, env: ParseEnv) -> str: + """获取文件的实际路径。""" + dst_prefix = get_dst_prefix(file_info, env) + return os.path.join(dst_prefix, os.path.basename(file_info.get("value"))) + + +def make_hash(filepath: str) -> str: + """计算文件的hash(sha256)值。""" + sha256_hash = hashlib.sha256() + with open(filepath, "rb") as file: + sha256_hash.update(file.read()) + + return sha256_hash.hexdigest() + + +def config_hash(parsed_result: FileInfoParsedResult, env: ParseEnv): + """配置hash值。""" + file_info = parsed_result.file_info + # 如果配置了configurable，需要计算文件的hash值 + if file_info and file_info["configurable"] == "TRUE": + src_target = get_dst_target(file_info, env) + hash_value = make_hash(src_target) + file_info["hash"] = hash_value + return parsed_result + + +def apply_func(func: Callable[[str], str], value: list[str] | set[str] | str) -> list[str] | set[str] | str: + """对一个字符串，或字符串序列，应用函数。""" + # 如：pkg_softlink列表 + if isinstance(value, list): + return list(map(func, value)) + # 如：feature集合 + if isinstance(value, set): + return set(map(func, value)) + return func(value) + + +REAL_PREFIX = "real:" + + +def join_dst_path(base: str, other: str) -> str: + """联结dst_path。""" + if other.startswith("real:"): + other = other[len(REAL_PREFIX) :] + return other + return os.path.join(base, other) + + +def evaluate_info(info: dict[str, str], loaded_block: LoadedBlockElement, env_dict: EnvDict) -> dict[str, str]: + """info元素求值。""" + dst_keys = ("dst_path",) + + replace_env_func = partial(replace_env, env_dict) + add_dst_path_func = partial(join_dst_path, loaded_block.dst_path) + + def upper_value(key: str, value: str) -> tuple[str, str]: + if key == "configurable": + return key, value.upper() + return key, value + + def add_dst_path(key: str, value: str) -> tuple[str, str]: + if key in dst_keys: + return key, apply_func(add_dst_path_func, value) + return key, value + + def replace_pkg_inner_softlink(key: str, value: str) -> tuple[str, str]: + if key == "pkg_inner_softlink": + return key, "NA" + return key, value + + def merge_feature(key: str, value: str) -> tuple[str, str]: + if key in ("chip", "feature"): + config_features = config_feature_to_set(value, key) + return key, config_features | getattr(loaded_block, f"{key}s") + return key, value + + def eval_value(_key: str, value: str) -> str: + if value is not None: + return apply_func(replace_env_func, value) + + eval_value_func = star_pipe( + upper_value, + add_dst_path, + replace_pkg_inner_softlink, + merge_feature, + eval_value, + ) + + return { + key: eval_value_func(key, value) + for key, value in itertools.chain( + # 默认值配置 + [ + ("dst_path", ""), + ("configurable", "FALSE"), + ("chip", None), + ("feature", None), + ("pkg_feature", None), + ], + info.items(), + ) + } + + +def parse_dir_info_elements( + loaded_block: LoadedBlockElement, default_config: dict[str, str], package_attr: PackageAttr, env: ParseEnv +) -> list[dict[str, str]]: + """解析dir_info元素。""" + dir_info_elements: list[ET.Element] = loaded_block.root_ele.findall("dir_info") + dir_infos = [] + for item in dir_info_elements: + dir_config = default_config.copy() + dir_config.update(item.attrib) + dir_config["module"] = dir_config.get("value") + for sub_item in list(item): + dir_info = dir_config.copy() + dir_info.update(sub_item.attrib) + dir_info = evaluate_info(dir_info, loaded_block, env.env_dict) + check_value(dir_info["value"], env.parse_option.package_check, package_attr) + dir_infos.append(dir_info) + + return dir_infos + + +def expand_dir(file_info: FileInfo, get_dst_target_func: Callable[[FileInfo], str]): + """ + 如果file_info中配置的路径是文件夹，需要展开到文件 + """ + file_info_list = [] + dir_info_list = [] + dst_target = get_dst_target_func(file_info) + + value_list = file_info.get("value").split("/") + target_name = value_list[-1] if value_list[-1] else value_list[-2] + + # 这里把当前目录也加入到dir_info_list中 + dir_info_copy = file_info.copy() + dir_info_copy["module"] = file_info.get("value") + dir_info_copy["value"] = os.path.join(file_info.get("install_path", ""), target_name) + + # 子目录的权限按照xml中subdir_mod配置，如果没有配置subdir_mod按照install_mod配置 + subdir_mod = file_info.get("subdir_mod", None) + if subdir_mod is not None: + dir_info_copy["install_mod"] = subdir_mod + # 被展开的当前目录不需要设置softlink + dir_info_copy["install_softlink"] = "NA" + dir_info_copy["pkg_inner_softlink"] = "NA" + dir_info_list.append(dir_info_copy) + + for root, dirs, files in os.walk(dst_target, followlinks=True): + # 不同操作系统上，os.walk遍历的结果顺序会略有不同，这里按字母排序，保证不同系统一致 + dirs.sort() + files.sort() + + dirs_to_remove = [] + for name in dirs: + dirname = os.path.join(root, name) + # 如果是指向目录的软连接，则按照文件处理，无需在安装时创建目录，只需要卸载时删除就行 + if os.path.islink(dirname) and not need_dereference(file_info): + copy_file_info = create_file_info(dirname, dst_target, file_info, name, target_name) + # 被展开的子文件不需要设置softlink + copy_file_info["install_softlink"] = "NA" + copy_file_info["pkg_inner_softlink"] = "NA" + file_info_list.append(copy_file_info) + dirs_to_remove.append(name) + continue + relative_dirname = os.path.relpath(dirname, dst_target) + dir_info_copy = file_info.copy() + dir_info_copy["module"] = file_info.get("value") + dir_info_copy["value"] = os.path.join(file_info.get("install_path", ""), target_name, relative_dirname) + # 被展开的子目录不需要设置softlink + dir_info_copy["install_softlink"] = "NA" + dir_info_copy["pkg_inner_softlink"] = "NA" + # 子目录的权限按照xml中subdir_mod配置，如果没有配置subdir_mod按照install_mod配置 + subdir_mod = file_info.get("subdir_mod", None) + if subdir_mod is not None: + dir_info_copy["install_mod"] = subdir_mod + dir_info_list.append(dir_info_copy) + for name in files: + filename = os.path.join(root, name) + copy_file_info = create_file_info(filename, dst_target, file_info, name, target_name) + file_info_list.append(copy_file_info) + + for name in dirs_to_remove: + dirs.remove(name) + return file_info_list, dir_info_list + + +def create_file_info(dirname, dst_target, file_info, name, target_name): + relative_filename = os.path.relpath(dirname, dst_target) + relative_dir_name = os.path.split(relative_filename)[0] + copy_file_info = file_info.copy() + copy_file_info["value"] = name + copy_file_info["src_path"] = os.path.join(file_info["src_path"], file_info["value"], relative_dir_name) + copy_file_info["dst_path"] = os.path.join(file_info["dst_path"], target_name, relative_dir_name) + copy_file_info["install_path"] = os.path.join(file_info.get("install_path", ""), target_name, relative_dir_name) + return copy_file_info + + +def expand_file_info_asterisk(parsed_result: FileInfoParsedResult, env: ParseEnv) -> Iterator[FileInfoParsedResult]: + """展开FileInfoParsedResult中的星号。""" + file_info = parsed_result.file_info + if check_contain_asterisk(file_info.get("value", "")): + dst_prefix = get_dst_prefix(file_info, env) + dst_targets = sorted(glob.glob(get_dst_target(file_info, env))) + if "exclude" in file_info: + exclude = list(map(methodcaller("strip"), file_info["exclude"].split(";"))) + else: + exclude = [] + for dst_target in dst_targets: + value = os.path.relpath(dst_target, dst_prefix) + if value in exclude: + continue + new_file_info = file_info.copy() + new_file_info["value"] = value + if "pkg_inner_softlink" in new_file_info: + # pkg_inner_softlink中的特殊变量$(FILE)替换为展开后的文件名 + pkg_inner_softlink = new_file_info["pkg_inner_softlink"] + new_file_info["pkg_inner_softlink"] = pkg_inner_softlink.replace( + "$(FILE)", os.path.basename(dst_target) + ) + yield parsed_result._replace(file_info=new_file_info) + else: + yield parsed_result + + +def trans_to_stream(item: Any) -> Iterator[Any]: + """转换为流。""" + yield item + + +def need_dereference(file_info: FileInfo) -> bool: + """是否需要解引用。""" + return "dereference" in file_info + + +def need_expand(file_info: FileInfo, get_dst_target_func: Callable[[FileInfo], str]) -> bool: + """是否需要展开子目录。""" + if file_info.get("entity") == "true": + return False + dst_target = get_dst_target_func(file_info) + if os.path.isdir(dst_target): + if need_dereference(file_info): + return True + return not os.path.islink(dst_target) + return False + + +def expand_file_info( + parsed_result: FileInfoParsedResult, use_move: bool, get_dst_target_func: Callable[[FileInfo], str] +) -> FileInfoParsedResult: + """展开FileInfoParsedResult中的目录。""" + file_info = parsed_result.file_info + if need_expand(file_info, get_dst_target_func): + # 如果当前是文件夹，需要展开计算 + expand_infos, dir_infos = expand_dir(file_info, get_dst_target_func) + # 实测发现，对于opp包，整体目录cp的安装速度要快于目录中各文件mv + # 可能的原因是，cp遍历目录的速度较快，并且目录中的文件都比较小。mv依赖shell迭代目录中的所有文件。 + return FileInfoParsedResult(merge_dict(file_info, {"is_dir": True}), [], dir_infos, expand_infos) + + if use_move: + return FileInfoParsedResult({}, [file_info], parsed_result.dir_infos, parsed_result.expand_infos) + + return parsed_result + + +def trans_file_info_to_result(file_info: FileInfo) -> FileInfoParsedResult: + """file_info转换为FileInfoParsedResult。""" + return FileInfoParsedResult(file_info, [], [], []) + + +def parse_file_element( + file_ele: ET.Element, + file_config: dict[str, str], + loaded_block: LoadedBlockElement, + package_attr: PackageAttr, + env: ParseEnv, +) -> Iterator[FileInfoParsedResult]: + """解析file元素。""" + file_info = merge_dict(file_config, file_ele.attrib) + file_info = evaluate_info(file_info, loaded_block, env.env_dict) + + if package_attr.get("expand_asterisk", False): + expand_asterisk_func = partial(expand_file_info_asterisk, env=env) + else: + expand_asterisk_func = trans_to_stream + + if "install_path" not in file_info: + file_info["install_path"] = "" + + trans_file_info_func = pipe( + trans_file_info_to_result, + expand_asterisk_func, + partial(map, partial(config_hash, env=env)), + partial( + map, + partial( + expand_file_info, use_move=loaded_block.use_move, get_dst_target_func=partial(get_dst_target, env=env) + ), + ), + ) + + yield from trans_file_info_func(file_info) + + +def parse_file_info_elements( + loaded_block: LoadedBlockElement, default_config: dict[str, str], package_attr: PackageAttr, env: ParseEnv +) -> Iterator[FileInfoParsedResult]: + """解析file_info元素。""" + file_info_elements: list[ET.Element] = loaded_block.root_ele.findall("file_info") + for file_info_ele in file_info_elements: + file_config = merge_dict(default_config, file_info_ele.attrib, {"module": file_info_ele.attrib.get("value")}) + + for sub_item in list(file_info_ele): + yield from parse_file_element(sub_item, file_config, loaded_block, package_attr, env) + + +def unique_infos(infos: Iterable) -> list[dict[str, str]]: + """infos去重。""" + cache: set[str] = set() + new_infos = [] + for info in infos: + if info["value"] in cache: + continue + cache.add(info["value"]) + new_infos.append(info) + + return new_infos + + +def parse_block_config(loaded_block: LoadedBlockElement, package_attr: PackageAttr, parse_env: ParseEnv): + """解析块配置。""" + default_config = copy.copy(loaded_block.attrs) + default_config.update(loaded_block.root_ele.attrib) + + dir_infos = parse_dir_info_elements( + loaded_block, + default_config, + package_attr, + parse_env, + ) + file_info_results = list( + chain( + parse_file_info_elements( + loaded_block, + default_config, + package_attr, + parse_env, + ) + ) + ) + + generate_infos = parse_generate_infos_by_loaded_block(loaded_block, default_config, parse_env.env_dict) + + return BlockConfig( + unique_infos(itertools.chain(dir_infos, flatten(result.dir_infos for result in file_info_results))), + list(flatten(map(attrgetter("move_infos"), file_info_results))), + list(flatten(map(attrgetter("expand_infos"), file_info_results))), + [result.file_info for result in file_info_results if result.file_info], + generate_infos, + ) + + +def make_loaded_block_element(root_ele: ET.Element, dst_path: str = "") -> LoadedBlockElement: + """创建加载后的块配置。""" + return LoadedBlockElement(root_ele, False, dst_path, set(), set(), {}) + + +def parse_block_element(block_ele: ET.Element, block_info_attr: dict[str, str]) -> BlockElement: + """解析单个块配置。""" + + def filter_attrs(attrs: dict[str, str]) -> dict[str, str]: + # block属性中过滤掉dst_path与block_conf_path + # dst_path由单独的参数传递 + # block中不需要block_conf_path + return {key: value for key, value in attrs.items() if key not in ("dst_path", "block_conf_path")} + + def with_merged_attrs(attrs: dict[str, str]) -> BlockElement: + name = attrs.get("name") + block_conf_path = attrs.get("block_conf_path") + + if not name: + raise BlockConfigError("block's name is not set!") + + if not block_conf_path: + raise BlockConfigError("block's conf_path is not set!") + + return BlockElement( + name=name, + block_conf_path=block_conf_path, + dst_path=attrs.get("dst_path", ""), + chips=config_feature_to_set(attrs.get("chip"), "chip"), + features=config_feature_to_set(attrs.get("feature"), "feature"), + attrs=filter_attrs(attrs), + ) + + return with_merged_attrs(merge_dict(block_info_attr, block_ele.attrib)) + + +def parse_block_info(block_info: ET.Element) -> list[BlockElement]: + """解析块配置。""" + + def parse_block_elements(block_elements: list[ET.Element]) -> list[BlockElement]: + return [parse_block_element(block_ele, block_info.attrib) for block_ele in block_elements] + + return parse_block_elements(list(block_info)) + + +def get_block_filepath(block_element: BlockElement) -> str: + """获取块配置路径。""" + return os.path.join( + pkg_utils.TOP_SOURCE_DIR, BLOCK_CONFIG_PATH, block_element.block_conf_path, f"{block_element.name}.xml" + ) + + +def load_block_element(package_attr: PackageAttr, block_element: BlockElement) -> LoadedBlockElement: + """加载块配置。""" + + def with_filepath(block_xml: str): + if not os.path.exists(block_xml): + raise BlockConfigError(f"block's config xml {block_xml} does not exist!") + + try: + return LoadedBlockElement( + root_ele=ET.parse(block_xml).getroot(), + use_move=package_attr.get("use_move", False), + **{name: getattr(block_element, name) for name in BLOCK_ELEMENT_PASS_THROUGH_ARGS}, + ) + except Exception: + raise BlockConfigError(f"dependent block configuration {block_xml} parse failed!") + + return with_filepath(get_block_filepath(block_element)) + + +def parse_blocks(root_ele: ET.Element, package_attr: PackageAttr, parse_env: ParseEnv) -> list[BlockConfig]: + """解析块列表。""" + return [ + parse_block_config(loaded_block, package_attr, parse_env) + for loaded_block in itertools.chain( + [make_loaded_block_element(root_ele)], + map( + partial(load_block_element, package_attr), + chain.from_iterable(map(parse_block_info, root_ele.findall("block_info"))), + ), + ) + ] + + +def read_version_info() -> tuple[str, str]: + version_path = os.path.join(pkg_utils.TOP_DIR, "version.info") + with open(version_path) as file: + line1 = file.readline().strip() + line2 = file.readline().strip() + version = line1.split("=")[1] + version_dir = line2.split("=")[1] + m = re.match(r"[.a-zA-Z0-9]+$", version) or re.match(r"[-a-zA-Z.0-9]+$", version) + if not m: + raise VersionFormatNotMatch() + + return version, version_dir + + +def parse_xml_config(filepath: str, delivery_dir: str, parse_option: ParseOption, args: Namespace) -> XmlConfig: + """解析打包xml配置。""" + try: + tree = ET.parse(filepath) + xml_root = tree.getroot() + except ET.ParseError as ex: + CommLog.cilog_error("xml parse %s failed: %s!", filepath, ex) + sys.exit(FAIL) + + default_config = xml_root.attrib.copy() + + package_attr = parse_package_attr(xml_root, args) + if args.version_dir: + version = args.version_dir + version_dir = args.version_dir + else: + version, version_dir = read_version_info() + if args.disable_multi_version: + version_dir = None + timestamp = get_timestamp(args) + try: + env_dict = parse_env_dict(parse_option.os_arch, package_attr, version, version_dir, timestamp) + except ParseOsArchError: + CommLog.cilog_error("os_arch %s is not correctly configured: %s!", parse_option.os_arch, filepath) + sys.exit(FAIL) + + parse_env = ParseEnv(env_dict, parse_option, delivery_dir, pkg_utils.TOP_SOURCE_DIR) + + blocks = parse_blocks(xml_root, package_attr, parse_env) + + if is_multi_version(version_dir): + fill_is_common_path_func = partial(fill_is_common_path, target_env=env_dict.get("TARGET_ENV")) + else: + fill_is_common_path_func = iter + + return XmlConfig(default_config, package_attr, None, blocks, version, None, PackerConfig(fill_is_common_path_func)) diff --git a/csrc/ascend/scripts/package/common/py/utils/comm_log.py b/csrc/ascend/scripts/package/common/py/utils/comm_log.py new file mode 100644 index 000000000..20b6f4b83 --- /dev/null +++ b/csrc/ascend/scripts/package/common/py/utils/comm_log.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import inspect +import logging +import time + +logging.basicConfig( + format="[%(asctime)s] [%(levelname)s] [%(pathname)s] [line:%(lineno)d] %(message)s", level=logging.INFO +) + + +class CommLog: + @staticmethod + def cilog_get_timestamp(): + return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + + @staticmethod + def cilog_print_element(cilog_element): + print("[" + cilog_element + "]", end=" ") + return + + @staticmethod + def cilog_logmsg(log_level, filename, line_no, log_msg, *log_paras): + log_timestamp = CommLog.cilog_get_timestamp() + CommLog.cilog_print_element(log_timestamp) + CommLog.cilog_print_element(log_level) + CommLog.cilog_print_element(filename) + CommLog.cilog_print_element(str(line_no)) + print(log_msg % log_paras[0]) + return + + @staticmethod + def cilog_error(log_msg, *log_paras): + frame = inspect.currentframe().f_back + line_no = frame.f_lineno + filename = frame.f_code.co_filename + CommLog.cilog_logmsg("ERROR", filename, line_no, log_msg, log_paras) + return + + @staticmethod + def cilog_warning(log_msg, *log_paras): + frame = inspect.currentframe().f_back + line_no = frame.f_lineno + filename = frame.f_code.co_filename + CommLog.cilog_logmsg("WARNING", filename, line_no, log_msg, log_paras) + return + + @staticmethod + def cilog_info(log_msg, *log_paras): + frame = inspect.currentframe().f_back + line_no = frame.f_lineno + filename = frame.f_code.co_filename + CommLog.cilog_logmsg("INFO", filename, line_no, log_msg, log_paras) + return diff --git a/csrc/ascend/scripts/package/common/py/utils/funcbase.py b/csrc/ascend/scripts/package/common/py/utils/funcbase.py new file mode 100644 index 000000000..7a03eddcd --- /dev/null +++ b/csrc/ascend/scripts/package/common/py/utils/funcbase.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""函数基础库。""" + +import operator +from collections.abc import Callable, Iterator +from typing import TypeVar + +A = TypeVar("A") + + +def constant(value: A) -> Callable[..., A]: + """常量值。""" + + def constant_inner(*_args, **_kwargs) -> A: + return value + + return constant_inner + + +def dispatch(*funcs): + """分派应用。""" + + def dispatch_inner(*args, **kwargs) -> Iterator: + return (func(*args, **kwargs) for func in funcs) + + return dispatch_inner + + +def pipe(*funcs): + """串联多个函数。""" + + def pipe_func(*args, **k_args): + result = funcs[0](*args, **k_args) + for func in funcs[1:]: + result = func(result) + return result + + return pipe_func + + +def identity(value: A) -> A: + """同一。""" + return value + + +def invoke(func, *args, **kwargs): + """调用。""" + return func(*args, **kwargs) + + +def side_effect(*funcs): + """调用函数，产生副作用，但不影响管道结果。""" + + def side_effect_func(arg): + for func in funcs: + # 不保留结果 + func(arg) + return arg + + return side_effect_func + + +def star_apply(func): + """列表展开再应用。""" + + def star_apply_func(arg): + return func(*arg) + + return star_apply_func + + +def any_(*funcs) -> Callable: + """高阶any。 + 注意，any有短路效果。""" + return pipe( + dispatch(*funcs), + any, + ) + + +def not_(func) -> Callable: + """高阶not。""" + return pipe(func, operator.not_) diff --git a/csrc/ascend/scripts/package/common/py/utils/pkg_utils.py b/csrc/ascend/scripts/package/common/py/utils/pkg_utils.py new file mode 100644 index 000000000..ad6026ace --- /dev/null +++ b/csrc/ascend/scripts/package/common/py/utils/pkg_utils.py @@ -0,0 +1,181 @@ +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""基础构件。""" + +import os +from collections.abc import Callable, Iterator +from itertools import chain, tee +from pathlib import Path +from typing import Optional, TypeVar + +TOP_DIR = str(Path(__file__).resolve().parents[5]) +TOP_SOURCE_DIR = TOP_DIR + "/scripts/" +DELIVERY_PATH = "build/_CPack_Packages/makeself_staging" +CONFIG_SCRIPT_PATH = "package" +BLOCK_CONFIG_PATH = "package/module" + +SUCCESS = 0 +FAIL = -1 + + +A = TypeVar("A") + + +class PackageError(Exception): + """打包异常基类。""" + + +class PackageConfigError(PackageError): + """打包配置错误异常。""" + + +class BlockConfigError(PackageError): + """块配置错误异常。""" + + +class ParseOsArchError(PackageError): + """解析os_arch失败异常。""" + + +class EnvNotSupported(PackageError): + """环境变量不支持异常。""" + + +class ContainAsteriskError(PackageError): + """包含星号异常。""" + + def __init__(self, value: str): + super().__init__() + self.value = value + + +class FilelistError(PackageError): + """文件列表异常。""" + + +class UnknownOperateTypeError(PackageError): + """未知的操作类型。""" + + +class PackageNameEmptyError(PackageError): + """包名为空错误。""" + + +class GenerateFilelistError(PackageError): + """生成文件列表文件异常。""" + + def __init__(self, filename: str): + super().__init__() + self.filename = filename + + +class IllegalVersionDir(PackageError): + """version_dir配置错误。""" + + +class CompressError(PackageError): + """打包错误。""" + + def __init__(self, package_name: str | None): + super().__init__(package_name) + self.package_name = package_name + + +def flatten(list_of_lists): + """Flatten one level of nesting""" + return chain.from_iterable(list_of_lists) + + +def merge_dict(base: dict, *news: dict): + """合并两个字典。""" + result = base.copy() + for new in news: + result.update(new) + return result + + +def star_pipe(*funcs): + """串联多个函数。解包结果。""" + + def pipe_func(*args, **k_args): + result = funcs[0](*args, **k_args) + for func in funcs[1:]: + # 解包元组或列表结果 + result = func(*result) + return result + + return pipe_func + + +def swap_args(func): + """交换函数前两个参数。""" + + def inner(fst, snd, *args, **k_args): + return func(snd, fst, *args, **k_args) + + return inner + + +def conditional_apply(predicate, func): + """条件下应用函数。""" + + def conditional_apply_func(arg): + if predicate(arg): + return func(arg) + return arg + + return conditional_apply_func + + +def pairwise(iterable): + """s -> (s0,s1), (s1,s2), (s2, s3), ...""" + a, b = tee(iterable) + next(b, None) + return zip(a, b) + + +def path_join(base: Optional, *others: str) -> Optional: + """路径联合。""" + if base is None: + return None + return os.path.join(base, *others) + + +def yield_if(data, predicate: Callable) -> Iterator: + """条件满足则产生。""" + if predicate(data): + yield data + + +def config_feature_to_set(feature_str: str, feature_type: str = "feature") -> set[str]: + """配置feature转换为集合。""" + if feature_str is None: + return set() + + if isinstance(feature_str, set): + return feature_str + + if feature_str == "": + raise PackageConfigError(f"Not allow to config {feature_type} empty.") + + features = set(feature_str.split(";")) + if "all" in features: + raise PackageConfigError(f"Not allow to config {feature_type} all.") + return features + + +def config_feature_to_string(features: set[str]) -> str: + """配置feature集合转换为字符串。""" + if not features: + return "all" + return ";".join(sorted(features)) diff --git a/csrc/ascend/scripts/package/common/py/version_info.py b/csrc/ascend/scripts/package/common/py/version_info.py new file mode 100644 index 000000000..5e5693ddc --- /dev/null +++ b/csrc/ascend/scripts/package/common/py/version_info.py @@ -0,0 +1,435 @@ +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import os +import xml.etree.ElementTree as ET +from functools import total_ordering +from pathlib import Path +from typing import NamedTuple + +import regex as re + + +class VersionInfoError(Exception): + """版本信息异常基类。""" + + +class VersionFormatNotMatch(VersionInfoError): + """版本格式未匹配。""" + + +class IntervalFormatNotMatch(VersionInfoError): + """区间格式未匹配。""" + + +class DuplicatedPkgConfig(VersionInfoError): + """解析版本配置失败。重复的包配置。""" + + def __init__(self, pkg_name): + super().__init__(pkg_name) + self.pkg_name = pkg_name + + +class ParseVersionFailed(VersionInfoError): + """解析版本失败。""" + + +class CollectRequiresFailed(VersionInfoError): + """收集包需求失败。""" + + def __init__(self, pkg_name, version_str, msg): + super().__init__(pkg_name, version_str, msg) + self.pkg_name = pkg_name + self.version_str = version_str + self.msg = msg + + +@total_ordering +class Version: + """版本号。""" + + def __init__(self, version): + self.version = version + + @classmethod + def match(cls, input_str): + """输入字符串是否匹配版本号模式。""" + m = re.match(r"[.a-zA-Z0-9]+$", input_str) + return bool(m) + + @classmethod + def parse(cls, input_str): + """解析版本号。""" + if not cls.match(input_str): + raise VersionFormatNotMatch() + + return cls(input_str) + + @classmethod + def try_convert_to_int_list(cls, str_list): + """尝试转换为int数组。""" + for idx, item in enumerate(str_list): + try: + int_item = int(item) + str_list[idx] = int_item + except ValueError: + pass + + def to_required_list(self): + """转换为版本需求字符串列表。""" + return [self.version] + + def __eq__(self, other): + """等于。""" + if not isinstance(other, self.__class__): + return False + return self.version == other.version + + def __lt__(self, other): + """小于。""" + if not isinstance(other, self.__class__): + return True + + self_list = self.version.split(".") + other_list = other.version.split(".") + + self.try_convert_to_int_list(self_list) + self.try_convert_to_int_list(other_list) + + self_tuple = tuple(self_list) + other_tuple = tuple(other_list) + + return self_tuple < other_tuple + + def __str__(self): + return self.version + + def __repr__(self): + return repr(self.version) + + +class Point(NamedTuple): + """区间端点。""" + + type_: int # 类型，0为闭区间，1为开区间 + value: Version + + +class Interval(NamedTuple): + """版本号区间。""" + + low: Point + high: Point + + @classmethod + def match(cls, input_str: str) -> bool: + """输入字符串是否匹配区间模式。""" + if not input_str.startswith("(") and not input_str.startswith("["): + return False + if not input_str.endswith(")") and not input_str.endswith("]"): + return False + input_str = input_str[1:-1] + return input_str.count(",") <= 1 + + @classmethod + def parse(cls, input_str): + """解析版本号区间。""" + if not cls.match(input_str): + raise IntervalFormatNotMatch() + + if input_str[0] == "[": + low_type = 0 + elif input_str[0] == "(": + low_type = 1 + else: + raise AssertionError("should not go here.") + + if input_str[-1] == "]": + high_type = 0 + elif input_str[-1] == ")": + high_type = 1 + else: + raise AssertionError("should not go here.") + + input_str = input_str[1:-1] + input_list = input_str.split(",") + low = input_list[0].strip() + if len(input_list) > 1: + high = input_list[1].strip() + else: + high = None + + if low: + low_version = Point(low_type, Version(low)) + else: + low_version = None + + if high: + high_version = Point(high_type, Version(high)) + else: + high_version = None + + return cls(low=low_version, high=high_version) + + def to_required_list(self): + """转换为版本需求字符串列表。""" + result = [] + + if self.low: + if self.low.type_ == 0: + operator = ">=" + else: + operator = ">" + required_str = f"{operator}{self.low.value.version}" + result.append(required_str) + + if self.high: + if self.high.type_ == 0: + operator = "<=" + else: + operator = "<" + required_str = f"{operator}{self.high.value.version}" + result.append(required_str) + + return result + + +class Require(NamedTuple): + """包需求。""" + + pkg_name: str + versions: list + + @classmethod + def _sort_key(cls, item) -> tuple: + """排序键。""" + if isinstance(item, Interval): + # 如果存在区间左值，则左值参与排序。 + if item.low: + return item.low.value, item.low.type_ + # 否则使用区间右值，由于开区间更小，所以type_取负。 + return item.high.value, -item.high.type_ + + return item, 0 + + @classmethod + def _sort_versions(cls, versions: list) -> bool: + """排序版本序列。""" + versions.sort(key=cls._sort_key) + return True + + @classmethod + def _to_required_list(cls, versions: list) -> list[str]: + """转换为版本需求字符串列表。""" + result = [] + for version in versions: + requires = version.to_required_list() + result.extend(requires) + + return result + + @classmethod + def _to_required_str(cls, versions: list) -> str: + """转换为版本需求字符串。""" + requires = cls._to_required_list(versions) + required_str = ", ".join(requires) + + return required_str + + def sort_versions(self) -> bool: + """排序版本序列。""" + return self._sort_versions(self.versions) + + def to_required_full_str(self) -> str: + """转换为版本需求字符串。""" + required_str = self._to_required_str(self.versions) + required_full_str = f'required_package_{self.pkg_name}_version="{required_str}"' + return required_full_str + + +class ItemElement(NamedTuple): + """item元素。""" + + name: str + version: str + + @classmethod + def parse(cls, item_ele: ET.Element, cur_ver: str): + """解析item元素。""" + name = item_ele.attrib["name"] + version = item_ele.attrib["version"].replace("$(CUR_VER)", cur_ver) + return cls(name=name, version=version) + + @classmethod + def skip(cls, item_ele: ET.Element): + """是否跳过item元素。""" + version = item_ele.attrib["version"] + return version.strip() == "" + + +class CompatibleElement(NamedTuple): + """compatible元素。""" + + items: list + + @classmethod + def parse(cls, compatible_ele: ET.Element, cur_ver: str): + """解析compatible元素。""" + items = [] + for item_ele in compatible_ele.findall("./item"): + if ItemElement.skip(item_ele): + continue + item = ItemElement.parse(item_ele, cur_ver) + items.append(item) + return cls(items=items) + + +def is_version_number(version: str) -> bool: + """字符串是否为版本号。""" + has_slash = "/" in version + return not has_slash and len(version.split(".")) >= 3 + + +class VersionXml(NamedTuple): + """版本配置。""" + + release_version: str + version_dir: str + packages: dict + + @classmethod + def match(cls, filepath: Path | str) -> bool: + """文件路径是否匹配版本信息文件。""" + return str(filepath).endswith(".xml") + + @classmethod + def parse_version(cls, version_str: str): + """解析版本配置。""" + ret = Interval.match(version_str) + if ret: + result = Interval.parse(version_str) + return result + + ret = Version.match(version_str) + if ret: + result = Version.parse(version_str) + return result + + raise ParseVersionFailed() + + def get_release_version(self): + """获取发布版本号。""" + return self.release_version + + def get_version_dir(self): + """获取多版本目录。""" + return self.version_dir + + def collect_requires(self, package: str) -> list[Require]: + """收集对应包的包需求列表。""" + requires = {} + + if package not in self.packages: + return [] + + compatible = self.packages[package] + + for item in compatible.items: + pkg_name = item.name + if pkg_name not in requires: + requires[pkg_name] = Require(pkg_name=pkg_name, versions=[]) + + version_str = item.version + try: + version = self.parse_version(version_str) + except ParseVersionFailed as ex: + msg = f"parse pkg {pkg_name} version {version_str} failed" + raise CollectRequiresFailed(pkg_name, version_str, msg) from ex + + requires[pkg_name].versions.append(version) + + result = [] + for pkg_name in sorted(requires.keys()): + requires[pkg_name].sort_versions() + result.append(requires[pkg_name]) + + return result + + +def get_version_dir(version_xml: VersionXml | None, disable_multi_version: bool, version_dir: str | None) -> str | None: + """获取版本目录名。""" + if disable_multi_version: + return None + + if version_dir: + return version_dir + + # 支持从version.xml中获取version_dir + if version_xml and version_xml.get_version_dir(): + return version_xml.get_version_dir() + + return None + + +def is_multi_version(version_dir: str) -> bool: + """是否多版本。""" + return bool(version_dir) + + +class VersionInfo(NamedTuple): + """版本信息。""" + + install_version_info: bool + install_version_info_attrib: dict[str, str] | None + itf_versions: list[str] + version: str + version_xml: VersionXml | None + timestamp: str | None + + +class VersionInfoFile(NamedTuple): + """生成的版本配置。""" + + version: str + itf_version_info: str | None = None + requires: list[Require] | None = None + version_dir: str | None = None + timestamp: str | None = None + + def _get_content(self) -> str: + """获取版本配置内容。""" + lines = [f"Version={self.version}"] + if self.version_dir: + lines.append(f"version_dir={self.version_dir}") + if self.timestamp: + lines.append(f"timestamp={self.timestamp}") + if self.itf_version_info: + lines.append(self.itf_version_info) + + if self.requires: + requires_str = [require.to_required_full_str() for require in self.requires] + lines.extend(requires_str) + + lines.append("") + + return "\n".join(lines) + + def save(self, target_path: Path | str): + """保存版本配置。""" + content = self._get_content() + + target_dir = os.path.dirname(target_path) + if not os.path.exists(target_dir): + os.makedirs(target_dir) + + with open(target_path, "w") as file: + file.write(content) diff --git a/csrc/ascend/scripts/package/common/sh/check_version_required.awk b/csrc/ascend/scripts/package/common/sh/check_version_required.awk new file mode 100644 index 000000000..2cca718f9 --- /dev/null +++ b/csrc/ascend/scripts/package/common/sh/check_version_required.awk @@ -0,0 +1,281 @@ +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +function strip(input) { + sub("^ +", "", input) + sub(" +$", "", input) + return input +} + +function check_compatible_le(version_arr, len_version_arr, require_arr, len_require_arr, i) { + for (i = 1; i <= len_require_arr; i++) { + if (require_arr[i] == "") { + continue + } + # len_version_arr lt len_require_arr + if (i > len_version_arr) { + return 1 + } + if (version_arr[i] < require_arr[i]) { + return 1 + } + if (version_arr[i] > require_arr[i]) { + return 0 + } + } + if (len_version_arr > len_require_arr) { + return 1 + } + # len_version_arr eq len_require_arr + return 1 +} + +function check_compatible_lt(version_arr, len_version_arr, require_arr, len_require_arr, i) { + for (i = 1; i <= len_require_arr; i++) { + if (require_arr[i] == "") { + continue + } + # len_version_arr lt len_require_arr + if (i > len_version_arr) { + return 1 + } + if (version_arr[i] < require_arr[i]) { + return 1 + } + if (version_arr[i] > require_arr[i]) { + return 0 + } + } + if (len_version_arr > len_require_arr) { + return 0 + } + # len_version_arr eq len_require_arr + return 0 +} + +function check_compatible_ge(version_arr, len_version_arr, require_arr, len_require_arr, i) { + for (i = 1; i <= len_require_arr; i++) { + if (require_arr[i] == "") { + continue + } + # len_version_arr lt len_require_arr + if (i > len_version_arr) { + return 0 + } + if (version_arr[i] < require_arr[i]) { + return 0 + } + if (version_arr[i] > require_arr[i]) { + return 1 + } + } + if (len_version_arr > len_require_arr) { + return 1 + } + # len_version_arr eq len_require_arr + return 1 +} + +function check_compatible_gt(version_arr, len_version_arr, require_arr, len_require_arr, i) { + for (i = 1; i <= len_require_arr; i++) { + if (require_arr[i] == "") { + continue + } + # len_version_arr lt len_require_arr + if (i > len_version_arr) { + return 0 + } + if (version_arr[i] < require_arr[i]) { + return 0 + } + if (version_arr[i] > require_arr[i]) { + return 1 + } + } + if (len_version_arr > len_require_arr) { + return 1 + } + # len_version_arr eq len_require_arr + return 0 +} + +function check_compatible_eq(version_arr, len_version_arr, require_arr, len_require_arr, i) { + for (i = 1; i <= len_require_arr; i++) { + if (require_arr[i] == "") { + continue + } + # len_version_arr lt len_require_arr + if (i > len_version_arr) { + return 0 + } + if (version_arr[i] != require_arr[i]) { + return 0 + } + } + if (len_version_arr > len_require_arr) { + return 1 + } + # len_version_arr eq len_require_arr + return 1 +} + +function check_compatible(version_arr, len_version_arr, require, require_arr, len_require_arr, pos) { + len_require_arr = split(require, require_arr, ".") + + pos = match(require_arr[1], /^>=/) + if (pos != 0) { + require_arr[1] = substr(require_arr[1], pos + RLENGTH) + return check_compatible_ge(version_arr, len_version_arr, require_arr, len_require_arr) + } + + pos = match(require_arr[1], /^>/) + if (pos != 0) { + require_arr[1] = substr(require_arr[1], pos + RLENGTH) + return check_compatible_gt(version_arr, len_version_arr, require_arr, len_require_arr) + } + + pos = match(require_arr[1], /^<=/) + if (pos != 0) { + require_arr[1] = substr(require_arr[1], pos + RLENGTH) + return check_compatible_le(version_arr, len_version_arr, require_arr, len_require_arr) + } + + pos = match(require_arr[1], /^/) + if (pos != 0) { + gt_require = 1 + lt_require = 0 + eq_require = 0 + } else { + pos = match(all_required_arr[i], /^> "${LOG_FILE}" +} + +# 设置日志参数 +set_comm_log() { + local pkg_name="$1" + local log_file="$2" + + LOG_PKG_NAME="${pkg_name}" + if [ "$log_file" != "" ]; then + LOG_FILE="${log_file}" + fi +} + +# 安全日志 +comm_log_operation() { + local cur_date="$(date +'%Y-%m-%d %H:%M:%S')" + local operation="$1" + local runfilename="$2" + local result="$3" + local installmode="$4" + local all_parma="$5" + local level="" + if [ "${operation}" = "Install" ]; then + level="SUGGESTION" + elif [ "${operation}" = "Upgrade" ]; then + level="MINOR" + elif [ "${operation}" = "Uninstall" ]; then + level="MAJOR" + else + level="UNKNOWN" + fi + + if [ ! -f "${COMM_OPERATION_LOGFILE}" ]; then + touch "${COMM_OPERATION_LOGFILE}" + chmod 640 "${COMM_OPERATION_LOGFILE}" + fi + + echo "${operation} ${level} ${COMM_USERNAME} ${cur_date} 127.0.0.1 ${runfilename} ${result} installmode=${installmode}; cmdlist=${all_parma}" >> "${COMM_OPERATION_LOGFILE}" +} + +## end module + +# 转换--install-for-all参数下文件权限 +comm_set_install_for_all_mod() { + local _outvar="$1" + local _mod="$2" + local _new_mod + + # ${parameter%word} Remove matching suffix pattern. + local _new_mod="${_mod%?}" + # ${parameter#word} Remove matching prefix pattern. + local _new_mod="${_new_mod}${_new_mod#${_new_mod%?}}" + + eval "${_outvar}=\"${_new_mod}\"" +} + +# 创建文件 +comm_create_file() { + local path="$1" + local mod="$2" + local own="$3" + local install_for_all="$4" + + if [ -d "$path" ]; then + comm_log "WARNING" "remove existed dir $path before create file." + rm -rf "$path" + fi + touch "$path" + if [ $? -ne 0 ]; then + comm_log "WARNING" "create file $path failed." + return 1 + fi + + if [ "$install_for_all" = "true" ] || [ "$install_for_all" = "y" ]; then + comm_set_install_for_all_mod "mod" "$mod" + fi + chmod "$mod" "$path" + if [ $? -ne 0 ]; then + comm_log "WARNING" "chmod path $path $mod failed." + return 1 + fi + + chown "$own" "$path" + if [ $? -ne 0 ]; then + comm_log "WARNING" "chown path $path $own failed." + return 1 + fi + return 0 +} + +# 创建目录 +comm_create_dir() { + local path="$1" + local mod="$2" + local own="$3" + local install_for_all="$4" + + if [ "$path" = "" ]; then + comm_log "WARNING" "dir path is empty" + return 1 + fi + + if [ ! -d "$path" ]; then + mkdir -p "$path" + if [ $? -ne 0 ]; then + comm_log "WARNING" "create dir $path failed." + return 1 + fi + fi + + if [ "$install_for_all" = "true" ] || [ "$install_for_all" = "y" ]; then + comm_set_install_for_all_mod "mod" "$mod" + fi + chmod "$mod" "$path" + if [ $? -ne 0 ]; then + comm_log "WARNING" "chmod path $path $mod failed." + return 1 + fi + + chown -f "$own" "$path" + if [ $? -ne 0 ]; then + comm_log "WARNING" "chown path $path $own failed." + return 1 + fi +} + +# 创建子包目录 +comm_create_package_dir() { + local package_dir="$1" + local own="$2" + local install_for_all="$3" + if [ ! -d "$package_dir" ]; then + comm_create_dir "$package_dir" "755" "$own" "$install_for_all" + fi +} + +# 创建install_info文件 +comm_create_install_info_by_path() { + local install_info="$1" + local own="$2" + if [ ! -f "$install_info" ]; then + comm_create_file "$install_info" "640" "$own" "false" + fi +} + +# 更新安装参数 +comm_update_install_param() { + local key="$1" + local val="$2" + local file="$3" + local param="" + if [ ! -f "${file}" ]; then + return 1 + fi + param="$(grep -i "^${key}=" "${file}")" + if [ "${param}" = "" ]; then + echo "${key}=${val}" >> "${file}" + else + sed -i "/^${key}=/Ic ${key}=${val}" "${file}" + fi +} + +# 获取安装参数 +# _outvar : [输出变量]，安装参数值 +# _file : install.info文件路径 +# _key : 参数名 +comm_get_install_param() { + local _outvar="$1" + local _file="$2" + local _key="$3" + local _result + + if [ ! -f "${_file}" ]; then + comm_log "WARNING" "file ${_file} doesn't exist in get install param." + return 1 + fi + _result="$(grep -i "^${_key}=" "${_file}" | cut -d"=" -f2-)" + eval "${_outvar}=\"${_result}\"" +} + +# 解析安装参数选项 +__parse_install_options() { + local -n install_path_param_ref="$1" + local -n install_type_ref="$2" + local -n username_ref="$3" + local -n usergroup_ref="$4" + local -n feature_type_ref="$5" + + while true + do + case "$1" in + --install-path-param=*) + install_path_param_ref="$(echo "$1" | cut -d"=" -f2-)" + shift + ;; + --install-type=*) + install_type_ref="$(echo "$1" | cut -d"=" -f2-)" + shift + ;; + --username=*) + username_ref="$(echo "$1" | cut -d"=" -f2-)" + shift + ;; + --usergroup=*) + usergroup_ref="$(echo "$1" | cut -d"=" -f2-)" + shift + ;; + --feature-type=*) + feature_type_ref="$(echo "$1" | cut -d"=" -f2-)" + shift + ;; + *) + break + ;; + esac + done +} + +# 更新安装参数文件 +comm_update_install_info() { + local install_path_param="" + local install_type="" + local username="" + local usergroup="" + local feature_type="" + local package install_info + + __parse_install_options "install_path_param" "install_type" "username" "usergroup" "feature_type" + + if [ $# -lt 2 ]; then + return 1 + fi + + package="$1" + install_info="$2" + + if [ "$install_path_param" != "" ]; then + comm_update_install_param "${package}_Install_Path_Param" "$install_path_param" "$install_info" + fi + if [ "$install_type" != "" ]; then + comm_update_install_param "${package}_Install_Type" "$install_type" "$install_info" + fi + if [ "$username" != "" ]; then + comm_update_install_param "${package}_UserName" "$username" "$install_info" + fi + if [ "$usergroup" != "" ]; then + comm_update_install_param "${package}_UserGroup" "$usergroup" "$install_info" + fi + if [ "$feature_type" != "" ]; then + comm_update_install_param "${package}_Feature_Type" "$feature_type" "$install_info" + fi +} + +# 开始安装前打印开始信息 +comm_start_log() { + local all_parma="$@" + local cur_date="$(date +'%Y-%m-%d %H:%M:%S')" + comm_log "INFO" "Start time:$cur_date" + comm_log "INFO" "LogFile:${COMM_LOGFILE}" + comm_log "INFO" "InputParams:$all_parma" + comm_log "INFO" "OperationLogFile:${COMM_OPERATION_LOGFILE}" +} + +# 安装结束退出前打印结束信息 +comm_exit_log() { + local cur_date="$(date +'%Y-%m-%d %H:%M:%S')" + comm_log "INFO" "End time:${cur_date}" + exit "$1" +} + +comm_print_usage() { + local runfilename="$1" + comm_log "INFO" "Please input this command for help: ./${runfilename} --help" +} + +# 判断安装路径是否合法 +comm_judgmentpath() { + local install_path="$1" + local pkg_name="$2" + check_install_path_valid "$install_path" + if [ $? -ne 0 ]; then + comm_log "ERROR" "The $pkg_name install_path $install_path is invalid, only characters in [a-z,A-Z,0-9,-,_] are supported!" + exit 1 + fi +} + +# 标准化安装路径 +# _outvar : [输出变量]，处理后的安装路径 +# _install_path : 安装路径 +comm_normalize_install_path() { + local _outvar="$1" + local _install_path="$2" + local _slashes_num _result + _slashes_num=$(echo "$_install_path" | grep -o '/' | wc -l) + if [ "$_slashes_num" -gt 1 ]; then + _result=$(echo "$_install_path" | sed "s/\/*$//g") + else + _result="$_install_path" + fi + eval "${_outvar}=\"${_result}\"" +} + +# 解析安装路径 +# _outvar : [输出变量]，处理后的安装路径 +# _install_path : 安装路径 +# _pkg_name : 包名 +comm_parse_install_path() { + local _outvar="$1" + local _install_path="$2" + local _pkg_name="$3" + + comm_judgmentpath "$_install_path" "$_pkg_name" + comm_normalize_install_path "$_outvar" "$_install_path" +} + +############### 错误函数 ############### +# 文件没有找到 +comm_err_file_or_directory_not_exist() { + comm_log "ERROR" "The file or directory doesn't exist, $1" + comm_exit_log 1 +} + +os_name() { + if [ ! -f "$1" ];then + HostOsName=unknown + HostOsFullName=unknown + HostOsVersion=unknown + return + fi + + HostOsName=$(cat "$1" | grep ^NAME= | awk -F "[\" ]" '{print $2}') + HostOsFullName=$(cat "$1" | grep ^NAME= | awk -F "\"" '{print $2}') + + if [ x"$HostOsName" = "x" ];then + HostOsName=$(cat "$1" | grep ^NAME= | awk -F "[=]" '{print $2}') + HostOsFullName=$(cat "$1" | grep ^NAME= | awk -F "\"" '{print $2}') + fi + if [ x"$HostOsName" = "x" ];then + HostOsName=unknown + HostOsFullName=unknown + fi + HostOsVersion=$(cat "$1" | grep ^VERSION_ID= | awk -F "\"" '{print $2}') + if [ x"$HostOsVersion" = "x" ];then + HostOsVersion=unknown + fi + return +} + +get_os_info() { + if [ -f /etc/os-release ];then + os_name /etc/os-release + elif [ -f /etc/centos-release ];then + HostOsName=CentOS + HostOsFullName="CentOS Linux" + HostOsVersion=$(cat /etc/centos-release | awk '{print $4}') + else + which lsb_release >/dev/null 2>&1 + if [ $? -eq 0 ];then + HostOsName=$(lsb_release -si) + HostOsFullName="${HostOsName}" + HostOsVersion=$(lsb_release -sr) + else + os_name /etc/issue + fi + fi + return +} + +get_system_info() { + get_os_info + HostArch=$(uname -m) + KernelVersion=$(uname -r) +} + +version_gt() { + if [ "$2"x = "x" ];then + return 0 + else + test "$(echo "$@" | tr " " "\n" | sort -V | head -n 1)" != "$1" + fi +} +version_le() { + if [ "$2"x = "x" ];then + return 0 + else + test "$(echo "$@" | tr " " "\n" | sort -V | head -n 1)" = "$1" + fi +} +version_lt() { + if [ "$2"x = "x" ];then + return 0 + else + test "$(echo "$@" | tr " " "\n" | sort -rV | head -n 1)" != "$1" + fi +} +version_ge() { + if [ "$2"x = "x" ];then + return 0 + else + test "$(echo "$@" | tr " " "\n" | sort -rV | head -n 1)" = "$1" + fi +} + +version_vaild() { + local ver_range="$1" #version range + local ver="$2" #version provider + local clean_ver_range=$(echo "$ver_range" | sed 's/"//g') # clean ver_range,去除引号啥的 + local new_ver_range=$(echo "$clean_ver_range" | sed 's/\[//' | sed 's/\]//' | sed 's/(//' | sed 's/)//') #处理过的version range,去除[],() 得到类似 1.0,2.0 + + local start=$(echo $new_ver_range | awk -F ',' '{print $1}') + local end=$(echo $new_ver_range | awk -F ',' '{print $2}') + + if echo $clean_ver_range | grep -Eq "^\[.*\]$" ; then + # 闭合区间 [a,b]={x|a<=x<=b} + if version_ge $ver $start && version_le $ver $end ; then + # pass + return 0 + fi + elif echo $clean_ver_range | grep -Eq "^\[.*\)$"; then + # 左闭右开 [a,b)={x|a<=x= require_ver 结果pass 否则 nopass。即依赖包的当前版本大于等于需求 + if version_ge $ver $clean_ver_range ;then + return 0 + fi + fi + #not pass + return 1 +} + +# 检查包接口版本。 +# 输出VerCheckStatus和ver_check_status变量。 +check_pkg_ver_deps() { + ver_path="$1" + req_pkg="$2" + req_ver_path="$3" + + ver_info_list=$(awk -F '[_=]' '$1=="required" && $2=="'$req_pkg'" {print $1":"$2":"$3":"$4":"$5}' "$ver_path") + + for line in ${ver_info_list} + do + itf=`echo $line | cut -d":" -f 3` + ver=`echo $line | cut -d":" -f 5` + req_ver=`awk -F '=' '$1=="'$itf'_version" {print $2}' "$req_ver_path"` + if ! version_vaild $ver $req_ver; then + VerCheckStatus=FAIL + ver_check_status=FAIL + return + fi + done + VerCheckStatus=SUCC + ver_check_status=SUCC +} + +# 检查安装路径是否合法。 +# install_path : 安装路径 +check_install_path_valid() { + local install_path="$1" + # 黑名单设置，不允许//，...这样的路径 + if echo "${install_path}" | grep -Eq '\/{2,}|\.{3,}'; then + return 1 + fi + # 白名单设置，只允许常见字符 + if echo "${install_path}" | grep -Eq '^\~?[a-zA-Z0-9./_-]*$'; then + return 0 + else + return 1 + fi +} + +# 创建相对软链接 +# mdc的release包场景，ln命令不支持-r参数，需要手动实现相对软链功能 +# 入参需要是规范（绝对）路径，完整路径。例如：目标路径不能是一个目录，否则相对软链计算结果不正确 +# src_path : 源路径 +# dst_path : 目标路径 +create_softlink_icp_relative() { + local src_path="$1" + local dst_path="$2" + local source="top${src_path}" + local target="top${dst_path}" + + # 若变量内容从尾向前的数据符合，则将符合的最短数据删除 + local common="${target%/*}" + # 若变量内容从头开始的数据符合，则将符合的最短数据删除 + local forward="${source#"$common"/}" + + local result="" + + while [ "${forward}" = "${source}" ]; do + common="$(dirname "$common")" + forward="${source#"$common"/}" + result="../${result}" + done + + result="${result}${forward}" + + ln -sfn "${result}" "${dst_path}" + if [ $? -ne 0 ]; then + log "ERROR" "create softlink relative from ${src_path} to ${dst_path} failed!" + return 1 + fi + return 0 +} + +# 创建软链接 +create_softlink() { + local _src_dir="$1" + local _dst_dir="$2" + local _sub_dir_src="$3" + local _sub_dir_dst="$4" + local _sub_dir_dst_new _dst_path ret total_ret=0 + + # 注意_sub_dir_src可能有通配符 + for _src_path in "${_src_dir}"/${_sub_dir_src} + do + if [ -z "$_sub_dir_dst" ]; then + _sub_dir_dst_new="$(basename "${_src_path}")" + else + _sub_dir_dst_new="${_sub_dir_dst}" + fi + _dst_path="${_dst_dir}/${_sub_dir_dst_new}" + if [ -L "${_dst_path}" ]; then + rm -f "${_dst_path}" + fi + # 目标取规范化路径 + _dst_path="$(readlink -f "${_dst_path}")" + if [ -d "${_dst_path}" ]; then + # 如果目标为目录，拼接上源文件名 + _dst_path="${_dst_path}/$(basename "${_src_path}")" + fi + + # 源取规范化路径 + _src_path="$(readlink -f "${_src_path}")" + + create_softlink_icp_relative "${_src_path}" "${_dst_path}" + ret="$?" && [ $ret -ne 0 ] && total_ret=$ret + done + + return $total_ret +} + +# 创建软链接，调用ln -r命令 +create_softlink_by_relative_ln() { + local _src_dir="$1" + local _dst_dir="$2" + local _sub_dir_src="$3" + local _sub_dir_dst="$4" + local _src_path="$_src_dir/$_sub_dir_src" + # 注意_src_path可能有通配符，处理流程中需要考虑通配场景。 + local _dst_path + if [ -z "$_sub_dir_dst" ]; then + _sub_dir_dst=$(basename $_src_path) + fi + _dst_path="$_dst_dir/$_sub_dir_dst" + if [ -L "${_dst_path}" ]; then + rm -f "${_dst_path}" + fi + ln -sr $_src_path $_dst_path +} + +# 如果源文件存在，则创建软链接 +create_softlink_if_exists() { + local _src_dir="$1" + local _sub_dir_src="$3" + local _src_path="$_src_dir/$_sub_dir_src" + + if [ -e $_src_path ]; then + create_softlink "$@" + fi +} + +# 目录是否为空 +# 空目录返回0 +# 目录不存在返回1 +# 目录非空或无权限访问返回2 +is_dir_empty() { + # 如果目录不存在，则返回 + if [ ! -d "$1" ]; then + return 1 + fi + # 否则检查目录是否为空 + # 2>&1重定向无权限访问报错信息 + if [ "$(ls -A "$1" 2>&1)" != "" ]; then + return 2 + fi + return 0 +} + +# 删除空目录 +remove_dir_if_empty() { + local dirpath="$1" + + if is_dir_empty "${dirpath}"; then + rm -rf "${dirpath}" + fi + return 0 +} + +# 删除目录 +remove_dir() { + local dirpath="$1" + + if [ -e "${dirpath}" ]; then + chmod u+w -R "${dirpath}" + fi + rm -rf "${dirpath}" + + return 0 +} + +# 删除软链接 +remove_softlink() { + local path="$1" + [ -L "${path}" ] && rm -f "${path}" +} + +# 处理pre_check +# pkg_name : 包名 +# pre_check_func : pre_check函数 +# standalone : 是否为独立的pre_check命令[y/n] +process_pre_check() { + local pkg_name="$1" + local pre_check_func="$2" + local standalone="$3" + + log "INFO" "${pkg_name} do pre_check started." + ${pre_check_func} + ret=$? + if [ $ret -ne 0 ]; then + log "WARNING" "${pkg_name} do pre check failed." + fi + log "INFO" "${pkg_name} do pre_check finished." + + if [ "${standalone}" = "y" ]; then + if [ $ret -ne 0 ]; then + exitInstallLog 1 + fi + exitInstallLog 0 + fi +} + +# 包是否在环境上安装 +# install_path : 安装路径 +# pkgname : 包名 +# varname : [输出变量]，包是否在环境上安装 +does_pkg_installed() { + local install_path="$1" + local pkgname="$2" + local varname="$3" + local result="false" + local pkg_path="${install_path}/share/info/${pkgname}" + + if [ -d "${pkg_path}" ]; then + result="true" + fi + + read -r "$varname" < --script-dir= --package= --logfile= [ --logstyle= ] [ --docker-root= ] +# +# --install-path= : 安装路径 +# --script-dir= : 当前脚本目录路径 +# --package= : 当前包名 +# --logfile= : ascend_install.log日志文件路径 +# --logstyle= : 日志风格 +# --docker-root= : docker根路径 +preinstall_process() { + local install_path="" + local docker_root="" + local script_dir="" + local package="" + local logfile="" + local real_install_path + local version_info_path + local ret + + __parse_preinstall_options "install_path" "docker_root" "script_dir" "package" "logfile" "LOG_STYLE" + + __validate_preinstall_params "$install_path" "$script_dir" "$package" "$logfile" || return 1 + + LOG_FILE="${logfile}" + _package_to_log_pkg_name "LOG_PKG_NAME" "${package}" + + if [ "${docker_root}" = "" ]; then + real_install_path="${install_path}" + else + real_install_path="${docker_root}/${install_path}" + fi + + _check_version_compatiable "$package" "${real_install_path}" "${script_dir}" + if [ $? -ne 0 ]; then + return 1 + fi + + return 0 +} + +# 安装前检查。 +# 返回码： +# 0: 成功 +# 1: 版本兼容性检查失败 +# 调用方式： +# preinstall_check --install-path= --script-dir= --package= --logfile= [ --logstyle= ] [ --docker-root= ] +# +# --install-path= : 安装路径 +# --script-dir= : 当前脚本目录路径 +# --package= : 当前包名 +# --logfile= : ascend_install.log日志文件路径 +# --logstyle= : 日志风格 +# --docker-root= : docker根路径 +preinstall_check() { + preinstall_process "$@" +} + +# 是否为多版本包 +# 读取version.info文件，检查本包是否为多版本包 +# _outvar : [输出变量]，是否为多版本包 +# _filepath : version.info文件路径 +is_multi_version_pkg() { + local _outvar="$1" + local _filepath="$2" + local _ret="false" + + if [ -f "${_filepath}" ]; then + grep "^version_dir=" "${_filepath}" > /dev/null + [ $? -eq 0 ] && _ret="true" + fi + eval "${_outvar}=\"${_ret}\"" +} + +# 文件列表的某一列 +_filelist_column() { + local _filepath="$1" + local _idx="$2" + tail -n +2 "$_filepath" | cut -d, -f${_idx} | sort | uniq +} + +# 移除comm特性 +_remove_comm_feature() { + local _outvar="$1" + local _feature="$2" + local _result + + _result="$(echo "$_feature" | sed 's/comm//g' | sed 's/,,/,/g' | sed 's/^,//g' | sed 's/,$//g')" + eval "${_outvar}=\"${_result}\"" +} + +# 特性列表转换为正则表达式 +_feature_to_regex() { + local _outvar="$1" + local _feature="$2" + local _result_ftr + _result_ftr="\b(($(echo "$_feature" | sed 's/,/)|(/g')))\b" + eval "${_outvar}=\"${_result_ftr}\"" +} + +# 是否包含指定feature +# 读取filelist.csv文件，检查是否包含指定feature +# _outvar : [输出变量](true/false)，是否包含指定feature +# _feature : 指定feature，可指定多个，以“,”分隔，如：dvpp,audio +# _filepath : filelist.csv文件路径 +contain_feature() { + local _outvar="$1" + local _feature="$2" + local _filepath="$3" + local _feature_regex _feature_cf + + # 输入的comm不作为特性，也就是说不支持--feature=comm + _remove_comm_feature "_feature_cf" "$_feature" + if [ "$_feature_cf" = "" ]; then + eval "${_outvar}=\"false\"" + return 0 + fi + + _feature_to_regex "_feature_regex" "$_feature_cf" + _filelist_column "$_filepath" 10 | grep -E "$_feature_regex" > /dev/null 2>&1 + if [ $? -eq 0 ]; then + eval "${_outvar}=\"true\"" + else + eval "${_outvar}=\"false\"" + fi +} + +# 向filelist.csv中写入1个条目 +# filepath : filelist.csv文件路径 +# operation : 操作类型 +# path_in_pkg : 包内路径 +# install_path : 安装路径 +# permission : 文件权限 +# owner_group : 文件属主，支持默认值DEFAULT +# install_type : 安装类型 +# softlink : 软链接 +# block : 所属块 +# feature : 所属特性，支持默认值DEFAULT +# chip : 所属芯片，支持默认值DEFAULT +# arch : 包架构 +add_fileitem() { + local filepath="$1" + local operation="$2" + local path_in_pkg="$3" + local install_path="$4" + local permission="$5" + local owner_group="$6" + local install_type="$7" + local softlink="$8" + local pkg_inner_softlink="$9" + local block="${10}" + local feature="${11}" + local chip="${12}" + local arch="${13}" + local is_common_path="N" + + if [ "$owner_group" = "DEFAULT" ]; then + owner_group="\\\\\$username:\\\\\$usergroup" + fi + + if [ "$feature" = "DEFAULT" ]; then + feature="all" + fi + + if [ "$chip" = "DEFAULT" ]; then + chip="all" + fi + + if echo "$install_path" | grep "^${arch}-linux/" > /dev/null; then + is_common_path="Y" + fi + + echo "NA,${operation},${path_in_pkg},${install_path},FALSE,${permission},${owner_group},${install_type},${softlink},${feature},${is_common_path},FALSE,NA,${block},${pkg_inner_softlink},${chip}" >> ${filepath} +} + +# 设置公共变量 +set_global_vars() { + local arch + local scene_filepath="${curpath}/../scene.info" + + if [ -f "${scene_filepath}" ]; then + get_scene_arch "arch" "${scene_filepath}" + if [ "${arch}" != "" ]; then + PKG_ARCH="${arch}" + fi + fi +} + +# 获取并发进程数 +get_thread_num() { + local _outvar="$1" + local _thread_num="$(cat /proc/cpuinfo | grep "^processor" | wc -l)" + + # 未取得CPU核数时，默认为1 + if [ "$_thread_num" = "" ] || [ "$_thread_num" = "0" ]; then + _thread_num="1" + fi + + # CPU核数2倍 + eval "${_outvar}=\"$((_thread_num*2))\"" +} + +# 初始化fifo，用于并发控制 +init_fifo() { + local _outvar="$1" + local _thread_num="$2" + local _tmppath_if tmp + local _skip_msg="skip init fifo" + + # 并发数为空时，默认为1 + if [ "$_thread_num" = "" ]; then + _thread_num="1" + comm_log "WARNING" "thread number is empty, use ${_thread_num}." + fi + # 并发数<=0时，默认为1 + if [ $_thread_num -le 0 ]; then + comm_log "WARNING" "thread number is ${_thread_num}, use 1." + _thread_num="1" + fi + + get_tmp_file _tmppath_if "fifo" + if [ -e "$_tmppath_if" ]; then + rm -rf "$_tmppath_if" + check_ret_warning "$?" "remove old file $_tmppath_if failed, ${_skip_msg}." + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + mkfifo "$_tmppath_if" + check_ret_warning "$?" "make fifo $_tmppath_if failed, ${_skip_msg}." + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 使用文件描述符8作并发控制 + exec 8<> "$_tmppath_if" # 创建文件描述符8 + check_ret_warning "$?" "open fifo $_tmppath_if file descriptor failed, ${_skip_msg}." + ret="$?" && [ $ret -ne 0 ] && return $ret + + rm -f "$_tmppath_if" + check_ret_warning "$?" "remove fifo $_tmppath_if failed, ${_skip_msg}." + ret="$?" && [ $ret -ne 0 ] && return $ret + + for tmp in $(seq $_thread_num); do + echo >&8 + check_ret_warning "$?" "echo fifo $_tmppath_if failed, ${_skip_msg}." + ret="$?" && [ $ret -ne 0 ] && return $ret + done + + eval "${_outvar}=\"${_tmppath_if}\"" +} + +# 根据参数执行 +exec_with_param() { + local param="$1" + shift 1 + local exec_mode fifo_path="$PARALLEL_FIFO" tmp + extract_1st "exec_mode" "$param" + if [ "$exec_mode" = "concurrency" ]; then + if [ "$fifo_path" = "none" ] || [ "$fifo_path" = "" ]; then + "$@" & + else + read tmp <&8 + { + "$@" + ret="$?" + echo >&8 + exit $ret + } & + fi + else + "$@" + fi +} diff --git a/csrc/ascend/scripts/package/common/sh/common_func_v2.inc b/csrc/ascend/scripts/package/common/sh/common_func_v2.inc new file mode 100644 index 000000000..f0a4a1078 --- /dev/null +++ b/csrc/ascend/scripts/package/common/sh/common_func_v2.inc @@ -0,0 +1,1729 @@ +#!/bin/sh +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +# 公共函数库 +# 总setenv文件权限 +SETENV_MOD="550" +# 总setenv文件可写状态权限 +SETENV_WRITEABLE_MOD="600" +# 默认特性配置 +DEFAULT_FEATURE_PARAM="all n all" +#export PS4='+ ${FUNCNAME[0]:+${FUNCNAME[0]}():} ${BASH_SOURCE}:${LINENO}: ' +#set -x +RESET_MOD="750" + +# db.info文件格式如下 +# --- +# CommonLib|atc,fwkacllib +# Compiler|atc,fwkacllib +# --- +# 使用"|"分隔模块名和包名列表，包名列表使用","分隔 +# 文件中块名保持以升序排序 +PKG_DB_INFO_RELPATH="var/ascend_package_db.info" +# 缓存mod文件名 +STASH_MOD_PATH="stash_mod.txt" +# stash_mod文件权限 +STASH_FILE_MOD="600" + +# 写日志 +log() { + local cur_date_="$(date +"%Y-%m-%d %H:%M:%S")" + local log_type_="${1}" + local msg_="${2}" + local log_format_="[Common] [${cur_date_}] [${log_type_}]: ${msg_}" + if [ "${log_type_}" = "INFO" ]; then + echo "${log_format_}" + elif [ "${log_type_}" = "WARNING" ]; then + echo "${log_format_}" + elif [ "${log_type_}" = "ERROR" ]; then + echo "${log_format_}" + elif [ "${log_type_}" = "DEBUG" ]; then + echo "${log_format_}" 1> /dev/null + fi +} + +# 返回列表长度 +__length_list() { + local list="$1" + local var="$2" + local list_item + local cnt=0 + + for list_item in ${list}; do + cnt=$((cnt+1)) + done + + eval "${var}=\"${cnt}\"" +} + +# 获取列表索引值 +__index_list() { + local list="$1" + shift + local list_item + local cnt=0 + + if [ $# -eq 0 ]; then + return 0 + fi + + for list_item in ${list}; do + if [ ${1} -eq ${cnt} ]; then + eval "${2}=\"${list_item}\"" + shift 2 + if [ $# -eq 0 ]; then + return 0 + fi + fi + cnt=$((cnt+1)) + done + + return 0 +} + +# 从列表中移除一项 +__remove_item_in_list() { + local to_removed="$1" + shift + local list="$*" + local list_item + local new_list + + for list_item in ${list}; do + if [ "${to_removed}" != "${list_item}" ]; then + if [ "${new_list}" = "" ]; then + new_list="${list_item}" + else + new_list="${new_list} ${list_item}" + fi + fi + done + echo "${new_list}" +} + +# 元素是否在列表中 +__item_in_list() { + local _outvar="$1" + local _item="$2" + shift 2 + local _list_item + local _matched="false" + + for _list_item in $*; do + if [ "${_item}" = "${_list_item}" ]; then + _matched="true" + break + fi + done + eval "${_outvar}=\"${_matched}\"" +} + +# 反转列表 +__reverse_list() { + local _outvar="$1" + local _list="$2" + local _new_list="" + local _list_item + + for _list_item in ${_list}; do + if [ "${_new_list}" = "" ]; then + _new_list="${_list_item}" + else + _new_list="${_list_item} ${_new_list}" + fi + done + eval "${_outvar}=\"${_new_list}\"" +} + +# 修改各文件及目录的属性 +change_own() { + local recursive="$3" + local option="" + local username="${USERNAME}" + local usergroup="${USERGROUP}" + if [ "$2" != "NA" ]; then + if [ "${recursive}" = "true" ]; then + option="-R" + fi + eval chown ${option} -h \"$2\" \"$1\" + if [ $? -ne 0 ]; then + log "ERROR" "$1 chown failed!" + return 1 + fi + fi +} + +# 获取install_for_all文件权限 +get_install_for_all_mod() { + local _outvar="$1" + local _mod="$2" + local _new_mod _other_mod + + _new_mod="${_mod%?}" + _other_mod="${_new_mod#${_new_mod%?}}" + _other_mod="$(($_other_mod & 5))" # other权限位移除写权限，仅支持普通用户运行 + + eval "${_outvar}=\"${_new_mod}${_other_mod}\"" +} + +# 修改各文件及目录的权限 +change_mod() { + local mod="$2" + local install_for_all="$3" + local recursive="$4" + local option="" new_mod + # 对于软连接，可能目标文件还没有拷贝进来，导致无法修改mod，这里过滤掉软连接 + if [ -L "$1" ]; then + return 0 + fi + if [ "$2" != "NA" ]; then + if [ "${recursive}" = "true" ]; then + option="-R" + fi + # 如果设置了install_for_all，则安装时other权限跟group权限对齐 + if [ "${install_for_all}" = "y" ]; then + get_install_for_all_mod new_mod "$mod" + chmod ${option} "${new_mod}" "$1" + else + chmod ${option} "$2" "$1" + fi + if [ $? -ne 0 ]; then + log "ERROR" "$1 chmod failed!" + return 1 + fi + fi + return 0 +} + +# 获取文件权限 +get_file_mod() { + local _outvar="$1" + local _options="" _ret + shift + + while true; do + case "$1" in + -L|--dereference) + _options="${_options} $1" + shift + ;; + *) + break + ;; + esac + done + + local _path="$1" + local _result + + _result="$(stat ${_options} -c %a "${_path}")" + _ret="$?" && [ $_ret -ne 0 ] && return $_ret + eval "${_outvar}=\"${_result}\"" +} + +# 检查路径是否为绝对路径 +__check_abs_path() { + local path="$1" + + if [ "${path#/}" != "${path}" ]; then + is_abs_path="true" + else + is_abs_path="false" + fi +} + +__set_abs_path() { + local install_path="$1" + local path="$2" + local varname="$3" + local is_abs_path + + __check_abs_path "${path}" + if [ "${is_abs_path}" != "true" ]; then + eval "${varname}=\"${install_path}/${path}\"" + else + eval "${varname}=\"${path}\"" + fi +} + +# 创建目录 +make_dir() { + local path="$1" + mkdir -p "${path}" + if [ $? -ne 0 ]; then + log "ERROR" "${path} mkdir failed!" + exit 1 + fi + return 0 +} + +# 检查install_path在docker_root之中 +check_install_path_in_docker_root() { + local install_path="$1" + local docker_root="$2" + + echo "${install_path}" | grep "^${docker_root}" > /dev/null 2>&1 + if [ $? -ne 0 ]; then + log "ERROR" "check install path ${install_path} in docker root ${docker_root} failed!" + return 1 + fi + return 0 +} + +# 移除路径右侧斜线(/) +rstrip_path() { + local _outvar="$1" + local _path="$2" + + _path="$(echo "${_path}" | sed "s/\/\+\$//g")" + eval "${_outvar}=\"${_path}\"" +} + +# 包是否在tools目录下 +is_package_under_tools() { + local _outvar="$1" + local _package="$2" + local _result_iput="false" + + if [ "${_package}" = "aoe" ] || [ "${_package}" = "nca" ] || [ "${_package}" = "amct_acl" ] || [ "${_package}" = "ncs" ]; then + _result_iput="true" + fi + + eval "${_outvar}=\"${_result_iput}\"" +} + +# 获取包目录名 +get_package_dir() { + local _outvar="$1" + local _package="$2" + local _is_under_tools _result + + if [ "${USE_SHARE_INFO}" == "y" ]; then + _result="share/info" + else + is_package_under_tools "_is_under_tools" "${_package}" + if [ "${_is_under_tools}" = "true" ]; then + _result="tools" + else + _result="" + fi + fi + eval "${_outvar}=\"${_result}\"" +} + +# 获取包目录路径 +get_package_dirpath() { + local _outvar="$1" + local _package="$2" + local _is_under_tools _result + + if [ "${USE_SHARE_INFO}" == "y" ]; then + _result="share/info/${_package}" + else + is_package_under_tools "_is_under_tools" "${_package}" + if [ "${_is_under_tools}" = "true" ]; then + _result="tools/${_package}" + else + _result="${_package}" + fi + fi + eval "${_outvar}=\"${_result}\"" +} + +# 获取包ascend_install.info路径 +get_package_install_info() { + local _outvar="$1" + local _install_path="$2" + local _version_dir="$3" + local _package="$4" + local _package_dirpath="" + + eval "${_outvar}=\"\"" + + get_package_dirpath "_package_dirpath" "${_package}" + + eval "${_outvar}=\"${_install_path}/${_version_dir}/${_package_dirpath}/ascend_install.info\"" +} + +# 获取包version.info路径 +get_package_version_info() { + local _outvar="$1" + local _install_path="$2" + local _version_dir="$3" + local _package="$4" + local _package_dirpath="" + + eval "${_outvar}=\"\"" + + get_package_dirpath "_package_dirpath" "${_package}" + + eval "${_outvar}=\"${_install_path}/${_version_dir}/${_package_dirpath}/version.info\"" +} + +# 获取包filelist.csv路径 +get_package_filelist() { + local _outvar="$1" + local _install_path="$2" + local _version_dir="$3" + local _package="$4" + local _package_dirpath="" + + eval "${_outvar}=\"\"" + + get_package_dirpath "_package_dirpath" "${_package}" + + eval "${_outvar}=\"${_install_path}/${_version_dir}/${_package_dirpath}/script/filelist.csv\"" +} + +# 获取包install_common_parser.sh路径 +get_package_install_common_parser() { + local _outvar="$1" + local _install_path="$2" + local _version_dir="$3" + local _package="$4" + local _package_dirpath="" + + eval "${_outvar}=\"\"" + + get_package_dirpath "_package_dirpath" "${_package}" + + eval "${_outvar}=\"${_install_path}/${_version_dir}/${_package_dirpath}/script/install_common_parser.sh\"" +} + +# 获取latest_manager的install_common_parser.sh路径 +get_latest_manager_install_common_parser() { + local _outvar="$1" + local _install_path="$2" + local _latest_dir="$3" + + eval "${_outvar}=\"${_install_path}/${_latest_dir}/var/manager/install_common_parser.sh\"" +} + +# 获取包script目录路径 +get_package_script_dirpath() { + local _outvar="$1" + local _install_path="$2" + local _version_dir="$3" + local _package="$4" + local _package_dirpath="" + + eval "${_outvar}=\"\"" + + get_package_dirpath "_package_dirpath" "${_package}" + + eval "${_outvar}=\"${_install_path}/${_version_dir}/${_package_dirpath}/script\"" +} + +# 检查参数不为空 +check_param_not_empty() { + local name="$1" + local error_msg="$2" + local value + + eval "value=\"\${${name}}\"" + + if [ "${value}" = "" ]; then + comm_log "ERROR" "$2" + return 1 + fi + + return 0 +} + +# 检查文件存在 +check_file_exists() { + local path="$1" + local error_msg="$2" + + if [ ! -f "${path}" ]; then + comm_log "ERROR" "$2" + return 1 + fi + + return 0 +} + +# 检查返回值是否为0 +check_ret_error() { + local ret="$1" + local msg="$2" + + if [ ${ret} -ne 0 ]; then + comm_log "ERROR" "${msg}" + return ${ret} + fi + + return 0 +} + +# 检查返回值是否为0 +check_ret_warning() { + local ret="$1" + local msg="$2" + + if [ ${ret} -ne 0 ]; then + comm_log "WARNING" "${msg}" + return ${ret} + fi + + return 0 +} + +# 获取真实路径 +get_realpath() { + local _outvar="$1" + local _path_gr="$2" + + _path_gr="$(readlink -f "${_path_gr}")" + eval "${_outvar}=\"${_path_gr}\"" +} + +# 检查返回值是否为0 +cleanup_if_error() { + local ret="$1" + local cleanup="$2" + + if [ ${ret} -ne 0 ]; then + eval "${cleanup}" + return ${ret} + fi + + return 0 +} + +# 获取临时目录 +get_tmp_root() { + local _outvar="$1" + local _result_gtr + + if [ -d "${HOME}" ]; then + _result_gtr="${HOME}" + elif [ $(id -u) -eq 0 ] && [ -d "/root" ]; then + _result_gtr="/root" + else + _result_gtr="${PWD}" + fi + + eval "${_outvar}=\"${_result_gtr}\"" + return 0 +} + +# 获取临时文件 +get_tmp_file() { + local _outvar="$1" + local _filename="$2" + local _tmp_file_gtf _result + + get_tmp_root "_tmp_file_gtf" + + _result=$(mktemp "$_tmp_file_gtf/${_filename}_XXXXXX") + check_ret_warning "$?" "mktemp $_tmp_file_gtf/${_filename}_XXXXXX failed." + ret="$?" && [ $ret -ne 0 ] && return $ret + + eval "${_outvar}=\"${_result}\"" + return 0 +} + +# 获取包架构 +get_scene_arch() { + local _outvar="$1" + local _scene_filepath="$2" + local _result + + _result="$(grep "^arch=" "${_scene_filepath}" | cut -d= -f2-)" + eval "${_outvar}=\"${_result}\"" +} + +# 打包feature参数 +pack_feature_param() { + local _outvar="$1" + local _feature_type="$2" + local _feature_exclude_all="$3" + local _chip="$4" + + eval "${_outvar}=\"${_feature_type} ${_feature_exclude_all} ${_chip}\"" +} + +# 解包feature参数 +# 注意，调用解包时参数不可加引号 +unpack_feature_param() { + local _feature_type_var="$1" + local _feature_exclude_all_var="$2" + local _chip_var="$3" + shift 3 + eval "${_feature_type_var}=\"${1}\"" + eval "${_feature_exclude_all_var}=\"${2}\"" + eval "${_chip_var}=\"${3}\"" +} + +# 展开version.info文件中参数 +expand_version_file() { + if [ "${VERSION_FILE}" = "" ]; then + return 0 + fi + get_version "VERSION" "${VERSION_FILE}" + get_version_dir "VERSION_DIR" "${VERSION_FILE}" +} + +# 提取第一项 +extract_1st() { + local _outvar="$1" + eval "${_outvar}=\"$2\"" +} + +# 提取第二项 +extract_2nd() { + local _outvar="$1" + eval "${_outvar}=\"$3\"" +} + +# 路径转为sed正则表达式 +path_to_regex() { + local _outvar="$1" + local _path="$2" + local _reslut_ptr="$(echo "${_path}" | sed "s#\/#\\\/#g")" + + eval "${_outvar}=\"${_reslut_ptr}\"" +} + +# 设置默认值 +set_default() { + local _outvar="$1" + local _input="$2" + local _default="$3" + + if [ "${_input}" = "" ]; then + eval "${_outvar}=\"${_default}\"" + else + eval "${_outvar}=\"${_input}\"" + fi +} + +# 标准化feature参数 +# 输入的feature参数中，如果有all字段（表示安装所有feature） +# 则将feature参数重置为all，在后续流程中安装所有feature +# 卸载流程中，feature与chip强制为all，保证卸载掉block中的所有文件 +normalize_feature() { + local _outvar="$1" + local _feature_nf="$2" + local _operation="$3" + + if [ "$_operation" = "uninstall" ]; then + eval "${_outvar}=\"all\"" + return 0 + fi + + if [ "$_feature_nf" = "" ]; then + eval "${_outvar}=\"all\"" + return 0 + fi + + case "${_feature_nf}" in + all,*) + _feature_nf="all" + ;; + *,all) + _feature_nf="all" + ;; + *,all,*) + _feature_nf="all" + ;; + esac + + eval "${_outvar}=\"${_feature_nf}\"" +} + +# 删除软连接 +remove_softlink_icp() { + local softlink="$1" + + if [ "${softlink}" = "" ]; then + return 0 + fi + + if [ "${softlink}" != "NA" ] && [ -L "${softlink}" ]; then + rm -f "${softlink}" + if [ $? -ne 0 ]; then + log "ERROR" "remove ${softlink} failed!" + return 1 + fi + fi + return 0 +} + +#移除文件 +remove_file() { + local target="$1" + local softlink="$2" + local ret + if [ -e "${target}" ] || [ -L "${target}" ]; then + rm -f "${target}" + if [ $? -ne 0 ]; then + log "ERROR" "remove ${target} failed!" + return 1 + fi + fi + remove_softlink_icp "${softlink}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# 创建文件夹 +create_folder() { + local install_path="$1" + local target="$2" + local softlinks_str="$3" + local ret target_abs + + __set_abs_path "${install_path}" "${target}" "target_abs" + + if [ ! -d "${target_abs}" ]; then + make_dir "${target_abs}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + change_mod "${target_abs}" "${RESET_MOD}" "" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + if [ "${softlinks_str}" != "NA" ]; then + create_softlink_by_install_path "${install_path}" "${target}" "${softlinks_str}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + return 0 +} + +# 创建目录 +create_dirs() { + local install_path="$1" + local line="$2" + local target + local target_abs + local softlinks_str + local ret + + __index_list "${line}" 1 "target" 4 "softlinks_str" + + __set_abs_path "${install_path}" "${target}" "target_abs" + + if [ -L "${target_abs}" ] ; then + rm -f "${target_abs}" + log "WARNING" "${target_abs} is an existing soft-link, deleted." + fi + create_folder "${install_path}" "${target}" "${softlinks_str}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + return 0 +} + +# 修改目录的权限和属组 +reset_mod_dirs() { + local install_path="$1" + local line="$2" + local mod + local target + local target_abs + local is_abs_path ret + + __index_list "${line}" 1 "target" + + __set_abs_path "${install_path}" "${target}" "target_abs" + + # 目录不存在时跳过 + if [ ! -d "${target_abs}" ]; then + return 0 + fi + # 只处理目录，没有处理目录的软链接 + change_mod "${target_abs}" "${RESET_MOD}" "" "false" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# 递归修改目录的权限和属组 +reset_mod_dirs_recursive() { + local install_path="$1" + local line="$2" + local mod + local target + local target_abs + local is_abs_path ret + + __index_list "${line}" 1 "target" + + __set_abs_path "${install_path}" "${target}" "target_abs" + + # 目录不存在时跳过 + if [ ! -d "${target_abs}" ]; then + return 0 + fi + # 只处理目录，没有处理目录的软链接 + change_mod "${target_abs}" "${RESET_MOD}" "" "true" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +__unpack_softlinks() { + local softlinks_str="$1" + local varname="$2" + local item + local temp + OLD_IFS="${IFS}" + IFS=";" + temp="" + for item in ${softlinks_str}; do + if [ "${temp}" = "" ]; then + temp="${item}" + else + temp="${temp} ${item}" + fi + done + IFS="${OLD_IFS}" + + eval "${varname}=\"${temp}\"" +} + +#移除文件夹 +remove_dir_icp() { + if [ -e "$1" ] || [ -L "$1" ]; then + rm -fr "$1" + if [ $? -ne 0 ]; then + log "ERROR" "$1 remove failed!" + return 1 + fi + fi + return 0 +} + +# 创建软链时，移除存在的目录 +remove_exists_dir_in_create_softlink() { + local dirpath="$1" + + if [ -d "${dirpath}" ] && [ ! -L "${dirpath}" ]; then + log "WARNING" "${dirpath} is an existing directory in create softlink, deleted." + change_mod "${dirpath}" "700" "n" "true" + remove_dir_icp "${dirpath}" + fi +} + +# 创建绝对软链接 +create_softlink_icp_absolute() { + local src_path="$1" + local dst_path="$2" + + remove_exists_dir_in_create_softlink "${dst_path}" + ln -sfn "${src_path}" "${dst_path}" + if [ $? -ne 0 ]; then + log "ERROR" "create softlink absolute from ${src_path} to ${dst_path} failed!" + return 1 + fi + return 0 +} + +# 该函数与common_func.inc脚本中的相同 +# install_common_parser.sh不一定能source到common_func.inc（原因见source common_func.inc的注释） +# 所以这里需要重复定义 +create_softlink_icp_relative() { + local src_path="$1" + local dst_path="$2" + local source="top${src_path}" + local target="top${dst_path}" + + # 若变量内容从尾向前的数据符合，则将符合的最短数据删除 + local common="${target%/*}" + # 若变量内容从头开始的数据符合，则将符合的最短数据删除 + local forward="${source#"$common"/}" + + local result="" + + while [ "${forward}" = "${source}" ]; do + common="$(dirname "$common")" + forward="${source#"$common"/}" + result="../${result}" + done + + result="${result}${forward}" + + remove_exists_dir_in_create_softlink "${dst_path}" + ln -sfn "${result}" "${dst_path}" + if [ $? -ne 0 ]; then + log "ERROR" "create softlink relative from ${src_path} to ${dst_path} failed!" + return 1 + fi + return 0 +} + +# 创建软连接 +create_softlink_by_install_path() { + local install_path="$1" + local target="$2" + local softlinks_str="$3" + local softlinks softlink + local target_abs + local softlink_abs + local is_abs_path + local ret + + if [ "${softlinks_str}" = "NA" ]; then + return 0 + fi + + __unpack_softlinks "${softlinks_str}" "softlinks" + + for softlink in ${softlinks}; do + __check_abs_path "${target}" + if [ "${is_abs_path}" = "true" ]; then + __set_abs_path "${install_path}" "${softlink}" "softlink_abs" + create_softlink_icp_absolute "${target}" "${softlink_abs}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + continue + fi + + __check_abs_path "${softlink}" + if [ "${is_abs_path}" = "true" ]; then + __set_abs_path "${install_path}" "${target}" "target_abs" + create_softlink_icp_absolute "${target_abs}" "${softlink}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + continue + fi + + create_softlink_icp_relative "${install_path}/${target}" "${install_path}/${softlink}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + done + + return 0 +} + +# 创建软链接。icp后缀为避免重名 +create_softlink_icp() { + local options="" + local is_relative="false" + + while true; do + case "$1" in + -r|--relative) + is_relative="true" + shift + ;; + -*) + log "ERROR" "unsupported option $1 in create softlink icp!" + return 1 + ;; + *) + break + ;; + esac + done + + local src_path="$1" + local dst_path="$2" + + if [ "${is_relative}" = "true" ]; then + create_softlink_icp_relative "${src_path}" "${dst_path}" + else + create_softlink_icp_absolute "${src_path}" "${dst_path}" + fi +} + +# 修改权限和属组 +change_mod_and_own(){ + local target="$1" + local mod="$2" + local own="$3" + local install_for_all="$4" + local recursive="$5" + local ret + + change_mod "${target}" "${mod}" "${install_for_all}" "${recursive}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + change_own "${target}" "${own}" "${recursive}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 修改目录的权限和属组 +change_mod_and_own_dirs() { + local install_path="$1" + local line="$2" + local target + local mod + local own + local is_abs_path + + __index_list "${line}" 1 "target" 2 "mod" 3 "own" + + __check_abs_path "${target}" + if [ "${is_abs_path}" != "true" ]; then + target="${install_path}/${target}" + fi + if [ ! -d "${target}" ]; then + return 0 + fi + # 只处理目录，没有处理目录的软链接 + change_mod_and_own "${target}" "${mod}" "${own}" "${INSTALL_FOR_ALL}" "false" +} + +# 创建stash_mod.txt文件 +create_stash_mod() { + local install_path="$1" + rm -f "${install_path}/${STASH_MOD_PATH}" + touch "${install_path}/${STASH_MOD_PATH}" + chmod ${STASH_FILE_MOD} "${install_path}/${STASH_MOD_PATH}" +} + +# 删除stash_mod.txt文件 +remove_stash_mod() { + local install_path="$1" + rm -f "${install_path}/${STASH_MOD_PATH}" +} + +# 修改目录的权限和属组 +restore_stash_mod() { + local install_path="$1" + local line="$2" + local target + local target_abs + local mod + local is_abs_path ret + + __index_list "${line}" 0 "target" 1 "mod" + + __set_abs_path "${install_path}" "${target}" "target_abs" + + # 目录不存在时跳过 + if [ ! -d "${target_abs}" ]; then + return 0 + fi + # 只处理目录，没有处理目录的软链接 + change_mod "${target_abs}" "${mod}" "${INSTALL_FOR_ALL}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# 修改目录的权限，暂存原文件权限 +reset_mod_dirs_with_stash_mod() { + local install_path="$1" + local line="$2" + local mod + local target + local target_abs + local is_abs_path ret + + __index_list "${line}" 1 "target" + + __set_abs_path "${install_path}" "${target}" "target_abs" + + # 目录不存在时跳过 + if [ ! -d "${target_abs}" ]; then + return 0 + fi + + get_file_mod "mod" "${target_abs}" + echo "${target}:${mod}" >> "${install_path}/${STASH_MOD_PATH}" + + # 只处理目录，没有处理目录的软链接 + change_mod "${target_abs}" "${RESET_MOD}" "" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# 删除软连接列表 +remove_softlinks() { + local install_path="$1" + local softlinks_str="$2" + local softlinks softlink softlink_abs ret + + __unpack_softlinks "${softlinks_str}" "softlinks" + + for softlink in ${softlinks}; do + __set_abs_path "${install_path}" "${softlink}" "softlink_abs" + remove_softlink_icp "${softlink_abs}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + done + return 0 +} + +# 删除安装文件夹 +remove_install_dirs() { + local install_path="$1" + local line="$2" + local target softlinks_str is_abs_path + + __index_list "${line}" 1 "target" 4 "softlinks_str" + + if [ "${target}" != "NA" ]; then + __check_abs_path "${target}" + if [ "${is_abs_path}" != "true" ]; then + target="${install_path}/${target}" + fi + if [ -d "${target}" ]; then + # 不同的blocks中，可能配置相同的dir_info。目录不为空时不删除。 + if [ "$(ls -A "${target}")" != "" ]; then + return 0 + fi + remove_dir_icp "${target}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + if [ "${softlinks_str}" != "NA" ]; then + remove_softlinks "${install_path}" "${softlinks_str}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + fi + return 0 +} + +# 获取所有的blocks +get_blocks_info() { + local _outvar="$1" + local _install_path="$2" + local _result _db_filepath="${_install_path}/${PKG_DB_INFO_RELPATH}" + + if [ ! -f "${_db_filepath}" ]; then + return 0 + fi + _result="$(cut -d'|' -f1 "${_db_filepath}" | xargs)" + eval "${_outvar}=\"${_result}\"" +} + +# 列表转为正则表达式 +trans_list_to_regex() { + local _outvar="$1" + local _value="$2" + # 移除右侧空白符，mawk无法编译右侧存在空白符时的正则表达式，如：^(EngineeringCommon|)$ + local _result="$(echo "$_value" | sed 's/\s\+$//g' | sed 's/ /|/g' | xargs printf "^(%s)$")" + eval "${_outvar}=\"${_result}\"" +} + +parse_filelist_core() { + awk -F, '{ print $3,$4,$6,$7,$9,$12,$13,$14,$15 }' +} + +# 过滤公共路径条目 +filter_common_dirs() { + awk -F, ' + $11 == "Y" { + print $0 + } + $11 == "YY" { + print $1 "," $2 "," $3 "," $4 "," $5 "," $6 "," $7 "," $8 ",NA," $10 "," $11 "," $12 "," $13 "," $14 "," $15 "," $16 + } + ' +} + +# 过滤包内软链 +filter_pkg_inner_softlink() { + awk -F, '$15 != "NA" { print $0 }' +} + +# 过滤块 +filter_blocks() { + local blocks="$1" + local blocks_reg + + trans_list_to_regex "blocks_reg" "$blocks" + awk -F, "\$14 ~ \"$blocks_reg\" {print \$0}" +} + +# 过滤不匹配块 +filterfalse_blocks() { + local blocks="$1" + local blocks_reg + + trans_list_to_regex "blocks_reg" "$blocks" + awk -F, "\$14 !~ \"$blocks_reg\" {print \$0}" +} + +# 过滤操作类型条目 +filter_operate_type() { + local operator_type="$1" + local operator_type_reg + + trans_list_to_regex "operator_type_reg" "$operator_type" + + # ~为部分匹配，正则表达式需要匹配头尾 + awk -F, "\$2 ~ \"$operator_type_reg\" {print \$0}" +} + +# 过滤安装类型条目 +filter_install_type() { + local install_type="$1" + local install_type_reg + + # 注意：特征串不能嵌套，否则会错误匹配 + if [ "$install_type" != "full" ] && [ "$install_type" != "debug" ]; then + install_type_reg="(all)|($install_type)" + else + install_type_reg="(all)|(docker)|(devel)|(run)" + fi + + awk -F, "\$8 ~ \"$install_type_reg\" {print \$0}" +} + +# 过滤特性参数 +filter_feature_param() { + local feature_param="$1" + local feature_type + local feature_exclude_all # feature是否为排除公共all文件 + local chip chip_list # 芯片类型 + local feature_list feature_type_list + + unpack_feature_param "feature_type" "feature_exclude_all" "chip" ${feature_param} + + if [ "$feature_type" != "all" ]; then + feature_list="$(echo $feature_type | tr ',' ' ')" + if [ "${feature_exclude_all}" = "y" ]; then + feature_type_list="${feature_list}" + else + feature_type_list="comm ${feature_list}" + fi + else + feature_type_list="all" + fi + + if [ "$chip" != "all" ]; then + chip_list="all $(echo $chip | tr ',' ' ')" + else + chip_list="all" + fi + + # filelist中的feature为all，表示所有场景下都安装这个文件 + # 输入参数中的feature为all，表示安装所有的feature + awk -v feature_type_list="${feature_type_list}" \ + -v chip_list="${chip_list}" ' + BEGIN { + FS= "," + split(feature_type_list, input_feature_type_arr, " ") + split(chip_list, input_chip_arr, " ") + } + + function match_feature_type(features_str) { + if (feature_type_list == "all") { + return 1 + } + matched_feature_type_tmp=0 + split(features_str, features, ";") + for(i in features) + { + for(j in input_feature_type_arr) { + if(input_feature_type_arr[j] == features[i]) { + matched_feature_type_tmp = 1 + break; + } + } + } + return matched_feature_type_tmp + } + + function match_chip(chip_str) { + if (chip_list == "all") { + return 1 + } + matched_chip_tmp=0 + split(chip_str, chips, ";") + for(i in chips) + { + for(j in input_chip_arr) { + if(input_chip_arr[j] == chips[i]) { + matched_chip_tmp = 1 + break; + } + } + } + return matched_chip_tmp + } + + { + matched_feature_type = match_feature_type($10); + if (matched_feature_type == 0) next; + + matched_chip = match_chip($16) + if (matched_chip == 0) next + + print $0 + }' +} + +# 读取fileist的第2行到最后行 +tail_filelist() { + local filelist_path="$1" + tail -n +2 "$filelist_path" +} + +# 解析filelist.csv脚本 +parse_filelist() { + local install_type="$1" + local operate_type="$2" + local filelist_path="$3" + local feature_param="$4" + local filter_type="$5" + local blocks="$6" + + if [ ! -f "$filelist_path" ]; then + log "ERROR" "filelist $filelist_path does not exist!" + exit 1 + fi + + # 注意：这里的filelist需要是全局变量。其它函数（如add_filelist_blocks_info）会引用这个变量。 + filelist=$( + parse_filelist_v2 "$install_type" "$operate_type" "$filelist_path" "$feature_param" \ + "$filter_type" "$blocks" + ) +} + +# 解析filelist.csv脚本 +parse_filelist_v2() { + local install_type="$1" + local operate_type="$2" + local filelist_path="$3" + local feature_param="$4" + local filter_type="$5" + local blocks="$6" + local filter_cmds="" + + if [ "$operate_type" != "all" ]; then + filter_cmds="$filter_cmds | filter_operate_type \"$operate_type\"" + fi + if printf "%s" "$filter_type" | grep -Eq "\"; then + filter_cmds="$filter_cmds | filter_common_dirs" + fi + if printf "%s" "$filter_type" | grep -Eq "\"; then + filter_cmds="$filter_cmds | filter_pkg_inner_softlink" + fi + if printf "%s" "$filter_type" | grep -Eq "\|\" && [ "$blocks" != "" ]; then + filter_cmds="$filter_cmds | filterfalse_blocks \"$blocks\"" + fi + if printf "%s" "$filter_type" | grep -Eq "\" && [ "$blocks" != "" ]; then + filter_cmds="$filter_cmds | filter_blocks \"$blocks\"" + fi + + tail_filelist "$filelist_path" \ + | filter_install_type "$install_type" \ + | filter_feature_param "$feature_param" \ + | eval cat $filter_cmds \ + | parse_filelist_core +} + + +pack_exec_params() { + local _outvar="$1" + local _exec_mode="$2" + + eval "${_outvar}=\"${_exec_mode}\"" +} + + +# 迭代filelist中的条目，执行操作 +foreach_filelist_exec() { + local filelist="$1" + local sort_filelist="$2" + local exec_mode="$3" + local exec_func="$4" + local install_path="$5" + local ret=0 tmp_ret exec_params + shift 5 + + pack_exec_params "exec_params" "${exec_mode}" + + if [ "${sort_filelist}" = "reverse" ]; then + # 对第二列文件路径做倒序排列，保证先删除子文件夹，再删除父文件夹 + # 不可以使用echo "${filelist}" + # 在dash中会消耗\\$username:\\$usergroup中的一个反斜线（bash与busybox的sh无此问题） + # 需要使用here document，防止shell的解析过程 + filelist=$(sort -k2,2 -b -r < /dev/null + else + log "ERROR" "sort_type param wrong! sort_type is ${sort_type}, exec_func is ${exec_func}" + exit 1 + fi + + while read line; do + array=${line} + __length_list "${array}" "len_array" + if [ ${len_array} -eq 0 ]; then + continue + fi + "${exec_func}" "${install_path}" "${line}" + tmp_ret="$?" && [ ${tmp_ret} -ne 0 ] && ret="${tmp_ret}" + done << EOF +${stashmod_list} +EOF + + return $ret +} + +# 创建目录并且设置权限 +make_dir_with_permission() { + local path="$1" + local mod="$2" + local username="$3" + local usergroup="$4" + local install_for_all="$5" + local ret + + make_dir "${path}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + change_mod "${path}" "${mod}" "${install_for_all}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + change_own "${path}" "${username}:${usergroup}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# blocks转换为db.info配置 +blocks_to_db_item() { + local package="$1" + local version_dir="$2" + + xargs printf "%s|[$package:$version_dir]\n" +} + +# 折叠第2个参数 +# 需要保证输入数据第1个参数已排序 +fold_2() { + local fs="$1" + local concat="$2" + local options="" + + if [ "$fs" != "" ]; then + options="-F$fs" + fi + + awk $options -v CONCAT="$concat" ' + NR == 1 { + field1 = $1 + field2 = $2 + } + NR != 1 { + if (field1 == $1) { + field2 = field2 CONCAT $2 + } else { + print field1 FS field2 + field1 = $1 + field2 = $2 + } + } + END { + if (field1) { + print field1 FS field2 + } + }' +} + +# 折叠第3个参数，保持第2个参数 +# 需要保证输入数据第1个参数已排序 +fold_3_keep_2() { + local fs="$1" + local concat="$2" + local options="" + + if [ "$fs" != "" ]; then + options="-F$fs" + fi + + awk $options -v CONCAT="$concat" ' + NR == 1 { + field1 = $1 + field2 = $2 + field3 = $3 + } + NR != 1 { + if (field1 == $1) { + field3 = field3 CONCAT $3 + } else { + print field1 FS field2 FS field3 + field1 = $1 + field2 = $2 + field3 = $3 + } + } + END { + if (field1) { + print field1 FS field2 FS field3 + } + }' +} + +# 根据第1列排序 +sort_1() { + local sep="$1" + local options="" + if [ "$sep" != "" ]; then + options="-t$sep" + fi + sort $options -k1,1 -s +} + +# 显示3,4列有差异的条目 +show_diff_3_4() { + awk '$3 != $4 { print $0 }' +} + +# 显示具有最少列数的条目 +show_min_nf() { + local min_nf="$1" + awk -v MIN_NF="$min_nf" 'NF >= MIN_NF { print $0 }' +} + +# 选取第3,2,1列 +select_fields_3_2_1() { + awk '{print $3, $2, $1}' +} + +# 选取第1列 +select_fields_1() { + awk '{print $1}' +} + +# 删除db.info配置 +del_db_items() { + local package="$1" + local version_dir="$2" + + sed "s/\\[$package:$version_dir\\]//g; /|\$/d" +} + +# 保留块最后一个配置 +# 输出：第1列（块），第2列（包名），第3列（版本目录） +remain_db_last_item() { + awk -F '[|:\\[\\]]+' '{ print $1, $(NF-2), $(NF-1) }' +} + +# 移除空白行 +remove_blank_line() { + sed '/^$/d' +} + +# 保证文件权限 +ensure_permission() { + local path="$1" + local mod="$2" + local username="$3" + local usergroup="$4" + local install_for_all="$5" + + change_mod "$path" "$mod" "$install_for_all" + ret=$? && [ $ret -ne 0 ] && return $ret + + change_own "$path" "$username:$usergroup" + ret=$? && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 创建文件 +touch_file() { + local path="$1" + local mod="$2" + local username="$3" + local usergroup="$4" + local install_for_all="$5" + local ret + + touch "$path" + ret=$? && [ $ret -ne 0 ] && return $ret + + ensure_permission "$path" "$mod" "$username" "$usergroup" "$install_for_all" + ret=$? && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 保证文件存在 +ensure_file() { + local path="$1" + local mod="$2" + local username="$3" + local usergroup="$4" + local install_for_all="$5" + local ret + + if [ ! -f "$path" ]; then + touch_file "$path" "$mod" "$username" "$usergroup" "$install_for_all" + ret=$? && [ $ret -ne 0 ] && return $ret + fi + + return 0 +} + +# 写文本文件 +write_text() { + local content="$1" + local filepath="$2" + + printf "%s\n" "$content" > "$filepath" +} + +# 修改文件权限并操作 +with_chmod() { + local path="$1" + local mod="$2" + local origin_mod ret + shift 2 + + origin_mod="$(stat -L -c "%a" "$path")" + chmod "$mod" "$path" + "$@" + ret="$?" + chmod "$origin_mod" "$path" + + return $ret +} + +# filelist中所有公共目录的块 +all_common_dirs_blocks_in_filelist() { + local install_type="$1" + local filelist_path="$2" + local feature_param="$3" + + parse_filelist_v2 "$install_type" "all" "$filelist_path" "$feature_param" "filter_common_dirs" "" \ + | cut -d' ' -f8 | sort | uniq +} + +# 是否为版本目录 +is_version_dirpath() { + local version_dirpath="$1" + if [ -d "$version_dirpath/share/info" ]; then + return 0 + fi + if [ -f "$version_dirpath/var/ascend_package_db.info" ]; then + return 0 + fi + return 1 +} diff --git a/csrc/ascend/scripts/package/common/sh/common_installer.inc b/csrc/ascend/scripts/package/common/sh/common_installer.inc new file mode 100644 index 000000000..f8b706fce --- /dev/null +++ b/csrc/ascend/scripts/package/common/sh/common_installer.inc @@ -0,0 +1,55 @@ +#!/bin/sh +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +# 通知latest管理器创建版本软链 +notify_latest_manager_create_version_softlink() { + local curpath install_path version_dir var_path + + set_comm_log "Notifier" + + curpath="$(dirname $(readlink -f "${BASH_SOURCE:-$0}"))" + install_path="$(readlink -f "$curpath/..")" + version_dir="$(basename "$curpath")" + var_path="$install_path/$LATEST_DIR/var" + + if [ ! -f "$var_path/manager.sh" ]; then + comm_log "ERROR" "$var_path/manager.sh doesn't exist!" + exit 2 + fi + + if ! "$var_path/manager.sh" --version-dir "$version_dir" create_version_softlink; then + comm_log "ERROR" "create version softlink failed!" + exit 1 + fi + return 0 +} + +# 通知latest管理器删除latest软链 +notify_latest_manager_remove_latest_softlink() { + local curpath install_path var_path + + set_comm_log "Notifier" + + curpath="$(dirname $(readlink -f "${BASH_SOURCE:-$0}"))" + install_path="$(readlink -f "$curpath/..")" + var_path="$install_path/$LATEST_DIR/var" + + if [ ! -f "$var_path/manager.sh" ]; then + comm_log "ERROR" "$var_path/manager.sh doesn't exist!" + exit 2 + fi + + if ! "$var_path/manager.sh" remove_latest_softlink; then + comm_log "ERROR" "remove latest softlink failed!" + exit 1 + fi + return 0 +} diff --git a/csrc/ascend/scripts/package/common/sh/common_interface.csh b/csrc/ascend/scripts/package/common/sh/common_interface.csh new file mode 100644 index 000000000..d56269687 --- /dev/null +++ b/csrc/ascend/scripts/package/common/sh/common_interface.csh @@ -0,0 +1,37 @@ +#!/bin/csh +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +set func_name = "$1" +switch ( "$func_name" ) + case "mk_custom_path": + if ( "`id -u`" == 0 ) then + exit 0 + endif + set file_path = "$2" + foreach line ("` cat $file_path `") + set custom_path = "`echo '$line' | cut --only-delimited -d= -f2`" + if ( "$custom_path" == "" ) then + continue + endif + set custom_path = "` eval echo $custom_path `" + if ( ! -d "$custom_path" ) then + mkdir -p "$custom_path" + if ( $status != 0 ) then + set cur_date = "`date +'%Y-%m-%d %H:%M:%S'`" + echo "[Common] [$cur_date] [ERROR]: create $custom_path failed." + exit 1 + endif + endif + end + breaksw + default: + breaksw +endsw diff --git a/csrc/ascend/scripts/package/common/sh/common_interface.fish b/csrc/ascend/scripts/package/common/sh/common_interface.fish new file mode 100644 index 000000000..b5369185d --- /dev/null +++ b/csrc/ascend/scripts/package/common/sh/common_interface.fish @@ -0,0 +1,33 @@ +#!/usr/bin/env fish +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +function mk_custom_path + set -l custom_file_path $argv[1] + if test (id -u) -eq 0 + return 0 + end + while read line + set -l _custom_path (echo "$line" | cut --only-delimited -d= -f2) + if test -z $_custom_path + continue + end + set -l _custom_path (eval echo "$_custom_path") + if not test -d $_custom_path + mkdir -p "$_custom_path" + if not test $status -eq 0 + set -l cur_date (date +"%Y-%m-%d %H:%M:%S") + echo "[Common] [$cur_date] [ERROR]: create $_custom_path failed." + return 1 + end + end + end < $custom_file_path + return 0 +end diff --git a/csrc/ascend/scripts/package/common/sh/common_interface.sh b/csrc/ascend/scripts/package/common/sh/common_interface.sh new file mode 100644 index 000000000..abadbdb07 --- /dev/null +++ b/csrc/ascend/scripts/package/common/sh/common_interface.sh @@ -0,0 +1,61 @@ +#!/bin/sh +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +mk_custom_path() { + if [ $(id -u) -eq 0 ]; then + return 0 + fi + local _custom_path_file="$1" + while read line || [ -n "$line" ] + do + local _custom_path="$(echo "$line" | cut --only-delimited -d= -f2)" + if [ -z "$_custom_path" ]; then + continue + fi + eval "_custom_path=$_custom_path" + if [ ! -d "$_custom_path" ]; then + mkdir -p "$_custom_path" + if [ $? -ne 0 ]; then + cur_date="$(date +"%Y-%m-%d %H:%M:%S")" + echo "[Common] [$cur_date] [ERROR]: create $_custom_path failed." + return 1 + fi + fi + done < $_custom_path_file + return 0 +} + +py_version_check(){ + local pyver_set="3.7 3.8 3.9 3.10 3.11 3.12" + local cur_date="$(date +"%Y-%m-%d %H:%M:%S")" + which python3 > /dev/null 2>&1 + if [ $? -eq 0 ]; then + local python_version="$(python3 --version 2>&1 | head -n 1)" + local python3_version=$(echo "$python_version" | sed -n 's/.*[^\.0-9]\([0-9]\+\.[0-9]\+\).*/\1/p') + if [ "x$python3_version" != "x" ]; then + for ver in $pyver_set; do + if [ "x$ver" = "x$python3_version" ]; then + return 0 + fi + done + + echo "[Common] [$cur_date] [WARNING]: $python_version is not in Python3.7.x, Python3.8.x, Python3.9.x, Python3.10.x, Python3.11.x, Python3.12.x" + return 1 + else + echo "[Common] [$cur_date] [WARNING]: $python_version cannot be identified as a standard version, please check manually." + return 1 + fi + else + echo "[Common] [$cur_date] [WARNING]: python3 is not found." + return 1 + fi +} + diff --git a/csrc/ascend/scripts/package/common/sh/install_common_parser.sh b/csrc/ascend/scripts/package/common/sh/install_common_parser.sh new file mode 100644 index 000000000..1fc4b439f --- /dev/null +++ b/csrc/ascend/scripts/package/common/sh/install_common_parser.sh @@ -0,0 +1,2192 @@ +#!/bin/sh +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +# run包安装解析公共脚本 +# 解析filelist.csv文件，完成目录创建，文件复制，权限设置，文件删除等操作。 + +# minirc场景下存在使用dash调用该脚本的情况 +#export PS4='+ ${FUNCNAME[0]:+${FUNCNAME[0]}():} ${BASH_SOURCE}:${LINENO}: ' +#set -x +# 总uninstall.sh文件权限 +TOTAL_UNINSTALL_MOD="500" +# db.info文件权限 +DB_INFO_MOD="640" +# 包架构 +PKG_ARCH="UNKNOWN" +# 复制文件是否并发 +PARALLEL="" +# 是否限制并发 +PARALLEL_LIMIT="" +# 控制并发的fifo文件路径 +PARALLEL_FIFO="none" +COPY_COMMAND="cp -rf" +MOVE_COMMAND="mv -f" + +curpath="$(dirname $(readlink -f "${BASH_SOURCE:-$0}"))" + +# 导入公共库 +common_func_v2_path="${curpath}/common_func_v2.inc" +. "${common_func_v2_path}" + +# 导入common_func.inc +# 不一定能source到common_func.inc +# 例如driver包中的install_common_parser.sh与common_func.inc，不在同一个目录下 +# 需要子包整改，才将该if语句移除 +common_func_path="${curpath}/common_func.inc" +if [ -f "${common_func_path}" ]; then + . "${common_func_path}" +fi + +# 导入script_operator.inc +script_operator_path="${curpath}/script_operator.inc" +if [ -f "${script_operator_path}" ]; then + . "${script_operator_path}" +fi + +# 导入version_cfg.inc +version_cfg_path="${curpath}/version_cfg.inc" +if [ -f "${version_cfg_path}" ]; then + . "${version_cfg_path}" +fi + +# 导入multi_version.inc +multi_version_path="${curpath}/multi_version.inc" +if [ -f "${multi_version_path}" ]; then + . "${multi_version_path}" +fi + +# 导入version_compatiable.inc +version_compatiable_path="${curpath}/version_compatiable.inc" +if [ -f "${version_compatiable_path}" ]; then + . "${version_compatiable_path}" +fi + +# 导入config.inc +config_path="${curpath}/config.inc" +if [ -f "${config_path}" ]; then + . "${config_path}" +fi + +# 导入cold_patch.sh +cold_patch_path="${curpath}/cold_patch.sh" +if [ -f "${cold_patch_path}" ]; then + . "${cold_patch_path}" +fi + +# bash执行时需要展开别名 +if [ -n "$BASH_SOURCE" ]; then + shopt -s expand_aliases +fi + +__pkg_in_pkgs() { + local matched + __item_in_list "matched" "$@" + echo "${matched}" +} + +__block_in_blocks() { + __item_in_list "matched" "$@" + echo "${matched}" +} + +# 包列表中移除包 +__remove_pkg_in_pkgs() { + __remove_item_in_list "$@" +} + +# 冒泡排序，O(n**2) +__sort_blocks() { + local blocks="$*" + local block + local sorted_blocks="" + local max + local i=0 + + __length_list "${blocks}" "len_blocks" + + while [ "$i" -lt ${len_blocks} ]; do + max="" + for block in ${blocks}; do + if [ "${block}" \> "${max}" ]; then + max="${block}" + fi + done + if [ "${sorted_blocks}" = "" ]; then + sorted_blocks="${max}" + else + sorted_blocks="${max} ${sorted_blocks}" + fi + blocks="$(__remove_item_in_list "${max}" "${blocks}")" + i=$((i+1)) + done + + echo "${sorted_blocks}" +} + +__unpack_block_item() { + local line="$1" + echo "${line}" | tr '|' ' ' +} + +__pack_block_item() { + local block_name="$1" + local block_pkgs_str="$2" + echo "${block_name}|${block_pkgs_str}" +} + +__unpack_block_pkgs() { + local block_pkgs_str="$1" + echo "${block_pkgs_str}" | tr ',' ' ' +} + +__pack_block_pkgs() { + echo "$*" | tr ' ' ',' +} + +check_install_path() { + local install_path="$1" + + if [ "${install_path}" = "" ]; then + log "ERROR" "install_path is empty!" + return 1 + fi +} + +# 添加pkg对应的block_info +add_pkg_blocks_info() { + local install_path="$1" + local pkg="$2" + shift; shift + # blocks为待加入db.info的块列表 + local blocks="$*" + local block_idx_value + + local db_filepath="${install_path}/${PKG_DB_INFO_RELPATH}" + local db_dirpath="$(dirname "${db_filepath}")" + + local db_filepath_new="${install_path}/${PKG_DB_INFO_RELPATH}~" + + local block_list + local block_name + local block_pkgs_str + local block_pkgs + local len_block_pkgs + local block_idx=0 + local len_blocks + + __length_list "${blocks}" "len_blocks" + if [ ${len_blocks} -eq 0 ]; then + return 0 + fi + + blocks="$(__sort_blocks "${blocks}")" + + if [ ! -d "${db_dirpath}" ]; then + mkdir -p "${db_dirpath}" + if [ $? -ne 0 ]; then + log "ERROR" "mkdir ${db_dirpath} failed!" + exit 1 + fi + fi + + if [ ! -f "${db_filepath}" ]; then + touch "${db_filepath}" + if [ $? -ne 0 ]; then + log "ERROR" "touch ${db_filepath} failed!" + exit 1 + fi + change_mod_and_own "${db_filepath}" "${DB_INFO_MOD}" "\$username:\$usergroup" "${INSTALL_FOR_ALL}" + fi + + rm -f "${db_filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "delete ${db_filepath_new} failed!" + exit 1 + fi + + while read line; do + block_list="$(__unpack_block_item "${line}")" + __index_list "${block_list}" 0 "block_name" + __index_list "${block_list}" 1 "block_pkgs_str" + + # 如果pkgs为空，则忽略该条目 + if [ "${block_pkgs_str}" = "" ]; then + continue + fi + + block_pkgs="$(__unpack_block_pkgs "${block_pkgs_str}")" + __length_list "${block_pkgs}" "len_block_pkgs" + if [ ${len_block_pkgs} -eq 0 ]; then + continue + fi + + if [ ${block_idx} -lt ${len_blocks} ]; then + __index_list "${blocks}" ${block_idx} "block_idx_value" + fi + + # blocks中的块名小于db中的当前块名，直接添加到db中。 + # 注意不能continue，因为还要处理db的当前条目。 + while [ ${block_idx} -lt ${len_blocks} ] && [ "${block_idx_value}" \< "${block_name}" ]; do + __pack_block_item "${block_idx_value}" "${pkg}" >> "${db_filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "write ${db_filepath_new} failed!" + exit 1 + fi + block_idx=$((block_idx+1)) + __index_list "${blocks}" ${block_idx} "block_idx_value" + done + + # 已经没有需要添加的块，保持db的当前条目。 + if [ ${block_idx} -eq ${len_blocks} ]; then + __pack_block_item "${block_name}" "${block_pkgs_str}" >> "${db_filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "write ${db_filepath_new} failed!" + exit 1 + fi + continue + fi + + # blocks中的块名大于db中的当前块名，保持db的当前条目。 + if [ "${block_idx_value}" \> "${block_name}" ]; then + __pack_block_item "${block_name}" "${block_pkgs_str}" >> "${db_filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "write ${db_filepath_new} failed!" + exit 1 + fi + continue + fi + + # blocks中的块名等于db中的当前块名，将pkg添加到当前条目pkg列表中。 + in_pkgs="$(__pkg_in_pkgs "${pkg}" "${block_pkgs}")" + if [ "${in_pkgs}" = "true" ]; then + __pack_block_item "${block_name}" "${block_pkgs_str}" >> "${db_filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "write ${db_filepath_new} failed!" + exit 1 + fi + else + __pack_block_item "${block_name}" "${block_pkgs_str},${pkg}" >> "${db_filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "write ${db_filepath_new} failed!" + exit 1 + fi + fi + block_idx=$((block_idx+1)) + done < "${db_filepath}" + + # blocks中剩余的块添加到db中。 + while [ ${block_idx} -lt ${len_blocks} ]; do + __index_list "${blocks}" ${block_idx} "block_idx_value" + echo "${block_idx_value}|${pkg}" >> "${db_filepath_new}" + block_idx=$((block_idx+1)) + done + + cp "${db_filepath_new}" "${db_filepath}" + if [ $? -ne 0 ]; then + log "ERROR" "replace ${db_filepath} failed!" + exit 1 + fi + + rm -f "${db_filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "delete ${db_filepath_new} failed!" + exit 1 + fi +} + +# 删除block_info中的pkg信息 +del_blocks_info_pkg() { + local install_path="$1" + local pkg="$2" + + local db_filepath="${install_path}/${PKG_DB_INFO_RELPATH}" + local db_dirpath="$(dirname "${db_filepath}")" + + local db_filepath_new="${install_path}/${PKG_DB_INFO_RELPATH}~" + + local block_list + local block_name + local block_pkgs_str + local block_pkgs + local len_block_pkgs + + if [ ! -f "${db_filepath}" ]; then + return 0 + fi + + rm -f "${db_filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "delete ${db_filepath_new} failed!" + exit 1 + fi + + while read line; do + block_list="$(__unpack_block_item "${line}")" + + __index_list "${block_list}" 0 "block_name" + __index_list "${block_list}" 1 "block_pkgs_str" + + block_pkgs="$(__unpack_block_pkgs "${block_pkgs_str}")" + + block_pkgs="$(__remove_pkg_in_pkgs "${pkg}" "${block_pkgs}")" + + __length_list "${block_pkgs}" "len_block_pkgs" + + if [ ${len_block_pkgs} -gt 0 ]; then + block_pkgs_str="$(__pack_block_pkgs "${block_pkgs}")" + __pack_block_item "${block_name}" "${block_pkgs_str}" >> "${db_filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "write ${db_filepath_new} failed!" + exit 1 + fi + fi + done < "${db_filepath}" + + # 当前条目如果没有pkg使用，则删除。 + if [ ! -f "${db_filepath_new}" ]; then + rm -f "${db_filepath}" + if [ $? -ne 0 ]; then + log "ERROR" "delete ${db_filepath} failed!" + exit 1 + fi + else + mv -f "${db_filepath_new}" "${db_filepath}" + if [ $? -ne 0 ]; then + log "ERROR" "replace ${db_filepath} failed!" + exit 1 + fi + fi + return 0 +} + +# 准备并且检查软链接路径 +prepare_and_check_softlink_path() { + local softlink_abs="$1" + local softlink_dir="$(dirname "${softlink_abs}")" + + if [ ! -d "${softlink_dir}" ]; then + mkdir -p "${softlink_dir}" + fi + + # 如果目标路径是个软链接，则移除 + if [ -L "${softlink_abs}" ]; then + rm -f "${softlink_abs}" + if [ $? -ne 0 ]; then + log "ERROR" "remove softlink ${softlink_abs} failed! (create relative softlink)" + exit 1 + fi + fi + + # 不允许目标路径已经是一个目录，防止软链接到错误的位置。 + if [ -d "${softlink_abs}" ]; then + log "ERROR" "softlink existed dir ${softlink_abs}!" + exit 1 + fi +} + +get_file_owner_group() { + local _outvar="$1" + local _path="$2" + local _result + + _result="$(stat -c %U "${_path}"):$(stat -c %G "${_path}")" + eval "${_outvar}=\"${_result}\"" +} + +# 根据pkg_inner_softlink创建软链接 +create_pkg_inner_softlink() { + local install_path="$1" + local line="$2" + local array + local target + local pkg_inner_softlink + local pkg_inner_softlink_list + local is_abs_path + + __index_list "${line}" 1 "target" 8 "pkg_inner_softlink" + + if [ "${pkg_inner_softlink}" != "NA" ]; then + create_softlink_by_install_path "${install_path}" "${target}" "${pkg_inner_softlink}" + fi + return 0 +} + +# 处理软链路径为已存在目录的情况 +# minirc场景，install_path目录下，存在include目录 +# 并且filelist.csv中配置了(x86_64|aarch64)-linux/include目录 +# 则将(x86_64|aarch64)-linux/include中的内容，软链接至include目录下 +deal_with_existed_dir() { + local install_path="$1" + local folder="$2" + local filelist_path="$3" + local install_path_regex="^(x86_64|aarch64)-linux/${folder}$" + local install_sub_path_regex="^(x86_64|aarch64)-linux/${folder}/[^/]+$" + + if [ -L "${install_path}/${folder}" ] || [ ! -d "${install_path}/${folder}" ]; then + return 0 + fi + + target_path=$(awk -v folder="${folder}" -v install_path_regex="${install_path_regex}" ' + BEGIN{ + FS= "," + } + { + if ($2 != "mkdir") next + + if ($4 !~ install_path_regex) next + + if ($9 != folder) next + + print $4 + }' "${filelist_path}") + + if [ "${target_path}" = "" ]; then + return 0 + fi + + awk -v folder="${folder}" -v install_path_regex="${install_path_regex}" -v install_sub_path_regex="${install_sub_path_regex}" ' + BEGIN{ + FS = ","; OFS = "," + } + { + if ($2 != "mkdir" && $2 != "copy" && $2 != "move") { + print $0 + next + } + if ($4 ~ install_path_regex) { + print $1, $2, $3, $4, $5, $6, $7, $8, "NA", $10, $11, $12, $13, $14, $15 + next + } + if ($4 !~ install_sub_path_regex) { + print $0 + next + } + softlink = $9 + if (softlink != "NA") { + print $0 + next + } + + z = split($4, filepath_list, "/") + softlink = folder "/" filepath_list[z] + print $1, $2, $3, $4, $5, $6, $7, $8, softlink, $10, $11, $12, $13, $14, $15 + }' "${filelist_path}" > "${filelist_path}.tmp" + if [ $? -ne 0 ]; then + log "ERROR" "modify filelist for ${install_path}/${folder} failed!" + return 1 + fi + + mv -f "${filelist_path}.tmp" "${filelist_path}" + if [ $? -ne 0 ]; then + log "ERROR" "replace filelist for ${install_path}/${folder} failed!" + return 1 + fi +} + +#执行创建目录动作 +do_create_dirs() { + local action="$1" + local install_type="$2" + local install_path="$3" + local filelist_path="$4" + local package="$5" + local feature_param="$6" + local ret + + if [ "${action}" = "resetmod" ]; then + if [ "${package}" != "" ]; then + del_blocks_info_pkg "${install_path}" "${package}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + create_stash_mod "${install_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_filelist "NA" "reset_mod_dirs_with_stash_mod" "${install_type}" "${install_path}" "mkdir" "${filelist_path}" "${feature_param}" "no" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + else + foreach_filelist "NA" "reset_mod_dirs" "${install_type}" "${install_path}" "mkdir" "${filelist_path}" "${feature_param}" "no" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + # 设置copy_entity路径的权限，解决删除时权限问题 + # filter_by_blocks过滤，防止变更不删除的文件（目录）的权限 + foreach_filelist "filter_by_blocks" "reset_mod_dirs_recursive" "${install_type}" "${install_path}" "copy_entity" "${filelist_path}" "${feature_param}" "no" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + elif [ "${action}" = "all" ] || [ "${action}" = "mkdir" ]; then + deal_with_existed_dir "$install_path" "include" "$filelist_path" + ret="$?" && [ $ret -ne 0 ] && return $ret + + deal_with_existed_dir "$install_path" "lib64" "$filelist_path" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 先重置目录权限配置，防止软链接时缺少权限 + foreach_filelist "NA" "reset_mod_dirs" "$install_type" "$install_path" "mkdir" "$filelist_path" "${feature_param}" "no" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 重置copy_entity权限，防止块复用场景，复制时缺少权限 + foreach_filelist "NA" "reset_mod_dirs_recursive" "$install_type" "$install_path" "copy_entity" "$filelist_path" "${feature_param}" "no" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_filelist "NA" "create_dirs" "$install_type" "$install_path" "mkdir" "$filelist_path" "${feature_param}" "no" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_filelist "filter_by_pkg_inner_softlink" "create_pkg_inner_softlink" "$install_type" "$install_path" "mkdir" "$filelist_path" "${feature_param}" "no" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + else + log "ERROR" "action wrong! action is ${action}" + return 1 + fi + + return 0 +} + +# 拷贝文件 +copy_file() { + local install_path="$1" + local source="$2" + local target="$3" + local softlink="$4" + local pkg_inner_softlink="$5" + local target_abs + local target_dir + + __set_abs_path "${install_path}" "${target}" "target_abs" + target_dir="$(dirname "${target_abs}")" + + if [ ! -e "${source}" ]; then + log "ERROR" "copy file source file ${source} doesn't exist!" + return 1 + fi + + ${COPY_COMMAND} "${source}" "${target_dir}" + if [ $? -ne 0 ]; then + log "ERROR" "${source} copy failed!" + return 1 + fi + if [ "${softlink}" != "NA" ]; then + create_softlink_by_install_path "${install_path}" "${target}" "${softlink}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + if [ "${pkg_inner_softlink}" != "NA" ]; then + create_softlink_by_install_path "${install_path}" "${target}" "${pkg_inner_softlink}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + return 0 +} + +# 移动文件 +move_file() { + local install_path="$1" + local source="$2" + local target="$3" + local softlink="$4" + local pkg_inner_softlink="$5" + local target_abs + local target_dir + + __set_abs_path "${install_path}" "${target}" "target_abs" + target_dir="$(dirname "${target_abs}")" + + if [ ! -e "${source}" ] && [ ! -L "${source}" ]; then + log "ERROR" "move file source file ${source} doesn't exist!" + return 1 + fi + + ${MOVE_COMMAND} "${source}" "${target_dir}" + if [ $? -ne 0 ]; then + log "ERROR" "${source} move failed!" + return 1 + fi + if [ "${softlink}" != "NA" ]; then + create_softlink_by_install_path "${install_path}" "${target}" "${softlink}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + if [ "${pkg_inner_softlink}" != "NA" ]; then + create_softlink_by_install_path "${install_path}" "${target}" "${pkg_inner_softlink}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + return 0 +} + +# 复制文件 +copy_files() { + local install_path="$1" + local line="$2" + local exec_params="$3" + local src_dir="$4" + local src src_abs target target_abs softlink configurable pkg_inner_softlink tmpdir action + local ret + + __index_list "${line}" 0 "src" 1 "target" 4 "softlink" 5 "configurable" 8 "pkg_inner_softlink" + if [ "${src_dir}" = "--move" ]; then + src_abs="${src}" + action="move_file" + elif [ "${src_dir}" != "" ]; then + __set_abs_path "${src_dir}" "${src}" "src_abs" + action="copy_file" + else + src_abs="${src}" + action="copy_file" + fi + __set_abs_path "${install_path}" "${target}" "target_abs" + + tmpdir="$(dirname "${target_abs}")" + if [ ! -d "${tmpdir}" ]; then + mkdir -p "${tmpdir}" + fi + + # 如果目标文件已经存在，而且是配置文件，则不执行覆盖操作 + if [ -e "${target_abs}" ] && [ "${configurable}" = "TRUE" ]; then + return 0 + fi + + # 源文件不是软链，并且目标文件是软链 + if [ ! -L "${src_abs}" ] && [ -L "${target_abs}" ] ; then + rm -f "${target_abs}" + log "WARNING" "${target_abs} is an existing softlink in copy files, deleted." + fi + + exec_with_param "$exec_params" "$action" "${install_path}" "${src_abs}" "${target}" "${softlink}" "${pkg_inner_softlink}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + return 0 +} + +__set_total_uninstall_path() { + local install_path="$1" + local varname="$2" + + eval "${varname}=\"${install_path}/cann_uninstall.sh\"" +} + +# 是否为hilinux环境 +__is_hilinux() { + local varname="$1" + + which lsattr > /dev/null 2>&1 + if [ $? -ne 0 ]; then + eval "${varname}=\"true\"" + else + eval "${varname}=\"false\"" + fi +} + +# 移除文件上的不可修改权限 +__remove_immutable() { + local file="$1" + local is_hilinux ret + + # 文件不存在则退出。 + if [ ! -f "${file}" ]; then + return 0 + fi + + __is_hilinux "is_hilinux" + if [ "${is_hilinux}" = "true" ]; then + return 0 + fi + + attr="$(lsattr "${file}" | cut -d' ' -f1 | grep -o "i")" + if [ "${attr}" != "" ]; then + chattr -i "${file}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + return 0 +} + +# 创建总uninstall.sh脚本 +__create_uninstall() { + local file="$1" + + if [ ! -f "${file}" ]; then + cat > ${file} < /dev/null 2>&1 + if [ $? -ne 0 ]; then + log "WARNING" "Delete file:${file} failed, please delete it by yourself." + fi + fi + return 0 +} + +# 向cann_uninstall.sh文件中添加uninstall_package命令 +add_cann_uninstall_script_dir() { + local install_path="$1" + local script_dir="$2" + local username="$3" + local usergroup="$4" + local install_for_all="$5" + local oldmod="" ret + + __set_total_uninstall_path "${install_path}" "total_uninstall_path" + + if [ -f "${install_path}/${script_dir}/uninstall.sh" ]; then + if [ -f "${total_uninstall_path}" ]; then + get_file_mod "oldmod" "${total_uninstall_path}" + else + __create_uninstall "${total_uninstall_path}" + fi + ret="$?" && [ $ret -ne 0 ] && return $ret + + __remove_immutable "${total_uninstall_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + change_own "${total_uninstall_path}" "${username}:${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + change_mod "${total_uninstall_path}" "${SETENV_WRITEABLE_MOD}" "" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __add_uninstall_package "${total_uninstall_path}" "${script_dir}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "${oldmod}" = "" ]; then + change_mod "${total_uninstall_path}" "${TOTAL_UNINSTALL_MOD}" "${install_for_all}" + else + change_mod "${total_uninstall_path}" "${oldmod}" "" + fi + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + return 0 +} + +# 向cann_uninstall.sh文件中添加子包命令 +add_cann_uninstall_package() { + local install_path="$1" + local package="$2" + local username="$3" + local usergroup="$4" + local install_for_all="$5" + local ret package_dirpath script_dir + + get_package_dirpath "package_dirpath" "${package}" + script_dir="${package_dirpath}/script" + + # opp_kernel包存在同时安装多种芯片包的场景 + # 确保cann_uninstall.sh脚本中只有一个opp_kernel的uninstall_package + del_cann_uninstall_script_dir "${install_path}" "${script_dir}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + add_cann_uninstall_script_dir "${install_path}" "${script_dir}" "${username}" "${usergroup}" "${install_for_all}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# cann_uninstall.sh文件中添加命令（命令入口） +do_add_cann_uninstall() { + local install_path="$1" + local script_dir="$2" + local username="$3" + local usergroup="$4" + local install_for_all="$5" + local ret + + check_param_not_empty "install_path" "need set package parameter in add cann uninstall!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + check_param_not_empty "script_dir" "need set script_dir parameter in add cann uninstall!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + check_param_not_empty "username" "need set username parameter in add cann uninstall!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + check_param_not_empty "usergroup" "need set usergroup parameter in add cann uninstall!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + # 删除cann_uninstall.sh文件中已存在的uninstall_package命令 + del_cann_uninstall_script_dir "${install_path}" "${script_dir}" + add_cann_uninstall_script_dir "${install_path}" "${script_dir}" "${username}" "${usergroup}" "${install_for_all}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# 删除cann_uninstall.sh文件中uninstall_package命令 +del_cann_uninstall_script_dir() { + local install_path="$1" + local script_dir="$2" + local oldmod="" ret + + __set_total_uninstall_path "${install_path}" "total_uninstall_path" + + if [ -f "${total_uninstall_path}" ]; then + __remove_immutable "${total_uninstall_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + get_file_mod "oldmod" "${total_uninstall_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + change_mod "${total_uninstall_path}" "${SETENV_WRITEABLE_MOD}" "" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __remove_uninstall_package "${total_uninstall_path}" "${script_dir}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + change_mod "${total_uninstall_path}" "${oldmod}" "" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __remove_uninstall_file_if_no_content "${total_uninstall_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + return 0 +} + +# 删除cann_uninstall.sh文件中子包命令 +del_cann_uninstall_package() { + local install_path="$1" + local package="$2" + local ret package_dirpath script_dir + + get_package_dirpath "package_dirpath" "${package}" + script_dir="${package_dirpath}/script" + + del_cann_uninstall_script_dir "${install_path}" "${script_dir}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# cann_uninstall.sh文件中删除命令（命令入口） +do_del_cann_uninstall() { + local install_path="$1" + local script_dir="$2" + local ret + + check_param_not_empty "install_path" "need set package parameter in del cann uninstall!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + check_param_not_empty "script_dir" "need set script_dir parameter in del cann uninstall!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + del_cann_uninstall_script_dir "${install_path}" "${script_dir}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# 删除ascend_install.info文件 +del_ascend_install_info() { + rm -f "$curpath/../ascend_install.info" +} + +# 执行拷贝动作 +do_copy_files() { + local install_type="$1" + local install_path="$2" + local filelist_path="$3" + local package="$4" + local feature_param="$5" + local total_uninstall_path exec_mode ret + + if [ "$PARALLEL" = "true" ]; then + exec_mode="concurrency" + else + exec_mode="normal" + fi + if [ "${USE_MOVE}" = "true" ]; then + foreach_filelist "NA" "copy_files" "$install_type" "$install_path" "move" "$filelist_path" "${feature_param}" \ + "no" "$exec_mode" "--move" + foreach_filelist "NA" "copy_files" "$install_type" "$install_path" "copy copy_entity" "$filelist_path" "${feature_param}" \ + "no" "$exec_mode" + else + foreach_filelist "NA" "copy_files" "$install_type" "$install_path" "copy copy_entity move" "$filelist_path" "${feature_param}" \ + "no" "$exec_mode" + fi + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "${package}" != "" ] && [ "${SET_CANN_UNINSTALL}" = "y" ]; then + add_cann_uninstall_package "${install_path}" "${package}" "${USERNAME}" "${USERGROUP}" "${INSTALL_FOR_ALL}" + fi +} + +# 修改文件的权限和属组 +change_mod_and_own_files() { + local install_path="$1" + local line="$2" + local exec_params="$3" + local target mod own is_abs_path + + __index_list "${line}" 1 "target" 2 "mod" 3 "own" + __check_abs_path "${target}" + if [ "${is_abs_path}" != "true" ]; then + target="${install_path}/${target}" + fi + if [ -d "${target}" ]; then + return 0 + fi + # 只处理文件，没有处理文件的软链接 + exec_with_param "$exec_params" change_mod_and_own "${target}" "${mod}" "${own}" "${INSTALL_FOR_ALL}" "false" +} + +# 递归修改文件的权限和属组 +change_mod_and_own_files_recursive() { + local install_path="$1" + local line="$2" + local exec_params="$3" + local target mod own is_abs_path + + __index_list "${line}" 1 "target" 2 "mod" 3 "own" + __check_abs_path "${target}" + if [ "${is_abs_path}" != "true" ]; then + target="${install_path}/${target}" + fi + if [ ! -e "${target}" ]; then + return 0 + fi + + exec_with_param "$exec_params" change_mod_and_own "${target}" "${mod}" "${own}" "${INSTALL_FOR_ALL}" "true" +} + +# filelist中的使用到的blocks，添加到db.info中 +add_filelist_blocks_info() { + local install_type="$1" + local install_path="$2" + local filelist_path="$3" + local feature_param="$4" + local package="$5" + local blocks="" + + parse_filelist "${install_type}" "copy copy_entity del mkdir move" "${filelist_path}" "${feature_param}" "NA" "" + filelist="$(echo "${filelist}" | cut -d' ' -f8 | sort | uniq)" + + while read line; do + block_name="${line}" + if [ "${block_name}" = "" ]; then + continue + fi + if [ "${blocks}" = "" ]; then + blocks="${block_name}" + else + blocks="${blocks} ${block_name}" + fi + done << EOF +${filelist} +EOF + add_pkg_blocks_info "${install_path}" "${package}" "${blocks}" +} + +# 修改文件和目录的权限 +do_chmod_file_dir() { + local install_type="$1" + local install_path="$2" + local filelist_path="$3" + local feature_param="$4" + local package="$5" + local ret + + foreach_filelist "NA" "change_mod_and_own_files" "$install_type" "$install_path" "copy del move" "$filelist_path" "${feature_param}" "no" "concurrency" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_filelist "NA" "change_mod_and_own_files_recursive" "$install_type" "$install_path" "copy_entity" "$filelist_path" "${feature_param}" "no" "concurrency" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_filelist "NA" "change_mod_and_own_dirs" "$install_type" "$install_path" "mkdir" "$filelist_path" "${feature_param}" "reverse" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "${package}" != "" ]; then + add_filelist_blocks_info "$install_type" "$install_path" "$filelist_path" "${feature_param}" "$package" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + return 0 +} + +# 删除安装生成的pkg_inner_softlinks +remove_install_pkg_inner_softlinks() { + local install_path="$1" + local line="$2" + local target + local pkg_inner_softlink + local pkg_inner_softlink_list + local pkg_inner_softlink_abs + local ret + + __index_list "${line}" 1 "target" 8 "pkg_inner_softlink" + if [ "${target}" != "NA" ]; then + if [ "${pkg_inner_softlink}" != "NA" ]; then + remove_softlinks "${install_path}" "${pkg_inner_softlink}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + fi + return 0 +} + +# 删除安装文件，入参为$1: install_type, $2: install_path $3:filelist_path +remove_install_files() { + local install_path="$1" + local line="$2" + local exec_params="$3" + local target softlink configurable hash_value is_abs_path + + __index_list "${line}" 1 "target" 4 "softlink" 5 "configurable" 6 "hash_value" + + if [ "${target}" != "NA" ]; then + __check_abs_path "${target}" + if [ "${is_abs_path}" != "true" ]; then + target="${install_path}/${target}" + fi + if [ -d "${target}" ] && [ ! -L "${target}" ]; then + return 0 + fi + if [ "${softlink}" != "NA" ]; then + __check_abs_path "${softlink}" + if [ "${is_abs_path}" != "true" ]; then + softlink="${install_path}/${softlink}" + fi + fi + # 配置文件不删除 + if [ "${configurable}" = "TRUE" ]; then + echo "${hash_value} ${target}" | sha256sum --check > /dev/null 2>&1 + if [ $? -ne 0 ]; then + log "WARNING" "${target} user configuration file has been modified, skip deleting." + return 0 + fi + fi + exec_with_param "$exec_params" remove_file "${target}" "${softlink}" + fi + return 0 +} + +# 递归删除安装文件 +remove_install_files_recursive() { + local install_path="$1" + local line="$2" + local exec_params="$3" + local target configurable hash_value is_abs_path + + __index_list "${line}" 1 "target" 5 "configurable" 6 "hash_value" + + if [ "${target}" != "NA" ]; then + __check_abs_path "${target}" + if [ "${is_abs_path}" != "true" ]; then + target="${install_path}/${target}" + fi + # 配置文件不删除 + if [ "${configurable}" = "TRUE" ]; then + echo "${hash_value} ${target}" | sha256sum --check > /dev/null 2>&1 + if [ $? -ne 0 ]; then + log "WARNING" "${target} user configuration file has been modified, skip deleting." + return 0 + fi + fi + exec_with_param "$exec_params" remove_dir_icp "${target}" + fi + return 0 +} + +# 删除安装文件夹生成的软链接 +remove_install_softlink() { + local install_path="$1" + local line="$2" + local target + local softlinks_str + local is_abs_path + + __index_list "${line}" 1 "target" 4 "softlinks_str" + + if [ "$target" != "NA" ]; then + if [ "${softlinks_str}" != "NA" ]; then + remove_softlinks "${install_path}" "${softlinks_str}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + fi + return 0 +} + +# 删除安装文件与目录等 +do_remove() { + local install_type="$1" + local install_path="$2" + local filelist_path="$3" + local package="$4" + local feature_param="$5" + local tmp_filelist_path + local total_uninstall_path + local oldmod + local func_before_remove + local func_after_remove_01 + local func_after_remove_02 + local ret + + get_tmp_file tmp_filelist_path "filelist" + cp -f "$filelist_path" "$tmp_filelist_path" + + if [ $? -ne 0 ]; then + log "ERROR" "cp -f $filelist_path $tmp_filelist_path failed!" + return 1 + fi + + if [ "${package}" != "" ]; then + # root帐户--uninstall时，存在不调用restoremod，直接调用remove的场景。（toolkit包） + func_before_remove="del_blocks_info_pkg \"${install_path}\" \"${package}\"" + if [ -f "${install_path}/${STASH_MOD_PATH}" ]; then + # 删除文件后，恢复目录权限配置 + func_after_remove_01="foreach_stashmod \"restore_stash_mod\" \"${install_path}\" \"reverse\"" + func_after_remove_02="remove_stash_mod \"${install_path}\"" + fi + + del_cann_uninstall_package "${install_path}" "${package}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "$REMOVE_INSTALL_INFO" = "y" ]; then + del_ascend_install_info + fi + fi + + eval "${func_before_remove}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # del元素不会出现pkg_inner_softlink + foreach_filelist "filter_by_pkg_inner_softlink" "remove_install_pkg_inner_softlinks" "$install_type" "$install_path" "mkdir copy copy_entity move" "$tmp_filelist_path" "${feature_param}" "no" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_filelist "filter_by_blocks" "remove_install_softlink" "$install_type" "$install_path" "mkdir copy copy_entity move" "$tmp_filelist_path" "${feature_param}" "no" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_filelist "filter_by_blocks" "remove_install_files" "$install_type" "$install_path" "copy del move" "$tmp_filelist_path" "${feature_param}" "no" "concurrency" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_filelist "filter_by_blocks" "remove_install_files_recursive" "$install_type" "$install_path" "copy_entity" "$tmp_filelist_path" "${feature_param}" "no" "concurrency" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_filelist "filter_by_blocks" "remove_install_dirs" "$install_type" "$install_path" "mkdir" "$tmp_filelist_path" "${feature_param}" "reverse" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + + eval "${func_after_remove_01}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + eval "${func_after_remove_02}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + rm -f "$tmp_filelist_path" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 打印安装信息 +print_install_content() { + local install_type="$1" + local install_path="$2" + local operate_type="$3" + local filelist_path="$4" + local feature_param="$5" + local blocks + + get_blocks_info "blocks" "$install_path" + parse_filelist "$install_type" "$operate_type" "$filelist_path" "$feature_param" "filter_by_blocks" "$blocks" + + # bash与dash的echo有差异，需要使用系统的/bin/echo + /bin/echo "$filelist" + exit 0 +} + +# 包名是否存在于db.info +pkg_in_dbinfo() { + local install_path="$1" + local pkg="$2" + + local db_filepath="${install_path}/${PKG_DB_INFO_RELPATH}" + + if [ -z "${install_path}" ]; then + echo "false" + exit 1 + fi + + if [ -z "${pkg}" ]; then + echo "false" + exit 2 + fi + + if [ ! -f "${db_filepath}" ]; then + echo "false" + return 0 + fi + + awk -v input_pkg="${pkg}" ' + BEGIN { + FS = "|" + found = "false" + } + { + split($2, pkg_list, ",") + + for (i in pkg_list) { + if (input_pkg == pkg_list[i]) { + found = "true" + exit + } + } + } + END { + print found + }' "${db_filepath}" +} + +# 展开自定义选项 +expand_custom_options() { + local _outvar="$1" + local _custom_options="$2" + + eval "${_outvar}=\"$(echo "${_custom_options}" | tr "," " ")\"" +} + +# 调用子包自定义脚本 +package_custom_script() { + local script_name="$1" + local install_path="$2" + local version_dir="$3" + local custom_options="$4" + local ret install_options="" + + if [ ! -f "${curpath}/${script_name}" ]; then + return 0 + fi + + install_options="--install-path=${install_path}" + + if [ "${version_dir}" != "" ]; then + install_options="${install_options} --version-dir=${version_dir}" + fi + + expand_custom_options "custom_options" "${custom_options}" + + chmod u+x ${curpath}/${script_name} + ${curpath}/${script_name} ${install_options} ${custom_options} + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 调用子包自定义安装脚本 +package_custom_install() { + local package="$1" + local install_path="$2" + local version_dir="$3" + local custom_options="$4" + local ret + + package_custom_script "${package}_custom_install.sh" "${install_path}" "${version_dir}" "${custom_options}" + check_ret_error "$?" "Run ${package} custom install failed!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# 调用子包自定义卸载脚本 +package_custom_uninstall() { + local package="$1" + local install_path="$2" + local version_dir="$3" + local custom_options="$4" + local ret + + package_custom_script "${package}_custom_uninstall.sh" "${install_path}" "${version_dir}" "${custom_options}" + check_ret_error "$?" "Run ${package} custom uninstall failed!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# 版本安装流程函数 +version_install() { + local install_type="$1" + local install_path="$2" + local filelist_path="$3" + local package="$4" + local feature_param="$5" + local version_dir="$6" + local username="$7" + local usergroup="$8" + local setenv="$9" + local is_upgrade="${10}" + local docker_root="${11}" + local custom_options="${12}" + local is_simple="${13}" + local install_path_full="" ret + local package_real + + if [ "${version_dir}" != "" ]; then + install_path_full="${install_path}/${version_dir}" + else + install_path_full="${install_path}" + fi + + if [ "${is_simple}" = "y" ]; then + package_real="" + else + package_real="${package}" + fi + + # 创建目录 + do_create_dirs "mkdir" "${install_type}" "${install_path_full}" "${filelist_path}" "${package_real}" "${feature_param}" + if [ $? -ne 0 ]; then + log "ERROR" "failed to create folder." + return 1 + fi + + # 拷贝目录与文件 + do_copy_files "${install_type}" "${install_path_full}" "${filelist_path}" "${package_real}" "${feature_param}" + if [ $? -ne 0 ]; then + log "ERROR" "failed to copy files." + return 1 + fi + + if [ "${is_simple}" != "y" ]; then + # set env + add_setenv "${install_path_full}" "${package_real}" "${setenv}" "${username}" "${usergroup}" "false" "${docker_root}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 调用组件自定义安装流程 + package_custom_install "${package_real}" "${install_path}" "${version_dir}" "${custom_options}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + # 文件与目录赋权 + do_chmod_file_dir "${install_type}" "${install_path_full}" "${filelist_path}" "${feature_param}" "${package_real}" + if [ $? -ne 0 ]; then + log "ERROR" "failed to chown files." + return 1 + fi + + return 0 +} + +# 版本卸载流程函数 +version_uninstall() { + local install_type="$1" + local install_path="$2" + local filelist_path="$3" + local package="$4" + local feature_param="$5" + local version_dir="$6" + local username="$7" + local docker_root="$8" + local custom_options="$9" + local is_simple="${10}" + local install_path_full="" ret total_ret=0 + local package_real + + if [ "${version_dir}" != "" ]; then + install_path_full="${install_path}/${version_dir}" + else + install_path_full="${install_path}" + fi + + if [ "${is_simple}" = "y" ]; then + package_real="" + else + package_real="${package}" + fi + + # 恢复权限 + do_create_dirs "resetmod" "${install_type}" "${install_path_full}" "${filelist_path}" "${package_real}" "${feature_param}" + if [ $? -ne 0 ]; then + log "ERROR" "failed to resetmod chmod." + return 1 + fi + + if [ "${is_simple}" != "y" ]; then + # unset env + del_setenv "${install_path_full}" "${package_real}" "${username}" "${docker_root}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 调用组件自定义卸载流程，失败流程不中断 + package_custom_uninstall "${package_real}" "${install_path}" "${version_dir}" "${custom_options}" + ret="$?" && [ $ret -ne 0 ] && total_ret="$ret" + fi + + # 删除文件和目录 + do_remove "${install_type}" "${install_path_full}" "${filelist_path}" "${package_real}" "${feature_param}" + if [ $? -ne 0 ]; then + log "ERROR" "failed to remove files and dirs." + return 1 + fi + + return ${total_ret} +} + +# 版本安装函数 +do_install() { + local install_type="$1" + local install_path="$2" + local filelist_path="$3" + local package="$4" + local feature_param="$5" + local docker_root="$6" + local is_simple="$7" + local install_path_real ret + + check_param_not_empty "package" "need set package parameter in install!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + _package_to_log_pkg_name "LOG_PKG_NAME" "${package}" + check_ret_error "$?" "Set log package name failed in install!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + expand_version_file + + # 获取filelist.csv文件真实路径 + get_realpath "filelist_path" "${filelist_path}" + # 获取install_path真实路径 + get_realpath "install_path" "${install_path}" + + if [ "${docker_root}" != "" ]; then + if [ "${WITH_DOCKER_ROOT_PREFIX}" = "y" ]; then + install_path_real="${install_path}" + else + install_path_real="${docker_root}${install_path}" + fi + else + install_path_real="${install_path}" + fi + + if [ "${VERSION}" != "" ] && [ "${VERSION_DIR}" != "" ]; then + multi_version_install "${install_type}" "${install_path_real}" "${filelist_path}" "${package}" "${feature_param}" \ + "${VERSION}" "${VERSION_DIR}" "${USERNAME}" "${USERGROUP}" "${SETENV}" "${IS_UPGRADE}" "${docker_root}" "${CUSTOM_OPTIONS}" \ + "${INSTALL_FOR_ALL}" + ret="$?" && [ $ret -ne 0 ] && return $ret + else + version_install "${install_type}" "${install_path_real}" "${filelist_path}" "${package}" "${feature_param}" \ + "" "${USERNAME}" "${USERGROUP}" "${SETENV}" "${IS_UPGRADE}" "${docker_root}" "${CUSTOM_OPTIONS}" "${is_simple}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + return 0 +} + +# 版本卸载函数 +do_uninstall() { + local install_type="$1" + local install_path="$2" + local filelist_path="$3" + local package="$4" + local feature_param="$5" + local docker_root="$6" + local is_simple="$7" + local install_path_real ret + + check_param_not_empty "package" "need set package parameter in uninstall!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + _package_to_log_pkg_name "LOG_PKG_NAME" "${package}" + check_ret_error "$?" "Set log package name failed in uninstall!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + expand_version_file + + # 获取filelist.csv文件真实路径 + get_realpath "filelist_path" "${filelist_path}" + # 获取install_path真实路径 + get_realpath "install_path" "${install_path}" + + if [ "${docker_root}" != "" ]; then + if [ "${WITH_DOCKER_ROOT_PREFIX}" = "y" ]; then + install_path_real="${install_path}" + else + install_path_real="${docker_root}${install_path}" + fi + else + install_path_real="${install_path}" + fi + + if [ "${VERSION}" != "" ] && [ "${VERSION_DIR}" != "" ]; then + multi_version_uninstall "${install_type}" "${install_path_real}" "${filelist_path}" "${package}" "${feature_param}" \ + "${VERSION}" "${VERSION_DIR}" "${USERNAME}" "${USERGROUP}" "${docker_root}" "${CUSTOM_OPTIONS}" "${IS_RECREATE_SOFTLINK}" + ret="$?" && [ $ret -ne 0 ] && return $ret + else + version_uninstall "${install_type}" "${install_path_real}" "${filelist_path}" "${package}" "${feature_param}" \ + "" "${USERNAME}" "${docker_root}" "${CUSTOM_OPTIONS}" "${is_simple}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + return 0 +} + +help_info() { + echo "Usage:" + echo "" + echo " ----------------------------------------------------------------------------------------------------" + echo "" + echo " $0 {--install,--uninstall,--mkdir,--makedir,--copy,--chmoddir,--restoremod,--remove}" + echo "" + echo " --install : Install package." + echo " --uninstall : Uninstall package." + echo " --mkdir : (Deprecated) Use --install instead. Create the install directories." + echo " --makedir : (Deprecated) Use --install instead. Create the install directories." + echo " --copy : (Deprecated) Use --install instead. Copy install files." + echo " --chmoddir : (Deprecated) Use --install instead. Set installed file's right." + echo " --restoremod : (Deprecated) Use --uninstall instead. Restore directories right." + echo " --remove : (Deprecated) Use --uninstall instead. Remove installed files." + echo "" + echo " Group options:" + echo " --install-path= : Specify install path." + echo " --username= : Specify install username." + echo " --usergroup= : Specify install usergroup." + echo " --install_for_all : Install for all users." + echo " Usually specified in root user install scene." + echo " --docker-root= : Specify docker root path." + echo " Install path is not contained docker root path." + echo " --chip= : Specify chip." + echo " Install files which chip matched or chip is all." + echo " --feature= : Specify feature. Default is all." + echo " Install files which feature matched or feature is comm." + echo " --feature-exclude-all : Switch feature mode." + echo " Install files which feature matched." + echo "" + echo " $0 --install [options] " + echo "" + echo " Command options:" + echo " --package= : Specify package name." + echo " --version= : Specify version. --version-file is recommended." + echo " --version-dir= : Specify version directory. --version-file is recommended." + echo " --version-file= : Specify version info file path." + echo " Common script will parse version and version_dir from version_file." + echo " If you specified this option, --version and --version-dir options" + echo " no longer need to be specified." + echo " --custom-options= : Specify custom options for package custom script." + echo " --setenv : Add \"source setenv.\${shell_type}\" to rcfile." + echo " --set-cann-uninstall : Add uninstall command into cann_uninstall.sh." + echo " Always specified this in multi-version install." + echo "" + echo " $0 --uninstall [options] " + echo "" + echo " Command options:" + echo " --package= : Specify package name." + echo " --version= : Specify version. --version-file is recommended." + echo " --version-dir= : Specify version directory. --version-file is recommended." + echo " --version-file= : Specify version info file path." + echo " Common script will parse version and version_dir from version_file." + echo " If you specified this option, --version and --version-dir options" + echo " no longer need to be specified." + echo " --custom-options= : Specify custom options for package custom script." + echo "" + echo " ----------------------------------------------------------------------------------------------------" + echo "" + echo " $0 {--add-cann-uninstall,--del-cann-uninstall}" + echo "" + echo " --add-cann-uninstall : Add uninstall command in cann_uninstall.sh." + echo " --del-cann-uninstall : Del uninstall command in cann_uninstall.sh." + echo "" + echo " Group options:" + echo " --install-path= : Specify install path." + echo "" + echo " $0 --add-cann-uninstall [options] " + echo "" + echo " Command options:" + echo " --username= : Specify install username." + echo " --usergroup= : Specify install usergroup." + echo " --install_for_all : Install for all users." + echo " Usually specified in root user install scene." + echo "" + echo " $0 --del-cann-uninstall [options] " + echo "" + echo " ----------------------------------------------------------------------------------------------------" + echo "" + echo " $0 {--add-env-rc,--del-env-rc}" + echo "" + echo " --add-env-rc : Add \"source setenv.\${shell_type}\" to rcfile." + echo " --del-env-rc : Del \"source setenv.\${shell_type}\" in rcfile." + echo "" + echo " Group options:" + echo " --username= : Specify install username." + echo " --usergroup= : Specify install usergroup." + echo " --docker-root= : Specify docker root path." + echo " Install path is not contained docker root path." + echo "" + echo " $0 --add-env-rc [options] " + echo "" + echo " Command options:" + echo " --package= : Specify package name." + echo " --setenv : Add \"source setenv.\${shell_type}\" to rcfile." + echo "" + echo " $0 --del-env-rc [options] " + echo "" + echo " ----------------------------------------------------------------------------------------------------" + echo "" + echo " $0 --pkg-in-dbinfo " + echo "" + echo " Does package in ascend_package_db.info. Echo true or false." + echo "" + echo " Command options:" + echo " --package= : Specify package name." + echo "" + echo " ----------------------------------------------------------------------------------------------------" + echo "" + echo " $0 --help" + echo "" + echo " Print help messages." +} + +# env rc相关命令 +env_rc_commands() { + local install_path + local setenv_filepath + local shell_type + + case "${OPERATE_TYPE}" in + "add-env-rc") + install_path="$1" + setenv_filepath="$2" + shell_type="$3" + add_env_rc "${install_path}" "${setenv_filepath}" "${PACKAGE}" "${shell_type}" "${SETENV}" "${USERNAME}" "${USERGROUP}" "false" "${DOCKER_ROOT}" + exit 0 + ;; + "del-env-rc") + install_path="$1" + setenv_filepath="$2" + shell_type="$3" + del_env_rc "${install_path}" "${setenv_filepath}" "${shell_type}" "${USERNAME}" "${DOCKER_ROOT}" + exit 0 + ;; + esac +} + +# db.info相关命令 +dbinfo_commands() { + local install_path + + case "${OPERATE_TYPE}" in + "pkg-in-dbinfo") + install_path="$1" + pkg_in_dbinfo "${install_path}" "${PACKAGE}" + exit 0 + esac +} + +# spc相关命令 +spc_commands() { + local install_path="$1" + local filelist_path="$2" + local filelist_spc_path="$3" + local version_dir="$4" + local ret + + case "${OPERATE_TYPE}" in + "spc_install") + install_patch "${install_path}" "${filelist_path}" "${filelist_spc_path}" "${version_dir}" + exit $? + ;; + "spc_rollback") + rollback_patch "${install_path}" "${filelist_path}" "${version_dir}" + exit 0 + ;; + "spc_uninstall") + uninstall_patch "${install_path}" "${filelist_path}" "${version_dir}" + exit 0 + ;; + esac +} + +# cann_uninstall总卸载脚本相关命令 +cann_uninstall_commands() { + local script_dir="$1" + local ret + + case "${OPERATE_TYPE}" in + "add-cann-uninstall") + do_add_cann_uninstall "${INSTALL_PATH}" "${script_dir}" "${USERNAME}" "${USERGROUP}" "${INSTALL_FOR_ALL}" + ret="$?" && [ ${ret} -ne 0 ] && exit 1 + exit 0 + ;; + "del-cann-uninstall") + do_del_cann_uninstall "${INSTALL_PATH}" "${script_dir}" + ret="$?" && [ ${ret} -ne 0 ] && exit 1 + exit 0 + ;; + esac +} + +# 多版本创建相关命令 +multi_version_commands() { + case "${OPERATE_TYPE}" in + "notify_create_softlink") + notify_latest_manager_create_softlink "$INSTALL_PATH/$LATEST_DIR/var" "$PACKAGE" "$VERSION" "$VERSION_DIR" \ + "$INSTALL_FOR_ALL" "$DOCKER_ROOT" + ret="$?" && [ $ret -ne 0 ] && exit 1 + exit 0 + ;; + "notify_remove_softlink") + notify_latest_manager_remove_softlink "$INSTALL_PATH/$LATEST_DIR/var" "$PACKAGE" "$VERSION" "$VERSION_DIR" \ + "$INSTALL_FOR_ALL" "$DOCKER_ROOT" + ret="$?" && [ $ret -ne 0 ] && exit 1 + exit 0 + ;; + esac +} + +# 正式命令 +formal_commands() { + local install_type="$1" + local install_path="$2" + local filelist_path="$3" + local input_feature input_chip + local ret docker_root + local feature_param thread_num + + if [ "$FEATURE" != "all" ]; then + normalize_feature "input_feature" "$FEATURE" "${OPERATE_TYPE}" + else + if [ $# = 3 ]; then + input_feature="all" + else + normalize_feature "input_feature" "$4" "${OPERATE_TYPE}" + fi + fi + normalize_feature "input_chip" "$CHIP" "${OPERATE_TYPE}" + + pack_feature_param "feature_param" "${input_feature}" "${FEATURE_EXCLUDE_ALL}" "${input_chip}" + + # 移除docker_root右侧的/ + rstrip_path "docker_root" "${DOCKER_ROOT}" + + if [ "$PARALLEL_LIMIT" = "true" ]; then + get_thread_num "thread_num" + init_fifo "PARALLEL_FIFO" "$thread_num" + fi + + case "${OPERATE_TYPE}" in + "install") + do_install "${install_type}" "${install_path}" "${filelist_path}" "${PACKAGE}" "${feature_param}" "${docker_root}" "n" + ret="$?" && [ $ret -ne 0 ] && return $ret + ;; + "uninstall") + # 统一使用full模式卸载文件。两个包共用部分block（如opp和opp_kernel），以不同的模式安装时（如：run/full） + # full模式的包先卸载，run模式的包后卸载，确保能够完整卸载。 + do_uninstall "full" "${install_path}" "${filelist_path}" "${PACKAGE}" "${feature_param}" "${docker_root}" + ret="$?" && [ $ret -ne 0 ] && return $ret + ;; + "simple_install") + # 简化安装模式 + do_install "${install_type}" "${install_path}" "${filelist_path}" "${PACKAGE}" "${feature_param}" "" "y" + ret="$?" && [ $ret -ne 0 ] && return $ret + ;; + "simple_uninstall") + # 简化卸载模式 + do_uninstall "full" "${install_path}" "${filelist_path}" "${PACKAGE}" "${feature_param}" "" "y" + ret="$?" && [ $ret -ne 0 ] && return $ret + ;; + "mkdir") + do_create_dirs "all" "$1" "$2" "$3" "${PACKAGE}" "${feature_param}" + ret="$?" && [ $ret -ne 0 ] && return $ret + ;; + "makedir") + do_create_dirs "mkdir" "$1" "$2" "$3" "${PACKAGE}" "${feature_param}" + ret="$?" && [ $ret -ne 0 ] && return $ret + ;; + "copy") + do_copy_files "$1" "$2" "$3" "${PACKAGE}" "${feature_param}" + ret="$?" && [ $ret -ne 0 ] && return $ret + ;; + "chmoddir") + do_chmod_file_dir "$install_type" "$2" "$3" "${feature_param}" "${PACKAGE}" + ret="$?" && [ $ret -ne 0 ] && return $ret + ;; + "restoremod") + do_create_dirs "resetmod" "$1" "$2" "$3" "${PACKAGE}" "${feature_param}" + ret="$?" && [ $ret -ne 0 ] && return $ret + ;; + "remove") + do_remove "$install_type" "$2" "$filelist_path" "${PACKAGE}" "${feature_param}" + ret="$?" && [ $ret -ne 0 ] && return $ret + ;; + "print") + print_install_content "$1" "$2" "${PRINT_OPERATE_TYPE}" "$3" "${feature_param}" + ret="$?" && [ $ret -ne 0 ] && return $ret + ;; + esac + + return 0 +} + +# 正式命令并加锁 +formal_commands_with_lock() { + local install_path="$2" + local lockfile docker_root install_path_real ret + + if [ "${OPERATE_TYPE}" = "install" ] || [ "${OPERATE_TYPE}" = "uninstall" ]; then + if [ "${WITH_DOCKER_ROOT_PREFIX}" = "y" ]; then + install_path_real="${install_path}" + else + # 移除docker_root右侧的/ + rstrip_path "docker_root" "${DOCKER_ROOT}" + install_path_real="${docker_root}${install_path}" + fi + else + install_path_real="${install_path}" + fi + lockfile="${install_path_real}/ascend.lock" + + if [ ! -d "${install_path_real}" ]; then + # 兼容覆盖安装场景，包安装目录被删除情况 + mkdir -p "${install_path_real}" + # 安装目录设置合适的权限 + change_mod "${install_path_real}" "750" "${INSTALL_FOR_ALL}" + check_ret_error "$?" "Change mod ${install_path_real} failed!" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + # 使用文件描述符9作为并发锁 + ( + flock -n 9 + if [ $? -ne 0 ]; then + log "ERROR" "Get ${install_path_real} lockfile failed! There may be another process also installing in the directory." + exit 1 + fi + formal_commands "$@" + ) 9> "${lockfile}" + ret="$?" + rm -f "${lockfile}" + + if [ $ret -ne 0 ]; then + exit 1 + fi +} + +# 第一个参数为source，标记为source场景，不执行主流程 +[ "$1" = "source" ] && return 0 + +# 全局变量 +# change_mod_and_own_files与change_mod_and_own_dirs函数中会使用 +INSTALL_FOR_ALL="" +SETENV="" +SET_CANN_UNINSTALL="" +IS_UPGRADE="" +IS_RECREATE_SOFTLINK="" +WITH_DOCKER_ROOT_PREFIX="" +FEATURE_EXCLUDE_ALL="n" +REMOVE_INSTALL_INFO="n" # 卸载时移除ascend_install.info文件 +USE_SHARE_INFO="n" +CHIP="all" +FEATURE="all" +INCREMENT="n" # 增量安装 + +OPERATE_TYPE="" +PACKAGE="" +USERNAME="" +USERGROUP="" +VERSION="" +VERSION_DIR="" +VERSION_FILE="" +CUSTOM_OPTIONS="" +# 宿主机打docker包场景，配置docker文件系统根路径 +DOCKER_ROOT="" +# 安装路径 +INSTALL_PATH="" +# 输入的latest_dir +INPUT_LATEST_DIR="" + +while true; do + case "$1" in + --spc-install | -i) + OPERATE_TYPE="spc_install" + shift + ;; + --spc-rollback | -l) + OPERATE_TYPE="spc_rollback" + shift + ;; + --spc-uninstall | -u) + OPERATE_TYPE="spc_uninstall" + shift + ;; + --install) + OPERATE_TYPE="install" + shift + ;; + --uninstall) + OPERATE_TYPE="uninstall" + shift + ;; + --simple-install) + OPERATE_TYPE="simple_install" + shift + ;; + --simple-uninstall) + OPERATE_TYPE="simple_uninstall" + shift + ;; + --copy | -c) + OPERATE_TYPE="copy" + shift + ;; + --mkdir | -m) + OPERATE_TYPE="mkdir" + shift + ;; + --makedir | -d) + OPERATE_TYPE="makedir" + shift + ;; + --chmoddir | -o) + OPERATE_TYPE="chmoddir" + shift + ;; + --restoremod | -e) + OPERATE_TYPE="restoremod" + shift + ;; + --remove | -r) + OPERATE_TYPE="remove" + shift + ;; + --add-cann-uninstall) + OPERATE_TYPE="add-cann-uninstall" + shift + ;; + --del-cann-uninstall) + OPERATE_TYPE="del-cann-uninstall" + shift + ;; + --add-env-rc) + OPERATE_TYPE="add-env-rc" + shift + ;; + --del-env-rc) + OPERATE_TYPE="del-env-rc" + shift + ;; + --pkg-in-dbinfo) + OPERATE_TYPE="pkg-in-dbinfo" + shift + ;; + --add-pkg-blocks) + OPERATE_TYPE="add-pkg-blocks" + shift + ;; + --del-pkg-blocks) + OPERATE_TYPE="del-pkg-blocks" + shift + ;; + --create-latest-softlink) + OPERATE_TYPE="notify_create_softlink" + shift + ;; + --create-package-latest-softlink) + OPERATE_TYPE="notify_create_softlink" + shift + ;; + --remove-latest-softlink) + OPERATE_TYPE="notify_remove_softlink" + shift + ;; + --remove-package-latest-softlink) + OPERATE_TYPE="notify_remove_softlink" + shift + ;; + --print-blocks) + OPERATE_TYPE="print-blocks" + shift + ;; + --print=*) + OPERATE_TYPE="print" + PRINT_OPERATE_TYPE="$(echo "$1" | cut -d"=" -f2)" + shift + ;; + --print| -p) + OPERATE_TYPE="print" + PRINT_OPERATE_TYPE="copy mkdir move" + shift + ;; + --install-path=*) + INSTALL_PATH=$(echo "$1" | cut -d"=" -f2) + shift + ;; + --version=*) + VERSION=$(echo "$1" | cut -d"=" -f2) + shift + ;; + --version-dir=*) + VERSION_DIR=$(echo "$1" | cut -d"=" -f2) + shift + ;; + --version-file=*) + VERSION_FILE=$(echo "$1" | cut -d"=" -f2) + shift + ;; + --latest-dir=*) + INPUT_LATEST_DIR=$(echo "$1" | cut -d"=" -f2) + shift + ;; + --username=*) + USERNAME=$(echo "$1" | cut -d"=" -f2) + shift + ;; + --usergroup=*) + USERGROUP=$(echo "$1" | cut -d"=" -f2) + shift + ;; + --custom-options=*) + CUSTOM_OPTIONS=$(echo "$1" | cut -d"=" -f2-) + shift + ;; + --arch=*) + PKG_ARCH=$(echo "$1" | cut -d"=" -f2-) + shift + if [ "$PKG_ARCH" = "" ]; then + log "ERROR" "The --arch option should not be an empty string." + exit 1 + fi + ;; + --upgrade) + IS_UPGRADE="y" + shift + ;; + --recreate-softlink) + IS_RECREATE_SOFTLINK="y" + shift + ;; + --install_for_all) + INSTALL_FOR_ALL="y" + shift + ;; + --setenv) + SETENV="y" + shift + ;; + --set-cann-uninstall) + SET_CANN_UNINSTALL="y" + shift + ;; + --package=*) + PACKAGE=$(echo "$1" | cut -d"=" -f2) + shift + ;; + --docker-root=*) + DOCKER_ROOT=$(echo "$1" | cut -d"=" -f2) + shift + ;; + --chip=*) + CHIP="$(echo "$1" | cut -d"=" -f2)" + shift + ;; + --feature=*) + FEATURE="$(echo "$1" | cut -d"=" -f2)" + shift + ;; + --with-docker-root-prefix) + WITH_DOCKER_ROOT_PREFIX="y" + shift + ;; + --feature-exclude-all) + FEATURE_EXCLUDE_ALL="y" + shift + ;; + --remove-install-info) + REMOVE_INSTALL_INFO="y" + shift + ;; + --use-share-info) + USE_SHARE_INFO="y" + shift + ;; + --increment) + INCREMENT="y" + shift + ;; + -h | --help) + help_info + exit 0 + ;; + -*) + echo Unrecognized input options : "$1" + help_info + exit 1 + ;; + *) + break + ;; + esac +done + +set_global_vars +env_rc_commands "$@" +dbinfo_commands "$@" +spc_commands "$@" +cann_uninstall_commands "$@" +multi_version_commands "$@" + +if [ $# -lt 3 ]; then + log "ERROR" "It's too few input params: $*" + exit 1 +fi + +formal_commands_with_lock "$@" +exit $? diff --git a/csrc/ascend/scripts/package/common/sh/multi_version.inc b/csrc/ascend/scripts/package/common/sh/multi_version.inc new file mode 100644 index 000000000..fe531feb9 --- /dev/null +++ b/csrc/ascend/scripts/package/common/sh/multi_version.inc @@ -0,0 +1,305 @@ +#!/bin/sh +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +# 多版本函数库 +# 创建版本目录 +create_version_dir() { + local install_path="$1" + local version_dir="$2" + local username="$3" + local usergroup="$4" + local install_for_all="$5" + local ret + + if [ ! -d "${install_path}/${version_dir}" ]; then + make_dir "${install_path}/${version_dir}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + # 约束版本号目录权限 + change_mod "${install_path}/${version_dir}" "750" "${install_for_all}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + change_own "${install_path}/${version_dir}" "${username}:${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 检查版本包安装结果 +check_version_install() { + local install_path="$1" + local version_dir="$2" + local package="$3" + local version_info_path + + get_package_version_info "version_info_path" "$install_path" "$version_dir" "$package" + + if [ ! -f "$version_info_path" ]; then + comm_log "ERROR" "$version_info_path doesn't exist in check version install!" + return 1 + fi + return 0 +} + +# 安装latest管理器 +install_latest_manager() { + local var_path="$1" + shift 1 + sh latest_manager/install.sh --install-path="$var_path" "$@" +} + +# 卸载latest管理器 +uninstall_latest_manager_for_upgrade() { + local var_path="$1" + ${var_path}/manager/uninstall.sh --upgrade +} + +# 升级latest管理器 +upgrade_latest_manager() { + local var_path="$1" + local latest_version_info latest_manager_version local_manager_version + + latest_version_info="$var_path/manager/version.info" + if [ -f "$latest_version_info" ]; then + get_version latest_manager_version "$latest_version_info" + get_version local_manager_version "latest_manager/version.info" + + set_default "latest_manager_version" "$latest_manager_version" "0" + set_default "local_manager_version" "$local_manager_version" "0" + + if [ "$latest_manager_version" -lt "$local_manager_version" ]; then + uninstall_latest_manager_for_upgrade "$var_path" + install_latest_manager "$var_path" "--upgrade" + fi + else + install_latest_manager "$var_path" + fi +} + +get_install_for_all_param() { + local _outvar="$1" + local _install_for_all="$2" + local _result="" + + if [ "$_install_for_all" = "y" ]; then + _result="--install-for-all" + fi + + eval "${_outvar}=\"${_result}\"" +} + +get_docker_root_param() { + local _outvar="$1" + local _docker_root="$2" + local _result="" + + if [ "$_docker_root" != "" ]; then + _result="--docker-root \\\"$_docker_root\\\"" + fi + + eval "${_outvar}=\"${_result}\"" +} + +# 通知latest管理器 +notify_latest_manager() { + local var_path="$1" + local package="$2" + local version="$3" + local version_dir="$4" + local install_for_all="$5" + local docker_root="$6" + local ext_params="$7" + local operation="$8" + local ret package_dir install_for_all_param docker_root_param + + if [ ! -f "$var_path/manager.sh" ]; then + return 1 + fi + + if [ "${USE_SHARE_INFO}" == "y" ]; then + ext_params="$ext_params --use-share-info" + fi + + get_package_dir package_dir "$package" + get_install_for_all_param "install_for_all_param" "$install_for_all" + get_docker_root_param "docker_root_param" "$docker_root" + + eval "\"$var_path/manager.sh\"" --version "\"$version\"" --version-dir "\"$version_dir\"" \ + --package "\"$package\"" --package-dir "\"$package_dir\"" \ + "$install_for_all_param" "$docker_root_param" "$ext_params" "$operation" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 通知latest管理器安装完成 +notify_latest_manager_installed() { + local local_manager_version + local ext_params + + get_version local_manager_version "latest_manager/version.info" + set_default "local_manager_version" "$local_manager_version" "0" + + if [ "$INCREMENT" = "y" ]; then + ext_params="--serial $local_manager_version --increment" + else + ext_params="--serial $local_manager_version" + fi + + notify_latest_manager "$@" "$ext_params" "package_installed" +} + +# 通知latest管理器创建软链 +# 老版本全部卸载后，latest回滚到新版本时，会用老版本的install_common_parser.sh， +# 调用新版本的--create-package-latest-softlink，会走到这个流程。 +notify_latest_manager_create_softlink() { + notify_latest_manager "$@" "" "package_create_softlink" +} + +# 通知latest管理器删除软链 +notify_latest_manager_remove_softlink() { + notify_latest_manager "$@" "" "package_remove_softlink" +} + +# 通知latest管理器准备卸载 +notify_latest_manager_pre_uninstall() { + notify_latest_manager "$@" "" "package_pre_uninstall" +} + +# 通知latest管理器卸载完成 +notify_latest_manager_uninstalled() { + local is_recreate_softlink="$1" + local recreate_softlink="" + local ret + shift 1 + + if [ "$is_recreate_softlink" = "y" ]; then + recreate_softlink="--recreate-softlink" + fi + + notify_latest_manager "$@" "$recreate_softlink" "package_uninstalled" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 多版本安装流程 +multi_version_install() { + local install_type="$1" + local install_path="$2" + local filelist_path="$3" + local package="$4" + local feature_param="$5" + local version="$6" + local version_dir="$7" + local username="$8" + local usergroup="$9" + local setenv="${10}" + local is_upgrade="${11}" + local docker_root="${12}" + local custom_options="${13}" + local install_for_all="${14}" + local ret total_ret="0" pkg_running_version version_pair_arr last_version last_version_dir + + create_version_dir "${install_path}" "${version_dir}" "${username}" "${usergroup}" "${install_for_all}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + version_install "${install_type}" "${install_path}" "${filelist_path}" "${package}" "${feature_param}" \ + "${version_dir}" "${username}" "${usergroup}" "${setenv}" "${is_upgrade}" "${docker_root}" "${custom_options}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + check_version_install "$install_path" "$version_dir" "$package" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return ${total_ret} +} + +# 删除latest下空目录 +del_empty_dirs_in_latest() { + local install_type="$1" + local install_path="$2" + local latest_dir="$3" + local filelist_path="$4" + local feature_param="$5" + local ret + + create_stash_mod "${install_path}/${latest_dir}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_filelist "filter_common_dirs" "reset_mod_dirs_with_stash_mod" "${install_type}" "${install_path}/${latest_dir}" "mkdir" \ + "${filelist_path}" "${feature_param}" "no" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_filelist "filter_common_dirs" "remove_install_dirs" "${install_type}" "${install_path}/${latest_dir}" "mkdir" \ + "${filelist_path}" "${feature_param}" "reverse" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_stashmod "restore_stash_mod" "${install_path}/${latest_dir}" "reverse" + ret="$?" && [ $ret -ne 0 ] && return $ret + + remove_stash_mod "${install_path}/${latest_dir}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 多版本卸载流程 +multi_version_uninstall() { + local install_type="$1" + local install_path="$2" + local filelist_path="$3" + local package="$4" + local feature_param="$5" + local version="$6" + local version_dir="$7" + local username="$8" + local usergroup="$9" + local docker_root="${10}" + local custom_options="${11}" + local is_recreate_softlink="${12}" + local tmp_root tmp_filelist_path + local ret total_ret=0 install_path_full="" is_running="" is_upgrade + local running_packages is_final_running="false" + + check_param_not_empty "usergroup" "need set usergroup parameter in multi version uninstall!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + get_tmp_root "tmp_root" + tmp_filelist_path=$(mktemp "$tmp_root/filelist_XXXXXX" || exit 1) + cp -f "${filelist_path}" "${tmp_filelist_path}" + + if [ $? -ne 0 ]; then + log "ERROR" "cp -f ${filelist_path} ${tmp_filelist_path} failed!" + exit 1 + fi + + del_tmp_filelist="rm -f \"${tmp_filelist_path}\"" + + version_uninstall "${install_type}" "${install_path}" "${filelist_path}" "${package}" "${feature_param}" \ + "${version_dir}" "${username}" "${docker_root}" "${custom_options}" + if [ $? -ne 0 ]; then + eval "${del_tmp_filelist}" + return 1 + fi + + # 删除临时文件tmp_filelist_path + eval "${del_tmp_filelist}" + + # 删除版本空目录 + is_dir_empty "${install_path}/${version_dir}" + if [ $? -eq 0 ]; then + remove_dir_icp "${install_path}/${version_dir}" + ret="$?" && [ $ret -ne 0 ] && total_ret="1" + fi + + return ${total_ret} +} \ No newline at end of file diff --git a/csrc/ascend/scripts/package/common/sh/script_operator.inc b/csrc/ascend/scripts/package/common/sh/script_operator.inc new file mode 100644 index 000000000..6436ff9dc --- /dev/null +++ b/csrc/ascend/scripts/package/common/sh/script_operator.inc @@ -0,0 +1,469 @@ +#!/bin/sh +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +# 公共脚本操作库 +ENV_SHELL_TYPES="bash csh fish" + +# 检查参数。 +__check_param() { + local name="$1" + local value + + value="$(eval echo \${${name}})" + + if [ "${value}" = "" ]; then + log "ERROR" "need set ${name} parameter!" + exit 1 + fi +} + +# 设置用户家目录路径 +__set_userpath() { + local username="$1" + local docker_root="$2" + + if [ -z "${docker_root}" ] || [ "${docker_root}" = "/" ]; then + userpath="$(eval echo "~${username}")" + else + if [ "${username}" = "root" ]; then + userpath="${docker_root}/root" + else + userpath="${docker_root}/home/${username}" + fi + fi +} + +# 获取setenv.[shell]路径正则 +get_setenv_path_regex() { + local _outvar="$1" + local _package="$2" + local _shell_type="$3" + # setenv传入路径正则规则 + # 必须以/开头（绝对路径） + # 必须以/${package}/bin/setenv.${shell_type}结尾 + local _path_regex="\/\(.\+\/\)\?${_package}\/bin\/setenv.${_shell_type}" + + eval "${_outvar}=\"${_path_regex}\"" +} + +# 创建setenv文件 +__create_setenv_file() { + local file="$1" + local shell_type="$2" + local add_multi_version_param="$3" + local install_path="$4" + local ret + + if [ ! -f "${file}" ]; then + echo "#!/usr/bin/env ${shell_type}" > ${file} + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "${add_multi_version_param}" = "true" ]; then + if [ "${shell_type}" = "bash" ]; then + echo "export ASCEND_HOME_PATH=\"${install_path}\"" >> "${file}" + ret="$?" && [ $ret -ne 0 ] && return $ret + elif [ "${shell_type}" = "fish" ]; then + echo "set -gx ASCEND_HOME_PATH \"${install_path}\"" >> "${file}" + ret="$?" && [ $ret -ne 0 ] && return $ret + elif [ "${shell_type}" = "csh" ]; then + echo "setenv ASCEND_HOME_PATH \"${install_path}\"" >> "${file}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + fi + fi + return 0 +} + +# 是否存在ASCEND_HOME_PATH环境变量 +has_ascend_home_path_env() { + local file="$1" + [ ! -f "${file}" ] && return 0 + grep "\" "${file}" > /dev/null 2>&1 +} + +# 添加ASCEND_HOME_PATH环境变量 +add_ascend_home_path_env() { + local file="$1" + local shell_type="$2" + local install_path="$3" + + if [ "${shell_type}" = "bash" ]; then + sed -i "1aexport ASCEND_HOME_PATH=\"${install_path}\"" "${file}" + elif [ "${shell_type}" = "fish" ]; then + sed -i "1aset -gx ASCEND_HOME_PATH \"${install_path}\"" "${file}" + elif [ "${shell_type}" = "csh" ]; then + sed -i "1asetenv ASCEND_HOME_PATH \"${install_path}\"" "${file}" + fi +} + +# 删除rcfile中的source +__remove_path_regex() { + local path_regex="$1" + local rcfile="$2" + + if [ -f "${rcfile}" ]; then + sed -i "/source ${path_regex}\( \"multi_version\"\)\?$/d" "${rcfile}" + if [ $? -ne 0 ]; then + log "ERROR" "remove ${rcfile} source command failed!" + exit 1 + fi + fi +} + +# 获取setenv.${shell_type}文件路径 +get_setenv_filepath() { + local _outvar="$1" + local _install_path="$2" + local _shell_type="$3" + + eval "${_outvar}=\"${_install_path}/bin/setenv.${_shell_type}\"" +} + +# setenv.[shell]总脚本中添加一条语句 +add_setenv_cmd() { + local install_path="$1" + local setenv_filepath="$2" + local package="$3" + local shell_type="$4" + local username="$5" + local usergroup="$6" + local add_multi_version_param="$7" + local path_regex config_path multi_version_param ret + + get_setenv_path_regex "path_regex" "${package}" "${shell_type}" + + get_setenv_filepath "config_path" "$install_path" "$shell_type" + __create_setenv_file "${config_path}" "${shell_type}" "${add_multi_version_param}" "${install_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + change_own "${config_path}" "${username}:${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + change_mod "${config_path}" "${SETENV_WRITEABLE_MOD}" "" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __remove_path_regex "${path_regex}" "${config_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "${add_multi_version_param}" = "true" ]; then + multi_version_param=" \"multi_version\"" + fi + + if [ "${shell_type}" = "bash" ] || [ "${shell_type}" = "fish" ]; then + echo "source ${setenv_filepath}${multi_version_param}" >> "${config_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret + elif [ "${shell_type}" = "csh" ]; then + echo "set argv=(\"${setenv_filepath}\"${multi_version_param}); source ${setenv_filepath}" >> "${config_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + change_mod "${config_path}" "${SETENV_MOD}" "${INSTALL_FOR_ALL}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 修改路径的own +__chown_path() { + local path="$1" + local username="$2" + local usergroup="$3" + + if [ "${username}" != "" ] && [ "${usergroup}" != "" ]; then + chown -h "${username}:${usergroup}" "${path}" + if [ $? -ne 0 ]; then + log "ERROR" "${path} chown failed!" + exit 1 + fi + fi + + return 0 +} + +# 创建配置所在目录 +__create_config_dir() { + local dir="$1" + local username="$2" + local usergroup="$3" + local ret + + if [ ! -d "${dir}" ]; then + mkdir -p "${dir}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __chown_path "${dir}" "${username}" "${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + return 0 +} + +# 创建配置文件 +__create_config_file() { + local file="$1" + local username="$2" + local usergroup="$3" + local ret + + if [ ! -f "${file}" ]; then + touch "${file}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __chown_path "${file}" "${username}" "${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + return 0 +} + +# 为bash添加rc配置 +add_bash_env_rc() { + local userpath="$1" + local username="$2" + local usergroup="$3" + local path_regex="$4" + local setenv_filepath="$5" + local config_path="${userpath}/.bashrc" + local ret + + __create_config_dir "${userpath}" "${username}" "${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __create_config_file "${config_path}" "${username}" "${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __remove_path_regex "${path_regex}" "${config_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + echo "source ${setenv_filepath}" >> "${config_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret +} + +# 为fish添加rc配置 +add_fish_env_rc() { + local userpath="$1" + local username="$2" + local usergroup="$3" + local path_regex="$4" + local setenv_filepath="$5" + local config_path="${userpath}/.config/fish/config.fish" + local ret + + __create_config_dir "${userpath}" "${username}" "${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __create_config_dir "${userpath}/.config" "${username}" "${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __create_config_dir "${userpath}/.config/fish" "${username}" "${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __create_config_file "${config_path}" "${username}" "${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __remove_path_regex "${path_regex}" "${config_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + echo "source ${setenv_filepath}" >> "${config_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret +} + +# 为csh添加rc配置 +add_csh_env_rc() { + local userpath="$1" + local username="$2" + local usergroup="$3" + local path_regex="$4" + local setenv_filepath="$5" + local config_path="${userpath}/.cshrc" + local ret + + __create_config_dir "${userpath}" "${username}" "${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __create_config_file "${config_path}" "${username}" "${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __remove_path_regex "${path_regex}" "${config_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + echo "set argv=(\"${setenv_filepath}\"); source ${setenv_filepath}" >> "${config_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret +} + +# 添加source setenv.[shell]脚本到rc文件中 +add_env_rc() { + local install_path="$1" + local setenv_filepath="$2" + local package="$3" + local shell_type="$4" + local setenv="$5" + local username="$6" + local usergroup="$7" + local add_multi_version_param="$8" + local docker_root="$9" + local path_suffix="${package}/bin/setenv.${shell_type}" + local path_regex + local matched + local userpath + local config_path + local ret + + __check_param "package" + __check_param "username" + __check_param "usergroup" + + echo "${shell_type}" | grep -E "^(bash|fish|csh)$" > /dev/null + if [ $? -ne 0 ]; then + log "ERROR" "shell type ${shell_type} not support!" + exit 1 + fi + + get_setenv_path_regex "path_regex" "${package}" "${shell_type}" + + matched="$(echo "${setenv_filepath}" | sed -n "/^${path_regex}$/p")" + if [ -z "${matched}" ]; then + log "ERROR" "setenv filepath is illegal, should endswith ${path_suffix}" + exit 1 + fi + + if [ "${setenv}" = "y" ]; then + __set_userpath "${username}" "${docker_root}" + case "${shell_type}" in + bash) add_bash_env_rc "${userpath}" "${username}" "${usergroup}" "${path_regex}" "${setenv_filepath}" ;; + fish) add_fish_env_rc "${userpath}" "${username}" "${usergroup}" "${path_regex}" "${setenv_filepath}" ;; + csh) add_csh_env_rc "${userpath}" "${username}" "${usergroup}" "${path_regex}" "${setenv_filepath}" ;; + esac + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + # 如果bin目录不存在，则不处理 + [ ! -d "${install_path}/bin" ] && return 0 + + add_setenv_cmd "${install_path}" "${setenv_filepath}" "${package}" "${shell_type}" "${username}" "${usergroup}" "${add_multi_version_param}" +} + +# 删除setenv文件，如果已经没有setenv内容 +__remove_setenv_file_if_no_content() { + local file="$1" + local shell_type="$2" + local num + + if [ ! -f "${file}" ]; then + return 0 + fi + + num=$(grep "setenv.${shell_type}" ${file} | wc -l) + if [ ${num} -eq 0 ]; then + rm -f "${file}" > /dev/null 2>&1 + if [ $? -ne 0 ]; then + log "WARNING" "Delete file:${file} failed, please delete it by yourself." + fi + fi +} + +# 从rc文件中删除source setenv.[shell] +del_env_rc() { + local install_path="$1" + local setenv_filepath="$2" + local shell_type="$3" + local username="$4" + local docker_root="$5" + local path_regex + local userpath + local config_path + local oldmod + + __check_param "username" + + echo "${shell_type}" | grep -E "^(bash|fish|csh)$" > /dev/null + if [ $? -ne 0 ]; then + log "ERROR" "shell type ${shell_type} not support!" + exit 1 + fi + + __set_userpath "${username}" "${docker_root}" + + # 将路径中的/转换为\/ + path_to_regex "path_regex" "${setenv_filepath}" + + if [ "${shell_type}" = "bash" ]; then + __remove_path_regex "${path_regex}" "${userpath}/.bashrc" + ret="$?" && [ $ret -ne 0 ] && return $ret + elif [ "${shell_type}" = "fish" ]; then + __remove_path_regex "${path_regex}" "${userpath}/.config/fish/config.fish" + ret="$?" && [ $ret -ne 0 ] && return $ret + elif [ "${shell_type}" = "csh" ]; then + __remove_path_regex "${path_regex}" "${userpath}/.cshrc" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + get_setenv_filepath "config_path" "$install_path" "$shell_type" + + if [ ! -f "${config_path}" ]; then + # 如果文件不存在，则不处理 + return 0 + fi + + get_file_mod "oldmod" "${config_path}" + change_mod "${config_path}" "${SETENV_WRITEABLE_MOD}" "" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __remove_path_regex "${path_regex}" "${config_path}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + change_mod "${config_path}" "${oldmod}" "" + ret="$?" && [ $ret -ne 0 ] && return $ret + + __remove_setenv_file_if_no_content "${config_path}" "${shell_type}" +} + +# 移除路径中的docker_root +remove_path_docker_root() { + local _outvar="$1" + local _path="$2" + local _docker_root="$3" + local _path_rpdr="$(echo "${_path}" | sed "s/^.\{${#_docker_root}\}//")" + + eval "${_outvar}=\"${_path_rpdr}\"" +} + +# 获取脚本路径与脚本真实路径（docker_root） +get_shell_path_and_shell_path_real() { + local _outvar="$1" + local _install_path="$2" + local _package="$3" + local _shell_filename="$4" + local _docker_root="$5" + local _package_dirpath _shell_path _shell_path_real + + get_package_dirpath "_package_dirpath" "${_package}" + _shell_path_real="${_install_path}/${_package_dirpath}/bin/${_shell_filename}" + + if [ "${_docker_root}" != "" ]; then + # 移除脚本路径中的docker_root前缀 + remove_path_docker_root "_shell_path" "${_shell_path_real}" "${_docker_root}" + else + _shell_path="${_shell_path_real}" + fi + + eval "${_outvar}=\"${_shell_path} ${_shell_path_real}\"" +} + + +# 设置环境变量 +add_setenv() { + return 0 +} + +# 删除环境变量 +del_setenv() { + return 0 +} diff --git a/csrc/ascend/scripts/package/common/sh/version_cfg.inc b/csrc/ascend/scripts/package/common/sh/version_cfg.inc new file mode 100644 index 000000000..b6bc5f2bf --- /dev/null +++ b/csrc/ascend/scripts/package/common/sh/version_cfg.inc @@ -0,0 +1,1075 @@ +#!/bin/sh +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +# 版本状态配置文件处理 +# 版本状态配置文件格式为 +# # version: 1.0 +# runtime_running_version=[1.80:CANN-1.80] +# compiler_running_version=[1.80:CANN-1.80] +# toolkit_running_version=[1.80:CANN-1.80] +# fwkplugin_running_version=[1.80:CANN-1.80] +# runtime_installed_version=[1.79:CANN-1.79][1.80:CANN-1.80] +# compiler_installed_version=[1.79:CANN-1.79][1.80:CANN-1.80][1.81:CANN-1.81] +# toolkit_installed_version=[1.79:CANN-1.79][1.80:CANN-1.80][1.81:CANN-1.81] +# fwkplugin_installed_version=[1.79:CANN-1.79][1.80:CANN-1.80][1.81:CANN-1.81] + +# 子包会source本文件，本文件中调用的函数，都需要在文件内找到定义 + +# latest目录名 +LATEST_DIR="latest" + +# 配置版本号 +_CFG_VERSION="1.0" + +# 获取列表索引值 +__index_list_by_cut() { + local _list_ilbc="$1" + local _cnt_ilbc _value_ilbc + shift + + while true; do + if [ $# -eq 0 ]; then + return 0 + fi + _cnt_ilbc="${1}" && _cnt_ilbc=$((_cnt_ilbc+1)) + _value_ilbc="$(echo "${_list_ilbc}" | cut -d' ' -f"${_cnt_ilbc}")" + eval "${2}=\"${_value_ilbc}\"" + shift 2 + done + + return 0 +} + +# 获取版本状态配置路径 +get_version_cfg_path() { + local _outvar="$1" + local _dirpath="$2" + eval "${_outvar}=\"${_dirpath}/version.cfg\"" +} + +# 获取版本状态配置新路径 +get_version_cfg_new_path() { + local _outvar="$1" + local _dirpath="$2" + eval "${_outvar}=\"${_dirpath}/version.cfg.new\"" +} + +# 解析配置行 +parse_cfg_line() { + local _outvar="$1" + local _cfg_line="$2" + local _cfg_items_pcl + + _cfg_items_pcl=$(echo "${_cfg_line}" | awk ' + function join(items, len, sep, result, i) { + if (len >= 1) { + result = items[1] + for (i = 2; i <= len; i++) { + result = sprintf("%s%s%s", result, sep, items[i]) + } + return result + } + return "" + } + + /^[A-Za-z0-9_-]+_version=/ { + split($0, line_tokens, "=") + len_attr_tokens = split(line_tokens[1], attr_tokens, "_") + len_package_names = 0 + for (i = 1; i < len_attr_tokens - 1; i++) { + len_package_names++ + package_names[len_package_names] = attr_tokens[i] + } + package_name = join(package_names, len_package_names, "_") + attr_type = attr_tokens[len_attr_tokens-1] + value = line_tokens[2] + + printf("%s %s %s", package_name, attr_type, value) + } + ') + eval "${_outvar}=\"${_cfg_items_pcl}\"" +} + +# 解析配置版本号列表 +unpack_versions() { + local _outvar="$1" + local _value="$2" + + local _versions="$(echo "${_value}" | sed 's/\[//g')" + _versions="$(echo "${_versions}" | sed 's/\]/ /g')" + _versions="$(echo ${_versions})" # strip + eval "${_outvar}=\"${_versions}\"" +} + +# 打包配置版本号列表 +pack_versions() { + local _outvar="$1" + local _versions="$2" + + local _value="$(echo "${_versions}" | sed 's/ /\]\[/g')" + eval "${_outvar}=\"[${_value}]\"" +} + +# 解析版本序对 +unpack_version_pair() { + local _outvar="$1" + local _value="$2" + + eval "${_outvar}=\"$(echo "${_value}" | tr ':' ' ')\"" +} + +# 打包运行版本配置 +pack_running_version() { + local _outvar="$1" + local _package="$2" + local _version="$3" + local _version_dir="$4" + + local _version_str + pack_versions "_version_str" "${_version}:${_version_dir}" + eval "${_outvar}=\"${_package}_running_version=${_version_str}\"" +} + +# 打包升级版本配置 +pack_upgrade_version() { + local _outvar="$1" + local _package="$2" + local _version="$3" + local _version_dir="$4" + + local _version_str + pack_versions "_version_str" "${_version}:${_version_dir}" + eval "${_outvar}=\"${_package}_upgrade_version=${_version_str}\"" +} + +# 检查配置版本 +check_cfg_version() { + local filepath="$1" + local version + + if [ ! -f "${filepath}" ]; then + return 0 + fi + + version=$(sed -n '1p' "${filepath}" | cut -d" " -f3) + if [ "${version}" != "${_CFG_VERSION}" ]; then + return 1 + fi + return 0 +} + +# 输出配置版本错误 +error_check_cfg_version() { + local filepath="$1" + log "ERROR" "check ${filepath} version failed!" +} + +# 检查配置版本并输出报错日志 +check_cfg_version_with_log() { + local filepath="$1" + local ret + + check_cfg_version "${filepath}" + ret="$?" + if [ ${ret} -ne 0 ]; then + error_check_cfg_version "${filepath}" + return ${ret} + fi + return 0 +} + +# 添加配置版本注释 +add_cfg_version_comment() { + local filepath="$1" + sed -i "1i # version: ${_CFG_VERSION}" "${filepath}" +} + +# 处理移除子包版本配置行 +__process_unset_package_line() { + local line="$1" + local package="$2" + local target_type="$3" + local filepath_new="$4" + local cfg_items package_name attr_type + + parse_cfg_line "cfg_items" "${line}" + [ "${cfg_items}" = "" ] && return + + __index_list_by_cut "${cfg_items}" 0 "package_name" 1 "attr_type" + + if [ "${attr_type}" = "${target_type}" ]; then + [ "${package_name}" != "${package}" ] && echo "${line}" >> "${filepath_new}" + else + echo "${line}" >> "${filepath_new}" + fi +} + +# 处理文件更新和清理 +__finalize_version_file() { + local filepath="$1" + local filepath_new="$2" + local cleanup="$3" + + if [ -f "${filepath_new}" ]; then + add_cfg_version_comment "${filepath_new}" + mv -f "${filepath_new}" "${filepath}" + if [ $? -ne 0 ]; then + log "ERROR" "replace ${filepath} failed!" + eval "${cleanup}" + return 1 + fi + else + rm -f "${filepath}" + if [ $? -ne 0 ]; then + log "ERROR" "delete ${filepath} failed!" + eval "${cleanup}" + return 1 + fi + fi + + eval "${cleanup}" + local ret="$?" && [ $ret -ne 0 ] && return $ret +} + +# 移除子包版本 +unset_package_version() { + local dirpath="$1" + local package="$2" + local target_type="$3" + local filepath filepath_new line cfg_items package_name attr_type + local cleanup="" ret mod + + get_version_cfg_path "filepath" "${dirpath}" + get_version_cfg_new_path "filepath_new" "${dirpath}" + + # 版本状态配置文件不存在则正常退出 + [ ! -f "${filepath}" ] && return 0 + + check_cfg_version_with_log "${filepath}" || return 1 + + # 保存当前目录的权限 + get_file_mod "mod" "${dirpath}" + + # 添加目录写权限 + chmod u+w "${dirpath}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 恢复权限 + cleanup="chmod ${mod} \"${dirpath}\"" + + rm -f "${filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "delete ${filepath_new} failed!" + eval "${cleanup}" + return 1 + fi + + while read line; do + __process_unset_package_line "${line}" "${package}" "${target_type}" "${filepath_new}" + done < "${filepath}" + + __finalize_version_file "${filepath}" "${filepath_new}" "${cleanup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + return 0 +} + +# 获取子包版本 +get_package_version_version_cfg() { + local _outvar="$1" + local _dirpath="$2" + local _package="$3" + local _target_type="$4" + local _filepath _version_gpv + + get_version_cfg_path "_filepath" "${_dirpath}" + + if [ ! -f "${_filepath}" ]; then + eval "${_outvar}=\"\"" + return 0 + fi + + check_cfg_version "${_filepath}" + if [ $? -ne 0 ]; then + eval "${_outvar}=\"\"" + return 1 + fi + + _version_gpv="$(grep "^${_package}_${_target_type}_version=" "${_filepath}" | cut -d= -f2-)" + if [ "${_version_gpv}" != "" ]; then + unpack_versions "_version_gpv" "${_version_gpv}" + fi + + eval "${_outvar}=\"${_version_gpv}\"" + return 0 +} + +# 处理单个配置文件 +__process_cfg_file() { + local filepath="$1" + local package="$2" + local running_version="$3" + local filepath_new="$4" + local cfg_items package_name attr_type line + local is_set="false" + + touch "${filepath}" + while read line; do + parse_cfg_line "cfg_items" "${line}" + [ "${cfg_items}" = "" ] && continue + + __index_list_by_cut "${cfg_items}" 0 "package_name" 1 "attr_type" + + if [ "${attr_type}" = "running" ]; then + if [ "${package_name}" = "${package}" ]; then + if [ "${is_set}" = "false" ]; then + echo "${running_version}" >> "${filepath_new}" + is_set="true" + fi + else + echo "${line}" >> "${filepath_new}" + fi + elif [ "${attr_type}" = "upgrade" ]; then + if [ "${is_set}" = "false" ]; then + echo "${running_version}" >> "${filepath_new}" + is_set="true" + fi + echo "${line}" >> "${filepath_new}" + elif [ "${attr_type}" = "installed" ]; then + if [ "${is_set}" = "false" ]; then + echo "${running_version}" >> "${filepath_new}" + is_set="true" + fi + echo "${line}" >> "${filepath_new}" + else + echo "${line}" >> "${filepath_new}" + fi + done < "${filepath}" + + if [ "${is_set}" = "false" ]; then + echo "${running_version}" >> "${filepath_new}" + is_set="true" + fi +} + +# 设置运行子包版本 +set_running_package_version() { + local dirpath="$1" + local package="$2" + local version="$3" + local version_dir="$4" + local filepath filepath_new line cfg_items package_name attr_type running_version + local is_set="false" cleanup="" ret mod + + get_version_cfg_path "filepath" "${dirpath}" + get_version_cfg_new_path "filepath_new" "${dirpath}" + + check_cfg_version_with_log "${filepath}" + if [ $? -ne 0 ]; then + return 1 + fi + + pack_running_version "running_version" "${package}" "${version}" "${version_dir}" + + # 保存当前目录的权限 + get_file_mod "mod" "${dirpath}" + + # 添加目录写权限 + chmod u+w "${dirpath}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 恢复权限 + cleanup="chmod ${mod} \"${dirpath}\"" + + rm -f "${filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "delete ${filepath_new} failed!" + eval "${cleanup}" + return 1 + fi + + __process_cfg_file "${filepath}" "${package}" "${running_version}" "${filepath_new}" + + add_cfg_version_comment "${filepath_new}" + mv -f "${filepath_new}" "${filepath}" + if [ $? -ne 0 ]; then + log "ERROR" "replace ${filepath} failed!" + eval "${cleanup}" + return 1 + fi + + eval "${cleanup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 移除子包运行版本 +unset_running_package() { + local dirpath="$1" + local package="$2" + + unset_package_version "${dirpath}" "${package}" "running" +} + +# 获取子包运行版本 +get_running_package_version() { + local _outvar="$1" + local _dirpath="$2" + local _package="$3" + + get_package_version_version_cfg "${_outvar}" "${_dirpath}" "${_package}" "running" +} + +# 处理升级版本文件操作和清理 +__finalize_upgrade_file() { + local filepath="$1" filepath_new="$2" is_set="$3" upgrade_version="$4" cleanup="$5" + + if [ "${is_set}" = "false" ]; then + echo "${upgrade_version}" >> "${filepath_new}" + is_set="true" + fi + + add_cfg_version_comment "${filepath_new}" + mv -f "${filepath_new}" "${filepath}" + if [ $? -ne 0 ]; then + log "ERROR" "replace ${filepath} failed!" + eval "${cleanup}" + return 1 + fi + + eval "${cleanup}" + local ret="$?" && [ $ret -ne 0 ] && return $ret + return 0 +} + +# 处理设置升级版本配置行 +__process_set_upgrade_line() { + local line="$1" + local package="$2" + local upgrade_version="$3" + local filepath_new="$4" + local -n is_set_ref="$5" + local cfg_items package_name attr_type + + parse_cfg_line "cfg_items" "${line}" + [ "${cfg_items}" = "" ] && return + + __index_list_by_cut "${cfg_items}" 0 "package_name" 1 "attr_type" + + if [ "${attr_type}" = "upgrade" ]; then + if [ "${package_name}" = "${package}" ]; then + if [ "${is_set_ref}" = "false" ]; then + echo "${upgrade_version}" >> "${filepath_new}" + is_set_ref="true" + fi + else + echo "${line}" >> "${filepath_new}" + fi + elif [ "${attr_type}" = "installed" ]; then + if [ "${is_set_ref}" = "false" ]; then + echo "${upgrade_version}" >> "${filepath_new}" + is_set_ref="true" + fi + echo "${line}" >> "${filepath_new}" + else + echo "${line}" >> "${filepath_new}" + fi +} + +# 设置子包升级版本 +set_upgrade_package_version() { + local dirpath="$1" + local package="$2" + local version="$3" + local version_dir="$4" + local filepath filepath_new line cfg_items package_name attr_type upgrade_version + local is_set="false" cleanup="" ret mod + + get_version_cfg_path "filepath" "${dirpath}" + get_version_cfg_new_path "filepath_new" "${dirpath}" + + check_cfg_version_with_log "${filepath}" || return 1 + + pack_upgrade_version "upgrade_version" "${package}" "${version}" "${version_dir}" + + # 保存当前目录的权限 + get_file_mod "mod" "${dirpath}" + + # 添加目录写权限 + chmod u+w "${dirpath}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 恢复权限 + cleanup="chmod ${mod} \"${dirpath}\"" + + rm -f "${filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "delete ${filepath_new} failed!" + eval "${cleanup}" + return 1 + fi + + touch "${filepath}" + while read line; do + __process_set_upgrade_line "${line}" "${package}" "${upgrade_version}" "${filepath_new}" "is_set" + done < "${filepath}" + + __finalize_upgrade_file "${filepath}" "${filepath_new}" "${is_set}" "${upgrade_version}" "${cleanup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 移除子包升级版本 +unset_upgrade_package() { + local dirpath="$1" + local package="$2" + unset_package_version "${dirpath}" "${package}" "upgrade" +} + +# 获取子包升级版本 +get_upgrade_package_version() { + local _outvar="$1" + local _dirpath="$2" + local _package="$3" + + get_package_version_version_cfg "${_outvar}" "${_dirpath}" "${_package}" "upgrade" +} + +# 包版本是否为某状态 +is_package_version_status() { + local _outvar="$1" + local _latest_path="$2" + local _package="$3" + local _version="$4" + local _version_dir="$5" + local _target_type="$6" + local _pkg_status_version + local _ret _result_ipvs="" + + eval "${_outvar}=\"\"" + + get_package_version_version_cfg "_pkg_status_version" "${_latest_path}" "${_package}" "${_target_type}" + _ret="$?" && [ $_ret -ne 0 ] && return $_ret + + if [ "${_pkg_status_version}" = "${_version}:${_version_dir}" ]; then + _result_ipvs="true" + else + _result_ipvs="false" + fi + + eval "${_outvar}=\"${_result_ipvs}\"" +} + +# 包版本是否running +is_package_version_running() { + local _outvar="$1" + local _latest_path="$2" + local _package="$3" + local _version="$4" + local _version_dir="$5" + + is_package_version_status "${_outvar}" "${_latest_path}" "${_package}" "${_version}" "${_version_dir}" "running" +} + +# 包版本是否upgrade +is_package_version_upgrade() { + local _outvar="$1" + local _latest_path="$2" + local _package="$3" + local _version="$4" + local _version_dir="$5" + + is_package_version_status "${_outvar}" "${_latest_path}" "${_package}" "${_version}" "${_version_dir}" "upgrade" +} + +# 添加安装子包版本 +add_installed_package_version() { + local dirpath="$1" + local package="$2" + local version="$3" + local version_dir="$4" + local filepath filepath_new line cfg_items package_name attr_type value versions matched version_str + local is_set="false" + + get_version_cfg_path "filepath" "${dirpath}" + get_version_cfg_new_path "filepath_new" "${dirpath}" + + check_cfg_version_with_log "${filepath}" + if [ $? -ne 0 ]; then + return 1 + fi + + rm -f "${filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "delete ${filepath_new} failed!" + return 1 + fi + + pack_versions "version_str" "${version}:${version_dir}" + + touch "${filepath}" + while read line; do + parse_cfg_line "cfg_items" "${line}" + [ "${cfg_items}" = "" ] && continue + + __index_list_by_cut "${cfg_items}" 0 "package_name" 1 "attr_type" 2 "value" + + if [ "${attr_type}" = "installed" ] && [ "${package_name}" = "${package}" ]; then + unpack_versions "versions" "${value}" + __item_in_list "matched" "${version}:${version_dir}" ${versions} + if [ "${matched}" = "true" ]; then + echo "${line}" >> "${filepath_new}" + else + echo "${line}${version_str}" >> "${filepath_new}" + fi + is_set="true" + else + echo "${line}" >> "${filepath_new}" + fi + done < "${filepath}" + + if [ "${is_set}" = "false" ]; then + echo "${package}_installed_version=${version_str}" >> "${filepath_new}" + is_set="true" + fi + + add_cfg_version_comment "${filepath_new}" + mv -f "${filepath_new}" "${filepath}" + if [ $? -ne 0 ]; then + log "ERROR" "replace ${filepath} failed!" + return 1 + fi +} + +__process_install_cfg_line() { + local package="$1" + local version="$2" + local version_dir="$3" + local filepath_new="$4" + local line cfg_items package_name attr_type value versions len_versions + + while read line; do + parse_cfg_line "cfg_items" "${line}" + [ "${cfg_items}" = "" ] && continue + + __index_list_by_cut "${cfg_items}" 0 "package_name" 1 "attr_type" 2 "value" + + if [ "${attr_type}" = "installed" ] && [ "${package_name}" = "${package}" ]; then + unpack_versions "versions" "${value}" + versions="$(__remove_item_in_list "${version}:${version_dir}" "${versions}")" + __length_list "${versions}" "len_versions" + if [ ${len_versions} -gt 0 ]; then + pack_versions "value" "${versions}" + echo "${package}_installed_version=${value}" >> "${filepath_new}" + fi + else + echo "${line}" >> "${filepath_new}" + fi + done < "${filepath}" +} + +# 删除安装子包版本 +del_installed_package_version() { + local dirpath="$1" + local package="$2" + local version="$3" + local version_dir="$4" + local filepath filepath_new line cfg_items package_name attr_type value versions len_versions + + get_version_cfg_path "filepath" "${dirpath}" + get_version_cfg_new_path "filepath_new" "${dirpath}" + + # 版本状态配置文件不存在则正常退出 + [ ! -f "${filepath}" ] && return 0 + + check_cfg_version_with_log "${filepath}" || return 1 + + rm -f "${filepath_new}" + if [ $? -ne 0 ]; then + log "ERROR" "delete ${filepath_new} failed!" + return 1 + fi + + __process_del_installed_line "${package}" "${version}" "${version_dir}" "${filepath_new}" + + if [ -f "${filepath_new}" ]; then + add_cfg_version_comment "${filepath_new}" + mv -f "${filepath_new}" "${filepath}" + if [ $? -ne 0 ]; then + log "ERROR" "replace ${filepath} failed!" + return 1 + fi + else + rm -f "${filepath}" + if [ $? -ne 0 ]; then + log "ERROR" "delete ${filepath} failed!" + return 1 + fi + fi + return 0 +} + +# 获取子包版本 +get_installed_package_versions() { + local _outvar="$1" + local _dirpath="$2" + local _package="$3" + local _filepath _value _new_versions_gipv + + get_version_cfg_path "_filepath" "${_dirpath}" + + if [ ! -f "${_filepath}" ]; then + eval "${_outvar}=\"\"" + return 0 + fi + + check_cfg_version "${_filepath}" + if [ $? -ne 0 ]; then + eval "${_outvar}=\"\"" + return 1 + fi + + _value="$(grep "^${_package}_installed_version=" "${_filepath}" | cut -d= -f2-)" + unpack_versions "_new_versions_gipv" "${_value}" + __reverse_list "_new_versions_gipv" "${_new_versions_gipv}" + + eval "${_outvar}=\"${_new_versions_gipv}\"" +} + +# 获取子包最后一次的安装版本 +get_package_last_installed_version() { + local _outvar="$1" + local _install_path="$2" + local _package="$3" + local _versions_gpliv="" _version_gpliv="" _ret + + eval "${_outvar}=\"\"" + + get_installed_package_versions "_versions_gpliv" "${_install_path}/${LATEST_DIR}" "${_package}" + _ret="$?" && [ ${_ret} -ne 0 ] && return ${_ret} + + [ "${_versions_gpliv}" = "" ] && return 0 + + __index_list_by_cut "${_versions_gpliv}" 0 "_version_gpliv" + + eval "${_outvar}=\"${_version_gpliv}\"" +} + + +# 获取子包最后一次的安装版本目录 +# 如果不存在上次安装版本，则返回空字符串 +get_package_last_installed_version_dir() { + local _outvar="$1" + local _install_path="$2" + local _package="$3" + local _version_pair _version_pair_arr _version_dir_gplivd="" _ret + + eval "${_outvar}=\"\"" + + get_package_last_installed_version "_version_pair" "${_install_path}" "${_package}" + _ret="$?" && [ ${_ret} -ne 0 ] && return ${_ret} + + unpack_version_pair "_version_pair_arr" "${_version_pair}" + __index_list_by_cut "${_version_pair_arr}" 1 "_version_dir_gplivd" + + eval "${_outvar}=\"${_version_dir_gplivd}\"" +} + +# 获取子包最后一次安装的install.info路径 +# 如果不存在上次安装版本，则返回空字符串 +get_package_last_install_info() { + local _outvar="$1" + local _install_path="$2" + local _package="$3" + local _version_dir="" _install_info="" _ret + + eval "${_outvar}=\"\"" + + get_package_last_installed_version_dir "_version_dir" "${_install_path}" "${_package}" + _ret="$?" && [ ${_ret} -ne 0 ] && return ${_ret} + + [ "${_version_dir}" = "" ] && return 0 + + get_package_install_info "_install_info" "${_install_path}" "${_version_dir}" "${_package}" + + eval "${_outvar}=\"${_install_info}\"" +} + +get_package_version_dir() { + local _outvar="$1" + local _install_path="$2" + local _package="$3" + local _target_type="$4" + local _version_pair _version_pair_arr _version_dir_gprvd="" _ret + + eval "${_outvar}=\"\"" + + get_package_version_version_cfg "_version_pair" "${_install_path}/${LATEST_DIR}" "${_package}" "${_target_type}" + _ret="$?" && [ ${_ret} -ne 0 ] && return ${_ret} + + if [ "${_version_pair}" = "" ] && [ "${_target_type}" = "upgrade" ]; then + get_package_version_version_cfg "_version_pair" "${_install_path}/${LATEST_DIR}" "${_package}" "running" + _ret="$?" && [ ${_ret} -ne 0 ] && return ${_ret} + fi + + [ "${_version_pair}" = "" ] && return 0 + + unpack_version_pair "_version_pair_arr" "${_version_pair}" + __index_list_by_cut "${_version_pair_arr}" 1 "_version_dir_gprvd" + + eval "${_outvar}=\"${_version_dir_gprvd}\"" +} + +get_package_install_message() { + local _outvar="$1" + local _install_path="$2" + local _package="$3" + local _target_type="$4" + local _version_dir_gpim="" _install_info_gpim="" _ret + + eval "${_outvar}=\"\"" + + get_package_version_dir "_version_dir_gpim" "${_install_path}" "${_package}" "${_target_type}" + _ret="$?" && [ ${_ret} -ne 0 ] && return ${_ret} + + [ "${_version_dir_gpim}" = "" ] && return 0 + + get_package_install_info "_install_info_gpim" "${_install_path}" "${_version_dir_gpim}" "${_package}" + + eval "${_outvar}=\"${_install_info_gpim}\"" +} + +# 获取子包的运行版本目录 +# 如果不存在运行版本，则返回空字符串 +get_package_running_version_dir() { + local _outvar="$1" + local _install_path="$2" + local _package="$3" + + get_package_version_dir "${_outvar}" "${_install_path}" "${_package}" "running" +} + +# 获取子包运行版本的install.info路径 +# 如果不存在运行版本，则返回空字符串 +get_package_running_install_info() { + local _outvar="$1" + local _install_path="$2" + local _package="$3" + + get_package_install_message "${_outvar}" "${_install_path}" "${_package}" "running" +} + +# 获取子包的升级版本目录 +# 如果不存在升级版本，则返回运行版本目录 +# 如果也不存在运行版本，则返回空字符串 +get_package_upgrade_version_dir() { + local _outvar="$1" + local _install_path="$2" + local _package="$3" + + get_package_version_dir "${_outvar}" "${_install_path}" "${_package}" "upgrade" +} + +# 获取子包升级版本的install.info路径 +# 如果不存在升级版本，则返回运行版本的install.info路径 +# 如果也不存在运行版本，则返回空字符串 +get_package_upgrade_install_info() { + local _outvar="$1" + local _install_path="$2" + local _package="$3" + + get_package_install_message "${_outvar}" "${_install_path}" "${_package}" "upgrade" +} + +# 获取所有运行包 +get_running_packages() { + local _outvar="$1" + local _dirpath="$2" + local _filepath _line _cfg_items _package_name_grp _attr_type_grp _packages="" + + eval "${_outvar}=\"\"" + + get_version_cfg_path "_filepath" "${_dirpath}" + + if [ ! -f "${_filepath}" ]; then + return 0 + fi + + check_cfg_version "${_filepath}" + if [ $? -ne 0 ]; then + return 1 + fi + + while read _line; do + parse_cfg_line "_cfg_items" "${_line}" + [ "${_cfg_items}" = "" ] && continue + + __index_list_by_cut "${_cfg_items}" 0 "_package_name_grp" 1 "_attr_type_grp" + + if [ "${_attr_type_grp}" = "running" ]; then + if [ "${_packages}" = "" ]; then + _packages="${_package_name_grp}" + else + _packages="${_packages} ${_package_name_grp}" + fi + fi + done < "${_filepath}" + + eval "${_outvar}=\"${_packages}\"" +} + +# 根据版本目录获取版本号 +get_version_by_version_dir() { + local _outvar="$1" + local _dirpath="$2" + local _version_dir="$3" + local _filepath _line _cfg_items _attr_type_gvbvd _value_gvbvd + local _version_pairs _version_pair _version_pair_arr _version_gvbvd _version_dir_gvbvd + + eval "${_outvar}=\"\"" + + get_version_cfg_path "_filepath" "${_dirpath}" + if [ ! -f "${_filepath}" ]; then + return 1 + fi + + check_cfg_version "${_filepath}" + if [ $? -ne 0 ]; then + return 1 + fi + + while read _line; do + parse_cfg_line "_cfg_items" "${_line}" + [ "${_cfg_items}" = "" ] && continue + + __index_list_by_cut "${_cfg_items}" 1 "_attr_type_gvbvd" 2 "_value_gvbvd" + + if [ "${_attr_type_gvbvd}" = "installed" ]; then + unpack_versions "_version_pairs" "${_value_gvbvd}" + for _version_pair in ${_version_pairs}; do + unpack_version_pair "_version_pair_arr" "${_version_pair}" + __index_list_by_cut "${_version_pair_arr}" 0 "_version_gvbvd" 1 "_version_dir_gvbvd" + if [ "${_version_dir_gvbvd}" = "${_version_dir}" ]; then + eval "${_outvar}=\"${_version_gvbvd}\"" + return 0 + fi + done + fi + done < "${_filepath}" + + return 1 +} + +# 根据版本号与版本目录获取所有安装包 +get_installed_packages_by_version_version_dir() { + local _outvar="$1" + local _dirpath="$2" + local _version="$3" + local _version_dir="$4" + local _filepath _line _cfg_items + local _package_gipbvvd _attr_type_gipbvvd _value_gipbvvd _version_pairs _in_list + local _packages_gipbvvd="" + + eval "${_outvar}=\"\"" + + get_version_cfg_path "_filepath" "${_dirpath}" + if [ ! -f "${_filepath}" ]; then + return 1 + fi + + check_cfg_version "${_filepath}" + if [ $? -ne 0 ]; then + return 1 + fi + + while read _line; do + parse_cfg_line "_cfg_items" "${_line}" + [ "${_cfg_items}" = "" ] && continue + + __index_list_by_cut "${_cfg_items}" 0 "_package_gipbvvd" 1 "_attr_type_gipbvvd" 2 "_value_gipbvvd" + + if [ "${_attr_type_gipbvvd}" = "installed" ]; then + unpack_versions "_version_pairs" "${_value_gipbvvd}" + __item_in_list "_in_list" "${_version}:${_version_dir}" ${_version_pairs} + if [ "${_in_list}" = "true" ]; then + if [ "${_packages_gipbvvd}" = "" ]; then + _packages_gipbvvd="${_package_gipbvvd}" + else + _packages_gipbvvd="${_packages_gipbvvd} ${_package_gipbvvd}" + fi + fi + fi + done < "${_filepath}" + + eval "${_outvar}=\"${_packages_gipbvvd}\"" +} + +# 获取所有安装包 +get_installed_packages() { + local _outvar="$1" + local _dirpath="$2" + local _filepath _line _cfg_items _package_name_gip _attr_type_gip _packages="" + + eval "${_outvar}=\"\"" + + get_version_cfg_path "_filepath" "${_dirpath}" + + if [ ! -f "${_filepath}" ]; then + return 0 + fi + + check_cfg_version "${_filepath}" + if [ $? -ne 0 ]; then + return 1 + fi + + while read _line; do + parse_cfg_line "_cfg_items" "${_line}" + [ "${_cfg_items}" = "" ] && continue + + __index_list_by_cut "${_cfg_items}" 0 "_package_name_gip" 1 "_attr_type_gip" + + if [ "${_attr_type_gip}" = "installed" ]; then + if [ "${_packages}" = "" ]; then + _packages="${_package_name_gip}" + else + _packages="${_packages} ${_package_name_gip}" + fi + fi + done < "${_filepath}" + + eval "${_outvar}=\"${_packages}\"" +} + +# version.cfg文件是否存在 +# 返回值0为真，1为假 +version_cfg_exists() { + local dirpath="$1" + local filepath + + get_version_cfg_path "filepath" "$dirpath" + if [ -f "$filepath" ]; then + return 0 + fi + return 1 +} + +# 显示version.cfg文件 +show_version_cfg() { + local dirpath="$1" + local filepath + + get_version_cfg_path "filepath" "$dirpath" + if [ -f "$filepath" ]; then + cat "$filepath" + fi +} diff --git a/csrc/ascend/scripts/package/common/sh/version_compatiable.inc b/csrc/ascend/scripts/package/common/sh/version_compatiable.inc new file mode 100644 index 000000000..a13960f2e --- /dev/null +++ b/csrc/ascend/scripts/package/common/sh/version_compatiable.inc @@ -0,0 +1,174 @@ +#!/bin/sh +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +# 获取版本号，用于版本兼容性对比 +get_package_compat_version() { + local _outvar="$1" + local _version_info_path="$2" + local _result + + if [ ! -f "${_version_info_path}" ]; then + read -r "$_outvar" <= 1) { + result = items[1] + for (i = 2; i <= len; i++) { + result = sprintf("%s%s%s", result, sep, items[i]) + } + return result + } + return "" + } + + /^required_package_.+_version=/ { + split($0, line_tokens, "=") + len_require_tokens = split(line_tokens[1], require_tokens, "_") + len_package_names = 0 + for (i = 3; i < len_require_tokens; i++) { + len_package_names++ + package_names[len_package_names] = require_tokens[i] + } + package_name = join(package_names, len_package_names, "_") + print package_name + } + ' "${_version_info_path}" +} + +# 获取需求包信息。 +_get_required_package_info() { + local _outvar="$1" + local _version_info_path="$2" + local _pkg_name="$3" + local _required_pkg_name + local _required_value + + if [ ! -f "${_version_info_path}" ]; then + eval "${_outvar}=\"\"" + return 1 + fi + + _required_pkg_name="required_package_${_pkg_name}_version" + _required_value="$(grep "^${_required_pkg_name}=" "${_version_info_path}" | cut -d= -f2-)" + # _required_value取得的值带有双引号，eval时不可以在外侧再添加双引号 + + eval "${_outvar}=${_required_value}" +} + +# 检查版本与需求版本兼容性。 +# 兼容返回0，不兼容返回1 +_check_version_required() { + local version="$1" + local require="$2" + local script_dir="$3" + local src_pkg="$4" + local dst_pkg="$5" + local result + + result=$(awk -f "${script_dir}/check_version_required.awk" -v version="${version}" -v all_required="${require}") + if [ "${result}" = "T" ]; then + return 0 + fi + + echo "Version compatibility check failed, $src_pkg required $dst_pkg version $require, but $dst_pkg version is $version!" + return 1 +} + +# 检查版本兼容性。 +_check_version_compatiable() { + local package="$1" + local install_path="$2" + local script_dir="$3" + local version_info_path="$script_dir/../version.info" + local err_msgs package_formal self_version src_pkg src_pkg_formal dst_pkg dst_pkg_formal + local installed_version_info pkg_version_info_path + + get_formal_package_name "package_formal" "$package" + get_package_compat_version "self_version" "$version_info_path" + + if [ -d "$install_path/share/info" ]; then + err_msgs="$( + ls "$install_path/share/info" | grep -v "^${package}$" | while read src_pkg; do + get_formal_package_name "src_pkg_formal" "$package" + installed_version_info="$install_path/share/info/$src_pkg/version.info" + grep "^required_package_${package_formal}_version=" "$installed_version_info" | cut -d= -f2- | tr -d '"' | while read required; do + _check_version_required "$self_version" "$required" "$script_dir" "$src_pkg_formal" "$package_formal" + done + done + )" + if [ "$err_msgs" != "" ]; then + while read err_msg; do + comm_log "ERROR" "$err_msg" + done < "$var_path/manager_refs") + else + echo "1" > "$var_path/manager_refs" + fi +} + +# 减少manager引用计数 +dec_manager_refs() { + local var_path="$1" + local cnt + + if [ -f "$var_path/manager_refs" ]; then + xargs expr -1 + < "$var_path/manager_refs" | (read -r arg; echo "$arg" > "$var_path/manager_refs") + + get_manager_refs cnt "$var_path" + if [ "$cnt" -le 0 ]; then + rm -f "$var_path/manager_refs" + fi + fi +} + +# 获取是否需要install_for_all +get_install_for_all() { + local _outvar="$1" + local _result + + if [ ! -f "$MANAGER_INFO" ]; then + eval "${_outvar}=\"n\"" + return 0 + fi + + _result="$(grep "^install_for_all=" "$MANAGER_INFO" | cut -d= -f2-)" + if [ "$_result" = "" ]; then + eval "${_outvar}=\"n\"" + return 0 + fi + + eval "${_outvar}=\"${_result}\"" +} + +# 过滤配置文件参数 +filter_info_param() { + local name="$1" + local value="$2" + + awk -v name="$name" -v value="$value" -F= ' + BEGIN { + MATCHED = 0 + OFS = "=" + } + $1 == name { + MATCHED = 1 + print $1, value + } + $1 != name { + print $0 + } + END { + if (MATCHED == 0) { + print name, value + } + } + ' +} + +# 修改配置文件参数 +modify_info_param() { + local filepath="$1" + local name="$2" + local value="$3" + local content + + if [ ! -f "$filepath" ]; then + return 1 + fi + + content="$(cat "$filepath" | filter_info_param "$name" "$value")" + with_chmod "$filepath" "700" write_text "$content" "$filepath" +} + +# 包创建软链事件 +package_create_softlink() { + local var_path="$1" + local version="$2" + local version_dir="$3" + local package="$4" + local docker_root="$5" + local ret install_path pkg_running_version install_for_all + + install_path="$(dirname "$(dirname "$var_path")")" + + get_running_package_version "pkg_running_version" "$install_path/$LATEST_DIR" "$package" + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "$pkg_running_version" != "" ]; then + unpack_version_pair "version_pair_arr" "$pkg_running_version" + __index_list "$version_pair_arr" 0 "last_version" 1 "last_version_dir" + compat_del_package_softlink_in_latest "$install_path" "$package" "$last_version" "$last_version_dir" \ + "$LATEST_DIR" "$USERNAME" "$docker_root" + check_ret_error "$?" "delete $package softlink in latest failed in package create softlink!" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + get_install_for_all "install_for_all" + ret="$?" && [ $ret -ne 0 ] && return $ret + INSTALL_FOR_ALL="$install_for_all" + + do_create_package_softlink_to_latest "$install_path" "$package" "$version" "$version_dir" "$LATEST_DIR" \ + "$USERNAME" "$USERGROUP" "$install_for_all" "$docker_root" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 包安装事件 +package_installed() { + local var_path="$1" + local version="$2" + local version_dir="$3" + local package="$4" + local install_for_all="$5" + local docker_root="$6" + local ret install_path pkg_running_version version_pair_arr last_version last_version_dir + local config_install_for_all + + if [ "$INCREMENT" != "y" ]; then + inc_manager_refs "$var_path" + fi + install_path="$(dirname "$(dirname "$var_path")")" + + create_total_create_softlink_script "$var_path/manager" "$install_path" "$version_dir" "$USERNAME" "$USERGROUP" + check_ret_warning "$?" "$package create total create softlink script in $version_dir in failed!" + + create_total_remove_softlink_script "$var_path/manager" "$install_path" "$LATEST_DIR" "$USERNAME" "$USERGROUP" + check_ret_warning "$?" "Create total remove softlink script in $LATEST_DIR in package installed failed!" + + create_platform_ini "$install_path" "$LATEST_DIR" "$USERNAME" "$USERGROUP" + check_ret_warning "$?" "Create platform.ini in $LATEST_DIR in package installed failed!" + + get_install_for_all "config_install_for_all" + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "$install_for_all" != "$config_install_for_all" ]; then + ensure_file "$MANAGER_INFO" "440" "$USERNAME" "$USERGROUP" "$install_for_all" + ret="$?" && [ $ret -ne 0 ] && return $ret + + modify_info_param "$MANAGER_INFO" "install_for_all" "$install_for_all" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + package_create_softlink "$var_path" "$version" "$version_dir" "$package" "$docker_root" + ret="$?" && [ $ret -ne 0 ] && return $ret + + add_installed_package_version "$install_path/$LATEST_DIR" "$package" "$version" "$version_dir" + check_ret_error "$?" "Add $package installed version $version $version_dir in multi version install failed!" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 包删除软链事件 +package_remove_softlink() { + local var_path="$1" + local version="$2" + local version_dir="$3" + local package="$4" + local docker_root="$5" + local ret install_path is_running install_for_all + + install_path="$(dirname "$(dirname "$var_path")")" + + # version参数可能为空 + # 1. 新版本升级老版本场景，断开兼容性软链时，会使用老版本的install_common_parser.sh， + # 调用新版本的--remove-package-latest-softlink，但不会传version参数。 + # 2. 新版本先装，老版本后装的场景，latest目录下remove_latest_softlink.sh为老版本， + # 调用remove_latest_softlink.sh时，不会传version参数。 + if [ "$version" = "" ]; then + get_version_by_version_dir "version" "$install_path/$LATEST_DIR" "$version_dir" + fi + + get_install_for_all "install_for_all" + ret="$?" && [ $ret -ne 0 ] && return $ret + INSTALL_FOR_ALL="$install_for_all" + + is_package_version_running "is_running" "$install_path/$LATEST_DIR" "$package" "$version" "$version_dir" + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "$is_running" = "true" ]; then + do_del_package_softlink_in_latest "$install_path" "$package" "$version" "$version_dir" "$LATEST_DIR" \ + "$USERNAME" "$docker_root" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + return 0 +} + +# 包卸载前事件 +package_pre_uninstall() { + return 0 +} + +# 包卸载事件 +package_uninstalled() { + local var_path="$1" + local version="$2" + local version_dir="$3" + local package="$4" + local is_recreate_softlink="$5" + local docker_root="$6" + local install_path is_upgrade running_packages + + install_path="$(dirname "$(dirname "$var_path")")" + + if ! is_package_version_upgrade "is_upgrade" "$install_path/$LATEST_DIR" "$package" "$version" "$version_dir"; then + return 1 + fi + if [ "$is_upgrade" = "true" ]; then + # 移除version.cfg中的upgrade配置 + if ! unset_upgrade_package "$install_path/$LATEST_DIR" "$package"; then + comm_log "ERROR" "Unset $package upgrade version in package uninstalled failed!" + return 1 + fi + fi + + if ! del_installed_package_version "$install_path/$LATEST_DIR" "$package" "$version" "$version_dir"; then + comm_log "ERROR" "Del $package installed version $version $version_dir in package uninstalled failed!" + return 1 + fi + + # 删除版本目录下创建软链总脚本 + if ! del_total_create_softlink_script "$install_path" "$LATEST_DIR" "$version" "$version_dir"; then + comm_log "WARNING" "Del total create softlink script in $version_dir in package uninstalled failed!" + fi + + if ! version_cfg_exists "$install_path/$LATEST_DIR"; then + # 删除latest目录下删除软链总脚本 + del_total_remove_softlink_script "$install_path" "$LATEST_DIR" + check_ret_warning "$?" "Del total remove softlink script in $LATEST_DIR in package uninstalled failed!" + + # 删除latest目录下platform.ini + del_platform_ini "$install_path" "$LATEST_DIR" + check_ret_warning "$?" "Del platform.ini in $LATEST_DIR in package uninstalled failed!" + fi + + if [ "$is_recreate_softlink" = "y" ]; then + # 多版本卸载时检查版本兼容性 + if ! recreate_compatiable_softlink_in_multi_version_uninstall "$install_path" "$LATEST_DIR" "$package" \ + "$USERNAME" "$USERGROUP" "$docker_root"; then + comm_log "ERROR" "Recreate ${package} compatible softlink in package uninstalled failed!" + return 1 + fi + fi + + dec_manager_refs "$var_path" + if ! manager_refs_exists "$var_path"; then + rm -f "$MANAGER_INFO" + sh "$var_path/manager/uninstall.sh" + fi + return 0 +} + +# 设置版本软链接 +create_version_softlink() { + local var_path="$1" + local version_dir="$2" + local install_path + local total_ret="0" version packages package version_pair version_pair_arr + local ret del_version del_version_dir install_for_all + + install_path="$(dirname "$(dirname "$var_path")")" + + get_version_by_version_dir "version" "$install_path/$LATEST_DIR" "$version_dir" + ret="$?" && [ $ret -ne 0 ] && return $ret + + get_installed_packages_by_version_version_dir "packages" "$install_path/$LATEST_DIR" "$version" "$version_dir" + if [ "$packages" = "" ]; then + return 1 + fi + + get_install_for_all "install_for_all" + ret="$?" && [ $ret -ne 0 ] && return $ret + INSTALL_FOR_ALL="$install_for_all" + + for package in ${packages}; do + if ! get_running_package_version "version_pair" "$install_path/$LATEST_DIR" "$package"; then + total_ret=1 + continue + fi + + if [ "$version_pair" != "" ]; then + unpack_version_pair "version_pair_arr" "$version_pair" + __index_list "$version_pair_arr" 0 "del_version" 1 "del_version_dir" + + if ! do_del_package_softlink_in_latest "$install_path" "$package" "$del_version" "$del_version_dir" "$LATEST_DIR" \ + "$USERNAME" ""; then + total_ret=1 + fi + fi + + if ! do_create_package_softlink_to_latest "$install_path" "$package" "$version" "$version_dir" "$LATEST_DIR" \ + "$USERNAME" "$USERGROUP" "$install_for_all" ""; then + total_ret=1 + fi + done + + return $total_ret +} + +# 删除运行软链接 +remove_latest_softlink() { + local var_path="$1" + local install_path + local total_ret="0" ret running_packages package version_pair version_pair_arr version version_dir + local install_for_all + + install_path="$(dirname "$(dirname "$var_path")")" + + get_running_packages "running_packages" "$install_path/$LATEST_DIR" + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "${running_packages}" = "" ]; then + return 0 + fi + + get_install_for_all "install_for_all" + ret="$?" && [ $ret -ne 0 ] && return $ret + INSTALL_FOR_ALL="$install_for_all" + + for package in ${running_packages}; do + if ! get_running_package_version "version_pair" "$install_path/$LATEST_DIR" "$package"; then + total_ret=1 + continue + fi + + unpack_version_pair "version_pair_arr" "$version_pair" + __index_list "$version_pair_arr" 0 "version" 1 "version_dir" + + if ! do_del_package_softlink_in_latest "$install_path" "$package" "$version" "$version_dir" "$LATEST_DIR" \ + "$USERNAME" ""; then + total_ret=1 + fi + done + + return $total_ret +} + +# 创建子包软链到latest目录下 +do_create_package_softlink_to_latest() { + local install_path="$1" + local package="$2" + local version="$3" + local version_dir="$4" + local latest_dir="$5" + local username="$6" + local usergroup="$7" + local install_for_all="$8" + local docker_root="$9" + local ret total_ret="0" install_type feature_type feature_param filelist_path install_info_path + local chip_type + + check_param_not_empty "package" "need set package parameter in create package softlink to latest!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + _package_to_log_pkg_name "LOG_PKG_NAME" "${package}" + check_ret_error "$?" "Set log package name failed in create package softlink to latest!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + if [ ! -d "${install_path}/${latest_dir}" ]; then + make_dir_with_permission "${install_path}/${latest_dir}" "750" "${username}" "${usergroup}" "${install_for_all}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + get_package_filelist "filelist_path" "${install_path}" "${version_dir}" "${package}" + get_package_install_info "install_info_path" "${install_path}" "${version_dir}" "${package}" + + check_file_exists "${install_info_path}" "${install_info_path} doesn't exist in create package softlink to latest!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + get_package_install_type "install_type" "${install_info_path}" "${package}" + check_ret_error "$?" "Get install_type from ascend_install.info failed in create package softlink to latest!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + get_package_feature_type "feature_type" "${install_info_path}" "${package}" "install" + check_ret_error "$?" "Get feature_type from ascend_install.info failed in create package softlink to latest!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + get_package_chip_type "chip_type" "${install_info_path}" "${package}" "install" + check_ret_error "$?" "Get chip_type from ascend_install.info failed in create package softlink to latest!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + pack_feature_param "feature_param" "${feature_type}" "n" "${chip_type}" + + create_common_dirs_softlink_to_latest "${install_type}" "${install_path}" "${package}" "${version_dir}" "${latest_dir}" \ + "${filelist_path}" "${feature_param}" "${username}" "${usergroup}" "${install_for_all}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 创建公共目录到latest目录下的软链接 + if [ "${USE_SHARE_INFO}" != "y" ]; then + # 创建版本包到latest目录下的软链接 + create_package_dir_softlink_to_latest "${install_path}" "${version_dir}" "${latest_dir}" "${package}" \ + "${username}" "${usergroup}" "${install_for_all}" + check_ret_error "$?" "Create $package $version_dir $latest_dir package dirs softlink to latest failed!" + ret="$?" && [ $ret -ne 0 ] && return $ret + else + if [ ! -d "${install_path}/${latest_dir}/share" ]; then + ln -srf "${install_path}/${version_dir}/share" "${install_path}/${latest_dir}/share" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + fi + + # latest目录下公共脚本添加条目 + add_latest_common_script "${install_path}/${latest_dir}" "${package}" "${username}" "${usergroup}" "${docker_root}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 更新version.cfg中的running配置 + set_running_package_version "${install_path}/${latest_dir}" "${package}" "${version}" "${version_dir}" + check_ret_error "$?" "Set ${package} running version ${version} ${version_dir} in create package softlink to latest failed!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + # 更新version.cfg中的upgrade配置 + set_upgrade_package_version "${install_path}/${latest_dir}" "${package}" "${version}" "${version_dir}" + check_ret_error "$?" "Set ${package} upgrade version ${version} ${version_dir} in create package softlink to latest failed!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return ${total_ret} +} + + +# [兼容老版本]创建子包软链到latest目录下 +compat_create_package_softlink_to_latest() { + local install_path="$1" + local package="$2" + local version="$3" + local version_dir="$4" + local latest_dir="$5" + local username="$6" + local usergroup="$7" + local install_for_all="$8" + local docker_root="$9" + local ret installer_path + + get_package_install_common_parser "installer_path" "${install_path}" "${version_dir}" "${package}" + check_file_exists "${installer_path}" "${installer_path} doesn't exist in create package softlink to latest!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + install_options="--create-latest-softlink" + install_options="${install_options} --install-path=\"${install_path}\" --package=\"${package}\" --version=\"${version}\" --version-dir=\"${version_dir}\"" + install_options="${install_options} --latest-dir=\"${latest_dir}\" --username=\"${username}\" --usergroup=\"${usergroup}\"" + if [ "${install_for_all}" = "y" ]; then + install_options="${install_options} --install_for_all" + fi + if [ "${docker_root}" != "" ]; then + install_options="${install_options} --docker-root=\"${docker_root}\"" + fi + + eval sh "${installer_path}" "${install_options}" + check_ret_error "$?" "Create ${version_dir} ${package} softlink to latest failed!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# 删除latest目录下子包的软链接 +do_del_package_softlink_in_latest() { + local install_path="$1" + local package="$2" + local version="$3" + local version_dir="$4" + local latest_dir="$5" + local username="$6" + local docker_root="$7" + local ret install_type feature_type feature_param package_dirpath filelist_path install_info_path + local chip_type + + _package_to_log_pkg_name "LOG_PKG_NAME" "${package}" + check_ret_error "$?" "Set log package name failed in del package softlink in latest!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + if [ ! -d "${install_path}/${latest_dir}" ]; then + comm_log "ERROR" "${install_path}/${latest_dir} doesn't exist in del package softlink in latest!" + return 1 + fi + + get_package_dirpath "package_dirpath" "$package" + filelist_path="$install_path/$version_dir/$package_dirpath/script/filelist.csv" + install_info_path="$install_path/$version_dir/$package_dirpath/ascend_install.info" + + check_file_exists "${install_info_path}" "${install_info_path} doesn't exist in del package softlink in latest!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + # 卸载时install_type强制为full + install_type="full" + + get_package_feature_type "feature_type" "${install_info_path}" "${package}" "uninstall" + check_ret_error "$?" "Get feature_type from ascend_install.info failed in del package softlink in latest!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + get_package_chip_type "chip_type" "${install_info_path}" "${package}" "uninstall" + check_ret_error "$?" "Get chip_type from ascend_install.info failed in del package softlink to latest!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + pack_feature_param "feature_param" "${feature_type}" "n" "${chip_type}" + + del_latest_common_script "${install_path}/${latest_dir}" "${package}" "${username}" "${docker_root}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 删除公共目录到latest目录下的软链接 + del_common_dirs_softlink_from_latest "${install_type}" "${install_path}" "${package}" "${version_dir}" "${latest_dir}" \ + "${filelist_path}" "${feature_param}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # latest目录下公共脚本删除条目 + if [ "${USE_SHARE_INFO}" != "y" ]; then + # 删除版本包到latest目录下的软链接 + del_package_dir_softlink_in_latest "${install_path}" "${latest_dir}" "${package}" + check_ret_error "$?" "Delete $package package dir softlink in latest in del package softlink in latest failed!" + ret="$?" && [ $ret -ne 0 ] && return $ret + else + if [ -L "${install_path}/${latest_dir}/share" ]; then + rm -f "${install_path}/${latest_dir}/share" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + fi + + # 删除latest下tools空目录。toolkit包创建latest软链时，latest下存在软链tools/simulator -> ${arch}-linux/simulator + # 通过filelist.csv文件，无法删除tools目录 + # 解决toolkit包安装卸载latest目录下残留tools目录问题 + remove_dir_if_empty "${install_path}/${latest_dir}/tools" + + # 移除version.cfg中的running配置 + unset_running_package "${install_path}/${latest_dir}" "${package}" + check_ret_error "$?" "Unset ${package} running version in del package softlink in latest failed!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# [兼容老版本]删除latest目录下子包的软链接 +compat_del_package_softlink_in_latest() { + local install_path="$1" + local package="$2" + local version="$3" + local version_dir="$4" + local latest_dir="$5" + local username="$6" + local docker_root="$7" + local ret installer_path + + get_package_install_common_parser "installer_path" "${install_path}" "${version_dir}" "${package}" + + check_file_exists "${installer_path}" "${installer_path} doesn't exist in del package softlink in latest!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + # 更新version.cfg中的upgrade配置 + set_upgrade_package_version "${install_path}/${latest_dir}" "${package}" "${version}" "${version_dir}" + check_ret_error "$?" "Set ${package} upgrade version ${version} ${version_dir} in create package softlink to latest failed!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + install_options="--remove-latest-softlink" + install_options="${install_options} --install-path=\"${install_path}\" --package=\"${package}\" --version-dir=\"${version_dir}\"" + install_options="${install_options} --latest-dir=\"${latest_dir}\" --username=\"${username}\"" + if [ "${docker_root}" != "" ]; then + install_options="${install_options} --docker-root=\"${docker_root}\"" + fi + + eval sh "${installer_path}" "${install_options}" + check_ret_error "$?" "Remove ${version_dir} ${package} softlink in latest failed!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# 获取包安装类型 +get_package_install_type() { + local _outvar="$1" + local _install_info="$2" + local _package="$3" + local _install_type_gpit="" + + eval "${_outvar}=\"\"" + + if [ ! -f "${_install_info}" ]; then + return 1 + fi + + _install_type_gpit="$(grep -i "^\(${_package}_\)\?install_type=" "${_install_info}" | cut -d"=" -f2-)" + if [ "${_install_type_gpit}" = "" ]; then + return 1 + fi + + eval "${_outvar}=\"${_install_type_gpit}\"" +} + +# 获取包特性参数 +get_package_feature_type() { + local _outvar="$1" + local _install_info="$2" + local _package="$3" + local _operation="$4" + local _feature_type_gpft="" + + eval "${_outvar}=\"\"" + + if [ ! -f "${_install_info}" ]; then + return 1 + fi + + if [ "${_package}" = "opp" ]; then + _feature_type_gpft="$(grep -i "^Opp_Install_Feature=" "${_install_info}" | cut -d"=" -f2-)" + else + _feature_type_gpft="$(grep -i "^\(${_package}_\)\?feature_type=" "${_install_info}" | cut -d"=" -f2-)" + fi + if [ "${_feature_type_gpft}" = "" ]; then + _feature_type_gpft="all" + fi + + normalize_feature "_feature_type_gpft" "${_feature_type_gpft}" "${_operation}" + + eval "${_outvar}=\"${_feature_type_gpft}\"" +} + +# 获取包芯片参数 +get_package_chip_type() { + local _outvar="$1" + local _install_info="$2" + local _package="$3" + local _operation="$4" + local _cihp_gpct="" + + eval "${_outvar}=\"\"" + + if [ ! -f "${_install_info}" ]; then + return 1 + fi + + if [ "${_package}" = "opp" ]; then + _cihp_gpct="$(grep -i "^Opp_Install_Chip=" "${_install_info}" | cut -d"=" -f2-)" + else + _cihp_gpct="$(grep -i "^\(${_package}_\)\?chip_type=" "${_install_info}" | cut -d"=" -f2-)" + fi + if [ "${_cihp_gpct}" = "" ]; then + _cihp_gpct="all" + fi + + normalize_feature "_cihp_gpct" "${_cihp_gpct}" "${_operation}" + + eval "${_outvar}=\"${_cihp_gpct}\"" +} + +# 创建latest软链接 +create_latest_softlink() { + local latest_dirpath="$1" + local line="$2" + local version_dirpath="$4" + local target latest_filepath version_filepath softlink + local ret + + __index_list "${line}" 1 "target" 4 "softlink" + + __set_abs_path "${latest_dirpath}" "${target}" "latest_filepath" + __set_abs_path "${version_dirpath}" "${target}" "version_filepath" + + if [ ! -e "$version_filepath" ]; then + # 如果源路径不存在，则跳过软链创建 + return 0 + fi + + if [ -d "${latest_filepath}" ]; then + remove_dir_icp "${latest_filepath}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + create_softlink_icp "-r" "${version_filepath}" "${latest_filepath}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "${softlink}" != "NA" ]; then + create_softlink_by_install_path "${latest_dirpath}" "${target}" "${softlink}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + + return 0 +} + +# 删除latest软链接 +del_latest_softlink() { + local latest_dirpath="$1" + local line="$2" + local target latest_filepath softlink + local ret + + __index_list "${line}" 1 "target" 4 "softlink" + + __set_abs_path "${latest_dirpath}" "${target}" "latest_filepath" + + if [ "${softlink}" != "NA" ]; then + remove_softlinks "${install_path}" "${softlink}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + + remove_softlink_icp "${latest_filepath}" "NA" + ret="$?" && [ $ret -ne 0 ] && return $ret + return 0 +} + +# 迁移latest/conf实体目录下的文件到版本conf目录下 +migrate_conf_files_from_latest_to_version() { + local install_type="$1" + local install_path="$2" + local version_dir="$3" + local latest_dir="$4" + local filelist_path="$5" + local feature_param="$6" + local latest_conf_path="${install_path}/${latest_dir}/conf" + local version_conf_path="${install_path}/${version_dir}/conf" + local ret + + if [ -d "${latest_conf_path}" ] && [ ! -L "${latest_conf_path}" ] && [ -d "${version_conf_path}" ]; then + foreach_filelist "filter_common_dirs" "reset_mod_dirs" "${install_type}" "${install_path}/${version_dir}" "mkdir" \ + "${filelist_path}" "${feature_param}" "no" "normal" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + is_dir_empty "${latest_conf_path}" + # 如果目录不为空，则迁移数据 + if [ $? -ne 0 ]; then + mv -f "${latest_conf_path}"/* "${version_conf_path}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + + remove_dir_icp "${latest_conf_path}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + foreach_filelist "filter_common_dirs" "change_mod_and_own_dirs" "${install_type}" "${install_path}/${version_dir}" "mkdir" \ + "${filelist_path}" "${feature_param}" "reverse" "normal" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + + return 0 +} + +# 按块创建公共目录软链接到latest目录下 +create_common_dirs_softlink_cross_version() { + local install_path="$1" + local package="$2" + local version_dir="$3" + local latest_dir="$4" + local blocks="$5" + local filelist_path install_info_path install_type feature_type chip_type feature_param + + get_package_filelist "filelist_path" "$install_path" "$version_dir" "$package" + get_package_install_info "install_info_path" "$install_path" "$version_dir" "$package" + + check_file_exists "$install_info_path" "$install_info_path doesn't exist in create common dirs softlink cross version!" + ret="$?" && [ $ret -ne 0 ] && return $ret + + get_package_install_type "install_type" "$install_info_path" "$package" + check_ret_error "$?" "Get install_type from ascend_install.info failed in create common dirs softlink cross version!" + ret="$?" && [ $ret -ne 0 ] && return $ret + + get_package_feature_type "feature_type" "$install_info_path" "$package" "install" + get_package_chip_type "chip_type" "$install_info_path" "$package" "install" + pack_feature_param "feature_param" "$feature_type" "n" "$chip_type" + + create_common_dirs_softlink_by_blocks "$install_type" "$install_path" "$package" "$version_dir" "$latest_dir" \ + "$filelist_path" "$feature_param" "$blocks" "" + check_ret_error "$?" "Create common dirs softlink by blocks failed in create common dirs softlink cross version!" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 按块创建公共目录软链接到latest目录下 +create_common_dirs_softlink_by_blocks() { + local install_type="$1" + local install_path="$2" + local package="$3" + local version_dir="$4" + local latest_dir="$5" + local filelist_path="$6" + local feature_param="$7" + local blocks="$8" + local custom_create_softlink="$9" + local ret package_dirpath filelist_path + + foreach_filelist_v2 "reset_mod_dirs" "${install_type}" "${install_path}/${latest_dir}" "mkdir" \ + "${filelist_path}" "${feature_param}" "filter_common_dirs,filter_blocks" "$blocks" "no" "normal" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + foreach_filelist_v2 "create_dirs" "${install_type}" "${install_path}/${latest_dir}" "mkdir" \ + "${filelist_path}" "${feature_param}" "filter_common_dirs,filter_blocks" "$blocks" "no" "normal" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + foreach_filelist_v2 "create_latest_softlink" "${install_type}" "${install_path}/${latest_dir}" "copy copy_entity move" \ + "${filelist_path}" "${feature_param}" "filter_common_dirs,filter_blocks" "$blocks" "no" "normal" "${install_path}/${version_dir}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + # 调用脚本创建自定义内容软链接到latest目录下 + if [ "$custom_create_softlink" != "" ] && [ -f "$custom_create_softlink" ]; then + chmod u+x "${custom_create_softlink}" + "${custom_create_softlink}" --install-path="${install_path}" --version-dir="${version_dir}" --latest-dir="${latest_dir}" + check_ret_error "$?" "Run ${package} custom create ${version_dir} softlink failed!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + + foreach_filelist_v2 "change_mod_and_own_dirs" "${install_type}" "${install_path}/${latest_dir}" "mkdir" \ + "${filelist_path}" "${feature_param}" "filter_common_dirs,filter_blocks" "$blocks" "reverse" "normal" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# 创建公共目录软链接到latest目录下 +create_common_dirs_softlink_to_latest() { + local install_type="$1" + local install_path="$2" + local package="$3" + local version_dir="$4" + local latest_dir="$5" + local filelist_path="$6" + local feature_param="$7" + local username="$8" + local usergroup="$9" + local install_for_all="${10}" + local db_info_path="$install_path/$latest_dir/var/ascend_package_db.info" + local ret package_dirpath custom_create_softlink db_info + + get_package_dirpath "package_dirpath" "$package" + + custom_create_softlink="${install_path}/${version_dir}/${package_dirpath}/script/${package}_custom_create_softlink.sh" + + migrate_conf_files_from_latest_to_version "${install_type}" "${install_path}" "${version_dir}" "${latest_dir}" \ + "${filelist_path}" "${feature_param}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + ensure_file "$db_info_path" "440" "$username" "$usergroup" "$install_for_all" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + db_info="$( + ( + cat "$db_info_path" + printf "\n" # 防止db.info缺失末尾换行符 + all_common_dirs_blocks_in_filelist "$install_type" "$filelist_path" "$feature_param" \ + | blocks_to_db_item "$package" "$version_dir" + ) | remove_blank_line | sort_1 "|" | fold_2 "|" "" + )" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + # 注意换行符 + ( + cat "$db_info_path" | remove_blank_line | remain_db_last_item + printf "%s\n" "$db_info" | remove_blank_line | remain_db_last_item + ) | sort_1 | fold_3_keep_2 "" " " | show_diff_3_4 | show_min_nf "4" | select_fields_3_2_1 | sort_1 | fold_3_keep_2 "" " " \ + | ( + total_ret=0 + while read -r tmp_version_dir tmp_package tmp_blocks; do + del_common_dirs_softlink_cross_version "$install_path" "$tmp_package" "$tmp_version_dir" "$latest_dir" \ + "EngineeringCommon $tmp_blocks" + ret="$?" && [ $ret -ne 0 ] && total_ret="1" + done + exit $total_ret + ) + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + # 安装时总是创建全部块的软链接 + create_common_dirs_softlink_by_blocks "$install_type" "$install_path" "$package" "$version_dir" "$latest_dir" \ + "$filelist_path" "$feature_param" "" "$custom_create_softlink" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + with_chmod "$db_info_path" "700" write_text "$db_info" "$db_info_path" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + + +# 跨版本按块删除公共目录到latest目录下的软链接 +del_common_dirs_softlink_cross_version() { + local install_path="$1" + local package="$2" + local version_dir="$3" + local latest_dir="$4" + local blocks="$5" + local filelist_path install_info_path install_type feature_type chip_type feature_param + + get_package_filelist "filelist_path" "$install_path" "$version_dir" "$package" + get_package_install_info "install_info_path" "$install_path" "$version_dir" "$package" + + check_file_exists "$install_info_path" "$install_info_path doesn't exist in del common dirs softlink cross version!" + ret="$?" && [ $ret -ne 0 ] && return $ret + + get_package_install_type "install_type" "$install_info_path" "$package" + check_ret_error "$?" "Get install_type from ascend_install.info failed in del common dirs softlink cross version!" + ret="$?" && [ $ret -ne 0 ] && return $ret + + get_package_feature_type "feature_type" "$install_info_path" "$package" "uninstall" + get_package_chip_type "chip_type" "$install_info_path" "$package" "uninstall" + pack_feature_param "feature_param" "$feature_type" "n" "$chip_type" + + del_common_dirs_softlink_by_blocks "$install_type" "$install_path" "$package" "$version_dir" "$latest_dir" \ + "$filelist_path" "$feature_param" "$blocks" "" + check_ret_error "$?" "Del common dirs softlink by blocks failed in del common dirs softlink cross version!" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 按块删除公共目录到latest目录下的软链接 +del_common_dirs_softlink_by_blocks() { + local install_type="$1" + local install_path="$2" + local package="$3" + local version_dir="$4" + local latest_dir="$5" + local filelist_path="$6" + local feature_param="$7" + local blocks="$8" + local custom_remove_softlink="$9" + local ret package_dirpath filelist_path + + create_stash_mod "${install_path}/${latest_dir}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_filelist_v2 "reset_mod_dirs_with_stash_mod" "${install_type}" "${install_path}/${latest_dir}" "mkdir" \ + "${filelist_path}" "${feature_param}" "filter_common_dirs,filter_blocks" "$blocks" "no" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 调用脚本删除latest目录下自定义内容软链接 + if [ "$custom_remove_softlink" != "" ] && [ -f "$custom_remove_softlink" ]; then + chmod u+x "${custom_remove_softlink}" + "${custom_remove_softlink}" --install-path="${install_path}" --version-dir="${version_dir}" --latest-dir="${latest_dir}" + check_ret_error "$?" "Run ${package} custom remove ${version_dir} softlink failed!" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + + foreach_filelist_v2 "del_latest_softlink" "${install_type}" "${install_path}/${latest_dir}" "copy copy_entity move" \ + "${filelist_path}" "${feature_param}" "filter_common_dirs,filter_blocks" "$blocks" "no" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_filelist_v2 "remove_install_dirs" "${install_type}" "${install_path}/${latest_dir}" "mkdir" \ + "${filelist_path}" "${feature_param}" "filter_common_dirs,filter_blocks" "$blocks" "reverse" "normal" + ret="$?" && [ $ret -ne 0 ] && return $ret + + foreach_stashmod "restore_stash_mod" "${install_path}/${latest_dir}" "reverse" + ret="$?" && [ $ret -ne 0 ] && return $ret + + remove_stash_mod "${install_path}/${latest_dir}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 删除公共目录到latest目录下的软链接 +del_common_dirs_softlink_from_latest() { + local install_type="$1" + local install_path="$2" + local package="$3" + local version_dir="$4" + local latest_dir="$5" + local filelist_path="$6" + local feature_param="$7" + local db_info_path="$install_path/$latest_dir/var/ascend_package_db.info" + local ret db_info_origin db_info diff_result blocks_to_remove package_dirpath custom_remove_softlink + + get_package_dirpath "package_dirpath" "$package" + custom_remove_softlink="${install_path}/${version_dir}/${package_dirpath}/script/${package}_custom_remove_softlink.sh" + + if [ -f "$db_info_path" ]; then + db_info_origin="$(cat "$db_info_path")""\n" + else + db_info_origin="" + fi + + db_info="$(printf "$db_info_origin" | del_db_items "$package" "$version_dir")" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 注意换行符和次序 + diff_result="$( + ( + printf "%s\n" "$db_info" | remove_blank_line | remain_db_last_item + printf "$db_info_origin" | remove_blank_line | remain_db_last_item + ) | sort_1 | fold_3_keep_2 "" " " | show_diff_3_4 + )" + blocks_to_remove="$(printf "%s\n" "$diff_result" | select_fields_1 | xargs)" + + # blocks_to_remove可以为空 + del_common_dirs_softlink_by_blocks "$install_type" "$install_path" "$package" "$version_dir" "$latest_dir" \ + "$filelist_path" "$feature_param" "EngineeringCommon $blocks_to_remove" "$custom_remove_softlink" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + printf "%s\n" "$diff_result" | show_min_nf "4" | select_fields_3_2_1 | sort_1 | fold_3_keep_2 "" " " \ + | ( + total_ret=0 + while read -r tmp_version_dir tmp_package tmp_blocks; do + create_common_dirs_softlink_cross_version "$install_path" "$tmp_package" "$tmp_version_dir" "$latest_dir" \ + "EngineeringCommon $tmp_blocks" + ret="$?" && [ $ret -ne 0 ] && total_ret="1" + done + exit $total_ret + ) + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + if [ "$db_info" != "" ]; then + with_chmod "$db_info_path" "700" write_text "$db_info" "$db_info_path" + else + rm -f "$db_info_path" + fi + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# latest目录下公共脚本添加条目 +add_latest_common_script() { + local latest_path="$1" + local package="$2" + local username="$3" + local usergroup="$4" + local docker_root="$5" + local ret mod + + # 保存bin目录的权限 + get_file_mod "mod" "-L" "${latest_path}/bin" + + # 给latest/bin目录提升权限到750 + chmod "750" "${latest_path}/bin" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 设置setenv + add_setenv "${latest_path}" "${package}" "NA" "${username}" "${usergroup}" "true" "${docker_root}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 恢复bin目录的权限 + chmod "${mod}" "${latest_path}/bin" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# latest目录下公共脚本删除条目 +del_latest_common_script() { + local latest_path="$1" + local package="$2" + local username="$3" + local docker_root="$4" + local ret + + if [ ! -d "${latest_path}/bin" ]; then + return 0 + fi + + # 保存bin目录的权限 + get_file_mod "mod" "-L" "${latest_path}/bin" + + # 给latest/bin目录提升权限到750 + chmod "750" "${latest_path}/bin" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # unsetenv + del_setenv "${latest_path}" "${package}" "${username}" "${docker_root}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 恢复bin目录的权限 + chmod "${mod}" "${latest_path}/bin" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 文件中添加一个空行 +add_blank_line_to_file() { + local filepath="$1" + + echo >> "${filepath}" +} + +# 总设置软链脚本是否存在 +# 返回值0为真，1为假 +total_create_softlink_script_exists() { + local install_path="$1" + local version_dir="$2" + + if [ -f "$install_path/$version_dir/$CREATE_VERSION_SOFTLINK" ]; then + return 0 + fi + return 1 +} + +# 创建总设置软链脚本 +create_total_create_softlink_script() { + local script_dir="$1" + local install_path="$2" + local version_dir="$3" + local username="$4" + local usergroup="$5" + local script_filepath="$install_path/$version_dir/$CREATE_VERSION_SOFTLINK" + local ret + + rm -f "$script_filepath" + cat "$script_dir/common_installer.inc" > "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + add_blank_line_to_file "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + sed -n '/^## module log/,/^## end module/ p' "$script_dir/common_func.inc" >> "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + add_blank_line_to_file "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + echo "LATEST_DIR=\"$LATEST_DIR\"" >> "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + add_blank_line_to_file "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + echo "notify_latest_manager_create_version_softlink" >> "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + change_mod "$script_filepath" "550" "" + ret="$?" && [ $ret -ne 0 ] && return $ret + + change_own "$script_filepath" "${username}:${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + + return 0 +} + +# 总删除软链脚本是否存在 +# 返回值0为真，1为假 +total_remove_softlink_script_exists() { + local install_path="$1" + local latest_dir="$2" + + if [ -f "$install_path/$latest_dir/$REMOVE_LATEST_SOFTLINK" ]; then + return 0 + fi + return 1 +} + +# 创建总删除软链脚本 +create_total_remove_softlink_script() { + local script_dir="$1" + local install_path="$2" + local latest_dir="$3" + local username="$4" + local usergroup="$5" + local script_filepath="$install_path/$latest_dir/$REMOVE_LATEST_SOFTLINK" + local ret + + rm -f "$script_filepath" + cat "$script_dir/common_installer.inc" > "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + add_blank_line_to_file "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + sed -n '/^## module log/,/^## end module/ p' "$script_dir/common_func.inc" >> "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + add_blank_line_to_file "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + echo "LATEST_DIR=\"$LATEST_DIR\"" >> "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + add_blank_line_to_file "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + echo "notify_latest_manager_remove_latest_softlink" >> "$script_filepath" + ret="$?" && [ $ret -ne 0 ] && return $ret + + change_mod "$script_filepath" "550" "" + ret="$?" && [ $ret -ne 0 ] && return $ret + + change_own "$script_filepath" "${username}:${usergroup}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 删除总创建软链脚本 +del_total_create_softlink_script() { + local install_path="$1" + local latest_dir="$2" + local version="$3" + local version_dir="$4" + local script_filepath="${install_path}/${version_dir}/${CREATE_VERSION_SOFTLINK}" + local ret version_cfg_path packages + + get_version_cfg_path "version_cfg_path" "${install_path}/${latest_dir}" + if [ ! -f "${version_cfg_path}" ]; then + remove_file "${script_filepath}" "NA" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 + fi + + get_installed_packages_by_version_version_dir "packages" "${install_path}/${latest_dir}" "${version}" "${version_dir}" + if [ "${packages}" = "" ]; then + remove_file "${script_filepath}" "NA" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + return 0 +} + +# 删除总删除软链脚本 +del_total_remove_softlink_script() { + local install_path="$1" + local latest_dir="$2" + local script_filepath="${install_path}/${latest_dir}/${REMOVE_LATEST_SOFTLINK}" + local ret version_cfg_path + + get_version_cfg_path "version_cfg_path" "${install_path}/${latest_dir}" + + if [ ! -f "${version_cfg_path}" ]; then + remove_file "${script_filepath}" "NA" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + return 0 +} + +# platform.ini配置是否存在 +# 返回值0为真，1为假 +platform_ini_exists() { + local install_path="$1" + local latest_dir="$2" + + if [ -f "$install_path/$latest_dir/platform.ini" ]; then + return 0 + fi + return 1 +} + +# 创建platform.ini配置 +create_platform_ini() { + local install_path="$1" + local latest_dir="$2" + local username="$3" + local usergroup="$4" + local config_dirpath="${install_path}/${latest_dir}" + local config_filepath="${config_dirpath}/platform.ini" + local ret mod + + # 配置文件存在则跳过 + if [ -f "${config_filepath}" ]; then + return 0 + fi + + # 保存当前目录的权限 + get_file_mod "mod" -L "${config_dirpath}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + # 恢复权限 + cleanup="chmod ${mod} \"${config_dirpath}\"" + + # 添加目录写权限 + chmod u+w "${config_dirpath}" + cleanup_if_error "$?" "${cleanup}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + touch "${config_filepath}" + cleanup_if_error "$?" "${cleanup}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + change_mod "${config_filepath}" "660" "" + cleanup_if_error "$?" "${cleanup}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + change_own "${config_filepath}" "${username}:${usergroup}" + cleanup_if_error "$?" "${cleanup}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + eval "${cleanup}" + + return 0 +} + +# 删除platform.ini配置 +del_platform_ini() { + local install_path="$1" + local latest_dir="$2" + local config_dirpath="${install_path}/${latest_dir}" + local config_filepath="${config_dirpath}/platform.ini" + local ret mod + + if [ ! -f "${config_filepath}" ]; then + return 0 + fi + + # 保存当前目录的权限 + get_file_mod "mod" -L "${config_dirpath}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + # 恢复权限 + cleanup="chmod ${mod} \"${config_dirpath}\"" + + # 添加目录写权限 + chmod u+w "${config_dirpath}" + cleanup_if_error "$?" "${cleanup}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + remove_file "${config_filepath}" "NA" + cleanup_if_error "$?" "${cleanup}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + eval "${cleanup}" + + return 0 +} + +# 获取aicpu创建软链脚本路径 +get_aicpu_custom_create_softlink_path() { + local _outvar="$1" + local _install_path="$2" + local _version_dir="$3" + local _chip_name="$4" + + eval "${_outvar}=\"${_install_path}/${_version_dir}/opp/${_chip_name}/aicpu/script/aicpu_custom_create_softlink.sh\"" +} + +# 重建软链时处理aicpu软链 +deal_with_aicpu_package() { + local install_path="$1" + local version_dir="$2" + local latest_dir="$3" + local chip_name + + for chip_name in "Ascend910" "Ascend310P" "Ascend310" "Ascend310RC" "Ascend"; do + get_aicpu_custom_create_softlink_path "aicpu_script" "${install_path}" "${version_dir}" "${chip_name}" + if [ -f "${aicpu_script}" ]; then + # 创建软链时，支持强制覆盖老版本软链。 + # 因为存在这样的场景：先卸载新版本的runtime包，再卸载新版本的aicpu_kernels包， + # 卸载新版本的runtime包时触发该流程，此时老版本的aicpu_kernels包软链还在存在于环境上。 + "${aicpu_script}" --install-path="${install_path}" --version-dir="${version_dir}" --latest-dir="${latest_dir}" + # 忽略aicpu创建软链报错 + fi + done + + return 0 +} + +# 恢复可兼容的子包软链接 +recreate_compatiable_softlink_sibling_package() { + local package="$1" + + # 判断是否有本包的running版本包 + get_running_package_version "version_pair" "$scope_install_path/$scope_latest_dir" "$package" + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "$version_pair" != "" ]; then + return 0 + fi + + if check_current_package_compatiable "$scope_install_path" "$scope_version" "$scope_version_dir" "$package"; then + compat_create_package_softlink_to_latest "$scope_install_path" "$package" "$scope_version" "$scope_version_dir" "$scope_latest_dir" \ + "$scope_username" "$scope_usergroup" "$scope_install_for_all" "$scope_docker_root" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + return 0 +} + +# 多版本卸载时检查版本兼容性 +recreate_compatiable_softlink_in_multi_version_uninstall() { + local scope_install_path="$1" + local scope_latest_dir="$2" + local package="$3" + local scope_username="$4" + local scope_usergroup="$5" + local scope_docker_root="$6" + local ret installed_versions version_pair version_pair_arr scope_version scope_version_dir + local scope_install_for_all version_cfg_path sibling_package + + # 切换目录避免使用uninstall.sh脚本，重建子包软链接时，找到不当前路径问题 + # sh: 0: getcwd() failed: No such file or directory + cd "${scope_install_path}" + + # 判断是否有本包的running版本包 + get_running_package_version "version_pair" "${scope_install_path}/${scope_latest_dir}" "${package}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "${version_pair}" != "" ]; then + return 0 + fi + + # 枚举version.cfg中本包installed版本 + get_installed_package_versions "installed_versions" "${scope_install_path}/${scope_latest_dir}" "${package}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + get_install_for_all "scope_install_for_all" + ret="$?" && [ $ret -ne 0 ] && return $ret + INSTALL_FOR_ALL="$scope_install_for_all" + + for version_pair in ${installed_versions}; do + unpack_version_pair "version_pair_arr" "${version_pair}" + __index_list "${version_pair_arr}" 0 "scope_version" 1 "scope_version_dir" + + if check_current_package_compatiable "$scope_install_path" "$scope_version" "$scope_version_dir" "$package"; then + compat_create_package_softlink_to_latest "${scope_install_path}" "${package}" "${scope_version}" "${scope_version_dir}" "${scope_latest_dir}" \ + "${scope_username}" "${scope_usergroup}" "${scope_install_for_all}" "${scope_docker_root}" + ret="$?" && [ $ret -ne 0 ] && return $ret + + # 恢复同版本下可兼容的其他子包软链接 + get_version_cfg_path "version_cfg_path" "$scope_latest_dir" + grep -F "installed_version=" "$version_cfg_path" | grep -F "[$scope_version:$scope_version_dir]" | cut -d= -f1 | sed 's/_installed_version$//' | while read sibling_package; do + if [ "$sibling_package" != "$package" ]; then + recreate_compatiable_softlink_sibling_package "$sibling_package" + fi + done + + if [ "$package" = "runtime" ]; then + deal_with_aicpu_package "${scope_install_path}" "${scope_version_dir}" "${scope_latest_dir}" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + + break + fi + done + + return 0 +} + +# 检查包版本兼容性 +check_package_compatiable() { + local install_path="$1" + local version_left="$2" + local version_dir_left="$3" + local package_left="$4" + local version_right="$5" + local version_dir_right="$6" + local package_right="$7" + local script_dir="$8" + local pkg_version_info_path_left pkg_version_info_path_right + local required_left required_right + + get_package_version_info "pkg_version_info_path_left" "$install_path" "$version_dir_left" "$package_left" + get_package_version_info "pkg_version_info_path_right" "$install_path" "$version_dir_right" "$package_right" + + if ! _get_required_package_info "required_left" "$pkg_version_info_path_right" "$package_left"; then + return 2 + fi + if ! _get_required_package_info "required_right" "$pkg_version_info_path_left" "$package_right"; then + return 2 + fi + + if [ "$required_left" = "" ] && [ "$required_right" = "" ]; then + return 0 + fi + + if _check_version_required "$version_left" "$required_left" "$script_dir" || _check_version_required "$version_right" "$required_right" "$script_dir"; then + return 0 + fi + return 1 +} + +# 检查当前包与latest版本兼容性 +check_current_package_compatiable() { + local install_path="$1" + local current_version="$2" + local current_version_dir="$3" + local current_package="$4" + local script_dir="$install_path/$LATEST_DIR/var/manager" + local ret running_packages package version_pair version version_dir + + get_running_packages "running_packages" "$install_path/$LATEST_DIR" + for package in ${running_packages}; do + if [ "$current_package" = "$package" ]; then + continue + fi + + get_running_package_version "version_pair" "$install_path/$LATEST_DIR" "$package" + ret="$?" && [ $ret -ne 0 ] && return $ret + + unpack_version_pair "version_pair_arr" "$version_pair" + __index_list "$version_pair_arr" 0 "version" 1 "version_dir" + + if ! check_package_compatiable "$install_path" "$current_version" "$current_version_dir" "$current_package" \ + "$version" "$version_dir" "$package" "$script_dir"; then + return 1 + fi + done + + return 0 +} + +# 多版本安装时检查版本兼容性 +check_compatiable_in_multi_version_install() { + local install_path="$1" + local username="$2" + local docker_root="$3" + local current_version="$4" + local current_version_dir="$5" + local current_package="$6" + local script_dir="$install_path/$LATEST_DIR/var/manager" + local ret running_packages package version_pair version version_dir + + get_running_packages "running_packages" "$install_path/$LATEST_DIR" + for package in ${running_packages}; do + if [ "$current_package" = "$package" ]; then + continue + fi + + get_running_package_version "version_pair" "$install_path/$LATEST_DIR" "$package" + ret="$?" && [ $ret -ne 0 ] && return $ret + + unpack_version_pair "version_pair_arr" "$version_pair" + __index_list "$version_pair_arr" 0 "version" 1 "version_dir" + + if ! check_package_compatiable "$install_path" "$current_version" "$current_version_dir" "$current_package" \ + "$version" "$version_dir" "$package" "$script_dir"; then + compat_del_package_softlink_in_latest "$install_path" "$package" "$version" "$version_dir" \ + "$LATEST_DIR" "$username" "$docker_root" + ret="$?" && [ $ret -ne 0 ] && return $ret + fi + done + + return 0 +} + +# 创建版本包目录软链接到latest目录下 +create_package_dir_softlink_to_latest() { + local install_path="$1" + local version_dir="$2" + local latest_dir="$3" + local package="$4" + local username="$5" + local usergroup="$6" + local install_for_all="$7" + local ret package_dirpath package_prefix + + get_package_dirpath "package_dirpath" "${package}" + package_prefix="$(dirname "${package_dirpath}")" + + if [ ! -d "${install_path}/${version_dir}/${package_dirpath}" ]; then + return 0 + fi + + if [ ! -d "${install_path}/${latest_dir}/${package_prefix}" ]; then + make_dir_with_permission "${install_path}/${latest_dir}/${package_prefix}" "750" "${username}" "${usergroup}" "${install_for_all}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + + create_softlink_icp "-r" "${install_path}/${version_dir}/${package_dirpath}" "${install_path}/${latest_dir}/${package_dirpath}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + return 0 +} + +# 删除latest目录下版本包目录的软链接 +del_package_dir_softlink_in_latest() { + local install_path="$1" + local latest_dir="$2" + local package="$3" + local ret package_dirpath package_prefix + + get_package_dirpath "package_dirpath" "${package}" + package_prefix="$(dirname "${package_dirpath}")" + + remove_softlink_icp "${install_path}/${latest_dir}/${package_dirpath}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + + if [ "${package_prefix}" != "." ]; then + remove_dir_if_empty "${install_path}/${latest_dir}/${package_prefix}" + ret="$?" && [ ${ret} -ne 0 ] && return ${ret} + fi + + return 0 +} + +# 生成包db.info条目列表 +packages_db_items() { + local install_path="$1" + local running_packages="$2" + local total_ret=0 ret running_packages package version_pair version_pair_arr version version_dir + local install_type feature_type feature_param filelist_path install_info_path + + for package in $running_packages; do + if ! get_running_package_version "version_pair" "$install_path/$LATEST_DIR" "$package"; then + total_ret=1 + continue + fi + + unpack_version_pair "version_pair_arr" "$version_pair" + __index_list "$version_pair_arr" 0 "version" 1 "version_dir" + + get_package_filelist "filelist_path" "$install_path" "$version_dir" "$package" + get_package_install_info "install_info_path" "$install_path" "$version_dir" "$package" + + check_file_exists "$install_info_path" "$install_info_path doesn't exist in create package softlink to latest!" + ret="$?" && [ $ret -ne 0 ] && total_ret=1 && continue + + get_package_install_type "install_type" "$install_info_path" "$package" + check_ret_error "$?" "Get install_type from ascend_install.info failed in create package softlink to latest!" + ret="$?" && [ $ret -ne 0 ] && total_ret=1 && continue + + get_package_feature_type "feature_type" "$install_info_path" "$package" "install" + check_ret_error "$?" "Get feature_type from ascend_install.info failed in create package softlink to latest!" + ret="$?" && [ $ret -ne 0 ] && total_ret=1 && continue + + get_package_chip_type "chip_type" "$install_info_path" "$package" "install" + check_ret_error "$?" "Get chip_type from ascend_install.info failed in create package softlink to latest!" + ret="$?" && [ $ret -ne 0 ] && total_ret=1 && continue + + pack_feature_param "feature_param" "${feature_type}" "n" "${chip_type}" + + all_common_dirs_blocks_in_filelist "$install_type" "$filelist_path" "$feature_param" \ + | blocks_to_db_item "$package" "$version_dir" + ret="$?" && [ $ret -ne 0 ] && total_ret=1 && continue + done + + return $total_ret +} + +# 生成运行包的ascend_package_db.info +generate_running_packages_db_info() { + local var_path="$1" + local ret install_path running_packages db_info install_for_all + local db_info_path="$var_path/ascend_package_db.info" + + install_path="$(dirname "$(dirname "$var_path")")" + + if [ -f "$db_info_path" ]; then + return 0 + fi + + get_running_packages "running_packages" "$install_path/$LATEST_DIR" + ret="$?" && [ $ret -ne 0 ] && return $ret + + if [ "${running_packages}" = "" ]; then + return 0 + fi + + db_info="$( + packages_db_items "$install_path" "$running_packages" \ + | remove_blank_line | sort_1 "|" | fold_2 "|" "" + )" + ret="$?" && [ $ret -ne 0 ] && return $ret + + get_install_for_all "install_for_all" + ret="$?" && [ $ret -ne 0 ] && return $ret + + ensure_file "$db_info_path" "440" "$USERNAME" "$USERGROUP" "$install_for_all" + ret="$?" && [ $ret -ne 0 ] && return $ret + + with_chmod "$db_info_path" "700" write_text "$db_info" "$db_info_path" + ret="$?" && [ $ret -ne 0 ] && return $ret + + return 0 +} + +# 升级setenv脚本 +upgrade_setenv() { + local latest_path="$1" + local shell_type config_path + for shell_type in ${ENV_SHELL_TYPES}; do + get_setenv_filepath "config_path" "$latest_path" "$shell_type" + add_ascend_home_path_env "$config_path" "$shell_type" "$latest_path" + done +} + +# latest下数据迁移 +migrate_latest_data() { + local var_path="$1" + local latest_path="$(dirname "$var_path")" + local ret config_path + + generate_running_packages_db_info "$var_path" + ret="$?" && [ $ret -ne 0 ] && return $ret + + get_setenv_filepath "config_path" "$latest_path" "bash" + if ! has_ascend_home_path_env "$config_path"; then + with_chmod "$latest_path/bin" "700" upgrade_setenv "$latest_path" + fi + + return 0 +} diff --git a/csrc/ascend/scripts/package/latest_manager/scripts/uninstall.sh b/csrc/ascend/scripts/package/latest_manager/scripts/uninstall.sh new file mode 100644 index 000000000..5ff78fa04 --- /dev/null +++ b/csrc/ascend/scripts/package/latest_manager/scripts/uninstall.sh @@ -0,0 +1,52 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +CURPATH=$(dirname $(readlink -f "$0")) +VARPATH="$(dirname "$CURPATH")" +USERNAME=$(id -un) +USERGROUP=$(id -gn) +common_func_path="$CURPATH/common_func.inc" +manager_func_path="$CURPATH/manager_func.sh" + +. "$common_func_path" +. "$manager_func_path" + +set_comm_log "Latest_manager" "$COMM_LOGFILE" + +IS_UPGRADE="n" + +while true +do + case "$1" in + --upgrade) + IS_UPGRADE="y" + shift + ;; + *) + break + ;; + esac +done + +if ! sh "$CURPATH/install_common_parser.sh" --package="latest_manager" --uninstall --username="$USERNAME" --usergroup="$USERGROUP" \ + --simple-uninstall "full" "$VARPATH" "$CURPATH/filelist.csv" "all"; then + comm_log "ERROR" "uninstall failed!" + exit 1 +fi + +if [ "$IS_UPGRADE" = "n" ]; then + remove_manager_refs "$VARPATH" +fi + +if ! remove_dir_if_empty "$VARPATH"; then + comm_log "ERROR" "uninstall failed!" + exit 1 +fi diff --git a/csrc/ascend/scripts/package/latest_manager/scripts/version.info b/csrc/ascend/scripts/package/latest_manager/scripts/version.info new file mode 100644 index 000000000..526e64acf --- /dev/null +++ b/csrc/ascend/scripts/package/latest_manager/scripts/version.info @@ -0,0 +1 @@ +Version=45 diff --git a/csrc/ascend/scripts/package/module/ascend/EngineeringCommon.xml b/csrc/ascend/scripts/package/module/ascend/EngineeringCommon.xml new file mode 100644 index 000000000..43baf819c --- /dev/null +++ b/csrc/ascend/scripts/package/module/ascend/EngineeringCommon.xml @@ -0,0 +1,14 @@ + + + + + + + + + + + + + + diff --git a/csrc/ascend/scripts/package/module/ascend/EngineeringFiles.xml b/csrc/ascend/scripts/package/module/ascend/EngineeringFiles.xml new file mode 100644 index 000000000..fc3174406 --- /dev/null +++ b/csrc/ascend/scripts/package/module/ascend/EngineeringFiles.xml @@ -0,0 +1,6 @@ + + + + + + diff --git a/csrc/ascend/scripts/package/module/ascend/OpsTransformer.xml b/csrc/ascend/scripts/package/module/ascend/OpsTransformer.xml new file mode 100644 index 000000000..518345cd1 --- /dev/null +++ b/csrc/ascend/scripts/package/module/ascend/OpsTransformer.xml @@ -0,0 +1,26 @@ + + + + + + + + + + + + + + + + + + + + + + diff --git a/csrc/ascend/scripts/package/module/ascend/OpsTransformerInc.xml b/csrc/ascend/scripts/package/module/ascend/OpsTransformerInc.xml new file mode 100644 index 000000000..dd8474080 --- /dev/null +++ b/csrc/ascend/scripts/package/module/ascend/OpsTransformerInc.xml @@ -0,0 +1,49 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/csrc/ascend/scripts/package/ops_transformer/ops_transformer.xml b/csrc/ascend/scripts/package/ops_transformer/ops_transformer.xml new file mode 100644 index 000000000..8ee96d257 --- /dev/null +++ b/csrc/ascend/scripts/package/ops_transformer/ops_transformer.xml @@ -0,0 +1,59 @@ + + + + cann + 910b + ops-transformer + run + share/info/ops_transformer/script/install.sh + share/info/ops_transformer/script/help.info + share/info/ops_transformer/script/cleanup.sh + + + $(ASCEND_VER) + $(TIMESTAMP_NO) + + + $(OS_NAME) + $(OS_VER) + $(ARCH) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/csrc/ascend/scripts/package/ops_transformer/scripts/cleanup.sh b/csrc/ascend/scripts/package/ops_transformer/scripts/cleanup.sh new file mode 100644 index 000000000..e810cb338 --- /dev/null +++ b/csrc/ascend/scripts/package/ops_transformer/scripts/cleanup.sh @@ -0,0 +1,14 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +set -e + +rm -rf $(pwd) diff --git a/csrc/ascend/scripts/package/ops_transformer/scripts/empty_package_scripts/cleanup.sh b/csrc/ascend/scripts/package/ops_transformer/scripts/empty_package_scripts/cleanup.sh new file mode 100644 index 000000000..7b842999b --- /dev/null +++ b/csrc/ascend/scripts/package/ops_transformer/scripts/empty_package_scripts/cleanup.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +exit 0 diff --git a/csrc/ascend/scripts/package/ops_transformer/scripts/empty_package_scripts/install.sh b/csrc/ascend/scripts/package/ops_transformer/scripts/empty_package_scripts/install.sh new file mode 100644 index 000000000..3672207dc --- /dev/null +++ b/csrc/ascend/scripts/package/ops_transformer/scripts/empty_package_scripts/install.sh @@ -0,0 +1,36 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +# error number and description + +if [ "$(id -u)" != "0" ]; then + _LOG_PATH=$(echo "${HOME}")"/var/log/ascend_seclog" + _INSTALL_LOG_FILE="${_LOG_PATH}/ascend_install.log" +else + _LOG_PATH="/var/log/ascend_seclog" + _INSTALL_LOG_FILE="${_LOG_PATH}/ascend_install.log" +fi + +# log functions +getdate() { + _cur_date=$(date +"%Y-%m-%d %H:%M:%S") + echo "${_cur_date}" +} + +logandprint() { + is_error_level=$(echo $1 | grep -E 'ERROR|WARN|INFO') + if [ "${is_quiet}" != "y" ] || [ "${is_error_level}" != "" ]; then + echo "[OpsTransformer] [$(getdate)] ""$1" + fi + echo "[OpsTransformer] [$(getdate)] ""$1" >>"${_INSTALL_LOG_FILE}" +} +logandprint "[INFO]: Opp package installed successfully! The new version takes effect immediately." + +exit 0 diff --git a/csrc/ascend/scripts/package/ops_transformer/scripts/help.info b/csrc/ascend/scripts/package/ops_transformer/scripts/help.info new file mode 100644 index 000000000..2e2d30750 --- /dev/null +++ b/csrc/ascend/scripts/package/ops_transformer/scripts/help.info @@ -0,0 +1,10 @@ + --full Install full mode + --install-path= Install product to specific dir path + --install-for-all Install for all user + --quiet Quiet install mode, skip human-computer interactions + --uninstall Uninstall product with compatible run package which is installed before + --install-path= Uninstall specific ops_transformer dir path + --upgrade Upgrade product immediately + --install-path= Upgrade specific ops_transformer dir path + --install-for-all Install for all user + --quiet Quiet install mode, skip human-computer interactions diff --git a/csrc/ascend/scripts/package/ops_transformer/scripts/install.sh b/csrc/ascend/scripts/package/ops_transformer/scripts/install.sh new file mode 100644 index 000000000..52c68dc1d --- /dev/null +++ b/csrc/ascend/scripts/package/ops_transformer/scripts/install.sh @@ -0,0 +1,617 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +# error number and description +OPERATE_FAILED="0x0001" +PARAM_INVALID="0x0002" +FILE_NOT_EXIST="0x0080" +FILE_NOT_EXIST_DES="File not found." +OPP_COMPATIBILITY_CEHCK_ERR="0x0092" +OPP_COMPATIBILITY_CEHCK_ERR_DES="OppTransformer compatibility check error." +PERM_DENIED="0x0093" +PERM_DENIED_DES="Permission denied." + +OPP_PLATFORM_DIR=ops_transformer +OPP_PLATFORM_UPPER=$(echo "${OPP_PLATFORM_DIR}" | tr '[:lower:]' '[:upper:]') +CURR_OPERATE_USER="$(id -nu 2>/dev/null)" +CURR_OPERATE_GROUP="$(id -ng 2>/dev/null)" +# defaults for general user +if [ "$(id -u)" != "0" ]; then + DEFAULT_INSTALL_PATH="${HOME}/Ascend" +else + IS_FOR_ALL="y" + DEFAULT_INSTALL_PATH="/usr/local/Ascend" +fi + +# run package's files info, CURR_PATH means current temp path +CURR_PATH=$(dirname $(readlink -f $0)) +INSTALL_SHELL_FILE="${CURR_PATH}/opp_install.sh" +RUN_PKG_INFO_FILE="${CURR_PATH}/../scene.info" +VERSION_INFO_FILE="${CURR_PATH}/../version.info" +COMMON_INC_FILE="${CURR_PATH}/common_func.inc" +VERCHECK_FILE="${CURR_PATH}/ver_check.sh" +VERSION_COMPAT_FUNC_PATH="${CURR_PATH}/version_compatiable.inc" +COMMON_FUNC_V2_PATH="${CURR_PATH}/common_func_v2.inc" +VERSION_CFG_PATH="${CURR_PATH}/version_cfg.inc" +OPP_COMMON_FILE="${CURR_PATH}/opp_common.sh" + +. "${VERSION_COMPAT_FUNC_PATH}" +. "${COMMON_INC_FILE}" +. "${COMMON_FUNC_V2_PATH}" +. "${VERSION_CFG_PATH}" +. "${OPP_COMMON_FILE}" + +ARCH_INFO=$(grep -e "arch" "$RUN_PKG_INFO_FILE" | cut --only-delimited -d"=" -f2-) +# 包内路径 +GRAPH_SO_PATH="${CURR_PATH}/../../../../${OPP_PLATFORM_DIR}/built-in/op_graph/lib/linux/${ARCH_INFO}/libopgraph_transformer.so" +HOST_SO_PATH="${CURR_PATH}/../../../../${OPP_PLATFORM_DIR}/built-in/op_impl/ai_core/tbe/op_host/lib/linux/${ARCH_INFO}/libophost_transformer.so" + +# defaults info determined by user's inputs +ASCEND_INSTALL_INFO="ascend_install.info" +TARGET_INSTALL_PATH="${DEFAULT_INSTALL_PATH}" #--input-path +TARGET_USERNAME="${CURR_OPERATE_USER}" +TARGET_USERGROUP="${CURR_OPERATE_GROUP}" +TARGET_VERSION_DIR="" # TARGET_INSTALL_PATH + PKG_VERSION_DIR +TARGET_SHARED_INFO_DIR="" + +# keys of infos in ascend_install.info +KEY_INSTALLED_UNAME="USERNAME" +KEY_INSTALLED_UGROUP="USERGROUP" +KEY_INSTALLED_TYPE="${OPP_PLATFORM_UPPER}_INSTALL_TYPE" +KEY_INSTALLED_PATH="${OPP_PLATFORM_UPPER}_INSTALL_PATH_VAL" +KEY_INSTALLED_VERSION="${OPP_PLATFORM_UPPER}_VERSION" +KEY_INSTALLED_FEATURE="${OPP_PLATFORM_UPPER}_INSTALL_FEATURE" +KEY_INSTALLED_CHIP="${OPP_PLATFORM_UPPER}_INSTALL_CHIP" + +# keys of infos in run package +KEY_RUNPKG_VERSION="Version" + +# init install cmd status, set default as n +CMD_LIST="$*" +IS_UNINSTALL=n +IS_INSTALL=n +IS_UPGRADE=n +IS_QUIET=n +IS_INPUT_PATH=n +IS_CHECK=n +IN_INSTALL_TYPE="" +IN_INSTALL_PATH="" +IS_DOCKER_INSTALL=n +IS_SETENV=n +DOCKER_ROOT="" +CONFLICT_CMD_NUMS=0 +IN_FEATURE="All" + +# log functions +# start info before shell executing +startlog() { + echo "[OpsTransformer] [$(getdate)] [INFO]: Start Time: $(getdate)" +} + +exitlog() { + echo "[OpsTransformer] [$(getdate)] [INFO]: End Time: $(getdate)" +} + +#check ascend_install.info for the change in code warning +get_installed_info() { + local key="$1" + local res="" + if [ -f "${INSTALL_INFO_FILE}" ]; then + chmod 644 "${INSTALL_INFO_FILE}" >/dev/null 2>&1 + res=$(cat ${INSTALL_INFO_FILE} | grep "${key}" | awk -F = '{print $2}') + fi + echo "${res}" +} + +clean_before_reinstall() { + local installed_path=$(get_installed_info "${KEY_INSTALLED_PATH}") + local existed_files=$(find ${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR} -type f -print 2>/dev/null) + if [ -z "${existed_files}" ]; then + logandprint "[INFO]: Directory is empty, directly install opp module." + return 0 + fi + + if [ "${IS_QUIET}" = "y" ]; then + logandprint "[WARNING]: Directory has file existed or installed opp\ + module, are you sure to keep installing opp module in it? y" + else + if [ ! -f "${INSTALL_INFO_FILE}" ]; then + logandprint "[INFO]: Directory has file existed, do you want to continue? [y/n]" + else + logandprint "[INFO]: Opp package has been installed on the path $(get_installed_info "${KEY_INSTALLED_PATH}"),\ + the version is $(get_installed_info "${KEY_INSTALLED_VERSION}"),\ + and the version of this package is ${RUN_PKG_VERSION}, do you want to continue? [y/n]" + fi + while true; do + read yn + if [ "$yn" = "n" ]; then + logandprint "[INFO]: Exit to install opp module." + exitlog + exit 0 + elif [ "$yn" = "y" ]; then + break + else + echo "[WARNING]: Input error, please input y or n to choose!" + fi + done + fi + + if [ "${installed_path}" = "${TARGET_VERSION_DIR}" ]; then + logandprint "[INFO]: Clean the installed opp module before install." + if [ ! -f "${UNINSTALL_SHELL_FILE}" ]; then + logandprint "[ERROR]: ERR_NO:${FILE_NOT_EXIST};ERR_DES:${FILE_NOT_EXIST_DES}.The file\ + (${UNINSTALL_SHELL_FILE}) not exists. Please set the correct install \ + path or clean the previous version opp install info (${INSTALL_INFO_FILE}) and then reinstall it." + return 1 + fi + bash "${UNINSTALL_SHELL_FILE}" "${TARGET_VERSION_DIR}" "upgrade" "${IS_QUIET}" ${IN_FEATURE} "${IS_DOCKER_INSTALL}" "${DOCKER_ROOT}" "$pkg_version_dir" + if [ "$?" != 0 ]; then + logandprint "[ERROR]: ERR_NO:${INSTALL_FAILED};ERR_DES:Clean the installed directory failed." + return 1 + fi + fi + return 0 +} + +select_last_dir_component() { + path="$1" + last_component=$(basename ${path}) + if [ "${last_component}" = "atc" ]; then + last_component="atc" + return + elif [ "${last_component}" = "fwkacllib" ]; then + last_component="fwkacllib" + return + elif [ "${last_component}" = "compiler" ]; then + last_component="compiler" + return + fi +} + +# check_version_file() { +# pkg_path="$1" +# component_ret="$2" +# run_pkg_path_temp=$(dirname "${pkg_path}") +# run_pkg_path_temp2=${run_pkg_path_temp%/*} +# run_pkg_path="${run_pkg_path_temp}""/${component_ret}" +# run_pkg_path_temp2=${run_pkg_path%/*} +# version_file="${run_pkg_path}""/version.info" +# version_file_tmp="${run_pkg_path_temp2}""/version.info" +# if [ -f "${version_file_tmp}" ]; then +# version_file=${version_file_tmp} +# fi +# if [ -f "${version_file}" ]; then +# echo "${version_file}" 2 >>/dev/null +# else +# logandprint "[ERROR]: ERR_NO:${FILE_NOT_EXIST}; The [${component_ret}] version.info in path [${pkg_path}] not exists." +# exitlog +# exit 1 +# fi +# return +# } + +check_opp_version_file() { + if [ -f "${CURR_PATH}/../../version.info" ]; then + opp_ver_info="${CURR_PATH}/../../version.info" + elif [ -f "${DEFAULT_INSTALL_PATH}/${OPP_PLATFORM_DIR}/share/info/version.info" ]; then + opp_ver_info="${DEFAULT_INSTALL_PATH}/${OPP_PLATFORM_DIR}/share/info/version.info" + else + logandprint "[ERROR]: ERR_NO:${FILE_NOT_EXIST}; The [${OPP_PLATFORM_DIR}] version.info not exists." + exitlog + exit 1 + fi + echo "find opp_ver_info: ${opp_ver_info}" + return +} + + +check_docker_path() { + docker_path="$1" + if [[ "${docker_path}" != /* ]]; then + echo "[OpsTransformer] [ERROR]: ERR_NO:${PARAM_INVALID};ERR_DES:Parameter --docker-root\ + must with absolute path that which is start with root directory /. Such as --docker-root=/${docker_path}" + exitlog + exit 1 + fi + if [ ! -d "${docker_path}" ]; then + echo "[OpsTransformer] [ERROR]: ERR_NO:${FILE_NOT_EXIST}; The directory:${docker_path} not exist, please create this directory." + exitlog + exit 1 + fi +} + +judgment_path() { + . "${COMMON_INC_FILE}" + check_install_path_valid "${1}" + if [ $? -ne 0 ]; then + echo "[OpsTransformer][ERROR]: The opp install path ${1} is invalid, only characters in [a-z,A-Z,0-9,-,_] are supported!" + exitlog + exit 1 + fi +} + +check_install_path() { + TARGET_INSTALL_PATH="$1" + # empty patch check + if [ "x${TARGET_INSTALL_PATH}" = "x" ]; then + echo "[OpsTransformer] [ERROR]: ERR_NO:${PARAM_INVALID};ERR_DES:Parameter --install-path\ + not support that the install path is empty." + exitlog + exit 1 + fi + # space check + if echo "x${TARGET_INSTALL_PATH}" | grep -q " "; then + echo "[OpsTransformer] [ERROR]: ERR_NO:${PARAM_INVALID};ERR_DES:Parameter --install-path\ + not support that the install path contains space character." + exitlog + exit 1 + fi + # delete last "/" + local temp_path="${TARGET_INSTALL_PATH}" + temp_path=$(echo "${temp_path%/}") + if [ x"${temp_path}" = "x" ]; then + temp_path="/" + fi + # convert relative path to absolute path + local prefix=$(echo "${temp_path}" | cut -d"/" -f1 | cut -d"~" -f1) + if [ "x${prefix}" = "x" ]; then + TARGET_INSTALL_PATH="${temp_path}" + else + prefix=$(echo "${RUN_PATH}" | cut -d"/" -f1 | cut -d"~" -f1) + if [ x"${prefix}" = "x" ]; then + TARGET_INSTALL_PATH="${RUN_PATH}/${temp_path}" + else + echo "[OpsTransformer] [ERROR]: ERR_NO:${PARAM_INVALID};ERR_DES: Run package path is invalid: $RUN_PATH" + exitlog + exit 1 + fi + fi + # convert '~' to home path + local home=$(echo "${TARGET_INSTALL_PATH}" | cut -d"~" -f1) + if [ "x${home}" = "x" ]; then + local temp_path_value=$(echo "${TARGET_INSTALL_PATH}" | cut -d"~" -f2) + if [ "$(id -u)" -eq 0 ]; then + TARGET_INSTALL_PATH="/root$temp_path_value" + else + local home_path=$(eval echo "${USER}") + home_path=$(echo "${home_path}%/") + TARGET_INSTALL_PATH="$home_path$temp_path_value" + fi + fi +} + +#get the dir of xxx.run +#opp_install_path_curr=`echo "$2" | cut -d"/" -f2- ` +# cut first two params from *.run +get_run_path() { + RUN_PATH=$(echo "$2" | cut -d"-" -f3-) + if [ x"${RUN_PATH}" = x"" ]; then + RUN_PATH=$(pwd) + else + # delete last "/" + RUN_PATH=$(echo "${RUN_PATH%/}") + if [ "x${RUN_PATH}" = "x" ]; then + # root path + RUN_PATH=$(pwd) + fi + fi +} + +get_opts() { + i=0 + while true + do + if [ "x$1" = "x" ]; then + break + fi + if [ "$(expr substr "$1" 1 2)" = "--" ]; then + i=$(expr $i + 1) + fi + if [ $i -gt 2 ]; then + break + fi + shift 1 + done + + if [ "$*" = "" ]; then + echo "[ERROR]: ERR_NO:${PARAM_INVALID}; ERR_DES:Unrecognized parameters.Try './xxx.run --help for more information.'" + exitlog + exit 1 + fi + + while true; do + # skip 2 parameters avoid run pkg and directory as input parameter + case "$1" in + --full) + IN_INSTALL_TYPE=$(echo ${1} | awk -F"--" '{print $2}') + IS_INSTALL="y" + CONFLICT_CMD_NUMS=$(expr $CONFLICT_CMD_NUMS + 1) + shift + ;; + --upgrade) + IS_UPGRADE="y" + CONFLICT_CMD_NUMS=$(expr $CONFLICT_CMD_NUMS + 1) + shift + ;; + --uninstall) + IS_UNINSTALL="y" + CONFLICT_CMD_NUMS=$(expr $CONFLICT_CMD_NUMS + 1) + shift + ;; + --install-path=*) + IS_INPUT_PATH="y" + IN_INSTALL_PATH=$(echo ${1} | cut -d"=" -f2-) + # check path + judgment_path "${IN_INSTALL_PATH}" + check_install_path "${IN_INSTALL_PATH}" + shift + ;; + --quiet) + IS_QUIET="y" + shift + ;; + --install-for-all) + IS_FOR_ALL="y" + shift + ;; + -*) + echo "[OpsTransformer] [ERROR]: ERR_NO:${PARAM_INVALID};ERR_DES:Unsupported parameters [$1],\ + operation execute failed. Please use [--help] to see the usage." + exitlog + exit 1 + ;; + *) + break + ;; + esac + done +} + +# pre-check +check_opts() { + if [ "${CONFLICT_CMD_NUMS}" != 1 ]; then + echo "[OpsTransformer] [ERROR]: ERR_NO:${PARAM_INVALID};ERR_DES:\ + only support one type: full/run/devel/upgrade/uninstall/check, operation execute failed!\ + Please use [--help] to see the usage." + exitlog + exit 1 + fi +} + +# init target_dir and log for install +init_env() { + # create log folder and log file + comm_init_log + + if is_version_dirpath "$TARGET_INSTALL_PATH"; then + pkg_version_dir="$(basename "$TARGET_INSTALL_PATH")" + TARGET_INSTALL_PATH="$(dirname "$TARGET_INSTALL_PATH")" + else + pkg_version_dir="cann" + fi + TARGET_VERSION_DIR="$TARGET_INSTALL_PATH/$pkg_version_dir" # Splicing docker-root and install-path + if [ "${IS_DOCKER_INSTALL}" = "y" ]; then + # delete last "/" + local temp_path_param="${DOCKER_ROOT}" + local temp_path_val=$(echo "${temp_path_param%/}") + if [ "x${temp_path_val}" = "x" ]; then + temp_path_val="/" + fi + TARGET_VERSION_DIR=${temp_path_val}${TARGET_VERSION_DIR} + fi + + TARGET_SHARED_INFO_DIR=${TARGET_VERSION_DIR}/share/info + UNINSTALL_SHELL_FILE="${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}/script/opp_uninstall.sh" + INSTALL_INFO_FILE="${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}/${ASCEND_INSTALL_INFO}" + + logandprint "[INFO]: Execute the opp run package." + logandprint "[INFO]: OperationLogFile path: ${COMM_LOGFILE}." + logandprint "[INFO]: Input params: $CMD_LIST" + + get_package_version "RUN_PKG_VERSION" "$VERSION_INFO_FILE" + local installed_version=$(get_installed_info "${KEY_INSTALLED_VERSION}") + if [ "${installed_version}" = "" ]; then + logandprint "[INFO]: Version of installing opp module is ${RUN_PKG_VERSION}." + else + if [ "${RUN_PKG_VERSION}" != "" ]; then + logandprint "[INFO]: Existed opp module version is ${installed_version},\ + the new opp module version is ${RUN_PKG_VERSION}." + fi + fi +} + +check_pre_install() { + local installed_user=$(get_installed_info "${KEY_INSTALLED_UNAME}") + local installed_group=$(get_installed_info "${KEY_INSTALLED_UGROUP}") + if [ "${installed_user}" != "" ] || [ "${installed_group}" != "" ]; then + if [ "${installed_user}" != "${TARGET_USERNAME}" ] || [ "${installed_group}" != "${TARGET_USERGROUP}" ]; then + logandprint "[ERROR]: The user and group are not same with last installation,\ + do not support overwriting installation!" + exitlog + exit 1 + fi + fi + + if [ "${IS_UPGRADE}" = "y" ]; then + if [ ! -e "${INSTALL_INFO_FILE}" ]; then + logandprint "[ERROR]: ERR_NO:${FILE_NOT_EXIST}; The directory:${TARGET_INSTALL_PATH} not install OpsTransformer, upgrade failed." + exitlog + exit 1 + fi + IN_INSTALL_TYPE=$(get_installed_info "${KEY_INSTALLED_TYPE}") + fi +} + +#Support the installation script when the specified path (relative path and absolute path) does not exist +mkdir_install_path() { + local base_dir=$(dirname ${TARGET_INSTALL_PATH}) + if [ ! -d ${base_dir} ]; then + logandprint "[ERROR]: ERR_NO:${FILE_NOT_EXIST}; The directory:${base_dir} not exist, please create this directory." + exitlog + exit 1 + fi + + if [ -d "${TARGET_INSTALL_PATH}" ]; then + test -w ${TARGET_INSTALL_PATH} >>/dev/null 2>&1 + if [ "$?" -ne 0 ]; then + #All paths exist with write permission + logandprint "[ERROR]: ERR_NO:${PERM_DENIED};ERR_DES:${PERM_DENIED_DES}. The ${TARGET_USERNAME} do\ + access ${TARGET_INSTALL_PATH} failed, please reset the directory to a right permission." + exit 1 + fi + else + test -w ${base_dir} >>/dev/null 2>&1 + if [ "$?" -ne 0 ]; then + #All paths exist with write permission + logandprint "[ERROR]: ERR_NO:${PERM_DENIED};ERR_DES:${PERM_DENIED_DES}. The ${TARGET_USERNAME} do\ + access ${base_dir} failed, please reset the directory to a right permission." + exit 1 + else + comm_create_dir "${TARGET_INSTALL_PATH}" "750" "${TARGET_USERNAME}:${TARGET_USERGROUP}" "${IS_FOR_ALL}" + fi + fi +} + +install_package() { + if [ "${IS_INSTALL}" = "n" ] && [ "${IS_UPGRADE}" = "n" ]; then + return + fi + + local architecture=$(uname -m) + local graph_so_dir_path="${TARGET_VERSION_DIR}/opp/built-in/op_graph/lib/linux/${ARCH_INFO}" + local host_so_dir_path="${TARGET_VERSION_DIR}/opp/built-in/op_impl/ai_core/tbe/op_host/lib/linux/${ARCH_INFO}" + # check platform + if [ "${architecture}" != "${ARCH_INFO}" ] ; then + logandprint "[INFO]: the architecture of the run package is inconsistent with that of the current environment. " + # 异构安装场景，拷贝so到指定目录 + if [ -d "${TARGET_VERSION_DIR}/opp/built-in/op_graph/lib/linux" ] ; then + chmod u+w ${TARGET_VERSION_DIR}/opp/built-in/op_graph/lib/linux + fi + if [ -d "${TARGET_VERSION_DIR}/opp/built-in/op_impl/ai_core/tbe/op_host/lib/linux" ] ; then + chmod u+w ${TARGET_VERSION_DIR}/opp/built-in/op_impl/ai_core/tbe/op_host/lib/linux + fi + mkdir -p ${graph_so_dir_path} + mkdir -p ${host_so_dir_path} + cp ${GRAPH_SO_PATH} ${graph_so_dir_path} + cp ${HOST_SO_PATH} ${host_so_dir_path} + chmod 755 ${graph_so_dir_path}/* + chmod 755 ${host_so_dir_path}/* + + chmod u-w ${TARGET_VERSION_DIR}/opp/built-in/op_graph/lib/linux + chmod u-w ${TARGET_VERSION_DIR}/opp/built-in/op_impl/ai_core/tbe/op_host/lib/linux + exit 0 + fi + + # use uninstall to clean the install folder + clean_before_reinstall + if [ "$?" != 0 ]; then + comm_log_operation "Install" "${IN_INSTALL_TYPE}" "OpsTransformer" "$?" "${CMD_LIST}" + fi + + bash "${INSTALL_SHELL_FILE}" "${TARGET_INSTALL_PATH}" "${TARGET_USERNAME}" "${TARGET_USERGROUP}" "${IN_FEATURE}" \ + "${IN_INSTALL_TYPE}" "${IS_FOR_ALL}" "${IS_SETENV}" "${IS_DOCKER_INSTALL}" "${DOCKER_ROOT}" "$pkg_version_dir" + if [ "$?" != 0 ]; then + comm_log_operation "Install" "${IN_INSTALL_TYPE}" "OpsTransformer" "$?" "${CMD_LIST}" + fi + if [ $(id -u) -eq 0 ]; then + chown -R "root":"root" "${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}/script" 2>/dev/null + chown "root":"root" "${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}" 2>/dev/null + chmod -R 555 "${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}/script" 2>/dev/null + chmod 444 "${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}/script/filelist.csv" 2>/dev/null + else + chmod -R 550 "${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}/script" 2>/dev/null + chmod 440 "${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}/script/filelist.csv" 2>/dev/null + fi + comm_log_operation "Install" "${IN_INSTALL_TYPE}" "OpsTransformer" "$?" "${CMD_LIST}" +} + +uninstall_package() { + if [ "${IS_UNINSTALL}" = "n" ]; then + return + fi + + if [ ! -f "${UNINSTALL_SHELL_FILE}" ]; then + logandprint "[ERROR]: ERR_NO:${FILE_NOT_EXIST};ERR_DES:The file\ + (${UNINSTALL_SHELL_FILE}) not exists. Please make sure that the opp module\ + installed in (${TARGET_VERSION_DIR}) and then set the correct install path." + uninstall_path=$(ls "${TARGET_INSTALL_PATH}" 2>/dev/null) + if [ "${uninstall_path}" = "" ]; then + rm -rf "${TARGET_INSTALL_PATH}" + fi + comm_log_operation "Uninstall" "${IN_INSTALL_TYPE}" "OpsTransformer" "$?" "${CMD_LIST}" + exit 0 + fi + + # 如果是异构卸载 + local architecture=$(uname -m) + if [ "${architecture}" != ${ARCH_INFO} ]; then + target_arch=${ARCH_INFO} + else + # 判断异构so是否存在，存在则删除 + if [ "${architecture}" = "x86_64" ]; then + target_arch="aarch64" + else + target_arch="x86_64" + fi + fi + local graph_so_path="${TARGET_VERSION_DIR}/opp/built-in/op_graph/lib/linux/${target_arch}/libopgraph_transformer.so" + local graph_so_dir_path="${TARGET_VERSION_DIR}/opp/built-in/op_graph/lib/linux/${target_arch}" + local host_so_path="${TARGET_VERSION_DIR}/opp/built-in/op_impl/ai_core/tbe/op_host/lib/linux/${target_arch}/libophost_transformer.so" + local host_so_dir_path="${TARGET_VERSION_DIR}/opp/built-in/op_impl/ai_core/tbe/op_host/lib/linux/${target_arch}" + if [ -f "${graph_so_path}" ]; then + rm -f "${graph_so_path}" + fi + if [ -f "${host_so_path}" ]; then + rm -f "${host_so_path}" + fi + # 判断目录是否存在且是否为空 + if [ -d "${graph_so_dir_path}" ]; then + if [ -z "$(ls -A "${graph_so_dir_path}")" ]; then + rm -rf "${graph_so_dir_path}" + fi + fi + if [ -d "${host_so_dir_path}" ]; then + if [ -z "$(ls -A "${host_so_dir_path}")" ]; then + rm -rf "${host_so_dir_path}" + fi + fi + if [ "${architecture}" != ${ARCH_INFO} ]; then + return + fi + + bash "${UNINSTALL_SHELL_FILE}" "${TARGET_INSTALL_PATH}" "uninstall" "${IS_QUIET}" ${IN_FEATURE} "${IS_DOCKER_INSTALL}" "${DOCKER_ROOT}" "$pkg_version_dir" + logandprint "[INFO]: Remove precheck info." + + comm_log_operation "Uninstall" "${IN_INSTALL_TYPE}" "OpsTransformer" "$?" "${CMD_LIST}" +} + +main() { + get_run_path "$@" + + startlog + + get_opts "$@" + + check_opts + + init_env + + check_pre_install + + mkdir_install_path + + install_package + + uninstall_package + +} + +main "$@" +exit 0 diff --git a/csrc/ascend/scripts/package/ops_transformer/scripts/opp_common.sh b/csrc/ascend/scripts/package/ops_transformer/scripts/opp_common.sh new file mode 100644 index 000000000..4bcf23c8b --- /dev/null +++ b/csrc/ascend/scripts/package/ops_transformer/scripts/opp_common.sh @@ -0,0 +1,271 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +if [ "$(id -u)" != "0" ]; then + _LOG_PATH=$(echo "${HOME}")"/var/log/ascend_seclog" + _INSTALL_LOG_FILE="${_LOG_PATH}/ascend_install.log" +else + _LOG_PATH="/var/log/ascend_seclog" + _INSTALL_LOG_FILE="${_LOG_PATH}/ascend_install.log" +fi + +# log functions +getdate() { + _cur_date=$(date +"%Y-%m-%d %H:%M:%S") + echo "${_cur_date}" +} + +logandprint() { + is_error_level=$(echo $1 | grep -E 'ERROR|WARN|INFO') + if [ "${is_quiet}" != "y" ] || [ "${is_error_level}" != "" ]; then + echo "[OpsTransformer] [$(getdate)] ""$1" + fi + echo "[OpsTransformer] [$(getdate)] ""$1" >>"${_INSTALL_LOG_FILE}" +} + +# create opapi soft link +createrelativelysoftlink() { + local src_path_="$1" + local dst_path_="$2" + local dst_parent_path_=$(dirname ${dst_path_}) + # echo "dst_parent_path_: ${dst_parent_path_}" + local relative_path_=$(realpath --relative-to="$dst_parent_path_" "$src_path_") + # echo "relative_path_: ${relative_path_}" + if [ -L "$2" ]; then + return 0 + fi + ln -s "${relative_path_}" "${dst_path_}" 2>/dev/null + if [ "$?" != "0" ]; then + return 1 + else + return 0 + fi +} + +createOpapiLatestSoftlink() { + targetPkg=$2 + if [ "${targetPkg}x" = "x" ]; then + #CHANGED + targetPkg=ops_transformer + fi + + osName="" + if [ -f "$1/$targetPkg/scene.info" ]; then + . $1/$targetPkg/scene.info + osName=${os} + fi + opapi_lib_path="$1/opp/built-in/op_impl/ai_core/tbe/op_api/lib/${osName}/${architecture}" + opapi_include_level1_path="$1/opp/built-in/op_impl/ai_core/tbe/op_api/include/aclnnop" + opapi_include_level2_path="${opapi_include_level1_path}/level2" + if [ ! -d ${opapi_lib_path} ] || [ ! -d ${opapi_include_level1_path} ] || [ ! -d ${opapi_include_level2_path} ]; then + return 3 + fi + if [ -d $(dirname $1)/latest/${architectureDir}/lib64 ]; then + for file_so in $(ls -1 $1/${architectureDir}/lib64 | grep -E "libaclnn_|libopapi.so"); do + latest_arch_lib64_src_path="$1/${architectureDir}/lib64/${file_so}" + latest_arch_lib64_dst_path="$(dirname $1)/latest/${architectureDir}/lib64/${file_so}" + if [ -f $latest_arch_lib64_dst_path ] || [ -L $latest_arch_lib64_dst_path ]; then + rm -fr "$latest_arch_lib64_dst_path" + fi + createrelativelysoftlink ${latest_arch_lib64_src_path} ${latest_arch_lib64_dst_path} + done + fi + + # second the headfiles with 1 and 2 level + if [ -d $1/${architectureDir}/include/aclnnop ]; then + for file_level1 in $(ls -1 -F ${opapi_include_level1_path} | grep -v [/$] | sed 's/\*$//'); do + latest_arch_include_src_path="${opapi_include_level1_path}/${file_level1}" + latest_arch_include_dst_path="$(dirname $1)/latest/${architectureDir}/include/aclnnop/${file_level1}" + if [ -f $latest_arch_include_dst_path ] || [ -L $latest_arch_include_dst_path ]; then + rm -fr "$latest_arch_include_dst_path" + fi + createrelativelysoftlink ${latest_arch_include_src_path} ${latest_arch_include_dst_path} + done + fi + + if [ -d $1/${architectureDir}/include/aclnnop/level2 ]; then + for file_level2 in $(ls -1 -F ${opapi_include_level2_path} | grep -v [/$] | sed 's/\*$//'); do + latest_arch_include_src_path="${opapi_include_level2_path}/${file_level2}" + latest_arch_include_dst_path="$(dirname $1)/latest/${architectureDir}/include/aclnnop/level2/${file_level2}" + if [ -f $latest_arch_include_dst_path ] || [ -L $latest_arch_include_dst_path ]; then + rm -fr "$latest_arch_include_dst_path" + fi + createrelativelysoftlink ${latest_arch_include_src_path} ${latest_arch_include_dst_path} + done + fi +} + +createOpapiSoftlink() { + osName="" + if [ -f "$1/opp/scene.info" ]; then + . $1/opp/scene.info + osName=${os} + fi + opapi_lib_path="$1/opp/built-in/op_impl/ai_core/tbe/op_api/lib/${osName}/${architecture}" + opapi_include_level1_path="$1/opp/built-in/op_impl/ai_core/tbe/op_api/include/aclnnop" + opapi_include_level2_path="${opapi_include_level1_path}/level2" + + if [ ! -d ${opapi_lib_path} ] || [ ! -d ${opapi_include_level1_path} ] || [ ! -d ${opapi_include_level2_path} ]; then + return 3 + fi + # first the libopapi.so + if [ -d $1/${architectureDir}/lib64 ]; then + for file_so in $(ls -1 ${opapi_lib_path} | grep "so"$); do + arch_lib64_src_path="${opapi_lib_path}/${file_so}" + arch_lib64_dst_path="$1/${architectureDir}/lib64/${file_so}" + if [ -f $arch_lib64_dst_path ] || [ -L $arch_lib64_dst_path ]; then + rm -fr "$arch_lib64_dst_path" + fi + createrelativelysoftlink ${arch_lib64_src_path} ${arch_lib64_dst_path} + done + fi + + if [ -d $1/opp/lib64 ]; then + for file_so in $(ls -1 $1/${architectureDir}/lib64 | grep -E "libaclnn_|libopapi.so"); do + opp_lib64_src_path="$1/${architectureDir}/lib64/${file_so}" + opp_lib64_dst_path="$1/opp/lib64/${file_so}" + if [ -f $opp_lib64_dst_path ] || [ -L $opp_lib64_dst_path ]; then + rm -fr "$opp_lib64_dst_path" + fi + createrelativelysoftlink ${opp_lib64_src_path} ${opp_lib64_dst_path} + done + fi + + # second the headfiles with 1 and 2 level + if [ -d $1/${architectureDir}/include/aclnnop ]; then + for file_level1 in $(ls -1 -F ${opapi_include_level1_path} | grep -v [/$] | sed 's/\*$//'); do + arch_include_src_path="${opapi_include_level1_path}/${file_level1}" + arch_include_dst_path="$1/${architectureDir}/include/aclnnop/${file_level1}" + if [ -f $arch_include_dst_path ] || [ -L $arch_include_dst_path ]; then + rm -fr "$arch_include_dst_path" + fi + createrelativelysoftlink ${arch_include_src_path} ${arch_include_dst_path} + + opp_include_src_path="${arch_include_dst_path}" + opp_include_dst_path="$1/opp/include/aclnnop/${file_level1}" + if [ -f $opp_include_dst_path ] || [ -L $opp_include_dst_path ]; then + rm -fr "$opp_include_dst_path" + fi + createrelativelysoftlink ${opp_include_src_path} ${opp_include_dst_path} + done + fi + + if [ -d $1/${architectureDir}/include/aclnnop/level2 ]; then + for file_level2 in $(ls -1 -F ${opapi_include_level2_path} | grep -v [/$] | sed 's/\*$//'); do + arch_include_src_path="${opapi_include_level2_path}/${file_level2}" + arch_include_dst_path="$1/${architectureDir}/include/aclnnop/level2/${file_level2}" + if [ -f $arch_include_dst_path ] || [ -L $arch_include_dst_path ]; then + rm -fr "$arch_include_dst_path" + fi + createrelativelysoftlink ${arch_include_src_path} ${arch_include_dst_path} + + opp_include_src_path="${arch_include_dst_path}" + opp_include_dst_path="$1/opp/include/aclnnop/level2/${file_level2}" + if [ -f $opp_include_dst_path ] || [ -L $opp_include_dst_path ]; then + rm -fr "$opp_include_dst_path" + fi + createrelativelysoftlink ${opp_include_src_path} ${opp_include_dst_path} + done + fi +} + +# remove opapi soft link +removeopapisoftlink() { + local path="$1" + if [ -L "$1" ]; then + rm -fr ${path} + return 0 + else + return 1 + fi +} + +latestSoftlinksRemove() { + targetdir=$1 + osName="" + if [ -f "$targetdir/opp/scene.info" ]; then + . $targetdir/opp/scene.info + osName=${os} + fi + opapi_lib_path="$targetdir/opp/built-in/op_impl/ai_core/tbe/op_api/lib/${osName}/${architecture}" + opapi_include_level1_path="$1/opp/built-in/op_impl/ai_core/tbe/op_api/include/aclnnop" + opapi_include_level2_path="${opapi_include_level1_path}/level2" + + if [ -d $(dirname $targetdir)/latest/${architectureDir}/lib64 ]; then + for file_so in $(ls -l "$(dirname $targetdir)/latest/${architectureDir}/lib64/" | grep -E "libaclnn_|libopapi.so"); do + latest_arch_lib64_path="$(dirname $targetdir)/latest/${architectureDir}/lib64/${file_so}" + removeopapisoftlink ${latest_arch_lib64_path} + done + fi + + # second the headfiles with 1 and 2 level + if [ -d $(dirname $targetdir)/latest/${architectureDir}/include/aclnnop ]; then + for file_level1 in $(ls -1 -F ${opapi_include_level1_path} | grep -v [/$] | sed 's/\*$//'); do + latest_arch_include_path="$(dirname $targetdir)/latest/${architectureDir}/include/aclnnop/${file_level1}" + removeopapisoftlink ${latest_arch_include_path} + done + fi + + if [ -d $(dirname $targetdir)/latest/${architectureDir}/include/aclnnop/level2 ]; then + for file_level2 in $(ls -1 -F ${opapi_include_level2_path} | grep -v [/$] | sed 's/\*$//'); do + latest_arch_include_path="$(dirname $targetdir)/latest/${architectureDir}/include/aclnnop/level2/${file_level2}" + removeopapisoftlink ${latest_arch_include_path} + done + fi +} + +softlinksRemove() { + targetdir=$1 + osName="" + if [ -f "$targetdir/opp/scene.info" ]; then + . $targetdir/opp/scene.info + osName=${os} + fi + opapi_lib_path="$targetdir/opp/built-in/op_impl/ai_core/tbe/op_api/lib/${osName}/${architecture}" + opapi_include_level1_path="$targetdir/opp/built-in/op_impl/ai_core/tbe/op_api/include/aclnnop" + opapi_include_level2_path="${opapi_include_level1_path}/level2" + + # first the libopapi.so + if [ -d $targetdir/${architectureDir}/lib64 ]; then + for file_so in $(ls -1 $targetdir/${architectureDir}/lib64 | grep -E "libaclnn_|libopapi.so"); do + arch_lib64_path="$targetdir/${architectureDir}/lib64/${file_so}" + removeopapisoftlink ${arch_lib64_path} + done + fi + + if [ -d $targetdir/opp/lib64 ]; then + for file_so in $(ls -l $targetdir/opp/lib64 | grep -E "libaclnn_|libopapi.so"); do + opp_lib64_path="$targetdir/opp/lib64/${file_so}" + removeopapisoftlink ${opp_lib64_path} + done + fi + + # second the headfiles with 1 and 2 level + if [ -d $targetdir/${architectureDir}/include/aclnnop ]; then + for file_level1 in $(ls -1 -F ${opapi_include_level1_path} | grep -v [/$] | sed 's/\*$//'); do + arch_include_path="$targetdir/${architectureDir}/include/aclnnop/${file_level1}" + removeopapisoftlink ${arch_include_path} + + opp_include_path="$targetdir/opp/include/aclnnop/${file_level1}" + removeopapisoftlink ${opp_include_path} + done + fi + + if [ -d $targetdir/${architectureDir}/include/aclnnop/level2 ]; then + for file_level2 in $(ls -1 -F ${opapi_include_level2_path} | grep -v [/$] | sed 's/\*$//'); do + arch_include_path="$targetdir/${architectureDir}/include/aclnnop/level2/${file_level2}" + removeopapisoftlink ${arch_include_path} + + opp_include_path="$targetdir/opp/include/aclnnop/level2/${file_level2}" + removeopapisoftlink ${opp_include_path} + done + fi +} diff --git a/csrc/ascend/scripts/package/ops_transformer/scripts/opp_custom_install.sh b/csrc/ascend/scripts/package/ops_transformer/scripts/opp_custom_install.sh new file mode 100644 index 000000000..fc88405a2 --- /dev/null +++ b/csrc/ascend/scripts/package/ops_transformer/scripts/opp_custom_install.sh @@ -0,0 +1,58 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +curpath=$(dirname $(readlink -f "$0")) +SCENE_FILE="${curpath}""/../scene.info" +OPP_COMMON="${curpath}""/opp_common.sh" +common_func_path="${curpath}/common_func.inc" +. "${OPP_COMMON}" +. "${common_func_path}" +# init arch +architecture=$(uname -m) +architectureDir="${architecture}-linux" + +while true; do + case "$1" in + --install-path=*) + install_path=$(echo "$1" | cut -d"=" -f2-) + shift + ;; + --version-dir=*) + version_dir=$(echo "$1" | cut -d"=" -f2) + shift + ;; + --latest-dir=*) + latest_dir=$(echo "$1" | cut -d"=" -f2) + shift + ;; + -*) + shift + ;; + *) + break + ;; + esac +done +get_version_dir "opp_kernel_version_dir" "$install_path/$version_dir/opp_kernel/version.info" + +if [ -z "$opp_kernel_version_dir" ]; then + # create op_api soft link + logandprint "[INFO]: Start create opapi softlinks." + createOpapiSoftlink "${install_path}/${version_dir}" + return_code=$? + if [ ${return_code} -eq 0 ]; then + logandprint "[INFO]: Create opapi softlinks successfully!" + elif [ ${return_code} -eq 3 ]; then + logandprint "[WARNING]: opapi source file does not exist!" + else + logandprint "[ERROR]: Create opapi softlinks failed!" + fi +fi diff --git a/csrc/ascend/scripts/package/ops_transformer/scripts/opp_custom_uninstall.sh b/csrc/ascend/scripts/package/ops_transformer/scripts/opp_custom_uninstall.sh new file mode 100644 index 000000000..f86596131 --- /dev/null +++ b/csrc/ascend/scripts/package/ops_transformer/scripts/opp_custom_uninstall.sh @@ -0,0 +1,55 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +curpath=$(dirname $(readlink -f "$0")) +SCENE_FILE="${curpath}""/../scene.info" +OPP_COMMON="${curpath}""/opp_common.sh" +common_func_path="${curpath}/common_func.inc" +. "${OPP_COMMON}" +. "${common_func_path}" +# init arch +architecture=$(uname -m) +architectureDir="${architecture}-linux" + +while true; do + case "$1" in + --install-path=*) + install_path=$(echo "$1" | cut -d"=" -f2-) + shift + ;; + --version-dir=*) + version_dir=$(echo "$1" | cut -d"=" -f2) + shift + ;; + --latest-dir=*) + latest_dir=$(echo "$1" | cut -d"=" -f2) + shift + ;; + -*) + shift + ;; + *) + break + ;; + esac +done +get_version_dir "opp_kernel_version_dir" "$install_path/$version_dir/opp_kernel/version.info" + +if [ -z "$opp_kernel_version_dir" ]; then + # before remove the oppkernel, remove the softlinks + logandprint "[INFO]: Start remove opapi softlinks." + softlinksRemove ${install_path}/${version_dir} + if [ $? -ne 0 ]; then + logandprint "[WARNING]: Remove opapi softlinks failed, some softlinks may not exist." + else + logandprint "[INFO]: Remove opapi softlinks successfully." + fi +fi diff --git a/csrc/ascend/scripts/package/ops_transformer/scripts/opp_install.sh b/csrc/ascend/scripts/package/ops_transformer/scripts/opp_install.sh new file mode 100644 index 000000000..75562cd4a --- /dev/null +++ b/csrc/ascend/scripts/package/ops_transformer/scripts/opp_install.sh @@ -0,0 +1,414 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +PARAM_INVALID="0x0002" +INSTALL_FAILED="0x0000" +INSTALL_FAILED_DES="Update successfully." +FILE_NOT_EXIST="0x0080" +FILE_NOT_EXIST_DES="File not found." +FILE_READ_FAILED="0x0082" +FILE_READ_FAILED_DES="File read failed." +FILE_WRITE_FAILED="0x0081" +FILE_WRITE_FAILED_DES="File write failed." +PERM_DENIED="0x0093" +PERM_DENIED_DES="Permission denied." + +# run package's files info +CURR_PATH=$(dirname $(readlink -f $0)) +VERSION_INFO_FILE="${CURR_PATH}/../version.info" +FILELIST_FILE="${CURR_PATH}/filelist.csv" +COMMON_PARSER_FILE="${CURR_PATH}/install_common_parser.sh" +SCENE_FILE="${CURR_PATH}/../scene.info" +ASCEND_INSTALL_INFO="ascend_install.info" + +ARCH_INFO=$(uname -m) +OPP_PLATFORM_DIR=ops_transformer +OPP_PLATFORM_UPPER=$(echo "${OPP_PLATFORM_DIR}" | tr '[:lower:]' '[:upper:]') + +TARGET_INSTALL_PATH="" +TARGET_MOULDE_DIR="" # TARGET_INSTALL_PATH + PKG_VERSION_DIR + OPP_PLATFORM_DIR +TARGET_VERSION_DIR="" # TARGET_INSTALL_PATH + PKG_VERSION_DIR +TARGET_SHARED_INFO_DIR="" # TARGET_INSTALL_PATH + PKG_VERSION_DIR + share/info +TARGET_OPP_BUILT_IN="" + +COMMON_INC_FILE="${CURR_PATH}/common_func.inc" +COMMON_FUNC_V2_PATH="${CURR_PATH}/common_func_v2.inc" +VERSION_CFG="${CURR_PATH}/version_cfg.inc" +OPP_COMMON_FILE="${CURR_PATH}/opp_common.sh" + +. "${COMMON_INC_FILE}" +. "${COMMON_FUNC_V2_PATH}" +. "${VERSION_CFG}" +. "${OPP_COMMON_FILE}" + +# keys of infos in ascend_install.info +KEY_INSTALLED_UNAME="USERNAME" +KEY_INSTALLED_UGROUP="USERGROUP" +KEY_INSTALLED_TYPE="${OPP_PLATFORM_UPPER}_INSTALL_TYPE" +KEY_INSTALLED_FEATURE="${OPP_PLATFORM_UPPER}_INSTALL_FEATURE" +KEY_INSTALLED_CHIP="${OPP_PLATFORM_UPPER}_INSTALL_CHIP" +KEY_INSTALLED_PATH="${OPP_PLATFORM_UPPER}_INSTALL_PATH_VAL" +KEY_INSTALLED_VERSION="${OPP_PLATFORM_UPPER}_VERSION" + +get_opts() { + TARGET_INSTALL_PATH="$1" + TARGET_USERNAME="$2" + TARGET_USERGROUP="$3" + IN_FEATURE="$4" + INSTALL_TYPE="$5" + IS_FOR_ALL="$6" + IS_SETENV="$7" + IS_DOCKER_INSTALL="$8" + DOCKER_ROOT="$9" + PKG_VERSION_DIR="${10}" + + if [ "${TARGET_INSTALL_PATH}" = "" ] || [ "${TARGET_USERNAME}" = "" ] || + [ "${TARGET_USERGROUP}" = "" ] || [ "${INSTALL_TYPE}" = "" ]; then + logandprint "[ERROR]: ERR_NO:${PARAM_INVALID};ERR_DES:Empty parameters is invalid for install." + exit 1 + fi + + INSTALL_FOR_ALL="" + if [ "${IS_FOR_ALL}" = "y" ]; then + INSTALL_FOR_ALL="--install_for_all" + fi +} + +init_install_env() { + get_package_version "RUN_PKG_VERSION" "$VERSION_INFO_FILE" + if [ "${PKG_VERSION_DIR}" = "" ]; then + TARGET_VERSION_DIR=${TARGET_INSTALL_PATH} + else + TARGET_VERSION_DIR=${TARGET_INSTALL_PATH}/${PKG_VERSION_DIR} + fi + TARGET_MOULDE_DIR=${TARGET_VERSION_DIR}/${OPP_PLATFORM_DIR} + TARGET_OPP_BUILT_IN=${TARGET_VERSION_DIR}/opp/built-in + TARGET_SHARED_INFO_DIR=${TARGET_VERSION_DIR}/share/info + INSTALL_INFO_FILE=${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}/${ASCEND_INSTALL_INFO} + + if [ "$(id -u)" != "0" ]; then + LOG_PATH_PERM="740" + LOG_FILE_PERM="640" + INSTALL_INFO_PERM="600" + else + LOG_PATH_PERM="750" + LOG_FILE_PERM="640" + INSTALL_INFO_PERM="644" + fi + + if [ "${IS_FOR_ALL}" = "y" ]; then + BUILTIN_PERM="555" + CUSTOM_PERM="755" + CREATE_DIR_PERM="755" + ONLYREAD_PERM="444" + else + BUILTIN_PERM="550" + CUSTOM_PERM="750" + CREATE_DIR_PERM="750" + ONLYREAD_PERM="440" + fi +} + +log_with_errorlevel() { + local ret_status="$1" + local level="$2" + local msg="$3" + if [ "${ret_status}" != 0 ]; then + if [ "${level}" = "error" ]; then + logandprint "${msg}" + exit 1 + else + logandprint "${msg}" + fi + fi +} + +get_installed_info() { + local key="$1" + local res="" + if [ -f "${INSTALL_INFO_FILE}" ]; then + res=$(cat ${INSTALL_INFO_FILE} | grep "${key}" | awk -F = '{print $2}') + fi + echo "${res}" +} + +update_install_info() { + local key_val="$1" + local val="$2" + local old_val=$(get_installed_info "${key_val}") + if [ -f "${INSTALL_INFO_FILE}" ]; then + if [ "x${old_val}" = "x" ]; then + echo "${key_val}=${val}" >>"${INSTALL_INFO_FILE}" + else + sed -i "/${key_val}/c ${key_val}=${val}" "${INSTALL_INFO_FILE}" + fi + else + echo "${key_val}=${val}" >"${INSTALL_INFO_FILE}" + fi +} + +update_install_infos() { + local uname="$1" + local ugroup="$2" + local type="$3" + local path="$4" + local version + get_package_version "version" "$VERSION_INFO_FILE" + comm_create_file "${INSTALL_INFO_FILE}" "${INSTALL_INFO_PERM}" "${TARGET_USERNAME}:${TARGET_USERGROUP}" "${IS_FOR_ALL}" + + update_install_info "${KEY_INSTALLED_UNAME}" "${uname}" + update_install_info "${KEY_INSTALLED_UGROUP}" "${ugroup}" + update_install_info "${KEY_INSTALLED_TYPE}" "${type}" + update_install_info "${KEY_INSTALLED_PATH}" "${path}" + update_install_info "${KEY_INSTALLED_VERSION}" "${version}" +} + +check_file_exist() { + local path_param="${1}" + if [ ! -f "${path_param}" ]; then + logandprint "[ERROR]: ERR_NO:${FILE_NOT_EXIST};ERR_DES:The file (${path_param}) does not existed." + exit 1 + fi +} + +check_env() { + check_file_exist "${FILELIST_FILE}" + check_file_exist "${COMMON_PARSER_FILE}" +} + +createsoftlink() { + local src_path="$1" + local dst_path="$2" + if [ -e "$dst_path" ]; then + if [ -L "$dst_path" ]; then + `rm -f $dst_path` + else + return 0 + fi + fi + ln -s "${src_path}" "${dst_path}" 2>/dev/null + log_with_errorlevel "$?" "error" "[ERROR]: ERR_NO:${PERM_DENIED};ERR_DES:${src_path} Create softlink to ${dst_path} failed." +} + +get_install_path() { + docker_root_tmp="$(echo "${DOCKER_ROOT}" | sed "s#/\+\$##g")" + docker_root_regex="$(echo "${docker_root_tmp}" | sed "s#\/#\\\/#g")" + relative_path_val=$(echo "${TARGET_VERSION_DIR}" | sed "s/^${docker_root_regex}//g" | sed "s/\/\+\$//g") + return +} + +setenv() { + logandprint "[INFO]: Set the environment path [ export ASCEND_OPP_PATH=${relative_path_val}/opp ]." + if [ "${IS_DOCKER_INSTALL}" = y ]; then + INSTALL_OPTION="--docker-root=${DOCKER_ROOT}" + else + INSTALL_OPTION="" + fi + if [ "${IS_SETENV}" = "y" ]; then + INSTALL_OPTION="${INSTALL_OPTION} --setenv" + fi +} + +# 创建单个文件的软链接，链接文件级别 +create_file_softlink() { + local src_file=$1 + local dst_file=$2 + local base_dir=$(dirname ${dst_file}) + + comm_create_dir "${base_dir}" "${CREATE_DIR_PERM}" "${TARGET_USERNAME}:${TARGET_USERGROUP}" "${IS_FOR_ALL}" + + local relative_file_path=$(realpath -s --relative-to="${base_dir}" "${src_file}") + # 创建软连接 + createsoftlink "${relative_file_path}" "${dst_file}" +} + +# 创建单个目录的软连接，链接目录级别 +create_dir_softlink() { + local src_dir=$1 + local dst_dir=$2 + local base_dir=$(dirname ${dst_dir}) + + comm_create_dir "${base_dir}" "${CREATE_DIR_PERM}" "${TARGET_USERNAME}:${TARGET_USERGROUP}" "${IS_FOR_ALL}" + + if [ ! -d "${src_dir}" ]; then + logandprint "[ERROR]: src dir ["${src_dir}"] not exists to create soft link." + fi + + # 获取相对路径 + relative_dir_path=$(realpath -s --relative-to="${base_dir}" "${src_dir}") + # 创建软连接 + createsoftlink "${relative_dir_path}" "${dst_dir}" +} + +# 创建目录下子目录的软连接，链接子目录级别 +create_softlink_for_dirs() { + local src_dir=$1 + local dst_dir=$2 + + if [ ! -d "${src_dir}" ]; then + logandprint "[ERROR]: src dir ["${src_dir}"] not exists to create soft link." + exit 1 + fi + + comm_create_dir "${dst_dir}" "${CREATE_DIR_PERM}" "${TARGET_USERNAME}:${TARGET_USERGROUP}" "${IS_FOR_ALL}" + + find "${src_dir}" -mindepth 1 -maxdepth 1 -type d -print0 | while IFS= read -r -d '' src_dir_path; do + local sub_dir_name=$(basename ${src_dir_path}) + local dst_dir_path="${dst_dir}/${sub_dir_name}" + # 计算目录相对路径 + local relative_dir_path=$(realpath -s --relative-to="${dst_dir}" "${src_dir_path}") + # 创建软连接 + createsoftlink "${relative_dir_path}" "${dst_dir_path}" + done +} + +# 创建某目录下所有文件的软链接，链接文件级别，要求目录中不能有子目录 +create_softlink_for_files() { + local src_dir=$1 + local dst_dir=$2 + local exclude_list="$3" + + if [ ! -d "${src_dir}" ]; then + logandprint "[ERROR]: src dir ["${src_dir}"] not exists, cannot create soft link." + exit 1 + fi + + comm_create_dir "${dst_dir}" "${CREATE_DIR_PERM}" "${TARGET_USERNAME}:${TARGET_USERGROUP}" "${IS_FOR_ALL}" + + find "${src_dir}" -mindepth 1 -maxdepth 1 -type f -print0 | while IFS= read -r -d '' src_file_path; do + # 文件名 + local file_name=$(basename ${src_file_path}) + if $(echo ${exclude_list} | grep -wq ${file_name}); then + continue + fi + local dst_file_path="${dst_dir}/${file_name}" + # 获取相对路径 + local relative_file_path=$(realpath --relative-to="${dst_dir}" "${src_file_path}") + # 创建软连接 + createsoftlink "${relative_file_path}" "${dst_file_path}" + done +} + +# 递归创建目录下所有文件的软链接，链接文件级别，目录中可以有子目录，对于子目录会创建对应目录不是链接 +create_softlink_for_files_and_dirs() { + local src_dir=$1 + local dst_dir=$2 + + if [ ! -d "${src_dir}" ]; then + logandprint "[ERROR]: src dir ["${src_dir}"] not exists, cannot create soft link." + exit 1 + fi + + comm_create_dir "${dst_dir}" "${CREATE_DIR_PERM}" "${TARGET_USERNAME}:${TARGET_USERGROUP}" "${IS_FOR_ALL}" + find "${src_dir}" -mindepth 1 -maxdepth 1 -type d -print0 | while IFS= read -r -d '' src_dir_path; do + local base_dir=$(basename ${src_dir_path}) + local dst_dir_path="${dst_dir}/${base_dir}" + create_softlink_for_files_and_dirs ${src_dir_path} ${dst_dir_path} + done + + create_softlink_for_files ${src_dir} ${dst_dir} +} + +add_init_py() { + local opp_builtin_mod="" + local built_in_impl_path=${TARGET_OPP_BUILT_IN}/op_impl/ai_core/tbe/impl/ops_transformer + if [ -d ${built_in_impl_path} ]; then + opp_builtin_mod=$(stat -c %a ${built_in_impl_path}) + if [ "$(id -u)" != 0 ] && [ ! -w "${built_in_impl_path}" ]; then + chmod u+w -R "${built_in_impl_path}" 2>/dev/null + fi + fi + touch ${built_in_impl_path}/__init__.py + + [ -d ${built_in_impl_path}/dynamic ] && touch ${built_in_impl_path}/dynamic/__init__.py + + if [ -n "${opp_builtin_mod}" ]; then + chmod ${opp_builtin_mod} -R "${built_in_impl_path}" 2>/dev/null + fi +} + +install_opp() { + logandprint "[INFO]: Begin install opp module." + comm_create_dir "${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}" "${CREATE_DIR_PERM}" "${TARGET_USERNAME}:${TARGET_USERGROUP}" "${IS_FOR_ALL}" + + setenv + + logandprint "[INFO]: Update the opp install info." + + update_install_infos "${TARGET_USERNAME}" "${TARGET_USERGROUP}" "${INSTALL_TYPE}" "${relative_path_val}" + log_with_errorlevel "$?" "error" "[ERROR]: ERR_NO:${INSTALL_FAILED};ERR_DES:Update opp install info failed." + + bash "${COMMON_PARSER_FILE}" --package="${OPP_PLATFORM_DIR}" --install --username="${TARGET_USERNAME}" \ + --usergroup="${TARGET_USERGROUP}" --set-cann-uninstall --version=$RUN_PKG_VERSION \ + --use-share-info --version-dir=$PKG_VERSION_DIR $INSTALL_OPTION ${INSTALL_FOR_ALL} "--feature=all" "--chip=all" \ + "${INSTALL_TYPE}" "${TARGET_INSTALL_PATH}" "${FILELIST_FILE}" + log_with_errorlevel "$?" "error" "[ERROR]: ERR_NO:${INSTALL_FAILED};ERR_DES:Install opp module files failed." + + logandprint "[INFO]: upgradePercentage:30%" + + add_init_py + + logandprint "[INFO]: upgradePercentage:50%" +} + +main() { + logandprint "[INFO]: Command opp_install" + + get_opts "$@" + + init_install_env + + get_package_upgrade_version_dir "upgrade_version_dir" "$TARGET_INSTALL_PATH" "${OPP_PLATFORM_DIR}" + get_package_last_installed_version "last_installed" "$TARGET_INSTALL_PATH" "${OPP_PLATFORM_DIR}" + last_installed_version=$(echo ${last_installed} | cut --only-delimited -d":" -f2-) + + get_install_path + + check_env + + install_opp + + # change log dir and file owner and rights + chmod "${LOG_PATH_PERM}" "${COMM_LOG_DIR}" 2>/dev/null + chmod "${LOG_FILE_PERM}" "${COMM_LOGFILE}" 2>/dev/null + chmod "${LOG_FILE_PERM}" "${COMM_OPERATION_LOGFILE}" 2>/dev/null + + if [ "$(id -u)" = "0" ]; then + chmod "${CUSTOM_PERM}" -R "${TARGET_OPP_BUILT_IN}" 2>/dev/null + else + chmod "${BUILTIN_PERM}" -R "${TARGET_OPP_BUILT_IN}" 2>/dev/null + fi + + chmod "${ONLYREAD_PERM}" "${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}/scene.info" 2>/dev/null + chmod "${ONLYREAD_PERM}" "${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}/version.info" 2>/dev/null + chmod "${ONLYREAD_PERM}" "${INSTALL_INFO_FILE}" 2>/dev/null + + # change installed folder's owner and group except aicpu + log_with_errorlevel "$?" "error" "[ERROR]: ERR_NO:${INSTALL_FAILED};ERR_DES:Change opp ownership failed.." + + logandprint "[INFO]: upgradePercentage:100%" + + logandprint "[INFO]: Installation information listed below:" + logandprint "[INFO]: Install path: (${TARGET_VERSION_DIR}/opp)" + logandprint "[INFO]: Install log file path: (${COMM_LOGFILE})" + logandprint "[INFO]: Operation log file path: (${COMM_OPERATION_LOGFILE})" + + if [ "${IS_SETENV}" != "y" ]; then + logandprint "[INFO]: Using requirements: when opp module install finished or \ +before you run the opp module, execute the command \ +[ export ASCEND_OPP_PATH=${TARGET_INSTALL_PATH}/cann/opp ] to set the environment path." + fi + + logandprint "[INFO]: Opp package installed successfully! The new version takes effect immediately." +} + +main "$@" +exit 0 diff --git a/csrc/ascend/scripts/package/ops_transformer/scripts/opp_uninstall.sh b/csrc/ascend/scripts/package/ops_transformer/scripts/opp_uninstall.sh new file mode 100644 index 000000000..04370789e --- /dev/null +++ b/csrc/ascend/scripts/package/ops_transformer/scripts/opp_uninstall.sh @@ -0,0 +1,225 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +OPERATE_FAILED="0x0001" +PARAM_INVALID="0x0002" +PARAM_INVALID_DES="Invalid input parameter." +FILE_NOT_EXIST="0x0080" +FILE_NOT_EXIST_DES="File not found." +FILE_READ_FAILED="0x0082" +FILE_READ_FAILED_DES="File read failed." + +CURR_PATH=$(dirname $(readlink -f $0)) +COMMON_INC_FILE="${CURR_PATH}/common_func.inc" +OPP_COMMON_FILE="${CURR_PATH}/opp_common.sh" + +. "${COMMON_INC_FILE}" +. "${OPP_COMMON_FILE}" + +ARCH_INFO=$(uname -m) +OPP_PLATFORM_DIR=ops_transformer +OPP_PLATFORM_UPPER=$(echo "${OPP_PLATFORM_DIR}" | tr '[:lower:]' '[:upper:]') +FILELIST_FILE="${CURR_PATH}/filelist.csv" +COMMON_PARSER_FILE="${CURR_PATH}/install_common_parser.sh" + +TARGET_INSTALL_PATH="" +TARGET_VERSION_DIR="${CURR_PATH}/../../../.." +TARGET_VERSION_DIR=$(readlink -f ${TARGET_VERSION_DIR}) # TARGET_INSTALL_PATH + PKG_VERSION_DIR +TARGET_MOULDE_DIR=${TARGET_VERSION_DIR}/${OPP_PLATFORM_DIR} # TARGET_INSTALL_PATH + PKG_VERSION_DIR + OPP_PLATFORM_DIR +TARGET_OPP_BUILT_IN=${TARGET_VERSION_DIR}/opp/built-in +TARGET_SHARED_INFO_DIR=${TARGET_VERSION_DIR}/share/info + +ASCEND_INSTALL_INFO="ascend_install.info" +# init log file path +INSTALL_INFO_FILE="${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}/${ASCEND_INSTALL_INFO}" + +VERSION_INFO_FILE="${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}/version.info" + +# keys of infos in ascend_install.info +KEY_INSTALLED_UNAME="USERNAME" +KEY_INSTALLED_UGROUP="USERGROUP" +KEY_INSTALLED_TYPE="${OPP_PLATFORM_UPPER}_INSTALL_TYPE" +KEY_INSTALLED_FEATURE="${OPP_PLATFORM_UPPER}_INSTALL_FEATURE" +KEY_INSTALLED_PATH="${OPP_PLATFORM_UPPER}_INSTALL_PATH_VAL" +KEY_INSTALLED_VERSION="${OPP_PLATFORM_UPPER}_VERSION" + +get_opts() { + INSTALLED_PATH="$1" + UNINSTALL_MODE="$2" + IS_QUIET="$3" + IN_FEATURE="$4" + IS_DOCKER_INSTALL="$5" + DOCKER_ROOT="$6" + PKG_VERSION_DIR="$7" + local paramter_num="$#" + + if [ "${paramter_num}" != 0 ]; then + if [ "${INSTALLED_PATH}" = "" ] || + [ "${UNINSTALL_MODE}" = "" ] || + [ "${IS_QUIET}" = "" ]; then + logandprint "[ERROR]: ERR_NO:${PARAM_INVALID};ERR_DES:Empty parameters is invalid\ +for call uninstall functions." + exit 1 + fi + fi +} + +get_docker_install_path() { + local docker_root_tmp="$(echo "${DOCKER_ROOT}" | sed "s#/\+\$##g")" + local docker_root_regex="$(echo "${docker_root_tmp}" | sed "s#\/#\\\/#g")" + relative_path_val=$(echo "${TARGET_VERSION_DIR}" | sed "s/^${docker_root_regex}//g" | sed "s/\/\+\$//g") + return +} + +log_with_errorlevel() { + local ret_status="$1" + local level="$2" + local msg="$3" + if [ "${ret_status}" != 0 ]; then + if [ "${level}" = "error" ]; then + logandprint "${msg}" + exit 1 + else + logandprint "${msg}" + fi + fi +} + +check_directory_exist() { + local path="${1}" + if [ ! -d "${path}" ]; then + logandprint "[ERROR]: ERR_NO:${FILE_NOT_EXIST};ERR_DES:Installation directory [${path}] does not exist, uninstall failed." + exit 1 + fi +} + +check_file_exist() { + local path_param="${1}" + if [ ! -f "${path_param}" ]; then + logandprint "[ERROR]: ERR_NO:${FILE_NOT_EXIST};ERR_DES:The file (${path_param}) does not existed." + exit 1 + fi +} + +check_installed_files() { + # check install folder existed + check_file_exist "${INSTALL_INFO_FILE}" + + check_file_exist "${FILELIST_FILE}" + + check_file_exist "${COMMON_PARSER_FILE}" + +} + +check_installed_type() { + local type="$1" + if [ "${type}" != "run" ] && + [ "${type}" != "full" ] && + [ "${type}" != "devel" ]; then + logandprint "[ERROR]: ERR_NO:${UNAME_NOT_EXIST};ERR_DES:Install type of opp module is not right!" + exit 1 + fi +} + +unsetenv() { + logandprint "[INFO]: Unset the environment path [ export ASCEND_OPP_PATH=${relative_path_val}/opp ]." + if [ "${IS_DOCKER_INSTALL}" = y ]; then + UNINSTALL_OPTION="--docker-root=${DOCKER_ROOT}" + else + UNINSTALL_OPTION="" + fi +} + +get_installed_info() { + local key="$1" + local res="" + if [ -f "${INSTALL_INFO_FILE}" ]; then + chmod 644 "${INSTALL_INFO_FILE}" >/dev/null 2>&1 + res=$(cat ${INSTALL_INFO_FILE} | grep "${key}" | awk -F = '{print $2}') + fi + echo "${res}" +} + +get_installed_param() { + INSTALLED_TYPE=$(get_installed_info "${KEY_INSTALLED_TYPE}") + TARGET_USERNAME=$(get_installed_info "${KEY_INSTALLED_UNAME}") + TARGET_USERGROUP=$(get_installed_info "${KEY_INSTALLED_UGROUP}") + get_package_version "RUN_PKG_VERSION" "$VERSION_INFO_FILE" + if [ "${PKG_VERSION_DIR}" = "" ]; then + TARGET_INSTALL_PATH=${TARGET_VERSION_DIR} + else + TARGET_INSTALL_PATH=$(readlink -f "${TARGET_VERSION_DIR}/../") + fi +} + +remove_module() { + chmod u+w ${TARGET_SHARED_INFO_DIR}/${OPP_PLATFORM_DIR}/scene.info + + logandprint "[INFO]: Delete the installed opp source files in (${TARGET_VERSION_DIR})." + + bash "${COMMON_PARSER_FILE}" --package="${OPP_PLATFORM_DIR}" --uninstall --remove-install-info \ + --username="${TARGET_USERNAME}" --usergroup="${TARGET_USERGROUP}" --version=$RUN_PKG_VERSION \ + --use-share-info --version-dir=$PKG_VERSION_DIR ${UNINSTALL_OPTION} "${INSTALLED_TYPE}" "${TARGET_INSTALL_PATH}" \ + "${FILELIST_FILE}" "${IN_FEATURE}" --recreate-softlink + log_with_errorlevel "$?" "error" "[ERROR]: ERR_NO:${OPERATE_FAILED};ERR_DES:Uninstall opp module failed." +} + +remove_init_py() { + local built_in_impl_path=${TARGET_OPP_BUILT_IN}/op_impl/ai_core/tbe/impl/ops_transformer + + [ -e ${built_in_impl_path}/__init__.py ] && rm ${built_in_impl_path}/__init__.py > /dev/null 2>&1 + + [ -e ${built_in_impl_path}/dynamic/__init__.py ] && rm ${built_in_impl_path}/dynamic/__init__.py > /dev/null 2>&1 +} + +remove_ops_transformer() { + if [ "$(id -u)" != 0 ] && [ ! -w "${TARGET_OPP_BUILT_IN}" ]; then + chmod u+w -R "${TARGET_OPP_BUILT_IN}" 2>/dev/null + fi + + remove_init_py + + remove_module + + if [ "${UNINSTALL_MODE}" != "upgrade" ]; then + logandprint "[INFO]: Delete the install info file (${INSTALL_INFO_FILE})." + rm -f "${INSTALL_INFO_FILE}" + log_with_errorlevel "$?" "warn" "[WARNING] Delete ops install info file failed, please delete it by yourself." + fi +} + +logandprint "[INFO]: Begin uninstall the opp module." + +main() { + get_opts "$@" + + get_docker_install_path + + check_installed_files + + get_installed_param + + check_installed_type "${INSTALLED_TYPE}" + + unsetenv + + remove_ops_transformer + + if [ "${UNINSTALL_MODE}" != "upgrade" ]; then + remove_dir_if_empty ${TARGET_VERSION_DIR} + fi + remove_dir_if_empty ${INSTALLED_PATH} + + logandprint "[INFO]: Opp package uninstalled successfully! Uninstallation takes effect immediately." +} + +main "$@" +exit 0 diff --git a/csrc/ascend/scripts/package/ops_transformer/scripts/uninstall.sh b/csrc/ascend/scripts/package/ops_transformer/scripts/uninstall.sh new file mode 100644 index 000000000..cc63f67d1 --- /dev/null +++ b/csrc/ascend/scripts/package/ops_transformer/scripts/uninstall.sh @@ -0,0 +1,84 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +CURR_PATH=$(dirname $(readlink -f $0)) + +# error number and description +FILE_NOT_EXIST="0x0080" +PERM_DENIED="0x0093" +PERM_DENIED_DES="Permission denied." +# log functions +getdate() { + _cur_date=$(date +"%Y-%m-%d %H:%M:%S") + echo "${_cur_date}" +} + +logandprint() { + is_error_level=$(echo $1 | grep -E 'ERROR|WARN|INFO') + if [ "${is_quiet}" != "y" ] || [ "${is_error_level}" != "" ]; then + echo "[OpsTransformer] [$(getdate)] ""$1" + fi + echo "[OpsTransformer] [$(getdate)] ""$1" >> "${_INSTALL_LOG_FILE}" +} + +if [ "$(id -u)" != "0" ]; then + _LOG_PATH=$(echo "${HOME}")"/var/log/ascend_seclog" + _INSTALL_LOG_FILE="${_LOG_PATH}/ascend_install.log" +else + _LOG_PATH="/var/log/ascend_seclog" + _INSTALL_LOG_FILE="${_LOG_PATH}/ascend_install.log" +fi + +# init install cmd status, set default as n +is_quiet=n +quiet_parameter="" +if [ "$#" != "0" ]; then + if [ "$1" = "--quiet" ] && [ "$#" = "1" ]; then + is_quiet=y + quiet_parameter="--quiet" + else + logandprint "Please use correct parameters, only support input nothing or only --quiet parameter." + exit 1 + fi +fi + +install_shell="${CURR_PATH}/install.sh" + +# shell exist check +if [ ! -f "${install_shell}" ]; then + logandprint "[ERROR]: ERR_NO:${FILE_NOT_EXIST};OpsTransformer module is not installed or some ops_transformer source files are lost.\ +If there are any residual files, please manually remove those files." + exit 1 +fi + +# shell execute perm check +if [ ! -x "${install_shell}" ]; then + logandprint "[ERROR]: ERR_NO:${PERM_DENIED};ERR_DES:The user do \ +not have the permission to execute this file, please reset the file \ +to a right permission." + exit 1 +fi + +installed_path="$(cd "${CURR_PATH}/../../../../"; pwd)" +parent_installed_path="$(cd "${installed_path}/../"; pwd)" +cd ~ +sh "${install_shell}" "--aa" "--aa" "--uninstall" "--install-path=${installed_path}" "${quiet_parameter}" +ret_status="$?" +if [ "${ret_status}" != "0" ]; then + exit 1 +fi +if [ -d "${parent_installed_path}" ];then + subdirs_param_install=$(ls "${parent_installed_path}" 2> /dev/null) + if [ "${subdirs_param_install}" = "" ]; then + rm -rf "${parent_installed_path}" + fi +fi +exit 0 diff --git a/csrc/ascend/scripts/package/ops_transformer/scripts/ver_check.sh b/csrc/ascend/scripts/package/ops_transformer/scripts/ver_check.sh new file mode 100644 index 000000000..2e50d08ed --- /dev/null +++ b/csrc/ascend/scripts/package/ops_transformer/scripts/ver_check.sh @@ -0,0 +1,173 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +req_ver_path=$1 +_CURR_PATH=$(dirname $(readlink -f $0)) +_DEFAULT_INSTALL_PATH="/usr/local/Ascend" +FILE_NOT_EXIST="0x0080" + +getdate() { + _cur_date=$(date +"%Y-%m-%d %H:%M:%S") + echo "${_cur_date}" +} + +logandprint() { + echo "[OpsTransformer] [$(getdate)] ""$1" +} + +check_path_pre() { + in_checkpath_0="$1" + in_checkpath_1=$(echo ${in_checkpath_0} | cut -d"=" -f2) + if [ "${in_checkpath_1}" = "" ]; then + logandprint "[WARNING]: please input correct path" + exit 1 + fi + arr=$(echo ${in_checkpath_1} | awk '{split($0,arr," ");for(i in arr) print arr[i]}') + index=0 + for i in $arr; do + id=$((${id:=-1} + 1)) + eval arr_$id=$i + index=$(expr $index + 1) + done + len="${index}" + b=0 + for i in $(seq 0 ${len}); do + select_last_dir_component "$(eval echo '$'arr_$i)" + ret=$last_component + if [ "${ret}" != "" ]; then + eval checked_path_temp_$b=$(eval echo '$'arr_$i) + b=$(expr $b + 1) + fi + done + check_all_path="" + for i in $(seq 0 ${len}); do + check_all_path="$(eval echo '$'checked_path_temp_$i) $check_all_path" + done + checked_path="${check_all_path}" + return +} + +select_last_dir_component() { + path="$1" + last_component=$(basename "${path}") + if [ "${last_component}" = "atc" ]; then + last_component="atc" + return + elif [ "${last_component}" = "fwkacllib" ]; then + last_component="fwkacllib" + return + elif [ "${last_component}" = "compiler" ]; then + last_component="compiler" + return + elif [ "${last_component}" = "fwkplugin" ]; then + last_component="fwkplugin" + return + else + last_component="atc or fwkacllib or compiler" + return + fi +} + +check_version_file() { + pkg_path="$1" + component_ret="$2" + run_pkg_path_temp=$(dirname "${pkg_path}") + run_pkg_path="${run_pkg_path_temp}""/${component_ret}" + version_file="${run_pkg_path}""/version.info" + if [ -f "${version_file}" ]; then + echo "${version_file}" 2 >>/dev/null + else + logandprint "[ERROR]: ERR_NO:${FILE_NOT_EXIST}; The [${component_ret}] version.info in path [${pkg_path}] not exists." + exit 1 + fi + return +} + +check_opp_version_file() { + if [ -f "${_CURR_PATH}/../../version.info" ]; then + ver_info="${_CURR_PATH}/../../version.info" + # ops_transformer/version.info -> ops_transformer + elif [ -f "${_DEFAULT_INSTALL_PATH}/ops_transformer/version.info" ]; then + ver_info="${_DEFAULT_INSTALL_PATH}/ops_transformer/version.info" + else + logandprint "[ERROR]: ERR_NO:${FILE_NOT_EXIST}; The [ops_transformer] version.info not exists." + fi + return +} + +check_relation() { + opp_ver_info="$1" + req_pkg_name="$2" + req_pkg_version="$3" + _COMMON_INC_FILE="${_CURR_PATH}/common_func.inc" + if [ -f "${_COMMON_INC_FILE}" ]; then + . "${_COMMON_INC_FILE}" + check_pkg_ver_deps "${opp_ver_info}" "${req_pkg_name}" "${req_pkg_version}" + ret_situation=$ver_check_status + else + logandprint "[ERROR]: ERR_NO:${FILE_NOT_EXIST}; The ${_COMMON_INC_FILE} not exists." + fi + return +} + +show_relation() { + relation_situation="$1" + req_pkg_name_val="$2" + req_pkg_path="$3" + if [ "$relation_situation" = "SUCC" ]; then + logandprint "[INFO]: Relationship of ops_transformer with ${req_pkg_name_val} in path ${req_pkg_path} check successfully" + return 0 + else + logandprint "[WARNING]: Relationship of ops_transformer with ${req_pkg_name_val} in path ${req_pkg_path} check failed. \ +do you want to continue. [y/n] " + while true; do + read yn + if [ "$yn" == "n" ]; then + echo "stop check!" + exit 1 + elif [ "$yn" = y ]; then + break + else + echo "[WARNING]: Input error, please input y or n to choose!" + fi + done + fi +} + +version_check() { + path_val="$1" + #get ops_transformer version + check_opp_version_file + ret_check_opp_version_file=$ver_info + #get checked path + check_path_pre "${path_val}" + ret_check_path_pre=$checked_path + if [ "${ret_check_path_pre}" != "" ]; then + for var in ${ret_check_path_pre}; do + # select_last_dir_component "${var}" + # component_ret=$last_component + #get atc or fwkacllib name + select_last_dir_component "${var}" + ret_last_component=$last_component + #get the version of atc/fwkacllib + check_version_file "${var}" "${ret_last_component}" + ret_check_version_file=$version_file + #check relation + check_relation "${ret_check_opp_version_file}" "${ret_last_component}" "${ret_check_version_file}" + ret_check_relation=$ret_situation + #show relation + show_relation "${ret_check_relation}" "${ret_last_component}" "${var}" + done + fi +} + +version_check "${req_ver_path}" +exit 0 diff --git a/csrc/ascend/scripts/package/package.py b/csrc/ascend/scripts/package/package.py new file mode 100644 index 000000000..c3eb8a65d --- /dev/null +++ b/csrc/ascend/scripts/package/package.py @@ -0,0 +1,823 @@ +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import argparse +import csv +import os +import shutil +import sys +import traceback +from argparse import Namespace +from collections import namedtuple +from collections.abc import Iterator +from datetime import datetime, timezone +from functools import partial +from itertools import chain +from typing import TextIO + +import regex as re +from common.py.filelist import ( + FileItem, + FileList, + check_filelist, + create_file_item, + generate_filelist, + get_transform_nested_path_func, +) +from common.py.packer import PackageName, create_makeself_pkg_params_factory, create_run_package_command, exec_pack_cmd +from common.py.pkg_parser import ParseOption, XmlConfig, get_cann_version_info, parse_xml_config +from common.py.utils import pkg_utils +from common.py.utils.comm_log import CommLog +from common.py.utils.funcbase import invoke, pipe +from common.py.utils.pkg_utils import ( + CONFIG_SCRIPT_PATH, + DELIVERY_PATH, + FAIL, + SUCCESS, + TOP_DIR, + CompressError, + ContainAsteriskError, + FilelistError, + GenerateFilelistError, + PackageNameEmptyError, + UnknownOperateTypeError, + path_join, +) + + +def get_comments(package_name: PackageName) -> str: + """获取run包注释。""" + comments = "_".join([package_name.chip_name.upper(), package_name.func_name.upper(), "RUN_PACKAGE"]) + return f'"{comments}"' + + +def get_compress_cmd(delivery_dir: str, pkg_args: Namespace, xml_config: XmlConfig) -> str: + """获取makeself压缩命令""" + suffix = xml_config.package_attr.get("suffix") + if suffix == "run": + package_name = PackageName(xml_config.package_attr, pkg_args, xml_config.version) + factory = create_makeself_pkg_params_factory( + pkg_args.pkg_output_dir, package_name.getvalue(), get_comments(package_name) + ) + params = factory(pkg_args.makeself_dir, xml_config.package_attr, pkg_args.independent_pkg) + pack_cmd, err_msg = create_run_package_command(params) + if err_msg: + CommLog.cilog_error(err_msg) + CommLog.cilog_error("create_run_command failed!") + raise CompressError(package_name.getvalue()) + if pkg_args.independent_pkg: + exec_pack_cmd(delivery_dir, pack_cmd, package_name.getvalue()) + else: + CommLog.cilog_error("the repack type '%s' is not support!", suffix) + sys.exit(FAIL) + try: + makeself_dir = os.path.join(TOP_DIR, "build/makeself.txt") + with open(makeself_dir, "w") as f: + f.write(pack_cmd) + except Exception as exception: + CommLog.cilog_error(f"save makeself.txt failed!{str(exception)}") + sys.exit(FAIL) + return package_name.getvalue() + + +def make_parse_option(args_: argparse.Namespace) -> ParseOption: + """创建解析参数。""" + + return ParseOption(args_.os_arch, args_.pkg_version, args_.build_type, args_.package_check, args_.ext_name) + + +PrivatePackageOption = namedtuple( + "PrivatePackageOption", + [ + "os_arch", + "package_suffix", + "not_in_name", + "pkg_version", + "ext_name", + "chip_name", + "func_name", + "version_dir", + "disable_multi_version", + "suffix", + ], +) + + +class PackageOption(PrivatePackageOption): + """打包配置参数。""" + + __slots__ = () # 优化内存，避免创建 __dict__ + + def __new__(cls, *package_option_args, **kwargs): + return super().__new__(cls, *package_option_args, **kwargs) + + +def generate_info_content(target_conf, ext_name) -> list[str]: + """生成info内容。""" + + def toolchain_llvm_config() -> Iterator[tuple[str, str]]: + if "llvm" in ext_name: + yield "toolchain", "llvm" + + content_list = [f"{key}={value}" for key, value in chain(target_conf["content"].items(), toolchain_llvm_config())] + return content_list + + +def generate_version_header_content(target_conf) -> Iterator[str]: + """生成version_header内容。""" + guard_name = target_conf["value"].replace(".", "_").upper() + yield f"#ifndef {guard_name}" + yield f"#define {guard_name}" + yield "" + for name, value in target_conf["content"].items(): + if name.endswith("_VERSION"): + version_infos = get_cann_version_info(name, value) + for version_name, version_value in version_infos: + yield f"#define {version_name} {version_value}" + else: + yield f"#define {name} {value}" + yield "" + yield f"#endif /* {guard_name} */" + yield "" + + +def generate_customized_file(target_conf, ext_name): + filepath = os.path.join(TOP_DIR, "build", target_conf.get("value")) + + generator = target_conf.get("generator", "info") + if generator == "version_header": + content_list = generate_version_header_content(target_conf) + else: + content_list = generate_info_content(target_conf, ext_name) + + file_content = "\n".join(content_list) + try: + with open(filepath, "w") as file: + file.write(file_content) + except Exception as ex: + CommLog.cilog_error(f"generate customized file {filepath} failed: {ex}!") + return FAIL + + return SUCCESS + + +def get_module(target_config) -> str: + """获取配置模块。""" + module = target_config.get("module", "NA") + return module if module else "NA" + + +def get_operation(operation, target_config) -> str: + """获取操作类型。""" + if operation in ("copy", "move") and target_config.get("entity") == "true": + return "copy_entity" + return operation + + +def get_permission(target_config) -> str: + """获取配置权限。""" + return target_config.get("install_mod", "NA") + + +def get_owner_group(target_config) -> str: + """获取配置属主。""" + # install_own的可能值为$username:$usergroup + # 防止变量在install_common_parser.sh中，被eval展开，添加\转义$ + # 由于awk会消耗1个\，所以需要2个转义符 + return target_config.get("install_own", "NA").replace("$", "\\\\$") + + +def get_install_type(target_config) -> str: + """获取安装类型。""" + return target_config.get("install_type", "NA") + + +def get_softlink(target_config) -> list[str]: + """获取配置软链。""" + softlink_str = target_config.get("install_softlink") + if not softlink_str: + return [] + return softlink_str.split(";") + + +def get_feature(target_config) -> set[str]: + """获取配置特性。""" + return target_config["feature"] + + +def get_chip(target_config) -> set[str]: + """获取配置芯片。""" + return target_config["chip"] + + +def get_configurable(target_config) -> str: + """获取配置是否为配置文件。""" + return target_config.get("configurable", "FALSE") + + +def get_hash_value(target_config) -> str: + """获取配置哈希值。""" + return target_config.get("hash", "NA") + + +def get_block(target_config) -> str: + """获取配置块信息。""" + return target_config.get("name", "NA") + + +def get_pkg_inner_softlink(target_config) -> list[str]: + """获取配置包内软链。""" + softlink_str = target_config.get("pkg_inner_softlink") + if not softlink_str: + return [] + return softlink_str.split(";") + + +def parse_install_info(infos: list, operate_type, filter_key) -> Iterator[FileItem]: + """根据配置解析生成安装信息。""" + for target_config in infos: + target_name = get_target_name(target_config) + if target_config.get("optional") == "true" and operate_type in ("copy", "move"): + path = os.path.join(TOP_DIR, DELIVERY_PATH, target_config.get("dst_path")) + value = os.path.join(TOP_DIR, DELIVERY_PATH, target_config.get("dst_path"), target_name) + if not os.path.exists(path): + continue + if not os.path.exists(value): + continue + if operate_type in ("copy", "move"): + relative_path_in_pkg = os.path.join(target_config.get("dst_path"), target_name) + relative_install_path = path_join(target_config.get("install_path"), target_name) + is_dir = target_config.get("is_dir", False) + elif operate_type == "mkdir": + relative_path_in_pkg = "NA" + relative_install_path = target_config.get("value") + is_dir = False + elif operate_type == "del": + relative_path_in_pkg = "NA" + relative_install_path = path_join(target_config.get("install_path"), target_name) + is_dir = False + else: + raise UnknownOperateTypeError(f"unknown operate type {operate_type}") + + if relative_install_path is None: + continue + + install_type = get_install_type(target_config) + if any(key in install_type for key in filter_key): + is_in_docker = "TRUE" + else: + is_in_docker = "FALSE" + + file_item = create_file_item( + get_module(target_config), + get_operation(operate_type, target_config), + relative_path_in_pkg, + relative_install_path, + is_in_docker, + get_permission(target_config), + get_owner_group(target_config), + install_type, + get_softlink(target_config), + get_feature(target_config), + "N", + get_configurable(target_config), + get_hash_value(target_config), + get_block(target_config), + get_pkg_inner_softlink(target_config), + get_chip(target_config), + is_dir, + ) + + yield file_item + + +def execute_repack_process( + xml_config: XmlConfig, + delivery_dir: str, + pkg_args: Namespace, + package_name: PackageName = None, + package_option: PackageOption = None, +): + """ + 功能描述: 执行打包流程(拷贝--->签名--->打包) + 返回值: SUCCESS/FAIL + """ + release_dir = os.path.join(delivery_dir, xml_config.default_config.get("name", "default")) + # 生成自定义文件 + for item in xml_config.generate_infos: + if generate_customized_file(item, package_option.ext_name): + return FAIL + + # 校验包中文件或目录大小 + if pkg_args.check_size == "True": + limit_list, tag = processing_csv_file( + release_dir, package_name.func_name, package_name.chip_name, pkg_args.build_type + ) + if not tag: + return FAIL + if limit_list: + abspath = os.path.abspath(release_dir) + replace_path = abspath + "/" + result = check_add_dir(replace_path, abspath, limit_list) + if not result: + return FAIL + try: + package_name = get_compress_cmd(pkg_args.pkg_output_dir, pkg_args, xml_config) + except CompressError: + return FAIL + + CommLog.cilog_info("package %s generate filelist.csv and makeself cmd successfully!", package_name) + return SUCCESS + + +def check_path_is_conflict(xml_config): + """ + 功能描述: 检查打包时安装路径与软连接路径是否冲突 + 参数: xml_config + 返回值: SUCCESS/FAIL + """ + install_path_list = set() + pkg_softlink_list = set() + for item in xml_config.package_content_list: + value_list = item.get("value").split("/") + target_name = value_list[-1] if value_list[-1] else value_list[-2] + if item.get("install_path"): + install_path_list.add(os.path.join(item["install_path"], target_name)) + if item.get("pkg_inner_softlink"): + pkg_softlink = item.get("pkg_inner_softlink") + pkg_softlink_list.add(pkg_softlink) + if install_path_list & pkg_softlink_list: + CommLog.cilog_info("intersection:{}".format(install_path_list & pkg_softlink_list)) + CommLog.cilog_info("path conflicting: pkg_inner_softlink dir equals install_path!!") + return FAIL + return SUCCESS + + +def checksum_value(limit_value, release_dir): + """ + 功能描叙: 校验传入的文件或目录大小是否合格 + 参数: + limit_value: limit.csv中的一行数据如[compiler/bin, 3976, 110%] + 返回值: True/False + """ + path = os.path.join(release_dir, limit_value[1]) + if len(limit_value) >= 7: + try: + max_value = int(limit_value[4]) + except ValueError: + CommLog.cilog_error(f"{path} configuration is not standard., Please check limit.csv.") + return True + else: + CommLog.cilog_error(f"{path} configuration is less than four, Please check limit.csv.") + return True + if not os.path.exists(path): + CommLog.cilog_warning(f"{path} doesn't exist, Please check limit.csv.") + return True + size = 0 + for root, dirs, files in os.walk(path): + size += os.path.getsize(root) + for f in files: + filepath = os.path.join(root, f) + if os.path.islink(filepath): + continue + if not os.path.exists(filepath): + continue + size += os.path.getsize(os.path.join(root, f)) + if size == 0: + size = os.path.getsize(path) + if size > max_value * 1024: + CommLog.cilog_error(f"\n{path} size {size} bytes exceeds maximum {max_value * 1024} bytes") + return False + return True + + +def processing_csv_file(release_dir, package_name, chip_name, build_type): + """ + 功能描叙: 处理limit.csv文件数据 + 返回值: [],True/[],False + """ + ret = True + limit_list = [] + product = os.path.basename(os.path.dirname(release_dir)) + limit_path = os.path.join(pkg_utils.TOP_SOURCE_DIR, CONFIG_SCRIPT_PATH, "common/limit.csv") + if not os.path.exists(limit_path): + CommLog.cilog_warning(f"{limit_path} doesn't exist.") + return limit_list, ret + with open(limit_path) as file: + reader = csv.reader(file) + next(reader) + for data in reader: + if not data: + CommLog.cilog_warning("The limit.csv file contains empty lines.") + continue + if is_match_line(package_name, chip_name, product, build_type, data): + if data[1][-1] == "/": + limit_list.append(data[1][:-1]) + else: + limit_list.append(data[1]) + res = checksum_value(data, release_dir) + if not res: + ret = False + return limit_list, ret + + +def is_match_line(package_name, chip_name, product, build_type, data): + return package_name == data[0] and chip_name == data[5] and product == data[6] and build_type == data[7].lower() + + +def check_add_dir(package_path, dirs, limit_list, ret=True): + """ + 功能描述: 校验新增目录 + 参数: path, limit_list + 返回值: False/True + """ + for limit_path in limit_list: + if dirs == os.path.join(os.path.split(dirs)[0], limit_path): + return ret + for dir_file in os.listdir(dirs): + path = os.path.join(dirs, dir_file) + relative_path = path.replace(package_path, "") + if os.path.isfile(path) and relative_path not in limit_list: + CommLog.cilog_error(f"{path} is not in limit.csv file and is newly added.") + ret = False + elif os.path.isdir(path) and relative_path not in limit_list: + ret = check_add_dir(package_path, path, limit_list, ret) + return ret + + +def get_target_name(target_conf) -> str: + """获取目标名。""" + rename = target_conf.get("rename") + if rename: + return rename + + value_list = target_conf.get("value").split("/") + target_name = value_list[-1] if value_list[-1] else value_list[-2] + return target_name + + +def gen_file_install_list(xml_config: XmlConfig, filter_key) -> tuple[FileList, FileList]: + """生成filelist列表。""" + file_install_list = [] + + dir_filelist = parse_install_info(xml_config.dir_install_list, "mkdir", filter_key) + move_filelist = parse_install_info(xml_config.move_content_list, "move", filter_key) + pkg_filelist = parse_install_info(xml_config.package_content_list, "copy", filter_key) + gen_filelist = parse_install_info(xml_config.generate_infos, "copy", filter_key) + # file_info中配置为文件夹，这里是被展开的文件,则需要单独删除 + del_filelist = parse_install_info(xml_config.expand_content_list, "del", filter_key) + collect_filelist = list(chain(dir_filelist, move_filelist, pkg_filelist, gen_filelist)) + collect_filelist = list(xml_config.packer_config.fill_is_common_path(collect_filelist)) + all_filelist = list(chain(collect_filelist, del_filelist)) + for file_item in all_filelist: + file_install_list.append(file_item) + + return file_install_list, [] + + +def generate_filelist_file_by_xml_config(xml_config: XmlConfig, filter_key: list[str], package_check: bool): + """生成文件列表文件。""" + check_move = xml_config.package_attr.get("use_move", False) + transform_nested_path_func = get_transform_nested_path_func(xml_config.package_attr.get("parallel") or check_move) + check_features = xml_config.package_attr.get("check_features", False) + + file_install_list, [] = invoke( + pipe( + gen_file_install_list, + partial(map, transform_nested_path_func), + tuple, + ), + xml_config, + filter_key, + ) + generate_filelist(file_install_list, "filelist.csv") + # 先生成再检查，有利于问题定位 + check_filelist(file_install_list, check_features, check_move) + + +def get_pkg_xml_relative_path(pkg_args: Namespace) -> str: + """获取包配置文件相对路径。""" + + def parts(): + yield CONFIG_SCRIPT_PATH + yield pkg_args.pkg_name + if pkg_args.chip_scenes: + yield pkg_args.chip_scenes + # 可以通過build_rule指定xml_file，而且优先级高于默认值 + if pkg_args.xml_file: + yield pkg_args.xml_file + else: + yield f"{pkg_args.pkg_name}.xml" + + return os.path.join(*parts()) + + +def write_config_inc_var(name: str, package_attr: dict, file: TextIO): + """向config.inc文件写入变量。""" + if name in package_attr: + value = str(package_attr[name]).lower() + file.write(f"{name.upper()}={value}\n") + + +def generate_config_inc(package_attr: dict): + """生成config.inc文件。""" + if "parallel" not in package_attr and "parallel_limit" not in package_attr and "use_move" not in package_attr: + return + year = datetime.now(timezone.utc).year + config_inc = os.path.join(TOP_DIR, "build", "config.inc") + header = [ + "#!/bin/sh\n", + "#----------------------------------------------------------------------------\n", + f"# Copyright Huawei Technologies Co., Ltd. 2023-{year}. All rights reserved.\n", + "#----------------------------------------------------------------------------\n", + "\n", + ] + if os.path.isfile(config_inc): + os.chmod(config_inc, 0o700) + with open(config_inc, "w", encoding="utf-8") as file: + file.writelines(header) + write_config_inc_var("parallel", package_attr, file) + write_config_inc_var("parallel_limit", package_attr, file) + write_config_inc_var("use_move", package_attr, file) + + os.chmod(config_inc, 0o500) + + +def update_version_info(new_version: str): + version_path = os.path.join(pkg_utils.TOP_DIR, "version.info") + with open(version_path) as file: + content = file.read() + content = re.sub(r"Version=.*", f"Version={new_version}", content) + content = re.sub(r"vension_dir=.*", f"version_dir={new_version}", content) + with open(version_path, "w") as file: + file.write(content) + + +def main(pkg_name="", xml_file="", main_args=None): + """ + 功能描述: 执行打包流程(解析配置--->生成文件列表--->执行拷贝/打包动作) + 参数: pkg_name, os_arch, type + 返回值: SUCCESS/FAIL + """ + delivery_dir = os.path.join(TOP_DIR, DELIVERY_PATH) + if not os.path.exists(delivery_dir): + return FAIL + + config_relative_path = get_pkg_xml_relative_path(main_args) + pkg_xml_file = os.path.join(pkg_utils.TOP_SOURCE_DIR, config_relative_path) + parse_option = make_parse_option(main_args) + if main_args.version_dir: + update_version_info(main_args.version_dir) + + try: + xml_config = parse_xml_config(pkg_xml_file, delivery_dir, parse_option, main_args) + except ContainAsteriskError as ex: + CommLog.cilog_error(f"Value contain '*' in {config_relative_path}. value is '{ex.value}'.") + return FAIL + + if pkg_name in ["driver", "firmware"]: + filter_key = ["all", "docker"] + elif pkg_name in ["aicpu_kernels_device", "aicpu_kernels_host"]: + filter_key = [] + else: + filter_key = ["all", "run"] + + # 生成filelist.csv安装列表文件 + try: + generate_filelist_file_by_xml_config( + xml_config, filter_key, main_args.package_check or xml_config.package_attr.get("package_check") + ) + except PackageNameEmptyError: + CommLog.cilog_error(f"package name is empty in {xml_file}, please check it") + return FAIL + except GenerateFilelistError as ex: + CommLog.cilog_error( + f"generate filelist {ex.filename} failed!", + ) + return FAIL + except FilelistError as ex: + CommLog.cilog_error("check filelist error! %s", str(ex)) + return FAIL + + generate_config_inc(xml_config.package_attr) + + if main_args.independent_pkg: + src_file_path = os.path.join(TOP_DIR, "build", "filelist.csv") + dst_file_path = os.path.join(main_args.pkg_output_dir, "share", "info", main_args.pkg_name, "script") + shutil.copy(src_file_path, dst_file_path) + + package_option = PackageOption( + main_args.os_arch, + main_args.package_suffix, + main_args.not_in_name, + main_args.pkg_version, + main_args.ext_name, + chip_name=main_args.chip_name, + func_name=main_args.func_name, + version_dir=main_args.version_dir, + disable_multi_version=main_args.disable_multi_version, + suffix=main_args.suffix, + ) + + package_name = PackageName(xml_config.package_attr, main_args, xml_config.version) + + # 检查install_path与pkg_inner_softlink路径是否冲突，若冲突则报错 + if check_path_is_conflict(xml_config) == FAIL: + return FAIL + + # 生成打包命令 + return execute_repack_process( + xml_config, delivery_dir, main_args, package_name=package_name, package_option=package_option + ) + + +def args_parse(): + """ + 功能描述 : 脚本入参解析 + 参数 : 调用脚本的传参 + 返回值 : 解析后的参数值 + """ + parser = argparse.ArgumentParser(description="This script is for package repack processing.") + parser.add_argument( + "-c", + "--chip_scenes", + metavar="chip_scenes", + required=False, + dest="chip_scenes", + nargs="?", + const="", + default="", + help="This parameter define chip id for package.", + ) + parser.add_argument( + "-n", "--pkg_name", metavar="pkg_name", required=False, help="This parameter define pkg_name for config_xml." + ) + parser.add_argument( + "-o", + "--os_arch", + metavar="os_arch", + required=False, + dest="os_arch", + nargs="?", + const="", + default=None, + help="This parameter define the package's os_arch", + ) + parser.add_argument( + "-t", + "--type", + metavar="type", + required=False, + dest="type", + nargs="?", + const="", + default="repack", + help="This parameter define this script's function", + ) + parser.add_argument( + "-i", + "--not_in_name", + metavar="not_in_name", + required=False, + dest="not_in_name", + nargs="?", + const="", + default="", + help="This parameter define the package's name not contain the element", + ) + parser.add_argument( + "-v", + "--pkg_version", + metavar="pkg_version", + required=False, + dest="pkg_version", + nargs="?", + const="", + default="", + help="This parameter define the version for package.", + ) + parser.add_argument( + "-e", + "--ext_name", + metavar="ext_name", + required=False, + dest="ext_name", + nargs="?", + const="", + default="", + help="This parameter define the package's ext_name", + ) + parser.add_argument( + "--package_suffix", + nargs="?", + const="none", + default="none", + help="This parameter define the package suffix, debug or none", + ) + parser.add_argument( + "--suffix", + metavar="suffix", + required=False, + dest="suffix", + nargs="?", + const="", + default=None, + help="This parameter define the package suffix, for example such as tar.gz", + ) + parser.add_argument( + "-b", + "--build_type", + metavar="build_type", + required=False, + dest="build_type", + nargs="?", + const="", + default="debug", + help="This parameter define release type of package", + ) + parser.add_argument( + "-x", + "--xml", + metavar="xml_file", + required=False, + dest="xml_file", + nargs="?", + const="", + default="", + help="This parameter define xml file", + ) + parser.add_argument( + "--chip_name", + metavar="chip_name", + required=False, + dest="chip_name", + nargs="?", + const=None, + default=None, + help="This parameter define package chip name, has higher priority than chip name in xml", + ) + parser.add_argument( + "--func_name", + metavar="func_name", + required=False, + dest="func_name", + nargs="?", + const=None, + default=None, + help="This parameter define package func name, has higher priority than func name in xml", + ) + parser.add_argument( + "--source_root", + metavar="source_root", + required=False, + dest="source_root", + nargs="?", + const="", + help="source root dir.", + ) + parser.add_argument( + "--makeself_dir", + metavar="makeself_dir", + required=False, + dest="makeself_dir", + nargs="?", + const="", + help="makeself dir.", + ) + parser.add_argument("--independent_pkg", action="store_true", help="Independent pkg.") + parser.add_argument("--pkg-output-dir", default="", help="Package dirpath.") + parser.add_argument("--version_dir", nargs="?", const="", default="", help="Set version dir.") + parser.add_argument("--tag", metavar="tag", nargs="?", const="", default="") + parser.add_argument("--disable-multi-version", action="store_true", help="Disable multi version.") + # 检查打包配置 + parser.add_argument("--package-check", action="store_true", help="check package config.") + parser.add_argument("--check_size", nargs="?", const="", default="", help="Check the size of a file or directory.") + parser.add_argument("--pkg-name-style", metavar="pkg_name_style", default="common", help="Package name style.") + return parser.parse_args() + + +if __name__ == "__main__": + CommLog.cilog_info("%s", " ".join(sys.argv)) + args = args_parse() + try: + if args.source_root: + pkg_utils.TOP_SOURCE_DIR = args.source_root + if args.build_type == "": + args.build_type = "debug" + else: + args.build_type = args.build_type.lower() + status = main(args.pkg_name, args.xml_file, main_args=args) + except Exception as e: + CommLog.cilog_error("exception is occurred (%s)!", e) + CommLog.cilog_info("%s", traceback.format_exc()) + status = FAIL + sys.exit(status) diff --git a/csrc/ascend/scripts/util/__init__.py b/csrc/ascend/scripts/util/__init__.py new file mode 100644 index 000000000..588ef7ff3 --- /dev/null +++ b/csrc/ascend/scripts/util/__init__.py @@ -0,0 +1,16 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import os +import sys + +PYF_PATH = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(PYF_PATH) diff --git a/csrc/ascend/scripts/util/build_opp_kernel_static.py b/csrc/ascend/scripts/util/build_opp_kernel_static.py new file mode 100644 index 000000000..63e2ec16e --- /dev/null +++ b/csrc/ascend/scripts/util/build_opp_kernel_static.py @@ -0,0 +1,779 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +""" +build_opp_kernel_static.py +""" + +import argparse +import concurrent.futures +import contextlib +import glob +import json +import logging as log +import multiprocessing +import os +import platform +import stat +import subprocess +import sys +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path + +import regex as re + + +class Const: + x86 = "x86_64" + arm = "aarch64" + + +def shell_exec(cmd, shell=False): + try: + ps = subprocess.Popen(cmd, shell) + ps.communicate(timeout=180) + except BaseException as e: + log.error("shell_exec error: %s", e) + sys.exit(1) + + +def shell_checkout_key_func(symbol_file, key_str): + process = subprocess.Popen(("cat", symbol_file), stdout=subprocess.PIPE) + awk_out = subprocess.check_output(("awk", "{print $8}"), stdin=process.stdout) + process.wait() + cppfilt = subprocess.check_output(("c++filt",), input=awk_out) + if key_str not in cppfilt.decode("utf-8"): + return b"" + grep_out = subprocess.check_output(("grep", key_str), input=cppfilt) + return grep_out.decode("utf-8") + + +def to_upper_camel_case(x) -> str: + """转大驼峰法命名""" + s = re.sub("_([a-zA-Z])", lambda m: (m.group(1).upper()), x.lower()) + return s[0].upper() + s[1:] + + +def generate_symbol(args): + library_file = args.library_file + symbol_file = args.symbol_file + if not os.path.exists(library_file): + raise FileExistsError(f"generate_symbol input library error, file <{library_file}> not exists.") + process = subprocess.Popen(("readelf", "-Ws", library_file), stdout=subprocess.PIPE) + output = process.communicate(timeout=180)[0].decode("utf-8") + with open(symbol_file, "w") as fd: + fd.write(output) + + +def parser_generate_symbol(subparsers): + generate_symbol_parser = subparsers.add_parser(name="GenerateSymbol", help="Generate Symbol file of input library") + generate_symbol_parser.add_argument( + "-l", "--library_file", type=str, required=False, dest="library_file", default="", help="Input the library file" + ) + generate_symbol_parser.add_argument( + "-s", + "--symbol_file", + type=str, + required=False, + dest="symbol_file", + default="", + help="The symbol file for output", + ) + generate_symbol_parser.set_defaults(func=generate_symbol) + + +class CompileOpStaticLib: + def __init__(self, ops_compile_files: dict, out_path: str, dist_index: int, arch: str): + self.ops_compile_files = ops_compile_files + self.out_path = out_path + self.part_index = dist_index + self.cpu_arch = arch + if self.cpu_arch not in [Const.x86, Const.arm]: + raise Exception(f"CompileOpStaticLib Error, input arch<{arch}> error...") + + def compile_link_single(self, file_path, file_o): + (dir_path, file_name) = os.path.split(file_path) + if self.cpu_arch == Const.x86: + shell_exec( + [ + "bash", + "-c", + f"cd {dir_path} && " + f"objcopy --input-target binary --output-target elf64-x86-64 " + f"--binary-architecture i386 " + f"{file_name} {file_o}", + ], + shell=False, + ) + elif self.cpu_arch == Const.arm and platform.machine() != Const.x86: + shell_exec( + [ + "bash", + "-c", + f"cd {dir_path} && " + f"objcopy --input-target binary " + f"--output-target elf64-littleaarch64 --binary-architecture aarch64 " + f"{file_name} {file_o}", + ], + shell=False, + ) + elif self.cpu_arch == Const.arm: + shell_exec( + [ + "bash", + "-c", + f"cd {dir_path} && " + f"aarch64-linux-gnu-objcopy --input-target binary " + f"--output-target elf64-littleaarch64 --binary-architecture aarch64 " + f"{file_name} {file_o}", + ], + shell=False, + ) + + def compile_link_o(self, out_path, file_path, is_need_path=True): + file_pre = os.path.basename(file_path).replace(".", "_").replace("-", "_") + path_o_prefix = os.path.join(out_path, f"data_{file_pre}_{self.cpu_arch}.o") + # 向json文件中写入"filePath"参数 + if is_need_path and file_path.name.endswith(".json"): + with open(file_path, encoding="UTF-8") as json_fd: + json_dict = json.load(json_fd) + soc = str(file_path).split("/binary/")[-1].split("/bin/")[0] + json_dict["filePath"] = os.path.join(soc, str(file_path).split("/bin/")[-1].split("/kernel/")[-1]) + if "opp/built-in/" in str(file_path): + json_dict["filePath"] = ( + str(file_path).split("/bin/")[-1].split("/kernel/")[-1].replace("/ops_transformer", "") + ) + file_path = os.path.join(out_path, os.path.basename(file_path)) + with open(file_path, "w", encoding="UTF-8") as new_json_fd: + new_json_fd.write(json.dumps(json_dict, indent=4)) + file_path = os.path.realpath(file_path) + self.compile_link_single(file_path, path_o_prefix) + return + + def compile_ops_part_o(self, out_path): + path_data_o = os.path.join(out_path, f"data_*{self.cpu_arch}.o") + path_data_o_list = glob.glob(path_data_o) + if not path_data_o_list: + return + (dir_path, ops_name) = os.path.split(out_path) + file_part_o = f"{ops_name}_{self.cpu_arch}_part{self.part_index}.o" # eg: floor_mod_aarch64_1.a + path_part_o = os.path.join(dir_path, file_part_o) + if self.cpu_arch == Const.x86 or (self.cpu_arch == Const.arm and platform.machine() != Const.x86): + shell_exec(["bash", "-c", f"cd {out_path} && ld -r {path_data_o} -o {path_part_o}"], shell=False) + if self.cpu_arch == Const.arm and platform.machine() == Const.x86: + shell_exec( + ["bash", "-c", f"cd {out_path} && aarch64-linux-gnu-ld -r {path_data_o} -o {path_part_o}"], shell=False + ) + return + + def exec_compile(self): + """ + 编译算子静态库 + :return: + """ + + def get_parallel_num() -> int: + """ + 获取多线程最大并发数量 + """ + num = multiprocessing.cpu_count() * 2 + if num == 0: + num = 16 + return num + + job_num = get_parallel_num() + for op in self.ops_compile_files: + compile_files = self.ops_compile_files[op].kernel_files + json_files = self.ops_compile_files[op].binary_config_files + runtime_kb_files = self.ops_compile_files[op].runtime_kb_files + op_out_path = os.path.join(self.out_path, op) + if not os.path.exists(op_out_path): + os.makedirs(op_out_path, exist_ok=True) + with concurrent.futures.ThreadPoolExecutor(max_workers=job_num) as executor: + for file in compile_files: + executor.submit(self.compile_link_o, op_out_path, file.resolve()) + for file in json_files: + executor.submit(self.compile_link_o, op_out_path, file.resolve(), False) + for file in runtime_kb_files: + executor.submit(self.compile_link_o, op_out_path, file.resolve(), False) + + with concurrent.futures.ThreadPoolExecutor(max_workers=job_num) as executor: + executor.submit(self.compile_ops_part_o, op_out_path) + return 0 + + +def compile_static_library(args): + index_num = args.index_num + cpu_aarch = args.cpu_aarch + + if cpu_aarch not in [Const.x86, Const.arm]: + raise Exception(f"Input cpu_aarch<{cpu_aarch}> Error, Please input parase.") + + ops_compile_files = GenOpResourceIni(args.soc_version, args.build_dir, args.jit).analyze_ops_files() + + csl = CompileOpStaticLib( + ops_compile_files, os.path.join(args.build_dir, f"bin_tmp/{args.soc_version}"), index_num, cpu_aarch + ) + ret = csl.exec_compile() + return ret + + +def parser_compile_static_library(subparsers): + """配置静态编译参数及执行信息""" + compile_lib_parser = subparsers.add_parser( + name="StaticCompile", help="Compile static libraries(.a) on distributed server" + ) + compile_lib_parser.add_argument( + "-s", + "--soc_version", + type=str, + required=True, + dest="soc_version", + help="Operator Name, eg: ascend910b, ascend310p", + ) + compile_lib_parser.add_argument( + "-b", "--build_dir", type=str, required=True, dest="build_dir", help="Input build dir for this project" + ) + compile_lib_parser.add_argument( + "-j", "--jit", action="store_true", dest="jit", help="Compile static libraries(.a) with cann package" + ) + compile_lib_parser.add_argument( + "-n", "--index_num", type=int, required=True, dest="index_num", help="Please input distributed compilation idx" + ) + compile_lib_parser.add_argument( + "-a", "--cpu_aarch", type=str, required=True, dest="cpu_aarch", help="Please input cpu aarch, eg:x86_64,aarch64" + ) + compile_lib_parser.set_defaults(func=compile_static_library) + + +@dataclass +class OpResource: + """算子资源""" + + # tiling 注册函数 + tiling_register: str = field(default=None) + extend_register: str = field(default_factory=list) + # InferShape 注册函数 + infer_shape_register: str = field(default=None) + # 知识库注册 + tuning_bank_key_register: str = field(default=None) + tuning_bank_parse_register: str = field(default=None) + tuning_tiling_helper: str = field(default=None) + # 二进制配置 + binary_config_files: list = field(default_factory=list) + # kernel 编译产出文件 + kernel_files: list = field(default_factory=list) + # 知识库文件 + runtime_kb_files: list = field(default_factory=list) + + +class GenOpResourceIni: + def __init__(self, soc_version: str, build_dir: str, build_with_package: bool): + self._soc_version = soc_version + self._build_dir = Path(build_dir) + opp_path = os.environ.get("ASCEND_OPP_PATH") + if build_with_package and opp_path: + opp_path = Path(opp_path) + self._binary_path = opp_path / "built-in/op_impl/ai_core/tbe/kernel" + self._tuning_basic_path = opp_path / "built-in/data/op" + ops_info = opp_path / "built-in/op_impl/ai_core/tbe/config" / self._soc_version + ops_info = list(ops_info.glob(f"aic-{self._soc_version}-ops-info-transformer.json")) + self._ops_info = ops_info[0] if len(ops_info) != 0 else None + else: + self._binary_path = self._build_dir / "binary" / self._soc_version / "bin" + self._tuning_basic_path = self._build_dir / "tbe/config" / self._soc_version + # transformer aic*.json 适配 + ops_info = self._build_dir / "custom/op_impl/ai_core/tbe/config" / self._soc_version + ops_info = list(ops_info.glob(f"aic-{self._soc_version}-ops-info*.json")) + self._ops_info = ops_info[0] if len(ops_info) != 0 else None + self._op_resource_path = self._build_dir / "autogen" / self._soc_version / "aclnnop_resource" + self._op_res: dict[str, OpResource] = defaultdict(OpResource) + self._l0op_list = [] + + TILING_REG_DECL_FMT = """ +namespace {namespace} {{ + extern gert::OpImplRegisterV2 {func_name}; +}} +""" + EXTEND_REG_DECL_FMT = """ +namespace {namespace} {{ + extern uint32_t {func_name}; +}} +""" + TILING_REG_RES_FUNC_FMT = """ +void * {op_type}TilingRegisterResource() {{ + return {reference_code}; +}} +""" + INFER_SHAPE_REG_DECL_FMT = """ +namespace {namespace} {{ + extern gert::OpImplRegisterV2 {func_name}; +}} +""" + INFER_SHAPE_REG_RES_FUNC_FMT = """ +void * {op_type}InferShapeRegisterResource() {{ + return {reference_code}; +}} +""" + TUNING_REG_DECL_FMT = """ +namespace {namespace} {{ + class {class_type}; + extern {class_type} {func_name}; +}} +""" + EXTLEND_REG_RES_FUNC_FMT = """ +void * {op_type}ExtendRegisterResource() {{ + static std::vector resource = {{{reference_code}}}; + return &resource; +}} +""" + TUNING_REG_RES_FUNC_FMT = """ +void * {op_type}TuningRegisterResource() {{ + static std::vector resource = {{{tuning_bank_key}, {tuning_bank_parse}, {tuning_helper}}}; + return &resource; +}} +""" + KERNEL_BINARY_RES_FUNC_FMT = """ +const OP_BINARY_RES& {op_type}KernelResource() {{ + static const OP_BINARY_RES resource = {{ + {binary_config_ref_code} + {kernel_files_ref_code} + }}; + return resource; +}} +""" + TUNING_KB_BINARY_RES_FUNC_FMT = """ +const OP_RUNTIME_KB_RES& {op_type}TuningResource() {{ + static const OP_RUNTIME_KB_RES resource = {{ + {reference_code} + }}; + return resource; +}} +""" + OP_RESOURCE_CPP_FMT = """/******************{op_type}算子的所有资源**********************/ +#include "register/op_impl_registry.h" +#include +#include +#include +#include +#include + +using OP_HOST_FUNC_HANDLE = std::vector; +using OP_RES = std::tuple; +using OP_BINARY_RES = std::vector; +using OP_RUNTIME_KB_RES = std::vector; +using OP_RESOURCES = std::map>; +namespace {op_type} {{ + auto initializer = StaticSpaceInitializer::GetInstance();; +}} + +// 资源声明 +// extend resource +{extend_declaration} +// Tiling +{tiling_declaration} +// InferShape +{infer_shape_declaration} +// Tuning +{tuning_bank_key_declaration} +{tuning_bank_parse_declaration} +{tuning_helper_declaration} +// kernel 二进制 +{binary_config_declaration} +{kernel_files_declaration} +// kb 二进制 +{tuning_kb_declaration} + +namespace l0op {{ +// 资源函数 +// Tiling register resource func +{tiling_reg_func} +// InferShape register resource func +{infer_shape_reg_func} +// Tuning register resource func +{tuning_reg_func} +// kernel resource func +{kernel_resource} +// Tuning resource func +{tuning_kb_resource} +}} + +// extend resource func +{extend_reg_func} + +""" + + @staticmethod + def _extract_op_symbol_pair(symbol_file: str, search_key: str, prefix: str, suffix: str): + symbol_ret = shell_checkout_key_func(symbol_file, search_key) + for symbol in symbol_ret.splitlines(): + symbol_name = symbol.split("::")[-1] + if not (symbol_name.startswith(prefix) and symbol_name.endswith(suffix)): + log.warning("symbol not satisfied with the format:%s%s, skip", prefix, suffix) + continue + op_type = symbol_name + if prefix: + op_type = op_type[len(prefix) :] + if suffix: + op_type = op_type[: -len(suffix)] + yield op_type, symbol + + @staticmethod + def _extract_op_symbol_pair_v2(symbol_file: str, search_key: str, prefix: str): + symbol_ret = shell_checkout_key_func(symbol_file, search_key) + for symbol in symbol_ret.splitlines(): + symbol_name = symbol.split("::")[-1] + if not (symbol_name.startswith(prefix)): + log.warning("symbol not satisfied with the format:%s, skip", prefix) + continue + op_type = symbol_name + if prefix: + op_type = op_type[len(prefix) :] + op_type = op_type.split("_")[0] + yield op_type, symbol + + @staticmethod + def _extract_register_symbol(register_symbol: str): + if not register_symbol: + return "", "", "nullptr" + + symbol_data = register_symbol.split("::") + namespace = "::".join(symbol_data[:-1]) + if "anonymous" in namespace: + return "", "", "nullptr" + func_name = symbol_data[-1] + reference_code = f"&{register_symbol}" + return namespace, func_name, reference_code + + @staticmethod + def _gen_binary_res_code(files): + declaration = "" + reference_code = "" + for binary_file in files: + # static not support supperkernel + if "relocatable" in binary_file.name: + continue + binary_name = binary_file.name.replace(".", "_").replace("-", "_") + declaration += f"""// {binary_file.name} +extern const uint8_t _binary_{binary_name}_start[]; +extern const uint8_t _binary_{binary_name}_end[]; +""" + reference_code += f"{{_binary_{binary_name}_start, _binary_{binary_name}_end}},\n" + return declaration, reference_code + + def gen_ops_ini_files(self): + self.analyze_ops_files() + self._analyze_symbols() + self._analyze_ops_l0op() + if not os.path.exists(self._op_resource_path): + os.makedirs(self._op_resource_path) + for op_type in self._l0op_list: + ini_content = self.generate_op_resouce_ini(op_type) + self._save_op_resource(op_type, ini_content) + for op_type in self._op_res: + if op_type in self._l0op_list: + continue + ini_content = self.generate_op_resouce_ini(op_type) + self._save_op_resource(op_type, ini_content) + + def generate_op_resouce_ini(self, op_type: str) -> str: + value_dict = { + "op_type": op_type, + } + value_dict.update(self._gen_register_resouce_code(op_type)) + value_dict.update(self._gen_tuning_register_resouce_code(op_type)) + value_dict.update(self._gen_binary_resource_code(op_type)) + # 处理特殊的共用算子kernel资源的算子 + sepical_ops = {"MatMulV2": "MatMul"} + if op_type in sepical_ops: + value_dict["kernel_files_declaration"] = "" + value_dict["kernel_resource"] = f""" +extern const OP_BINARY_RES& {sepical_ops[op_type]}KernelResource(); +const OP_BINARY_RES& {op_type}KernelResource() {{ + return {sepical_ops[op_type]}KernelResource(); +}} +""" + return self.OP_RESOURCE_CPP_FMT.format_map(value_dict) + + def analyze_ops_files(self): + if not self._ops_info: + return self._op_res + with open(self._ops_info) as autogen_fd: + ops_info_json = json.load(autogen_fd) + + for ops in ops_info_json: + if "opFile" in ops_info_json[ops]: + json_file = f"{ops_info_json[ops]['opFile']['value']}.json" + else: + o_lists = list(Path(self._binary_path).rglob(f"{self._soc_version}/**/*{ops}*.o")) + if len(o_lists) == 0: + continue + else: + json_file = f"{os.path.basename(os.path.dirname(o_lists[0]))}.json" + # json_path = self._binary_path / "config" / self._soc_version / json_file + json_path = self._binary_path / json_file + if "opp/built-in/" in str(self._binary_path): + json_path = self._binary_path / "config" / self._soc_version / "ops_transformer" / json_file + if not os.path.exists(json_path): + continue + with open(json_path) as op_json_fd: + op_json_content = json.load(op_json_fd) + if "binList" not in op_json_content or len(op_json_content["binList"]) == 0: + continue + # 算子.json内 kernel json路径适配 + bin_json_file = ( + self._binary_path / op_json_content["binList"][0]["binInfo"]["jsonFilePath"].split("/", 1)[1] + ) + if "opp/built-in/" in str(self._binary_path): + bin_json_file = ( + self._binary_path + / self._soc_version + / "ops_transformer" + / op_json_content["binList"][0]["binInfo"]["jsonFilePath"].split("/", 1)[1] + ) + ops_path = os.path.dirname(bin_json_file) + self._op_res[ops].binary_config_files.append(json_path) + self._op_res[ops].kernel_files.extend(sorted(Path(ops_path).iterdir())) + for kb_json in list(Path(self._tuning_basic_path).rglob("*_AiCore_*_runtime_kb.json")): + ops = kb_json.name.split("_AiCore_")[-1].split("_runtime_kb")[0] + self._op_res[ops].runtime_kb_files.append(kb_json) + self._op_res[ops].runtime_kb_files.sort(key=lambda p: p.name) + return self._op_res + + def _analyze_ops_l0op(self): + opapi_symbol = self._build_dir / "opapi_transformer.txt" + if not os.path.exists(opapi_symbol): + return + # infershape + for op_type, _ in self._extract_op_symbol_pair(opapi_symbol, "_kernelName_Be_Defined_Multi_Times__", "", ""): + self._l0op_list.append(op_type.split("_kernelName_")[0]) + self._l0op_list.sort() + + def _save_op_resource(self, op_type, res_content): + res_cpp_file = self._op_resource_path / f"{op_type}_op_resource.cpp" + with contextlib.suppress(FileNotFoundError): + res_cpp_file.unlink() + + flags = os.O_WRONLY | os.O_CREAT + modes = stat.S_IWUSR | stat.S_IRUSR + with os.fdopen(os.open(res_cpp_file, flags, modes), "w") as fd: + fd.write(res_content) + + def _analyze_symbols(self): + # ophost txt 适配 + ophost_symbol = self._build_dir / "ophost_transformer.txt" + if not os.path.exists(ophost_symbol): + return + # infershape + for op_type, symbol in self._extract_op_symbol_pair( + ophost_symbol, "op_impl_register_infershape_", "op_impl_register_infershape_", "" + ): + self._op_res[op_type].infer_shape_register = symbol + # tiling + for op_type, symbol in self._extract_op_symbol_pair( + ophost_symbol, "op_impl_register_optiling_", "op_impl_register_optiling_", "" + ): + self._op_res[op_type].tiling_register = symbol + for op_type, symbol in self._extract_op_symbol_pair_v2( + ophost_symbol, "op_impl_register_template_", "op_impl_register_template_" + ): + self._op_res[op_type].extend_register.append(symbol) + # 知识库 + for op_type, symbol in self._extract_op_symbol_pair( + ophost_symbol, "BankKeyRegistryInterf", "g_", "BankKeyRegistryInterf" + ): + self._op_res[op_type].tuning_bank_key_register = symbol + for op_type, symbol in self._extract_op_symbol_pair(ophost_symbol, "BankParseInterf", "g_", "BankParseInterf"): + self._op_res[op_type].tuning_bank_parse_register = symbol + for op_type, symbol in self._extract_op_symbol_pair( + ophost_symbol, "g_tuning_tiling_", "g_tuning_tiling_", "Helper" + ): + self._op_res[op_type].tuning_tiling_helper = symbol + + def _gen_register_resouce_code(self, op_type: str): + """注册函数""" + # Tiling + namespace, func_name, reference_code = self._extract_register_symbol(self._op_res[op_type].tiling_register) + symbol_map = { + "op_type": op_type, + "namespace": namespace, + "func_name": func_name, + "reference_code": reference_code, + } + tiling_declaration = self.TILING_REG_DECL_FMT.format_map(symbol_map) if func_name else "" + tiling_reg_func = self.TILING_REG_RES_FUNC_FMT.format_map(symbol_map) if func_name else "" + + reference_code_list = [] + extend_declaration = "" + for symbol in self._op_res[op_type].extend_register: + namespace, func_name, reference_code = self._extract_register_symbol(symbol) + if func_name: + extend_declaration += self.EXTEND_REG_DECL_FMT.format(namespace=namespace, func_name=func_name) + reference_code_list.append(reference_code) + reference_code = ", ".join(reference_code_list) + extend_reg_func = self.EXTLEND_REG_RES_FUNC_FMT.format(op_type=op_type, reference_code=reference_code) + + # InferShape + namespace, func_name, reference_code = self._extract_register_symbol(self._op_res[op_type].infer_shape_register) + symbol_map = { + "op_type": op_type, + "namespace": namespace, + "func_name": func_name, + "reference_code": reference_code, + } + infer_shape_declaration = self.INFER_SHAPE_REG_DECL_FMT.format_map(symbol_map) if func_name else "" + infer_shape_reg_func = self.INFER_SHAPE_REG_RES_FUNC_FMT.format_map(symbol_map) if func_name else "" + + return { + "tiling_declaration": tiling_declaration, + "infer_shape_declaration": infer_shape_declaration, + "tiling_reg_func": tiling_reg_func, + "infer_shape_reg_func": infer_shape_reg_func, + "extend_reg_func": extend_reg_func, + "extend_declaration": extend_declaration, + } + + def _gen_tuning_register_resouce_code(self, op_type: str): + """知识库注册函数""" + # Tuning + namespace, func_name, tuning_bank_key_ref_code = self._extract_register_symbol( + self._op_res[op_type].tuning_bank_key_register + ) + tuning_bank_key_declaration = ( + self.TUNING_REG_DECL_FMT.format( + namespace=namespace, + class_type="OpBankKeyFuncRegistryV2", + func_name=func_name, + ) + if func_name + else "" + ) + namespace, func_name, tuning_bank_parse_ref_code = self._extract_register_symbol( + self._op_res[op_type].tuning_bank_parse_register + ) + tuning_bank_parse_declaration = ( + self.TUNING_REG_DECL_FMT.format( + namespace=namespace, + class_type="OpBankKeyFuncRegistryV2", + func_name=func_name, + ) + if func_name + else "" + ) + namespace, func_name, tuning_helper_ref_code = self._extract_register_symbol( + self._op_res[op_type].tuning_tiling_helper + ) + tuning_helper_declaration = ( + self.TUNING_REG_DECL_FMT.format( + namespace=namespace, + class_type=f"{op_type}ClassHelper", + func_name=func_name, + ) + if func_name + else "" + ) + tuning_reg_func = self.TUNING_REG_RES_FUNC_FMT.format( + op_type=op_type, + tuning_bank_key=tuning_bank_key_ref_code, + tuning_bank_parse=tuning_bank_parse_ref_code, + tuning_helper=tuning_helper_ref_code, + ) + return { + "tuning_bank_key_declaration": tuning_bank_key_declaration, + "tuning_bank_parse_declaration": tuning_bank_parse_declaration, + "tuning_helper_declaration": tuning_helper_declaration, + "tuning_reg_func": tuning_reg_func, + } + + def _gen_binary_resource_code(self, op_type: str) -> str: + """二进制""" + # kernel + binary_config_declaration, binary_config_ref_code = self._gen_binary_res_code( + self._op_res[op_type].binary_config_files + ) + kernel_files_declaration, kernel_files_ref_code = self._gen_binary_res_code(self._op_res[op_type].kernel_files) + kernel_resource = ( + self.KERNEL_BINARY_RES_FUNC_FMT.format( + op_type=op_type, + binary_config_ref_code=binary_config_ref_code, + kernel_files_ref_code=kernel_files_ref_code, + ) + if kernel_files_ref_code + else "" + ) + # 知识库 + tuning_kb_declaration, tuning_kb_ref_code = self._gen_binary_res_code(self._op_res[op_type].runtime_kb_files) + tuning_kb_resource = self.TUNING_KB_BINARY_RES_FUNC_FMT.format( + op_type=op_type, + reference_code=tuning_kb_ref_code, + ) + return { + "binary_config_declaration": binary_config_declaration, + "kernel_files_declaration": kernel_files_declaration, + "tuning_kb_declaration": tuning_kb_declaration, + "kernel_resource": kernel_resource, + "tuning_kb_resource": tuning_kb_resource, + } + + +def generate_op_resource_h_file(args): + soc_version: str = args.soc_version + build_dir = args.build_dir + + gen_ini = GenOpResourceIni(soc_version, build_dir, args.jit) + gen_ini.gen_ops_ini_files() + return + + +def parser_generate_op_resource_h_file(subparsers): + gen_resource_ini_parser = subparsers.add_parser( + name="GenStaticOpResourceIni", help="Generate xxx_op_resource.h on consolidation server" + ) + gen_resource_ini_parser.add_argument( + "-s", + "--soc_version", + type=str, + required=True, + dest="soc_version", + help="Operator Name, eg: ascend910b, ascend310p", + ) + gen_resource_ini_parser.add_argument( + "-b", "--build_dir", type=str, required=True, dest="build_dir", help="Input build dir for this project" + ) + gen_resource_ini_parser.add_argument( + "-j", "--jit", action="store_true", dest="jit", help="Generate xxx_op_resource.h with cann package" + ) + gen_resource_ini_parser.set_defaults(func=generate_op_resource_h_file) + + +def execute_argus_parse_func(): + parser = argparse.ArgumentParser() + + subparsers = parser.add_subparsers(help="Subparsers Commands") + + """ 配置静态编译参数及执行信息 """ + parser_compile_static_library(subparsers) + + """ 配置头文件生成功能参数及执行信息 """ + parser_generate_op_resource_h_file(subparsers) + + """ 生成指定库的symbol文件 """ + parser_generate_symbol(subparsers) + + """ 执行函数功能 """ + args = parser.parse_args() + args.func(args) + + +if __name__ == "__main__": + execute_argus_parse_func() + exit(0) diff --git a/csrc/ascend/scripts/util/const_var.py b/csrc/ascend/scripts/util/const_var.py new file mode 100644 index 000000000..86b41fbba --- /dev/null +++ b/csrc/ascend/scripts/util/const_var.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +# ---------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ---------------------------------------------------------------------------- + +""" +Function: +The replay function entry +Copyright Information: +Huawei Technologies Co., Ltd. All Rights Reserved © 2020 +""" + +import os +import stat + +REPLAY_BATCH = "batch" +REPLAY_ITERATE = "iterate" +CFG_IMPL_DIR = "impl_dir" +CFG_OUT_DIR = "out_dir" +AUTO_GEN_DIR = "auto_gen_dir" +WFLAGS = os.O_WRONLY | os.O_CREAT | os.O_TRUNC +WMODES = stat.S_IWUSR | stat.S_IRUSR +SOC_MAP_EXT = { + "ascend310p": "Ascend310P3", + "ascend310b": "Ascend310B1", + "ascend910": "Ascend910A", + "ascend910b": "Ascend910B1", + "ascend910_93": "Ascend910_9391", + "ascend950": "Ascend950", + "kirinx90": "KirinX90", +} +BIN_CMD = "asc_opc $1 --main_func={fun} --input_param={param} --soc_version={soc} \ +--output=$2 --impl_mode={impl} --simplified_key_mode=0 --op_mode=dynamic\n" +SET_PLOG_LEVEL_ERROR = "export ASCEND_GLOBAL_LOG_LEVEL=3\n" +SET_PLOG_STDOUT = "export ASCEND_SLOG_PRINT_TO_STDOUT=1\n" +SRC_ENV = """ +while true; do + case "$1" in + --kernel-src=*) + export BUILD_KERNEL_SRC=$(echo "$1" | cut -d"=" -f2-) + shift + ;; + -*) + shift + ;; + *) + break + ;; + esac +done +""" +CHK_CMD = """ +if ! test -f $2/{res_file} ; then + echo "$2/{res_file} not generated!" + exit 1 +fi +""" +ATTR_DEF_VAL = { + "str": "", + "int": 0, + "float": 0.0, + "bool": False, + "list_bool": [], + "list_int": [], + "list_float": [], + "list_list_int": [[]], +} + + +def conv_soc_ver(ver: str): + return SOC_MAP_EXT.get(ver) diff --git a/csrc/ascend/scripts/util/insert_op_info.py b/csrc/ascend/scripts/util/insert_op_info.py new file mode 100644 index 000000000..63e78b6aa --- /dev/null +++ b/csrc/ascend/scripts/util/insert_op_info.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 +# ---------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ---------------------------------------------------------------------------- + +import json +import os +import sys + +import const_var + +if __name__ == "__main__": + if len(sys.argv) != 3: + print(sys.argv) + print("argv error, inert_op_info.py your_op_file lib_op_file") + sys.exit(2) + + with open(sys.argv[1]) as load_f: + insert_operator = json.load(load_f) + + all_operators = {} + if os.path.exists(sys.argv[2]): + if os.path.getsize(sys.argv[2]) != 0: + with open(sys.argv[2]) as load_f: + all_operators = json.load(load_f) + + for k in insert_operator: + if k in all_operators: + print("replace op:[", k, "] success") + else: + print("insert op:[", k, "] success") + all_operators[k] = insert_operator[k] + + with os.fdopen(os.open(sys.argv[2], const_var.WFLAGS, const_var.WMODES), "w") as json_file: + json_file.write(json.dumps(all_operators, indent=4)) diff --git a/csrc/ascend/scripts/util/merge_aicpu_info_json.sh b/csrc/ascend/scripts/util/merge_aicpu_info_json.sh new file mode 100644 index 000000000..a40c55947 --- /dev/null +++ b/csrc/ascend/scripts/util/merge_aicpu_info_json.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +echo "$@" +top_dir=$1 +output_json=$2 +first_json=$3 +shift 3 + +merge_json_tool="${top_dir}/scripts/util/insert_op_info.py" +>$output_json + +if [[ -f "$first_json" ]] +then + cp -f $first_json $output_json +else + echo "[ERROR] ${first_json} is not a file" + exit 1 +fi + +for single_json in "$@" +do + if [[ -f "${single_json}" ]] + then + python3 ${top_dir}/scripts/util/insert_op_info.py ${single_json} ${output_json} + else + echo "[ERROR] ${single_json} is not a file" + fi +done diff --git a/csrc/ascend/scripts/util/merge_proto.py b/csrc/ascend/scripts/util/merge_proto.py new file mode 100644 index 000000000..f34e45fef --- /dev/null +++ b/csrc/ascend/scripts/util/merge_proto.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import argparse +import os +import sys + +import regex as re + + +def match_op_proto(file_path): + with open(file_path, encoding="utf-8") as f: + content = f.read() + + op_def_pattern = re.compile(r"REG_OP\((.+)\).*OP_END_FACTORY_REG\(\1\)", re.DOTALL) + match = op_def_pattern.search(content) + + if match: + op_name = match.group(1) + op_def = match.group(0) + return op_name, op_def + else: + return None, None + + +def merge_op_proto(protos_path, output_file): + op_defs = [] + for proto_path in protos_path: + if not proto_path.endswith("_proto.h"): + continue + print(f"proto_path: {proto_path}") + op_name, op_def = match_op_proto(proto_path) + if op_def: + op_defs.append(op_def) + + # merge op_proto + merged_content = f"""#ifndef OP_TRANSFORMER_PROTO_H_ +#define OP_TRANSFORMER_PROTO_H_ + +#include "graph/operator_reg.h" +#include "register/op_impl_registry.h" + +namespace ge{{ + +{os.linesep.join([f"{op_def}{os.linesep}" for op_def in op_defs])} +}} // namespace ge + +#endif // OP_TRANSFORMER_PROTO_H_ +""" + + with open(output_file, "w", encoding="utf-8") as f: + f.write(merged_content) + + print(f"merged op transformer proto file: {output_file}") + + +def parse_args(argv): + parser = argparse.ArgumentParser() + parser.add_argument("protos", nargs="+") + parser.add_argument("--output-file", nargs=1, default=None) + return parser.parse_args(argv) + + +if __name__ == "__main__": + args = parse_args(sys.argv) + + protos_path = args.protos[1:] + output_file = args.output_file[0] + merge_op_proto(protos_path, output_file) diff --git a/csrc/ascend/scripts/util/modify_gen_aclnn.py b/csrc/ascend/scripts/util/modify_gen_aclnn.py new file mode 100644 index 000000000..d3ee3a29d --- /dev/null +++ b/csrc/ascend/scripts/util/modify_gen_aclnn.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------------------------------------------------- +# Copyright (c) 2025 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import logging as log +import os +import subprocess +import sys +from pathlib import Path + + +def shell_exec(cmd, shell=False): + try: + ps = subprocess.Popen(cmd, shell) + ps.communicate(timeout=180) + except BaseException as e: + log.error("shell_exec error: %s", e) + sys.exit(1) + + +def search_file(aclnn_cpp): + op_type = None + index = 0 + with open(aclnn_cpp) as f: + for line in f.readlines(): + index = index + 1 + if "_op_resource.h" in line: + op_type = line.replace('_op_resource.h"', "").replace('#include "', "").strip() + if "EXTERN_OP_RESOURCE" in line or "namespace op {" in line: + break + return (op_type, index) + + +def modify_gen_aclnn(build_path): + auto_gen_cpps = Path(os.path.join(build_path, "autogen")).rglob("aclnn*.cpp") + for aclnn_cpp in auto_gen_cpps: + (op_type, index) = search_file(aclnn_cpp) + if op_type: + shell_exec( + ["bash", "-c", f"""sed -i 's/{op_type}_op_resource.h/op_resource.h/g' {aclnn_cpp}"""], shell=False + ) + shell_exec( + ["bash", "-c", f"""sed -i 's/{op_type}_RESOURCES/AUTO_GEN_OP_RESOURCE({op_type})/g' {aclnn_cpp}"""], + shell=False, + ) + shell_exec(["bash", "-c", f"sed -i '{index}i\\EXTERN_OP_RESOURCE({op_type})' {aclnn_cpp}"], shell=False) + return + + +if __name__ == "__main__": + modify_gen_aclnn(sys.argv[1]) diff --git a/csrc/ascend/third_party/catlass b/csrc/ascend/third_party/catlass new file mode 160000 index 000000000..0843e90ab --- /dev/null +++ b/csrc/ascend/third_party/catlass @@ -0,0 +1 @@ +Subproject commit 0843e90ab3edaeccd149e19f01213575ab615770 diff --git a/csrc/ascend/third_party/pto-isa b/csrc/ascend/third_party/pto-isa new file mode 160000 index 000000000..6e076eb94 --- /dev/null +++ b/csrc/ascend/third_party/pto-isa @@ -0,0 +1 @@ +Subproject commit 6e076eb943c42f5ab3f204a0519d8d3a4f4266b8 diff --git a/csrc/ascend/torch_binding.cpp b/csrc/ascend/torch_binding.cpp new file mode 100644 index 000000000..35cada2ae --- /dev/null +++ b/csrc/ascend/torch_binding.cpp @@ -0,0 +1,2697 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include "acl/acl.h" +#include "acl/acl_rt.h" +#include +#include +#include +#include "torch_npu/csrc/core/npu/NPUGuard.h" +#include +#include "ops.h" +#include "utils.h" +#include "aclnn_torch_adapter/op_api_common.h" +#include "moe/add_rms_norm_bias/add_rms_norm_bias_torch_adpt.h" +#include "moe/apply_top_k_top_p_custom/apply_top_k_top_p_custom_torch_adpt.h" +#ifdef VLLM_ENABLE_ATB_AND_DIRECT_KERNELS +#include "batch_matmul_transpose/batch_matmul_transpose_torch_adpt.h" +#include "mla_preprocess/mla_preprocess_torch_adpt.h" +#endif +#include "mc2/dispatch_ffn_combine/dispatch_ffn_combine_torch_adpt.h" +#include "gmm/grouped_matmul_swiglu_quant_weight_nz_tensor_list/grouped_matmul_swiglu_quant_torch_adpt.h" +#include "gmm/grouped_matmul_swiglu_quant_v2/grouped_matmul_swiglu_quant_v2_torch_adpt.h" +#include "attention/lightning_indexer/lightning_indexer_torch_adpt.h" +#include "mc2/matmul_allreduce_add_rmsnorm/matmul_allreduce_add_rmsnorm_torch_adpt.h" +#include "moe/moe_gating_top_k/moe_gating_top_k_torch_adpt.h" +#include "moe/moe_init_routing_custom/moe_init_routing_custom_torch_adpt.h" +#include "attention/sparse_flash_attention/sparse_flash_attention_torch_adpt.h" +#include "attention/kv_quant_sparse_flash_attention/kv_quant_sparse_flash_attention_torch_adpt.h" +#include "attention/lightning_indexer_quant/lightning_indexer_quant_torch_adpt.h" +#include "attention/ngram_spec_decode/ngram_spec_decode_torch_adpt.h" +#include "attention/recurrent_gated_delta_rule/recurrent_gated_delta_rule_torch_adpt.h" +#include "attention/store_kv_block/store_kv_block_torch_adpt.h" +#include "attention/fused_gdn_gating/fused_gdn_gating_torch_adpt.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace vllm_fl { + +namespace { + +constexpr int64_t DSA_SLOT_MAPPING_FLAT = 1; +constexpr int64_t DSA_SLOT_MAPPING_BLOCK_OFFSET = 2; + +struct DevicePrintPayload { + std::string message; + at::Tensor host_tensor_snapshot; +}; + +std::mutex& get_device_print_mutex() +{ + static std::mutex device_print_mutex; + return device_print_mutex; +} + +void device_print_callback(void* args) +{ + // device_print is a debug-only helper. We intentionally do not reclaim the + // callback payload here because aclgraph replay may re-execute the same host + // callback payload multiple times. Freeing it on first execution would make + // later replays dereference a dangling pointer. + auto* payload = static_cast(args); + if (payload == nullptr) { + return; + } + + std::lock_guard guard(get_device_print_mutex()); + if (!payload->message.empty()) { + std::cout << payload->message; + } + + if (payload->host_tensor_snapshot.defined()) { + if (!payload->message.empty()) { + std::cout << std::endl; + } + at::print(std::cout, payload->host_tensor_snapshot.contiguous(), 120); + } + + std::cout << std::endl; + std::cout.flush(); +} + +void enqueue_device_print(std::unique_ptr payload, + aclrtStream stream) +{ + auto* raw_payload = payload.release(); + const aclError ret = aclrtLaunchHostFunc(stream, device_print_callback, + raw_payload); + if (ret != ACL_SUCCESS) { + delete raw_payload; + } + TORCH_CHECK(ret == ACL_SUCCESS, "aclrtLaunchHostFunc failed, error code: ", ret); +} + +} + +void swap_blocks_batch(const torch::Tensor& src_ptrs, + const torch::Tensor& dst_ptrs, + const torch::Tensor& sizes, + int64_t direction) { + + TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU"); + TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU"); + TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU"); + TORCH_CHECK(src_ptrs.dtype() == torch::kInt64, "src_ptrs must be int64"); + TORCH_CHECK(dst_ptrs.dtype() == torch::kInt64, "dst_ptrs must be int64"); + TORCH_CHECK(sizes.dtype() == torch::kInt64, "sizes must be int64"); + + const int64_t n = src_ptrs.size(0); + TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs"); + TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs"); + + if (n == 0) return; + + const int64_t* src_data = src_ptrs.data_ptr(); + const int64_t* dst_data = dst_ptrs.data_ptr(); + const int64_t* size_data = sizes.data_ptr(); + + aclrtStream stream = c10_npu::getCurrentNPUStream().stream(); + + aclrtMemcpyKind memcpy_kind; + switch (direction) { + case 0: + memcpy_kind = ACL_MEMCPY_HOST_TO_DEVICE; + break; + case 1: + memcpy_kind = ACL_MEMCPY_DEVICE_TO_HOST; + break; + case 2: + memcpy_kind = ACL_MEMCPY_DEVICE_TO_DEVICE; + break; + default: + TORCH_CHECK(false, + "swap_blocks_batch: invalid direction ", direction, + " (expected 0=H2D, 1=D2H, 2=D2D)"); + } + + // ========================================================================= + // path 1: aclrtMemcpyBatchAsync (CANN 8.5+) + // ========================================================================= +#if defined(CANN_MEMCPY_BATCH_ASYNC) + if (memcpy_kind != ACL_MEMCPY_DEVICE_TO_DEVICE) { + static_assert(sizeof(void*) == sizeof(int64_t), + "void* and int64_t must be the same size"); + static_assert(sizeof(size_t) == sizeof(int64_t), + "size_t and int64_t must be the same size"); + + void** dst_arr = reinterpret_cast( + const_cast(dst_data)); + void** src_arr = reinterpret_cast( + const_cast(src_data)); + size_t* size_arr = reinterpret_cast( + const_cast(size_data)); + size_t* dest_maxs = size_arr; + + // aclrtMemcpyBatchAttr uses srcLoc/dstLoc (aclrtMemLocation) + // to specify memory locations, not aclrtMemcpyKind. + int32_t device_id = 0; + aclrtGetDevice(&device_id); + + aclrtMemLocation host_loc = {}; + host_loc.type = ACL_MEM_LOCATION_TYPE_HOST; + host_loc.id = 0; + + aclrtMemLocation device_loc = {}; + device_loc.type = ACL_MEM_LOCATION_TYPE_DEVICE; + device_loc.id = device_id; + + aclrtMemcpyBatchAttr attr = {}; + if (memcpy_kind == ACL_MEMCPY_HOST_TO_DEVICE) { + attr.srcLoc = host_loc; + attr.dstLoc = device_loc; + } else { // ACL_MEMCPY_DEVICE_TO_HOST + attr.srcLoc = device_loc; + attr.dstLoc = host_loc; + } + + size_t attrs_index = 0; + size_t fail_index = 0; + + aclError result = aclrtMemcpyBatchAsync( + dst_arr, dest_maxs, src_arr, size_arr, + static_cast(n), + &attr, &attrs_index, 1, + &fail_index, stream); + + TORCH_CHECK(result == ACL_SUCCESS, + "aclrtMemcpyBatchAsync failed at index ", fail_index, + " with error code ", result); + return; + } +#endif + + // ========================================================================= + // path 2: aclrtMemcpyAsync + // ========================================================================= + for (int64_t i = 0; i < n; i++) { + void* dst = reinterpret_cast(dst_data[i]); + const void* src = reinterpret_cast(src_data[i]); + size_t copy_size = static_cast(size_data[i]); + + aclError ret = aclrtMemcpyAsync( + dst, + copy_size, + src, + copy_size, + memcpy_kind, + stream); + + TORCH_CHECK(ret == ACL_SUCCESS, + "aclrtMemcpyAsync failed at index ", i, + " with error code ", ret, + ", src=", src_data[i], + ", dst=", dst_data[i], + ", size=", size_data[i]); + } +} + +#ifdef VLLM_ENABLE_ATB_AND_DIRECT_KERNELS +// Direct kernel wrappers depend on vllm_ascend_kernels, which is skipped on +// 310P and A5 builds. +void swap_blocks_impl(torch::Tensor& src, torch::Tensor& dst, + const torch::Tensor& block_mapping, aclrtStream stream) +{ + torch::Device src_device = src.device(); + torch::Device dst_device = dst.device(); + aclrtMemcpyKind memcpy_type; + + if ((!src_device.is_cpu()) && (!dst_device.is_cpu())) { + TORCH_CHECK(src_device.index() == dst_device.index(), + "src and dst must be on the same npu"); + memcpy_type = ACL_MEMCPY_DEVICE_TO_DEVICE; + } else if ((!src_device.is_cpu()) && dst_device.is_cpu()) { + memcpy_type = ACL_MEMCPY_DEVICE_TO_HOST; + } else if (src_device.is_cpu() && (!dst_device.is_cpu())) { + memcpy_type = ACL_MEMCPY_HOST_TO_DEVICE; + } else { + TORCH_CHECK(false, "Invalid device combination, src tensor device: ", src_device, ", dst tensor device: ", dst_device); + } + + TORCH_CHECK(block_mapping.device().is_cpu(), "block_mapping must be on CPU"); + + char* src_ptr = static_cast(src.data_ptr()); + char* dst_ptr = static_cast(dst.data_ptr()); + + const int64_t block_size_in_bytes = src.element_size() * src.stride(0); + + const int64_t num_blocks = block_mapping.size(0); + const int64_t max_src_block = src.size(0); + const int64_t max_dst_block = dst.size(0); + for (size_t i = 0; i < num_blocks; i++) { + int64_t src_block_number = block_mapping[i][0].item(); + int64_t dst_block_number = block_mapping[i][1].item(); + TORCH_CHECK(src_block_number >= 0 && src_block_number <= max_src_block, + "src block index ", src_block_number, " out of range (max: ", max_src_block, ")"); + TORCH_CHECK(dst_block_number >= 0 && dst_block_number <= max_dst_block, + "dst block index ", dst_block_number, " out of range (max: ", max_dst_block, ")"); + + int64_t src_offset = src_block_number * block_size_in_bytes; + int64_t dst_offset = dst_block_number * block_size_in_bytes; + + aclrtMemcpyAsync(dst_ptr + dst_offset, block_size_in_bytes, + src_ptr + src_offset, block_size_in_bytes, + memcpy_type, stream); + } +} + +void swap_blocks(torch::Tensor &x, torch::Tensor &y, const torch::Tensor &z) +{ + + const c10_npu::OptionalNPUGuard npuGuard( + (!x.device().is_cpu()) ? x.device() : y.device() + ); + aclrtStream stream = c10_npu::getCurrentNPUStream().stream(); + swap_blocks_impl(x, y, z, stream); + return; +} + +AscendType get_dtype_from_torch(at::ScalarType scalarType) +{ + if (scalarType == at::ScalarType::Float) { + return AscendType::FP32; + } else if (scalarType == at::ScalarType::BFloat16) { + return AscendType::BF16; + } else { + return AscendType::FP16; + } +} + +void bgmv_shrink(at::Tensor &x, at::Tensor &weight, at::Tensor &indices, at::Tensor &y, double scale) +{ + at::ScalarType scalar_type = x.scalar_type(); + TORCH_CHECK(scalar_type == torch::kHalf || scalar_type == torch::kBFloat16, "only support half and bf16"); + TORCH_CHECK(x.dim() == 2, "x should be [batch_size, hidden_in]"); + TORCH_CHECK(weight.dim() == 3 || weight.dim() == 4, + "weight should be [num_loras, hidden_out, hidden_in] or [num_loras, 1, hidden_out, hidden_in]"); + TORCH_CHECK(y.dim() == 2, "y should be [batch_size, hidden_out]"); + TORCH_CHECK(indices.dim() == 1, "indices should be [batch_size]"); + TORCH_CHECK(x.size(0) == y.size(0) && x.size(0) == indices.size(0), + "the first dimension of x, y, indices should be same"); + TORCH_CHECK(x.size(1) > y.size(1), "hidden in should be greater than hidden out"); + void* x_ptr = x.data_ptr(); + void* weight_ptr = weight.data_ptr(); + void* indices_ptr = indices.data_ptr(); + int indices_size = indices.size(0); + void* y_ptr = y.data_ptr(); + int batch_size = x.size(0); + int input_hidden_token = x.size(1); + uint32_t lora_rank = y.size(1); + float scale_f = static_cast(scale); + aclrtStream stream = c10_npu::getCurrentNPUStream().stream(); + at_npu::native::OpCommand cmd; + cmd.Name("bgmv_shrink"); + cmd.SetCustomHandler([scalar_type, stream, x_ptr, weight_ptr, indices_ptr, indices_size, y_ptr, batch_size, input_hidden_token, + lora_rank, scale_f]() -> int { + auto dtype = get_dtype_from_torch(scalar_type); + int device_id = 0; + int64_t aiv_num = 0; + TORCH_CHECK(aclGetDeviceCapability(device_id, ACL_DEVICE_INFO_VECTOR_CORE_NUM, &aiv_num) == ACL_SUCCESS); + int num_tokens_per_core = (batch_size + aiv_num - 1) / aiv_num; + TORCH_CHECK("num_tokens_per_core != 0", "num_tokens_per_core should not be 0"); + bgmv_shrink_impl(dtype, stream, x_ptr, weight_ptr, indices_ptr, indices_size, y_ptr, batch_size, num_tokens_per_core, + input_hidden_token, lora_rank, scale_f); + return 0; + }); + cmd.Run(); + return; +} + +at::Tensor bgmv_expand(at::Tensor &x, at::Tensor &weight, at::Tensor &indices, at::Tensor &y, + int64_t slice_offset, int64_t slice_size) +{ + at::ScalarType scalar_type = y.scalar_type(); + TORCH_CHECK(scalar_type == torch::kHalf || scalar_type == torch::kBFloat16, "only support half and bf16"); + TORCH_CHECK(x.dim() == 2, "x should be [batch_size, hidden_in]"); + TORCH_CHECK(weight.dim() == 3 || weight.dim() == 4, + "weight should be [num_loras, hidden_out, hidden_in] or [num_loras, 1, hidden_out, hidden_in]"); + TORCH_CHECK(y.dim() == 2, "y should be [batch_size, hidden_out]"); + TORCH_CHECK(indices.dim() == 1, "indices should be [batch_size]"); + TORCH_CHECK(x.size(0) == y.size(0) && x.size(0) == indices.size(0), + "the first dimension of x, y, indices should be same"); + TORCH_CHECK(x.size(1) <= slice_size, "hidden in should be smaller than hidden out"); + TORCH_CHECK(slice_offset >= 0, "slice offset should be no smaller than 0"); + TORCH_CHECK((slice_size + slice_offset) <= y.size(1), + "slice_size + slice_offset should be smaller than the second dimension of y") + + at::Tensor y_out = y; + void* x_ptr = x.data_ptr(); + void* weight_ptr = weight.data_ptr(); + void* indices_ptr = indices.data_ptr(); + int indices_size = indices.size(0); + void* y_ptr = y.data_ptr(); + void* y_out_ptr = y_out.data_ptr(); + int batch_size = x.size(0); + int lora_rank = x.size(1); + int output_full_dim = y.size(1); + aclrtStream stream = c10_npu::getCurrentNPUStream().stream(); + at_npu::native::OpCommand cmd; + cmd.Name("bgmv_expand"); + cmd.SetCustomHandler([scalar_type, stream, x_ptr, weight_ptr, indices_ptr, indices_size, y_ptr, y_out_ptr, batch_size, lora_rank, + slice_offset, slice_size, output_full_dim]() -> int { + auto dtype = get_dtype_from_torch(scalar_type); + int device_id = 0; + int64_t aiv_num = 0; + TORCH_CHECK(aclGetDeviceCapability(device_id, ACL_DEVICE_INFO_VECTOR_CORE_NUM, &aiv_num) == ACL_SUCCESS); + int num_tokens_per_core = (batch_size + aiv_num - 1) / aiv_num; + TORCH_CHECK("num_tokens_per_core != 0", "num_tokens_per_core should not be 0"); + bgmv_expand_impl(dtype, stream, x_ptr, weight_ptr, indices_ptr, indices_size, y_ptr, y_out_ptr, batch_size, + num_tokens_per_core, lora_rank, slice_size, slice_offset, output_full_dim); + return 0; + }); + cmd.Run(); + return y_out; +} + +void sgmv_shrink(at::Tensor &x, at::Tensor &weight, at::Tensor &lora_indices, at::Tensor &seq_len, + at::Tensor &y, double scale) +{ + at::ScalarType scalar_type = x.scalar_type(); + TORCH_CHECK(scalar_type == torch::kHalf || scalar_type == torch::kBFloat16, "only support half and bf16"); + TORCH_CHECK(x.dim() == 2, "x should be [batch_size, hidden_in]"); + TORCH_CHECK(weight.dim() == 3 || weight.dim() == 4, + "weight should be [num_loras, hidden_out, hidden_in] or [num_loras, 1, hidden_out, hidden_in]"); + TORCH_CHECK(y.dim() == 2, "y should be [batch_size, hidden_out]"); + TORCH_CHECK(x.size(1) > y.size(1), "hidden in should be greater than hidden out"); + void* x_ptr = x.data_ptr(); + void* weight_ptr = weight.data_ptr(); + void* lora_indices_ptr = lora_indices.data_ptr(); + void* seq_len_ptr = seq_len.data_ptr(); + int lora_indices_size = lora_indices.size(0); + int seq_len_size = seq_len.size(0); + void* y_ptr = y.data_ptr(); + int batch_size = x.size(0); + int input_hidden_token = x.size(1); + uint32_t lora_rank = y.size(1); + float scale_f = static_cast(scale); + aclrtStream stream = c10_npu::getCurrentNPUStream().stream(); + at_npu::native::OpCommand cmd; + cmd.Name("sgmv_shrink"); + cmd.SetCustomHandler([scalar_type, stream, x_ptr, weight_ptr, lora_indices_ptr, lora_indices_size, + seq_len_ptr, seq_len_size, y_ptr, + batch_size, input_hidden_token, lora_rank, scale_f]() -> int { + auto dtype = get_dtype_from_torch(scalar_type); + int device_id = 0; + int64_t aiv_num = 0; + TORCH_CHECK(aclGetDeviceCapability(device_id, ACL_DEVICE_INFO_VECTOR_CORE_NUM, &aiv_num) == ACL_SUCCESS); + int num_tokens_per_core = (batch_size + aiv_num - 1) / aiv_num; + TORCH_CHECK("num_tokens_per_core != 0", "num_tokens_per_core should not be 0"); + sgmv_shrink_impl(dtype, stream, x_ptr, weight_ptr, lora_indices_ptr, lora_indices_size, seq_len_ptr, seq_len_size, + y_ptr, batch_size, + num_tokens_per_core, input_hidden_token, lora_rank, scale_f); + return 0; + }); + cmd.Run(); + return; +} + +at::Tensor sgmv_expand(at::Tensor &x, at::Tensor &weight, at::Tensor &lora_indices, at::Tensor &seq_len, + at::Tensor &y, int64_t slice_offset, int64_t slice_size) +{ + at::ScalarType scalar_type = y.scalar_type(); + TORCH_CHECK(scalar_type == torch::kHalf || scalar_type == torch::kBFloat16, "only support half and bf16"); + TORCH_CHECK(x.dim() == 2, "x should be [batch_size, hidden_in]"); + TORCH_CHECK(weight.dim() == 3 || weight.dim() == 4, + "weight should be [num_loras, hidden_out, hidden_in] or [num_loras, 1, hidden_out, hidden_in]"); + TORCH_CHECK(y.dim() == 2, "y should be [batch_size, hidden_out]"); + TORCH_CHECK(x.size(1) <= slice_size, "hidden in should be smaller than hidden out"); + TORCH_CHECK(slice_offset >= 0, "slice offset should be no smaller than 0"); + TORCH_CHECK((slice_size + slice_offset) <= y.size(1), + "slice_size + slice_offset should be smaller than the second dimension of y") + + at::Tensor y_out = y; + void* x_ptr = x.data_ptr(); + void* weight_ptr = weight.data_ptr(); + void* lora_indices_ptr = lora_indices.data_ptr(); + void* seq_len_ptr = seq_len.data_ptr(); + int lora_indices_size = lora_indices.size(0); + int seq_len_size = seq_len.size(0); + void* y_ptr = y.data_ptr(); + void* y_out_ptr = y_out.data_ptr(); + int batch_size = x.size(0); + int lora_rank = x.size(1); + int output_full_dim = y.size(1); + aclrtStream stream = c10_npu::getCurrentNPUStream().stream(); + at_npu::native::OpCommand cmd; + cmd.Name("sgmv_expand"); + cmd.SetCustomHandler([scalar_type, stream, x_ptr, weight_ptr, lora_indices_ptr, lora_indices_size, seq_len_ptr, seq_len_size, y_ptr, y_out_ptr, + batch_size, lora_rank, slice_offset, slice_size, output_full_dim]() -> int { + auto dtype = get_dtype_from_torch(scalar_type); + int device_id = 0; + int64_t aiv_num = 0; + TORCH_CHECK(aclGetDeviceCapability(device_id, ACL_DEVICE_INFO_VECTOR_CORE_NUM, &aiv_num) == ACL_SUCCESS); + int num_tokens_per_core = (batch_size + aiv_num - 1) / aiv_num; + TORCH_CHECK("num_tokens_per_core != 0", "num_tokens_per_core should not be 0"); + sgmv_expand_impl(dtype, stream, x_ptr, weight_ptr, lora_indices_ptr, lora_indices_size, seq_len_ptr, seq_len_size, y_ptr, y_out_ptr, + batch_size, num_tokens_per_core, lora_rank, slice_size, slice_offset, output_full_dim); + return 0; + }); + cmd.Run(); + return y_out; +} +#endif + +at::Tensor convert_hamming_dist_top_k_output(const at::Tensor &hashq, + const at::Tensor &hashkCache, + const c10::optional& indices) { + if (indices.has_value()) { + return indices.value(); + } + uint32_t MAX_BLOCK_PER_REQ_INHSA = 512; + + auto n_bs = hashq.size(0); + auto n_kv_heads = hashkCache.size(1); + auto n_max_kv = MAX_BLOCK_PER_REQ_INHSA; + at::Tensor res = at::empty({n_bs, n_kv_heads, n_max_kv}, torch::TensorOptions().dtype(torch::kInt32).device(hashq.device())); + return res; +} + +at::Tensor npu_hamming_dist_top_k(const at::Tensor &hashq, + const at::Tensor &hashkCache, + const at::Tensor& hashkCacheRope, + const at::Tensor &topN, + const at::Tensor &seqLen, + const c10::optional &chunkSize, + const c10::optional maxSeqLen, + const c10::optional sink, + const c10::optional recent, + const c10::optional supportOffload, + const c10::optional &blockTable, + const c10::optional &mask, + const c10::optional& indices) { + + auto&& maxSeqLen_ = maxSeqLen.value_or(0); + auto&& sink_ = sink.value_or(0); + auto&& recent_ = recent.value_or(0); + auto&& supportOffload_ = supportOffload.value_or(0); + + at::Tensor out = convert_hamming_dist_top_k_output(hashq, hashkCache, indices); + EXEC_NPU_CMD(aclnnHammingDistTopK, hashq, hashkCache, topN, seqLen, chunkSize, blockTable, indices, hashkCacheRope, mask, maxSeqLen_, sink_, recent_, supportOffload_, out); + return out; +} + +at::Tensor npu_reshape_and_cache_bnsd(const at::Tensor& hashq, + const at::Tensor& hashkCache, + const at::Tensor& slotMapping, + const at::Tensor& seqLen, + const at::Tensor& hashkCacheOut) { + EXEC_NPU_CMD(aclnnReshapeAndCacheBnsd, hashq, hashkCache, slotMapping, seqLen, hashkCacheOut); + return hashkCacheOut; +} + +at::Tensor npu_sign_bits_pack(const at::Tensor& input, + const int64_t size) { + int64_t ySize = (input.size(0) + 7) / 8; + int64_t outDim = 0; + if (size != 0) { + outDim = ySize / size; + } + + at::Tensor out = torch::empty({size, outDim}, torch::TensorOptions().dtype(torch::kUInt8).device(input.device())); + EXEC_NPU_CMD(aclnnSignBitsPack, input, size, out); + return out; +} + +std::tuple npu_gemma_rms_norm( + const at::Tensor& x, + const at::Tensor& gamma, + double epsilon) +{ + int64_t dim_x = x.dim(); + int64_t dim_gamma = gamma.dim(); + int64_t diff = dim_x - dim_gamma; + std::vector new_shape; + at::Tensor rstd; + if (diff > 0) { + new_shape.reserve(dim_x); + auto x_sizes = x.sizes(); + for (int64_t i = 0; i < diff; ++i) { + new_shape.push_back(x_sizes[i]); + } + for (int64_t i = 0; i < dim_gamma; ++i) { + new_shape.push_back(1); + } + } else { + new_shape.assign(dim_x, 1); + } + rstd = at::empty(new_shape, x.options().dtype(at::kFloat)); + at::Tensor y = at::empty(x.sizes(), x.options()); + EXEC_NPU_CMD(aclnnGemmaRmsNorm, x, gamma, epsilon, y, rstd); + return std::tuple(y, rstd); +} + +void transpose_kv_cache_by_block( + const at::TensorList &kCache, + const at::TensorList &vCache, + const at::Tensor &blockIDs, + int64_t blockSize, + int64_t headNum, + int64_t headDim, + int64_t splitNum, + int64_t layerNum) +{ + + EXEC_NPU_CMD(aclnnTransposeKvCacheByBlock, kCache, vCache, blockIDs, + blockSize, headNum, headDim, splitNum, layerNum); + +} + +void device_print(c10::string_view msg) +{ + auto payload = std::make_unique(); + payload->message = std::string(msg); + enqueue_device_print(std::move(payload), c10_npu::getCurrentNPUStream().stream()); +} + +void device_print(const at::Tensor& tensor) +{ + TORCH_CHECK(tensor.defined(), "tensor must be defined"); + TORCH_CHECK( + tensor.device().is_cpu() || + tensor.device().type() == c10::DeviceType::PrivateUse1, + "device_print only supports CPU and NPU tensors, but got device ", + tensor.device()); + + auto payload = std::make_unique(); + if (tensor.device().is_cpu()) { + payload->host_tensor_snapshot = tensor.contiguous().clone(); + enqueue_device_print(std::move(payload), + c10_npu::getCurrentNPUStream().stream()); + return; + } + + const c10_npu::OptionalNPUGuard npu_guard(tensor.device()); + aclrtStream stream = c10_npu::getCurrentNPUStream().stream(); + at::Tensor contiguous_tensor = tensor.contiguous(); + payload->host_tensor_snapshot = at::empty_like( + contiguous_tensor, + contiguous_tensor.options().device(at::kCPU).pinned_memory(true)); + + const size_t num_bytes = contiguous_tensor.numel() * + contiguous_tensor.element_size(); + const aclError memcpy_ret = aclrtMemcpyAsync( + payload->host_tensor_snapshot.data_ptr(), num_bytes, + contiguous_tensor.data_ptr(), num_bytes, ACL_MEMCPY_DEVICE_TO_HOST, stream); + TORCH_CHECK(memcpy_ret == ACL_SUCCESS, + "aclrtMemcpyAsync failed, error code: ", memcpy_ret); + + // The D2H copy and host callback are queued on the same stream so the + // callback prints only after the host snapshot is ready. + enqueue_device_print(std::move(payload), stream); +} + +std::tuple +npu_copy_and_expand_eagle_inputs( + const at::Tensor &target_token_ids, + const at::Tensor &target_positions, + const at::Tensor &next_token_ids, + const at::Tensor &query_start_loc, + const at::Tensor &query_end_loc, + int64_t padding_token_id, + int64_t parallel_drafting_token_id, + int64_t num_padding_slots_per_request, + bool shift_input_ids, + int64_t total_draft_tokens) +{ + int64_t total_input_tokens = target_token_ids.size(0); + int64_t num_reqs = query_start_loc.size(0) - 1; + + auto device = target_token_ids.device(); + at::Tensor out_input_ids = at::zeros({total_draft_tokens}, at::dtype(at::kInt).device(device)); + at::Tensor out_positions = at::zeros({total_draft_tokens}, at::dtype(at::kInt).device(device)); + at::Tensor out_is_rejected_token_mask = at::zeros({total_draft_tokens}, at::dtype(at::kChar).device(device)); + at::Tensor out_is_masked_token_mask = at::zeros({total_draft_tokens}, at::dtype(at::kChar).device(device)); + at::Tensor out_new_token_indices = at::zeros({num_reqs * num_padding_slots_per_request}, at::dtype(at::kInt).device(device)); + at::Tensor out_hidden_state_mapping = at::zeros({total_input_tokens}, at::dtype(at::kInt).device(device)); + + EXEC_NPU_CMD(aclnnCopyAndExpandEagleInputs, + target_token_ids, target_positions, next_token_ids, query_start_loc, query_end_loc, + padding_token_id, parallel_drafting_token_id, num_padding_slots_per_request, + shift_input_ids, total_input_tokens, + out_input_ids, out_positions, out_is_rejected_token_mask, out_is_masked_token_mask, + out_new_token_indices, out_hidden_state_mapping); + + return {out_input_ids, out_positions, out_is_rejected_token_mask, out_is_masked_token_mask, + out_new_token_indices, out_hidden_state_mapping}; +} + +at::Tensor npu_causal_conv1d_custom( + const at::Tensor& output, + const at::Tensor& x, + const at::Tensor& weight, + const at::Tensor& conv_state, + const c10::optional& bias_opt, + const c10::optional& query_start_loc_opt, + const c10::optional& cache_indices_opt, + const c10::optional& initial_state_mode_opt, + const c10::optional& num_accepted_tokens_opt, + int64_t activation_mode, + int64_t pad_slot_id, + int64_t run_mode) +{ + EXEC_NPU_CMD(aclnnCausalConv1d, + x, + weight, + bias_opt, + conv_state, + query_start_loc_opt, + cache_indices_opt, + initial_state_mode_opt, + num_accepted_tokens_opt, + activation_mode, + pad_slot_id, + run_mode, + output + ); + + return output; +} + +// It is expected that further improvements will be made after it is incorporated into CANN on June 30th. +std::vector moe_grouped_matmul( + at::Tensor x, + at::Tensor weight, + const at::Tensor& group_list, + int64_t split_item, + int64_t group_type, + int64_t group_list_type +) +{ + bool transpose_weight = false; + bool weight_nz = true; + + at::TensorList x_list = at::TensorList(x); + at::TensorList weight_list = at::TensorList(weight); + std::vector y; + c10::TensorOptions options = x_list[0].options().dtype(x[0].scalar_type()); + auto m = x_list[0].sizes()[0]; + auto n = weight_list[0].sizes()[1]; + if (!transpose_weight) { + n = weight_list[0].sizes()[2]; + } + at::Tensor y_0 = at::empty(at::IntArrayRef{m, n}, options); + y.emplace_back(y_0); + at::TensorList result = at::TensorList(y); + + EXEC_NPU_CMD(aclnnMoeGroupedMatmulWeightNz, + x_list, weight_list, group_list, transpose_weight, result); + + return y; +} + +std::tuple moe_gating_top_k_hash( + const at::Tensor& x, + int64_t k, + const c10::optional& bias_opt, + const c10::optional& input_ids_opt, + const c10::optional& tid2eid_opt, + int64_t k_group, + int64_t group_count, + double routed_scaling_factor, + double eps, + int64_t group_select_mode, + int64_t renorm, + int64_t norm_type, + bool out_flag) +{ + + TORCH_CHECK(x.dim() == 2, "x must be 2D, but got dim=", x.dim()); + TORCH_CHECK( + x.scalar_type() == at::kHalf || x.scalar_type() == at::kFloat || x.scalar_type() == at::kBFloat16, + "x dtype must be float16/float32/bfloat16, but got ", x.scalar_type()); + + TORCH_CHECK(k > 0, "k must be > 0, but got k=", k); + TORCH_CHECK(k_group >= 1, "k_group must be >= 1, but got k_group=", k_group); + TORCH_CHECK(group_count >= 1, "group_count must be >= 1, but got group_count=", group_count); + + TORCH_CHECK(group_select_mode == 0 || group_select_mode == 1, + "group_select_mode must be 0 or 1, but got ", group_select_mode); + TORCH_CHECK(renorm == 0, + "renorm can only be 0 currently, but got ", renorm); + TORCH_CHECK(norm_type == 0 || norm_type == 1 || norm_type ==2, + "norm_type must be 0 (softmax) or 1 (sigmoid) or 2 (softplus), but got ", norm_type); + + TORCH_CHECK(eps > 0.0, "eps must be > 0, but got ", eps); + TORCH_CHECK(routed_scaling_factor > 0.0, + "routed_scaling_factor must be > 0, but got ", routed_scaling_factor); + + const auto sizes = x.sizes(); + const int64_t rows = sizes[0]; + const int64_t expert_num = sizes[1]; + + TORCH_CHECK(expert_num > 0, "expert_num must be > 0"); + TORCH_CHECK(expert_num <= 2048, + "expert_num (E) must be <= 2048, but got ", expert_num); + + if (bias_opt.has_value() && bias_opt->defined()) { + const auto& bias = *bias_opt; + TORCH_CHECK(bias.dim() == 1, "bias must be 1D, but got dim=", bias.dim()); + TORCH_CHECK(bias.size(0) == expert_num, + "bias.size(0) must equal expert_num. bias.size(0)=", + bias.size(0), ", expert_num=", expert_num); + TORCH_CHECK(bias.scalar_type() == x.scalar_type(), + "bias dtype must equal x dtype. x=", x.scalar_type(), + ", bias=", bias.scalar_type()); + } + + if (input_ids_opt.has_value() && input_ids_opt->defined()) { + const auto& input_ids = *input_ids_opt; + TORCH_CHECK(input_ids.scalar_type() == at::kInt || input_ids.scalar_type() == at::kLong, + "input_ids dtype must be int32 or int64, but got ", input_ids.scalar_type()); + TORCH_CHECK(input_ids.numel() == rows, + "input_ids.numel() must equal x.size(0). input_ids.numel()=", + input_ids.numel(), ", rows=", rows); + } + + if (tid2eid_opt.has_value() && tid2eid_opt->defined()) { + const auto& tid2eid = *tid2eid_opt; + TORCH_CHECK(tid2eid.scalar_type() == at::kInt || tid2eid.scalar_type() == at::kLong, + "tid2eid dtype must be int32 or int64, but got ", tid2eid.scalar_type()); + TORCH_CHECK(tid2eid.dim() >= 1, "tid2eid must have dim>=1, but got dim=", tid2eid.dim()); + } + + const at::Tensor& bias = c10::value_or_else(bias_opt, [] { return at::Tensor(); }); + const at::Tensor& input_ids = c10::value_or_else(input_ids_opt, [] { return at::Tensor(); }); + const at::Tensor& tid2eid = c10::value_or_else(tid2eid_opt, [] { return at::Tensor(); }); + + at::Tensor y = at::empty({rows, k}, x.options()); + at::Tensor expert_idx = at::empty({rows, k}, x.options().dtype(at::kInt)); + at::Tensor out = at::empty({rows, expert_num}, x.options().dtype(at::kFloat)); + + EXEC_NPU_CMD(aclnnMoeGatingTopKHash, + x, + bias, + input_ids, + tid2eid, + k, + k_group, + group_count, + routed_scaling_factor, + eps, + group_select_mode, + renorm, + norm_type, + out_flag, + y, + expert_idx, + out); + + return {y, expert_idx, out}; +} + +std::vector is_contiguous_axes(const at::Tensor &tensor) +{ + auto sizes = tensor.sizes(); + auto strides = tensor.strides(); + int64_t ndim = sizes.size(); + + if (ndim == 0) { + return {}; + } + std::vector result(ndim, false); + + std::vector contiguous_stride(ndim, 1); + for (int64_t i = ndim - 2; i >= 0; i--) { + contiguous_stride[i] = contiguous_stride[i + 1] * sizes[i + 1]; + } + + + for (int64_t i = 0; i < ndim; i++) { + result[i] = (strides[i] == contiguous_stride[i]); + } + return result; +} + +std::tuple construct_compressor_output_tensor(const at::Tensor &x, const at::Tensor &norm_weight, + const at::Tensor &rope_sin, int64_t cmp_ratio, int64_t coff) +{ + constexpr int DIM_3 = 3; + auto x_dim = x.dim(); + at::SmallVector cmp_kv_size; + at::Tensor cmp_kv; + auto cmp_s = 0; + if (x_dim == DIM_3) { + cmp_s = (x.size(1) + cmp_ratio - 1) / cmp_ratio; + cmp_kv_size = {x.size(0), cmp_s, norm_weight.size(0)}; + } else { + cmp_s = rope_sin.size(0); + cmp_kv_size = {cmp_s, norm_weight.size(0)}; + } + + cmp_kv = at::empty(cmp_kv_size, x.options().dtype(x.dtype())); + + return std::tuple(cmp_kv); +} + + +std::tuple compressor(const at::Tensor &x, const at::Tensor &wkv, const at::Tensor &wgate, + at::Tensor &state_cache, const at::Tensor &ape, const at::Tensor &norm_weight, + const at::Tensor &rope_sin, const at::Tensor &rope_cos, + const c10::optional &state_block_table, + const c10::optional &cu_seqlens, const c10::optional &seqused, + const c10::optional &start_pos, int64_t rope_head_dim, int64_t cmp_ratio, + int64_t coff, double norm_eps, int64_t rotary_mode, int64_t cache_mode) +{ + constexpr int CONTINUOUS = 1; + constexpr int32_t DIM_1 = 1; + constexpr int32_t DIM_2 = 2; + constexpr int32_t DIM_3 = 3; + constexpr int32_t VALUE_0 = 0; + auto x_dim = x.dim(); + TORCH_CHECK(x_dim == DIM_2 || x_dim == DIM_3, "x dim num[", x_dim, "] should be 2 or 3"); + + TORCH_CHECK(norm_weight.defined(), "Check norm_weight != nullptr failed"); + auto norm_weight_dim = norm_weight.dim(); + TORCH_CHECK(norm_weight_dim == DIM_1, "norm_weight dim num[", norm_weight_dim, "] should be 1"); + + TORCH_CHECK(rope_sin.defined(), "Check rope_sin != nullptr failed"); + auto rope_sin_dim = rope_sin.dim(); + TORCH_CHECK(rope_sin_dim == x_dim, "rope_sin dim num[", rope_sin_dim, "] should be equal to x dim num[", x_dim, + "]"); + + TORCH_CHECK(cmp_ratio > VALUE_0, "cmp_ratio should be greater than 0"); + + std::tuple output = construct_compressor_output_tensor(x, norm_weight, rope_sin, cmp_ratio, coff); + at::Tensor cmp_kv = std::get<0>(output); + + auto state_cache_dim = state_cache.dim(); + TORCH_CHECK(state_cache_dim == DIM_3, "state_cache dim num[", state_cache_dim, "] should be 3"); + auto contiguous_axes_result = is_contiguous_axes(state_cache); + // if (cache_mode == CONTINUOUS) { + // TORCH_CHECK(contiguous_axes_result[0] && contiguous_axes_result[1] && contiguous_axes_result[2], + // "when cache_mode == ", cache_mode, ", state_cache must be contiguous on all axes"); + // } + int64_t state_cache_stride_dim0 = state_cache.stride(0); + + EXEC_NPU_CMD(aclnnCompressor, x, wkv, wgate, state_cache, ape, norm_weight, rope_sin, rope_cos, + state_block_table, cu_seqlens, seqused, start_pos, rope_head_dim, cmp_ratio, coff, norm_eps, + rotary_mode, cache_mode, state_cache_stride_dim0, cmp_kv); + + return std::tuple(cmp_kv); +} + +void check_compressor_metadata_common( + const at::Tensor &rope_cos, const at::Tensor &rope_sin, const at::Tensor &cu_seqlens, + const at::Tensor &start_pos, const at::Tensor &kv_block_table, int64_t kv_block_size, + int64_t slot_mapping_format, int64_t compress_ratio, int64_t num_reqs_actual) +{ + constexpr int64_t DIM_2 = 2; + constexpr int64_t VALUE_0 = 0; + + TORCH_CHECK(rope_cos.defined() && rope_sin.defined(), "rope_cos and rope_sin should be defined"); + TORCH_CHECK(rope_cos.dim() == DIM_2 && rope_sin.dim() == DIM_2, + "rope_cos and rope_sin should be 2D tensors"); + TORCH_CHECK(rope_cos.scalar_type() == rope_sin.scalar_type(), + "rope_cos and rope_sin should have same dtype"); + TORCH_CHECK(rope_cos.size(0) == rope_sin.size(0) && rope_cos.size(1) == rope_sin.size(1), + "rope_cos and rope_sin should have same shape"); + TORCH_CHECK(rope_cos.size(0) > VALUE_0 && rope_cos.size(1) > VALUE_0, + "rope_cos shape should be non-empty"); + TORCH_CHECK(cu_seqlens.defined() && cu_seqlens.dim() == 1, "cu_seqlens should be a 1D tensor"); + TORCH_CHECK(start_pos.defined() && start_pos.dim() == 1, "start_pos should be a 1D tensor"); + TORCH_CHECK(kv_block_table.defined() && kv_block_table.dim() == DIM_2, "kv_block_table should be a 2D tensor"); + TORCH_CHECK(kv_block_size > VALUE_0, "kv_block_size should be greater than 0"); + TORCH_CHECK(compress_ratio > VALUE_0, "compress_ratio should be greater than 0"); + TORCH_CHECK(slot_mapping_format == DSA_SLOT_MAPPING_BLOCK_OFFSET || slot_mapping_format == DSA_SLOT_MAPPING_FLAT, + "slot_mapping_format should be 1(flat) or 2(block_offset), but got ", slot_mapping_format); + TORCH_CHECK(num_reqs_actual > VALUE_0, "num_reqs_actual should be greater than 0"); + TORCH_CHECK(cu_seqlens.size(0) > num_reqs_actual, + "cu_seqlens dim0 should be greater than num_reqs_actual"); + TORCH_CHECK(start_pos.size(0) >= num_reqs_actual, + "start_pos dim0 should be greater than or equal to num_reqs_actual"); + TORCH_CHECK(kv_block_table.size(0) >= num_reqs_actual, + "kv_block_table dim0 should be greater than or equal to num_reqs_actual"); +} + +void check_compressor_metadata_outputs( + const at::Tensor &rope_cos, const at::Tensor &compress_cos, const at::Tensor &compress_sin, + const at::Tensor &slot_mapping, int64_t slot_mapping_format) +{ + constexpr int64_t DIM_2 = 2; + constexpr int64_t VALUE_0 = 0; + + TORCH_CHECK(compress_cos.defined() && compress_sin.defined() && slot_mapping.defined(), + "compress_cos, compress_sin, and slot_mapping should be defined"); + TORCH_CHECK(compress_cos.dim() >= DIM_2, "compress_cos dim num should be at least 2"); + TORCH_CHECK(compress_sin.dim() == compress_cos.dim(), "compress_cos and compress_sin should have same dim num"); + TORCH_CHECK(compress_cos.size(0) > VALUE_0, "compress_cos dim0 should be greater than 0"); + TORCH_CHECK(compress_cos.size(compress_cos.dim() - 1) == rope_cos.size(1), + "compress_cos last dim should match rope dim"); + for (int64_t dim_idx = 0; dim_idx < compress_cos.dim(); ++dim_idx) { + TORCH_CHECK(compress_sin.size(dim_idx) == compress_cos.size(dim_idx), + "compress_cos and compress_sin should have same shape"); + } + TORCH_CHECK(compress_cos.scalar_type() == rope_cos.scalar_type() && + compress_sin.scalar_type() == rope_cos.scalar_type(), + "compress outputs should have same dtype as rope_cos"); + TORCH_CHECK(slot_mapping.scalar_type() == at::kInt, "slot_mapping dtype should be int32"); + if (slot_mapping_format == DSA_SLOT_MAPPING_BLOCK_OFFSET) { + TORCH_CHECK(slot_mapping.dim() == DIM_2 && slot_mapping.size(0) == compress_cos.size(0) && + slot_mapping.size(1) == DIM_2, + "block_offset slot_mapping should have shape [num_rows, 2]"); + } else { + TORCH_CHECK(slot_mapping.dim() == 1 && slot_mapping.size(0) == compress_cos.size(0), + "flat slot_mapping should have shape [num_rows]"); + } +} + +std::tuple compressor_metadata( + const at::Tensor &rope_cos, const at::Tensor &rope_sin, const at::Tensor &cu_seqlens, + const at::Tensor &start_pos, const at::Tensor &kv_block_table, int64_t kv_block_size, + int64_t slot_mapping_format, int64_t compress_ratio, int64_t num_compressed_tokens, int64_t num_reqs_actual) +{ + constexpr int64_t VALUE_0 = 0; + + check_compressor_metadata_common( + rope_cos, rope_sin, cu_seqlens, start_pos, kv_block_table, kv_block_size, slot_mapping_format, compress_ratio, + num_reqs_actual); + TORCH_CHECK(num_compressed_tokens > VALUE_0, "num_compressed_tokens should be greater than 0"); + + at::SmallVector rope_output_size = {num_compressed_tokens, 1, 1, rope_cos.size(1)}; + at::Tensor compress_cos = at::empty(rope_output_size, rope_cos.options()); + at::Tensor compress_sin = at::empty(rope_output_size, rope_sin.options()); + + at::SmallVector slot_mapping_size; + if (slot_mapping_format == DSA_SLOT_MAPPING_BLOCK_OFFSET) { + slot_mapping_size = {num_compressed_tokens, 2}; + } else { + slot_mapping_size = {num_compressed_tokens}; + } + at::Tensor slot_mapping = at::empty(slot_mapping_size, kv_block_table.options().dtype(at::kInt)); + + EXEC_NPU_CMD(aclnnCompressorMetadata, rope_cos, rope_sin, cu_seqlens, start_pos, kv_block_table, + kv_block_size, slot_mapping_format, compress_ratio, num_reqs_actual, compress_cos, compress_sin, + slot_mapping); + return std::make_tuple(compress_cos, compress_sin, slot_mapping); +} + +std::tuple compressor_metadata_out( + const at::Tensor &rope_cos, const at::Tensor &rope_sin, const at::Tensor &cu_seqlens, + const at::Tensor &start_pos, const at::Tensor &kv_block_table, int64_t kv_block_size, + int64_t slot_mapping_format, int64_t compress_ratio, int64_t num_reqs_actual, at::Tensor &compress_cos, + at::Tensor &compress_sin, at::Tensor &slot_mapping) +{ + check_compressor_metadata_common( + rope_cos, rope_sin, cu_seqlens, start_pos, kv_block_table, kv_block_size, slot_mapping_format, compress_ratio, + num_reqs_actual); + check_compressor_metadata_outputs(rope_cos, compress_cos, compress_sin, slot_mapping, slot_mapping_format); + + EXEC_NPU_CMD(aclnnCompressorMetadata, rope_cos, rope_sin, cu_seqlens, start_pos, kv_block_table, + kv_block_size, slot_mapping_format, compress_ratio, num_reqs_actual, compress_cos, compress_sin, + slot_mapping); + return std::make_tuple(compress_cos, compress_sin, slot_mapping); +} + +std::tuple construct_quant_lightning_indexer_output_tensor(const at::Tensor& query, const at::Tensor& key, + int64_t sparse_count, std::string query_layout_str, + std::string key_layout_str, bool return_value) +{ + constexpr int64_t SIZE = 8; + constexpr int64_t DIM_0 = 0; + constexpr int64_t DIM_1 = 1; + constexpr int64_t DIM_2 = 2; + constexpr int64_t DIM_3 = 3; + at::SmallVector output_size; + for (size_t i = 0; i < query.sizes().size(); i++) { + TORCH_CHECK(query.size(i) > 0, "All values within query's shape should be greater " + "than 0, but shape[", i, "] is ", query.size(i)); + } + for (size_t i = 0; i < key.sizes().size(); i++) { + TORCH_CHECK(key.size(i) > 0, "All values within key's shape should be greater " + "than 0, but shape[", i, "] is ", key.size(i)); + } + TORCH_CHECK(sparse_count > 0, "sparse count should be greater than 0, but now is ", sparse_count); + int64_t keyHeadNum = (key_layout_str == "TND")? key.size(DIM_1) : key.size(DIM_2); + if (query_layout_str == "BSND") { + output_size = {query.size(DIM_0), query.size(DIM_1), keyHeadNum, sparse_count}; + } else { + output_size = {query.size(DIM_0), keyHeadNum, sparse_count}; + } + at::Tensor sparse_indices_out = at::empty(output_size, query.options().dtype(at::kInt)); + at::Tensor sparse_values_out; + if (return_value) { + sparse_values_out = at::empty(output_size, query.options().dtype(at::kFloat)); + } else { + sparse_values_out = at::empty({0}, query.options().dtype(at::kFloat)); + } + + return std::tuple(sparse_indices_out, sparse_values_out); +} + +std::tuple npu_vllm_quant_lightning_indexer_npu( + const at::Tensor &query, const at::Tensor &key, const at::Tensor &weights, + const at::Tensor &query_dequant_scale, const at::Tensor &key_dequant_scale, + int64_t query_quant_mode, int64_t key_quant_mode, + const c10::optional &actual_seq_lengths_query, + const c10::optional &actual_seq_lengths_key, + const c10::optional &block_table, + const c10::optional &metadata, + c10::string_view layout_query, c10::string_view layout_key, int64_t sparse_count, + int64_t sparse_mode, int64_t pre_tokens, int64_t next_tokens, int64_t cmp_ratio, bool return_value) +{ + std::string query_layout_str = std::string(layout_query); + std::string key_layout_str = std::string(layout_key); + + std::tuple quant_lightning_indexer_output = construct_quant_lightning_indexer_output_tensor( + query, key, sparse_count, query_layout_str, key_layout_str, return_value); + at::Tensor sparse_indices_out = std::get<0>(quant_lightning_indexer_output); + at::Tensor sparse_values_out = std::get<1>(quant_lightning_indexer_output); + char *query_layout_ptr = const_cast(query_layout_str.c_str()); + char *key_layout_ptr = const_cast(key_layout_str.c_str()); + int64_t stride = key.stride(0); + int64_t scale_stride = key_dequant_scale.stride(0); + + if (key_layout_str == "PA_BSND") { + auto contiguous_axes_result_key = is_contiguous_axes(key); + TORCH_CHECK(contiguous_axes_result_key[1] && contiguous_axes_result_key[2], + "key must be contiguous on all axes except axis 0"); + auto contiguous_axes_result_key_scale = is_contiguous_axes(key_dequant_scale); + TORCH_CHECK(contiguous_axes_result_key_scale[1] && contiguous_axes_result_key_scale[2], + "key_dequant_scale must be contiguous on all axes except axis 0"); + } + + EXEC_NPU_CMD(aclnnVllmQuantLightningIndexer, query, + key, weights, query_dequant_scale, key_dequant_scale, actual_seq_lengths_query, actual_seq_lengths_key, + block_table, metadata, query_quant_mode, key_quant_mode, query_layout_ptr, key_layout_ptr, sparse_count, sparse_mode, + pre_tokens, next_tokens, cmp_ratio, return_value, stride, scale_stride, sparse_indices_out, sparse_values_out); + + + return std::tuple(sparse_indices_out, sparse_values_out); +} + +std::tuple construct_output_tensor(const at::Tensor &q, std::string layout, + bool return_softmax_lse) +{ + for (size_t i = 0; i < q.sizes().size(); i++) { + TORCH_CHECK(q.size(i) > 0, + "All values within query's shape should be greater " + "than 0, but shape[", + i, + "] is ", + q.size(i)); + } + at::Tensor output = at::empty(q.sizes(), q.options().dtype(q.dtype())); + at::Tensor softmax_lse; + if (return_softmax_lse) { + std::vector lse_sizes(q.sizes().begin(), q.sizes().end()); + lse_sizes.back() = 1; + softmax_lse = at::empty(lse_sizes, q.options().dtype(c10::ScalarType::Float)); + } else { + softmax_lse = at::empty({0}, q.options().dtype(c10::ScalarType::Float)); + } + return std::tuple(output, softmax_lse); +} + +std::tuple npu_sparse_attn_sharedkv_npu(const at::Tensor &q, const c10::optional &ori_kv, + const c10::optional &cmp_kv, const c10::optional &ori_sparse_indices, + const c10::optional &cmp_sparse_indices, const c10::optional &ori_block_table, + const c10::optional &cmp_block_table, const c10::optional &cu_seqlens_q, + const c10::optional &cu_seqlens_ori_kv, const c10::optional &cu_seqlens_cmp_kv, + const c10::optional &seqused_q, const c10::optional &seqused_kv, + const c10::optional &sinks, const c10::optional &metadata, + double softmax_scale, int64_t cmp_ratio, int64_t ori_mask_mode, int64_t cmp_mask_mode, int64_t ori_win_left, + int64_t ori_win_right, c10::string_view layout_q, c10::string_view layout_kv, bool return_softmax_lse) +{ + std::string layout_q_str = std::string(layout_q); + std::string layout_kv_str = std::string(layout_kv); + std::tuple output = construct_output_tensor(q, layout_q_str, return_softmax_lse); + at::Tensor attn_out = std::get<0>(output); + at::Tensor softmax_lse = std::get<1>(output); + int64_t ori_kv_stride = 0; + int64_t cmp_kv_stride = 0; + if (ori_kv.has_value()){ + const at::Tensor& tmp_kv = *ori_kv; + ori_kv_stride = tmp_kv.stride(0); + } + if (cmp_kv.has_value()){ + const at::Tensor& tmp_kv = *cmp_kv; + cmp_kv_stride = tmp_kv.stride(0); + } + + char *layout_q_ptr = const_cast(layout_q_str.c_str()); + char *layout_kv_ptr = const_cast(layout_kv_str.c_str()); + EXEC_NPU_CMD(aclnnSparseAttnSharedkv, q, ori_kv, cmp_kv, ori_sparse_indices, cmp_sparse_indices, + ori_block_table, cmp_block_table, cu_seqlens_q, cu_seqlens_ori_kv, cu_seqlens_cmp_kv, seqused_q, seqused_kv, sinks, + metadata, softmax_scale, cmp_ratio, ori_mask_mode, cmp_mask_mode, ori_kv_stride, cmp_kv_stride, ori_win_left, ori_win_right, layout_q_ptr, + layout_kv_ptr, return_softmax_lse, attn_out, softmax_lse); + return std::tuple(attn_out, softmax_lse); +} + +auto get_valid_tensor = [](const c10::optional &tensor_opt, at::Device device) { + return tensor_opt.has_value() ? tensor_opt : torch::empty({0}, torch::dtype(torch::kInt32).device(device)); +}; + +at::Tensor npu_sparse_attn_sharedkv_metadata_npu( + int64_t num_heads_q, + int64_t num_heads_kv, + int64_t head_dim, + const c10::optional &cu_seqlens_q, + const c10::optional &cu_seqlens_ori_kv, + const c10::optional &cu_seqlens_cmp_kv, + const c10::optional &seqused_q, + const c10::optional &seqused_kv, + int64_t batch_size, + int64_t max_seqlen_q, + int64_t max_seqlen_kv, + int64_t ori_topk, + int64_t cmp_topk, + int64_t cmp_ratio, + int64_t ori_mask_mode, + int64_t cmp_mask_mode, + int64_t ori_win_left, + int64_t ori_win_right, + c10::string_view layout_q, + c10::string_view layout_kv, + bool has_ori_kv, + bool has_cmp_kv, + const c10::string_view device) +{ + constexpr int64_t OUTPUT_SIZE = 1024; + at::Device output_device = at::Device(std::string(device)); + if (cu_seqlens_q.has_value()) { + output_device = cu_seqlens_q.value().device(); + } else if (cu_seqlens_ori_kv.has_value()) { + output_device = cu_seqlens_ori_kv.value().device(); + } else if (cu_seqlens_cmp_kv.has_value()) { + output_device = cu_seqlens_cmp_kv.value().device(); + } else if (seqused_q.has_value()) { + output_device = seqused_q.value().device(); + } else if (seqused_kv.has_value()) { + output_device = seqused_kv.value().device(); + } + at::Tensor output = torch::empty({OUTPUT_SIZE}, torch::dtype(torch::kInt32).device(output_device)); + + auto cu_seqlens_q_val = get_valid_tensor(cu_seqlens_q, output_device); + auto cu_seqlens_ori_kv_val = get_valid_tensor(cu_seqlens_ori_kv, output_device); + auto cu_seqlens_cmp_kv_val = get_valid_tensor(cu_seqlens_cmp_kv, output_device); + auto seqused_q_val = get_valid_tensor(seqused_q, output_device); + auto seqused_kv_val = get_valid_tensor(seqused_kv, output_device); + + std::string layout_q_str = std::string(layout_q); + std::string layout_kv_str = std::string(layout_kv); + char *layout_q_ptr = const_cast(layout_q_str.c_str()); + char *layout_kv_ptr = const_cast(layout_kv_str.c_str()); + + EXEC_NPU_CMD(aclnnSparseAttnSharedkvMetadata, cu_seqlens_q_val, cu_seqlens_ori_kv_val, cu_seqlens_cmp_kv_val, seqused_q_val, + seqused_kv_val, num_heads_q, num_heads_kv, head_dim, batch_size, max_seqlen_q, max_seqlen_kv, ori_topk, cmp_topk, + cmp_ratio, ori_mask_mode, cmp_mask_mode, ori_win_left, ori_win_right, layout_q_ptr, + layout_kv_ptr, has_ori_kv, has_cmp_kv, output); + return output; +} + +at::Tensor npu_vllm_quant_lightning_indexer_metadata_npu( + int64_t num_heads_q, int64_t num_heads_k, int64_t head_dim, int64_t query_quant_mode, int64_t key_quant_mode, + const c10::optional &actual_seq_lengths_query, const c10::optional &actual_seq_lengths_key, int64_t batch_size, + int64_t max_seqlen_q, int64_t max_seqlen_k, const c10::string_view layout_query, c10::string_view layout_key, int64_t sparse_count, + int64_t sparse_mode, int64_t pre_tokens, int64_t next_tokens, int64_t cmp_ratio, const c10::string_view device) +{ + constexpr int64_t OUTPUT_SIZE = 1024; + at::Device output_device = at::Device(std::string(device)); + if (actual_seq_lengths_query.has_value()) { + output_device = actual_seq_lengths_query.value().device(); + } else if (actual_seq_lengths_key.has_value()) { + output_device = actual_seq_lengths_key.value().device(); + } + + at::Tensor output = torch::empty({OUTPUT_SIZE}, torch::dtype(torch::kInt32).device(output_device)); + auto actual_seq_lengths_query_val = get_valid_tensor(actual_seq_lengths_query, output_device); + auto actual_seq_lengths_key_val = get_valid_tensor(actual_seq_lengths_key, output_device); + + std::string layout_query_str = std::string(layout_query); + char *layout_query_ptr = const_cast(layout_query_str.c_str()); + std::string layout_key_str = std::string(layout_key); + char *layout_key_ptr = const_cast(layout_key_str.c_str()); + + EXEC_NPU_CMD(aclnnVllmQuantLightningIndexerMetadata, actual_seq_lengths_query_val, actual_seq_lengths_key_val, + num_heads_q, num_heads_k, head_dim, query_quant_mode, key_quant_mode, batch_size, + max_seqlen_q, max_seqlen_k, layout_query_ptr, layout_key_ptr, sparse_count, + sparse_mode, pre_tokens, next_tokens, cmp_ratio, output); + + return output; +} + +at::Tensor construct_hc_post_output_tensor(const at::Tensor& residual) +{ + constexpr int64_t SIZE = 8; + constexpr int64_t DIM_0 = 0; + constexpr int64_t DIM_1 = 1; + constexpr int64_t DIM_2 = 2; + constexpr int64_t DIM_3 = 3; + at::SmallVector output_size = {residual.size(DIM_0), residual.size(DIM_1), residual.size(DIM_2), residual.size(DIM_3)}; + at::Tensor out = at::empty(output_size, residual.options().dtype(residual.dtype())); + return out; +} + +// step1，工具函数，检查输入shape +void check_hc_post_shape_and_dtype(const at::Tensor& x, const at::Tensor& residual, const at::Tensor& post, const at::Tensor& com) { + // check x shape: [b, s, d] + TORCH_CHECK(x.dim() == 3, "Input tensor x's dim num should be 3, actual ", x.dim(), "."); + for (size_t i = 0; i < 3; i++) { + TORCH_CHECK(x.size(i) > 0, "Input tensor x's shape should be positive, but x.shape[", i, "] is :", x.size(i), "."); + } + auto batch = x.size(0); + auto sequence = x.size(1); + auto d = x.size(2); + // check residual: [b, s, hc, d] + TORCH_CHECK(residual.dim() == 4, "Input tensor residual's dim num should be 4, actual ", residual.dim(), "."); + auto hc = residual.size(2); + TORCH_CHECK(hc > 0, "The hc of residual should be positive, actual ", hc, "."); + TORCH_CHECK(residual.size(0) == batch, "The residual.shape[0] should be batch, actual residual.shape[0] is ", residual.size(0), ", batch is ", batch, "."); + TORCH_CHECK(residual.size(1) == sequence, "The residual.shape[1] should be sequence, actual residual.shape[1] is ", residual.size(1), ", sequence is ", sequence, "."); + TORCH_CHECK(residual.size(3) == d, "The residual.shape[3] should be d, actual residual.shape[3] is ", residual.size(3), ", d is ", d, "."); + // check post [b, s, hc] + TORCH_CHECK(post.dim() == 3, "Input tensor post's dim num should be 3, actual ", post.dim(), "."); + TORCH_CHECK(post.size(0) == batch, "The post.shape[0] should be batch, actual post.shape[0] is ", post.size(0), ", batch is ", batch, "."); + TORCH_CHECK(post.size(1) == sequence, "The post.shape[1] should be sequence, actual post.shape[1] is ", post.size(1), ", sequence is ", sequence, "."); + TORCH_CHECK(post.size(2) == hc, "The post.shape[2] should be hc, actual post.shape[2] is ", post.size(2), ", hc is ", hc, "."); + // check com: [b, s, hc, hc] + TORCH_CHECK(com.dim() == 4, "Input tensor com's dim num should be 4, actual ", com.dim(), "."); + TORCH_CHECK(com.size(0) == batch, "The com.shape[0] should be batch, actual com.shape[0] is ", com.size(0), ", batch is ", batch, "."); + TORCH_CHECK(com.size(1) == sequence, "The com.shape[1] should be sequence, actual com.shape[1] is ", com.size(1), ", sequence is ", sequence, "."); + TORCH_CHECK(com.size(2) == hc, "The com.shape[2] should be hc, actual com.shape[2] is ", com.size(2), ", hc is ", hc, "."); + TORCH_CHECK(com.size(3) == hc, "The com.shape[3] should be hc, actual com.shape[3] is ", com.size(3), ", hc is ", hc, "."); + // check dtype + TORCH_CHECK(x.dtype() == at::kFloat || x.dtype() == at::kHalf || x.dtype() == at::kBFloat16, + "x should be FLOAT16, BFLOAT16, or FLOAT32."); + TORCH_CHECK(residual.dtype() == x.dtype(), "x's dtype should be equal to residual's dtype."); + TORCH_CHECK(post.dtype() == at::kFloat || post.dtype() == at::kHalf || post.dtype() == at::kBFloat16, + "post should be FLOAT16, BFLOAT16, or FLOAT32."); + TORCH_CHECK(com.dtype() == post.dtype(), "com's dtype should be equal to post's dtype."); +} + +at::Tensor npu_hc_post_npu( + const at::Tensor& x, + const at::Tensor& residual, + const at::Tensor& post, + const at::Tensor& comb) +{ + check_hc_post_shape_and_dtype(x, residual, post, comb); + // construct the output tensor + at::Tensor out = construct_hc_post_output_tensor(residual); + EXEC_NPU_CMD(aclnnHcPost, x, residual, post, comb, out); + return out; +} + +constexpr int64_t HC_PRE_HC_LIMIT = 4; +constexpr int64_t HC_PRE_D_LIMIT = 4096; +constexpr int64_t HC_PRE_D_LIMIT_EXTEND = 7168; +constexpr int64_t HC_PRE_MIX_HC_LIMIT = 24; + +std::tuple construct_hc_pre_output_tensor(const at::Tensor& x, int64_t hc_mult) +{ + auto xDims = x.dim(); + at::SmallVector y_size; + at::SmallVector post_size; + at::SmallVector comb_frag_size; + if (xDims == 4) { + auto batch = x.size(0); + auto size = x.size(1); + auto d = x.size(3); + y_size = {batch, size, d}; + post_size = {batch, size, hc_mult}; + comb_frag_size = {batch, size, hc_mult, hc_mult}; + } else if (xDims == 3){ + auto bs = x.size(0); + auto d = x.size(2); + y_size = {bs, d}; + post_size = {bs, hc_mult}; + comb_frag_size = {bs, hc_mult, hc_mult}; + } + + at::Tensor y = at::empty(y_size, x.options().dtype(at::kBFloat16)); + at::Tensor post = at::empty(post_size, x.options().dtype(at::kFloat)); + at::Tensor comb_frag = at::empty(comb_frag_size, x.options().dtype(at::kFloat)); + + return std::tuple(y, post, comb_frag); +} + +at::Tensor construct_hc_pre_rsqrt_output_tensor(const at::Tensor& x, float epsilon=1e-6) +{ + constexpr int64_t SIZE = 8; + TORCH_CHECK(epsilon >= 0, "epsilon should be greater than 0."); + + auto options = x.options(); + auto xDims = x.dim(); + c10::SmallVector yOut_shape; + for (size_t i = 0; i < xDims - 2; i++) { + yOut_shape.push_back(x.sizes()[i]); + } + yOut_shape.push_back(1); + at::Tensor yOut = at::empty(yOut_shape, options.dtype(at::kFloat)); + + return yOut; +} + +void check_hc_pre_shape_and_dtype( + const at::Tensor& x, + const at::Tensor& hc_fn, + const at::Tensor& hc_scale, + const at::Tensor& hc_base, + int64_t hc_mult) +{ + constexpr int64_t HC_SCALE_SIZE = 3; + auto x_dims = x.dim(); + TORCH_CHECK(x_dims == 3 || x_dims == 4, "Input tensor x's dim num should be 3 or 4, actual ", x_dims, "."); + for (auto i = 0; i < x_dims; i++) { + TORCH_CHECK(x.size(i) > 0, "Input tensor x's shape should be positive, but x.shape[", i, "] is ", + x.size(i), "."); + } + + auto hc = x_dims == 4 ? x.size(2) : x.size(1); + auto d = x_dims == 4 ? x.size(3) : x.size(2); + TORCH_CHECK(hc_mult == HC_PRE_HC_LIMIT, "hc_mult only supports ", HC_PRE_HC_LIMIT, ", actual ", hc_mult, "."); + TORCH_CHECK(hc == HC_PRE_HC_LIMIT, "The hc of x only supports ", HC_PRE_HC_LIMIT, ", actual ", hc, "."); + TORCH_CHECK(d == HC_PRE_D_LIMIT || d == HC_PRE_D_LIMIT_EXTEND, "The d of x only supports ", HC_PRE_D_LIMIT, + " or ", HC_PRE_D_LIMIT_EXTEND, ", actual ", d, "."); + TORCH_CHECK(hc_fn.dim() == 2, "Input tensor hc_fn's dim num should be 2, actual ", hc_fn.dim(), "."); + TORCH_CHECK(hc_fn.size(0) == HC_PRE_MIX_HC_LIMIT, "The hc_fn.shape[0] only supports ", + HC_PRE_MIX_HC_LIMIT, ", actual ", hc_fn.size(0), "."); + TORCH_CHECK(hc_fn.size(1) == hc * d, "The hc_fn.shape[1] should be hc * d, actual hc_fn.shape[1] is ", + hc_fn.size(1), ", hc is ", hc, ", d is ", d, "."); + TORCH_CHECK(hc_scale.dim() == 1, "Input tensor hc_scale's dim num should be 1, actual ", hc_scale.dim(), "."); + TORCH_CHECK(hc_scale.size(0) == HC_SCALE_SIZE, "Input tensor hc_scale's shape should be [", HC_SCALE_SIZE, + "], actual [", hc_scale.size(0), "]."); + TORCH_CHECK(hc_base.dim() == 1, "Input tensor hc_base's dim num should be 1, actual ", hc_base.dim(), "."); + TORCH_CHECK(hc_base.size(0) == HC_PRE_MIX_HC_LIMIT, "The hc_base.shape[0] only supports ", + HC_PRE_MIX_HC_LIMIT, ", actual ", hc_base.size(0), "."); + + TORCH_CHECK(x.dtype() == at::kBFloat16, "x's dtype should be BFLOAT16."); + TORCH_CHECK(hc_fn.dtype() == at::kFloat, "hc_fn's dtype should be FLOAT32."); + TORCH_CHECK(hc_scale.dtype() == at::kFloat, "hc_scale's dtype should be FLOAT32."); + TORCH_CHECK(hc_base.dtype() == at::kFloat, "hc_base's dtype should be FLOAT32."); +} + +std::tuple run_hc_pre_composite( + const at::Tensor& x, const at::Tensor& hc_fn, const at::Tensor& hc_scale, const at::Tensor& hc_base, + int64_t hc_mult, int64_t hc_sinkhorn_iters, double norm_eps, double hc_eps) +{ + auto xDims = x.dim(); + auto rsqrt = construct_hc_pre_rsqrt_output_tensor(x, norm_eps); + EXEC_NPU_CMD(aclnnHcPreInvRms, x, norm_eps, rsqrt); + + auto original_type = x.dtype(); + at::Tensor x_float = x.to(at::kFloat); + at::Tensor x_flattened = x_float.flatten(2, -1); + if (xDims == 3) { + x_flattened = x_float.flatten(1, -1); + } + auto mixes = at::linear(x_flattened, hc_fn); + + auto output_tensors = construct_hc_pre_output_tensor(x, hc_mult); + at::Tensor y = std::get<0>(output_tensors); + at::Tensor post = std::get<1>(output_tensors); + at::Tensor comb_frag = std::get<2>(output_tensors); + EXEC_NPU_CMD(aclnnHcPreSinkhorn, mixes, rsqrt, hc_scale, hc_base, x, hc_mult, hc_sinkhorn_iters, hc_eps, + y, post, comb_frag); + y = y.to(original_type); + + return std::tuple(y, post, comb_frag); +} + +std::tuple run_hc_pre_fusion( + const at::Tensor& x, const at::Tensor& hc_fn, const at::Tensor& hc_scale, const at::Tensor& hc_base, + int64_t hc_mult, int64_t hc_sinkhorn_iters, double norm_eps, double hc_eps) +{ + auto output_tensors = construct_hc_pre_output_tensor(x, hc_mult); + at::Tensor y = std::get<0>(output_tensors); + at::Tensor post = std::get<1>(output_tensors); + at::Tensor comb_frag = std::get<2>(output_tensors); + EXEC_NPU_CMD(aclnnHcPre, x, hc_fn, hc_scale, hc_base, hc_mult, hc_sinkhorn_iters, hc_eps, norm_eps, + y, post, comb_frag); + + return std::tuple(y, post, comb_frag); +} + +std::tuple npu_hc_pre_npu( + const at::Tensor& x, const at::Tensor& hc_fn, const at::Tensor& hc_scale, const at::Tensor& hc_base, + int64_t hc_mult, int64_t hc_sinkhorn_iters, double norm_eps, double hc_eps) +{ + check_hc_pre_shape_and_dtype(x, hc_fn, hc_scale, hc_base, hc_mult); + return run_hc_pre_composite(x, hc_fn, hc_scale, hc_base, hc_mult, hc_sinkhorn_iters, norm_eps, hc_eps); +} + +std::tuple npu_hc_pre_v2_npu( + const at::Tensor& x, const at::Tensor& hc_fn, const at::Tensor& hc_scale, const at::Tensor& hc_base, + int64_t hc_mult, int64_t hc_sinkhorn_iters, double norm_eps, double hc_eps) +{ + check_hc_pre_shape_and_dtype(x, hc_fn, hc_scale, hc_base, hc_mult); + return run_hc_pre_fusion(x, hc_fn, hc_scale, hc_base, hc_mult, hc_sinkhorn_iters, norm_eps, hc_eps); +} + +at::Tensor construct_hc_pre_inv_rms_output_tensor(const at::Tensor& x, float epsilon=1e-20) +{ + constexpr int64_t SIZE = 8; + TORCH_CHECK(epsilon >= 0, "epsilon should be greater than 0."); + + auto options = x.options(); + auto xDims = x.dim(); + c10::SmallVector yOut_shape; + for (auto i = 0; i < xDims - 2; i++) { + yOut_shape.push_back(x.sizes()[i]); + } + yOut_shape.push_back(1); + at::Tensor yOut = at::empty(yOut_shape, options.dtype(at::kFloat)); + + return yOut; +} + +at::Tensor npu_hc_pre_inv_rms_npu(const at::Tensor& x, double epsilon=1e-20) +{ + TORCH_CHECK(x.numel() > 0, "Input tensor x should not be empty."); + TORCH_CHECK(epsilon >= 0, "epsilon should be greater than 0."); + + TORCH_CHECK(x.dtype() == at::kFloat || x.dtype() == at::kHalf || x.dtype() == at::kBFloat16, + "x should be FLOAT16, BFLOAT16, or FLOAT32."); + + at::Tensor yOut; + yOut = construct_hc_pre_inv_rms_output_tensor(x, epsilon); + + EXEC_NPU_CMD(aclnnHcPreInvRms, x, epsilon, yOut); + + return yOut; +} + +std::tuple construct_hc_pre_sinkhorn_output_tensor(const at::Tensor& mixes, const at::Tensor& x, int64_t hc_mult) +{ + auto xDims = x.dim(); + at::SmallVector y_size; + at::SmallVector post_size; + at::SmallVector comb_frag_size; + if (xDims == 4) { + auto batch = x.size(0); + auto size = x.size(1); + auto d = x.size(3); + y_size = {batch, size, d}; + post_size = {batch, size, hc_mult}; + comb_frag_size = {batch, size, hc_mult, hc_mult}; + } else if (xDims == 3){ + auto bs = x.size(0); + auto d = x.size(2); + y_size = {bs, d}; + post_size = {bs, hc_mult}; + comb_frag_size = {bs, hc_mult, hc_mult}; + } + + at::Tensor y = at::empty(y_size, x.options().dtype(at::kBFloat16)); + at::Tensor post = at::empty(post_size, x.options().dtype(at::kFloat)); + at::Tensor comb_frag = at::empty(comb_frag_size, x.options().dtype(at::kFloat)); + + return std::tuple(y, post, comb_frag); +} + +std::tuple npu_hc_pre_sinkhorn_npu( + const at::Tensor& mixes, const at::Tensor& rsqrt, const at::Tensor& hc_scale, const at::Tensor& hc_base, + const at::Tensor& x, int64_t hc_mult, int64_t hc_sinkhorn_iters, double hc_eps) +{ + auto output_tensors = construct_hc_pre_sinkhorn_output_tensor(mixes, x, hc_mult); + at::Tensor y = std::get<0>(output_tensors); + at::Tensor post = std::get<1>(output_tensors); + at::Tensor comb_frag = std::get<2>(output_tensors); + + EXEC_NPU_CMD(aclnnHcPreSinkhorn, mixes, rsqrt, hc_scale, hc_base, x, hc_mult, hc_sinkhorn_iters, hc_eps, + y, post, comb_frag); + + return std::tuple(y, post, comb_frag); +} + +void inplace_partial_rotary_mul_npu(at::Tensor & x, const at::Tensor &r1, const at::Tensor &r2, c10::string_view rotary_mode, at::IntArrayRef partial_slice) +{ + constexpr int BSND_DIM_NUM = 4; + static const std::unordered_map mode_map = { + {"half", 0}, + {"interleave", 1}, + {"quarter", 2}, + {"interleave-half", 3} + }; + std::string rotary_mode_str = std::string(rotary_mode); + auto it = mode_map.find(rotary_mode_str); + if (it == mode_map.end()) + { + return; + } + auto origin_dim_num = x.dim(); + TORCH_CHECK(origin_dim_num == BSND_DIM_NUM, "Input tensor x's dim num should be 4, actual ", origin_dim_num, "."); + EXEC_NPU_CMD(aclnnInplacePartialRotaryMul, x, r1, r2, it->second, partial_slice); +} + +std::tuple npu_rms_norm_dynamic_quant_npu( + const at::Tensor& x, + const at::Tensor& gamma, + const c10::optional& smooth_scale, + const c10::optional& beta, + double epsilon) +{ + constexpr int32_t SIZE = 8; + TORCH_CHECK(x.numel() > 0, "Input tensor x should not be empty."); + TORCH_CHECK(gamma.numel() > 0, "Input tensor gamma should not be empty."); + TORCH_CHECK(gamma.dim() == 1 && gamma.size(0) == x.size(-1), "gamma dim are not equal to last dim of x shape."); + TORCH_CHECK(epsilon > 0, "epsilon should be greater than 0."); + TORCH_CHECK(x.dtype() == at::kHalf || x.dtype() == at::kBFloat16, "x should be FLOAT16, BFLOAT16."); + + at::Tensor smooth_scale2{nullptr}; + auto options = x.options(); + at::Tensor y_out = at::empty_like(x, options.dtype(at::kChar)); + at::Tensor y2_out = at::empty({1}, options.dtype(at::kChar)); + + c10::SmallVector scale_out_shape; + for (size_t i = 0; i < x.sizes().size() - 1; i++) { + scale_out_shape.push_back(x.sizes()[i]); + } + at::Tensor scale_out = at::empty(scale_out_shape, options.dtype(at::kFloat)); + at::Tensor scale2_out = at::empty_like(scale_out); + std::array* output_mask = nullptr; + int64_t* dst_type = nullptr; + + EXEC_NPU_CMD(aclnnRmsNormDynamicQuant, x, gamma, smooth_scale, smooth_scale2, beta, epsilon, output_mask, dst_type, + y_out, y2_out, scale_out, scale2_out); + + return std::make_tuple(y_out, scale_out); +} + +void indexer_compress_epilog_npu( + at::Tensor& indexer_compress_cache, + at::Tensor& indexer_compress_cache_scale, + const at::Tensor& x, + const at::Tensor& slot_mapping, + int64_t quant_mode = 1, + bool round_scale = true) +{ + EXEC_NPU_CMD(aclnnIndexerCompressEpilog, indexer_compress_cache, indexer_compress_cache_scale, x, + slot_mapping, quant_mode, round_scale); +} + +void validate_kv_compress_epilog_inputs( + const at::Tensor& x, + const at::Tensor& slot_mapping, + at::Tensor& kv_compress_cache) +{ + TORCH_CHECK(x.dim() == 2, "x must be 2D tensor, but got dimensions: ", x.dim()); + TORCH_CHECK(x.size(0) > 0 && x.size(1) > 0, + "x dimensions must be positive, but got: [", x.size(0), ", ", x.size(1), "]"); + TORCH_CHECK(slot_mapping.dim() == 1, + "slot_mapping must be 1D tensor, but got dimensions: ", slot_mapping.dim()); + TORCH_CHECK(slot_mapping.size(0) == x.size(0), + "slot_mapping size must equal x's first dimension, but got slot_mapping_size=", + slot_mapping.size(0), ", x.dim(0)=", x.size(0)); + if (kv_compress_cache.dim() == 4) { + TORCH_CHECK(kv_compress_cache.size(2) == 1, + "kv_compress_cache 4D tensor requires headnum (dim 2) == 1, but got ", + kv_compress_cache.size(2)); + } + TORCH_CHECK(x.dtype() == at::kBFloat16, "x must be BF16, but got ", x.dtype()); + TORCH_CHECK(slot_mapping.dtype() == at::kInt || slot_mapping.dtype() == at::kLong, + "slot_mapping must be INT32 or INT64, but got ", slot_mapping.dtype()); + TORCH_CHECK(kv_compress_cache.dtype() == at::ScalarType::Float8_e5m2 || + kv_compress_cache.dtype() == at::ScalarType::Float8_e4m3fn, + "kv_compress_cache must be FP8_E5M2 or FP8_E4M3, but got ", kv_compress_cache.dtype()); +} + +void kv_compress_epilog_npu( + at::Tensor& kv_compress_cache, + const at::Tensor& x, + const at::Tensor& slot_mapping, + int64_t quant_group_size, + int64_t quant_mode, + bool round_scale_flag, + int64_t layout) +{ + validate_kv_compress_epilog_inputs(x, slot_mapping, kv_compress_cache); + + at::Tensor cache = kv_compress_cache; + if (cache.dim() == 4) { + cache = cache.squeeze(2); + } + + int64_t round_scale = round_scale_flag ? 1 : 0; + int64_t cache_stride = cache.stride(0); + EXEC_NPU_CMD(aclnnKvCompressEpilog, cache, x, slot_mapping, quant_group_size, quant_mode, round_scale, + layout, cache_stride); +} + +std::tuple npu_kv_quant_sparse_attn_sharedkv_npu( + const at::Tensor& q, + int64_t kv_quant_mode, + const c10::optional& ori_kv, + const c10::optional& cmp_kv, + const c10::optional& ori_sparse_indices, + const c10::optional& cmp_sparse_indices, + const c10::optional& ori_block_table, + const c10::optional& cmp_block_table, + const c10::optional& cu_seqlens_q, + const c10::optional& cu_seqlens_ori_kv, + const c10::optional& cu_seqlens_cmp_kv, + const c10::optional& seqused_q, + const c10::optional& seqused_kv, + const c10::optional& sinks, + const c10::optional& metadata, + int64_t tile_size, + int64_t rope_head_dim, + double softmax_scale, + int64_t cmp_ratio, + int64_t ori_mask_mode, + int64_t cmp_mask_mode, + int64_t ori_win_left, + int64_t ori_win_right, + c10::string_view layout_q, + c10::string_view layout_kv, + bool return_softmax_lse) +{ + std::string layout_q_str = std::string(layout_q); + std::string layout_kv_str = std::string(layout_kv); + auto output = construct_output_tensor(q, layout_q_str, return_softmax_lse); + at::Tensor attn_out = std::get<0>(output); + at::Tensor softmax_lse = std::get<1>(output); + + char* layout_q_ptr = const_cast(layout_q_str.c_str()); + char* layout_kv_ptr = const_cast(layout_kv_str.c_str()); + int64_t ori_kv_stride0 = 0; + int64_t cmp_kv_stride0 = 0; + if (ori_kv.has_value() && ori_kv.value().defined()) { + ori_kv_stride0 = ori_kv.value().stride(0); + } + if (cmp_kv.has_value() && cmp_kv.value().defined()) { + cmp_kv_stride0 = cmp_kv.value().stride(0); + } + + EXEC_NPU_CMD(aclnnKvQuantSparseAttnSharedkv, q, ori_kv, cmp_kv, ori_sparse_indices, cmp_sparse_indices, + ori_block_table, cmp_block_table, cu_seqlens_q, cu_seqlens_ori_kv, cu_seqlens_cmp_kv, + seqused_q, seqused_kv, sinks, metadata, kv_quant_mode, tile_size, rope_head_dim, + softmax_scale, cmp_ratio, ori_mask_mode, cmp_mask_mode, ori_win_left, ori_win_right, + layout_q_ptr, layout_kv_ptr, ori_kv_stride0, cmp_kv_stride0, return_softmax_lse, + attn_out, softmax_lse); + return std::tuple(attn_out, softmax_lse); +} + +at::Tensor npu_kv_quant_sparse_attn_sharedkv_metadata_npu( + int64_t num_heads_q, + int64_t num_heads_kv, + int64_t head_dim, + int64_t kv_quant_mode, + const c10::optional& cu_seqlens_q, + const c10::optional& cu_seqlens_ori_kv, + const c10::optional& cu_seqlens_cmp_kv, + const c10::optional& seqused_q, + const c10::optional& seqused_kv, + int64_t batch_size, + int64_t max_seqlen_q, + int64_t max_seqlen_kv, + int64_t ori_topk, + int64_t cmp_topk, + int64_t tile_size, + int64_t rope_head_dim, + int64_t cmp_ratio, + int64_t ori_mask_mode, + int64_t cmp_mask_mode, + int64_t ori_win_left, + int64_t ori_win_right, + c10::string_view layout_q, + c10::string_view layout_kv, + bool has_ori_kv, + bool has_cmp_kv, + const c10::string_view device) +{ + constexpr int64_t OUTPUT_SIZE = 1024; + at::Device output_device = at::Device(std::string(device)); + if (cu_seqlens_q.has_value()) { + output_device = cu_seqlens_q.value().device(); + } else if (cu_seqlens_ori_kv.has_value()) { + output_device = cu_seqlens_ori_kv.value().device(); + } else if (cu_seqlens_cmp_kv.has_value()) { + output_device = cu_seqlens_cmp_kv.value().device(); + } else if (seqused_q.has_value()) { + output_device = seqused_q.value().device(); + } else if (seqused_kv.has_value()) { + output_device = seqused_kv.value().device(); + } + at::Tensor output = torch::empty({OUTPUT_SIZE}, torch::dtype(torch::kInt32).device(output_device)); + + auto cu_seqlens_q_val = get_valid_tensor(cu_seqlens_q, output_device); + auto cu_seqlens_ori_kv_val = get_valid_tensor(cu_seqlens_ori_kv, output_device); + auto cu_seqlens_cmp_kv_val = get_valid_tensor(cu_seqlens_cmp_kv, output_device); + auto seqused_q_val = get_valid_tensor(seqused_q, output_device); + auto seqused_kv_val = get_valid_tensor(seqused_kv, output_device); + + std::string layout_q_str = std::string(layout_q); + std::string layout_kv_str = std::string(layout_kv); + char* layout_q_ptr = const_cast(layout_q_str.c_str()); + char* layout_kv_ptr = const_cast(layout_kv_str.c_str()); + + EXEC_NPU_CMD(aclnnKvQuantSparseAttnSharedkvMetadata, cu_seqlens_q_val, cu_seqlens_ori_kv_val, + cu_seqlens_cmp_kv_val, seqused_q_val, seqused_kv_val, num_heads_q, num_heads_kv, + head_dim, batch_size, max_seqlen_q, max_seqlen_kv, ori_topk, cmp_topk, kv_quant_mode, + tile_size, rope_head_dim, cmp_ratio, ori_mask_mode, cmp_mask_mode, ori_win_left, + ori_win_right, layout_q_ptr, layout_kv_ptr, has_ori_kv, has_cmp_kv, output); + return output; +} + +int64_t get_type_code(at::ScalarType dst_type) +{ + switch (dst_type) { + case at::ScalarType::Float8_e5m2: + return 35; + case at::ScalarType::Float8_e4m3fn: + return 36; + case at::ScalarType::Half: + return 1; + case at::ScalarType::BFloat16: + return 27; + default: + TORCH_CHECK(false, "Unsupported dtype: ", dst_type); + } + return 0; +} + +std::tuple construct_swiglu_group_quant_output_tensor( + const at::Tensor& x, + int64_t dst_type, + int64_t quant_mode, + bool ue8m0_scale) +{ + constexpr int64_t SIZE = 8; + constexpr int64_t SWIGLU_FACTOR = 2; + constexpr int64_t PER_BLOCK_FP16 = 128; + constexpr int64_t PER_MX_FP16 = 32; + constexpr int64_t MX_SCALE_ALIGN_FACTOR = 2; + constexpr int64_t GROUP_QUANT = 1; + constexpr int64_t MX_QUANT = 2; + constexpr int64_t FP8_QUANT = 3; + + at::SmallVector y_size(x.sizes().begin(), x.sizes().end()); + for (size_t i = 0; i < x.sizes().size(); i++) { + TORCH_CHECK(x.size(i) >= 0, "All values within x's shape should be non-negative, but shape[", + i, "] is ", x.size(i)); + } + TORCH_CHECK(x.dtype() == at::kHalf || x.dtype() == at::kBFloat16, + "x should be FLOAT16 or BFLOAT16."); + int64_t x_last_dim = x.sizes().back(); + TORCH_CHECK(quant_mode == GROUP_QUANT || quant_mode == MX_QUANT || quant_mode == FP8_QUANT, + "Unsupported quant mode, only support ", GROUP_QUANT, " or ", MX_QUANT, " or ", FP8_QUANT, "."); + if (quant_mode == GROUP_QUANT || quant_mode == FP8_QUANT) { + TORCH_CHECK(x_last_dim % 256 == 0, + "In group quant, the last dim of x should be divisible by 256, actual ", x_last_dim, "."); + } else { + TORCH_CHECK(x_last_dim % 128 == 0, + "In mx quant, the last dim of x should be divisible by 128, actual ", x_last_dim, "."); + } + + y_size.back() = y_size.back() / SWIGLU_FACTOR; + int64_t y_last_dim = y_size.back(); + auto y_dtype = dst_type == 35 ? at::kFloat8_e5m2 : at::kFloat8_e4m3fn; + at::Tensor y = at::empty(y_size, x.options().dtype(y_dtype)); + + at::SmallVector scale_size(y_size.begin(), y_size.end()); + if (quant_mode == GROUP_QUANT || quant_mode == FP8_QUANT) { + scale_size.back() = (y_last_dim + PER_BLOCK_FP16 - 1) / PER_BLOCK_FP16; + } else if (quant_mode == MX_QUANT) { + int64_t scale_last_dim = (y_last_dim + PER_MX_FP16 - 1) / PER_MX_FP16; + scale_last_dim = (scale_last_dim + MX_SCALE_ALIGN_FACTOR - 1) / MX_SCALE_ALIGN_FACTOR; + scale_size.back() = scale_last_dim; + scale_size.push_back(MX_SCALE_ALIGN_FACTOR); + } + + auto scale_type = at::kFloat; + if (quant_mode == MX_QUANT || (quant_mode == FP8_QUANT && ue8m0_scale)) { + scale_type = at::kFloat8_e8m0fnu; + } + at::Tensor scale = at::empty(scale_size, x.options().dtype(scale_type)); + at::Tensor y_origin = at::empty(y_size, x.options().dtype(x.dtype())); + + return std::tuple(y, scale, y_origin); +} + +std::tuple npu_swiglu_group_quant_npu( + const at::Tensor& x, + const c10::optional& topk_weight, + const c10::optional& group_index, + at::ScalarType dst_type = at::ScalarType::Float8_e4m3fn, + int64_t quant_mode = 1, + int64_t group_size = 128, + bool round_scale = false, + bool ue8m0_scale = false, + bool output_origin = false, + int64_t group_list_type = 0, + double clamp_value = 0.0) +{ + int64_t dst_type_code = get_type_code(dst_type); + auto output_tensors = construct_swiglu_group_quant_output_tensor(x, dst_type_code, quant_mode, ue8m0_scale); + at::Tensor y = std::get<0>(output_tensors); + at::Tensor scale = std::get<1>(output_tensors); + at::Tensor y_origin = std::get<2>(output_tensors); + + EXEC_NPU_CMD(aclnnSwigluGroupQuant, x, topk_weight, group_index, dst_type_code, quant_mode, group_size, + round_scale, ue8m0_scale, output_origin, group_list_type, clamp_value, y, scale, y_origin); + + return std::tuple(y, scale, y_origin); +} + +std::tuple construct_load_index_kv_cache_output_tensor( + const at::Tensor& kv_cache, + const at::Tensor& slot_mapping) +{ + constexpr int64_t KV_LAST_DIM = 128; + int64_t n = slot_mapping.size(0); + + at::Tensor kv = at::empty({n, KV_LAST_DIM}, kv_cache.options().dtype(at::kFloat8_e4m3fn)); + at::Tensor kv_scale = at::empty({n}, kv_cache.options().dtype(at::kFloat)); + + return std::tuple(kv, kv_scale); +} + +std::tuple npu_load_index_kv_cache_npu( + const at::Tensor& kv_cache, + const at::Tensor& slot_mapping) +{ + auto output_tensors = construct_load_index_kv_cache_output_tensor(kv_cache, slot_mapping); + at::Tensor kv = std::get<0>(output_tensors); + at::Tensor kv_scale = std::get<1>(output_tensors); + + int64_t kv_cache_stride = kv_cache.stride(0); + EXEC_NPU_CMD(aclnnLoadIndexKvCache, kv_cache, slot_mapping, kv_cache_stride, kv, kv_scale); + + return std::tuple(kv, kv_scale); +} + +void indexer_compress_epilog_v2_npu( + at::Tensor& indexer_compress_cache, + const at::Tensor& x, + const at::Tensor& slot_mapping, + int64_t layout = 2) +{ + int64_t indexer_compress_cache_stride = indexer_compress_cache.stride(0); + EXEC_NPU_CMD(aclnnIndexerCompressEpilogV2, indexer_compress_cache, x, slot_mapping, layout, + indexer_compress_cache_stride); +} + +std::tuple npu_dequant_swiglu_quant( + const at::Tensor& x, + const c10::optional& weight_scale, + const c10::optional& activation_scale, + const c10::optional& bias, + const c10::optional& quant_scale, + const c10::optional& quant_offset, + const c10::optional& group_index, + bool activate_left, + int64_t quant_mode, + int64_t swiglu_mode, + double clamp_limit, + double glu_alpha, + double glu_bias) +{ + TORCH_CHECK(x.dim() > 1, "x dim should larger than 1"); + TORCH_CHECK(quant_mode == 0 || quant_mode == 1, "quant_mode only support 0 or 1, but got ", quant_mode); + TORCH_CHECK(swiglu_mode == 0 || swiglu_mode == 1, "swiglu_mode only support 0 or 1, but got ", swiglu_mode); + TORCH_CHECK(std::isfinite(clamp_limit) && clamp_limit >= 0.0, "clamp_limit should be positive finite"); + TORCH_CHECK(std::isfinite(glu_alpha), "glu_alpha should be finite"); + TORCH_CHECK(std::isfinite(glu_bias), "glu_bias should be finite"); + TORCH_CHECK(x.size(x.dim() - 1) % 2 == 0, "x last dim should be even"); + + c10::SmallVector y_size; + c10::SmallVector scale_size; + for (int64_t i = 0; i < x.dim() - 1; ++i) { + y_size.push_back(x.size(i)); + scale_size.push_back(x.size(i)); + } + y_size.push_back(x.size(x.dim() - 1) / 2); + + at::Tensor y = at::empty(y_size, x.options().dtype(c10::ScalarType::Char)); + at::Tensor scale = at::empty(scale_size, x.options().dtype(c10::ScalarType::Float)); + + std::string quant_mode_str = quant_mode == 1 ? "dynamic" : "static"; + char* quant_mode_ptr = const_cast(quant_mode_str.c_str()); + + const at::Tensor& weight_scale_value = c10::value_or_else(weight_scale, [] { return at::Tensor(); }); + const at::Tensor& activation_scale_opt = c10::value_or_else(activation_scale, [] { return at::Tensor(); }); + const at::Tensor& bias_opt = c10::value_or_else(bias, [] { return at::Tensor(); }); + const at::Tensor& quant_scale_opt = c10::value_or_else(quant_scale, [] { return at::Tensor(); }); + const at::Tensor& quant_offset_opt = c10::value_or_else(quant_offset, [] { return at::Tensor(); }); + const at::Tensor& group_index_opt = c10::value_or_else(group_index, [] { return at::Tensor(); }); + + static const bool is_v2_available = + GetOpApiFuncAddr("aclnnDequantSwigluQuantV2") != nullptr && + GetOpApiFuncAddr("aclnnDequantSwigluQuantV2GetWorkspaceSize") != nullptr; + + if (swiglu_mode == 0 && !is_v2_available) { + EXEC_NPU_CMD(aclnnDequantSwigluQuant, x, weight_scale_value, activation_scale_opt, bias_opt, quant_scale_opt, + quant_offset_opt, group_index_opt, activate_left, quant_mode_ptr, y, scale); + } else { + int64_t dst_type = 2; + char* round_mode = const_cast("rint"); + int64_t activate_dim = -1; + EXEC_NPU_CMD(aclnnDequantSwigluQuantV2, x, weight_scale_value, activation_scale_opt, bias_opt, quant_scale_opt, + quant_offset_opt, group_index_opt, activate_left, quant_mode_ptr, dst_type, round_mode, + activate_dim, swiglu_mode, clamp_limit, glu_alpha, glu_bias, y, scale); + } + + return std::make_tuple(y, scale); +} + +void npu_scatter_nd_update_v2( + at::Tensor& var, + const at::Tensor& indices, + const at::Tensor& update) +{ + // construct the output tensor + at::IntArrayRef var_stride = var.strides(); + EXEC_NPU_CMD(aclnnScatterNdUpdateV2, var, indices, update, var_stride); + return; +} + +std::tuple chunk_gated_delta_rule_fwd_h( + const at::Tensor & k, + const at::Tensor & w, + const at::Tensor & u, + const c10::optional & g, + const c10::optional & gk, + const c10::optional & initial_state, + c10::optional output_final_state, + c10::optional chunk_size, + c10::optional save_new_value, + c10::optional cu_seqlens, + c10::optional chunk_indices, + c10::optional use_exp2, + c10::optional transpose_state_layout) +{ + bool output_final_state_ = output_final_state.has_value() ? output_final_state.value() : false; + const at::Tensor &initial_state_ = c10::value_or_else(initial_state, [] { return at::Tensor(); }); + int64_t chunk_size_ = chunk_size.has_value() ? chunk_size.value() : 64; + const at::Tensor &g_ = c10::value_or_else(g, [] { return at::Tensor(); }); + const at::Tensor &gk_ = c10::value_or_else(gk, [] { return at::Tensor(); }); + + auto k_sizes = k.sizes(); + auto u_sizes = u.sizes(); + int K = k_sizes[3]; + int B = k_sizes[0]; + int T = k_sizes[2]; + int HV = u_sizes[1]; + int V = u_sizes[3]; + + int NT = 0; + if (chunk_indices.has_value()) { + auto chunk_indices_ref = chunk_indices.value(); + NT = chunk_indices_ref.size() / 2; + } else { + NT = (T + chunk_size_ - 1) / chunk_size_; + } + + at::Tensor h_out = at::zeros({B, HV, NT, K, V}, k.options()); + at::Tensor v_new_out = at::zeros(u.sizes(), u.options()); + at::Tensor final_state_out; + if (output_final_state_) { + int N = cu_seqlens.has_value() ? cu_seqlens->size() - 1 : B; + auto state_options = initial_state.has_value() ? initial_state->options() : h_out.options(); + final_state_out = at::empty({N, HV, K, V}, state_options); + } else { + final_state_out = at::empty({1}, k.options()); + } + + bool save_new_value_ = save_new_value.value_or(true); + bool use_exp2_ = use_exp2.value_or(false); + bool transpose_state_layout_ = transpose_state_layout.value_or(false); + + EXEC_NPU_CMD( + aclnnChunkGatedDeltaRuleFwdH, + k, w, u, g_, + gk_, initial_state_, output_final_state_, chunk_size_, save_new_value_, + cu_seqlens, chunk_indices, use_exp2_, transpose_state_layout_, + h_out, v_new_out, final_state_out + ); + + if (output_final_state_) { + return std::make_tuple(h_out, v_new_out, final_state_out); + } else { + return std::make_tuple(h_out, v_new_out, at::Tensor()); + } +} + +at::Tensor chunk_fwd_o( + const at::Tensor & q, + const at::Tensor & k, + const at::Tensor & v, + const at::Tensor & h, + double scale, + const c10::optional & g, + const c10::optional & g_gamma, + c10::optional cu_seqlens, + c10::optional chunk_indices, + c10::optional chunk_size, + c10::optional transpose_state_layout) +{ + at::Tensor o = at::zeros(v.sizes(), v.options()); + int64_t chunk_size_ = chunk_size.has_value() ? chunk_size.value() : 64; + const at::Tensor &g_ = c10::value_or_else(g, [] { return at::Tensor(); }); + (void)g_gamma; + (void)transpose_state_layout; + + EXEC_NPU_CMD( + aclnnChunkFwdO, + q, k, v, h, g_, + cu_seqlens, chunk_indices, scale, chunk_size_, + o + ); + return o; +} + +std::vector get_npu_storage_shape(const at::Tensor& tensor) +{ + TORCH_CHECK( + tensor.is_privateuseone(), + "get_npu_storage_shape only supports NPU tensors, but got device ", + tensor.device()); + const auto& desc = NPUBridge::GetNpuStorageImplDesc(tensor); + return std::vector(desc.storage_sizes_.begin(), desc.storage_sizes_.end()); +} + + +} // namespace vllm_fl + +#ifdef ASCEND_PLATFORM_310P +// Pybind on Ascend 310P +TORCH_LIBRARY_EXPAND(CONCAT(_C, _ascend), ops) +{ + ops.def( + "npu_causal_conv1d_310(Tensor x, " + " Tensor weight, " + " Tensor? bias, " + " Tensor conv_states, " + " Tensor? query_start_loc, " + " Tensor? cache_indices, " + " Tensor? initial_state_mode, " + " Tensor? num_accepted_tokens, " + " int activation_mode, " + " int pad_slot_id, " + " int run_mode) -> (Tensor output)"); + ops.impl("npu_causal_conv1d_310", torch::kPrivateUse1, &vllm_fl::npu_causal_conv1d_310); + + ops.def( + "npu_recurrent_gated_delta_rule_310(Tensor query, " + " Tensor key, " + " Tensor value, " + " Tensor beta, " + " Tensor state, " + " Tensor actual_seq_lengths, " + " Tensor ssm_state_indices, " + " Tensor? g, " + " Tensor? gk, " + " Tensor? num_accepted_tokens, " + " float scale_value=1.0) -> (Tensor output)"); + ops.impl("npu_recurrent_gated_delta_rule_310", torch::kPrivateUse1, &vllm_fl::npu_recurrent_gated_delta_rule_310); + + ops.def( + "chunk_gated_delta_rule_fwd_h(Tensor k, Tensor w, Tensor u, Tensor? g=None, *, Tensor? gk=None, Tensor? initial_state=None, bool? output_final_state=False, int? chunk_size=None, bool? save_new_value=True, int[]? cu_seqlens=None, int[]? chunk_indices=None, bool? use_exp2=False, bool? transpose_state_layout=False) -> (Tensor h_out, Tensor v_new_out, Tensor final_state_out)" + ); + ops.impl("chunk_gated_delta_rule_fwd_h", torch::kPrivateUse1, &vllm_fl::chunk_gated_delta_rule_fwd_h); + + ops.def( + "chunk_fwd_o(Tensor q, Tensor k, Tensor v, Tensor h, float scale, *, Tensor? g=None, Tensor? g_gamma=None, int[]? cu_seqlens=None, int[]? chunk_indices=None, int? chunk_size=None, bool? transpose_state_layout=False) -> Tensor" + ); + ops.impl("chunk_fwd_o", torch::kPrivateUse1, &vllm_fl::chunk_fwd_o); +} +#else +// Pybind on other platform +TORCH_LIBRARY_EXPAND(CONCAT(_C, _ascend), ops) +{ + + // vLLM-Ascend custom ops + // Gemma RmsNorm + ops.def( + "npu_gemma_rms_norm(Tensor x, " + "Tensor gamma, " + "float epsilon=1e-6)" + "-> (Tensor y ,Tensor rstd)" + ); + ops.impl("npu_gemma_rms_norm", torch::kPrivateUse1, &vllm_fl::npu_gemma_rms_norm); + + ops.def( + "npu_recurrent_gated_delta_rule(Tensor query, " + " Tensor key, " + " Tensor value, " + " Tensor(a!) state, " + " *, " + " Tensor? beta=None, " + " float? scale=None, " + " Tensor? actual_seq_lengths=None, " + " Tensor? ssm_state_indices=None, " + " Tensor? num_accepted_tokens=None, " + " Tensor? g=None, " + " Tensor? gk=None) -> Tensor"); + ops.impl("npu_recurrent_gated_delta_rule", torch::kPrivateUse1, &vllm_fl::npu_recurrent_gated_delta_rule); + +#ifdef VLLM_ENABLE_ATB_AND_DIRECT_KERNELS + // Direct kernel custom ops + ops.def("bgmv_shrink(Tensor! x, Tensor! weight, Tensor! indices, Tensor! y, float scale) -> ()"); + ops.impl("bgmv_shrink", torch::kPrivateUse1, &vllm_fl::bgmv_shrink); + + ops.def( + "bgmv_expand(Tensor! x, Tensor! weight, Tensor! indices, Tensor! y," + " int slice_offset, int slice_size) -> Tensor"); + ops.impl("bgmv_expand", torch::kPrivateUse1, &vllm_fl::bgmv_expand); + + ops.def("sgmv_shrink(Tensor! x, Tensor! weight, Tensor! lora_indices, Tensor! seq_len, Tensor! y, float scale) -> ()"); + ops.impl("sgmv_shrink", torch::kPrivateUse1, &vllm_fl::sgmv_shrink); + + ops.def( + "sgmv_expand(Tensor! x, Tensor! weight, Tensor! lora_indices, Tensor! seq_len, Tensor! y," + " int slice_offset, int slice_size) -> Tensor"); + ops.impl("sgmv_expand", torch::kPrivateUse1, &vllm_fl::sgmv_expand); + + ops.def( + "mla_preprocess(Tensor hiddenState, Tensor wdqkv," + " Tensor? descale0, Tensor gamma1, Tensor? beta1, Tensor wuq, Tensor? descale1," + " Tensor gamma2, Tensor cos, Tensor sin, Tensor wuk, Tensor kv_cache," + " Tensor kv_cache_rope, Tensor slotmapping, Tensor? quant_scale0," + " Tensor? quant_offset0, Tensor? bias0, Tensor? quant_scale1, Tensor? quant_offset1," + " Tensor? bias1, Tensor? ctkv_scale, Tensor? q_nope_scale, str? cache_mode," + " str? quant_mode, bool? enable_inner_out, Tensor! q_out0, Tensor! kv_cache_out0, Tensor! q_out1," + " Tensor! kv_cache_out1, Tensor! inner_out) -> (Tensor q_out0, Tensor kv_cache_out0," + " Tensor q_out1, Tensor kv_cache_out1, Tensor inner_out)" + ); + ops.impl("mla_preprocess", torch::kPrivateUse1, &vllm_fl::mla_preprocess); + + //batch_matmul ops refer to sgl-kernel-npu + ops.def( + "batch_matmul_transpose(Tensor tensor_a, Tensor tensor_b, Tensor tensor_c, str? format_mode=None, str? quant_mode=None) -> ()"); + ops.impl("batch_matmul_transpose", torch::kPrivateUse1, &vllm_fl::batch_matmul_transpose); + + ops.def("swap_blocks(Tensor! x, Tensor! y, Tensor z) -> ()"); + ops.impl("swap_blocks", torch::kPrivateUse1, &vllm_fl::swap_blocks); +#endif + + // swap_blocks_batch takes CPU tensors (int64 pointer/size arrays), not NPU + // tensors, so dispatch must be registered on the CPU backend. The function + // internally submits async memcpy on the current NPU stream. + ops.def("swap_blocks_batch(Tensor x, Tensor y, Tensor z, int direction) -> ()"); + ops.impl("swap_blocks_batch", torch::kCPU, &vllm_fl::swap_blocks_batch); + ops.def("device_print(str msg) -> ()"); + ops.impl("device_print", c10::DispatchKey::CompositeExplicitAutograd, + static_cast(&vllm_fl::device_print)); + + ops.def("device_print_tensor(Tensor tensor) -> ()"); + ops.impl("device_print_tensor", c10::DispatchKey::CompositeExplicitAutograd, + static_cast(&vllm_fl::device_print)); + + ops.def("get_npu_storage_shape(Tensor tensor) -> int[]"); + ops.impl("get_npu_storage_shape", c10::DispatchKey::CompositeExplicitAutograd, + &vllm_fl::get_npu_storage_shape); + + ops.def( + "grouped_matmul_swiglu_quant(Tensor x, Tensor weight, Tensor weight_scale, Tensor x_scale," + " Tensor group_list, *, Tensor? bias=None," + " Tensor? offset=None, float swiglu_limit=0.0) ->" + " (Tensor output, Tensor output_scale, Tensor output_offset)"); + ops.impl("grouped_matmul_swiglu_quant", torch::kPrivateUse1, &vllm_fl::grouped_matmul_swiglu_quant); + + ops.def( + "grouped_matmul_swiglu_quant_weight_nz(Tensor x, Tensor weight, Tensor weight_scale, Tensor x_scale," + " Tensor group_list, *, Tensor? bias=None," + " Tensor? offset=None, float swiglu_limit=0.0) -> " + " (Tensor output, Tensor output_scale, Tensor output_offset)"); + ops.impl("grouped_matmul_swiglu_quant_weight_nz", torch::kPrivateUse1, &vllm_fl::grouped_matmul_swiglu_quant_weight_nz); + + ops.def( + "grouped_matmul_swiglu_quant_weight_nz_tensor_list(Tensor x, Tensor[] weight, Tensor[] weight_scale, Tensor x_scale," + " Tensor group_list, *," + " Tensor? bias=None, Tensor? offset=None, float swiglu_limit=0.0) ->" + " (Tensor output, Tensor output_scale, Tensor output_offset)" + ); + ops.impl("grouped_matmul_swiglu_quant_weight_nz_tensor_list", torch::kPrivateUse1, &vllm_fl::grouped_matmul_swiglu_quant_weight_nz_tensor_list); + + ops.def( + "grouped_matmul_swiglu_quant_v2(Tensor x, Tensor[] weight, Tensor[] weight_scale, Tensor x_scale, Tensor group_list, Tensor? smooth_scale=None," + " Tensor[]? weight_assist_matrix=None, Tensor? bias=None, int? dequant_mode=0, int? dequant_dtype=0, int? quant_mode=0," + " int? quant_dtype=0, bool transpose_weight=False, int group_list_type=0, int[2] tuning_config=[],float swiglu_limit=0.0) ->" + " (Tensor output, Tensor output_scale)" + ); + ops.impl("grouped_matmul_swiglu_quant_v2", torch::kPrivateUse1, &vllm_fl::grouped_matmul_swiglu_quant_v2); + + ops.def( + "npu_lightning_indexer(" + "Tensor query, Tensor key, Tensor weights, " + "*, " + "Tensor? actual_seq_lengths_query=None, " + "Tensor? actual_seq_lengths_key=None, " + "Tensor? block_table=None, " + "str layout_query=\"BSND\", str layout_key=\"BSND\", " + "int sparse_count=2048, int sparse_mode=3, " + "int pre_tokens=9223372036854775807, " + "int next_tokens=9223372036854775807, " + "bool return_value=False" + ") -> (Tensor sparse_indices, Tensor sparse_values)" + ); + ops.impl("npu_lightning_indexer", torch::kPrivateUse1, &vllm_fl::npu_lightning_indexer); + + ops.def( + "npu_sparse_flash_attention(Tensor query, Tensor key, Tensor value," + " Tensor sparse_indices, float scale_value, *," + " Tensor? block_table=None, Tensor? actual_seq_lengths_query=None," + " Tensor? actual_seq_lengths_kv=None, Tensor? query_rope=None," + " Tensor? key_rope=None, int sparse_block_size=1," + " str layout_query='BSND', str layout_kv='BSND'," + " int sparse_mode=3, int pre_tokens=9223372036854775807," + " int next_tokens=9223372036854775807, int attention_mode=2," + " bool return_softmax_lse=False) -> (Tensor attention_out, Tensor softmax_max, Tensor softmax_sum)" + ); + ops.impl("npu_sparse_flash_attention", torch::kPrivateUse1, &vllm_fl::npu_sparse_flash_attention); + + ops.def( + "npu_kv_quant_sparse_flash_attention(Tensor query, Tensor key, Tensor value," + " Tensor sparse_indices, float scale_value, *," + " int key_quant_mode=1, int value_quant_mode=1," + " Tensor? key_dequant_scale=None," + " Tensor? value_dequant_scale=None," + " Tensor? block_table=None," + " Tensor? actual_seq_lengths_query=None," + " Tensor? actual_seq_lengths_kv=None," + " int sparse_block_size=1," + " str layout_query='BSND', str layout_kv='BSND'," + " int sparse_mode=3," + " int pre_tokens=9223372036854775807," + " int next_tokens=9223372036854775807," + " int attention_mode=2," + " int quant_scale_repo_mode=1," + " int tile_size=128," + " int rope_head_dim=64," + " bool return_softmax_lse=False)" + " -> (Tensor attention_out, Tensor softmax_max, Tensor softmax_sum)" + ); + ops.impl("npu_kv_quant_sparse_flash_attention", torch::kPrivateUse1, + &vllm_fl::npu_kv_quant_sparse_flash_attention); + + ops.def( + "dispatch_ffn_combine(Tensor x, Tensor[] weight1, Tensor[] weight2, Tensor expert_idx," + " Tensor[] scale1, Tensor[] scale2, Tensor[] bias1, Tensor[] bias2, Tensor probs, str group," + " int max_output_size, Tensor! out, Tensor! expert_token_nums, Tensor? x_active_mask=None, float swiglu_limit=1000000.0) -> (Tensor out, Tensor expert_token_nums)" + ); + ops.impl("dispatch_ffn_combine", torch::kPrivateUse1, &vllm_fl::dispatch_ffn_combine); + + ops.def("matmul_allreduce_add_rmsnorm(Tensor x1, Tensor x2, Tensor residual, Tensor gamma, \ + str groupTp, int tpRankSize, int tpRankId, float epsilon, bool isTransB, bool isGatherAddOut) -> (Tensor output, Tensor add_out)"); + ops.impl("matmul_allreduce_add_rmsnorm", torch::kPrivateUse1, &vllm_fl::matmul_allreduce_add_rmsnorm); + + ops.def( + "npu_moe_init_routing_custom(Tensor x, Tensor expert_idx, *, Tensor? scale=None, Tensor? offset=None, int active_num=-1, " + " int expert_capacity=-1, int expert_num=-1, int drop_pad_mode=0, int expert_tokens_num_type=0, " + " bool expert_tokens_num_flag=False, int quant_mode=0, int[2] active_expert_range=[], " + " int row_idx_type=0) -> (Tensor, Tensor, Tensor, Tensor)" + ); + ops.impl("npu_moe_init_routing_custom", torch::kPrivateUse1, &vllm_fl::npu_moe_init_routing_custom); + // vLLM-Ascend custom ops + ops.def( + "moe_gating_top_k(Tensor x, " + "int k, " + "int k_group, " + "int group_count, " + "int group_select_mode, " + "int renorm, " + "int norm_type, " + "bool out_flag, " + "float routed_scaling_factor, " + "float eps," + "Tensor? bias_opt=None)" + + "-> (Tensor y ,Tensor expert_idx, Tensor out)" + ); + ops.impl("moe_gating_top_k", torch::kPrivateUse1,&vllm_fl::moe_gating_top_k); + + ops.def( + "npu_add_rms_norm_bias(Tensor x1, " + "Tensor x2, " + "Tensor gamma, " + "Tensor? beta=None, " + "float epsilon=1e-6)" + "-> (Tensor y ,Tensor rstd, Tensor x)" + ); + ops.impl("npu_add_rms_norm_bias", torch::kPrivateUse1, &vllm_fl::npu_add_rms_norm_bias); + + ops.def("npu_apply_top_k_top_p(Tensor logits, Tensor? p=None, Tensor? k=None) -> Tensor"); + ops.impl("npu_apply_top_k_top_p", torch::kPrivateUse1, &vllm_fl::npu_apply_top_k_top_p); + + ops.def( + "npu_hamming_dist_top_k(Tensor q, Tensor k_comp, Tensor k_comp_rope, Tensor k," + " Tensor seq_len, Tensor? chunk_size=None," + " int? max_seq_len=None, int? sink=None, int? recent=None, int? support_offload=None," + " Tensor? key_block_table=None, Tensor? mask=None, Tensor? indices=None) -> Tensor" + ); + ops.impl("npu_hamming_dist_top_k", torch::kPrivateUse1, &vllm_fl::npu_hamming_dist_top_k); + + ops.def( + "npu_reshape_and_cache_bnsd(Tensor q, Tensor k_comp, Tensor slot_mapping, Tensor seq_len, Tensor k_out) -> Tensor" + ); + ops.impl("npu_reshape_and_cache_bnsd", torch::kPrivateUse1, &vllm_fl::npu_reshape_and_cache_bnsd); + + + ops.def( + "transpose_kv_cache_by_block(Tensor[] kCache, Tensor[] vCache, Tensor blockIDs, int blockSize, int headNum, int headDim, int splitNum, int layerNum) -> ()" + ); + ops.impl("transpose_kv_cache_by_block", torch::kPrivateUse1, &vllm_fl::transpose_kv_cache_by_block); + + ops.def( + "npu_copy_and_expand_eagle_inputs(Tensor target_token_ids, Tensor target_positions, " + "Tensor next_token_ids, Tensor query_start_loc, Tensor query_end_loc, " + "int padding_token_id, int parallel_drafting_token_id, int num_padding_slots_per_request, " + "bool shift_input_ids, int total_draft_tokens) -> " + "(Tensor out_input_ids, Tensor out_positions, Tensor out_is_rejected_token_mask, " + "Tensor out_is_masked_token_mask, Tensor out_new_token_indices, Tensor out_hidden_state_mapping)" + ); + ops.impl("npu_copy_and_expand_eagle_inputs", torch::kPrivateUse1, &vllm_fl::npu_copy_and_expand_eagle_inputs); + ops.def( + "npu_causal_conv1d_custom(Tensor output, Tensor x, " + " Tensor weight, " + " Tensor conv_state, " + " Tensor? bias_opt, " + " Tensor? query_start_loc_opt, " + " Tensor? cache_indices_opt, " + " Tensor? initial_state_mode_opt, " + " Tensor? num_accepted_tokens_opt, " + " int activation_mode, " + " int pad_slot_id, " + " int run_mode" + ") -> (Tensor output)"); + ops.impl("npu_causal_conv1d_custom", torch::kPrivateUse1, &vllm_fl::npu_causal_conv1d_custom); + ops.def( + "moe_grouped_matmul(" + "Tensor x," + "Tensor weight," + "Tensor group_list," + "int split_item," + "int group_type," + "int group_list_type)" + + "-> Tensor[]" + ); + ops.impl("moe_grouped_matmul", torch::kPrivateUse1,&vllm_fl::moe_grouped_matmul); + + ops.def( + "moe_gating_top_k_hash(" + "Tensor x, " + "int k, " + "Tensor? bias=None, " + "Tensor? input_ids=None, " + "Tensor? tid2eid=None, " + "int k_group=1, " + "int group_count=1, " + "float routed_scaling_factor=1.0, " + "float eps=1e-20, " + "int group_select_mode=0, " + "int renorm=0, " + "int norm_type=0, " + "bool out_flag=False" + ") -> (Tensor y, Tensor expert_idx, Tensor out)" + ); + ops.impl("moe_gating_top_k_hash", torch::kPrivateUse1,&vllm_fl::moe_gating_top_k_hash); + + ops.def( + "compressor(" + "Tensor x, Tensor wkv, Tensor wgate, " + "Tensor(a!) state_cache, Tensor ape, Tensor norm_weight, " + "Tensor rope_sin, Tensor rope_cos, " + "Tensor? state_block_table, Tensor? cu_seqlens, " + "Tensor? seqused, Tensor? start_pos, " + "int rope_head_dim, int cmp_ratio, int coff, " + "float norm_eps, int rotary_mode, int cache_mode" + ") -> Tensor" + ); + ops.impl("compressor", torch::kPrivateUse1, &vllm_fl::compressor); + + ops.def( + "compressor_metadata(" + "Tensor rope_cos, Tensor rope_sin, " + "Tensor cu_seqlens, Tensor start_pos, Tensor kv_block_table, " + "int kv_block_size, int slot_mapping_format, int compress_ratio, int num_compressed_tokens, " + "int num_reqs_actual" + ") -> (Tensor, Tensor, Tensor)" + ); + ops.impl("compressor_metadata", torch::kPrivateUse1, &vllm_fl::compressor_metadata); + + ops.def( + "compressor_metadata_out(" + "Tensor rope_cos, Tensor rope_sin, " + "Tensor cu_seqlens, Tensor start_pos, Tensor kv_block_table, " + "int kv_block_size, int slot_mapping_format, int compress_ratio, int num_reqs_actual, " + "Tensor(a!) compress_cos, Tensor(b!) compress_sin, Tensor(c!) slot_mapping" + ") -> (Tensor(a!), Tensor(b!), Tensor(c!))" + ); + ops.impl("compressor_metadata_out", torch::kPrivateUse1, &vllm_fl::compressor_metadata_out); + + ops.def( + "npu_vllm_quant_lightning_indexer(" + "Tensor query, Tensor key, Tensor weights, " + "Tensor query_dequant_scale, Tensor key_dequant_scale, " + "int query_quant_mode=0, int key_quant_mode=0, " + "Tensor? actual_seq_lengths_query=None, " + "Tensor? actual_seq_lengths_key=None, " + "Tensor? block_table=None, " + "Tensor? metadata=None, " + "str layout_query=\"BSND\", str layout_key=\"BSND\", " + "int sparse_count=2048, int sparse_mode=3, " + "int pre_tokens=9223372036854775807, " + "int next_tokens=9223372036854775807, " + "int cmp_ratio=1, bool return_value=False" + ") -> (Tensor sparse_indices, Tensor sparse_values)" + ); + ops.impl("npu_vllm_quant_lightning_indexer", torch::kPrivateUse1, &vllm_fl::npu_vllm_quant_lightning_indexer_npu); + + ops.def( + "npu_sparse_attn_sharedkv(" + "Tensor q, *, " + "Tensor? ori_kv=None, " + "Tensor? cmp_kv=None, " + "Tensor? ori_sparse_indices=None, " + "Tensor? cmp_sparse_indices=None, " + "Tensor? ori_block_table=None, " + "Tensor? cmp_block_table=None, " + "Tensor? cu_seqlens_q=None, " + "Tensor? cu_seqlens_ori_kv=None, " + "Tensor? cu_seqlens_cmp_kv=None, " + "Tensor? seqused_q=None, " + "Tensor? seqused_kv=None, " + "Tensor? sinks=None, " + "Tensor? metadata=None, " + "float softmax_scale=0, " + "int cmp_ratio=0, " + "int ori_mask_mode=4, " + "int cmp_mask_mode=3, " + "int ori_win_left=128, " + "int ori_win_right=0, " + "str layout_q=\"BSND\", " + "str layout_kv=\"PA_ND\", " + "bool return_softmax_lse=False" + ") -> (Tensor out, Tensor softmax_lse)" + ); + ops.impl("npu_sparse_attn_sharedkv", torch::kPrivateUse1, &vllm_fl::npu_sparse_attn_sharedkv_npu); + + ops.def( + "npu_sparse_attn_sharedkv_metadata(" + "int num_heads_q, " + "int num_heads_kv, " + "int head_dim, " + "Tensor? cu_seqlens_q=None, " + "Tensor? cu_seqlens_ori_kv=None, " + "Tensor? cu_seqlens_cmp_kv=None, " + "Tensor? seqused_q=None, " + "Tensor? seqused_kv=None, " + "int batch_size=0, " + "int max_seqlen_q=0, " + "int max_seqlen_kv=0, " + "int ori_topk=0, " + "int cmp_topk=0, " + "int cmp_ratio=4, " + "int ori_mask_mode=4, " + "int cmp_mask_mode=3, " + "int ori_win_left=128, " + "int ori_win_right=0, " + "str layout_q=\"BSND\", " + "str layout_kv=\"PA_ND\", " + "bool has_ori_kv=True, " + "bool has_cmp_kv=True, " + "str device=\"npu\"" + ") -> (Tensor metadata)" + ); + ops.impl("npu_sparse_attn_sharedkv_metadata", torch::kPrivateUse1, &vllm_fl::npu_sparse_attn_sharedkv_metadata_npu); + + ops.def( + "npu_vllm_quant_lightning_indexer_metadata(" + "int num_heads_q, " + "int num_heads_k, " + "int head_dim, " + "int query_quant_mode, " + "int key_quant_mode, " + "Tensor? actual_seq_lengths_query=None, " + "Tensor? actual_seq_lengths_key=None, " + "int batch_size=0, " + "int max_seqlen_q=0, " + "int max_seqlen_k=0, " + "str layout_query=\"BSND\", " + "str layout_key=\"BSND\", " + "int sparse_count=2048, " + "int sparse_mode=3, " + "int pre_tokens=9223372036854775807, " + "int next_tokens=9223372036854775807, " + "int cmp_ratio=1, " + "str device=\"npu\"" + ") -> (Tensor metadata)" + ); + ops.impl("npu_vllm_quant_lightning_indexer_metadata", torch::kPrivateUse1, &vllm_fl::npu_vllm_quant_lightning_indexer_metadata_npu); + + ops.def( + "npu_hc_post(" + "Tensor x, " + "Tensor residual, " + "Tensor post, " + "Tensor comb" + ") -> (Tensor out)" + ); + ops.impl("npu_hc_post", torch::kPrivateUse1, &vllm_fl::npu_hc_post_npu); + + ops.def( + "npu_hc_pre(" + "Tensor x, Tensor hc_fn, Tensor hc_scale, Tensor hc_base, " + "int hc_mult, int hc_sinkhorn_iters, " + "float norm_eps, float hc_eps" + ") -> (Tensor out0, Tensor out1, Tensor out2)" + ); + ops.impl("npu_hc_pre", torch::kPrivateUse1, &vllm_fl::npu_hc_pre_npu); + + ops.def( + "npu_hc_pre_v2(" + "Tensor x, Tensor hc_fn, Tensor hc_scale, Tensor hc_base, " + "int hc_mult, int hc_sinkhorn_iters, " + "float norm_eps, float hc_eps" + ") -> (Tensor out0, Tensor out1, Tensor out2)" + ); + ops.impl("npu_hc_pre_v2", torch::kPrivateUse1, &vllm_fl::npu_hc_pre_v2_npu); + + ops.def( + "npu_hc_pre_inv_rms(" + "Tensor x, float epsilon=1e-20" + ") -> (Tensor out)" + ); + ops.impl("npu_hc_pre_inv_rms", torch::kPrivateUse1, &vllm_fl::npu_hc_pre_inv_rms_npu); + + ops.def( + "npu_hc_pre_sinkhorn(" + "Tensor mixes, Tensor rsqrt, Tensor hc_scale, Tensor hc_base, Tensor x, " + "int hc_mult, int hc_sinkhorn_iters, float hc_eps" + ") -> (Tensor out0, Tensor out1, Tensor out2)" + ); + ops.impl("npu_hc_pre_sinkhorn", torch::kPrivateUse1, &vllm_fl::npu_hc_pre_sinkhorn_npu); + + ops.def( + "inplace_partial_rotary_mul(" + "Tensor(a!) x, Tensor r1, Tensor r2, str rotary_mode, int[] partial_slice" + ") -> ()" + ); + ops.impl("inplace_partial_rotary_mul", torch::kPrivateUse1, &vllm_fl::inplace_partial_rotary_mul_npu); + + + + + ops.def( + "npu_dequant_swiglu_quant(" + "Tensor x, *, " + "Tensor? weight_scale=None, " + "Tensor? activation_scale=None, " + "Tensor? bias=None, " + "Tensor? quant_scale=None, " + "Tensor? quant_offset=None, " + "Tensor? group_index=None, " + "bool activate_left=True, " + "int quant_mode=0, " + "int swiglu_mode=0, " + "float clamp_limit=0.0, " + "float glu_alpha=1.0, " + "float glu_bias=0.0" + ") -> (Tensor y, Tensor scale)" + ); + ops.impl("npu_dequant_swiglu_quant", torch::kPrivateUse1, &vllm_fl::npu_dequant_swiglu_quant); + + ops.def( + "npu_scatter_nd_update_v2(" + "Tensor(a!) var, Tensor indices, Tensor update" + ") -> ()" + ); + ops.impl("npu_scatter_nd_update_v2", torch::kPrivateUse1, &vllm_fl::npu_scatter_nd_update_v2); + + // This operator is planned to be integrated into PTA in the near future. + // Once that happens, the implementation in csrc will be removed. + ops.def( + "npu_lightning_indexer_quant(Tensor query, Tensor key, Tensor weights, Tensor query_dequant_scale, " + " Tensor key_dequant_scale, *, Tensor? actual_seq_lengths_query=None, " + " Tensor? actual_seq_lengths_key=None, Tensor? block_table=None, " + " int query_quant_mode=0, int key_quant_mode=0, " + " str layout_query='BSND', str layout_key='BSND'," + " int sparse_count=2048, int sparse_mode=3) -> Tensor" + ); + ops.impl("npu_lightning_indexer_quant", torch::kPrivateUse1, &vllm_fl::npu_lightning_indexer_quant); + // N-gram spec decode + ops.def( + "npu_ngram_spec_decode(Tensor(a!) token_ids, Tensor num_tokens_no_spec, " + "Tensor sampled_token_ids, Tensor discard_request_mask, " + "int vocab_size, int min_n, int max_n, int k) -> " + "(Tensor token_ids, Tensor next_token_ids, Tensor draft_token_ids, Tensor num_valid_draft_tokens)" + ); + ops.impl("npu_ngram_spec_decode", torch::kPrivateUse1, + &vllm_fl::npu_ngram_spec_decode); + + ops.def( + "chunk_gated_delta_rule_fwd_h(Tensor k, Tensor w, Tensor u, Tensor? g=None, *, Tensor? gk=None, Tensor? initial_state=None, bool? output_final_state=False, int? chunk_size=None, bool? save_new_value=True, int[]? cu_seqlens=None, int[]? chunk_indices=None, bool? use_exp2=False, bool? transpose_state_layout=False) -> (Tensor h_out, Tensor v_new_out, Tensor final_state_out)" + ); + ops.impl("chunk_gated_delta_rule_fwd_h", torch::kPrivateUse1, &vllm_fl::chunk_gated_delta_rule_fwd_h); + + ops.def( + "chunk_fwd_o(Tensor q, Tensor k, Tensor v, Tensor h, float scale, *, Tensor? g=None, Tensor? g_gamma=None, int[]? cu_seqlens=None, int[]? chunk_indices=None, int? chunk_size=None, bool? transpose_state_layout=False) -> Tensor" + ); + ops.impl("chunk_fwd_o", torch::kPrivateUse1, &vllm_fl::chunk_fwd_o); + + ops.def( + "store_kv_block(Tensor key_in, Tensor key_cache_in, Tensor group_len, Tensor group_key_idx,Tensor group_key_cache_idx, int block_size=0) -> ()" + ); + ops.impl("store_kv_block", torch::kPrivateUse1, &vllm_fl::store_kv_block); + + // Fused GDN gating. + ops.def( + "npu_fused_gdn_gating(Tensor A_log, " + " Tensor a, " + " Tensor b, " + " Tensor dt_bias, " + " float beta=1.0, " + " float threshold=20.0) -> (Tensor g, Tensor beta_output)"); + ops.impl("npu_fused_gdn_gating", torch::kPrivateUse1, &vllm_fl::npu_fused_gdn_gating); +} +#endif diff --git a/csrc/ascend/torch_binding_meta.cpp b/csrc/ascend/torch_binding_meta.cpp new file mode 100644 index 000000000..dd8cc73b2 --- /dev/null +++ b/csrc/ascend/torch_binding_meta.cpp @@ -0,0 +1,1802 @@ +#include +#include +#include +#include +#include +#include +#include "utils.h" +/* + * How to write a meta implementation for a custom operator (meta kernel): + * + * Meta implementations are used for shape and dtype inference, tracing, and export. + * They do NOT perform any real computation or allocate device memory. + * Instead, they return empty tensors with the correct shapes, dtypes, and device types. + * + * Steps to write a meta implementation: + * 1. The function signature should match the operator's schema, but only use the arguments + * necessary to infer output shapes and dtypes. + * 2. Use input tensor shapes, dtypes, and any relevant arguments to compute the output shapes. + * 3. Return empty tensors (e.g., at::empty_symint, at::empty_like) with the correct shape and dtype. + * 4. Do NOT perform any real computation or data movement. + * 5. Register the meta implementation with the "Meta" dispatch key using TORCH_LIBRARY_IMPL or similar. + * + * Example: + * std::tuple my_op_meta( + * at::Tensor &input, int64_t some_param) { + * // Infer output shape based on input and parameters + * auto out_shape = ...; + * at::Tensor out = at::empty_symint(out_shape, input.options()); + * // Return empty tensor(s) with correct shape/dtype + * return {out, ...}; + * } + * + * See below for real examples. + */ + +namespace vllm_fl { +namespace meta { +const int64_t INT4_NUMS_IN_INT32 = 8; +constexpr int64_t DSA_SLOT_MAPPING_FLAT = 1; +constexpr int64_t DSA_SLOT_MAPPING_BLOCK_OFFSET = 2; + +c10::SymInt ceil_div(const c10::SymInt& value, int64_t divisor) +{ + return (value + c10::SymInt(divisor - 1)) / c10::SymInt(divisor); +} + +#ifdef VLLM_ENABLE_ATB_AND_DIRECT_KERNELS +at::Tensor bgmv_expand_meta(at::Tensor &x, at::Tensor &weight, at::Tensor &indices, at::Tensor &y, + int64_t slice_offset, int64_t slice_size) { + at::Tensor y_out = at::empty_like(y); + return y_out; +} + +at::Tensor sgmv_expand_meta(at::Tensor &x, at::Tensor &weight, at::Tensor &lora_indices, at::Tensor &seq_len, + at::Tensor &y, int64_t slice_offset, int64_t slice_size) { + at::Tensor y_out = at::empty_like(y); + return y_out; +} + +std::tuple mla_preprocess( + const at::Tensor &hiddenState, + const at::Tensor &wdqkv, + const c10::optional &descale0, + const at::Tensor &gamma1, + const c10::optional &beta1, + const at::Tensor &wuq, + const c10::optional &descale1, + const at::Tensor &gamma2, + const at::Tensor &cos, + const at::Tensor &sin, + const at::Tensor &wuk, + const at::Tensor &kv_cache, + const at::Tensor &kv_cache_rope, + const at::Tensor &slotmapping, + const c10::optional &quant_scale0, + const c10::optional &quant_offset0, + const c10::optional &bias0, + const c10::optional &quant_scale1, + const c10::optional &quant_offset1, + const c10::optional &bias1, + const c10::optional &ctkv_scale, + const c10::optional &q_nope_scale, + c10::optional cache_mode, + c10::optional quant_mode, + c10::optional enable_inner_out, + at::Tensor &q_out0, + at::Tensor &kv_cache_out0, + at::Tensor &q_out1, + at::Tensor &kv_cache_out1, + at::Tensor &inner_out + ) +{ + return {q_out0, kv_cache_out0, q_out1, kv_cache_out1, inner_out}; +} + +void batch_matmul_transpose(const at::Tensor &tensor_a, const at::Tensor &tensor_b, at::Tensor &tensor_c, + c10::optional format_mode, + c10::optional quant_mode) +{ + return; +} +#endif + +void device_print_meta(c10::string_view msg) +{ + (void)msg; +} + +void device_print_tensor_meta(const at::Tensor& tensor) +{ + (void)tensor; +} + +std::tuple grouped_matmul_swiglu_quant( + const at::Tensor &x, const at::Tensor &weight, const at::Tensor &weight_scale, const at::Tensor &x_scale, + const at::Tensor &group_list, const c10::optional &bias, const c10::optional &offset, + double swiglu_limit) +{ + auto m = x.sym_size(0); + auto n = weight.sym_size(2); + bool is_a8w4 = x.dtype() == at::kChar && weight.dtype() == at::kInt; + if (is_a8w4) { + n *= c10::SymInt(INT4_NUMS_IN_INT32); + } + c10::SymDimVector output_shape = {m, n / c10::SymInt(2)}; + c10::SymDimVector scale_shape = {m}; + c10::SymDimVector scalar_shape; + at::Tensor output = at::empty_symint(output_shape, x.options().dtype(c10::ScalarType::Char)); + at::Tensor output_scale = at::empty_symint(scale_shape, x.options().dtype(c10::ScalarType::Float)); + at::Tensor output_offset = at::empty_symint(scalar_shape, x.options().dtype(c10::ScalarType::Float)); + return {output, output_scale, output_offset}; +} + +std::tuple grouped_matmul_swiglu_quant_weight_nz_tensor_list_meta( + const at::Tensor & x, + const at::TensorList & weight, + const at::TensorList & weight_scale, + const at::Tensor & x_scale, + const at::Tensor & group_list, + const c10::optional & bias, + const c10::optional & offset, + double swiglu_limit) +{ + auto m = x.sym_size(0); + auto n = weight[0].sym_size(1); + + c10::SymDimVector output_shape = {m, n / c10::SymInt(2)}; + c10::SymDimVector scale_shape = {m}; + at::Tensor output = at::empty_symint(output_shape, x.options().dtype(c10::ScalarType::Char)); + at::Tensor output_scale = at::empty_symint(scale_shape, x.options().dtype(c10::ScalarType::Float)); + at::Tensor output_offset = at::empty_symint(scale_shape, x.options().dtype(c10::ScalarType::Float)); + + return std::tuple(output, output_scale, output_offset); +} + +std::tuple grouped_matmul_swiglu_quant_v2_meta( + const at::Tensor & x, + const at::TensorList &weight, + const at::TensorList &weight_scale, + const at::Tensor & x_scale, + const at::Tensor & group_list, + const c10::optional & smooth_scale, + const c10::optional weight_assist_matrix, + const c10::optional & bias, + c10::optional dequant_mode, + c10::optional dequant_dtype, + c10::optional quant_mode, + c10::optional quant_dtype, + bool transpose_weight, + int64_t group_list_type, + at::IntArrayRef tuning_config, + double swiglu_limit) +{ + + auto m = x.sym_size(0); + auto n = weight_scale[0].sym_size(weight_scale[0].dim() - 1); + + c10::SymDimVector output_shape = {m, n / c10::SymInt(2)}; + c10::SymDimVector scale_shape = {m}; + at::Tensor output = at::empty_symint(output_shape, x.options().dtype(at::kChar)); + at::Tensor output_scale = at::empty_symint(scale_shape, x.options().dtype(at::kFloat)); + + + + return std::tuple(output, output_scale); +} +std::tuple dispatch_ffn_combine_meta( + const at::Tensor& x, + const at::TensorList& weight1, + const at::TensorList& weight2, + const at::Tensor& expert_idx, + const at::TensorList& scale1, + const at::TensorList& scale2, + const at::TensorList& bias1, + const at::TensorList& bias2, + const at::Tensor& probs, + c10::string_view group, + int64_t max_output_size, + at::Tensor& out, + at::Tensor& expert_token_nums, + const c10::optional &x_active_mask, + double swiglu_limit +) { + return {out, expert_token_nums}; +} + +std::tuple npu_lightning_indexer_meta( + const at::Tensor &query, const at::Tensor &key, const at::Tensor &weights, + const c10::optional &actual_seq_lengths_query, + const c10::optional &actual_seq_lengths_key, + const c10::optional &block_table, c10::string_view layout_query, + c10::string_view layout_key, int64_t sparse_count, int64_t sparse_mode, + int64_t pre_tokens, int64_t next_tokens, bool return_value) +{ + constexpr int64_t DIM_0 = 0; + constexpr int64_t DIM_1 = 1; + constexpr int64_t DIM_2 = 2; + + TORCH_CHECK(sparse_count > 0, "sparse count should be greater than 0, but now is ", sparse_count); + + std::string query_layout_str = std::string(layout_query); + std::string key_layout_str = std::string(layout_key); + c10::SymDimVector output_size; + if (query_layout_str == "BSND") { + output_size = {query.sym_size(DIM_0), query.sym_size(DIM_1), key.sym_size(DIM_2), c10::SymInt(sparse_count)}; + } else { + int n_dim_index = 0; + n_dim_index = (key_layout_str == "TND") ? DIM_1 : DIM_2; + output_size = {query.sym_size(DIM_0), key.sym_size(n_dim_index), c10::SymInt(sparse_count)}; + } + // construct the output tensor + at::Tensor sparse_indices_out = at::empty_symint(output_size, query.options().dtype(at::kInt)); + at::Tensor sparse_values_out; + if (return_value) { + sparse_values_out = at::empty_symint(output_size, query.options().dtype(query.dtype())); + } else { + sparse_values_out = at::empty_symint(c10::SymDimVector{c10::SymInt(0)}, query.options().dtype(query.dtype())); + } + return std::tuple(sparse_indices_out, sparse_values_out); +} + +std::tuple npu_sparse_flash_attention_meta( + const at::Tensor &query, const at::Tensor &key, const at::Tensor &value, + const at::Tensor &sparse_indices, double scale_value, + const c10::optional &block_table, + const c10::optional &actual_seq_lengths_query, + const c10::optional &actual_seq_lengths_kv, + const c10::optional &query_rope, + const c10::optional &key_rope, int64_t sparse_block_size, + c10::string_view layout_query, c10::string_view layout_kv, + int64_t sparse_mode, int64_t pre_tokens, int64_t next_tokens, + int64_t attention_mode, bool return_softmax_lse) +{ + constexpr int64_t DIM_0 = 0; + constexpr int64_t DIM_1 = 1; + constexpr int64_t DIM_2 = 2; + constexpr int64_t DIM_3 = 3; + constexpr int64_t DIM_4 = 4; + + std::string layout_query_str = std::string(layout_query); + TORCH_CHECK(layout_query_str == "BSND" || layout_query_str == "TND", + "The layout of query only support BSND and TND, but got ", + layout_query_str); + c10::SymDimVector output_size; + if (layout_query_str == "TND") { + TORCH_CHECK(query.dim() == DIM_3, + "When the layout of query is TND, the query dimension must be 3, but got ", + query.dim()); + output_size = {query.sym_size(DIM_0), query.sym_size(DIM_1), query.sym_size(DIM_2)}; + } else { + TORCH_CHECK(query.dim() == DIM_4, + "When the layout of query is BSND, the query dimension must be 4, but got ", + query.dim()); + output_size = {query.sym_size(DIM_0), query.sym_size(DIM_1), query.sym_size(DIM_2), query.sym_size(DIM_3)}; + } + + at::Tensor output = at::empty_symint(output_size, query.options().dtype(query.dtype())); + c10::SymDimVector softmax_size; + if (return_softmax_lse) { + if (query.dim() == DIM_3) { + const auto layout_kv_str = std::string(layout_kv); + const auto kv_head_num = + layout_kv_str == "PA_BSND" ? key.sym_size(DIM_2) : key.sym_size(DIM_1); + softmax_size = { + kv_head_num, + query.sym_size(DIM_0), + query.sym_size(DIM_1) / kv_head_num, + }; + } else { + softmax_size = { + query.sym_size(DIM_0), + key.sym_size(DIM_2), + query.sym_size(DIM_1), + query.sym_size(DIM_2) / key.sym_size(DIM_2), + }; + } + } else { + softmax_size = {c10::SymInt(0)}; + } + + at::Tensor softmax_max = at::empty_symint(softmax_size, query.options().dtype(at::kFloat)); + at::Tensor softmax_sum = at::empty_symint(softmax_size, query.options().dtype(at::kFloat)); + return std::tuple(output, softmax_max, softmax_sum); +} + +std::tuple npu_kv_quant_sparse_flash_attention_meta( + const at::Tensor &query, + const at::Tensor &key, + const at::Tensor &value, + const at::Tensor &sparse_indices, + double scale_value, + int64_t key_quant_mode, + int64_t value_quant_mode, + const c10::optional &key_dequant_scale, + const c10::optional &value_dequant_scale, + const c10::optional &block_table, + const c10::optional &actual_seq_lengths_query, + const c10::optional &actual_seq_lengths_kv, + int64_t sparse_block_size, + c10::string_view layout_query, + c10::string_view layout_kv, + int64_t sparse_mode, + int64_t pre_tokens, + int64_t next_tokens, + int64_t attention_mode, + int64_t quant_scale_repo_mode, + int64_t tile_size, + int64_t rope_head_dim, + bool return_softmax_lse) +{ + constexpr int64_t DIM_0 = 0; + constexpr int64_t DIM_1 = 1; + constexpr int64_t DIM_2 = 2; + constexpr int64_t DIM_3 = 3; + constexpr int64_t DIM_4 = 4; + + std::string layout_query_str = std::string(layout_query); + std::string layout_kv_str = std::string(layout_kv); + TORCH_CHECK(layout_query_str == "BSND" || layout_query_str == "TND", + "The layout of query only support BSND and TND, but got ", + layout_query_str); + c10::SymDimVector output_size; + if (layout_query_str == "BSND") { + TORCH_CHECK(query.dim() == DIM_4, + "When the layout of query is BSND, the query dimension must be 4, but got ", + query.dim()); + output_size = {query.sym_size(DIM_0), query.sym_size(DIM_1), query.sym_size(DIM_2), + query.sym_size(DIM_3) - c10::SymInt(rope_head_dim)}; + } else { + TORCH_CHECK(query.dim() == DIM_3, + "When the layout of query is TND, the query dimension must be 3, but got ", + query.dim()); + output_size = {query.sym_size(DIM_0), query.sym_size(DIM_1), + query.sym_size(DIM_2) - c10::SymInt(rope_head_dim)}; + } + + at::Tensor output = at::empty_symint(output_size, query.options().dtype(query.dtype())); + c10::SymDimVector softmax_size; + if (return_softmax_lse) { + if (query.dim() == DIM_3) { + const c10::SymInt kv_head_dim = + layout_kv_str == "PA_BSND" ? key.sym_size(DIM_2) : key.sym_size(DIM_1); + softmax_size = {kv_head_dim, query.sym_size(DIM_0), + query.sym_size(DIM_1) / kv_head_dim}; + } else { + softmax_size = { + query.sym_size(DIM_0), key.sym_size(DIM_2), query.sym_size(DIM_1), + query.sym_size(DIM_2) / key.sym_size(DIM_2)}; + } + } else { + softmax_size = {c10::SymInt(0)}; + } + + at::Tensor softmax_max = at::empty_symint(softmax_size, query.options().dtype(at::kFloat)); + at::Tensor softmax_sum = at::empty_symint(softmax_size, query.options().dtype(at::kFloat)); + return std::tuple(output, softmax_max, softmax_sum); +} + +std::tuple matmul_allreduce_add_rmsnorm_meta( + const at::Tensor &x1, + const at::Tensor &x2, + const at::Tensor &residual, + const at::Tensor &gamma, + c10::string_view group_tp, + int64_t tp_rank_size, + int64_t tp_rank_id, + double epsilon, + bool is_trans_b, + bool is_gather_add_out) + { + at::Tensor output = at::empty_like(residual); + at::Tensor add_out = at::empty_like(residual); + + return {output, add_out}; + } + +std::tuple npu_moe_init_routing_custom_meta( + const at::Tensor &x, const at::Tensor &expert_idx, + const c10::optional &scale, const c10::optional &offset, int64_t active_num, + int64_t expert_capacity, int64_t expert_num, int64_t drop_pad_mode, int64_t expert_tokens_num_type, + bool expert_tokens_num_flag, int64_t quant_mode, at::IntArrayRef active_expert_range, int64_t row_idx_type) +{ + constexpr int64_t DIM_X = 2; + constexpr int64_t DIM_EXPERT_IDX = 2; + constexpr int64_t LENGTH_ACTIVE_EXPERT_RANGE = 2; + constexpr int64_t EXPERT_TOKENS_COUNT = 1; + constexpr int64_t EXPERT_TOKENS_KEY_VALUE = 2; + constexpr int64_t QUANT_MODE_UNQUANT = -1; + constexpr int64_t QUANT_MODE_DYNAMIC_QUANT = 1; + constexpr int64_t CUMSUM = 0; + constexpr int64_t COUNT = 1; + constexpr int64_t KEY_VALUE = 2; + + if (active_expert_range.empty()) { + active_expert_range = at::IntArrayRef({0, expert_num}); + } + + int64_t x_dim = x.dim(); + TORCH_CHECK(x_dim == DIM_X, "The x should be ", DIM_X, + "-Dimension, current is ", x_dim, "-Dimension."); + + int64_t expert_idx_dim = expert_idx.dim(); + TORCH_CHECK(expert_idx_dim == DIM_EXPERT_IDX, "The expert_idx should be ", DIM_EXPERT_IDX, + "-Dimension, current is ", expert_idx_dim, "-Dimension."); + + // symbolic-meta-ok: active_expert_range is an IntArrayRef schema argument, not a Tensor shape. + int64_t active_expert_range_length = active_expert_range.size(); + TORCH_CHECK(active_expert_range_length == LENGTH_ACTIVE_EXPERT_RANGE, "The active_expert_range should be ", LENGTH_ACTIVE_EXPERT_RANGE, + "-Dimension, current is ", expert_idx_dim, "-Dimension."); + + int expert_length = active_expert_range[1] - active_expert_range[0]; + auto bs = x.sym_size(0); + auto h = x.sym_size(1); + auto k = expert_idx.sym_size(1); + c10::SymInt expanded_scale_len(0); + at::Tensor expanded_x; + + if (drop_pad_mode == 1) { // Drop/Pad + c10::SymDimVector expanded_x_shape = {c10::SymInt(expert_num), c10::SymInt(expert_capacity), h}; + if (quant_mode == QUANT_MODE_UNQUANT) { + expanded_x = at::empty_symint(expanded_x_shape, x.options()); + } else { + expanded_x = at::empty_symint(expanded_x_shape, x.options().dtype(at::kChar)); + } + expanded_scale_len = c10::SymInt(expert_num * expert_capacity); + } else { // Dropless / Active + if (active_num > 0) { // Active + c10::SymInt num_out_tokens = (bs * k).min(c10::SymInt(active_num)); + c10::SymDimVector expanded_x_shape = {num_out_tokens, h}; + if (quant_mode == QUANT_MODE_UNQUANT) { + expanded_x = at::empty_symint(expanded_x_shape, x.options()); + } else { + expanded_x = at::empty_symint(expanded_x_shape, x.options().dtype(at::kChar)); + } + expanded_scale_len = num_out_tokens; + } else { // Dropless + c10::SymDimVector expanded_x_shape = {bs * k, h}; + if (quant_mode == QUANT_MODE_UNQUANT) { + expanded_x = at::empty_symint(expanded_x_shape, x.options()); + } else { + expanded_x = at::empty_symint(expanded_x_shape, x.options().dtype(at::kChar)); + } + expanded_scale_len = bs * k; + } + } + + c10::SymDimVector expanded_row_idx_shape = {bs * k}; + at::Tensor expanded_row_idx = at::empty_symint(expanded_row_idx_shape, expert_idx.options()); + at::Tensor expert_tokens_count_or_cumsum; + if (expert_tokens_num_type >= CUMSUM && expert_tokens_num_type <= COUNT) { + // expert_tokens_count_or_cumsum in [end-start, ] + expert_tokens_count_or_cumsum = at::empty_symint( + c10::SymDimVector{c10::SymInt(expert_length)}, x.options().dtype(at::kLong)); + } else if (expert_tokens_num_type == KEY_VALUE) { + // key_value in [2, end-start] + expert_tokens_count_or_cumsum = at::empty_symint( + c10::SymDimVector{c10::SymInt(expert_num), c10::SymInt(2)}, x.options().dtype(at::kLong)); + } + + at::Tensor expanded_scale = at::empty_symint( + c10::SymDimVector{expanded_scale_len}, x.options().dtype(at::kFloat)); + return {expanded_x, expanded_row_idx, expert_tokens_count_or_cumsum, expanded_scale}; +} +std::tuple moe_gating_top_k_meta( + const at::Tensor& x, + int64_t k, + int64_t k_group, + int64_t group_count, + int64_t group_select_mode, + int64_t renorm, + int64_t norm_type, + bool out_flag, + double routed_scaling_factor, + double eps, + const c10::optional& bias_opt + + ) +{ + TORCH_CHECK(x.dim() == 2, "The x should be 2D"); + TORCH_CHECK( + x.scalar_type() == at::kHalf || x.scalar_type() == at::kFloat || x.scalar_type() == at::kBFloat16, + "float16、float32 or bfloat16 tensor expected but got a tensor with dtype: ", + x.scalar_type()); + + auto rows = x.sym_size(0); + auto expert_num = x.sym_size(1); + const at::Tensor &bias = c10::value_or_else(bias_opt, [] { return at::Tensor(); }); + if (bias.defined()) { + TORCH_CHECK(x.scalar_type() == bias.scalar_type(), "The dtype of x and bias should be same"); + TORCH_CHECK(bias.dim() == 1, "The bias should be 1D"); + } + at::Tensor y = at::empty_symint(c10::SymDimVector{rows, c10::SymInt(k)}, x.options()); + at::Tensor expert_idx = at::empty_symint(c10::SymDimVector{rows, c10::SymInt(k)}, x.options().dtype(at::kInt)); + at::Tensor out = at::empty_symint(c10::SymDimVector{rows, expert_num}, x.options().dtype(at::kFloat)); + + return std::tuple(y,expert_idx,out); +} + +std::tuple npu_add_rms_norm_bias_meta( + const at::Tensor& x1, + const at::Tensor& x2, + const at::Tensor& gamma, + const c10::optional &beta, + double epsilon) +{ + int64_t dim_x = x1.dim(); + int64_t dim_gamma = gamma.dim(); + int64_t diff = dim_x - dim_gamma; + c10::SymDimVector new_shape; + at::Tensor rstd; + + if (diff > 0) { + new_shape.reserve(dim_x); + auto x1_sizes = x1.sym_sizes(); + for (int64_t i = 0; i < diff; ++i) { + new_shape.push_back(x1_sizes[i]); + } + for (int64_t i = 0; i < dim_gamma; ++i) { + new_shape.push_back(c10::SymInt(1)); + } + } else { + new_shape.assign(dim_x, c10::SymInt(1)); + } + rstd = at::empty_symint(new_shape, x1.options().dtype(at::kFloat)); + at::Tensor y = at::empty_symint(x1.sym_sizes(), x1.options()); + at::Tensor x = at::empty_symint(x1.sym_sizes(), x1.options()); + return std::tuple(y, rstd, x); +} + +at::Tensor npu_reshape_and_cache_bnsd_meta(const at::Tensor& hashq, + const at::Tensor& hashkCache, + const at::Tensor& slotMapping, + const at::Tensor& seqLen, + const at::Tensor& hashkCacheOut) { + at::Tensor output = at::empty_symint( + hashkCache.sym_sizes(), hashkCache.options().dtype(hashkCache.dtype()).device(hashkCache.device())); + return output; +} + + +at::Tensor npu_hamming_dist_top_k_meta(const at::Tensor &hashq, + const at::Tensor &hashkCache, + const at::Tensor& hashkCacheRope, + const at::Tensor &topN, + const at::Tensor &seqLen, + const c10::optional &chunkSize, + const c10::optional maxSeqLen, + const c10::optional sink, + const c10::optional recent, + const c10::optional supportOffload, + const c10::optional &blockTable, + const c10::optional &mask, + const c10::optional& indices) { + if (indices.has_value()) { + return at::empty_like(indices.value()); + } + uint32_t MAX_BLOCK_PER_REQ_INHSA = 512; + + auto n_bs = hashq.sym_size(0); + auto n_kv_heads = hashkCache.sym_size(1); + auto n_max_kv = MAX_BLOCK_PER_REQ_INHSA; + at::Tensor out = at::empty_symint( + c10::SymDimVector{n_bs, n_kv_heads, c10::SymInt(n_max_kv)}, + torch::TensorOptions().dtype(torch::kInt32).device(hashq.device())); + return out; +} + +at::Tensor npu_sign_bits_pack_meta(const at::Tensor& input, + const int64_t size) { + auto ySize = ceil_div(input.sym_size(0), 8); + c10::SymInt outDim(0); + if (size != 0) { + outDim = ySize / c10::SymInt(size); + } + + at::Tensor out = at::empty_symint( + c10::SymDimVector{c10::SymInt(size), outDim}, + torch::TensorOptions().dtype(torch::kUInt8).device(input.device())); + return out; +} + +std::tuple npu_gemma_rms_norm_meta( + const at::Tensor& x, + const at::Tensor& gamma, + double epsilon) +{ + int64_t dim_x = x.dim(); + int64_t dim_gamma = gamma.dim(); + int64_t diff = dim_x - dim_gamma; + c10::SymDimVector new_shape; + at::Tensor rstd; + if (diff > 0) { + new_shape.reserve(dim_x); + auto x_sizes = x.sym_sizes(); + for (int64_t i = 0; i < diff; ++i) { + new_shape.push_back(x_sizes[i]); + } + for (int64_t i = 0; i < dim_gamma; ++i) { + new_shape.push_back(c10::SymInt(1)); + } + } else { + new_shape.assign(dim_x, c10::SymInt(1)); + } + rstd = at::empty_symint(new_shape, x.options().dtype(at::kFloat)); + at::Tensor y = at::empty_symint(x.sym_sizes(), x.options()); + return std::tuple(y, rstd); +} + +void transpose_kv_cache_by_block_meta( + const at::TensorList &k_cache, + const at::TensorList &v_cache, + const at::Tensor &block_ids, + int64_t block_size, + int64_t head_num, + int64_t head_dim, + int64_t split_num, + int64_t layer_num) +{ + return; +} + +std::tuple +npu_copy_and_expand_eagle_inputs_meta( + const at::Tensor &target_token_ids, + const at::Tensor &target_positions, + const at::Tensor &next_token_ids, + const at::Tensor &query_start_loc, + const at::Tensor &query_end_loc, + int64_t padding_token_id, + int64_t parallel_drafting_token_id, + int64_t num_padding_slots_per_request, + bool shift_input_ids, + int64_t total_draft_tokens) +{ + auto total_input_tokens = target_token_ids.sym_size(0); + auto num_reqs = query_start_loc.sym_size(0) - c10::SymInt(1); + + c10::SymDimVector draft_shape = {c10::SymInt(total_draft_tokens)}; + at::Tensor out_input_ids = at::empty_symint(draft_shape, target_token_ids.options()); + at::Tensor out_positions = at::empty_symint(draft_shape, target_token_ids.options()); + at::Tensor out_is_rejected_token_mask = at::empty_symint(draft_shape, target_token_ids.options().dtype(at::kChar)); + at::Tensor out_is_masked_token_mask = at::empty_symint(draft_shape, target_token_ids.options().dtype(at::kChar)); + at::Tensor out_new_token_indices = at::empty_symint( + c10::SymDimVector{num_reqs * c10::SymInt(num_padding_slots_per_request)}, target_token_ids.options()); + at::Tensor out_hidden_state_mapping = at::empty_symint( + c10::SymDimVector{total_input_tokens}, target_token_ids.options()); + + return {out_input_ids, out_positions, out_is_rejected_token_mask, out_is_masked_token_mask, + out_new_token_indices, out_hidden_state_mapping}; +} + +at::Tensor npu_causal_conv1d_custom_meta( + const at::Tensor& output, + const at::Tensor& x, + const at::Tensor& weight, + const at::Tensor& conv_state, + const c10::optional& bias_opt, + const c10::optional& query_start_loc_opt, + const c10::optional& cache_indices_opt, + const c10::optional& initial_state_mode_opt, + const c10::optional& num_accepted_tokens_opt, + int64_t activation_mode, + int64_t pad_slot_id, + int64_t run_mode) +{ + return output; +} + +at::Tensor npu_causal_conv1d_310_meta( + const at::Tensor& x, + const at::Tensor& weight, + const c10::optional& bias, + const at::Tensor& conv_states, + const c10::optional& query_start_loc, + const c10::optional& cache_indices, + const c10::optional& initial_state_mode, + const c10::optional& num_accepted_tokens, + int64_t activation_mode, + int64_t pad_slot_id, + int64_t run_mode) +{ + + at::Tensor output = at::empty_symint(x.sym_sizes(), x.options()); + return output; +} + +at::Tensor npu_recurrent_gated_delta_rule_310_meta( + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& beta, + at::Tensor& state, + const at::Tensor& actual_seq_lengths, + const at::Tensor& ssm_state_indices, + const c10::optional& g, + const c10::optional& gk, + const c10::optional& num_accepted_tokens, + double scale_value) +{ + + at::Tensor output = at::empty_symint(value.sym_sizes(), value.options()); + return output; +} + +at::Tensor npu_recurrent_gated_delta_rule_meta( + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + at::Tensor& state, + const c10::optional& beta, + const c10::optional scale, + const c10::optional& actual_seq_lengths, + const c10::optional& ssm_state_indices, + const c10::optional& num_accepted_tokens, + const c10::optional& g, + const c10::optional& gk) +{ + + auto options = value.options().dtype(at::ScalarType::BFloat16); + at::Tensor output = at::empty_symint(value.sym_sizes(), options); + return output; +} + +std::tuple npu_fused_gdn_gating_meta( + const at::Tensor& A_log, + const at::Tensor& a, + const at::Tensor& b, + const at::Tensor& dt_bias, + double beta, + double threshold) +{ + (void)beta; + (void)threshold; + auto batch = a.sym_size(0); + auto num_heads = a.sym_size(1); + + at::Tensor g = at::empty_symint( + c10::SymDimVector{c10::SymInt(1), batch, num_heads}, a.options().dtype(c10::kFloat)); + at::Tensor beta_output = at::empty_symint( + c10::SymDimVector{c10::SymInt(1), batch, num_heads}, b.options()); + + return std::make_tuple(g, beta_output); +} + +std::vector moe_grouped_matmul_meta( + at::Tensor x, + at::Tensor weight, + const at::Tensor& group_list, + int64_t split_item, + int64_t group_type, + int64_t group_list_type +) +{ + bool transpose_weight = false; + bool weight_nz = true; + + at::TensorList x_list = at::TensorList(x); + at::TensorList weight_list = at::TensorList(weight); + std::vector y; + c10::TensorOptions options = x[0].options().dtype(x[0].scalar_type()); + auto m = x[0].sym_size(0); + auto n = weight[0].sym_size(1); + if (!transpose_weight) { + n = weight[0].sym_size(2); + } + at::Tensor y_0 = at::empty_symint(c10::SymDimVector{m, n}, options); + y.emplace_back(y_0); + at::TensorList result = at::TensorList(y); + + return y; +} + +std::tuple moe_gating_top_k_hash_meta( + const at::Tensor& x, + int64_t k, + const c10::optional& bias_opt, + const c10::optional& input_ids_opt, + const c10::optional& tid2eid_opt, + int64_t k_group, + int64_t group_count, + double routed_scaling_factor, + double eps, + int64_t group_select_mode, + int64_t renorm, + int64_t norm_type, + bool out_flag) +{ + TORCH_CHECK(x.dim() == 2, "x must be 2D, but got dim=", x.dim()); + TORCH_CHECK( + x.scalar_type() == at::kHalf || x.scalar_type() == at::kFloat || x.scalar_type() == at::kBFloat16, + "x dtype must be float16/float32/bfloat16, but got ", x.scalar_type()); + + TORCH_CHECK(k > 0, "k must be > 0, but got k=", k); + TORCH_CHECK(k_group >= 1, "k_group must be >= 1, but got k_group=", k_group); + TORCH_CHECK(group_count >= 1, "group_count must be >= 1, but got group_count=", group_count); + + TORCH_CHECK(group_select_mode == 0 || group_select_mode == 1, + "group_select_mode must be 0 or 1, but got ", group_select_mode); + TORCH_CHECK(renorm == 0, + "renorm can only be 0 currently, but got ", renorm); + TORCH_CHECK(norm_type == 0 || norm_type == 1 || norm_type ==2, + "norm_type must be 0 (softmax) or 1 (sigmoid) or 2 (softplus), but got ", norm_type); + + TORCH_CHECK(eps > 0.0, "eps must be > 0, but got ", eps); + TORCH_CHECK(routed_scaling_factor > 0.0, + "routed_scaling_factor must be > 0, but got ", routed_scaling_factor); + + auto rows = x.sym_size(0); + auto expert_num = x.sym_size(1); + + if (bias_opt.has_value() && bias_opt->defined()) { + const auto& bias = *bias_opt; + TORCH_CHECK(bias.dim() == 1, "bias must be 1D, but got dim=", bias.dim()); + TORCH_CHECK(bias.scalar_type() == x.scalar_type(), + "bias dtype must equal x dtype. x=", x.scalar_type(), + ", bias=", bias.scalar_type()); + } + + if (input_ids_opt.has_value() && input_ids_opt->defined()) { + const auto& input_ids = *input_ids_opt; + TORCH_CHECK(input_ids.scalar_type() == at::kInt || input_ids.scalar_type() == at::kLong, + "input_ids dtype must be int32 or int64, but got ", input_ids.scalar_type()); + } + + if (tid2eid_opt.has_value() && tid2eid_opt->defined()) { + const auto& tid2eid = *tid2eid_opt; + TORCH_CHECK(tid2eid.scalar_type() == at::kInt || tid2eid.scalar_type() == at::kLong, + "tid2eid dtype must be int32 or int64, but got ", tid2eid.scalar_type()); + TORCH_CHECK(tid2eid.dim() >= 1, "tid2eid must have dim>=1, but got dim=", tid2eid.dim()); + } + + at::Tensor y = at::empty_symint(c10::SymDimVector{rows, c10::SymInt(k)}, x.options()); + at::Tensor expert_idx = at::empty_symint(c10::SymDimVector{rows, c10::SymInt(k)}, x.options().dtype(at::kInt)); + at::Tensor out = at::empty_symint(c10::SymDimVector{rows, expert_num}, x.options().dtype(at::kFloat)); + + return {y, expert_idx, out}; +} + +std::tuple construct_compressor_output_tensor(const at::Tensor &x, const at::Tensor &norm_weight, + const at::Tensor &rope_sin, int64_t cmp_ratio, int64_t coff) +{ + constexpr int DIM_3 = 3; + auto x_dim = x.dim(); + c10::SymDimVector cmp_kv_size; + at::Tensor cmp_kv; + c10::SymInt cmp_s(0); + if (x_dim == DIM_3) { + cmp_s = ceil_div(x.sym_size(1), cmp_ratio); + cmp_kv_size = {x.sym_size(0), cmp_s, norm_weight.sym_size(0)}; + } else { + cmp_s = rope_sin.sym_size(0); + cmp_kv_size = {cmp_s, norm_weight.sym_size(0)}; + } + + cmp_kv = at::empty_symint(cmp_kv_size, x.options().dtype(x.dtype())); + + return std::tuple(cmp_kv); +} + +std::tuple +compressor_meta(const at::Tensor &x, const at::Tensor &wkv, const at::Tensor &wgate, at::Tensor &state_cache, + const at::Tensor &ape, const at::Tensor &norm_weight, const at::Tensor &rope_sin, + const at::Tensor &rope_cos, const c10::optional &state_block_table, + const c10::optional &cu_seqlens, const c10::optional &seqused, + const c10::optional &start_pos, int64_t rope_head_dim, int64_t cmp_ratio, int64_t coff, + double norm_eps, int64_t rotary_mode, int64_t cache_mode) +{ + // construct the output tensor + auto x_dim = x.dim(); + auto norm_weight_dim = norm_weight.dim(); + auto rope_sin_dim = rope_sin.dim(); + + std::tuple output = construct_compressor_output_tensor(x, norm_weight, rope_sin, cmp_ratio, coff); + + return output; +} + +std::tuple compressor_metadata_meta( + const at::Tensor &rope_cos, const at::Tensor &rope_sin, const at::Tensor &cu_seqlens, + const at::Tensor &start_pos, const at::Tensor &kv_block_table, int64_t kv_block_size, + int64_t slot_mapping_format, int64_t compress_ratio, int64_t num_compressed_tokens, int64_t num_reqs_actual) +{ + constexpr int64_t VALUE_0 = 0; + + TORCH_CHECK(rope_cos.dim() == 2 && rope_sin.dim() == 2, "rope_cos and rope_sin should be 2D tensors"); + TORCH_CHECK(rope_cos.scalar_type() == rope_sin.scalar_type(), + "rope_cos and rope_sin should have same dtype"); + TORCH_CHECK(kv_block_size > VALUE_0, "kv_block_size should be greater than 0"); + TORCH_CHECK(compress_ratio > VALUE_0, "compress_ratio should be greater than 0"); + TORCH_CHECK(num_compressed_tokens > VALUE_0, "num_compressed_tokens should be greater than 0"); + TORCH_CHECK(num_reqs_actual > VALUE_0, "num_reqs_actual should be greater than 0"); + + c10::SymDimVector rope_output_size = { + c10::SymInt(num_compressed_tokens), c10::SymInt(1), c10::SymInt(1), rope_cos.sym_size(1)}; + at::Tensor compress_cos = at::empty_symint(rope_output_size, rope_cos.options()); + at::Tensor compress_sin = at::empty_symint(rope_output_size, rope_sin.options()); + + c10::SymDimVector slot_mapping_size; + if (slot_mapping_format == DSA_SLOT_MAPPING_BLOCK_OFFSET) { + slot_mapping_size = {c10::SymInt(num_compressed_tokens), c10::SymInt(2)}; + } else { + TORCH_CHECK(slot_mapping_format == DSA_SLOT_MAPPING_FLAT, + "slot_mapping_format should be 1(flat) or 2(block_offset), but got ", slot_mapping_format); + slot_mapping_size = {c10::SymInt(num_compressed_tokens)}; + } + at::Tensor slot_mapping = at::empty_symint(slot_mapping_size, kv_block_table.options().dtype(at::kInt)); + return std::make_tuple(compress_cos, compress_sin, slot_mapping); +} + +std::tuple construct_quant_lightning_indexer_output_tensor(const at::Tensor& query, const at::Tensor& key, + int64_t sparse_count, std::string query_layout_str, + std::string key_layout_str, bool return_value) +{ + constexpr int64_t DIM_0 = 0; + constexpr int64_t DIM_1 = 1; + constexpr int64_t DIM_2 = 2; + c10::SymDimVector output_size; + TORCH_CHECK(sparse_count > 0, "sparse count should be greater than 0, but now is ", sparse_count); + c10::SymInt keyHeadNum = (key_layout_str == "TND") ? key.sym_size(DIM_1) : key.sym_size(DIM_2); + if (query_layout_str == "BSND") { + output_size = {query.sym_size(DIM_0), query.sym_size(DIM_1), keyHeadNum, c10::SymInt(sparse_count)}; + } else { + output_size = {query.sym_size(DIM_0), keyHeadNum, c10::SymInt(sparse_count)}; + } + at::Tensor sparse_indices_out = at::empty_symint(output_size, query.options().dtype(at::kInt)); + at::Tensor sparse_values_out; + if (return_value) { + sparse_values_out = at::empty_symint(output_size, query.options().dtype(at::kFloat)); + } else { + sparse_values_out = at::empty_symint(c10::SymDimVector{c10::SymInt(0)}, query.options().dtype(at::kFloat)); + } + + return std::tuple(sparse_indices_out, sparse_values_out); +} + +std::tuple npu_vllm_quant_lightning_indexer_meta( + const at::Tensor &query, const at::Tensor &key, const at::Tensor &weights, + const at::Tensor &query_dequant_scale, const at::Tensor &key_dequant_scale, + int64_t query_quant_mode, int64_t key_quant_mode, + const c10::optional &actual_seq_lengths_query, + const c10::optional &actual_seq_lengths_key, + const c10::optional &block_table, + const c10::optional &metadata, + c10::string_view layout_query, c10::string_view layout_key, int64_t sparse_count, + int64_t sparse_mode, int64_t pre_tokens, int64_t next_tokens, int64_t cmp_ratio, bool return_value) +{ + std::string query_layout_str = std::string(layout_query); + std::string key_layout_str = std::string(layout_key); + std::tuple quant_lightning_indexer_output = construct_quant_lightning_indexer_output_tensor( + query, key, sparse_count, query_layout_str, key_layout_str, return_value); + at::Tensor sparse_indices_out = std::get<0>(quant_lightning_indexer_output); + at::Tensor sparse_values_out = std::get<1>(quant_lightning_indexer_output); + + return std::tuple(sparse_indices_out, sparse_values_out); +} + +std::tuple construct_output_tensor(const at::Tensor &q, std::string layout, + bool return_softmax_lse) +{ + at::Tensor output = at::empty_symint(q.sym_sizes(), q.options().dtype(q.dtype())); + at::Tensor softmax_lse; + if (return_softmax_lse) { + c10::SymDimVector lse_sizes(q.sym_sizes().begin(), q.sym_sizes().end()); + lse_sizes.back() = c10::SymInt(1); + softmax_lse = at::empty_symint(lse_sizes, q.options().dtype(c10::ScalarType::Float)); + } else { + softmax_lse = at::empty_symint(c10::SymDimVector{c10::SymInt(0)}, q.options().dtype(c10::ScalarType::Float)); + } + return std::tuple(output, softmax_lse); +} + +std::tuple npu_sparse_attn_sharedkv_meta(const at::Tensor &q, const c10::optional &ori_kv, + const c10::optional &cmp_kv, const c10::optional &ori_sparse_indices, + const c10::optional &cmp_sparse_indices, const c10::optional &ori_block_table, + const c10::optional &cmp_block_table, const c10::optional &cu_seqlens_q, + const c10::optional &cu_seqlens_ori_kv, const c10::optional &cu_seqlens_cmp_kv, + const c10::optional &seqused_q, const c10::optional &seqused_kv, + const c10::optional &sinks, const c10::optional &metadata, + double softmax_scale, int64_t cmp_ratio, int64_t ori_mask_mode, int64_t cmp_mask_mode, int64_t ori_win_left, + int64_t ori_win_right, c10::string_view layout_q, c10::string_view layout_kv, bool return_softmax_lse) +{ + std::string layout_q_str = std::string(layout_q); + std::tuple output = construct_output_tensor(q, layout_q_str, return_softmax_lse); + + return output; +} + +at::Tensor npu_sparse_attn_sharedkv_metadata_meta( + int64_t num_heads_q, + int64_t num_heads_kv, + int64_t head_dim, + const c10::optional &cu_seqlens_q, + const c10::optional &cu_seqlens_ori_kv, + const c10::optional &cu_seqlens_cmp_kv, + const c10::optional &seqused_q, + const c10::optional &seqused_kv, + int64_t batch_size, + int64_t max_seqlen_q, + int64_t max_seqlen_kv, + int64_t ori_topk, + int64_t cmp_topk, + int64_t cmp_ratio, + int64_t ori_mask_mode, + int64_t cmp_mask_mode, + int64_t ori_win_left, + int64_t ori_win_right, + c10::string_view layout_q, + c10::string_view layout_kv, + bool has_ori_kv, + bool has_cmp_kv, + const c10::string_view device) +{ + constexpr int64_t OUTPUT_SIZE = 1024; + at::Tensor output; + if (cu_seqlens_q.has_value()) { + output = at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(cu_seqlens_q.value().device())); + } else if (cu_seqlens_ori_kv.has_value()) { + output = at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(cu_seqlens_ori_kv.value().device())); + } else if (cu_seqlens_cmp_kv.has_value()) { + output = at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(cu_seqlens_cmp_kv.value().device())); + } else if (seqused_q.has_value()) { + output = at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(seqused_q.value().device())); + } else if (seqused_kv.has_value()) { + output = at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(seqused_kv.value().device())); + } else { + auto deviceOri = at::Device(std::string(device)); + std::string device_str = "meta"; + if (deviceOri.has_index()) { + device_str += ":"; + device_str += std::to_string(deviceOri.index()); + } + output = at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(at::Device(device_str))); + } + return output; +} + +at::Tensor npu_vllm_quant_lightning_indexer_metadata_meta( + int64_t num_heads_q, int64_t num_heads_k, int64_t head_dim, int64_t query_quant_mode, int64_t key_quant_mode, + const c10::optional &actual_seq_lengths_query, const c10::optional &actual_seq_lengths_key, int64_t batch_size, + int64_t max_seqlen_q, int64_t max_seqlen_k, const c10::string_view layout_query, c10::string_view layout_key, int64_t sparse_count, + int64_t sparse_mode, int64_t pre_tokens, int64_t next_tokens, int64_t cmp_ratio, const c10::string_view device) +{ + constexpr int64_t OUTPUT_SIZE = 1024; + at::Tensor output; + if (actual_seq_lengths_query.has_value()) { + output = at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(actual_seq_lengths_query.value().device())); + } else if (actual_seq_lengths_key.has_value()) { + output = at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(actual_seq_lengths_key.value().device())); + } else { + auto deviceOri = at::Device(std::string(device)); + std::string device_str = "meta"; + if (deviceOri.has_index()) { + device_str += ":"; + device_str += std::to_string(deviceOri.index()); + } + output = at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(at::Device(device_str))); + } + + return output; +} + +at::Tensor construct_hc_post_output_tensor(const at::Tensor& residual) +{ + c10::SymIntArrayRef output_size = residual.sym_sizes(); + at::Tensor out = at::empty_symint(output_size, residual.options().dtype(residual.dtype())); + return out; +} + +at::Tensor npu_hc_post_meta( + const at::Tensor& x, + const at::Tensor& residual, + const at::Tensor& post, + const at::Tensor& comb) +{ + at::Tensor outputs = construct_hc_post_output_tensor(residual); + return outputs; +} + +std::tuple construct_hc_pre_output_tensor(const at::Tensor& x, int64_t hc_mult) +{ + auto xDims = x.dim(); + at::SmallVector y_size; + at::SmallVector post_size; + at::SmallVector comb_frag_size; + if (xDims == 4) { + auto batch = x.sym_size(0); + auto size = x.sym_size(1); + auto d = x.sym_size(3); + y_size = {batch, size, d}; + post_size = {batch, size, hc_mult}; + comb_frag_size = {batch, size, hc_mult, hc_mult}; + } else if (xDims == 3){ + auto bs = x.sym_size(0); + auto d = x.sym_size(2); + y_size = {bs, d}; + post_size = {bs, hc_mult}; + comb_frag_size = {bs, hc_mult, hc_mult}; + } + + at::Tensor y = at::empty_symint(c10::SymIntArrayRef(y_size), x.options().dtype(at::kBFloat16)); + at::Tensor post = at::empty_symint(c10::SymIntArrayRef(post_size), x.options().dtype(at::kFloat)); + at::Tensor comb_frag = at::empty_symint(c10::SymIntArrayRef(comb_frag_size), x.options().dtype(at::kFloat)); + + return std::tuple(y, post, comb_frag); +} + +at::Tensor construct_hc_pre_rsqrt_output_tensor(const at::Tensor& x, float epsilon=1e-6) +{ + TORCH_CHECK(epsilon >= 0, "epsilon should be greater than 0."); + + auto options = x.options(); + auto xDims = x.dim(); + c10::SymDimVector yOut_shape; + for (size_t i = 0; i < xDims - 2; i++) { + yOut_shape.push_back(x.sym_size(i)); + } + yOut_shape.push_back(c10::SymInt(1)); + at::Tensor yOut = at::empty_symint(yOut_shape, options.dtype(at::kFloat)); + + return yOut; +} + +std::tuple npu_hc_pre_meta( + const at::Tensor& x, const at::Tensor& hc_fn, const at::Tensor& hc_scale, const at::Tensor& hc_base, + int64_t hc_mult, int64_t hc_sinkhorn_iters, double norm_eps, double hc_eps) +{ + auto output_tensors = construct_hc_pre_output_tensor(x, hc_mult); + at::Tensor y = std::get<0>(output_tensors); + at::Tensor post = std::get<1>(output_tensors); + at::Tensor comb_frag = std::get<2>(output_tensors); + + return std::tuple(y, post, comb_frag); +} + +at::Tensor construct_hc_pre_inv_rms_output_tensor(const at::Tensor& x, float epsilon=1e-20) +{ + TORCH_CHECK(epsilon >= 0, "epsilon should be greater than 0."); + + auto options = x.options(); + auto xDims = x.dim(); + c10::SymDimVector yOut_shape; + for (auto i = 0; i < xDims - 2; i++) { + yOut_shape.push_back(x.sym_size(i)); + } + yOut_shape.push_back(c10::SymInt(1)); + at::Tensor yOut = at::empty_symint(yOut_shape, options.dtype(at::kFloat)); + + return yOut; +} + +at::Tensor npu_hc_pre_inv_rms_meta(const at::Tensor& x, double epsilon=1e-20) +{ + TORCH_CHECK(epsilon >= 0, "epsilon should be greater than 0."); + + at::Tensor yOut; + yOut = construct_hc_pre_inv_rms_output_tensor(x, epsilon); + + return yOut; +} + +std::tuple construct_hc_pre_sinkhorn_output_tensor(const at::Tensor& mixes, const at::Tensor& x, int64_t hc_mult) +{ + auto xDims = x.dim(); + c10::SymDimVector y_size; + c10::SymDimVector post_size; + c10::SymDimVector comb_frag_size; + if (xDims == 4) { + auto batch = x.sym_size(0); + auto size = x.sym_size(1); + auto d = x.sym_size(3); + y_size = {batch, size, d}; + post_size = {batch, size, c10::SymInt(hc_mult)}; + comb_frag_size = {batch, size, c10::SymInt(hc_mult), c10::SymInt(hc_mult)}; + } else if (xDims == 3){ + auto bs = x.sym_size(0); + auto d = x.sym_size(2); + y_size = {bs, d}; + post_size = {bs, c10::SymInt(hc_mult)}; + comb_frag_size = {bs, c10::SymInt(hc_mult), c10::SymInt(hc_mult)}; + } + + at::Tensor y = at::empty_symint(y_size, x.options().dtype(at::kBFloat16)); + at::Tensor post = at::empty_symint(post_size, x.options().dtype(at::kFloat)); + at::Tensor comb_frag = at::empty_symint(comb_frag_size, x.options().dtype(at::kFloat)); + + return std::tuple(y, post, comb_frag); +} + +std::tuple npu_hc_pre_sinkhorn_meta( + const at::Tensor& mixes, const at::Tensor& rsqrt, const at::Tensor& hc_scale, const at::Tensor& hc_base, + const at::Tensor& x, int64_t hc_mult, int64_t hc_sinkhorn_iters, double hc_eps) +{ + auto output_tensors = construct_hc_pre_sinkhorn_output_tensor(mixes, x, hc_mult); + at::Tensor y = std::get<0>(output_tensors); + at::Tensor post = std::get<1>(output_tensors); + at::Tensor comb_frag = std::get<2>(output_tensors); + + return std::tuple(y, post, comb_frag); +} + +void inplace_partial_rotary_mul_meta( + at::Tensor &x, + const at::Tensor &r1, + const at::Tensor &r2, + c10::string_view rotary_mode, + at::IntArrayRef partial_slice) +{ + auto origin_dim_num = x.dim(); + return; +} + +std::tuple npu_rms_norm_dynamic_quant_meta( + const at::Tensor& x, + const at::Tensor& gamma, + const c10::optional& smooth_scale, + const c10::optional& beta, + double epsilon) +{ + at::Tensor y_out = at::empty_like(x); + auto options = x.options(); + c10::SymDimVector scale_out_shape; + for (size_t i = 0; i < x.dim() - 1; i++) { + scale_out_shape.push_back(x.sym_size(i)); + } + at::Tensor scale_out = at::empty_symint(scale_out_shape, options.dtype(at::kFloat)); + + return std::make_tuple(y_out, scale_out); +} + +void indexer_compress_epilog_meta( + at::Tensor& indexer_compress_cache, + at::Tensor& indexer_compress_cache_scale, + const at::Tensor& x, + const at::Tensor& slot_mapping, + int64_t quant_mode = 1, + bool round_scale = true) +{ + return; +} + +void kv_compress_epilog_meta( + at::Tensor& kv_compress_cache, + const at::Tensor& x, + const at::Tensor& slot_mapping, + int64_t quant_group_size, + int64_t quant_mode, + bool round_scale_flag, + int64_t layout) +{ + return; +} + +std::tuple npu_kv_quant_sparse_attn_sharedkv_meta( + const at::Tensor& q, + int64_t kv_quant_mode, + const c10::optional& ori_kv, + const c10::optional& cmp_kv, + const c10::optional& ori_sparse_indices, + const c10::optional& cmp_sparse_indices, + const c10::optional& ori_block_table, + const c10::optional& cmp_block_table, + const c10::optional& cu_seqlens_q, + const c10::optional& cu_seqlens_ori_kv, + const c10::optional& cu_seqlens_cmp_kv, + const c10::optional& seqused_q, + const c10::optional& seqused_kv, + const c10::optional& sinks, + const c10::optional& metadata, + int64_t tile_size, + int64_t rope_head_dim, + double softmax_scale, + int64_t cmp_ratio, + int64_t ori_mask_mode, + int64_t cmp_mask_mode, + int64_t ori_win_left, + int64_t ori_win_right, + c10::string_view layout_q, + c10::string_view layout_kv, + bool return_softmax_lse) +{ + std::string layout_q_str = std::string(layout_q); + return construct_output_tensor(q, layout_q_str, return_softmax_lse); +} + +at::Tensor npu_kv_quant_sparse_attn_sharedkv_metadata_meta( + int64_t num_heads_q, + int64_t num_heads_kv, + int64_t head_dim, + int64_t kv_quant_mode, + const c10::optional& cu_seqlens_q, + const c10::optional& cu_seqlens_ori_kv, + const c10::optional& cu_seqlens_cmp_kv, + const c10::optional& seqused_q, + const c10::optional& seqused_kv, + int64_t batch_size, + int64_t max_seqlen_q, + int64_t max_seqlen_kv, + int64_t ori_topk, + int64_t cmp_topk, + int64_t tile_size, + int64_t rope_head_dim, + int64_t cmp_ratio, + int64_t ori_mask_mode, + int64_t cmp_mask_mode, + int64_t ori_win_left, + int64_t ori_win_right, + c10::string_view layout_q, + c10::string_view layout_kv, + bool has_ori_kv, + bool has_cmp_kv, + const c10::string_view device) +{ + constexpr int64_t OUTPUT_SIZE = 1024; + if (cu_seqlens_q.has_value()) { + return at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(cu_seqlens_q.value().device())); + } + if (cu_seqlens_ori_kv.has_value()) { + return at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(cu_seqlens_ori_kv.value().device())); + } + if (cu_seqlens_cmp_kv.has_value()) { + return at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(cu_seqlens_cmp_kv.value().device())); + } + if (seqused_q.has_value()) { + return at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(seqused_q.value().device())); + } + if (seqused_kv.has_value()) { + return at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(seqused_kv.value().device())); + } + + auto device_ori = at::Device(std::string(device)); + std::string device_str = "meta"; + if (device_ori.has_index()) { + device_str += ":"; + device_str += std::to_string(device_ori.index()); + } + return at::empty_symint( + c10::SymDimVector{c10::SymInt(OUTPUT_SIZE)}, + torch::dtype(torch::kInt32).device(at::Device(device_str))); +} + +int64_t get_type_code(at::ScalarType dst_type) +{ + switch (dst_type) { + case at::ScalarType::Float8_e5m2: + return 35; + case at::ScalarType::Float8_e4m3fn: + return 36; + case at::ScalarType::Half: + return 1; + case at::ScalarType::BFloat16: + return 27; + default: + TORCH_CHECK(false, "Unsupported dtype: ", dst_type); + } + return 0; +} + +std::tuple construct_swiglu_group_quant_output_tensor( + const at::Tensor& x, + int64_t dst_type, + int64_t quant_mode, + bool ue8m0_scale) +{ + constexpr int64_t SWIGLU_FACTOR = 2; + constexpr int64_t PER_BLOCK_FP16 = 128; + constexpr int64_t PER_MX_FP16 = 32; + constexpr int64_t MX_SCALE_ALIGN_FACTOR = 2; + constexpr int64_t GROUP_QUANT = 1; + constexpr int64_t MX_QUANT = 2; + constexpr int64_t FP8_QUANT = 3; + + c10::SymDimVector y_size(x.sym_sizes().begin(), x.sym_sizes().end()); + TORCH_CHECK(x.dtype() == at::kHalf || x.dtype() == at::kBFloat16, + "x should be FLOAT16 or BFLOAT16."); + TORCH_CHECK(quant_mode == GROUP_QUANT || quant_mode == MX_QUANT || quant_mode == FP8_QUANT, + "Unsupported quant mode, only support ", GROUP_QUANT, " or ", MX_QUANT, " or ", FP8_QUANT, "."); + + y_size.back() = y_size.back() / c10::SymInt(SWIGLU_FACTOR); + c10::SymInt y_last_dim = y_size.back(); + auto y_dtype = dst_type == 35 ? at::kFloat8_e5m2 : at::kFloat8_e4m3fn; + at::Tensor y = at::empty_symint(y_size, x.options().dtype(y_dtype)); + + c10::SymDimVector scale_size(y_size.begin(), y_size.end()); + if (quant_mode == GROUP_QUANT || quant_mode == FP8_QUANT) { + scale_size.back() = ceil_div(y_last_dim, PER_BLOCK_FP16); + } else if (quant_mode == MX_QUANT) { + c10::SymInt scale_last_dim = ceil_div(y_last_dim, PER_MX_FP16); + scale_last_dim = ceil_div(scale_last_dim, MX_SCALE_ALIGN_FACTOR); + scale_size.back() = scale_last_dim; + scale_size.push_back(c10::SymInt(MX_SCALE_ALIGN_FACTOR)); + } + + auto scale_type = at::kFloat; + if (quant_mode == MX_QUANT || (quant_mode == FP8_QUANT && ue8m0_scale)) { + scale_type = at::kFloat8_e8m0fnu; + } + at::Tensor scale = at::empty_symint(scale_size, x.options().dtype(scale_type)); + at::Tensor y_origin = at::empty_symint(y_size, x.options().dtype(x.dtype())); + + return std::tuple(y, scale, y_origin); +} + +std::tuple npu_swiglu_group_quant_meta( + const at::Tensor& x, + const c10::optional& topk_weight, + const c10::optional& group_index, + at::ScalarType dst_type = at::ScalarType::Float8_e4m3fn, + int64_t quant_mode = 1, + int64_t group_size = 128, + bool round_scale = false, + bool ue8m0_scale = false, + bool output_origin = false, + int64_t group_list_type = 0, + double clamp_value = 0.0) +{ + int64_t dst_type_code = get_type_code(dst_type); + return construct_swiglu_group_quant_output_tensor(x, dst_type_code, quant_mode, ue8m0_scale); +} + +std::tuple construct_load_index_kv_cache_output_tensor( + const at::Tensor& kv_cache, + const at::Tensor& slot_mapping) +{ + constexpr int64_t KV_LAST_DIM = 128; + auto n = slot_mapping.sym_size(0); + + at::Tensor kv = at::empty_symint( + c10::SymDimVector{n, c10::SymInt(KV_LAST_DIM)}, kv_cache.options().dtype(at::kFloat8_e4m3fn)); + at::Tensor kv_scale = at::empty_symint(c10::SymDimVector{n}, kv_cache.options().dtype(at::kFloat)); + + return std::tuple(kv, kv_scale); +} + +std::tuple npu_load_index_kv_cache_meta( + const at::Tensor& kv_cache, + const at::Tensor& slot_mapping) +{ + return construct_load_index_kv_cache_output_tensor(kv_cache, slot_mapping); +} + +void indexer_compress_epilog_v2_meta( + at::Tensor& indexer_compress_cache, + const at::Tensor& x, + const at::Tensor& slot_mapping, + int64_t layout = 2) +{ + return; +} + +std::tuple npu_dequant_swiglu_quant_meta( + const at::Tensor& x, + const c10::optional& weight_scale, + const c10::optional& activation_scale, + const c10::optional& bias, + const c10::optional& quant_scale, + const c10::optional& quant_offset, + const c10::optional& group_index, + bool activate_left, + int64_t quant_mode, + int64_t swiglu_mode, + double clamp_limit, + double glu_alpha, + double glu_bias) +{ + c10::SymDimVector y_size; + c10::SymDimVector scale_size; + for (int64_t i = 0; i < x.dim() - 1; ++i) { + y_size.push_back(x.sym_size(i)); + scale_size.push_back(x.sym_size(i)); + } + y_size.push_back(x.sym_size(x.dim() - 1) / c10::SymInt(2)); + + at::Tensor y = at::empty_symint(y_size, x.options().dtype(c10::ScalarType::Char)); + at::Tensor scale = at::empty_symint(scale_size, x.options().dtype(c10::ScalarType::Float)); + return {y, scale}; +} + +at::Tensor npu_lightning_indexer_quant_meta( + const at::Tensor &query, const at::Tensor &key, const at::Tensor &weights, + const at::Tensor &query_dequant_scale, const at::Tensor &key_dequant_scale, + const c10::optional &actual_seq_lengths_query, + const c10::optional &actual_seq_lengths_key, + const c10::optional &block_table, int64_t query_quant_mode, int64_t key_quant_mode, + c10::string_view layout_query, c10::string_view layout_key, int64_t sparse_count, int64_t sparse_mode) +{ + std::string query_layout_str = std::string(layout_query); + std::string key_layout_str = std::string(layout_key); + + const int DIM_0 = 0; + const int DIM_1 = 1; + const int DIM_2 = 2; + + c10::SymDimVector output_size; + TORCH_CHECK(sparse_count > 0, "sparse count should be greater than 0, but now is ", sparse_count); + c10::SymInt keyHeadNum = (key_layout_str == "TND") ? key.sym_size(DIM_1) : key.sym_size(DIM_2); + if (query_layout_str == "BSND") { + output_size = {query.sym_size(DIM_0), query.sym_size(DIM_1), keyHeadNum, c10::SymInt(sparse_count)}; + } else { + output_size = {query.sym_size(DIM_0), keyHeadNum, c10::SymInt(sparse_count)}; + } + at::Tensor lightning_indexer_quant_output = at::empty_symint(output_size, query.options().dtype(at::kInt)); + + return lightning_indexer_quant_output; +} + +void npu_scatter_nd_update_v2_meta( + at::Tensor& var, + const at::Tensor& indices, + const at::Tensor& update) +{ + return; +} + +// N-gram spec decode meta +std::tuple npu_ngram_spec_decode_meta( + at::Tensor &token_ids, + const at::Tensor &num_tokens_no_spec, + const at::Tensor &sampled_token_ids, + const at::Tensor &discard_request_mask, + int64_t vocab_size, + int64_t min_n, + int64_t max_n, + int64_t k) +{ + auto batch_size = token_ids.sym_size(0); + at::Tensor next_token_ids = at::empty_symint(c10::SymDimVector{batch_size}, token_ids.options()); + at::Tensor draft_token_ids = at::empty_symint( + c10::SymDimVector{batch_size, c10::SymInt(k)}, token_ids.options()); + at::Tensor num_valid_draft_tokens = at::empty_symint(c10::SymDimVector{batch_size}, token_ids.options()); + return std::make_tuple(token_ids, next_token_ids, draft_token_ids, num_valid_draft_tokens); +} + +std::tuple chunk_gated_delta_rule_fwd_h_meta( + const at::Tensor & k, + const at::Tensor & w, + const at::Tensor & u, + const c10::optional & g, + const c10::optional & gk, + const c10::optional & initial_state, + c10::optional output_final_state, + c10::optional chunk_size, + c10::optional save_new_value, + c10::optional cu_seqlens, + c10::optional chunk_indices, + c10::optional use_exp2, + c10::optional transpose_state_layout) +{ + bool output_final_state_ = output_final_state.has_value() ? output_final_state.value() : false; + const at::Tensor &initial_state_ = c10::value_or_else(initial_state, [] { return at::Tensor(); }); + int64_t chunk_size_ = chunk_size.has_value() ? chunk_size.value() : 64; + const at::Tensor &g_ = c10::value_or_else(g, [] { return at::Tensor(); }); + const at::Tensor &gk_ = c10::value_or_else(gk, [] { return at::Tensor(); }); + + auto K = k.sym_size(3); + auto B = k.sym_size(0); + auto T = k.sym_size(2); + auto HV = u.sym_size(1); + auto V = u.sym_size(3); + + c10::SymInt NT(0); + if (chunk_indices.has_value()) { + auto chunk_indices_ref = chunk_indices.value(); + // symbolic-meta-ok: chunk_indices is an IntArrayRef schema argument, not a Tensor shape. + NT = c10::SymInt(chunk_indices_ref.size() / 2); + } else { + NT = ceil_div(T, chunk_size_); + } + + at::Tensor h_out = at::empty_symint(c10::SymDimVector{B, HV, NT, K, V}, k.options()); + at::Tensor v_new_out = at::empty_symint(u.sym_sizes(), u.options()); + at::Tensor final_state_out; + if (output_final_state_) { + c10::SymInt N = cu_seqlens.has_value() ? c10::SymInt(cu_seqlens->size() - 1) : B; + auto state_options = initial_state.has_value() ? initial_state->options() : h_out.options(); + final_state_out = at::empty_symint(c10::SymDimVector{N, HV, K, V}, state_options); + } else { + final_state_out = at::empty_symint(c10::SymDimVector{c10::SymInt(1)}, k.options()); + } + + bool save_new_value_ = save_new_value.value_or(true); + bool use_exp2_ = use_exp2.value_or(false); + bool transpose_state_layout_ = transpose_state_layout.value_or(false); + + if (output_final_state_) { + return std::make_tuple(h_out, v_new_out, final_state_out); + } else { + return std::make_tuple(h_out, v_new_out, at::Tensor()); + } +} + +at::Tensor chunk_fwd_o_meta( + const at::Tensor & q, + const at::Tensor & k, + const at::Tensor & v, + const at::Tensor & h, + double scale, + const c10::optional & g, + const c10::optional & g_gamma, + c10::optional cu_seqlens, + c10::optional chunk_indices, + c10::optional chunk_size, + c10::optional transpose_state_layout) +{ + at::Tensor o = at::empty_symint(v.sym_sizes(), v.options()); + int64_t chunk_size_ = chunk_size.has_value() ? chunk_size.value() : 64; + const at::Tensor &g_ = c10::value_or_else(g, [] { return at::Tensor(); }); + (void)g_gamma; + (void)transpose_state_layout; + + return o; +} + +void store_kv_block_metadata( + const at::Tensor &slot_mapping_npu, + const at::Tensor &group_len, + const at::Tensor &group_key_idx, + const at::Tensor &group_key_cache_idx, + int64_t block_size) + { + return; + } + +void store_kv_block( + const at::Tensor &key_in, + const at::Tensor &key_cache_in, + const at::Tensor &group_len, + const at::Tensor &group_key_idx, + const at::Tensor &group_key_cache_idx, + int64_t block_size) +{ + return; + +} + +} // namespace meta +} // namespace vllm_fl + +// Register the meta implementations of the custom kernels for symbolic tracing, this will also +// the custom kernel been captured into aclgraph +#ifdef ASCEND_PLATFORM_310P +// Pybind on Ascend 310P +namespace { +TORCH_LIBRARY_IMPL_EXPAND(CONCAT(_C, _ascend), Meta, ops) { + // causal_conv1d_310 + ops.impl("npu_causal_conv1d_310", &vllm_fl::meta::npu_causal_conv1d_310_meta); + // npu_recurrent_gated_delta_rule_310 + ops.impl("npu_recurrent_gated_delta_rule_310", &vllm_fl::meta::npu_recurrent_gated_delta_rule_310_meta); + // chunk_gated_delta_rule_fwd_h + ops.impl("chunk_gated_delta_rule_fwd_h", &vllm_fl::meta::chunk_gated_delta_rule_fwd_h_meta); + // chunk_fwd_o + ops.impl("chunk_fwd_o", &vllm_fl::meta::chunk_fwd_o_meta); +} +} +#else +// Pybind on other platform +namespace { +TORCH_LIBRARY_IMPL_EXPAND(CONCAT(_C, _ascend), Meta, ops) { + //Gemma rmsnorm meta implementation + ops.impl("npu_gemma_rms_norm", &vllm_fl::meta::npu_gemma_rms_norm_meta); + // recurrent_gated_delta_rule meta implementation + ops.impl("npu_recurrent_gated_delta_rule", &vllm_fl::meta::npu_recurrent_gated_delta_rule_meta); + // Launch host print from device + ops.impl("device_print", &vllm_fl::meta::device_print_meta); + // launch host print from device for tensors + ops.impl("device_print_tensor", &vllm_fl::meta::device_print_tensor_meta); +#ifdef VLLM_ENABLE_ATB_AND_DIRECT_KERNELS + // Direct kernel meta implementations + // Bgmv expand + ops.impl("bgmv_expand", &vllm_fl::meta::bgmv_expand_meta); + // Sgmv expand + ops.impl("sgmv_expand", &vllm_fl::meta::sgmv_expand_meta); + // MLA preprocess + ops.impl("mla_preprocess", &vllm_fl::meta::mla_preprocess); + // batch_matmul_transpose + ops.impl("batch_matmul_transpose", &vllm_fl::meta::batch_matmul_transpose); +#endif + // grouped_matmul_swiglu_quant_weight_nz meta implementation + ops.impl("grouped_matmul_swiglu_quant_weight_nz", &vllm_fl::meta::grouped_matmul_swiglu_quant); + // grouped_matmul_swiglu_quant meta implementation + ops.impl("grouped_matmul_swiglu_quant", &vllm_fl::meta::grouped_matmul_swiglu_quant); + // Grouped matmul swiglu quant weight nz tensor list + ops.impl("grouped_matmul_swiglu_quant_weight_nz_tensor_list", &vllm_fl::meta::grouped_matmul_swiglu_quant_weight_nz_tensor_list_meta); + // Grouped matmul swiglu quant v2 + ops.impl("grouped_matmul_swiglu_quant_v2", &vllm_fl::meta::grouped_matmul_swiglu_quant_v2_meta); + // Lightning indexer + ops.impl("npu_lightning_indexer", &vllm_fl::meta::npu_lightning_indexer_meta); + // Sparse flash attention + ops.impl("npu_sparse_flash_attention", &vllm_fl::meta::npu_sparse_flash_attention_meta); + ops.impl("npu_kv_quant_sparse_flash_attention", + &vllm_fl::meta::npu_kv_quant_sparse_flash_attention_meta); + // MoE dispatch-ffn-combine + ops.impl("dispatch_ffn_combine", &vllm_fl::meta::dispatch_ffn_combine_meta); + // matmul allreduce add rmsnorm + ops.impl("matmul_allreduce_add_rmsnorm", &vllm_fl::meta::matmul_allreduce_add_rmsnorm_meta); + // moe_init_routing_custom + ops.impl("npu_moe_init_routing_custom", &vllm_fl::meta::npu_moe_init_routing_custom_meta); + // Moe_gating_top_k + ops.impl("moe_gating_top_k", &vllm_fl::meta::moe_gating_top_k_meta); + // Add_Rms_Norm_Bias + ops.impl("npu_add_rms_norm_bias", &vllm_fl::meta::npu_add_rms_norm_bias_meta); + // transpose_kv_cache_by_block + ops.impl("transpose_kv_cache_by_block", &vllm_fl::meta::transpose_kv_cache_by_block_meta); + // hamming_dist_top_k + ops.impl("npu_hamming_dist_top_k", &vllm_fl::meta::npu_hamming_dist_top_k_meta); + // reshape_and_cache_bnsd + ops.impl("npu_reshape_and_cache_bnsd", &vllm_fl::meta::npu_reshape_and_cache_bnsd_meta); + // npu_sign_bits_pack + ops.impl("npu_sign_bits_pack", &vllm_fl::meta::npu_sign_bits_pack_meta); + // CopyAndExpandEagleInputs + ops.impl("npu_copy_and_expand_eagle_inputs", &vllm_fl::meta::npu_copy_and_expand_eagle_inputs_meta); + // causal_conv1d_fn + ops.impl("npu_causal_conv1d_custom", &vllm_fl::meta::npu_causal_conv1d_custom_meta); + // moe_grouped_matmul + ops.impl("moe_grouped_matmul", &vllm_fl::meta::moe_grouped_matmul_meta); + ops.impl("moe_gating_top_k_hash", &vllm_fl::meta::moe_gating_top_k_hash_meta); + ops.impl("compressor", &vllm_fl::meta::compressor_meta); + ops.impl("compressor_metadata", &vllm_fl::meta::compressor_metadata_meta); + ops.impl("npu_vllm_quant_lightning_indexer", &vllm_fl::meta::npu_vllm_quant_lightning_indexer_meta); + ops.impl("npu_vllm_quant_lightning_indexer_metadata", &vllm_fl::meta::npu_vllm_quant_lightning_indexer_metadata_meta); + ops.impl("npu_sparse_attn_sharedkv", &vllm_fl::meta::npu_sparse_attn_sharedkv_meta); + ops.impl("npu_sparse_attn_sharedkv_metadata", &vllm_fl::meta::npu_sparse_attn_sharedkv_metadata_meta); + ops.impl("npu_hc_post", &vllm_fl::meta::npu_hc_post_meta); + ops.impl("npu_hc_pre", &vllm_fl::meta::npu_hc_pre_meta); + ops.impl("npu_hc_pre_v2", &vllm_fl::meta::npu_hc_pre_meta); + ops.impl("npu_hc_pre_inv_rms", &vllm_fl::meta::npu_hc_pre_inv_rms_meta); + ops.impl("npu_hc_pre_sinkhorn", &vllm_fl::meta::npu_hc_pre_sinkhorn_meta); + ops.impl("inplace_partial_rotary_mul", &vllm_fl::meta::inplace_partial_rotary_mul_meta); + ops.impl("npu_rms_norm_dynamic_quant", &vllm_fl::meta::npu_rms_norm_dynamic_quant_meta); + ops.impl("indexer_compress_epilog", &vllm_fl::meta::indexer_compress_epilog_meta); + ops.impl("kv_compress_epilog", &vllm_fl::meta::kv_compress_epilog_meta); + ops.impl("npu_kv_quant_sparse_attn_sharedkv", &vllm_fl::meta::npu_kv_quant_sparse_attn_sharedkv_meta); + ops.impl("npu_kv_quant_sparse_attn_sharedkv_metadata", + &vllm_fl::meta::npu_kv_quant_sparse_attn_sharedkv_metadata_meta); + ops.impl("npu_swiglu_group_quant", &vllm_fl::meta::npu_swiglu_group_quant_meta); + ops.impl("npu_load_index_kv_cache", &vllm_fl::meta::npu_load_index_kv_cache_meta); + ops.impl("indexer_compress_epilog_v2", &vllm_fl::meta::indexer_compress_epilog_v2_meta); + ops.impl("npu_dequant_swiglu_quant", &vllm_fl::meta::npu_dequant_swiglu_quant_meta); + ops.impl("npu_scatter_nd_update_v2", &vllm_fl::meta::npu_scatter_nd_update_v2_meta); + // Lightning indexer quant + ops.impl("npu_lightning_indexer_quant", &vllm_fl::meta::npu_lightning_indexer_quant_meta); + // N-gram spec decode + ops.impl("npu_ngram_spec_decode", &vllm_fl::meta::npu_ngram_spec_decode_meta); + // chunk_gated_delta_rule_fwd_h + ops.impl("chunk_gated_delta_rule_fwd_h", &vllm_fl::meta::chunk_gated_delta_rule_fwd_h_meta); + // chunk_fwd_o + ops.impl("chunk_fwd_o", &vllm_fl::meta::chunk_fwd_o_meta); + // store_kv_block + ops.impl("store_kv_block_pre", &vllm_fl::meta::store_kv_block_metadata); + ops.impl("store_kv_block", &vllm_fl::meta::store_kv_block); + // npu_fused_gdn_gating + ops.impl("npu_fused_gdn_gating", &vllm_fl::meta::npu_fused_gdn_gating_meta); +} +} +#endif diff --git a/csrc/ascend/utils.h b/csrc/ascend/utils.h new file mode 100644 index 000000000..a692b87f2 --- /dev/null +++ b/csrc/ascend/utils.h @@ -0,0 +1,55 @@ +#pragma once + +#include "kernels/types.h" +#include +#include + +#define _CONCAT(A, B) A##B +#define CONCAT(A, B) _CONCAT(A, B) + +#define _STRINGIFY(A) #A +#define STRINGIFY(A) _STRINGIFY(A) + +// A version of the TORCH_LIBRARY macro that expands the NAME, i.e. so NAME +// could be a macro instead of a literal token. +#define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE) + +// A version of the TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME +// could be a macro instead of a literal token. +#define TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) \ + TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE) + +// REGISTER_EXTENSION allows the shared library to be loaded and initialized +// via python's import statement. +#define REGISTER_EXTENSION(NAME) \ + PyMODINIT_FUNC CONCAT(PyInit_, NAME)() { \ + static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, \ + STRINGIFY(NAME), nullptr, 0, nullptr}; \ + return PyModule_Create(&module); \ + } + +class TrochBindException : public std::exception +{ +private: + std::string message = {}; + +public: + explicit TrochBindException(const char *name, const char *file, const int line, const std::string &error) + { + message = std::string("Failed: ") + name + " error " + file + ":" + std::to_string(line) + + " error message or error code is '" + error + "'"; + } + + const char *what() const noexcept override + { + return message.c_str(); + } +}; + +#define TORCH_BIND_ASSERT(cond) \ + ; \ + do { \ + if (not(cond)) { \ + throw TrochBindException("Assertion", __FILE__, __LINE__, #cond); \ + } \ + } while (0) diff --git a/csrc/ascend/utils/CMakeLists.txt b/csrc/ascend/utils/CMakeLists.txt new file mode 100644 index 000000000..db468cb2b --- /dev/null +++ b/csrc/ascend/utils/CMakeLists.txt @@ -0,0 +1,48 @@ +# Copyright (c) 2024 Huawei Technologies Co., Ltd. +# This file is a part of the CANN Open Software. +# Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ====================================================================================================================== + +add_library(ops_utils_tiling_headers INTERFACE) + +target_include_directories(ops_utils_tiling_headers INTERFACE + $ + $<$:$> + $<$:$> + $<$:$> + $<$:$> + $ +) + +target_compile_definitions(ops_utils_tiling_headers INTERFACE + OPS_UTILS_LOG_SUB_MOD_NAME="OP_TILING" + OPS_UTILS_LOG_PACKAGE_TYPE=$,"[Custom]",""> +) + +add_library(ops_utils_proto_headers INTERFACE) + +target_include_directories(ops_utils_proto_headers INTERFACE + $ + $<$:$> + $<$:$> + $<$:$> + $ +) + +target_compile_definitions(ops_utils_proto_headers INTERFACE + OPS_UTILS_LOG_SUB_MOD_NAME="OP_PROTO" + OPS_UTILS_LOG_PACKAGE_TYPE=$,"[Custom]",""> +) + +if(NOT BUILD_OPEN_PROJECT) + install_package( + PACKAGE ops_adv + TARGETS ops_utils_tiling_headers ops_utils_proto_headers + DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/inc/ + DESTINATION include/ops_adv/utils + ) +endif() diff --git a/csrc/ascend/utils/inc/aclnn_util.h b/csrc/ascend/utils/inc/aclnn_util.h new file mode 100644 index 000000000..472ea4dbb --- /dev/null +++ b/csrc/ascend/utils/inc/aclnn_util.h @@ -0,0 +1,14 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef OP_API_INC_ACLNN_UTIL_H +#define OP_API_INC_ACLNN_UTIL_H + +#define ACLNN_API __attribute__((visibility("default"))) +#endif // OP_API_INC_ACLNN_UTIL_H \ No newline at end of file diff --git a/csrc/ascend/utils/inc/error/ops_error.h b/csrc/ascend/utils/inc/error/ops_error.h new file mode 100644 index 000000000..fbb5c295c --- /dev/null +++ b/csrc/ascend/utils/inc/error/ops_error.h @@ -0,0 +1,25 @@ +/** + * Copyright (c) 2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file ops_error.h + * \brief + */ + +#pragma once + +#include "log/ops_log.h" + +/* 基础报错 */ +#define OPS_REPORT_VECTOR_INNER_ERR(OPS_DESC, ...) OPS_INNER_ERR_STUB("E89999", OPS_DESC, __VA_ARGS__) +#define OPS_REPORT_CUBE_INNER_ERR(OPS_DESC, ...) OPS_INNER_ERR_STUB("E69999", OPS_DESC, __VA_ARGS__) + +/* 条件报错 */ +#define OPS_ERR_IF(COND, LOG_FUNC, EXPR) OPS_LOG_STUB_IF(COND, LOG_FUNC, EXPR) diff --git a/csrc/ascend/utils/inc/fallback.h b/csrc/ascend/utils/inc/fallback.h new file mode 100644 index 000000000..eb19050d3 --- /dev/null +++ b/csrc/ascend/utils/inc/fallback.h @@ -0,0 +1,497 @@ +/** + * Copyright (c) 2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file fallback.h + * \brief + */ + +#ifndef ACLNNFALLBACK_OPAPI_H_ +#define ACLNNFALLBACK_OPAPI_H_ + +#include + +#include +#include +#include +#include + +#include "aclnn/aclnn_base.h" +#include "fallback_comm.h" +#include "error/ops_error.h" +#include "runtime/base.h" + +namespace fallback { +using namespace std; +using namespace gert; +using namespace ge; +using namespace std; + +namespace std_utils { + template + struct index_sequence {}; + + template + struct make_index_sequence_helper : make_index_sequence_helper {}; + + template + struct make_index_sequence_helper<0, Is...> { + using type = index_sequence; + }; + + template + using make_index_sequence = typename make_index_sequence_helper::type; +} + +using aclOpExecutor = struct aclOpExecutor; +using aclTensor = struct aclTensor; +using aclScalar = struct aclScalar; +using aclIntArray = struct aclIntArray; +using aclFloatArray = struct aclFloatArray; +using aclBoolArray = struct aclBoolArray; +using aclTensorList = struct aclTensorList; + +using _aclCreateTensor = aclTensor* (*)(const int64_t* view_dims, uint64_t view_dims_num, aclDataType data_type, + const int64_t* stride, int64_t offset, aclFormat format, + const int64_t* storage_dims, uint64_t storage_dims_num, void* tensor_data); + +using _aclCreateScalar = aclScalar* (*)(void* value, aclDataType data_type); +using _aclCreateIntArray = aclIntArray* (*)(const int64_t* value, uint64_t size); +using _aclCreateFloatArray = aclFloatArray* (*)(const float* value, uint64_t size); +using _aclCreateBoolArray = aclBoolArray* (*)(const bool* value, uint64_t size); +using _aclCreateTensorList = aclTensorList* (*)(const aclTensor* const* value, uint64_t size); + +using _aclDestroyTensor = int (*)(const aclTensor* tensor); +using _aclDestroyScalar = int (*)(const aclScalar* scalar); +using _aclDestroyIntArray = int (*)(const aclIntArray* array); +using _aclDestroyFloatArray = int (*)(const aclFloatArray* array); +using _aclDestroyBoolArray = int (*)(const aclBoolArray* array); +using _aclDestroyTensorList = int (*)(const aclTensorList* array); + +#define GET_OP_API_FUNC(apiName) reinterpret_cast<_##apiName>(GetOpApiFuncAddr(#apiName)) + +inline const char* GetOpApiLibName(void) { + return "libopapi.so"; +} + +inline const char* GetCustOpApiLibName(void) { + return "libcust_opapi.so"; +} + +inline void* GetOpApiFuncAddrInLib(void* handler, const char* libName, const char* apiName) { + auto funcAddr = dlsym(handler, apiName); + if (funcAddr == nullptr) { + OPS_LOG_W("aclnnfallback", "dlsym %s from %s failed, error:%s.", apiName, libName, dlerror()); + } + return funcAddr; +} + +inline void* GetOpApiLibHandler(const char* libName) { + auto handler = dlopen(libName, RTLD_LAZY); + if (handler == nullptr) { + OPS_LOG_W("aclnnfallback", "dlopen %s failed, error:%s.", libName, dlerror()); + } + return handler; +} + +inline void* GetAclnnArrdByApiName(const char *apiName) { + vector libs = {"libaclnn_ops_infer.so", "libaclnn_ops_train.so", "libaclnn_math.so", + "libaclnn_rand.so", "libaclnn_sparse.so", "libaclnn_fft.so"}; + for (const auto &libName : libs) { + static auto libHandler = GetOpApiLibHandler(libName.c_str()); + if (libHandler != nullptr) { + auto funcAddr = GetOpApiFuncAddrInLib(libHandler, libName.c_str(), apiName); + if (funcAddr != nullptr) { + return funcAddr; + } + } + } + OPS_LOG_E("aclnnfallback", "api %s can't find in any aclnn lib.", apiName); + return nullptr; +} + +inline void* GetOpApiFuncAddr(const char* apiName) { + static auto custOpApiHandler = GetOpApiLibHandler(GetCustOpApiLibName()); + if (custOpApiHandler != nullptr) { + auto funcAddr = GetOpApiFuncAddrInLib(custOpApiHandler, GetCustOpApiLibName(), apiName); + if (funcAddr != nullptr) { + return funcAddr; + } + } + + static auto opApiHandler = GetOpApiLibHandler(GetOpApiLibName()); + if (opApiHandler != nullptr) { + auto funcAddr = GetOpApiFuncAddrInLib(opApiHandler, GetOpApiLibName(), apiName); + if (funcAddr != nullptr) { + return funcAddr; + } + } + OPS_LOG_D("aclnnfallback", "opapi lib is not exist,will use aclnn lib."); + return GetAclnnArrdByApiName(apiName); +} + +inline aclTensor* ConvertType(aclTensor* ge_tensor) { + return ge_tensor; +} + +inline aclIntArray* ConvertType(const std::vector &arr) { + if (arr.empty()) { + return nullptr; + } + static const auto aclCreateIntArray = GET_OP_API_FUNC(aclCreateIntArray); + auto array = aclCreateIntArray(arr.data(), arr.size()); + return array; +} + +inline aclDataType GetConvertType(const gert::Tensor* ge_tensor) { + // convert data type + auto dataType_ge = ge_tensor->GetDataType(); + auto dataType = aclDataType::ACL_FLOAT16; + if (dataType_ge == DT_FLOAT) { + dataType = aclDataType::ACL_FLOAT; + } else if (dataType_ge == DT_BF16) { + dataType = aclDataType::ACL_BF16; + } else if (dataType_ge == DT_BOOL) { + dataType = aclDataType::ACL_BOOL; + } else if (dataType_ge == DT_INT64) { + dataType = aclDataType::ACL_INT64; + } else if (dataType_ge == DT_INT32) { + dataType = aclDataType::ACL_INT32; + } else if (dataType_ge == DT_UINT64) { + dataType = aclDataType::ACL_UINT64; + } else if (dataType_ge == DT_UINT32) { + dataType = aclDataType::ACL_UINT32; + } else if (dataType_ge == DT_INT8) { + dataType = aclDataType::ACL_INT8; + } else if (dataType_ge == DT_UINT8) { + dataType = aclDataType::ACL_UINT8; + } else if (dataType_ge == DT_INT4) { + dataType = aclDataType::ACL_INT4; + } else { + dataType = aclDataType::ACL_FLOAT16; + } + + return dataType; +} + +inline aclTensor* ConvertType(const gert::Tensor* ge_tensor) { + if (ge_tensor == nullptr) { + return nullptr; + } + + static const auto aclCreateTensor = GET_OP_API_FUNC(aclCreateTensor); + OPS_ERR_IF(aclCreateTensor == nullptr, OPS_LOG_E("aclnnfallback", "aclCreateTensor nullptr"), return nullptr); + + void* device_addr = nullptr; + auto tensor_place = ge_tensor->GetPlacement(); + device_addr = const_cast(ge_tensor->GetAddr()); + + auto dataType = GetConvertType(ge_tensor); + + OPS_LOG_D("aclnnfallback", "aclCreateTensor: tensor type is %d", dataType); + + // convert shape + auto gert_shape = ge_tensor->GetStorageShape(); + std::vector shape; + for (size_t i = 0; i < gert_shape.GetDimNum(); ++i) { + shape.push_back(gert_shape.GetDim(i)); + } + + // 计算连续tensor的strides + std::vector strides(shape.size(), 1); + for (int64_t i = shape.size() - 2; i >= 0; i--) { + strides[i] = shape[i + 1] * strides[i + 1]; + } + + aclTensor* out = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), + 0, aclFormat::ACL_FORMAT_ND, + shape.data(), shape.size(), device_addr); + + OPS_ERR_IF(out == nullptr, + OPS_LOG_E("aclnnfallback", "out nullptr"), return nullptr); + + return out; +} + +inline aclTensorList* ConvertType(std::vector& ge_tenserList) { + OPS_ERR_IF(ge_tenserList.size() == 0, + OPS_LOG_E("aclnnfallback", "ge_tenserList size 0"), return nullptr); + + static const auto aclCreateTensorList = GET_OP_API_FUNC(aclCreateTensorList); + OPS_ERR_IF(aclCreateTensorList == nullptr, + OPS_LOG_E("aclnnfallback", "ge_tenserList size 0"), return nullptr); + + std::vector tmp; + for (size_t i = 0; i < ge_tenserList.size(); i++) { + auto t_acl = ConvertType(ge_tenserList[i]); + tmp.push_back(t_acl); + } + + aclTensorList* tensorList = aclCreateTensorList(tmp.data(), tmp.size()); + return tensorList; +} + +template +inline aclScalar* ConvertScalarType(T value) { + static const auto aclCreateScalar = GET_OP_API_FUNC(aclCreateScalar); + OPS_ERR_IF(aclCreateScalar == nullptr, + OPS_LOG_E("aclnnfallback", "aclCreateScalar nullptr"), return nullptr); + if (typeid(value) == typeid(float)) { + return aclCreateScalar(&value, aclDataType::ACL_FLOAT); + } + return nullptr; +} + +template +T ConvertType(T value) { + return value; +} + +inline aclTensor* ConvertMmType(const gert::Tensor* ge_tensor, bool transpose, bool enable_NZ=false) { + if (ge_tensor == nullptr) { + return nullptr; + } + auto gert_shape = ge_tensor->GetStorageShape(); + if (gert_shape.GetDimNum() <= 1) { + return ConvertType(ge_tensor); + } + + static const auto aclCreateTensor = GET_OP_API_FUNC(aclCreateTensor); + OPS_ERR_IF(aclCreateTensor == nullptr, OPS_LOG_E("aclnnfallback", "aclCreateTensor nullptr"), return nullptr); + + void* device_addr = const_cast(ge_tensor->GetAddr()); + // convert data type + auto dataType_ge = ge_tensor->GetDataType(); + auto dataType = ToAclDataType(dataType_ge); + // convert shape + std::vector shape; + for (size_t i = 0; i < gert_shape.GetDimNum(); ++i) { + shape.push_back(gert_shape.GetDim(i)); + } + // 计算连续tensor的strides + std::vector strides(shape.size(), 1); + for (int64_t i = shape.size() - 2; i >= 0; i--) { + strides[i] = shape[i + 1] * strides[i + 1]; + } + + auto viewShape = shape; + // 对于transpose后的tensor对后两维度进行strides, viewShape转换 + if (transpose) { + // dimM 为倒数第二维， dimN 为倒数第一维度 + auto dimM = shape.size() - 2; + auto dimN = shape.size() - 1; + auto swap = strides[dimN]; + strides[dimN] = strides[dimM]; + strides[dimM] = swap; + // 修改viewShape + viewShape[dimN] = shape[dimM]; + viewShape[dimM] = shape[dimN]; + } + auto acl_format = aclFormat::ACL_FORMAT_ND; + if (enable_NZ && GetPrimaryFormat(ge_tensor->GetStorageFormat()) == ge::Format::FORMAT_FRACTAL_NZ) { + acl_format = aclFormat::ACL_FORMAT_FRACTAL_NZ; + } + aclTensor* out = aclCreateTensor(viewShape.data(), shape.size(), dataType, strides.data(), + 0, acl_format, shape.data(), shape.size(), device_addr); + OPS_ERR_IF(out == nullptr, OPS_LOG_E("aclnnfallback", "out nullptr"), return nullptr); + + return out; +} + +inline void Release(aclTensor* p) { + static const auto aclDestroyTensor = GET_OP_API_FUNC(aclDestroyTensor); + OPS_ERR_IF(aclDestroyTensor == nullptr, + OPS_LOG_E("aclnnfallback", "aclDestroyTensor is null"), return); + aclDestroyTensor(p); +} + +inline void Release(aclScalar* p) { + static const auto aclDestroyScalar = GET_OP_API_FUNC(aclDestroyScalar); + OPS_ERR_IF(aclDestroyScalar == nullptr, + OPS_LOG_E("aclnnfallback", "aclDestroyScalar is null"), return); + aclDestroyScalar(p); +} + +inline void Release(aclIntArray* p) { + static const auto aclDestroyIntArray = GET_OP_API_FUNC(aclDestroyIntArray); + OPS_ERR_IF(aclDestroyIntArray == nullptr, + OPS_LOG_E("aclnnfallback", "aclDestroyIntArray is null"), return); + aclDestroyIntArray(p); +} + +inline void Release(aclBoolArray* p) { + static const auto aclDestroyBoolArray = GET_OP_API_FUNC(aclDestroyBoolArray); + OPS_ERR_IF(aclDestroyBoolArray == nullptr, + OPS_LOG_E("aclnnfallback", "aclDestroyBoolArray is null"), return); + aclDestroyBoolArray(p); +} + +inline void Release(aclTensorList* p) { + static const auto aclDestroyTensorList = GET_OP_API_FUNC(aclDestroyTensorList); + OPS_ERR_IF(aclDestroyTensorList == nullptr, + OPS_LOG_E("aclnnfallback", "aclDestroyTensorList is null"), return); + aclDestroyTensorList(p); +} + +template +void Release(T value) { + (void)value; +} + +template +void CallRelease(Tuple t, std_utils::index_sequence) { + (void)std::initializer_list{(Release(std::get(t)), 0)...}; +} + +template +void ReleaseConvertTypes(Tuple& t) { + static constexpr auto size = std::tuple_size::value; + CallRelease(t, std_utils::make_index_sequence{}); +} + +template +auto ConvertTypes(Ts&... args) -> decltype(std::make_tuple(ConvertType(args)...)) { + auto tp = std::make_tuple(ConvertType(args)...); + return tp; +} + +template +auto call(Function f, Tuple t, std_utils::index_sequence) -> int { + return f(std::get(t)...); +} + +template +auto call(Function f, Tuple t) -> int { + static constexpr auto size = std::tuple_size::value; + return call(f, t, std_utils::make_index_sequence{}); +} + +template +auto ConvertToOpApiFunc(const Tuple& params, void* opApiAddr, std_utils::index_sequence) + -> int (*)(typename std::decay(params))>::type...) { + using OpApiFunc = int (*)(typename std::decay(params))>::type...); + auto func = reinterpret_cast(opApiAddr); + return func; +} + +template +auto ConvertToOpApiFunc(const Tuple& params, void* opApiAddr) + -> typename std::enable_if::value != 0, + decltype(ConvertToOpApiFunc(params, opApiAddr, std_utils::make_index_sequence::value>{}))>::type { + static constexpr auto size = std::tuple_size::value; + return ConvertToOpApiFunc(params, opApiAddr, std_utils::make_index_sequence{}); +} + +template +class ConvertedParams { + public: + ConvertedParams(Tuple&& convertedParams) : convertedParams_(std::move(convertedParams)){}; + ConvertedParams(ConvertedParams&& other) : convertedParams_(std::move(other.convertedParams_)) { + other.validParams_ = false; + }; + ConvertedParams& operator=(ConvertedParams&& other) { + if (this == &other) { + return *this; + } + + convertedParams_ = std::move(other.convertedParams_); + validParams_ = true; + other.validParams_ = false; + return *this; + } + + ConvertedParams() = delete; + ConvertedParams(const ConvertedParams& other) = delete; + ConvertedParams& operator=(const ConvertedParams& other) = delete; + + ~ConvertedParams() { + if (validParams_) { + ReleaseConvertTypes(convertedParams_); + } + } + + const Tuple& GetConvertedParams() const { + return convertedParams_; + } + + private: + Tuple convertedParams_; + bool validParams_{true}; +}; + +using InitHugeMemThreadLocal = int (*)(void*, bool); +using UnInitHugeMemThreadLocal = void (*)(void*, bool); +using ReleaseHugeMem = void (*)(void*, bool); +using PTAGetExecCache = aclOpExecutor* (*)(uint64_t, uint64_t*); +using InitPTACacheThreadLocal = void (*)(); +using SetPTAHashKey = void (*)(uint64_t); +using CanUsePTACache = bool (*)(const char*); + +using ResetCacheThreadLocal = void (*)(); + +#define EXEC_OPAPI_CMD(aclnn_api, ...) \ + ({ \ + static auto ret = GRAPH_SUCCESS; \ + do { \ + static const auto ResetCacheThreadLocalAddr = GetOpApiFuncAddr("ResetCacheThreadLocal"); \ + static const auto getWorkspaceSizeFuncAddr = GetOpApiFuncAddr(#aclnn_api "GetWorkspaceSize"); \ + static const auto opApiFuncAddr = GetOpApiFuncAddr(#aclnn_api); \ + if (getWorkspaceSizeFuncAddr == nullptr || opApiFuncAddr == nullptr || ResetCacheThreadLocalAddr == nullptr) { \ + OPS_LOG_E("aclnnfallback", "%s or %s not in %s or %s or ResetCacheThreadLocal not found.", \ + #aclnn_api "GetWorkspaceSize", #aclnn_api, GetOpApiLibName(), GetOpApiLibName()); \ + ret = GRAPH_FAILED; \ + break; \ + } \ + auto ResetCacheThreadLocalFunc = reinterpret_cast(ResetCacheThreadLocalAddr); \ + ResetCacheThreadLocalFunc(); \ + uint64_t workspace_size = 0; \ + uint64_t* workspace_size_addr = &workspace_size; \ + aclOpExecutor* executor = nullptr; \ + aclOpExecutor** executor_addr = &executor; \ + auto converted_params = ConvertTypes(__VA_ARGS__, workspace_size_addr, executor_addr); \ + static auto getWorkspaceSizeFunc = ConvertToOpApiFunc(converted_params, getWorkspaceSizeFuncAddr); \ + auto workspace_status = call(getWorkspaceSizeFunc, converted_params); \ + if (workspace_status != 0) { \ + OPS_LOG_E("aclnnfallback", "call %s failed:", #aclnn_api); \ + ret = GRAPH_FAILED; \ + break; \ + } \ + void* workspace_addr = nullptr; \ + if (workspace_size > 0) { \ + workspace_addr = host_api_ctx->MallocWorkspace(workspace_size); \ + if (workspace_addr == nullptr) { \ + OPS_LOG_E("aclnnfallback", "call %s allocate workspace failed", #aclnn_api); \ + ret = GRAPH_FAILED; \ + break; \ + } \ + } \ + auto acl_stream = host_api_ctx->GetStream(); \ + auto acl_call = [converted_params, workspace_addr, workspace_size, host_api_ctx, acl_stream, \ + executor]() -> int { \ + using OpApiFunc = int (*)(void*, uint64_t, aclOpExecutor*, const aclrtStream); \ + OpApiFunc opApiFunc = reinterpret_cast(opApiFuncAddr); \ + auto api_ret = opApiFunc(workspace_addr, workspace_size, executor, acl_stream); \ + ReleaseConvertTypes(converted_params); \ + host_api_ctx->FreeWorkspace(); \ + if (api_ret != 0) { \ + OPS_LOG_E("aclnnfallback", "call %s allocate workspace failed api_ret: %d", #aclnn_api, api_ret); \ + return GRAPH_FAILED; \ + } \ + return api_ret; \ + }; \ + \ + ret = acl_call(); \ + } while (false); \ + (ret); \ + }) + +} // namespace fallback + +#endif // ACLNNFALLBACK_OPAPI_H_ diff --git a/csrc/ascend/utils/inc/fallback_comm.h b/csrc/ascend/utils/inc/fallback_comm.h new file mode 100644 index 000000000..a2dd5cfd1 --- /dev/null +++ b/csrc/ascend/utils/inc/fallback_comm.h @@ -0,0 +1,38 @@ +/** + * Copyright (c) 2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file fallback_comm.h + * \brief + */ + +#ifndef INC_EXTERNAL_GRAPH_FALLBACK_COMMON_H_ +#define INC_EXTERNAL_GRAPH_FALLBACK_COMMON_H_ + +#include "aclnn/aclnn_base.h" +#include "exe_graph/runtime/op_execute_context.h" +#include "exe_graph/runtime/tensor.h" +#include "register/op_impl_registry.h" +#include "runtime/base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +namespace fallback { + +aclDataType ToAclDataType(ge::DataType dtype); +} // namespace fallback + +#ifdef __cplusplus +} +#endif + +#endif // INC_EXTERNAL_GRAPH_FALLBACK_COMMON_H_ diff --git a/csrc/ascend/utils/inc/kernel/comm_args.h b/csrc/ascend/utils/inc/kernel/comm_args.h new file mode 100644 index 000000000..3aadb840e --- /dev/null +++ b/csrc/ascend/utils/inc/kernel/comm_args.h @@ -0,0 +1,72 @@ +#ifndef COMM_ARGS_H +#define COMM_ARGS_H +#include + +#define FORCE_INLINE_AICORE __attribute__((always_inline)) inline __aicore__ +#include "kernel_operator.h" + +namespace Moe { +constexpr int CAM_MAX_RANK_SIZE = 384; // Maximum number of NPU cards supported by the communication library + +constexpr int64_t IPC_BUFF_MAX_SIZE = 100 * 1024 * 1024; +constexpr int64_t IPC_DATA_OFFSET = 2 * 1024 * 1024; // First 2MB as flag, then 100MB as data storage +constexpr int64_t PING_PONG_SIZE = 2; +constexpr int64_t UB_SINGLE_DMA_SIZE_MAX = 190 * 1024; +constexpr int64_t SMALL_DATA_SIZE = 1 * 1024 * 1024; +constexpr int64_t UB_SINGLE_PING_PONG_ADD_SIZE_MAX = UB_SINGLE_DMA_SIZE_MAX / 2; +constexpr int UB_ALIGN_SIZE = 32; +constexpr int64_t MAGIC_ALIGN_COUNT = UB_ALIGN_SIZE / sizeof(int32_t); + +constexpr uint8_t COMM_NUM = 2; // Size of communication domain +constexpr uint8_t COMM_EP_IDX = 0; +constexpr uint8_t COMM_TP_IDX = 1; + +constexpr int DFX_COUNT = 50; +constexpr int64_t WAIT_SUCCESS = 112233445566; +constexpr int64_t IPC_CHUNK_FLAG = 0; // Start offset for send recv, chunk flag region +constexpr int64_t MAX_WAIT_ROUND_UNIT = 10 * 1000 * 1000; // Threshold for waiting to get Flag under normal conditions within the same SIO + +constexpr static int32_t UB_HEAD_OFFSET = 96; +constexpr static int32_t UB_MID_OFFSET = UB_HEAD_OFFSET + UB_SINGLE_PING_PONG_ADD_SIZE_MAX + UB_ALIGN_SIZE; +constexpr static int64_t UB_FLAG_SIZE = 2 * 1024; +constexpr static int64_t MAX_CORE_NUM = 48; +constexpr static uint64_t STATE_WIN_OFFSET = 900 * 1024; +constexpr static int64_t COMPARE_ALIGN_SIZE = 256; + +constexpr static int64_t UB_SINGLE_TOTAL_SIZE_MAX = 192 * 1024; +constexpr static int64_t START_OFFSET_FOR_SHARE = 512; + +enum Op : int { + COPYONLY = -1, + ADD = 0, + MUL = 1, + MAX = 2, + MIN = 3 +}; + +struct CommArgs { + int rank = 0; // attr rank_id, global rank + int localRank = -1; + int rankSize = 0; // global rank size + int localRankSize = -1; // This parameter refers to the number of cards interconnected in fullmesh + uint32_t extraFlag = 0; // 32 bit map, the specific meaning of each bit is above in this file + int testFlag = 0; + GM_ADDR peerMems[CAM_MAX_RANK_SIZE] = {}; // Buffer obtained from initialization, all allreduce is the same parameter + /** + * @param sendCountMatrix One-dimensional array with a size of rankSize*rankSize + * eg: The value of sendCountMatrix[1] corresponds to the [0][1] of the two-dimensional array, indicating the number of data that card 0 needs to send to card 1 + */ + int64_t sendCountMatrix[CAM_MAX_RANK_SIZE * CAM_MAX_RANK_SIZE] = {}; // for all2allvc + int64_t sendCounts[CAM_MAX_RANK_SIZE] = {}; // for all2allv + int64_t sdispls[CAM_MAX_RANK_SIZE] = {}; // for all2allv + int64_t recvCounts[CAM_MAX_RANK_SIZE] = {}; // for all2allv + int64_t rdispls[CAM_MAX_RANK_SIZE] = {}; // for all2allv + int64_t batchSize; + int64_t hiddenSize; + int64_t topk; + int64_t sharedExpertRankNum; + int64_t expertNumPerRank; + int64_t dfx[DFX_COUNT] = {}; +}; +} +#endif // COMM_ARGS_H diff --git a/csrc/ascend/utils/inc/kernel/data_copy.h b/csrc/ascend/utils/inc/kernel/data_copy.h new file mode 100644 index 000000000..d9490e1ca --- /dev/null +++ b/csrc/ascend/utils/inc/kernel/data_copy.h @@ -0,0 +1,68 @@ +#ifndef CAM_DATACOPY_GM2GM_H +#define CAM_DATACOPY_GM2GM_H +#include +#include "comm_args.h" + +using namespace AscendC; +using namespace Moe; + +template +FORCE_INLINE_AICORE void SetAtomicOpType(int op) +{ + switch (op) { + case ADD: + AscendC::SetAtomicAdd(); + break; + case MUL: + // Ignore setting the atomic register when performing mul + break; + case MAX: + AscendC::SetAtomicMax(); + break; + case MIN: + AscendC::SetAtomicMin(); + break; + default: + AscendC::SetAtomicNone(); + } +} + +template +FORCE_INLINE_AICORE void CpUB2GM(__gm__ T *gmAddr, __ubuf__ T *ubAddr, uint32_t size) +{ + LocalTensor ubTensor; + GlobalTensor gmTensor; + DataCopyExtParams dataCopyParams(1, size, 0, 0, 0); + ubTensor.address_.logicPos = static_cast(TPosition::VECIN); + ubTensor.address_.bufferAddr = reinterpret_cast(ubAddr); + gmTensor.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(gmAddr)); + DataCopyPad(gmTensor, ubTensor, dataCopyParams); +} + +template +FORCE_INLINE_AICORE void CpGM2UB(__ubuf__ T *ubAddr, __gm__ T *gmAddr, uint32_t size) +{ + LocalTensor ubTensor; + GlobalTensor gmTensor; + DataCopyExtParams dataCopyParams(1, size, 0, 0, 0); + ubTensor.address_.logicPos = static_cast(TPosition::VECIN); + ubTensor.address_.bufferAddr = reinterpret_cast(ubAddr); + gmTensor.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(gmAddr)); + DataCopyPadExtParams padParams; + DataCopyPad(ubTensor, gmTensor, dataCopyParams, padParams); +} + +template +FORCE_INLINE_AICORE void CopyUB2UB(__ubuf__ T *dst, __ubuf__ T *src, const uint32_t calCount) +{ + LocalTensor srcTensor; + LocalTensor dstTensor; + TBuffAddr srcAddr, dstAddr; + srcAddr.bufferAddr = reinterpret_cast(src); + dstAddr.bufferAddr = reinterpret_cast(dst); + srcTensor.SetAddr(srcAddr); + dstTensor.SetAddr(dstAddr); + DataCopy(dstTensor, srcTensor, calCount); +} + +#endif // CAM_DATACOPY_GM2GM_H \ No newline at end of file diff --git a/csrc/ascend/utils/inc/kernel/dropmask.h b/csrc/ascend/utils/inc/kernel/dropmask.h new file mode 100644 index 000000000..13ed9c350 --- /dev/null +++ b/csrc/ascend/utils/inc/kernel/dropmask.h @@ -0,0 +1,121 @@ +/** + * Copyright (c) 2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dropmask.h + * \brief + */ + +#ifndef DROPMASK_H +#define DROPMASK_H + +#include "util.h" + +using AscendC::DROPOUT_MODE_BIT_MISALIGN; +using AscendC::DropOutShapeInfo; +using AscendC::DropOut; + +struct DropMaskInfo { + // for compute dropout mask offset + // 参数按B N G S1 S2全部切分设置进行偏移计算，没有切分的轴对应的参数设置为合适的0或者原始值 + int64_t n2G; // n2 * g + int64_t gSize; // g + int64_t s1Size; // s1 + int64_t s2Size; // s2 + int64_t gOutIdx; // g out index + int64_t bSSOffset; // boidx * s1 * s2 ===bSSOffset + int64_t n2OutIdx; // n out index + int64_t s1OutIdx; // s1 out index ===s1oIdx + int64_t s1InnerIdx; // s1 inner index, 配比 ===loopIdx + int64_t s1BaseSize; // S1基本块大小 + int64_t splitS1BaseSize; // s1 split size ===vec1S1BaseSize + int64_t s2StartIdx; // s2 start index + int64_t s2Idx; // s2 index =====s2LoopCount + int64_t s2BaseNratioSize; // s2的配比长度: s2BaseSize(S2基本块大小) * nRatio + + // for copy in dropout mask + uint32_t s1CopySize; + uint32_t s2CopySize; + int64_t s2TotalSize; + + // for compute dropout mask + uint32_t firstAxis; + uint32_t lstAxis; + uint32_t maskLstAxis; + int64_t vecCoreOffset = 0; + float keepProb; + + bool boolMode; +}; + +template +__aicore__ inline int64_t ComputeDropOffset(DropMaskInfo &dropMaskInfo) +{ + if constexpr (hasDrop == true) { + // boidx * n2 * g* s1 * s2 + int64_t bOffset = dropMaskInfo.bSSOffset * dropMaskInfo.n2G; + // n2oIdx * g * s1 *s2 + int64_t n2Offset = dropMaskInfo.n2OutIdx * dropMaskInfo.gSize * dropMaskInfo.s1Size * dropMaskInfo.s2Size; + // goIdx * s1 * s2 + int64_t gOffset = dropMaskInfo.gOutIdx * dropMaskInfo.s1Size * dropMaskInfo.s2Size; + // s1oIdx * s1BaseSize * s2Size + s1innerindex * vec1S1BaseSize * s2Size + int64_t s1Offset = (dropMaskInfo.s1OutIdx * dropMaskInfo.s1BaseSize + dropMaskInfo.vecCoreOffset + + dropMaskInfo.s1InnerIdx * dropMaskInfo.splitS1BaseSize) * dropMaskInfo.s2Size; + // s2StartIdx + s2index * s2BaseNratioSize + int64_t s2Offset = dropMaskInfo.s2StartIdx + dropMaskInfo.s2Idx * dropMaskInfo.s2BaseNratioSize; + return bOffset + n2Offset + gOffset + s1Offset + s2Offset; + } else { + return 0; + } +} + +template +__aicore__ inline void CopyInDropMask(LocalTensor&dstTensor, GlobalTensor& srcBoolTensor, + GlobalTensor& srcByteTensor, DropMaskInfo &dropMaskInfo, int64_t alignedSize = blockBytes) +{ + if constexpr (hasDrop == true) { + int64_t dropMaskOffset = ComputeDropOffset(dropMaskInfo); + if (unlikely(dropMaskInfo.boolMode)) { + BoolCopyIn(dstTensor, srcBoolTensor, dropMaskOffset, + dropMaskInfo.s1CopySize, dropMaskInfo.s2CopySize, dropMaskInfo.s2TotalSize, alignedSize); + } else { + Bit2Int8CopyIn(dstTensor, srcByteTensor, dropMaskOffset, 1, + dropMaskInfo.s1CopySize, dropMaskInfo.s2CopySize, dropMaskInfo.s2TotalSize, alignedSize); + } + return; + } +} + +template +__aicore__ inline void ComputeDropMask(LocalTensor& dstTensor, LocalTensor& srcTensor, + LocalTensor& dropoutBuffer, LocalTensor& tmpDropBuffer, DropMaskInfo &dropMaskInfo) +{ + if constexpr (hasDrop == true) { + DropOutShapeInfo dropOutShapeInfo; + dropOutShapeInfo.firstAxis = dropMaskInfo.firstAxis; + dropOutShapeInfo.srcLastAxis = dropMaskInfo.lstAxis; + + if (unlikely(dropMaskInfo.boolMode)) { + dropOutShapeInfo.maskLastAxis = CeilDiv(dropMaskInfo.maskLstAxis, blockBytes) * blockBytes; + DropOut(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, dropMaskInfo.keepProb, dropOutShapeInfo); + } else { + dropOutShapeInfo.maskLastAxis = CeilDiv(dropMaskInfo.maskLstAxis / byteBitRatio, blockBytes) * blockBytes; + if (likely(dropMaskInfo.lstAxis / byteBitRatio % blockBytes == 0)) { + DropOut(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, dropMaskInfo.keepProb, dropOutShapeInfo); + } else { + DropOut(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, + dropMaskInfo.keepProb, dropOutShapeInfo); + } + } + return; + } +} + +#endif // DROPMASK_H diff --git a/csrc/ascend/utils/inc/kernel/moe_distribute_base.h b/csrc/ascend/utils/inc/kernel/moe_distribute_base.h new file mode 100644 index 000000000..b74df9a55 --- /dev/null +++ b/csrc/ascend/utils/inc/kernel/moe_distribute_base.h @@ -0,0 +1,288 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file moe_distribute_base.h + * \brief + */ + +#ifndef MOE_DISTRIBUTE_BASE_H +#define MOE_DISTRIBUTE_BASE_H + +constexpr uint32_t LOCAL_NOTIFY_MAX_NUM = 64; +constexpr uint32_t LOCAL_STREAM_MAX_NUM = 19; +constexpr uint32_t AICPU_OP_NOTIFY_MAX_NUM = 2; +constexpr uint32_t AICPU_MAX_RANK_NUM = 128 * 1024; + +struct HcclSignalInfo { + uint64_t resId; + uint64_t addr; + uint32_t devId; + uint32_t tsId; + uint32_t rankId; + uint32_t flag; +}; + +struct ListCommon { + uint64_t nextHost; + uint64_t preHost; + uint64_t nextDevice; + uint64_t preDevice; +}; + +struct HcclStreamInfo { + int32_t streamIds; + uint32_t sqIds; + uint32_t cqIds; + uint32_t logicCqids; +}; + +struct LocalResInfoV2 { + uint32_t streamNum; + uint32_t signalNum; + HcclSignalInfo localSignals[LOCAL_NOTIFY_MAX_NUM]; + HcclStreamInfo streamInfo[LOCAL_STREAM_MAX_NUM]; + HcclStreamInfo mainStreamInfo; + HcclSignalInfo aicpuOpNotify[AICPU_OP_NOTIFY_MAX_NUM]; + ListCommon nextTagRes; // HccltagLocalResV2 +}; + +enum class rtFloatOverflowMode_t { + RT_OVERFLOW_MODE_SATURATION = 0, + RT_OVERFLOW_MODE_INFNAN, + RT_OVERFLOW_MODE_UNDEF, +}; + +struct AlgoTopoInfo { + uint32_t userRank; + uint32_t userRankSize; + int32_t deviceLogicId; + bool isSingleMeshAggregation; + uint32_t deviceNumPerAggregation; + uint32_t superPodNum; + uint32_t devicePhyId; + uint32_t topoType; // TopoType + uint32_t deviceType; + uint32_t serverNum; + uint32_t meshAggregationRankSize; + uint32_t multiModuleDiffDeviceNumMode; + uint32_t multiSuperPodDiffServerNumMode; + uint32_t realUserRank; + bool isDiffDeviceModule; + bool isDiffDeviceType; + uint32_t gcdDeviceNumPerAggregation; + uint32_t moduleNum; + uint32_t isUsedRdmaRankPairNum; + uint64_t isUsedRdmaRankPair; + uint32_t pairLinkCounterNum; + uint64_t pairLinkCounter; + uint32_t nicNum; + uint64_t nicList; + uint64_t complanRankLength; + uint64_t complanRank; + uint64_t bridgeRankNum; + uint64_t bridgeRank; + uint64_t serverAndsuperPodRankLength; + uint64_t serverAndsuperPodRank; +}; + +struct HcclOpConfig { + uint8_t deterministic; + uint8_t retryEnable; + uint8_t highPerfEnable; + uint8_t padding[5]; + uint8_t linkTimeOut[8]; + uint64_t notifyWaitTime; + uint32_t retryHoldTime; + uint32_t retryIntervalTime; + bool interHccsDisable = false; + rtFloatOverflowMode_t floatOverflowMode = rtFloatOverflowMode_t::RT_OVERFLOW_MODE_UNDEF; + uint32_t multiQpThreshold = 512; +}; + +struct HcclMC2WorkSpace { + uint64_t workSpace; + uint64_t workSpaceSize; +}; + +struct RemoteResPtr { + uint64_t nextHostPtr; + uint64_t nextDevicePtr; +}; + +struct HDCommunicateParams { + uint64_t hostAddr { 0 }; + uint64_t deviceAddr { 0 }; + uint64_t readCacheAddr { 0 }; + uint32_t devMemSize{ 0 }; + uint32_t buffLen{ 0 }; + uint32_t flag{ 0 }; +}; + +struct HcclRankRelationResV2 { + uint32_t remoteUsrRankId; + uint32_t remoteWorldRank; + uint64_t windowsIn; + uint64_t windowsOut; + uint64_t windowsExp; + ListCommon nextTagRes; +}; + +struct HcclOpResParam { + HcclMC2WorkSpace mc2WorkSpace; + uint32_t localUsrRankId; // usrrankid + uint32_t rankSize; + uint64_t winSize; + uint64_t localWindowsIn; + uint64_t localWindowsOut; + char hcomId[128]; + uint64_t winExpSize; + uint64_t localWindowsExp; + uint32_t rWinStart; + uint32_t rWinOffset; + uint64_t version; + LocalResInfoV2 localRes; + AlgoTopoInfo topoInfo; + + HcclOpConfig config; + uint64_t hostStateInfo; + uint64_t aicpuStateInfo; + uint64_t lockAddr; + uint32_t rsv[16]; + uint32_t notifysize; + uint32_t remoteResNum; + RemoteResPtr remoteRes[AICPU_MAX_RANK_NUM]; + + // communicate retry + HDCommunicateParams kfcControlTransferH2DParams; + HDCommunicateParams kfcStatusTransferD2HParams; + uint64_t tinyMem; // for all2all + uint64_t tinyMemSize; + uint64_t zeroCopyHeadPtr; + uint64_t zeroCopyTailPtr; + uint64_t zeroCopyRingBuffer; + uint64_t zeroCopyIpcPtrs[16]; + uint32_t zeroCopyDevicePhyId[16]; + + bool utraceStatusFlag; +}; + +// Transport +enum class HcclAiRMAMemType : uint32_t { + LOCAL_INPUT = 0, + REMOTE_INPUT, + LOCAL_OUTPUT, + REMOTE_OUTPUT, + MAX_NUM +}; + +struct HcclAiRMAMemInfo { + uint32_t memMaxNum{0}; + uint32_t sizeOfMemDetails{0}; + uint64_t memDetailPtr{0}; +}; + +// Transport QP/Mem +struct HcclAiRMAInfo { + uint32_t curRankId{0}; + uint32_t rankNum{0}; + uint32_t qpNum{0}; + uint32_t sizeOfAiRMAWQ{0}; // sizeof(HcclAiRMAWQ) + uint32_t sizeOfAiRMACQ{0}; // sizeof(HcclAiRMACQ) + uint32_t sizeOfAiRMAMem{0}; // sizeof(HcclAiRMAMemInfo) + uint64_t sqPtr{0}; + uint64_t scqPtr{0}; + uint64_t rqPtr{0}; + uint64_t rcqPtr{0}; + uint64_t memPtr{0}; +}; + +struct HcclA2CombineOpParam { + uint64_t workSpace; // Address for communication between client and server, + // hccl requests and clears + uint64_t workSpaceSize; // Space for communication between client and server + uint32_t rankId; // id of this rank + uint32_t rankNum; // num of ranks in this comm group + uint64_t winSize; // size of each windows memory + uint64_t windowsIn[AscendC::HCCL_MAX_RANK_NUM]; // windows address for input, windowsIn[rankId] corresponds + // to the local card address, + // and others are cross-card mapping addresses. + uint64_t windowsOut[AscendC::HCCL_MAX_RANK_NUM]; // windows address for output, windowsOut[rankId] corresponds + // to the local card address, + // and others are cross-card mapping addresses. + uint8_t res[8328]; + uint8_t multiFlag; + __gm__ AscendC::IbVerbsData *data; + uint64_t dataSize; + + uint64_t sizeOfAiRMAInfo; + uint64_t aiRMAInfo; +}; + +enum class DBMode : int32_t { + INVALID_DB = -1, + HW_DB = 0, + SW_DB +}; + +struct HcclAiRMAWQ { + uint32_t wqn{0}; + uint64_t bufAddr{0}; + uint32_t wqeSize{0}; + uint32_t depth{0}; + uint64_t headAddr{0}; + uint64_t tailAddr{0}; + DBMode dbMode{DBMode::INVALID_DB}; // 0-hw/1-sw + uint64_t dbAddr{0}; + uint32_t sl{0}; +}; + +struct HcclAiRMACQ { + uint32_t cqn{0}; + uint64_t bufAddr{0}; + uint32_t cqeSize{0}; + uint32_t depth{0}; + uint64_t headAddr{0}; + uint64_t tailAddr{0}; + DBMode dbMode{DBMode::INVALID_DB}; // 0-hw/1-sw + uint64_t dbAddr{0}; +}; + +struct hns_roce_rc_sq_wqe { + uint32_t byte_4; + uint32_t msg_len; + uint32_t immtdata; + uint32_t byte_16; + uint32_t byte_20; + uint32_t rkey; + uint64_t remoteVA; +}; + + +struct hns_roce_lite_wqe_data_seg { + uint32_t len; + uint32_t lkey; + uint64_t localVA; +}; + +__aicore__ inline void cacheWriteThrough(__gm__ uint8_t* sourceAddr, uint64_t length) { + __gm__ uint8_t* start = + (__gm__ uint8_t*)((uint64_t)sourceAddr / AscendC::CACHE_LINE_SIZE * AscendC::CACHE_LINE_SIZE); + __gm__ uint8_t* end = + (__gm__ uint8_t*)(((uint64_t)sourceAddr + length) / AscendC::CACHE_LINE_SIZE * AscendC::CACHE_LINE_SIZE); + AscendC::GlobalTensor global; + global.SetGlobalBuffer(start); + for (uint32_t i = 0; i <= end - start; i += AscendC::CACHE_LINE_SIZE) { + AscendC::DataCacheCleanAndInvalid(global[i]); + } +} + +#endif // MOE_DISTRIBUTE_BASE_H \ No newline at end of file diff --git a/csrc/ascend/utils/inc/kernel/pse.h b/csrc/ascend/utils/inc/kernel/pse.h new file mode 100644 index 000000000..582ad3e53 --- /dev/null +++ b/csrc/ascend/utils/inc/kernel/pse.h @@ -0,0 +1,483 @@ +/** + * Copyright (c) 2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file pse.h + * \brief + */ + +#ifndef FLASH_ATTENTION_SCORE_PSE_H +#define FLASH_ATTENTION_SCORE_PSE_H + +#include "kernel_operator.h" +#include "util.h" + +constexpr static int64_t pseS1S2 = 0; +constexpr static int64_t pse1S2 = 1; +constexpr static int64_t pseSlopeBn = 2; +constexpr static int64_t pseSlopeN = 3; + +constexpr static uint8_t pseEncodeALibiS2Full = 0x11; + +enum class PseTypeEnum { + PSE_OUTER_MUL_ADD_TYPE = 0, // default + PSE_OUTER_ADD_MUL_TYPE, + PSE_INNER_MUL_ADD_TYPE, + PSE_INNER_MUL_ADD_SQRT_TYPE, + PSE_INVALID_TYPE +}; + +struct PseInfo { + int64_t blockCount; + int64_t bSSOffset; // boidx * s1 * s2 + int64_t boIdx; + int64_t gSize; + int64_t goIdx; + int64_t loopIdx; + int64_t n2G; + int64_t n2oIdx; + int64_t pseBSize; + int64_t pseS1Size; // for alibi + int64_t pseS2ComputeSize; // for alibi, do not need assignment + int64_t pseS2Size; // for alibi + uint32_t pseShapeType; + int64_t readS2Size; // for alibi, do not need assignment + int64_t s1BaseSize; + int64_t s1Size; + int64_t s1oIdx; + int64_t s2AlignedSize; + int64_t s2BaseNratioSize; + int64_t s2LoopCount; + int64_t s2RealSize; + int64_t s2Size; + int64_t s2SizeAcc; // accumulated sum of s2 size + int64_t s2StartIdx; + int64_t vec1S1BaseSize; + int64_t vec1S1RealSize; + uint32_t pseEncodeType; // for distinguish alibi + uint32_t pseType; // 0: outer, mul-add 1:outer, add-mul 2:inner, mul-add 3:inner, mul-add-sqrt + int64_t pseAlibiBaseS1; + int64_t pseAlibiBaseS2; + int64_t qStartIdx; + int64_t kvStartIdx; + int64_t vecCoreOffset = 0; + bool needCast; + bool align8 = false; + bool pseEndogenous = false; +}; + +template +__aicore__ inline void DataCopyInCommon(LocalTensor &dstTensor, GlobalTensor &srcTensor, int64_t offset, + int64_t s1Size, int64_t s2Size, int64_t actualS2Len, int32_t dtypeSize, + int32_t alignedS2Size) +{ + if constexpr (hasPse == true) { + uint32_t shapeArray[] = {static_cast(s1Size), static_cast(alignedS2Size)}; + dstTensor.SetShapeInfo(ShapeInfo(2, shapeArray, DataFormat::ND)); + dstTensor.SetSize(s1Size * alignedS2Size); + DataCopyParams dataCopyParams; + dataCopyParams.blockCount = s1Size; + dataCopyParams.blockLen = CeilDiv(s2Size * dtypeSize, blockBytes); // 单位32B + dataCopyParams.dstStride = alignedS2Size * dtypeSize / blockBytes - dataCopyParams.blockLen; // gap + if (actualS2Len * dtypeSize % blockBytes == 0) { + dataCopyParams.srcStride = + (actualS2Len * dtypeSize - dataCopyParams.blockLen * blockBytes) / blockBytes; // srcGap + DataCopy(dstTensor, srcTensor[offset], dataCopyParams); + } else { + dataCopyParams.blockLen = s2Size * dtypeSize; // 单位Byte + dataCopyParams.srcStride = (actualS2Len * dtypeSize - dataCopyParams.blockLen); + dataCopyParams.dstStride = (alignedS2Size - s2Size) * dtypeSize / blockBytes; + DataCopyPadParams dataCopyPadParams; + dataCopyPadParams.isPad = false; + DataCopyPad(dstTensor, srcTensor[offset], dataCopyParams, dataCopyPadParams); + } + } +} + +template +__aicore__ inline void DataCopyIn(LocalTensor &dstTensor, GlobalTensor &srcTensor, int64_t offset, + int64_t s1Size, int64_t s2Size, int64_t actualS2Len, int64_t alignedSize = 16) +{ + if constexpr (hasPse == true) { + int32_t dtypeSize = sizeof(INPUT_T); + int32_t alignedS2Size = CeilDiv(s2Size, alignedSize) * alignedSize; + DataCopyInCommon(dstTensor, srcTensor, offset, s1Size, s2Size, + actualS2Len, dtypeSize, alignedS2Size); + } +} + +template +__aicore__ inline void DataCopyInAlign8(LocalTensor &dstTensor, GlobalTensor &srcTensor, int64_t offset, + int64_t s1Size, int64_t s2Size, int64_t actualS2Len) +{ + if constexpr (hasPse == true) { + int32_t dtypeSize = sizeof(INPUT_T); + if (dtypeSize == 0){ + return; + } + int32_t alignedS2Size = CeilDiv(s2Size, 32 / dtypeSize) * (32 / dtypeSize); + DataCopyInCommon(dstTensor, srcTensor, offset, s1Size, s2Size, + actualS2Len, dtypeSize, alignedS2Size); + } +} + +/* +dst = BroadcastAdd(src0, src1) +src0 shape: (s1, s2) +src1 shape: (1, s2) +dst shape: (s1, s2) +*/ +template +__aicore__ inline void BroadcastAdd(const LocalTensor &src0Tensor, const LocalTensor &src1Tensor, + int64_t src0Offset, int32_t src1Size, int32_t repeatTimes) +{ + if constexpr (hasPse == true) { + /* Total data number of single step should be smaller than 256bytes. + * If larger, we need to do add multiple times. */ + int32_t innerLoop = src1Size / repeatMaxSize; // s2轴整块计算次数 + int32_t innerRemain = src1Size % repeatMaxSize; // s2轴尾块计算量 + BinaryRepeatParams binaryRepeatParams; + binaryRepeatParams.src0BlkStride = 1; + binaryRepeatParams.src0RepStride = src1Size / blockSize; + binaryRepeatParams.src1BlkStride = 1; + binaryRepeatParams.src1RepStride = 0; + binaryRepeatParams.dstRepStride = binaryRepeatParams.src0RepStride; + binaryRepeatParams.blockNumber = binaryRepeatParams.src0RepStride; + + for (int32_t j = 0; j < innerLoop; j++) { + auto innerOffset = j * repeatMaxSize; + auto ubOffset = src0Offset + innerOffset; + Add(src0Tensor[ubOffset], src0Tensor[ubOffset], src1Tensor[innerOffset], repeatMaxSize, repeatTimes, + binaryRepeatParams); + } + if (innerRemain > 0) { + auto innerOffset = innerLoop * repeatMaxSize; + auto ubOffset = src0Offset + innerOffset; + Add(src0Tensor[ubOffset], src0Tensor[ubOffset], src1Tensor[innerOffset], innerRemain, repeatTimes, + binaryRepeatParams); + } + } +} + +template +__aicore__ inline void PseBroadcastAdd(int32_t s1Size, int32_t s2Size, int32_t computeSize, const LocalTensor &pseUb, + const LocalTensor &dstTensor, uint32_t pseShapeType) +{ + if constexpr (hasPse == true) { + if (pseShapeType == pseS1S2 || pseShapeType == pseSlopeBn || pseShapeType == pseSlopeN) { + Add(dstTensor, dstTensor, pseUb, computeSize); + } else { + /* Total repeated times should be <= repeatMaxTimes. If larger, + * we need to do multiple inner loops. */ + int32_t s1OuterLoop = s1Size / repeatMaxTimes; + int32_t s1OuterRemain = s1Size % repeatMaxTimes; + for (int32_t s1OuterIdx = 0; s1OuterIdx < s1OuterLoop; s1OuterIdx++) { + int32_t s1OuterOffset = s1OuterIdx * repeatMaxTimes * s2Size; + BroadcastAdd(dstTensor, pseUb, s1OuterOffset, s2Size, repeatMaxTimes); + } + if (s1OuterRemain > 0) { + int32_t s1OuterOffset = s1OuterLoop * repeatMaxTimes * s2Size; + BroadcastAdd(dstTensor, pseUb, s1OuterOffset, s2Size, s1OuterRemain); + } + } + } +} +template __aicore__ inline int64_t PseComputeOffset(PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + int64_t bOffset = 0; + int64_t n2Offset = 0; + int64_t s1Offset = 0; + int64_t s2Offset = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize; + int64_t gOffset = 0; + if (pseInfo.pseShapeType == pseS1S2) { + // b, n2, g, s1, s2 + bOffset = pseInfo.bSSOffset * pseInfo.n2G; + n2Offset = pseInfo.n2oIdx * pseInfo.gSize * pseInfo.s1Size * pseInfo.s2Size; + gOffset = pseInfo.goIdx * pseInfo.s1Size * pseInfo.s2Size; + s1Offset = (pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + pseInfo.loopIdx * pseInfo.vec1S1BaseSize) * pseInfo.s2Size; + } else if (pseInfo.pseShapeType == pse1S2) { + // b, n2, g, 1, s2 + bOffset = pseInfo.s2SizeAcc * pseInfo.n2G; + n2Offset = pseInfo.n2oIdx * pseInfo.gSize * pseInfo.s2Size; + gOffset = pseInfo.goIdx * pseInfo.s2Size; + } + if (pseInfo.pseBSize == 1) { + bOffset = 0; + } + return bOffset + n2Offset + gOffset + s1Offset + s2Offset; + } else { + return 0; + } +} + +template __aicore__ inline int64_t PseAlibiComputeOffset(PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + int64_t bOffset = (pseInfo.boIdx % pseInfo.pseBSize) * pseInfo.n2G * pseInfo.pseS2Size * pseInfo.pseS1Size; + int64_t n2Offset = pseInfo.n2oIdx * pseInfo.gSize * pseInfo.pseS2Size * pseInfo.pseS1Size; + int64_t gOffset = pseInfo.goIdx * pseInfo.pseS2Size * pseInfo.pseS1Size; + int64_t row = pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + pseInfo.loopIdx * pseInfo.vec1S1BaseSize; + int64_t column = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize; + int64_t m = 0; + int64_t k = 0; + if constexpr (layOutType != LayOutTypeEnum::LAYOUT_TND) { + int64_t threshold = pseInfo.s1Size - pseInfo.pseS1Size; + if (row >= threshold) { + m = row - threshold; + k = column; + } else { + m = row % pseInfo.pseS1Size; + k = pseInfo.pseS2Size - (row - column) - (pseInfo.pseS1Size - m); + } + } else { + int64_t threshold = pseInfo.pseS2Size - pseInfo.pseS1Size; + int64_t posVal = row - column - threshold; + if (threshold >= 0) { + if (posVal >= 0) { + m = posVal; + k = 0; + } else { + m = 0; + k = -posVal; + } + } else { + m = posVal; + k = 0; + } + } + int64_t s1Offset = m * pseInfo.pseS2Size; + int64_t s2Offset = k; + pseInfo.readS2Size = Min(pseInfo.s2AlignedSize, pseInfo.pseS2Size - k); + pseInfo.pseS2ComputeSize = Align(pseInfo.readS2Size); + + return bOffset + n2Offset + gOffset + s1Offset + s2Offset; + } else { + return 0; + } +} + +template __aicore__ inline bool NeedPseAlibiCompute(PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + // Alibi编码只计算下三角 + if (pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + (pseInfo.loopIdx + 1) * pseInfo.vec1S1BaseSize <= + pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize) { + return false; + } + return true; + } else { + return false; + } +} + +template +__aicore__ inline void PseAlibiCopyIn(LocalTensor &dstTensor, LocalTensor &tmpTensor, + GlobalTensor &srcTensor, PseInfo &pseInfo, int64_t alignedSize = 16) +{ + if constexpr (hasPse == true) { + if (!NeedPseAlibiCompute(pseInfo)) { + return; + } + int64_t offset = PseAlibiComputeOffset(pseInfo); + if constexpr (IsSameType::value) { + if (!pseInfo.align8){ + DataCopyIn(dstTensor, srcTensor, offset, pseInfo.vec1S1RealSize, pseInfo.readS2Size, + pseInfo.pseS2Size, alignedSize); + } else { + DataCopyInAlign8(dstTensor, srcTensor, offset, pseInfo.vec1S1RealSize, + pseInfo.readS2Size, pseInfo.pseS2Size); + } + return; + } + + DataCopyIn(tmpTensor, srcTensor, offset, pseInfo.vec1S1RealSize, pseInfo.readS2Size, + pseInfo.pseS2Size, alignedSize); + if (pseInfo.needCast) { + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(dstTensor, tmpTensor, RoundMode::CAST_NONE, pseInfo.vec1S1RealSize * pseInfo.pseS2ComputeSize); + } + return; + } +} + +template +__aicore__ inline void PseSlopeCopyIn(LocalTensor &dstTensor, LocalTensor &helpTensor, + __gm__ uint8_t *pseSlope, GlobalTensor &alibiGm, PseInfo &pseInfo, + int64_t alignedSize = 16) { + if constexpr (hasPse == true) { + int64_t bOffset = 0; + int64_t n2Offset = pseInfo.n2oIdx * pseInfo.gSize; + int64_t gOffset = pseInfo.goIdx; + + if (pseInfo.pseShapeType == pseSlopeBn) { + bOffset = pseInfo.boIdx * pseInfo.n2G; + } + int64_t offset = bOffset + n2Offset + gOffset; + + DataCopyIn(helpTensor, alibiGm, 0, pseInfo.vec1S1RealSize, + pseInfo.s2RealSize, pseInfo.pseAlibiBaseS2, alignedSize); + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + + if (pseInfo.needCast) { + int64_t computeSize = pseInfo.vec1S1RealSize * pseInfo.s2AlignedSize; + Cast(dstTensor, helpTensor, RoundMode::CAST_NONE, computeSize); + AscendC::PipeBarrier(); + + int64_t s1Offset = pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + pseInfo.loopIdx * pseInfo.vec1S1BaseSize; + int64_t s2Offset = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize; + + float posShift = float(s2Offset + pseInfo.kvStartIdx - s1Offset - pseInfo.qStartIdx); + + Adds(dstTensor, dstTensor, posShift, computeSize); + AscendC::PipeBarrier(); + Abs(dstTensor, dstTensor, computeSize); + AscendC::PipeBarrier(); + float slopes = ((__gm__ T *)pseSlope)[offset] * -1; + if (pseInfo.pseType == (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_SQRT_TYPE) { + Sqrt(dstTensor, dstTensor, computeSize); + AscendC::PipeBarrier(); + } + Muls(dstTensor, dstTensor, slopes, computeSize); + AscendC::PipeBarrier(); + } + } +} + +template +__aicore__ inline void PseSlopeCast(LocalTensor &dstTensor, LocalTensor &helpTensor, + __gm__ uint8_t *pseSlope, PseInfo &pseInfo) { + if constexpr (hasPse == true) { + int64_t bOffset = 0; + int64_t n2Offset = pseInfo.n2oIdx * pseInfo.gSize; + int64_t gOffset = pseInfo.goIdx; + + if (pseInfo.pseShapeType == pseSlopeBn) { + bOffset = pseInfo.boIdx * pseInfo.n2G; + } + int64_t offset = bOffset + n2Offset + gOffset; + int64_t computeSize = pseInfo.vec1S1RealSize * pseInfo.s2AlignedSize; + Cast(dstTensor, helpTensor, RoundMode::CAST_NONE, computeSize); + AscendC::PipeBarrier(); + + int64_t s1Offset = pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset + + pseInfo.loopIdx * pseInfo.vec1S1BaseSize; + int64_t s2Offset = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize; + + float posShift = float(s2Offset + pseInfo.kvStartIdx - s1Offset - pseInfo.qStartIdx); + + Adds(dstTensor, dstTensor, posShift, computeSize); + AscendC::PipeBarrier(); + Abs(dstTensor, dstTensor, computeSize); + AscendC::PipeBarrier(); + float slopes = ((__gm__ T *)pseSlope)[offset] * -1; + if (pseInfo.pseType == (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_SQRT_TYPE) { + Sqrt(dstTensor, dstTensor, computeSize); + AscendC::PipeBarrier(); + } + Muls(dstTensor, dstTensor, slopes, computeSize); + AscendC::PipeBarrier(); + } +} + +template +__aicore__ inline void PseCopyIn(LocalTensor &dstTensor, LocalTensor &tmpTensor, + GlobalTensor &srcTensor, PseInfo &pseInfo, int64_t alignedSize = 16) +{ + if constexpr (hasPse == true) { + if (pseInfo.pseEncodeType == pseEncodeALibiS2Full) { + return PseAlibiCopyIn(dstTensor, tmpTensor, srcTensor, pseInfo, alignedSize); + } + int64_t offset = PseComputeOffset(pseInfo); + int64_t s1Size = pseInfo.pseShapeType == pse1S2 ? (pseInfo.blockCount == 0 ? 1 : pseInfo.blockCount) : + pseInfo.vec1S1RealSize; + + if constexpr (IsSameType::value) { + if (!pseInfo.align8){ + DataCopyIn(dstTensor, srcTensor, offset, s1Size, pseInfo.s2RealSize, + pseInfo.s2Size, alignedSize); + } else { + DataCopyInAlign8(dstTensor, srcTensor, offset, s1Size, pseInfo.s2RealSize, pseInfo.s2Size); + } + return; + } + DataCopyIn(tmpTensor, srcTensor, offset, s1Size, pseInfo.s2RealSize, pseInfo.s2Size, + alignedSize); + if (pseInfo.needCast) { + event_t eventIdMte2ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eventIdMte2ToV); + WaitFlag(eventIdMte2ToV); + Cast(dstTensor, tmpTensor, RoundMode::CAST_NONE, s1Size * pseInfo.s2AlignedSize); + } + return; + } +} + +template +__aicore__ inline void PseAlibiCompute(LocalTensor &dstTensor, LocalTensor &pseTensor, PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + if (!NeedPseAlibiCompute(pseInfo)) { + return; + } + Add(dstTensor, dstTensor, pseTensor, pseInfo.vec1S1RealSize * pseInfo.pseS2ComputeSize); + return; + } +} + +template +__aicore__ inline void PseCompute(LocalTensor &dstTensor, LocalTensor &pseTensor, PseInfo &pseInfo) +{ + if constexpr (hasPse == true) { + if (pseInfo.pseEncodeType == pseEncodeALibiS2Full) { + return PseAlibiCompute(dstTensor, pseTensor, pseInfo); + } + int64_t computeSize = (pseInfo.pseShapeType == pseS1S2 || pseInfo.pseShapeType == pseSlopeBn || + pseInfo.pseShapeType == pseSlopeN) + ? pseInfo.vec1S1RealSize * pseInfo.s2AlignedSize + : pseInfo.s2AlignedSize; + PseBroadcastAdd(pseInfo.vec1S1RealSize, pseInfo.s2AlignedSize, computeSize, pseTensor, + dstTensor, pseInfo.pseShapeType); + return; + } +} + +template +__aicore__ inline void PseInnerAlibiCreate(GlobalTensor &dstTensor, LocalTensor &helpTensor, PseInfo &pseInfo) { + if constexpr (hasPse == true) { + if (pseInfo.pseType != (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_TYPE && pseInfo.pseType != (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_SQRT_TYPE) { + return; + } + event_t eventIdMte3ToV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V)); + event_t eventIdMte3ToS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S)); + event_t eventIdVToMte3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + float tmpValue = -1.0; + + for (int64_t i = 0; i < pseInfo.pseAlibiBaseS1; i++) { + CreateVecIndex(helpTensor, (half)(i * tmpValue), pseInfo.pseAlibiBaseS2); + SetFlag(eventIdVToMte3); + WaitFlag(eventIdVToMte3); + DataCopy(dstTensor[i * pseInfo.pseAlibiBaseS2], helpTensor, pseInfo.pseAlibiBaseS2); + SetFlag(eventIdMte3ToV); + WaitFlag(eventIdMte3ToV); + SetFlag(eventIdMte3ToS); + WaitFlag(eventIdMte3ToS); + } + } +} +#endif diff --git a/csrc/ascend/utils/inc/kernel/sync_collectives.h b/csrc/ascend/utils/inc/kernel/sync_collectives.h new file mode 100644 index 000000000..9653e21a8 --- /dev/null +++ b/csrc/ascend/utils/inc/kernel/sync_collectives.h @@ -0,0 +1,426 @@ +#ifndef SYNC_COLLECTIVES_H +#define SYNC_COLLECTIVES_H + +#include "comm_args.h" + +using namespace AscendC; +using namespace Moe; + +// Synchronization flag occupies length +constexpr int64_t FLAG_UNIT_INT_NUM = 4; +// Memory size occupied by each synchronization unit (Bytes) +constexpr int64_t SYNC_UNIT_SIZE = FLAG_UNIT_INT_NUM * sizeof(int64_t); +// High-order offset when using magic as a comparison value +constexpr int64_t MAGIC_OFFSET = 32; +constexpr int64_t MAGIC_MASK = ~((1LL << MAGIC_OFFSET) - 1); + +class SyncCollectives { +public: + __aicore__ inline SyncCollectives() {} + + __aicore__ inline void Init(int rank, int rankSize, GM_ADDR *shareAddrs, TBuf &tBuf) + { + this->rank = rank; + this->rankSize = rankSize; + this->shareAddrs = shareAddrs; + this->blockIdx = GetBlockIdx(); + this->blockNum = GetBlockNum(); + // Length of a single indicator segment + segmentCount = GetBlockNum() * FLAG_UNIT_INT_NUM; + // Initialize the intra-card/inter-card synchronization address corresponding to the current core. + localSyncAddr = (__gm__ int64_t*)(shareAddrs[rank]); + basicSyncAddr = (__gm__ int64_t*)(shareAddrs[rank]) + GetBlockIdx() * FLAG_UNIT_INT_NUM; + blockOuterSyncAddr = (__gm__ int64_t*)(shareAddrs[rank]) + segmentCount + GetBlockIdx() * FLAG_UNIT_INT_NUM; + this->tBuf = tBuf; + } + + __aicore__ inline void SetSyncFlag(int32_t magic, int32_t value, int32_t eventID) + { + int64_t v = MergeMagicWithValue(magic, value); + SetFlag(localSyncAddr + eventID * FLAG_UNIT_INT_NUM, v); + } + + /** + * @brief Set the flag for the specified eventID of the designated card, with the value being a combination of magic and value. + * @param magic The operator batch, which will be combined into the high 32 bits of the flag value to be set. + * @param value The specific value to be set, which will be the low 32 bits of the flag value to be set. + * @param eventID Physically, it is an offset from the shared memory base address (requires scaling, not an absolute value). + * @param rank This rank is the rankId corresponding to the peerMems array in the CommArgs structure, not a global or local id. + * (Local is not applicable in the 91093 scenario, and global is not applicable in the 910B multi-machine scenario.) + */ + __aicore__ inline void SetSyncFlag(int32_t magic, int32_t value, int32_t eventID, int32_t rank) + { + int64_t v = MergeMagicWithValue(magic, value); + SetFlag((__gm__ int64_t*)(shareAddrs[rank]) + eventID * FLAG_UNIT_INT_NUM, v); + } + + __aicore__ inline int32_t CalEventIdByMulBlockNum(int32_t blockMultiplier, int32_t targetCoreId) + { + return (blockMultiplier * blockNum) + targetCoreId; + } + + /** + * @brief Wait for the flag of the specified eventID on the specified card to become a value + * composed of the combination of magic and value. + * @param magic The operator batch, which will be combined into the high 32 bits of the flag + * value to be wait. + * @param value The specific value to be wait, which will be the low 32 bits of the flag + * value to be wait. + * @param eventID Physically, it is an offset from the shared memory base address (requires + * scaling, not an absolute value). + * @param rank This rank is the rankId corresponding to the peerMems array in the CommArgs + * structure, not a global or local id. (Local is not applicable in the 91093 + * scenario, and global is not applicable in the 910B multi-machine scenario.) + */ + __aicore__ inline void WaitSyncFlag(int32_t magic, int32_t value, int32_t eventID, int32_t rank) + { + int64_t v = MergeMagicWithValue(magic, value); + WaitOneRankPartFlag((__gm__ int64_t*)(shareAddrs[rank]) + eventID * FLAG_UNIT_INT_NUM, 1, v); + } + + __aicore__ inline void WaitSyncFlag(int32_t magic, int32_t value, int32_t eventID) + { + int64_t v = MergeMagicWithValue(magic, value); + WaitOneRankPartFlag((__gm__ int64_t*)(shareAddrs[this->rank]) + eventID * FLAG_UNIT_INT_NUM, 1, v); + } + + /** + * @brief Wait for the flags starting from the specified eventID on the specified card to become + * a value composed of the combination of magic and value.
+ * Note: [eventID, eventID + flagNum) + */ + __aicore__ inline void WaitSyncFlag(int32_t magic, int32_t value, int32_t eventID, int32_t rank, int64_t flagNum) + { + int64_t v = MergeMagicWithValue(magic, value); + WaitOneRankPartFlag((__gm__ int64_t*)(shareAddrs[rank]) + eventID * FLAG_UNIT_INT_NUM, flagNum, v); + } + + // Set inner-card synchronization flag (memory A) + __aicore__ inline void SetInnerFlag(int32_t magic, int32_t eventID) + { + int64_t value = MergeMagicWithValue(magic, eventID); + SetFlag(basicSyncAddr, value); + } + + __aicore__ inline void SetInnerFlag(int32_t magic, int32_t eventID, int64_t setRank, int64_t setBlock) + { + int64_t value = MergeMagicWithValue(magic, eventID); + SetFlag((__gm__ int64_t*)(shareAddrs[setRank]) + setBlock * FLAG_UNIT_INT_NUM, value); + } + + // Wait for a single inner-card synchronization flag (memory A) + __aicore__ inline void WaitInnerFlag(int32_t magic, int32_t eventID, int64_t waitRank, int64_t waitBlock) + { + int64_t value = MergeMagicWithValue(magic, eventID); + WaitOneRankPartFlag((__gm__ int64_t*)(shareAddrs[waitRank]) + waitBlock * FLAG_UNIT_INT_NUM, 1, value); + } + + // Wait for all inner-card synchronization flags within the entire rank (memory A) + __aicore__ inline void WaitRankInnerFlag(int32_t magic, int32_t eventID, int64_t waitRank) + { + int64_t value = MergeMagicWithValue(magic, eventID); + WaitOneRankAllFlag((__gm__ int64_t*)(shareAddrs[waitRank]), value); + } + + // Check all inner-card synchronization flags within the entire rank (memory A) + __aicore__ inline bool CheckRankInnerFlag(int32_t magic, int32_t eventID, int64_t waitRank) + { + int64_t value = MergeMagicWithValue(magic, eventID); + return CheckOneRankAllFlag((__gm__ int64_t*)(shareAddrs[waitRank]), value); + } + + // Set inter-card synchronization flag (memory B) + __aicore__ inline void SetOuterFlag(int32_t magic, int32_t eventID) + { + int64_t value = MergeMagicWithValue(magic, eventID); + SetFlag(blockOuterSyncAddr, value); + } + + __aicore__ inline void SetOuterFlag(int32_t magic, int32_t eventID, int64_t setRank, int64_t setBlock) + { + __gm__ int64_t* flagAddr = GetOuterFlagAddr(setRank, setBlock); + int64_t value = MergeMagicWithValue(magic, eventID); + SetFlag(flagAddr, value); + } + + // Wait for a single inter-card synchronization flag (memory B) + __aicore__ inline void WaitOuterFlag(int32_t magic, int32_t eventID, int64_t waitRank, int64_t waitBlock) + { + int64_t value = MergeMagicWithValue(magic, eventID); + __gm__ int64_t* flagAddr = GetOuterFlagAddr(waitRank, waitBlock); + WaitOneRankPartFlag(flagAddr, 1, value); + } + + // Wait for all inter-card synchronization flags within the entire rank (memory B) + __aicore__ inline void WaitOneRankOuterFlag(int32_t magic, int32_t eventID, int64_t rank) + { + int64_t value = MergeMagicWithValue(magic, eventID); + __gm__ int64_t* flagAddr; + flagAddr = GetOuterFlagAddr(rank, 0); + WaitOneRankPartFlag(flagAddr, blockNum, value); + } + + // Wait for flagNum inter-card synchronization flags starting from startBlock for all ranks (memory B) + __aicore__ inline void WaitAllRankPartOuterFlag(int32_t magic, int32_t eventID, int64_t startBlock, int64_t flagNum) + { + int64_t value = MergeMagicWithValue(magic, eventID); + __gm__ int64_t* flagAddr; + int waitRank; + for (auto r = 0; r < rankSize; ++r) { + waitRank = (rank + r) % rankSize; // Offset reading of rank flags to prevent performance impact from concurrent copying by multiple cores + flagAddr = GetOuterFlagAddr(waitRank, startBlock); + WaitOneRankPartFlag(flagAddr, flagNum, value); + } + } + + // Check flagNum inter-card synchronization flags starting from startBlock for all ranks (memory B) + __aicore__ inline bool CheckAllRankPartOuterFlag(int32_t magic, int32_t eventID, int64_t startBlock, + int64_t flagNum) + { + int64_t value = MergeMagicWithValue(magic, eventID); + __gm__ int64_t* flagAddr; + int waitRank; + for (auto r = 0; r < rankSize; ++r) { + waitRank = (rank + r) % rankSize; // Offset reading of rank flags to prevent performance impact from concurrent copying by multiple cores + flagAddr = GetOuterFlagAddr(waitRank, startBlock); + if (!CheckOneRankPartFlag(flagAddr, flagNum, value)) { + return false; + } + } + return true; + } + + // Wait for all inter-card synchronization flags for all ranks, full rank synchronization (memory B) + __aicore__ inline void WaitAllRankOuterFlag(int32_t magic, int32_t eventID) + { + WaitAllRankPartOuterFlag(magic, eventID, 0, blockNum); + } + + // Check all inter-card synchronization flags for all ranks, full rank synchronization (memory B) + __aicore__ inline bool CheckAllRankOuterFlag(int32_t magic, int32_t eventID) + { + return CheckAllRankPartOuterFlag(magic, eventID, 0, blockNum); + } + + // Low-level interface, set synchronization flag + __aicore__ inline void SetFlag(__gm__ int64_t* setAddr, int64_t setValue) + { + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + GlobalTensor globalSet; + globalSet.SetGlobalBuffer(setAddr, FLAG_UNIT_INT_NUM); + LocalTensor localSet = tBuf.GetWithOffset(1, 0); + localSet.SetValue(0, setValue); + + // Copy global synchronization flag to local + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); // Wait for SetValue to complete + DataCopy(globalSet, localSet, FLAG_UNIT_INT_NUM); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); // Wait for UB->GM to complete + } + + // Low-level interface, wait for synchronization flag + __aicore__ inline void WaitFlag(__gm__ int64_t* waitAddr, int64_t waitValue) + { + WaitOneRankPartFlag(waitAddr, 1, waitValue); + } + + // Read a flag, return an immediate number + __aicore__ inline int64_t GetFlag(__gm__ int64_t* waitAddr) + { + GlobalTensor globalWait; + globalWait.SetGlobalBuffer(waitAddr, FLAG_UNIT_INT_NUM); + LocalTensor localWait = tBuf.GetWithOffset(1, 0); + // Copy global to local + DataCopy(localWait, globalWait, FLAG_UNIT_INT_NUM); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); // Wait for GM->UB + + int64_t res = localWait.GetValue(0); + return res; + } + + // Get multiple consecutive synchronization flags within a single card + __aicore__ inline void WaitOneRankPartOuterFlag(int32_t magic, int32_t eventID, int64_t waitRank, + int64_t startBlock, int64_t flagNum) + { + int64_t value = MergeMagicWithValue(magic, eventID); + __gm__ int64_t* flagAddr; + flagAddr = GetOuterFlagAddr(waitRank, startBlock); + WaitOneRankPartFlag(flagAddr, flagNum, value); + } + + // Get synchronization flag within a single card (memory A) + __aicore__ inline int64_t GetInnerFlag(int64_t waitRank, int64_t waitBlock) + { + return GetFlag((__gm__ int64_t*)(shareAddrs[waitRank]) + waitBlock * FLAG_UNIT_INT_NUM); + } + + __aicore__ inline int64_t GetOuterFlag(int64_t waitRank, int64_t waitBlock) + { + return GetFlag((__gm__ int64_t*)(shareAddrs[waitRank]) + segmentCount + waitBlock * FLAG_UNIT_INT_NUM); + } + + // In the rank Chunk Flag area, return success if the destRank chunk Flag value is 0, otherwise fail + __aicore__ inline int64_t GetChunkFlag(int64_t rank, int64_t destRank, int64_t magic, int64_t timeout) + { + int64_t value = MergeMagicWithValue(magic, 0); + int64_t status = GetChunkFlagValue((__gm__ int64_t*)(shareAddrs[rank]) + + IPC_CHUNK_FLAG + destRank * FLAG_UNIT_INT_NUM, value, timeout); + return status; + } + + // Set the destRank chunk Flag value in the rank Chunk Flag area to value + __aicore__ inline void SetChunkFlag(int64_t rank, int64_t destRank, int64_t magic, int64_t eventId) + { + int64_t value = MergeMagicWithValue(magic, eventId); + SetFlag((__gm__ int64_t*)(shareAddrs[rank]) + IPC_CHUNK_FLAG + destRank * FLAG_UNIT_INT_NUM, value); + } + + __aicore__ inline int64_t GetChunkRecvLen(int64_t rank, int64_t destRank, int64_t magic, int64_t timeout) + { + int64_t len = GetChunkFlagValue((__gm__ int64_t*)(shareAddrs[rank]) + IPC_CHUNK_FLAG + + destRank * FLAG_UNIT_INT_NUM, 0, timeout, true, magic); + return len; + } + +private: + __aicore__ inline int64_t MergeMagicWithValue(int32_t magic, int32_t value) + { + // Merge magic as the high bits and eventID as the low bits into a value for comparison + return (static_cast(static_cast(magic)) << MAGIC_OFFSET) | static_cast(value); + } + + __aicore__ inline __gm__ int64_t* GetInnerFlagAddr(int64_t flagRank, int64_t flagBlock) + { + return (__gm__ int64_t*)(shareAddrs[flagRank]) + flagBlock * FLAG_UNIT_INT_NUM; + } + + __aicore__ inline __gm__ int64_t* GetOuterFlagAddr(int64_t flagRank, int64_t flagBlock) + { + return (__gm__ int64_t*)(shareAddrs[flagRank]) + segmentCount + flagBlock * FLAG_UNIT_INT_NUM; + } + + // Wait for a part of synchronization flags within a rank + __aicore__ inline void WaitOneRankPartFlag(__gm__ int64_t* waitAddr, int64_t flagNum, int64_t checkValue) + { + GlobalTensor globalWait; + globalWait.SetGlobalBuffer(waitAddr, flagNum * FLAG_UNIT_INT_NUM); + LocalTensor localWait = tBuf.GetWithOffset(flagNum * FLAG_UNIT_INT_NUM, 0); + bool isSync = true; + int64_t checkedFlagNum = 0; + do { + // Copy global synchronization flags to local + DataCopy(localWait, globalWait[checkedFlagNum * FLAG_UNIT_INT_NUM], + (flagNum - checkedFlagNum) * FLAG_UNIT_INT_NUM); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); // Wait for GM->UB + + // Check if the synchronization flags are equal to checkValue + isSync = true; + int64_t remainToCheck = flagNum - checkedFlagNum; + for (auto i = 0; i < remainToCheck; ++i) { + // Continue waiting if any core has not reached the checkValue phase + int64_t v = localWait.GetValue(i * FLAG_UNIT_INT_NUM); + if ((v & MAGIC_MASK) != (checkValue & MAGIC_MASK) || v < checkValue) { + isSync = false; + checkedFlagNum += i; + break; + } + } + } while (!isSync); + } + + // Wait for all synchronization flags within a rank + __aicore__ inline void WaitOneRankAllFlag(__gm__ int64_t* waitAddr, int64_t checkValue) + { + WaitOneRankPartFlag(waitAddr, blockNum, checkValue); + } + + // Check partial synchronization flags within a rank, copy only once + __aicore__ inline bool CheckOneRankPartFlag(__gm__ int64_t* waitAddr, int64_t flagNum, int64_t checkValue) + { + GlobalTensor globalWait; + globalWait.SetGlobalBuffer(waitAddr, flagNum * FLAG_UNIT_INT_NUM); + LocalTensor localWait = tBuf.GetWithOffset(flagNum * FLAG_UNIT_INT_NUM, 0); + // Copy global synchronization flags to local + DataCopy(localWait, globalWait, flagNum * FLAG_UNIT_INT_NUM); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); // Wait for GM->UB + // Check if the synchronization flags are equal to checkValue + bool isSync = true; + for (auto i = 0; i < flagNum; ++i) { + // Continue waiting if any core has not reached the checkValue phase + int64_t v = localWait.GetValue(i * FLAG_UNIT_INT_NUM); + if ((v & MAGIC_MASK) != (checkValue & MAGIC_MASK) || v < checkValue) { + isSync = false; + break; + } + } + return isSync; + } + + __aicore__ inline int64_t GetChunkFlagValue(__gm__ int64_t* waitAddr, int64_t checkValue, int64_t timeout, + bool checkNonZero = false, int64_t magic = 0) + { + GlobalTensor globalWait; + globalWait.SetGlobalBuffer(waitAddr, FLAG_UNIT_INT_NUM); + LocalTensor localWait = tBuf.GetWithOffset(FLAG_UNIT_INT_NUM, 0); + bool isSync = true; + + int64_t waitTimes = 0; + int64_t v = 0; + + do { + // Copy global sync flag to local + DataCopy(localWait, globalWait[0], FLAG_UNIT_INT_NUM); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); // Wait for GM->UB + + isSync = true; + v = localWait.GetValue(0); + if (checkNonZero) { + // Non-zero check mode + if (((v & MAGIC_MASK) == (static_cast(magic) << MAGIC_OFFSET)) && (v & 0xFFFFFFFF)) { + return v & 0xFFFFFFFF; // Return lower 32 bits when non-zero + } + } else { + // Exact value check mode + if (v == checkValue) { + return WAIT_SUCCESS; + } + } + + isSync = false; + waitTimes++; + + if (timeout > INT64_MAX / MAX_WAIT_ROUND_UNIT || waitTimes >= (timeout * MAX_WAIT_ROUND_UNIT)) { + isSync = true; + return v; // Return the read flag value + } + } while (!isSync); + + return checkNonZero ? 0 : v; + } + + // Check all sync flags within a rank, copy only once + __aicore__ inline bool CheckOneRankAllFlag(__gm__ int64_t* waitAddr, int64_t checkValue) + { + return CheckOneRankPartFlag(waitAddr, blockNum, checkValue); + } + int rank; + int rankSize; + int blockIdx; + int blockNum; + GM_ADDR *shareAddrs; + int64_t segmentCount; // Length of a single sync flag segment (count in int64_t) + __gm__ int64_t* localSyncAddr; + __gm__ int64_t* basicSyncAddr; // Intra-card sync flag address for the current block + __gm__ int64_t* blockOuterSyncAddr; // Inter-card sync flag address for the current block + TBuf tBuf; +}; + +#endif // SYNC_COLLECTIVES_H \ No newline at end of file diff --git a/csrc/ascend/utils/inc/kernel/util.h b/csrc/ascend/utils/inc/kernel/util.h new file mode 100644 index 000000000..2c7d20893 --- /dev/null +++ b/csrc/ascend/utils/inc/kernel/util.h @@ -0,0 +1,144 @@ +/** + * Copyright (c) 2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file util.h + * \brief + */ + +#ifndef FLASH_ATTENTION_UTIL_H +#define FLASH_ATTENTION_UTIL_H + +constexpr int32_t blockBytes = 32; +constexpr int32_t byteBitRatio = 8; +constexpr int64_t prefixAttenMaskDownHeight = 1024; +constexpr static int32_t blockSize = blockBytes / 4; // 4 means sizeof(T) +constexpr static int32_t repeatMaxBytes = 256; +constexpr static int32_t repeatMaxTimes = 255; +constexpr static int32_t repeatMaxSize = repeatMaxBytes / 4; // 4 means sizeof(T) + +using AscendC::LocalTensor; +using AscendC::GlobalTensor; +using AscendC::DataFormat; +using AscendC::ShapeInfo; +using AscendC::DataCopyParams; +using AscendC::DataCopyPadParams; +using AscendC::BinaryRepeatParams; +using AscendC::IsSameType; +using AscendC::HardEvent; +using AscendC::SetFlag; +using AscendC::WaitFlag; + +enum class LayOutTypeEnum { None = 0, LAYOUT_BSH = 1, LAYOUT_SBH = 2, LAYOUT_BNSD = 3, LAYOUT_TND = 4, LAYOUT_NTD_TND = 5}; + +namespace math { +template __aicore__ inline T Ceil(T a, T b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +template __aicore__ inline T Align(T a, T b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b * b; +} +} + +template +__aicore__ inline T1 CeilDiv(T1 a, T2 b) +{ + if (b == 0) { + return 0; + } + return (a + b - 1) / b; +} + +template +__aicore__ inline T1 Max(T1 a, T2 b) +{ + return (a > b) ? (a) : (b); +} + +template +__aicore__ inline T1 Min(T1 a, T2 b) +{ + return (a > b) ? (b) : (a); +} + +__aicore__ inline void BoolCopyIn(LocalTensor &dstTensor, GlobalTensor &srcTensor, + int64_t srcOffset, uint32_t s1Size, uint32_t s2Size, int64_t totalS2Size, int64_t alignedSize = blockBytes) +{ + uint32_t alignedS2Size = CeilDiv(s2Size, alignedSize) * alignedSize; + uint32_t shapeArray[] = {s1Size, alignedS2Size}; + dstTensor.SetShapeInfo(ShapeInfo(2, shapeArray, DataFormat::ND)); + dstTensor.SetSize(s1Size * alignedS2Size); + DataCopyParams dataCopyParams; + dataCopyParams.blockCount = s1Size; + dataCopyParams.dstStride = 0; + if (totalS2Size == blockBytes && alignedSize == 64) { // totalS2Size < 64 && totalS2Size % blockBytes == 0 + dataCopyParams.dstStride = 1; + alignedSize = blockBytes; + alignedS2Size = CeilDiv(s2Size, blockBytes) * blockBytes; + } + if (totalS2Size % alignedSize == 0) { + dataCopyParams.blockLen = alignedS2Size / blockBytes; + dataCopyParams.srcStride = (totalS2Size - alignedS2Size) / blockBytes; + DataCopy(dstTensor, srcTensor[srcOffset], dataCopyParams); + } else { + dataCopyParams.blockLen = s2Size; + dataCopyParams.srcStride = totalS2Size - s2Size; + DataCopyPadParams dataCopyPadParams; + dataCopyPadParams.isPad = true; + dataCopyPadParams.rightPadding = Min(alignedS2Size - s2Size, blockBytes); + dataCopyPadParams.paddingValue = 1; + DataCopyPad(dstTensor, srcTensor[srcOffset], dataCopyParams, dataCopyPadParams); + } +} + +__aicore__ inline void Bit2Int8CopyIn(LocalTensor &dstTensor, GlobalTensor &srcTensor, + int64_t srcOffset, uint32_t batchSize, uint32_t s1BaseSize, uint32_t s2BaseSize, int64_t s2TotalSize, + int64_t alignedSize = blockBytes) +{ + uint32_t alignedS2Size = CeilDiv(s2BaseSize / byteBitRatio, alignedSize) * alignedSize; + uint32_t shapeArray[] = {batchSize * s1BaseSize, alignedS2Size}; + dstTensor.SetShapeInfo(ShapeInfo(2, shapeArray, DataFormat::ND)); + dstTensor.SetSize(batchSize * s1BaseSize * alignedS2Size); + DataCopyParams dataCopyParams; + dataCopyParams.blockCount = batchSize * s1BaseSize; + dataCopyParams.blockLen = CeilDiv(s2BaseSize / byteBitRatio, blockBytes); + dataCopyParams.dstStride = 0; + if (s2TotalSize / byteBitRatio % alignedSize == 0 && s2BaseSize / byteBitRatio % alignedSize == 0) { + dataCopyParams.srcStride = + (s2TotalSize / byteBitRatio - dataCopyParams.blockLen * blockBytes) / blockBytes; + DataCopy(dstTensor, srcTensor[srcOffset / byteBitRatio], dataCopyParams); + } else { + dataCopyParams.blockLen = CeilDiv(s2BaseSize , byteBitRatio); + dataCopyParams.srcStride = (s2TotalSize - s2BaseSize) / byteBitRatio; + DataCopyPadParams dataCopyPadParams; + dataCopyPadParams.isPad = true; + dataCopyPadParams.rightPadding = 0; + dataCopyPadParams.paddingValue = 0; + DataCopyPad(dstTensor, srcTensor[srcOffset / byteBitRatio], dataCopyParams, dataCopyPadParams); + } +} + +__aicore__ inline int32_t Align(int32_t shape) +{ + int32_t alignFactor = 16; + int32_t alignedSize = CeilDiv(shape, alignFactor) * alignFactor; + return alignedSize; +} + +#endif // FLASH_ATTENTION_UTIL_H diff --git a/csrc/ascend/utils/inc/log/inner/dfx_base.h b/csrc/ascend/utils/inc/log/inner/dfx_base.h new file mode 100644 index 000000000..0fd1edb4e --- /dev/null +++ b/csrc/ascend/utils/inc/log/inner/dfx_base.h @@ -0,0 +1,190 @@ +/** + * Copyright (c) 2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dfx_base.h + * \brief 外部模块不应直接引用本头文件 + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ops { +namespace utils { + +class LogBase { +public: + static constexpr const int MAX_LOG_LEN = 16000; + static constexpr const int MSG_HDR_LEN = 200; + + static inline uint64_t GetTid() + { + return static_cast(syscall(__NR_gettid)); + } + + static inline const char *GetStr(const std::string &str) + { + return str.c_str(); + } + + static inline const char *GetStr(const char *str) + { + return str; + } + + static inline const std::string &GetOpInfo(const std::string &str) + { + return str; + } + + static inline const char *GetOpInfo(const char *str) + { + return str; + } + + static inline std::string GetOpInfo(const gert::TilingContext *context) + { + return GetOpInfoFromContext(context); + } + + static inline std::string GetOpInfo(const gert::TilingParseContext *context) + { + return GetOpInfoFromContext(context); + } + + static inline std::string GetOpInfo(const gert::InferShapeContext *context) + { + return GetOpInfoFromContext(context); + } + + static inline std::string GetOpInfo(const gert::InferDataTypeContext *context) + { + return GetOpInfoFromContext(context); + } + +private: + template static inline std::string GetOpInfoFromContext(T context) + { + if (context == nullptr) { + return "nil:nil"; + } + std::string opInfo = context->GetNodeType() != nullptr ? context->GetNodeType() : "nil"; + opInfo += ":"; + opInfo += context->GetNodeName() != nullptr ? context->GetNodeName() : "nil"; + return opInfo; + } +}; + +} // namespace utils + +template +std::string Shape2String(const T& shape) { + std::ostringstream oss; + oss << "["; + if (shape.GetDimNum() > 0) { + for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) { + oss << shape.GetDim(i) << ", "; + } + oss << shape.GetDim(shape.GetDimNum() - 1); + } + oss << "]"; + return oss.str(); +} +} // namespace ops + +// 使用本宏前需预定义标识子模块名称的 OPS_UTILS_LOG_SUB_MOD_NAME +// 如: #define OPS_UTILS_LOG_SUB_MOD_NAME "OP_TILING" 或通过 CMake 传递预定义宏 +#define OPS_LOG_STUB(MOD_ID, LOG_LEVEL, OPS_DESC, FMT, ...) \ + do { \ + if (AlogCheckDebugLevel(static_cast(MOD_ID), (LOG_LEVEL)) == 1) { \ + AlogRecord(static_cast(MOD_ID), DLOG_TYPE_DEBUG, (LOG_LEVEL), \ + "[%s:%d][%s]%s[%s][%lu] OpName:[%s] " #FMT, \ + __FILE__, __LINE__, (OPS_UTILS_LOG_SUB_MOD_NAME), \ + (OPS_UTILS_LOG_PACKAGE_TYPE), __FUNCTION__, ops::utils::LogBase::GetTid(), \ + ops::utils::LogBase::GetStr(ops::utils::LogBase::GetOpInfo(OPS_DESC)), ##__VA_ARGS__); \ + } \ + }while (0) + +#define OPS_LOG_STUB_IF(COND, LOG_FUNC, EXPR) \ + static_assert(std::is_same::type>::value, "condition should be bool"); \ + do { \ + if (__builtin_expect((COND), 0)) { \ + LOG_FUNC; \ + EXPR; \ + } \ + } while (0) + +#define OPS_INNER_ERR_STUB(ERR_CODE_STR, OPS_DESC, FMT, ...) \ + do { \ + OPS_LOG_STUB(OP, DLOG_ERROR, OPS_DESC, FMT, ##__VA_ARGS__); \ + REPORT_INNER_ERR_MSG(ERR_CODE_STR, FMT, ##__VA_ARGS__); \ + } while (0) + +#define OPS_CALL_ERR_STUB(ERR_CODE_STR, OPS_DESC, FMT, ...) \ + do { \ + OPS_LOG_STUB(OP, DLOG_ERROR, OPS_DESC, FMT, ##__VA_ARGS__); \ + REPORT_INNER_ERR_MSG(ERR_CODE_STR, FMT, ##__VA_ARGS__); \ + } while (0) + +#define OPS_LOG_STUB_D(OPS_DESC, FMT, ...) OPS_LOG_STUB(OP, DLOG_DEBUG, OPS_DESC, FMT, ##__VA_ARGS__) +#define OPS_LOG_STUB_I(OPS_DESC, FMT, ...) OPS_LOG_STUB(OP, DLOG_INFO, OPS_DESC, FMT, ##__VA_ARGS__) +#define OPS_LOG_STUB_W(OPS_DESC, FMT, ...) OPS_LOG_STUB(OP, DLOG_WARN, OPS_DESC, FMT, ##__VA_ARGS__) +#define OPS_LOG_STUB_E(OPS_DESC, FMT, ...) OPS_LOG_STUB(OP, DLOG_ERROR, OPS_DESC, FMT, ##__VA_ARGS__) +#define OPS_LOG_STUB_EVENT(OPS_DESC, FMT, ...) OPS_LOG_STUB(OP, DLOG_EVENT, OPS_DESC, FMT, ##__VA_ARGS__) + +#define OPS_LOG_STUB_FULL(LEVEL, OPS_DESC, FMT, ...) \ + do { \ + if (0 == AlogCheckDebugLevel(OP, (LEVEL))) { \ + break; \ + } \ + char msgbufxyz[ops::utils::LogBase::MAX_LOG_LEN]; \ + size_t msgmaxlen = (MSG_LENGTH - ops::utils::LogBase::MSG_HDR_LEN); \ + int rettmp = snprintf_s(msgbufxyz, sizeof(msgbufxyz), sizeof(msgbufxyz) - 1, FMT, ##__VA_ARGS__); \ + if (rettmp == -1) { \ + msgbufxyz[sizeof(msgbufxyz) - 1] = '\0'; \ + } \ + size_t msglength = std::strlen(msgbufxyz); \ + if (msglength < msgmaxlen) { \ + OPS_LOG_STUB(OP, (LEVEL), (OPS_DESC), "%s", msgbufxyz); \ + break; \ + } \ + char *msgchunkbegin = msgbufxyz; \ + char *msgchunkend = nullptr; \ + while (msgchunkbegin < msgbufxyz + msglength) { \ + if (msgchunkbegin[0] == '\n') { \ + OPS_LOG_STUB(OP, (LEVEL), (OPS_DESC), ""); \ + msgchunkbegin += 1; \ + continue; \ + } \ + msgchunkend = std::strchr(msgchunkbegin, '\n'); \ + if (msgchunkend == nullptr) { \ + msgchunkend = msgchunkbegin + std::strlen(msgchunkbegin); \ + } \ + while (msgchunkend > msgchunkbegin) { \ + std::string msgchunk(msgchunkbegin, \ + std::min(msgmaxlen, static_cast(msgchunkend - msgchunkbegin))); \ + OPS_LOG_STUB(OP, (LEVEL), (OPS_DESC), "%s", msgchunk.c_str()); \ + msgchunkbegin += msgchunk.size(); \ + } \ + msgchunkbegin += 1; \ + } \ + } while (0) diff --git a/csrc/ascend/utils/inc/log/ops_log.h b/csrc/ascend/utils/inc/log/ops_log.h new file mode 100644 index 000000000..e7653a89b --- /dev/null +++ b/csrc/ascend/utils/inc/log/ops_log.h @@ -0,0 +1,59 @@ +/** + * Copyright (c) 2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file ops_log.h + * \brief + */ + +#pragma once + +#include "log/inner/dfx_base.h" + +/* 基础日志 */ +#define OPS_LOG_D(OPS_DESC, ...) OPS_LOG_STUB_D(OPS_DESC, __VA_ARGS__) +#define OPS_LOG_I(OPS_DESC, ...) OPS_LOG_STUB_I(OPS_DESC, __VA_ARGS__) +#define OPS_LOG_W(OPS_DESC, ...) OPS_LOG_STUB_W(OPS_DESC, __VA_ARGS__) +#define OPS_LOG_E(OPS_DESC, ...) OPS_INNER_ERR_STUB("EZ9999", OPS_DESC, __VA_ARGS__) +#define OPS_LOG_E_WITHOUT_REPORT(OPS_DESC, ...) OPS_LOG_STUB_E(OPS_DESC, __VA_ARGS__) +#define OPS_LOG_EVENT(OPS_DESC, ...) OPS_LOG_STUB_EVENT(OPS_DESC, __VA_ARGS__) + +/* 全量日志 + * 输出超长日志, 若日志超长, 则会被分为多行输出 */ +#define OPS_LOG_FULL(LEVEL, OPS_DESC, ...) OPS_LOG_STUB_FULL(LEVEL, OPS_DESC, __VA_ARGS__) +#define OPS_LOG_D_FULL(OPS_DESC, ...) OPS_LOG_STUB_FULL(DLOG_DEBUG, OPS_DESC, __VA_ARGS__) +#define OPS_LOG_I_FULL(OPS_DESC, ...) OPS_LOG_STUB_FULL(DLOG_INFO, OPS_DESC, __VA_ARGS__) +#define OPS_LOG_W_FULL(OPS_DESC, ...) OPS_LOG_STUB_FULL(DLOG_WARN, OPS_DESC, __VA_ARGS__) + +/* 条件日志 */ +#define OPS_LOG_D_IF(COND, OP_DESC, EXPR, ...) OPS_LOG_STUB_IF(COND, OPS_LOG_D(OP_DESC, __VA_ARGS__), EXPR) +#define OPS_LOG_I_IF(COND, OP_DESC, EXPR, ...) OPS_LOG_STUB_IF(COND, OPS_LOG_I(OP_DESC, __VA_ARGS__), EXPR) +#define OPS_LOG_W_IF(COND, OP_DESC, EXPR, ...) OPS_LOG_STUB_IF(COND, OPS_LOG_W(OP_DESC, __VA_ARGS__), EXPR) +#define OPS_LOG_E_IF(COND, OP_DESC, EXPR, ...) OPS_LOG_STUB_IF(COND, OPS_LOG_E(OP_DESC, __VA_ARGS__), EXPR) +#define OPS_LOG_EVENT_IF(COND, OP_DESC, EXPR, ...) OPS_LOG_STUB_IF(COND, OPS_LOG_EVENT(OP_DESC, __VA_ARGS__), EXPR) + +#define OPS_LOG_E_IF_NULL(OPS_DESC, PTR, EXPR) \ + if (__builtin_expect((PTR) == nullptr, 0)) { \ + OPS_LOG_STUB_E(OPS_DESC, "%s is nullptr!", #PTR); \ + OPS_CALL_ERR_STUB("EZ9999", OPS_DESC, "%s is nullptr!", #PTR); \ + EXPR; \ + } + +#define OPS_CHECK(COND, LOG_FUNC, EXPR) \ + if (COND) { \ + LOG_FUNC; \ + EXPR; \ + } + +#define OP_CHECK(COND, LOG_FUNC, EXPR) \ + if (COND) { \ + LOG_FUNC; \ + EXPR; \ + } diff --git a/csrc/ascend/utils/inc/tiling/data_copy_transpose_tiling.h b/csrc/ascend/utils/inc/tiling/data_copy_transpose_tiling.h new file mode 100644 index 000000000..7e8d15d7f --- /dev/null +++ b/csrc/ascend/utils/inc/tiling/data_copy_transpose_tiling.h @@ -0,0 +1,47 @@ +/** + * Copyright (c) 2023-2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file data_copy_transpose_tiling.h + * \brief + */ + +#pragma once + +#include +#include +#include "data_copy_transpose_tiling_def.h" + +namespace optiling { + +inline void GetDataCopyTransposeTiling(const ge::Shape &dstShape, const ge::Shape &srcShape, const uint32_t typeSize, + optiling::CopyTransposeTiling &tiling) +{ + std::vector dstShapeInfo = dstShape.GetDims(); + std::vector srcShapeInfo = srcShape.GetDims(); + + tiling.set_dstShapeB(dstShapeInfo[0]); + tiling.set_dstShapeN(dstShapeInfo[1]); + tiling.set_dstShapeS(dstShapeInfo[2]); + tiling.set_dstShapeH(dstShapeInfo[3]); + tiling.set_dstShapeHN(tiling.get_dstShapeH() / tiling.get_dstShapeN()); + + tiling.set_srcShapeB(srcShapeInfo[0]); + tiling.set_srcShapeN(srcShapeInfo[1]); + tiling.set_srcShapeS(srcShapeInfo[2]); + tiling.set_srcShapeHN(srcShapeInfo[3]); + tiling.set_originalShapeNLen(tiling.get_srcShapeHN() * typeSize); + tiling.set_shapeSHValue(tiling.get_dstShapeS() * tiling.get_dstShapeH()); + tiling.set_shapeNsValue(tiling.get_dstShapeN() * tiling.get_dstShapeS()); + tiling.set_shapeNsnValue(tiling.get_dstShapeN() * tiling.get_srcShapeS() * tiling.get_srcShapeN()); + tiling.set_shapeBHValue(tiling.get_dstShapeB() * tiling.get_dstShapeH()); +} + +} // namespace optiling diff --git a/csrc/ascend/utils/inc/tiling/data_copy_transpose_tiling_def.h b/csrc/ascend/utils/inc/tiling/data_copy_transpose_tiling_def.h new file mode 100644 index 000000000..510b5cdad --- /dev/null +++ b/csrc/ascend/utils/inc/tiling/data_copy_transpose_tiling_def.h @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file data_copy_transpose_tiling_def.h + * \brief + */ + +#pragma once + +#include +#include + +namespace optiling { + +BEGIN_TILING_DATA_DEF(CopyTransposeTiling) +TILING_DATA_FIELD_DEF(uint32_t, dstShapeB); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeN); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeS); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeHN); +TILING_DATA_FIELD_DEF(uint32_t, dstShapeH); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeB); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeN); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeS); +TILING_DATA_FIELD_DEF(uint32_t, srcShapeHN); +TILING_DATA_FIELD_DEF(uint32_t, originalShapeNLen); +TILING_DATA_FIELD_DEF(uint32_t, shapeSHValue); +TILING_DATA_FIELD_DEF(uint32_t, shapeNsValue); +TILING_DATA_FIELD_DEF(uint32_t, shapeNsnValue); +TILING_DATA_FIELD_DEF(uint32_t, invalidParamCopyTransposeTiling); +TILING_DATA_FIELD_DEF(uint32_t, shapeBHValue); +TILING_DATA_FIELD_DEF(uint32_t, paramsAlign); +END_TILING_DATA_DEF; +REGISTER_TILING_DATA_CLASS(CopyTransposeTilingOp, CopyTransposeTiling) + +} // namespace optiling diff --git a/csrc/ascend/utils/inc/tiling/tiling_base.h b/csrc/ascend/utils/inc/tiling/tiling_base.h new file mode 100644 index 000000000..9776d90c8 --- /dev/null +++ b/csrc/ascend/utils/inc/tiling/tiling_base.h @@ -0,0 +1,225 @@ +/** + * Copyright (c) 2023-2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_base.h + * \brief + */ + +#pragma once + +#include +#include +#include +#include +#include "log/ops_log.h" + +#ifdef ASCENDC_OP_TEST +#define ASCENDC_EXTERN_C extern "C" +#else +#define ASCENDC_EXTERN_C +#endif + +namespace optiling { + +struct AiCoreParams { + uint64_t ubSize; + uint64_t blockDim; + uint64_t aicNum; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; +}; + +struct FlashAttentionScoreGradCompileInfo { + uint32_t aivNum; + uint32_t aicNum; + uint64_t ubSize; + uint64_t l1Size; + uint64_t l0aSize; + uint64_t l0bSize; + uint64_t l0cSize; + uint64_t l2CacheSize; + int64_t coreNum; +}; + +class TilingBaseClass { +public: + TilingBaseClass() = default; + + explicit TilingBaseClass(gert::TilingContext *context) : context_(context) + { + } + + virtual ~TilingBaseClass() = default; + + // Tiling执行框架 + // 1、GRAPH_SUCCESS: 成功，并且不需要继续执行后续Tiling类的实现 + // 2、GRAPH_FAILED: 失败，中止整个Tiling流程 + // 3、GRAPH_PARAM_INVALID: 本类不支持，需要继续往下执行其他Tiling类的实现 + ge::graphStatus DoTiling() + { + auto ret = GetShapeAttrsInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetPlatformInfo(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + if (!IsCapable()) { + return ge::GRAPH_PARAM_INVALID; + } + ret = DoOpTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = DoLibApiTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = GetWorkspaceSize(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + ret = PostTiling(); + if (ret != ge::GRAPH_SUCCESS) { + return ret; + } + context_->SetTilingKey(GetTilingKey()); + DumpTilingInfo(); + return ge::GRAPH_SUCCESS; + } + + // 更新 context + virtual void Reset(gert::TilingContext *context) + { + context_ = context; + } + +protected: + virtual bool IsCapable() = 0; + // 1、获取平台信息比如CoreNum、UB/L1/L0C资源大小 + virtual ge::graphStatus GetPlatformInfo() = 0; + // 2、获取INPUT/OUTPUT/ATTR信息 + virtual ge::graphStatus GetShapeAttrsInfo() = 0; + // 3、计算数据切分TilingData + virtual ge::graphStatus DoOpTiling() = 0; + // 4、计算高阶API的TilingData + virtual ge::graphStatus DoLibApiTiling() = 0; + // 5、计算TilingKey + [[nodiscard]] virtual uint64_t GetTilingKey() const = 0; + // 6、计算Workspace 大小 + virtual ge::graphStatus GetWorkspaceSize() = 0; + // 7、保存Tiling数据 + virtual ge::graphStatus PostTiling() = 0; + // 8、Dump Tiling数据 + virtual void DumpTilingInfo() + { + int32_t enable = AlogCheckDebugLevel(static_cast(OP), DLOG_DEBUG); + if (enable != 1) { + return; + } + auto buf = (uint32_t *)context_->GetRawTilingData()->GetData(); + auto bufLen = context_->GetRawTilingData()->GetDataSize(); + std::ostringstream oss; + oss << "Start to dump tiling info. tilingkey:" << GetTilingKey() << ", tiling data size:" << bufLen + << ", content:"; + for (size_t i = 0; i < bufLen / sizeof(uint32_t); i++) { + oss << *(buf + i) << ","; + if (oss.str().length() > 640) { // Split according to 640 to avoid truncation + OPS_LOG_D(context_, "%s", oss.str().c_str()); + oss.str(""); + } + } + OPS_LOG_D(context_, "%s", oss.str().c_str()); + } + + static uint32_t CalcTschBlockDim(uint32_t sliceNum, uint32_t aicCoreNum, uint32_t aivCoreNum) + { + uint32_t ration; + if (aicCoreNum == 0 || aivCoreNum == 0 || aicCoreNum > aivCoreNum) { + return sliceNum; + } + ration = aivCoreNum / aicCoreNum; + return (sliceNum + (ration - 1)) / ration; + } + + template [[nodiscard]] std::string GetShapeDebugStr(const T &shape) const + { + std::ostringstream oss; + oss << "["; + if (shape.GetDimNum() > 0) { + for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) { + oss << shape.GetDim(i) << ", "; + } + oss << shape.GetDim(shape.GetDimNum() - 1); + } + oss << "]"; + return oss.str(); + } + + [[nodiscard]] std::string GetTensorDebugStr(const gert::StorageShape *shape, + const gert::CompileTimeTensorDesc *tensor) + { + if (shape == nullptr || tensor == nullptr) { + return "nil "; + } + std::ostringstream oss; + oss << "(dtype: " << ge::TypeUtils::DataTypeToSerialString(tensor->GetDataType()) << "),"; + oss << "(shape:" << GetShapeDebugStr(shape->GetStorageShape()) << "),"; + oss << "(ori_shape:" << GetShapeDebugStr(shape->GetOriginShape()) << "),"; + oss << "(format: " + << ge::TypeUtils::FormatToSerialString( + static_cast(ge::GetPrimaryFormat(tensor->GetStorageFormat()))) + << "),"; + oss << "(ori_format: " << ge::TypeUtils::FormatToSerialString(tensor->GetOriginFormat()) << ") "; + return oss.str(); + } + + [[nodiscard]] std::string GetTilingContextDebugStr() + { + std::ostringstream oss; + for (size_t i = 0; i < context_->GetComputeNodeInfo()->GetInputsNum(); ++i) { + oss << "input" << i << ": "; + oss << GetTensorDebugStr(context_->GetInputShape(i), context_->GetInputDesc(i)); + } + + for (size_t i = 0; i < context_->GetComputeNodeInfo()->GetOutputsNum(); ++i) { + oss << "output" << i << ": "; + oss << GetTensorDebugStr(context_->GetOutputShape(i), context_->GetOutputDesc(i)); + } + return oss.str(); + } + + [[nodiscard]] std::string GetTilingDataDebugStr() const + { + auto rawTilingData = context_->GetRawTilingData(); + auto rawTilingDataSize = rawTilingData->GetDataSize(); + auto data = reinterpret_cast(rawTilingData->GetData()); + size_t len = rawTilingDataSize / sizeof(int32_t); + std::ostringstream oss; + for (size_t i = 0; i < len; i++) { + oss << data[i] << ", "; + } + return oss.str(); + } + +protected: + gert::TilingContext *context_ = nullptr; + std::unique_ptr ascendcPlatform_{nullptr}; + uint32_t blockDim_{0}; + uint64_t workspaceSize_{0}; + uint64_t tilingKey_{0}; + AiCoreParams aicoreParams_{0, 0, 0, 0, 0, 0, 0}; +}; + +} // namespace optiling diff --git a/csrc/ascend/utils/inc/tiling/tiling_templates_registry.h b/csrc/ascend/utils/inc/tiling/tiling_templates_registry.h new file mode 100644 index 000000000..53fc590aa --- /dev/null +++ b/csrc/ascend/utils/inc/tiling/tiling_templates_registry.h @@ -0,0 +1,162 @@ +/** + * Copyright (c) 2023-2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_templates_registry.h + * \brief + */ + +#pragma once + +#include +#include +#include +#include +#include "tiling/tiling_base.h" +#include "log/ops_log.h" +#include "error/ops_error.h" + +namespace optiling { + +template std::unique_ptr TILING_CLASS(gert::TilingContext *context) +{ + return std::unique_ptr(new (std::nothrow) T(context)); +} + +using TilingClassCase = std::unique_ptr (*)(gert::TilingContext *); + +class TilingCases { +public: + explicit TilingCases(std::string op_type) : op_type_(std::move(op_type)) + { + } + + template void AddTiling(int32_t priority) + { + OPS_ERR_IF(cases_.find(priority) != cases_.end(), + OPS_REPORT_VECTOR_INNER_ERR(op_type_, "There are duplicate registrations."), return); + cases_[priority] = TILING_CLASS; + OPS_ERR_IF( + cases_[priority] == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(op_type_, "Register op tiling func failed, please check the class name."), + return); + } + + const std::map &GetTilingCases() + { + return cases_; + } + +private: + std::map cases_; + const std::string op_type_; +}; + +class TilingRegistry { +public: + TilingRegistry() = default; + +#ifdef ASCENDC_OP_TEST + static TilingRegistry &GetInstance(); +#else + static TilingRegistry &GetInstance() + { + static TilingRegistry registry_impl_; + return registry_impl_; + } +#endif + + std::shared_ptr RegisterOp(const std::string &op_type) + { + if (registry_map_.find(op_type) == registry_map_.end()) { + registry_map_[op_type] = std::shared_ptr(new (std::nothrow) TilingCases(op_type)); + } + OPS_ERR_IF(registry_map_[op_type] == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(op_type, "Register tiling func failed, please check the class name."), + return nullptr); + return registry_map_[op_type]; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext *context) + { + const char *op_type = context->GetNodeType(); + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type); + for (auto it = tilingTemplateRegistryMap.begin(); it != tilingTemplateRegistryMap.end(); ++it) { + auto tilingTemplate = it->second(context); + if (tilingTemplate != nullptr) { + ge::graphStatus status = tilingTemplate->DoTiling(); + if (status != ge::GRAPH_PARAM_INVALID) { + OPS_LOG_D(context, "Do general op tiling success priority=%d", it->first); + return status; + } + OPS_LOG_D(context, "Ignore general op tiling priority=%d", it->first); + } + } + OPS_REPORT_VECTOR_INNER_ERR(op_type, "Do op tiling failed, no valid template is found."); + return ge::GRAPH_FAILED; + } + + ge::graphStatus DoTilingImpl(gert::TilingContext *context, const std::vector &priorities) + { + const char *op_type = context->GetNodeType(); + auto tilingTemplateRegistryMap = GetTilingTemplates(op_type); + for (auto priorityId : priorities) { + auto templateFunc = tilingTemplateRegistryMap[priorityId](context); + if (templateFunc != nullptr) { + ge::graphStatus status = templateFunc->DoTiling(); + if (status == ge::GRAPH_SUCCESS) { + OPS_LOG_D(context, "Do general op tiling success priority=%d", priorityId); + return status; + } + OPS_LOG_D(context, "Ignore general op tiling priority=%d", priorityId); + } + } + return ge::GRAPH_FAILED; + } + + const std::map &GetTilingTemplates(const std::string &op_type) + { + OPS_ERR_IF(registry_map_.find(op_type) == registry_map_.end(), + OPS_REPORT_VECTOR_INNER_ERR(op_type, "Get op tiling func failed, please check the op name."), + return empty_tiling_case_); + return registry_map_[op_type]->GetTilingCases(); + } + +private: + std::map> registry_map_; + const std::map empty_tiling_case_ {}; +}; + +class Register { +public: + explicit Register(std::string op_type) : op_type_(std::move(op_type)) + { + } + + template Register &tiling(int32_t priority) + { + auto tilingCases = TilingRegistry::GetInstance().RegisterOp(op_type_); + OPS_ERR_IF(tilingCases == nullptr, + OPS_REPORT_VECTOR_INNER_ERR(op_type_, "Register op tiling failed, please the op name."), + return *this); + tilingCases->AddTiling(priority); + return *this; + } + +private: + const std::string op_type_; +}; + +// op_type: 算子名称， class_name: 注册的 tiling 类, +// priority: tiling 类的优先级, 越小表示优先级越高, 即被选中的概率越大 +#define REGISTER_TILING_TEMPLATE(op_type, class_name, priority) \ + static Register VAR_UNUSED##op_type_##class_name##priority_register = Register(op_type).tiling(priority) + +} // namespace optiling diff --git a/csrc/ascend/utils/inc/tiling/tiling_type.h b/csrc/ascend/utils/inc/tiling/tiling_type.h new file mode 100644 index 000000000..dfdd67f69 --- /dev/null +++ b/csrc/ascend/utils/inc/tiling/tiling_type.h @@ -0,0 +1,136 @@ +/** + * Copyright (c) 2023-2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file tiling_type.h + * \brief + */ + +#pragma once + +#include + +namespace optiling { + +enum class AxisEnum { + B = 0, + N2 = 1, + G = 2, + S1 = 3, + S2 = 4, + D = 5, + NONE = 9, +}; + +enum class DtypeEnum { + FLOAT16 = 0, + FLOAT32 = 1, + BFLOAT16 = 2, + FLOAT16_PRECISION = 3, +}; + +enum class PerformanceOrientedEnum { + BIG_BUFFER = 1, + BIG_DOUBLE_BUFFER = 2, +}; + +enum class MatmulConfig { + NULL_CONFIG = 0, + NORMAL_CONFIG = 1, + MDL_CONFIG = 2 +}; + +enum class PseConfig { + NO_PSE = 0, + EXIST_PSE = 1 +}; + +enum class AttenMaskConfig { + NO_ATTEN_MASK = 0, + EXIST_ATTEN_MASK = 1 +}; + +enum class DropOutConfig { + NO_DROP_OUT = 0, + EXIST_DROP_OUT = 1 +}; + +enum class CubeFormatEnum { + ND = 0, + NZ = 1 +}; +enum class LayoutEnum { + BSND = 0, + SBND = 1, + BNSD = 2, + TND = 3 +}; + +enum class CubeInputSourceEnum { + GM = 0, + L1 = 1 +}; + +enum class OptionEnum { + DISABLE = 0, + ENABLE = 1 +}; + +enum class SparseEnum { + ALL = 0, + NONE = 1, + ANY = 2, + CAUSAL = 3, + BAND = 4, + PREFIX = 5, + BAND_COMPRESS = 6, + RIGHT_DOWN_CAUSAL = 7, + RIGHT_DOWN_CAUSAL_BAND = 8, + BAND_LEFT_UP_CAUSAL = 9 +}; + +constexpr uint64_t RecursiveSum() +{ + return 0; +} + +template constexpr uint64_t RecursiveSum(T templateId, Args... templateIds) +{ + return static_cast(templateId) + 10 * RecursiveSum(templateIds...); +} + +// TilingKey 的生成规则： +// FlashAttentionScore/FlashAttentionScoreGrad 十进制位组装tiling key，包含以下关键参数，从低位到高位依次是：Ub0, Ub1, +// Block, DataType, Format, Sparse, 特化模板 Ub0、Ub1: +// 表示Ub核内切分的轴，使用枚举AxisEnum表示，因为我们允许最多切分两根轴，所以存在UB0和UB1，如果没有UB核内切分， +// 那么填AXIS_NONE。UB0和UB1各占一个十进制位; +// Block: 表示UB用来分核的轴，使用枚举AxisEnum表示，占一个十进制位; +// DataType: 表示当前tiling key支持的输入输出的数据类型，使用枚举SupportedDtype来表示，占一个十进制位 +// Format: 表示当前tiling key支持的Format, 使用枚举InputLayout表示，占一个十进制位 +// Sparse: 表示当前tiling key是否支持Sparse，使用枚举SparseCapability表示，占一个十进制位 +// 其余特化场景，定义自己的位域和值 +// usage: get tilingKey from inputted types +// uint64_t tilingKey = GET_FLASHATTENTION_TILINGKEY(AxisEnum::AXIS_S1, AxisEnum::AXIS_S2, AxisEnum::AXIS_N2, +// SupportedDtype::FLOAT32, InputLayout::BSH, SparseCapability::SUPPORT_ALL) + +constexpr uint64_t TILINGKEYOFFSET = uint64_t(10000000000000000000UL); // 10^19 +template constexpr uint64_t GET_TILINGKEY(Args... templateIds) +{ + return TILINGKEYOFFSET + RecursiveSum(templateIds...); +} + +// usage: get tilingKey from inputted types +// uint64_t tilingKey = TILINGKEY(S2, S1, N2, FLOAT32, BSND, ALL) + +#define TILINGKEY(ub2, ub1, block, dtype, layout, sparse) \ + (GET_TILINGKEY(AxisEnum::ub2, AxisEnum::ub1, AxisEnum::block, DtypeEnum::dtype, LayoutEnum::layout, \ + SparseEnum::sparse)) + +} // namespace optiling diff --git a/csrc/ascend/utils/src/fallback_comm.cpp b/csrc/ascend/utils/src/fallback_comm.cpp new file mode 100644 index 000000000..949cb7289 --- /dev/null +++ b/csrc/ascend/utils/src/fallback_comm.cpp @@ -0,0 +1,53 @@ +/** + * Copyright (c) 2024 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file fallback_comm.cpp + * \brief + */ + +#include "fallback_comm.h" + +#include +#include +#include +#include + +#include "aclnn/aclnn_base.h" +#include "runtime/base.h" + +#ifdef __cplusplus +extern "C" { +#endif + +namespace fallback { +using namespace std; +using namespace gert; +using namespace ge; + +aclDataType ToAclDataType(ge::DataType dtype) { + static const std::vector CANN_CONVERT_TO_ACL_DataType_LIST = { + ge::DataType::DT_FLOAT, ge::DataType::DT_FLOAT16, ge::DataType::DT_INT8, ge::DataType::DT_INT32, + ge::DataType::DT_UINT8, ge::DataType::DT_INT16, ge::DataType::DT_UINT16, ge::DataType::DT_UINT32, + ge::DataType::DT_INT64, ge::DataType::DT_DOUBLE, ge::DataType::DT_BOOL, ge::DataType::DT_STRING, + ge::DataType::DT_COMPLEX64, ge::DataType::DT_COMPLEX128, ge::DataType::DT_BF16, ge::DataType::DT_UINT64, + ge::DataType::DT_INT4}; + auto iter = std::find(CANN_CONVERT_TO_ACL_DataType_LIST.begin(), CANN_CONVERT_TO_ACL_DataType_LIST.end(), dtype); + if (iter == CANN_CONVERT_TO_ACL_DataType_LIST.end()) { + return aclDataType::ACL_DT_UNDEFINED; + } + return static_cast(dtype); +} + +} // namespace fallback + +#ifdef __cplusplus +} +#endif diff --git a/csrc/ascend/version.info b/csrc/ascend/version.info new file mode 100644 index 000000000..acb557b85 --- /dev/null +++ b/csrc/ascend/version.info @@ -0,0 +1,11 @@ +Version=9.0.1 +version_dir=cann +required_package_runtime_version=">=9.0.1" +required_package_opbase_version=">=9.0.1" +required_package_asc-devkit_version=">=9.0.1" +required_package_ge-compiler_version=">=9.0.1" +required_package_metadef_version=">=9.0.1" +required_package_ge-executor_version=">=9.0.1" +required_package_bisheng-compiler_version=">=9.0.1" +required_package_tbe-kit_version=">=9.0.1" +required_package_asc-tools_version=">=9.0.1" \ No newline at end of file diff --git a/setup.py b/setup.py new file mode 100644 index 000000000..b66f6c36c --- /dev/null +++ b/setup.py @@ -0,0 +1,277 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# +# vllm-plugin-FL: vLLM Federated Learning Plugin +# +# This setup script builds the vllm_fl._C / vllm_fl._C_ascend vendor-specific +# C++ extension via CMake. It supports CUDA and Ascend backends controlled by +# the VLLM_VENDOR environment variable. The build pipeline: +# 1. Detects available tooling (cmake, ninja, sccache/ccache) +# 2. Configures and compiles the C++/CUDA/Ascend sources under csrc/ +# 3. Copies the resulting shared library (.so/.pyd) to the package directory + +from __future__ import annotations + +import glob +import logging +import os +import shutil +import subprocess +import sys +from pathlib import Path +from shutil import which + +from setuptools import Extension, setup +from setuptools.command.build_ext import build_ext + +ROOT_DIR = Path(__file__).parent.resolve() +logger = logging.getLogger(__name__) + +VLLM_VENDOR = os.environ.get("VLLM_VENDOR", "").lower() +MAX_JOBS = os.environ.get("MAX_JOBS") +NVCC_THREADS = os.environ.get("NVCC_THREADS") +CMAKE_BUILD_TYPE = os.environ.get("CMAKE_BUILD_TYPE") +VERBOSE = os.environ.get("VERBOSE", "0") == "1" + +SUPPORTED_VENDORS = ("cuda", "ascend") + + +def _is_cuda() -> bool: + return VLLM_VENDOR == "cuda" + + +def _is_ascend() -> bool: + return VLLM_VENDOR == "ascend" + + +def _which(name: str) -> bool: + return which(name) is not None + + +def _get_torch_npu_path() -> str: + """Return the directory containing the installed torch_npu package.""" + try: + import torch_npu + return str(Path(torch_npu.__file__).parent) + except Exception as exc: + raise RuntimeError( + "torch_npu is required for Ascend builds but could not be imported." + ) from exc + + +def _get_ascend_home_path() -> str: + """Return the ASCEND_HOME_PATH (CANN toolkit installation root).""" + ascend_home = os.environ.get("ASCEND_HOME_PATH") + if ascend_home: + return ascend_home + # Common default location. + default = "/usr/local/Ascend/ascend-toolkit/latest" + if os.path.isdir(default): + return default + raise RuntimeError( + "ASCEND_HOME_PATH is not set and the default CANN path does not exist. " + "Please set ASCEND_HOME_PATH to your CANN toolkit installation directory." + ) + + +class CMakeExtension(Extension): + def __init__(self, name: str, cmake_lists_dir: str) -> None: + super().__init__(name, sources=[]) + self.cmake_lists_dir = os.path.abspath(cmake_lists_dir) + + +class CMakeBuildExt(build_ext): + did_config: dict[str, bool] = {} + + def run(self) -> None: + self.build_extensions() + + def compute_num_jobs(self) -> tuple[int, int | None]: + if MAX_JOBS is not None: + num_jobs = int(MAX_JOBS) + logger.info("Using MAX_JOBS=%d as the number of jobs.", num_jobs) + else: + try: + num_jobs = len(os.sched_getaffinity(0)) + except AttributeError: + num_jobs = os.cpu_count() or 1 + + nvcc_threads = None + if _is_cuda() and NVCC_THREADS is not None: + nvcc_threads = int(NVCC_THREADS) + logger.info("Using NVCC_THREADS=%d.", nvcc_threads) + num_jobs = max(1, num_jobs // nvcc_threads) + + return num_jobs, nvcc_threads + + def configure(self, ext: CMakeExtension) -> None: + if CMakeBuildExt.did_config.get(ext.cmake_lists_dir): + return + + CMakeBuildExt.did_config[ext.cmake_lists_dir] = True + cfg = CMAKE_BUILD_TYPE or ("Debug" if self.debug else "RelWithDebInfo") + cmake_args = [ + f"-DCMAKE_BUILD_TYPE={cfg}", + f"-DVLLM_VENDOR={VLLM_VENDOR}", + f"-DVLLM_PYTHON_EXECUTABLE={sys.executable}", + ] + + if VERBOSE: + cmake_args.append("-DCMAKE_VERBOSE_MAKEFILE=ON") + + if _which("sccache"): + cmake_args += [ + "-DCMAKE_C_COMPILER_LAUNCHER=sccache", + "-DCMAKE_CXX_COMPILER_LAUNCHER=sccache", + "-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache", + ] + elif _which("ccache"): + cmake_args += [ + "-DCMAKE_C_COMPILER_LAUNCHER=ccache", + "-DCMAKE_CXX_COMPILER_LAUNCHER=ccache", + "-DCMAKE_CUDA_COMPILER_LAUNCHER=ccache", + ] + + num_jobs, nvcc_threads = self.compute_num_jobs() + if nvcc_threads: + cmake_args.append(f"-DNVCC_THREADS={nvcc_threads}") + + if _is_ascend(): + cmake_args += [ + f"-DTORCH_NPU_PATH={_get_torch_npu_path()}", + f"-DASCEND_HOME_PATH={_get_ascend_home_path()}", + ] + soc_version = os.environ.get("SOC_VERSION") + if soc_version: + cmake_args.append(f"-DSOC_VERSION={soc_version}") + if os.environ.get("BUILD_PTO_CHUNK_GDN"): + cmake_args.append("-DBUILD_PTO_CHUNK_GDN=ON") + + build_tool = [] + # AscendC kernel auto-codegen currently assumes a Makefile generator + # (link.txt, etc.). Do not use Ninja for Ascend builds even if available. + if _is_cuda() and _which("ninja"): + build_tool = ["-G", "Ninja"] + cmake_args += [ + "-DCMAKE_JOB_POOL_COMPILE:STRING=compile", + f"-DCMAKE_JOB_POOLS:STRING=compile={num_jobs}", + ] + + extra_cmake_args = os.environ.get("CMAKE_ARGS") + if extra_cmake_args: + cmake_args += extra_cmake_args.split() + + subprocess.check_call( + ["cmake", ext.cmake_lists_dir, *build_tool, *cmake_args], + cwd=self.build_temp, + ) + + def _apply_cann_extract_host_stub_patch(self) -> Path | None: + """Apply a temporary workaround patch for CANN 8.5.x extract_host_stub.py.""" + if not _is_ascend(): + return None + script = ROOT_DIR / "csrc/ascend/patch_cann_extract_host_stub.py" + if not script.is_file(): + raise RuntimeError(f"CANN patch script not found: {script}") + backup_path = Path(self.build_temp) / "extract_host_stub.py.vllm_fl_bak" + subprocess.check_call( + [sys.executable, str(script), "apply", "--backup", str(backup_path)] + ) + return backup_path + + def _restore_cann_extract_host_stub_patch(self, backup_path: Path | None) -> None: + """Restore the original CANN extract_host_stub.py from backup.""" + if backup_path is None: + return + script = ROOT_DIR / "csrc/ascend/patch_cann_extract_host_stub.py" + if not script.is_file(): + return + subprocess.run( + [sys.executable, str(script), "restore", "--backup", str(backup_path)], + check=False, + ) + + def build_extensions(self) -> None: + try: + subprocess.check_output(["cmake", "--version"], stderr=subprocess.STDOUT) + except (OSError, subprocess.CalledProcessError) as exc: + raise RuntimeError( + "CMake is required to build vllm_fl extensions. " + "Install cmake and run with VLLM_VENDOR=cuda or VLLM_VENDOR=ascend." + ) from exc + + os.makedirs(self.build_temp, exist_ok=True) + + targets = [] + for ext in self.extensions: + self.configure(ext) + targets.append(ext.name.split(".")[-1]) + + num_jobs, _ = self.compute_num_jobs() + build_args = [ + "--build", + ".", + f"-j={num_jobs}", + *[f"--target={name}" for name in targets], + ] + + cann_patch_backup: Path | None = None + try: + cann_patch_backup = self._apply_cann_extract_host_stub_patch() + subprocess.check_call(["cmake", *build_args], cwd=self.build_temp) + + for ext in self.extensions: + dest_path = Path(self.get_ext_fullpath(ext.name)).absolute() + dest_path.parent.mkdir(parents=True, exist_ok=True) + target_name = ext.name.split(".")[-1] + patterns = [ + f"{self.build_temp}/{VLLM_VENDOR}/{target_name}*.so", + f"{self.build_temp}/{VLLM_VENDOR}/{target_name}*.pyd", + f"{self.build_temp}/{target_name}*.so", + f"{self.build_temp}/{target_name}*.pyd", + ] + built_ext = next( + (match for pattern in patterns for match in glob.glob(pattern)), + None, + ) + if built_ext is None: + raise RuntimeError( + f"Could not find built extension {target_name} in {self.build_temp}" + ) + shutil.copy2(built_ext, dest_path) + + # Ascend builds produce an additional shared AscendC kernel library + # that the extension links against. Copy it next to the extension so + # the $ORIGIN rpath can resolve it at load time. + if _is_ascend(): + kernel_lib_patterns = [ + f"{self.build_temp}/lib/libvllm_fl_kernels*.so", + f"{self.build_temp}/**/libvllm_fl_kernels*.so", + ] + for pattern in kernel_lib_patterns: + for kernel_lib in glob.glob(pattern, recursive=True): + shutil.copy2(kernel_lib, dest_path.parent / Path(kernel_lib).name) + break + else: + continue + break + finally: + self._restore_cann_extract_host_stub_patch(cann_patch_backup) + + +ext_modules = [] +if VLLM_VENDOR: + if VLLM_VENDOR not in SUPPORTED_VENDORS: + raise ValueError( + f"Unsupported vendor: {VLLM_VENDOR}. " + f"Supported vendors: {', '.join(SUPPORTED_VENDORS)}" + ) + ext_name = "vllm_fl._C_ascend" if _is_ascend() else "vllm_fl._C" + ext_modules.append( + CMakeExtension(name=ext_name, cmake_lists_dir=str(ROOT_DIR / "csrc")) + ) + + +setup( + ext_modules=ext_modules, + cmdclass={"build_ext": CMakeBuildExt} if ext_modules else {}, +) diff --git a/tests/custom_ops_tests/README.md b/tests/custom_ops_tests/README.md new file mode 100644 index 000000000..74edef688 --- /dev/null +++ b/tests/custom_ops_tests/README.md @@ -0,0 +1,259 @@ +# vLLM-FL Ascend Custom Ops Tests + +本目录包含 `vllm-plugin-FL` 在 Ascend NPU 上的自定义算子连接测试。 + +## 1. `csrc/ascend/` 下的两种算子接入方式 + +`csrc/ascend/` 目前承载两类 Ascend 自定义算子，它们的编译、安装和运行加载方式完全不同： + +### 1.1 CANN Framework 算子（aclnn 路径） + +- **源码位置**：`csrc/ascend///` + - 例如 `csrc/ascend/moe/causal_conv1d/`、`csrc/ascend/attention/fused_gdn_gating/`。 +- **构建工具链**：CANN `op_host` / `op_kernel` / `aclnn` 工具链。 +- **产物**：自解压 `.run` 算子包，例如 `csrc/ascend/build/cann-ops-transformer-custom_linux-aarch64.run`。 +- **安装位置**：默认隔离安装到项目目录 `vllm_fl/_cann_ops_custom/vendors/custom_transformer/`，不污染系统 CANN。 + - 该目录是构建产物，**不应提交到版本控制**。 +- **运行时加载**： + - C++ torch extension `vllm_fl._C_ascend` 注册 `torch.ops._C_ascend.*` schema； + - `vllm_fl.utils.enable_custom_op()` 会自动发现已安装的 `_cann_ops_custom` 包，设置 `ASCEND_CUSTOM_OPP_PATH` 和 `LD_LIBRARY_PATH`； + - 即使 `set_env.bash` 中写入了安装时的绝对路径，运行时也以 `vllm_fl` 包的实际位置为准，因此安装目录可以被移动。 + +### 1.2 PTO GDN 预编译算子（Bisheng 路径） + +- **源码位置**：`csrc/ascend/pto_chunk_gdn/`。 +- **构建工具链**：Bisheng C++ 编译器（`-xcce --cce-aicore-arch=dav-c220`），直接编译 `.cpp` 为 AI Core `.so`。 +- **依赖头库**：`csrc/ascend/third_party/pto-isa/`。 +- **产物**：多个 `mega_kernel_H*_Hg*_D*_C*.so`。 +- **安装位置**： + - 预编译模式：安装到 Python site-packages 下的 `vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib/`； + - JIT 模式：首次调用时由 `vllm_fl/ops/pto_chunk_gdn/compile.py` 自动编译并缓存到同一目录。 +- **运行时加载**：Python 代码通过 `ctypes.CDLL` / `torch.ops.load_library` 直接加载 `.so`，不经过 CANN `opp/vendors` 路径。 + +## 2. 环境准备 + +```bash +# 1. 激活 CANN 环境 +source /usr/local/Ascend/ascend-toolkit/set_env.sh + +# 2. 初始化 catlass 和 pto-isa 源码子模块（这两个第三方库不随主仓库一起拉取） +git submodule update --init --recursive csrc/ascend/third_party/catlass +git submodule update --init --recursive csrc/ascend/third_party/pto-isa + +# 3. （可选）固定子模块版本，确保与上游来源一致 +# catlass 固定到 vllm-ascend 使用的 commit +# pto-isa 固定到 PR #8872 使用的 commit（如有需要） +cd csrc/ascend/third_party/catlass +git checkout 41bf90da655bba3c66d0acd7e00abe33960ecfd6 +cd ../../.. + +cd csrc/ascend/third_party/pto-isa +# 如需固定版本，替换为 PR #8872 对应的 commit +# git checkout +cd ../../.. + +# 4. 确认环境变量（根据实际安装路径调整） +export ASCEND_HOME_PATH=/usr/local/Ascend/cann-9.0.0 +export SOC_VERSION=ascend910b # 根据实际芯片调整 +``` + +## 3. 一键构建并运行（推荐） + +项目提供了 `tests/custom_ops_tests/build_and_run.sh`，在激活 CANN 环境并初始化子模块后，可以一键完成扩展编译、CANN framework 算子检查以及全部测试： + +```bash +bash tests/custom_ops_tests/build_and_run.sh +``` + +脚本会自动完成以下步骤： + +1. 检查 `ASCEND_HOME_PATH` / `ASCEND_TOOLKIT_HOME` 是否已设置； +2. 检测 CANN 版本，要求 **CANN 9.0.0 及以上**，否则终止并提示安装； +3. 询问是否配置 GitHub / pip 镜像（分别对应 `ghfast.top` 和清华源）； +4. 检查并初始化 `catlass`、`pto-isa` 子模块； +5. 编译/安装 Python 包与 torch extension； +6. （可选）编译并安装 CANN framework 算子包； +7. 检查 CANN framework 算子包是否已安装； +8. source `set_env.bash` 并设置 `ASCEND_CUSTOM_OPP_PATH` / `LD_LIBRARY_PATH`； +9. 检查 `FlagGems`，如果没有则从 `https://github.com/flagos-ai/FlagGems` clone 并以 editable 模式安装； +10. 依次运行 `tests/custom_ops_tests/test_*.py`。 + +### 常用选项 + +| 选项 | 说明 | +|---|---| +| `--build-ops` | 从源码编译并安装 CANN framework 算子包（执行 `csrc/ascend/build_aclnn.sh`）。 | +| `--soc ` | 指定芯片版本，如 `ascend910b`、`ascend910_93`；优先级高于环境变量 `SOC_VERSION`，默认由 `build_aclnn.sh` 决定。 | +| `--editable` / `-e` | 使用 `VLLM_VENDOR=ascend pip install --no-build-isolation -e .` 以 editable 模式安装（当前环境必须已有 `torch_npu`），而不执行 `build_ext --inplace`。 | +| `-h` / `--help` | 显示帮助信息。 | + +示例： + +```bash +# 默认：只编译 torch extension、检查已安装的算子包、跑测试 +bash tests/custom_ops_tests/build_and_run.sh + +# 一键完成所有事情：编译 extension + 编译安装 framework 算子 + 跑测试 +bash tests/custom_ops_tests/build_and_run.sh --build-ops + +# 指定 910C 芯片并启用 editable install +bash tests/custom_ops_tests/build_and_run.sh --build-ops --soc ascend910_93 --editable +``` + +> **注意**：脚本不会替你安装 CANN toolkit。如果第 7 步报错且你没有加 `--build-ops`，请先安装 CANN framework 算子包（见第 5 节），或者重新运行脚本并加上 `--build-ops`。 + +## 4. 编译并安装 torch extension `vllm_fl._C_ascend` + +该 extension 把 C++ 算子实现注册到 `torch.ops._C_ascend`，同时包含 `camem_allocator` 等基础设施。 + +```bash +VLLM_VENDOR=ascend python setup.py build_ext --inplace +``` + +完成后会在项目根目录生成： + +```text +vllm_fl/_C_ascend.cpython-311-aarch64-linux-gnu.so +vllm_fl/libvllm_fl_kernels.so +``` + +测试脚本中通过 `import vllm_fl._C_ascend` 加载 extension，随后即可调用 `torch.ops._C_ascend.*`。 + +## 5. 安装 CANN framework 算子包 + +如果已经存在构建好的 `.run` 包（例如 `csrc/ascend/build/cann-ops-transformer-custom_linux-aarch64.run`），可以直接安装： + +```bash +bash csrc/ascend/build/cann-ops-transformer-custom_linux-aarch64.run \ + --install-path="$(pwd)/vllm_fl/_cann_ops_custom" +``` + +`.run` 包会把算子安装到指定的 `--install-path` 下，生成： + +```text +vllm_fl/_cann_ops_custom/vendors/custom_transformer/ +├── bin/set_env.bash # 安装时生成的环境脚本 +├── op_api/include/aclnn/ # aclnn 头文件 +├── op_api/lib/libcust_opapi.so +├── op_proto/ +└── op_impl/ +``` + +> **注意**：`set_env.bash` 会记录安装时的绝对路径。测试脚本不依赖该脚本，而是调用 `vllm_fl.utils.enable_custom_op()` 根据 `vllm_fl` 包的实际位置自动设置环境变量，因此安装目录可以被移动。 + +如果需要从头编译 `.run` 包，执行： + +```bash +bash csrc/ascend/build_aclnn.sh +# 例如：bash csrc/ascend/build_aclnn.sh ascend910b +``` + +构建完成后会生成 `csrc/ascend/build/cann-ops-transformer-custom_linux-aarch64.run` 并自动安装到 `vllm_fl/_cann_ops_custom/`。 + +如果希望清理构建过程中下载的第三方库（`abseil-cpp`、`ascend_protobuf`、`json`、`pkg` 缓存），保留源码子模块 `catlass` 和 `pto-isa`，可以加上 `--clean-third-party`： + +```bash +bash csrc/ascend/build_aclnn.sh ascend910b --clean-third-party +``` + +## 6. PTO GDN 算子的两种使用方式 + +### 方式 A：预编译（推荐生产环境） + +```bash +VLLM_VENDOR=ascend BUILD_PTO_CHUNK_GDN=ON python setup.py build_ext --inplace + +# 显式编译 PTO megakernel +cmake --build build/temp.linux-aarch64-cpython-311 \ + --target pto_chunk_gdn_kernels -j$(nproc) +``` + +产物会安装到当前 Python 环境 site-packages 下的： + +```text +vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib/ +├── mega_kernel_H16_Hg8_D128_C128.so +├── mega_kernel_H16_Hg16_D128_C128.so +└── ... +``` + +> 在 editable install（`pip install -e .`）下，`_PACKAGE_ROOT` 等于仓库根目录，因此也会写到仓库内的 `vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib/`。 + +### 方式 B：JIT 首次编译（开发调试用） + +不预编译，直接运行 `tests/custom_ops_tests/test_pto_chunk_gdn.py`。`vllm_fl/ops/pto_chunk_gdn/compile.py` 会： + +1. 自动查找 `csrc/ascend/third_party/pto-isa`； +2. 调用系统 `bisheng` 编译对应配置的 `mega_kernel_*.so`； +3. 缓存到 `vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib/`； +4. 后续调用直接复用缓存。 + +## 7. 目录结构总览 + +```text +csrc/ascend/ +├── CMakeLists.txt # 构建 _C_ascend + 分发 pto_chunk_gdn +├── torch_binding.cpp # torch.ops._C_ascend 注册 +├── torch_binding_meta.cpp # meta kernel 注册 +├── camem_allocator.cpp # NPU 显存分配器 +├── build.sh # CANN framework 算子构建脚本 +├── build_aclnn.sh # 一键打包 + 安装 .run +├── build/ # 构建产物（包含 .run 包） +│ └── cann-ops-transformer-custom_linux-aarch64.run +├── // # CANN framework 算子源码 +│ └── op_host/op_kernel/... +├── pto_chunk_gdn/ # PTO GDN megakernel 源码 +│ ├── CMakeLists.txt +│ ├── mega_kernel.cpp +│ └── include/ +└── third_party/ + ├── catlass/ # CANN 算子依赖 + └── pto-isa/ # PTO 算子依赖 +``` + +## 8. 如何执行测试 + +测试脚本会自动调用 `vllm_fl.utils.enable_custom_op()` 设置 CANN 自定义算子环境，**不需要手动 `source set_env.bash`**。 + +### 7.1 逐个运行 + +```bash +# CANN framework 算子 +python tests/custom_ops_tests/test_causal_conv1d.py +python tests/custom_ops_tests/test_fused_gdn_gating.py +python tests/custom_ops_tests/test_gemma_rms_norm.py +python tests/custom_ops_tests/test_recurrent_gated_delta_rule.py +python tests/custom_ops_tests/test_chunk_gated_delta_rule_fwd_h.py + +# PTO GDN 算子（首次运行会触发 Bisheng JIT 编译） +python tests/custom_ops_tests/test_pto_chunk_gdn.py +``` + +### 7.2 批量运行 + +```bash +for f in tests/custom_ops_tests/test_*.py; do + echo "=== $f ===" + python "$f" 2>&1 | tail -3 +done +``` + +### 7.3 常见失败原因 + +| 现象 | 原因 | 解决 | +|---|---|---| +| `AttributeError: '_OpNamespace' '_C_ascend' object has no attribute 'xxx'` | `vllm_fl._C_ascend` 未编译或算子未注册 | 重新执行 `VLLM_VENDOR=ascend python setup.py build_ext --inplace` | +| `aclnnXxx ... not in libopapi.so` | `_cann_ops_custom` 未安装 | 执行 `bash csrc/ascend/build/cann-ops-transformer-custom_linux-aarch64.run --install-path="$(pwd)/vllm_fl/_cann_ops_custom"` | +| `ImportError: dynamic module does not define module export function (PyInit__C_ascend)` | `camem_allocator.cpp` 里的 PyInit 函数名与 extension 名不匹配 | 检查 `csrc/ascend/camem_allocator.cpp` 是否为 `PyInit__C_ascend` | +| PTO 测试提示找不到 `pto-isa` | 子模块未初始化或路径错误 | `git submodule update --init --recursive csrc/ascend/third_party/pto-isa` | + +## 9. 测试脚本说明 + +| 测试脚本 | 对应算子 | 接入方式 | +|---|---|---| +| `test_causal_conv1d.py` | `npu_causal_conv1d_custom` | CANN framework | +| `test_fused_gdn_gating.py` | `npu_fused_gdn_gating` | CANN framework | +| `test_gemma_rms_norm.py` | `npu_gemma_rms_norm` | CANN framework | +| `test_recurrent_gated_delta_rule.py` | `npu_recurrent_gated_delta_rule` | CANN framework | +| `test_chunk_gated_delta_rule_fwd_h.py` | `chunk_gated_delta_rule_fwd_h` | CANN framework | +| `test_pto_chunk_gdn.py` | PTO GDN megakernel | Bisheng PTO | diff --git a/tests/custom_ops_tests/build_and_run.sh b/tests/custom_ops_tests/build_and_run.sh new file mode 100755 index 000000000..1602a6488 --- /dev/null +++ b/tests/custom_ops_tests/build_and_run.sh @@ -0,0 +1,407 @@ +#!/bin/bash +# Copyright (c) 2026 BAAI. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# One-shot build and test runner for vllm-plugin-FL Ascend custom ops. +# +# This script assumes the caller has already activated the CANN environment +# (e.g. source /usr/local/Ascend/ascend-toolkit/set_env.sh). It will verify +# that CANN >= 9.0.0 is available, then build the torch extension, optionally +# build/install the CANN framework operator package, and finally run all tests +# under tests/custom_ops_tests/. + +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd) +ROOT_DIR=$(cd "$SCRIPT_DIR/../.." && pwd) +cd "$ROOT_DIR" + +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +BLUE='\033[0;34m' +NC='\033[0m' + +log_info() { echo -e "${GREEN}[INFO]${NC} $*"; } +log_warn() { echo -e "${YELLOW}[WARN]${NC} $*"; } +log_error() { echo -e "${RED}[ERROR]${NC} $*"; } +log_step() { echo -e "${BLUE}[STEP]${NC} $*"; } + +usage() { + cat < SOC version passed to build_aclnn.sh, e.g. ascend910b + or ascend910_93. Defaults to \$SOC_VERSION, or + ascend910b if neither is set. + --editable, -e Install the package in editable mode + (VLLM_VENDOR=ascend pip install -e .) instead of + running VLLM_VENDOR=ascend python setup.py build_ext + --inplace. + -h, --help Show this help message. + +Examples: + # Default: build extension, check installed ops, run tests + bash tests/custom_ops_tests/build_and_run.sh + + # Also compile and install CANN framework operators + bash tests/custom_ops_tests/build_and_run.sh --build-ops + + # Compile operators for ascend910_93 and install editable + bash tests/custom_ops_tests/build_and_run.sh --build-ops --soc ascend910_93 --editable +EOF +} + +BUILD_OPS=0 +SOC_VERSION_ARG="" +EDITABLE=0 + +while [[ $# -gt 0 ]]; do + case "$1" in + --build-ops) + BUILD_OPS=1 + ;; + --soc) + if [[ $# -lt 2 ]]; then + log_error "--soc requires a value." + usage + exit 1 + fi + SOC_VERSION_ARG="$2" + shift + ;; + --editable|-e) + EDITABLE=1 + ;; + -h|--help) + usage + exit 0 + ;; + *) + log_error "Unknown option: $1" + usage + exit 1 + ;; + esac + shift +done + +# ----------------------------------------------------------------------------- +# 1. CANN version check +# ----------------------------------------------------------------------------- +check_cann() { + log_step "Checking CANN environment ..." + + local ascend_home="${ASCEND_HOME_PATH:-${ASCEND_TOOLKIT_HOME:-}}" + if [[ -z "$ascend_home" ]]; then + log_error "ASCEND_HOME_PATH or ASCEND_TOOLKIT_HOME is not set." + log_error "Please activate the CANN environment first, for example:" + log_error " source /usr/local/Ascend/ascend-toolkit/set_env.sh" + exit 1 + fi + + # Follow symlinks such as .../ascend-toolkit/latest + ascend_home=$(readlink -f "$ascend_home" 2>/dev/null || echo "$ascend_home") + + if [[ ! -d "$ascend_home" ]]; then + log_error "CANN path does not exist: $ascend_home" + exit 1 + fi + + log_info "CANN home: $ascend_home" + + local version="" + + # Try to extract version from the directory name, e.g. /usr/local/Ascend/cann-9.0.0 + version=$(echo "$ascend_home" | grep -oE '[0-9]+\.[0-9]+(\.[0-9]+)?' | head -n1 || true) + + # Fallback: look for version.info or ascend_toolkit_install.info + if [[ -z "$version" && -f "$ascend_home/version.info" ]]; then + version=$(grep -oE '[0-9]+\.[0-9]+(\.[0-9]+)?' "$ascend_home/version.info" | head -n1 || true) + fi + if [[ -z "$version" && -f "$ascend_home/ascend_toolkit_install.info" ]]; then + version=$(grep -oE '[0-9]+\.[0-9]+(\.[0-9]+)?' "$ascend_home/ascend_toolkit_install.info" | head -n1 || true) + fi + + if [[ -z "$version" ]]; then + log_error "Cannot detect CANN version from $ascend_home." + log_error "Make sure you have sourced the correct CANN set_env.sh." + exit 1 + fi + + log_info "Detected CANN version: $version" + + # Compare major.minor against 9.0 + local major minor + major=$(echo "$version" | cut -d. -f1) + minor=$(echo "$version" | cut -d. -f2) + + if [[ "$major" -lt 9 ]] || { [[ "$major" -eq 9 && "$minor" -lt 0 ]]; }; then + log_error "CANN version $version is too old. This project requires CANN 9.0.0 or higher." + log_error "Please install a compatible CANN toolkit and try again." + exit 1 + fi + + log_info "CANN version check passed." +} + +# ----------------------------------------------------------------------------- +# 2. Optional environment configuration +# ----------------------------------------------------------------------------- +ask_yes_no() { + local prompt="$1" + local response + read -rp "$prompt [y/N]: " response + case "$response" in + [Yy]*) return 0 ;; + *) return 1 ;; + esac +} + +configure_git_mirror() { + log_step "Configuring git mirror ..." + git config --global url."https://ghfast.top/https://github.com/".insteadOf "https://github.com/" + log_info "Git mirror configured." +} + +configure_pip_mirror() { + log_step "Configuring pip mirror ..." + pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple + log_info "Pip mirror configured." +} + +# ----------------------------------------------------------------------------- +# 3. Submodule check +# ----------------------------------------------------------------------------- +check_submodules() { + log_step "Checking source submodules ..." + + local missing=0 + for sub in csrc/ascend/third_party/catlass csrc/ascend/third_party/pto-isa; do + if [[ ! -e "$sub/.git" ]]; then + log_warn "Submodule $sub is not initialized." + missing=1 + fi + done + + if [[ "$missing" -eq 1 ]]; then + log_info "Initializing catlass and pto-isa submodules ..." + git submodule update --init --recursive csrc/ascend/third_party/catlass + git submodule update --init --recursive csrc/ascend/third_party/pto-isa + else + log_info "Submodules are already initialized." + fi +} + +# ----------------------------------------------------------------------------- +# 4. Build/install the Python package +# ----------------------------------------------------------------------------- +build_extension() { + log_step "Building vllm_fl._C_ascend torch extension ..." + + if [[ ! -f "setup.py" ]]; then + log_error "setup.py not found in $ROOT_DIR. Are you running this script from the project root?" + exit 1 + fi + + VLLM_VENDOR=ascend python setup.py build_ext --inplace + + log_info "Torch extension built successfully." +} + +install_editable() { + log_step "Installing vllm-plugin-FL in editable mode ..." + + if [[ ! -f "setup.py" ]]; then + log_error "setup.py not found in $ROOT_DIR." + exit 1 + fi + + if ! command -v pip >/dev/null 2>&1; then + log_error "pip is not available in the current environment." + exit 1 + fi + + # This project imports torch_npu at build time, which is not available in + # PEP 517 isolated build environments. Use --no-build-isolation so the + # current Python environment (where torch_npu is installed) is used. + VLLM_VENDOR=ascend pip install --no-build-isolation -e . + + log_info "Editable install complete." +} + +# ----------------------------------------------------------------------------- +# 5. Build/install CANN framework operators (optional) +# ----------------------------------------------------------------------------- +build_ops() { + log_step "Building and installing CANN framework operators ..." + + local soc_args=() + if [[ -n "$SOC_VERSION_ARG" ]]; then + soc_args=("$SOC_VERSION_ARG") + elif [[ -n "${SOC_VERSION:-}" ]]; then + soc_args=("$SOC_VERSION") + fi + + if [[ ${#soc_args[@]} -gt 0 ]]; then + log_info "SOC version: ${soc_args[0]}" + else + log_info "SOC version: (default from build_aclnn.sh)" + fi + + bash csrc/ascend/build_aclnn.sh "${soc_args[@]}" + + log_info "CANN framework operators built and installed." +} + +# ----------------------------------------------------------------------------- +# 6. Check CANN framework operator package +# ----------------------------------------------------------------------------- +check_cann_framework_ops() { + log_step "Checking CANN framework operator package ..." + + local vendor_dir="vllm_fl/_cann_ops_custom/vendors/custom_transformer" + if [[ ! -d "$vendor_dir" ]]; then + log_error "CANN framework operators are not installed at vllm_fl/_cann_ops_custom/" + log_error "" + log_error "If you already have a built .run package, install it with:" + log_error " bash csrc/ascend/build/cann-ops-transformer-custom_linux-aarch64.run \\" + log_error " --install-path=\$(pwd)/vllm_fl/_cann_ops_custom" + log_error "" + log_error "Or build from source by adding --build-ops to this script:" + log_error " bash tests/custom_ops_tests/build_and_run.sh --build-ops [--soc ]" + exit 1 + fi + + log_info "CANN framework operators found." +} + +# ----------------------------------------------------------------------------- +# 7. Set up CANN custom-op runtime environment +# ----------------------------------------------------------------------------- +setup_cann_op_env() { + log_step "Setting up CANN custom-op runtime environment ..." + + local vendor_dir="${ROOT_DIR}/vllm_fl/_cann_ops_custom/vendors/custom_transformer" + local set_env_script="${vendor_dir}/bin/set_env.bash" + + # The package-provided set_env.bash appends to these variables, so make + # sure they are defined before sourcing it. This avoids an unbound + # variable failure when this script is run with `set -u`. + export ASCEND_CUSTOM_OPP_PATH="${ASCEND_CUSTOM_OPP_PATH:-}" + export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}" + + if [[ -f "$set_env_script" ]]; then + log_info "Sourcing ${set_env_script} ..." + # shellcheck source=/dev/null + source "$set_env_script" + fi + + # Always override the two path variables with the actual install location, + # in case the package has been relocated since installation. + export ASCEND_CUSTOM_OPP_PATH="$vendor_dir" + export LD_LIBRARY_PATH="${vendor_dir}/op_api/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + + log_info "ASCEND_CUSTOM_OPP_PATH=$ASCEND_CUSTOM_OPP_PATH" +} + +# ----------------------------------------------------------------------------- +# 8. Check / install FlagGems +# ----------------------------------------------------------------------------- +check_flag_gems() { + log_step "Checking FlagGems ..." + + if python -c "import flag_gems" >/dev/null 2>&1; then + log_info "FlagGems is already installed." + return 0 + fi + + log_warn "FlagGems not found. Installing from source ..." + + local parent_dir + parent_dir=$(dirname "$ROOT_DIR") + cd "$parent_dir" + + if [[ ! -d "FlagGems" ]]; then + log_info "Cloning FlagGems ..." + git clone https://github.com/flagos-ai/FlagGems + else + log_info "FlagGems directory already exists, skipping clone." + fi + + cd FlagGems + pip install --no-build-isolation -e . + log_info "FlagGems installed." + + cd "$ROOT_DIR" +} + +# ----------------------------------------------------------------------------- +# 9. Run tests +# ----------------------------------------------------------------------------- +run_tests() { + log_step "Running custom ops tests ..." + + local test_dir="tests/custom_ops_tests" + local failed=0 + + if [[ ! -d "$test_dir" ]]; then + log_error "Test directory not found: $test_dir" + exit 1 + fi + + for test in "$test_dir"/test_*.py; do + if [[ ! -f "$test" ]]; then + continue + fi + log_info "Running $(basename "$test") ..." + if ! python "$test"; then + log_error "$(basename "$test") FAILED" + failed=1 + fi + done + + if [[ "$failed" -ne 0 ]]; then + log_error "Some tests failed." + exit 1 + fi + + log_info "All tests passed." +} + +# ----------------------------------------------------------------------------- +# Main +# ----------------------------------------------------------------------------- +main() { + log_info "Working directory: $ROOT_DIR" + check_cann + + if ask_yes_no "Configure git mirror for GitHub (ghfast.top)?"; then + configure_git_mirror + fi + if ask_yes_no "Configure pip mirror (Tsinghua)?"; then + configure_pip_mirror + fi + + check_submodules + if [[ "$EDITABLE" -eq 1 ]]; then + install_editable + else + build_extension + fi + if [[ "$BUILD_OPS" -eq 1 ]]; then + build_ops + fi + check_cann_framework_ops + setup_cann_op_env + check_flag_gems + run_tests + log_info "Done." +} + +main diff --git a/tests/custom_ops_tests/test_causal_conv1d.py b/tests/custom_ops_tests/test_causal_conv1d.py new file mode 100644 index 000000000..803e173e0 --- /dev/null +++ b/tests/custom_ops_tests/test_causal_conv1d.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Minimal connectivity test for the npu_causal_conv1d_custom framework op. +# +# NOTE: this test automatically enables the CANN custom-op environment by +# discovering the vllm_fl/_cann_ops_custom package installed next to vllm_fl. +# No manual `source set_env.bash` is required. + +import sys + +from vllm_fl.utils import enable_custom_op + +if not enable_custom_op(): + print( + "ERROR: vllm_fl/_cann_ops_custom is not installed.\n" + "Please build and install the CANN framework operators first, e.g.:\n" + " bash csrc/ascend/build_aclnn.sh ", + file=sys.stderr, + ) + sys.exit(1) + +import torch +import torch_npu + +# Load the Ascend C++ extension that registers torch.ops._C_ascend. +import vllm_fl._C_ascend # noqa: F401 + +DEVICE = "npu:0" + + +def test_causal_conv1d_varlen(): + """Varlen prefill mode (run_mode=0) with 2 sequences.""" + batch = 2 + seqlen = 16 + dim = 16 # must be aligned to 16 + width = 4 # in [2, 4] + state_len = width - 1 + num_cache_lines = batch + + total_seqlen = batch * seqlen + + x = torch.randn(total_seqlen, dim, dtype=torch.float16, device=DEVICE) + weight = torch.randn(width, dim, dtype=torch.float16, device=DEVICE) + conv_state = torch.randn(num_cache_lines, state_len, dim, dtype=torch.float16, device=DEVICE) + bias = torch.randn(dim, dtype=torch.float16, device=DEVICE) + query_start_loc = torch.tensor([0, seqlen, total_seqlen], dtype=torch.int32, device=DEVICE) + cache_indices = torch.arange(batch, dtype=torch.int32, device=DEVICE) + initial_state_mode = torch.zeros(batch, dtype=torch.int32, device=DEVICE) + + output = torch.empty_like(x) + + result = torch.ops._C_ascend.npu_causal_conv1d_custom( + output, + x, + weight, + conv_state, + bias, + query_start_loc, + cache_indices, + initial_state_mode, + None, # num_accepted_tokens + activation_mode=1, + pad_slot_id=-1, + run_mode=0, + ) + + assert result.shape == x.shape + assert result.device == x.device + + +if __name__ == "__main__": + test_causal_conv1d_varlen() + print("npu_causal_conv1d_custom varlen test passed") diff --git a/tests/custom_ops_tests/test_chunk_gated_delta_rule_fwd_h.py b/tests/custom_ops_tests/test_chunk_gated_delta_rule_fwd_h.py new file mode 100644 index 000000000..e3867b0db --- /dev/null +++ b/tests/custom_ops_tests/test_chunk_gated_delta_rule_fwd_h.py @@ -0,0 +1,73 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Minimal connectivity test for the chunk_gated_delta_rule_fwd_h framework op. +# +# NOTE: this test automatically enables the CANN custom-op environment by +# discovering the vllm_fl/_cann_ops_custom package installed next to vllm_fl. +# No manual `source set_env.bash` is required. + +import sys + +from vllm_fl.utils import enable_custom_op + +if not enable_custom_op(): + print( + "ERROR: vllm_fl/_cann_ops_custom is not installed.\n" + "Please build and install the CANN framework operators first, e.g.:\n" + " bash csrc/ascend/build_aclnn.sh ", + file=sys.stderr, + ) + sys.exit(1) + +import torch +import torch_npu + +import vllm_fl._C_ascend # noqa: F401 + +DEVICE = "npu:0" + + +def test_chunk_gated_delta_rule_fwd_h(): + """Basic chunk_gated_delta_rule_fwd_h connectivity test.""" + b = 2 + h = 2 + t = 64 + k = 32 + hv = 4 + v = 32 + chunk_size = 64 + + k_tensor = torch.randn(b, h, t, k, dtype=torch.bfloat16, device=DEVICE) + w_tensor = torch.randn(b, h, t, k, dtype=torch.bfloat16, device=DEVICE) + u_tensor = torch.randn(b, hv, t, v, dtype=torch.bfloat16, device=DEVICE) + g_tensor = torch.randn(b, h, t, k, dtype=torch.float32, device=DEVICE) + initial_state = torch.randn(b, hv, k, v, dtype=torch.bfloat16, device=DEVICE) + + cu_seqlens = [0, t, b * t] + + h_out, v_new_out, final_state_out = torch.ops._C_ascend.chunk_gated_delta_rule_fwd_h( + k_tensor, + w_tensor, + u_tensor, + g=g_tensor, + gk=None, + initial_state=initial_state, + output_final_state=True, + chunk_size=chunk_size, + save_new_value=True, + cu_seqlens=cu_seqlens, + chunk_indices=None, + use_exp2=False, + transpose_state_layout=False, + ) + + nt = (t + chunk_size - 1) // chunk_size + assert h_out.shape == (b, hv, nt, k, v) + assert v_new_out.shape == u_tensor.shape + assert final_state_out.shape == (b, hv, k, v) + + +if __name__ == "__main__": + test_chunk_gated_delta_rule_fwd_h() + print("chunk_gated_delta_rule_fwd_h test passed") diff --git a/tests/custom_ops_tests/test_fused_gdn_gating.py b/tests/custom_ops_tests/test_fused_gdn_gating.py new file mode 100644 index 000000000..1d733e37d --- /dev/null +++ b/tests/custom_ops_tests/test_fused_gdn_gating.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Minimal connectivity test for the npu_fused_gdn_gating framework op. +# +# NOTE: this test automatically enables the CANN custom-op environment by +# discovering the vllm_fl/_cann_ops_custom package installed next to vllm_fl. +# No manual `source set_env.bash` is required. + +import sys + +from vllm_fl.utils import enable_custom_op + +if not enable_custom_op(): + print( + "ERROR: vllm_fl/_cann_ops_custom is not installed.\n" + "Please build and install the CANN framework operators first, e.g.:\n" + " bash csrc/ascend/build_aclnn.sh ", + file=sys.stderr, + ) + sys.exit(1) + +import torch +import torch_npu + +import vllm_fl._C_ascend # noqa: F401 + +DEVICE = "npu:0" + + +def test_fused_gdn_gating(): + """Basic fused GDN gating connectivity test.""" + batch = 4 + num_heads = 8 + + A_log = torch.randn(num_heads, dtype=torch.bfloat16, device=DEVICE) + a = torch.randn(batch, num_heads, dtype=torch.bfloat16, device=DEVICE) + b = torch.randn(batch, num_heads, dtype=torch.bfloat16, device=DEVICE) + dt_bias = torch.randn(num_heads, dtype=torch.bfloat16, device=DEVICE) + + g, beta_output = torch.ops._C_ascend.npu_fused_gdn_gating( + A_log, a, b, dt_bias, beta=1.0, threshold=20.0 + ) + + assert g.shape == (1, batch, num_heads) + assert beta_output.shape == (1, batch, num_heads) + assert g.device == A_log.device + + +if __name__ == "__main__": + test_fused_gdn_gating() + print("npu_fused_gdn_gating test passed") diff --git a/tests/custom_ops_tests/test_gemma_rms_norm.py b/tests/custom_ops_tests/test_gemma_rms_norm.py new file mode 100644 index 000000000..2c669ac03 --- /dev/null +++ b/tests/custom_ops_tests/test_gemma_rms_norm.py @@ -0,0 +1,45 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Minimal connectivity test for the npu_gemma_rms_norm framework op. +# +# NOTE: this test automatically enables the CANN custom-op environment by +# discovering the vllm_fl/_cann_ops_custom package installed next to vllm_fl. +# No manual `source set_env.bash` is required. + +import sys + +from vllm_fl.utils import enable_custom_op + +if not enable_custom_op(): + print( + "ERROR: vllm_fl/_cann_ops_custom is not installed.\n" + "Please build and install the CANN framework operators first, e.g.:\n" + " bash csrc/ascend/build_aclnn.sh ", + file=sys.stderr, + ) + sys.exit(1) + +import torch +import torch_npu + +import vllm_fl._C_ascend # noqa: F401 + +DEVICE = "npu:0" + + +def test_gemma_rms_norm(): + """Basic GemmaRMSNorm connectivity test.""" + x = torch.randn(16, 128, dtype=torch.float16, device=DEVICE) + gamma = torch.randn(128, dtype=torch.float16, device=DEVICE) + + y, rstd = torch.ops._C_ascend.npu_gemma_rms_norm(x, gamma, 1e-6) + + assert y.shape == x.shape + assert y.device == x.device + assert rstd.shape == (x.size(0), 1) + + +if __name__ == "__main__": + test_gemma_rms_norm() + print("npu_gemma_rms_norm test passed") diff --git a/tests/custom_ops_tests/test_moe_fused_custom.py b/tests/custom_ops_tests/test_moe_fused_custom.py new file mode 100644 index 000000000..cda5d968e --- /dev/null +++ b/tests/custom_ops_tests/test_moe_fused_custom.py @@ -0,0 +1,156 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Correctness test for the AscendC fused-MoE custom ops used by the FL MoE path: +# * torch.ops._C_ascend.moe_gating_top_k (fused router: softmax + topk + renorm) +# * torch.ops._C_ascend.npu_moe_init_routing_custom (token expand/sort by expert) +# +# plus an end-to-end parity check of the new `_ascendc_fused_experts_impl` +# (custom routing + pre-transposed weights) against the legacy +# `_torch_fused_experts_impl` (npu_moe_init_routing_v2 + per-call transpose). +# +# NOTE: the custom `moe_grouped_matmul` (NZ-weight GMM) is intentionally not +# covered and not used by the plugin: its kernel crashes the aicore on +# ascend910b even for trivial single-expert inputs (CCU instruction address +# check error). The grouped matmuls stay on torch_npu.npu_grouped_matmul. + +import sys + +from vllm_fl.utils import enable_custom_op + +if not enable_custom_op(): + print( + "ERROR: vllm_fl/_cann_ops_custom is not installed.\n" + "Please build and install the CANN framework operators first, e.g.:\n" + " bash csrc/ascend/build_aclnn.sh ", + file=sys.stderr, + ) + sys.exit(1) + +import torch +import torch_npu + +import vllm_fl._C_ascend # noqa: F401 +from vllm_fl.dispatch.backends.vendor.ascend.impl.fused_moe import ( + _ascendc_fused_experts_impl, + _torch_fused_experts_impl, + fused_topk_ascend, +) + +DEVICE = "npu:0" + + +def _ref_topk_softmax(gating_output: torch.Tensor, topk: int, renormalize: bool): + """fp32 softmax -> topk -> optional L1 renorm (vLLM fused_topk semantics).""" + probs = torch.softmax(gating_output.to(torch.float32), dim=-1) + topk_weights, topk_ids = torch.topk(probs, topk, dim=-1) + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + return topk_weights, topk_ids.to(torch.int32) + + +def test_moe_gating_top_k(): + """moe_gating_top_k vs fp32 torch reference (Qwen3.6 shape: 256 experts, top-8).""" + torch.manual_seed(0) + num_tokens, num_experts, topk = 129, 256, 8 + + for in_dtype in (torch.float32, torch.bfloat16): + logits = torch.randn(num_tokens, num_experts, dtype=in_dtype, device=DEVICE) + + for renormalize in (True, False): + y, expert_idx, _ = torch.ops._C_ascend.moe_gating_top_k( + logits, + topk, + 1, # k_group + 1, # group_count + 0, # group_select_mode + int(renormalize), # renorm + 0, # norm_type: softmax + False, # out_flag + 1.0, # routed_scaling_factor + 1e-20, # eps + None, # bias_opt + ) + ref_w, ref_ids = _ref_topk_softmax(logits, topk, renormalize) + + assert y.dtype == logits.dtype, (y.dtype, logits.dtype) + assert expert_idx.dtype == torch.int32 + assert y.shape == (num_tokens, topk) + # Same experts selected (order within top-k may differ for ties). + for i in range(num_tokens): + assert set(expert_idx[i].tolist()) == set(ref_ids[i].tolist()), ( + f"token {i}: ids {expert_idx[i].tolist()} vs ref {ref_ids[i].tolist()}" + ) + # Gather ref weights in the op's id order and compare values. + ref_w_reorder = torch.gather( + torch.softmax(logits.to(torch.float32), dim=-1), 1, expert_idx.long() + ) + if renormalize: + ref_w_reorder = ref_w_reorder / ref_w_reorder.sum(dim=-1, keepdim=True) + # Kernel computes in reduced precision internally; allow bf16-level + # tolerance (fp32 input: kernel still rounds to ~bf16 precision). + rtol, atol = (1e-2, 1e-3) if in_dtype == torch.float32 else (2e-2, 1e-2) + torch.testing.assert_close( + y.to(torch.float32), ref_w_reorder, rtol=rtol, atol=atol + ) + + # fused_topk_ascend wrapper: fp32 weights out, ids dtype honored. + hidden = torch.randn(num_tokens, 64, dtype=torch.bfloat16, device=DEVICE) + logits = torch.randn(num_tokens, num_experts, dtype=torch.bfloat16, device=DEVICE) + w, ids, _ = fused_topk_ascend(hidden, logits, topk, True, torch.int64) + assert w.dtype == torch.float32 + assert ids.dtype == torch.int64 + assert w.shape == (num_tokens, topk) + torch.testing.assert_close( + w.sum(dim=-1), torch.ones(num_tokens, device=DEVICE), rtol=1e-3, atol=1e-3 + ) + print("moe_gating_top_k test passed") + + +def test_fused_experts_parity(): + """_ascendc_fused_experts_impl vs legacy _torch_fused_experts_impl.""" + torch.manual_seed(1) + num_tokens, hidden, intermediate = 33, 256, 128 + global_num_experts, topk = 16, 4 + ep_rank, num_local_experts = 1, 8 # exercise the expert_map path + + hidden_states = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=DEVICE) + w1 = torch.randn(num_local_experts, 2 * intermediate, hidden, dtype=torch.bfloat16, device=DEVICE) * 0.02 + w2 = torch.randn(num_local_experts, hidden, intermediate, dtype=torch.bfloat16, device=DEVICE) * 0.02 + + expert_map = torch.full((global_num_experts,), -1, dtype=torch.int32, device=DEVICE) + expert_map[ep_rank * num_local_experts:(ep_rank + 1) * num_local_experts] = torch.arange( + num_local_experts, dtype=torch.int32, device=DEVICE + ) + topk_ids = torch.randint(0, global_num_experts, (num_tokens, topk), dtype=torch.int32, device=DEVICE) + topk_weights = torch.rand(num_tokens, topk, dtype=torch.float32, device=DEVICE) + + # Legacy path (original weight layout [E, 2I, H] / [E, H, I]). + ref = _torch_fused_experts_impl( + hidden_states, w1, w2, topk_weights, topk_ids, + global_num_experts=global_num_experts, expert_map=expert_map, + ) + + # AscendC path (pre-transposed layout [E, H, 2I] / [E, I, H]). + w1_t = w1.transpose(1, 2).contiguous() + w2_t = w2.transpose(1, 2).contiguous() + out = _ascendc_fused_experts_impl( + hidden_states, w1_t, w2_t, topk_weights, topk_ids, + global_num_experts=global_num_experts, expert_map=expert_map, + ) + + assert out.shape == ref.shape == hidden_states.shape + torch.testing.assert_close(out.to(torch.float32), ref.to(torch.float32), rtol=2e-2, atol=2e-2) + + # Also without expert_map (dense path). + topk_ids_local = torch.randint(0, num_local_experts, (num_tokens, topk), dtype=torch.int32, device=DEVICE) + ref = _torch_fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids_local) + out = _ascendc_fused_experts_impl(hidden_states, w1_t, w2_t, topk_weights, topk_ids_local) + torch.testing.assert_close(out.to(torch.float32), ref.to(torch.float32), rtol=2e-2, atol=2e-2) + print("fused_experts parity test passed") + + +if __name__ == "__main__": + test_moe_gating_top_k() + test_fused_experts_parity() + print("All MoE custom-op tests passed") diff --git a/tests/custom_ops_tests/test_pto_chunk_gdn.py b/tests/custom_ops_tests/test_pto_chunk_gdn.py new file mode 100644 index 000000000..d5f181a9d --- /dev/null +++ b/tests/custom_ops_tests/test_pto_chunk_gdn.py @@ -0,0 +1,68 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Minimal integration test for the PTO GDN megakernel. +# Verifies that the 6-stage fused AscendC/PTO kernel can be compiled and +# launched through vllm-plugin-FL's Python bindings. + +import os +import sys +import tempfile + +import torch +import torch_npu + +# Allow importing the in-tree vllm_fl without installation. +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +sys.path.insert(0, ROOT) + + +def main() -> int: + torch.npu.config.allow_internal_format = True + + from vllm_fl.utils import enable_custom_op + + enable_custom_op() + + from vllm_fl.ops.pto_chunk_gdn.mega_kernel import run_mega_kernel + + device = torch.device("npu:0") + B, T, Hg, H, D = 1, 128, 8, 16, 128 + + q = torch.randn(B, T, Hg, D, dtype=torch.float16, device=device) + k = torch.randn(B, T, Hg, D, dtype=torch.float16, device=device) + v = torch.randn(B, T, H, D, dtype=torch.float16, device=device) + g_in = torch.randn(B, T, H, dtype=torch.float32, device=device) + beta = torch.randn(B, T, H, dtype=torch.float16, device=device) + cu_seqlens = torch.tensor([0, T], dtype=torch.int32, device=device) + stream = torch.npu.current_stream()._as_parameter_ + + print("Compiling / loading PTO GDN megakernel ...") + try: + out = run_mega_kernel( + q, + k, + v, + g_in, + beta, + cu_seqlens, + stream=stream, + chunk_size=128, + scale=D**-0.5, + key_heads=Hg, + ) + except Exception as exc: + print(f"FAIL: {type(exc).__name__}: {exc}") + return 1 + + expected_shape = (B, T, H, D) + if out.shape != expected_shape: + print(f"FAIL: output shape {tuple(out.shape)} != expected {expected_shape}") + return 1 + + print(f"OK: PTO GDN megakernel produced output shape {tuple(out.shape)}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/custom_ops_tests/test_recurrent_gated_delta_rule.py b/tests/custom_ops_tests/test_recurrent_gated_delta_rule.py new file mode 100644 index 000000000..028d16528 --- /dev/null +++ b/tests/custom_ops_tests/test_recurrent_gated_delta_rule.py @@ -0,0 +1,66 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Minimal connectivity test for the npu_recurrent_gated_delta_rule framework op. +# +# NOTE: this test automatically enables the CANN custom-op environment by +# discovering the vllm_fl/_cann_ops_custom package installed next to vllm_fl. +# No manual `source set_env.bash` is required. + +import sys + +from vllm_fl.utils import enable_custom_op + +if not enable_custom_op(): + print( + "ERROR: vllm_fl/_cann_ops_custom is not installed.\n" + "Please build and install the CANN framework operators first, e.g.:\n" + " bash csrc/ascend/build_aclnn.sh ", + file=sys.stderr, + ) + sys.exit(1) + +import torch +import torch_npu + +import vllm_fl._C_ascend # noqa: F401 + +DEVICE = "npu:0" + + +def test_recurrent_gated_delta_rule(): + """Basic recurrent gated delta rule connectivity test.""" + t = 16 + nk = 2 + nv = 4 # must be a multiple of nk + dk = 32 + dv = 32 + batch = 2 + + query = torch.randn(t, nk, dk, dtype=torch.bfloat16, device=DEVICE) + key = torch.randn(t, nk, dk, dtype=torch.bfloat16, device=DEVICE) + value = torch.randn(t, nv, dv, dtype=torch.bfloat16, device=DEVICE) + beta = torch.randn(t, nv, dtype=torch.bfloat16, device=DEVICE) + state = torch.randn(batch, nv, dv, dk, dtype=torch.bfloat16, device=DEVICE) + + actual_seq_lengths = torch.full((batch,), t, dtype=torch.int32, device=DEVICE) + ssm_state_indices = torch.arange(batch, dtype=torch.int32, device=DEVICE) + + output = torch.ops._C_ascend.npu_recurrent_gated_delta_rule( + query, + key, + value, + state, + beta=beta, + scale=dk ** -0.5, + actual_seq_lengths=actual_seq_lengths, + ssm_state_indices=ssm_state_indices, + ) + + assert output.shape == value.shape + assert output.device == value.device + + +if __name__ == "__main__": + test_recurrent_gated_delta_rule() + print("npu_recurrent_gated_delta_rule test passed") diff --git a/vllm_fl/__init__.py b/vllm_fl/__init__.py index cdf6d7cbe..e1ed778b5 100644 --- a/vllm_fl/__init__.py +++ b/vllm_fl/__init__.py @@ -5,6 +5,52 @@ import os import logging + +def _bootstrap_cann_custom_op_env(): + """Make the packaged CANN custom-op package discoverable at runtime. + + The AscendC custom ops under ``vllm_fl/_cann_ops_custom`` require + ``ASCEND_CUSTOM_OPP_PATH`` to be set *before* the CANN runtime is + initialized (the op registration/infershape scan happens at aclInit + time), and ``libcust_opapi.so`` to be resolvable by bare-name + ``dlopen`` (which only consults the startup-time ``LD_LIBRARY_PATH``). + Running this at vllm_fl package import time — before torch_npu/CANN + initialization — makes the ops work without sourcing ``set_env.bash``. + The libcust preload uses RTLD_LOCAL: RTLD_GLOBAL causes a double-free + at process teardown. + """ + vendor_dir = os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "_cann_ops_custom", + "vendors", + "custom_transformer", + ) + if not os.path.isdir(vendor_dir): + return + opp_path = os.environ.get("ASCEND_CUSTOM_OPP_PATH", "") + if vendor_dir not in opp_path: + os.environ["ASCEND_CUSTOM_OPP_PATH"] = vendor_dir + ( + ":" + opp_path if opp_path else "" + ) + lib_dir = os.path.join(vendor_dir, "op_api", "lib") + if lib_dir not in os.environ.get("LD_LIBRARY_PATH", ""): + os.environ["LD_LIBRARY_PATH"] = lib_dir + ( + ":" + os.environ["LD_LIBRARY_PATH"] + if os.environ.get("LD_LIBRARY_PATH") + else "" + ) + try: + import ctypes + + ctypes.CDLL( + os.path.join(lib_dir, "libcust_opapi.so"), mode=ctypes.RTLD_LOCAL + ) + except OSError: + pass + + +_bootstrap_cann_custom_op_env() + from vllm_fl.utils import get_op_config as _get_op_config from . import version as version # PyTorch-style: vllm_fl.version.git_version @@ -167,4 +213,36 @@ def register_model(): "vllm_fl.models.bge_m3:BgeM3EmbeddingModel", ) except Exception as e: - logger.error(f"Register BgeM3EmbeddingModel error: {str(e)}") \ No newline at end of file + logger.error(f"Register BgeM3EmbeddingModel error: {str(e)}") + + # Ensure Qwen3NextMTP draft configs are treated as text-only even when + # ModelConfig is reconstructed in worker processes before the platform + # post-init patch is active. + try: + from vllm_fl.dispatch.backends.vendor.ascend.patches.patch_qwen3_mtp import ( + patch_qwen3_next_mtp_multimodal_flag, + ) + patch_qwen3_next_mtp_multimodal_flag() + except Exception as e: + logger.error(f"Patch Qwen3NextMTP multimodal flag error: {str(e)}") + + # Use Ascend-compatible rejection-sampler helpers for speculative decoding. + try: + from vllm_fl.dispatch.backends.vendor.ascend.patches.patch_qwen3_mtp import ( + patch_ascend_rejection_sampler, + ) + patch_ascend_rejection_sampler() + except Exception as e: + logger.error(f"Patch Ascend rejection sampler error: {str(e)}") + + # Honor VLLM_CUDAGRAPH_* environment variables so that scripts can limit + # NPU graph capture sizes without passing --compilation-config. + try: + from vllm_fl.dispatch.backends.vendor.ascend.patches.patch_qwen3_mtp import ( + patch_cudagraph_env_vars, + ) + patch_cudagraph_env_vars() + except Exception as e: + logger.error(f"Apply cudagraph env patch error: {str(e)}") + + return "vllm_fl.platform.PlatformFL" diff --git a/vllm_fl/_version.py b/vllm_fl/_version.py new file mode 100644 index 000000000..4cd74d1c3 --- /dev/null +++ b/vllm_fl/_version.py @@ -0,0 +1,21 @@ +# This file is generated by setuptools-scm at build time. Do not edit. +# At runtime vllm_fl/version.py imports git_version and git_date from here. +from __future__ import annotations +import re as _re + +__version__: str = "0.1.0+vllm0.13.0.ga3644b2d9.d20260718" +version: str = __version__ +_node_date: str = "2026-07-01" + +_sha_m = _re.search(r"(?:^|[.+-])g(?P[0-9a-f]{7,40})(?:[.+-]|$)", __version__) +_date_m = _re.search(r"(?:^|[.+-])d(?P\d{8})(?:[.+-]|$)", __version__) + +git_version: str = _sha_m.group("sha") if _sha_m else "Unknown" +git_date: str = ( + _node_date if _node_date not in ("", "None") + else ( + _date_m.group("ymd")[0:4] + "-" + _date_m.group("ymd")[4:6] + "-" + _date_m.group("ymd")[6:8] + if _date_m else "Unknown" + ) +) +git_info: dict[str, str] = {"id": git_version, "date": git_date} diff --git a/vllm_fl/compilation/graph.py b/vllm_fl/compilation/graph.py index bae44ffdf..d5ea4bde8 100644 --- a/vllm_fl/compilation/graph.py +++ b/vllm_fl/compilation/graph.py @@ -5,7 +5,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import dataclasses -from collections import Counter from collections.abc import Callable from contextlib import ExitStack from typing import Any, Optional @@ -24,14 +23,78 @@ logger = init_logger(__name__) +# --------------------------------------------------------------------------- # +# Backend-specific extension registry for GraphWrapper. +# +# A vendor backend can register a mixin class that implements any subset of the +# hook methods below. The mixin is instantiated once per GraphWrapper and +# receives the same hook calls during capture/replay. This keeps the generic +# framework code free of device-specific details. +# --------------------------------------------------------------------------- # +_graph_wrapper_backend_registry: dict[str, type] = {} + + +def register_graph_wrapper_backend(device_type: str, backend_cls: type) -> None: + """Register a backend-specific mixin for GraphWrapper. + + The mixin class may implement any of the following hook methods: + + - before_capture(self, entry: GraphEntry, args, kwargs) -> None + Called after input_addresses are recorded but before the graph capture + context is entered. Useful for stream synchronization or patching the + runtime environment. + + - wrap_capture_context(self, entry: GraphEntry, stack: ExitStack) -> None + Called inside the ExitStack used around graph capture. The backend can + enter additional context managers (e.g. disable gc / empty_cache). + + - after_capture(self, entry: GraphEntry, output, args, kwargs) -> Any + Called after the graph capture block exits but before the entry is + stored. The return value replaces `output` if not None. + + - before_replay(self, entry: GraphEntry, args, kwargs) -> None + Called before graph replay. Useful for device synchronization. + + - capture_error_handler(self, exc: BaseException) -> None + Called when graph capture raises. May translate the exception or raise + a more informative error. + + - weak_ref_tensors(self, tensor: Any) -> Any + Return a weak-ref version of tensors for the current backend. If not + implemented, falls back to the generic helper. + """ + _graph_wrapper_backend_registry[device_type] = backend_cls + logger.info_once("Registered graph wrapper backend for device_type=%s: %s", + device_type, backend_cls.__name__) + + +def get_graph_wrapper_backend(device_type: str) -> Optional[type]: + return _graph_wrapper_backend_registry.get(device_type) + def weak_ref_tensors(tensor: Any) -> Any: + backend = _get_active_backend() + if backend is not None and hasattr(backend, "weak_ref_tensors"): + return backend.weak_ref_tensors(tensor) if current_platform.device_type == "cuda": from vllm.utils.torch_utils import weak_ref_tensors return weak_ref_tensors(tensor) - else: - ### TODO: add csrc npu custom op - return tensor + # TODO: add csrc npu custom op when available + return tensor + + +# Per-wrapper active backend instance. This is set by GraphWrapper.__init__ and +# used by the weak_ref_tensors helper so callers do not need to pass `self`. +_active_backend_instance: Optional[Any] = None + + +def _set_active_backend(backend: Optional[Any]) -> None: + global _active_backend_instance + _active_backend_instance = backend + + +def _get_active_backend() -> Optional[Any]: + return _active_backend_instance class Graph: @@ -46,6 +109,7 @@ class Graph: else: raise NotImplementedError("not support graph") + @dataclasses.dataclass class GraphEntry: batch_descriptor: BatchDescriptor @@ -56,6 +120,7 @@ class GraphEntry: # during capture, and check if they are the same during replay input_addresses: Optional[list[int]] = None + @dataclasses.dataclass class GraphOptions: debug_log_enable: bool = True @@ -92,6 +157,13 @@ def __init__(self, # cudagraphs for. self.concrete_graph_entries: dict[BatchDescriptor, GraphEntry] = {} + # Instantiate backend-specific mixin if one has been registered. + backend_cls = get_graph_wrapper_backend(current_platform.device_type) + if backend_cls is not None: + self.backend = backend_cls(self) + else: + self.backend = None + def __getattr__(self, key: str): # allow accessing the attributes of the runnable. if hasattr(self.runnable, key): @@ -150,6 +222,14 @@ def __call__(self, *args, **kwargs): entry.input_addresses = input_addresses graph = Graph.graph() + _set_active_backend(self.backend) + + # Give the backend a chance to run pre-capture logic (e.g. stream + # sync, offloader sync). + if self.backend is not None and hasattr(self.backend, + "before_capture"): + self.backend.before_capture(entry, args, kwargs) + with ExitStack() as stack: if self.graph_options.gc_disable: # during every model forward for piecewise graph @@ -160,29 +240,54 @@ def __call__(self, *args, **kwargs): # and disable gc for the rest of the graphs. stack.enter_context(patch("gc.collect", lambda: None)) stack.enter_context( - patch("vllm_fl.platform.PlatformFL.empty_cache", lambda: None) + patch("vllm_fl.platform.PlatformFL.empty_cache", + lambda: None) ) - set_graph_pool_id(self.graph_pool) - - # mind-exploding: carefully manage the reference and memory. - with current_platform.torch_device_fn.graph(graph, pool=self.graph_pool): - # `output` is managed by pytorch's cudagraph pool - output = self.runnable(*args, **kwargs) - if self.graph_options.weak_ref_output: - # by converting it to weak ref, - # the original `output` will immediately be released - # to save memory. It is only safe to do this for - # the last graph in piecewise cuadgraph mode, because - # the output of the last graph will not be used by - # any other cuda graph. - output = weak_ref_tensors(output) + # Backend-specific context wrappers (e.g. disable NPU + # empty_cache when it lives on a different module path). + if self.backend is not None and hasattr( + self.backend, "wrap_capture_context"): + self.backend.wrap_capture_context(entry, stack) + + set_graph_pool_id(self.graph_pool) + + # mind-exploding: carefully manage the reference and memory. + try: + with current_platform.torch_device_fn.graph( + graph, pool=self.graph_pool): + # `output` is managed by pytorch's cudagraph pool + output = self.runnable(*args, **kwargs) + if self.graph_options.weak_ref_output: + # by converting it to weak ref, + # the original `output` will immediately be released + # to save memory. It is only safe to do this for + # the last graph in piecewise cuadgraph mode, because + # the output of the last graph will not be used by + # any other cuda graph. + output = weak_ref_tensors(output) + except BaseException as exc: + if self.backend is not None and hasattr( + self.backend, "capture_error_handler"): + self.backend.capture_error_handler(exc) + raise + + # Backend-specific post-capture logic (weak-ref workspaces, + # offloader join, etc.). The return value may replace `output`. + if self.backend is not None and hasattr(self.backend, + "after_capture"): + backend_output = self.backend.after_capture( + entry, output, args, kwargs) + if backend_output is not None: + output = backend_output entry.output = weak_ref_tensors(output) entry.graph = graph compilation_counter.num_cudagraph_captured += 1 + _set_active_backend(None) + # important: we need to return the output, rather than # the weak ref of the output, so that pytorch can correctly # manage the memory during graph capture @@ -199,6 +304,15 @@ def __call__(self, *args, **kwargs): f"got {new_input_addresses}" ) - current_platform.torch_device_fn.synchronize() + _set_active_backend(self.backend) + + if self.backend is not None and hasattr(self.backend, + "before_replay"): + self.backend.before_replay(entry, args, kwargs) + else: + current_platform.torch_device_fn.synchronize() + entry.graph.replay() + + _set_active_backend(None) return entry.output diff --git a/vllm_fl/dispatch/backends/vendor/ascend/__init__.py b/vllm_fl/dispatch/backends/vendor/ascend/__init__.py index 1200f348d..78c268d3f 100644 --- a/vllm_fl/dispatch/backends/vendor/ascend/__init__.py +++ b/vllm_fl/dispatch/backends/vendor/ascend/__init__.py @@ -6,7 +6,9 @@ from .ascend import AscendBackend from .patch import patch_mamba_config +from .patches.patch_qwen3_mtp import patch_qwen3_mtp_platform patch_mamba_config() +patch_qwen3_mtp_platform() __all__ = ["AscendBackend"] diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/__init__.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/compiler_interface.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/compiler_interface.py new file mode 100644 index 000000000..a6fe93e27 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/compilation/compiler_interface.py @@ -0,0 +1,75 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Adapted from https://github.com/vllm-project/vllm-ascend/blob/main/vllm_ascend/compilation/compiler_interface.py +# Below is the original copyright: +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Ascend-specific compiler interface for vllm-plugin-FL. + +This module provides the CompilerInterface subclass used when Ascend graph +compilation (npugraph_ex / torchair) is enabled. By default vllm-plugin-FL +falls back to eager mode on NPU, so this class is only instantiated when the +user explicitly enables `ascend_compilation_config.enable_npugraph_ex`. +""" + +from __future__ import annotations + +import logging +from typing import Any, Callable + +import torch + +from vllm.compilation.compiler_interface import CompilerInterface +from vllm.config import VllmConfig + +logger = logging.getLogger(__name__) + + +class AscendCompiler(CompilerInterface): + """ + Ascend compiler interface. + + When npugraph_ex is available, this compiler delegates graph compilation to + it. Otherwise it raises a clear error directing users to eager mode. + """ + + def __init__(self) -> None: + super().__init__() + self._nge = None + try: + import npugraph_ex as nge + self._nge = nge + self._use_npugraph_ex = True + except ImportError: + try: + import torchair as nge + self._nge = nge + self._use_npugraph_ex = False + except ImportError as exc: + raise ImportError( + "npugraph_ex or torchair is required for AscendCompiler. " + "Either install it or disable ascend_compilation_config." + ) from exc + + def compute_hash(self, vllm_config: VllmConfig) -> str: + import torch_npu + return f"{torch_npu.__version__}_{self._use_npugraph_ex}" + + def initialize_cache(self, cache_dir: str, *args, **kwargs) -> None: + logger.info("AscendCompiler cache dir: %s", cache_dir) + + def compile( + self, + graph: Callable, + example_inputs: list[torch.Tensor], + additional_inductor_config: dict, + rank: int = 0, + ) -> Any: + raise NotImplementedError( + "AscendCompiler.compile is not yet implemented in vllm-plugin-FL. " + "Use backend='eager' (default on NPU) for cudagraph-only execution.") + + def load(self, path: str) -> Any: + raise NotImplementedError( + "AscendCompiler.load is not yet implemented in vllm-plugin-FL.") diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/graph_fusion_pass_manager.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/graph_fusion_pass_manager.py new file mode 100644 index 000000000..8f85ec9eb --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/compilation/graph_fusion_pass_manager.py @@ -0,0 +1,91 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from torch import fx as fx +from vllm.compilation.inductor_pass import get_pass_context +from vllm.compilation.vllm_inductor_pass import VllmInductorPass +from vllm.config import VllmConfig + + +def _is_310p() -> bool: + """Return True if the current Ascend device is a 310P. + + In vllm-plugin-FL we do not depend on vllm-ascend internals, so default to + False. Vendors can override this helper if they need 310P-specific pass + disabling. + """ + try: + from vllm_ascend.utils import is_310p + return is_310p() + except Exception: + return False + + +class GraphFusionPassManager: + """ + A pass manager for graph fusion passes. + It handles the configuration and execution of passes. + The counterpart in vllm is PostGradPassManager. Since torch_npu + does not support triton for now, we define our own pass manager. + """ + + def __init__(self): + self.passes: list[VllmInductorPass] = [] + + def __call__(self, graph: fx.Graph) -> fx.Graph: + compile_range = get_pass_context().compile_range + + for pass_ in self.passes: + if pass_.is_applicable_for_range(compile_range): + pass_(graph) + graph.recompile() + return graph + + def add(self, pass_: VllmInductorPass): + assert isinstance(pass_, VllmInductorPass) + self.passes.append(pass_) + + def configure(self, config: VllmConfig): + # By default, we enable the graph fusion and quantization fusion pass. + self.ascend_compilation_config: dict = config.additional_config.get("ascend_compilation_config", {}) + if self.ascend_compilation_config.get("fuse_norm_quant", True) and not _is_310p(): + from .passes.norm_quant_fusion_pass import AddRMSNormQuantFusionPass + + self.passes.append(AddRMSNormQuantFusionPass(config)) + + if self.ascend_compilation_config.get("fuse_qknorm_rope", True): + from .passes.qknorm_rope_fusion_pass import QKNormRopeFusionPass + + self.passes.append(QKNormRopeFusionPass(config)) + + if self.ascend_compilation_config.get("fuse_allreduce_rms", True): + from .passes.allreduce_rmsnorm_fusion_pass import MatmulAllReduceAddRMSNormPass + + self.passes.append(MatmulAllReduceAddRMSNormPass(config)) + + if self.ascend_compilation_config.get("fuse_muls_add", True) and not _is_310p(): + from .passes.muls_add_pass import MulsAddFusionPass + + self.passes.append(MulsAddFusionPass(config)) + + if config.compilation_config.pass_config.enable_sp: + from .passes.sequence_parallelism import SequenceParallelismPass + from .passes.sequence_parallelism_moe import SequenceParallelismMoePass + + self.passes.append(SequenceParallelismPass(config)) + self.passes.append(SequenceParallelismMoePass(config)) diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/__init__.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/allgather_chunk_noop_pass.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/allgather_chunk_noop_pass.py new file mode 100644 index 000000000..68b2aae66 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/allgather_chunk_noop_pass.py @@ -0,0 +1,40 @@ +import torch +import torch._inductor.pattern_matcher as pm +from torch._inductor.pattern_matcher import PatternMatcherPass +from vllm.compilation.vllm_inductor_pass import VllmInductorPass +from vllm.config import VllmConfig +from vllm.distributed import get_tensor_model_parallel_world_size, get_tp_group +from vllm.logger import logger + + +class AllGatherChunkNoOpCleanupPass(VllmInductorPass): + """Fold all_gather + sequence_parallel_chunk_impl into identity.""" + + def __init__(self, config: VllmConfig): + super().__init__(config) + self.tp_group = get_tp_group() + self.tp_size = get_tensor_model_parallel_world_size() + self.patterns: PatternMatcherPass = PatternMatcherPass(pass_name="npu_allgather_chunk_noop_cleanup_pass") + self._register_patterns() + + def _all_gather(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.vllm.all_gather(x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name) + + def _empty(self, *args, **kwargs): + return torch.empty(*args, dtype=self.model_dtype, device=self.device, **kwargs) + + def _register_patterns(self) -> None: + def pattern(input: torch.Tensor) -> torch.Tensor: + gathered = self._all_gather(input) + return torch.ops.vllm.sequence_parallel_chunk_impl(gathered) + + def replacement(input: torch.Tensor) -> torch.Tensor: + return input + + pm.register_replacement(pattern, replacement, [self._empty(8, 16)], pm.fwd_only, self.patterns) + + def __call__(self, graph: torch.fx.Graph) -> None: + self.begin() + matched_count = self.patterns.apply(graph) + logger.debug("AllGatherChunkNoOpCleanupPass replaced %s patterns", matched_count) + self.end_and_log() diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/allreduce_rmsnorm_fusion_pass.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/allreduce_rmsnorm_fusion_pass.py new file mode 100644 index 000000000..ba2b7c1c3 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/allreduce_rmsnorm_fusion_pass.py @@ -0,0 +1,159 @@ +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import torch +from torch._inductor.pattern_matcher import PatternMatcherPass, PatternPrettyPrinter +from vllm.compilation.vllm_inductor_pass import VllmInductorPass +from vllm.config import VllmConfig +from vllm.config.compilation import Range +from vllm.distributed import get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce +from vllm.distributed.parallel_state import get_tp_group +from vllm.logger import logger + +from vllm_ascend.compilation.passes.base_pattern import BasePattern + +# computation-communication tiling block is 512 +ALLREDUCE_NORM_FUSE_THRESHOLD = 512 + + +class MiddleLayerMatmulAllReduceAddRMSNormPattern(BasePattern): + """ + recognizing the Matmul+AllReduce+AddRMSNorm computation pattern + AllReduce is optimized in the fusion operator to a two-stage communication of ReduceScatter+AllGather + """ + + def __init__(self, vllm_config, eps=1e-6): + self.vllm_config = vllm_config + self.eps = eps + device_group = get_tp_group().device_group + backend = device_group._get_backend(torch.device("npu")) + self.local_rank = torch.distributed.get_rank(group=device_group) + self.tp_group_name = backend.get_hccl_comm_name(self.local_rank) + self.tp_size = get_tensor_model_parallel_world_size() + + def get_inputs(self): + batch_size, seq_len = 2, 4 + hidden_size = 4096 + x = torch.randn(batch_size, seq_len, hidden_size, device="npu") + weight = torch.randn(hidden_size, hidden_size, device="npu") + residual = torch.randn(batch_size, seq_len, hidden_size, device="npu") + rms_norm_weight = torch.randn(hidden_size, device="npu") + return [x, weight, residual, rms_norm_weight] + + def get_pattern(self): + def pattern(x, weight, residual, rms_norm_weight): + mm = torch.ops.vllm.unquantized_gemm(x, weight, None) + all_reduce_ = tensor_model_parallel_all_reduce(mm) + chunked_residual = torch.ops.vllm.maybe_chunk_residual(all_reduce_, residual) + output = torch.ops._C_ascend.npu_add_rms_norm_bias(all_reduce_, chunked_residual, rms_norm_weight, None) + out0 = output[0] + out1 = output[2] + return out0, out1 + + return pattern + + def get_replacement(self): + def replacement(x, weight, residual, rms_norm_weight): + out0, out1 = torch.ops._C_ascend.matmul_allreduce_add_rmsnorm( + x, + weight, + residual, + rms_norm_weight, + self.tp_group_name, + self.tp_size, + self.local_rank, + self.eps, + True, + False, + ) + return out0, out1 + + return replacement + + +class LastLayerMatmulAllReduceAddRMSNormPattern(BasePattern): + def __init__(self, vllm_config, eps=1e-6): + super().__init__(vllm_config, eps) + device_group = get_tp_group().device_group + backend = device_group._get_backend(torch.device("npu")) + self.local_rank = torch.distributed.get_rank(group=device_group) + self.tp_group_name = backend.get_hccl_comm_name(self.local_rank) + self.tp_size = get_tensor_model_parallel_world_size() + + def get_inputs(self): + batch_size, seq_len = 2, 4 + hidden_size = 4096 + x = torch.randn(batch_size, seq_len, hidden_size, device="npu") + weight = torch.randn(hidden_size, hidden_size, device="npu") + residual = torch.randn(batch_size, seq_len, hidden_size, device="npu") + rms_norm_weight = torch.randn(hidden_size, device="npu") + return [x, weight, residual, rms_norm_weight] + + def get_pattern(self): + def pattern(x, weight, residual, rms_norm_weight): + mm = torch.ops.vllm.unquantized_gemm(x, weight, None) + all_reduce_ = tensor_model_parallel_all_reduce(mm) + chunked_residual = torch.ops.vllm.maybe_chunk_residual(all_reduce_, residual) + output = torch.ops._C_ascend.npu_add_rms_norm_bias(all_reduce_, chunked_residual, rms_norm_weight, None) + return output[0] + + return pattern + + def get_replacement(self): + def replacement(x, weight, residual, rms_norm_weight): + out0, _ = torch.ops._C_ascend.matmul_allreduce_add_rmsnorm( + x, + weight, + residual, + rms_norm_weight, + self.tp_group_name, + self.tp_size, + self.local_rank, + self.eps, + True, + False, + ) + return out0 + + return replacement + + +class MatmulAllReduceAddRMSNormPass(VllmInductorPass): + def __init__(self, vllm_config: VllmConfig): + super().__init__(vllm_config) + self.pattern_match_passes: PatternMatcherPass = PatternMatcherPass(pass_name="allreduce_rmsnorm_fusion_pass") + + MiddleLayerMatmulAllReduceAddRMSNormPattern(vllm_config).register(self.pattern_match_passes) + LastLayerMatmulAllReduceAddRMSNormPattern(vllm_config).register(self.pattern_match_passes) + + def __call__(self, graph: torch.fx.Graph): + self.begin() + self.matched_count = self.pattern_match_passes.apply(graph) + pattern_idx = 0 + for pattern_entry in self.pattern_match_passes.patterns.values(): + for p in pattern_entry: + p_str = PatternPrettyPrinter.run(p.pattern) + logger.debug("Pattern %d: %s", pattern_idx, p_str) + pattern_idx += 1 + logger.debug("Replaced %s patterns", self.matched_count) + self.end_and_log() + + def is_applicable_for_range(self, compile_range: Range) -> bool: + """ + Check if the pass is applicable for the current configuration. + """ + applicable = compile_range.start > ALLREDUCE_NORM_FUSE_THRESHOLD + return applicable diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/base_pattern.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/base_pattern.py new file mode 100644 index 000000000..f91c2fd56 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/base_pattern.py @@ -0,0 +1,63 @@ +from abc import ABC, abstractmethod +from collections.abc import Callable + +import torch +import torch._inductor.pattern_matcher as pm +from torch._inductor.pattern_matcher import PatternMatcherPass +from vllm.config import VllmConfig + +try: + import npugraph_ex as nge +except ImportError: + import torchair as nge + +from vllm_ascend.compilation.passes.utils.npugraph_ex_utils_check import extra_stream_scope_check + +# Global set to track registered patterns and prevent duplicates +_registered_patterns: set[str] = set() + + +class BasePattern(ABC): + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + self.vllm_config = vllm_config + self.dtype = vllm_config.model_config.dtype + self.eps = eps + + @abstractmethod + def get_inputs(self) -> list[torch.Tensor]: + pass + + @abstractmethod + def get_pattern(self) -> Callable: + pass + + @abstractmethod + def get_replacement(self) -> Callable: + pass + + def get_extra_stream_scope_check(self): + return extra_stream_scope_check + + def register(self, pm_pass: PatternMatcherPass) -> None: + # Create a unique identifier for this pattern based on class name and eps + pattern_id = f"{self.__class__.__name__}_{self.eps}" + + # Skip registration if this pattern has already been registered globally + if pattern_id in _registered_patterns: + return + + pattern_fn = self.get_pattern() + replacement_fn = self.get_replacement() + example_inputs = self.get_inputs() + + pm.register_replacement(pattern_fn, replacement_fn, example_inputs, pm.fwd_only, pm_pass) + + nge.register_replacement( + search_fn=pattern_fn, + replace_fn=replacement_fn, + example_inputs=example_inputs, + extra_check=self.get_extra_stream_scope_check(), + ) + + # Mark this pattern as registered + _registered_patterns.add(pattern_id) diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/muls_add_pass.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/muls_add_pass.py new file mode 100644 index 000000000..1cbf8c3de --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/muls_add_pass.py @@ -0,0 +1,110 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from __future__ import annotations + +import torch +from torch._inductor.pattern_matcher import PatternMatcherPass +from vllm.compilation.vllm_inductor_pass import VllmInductorPass +from vllm.config import VllmConfig +from vllm.config.compilation import Range +from vllm.logger import logger + +from vllm_ascend.compilation.passes.base_pattern import BasePattern + + +class MulsAddPattern(BasePattern): + """ + Pattern that matches an element-wise mul + add sequence: + tmp = x * scale + out = tmp + y + and replaces it with a call to the muls_add_triton kernel. + """ + + def __init__(self, vllm_config: VllmConfig, scale: float = 1.0): + super().__init__(vllm_config) + self.scale = scale + + def get_inputs(self) -> list[torch.Tensor]: + """ + Generate example inputs for the MulsAddPattern. + + The exact shapes are not important for pattern matching; they only + provide meta information for the pattern matcher. + """ + x = torch.randn(2, 2048, device="npu", dtype=self.dtype) + y = torch.randn(2, 2048, device="npu", dtype=self.dtype) + # Only tensor inputs are needed here. The scalar scale is stored on the + # pattern instance (self.scale) instead of being passed as an input. + return [x, y] + + def get_pattern(self): + def pattern(x: torch.Tensor, y: torch.Tensor): + """ + Pattern for element-wise x * scale + y. + """ + tmp = x * self.scale + out = tmp + y + return out + + return pattern + + def get_replacement(self): + def replacement(x: torch.Tensor, y: torch.Tensor): + """ + Replacement that calls the muls_add_triton kernel using the + class-level scalar self.scale. + """ + return torch.ops.vllm.muls_add(x, y, self.scale) + + return replacement + + +class MulsAddFusionPass(VllmInductorPass): + """ + A fusion pass that replaces simple element-wise x * scale + y patterns + with the Triton-based muls_add_triton kernel on Ascend. + """ + + def __init__(self, vllm_config: VllmConfig): + super().__init__(vllm_config) + self.pattern_match_passes: PatternMatcherPass = PatternMatcherPass(pass_name="muls_add_fusion_pass") + + # For now we enable this pass for all floating-point dtypes that the + # model is configured to use. + dtype = vllm_config.model_config.dtype + if dtype not in (torch.float16, torch.bfloat16, torch.float32): + logger.debug("MulsAdd fusion not enabled: unsupported dtype %s", dtype) + return + + routed_scaling_factor = getattr(vllm_config.model_config.hf_text_config, "routed_scaling_factor", 1.0) + MulsAddPattern(vllm_config, scale=routed_scaling_factor).register(self.pattern_match_passes) + + def __call__(self, graph: torch.fx.Graph) -> None: # type: ignore[override] + self.begin() + self.matched_count = self.pattern_match_passes.apply(graph) + logger.debug("Fused %s muls_add patterns", self.matched_count) + self.end_and_log() + + def is_applicable_for_range(self, compile_range: Range) -> bool: + """ + Check if the pass is applicable for the current configuration. + + For now, muls_add fusion is always allowed for the selected ranges. + This hook exists so that we can add more fine-grained range control + in the future if needed. + """ + return True diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/noop_elimination.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/noop_elimination.py new file mode 100644 index 000000000..9f58605c5 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/noop_elimination.py @@ -0,0 +1,62 @@ +from collections.abc import Iterable + +import torch +import torch.fx +from torch import SymInt +from torch.fx.experimental.symbolic_shapes import statically_known_true +from vllm.compilation.vllm_inductor_pass import VllmInductorPass +from vllm.logger import logger + + +class NoOpEliminationPass(VllmInductorPass): + """Remove no-op view/reshape nodes after pattern rewrites.""" + + def __call__(self, graph: torch.fx.Graph) -> None: + fx_graph = graph.graph if hasattr(graph, "graph") else graph + removed = 0 + for node in list(fx_graph.nodes): + if not self._is_view_like(node): + continue + + input_node = node.args[0] + if not isinstance(input_node, torch.fx.Node): + continue + + input_meta = input_node.meta.get("val") + output_meta = node.meta.get("val") + if input_meta is None or output_meta is None: + continue + + input_shape = getattr(input_meta, "shape", None) + output_shape = getattr(output_meta, "shape", None) + if input_shape is None or output_shape is None: + continue + + if self._all_dims_equivalent(input_shape, output_shape): + node.replace_all_uses_with(input_node) + fx_graph.erase_node(node) + removed += 1 + + logger.debug("NoOpEliminationPass removed %s no-op views", removed) + + @staticmethod + def _is_view_like(node: torch.fx.Node) -> bool: + return (node.op == "call_method" and node.target in {"view", "reshape"}) or ( + node.op == "call_function" + and node.target + in { + torch.ops.aten.view.default, + torch.ops.aten.reshape.default, + } + ) + + @staticmethod + def _dims_equivalent(dim: int | SymInt, i_dim: int | SymInt) -> bool: + return statically_known_true(dim == i_dim) # type: ignore[no-any-return] + + def _all_dims_equivalent(self, dims: Iterable[int | SymInt], i_dims: Iterable[int | SymInt]) -> bool: + dims_ = list(dims) + i_dims_ = list(i_dims) + if len(dims_) != len(i_dims_): + return False + return all(self._dims_equivalent(s, i_s) for s, i_s in zip(dims_, i_dims_)) diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/norm_quant_fusion_pass.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/norm_quant_fusion_pass.py new file mode 100644 index 000000000..f1f6ecf79 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/norm_quant_fusion_pass.py @@ -0,0 +1,513 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import torch +from torch._inductor.pattern_matcher import PatternMatcherPass +from vllm.compilation.vllm_inductor_pass import VllmInductorPass +from vllm.config import VllmConfig +from vllm.config.compilation import Range +from vllm.logger import logger + +from vllm_ascend.compilation.passes.base_pattern import BasePattern +from vllm_ascend.utils import enable_custom_op + + +class AddRMSNormQuantPattern(BasePattern): + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(vllm_config, eps) + + def get_inputs(self): + """ + Generate example inputs for the AddRMSNormQuant fusion pattern. + """ + rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype) + residual = torch.randn(2, 4, device="npu", dtype=self.dtype) + rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype) + scale = torch.ones(4, device="npu", dtype=self.dtype) + scale_reciprocal = torch.ones(4, device="npu", dtype=self.dtype) + offset = torch.zeros(4, device="npu", dtype=self.dtype) + return [rms_norm_input, residual, rms_norm_weight, scale, scale_reciprocal, offset] + + def get_pattern(self): + def pattern( + rms_norm_input: torch.Tensor, + residual: torch.Tensor, + rms_norm_weight: torch.Tensor, + scale: torch.Tensor, + scale_reciprocal: torch.Tensor, + offset: torch.Tensor, + ): + """ + Pattern for AddRMSNormQuant fusion. + """ + output = torch.ops._C_ascend.npu_add_rms_norm_bias( + rms_norm_input, residual, rms_norm_weight, None, self.eps + ) + out0 = output[0] + out1 = output[2] + quantized_output = torch.ops.vllm.quantize(out0, scale, scale_reciprocal, offset) + return quantized_output, out1 + + return pattern + + def get_replacement(self): + def replacement( + rms_norm_input: torch.Tensor, + residual: torch.Tensor, + rms_norm_weight: torch.Tensor, + scale: torch.Tensor, + scale_reciprocal: torch.Tensor, + offset: torch.Tensor, + ): + """ + Replacement for the AddRMSNormQuant fusion. + """ + output = torch.ops.npu.npu_add_rms_norm_quant( + rms_norm_input, residual, rms_norm_weight, scale, offset, epsilon=self.eps + ) + quantized_output = output[0] + out1 = output[2] + return quantized_output, out1 + + return replacement + + +class AddRMSNormQuantPatternWithBias(BasePattern): + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(vllm_config, eps) + + def get_inputs(self): + """ + Generate example inputs for the AddRMSNormQuant fusion pattern. + """ + rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype) + residual = torch.randn(2, 4, device="npu", dtype=self.dtype) + rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype) + rmsnorm_bias = torch.randn(4, device="npu", dtype=self.dtype) + scale = torch.ones(4, device="npu", dtype=self.dtype) + scale_reciprocal = torch.ones(4, device="npu", dtype=self.dtype) + offset = torch.zeros(4, device="npu", dtype=self.dtype) + return [rms_norm_input, residual, rms_norm_weight, scale, scale_reciprocal, offset, rmsnorm_bias] + + def get_pattern(self): + def pattern( + rms_norm_input: torch.Tensor, + residual: torch.Tensor, + rms_norm_weight: torch.Tensor, + scale: torch.Tensor, + scale_reciprocal: torch.Tensor, + offset: torch.Tensor, + bias: torch.Tensor, + ): + """ + Pattern for AddRMSNormQuant fusion. + """ + output = torch.ops._C_ascend.npu_add_rms_norm_bias( + rms_norm_input, residual, rms_norm_weight, bias, self.eps + ) + out0 = output[0] + out1 = output[2] + quantized_output = torch.ops.vllm.quantize(out0, scale, scale_reciprocal, offset) + return quantized_output, out1 + + return pattern + + def get_replacement(self): + def replacement( + rms_norm_input: torch.Tensor, + residual: torch.Tensor, + rms_norm_weight: torch.Tensor, + scale: torch.Tensor, + scale_reciprocal: torch.Tensor, + offset: torch.Tensor, + bias: torch.Tensor, + ): + """ + Replacement for the AddRMSNormQuant fusion. + """ + output = torch.ops.npu.npu_add_rms_norm_quant( + rms_norm_input, residual, rms_norm_weight, scale, offset, epsilon=self.eps, beta=bias + ) + quantized_output = output[0] + out1 = output[2] + return quantized_output, out1 + + return replacement + + +class AddRMSNormQuantSPPattern(BasePattern): + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(vllm_config, eps) + + def get_inputs(self): + """ + Generate example inputs for the AddRMSNormQuant fusion pattern. + """ + rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype) + residual = torch.randn(2, 4, device="npu", dtype=self.dtype) + rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype) + scale = torch.ones(4, device="npu", dtype=self.dtype) + scale_reciprocal = torch.ones(4, device="npu", dtype=self.dtype) + offset = torch.zeros(4, device="npu", dtype=self.dtype) + return [rms_norm_input, residual, rms_norm_weight, scale, scale_reciprocal, offset] + + def get_pattern(self): + def pattern( + rms_norm_input: torch.Tensor, + residual: torch.Tensor, + rms_norm_weight: torch.Tensor, + scale: torch.Tensor, + scale_reciprocal: torch.Tensor, + offset: torch.Tensor, + ): + """ + Pattern for AddRMSNormQuant fusion. + """ + output = torch.ops._C_ascend.npu_add_rms_norm_bias( + rms_norm_input, residual, rms_norm_weight, None, self.eps + ) + out0 = output[0] + out1 = output[2] + out0 = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(out0, True) + quantized_output = torch.ops.vllm.quantize(out0, scale, scale_reciprocal, offset) + return quantized_output, out1 + + return pattern + + def get_replacement(self): + def replacement( + rms_norm_input: torch.Tensor, + residual: torch.Tensor, + rms_norm_weight: torch.Tensor, + scale: torch.Tensor, + scale_reciprocal: torch.Tensor, + offset: torch.Tensor, + ): + """ + Replacement for the AddRMSNormQuant fusion. + """ + output = torch.ops.npu.npu_add_rms_norm_quant( + rms_norm_input, residual, rms_norm_weight, scale, offset, epsilon=self.eps + ) + quantized_output = output[0] + out1 = output[2] + quantized_output = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(quantized_output, True) + return quantized_output, out1 + + return replacement + + +class AddRMSNormQuantSPPatternWithBias(BasePattern): + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(vllm_config, eps) + + def get_inputs(self): + """ + Generate example inputs for the AddRMSNormQuant fusion pattern. + """ + rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype) + residual = torch.randn(2, 4, device="npu", dtype=self.dtype) + rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype) + rmsnorm_bias = torch.randn(4, device="npu", dtype=self.dtype) + scale = torch.ones(4, device="npu", dtype=self.dtype) + scale_reciprocal = torch.ones(4, device="npu", dtype=self.dtype) + offset = torch.zeros(4, device="npu", dtype=self.dtype) + return [rms_norm_input, residual, rms_norm_weight, scale, scale_reciprocal, offset, rmsnorm_bias] + + def get_pattern(self): + def pattern( + rms_norm_input: torch.Tensor, + residual: torch.Tensor, + rms_norm_weight: torch.Tensor, + scale: torch.Tensor, + scale_reciprocal: torch.Tensor, + offset: torch.Tensor, + bias: torch.Tensor, + ): + """ + Pattern for AddRMSNormQuant fusion. + """ + output = torch.ops._C_ascend.npu_add_rms_norm_bias( + rms_norm_input, residual, rms_norm_weight, bias, self.eps + ) + out0 = output[0] + out1 = output[2] + out0 = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(out0, True) + quantized_output = torch.ops.vllm.quantize(out0, scale, scale_reciprocal, offset) + return quantized_output, out1 + + return pattern + + def get_replacement(self): + def replacement( + rms_norm_input: torch.Tensor, + residual: torch.Tensor, + rms_norm_weight: torch.Tensor, + scale: torch.Tensor, + scale_reciprocal: torch.Tensor, + offset: torch.Tensor, + bias: torch.Tensor, + ): + """ + Replacement for the AddRMSNormQuant fusion. + """ + output = torch.ops.npu.npu_add_rms_norm_quant( + rms_norm_input, residual, rms_norm_weight, scale, offset, epsilon=self.eps, beta=bias + ) + quantized_output = output[0] + out1 = output[2] + quantized_output = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(quantized_output, True) + return quantized_output, out1 + + return replacement + + +class AddRMSNormDynamicQuantPattern(BasePattern): + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(vllm_config, eps) + + def get_inputs(self): + """ + Generate example inputs for the AddRMSNormQuant fusion pattern. + """ + rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype) + residual = torch.randn(2, 4, device="npu", dtype=self.dtype) + rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype) + return [rms_norm_input, residual, rms_norm_weight] + + def get_pattern(self): + def pattern(rms_norm_input: torch.Tensor, residual: torch.Tensor, rms_norm_weight: torch.Tensor): + """ + Pattern for AddRMSNormQuant fusion. + """ + output = torch.ops.npu.npu_add_rms_norm(rms_norm_input, residual, rms_norm_weight, self.eps) + out0 = output[0] + out1 = output[2] + quantized_output = torch.ops.npu.npu_dynamic_quant(out0) + return quantized_output[0], quantized_output[1], out1 + + return pattern + + def get_replacement(self): + def replacement(rms_norm_input: torch.Tensor, residual: torch.Tensor, rms_norm_weight: torch.Tensor): + """ + Replacement for the AddRMSNormQuant fusion. + """ + output = torch.ops.npu.npu_add_rms_norm_dynamic_quant( + rms_norm_input, residual, rms_norm_weight, epsilon=self.eps, output_mask=[True, False] + ) + return ( + output[0], + output[3], + output[2], + ) + + return replacement + + +class AddRMSNormDynamicQuantPatternWithBias(BasePattern): + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(vllm_config, eps) + + def get_inputs(self): + """ + Generate example inputs for the AddRMSNormQuant fusion pattern. + """ + rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype) + residual = torch.randn(2, 4, device="npu", dtype=self.dtype) + rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype) + rmsnorm_bias = torch.randn(4, device="npu", dtype=self.dtype) + return [rms_norm_input, residual, rms_norm_weight, rmsnorm_bias] + + def get_pattern(self): + def pattern( + rms_norm_input: torch.Tensor, + residual: torch.Tensor, + rms_norm_weight: torch.Tensor, + bias: torch.Tensor, + ): + """ + Pattern for AddRMSNormQuant fusion. + """ + output = torch.ops._C_ascend.npu_add_rms_norm_bias( + rms_norm_input, residual, rms_norm_weight, bias, self.eps + ) + out0 = output[0] + out1 = output[2] + quantized_output = torch.ops.npu.npu_dynamic_quant(out0) + return quantized_output[0], quantized_output[1], out1 + + return pattern + + def get_replacement(self): + def replacement( + rms_norm_input: torch.Tensor, + residual: torch.Tensor, + rms_norm_weight: torch.Tensor, + bias: torch.Tensor, + ): + """ + Replacement for the AddRMSNormQuant fusion. + """ + output = torch.ops.npu.npu_add_rms_norm_dynamic_quant( + rms_norm_input, residual, rms_norm_weight, epsilon=self.eps, output_mask=[True, False], beta=bias + ) + return ( + output[0], + output[3], + output[2], + ) + + return replacement + + +class AddRMSNormDynamicQuantSPPattern(BasePattern): + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(vllm_config, eps) + + def get_inputs(self): + """ + Generate example inputs for the AddRMSNormQuant fusion pattern. + """ + rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype) + residual = torch.randn(2, 4, device="npu", dtype=self.dtype) + rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype) + return [rms_norm_input, residual, rms_norm_weight] + + def get_pattern(self): + def pattern(rms_norm_input: torch.Tensor, residual: torch.Tensor, rms_norm_weight: torch.Tensor): + """ + Pattern for AddRMSNormQuant fusion. + """ + output = torch.ops.npu.npu_add_rms_norm(rms_norm_input, residual, rms_norm_weight, self.eps) + out0 = output[0] + out1 = output[2] + out0 = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(out0, True) + quantized_output = torch.ops.npu.npu_dynamic_quant(out0) + return quantized_output[0], quantized_output[1], out1 + + return pattern + + def get_replacement(self): + def replacement(rms_norm_input: torch.Tensor, residual: torch.Tensor, rms_norm_weight: torch.Tensor): + """ + Replacement for the AddRMSNormQuant fusion. + """ + output = torch.ops.npu.npu_add_rms_norm_dynamic_quant( + rms_norm_input, residual, rms_norm_weight, epsilon=self.eps, output_mask=[True, False] + ) + out3 = output[3] + quantized_output = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(output[0], True) + out3 = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(out3, True) + return quantized_output, out3, output[2] + + return replacement + + +class AddRMSNormDynamicQuantSPPatternWithBias(BasePattern): + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(vllm_config, eps) + + def get_inputs(self): + """ + Generate example inputs for the AddRMSNormQuant fusion pattern. + """ + rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype) + residual = torch.randn(2, 4, device="npu", dtype=self.dtype) + rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype) + rmsnorm_bias = torch.randn(4, device="npu", dtype=self.dtype) + return [rms_norm_input, residual, rms_norm_weight, rmsnorm_bias] + + def get_pattern(self): + def pattern( + rms_norm_input: torch.Tensor, + residual: torch.Tensor, + rms_norm_weight: torch.Tensor, + bias: torch.Tensor, + ): + """ + Pattern for AddRMSNormQuant fusion. + """ + output = torch.ops._C_ascend.npu_add_rms_norm_bias( + rms_norm_input, residual, rms_norm_weight, bias, self.eps + ) + out0 = output[0] + out1 = output[2] + out0 = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(out0, True) + quantized_output = torch.ops.npu.npu_dynamic_quant(out0) + return quantized_output[0], quantized_output[1], out1 + + return pattern + + def get_replacement(self): + def replacement( + rms_norm_input: torch.Tensor, + residual: torch.Tensor, + rms_norm_weight: torch.Tensor, + bias: torch.Tensor, + ): + """ + Replacement for the AddRMSNormQuant fusion. + """ + output = torch.ops.npu.npu_add_rms_norm_dynamic_quant( + rms_norm_input, residual, rms_norm_weight, epsilon=self.eps, output_mask=[True, False], beta=bias + ) + out3 = output[3] + quantized_output = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(output[0], True) + out3 = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(out3, True) + return quantized_output, out3, output[2] + + return replacement + + +class AddRMSNormQuantFusionPass(VllmInductorPass): + """ + A pass for fusing AddRMSNorm and W8A8 quantization operations on Ascend. + """ + + def __init__(self, vllm_config: VllmConfig): + super().__init__(vllm_config) + self.pattern_match_passes: PatternMatcherPass = PatternMatcherPass(pass_name="rmsnorm_quant_fusion_pass") + + dtype = vllm_config.model_config.dtype + if dtype not in (torch.bfloat16, torch.float16): + logger.debug("Quant fusion not enabled: unsupported dtype %s", dtype) + return + + common_epsilons = [1e-5, 1e-6] + for eps in common_epsilons: + AddRMSNormDynamicQuantPattern(vllm_config, eps=eps).register(self.pattern_match_passes) + AddRMSNormDynamicQuantSPPattern(vllm_config, eps=eps).register(self.pattern_match_passes) + if enable_custom_op(): + AddRMSNormQuantPattern(vllm_config, eps=eps).register(self.pattern_match_passes) + AddRMSNormQuantSPPattern(vllm_config, eps=eps).register(self.pattern_match_passes) + AddRMSNormQuantPatternWithBias(vllm_config, eps=eps).register(self.pattern_match_passes) + AddRMSNormQuantSPPatternWithBias(vllm_config, eps=eps).register(self.pattern_match_passes) + AddRMSNormDynamicQuantPatternWithBias(vllm_config, eps=eps).register(self.pattern_match_passes) + AddRMSNormDynamicQuantSPPatternWithBias(vllm_config, eps=eps).register(self.pattern_match_passes) + + def __call__(self, graph: torch.fx.Graph): + self.begin() + self.matched_count = self.pattern_match_passes.apply(graph) + logger.debug("Replaced %s patterns", self.matched_count) + self.end_and_log() + + def is_applicable_for_range(self, compile_range: Range) -> bool: + """ + Check if the pass is applicable for the current configuration. + """ + return True diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/qknorm_rope_fusion_pass.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/qknorm_rope_fusion_pass.py new file mode 100644 index 000000000..622a30207 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/qknorm_rope_fusion_pass.py @@ -0,0 +1,244 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import torch +from torch._inductor.pattern_matcher import PatternMatcherPass, PatternPrettyPrinter +from vllm.compilation.vllm_inductor_pass import VllmInductorPass +from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.config.compilation import Range +from vllm.logger import logger +from vllm.model_executor.layers.attention import Attention + +from vllm_ascend.compilation.passes.base_pattern import BasePattern +from vllm_ascend.device.device_op import DeviceOperator +from vllm_ascend.utils import get_rope_dim + + +class QKNormRopeFusionPattern(BasePattern): + def __init__(self, vllm_config, head_dim, num_heads, num_kv_heads, eps=1e-6): + super().__init__(vllm_config, eps) + self.head_dim = head_dim + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.device = vllm_config.device_config.device if vllm_config.device_config else None + self.rope_dim = get_rope_dim(vllm_config) + + def get_inputs(self): + T = 5 + max_position_embeddings = 16384 + qkv = torch.empty(T, self.q_size + 2 * self.kv_size, dtype=torch.bfloat16, device="npu") + q_weight = torch.empty(self.head_dim, dtype=torch.bfloat16, device="npu") + k_weight = torch.empty(self.head_dim, dtype=torch.bfloat16, device="npu") + cos_sin_cache = torch.empty(max_position_embeddings, self.head_dim, dtype=torch.bfloat16, device="npu") + positions = torch.ones(T, dtype=torch.int64, device="npu") + return [qkv, q_weight, k_weight, cos_sin_cache, positions] + + def get_pattern(self): + def pattern( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + ): + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim) + q_norm_out, _ = torch.ops.npu.npu_rms_norm(q_by_head, q_weight, self.eps) + + k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim) + k_norm_out, _ = torch.ops.npu.npu_rms_norm(k_by_head, k_weight, self.eps) + + q_flat = q_norm_out.view(q.shape) + k_flat = k_norm_out.view(k.shape) + q_rope, k_rope = torch.ops.vllm.npu_rotary_embedding( + positions, q_flat, k_flat, cos_sin_cache, self.head_dim, self.rope_dim, True + ) + + return q_rope, k_rope, v + + return pattern + + def get_replacement(self): + def replacement( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + ): + results = DeviceOperator.split_qkv_rmsnorm_rope( + input=qkv, + q_weight=q_weight, + k_weight=k_weight, + q_hidden_size=self.q_size, + kv_hidden_size=self.kv_size, + head_dim=self.head_dim, + eps=self.eps, + q_bias=None, + k_bias=None, + cos_sin_cache=cos_sin_cache, + positions=positions, + ) + + return results + + return replacement + + +class QKNormRopeFusionPatternWithBias(BasePattern): + def __init__(self, vllm_config, head_dim, num_heads, num_kv_heads, eps=1e-6): + super().__init__(vllm_config, eps) + self.head_dim = head_dim + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.device = vllm_config.device_config.device if vllm_config.device_config else None + self.rope_dim = get_rope_dim(vllm_config) + + def get_inputs(self): + T = 5 + max_position_embeddings = 16384 + qkv = torch.empty(T, self.q_size + 2 * self.kv_size, dtype=torch.bfloat16, device="npu") + q_weight = torch.empty(self.head_dim, dtype=torch.bfloat16, device="npu") + k_weight = torch.empty(self.head_dim, dtype=torch.bfloat16, device="npu") + q_bias = torch.empty(self.head_dim, dtype=torch.bfloat16, device="npu") + k_bias = torch.empty(self.head_dim, dtype=torch.bfloat16, device="npu") + cos_sin_cache = torch.empty(max_position_embeddings, self.head_dim, dtype=torch.bfloat16, device="npu") + positions = torch.ones(T, dtype=torch.int64, device="npu") + + return [qkv, q_weight, k_weight, q_bias, k_bias, cos_sin_cache, positions] + + def get_pattern(self): + def pattern( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + q_bias: torch.Tensor, + k_bias: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + ): + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim) + q_norm_out, _ = torch.ops.npu.npu_rms_norm(q_by_head, q_weight, self.eps) + q_normed = q_norm_out + q_bias + + k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim) + k_norm_out, _ = torch.ops.npu.npu_rms_norm(k_by_head, k_weight, self.eps) + k_normed = k_norm_out + k_bias + + q_flat = q_normed.view(q.shape) + k_flat = k_normed.view(k.shape) + q_rope, k_rope = torch.ops.vllm.npu_rotary_embedding( + positions, q_flat, k_flat, cos_sin_cache, self.head_dim, self.rope_dim, True + ) + + return q_rope, k_rope, v + + return pattern + + def get_replacement(self): + def replacement( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + q_bias: torch.Tensor, + k_bias: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + ): + results = DeviceOperator.split_qkv_rmsnorm_rope( + input=qkv, + q_weight=q_weight, + k_weight=k_weight, + q_hidden_size=self.q_size, + kv_hidden_size=self.kv_size, + head_dim=self.head_dim, + eps=self.eps, + q_bias=q_bias, + k_bias=k_bias, + cos_sin_cache=cos_sin_cache, + positions=positions, + ) + return results + + return replacement + + +class QKNormRopeFusionPass(VllmInductorPass): + """ + A pass for fusing QKV split and RMSNorm operations into a single qk_rmsnorm operator. + """ + + def __init__(self, vllm_config: VllmConfig): + super().__init__(vllm_config) + self.pattern_match_passes: PatternMatcherPass = PatternMatcherPass(pass_name="qknorm_rope_fusion_pass") + + dtype = vllm_config.model_config.dtype + if dtype not in (torch.bfloat16,): + logger.debug("QKNorm and Rope fusion not enabled: unsupported dtype %s", dtype) + return + + # use one attn layer to get meta (such as head_dim) for QKNormRopeFusionPattern + attn_layers: dict[str, Attention] = get_layers_from_vllm_config(vllm_config, Attention) + if len(attn_layers) == 0: + logger.debug("QKNorm and Rope fusion enabled, but no Attention layers were discovered.") + return + layer = next(iter(attn_layers.values())) + for epsilon in [1e-6, 1e-5]: + if layer.head_size != 128: + logger.debug("QKNorm and Rope fusion not enabled: head_dim %d is not equal of 128", layer.head_size) + continue + QKNormRopeFusionPattern( + vllm_config=vllm_config, + head_dim=layer.head_size, + num_heads=layer.num_heads, + num_kv_heads=layer.num_kv_heads, + eps=epsilon, + ).register(self.pattern_match_passes) + + QKNormRopeFusionPatternWithBias( + vllm_config=vllm_config, + head_dim=layer.head_size, + num_heads=layer.num_heads, + num_kv_heads=layer.num_kv_heads, + eps=epsilon, + ).register(self.pattern_match_passes) + + def __call__(self, graph: torch.fx.Graph): + self.begin() + self.matched_count = self.pattern_match_passes.apply(graph) + logger.debug("Fused %s QKNorm and Rope patterns", self.matched_count) + logger.debug("Patterns registered for replacement:") + pattern_idx = 0 + for pattern_entry in self.pattern_match_passes.patterns.values(): + for p in pattern_entry: + p_str = PatternPrettyPrinter.run(p.pattern) + logger.debug("Pattern %d: %s", pattern_idx, p_str) + pattern_idx += 1 + self.end_and_log() + + def is_applicable_for_range(self, compile_range: Range) -> bool: + """ + Check if the pass is applicable for the current configuration. + """ + return True diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/sequence_parallelism.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/sequence_parallelism.py new file mode 100644 index 000000000..db9043c45 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/sequence_parallelism.py @@ -0,0 +1,234 @@ +import torch +import torch._inductor.pattern_matcher as pm +from torch._inductor.pattern_matcher import PatternMatcherPass +from vllm.compilation.vllm_inductor_pass import VllmInductorPass +from vllm.config import VllmConfig +from vllm.config.utils import Range +from vllm.distributed import get_tensor_model_parallel_world_size, get_tp_group, tensor_model_parallel_all_reduce +from vllm.logger import logger + +from vllm_ascend.compilation.passes.noop_elimination import NoOpEliminationPass +from vllm_ascend.utils import is_moe_model + +SP_MIN_TOKEN_NUM_DEFAULT = 1000 + + +def get_sp_min_token_num(config: VllmConfig) -> int: + if is_moe_model(config): + return 1 + + return SP_MIN_TOKEN_NUM_DEFAULT + + +class _SequenceParallelPatternHelper: + """Helper for sequence parallelism patterns. + + Provides TP communication helper methods: _all_reduce, _reduce_scatter, + _all_gather, and tensor creation utilities. + """ + + def __init__( + self, + epsilon: float, + dtype: torch.dtype, + device: str, + ): + self.eps = epsilon + self.dtype = dtype + self.device = device + self.tp_group = get_tp_group() + self.tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_tp_group().rank_in_group + + def _all_reduce(self, x: torch.Tensor) -> torch.Tensor: + return tensor_model_parallel_all_reduce(x) + + def _reduce_scatter(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.vllm.reduce_scatter(x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name) + + def _all_gather(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.vllm.all_gather(x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name) + + def empty(self, *args, **kws): + return torch.empty(*args, dtype=self.dtype, device="npu", **kws) + + +class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper): + """Replaces all_reduce + AddRMSNormBias with reduce_scatter + AddRMSNormBias + + all_gather for middle-layer sequence parallelism.""" + + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device()) + + def empty(self, *args, **kws): + return torch.empty(*args, dtype=self.dtype, device="npu", **kws) + + def get_inputs(self): + """ + Generate example inputs. + """ + input = self.empty(8, 16) + weight = self.empty(16) + residual = self.empty(8, 16) + return [input, weight, residual] + + def register(self, pm_pass: PatternMatcherPass): + def pattern( + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + x = self._all_reduce(input) + result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(x, residual, weight, None, self.eps) + + return result, residual + + def replacement( + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + reduce_scatter = self._reduce_scatter(input) + residual = torch.ops.vllm.maybe_chunk_residual(reduce_scatter, residual) + result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias( + reduce_scatter, residual, weight, None, self.eps + ) + all_gather = self._all_gather(result) + return all_gather, residual + + pm.register_replacement(pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass) + + +class LastAllReduceRMSNormPattern(_SequenceParallelPatternHelper): + """Same as MiddleAllReduceRMSNormPattern but for the last layer + (no residual backprop).""" + + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device()) + + def get_inputs(self): + input = self.empty(8, 16) + weight = self.empty(16) + residual = self.empty(8, 16) + return [input, weight, residual] + + def register(self, pm_pass: PatternMatcherPass): + def pattern( + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + ) -> torch.Tensor: + x = self._all_reduce(input) + result, _, _ = torch.ops._C_ascend.npu_add_rms_norm_bias(x, residual, weight, None, self.eps) + + return result + + def replacement( + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + ) -> torch.Tensor: + reduce_scatter = self._reduce_scatter(input) + residual = torch.ops.vllm.maybe_chunk_residual(reduce_scatter, residual) + result, _, _ = torch.ops._C_ascend.npu_add_rms_norm_bias(reduce_scatter, residual, weight, None, self.eps) + all_gather = self._all_gather(result) + return all_gather + + pm.register_replacement(pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass) + + +class Qwen3VLMiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper): + """For Qwen3-VL middle layers with hidden_states + deepstack_input_embeds add. + + Replaces all_reduce + add + AddRMSNormBias with reduce_scatter + + chunk(deepstack_input_embeds) + add + AddRMSNormBias + all_gather. + """ + + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device()) + + def get_inputs(self): + input = self.empty(8, 16) + weight = self.empty(16) + residual = self.empty(8, 16) + deepstack_input_embeds = self.empty(8, 16) + return [input, weight, residual, deepstack_input_embeds] + + def register(self, pm_pass: PatternMatcherPass): + def pattern( + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + deepstack_input_embeds: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + x = self._all_reduce(input) + add_ = x + deepstack_input_embeds + result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(add_, residual, weight, None, self.eps) + + return result, residual + + def replacement( + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + deepstack_input_embeds: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + reduce_scatter = self._reduce_scatter(input) + chunk = deepstack_input_embeds.chunk(self.tp_size)[self.tp_rank] + add_ = reduce_scatter + chunk + residual = torch.ops.vllm.maybe_chunk_residual(reduce_scatter, residual) + result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(add_, residual, weight, None, self.eps) + all_gather = self._all_gather(result) + return all_gather, residual + + pm.register_replacement(pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass) + + +class SequenceParallelismPass(VllmInductorPass): + """Sequence parallelism compilation pass. + + Registers and applies the above patterns. Runs noop cleanup first, then + uses token range to determine whether to enable SP. + """ + + def __init__(self, config: VllmConfig): + super().__init__(config) + + self.patterns: PatternMatcherPass = PatternMatcherPass(pass_name="npu_sequence_parallelism_pass") + self.noop_cleanup = NoOpEliminationPass(config) + + for epsilon in [1e-5, 1e-6]: + MiddleAllReduceRMSNormPattern(config, epsilon).register(self.patterns) + + LastAllReduceRMSNormPattern(config, epsilon).register(self.patterns) + + Qwen3VLMiddleAllReduceRMSNormPattern(config, epsilon).register(self.patterns) + + self.min_tokens = get_sp_min_token_num(config) + + def __call__(self, graph: torch.fx.Graph): + self.begin() + self.noop_cleanup(graph) # Eliminate redundant view-like operations + logger.debug("after noop_cleanup %s", graph.graph) + self.matched_count = self.patterns.apply(graph) + logger.debug("Replaced %s patterns", self.matched_count) + logger.debug("after apply replacement %s", graph.graph) + + from torch._inductor.pattern_matcher import PatternPrettyPrinter + + pattern_idx = 0 + for pattern_entry in self.patterns.patterns.values(): + for p in pattern_entry: + p_str = PatternPrettyPrinter.run(p.pattern) + logger.debug("Pattern %d: %s", pattern_idx, p_str) + pattern_idx += 1 + + self.end_and_log() + + def is_applicable_for_range(self, compile_range: Range) -> bool: + """ + Check if the pass is applicable for the current configuration. + """ + applicable = compile_range.start >= self.min_tokens + logger.debug("SequenceParallelismPass compile_range=%r applicable=%r", compile_range, applicable) + return applicable diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/sequence_parallelism_moe.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/sequence_parallelism_moe.py new file mode 100644 index 000000000..3ff16aef0 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/sequence_parallelism_moe.py @@ -0,0 +1,204 @@ +import torch +import torch._inductor.pattern_matcher as pm +from torch._inductor.pattern_matcher import PatternMatcherPass +from vllm.compilation.vllm_inductor_pass import PatternPrettyPrinter, VllmInductorPass +from vllm.config import VllmConfig +from vllm.config.utils import Range +from vllm.logger import logger + +from vllm_ascend.compilation.passes.sequence_parallelism import ( + _SequenceParallelPatternHelper, + get_sp_min_token_num, +) + + +class MiddleLayerAllgatherAddRMSNormPattern(_SequenceParallelPatternHelper): + """Replaces all_gather + slice + AddRMSNormBias with AddRMSNormBias + + all_gather to avoid middle-layer shape mismatch.""" + + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device()) + + def get_inputs(self): + input = self.empty(5, 16) + weight = self.empty(16) + residual = self.empty(8, 16) + # num_tokens = 8 + return [input, weight, residual] + + def get_scalar_inputs(self): + return {"num_tokens": 8} + + def register(self, pm_pass: PatternMatcherPass): + def pattern( + input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor, num_tokens + ) -> tuple[torch.Tensor, torch.Tensor]: + all_gather = self._all_gather(input) + x_sliced = all_gather[:num_tokens] + result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(x_sliced, residual, weight, None, self.eps) + + return result, residual + + def replacement( + input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor, num_tokens + ) -> tuple[torch.Tensor, torch.Tensor]: + residual = torch.ops.vllm.maybe_chunk_residual(input, residual) + result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(input, residual, weight, None, self.eps) + all_gather = self._all_gather(result) + return all_gather, residual + + pm.register_replacement( + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass, scalar_workaround=self.get_scalar_inputs() + ) + + +class LastLayerAllgatherRMSNormPattern(_SequenceParallelPatternHelper): + """Same as MiddleLayerAllgatherAddRMSNormPattern but for the last layer (no residual) + all_gather + RMSNorm fusion.""" + + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device()) + + def get_inputs(self): + input = self.empty(5, 16) + weight = self.empty(16) + residual = self.empty(8, 16) + return [input, weight, residual] + + def get_scalar_inputs(self): + return {"num_tokens": 8} + + def register(self, pm_pass: PatternMatcherPass): + def pattern( + input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor, num_tokens + ) -> tuple[torch.Tensor, torch.Tensor]: + all_gather = self._all_gather(input) + x_sliced = all_gather[:num_tokens] + result, _, _ = torch.ops._C_ascend.npu_add_rms_norm_bias(x_sliced, residual, weight, None, self.eps) + + return result + + def replacement( + input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor, num_tokens + ) -> tuple[torch.Tensor, torch.Tensor]: + residual = torch.ops.vllm.maybe_chunk_residual(input, residual) + result, _, _ = torch.ops._C_ascend.npu_add_rms_norm_bias(input, residual, weight, None, self.eps) + all_gather = self._all_gather(result) + return all_gather + + pm.register_replacement( + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass, scalar_workaround=self.get_scalar_inputs() + ) + + +class Qwen3VLMiddleLayerAllgatherAddRMSNormPattern(_SequenceParallelPatternHelper): + """Replaces all_gather + slice + add + AddRMSNormBias with add(chunk) + + AddRMSNormBias + all_gather for Qwen3-VL-style all_gather path.""" + + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device()) + + def get_inputs(self): + input = self.empty(5, 16) + weight = self.empty(16) + residual = self.empty(8, 16) + deepstack_input_embeds = self.empty(8, 16) + return [input, weight, residual, deepstack_input_embeds] + + def get_scalar_inputs(self): + return {"num_tokens": 8} + + def register(self, pm_pass: PatternMatcherPass): + def pattern( + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + deepstack_input_embeds: torch.Tensor, + num_tokens, + ) -> tuple[torch.Tensor, torch.Tensor]: + all_gather = self._all_gather(input) + x_sliced = all_gather[:num_tokens] + add_ = x_sliced + deepstack_input_embeds + result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(add_, residual, weight, None, self.eps) + + return result, residual + + def replacement( + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + deepstack_input_embeds: torch.Tensor, + num_tokens, + ) -> tuple[torch.Tensor, torch.Tensor]: + chunk = deepstack_input_embeds.chunk(self.tp_size)[self.tp_rank] + add_ = input + chunk + residual = torch.ops.vllm.maybe_chunk_residual(input, residual) + result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(add_, residual, weight, None, self.eps) + all_gather = self._all_gather(result) + return all_gather, residual + + pm.register_replacement( + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass, scalar_workaround=self.get_scalar_inputs() + ) + + +class AllGatherChunkNoOpPattern(_SequenceParallelPatternHelper): + """Folds all_gather + sequence_parallel_chunk_impl into identity (no-op).""" + + def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6): + super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device()) + + def get_inputs(self): + return [self.empty(8, 16)] + + def register(self, pm_pass: PatternMatcherPass): + def pattern(input: torch.Tensor) -> torch.Tensor: + gathered = self._all_gather(input) + return torch.ops.vllm.sequence_parallel_chunk_impl(gathered) + + def replacement(input: torch.Tensor) -> torch.Tensor: + return input + + pm.register_replacement(pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass) + + +class SequenceParallelismMoePass(VllmInductorPass): + """Sequence parallelism AllGather epilogue pass. + + Applies AllGather-based patterns: MiddleLayerAllgatherAddRMSNormPattern, + LastLayerAllgatherRMSNormPattern, Qwen3VLMiddleLayerAllgatherAddRMSNormPattern, + and AllGatherChunkNoOpPattern (all_gather + sequence_parallel_chunk_impl -> identity). + """ + + def __init__(self, config: VllmConfig): + super().__init__(config) + + self.patterns: PatternMatcherPass = PatternMatcherPass(pass_name="npu_sequence_parallelism_allgather_ep_pass") + + for epsilon in [1e-5, 1e-6]: + MiddleLayerAllgatherAddRMSNormPattern(config, epsilon).register(self.patterns) + LastLayerAllgatherRMSNormPattern(config, epsilon).register(self.patterns) + Qwen3VLMiddleLayerAllgatherAddRMSNormPattern(config, epsilon).register(self.patterns) + + AllGatherChunkNoOpPattern(config).register(self.patterns) + + self.min_tokens = get_sp_min_token_num(config) + + def __call__(self, graph: torch.fx.Graph): + self.begin() + logger.debug("before apply replacement %s", str(graph)) + self.matched_count = self.patterns.apply(graph) + logger.debug("after apply replacement %s", str(graph)) + logger.debug("SequenceParallelismMoePass replaced %s patterns", self.matched_count) + pattern_idx = 0 + for pattern_entry in self.patterns.patterns.values(): + for p in pattern_entry: + p_str = PatternPrettyPrinter.run(p.pattern) + logger.debug("Pattern %d: %s", pattern_idx, p_str) + pattern_idx += 1 + self.end_and_log() + + def is_applicable_for_range(self, compile_range: Range) -> bool: + applicable = compile_range.start >= self.min_tokens + logger.debug("SequenceParallelismMoePass compile_range=%r applicable=%r", compile_range, applicable) + return applicable diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/utils/__init__.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/utils/npugraph_ex_utils_check.py b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/utils/npugraph_ex_utils_check.py new file mode 100644 index 000000000..a91d08426 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/compilation/passes/utils/npugraph_ex_utils_check.py @@ -0,0 +1,75 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from torch._inductor.pattern_matcher import Match +from vllm.logger import logger + + +def extra_stream_scope_check(match: Match) -> bool: + """ + Checks if all nodes in the same stream. + """ + non_default_streams = set() + has_default = False + + for node in match.nodes: + if node.op == "call_function": + current_stream = node.meta.get("stream_label") + if current_stream is None: + has_default = True + else: + non_default_streams.add(current_stream) + if len(non_default_streams) > 1: + logger.debug( + "Cross-stream operation detected in pattern match for AddRMSNormQuant. " + "Multiple streams found: %s. Fusion is not supported for cross-stream operations.", + non_default_streams, + ) + return False + + if has_default and len(non_default_streams) > 0: + logger.debug( + "Cross-stream operation detected in pattern match for AddRMSNormQuant. " + "Multiple streams found: %s. Fusion is not supported for cross-stream operations.", + non_default_streams, + ) + return False + + return True + + +_register_patterns = set() + + +def check_and_register_fusion_pass(pattern_class: type, **kwargs): + global _register_patterns + eps = kwargs.get("eps", 1e-6) + pattern_key = str(pattern_class.__name__) + str(eps) + if pattern_key in _register_patterns: + return + + pattern = pattern_class(**kwargs) + try: + pattern.register() + _register_patterns.add(pattern_key) + except RuntimeError as e: + if "Duplicate pattern" in str(e): + logger.warning("Pattern %s eps %s has been registered", pattern_class.__name__, eps) + _register_patterns.add(pattern_key) + else: + raise e diff --git a/vllm_fl/dispatch/backends/vendor/ascend/impl/attention.py b/vllm_fl/dispatch/backends/vendor/ascend/impl/attention.py index 7424ec37a..9bd10e94d 100644 --- a/vllm_fl/dispatch/backends/vendor/ascend/impl/attention.py +++ b/vllm_fl/dispatch/backends/vendor/ascend/impl/attention.py @@ -35,15 +35,28 @@ AttentionType, ) from vllm.config import VllmConfig, get_current_vllm_config +from vllm.forward_context import get_forward_context from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.utils import AttentionCGSupport, CommonAttentionMetadata from vllm_fl.dispatch.backends.vendor.ascend.impl.attention_mask import ( AttentionMaskBuilder, ) +from vllm_fl.dispatch.backends.vendor.ascend.patches.patch_graph import ( + get_draft_graph_params, + get_draft_graph_prefill_params, + get_graph_params, + update_draft_graph_params_workspaces, + update_graph_params_workspaces, + weak_ref_tensors, +) logger = logging.getLogger(__name__) +# Ascend npu_fused_infer_attention_score with input_layout="TND" only supports +# head_dim in {64, 128, 192} (plus the special qD=kD=192, vD=128 case). +_PFA_TND_SUPPORTED_HEAD_DIMS = frozenset({64, 128, 192}) + # Check torch_npu availability and setup NPU compatibility _TORCH_NPU_AVAILABLE = False try: @@ -338,6 +351,17 @@ def build( model_runner_type=self.model_config.runner_type, ) + def build_for_drafting( + self, + common_attn_metadata, + draft_index: int, + ): + """Build attention metadata for the draft model during speculative decoding.""" + return self.build( + common_prefix_len=0, + common_attn_metadata=common_attn_metadata, + ) + def _determine_attn_state( self, num_decodes: int, @@ -539,6 +563,69 @@ def __init__( self.num_queries_per_kv = self.num_heads // self.num_kv_heads self.key_cache = None self.value_cache = None + self._use_fusion_fallback = self.head_size not in _PFA_TND_SUPPORTED_HEAD_DIMS + if self._use_fusion_fallback: + logger.info( + "AscendAttentionBackendImpl: head_size=%d is not supported by " + "npu_fused_infer_attention_score TND layout, falling back to " + "head-split attention (2x%d).", self.head_size, + self.head_size // 2) + + @classmethod + def update_graph_params( + cls, + update_stream, + forward_context, + num_tokens: int, + vllm_config: VllmConfig, + speculative_config=None, + num_dcp_pcp_tokens=None, + draft_attn_metadatas=None, + ) -> None: + """Update captured paged-attention tasks with this step's seq_lens.""" + graph_params = get_graph_params() + if graph_params is None or not graph_params.handles.get(num_tokens): + return + + with torch.npu.stream(update_stream): + captured_params = graph_params.attn_params[num_tokens] + handles = graph_params.handles[num_tokens] + events = graph_params.events[num_tokens] + if not (len(captured_params) == len(handles) == len(events)): + raise RuntimeError( + "Mismatched FULL ACL graph attention metadata: " + f"params={len(captured_params)}, handles={len(handles)}, " + f"events={len(events)}, num_tokens={num_tokens}" + ) + for param, handle, event in zip(captured_params, handles, events): + ( + layer_name, + query, + key_cache, + value_cache, + num_kv_heads, + num_heads, + scale, + block_table, + output, + ) = param + attn_metadata = forward_context.attn_metadata[layer_name] + + torch.npu.graph_task_update_begin(update_stream, handle) + torch_npu._npu_paged_attention( + query=query, + key_cache=key_cache, + value_cache=value_cache, + num_kv_heads=num_kv_heads, + num_heads=num_heads, + scale_value=scale, + block_table=block_table, + context_lens=attn_metadata.seq_lens, + out=output, + workspace=graph_params.workspaces.get(num_tokens), + ) + torch.npu.graph_task_update_end(update_stream) + event.record(update_stream) def _get_fia_params( self, @@ -547,9 +634,8 @@ def _get_fia_params( attn_metadata: AscendMetadata, ): """Get parameters for fused_infer_attention.""" - + block_size = 128 if attn_metadata.attn_state == AscendAttentionState.PrefillNoCache: - block_size = 128 block_table = None actual_seq_lengths_kv = attn_metadata.actual_seq_lengths_q elif attn_metadata.attn_state == AscendAttentionState.PrefillCacheHit: @@ -563,8 +649,16 @@ def _get_fia_params( # num_block, block_size, _, _ = self.key_cache.shape # key = self.key_cache.view(num_block, block_size, -1) # value = self.value_cache.view(num_block, block_size, -1) - key = self.key_cache.view(-1, block_size, 256) - value = self.value_cache.view(-1, block_size, 256) + + #key = self.key_cache.view(-1, block_size, 256) + #value = self.value_cache.view(-1, block_size, 256) + #block_table = attn_metadata.block_tables + #actual_seq_lengths_kv = attn_metadata.seq_lens_list + + # Get block_size from key_cache shape + num_block, block_size, _, _ = self.key_cache.shape + key = self.key_cache.view(num_block, block_size, -1) + value = self.value_cache.view(num_block, block_size, -1) block_table = attn_metadata.block_tables actual_seq_lengths_kv = attn_metadata.seq_lens_list else: @@ -572,8 +666,16 @@ def _get_fia_params( # num_block, block_size, _, _ = self.key_cache.shape # key = self.key_cache.view(num_block, block_size, -1) # value = self.value_cache.view(num_block, block_size, -1) - key = self.key_cache.view(-1, block_size, 256) - value = self.value_cache.view(-1, block_size, 256) + + #key = self.key_cache.view(-1, block_size, 256) + #value = self.value_cache.view(-1, block_size, 256) + #block_table = attn_metadata.block_tables + #actual_seq_lengths_kv = attn_metadata.seq_lens_list + + # Get block_size from key_cache shape + num_block, block_size, _, _ = self.key_cache.shape + key = self.key_cache.view(num_block, block_size, -1) + value = self.value_cache.view(num_block, block_size, -1) block_table = attn_metadata.block_tables actual_seq_lengths_kv = attn_metadata.seq_lens_list @@ -605,6 +707,130 @@ def reshape_and_cache( ) return key, value + def _gather_kv_cache( + self, + attn_metadata: AscendMetadata, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Gather paged KV cache into dense (T_kv, num_kv_heads, head_size). + + Only the blocks actually referenced by each request's sequence length + are gathered. The block table is padded to ``max_num_blocks_per_req``, + so gathering the whole table would allocate ``batch_size`` times more + memory than necessary and can OOM on long-context configs. + """ + if self.key_cache is None or self.value_cache is None: + raise RuntimeError( + "key_cache/value_cache is not initialized for KV gather") + + seq_lens = attn_metadata.seq_lens_list + batch_size = len(seq_lens) + block_size = self.key_cache.shape[1] + block_table = attn_metadata.block_tables[:batch_size] + + k_list: List[torch.Tensor] = [] + v_list: List[torch.Tensor] = [] + for i, sl in enumerate(seq_lens): + num_blocks = (sl + block_size - 1) // block_size + block_ids = block_table[i, :num_blocks].long() + gathered_k = self.key_cache[block_ids].view( + -1, self.num_kv_heads, self.head_size)[:sl] + gathered_v = self.value_cache[block_ids].view( + -1, self.num_kv_heads, self.head_size)[:sl] + k_list.append(gathered_k) + v_list.append(gathered_v) + + return torch.cat(k_list, dim=0), torch.cat(v_list, dim=0) + + def _split_heads_for_pfa( + self, + x: torch.Tensor, + num_heads: int, + ) -> torch.Tensor: + """Split head dimension to a supported PFA TND head size. + + Ascend's npu_fused_infer_attention_score TND layout only supports + head_dim in {64, 128, 192}. For unsupported head sizes such as 256, + we split each head into two heads of half the dimension and double the + head count, which is mathematically equivalent for self-attention. + + Accepts either 2D input [T, num_heads * head_size] or 3D input + [T, num_heads, head_size]. Output shape: [T, num_heads * 2, + head_size // 2]. + """ + head_size_128 = self.head_size // 2 + if x.dim() == 2: + x = x.view(-1, num_heads, self.head_size) + return x.view(-1, num_heads, 2, head_size_128).reshape( + -1, num_heads * 2, head_size_128) + + def _merge_heads_from_pfa( + self, + x: torch.Tensor, + ) -> torch.Tensor: + """Merge split heads back to original head dimension. + + Input shape: [T, num_heads * 2, head_size // 2] + Output shape: [T, num_heads * head_size] + """ + head_size_128 = self.head_size // 2 + return x.view(-1, self.num_heads, 2, head_size_128).reshape( + -1, self.num_heads * self.head_size) + + def _forward_fusion_attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AscendMetadata, + output: torch.Tensor, + ) -> torch.Tensor: + """Forward pass for unsupported head sizes using head-split PFA. + + Splits head_dim into two supported heads and runs + npu_fused_infer_attention_score (PFA) with the standard TND path. + This keeps the implementation compatible with ACL graph capture while + avoiding the per-batch dense causal mask limitation of + npu_fusion_attention. + """ + num_tokens = attn_metadata.actual_seq_lengths_q[-1] + query = query[:num_tokens] + + if attn_metadata.attn_state == AscendAttentionState.PrefillNoCache: + key = key[:num_tokens] + value = value[:num_tokens] + actual_seq_kvlen = attn_metadata.actual_seq_lengths_q + else: + key, value = self._gather_kv_cache(attn_metadata) + actual_seq_kvlen = [ + sum(attn_metadata.seq_lens_list[:i + 1]) + for i in range(len(attn_metadata.seq_lens_list)) + ] + + query = self._split_heads_for_pfa(query, self.num_heads) + key = self._split_heads_for_pfa(key, self.num_kv_heads) + value = self._split_heads_for_pfa(value, self.num_kv_heads) + + attn_output, _ = torch_npu.npu_fused_infer_attention_score( + query=query, + key=key, + value=value, + atten_mask=attn_metadata.attn_mask, + block_table=None, + input_layout="TND", + block_size=AscendAttentionBackend.get_supported_block_size()[0], + actual_seq_lengths=attn_metadata.actual_seq_lengths_q, + actual_seq_lengths_kv=actual_seq_kvlen, + num_key_value_heads=self.num_kv_heads * 2, + num_heads=self.num_heads * 2, + scale=self.scale, + sparse_mode=3, + ) + + attn_output = self._merge_heads_from_pfa(attn_output) + attn_output = attn_output.view(num_tokens, self.num_heads, self.head_size) + output[:num_tokens] = attn_output[:num_tokens] + return output + def forward_fused_infer_attention( self, query: torch.Tensor, @@ -614,6 +840,11 @@ def forward_fused_infer_attention( output: torch.Tensor, ) -> torch.Tensor: """Forward pass using fused_infer_attention_score.""" + if (self._use_fusion_fallback + and attn_metadata.attn_state != AscendAttentionState.DecodeOnly): + return self._forward_fusion_attention( + query, key, value, attn_metadata, output) + key, value, block_size, block_table, actual_seq_lengths_kv = \ self._get_fia_params(key, value, attn_metadata) @@ -652,8 +883,73 @@ def forward_paged_attention( query: torch.Tensor, attn_metadata: AscendMetadata, output: Optional[torch.Tensor] = None, + layer_name: Optional[str] = None, ) -> torch.Tensor: """Forward pass using paged attention for decode.""" + forward_context = get_forward_context() + num_tokens = query.shape[0] + if getattr(forward_context, "capturing", False): + graph_params = get_graph_params() + if graph_params is None: + raise RuntimeError( + "ACL graph parameters were not initialized before capture" + ) + if layer_name is None: + raise RuntimeError( + "Attention layer name is required for FULL ACL graph capture" + ) + + workspace = graph_params.workspaces.get(num_tokens) + if workspace is None: + workspace = torch_npu._npu_paged_attention_get_workspace( + query=query, + key_cache=self.key_cache, + value_cache=self.value_cache, + num_kv_heads=self.num_kv_heads, + num_heads=self.num_heads, + scale_value=self.scale, + block_table=attn_metadata.block_tables, + context_lens=attn_metadata.seq_lens, + out=output, + ) + update_graph_params_workspaces(num_tokens, workspace) + + stream = torch.npu.current_stream() + event = torch.npu.ExternalEvent() + event.wait(stream) + event.reset(stream) + graph_params.events[num_tokens].append(event) + graph_params.attn_params[num_tokens].append( + ( + layer_name, + weak_ref_tensors(query), + weak_ref_tensors(self.key_cache), + weak_ref_tensors(self.value_cache), + self.num_kv_heads, + self.num_heads, + self.scale, + weak_ref_tensors(attn_metadata.block_tables), + weak_ref_tensors(output), + ) + ) + + torch.npu.graph_task_group_begin(stream) + torch_npu._npu_paged_attention( + query=query, + key_cache=self.key_cache, + value_cache=self.value_cache, + num_kv_heads=self.num_kv_heads, + num_heads=self.num_heads, + scale_value=self.scale, + block_table=attn_metadata.block_tables, + context_lens=attn_metadata.seq_lens, + out=output, + workspace=workspace, + ) + handle = torch.npu.graph_task_group_end(stream) + graph_params.handles[num_tokens].append(handle) + return output + torch_npu._npu_paged_attention( query=query, key_cache=self.key_cache, @@ -713,6 +1009,7 @@ def forward_impl( kv_cache: Tuple[torch.Tensor], attn_metadata: AscendMetadata, output: torch.Tensor, + layer_name: Optional[str] = None, ): """Forward implementation dispatching to appropriate attention method.""" num_tokens = query.shape[0] @@ -720,7 +1017,8 @@ def forward_impl( # Use paged attention for decode-only state if (attn_metadata.attn_state == AscendAttentionState.DecodeOnly and self.sliding_window is None): - output = self.forward_paged_attention(query, attn_metadata, output) + output = self.forward_paged_attention( + query, attn_metadata, output, layer_name=layer_name) else: output = self.forward_fused_infer_attention( query, key, value, attn_metadata, output) @@ -778,7 +1076,12 @@ def forward( return output.fill_(0) # Reshape and cache KV - if attn_metadata != AscendAttentionState.DecodeOnly: + # NOTE: compare attn_state (not the metadata object itself) — the + # previous `attn_metadata != AscendAttentionState.DecodeOnly` was + # always True and forced a full-cache .contiguous() copy on every + # call. With the contiguous (2, num_blocks, ...) cache layout this + # is a no-op; only non-decode states ever need it. + if attn_metadata.attn_state != AscendAttentionState.DecodeOnly: kv_cache = [i.contiguous() for i in kv_cache] if key is not None and value is not None: key = key.contiguous() @@ -794,25 +1097,216 @@ def forward( # Standard forward output = self.forward_impl( - query, key, value, kv_cache, attn_metadata, output) + query, + key, + value, + kv_cache, + attn_metadata, + output, + layer_name=layer.layer_name, + ) return output -# MLA Backend placeholder - can be extended later -class AscendMLABackend: + +# MLA Backend +class AscendMLABackend(AttentionBackend): """ - Ascend MLA (Multi-head Latent Attention) backend placeholder. + Ascend MLA (Multi-head Latent Attention) backend. - This is a minimal implementation. Full MLA support would require - additional implementation based on the specific MLA algorithm. + This is a structural backend that participates in the Ascend ACL graph + parameter upgrade mechanism. The full MLA forward path can be extended + later; for now the class provides the backend contract and graph-param + hooks required by the model runner. """ - def __init__(self, *args, **kwargs): - raise NotImplementedError( - "Ascend MLA attention backend is not yet fully implemented. " - "Please use standard attention backend by setting use_mla=False" + accept_output_buffer: bool = True + + @staticmethod + def get_name() -> str: + return "ASCEND_MLA_FL" + + @staticmethod + def get_impl_cls() -> Type["AscendMLABackendImpl"]: + return AscendMLABackendImpl + + @staticmethod + def get_builder_cls() -> Type["AscendMLAMetadataBuilder"]: + return AscendMLAMetadataBuilder + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> Tuple[int, ...]: + # MLA uses a compressed KV cache shape. + return (num_blocks, block_size, num_kv_heads, head_size) + + @staticmethod + def swap_blocks( + src_kv_cache: List[torch.Tensor], + dst_kv_cache: List[torch.Tensor], + src_to_dst: torch.Tensor, + ) -> None: + src_indices = src_to_dst[:, 0] + dst_indices = src_to_dst[:, 1] + for src, dst in zip(src_kv_cache, dst_kv_cache): + dst[dst_indices] = src[src_indices].to(dst.device) + + @staticmethod + def copy_blocks( + kv_caches: List[torch.Tensor], + src_to_dists: torch.Tensor, + ) -> None: + src_indices = src_to_dists[:, 0] + dst_indices = src_to_dists[:, 1] + for kv_cache in kv_caches: + kv_cache[dst_indices] = kv_cache[src_indices] + + @staticmethod + def get_supported_block_size() -> list[int]: + return [128] + + +@dataclass +class AscendMLAMetadata: + """Per-layer metadata placeholder for Ascend MLA attention.""" + + attn_mask: Optional[torch.Tensor] = None + slot_mapping: torch.Tensor = None + block_tables: torch.Tensor = None + seq_lens: torch.Tensor = None + seq_lens_list: List[int] = None + query_start_loc: torch.Tensor = None + max_query_len: Optional[int] = None + num_actual_tokens: int = 0 + + +class AscendMLAMetadataBuilder: + """Metadata builder for Ascend MLA attention.""" + + aclgraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS + reorder_batch_threshold: ClassVar[int] = 1 + + def __init__( + self, + kv_cache_spec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ): + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config + self.device = device + + @staticmethod + def get_cudagraph_support(vllm_config, kv_cache_spec) -> AttentionCGSupport: + return AttentionCGSupport.ALWAYS + + def build( + self, + common_prefix_len: int, + common_attn_metadata, + model: Optional[nn.Module] = None, + ): + num_reqs = common_attn_metadata.num_reqs + num_actual_tokens = common_attn_metadata.num_actual_tokens + return AscendMLAMetadata( + slot_mapping=common_attn_metadata.slot_mapping[:num_actual_tokens], + block_tables=common_attn_metadata.block_table_tensor, + seq_lens=common_attn_metadata.seq_lens_cpu[:num_reqs], + seq_lens_list=common_attn_metadata.seq_lens_cpu[:num_reqs].tolist(), + query_start_loc=common_attn_metadata.query_start_loc_cpu[:num_reqs + 1], + max_query_len=common_attn_metadata.max_query_len, + num_actual_tokens=num_actual_tokens, ) + def build_for_graph_capture( + self, + common_attn_metadata, + model: Optional[nn.Module] = None, + ): + return self.build(0, common_attn_metadata, model) + + def reorder_batch(self, input_batch, scheduler_output) -> bool: + return False + + def use_cascade_attention(self, *args, **kwargs) -> bool: + return False + + +class AscendMLABackendImpl(AttentionImpl): + """ + Ascend MLA attention implementation placeholder. + + Provides the graph-param upgrade hook; full MLA forward kernel integration + can be added later without changing the framework contract. + """ + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: Optional[List[float]], + sliding_window: Optional[int], + kv_cache_dtype: str, + logits_soft_cap: Optional[float], + attn_type: str, + kv_sharing_target_layer_name: Optional[str], + **kwargs, + ) -> None: + self.num_heads = num_heads + self.head_size = head_size + self.scale = float(scale) + self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads + self.kv_cache_dtype = kv_cache_dtype + self.sliding_window = sliding_window + self.attn_type = attn_type + + @classmethod + def update_graph_params( + cls, + update_stream, + forward_context, + num_tokens: int, + vllm_config: VllmConfig, + speculative_config=None, + num_dcp_pcp_tokens=None, + draft_attn_metadatas=None, + ) -> None: + """Update graph parameters for Ascend MLA ACL graph capture.""" + logger.debug( + "Updating graph params for AscendMLABackendImpl num_tokens=%s", + num_tokens) + params = get_graph_params() + if params is not None and num_tokens in params.workspaces: + if params.workspaces[num_tokens] is None: + params.workspaces[num_tokens] = True # marker + + def forward( + self, + layer: AttentionLayer, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: Tuple[torch.Tensor], + attn_metadata: AscendMLAMetadata, + output: Optional[torch.Tensor] = None, + output_scale: Optional[torch.Tensor] = None, + output_block_scale: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if output is None: + output = torch.empty_like(query) + # TODO: integrate native Ascend MLA kernel. + raise NotImplementedError( + "AscendMLABackendImpl.forward is not yet implemented. " + "Use standard attention backend for now.") + __all__ = [ "AscendAttentionBackend", @@ -821,5 +1315,8 @@ def __init__(self, *args, **kwargs): "AscendMetadata", "AscendAttentionState", "AscendMLABackend", + "AscendMLAMetadataBuilder", + "AscendMLABackendImpl", + "AscendMLAMetadata", "is_torch_npu_available", ] diff --git a/vllm_fl/dispatch/backends/vendor/ascend/impl/fla/fused_recurrent.py b/vllm_fl/dispatch/backends/vendor/ascend/impl/fla/fused_recurrent.py new file mode 100644 index 000000000..eaaa31c88 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/impl/fla/fused_recurrent.py @@ -0,0 +1,223 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Adapted from https://github.com/vllm-project/vllm-ascend/blob/main/vllm_ascend/ops/triton/fla/sigmoid_gating.py +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# ruff: noqa: E501 +# mypy: ignore-errors + +import torch +from vllm.triton_utils import tl, triton + + +@triton.heuristics( + { + "USE_INITIAL_STATE": lambda args: args["h0_source"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.jit(do_not_specialize=["T"]) +def fused_recurrent_delta_rule_update_kernel( + q, + k, + v, + g, + beta, + o, + h0_source, + h0_indices, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """ + Recurrent delta-rule decode update: q/k L2 norm + state update + output + in a single launch. The sigmoid gating (g/beta) is computed outside by + the AscendC ``npu_fused_gdn_gating`` op — the in-kernel gating section + of the upstream fused_sigmoid_gating_delta_rule_update kernel is + miscompiled by the Ascend Triton pipeline in this environment. + """ + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hv = i_nh // HV - 1, i_nh % HV + # The grid carries one dummy group of HV programs up front: the first + # scheduled program on this Ascend Triton pipeline sporadically + # produces a corrupted state/output tile (verified on 910B4-1 with the + # CANN 8.5.0 bishengir pipeline; every other program is exact), so it + # is made to return immediately and all real sequences are shifted to + # healthy programs. + if i_n < 0: + return + i_h = i_hv // (HV // H) + + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + p_g = g + bos * HV + i_hv + p_beta = beta + bos * HV + i_hv + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + + # NOTE: the FL ssm state cache is laid out per slot as (HV, V, K) + # (v-major), as required by the AscendC recurrent_gated_delta_rule + # kernel, while the upstream kernel this is adapted from assumes the + # FLA (HV, K, V) layout. The state tile is therefore kept in the + # native (BV, BK) v-major orientation end to end. + mask_h_t = mask_v[:, None] & mask_k[None, :] + + b_h = tl.zeros([BV, BK], dtype=tl.float32) + if USE_INITIAL_STATE: + idx = tl.load(h0_indices + i_n).to(tl.int64) + # if idx >= 0: + tmp0 = tl.where(idx < 0, 0, idx) + p_h0 = h0_source + tmp0 * HV * K * V + i_hv * K * V + o_v[:, None] * K + o_k[None, :] + temp1 = tl.load(p_h0, mask=mask_h_t, other=0).to(tl.float32) + temp2 = tl.zeros_like(temp1) + b_h += tl.where(idx < 0, temp2, temp1) + + for i in range(0, T): + # Load inputs + b_q = tl.load(p_q + i * H * K, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k + i * H * K, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v + i * HV * V, mask=mask_v, other=0).to(tl.float32) + b_g = tl.load(p_g + i * HV).to(tl.float32) + b_beta = tl.load(p_beta + i * HV).to(tl.float32) + + # Apply L2 normalization if enabled + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q)) + 1e-6) + b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k)) + 1e-6) + + b_q = b_q * scale + + # Apply gating to hidden state: h *= exp(g) + b_h *= tl.exp(b_g) + + # Delta rule: v -= sum(h * k, dim=1) + b_v -= tl.sum(b_h * b_k[None, :], 1) + + # Apply beta gating: v *= beta + b_v *= b_beta + + # Update hidden state: h += v[:, None] * k[None, :] + b_h += b_v[:, None] * b_k[None, :] + + # Compute output: o = sum(h * q, dim=1) + b_o = tl.sum(b_h * b_q[None, :], 1) + tl.store(p_o + i * HV * V, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + # Store final state back to h0_source with bounds checking + if USE_INITIAL_STATE: + idx = tl.load(h0_indices + i_n).to(tl.int64) + if idx >= 0: + p_h0 = h0_source + idx * HV * K * V + i_hv * K * V + o_v[:, None] * K + o_k[None, :] + tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h_t) + + +def fused_recurrent_delta_rule_update( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + scale: float = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor = None, +): + """ + Fused triton implementation of the recurrent gated delta-rule decode + update: q/k L2 norm + state update + output in a single kernel launch. + + ``g`` (log-space forget gate) and ``beta`` are precomputed (by the + AscendC ``npu_fused_gdn_gating`` op) with shape ``[B * T, HV]``. + + Note: ``initial_state_source`` must use the FL per-slot (HV, V, K) state + layout (v-major), consistent with the AscendC GDN kernels. + """ + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 64) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + assert NK == 1, "NK > 1 is not supported yet" + num_stages = 3 + num_warps = 1 + + if scale is None: + scale = k.shape[-1] ** -0.5 + else: + assert scale > 0, "scale must be positive" + + if cu_seqlens is not None: + # One extra group of dummy programs is prepended in the grid (the + # kernel returns immediately for them) to absorb the corrupted + # first scheduled program; no tensor padding needed. + N += 1 + + o = q.new_empty(NK, *v.shape) + grid = (NK, NV, N * HV) + + if not initial_state_indices.is_contiguous(): + initial_state_indices = initial_state_indices.contiguous() + if not initial_state_source.is_contiguous(): + initial_state_source = initial_state_source.contiguous() + if not cu_seqlens.is_contiguous(): + cu_seqlens = cu_seqlens.contiguous() + + fused_recurrent_delta_rule_update_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + o=o, + h0_source=initial_state_source, + h0_indices=initial_state_indices, + cu_seqlens=cu_seqlens, + scale=scale, + T=T, + B=B, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + num_warps=num_warps, + num_stages=num_stages, + ) + o = o.squeeze(0) + return o diff --git a/vllm_fl/dispatch/backends/vendor/ascend/impl/fused_moe.py b/vllm_fl/dispatch/backends/vendor/ascend/impl/fused_moe.py index f217627fb..62eac26ff 100644 --- a/vllm_fl/dispatch/backends/vendor/ascend/impl/fused_moe.py +++ b/vllm_fl/dispatch/backends/vendor/ascend/impl/fused_moe.py @@ -4,15 +4,121 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import logging +import os from typing import Optional import torch import torch.nn.functional as F import torch_npu -from flag_gems.runtime.backend._ascend import fused +logger = logging.getLogger(__name__) -def _torch_fused_experts_impl( +_ASCENDC_MOE_AVAILABLE: bool | None = None + + +def ascendc_moe_available() -> bool: + """Whether the AscendC fused-MoE custom ops can be used. + + Requires the packaged CANN custom-op package (same one the GDN patch + bootstraps) and the ``_C_ascend`` torch bindings. Note: the custom + ``moe_grouped_matmul`` kernel is *not* used on purpose — it crashes the + aicore on ascend910b even for trivial single-expert inputs (CCU + instruction address check error), so the grouped matmuls stay on + ``torch_npu.npu_grouped_matmul``. Set ``VLLM_FL_DISABLE_ASCENDC_MOE=1`` + to keep the FlagGems/torch_npu path. + """ + global _ASCENDC_MOE_AVAILABLE + if _ASCENDC_MOE_AVAILABLE is not None: + return _ASCENDC_MOE_AVAILABLE + if os.environ.get("VLLM_FL_DISABLE_ASCENDC_MOE", "0") == "1": + logger.info("VLLM_FL_DISABLE_ASCENDC_MOE=1, keep FlagGems/torch_npu MoE path") + _ASCENDC_MOE_AVAILABLE = False + return False + try: + from ..patches.patch_qwen3_6_gdn import _bootstrap_custom_op_env + + if not _bootstrap_custom_op_env(): + _ASCENDC_MOE_AVAILABLE = False + return False + except Exception as e: + logger.warning("CANN custom op bootstrap failed: %s", e) + _ASCENDC_MOE_AVAILABLE = False + return False + missing = [ + name + for name in ("moe_gating_top_k", "npu_moe_init_routing_custom") + if not hasattr(torch.ops._C_ascend, name) + ] + if missing: + logger.warning("torch.ops._C_ascend missing ops %s; keep torch_npu MoE path", missing) + _ASCENDC_MOE_AVAILABLE = False + return False + _ASCENDC_MOE_AVAILABLE = True + return True + + +def fused_topk_ascend( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + indices_type: torch.dtype | None = None, +) -> tuple[torch.Tensor, torch.Tensor, None]: + """Drop-in replacement for vllm_fl ``fused_topk`` (softmax routing only). + + Uses the fused AscendC ``moe_gating_top_k`` kernel (softmax + top-k + + optional L1 renorm in one launch) instead of the FlagGems Triton + ``topk_softmax``. Applies to the plain ``RoutingMethodType.Renormalize`` + routing used by Qwen3.5/Qwen3.6 (no grouped top-k, no correction bias); + other routing modes never reach ``fused_topk``. + """ + assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch" + + # fp32 input keeps the softmax/renorm precision identical to the + # FlagGems path and makes the returned topk_weights fp32 as before. + topk_weights, topk_ids, _ = torch.ops._C_ascend.moe_gating_top_k( + gating_output.to(torch.float32), + topk, + 1, # k_group: no grouped top-k + 1, # group_count: no grouped top-k + 0, # group_select_mode + int(renormalize), # renorm: 1 = L1 renorm of the top-k weights + 0, # norm_type: softmax + False, # out_flag + 1.0, # routed_scaling_factor + 1e-20, # eps + None, # bias_opt + ) + if indices_type is not None: + topk_ids = topk_ids.to(indices_type) + return topk_weights, topk_ids, None + + +def convert_moe_weights_pretransposed(layer: torch.nn.Module) -> None: + """Transpose FusedMoE expert weights once at load time. + + The legacy path pays a ``w.transpose(1, 2).contiguous()`` copy for both + grouped matmuls on every forward (e.g. 2 x 128MB per layer per step for + Qwen3.6-35B-A3B TP=4). Store them pre-transposed instead: + + w13_weight: [E, 2*intermediate, hidden] -> [E, hidden, 2*intermediate] + w2_weight: [E, hidden, intermediate] -> [E, intermediate, hidden] + + Idempotent; only meaningful together with ``_ascendc_fused_experts_impl`` + (which detects the layout from the weight shapes). + """ + w13 = layer.w13_weight.data + w2 = layer.w2_weight.data + # After conversion w13.size(1) == hidden == w2.size(2); before, they are + # 2*intermediate and intermediate respectively. + if w13.size(1) == w2.size(2): + return + layer.w13_weight.data = w13.transpose(1, 2).contiguous() + layer.w2_weight.data = w2.transpose(1, 2).contiguous() + + +def _ascendc_fused_experts_impl( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, @@ -22,80 +128,245 @@ def _torch_fused_experts_impl( activation: str = "silu", apply_router_weight_on_input: bool = False, global_num_experts: int = -1, - expert_map: torch.Tensor|None = None, + expert_map: torch.Tensor | None = None, ) -> torch.Tensor: - """Pure PyTorch implementation of fused MoE experts for NPU. + """AscendC fused MoE experts implementation. + + Same structure as ``_torch_fused_experts_impl`` but the token + routing/permute runs on the AscendC ``npu_moe_init_routing_custom`` + kernel and the weights are pre-transposed (see + ``convert_moe_weights_pretransposed``), so the grouped matmuls consume + them directly: + + w1: [E, hidden, 2*intermediate] (pre-transposed) + w2: [E, intermediate, hidden] (pre-transposed) - This avoids the Triton fused_moe_kernel which has compatibility issues - on Ascend NPU hardware. + ``expanded_row_idx`` (row_idx_type=0) is the gather map expected by + ``npu_moe_token_unpermute`` (``|.|`` guards inactive rows), so the + scatter-based inverse permutation of the legacy path is not needed. + All operations are graph-safe. """ num_tokens, hidden_dim = hidden_states.size() - E, N, _ = w1.size() # w1: [E, N, K_in] - K = w2.size(1) # w2: [E, K_out, N//2] + E = w1.size(0) top_k = topk_ids.size(1) if global_num_experts == -1: - global_num_experts = E + global_num_experts = len(expert_map) if expert_map is not None else E - if inplace: - out_hidden_states = hidden_states - else: - out_hidden_states = torch.zeros_like(hidden_states) - - # Map global expert ids to local expert ids + # Map global expert ids to local expert ids. Out-of-range entries are + # clamped to 0 and masked to zero weight so the graph shape stays static. if expert_map is not None: - local_topk_ids = expert_map[topk_ids.long()] + mask = expert_map[topk_ids.long()] != -1 + local_topk_ids = expert_map[topk_ids.long()].clamp(min=0) + topk_weights = topk_weights * mask.to(topk_weights.dtype) else: local_topk_ids = topk_ids.long() - # Process each expert - for expert_idx in range(E): - # Find which (token, k) pairs are assigned to this expert - mask = (local_topk_ids == expert_idx) # [num_tokens, top_k] - if not mask.any(): - continue + # Expand tokens according to top-k expert assignment and sort them by + # expert. expanded_row_idx maps each sorted row back to the original + # flat (token*topk + k) position (row_idx_type=0). + expanded_x, expanded_row_idx, expert_token_count, _ = ( + torch.ops._C_ascend.npu_moe_init_routing_custom( + hidden_states, + local_topk_ids.to(torch.int32), + active_num=num_tokens * top_k, + expert_num=global_num_experts, + drop_pad_mode=0, + expert_tokens_num_type=1, # count mode + expert_tokens_num_flag=True, + quant_mode=-1, + active_expert_range=[0, E], + row_idx_type=0, + ) + ) + + # Apply router weight on the expanded tokens if requested. + if apply_router_weight_on_input: + expanded_weights = ( + topk_weights.view(-1)[expanded_row_idx.abs().long()] + .unsqueeze(-1) + .to(expanded_x.dtype) + ) + expanded_x = expanded_x * expanded_weights + probs = None + else: + probs = topk_weights + + # Gate-up grouped matmul on the pre-transposed weight: + # x [total_tokens*top_k, hidden], weight [E, hidden, 2*intermediate]. + gate_up = torch_npu.npu_grouped_matmul( + [expanded_x], + [w1], + group_list=expert_token_count, + split_item=2, + group_type=0, + group_list_type=1, + )[0] + + # Activation. + if activation == "silu": + gate_up = torch_npu.npu_swiglu(gate_up) + elif activation == "gelu": + d = gate_up.shape[-1] // 2 + gate_up = F.gelu(gate_up[..., :d]) * gate_up[..., d:] + elif activation == "silu_no_mul": + gate_up = F.silu(gate_up) + elif activation == "gelu_no_mul": + gate_up = F.gelu(gate_up) + else: + raise ValueError(f"Unsupported FusedMoe activation: {activation}.") + + # Down grouped matmul on the pre-transposed weight: + # x [total_tokens*top_k, intermediate], weight [E, intermediate, hidden]. + down = torch_npu.npu_grouped_matmul( + [gate_up], + [w2], + group_list=expert_token_count, + split_item=2, + group_type=0, + group_list_type=1, + )[0] + + # Scatter/sum the expert outputs back to the token dimension. + out = torch_npu.npu_moe_token_unpermute( + permuted_tokens=down, + sorted_indices=expanded_row_idx.abs(), + probs=probs, + ) + + if inplace: + hidden_states.copy_(out) + return hidden_states + return out + + +def _torch_fused_experts_impl( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + inplace: bool = False, + activation: str = "silu", + apply_router_weight_on_input: bool = False, + global_num_experts: int = -1, + expert_map: torch.Tensor | None = None, +) -> torch.Tensor: + """Ascend NPU native fused MoE experts implementation. + + This implementation replaces the previous pure-PyTorch loop with + ``npu_moe_init_routing_v2`` + ``npu_grouped_matmul`` + + ``npu_moe_token_unpermute``. All operations are graph-safe (no host-side + synchronization or data-dependent control flow), so the whole MoE layer can + be captured inside ``torch.npu.NPUGraph`` / ``torch.compile(fullgraph=True)``. + + Weight layout expected from vLLM ``FusedMoE``: + - w1: [E, 2*intermediate, hidden] + - w2: [E, hidden, intermediate] + For grouped matmul we transpose the last two dims so each expert weight + becomes [hidden, out_features]. + """ + num_tokens, hidden_dim = hidden_states.size() + E = w1.size(0) + top_k = topk_ids.size(1) - # Get token indices and their k-slot indices - token_indices, k_indices = torch.where(mask) + if global_num_experts == -1: + global_num_experts = len(expert_map) if expert_map is not None else E - # Gather the hidden states for these tokens - expert_input = hidden_states[token_indices] # [n, hidden_dim] + # Map global expert ids to local expert ids. Out-of-range entries are + # clamped to 0 and masked to zero weight so the graph shape stays static. + if expert_map is not None: + mask = expert_map[topk_ids.long()] != -1 + local_topk_ids = expert_map[topk_ids.long()].clamp(min=0) + topk_weights = topk_weights * mask.to(topk_weights.dtype) + else: + local_topk_ids = topk_ids.long() + + # Expand tokens according to top-k expert assignment and sort them by + # expert. row_idx maps each sorted row back to the original flat + # (token*topk + k) position. + expanded_x, row_idx, expert_token_count, _ = torch_npu.npu_moe_init_routing_v2( + hidden_states, + local_topk_ids.to(torch.int32), + active_num=num_tokens * top_k, + expert_num=global_num_experts, + expert_tokens_num_type=1, # count mode + expert_tokens_num_flag=True, + quant_mode=-1, + active_expert_range=[0, E], + row_idx_type=1, + ) - # Apply router weight on input if needed - if apply_router_weight_on_input: - weights = topk_weights[token_indices, k_indices].unsqueeze(-1) - expert_input = expert_input * weights.to(expert_input.dtype) + # Apply router weight on the expanded tokens if requested. + if apply_router_weight_on_input: + expanded_weights = ( + topk_weights.view(-1)[row_idx.long()] + .unsqueeze(-1) + .to(expanded_x.dtype) + ) + expanded_x = expanded_x * expanded_weights + probs = None + else: + probs = topk_weights - # First matmul: expert_input @ w1[expert_idx].T - # w1[expert_idx] shape: [N, hidden_dim], result: [n, N] - gate_up = torch.mm(expert_input, w1[expert_idx].t()) + # npu_moe_token_unpermute expects ``sorted_indices`` as the *gather* index + # (sorted_position -> original flat position). Compute the inverse of + # row_idx on-device. + sorted_indices = torch.empty_like(row_idx) + sorted_indices.scatter_( + 0, + row_idx.long(), + torch.arange(row_idx.numel(), device=row_idx.device, dtype=torch.int32), + ) - # Activation (pure PyTorch to avoid Triton kernel issues on NPU) - if activation == "silu": - d = gate_up.shape[-1] // 2 - gate_up = F.silu(gate_up[..., :d]) * gate_up[..., d:] - elif activation == "gelu": - gate_up = torch_npu.npu_gelu_mul(gate_up) - elif activation == "silu_no_mul": - gate_up = F.silu(gate_up) - elif activation == "gelu_no_mul": - gate_up = torch_npu.npu_gelu(gate_up) - else: - raise ValueError(f"Unsupported FusedMoe activation: {activation}.") + # Gate-up grouped matmul: x shape [total_tokens*top_k, hidden] + # weight shape [E, hidden, 2*intermediate]. + w1_t = w1.transpose(1, 2).contiguous() + gate_up = torch_npu.npu_grouped_matmul( + [expanded_x], + [w1_t], + group_list=expert_token_count, + split_item=2, + group_type=0, + group_list_type=1, + )[0] - # Second matmul: activated @ w2[expert_idx].T - # w2[expert_idx] shape: [K_out, N//2], result: [n, K_out] - expert_output = torch.mm(gate_up, w2[expert_idx].t()) + # Activation. + if activation == "silu": + gate_up = torch_npu.npu_swiglu(gate_up) + elif activation == "gelu": + d = gate_up.shape[-1] // 2 + gate_up = F.gelu(gate_up[..., :d]) * gate_up[..., d:] + elif activation == "silu_no_mul": + gate_up = F.silu(gate_up) + elif activation == "gelu_no_mul": + gate_up = F.gelu(gate_up) + else: + raise ValueError(f"Unsupported FusedMoe activation: {activation}.") - # Apply router weight on output if not applied on input - if not apply_router_weight_on_input: - weights = topk_weights[token_indices, k_indices].unsqueeze(-1) - expert_output = expert_output * weights.to(expert_output.dtype) + # Down grouped matmul: x shape [total_tokens*top_k, intermediate] + # weight shape [E, intermediate, hidden]. + w2_t = w2.transpose(1, 2).contiguous() + down = torch_npu.npu_grouped_matmul( + [gate_up], + [w2_t], + group_list=expert_token_count, + split_item=2, + group_type=0, + group_list_type=1, + )[0] - # Accumulate results - out_hidden_states.index_add_(0, token_indices, expert_output) + # Scatter/sum the expert outputs back to the token dimension. + out = torch_npu.npu_moe_token_unpermute( + permuted_tokens=down, + sorted_indices=sorted_indices, + probs=probs, + ) - return out_hidden_states + if inplace: + hidden_states.copy_(out) + return hidden_states + return out def fused_experts_impl( @@ -124,6 +395,30 @@ def fused_experts_impl( w1_bias: Optional[torch.Tensor] = None, w2_bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: + # Unquantized weights that went through convert_moe_weights_pretransposed + # are stored as [E, hidden, ...] (hidden_states.size(1) == w1.size(1)); + # the AscendC custom-op path consumes them directly. + is_unquantized = not ( + use_fp8_w8a8 or use_int8_w8a8 or use_int8_w8a16 or use_int4_w4a16 + ) + if ( + is_unquantized + and hidden_states.size(1) == w1.size(1) + and ascendc_moe_available() + ): + return _ascendc_fused_experts_impl( + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=inplace, + activation=activation, + apply_router_weight_on_input=apply_router_weight_on_input, + global_num_experts=global_num_experts, + expert_map=expert_map, + ) + # Check constraints. if use_int4_w4a16: assert hidden_states.size(1) // 2 == w1.size(2), "Hidden size mismatch" @@ -138,8 +433,9 @@ def fused_experts_impl( assert w2.stride(-1) == 1, "Stride of last dimension must be 1" assert hidden_states.dtype in [torch.float32, torch.float16, torch.bfloat16] - # Use pure-torch implementation on NPU to avoid Triton kernel - # compatibility issues with the Ascend backend. + # Quantized MoE is not handled here; fall back would require a different + # code path. For the graph-mode Qwen3.6-35B-A3B serving scenario we only + # need the unquantized path. return _torch_fused_experts_impl( hidden_states=hidden_states, w1=w1, diff --git a/vllm_fl/dispatch/backends/vendor/ascend/impl/linearnorm/__init__.py b/vllm_fl/dispatch/backends/vendor/ascend/impl/linearnorm/__init__.py new file mode 100644 index 000000000..d9ef92054 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/impl/linearnorm/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2026 BAAI. All rights reserved. diff --git a/vllm_fl/dispatch/backends/vendor/ascend/impl/linearnorm/layernorm_gated.py b/vllm_fl/dispatch/backends/vendor/ascend/impl/linearnorm/layernorm_gated.py new file mode 100644 index 000000000..077f5c4d0 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/impl/linearnorm/layernorm_gated.py @@ -0,0 +1,246 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Adapted from https://github.com/vllm-project/vllm-ascend/blob/main/vllm_ascend/ops/triton/layernorm_gated.py +# Adapt from https://github.com/fla-org/flash-linear-attention/blob/main/fla/modules/layernorm_gated.py +# Copyright (c) 2024, Tri Dao. +# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html +# mypy: ignore-errors + +import torch +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import direct_register_custom_op + + +@triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["Z"] is not None}) +@triton.jit(do_not_specialize=["stride_x_row", "stride_y_row", "stride_z_row", "M", "N", "eps"]) +def _layer_norm_fwd_1pass_kernel_npu( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Z, # pointer to the other branch + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_y_row, + stride_z_row, + M, # number of rows in X + N, # number of columns in X + eps, # epsilon to avoid division by zero + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_Z: tl.constexpr, + NORM_BEFORE_GATE: tl.constexpr, + IS_RMS_NORM: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + pid_m = tl.program_id(0) + group = tl.program_id(1) + if not IS_RMS_NORM: + Mean += group * M + Rstd += group * M + W += group * N + if HAS_BIAS: + B += group * N + + # Compute row indices for this program + rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.arange(0, BLOCK_N) + + # Mask for valid rows and cols + row_mask = rows < M + col_mask = cols < N + + # Load weight once (broadcasted over rows) + w = tl.load(W + cols, mask=col_mask).to(tl.float32) + if HAS_BIAS: + b = tl.load(B + cols, mask=col_mask).to(tl.float32) + + # Load X: shape [BLOCK_M, BLOCK_N] + x_ptrs = X + rows[:, None] * stride_x_row + cols[None, :] + group * N + x = tl.load(x_ptrs, mask=row_mask[:, None] & col_mask[None, :]).to(tl.float32) + + # Load Z if needed + if HAS_Z: + z_ptrs = Z + rows[:, None] * stride_z_row + cols[None, :] + group * N + z = tl.load(z_ptrs, mask=row_mask[:, None] & col_mask[None, :]).to(tl.float32) + if not NORM_BEFORE_GATE: + x *= z * tl.sigmoid(z) + + # Compute statistics per row + if not IS_RMS_NORM: + mean = tl.sum(x, axis=1) / N # [BLOCK_M] + xbar = tl.where(col_mask[None, :], x - mean[:, None], 0.0) + var = tl.sum(xbar * xbar, axis=1) / N + tl.store(Mean + rows, mean, mask=row_mask) + else: + xbar = tl.where(col_mask[None, :], x, 0.0) + var = tl.sum(xbar * xbar, axis=1) / N + + rstd = 1.0 / tl.sqrt(var + eps) # [BLOCK_M] + tl.store(Rstd + rows, rstd, mask=row_mask) + + # Normalize + if not IS_RMS_NORM: + x_hat = (x - mean[:, None]) * rstd[:, None] + else: + x_hat = x * rstd[:, None] + + y = x_hat * w[None, :] + if HAS_BIAS: + y += b[None, :] + + # Post-gate + if HAS_Z and NORM_BEFORE_GATE: + y *= z * tl.sigmoid(z) + + # Store output + y_ptrs = Y + rows[:, None] * stride_y_row + cols[None, :] + group * N + tl.store(y_ptrs, y, mask=row_mask[:, None] & col_mask[None, :]) + + +def layer_norm_fwd_npu( + x, + weight, + bias, + eps, + z=None, + out=None, + group_size=None, + norm_before_gate=True, + is_rms_norm=False, +): + M, N = x.shape + if group_size is None: + group_size = N + assert N % group_size == 0 + ngroups = N // group_size + + assert x.stride(-1) == 1 + if z is not None: + assert z.stride(-1) == 1 + assert z.shape == (M, N) + assert weight.shape == (N,) + assert weight.stride(-1) == 1 + if bias is not None: + assert bias.stride(-1) == 1 + assert bias.shape == (N,) + # allocate output + if out is not None: + assert out.shape == x.shape + else: + out = torch.empty_like(x) + assert out.stride(-1) == 1 + mean = torch.empty((ngroups * M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None + rstd = torch.empty((ngroups * M,), dtype=torch.float32, device=x.device) + + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size)) + if group_size > BLOCK_N: + raise RuntimeError("Feature dim too large.") + + # Choose BLOCK_M: e.g., 16, 32, 64 — depends on NPU vector core capacity. + # NOTE: 64 overflows the 192KB unified buffer on 910B with the + # multi-buffering enabled by the Ascend Triton pipeline here + # (requires 1710080 bits > 1572864 bits available), keep 32. + BLOCK_M = 32 # Tune this based on your NPU's register/shared memory + + # Now grid is (num blocks over M, num groups) + grid = (triton.cdiv(M, BLOCK_M), ngroups) + _layer_norm_fwd_1pass_kernel_npu[grid]( + x, + out, + weight, + bias, + z, + mean, + rstd, + x.stride(0), + out.stride(0), + z.stride(0) if z is not None else 0, + M, + group_size, + eps, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + NORM_BEFORE_GATE=norm_before_gate, + IS_RMS_NORM=is_rms_norm, + # Remove multibuffer if not needed + ) + return out, mean, rstd + + +def rmsnorm_gated_fwd( + x: torch.Tensor, + weight: torch.Tensor, + z: torch.Tensor | None = None, + eps: float = 1e-5, + group_size: int | None = None, + norm_before_gate: bool = True, +) -> torch.Tensor: + """Fused RMSNormGated forward (inference only). + + If z is not None: out = norm(x) * silu(z) when norm_before_gate, + else out = norm(x * silu(z)). + """ + x_shape_og = x.shape + # reshape input data into 2D tensor + x = x.reshape(-1, x.shape[-1]) + if x.stride(-1) != 1: + x = x.contiguous() + if z is not None: + assert z.shape == x_shape_og + z = z.reshape(-1, z.shape[-1]) + if z.stride(-1) != 1: + z = z.contiguous() + y, _, _ = layer_norm_fwd_npu( + x, + weight.contiguous(), + None, + eps, + z=z, + group_size=group_size, + norm_before_gate=norm_before_gate, + is_rms_norm=True, + ) + return y.reshape(x_shape_og) + + +def _rmsnorm_gated_fwd_fake( + x: torch.Tensor, + weight: torch.Tensor, + z: torch.Tensor | None = None, + eps: float = 1e-5, + group_size: int | None = None, + norm_before_gate: bool = True, +) -> torch.Tensor: + return torch.empty_like(x) + + +# Register the fused kernel as a vLLM custom op so Dynamo treats it as an +# opaque node: torch 2.8 cannot trace user-defined Triton kernel launches +# on Ascend (NPUDriver.get_current_device is skip-listed), so calling the +# Triton wrapper directly from compiled model code crashes tracing. +direct_register_custom_op( + op_name="rmsnorm_gated_fwd", + op_func=rmsnorm_gated_fwd, + mutates_args=[], + fake_impl=_rmsnorm_gated_fwd_fake, +) + + +def rmsnorm_gated_oot( + x: torch.Tensor, + weight: torch.Tensor, + z: torch.Tensor | None = None, + eps: float = 1e-5, + group_size: int | None = None, + norm_before_gate: bool = True, +) -> torch.Tensor: + """Dynamo-safe entry point: dispatch through the registered custom op + when available (graph mode), otherwise call the Triton wrapper.""" + op = getattr(torch.ops.vllm, "rmsnorm_gated_fwd", None) + if op is not None: + return op(x, weight, z, eps, group_size, norm_before_gate) + return rmsnorm_gated_fwd(x, weight, z, eps, group_size, norm_before_gate) diff --git a/vllm_fl/dispatch/backends/vendor/ascend/impl/linearnorm/split_qkv_rmsnorm_mrope.py b/vllm_fl/dispatch/backends/vendor/ascend/impl/linearnorm/split_qkv_rmsnorm_mrope.py new file mode 100644 index 000000000..f6e6e1ca4 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/impl/linearnorm/split_qkv_rmsnorm_mrope.py @@ -0,0 +1,429 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# Copyright (c) 2026 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Adapted from https://github.com/vllm-project/vllm-ascend/blob/main/vllm_ascend/ops/triton/linearnorm/split_qkv_rmsnorm_mrope.py +# mypy: ignore-errors + +import torch +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import direct_register_custom_op + +from ..triton_utils import extract_slice, get_vectorcore_num, insert_slice + + +@triton.jit( + do_not_specialize=["num_tokens", "front_core_num", "num_tokens_each_front_core", "num_tokens_each_tail_core"] +) +def split_qkv_rmsnorm_mrope_kernel( + in_qkv_ptr: torch.Tensor, + q_weight_ptr: torch.Tensor, + q_bias_ptr: torch.Tensor, + k_weight_ptr: torch.Tensor, + k_bias_ptr: torch.Tensor, + cos_sin_ptr: torch.Tensor, + out_q_ptr: torch.Tensor, + out_k_ptr: torch.Tensor, + out_v_ptr: torch.Tensor, + out_gate_ptr: torch.Tensor, + num_tokens, + front_core_num, + num_tokens_each_front_core, + num_tokens_each_tail_core, + num_q_heads: tl.constexpr, + num_kv_heads: tl.constexpr, + head_size: tl.constexpr, + q_size: tl.constexpr, + kv_size: tl.constexpr, + eps: tl.constexpr, + mrope_section_t, + mrope_section_h, + mrope_section_w, + has_bias: tl.constexpr, + is_interleaved: tl.constexpr, + rope_dim: tl.constexpr, + half_rope_dim: tl.constexpr, + IS_PARTIAL_ROPE: tl.constexpr, + gate_size: tl.constexpr, +): + block_idx = tl.program_id(0) + + loop_num = num_tokens_each_front_core + if block_idx >= front_core_num: + loop_num = num_tokens_each_tail_core + + block_offset = num_tokens_each_front_core * block_idx + if block_idx >= front_core_num: + block_offset = ( + num_tokens_each_front_core * front_core_num + (block_idx - front_core_num) * num_tokens_each_tail_core + ) + + q_rmsnorm_weight = tl.load(q_weight_ptr + tl.arange(0, head_size)) + k_rmsnorm_weight = tl.load(k_weight_ptr + tl.arange(0, head_size)) + + if has_bias: + q_bias = tl.load(q_bias_ptr + tl.arange(0, head_size)) + k_bias = tl.load(k_bias_ptr + tl.arange(0, head_size)) + + for index in range(loop_num): + ## load ## + # q + in_q_offset = in_qkv_ptr + (block_offset + index) * (q_size + gate_size + 2 * kv_size) + if gate_size > 0: + in_q_gate_tensor = ( + tl.load(in_q_offset + tl.arange(0, q_size + gate_size)) + .to(tl.float32) + .reshape(num_q_heads, head_size * 2) + ) + in_q_tensor = extract_slice( + in_q_gate_tensor, + offsets=(0, 0), + sizes=(num_q_heads, head_size), + strides=(1, 1), + ) + in_gate_tensor = extract_slice( + in_q_gate_tensor, + offsets=(0, head_size), + sizes=(num_q_heads, head_size), + strides=(1, 1), + ).reshape(q_size) + else: + in_q_tensor = tl.load(in_q_offset + tl.arange(0, q_size)).to(tl.float32).reshape(num_q_heads, head_size) + + # k + in_k_offset = in_q_offset + q_size + gate_size + in_k_tensor = tl.load(in_k_offset + tl.arange(0, kv_size)).to(tl.float32).reshape(num_kv_heads, head_size) + # v + in_v_offset = in_k_offset + kv_size + in_v_tensor = tl.load(in_v_offset + tl.arange(0, kv_size)) + + # cos, sin + cos_offsets = tl.arange(0, half_rope_dim) + if is_interleaved: + h_mask = ((cos_offsets % 3) == 1) & (cos_offsets <= 3 * mrope_section_h) + w_mask = ((cos_offsets % 3) == 2) & (cos_offsets <= 3 * mrope_section_w) + t_mask = ~(h_mask | w_mask) + else: + t_mask = cos_offsets < mrope_section_t + h_mask = (mrope_section_t - 1 < cos_offsets) & (cos_offsets < mrope_section_t + mrope_section_h) + w_mask = (mrope_section_t + mrope_section_h - 1 < cos_offsets) & ( + cos_offsets < mrope_section_t + mrope_section_h + mrope_section_w + ) + + t_cos_offset = cos_sin_ptr + (block_offset + index) * rope_dim + h_cos_offset = t_cos_offset + num_tokens * rope_dim + w_cos_offset = h_cos_offset + num_tokens * rope_dim + + t_sin_offset = cos_sin_ptr + (block_offset + index) * rope_dim + half_rope_dim + h_sin_offset = t_sin_offset + num_tokens * rope_dim + w_sin_offset = h_sin_offset + num_tokens * rope_dim + + t_cos_tensor = tl.load(t_cos_offset + cos_offsets, mask=t_mask, other=0) + h_cos_tensor = tl.load(h_cos_offset + cos_offsets, mask=h_mask, other=0) + w_cos_tensor = tl.load(w_cos_offset + cos_offsets, mask=w_mask, other=0) + t_sin_tensor = tl.load(t_sin_offset + cos_offsets, mask=t_mask, other=0) + h_sin_tensor = tl.load(h_sin_offset + cos_offsets, mask=h_mask, other=0) + w_sin_tensor = tl.load(w_sin_offset + cos_offsets, mask=w_mask, other=0) + + cos_tensor = (t_cos_tensor + h_cos_tensor + w_cos_tensor).to(tl.float32).reshape(1, half_rope_dim) + cos_tensor = tl.broadcast_to(cos_tensor, (2, half_rope_dim)).reshape(1, rope_dim) + + sin_tensor = (t_sin_tensor + h_sin_tensor + w_sin_tensor).to(tl.float32).reshape(1, half_rope_dim) + sin_tensor = tl.broadcast_to(sin_tensor, (2, half_rope_dim)).reshape(1, rope_dim) + + ## compute ## + # q-rmsnorm + squares = in_q_tensor * in_q_tensor + variances = tl.sum(squares, axis=1) / head_size + reciprocal_std = (1 / tl.sqrt(variances + eps)).reshape(num_q_heads, 1) + q_normalized = in_q_tensor * reciprocal_std + q_normalized = q_normalized * q_rmsnorm_weight + if has_bias: + q_normalized = q_normalized + q_bias + + # k-rmsnorm + squares = in_k_tensor * in_k_tensor + variances = tl.sum(squares, axis=1) / head_size + reciprocal_std = (1 / tl.sqrt(variances + eps)).reshape(num_kv_heads, 1) + k_normalized = in_k_tensor * reciprocal_std + k_normalized = k_normalized * k_rmsnorm_weight + if has_bias: + k_normalized = k_normalized + k_bias + + # q-mrope + x1 = extract_slice( + q_normalized, + offsets=(0, 0), + sizes=(num_q_heads, half_rope_dim), + strides=(1, 1), + ) + x2 = extract_slice( + q_normalized, + offsets=(0, half_rope_dim), + sizes=(num_q_heads, half_rope_dim), + strides=(1, 1), + ) + cat_x = tl.zeros((num_q_heads, rope_dim), dtype=tl.float32) + cat_x = insert_slice( + cat_x, + -x2, + offsets=(0, 0), + sizes=(num_q_heads, half_rope_dim), + strides=(1, 1), + ) + cat_x = insert_slice( + cat_x, + x1, + offsets=(0, half_rope_dim), + sizes=(num_q_heads, half_rope_dim), + strides=(1, 1), + ) + if IS_PARTIAL_ROPE: + orig_qk = extract_slice( + q_normalized, + offsets=(0, 0), + sizes=(num_q_heads, rope_dim), + strides=(1, 1), + ) + else: + orig_qk = q_normalized + roped_q = cat_x * sin_tensor + orig_qk * cos_tensor + + # k-mrope + y1 = extract_slice( + k_normalized, + offsets=(0, 0), + sizes=(num_kv_heads, half_rope_dim), + strides=(1, 1), + ) + y2 = extract_slice( + k_normalized, + offsets=(0, half_rope_dim), + sizes=(num_kv_heads, half_rope_dim), + strides=(1, 1), + ) + cat_y = tl.zeros((num_kv_heads, rope_dim), dtype=tl.float32) + cat_y = insert_slice( + cat_y, + -y2, + offsets=(0, 0), + sizes=(num_kv_heads, half_rope_dim), + strides=(1, 1), + ) + cat_y = insert_slice( + cat_y, + y1, + offsets=(0, half_rope_dim), + sizes=(num_kv_heads, half_rope_dim), + strides=(1, 1), + ) + if IS_PARTIAL_ROPE: + orig_qk = extract_slice( + k_normalized, + offsets=(0, 0), + sizes=(num_kv_heads, rope_dim), + strides=(1, 1), + ) + else: + orig_qk = k_normalized + roped_k = cat_y * sin_tensor + orig_qk * cos_tensor + + if IS_PARTIAL_ROPE: + q_normalized = insert_slice( + q_normalized, + roped_q, + offsets=(0, 0), + sizes=(num_q_heads, rope_dim), + strides=(1, 1), + ) + k_normalized = insert_slice( + k_normalized, + roped_k, + offsets=(0, 0), + sizes=(num_kv_heads, rope_dim), + strides=(1, 1), + ) + else: + q_normalized = roped_q + k_normalized = roped_k + + ## store ## + # out_q + out_q_offset = out_q_ptr + (block_offset + index) * q_size + out_q_indices = tl.arange(0, q_size) + tl.store(out_q_offset + out_q_indices, q_normalized.reshape(q_size)) + + # out_k + out_k_offset = out_k_ptr + (block_offset + index) * kv_size + out_k_indices = tl.arange(0, kv_size) + tl.store(out_k_offset + out_k_indices, k_normalized.reshape(kv_size)) + + # out_v + out_v_offset = out_v_ptr + (block_offset + index) * kv_size + tl.store(out_v_offset + tl.arange(0, kv_size), in_v_tensor) + + # out_gate + if gate_size > 0: + out_gate_offset = out_gate_ptr + (block_offset + index) * gate_size + tl.store(out_gate_offset + tl.arange(0, gate_size), in_gate_tensor) + + +def triton_split_qkv_rmsnorm_mrope( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin: torch.Tensor, + num_q_heads: int, + num_kv_heads: int, + head_size: int, + eps: float, + mrope_section: list[int], + is_interleaved: bool, + rope_dim: int | None = None, + q_bias: torch.Tensor | None = None, + k_bias: torch.Tensor | None = None, + has_gate: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + core_num = get_vectorcore_num() + + q_size = num_q_heads * head_size + kv_size = num_kv_heads * head_size + num_tokens = qkv.shape[0] + + gate_size = q_size if has_gate else 0 + + if rope_dim is None: + rope_dim = head_size + IS_PARTIAL_ROPE = rope_dim != head_size + + front_core_num = core_num + if num_tokens % core_num != 0: + front_core_num = num_tokens % core_num + + num_tokens_each_front_core = (num_tokens + core_num - 1) // core_num + + tail_core_num = 0 + if num_tokens > core_num: + tail_core_num = core_num - front_core_num + + num_tokens_each_tail_core = num_tokens // core_num + + q_output = torch.empty(num_tokens, q_size, device=qkv.device, dtype=qkv.dtype) + k_output = torch.empty(num_tokens, kv_size, device=qkv.device, dtype=qkv.dtype) + v_output = torch.empty(num_tokens, kv_size, device=qkv.device, dtype=qkv.dtype) + gate_output = torch.empty(num_tokens, gate_size, device=qkv.device, dtype=qkv.dtype) + + total_core = front_core_num + tail_core_num + block_dim = core_num + if total_core < core_num: + block_dim = total_core + + has_bias = q_bias is not None + + split_qkv_rmsnorm_mrope_kernel[(block_dim,)]( + qkv, + q_weight, + q_bias, + k_weight, + k_bias, + cos_sin, + q_output, + k_output, + v_output, + gate_output, + num_tokens, + front_core_num, + num_tokens_each_front_core, + num_tokens_each_tail_core, + num_q_heads, + num_kv_heads, + head_size, + q_size, + kv_size, + eps, + mrope_section[0], + mrope_section[1], + mrope_section[2], + has_bias, + is_interleaved, + rope_dim, + rope_dim // 2, + IS_PARTIAL_ROPE, + gate_size, + ) + + return q_output, k_output, v_output, gate_output + + +def triton_split_qkv_rmsnorm_mrope_fake( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin: torch.Tensor, + num_q_heads: int, + num_kv_heads: int, + head_size: int, + eps: float, + mrope_section: list[int], + is_interleaved: bool, + rope_dim: int | None = None, + q_bias: torch.Tensor | None = None, + k_bias: torch.Tensor | None = None, + has_gate: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + num_tokens = qkv.shape[0] + q_size = num_q_heads * head_size + kv_size = num_kv_heads * head_size + gate_size = q_size if has_gate else 0 + + q_output = torch.empty( + num_tokens, + q_size, + device=qkv.device, + dtype=qkv.dtype, + ) + + k_output = torch.empty( + num_tokens, + kv_size, + device=qkv.device, + dtype=qkv.dtype, + ) + + v_output = torch.empty( + num_tokens, + kv_size, + device=qkv.device, + dtype=qkv.dtype, + ) + + gate_output = torch.empty( + num_tokens, + gate_size, + device=qkv.device, + dtype=qkv.dtype, + ) + + return q_output, k_output, v_output, gate_output + + +if not hasattr(torch.ops.vllm, "triton_split_qkv_rmsnorm_mrope"): + direct_register_custom_op( + op_name="triton_split_qkv_rmsnorm_mrope", + op_func=triton_split_qkv_rmsnorm_mrope, + fake_impl=triton_split_qkv_rmsnorm_mrope_fake, + mutates_args=[], + dispatch_key="PrivateUse1", + ) diff --git a/vllm_fl/dispatch/backends/vendor/ascend/impl/triton_utils.py b/vllm_fl/dispatch/backends/vendor/ascend/impl/triton_utils.py index 0d692bb29..14cc02c42 100644 --- a/vllm_fl/dispatch/backends/vendor/ascend/impl/triton_utils.py +++ b/vllm_fl/dispatch/backends/vendor/ascend/impl/triton_utils.py @@ -4,10 +4,48 @@ from typing import Any, Dict import torch -from vllm.triton_utils import HAS_TRITON, triton +from vllm.triton_utils import HAS_TRITON, tl, triton _NUM_AICORE = -1 _NUM_VECTORCORE = -1 +_extension_module = None + +if HAS_TRITON: + try: + import triton.language.extra.cann.extension as _extension_module # type: ignore + except ImportError: + _extension_module = None + + +def _resolve_triton_ascend_op(op_name: str): + if not HAS_TRITON: + raise RuntimeError( + f"Triton op '{op_name}' cannot be resolved because HAS_TRITON is False" + ) + + if _extension_module is not None: + extension_op = getattr(_extension_module, op_name, None) + if extension_op is not None: + return extension_op + + tl_op = getattr(tl, op_name, None) + if tl_op is not None: + return tl_op + + raise RuntimeError( + f"Failed to resolve Triton op '{op_name}': " + "neither triton.language.extra.cann.extension nor triton.language provides it." + ) + + +if HAS_TRITON: + insert_slice = _resolve_triton_ascend_op("insert_slice") + extract_slice = _resolve_triton_ascend_op("extract_slice") + get_element = _resolve_triton_ascend_op("get_element") +else: + insert_slice = None + extract_slice = None + get_element = None def init_device_properties_triton(): diff --git a/vllm_fl/dispatch/backends/vendor/ascend/patch.py b/vllm_fl/dispatch/backends/vendor/ascend/patch.py index ff2a101e9..9216fff36 100644 --- a/vllm_fl/dispatch/backends/vendor/ascend/patch.py +++ b/vllm_fl/dispatch/backends/vendor/ascend/patch.py @@ -18,6 +18,12 @@ def apply_ascend_patches(): patch_fla_ops() patch_op_cls() patch_fused_moe() + patch_qwen3_5_attention() + patch_qwen3_6_gdn() + patch_qwen3_mtp() + patch_graph() + patch_npugraph_ex() + patch_dynamo_safe_ops() def patch_mamba_config(): """Patch HybridAttentionMambaModelConfig for Ascend.""" @@ -44,9 +50,26 @@ def patch_causal_conv1d(): logger.warning("Failed to patch causal_conv1d ops: %s", e) def patch_fused_moe(): - """Patch fused MoE ops with Ascend implementations.""" + """Patch fused MoE ops with Ascend implementations. + + Always replaces ``fused_experts_impl`` (it dispatches between the AscendC + custom-op path and the legacy torch_npu path at call time based on the + weight layout). When the AscendC MoE custom ops are available + (``ascendc_moe_available``), additionally: + + * replace ``fused_topk`` with the fused AscendC ``moe_gating_top_k`` + kernel (softmax + top-k + renorm in one launch); + * wrap ``UnquantizedFusedMoEMethod.process_weights_after_loading`` so + expert weights are stored pre-transposed, removing the per-forward + ``transpose(1, 2).contiguous()`` copies of the legacy path. + """ # TODO ops' triton implementation is not ready yet - from .impl.fused_moe import fused_experts_impl + from .impl.fused_moe import ( + ascendc_moe_available, + convert_moe_weights_pretransposed, + fused_experts_impl, + fused_topk_ascend, + ) try: import vllm_fl.ops.fused_moe.fused_moe as fused_moe_lib @@ -55,6 +78,88 @@ def patch_fused_moe(): logger.info("Patched fused_moe for Ascend") except Exception as e: logger.warning("Failed to patch fused_moe ops: %s", e) + fused_moe_lib = None + + if fused_moe_lib is None or not ascendc_moe_available(): + return + + try: + fused_moe_lib.fused_topk = fused_topk_ascend + logger.info("Patched fused_topk with AscendC moe_gating_top_k") + except Exception as e: + logger.warning("Failed to patch fused_topk: %s", e) + + try: + from vllm.model_executor.layers.fused_moe.layer import ( + UnquantizedFusedMoEMethod, + ) + + orig_process_weights = UnquantizedFusedMoEMethod.process_weights_after_loading + + def process_weights_after_loading_pretransposed(self, layer): + orig_process_weights(self, layer) + convert_moe_weights_pretransposed(layer) + + UnquantizedFusedMoEMethod.process_weights_after_loading = ( + process_weights_after_loading_pretransposed + ) + logger.info("Patched MoE weight loading with pre-transposed layout for AscendC ops") + except Exception as e: + logger.warning("Failed to patch MoE process_weights_after_loading: %s", e) + +def patch_qwen3_5_attention(): + """Patch Qwen3.5/Qwen3.6 attention to use the fused Ascend kernel.""" + try: + from .patches.patch_qwen3_5 import patch_qwen3_5_attention as _do_patch + + _do_patch() + except Exception as e: + logger.warning("Failed to patch Qwen3NextAttention for Ascend: %s", e) + + +def patch_qwen3_6_gdn(): + """Patch Qwen3.5/Qwen3.6 GatedDeltaNet and GemmaRMSNorm with AscendC ops. + + Falls back to the existing Triton path when the CANN custom-op package + is not available at runtime. + """ + try: + from .patches.patch_qwen3_6_gdn import patch_qwen3_6_gdn as _do_patch + + _do_patch() + except Exception as e: + logger.warning("Failed to patch Qwen3.6 GDN AscendC ops: %s", e) + + +def patch_qwen3_mtp(): + """Patch Qwen3.5/Qwen3.6 Multi-Token Prediction for Ascend.""" + try: + from .patches.patch_qwen3_mtp import patch_qwen3_mtp as _do_patch + + _do_patch() + except Exception as e: + logger.warning("Failed to patch Qwen3 MTP for Ascend: %s", e) + + +def patch_graph(): + """Patch GraphWrapper with Ascend ACL graph behavior.""" + try: + from .patches.patch_graph import patch_graph as _do_patch + + _do_patch() + except Exception as e: + logger.warning("Failed to patch GraphWrapper for Ascend: %s", e) + + +def patch_npugraph_ex(): + """Patch npugraph_ex/torchair ValuePack handling.""" + try: + from .patches.patch_npugraph_ex import patch_npugraph_ex as _do_patch + + _do_patch() + except Exception as e: + logger.warning("Failed to patch npugraph_ex for Ascend: %s", e) + def patch_fla_ops(): """Patch FLA ops and fused_gdn_gating with Ascend implementations.""" @@ -108,6 +213,101 @@ def patch_op_cls(): except Exception as e: logger.warning("Failed to patch MMEncoderAttention: %s", e) +def patch_dynamo_safe_ops(): + """Bypass the FL dispatch manager for OOT ops when running on NPU. + + vLLM v1 graph mode compiles model forward with torch.compile(fullgraph=True). + The FL dispatch manager's call_op uses Python RLock/context managers that + Dynamo cannot trace, so on NPU we replace the OOT forward methods with + direct calls to the Ascend (or reference PyTorch) implementations. This + keeps the dispatch manager available for eager-mode / non-graph use. + """ + try: + from vllm.platforms import current_platform + + if getattr(current_platform, "device_type", None) != "npu": + return + + from vllm_fl.ops.activation import GeluAndMulFL, SiluAndMulFL + from vllm_fl.ops.layernorm import RMSNormFL + from vllm_fl.ops.rotary_embedding import RotaryEmbeddingFL + + from .impl.activation import silu_and_mul_ascend + from .impl.normalization import rms_norm_ascend + from .impl.rotary import rotary_embedding_ascend + from vllm_fl.dispatch.backends.reference.impl.activation import ( + gelu_and_mul_torch, + silu_and_mul_torch, + ) + + SiluAndMulFL.forward_oot = silu_and_mul_ascend + GeluAndMulFL.forward_oot = gelu_and_mul_torch + RMSNormFL.forward_oot = rms_norm_ascend + RotaryEmbeddingFL.forward_oot = _make_rotary_forward_oot( + rotary_embedding_ascend + ) + + logger.info("Patched FL OOT ops for Dynamo-safe NPU execution") + except Exception as e: + logger.warning("Failed to patch FL OOT ops for NPU: %s", e) + + +def _make_rotary_forward_oot(rotary_impl): + """Build a RotaryEmbeddingFL.forward_oot that calls ``rotary_impl`` directly. + + The original forward_oot reshapes query/key and extracts cos/sin before + calling the operator; we keep that logic and only bypass call_op. + """ + + def forward_oot( + self, + positions: "torch.Tensor", + query: "torch.Tensor", + key: "torch.Tensor | None" = None, + ) -> tuple["torch.Tensor", "torch.Tensor | None"]: + # Use a local tensor instead of assigning back to the buffer; assigning + # to self.cos_sin_cache inside forward is forbidden when cudagraph is + # used inside torch.compile. + cos_sin_cache = self.cos_sin_cache.to(positions.device) + positions = positions.flatten() + num_tokens = positions.shape[0] + + query_shape = query.shape + key_shape = key.shape + query = query.view(num_tokens, -1, self.head_size) + key = key.view(num_tokens, -1, self.head_size) + + query_rot = query[..., : self.rotary_dim] + key_rot = key[..., : self.rotary_dim] + if self.rotary_dim < self.head_size: + query_pass = query[..., self.rotary_dim :] + key_pass = key[..., self.rotary_dim :] + + cos, sin = cos_sin_cache.chunk(2, dim=-1) + + q_embed, k_embed = rotary_impl( + self, + query_rot, + key_rot, + cos, + sin, + positions, + not self.is_neox_style, + True, + ) + + if self.rotary_dim < self.head_size: + query = torch.cat((q_embed, query_pass), dim=-1).reshape(query_shape) + key = torch.cat((k_embed, key_pass), dim=-1).reshape(key_shape) + else: + query = q_embed.reshape(query_shape) + key = k_embed.reshape(key_shape) + + return query, key + + return forward_oot + + def refresh_block_size(vllm_config, block_size = 128): """ Refresh the block size in cache config. diff --git a/vllm_fl/dispatch/backends/vendor/ascend/patches/README.md b/vllm_fl/dispatch/backends/vendor/ascend/patches/README.md new file mode 100644 index 000000000..f900d866c --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/patches/README.md @@ -0,0 +1,22 @@ +# Ascend-specific patches + +This directory contains runtime monkey-patches that adapt upstream vLLM model +code for Huawei Ascend NPUs. These patches are applied automatically when the +`fl` plugin is loaded on an Ascend platform (`VLLM_FL_PLATFORM=ascend`). + +## Patch list + +| File | Target | Purpose | +|------|--------|---------| +| `patch_mamba_config.py` | `HybridAttentionMambaModelConfig.verify_and_update_config` | Aligns attention/mamba block sizes with Ascend requirements. | +| `patch_multimodal_merge.py` | `vllm.model_executor.models.utils.merge_multimodal_embeddings` | In-place merge of multimodal embeddings on NPU. | +| `patch_qwen3_5.py` | `vllm.model_executor.models.qwen3_next.Qwen3NextAttention.forward` | Fuses the Q/K/V split, RMSNorm and M-RoPE of Qwen3.5/Qwen3.6 full-attention layers into a single Triton kernel (`torch.ops.vllm.triton_split_qkv_rmsnorm_mrope`). | +| `patch_qwen3_6_gdn.py` | `vllm.model_executor.models.qwen3_next.Qwen3NextGatedDeltaNet._forward_core` / `get_state_shape`, `vllm.model_executor.layers.layernorm.GemmaRMSNorm.forward_oot`, `vllm_fl.worker.model_runner.ModelRunnerFL._reshape_kv_cache_tensors`, `vllm.v1.attention.backends.gdn_attn.GDNAttentionMetadataBuilder.build` | Routes the Qwen3.5/Qwen3.6 GDN core computation to the AscendC fused kernels (`torch.ops._C_ascend.npu_causal_conv1d_custom`, `npu_fused_gdn_gating`, `npu_recurrent_gated_delta_rule`) and GemmaRMSNorm to `npu_gemma_rms_norm` / `npu_add_rms_norm_bias`, mirroring vllm-ascend's `AscendGatedDeltaNetAttention`. Stores the ssm state in the kernel-native `(Hv, Dv, Dk)` layout via `get_state_shape`. Fresh (all-zero `initial_state`) prefill batches run the fused PTO/Bisheng megakernel (`vllm_fl/ops/pto_chunk_gdn`, vllm-ascend PR #8872, 6 GDN stages in one launch); the PTO-vs-Triton decision and chunk counting are made from CPU-side metadata flags attached by the builder wrap, so no per-layer device syncs are introduced. Other prefill batches keep the Triton `chunk_gated_delta_rule` with transposes at the boundary. Also regroups the mamba KV-cache views into dense per-state tensors (the AscendC kernels address state caches assuming dense layout; the Triton path is layout-agnostic and unaffected). The CANN custom-op environment is bootstrapped automatically (`ASCEND_CUSTOM_OPP_PATH` pointed at the packaged `_cann_ops_custom` vendor dir); the patch falls back to the Triton path only when the `_C_ascend` bindings/op package are unavailable, `VLLM_FL_DISABLE_ASCENDC_GDN=1` is set, or `VLLM_FL_DISABLE_PTO_GDN=1` disables just the PTO megakernel. | +| `patch_qwen3_mtp.py` | `vllm.v1.worker.utils.bind_kv_cache`, `Qwen3NextMultiTokenPredictor.forward`, `Qwen3_5MultiTokenPredictor.forward`, `SpeculativeConfig.hf_config_override`, `MRotaryEmbedding.forward_native`, `Qwen3NextMTP.load_weights` | Enables Multi-Token Prediction (MTP) on Ascend by allowing multiple attention layers to share a layer index in KV-cache binding and by forcing the MTP drafter to use local token embeddings on the last PP rank. Maps `qwen3_5`/`qwen3_5_moe` draft configs to a real `Qwen3NextConfig` (supplying `decoder_sparse_step`, `mlp_only_layers`, etc.) so that MoE checkpoints such as Qwen3.6-35B-A3B can load the upstream `Qwen3NextMTP` drafter. Splits packed `experts.gate_up_proj`/`experts.down_proj` MoE weights into 2-D per-expert `gate_proj`/`up_proj`/`down_proj` slices expected by upstream `FusedMoE.make_expert_params_mapping`. Also makes M-RoPE's fallback native path graph-safe for dynamic token counts. | + +## Adding a new patch + +1. Create `patches/patch_.py`. +2. Expose an idempotent `patch_()` function. +3. Register the function in `../patch.py::apply_ascend_patches()`. +4. Update this README with the new entry. diff --git a/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_graph.py b/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_graph.py new file mode 100644 index 000000000..f6fbbc61f --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_graph.py @@ -0,0 +1,440 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Adapted from https://github.com/vllm-project/vllm-ascend/blob/main/vllm_ascend/compilation/acl_graph.py +# Below is the original copyright: +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Ascend-specific ACL graph extensions for vllm-plugin-FL. + +This module is intentionally separated from the generic graph wrapper so that +Ascend behavior (stream sync, graph-param workspaces, capture-error diagnosis, +etc.) is injected at runtime rather than hard-coded into the multi-hardware +framework. +""" + +from __future__ import annotations + +import dataclasses +import logging +import os +import weakref +from dataclasses import dataclass +from typing import Any, ClassVar, Optional +from unittest.mock import patch + +import torch + +from vllm.compilation.counter import compilation_counter +from vllm.config import CUDAGraphMode, VllmConfig +from vllm.forward_context import BatchDescriptor, get_forward_context +from vllm.platforms import current_platform + +from vllm_fl.compilation.graph import register_graph_wrapper_backend + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- # +# Stream-resource capture error diagnostics (CANN error 207008) +# --------------------------------------------------------------------------- # +_STREAM_RESOURCE_ERROR_CODE = "207008" +_STREAM_RESOURCE_ERROR_MARKERS = ( + "insufficient_stream_resources", + "stream resources are insufficient", +) +_STREAM_RESOURCE_GUIDANCE = ( + "ACL graph capture failed with a known stream-resource exhaustion " + "signature. Consider upgrading to a newer HDK/CANN stack, reducing " + "cudagraph_capture_sizes, lowering max_cudagraph_capture_size, preferring " + "FULL or FULL_DECODE_ONLY for mostly uniform decode workloads, or " + "temporarily disabling graph mode to confirm the failure is capture-related." +) + + +def _is_stream_resource_capture_error(exc: RuntimeError) -> bool: + message = str(exc) + lowered_message = message.lower() + has_error_code = _STREAM_RESOURCE_ERROR_CODE in message + has_stream_resource_marker = any( + marker in lowered_message for marker in _STREAM_RESOURCE_ERROR_MARKERS) + return has_stream_resource_marker or (has_error_code + and "stream resource" in lowered_message) + + +def _raise_stream_resource_capture_error(exc: RuntimeError) -> None: + raise RuntimeError( + f"{_STREAM_RESOURCE_GUIDANCE}\nOriginal error:\n{exc}") from exc + + +# --------------------------------------------------------------------------- # +# Graph parameter bookkeeping for attention/MLA workspace reuse across captures +# --------------------------------------------------------------------------- # +@dataclass +class GraphParams: + events: dict[int, list[torch.npu.ExternalEvent]] + workspaces: dict[int, torch.Tensor] + handles: dict[int, list[Any]] + attn_params: dict[int, list[tuple]] + conv1d_params: dict[int, list[tuple]] # for causal conv1d params + conv1d_handles: dict[int, list[Any]] # for causal conv1d params handles + conv1d_events: dict[int, list[torch.npu.ExternalEvent]] # for causal conv1d params events + + +_graph_params: Optional[GraphParams] = None +_draft_graph_params: Optional[GraphParams] = None +_draft_graph_prefill_params: Optional[GraphParams] = None + +def reset_graph_params() -> None: + global _graph_params, _draft_graph_params, _draft_graph_prefill_params + _graph_params = None + _draft_graph_params = None + _draft_graph_prefill_params = None + + +def _make_empty_graph_params(capture_sizes: list[int]) -> GraphParams: + return GraphParams( + {size: [] for size in capture_sizes}, + {size: None for size in capture_sizes}, + {size: [] for size in capture_sizes}, + {size: [] for size in capture_sizes}, + {size: [] for size in capture_sizes}, + {size: [] for size in capture_sizes}, + {size: [] for size in capture_sizes}, + ) + + +def set_graph_params(aclgraph_capture_sizes: list[int]) -> None: + global _graph_params + if _graph_params is not None: + raise ValueError("Graph parameters have already been set!") + _graph_params = _make_empty_graph_params(aclgraph_capture_sizes) + + +def ensure_graph_params(aclgraph_capture_sizes: list[int]) -> GraphParams: + """Initialize graph bookkeeping and add any newly resolved capture sizes. + + The final ACL graph sizes are not known until attention backends have + resolved the cudagraph mode. This helper is deliberately idempotent so a + worker can run that resolution more than once without losing graph handles + that have already been registered. + """ + global _graph_params + if _graph_params is None: + _graph_params = _make_empty_graph_params(aclgraph_capture_sizes) + return _graph_params + + for size in aclgraph_capture_sizes: + _graph_params.events.setdefault(size, []) + _graph_params.workspaces.setdefault(size, None) + _graph_params.handles.setdefault(size, []) + _graph_params.attn_params.setdefault(size, []) + _graph_params.conv1d_params.setdefault(size, []) + _graph_params.conv1d_handles.setdefault(size, []) + _graph_params.conv1d_events.setdefault(size, []) + return _graph_params + + +def update_graph_params_workspaces(num_tokens: int, workspace: torch.Tensor) -> None: + global _graph_params + if _graph_params is not None: + _graph_params.workspaces[num_tokens] = workspace + + +def get_graph_params() -> Optional[GraphParams]: + return _graph_params + + +def set_draft_graph_params(aclgraph_capture_sizes: list[int]) -> None: + global _draft_graph_params + if _draft_graph_params is not None: + raise ValueError("DraftGraph parameters have already been set!") + _draft_graph_params = _make_empty_graph_params(aclgraph_capture_sizes) + + +def update_draft_graph_params_workspaces(num_tokens: int, workspace: Any) -> None: + global _draft_graph_params + if _draft_graph_params is not None: + _draft_graph_params.workspaces[num_tokens] = workspace + + +def get_draft_graph_params() -> Optional[GraphParams]: + return _draft_graph_params + + +def set_draft_graph_prefill_params(aclgraph_capture_sizes: list[int]) -> None: + global _draft_graph_prefill_params + if _draft_graph_prefill_params is not None: + raise ValueError("DraftGraph prefill parameters have already been set!") + _draft_graph_prefill_params = _make_empty_graph_params(aclgraph_capture_sizes) + + +def update_draft_graph_prefill_params_workspaces(num_tokens: int, + workspace: Any) -> None: + global _draft_graph_prefill_params + if _draft_graph_prefill_params is not None: + _draft_graph_prefill_params.workspaces[num_tokens] = workspace + + +def get_draft_graph_prefill_params() -> Optional[GraphParams]: + return _draft_graph_prefill_params + + +def weak_ref_tensors(tensor: Any) -> Any: + """Convert tensors to weak references to save memory during graph replay.""" + from vllm_fl.compilation.graph import weak_ref_tensors as _generic_weak_ref + return _generic_weak_ref(tensor) + + +def weak_ref_workspaces(params: Optional[GraphParams]) -> None: + if params is None: + return + for num_tokens in params.workspaces: + if params.workspaces[num_tokens] is None: + continue + params.workspaces[num_tokens] = weak_ref_tensors( + params.workspaces[num_tokens]) + + +def _has_runtime_graph_updates( + params: Optional[GraphParams], num_tokens: Optional[int] +) -> bool: + """Whether replay can race with graph-task parameter updates. + + The host-blocking replay barrier inherited from vLLM-Ascend is only + required when captured attention/conv1d graph tasks are updated on a + separate stream. Empty maps let graph paths without such tasks skip it. + """ + if params is None: + return False + + runtime_update_maps = ( + params.events, + params.handles, + params.attn_params, + params.conv1d_events, + params.conv1d_handles, + params.conv1d_params, + ) + if num_tokens is None: + return any( + any(bucket for bucket in values.values()) + for values in runtime_update_maps + ) + return any(values.get(num_tokens) for values in runtime_update_maps) + + +def update_full_graph_params( + attn_backends, + update_stream, + forward_context, + num_tokens: int, + vllm_config: VllmConfig, + speculative_config=None, + num_dcp_pcp_tokens: Optional[int] = None, + draft_attn_metadatas=None, +) -> None: + """Dispatch graph-param updates to the attention backend and GDN conv1d.""" + if not isinstance(attn_backends, (list, tuple, set)): + attn_backends = (attn_backends,) + + updated_impls = set() + for attn_backend in attn_backends: + try: + impl_cls = attn_backend.get_impl_cls() + except NotImplementedError: + # Hybrid state-space/GDN backends may implement attention through + # custom ops without exposing the standard AttentionImpl hook. + # Their graph-task updates are handled by the conv1d path below. + continue + if impl_cls in updated_impls or not hasattr(impl_cls, "update_graph_params"): + continue + updated_impls.add(impl_cls) + impl_cls.update_graph_params( + update_stream, + forward_context, + num_tokens, + vllm_config, + speculative_config, + num_dcp_pcp_tokens, + draft_attn_metadatas, + ) + + has_attention_tasks = any( + params is not None and params.handles.get(num_tokens) + for params in ( + get_graph_params(), + get_draft_graph_params(), + get_draft_graph_prefill_params(), + ) + ) + if has_attention_tasks and not updated_impls: + raise RuntimeError( + "FULL ACL graph has captured attention task handles, but no " + "attention graph-parameter updater was resolved. Replaying would " + "deadlock while waiting for an unrecorded ExternalEvent." + ) + + # Optional GDN conv1d update (only available when vllm-ascend gdn is present). + try: + from vllm_ascend.ops.gdn import update_conv1d_graph_params + update_conv1d_graph_params( + update_stream, + forward_context, + num_tokens, + vllm_config, + getattr(forward_context, "is_draft_model", False), + draft_attn_metadatas, + ) + except Exception: + pass + + +# --------------------------------------------------------------------------- # +# Ascend backend mixin for GraphWrapper +# --------------------------------------------------------------------------- # +@dataclasses.dataclass +class _ACLGraphEntry: + """Internal entry used by the mixin; mirrors generic GraphEntry fields.""" + batch_descriptor: BatchDescriptor + aclgraph: Any | None = None + output: Any | None = None + input_addresses: Optional[list[int]] = None + + +class ACLGraphBackendMixin: + """ + Backend-specific mixin that supplies Ascend ACL graph behavior to the + generic GraphWrapper. + + The mixin is instantiated once per GraphWrapper and receives hook calls + during capture and replay. It mirrors the workflow of + `vllm_ascend.compilation.acl_graph.ACLGraphWrapper` but stays out of the + generic code path. + """ + + _all_instances: ClassVar[weakref.WeakSet["ACLGraphBackendMixin"]] = weakref.WeakSet() + + @classmethod + def clear_all_graphs(cls) -> None: + for instance in list(cls._all_instances): + instance.wrapper.concrete_graph_entries.clear() + + def __init__(self, wrapper): + self.wrapper = wrapper + self.vllm_config = wrapper.vllm_config + self.runtime_mode = wrapper.runtime_mode + self.aclgraph_options = wrapper.graph_options + self.use_eagle = getattr(wrapper, "use_eagle", False) + self.enable_enpu = getattr(wrapper, "enable_enpu", False) + # Emergency rollback/A-B switch for stacks that update graph tasks + # outside the GraphParams bookkeeping used below. + self.force_replay_sync = os.environ.get( + "VLLM_FL_FORCE_ACLGRAPH_REPLAY_SYNC", "0" + ).lower() in ("1", "true", "yes", "on") + self.is_debugging_mode = wrapper.is_debugging_mode + self._runnable_str = str( + wrapper.runnable) if self.is_debugging_mode else None + ACLGraphBackendMixin._all_instances.add(self) + + def _is_stream_resource_capture_error(self, exc: RuntimeError) -> bool: + return _is_stream_resource_capture_error(exc) + + def _sync_offloader_before_capture(self) -> None: + try: + from vllm.model_executor.offloader.base import get_offloader + get_offloader().sync_prev_onload() + except Exception: + pass + + def _join_offloader_after_forward(self) -> None: + try: + from vllm.model_executor.offloader.base import get_offloader + get_offloader().join_after_forward() + except Exception: + pass + + def before_capture(self, entry, args, kwargs) -> None: + self._sync_offloader_before_capture() + forward_context = get_forward_context() + self._previous_capturing = getattr(forward_context, "capturing", False) + forward_context.capturing = True + + def wrap_capture_context(self, entry, stack) -> None: + # For NPU, torch.npu.empty_cache is the function that needs to be + # disabled when gc_disable is enabled. The generic wrapper already + # patches PlatformFL.empty_cache; patch torch.npu.empty_cache as well. + if self.aclgraph_options.gc_disable: + stack.enter_context(patch("torch.npu.empty_cache", lambda: None)) + + def after_capture(self, entry, output, args, kwargs) -> Any: + self._join_offloader_after_forward() + + forward_context = get_forward_context() + forward_context.capturing = getattr( + self, "_previous_capturing", False + ) + + # Convert attention workspace tensors to weak refs to save memory. + weak_ref_workspaces(get_graph_params()) + weak_ref_workspaces(get_draft_graph_params()) + weak_ref_workspaces(get_draft_graph_prefill_params()) + + # The generic wrapper will weak-ref the output again; return the + # original output so PyTorch can manage memory correctly during capture. + return output + + def capture_error_handler(self, exc: BaseException) -> None: + try: + forward_context = get_forward_context() + forward_context.capturing = getattr( + self, "_previous_capturing", False + ) + except Exception: + pass + if isinstance(exc, RuntimeError) and self._is_stream_resource_capture_error(exc): + _raise_stream_resource_capture_error(exc) + + def before_replay(self, entry, args, kwargs) -> None: + # In async scheduling or multi-threaded scenarios, ensure host-side + # attention-param updates stay ordered with graph execution. + # When enable_enpu is on, model_runner orders update vs replay; skip. + # When FULL + EAGLE draft (merge path), replay does not need barrier. + is_draft_eagle = False + try: + from vllm_ascend.ascend_forward_context import _EXTRA_CTX + is_draft_eagle = _EXTRA_CTX.is_draft_model and self.use_eagle + except Exception: + pass + + need_sync = self.runtime_mode == CUDAGraphMode.FULL and not is_draft_eagle + num_tokens = getattr(entry.batch_descriptor, "num_tokens", None) + has_runtime_updates = any( + _has_runtime_graph_updates(params, num_tokens) + for params in ( + get_graph_params(), + get_draft_graph_params(), + get_draft_graph_prefill_params(), + ) + ) + if ( + not self.enable_enpu + and need_sync + and (self.force_replay_sync or has_runtime_updates) + ): + torch.npu.current_stream().synchronize() + + def weak_ref_tensors(self, tensor: Any) -> Any: + # Ascend does not yet have a dedicated weak-ref csrc op; fall back to + # the generic implementation which currently returns the tensor as-is. + return tensor + + +def patch_graph() -> None: + """Register the Ascend ACL graph backend mixin.""" + if current_platform.device_type != "npu": + logger.info( + "Skipping ACL graph patch: current platform is not NPU (%s)", + current_platform.device_type) + return + register_graph_wrapper_backend("npu", ACLGraphBackendMixin) + logger.info("Registered Ascend ACL graph backend mixin for GraphWrapper") diff --git a/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_npugraph_ex.py b/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_npugraph_ex.py new file mode 100644 index 000000000..fdbd3611e --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_npugraph_ex.py @@ -0,0 +1,149 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Adapted from https://github.com/vllm-project/vllm-ascend/blob/main/vllm_ascend/patch/worker/patch_npugraph_ex_triton.py +# Below is the original copyright: +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +npugraph_ex ValuePack patch for Triton scenarios. + +This module patches npugraph_ex/torchair internals so that ValuePack inputs are +unpacked correctly when Triton kernels are present in the captured graph. It is +applied at worker startup only when npugraph_ex is available. +""" + +import importlib +import logging +import sys + +import torch +from torch._subclasses.fake_tensor import FakeTensor + +logger = logging.getLogger(__name__) + +_PATCHED = False + + +def patch_npugraph_ex() -> None: + """Apply npugraph_ex/torchair ValuePack patches.""" + global _PATCHED + if _PATCHED: + return + _PATCHED = True + + try: + import npugraph_ex as nge + from npugraph_ex.core._concrete_graph import _is_symlist + from npugraph_ex.npu_fx_compiler import _unpack_meta_list + _USE_NPUGRAPH_EX = True + except ImportError: + try: + import torchair as nge + from torchair.core._concrete_graph import _is_symlist + from torchair.npu_fx_compiler import _unpack_meta_list + _USE_NPUGRAPH_EX = False + except ImportError: + logger.info("npugraph_ex/torchair not available; skipping patch") + return + + class ValuePack: + def __init__(self, meta, npu_meta=None) -> None: + self._meta = meta + self._npu_meta = meta if npu_meta is None else npu_meta + + @property + def meta(self): + return self._meta + + @property + def npu(self): + return self._npu_meta + + def __getitem__(self, key): + if isinstance(self._meta, dict): + return self._meta.get(key) + raise ValueError( + f"Unsupported meta type for ValuePack __getitem__, " + f"key:{key}, type: {type(self._meta)}") + + def __repr__(self) -> str: + if isinstance(self._meta, FakeTensor): + meta_str = f"FakeTensor(dtype={self._meta.dtype}, size={list(self._meta.size())}" + elif isinstance(self._meta, torch.Tensor): + meta_str = f"torch.Tensor(dtype={self._meta.dtype}, size={list(self._meta.size())}" + elif isinstance(self._meta, torch.SymInt): + meta_str = f"torch.SymInt({self._meta})" + else: + try: + meta_str = f"{type(self._meta)}({self._meta})" + except Exception: + meta_str = f"{type(self._meta)}" + return f"Pack(meta:{meta_str} npu:{self._npu_meta})" + + def _unpack_meta(args, kwargs): + unpacked_args = [] + unpacked_kwargs = {} + + def _get_meta_part(arg): + if isinstance(arg, (list, tuple)) and any( + isinstance(v, ValuePack) for v in arg): + return _unpack_meta_list(arg) + elif isinstance(arg, dict): + return { + k: v.meta if isinstance(v, ValuePack) else v + for k, v in arg.items() + } + elif isinstance(arg, ValuePack): + return arg.meta + else: + return arg + + for arg in args: + unpacked_args.append(_get_meta_part(arg)) + + for key, value in kwargs.items(): + unpacked_kwargs[key] = _get_meta_part(value) + + return list(unpacked_args), unpacked_kwargs + + def _unpack_npu(self, args, kwargs): + unpacked = [] + unpacked_kwargs = {} + + def _get_npu_part(arg): + if isinstance(arg, (list, tuple)) and len(arg): + if _is_symlist(arg): + arg = self._graph.parse_symlist(arg) + else: + arg = [(v.npu if isinstance(v, ValuePack) else v) + for v in arg] + return arg + elif isinstance(arg, dict): + return { + k: v.npu if isinstance(v, ValuePack) else v + for k, v in arg.items() + } + elif isinstance(arg, ValuePack): + return arg.npu + else: + return arg + + for arg in args: + unpacked.append(_get_npu_part(arg)) + + for key, value in kwargs.items(): + unpacked_kwargs[key] = _get_npu_part(value) + + return unpacked, unpacked_kwargs + + nge.core._concrete_graph.ValuePack = ValuePack + # The ValuePack class is referenced in the npu_fx_compiler module (and + # fx_summary for torchair), and after the patch these modules need to be + # reloaded. + if not _USE_NPUGRAPH_EX: + importlib.reload(sys.modules["torchair.fx_summary"]) + pkg_prefix = "npugraph_ex" if _USE_NPUGRAPH_EX else "torchair" + importlib.reload(sys.modules[f"{pkg_prefix}.npu_fx_compiler"]) + nge.npu_fx_compiler._unpack_meta = _unpack_meta + nge.npu_fx_compiler._NpuGraphConverter._unpack_npu = _unpack_npu + logger.info("Patched npugraph_ex/torchair ValuePack handling") diff --git a/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_qwen3_5.py b/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_qwen3_5.py new file mode 100644 index 000000000..f7e455749 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_qwen3_5.py @@ -0,0 +1,105 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# Copyright (c) 2026 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# mypy: ignore-errors + +"""Ascend-specific patch for Qwen3.5/Qwen3.6 attention. + +Replaces the upstream ``Qwen3NextAttention.forward`` with a fused +``triton_split_qkv_rmsnorm_mrope`` implementation on Ascend NPUs. +""" + +import logging + +import torch + +from vllm.model_executor.models.qwen3_next import Qwen3NextAttention + +# Importing this module registers ``torch.ops.vllm.triton_split_qkv_rmsnorm_mrope`` +# if it has not been registered by another backend (e.g. vllm-ascend). +from ..impl.linearnorm import split_qkv_rmsnorm_mrope # noqa: F401 + +logger = logging.getLogger(__name__) + + +class AscendQwen3NextAttention(Qwen3NextAttention): + def forward( + self, + positions: torch.Tensor, + output: torch.Tensor, + hidden_states: torch.Tensor, + ): + qkv, _ = self.qkv_proj(hidden_states) + + if "qwen3_5" in self.config.model_type: + cos_sin = self.rotary_emb.cos_sin_cache[positions] + if cos_sin.device != qkv.device: + cos_sin = cos_sin.to(qkv.device) + if cos_sin.dtype != qkv.dtype: + cos_sin = cos_sin.to(qkv.dtype) + + q, k, v, gate = torch.ops.vllm.triton_split_qkv_rmsnorm_mrope( + qkv=qkv, + q_weight=1.0 + self.q_norm.weight, + k_weight=1.0 + self.k_norm.weight, + cos_sin=cos_sin, + num_q_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + eps=self.config.rms_norm_eps, + mrope_section=self.rotary_emb.mrope_section, + is_interleaved=self.rotary_emb.mrope_interleaved, + rope_dim=self.rotary_emb.rotary_dim, + has_gate=self.attn_output_gate, + ) + else: + if self.attn_output_gate: + q_gate, k, v = qkv.split( + [self.q_size * 2, self.kv_size, self.kv_size], dim=-1 + ) + orig_shape = q_gate.shape[:-1] + q_gate = q_gate.view(*orig_shape, self.num_heads, -1) + q, gate = torch.chunk(q_gate, 2, dim=-1) + q = q.reshape(*orig_shape, -1) + gate = gate.reshape(*orig_shape, -1) + else: + q, k, v = qkv.split( + [self.q_size, self.kv_size, self.kv_size], dim=-1 + ) + gate = None + + q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view( + -1, self.num_heads * self.head_dim + ) + k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view( + -1, self.num_kv_heads * self.head_dim + ) + + q, k = self.rotary_emb(positions, q, k) + + attn_output = self.attn(q, k, v) + + if self.attn_output_gate: + gate = torch.sigmoid(gate) + attn_output = attn_output * gate + + output[:], _ = self.o_proj(attn_output) + + +def patch_qwen3_5_attention() -> None: + """Apply the Ascend Qwen3.5/Qwen3.6 attention patch.""" + Qwen3NextAttention.forward = AscendQwen3NextAttention.forward + logger.info("Patched Qwen3NextAttention for Ascend (split_qkv_rmsnorm_mrope)") diff --git a/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_qwen3_6_gdn.py b/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_qwen3_6_gdn.py new file mode 100644 index 000000000..4d884f5d9 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_qwen3_6_gdn.py @@ -0,0 +1,779 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# Copyright (c) 2026 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# mypy: ignore-errors + +"""AscendC fused-op patch for Qwen3.5/Qwen3.6 GatedDeltaNet (GDN) layers. + +Ports the vllm-ascend ``AscendGatedDeltaNetAttention`` integration +(``vllm_ascend/ops/gdn.py`` and ``vllm_ascend/ops/layernorm.py``) to the +FL plugin on vLLM 0.13: + +* ``npu_causal_conv1d_custom`` replaces the Triton ``causal_conv1d_fn`` / + ``causal_conv1d_update`` calls inside ``Qwen3NextGatedDeltaNet._forward_core``. + Its ``(width, dim)`` weight layout is materialized once after checkpoint + loading instead of being transposed implicitly on every forward. +* ``npu_fused_gdn_gating`` replaces the Triton ``fused_gdn_gating``. +* ``npu_recurrent_gated_delta_rule`` replaces ``fused_recurrent_gated_delta_rule`` + on the (speculative-)decode paths. The chunked prefill path keeps the + existing Triton ``chunk_gated_delta_rule`` (already patched to the Ascend + implementation by ``patch_fla_ops``). +* ``npu_gemma_rms_norm`` / ``npu_add_rms_norm_bias`` back + ``GemmaRMSNorm.forward_oot``. +* Non-speculative decode batches run the q/k L2 norm + delta-rule state + update as one fused Triton kernel (``fused_recurrent_delta_rule_update``, + adapted from vllm-ascend's fused_sigmoid_gating_delta_rule_update) on top + of the AscendC ``npu_fused_gdn_gating`` op, replacing the separate + 2x l2norm_fwd + npu_recurrent_gated_delta_rule calls. The in-kernel + sigmoid-gating section of the upstream kernel is miscompiled by the + Ascend Triton pipeline in this environment, so the gating stays on the + AscendC op. Set ``VLLM_FL_DISABLE_FUSED_DECODE_GDN=1`` to fall back to + the AscendC recurrent op. +* ``RMSNormGated.forward_oot`` runs the fused Triton + ``layer_norm_fwd_1pass`` kernel (ported from vllm-ascend) instead of the + decomposed eager ``forward_native`` chain. +* Fresh (all-zero ``initial_state``) prefill batches run the fused + PTO/Bisheng megakernel (vllm-ascend PR #8872 port, + ``vllm_fl/ops/pto_chunk_gdn``): all six GDN stages in a single launch. + The PTO-vs-Triton decision and the megakernel chunk counting are made + from CPU-side flags attached to ``GDNAttentionMetadata`` by the builder + wrap — no per-layer device→host syncs (the naive wrapper approach costs + two syncs per GDN layer per prefill step, which regresses batch64 + serving). Other prefill batches keep the Triton chunk kernel. + ``VLLM_FL_DISABLE_PTO_GDN=1`` disables the megakernel. + +Layout notes (must stay consistent with the kernels): + +* conv_state: vLLM 0.13 allocates the GDN conv cache as + ``(state_len, conv_dim)`` per slot, which is exactly what the AscendC + kernel expects, so the cache is passed through *without* the transpose + used by the Triton path. +* ssm_state: the AscendC ``recurrent_gated_delta_rule`` kernel expects the + state in ``(Hv, Dv, Dk)`` layout (see + ``csrc/ascend/attention/recurrent_gated_delta_rule``), while vLLM 0.13 + allocates ``(Hv, Dk, Dv)``. ``get_state_shape`` is therefore patched to + swap the last two dims, and the chunked-prefill path transposes the + initial/final state at the boundary. +* ``actual_seq_lengths`` of ``npu_recurrent_gated_delta_rule`` follows the + cu_seqlens convention ``[0, len_1, ..., len_B]`` (batch = numel - 1). +* mamba KV-cache: upstream stores conv/ssm states interleaved inside one + (padded) page per block; the AscendC kernels address the state caches + assuming dense per-state tensors, so ``_reshape_kv_cache_tensors`` is + wrapped to regroup the state views into dense per-state tensors over the + same raw storage (transparent to the Triton fallback path, which uses + explicit strides). + +The patch bootstraps the CANN custom-op environment automatically +(``ASCEND_CUSTOM_OPP_PATH`` pointing at the packaged +``_cann_ops_custom/vendors/custom_transformer``) and is only skipped when +the ``_C_ascend`` bindings or the op package are unavailable; otherwise the +existing Triton path is kept. Set ``VLLM_FL_DISABLE_ASCENDC_GDN=1`` to +force the Triton path. +""" + +import logging +import math +import os + +import torch +from vllm.attention.backends.abstract import AttentionMetadata +from vllm.attention.backends.utils import PAD_SLOT_ID +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNormGated +from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator +from vllm.model_executor.models.qwen3_next import Qwen3NextGatedDeltaNet +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata +from vllm.v1.kv_cache_interface import MambaSpec + +import vllm.model_executor.models.qwen3_next as _qwen3_next_lib + +from ..impl.fla.fused_recurrent import fused_recurrent_delta_rule_update +from ..impl.fla.l2norm import l2norm_fwd +from ..impl.linearnorm.layernorm_gated import rmsnorm_gated_oot + +logger = logging.getLogger(__name__) + +_CUSTOM_OPP_MARKER = "custom_transformer" +_REQUIRED_OPS = ( + "npu_causal_conv1d_custom", + "npu_fused_gdn_gating", + "npu_recurrent_gated_delta_rule", + "npu_gemma_rms_norm", + "npu_add_rms_norm_bias", +) + + +def _bootstrap_custom_op_env() -> bool: + """Make the packaged CANN custom-op package discoverable at runtime. + + Same idea as vllm-ascend's ``bootstrap_custom_op_env``: prepend the + packaged ``_cann_ops_custom/vendors/custom_transformer`` dir to + ``ASCEND_CUSTOM_OPP_PATH`` so users do not have to source + ``set_env.bash`` before launching the server. The OPP path is scanned + lazily by the AscendCL runtime at the first custom-op call, so setting + it here (before any op invocation) is sufficient; the variable is also + inherited by spawned worker processes. + + Additionally preload ``libcust_opapi.so`` by absolute path: the aclnn + adapter resolves custom symbols via ``dlopen("libcust_opapi.so")`` by + bare name, which only searches the *startup-time* ``LD_LIBRARY_PATH`` + (glibc caches it) — but finds the already-loaded library by SONAME + after a preload. ``RTLD_LOCAL`` is used on purpose: ``RTLD_GLOBAL`` + leads to a double-free at process teardown. + """ + try: + import vllm_fl._C_ascend as _ext # noqa: F401 + except Exception as e: + logger.warning("Failed to import vllm_fl._C_ascend: %s", e) + return False + vendor_dir = os.path.join( + os.path.dirname(_ext.__file__), "_cann_ops_custom", "vendors", _CUSTOM_OPP_MARKER + ) + if not os.path.isdir(vendor_dir): + logger.warning("CANN custom op package not found at %s", vendor_dir) + return False + opp_path = os.environ.get("ASCEND_CUSTOM_OPP_PATH", "") + if vendor_dir not in opp_path: + os.environ["ASCEND_CUSTOM_OPP_PATH"] = ( + vendor_dir + (":" + opp_path if opp_path else "") + ) + ld_path = os.environ.get("LD_LIBRARY_PATH", "") + lib_dir = os.path.join(vendor_dir, "op_api", "lib") + if lib_dir not in ld_path: + os.environ["LD_LIBRARY_PATH"] = lib_dir + (":" + ld_path if ld_path else "") + try: + import ctypes + + ctypes.CDLL( + os.path.join(lib_dir, "libcust_opapi.so"), mode=ctypes.RTLD_LOCAL + ) + except OSError as e: + logger.warning("Failed to preload libcust_opapi.so: %s", e) + return False + return True + + +def _ascendc_ops_available() -> bool: + """Check the CANN custom-op env and the ``_C_ascend`` torch bindings.""" + if os.environ.get("VLLM_FL_DISABLE_ASCENDC_GDN", "0") == "1": + logger.info("VLLM_FL_DISABLE_ASCENDC_GDN=1, keep Triton GDN path") + return False + if _CUSTOM_OPP_MARKER not in os.environ.get("ASCEND_CUSTOM_OPP_PATH", ""): + if not _bootstrap_custom_op_env(): + logger.warning( + "CANN custom op environment is not set and auto-bootstrap " + "failed; keep Triton GDN path" + ) + return False + try: + import vllm_fl._C_ascend # noqa: F401 + except Exception as e: + logger.warning("Failed to import vllm_fl._C_ascend: %s; keep Triton GDN path", e) + return False + missing = [name for name in _REQUIRED_OPS if not hasattr(torch.ops._C_ascend, name)] + if missing: + logger.warning("torch.ops._C_ascend missing ops %s; keep Triton GDN path", missing) + return False + return True + + +def _build_actual_seq_lengths( + query_start_loc: torch.Tensor, + num_sequences: int, +) -> torch.Tensor: + """Build ``[0, len_1, ..., len_B]`` cu-seqlens style actual_seq_lengths.""" + actual_seq_lengths = torch.empty_like(query_start_loc[: num_sequences + 1]) + actual_seq_lengths[:1].copy_(query_start_loc[:1]) + torch.sub( + query_start_loc[1 : num_sequences + 1], + query_start_loc[:num_sequences], + out=actual_seq_lengths[1:], + ) + return actual_seq_lengths + + +def _fused_decode_gdn_enabled() -> bool: + """Whether to use the fused Triton decode kernel (q/k L2 norm + + recurrent delta-rule state update in a single launch, adapted from + vllm-ascend) for non-speculative decode batches.""" + return os.environ.get("VLLM_FL_DISABLE_FUSED_DECODE_GDN", "0") != "1" + + +def _cache_conv1d_weight_transposed(layer: Qwen3NextGatedDeltaNet) -> None: + """Materialize the AscendC conv weight layout once after weight loading.""" + weight = layer.conv1d.weight + conv_weights = weight.detach().view(weight.size(0), weight.size(2)) + cached_weight = conv_weights.transpose(0, 1).contiguous() + if "_ascendc_conv_weights_t" in layer._buffers: + layer._ascendc_conv_weights_t = cached_weight + elif hasattr(layer, "_ascendc_conv_weights_t"): + layer._ascendc_conv_weights_t = cached_weight + else: + # Qwen3_5GatedDeltaNet deliberately skips the upstream GDN __init__, + # so it does not receive the loader-time buffer registration below. + layer.register_buffer( + "_ascendc_conv_weights_t", cached_weight, persistent=False + ) + + +def _patch_gdn_conv_weight_loader() -> None: + """Cache ``(width, dim)`` conv weights instead of transposing per token.""" + orig_init = Qwen3NextGatedDeltaNet.__init__ + if getattr(orig_init, "_vllm_fl_ascendc_conv_cache", False): + return + + def init_with_conv_cache(self, *args, **kwargs): + orig_init(self, *args, **kwargs) + self.register_buffer( + "_ascendc_conv_weights_t", None, persistent=False + ) + orig_weight_loader = self.conv1d.weight.weight_loader + + def weight_loader_with_conv_cache(param, *loader_args, **loader_kwargs): + result = orig_weight_loader(param, *loader_args, **loader_kwargs) + _cache_conv1d_weight_transposed(self) + return result + + self.conv1d.weight.weight_loader = weight_loader_with_conv_cache + + init_with_conv_cache._vllm_fl_ascendc_conv_cache = True + Qwen3NextGatedDeltaNet.__init__ = init_with_conv_cache + + +# --------------------------------------------------------------------------- +# PTO megakernel for the chunked-prefill path (vllm-ascend PR #8872 port) +# --------------------------------------------------------------------------- +_PTO_AVAILABLE: bool | None = None + + +def _pto_available() -> bool: + """Whether the PTO megakernel can be used for fresh prefill batches.""" + global _PTO_AVAILABLE + if _PTO_AVAILABLE is not None: + return _PTO_AVAILABLE + if os.environ.get("VLLM_FL_DISABLE_PTO_GDN", "0") == "1": + logger.info("VLLM_FL_DISABLE_PTO_GDN=1, keep Triton chunk path") + _PTO_AVAILABLE = False + return False + try: + from vllm_fl.ops.pto_chunk_gdn.mega_kernel import run_mega_kernel # noqa: F401 + + _PTO_AVAILABLE = True + except Exception as e: + logger.warning("PTO chunk_gated_delta_rule unavailable: %s", e) + _PTO_AVAILABLE = False + return _PTO_AVAILABLE + + +def _pto_prefill_usable(attn_metadata) -> bool: + """PTO applies only when *every* sequence in the batch starts from zero + state (the megakernel cannot consume a non-zero initial_state). The flag + is computed on CPU by the metadata builder, so this check is sync-free.""" + if not _pto_available(): + return False + return getattr(attn_metadata, "any_initial_state_cpu", True) is False + + +def _patch_gdn_metadata_host_flags() -> None: + """Attach CPU-side prefill flags to ``GDNAttentionMetadata``. + + ``any_initial_state_cpu`` and ``cu_seqlens_host`` let ``_forward_core`` + pick PTO vs Triton and size the megakernel workspaces without any + device→host synchronization (the generic PTO wrapper previously paid one + ``torch.any(initial_state != 0)`` plus one ``.cpu().tolist()`` sync per + GDN layer per prefill step — a measurable regression at batch64). + """ + from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder + + orig_build = GDNAttentionMetadataBuilder.build + + def build_with_host_flags( + self, + common_prefix_len, + common_attn_metadata, + *args, + **kwargs, + ): + attn_metadata = orig_build( + self, common_prefix_len, common_attn_metadata, *args, **kwargs + ) + if attn_metadata.num_prefills > 0: + context_lens_cpu = common_attn_metadata.num_computed_tokens_cpu + attn_metadata.any_initial_state_cpu = bool((context_lens_cpu > 0).any()) + qsl_cpu = common_attn_metadata.query_start_loc_cpu + attn_metadata.cu_seqlens_host = tuple(int(x) for x in qsl_cpu.tolist()) + return attn_metadata + + GDNAttentionMetadataBuilder.build = build_with_host_flags + logger.info("Patched GDNAttentionMetadataBuilder with CPU prefill flags for PTO") + + +def _chunk_gdn_pto(q, k, v, g, beta, cu_seqlens, attn_metadata): + """Run the PTO megakernel for a fresh (all-zero initial state) prefill batch.""" + from vllm_fl.ops.pto_chunk_gdn.mega_kernel import run_mega_kernel + + Hg, D = q.shape[2], q.shape[3] + q16 = l2norm_fwd(q.to(torch.float16)) + k16 = l2norm_fwd(k.to(torch.float16)) + cu32 = cu_seqlens if cu_seqlens.dtype == torch.int32 else cu_seqlens.to(torch.int32) + lens_host = getattr(attn_metadata, "cu_seqlens_host", None) + total_chunks = None + if lens_host is not None: + total_chunks = sum( + (lens_host[i + 1] - lens_host[i] + 127) // 128 + for i in range(len(lens_host) - 1) + ) + o, fs = run_mega_kernel( + q16, + k16, + v.to(torch.float16), + g.float(), + beta.to(torch.float16), + cu32.contiguous(), + stream=torch.npu.current_stream()._as_parameter_, + chunk_size=128, + scale=D**-0.5, + key_heads=Hg, + return_final_state=True, + total_chunks=total_chunks, + ) + return o.to(q.dtype), fs.to(q.dtype) + + +def _patch_mamba_cache_dense_layout() -> None: + """Rebuild mamba KV-cache views as dense per-state tensors. + + Upstream vLLM lays out the conv/ssm states of a mamba block interleaved + inside one (padded) page, so the per-state views have a first-dim stride + larger than the dense block size. The AscendC kernels address the state + cache assuming dense per-state tensors (verified: in-place state updates + land at wrong offsets with the paged views), while the Triton kernels + take explicit strides and work with either layout. Wrap + ``ModelRunnerFL._reshape_kv_cache_tensors`` so that, after the original + reshape, every MambaSpec layer's state views are rebuilt as dense, + grouped per-state views over the same raw storage. This is semantically + transparent for all other consumers (they index the views by block id). + """ + from vllm.utils.torch_utils import get_dtype_size + + from vllm_fl.worker.model_runner import ModelRunnerFL + + orig_reshape = ModelRunnerFL._reshape_kv_cache_tensors + + def _reshape_kv_cache_tensors_dense_mamba( + self, + kv_cache_config, + kv_cache_raw_tensors, + kernel_block_sizes, + ): + kv_caches = orig_reshape( + self, kv_cache_config, kv_cache_raw_tensors, kernel_block_sizes + ) + for group in self._kv_cache_spec_attn_group_iterator(): + kv_cache_spec = group.kv_cache_spec + if not isinstance(kv_cache_spec, MambaSpec): + continue + if group.kv_cache_group_id == len(kernel_block_sizes): + continue + for layer_name in group.layer_names: + if layer_name in self.runner_only_attn_layers: + continue + raw_tensor = kv_cache_raw_tensors[layer_name] + num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes + state_tensors = [] + storage_offset_bytes = 0 + raw_u8 = raw_tensor.view(torch.uint8) + for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes): + dtype_size = get_dtype_size(dtype) + num_bytes = num_blocks * math.prod(shape) * dtype_size + tensor = ( + raw_u8[storage_offset_bytes : storage_offset_bytes + num_bytes] + .view(dtype) + .view(num_blocks, *shape) + ) + state_tensors.append(tensor) + storage_offset_bytes += num_bytes + kv_caches[layer_name] = state_tensors + return kv_caches + + ModelRunnerFL._reshape_kv_cache_tensors = _reshape_kv_cache_tensors_dense_mamba + logger.info("Patched mamba KV-cache views to dense per-state layout for AscendC GDN ops") + + +class AscendCGatedDeltaNet(Qwen3NextGatedDeltaNet): + """GDN layer backed by the AscendC fused kernels (eager mode).""" + + def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: + conv_state_shape, temporal_state_shape = ( + MambaStateShapeCalculator.gated_delta_net_state_shape( + self.tp_size, + self.num_k_heads, + self.num_v_heads, + self.head_k_dim, + self.head_v_dim, + self.conv_kernel_size, + self.num_spec, + ) + ) + # The AscendC recurrent_gated_delta_rule kernel keeps the ssm state + # in (Hv, Dv, Dk) layout; vLLM 0.13 uses (Hv, Dk, Dv). + num_v_heads, head_k_dim, head_v_dim = temporal_state_shape + return conv_state_shape, (num_v_heads, head_v_dim, head_k_dim) + + def _forward_core( + self, + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + core_attn_out: torch.Tensor, + ): + """ + Core attention computation (called by custom op). + """ + forward_context = get_forward_context() + attn_metadata: AttentionMetadata = forward_context.attn_metadata + + if attn_metadata is None: + # V1 profile run + return + + assert isinstance(attn_metadata, dict) + attn_metadata = attn_metadata[self.prefix] + assert isinstance(attn_metadata, GDNAttentionMetadata) + has_initial_state = attn_metadata.has_initial_state + spec_query_start_loc = attn_metadata.spec_query_start_loc + non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc + spec_sequence_masks = attn_metadata.spec_sequence_masks + spec_token_indx = attn_metadata.spec_token_indx + non_spec_token_indx = attn_metadata.non_spec_token_indx + spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor # noqa: E501 + non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501 + self_kv_cache = self.kv_cache[forward_context.virtual_engine] + # conv cache is already (slot, state_len, conv_dim): pass through. + conv_state = self_kv_cache[0] + ssm_state = self_kv_cache[1] + num_actual_tokens = attn_metadata.num_actual_tokens + num_accepted_tokens = attn_metadata.num_accepted_tokens + + mixed_qkv = mixed_qkv[:num_actual_tokens] + b = b[:num_actual_tokens] + a = a[:num_actual_tokens] + + # 1. Convolution sequence transformation + # The AscendC kernel expects (width, dim). The loader normally creates + # this contiguous layout once; the fallback also covers non-standard + # loading paths that bypass the parameter's weight_loader. + conv_weights_t = getattr(self, "_ascendc_conv_weights_t", None) + if conv_weights_t is None: + _cache_conv1d_weight_transposed(self) + conv_weights_t = self._ascendc_conv_weights_t + activation_mode = 1 if self.activation else 0 + + if spec_sequence_masks is not None: + if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0: + mixed_qkv_spec = mixed_qkv + mixed_qkv_non_spec = None + else: + mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx) + mixed_qkv_non_spec = mixed_qkv.index_select(0, non_spec_token_indx) + else: + mixed_qkv_spec = None + mixed_qkv_non_spec = mixed_qkv + + # 1.1: Process the multi-query part + if spec_sequence_masks is not None: + spec_num_rows = spec_query_start_loc.size(0) - 1 + mixed_qkv_spec_out = torch.empty_like(mixed_qkv_spec) + torch.ops._C_ascend.npu_causal_conv1d_custom( + mixed_qkv_spec_out, + mixed_qkv_spec, + conv_weights_t, + conv_state, + self.conv1d.bias, + spec_query_start_loc, + spec_state_indices_tensor[:spec_num_rows], + None, # initial_state_mode + num_accepted_tokens, + activation_mode, + PAD_SLOT_ID, + 1, # run_mode: decode/speculative update + ) + mixed_qkv_spec = mixed_qkv_spec_out + + # 1.2: Process the remaining part + if attn_metadata.num_prefills > 0: + non_spec_num_rows = non_spec_query_start_loc.size(0) - 1 + mixed_qkv_non_spec_out = torch.empty_like(mixed_qkv_non_spec) + torch.ops._C_ascend.npu_causal_conv1d_custom( + mixed_qkv_non_spec_out, + mixed_qkv_non_spec, + conv_weights_t, + conv_state, + self.conv1d.bias, + non_spec_query_start_loc, + non_spec_state_indices_tensor[:non_spec_num_rows], + has_initial_state, # initial_state_mode + None, # num_accepted_tokens + activation_mode, + PAD_SLOT_ID, + 0, # run_mode: varlen prefill + ) + mixed_qkv_non_spec = mixed_qkv_non_spec_out + elif attn_metadata.num_decodes > 0: + mixed_qkv_non_spec_out = torch.empty_like(mixed_qkv_non_spec) + torch.ops._C_ascend.npu_causal_conv1d_custom( + mixed_qkv_non_spec_out, + mixed_qkv_non_spec, + conv_weights_t, + conv_state, + self.conv1d.bias, + non_spec_query_start_loc, + non_spec_state_indices_tensor[: attn_metadata.num_decodes], + None, # initial_state_mode + None, # num_accepted_tokens + activation_mode, + PAD_SLOT_ID, + 1, # run_mode: decode update + ) + mixed_qkv_non_spec = mixed_qkv_non_spec_out + else: + mixed_qkv_non_spec = None + + query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) + query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv( + mixed_qkv_non_spec + ) + + # 2. Recurrent attention + g, beta = torch.ops._C_ascend.npu_fused_gdn_gating( + self.A_log, a, b, self.dt_bias.to(self.A_log.dtype) + ) + + if spec_sequence_masks is not None: + if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0: + g_spec = g + beta_spec = beta + g_non_spec = None + beta_non_spec = None + else: + g_spec = g.index_select(1, spec_token_indx) + beta_spec = beta.index_select(1, spec_token_indx) + g_non_spec = g.index_select(1, non_spec_token_indx) + beta_non_spec = beta.index_select(1, non_spec_token_indx) + else: + g_spec = None + beta_spec = None + g_non_spec = g + beta_non_spec = beta + + # 2.1: Process the multi-query part + if spec_sequence_masks is not None: + actual_seq_lengths = _build_actual_seq_lengths( + spec_query_start_loc, attn_metadata.num_spec_decodes + ) + query_spec = l2norm_fwd(query_spec) + key_spec = l2norm_fwd(key_spec) + # The AscendC kernel does not apply the q/k L2 norm in-kernel, + # and writes the updated state back in place. + core_attn_out_spec = torch.ops._C_ascend.npu_recurrent_gated_delta_rule( + query=query_spec.squeeze(0), + key=key_spec.squeeze(0), + value=value_spec.squeeze(0), + g=g_spec.squeeze(0), + beta=beta_spec.squeeze(0), + state=ssm_state, + scale=key_spec.shape[-1] ** -0.5, + actual_seq_lengths=actual_seq_lengths, + ssm_state_indices=spec_state_indices_tensor.flatten(), + num_accepted_tokens=num_accepted_tokens.to(torch.int32), + ).unsqueeze(0) + else: + core_attn_out_spec, last_recurrent_state = None, None + + # 2.2: Process the remaining part + if attn_metadata.num_prefills > 0: + if _pto_prefill_usable(attn_metadata): + # Fresh prefill batch: the fused PTO megakernel runs all six + # GDN stages in a single launch. The decision is made from + # CPU-side metadata (no device sync). + core_attn_out_non_spec, last_recurrent_state = _chunk_gdn_pto( + query_non_spec, + key_non_spec, + value_non_spec, + g_non_spec, + beta_non_spec, + non_spec_query_start_loc, + attn_metadata, + ) + else: + # Chunked prefill stays on the (Ascend Triton) chunk kernel, which + # uses the FLA (Hv, Dk, Dv) state layout: transpose at the boundary. + initial_state = ( + ssm_state[non_spec_state_indices_tensor].transpose(-1, -2).contiguous() + ) + initial_state[~has_initial_state, ...] = 0 + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = _qwen3_next_lib.chunk_gated_delta_rule( + q=query_non_spec, + k=key_non_spec, + v=value_non_spec, + g=g_non_spec, + beta=beta_non_spec, + initial_state=initial_state, + output_final_state=True, + cu_seqlens=non_spec_query_start_loc, + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + # Init cache + ssm_state[non_spec_state_indices_tensor] = ( + last_recurrent_state.transpose(-1, -2).contiguous().to(ssm_state.dtype) + ) + elif attn_metadata.num_decodes > 0: + if _fused_decode_gdn_enabled(): + # One fused Triton launch: q/k L2 norm + delta-rule state + # update (state kept in the (Hv, Dv, Dk) AscendC layout, + # which the kernel accesses in its native v-major + # orientation). The sigmoid gating above stays on the + # AscendC npu_fused_gdn_gating op. + core_attn_out_non_spec = fused_recurrent_delta_rule_update( + q=query_non_spec.contiguous(), + k=key_non_spec.contiguous(), + v=value_non_spec.contiguous(), + g=g_non_spec.squeeze(0).contiguous(), + beta=beta_non_spec.squeeze(0).contiguous(), + initial_state_source=ssm_state, + initial_state_indices=non_spec_state_indices_tensor[ + : attn_metadata.num_decodes + ], + cu_seqlens=non_spec_query_start_loc, + use_qk_l2norm_in_kernel=True, + ) + else: + actual_seq_lengths = _build_actual_seq_lengths( + non_spec_query_start_loc, attn_metadata.num_decodes + ) + query_non_spec = l2norm_fwd(query_non_spec) + key_non_spec = l2norm_fwd(key_non_spec) + core_attn_out_non_spec = torch.ops._C_ascend.npu_recurrent_gated_delta_rule( + query=query_non_spec.squeeze(0), + key=key_non_spec.squeeze(0), + value=value_non_spec.squeeze(0), + g=g_non_spec.squeeze(0), + beta=beta_non_spec.squeeze(0), + state=ssm_state, + scale=key_non_spec.shape[-1] ** -0.5, + actual_seq_lengths=actual_seq_lengths, + ssm_state_indices=non_spec_state_indices_tensor[ + : attn_metadata.num_decodes + ], + ).unsqueeze(0) + else: + core_attn_out_non_spec, last_recurrent_state = None, None + + # 3. Merge core attention output + if spec_sequence_masks is not None and core_attn_out_non_spec is not None: + merged_out = torch.empty( + (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), + dtype=core_attn_out_non_spec.dtype, + device=core_attn_out_non_spec.device, + ) + merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec) + merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) + core_attn_out[:num_actual_tokens] = merged_out.squeeze(0) + elif spec_sequence_masks is not None: + core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) + else: + core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0) + + +class AscendCGemmaRMSNorm(GemmaRMSNorm): + """GemmaRMSNorm backed by the AscendC ``npu_gemma_rms_norm`` kernel.""" + + def forward_oot( + self, + x: torch.Tensor, + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is not None: + x, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias( + x, residual, 1.0 + self.weight, None, self.variance_epsilon + ) + return x, residual + + # npu_gemma_rms_norm implements the Gemma (1 + weight) convention + # internally, so the raw weight is passed (same as vllm-ascend). + x, _ = torch.ops._C_ascend.npu_gemma_rms_norm( + x, self.weight, self.variance_epsilon + ) + return x + + +class AscendCRMSNormGated(RMSNormGated): + """RMSNormGated backed by the fused Triton ``layer_norm_fwd_1pass`` kernel. + + The upstream OOT fallback (``forward_native``) decomposes the gated RMS + norm into a long chain of eager ops (silu/pow/mean/rsqrt/mul), which is + a major decode-stage cost for GDN layers; the fused kernel runs the + whole norm (+ SiLU gating) in one launch, same as vllm-ascend's + ``AscendRMSNormGated``. + """ + + def forward_oot( + self, + x: torch.Tensor, + z: torch.Tensor | None = None, + ) -> torch.Tensor: + if z is None: + return RMSNormGated.forward_native(self, x, z) + return rmsnorm_gated_oot( + x, + self.weight, + z, + eps=self.eps, + group_size=self.group_size, + norm_before_gate=self.norm_before_gate, + ) + + +def patch_qwen3_6_gdn() -> bool: + """Apply the AscendC GDN patch for Qwen3.5/Qwen3.6. + + Returns True when the AscendC kernels were wired in; otherwise the + upstream/Triton implementations are kept. + """ + if not _ascendc_ops_available(): + return False + + _patch_gdn_conv_weight_loader() + Qwen3NextGatedDeltaNet.get_state_shape = AscendCGatedDeltaNet.get_state_shape + Qwen3NextGatedDeltaNet._forward_core = AscendCGatedDeltaNet._forward_core + _patch_mamba_cache_dense_layout() + GemmaRMSNorm.forward_oot = AscendCGemmaRMSNorm.forward_oot + RMSNormGated.forward_oot = AscendCRMSNormGated.forward_oot + if _pto_available(): + _patch_gdn_metadata_host_flags() + logger.info( + "Patched Qwen3NextGatedDeltaNet and GemmaRMSNorm/RMSNormGated for Ascend " + "(AscendC causal_conv1d / fused_gdn_gating / recurrent_gated_delta_rule " + "/ gemma_rms_norm, fused Triton delta-rule decode update: %s, " + "PTO megakernel for fresh prefill: %s)", + "on" if _fused_decode_gdn_enabled() else "off", + "on" if _pto_available() else "off", + ) + return True diff --git a/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_qwen3_mtp.py b/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_qwen3_mtp.py new file mode 100644 index 000000000..4a05364f1 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/ascend/patches/patch_qwen3_mtp.py @@ -0,0 +1,1120 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Adapted from the vllm-ascend project. +# mypy: ignore-errors + +"""Ascend-specific patches for Qwen3.5/Qwen3.6 Multi-Token Prediction (MTP). + +These patches are needed because: + +1. The upstream ``vllm.v1.worker.utils.bind_kv_cache`` raises + ``NotImplementedError`` when a single layer index maps to more than one + attention module and the platform is not CUDA/XPU/CPU. With an MTP drafter + model, both the target model and the drafter model contain attention layers + with the same layer index, which triggers this error on Ascend NPUs. + +2. The upstream ``Qwen3NextMultiTokenPredictor.forward`` consumes PP + intermediate tensors on non-first PP ranks. On Ascend the MTP drafter is + placed on the last PP rank and should instead always combine local token + embeddings with the target hidden states passed in from the base model. +""" + +import copy +import logging +import re +from collections import defaultdict + +import numpy as np +import torch +import vllm.v1.worker.utils as worker_utils +from vllm.distributed.parallel_state import get_pp_group +from vllm.sequence import IntermediateTensors +from vllm.v1.attention.backends.utils import CommonAttentionMetadata + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Patch 1: bind_kv_cache +# --------------------------------------------------------------------------- + +def _ascend_bind_kv_cache( + kv_caches: dict[str, torch.Tensor], + forward_context: dict, + runner_kv_caches: list[torch.Tensor], + num_attn_module: int = 1, +) -> None: + """Bind the allocated KV cache to both ModelRunner and forward context. + + This is the Ascend-safe variant of ``vllm.v1.worker.utils.bind_kv_cache``. + When multiple attention layers share the same layer index (e.g. target model + + MTP drafter model), the upstream implementation raises + ``NotImplementedError`` on non-CUDA platforms. Here we follow the CUDA path + and keep only the first layer's KV cache in ``runner_kv_caches`` while + still binding every layer's cache to ``forward_context``. + """ + assert len(runner_kv_caches) == 0 + + index2name = defaultdict(list) + for layer_name in kv_caches: + index2name[worker_utils.extract_layer_index(layer_name, num_attn_module)].append( + layer_name + ) + + for layer_index in sorted(index2name.keys()): + layer_names = index2name[layer_index] + # Typical encoder-decoder / MTP drafter case: multiple attention modules + # share the same layer index. Keep the first one for the runner list and + # bind all of them to the forward context below. + layer_name = layer_names[0] + runner_kv_caches.append(kv_caches[layer_name]) + + for layer_name, kv_cache in kv_caches.items(): + forward_context[layer_name].kv_cache = [kv_cache] + + +def patch_bind_kv_cache() -> None: + """Replace ``vllm.v1.worker.utils.bind_kv_cache`` with the Ascend version.""" + worker_utils.bind_kv_cache = _ascend_bind_kv_cache + logger.info("Patched vllm.v1.worker.utils.bind_kv_cache for Ascend MTP") + + +# --------------------------------------------------------------------------- +# Patch 2: Qwen3NextMultiTokenPredictor.forward +# --------------------------------------------------------------------------- + +def _ascend_qwen3_next_mtp_forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, +) -> torch.Tensor: + """MTP drafter forward that always uses local embeddings on the last PP rank. + + The upstream implementation only builds local embeddings on the first PP + rank and consumes ``intermediate_tensors`` elsewhere. For Ascend the MTP + drafter runs on the last PP rank together with the base model, so the last + rank should always combine token embeddings with the target hidden states + instead of receiving them through PP intermediate tensors. + """ + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + assert hidden_states.shape[-1] == inputs_embeds.shape[-1] + inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds) + hidden_states = self.pre_fc_norm_hidden(hidden_states) + hidden_states = torch.cat([inputs_embeds, hidden_states], dim=-1) + hidden_states = self.fc(hidden_states) + residual = None + + current_step_idx = spec_step_idx % self.num_mtp_layers + hidden_states, residual = self.layers[current_step_idx]( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + { + "hidden_states": hidden_states, + "residual": residual, + } + ) + + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +def _disable_torch_compile_for_drafter( + cls, +) -> None: + """Do not compile the MTP drafter model on Ascend. + + The drafter is invoked with slices of persistent CPU-sized buffers + (``input_ids``, ``hidden_states``). Under ``torch.compile`` these buffers + cause shape specialization / constraint violations because their storage + size is known at trace time. Running the drafter eagerly avoids the issue + while the target model still uses the Ascend graph backend. + """ + ignore_key = "_ignore_compile_vllm" + setattr(cls, ignore_key, True) + logger.info("Disabled torch.compile for %s on Ascend", cls.__name__) + + +def patch_qwen3_next_mtp() -> None: + """Patch ``Qwen3NextMultiTokenPredictor.forward`` and weight loading.""" + try: + from vllm.model_executor.model_loader.base_loader import BaseModelLoader + from vllm.model_executor.models.qwen3_next import QwenNextMixtureOfExperts + from vllm.model_executor.models.qwen3_next_mtp import ( + Qwen3NextMTP, + Qwen3NextMultiTokenPredictor, + ) + except ImportError: + logger.debug("Qwen3NextMTP not available, skip MTP patch") + return + + _disable_torch_compile_for_drafter(Qwen3NextMultiTokenPredictor) + _disable_torch_compile_for_drafter(Qwen3NextMTP) + + _orig_load_model = BaseModelLoader.load_model + + def _ascend_base_load_model( + self, vllm_config, model_config + ): + """Use eager safetensors loading for qwen3_next_mtp draft models. + + The memory-mapped lazy tensors produced by the default safetensors + iterator can become ``UntypedStorage`` objects when they are passed + through the MTP weight-loading pipeline on Ascend NPUs. Loading the + files eagerly into CPU memory before filtering avoids this problem. + + ``initialize_model`` passes the top-level ``vllm_config.model_config`` + to the model's ``__init__``, but when loading the draft model the real + config is the ``model_config`` argument. We temporarily replace + ``vllm_config.model_config`` so that draft-only rewrites do not touch + the target config. + """ + if getattr(model_config.hf_config, "model_type", None) == "qwen3_next_mtp": + load_config = vllm_config.load_config + original_strategy = load_config.safetensors_load_strategy + load_config.safetensors_load_strategy = "eager" + original_model_config = vllm_config.model_config + vllm_config.model_config = model_config + try: + return _orig_load_model(self, vllm_config, model_config) + finally: + vllm_config.model_config = original_model_config + load_config.safetensors_load_strategy = original_strategy + return _orig_load_model(self, vllm_config, model_config) + + _orig_set_moe_parameters = QwenNextMixtureOfExperts.set_moe_parameters + + def _ascend_set_moe_parameters(self): + try: + _orig_set_moe_parameters(self) + except RuntimeError as e: + # Dense qwen3.5/qwen3.5 models do not have any MoE layers in the MTP + # module. The upstream mixin raises in that case, so we swallow the + # specific error and leave the MoE-related attributes unset. + if "No Qwen3Next layer found" in str(e): + logger.debug( + "Skipping set_moe_parameters for dense Ascend MTP model" + ) + else: + raise + + _orig_load_weights = Qwen3NextMTP.load_weights + + def _ascend_qwen3_next_mtp_load_weights( + self, weights: "Iterable[tuple[str, torch.Tensor]]" + ) -> "set[str]": + """Load Qwen3NextMTP weights from qwen3.5/qwen3.5_moe checkpoints. + + Upstream ``Qwen3NextMTP.load_weights`` expects the qwen3_next layout + where shared embedding/head weights live at ``model.embed_tokens`` + and ``lm_head``. Qwen3.5/Qwen3.6 checkpoints wrap them under + ``model.language_model.*``; strip that prefix before delegating to the + original loader. + + For MoE checkpoints the routed expert weights are stored as packed + tensors (``experts.gate_up_proj`` / ``experts.down_proj``) while the + upstream ``Qwen3NextMultiTokenPredictor`` loader expects per-expert + slices (``experts.{eid}.gate_proj`` / ``up_proj`` / ``down_proj``) so + that ``FusedMoE.make_expert_params_mapping`` can route them to the + packed ``w13_weight`` / ``w2_weight`` model parameters. We split the + checkpoint tensors into 2-D per-expert slices before delegating. + """ + + # Only the MTP drafter weights carry the ``mtp.`` prefix; target-model + # MoE weights live under ``model.language_model.*`` and are consumed by + # the target-model loader, so we must not rename them here. + _MTP_PACKED_EXPERT_RE = re.compile( + r"^mtp\.layers\.(\d+)\.mlp\.experts\.(gate_up_proj|down_proj)$" + ) + + def _remap_language_model_prefix(weights): + for name, tensor in weights: + if name.startswith("model.language_model.embed_tokens"): + # The shared token embedding lives under the multimodal + # ``language_model`` wrapper in qwen3.5/qwen3.5_moe + # checkpoints, but Qwen3NextMTP keeps it inside ``model``. + name = name.replace( + "model.language_model.embed_tokens", + "model.embed_tokens", + 1, + ) + elif name.startswith("model.language_model."): + name = name[len("model.language_model."):] + yield name, tensor + + def _unpack_mtp_moe_packed_weights(weights): + """Split packed MoE expert tensors into 2-D per-expert slices. + + Qwen3.5/Qwen3.6 MoE checkpoints store routed experts as + ``experts.gate_up_proj`` (shape ``(E, 2*I, H)``) and + ``experts.down_proj`` (shape ``(E, H, I)``). The upstream + ``Qwen3NextMultiTokenPredictor`` loader expects per-expert + ``gate_proj``/``up_proj``/``down_proj`` slices, which + ``FusedMoE.make_expert_params_mapping`` routes to the packed + ``w13_weight``/``w2_weight`` model parameters. + """ + # Materialise the iterator: we need to inspect down_proj shapes + # before we can safely split gate_up_proj. + weight_list = list(weights) + + # Map layer index -> intermediate_size (I) inferred from down_proj. + layer_to_intermediate_size: dict[int, int] = {} + for name, tensor in weight_list: + match = _MTP_PACKED_EXPERT_RE.match(name) + if match is None: + continue + proj_type = match.group(2) + if proj_type != "down_proj" or tensor.ndim != 3: + continue + # down_proj shape is (E, H, I); the smallest non-expert dim is I. + layer_idx = int(match.group(1)) + layer_to_intermediate_size[layer_idx] = min( + tensor.shape[1], tensor.shape[2] + ) + + for name, tensor in weight_list: + match = _MTP_PACKED_EXPERT_RE.match(name) + if match is None or tensor.ndim != 3: + yield name, tensor + continue + + layer_idx = int(match.group(1)) + proj_type = match.group(2) + num_experts = tensor.shape[0] + + if proj_type == "gate_up_proj": + intermediate_size = layer_to_intermediate_size.get(layer_idx) + if intermediate_size is None: + logger.warning( + "No down_proj found for layer %s; cannot split " + "gate_up_proj, passing through unchanged", + layer_idx, + ) + yield name, tensor + continue + + packed_size = 2 * intermediate_size + packed_dim = None + for dim, size in enumerate(tensor.shape): + if dim == 0: + continue + if size == packed_size: + packed_dim = dim + break + + if packed_dim is None: + logger.warning( + "Unable to find packed dim (expected size %s) " + "for %s with shape %s; passing through unchanged", + packed_size, + name, + tensor.shape, + ) + yield name, tensor + continue + + gate_tensor, up_tensor = tensor.chunk(2, dim=packed_dim) + for expert_id in range(num_experts): + yield ( + f"mtp.layers.{layer_idx}.mlp.experts." + f"{expert_id}.gate_proj.weight", + gate_tensor[expert_id], + ) + yield ( + f"mtp.layers.{layer_idx}.mlp.experts." + f"{expert_id}.up_proj.weight", + up_tensor[expert_id], + ) + elif proj_type == "down_proj": + for expert_id in range(num_experts): + yield ( + f"mtp.layers.{layer_idx}.mlp.experts." + f"{expert_id}.down_proj.weight", + tensor[expert_id], + ) + else: + yield name, tensor + + return _orig_load_weights( + self, + _unpack_mtp_moe_packed_weights( + _remap_language_model_prefix(weights) + ), + ) + + _orig_init = Qwen3NextMTP.__init__ + + def _ascend_qwen3_next_mtp_init( + self, *, vllm_config, prefix: str = "" + ): + """Rewrite qwen3.5/qwen3.5_moe draft configs before building the model. + + ``ModelConfig`` may be reconstructed in worker processes without going + through the SpeculativeConfig hf_config_override path. Rewriting here + guarantees that ``Qwen3NextMTP`` always sees a flat ``qwen3_next_mtp`` + config with the vocabulary size and layer fields it expects. + + We deep-copy the config before rewriting so that the target model's + ``hf_config`` is not corrupted when the draft and target configs share + the same object (e.g. through ``AutoConfig`` caching). + """ + cfg = vllm_config.model_config.hf_config + if cfg.model_type in ("qwen3_5", "qwen3_5_moe"): + cfg = _rewrite_qwen3_5_config_for_mtp(copy.deepcopy(cfg)) + vllm_config.model_config.hf_config = cfg + return _orig_init(self, vllm_config=vllm_config, prefix=prefix) + + BaseModelLoader.load_model = _ascend_base_load_model + QwenNextMixtureOfExperts.set_moe_parameters = _ascend_set_moe_parameters + Qwen3NextMultiTokenPredictor.forward = _ascend_qwen3_next_mtp_forward + Qwen3NextMTP.__init__ = _ascend_qwen3_next_mtp_init + Qwen3NextMTP.load_weights = _ascend_qwen3_next_mtp_load_weights + logger.info("Patched Qwen3NextMTP loader/init/forward/load_weights for Ascend") + + +# --------------------------------------------------------------------------- +# Patch 3: Qwen3_5MultiTokenPredictor.forward (if present in this vLLM version) +# --------------------------------------------------------------------------- + +def patch_qwen3_5_mtp() -> None: + """Patch ``Qwen3_5MultiTokenPredictor.forward`` for Ascend if it exists.""" + try: + from vllm.model_executor.models.qwen3_5_mtp import ( + Qwen3_5MultiTokenPredictor, + ) + except ImportError: + logger.debug("Qwen3_5MultiTokenPredictor not available, skip MTP patch") + return + + Qwen3_5MultiTokenPredictor.forward = _ascend_qwen3_next_mtp_forward + logger.info("Patched Qwen3_5MultiTokenPredictor.forward for Ascend") + + +# --------------------------------------------------------------------------- +# Patch 4: SpeculativeConfig hf_config_override for qwen3.5/qwen3.6 MTP +# --------------------------------------------------------------------------- + +_SPECULATIVE_OVERRIDE_PATCHED = False +_ORIG_HF_CONFIG_OVERRIDE = None + + +def _rewrite_qwen3_5_config_for_mtp(hf_config): + """Flatten a qwen3.5/qwen3.5_moe config and return a ``Qwen3NextConfig``. + + Upstream ``Qwen3NextMTP`` / ``Qwen3NextMultiTokenPredictor`` instantiate + ``Qwen3NextDecoderLayer``, which expects a real ``Qwen3NextConfig`` object + (it reads ``decoder_sparse_step``, ``mlp_only_layers`` and uses the default + ``layer_types`` generator). Simply mutating the model_type on the original + ``Qwen3_5MoeConfig`` is not enough for MoE checkpoints, so we build a fresh + ``Qwen3NextConfig`` from the flattened text fields. + """ + from vllm.transformers_utils.configs import Qwen3NextConfig + + original_model_type = hf_config.model_type + + # Qwen3.5/Qwen3.6 checkpoints store language-model-specific fields in + # ``text_config`` (e.g. vocab_size, hidden_size, num_hidden_layers). + # ``Qwen3NextMTP`` expects these fields directly on the config object, + # so flatten them before rewriting the model type. + text_cfg = getattr(hf_config, "text_config", None) + if text_cfg is not None: + text_items = ( + text_cfg.items() + if isinstance(text_cfg, dict) + else text_cfg.__dict__.items() + ) + for key, value in text_items: + if key in ("model_type", "architectures", "text_config"): + continue + if not hasattr(hf_config, key) or getattr( + hf_config, key, None + ) is None: + setattr(hf_config, key, value) + + # Dense models do not define ``num_experts``; set a safe default so + # the generic Qwen3Next MoE code path skips expert mapping. + if not hasattr(hf_config, "num_experts"): + hf_config.num_experts = 0 + + # Fields accepted by Qwen3NextConfig.__init__. We copy whatever exists on + # the original config and supply sane defaults for the rest. + qwen3_next_fields = { + "vocab_size", + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "hidden_act", + "max_position_embeddings", + "initializer_range", + "rms_norm_eps", + "use_cache", + "tie_word_embeddings", + "rope_parameters", + "attention_bias", + "attention_dropout", + "head_dim", + "linear_conv_kernel_dim", + "linear_key_head_dim", + "linear_value_head_dim", + "linear_num_key_heads", + "linear_num_value_heads", + "decoder_sparse_step", + "moe_intermediate_size", + "shared_expert_intermediate_size", + "num_experts_per_tok", + "num_experts", + "norm_topk_prob", + "output_router_logits", + "router_aux_loss_coef", + "mlp_only_layers", + "layer_types", + } + + kwargs = {} + for key in qwen3_next_fields: + if hasattr(hf_config, key): + kwargs[key] = getattr(hf_config, key) + + # Defaults required by Qwen3NextDecoderLayer but absent from qwen3.5_moe. + kwargs.setdefault("decoder_sparse_step", 1) + kwargs.setdefault("mlp_only_layers", []) + kwargs.setdefault( + "intermediate_size", + getattr(hf_config, "moe_intermediate_size", kwargs.get("hidden_size", 0) * 4), + ) + + new_config = Qwen3NextConfig(**kwargs) + + # Preserve extra fields that upstream MTP code may rely on. + for key in ( + "n_predict", + "num_nextn_predict_layers", + "mtp_num_hidden_layers", + "mtp_use_dedicated_embeddings", + "dtype", + "image_token_index", + ): + if hasattr(hf_config, key) and not hasattr(new_config, key): + setattr(new_config, key, getattr(hf_config, key)) + + new_config.model_type = "qwen3_next_mtp" + n_predict = getattr( + new_config, "num_nextn_predict_layers", None + ) or getattr(new_config, "mtp_num_hidden_layers", None) + new_config.n_predict = n_predict + new_config.architectures = ["Qwen3NextMTP"] + + logger.info( + "Rewrote %s draft config to Qwen3NextConfig for Ascend MTP", + original_model_type, + ) + return new_config + + +def _ascend_hf_config_override(hf_config): + """Map qwen3.5/qwen3.6 configs to the upstream qwen3_next_mtp type. + + Defined at module level so that the callable can be pickled and sent to + worker processes (a closure capturing a local variable would not survive + the ``ModelConfig`` serialization). + """ + global _ORIG_HF_CONFIG_OVERRIDE + if _ORIG_HF_CONFIG_OVERRIDE is None: + from vllm.config.speculative import SpeculativeConfig + + _ORIG_HF_CONFIG_OVERRIDE = SpeculativeConfig.hf_config_override + + # The draft config may share the same hf_config object with the target + # config (e.g. through AutoConfig caching). Deep-copy before rewriting so + # the target model keeps its original qwen3.5/qwen3.5_moe type. + hf_config = copy.deepcopy(hf_config) + hf_config = _ORIG_HF_CONFIG_OVERRIDE(hf_config) + if hf_config.model_type in ("qwen3_5", "qwen3_5_moe"): + hf_config = _rewrite_qwen3_5_config_for_mtp(hf_config) + return hf_config + + +def patch_speculative_config_override() -> None: + """Install the module-level ``_ascend_hf_config_override``. + + Upstream vLLM only auto-detects ``qwen3_next`` models as MTP. Qwen3.5 and + Qwen3.6 checkpoints use ``qwen3_5`` / ``qwen3_5_moe`` model_type but share + the same MTP structure, so we rewrite the draft model config to + ``qwen3_next_mtp`` so that ``Qwen3NextMTP`` / ``Qwen3NextMultiTokenPredictor`` + are used. This lets the Ascend MTP forward patch above take effect. + + This patch must be applied at platform-plugin import time, before the main + process creates ``SpeculativeConfig`` (and therefore ``draft_model_config``). + """ + global _SPECULATIVE_OVERRIDE_PATCHED, _ORIG_HF_CONFIG_OVERRIDE + if _SPECULATIVE_OVERRIDE_PATCHED: + return + + try: + from vllm.config.speculative import SpeculativeConfig + except ImportError: + logger.debug("SpeculativeConfig not available, skip override patch") + return + + _ORIG_HF_CONFIG_OVERRIDE = SpeculativeConfig.hf_config_override + SpeculativeConfig.hf_config_override = _ascend_hf_config_override + _SPECULATIVE_OVERRIDE_PATCHED = True + logger.info("Patched SpeculativeConfig.hf_config_override for Ascend MTP") + + +def _normalize_qwen3_5_config(hf_config): + """Add aliases that upstream vLLM expects for qwen3.5/qwen3.5_moe configs. + + The multimodal target model config uses ``image_token_id`` but the + speculative-decoding drafter code looks for ``image_token_index``. + """ + if not hasattr(hf_config, "image_token_index") and hasattr( + hf_config, "image_token_id" + ): + hf_config.image_token_index = hf_config.image_token_id + + +def patch_model_config_for_qwen3_mtp() -> None: + """Patch ``ModelConfig.__post_init__`` for qwen3.5/qwen3.5_moe MTP. + + The ``hf_overrides`` callable attached to the draft ``ModelConfig`` is not + always preserved when the config is sent to worker processes. This wrapper + rewrites the config after the upstream post-init: + + - For target configs, it just adds the ``image_token_index`` alias. + - For draft configs, it flattens ``text_config`` and rewrites the model type + to ``qwen3_next_mtp`` so the correct model class is resolved. + """ + try: + from vllm.config import ModelConfig + from vllm.transformers_utils.config import get_hf_text_config + except ImportError: + logger.debug("ModelConfig not available, skip model config patch") + return + + if getattr(patch_model_config_for_qwen3_mtp, "_patched", False): + return + + _orig_post_init = ModelConfig.__post_init__ + + def _ascend_model_config_post_init(self, *args, **kwargs): + _orig_post_init(self, *args, **kwargs) + if self.hf_config.model_type not in ("qwen3_5", "qwen3_5_moe"): + return + + # Deep-copy so that target and draft configs do not share the same + # hf_config object after one of them is rewritten. + self.hf_config = copy.deepcopy(self.hf_config) + + _normalize_qwen3_5_config(self.hf_config) + + logger.info( + "[Ascend MTP] ModelConfig post-init: runner=%s runner_type=%s " + "model_type=%s architectures=%s", + self.runner, + self.runner_type, + self.hf_config.model_type, + getattr(self.hf_config, "architectures", None), + ) + + if self.runner == "draft": + self.hf_config = _rewrite_qwen3_5_config_for_mtp(self.hf_config) + self.hf_text_config = get_hf_text_config(self.hf_config) + # Re-resolve the model class now that the architecture has changed. + registry = self.registry + model_info, arch = registry.inspect_model_cls(self.architectures, self) + self._model_info = model_info + self._architecture = arch + logger.info("Resolved architecture: %s", arch) + + ModelConfig.__post_init__ = _ascend_model_config_post_init + patch_model_config_for_qwen3_mtp._patched = True + logger.info("Patched ModelConfig.__post_init__ for Ascend MTP") + + +def patch_qwen3_next_mtp_multimodal_flag() -> None: + """Force ``Qwen3NextMTP`` configs to be treated as text-only. + + ``ModelConfig`` may be reconstructed in worker processes before the + platform-level post-init patch is active, leaving the stale ``_model_info`` + from the original qwen3.5/qwen3.5_moe config. That can cause the multimodal + registry to believe the draft model supports multimodal inputs. Since + ``Qwen3NextMTP`` is text-only, override the property. + """ + try: + from vllm.config import ModelConfig + except ImportError: + logger.debug("ModelConfig not available, skip multimodal flag patch") + return + + if getattr(ModelConfig, "_ascend_qwen3_next_mtp_mm_patched", False): + return + + _orig = ModelConfig.is_multimodal_model.fget + + @property + def _ascend_is_multimodal_model(self): + if getattr(self.hf_config, "model_type", None) == "qwen3_next_mtp": + return False + return _orig(self) + + ModelConfig.is_multimodal_model = _ascend_is_multimodal_model + ModelConfig._ascend_qwen3_next_mtp_mm_patched = True + logger.info("Patched ModelConfig.is_multimodal_model for Ascend MTP") + + +def _lazy_init_triton_device_props() -> None: + """Initialize Ascend triton device properties if they are not ready yet.""" + try: + from vllm_ascend.ops.triton.triton_utils import init_device_properties_triton + init_device_properties_triton() + except Exception: + pass + + +def _wrap_with_triton_init(fn): + """Call device-property init before delegating to ``fn``.""" + + def wrapper(*args, **kwargs): + _lazy_init_triton_device_props() + return fn(*args, **kwargs) + + return wrapper + + +def patch_ascend_rejection_sampler() -> None: + """Replace upstream rejection-sampler helpers with Ascend implementations. + + The upstream Triton kernels used by ``vllm.v1.sample.rejection_sampler`` + are not compatible with Ascend NPUs. vllm-ascend provides Ascend-tuned + versions of the same module-level helpers; swap them in so that speculative + decoding / MTP works on Ascend through vllm-plugin-FL. + """ + try: + import vllm.v1.sample.rejection_sampler as rs + from vllm_ascend.sample.rejection_sampler import ( + apply_sampling_constraints, + expand_batch_to_tokens, + rejection_sample, + ) + except Exception as e: + logger.debug("Ascend rejection sampler helpers not available: %s", e) + return + + if getattr(patch_ascend_rejection_sampler, "_patched", False): + return + + rs.apply_sampling_constraints = _wrap_with_triton_init(apply_sampling_constraints) + rs.rejection_sample = _wrap_with_triton_init(rejection_sample) + rs.expand_batch_to_tokens = _wrap_with_triton_init(expand_batch_to_tokens) + patch_ascend_rejection_sampler._patched = True + logger.info("Patched rejection sampler helpers for Ascend") + + +def patch_ascend_eagle_proposer() -> None: + """Replace upstream Eagle padded-batch Triton kernels with PyTorch. + + The upstream ``eagle_prepare_next_token_padded_kernel`` and + ``eagle_prepare_inputs_padded_kernel`` use Triton patterns that fail to + compile on Ascend NPUs. Re-implement the same logic with standard PyTorch + ops so that EAGLE / MTP padded-batch speculative decoding works through + vllm-plugin-FL. + """ + try: + import vllm.v1.spec_decode.eagle as eagle_mod + except Exception as e: + logger.debug("Ascend EagleProposer patch not applicable: %s", e) + return + + if getattr(patch_ascend_eagle_proposer, "_patched", False): + return + + def _prepare_next_token_ids_padded( + self, + common_attn_metadata, + sampled_token_ids: torch.Tensor, + requests, + gpu_input_batch, + discard_request_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Precompute backup tokens for when there is no valid next token. + num_reqs = gpu_input_batch.num_reqs + self.backup_next_token_ids.np[:num_reqs] = np.array( + [ + requests[gpu_input_batch.req_ids[i]].get_token_id( + common_attn_metadata.seq_lens_cpu[i].item() + ) + for i in range(num_reqs) + ], + dtype=np.int32, + ) + self.backup_next_token_ids.copy_to_gpu(num_reqs) + backup_tokens_gpu = self.backup_next_token_ids.gpu + + batch_size, num_tokens = sampled_token_ids.shape + device = sampled_token_ids.device + + assert discard_request_mask.dtype == torch.bool + assert backup_tokens_gpu.dtype == torch.int32 + + vocab_size = gpu_input_batch.vocab_size + discard_mask = discard_request_mask[:batch_size] + + # Valid sampled tokens are in [0, vocab_size); -1 means rejected. + valid_mask = (sampled_token_ids != -1) & (sampled_token_ids < vocab_size) + valid_count = valid_mask.sum(dim=1).to(torch.int32) + valid_count = torch.where( + discard_mask, torch.zeros_like(valid_count), valid_count + ) + + # Find the last valid token index in each row. + positions = torch.arange( + num_tokens, device=device, dtype=torch.int64 + ).unsqueeze(0).expand(batch_size, -1) + last_valid_pos = torch.where( + valid_mask, + positions, + torch.full_like(positions, -1), + ).max(dim=1).values + last_valid_pos_safe = torch.clamp(last_valid_pos, min=0) + + selected = sampled_token_ids.gather( + 1, last_valid_pos_safe.unsqueeze(1) + ).squeeze(1).to(torch.int32) + has_valid = (valid_count > 0) & (~discard_mask) + + next_token_ids = torch.where( + has_valid, + selected, + backup_tokens_gpu[:batch_size], + ) + return next_token_ids, valid_count + + def _prepare_inputs_padded( + self, + common_attn_metadata: CommonAttentionMetadata, + spec_decode_metadata, + valid_sampled_tokens_count: torch.Tensor, + ) -> tuple[CommonAttentionMetadata, torch.Tensor]: + num_reqs = common_attn_metadata.num_reqs + device = valid_sampled_tokens_count.device + + token_indices_to_sample = torch.empty( + (num_reqs,), dtype=torch.int32, device=device + ) + + cu_num_draft_tokens = spec_decode_metadata.cu_num_draft_tokens.to( + torch.int32 + ) + num_draft_tokens = cu_num_draft_tokens.clone() + if num_reqs > 1: + num_draft_tokens[1:] -= cu_num_draft_tokens[:-1] + + num_rejected_tokens = torch.where( + num_draft_tokens > 0, + num_draft_tokens + 1 - valid_sampled_tokens_count, + torch.zeros_like(num_draft_tokens), + ) + + q_last_tok_idx = common_attn_metadata.query_start_loc[1:] - 1 + token_indices_to_sample = q_last_tok_idx - num_rejected_tokens + token_indices_to_sample = token_indices_to_sample.to(torch.int32) + + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu + new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + total_num_tokens = query_start_loc_cpu[-1].item() + + spec_common_attn_metadata = CommonAttentionMetadata( + query_start_loc=common_attn_metadata.query_start_loc, + seq_lens=common_attn_metadata.seq_lens, + query_start_loc_cpu=query_start_loc_cpu, + _seq_lens_cpu=common_attn_metadata._seq_lens_cpu, + _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu, + num_reqs=common_attn_metadata.num_reqs, + num_actual_tokens=total_num_tokens, + max_query_len=new_query_len_per_req.max().item(), + max_seq_len=common_attn_metadata.seq_lens_cpu.max().item(), + block_table_tensor=common_attn_metadata.block_table_tensor, + slot_mapping=common_attn_metadata.slot_mapping[:total_num_tokens], + causal=True, + dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens, + ) + return spec_common_attn_metadata, token_indices_to_sample + + eagle_mod.EagleProposer.prepare_next_token_ids_padded = ( + _prepare_next_token_ids_padded + ) + eagle_mod.EagleProposer.prepare_inputs_padded = _prepare_inputs_padded + patch_ascend_eagle_proposer._patched = True + logger.info("Patched EagleProposer padded helpers for Ascend") + + +def patch_mrope_for_graph_mode_mtp() -> None: + """Make ``MRotaryEmbedding.forward_native`` graph-safe for text positions. + + The upstream implementation uses ``query.view(num_tokens, -1, self.head_size)`` + where ``num_tokens = positions.shape[-1]``. Under ``torch.compile`` with + dynamic batch sizes this explicit symbolic dimension causes a data-dependent + guard inside the view meta. The MTP drafter model is compiled together with + the target model on Ascend, so we replace the view with a shape that only + uses concrete static dimensions (``num_heads`` and ``head_size``) and an + inferred leading dimension. + """ + try: + from vllm.model_executor.layers.rotary_embedding.mrope import ( + MRotaryEmbedding, + ) + except ImportError: + logger.debug("MRotaryEmbedding not available, skip graph-safe patch") + return + + if getattr(patch_mrope_for_graph_mode_mtp, "_patched", False): + return + + _orig_forward_native = MRotaryEmbedding.forward_native + + def _graph_safe_forward_native( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor | None = None, + offsets: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + assert positions.ndim == 1 or positions.ndim == 2 + assert key is not None + + # The multimodal (positions.ndim == 2) path is not exercised by the MTP + # drafter and is left unchanged to avoid altering target-model behavior. + if positions.ndim == 2: + return _orig_forward_native(self, positions, query, key, offsets) + + self._match_cos_sin_cache_dtype(query) + cos_sin = self.cos_sin_cache[positions] + cos, sin = cos_sin.chunk(2, dim=-1) + + query_shape = query.shape + num_heads = query_shape[-1] // self.head_size + query = query.view(-1, num_heads, self.head_size) + query_rot = query[..., : self.rotary_dim] + query_pass = query[..., self.rotary_dim :] + query_rot = self.apply_rotary_emb.forward_native(query_rot, cos, sin) + query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape) + + key_shape = key.shape + num_kv_heads = key_shape[-1] // self.head_size + key = key.view(-1, num_kv_heads, self.head_size) + key_rot = key[..., : self.rotary_dim] + key_pass = key[..., self.rotary_dim :] + key_rot = self.apply_rotary_emb.forward_native(key_rot, cos, sin) + key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) + + return query, key + + MRotaryEmbedding.forward_native = _graph_safe_forward_native + patch_mrope_for_graph_mode_mtp._patched = True + logger.info("Patched MRotaryEmbedding.forward_native for graph-safe MTP") + + +def patch_qwen3_mtp_platform() -> None: + """Apply only the platform-level MTP patch. + + ``SpeculativeConfig`` is constructed in the main (API) process before worker + processes are spawned, so the hf_config_override patch must be installed + during platform plugin registration. ``ModelConfig`` is reconstructed in + worker processes, so its post-init is also patched here. + """ + patch_speculative_config_override() + patch_model_config_for_qwen3_mtp() + patch_qwen3_next_mtp_multimodal_flag() + + +def patch_mamba_speculative_support() -> None: + """Allow qwen3.5/qwen3.5_moe models with Mamba layers to use MTP. + + Upstream vLLM only permits ``qwen3_next`` models to combine Mamba layers + with speculative decoding. The Ascend MTP patches above make qwen3.5/3.6 + work the same way, so expand the allow-list. + """ + try: + from vllm.model_executor.layers.mamba.abstract import MambaBase + from vllm.v1.kv_cache_interface import MambaSpec + except ImportError: + logger.debug("MambaBase not available, skip speculative support patch") + return + + _orig_get_kv_cache_spec = MambaBase.get_kv_cache_spec + + def _ascend_mamba_get_kv_cache_spec(self, vllm_config): + model_type = vllm_config.model_config.hf_config.model_type + architectures = getattr(vllm_config.model_config.hf_config, "architectures", None) + logger.info( + "[Ascend MTP] MambaBase.get_kv_cache_spec called for model_type=%s " + "architectures=%s speculative_config=%s", + model_type, architectures, vllm_config.speculative_config is not None, + ) + if ( + vllm_config.speculative_config is not None + and model_type + not in ("qwen3_next", "qwen3_next_mtp", "qwen3_5", "qwen3_5_moe") + ): + logger.info( + "[Ascend MTP] model_type %s not in allow-list, delegating to upstream", + model_type, + ) + return _orig_get_kv_cache_spec(self, vllm_config) + + mamba_block_size = vllm_config.cache_config.mamba_block_size + page_size_padded = vllm_config.cache_config.mamba_page_size_padded + return MambaSpec( + shapes=self.get_state_shape(), + dtypes=self.get_state_dtype(), + block_size=mamba_block_size, + page_size_padded=page_size_padded, + mamba_type=self.mamba_type, + num_speculative_blocks=( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ), + ) + + MambaBase.get_kv_cache_spec = _ascend_mamba_get_kv_cache_spec + logger.info("Patched MambaBase.get_kv_cache_spec for Ascend MTP") + + +def patch_qwen3_5_mtp_post_init() -> None: + """Clear stale target-model ``_model_info`` after draft config is cloned. + + When the speculative config clones the base config into a draft config, + the cloned ``ModelConfig`` object still carries the target model's cached + ``_model_info`` and ``_architecture``. For ``Qwen3NextMTP`` this causes + vLLM to resolve the draft architecture as the target architecture. Clear + the cached values so that the draft model is re-resolved correctly. + """ + try: + from vllm.config import ModelConfig + except ImportError: + logger.debug("ModelConfig not available, skip qwen3.5 MTP post-init patch") + return + + _orig_post_init = ModelConfig.__post_init__ + + def _ascend_qwen3_5_mtp_post_init(self, *args, **kwargs): + _orig_post_init(self, *args, **kwargs) + if getattr(self.hf_config, "model_type", None) == "qwen3_next_mtp": + return + speculative_config = getattr(self, "speculative_config", None) + if ( + speculative_config is not None + and getattr(speculative_config, "_draft_hf_config", None) is not None + and speculative_config._draft_hf_config.model_type == "qwen3_next_mtp" + ): + self._model_info = None + self._architecture = None + + ModelConfig.__post_init__ = _ascend_qwen3_5_mtp_post_init + logger.info("Patched ModelConfig.__post_init__ for qwen3.5 MTP draft") + + +def patch_qwen3_mtp_config_overrides() -> None: + """Apply ``SpeculativeConfig.hf_config_override`` without mutating target config. + + Upstream vLLM applies the override dict to the shared ``hf_config`` object, + which means the target model's config is also rewritten (e.g. + ``model_type='qwen3_next_mtp'``). Apply the override only to the draft + config copy that ``SpeculativeConfig`` keeps internally. + """ + try: + from vllm.config import ModelConfig + except ImportError: + logger.debug("ModelConfig not available, skip MTP config override patch") + return + + _orig = ModelConfig.__post_init__ + + def _ascend_mtp_config_overrides(self, *args, **kwargs): + speculative_config = getattr(self, "speculative_config", None) + if ( + speculative_config is not None + and getattr(speculative_config, "hf_config_override", None) + and getattr(speculative_config, "_draft_hf_config", None) is not None + ): + override = speculative_config.hf_config_override + for key, value in override.items(): + if hasattr(speculative_config._draft_hf_config, key): + setattr(speculative_config._draft_hf_config, key, value) + _orig(self, *args, **kwargs) + + ModelConfig.__post_init__ = _ascend_mtp_config_overrides + logger.info("Patched ModelConfig.__post_init__ for MTP config overrides") + + +def patch_cudagraph_env_vars() -> None: + """Honor VLLM_CUDAGRAPH_* env vars to limit NPU graph capture sizes. + + vLLM's ``--compilation-config`` is not always exposed in entrypoint + scripts, so we provide environment variables as a lightweight alternative: + + - ``VLLM_CUDAGRAPH_CAPTURE_SIZES``: comma-separated list of batch sizes to + capture (e.g. ``1,2,4,8,16,32,48,64,80,96``). + - ``VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE``: integer upper bound (optional). + """ + import os + + cg_sizes_env = os.environ.get("VLLM_CUDAGRAPH_CAPTURE_SIZES") + max_size_env = os.environ.get("VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE") + if not cg_sizes_env and not max_size_env: + return + + try: + from vllm.config.vllm import VllmConfig + except ImportError: + logger.debug("VllmConfig not available, skip cudagraph env patch") + return + + _orig_post_init = VllmConfig.__post_init__ + + def _ascend_cudagraph_env_post_init(self, *args, **kwargs): + _cg_sizes_env = os.environ.get("VLLM_CUDAGRAPH_CAPTURE_SIZES") + _max_size_env = os.environ.get("VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE") + if _cg_sizes_env: + sizes = [ + int(x.strip()) + for x in _cg_sizes_env.split(",") + if x.strip() + ] + self.compilation_config.cudagraph_capture_sizes = sizes + self.compilation_config.max_cudagraph_capture_size = sizes[-1] + if _max_size_env: + self.compilation_config.max_cudagraph_capture_size = int(_max_size_env) + _orig_post_init(self, *args, **kwargs) + + VllmConfig.__post_init__ = _ascend_cudagraph_env_post_init + logger.info( + "Patched VllmConfig.__post_init__ to honor VLLM_CUDAGRAPH_* env vars" + ) + + +def patch_qwen3_mtp() -> None: + """Apply all Ascend MTP-related runtime patches.""" + patch_bind_kv_cache() + patch_qwen3_next_mtp() + patch_qwen3_5_mtp() + patch_mamba_speculative_support() + patch_qwen3_next_mtp_multimodal_flag() + patch_qwen3_5_mtp_post_init() + patch_qwen3_mtp_config_overrides() + patch_ascend_rejection_sampler() + patch_ascend_eagle_proposer() + patch_mrope_for_graph_mode_mtp() + patch_cudagraph_env_vars() diff --git a/vllm_fl/dispatch/policy.py b/vllm_fl/dispatch/policy.py index e7f06139f..80719a880 100644 --- a/vllm_fl/dispatch/policy.py +++ b/vllm_fl/dispatch/policy.py @@ -14,6 +14,7 @@ from dataclasses import dataclass, field from typing import Any, Dict, FrozenSet, List, Optional, Set, Tuple +import torch from vllm_fl.utils import get_op_config @@ -166,6 +167,10 @@ def __init__(self): "vllm_fl_selection_policy", default=None, ) + # Thread-local mirror of the context var. Dynamo cannot trace + # ContextVar.get(), but plain attribute access on thread-local storage + # is fine. + self._thread_policy = threading.local() @classmethod def get_instance(cls): @@ -188,7 +193,30 @@ def bump_policy_epoch(self) -> int: return self._policy_epoch def get_policy(self) -> SelectionPolicy: - """Get the current effective policy (context or global).""" + """Get the current effective policy (context or global). + + Uses a thread-local mirror of the context var so that the hot path + does not need ContextVar.get(), which torch.compile/Dynamo cannot + trace. The ContextVar is still updated by _PolicyContext for + correctness in non-compiled code. + """ + # Fast, Dynamo-friendly path: check the thread-local mirror first. + thread_policy = getattr(self._thread_policy, "current", None) + if thread_policy is not None: + return thread_policy + + # When being traced by torch.compile, avoid ContextVar.get() entirely + # because Dynamo cannot trace it. Context overrides are mirrored in + # thread-local storage above, so reaching here under compilation means + # we should use the global policy. + if torch._dynamo.is_compiling(): + if self._global_policy is None: + with self._global_policy_lock: + if self._global_policy is None: + self._global_policy = self._policy_from_env() + return self._global_policy + + # Non-compiled path: keep ContextVar for correctness. ctx_policy = self._policy_var.get() if ctx_policy is not None: return ctx_policy @@ -481,6 +509,9 @@ def __init__(self, manager: PolicyManager, policy: SelectionPolicy): def __enter__(self) -> "_PolicyContext": policy_var = self._manager._get_policy_var() self._token = policy_var.set(self._policy) + # Mirror the context override in thread-local storage for the + # Dynamo-friendly fast path in get_policy. + self._manager._thread_policy.current = self._policy self._manager.bump_policy_epoch() return self @@ -488,6 +519,7 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> None: if self._token is not None: policy_var = self._manager._get_policy_var() policy_var.reset(self._token) + self._manager._thread_policy.current = None self._manager.bump_policy_epoch() diff --git a/vllm_fl/ops/pto_chunk_gdn/__init__.py b/vllm_fl/ops/pto_chunk_gdn/__init__.py new file mode 100644 index 000000000..7b76919bb --- /dev/null +++ b/vllm_fl/ops/pto_chunk_gdn/__init__.py @@ -0,0 +1,38 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""PTO chunk-GDN megakernel module for vLLM-Ascend. + +Provides a Bisheng-JIT-compiled fused Ascend NPU megakernel for the chunk +GatedDeltaNet (GDN) recurrent layer used in Qwen3.5 / Qwen3.6 models, +replacing the default Triton baseline during prefill. + +Enable via: + VLLM_ASCEND_PTO_CHUNK_GDN=1 (set before launching vLLM) + +The decode phase always uses the original Triton implementation. +""" +from vllm_fl.ops.pto_chunk_gdn.compile import BLOCK_DIM, PTO_LIB_PATH +from vllm_fl.ops.pto_chunk_gdn.mega_kernel import run_mega_kernel +from vllm_fl.ops.pto_chunk_gdn.worker_hook import apply_pto_gdn_patch, is_pto_gdn_patch_active + +__all__ = [ + "BLOCK_DIM", + "PTO_LIB_PATH", + "run_mega_kernel", + "apply_pto_gdn_patch", + "is_pto_gdn_patch_active", +] diff --git a/vllm_fl/ops/pto_chunk_gdn/chunk_gated_delta_wrapper.py b/vllm_fl/ops/pto_chunk_gdn/chunk_gated_delta_wrapper.py new file mode 100644 index 000000000..fc2c8acfe --- /dev/null +++ b/vllm_fl/ops/pto_chunk_gdn/chunk_gated_delta_wrapper.py @@ -0,0 +1,139 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""PTO megakernel drop-in for ``chunk_gated_delta_rule`` — prefill only. + +Activated via ``VLLM_ASCEND_PTO_CHUNK_GDN=1``. Falls back transparently to +the Triton implementation for: + - Non-zero ``initial_state`` (decode step with prior recurrent state) + - Missing ``cu_seqlens`` (non-varlen path) + - Multi-device pipeline-parallel (PCP) groups (world_size > 1) + - Mismatched Q/K/V head dimensions + - Non-NPU device + +GQA is supported: if ``v.shape[2] > q.shape[2]`` the GQA path is taken. + +The PTO path only executes for **prefill** (zero initial_state, cu_seqlens +present). Decode transparently falls back to Triton so normal inference +(prefill + decode) works correctly end-to-end. +""" +from __future__ import annotations + +import torch +from einops import rearrange + + +def _needs_triton_fallback( + initial_state: torch.Tensor | None, + cu_seqlens: torch.LongTensor | None, +) -> bool: + if initial_state is not None and torch.any(initial_state != 0): + return True + return cu_seqlens is None + + +@torch.compiler.disable +def chunk_gated_delta_rule_pto( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + prebuilt_meta=None, + head_first: bool = False, + use_qk_l2norm_in_kernel: bool = False, + *, + _triton_impl, +): + """PTO megakernel drop-in for ``chunk_gated_delta_rule``. + + Runs the fused Bisheng megakernel for prefill. Falls back to Triton for + decode (``initial_state != 0``) so end-to-end inference is unaffected. + """ + assert q.dtype == k.dtype == v.dtype + assert q.dtype != torch.float32, "Use bfloat16 or float16, not float32." + assert beta.ndim == 3, "beta must be [B, T, H] (head_first=False)." + + if head_first: + q, k, v, beta, g = (rearrange(x, "b h t ... -> b t h ...") for x in (q, k, v, beta, g)) + + if scale is None: + scale = float(k.shape[-1] ** -0.5) + + if use_qk_l2norm_in_kernel: + from vllm_fl.dispatch.backends.vendor.ascend.impl.fla.l2norm import l2norm_fwd + q, k = l2norm_fwd(q), l2norm_fwd(k) + + def _triton(): + return _triton_impl( + q, k, v, g, beta, + scale=scale, initial_state=initial_state, + output_final_state=output_final_state, cu_seqlens=cu_seqlens, + head_first=False, + use_qk_l2norm_in_kernel=False, + ) + + if q.device.type != "npu": + return _triton() + if _needs_triton_fallback(initial_state, cu_seqlens): + return _triton() + if q.shape[3] != v.shape[3]: + return _triton() + if v.shape[2] != q.shape[2] and v.shape[2] % q.shape[2] != 0: + return _triton() + try: + from vllm.distributed import get_pcp_group + if get_pcp_group().world_size > 1: + return _triton() + except Exception: + pass + + from vllm_fl.ops.pto_chunk_gdn.mega_kernel import run_mega_kernel + + kh = q.shape[2] + cu32 = cu_seqlens.to(torch.int32).contiguous() + stream = torch.npu.current_stream()._as_parameter_ + + with torch.autograd.profiler.record_function("PTO_mega_kernel"): + result = run_mega_kernel( + q.to(torch.float16), + k.to(torch.float16), + v.to(torch.float16), + g.float(), + beta.to(torch.float16), + cu32, + stream=stream, + chunk_size=128, + scale=scale, + key_heads=kh, + return_final_state=output_final_state, + ) + + if output_final_state: + o, fs = result + o = o.to(q.dtype) + fs = fs.to(q.dtype) + else: + o = result.to(q.dtype) # type: ignore[assignment] + fs = None + + if head_first: + o = rearrange(o, "b t h ... -> b h t ...") + return o, fs diff --git a/vllm_fl/ops/pto_chunk_gdn/compile.py b/vllm_fl/ops/pto_chunk_gdn/compile.py new file mode 100644 index 000000000..5dd192a42 --- /dev/null +++ b/vllm_fl/ops/pto_chunk_gdn/compile.py @@ -0,0 +1,182 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Bisheng JIT compilation for the PTO GDN megakernel on Ascend NPU. + +The megakernel is compiled on first use and cached under +``vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib/``. +Re-compilation is triggered when the C++ source mtime changes. + +Environment variables: + PTO_LIB_PATH Path to pto-isa header directory (contains ``include/``). + Auto-detected from ``csrc/ascend/third_party/pto-isa`` in the + package source tree, then ``/sources/pto-isa`` fallback. + ASCEND_TOOLKIT_HOME Ascend toolkit root (required). + GDN_NPU_DEVICE NPU device for ``cube_core_num`` query (default ``npu:0``). + VERBOSE_COMPILE Set to ``1`` to print the full bisheng command. + PTO_DYNAMIC_EXTRA_FLAGS Extra flags appended to every bisheng invocation. +""" +from __future__ import annotations + +import os +import subprocess +from functools import lru_cache +from pathlib import Path + +import torch + +# --------------------------------------------------------------------------- +# Paths — resolved relative to this file's installed location +# --------------------------------------------------------------------------- +_THIS_DIR = Path(__file__).resolve().parent +_VLLM_FL_DIR = _THIS_DIR.parent.parent # vllm_fl/ +_PACKAGE_ROOT = _VLLM_FL_DIR.parent # site-packages root + +# C++ sources live in csrc/ascend/pto_chunk_gdn/ inside the source tree. +# In editable installs (pip install -e .) _PACKAGE_ROOT == repo root. +_CSRC_PTO = _PACKAGE_ROOT / "csrc" / "ascend" / "pto_chunk_gdn" +if not _CSRC_PTO.is_dir(): + _CSRC_PTO = _THIS_DIR / "csrc" + +KERNELS_PTO: str = str(_CSRC_PTO) +KERNEL_INCLUDE: str = str(_CSRC_PTO / "include") + +# Compiled .so cache — writable at runtime inside the installed package. +_COMPILED_DIR = _THIS_DIR / "kernels" / "compiled_lib" +COMPILED_DIR: str = str(_COMPILED_DIR) + +_DRIVER_INC = os.environ.get("ASCEND_DRIVER_PATH", "/usr/local/Ascend/driver/kernel/inc") + +ASCEND_TOOLKIT_HOME: str = ( + os.environ.get("ASCEND_TOOLKIT_HOME") or os.environ.get("ASCEND_HOME_PATH", "") +) +if not ASCEND_TOOLKIT_HOME: + raise RuntimeError( + "ASCEND_TOOLKIT_HOME (or ASCEND_HOME_PATH) must be set to the Ascend toolkit root." + ) + + +def _resolve_pto_lib_path() -> str: + if "PTO_LIB_PATH" in os.environ: + return os.environ["PTO_LIB_PATH"] + submodule = _PACKAGE_ROOT / "csrc" / "ascend" / "third_party" / "pto-isa" + if (submodule / "include").is_dir(): + os.environ["PTO_LIB_PATH"] = str(submodule) + return str(submodule) + raise RuntimeError( + "PTO ISA headers not found. Please either:\n" + " 1. Initialize the submodule: " + "git submodule update --init --recursive csrc/ascend/third_party/pto-isa\n" + " 2. Set PTO_LIB_PATH to the pto-isa directory (contains include/)." + ) + + +PTO_LIB_PATH: str = _resolve_pto_lib_path() + +# --------------------------------------------------------------------------- +# Hardware: query cube_core_num +# --------------------------------------------------------------------------- +_npu_dev = os.environ.get("GDN_NPU_DEVICE", "npu:0") +try: + BLOCK_DIM: int = int( + getattr(torch.npu.get_device_properties(_npu_dev), "cube_core_num", 20) + ) +except (RuntimeError, AssertionError): + BLOCK_DIM = 24 + + +# --------------------------------------------------------------------------- +# Compilation +# --------------------------------------------------------------------------- + +def _common_flags( + *, + num_heads: int, + key_heads: int, + hidden_size: int, + chunk_size: int, +) -> list[str]: + flags = [ + "-fPIC", "-shared", "-xcce", "-DMEMORY_BASE", "-O2", "-std=gnu++17", + "--cce-aicore-arch=dav-c220", + "-mllvm", "-cce-aicore-stack-size=0x8000", + "-mllvm", "-cce-aicore-function-stack-size=0x8000", + "-mllvm", "-cce-aicore-record-overflow=true", + "-mllvm", "-cce-aicore-dcci-insert-for-scalar=false", + "-Wno-macro-redefined", "-Wno-ignored-attributes", + f"-I{KERNEL_INCLUDE}", + f"-I{os.path.join(PTO_LIB_PATH, 'include')}", + f"-I{ASCEND_TOOLKIT_HOME}/include", + f"-I{ASCEND_TOOLKIT_HOME}/pkg_inc", + f"-I{ASCEND_TOOLKIT_HOME}/pkg_inc/runtime", + f"-I{ASCEND_TOOLKIT_HOME}/pkg_inc/profiling", + f"-DGDN_H={num_heads}", + f"-DGDN_HG={key_heads}", + f"-DGDN_D={hidden_size}", + f"-DGDN_C={chunk_size}", + ] + if os.path.isdir(_DRIVER_INC): + flags.append(f"-I{_DRIVER_INC}") + extra = os.environ.get("PTO_DYNAMIC_EXTRA_FLAGS", "").split() + flags.extend(extra) + return flags + + +@lru_cache(maxsize=None) +def compile_mega_kernel( + *, + num_heads: int = 16, + key_heads: int | None = None, + hidden_size: int = 128, + chunk_size: int = 128, + cpp_mtime_ns: int = 0, +) -> str: + """Compile the fused PTO GDN megakernel and return the ``.so`` path. + + Args: + num_heads: Number of value heads H. + key_heads: Number of Q/K heads Hg (GQA; defaults to H if None). + hidden_size: Head dimension D. + chunk_size: Chunk size C (must be 128). + cpp_mtime_ns: Source file mtime for cache invalidation. + + Returns: + Absolute path to the compiled ``.so`` file. + """ + kh = key_heads if key_heads is not None else num_heads + os.makedirs(COMPILED_DIR, exist_ok=True) + cpp_path = os.path.join(KERNELS_PTO, "mega_kernel.cpp") + lib_path = os.path.join( + COMPILED_DIR, + f"mega_kernel_H{num_heads}_Hg{kh}_D{hidden_size}_C{chunk_size}.so", + ) + if os.path.exists(lib_path): + # Reuse the cached build: avoids every TP worker recompiling into the + # same output path on each server start. + return lib_path + flags = _common_flags( + num_heads=num_heads, key_heads=kh, hidden_size=hidden_size, chunk_size=chunk_size + ) + cmd = ["bisheng", *flags, cpp_path, "-o", lib_path] + if os.environ.get("VERBOSE_COMPILE"): + print("compile:", " ".join(cmd)) + import logging + logging.getLogger(__name__).info( + "[pto_chunk_gdn] Compiling mega_kernel H=%d Hg=%d D=%d C=%d …", + num_heads, kh, hidden_size, chunk_size, + ) + subprocess.run(cmd, check=True, timeout=600) + return lib_path diff --git a/vllm_fl/ops/pto_chunk_gdn/mega_kernel.py b/vllm_fl/ops/pto_chunk_gdn/mega_kernel.py new file mode 100644 index 000000000..b9617aa24 --- /dev/null +++ b/vllm_fl/ops/pto_chunk_gdn/mega_kernel.py @@ -0,0 +1,215 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Fused PTO GDN megakernel: all six GDN pipeline stages in one NPU launch. + +Fuses cumsum → scaled_dot_kkt → solve_tril → wy_fast → chunk_h → chunk_o +into a single ``call_kernel`` dispatch, eliminating Python-level inter-stage +synchronization overhead. + +GQA is supported: Q/K use ``Hg`` heads while V/gates use ``H ≥ Hg`` value heads +(``H % Hg == 0``). + +Usage:: + + from vllm_ascend.ops.pto_chunk_gdn.mega_kernel import run_mega_kernel + + o = run_mega_kernel(q, k, v, g, beta, cu_seqlens, stream=stream, + chunk_size=128, scale=head_dim**-0.5, key_heads=Hg) +""" +from __future__ import annotations + +import ctypes +import os +from functools import lru_cache + +import torch + +from vllm_fl.ops.pto_chunk_gdn.compile import BLOCK_DIM, KERNELS_PTO, compile_mega_kernel + + +# --------------------------------------------------------------------------- +# Small utilities (self-contained, no external module dependency) +# --------------------------------------------------------------------------- + +def _vp(t: torch.Tensor | None) -> ctypes.c_void_p: + if t is None: + return ctypes.c_void_p() + return ctypes.c_void_p(t.data_ptr()) + + +@lru_cache(maxsize=48) +def _precomputed_minus_identity( + device_ty: str, device_index: int, chunk_size: int +) -> torch.Tensor: + """``[C, C]`` fp16 buffer with diagonal ``-1``, cached per (device, C).""" + idx = max(device_index, 0) + dev = torch.device(device_ty, idx) if device_ty != "cpu" else torch.device("cpu") + t = torch.zeros(chunk_size, chunk_size, device=dev, dtype=torch.float16) + t.fill_diagonal_(-1) + return t + + +@lru_cache(maxsize=48) +def _causal_masks( + device_ty: str, device_index: int, chunk_size: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Lower-triangle and full causal masks, cached per (device, C).""" + idx = max(device_index, 0) + dev = torch.device(device_ty, idx) if device_ty != "cpu" else torch.device("cpu") + m_lower = torch.tril(torch.ones(chunk_size, chunk_size, device=dev), diagonal=-1).float() + m_full = torch.tril(torch.ones(chunk_size, chunk_size, device=dev), diagonal=0).float() + return m_lower, m_full + + +def _total_chunks(N_seq: int, cu_seqlens: torch.Tensor, chunk_size: int) -> int: + """Total number of chunks across all sequences in the varlen batch.""" + cu = cu_seqlens.cpu().tolist() + return sum( + (cu[i + 1] - cu[i] + chunk_size - 1) // chunk_size + for i in range(N_seq) + ) + + +# --------------------------------------------------------------------------- +# Kernel loading +# --------------------------------------------------------------------------- + +@lru_cache(maxsize=None) +def _load_mega_kernel( + *, + num_heads: int, + key_heads: int, + hidden_size: int, + chunk_size: int, +) -> ctypes.CDLL: + mtime = os.stat(os.path.join(KERNELS_PTO, "mega_kernel.cpp")).st_mtime_ns + lib_path = compile_mega_kernel( + num_heads=num_heads, key_heads=key_heads, + hidden_size=hidden_size, chunk_size=chunk_size, + cpp_mtime_ns=mtime, + ) + lib = ctypes.CDLL(os.path.abspath(lib_path)) + lib.call_kernel.argtypes = ( + [ctypes.c_uint32, ctypes.c_void_p] + + [ctypes.c_void_p] * 28 + + [ctypes.c_int64, ctypes.c_int64, ctypes.c_int64, ctypes.c_uint32] + ) + lib.call_kernel.restype = None + return lib + + +# --------------------------------------------------------------------------- +# Public entry point +# --------------------------------------------------------------------------- + +def run_mega_kernel( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_in: torch.Tensor, + beta: torch.Tensor, + cu_seqlens: torch.Tensor, + *, + stream, + chunk_size: int = 128, + scale: float = 1.0, + block_dim: int | None = None, + key_heads: int | None = None, + return_final_state: bool = False, + total_chunks: int | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Run all six GDN pipeline stages in a single fused NPU kernel launch. + + Args: + q, k: ``[B, T, Hg, D]`` fp16 query / key tensors. + v: ``[B, T, H, D]`` fp16 value tensor (H ≥ Hg, H % Hg == 0). + g_in: ``[B, T, H]`` float32 pre-cumsum gate logits. + beta: ``[B, T, H]`` fp16 gate bias. + cu_seqlens: ``int32`` cumulative sequence-length boundaries ``[0, …, T]``. + stream: NPU stream handle (``torch.npu.current_stream()._as_parameter_``). + chunk_size: Tile side length C (must be 128). + scale: Output scale (typically ``head_dim ** -0.5``). + block_dim: AI-Core block count; auto-detected from device if None. + key_heads: Q/K head count Hg (inferred from ``q`` if None). + return_final_state: Also return ``[N_seq, H, D, D]`` final recurrent states. + total_chunks: Precomputed chunk count; when given, skips the + ``cu_seqlens.cpu()`` host sync in ``_total_chunks``. + + Returns: + ``O * scale`` of shape ``[B, T, H, D]`` fp16, and optionally final + recurrent state ``[N_seq, H, D, D]`` fp16. + """ + dev = q.device + kh = key_heads if key_heads is not None else q.shape[2] + H, D = v.shape[2], q.shape[3] + C = chunk_size + T = q.shape[1] + N_seq = int(cu_seqlens.numel()) - 1 + bd = block_dim or BLOCK_DIM + + if cu_seqlens.dtype != torch.int32: + cu_seqlens = cu_seqlens.to(torch.int32) + + dt, di = dev.type, dev.index if dev.index is not None else -1 + msk_lower, msk_full = _causal_masks(dt, di, C) + minus_identity = _precomputed_minus_identity(dt, di, C) + + tc = ( + total_chunks + if total_chunks is not None + else _total_chunks(N_seq, cu_seqlens, C) + ) + num_matrices = tc * H + + g_sum = torch.empty(1, T, H, device=dev, dtype=torch.float32) + g_t = torch.empty(H, T, device=dev, dtype=torch.float32) + beta_t = torch.empty(H, T, device=dev, dtype=torch.float16) + A = torch.zeros(1, T, H, C, device=dev, dtype=torch.float16) + A_inv_f32 = torch.zeros(1, T, H, C, device=dev, dtype=torch.float32) + A_inv = torch.zeros(1, T, H, C, device=dev, dtype=torch.float16) + w = torch.empty_like(v) + u = torch.empty_like(v) + s = torch.zeros(tc * H, D, D, device=dev, dtype=torch.float16) + v_new = torch.empty_like(v) + fs = torch.zeros(N_seq * H, D, D, device=dev, dtype=torch.float16) + kkt_ws = torch.zeros(bd * 2, C, C, device=dev, dtype=torch.float16) + wy_ws_a1 = torch.zeros(bd, C, C, device=dev, dtype=torch.float16) + wy_ws_a2 = torch.zeros(bd, C, C, device=dev, dtype=torch.float16) + h_ws = torch.zeros(bd * 4, D, D, device=dev, dtype=torch.float16) + o_ws_qk = torch.zeros(bd, C, C, device=dev, dtype=torch.float16) + o_ws_qs = torch.zeros(bd, C, D, device=dev, dtype=torch.float16) + o_ws_gated = torch.zeros(bd, C, C, device=dev, dtype=torch.float16) + o_out = torch.empty_like(v) + + lib = _load_mega_kernel(num_heads=H, key_heads=kh, hidden_size=D, chunk_size=C) + lib.call_kernel( + bd, stream, + _vp(q), _vp(k), _vp(v), _vp(g_in), _vp(beta), + _vp(msk_lower), _vp(msk_full), _vp(minus_identity), _vp(cu_seqlens), + _vp(o_out), + _vp(g_sum), _vp(g_t), _vp(beta_t), + _vp(A), _vp(A_inv_f32), _vp(A_inv), + _vp(w), _vp(u), _vp(s), _vp(v_new), _vp(fs), + _vp(kkt_ws), _vp(wy_ws_a1), _vp(wy_ws_a2), _vp(h_ws), + _vp(o_ws_qk), _vp(o_ws_qs), _vp(o_ws_gated), + N_seq, T, T, num_matrices, + ) + + o_scaled = o_out * scale + if return_final_state: + return o_scaled, fs.view(N_seq, H, D, D) + return o_scaled diff --git a/vllm_fl/ops/pto_chunk_gdn/worker_hook.py b/vllm_fl/ops/pto_chunk_gdn/worker_hook.py new file mode 100644 index 000000000..77d32e440 --- /dev/null +++ b/vllm_fl/ops/pto_chunk_gdn/worker_hook.py @@ -0,0 +1,84 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Native worker hook: activate the PTO megakernel for ``chunk_gated_delta_rule``. + +Called from ``vllm_ascend.patch.worker.__init__`` when +``VLLM_ASCEND_PTO_CHUNK_GDN=1``. Must run **after** the Triton patches +install the baseline implementation and **before** any model module imports +``gdn.py``. +""" +from __future__ import annotations + +import logging +import sys + +_log = logging.getLogger(__name__) +_PATCH_ACTIVE = False + + +def apply_pto_gdn_patch() -> None: + """Replace ``chunk_gated_delta_rule`` with the PTO megakernel. + + Patches three locations so all existing import paths see the PTO version: + 1. ``vllm_ascend.ops.triton.fla.chunk`` — primary defining module. + 2. ``vllm.model_executor.layers.fla.ops`` — vLLM public FLA namespace. + 3. ``vllm_ascend.ops.gdn`` if already imported (refreshes module attribute). + """ + global _PATCH_ACTIVE + + import vllm.model_executor.layers.fla.ops as fla_ops + import vllm_ascend.ops.triton.fla.chunk as _ascend_chunk_mod + from vllm_ascend.ops.pto_chunk_gdn.chunk_gated_delta_wrapper import ( + chunk_gated_delta_rule_pto, + ) + + triton_impl = _ascend_chunk_mod.chunk_gated_delta_rule + + def _pto_bound( + q, k, v, g, beta, + scale=None, initial_state=None, output_final_state=False, + cu_seqlens=None, prebuilt_meta=None, head_first=False, + use_qk_l2norm_in_kernel=False, + ): + return chunk_gated_delta_rule_pto( + q, k, v, g, beta, + scale=scale, initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, prebuilt_meta=prebuilt_meta, + head_first=head_first, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + _triton_impl=triton_impl, + ) + + _pto_bound.__name__ = "chunk_gated_delta_rule" + _pto_bound._vllm_ascend_pto_gdn = True + + _ascend_chunk_mod.chunk_gated_delta_rule = _pto_bound + fla_ops.chunk_gated_delta_rule = _pto_bound + + _gdn_mod = sys.modules.get("vllm_ascend.ops.gdn") + if _gdn_mod is not None and hasattr(_gdn_mod, "chunk_gated_delta_rule"): + _gdn_mod.chunk_gated_delta_rule = _pto_bound + + _PATCH_ACTIVE = True + _log.warning( + "PTO GDN megakernel active (chunk_size=128, Ascend 910B)." + ) + + +def is_pto_gdn_patch_active() -> bool: + return _PATCH_ACTIVE diff --git a/vllm_fl/platform.py b/vllm_fl/platform.py index 535a8c9e2..2085766d1 100644 --- a/vllm_fl/platform.py +++ b/vllm_fl/platform.py @@ -206,6 +206,34 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: ) compilation_config.cudagraph_mode = CUDAGraphMode.NONE + # -------------------------------------------------------- + # NPU specific: inductor is not supported by default, but keep + # cudagraph mode unchanged so FULL_DECODE_ONLY can be profiled. + # If ascend_compilation_config.enable_npugraph_ex is set, leave the + # backend alone so that the AscendCompiler can be used. + if cls.device_type == "npu": + ascend_compilation_config = ( + vllm_config.additional_config or {} + ).get("ascend_compilation_config", {}) + enable_npugraph_ex = ascend_compilation_config.get( + "enable_npugraph_ex", False) + + # The default post-grad fusion passes are CUDA-specific and will + # reference undefined pass classes on NPU. Disable them here. + compilation_config.pass_config.fuse_norm_quant = False + compilation_config.pass_config.fuse_act_quant = False + compilation_config.pass_config.fuse_attn_quant = False + compilation_config.pass_config.fuse_allreduce_rms = False + + backend = getattr(compilation_config, "backend", "") + if backend in ("", "inductor") and not enable_npugraph_ex: + compilation_config.backend = "eager" + logger.warning( + "NPU does not support torch inductor compilation. " + "Switching backend from '%s' to 'eager'.", + backend or "", + ) + # -------------------------------------------------------- # maca specific config updates if cls.vendor_name == "metax": @@ -296,6 +324,45 @@ def get_device_communicator_cls(cls) -> str: def get_static_graph_wrapper_cls(cls) -> str: return "vllm_fl.compilation.graph.GraphWrapper" + @classmethod + def get_pass_manager_cls(cls) -> str: + # Only use the Ascend graph-fusion pass manager when npugraph_ex is + # explicitly enabled; otherwise keep the upstream default so that + # ordinary graph-mode serving works with the installed vllm version. + if cls.device_type == "npu" and cls._is_npugraph_ex_enabled(): + return "vllm_fl.dispatch.backends.vendor.ascend.compilation.graph_fusion_pass_manager.GraphFusionPassManager" + return "vllm.compilation.pass_manager.PostGradPassManager" + + @classmethod + def get_compile_backend(cls) -> str: + if cls.device_type == "npu" and cls._is_npugraph_ex_enabled(): + return "vllm_fl.dispatch.backends.vendor.ascend.compilation.compiler_interface.AscendCompiler" + return "vllm.compilation.compiler_interface.EagerAdaptor" + + @classmethod + def _is_npugraph_ex_enabled(cls) -> bool: + """Check whether npugraph_ex compilation is explicitly enabled. + + npugraph_ex is an advanced Ascend compilation path that requires + compatible vllm/torchair versions. It is disabled by default and must + be turned on explicitly via the environment variable + VLLM_FL_ENABLE_NPUGRAPH_EX=1 or via + additional_config.ascend_compilation_config.enable_npugraph_ex=true. + """ + import os + if os.environ.get("VLLM_FL_ENABLE_NPUGRAPH_EX", "0") == "1": + return True + try: + from vllm.config import get_current_vllm_config + vllm_config = get_current_vllm_config() + if vllm_config is None: + return False + ascend_cfg = (vllm_config.additional_config or {}).get( + "ascend_compilation_config", {}) + return ascend_cfg.get("enable_npugraph_ex", False) + except Exception: + return False + @classmethod def support_static_graph_mode(cls) -> bool: if cls.vendor_name in ["nvidia", "ascend", "metax", "kunlunxin", "hygon", "txda"]: diff --git a/vllm_fl/profiler/__init__.py b/vllm_fl/profiler/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_fl/profiler/ascend.py b/vllm_fl/profiler/ascend.py new file mode 100644 index 000000000..952848c47 --- /dev/null +++ b/vllm_fl/profiler/ascend.py @@ -0,0 +1,119 @@ +# Copyright (c) 2025 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0. +"""Ascend NPU profiler wrapper.""" + +from __future__ import annotations + +import os + +import torch +from typing_extensions import override + +from vllm.config import ProfilerConfig +from vllm.logger import init_logger +from vllm.profiler.wrapper import WorkerProfiler + +logger = init_logger(__name__) + + +class AscendTorchProfilerWrapper(WorkerProfiler): + """WorkerProfiler backed by torch_npu.profiler.profile.""" + + def __init__( + self, + profiler_config: ProfilerConfig, + worker_name: str, + local_rank: int, + ) -> None: + super().__init__(profiler_config) + self.local_rank = local_rank + self.worker_name = worker_name + self.profiler_config = profiler_config + self.profiler = None + + if profiler_config.profiler != "torch": + raise RuntimeError( + f"Unrecognized profiler: {profiler_config.profiler}" + ) + if not profiler_config.torch_profiler_dir: + raise RuntimeError("torch_profiler_dir cannot be empty.") + + import torch_npu.profiler # type: ignore + + self.torch_npu_profiler = torch_npu.profiler + self.experimental_config = torch_npu.profiler._ExperimentalConfig( + export_type=torch_npu.profiler.ExportType.Text, + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + msprof_tx=False, + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, + l2_cache=False, + op_attr=False, + data_simplification=True, + record_op_args=False, + gc_detect_threshold=None, + ) + + if local_rank in (None, 0): + logger.info_once( + "Ascend NPU profiling enabled. Trace root: %s", + profiler_config.torch_profiler_dir, + scope="local", + ) + + def _case_profiler_dir(self) -> str: + root = self.profiler_config.torch_profiler_dir + marker = os.path.join(root, ".current_case_dir") + + try: + case_dir = open(marker, encoding="utf-8").read().strip() + except OSError: + case_dir = "" + + if not case_dir: + return root + if not os.path.isabs(case_dir): + case_dir = os.path.join(root, case_dir) + + os.makedirs(case_dir, exist_ok=True) + return case_dir + + def _create_profiler(self): + trace_dir = self._case_profiler_dir() + + if self.local_rank in (None, 0): + logger.info( + "Ascend NPU profiling trace dir: %s", + trace_dir, + ) + + return self.torch_npu_profiler.profile( + activities=[ + self.torch_npu_profiler.ProfilerActivity.CPU, + self.torch_npu_profiler.ProfilerActivity.NPU, + ], + with_stack=False, + profile_memory=self.profiler_config.torch_profiler_with_memory, + with_modules=self.profiler_config.torch_profiler_with_stack, + experimental_config=self.experimental_config, + on_trace_ready=self.torch_npu_profiler.tensorboard_trace_handler( + trace_dir, + worker_name=self.worker_name, + ), + ) + + @override + def _start(self) -> None: + self.profiler = self._create_profiler() + self.profiler.start() + + @override + def _stop(self) -> None: + if self.profiler is None: + return + self.profiler.stop() + self.profiler = None + + @override + def annotate_context_manager(self, name: str): + return torch.profiler.record_function(name) diff --git a/vllm_fl/utils.py b/vllm_fl/utils.py index 54cb4b55c..9c4547bf9 100644 --- a/vllm_fl/utils.py +++ b/vllm_fl/utils.py @@ -1,18 +1,23 @@ # Copyright (c) 2025 BAAI. All rights reserved. -# -# 2026 - Modified by Kunlunxin, Inc. All Rights Reserved. import json import os +import re +import subprocess +from pathlib import Path from typing import Optional, Tuple import flag_gems -from flag_gems.runtime.backend.device import DeviceDetector +try: + # FlagGems<=5.0.2: DeviceDetector lives in device. + from flag_gems.runtime.backend.device import DeviceDetector +except (ImportError, FileNotFoundError): + # FlagGems>5.0.2: DeviceDetector lives in device_finder. + from flag_gems.runtime.backend.device_finder import DeviceDetector from flag_gems.runtime import backend _OP_CONFIG: Optional[dict[str, str]] = None - # Mapping used by dispatch registration to resolve the current runtime platform # into a backend directory under dispatch/backends/vendor. # @@ -27,7 +32,7 @@ # Source: runtime platform detection (current_platform.device_name). # # Values are normalized to lowercase and matched against available backend -# subdirectories (for example, cuda/ascend/metax/iluvatar/mthreads/kunlunxin). +# subdirectories (for example, cuda/ascend/metax/iluvatar/mthreads). VENDOR_DEVICE_MAP: dict[str, dict[str, str]] = { # Registered backend: vendor/cuda "nvidia": {"device_type": "cuda", "device_name": "nvidia"}, @@ -39,12 +44,14 @@ "metax": {"device_type": "cuda", "device_name": "metax"}, # Registered backend: vendor/musa "mthreads": {"device_type": "musa", "device_name": "musa"}, - # Registered backend: vendor/kunlunxin - "kunlunxin": {"device_type": "cuda", "device_name": "kunlunxin"}, + # Registered backend: vendor/sunrise + "sunrise": {"device_type": "ptpu", "device_name": "ptpu"}, # Registered backend: vendor/hygon "hygon": {"device_type": "cuda", "device_name": "cuda"}, # Registered backend: vendor/txda "tsingmicro": {"device_type": "txda", "device_name": "txda"}, + # Registered backend: vendor/thead (PPU) + "thead": {"device_type": "cuda", "device_name": "thead"}, } @@ -213,7 +220,7 @@ def get_op_config() -> Optional[dict[str, str]]: class DeviceInfo: def __init__(self): self.device = DeviceDetector() - self.supported_device = ["nvidia", "ascend", "metax", "mthreads", "kunlunxin", "hygon"] + self.supported_device = ["nvidia", "ascend", "metax", "mthreads", "sunrise", "thead"] backend.set_torch_backend_device_fn(self.device.vendor_name) @property @@ -337,5 +344,85 @@ def is_oot_enabled() -> bool: return enabled_str.lower() in ("1", "true") +_CUSTOM_OP_ENABLED = None + + +def enable_custom_op() -> bool: + """ + Enable vllm-plugin-FL CANN custom ops deployed under vllm_fl/_cann_ops_custom. + + Mirrors vllm_ascend.utils.enable_custom_op: sources the generated + set_env.bash (or falls back to setting ASCEND_CUSTOM_OPP_PATH and + LD_LIBRARY_PATH directly) so that the aclnn custom-op symbols can be + discovered at runtime. + + The absolute paths baked into set_env.bash during installation are not + trusted blindly, because the installed package may have been moved to a + different prefix. We always override the two critical variables with + paths derived from the actual vllm_fl package location. + + Returns: + True if custom-op environment was configured, False otherwise. + """ + global _CUSTOM_OP_ENABLED + if _CUSTOM_OP_ENABLED is not None: + return _CUSTOM_OP_ENABLED + + custom_ops_root = Path(__file__).parent / "_cann_ops_custom" + if not custom_ops_root.is_dir(): + _CUSTOM_OP_ENABLED = False + return _CUSTOM_OP_ENABLED + + vendor_dir = custom_ops_root / "vendors" / "custom_transformer" + set_env_script = vendor_dir / "bin" / "set_env.bash" + + if set_env_script.is_file(): + try: + # Source the script in a subshell and capture exported variables. + output = subprocess.check_output( + ["bash", "-c", f"source '{set_env_script}' && env"], + text=True, + stderr=subprocess.DEVNULL, + ) + for line in output.splitlines(): + if "=" not in line: + continue + key, _, value = line.partition("=") + # Preserve any extra environment variables exported by the + # script, but do NOT keep ASCEND_CUSTOM_OPP_PATH or + # LD_LIBRARY_PATH from set_env.bash: they contain absolute + # paths that become stale when the package is relocated. + if key and key not in ( + "ASCEND_CUSTOM_OPP_PATH", + "LD_LIBRARY_PATH", + ): + os.environ[key] = value + except subprocess.CalledProcessError: + # Fall back to setting the two known variables directly. + pass + + # Always use paths derived from the actual installed location. + _set_custom_op_env_paths(vendor_dir) + + _CUSTOM_OP_ENABLED = True + return _CUSTOM_OP_ENABLED + + +def _set_custom_op_env_paths(vendor_dir: Path) -> None: + """Set ASCEND_CUSTOM_OPP_PATH and LD_LIBRARY_PATH for the custom op package.""" + vendor_str = str(vendor_dir) + lib_dir = str(vendor_dir / "op_api" / "lib") + + old_opp_path = os.environ.get("ASCEND_CUSTOM_OPP_PATH", "") + os.environ["ASCEND_CUSTOM_OPP_PATH"] = ( + f"{vendor_str}:{old_opp_path}" if old_opp_path else vendor_str + ) + + old_ld_path = os.environ.get("LD_LIBRARY_PATH", "") + os.environ["LD_LIBRARY_PATH"] = ( + f"{lib_dir}:{old_ld_path}" if old_ld_path else lib_dir + ) + + if __name__ == "__main__": device = DeviceInfo() diff --git a/vllm_fl/worker/model_runner.py b/vllm_fl/worker/model_runner.py index 6720cbcbb..e215bc324 100644 --- a/vllm_fl/worker/model_runner.py +++ b/vllm_fl/worker/model_runner.py @@ -53,6 +53,7 @@ ) from vllm.forward_context import ( BatchDescriptor, + get_forward_context, set_forward_context, ) from vllm.logger import init_logger @@ -99,7 +100,8 @@ from vllm.utils.platform_utils import is_pin_memory_available from vllm.platforms import current_platform -if current_platform.dist_backend == "flagcx": + +if current_platform.dist_backend in ("flagcx", "hccl"): @contextmanager def graph_capture(device: torch.device): """ @@ -602,6 +604,7 @@ def __init__( # Cudagraph dispatcher for runtime cudagraph dispatching. self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config) + self.aclgraph_update_stream = None self.mm_budget = ( MultiModalBudget( @@ -1621,7 +1624,16 @@ def _get_block_table_and_slot_mapping(kv_cache_gid: int): # Fill unused with -1. Needed for reshape_and_cache in full cuda # graph mode. `blk_table_tensor` -1 to match mamba PAD_SLOT_ID slot_mapping[num_tokens:num_tokens_padded].fill_(-1) - blk_table_tensor[num_reqs:num_reqs_padded].fill_(-1) + # Pad block-table rows with block id 0 rather than -1: when a + # decode batch shrinks below the captured size (e.g. the first + # request finishes at max concurrency), the full-attention FIA + # kernel receives the padded rows and dereferences the block ids + # even for seq_len==0 rows, and block id -1 makes it fault + # (fftsplus aicore error / CCU instruction address check error, + # surfacing as aclrtSynchronizeEvent 507011). The mamba/GDN pad + # slots keep their PAD_SLOT_ID=-1 convention via the metadata + # builder's own fills, so this only affects the attention path. + blk_table_tensor[num_reqs:num_reqs_padded].fill_(0) return blk_table_tensor, slot_mapping @@ -2078,6 +2090,7 @@ def _calc_spec_decode_metadata( draft_token_ids=draft_token_ids, num_draft_tokens=num_draft_tokens.tolist(), cu_num_draft_tokens=cu_num_draft_tokens, + cu_num_sampled_tokens=cu_num_sampled_tokens, target_logits_indices=target_logits_indices, bonus_logits_indices=bonus_logits_indices, logits_indices=logits_indices, @@ -3129,6 +3142,23 @@ def execute_model( inputs_embeds=inputs_embeds, **model_kwargs, ) + if ( + current_platform.device_type == "npu" + and cudagraph_mode == CUDAGraphMode.FULL + and self.aclgraph_update_stream is not None + ): + from vllm_fl.dispatch.backends.vendor.ascend.patches.patch_graph import ( + update_full_graph_params, + ) + + update_full_graph_params( + [group.backend for group in self._attn_group_iterator()], + self.aclgraph_update_stream, + get_forward_context(), + num_tokens_padded, + self.vllm_config, + self.speculative_config, + ) with record_function_or_nullcontext("gpu_model_runner: postprocess"): if self.use_aux_hidden_state_outputs: @@ -3768,6 +3798,8 @@ def load_model(self, eep_scale_up: bool = False) -> None: cudagraph_mode = self.compilation_config.cudagraph_mode assert cudagraph_mode is not None if cudagraph_mode.has_full_cudagraphs() and not self.parallel_config.enable_dbo: + if current_platform.device_type == "npu": + self.aclgraph_update_stream = torch.npu.Stream() self.model = GraphWrapper( self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL ) @@ -4165,6 +4197,24 @@ def _dummy_run( # shorter sequence lengths to run faster. # TODO(luka) better system for describing dummy batches seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] + elif ( + current_platform.device_type == "npu" + and is_graph_capturing + and cudagraph_runtime_mode == CUDAGraphMode.FULL + ): + # _npu_paged_attention only returns its maximum workspace at a + # sufficiently large context length. Capturing with decode + # query length 1 works until the runtime context crosses an + # internal tiling boundary, where graph-task update then fails + # with CANN 507000. This matches vllm-ascend's FULL graph + # capture strategy; inference still updates the task with the + # real per-request sequence lengths. + seq_lens = min(6144, self.max_model_len) + logger.info_once( + "Using seq_len=%d to capture the maximum Ascend " + "paged-attention workspace", + seq_lens, + ) else: seq_lens = max_query_len # type: ignore[assignment] self.seq_lens.np[:num_reqs] = seq_lens @@ -4969,6 +5019,15 @@ def _check_and_update_cudagraph_mode( # Trigger cudagraph dispatching keys initialization after # resolved cudagraph mode. self.compilation_config.cudagraph_mode = cudagraph_mode + if ( + current_platform.device_type == "npu" + and cudagraph_mode.has_full_cudagraphs() + ): + from vllm_fl.dispatch.backends.vendor.ascend.patches.patch_graph import ( + ensure_graph_params, + ) + + ensure_graph_params(list(self.cudagraph_batch_sizes)) self.cudagraph_dispatcher.initialize_cudagraph_keys( cudagraph_mode, self.uniform_decode_query_len ) @@ -5298,28 +5357,21 @@ def _update_hybrid_attention_mamba_layout( self, kv_caches: dict[str, torch.Tensor] ) -> None: """ - Update the layout of attention layers from (2, num_blocks, ...) to - (num_blocks, 2, ...). - - Args: - kv_caches: The KV cache buffer of each layer. + Keep the attention KV cache in the contiguous (2, num_blocks, ...) + layout. + + Upstream vLLM re-strides the cache into an interleaved + (num_blocks, 2, ...) layout here, which makes ``kv_cache[0]`` / + ``kv_cache[1]`` non-contiguous views. The Ascend attention impl + (``vllm_fl/dispatch/backends/vendor/ascend/impl/attention.py``) + requires contiguous k/v caches and was paying one full-cache + ``.contiguous()`` copy per attention layer per step (~1ms per copy + at typical cache sizes, seen as the ``aclnnInplaceCopy_SliceAiCore`` + hotspot in profiles). Skipping the re-stride keeps the views + contiguous; all accesses go through the same views, so this is + semantically transparent. """ - - for group in self._kv_cache_spec_attn_group_iterator(): - kv_cache_spec = group.kv_cache_spec - for layer_name in group.layer_names: - kv_cache = kv_caches[layer_name] - if isinstance(kv_cache_spec, AttentionSpec) and kv_cache.shape[0] == 2: - assert kv_cache.shape[1] != 2, ( - "Fail to determine whether the layout is " - "(2, num_blocks, ...) or (num_blocks, 2, ...) for " - f"a tensor of shape {kv_cache.shape}" - ) - hidden_size = kv_cache.shape[2:].numel() - kv_cache.as_strided_( - size=kv_cache.shape, - stride=(hidden_size, 2 * hidden_size, *kv_cache.stride()[2:]), - ) + return def initialize_kv_cache_tensors( self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] diff --git a/vllm_fl/worker/worker.py b/vllm_fl/worker/worker.py index 1d1358ee9..9e9c55d57 100644 --- a/vllm_fl/worker/worker.py +++ b/vllm_fl/worker/worker.py @@ -221,12 +221,26 @@ def __init__( profiler_config = vllm_config.profiler_config if profiler_config.profiler == "torch": worker_name = f"{vllm_config.instance_id}-rank-{self.rank}" - self.profiler = TorchProfilerWrapper( - profiler_config, - worker_name=worker_name, - local_rank=self.local_rank, - activities=["CPU", "CUDA"], - ) + try: + import torch_npu # noqa: F401 + is_npu = torch.npu.is_available() + except ImportError: + is_npu = False + + if is_npu: + from vllm_fl.profiler.ascend import AscendTorchProfilerWrapper + self.profiler = AscendTorchProfilerWrapper( + profiler_config, + worker_name=worker_name, + local_rank=self.local_rank, + ) + else: + self.profiler = TorchProfilerWrapper( + profiler_config, + worker_name=worker_name, + local_rank=self.local_rank, + activities=["CPU", "CUDA"], + ) else: self.profiler = None